@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 5372e72a98

[Perf][DSv4.1] Fold the mHC post block into the delayed pre projection (#56633)

Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>

Yongye Zhu committed Sep 14, 2026 at 11:57 UTC 5372e72a9884c141874465f0eb8054a4a81da6bb
8 files changed +837 -92
tests/distributed/test_engram_dp_shard.py
+5 -2
@@ -491,7 +491,9 @@ def _check_dummy_hash_model_forward(
491 state.forward = lambda *args: ids.unsqueeze(1)
492
493 class Decoder(SimpleNamespace):
494 - def __call__(self, hidden, positions, input_ids, *args):
494 + # Engram hashes and the mask stay the last positional arguments;
495 + # **kwargs absorbs the decoder's keyword-only flags.
496 + def __call__(self, hidden, positions, input_ids, *args, **kwargs):
497 hashes, keep = args[-2:]
498 if dp_shared_memory and dp_rank == 1:
499 assert hashes is None and keep is None
@@ -504,7 +506,8 @@ def _check_dummy_hash_model_forward(
506 if dp_rank == 1:
507 assert torch.all(hashes == engram_ops.DEAD_ID)
508 output = engram.embed(hashes[:, 0])
507 - return output, None, None, None, None
509 + # Trailing None is previous_aux; this stub captures no aux states.
510 + return output, None, None, None, None, None
511
512 model = SimpleNamespace(
513 use_mega_moe=False,
tests/kernels/test_mhc_jit_warmup.py
+38 -9
@@ -17,6 +17,8 @@ from vllm.model_executor.kernels.mhc.tilelang_kernels import (
17 MhcFusedTileLangKernel,
18 MhcPostTileLangKernel,
19 MhcPreBigFuseTileLangKernel,
20 + mhc_fused_post_pre_split_config,
21 + require_fused_post_pre_config,
22 )
23 from vllm.model_executor.warmup import jit_warmup_tilelang_helper
24
@@ -113,17 +115,32 @@ def test_mhc_pre_big_fuse_dispatch_matches_legacy_runtime_config(
115
116
117 @pytest.mark.parametrize(
116 - ("num_tokens", "hidden_size", "expected_n_splits", "expected_tile_n"),
117 - [(4, 4096, 8, 2), (4, 8192, 4, 2), (8, 4096, 4, 3)],
118 + ("num_tokens", "hidden_size", "expected_tile_n", "expected_n_splits"),
119 + [
120 + (1, 4096, 2, 8),
121 + (4, 8192, 2, 8),
122 + (8, 5120, 2, 8),
123 + (16, 5120, 6, 8),
124 + (32, 7168, 6, 8),
125 + ],
126 )
119 -def test_mhc_fused_dispatch_matches_legacy_runtime_config(
127 +def test_mhc_fused_dispatch_matches_the_launch_config(
128 num_tokens: int,
129 hidden_size: int,
122 - expected_n_splits: int,
130 expected_tile_n: int,
131 + expected_n_splits: int,
132 ) -> None:
133 + """The compile key must be the config the launch will actually use.
134 +
135 + Both read mhc_fused_post_pre_split_config, so this pins the tuned bands
136 + and guards against the two drifting apart again.
137 + """
138 kernel = MhcFusedTileLangKernel()
139 + tile_n, n_splits, n_thr = require_fused_post_pre_config(
140 + num_tokens, hidden_size, hc_mult=4
141 + )
142
143 + assert (tile_n, n_splits) == (expected_tile_n, expected_n_splits)
144 assert kernel.dispatch(
145 num_tokens=num_tokens,
146 hidden_size=hidden_size,
@@ -131,11 +148,20 @@ def test_mhc_fused_dispatch_matches_legacy_runtime_config(
148 ) == kernel.CompileKey(
149 hidden_size=hidden_size,
150 hc_mult=4,
134 - n_splits=expected_n_splits,
135 - tile_n=expected_tile_n,
151 + n_splits=n_splits,
152 + tile_n=tile_n,
153 + n_thr=n_thr,
154 )
155
156
157 +def test_mhc_fused_declines_shapes_it_cannot_tile() -> None:
158 + """Above the token cutoff, and for a hidden size the block cannot split."""
159 + assert mhc_fused_post_pre_split_config(33, 5120, 4) is None
160 + assert mhc_fused_post_pre_split_config(1, 5137, 4) is None
161 + with pytest.raises(ValueError, match="does not cover num_tokens"):
162 + require_fused_post_pre_config(33, 5120, 4)
163 +
164 +
165 @pytest.mark.parametrize(
166 ("kernel", "compile_key"),
167 [
@@ -235,8 +261,9 @@ def test_mhc_fused_dispatch_matches_legacy_runtime_config(
261 MhcFusedTileLangKernel.CompileKey(
262 hidden_size=4096,
263 hc_mult=4,
238 - n_splits=4,
239 - tile_n=3,
264 + n_splits=8,
265 + tile_n=6,
266 + n_thr=128,
267 ),
268 ),
269 (
@@ -246,6 +273,7 @@ def test_mhc_fused_dispatch_matches_legacy_runtime_config(
273 hc_mult=4,
274 n_splits=8,
275 tile_n=2,
276 + n_thr=128,
277 ),
278 ),
279 (
@@ -253,8 +281,9 @@ def test_mhc_fused_dispatch_matches_legacy_runtime_config(
281 MhcFusedTileLangKernel.CompileKey(
282 hidden_size=8192,
283 hc_mult=4,
256 - n_splits=4,
284 + n_splits=8,
285 tile_n=2,
286 + n_thr=128,
287 ),
288 ),
289 (
tests/kernels/test_mhc_kernels.py
+244 -4
@@ -15,8 +15,12 @@ from vllm.model_executor.kernels.mhc.tilelang import (
15 )
16 from vllm.model_executor.kernels.mhc.tilelang_kernels import (
17 _HC_PRENORM_GEMM_TILELANG_KERNEL,
18 + mhc_fused_post_pre_split_config,
19 +)
20 +from vllm.model_executor.kernels.mhc.torch import (
21 + mhc_post_torch,
22 + mhc_pre_delayed_torch,
23 )
19 -from vllm.model_executor.kernels.mhc.torch import mhc_pre_delayed_torch
24 from vllm.model_executor.kernels.mhc.triton import hc_collapse_triton
25 from vllm.model_executor.layers.mhc import (
26 HAS_AITER_MHC,
@@ -115,9 +119,9 @@ def test_v41_dspark_head_collapses_with_last_ffn_mix(num_tokens, monkeypatch):
119 carried_mixes = []
120
121 def make_layer(mix):
118 - def forward(hidden, positions, ids, pre, post, res, residual):
122 + def forward(hidden, positions, ids, pre, post, res, residual, **kwargs):
123 carried_mixes.append(pre)
120 - return hidden, None, None, None, mix
124 + return hidden, None, None, None, mix, None
125
126 return forward
127
@@ -235,6 +239,171 @@ def test_deepseek_v41_mhc_pre_delayed(
239 torch.testing.assert_close(actual[2], expected[2], atol=1.6e-2, rtol=1e-2)
240
241
242 +def mhc_fused_post_pre_delayed_ref(
243 + x: torch.Tensor,
244 + residual: torch.Tensor,
245 + post_layer_mix: torch.Tensor,
246 + comb_res_mix: torch.Tensor,
247 + *mix_args,
248 + pre_mix: torch.Tensor | None,
249 +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
250 + """Project the post result before it is rounded to BF16.
251 +
252 + The fused kernel keeps the updated residual streams in registers, so its
253 + projection sees the FP32 post result while the streams it stores are the
254 + BF16 rounding of it.
255 + """
256 + post = torch.einsum("tij,tih->tjh", comb_res_mix, residual.float())
257 + post = post + post_layer_mix * x.float().unsqueeze(-2)
258 + return mhc_pre_delayed_torch(
259 + post.bfloat16(), *mix_args, pre_mix=pre_mix, x=post.flatten(1)
260 + )
261 +
262 +
263 +@pytest.mark.skipif(not HAS_TILELANG_MHC, reason="TileLang MHC support required")
264 +@pytest.mark.parametrize("num_tokens", [0, 1, 8, 128])
265 +@pytest.mark.parametrize("hidden_size", [4096, 7168])
266 +@pytest.mark.parametrize("carried", [False, True])
267 +def test_deepseek_v41_mhc_fused_post_pre_delayed(num_tokens, hidden_size, carried):
268 + """Fold the post block into the delayed pre without changing its outputs."""
269 + set_random_seed(0)
270 + hc_mult = 4
271 + mix_size = hc_mult * (hc_mult + 2)
272 + x = torch.randn(num_tokens, hidden_size, dtype=torch.bfloat16, device=DEVICE)
273 + residual = torch.randn(
274 + num_tokens, hc_mult, hidden_size, dtype=torch.bfloat16, device=DEVICE
275 + )
276 + post_layer_mix = torch.rand(num_tokens, hc_mult, 1, device=DEVICE)
277 + comb_res_mix = torch.rand(num_tokens, hc_mult, hc_mult, device=DEVICE)
278 + fn = torch.randn(mix_size, hc_mult * hidden_size, device=DEVICE) * 0.02
279 + scale = torch.tensor([0.5, 0.25, 1.0], device=DEVICE)
280 + base = torch.randn(mix_size, device=DEVICE)
281 + pre_mix = torch.rand(num_tokens, hc_mult, device=DEVICE) if carried else None
282 + weight = torch.empty(hidden_size, dtype=torch.bfloat16, device=DEVICE)
283 + weight.uniform_(0.5, 1.5)
284 + mix_args = (fn, scale, base, 1e-20, 1e-6, 1e-6, 2.0, 20)
285 +
286 + residual_cur, post, comb, layer_input, next_pre, aux = (
287 + torch.ops.vllm.mhc_fused_post_pre_delayed_tilelang(
288 + x, residual, post_layer_mix, comb_res_mix, *mix_args, pre_mix, weight, 1e-6
289 + )
290 + )
291 + assert aux.shape == (0, hidden_size)
292 + assert [
293 + residual_cur.shape,
294 + post.shape,
295 + comb.shape,
296 + layer_input.shape,
297 + next_pre.shape,
298 + ] == [
299 + (num_tokens, hc_mult, hidden_size),
300 + (num_tokens, hc_mult, 1),
301 + (num_tokens, hc_mult, hc_mult),
302 + (num_tokens, hidden_size),
303 + (num_tokens, hc_mult),
304 + ]
305 + if num_tokens == 0:
306 + return
307 +
308 + # The post mapping and the collapse are what the unfused kernels produce.
309 + residual_ref = torch.ops.vllm.mhc_post_tilelang(
310 + x, residual, post_layer_mix, comb_res_mix
311 + )
312 + layer_input_ref = mhc_pre_delayed_tilelang(
313 + residual_ref, *mix_args, pre_mix=pre_mix, norm_weight=weight, norm_eps=1e-6
314 + )[2]
315 + torch.testing.assert_close(residual_cur, residual_ref, atol=0, rtol=0)
316 + torch.testing.assert_close(layer_input, layer_input_ref, atol=0, rtol=0)
317 +
318 + post_ref, comb_ref, _, next_pre_ref = mhc_fused_post_pre_delayed_ref(
319 + x, residual, post_layer_mix, comb_res_mix, *mix_args, pre_mix=pre_mix
320 + )
321 + # Above the fused kernel's token count the projection falls back to a
322 + # split-k TF32 GEMM over the stored BF16 streams.
323 + tol = (
324 + 1e-5
325 + if mhc_fused_post_pre_split_config(num_tokens, hidden_size, hc_mult) is not None
326 + else 5e-3
327 + )
328 + for actual, expected in (
329 + (post, post_ref),
330 + (comb, comb_ref),
331 + (next_pre, next_pre_ref),
332 + ):
333 + torch.testing.assert_close(actual, expected, atol=tol, rtol=tol)
334 +
335 + # Folding the stream mean into the collapse must not change it, or anything
336 + # else the op returns.
337 + captured = torch.ops.vllm.mhc_fused_post_pre_delayed_tilelang(
338 + x,
339 + residual,
340 + post_layer_mix,
341 + comb_res_mix,
342 + *mix_args,
343 + pre_mix,
344 + weight,
345 + 1e-6,
346 + True,
347 + )
348 + torch.testing.assert_close(captured[5], residual_cur.mean(dim=1), atol=0, rtol=0)
349 + unchanged = (residual_cur, post, comb, layer_input, next_pre)
350 + for with_aux, without in zip(captured[:5], unchanged, strict=True):
351 + torch.testing.assert_close(with_aux, without, atol=0, rtol=0)
352 +
353 + # The unnormalized epilogue has its own fused aux path.
354 + bare = torch.ops.vllm.mhc_fused_post_pre_delayed_tilelang(
355 + x, residual, post_layer_mix, comb_res_mix, *mix_args, pre_mix, None, 1e-6, True
356 + )
357 + torch.testing.assert_close(bare[5], bare[0].mean(dim=1), atol=0, rtol=0)
358 +
359 + # Decode replays this op from a captured graph.
360 + graph = torch.cuda.CUDAGraph()
361 + with torch.cuda.graph(graph):
362 + captured = torch.ops.vllm.mhc_fused_post_pre_delayed_tilelang(
363 + x, residual, post_layer_mix, comb_res_mix, *mix_args, pre_mix, weight, 1e-6
364 + )
365 + graph.replay()
366 + eager = (residual_cur, post, comb, layer_input, next_pre, aux)
367 + for replayed, expected in zip(captured, eager, strict=True):
368 + torch.testing.assert_close(replayed, expected, atol=0, rtol=0)
369 +
370 +
371 +@pytest.mark.skipif(not HAS_TILELANG_MHC, reason="TileLang MHC support required")
372 +@pytest.mark.parametrize("carried", [False, True])
373 +@pytest.mark.parametrize("capture_aux", [False, True])
374 +def test_mhc_fused_post_pre_delayed_custom_op_supports_compile(carried, capture_aux):
375 + set_random_seed(0)
376 + x = torch.randn(2, 5120, dtype=torch.bfloat16, device=DEVICE)
377 + residual = torch.randn(2, 4, 5120, dtype=torch.bfloat16, device=DEVICE)
378 + post_layer_mix = torch.rand(2, 4, 1, device=DEVICE)
379 + comb_res_mix = torch.rand(2, 4, 4, device=DEVICE)
380 + fn = torch.randn(24, 20480, device=DEVICE) * 0.02
381 + pre_mix = torch.rand(2, 4, device=DEVICE) if carried else None
382 + scale = torch.ones(3, device=DEVICE)
383 + base = torch.zeros(24, device=DEVICE)
384 + torch.library.opcheck(
385 + torch.ops.vllm.mhc_fused_post_pre_delayed_tilelang.default,
386 + (
387 + x,
388 + residual,
389 + post_layer_mix,
390 + comb_res_mix,
391 + fn,
392 + scale,
393 + base,
394 + 1e-20,
395 + 1e-6,
396 + 1e-6,
397 + 2.0,
398 + 20,
399 + pre_mix,
400 + None,
401 + 1e-6,
402 + capture_aux,
403 + ),
404 + )
405 +
406 +
407 @pytest.mark.skipif(not HAS_TILELANG_MHC, reason="TileLang MHC support required")
408 @pytest.mark.parametrize("entry", ["broadcast", "pipeline", "residual", "engram"])
409 def test_deepseek_v41_decoder_mixes_match_torch(
@@ -295,15 +464,86 @@ def test_deepseek_v41_decoder_mixes_match_torch(
464 post, res, collapsed, pre = mhc_pre_delayed_torch(*args, **kwargs)
465 return post, res, decoder.attn_norm(collapsed), pre
466
467 + def fused_reference(x, residual, post_mix, res_mix, *args, capture_aux=False, **kw):
468 + residual = mhc_post_torch(x, residual, post_mix, res_mix)
469 + aux = residual.mean(dim=1) if capture_aux else residual.new_empty(0)
470 + return residual, *reference(residual, *args, **kw), aux
471 +
472 monkeypatch.setattr(
473 "vllm.models.deepseek_v41.nvidia.model.mhc_pre_delayed_tilelang",
474 reference,
475 )
476 + monkeypatch.setattr(
477 + "vllm.models.deepseek_v41.nvidia.model.mhc_fused_post_pre_delayed_tilelang",
478 + fused_reference,
479 + )
480 expected = decoder(x, positions, None, **kwargs)
303 - for result, ref in zip(actual, expected, strict=True):
481 + assert actual[-1] is None and expected[-1] is None
482 + for result, ref in zip(actual[:-1], expected[:-1], strict=True):
483 torch.testing.assert_close(result, ref, atol=2e-2, rtol=1e-2)
484
485
486 +@pytest.mark.skipif(not HAS_TILELANG_MHC, reason="TileLang MHC support required")
487 +@pytest.mark.parametrize("entry", ["fused", "engram"])
488 +def test_deepseek_v41_capture_previous_aux(entry, monkeypatch, default_vllm_config):
489 + """Read the aux hidden state back out of the seam that already computed it.
490 +
491 + The draft model's input is the mean over hc streams of the previous
492 + layer's post. That post now runs inside this layer's fused pre, so the
493 + captured value must equal a standalone ``mhc_post_tilelang`` on the same
494 + inputs -- and Engram must not be folded into it, since the injection
495 + happens after the aux consumers read the stream.
496 + """
497 + set_random_seed(0)
498 + decoder = DeepseekV41DecoderLayer.__new__(DeepseekV41DecoderLayer)
499 + nn.Module.__init__(decoder)
500 + decoder.hc_mult = 4
501 + decoder.hc_sinkhorn_iters = 20
502 + decoder.hc_eps = 1e-6
503 + decoder.rms_norm_eps = 1e-6
504 + decoder.hc_post_alpha = 2.0
505 + decoder.use_sequence_parallel = False
506 + decoder.engram = None
507 + from vllm.model_executor.layers.layernorm import RMSNorm
508 +
509 + decoder.attn_norm = decoder.ffn_norm = RMSNorm(5120, 1e-6).to(
510 + device=DEVICE, dtype=torch.bfloat16
511 + )
512 + decoder.attn = lambda positions, x, _: x * 0.5
513 + decoder.ffn = lambda x, input_ids: x * 0.25
514 + with torch.device(DEVICE):
515 + decoder.hc_attn_fn = torch.randn(24, 20480) * 0.02
516 + decoder.hc_ffn_fn = torch.randn(24, 20480) * 0.02
517 + decoder.hc_attn_scale = decoder.hc_ffn_scale = torch.ones(3)
518 + decoder.hc_attn_base = torch.randn(24)
519 + decoder.hc_ffn_base = torch.randn(24)
520 + x = torch.randn(3, 5120, dtype=torch.bfloat16)
521 + kwargs = dict(
522 + pre_mix=torch.rand(3, 4),
523 + residual=torch.randn(3, 4, 5120, dtype=torch.bfloat16),
524 + post_mix=torch.rand(3, 4, 1),
525 + res_mix=torch.rand(3, 4, 4),
526 + )
527 + if entry == "engram":
528 +
529 + class FakeEngram(nn.Module):
530 + layer_hash_index = 0
531 +
532 + def forward(self, residual, hashes, mask):
533 + return residual + 0.125
534 +
535 + decoder.engram = FakeEngram()
536 + kwargs["engram_hashes"] = torch.zeros(3, 1, 1, dtype=torch.int32)
537 + positions = torch.arange(3)
538 +
539 + expected = torch.ops.vllm.mhc_post_tilelang(
540 + x, kwargs["residual"], kwargs["post_mix"], kwargs["res_mix"]
541 + ).mean(dim=1)
542 + previous_aux = decoder(x, positions, None, **kwargs, capture_previous_aux=True)[-1]
543 + torch.testing.assert_close(previous_aux, expected, atol=0, rtol=0)
544 + assert decoder(x, positions, None, **kwargs)[-1] is None
545 +
546 +
547 @pytest.mark.skipif(not HAS_TILELANG_MHC, reason="TileLang MHC support required")
548 @pytest.mark.parametrize("carried", [False, True])
549 def test_mhc_pre_delayed_custom_op_supports_compile(carried):
vllm/model_executor/kernels/mhc/tilelang.py
+283 -2
@@ -194,6 +194,7 @@ def mhc_pre_delayed_tilelang(
194 norm_weight,
195 pre_mix if pre_mix is not None else post,
196 next_pre_mix,
197 + layer_input,
198 hidden_size=hidden_size,
199 rms_eps=rms_eps,
200 hc_pre_eps=hc_pre_eps,
@@ -218,6 +219,7 @@ def mhc_pre_delayed_tilelang(
219 layer_input,
220 pre_mix if pre_mix is not None else post,
221 next_pre_mix,
222 + layer_input,
223 hidden_size,
224 rms_eps,
225 hc_pre_eps,
@@ -263,6 +265,278 @@ def _mhc_pre_delayed_tilelang_fake(
265 )
266
267
268 +def mhc_fused_post_pre_delayed_tilelang(
269 + x: torch.Tensor,
270 + residual: torch.Tensor,
271 + post_layer_mix: torch.Tensor,
272 + comb_res_mix: torch.Tensor,
273 + fn: torch.Tensor,
274 + hc_scale: torch.Tensor,
275 + hc_base: torch.Tensor,
276 + rms_eps: float,
277 + hc_pre_eps: float,
278 + hc_sinkhorn_eps: float,
279 + hc_post_mult_value: float,
280 + sinkhorn_repeat: int,
281 + pre_mix: torch.Tensor | None = None,
282 + norm_weight: torch.Tensor | None = None,
283 + norm_eps: float = 1e-6,
284 + capture_aux: bool = False,
285 +) -> tuple[
286 + torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor
287 +]:
288 + """Run one mHC post block followed by the next delayed mHC pre block.
289 +
290 + Within the fused kernel's token range the post mapping is folded into the
291 + pre-norm GEMM, so the updated residual streams feed the projection from
292 + registers instead of a second pass over global memory. Above it this runs
293 + the same post kernel and split-k GEMM as the unfused pair.
294 +
295 + Args:
296 + x: BF16 sublayer output of shape (tokens, hidden_size).
297 + residual: BF16 residual streams of shape (tokens, hc_mult, hidden_size).
298 + post_layer_mix: FP32 post coefficients of shape (tokens, hc_mult, 1).
299 + comb_res_mix: FP32 residual coefficients, (tokens, hc_mult, hc_mult).
300 + fn: FP32 projection of shape (hc_mult * (hc_mult + 2), input_size).
301 + hc_scale: FP32 scales of shape (3,).
302 + hc_base: FP32 bias of shape (hc_mult * (hc_mult + 2),).
303 + rms_eps: RMS normalization epsilon.
304 + hc_pre_eps: Pre-mix epsilon.
305 + hc_sinkhorn_eps: Sinkhorn epsilon.
306 + hc_post_mult_value: Post-mix multiplier.
307 + sinkhorn_repeat: Number of Sinkhorn iterations.
308 + pre_mix: FP32 coefficients from the previous sublayer, or None to
309 + select residual stream zero.
310 + norm_weight: Optional BF16 RMSNorm weight for the collapsed input.
311 + norm_eps: RMSNorm epsilon for the collapsed input.
312 + capture_aux: Also return the mean over the post-mapped streams, which
313 + draft models consume as the target's hidden state. It is folded
314 + into the collapse, which already reads those streams.
315 +
316 + Returns:
317 + The post-mapped residual streams, the post and residual coefficients,
318 + the optionally normalized BF16 layer input, the next FP32 pre-mix, and
319 + the BF16 stream mean (empty unless capture_aux), with shapes
320 + (tokens, hc_mult, hidden_size), (tokens, hc_mult, 1),
321 + (tokens, hc_mult, hc_mult), (tokens, hidden_size), (tokens, hc_mult),
322 + and (tokens, hidden_size).
323 + """
324 + from vllm.model_executor.kernels.mhc.tilelang_kernels import (
325 + _HC_PRENORM_GEMM_TILELANG_KERNEL,
326 + _MHC_FUSED_TILELANG_KERNEL,
327 + _MHC_POST_TILELANG_KERNEL,
328 + mhc_fused_post_pre_split_config,
329 + mhc_pre_big_fuse_tilelang,
330 + )
331 + from vllm.model_executor.kernels.mhc.warmup import (
332 + MHC_PRE_NORM_KERNEL,
333 + compute_mhc_pre_num_splits,
334 + )
335 + from vllm.utils.deep_gemm import (
336 + is_deep_gemm_supported,
337 + tf32_hc_prenorm_gemm,
338 + )
339 +
340 + assert residual.ndim == 3 and residual.dtype == torch.bfloat16
341 + assert residual.is_contiguous()
342 + num_tokens, hc_mult, hidden_size = residual.shape
343 + input_size = hc_mult * hidden_size
344 + mix_size = hc_mult * (hc_mult + 2)
345 + assert x.shape == (num_tokens, hidden_size) and x.dtype == torch.bfloat16
346 + assert x.is_contiguous()
347 + assert post_layer_mix.shape[:2] == (num_tokens, hc_mult)
348 + assert post_layer_mix.dtype == torch.float32 and post_layer_mix.is_contiguous()
349 + assert comb_res_mix.shape == (num_tokens, hc_mult, hc_mult)
350 + assert comb_res_mix.dtype == torch.float32 and comb_res_mix.is_contiguous()
351 + assert fn.shape == (mix_size, input_size) and fn.dtype == torch.float32
352 + assert hc_scale.shape == (3,) and hc_scale.dtype == torch.float32
353 + assert hc_base.shape == (mix_size,) and hc_base.dtype == torch.float32
354 + if pre_mix is not None:
355 + assert pre_mix.shape == (num_tokens, hc_mult)
356 + assert pre_mix.dtype == torch.float32 and pre_mix.is_contiguous()
357 +
358 + next_pre_mix = torch.empty(
359 + num_tokens, hc_mult, dtype=torch.float32, device=residual.device
360 + )
361 + post = torch.empty_like(next_pre_mix)
362 + comb = torch.empty(
363 + num_tokens, hc_mult * hc_mult, dtype=torch.float32, device=residual.device
364 + )
365 + layer_input = torch.empty(
366 + num_tokens, hidden_size, dtype=torch.bfloat16, device=residual.device
367 + )
368 + aux = torch.empty(
369 + num_tokens if capture_aux else 0,
370 + hidden_size,
371 + dtype=torch.bfloat16,
372 + device=residual.device,
373 + )
374 + if num_tokens == 0:
375 + return (
376 + torch.empty_like(residual),
377 + post.unsqueeze(-1),
378 + comb.view(num_tokens, hc_mult, hc_mult),
379 + layer_input,
380 + next_pre_mix,
381 + aux,
382 + )
383 +
384 + fused_config = mhc_fused_post_pre_split_config(num_tokens, hidden_size, hc_mult)
385 + if fused_config is not None:
386 + mixes, sqrsum, residual_cur = _MHC_FUSED_TILELANG_KERNEL(
387 + comb_res_mix,
388 + residual,
389 + post_layer_mix.view(num_tokens, hc_mult),
390 + x,
391 + fn.view(mix_size, hc_mult, hidden_size),
392 + hc_mult,
393 + hidden_size,
394 + mix_size,
395 + )
396 + else:
397 + residual_cur = torch.empty_like(residual)
398 + _MHC_POST_TILELANG_KERNEL(
399 + comb_res_mix,
400 + residual,
401 + post_layer_mix.view(num_tokens, hc_mult),
402 + x,
403 + residual_cur,
404 + hc_mult,
405 + hidden_size,
406 + )
407 + # The delayed epilogue is compiled per bucketed split count, so the
408 + # projection has to use the same bucket rather than the raw estimate.
409 + use_deep_gemm = is_deep_gemm_supported()
410 + n_splits = (
411 + compute_mhc_pre_num_splits(input_size, num_tokens) if use_deep_gemm else 1
412 + )
413 + mixes = torch.empty(
414 + n_splits, num_tokens, mix_size, dtype=torch.float32, device=residual.device
415 + )
416 + sqrsum = torch.empty(
417 + n_splits, num_tokens, dtype=torch.float32, device=residual.device
418 + )
419 + residual_cur_2d = residual_cur.view(num_tokens, input_size)
420 + if use_deep_gemm:
421 + tf32_hc_prenorm_gemm(residual_cur_2d, fn, mixes, sqrsum, n_splits)
422 + else:
423 + _HC_PRENORM_GEMM_TILELANG_KERNEL(
424 + residual_cur_2d,
425 + fn,
426 + mixes,
427 + sqrsum,
428 + input_size,
429 + 1,
430 + )
431 +
432 + outputs = (
433 + residual_cur,
434 + post.unsqueeze(-1),
435 + comb.view(num_tokens, hc_mult, hc_mult),
436 + layer_input,
437 + next_pre_mix,
438 + aux,
439 + )
440 + if norm_weight is not None:
441 + assert norm_weight.shape == (hidden_size,)
442 + assert norm_weight.dtype == torch.bfloat16 and norm_weight.is_contiguous()
443 + MHC_PRE_NORM_KERNEL(
444 + mixes,
445 + sqrsum,
446 + hc_scale,
447 + hc_base,
448 + residual_cur,
449 + post,
450 + comb,
451 + layer_input,
452 + norm_weight,
453 + pre_mix if pre_mix is not None else post,
454 + next_pre_mix,
455 + aux if capture_aux else layer_input,
456 + hidden_size=hidden_size,
457 + rms_eps=rms_eps,
458 + hc_pre_eps=hc_pre_eps,
459 + hc_sinkhorn_eps=hc_sinkhorn_eps,
460 + hc_post_mult_value=hc_post_mult_value,
461 + sinkhorn_repeat=sinkhorn_repeat,
462 + norm_eps=norm_eps,
463 + hc_mult=hc_mult,
464 + use_pre_mix_in=pre_mix is not None,
465 + save_pre_mix=True,
466 + rms_numel=input_size,
467 + write_aux=capture_aux,
468 + )
469 + return outputs
470 + mhc_pre_big_fuse_tilelang(
471 + mixes,
472 + sqrsum,
473 + hc_scale,
474 + hc_base,
475 + residual_cur,
476 + post,
477 + comb,
478 + layer_input,
479 + pre_mix if pre_mix is not None else post,
480 + next_pre_mix,
481 + aux if capture_aux else layer_input,
482 + hidden_size,
483 + rms_eps,
484 + hc_pre_eps,
485 + hc_sinkhorn_eps,
486 + hc_post_mult_value,
487 + sinkhorn_repeat,
488 + mixes.shape[0],
489 + hc_mult,
490 + use_pre_mix_in=pre_mix is not None,
491 + save_pre_mix=True,
492 + rms_numel=input_size,
493 + write_aux=capture_aux,
494 + )
495 + return outputs
496 +
497 +
498 +def _mhc_fused_post_pre_delayed_tilelang_fake(
499 + x: torch.Tensor,
500 + residual: torch.Tensor,
501 + post_layer_mix: torch.Tensor,
502 + comb_res_mix: torch.Tensor,
503 + fn: torch.Tensor,
504 + hc_scale: torch.Tensor,
505 + hc_base: torch.Tensor,
506 + rms_eps: float,
507 + hc_pre_eps: float,
508 + hc_sinkhorn_eps: float,
509 + hc_post_mult_value: float,
510 + sinkhorn_repeat: int,
511 + pre_mix: torch.Tensor | None = None,
512 + norm_weight: torch.Tensor | None = None,
513 + norm_eps: float = 1e-6,
514 + capture_aux: bool = False,
515 +) -> tuple[
516 + torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor
517 +]:
518 + num_tokens, hc_mult, hidden_size = residual.shape
519 + return (
520 + torch.empty_like(residual),
521 + torch.empty(
522 + num_tokens, hc_mult, 1, dtype=torch.float32, device=residual.device
523 + ),
524 + torch.empty(
525 + num_tokens, hc_mult, hc_mult, dtype=torch.float32, device=residual.device
526 + ),
527 + torch.empty(
528 + num_tokens, hidden_size, dtype=torch.bfloat16, device=residual.device
529 + ),
530 + torch.empty(num_tokens, hc_mult, dtype=torch.float32, device=residual.device),
531 + torch.empty(
532 + num_tokens if capture_aux else 0,
533 + hidden_size,
534 + dtype=torch.bfloat16,
535 + device=residual.device,
536 + ),
537 + )
538 +
539 +
540 def mhc_pre_tilelang(
541 residual: torch.Tensor,
542 fn: torch.Tensor,
@@ -574,6 +848,7 @@ def mhc_fused_post_pre_tilelang(
848 _MHC_FUSED_TILELANG_KERNEL,
849 _MHC_POST_TILELANG_KERNEL,
850 _MHC_PRE_BIG_FUSE_TILELANG_KERNEL,
851 + mhc_fused_post_pre_split_config,
852 )
853
854 assert residual.dtype == torch.bfloat16
@@ -614,7 +889,7 @@ def mhc_fused_post_pre_tilelang(
889 post_layer_mix_flat = post_layer_mix.view(num_tokens, hc_mult)
890 comb_res_mix_flat = comb_res_mix.view(num_tokens, hc_mult, hc_mult)
891
617 - use_small_fma = num_tokens <= 16
892 + fused_config = mhc_fused_post_pre_split_config(num_tokens, hidden_size, hc_mult)
893
894 post_mix_cur = torch.empty(
895 num_tokens,
@@ -635,7 +910,7 @@ def mhc_fused_post_pre_tilelang(
910 device=residual.device,
911 )
912
638 - if use_small_fma:
913 + if fused_config is not None:
914 gemm_out_mul, gemm_out_sqrsum, residual_cur = _MHC_FUSED_TILELANG_KERNEL(
915 comb_res_mix_flat,
916 residual_flat,
@@ -815,6 +1090,12 @@ direct_register_custom_op(
1090 mutates_args=[],
1091 fake_impl=_mhc_pre_delayed_tilelang_fake,
1092 )
1093 +direct_register_custom_op(
1094 + op_name="mhc_fused_post_pre_delayed_tilelang",
1095 + op_func=mhc_fused_post_pre_delayed_tilelang,
1096 + mutates_args=[],
1097 + fake_impl=_mhc_fused_post_pre_delayed_tilelang_fake,
1098 +)
1099 direct_register_custom_op(
1100 op_name="mhc_pre_tilelang",
1101 op_func=mhc_pre_tilelang,
vllm/model_executor/kernels/mhc/tilelang_kernels.py
+124 -16
@@ -37,8 +37,67 @@ def compute_num_split(block_k: int, k: int | None, grid_size: int) -> int:
37 return split_k
38
39
40 -def _mhc_fused_n_splits(num_tokens: int, hidden_size: int) -> int:
41 - return 8 if num_tokens < 8 and hidden_size <= 4096 else 4
40 +# ``mhc_fused_tilelang`` splits the hidden size across a fixed thread block, so
41 +# every (tile_n, n_thr) pair is one more kernel to compile. Tuned on GB300
42 +# (hc_mult 4, hidden_size 5120 and 7168) against the separate post + split-k
43 +# GEMM it replaces; the bands are coarse because the loss against a per-shape
44 +# optimum is ~1-2%, while one fixed config gives up to 18% at the top of the
45 +# range.
46 +_FUSED_POST_PRE_N_THR = 128
47 +_FUSED_POST_PRE_MAX_TOKENS = 32
48 +
49 +
50 +def mhc_fused_post_pre_split_config(
51 + num_tokens: int, hidden_size: int, hc_mult: int
52 +) -> tuple[int, int, int] | None:
53 + """Pick ``(tile_n, n_splits, n_thr)`` for the fused post + pre-norm GEMM.
54 +
55 + Returns None when a separate post kernel followed by a split-k GEMM is the
56 + better choice, or when the hidden size does not divide evenly across the
57 + fused kernel's block. One source of truth for the dispatch decision, the
58 + compile key and the launch, which must agree.
59 + """
60 + if num_tokens > _FUSED_POST_PRE_MAX_TOKENS:
61 + return None
62 + n_thr = _FUSED_POST_PRE_N_THR
63 + if hidden_size % n_thr:
64 + return None
65 + # The projection tiles supply the parallelism a small token count cannot,
66 + # so the tile grows with the batch: 1.3x at one token and 1.03x at 32
67 + # against a split-k GEMM given the split its own estimator computes.
68 + if num_tokens < 16:
69 + tile_n, n_splits = 2, 8
70 + else:
71 + tile_n, n_splits = 6, 8
72 + # The kernel drops projection tiles and k-slices it cannot fill evenly.
73 + mix_size = hc_mult * (hc_mult + 2)
74 + while tile_n > 1 and mix_size % tile_n:
75 + tile_n -= 1
76 + while n_splits > 1 and hidden_size % (n_splits * n_thr):
77 + n_splits //= 2
78 + return tile_n, n_splits, n_thr
79 +
80 +
81 +def require_fused_post_pre_config(
82 + num_tokens: int, hidden_size: int, hc_mult: int
83 +) -> tuple[int, int, int]:
84 + """The config for a shape the caller has already committed to fusing."""
85 + config = mhc_fused_post_pre_split_config(num_tokens, hidden_size, hc_mult)
86 + if config is None:
87 + raise ValueError(
88 + "the fused mHC post + pre-norm GEMM does not cover num_tokens="
89 + f"{num_tokens} at hidden_size={hidden_size}"
90 + )
91 + return config
92 +
93 +
94 +def mhc_fused_post_pre_splits(hidden_size: int, hc_mult: int) -> tuple[int, ...]:
95 + """Every split-k factor the fused post + pre-norm GEMM path can pick."""
96 + configs = (
97 + mhc_fused_post_pre_split_config(num_tokens, hidden_size, hc_mult)
98 + for num_tokens in range(1, _FUSED_POST_PRE_MAX_TOKENS + 1)
99 + )
100 + return tuple(sorted({config[1] for config in configs if config is not None}))
101
102
103 @tilelang_jit
@@ -53,6 +112,7 @@ def mhc_pre_big_fuse_tilelang(
112 layer_input,
113 pre_mix_in,
114 pre_mix_out,
115 + aux_out,
116 hidden_size: int,
117 rms_eps: float,
118 hc_pre_eps: float,
@@ -64,6 +124,7 @@ def mhc_pre_big_fuse_tilelang(
124 use_pre_mix_in: bool = False,
125 save_pre_mix: bool = False,
126 rms_numel: int = 0,
127 + write_aux: bool = False,
128 ):
129 """Fuse coefficient generation and residual collapse after the projection.
130
@@ -88,6 +149,7 @@ def mhc_pre_big_fuse_tilelang(
149
150 pre_mix_in: T.Tensor[[num_tokens, hc_mult], T.float32] # type: ignore[no-redef, valid-type]
151 pre_mix_out: T.Tensor[[num_tokens, hc_mult], T.float32] # type: ignore[no-redef, valid-type]
152 + aux_out: T.Tensor[[num_tokens, hidden_size], T.bfloat16] # type: ignore[no-redef, valid-type]
153
154 with T.Kernel(num_tokens, threads=96) as i:
155 if ENABLE_PDL:
@@ -188,11 +250,24 @@ def mhc_pre_big_fuse_tilelang(
250
251 ol = T.alloc_fragment(hidden_block, T.float32)
252 T.clear(ol)
253 + if write_aux:
254 + # Aux consumers (draft models) take the plain stream mean
255 + # of the same residual this collapse reads.
256 + aux = T.alloc_fragment(hidden_block, T.float32)
257 + T.clear(aux)
258
259 for i_hc in T.serial(hc_mult):
260 pre = pre_mix_shared[i_hc]
261 for i1_h in T.Parallel(hidden_block):
262 ol[i1_h] += pre * xl[i_hc, i1_h]
263 + if write_aux:
264 + aux[i1_h] += xl[i_hc, i1_h]
265 +
266 + if write_aux:
267 + for i1_h in T.Parallel(hidden_block):
268 + aux_out[i, i0_h * hidden_block + i1_h] = T.bfloat16(
269 + aux[i1_h] / hc_mult
270 + )
271
272 T.copy(ol, layer_input[i, i0_h * hidden_block])
273
@@ -216,6 +291,7 @@ def mhc_pre_big_fuse_with_norm_tilelang(
291 norm_weight,
292 pre_mix_in,
293 pre_mix_out,
294 + aux_out,
295 hidden_size: int,
296 rms_eps: float,
297 hc_pre_eps: float,
@@ -229,6 +305,7 @@ def mhc_pre_big_fuse_with_norm_tilelang(
305 use_pre_mix_in: bool = False,
306 save_pre_mix: bool = False,
307 rms_numel: int = 0,
308 + write_aux: bool = False,
309 ):
310 num_tokens = T.dynamic("num_tokens")
311 hc_mult3 = hc_mult * (2 + hc_mult)
@@ -250,6 +327,7 @@ def mhc_pre_big_fuse_with_norm_tilelang(
327
328 pre_mix_in: T.Tensor[[num_tokens, hc_mult], T.float32] # type: ignore[no-redef, valid-type]
329 pre_mix_out: T.Tensor[[num_tokens, hc_mult], T.float32] # type: ignore[no-redef, valid-type]
330 + aux_out: T.Tensor[[num_tokens, hidden_size], T.bfloat16] # type: ignore[no-redef, valid-type]
331
332 with T.Kernel(num_tokens, threads=96) as i:
333 rms = T.alloc_fragment(1, T.float32)
@@ -345,11 +423,24 @@ def mhc_pre_big_fuse_with_norm_tilelang(
423
424 ol = T.alloc_fragment(hidden_block, T.float32)
425 T.clear(ol)
426 + if write_aux:
427 + # Aux consumers (draft models) take the plain stream mean
428 + # of the same residual this collapse reads.
429 + aux = T.alloc_fragment(hidden_block, T.float32)
430 + T.clear(aux)
431
432 for i_hc in T.serial(hc_mult):
433 pre = pre_mix_shared[i_hc]
434 for i1_h in T.Parallel(hidden_block):
435 ol[i1_h] += pre * xl[i_hc, i1_h]
436 + if write_aux:
437 + aux[i1_h] += xl[i_hc, i1_h]
438 +
439 + if write_aux:
440 + for i1_h in T.Parallel(hidden_block):
441 + aux_out[i, i0_h * hidden_block + i1_h] = T.bfloat16(
442 + aux[i1_h] / hc_mult
443 + )
444
445 if save_pre_mix:
446 # Keep the BF16 boundary before the delayed input RMSNorm.
@@ -1211,9 +1302,17 @@ class MhcPreBigFuseTileLangKernel(
1302 middle_args = (residual_out, *output_args, norm_weight)
1303 elif compile_key.use_norm_weight:
1304 assert norm_weight is not None
1214 - middle_args = (*output_args, norm_weight, post_mix, post_mix)
1305 + # layer_input stands in for the aux buffer on both epilogues: same
1306 + # shape and dtype, and write_aux is off here so it is never written.
1307 + middle_args = (
1308 + *output_args,
1309 + norm_weight,
1310 + post_mix,
1311 + post_mix,
1312 + layer_input,
1313 + )
1314 else:
1216 - middle_args = (*output_args, post_mix, post_mix)
1315 + middle_args = (*output_args, post_mix, post_mix, layer_input)
1316 norm_eps_args = (compile_key.norm_eps,) if compile_key.use_norm_weight else ()
1317 return (
1318 *common_args,
@@ -1247,8 +1346,14 @@ class MhcPreBigFuseTileLangKernel(
1346 if use_pre_gemm_splits
1347 else n_splits
1348 )
1250 - fused_n_splits = _mhc_fused_n_splits(num_tokens, hidden_size)
1251 - actual_n_splits = fused_n_splits if use_fused_tilelang else pre_gemm_n_splits
1349 + # The epilogue reduces over whatever the fused kernel split the GEMM
1350 + # into, so both have to read the same config.
1351 + fused_config = mhc_fused_post_pre_split_config(num_tokens, hidden_size, hc_mult)
1352 + actual_n_splits = (
1353 + fused_config[1]
1354 + if use_fused_tilelang and fused_config is not None
1355 + else pre_gemm_n_splits
1356 + )
1357 actual_norm_eps = broadcast_norm_eps if is_broadcast else norm_eps
1358 actual_use_norm_weight = use_norm_weight or is_broadcast
1359 return self.CompileKey(
@@ -1540,6 +1645,7 @@ class MhcFusedTileLangKernel(
1645 hc_mult: int
1646 n_splits: int
1647 tile_n: int
1648 + n_thr: int
1649
1650 @staticmethod
1651 def kernel() -> Any:
@@ -1552,14 +1658,15 @@ class MhcFusedTileLangKernel(
1658 hidden_size: int,
1659 hc_mult: int,
1660 ) -> CompileKey:
1555 - # TODO(gnovack): investigate autotuning these heuristics
1556 - tile_n = 2 if num_tokens < 8 else 3
1557 - n_splits = _mhc_fused_n_splits(num_tokens, hidden_size)
1661 + # Subscripts rather than an unpack: the warmup tracer parses this body
1662 + # and allows only plain assignments before the return.
1663 + config = require_fused_post_pre_config(num_tokens, hidden_size, hc_mult)
1664 return self.CompileKey(
1665 hidden_size=hidden_size,
1666 hc_mult=hc_mult,
1561 - n_splits=n_splits,
1562 - tile_n=tile_n,
1667 + n_splits=config[1],
1668 + tile_n=config[0],
1669 + n_thr=config[2],
1670 )
1671
1672 def get_warmup_keys(
@@ -1574,11 +1681,11 @@ class MhcFusedTileLangKernel(
1681 num_tokens=WarmupIntRange(1, max_tokens + 1),
1682 hidden_size=hidden_size,
1683 hc_mult=hc_mult,
1577 - _when=lambda *, num_tokens: num_tokens <= 16,
1684 + _when=lambda *, num_tokens: num_tokens <= _FUSED_POST_PRE_MAX_TOKENS,
1685 )
1686
1687 def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]:
1581 - num_tokens = 1 if compile_key.tile_n == 2 else 8
1688 + num_tokens = 1 if compile_key.tile_n == 2 else 16
1689 hidden_size = compile_key.hidden_size
1690 hc_mult = compile_key.hc_mult
1691 hc_mult3 = hc_mult * (2 + hc_mult)
@@ -1617,8 +1724,9 @@ class MhcFusedTileLangKernel(
1724 hc_mult3: int,
1725 ) -> TileLangLaunchSpec:
1726 num_tokens = residual_in.shape[0]
1620 - tile_n = 2 if num_tokens < 8 else 3
1621 - n_splits = _mhc_fused_n_splits(num_tokens, hidden_size)
1727 + tile_n, n_splits, n_thr = require_fused_post_pre_config(
1728 + num_tokens, hidden_size, hc_mult
1729 + )
1730 yp_out = residual_in.new_empty(
1731 (n_splits, num_tokens, hc_mult3), dtype=torch.float32
1732 )
@@ -1639,7 +1747,7 @@ class MhcFusedTileLangKernel(
1747 hidden_size,
1748 hc_mult3,
1749 ),
1642 - dict(tile_n=tile_n, split_k=n_splits),
1750 + dict(n_thr=n_thr, tile_n=tile_n),
1751 (yp_out, rp_out, residual_out),
1752 )
1753
vllm/model_executor/kernels/mhc/warmup.py
+16 -6
@@ -36,23 +36,33 @@ class MHCPreNormKernel(VllmJitKernel["MHCPreNormKernel.CompileKey"]):
36 use_pre_mix_in: bool
37 rms_numel: int
38 save_pre_mix: bool = True
39 + write_aux: bool = False
40
41 kernel: Any = staticmethod(mhc_pre_big_fuse_with_norm_tilelang)
42
43 def dispatch(self, *, n_splits, **fields) -> CompileKey: # type: ignore[override]
44 return self.CompileKey(n_splits=n_splits, **fields)
45
45 - def get_warmup_keys(self, *, max_tokens: int, **fields) -> list[CompileKey]:
46 - # The split heuristic changes only at 64-token boundaries.
47 - splits = (
48 - sorted(
46 + def get_warmup_keys(
47 + self,
48 + *,
49 + max_tokens: int,
50 + extra_splits: Iterable[int] = (),
51 + **fields,
52 + ) -> list[CompileKey]:
53 + # The split heuristic changes only at 64-token boundaries. Callers that
54 + # reach this kernel through another GEMM (the fused post + pre-norm
55 + # GEMM) pass the splits that GEMM picks in extra_splits.
56 + splits = sorted(
57 + set(extra_splits)
58 + | (
59 {
60 compute_mhc_pre_num_splits(fields["rms_numel"], num_tokens)
61 for num_tokens in range(1, max_tokens + 1, 64)
62 }
63 + if is_deep_gemm_supported()
64 + else {1}
65 )
54 - if is_deep_gemm_supported()
55 - else [1]
66 )
67 return self._trace_dispatch(self.dispatch)(n_splits=splits, **fields)
68
vllm/models/deepseek_v41/nvidia/dspark.py
+1 -1
@@ -207,7 +207,7 @@ class DSparkDeepseekV4Model(nn.Module):
207
208 residual = post_mix = res_mix = pre_mix = None
209 for layer in self.layers:
210 - hidden_states, residual, post_mix, res_mix, pre_mix = layer(
210 + hidden_states, residual, post_mix, res_mix, pre_mix, _ = layer(
211 hidden_states,
212 positions,
213 input_ids,
vllm/models/deepseek_v41/nvidia/model.py
+126 -52
@@ -20,6 +20,7 @@ from vllm.distributed import (
20 from vllm.forward_context import get_forward_context, is_forward_context_available
21 from vllm.logger import init_logger
22 from vllm.model_executor.kernels.mhc.tilelang import (
23 + mhc_fused_post_pre_delayed_tilelang,
24 mhc_post_tilelang,
25 mhc_pre_delayed_tilelang,
26 )
@@ -214,17 +215,45 @@ class DeepseekV4DecoderLayer(nn.Module):
215 self.hc_eps = config.hc_eps
216 self.hc_post_alpha = 2.0
217 if vllm_config.kernel_config.enable_jit_warmup and current_platform.is_cuda():
218 + from vllm.model_executor.kernels.mhc.tilelang_kernels import (
219 + mhc_fused_post_pre_splits,
220 + )
221 from vllm.model_executor.kernels.mhc.warmup import MHC_PRE_NORM_KERNEL
222
223 max_tokens = vllm_config.scheduler_config.max_num_batched_tokens
224 if self.use_sequence_parallel:
225 tp_size = vllm_config.parallel_config.tensor_parallel_size
226 max_tokens = (max_tokens + tp_size - 1) // tp_size
223 - for input_size, use_pre_mix in (
224 - (self.hidden_size, False),
225 - (self.hc_mult * self.hidden_size, False),
226 - (self.hc_mult * self.hidden_size, True),
227 - ):
227 + # The epilogue compiles per projection width and per pre-mix mode.
228 + # The first layer projects the broadcast embedding, so it reads one
229 + # hidden_size-wide row and selects stream zero; every later sublayer
230 + # projects the full hc stream, and collapses it with the pre-mix the
231 + # previous sublayer carried in. Those later shapes also arrive
232 + # through the fused post + pre-norm GEMM, which picks split-k
233 + # factors the token sweep below never produces.
234 + broadcast_embedding = {
235 + "rms_numel": self.hidden_size,
236 + "use_pre_mix_in": False,
237 + "extra_splits": (),
238 + }
239 + hc_stream = {
240 + "rms_numel": self.hc_mult * self.hidden_size,
241 + "extra_splits": mhc_fused_post_pre_splits(
242 + self.hidden_size, self.hc_mult
243 + ),
244 + }
245 + variants = [
246 + broadcast_embedding,
247 + {**hc_stream, "use_pre_mix_in": False},
248 + {**hc_stream, "use_pre_mix_in": True},
249 + ]
250 + if vllm_config.speculative_config is not None:
251 + # Draft setups read aux hidden states out of the same collapse,
252 + # which is its own specialization of the kernel.
253 + variants.append(
254 + {**hc_stream, "use_pre_mix_in": True, "write_aux": True}
255 + )
256 + for variant in variants:
257 MHC_PRE_NORM_KERNEL.register_warmup(
258 max_tokens=max_tokens,
259 hidden_size=self.hidden_size,
@@ -235,8 +264,7 @@ class DeepseekV4DecoderLayer(nn.Module):
264 sinkhorn_repeat=self.hc_sinkhorn_iters,
265 norm_eps=self.rms_norm_eps,
266 hc_mult=self.hc_mult,
238 - use_pre_mix_in=use_pre_mix,
239 - rms_numel=input_size,
267 + **variant,
268 )
269 mix_hc = (2 + self.hc_mult) * self.hc_mult
270 hc_dim = self.hc_mult * self.hidden_size
@@ -295,7 +323,17 @@ class DeepseekV4DecoderLayer(nn.Module):
323 residual: torch.Tensor | None = None,
324 engram_hashes: torch.Tensor | None = None,
325 engram_mask: torch.Tensor | None = None,
298 - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
326 + *,
327 + capture_previous_aux: bool = False,
328 + ) -> tuple[
329 + torch.Tensor,
330 + torch.Tensor,
331 + torch.Tensor,
332 + torch.Tensor,
333 + torch.Tensor,
334 + torch.Tensor | None,
335 + ]:
336 + previous_aux: torch.Tensor | None = None
337 # The reference collapses each sublayer's input with the *previous*
338 # sublayer's pre-mix: attention uses the pre-mix carried in (identity
339 # for the first layer), the FFN uses this layer's attention pre-mix.
@@ -335,17 +373,19 @@ class DeepseekV4DecoderLayer(nn.Module):
373 norm_weight=self.attn_norm.weight,
374 norm_eps=self.attn_norm.variance_epsilon,
375 )
338 - else:
339 - residual = mhc_post_tilelang(x, residual, post_mix, res_mix)
340 - if self.engram is not None and engram_hashes is not None:
341 - # Engram injection happens between the previous sublayer's
342 - # post and this block's pre, on the full hc stream, so the
343 - # mix coefficients see the injected stream.
344 - residual = self.engram(
345 - residual,
346 - engram_hashes[:, self.engram.layer_hash_index],
347 - engram_mask,
348 - )
376 + elif self.engram is not None and engram_hashes is not None:
377 + # Engram injection happens between the previous sublayer's post
378 + # and this block's pre, on the full hc stream, so the mix
379 + # coefficients see the injected stream. The injection also keeps
380 + # the post out of the pre-norm GEMM's fused prologue.
381 + previous_post = mhc_post_tilelang(x, residual, post_mix, res_mix)
382 + if capture_previous_aux:
383 + previous_aux = previous_post.mean(dim=1)
384 + residual = self.engram(
385 + previous_post,
386 + engram_hashes[:, self.engram.layer_hash_index],
387 + engram_mask,
388 + )
389 post_mix, res_mix, x, attn_pre = mhc_pre_delayed_tilelang(
390 residual,
391 self.hc_attn_fn,
@@ -360,6 +400,31 @@ class DeepseekV4DecoderLayer(nn.Module):
400 norm_weight=self.attn_norm.weight,
401 norm_eps=self.attn_norm.variance_epsilon,
402 )
403 + else:
404 + # The collapse already reads the post-mapped streams, so the mean
405 + # aux consumers want comes out of the same kernel.
406 + residual, post_mix, res_mix, x, attn_pre, aux = (
407 + mhc_fused_post_pre_delayed_tilelang(
408 + x,
409 + residual,
410 + post_mix,
411 + res_mix,
412 + self.hc_attn_fn,
413 + self.hc_attn_scale,
414 + self.hc_attn_base,
415 + self.rms_norm_eps,
416 + self.hc_eps,
417 + self.hc_eps,
418 + self.hc_post_alpha,
419 + self.hc_sinkhorn_iters,
420 + pre_mix=pre_mix,
421 + norm_weight=self.attn_norm.weight,
422 + norm_eps=self.attn_norm.variance_epsilon,
423 + capture_aux=capture_previous_aux,
424 + )
425 + )
426 + if capture_previous_aux:
427 + previous_aux = aux
428
429 if self.use_sequence_parallel:
430 x = sp_all_gather(x)[: positions.shape[0]]
@@ -368,23 +433,27 @@ class DeepseekV4DecoderLayer(nn.Module):
433 if self.use_sequence_parallel:
434 x = sp_reduce_scatter(x)
435
371 - residual = mhc_post_tilelang(x, residual, post_mix, res_mix)
372 - post_mix, res_mix, x, ffn_pre = mhc_pre_delayed_tilelang(
373 - residual,
374 - self.hc_ffn_fn,
375 - self.hc_ffn_scale,
376 - self.hc_ffn_base,
377 - self.rms_norm_eps,
378 - self.hc_eps,
379 - self.hc_eps,
380 - self.hc_post_alpha,
381 - self.hc_sinkhorn_iters,
382 - pre_mix=attn_pre,
383 - norm_weight=self.ffn_norm.weight,
384 - norm_eps=self.ffn_norm.variance_epsilon,
436 + residual, post_mix, res_mix, x, ffn_pre, _ = (
437 + mhc_fused_post_pre_delayed_tilelang(
438 + x,
439 + residual,
440 + post_mix,
441 + res_mix,
442 + self.hc_ffn_fn,
443 + self.hc_ffn_scale,
444 + self.hc_ffn_base,
445 + self.rms_norm_eps,
446 + self.hc_eps,
447 + self.hc_eps,
448 + self.hc_post_alpha,
449 + self.hc_sinkhorn_iters,
450 + pre_mix=attn_pre,
451 + norm_weight=self.ffn_norm.weight,
452 + norm_eps=self.ffn_norm.variance_epsilon,
453 + )
454 )
455 x = self.ffn(x, input_ids)
387 - return x, residual, post_mix, res_mix, ffn_pre
456 + return x, residual, post_mix, res_mix, ffn_pre, previous_aux
457
458
459 class DeepseekV4Model(nn.Module, EagleModelMixin):
@@ -631,13 +700,14 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):
700 if not get_pp_group().is_first_rank:
701 assert intermediate_tensors is not None
702 pre_mix = intermediate_tensors["pre_mix"]
634 - aux_hidden_states: list[torch.Tensor] = []
635 - final_aux_recon: torch.Tensor | None = None # avoid duplicate mhc_post call
703 + # Every layer's post runs inside the next layer's fused pre, so aux
704 + # hidden states are read back from there instead of recomputed.
705 + aux_hidden_by_layer: dict[int, torch.Tensor] = {}
706 for idx, layer in enumerate(
707 islice(self.layers, self.start_layer, self.end_layer),
708 start=self.start_layer,
709 ):
640 - hidden_states, residual, post_mix, res_mix, pre_mix = layer(
710 + hidden_states, residual, post_mix, res_mix, pre_mix, previous_aux = layer(
711 hidden_states,
712 positions,
713 input_ids,
@@ -647,25 +717,29 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):
717 residual,
718 engram_hashes,
719 engram_mask,
720 + capture_previous_aux=idx in self.aux_hidden_state_layers,
721 )
651 - if idx + 1 in self.aux_hidden_state_layers:
652 - # Reconstruct the aux hidden state for draft models
653 - aux_recon = mhc_post_tilelang(
654 - hidden_states, residual, post_mix, res_mix
655 - )
656 - aux_hidden_state = aux_recon.mean(dim=1)
722 + if previous_aux is not None:
723 + # idx is the one-based id of the layer whose post this is.
724 if self.use_sequence_parallel:
658 - aux_hidden_state = sp_all_gather(aux_hidden_state)[:full_num_tokens]
659 - aux_hidden_states.append(aux_hidden_state)
660 - final_aux_recon = aux_recon
725 + previous_aux = sp_all_gather(previous_aux)[:full_num_tokens]
726 + aux_hidden_by_layer[idx] = previous_aux
727 if layer is not None:
662 - # Reuse if the last layer was captured as an aux hidden state
728 + # The last layer has no successor to fold its post into.
729 + hidden_states = mhc_post_tilelang(
730 + hidden_states, residual, post_mix, res_mix
731 + )
732 if self.end_layer in self.aux_hidden_state_layers:
664 - hidden_states = final_aux_recon
665 - else:
666 - hidden_states = mhc_post_tilelang(
667 - hidden_states, residual, post_mix, res_mix
668 - )
733 + final_aux = hidden_states.mean(dim=1)
734 + if self.use_sequence_parallel:
735 + final_aux = sp_all_gather(final_aux)[:full_num_tokens]
736 + aux_hidden_by_layer[self.end_layer] = final_aux
737 +
738 + aux_hidden_states = [
739 + aux_hidden_by_layer[layer_id]
740 + for layer_id in self.aux_hidden_state_layers
741 + if layer_id in aux_hidden_by_layer
742 + ]
743
744 if not get_pp_group().is_last_rank:
745 return IntermediateTensors(