[K3 Perf] Optimize k3 dspark fused kv, 4.5~4.6x kernel performance improvement (#50585)
Signed-off-by: yewentao256 <zhyanwentao@126.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Wentao Ye committed
Aug 7, 2026 at 15:16 UTC
56a4b63d44c71784de75e6f2d6cbe28adf74635f
2 files changed
+104
-88
tests/models/test_dspark_mla.py
+48
-1
@@ -107,6 +107,7 @@ def test_dspark_markov_head_is_replicated(
107
@pytest.mark.cpu_test
108
def test_k3_dspark_uses_replicated_markov_head(monkeypatch: pytest.MonkeyPatch):
109
markov_head_calls = []
110
+ context_kv_proj_calls = []
111
112
class DummyModule(nn.Module):
113
def __init__(self, *args, **kwargs):
@@ -116,8 +117,13 @@ def test_k3_dspark_uses_replicated_markov_head(monkeypatch: pytest.MonkeyPatch):
117
markov_head_calls.append((args, kwargs))
118
return DummyModule()
119
120
+ def make_context_kv_proj(*args, **kwargs):
121
+ context_kv_proj_calls.append((args, kwargs))
122
+ return DummyModule()
123
+
124
monkeypatch.setattr(dspark_mla, "get_draft_quant_config", lambda _: None)
125
monkeypatch.setattr(dspark_mla, "ReplicatedLinear", DummyModule)
126
+ monkeypatch.setattr(dspark_mla, "MergedColumnParallelLinear", make_context_kv_proj)
127
monkeypatch.setattr(dspark_mla, "RMSNorm", DummyModule)
128
monkeypatch.setattr(dspark_mla, "K3DSparkDecoderLayer", DummyModule)
129
monkeypatch.setattr(dspark_mla, "DSparkMarkovHead", make_markov_head)
@@ -126,6 +132,8 @@ def test_k3_dspark_uses_replicated_markov_head(monkeypatch: pytest.MonkeyPatch):
132
target_hidden_size=16,
133
num_target_layers=2,
134
hidden_size=8,
135
+ kv_lora_rank=3,
136
+ qk_rope_head_dim=1,
137
rms_norm_eps=1e-6,
138
num_hidden_layers=1,
139
vocab_size=128,
@@ -135,9 +143,48 @@ def test_k3_dspark_uses_replicated_markov_head(monkeypatch: pytest.MonkeyPatch):
143
vllm_config = SimpleNamespace(
144
speculative_config=SimpleNamespace(
145
draft_model_config=SimpleNamespace(hf_config=config)
138
- )
146
+ ),
147
+ scheduler_config=SimpleNamespace(max_num_batched_tokens=16),
148
)
149
150
K3DSparkModel(vllm_config=vllm_config, start_layer_id=0, prefix="model")
151
152
assert len(markov_head_calls) == 1
153
+ assert context_kv_proj_calls == [
154
+ (
155
+ (8, [4]),
156
+ {
157
+ "bias": False,
158
+ "return_bias": False,
159
+ "quant_config": None,
160
+ "prefix": "model.layers.0.self_attn.fused_qkv_a_proj",
161
+ "disable_tp": True,
162
+ },
163
+ )
164
+ ]
165
+
166
+
167
+def test_context_kv_weights_are_loaded_as_merged_linear_shards():
168
+ weights = [
169
+ (
170
+ "layers.0.self_attn.kv_a_proj_with_mqa.weight_packed",
171
+ torch.arange(4),
172
+ ),
173
+ (
174
+ "layers.1.self_attn.kv_a_proj_with_mqa.weight_scale",
175
+ torch.tensor(0.5),
176
+ ),
177
+ ]
178
+
179
+ duplicated = dspark_mla._duplicate_context_kv_weights(weights, 2)
180
+ mapped = list(K3DSparkForCausalLM.hf_to_vllm_mapper.apply(duplicated))
181
+
182
+ assert [name for name, _ in mapped] == [
183
+ "model.layers.0.self_attn.fused_qkv_a_proj.weight_packed",
184
+ "model.context_kv_proj.weight_packed",
185
+ "model.layers.1.self_attn.fused_qkv_a_proj.weight_scale",
186
+ "model.context_kv_proj.weight_scale",
187
+ ]
188
+ assert [weight.shard_id for _, weight in mapped] == [1, 0, 1, 1]
189
+ assert mapped[0][1].data_ptr() == mapped[1][1].data_ptr()
190
+ assert mapped[2][1].data_ptr() == mapped[3][1].data_ptr()
vllm/models/kimi_k3/nvidia/dspark_mla.py
+56
-87
@@ -6,12 +6,14 @@ from collections.abc import Iterable
6
7
import torch
8
import torch.nn as nn
9
-import torch.nn.functional as F
9
10
import vllm._custom_ops as ops
11
from vllm.config import VllmConfig
12
from vllm.model_executor.layers.layernorm import RMSNorm
14
-from vllm.model_executor.layers.linear import ReplicatedLinear
13
+from vllm.model_executor.layers.linear import (
14
+ MergedColumnParallelLinear,
15
+ ReplicatedLinear,
16
+)
17
from vllm.model_executor.layers.logits_processor import LogitsProcessor
18
from vllm.model_executor.models.qwen3_dspark import DSparkMarkovHead
19
from vllm.model_executor.models.utils import (
@@ -24,6 +26,29 @@ from vllm.models.common.ops.fused_allreduce_rms_norm import fused_allreduce_rms_
26
from vllm.models.kimi_k3.nvidia.mla import MultiHeadLatentAttention
27
from vllm.models.kimi_k3.nvidia.model import KimiMLP
28
from vllm.utils.torch_utils import is_quantized_kv_cache
29
+from vllm.v1.worker.workspace import current_workspace_manager
30
+
31
+
32
+def _duplicate_context_kv_weights(
33
+ weights: Iterable[tuple[str, torch.Tensor]], num_layers: int
34
+) -> Iterable[tuple[str, torch.Tensor]]:
35
+ """Load each layer's KV projection into the cross-layer linear."""
36
+ for name, weight in weights:
37
+ yield name, weight
38
+ layer_prefix, marker, param_name = name.partition(
39
+ ".self_attn.kv_a_proj_with_mqa."
40
+ )
41
+ if not marker:
42
+ continue
43
+ layer_idx_str = layer_prefix.rsplit(".", 1)[-1]
44
+ if not layer_idx_str.isdecimal():
45
+ continue
46
+ layer_idx = int(layer_idx_str)
47
+ if layer_idx >= num_layers:
48
+ continue
49
+ fused_weight = weight.detach()
50
+ fused_weight.shard_id = layer_idx
51
+ yield f"context_kv_proj.{param_name}", fused_weight
52
53
54
class K3DSparkDecoderLayer(nn.Module):
@@ -139,6 +164,19 @@ class K3DSparkModel(nn.Module):
164
for layer_idx in range(self.config.num_hidden_layers)
165
]
166
)
167
+ kv_width = self.config.kv_lora_rank + self.config.qk_rope_head_dim
168
+ self.context_kv_proj = MergedColumnParallelLinear(
169
+ self.config.hidden_size,
170
+ [kv_width] * self.config.num_hidden_layers,
171
+ bias=False,
172
+ return_bias=False,
173
+ quant_config=self.quant_config,
174
+ prefix=maybe_prefix(
175
+ prefix,
176
+ f"layers.{start_layer_id}.self_attn.fused_qkv_a_proj",
177
+ ),
178
+ disable_tp=True,
179
+ )
180
self.final_norm = RMSNorm(self.config.hidden_size, eps=self.config.rms_norm_eps)
181
self.markov_head = DSparkMarkovHead(
182
self.config.vocab_size,
@@ -146,7 +184,6 @@ class K3DSparkModel(nn.Module):
184
self.config.markov_rank,
185
prefix=maybe_prefix(prefix, "markov_head"),
186
)
149
- self._context_kv_fusion_available: bool | None = None
187
self._max_num_context_tokens = (
188
vllm_config.scheduler_config.max_num_batched_tokens
189
)
@@ -166,76 +203,19 @@ class K3DSparkModel(nn.Module):
203
context_slot_mapping: torch.Tensor | list[torch.Tensor | None] | None = None,
204
) -> None:
205
"""Project target-derived context into each draft layer's latent cache."""
169
- if self._context_kv_fusion_available is None:
170
- self._build_fused_context_kv_buffers()
171
- if self._context_kv_fusion_available:
172
- self._precompute_fused_context_kv(
173
- context_states, context_positions, context_slot_mapping
174
- )
175
- return
176
-
177
- # Quantized fallback. Directly invoking the projection modules preserves
178
- # their quantization methods, at the cost of also computing unused Q rows.
179
- for layer_idx, layer in enumerate(self.layers):
180
- attn = layer.self_attn
181
- assert attn.fused_qkv_a_proj is not None
182
- assert attn.q_lora_rank is not None
183
- assert attn.rotary_emb is not None
184
- qkv_lora = attn.fused_qkv_a_proj(context_states)[0]
185
- kv_lora = qkv_lora[..., attn.q_lora_rank :]
186
- kv_c, k_pe = kv_lora.split(
187
- [attn.kv_lora_rank, attn.qk_rope_head_dim], dim=-1
188
- )
189
- kv_c = attn.kv_a_layernorm(kv_c)
190
- k_pe = k_pe.unsqueeze(1)
191
- # DeepSeek YaRN's FlashInfer path requires paired Q/K tensors.
192
- # The vLLM CUDA op supports rotating one tensor in place and
193
- # consumes the same (possibly scaled fp32) cos/sin cache.
194
- rotary_emb = attn.rotary_emb
195
- ops.rotary_embedding(
196
- context_positions,
197
- k_pe,
198
- None,
199
- rotary_emb.head_size,
200
- rotary_emb.cos_sin_cache,
201
- rotary_emb.is_neox_style,
202
- )
203
-
204
- slot_mapping = (
205
- context_slot_mapping[layer_idx]
206
- if isinstance(context_slot_mapping, (list, tuple))
207
- else context_slot_mapping
208
- )
209
- if slot_mapping is None:
210
- continue
211
- attn.impl.do_kv_cache_update(
212
- kv_c,
213
- k_pe,
214
- attn.kv_cache,
215
- slot_mapping,
216
- attn.kv_cache_dtype,
217
- attn._k_scale,
218
- )
219
-
220
- def _build_fused_context_kv_buffers(self) -> None:
221
- """Build a cross-layer KV-only A projection after checkpoint loading."""
222
- if self.quant_config is not None:
223
- self._context_kv_fusion_available = False
224
- return
206
+ if not hasattr(self, "_num_context_layers"):
207
+ self._build_fused_context_kv_metadata()
208
+ self._precompute_fused_context_kv(
209
+ context_states, context_positions, context_slot_mapping
210
+ )
211
212
+ def _build_fused_context_kv_metadata(self) -> None:
213
+ """Build cross-layer metadata after checkpoint loading."""
214
attentions = [layer.self_attn for layer in self.layers]
227
- if not attentions or any(
228
- attn.fused_qkv_a_proj is None
229
- or not hasattr(attn.fused_qkv_a_proj, "weight")
230
- for attn in attentions
231
- ):
232
- self._context_kv_fusion_available = False
233
- return
234
-
215
+ assert attentions
216
attn0 = attentions[0]
217
assert attn0.q_lora_rank is not None
218
kv_width = attn0.kv_lora_rank + attn0.qk_rope_head_dim
238
- kv_weights = []
219
for attn in attentions:
220
assert attn.q_lora_rank is not None
221
assert (
@@ -245,16 +225,6 @@ class K3DSparkModel(nn.Module):
225
and attn.kv_a_layernorm.variance_epsilon
226
== attn0.kv_a_layernorm.variance_epsilon
227
), "All MLA DSpark layers must share their latent KV geometry."
248
- kv_weights.append(
249
- attn.fused_qkv_a_proj.weight.detach().narrow(
250
- 0, attn.q_lora_rank, kv_width
251
- )
252
- )
253
-
254
- # [L * (kv_lora_rank + rope_dim), hidden_size]. The underlying fused
255
- # A weights are replicated (`disable_tp=True`), so this is valid on
256
- # every TP rank without communication.
257
- self._fused_context_kv_weight = torch.cat(kv_weights, dim=0)
228
self._context_kv_norm_weights = torch.stack(
229
[attn.kv_a_layernorm.weight.detach() for attn in attentions], dim=0
230
).contiguous()
@@ -263,12 +233,6 @@ class K3DSparkModel(nn.Module):
233
self._context_kv_lora_rank = attn0.kv_lora_rank
234
self._context_rope_dim = attn0.qk_rope_head_dim
235
self._context_rms_norm_eps = attn0.kv_a_layernorm.variance_epsilon
266
- self._context_positions_repeated = torch.empty(
267
- self._num_context_layers * self._max_num_context_tokens,
268
- dtype=torch.int64,
269
- device=self._fused_context_kv_weight.device,
270
- )
271
- self._context_kv_fusion_available = True
236
237
def _precompute_fused_context_kv(
238
self,
@@ -281,7 +245,7 @@ class K3DSparkModel(nn.Module):
245
246
# One KV-only GEMM replaces five full Q+KV GEMMs. For K3 this projects
247
# 5*576 rows rather than 5*2112 rows (72.7% fewer A-projection FLOPs).
284
- all_kv = F.linear(context_states, self._fused_context_kv_weight)
248
+ all_kv = self.context_kv_proj(context_states)
249
all_kv = all_kv.view(num_ctx, num_layers, self._context_kv_width)
250
all_kv_c = all_kv[..., : self._context_kv_lora_rank]
251
all_k_pe = all_kv[..., self._context_kv_lora_rank :]
@@ -299,7 +263,10 @@ class K3DSparkModel(nn.Module):
263
264
all_k_pe = all_k_pe.permute(1, 0, 2).contiguous()
265
all_k_pe_flat = all_k_pe.view(num_layers * num_ctx, 1, self._context_rope_dim)
302
- repeated_positions = self._context_positions_repeated[: num_layers * num_ctx]
266
+ (repeated_positions,) = current_workspace_manager().get_simultaneous(
267
+ ((num_layers * self._max_num_context_tokens,), torch.int64),
268
+ )
269
+ repeated_positions = repeated_positions[: num_layers * num_ctx]
270
repeated_positions.view(num_layers, num_ctx).copy_(context_positions)
271
# Keep the single-tensor context RoPE on vLLM's optimized CUDA op;
272
# DeepSeek YaRN's FlashInfer wrapper assumes a non-null key tensor.
@@ -516,6 +483,8 @@ class K3DSparkForCausalLM(nn.Module):
483
self,
484
skip_substrs=list(self.checkpoint_skip_substrs),
485
)
486
+ # read: 1. all weights. 2. context kv weights
487
+ weights = _duplicate_context_kv_weights(weights, len(self.model.layers))
488
loaded_weights = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
520
- self.model._build_fused_context_kv_buffers()
489
+ self.model._build_fused_context_kv_metadata()
490
return loaded_weights