Prhub

#52795 [Spec Decode] Enable adaptive verification on DSv4 + sm90

原始 PR 作者 LucasWilkinson 合并时间 2026-08-21 10:44 文件变更 1 提交数 3 评论 8 代码增减 +28 / -4

执行摘要

SM90 启用 DSv4 自适应验证,解码改扁平化元数据

adaptive verification 会按置信度把每请求草稿验证长度裁剪到不同数值,而 CPU 侧元数据只有统一占位长度;SM90 上 DeepGEMM 没有可消费 device 侧每请求 query 长度的 varlen paged-MQA 内核,导致 DSV4 在 Hopper 上启用该功能时启动即失败(见 #47808 评论)。同时 #47808 的动机指出固定 k 投机在高并发下会跌到非投机基线之下,自适应验证是维持投机收益的核心手段,因此需要把这条路径带到 SM90。

值得精读。虽然只有 28 行改动,但它展示了一个典型的硬件能力分层决策:SM100 走原生 varlen、SM90 走 device 扁平化、单请求走 uniform 特判,并用 _use_flattening 统一入口。对从事投机解码、attention 后端与 CUDA graph 相关工作的人尤其有参考价值。阅读时可重点关注 _prepare_decode_tensors 中 adaptive 分支与 uniform 分支的切换条件,以及 PR body 里的性能剖析方法(profiler 对比 + 元数据耗时拆解)。

讨论亮点

唯一的 review 交锋来自 benchislett 的 nit:在 _use_flattening 的 diff 上指出,原先 1 + (speculative_config.num_speculative_tokens if speculative_config else 0) 的写法可以直接写为 1 + vllm_config.num_speculative_tokens,因为配置归一化后已包含 None 语义。作者在下一提交 b505462ef2 即按此修改并回复确认。另一个重要讨论集中在 PR body 的性能剖析:并发 1 下 device-ragged 元数据准备耗时 115.08 us,而 uniform 仅 24.14 us;加上 AdaptiveVerificationManager.record_confidences 中每步约 0.706 ms 的阻塞同步,解释了大部分单并发步延迟差距。

实现拆解

  1. 平台能力探测:在 vllm/v1/attention/backends/mla/indexer.py 新增 _supports_flattened_device_query_lens(),仅在 CUDA、SM90 且启用 DeepGEMM 时返回 True;supports_device_cpu_query_lens_mismatch() 从仅依赖 _supports_varlen_paged_mqa_logits() 改为两者取或,向调度层声明 device 与 CPU query 长度可以不一致。
  2. 统一扁平化决策:新增 _use_flattening(vllm_config),把原先 __init__ 里 self.use_flattening = not _supports_native_decode(next_n) 的逻辑收敛为单一函数,并追加“adaptive verification 开启且 SM90 可用则强制扁平化”的分支,避免 CPU 占位长度与 device 真实长度不一致。
  3. 放开 CUDA graph 捕获:get_cudagraph_support() 在 _supports_varlen_paged_mqa_logits() 基础上并入 _use_flattening(),使 SM90 扁平化路径返回 AttentionCGSupport.ALWAYS,保证 adaptive 裁剪后各请求长度不同也能走 decode CUDA graph。
  4. 单请求特判:_prepare_decode_tensors() 增加 adaptive 检测,当只有一个 decode 请求时(num_decodes == 1 或非 adaptive)且长度均匀时仍走便宜的 uniform 扁平化元数据路径,因为此时 CPU 与 device 的请求边界必然一致。
  5. 验证与配套:PR 自带完整测试计划(indexer、adaptive verification、GPU batch ordering、DeepGEMM 内核测试共 30 个用例通过)及 4x H200 的 GSM8K 评测(94.996% 准确率)与吞吐基准,但未新增任何测试文件,属于纯源码变更。
文件 模块 状态 重要度
vllm/v1/attention/backends/mla/indexer.py 注意力后端 modified 7.21

关键符号

_supports_flattened_device_query_lens _use_flattening supports_device_cpu_query_lens_mismatch get_cudagraph_support _prepare_decode_tensors

关键源码片段

vllm/v1/attention/backends/mla/indexer.py core-logic

唯一变更文件,是 DeepSeek V4 indexer 的 attention 元数据构建核心路径;本次改动新增 SM90 扁平化能力探测并统一 use_flattening 决策,使 adaptive verification 在 Hopper 上可用。

# 核心决策函数:SM90 上没有 SM100 那样的 varlen paged-MQA 内核,
# 但可以在 device 侧构建扁平化元数据,把每个 query 变成单 token 行,
# 从而让“每个请求被裁剪到不同长度”的 adaptive verification 也能工作。
def _supports_flattened_device_query_lens() -> bool:
    return (
        current_platform.is_cuda()
        and current_platform.is_device_capability_family(90)
        and has_deep_gemm()
    )
​
​
# 统一扁平化决策入口,替代原先 __init__ 里分散的写法:
# 1) 不支持 native decode 时本来就要扁平化;
# 2) SM90 上开启 adaptive verification 时也必须扁平化,
# 因为每个验证步各请求的 query 长度都可能不同,CPU 只有占位长度。
def _use_flattening(vllm_config: VllmConfig) -> bool:
    speculative_config = vllm_config.speculative_config
    next_n = 1 + vllm_config.num_speculative_tokens
    return not _supports_native_decode(next_n) or (
        speculative_config is not None
        and speculative_config.enable_adaptive_verification
        and _supports_flattened_device_query_lens()
    )
​
​
class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
    @classmethod
    def get_cudagraph_support(cls, vllm_config: VllmConfig,
                              kv_cache_spec: KVCacheSpec) -> AttentionCGSupport:
        # SM90 扁平化路径同样允许每请求不同长度,因此也要进入 ALWAYS
        # 级别的 CUDA graph 捕获,否则 adaptive 裁剪会破坏图重放。
        if _supports_varlen_paged_mqa_logits() or _use_flattening(vllm_config):
            return AttentionCGSupport.ALWAYS
        return AttentionCGSupport.UNIFORM_BATCH
​
    def __init__(self, *args, block_table_width: int, **kwargs) -> None:
        super().__init__(*args, **kwargs)
        # ...(其余配置初始化不变)
        next_n = self.num_speculative_tokens + 1
        self.decode_threshold = next_n
        self.reorder_batch_threshold = None
        # 统一走 _use_flattening,使 SM90 + adaptive verification 强制扁平化
        self.use_flattening = _use_flattening(self.vllm_config)
        self.supports_varlen = _supports_varlen_paged_mqa_logits()
​
    def _prepare_decode_tensors(self, ...):
        # 节选:adaptive 分支与 uniform 分支的切换条件
        spec_config = self.vllm_config.speculative_config
        adaptive = bool(spec_config and spec_config.enable_adaptive_verification)
        min_decode_len = int(decode_lens_cpu.min().item())
        if not use_native:
            assert self.decode_seq_lens_buffer.dim() == 1
            if (
                not self.supports_varlen
                # 只有一个 decode 请求时,CPU 与 device 的请求边界必然一致,
                # 复用便宜的 uniform 扁平化元数据即可,避免 device-ragged 开销。
                and (num_decodes == 1 or not adaptive)
                and min_decode_len == max_decode_len
                # ...(其余均匀性条件省略)
            ):
                pass # 走 uniform 扁平化路径
            # 否则走 device-ragged 扁平化路径:从 device 侧 decode 长度构建单 token 行

评论区精华

next_n 计算的简化写法 style

benchislett 在 `_use_flattening` 的 diff 上留言:'nit: pretty sure `1 + vllm_config.num_speculative_tokens` does what you want here',指出原先 `1 + (speculative_config.num_speculative_tokens if speculative_config else 0)` 的写法冗余。

结论:作者在 b505462ef2 中改为 `1 + vllm_config.num_speculative_tokens`,并回复 'Done in b505462ef2'。 · 已解决

单并发下 adaptive 吞吐回归与缓解 性能

PR body 剖析并发 1 下固定与自适应验证的步延迟差距:目标图执行时间几乎一致(10.181 vs 10.185 ms),差距来自 device-ragged 元数据准备(115.08 us vs 24.14 us)与 `AdaptiveVerificationManager.record_confidences` 中每步约 0.706 ms 的阻塞同步;单请求复用 uniform 元数据后吞吐从 236.9 提到 240.2 tokens/s。

结论:单并发回归由约 0.706 ms 的同步主导,作者保留为已知限制,并验证高并发收益显著(c=64 时 +13.85%)。 · 待处理

风险与影响

  1. 低并发回归:并发 1 时吞吐 -6.35%(240.2 vs 256.5 tokens/s),源于 device-ragged 元数据准备(115.08 us vs 24.14 us)与每步 confidence-copy 事件同步(约 0.706 ms 关键路径);PR 已通过单请求复用 uniform 路径缓解(+1.4%),但未消除。
  2. CUDA graph 捕获开销:TP2/DP2 对比中捕获时间 121 s / 4.71 GiB(adaptive)vs 52 s / 2.61 GiB(fixed),启动延迟和显存占用明显上升。
  3. 测试覆盖缺口:未新增测试文件,依赖现有 test_adaptive_verification.py、test_indexer_native_next_n.py 等用例与人工 4x H200 验证,缺少自动化覆盖 CUDA graph 捕获路径的长期保护。
  4. 核心路径影响:改动位于 decode attention 元数据构建主路径,DSV4 在 SM90 上启用 adaptive verification 时进入新的 flattening 分支;SM100 原生路径未改动。
  5. AI 辅助生成代码:主体代码由 OpenAI Codex 生成,虽有作者逐行审查承诺,仍存在引入隐蔽行为差异的风险。

用户侧:DSV4 系列在 H100/H200 上可用自适应验证(修复此前启动失败),高并发(16/64)吞吐提升 5.78%/13.85%,接受长度约 4.1,GSM8K 5-shot 精度 94.996% 无退化;单并发用户会看到约 6% 回归。系统侧:decode CUDA graph 捕获时长与峰值显存增加,启动阶段耗时变长。团队侧:确立了“SM90 无 varlen 内核时用 device 长度扁平化”的通用模式,为后续在更多平台/模型上启用自适应验证提供模板;indexer.py 是多个投机解码后端的共享关注点,后续改动需注意该文件。

解码元数据核心路径变更 低并发吞吐回归约 6% CUDA graph 捕获开销显著增加 无新增测试文件 AI 辅助生成代码需人工复核

关联 Issue

#47808 [Spec Decode] DSpark confidence-scheduled verification
#51972 [BugFix] Clamp indexer scheduler-metadata seq_lens to avoid OOB read in DeepGEMM paged-MQA metadata kernel

完整报告

参与讨论