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, 02:15 UTC
[Bugfix][Perf] Indexer upcast WK to BF16 for fusion (#38928)
Signed-off-by: Benjamin Chislett <bchislett@nvidia.com>
Benjamin Chislett committed
Apr 15, 2026 at 16:39 UTC
ac3dac545b28ea6cf847e0044859e58f33d4f8b9
2 files changed
+84
-64
vllm/model_executor/models/deepseek_mtp.py
+14
-11
@@ -30,6 +30,7 @@ from .deepseek_v2 import (
30
DeepseekV2DecoderLayer,
31
DeepseekV2MixtureOfExperts,
32
DeepseekV2MoE,
33
+ _try_load_fp8_indexer_wk,
34
get_spec_layer_idx_from_weight_name,
35
)
36
from .utils import maybe_prefix
@@ -190,10 +191,6 @@ class DeepSeekMTP(nn.Module, DeepseekV2MixtureOfExperts):
191
)
192
# Set MoE hyperparameters
193
self.set_moe_parameters()
193
- self.is_fp4_ckpt = (
194
- self.quant_config is not None
195
- and self.quant_config.get_name() == "modelopt_fp4"
196
- )
194
195
def set_moe_parameters(self):
196
self.expert_weights = []
@@ -248,13 +245,12 @@ class DeepSeekMTP(nn.Module, DeepseekV2MixtureOfExperts):
245
("fused_qkv_a_proj", "kv_a_proj_with_mqa", 1),
246
]
247
251
- if self.is_fp4_ckpt:
252
- # Fused indexer wk + weights_proj (shard 0 = wk, shard 1 = weights_proj)
253
- indexer_fused_mapping = [
254
- ("wk_weights_proj", "wk", 0),
255
- ("wk_weights_proj", "weights_proj", 1),
256
- ]
257
- stacked_params_mapping.extend(indexer_fused_mapping)
248
+ # Fused indexer wk + weights_proj (shard 0 = wk, shard 1 = weights_proj)
249
+ indexer_fused_mapping = [
250
+ ("wk_weights_proj", "wk", 0),
251
+ ("wk_weights_proj", "weights_proj", 1),
252
+ ]
253
+ stacked_params_mapping.extend(indexer_fused_mapping)
254
255
expert_params_mapping = SharedFusedMoE.make_expert_params_mapping(
256
self,
@@ -271,6 +267,7 @@ class DeepSeekMTP(nn.Module, DeepseekV2MixtureOfExperts):
267
268
params_dict = dict(self.named_parameters())
269
loaded_params: set[str] = set()
270
+ _pending_wk_fp8: dict = {} # FP8 indexer wk dequant buffer
271
for name, loaded_weight in weights:
272
if "rotary_emb.inv_freq" in name:
273
continue
@@ -281,6 +278,12 @@ class DeepSeekMTP(nn.Module, DeepseekV2MixtureOfExperts):
278
rocm_aiter_moe_shared_expert_enabled and ("mlp.shared_experts" in name)
279
)
280
name = self._rewrite_spec_layer_name(spec_layer, name)
281
+
282
+ if _try_load_fp8_indexer_wk(
283
+ name, loaded_weight, _pending_wk_fp8, params_dict, loaded_params
284
+ ):
285
+ continue
286
+
287
for param_name, weight_name, shard_id in stacked_params_mapping:
288
# Skip non-stacked layers and experts (experts handled below).
289
if weight_name not in name:
vllm/model_executor/models/deepseek_v2.py
+70
-53
@@ -66,6 +66,10 @@ from vllm.model_executor.layers.quantization import QuantizationConfig
66
from vllm.model_executor.layers.quantization.utils.fp8_utils import (
67
per_token_group_quant_fp8,
68
)
69
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
70
+ GroupShape,
71
+ scaled_dequantize,
72
+)
73
from vllm.model_executor.layers.rotary_embedding import get_rope
74
from vllm.model_executor.layers.sparse_attn_indexer import (
75
SparseAttnIndexer,
@@ -628,10 +632,6 @@ class Indexer(nn.Module):
632
self.vllm_config = vllm_config
633
self.config = config
634
self.quant_config = quant_config
631
- self.is_fp4_ckpt = (
632
- self.quant_config is not None
633
- and self.quant_config.get_name() == "modelopt_fp4"
634
- )
635
# self.indexer_cfg = config.attn_module_list_cfg[0]["attn_index"]
636
self.topk_tokens = config.index_topk
637
self.n_head = config.index_n_heads # 64
@@ -646,36 +646,16 @@ class Indexer(nn.Module):
646
quant_config=quant_config,
647
prefix=f"{prefix}.wq_b",
648
)
649
- if self.is_fp4_ckpt:
650
- # Fused wk + weights_proj: single GEMM producing [head_dim + n_head].
651
- # weights_proj does not get quantized,
652
- # so we run both with quant_config=None
653
- # wk may be upcasted from the default quant;
654
- # experiments show fusion is always faster unless WK proj is in FP4,
655
- # which is not the case for all known quants.
656
- self.wk_weights_proj = MergedColumnParallelLinear(
657
- hidden_size,
658
- [self.head_dim, self.n_head],
659
- bias=False,
660
- quant_config=None,
661
- disable_tp=True,
662
- prefix=f"{prefix}.wk_weights_proj",
663
- )
664
- else:
665
- self.wk = ReplicatedLinear(
666
- hidden_size,
667
- self.head_dim,
668
- bias=False,
669
- quant_config=quant_config,
670
- prefix=f"{prefix}.wk",
671
- )
672
- self.weights_proj = ReplicatedLinear(
673
- hidden_size,
674
- self.n_head,
675
- bias=False,
676
- quant_config=None,
677
- prefix=f"{prefix}.weights_proj",
678
- )
649
+ # Fused wk + weights_proj: single GEMM producing [head_dim + n_head].
650
+ # FP8 wk weights are upcasted to BF16 during loading to maintain fusion.
651
+ self.wk_weights_proj = MergedColumnParallelLinear(
652
+ hidden_size,
653
+ [self.head_dim, self.n_head],
654
+ bias=False,
655
+ quant_config=None,
656
+ disable_tp=True,
657
+ prefix=f"{prefix}.wk_weights_proj",
658
+ )
659
self.k_norm = LayerNorm(self.head_dim, eps=1e-6)
660
self.softmax_scale = self.head_dim**-0.5
661
@@ -716,14 +696,10 @@ class Indexer(nn.Module):
696
q_pe, q_nope = torch.split(
697
q, [self.rope_dim, self.head_dim - self.rope_dim], dim=-1
698
)
719
- if self.is_fp4_ckpt:
720
- # Fused wk + weights_proj: one GEMM, then split
721
- kw, _ = self.wk_weights_proj(hidden_states)
722
- k = kw[:, : self.head_dim]
723
- weights = kw[:, self.head_dim :]
724
- else:
725
- k, _ = self.wk(hidden_states)
726
- weights, _ = self.weights_proj(hidden_states)
699
+ # Fused wk + weights_proj: one GEMM, then split
700
+ kw, _ = self.wk_weights_proj(hidden_states)
701
+ k = kw[:, : self.head_dim]
702
+ weights = kw[:, self.head_dim :]
703
704
k = self.k_norm(k)
705
k_pe, k_nope = torch.split(
@@ -761,6 +737,46 @@ class Indexer(nn.Module):
737
return self.indexer_op(hidden_states, q_fp8, k, weights)
738
739
740
+def _try_load_fp8_indexer_wk(name, tensor, buf, params_dict, loaded_params):
741
+ """
742
+ We fuse the WK and weights_proj projections, but in some checkpoints WK is stored
743
+ in FP8 with a separate weight_scale_inv, while weights_proj is stored in BF16.
744
+ Upcasting to BF16 during loading enables the fusion. This function loads the FP8 WK
745
+ weights and scale, and when both are available, dequantizes to BF16 and stores into
746
+ the fused wk_weights_proj.weight parameter.
747
+ """
748
+ if "indexer.wk." not in name or "wk_weights" in name:
749
+ return False # Weight is not an isolated WK weight for the indexer, ignore.
750
+ is_weight = name.endswith(".weight") and tensor.dtype == torch.float8_e4m3fn
751
+ is_scale = "weight_scale_inv" in name
752
+ if not is_weight and not is_scale:
753
+ return False # WK is not in FP8 format, ignore.
754
+ # Buffer this tensor (weight or scale) until both have arrived.
755
+ layer_prefix = name.rsplit(".wk.", 1)[0] # e.g. "model.layers.0.self_attn.indexer"
756
+ entry = buf.setdefault(layer_prefix, {})
757
+ entry["weight" if is_weight else "scale"] = tensor
758
+ if "weight" not in entry or "scale" not in entry:
759
+ return True # still waiting for the other param
760
+
761
+ # We have both weight and scale: dequantize FP8 to BF16.
762
+ weight_fp8, scale_inv = entry["weight"], entry["scale"]
763
+ del buf[layer_prefix]
764
+ block_size = weight_fp8.shape[1] // scale_inv.shape[1]
765
+ weight_bf16 = scaled_dequantize(
766
+ weight_fp8,
767
+ scale_inv,
768
+ group_shape=GroupShape(block_size, block_size),
769
+ out_dtype=torch.bfloat16,
770
+ )
771
+
772
+ # Load the dequantized weight into shard 0 of the fused buffer.
773
+ fused_name = f"{layer_prefix}.wk_weights_proj.weight"
774
+ param = params_dict[fused_name]
775
+ param.weight_loader(param, weight_bf16, 0)
776
+ loaded_params.add(fused_name)
777
+ return True
778
+
779
+
780
def _min_latency_fused_qkv_a_proj_impl(
781
input_: torch.Tensor,
782
weight: torch.Tensor,
@@ -1344,10 +1360,6 @@ class DeepseekV2ForCausalLM(
1360
quant_config = vllm_config.quant_config
1361
self.config = config
1362
self.quant_config = quant_config
1347
- self.is_fp4_ckpt = (
1348
- self.quant_config is not None
1349
- and self.quant_config.get_name() == "modelopt_fp4"
1350
- )
1363
1364
qk_nope_head_dim = getattr(config, "qk_nope_head_dim", 0)
1365
qk_rope_head_dim = getattr(config, "qk_rope_head_dim", 0)
@@ -1473,13 +1485,13 @@ class DeepseekV2ForCausalLM(
1485
("qkv_proj", "k_proj", "k"),
1486
("qkv_proj", "v_proj", "v"),
1487
]
1476
- if self.is_fp4_ckpt:
1477
- # Fused indexer wk + weights_proj (shard 0 = wk, shard 1 = weights_proj)
1478
- indexer_fused_mapping = [
1479
- ("wk_weights_proj", "wk", 0),
1480
- ("wk_weights_proj", "weights_proj", 1),
1481
- ]
1482
- stacked_params_mapping.extend(indexer_fused_mapping)
1488
+ # Fused indexer wk + weights_proj (shard 0 = wk, shard 1 = weights_proj)
1489
+ _pending_wk_fp8: dict = {} # When WK is in FP8, we dequant to BF16 for fusion
1490
+ indexer_fused_mapping = [
1491
+ ("wk_weights_proj", "wk", 0),
1492
+ ("wk_weights_proj", "weights_proj", 1),
1493
+ ]
1494
+ stacked_params_mapping.extend(indexer_fused_mapping)
1495
1496
if self.use_mha:
1497
stacked_params_mapping.extend(mha_params_mapping)
@@ -1516,6 +1528,11 @@ class DeepseekV2ForCausalLM(
1528
rocm_aiter_moe_shared_expert_enabled and ("mlp.shared_experts" in name)
1529
)
1530
1531
+ if _try_load_fp8_indexer_wk(
1532
+ name, loaded_weight, _pending_wk_fp8, params_dict, loaded_params
1533
+ ):
1534
+ continue
1535
+
1536
for param_name, weight_name, shard_id in stacked_params_mapping:
1537
# Skip non-stacked layers and experts (experts handled below).
1538
if weight_name not in name: