diff --git a/tests/models/kimi_k3/test_kda_metadata.py b/tests/models/kimi_k3/test_kda_metadata.py index 8431705ec6..ca184467f0 100644 --- a/tests/models/kimi_k3/test_kda_metadata.py +++ b/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") diff --git a/vllm/models/kimi_k3/nvidia/kda_metadata.py b/vllm/models/kimi_k3/nvidia/kda_metadata.py index 748e5219a7..2920bf62a6 100644 --- a/vllm/models/kimi_k3/nvidia/kda_metadata.py +++ b/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: