Skip to content

perf(glm): fuse MTP input normalization and concatenation - #1712

Open
shihaobai wants to merge 1 commit into
mainfrom
bsh/glm-mtp-fused-norm
Open

shihaobai wants to merge 1 commit into
mainfrom
bsh/glm-mtp-fused-norm

Conversation

@shihaobai

@shihaobai shihaobai commented Nov 12, 2026 •

Copy link
Copy Markdown
Collaborator

Fuse the three GLM MTP input RMSNorms and concatenation before eh_proj into one Triton kernel. Preserve main-norm intermediate rounding and both in-place input updates. Use the fused path for 4096-wide BF16/FP16 CUDA inputs with the default norm warp count; retain the existing path otherwise.

Validated on H100 and retested after rebasing onto main db192ec: 46 GPU tests pass, covering rounding, strided buffers, CUDA graph replay and fallback. All 18 actual-checkpoint projection comparisons are byte-exact (1–8192 tokens). Norm+concat is 2.17–2.38x faster under CUDA graph at 1–64 tokens; including eh_proj, the gain is 1.18–1.19x. The 8192-token eager consumer is 1.12x faster. No end-to-end TPM claim.

@shihaobai
shihaobai force-pushed the bsh/glm-mtp-fused-norm branch from 05e3762 to c2e7e21 Compare November 12, 2026 11:12
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant