[Bugfix] Package the SM90 DeepGEMM mHC register lifetime fix

[mengziheng]

Purpose

Package the SM90 mHC prenorm GEMM register-lifetime fix from deepseek-ai/DeepGEMM#448, authored by jason (leevan). In a multi-stage K loop, the next iteration can overwrite register A operands before the previous asynchronous WGMMA has finished. Repeated calls with identical inputs can therefore produce different outputs. The upstream fix waits immediately after committing the WGMMA batch.

This integrates that existing fix into vLLM's packaged JIT header until the pinned DeepGEMM dependency contains it; it is not a new kernel implementation. DeepGEMM #448 is currently open. Remove the temporary integration once the dependency pin includes the fix.

This correctness change was split out of vLLM #58986, which now contains only the H20 block-FP8 MoE tuning. This branch starts from main 07f1ba2ec66776650a490630ddfa9e3c92713c33 and does not depend on that performance PR. The corresponding v0.30 backport remains mengziheng/vllm#3. There is no throughput improvement claim for this correctness repair.

Changes

• Generate a patched header during CMake configuration and install it without modifying a supplied DEEPGEMM_SRC_DIR. Exclude the original header from the directory install to prevent a same-timestamp overwrite.
• Use the same helper in the standalone DeepGEMM installer. Recognized patched input is idempotent; unrecognized source fails without overwriting the destination.
• Extend the existing mHC suite with six SM90 regressions: 256/512/1024 tokens, different split-K counts, default and side streams, independent FP32 references and exact repeatability. Skip when DeepGEMM is unavailable.

Validation

The following GPU results were recorded on H20/v0.30 before this PR split. The main test path is tests/kernels/mhc/test_mhc_kernels.py; the v0.30 suite used tests/kernels/test_mhc_kernels.py.

.venv/bin/python -m pytest tests/kernels/test_mhc_kernels.py \
  -k 'sm90_mhc_prenorm_preserves_register_operands or test_mhc_fused_post_pre' -v
.venv/bin/python -m pytest tests/kernels/test_mhc_kernels.py \
  -k sm90_mhc_prenorm_preserves_register_operands -v

• The pinned original header fails the new regression; the fixed header passes all six cases. The final test including its DeepGEMM availability guard returned 6 passed, 212 deselected.
• The broader selection returned 17 passed, 9 skipped, 1 failed. The existing H=7168/M=128 BF16 reference case fails with both original and repaired headers. No tolerance was loosened; this is not a full-suite pass.
• A fresh-process audit of the complete mHC operation covers 12 shape/seed groups with five repeats each. Repaired results are exactly repeatable; the original produces inconsistent results at 256/512 tokens. This does not establish whole-model token determinism.
• An installation fixture using the actual CMake integration verified installed bytes equal the upstream fix, repeated configure/install preserves it, and supplied source content remains unchanged. Standalone idempotence and rejection of unknown input without destination overwrite also passed. A full wheel build was not performed.

Fresh local checks for this split:

pre-commit run --from-ref 07f1ba2ec66776650a490630ddfa9e3c92713c33 --to-ref HEAD
git diff --check 07f1ba2ec66776650a490630ddfa9e3c92713c33 HEAD

All applicable pre-commit checks pass. The three integration/helper files are byte-identical to the previously tested implementation, and the new regression's AST is identical. A local helper check again passed original-to-fixed conversion, idempotence, source preservation and unknown-input rejection. The generated header SHA256 is b90beeaed11e279d3451ad3d438b8a64a41fd003c4b7f4d5e0a8fb79a20ad80b. GPU tests were not rerun on this latest-main branch; these checks establish split equivalence, not a full main-wheel or full-model validation.

Model evaluation

Historical validation uses the same GLM-5.3-Flash checkpoint on 4×H20, TP4, FP8 weights/BF16 activations and MTP3, with v0.30 native extensions and the patched JIT header. The corrected control without the MoE tuning and the corrected candidate with the tuning each score 121/128 on a GSM8K first-128, five-shot raw-completion diagnostic (greedy, concurrency 8, output cap 768). Each has 128 successful requests and no truncation. Only 39/128 output texts match exactly; each arm uniquely answers three questions correctly. This subset is not an accuracy-improvement claim or proof of full-model equivalence.

The corrected control completed 1,152/1,152 measured serving requests successfully; all three recorded control/candidate blocks completed 3,456/3,456 in total. The corrected candidate also passed six context-retrieval checks at 511/2047/2048/2049/4097/8193 tokens. Full serving settings, paired measurements and limitations remain in #58986. Those performance differences belong to the MoE tuning, with this fix present in both arms.

B200, production traffic, long soak tests and full-model serving on latest main were not validated.

Related work

The open-PR duplicate search was refreshed before splitting. #50645 and #53055 address missing-DeepGEMM fallbacks, #52165 changes detection, #48619 ports the SM100 kernel to CuTeDSL, and #41834/#42856 concern SM12x support. #55184 repairs sparse MLA on SM8x; #54955 is an opt-in batch-invariant execution preview using explicit mHC operations. Neither addresses packaging this register-lifetime fix in the default SM90 path. No duplicate integration was found. The same code has been removed from #58986 so the two community PRs do not duplicate it. No second DeepGEMM fix PR is proposed.

AI assistance was used for investigation, integration, tests and this description. Human review has not been asserted.

[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.

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

PRs do not trigger a full CI run by default. Reviewers with write access and configured trusted contributors can comment /ci run for upstream CI or /amd-ci run for AMD CI only whenever CI signals are needed.

Once the PR is approved or has the ready label, the PR author can also use the corresponding /ci run, /ci retry, and /ci cancel commands, or their /amd-ci variants. New commits do not start upstream CI automatically.

If you have any questions, please reach out to us on Slack at https://slack.vllm.ai.

Agent Guidelines

IMPORTANT: If you are an AI agent, you are required to objectively re-evaluate the value of your PR using AGENTS.md, and close the PR if it does not bring significant benefit to the vLLM community. Failure to do so may result in an immediate ban.

🚀