[Kimi Bug] Fix `cannot access local variable 'active_non_spec_mask_cpu'` (#54781)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Sep 1, 2026 at 17:08 UTC
fc72fc39ace2d5f42346880d3f8d27463a5c7079
2 files changed
+145
-1
tests/models/kimi_k3/test_kda_metadata.py
+143
-1
@@ -2,7 +2,8 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
from dataclasses import fields
5
-from unittest.mock import Mock
5
+from types import SimpleNamespace
6
+from unittest.mock import Mock, patch
7
8
import pytest
9
import torch
@@ -21,6 +22,7 @@ from vllm.models.kimi_k3.nvidia.kda_metadata import (
22
_mamba_get_block_table_tensor,
23
stage_spec_decode_metadata,
24
)
25
+from vllm.models.kimi_k3.nvidia.model import KimiLinearForCausalLM
26
from vllm.v1.attention.backend import AttentionMetadataBuilder
27
from vllm.v1.attention.backends.gdn_attn import (
28
GDNAttentionBackend,
@@ -30,11 +32,13 @@ from vllm.v1.attention.backends.gdn_attn import (
32
from vllm.v1.attention.backends.recoverssm_metadata import (
33
RecoverSSMPostprocessMetadata,
34
)
35
+from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum
36
from vllm.v1.attention.backends.utils import (
37
NULL_BLOCK_ID,
38
mamba_get_block_table_tensor,
39
)
40
from vllm.v1.kv_cache_interface import MambaSpec
41
+from vllm.v1.worker.mamba_utils import validate_mamba_state_copy_funcs
42
43
BLOCK_SIZE = 16
44
DEVICE = torch.device("cpu")
@@ -123,6 +127,144 @@ def _make_builder(
127
return builder
128
129
130
+def test_kda_recoverssm_startup_metadata_flow_without_model(monkeypatch):
131
+ """Exercise KDA RecoverSSM startup without loading Kimi-K3 weights."""
132
+ monkeypatch.setattr("vllm.utils.torch_utils.PIN_MEMORY", False)
133
+ monkeypatch.setattr("vllm.v1.attention.backends.utils.PIN_MEMORY", False)
134
+ layout_config = SimpleNamespace(
135
+ model_config=SimpleNamespace(
136
+ dtype=torch.bfloat16,
137
+ hf_config=SimpleNamespace(
138
+ linear_attn_config={
139
+ "num_heads": 4,
140
+ "head_dim": 32,
141
+ "short_conv_kernel_size": 4,
142
+ }
143
+ ),
144
+ ),
145
+ cache_config=SimpleNamespace(
146
+ mamba_cache_dtype="auto",
147
+ use_kda_recoverssm=True,
148
+ ),
149
+ parallel_config=SimpleNamespace(tensor_parallel_size=1),
150
+ speculative_config=SimpleNamespace(num_speculative_tokens=2),
151
+ )
152
+ kv_cache_spec = MambaSpec(
153
+ block_size=BLOCK_SIZE,
154
+ shapes=KimiLinearForCausalLM.get_mamba_state_shape_from_config(layout_config),
155
+ dtypes=KimiLinearForCausalLM.get_mamba_state_dtype_from_config(layout_config),
156
+ mamba_type=MambaAttentionBackendEnum.GDN_ATTN,
157
+ mamba_cache_mode="align",
158
+ num_prefill_checkpoint_blocks=1,
159
+ )
160
+
161
+ # This is the same compatibility check performed while initializing the
162
+ # model runner. RecoverSSM adds two workspace states that are not copied.
163
+ validate_mamba_state_copy_funcs(
164
+ {kv_cache_spec: [0]},
165
+ {
166
+ MambaAttentionBackendEnum.GDN_ATTN: (
167
+ KimiLinearForCausalLM.get_mamba_state_copy_func()
168
+ )
169
+ },
170
+ )
171
+ assert len(kv_cache_spec.shapes) == 4
172
+
173
+ builder_config = SimpleNamespace(
174
+ model_config=SimpleNamespace(
175
+ hf_text_config=SimpleNamespace(linear_key_head_dim=32)
176
+ ),
177
+ cache_config=SimpleNamespace(
178
+ mamba_cache_mode="align",
179
+ use_kda_recoverssm=True,
180
+ ),
181
+ parallel_config=SimpleNamespace(decode_context_parallel_size=1),
182
+ speculative_config=SimpleNamespace(
183
+ num_speculative_tokens=2,
184
+ parallel_drafting=False,
185
+ ),
186
+ compilation_config=SimpleNamespace(
187
+ cudagraph_mode=CUDAGraphMode.NONE,
188
+ max_cudagraph_capture_size=None,
189
+ static_forward_context={},
190
+ ),
191
+ scheduler_config=SimpleNamespace(max_num_seqs=4),
192
+ additional_config={},
193
+ num_speculative_tokens=2,
194
+ use_v2_model_runner=False,
195
+ )
196
+ builder = KimiK3KDAMetadataBuilder(
197
+ kv_cache_spec=kv_cache_spec,
198
+ layer_names=["layer.0"],
199
+ vllm_config=builder_config,
200
+ device=DEVICE,
201
+ )
202
+
203
+ # An all-prefill speculative batch used to leave an all-false spec mask
204
+ # alive, then access active_non_spec_mask_cpu before it was initialized.
205
+ prefill_common = create_common_attn_metadata(
206
+ BatchSpec(seq_lens=[50], query_lens=[50]),
207
+ BLOCK_SIZE,
208
+ DEVICE,
209
+ arange_block_indices=True,
210
+ ).replace(is_prefilling=torch.tensor([True]))
211
+ builder.mamba_aligned_state_indices = mamba_get_block_table_tensor(
212
+ prefill_common.block_table_tensor,
213
+ prefill_common.seq_lens,
214
+ kv_cache_spec,
215
+ "align",
216
+ )
217
+ prefill_metadata = builder.build(
218
+ 0,
219
+ prefill_common,
220
+ num_decode_draft_tokens_cpu=torch.tensor([-1], dtype=torch.int32),
221
+ num_accepted_tokens=torch.ones(1, dtype=torch.int32),
222
+ )
223
+ assert prefill_metadata.num_prefills == 1
224
+ assert prefill_metadata.num_spec_decodes == 0
225
+ assert prefill_metadata.checkpoint is not None
226
+
227
+ # Creating the commit context is the first metadata path that consumes the
228
+ # builder's layer_names. Mock only the GPU-cache-dependent constructor.
229
+ builder.recoverssm_context = None
230
+ forward_layer = Mock()
231
+ builder.vllm_config.compilation_config.static_forward_context["layer.0"] = (
232
+ forward_layer
233
+ )
234
+ spec_common = create_common_attn_metadata(
235
+ BatchSpec(seq_lens=[20], query_lens=[3]),
236
+ BLOCK_SIZE,
237
+ DEVICE,
238
+ arange_block_indices=True,
239
+ ).replace(is_prefilling=torch.tensor([False]))
240
+ builder.mamba_aligned_state_indices = mamba_get_block_table_tensor(
241
+ spec_common.block_table_tensor,
242
+ spec_common.seq_lens,
243
+ kv_cache_spec,
244
+ "align",
245
+ )
246
+ recoverssm_context = Mock()
247
+ with patch(
248
+ "vllm.models.kimi_k3.nvidia.ops.recoverssm.KDARecoverSSMCommitContext.create",
249
+ return_value=recoverssm_context,
250
+ ) as create_context:
251
+ spec_metadata = builder.build(
252
+ 0,
253
+ spec_common,
254
+ num_decode_draft_tokens_cpu=torch.tensor([2], dtype=torch.int32),
255
+ num_accepted_tokens=torch.ones(1, dtype=torch.int32),
256
+ )
257
+
258
+ assert spec_metadata.num_spec_decodes == 1
259
+ assert spec_metadata.recoverssm_commit is not None
260
+ assert spec_metadata.recoverssm_context is recoverssm_context
261
+ create_context.assert_called_once_with(
262
+ [forward_layer],
263
+ spec_query_len=3,
264
+ max_num_reqs=builder.vllm_config.scheduler_config.max_num_seqs,
265
+ )
266
+
267
+
268
@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
269
def test_internal_checkpoint_metadata_targets_last_aligned_boundary():
270
device = torch.device("cuda")
vllm/models/kimi_k3/nvidia/kda_metadata.py
+2
@@ -409,6 +409,8 @@ class KimiK3KDAMetadataBuilder(GDNAttentionMetadataBuilder):
409
num_spec_decodes = 0
410
else:
411
num_spec_decodes = spec_sequence_masks_cpu.sum().item()
412
+ if num_spec_decodes == 0:
413
+ spec_sequence_masks_cpu = None
414
415
spec_request_indices = None
416
if num_spec_decodes == 0: