FlashSampling：快速且内存高效的精确采样

将采样融合到LM-head矩阵乘法中，避免logits物化，实现精确且高效的分类采样

前置知识

Gumbel-Max技巧

精确分类采样的数学技巧，给定 logits $e_\ell \in \mathbb{R}^V$，通过添加独立同分布的 Gumbel 噪声 $g_i \sim \text{Gumbel}(0,1)$ 后取 argmax 得到样本：$i^* = \arg\max_{i \in [V]} e_\ell_i + g_i$。这个结果等价于先计算 softmax 概率再采样，但不需要显式计算归一化常数或概率分布，只需要找到最大的扰动 logit 即可

这是 FlashSampling 的理论基础，使得精确采样可以避免昂贵的 softmax 计算，将采样简化为寻找最大值的问题

Tensor Parallelism (TP)

分布式训练和推理的并行策略，将模型的某些层（如 LM-head 权重矩阵 $W \in \mathbb{R}^{V \times D}$）沿词汇表维度 $V$ 分片到多个 GPU 上。每个 GPU 只计算完整 logits 矩阵的一部分，采样前需要通过 all-gather 操作同步所有分片，这在解码阶段成为性能瓶颈

FlashSampling 的一个核心优势是在多 GPU 设置下避免昂贵的 all-gather 操作，通过 P2P 写入将通信与计算重叠

Kernel Fusion

GPU 编程优化技术，将多个连续的计算内核合并成一个内核，避免中间结果写入 HBM（高带宽内存）。HBM 访问比 SRAM 慢约 100 倍，因此减少 HBM 读写可以显著提升性能。融合通常将矩阵乘法与其后续操作（如激活函数、归约）合并到同一个内核的 epilogue 阶段

FlashSampling 的核心创新就是将采样融合到 LM-head 的 GEMM 操作中，避免将完整的 logits 张量 $[B, V]$ 写入 HBM 再读回

研究动机

现代 LLM 推理栈在自回归解码过程中需要重复进行分类采样，词汇表规模通常达到数万甚至数十万（如 Qwen3-8B 的 $V=151,936$）。传统采样流水线包括三个独立步骤：先用 GEMM 计算 logits $Y = HW^\top \in \mathbb{R}^{B \times V}$，然后将 logits 写入 HBM，最后启动额外的采样内核进行归一化（softmax）和采样。单 GPU 上采样可占用超过 10% 的 token 生成时间，在 tensor-parallel 设置下，由于需要通过 all-gather 收集跨 GPU 的 logits，这一比例升至 20-38%。解码阶段当激活批次较小时，LM-head 投影是内存带宽受限的，物化完整的 logits 张量并启动额外的采样内核增加了无用的内存访问，因为这些 logits 在采样后立即被丢弃

本文的目标是设计一个将采样融合到 LM-head 矩阵乘法 epilogue 中的采样原语，避免将完整的 logits 张量 $[B, V]$ 物化到 HBM。在多 GPU 设置下，将 GPU 间通信与 matmul 计算和 HBM 加载重叠，实现近乎理想的扩展性。方法必须保持精确性，即从目标分类分布中采样而不引入任何近似，同时支持在线和分布式设置

与已有工作不同的是，现有工作主要关注优化注意力机制（如 FlashAttention 避免物化注意力矩阵）或训练时的损失计算，但推理时的采样步骤尚未得到系统优化。高效采样方法如 FlashInfer 仍然操作已物化的 logits，而像 SIMPLE 这样的方法将采样卸载到 CPU 虽然减少了 GPU 开销，但引入了 CPU-GPU 通信延迟。FlashSampling 的独特之处在于从数学角度重新审视精确采样——认识到精确采样不需要显式的 softmax，只需找到最大的扰动 logit，这为内核级融合提供了理论基础。与近期近似采样方法不同，FlashSampling 保持分布精确性，这对生产环境中的 LLM 服务至关重要

核心方法

FlashSampling 采用两阶段设计，核心思想是逐 tile（分块）计算 logits 并采样，而不是先物化完整的 logits 矩阵。Stage 1 是融合的 matmul-plus-sample 内核，在片上 SRAM 中计算一个批次 tile 和一个词汇表 tile 的 logits，应用确定性变换（温度缩放、偏置、掩码），添加 Gumbel 噪声 $g_i = -\log(-\log u_i)$，其中 $u_i \sim (0,1)$，然后只保留每行在每个词汇表 tile 中的局部最大值及其索引。Stage 2 是轻量级归约，在跨所有词汇表 tile 的候选者中找到每行的全局最大值，得到最终的采样索引。这个设计避免了 softmax、归一化常数和前缀和计算，因为精确采样只需要最大扰动 logit 的索引

核心创新点有两个层次：首先，基于 Gumbel-Max 技巧，认识到精确分类采样不需要计算概率分布，只需要找到 $\arg\max_i e_\ell_i + g_i$，这消除了 softmax 和归一化步骤；其次，在多 GPU tensor-parallel 设置下，用流式点对点（P2P）写入替代 all-gather 集合通信，每个 GPU 在 matmul epilogue 中将其每个 tile 的局部最大值写入其他 GPU，而不是等待 GEMM 完成后再通信。这种设计将 GPU 间通信与 matmul 计算和 HBM 加载重叠，实现了近乎理想的扩展。与需要近似的方法不同，FlashSampling 通过 argmax 在 vocabulary tiles 上的分解保持路径精确性，通过 group log-masses 的分层分解保持分布精确性

