Modular and Efficient Approximate Nearest Neighbors Search with Vectorium and kANNolo
做向量检索时,换一种 embedding 形态或量化方式,往往就得换一个库。现有的 ANN 库要么只支持一种组合,要么每种组合各写一份代码,没有一个 ANN 库在同一套接口下覆盖稠密、稀疏和多向量三种 embedding。
今日场景今天的基础设施赛道大多在算推理服务的细账。投机解码有两篇:HybridSpec 发现 GDN 和 MoE 混合的模型验证成本涨得特别快,另一篇提醒投机解码可能悄悄改变带约束生成的输出。FlagRot 让 KV cache 里每个 K 向量的每一小段自己决定要不要旋转,GC-MoE 和 EpiLoRA 分别把 MoE 的通信和 LoRA 的计算藏进主计算里。训练这边,LEAP 在任务提交前预测它在每种加速卡上的吞吐,用来选卡。 — 评分 133 篇,入选 10 篇,分数 3.70 – 4.55
同一个训练任务往往能跑在好几种加速卡上,跑在哪种上最划算,取决于它在每种卡上的吞吐。任务放错了卡照样能跑完、不报错,多花的钱没人察觉。
做向量检索时,换一种 embedding 形态或量化方式,往往就得换一个库。现有的 ANN 库要么只支持一种组合,要么每种组合各写一份代码,没有一个 ANN 库在同一套接口下覆盖稠密、稀疏和多向量三种 embedding。
长上下文推理的瓶颈在 KV cache,把其中的 K 向量压到很低的位宽能省下大量显存,但误差难控。常见做法是量化前先把 K 向量旋转一下,默认旋转总有好处。
投机解码号称不改变输出分布,但这个保证有个前提。语法约束、重复惩罚、水印这类组件会在采样前修改 token 分数,它们在每个验证位置拿到的前缀和状态必须是对的。作者检查了 Transformers 等推理库的部分版本,确认了两种出错方式:这类组件读到了还没验证通过的草稿 token,或者推进了内部状态,草稿被拒后却没有回滚。作者在一个训练好的模型上开了无分类器引导,选定一个 prompt,统计第一个输出 token 是否为某个事先定好的 token。普通解码和从 prompt 里找草稿的投机解码相比,这个概率至少相差 60.26 个百分点。用投机解码加速带约束的生成时,输出可能已经悄悄变了。
循环语言模型每生成一个 token,都要把同一组层重复跑好几遍,推理引擎里每层的固定开销也就跟着成倍增加。作者在一个开源的循环模型推理引擎上测了 Ouro 和 Huginn 两个模型。单请求时,把每轮循环各提交一次的 CUDA Graph 合成一次提交,吞吐几乎不变。真正的开销是 Ouro 每个 token 要跑的 5689 个小 kernel。把这些小 kernel 融合起来,同时保证输出逐位不变,吞吐提高到原来的 1.11 到 1.59 倍。作者按是否逐位一致给每项优化分级:一种近似融合没有更快,还在 GSM8K 上少答对了几道题,所以没有采用。
投机解码里,草稿越来越便宜,每一轮的时间主要花在目标模型验证草稿上。一次验证多少个草稿 token 是关键选择:验得多能多接受,但第一个被拒之后的计算全白做。目标模型一旦混合了 Gated DeltaNet(GDN,一种带循环状态的线性注意力层)和 MoE,这笔账就变了。在 Qwen 的一个这类模型上,验证长度从 4 加到 16,验证耗时涨了 71%,稠密模型只涨几个百分点。HybridSpec 每一轮都在 GPU 上重新选验证长度,依据是草稿 token 的预估接受率和目标模型的实测延迟。在三个这类模型上,它比 DFlash 这种并行出草稿的方法再快三成左右。
给一个大模型找最合适的推理部署配置,要选引擎、怎么切到多张 GPU、什么精度和调度参数,平时只能一个个试,每次占一到几个节点好几分钟。ROCm InferaSim 抓住了一点:batch 大小、prompt 长度、GPU 数量这类选择只改变每张卡的工作量,不改变实际跑的 kernel。所以每组不同的 kernel 只需在一个节点上实测一次,作者把这次实测叫锚点,再用解析公式推算其他配置的性能。在 AMD GPU 上,实测这些锚点花了 9.1 GPU 小时,直接测它们覆盖的配置要 204 GPU 小时。预测出的每个输出 token 耗时,中位误差是 8.1%。调参前先用预测筛掉大部分候选,能省下大部分试机成本。
基础模型换成 FP8 或更低精度后,主矩阵乘变便宜了,LoRA 适配器却成了每步解码里越来越大的开销。多租户推理里,一份基础模型同时挂着很多租户的适配器,它们还在主矩阵乘旁边单独跑一串小 kernel。作者注意到,主矩阵乘只有在最后写出结果时,才用到 LoRA 降维那一步的输出。EpiLoRA 于是先启动降维计算,主矩阵乘同时读取权重,到写结果时,在同一个 kernel 里算出每个租户的升维结果并加上去。在 Hopper GPU 上,单个适配器解码时每个 token 的耗时比 vLLM 最快的配置低最多 24%。同样的低精度权重也能直接用来做 LoRA 强化学习,27B 的模型一张 48GB 的卡就能训。
MoE 解码时,每张卡上只有几个到几百个 token,分组矩阵乘受带宽限制,本来就用不满所有 SM。可现有的融合算子仍让 token 分发和矩阵乘挤在同一批 SM 上,互相抢寄存器和共享内存。GC-MoE 用 CUDA Green Contexts 在驱动层把 GPU 切成一小块通信分区和一大块计算分区,两边各跑自己的 kernel、各自调优,第一批 token 一到,计算就能开始。在 H20 上解码 DeepSeek 的一个 MoE 模型,它比现有的 DeepGEMM 融合算子平均快 18.8%。MoE 解码时,把通信和计算分到不同的 SM 上,比让它们在同一个 SM 里抢资源更划算。
分布式训练里,一张 GPU 出了硬件故障,最先报错的往往是其他节点上的 NCCL。日志里的异常只说明任务在哪里停下,不说明故障从哪里开始。NVIDIA 开源了内部用了多年的故障归因系统 NVIDIA Attribution。它把一个任务相关的应用日志、系统日志、GPU 驱动消息、调度器状态和监控指标收在一起。然后按固定规则比对各个事件的先后和关联:某个节点先报 GPU 硬件错误、随后其他节点才出 NCCL 报错,就把前者判为根因,后者记为症状。最后它给出根因、受影响的节点和证据。文中的例子里,三条 NCCL 报错最终归到某个节点上的一次 NVLink 信号故障。按 NVIDIA 的说法,它每周分析数百万个任务。这样告警和修复可以直接对准出故障的机器,不被下游报错带偏。