Skip to content

Sparse Load for LoopK Backward#269

Open
KevinZeng08 wants to merge 22 commits intoSandAI-org:mainfrom
KevinZeng08:bwd-sparse-load
Open

Sparse Load for LoopK Backward#269
KevinZeng08 wants to merge 22 commits intoSandAI-org:mainfrom
KevinZeng08:bwd-sparse-load

Conversation

@KevinZeng08
Copy link
Copy Markdown
Contributor

@KevinZeng08 KevinZeng08 commented Feb 25, 2026 •
edited
Loading

DONE

  • implement sparse load with loopk backward for small k block sizes
  • correctness test
  • refactor fwd sparse load
  • range merge for loopk backward, same as Range Merge for LoopK Backward #240

TODO

  • optimize perf
  • code refactor

@KevinZeng08
range merge for loopk bwd, TODO: ffa test, perf profile, sparse load
c780ed0
@KevinZeng08
update test
44503f3
@KevinZeng08
update bench
163f6f1
@KevinZeng08
Merge branch 'main' of https://github.com/SandAI-org/MagiAttention in…
a790382 
…to ffa-bwd
@KevinZeng08
update test
03bd966
@KevinZeng08
update bench
22c6cd1
@KevinZeng08
init bwd sparse load interface
c73fea0
@KevinZeng08
refactor load with block meta
bfb3891
@KevinZeng08
refactor mma & epilogue with block meta
3455530
@KevinZeng08
remove unroll, restore previous perf
8f360b5
@KevinZeng08
refactor fwd sparse load
1371c89
@KevinZeng08
init bwd sparse load pipeline
2ea1b1c
@KevinZeng08
add sparse load logic for load & mma, TODO: store_dkv scatter store, …
8f24d1f 
…boundary mask
@KevinZeng08
add scatter store for dK/dV, fix boundary mask, TODO: pass test, perf…
6b6195c 
… test
@KevinZeng08
fix case for 64 head dim
10629a7
@KevinZeng08
perf: simply tune register
68ef8e4
@KevinZeng08
remove todo
0bd7c41
@KevinZeng08
Merge branch 'main' of https://github.com/SandAI-org/MagiAttention in…
7951bca 
…to bwd-sparse-load
@KevinZeng08
skip packgqa specific test in block sparse test
c9ea985
@KevinZeng08
refactor block meta
be19c43
@KevinZeng08
Merge branch 'main' of https://github.com/SandAI-org/MagiAttention in…
b30b9ca 
…to bwd-sparse-load
@KevinZeng08 KevinZeng08 changed the title [WIP] Sparse Load for LoopK Backward Sparse Load for LoopK Backward Mar 9, 2026
@KevinZeng08 KevinZeng08 closed this Mar 9, 2026
@KevinZeng08 KevinZeng08 reopened this Mar 9, 2026
@KevinZeng08 KevinZeng08 closed this Mar 9, 2026
@KevinZeng08 KevinZeng08 reopened this Mar 9, 2026
@KevinZeng08 KevinZeng08 closed this Mar 9, 2026
@KevinZeng08 KevinZeng08 reopened this Mar 9, 2026
@Strivin0311 Strivin0311 closed this Mar 9, 2026
@Strivin0311 Strivin0311 reopened this Mar 9, 2026
@codecov
Copy link
Copy Markdown

codecov bot commented Mar 10, 2026

Codecov Report

❌ Patch coverage is 70.83333% with 7 lines in your changes missing coverage. Please review.

Files with missing lines Patch % Lines
magi_attention/utils/sparse_utils.py 0.00% 4 Missing ⚠️
magi_attention/functional/flex_flash_attn.py 84.21% 3 Missing ⚠️

📢 Thoughts on this report? Let us know!

@Strivin0311 Strivin0311 requested a review from littsk March 10, 2026 05:17
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@littsk littsk Awaiting requested review from littsk

At least 0 approving reviews are required to merge this pull request.

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@KevinZeng08 @Strivin0311