方法步骤详情

FlashSampling 的完整算法包括两个阶段。Stage 1（融合内核）：对于每个批次 tile $\mathcal{B}$ 和词汇表 tile $\mathcal{V}_t$ 并行执行：(1) 在片上计算矩阵乘法 $Y_{b,i}^{(t)} = \sum_{d=1}^{D} H_{b,d} W_{i,d}$，其中 $(b,i) \in \mathcal{B} \times \mathcal{V}_t$；(2) 对每个输出元素应用确定性变换 $\tilde{Y}_{b,i}^{(t)} = \text{transform}(Y_{b,i}^{(t)})$，包括温度缩放、偏置、掩码；(3) 采样均匀变量 $u_{b,i} \in (0,1)$ 并计算 Gumbel 噪声 $g_{b,i} = -\log(-\log u_{b,i})$；(4) 计算扰动分数 $s_{b,i} = \tilde{Y}_{b,i}^{(t)} + g_{b,i}$；(5) 对每行 $b \in \mathcal{B}$，找到该 tile 中的最大值及其全局词汇表索引 $(m_b^{(t)}, \text{idx}_b^{(t)}) = \max_{i \in \mathcal{V}_t} s_{b,i}$；(6) 将 $(m_b^{(t)}, \text{idx}_b^{(t)})$ 写入 HBM；(7) 如果是多 GPU 设置，通过 P2P 写入将候选者广播到其他 ranks；(8) 如果是多 GPU 设置，跨 rank 设置屏障（P2P 写入不是集合操作）。Stage 2（归约）：对于每行 $b \in \{1, \ldots, B\}$，找到最大值的 tile $t^* = \arg\max_t m_b^{(t)}$，然后输出 $i_b^* = \text{idx}_b^{(t^*)}$

技术新颖性

FlashSampling 的技术新颖性体现在三个方面。第一，首次将精确采样完全融合到 LLM 解码的 LM-head matmul epilogue 中，避免了 logits 张量的 HBM 物化，这是 IO 感知内核融合思想在推理时采样领域的首次应用。第二，提出了基于 Gumbel-Max 分解的精确采样框架，区别于之前的工作，FlashSampling 清晰区分了两个精确性维度：融合的 tile 内核通过 argmax 在 vocabulary tiles 上的分解保持路径精确性，而 grouped、在线和分布式变体通过 group log-masses 的分层分解保持分布精确性。第三，在多 GPU tensor-parallel 设置下，用流式 P2P 写入替代 all-gather 集合通信，这种通信模式不是集合操作，可以在 matmul 计算的同时进行，实现了计算与通信的重叠，这是现有 LLM 采样框架中没有的技术

**Figure 1**: Conventional multinomial sampling (left) materializes the full [B, V] logits tensor in HBM between the matmul and the sampler. FlashSampling (right) fuses sampling into the matmul epilogue, followed by a lightweight reduction over vocabulary tiles.

实验结果

论文在四个 NVIDIA 数据中心 GPU（H100、H200、B200、B300）上进行了内核级微基准测试和端到端 vLLM 集成评估。内核级测试显示 FlashSampling 在小到中等批次规模下（$B \leq 64$）在所有 GPU 上都优于所有基线。与 Multinomial Sampling（使用 torch.compile 优化）相比，峰值加速比达到 2.23×（B300 GPU），与 FlashInfer FI1（top-k/top-p）相比达到 1.74×（B200 GPU）。与 FI2（FlashInfer 的 Gumbel-Max）相比加速比较小，因为 FI2 也使用 Gumbel-Max 采样，这表明 FlashSampling 的优势主要来自消除 logits 物化和额外的采样内核开销。多 GPU 测试（TP=2,4,8）显示 FlashSampling 在内存受限的低到中等批次规模（$B=16,64$）下优于所有基线，在计算受限的大批次规模（$B=256$）下接近理想加速比。表 1 显示 FlashSampling 的采样时间只占内核时间的 2-6%，而基线方法随批次规模增长到 12-29%。端到端 vLLM 测试在 B200 GPU 上使用 AIME22-24 数据集，四个模型（Qwen3-1.7B TP1、Qwen3-8B TP1、Qwen3-32B TP2、Llama-3.3-70B TP2）的 TPOT 降低峰值分别为 10.2%（$B=64$）、8.7%（$B=16$）、2.9%（$B=64$）、2.7%（$B=16$）。加速幅度与模型在 LM head 上花费的解码时间比例成正比，小模型因为 attention 和 FFN 层占比小而加速更大

**Table 1**: Sampling as a percentage of total kernel time. A high fraction spent on sampling rather than matmul is an indicator of inefficient sampling implementation.

**Figure 2**: Relative performance on B200. Left: Versus Multinomial Sampling. Right: Versus FlashInfer FI1 and FI2.

