[DeepSeek-V4] Support TEP=16 for the block-FP8 shared expert (#46001)
Signed-off-by: Jeff Ma <jeffjma@umich.edu> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Jeff (Junze) Ma committed
Jun 18, 2026 at 20:10 UTC
2a6c6b94293edb54bff8088a5d64b703aac187ff
1 file changed
+44
vllm/models/deepseek_v4/nvidia/model.py
+44
@@ -64,6 +64,7 @@ from vllm.models.deepseek_v4.nvidia.flashinfer_sparse import (
64
from vllm.models.deepseek_v4.nvidia.flashmla import DeepseekV4FlashMLAAttention
65
from vllm.models.deepseek_v4.nvidia.ops.prepare_megamoe import prepare_megamoe_inputs
66
from vllm.sequence import IntermediateTensors
67
+from vllm.utils.math_utils import cdiv
68
from vllm.v1.attention.backends.registry import AttentionBackendEnum
69
70
@@ -85,6 +86,15 @@ class DeepseekV4MLP(nn.Module):
86
# across the ranks within the tp_group. In this case the weights are
87
# replicated and no collective ops are needed.
88
# Otherwise we use standard TP with an allreduce at the end.
89
+ #
90
+ # Block-FP8 shards in whole 128-blocks; cdiv rounds the per-rank block
91
+ # count up so the linear's even TP split stays block-aligned, with the
92
+ # trailing ranks zero-filled by load_weights.
93
+ block_size = getattr(quant_config, "weight_block_size", None)
94
+ if block_size is not None and not is_sequence_parallel:
95
+ tp_size = get_tensor_model_parallel_world_size()
96
+ n_local = cdiv(intermediate_size // block_size[0], tp_size)
97
+ intermediate_size = n_local * block_size[0] * tp_size
98
self.gate_up_proj = MergedColumnParallelLinear(
99
hidden_size,
100
[intermediate_size] * 2,
@@ -892,6 +902,8 @@ class DeepseekV4Model(nn.Module):
902
config = vllm_config.model_config.hf_config
903
quant_config = vllm_config.quant_config
904
self.config = config
905
+ self.quant_config = quant_config
906
+ self.parallel_config = vllm_config.parallel_config
907
self.use_mega_moe = (
908
vllm_config.kernel_config.moe_backend == "deep_gemm_mega_moe"
909
)
@@ -1080,7 +1092,17 @@ class DeepseekV4Model(nn.Module):
1092
# Pre-compute expert mapping ONCE.
1093
expert_mapping = self.get_expert_mapping()
1094
1095
+ # Block-FP8 shared experts: pad the intermediate up to the TP-uniform
1096
+ # block count so the standard loaders below slice it evenly (trailing
1097
+ # ranks land on the zero pad). SP / unquantized ones need no padding.
1098
+ pad_shared_expert = (
1099
+ getattr(self.quant_config, "weight_block_size", None) is not None
1100
+ and not self.parallel_config.use_sequence_parallel_moe
1101
+ )
1102
+
1103
for name, loaded_weight in weights:
1104
+ if pad_shared_expert and ".shared_experts." in name:
1105
+ loaded_weight = self._pad_shared_expert_weight(name, loaded_weight)
1106
for param_name, weight_name, shard_id in stacked_params_mapping:
1107
# Skip non-stacked layers and experts (experts handled below).
1108
if ".experts." in name:
@@ -1155,6 +1177,28 @@ class DeepseekV4Model(nn.Module):
1177
1178
return loaded_params
1179
1180
+ def _pad_shared_expert_weight(
1181
+ self, name: str, loaded_weight: torch.Tensor
1182
+ ) -> torch.Tensor:
1183
+ """Zero-pad a block-FP8 shared-expert weight/scale on its intermediate
1184
+ axis so the standard TP loaders split it into even, block-aligned shards
1185
+ (trailing ranks get the zero pad). gate (w1)/up (w3) [I, H] pad dim 0;
1186
+ down (w2 -> down_proj) [H, I] pads dim 1.
1187
+ """
1188
+ block_size = getattr(self.quant_config, "weight_block_size", None)
1189
+ assert block_size is not None
1190
+ # Round the intermediate axis up to a whole number of TP shards. The axis
1191
+ # is in elements for weights (step = block) and in blocks for scales.
1192
+ step = 1 if name.endswith("weight_scale_inv") else block_size[0]
1193
+ dim = 1 if ".down_proj." in name else 0
1194
+ mult = get_tensor_model_parallel_world_size() * step
1195
+ pad = cdiv(loaded_weight.shape[dim], mult) * mult - loaded_weight.shape[dim]
1196
+ if pad == 0:
1197
+ return loaded_weight
1198
+ pad_shape = list(loaded_weight.shape)
1199
+ pad_shape[dim] = pad
1200
+ return torch.cat([loaded_weight, loaded_weight.new_zeros(pad_shape)], dim=dim)
1201
+
1202
def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
1203
first_layer = next(iter(islice(self.layers, self.start_layer, self.end_layer)))
1204
if first_layer.ffn.use_mega_moe: