[refactor][model] Kimi-K3 DSpark MLA common base class

.buildkite/test-amd.yaml

@@ -4869,8 +4869,9 @@ steps:
   - tests/models/test_hyv4_rocm.py
   - tests/models/test_vision.py
   - tests/models/test_deepseek_v4_vl_rocm.py
+  - tests/models/test_dspark_mla.py
   commands:
-  - pytest -v -s models/test_terratorch.py models/transformers/test_backend.py models/test_registry.py models/test_deepseek_v4_vl_rocm.py
+  - pytest -v -s models/test_terratorch.py models/transformers/test_backend.py models/test_registry.py models/test_deepseek_v4_vl_rocm.py models/test_dspark_mla.py
   - VLLM_ROCM_USE_AITER=1 pytest -v -s models/test_hyv4_rocm.py -m 'not distributed'
   - pytest -v -s models/test_vision.py::test_simple_mrope_vision_model_spatial_merge
 

.buildkite/test_areas/models_basic.yaml

@@ -70,8 +70,9 @@ steps:
   - tests/models/test_deepseek_v41_replay_batch.py
   - tests/models/transformers/test_layer_registry_e2e.py
   - tests/models/test_deepseek_v41_replay_start.py
+  - tests/models/test_dspark_mla.py
   commands:
-    - pytest -v -s models/test_terratorch.py models/transformers/test_backend.py models/test_registry.py models/test_deepseek_v4_vl_rocm.py models/transformers/test_layer_registry_e2e.py models/test_deepseek_v41_replay_start.py models/test_deepseek_v41_decoder_replay_layers.py models/test_deepseek_v41_replay_batch.py
+    - pytest -v -s models/test_terratorch.py models/transformers/test_backend.py models/test_registry.py models/test_deepseek_v4_vl_rocm.py models/transformers/test_layer_registry_e2e.py models/test_deepseek_v41_replay_start.py models/test_deepseek_v41_decoder_replay_layers.py models/test_deepseek_v41_replay_batch.py models/test_dspark_mla.py
   mirror:
     amd:
       label: ":amd: (MI355) Basic Models (Other)"
@@ -91,8 +92,9 @@ steps:
       - tests/models/test_deepseek_v4_dspark_rocm.py
       - tests/models/test_deepseek_v4_rocm_wo_a.py
       - tests/models/test_deepseek_v4_rocm_compressor_gemm_fusion.py
+      - tests/models/test_dspark_mla.py
       commands:
-      - pytest -v -s models/test_terratorch.py models/transformers/test_backend.py models/test_registry.py models/test_deepseek_v4_vl_rocm.py models/transformers/test_layer_registry_e2e.py models/test_deepseek_v41_replay_start.py models/test_deepseek_v41_decoder_replay_layers.py models/test_deepseek_v41_replay_batch.py models/test_deepseek_v4_dspark_rocm.py models/test_deepseek_v4_rocm_wo_a.py models/test_deepseek_v4_rocm_compressor_gemm_fusion.py
+      - pytest -v -s models/test_terratorch.py models/transformers/test_backend.py models/test_registry.py models/test_deepseek_v4_vl_rocm.py models/transformers/test_layer_registry_e2e.py models/test_deepseek_v41_replay_start.py models/test_deepseek_v41_decoder_replay_layers.py models/test_deepseek_v41_replay_batch.py models/test_deepseek_v4_dspark_rocm.py models/test_deepseek_v4_rocm_wo_a.py models/test_deepseek_v4_rocm_compressor_gemm_fusion.py models/test_dspark_mla.py
       - VLLM_ROCM_USE_AITER=1 pytest -v -s models/test_hyv4_rocm.py -m 'not distributed'
 
 - label: ":nvidia: (B200) Inkling"

tests/models/test_dspark_mla.py

@@ -12,10 +12,102 @@
 from vllm.model_executor.models.qwen3_dspark import DSparkMarkovHead
 from vllm.model_executor.models.registry import ModelRegistry
 from vllm.models.deepseek_v4.nvidia import dspark as dsv4_dspark
+from vllm.models.kimi_k3.common import dspark_mla as common_dspark_mla
 from vllm.models.kimi_k3.nvidia import dspark_mla
 from vllm.models.kimi_k3.nvidia.dspark_mla import K3DSparkForCausalLM, K3DSparkModel
 
 
+def test_nvidia_dspark_binding_uses_multi_head_latent_attention():
+    from vllm.models.kimi_k3.nvidia.mla import MultiHeadLatentAttention
+
+    assert dspark_mla.MultiHeadLatentAttention is MultiHeadLatentAttention
+
+
+def test_default_mla_hooks_return_their_inputs():
+    from vllm.model_executor.layers.attention.mla_attention import MLAAttention
+
+    kv = torch.zeros(2, 4)
+    k_pe = torch.zeros(2, 1, 2)
+    slots = torch.zeros(2, dtype=torch.int64)
+    q = torch.zeros(2, 2, 4)
+    ql = torch.zeros(2, 2, 4)
+    q_pe = torch.zeros(2, 2, 2)
+    layer = SimpleNamespace(
+        kv_cache_dtype="auto",
+        impl=SimpleNamespace(supports_quant_query_input=False),
+    )
+
+    out_kv, out_k_pe, out_slots = MLAAttention._prepare_kv_cache_update(
+        layer, kv, k_pe, slots, None
+    )
+    assert out_kv is kv
+    assert out_k_pe is k_pe
+    assert out_slots is slots
+
+    out_q, out_mha_k_pe = MLAAttention._prepare_mha_inputs(layer, q, k_pe)
+    assert out_q is q
+    assert out_mha_k_pe is k_pe
+
+    formed = MLAAttention._form_decode_q(layer, ql, q_pe, kv, k_pe, kv, None, 2)
+    assert formed[0] is ql
+    assert formed[1] is q_pe
+
+
+def test_kv_cache_layer_defaults_to_the_attention_module():
+    attn = SimpleNamespace(layer_name="model.layers.0.self_attn", mla_attn=object())
+    assert K3DSparkModel.kv_cache_layer(SimpleNamespace(), attn) is attn
+
+    class NestedCacheOwner(K3DSparkModel):
+        def kv_cache_layer(self, attn):
+            return attn.mla_attn
+
+    assert NestedCacheOwner.kv_cache_layer(SimpleNamespace(), attn) is attn.mla_attn
+
+
+def test_wrapper_uses_mla_attn_cls():
+    from vllm.model_executor.layers.mla import (
+        MLAModules,
+        MultiHeadLatentAttentionWrapper,
+    )
+
+    class DummyAttn(nn.Module):
+        def __init__(self, *args, **kwargs):
+            super().__init__()
+            self.prefix = kwargs["prefix"]
+
+    class Sub(MultiHeadLatentAttentionWrapper):
+        mla_attn_cls = DummyAttn
+
+    modules = MLAModules(
+        kv_a_layernorm=nn.Identity(),
+        kv_b_proj=nn.Identity(),
+        rotary_emb=None,
+        o_proj=nn.Identity(),
+        fused_qkv_a_proj=None,
+        kv_a_proj_with_mqa=nn.Identity(),
+        q_a_layernorm=None,
+        q_b_proj=None,
+        q_proj=nn.Identity(),
+        indexer=None,
+        is_sparse=False,
+        topk_indices_buffer=None,
+    )
+    wrapper = Sub(
+        hidden_size=8,
+        num_heads=2,
+        scale=1.0,
+        qk_nope_head_dim=4,
+        qk_rope_head_dim=2,
+        v_head_dim=4,
+        q_lora_rank=None,
+        kv_lora_rank=4,
+        mla_modules=modules,
+        prefix="model.layers.0.self_attn",
+    )
+    assert isinstance(wrapper.mla_attn, DummyAttn)
+    assert wrapper.mla_attn.prefix == "model.layers.0.self_attn.attn"
+
+
 def test_dspark_mla_uses_compile_free_model_entrypoint():
     assert ModelRegistry._try_load_model_cls("K3DSparkModel") is K3DSparkForCausalLM
     assert not issubclass(K3DSparkModel, TorchCompileWithNoGuardsWrapper)
@@ -121,12 +213,14 @@ def make_context_kv_proj(*args, **kwargs):
         context_kv_proj_calls.append((args, kwargs))
         return DummyModule()
 
-    monkeypatch.setattr(dspark_mla, "get_draft_quant_config", lambda _: None)
-    monkeypatch.setattr(dspark_mla, "ReplicatedLinear", DummyModule)
-    monkeypatch.setattr(dspark_mla, "MergedColumnParallelLinear", make_context_kv_proj)
-    monkeypatch.setattr(dspark_mla, "RMSNorm", DummyModule)
-    monkeypatch.setattr(dspark_mla, "K3DSparkDecoderLayer", DummyModule)
-    monkeypatch.setattr(dspark_mla, "DSparkMarkovHead", make_markov_head)
+    monkeypatch.setattr(common_dspark_mla, "get_draft_quant_config", lambda _: None)
+    monkeypatch.setattr(common_dspark_mla, "ReplicatedLinear", DummyModule)
+    monkeypatch.setattr(
+        common_dspark_mla, "MergedColumnParallelLinear", make_context_kv_proj
+    )
+    monkeypatch.setattr(common_dspark_mla, "RMSNorm", DummyModule)
+    monkeypatch.setattr(K3DSparkModel, "decoder_layer_cls", DummyModule)
+    monkeypatch.setattr(common_dspark_mla, "DSparkMarkovHead", make_markov_head)
 
     config = SimpleNamespace(
         target_hidden_size=16,

vllm/model_executor/layers/attention/mla_attention.py

@@ -310,6 +310,7 @@
 )
 
 if TYPE_CHECKING:
+    from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding
     from vllm.v1.attention.backends.mla.index_group import (
         SparseMLAIndexGroupBuilder,
     )
@@ -463,6 +464,7 @@ def __init__(
         non_causal_multi_token_decode: bool = False,
         sliding_window: int | None = None,
         prefill_backend_cls: type[MLAPrefillBackend] | None = None,
+        rotary_emb: "RotaryEmbedding | None" = None,
         **extra_impl_args,
     ):
         super().__init__()
@@ -480,6 +482,10 @@ def __init__(
         self.layer_name = prefix
         self.indexer = indexer
         self.non_causal_multi_token_decode = non_causal_multi_token_decode
+        # Stored for a platform decode epilogue under a name subclasses do not
+        # already use. The default path does not read it.
+        self._epilogue_rotary_emb = rotary_emb
+        self._rope_positions: torch.Tensor | None = None
         self.sliding_window = sliding_window
         self.num_kv_heads = 1
         self.qk_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim
@@ -794,6 +800,9 @@ def update_kv_cache(
             )
             kv_c_normed = kv_c_normed[:num_rows]
             k_pe = k_pe[:num_rows]
+        kv_c_normed, k_pe, slot_mapping = self._prepare_kv_cache_update(
+            kv_c_normed, k_pe, slot_mapping, attn_metadata
+        )
         self.impl.do_kv_cache_update(  # type: ignore[attr-defined]
             kv_c_normed,
             k_pe,
@@ -803,6 +812,61 @@ def update_kv_cache(
             k_scale,
         )
 
+    def _prepare_kv_cache_update(
+        self,
+        kv_c_normed: torch.Tensor,
+        k_pe: torch.Tensor,
+        slot_mapping: torch.Tensor,
+        attn_metadata: "MLACommonMetadata | None",
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Return the tensors written by ``do_kv_cache_update``.
+
+        The default returns its inputs unchanged. An override may drop rows it
+        will write later from ``_form_decode_q`` and may read
+        ``self._rope_positions``.
+        """
+        del attn_metadata
+        return kv_c_normed, k_pe, slot_mapping
+
+    def _prepare_mha_inputs(
+        self,
+        q: torch.Tensor,
+        k_pe: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Return the prefill query and key passed to ``forward_mha``.
+
+        The default returns its inputs unchanged. An override may read
+        ``self._rope_positions``.
+        """
+        return q, k_pe
+
+    def _form_decode_q(
+        self,
+        mqa_ql_nope: torch.Tensor,
+        mqa_q_pe: torch.Tensor,
+        k_c_normed: torch.Tensor,
+        k_pe: torch.Tensor,
+        kv_cache: torch.Tensor,
+        attn_metadata: "MLACommonMetadata",
+        num_mqa_tokens: int,
+    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
+        """Form the decode query after the absorb BMM.
+
+        The default is the existing concat and does not write the cache.
+        Overrides receive the latent KV, cache, and ``num_mqa_tokens`` so a
+        platform epilogue can replace this without editing ``forward_impl``.
+        ``self._rope_positions`` and ``self._epilogue_rotary_emb`` are available there.
+        """
+        del k_c_normed, k_pe, kv_cache, attn_metadata, num_mqa_tokens
+        fp8_attention = is_quantized_kv_cache(self.kv_cache_dtype)
+        if fp8_attention and self.impl.supports_quant_query_input:
+            assert mqa_ql_nope.shape[0] == mqa_q_pe.shape[0]
+            assert mqa_ql_nope.shape[1] == mqa_q_pe.shape[1]
+            return self._decode_concat_quant_fp8_op(
+                mqa_ql_nope, mqa_q_pe, self._q_scale
+            )
+        return (mqa_ql_nope, mqa_q_pe)
+
     def prepare_kv_cache_update(
         self, attn_metadata: "MLACommonMetadata | None"
     ) -> None:
@@ -817,7 +881,15 @@ def forward(
         k_pe: torch.Tensor,
         output_shape: torch.Size | None = None,
         q_dcp_replicated: torch.Tensor | None = None,
+        positions: torch.Tensor | None = None,
     ) -> torch.Tensor:
+        # Assigned on every call, including None, so a previous forward cannot
+        # leave a stale tensor for the compiled KV-update op to read.
+        # CUDA-graph replay sees new values only while this object stays the
+        # persistent buffer the runner updates in place (InputBuffers.positions,
+        # captured as input_buffers.positions[:num_tokens]). Replacing the
+        # tensor freezes the captured pointer.
+        self._rope_positions = positions
         if self.use_direct_call:
             forward_context: ForwardContext = get_forward_context()
             attn_metadata_raw = forward_context.attn_metadata
@@ -831,6 +903,9 @@ def forward(
             else:
                 attn_metadata = attn_metadata_raw
             self_kv_cache = self.kv_cache
+            # Same replay rule as _rope_positions: forward_context.slot_mapping
+            # is the persistent buffer the runner updates in place. Replacing
+            # that tensor freezes the captured pointer.
             slot_mapping = forward_context.slot_mapping
 
             assert isinstance(slot_mapping, dict), (
@@ -999,10 +1074,13 @@ def forward_impl(
                 mha_output = output
                 mha_output_scale = None
 
+            mha_q, mha_k_pe = self._prepare_mha_inputs(
+                q[num_mqa_tokens:], k_pe[num_mqa_tokens:]
+            )
             self.impl.forward_mha(  # type: ignore[attr-defined]
-                q[num_mqa_tokens:],
+                mha_q,
                 k_c_normed[num_mqa_tokens:],
-                k_pe[num_mqa_tokens:],
+                mha_k_pe,
                 kv_cache,
                 attn_metadata,
                 self._k_scale,
@@ -1087,14 +1165,15 @@ def forward_impl(
                     mqa_ql_nope = mqa_q_nope.new_empty((B, N, L))
                     torch.bmm(mqa_q_nope, W_UK_T, out=mqa_ql_nope.transpose(0, 1))
 
-            if fp8_attention and self.impl.supports_quant_query_input:
-                assert mqa_ql_nope.shape[0] == mqa_q_pe.shape[0]
-                assert mqa_ql_nope.shape[1] == mqa_q_pe.shape[1]
-                mqa_q = self._decode_concat_quant_fp8_op(
-                    mqa_ql_nope, mqa_q_pe, self._q_scale
-                )
-            else:
-                mqa_q = (mqa_ql_nope, mqa_q_pe)
+            mqa_q = self._form_decode_q(
+                mqa_ql_nope,
+                mqa_q_pe,
+                k_c_normed,
+                k_pe,
+                kv_cache,
+                attn_metadata,
+                num_mqa_tokens,
+            )
             # concatenate nope + pe -> (B, N, L + P) (fp8 op above may have fused)
             if self.impl.dcp_world_size > 1:
                 assert self.dcp_manager is not None

vllm/model_executor/layers/mla.py

@@ -60,6 +60,9 @@ class MultiHeadLatentAttentionWrapper(PluggableLayer):
 
     # --8<-- [end:multi_head_latent_attention]
 
+    # Subclasses replace this to supply a platform decode epilogue.
+    mla_attn_cls: type[MLAAttention] = MLAAttention
+
     def __init__(
         self,
         hidden_size: int,
@@ -119,7 +122,7 @@ def __init__(
             self.topk_tokens = self.indexer.topk_tokens
             self.topk_indices_buffer = mla_modules.topk_indices_buffer
 
-        self.mla_attn = MLAAttention(
+        self.mla_attn = self.mla_attn_cls(
             num_heads=self.num_heads,
             scale=scale,
             qk_nope_head_dim=self.qk_nope_head_dim,