[Bugfix][DSV4.1] Restrict mHC overlap to full CUDA graphs (#57874)
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai> Co-authored-by: OpenAI Codex <noreply@openai.com>
Woosuk Kwon committed
Sep 21, 2026 at 01:16 UTC
67513c8b67e2439934ce7c0ba5a18c2e859f5c53
4 files changed
+45
-11
tests/kernels/test_mhc_kernels.py
+21
-2
@@ -539,7 +539,7 @@ def test_mhc_fused_post_pre_delayed_custom_op_supports_compile(carried, capture_
539
@pytest.mark.skipif(not HAS_TILELANG_MHC, reason="TileLang MHC support required")
540
@pytest.mark.parametrize("entry", ["broadcast", "pipeline", "residual", "engram"])
541
@pytest.mark.parametrize(
542
- "mhc_mode", ["disabled", "overlap", "piecewise", "large_batch"]
542
+ "mhc_mode", ["disabled", "eager", "overlap", "piecewise", "large_batch"]
543
)
544
def test_deepseek_v41_decoder_mixes_match_torch(
545
entry, mhc_mode, monkeypatch, default_vllm_config
@@ -568,7 +568,7 @@ def test_deepseek_v41_decoder_mixes_match_torch(
568
):
569
pytest.skip("SM100 DeepGEMM required for overlap")
570
decoder.mhc_stream = torch.cuda.Stream()
571
- if mhc_mode in ("piecewise", "large_batch"):
571
+ if mhc_mode in ("eager", "piecewise", "large_batch"):
572
573
def unexpected_overlap(*args, **kwargs):
574
pytest.fail("unsupported execution must retain native mHC")
@@ -624,8 +624,27 @@ def test_deepseek_v41_decoder_mixes_match_torch(
624
from vllm.forward_context import set_forward_context
625
626
mode = CUDAGraphMode.PIECEWISE if mhc_mode == "piecewise" else CUDAGraphMode.NONE
627
+ overlap_calls = []
628
+ if mhc_mode == "overlap":
629
+ from vllm.models.deepseek_v41.nvidia.ops.mhc import mhc_pre_delayed_overlap
630
+
631
+ def checked_overlap(*args, **kwargs):
632
+ overlap_calls.append(torch.cuda.is_current_stream_capturing())
633
+ return mhc_pre_delayed_overlap(*args, **kwargs)
634
+
635
+ for module in ("model", "ops.mhc"):
636
+ monkeypatch.setattr(
637
+ f"vllm.models.deepseek_v41.nvidia.{module}.mhc_pre_delayed_overlap",
638
+ checked_overlap,
639
+ )
640
with set_forward_context(None, default_vllm_config, cudagraph_runtime_mode=mode):
641
actual = decoder(x, positions, None, **kwargs)
642
+ if mhc_mode == "overlap":
643
+ graph = torch.cuda.CUDAGraph()
644
+ with torch.cuda.graph(graph):
645
+ actual = decoder(x, positions, None, **kwargs)
646
+ graph.replay()
647
+ assert overlap_calls and all(overlap_calls)
648
649
def reference(*args, norm_weight, norm_eps, **kwargs):
650
post, res, collapsed, pre = mhc_pre_delayed_torch(*args, **kwargs)
vllm/models/deepseek_v41/nvidia/model.py
+2
-1
@@ -385,8 +385,9 @@ class DeepseekV4DecoderLayer(nn.Module):
385
if mhc_stream is not None and (
386
in_piecewise_cudagraph()
387
or not 0 < positions.shape[0] <= MHC_OVERLAP_MAX_TOKENS
388
+ or not torch.cuda.is_current_stream_capturing()
389
):
389
- # A side stream cannot remain unjoined across breakable graph segments.
390
+ # Use overlap only in FULL graphs; eager warmup initializes Mega mHC.
391
mhc_stream = None
392
mhc_pre = (
393
partial(mhc_pre_delayed_overlap, stream=mhc_stream)
vllm/models/deepseek_v41/nvidia/ops/mega_mhc.py
+16
@@ -28,6 +28,22 @@ def is_mega_mhc_supported(hidden_size: int, hc_mult: int) -> bool:
28
return deep_gemm is not None and callable(getattr(deep_gemm, "mega_mhc", None))
29
30
31
+def can_use_mega_mhc(
32
+ x: torch.Tensor,
33
+ residual: torch.Tensor,
34
+ pre_mix: torch.Tensor | None,
35
+ norm_weight: torch.Tensor | None,
36
+ capture_aux: bool,
37
+) -> bool:
38
+ return (
39
+ pre_mix is not None
40
+ and norm_weight is not None
41
+ and not capture_aux
42
+ and x.shape[0] <= 1 << 20
43
+ and is_mega_mhc_supported(x.shape[1], residual.shape[1])
44
+ )
45
+
46
+
47
def mhc_shifted_post_pre_deep_gemm(
48
x: torch.Tensor,
49
residual: torch.Tensor,
vllm/models/deepseek_v41/nvidia/ops/mhc.py
+6
-8
@@ -16,7 +16,10 @@ from vllm.model_executor.kernels.mhc.tilelang import (
16
from vllm.platforms import current_platform
17
from vllm.utils.deep_gemm import is_deep_gemm_supported
18
19
-from .mega_mhc import is_mega_mhc_supported, mhc_shifted_post_pre_deep_gemm
19
+from .mega_mhc import (
20
+ can_use_mega_mhc,
21
+ mhc_shifted_post_pre_deep_gemm,
22
+)
23
24
if TYPE_CHECKING:
25
from vllm.distributed.device_communicators.cuda_communicator import CudaCommunicator
@@ -223,13 +226,8 @@ def mhc_shifted_post_pre(
226
)
227
return residual, *pre_outputs, aux
228
226
- if (
227
- pre_mix is not None
228
- and norm_weight is not None
229
- and not capture_aux
230
- and x.shape[0] <= 1 << 20
231
- and is_mega_mhc_supported(x.shape[1], residual.shape[1])
232
- ):
229
+ if can_use_mega_mhc(x, residual, pre_mix, norm_weight, capture_aux):
230
+ assert pre_mix is not None and norm_weight is not None
231
outputs = mhc_shifted_post_pre_deep_gemm(
232
x,
233
residual,