[Perf][Kernel] Integrate Mega-mHC from DeepGEMM for DeepSeek V4.1 (reopen of #56255) (#56962)
Signed-off-by: Yongye Zhu <zyy1102000@gmail.com> Signed-off-by: Juntian Liu <Juntianl777@gmail.com> Co-authored-by: Yongye Zhu <zyy1102000@gmail.com> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com> Co-authored-by: Thien Tran <gau.nernst@yahoo.com.sg>
Juntian Liu committed
Sep 15, 2026 at 04:01 UTC
bb5507741155a698750bc422cda1f7c57f7d114d
6 files changed
+316
-41
tests/kernels/quantization/test_block_fp8.py
+6
-1
@@ -27,6 +27,7 @@ from vllm.platforms import current_platform
27
from vllm.utils.deep_gemm import (
28
fp8_gemm_nt,
29
get_tma_aligned_size,
30
+ is_deep_gemm_e8m0_used,
31
per_block_cast_to_fp8,
32
should_use_deepgemm_for_fp8_linear,
33
)
@@ -291,7 +292,11 @@ def test_w8a8_block_fp8_deep_gemm_matmul(M, N, K, block_size, out_dtype, seed):
292
A_fp8, As_fp8 = per_token_group_quant_fp8(
293
A_fp32, block_size[1], column_major_scales=True, tma_aligned_scales=True
294
)
294
- B_fp8, Bs_fp8 = per_block_cast_to_fp8(B_fp32, block_size=block_size)
295
+ B_fp8, Bs_fp8 = per_block_cast_to_fp8(
296
+ B_fp32,
297
+ block_size=block_size,
298
+ use_ue8m0=is_deep_gemm_e8m0_used(),
299
+ )
300
301
As = As_fp8.to(torch.float32)
302
Bs = Bs_fp8.to(torch.float32)
tests/kernels/test_mhc_kernels.py
+104
-1
@@ -11,6 +11,7 @@ import vllm.model_executor.kernels.mhc # noqa: F401
11
import vllm.model_executor.layers.mhc as mhc_layers
12
from vllm.model_executor.kernels.mhc.tilelang import (
13
_torch_hc_prenorm_gemm,
14
+ mhc_post_tilelang,
15
mhc_pre_delayed_tilelang,
16
)
17
from vllm.model_executor.kernels.mhc.tilelang_kernels import (
@@ -39,6 +40,10 @@ from vllm.models.deepseek_v4.nvidia.model import (
40
from vllm.models.deepseek_v41.nvidia.model import (
41
DeepseekV4DecoderLayer as DeepseekV41DecoderLayer,
42
)
43
+from vllm.models.deepseek_v41.nvidia.ops.mega_mhc import (
44
+ is_mega_mhc_supported,
45
+ mhc_shifted_post_pre_deep_gemm,
46
+)
47
from vllm.platforms import current_platform
48
from vllm.utils.torch_utils import set_random_seed
49
@@ -495,7 +500,7 @@ def test_deepseek_v41_decoder_mixes_match_torch(
500
reference,
501
)
502
monkeypatch.setattr(
498
- "vllm.models.deepseek_v41.nvidia.model.mhc_fused_post_pre_delayed_tilelang",
503
+ "vllm.models.deepseek_v41.nvidia.model.mhc_shifted_post_pre",
504
fused_reference,
505
)
506
expected = decoder(x, positions, None, **kwargs)
@@ -580,6 +585,104 @@ def test_mhc_pre_delayed_custom_op_supports_compile(carried):
585
)
586
587
588
+@pytest.mark.skipif(not HAS_TILELANG_MHC, reason="TileLang MHC support required")
589
+@pytest.mark.skipif(
590
+ not current_platform.is_device_capability_family(100),
591
+ reason="DeepGEMM Mega mHC requires SM100-family CUDA",
592
+)
593
+@pytest.mark.parametrize("num_tokens", [0, 1, 17, 128, 1024])
594
+@pytest.mark.parametrize("hidden_size", [5120, 7168])
595
+def test_deep_gemm_mega_mhc_correctness(num_tokens, hidden_size):
596
+ set_random_seed(0)
597
+ hc_mult = 4
598
+ if not is_mega_mhc_supported(hidden_size, 4):
599
+ pytest.skip("DeepGEMM Mega mHC is not available")
600
+
601
+ x = torch.randn(num_tokens, hidden_size, dtype=torch.bfloat16, device=DEVICE)
602
+ residual = torch.randn(
603
+ num_tokens, hc_mult, hidden_size, dtype=torch.bfloat16, device=DEVICE
604
+ )
605
+ previous_mix = torch.rand(num_tokens, hc_mult, device=DEVICE)
606
+ post_mix = torch.rand(num_tokens, hc_mult, 1, device=DEVICE)
607
+ res_mix = torch.rand(num_tokens, hc_mult, hc_mult, device=DEVICE)
608
+ fn = torch.randn(24, hc_mult * hidden_size, device=DEVICE) * 0.01
609
+ scale = torch.randn(3, device=DEVICE) * 0.1
610
+ base = torch.randn(24, device=DEVICE) * 0.1
611
+ norm_weight = torch.empty(
612
+ hidden_size, dtype=torch.bfloat16, device=DEVICE
613
+ ).uniform_(0.9, 1.1)
614
+
615
+ expected_residual = (
616
+ mhc_post_tilelang(x, residual, post_mix, res_mix) if num_tokens else residual
617
+ )
618
+ (
619
+ expected_post_mix,
620
+ expected_res_mix,
621
+ expected_y_bf16,
622
+ expected_previous_mix,
623
+ ) = mhc_pre_delayed_tilelang(
624
+ expected_residual,
625
+ fn,
626
+ scale,
627
+ base,
628
+ 2e-5,
629
+ 3e-4,
630
+ 2e-6,
631
+ 1.25,
632
+ 10,
633
+ pre_mix=previous_mix,
634
+ norm_weight=norm_weight,
635
+ norm_eps=7e-6,
636
+ )
637
+ args = (
638
+ x,
639
+ residual,
640
+ previous_mix,
641
+ post_mix,
642
+ res_mix,
643
+ fn,
644
+ scale,
645
+ base,
646
+ 2e-5,
647
+ 3e-4,
648
+ 1.25,
649
+ 2e-6,
650
+ 10,
651
+ norm_weight,
652
+ 7e-6,
653
+ )
654
+ actual = mhc_shifted_post_pre_deep_gemm(*args)
655
+ expected = (
656
+ expected_residual,
657
+ expected_post_mix,
658
+ expected_res_mix,
659
+ expected_y_bf16,
660
+ expected_previous_mix,
661
+ )
662
+ for i, (result, ref) in enumerate(zip(actual, expected, strict=True)):
663
+ # Match the existing mHC BF16 tolerance: post rounding differences
664
+ # can be amplified by cancellation in the carried collapse.
665
+ atol, rtol = (1.6e-2, 1e-2) if i in (0, 3) else (1e-6, 1e-3)
666
+ torch.testing.assert_close(result, ref, atol=atol, rtol=rtol)
667
+
668
+ if num_tokens:
669
+ stream = torch.cuda.Stream()
670
+ stream.wait_stream(torch.cuda.current_stream())
671
+ with torch.cuda.stream(stream):
672
+ # DeepGEMM initializes barriers per stream before capture.
673
+ mhc_shifted_post_pre_deep_gemm(*args)
674
+ graph = torch.cuda.CUDAGraph()
675
+ with torch.cuda.graph(graph, stream=stream):
676
+ captured = mhc_shifted_post_pre_deep_gemm(*args)
677
+ for multiplier in (0.5, -1.0):
678
+ x.mul_(multiplier)
679
+ graph.replay()
680
+ eager = mhc_shifted_post_pre_deep_gemm(*args)
681
+ for result, ref in zip(captured, eager, strict=True):
682
+ torch.testing.assert_close(result, ref, atol=0, rtol=0)
683
+ torch.cuda.current_stream().wait_stream(stream)
684
+
685
+
686
def sinkhorn_normalize_ref(x: torch.Tensor, repeat: int, eps: float) -> torch.Tensor:
687
x = x.softmax(-1) + eps
688
x = x / (x.sum(-2, keepdim=True) + eps)
vllm/models/deepseek_v41/nvidia/model.py
+34
-38
@@ -20,7 +20,6 @@ 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,
23
mhc_post_tilelang,
24
mhc_pre_delayed_tilelang,
25
)
@@ -79,6 +78,7 @@ from vllm.v1.worker.ubatching import dbo_current_ubatch_id
78
from ..common.engram import EngramLayout, NgramHashState
79
from ..common.mm_preprocess import IMAGE_SENTINEL_BASE_ID, image_sentinel_mask
80
from .engram import Engram, gather_engram_hashes
81
+from .ops.mega_mhc import mhc_shifted_post_pre
82
83
if typing.TYPE_CHECKING:
84
from vllm.v1.attention.backends.mla.sparse_swa import DeepseekSparseSWAMetadata
@@ -403,25 +403,23 @@ class DeepseekV4DecoderLayer(nn.Module):
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
- )
406
+ residual, post_mix, res_mix, x, attn_pre, aux = mhc_shifted_post_pre(
407
+ x,
408
+ residual,
409
+ post_mix,
410
+ res_mix,
411
+ self.hc_attn_fn,
412
+ self.hc_attn_scale,
413
+ self.hc_attn_base,
414
+ self.rms_norm_eps,
415
+ self.hc_eps,
416
+ self.hc_eps,
417
+ self.hc_post_alpha,
418
+ self.hc_sinkhorn_iters,
419
+ pre_mix=pre_mix,
420
+ norm_weight=self.attn_norm.weight,
421
+ norm_eps=self.attn_norm.variance_epsilon,
422
+ capture_aux=capture_previous_aux,
423
)
424
if capture_previous_aux:
425
previous_aux = aux
@@ -433,24 +431,22 @@ class DeepseekV4DecoderLayer(nn.Module):
431
if self.use_sequence_parallel:
432
x = sp_reduce_scatter(x)
433
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
- )
434
+ residual, post_mix, res_mix, x, ffn_pre, _ = mhc_shifted_post_pre(
435
+ x,
436
+ residual,
437
+ post_mix,
438
+ res_mix,
439
+ self.hc_ffn_fn,
440
+ self.hc_ffn_scale,
441
+ self.hc_ffn_base,
442
+ self.rms_norm_eps,
443
+ self.hc_eps,
444
+ self.hc_eps,
445
+ self.hc_post_alpha,
446
+ self.hc_sinkhorn_iters,
447
+ pre_mix=attn_pre,
448
+ norm_weight=self.ffn_norm.weight,
449
+ norm_eps=self.ffn_norm.variance_epsilon,
450
)
451
x = self.ffn(x, input_ids)
452
return x, residual, post_mix, res_mix, ffn_pre, previous_aux
vllm/models/deepseek_v41/nvidia/ops/__init__.py
new
+2
@@ -0,0 +1,2 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
vllm/models/deepseek_v41/nvidia/ops/mega_mhc.py
new
+156
@@ -0,0 +1,156 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import functools
5
+
6
+import torch
7
+
8
+from vllm.model_executor.kernels.mhc.tilelang import (
9
+ mhc_fused_post_pre_delayed_tilelang,
10
+)
11
+from vllm.platforms import current_platform
12
+from vllm.utils.deep_gemm import (
13
+ _import_deep_gemm,
14
+ is_deep_gemm_supported,
15
+ mega_mhc,
16
+)
17
+
18
+
19
+@functools.cache
20
+def is_mega_mhc_supported(hidden_size: int, hc_mult: int) -> bool:
21
+ """Return whether DeepGEMM can run the DSV4.1 shifted mHC kernel."""
22
+ if not (
23
+ is_deep_gemm_supported()
24
+ and current_platform.is_device_capability_family(100)
25
+ and hidden_size > 0
26
+ and hidden_size % 1024 == 0
27
+ and hc_mult == 4
28
+ ):
29
+ return False
30
+ deep_gemm = _import_deep_gemm()
31
+ return deep_gemm is not None and callable(getattr(deep_gemm, "mega_mhc", None))
32
+
33
+
34
+def mhc_shifted_post_pre_deep_gemm(
35
+ x: torch.Tensor,
36
+ residual: torch.Tensor,
37
+ shifted_prev_mix: torch.Tensor,
38
+ post_mix: torch.Tensor,
39
+ comb_res_mix: torch.Tensor,
40
+ fn: torch.Tensor,
41
+ mix_scales: torch.Tensor,
42
+ mix_bases: torch.Tensor,
43
+ hc_norm_eps: float,
44
+ hc_pre_eps: float,
45
+ hc_post_scale: float,
46
+ sinkhorn_eps: float,
47
+ num_sinkhorn_iters: int,
48
+ rmsnorm_weight: torch.Tensor,
49
+ rmsnorm_eps: float,
50
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
51
+ """Run DSV4.1 shifted post, next pre, and BF16 RMSNorm with Mega mHC."""
52
+ num_tokens, hidden_size = x.shape
53
+ hc_mult = residual.shape[1]
54
+ new_residual = torch.empty_like(residual)
55
+ new_prev_mix = x.new_empty(num_tokens, hc_mult, 1, dtype=torch.float32)
56
+ new_post_mix = torch.empty_like(post_mix)
57
+ new_comb_res_mix = torch.empty_like(comb_res_mix)
58
+ y_bf16 = x.new_empty(num_tokens, hidden_size, dtype=torch.bfloat16)
59
+ mega_mhc(
60
+ x=x,
61
+ residual=residual,
62
+ shifted_prev_mix=shifted_prev_mix.unsqueeze(-1),
63
+ post_mix=post_mix,
64
+ comb_res_mix=comb_res_mix,
65
+ fn=fn,
66
+ mix_scales=mix_scales,
67
+ mix_bases=mix_bases,
68
+ hc_mult=hc_mult,
69
+ hc_norm_eps=hc_norm_eps,
70
+ hc_pre_eps=hc_pre_eps,
71
+ hc_post_scale=hc_post_scale,
72
+ sinkhorn_eps=sinkhorn_eps,
73
+ num_sinkhorn_iters=num_sinkhorn_iters,
74
+ rmsnorm_weight=rmsnorm_weight,
75
+ rmsnorm_eps=rmsnorm_eps,
76
+ rmsnorm_scale=1.0,
77
+ new_residual=new_residual,
78
+ new_prev_mix=new_prev_mix,
79
+ new_post_mix=new_post_mix,
80
+ new_comb_res_mix=new_comb_res_mix,
81
+ y_bf16=y_bf16,
82
+ )
83
+ return (
84
+ new_residual,
85
+ new_post_mix,
86
+ new_comb_res_mix,
87
+ y_bf16,
88
+ new_prev_mix.squeeze(-1),
89
+ )
90
+
91
+
92
+def mhc_shifted_post_pre(
93
+ x: torch.Tensor,
94
+ residual: torch.Tensor,
95
+ post_layer_mix: torch.Tensor,
96
+ comb_res_mix: torch.Tensor,
97
+ fn: torch.Tensor,
98
+ hc_scale: torch.Tensor,
99
+ hc_base: torch.Tensor,
100
+ rms_eps: float,
101
+ hc_pre_eps: float,
102
+ hc_sinkhorn_eps: float,
103
+ hc_post_mult_value: float,
104
+ sinkhorn_repeat: int,
105
+ pre_mix: torch.Tensor | None = None,
106
+ norm_weight: torch.Tensor | None = None,
107
+ norm_eps: float = 1e-6,
108
+ capture_aux: bool = False,
109
+) -> tuple[
110
+ torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor
111
+]:
112
+ """Use Mega-mHC for carried mixing, retaining TileLang for aux capture."""
113
+ if (
114
+ pre_mix is not None
115
+ and norm_weight is not None
116
+ and not capture_aux
117
+ and x.shape[0] <= 1 << 20
118
+ and is_mega_mhc_supported(x.shape[1], residual.shape[1])
119
+ ):
120
+ outputs = mhc_shifted_post_pre_deep_gemm(
121
+ x,
122
+ residual,
123
+ pre_mix,
124
+ post_layer_mix,
125
+ comb_res_mix,
126
+ fn,
127
+ hc_scale,
128
+ hc_base,
129
+ rms_eps,
130
+ hc_pre_eps,
131
+ hc_post_mult_value,
132
+ hc_sinkhorn_eps,
133
+ sinkhorn_repeat,
134
+ norm_weight,
135
+ norm_eps,
136
+ )
137
+ return *outputs, x.new_empty(0, x.shape[1])
138
+
139
+ return mhc_fused_post_pre_delayed_tilelang(
140
+ x,
141
+ residual,
142
+ post_layer_mix,
143
+ comb_res_mix,
144
+ fn,
145
+ hc_scale,
146
+ hc_base,
147
+ rms_eps,
148
+ hc_pre_eps,
149
+ hc_sinkhorn_eps,
150
+ hc_post_mult_value,
151
+ sinkhorn_repeat,
152
+ pre_mix=pre_mix,
153
+ norm_weight=norm_weight,
154
+ norm_eps=norm_eps,
155
+ capture_aux=capture_aux,
156
+ )
vllm/utils/deep_gemm.py
+14
-1
@@ -159,6 +159,7 @@ _fp8_fp4_mqa_logits_impl: Callable[..., Any] | None = None
159
_fp8_fp4_paged_mqa_logits_impl: Callable[..., Any] | None = None
160
_get_paged_mqa_logits_metadata_impl: Callable[..., Any] | None = None
161
_tf32_hc_prenorm_gemm_impl: Callable[..., Any] | None = None
162
+_mega_mhc_impl: Callable[..., Any] | None = None
163
_get_mn_major_tma_aligned_tensor_impl: Callable[..., Any] | None = None
164
_get_mk_alignment_for_contiguous_layout_impl: Callable[..., Any] | None = None
165
_get_theoretical_mk_alignment_for_contiguous_layout_impl: Callable[..., Any] | None = (
@@ -231,7 +232,7 @@ def _lazy_init() -> None:
232
global _grouped_impl, _grouped_masked_impl, _grouped_fp4_impl
233
global _fp8_fp4_mqa_logits_impl, _fp8_fp4_paged_mqa_logits_impl
234
global _get_paged_mqa_logits_metadata_impl
234
- global _tf32_hc_prenorm_gemm_impl
235
+ global _tf32_hc_prenorm_gemm_impl, _mega_mhc_impl
236
global _get_mn_major_tma_aligned_tensor_impl
237
global _get_mk_alignment_for_contiguous_layout_impl
238
global _get_theoretical_mk_alignment_for_contiguous_layout_impl
@@ -251,6 +252,7 @@ def _lazy_init() -> None:
252
or _fp8_fp4_paged_mqa_logits_impl is not None
253
or _get_paged_mqa_logits_metadata_impl is not None
254
or _tf32_hc_prenorm_gemm_impl is not None
255
+ or _mega_mhc_impl is not None
256
or _get_mk_alignment_for_contiguous_layout_impl is not None
257
or _transform_sf_into_required_layout_impl is not None
258
or _pack_ue8m0_to_int_impl is not None
@@ -294,6 +296,7 @@ def _lazy_init() -> None:
296
_dg, "get_paged_mqa_logits_metadata", None
297
)
298
_tf32_hc_prenorm_gemm_impl = getattr(_dg, "tf32_hc_prenorm_gemm", None)
299
+ _mega_mhc_impl = getattr(_dg, "mega_mhc", None)
300
_get_mn_major_tma_aligned_tensor_impl = getattr(
301
_dg, "get_mn_major_tma_aligned_tensor", None
302
)
@@ -704,6 +707,15 @@ def tf32_hc_prenorm_gemm(
707
)
708
709
710
+def mega_mhc(*args: Any, **kwargs: Any) -> None:
711
+ """Run DeepGEMM Mega mHC with caller-owned output tensors."""
712
+ _lazy_init()
713
+ if _mega_mhc_impl is None:
714
+ _missing()
715
+ return
716
+ _mega_mhc_impl(*args, **kwargs)
717
+
718
+
719
def _ceil_to_ue8m0(x: torch.Tensor):
720
return torch.pow(2.0, torch.ceil(torch.log2(x.abs())))
721
@@ -797,6 +809,7 @@ __all__ = [
809
"per_block_cast_to_fp8",
810
"is_deep_gemm_e8m0_used",
811
"is_deep_gemm_supported",
812
+ "mega_mhc",
813
"get_num_sms",
814
"set_num_sms",
815
"should_use_deepgemm_for_fp8_linear",