跳转至

ReSET: Accurate Latency-Critical NVFP4 Reasoning via Step-Aware Temperature Scaling

会议: NeurIPS 2026
arXiv: 2606.13233
OpenReview: https://openreview.net/forum?id=devQEPcZzG
代码: https://github.com/aiha-lab/ReSET
领域: 其他
关键词: 大模型推理, NVFP4量化, 自适应温度缩放, 延迟敏感解码, CUDA内核优化

一句话总结

针对低比特NVFP4量化在长链推理中符号token易错采样及小batch解码硬件利用率崩溃两大瓶颈,ReSET提出结合全局基线与步骤级相对熵的动态温度调节机制,并配合专为极小Batch设计的CUDA-core NVFP4 GEMV算子,在B200上实现推理精度全面超越PTQ基线并取得最高2倍端到端解码加速。

研究背景与动机

大语言推理模型(LRMs,如 DeepSeek-R1、o1 与 Qwen3)通过生成极长的中间思维链(CoT,通常超过万个 token)大幅提升了数学和逻辑推理能力。然而,这种依靠测试时计算扩展的范式将计算开销完全转移到了自回归生成的单 token 解码延迟(TPOT)上。长序列解码在硬件底层面临两重严苛束缚:一是每生成一个 token 都需要全量访存加载权重带来的内存带宽瓶颈,二是长上下文累积的庞大浮点运算量。NVIDIA Blackwell 架构原生支持的微缩放低精度格式 NVFP4(W4A4,细粒度分组大小为 16),理论上提供了相较于 BF16 约 4 倍的峰值算力与 4 倍的显存压缩比,成为突破推理服务吞吐与延迟瓶颈的最具吸引力的途径。

然而,将 NVFP4 直接部署到长链推理模型时面临两大现实阻碍。第一是精度灾难:多步推理过程中充斥着数字、运算符和结构标记等局部强约束的符号 token,量化引起的概率扰动会在此类低熵位置诱发错误采样,并通过漫长推理轨迹级联放大导致最终答案崩溃;现有的训练后量化(PTQ)方法大多专注于压低本就极小的权重重构误差,但权重误差与最终推理精度并无直接相关性,且在线激活旋转等操作带来了不可忽视的运行时开销。第二是小 Batch 场景下的硬件利用率崩溃:Blackwell 架构的 NVFP4 Tensor Core 指令(tcgen05.mma)在 token 维度被硬件强制绑定为 \(M=128\) 的固定计算分块,而在严格满足 TPOT 延迟服务等级协议(SLO)、显存 KV Cache 压力与长上下文干扰的真实在线推理服务中,可行解码 batch size 往往被压制在 \(M \le 8\)。这导致在这一延迟关键区间,Tensor Core 的分块占用率不足 6.25%,实测硬件利用率甚至低于 1%,NVFP4 标称的 4 倍算力优势完全丧失。

更深入的不确定性分析表明,现存以 token 熵为单一阈值的解码控制存在严重的粒度错配:token 的绝对熵值并非仅由局部语法决定,而是被其所属推理步骤(Reasoning Step)的整体不确定性主导。当模型进入高不确定性的思考步骤时,原本完全确定性的符号 token(如数字)其绝对熵也会水涨船高,进而被全局固定阈值错误跳过、无法施加防错采样锐化;同时,量化还会使得不确定步骤中的发散性分支 token 出现概率过度坍塌。核心 idea:根据在线估算的步骤级相对不确定性动态切换全局与步骤相对熵阈值,对低熵符号决策实施精确锐化并对高熵分支恢复探索多样性,同时配合为 \(M \le 8\) 量身定制的无共享内存依赖 CUDA-core NVFP4 矩阵向量乘内核,软硬协同攻克量化长链推理的精度与延迟双重瓶颈。

方法详解

整体框架

ReSET 的端到端推理管线主要包含两个协同轴线:在算法侧,设计因果在线步骤熵估计器与步骤感知温度缩放(Step-Aware Threshold, SAT)策略,在几乎零计算开销(每 step 仅标量递推与单次分支)下动态调控采样分布;在系统底层,根据解码 batch 大小 \(M\) 实施动态内核分发,在小 batch 延迟敏感阶段调用定制的 CUDA-core GEMV 内核,彻底绕过 Tensor Core 的固定分块填充损耗。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入上下文与当前生成的Token序列"] --> B["自回归解码投影与内核运行时分发"]
    B --> C["CUDA-core极小MM GEMV内核<br/>多Token CTA融合 + 寄存器级反量化"]
    C --> D["Token级熵计算与在线步骤熵混合估计"]
    D --> E["步骤感知阈值判定与双向温度缩放"]
    E --> F["自适应采样输出下一个Token并更新推理轨迹"]

关键设计

1. 步骤感知温度缩放:打破固定阈值对不确定步骤符号决策的误判 在传统基于 token 熵的采样方法中,通常预设一个全局阈值 \(\tau_0\)(如模型在校准集生成的第 80 百分位 token 熵),将低于阈值的 token 视为确定性决策并降低温度。然而,当推理模型处于高度复杂的思考步骤时,整个步骤的基础不确定性升高,使得原本应严格准确的局部符号(例如推导公式中的常数“7”或运算符“+”)的绝对熵值被迫漂移至 \(\tau_0\) 之上,导致固定阈值失效。针对这一痛点,ReSET 引入了全局均值 \(\bar{H} = \frac{1}{t+1}\sum_{i=0}^t H_i\) 与步骤在线熵估计 \(\hat{H}_{\text{step}}(t)\),提出了步骤感知阈值(SAT)策略: $$ \tau_t = \begin{cases} \tau_0, & \hat{H}{\text{step}}(t) \le \bar{H} \quad (\text{置信步骤}) \ \hat{H} $$ 在置信步骤中,步骤整体熵极低,退回全局经验阈值 }}(t), & \hat{H}_{\text{step}}(t) > \bar{H} \quad (\text{不确定步骤}) \end{cases\(\tau_0\) 能避免阈值过低导致的漏锐化;在不确定步骤中,阈值自适应上浮为当前步骤的平均熵 \(\hat{H}_{\text{step}}(t)\),使得判断准则转化为“该 token 相对其所在步骤是否属于确定性决策”。若当前 token 熵 \(H_t < \tau_t\),则赋予极低温度 \(T_{\text{low}}\)(如 0.1)压制量化导致的非 top-1 错误采样;反之,若 \(H_t \ge \tau_t\),则赋予升高温度 \(T_{\text{high}}=1.0\),有效补偿量化引发的高熵分支 token(如转折连词“Alternatively”)概率坍塌,恢复多路径探索能力。

2. 混合在线步骤熵因果估计器:平衡步骤初期的方差与中后期的低偏差 SAT 的核心输入是当前推理步骤的熵统计量 \(\hat{H}_{\text{step}}(t)\),但离线真实平均需要获知整段由双换行符(\n\n)划定的完整推理单元,而自回归生成只能依据历史已生成 token。针对这一因果约束,ReSET 针对步骤内部的不同阶段构建了双机制混合估计器(HSE): $$ \hat{H}_{\text{step}}(t) = \begin{cases} \frac{1}{|\mathcal{W}t|}\sumt} H_i, & t - t_0 < w \quad (\text{步骤初期初始化}) \ \frac{1}{t - t_0 + 1}\sum $$ 其中 }^t H_i, & t - t_0 \ge w \quad (\text{步骤中后期累计}) \end{cases\(t_0\) 为当前步骤的起始 token 索引,\(w\) 为滑动窗口大小(默认设为 32)。在刚跨入新步骤的前 \(w\) 个 token 内,若直接采用步骤内单边均值,样本过少且多为格式化引导词,方差极大;因此采用跨越边界的历史 \(w\) 长度滑动窗口 \(\mathcal{W}_t\) 提供稳定的初始基准。一旦步骤内生成的 token 达到 \(w\) 个,立即无缝平滑切换为严格的步骤内累计均值,以零偏差捕获当前步骤独有的不确定性特征。整个统计仅需维护一个全局累加和与步骤内部累加量,每个解码步的额外开销仅为微秒级的常数时间标量操作。

3. 极小Batch专用的 CUDA-core NVFP4 GEMV 内核:打破 Tensor Core 硬件瓦解瓶颈 针对真实 serving 中 \(M \le 8\) 导致 Tensor Core 无法被填满、有效算力利用率不足 1% 的硬件物理限制,ReSET 抛弃了强制将 \(M\) padding 至 128 的 Tensor Core CUTLASS 路径,从底层构建了纯 CUDA-core NVFP4 向量矩阵乘内核。该内核针对性地攻克了三大底层挑战: 首先,多 Token CTA 融合(Multi-token CTA fusion):针对独立线程块重复访存权重的问题,在单个 CUDA 线程块(CTA)内聚合多个活跃解码 token,共享并仅加载一次来自 HBM 的 NVFP4 权重切片,在不扩增 \(M\) 的前提下实现了跨 token 的权重访存复用; 其次,多累加器线程设计(Multi-accumulator threading):单个线程同时负责 2 行独立权重与对应的输出累加,构造出两条正交的 hfma2 计算流水线,使得 Warp 调度器能够交错发射反量化与乘加指令,彻底隐藏 FP4 解包延迟并激发出极高的指令级并行度; 最后,全寄存器级反量化流水(Register-only dequantization):利用 cvt.rn.f16x2.e2m1x2 指令将打包的 E2M1 FP4 数据解包为 FP16,并直接在寄存器中与 E4M3 块缩放及全局 FP32 缩放因子相乘,随后直接输入 hfma2 运算,完全不占用共享内存(Shared Memory),消除了中间数据落盘与线程同步屏障,将线性层激活量化开销由 4.55–5.08 \(\mu s\) 降至 1.65–1.81 \(\mu s\)。

损失函数 / 训练策略

ReSET 属于完全免训练(Training-free)的测试时解码控制与内核优化技术。模型权重直接采用 NVIDIA ModelOpt 工具库完成离线 PTQ 转换(E2M1 权重/激活,E4M3 块缩放,block size 为 16,全局缩放因子在 C4 数据集上校准),KV Cache 保持 BF16 精度。基准阈值 \(\tau_0\) 的标定仅需在 NuminaMath-1.5 验证集上抽取 5 道题目,统计生成 token 的 80% 分位数(如 R1-7B 为 0.6446,Qwen3-32B 为 0.5363),即可全自动泛化至所有下游任务,无需任何任务级超参搜索。

实验关键数据

主实验

在 NVIDIA 单张 B200 GPU 上,针对 DeepSeek-R1-Distill-Qwen 与 Qwen3 两个系列共 5 个尺寸模型,在 AIME-120、GPQA-Diamond 与 LiveCodeBench 三大高难度长链推理基准上进行 8 次随机种子评测,全面对比 BF16 基线及主流 PTQ 量化方案(RTN、BRQ、4/6、MR-GPTQ):

基准任务 方法 R1-7B R1-14B Qwen3-8B Qwen3-14B Qwen3-32B 平均分 (Avg)
AIME-120 BF16 Baseline 45.7 57.4 70.4 76.1 75.8 65.1
NVFP4 (RTN) 39.6 52.4 62.5 70.4 74.4 59.9
BRQ (Hadamard 16) 41.4 49.8 53.8 66.9 73.0 57.0
4/6 (Adaptive Scale) 41.1 53.1 64.0 70.1 74.8 60.6
MR-GPTQ 39.6 50.6 65.2 71.0 73.3 60.0
ReSET (本文) 43.8 54.0 64.9 72.1 77.5 62.5
GPQA-Diamond BF16 Baseline 48.3 58.4 54.4 61.3 64.1 57.3
NVFP4 (RTN) 47.1 53.7 50.7 57.4 62.8 54.3
BRQ 45.5 56.7 53.5 58.2 60.8 55.0
4/6 44.6 55.4 51.9 57.6 64.6 54.8
MR-GPTQ 43.4 57.4 52.4 60.9 60.6 54.9
ReSET (本文) 46.0 57.6 53.2 58.5 62.9 55.6
LiveCodeBench BF16 Baseline 28.2 38.5 43.2 46.5 49.9 41.3
NVFP4 (RTN) 29.5 37.1 36.4 46.7 46.5 39.2
BRQ 27.6 37.7 36.2 39.0 45.1 37.1
4/6 27.0 36.4 35.8 45.3 45.7 38.0
MR-GPTQ 27.2 34.9 42.9 47.2 45.1 39.5
ReSET (本文) 28.4 37.9 42.1 46.1 46.7 40.2

在 Qwen3-32B 上的单投影层内核延迟(\(M=1\), 单位 \(\mu s\)):

投影层类型 维度 (\(N \times K\)) vLLM (CUTLASS) Ours (CUDA-core GEMV) 加速比
QKV 投影 \(10240 \times 5120\) 13.41 6.62 2.03×
Out 投影 \(5120 \times 8192\) 12.40 7.73 1.60×
Gate-Up 投影 \(25600 \times 5120\) 25.89 13.25 1.96×
Down 投影 \(5120 \times 25600\) 33.26 15.00 2.22×

消融实验

针对温度控制机制在 AIME-120 上的策略消融,以及 CUDA-core 内核各优化点在 Qwen3-32B Down 投影上的消融验证:

实验组别 评估配置 核心指标表现 结果分析与说明
采样策略消融 (AIME-120 Avg) NVFP4 基准 (\(T=0.6\)) 59.9 默认固定温度采样基线
全局固定阈值 (Fixed Threshold) 60.9 (+1.0) 忽略步骤不确定性,高熵步骤符号错误依旧频繁
滑动窗口均值 (Sliding Window) 60.8 (+0.9) 仅靠短时滑动窗口缺乏推理步骤边界全局感知
低熵步骤全量统一锐化 (All \(T_{\text{low}}\)) 60.6 (+0.7) 粗暴压低所有词温度,抹杀语义灵活性
截断采样搜索 (Top-\(p\) / Min-\(p\)) 58.2 ~ 60.3 静态截断无法区分局部符号决策与全局不确定性
ReSET 完整方案 62.5 (+2.6) 动态步骤阈值实现高置信保护与发散探索完美平衡
内核设计消融 (\(M=8\) Down_proj \(\mu s\)) 完整定制内核 (C1+C2+C3) 58.3 达成本文最佳延迟
去除多 Token CTA 融合 (w/o C1) 70.9 (+21.6%) 缺乏 CTA 内部跨 Token 权重复用,显存带宽受限
去除多累加器设计 (w/o C2) 64.6 (+10.8%) 单累加链导致 hfma2 存在长依赖空转
同时去除 C1 与 C2 76.5 (+31.2%) 严重退化为朴素单线程向量点积实现

关键发现

  • 步骤上下文主导了局部 token 的不确定性分布:即便像数字这类确定性符号,在高不确定性步骤中的平均绝对熵也会大幅上涨并超越全局固定阈值 \(\tau_0\);ReSET 动态将其与步骤局部均值比较,成功在不确定步骤中恢复了对局部关键符号的保护。
  • 高熵分支决策需要反向提高温度以对抗量化坍塌:实验表明,将不确定步骤的高熵分支温度提升至 \(T_{\text{high}}=1.0\),模型在全部尺寸上的平均得分均获得一致提高;这证明 NVFP4 量化会过分压缩高方差分支的探索空间,适度升温能重新激活多样化探索。
  • 算子级优化在小 Batch 场景下完全颠覆了 Tensor Core 的理论优势:由于 \(M \le 8\) 时硬件分块填充率不足 6.25%,定制的纯 CUDA-core GEMV 加上全寄存器解包管线,反而比使用 Blackwell 硬件 Tensor Core 提速 1.57–2.49 倍,并在端到端推理中取得逼近 2 倍(Qwen3-32B 达 1.97×)于 BF16 的总延迟优势。

亮点与洞察

  • 解耦推理步骤与局部决策的不确定性层级:指出以往 LLM 解码控温方法将“Token 熵”与“步骤不确定性”混为一谈的根本误区,以极其朴素的二阶均值比较(全局均值 vs 局部步骤均值)完美化解了长链推理中的量化误差累积。
  • 反直觉的软硬协同系统工程直觉:在 Blackwell 具备最强 NVFP4 Tensor Core 的前提下,敏锐抓住了延迟关键型长推理任务中 \(M \le 8\) 的物理本质,大胆返璞归真采用优化到极致的 CUDA-core 寄存器反量化路线,攻克了系统领域的“纸面峰值算力陷阱”。
  • 极具通用性的无侵入设计:ReSET 既不破坏已有预训练模型权重,也不需要引入高延迟的在线旋转投影(如 Hadamard 变换),仅凭一个仅需 5 个校准样本的标量温度策略即可无缝嫁接于任何量化模型之上。

局限与展望

  • 推理步骤边界的简易启发式切分依赖:当前系统依赖双换行符(\n\n)作为推理步骤的物理切分边界,尽管在绝大多数显式 CoT 模型(如 DeepSeek-R1、Qwen3)中高度规范且有效,但在没有严格结构换行标记的模型输出中可能出现步骤界定失真。
  • 多卡分布式张量并行扩展评估缺失:论文当前的 GEMV 内核实现主要针对单卡单 GPU 的本地线性投影优化,在更大尺寸模型(如 70B+)跨卡 Tensor Parallelism 的通信掩盖上仍有待进一步测试与融合设计。

相关工作与启发

  • vs PTQ 权重误差优化(如 MR-GPTQ / 4/6):传统方法试图通过二阶信息微调权重或优化动态 block scale 来压缩数值误差,然而 NVFP4 分组较小本身重构误差已极低;ReSET 证明量化对推理的伤害在于采样概率扰动而非单纯的张量重建,从解码采样层介入投入产出比极高。
  • vs 旋转变换方案(如 BRQ / QuIP#):引入随机 Hadamard 变换虽然消除了激活异常值,但在小 Batch 推理中强行引入了多次额外的在线旋转开销,使原本紧张的微秒级解码预算雪上加霜;ReSET 避免了任何运行期矩阵旋转,直接以原生权重运行,保证了极限吞吐。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ 首次系统性揭示步骤级不确定性对量化采样崩塌的支配机制,并提出打破常规的 CUDA-core 小批次 NVFP4 GEMV 实现。
  • 实验充分度: ⭐⭐⭐⭐⭐ 涵盖 5 款主流推理模型、3 大权威硬核基准,且同时提供完备的内核级微秒基准测试与长文本端到端加速评测。
  • 写作质量: ⭐⭐⭐⭐⭐ 逻辑层层推进,从现象观察、理论机制归纳到系统工程落地的闭环非常精彩流畅。
  • 价值: ⭐⭐⭐⭐⭐ 直击下一代 Blackwell 芯片在线长链大模型推理落地的真实核心痛点,具有极高的工业界工程借鉴与落地价值。