From a8a4272f8ac40f6d4a98a79b842891c228cbd9fc Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Stig-Arne=20Gr=C3=B6nroos?= Date: Fri, 2 Oct 2026 10:28:11 +0300 Subject: [PATCH 1/7] [Model] Move Kimi-K3 DSpark MLA to the common package MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Keep the NVIDIA module as a re-export so the registry path is unchanged. Signed-off-by: Stig-Arne Grönroos Co-authored-by: Cursor --- vllm/models/kimi_k3/common/dspark_mla.py | 537 ++++++++++++++++++++++ vllm/models/kimi_k3/nvidia/dspark_mla.py | 545 +---------------------- 2 files changed, 551 insertions(+), 531 deletions(-) create mode 100644 vllm/models/kimi_k3/common/dspark_mla.py diff --git a/vllm/models/kimi_k3/common/dspark_mla.py b/vllm/models/kimi_k3/common/dspark_mla.py new file mode 100644 index 000000000000..37207d7bc9cf --- /dev/null +++ b/vllm/models/kimi_k3/common/dspark_mla.py @@ -0,0 +1,537 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""K3 dense MLA draft model for DSpark speculative decoding.""" + +from collections.abc import Iterable + +import torch +import torch.nn as nn + +import vllm._custom_ops as ops +from vllm.config import VllmConfig +from vllm.model_executor.layers.layernorm import RMSNorm +from vllm.model_executor.layers.linear import ( + MergedColumnParallelLinear, + ReplicatedLinear, +) +from vllm.model_executor.layers.logits_processor import LogitsProcessor +from vllm.model_executor.models.qwen3_dspark import ( + DSparkConfidenceHead, + DSparkMarkovHead, +) +from vllm.model_executor.models.utils import ( + AutoWeightsLoader, + WeightsMapper, + get_draft_quant_config, + maybe_prefix, +) +from vllm.models.common.ops.fused_allreduce_rms_norm import fused_allreduce_rms_norm +from vllm.models.kimi_k3.nvidia.mla import MultiHeadLatentAttention +from vllm.models.kimi_k3.nvidia.model import KimiMLP +from vllm.utils.torch_utils import is_quantized_kv_cache +from vllm.v1.worker.workspace import current_workspace_manager + +_GROUPED_KV_CACHE_DTYPES = frozenset( + {"auto", "bfloat16", "fp8", "fp8_e4m3", "fp8_e5m2"} +) + + +def _duplicate_context_kv_weights( + weights: Iterable[tuple[str, torch.Tensor]], num_layers: int +) -> Iterable[tuple[str, torch.Tensor]]: + """Load each layer's KV projection into the cross-layer linear.""" + for name, weight in weights: + yield name, weight + layer_prefix, marker, param_name = name.partition( + ".self_attn.kv_a_proj_with_mqa." + ) + if not marker: + continue + layer_idx_str = layer_prefix.rsplit(".", 1)[-1] + if not layer_idx_str.isdecimal(): + continue + layer_idx = int(layer_idx_str) + if layer_idx >= num_layers: + continue + fused_weight = weight.detach() + fused_weight.shard_id = layer_idx + yield f"context_kv_proj.{param_name}", fused_weight + + +class K3DSparkDecoderLayer(nn.Module): + def __init__( + self, + *, + vllm_config: VllmConfig, + config, + layer_idx: int, + start_layer_id: int, + prefix: str, + ) -> None: + super().__init__() + quant_config = get_draft_quant_config(vllm_config) + self.self_attn = MultiHeadLatentAttention( + config=config, + hidden_size=config.hidden_size, + num_heads=config.num_attention_heads, + qk_nope_head_dim=config.qk_nope_head_dim, + qk_rope_head_dim=config.qk_rope_head_dim, + v_head_dim=config.v_head_dim, + q_lora_rank=config.q_lora_rank, + kv_lora_rank=config.kv_lora_rank, + cache_config=vllm_config.cache_config, + quant_config=quant_config, + prefix=maybe_prefix( + prefix, f"layers.{start_layer_id + layer_idx}.self_attn" + ), + use_rope=True, + non_causal_multi_token_decode=True, + ) + # Both row-parallel outputs stay un-reduced; their all-reduces are fused + # into the RMSNorm that follows via fused_allreduce_rms_norm. + self.self_attn.o_proj.reduce_results = False + self.mlp = KimiMLP( + hidden_size=config.hidden_size, + intermediate_size=config.intermediate_size, + hidden_act=config.hidden_act, + quant_config=quant_config, + reduce_results=False, + prefix=maybe_prefix(prefix, f"layers.{start_layer_id + layer_idx}.mlp"), + ) + self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + residual: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor]: + if residual is None: + # First layer: hidden_states is the (already reduced) embedding. + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + else: + hidden_states, residual = fused_allreduce_rms_norm( + hidden_states, residual, self.input_layernorm + ) + + hidden_states = self.self_attn( + positions=positions, + hidden_states=hidden_states, + ) + hidden_states, residual = fused_allreduce_rms_norm( + hidden_states, residual, self.post_attention_layernorm + ) + # The MLP output is reduced by the next layer's input_layernorm (or by + # the model's final_norm). + hidden_states = self.mlp(hidden_states) + return hidden_states, residual + + +class K3DSparkModel(nn.Module): + def __init__( + self, + *, + vllm_config: VllmConfig, + start_layer_id: int, + prefix: str, + ) -> None: + super().__init__() + assert vllm_config.speculative_config is not None + self.config = vllm_config.speculative_config.draft_model_config.hf_config + self.quant_config = get_draft_quant_config(vllm_config) + + # The frozen target embedding is aliased after the draft checkpoint loads. + self.embed_tokens: nn.Module | None = None + + self.context_proj = ReplicatedLinear( + self.config.target_hidden_size * self.config.num_target_layers, + self.config.hidden_size, + bias=False, + return_bias=False, + quant_config=self.quant_config, + prefix=maybe_prefix(prefix, "context_proj"), + ) + self.context_norm = RMSNorm( + self.config.hidden_size, eps=self.config.rms_norm_eps + ) + + self.layers = nn.ModuleList( + [ + K3DSparkDecoderLayer( + vllm_config=vllm_config, + config=self.config, + layer_idx=layer_idx, + start_layer_id=start_layer_id, + prefix=prefix, + ) + for layer_idx in range(self.config.num_hidden_layers) + ] + ) + kv_width = self.config.kv_lora_rank + self.config.qk_rope_head_dim + self.context_kv_proj = MergedColumnParallelLinear( + self.config.hidden_size, + [kv_width] * self.config.num_hidden_layers, + bias=False, + return_bias=False, + quant_config=self.quant_config, + prefix=maybe_prefix( + prefix, + f"layers.{start_layer_id}.self_attn.fused_qkv_a_proj", + ), + disable_tp=True, + ) + self.final_norm = RMSNorm(self.config.hidden_size, eps=self.config.rms_norm_eps) + self.markov_head = DSparkMarkovHead( + self.config.vocab_size, + self.config.draft_vocab_size, + self.config.markov_rank, + prefix=maybe_prefix(prefix, "markov_head"), + ) + self.confidence_head: DSparkConfidenceHead | None = None + if getattr(self.config, "enable_confidence_head", False): + with_markov = getattr(self.config, "confidence_head_with_markov", False) + input_dim = self.config.hidden_size + ( + self.config.markov_rank if with_markov else 0 + ) + self.confidence_head = DSparkConfidenceHead( + input_dim, + prefix=maybe_prefix(prefix, "confidence_head"), + bias=True, + with_markov=with_markov, + ) + self._max_num_context_tokens = ( + vllm_config.scheduler_config.max_num_batched_tokens + ) + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + assert self.embed_tokens is not None + return self.embed_tokens(input_ids) + + def combine_hidden_states(self, hidden_states: torch.Tensor) -> torch.Tensor: + return self.context_norm(self.context_proj(hidden_states)) + + @torch.inference_mode() + def precompute_and_store_context_kv( + self, + context_states: torch.Tensor, + context_positions: torch.Tensor, + context_slot_mapping: torch.Tensor | list[torch.Tensor | None] | None = None, + ) -> None: + """Project target-derived context into each draft layer's latent cache.""" + if not hasattr(self, "_num_context_layers"): + self._build_fused_context_kv_metadata() + self._precompute_fused_context_kv( + context_states, context_positions, context_slot_mapping + ) + + def _build_fused_context_kv_metadata(self) -> None: + """Build cross-layer metadata after checkpoint loading.""" + attentions = [layer.self_attn for layer in self.layers] + assert attentions + attn0 = attentions[0] + assert attn0.q_lora_rank is not None + kv_width = attn0.kv_lora_rank + attn0.qk_rope_head_dim + for attn in attentions: + assert attn.q_lora_rank is not None + assert ( + attn.q_lora_rank == attn0.q_lora_rank + and attn.kv_lora_rank == attn0.kv_lora_rank + and attn.qk_rope_head_dim == attn0.qk_rope_head_dim + and attn.kv_a_layernorm.variance_epsilon + == attn0.kv_a_layernorm.variance_epsilon + ), "All MLA DSpark layers must share their latent KV geometry." + self._context_kv_norm_weights = torch.stack( + [attn.kv_a_layernorm.weight.detach() for attn in attentions], dim=0 + ).contiguous() + self._num_context_layers = len(attentions) + self._context_kv_width = kv_width + self._context_kv_lora_rank = attn0.kv_lora_rank + self._context_rope_dim = attn0.qk_rope_head_dim + self._context_rms_norm_eps = attn0.kv_a_layernorm.variance_epsilon + self._context_kv_scales: torch.Tensor | None = None + if attn0.kv_cache_dtype in _GROUPED_KV_CACHE_DTYPES and is_quantized_kv_cache( + attn0.kv_cache_dtype + ): + self._context_kv_scales = torch.stack( + [attn._k_scale.reshape(()) for attn in attentions] + ) + + def _precompute_fused_context_kv( + self, + context_states: torch.Tensor, + context_positions: torch.Tensor, + context_slot_mapping: torch.Tensor | list[torch.Tensor | None] | None, + ) -> None: + num_ctx = context_states.shape[0] + num_layers = self._num_context_layers + + # One KV-only GEMM replaces five full Q+KV GEMMs. For K3 this projects + # 5*576 rows rather than 5*2112 rows (72.7% fewer A-projection FLOPs). + all_kv = self.context_kv_proj(context_states) + all_kv = all_kv.view(num_ctx, num_layers, self._context_kv_width) + all_kv_c = all_kv[..., : self._context_kv_lora_rank] + all_k_pe = all_kv[..., self._context_kv_lora_rank :] + + # Layer-major layout lets the 2-D RMSNorm weights select a distinct row + # for each draft layer in one grouped kernel. + all_kv_c = all_kv_c.permute(1, 0, 2).contiguous() + all_kv_c_normed = torch.empty_like(all_kv_c) + ops.rms_norm( + all_kv_c_normed, + all_kv_c, + self._context_kv_norm_weights, + self._context_rms_norm_eps, + ) + + all_k_pe = all_k_pe.permute(1, 0, 2).contiguous() + all_k_pe_flat = all_k_pe.view(num_layers * num_ctx, 1, self._context_rope_dim) + (repeated_positions,) = current_workspace_manager().get_simultaneous( + ((num_layers * self._max_num_context_tokens,), torch.int64), + ) + repeated_positions = repeated_positions[: num_layers * num_ctx] + repeated_positions.view(num_layers, num_ctx).copy_(context_positions) + # Keep the single-tensor context RoPE on vLLM's optimized CUDA op; + # DeepSeek YaRN's FlashInfer wrapper assumes a non-null key tensor. + rotary_emb = self.layers[0].self_attn.rotary_emb + assert rotary_emb is not None + ops.rotary_embedding( + repeated_positions, + all_k_pe_flat, + None, + rotary_emb.head_size, + rotary_emb.cos_sin_cache, + rotary_emb.is_neox_style, + ) + all_k_pe = all_k_pe_flat.view(num_layers, num_ctx, 1, self._context_rope_dim) + + if context_slot_mapping is None: + return + + cache_layers = [layer.self_attn for layer in self.layers] + cache_dtype = cache_layers[0].kv_cache_dtype + if ( + cache_dtype in _GROUPED_KV_CACHE_DTYPES + and all(cl.kv_cache_dtype == cache_dtype for cl in cache_layers) + and self._has_uniform_block_layout(cache_layers) + and ( + isinstance(context_slot_mapping, torch.Tensor) + or all(s is not None for s in context_slot_mapping) + ) + ): + # Grouped context KV insert assumes that all layers share the same + # cache dtype and block layout. + + if isinstance(context_slot_mapping, (list, tuple)): + per_layer_slot_mappings = [ + s for s in context_slot_mapping if s is not None + ] + if len({s.data_ptr() for s in per_layer_slot_mappings}) == 1: + # All rows alias to the same slot mapping. + slot_mapping = ( + per_layer_slot_mappings[0].unsqueeze(0).expand(num_layers, -1) + ) + else: + slot_mapping = torch.stack(per_layer_slot_mappings, dim=0) + else: + # Broadcast the single shared context_slot_mapping tensor. + slot_mapping = context_slot_mapping.unsqueeze(0).expand(num_layers, -1) + + ref_cache = cache_layers[0].kv_cache + ops.concat_and_cache_mla_grouped( + all_kv_c_normed, + all_k_pe.squeeze(2), + self._get_context_kv_cache_ptrs(cache_layers), + slot_mapping, + ref_cache.size(1), + ref_cache.stride(0), + ref_cache.stride(1), + self._context_kv_scales, + cache_dtype, + ) + return + + for layer_idx, layer in enumerate(self.layers): + slot_mapping = ( + context_slot_mapping[layer_idx] + if isinstance(context_slot_mapping, (list, tuple)) + else context_slot_mapping + ) + if slot_mapping is None: + continue + attn = layer.self_attn + attn.impl.do_kv_cache_update( + all_kv_c_normed[layer_idx], + all_k_pe[layer_idx], + attn.kv_cache, + slot_mapping, + attn.kv_cache_dtype, + attn._k_scale, + ) + + def _has_uniform_block_layout( + self, + cache_layers: list[MultiHeadLatentAttention], + ) -> bool: + key = tuple(cl.kv_cache.data_ptr() for cl in cache_layers) + if getattr(self, "_kv_cache_ptrs_key", None) != key: + assert not torch.cuda.is_current_stream_capturing() + self._kv_cache_ptrs_key = key + if hasattr(self, "_context_cache_ptrs"): + del self._context_cache_ptrs + ref_cache = cache_layers[0].kv_cache + self._layers_share_kv_block_layout = all( + cl.kv_cache.size(1) == ref_cache.size(1) + and cl.kv_cache.stride(0) == ref_cache.stride(0) + and cl.kv_cache.stride(1) == ref_cache.stride(1) + for cl in cache_layers + ) + return self._layers_share_kv_block_layout + + def _get_context_kv_cache_ptrs( + self, + cache_layers: list[MultiHeadLatentAttention], + ) -> torch.Tensor: + # The per-layer KV cache base pointers are stable after allocation, so + # build the pointer array once and return it on every call. + if not hasattr(self, "_context_cache_ptrs"): + ref_cache = cache_layers[0].kv_cache + cache_ptrs = torch.tensor( + [cl.kv_cache.data_ptr() for cl in cache_layers], + dtype=torch.int64, + device=ref_cache.device, + ) + self._context_cache_ptrs = cache_ptrs + return self._context_cache_ptrs + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + inputs_embeds: torch.Tensor | None = None, + ) -> torch.Tensor: + if inputs_embeds is None: + inputs_embeds = self.embed_input_ids(input_ids) + + hidden_states = inputs_embeds + residual = None + for layer in self.layers: + hidden_states, residual = layer( + positions=positions, + hidden_states=hidden_states, + residual=residual, + ) + hidden_states, _ = fused_allreduce_rms_norm( + hidden_states, residual, self.final_norm + ) + return hidden_states + + +class K3DSparkForCausalLM(nn.Module): + has_own_embed_tokens = False + has_own_lm_head = False + draft_id_to_target_id = None + hf_to_vllm_mapper = WeightsMapper( + # confidence_head is training-only. The frozen target embedding and LM + # head are shared after this draft-specific checkpoint is loaded. + orig_to_new_substr={ + "confidence_head": None, + "embed_tokens": None, + "lm_head": None, + }, + orig_to_new_prefix={"": "model."}, + orig_to_new_stacked={ + ".gate_proj": (".gate_up_proj", 0), + ".up_proj": (".gate_up_proj", 1), + ".q_a_proj": (".fused_qkv_a_proj", 0), + ".kv_a_proj_with_mqa": (".fused_qkv_a_proj", 1), + }, + ) + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: + super().__init__() + assert vllm_config.speculative_config is not None + self.draft_model_config = vllm_config.speculative_config.draft_model_config + self.config = self.draft_model_config.hf_config + if getattr(self.config, "enable_confidence_head", False): + # The class mapper drops the head as training-only; keep it for drafts + # that actually built one, since adaptive verification reads it. + self.hf_to_vllm_mapper = self.hf_to_vllm_mapper | WeightsMapper( + orig_to_new_substr={"confidence_head": "confidence_head"} + ) + target_layer_num = vllm_config.model_config.get_total_num_hidden_layers() + self.model = K3DSparkModel( + vllm_config=vllm_config, + start_layer_id=target_layer_num, + prefix=maybe_prefix(prefix, "model"), + ) + + # Assigned by load_dspark_model from the target. Keeping no placeholder + # avoids a transient full-vocabulary allocation for this 163k-vocab model. + self.lm_head: nn.Module | None = None + logit_scale = getattr(self.config, "logit_scale", 1.0) + self.logits_processor = LogitsProcessor( + self.config.draft_vocab_size, scale=logit_scale + ) + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.model.embed_input_ids(input_ids) + + def combine_hidden_states(self, hidden_states: torch.Tensor) -> torch.Tensor: + return self.model.combine_hidden_states(hidden_states) + + def get_draft_kv_cache_layer_names(self) -> list[str]: + return [layer.self_attn.layer_name for layer in self.model.layers] + + def precompute_and_store_context_kv( + self, + context_states: torch.Tensor, + context_positions: torch.Tensor, + context_slot_mapping: torch.Tensor | list[torch.Tensor | None] | None = None, + ) -> None: + self.model.precompute_and_store_context_kv( + context_states, context_positions, context_slot_mapping + ) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + inputs_embeds: torch.Tensor | None = None, + ) -> torch.Tensor: + return self.model(input_ids, positions, inputs_embeds) + + def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor: + assert self.lm_head is not None + return self.logits_processor(self.lm_head, hidden_states) + + def compute_draft_logits(self, hidden_states: torch.Tensor) -> torch.Tensor: + return self.compute_logits(hidden_states) + + def map_draft_to_target(self, draft_ids: torch.Tensor) -> torch.Tensor: + return draft_ids + + def markov_embed(self, token_ids: torch.Tensor) -> torch.Tensor: + return self.model.markov_head.embed(token_ids) + + def markov_bias(self, markov_embed: torch.Tensor) -> torch.Tensor: + return self.model.markov_head.bias(markov_embed, self.logits_processor) + + def compute_confidence( + self, head_hidden: torch.Tensor, markov_embed: torch.Tensor + ) -> torch.Tensor: + """Per-position acceptance probability for each drafted token.""" + assert self.model.confidence_head is not None + return torch.sigmoid(self.model.confidence_head(head_hidden, markov_embed)) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + loader = AutoWeightsLoader(self) + # read: 1. all weights. 2. context kv weights + weights = _duplicate_context_kv_weights(weights, len(self.model.layers)) + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) + + def process_weights_after_loading(self) -> None: + self.model._build_fused_context_kv_metadata() diff --git a/vllm/models/kimi_k3/nvidia/dspark_mla.py b/vllm/models/kimi_k3/nvidia/dspark_mla.py index 37207d7bc9cf..c360a23b0657 100644 --- a/vllm/models/kimi_k3/nvidia/dspark_mla.py +++ b/vllm/models/kimi_k3/nvidia/dspark_mla.py @@ -1,537 +1,20 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""K3 dense MLA draft model for DSpark speculative decoding.""" +"""NVIDIA entry point for the Kimi-K3 DSpark MLA draft. -from collections.abc import Iterable +The implementation lives in ``vllm.models.kimi_k3.common.dspark_mla``. +""" -import torch -import torch.nn as nn - -import vllm._custom_ops as ops -from vllm.config import VllmConfig -from vllm.model_executor.layers.layernorm import RMSNorm -from vllm.model_executor.layers.linear import ( - MergedColumnParallelLinear, - ReplicatedLinear, -) -from vllm.model_executor.layers.logits_processor import LogitsProcessor -from vllm.model_executor.models.qwen3_dspark import ( - DSparkConfidenceHead, - DSparkMarkovHead, +from vllm.models.kimi_k3.common.dspark_mla import ( + K3DSparkDecoderLayer, + K3DSparkForCausalLM, + K3DSparkModel, + _duplicate_context_kv_weights, ) -from vllm.model_executor.models.utils import ( - AutoWeightsLoader, - WeightsMapper, - get_draft_quant_config, - maybe_prefix, -) -from vllm.models.common.ops.fused_allreduce_rms_norm import fused_allreduce_rms_norm -from vllm.models.kimi_k3.nvidia.mla import MultiHeadLatentAttention -from vllm.models.kimi_k3.nvidia.model import KimiMLP -from vllm.utils.torch_utils import is_quantized_kv_cache -from vllm.v1.worker.workspace import current_workspace_manager - -_GROUPED_KV_CACHE_DTYPES = frozenset( - {"auto", "bfloat16", "fp8", "fp8_e4m3", "fp8_e5m2"} -) - - -def _duplicate_context_kv_weights( - weights: Iterable[tuple[str, torch.Tensor]], num_layers: int -) -> Iterable[tuple[str, torch.Tensor]]: - """Load each layer's KV projection into the cross-layer linear.""" - for name, weight in weights: - yield name, weight - layer_prefix, marker, param_name = name.partition( - ".self_attn.kv_a_proj_with_mqa." - ) - if not marker: - continue - layer_idx_str = layer_prefix.rsplit(".", 1)[-1] - if not layer_idx_str.isdecimal(): - continue - layer_idx = int(layer_idx_str) - if layer_idx >= num_layers: - continue - fused_weight = weight.detach() - fused_weight.shard_id = layer_idx - yield f"context_kv_proj.{param_name}", fused_weight - - -class K3DSparkDecoderLayer(nn.Module): - def __init__( - self, - *, - vllm_config: VllmConfig, - config, - layer_idx: int, - start_layer_id: int, - prefix: str, - ) -> None: - super().__init__() - quant_config = get_draft_quant_config(vllm_config) - self.self_attn = MultiHeadLatentAttention( - config=config, - hidden_size=config.hidden_size, - num_heads=config.num_attention_heads, - qk_nope_head_dim=config.qk_nope_head_dim, - qk_rope_head_dim=config.qk_rope_head_dim, - v_head_dim=config.v_head_dim, - q_lora_rank=config.q_lora_rank, - kv_lora_rank=config.kv_lora_rank, - cache_config=vllm_config.cache_config, - quant_config=quant_config, - prefix=maybe_prefix( - prefix, f"layers.{start_layer_id + layer_idx}.self_attn" - ), - use_rope=True, - non_causal_multi_token_decode=True, - ) - # Both row-parallel outputs stay un-reduced; their all-reduces are fused - # into the RMSNorm that follows via fused_allreduce_rms_norm. - self.self_attn.o_proj.reduce_results = False - self.mlp = KimiMLP( - hidden_size=config.hidden_size, - intermediate_size=config.intermediate_size, - hidden_act=config.hidden_act, - quant_config=quant_config, - reduce_results=False, - prefix=maybe_prefix(prefix, f"layers.{start_layer_id + layer_idx}.mlp"), - ) - self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) - self.post_attention_layernorm = RMSNorm( - config.hidden_size, eps=config.rms_norm_eps - ) - - def forward( - self, - positions: torch.Tensor, - hidden_states: torch.Tensor, - residual: torch.Tensor | None, - ) -> tuple[torch.Tensor, torch.Tensor]: - if residual is None: - # First layer: hidden_states is the (already reduced) embedding. - residual = hidden_states - hidden_states = self.input_layernorm(hidden_states) - else: - hidden_states, residual = fused_allreduce_rms_norm( - hidden_states, residual, self.input_layernorm - ) - - hidden_states = self.self_attn( - positions=positions, - hidden_states=hidden_states, - ) - hidden_states, residual = fused_allreduce_rms_norm( - hidden_states, residual, self.post_attention_layernorm - ) - # The MLP output is reduced by the next layer's input_layernorm (or by - # the model's final_norm). - hidden_states = self.mlp(hidden_states) - return hidden_states, residual - - -class K3DSparkModel(nn.Module): - def __init__( - self, - *, - vllm_config: VllmConfig, - start_layer_id: int, - prefix: str, - ) -> None: - super().__init__() - assert vllm_config.speculative_config is not None - self.config = vllm_config.speculative_config.draft_model_config.hf_config - self.quant_config = get_draft_quant_config(vllm_config) - - # The frozen target embedding is aliased after the draft checkpoint loads. - self.embed_tokens: nn.Module | None = None - - self.context_proj = ReplicatedLinear( - self.config.target_hidden_size * self.config.num_target_layers, - self.config.hidden_size, - bias=False, - return_bias=False, - quant_config=self.quant_config, - prefix=maybe_prefix(prefix, "context_proj"), - ) - self.context_norm = RMSNorm( - self.config.hidden_size, eps=self.config.rms_norm_eps - ) - - self.layers = nn.ModuleList( - [ - K3DSparkDecoderLayer( - vllm_config=vllm_config, - config=self.config, - layer_idx=layer_idx, - start_layer_id=start_layer_id, - prefix=prefix, - ) - for layer_idx in range(self.config.num_hidden_layers) - ] - ) - kv_width = self.config.kv_lora_rank + self.config.qk_rope_head_dim - self.context_kv_proj = MergedColumnParallelLinear( - self.config.hidden_size, - [kv_width] * self.config.num_hidden_layers, - bias=False, - return_bias=False, - quant_config=self.quant_config, - prefix=maybe_prefix( - prefix, - f"layers.{start_layer_id}.self_attn.fused_qkv_a_proj", - ), - disable_tp=True, - ) - self.final_norm = RMSNorm(self.config.hidden_size, eps=self.config.rms_norm_eps) - self.markov_head = DSparkMarkovHead( - self.config.vocab_size, - self.config.draft_vocab_size, - self.config.markov_rank, - prefix=maybe_prefix(prefix, "markov_head"), - ) - self.confidence_head: DSparkConfidenceHead | None = None - if getattr(self.config, "enable_confidence_head", False): - with_markov = getattr(self.config, "confidence_head_with_markov", False) - input_dim = self.config.hidden_size + ( - self.config.markov_rank if with_markov else 0 - ) - self.confidence_head = DSparkConfidenceHead( - input_dim, - prefix=maybe_prefix(prefix, "confidence_head"), - bias=True, - with_markov=with_markov, - ) - self._max_num_context_tokens = ( - vllm_config.scheduler_config.max_num_batched_tokens - ) - - def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: - assert self.embed_tokens is not None - return self.embed_tokens(input_ids) - - def combine_hidden_states(self, hidden_states: torch.Tensor) -> torch.Tensor: - return self.context_norm(self.context_proj(hidden_states)) - - @torch.inference_mode() - def precompute_and_store_context_kv( - self, - context_states: torch.Tensor, - context_positions: torch.Tensor, - context_slot_mapping: torch.Tensor | list[torch.Tensor | None] | None = None, - ) -> None: - """Project target-derived context into each draft layer's latent cache.""" - if not hasattr(self, "_num_context_layers"): - self._build_fused_context_kv_metadata() - self._precompute_fused_context_kv( - context_states, context_positions, context_slot_mapping - ) - - def _build_fused_context_kv_metadata(self) -> None: - """Build cross-layer metadata after checkpoint loading.""" - attentions = [layer.self_attn for layer in self.layers] - assert attentions - attn0 = attentions[0] - assert attn0.q_lora_rank is not None - kv_width = attn0.kv_lora_rank + attn0.qk_rope_head_dim - for attn in attentions: - assert attn.q_lora_rank is not None - assert ( - attn.q_lora_rank == attn0.q_lora_rank - and attn.kv_lora_rank == attn0.kv_lora_rank - and attn.qk_rope_head_dim == attn0.qk_rope_head_dim - and attn.kv_a_layernorm.variance_epsilon - == attn0.kv_a_layernorm.variance_epsilon - ), "All MLA DSpark layers must share their latent KV geometry." - self._context_kv_norm_weights = torch.stack( - [attn.kv_a_layernorm.weight.detach() for attn in attentions], dim=0 - ).contiguous() - self._num_context_layers = len(attentions) - self._context_kv_width = kv_width - self._context_kv_lora_rank = attn0.kv_lora_rank - self._context_rope_dim = attn0.qk_rope_head_dim - self._context_rms_norm_eps = attn0.kv_a_layernorm.variance_epsilon - self._context_kv_scales: torch.Tensor | None = None - if attn0.kv_cache_dtype in _GROUPED_KV_CACHE_DTYPES and is_quantized_kv_cache( - attn0.kv_cache_dtype - ): - self._context_kv_scales = torch.stack( - [attn._k_scale.reshape(()) for attn in attentions] - ) - - def _precompute_fused_context_kv( - self, - context_states: torch.Tensor, - context_positions: torch.Tensor, - context_slot_mapping: torch.Tensor | list[torch.Tensor | None] | None, - ) -> None: - num_ctx = context_states.shape[0] - num_layers = self._num_context_layers - - # One KV-only GEMM replaces five full Q+KV GEMMs. For K3 this projects - # 5*576 rows rather than 5*2112 rows (72.7% fewer A-projection FLOPs). - all_kv = self.context_kv_proj(context_states) - all_kv = all_kv.view(num_ctx, num_layers, self._context_kv_width) - all_kv_c = all_kv[..., : self._context_kv_lora_rank] - all_k_pe = all_kv[..., self._context_kv_lora_rank :] - - # Layer-major layout lets the 2-D RMSNorm weights select a distinct row - # for each draft layer in one grouped kernel. - all_kv_c = all_kv_c.permute(1, 0, 2).contiguous() - all_kv_c_normed = torch.empty_like(all_kv_c) - ops.rms_norm( - all_kv_c_normed, - all_kv_c, - self._context_kv_norm_weights, - self._context_rms_norm_eps, - ) - - all_k_pe = all_k_pe.permute(1, 0, 2).contiguous() - all_k_pe_flat = all_k_pe.view(num_layers * num_ctx, 1, self._context_rope_dim) - (repeated_positions,) = current_workspace_manager().get_simultaneous( - ((num_layers * self._max_num_context_tokens,), torch.int64), - ) - repeated_positions = repeated_positions[: num_layers * num_ctx] - repeated_positions.view(num_layers, num_ctx).copy_(context_positions) - # Keep the single-tensor context RoPE on vLLM's optimized CUDA op; - # DeepSeek YaRN's FlashInfer wrapper assumes a non-null key tensor. - rotary_emb = self.layers[0].self_attn.rotary_emb - assert rotary_emb is not None - ops.rotary_embedding( - repeated_positions, - all_k_pe_flat, - None, - rotary_emb.head_size, - rotary_emb.cos_sin_cache, - rotary_emb.is_neox_style, - ) - all_k_pe = all_k_pe_flat.view(num_layers, num_ctx, 1, self._context_rope_dim) - - if context_slot_mapping is None: - return - - cache_layers = [layer.self_attn for layer in self.layers] - cache_dtype = cache_layers[0].kv_cache_dtype - if ( - cache_dtype in _GROUPED_KV_CACHE_DTYPES - and all(cl.kv_cache_dtype == cache_dtype for cl in cache_layers) - and self._has_uniform_block_layout(cache_layers) - and ( - isinstance(context_slot_mapping, torch.Tensor) - or all(s is not None for s in context_slot_mapping) - ) - ): - # Grouped context KV insert assumes that all layers share the same - # cache dtype and block layout. - - if isinstance(context_slot_mapping, (list, tuple)): - per_layer_slot_mappings = [ - s for s in context_slot_mapping if s is not None - ] - if len({s.data_ptr() for s in per_layer_slot_mappings}) == 1: - # All rows alias to the same slot mapping. - slot_mapping = ( - per_layer_slot_mappings[0].unsqueeze(0).expand(num_layers, -1) - ) - else: - slot_mapping = torch.stack(per_layer_slot_mappings, dim=0) - else: - # Broadcast the single shared context_slot_mapping tensor. - slot_mapping = context_slot_mapping.unsqueeze(0).expand(num_layers, -1) - - ref_cache = cache_layers[0].kv_cache - ops.concat_and_cache_mla_grouped( - all_kv_c_normed, - all_k_pe.squeeze(2), - self._get_context_kv_cache_ptrs(cache_layers), - slot_mapping, - ref_cache.size(1), - ref_cache.stride(0), - ref_cache.stride(1), - self._context_kv_scales, - cache_dtype, - ) - return - - for layer_idx, layer in enumerate(self.layers): - slot_mapping = ( - context_slot_mapping[layer_idx] - if isinstance(context_slot_mapping, (list, tuple)) - else context_slot_mapping - ) - if slot_mapping is None: - continue - attn = layer.self_attn - attn.impl.do_kv_cache_update( - all_kv_c_normed[layer_idx], - all_k_pe[layer_idx], - attn.kv_cache, - slot_mapping, - attn.kv_cache_dtype, - attn._k_scale, - ) - - def _has_uniform_block_layout( - self, - cache_layers: list[MultiHeadLatentAttention], - ) -> bool: - key = tuple(cl.kv_cache.data_ptr() for cl in cache_layers) - if getattr(self, "_kv_cache_ptrs_key", None) != key: - assert not torch.cuda.is_current_stream_capturing() - self._kv_cache_ptrs_key = key - if hasattr(self, "_context_cache_ptrs"): - del self._context_cache_ptrs - ref_cache = cache_layers[0].kv_cache - self._layers_share_kv_block_layout = all( - cl.kv_cache.size(1) == ref_cache.size(1) - and cl.kv_cache.stride(0) == ref_cache.stride(0) - and cl.kv_cache.stride(1) == ref_cache.stride(1) - for cl in cache_layers - ) - return self._layers_share_kv_block_layout - - def _get_context_kv_cache_ptrs( - self, - cache_layers: list[MultiHeadLatentAttention], - ) -> torch.Tensor: - # The per-layer KV cache base pointers are stable after allocation, so - # build the pointer array once and return it on every call. - if not hasattr(self, "_context_cache_ptrs"): - ref_cache = cache_layers[0].kv_cache - cache_ptrs = torch.tensor( - [cl.kv_cache.data_ptr() for cl in cache_layers], - dtype=torch.int64, - device=ref_cache.device, - ) - self._context_cache_ptrs = cache_ptrs - return self._context_cache_ptrs - - def forward( - self, - input_ids: torch.Tensor, - positions: torch.Tensor, - inputs_embeds: torch.Tensor | None = None, - ) -> torch.Tensor: - if inputs_embeds is None: - inputs_embeds = self.embed_input_ids(input_ids) - - hidden_states = inputs_embeds - residual = None - for layer in self.layers: - hidden_states, residual = layer( - positions=positions, - hidden_states=hidden_states, - residual=residual, - ) - hidden_states, _ = fused_allreduce_rms_norm( - hidden_states, residual, self.final_norm - ) - return hidden_states - - -class K3DSparkForCausalLM(nn.Module): - has_own_embed_tokens = False - has_own_lm_head = False - draft_id_to_target_id = None - hf_to_vllm_mapper = WeightsMapper( - # confidence_head is training-only. The frozen target embedding and LM - # head are shared after this draft-specific checkpoint is loaded. - orig_to_new_substr={ - "confidence_head": None, - "embed_tokens": None, - "lm_head": None, - }, - orig_to_new_prefix={"": "model."}, - orig_to_new_stacked={ - ".gate_proj": (".gate_up_proj", 0), - ".up_proj": (".gate_up_proj", 1), - ".q_a_proj": (".fused_qkv_a_proj", 0), - ".kv_a_proj_with_mqa": (".fused_qkv_a_proj", 1), - }, - ) - - def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: - super().__init__() - assert vllm_config.speculative_config is not None - self.draft_model_config = vllm_config.speculative_config.draft_model_config - self.config = self.draft_model_config.hf_config - if getattr(self.config, "enable_confidence_head", False): - # The class mapper drops the head as training-only; keep it for drafts - # that actually built one, since adaptive verification reads it. - self.hf_to_vllm_mapper = self.hf_to_vllm_mapper | WeightsMapper( - orig_to_new_substr={"confidence_head": "confidence_head"} - ) - target_layer_num = vllm_config.model_config.get_total_num_hidden_layers() - self.model = K3DSparkModel( - vllm_config=vllm_config, - start_layer_id=target_layer_num, - prefix=maybe_prefix(prefix, "model"), - ) - - # Assigned by load_dspark_model from the target. Keeping no placeholder - # avoids a transient full-vocabulary allocation for this 163k-vocab model. - self.lm_head: nn.Module | None = None - logit_scale = getattr(self.config, "logit_scale", 1.0) - self.logits_processor = LogitsProcessor( - self.config.draft_vocab_size, scale=logit_scale - ) - - def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: - return self.model.embed_input_ids(input_ids) - - def combine_hidden_states(self, hidden_states: torch.Tensor) -> torch.Tensor: - return self.model.combine_hidden_states(hidden_states) - - def get_draft_kv_cache_layer_names(self) -> list[str]: - return [layer.self_attn.layer_name for layer in self.model.layers] - - def precompute_and_store_context_kv( - self, - context_states: torch.Tensor, - context_positions: torch.Tensor, - context_slot_mapping: torch.Tensor | list[torch.Tensor | None] | None = None, - ) -> None: - self.model.precompute_and_store_context_kv( - context_states, context_positions, context_slot_mapping - ) - - def forward( - self, - input_ids: torch.Tensor, - positions: torch.Tensor, - inputs_embeds: torch.Tensor | None = None, - ) -> torch.Tensor: - return self.model(input_ids, positions, inputs_embeds) - - def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor: - assert self.lm_head is not None - return self.logits_processor(self.lm_head, hidden_states) - - def compute_draft_logits(self, hidden_states: torch.Tensor) -> torch.Tensor: - return self.compute_logits(hidden_states) - - def map_draft_to_target(self, draft_ids: torch.Tensor) -> torch.Tensor: - return draft_ids - - def markov_embed(self, token_ids: torch.Tensor) -> torch.Tensor: - return self.model.markov_head.embed(token_ids) - - def markov_bias(self, markov_embed: torch.Tensor) -> torch.Tensor: - return self.model.markov_head.bias(markov_embed, self.logits_processor) - - def compute_confidence( - self, head_hidden: torch.Tensor, markov_embed: torch.Tensor - ) -> torch.Tensor: - """Per-position acceptance probability for each drafted token.""" - assert self.model.confidence_head is not None - return torch.sigmoid(self.model.confidence_head(head_hidden, markov_embed)) - - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loader = AutoWeightsLoader(self) - # read: 1. all weights. 2. context kv weights - weights = _duplicate_context_kv_weights(weights, len(self.model.layers)) - return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) - def process_weights_after_loading(self) -> None: - self.model._build_fused_context_kv_metadata() +__all__ = [ + "K3DSparkDecoderLayer", + "K3DSparkForCausalLM", + "K3DSparkModel", + "_duplicate_context_kv_weights", +] From fa1dd51180101559a26142abbac71b7ddd8a4e0a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Stig-Arne=20Gr=C3=B6nroos?= Date: Fri, 2 Oct 2026 10:43:30 +0300 Subject: [PATCH 2/7] [Model] Parameterize Kimi-K3 DSpark attention and KV-cache owner MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit NVIDIA keeps MultiHeadLatentAttention as the cache owner. A platform override can point context-KV insert and draft layer names at another module. Signed-off-by: Stig-Arne Grönroos Co-authored-by: Cursor --- tests/models/test_dspark_mla.py | 26 +++++-- vllm/models/kimi_k3/common/dspark_mla.py | 94 ++++++++++++++++-------- vllm/models/kimi_k3/nvidia/dspark_mla.py | 64 ++++++++++++++-- 3 files changed, 140 insertions(+), 44 deletions(-) diff --git a/tests/models/test_dspark_mla.py b/tests/models/test_dspark_mla.py index 53a93b940cd5..085a28b9424c 100644 --- a/tests/models/test_dspark_mla.py +++ b/tests/models/test_dspark_mla.py @@ -12,10 +12,22 @@ 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_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_dspark_mla_uses_compile_free_model_entrypoint(): assert ModelRegistry._try_load_model_cls("K3DSparkModel") is K3DSparkForCausalLM assert not issubclass(K3DSparkModel, TorchCompileWithNoGuardsWrapper) @@ -121,12 +133,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, diff --git a/vllm/models/kimi_k3/common/dspark_mla.py b/vllm/models/kimi_k3/common/dspark_mla.py index 37207d7bc9cf..7028810e4c89 100644 --- a/vllm/models/kimi_k3/common/dspark_mla.py +++ b/vllm/models/kimi_k3/common/dspark_mla.py @@ -1,8 +1,14 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""K3 dense MLA draft model for DSpark speculative decoding.""" +"""Platform-neutral Kimi-K3 DSpark MLA draft. + +Attention construction and the MLP class are supplied by the platform module. +This module must not import ``vllm.models.kimi_k3.amd`` or +``vllm.models.kimi_k3.nvidia``. +""" from collections.abc import Iterable +from typing import Any, Protocol import torch import torch.nn as nn @@ -26,8 +32,6 @@ maybe_prefix, ) from vllm.models.common.ops.fused_allreduce_rms_norm import fused_allreduce_rms_norm -from vllm.models.kimi_k3.nvidia.mla import MultiHeadLatentAttention -from vllm.models.kimi_k3.nvidia.model import KimiMLP from vllm.utils.torch_utils import is_quantized_kv_cache from vllm.v1.worker.workspace import current_workspace_manager @@ -36,6 +40,20 @@ ) +class _KVCacheLayer(Protocol): + """Attention module that owns the draft KV cache. + + NVIDIA ``MultiHeadLatentAttention`` is this object. On AMD the wrapper is + not; ``kv_cache_layer`` returns the inner ``MLAAttention``. + """ + + kv_cache: torch.Tensor + kv_cache_dtype: str + _k_scale: torch.Tensor + impl: Any + layer_name: str + + def _duplicate_context_kv_weights( weights: Iterable[tuple[str, torch.Tensor]], num_layers: int ) -> Iterable[tuple[str, torch.Tensor]]: @@ -70,32 +88,20 @@ def __init__( ) -> None: super().__init__() quant_config = get_draft_quant_config(vllm_config) - self.self_attn = MultiHeadLatentAttention( + self.self_attn = self.build_self_attn( + vllm_config=vllm_config, config=config, - hidden_size=config.hidden_size, - num_heads=config.num_attention_heads, - qk_nope_head_dim=config.qk_nope_head_dim, - qk_rope_head_dim=config.qk_rope_head_dim, - v_head_dim=config.v_head_dim, - q_lora_rank=config.q_lora_rank, - kv_lora_rank=config.kv_lora_rank, - cache_config=vllm_config.cache_config, quant_config=quant_config, prefix=maybe_prefix( prefix, f"layers.{start_layer_id + layer_idx}.self_attn" ), - use_rope=True, - non_causal_multi_token_decode=True, ) # Both row-parallel outputs stay un-reduced; their all-reduces are fused # into the RMSNorm that follows via fused_allreduce_rms_norm. self.self_attn.o_proj.reduce_results = False - self.mlp = KimiMLP( - hidden_size=config.hidden_size, - intermediate_size=config.intermediate_size, - hidden_act=config.hidden_act, + self.mlp = self.build_mlp( + config=config, quant_config=quant_config, - reduce_results=False, prefix=maybe_prefix(prefix, f"layers.{start_layer_id + layer_idx}.mlp"), ) self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) @@ -103,6 +109,19 @@ def __init__( config.hidden_size, eps=config.rms_norm_eps ) + def build_self_attn( + self, + *, + vllm_config: VllmConfig, + config, + quant_config, + prefix: str, + ) -> nn.Module: + raise NotImplementedError + + def build_mlp(self, *, config, quant_config, prefix: str) -> nn.Module: + raise NotImplementedError + def forward( self, positions: torch.Tensor, @@ -132,6 +151,17 @@ def forward( class K3DSparkModel(nn.Module): + decoder_layer_cls: type[K3DSparkDecoderLayer] + + def kv_cache_layer(self, attn: nn.Module) -> nn.Module: + """Return the module that owns this layer's KV cache. + + The default is ``attn`` itself, which is NVIDIA + ``MultiHeadLatentAttention``. AMD overrides this to return the inner + ``MLAAttention``. All cache reads go through this method. + """ + return attn + def __init__( self, *, @@ -161,7 +191,7 @@ def __init__( self.layers = nn.ModuleList( [ - K3DSparkDecoderLayer( + self.decoder_layer_cls( vllm_config=vllm_config, config=self.config, layer_idx=layer_idx, @@ -253,11 +283,12 @@ def _build_fused_context_kv_metadata(self) -> None: self._context_rope_dim = attn0.qk_rope_head_dim self._context_rms_norm_eps = attn0.kv_a_layernorm.variance_epsilon self._context_kv_scales: torch.Tensor | None = None - if attn0.kv_cache_dtype in _GROUPED_KV_CACHE_DTYPES and is_quantized_kv_cache( - attn0.kv_cache_dtype + cache0 = self.kv_cache_layer(attn0) + if cache0.kv_cache_dtype in _GROUPED_KV_CACHE_DTYPES and is_quantized_kv_cache( + cache0.kv_cache_dtype ): self._context_kv_scales = torch.stack( - [attn._k_scale.reshape(()) for attn in attentions] + [self.kv_cache_layer(attn)._k_scale.reshape(()) for attn in attentions] ) def _precompute_fused_context_kv( @@ -311,7 +342,7 @@ def _precompute_fused_context_kv( if context_slot_mapping is None: return - cache_layers = [layer.self_attn for layer in self.layers] + cache_layers = [self.kv_cache_layer(layer.self_attn) for layer in self.layers] cache_dtype = cache_layers[0].kv_cache_dtype if ( cache_dtype in _GROUPED_KV_CACHE_DTYPES @@ -354,7 +385,7 @@ def _precompute_fused_context_kv( ) return - for layer_idx, layer in enumerate(self.layers): + for layer_idx, attn in enumerate(cache_layers): slot_mapping = ( context_slot_mapping[layer_idx] if isinstance(context_slot_mapping, (list, tuple)) @@ -362,7 +393,6 @@ def _precompute_fused_context_kv( ) if slot_mapping is None: continue - attn = layer.self_attn attn.impl.do_kv_cache_update( all_kv_c_normed[layer_idx], all_k_pe[layer_idx], @@ -374,7 +404,7 @@ def _precompute_fused_context_kv( def _has_uniform_block_layout( self, - cache_layers: list[MultiHeadLatentAttention], + cache_layers: list[_KVCacheLayer], ) -> bool: key = tuple(cl.kv_cache.data_ptr() for cl in cache_layers) if getattr(self, "_kv_cache_ptrs_key", None) != key: @@ -393,7 +423,7 @@ def _has_uniform_block_layout( def _get_context_kv_cache_ptrs( self, - cache_layers: list[MultiHeadLatentAttention], + cache_layers: list[_KVCacheLayer], ) -> torch.Tensor: # The per-layer KV cache base pointers are stable after allocation, so # build the pointer array once and return it on every call. @@ -431,6 +461,7 @@ def forward( class K3DSparkForCausalLM(nn.Module): + model_cls: type[K3DSparkModel] has_own_embed_tokens = False has_own_lm_head = False draft_id_to_target_id = None @@ -463,7 +494,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: orig_to_new_substr={"confidence_head": "confidence_head"} ) target_layer_num = vllm_config.model_config.get_total_num_hidden_layers() - self.model = K3DSparkModel( + self.model = self.model_cls( vllm_config=vllm_config, start_layer_id=target_layer_num, prefix=maybe_prefix(prefix, "model"), @@ -484,7 +515,10 @@ def combine_hidden_states(self, hidden_states: torch.Tensor) -> torch.Tensor: return self.model.combine_hidden_states(hidden_states) def get_draft_kv_cache_layer_names(self) -> list[str]: - return [layer.self_attn.layer_name for layer in self.model.layers] + return [ + self.model.kv_cache_layer(layer.self_attn).layer_name + for layer in self.model.layers + ] def precompute_and_store_context_kv( self, diff --git a/vllm/models/kimi_k3/nvidia/dspark_mla.py b/vllm/models/kimi_k3/nvidia/dspark_mla.py index c360a23b0657..5e7f0f3a563d 100644 --- a/vllm/models/kimi_k3/nvidia/dspark_mla.py +++ b/vllm/models/kimi_k3/nvidia/dspark_mla.py @@ -1,16 +1,20 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""NVIDIA entry point for the Kimi-K3 DSpark MLA draft. - -The implementation lives in ``vllm.models.kimi_k3.common.dspark_mla``. -""" +"""NVIDIA binding for the Kimi-K3 DSpark MLA draft.""" +from vllm.config import VllmConfig +from vllm.models.kimi_k3.common.dspark_mla import ( + K3DSparkDecoderLayer as _K3DSparkDecoderLayer, +) from vllm.models.kimi_k3.common.dspark_mla import ( - K3DSparkDecoderLayer, - K3DSparkForCausalLM, - K3DSparkModel, - _duplicate_context_kv_weights, + K3DSparkForCausalLM as _K3DSparkForCausalLM, ) +from vllm.models.kimi_k3.common.dspark_mla import ( + K3DSparkModel as _K3DSparkModel, +) +from vllm.models.kimi_k3.common.dspark_mla import _duplicate_context_kv_weights +from vllm.models.kimi_k3.nvidia.mla import MultiHeadLatentAttention +from vllm.models.kimi_k3.nvidia.model import KimiMLP __all__ = [ "K3DSparkDecoderLayer", @@ -18,3 +22,47 @@ "K3DSparkModel", "_duplicate_context_kv_weights", ] + + +class K3DSparkDecoderLayer(_K3DSparkDecoderLayer): + def build_self_attn( + self, + *, + vllm_config: VllmConfig, + config, + quant_config, + prefix: str, + ): + return MultiHeadLatentAttention( + config=config, + hidden_size=config.hidden_size, + num_heads=config.num_attention_heads, + qk_nope_head_dim=config.qk_nope_head_dim, + qk_rope_head_dim=config.qk_rope_head_dim, + v_head_dim=config.v_head_dim, + q_lora_rank=config.q_lora_rank, + kv_lora_rank=config.kv_lora_rank, + cache_config=vllm_config.cache_config, + quant_config=quant_config, + prefix=prefix, + use_rope=True, + non_causal_multi_token_decode=True, + ) + + def build_mlp(self, *, config, quant_config, prefix: str): + return KimiMLP( + hidden_size=config.hidden_size, + intermediate_size=config.intermediate_size, + hidden_act=config.hidden_act, + quant_config=quant_config, + reduce_results=False, + prefix=prefix, + ) + + +class K3DSparkModel(_K3DSparkModel): + decoder_layer_cls = K3DSparkDecoderLayer + + +class K3DSparkForCausalLM(_K3DSparkForCausalLM): + model_cls = K3DSparkModel From 8b96e61940e3b6bba1262fd0e2c871df9be84051 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Stig-Arne=20Gr=C3=B6nroos?= Date: Fri, 2 Oct 2026 10:56:47 +0300 Subject: [PATCH 3/7] [Attention] Add defaulted MLA hooks for a post-absorb decode epilogue MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The defaults keep today's concat and cache write. A subclass can replace the decode query and the KV rows it writes without a new custom-op argument. Signed-off-by: Stig-Arne Grönroos Co-authored-by: Cursor --- tests/models/test_dspark_mla.py | 44 +++++++++ .../layers/attention/mla_attention.py | 91 +++++++++++++++++-- vllm/model_executor/layers/mla.py | 5 +- 3 files changed, 129 insertions(+), 11 deletions(-) diff --git a/tests/models/test_dspark_mla.py b/tests/models/test_dspark_mla.py index 085a28b9424c..8533f994e976 100644 --- a/tests/models/test_dspark_mla.py +++ b/tests/models/test_dspark_mla.py @@ -28,6 +28,50 @@ def kv_cache_layer(self, 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) diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index 2a3048096c98..01d723fd911e 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/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,9 @@ 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. The default path does not read it. + self.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 +799,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 +811,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.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 +880,11 @@ 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. + self._rope_positions = positions if self.use_direct_call: forward_context: ForwardContext = get_forward_context() attn_metadata_raw = forward_context.attn_metadata @@ -999,10 +1066,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 +1157,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 diff --git a/vllm/model_executor/layers/mla.py b/vllm/model_executor/layers/mla.py index 9dd5ad67736c..f35907fe83f4 100644 --- a/vllm/model_executor/layers/mla.py +++ b/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, From 542a75a30f669482aeb989f5e2148baecf935cb8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Stig-Arne=20Gr=C3=B6nroos?= Date: Fri, 2 Oct 2026 13:46:51 +0300 Subject: [PATCH 4/7] [Model] Satisfy ruff UP038 on the moved DSpark slot-mapping checks MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The common module is a new file, so pre-commit lints it in full. The (list, tuple) form was already on main and fails UP038 once the file moves. Signed-off-by: Stig-Arne Grönroos Co-authored-by: Cursor --- vllm/models/kimi_k3/common/dspark_mla.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/vllm/models/kimi_k3/common/dspark_mla.py b/vllm/models/kimi_k3/common/dspark_mla.py index 7028810e4c89..9645af1d9a91 100644 --- a/vllm/models/kimi_k3/common/dspark_mla.py +++ b/vllm/models/kimi_k3/common/dspark_mla.py @@ -356,7 +356,7 @@ def _precompute_fused_context_kv( # Grouped context KV insert assumes that all layers share the same # cache dtype and block layout. - if isinstance(context_slot_mapping, (list, tuple)): + if isinstance(context_slot_mapping, list | tuple): per_layer_slot_mappings = [ s for s in context_slot_mapping if s is not None ] @@ -388,7 +388,7 @@ def _precompute_fused_context_kv( for layer_idx, attn in enumerate(cache_layers): slot_mapping = ( context_slot_mapping[layer_idx] - if isinstance(context_slot_mapping, (list, tuple)) + if isinstance(context_slot_mapping, list | tuple) else context_slot_mapping ) if slot_mapping is None: From 8b74ac7a40da16a3b4a905b87b52aad113e08561 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Stig-Arne=20Gr=C3=B6nroos?= Date: Fri, 2 Oct 2026 13:49:37 +0300 Subject: [PATCH 5/7] [CI] Run the Kimi-K3 DSpark MLA tests in Basic Models MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The file was not listed in Buildkite, so the refactor would not be collected on either platform. Signed-off-by: Stig-Arne Grönroos Co-authored-by: Cursor --- .buildkite/test-amd.yaml | 3 ++- .buildkite/test_areas/models_basic.yaml | 6 ++++-- 2 files changed, 6 insertions(+), 3 deletions(-) diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index a3f7dbd83874..229cd50eb304 100644 --- a/.buildkite/test-amd.yaml +++ b/.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 diff --git a/.buildkite/test_areas/models_basic.yaml b/.buildkite/test_areas/models_basic.yaml index a25c9a78a86c..b9235874447c 100644 --- a/.buildkite/test_areas/models_basic.yaml +++ b/.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" From 4a789d2308e4f28da7e560abfea992898c0b337f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Stig-Arne=20Gr=C3=B6nroos?= Date: Fri, 2 Oct 2026 14:13:43 +0300 Subject: [PATCH 6/7] [Attention] Store the MLA epilogue RoPE off rotary_emb MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit DeepseekV32Attention already assigns a concrete RotaryEmbedding to that name. Typing the base attribute as optional made mypy reject its cache reads. Signed-off-by: Stig-Arne Grönroos Co-authored-by: Cursor --- vllm/model_executor/layers/attention/mla_attention.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index 01d723fd911e..ae06ba5c926e 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -482,8 +482,9 @@ 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. The default path does not read it. - self.rotary_emb = rotary_emb + # 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 @@ -854,7 +855,7 @@ def _form_decode_q( 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.rotary_emb`` are available there. + ``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) From e1b19057d53ba6c4672f304c19080575594a0d6d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Stig-Arne=20Gr=C3=B6nroos?= Date: Thu, 8 Oct 2026 15:51:35 +0300 Subject: [PATCH 7/7] [Attention] Document MLA CUDA-graph buffer lifetime MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replay reads the captured positions and slot-mapping allocations, so a replacement tensor freezes the pointer. Lock the NVIDIA binding class and the default hook returns. Signed-off-by: Stig-Arne Grönroos Co-authored-by: Cursor --- tests/models/test_dspark_mla.py | 36 +++++++++++++++++++ .../layers/attention/mla_attention.py | 7 ++++ 2 files changed, 43 insertions(+) diff --git a/tests/models/test_dspark_mla.py b/tests/models/test_dspark_mla.py index 8533f994e976..89503279e3fc 100644 --- a/tests/models/test_dspark_mla.py +++ b/tests/models/test_dspark_mla.py @@ -17,6 +17,42 @@ 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 diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index ae06ba5c926e..cdc25ece7ff0 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -885,6 +885,10 @@ def forward( ) -> 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() @@ -899,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), (