Flash-KMeans：快速且内存高效的精确 K-Means，可在单张 GPU 进行亿级数据的聚类

在当前的人工智能领域LLM 及其生成能力几乎独占了所有焦点。但再精密的 RAG Pipeline能力上限也取决于那个沉默的引擎搜索与聚类层。聚类不只是一项经典的数据科学任务它是组织高维向量空间的核心机制——让 LLM 能在数十亿条文档和参数的海洋中定位正确的记忆。随着数据集规模持续扩大沿用数十年的标准算法已经撞上了墙。本文介绍 Flash-KMeans是一个近期提出的框架它受 Flash最小化数据移动的启发论文给出了一种执行精确 K-Means 的方案速度更快内存效率也远优于 FAISS 等行业标准实现。这不只是性能上的边际改进而是在聚类过程中处理数据物化materialization方式上的根本性转变。参数化方法与非参数化方法在介绍现代检索机制之前需要先回到统计学习的视角参数化Parametric模型与非参数化Non-parametric模型。参数化方法此类模型假设底层数据服从某种固定的函数形式或分布由有限个参数定义例如线性回归、逻辑回归。模型复杂度是常数无论数据量多大都不会增长。非参数化方法K-近邻K-Nearest NeighborsKNN等方法不做此类假设。数据点本身就是参数随着数据集扩大模型捕捉复杂非线性模式的能力也随之增长。智能即对称与压缩从根本上说智能是在混沌中发现对称、在不失本质的前提下压缩信息的能力。原始数据具有高熵且充满噪声智能的运作方式是识别不同数据点之间反复出现的模式与对称性并将其映射为统一的概念。这正是聚类不可或缺的原因——通过将相似信息归为一组把庞大、难以管理的数据集转化为一个压缩后的代表性中心点centroid集合。每一次聚类都是一种有损压缩保留了最关键的语义关系。若缺乏分类和凝练的能力系统就会被迫将每次新的经历视为独特事件泛化推理也就无从谈起。聚类让人类得以通过抽象概念而非原始比特来认知世界。KNN 究竟如何工作要理解 Flash-KMeans 的优化之处必须先了解其前驱算法 K-近邻KNN的原始运作机制。K-Means 是一个迭代聚类过程每次迭代的核心都依赖最近邻搜索将数据点分配到最近的中心点。算法流程KNN 的逻辑在数学上十分优雅但计算开销极大。对于每个查询点算法执行以下步骤距离计算计算查询点与数据集中每个点之间的距离通常为欧氏距离或余弦距离。排序将这些距离按升序排列。选择找出距离最小的K个点。决策在分类任务中对这K个邻居进行多数投票在聚类任务中将该点分配到对应的邻域。时间与内存复杂度分析在研究环境中算法的评估标准是随数据增长的扩展性。设N为数据点数量D为维度特征数量KNN 是非参数化方法不会学习简化模型而是在数据集的全部权重空间中搜索。搜索时计算到每个点的距离O(N*D)会产生巨大开销。内存则需求很大为了在 GPU 上执行这些计算通常需要将中间距离矩阵和聚类分配结果物化存储在内存中。随着 N 和 D 增大VRAM 最终会耗尽导致崩溃或性能骤降——这正是 Flash-KMeans 要解决的问题。内存分配与速度问题KNN 和 K-Means 的理论基础是扎实的但在 GPU 上的实现却撞上了物理壁垒。在高性能计算中瓶颈往往不是原始浮点运算FLOPs而是硬件管理数据流动的方式。分配阶段的 IO 瓶颈标准分配阶段GPU 需要计算每个数据点与每个中心点之间的距离传统做法是在 GPU 全局内存中物化一个巨大的距离矩阵。全局内存VRAM比 GPU 片上 SRAM 慢得多系统因此变为 IO 受限IO-bound——处理核心空转等待距离数据从慢速 VRAM 移动到快速 SRAM。随着 K聚类数或 N数据点数的增加数据移动量急剧增大吞吐量被严重拖累。更新阶段的原子写入竞争分配完成后需要对所有分配到同一中心点的数据点坐标求平均以计算新中心点坐标。问题在于多个 GPU 线程各自处理不同的数据点可能同时尝试更新同一个中心点。这会产生原子写入竞争atomic write contention——为防止数据损坏硬件必须将这些写操作串行化或使用原子加法操作当数百个线程争抢同一内存地址时速度会极其低下。系统级约束除算法本身外还面临硬性的系统级限制VRAM 上限若数据集或距离矩阵超过 GPU 内存例如 H100 的 80GB进程会完全失败或退而依赖系统内存RAM而后者的速度慢了几个数量级。实时系统的延迟在大规模推荐引擎或自动欺诈检测系统等生产场景中聚类必须在毫秒级完成才能支撑实时推断。为何 KNN 必须足够快理论上O(N*D) 的复杂度在小规模下还算可以接受。但在高频交易、实时推荐引擎或大规模法律文档分析等生产场景中这是灾难性的瓶颈。要理解为何如此执着于优化需要先了解当前如何对数据进行分区以高效检索。空间分区Voronoi 与 Delaunay为避免将查询与每个点逐一比较可以利用几何结构来压缩搜索空间Voronoi 图根据到一组种子中心点的距离将空间划分为单元格。每个单元格包含距其种子最近的所有点从而能立即缩小搜索范围。Delaunay 三角剖分Voronoi 图的对偶结构。将共享边界的中心点相连形成一个图从而能以数学上高效的方式从一个区域跳转到另一个区域。在高维向量空间中数据通过邻域划分来组织依赖这两种相互交织的结构Voronoi 图国家层级表示节点层级的主权。每个中心点“首都”定义一条边界。若查询落在法国边界内搜索范围立即缩小到该节点无需遍历世界其他地方。Delaunay 三角剖分菜系层级表示边层级的连通性。这些是相邻国家特征之间的桥梁可以是菜系、最受欢迎的运动、GDP 等任意特征。假设已到达法国节点但查询的是地中海意面Voronoi 边界无法引导移动方向此时需要看边Edges。由于法国和意大利共享边界它们之间有一条 Delaunay 边代表一个共同特征例如菜系。这种双系统的威力在于 Next Hop 逻辑利用 Voronoi 单元格确定当前位置例如法国。使用 Delaunay 查看可以去往何处。若搜索的期望条件是食物边会引导到意大利。跳过德国因为虽然它是邻居但不共享菜系这条边。标准 K-Means 算法中每次迭代都依赖 Voronoi 逻辑运作。在分配阶段算法必须确定每个数据点属于哪个国家Voronoi 单元格。若有十亿个点和 10,000 个中心点算法就必须测量每个点到每个首都的距离产生巨大的 O(NK) 计算开销。标准 K-Means 本质上就是反复重绘这些国家边界Voronoi 单元格直到首都中心点完全居中于各自领土的过程。*检索的核心HNSW现代向量数据库依赖层级可导航小世界Hierarchical Navigable Small WorldsHNSW执行近似最近邻搜索。HNSW 可以理解为一个多层图——高维空间的跳表skip-list。上层节点较少拥有长距离连接用于快速遍历下层提供更细粒度的密度。在每一层算法执行局部 KNN 搜索以找到下一个入口点。若底层 KNN 速度慢整个图遍历就会滞后RAG 整个检索组件的性能随之下降。从确定性到变分推断KNN 方法本质上是变分推断Variational Inference的离散化变体。KNN 是确定性的硬边界一个点要么属于某个聚类要么不属于是一种严格的几何分配。高斯混合模型Gaussian Mixture ModelsGMMs则是概率性的软边界将世界表示为重叠的分布一个点以某个概率属于多个聚类通过期望最大化Expectation-MaximizationEM算法进行优化。EM 本质上是变分推断VI的一个离散化特例。VI 是几乎所有生成模型包括 VAE 以及现代 LLM 的损失函数的数学核心因此简单的 K-Means 分配实际上是当今最先进 AI 所使用的复杂概率推理的简化硬版本。KNN 必须足够快因为它是理解信息分布的基础单元。理想情形要设计更好的系统必须先精确理解标准实现在哪里失败。FAISS、scikit-learn 等库中的标准 GPU 加速 K-Means 建立在 Lloyd 算法之上该算法在两个不同阶段之间迭代分配Assignment和更新Update。Lloyd 算法与标准 GPU 实现分配阶段使用距离度量通常为欧氏距离将每个数据点与所有K个中心点进行比较找出最近邻。更新阶段取分配到同一中心点的所有数据点的均值重新计算中心点坐标。数据物化标准实现中GPU 会将一个大型分配向量物化写入全局内存以追踪每个点属于哪个聚类。这就是 K-Means 目前的实现方式。现代 K-Means 实现基于 Lloyd 算法而该算法在 GPU 上的扩展性并不理想。Kernel 级瓶颈标准实现无法达到理想情形根源在于两个主要的 Kernel 级问题内存带宽IO 墙分配阶段本质上是 IO 受限的。标准实现花在将距离数据从慢速 VRAM 移动到 GPU 核心上的时间多于实际执行计算的时间。原子写入竞争更新阶段处理不同数据点的多个线程经常同时尝试更新同一个中心点例如算法中的 atomic_add。硬件被迫将这些写操作串行化形成巨大的队列阻塞整个 Pipeline。上述瓶颈在实际场景中体现为可扩展性的硬性限制VRAM 过度消耗存储大型中间分配矩阵意味着随着K增大内存占用线性增长往往超过顶级 A100 或 H100 GPU 的 80GB 限制。计算利用率低系统阻塞在等待内存 IO 上高性能 Tensor Core 或 CUDA Core 大量闲置执行时间线中出现低效的空泡。理想的实现应完全避免将中间分配结果写入全局内存将数据保持在最快的内存层——片上 SRAM只在最终将聚合结果提交到主内存。为弥合这一差距需要从通用的现成聚类库转向系统级协同设计——从根本上重写算法 Kernel以尊重硅片的物理层级结构优先利用片上速度而非全局容量。下一节将介绍 Flash-KMeans 如何通过两项关键架构创新将这些硬件瓶颈转化为并行优势。Flash-KMeansFlash-KMeans 的创新在于其算法-系统协同设计Algorithm-System Co-design重新设计了 K-Means 算法的执行流程以绕过 GPU 硬件的物理限制。该框架引入了两个关键组件FlashAssign 和 Sort-Inverse Update以及面向实际部署的系统级优化。下面逐一介绍论文中提及的各个组件。FlashAssign基于 Online Argmin 的无物化分配为解决传统实现中严重的高带宽内存High Bandwidth MemoryHBM流量开销——即物化巨大 N×K 距离矩阵所造成的 IO 墙——论文引入了 FlashAssign。FlashAssign 将距离计算与规约reduction融合为一个单一的流式过程确保完整的距离矩阵从不在内存中显式构建。Online ArgminFlashAssign 的核心是 Online Argmin 更新。传统做法是 GPU 计算所有距离、存储后再找最小值。FlashAssign 反其道而行运行状态对于每个数据点 x_iKernel 只在局部寄存器中维护两个运行状态当前最小距离m_i和对应的中心点索引a_i。分块扫描中心点按 Tile 进行扫描。对于每个 TileKernel 在片上计算局部距离找到该 Tile 内的局部最小值并与运行状态m_i, a_i比较更新。寄存器优先于 HBM在所有中心点 Tile 上重复此过程全局最小值完全在片上通过寄存器而非全局内存确定。延迟隐藏Tiling 与异步预取为确保 GPU 核心不因等待数据而空转FlashAssign 在数据点和中心点两个维度上采用二维 Tiling 策略。实现中使用双缓冲Double Buffering模式GPU 在对当前中心点 Tilet执行距离计算的同时异步从 HBM 将下一个 Tilet1预取到另一块片上缓冲区将内存延迟与计算重叠保持高吞吐量。整个过程遵循严格的流式逻辑初始化将片上状态初始化为 m ∞并预计算范数。流式循环顺序扫描各中心点 Tile。计算与更新在片上计算局部距离求 Tile 内局部最小值并通过 Online Argmin 逻辑更新全局运行状态。最终写入所有中心点处理完毕后只将最终分配结果 a 写入 HBM。通过融合上述步骤FlashAssign 从根本上改变了数据流动方式将主导性的 IO 复杂度从 O(NK) 降低到 O(Nd Kd)有效消除了困扰标准 FAISS 式实现的 2·Θ(NK) HBM 流量开销。Sort-Inverse Update低竞争的中心点聚合为解决中心点更新阶段严重的写入串行化问题论文提出了 Sort-Inverse Update。传统实现采用 Scatter 式更新直接将 Token 散射到对应的中心点当多线程同时更新同一中心点时会产生严重的原子竞争。Flash-KMeans 不再与不规则的映射关系对抗而是将 Token 到聚类的更新转换为聚类到 Token 的聚集gather。Argsort 操作框架首先对分配向量a执行argsort得到置换索引sorted_idx。逻辑分组这会生成一个已排序的聚类 ID 序列相同的聚类 ID 自然聚集到连续的段中。内存效率排序只作用于一维分配向量数据点矩阵 X 的物理顺序在内存中从不改变。段级本地聚合逆映射构建完成后实际规约从全局内存转移到 GPU 快速片上内存。每个线程块CTA处理排序序列的一段连续块并识别段边界每个段的局部求和与计数完全在寄存器或共享内存中累积CTA 只在段边界处向 HBM 发出全局atomic_add操作而非对每个 Token 都发出。复杂度与竞争分析这种重组方式大幅减少了所需的原子操作数量。标准更新中原子操作规模为 O(Nd)Sort-Inverse 设计中原子合并次数降低至(O((K \lceil N/B_N \rceil)d))理论上的减少直接转化为写入竞争的消除实现了无竞争的内存写入从而加速规约阶段。算法分解排序通过argsort(a)计算sorted_idx构建已排序的聚类 ID。初始化将中心点求和s与计数n初始化为零。识别段遍历排序序列的各块找出连续的相同聚类 ID。聚集与累积从原始顺序中收集 Token 特征在片上累积局部求和与计数。原子合并在每个段边界处向 HBM 发出一次atomic_add。最终更新按 c_k s_k / n_k 重新计算中心点。高效的算法-系统协同设计方法论的最后一层关注可部署性。高性能 Kernel 若无法扩展到单张 GPU 内存之外或每次数据规模变化都需数小时调优就毫无实用价值。论文通过两项协同设计策略解决了这些实际问题。通过分块流式重叠处理大规模数据数据集超出可用 VRAM 是实际系统中的常态。Flash-KMeans 采用流式模式将数据从 CPU主机分阶段传输到 GPU设备数据分割为若干块利用 CUDA 流在 GPU 处理当前块的分配和更新 Kernel 的同时将下一块异步拷贝到 GPU双缓冲确保 PCIe 总线与 GPU 核心同时保持活跃隐藏从系统内存到 GPU 的数据传输延迟。缓存感知的编译启发式高性能机器学习中一项重要的隐性开销是编译时间。不同的数据规模不同的 N、D 或 K通常需要穷举式自动调优来找到最佳 Kernel 配置。Flash-KMeans 使用缓存感知的启发式方法直接根据硬件物理特性具体为 L1 和 L2 Cache 大小选择配置即使跨越不同硬件架构或数据规模动态变化也能几乎即时达到接近最优的性能水平Fast Time-to-First-Run。得益于中心点分配和分块方面的所有优化Flash-KMeans 得以实现。下面看整个开发成果的效果如何。结果与成效Flash-KMeans 的性能以业界标准基准 FAISS 为主要对比对象——后者被广泛视为 GPU 加速向量搜索与聚类的最先进实现。评测与基准测试的硬件环境为 NVIDIA H20080GBGPU重点考量执行速度与内存占用两个维度。结果表明解决 IO 和竞争瓶颈之后Flash-KMeans 重新定义了精确 K-Means 的效率边界。端到端加速研究人员考察了三种典型场景这些场景历来是 GPU 实现的难点大 N 与大 K内存密集型这是OOM内存溢出危险区。对于 N1M、K64K 的工作负载标准 PyTorch 实现在尝试物化巨型距离矩阵时直接失败。Flash-KMeans 在此取得了最显著的绝对加速性能比最佳替代方案fastkmeans高出 5.4× 以上。大 N 与小 K计算密集型在超大规模数据集N8M上进行搜索时延迟通常由原始距离计算主导。Flash-KMeans 将端到端延迟降低了 94.4%相比传统实现达到 17.9× 加速。小 N 与小 K框架开销场景即使在较小的高度批处理场景B32下——框架和 Kernel 启动开销通常会掩盖算法增益——Flash-KMeans 仍保持稳健实现高达 15.3× 的加速比。Kernel 级效率为验证加速效果确实源于消除特定硬件瓶颈研究人员考察了两个核心阶段各自的性能。FlashAssign分配阶段在高压配置N1M、K8192下执行时间从 122.5ms 压缩至仅 5.8ms相对于标准分配方法达到 21.2× 加速证实了无物化的 Online Argmin 方法有效绕过了 HBM 流量开销。Sort-Inverse Update规约阶段在大规模工作负载N33M下通过将无序的高竞争原子 Scatter 转化为规整的段级规约实现了高达 6.3× 的加速比。大规模核外Out-of-Core处理研究人员对数据集严重超出 GPU VRAM 容量的工作负载进行了基准测试规模扩展至十亿个点。在极端工作负载N10⁹、K32768下标准 PyTorch 实现因 OOM 错误立即崩溃Flash-KMeans 仅用 41.4 秒完成一次迭代而最健壮的基准方案fastkmeans需要 261.8 秒端到端加速比为 6.3× 到 10.5×。这证明了该框架的流水线执行在限制峰值内存占用的同时对那些装不进显卡的数据实现了数量级的加速。快速首次运行高性能机器学习中一项隐性成本是为特定数据规模寻找最优 Kernel 配置而进行的穷举式自动调优。在大规模数据N8M、K64K下穷举调优需要超过 325 秒才能找到最优配置缓存感知编译启发式通过分析推导在不到 2.5 秒内得出接近最优的配置首次运行等待时间减少高达 175×。最关键的是启发式方法与穷举最优解的性能差距在 0.3% 以内——无需等待数分钟的调优即可获得最佳速度。质量与加速比这些加速提升并不以精度为代价。与乘积量化Product Quantization等用精度换时间的近似方法不同Flash-KMeans 在数学上保持精确。平均 12.5× 的加速比不是通过在数学上走捷径实现的而是通过削减 IO 开销实现的——与 Lloyd 算法完全相同的数学收敛性速度大幅提升内存占用仅为原来的一小部分。这些实验证明Flash-KMeans 是首个能在单张 GPU 上扩展到超大 K 值、同时保持峰值硬件效率的精确 K-Means 实现。它将内存密集型的聚类任务转化为计算受限的流式过程。总结Flash-KMeans 将聚类重新定义为流式效率的函数而非静态内存容量的函数。通过摒弃传统 GPU 实现中重物化的旧范式将算法视为硅片与逻辑的一体化协同设计整个系统更接近于一台真正尊重硬件物理层级结构的机器而不只是被动地存储数据。本文从信息假设出发理解智能如何从高维数据的低维对称性中涌现。建立了参数化与非参数化的分野重访了确定性 KNN 与变分推断概率世界之间的理论桥梁将 K-Means 定位为信息组织的原子单元。随后识别了历来作为 FAISS 等标准库架构瓶颈的 IO 受限瓶颈和原子写入竞争深入探讨了 Flash-KMeans 的方法论——FlashAssign 以无物化的 Online Argmin 取代 O(NK) 内存开销Sort-Inverse Update 将无序 Scatter 转化为规整的段级规约两者共同实现理想情形。最终实验结果表明Lloyd 算法的精确实现达到了 12.5× 的加速比证明在计算成本低廉而带宽才是瓶颈的当下内存效率才是扩展的真正前沿。传统聚类的将每个中间距离外化到慢速全局内存的需求是一种可以绕过的低效。将中心点计算与 HBM 流量解耦得到的不只是更快的工具而是一类新的高性能算法不再与硅片对抗而是对数据本身的结构对称性建模。论文https://avoid.overfit.cn/post/8ed38476482d4192861b5ec68fd5ea45作者Vizuara AI