feat(minimax): integrate FlashInfer MSA on Blackwell by yyihuang · Pull Request #35846 · sgl-project/sglang

yyihuang · 2026-08-21T09:52:20Z

Summary

route MiniMax-M3 sparse attention through the FlashInfer MSA source backend on SM100/SM103
preserve the existing fallback and CUDA Graph contracts while adding deterministic route validation
add a reproducible TP4 A/B harness with GPQA Diamond, LongBench-v2, and fixed-request concurrency sweeps

Depends on flashinfer-ai/flashinfer#4355.

Validation

4x GB300 offline preflight: passed
SGLang MSA contract tests: 20 passed
required SM103 decode, long-prefill, and TopK-16 routes: present and loadable
MiniMax-M3-MXFP8 model: all 31 shards loaded with TP4
both A/B arms used a fresh server and cache, passed fixed-parity and backend-route audits, and completed every measured request without failure or retry

Correctness

One source-to-export round used seed 20260819, temperature 0, top-p 1, one request thread, and deterministic NCCL Ring/Simple communication.

Evaluation	Existing source	Export	Delta	Threshold	Result
GPQA Diamond	148/198 (0.747475)	149/198 (0.752525)	+1 question	at least -1 question	Pass
balanced LongBench-v2	0.62	0.69	+0.07	at least -0.02	Pass

Each arm completed 298/298 real /v1/chat/completions requests. GPQA evaluator runtimes were 10,788.10 s (source) and 10,518.17 s (export); LongBench-v2 runtimes were 2,892.19 s and 3,226.26 s. Total physical turnaround was 8:41:18.

Fixed-request serving performance

One fresh-server source-to-export round used 8,192 input plus 1,024 output tokens, 256 prompts at each concurrency, and no measured-call warmup.

Concurrency	Existing source tok/s	Export tok/s	Export speedup	Source median TTFT ms	Export median TTFT ms
1	118.690021	120.097017	1.011854x	271.597387	198.826976
8	607.231454	628.450218	1.034943x	1437.199334	1018.541559
32	1318.364194	1426.237567	1.081824x	4427.318413	3356.341827
128	2168.235875	2582.798452	1.191198x	17291.430353	12113.521719

Both arms completed 1,024/1,024 requests. Measured serving runtime was 2,960.09 s for source and 2,885.19 s for export; physical turnaround was 2:41:36 and Slurm runtime was 2:41:12.

CI States

Latest PR Test (Base): ❌ Run #32524726753
Latest PR Test (Extra): ❌ Run #32524726746
Latest PR Test (AMD ROCm 7.2): ❌ Run #32524726692

…#4355) ## Summary This pull request adds source-distributed Blackwell MSA kernels behind the existing `flashinfer.msa_ops` APIs. It covers sparse prefill, sparse decode, exact TopK-16 selection, architecture-specific routing for compute capability 10.0 and 10.3, JIT/AOT registration, workspace and CUDA Graph support, tests, and documentation. The final source inventory contains 75 units: 38 for SM100a and 37 for SM103a. The exported sources pass native replay, public API, direct reducer, exact-route, formatting, hash, and compute-sanitizer gates on both architectures. Related to #4254. SGLang integration: sgl-project/sglang#35846. ## Validation - Read-only source replay: 38/38 native compilations on B200 and 37/37 on GB300, with no stale generated files. - B200 and GB300 final gates: 42/42 CPU/source-manifest/JIT/routing checks, 27/27 public GPU API checks, 3/3 direct reducers, and 4/4 exact-route checks on each architecture. - Compute Sanitizer synccheck: 0 errors for direct reducers and exact routes on both B200 and GB300. - The final head preserves the measured source payload: all 151 CUDA-source entries are byte-identical to the measured parent; the changed Python files are AST-equivalent. ## Kernel A/B One outer round per architecture in one process and at the same shapes. Every arm visit used 10 warmups and 30 CUPTI cold-L2 samples. ### NVIDIA B200 (`sm_100a`) | Shape | Existing source ms | Export ms | MiniMax baseline ms | Baseline/source | Baseline/export | |---|---:|---:|---:|---:|---:| | prefill_bf16_b1_q4096_kv4096_h64 | 0.4762895 | 0.475393 | 0.4869615 | 1.02240654x | 1.0243346x | | decode_bf16_b128_q1_kv4096_h64 | 0.1007845 | 0.1004 | 0.4752805 | 4.71580947x | 4.73386952x | | speculative_bf16_b128_q4_kv4096_h64 | 0.1864325 | 0.18544 | 0.786034 | 4.21618548x | 4.23875108x | | mtp_bf16_b128_q16_kv4096_h64 | 0.656113 | 0.656001 | 1.16717 | 1.77891613x | 1.77921985x | | decode_fp16_b128_q1_kv4096_h64 | 0.1150885 | 0.114272 | N/A | N/A | N/A | | decode_fp8_b128_q1_kv4096_h64 | 0.2856005 | 0.284353 | 0.5337295 | 1.8687975x | 1.8769962x | | official_decode_bf16_b32_q8_kv8192_h64_hkv4_k16_paged | 0.141569 | 0.141201 | 0.362833 | 2.56294104x | 2.56962061x | | official_decode_bf16_b64_q8_kv65536_h64_hkv4_k32_paged | 0.633009 | 0.635281 | 3.258534 | 5.14768984x | 5.1292798x | | official_prefill_mixed_fp8_b3_q1024_kv8192_h32_hkv2_k8_flat | 0.109632 | 0.108544 | 0.1572 | 1.43388792x | 1.44826061x | | official_prefill_bf16_b3_q4096_kv8192_h8_hkv2_k4_paged | 0.072384 | 0.07152 | 0.107904 | 1.49071618x | 1.50872483x | | coverage_decode_fp16_b32_q4_kv8192_h64_hkv4_k16_paged | 0.0973445 | 0.096896 | N/A | N/A | N/A | | coverage_decode_mixed_fp8_b32_q1_kv8192_h64_hkv4_k16_flat | 0.074336 | 0.074336 | 0.1868005 | 2.51292106x | 2.51292106x | | boundary_decode_bf16_b2_q1_kv257_h8_hkv1_k4_paged | 0.008608 | 0.007328 | 0.028416 | 3.30111524x | 3.87772926x | Over the 11 baseline-comparable rows, baseline/source geomean is 2.408351x and baseline/export geomean is 2.452001x. Across all 13 rows, the export/source latency geomean is 0.984027, or a 1.016232x export speedup. Physical benchmark turnaround was 2,306 s; the Slurm runtime was 2,305 s. ### NVIDIA GB300 (`sm_103a`) | Shape | Existing source ms | Export ms | MiniMax baseline ms | Baseline/source | Baseline/export | |---|---:|---:|---:|---:|---:| | prefill_bf16_b1_q4096_kv4096_h64 | 0.2954575 | 0.2940975 | 0.426866 | 1.44476278x | 1.45144382x | | decode_bf16_b128_q1_kv4096_h64 | 0.099761 | 0.0990245 | 0.431266 | 4.32299195x | 4.35514443x | | speculative_bf16_b128_q4_kv4096_h64 | 0.180529 | 0.180193 | 0.709075 | 3.9277623x | 3.93508627x | | mtp_bf16_b128_q16_kv4096_h64 | 0.555027 | 0.5514425 | 1.038181 | 1.8705054x | 1.8826641x | | decode_fp16_b128_q1_kv4096_h64 | 0.111696 | 0.110816 | N/A | N/A | N/A | | decode_fp8_b128_q1_kv4096_h64 | 0.269329 | 0.268625 | 0.4771545 | 1.77164175x | 1.77628478x | | official_decode_bf16_b32_q8_kv8192_h64_hkv4_k16_paged | 0.133921 | 0.1337925 | 0.3275855 | 2.44611002x | 2.44845937x | | official_decode_bf16_b64_q8_kv65536_h64_hkv4_k32_paged | 0.627427 | 0.629475 | 2.9496455 | 4.70117719x | 4.68588189x | | official_prefill_mixed_fp8_b3_q1024_kv8192_h32_hkv2_k8_flat | 0.096641 | 0.089968 | 0.1379045 | 1.42697716x | 1.53281722x | | official_prefill_bf16_b3_q4096_kv8192_h8_hkv2_k4_paged | 0.060592 | 0.05896 | 0.104896 | 1.73118564x | 1.77910448x | | coverage_decode_fp16_b32_q4_kv8192_h64_hkv4_k16_paged | 0.094481 | 0.0938885 | N/A | N/A | N/A | | coverage_decode_mixed_fp8_b32_q1_kv8192_h64_hkv4_k16_flat | 0.0699365 | 0.0697765 | 0.168369 | 2.40745533x | 2.41297572x | | boundary_decode_bf16_b2_q1_kv257_h8_hkv1_k4_paged | 0.007968 | 0.007104 | 0.03552 | 4.45783133x | 5.00000000x | Over the 11 baseline-comparable rows, baseline/source geomean is 2.508831x and baseline/export geomean is 2.563379x. Across all 13 rows, the export/source latency geomean is 0.980892, or a 1.019480x export speedup. Physical benchmark turnaround was 2,450 s; the Slurm runtime was 2,442 s. ## SGLang end-to-end correctness MiniMax-M3-MXFP8 ran on 4x GB300 with TP4, one fresh-server source-to-export round, fresh cache per arm, seed 20260819, temperature 0, top-p 1, and one request thread. Both arms completed 298/298 real `/v1/chat/completions` requests with zero failures, retries, or measured-window compilation events, and both passed the required backend-route audit. | Evaluation | Existing source | Export | Delta | Threshold | Result | |---|---:|---:|---:|---:|---:| | GPQA Diamond | 148/198 (0.747475) | 149/198 (0.752525) | +1 question | at least -1 question | Pass | | balanced LongBench-v2 | 0.62 | 0.69 | +0.07 | at least -0.02 | Pass | GPQA evaluator runtimes were 10,788.10 s (source) and 10,518.17 s (export); LongBench-v2 runtimes were 2,892.19 s and 3,226.26 s. The complete job turnaround was 8:41:18. ## SGLang fixed-request serving A/B The same model and hardware ran one fresh-server source-to-export round using 8,192 input plus 1,024 output tokens, 256 prompts at each concurrency, and no measured-call warmup. Both arms completed 1,024/1,024 requests without failure and passed startup, fixed-parity, route, and measured-window audits. | Concurrency | Existing source tok/s | Export tok/s | Export speedup | Source median TTFT ms | Export median TTFT ms | |---:|---:|---:|---:|---:|---:| | 1 | 118.690021 | 120.097017 | 1.011854x | 271.597387 | 198.826976 | | 8 | 607.231454 | 628.450218 | 1.034943x | 1437.199334 | 1018.541559 | | 32 | 1318.364194 | 1426.237567 | 1.081824x | 4427.318413 | 3356.341827 | | 128 | 2168.235875 | 2582.798452 | 1.191198x | 17291.430353 | 12113.521719 | Measured serving runtime was 2,960.09 s for source and 2,885.19 s for export. Physical turnaround was 2:41:36; the Slurm runtime was 2:41:12.  ## Summary by CodeRabbit * **New Features** * Added Blackwell MSA support for SM100 and SM103 architectures. * Added sparse prefill, decode, and top-k selection across flat and paged layouts. * Added support for BF16, FP16, FP8 key/value formats, GQA variants, and long-prefill workflows. * Added architecture-aware packed-KV detection and workspace support for CUDA graph capture. * **Documentation** * Expanded MSA guidance with supported architectures, formats, constraints, and workspace requirements. * **Tests** * Added comprehensive correctness, routing, capability, JIT, source validation, and benchmark coverage.  --------- Co-authored-by: Yingyi Huang <averyh@nvidia.com> Co-authored-by: Zihao Ye <zihaoy@nvidia.com>

averyhNV added 29 commits August 21, 2026 01:23

feat(minimax-m3): integrate FlashInfer source MSA

4730128

test(minimax-m3): make MSA A/B validation fail closed

eb7fbac

fix(minimax-m3): preserve MSA capture addresses

1bdc42f

bench(minimax-m3): automate three-run MSA gate

bf8d382

bench(minimax-m3): normalize LongBench domains

77242a7

fix(bench): validate MSA package source path

8a35aa6

test(bench): pin the SGLang kernel dependency

f5a66c7

test(bench): load exact Torch kernel ABI in preflight

783f322

test(bench): exercise the DeepGEMM Torch ABI

f6ac665

bench(minimax-m3): bound long-context accuracy slice

8c43562

test(minimax-m3): pin the baseline FFI runtime

81bbfab

bench(minimax-m3): enforce balanced long-context slice

c3d365f

bench(minimax-m3): cover balanced 512K contexts

b71548f

bench(minimax-m3): serialize baseline JIT preparation

7e6d709

bench(minimax-m3): keep serving sweeps offline

eaf21b0

bench(minimax-m3): preserve exact serving lengths

144742f

fix(minimax-m3): stage contiguous MSA queries

06cc6e7

fix(minimax-m3): keep MSA selector local

a8ee14f

fix(bench): satisfy repository checks

ae67038

test(minimax-m3): make MSA accuracy margins explicit

62bfdd6

test: add MiniMax MSA suite entry point

073c862

style: format MiniMax MSA resume check

c4d46c8

fix(minimax-m3): fence MSA page-table replay

8c780ee

style: format MSA replay test import

a2233b7

style: format MiniMax sparse backend

acd65e4

fix(bench): tolerate transient health responses

986547b

fix(bench): compare visible fixed answers

baa3997

docs(bench): define discrete accuracy margins

3814380

test(minimax-m3): pin generated MSA routes

f6f8ee0

github-actions Bot added the documentation Improvements or additions to documentation label Aug 21, 2026

averyhNV added 8 commits August 21, 2026 03:05

style(minimax): apply Black formatting

64a241d

test(minimax-m3): retain per-example GPQA evidence

476560c

test(minimax-m3): register GPQA evidence checks

e3415f1

style(minimax-m3): normalize GPQA registry spacing

7434fbb

test(minimax-m3): make integration gate fail closed

a1f6219

style(minimax-m3): apply lint formatting

0367754

test(minimax-m3): use one pair for speed modes

b67f80d

style(minimax-m3): apply import ordering

8b09f95

yyihuang mentioned this pull request Aug 23, 2026

feat(cake_msa): add Blackwell minimax sparse attention source kernels flashinfer-ai/flashinfer#4355

Merged

This was referenced Sep 7, 2026

feat(vibecuda_msa): add SM100/SM103 sparse attention backend flashinfer-ai/flashinfer#4817

Open

feat(minimax): integrate VibeCUDA MSA on Blackwell #39233

Draft

Navigation Menu

Uh oh!

feat(minimax): integrate FlashInfer MSA on Blackwell - #35846

feat(minimax): integrate FlashInfer MSA on Blackwell#35846
yyihuang wants to merge 37 commits into
sgl-project:mainfrom
yyihuang:codex/msa-flashinfer-e2e-v2-20260821

yyihuang commented Aug 21, 2026 •

edited

Loading

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

Conversation

yyihuang commented Aug 21, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Validation

Correctness

Fixed-request serving performance

CI States

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

yyihuang commented Aug 21, 2026 •

edited

Loading