@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / b06813e872
Fork 0

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 2, 22:00 UTC

[Kernel] Add mhc_pre_big_fuse_with_norm_tilelang (#43474)

Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai>

Jee Jee Li committed May 25, 2026 at 09:19 UTC b06813e87207e15b133e903d641e03f237d85b17
5 files changed +360 -91
vllm/_tilelang_ops.py
+167
@@ -177,6 +177,173 @@ def mhc_pre_big_fuse_tilelang(
177 T.pdl_trigger()
178
179
180 +# Copied from https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/layers/mhc.py#L478
181 +
182 +
183 +@tilelang.jit(
184 + pass_configs={
185 + tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True,
186 + tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True,
187 + tilelang.PassConfigKey.TL_PTXAS_REGISTER_USAGE_LEVEL: 10,
188 + },
189 +)
190 +def mhc_pre_big_fuse_with_norm_tilelang(
191 + gemm_out_mul,
192 + gemm_out_sqrsum,
193 + hc_scale,
194 + hc_base,
195 + residual,
196 + post_mix,
197 + comb_mix,
198 + layer_input,
199 + norm_weight,
200 + hidden_size: int,
201 + rms_eps: float,
202 + hc_pre_eps: float,
203 + hc_sinkhorn_eps: float,
204 + hc_post_mult_value: float,
205 + sinkhorn_repeat: int,
206 + norm_eps: float,
207 + n_splits: int = 16,
208 + hc_mult: int = 4,
209 + gemm_last_dim: int = -1,
210 +):
211 + num_tokens = T.dynamic("num_tokens")
212 + hc_mult3 = hc_mult * (2 + hc_mult)
213 + if gemm_last_dim < 0:
214 + gemm_last_dim = hc_mult3
215 + hidden_block = math.gcd(1024, hidden_size)
216 +
217 + gemm_out_mul: T.Tensor[[n_splits, num_tokens, gemm_last_dim], T.float32] # type: ignore[no-redef, valid-type]
218 + gemm_out_sqrsum: T.Tensor[[n_splits, num_tokens], T.float32] # type: ignore[no-redef, valid-type]
219 + hc_scale: T.Tensor[[3], T.float32] # type: ignore[no-redef, valid-type]
220 + hc_base: T.Tensor[[hc_mult3], T.float32] # type: ignore[no-redef, valid-type]
221 + residual: T.Tensor[[num_tokens, hc_mult, hidden_size], T.bfloat16] # type: ignore[no-redef, valid-type]
222 + post_mix: T.Tensor[[num_tokens, hc_mult], T.float32] # type: ignore[no-redef, valid-type]
223 + comb_mix: T.Tensor[[num_tokens, hc_mult * hc_mult], T.float32] # type: ignore[no-redef, valid-type]
224 + layer_input: T.Tensor[[num_tokens, hidden_size], T.bfloat16] # type: ignore[no-redef, valid-type]
225 + norm_weight: T.Tensor[[hidden_size], T.bfloat16] # type: ignore[no-redef, valid-type]
226 +
227 + with T.Kernel(num_tokens, threads=96) as i:
228 + rms = T.alloc_fragment(1, T.float32)
229 + mixes = T.alloc_fragment(hc_mult3, T.float32)
230 + T.clear(mixes)
231 + rms[0] = 0
232 +
233 + T.pdl_sync()
234 +
235 + for i_split in T.serial(n_splits):
236 + rms[0] += gemm_out_sqrsum[i_split, i]
237 + rms[0] = T.rsqrt(rms[0] / (hc_mult * hidden_size) + rms_eps)
238 + for j in T.Parallel(hc_mult3):
239 + mixes[j] = 0
240 + for i_split in T.serial(n_splits):
241 + mixes[j] += gemm_out_mul[i_split, i, j]
242 + mixes[j] *= rms[0]
243 + mixes_shared = T.alloc_shared(hc_mult3, T.float32)
244 + T.copy(mixes, mixes_shared)
245 +
246 + if T.get_thread_binding() < 32:
247 + cm = T.alloc_fragment((hc_mult, hc_mult), T.float32)
248 + for j in T.Parallel(hc_mult):
249 + post_mix[i, j] = (
250 + T.sigmoid(
251 + mixes_shared[j + hc_mult] * hc_scale[1] + hc_base[j + hc_mult]
252 + )
253 + * hc_post_mult_value
254 + )
255 + for j, k in T.Parallel(hc_mult, hc_mult):
256 + cm[j, k] = (
257 + mixes_shared[j * hc_mult + k + hc_mult * 2] * hc_scale[2]
258 + + hc_base[j * hc_mult + k + hc_mult * 2]
259 + )
260 +
261 + row_sum = T.alloc_fragment(hc_mult, T.float32)
262 + col_sum = T.alloc_fragment(hc_mult, T.float32)
263 +
264 + row_max = T.alloc_fragment(hc_mult, T.float32)
265 + T.reduce_max(cm, row_max, dim=1)
266 + for j, k in T.Parallel(hc_mult, hc_mult):
267 + cm[j, k] = T.exp(cm[j, k] - row_max[j])
268 + T.reduce_sum(cm, row_sum, dim=1)
269 + for j, k in T.Parallel(hc_mult, hc_mult):
270 + cm[j, k] = cm[j, k] / row_sum[j] + hc_sinkhorn_eps
271 +
272 + T.reduce_sum(cm, col_sum, dim=0)
273 + for j, k in T.Parallel(hc_mult, hc_mult):
274 + cm[j, k] = cm[j, k] / (col_sum[k] + hc_sinkhorn_eps)
275 +
276 + for _ in T.serial(sinkhorn_repeat - 1):
277 + T.reduce_sum(cm, row_sum, dim=1)
278 + for j, k in T.Parallel(hc_mult, hc_mult):
279 + cm[j, k] = cm[j, k] / (row_sum[j] + hc_sinkhorn_eps)
280 +
281 + T.reduce_sum(cm, col_sum, dim=0)
282 + for j, k in T.Parallel(hc_mult, hc_mult):
283 + cm[j, k] = cm[j, k] / (col_sum[k] + hc_sinkhorn_eps)
284 +
285 + for j, k in T.Parallel(hc_mult, hc_mult):
286 + comb_mix[i, j * hc_mult + k] = cm[j, k]
287 + else:
288 + pre_mix_shared = T.alloc_shared(hc_mult, T.float32)
289 + for j in T.Parallel(hc_mult):
290 + pre_mix_shared[j] = (
291 + T.sigmoid(
292 + mixes_shared[j] * hc_scale[0] + hc_base[j],
293 + )
294 + + hc_pre_eps
295 + )
296 +
297 + # Pass 1: stash unnormalized weighted-sum output in shared memory
298 + # as bf16 (matches the rounding that RMSNorm would see) while
299 + # accumulating the per-position squared sum.
300 + output_shared = T.alloc_shared(hidden_size, T.bfloat16)
301 + sumsq_per_pos = T.alloc_fragment(hidden_block, T.float32)
302 + T.clear(sumsq_per_pos)
303 +
304 + for i0_h in T.Pipelined(hidden_size // hidden_block, num_stages=3):
305 + xs = T.alloc_shared((hc_mult, hidden_block), T.bfloat16)
306 + xl = T.alloc_fragment((hc_mult, hidden_block), T.float32)
307 + T.copy(residual[i, 0, i0_h * hidden_block], xs)
308 + T.copy(xs, xl)
309 +
310 + ol = T.alloc_fragment(hidden_block, T.float32)
311 + T.clear(ol)
312 +
313 + for i_hc in T.serial(hc_mult):
314 + pre = pre_mix_shared[i_hc]
315 + for i1_h in T.Parallel(hidden_block):
316 + ol[i1_h] += pre * xl[i_hc, i1_h]
317 +
318 + for i1_h in T.Parallel(hidden_block):
319 + sumsq_per_pos[i1_h] += ol[i1_h] * ol[i1_h]
320 + output_shared[i0_h * hidden_block + i1_h] = T.bfloat16(ol[i1_h])
321 +
322 + sumsq = T.alloc_fragment(1, T.float32)
323 + T.reduce_sum(sumsq_per_pos, sumsq, dim=0)
324 + rsqrt_norm = T.alloc_fragment(1, T.float32)
325 + rsqrt_norm[0] = T.rsqrt(sumsq[0] / hidden_size + norm_eps)
326 +
327 + # Pass 2: scale by rsqrt * norm_weight and write the result to HBM.
328 + for i0_h in T.Pipelined(hidden_size // hidden_block, num_stages=2):
329 + w_shared = T.alloc_shared(hidden_block, T.bfloat16)
330 + w_local = T.alloc_fragment(hidden_block, T.float32)
331 + T.copy(norm_weight[i0_h * hidden_block], w_shared)
332 + T.copy(w_shared, w_local)
333 +
334 + ol = T.alloc_fragment(hidden_block, T.float32)
335 + for i1_h in T.Parallel(hidden_block):
336 + ol[i1_h] = (
337 + output_shared[i0_h * hidden_block + i1_h]
338 + * rsqrt_norm[0]
339 + * w_local[i1_h]
340 + )
341 +
342 + T.copy(ol, layer_input[i, i0_h * hidden_block])
343 +
344 + T.pdl_trigger()
345 +
346 +
347 @tilelang.jit(
348 pass_configs={
349 tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True,
vllm/model_executor/kernels/mhc/tilelang.py
+114 -36
@@ -16,6 +16,8 @@ def mhc_pre_tilelang(
16 hc_post_mult_value: float,
17 sinkhorn_repeat: int,
18 n_splits: int = 1,
19 + norm_weight: torch.Tensor | None = None,
20 + norm_eps: float = 1e-6,
21 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
22 """
23 Forward pass for mHC pre block.
@@ -31,6 +33,11 @@ def mhc_pre_tilelang(
33 hc_post_mult_value: post-mix multiplier value
34 sinkhorn_repeat: number of sinkhorn iterations
35 n_splits: split-k factor;
36 + norm_weight: optional RMSNorm weight, shape (hidden_size,), dtype
37 + torch.bfloat16. When provided, RMSNorm is fused into the
38 + layer_input write path of the big_fuse kernel.
39 + norm_eps: epsilon for the fused RMSNorm; only consulted when
40 + norm_weight is given.
41
42 Returns:
43 post_mix: shape (..., hc_mult), dtype torch.float32
@@ -40,6 +47,7 @@ def mhc_pre_tilelang(
47 from vllm._tilelang_ops import (
48 compute_num_split,
49 mhc_pre_big_fuse_tilelang,
50 + mhc_pre_big_fuse_with_norm_tilelang,
51 )
52 from vllm.utils.deep_gemm import tf32_hc_prenorm_gemm
53 from vllm.utils.math_utils import cdiv
@@ -60,6 +68,13 @@ def mhc_pre_tilelang(
68 assert hc_scale.shape == (3,)
69 assert hc_base.shape == (hc_mult3,)
70
71 + if norm_weight is not None:
72 + assert norm_weight.shape == (hidden_size,)
73 + if norm_weight.dtype != torch.bfloat16:
74 + norm_weight = norm_weight.to(torch.bfloat16)
75 + if not norm_weight.is_contiguous():
76 + norm_weight = norm_weight.contiguous()
77 +
78 outer_shape = residual.shape[:-2]
79
80 residual_flat = residual.view(-1, hc_mult, hidden_size)
@@ -95,24 +110,46 @@ def mhc_pre_tilelang(
110 n_splits,
111 )
112
98 - mhc_pre_big_fuse_tilelang(
99 - gemm_out_mul,
100 - gemm_out_sqrsum,
101 - hc_scale,
102 - hc_base,
103 - residual_flat,
104 - post_mix,
105 - comb_mix,
106 - layer_input,
107 - hidden_size,
108 - rms_eps,
109 - hc_pre_eps,
110 - hc_sinkhorn_eps,
111 - hc_post_mult_value,
112 - sinkhorn_repeat,
113 - n_splits,
114 - hc_mult,
115 - )
113 + if norm_weight is None:
114 + mhc_pre_big_fuse_tilelang(
115 + gemm_out_mul,
116 + gemm_out_sqrsum,
117 + hc_scale,
118 + hc_base,
119 + residual_flat,
120 + post_mix,
121 + comb_mix,
122 + layer_input,
123 + hidden_size,
124 + rms_eps,
125 + hc_pre_eps,
126 + hc_sinkhorn_eps,
127 + hc_post_mult_value,
128 + sinkhorn_repeat,
129 + n_splits,
130 + hc_mult,
131 + )
132 + else:
133 + mhc_pre_big_fuse_with_norm_tilelang(
134 + gemm_out_mul,
135 + gemm_out_sqrsum,
136 + hc_scale,
137 + hc_base,
138 + residual_flat,
139 + post_mix,
140 + comb_mix,
141 + layer_input,
142 + norm_weight,
143 + hidden_size,
144 + rms_eps,
145 + hc_pre_eps,
146 + hc_sinkhorn_eps,
147 + hc_post_mult_value,
148 + sinkhorn_repeat,
149 + norm_eps,
150 + n_splits,
151 + hc_mult,
152 + )
153
154 return (
155 post_mix.view(*outer_shape, hc_mult, 1),
@@ -132,6 +169,8 @@ def _mhc_pre_tilelang_fake(
169 hc_post_mult_value: float,
170 sinkhorn_repeat: int,
171 n_splits: int = 1,
172 + norm_weight: torch.Tensor | None = None,
173 + norm_eps: float = 0.0,
174 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
175 hc_mult = residual.shape[-2]
176 hidden_size = residual.shape[-1]
@@ -198,10 +237,16 @@ def mhc_fused_post_pre_tilelang(
237 sinkhorn_repeat: int,
238 n_splits: int = 1,
239 tile_n: int = 1,
240 + norm_weight: torch.Tensor | None = None,
241 + norm_eps: float = 0.0,
242 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
243 """
244 Run one MHC post block followed by the next MHC pre block.
245
246 + When ``norm_weight`` is provided, the layer_input_cur output is the
247 + RMSNorm'd activation (fused into the kernel); otherwise it is the
248 + raw pre-norm activation as before.
249 +
250 Returns:
251 residual_cur: post-mapped residual, shape (..., hc_mult, hidden_size)
252 post_mix_cur: shape (..., hc_mult, 1)
@@ -214,6 +259,7 @@ def mhc_fused_post_pre_tilelang(
259 mhc_fused_tilelang,
260 mhc_post_tilelang,
261 mhc_pre_big_fuse_tilelang,
262 + mhc_pre_big_fuse_with_norm_tilelang,
263 )
264 from vllm.utils.math_utils import cdiv
265
@@ -242,6 +288,13 @@ def mhc_fused_post_pre_tilelang(
288 assert hc_scale.shape == (3,)
289 assert hc_base.shape == (hc_mult3,)
290
291 + if norm_weight is not None:
292 + assert norm_weight.shape == (hidden_size,)
293 + if norm_weight.dtype != torch.bfloat16:
294 + norm_weight = norm_weight.to(torch.bfloat16)
295 + if not norm_weight.is_contiguous():
296 + norm_weight = norm_weight.contiguous()
297 +
298 assert n_splits in (1, 2, 4, 8)
299 assert hidden_size % n_splits == 0
300
@@ -332,24 +385,46 @@ def mhc_fused_post_pre_tilelang(
385 n_splits,
386 )
387
335 - mhc_pre_big_fuse_tilelang(
336 - gemm_out_mul,
337 - gemm_out_sqrsum,
338 - hc_scale,
339 - hc_base,
340 - residual_cur,
341 - post_mix_cur,
342 - comb_mix_cur,
343 - layer_input_cur,
344 - hidden_size,
345 - rms_eps,
346 - hc_pre_eps,
347 - hc_sinkhorn_eps,
348 - hc_post_mult_value,
349 - sinkhorn_repeat,
350 - n_splits,
351 - hc_mult,
352 - )
388 + if norm_weight is None:
389 + mhc_pre_big_fuse_tilelang(
390 + gemm_out_mul,
391 + gemm_out_sqrsum,
392 + hc_scale,
393 + hc_base,
394 + residual_cur,
395 + post_mix_cur,
396 + comb_mix_cur,
397 + layer_input_cur,
398 + hidden_size,
399 + rms_eps,
400 + hc_pre_eps,
401 + hc_sinkhorn_eps,
402 + hc_post_mult_value,
403 + sinkhorn_repeat,
404 + n_splits,
405 + hc_mult,
406 + )
407 + else:
408 + mhc_pre_big_fuse_with_norm_tilelang(
409 + gemm_out_mul,
410 + gemm_out_sqrsum,
411 + hc_scale,
412 + hc_base,
413 + residual_cur,
414 + post_mix_cur,
415 + comb_mix_cur,
416 + layer_input_cur,
417 + norm_weight,
418 + hidden_size,
419 + rms_eps,
420 + hc_pre_eps,
421 + hc_sinkhorn_eps,
422 + hc_post_mult_value,
423 + sinkhorn_repeat,
424 + norm_eps,
425 + n_splits,
426 + hc_mult,
427 + )
428
429 return (
430 residual_cur.view(*outer_shape, hc_mult, hidden_size),
@@ -373,6 +448,9 @@ def _mhc_fused_post_pre_tilelang_fake(
448 hc_post_mult_value: float,
449 sinkhorn_repeat: int,
450 n_splits: int = 1,
451 + tile_n: int = 1,
452 + norm_weight: torch.Tensor | None = None,
453 + norm_eps: float = 0.0,
454 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
455 hc_mult = residual.shape[-2]
456 hidden_size = residual.shape[-1]
vllm/model_executor/layers/mhc.py
+10
@@ -34,6 +34,8 @@ class MHCPreOp(CustomOp):
34 hc_post_mult_value: float,
35 sinkhorn_repeat: int,
36 n_splits: int = 1,
37 + norm_weight: torch.Tensor | None = None,
38 + norm_eps: float = 0.0,
39 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
40 return torch.ops.vllm.mhc_pre_tilelang(
41 residual,
@@ -46,6 +48,8 @@ class MHCPreOp(CustomOp):
48 hc_post_mult_value,
49 sinkhorn_repeat,
50 n_splits,
51 + norm_weight,
52 + norm_eps,
53 )
54
55 def forward_hip(
@@ -60,6 +64,8 @@ class MHCPreOp(CustomOp):
64 hc_post_mult_value: float,
65 sinkhorn_repeat: int,
66 n_splits: int = 1,
67 + norm_weight: torch.Tensor | None = None,
68 + norm_eps: float = 0.0,
69 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
70 # TODO: Reenable aiter after we are at the aiter
71 # version that has this bugfix
@@ -262,6 +268,8 @@ class MHCFusedPostPreOp(CustomOp):
268 sinkhorn_repeat: int,
269 n_splits: int = 1,
270 tile_n: int = 1,
271 + norm_weight: torch.Tensor | None = None,
272 + norm_eps: float = 0.0,
273 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
274 return torch.ops.vllm.mhc_fused_post_pre_tilelang(
275 x,
@@ -278,6 +286,8 @@ class MHCFusedPostPreOp(CustomOp):
286 sinkhorn_repeat,
287 n_splits,
288 tile_n,
289 + norm_weight,
290 + norm_eps,
291 )
292
293 def forward_hip(self, *args, **kwargs):
vllm/models/deepseek_v4/nvidia/model.py
+66 -47
@@ -22,9 +22,7 @@ from vllm.model_executor.layers.fused_moe import FusedMoE
22 from vllm.model_executor.layers.fused_moe.router.fused_topk_bias_router import (
23 fused_topk_bias,
24 )
25 -from vllm.model_executor.layers.fused_moe.router.norm_gate_linear import (
26 - NormGateLinear,
27 -)
25 +from vllm.model_executor.layers.fused_moe.router.gate_linear import GateLinear
26 from vllm.model_executor.layers.layernorm import RMSNorm
27 from vllm.model_executor.layers.linear import (
28 ColumnParallelLinear,
@@ -655,23 +653,23 @@ class DeepseekV4MoE(nn.Module):
653 "deep_gemm_mega_moe for this checkpoint."
654 )
655
658 - # Fused RMSNorm + gate: owns both ffn_norm and the gate matmul.
659 - self.norm_gate = NormGateLinear(
660 - hidden_size=config.hidden_size,
661 - num_experts=config.n_routed_experts,
662 - rms_eps=config.rms_norm_eps,
663 - prefix=f"{prefix}.norm_gate",
656 + self.gate = GateLinear(
657 + input_size=config.hidden_size,
658 + output_size=config.n_routed_experts,
659 + bias=False,
660 + out_dtype=torch.float32,
661 + prefix=f"{prefix}.gate",
662 )
665 - # Routing-side tensors live on ``norm_gate`` directly (not on the
666 - # inner gate); they are initialized to None in NormGatedLinear and
667 - # populated below depending on the MoE variant.
663 +
664 + self.gate.e_score_correction_bias = None
665 + self.gate.tid2eid = None
666 is_hash_moe = extract_layer_index(prefix) < config.num_hash_layers
667 self.hash_indices_dtype = torch.int64 if self.use_mega_moe else torch.int32
668 if is_hash_moe:
669 # hash MoE doesn't use e_score_correction_bias
670 # Use randint instead of empty to avoid garbage values causing
671 # invalid memory access in dummy mode (--load-format="dummy")
674 - self.norm_gate.tid2eid = nn.Parameter(
672 + self.gate.tid2eid = nn.Parameter(
673 torch.randint(
674 0,
675 config.n_routed_experts,
@@ -681,7 +679,7 @@ class DeepseekV4MoE(nn.Module):
679 requires_grad=False,
680 )
681 elif getattr(config, "topk_method", None) == "noaux_tc":
684 - self.norm_gate.e_score_correction_bias = nn.Parameter(
682 + self.gate.e_score_correction_bias = nn.Parameter(
683 torch.empty(config.n_routed_experts, dtype=torch.float32),
684 requires_grad=False,
685 )
@@ -744,9 +742,10 @@ class DeepseekV4MoE(nn.Module):
742 self.n_local_experts = config.n_routed_experts // self.tp_size
743 self.experts_start_idx = self.tp_rank * self.n_local_experts
744 self.experts_end_idx = self.experts_start_idx + self.n_local_experts
747 - # We don't pass `gate` into FusedMoE
745 +
746 self.experts = FusedMoE(
747 shared_experts=self.shared_experts,
748 + gate=self.gate,
749 num_experts=config.n_routed_experts,
750 top_k=config.num_experts_per_tok,
751 hidden_size=config.hidden_size,
@@ -756,8 +755,8 @@ class DeepseekV4MoE(nn.Module):
755 prefix=f"{prefix}.experts",
756 scoring_func=self.scoring_func,
757 routed_scaling_factor=self.routed_scaling_factor,
759 - e_score_correction_bias=self.norm_gate.e_score_correction_bias,
760 - hash_indices_table=self.norm_gate.tid2eid,
758 + e_score_correction_bias=self.gate.e_score_correction_bias,
759 + hash_indices_table=self.gate.tid2eid,
760 swiglu_limit=self.swiglu_limit,
761 router_logits_dtype=torch.float32,
762 )
@@ -765,40 +764,40 @@ class DeepseekV4MoE(nn.Module):
764 def forward(
765 self, hidden_states: torch.Tensor, input_ids: torch.Tensor | None = None
766 ) -> torch.Tensor:
768 - if self.norm_gate.tid2eid is not None and input_ids is None:
767 + if self.gate.tid2eid is not None and input_ids is None:
768 raise ValueError("DeepSeek V4 hash MoE routing requires input_ids.")
769
770 if not self.use_mega_moe:
771 return self._forward_fused_moe(hidden_states, input_ids)
772
773 org_shape = hidden_states.shape
775 - normed_x, router_logits = self.norm_gate(hidden_states)
774 + router_logits, _ = self.gate(hidden_states)
775 topk_weights, topk_ids = fused_topk_bias(
777 - hidden_states=normed_x,
776 + hidden_states=hidden_states,
777 gating_output=router_logits,
778 scoring_func=self.scoring_func,
780 - e_score_correction_bias=self.norm_gate.e_score_correction_bias.data
781 - if self.norm_gate.e_score_correction_bias is not None
779 + e_score_correction_bias=self.gate.e_score_correction_bias.data
780 + if self.gate.e_score_correction_bias is not None
781 else None,
782 topk=self.n_activated_experts,
783 renormalize=self.renormalize,
784 indices_type=self.hash_indices_dtype,
785 input_tokens=input_ids,
787 - hash_indices_table=self.norm_gate.tid2eid,
786 + hash_indices_table=self.gate.tid2eid,
787 routed_scaling_factor=self.routed_scaling_factor,
788 )
789 activation_clamp = (
790 float(self.swiglu_limit) if self.swiglu_limit is not None else None
791 )
792 final_hidden_states = self.experts(
794 - normed_x,
793 + hidden_states,
794 topk_weights,
795 topk_ids,
796 activation_clamp=activation_clamp,
797 )
798
799 if self.shared_experts is not None:
801 - shared_output = self.shared_experts(normed_x)
800 + shared_output = self.shared_experts(hidden_states)
801 final_hidden_states += shared_output
802
803 return final_hidden_states.view(org_shape)
@@ -806,14 +805,21 @@ class DeepseekV4MoE(nn.Module):
805 def _forward_fused_moe(
806 self, hidden_states: torch.Tensor, input_ids: torch.Tensor | None = None
807 ) -> torch.Tensor:
809 - assert not self.experts.is_internal_router
808 org_shape = hidden_states.shape
811 - normed_x, router_logits = self.norm_gate(hidden_states)
812 - final_hidden_states = self.experts(
813 - hidden_states=normed_x,
814 - router_logits=router_logits,
815 - input_ids=input_ids,
816 - )
809 + if self.experts.is_internal_router:
810 + # In this case, the gate/router runs inside the FusedMoE class
811 + final_hidden_states = self.experts(
812 + hidden_states=hidden_states,
813 + router_logits=hidden_states,
814 + input_ids=input_ids,
815 + )
816 + else:
817 + router_logits, _ = self.gate(hidden_states)
818 + final_hidden_states = self.experts(
819 + hidden_states=hidden_states,
820 + router_logits=router_logits,
821 + input_ids=input_ids,
822 + )
823
824 return final_hidden_states.view(org_shape)
825
@@ -1024,8 +1030,7 @@ class DeepseekV4DecoderLayer(nn.Module):
1030 self.ffn = DeepseekV4MoE(vllm_config, prefix=f"{prefix}.ffn")
1031
1032 self.attn_norm = RMSNorm(self.hidden_size, self.rms_norm_eps)
1027 - # ``ffn_norm`` is owned by ``self.ffn.norm_gate`` (fused with the
1028 - # router gate matmul); see ``NormGatedLinear``.
1033 + self.ffn_norm = RMSNorm(self.hidden_size, self.rms_norm_eps)
1034 self.hc_mult = config.hc_mult
1035 self.hc_sinkhorn_iters = config.hc_sinkhorn_iters
1036 self.hc_eps = config.hc_eps
@@ -1084,6 +1089,8 @@ class DeepseekV4DecoderLayer(nn.Module):
1089 hc_fn: torch.Tensor,
1090 hc_scale: torch.Tensor,
1091 hc_base: torch.Tensor,
1092 + norm_weight: torch.Tensor | None = None,
1093 + norm_eps: float = 1e-6,
1094 ):
1095 post_mix, res_mix, layer_input = self.mhc_pre(
1096 residual=x,
@@ -1095,6 +1102,8 @@ class DeepseekV4DecoderLayer(nn.Module):
1102 hc_sinkhorn_eps=self.hc_eps,
1103 hc_post_mult_value=self.hc_post_alpha,
1104 sinkhorn_repeat=self.hc_sinkhorn_iters,
1105 + norm_weight=norm_weight,
1106 + norm_eps=norm_eps,
1107 )
1108 return layer_input, post_mix, res_mix
1109
@@ -1116,11 +1125,18 @@ class DeepseekV4DecoderLayer(nn.Module):
1125 res_mix: torch.Tensor | None = None,
1126 residual: torch.Tensor | None = None,
1127 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
1128 + attn_norm_weight = self.attn_norm.weight.data
1129 + attn_norm_eps = self.attn_norm.variance_epsilon
1130 if residual is None:
1131 # Run standalone hc_pre on first layer
1132 residual = x
1133 x, post_mix, res_mix = self.hc_pre(
1123 - x, self.hc_attn_fn, self.hc_attn_scale, self.hc_attn_base
1134 + x,
1135 + self.hc_attn_fn,
1136 + self.hc_attn_scale,
1137 + self.hc_attn_base,
1138 + norm_weight=attn_norm_weight,
1139 + norm_eps=attn_norm_eps,
1140 )
1141 else:
1142 residual, post_mix, res_mix, x = self.mhc_fused_post_pre(
@@ -1136,11 +1152,17 @@ class DeepseekV4DecoderLayer(nn.Module):
1152 self.hc_eps,
1153 self.hc_post_alpha,
1154 self.hc_sinkhorn_iters,
1155 + n_splits=1,
1156 + tile_n=1,
1157 + norm_weight=attn_norm_weight,
1158 + norm_eps=attn_norm_eps,
1159 )
1160
1141 - x = self.attn_norm(x)
1161 + # attn_norm is fused into hc_pre / mhc_fused_post_pre above.
1162 x = self.attn(positions, x, None)
1163
1164 + ffn_norm_weight = self.ffn_norm.weight.data
1165 + ffn_norm_eps = self.ffn_norm.variance_epsilon
1166 residual, post_mix, res_mix, x = self.mhc_fused_post_pre(
1167 x,
1168 residual,
@@ -1154,9 +1176,13 @@ class DeepseekV4DecoderLayer(nn.Module):
1176 self.hc_eps,
1177 self.hc_post_alpha,
1178 self.hc_sinkhorn_iters,
1179 + n_splits=1,
1180 + tile_n=1,
1181 + norm_weight=ffn_norm_weight,
1182 + norm_eps=ffn_norm_eps,
1183 )
1158 - # ffn_norm is now folded into self.ffn.norm_gate; ffn() takes
1159 - # the pre-norm activation directly.
1184 + # ffn_norm is fused into mhc_fused_post_pre above; ffn() takes the
1185 + # already-normed activation directly.
1186 x = self.ffn(x, input_ids)
1187 return x, residual, post_mix, res_mix
1188
@@ -1183,8 +1209,7 @@ class DeepseekV4DecoderLayer(nn.Module):
1209 x, post, comb = self.hc_pre(
1210 x, self.hc_ffn_fn, self.hc_ffn_scale, self.hc_ffn_base
1211 )
1186 - # ffn_norm is now folded into self.ffn.norm_gate; ffn() takes
1187 - # the pre-norm activation directly.
1212 + x = self.ffn_norm(x)
1213 x = self.ffn(x, input_ids)
1214 x = self.hc_post(x, residual, post, comb)
1215 return x, None, None, None
@@ -1534,13 +1559,7 @@ def _make_deepseek_v4_weights_mapper(expert_dtype: str) -> WeightsMapper:
1559 orig_to_new_suffix={
1560 "head.weight": "lm_head.weight",
1561 "embed.weight": "embed_tokens.weight",
1537 - # Pre-MoE norm + gate are now owned by ``DeepseekV4MoE.norm_gate``
1538 - # (see NormGatedLinear).
1539 - ".ffn_norm.weight": ".ffn.norm_gate.norm.weight",
1540 - ".ffn.gate.weight": ".ffn.norm_gate.gate.weight",
1541 - ".ffn.gate.bias": ".ffn.norm_gate.e_score_correction_bias",
1542 - # Hash MoE table also moved off the inner gate.
1543 - ".ffn.gate.tid2eid": ".ffn.norm_gate.tid2eid",
1562 + ".ffn.gate.bias": ".ffn.gate.e_score_correction_bias",
1563 },
1564 orig_to_new_substr={
1565 ".attn.compressor.": ".attn.mla_attn.compressor.",
vllm/models/deepseek_v4/nvidia/mtp.py
+3 -8
@@ -292,11 +292,6 @@ class DeepSeekV4MTP(nn.Module):
292 ".emb.tok_emb.weight": ".embed_tokens.weight",
293 ".head.weight": ".shared_head.head.weight",
294 ".norm.weight": ".shared_head.norm.weight",
295 - # Pre-MoE norm + gate are now owned by
296 - # ``DeepseekV4MoE.norm_gate`` (see NormGatedLinear).
297 - ".ffn_norm.weight": ".ffn.norm_gate.norm.weight",
298 - ".ffn.gate.weight": ".ffn.norm_gate.gate.weight",
299 - ".ffn.gate.tid2eid": ".ffn.norm_gate.tid2eid",
295 }
296
297 def _remap_weight_name(name: str) -> str:
@@ -444,11 +439,11 @@ class DeepSeekV4MTP(nn.Module):
439 ".shared_experts.w2", ".shared_experts.down_proj"
440 )
441 if name.endswith(".ffn.gate.bias"):
447 - # ``e_score_correction_bias`` lives on
448 - # ``norm_gate`` directly (not on the inner gate).
442 + # ``e_score_correction_bias`` lives on the gate
443 + # under a different attribute name.
444 name = name.replace(
445 ".ffn.gate.bias",
451 - ".ffn.norm_gate.e_score_correction_bias",
446 + ".ffn.gate.e_score_correction_bias",
447 )
448 param = params_dict[name]
449 weight_loader = getattr(