@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 41f9104fa6

[DSv4.1] Integrate Mega-Gate from DeepGEMM (#56266)

Signed-off-by: zjy0516 <riverclouds.zhu@qq.com> Signed-off-by: Thien Tran <gau.nernst@yahoo.com.sg> Co-authored-by: OpenAI Codex <codex@openai.com> Co-authored-by: zjy0516 <riverclouds.zhu@qq.com>

Thien Tran committed Sep 18, 2026 at 03:06 UTC 41f9104fa649aed2c26af8a459210619ad9514fe

9 files changed +369 -35

.buildkite/test_areas/kernels.yaml

@@ -120,6 +120,7 @@ steps:

120 - csrc/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu

121 - vllm/models/deepseek_v4/common/ops/

122 - vllm/models/deepseek_v4/nvidia/

123 + - vllm/utils/deep_gemm.py

124 - tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py

125 - tests/models/test_deepseek_v4_mega_moe.py

126 commands:

tests/kernels/test_mhc_kernels.py

+3 -2

@@ -132,6 +132,7 @@ def test_v41_dspark_head_collapses_with_last_ffn_mix(num_tokens, monkeypatch):

132

133 monkeypatch.setattr(dspark, "mhc_post_tilelang", lambda *args: streams)

134 draft = SimpleNamespace(

135 + use_mega_moe=False,

136 use_sequence_parallel=False,

137 hc_mult=hc_mult,

138 layers=[make_layer(mix) for mix in mixes],

@@ -452,7 +453,7 @@ def test_deepseek_v41_decoder_mixes_match_torch(

453 device=DEVICE, dtype=torch.bfloat16

454 )

455 decoder.attn = lambda positions, x, _: x * 0.5

455 - decoder.ffn = lambda x, input_ids: x * 0.25

456 + decoder.ffn = lambda x, input_ids, mega_gate_metadata=None: x * 0.25

457 with torch.device(DEVICE):

458 decoder.hc_attn_fn = torch.randn(24, 20480) * 0.02

459 decoder.hc_ffn_fn = torch.randn(24, 20480) * 0.02

@@ -536,7 +537,7 @@ def test_deepseek_v41_capture_previous_aux(entry, monkeypatch, default_vllm_conf

537 device=DEVICE, dtype=torch.bfloat16

538 )

539 decoder.attn = lambda positions, x, _: x * 0.5

539 - decoder.ffn = lambda x, input_ids: x * 0.25

540 + decoder.ffn = lambda x, input_ids, mega_gate_metadata=None: x * 0.25

541 with torch.device(DEVICE):

542 decoder.hc_attn_fn = torch.randn(24, 20480) * 0.02

543 decoder.hc_ffn_fn = torch.randn(24, 20480) * 0.02

tests/models/test_deepseek_v4_mega_moe.py

+138 -12

@@ -15,6 +15,7 @@ from vllm.models.deepseek_v4.nvidia.model import (

15 DeepseekV4MegaMoEExperts,

16 DeepseekV4MoE,

17 make_deepseek_v4_expert_params_mapping,

18 + prepare_mega_gate_routing_metadata,

19 )

20 from vllm.models.deepseek_v4.nvidia.mtp import DeepSeekV4MTP

21 from vllm.models.deepseek_v4.nvidia.ops.prepare_megamoe import prepare_megamoe_inputs

@@ -22,6 +23,7 @@ from vllm.models.deepseek_v41.common.mm_preprocess import IMAGE_SENTINEL_BASE_ID

23 from vllm.models.deepseek_v41.nvidia.model import DeepseekV4MoE as DeepseekV41MoE

24 from vllm.platforms import current_platform

25 from vllm.transformers_utils.configs.deepseek_v41 import DeepseekV41Config

26 +from vllm.utils.torch_utils import set_default_torch_dtype

28 pytestmark = pytest.mark.skipif(

29 not current_platform.is_cuda(),

@@ -33,6 +35,7 @@ pytestmark = pytest.mark.skipif(

35 def v41_moe_config(dist_init):

36 return SimpleNamespace(

37 model_config=SimpleNamespace(

38 + dtype=torch.bfloat16,

39 hf_config=DeepseekV41Config(

40 text_config=dict(

41 hidden_size=128,

@@ -63,28 +66,52 @@ def v41_moe_config(dist_init):

66 )

69 +@pytest.mark.parametrize("use_cudagraph", [False, True])

70 +@pytest.mark.parametrize("above_threshold", [False, True])

71 @pytest.mark.parametrize("vision", [False, True])

67 -@pytest.mark.parametrize("layer_id,num_experts,top_k", [(0, 8, 2), (2, 4, 3)])

72 +@pytest.mark.parametrize("layer_id,num_experts,top_k", [(0, 384, 6), (2, 128, 3)])

73 def test_deepseek_v41_moe_routes_without_hash_table(

69 - v41_moe_config, monkeypatch, vision, layer_id, num_experts, top_k

74 + v41_moe_config,

75 + monkeypatch,

76 + vision,

77 + layer_id,

78 + num_experts,

79 + top_k,

80 + above_threshold,

81 + use_cudagraph,

82 ):

83 """Main and draft layers select experts by score and preserve image routing."""

84 + if not current_platform.is_device_capability_family(100):

85 + pytest.skip("DeepGEMM Mega Gate requires SM100")

86 +

87 config = v41_moe_config.model_config.hf_config

88 + config.hidden_size = 256

89 + config.n_routed_experts = 384

90 + config.num_experts_per_tok = 6

91 + config.dspark_n_routed_experts = 128

92 + num_tokens = (1 if num_experts == 128 else 16) + int(above_threshold)

93 config.vision_n_layers = int(vision)

74 - with torch.device("cuda"):

94 + with (

95 + set_default_torch_dtype(v41_moe_config.model_config.dtype),

96 + torch.device("cuda"),

97 + ):

98 moe = DeepseekV41MoE(v41_moe_config, prefix=f"model.layers.{layer_id}.ffn")

76 - hidden_states = torch.randn(4, config.hidden_size)

99 + hidden_states = torch.randn(

100 + num_tokens, config.hidden_size, dtype=torch.bfloat16

101 + )

102 input_ids = (

78 - torch.tensor([42, IMAGE_SENTINEL_BASE_ID, IMAGE_SENTINEL_BASE_ID, 129257])

103 + torch.tensor(

104 + [42, IMAGE_SENTINEL_BASE_ID, IMAGE_SENTINEL_BASE_ID, 129257]

105 + ).repeat((num_tokens + 3) // 4)[:num_tokens]

106 if vision

107 else None

108 )

109 assert moe.gate.tid2eid is None

110 assert moe.gate.weight.shape == (num_experts, config.hidden_size)

111 assert isinstance(moe.experts, DeepseekV4MegaMoEExperts)

85 - assert moe.experts.w13_weight.shape == (num_experts, 256, 64)

112 + assert moe.experts.w13_weight.shape == (num_experts, 256, 128)

113 assert moe.experts.top_k == top_k

87 - assert (config.n_routed_experts, config.num_experts_per_tok) == (8, 2)

114 + assert (config.n_routed_experts, config.num_experts_per_tok) == (384, 6)

115

116 with torch.no_grad():

117 moe.gate.weight.normal_(std=0.01)

@@ -93,7 +120,7 @@ def test_deepseek_v41_moe_routes_without_hash_table(

120 moe.gate.bias_vl.copy_(-moe.gate.e_score_correction_bias)

121

122 scores = torch.nn.functional.softplus(

96 - torch.nn.functional.linear(hidden_states, moe.gate.weight)

123 + torch.mm(hidden_states, moe.gate.weight.t(), out_dtype=torch.float32)

124 ).sqrt()

125 bias = moe.gate.e_score_correction_bias

126 if vision:

@@ -105,15 +132,49 @@ def test_deepseek_v41_moe_routes_without_hash_table(

132 dim=-1, keepdim=True

133 )

134

135 + routed = {}

136 +

137 def check_routing(x, weights, ids, *, activation_clamp):

109 - torch.testing.assert_close(ids, expected_ids)

110 - torch.testing.assert_close(weights, expected_weights)

138 + routed["ids"], routed["weights"] = ids, weights

139 assert activation_clamp == config.swiglu_limit

140 return x.clone()

141

142 monkeypatch.setattr(moe.experts, "forward", check_routing)

143 monkeypatch.setattr(moe.shared_experts, "forward", lambda x: 2 * x)

116 - torch.testing.assert_close(moe(hidden_states, input_ids), 3 * hidden_states)

144 + routing_input_ids = (

145 + input_ids

146 + if input_ids is not None

147 + else torch.zeros(num_tokens, dtype=torch.int64, device="cuda")

148 + )

149 + metadata = prepare_mega_gate_routing_metadata(

150 + routing_input_ids,

151 + has_hash_routing=False,

152 + image_sentinel_base_id=IMAGE_SENTINEL_BASE_ID if vision else None,

153 + )

154 + if use_cudagraph:

155 + stream = torch.cuda.Stream()

156 + stream.wait_stream(torch.cuda.current_stream())

157 + with torch.cuda.stream(stream):

158 + for _ in range(3):

159 + moe(hidden_states, input_ids, metadata)

160 + stream.synchronize()

161 + graph = torch.cuda.CUDAGraph()

162 + with torch.cuda.graph(graph, stream=stream):

163 + output = moe(hidden_states, input_ids, metadata)

164 + graph.replay()

165 + torch.accelerator.synchronize()

166 + else:

167 + output = moe(hidden_states, input_ids, metadata)

168 +

169 + ids, weights = routed["ids"], routed["weights"]

170 + torch.testing.assert_close(ids, expected_ids)

171 + torch.testing.assert_close(

172 + weights,

173 + expected_weights,

174 + rtol=1e-3,

175 + atol=1e-4,

176 + )

177 + torch.testing.assert_close(output, 3 * hidden_states)

178

179

180 @pytest.mark.parametrize(

@@ -156,6 +217,67 @@ def test_deepseek_v4_moe_preserves_configured_hash_layers(v41_moe_config):

217 moe(torch.zeros(1, 128))

218

219

220 +@pytest.mark.parametrize("num_tokens", [16, 17])

221 +def test_deepseek_v4_mega_gate_hash_routing_correctness(

222 + v41_moe_config, monkeypatch, num_tokens

223 +):

224 + if not current_platform.is_device_capability_family(100):

225 + pytest.skip("DeepGEMM Mega Gate requires SM100")

226 +

227 + config = v41_moe_config.model_config.hf_config

228 + config.hidden_size = 256

229 + config.num_hash_layers = 1

230 + config.vocab_size = 32

231 + with (

232 + set_default_torch_dtype(v41_moe_config.model_config.dtype),

233 + torch.device("cuda"),

234 + ):

235 + moe = DeepseekV4MoE(

236 + v41_moe_config,

237 + prefix="model.layers.0.ffn",

238 + num_hash_layers=config.num_hash_layers,

239 + )

240 + hidden_states = torch.randn(

241 + num_tokens, config.hidden_size, dtype=torch.bfloat16

242 + )

243 + input_ids = torch.arange(num_tokens)

244 +

245 + token_ids = torch.arange(config.vocab_size, device="cuda")

246 + fixed_ids = torch.stack(

247 + (

248 + token_ids % config.n_routed_experts,

249 + (token_ids + 3) % config.n_routed_experts,

250 + ),

251 + dim=1,

252 + )

253 + with torch.no_grad():

254 + moe.gate.weight.normal_(std=0.01)

255 + moe.gate.tid2eid.copy_(fixed_ids)

256 +

257 + expected_ids = fixed_ids[input_ids]

258 + scores = torch.nn.functional.softplus(

259 + torch.mm(hidden_states, moe.gate.weight.t(), out_dtype=torch.float32)

260 + ).sqrt()

261 + expected_weights = scores.gather(1, expected_ids)

262 + expected_weights *= config.routed_scaling_factor / expected_weights.sum(

263 + dim=-1, keepdim=True

264 + )

265 +

266 + def check_routing(x, weights, ids, *, activation_clamp):

267 + torch.testing.assert_close(ids, expected_ids)

268 + torch.testing.assert_close(weights, expected_weights, rtol=1e-3, atol=1e-4)

269 + return x.clone()

270 +

271 + monkeypatch.setattr(moe.experts, "forward", check_routing)

272 + monkeypatch.setattr(moe.shared_experts, "forward", torch.zeros_like)

273 + metadata = prepare_mega_gate_routing_metadata(

274 + input_ids,

275 + has_hash_routing=True,

276 + image_sentinel_base_id=None,

277 + )

278 + torch.testing.assert_close(moe(hidden_states, input_ids, metadata), hidden_states)

279 +

280 +

281 def test_deepseek_v4_mega_moe_expert_mapping():

282 mapping = make_deepseek_v4_expert_params_mapping(2)

283

@@ -513,6 +635,7 @@ def test_deepseek_v4_mega_moe_does_not_double_add_fused_shared_expert(

635 monkeypatch, fused

636 ):

637 class FakeGate(torch.nn.Module):

638 + weight = torch.empty(2, 128)

639 tid2eid = None

640 e_score_correction_bias = None

641

@@ -760,7 +883,8 @@ def test_deepseek_v4_drafter_pwal_hooks_finalize_mega_moe():

883 not torch.cuda.is_available(),

884 reason="DeepSeek V4 MegaMoE fused input staging requires CUDA.",

885 )

763 -def test_deepseek_v4_mega_moe_fused_input_staging_masks_padding():

886 +@pytest.mark.parametrize("nonfinite_padding", [False, True])

887 +def test_deepseek_v4_mega_moe_fused_input_staging_masks_padding(nonfinite_padding):

888 from vllm.third_party.deep_gemm.utils import per_token_cast_to_fp8

889

890 device = torch.device("cuda")

@@ -796,6 +920,8 @@ def test_deepseek_v4_mega_moe_fused_input_staging_masks_padding():

920 [False, True, False, False, True, False, True],

921 device=device,

922 )

923 + if nonfinite_padding:

924 + topk_weights[is_padding] = float("nan")

925

926 ref_x, ref_x_sf = per_token_cast_to_fp8(

927 hidden_states,

vllm/models/deepseek_v4/nvidia/dspark.py

+14

@@ -17,6 +17,7 @@ import torch.nn as nn

18 import vllm.envs as envs

19 from vllm.config import VllmConfig, get_current_vllm_config

20 +from vllm.config.kernel import MEGA_MOE_BACKENDS

21 from vllm.distributed import (

22 get_tensor_model_parallel_rank,

23 get_tensor_model_parallel_world_size,

@@ -51,12 +52,14 @@ from vllm.models.common.ops.sequence_parallel import (

52 sp_padding_mask,

53 sp_shard,

54 )

55 +from vllm.models.deepseek_v4.common.mm_preprocess import IMAGE_SENTINEL_BASE_ID

57 from .model import (

58 DeepseekV4DecoderLayer,

59 DeepseekV4Model,

60 _use_sequence_parallel,

61 make_deepseek_v4_expert_params_mapping,

62 + prepare_mega_gate_routing_metadata,

63 )

65 logger = init_logger(__name__)

@@ -96,6 +99,7 @@ class DSparkDeepseekV4Model(nn.Module):

99 self.rms_norm_eps = config.rms_norm_eps

100 self.num_hidden_layers = config.num_hidden_layers

101 self.target_layer_ids = tuple(config.dspark_target_layer_ids)

102 + self.use_mega_moe = vllm_config.kernel_config.moe_backend in MEGA_MOE_BACKENDS

103 self.use_sequence_parallel = _use_sequence_parallel(vllm_config)

104

105 self.num_dspark_layers = getattr(config, "n_mtp_layers", None) or 3

@@ -235,6 +239,15 @@ class DSparkDeepseekV4Model(nn.Module):

239 )

240 inputs_embeds = sp_shard(inputs_embeds)

241 input_ids = sp_shard(input_ids)

242 + mega_gate_metadata = None

243 + if self.use_mega_moe:

244 + mega_gate_metadata = prepare_mega_gate_routing_metadata(

245 + input_ids,

246 + has_hash_routing=False,

247 + image_sentinel_base_id=IMAGE_SENTINEL_BASE_ID

248 + if getattr(self.config, "vision_n_layers", 0) > 0

249 + else None,

250 + )

251 # Expand to hc_mult copies for hyper-connections ([T, H] -> [T, hc, H]).

252 hidden_states = inputs_embeds.unsqueeze(-2).repeat(1, self.hc_mult, 1)

253

@@ -247,6 +260,7 @@ class DSparkDeepseekV4Model(nn.Module):

260 post_mix,

261 res_mix,

262 residual,

263 + mega_gate_metadata=mega_gate_metadata,

264 )

265 hidden_states = mhc_post_tilelang(hidden_states, residual, post_mix, res_mix)

266 if self.use_sequence_parallel:

vllm/models/deepseek_v4/nvidia/model.py

+110 -19

@@ -98,6 +98,37 @@ from ..common.mm_preprocess import IMAGE_SENTINEL_BASE_ID

98 logger = init_logger(__name__)

100

101 +class MegaGateRoutingMetadata(typing.NamedTuple):

102 + safe_hash_input_ids: torch.Tensor | None

103 + image_token_mask: torch.Tensor | None

104 + hash_token_mask: torch.Tensor | None

105 +

106 +

107 +def prepare_mega_gate_routing_metadata(

108 + input_ids: torch.Tensor,

109 + *,

110 + has_hash_routing: bool,

111 + image_sentinel_base_id: int | None,

112 +) -> MegaGateRoutingMetadata:

113 + if image_sentinel_base_id is None:

114 + return MegaGateRoutingMetadata(

115 + input_ids if has_hash_routing else None,

116 + None,

117 + torch.ones_like(input_ids, dtype=torch.bool) if has_hash_routing else None,

118 + )

119 +

120 + image_token_mask = (input_ids >= image_sentinel_base_id) & (

121 + input_ids < image_sentinel_base_id + 5

122 + )

123 + if not has_hash_routing:

124 + return MegaGateRoutingMetadata(None, image_token_mask, None)

125 + return MegaGateRoutingMetadata(

126 + torch.where(image_token_mask, 0, input_ids),

127 + image_token_mask,

128 + ~image_token_mask,

129 + )

130 +

131 +

132 class DeepseekV4MLP(nn.Module):

133 def __init__(

134 self,

@@ -803,6 +834,10 @@ class DeepseekV4MoE(nn.Module):

834 raise NotImplementedError(

835 "DeepSeek V4 MegaMoE currently supports sqrtsoftplus routing only."

836 )

837 + if self.use_mega_moe and not self.renormalize:

838 + raise NotImplementedError(

839 + "DeepSeek V4 MegaMoE requires normalized top-k probabilities."

840 + )

841 if self.use_mega_moe and getattr(config, "expert_dtype", "fp4") != "fp4":

842 raise NotImplementedError(

843 "DeepSeek V4 MegaMoE only supports fp4 experts; got expert_dtype="

@@ -1000,7 +1035,10 @@ class DeepseekV4MoE(nn.Module):

1035 )

1036

1037 def forward(

1003 - self, hidden_states: torch.Tensor, input_ids: torch.Tensor | None = None

1038 + self,

1039 + hidden_states: torch.Tensor,

1040 + input_ids: torch.Tensor | None = None,

1041 + mega_gate_metadata: MegaGateRoutingMetadata | None = None,

1042 ) -> torch.Tensor:

1043 if self.gate.tid2eid is not None and input_ids is None:

1044 raise ValueError("DeepSeek V4 hash MoE routing requires input_ids.")

@@ -1013,23 +1051,61 @@ class DeepseekV4MoE(nn.Module):

1051 return self._forward_fused_moe(hidden_states, input_ids)

1052

1053 org_shape = hidden_states.shape

1016 - router_logits, _ = self.gate(hidden_states)

1017 - topk_weights, topk_ids = fused_topk_bias(

1018 - hidden_states=hidden_states,

1019 - gating_output=router_logits,

1020 - scoring_func=self.scoring_func,

1021 - e_score_correction_bias=self.gate.e_score_correction_bias.data

1022 - if self.gate.e_score_correction_bias is not None

1023 - else None,

1024 - topk=self.n_activated_experts,

1025 - renormalize=self.renormalize,

1026 - indices_type=self.hash_indices_dtype,

1027 - input_tokens=input_ids,

1028 - hash_indices_table=self.gate.tid2eid,

1029 - routed_scaling_factor=self.routed_scaling_factor,

1030 - bias_vl=bias_vl.data if bias_vl is not None else None,

1031 - image_sentinel_lo=self.image_sentinel_lo if bias_vl is not None else 0,

1032 - )

1054 + # Small local padded batches favor GateLinear; 128-expert gates cross earlier.

1055 + gate_threshold = 1 if self.gate.weight.shape[0] == 128 else 16

1056 + if hidden_states.shape[0] <= gate_threshold:

1057 + router_logits, _ = self.gate(hidden_states)

1058 + topk_weights, topk_ids = fused_topk_bias(

1059 + hidden_states=hidden_states,

1060 + gating_output=router_logits,

1061 + scoring_func=self.scoring_func,

1062 + e_score_correction_bias=self.gate.e_score_correction_bias.data

1063 + if self.gate.e_score_correction_bias is not None

1064 + else None,

1065 + topk=self.n_activated_experts,

1066 + renormalize=self.renormalize,

1067 + indices_type=self.hash_indices_dtype,

1068 + input_tokens=input_ids,

1069 + hash_indices_table=self.gate.tid2eid,

1070 + routed_scaling_factor=self.routed_scaling_factor,

1071 + bias_vl=bias_vl.data if bias_vl is not None else None,

1072 + image_sentinel_lo=self.image_sentinel_lo if bias_vl is not None else 0,

1073 + )

1074 + else:

1075 + from vllm.utils.deep_gemm import bf16_mega_gate

1076 +

1077 + if mega_gate_metadata is None:

1078 + raise RuntimeError(

1079 + "MegaMoE routing metadata must be prepared once by the model."

1080 + )

1081 +

1082 + unmapped_topk_idx = None

1083 + fix_routing_mask = None

1084 + if self.gate.tid2eid is not None:

1085 + if mega_gate_metadata.safe_hash_input_ids is None:

1086 + raise RuntimeError("Hash routing requires prepared input IDs.")

1087 + if mega_gate_metadata.hash_token_mask is None:

1088 + raise RuntimeError("Hash routing requires a prepared routing mask.")

1089 + unmapped_topk_idx = self.gate.tid2eid[

1090 + mega_gate_metadata.safe_hash_input_ids

1091 + ]

1092 + fix_routing_mask = mega_gate_metadata.hash_token_mask

1093 +

1094 + topk_weights, topk_ids = bf16_mega_gate(

1095 + hidden_states,

1096 + self.gate.weight,

1097 + self.n_activated_experts,

1098 + scoring_func=self.scoring_func,

1099 + routed_scaling_factor=self.routed_scaling_factor,

1100 + ep_rank=self.ep_rank,

1101 + bias=self.gate.e_score_correction_bias.data

1102 + if self.gate.e_score_correction_bias is not None

1103 + else None,

1104 + image_bias=bias_vl.data if bias_vl is not None else None,

1105 + image_token_mask=mega_gate_metadata.image_token_mask,

1106 + fix_routing_mask=fix_routing_mask,

1107 + unmapped_topk_idx=unmapped_topk_idx,

1108 + )

1109 activation_clamp = (

1110 float(self.swiglu_limit) if self.swiglu_limit is not None else None

1111 )

@@ -1248,6 +1324,7 @@ class DeepseekV4DecoderLayer(nn.Module):

1324 post_mix: torch.Tensor | None = None,

1325 res_mix: torch.Tensor | None = None,

1326 residual: torch.Tensor | None = None,

1327 + mega_gate_metadata: MegaGateRoutingMetadata | None = None,

1328 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:

1329 attn_norm_weight = self.attn_norm.weight.data

1330 attn_norm_eps = self.attn_norm.variance_epsilon

@@ -1332,7 +1409,7 @@ class DeepseekV4DecoderLayer(nn.Module):

1409 norm_eps=ffn_norm_eps,

1410 )

1411

1335 - x = self.ffn(x, input_ids)

1412 + x = self.ffn(x, input_ids, mega_gate_metadata)

1413 return x, residual, post_mix, res_mix

1414

1415

@@ -1394,6 +1471,9 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):

1471 ),

1472 prefix=f"{prefix}.layers",

1473 )

1474 + self.has_local_hash_moe = self.start_layer < min(

1475 + self.end_layer, config.num_hash_layers

1476 + )

1477

1478 if get_pp_group().is_last_rank:

1479 self.norm = RMSNorm(config.hidden_size, self.rms_norm_eps)

@@ -1496,6 +1576,16 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):

1576 hidden_states = sp_shard(hidden_states)

1577 input_ids = sp_shard(input_ids)

1578

1579 + mega_gate_metadata = None

1580 + if self.use_mega_moe:

1581 + mega_gate_metadata = prepare_mega_gate_routing_metadata(

1582 + input_ids,

1583 + has_hash_routing=self.has_local_hash_moe,

1584 + image_sentinel_base_id=IMAGE_SENTINEL_BASE_ID

1585 + if getattr(self.config, "vision_n_layers", 0) > 0

1586 + else None,

1587 + )

1588 +

1589 residual, post_mix, res_mix = None, None, None

1590 remote_aux = self.collect_remote_aux_hidden_states(intermediate_tensors)

1591 aux_hidden_states: list[torch.Tensor] = []

@@ -1511,6 +1601,7 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):

1601 post_mix,

1602 res_mix,

1603 residual,

1604 + mega_gate_metadata=mega_gate_metadata,

1605 )

1606 if idx + 1 in self.aux_hidden_state_layers:

1607 # Reconstruct the aux hidden state for draft models

vllm/models/deepseek_v4/nvidia/mtp.py

+20 -1

@@ -20,6 +20,7 @@ import torch.nn as nn

21 import vllm.envs as envs

22 from vllm.config import VllmConfig

23 +from vllm.config.kernel import MEGA_MOE_BACKENDS

24 from vllm.distributed import (

25 get_tensor_model_parallel_rank,

26 get_tensor_model_parallel_world_size,

@@ -51,6 +52,7 @@ from vllm.models.common.ops.sequence_parallel import (

52 sp_padding_mask,

53 sp_shard,

54 )

55 +from vllm.models.deepseek_v4.common.mm_preprocess import IMAGE_SENTINEL_BASE_ID

56 from vllm.models.deepseek_v4.common.ops.fused_mtp_input_rmsnorm import (

57 _FUSED_MTP_INPUT_RMSNORM_KERNEL,

58 _MTP_SHARED_HEAD_RMSNORM_KERNEL,

@@ -62,6 +64,7 @@ from .model import (

64 DeepseekV4Model,

65 _use_sequence_parallel,

66 make_deepseek_v4_expert_params_mapping,

67 + prepare_mega_gate_routing_metadata,

68 )

70 logger = init_logger(__name__)

@@ -92,6 +95,7 @@ class DeepSeekV4MultiTokenPredictorLayer(nn.Module):

95 self.config = config

96 quant_config = vllm_config.quant_config

97 self.rms_norm_eps = config.rms_norm_eps

98 + self.use_mega_moe = vllm_config.kernel_config.moe_backend in MEGA_MOE_BACKENDS

100 self.enorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

101 self.hnorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

@@ -181,8 +185,23 @@ class DeepSeekV4MultiTokenPredictorLayer(nn.Module):

185 hidden_states = self.h_proj(previous_hidden_states) + self.e_proj(

186 inputs_embeds

187 ).unsqueeze(-2)

188 + mega_gate_metadata = None

189 + if self.use_mega_moe:

190 + routing_input_ids = input_ids

191 + if self.mtp_block.use_sequence_parallel:

192 + routing_input_ids = sp_shard(routing_input_ids)

193 + mega_gate_metadata = prepare_mega_gate_routing_metadata(

194 + routing_input_ids,

195 + has_hash_routing=False,

196 + image_sentinel_base_id=IMAGE_SENTINEL_BASE_ID

197 + if getattr(self.config, "vision_n_layers", 0) > 0

198 + else None,

199 + )

200 hidden_states, residual, post_mix, res_mix = self.mtp_block(

185 - positions=positions, x=hidden_states, input_ids=input_ids

201 + positions=positions,

202 + x=hidden_states,

203 + input_ids=input_ids,

204 + mega_gate_metadata=mega_gate_metadata,

205 )

206 hidden_states = mhc_post_tilelang(hidden_states, residual, post_mix, res_mix)

207 if self.mtp_block.use_sequence_parallel:

vllm/models/deepseek_v41/nvidia/dspark.py

+14

@@ -18,6 +18,7 @@ import torch.nn as nn

19 import vllm.envs as envs

20 from vllm.config import VllmConfig, get_current_vllm_config

21 +from vllm.config.kernel import MEGA_MOE_BACKENDS

22 from vllm.distributed import (

23 get_tensor_model_parallel_rank,

24 get_tensor_model_parallel_world_size,

@@ -50,12 +51,14 @@ from vllm.models.common.ops.sequence_parallel import (

51 sp_shard,

52 )

54 +from ..common.mm_preprocess import IMAGE_SENTINEL_BASE_ID

55 from .model import (

56 DeepseekV4DecoderLayer,

57 DeepseekV4Model,

58 _linear_scale_param_name,

59 _use_sequence_parallel,

60 make_deepseek_v4_expert_params_mapping,

61 + prepare_mega_gate_routing_metadata,

62 )

64 logger = init_logger(__name__)

@@ -77,6 +80,7 @@ class DSparkDeepseekV4Model(nn.Module):

80 self.rms_norm_eps = config.rms_norm_eps

81 self.num_hidden_layers = config.num_hidden_layers

82 self.target_layer_ids = tuple(config.dspark_target_layer_ids)

83 + self.use_mega_moe = vllm_config.kernel_config.moe_backend in MEGA_MOE_BACKENDS

84 self.use_sequence_parallel = _use_sequence_parallel(vllm_config)

86 self.num_dspark_layers = (

@@ -202,6 +206,15 @@ class DSparkDeepseekV4Model(nn.Module):

206 )

207 inputs_embeds = sp_shard(inputs_embeds)

208 input_ids = sp_shard(input_ids)

209 + mega_gate_metadata = None

210 + if self.use_mega_moe:

211 + mega_gate_metadata = prepare_mega_gate_routing_metadata(

212 + input_ids,

213 + has_hash_routing=False,

214 + image_sentinel_base_id=IMAGE_SENTINEL_BASE_ID

215 + if getattr(self.config, "vision_n_layers", 0) > 0

216 + else None,

217 + )

218 # Expand to hc_mult copies for hyper-connections ([T, H] -> [T, hc, H]).

219 hidden_states = inputs_embeds.unsqueeze(-2).repeat(1, self.hc_mult, 1)

220

@@ -215,6 +228,7 @@ class DSparkDeepseekV4Model(nn.Module):

228 post_mix,

229 res_mix,

230 residual,

231 + mega_gate_metadata=mega_gate_metadata,

232 )

233 hidden_states = mhc_post_tilelang(hidden_states, residual, post_mix, res_mix)

234 # Collapse the hc copies with the pre-mix from the last layer's FFN

vllm/models/deepseek_v41/nvidia/model.py

+15 -1

@@ -61,7 +61,9 @@ from vllm.models.deepseek_v4.nvidia.model import (

61 DeepseekV4MoE as DeepseekV4MoEBase,

62 )

63 from vllm.models.deepseek_v4.nvidia.model import (

64 + MegaGateRoutingMetadata,

65 make_deepseek_v4_expert_params_mapping,

66 + prepare_mega_gate_routing_metadata,

67 )

68 from vllm.models.deepseek_v41.attention import DeepseekV4Attention

69 from vllm.models.deepseek_v41.nvidia.flash_mla_mega_attn import (

@@ -339,6 +341,7 @@ class DeepseekV4DecoderLayer(nn.Module):

341 engram_mask: torch.Tensor | None = None,

342 *,

343 capture_previous_aux: bool = False,

344 + mega_gate_metadata: MegaGateRoutingMetadata | None = None,

345 ) -> tuple[

346 torch.Tensor,

347 torch.Tensor,

@@ -462,7 +465,7 @@ class DeepseekV4DecoderLayer(nn.Module):

465 norm_weight=self.ffn_norm.weight,

466 norm_eps=self.ffn_norm.variance_epsilon,

467 )

465 - x = self.ffn(x, input_ids)

468 + x = self.ffn(x, input_ids, mega_gate_metadata)

469 return x, residual, post_mix, res_mix, ffn_pre, previous_aux

470

471

@@ -705,6 +708,16 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):

708 hidden_states = sp_shard(hidden_states)

709 input_ids = sp_shard(input_ids)

710

711 + mega_gate_metadata = None

712 + if self.use_mega_moe:

713 + mega_gate_metadata = prepare_mega_gate_routing_metadata(

714 + input_ids,

715 + has_hash_routing=False,

716 + image_sentinel_base_id=IMAGE_SENTINEL_BASE_ID

717 + if getattr(self.config, "vision_n_layers", 0) > 0

718 + else None,

719 + )

720 +

721 residual, post_mix, res_mix = None, None, None

722 pre_mix: torch.Tensor | None = None

723 if not get_pp_group().is_first_rank:

@@ -728,6 +741,7 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):

741 engram_hashes,

742 engram_mask,

743 capture_previous_aux=idx in self.aux_hidden_state_layers,

744 + mega_gate_metadata=mega_gate_metadata,

745 )

746 if previous_aux is not None:

747 # idx is the one-based id of the layer whose post this is.

vllm/utils/deep_gemm.py

+54

@@ -164,6 +164,7 @@ _get_sparse_mqa_logits_metadata_impl: Callable[..., Any] | None = None

164 _get_paged_sparse_mqa_logits_metadata_impl: Callable[..., Any] | None = None

165 _tf32_hc_prenorm_gemm_impl: Callable[..., Any] | None = None

166 _mega_mhc_impl: Callable[..., Any] | None = None

167 +_bf16_mega_gate_impl: Callable[..., Any] | None = None

168 _get_mn_major_tma_aligned_tensor_impl: Callable[..., Any] | None = None

169 _get_mk_alignment_for_contiguous_layout_impl: Callable[..., Any] | None = None

170 _get_theoretical_mk_alignment_for_contiguous_layout_impl: Callable[..., Any] | None = (

@@ -240,6 +241,7 @@ def _lazy_init() -> None:

241 global _get_sparse_mqa_logits_metadata_impl

242 global _get_paged_sparse_mqa_logits_metadata_impl

243 global _tf32_hc_prenorm_gemm_impl, _mega_mhc_impl

244 + global _bf16_mega_gate_impl

245 global _get_mn_major_tma_aligned_tensor_impl

246 global _get_mk_alignment_for_contiguous_layout_impl

247 global _get_theoretical_mk_alignment_for_contiguous_layout_impl

@@ -260,6 +262,7 @@ def _lazy_init() -> None:

262 or _get_paged_mqa_logits_metadata_impl is not None

263 or _tf32_hc_prenorm_gemm_impl is not None

264 or _mega_mhc_impl is not None

265 + or _bf16_mega_gate_impl is not None

266 or _get_mk_alignment_for_contiguous_layout_impl is not None

267 or _transform_sf_into_required_layout_impl is not None

268 or _pack_ue8m0_to_int_impl is not None

@@ -315,6 +318,7 @@ def _lazy_init() -> None:

318 )

319 _tf32_hc_prenorm_gemm_impl = getattr(_dg, "tf32_hc_prenorm_gemm", None)

320 _mega_mhc_impl = getattr(_dg, "mega_mhc", None)

321 + _bf16_mega_gate_impl = getattr(_dg, "bf16_mega_gate", None)

322 _get_mn_major_tma_aligned_tensor_impl = getattr(

323 _dg, "get_mn_major_tma_aligned_tensor", None

324 )

@@ -482,6 +486,55 @@ def cublaslt_gemm_nt(*args, **kwargs):

486 return _cublaslt_gemm_nt_impl(*args, **kwargs)

487

488

489 +def bf16_mega_gate(

490 + x: torch.Tensor,

491 + weight: torch.Tensor,

492 + num_topk: int,

493 + *,

494 + scoring_func: str,

495 + routed_scaling_factor: float,

496 + ep_rank: int,

497 + bias: torch.Tensor | None = None,

498 + image_bias: torch.Tensor | None = None,

499 + image_token_mask: torch.Tensor | None = None,

500 + fix_routing_mask: torch.Tensor | None = None,

501 + unmapped_topk_idx: torch.Tensor | None = None,

502 +) -> tuple[torch.Tensor, torch.Tensor]:

503 + """Run DeepGEMM Mega Gate.

504 +

505 + The output order follows vLLM's router contract: weights, then indices.

506 + Passing PyTorch-owned outputs to DeepGEMM makes their storage part of CUDA

507 + graph capture.

508 + """

509 + _lazy_init()

510 + if _bf16_mega_gate_impl is None:

511 + raise RuntimeError(

512 + "DeepGEMM MegaMoE requires a DeepGEMM build with bf16_mega_gate."

513 + )

514 +

515 + topk_idx = torch.empty((x.shape[0], num_topk), dtype=torch.int64, device=x.device)

516 + topk_weights = torch.empty(

517 + (x.shape[0], num_topk), dtype=torch.float32, device=x.device

518 + )

519 + _bf16_mega_gate_impl(

520 + x,

521 + weight,

522 + num_topk,

523 + use_shared_as_routed=False,

524 + num_shared_experts=0,

525 + routed_scaling_factor=routed_scaling_factor,

526 + ep_rank=ep_rank,

527 + scoring_func=scoring_func,

528 + bias=bias,

529 + image_bias=image_bias,

530 + image_token_mask=image_token_mask,

531 + fix_routing_mask=fix_routing_mask,

532 + unmapped_topk_idx=unmapped_topk_idx,

533 + out=(topk_idx, topk_weights),

534 + )

535 + return topk_weights, topk_idx

536 +

537 +

538 def fp8_gemm_nt(*args, **kwargs):

539 _lazy_init()

540 if _fp8_gemm_nt_impl is None:

@@ -988,6 +1041,7 @@ def should_use_deepgemm_for_fp8_linear(

1041

1042 __all__ = [

1043 "calc_diff",

1044 + "bf16_mega_gate",

1045 "DeepGemmQuantScaleFMT",

1046 "fp8_gemm_nt",

1047 "fp8_einsum",