[DSv4.1] Integrate Mega-mHC from DeepGEMM

[gau-nernst]

Purpose

Integrate Mega-mHC from DeepGEMM (deepseek-ai/DeepGEMM#432)

• Temporarily change DeepGEMM commit to main branch. This removes sm120 support. We will add back sm120 support once nv_dev branch has rebased onto latest main.
• Add elfutils-devel to Dockerfile since DeepJIT depends on it
• New DeepGEMM no longer allows FP32 scales on SM100, hence the change in tests/kernels/quantization/test_block_fp8.py is necessary
• The Mega-mHC doesn't support engram. Hence, when there is engram, it uses the existing 2 tilelang kernels.
• Technically Mega-mHC supports DSv4 and GLM5-Next as well (non-shift mode), but this PR only hooks it up to DSv4.1

Test Plan

Updated unit tests in tests/kernels/test_mhc_kernels.py

GSM8K on GB300 TP4, 1 epoch

    "exact_match/flexible-extract": 0.9295,
    "exact_match/strict-match": 0.9303

Test Result

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[ywang96]

Thanks for the PR — the integration is cleanly done. I reviewed it against the DeepGEMM binding at the pinned commit (csrc/apis/mega_mhc.hpp @ 39d8c4c), the surrounding DSv4.1 model code, and vLLM's cudagraph capture flow. Verdict: approve with comments — the op wiring itself is correct; the blockers are the DeepGEMM pin bump's blast radius and one test-robustness gap.

Correctness findings

Major

1. SM120 regression is global, not DSv4.1-scoped (vllm/platforms/cuda.py:705-710, cmake/external_projects/deepgemm.cmake:75-79). Dropping is_device_capability_family(120) from support_deep_gemm() silently disables all DeepGEMM paths (FP8 block GEMM, paged MQA logits, …) on SM120 parts (RTX 5090/PRO 6000, DGX Spark), not just mHC. Understood this is temporary pending the nv_dev rebase — please add a TODO with a tracking issue, and make sure this never propagates to main as-is.

2. The pin bump changes behavior for every DeepGEMM consumer. 2.8.0 drops FP32-scale support on SM100 (hence the use_ue8m0=is_deep_gemm_e8m0_used() fix in tests/kernels/quantization/test_block_fp8.py:295) and restructures the _C API surface. vLLM's lazy getattr(…, None) wrappers convert any renamed/removed symbol into a runtime _missing() error on a hot path rather than an import failure, and only the one test that broke was updated. Can you confirm which CI lanes exercise the pre-existing DeepGEMM paths (SM90 and SM100) against the new pin?

Minor

3. CUDA-graph capture contract is satisfied today but fragile. mega_mhc lazily creates per-stream split barriers and hard-asserts CaptureStatus::None if first invoked on a stream during capture (DeepGEMM csrc/apis/mega_mhc.hpp:26-39; DeepGEMM's own test_mega_mhc.py asserts capture fails without a same-stream warmup). vLLM captures with stream=current_stream() (vllm/compilation/cuda_graph.py:313-316) and the profile/warmup dummy runs execute on that same stream, so this works — but any future flow that captures on a fresh stream dies at startup with a cryptic assert. A one-line comment in ops/mega_mhc.py ("first call on a stream must happen outside graph capture") would document the invariant.

4. New test is missing the tilelang guard (tests/kernels/test_mhc_kernels.py:321-397). test_deep_gemm_mega_mhc_correctness computes its reference via mhc_post_tilelang/mhc_pre_delayed_tilelang but only skips on SM100/DeepGEMM availability; every other test in the file also gates on HAS_TILELANG_MHC. On an SM100 machine with DeepGEMM but without tilelang this fails with ImportError instead of skipping. Please add @pytest.mark.skipif(not HAS_TILELANG_MHC, …).

5. mhc_shifted_post_pre is not a registered custom op, unlike the tilelang kernels it dispatches to (direct_register_custom_op, vllm/model_executor/kernels/mhc/tilelang.py:959-975). Fine today since DSv4.1 doesn't go through torch.compile, but if the model ever opts into piecewise compilation the mega path graph-breaks while the fallback doesn't. Worth a comment in the file.

Nit

6. vllm/models/deepseek_v4_1/nvidia/model.py:363 — assert pre_mix is not None converts a case tilelang previously tolerated (None → select residual stream zero) into a hard failure. It's unreachable in current flows (once residual is set, pre_mix always is too), so it's fine as an invariant — but since the mega path would otherwise die with an opaque AttributeError on None.unsqueeze, consider asserting inside mhc_shifted_post_pre itself rather than relying on callers.

Verified OK (checked, no action needed)

• Keyword names and semantics in the mega_mhc(...) call match the pinned pybind signature exactly, including the swapped hc_post_scale/sinkhorn_eps tail vs. tilelang's ordering; the test's distinct eps values (2e-5/3e-4/2e-6/1.25) would catch a swap.
• The support gate mirrors the kernel's host asserts: hidden % 1024 == 0 (= kDefaultNumSplits × BLOCK_K), hc_mult == 4 (= kNumRoutes), SM100-family check matches DG_HOST_ASSERT(arch_major == 10); kNumMaxTokens = 2^20 is a non-issue.
• Output shapes/dtypes (new_prev_mix.squeeze(-1) → [T, hc] FP32, etc.) match the tilelang contract; engram gating preserves the old post→engram→pre ordering; the tilelang fallback is argument-for-argument identical to the old call sequence.
• elfutils-devel is build-stage-only, but DeepJIT dlopens libdw.so.1 lazily and degrades gracefully, so the Ubuntu runtime stage doesn't need it; DeepJIT's runtime nvcc requirement is covered by cuda-nvcc-* in the runtime image (docker/Dockerfile:798).
• git submodule update --init --recursive in tools/install_deepgemm.sh correctly syncs the changed submodule set (fmt→deep_jit) after checkout.

Style findings

Major

1. Docstrings below house standard on the new public ops (vllm/models/deepseek_v4_1/nvidia/ops/mega_mhc.py:36,94 — one-liners at :53 and :110 for 16/14-parameter functions, vs. the full Args:/Returns: with shapes and dtypes on the kernels they wrap, vllm/model_executor/kernels/mhc/tilelang.py:90-127,622-648). This matters more than usual here because the wrapper renames every parameter across the boundary (rms_eps→hc_norm_eps, hc_post_alpha→hc_post_scale, norm_eps→rmsnorm_eps, pre_mix→shifted_prev_mix, res_mix→comb_res_mix); without arg docs the semantics require opening DeepGEMM's C++ header. The 5-tuple return order is also undocumented.

Minor

2. Commented-out dead code: vllm/platforms/cuda.py:709 leaves # or cls.is_device_capability_family(120) dangling inside a parenthesized return with no explanation; cmake/external_projects/deepgemm.cmake:75-79 comments out four lines (that one at least has a restore note at :75). Prefer a TODO(<issue>) tied to the nv_dev rebase, or deletion.
3. A third naming vocabulary for the same quantities: model attrs say hc_post_alpha/hc_attn_scale/pre_mix, tilelang says hc_post_mult_value/hc_scale/pre_mix, the new module + DeepGEMM say hc_post_scale/mix_scales/shifted_prev_mix — and the two halves of the new module disagree with each other (previous_mix/res_mix/scale at mega_mhc.py:94-109 forwarding into shifted_prev_mix/comb_res_mix/mix_scales at :36-52).
4. 16-argument positional forwarding at the hazard point (mega_mhc.py:112-128): the float tail of mhc_shifted_post_pre_deep_gemm (hc_post_scale before sinkhorn_eps) is ordered differently from the tilelang fallback right below it — exactly where an order-swap regression would hide. The DeepGEMM call above (:61-83) is all-keyword; make the internal call keyword-based too, at least for the scalar tail.
5. Private cross-module import: from vllm.utils.deep_gemm import _import_deep_gemm (mega_mhc.py:15). The public has_deep_gemm from vllm.utils.import_utils (already used at vllm/utils/deep_gemm.py:115) plus a direct import deep_gemm would express the same availability check.
6. Wrapper deviates from file idiom (vllm/utils/deep_gemm.py:706-712): mega_mhc(*args: Any, **kwargs: Any) -> None sits among fully-typed wrappers (e.g. tf32_hc_prenorm_gemm, :680), so a 20-keyword signature gets no mypy checking at call sites; and _missing(); return leaves a dead statement — every other wrapper writes return _missing().

Nit

7. tests/kernels/test_mhc_kernels.py:330 — is_mega_mhc_supported(hidden_size, 4) uses a literal despite hc_mult = 4 bound on the line above.
8. Dual import of the same module (test_mhc_kernels.py:12 module-form + :37-40 from-form) — needed for the monkeypatch target, but call sites could uniformly use mega_mhc_ops.*.
9. Bare positional eps literals (2e-5, 3e-4, 2e-6, 1.25) in the new test (:355-383) — a short comment naming which is which protects the exact ordering property the test verifies.

Checked and cleared on style: module placement follows existing family precedent (deepseek_v4/nvidia/ops/o_proj.py, imported cross-model by deepseek_v4_1/nvidia/flashmla.py); import order is isort-clean; nothing trips the repo's ruff rule set or the 88-char limit (longest added line is exactly 88); mypy check_untyped_defs is satisfied; license-only ops/__init__.py and SPDX headers match convention; no helper duplication (the existing mhc_fused_post_pre_tilelang has no carried pre_mix, so composing via mhc_pre_delayed_tilelang is correct).

Merge guidance

Blocking for me: the SM120 TODO/tracking (finding 1), confirmation that CI covers the pre-existing DeepGEMM paths on the new pin (finding 2), and the HAS_TILELANG_MHC guard on the new test (finding 4). Everything else — docstrings, keyword-forwarding, the capture-stream comment, naming — is worthwhile but can land as follow-ups if you're in a hurry.

[ivanusto]

Thanks for the Mega-mHC integration. The is_mega_mhc_supported() gate with the TileLang fallback on the V4.1 path looks right. I want to flag what the pin move does to SM12x before this lands, from 2x DGX Spark (GB10, sm_121), where DeepSeek-V4-Flash currently serves on the nv_dev pin.

What changes for SM12x in this PR

• The pin moves from nv_dev 8b1392b to DeepGEMM main. DeepGEMM main has no SM12x sources, and nv_dev (latest commit 572557e, 09-08) is 108 commits ahead of main and 1 behind, so it has not been rebased onto the 26/09 release yet.
• deepgemm.cmake drops 12.0f, 12.0a and 12.1a from DEEPGEMM_SUPPORT_ARCHS.
• support_deep_gemm() comments out is_device_capability_family(120).

Paths on current main that still call DeepGEMM on SM12x regardless of that gate

1. mhc_pre_broadcast_tilelang passes use_tilelang_fallback=False to _hc_prenorm_gemm_outputs, so use_deep_gemm = is_deep_gemm_supported() or not use_tilelang_fallback is always true and tf32_hc_prenorm_gemm is called unconditionally. The TileLang GEMM asserts x.shape[1] == hc_mult * hidden_size, so the (T, H) broadcast input cannot take that fallback. [Bugfix][Kernels] Guard DeepGEMM in mhc_pre_broadcast_tilelang with a torch fallback #53055 is adding a torch fallback for exactly this call.
2. deep_gemm_fp8_o_proj calls fp8_einsum without consulting is_deep_gemm_supported(), from both _o_proj call sites in vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py.

With no SM12x kernels, both of those are left with nothing to call.

What happened when only the gate was turned off on sm_121

I have not built this PR, but I did measure the support_deep_gemm() half in isolation, in an earlier revision of #53055: 2x GB10, TP=2, DeepSeek-V4-Flash-0731, nightly 0.26.1rc1.dev1219 on pin 8b1392b.

configuration	result
unmodified	serves, all 984 DeepGEMM warmup kernels pass, about 27 tok/s decode, coherent output
sm_121 excluded from support_deep_gemm()	engine init fails in profile_run: the CUTLASS block kernel receives E8M0 scales, scaled_mm_helper.hpp:17
same, plus VLLM_DISABLED_KERNELS=CutlassFp8BlockScaledMMKernel	fails in deep_gemm_fp8_o_proj: fp8_einsum asserts layout.hpp:40: t.dim() == N

So on GB10 today, DeepGEMM from nv_dev is what makes V4-Flash serve, and this PR removes both the kernels and the gate at once.

Ask

Could SM12x keep a working path while the rebase is pending? The cmake comment already anticipates restoring SM120 once the pin is on nv_dev, so holding the arch and gate removal until then, or keeping a separate pin for 12.x, would avoid a gap. Happy to test any variant on 2x GB10.

[gau-nernst]

@ivanusto We are in the process of rebasing nv_dev DeepGEMM onto latest DeepGEMM to preserve sm120 support. This PR won't be merged until then, so don't worry

[ivanusto]

@gau-nernst thank you, that is good to hear. When the rebased nv_dev is up, I am happy to run it on 2x GB10 (sm_121): build, DeepGEMM warmup, and a DeepSeek-V4-Flash TP=2 serve. Just ping me with the commit.