Skip to content

refactor: separate sparse FlashMLA attention - #4892

Closed
CUHKSZzxy wants to merge 23 commits into
InternLM:mainfrom
CUHKSZzxy:refactor/split-flashmla-attention
Closed

refactor: separate sparse FlashMLA attention#4892
CUHKSZzxy wants to merge 23 commits into
InternLM:mainfrom
CUHKSZzxy:refactor/split-flashmla-attention

Conversation

@CUHKSZzxy

@CUHKSZzxy CUHKSZzxy commented Aug 21, 2026

Copy link
Copy Markdown
Collaborator

Summary

  • Separate dense MLA and sparse DSA FlashMLA execution into dedicated implementations.
  • Keep shared KV-cache fill, flattening, and paged FlashMLA primitives in the dense base while moving sparse index mapping and kernels to sparse_mla.py.
  • Select the sparse implementation from the CUDA attention builder and derive scheduler metadata from the selected implementation.

Validation

  • Pre-commit checks passed.
  • Focused FlashMLA tests passed (7 passed).
  • GLM-5.2 FP8 TP8 with FP8 KV cache, MTP5, and symmetric-memory all-reduce completed 80/80 SWE-Smith requests without CUDA errors, OOMs, graph breaks, or recompilation warnings.

Assistance

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

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