[Perf] Fuse the glm5_next mHC attn->MLP boundary

[mmangkad]

Motivation

GLM-5.3-Flash crosses 90 mHC post-then-pre boundaries per forward (45 layers, twice each). Each ran as three launches — mhc_post, the tf32 pre-norm GEMM, mhc_pre_big_fuse — which at a 6-token verify batch is launch cost, not work.

mhc_fused_post_pre and apply_mhc_post_pre_boundary already exist and are live for DeepSeek-V4. glm5_next never called them: it routes through MHCLayerCommunicator, not DeepSeek-V4's layer code.

Modifications

MHCState.hc_ffn_post_pre	optional callable tried first in attn_to_mlp; None keeps the unfused chain
Glm5NextDecoderLayer.hc_ffn_post_pre	GLM adapter onto the shared dispatcher
_MHC_FUSED_BOUNDARY_MAX_TOKENS = 16	conservative cap; declines above it

Folds mhc_post into the pre-norm GEMM — mhc_pre_big_fuse still launches separately, so two launches instead of three. 45 of 90 boundaries fuse; the rest need DeepSeek-V4's cross-layer deferral. fn_transpose=True matches DeepSeek-V4's call sites but is read only by the ROCm Triton tier, so it is unexercised on CUDA. Scope is glm5_next with mhc=True (config default is False).

Why 16

CUDA-graph replays, GLM-5.3-Flash's hc_mult=4 / hidden_size=4096:

tokens	≤6	8–16	24	32	≥33
DeepGEMM on	1.40x	1.17x	1.00x	0.88x	1.00x
DeepGEMM off	1.37x	1.17x	1.00x	0.87x	0.21x

Past the crossover the fused pre-norm GEMM drops the split-K kernel mhc_pre uses to 2048 tokens. 16 is conservative — parity is at 24, 17–23 unmeasured. The gate is in the GLM callback, not the shared dispatcher, which also serves DeepSeek-V4 at hc_hidden_size=28672 — these numbers cover one shape on one hardware config.

Verify batch is bs * speculative_num_draft_tokens: c=1 is 6 tokens and fuses, c=64 is 384 and declines.

Accuracy Tests

GPQA Diamond, thinking on: 92.42% (183/198), 0 truncated, 0 errors. temperature=1.0, top_p=0.95, max_tokens=131072, 1 repeat.

Speed Tests and Profiling

GLM-5.3-Flash, tp4 on 4x GB300, 8192 in / 1024 out at c=1. 12 runs: 3 per arm across 4 server sessions, arms interleaved. The metric is verify steps per second, tok/s divided by accept length — the decode-loop rate, which is what a per-forward change moves. Raw tok/s scales with accept length, which clusters per server process.

session	fused	unfused
1	79.07	78.55	+0.66%
2	79.82	77.99	+2.34%

Faster in both, with each session's run ranges disjoint.

At c=64 the gate declines and throughput is unchanged: 3466.9 vs 3456.3 tok/s.

Checklist

• pre-commit clean
• 28 cases in test_mhc_kernels.py; reverting communicator_mhc.py fails 8, moving the cutoff to 64 or 4 fails 2 each
• Accuracy and speed measured above

---

CI States

Latest PR Test (Base): ❌ Run #34830061832
Latest PR Test (Extra): ❌ Run #34830061740
Latest PR Test (AMD ROCm 10): ❌ Run #34830061928

[mmangkad]

/rerun-test test/registered/kernels/ops/layernorm/test_mhc_kernels.py test/registered/e2e/models/test_glm53_flash_b200.py test/registered/e2e/models/test_glm53_flash_h200.py

[github-actions]

Results for /rerun-test test/registered/kernels/ops/layernorm/test_mhc_kernels.py test/registered/e2e/models/test_glm53_flash_b200.py test/registered/e2e/models/test_glm53_flash_h200.py:

🚀 1-gpu-h100 (1 test): ✅ View workflow run

cd test/ && python3 registered/kernels/ops/layernorm/test_mhc_kernels.py

🚀 4-gpu-b200 (1 test): ✅ View workflow run

cd test/ && python3 registered/e2e/models/test_glm53_flash_b200.py

🚀 8-gpu-h200 (1 test): ✅ View workflow run

cd test/ && python3 registered/e2e/models/test_glm53_flash_h200.py