执行摘要

• 一句话：修复 Kimi-K3 RecoverSSM 投机解码 prefill 路径的变量未定义崩溃。
• 推荐动作：该 PR 值得精读。它是一个针对特定模型运行时错误的精准修复，代码改动小但逻辑清晰，且补充了完整的单元测试覆盖。修复方案体现了对变量作用域和控制流状态的细致管理，可作为处理类似条件变量问题的参考。

功能与动机

运行 vllm serve moonshotai/Kimi-K3 启用 RecoverSSM 投机解码时，引擎核心抛出 RuntimeError: cannot access local variable 'active_non_spec_mask_cpu'，导致服务失败。错误发生在 prefill 请求未调度草稿令牌（num_decode_draft_tokens_cpu 为 -1）的路径上，变量未初始化。

实现拆解

1. 核心逻辑修复：在 vllm/models/kimi_k3/nvidia/kda_metadata.py 的 KimiK3KDAMetadataBuilder.build 方法中，当 num_spec_decodes == 0 时，将 spec_sequence_masks_cpu 显式设为 None。这确保后续检查 spec_sequence_masks_cpu is not None 时返回 False，避免访问未定义的 active_non_spec_mask_cpu。该修改仅增加两行控制流代码。
2. 测试覆盖：在 tests/models/kimi_k3/test_kda_metadata.py 中新增单元测试 test_kda_recoverssm_startup_metadata_flow_without_model。该测试模拟不加载模型权重的 KDA RecoverSSM 启动流程，先构建一个纯 prefill 批次验证崩溃路径，再构建投机解码批次验证元数据构建正确性。
3. 测试配套：测试通过 monkeypatch 禁用 PIN_MEMORY，并使用 SimpleNamespace 模拟配置对象，避免实际模型加载。还 patch 了 KDARecoverSSMCommitContext.create 以隔离 GPU 依赖。

关键文件：

• vllm/models/kimi_k3/nvidia/kda_metadata.py（模块 元数据构建；类别 source；类型 core-logic；符号 build）: 修复的核心文件，在元数据构建方法中添加控制流调整以解决变量未定义错误。
• tests/models/kimi_k3/test_kda_metadata.py（模块 元数据测试；类别 test；类型 test-coverage；符号 test_kda_recoverssm_startup_metadata_flow_without_model）: 新增单元测试覆盖崩溃路径，验证修复效果并确保元数据构建在 prefill 和投机解码场景下的正确性。

关键符号：build, test_kda_recoverssm_startup_metadata_flow_without_model

关键源码片段

vllm/models/kimi_k3/nvidia/kda_metadata.py

修复的核心文件，在元数据构建方法中添加控制流调整以解决变量未定义错误。

        # 修复：当无投机解码请求时，清除 spec_sequence_masks_cpu 以避免后续访问未定义变量
        if not self.use_spec_decode or num_decode_draft_tokens_cpu is None:
            spec_sequence_masks_cpu = None
            num_spec_decodes = 0
        else:
            spec_sequence_masks_cpu = num_decode_draft_tokens_cpu >= 0
            if self.use_recoverssm:
                assert m.is_prefilling is not None
                assert m.is_prefilling.device.type == "cpu"
                active_decode_mask_cpu = (~m.is_prefilling) & (
                    query_start_loc_cpu.diff() > 0
                )
                spec_sequence_masks_cpu |= active_decode_mask_cpu
            # 原有逻辑：当无草稿令牌且非 RecoverSSM 时清除
            if (
                not self.use_recoverssm
                and num_decode_draft_tokens_cpu[spec_sequence_masks_cpu].sum().item()
                == 0
            ):
                spec_sequence_masks_cpu = None
                num_spec_decodes = 0
            else:
                num_spec_decodes = spec_sequence_masks_cpu.sum().item()
                # 新增修复：即使 RecoverSSM 路径，若 num_spec_decodes 为 0 也清除
                if num_spec_decodes == 0:
                    spec_sequence_masks_cpu = None
        # 后续逻辑：spec_sequence_masks_cpu 为 None 时，不会尝试访问 active_non_spec_mask_cpu
        spec_request_indices = None
        if num_spec_decodes == 0:
            # ... 进入非投机解码分支
        else:
            # ... 进入投机解码分支，此时 spec_sequence_masks_cpu 保证非 None

tests/models/kimi_k3/test_kda_metadata.py

新增单元测试覆盖崩溃路径，验证修复效果并确保元数据构建在 prefill 和投机解码场景下的正确性。

def test_kda_recoverssm_startup_metadata_flow_without_model(monkeypatch):
    """Exercise KDA RecoverSSM startup without loading Kimi-K3 weights."""
    monkeypatch.setattr("vllm.utils.torch_utils.PIN_MEMORY", False)
    monkeypatch.setattr("vllm.v1.attention.backends.utils.PIN_MEMORY", False)
    # 使用 SimpleNamespace 模拟配置对象，避免实际模型加载
    layout_config = SimpleNamespace(
        model_config=SimpleNamespace(
            dtype=torch.bfloat16,
            hf_config=SimpleNamespace(
                linear_attn_config={
                    "num_heads": 4,
                    "head_dim": 32,
                    "short_conv_kernel_size": 4,
                }
            ),
        ),
        cache_config=SimpleNamespace(
            mamba_cache_dtype="auto",
            use_kda_recoverssm=True,
        ),
        parallel_config=SimpleNamespace(tensor_parallel_size=1),
        speculative_config=SimpleNamespace(num_speculative_tokens=2),
    )
    kv_cache_spec = MambaSpec(
        block_size=BLOCK_SIZE,
        shapes=KimiLinearForCausalLM.get_mamba_state_shape_from_config(layout_config),
        dtypes=KimiLinearForCausalLM.get_mamba_state_dtype_from_config(layout_config),
        mamba_type=MambaAttentionBackendEnum.GDN_ATTN,
        mamba_cache_mode="align",
        num_prefill_checkpoint_blocks=1,
    )
    # 构建元数据 builder
    builder = KimiK3KDAMetadataBuilder(
        kv_cache_spec=kv_cache_spec,
        layer_names=["layer.0"],
        vllm_config=builder_config,
        device=DEVICE,
    )
    # 模拟纯 prefill 批次：此路径曾触发崩溃
    prefill_common = create_common_attn_metadata(
        BatchSpec(seq_lens=[50], query_lens=[50]),
        BLOCK_SIZE,
        DEVICE,
        arange_block_indices=True,
    ).replace(is_prefilling=torch.tensor([True]))
    # 构建 prefill 元数据，验证不抛出异常且 num_spec_decodes == 0
    prefill_metadata = builder.build(
        0,
        prefill_common,
        num_decode_draft_tokens_cpu=torch.tensor([-1], dtype=torch.int32),
        num_accepted_tokens=torch.ones(1, dtype=torch.int32),
    )
    assert prefill_metadata.num_prefills == 1
    assert prefill_metadata.num_spec_decodes == 0
    # 模拟投机解码批次，验证正确进入投机解码路径
    spec_common = create_common_attn_metadata(
        BatchSpec(seq_lens=[20], query_lens=[3]),
        BLOCK_SIZE,
        DEVICE,
        arange_block_indices=True,
    ).replace(is_prefilling=torch.tensor([False]))
    spec_metadata = builder.build(
        0,
        spec_common,
        num_decode_draft_tokens_cpu=torch.tensor([2], dtype=torch.int32),
        num_accepted_tokens=torch.ones(1, dtype=torch.int32),
    )
    assert spec_metadata.num_spec_decodes == 1
    assert spec_metadata.recoverssm_commit is not None

评论区精华

1. 崩溃路径分析：Manny7717 详细指出了崩溃的具体代码路径：当 use_recoverssm=True 且无草稿令牌时，spec_sequence_masks_cpu 为全 False 张量而非 None，导致后续分支未定义 active_non_spec_mask_cpu。
2. 修复正确性确认：Manny7717 认证修复方案：将 spec_sequence_masks_cpu 置为 None 可使后续条件判断跳过错误路径，逻辑与非 RecoverSSM 路径一致。
3. 测试需求讨论：mgoin 要求添加单元测试强制覆盖该情况，并指出应限制在 KDA 模块和 RecoverSSM 场景；yewentao256 询问 B300 资源用于集成测试，但最终采用单元测试方案。

• 是否添加单元测试以强制覆盖 prefill 路径崩溃场景 (testing): 决定添加单元测试，通过模拟配置和 patch 避免 GPU 依赖，覆盖崩溃路径。

风险与影响

• 风险：

  1. 回归风险：修复仅针对 num_spec_decodes == 0 的分支，不影响正常投机解码路径；测试覆盖了 prefill 和投机解码两种场景，回归风险低。
  2. 兼容性风险：变更限于 Kimi-K3 模型的 RecoverSSM 特定代码，不影响其他模型或后端。
  3. 性能影响：新增的条件检查仅在 num_spec_decodes == 0 时执行，无性能开销。
     - 影响：用户影响：修复了 Kimi-K3 模型在启用 RecoverSSM 投机解码时的启动崩溃，使服务能正常运行。
     系统影响：提高了 vLLM 对 Mamba 类模型投机解码路径的鲁棒性，确保元数据构建在边缘情况下的正确性。
     团队影响：为类似投机解码场景提供了测试模板，强调了对变量初始化状态的严格管理。

• 风险标记：控制流变更, 测试覆盖

关联脉络

• PR #52724 [Attention] Enable adaptive verification for FLASHINFER_MLA_SPARSE_DSV4: 同属投机解码路径的元数据处理，涉及验证逻辑调整。
• PR #53388 [Feature][Spec] Support disabling trailing prefix-cache block dropping: 投机解码功能演进，涉及配置选项和路径控制。