Skip to content

perf(glm): write prefill MLA projections in token-major layout - #1710

Open
shihaobai wants to merge 1 commit into
mainfrom
bsh/glm-prefill-mla-token-major-bmm
Open

shihaobai wants to merge 1 commit into
mainfrom
bsh/glm-prefill-mla-token-major-bmm

Conversation

@shihaobai

Copy link
Copy Markdown
Collaborator

Write GLM sparse MLA Q/V projections directly in token-major layout to avoid the following layout copy. Enable for supported BF16 Hopper shapes in eager Prefill with one TP rank per DP; keep the reference path for short cold inputs and other configurations.

Validation against main dc53088 on H100: 15 GPU tests pass, plus byte-exact checks using checkpoint layers 3/43, two seeds and 1024/1659/4096/8192 rows. At 8K rows, Q takes 0.981 → 0.315 ms and V 0.627 → 0.294 ms. These are projection-plus-layout timings, not system TPM gains.

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