**Figure 3**: Runtime (lower is better) of FlashSampling and the three baselines at TP=1, 2, 4, 8. The batch sizes (16, 64, 256) are selected to include compute- and memory-bound decode regimes.

**Figure 4**: Sampling runtime (left) and matmul runtime (in µs) vs batch size. Lower is better.

**Figure 5**: TPOT vs concurrency on B200 across four model sizes. Qwen3-1.7B (TP1) and Qwen3-8B (TP1) see the largest reductions (up to 10% and 7–9%, respectively). Qwen3-32B (TP2) and Llama-3.3-70B (TP2) achieve smaller but consistent gains (peaks of 2.9% and 2.7%, respectively), since attention and FFN dominate decode time at these scales.

查看结构化数据

任务	指标	本文	基线	提升
LLM解码采样	相对加速比（vs Multinomial Sampling）	1.23-2.23×（B=1-64）	1.0×	最高2.23×，小批次优势更明显
LLM解码采样	相对加速比（vs FlashInfer FI1）	1.16-1.74×（B=1-64）	1.0×	最高1.74×
vLLM端到端推理（Qwen3-1.7B）	TPOT降低	最高10.2%	0%	小模型加速显著
多GPU扩展性（TP=8, B=16）	相对理想加速	~2.0×（接近理想）	1.8-1.9×	通信与计算重叠

局限与改进

作者承认的主要局限性是 Triton matmul 实现在大批次规模下效率低于 cuBLAS。当 $B$ 增长到 256 时，工作负载从内存受限转向计算受限，FlashSampling 的 Triton matmul 效率损失部分抵消了采样节省。这是一个可移植性权衡：Triton 是平台无关的（支持 AMD、Intel GPU），而 CUTLASS 实现会更高效但仅限于 NVIDIA GPU。其他局限性包括 top-k 扩展虽然在附录 D 中证明正确但尚未实现，低精度输入（FP8、MXFP4）支持尚未添加，尽管内核结构可以自然支持这些格式。作者还观察到多 GPU 基准测试中不同 Modal 云硬件节点间存在较大运行时差异，即使 GPU 通过 NVLink 连接，这表明实际部署性能可能与理想环境有所不同

独立分析的弱点

FlashSampling 在计算受限的大批次场景下优势减弱，因为 Triton matmul 的效率低于 cuBLAS。改进方向是使用平台特定的优化库（如 CUTLASS）实现 matmul 部分，或者混合使用 cuBLAS 和自定义 epilogue。另一个限制是当前实现只支持单个 token 采样，不支持 top-k 或采样 without replacement（虽然理论上通过 Gumbel-Top-k 可以实现）。改进方向是实现完整的 Gumbel-Top-k 算法，这在 Beam Search 等场景中很有用。此外，当前实现假设每行至少有一个有限的 logit，这在某些掩码场景下可能不成立，需要更鲁棒的边界处理。多 GPU 设置下，P2P 写入模式需要 NVLink 支持，在没有 NVLink 的平台上性能可能下降，可以考虑 fallback 到优化后的集合通信

未来方向

作者提出的未来工作方向包括实现 top-k 扩展（理论上已证明正确），添加低精度输入支持（FP8、MXFP4），以及探索训练阶段的应用。基于 FlashSampling 的成果，可以延伸的研究方向包括：将同样的融合思想应用于其他推理阶段的采样操作（如分层采样、条件采样）；探索与量化技术的结合，在低精度权重和激活下实现高效采样；研究动态批次大小下的自适应 tile 大小选择，进一步优化不同负载下的性能；将 FlashSampling 集成到更多推理框架（如 TensorRT-LLM、TGI）中，验证其在不同栈上的性能；探索在训练时采样（如对比学习、负采样）中的应用，这可能带来更大的性能提升

复现评估

论文提供了良好的复现支持。代码已开源（GitHub：FlashSampling/FlashSampling），使用 Triton 实现，匿名补充 zip 包包含实现源代码和实验脚本。内核级微基准测试在四个 NVIDIA 数据中心 GPU 上运行，使用 Modal 云服务（至少 16 CPU 核心）。软件栈包括 PyTorch 2.11.0、CUDA 13.0、Triton 3.6、FlashInfer 0.6.9。输入和权重使用 BF16 格式。实验配置详细记录：主要测试使用 $D=4,096, V=151,936$，批次规模 $B \in \{1,2,4,8,16,32,64,128,256\}$，附加测试使用更大的 $D=8,192, V=128,256$。所有内核在计时前预热 25 次迭代，使用 NVIDIA CUPTI（单 GPU）和 CUDA events（多 GPU）测量运行时。端到端 vLLM 测试在 B200 GPU 上运行，使用 AIME22-24 数据集，每个批次规模运行 5 次并报告中位数。多 GPU 微基准测试总耗时不到 5 小时，主要挑战是硬件节点间的性能差异。作者使用卡方拟合优度验证内核级正确性，使用 GSM8K 数据集验证端到端正确性（89.4% vs 89.6% 基准，p=0.776）

FlashSampling：快速且内存高效的精确采样 FlashSampling: Fast and Memory-Efficient Exact Sampling