arXiv:2603.15854cs.LGcs.AI2026-03被引 2

将采样融合进LM头计算,无需显存存储概率,大幅提速推理。

FlashSampling: Fast and Memory-Efficient Exact Sampling

  • 采样与矩阵乘法融合,逐块计算并只保留最大值,避免生成完整logits
  • 在4种数据中心GPU上实现核级加速,端到端每输出一个词元快10%
  • 适用于大规模分布式推理,支持多卡高效通信重叠计算

从类别分布中采样在数学上简单,但在大词汇量解码时常引发额外内存流量和后续核函数调用。我们提出FlashSampling,一种精确采样原语,将采样过程融入LM头的矩阵乘法中,且从不在HBM中显式存储logits张量。方法为:在芯片上逐块计算logits,添加Gumbel噪声,每行和每词汇块仅保留一个最大值,最后对块进行小规模归约。在张量并行解码中,FlashSampling用流式点对点写替代logits的all-gather操作,实现GPU间通信与计算、HBM加载的重叠,最多支持8个GPU,大批次下接近理想扩展性。该内核是精确的,因为argmax可分解于分区;分组变体在在线和张量并行场景下通过层次化因子分解保证精确性。FlashSampling在4种数据中心GPU(H100、H200、B200、B300)上的解码工作负载中展示出核级加速,在vLLM端到端实验中,使每输出一个词元的时间减少最多10%。结果表明,精确采样可直接集成在矩阵乘法内部,将带宽瓶颈的采样步骤高效整合为计算尾部。

原文摘要 · Abstract (English)

Sampling from a categorical distribution is mathematically simple, but in large-vocabulary decoding, it often triggers extra memory traffic and extra kernels after the LM head. We present FlashSampling, an exact sampling primitive that fuses sampling into the LM-head matmul and never materializes the logits tensor in HBM. The method is simple: compute logits tile-by-tile on chip, add Gumbel noise, keep only one maximizer per row and per vocabulary tile, and finish with a small reduction over tiles. In tensor-parallel decoding, FlashSampling replaces the all-gather of logits with streaming peer-to-peer writes: This overlaps GPU-to-GPU communication with computation and HBM loads across up to 8 GPUs, with near-ideal scaling at large batch sizes. Our kernel is exact because argmax decomposes over partitions; grouped variants for online and tensor-parallel settings are exact by hierarchical factorization of the categorical distribution. FlashSampling demonstrates kernel-level speedups on decode workloads across 4 different datacenter GPUs (H100, H200, B200, B300), and in end-to-end vLLM experiments, it reduces time per output token by up to $10\%$ on the models we test. These results show that exact sampling, with no approximation, can be integrated into the matmul itself, consolidating the bandwidth-bound sampling step in an efficient epilogue.

推理加速采样优化GPU并行

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。