Skip to content

fix: bound DSA prefill score memory - #4896

Draft
CUHKSZzxy wants to merge 3 commits into
InternLM:mainfrom
CUHKSZzxy:fix/glm-dsa-logits-chunking
Draft

fix: bound DSA prefill score memory#4896
CUHKSZzxy wants to merge 3 commits into
InternLM:mainfrom
CUHKSZzxy:fix/glm-dsa-logits-chunking

Conversation

@CUHKSZzxy

Copy link
Copy Markdown
Collaborator

Depends on #4853. Supersedes #4892 by carrying its sparse FlashMLA refactor together with the DSA memory fix.

Summary

  • Separate dense MLA and sparse DSA FlashMLA execution into dedicated implementations, with sparse implementation selection handled by the CUDA attention builder.
  • Bound the FP32 DeepGEMM DSA prefill score tensor by chunking query rows under a configurable 512 MiB budget while flattening and reusing the indexer K cache once.
  • Reserve the same DSA score workspace before automatic KV-cache sizing so cache allocation cannot consume the transient runtime budget.
  • Keep score computation and top-k selection within consistent prefill and decode/fallback paths, and add coverage for chunking, result parity, and runtime-memory estimation.

Validation

  • Pre-commit checks passed.
  • Targeted DSA and executor tests passed (31 passed, 3 CUDA-dependent skipped).
  • Focused FlashMLA tests passed (7 passed).
  • Full GLM-5.2 LCB GPU reproduction is pending.

Assistance

Assisted with Codex + GPT-5.6-Sol xHigh, reviewed manually

@CUHKSZzxy
CUHKSZzxy force-pushed the fix/glm-dsa-logits-chunking branch from 5bb1923 to 69dd860 Compare August 24, 2026 04:43
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