diff --git a/tests/models/test_dspark_mla.py b/tests/models/test_dspark_mla.py index 28dcd8ebf4..3188f0fadd 100644 --- a/tests/models/test_dspark_mla.py +++ b/tests/models/test_dspark_mla.py @@ -107,6 +107,7 @@ def test_dspark_markov_head_is_replicated( @pytest.mark.cpu_test def test_k3_dspark_uses_replicated_markov_head(monkeypatch: pytest.MonkeyPatch): markov_head_calls = [] + context_kv_proj_calls = [] class DummyModule(nn.Module): def __init__(self, *args, **kwargs): @@ -116,8 +117,13 @@ def test_k3_dspark_uses_replicated_markov_head(monkeypatch: pytest.MonkeyPatch): markov_head_calls.append((args, kwargs)) return DummyModule() + 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) @@ -126,6 +132,8 @@ def test_k3_dspark_uses_replicated_markov_head(monkeypatch: pytest.MonkeyPatch): target_hidden_size=16, num_target_layers=2, hidden_size=8, + kv_lora_rank=3, + qk_rope_head_dim=1, rms_norm_eps=1e-6, num_hidden_layers=1, vocab_size=128, @@ -135,9 +143,48 @@ def test_k3_dspark_uses_replicated_markov_head(monkeypatch: pytest.MonkeyPatch): vllm_config = SimpleNamespace( speculative_config=SimpleNamespace( draft_model_config=SimpleNamespace(hf_config=config) - ) + ), + scheduler_config=SimpleNamespace(max_num_batched_tokens=16), ) K3DSparkModel(vllm_config=vllm_config, start_layer_id=0, prefix="model") assert len(markov_head_calls) == 1 + assert context_kv_proj_calls == [ + ( + (8, [4]), + { + "bias": False, + "return_bias": False, + "quant_config": None, + "prefix": "model.layers.0.self_attn.fused_qkv_a_proj", + "disable_tp": True, + }, + ) + ] + + +def test_context_kv_weights_are_loaded_as_merged_linear_shards(): + weights = [ + ( + "layers.0.self_attn.kv_a_proj_with_mqa.weight_packed", + torch.arange(4), + ), + ( + "layers.1.self_attn.kv_a_proj_with_mqa.weight_scale", + torch.tensor(0.5), + ), + ] + + duplicated = dspark_mla._duplicate_context_kv_weights(weights, 2) + mapped = list(K3DSparkForCausalLM.hf_to_vllm_mapper.apply(duplicated)) + + assert [name for name, _ in mapped] == [ + "model.layers.0.self_attn.fused_qkv_a_proj.weight_packed", + "model.context_kv_proj.weight_packed", + "model.layers.1.self_attn.fused_qkv_a_proj.weight_scale", + "model.context_kv_proj.weight_scale", + ] + assert [weight.shard_id for _, weight in mapped] == [1, 0, 1, 1] + assert mapped[0][1].data_ptr() == mapped[1][1].data_ptr() + assert mapped[2][1].data_ptr() == mapped[3][1].data_ptr() diff --git a/vllm/models/kimi_k3/nvidia/dspark_mla.py b/vllm/models/kimi_k3/nvidia/dspark_mla.py index 64714496a1..81715e02c9 100644 --- a/vllm/models/kimi_k3/nvidia/dspark_mla.py +++ b/vllm/models/kimi_k3/nvidia/dspark_mla.py @@ -6,12 +6,14 @@ from collections.abc import Iterable import torch import torch.nn as nn -import torch.nn.functional as F 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 ReplicatedLinear +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 DSparkMarkovHead from vllm.model_executor.models.utils import ( @@ -24,6 +26,29 @@ from vllm.models.common.ops.fused_allreduce_rms_norm import fused_allreduce_rms_ 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 + + +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): @@ -139,6 +164,19 @@ class K3DSparkModel(nn.Module): 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, @@ -146,7 +184,6 @@ class K3DSparkModel(nn.Module): self.config.markov_rank, prefix=maybe_prefix(prefix, "markov_head"), ) - self._context_kv_fusion_available: bool | None = None self._max_num_context_tokens = ( vllm_config.scheduler_config.max_num_batched_tokens ) @@ -166,76 +203,19 @@ class K3DSparkModel(nn.Module): context_slot_mapping: torch.Tensor | list[torch.Tensor | None] | None = None, ) -> None: """Project target-derived context into each draft layer's latent cache.""" - if self._context_kv_fusion_available is None: - self._build_fused_context_kv_buffers() - if self._context_kv_fusion_available: - self._precompute_fused_context_kv( - context_states, context_positions, context_slot_mapping - ) - return - - # Quantized fallback. Directly invoking the projection modules preserves - # their quantization methods, at the cost of also computing unused Q rows. - for layer_idx, layer in enumerate(self.layers): - attn = layer.self_attn - assert attn.fused_qkv_a_proj is not None - assert attn.q_lora_rank is not None - assert attn.rotary_emb is not None - qkv_lora = attn.fused_qkv_a_proj(context_states)[0] - kv_lora = qkv_lora[..., attn.q_lora_rank :] - kv_c, k_pe = kv_lora.split( - [attn.kv_lora_rank, attn.qk_rope_head_dim], dim=-1 - ) - kv_c = attn.kv_a_layernorm(kv_c) - k_pe = k_pe.unsqueeze(1) - # DeepSeek YaRN's FlashInfer path requires paired Q/K tensors. - # The vLLM CUDA op supports rotating one tensor in place and - # consumes the same (possibly scaled fp32) cos/sin cache. - rotary_emb = attn.rotary_emb - ops.rotary_embedding( - context_positions, - k_pe, - None, - rotary_emb.head_size, - rotary_emb.cos_sin_cache, - rotary_emb.is_neox_style, - ) - - 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.impl.do_kv_cache_update( - kv_c, - k_pe, - attn.kv_cache, - slot_mapping, - attn.kv_cache_dtype, - attn._k_scale, - ) - - def _build_fused_context_kv_buffers(self) -> None: - """Build a cross-layer KV-only A projection after checkpoint loading.""" - if self.quant_config is not None: - self._context_kv_fusion_available = False - return + 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] - if not attentions or any( - attn.fused_qkv_a_proj is None - or not hasattr(attn.fused_qkv_a_proj, "weight") - for attn in attentions - ): - self._context_kv_fusion_available = False - return - + assert attentions attn0 = attentions[0] assert attn0.q_lora_rank is not None kv_width = attn0.kv_lora_rank + attn0.qk_rope_head_dim - kv_weights = [] for attn in attentions: assert attn.q_lora_rank is not None assert ( @@ -245,16 +225,6 @@ class K3DSparkModel(nn.Module): and attn.kv_a_layernorm.variance_epsilon == attn0.kv_a_layernorm.variance_epsilon ), "All MLA DSpark layers must share their latent KV geometry." - kv_weights.append( - attn.fused_qkv_a_proj.weight.detach().narrow( - 0, attn.q_lora_rank, kv_width - ) - ) - - # [L * (kv_lora_rank + rope_dim), hidden_size]. The underlying fused - # A weights are replicated (`disable_tp=True`), so this is valid on - # every TP rank without communication. - self._fused_context_kv_weight = torch.cat(kv_weights, dim=0) self._context_kv_norm_weights = torch.stack( [attn.kv_a_layernorm.weight.detach() for attn in attentions], dim=0 ).contiguous() @@ -263,12 +233,6 @@ class K3DSparkModel(nn.Module): 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_positions_repeated = torch.empty( - self._num_context_layers * self._max_num_context_tokens, - dtype=torch.int64, - device=self._fused_context_kv_weight.device, - ) - self._context_kv_fusion_available = True def _precompute_fused_context_kv( self, @@ -281,7 +245,7 @@ class K3DSparkModel(nn.Module): # 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 = F.linear(context_states, self._fused_context_kv_weight) + 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 :] @@ -299,7 +263,10 @@ class K3DSparkModel(nn.Module): 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 = self._context_positions_repeated[: num_layers * num_ctx] + (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. @@ -516,6 +483,8 @@ class K3DSparkForCausalLM(nn.Module): self, skip_substrs=list(self.checkpoint_skip_substrs), ) + # read: 1. all weights. 2. context kv weights + weights = _duplicate_context_kv_weights(weights, len(self.model.layers)) loaded_weights = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) - self.model._build_fused_context_kv_buffers() + self.model._build_fused_context_kv_metadata() return loaded_weights