Skip to content

perf(kda): parallelize diagonal solves in Prefill - #1703

Open
shihaobai wants to merge 1 commit into
mainfrom
bsh/glm-kda-split-triangular
Open

shihaobai wants to merge 1 commit into
mainfrom
bsh/glm-kda-split-triangular

Conversation

@shihaobai

@shihaobai shihaobai commented Nov 11, 2026 •

Copy link
Copy Markdown
Collaborator

Parallelize diagonal solves for packed KDA Prefill, selecting by chunk occupancy instead of GPU or request-count whitelists.

H200, 8K tokens and 64 heads: eager KDA takes 4.392 → 3.899 ms, with 32 MiB additional solver workspace. The full-forward intermediate allocation peak remains 1412 MiB. CUDA correctness tests, 26 unit tests and memcheck pass; SM80/89/90/100/120 compilation also passes.

@shihaobai
shihaobai marked this pull request as ready for review November 11, 2026 14:04
@shihaobai
shihaobai force-pushed the bsh/glm-kda-split-triangular branch from dae6c2d to 8c815ee Compare November 11, 2026 14:09
@shihaobai shihaobai changed the title perf(kda): parallelize diagonal solves in Hopper Prefill perf(kda): parallelize diagonal solves in Prefill Nov 11, 2026
Solve the four diagonal blocks independently and retain FP32 inverse tiles through merging. Select by chunk occupancy to bound scratch, with CUDA regression coverage for packed layouts, dtypes, precision modes and fallback behavior.
@shihaobai
shihaobai force-pushed the bsh/glm-kda-split-triangular branch from 4d70ee5 to 87deeb0 Compare November 11, 2026 15:24
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