[Kimi Bug] Fix cannot access local variable 'active_non_spec_mask_cpu' (#54781)

Signed-off-by: yewentao256 <zhyanwentao@126.com>

tests/models/kimi_k3/test_kda_metadata.py

@@ -2,7 +2,8 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 from dataclasses import fields
-from unittest.mock import Mock
+from types import SimpleNamespace
+from unittest.mock import Mock, patch
 
 import pytest
 import torch
@@ -21,6 +22,7 @@ from vllm.models.kimi_k3.nvidia.kda_metadata import (
     _mamba_get_block_table_tensor,
     stage_spec_decode_metadata,
 )
+from vllm.models.kimi_k3.nvidia.model import KimiLinearForCausalLM
 from vllm.v1.attention.backend import AttentionMetadataBuilder
 from vllm.v1.attention.backends.gdn_attn import (
     GDNAttentionBackend,
@@ -30,11 +32,13 @@ from vllm.v1.attention.backends.gdn_attn import (
 from vllm.v1.attention.backends.recoverssm_metadata import (
     RecoverSSMPostprocessMetadata,
 )
+from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum
 from vllm.v1.attention.backends.utils import (
     NULL_BLOCK_ID,
     mamba_get_block_table_tensor,
 )
 from vllm.v1.kv_cache_interface import MambaSpec
+from vllm.v1.worker.mamba_utils import validate_mamba_state_copy_funcs
 
 BLOCK_SIZE = 16
 DEVICE = torch.device("cpu")
@@ -123,6 +127,144 @@ def _make_builder(
     return builder
 
 
+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)
+    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,
+    )
+
+    # This is the same compatibility check performed while initializing the
+    # model runner. RecoverSSM adds two workspace states that are not copied.
+    validate_mamba_state_copy_funcs(
+        {kv_cache_spec: [0]},
+        {
+            MambaAttentionBackendEnum.GDN_ATTN: (
+                KimiLinearForCausalLM.get_mamba_state_copy_func()
+            )
+        },
+    )
+    assert len(kv_cache_spec.shapes) == 4
+
+    builder_config = SimpleNamespace(
+        model_config=SimpleNamespace(
+            hf_text_config=SimpleNamespace(linear_key_head_dim=32)
+        ),
+        cache_config=SimpleNamespace(
+            mamba_cache_mode="align",
+            use_kda_recoverssm=True,
+        ),
+        parallel_config=SimpleNamespace(decode_context_parallel_size=1),
+        speculative_config=SimpleNamespace(
+            num_speculative_tokens=2,
+            parallel_drafting=False,
+        ),
+        compilation_config=SimpleNamespace(
+            cudagraph_mode=CUDAGraphMode.NONE,
+            max_cudagraph_capture_size=None,
+            static_forward_context={},
+        ),
+        scheduler_config=SimpleNamespace(max_num_seqs=4),
+        additional_config={},
+        num_speculative_tokens=2,
+        use_v2_model_runner=False,
+    )
+    builder = KimiK3KDAMetadataBuilder(
+        kv_cache_spec=kv_cache_spec,
+        layer_names=["layer.0"],
+        vllm_config=builder_config,
+        device=DEVICE,
+    )
+
+    # An all-prefill speculative batch used to leave an all-false spec mask
+    # alive, then access active_non_spec_mask_cpu before it was initialized.
+    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]))
+    builder.mamba_aligned_state_indices = mamba_get_block_table_tensor(
+        prefill_common.block_table_tensor,
+        prefill_common.seq_lens,
+        kv_cache_spec,
+        "align",
+    )
+    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
+    assert prefill_metadata.checkpoint is not None
+
+    # Creating the commit context is the first metadata path that consumes the
+    # builder's layer_names. Mock only the GPU-cache-dependent constructor.
+    builder.recoverssm_context = None
+    forward_layer = Mock()
+    builder.vllm_config.compilation_config.static_forward_context["layer.0"] = (
+        forward_layer
+    )
+    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]))
+    builder.mamba_aligned_state_indices = mamba_get_block_table_tensor(
+        spec_common.block_table_tensor,
+        spec_common.seq_lens,
+        kv_cache_spec,
+        "align",
+    )
+    recoverssm_context = Mock()
+    with patch(
+        "vllm.models.kimi_k3.nvidia.ops.recoverssm.KDARecoverSSMCommitContext.create",
+        return_value=recoverssm_context,
+    ) as create_context:
+        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
+    assert spec_metadata.recoverssm_context is recoverssm_context
+    create_context.assert_called_once_with(
+        [forward_layer],
+        spec_query_len=3,
+        max_num_reqs=builder.vllm_config.scheduler_config.max_num_seqs,
+    )
+
+
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
 def test_internal_checkpoint_metadata_targets_last_aligned_boundary():
     device = torch.device("cuda")

vllm/models/kimi_k3/nvidia/kda_metadata.py

@@ -409,6 +409,8 @@ class KimiK3KDAMetadataBuilder(GDNAttentionMetadataBuilder):
                 num_spec_decodes = 0
             else:
                 num_spec_decodes = spec_sequence_masks_cpu.sum().item()
+                if num_spec_decodes == 0:
+                    spec_sequence_masks_cpu = None
 
         spec_request_indices = None
         if num_spec_decodes == 0: