A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 5, 16:00 UTC
[Perf] DSV3.2 Indexer Fused Weights Projection (#38684)
Signed-off-by: Benjamin Chislett <bchislett@nvidia.com>
Benjamin Chislett committed
Apr 1, 2026 at 23:34 UTC
5f96f9aff10fec0c31eb61f7587b3789400f9c17
2 files changed
+25
-14
vllm/model_executor/models/deepseek_mtp.py
+3
@@ -241,6 +241,9 @@ class DeepSeekMTP(nn.Module, DeepseekV2MixtureOfExperts):
241
("gate_up_proj", "up_proj", 1),
242
("fused_qkv_a_proj", "q_a_proj", 0),
243
("fused_qkv_a_proj", "kv_a_proj_with_mqa", 1),
244
+ # Fused indexer wk + weights_proj
245
+ ("wk_weights_proj", "wk", 0),
246
+ ("wk_weights_proj", "weights_proj", 1),
247
]
248
249
expert_params_mapping = SharedFusedMoE.make_expert_params_mapping(
vllm/model_executor/models/deepseek_v2.py
+22
-14
@@ -639,21 +639,19 @@ class Indexer(nn.Module):
639
quant_config=quant_config,
640
prefix=f"{prefix}.wq_b",
641
)
642
- self.wk = ReplicatedLinear(
642
+ # Fused wk + weights_proj: single GEMM producing [head_dim + n_head].
643
+ # weights_proj does not get quantized, so we run both with quant_config=None
644
+ # wk may be upcasted from the default quant; experiments show fusion is always
645
+ # faster unless WK proj is in FP4, which is not the case for all known quants.
646
+ self.wk_weights_proj = MergedColumnParallelLinear(
647
hidden_size,
644
- self.head_dim,
645
- bias=False,
646
- quant_config=quant_config,
647
- prefix=f"{prefix}.wk",
648
- )
649
- self.k_norm = LayerNorm(self.head_dim, eps=1e-6)
650
- self.weights_proj = ReplicatedLinear(
651
- hidden_size,
652
- self.n_head,
648
+ [self.head_dim, self.n_head],
649
bias=False,
650
quant_config=None,
655
- prefix=f"{prefix}.weights_proj",
651
+ disable_tp=True,
652
+ prefix=f"{prefix}.wk_weights_proj",
653
)
654
+ self.k_norm = LayerNorm(self.head_dim, eps=1e-6)
655
self.softmax_scale = self.head_dim**-0.5
656
657
self.scale_fmt = "ue8m0"
@@ -694,7 +692,11 @@ class Indexer(nn.Module):
692
q, [self.rope_dim, self.head_dim - self.rope_dim], dim=-1
693
)
694
697
- k, _ = self.wk(hidden_states)
695
+ # Fused wk + weights_proj: one GEMM, then split
696
+ kw, _ = self.wk_weights_proj(hidden_states)
697
+ k = kw[:, : self.head_dim]
698
+ weights_raw = kw[:, self.head_dim :]
699
+
700
k = self.k_norm(k)
701
k_pe, k_nope = torch.split(
702
k, [self.rope_dim, self.head_dim - self.rope_dim], dim=-1
@@ -723,9 +725,8 @@ class Indexer(nn.Module):
725
q_fp8 = q_fp8.view(-1, self.n_head, self.head_dim)
726
q_scale = q_scale.view(-1, self.n_head, 1)
727
726
- weights, _ = self.weights_proj(hidden_states)
728
weights = (
728
- weights.unsqueeze(-1) * q_scale * self.softmax_scale * self.n_head**-0.5
729
+ weights_raw.unsqueeze(-1) * q_scale * self.softmax_scale * self.n_head**-0.5
730
)
731
weights = weights.squeeze(-1)
732
@@ -1438,6 +1439,13 @@ class DeepseekV2ForCausalLM(
1439
("qkv_proj", "k_proj", "k"),
1440
("qkv_proj", "v_proj", "v"),
1441
]
1442
+ # Fused indexer wk + weights_proj (shard 0 = wk, shard 1 = weights_proj)
1443
+ indexer_fused_mapping = [
1444
+ ("wk_weights_proj", "wk", 0),
1445
+ ("wk_weights_proj", "weights_proj", 1),
1446
+ ]
1447
+ stacked_params_mapping.extend(indexer_fused_mapping)
1448
+
1449
if self.use_mha:
1450
stacked_params_mapping.extend(mha_params_mapping)
1451
else: