Skip to content

[Perf] minimax-sparse: Speed up MiniMax-M3 sparse attention on Hopper - #46135

Merged
ispobock merged 3 commits into
mainfrom
m3-sparse-hopper
Oct 25, 2026
Merged

[Perf] minimax-sparse: Speed up MiniMax-M3 sparse attention on Hopper#46135
ispobock merged 3 commits into
mainfrom
m3-sparse-hopper

Conversation

@ispobock

@ispobock ispobock commented Oct 23, 2026 •
edited by github-actions Bot
Loading

Copy link
Copy Markdown
Collaborator

Summary

Faster MiniMax-M3 sparse attention on Hopper: Triton decode and prefill kernels with tiled block pointers (Triton >= 3.9), fp16-widened decode for fp8 KV on the paged path, and a fused bf16-to-fp8 cast in the fused K/V/index store.

Changes to existing code

  • On Hopper with fp8_e4m3 KV, the sparse index K cache is stored in fp8 (it was bf16 unless fp8 attention GEMMs were on). Other platforms keep the previous dtype.
  • The fused K/V/index store also handles an unscaled bf16-to-fp8_e4m3 cast (CUDA only); other inputs use the separate stores as before.

Verification

  • Kernel tests on H200: test_minimax_sparse_fp8_attn_gemm.py 88 passed; the fused store cast is bit-exact against x.float().clamp(-448, 448).to(torch.float8_e4m3fn).
  • MiniMax-M3.1-Flash on 8xH200, H200 cookbook recipe plus --kv-cache-dtype fp8_e4m3, sgl-eval 0.1.3, this PR vs. fp16 widening disabled:
GSM8K (greedy, 2 runs) AIME26 (avg@8)
this PR 96.21 / 96.06 91.25
widening off 96.36 / 96.51 93.75

The between-arm per-question divergence and correctness flips are within the run-to-run noise of either arm.

Cut from #45925 (MiniMax-M3.1-Flash support) at 10e655a.

CI States

Latest PR Test (Base): 🚫 Run #40114646995
Latest PR Test (Extra): ✅ Run #40114646930
Latest PR Test (AMD ROCm 10): ❌ Run #40114646909

@github-actions

This comment has been minimized.

Sign in to view
@ispobock

ispobock commented Oct 23, 2026

Copy link
Copy Markdown
Collaborator Author

/rerun-test test/registered/kernels/ops/attention/test_minimax_sparse_flash_with_topk_idx.py test/registered/kernels/ops/attention/test_minimax_sparse_decode_paged_tile.py test/registered/kernels/ops/attention/test_minimax_sparse_fp8_attn_gemm.py test/registered/kernels/ops/attention/test_minimax_decode_topk_page_table.py test/registered/kernels/ops/attention/test_minimax_sparse_gqa.py test/registered/kernels/ops/kvcache/test_minimax_store_kv_index.py

@github-actions

github-actions Bot commented Oct 23, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

🚀 1-gpu-h100 (5 tests): ❌ View workflow run 

cd test/ && python3 registered/kernels/ops/attention/test_minimax_sparse_flash_with_topk_idx.py
cd test/ && python3 registered/kernels/ops/attention/test_minimax_sparse_decode_paged_tile.py
cd test/ && python3 registered/kernels/ops/attention/test_minimax_sparse_fp8_attn_gemm.py
cd test/ && python3 registered/kernels/ops/attention/test_minimax_sparse_gqa.py
cd test/ && python3 registered/kernels/ops/kvcache/test_minimax_store_kv_index.py

🚀 4-gpu-b200 (6 tests): ✅ View workflow run 

cd test/ && python3 registered/kernels/ops/attention/test_minimax_sparse_flash_with_topk_idx.py
cd test/ && python3 registered/kernels/ops/attention/test_minimax_sparse_decode_paged_tile.py
cd test/ && python3 registered/kernels/ops/attention/test_minimax_sparse_fp8_attn_gemm.py
cd test/ && python3 registered/kernels/ops/attention/test_minimax_decode_topk_page_table.py
cd test/ && python3 registered/kernels/ops/attention/test_minimax_sparse_gqa.py
cd test/ && python3 registered/kernels/ops/kvcache/test_minimax_store_kv_index.py

@ispobock @zijiexia @zcnrex
speed up m3 sparse attention on hopper
691a025 
Co-authored-by: zijiexia <37504505+zijiexia@users.noreply.github.com>
Co-authored-by: Chunan Zeng <zcnrex@gmail.com>
@ispobock

ispobock commented Oct 23, 2026

Copy link
Copy Markdown
Collaborator Author

/rerun-test test/registered/kernels/ops/attention/test_minimax_sparse_flash_with_topk_idx.py test/registered/kernels/ops/attention/test_minimax_sparse_decode_paged_tile.py test/registered/kernels/ops/attention/test_minimax_sparse_fp8_attn_gemm.py test/registered/kernels/ops/attention/test_minimax_decode_topk_page_table.py test/registered/kernels/ops/attention/test_minimax_sparse_gqa.py test/registered/kernels/ops/kvcache/test_minimax_store_kv_index.py

@github-actions

github-actions Bot commented Oct 23, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

🚀 1-gpu-h100 (5 tests): ❌ View workflow run 

cd test/ && python3 registered/kernels/ops/attention/test_minimax_sparse_flash_with_topk_idx.py
cd test/ && python3 registered/kernels/ops/attention/test_minimax_sparse_decode_paged_tile.py
cd test/ && python3 registered/kernels/ops/attention/test_minimax_sparse_fp8_attn_gemm.py
cd test/ && python3 registered/kernels/ops/attention/test_minimax_sparse_gqa.py
cd test/ && python3 registered/kernels/ops/kvcache/test_minimax_store_kv_index.py

🚀 4-gpu-b200 (6 tests): ✅ View workflow run 

cd test/ && python3 registered/kernels/ops/attention/test_minimax_sparse_flash_with_topk_idx.py
cd test/ && python3 registered/kernels/ops/attention/test_minimax_sparse_decode_paged_tile.py
cd test/ && python3 registered/kernels/ops/attention/test_minimax_sparse_fp8_attn_gemm.py
cd test/ && python3 registered/kernels/ops/attention/test_minimax_decode_topk_page_table.py
cd test/ && python3 registered/kernels/ops/attention/test_minimax_sparse_gqa.py
cd test/ && python3 registered/kernels/ops/kvcache/test_minimax_store_kv_index.py

@ispobock ispobock added run-ci CI: run the baseline test suite on this PR run-ci-extra CI: also run the extra suite (requires run-ci) bypass-fail-fast CI: a failing job no longer aborts its siblings (lint still gates) labels Oct 24, 2026
@ispobock

ispobock commented Oct 24, 2026

Copy link
Copy Markdown
Collaborator Author

/rerun-test test/registered/kernels/ops/attention/test_minimax_sparse_fp8_attn_gemm.py test/registered/kernels/ops/kvcache/test_minimax_store_kv_index.py test/registered/kernels/benchmark/kvcache/bench_minimax_store_kv_index.py

@github-actions

github-actions Bot commented Oct 24, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

🚀 1-gpu-h100 (3 tests): ✅ View workflow run 

cd test/ && python3 registered/kernels/ops/attention/test_minimax_sparse_fp8_attn_gemm.py
cd test/ && python3 registered/kernels/ops/kvcache/test_minimax_store_kv_index.py
cd test/ && python3 registered/kernels/benchmark/kvcache/bench_minimax_store_kv_index.py

🚀 4-gpu-b200 (2 tests): ✅ View workflow run 

