介紹 Flash-KMeans：一款感知 I/O 的精確 K-Means，在 GPU 上比 FAISS 快超過 200 倍

k-means has been an offline tool for decades. You run it once to preprocess data, then move on. A team of researchers from UC Berkeley and UT Austin released Flash-KMeans, a new open-source library that targets a different setting. Modern AI pipelines now call k-means inside training and inference loops. At that frequency, latency per call matters more than theoretical FLOPs. Flash-KMeans is an IO-aware implementation of standard Lloyd’s k-means. It does not change the math, and it does not approximate. It only restructures how the algorithm moves data on a GPU. On an NVIDIA H200, the research team reported up to 17.9× end-to-end speedup over the best baseline. Against NVIDIA cuML they report 33×. Against FAISS they report over 200×. What is Flash-KMeans Flash-KMeans is a batched k-means library written in Triton GPU kernels. It ships under Apache 2.0 and installs with pip install flash-kmeans. The output is mathematically identical to standard Lloyd’s k-means. The speedup comes from kernel-level dataflow, not from skipping work. That separates it from algorithmic methods like triangle-inequality pruning or coreset sampling. A standard Lloyd iteration has two stages. The assignment stage computes each point’s distance to every centroid, then picks the nearest. The update stage averages the points in each cluster to form new centroids. Both stages are simple arithmetic. On GPUs, both are bottlenecked by memory, not compute. The Two Bottlenecks It Attacks The first bottleneck is the assignment stage. Standard code builds a full distance matrix D of shape N×K in High Bandwidth Memory (HBM). It writes the matrix, then reads it back to run argmin. For N=65536, K=1024, d=128, B=32, the distance math takes 2.6ms. Writing and consuming D takes about 23ms. The matrix is the cost, not the arithmetic. Flash-KMeans replaces this with FlashAssign. The design borrows from FlashAttention. FlashAssign streams tiles of points and centroids from HBM into on-chip SRAM. It fuses distance computation with an online argmin. The full N×K matrix is never materialized. This cuts the dominant IO complexity from O(NK) to O(Nd + Kd). At the kernel level, FlashAssign reaches up to 21.2×. In one case it cut assignment from 122.5ms to 5.8ms. The second bottleneck is the centroid update stage. Standard code uses scatter-style atomic adds. Each thread adds its point into a shared sum buffer keyed by cluster id. Many threads hit the same ‘hot’ cluster at once. That causes atomic contention and hardware serialization. The research team measured only 50 GB/s effective bandwidth here on an H200. Flash-KMeans replaces this with Sort-Inverse Update. It sorts the 1D assignment vector by cluster id using argsort. Identical cluster ids then form contiguous segments. Each thread block reduces a segment on-chip, then issues one atomic add per segment. The heavy point matrix is never physically permuted. Atomic operations drop from (O((K+NBN)d))(O((K + \frac{N}{B_N})d)) . The kernel reaches up to 6.3×. Benchmark The research team test it on an H200 with CUDA 12.8, FP16 data, and d=128. They sweep N, K, and batch size B. They compare against four optimized baselines: fast_pytorch_kmeans, fastkmeans, cuML, and FAISS. ComparisonReported speedupWorkload contextEnd-to-end vs best baselineup to 17.9×N=8M, K=1024 (large N, small K)vs NVIDIA cuML33×industry libraryvs FAISSover 200×industry libraryFlashAssign kernelup to 21.2×N=1M, K=8192 (assignment)Sort-Inverse Update kernelup to 6.3×N=33M, K=4096 (update)Out-of-core, large scaleup to 10.5×N=400M, K=16384 vs fastkmeans One failure mode matters for context. Standard PyTorch implementations run out of memory in large-K regimes. They cannot materialize the N×K matrix. FAISS is the industry-standard library under many production vector-search systems. The library also runs out-of-core. On one billion points (K=32768, d=128), it finishes an iteration in 41.4s, against 261.8s for the baseline. It uses chunked stream overlap to hide PCIe transfer behind compute. A cache-aware compile heuristic also cuts tuning overhead by up to 175×, within 0.3% of tuned speed. MTP Interactive Explainer #mtp-fk-demo *{box-sizing:border-box!important;margin:0;padding:0} #mtp-fk-demo{ background:#111!important;color:#e8e8e8!important; border:1px solid #222!important;border-radius:14px!important; font-family:-apple-system,BlinkMacSystemFont,"Segoe UI",Roboto,Helvetica,Arial,sans-serif!important; max-width:980px;margin:24px auto!important;padding:0!important;overflow:hidden!important; line-height:1.5!important } #mtp-fk-demo hr,#mtp-fk-demo p:empty,#mtp-fk-demo del,#mtp-fk-demo s{display:none!important} #mtp-fk-demo .fk-head{padding:22px 24px 14px!important;border-bottom:1px solid #1d1d1d!important} #mtp-fk-demo .fk-eyebrow{color:#76B900!important;font-size:11px!important;letter-spacing:.16em!important;text-transform:uppercase!important;font-weight:700!important} #mtp-fk-demo h2.fk-title{color:#fff!important;font-size:23px!important;font-weight:800!important;margin:6px 0 4px!important;letter-spacing:-.01em!important} #mtp-fk-demo .fk-sub{color:#9a9a9a!important;font-size:13.5px!important;max-width:680px!important} #mtp-fk-demo .fk-stats{display:grid!important;grid-template-columns:repeat(4,1fr)!important;gap:1px!important;background:#1d1d1d!important;border-bottom:1px solid #1d1d1d!important} #mtp-fk-demo .fk-stat{background:#0c0c0c!important;padding:14px 16px!important} #mtp-fk-demo .fk-stat b{color:#76B900!important;font-size:21px!important;font-weight:800!important;display:block!important;font-variant-numeric:tabular-nums!important} #mtp-fk-demo .fk-stat span{color:#8a8a8a!important;font-size:11px!important;display:block!important;margin-top:2px!important} #mtp-fk-demo .fk-tabs{display:flex!important;gap:6px!important;padding:14px 24px 0!important;flex-wrap:wrap!important} #mtp-fk-demo .fk-tab{ background:#181818!important;color:#bdbdbd!important;border:1px solid #262626!important; border-radius:8px 8px 0 0!important;padding:9px 15px!important;font-size:13px!important;font-weight:600!important; cursor:pointer!important;transition:all .15s!important } #mtp-fk-demo .fk-tab:hover{color:#fff!important;border-color:#3a3a3a!important} #mtp-fk-demo .fk-tab.on{background:#0c0c0c!important;color:#76B900!important;border-color:#76B900!important;border-bottom-color:#0c0c0c!important} #mtp-fk-demo .fk-panel{display:none!important;padding:18px 24px 22px!important} #mtp-fk-demo .fk-panel.on{display:block!important} #mtp-fk-demo .fk-row{display:flex!important;gap:18px!important;flex-wrap:wrap!important;align-items:flex-start!important} #mtp-fk-demo .fk-canvaswrap{flex:1 1 360px!important;min-width:280px!important} #mtp-fk-demo canvas{width:100%!important;display:block!important;background:#0a0a0a!important;border:1px solid #1d1d1d!important;border-radius:10px!important} #mtp-fk-demo .fk-side{flex:1 1 220px!important;min-width:200px!important} #mtp-fk-demo .fk-ctl{margin-bottom:14px!important} #mtp-fk-demo .fk-ctl label{display:flex!important;justify-content:space-between!important;font-size:12px!important;color:#bdbdbd!important;margin-bottom:6px!important} #mtp-fk-demo .fk-ctl label em{color:#76B900!important;font-style:normal!important;font-weight:700!important;font-variant-numeric:tabular-nums!important} #mtp-fk-demo input[type=range]{-webkit-appearance:none;appearance:none;width:100%!important;height:5px!important;background:#262626!important;border-radius:4px!important;outline:none!important} #mtp-fk-demo input[type=range]::-webkit-slider-thumb{-webkit-appearance:none;appearance:none;width:16px;height:16px;border-radius:50%;background:#76B900;cursor:pointer;border:2px solid #0c0c0c} #mtp-fk-demo input[type=range]::-moz-range-thumb{width:16px;height:16px;border-radius:50%;background:#76B900;cursor:pointer;border:2px solid #0c0c0c} #mtp-fk-demo .fk-btns{display:flex!important;gap:8px!important;flex-wrap:wrap!important;margin:4px 0 12px!important} #mtp-fk-demo button.fk-b{ background:#7