[WIP][Kernel][DSv4.1] Add compensated BF16 mHC prenorm for SM90

[xijiaat]

Purpose

Draft a compensated BF16 mHC prenorm primitive for Hopper, adapted from SGLang's hc_mix_stats_bf16x3 in sgl-project/sglang#39664 and sgl-project/sglang#41251. It splits the FP32 mixing weight into three BF16 components, accumulates their projections separately in FP32, and computes the input squared norm in the same pass.

This initial draft contains a kernel, correctness tests and a comparison benchmark. Serving dispatch is intentionally not wired yet. No kernel or end-to-end speedup, numerical equivalence, or production-readiness claim is made. Integration and default selection depend on real SM90 measurements against the current fused mHC path, not only against a standalone GEMM.

Implementation

• Accept caller-owned prepared weights and output buffers; no allocation or CPU/GPU synchronization in the projection wrapper.
• Write [splits, tokens, mix] projection and [splits, tokens] squared-norm partials, compatible with the existing prenorm epilogue layout.
• Support split counts 1/4/16, token tiles 32/64/128, empty batches, masked tails, and contiguous offset views. These are tuning candidates, not validated performance defaults.
• Reuse the existing mHC correctness suite. Add FP64 reference, cancellation/weight-reconstruction, output-boundary and changed-input/weight CUDA graph tests.
• Add a CUPTI benchmark that checks correctness before timing and sweeps both DeepGEMM and candidate configurations with CUDA graphs and cold L2. Weight preparation, allocations and the common epilogue are excluded from the reported prenorm timing.

The prepared weights occupy 6 bytes per original FP32 element, in addition to any retained original weight. The caller must rebuild them after weight updates. Different accumulation orders can change outputs; batch invariance is not promised. SGLang's bundled serving gains are not an estimate for this isolated kernel or for vLLM.

Duplicate check and provenance

Checked #57448 and public open/closed PR searches for mHC, compensated, bf16x3, SM90, and Hopper on 2026-10-09. No open PR implementing this compensated mHC projection was found.

• [WIP][DSv4.1] Opt in to mHC coefficient overlap on SM90 #60768 only opts SM90 into the existing TF32 coefficient-overlap path; this is an independent kernel draft, not a duplicate or dependency.
• [Perf][DSV4.1] Overlap mHC coefficients for small TP batches #57603 implements small-batch overlap; [Perf][DSv4.1] Fold the mHC post block into the delayed pre projection #56633 fuses post/pre; [Perf][Kernel] Integrate Mega-mHC from DeepGEMM for DeepSeek V4.1 (reopen of #56255) #56962 integrates Mega-mHC. They are relevant baselines, not this BF16x3 projection.
• [Perf] Faster BF16x3 router GEMM at large batch sizes on SM100 #60376 is a router-GEMM optimization on SM100, not mHC prenorm on SM90.
• [Bugfix] Package the SM90 DeepGEMM mHC register lifetime fix #59129 concerns the separate DeepGEMM SM90 correctness fix and must be accounted for when benchmarking the baseline.

The SGLang kernel's Apache-2.0 attribution is retained. Reviewed source: python/sglang/kernels/ops/layernorm/mhc.py on SGLang main around 5d8e98b7fdd83907616b1ff8ca68229f2ae95b34. This adapts its numerical approach to vLLM's output layout and caller-owned workspace; it does not copy SGLang's runtime dispatch or Sinkhorn implementation.

Validation completed

CPU-only Linux container, Python 3.12.3, PyTorch 2.13.0+cu130, Triton 3.7.1. Source under test was overlaid into an isolated source checkout; dependencies came from the existing vLLM v0.28.0 image, not a newly built wheel. Container used runc, no GPU device nodes, empty CUDA_VISIBLE_DEVICES, NVIDIA_VISIBLE_DEVICES=void, no network, 2 CPU cores and 4 GiB RAM. torch.cuda.is_available() was false.

VLLM_TARGET_DEVICE=cpu PYTHONPATH=$PWD PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 \
  .venv/bin/python -m pytest --confcutdir=tests/kernels/mhc -o addopts= \
  -q tests/kernels/mhc/test_mhc_kernels.py -k bf16x3

3 passed, 9 GPU cases skipped, 320 deselected. CPU checks cover random, cancellation and scaled FP32 weight decomposition, including preserving the original weight. Source-only version/extension warnings and torch deprecation warnings were emitted.

18 SM90 offline compilations passed: explicit GPUTarget("cuda", 90, 32), K=5120/20480, splits=1/4/16, BLOCK_M=32/64/128, N=24, BLOCK_N=32, BLOCK_K=64, four warps and three stages. PTX/cubin were generated without loading or launching a CUDA kernel. This proves compilation only, not GPU numerical correctness or speed.

All applicable pre-commit hooks passed, including Ruff, mypy, SPDX and Buildkite test tethering; git diff --check passed. The existing mHC CI job exercises the suite on B200; actual SM90 execution remains a validation gap.

Before ready for review

• Run the 9 CUDA correctness/replay cases on SM90, including bounds checking.
• Run benchmarks/kernels/benchmark_mhc_bf16x3.py on H20/H200 and retain all configurations, not only winners.
• Integrate opt-in dispatch with weight-update handling and startup warmup only if the kernel is competitive against the actual fused post/pre baseline.
• Compare the complete mHC boundary and full serving performance; evaluate model quality after integration.
• Complete human line-by-line review.

AI assistance: implemented and checked with OpenAI Codex. This draft was opened at the human author's explicit request to stage the implementation while GPU resources are being arranged. GPU execution, model evaluations, serving benchmarks and human review are pending; this is not ready to merge.

[xijiaat]

Closing this draft after H20 (SM90) validation. The current implementation is not a performance win in the tested shapes, and its original numerical suite is not fully passing.

• GPU suite: 11 passed, 1 failed. The failing case is test_bf16x3_prenorm_partials_match_fp64[33-20480-1-64], with 4/792 projection elements outside the unchanged tolerance (rtol=5e-5, atol=2e-5).
• A projection-only screen over M={1,32,256,8192}, K={5120,20480}, splits={1,4,16}, and candidate BLOCK_M={32,64,128} found no speedup. After filtering candidate configurations for precision and repeatability, the best BF16x3 configurations were approximately 1.69–3.24x slower than the tuned DeepGEMM baseline in hot-cache CUDA Event measurements. Explicit L2-flush retests were also slower.
• For example, M=8192/K=20480 measured 763.1 us for BF16x3 versus 250.8 us for DeepGEMM. The baseline included the SM90 register-A lifetime correction; baseline TF32 error was reported separately from the stricter candidate precision gate.
• These are isolated kernel measurements using the PR source with an existing vLLM image's native dependencies, not an end-to-end serving evaluation. No conclusion about model quality or other GPU architectures is implied.

Serving dispatch was never enabled. I will retain the branch and experimental results, but will not pursue integration of this version without a materially better implementation. Thank you.