cd test/ && python3 registered/kernels/ops/attention/test_minimax_sparse_fp8_attn_gemm.py
cd test/ && python3 registered/kernels/ops/kvcache/test_minimax_store_kv_index.py

@ispobock

ispobock commented Oct 24, 2026

Copy link
Copy Markdown
Collaborator Author

/run-full-ci

@ispobock

ispobock commented Oct 25, 2026

Copy link
Copy Markdown
Collaborator Author

/rerun-failed-ci

@ispobock
ispobock merged commit db7b66c into main Oct 25, 2026
332 of 354 checks passed
@ispobock
ispobock deleted the m3-sparse-hopper branch October 25, 2026 03:28
@thanhhao98 thanhhao98 mentioned this pull request Oct 26, 2026
Open
4 of 5 tasks
@akhilg-nv akhilg-nv mentioned this pull request Oct 26, 2026
Merged
3 of 5 tasks
akhilg-nv added a commit to akhilg-nv/sglang that referenced this pull request Oct 26, 2026
@akhilg-nv @claude
Merge branch 'main' into fla-triton38-compat
c180a64 
Conflict in minimax_sparse/decode/topk_sparse.py: main (sgl-project#46135) ported the
MiniMax sparse kernels off block pointers and replaced the stale block_size
autotune key with BLOCK_SIZE_N, which supersedes this branch's key fix, so
take main's version of the file.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
akhilg-nv added a commit to akhilg-nv/sglang that referenced this pull request Oct 26, 2026
@akhilg-nv @claude
Drop the MiniMax sparse import of the block-pointer shim
98b3487 
Main (sgl-project#46135) ported the MiniMax sparse kernels off block pointers, so
minimax_sparse no longer needs the shim. Only the FLA kernels use it now.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
@waqahmed-amd-fi waqahmed-amd-fi mentioned this pull request Oct 26, 2026
Draft
5 tasks done
@ispobock ispobock mentioned this pull request Oct 27, 2026
Open
13 tasks done
@jiaqiang-dot-liu jiaqiang-dot-liu mentioned this pull request Oct 28, 2026
Closed
4 of 5 tasks
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@merrymercy merrymercy Awaiting requested review from merrymercy merrymercy is a code owner
@Ying1123 Ying1123 Awaiting requested review from Ying1123 Ying1123 is a code owner
@hnyls2002 hnyls2002 Awaiting requested review from hnyls2002 hnyls2002 is a code owner
@xiezhq-hermann xiezhq-hermann Awaiting requested review from xiezhq-hermann xiezhq-hermann is a code owner
@hanming-lu hanming-lu Awaiting requested review from hanming-lu hanming-lu is a code owner
@yizhang2077 yizhang2077 Awaiting requested review from yizhang2077 yizhang2077 is a code owner
@hzh0425 hzh0425 Awaiting requested review from hzh0425 hzh0425 is a code owner
@alphabetc1 alphabetc1 Awaiting requested review from alphabetc1 alphabetc1 is a code owner
@huangtingwei9988 huangtingwei9988 Awaiting requested review from huangtingwei9988 huangtingwei9988 is a code owner
@DarkSharpness DarkSharpness Awaiting requested review from DarkSharpness DarkSharpness is a code owner
@HaiShaw HaiShaw Awaiting requested review from HaiShaw HaiShaw is a code owner
@BBuf BBuf Awaiting requested review from BBuf BBuf is a code owner
@celve celve Awaiting requested review from celve celve is a code owner
@HydraQYH HydraQYH Awaiting requested review from HydraQYH HydraQYH is a code owner
@yuan-luo yuan-luo Awaiting requested review from yuan-luo yuan-luo is a code owner

Assignees

No one assigned

Labels

bypass-fail-fast CI: a failing job no longer aborts its siblings (lint still gates) jit-kernel memory-pool run-ci CI: run the baseline test suite on this PR run-ci-extra CI: also run the extra suite (requires run-ci)

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

1 participant

@ispobock