feat(cake_alpha_moe): add optimized Blackwell W8A8 expert up/down compute - #4287
Merged
Merged
Conversation
…_sm100) Add a frozen device TU generated from a Loom schedule fusing the full FP8 W8A8 block-scale MoE expert computation — routed gate/up projection, SwiGLU with per-token FP8 requantization, down projection, and asynchronous BF16 reduce-add into a caller-owned output accumulator — into one kernel over the standard vLLM/SGLang moe_align_block_size routing plan. New public API flashinfer.fused_moe.alphamoe_fp8_block_scale_aligned_moe plus the offline gate/up weight interleaving helper. - csrc/alphamoe_sm100/: generated device TU (verbatim, clang-format off) + a binding TU (typedef isolation) doing validation, the three TMA descriptor encodes, and the launch, mirroring the internal host launcher field for field - one kernel identity across runtime M/K/N/E/top_k/block_m; N=256/512/1024 tile through grid.y; blocks past the device-side num_tokens_post_padded extent exit untouched, so worst-case plan buffers are fine - out is a caller-owned BF16 accumulator (cp.reduce.async.bulk add); the wrapper zero-allocates it when omitted - tests: parity vs an independent torch oracle (torch.equal) on six internal-contract rows including M=1 decode, DeepSeek E257/top9 and Qwen E512/top10 coordinates, hot-expert imbalance, and N=1024; plan-extent guard with garbage-filled worst-case buffers; accumulation semantics
xslingcn
requested review from
IwakuraRein,
aleozlx,
jiahanc and
yzh119
as code owners
July 31, 2026 03:12
Contributor
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
xslingcn
requested review from
Aneureka,
bkryu,
cyx-6,
dhiraj113,
feih-nv,
jimmyzho,
kahyunnam,
nv-yunzheq,
qiching,
qsang-nv,
saltyminty,
samuellees,
sricketts,
yongwww and
yyihuang
as code owners
July 31, 2026 03:12
Contributor
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Advanced Run ID: 📒 Files selected for processing (1)
🚧 Files skipped from review as they are similar to previous changes (1)
Included review availability: Your plan provides up to 8 included reviews per hour; 4 remain after this review. 📝 WalkthroughWalkthroughAdded an SM100/SM103 fused W8A8 block-scale MoE kernel with TMA loading, SwiGLU, FP8 requantization, and BF16 reduction. Added Python, JIT, AOT, and trace integrations with correctness and contract tests. ChangesAlphaMoE SM100 integration
Priority: ➖ Normal Estimated code review effort: 5 (Critical) | ~120 minutes Change: Feature Sequence Diagram(s)sequenceDiagram
participant Caller
participant PythonAPI as alphamoe_fp8_block_scale_aligned_moe
participant Run as alphamoe_sm100::Run
participant Kernel as kernel_alpha_moe_w8a8_up_down
participant Output as BF16 output tensor
Caller->>PythonAPI: Provide tensors and routing parameters
PythonAPI->>Run: Invoke registered fused operation
Run->>Run: Validate tensors and encode TMA descriptors
Run->>Kernel: Launch configured SM100/SM103 kernel
Kernel->>Output: Reduce routed results into output
Kernel-->>PythonAPI: Complete CUDA operation
PythonAPI-->>Caller: Return output tensor
Merge Risk: 🟡 Moderate · up to Large output accumulators can overflow the reduction offset and write outside the output tensor. Resolve this bounds check before merging. 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 1🛠️ Fix failing CI checks 💡
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 8
🧹 Nitpick comments (3)
tests/moe/test_alphamoe_sm100.py (1)
418-425: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueGuard the exactness assumption of the truncated-extent check.
The
torch.equalassertion holds only whileCONTRACT_CASES[1]keepsn == 256and one expert per block. If a maintainer changes that row, the assertion becomes order-sensitive and flaky, and the failure will not point at the case list. Add an explicit precondition so the test fails with a clear reason instead.🧪 Proposed guard
truncated = torch.tensor([block_m], dtype=torch.int32, device="cuda") + # Exactness below requires a single 128-wide intermediate block. + assert n == 256, "this exact-equality check assumes N=256 (one intermediate block)" out = _launch(case, num_tokens_post_padded=truncated)🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/moe/test_alphamoe_sm100.py` around lines 418 - 425, Add an explicit precondition immediately before the truncated-extent assertion in the test, validating that CONTRACT_CASES[1] still has n == 256 and one expert per block. Use a clear failure message identifying the exactness assumption, then retain the existing torch.equal check unchanged.flashinfer/fused_moe/alphamoe_sm100.py (1)
69-104: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd
register_fake_opsupport for this custom op before relying on it in torch.compile paths.
fp8_block_scale_aligned_moeis registered as a custom op, butregister_fake_opcurrently only returnslambda x: xinflashinfer/utils.py, so this alias does not register a meta implementation. Register and import a matching fake op, then wire it through the same conditional path as the real op.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@flashinfer/fused_moe/alphamoe_sm100.py` around lines 69 - 104, Add a matching fake/meta implementation for fp8_block_scale_aligned_moe and register it through the existing register_fake_op mechanism in flashinfer.utils, updating the lambda-based handling as needed. Import the fake registration and wire it alongside the existing `@register_custom_op` declaration so torch.compile can resolve this alias through the same conditional path as the real operation.flashinfer/jit/fused_moe.py (1)
339-345: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAvoid unconditionally adding SM103a to the AlphaMoE spec.
gen_alphamoe_sm100_moduleappends a hardcoded SM103a gencode target tosm100a_nvcc_flags, so this module always builds both SM100a and SM103a and cannot be limited bysupported_major_versions=[10]or an SM100-onlyFLASHINFER_CUDA_ARCH_LIST. Usecurrent_compilation_context.get_nvcc_flags_list(...)here as the sibling generators do; add the SM103a target only under a condition that matches the requested compilation context.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@flashinfer/jit/fused_moe.py` around lines 339 - 345, Update gen_alphamoe_sm100_module to obtain NVCC flags through current_compilation_context.get_nvcc_flags_list, matching the sibling generators, instead of unconditionally appending the SM103a gencode target to sm100a_nvcc_flags. Include the SM103a target only when the requested compilation context supports it, so supported_major_versions=[10] and SM100-only architecture settings build only SM100a.Source: Linters/SAST tools
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@csrc/alphamoe_sm100.cu`:
- Around line 773-784: Guard the per-token activation scale against zero in the
generator’s assignment to smem_act_scale, clamping the decoded scale to a
positive FLT_MIN floor before inversion. Mirror the same floor in the Torch
oracle in test_alphamoe_sm100.py so generated-kernel and reference behavior
match, while preserving the generated kernel body.
- Around line 1212-1215: Update CheckInputs to validate that
sorted_token_ids.numel() is at least expert_ids.numel() multiplied by block_m,
matching the x-grid launched by the kernel and its eight-ID CTA reads. Add this
host-side bound check alongside the existing expert_ids validation, before
launching the kernel.
- Around line 1298-1308: Update Run to create an ffi::CUDADeviceGuard for
hidden_states.device().device_id before CheckInputs and all CUDA setup or launch
operations, including cudaFuncSetAttribute and LaunchKernel, so execution
targets the tensor’s device.
In `@flashinfer/fused_moe/__init__.py`:
- Around line 86-89: Add alphamoe_fp8_block_scale_aligned_moe and
alphamoe_interleave_gated_weights to the __all__ declaration in
flashinfer.fused_moe, alongside the other public MoE APIs, so wildcard imports
and public API tooling expose both imported functions.
In `@flashinfer/fused_moe/alphamoe_sm100.py`:
- Around line 190-194: Update the Notes docstring in the relevant AlphaMoE SM100
definition to reference the actual translation unit csrc/alphamoe_sm100.cu
instead of the nonexistent csrc/alphamoe_sm100/ directory, leaving the device
requirement and other documentation unchanged.
- Around line 109-125: Add the package-standard API metadata to
alphamoe_fp8_block_scale_aligned_moe, including
supported_compute_capability([100, 103]), and add flashinfer_api(trace=...) to
alphamoe_interleave_gated_weights using the benchmark trace definition. Update
gen_alphamoe_sm100_module and its AOT registration so both exported entry points
are covered when FLASHINFER_DISABLE_JIT is enabled, with unsupported compute
capabilities rejected through the decorator.
In `@tests/moe/test_alphamoe_sm100.py`:
- Around line 405-407: Rename the unused label unpack target to _label in both
CONTRACT_CASES unpacking sites: tests/moe/test_alphamoe_sm100.py lines 405-407
for CONTRACT_CASES[1] and lines 431-433 for CONTRACT_CASES[0]. Leave the other
unpacked values unchanged.
- Around line 252-254: Move the `w1` and `w2` dequantization out of the pre-loop
setup and into the existing per-expert loop after the empty `pair_indices`
check. Dequantize only `case["w1"][expert]` and `case["w2"][expert]` with that
expert’s scale slice, then use these tensors for the expert’s matrix operations,
removing the full all-expert float32 materializations.
---
Nitpick comments:
In `@flashinfer/fused_moe/alphamoe_sm100.py`:
- Around line 69-104: Add a matching fake/meta implementation for
fp8_block_scale_aligned_moe and register it through the existing
register_fake_op mechanism in flashinfer.utils, updating the lambda-based
handling as needed. Import the fake registration and wire it alongside the
existing `@register_custom_op` declaration so torch.compile can resolve this alias
through the same conditional path as the real operation.
In `@flashinfer/jit/fused_moe.py`:
- Around line 339-345: Update gen_alphamoe_sm100_module to obtain NVCC flags
through current_compilation_context.get_nvcc_flags_list, matching the sibling
generators, instead of unconditionally appending the SM103a gencode target to
sm100a_nvcc_flags. Include the SM103a target only when the requested compilation
context supports it, so supported_major_versions=[10] and SM100-only
architecture settings build only SM100a.
In `@tests/moe/test_alphamoe_sm100.py`:
- Around line 418-425: Add an explicit precondition immediately before the
truncated-extent assertion in the test, validating that CONTRACT_CASES[1] still
has n == 256 and one expert per block. Use a clear failure message identifying
the exactness assumption, then retain the existing torch.equal check unchanged.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: 4480f683-afad-4774-9365-ad0ac8be6941
📒 Files selected for processing (6)
csrc/alphamoe_sm100.cuflashinfer/fused_moe/__init__.pyflashinfer/fused_moe/alphamoe_sm100.pyflashinfer/jit/__init__.pyflashinfer/jit/fused_moe.pytests/moe/test_alphamoe_sm100.py
4 tasks
Collaborator
|
/bot run |
Collaborator
Collaborator
|
[FAILED] Pipeline #60909301 — 9/18 executed test jobs passed Compared with nightly #60712014. Unit Tests
✅ Pass · 🟡 Old failure · ❌ New failure · ⏱ Test timeout · Multi-GPU and Multi-Node Tests — 5/6 passed
Failure detailsNew relative to nightly (attribution uncertain)
Pre-existing failures
Timeouts, infrastructure, or incomplete jobs
|
Collaborator
|
/bot run run tests/moe/test_alphamoe_sm100.py |
Collaborator
|
Invalid |
Collaborator
|
/bot run tests/moe/test_alphamoe_sm100.py |
Collaborator
|
@flashinfer-bot run |
Collaborator
Collaborator
|
[FAILED] Pipeline #67588326 — 12/19 executed test jobs passed Compared with nightly #67341591 (different CI configuration). Unit Tests
✅ Pass · 🟡 Old failure · ❌ New failure · ⏱ Test timeout · Multi-GPU and Multi-Node Tests — 6/6 passed
Failure detailsPR-related regressions
Timeouts, infrastructure, or incomplete jobs
|
yyihuang
changed the title
yyihuang
changed the title
yyihuang
changed the title
yyihuang
changed the title
Collaborator
|
@flashinfer-bot run |
Collaborator
|
/bot run tests/moe/test_alphamoe_sm100.py |
Collaborator
yyihuang
approved these changes
Sep 13, 2026
yzh119
approved these changes
Sep 13, 2026
Collaborator
|
[FAILED] Pipeline #67614749 — 12/19 executed test jobs passed Compared with nightly #67341591 (different CI configuration). Unit Tests
✅ Pass · 🟡 Old failure · ❌ New failure · ⏱ Test timeout · Multi-GPU and Multi-Node Tests — 6/6 passed
Failure detailsPR-related regressions
Timeouts, infrastructure, or incomplete jobs
|
Collaborator
|
/bot run tests/moe/test_alphamoe_sm100.py |
Collaborator
Collaborator
|
[SUCCESS] Pipeline #67617789: 18/19 executed test jobs passed |
yyihuang
enabled auto-merge (squash)
September 13, 2026 08:22
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
6 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
📌 Description
Add the SM100a/SM103a AlphaMoE W8A8 compute API
alphamoe_fp8_block_scale_aligned_moeand offlinealphamoe_interleave_gated_weightshelper. The compute call fuses gate/up projection, SwiGLU, intermediate FP8 requantization and down projection over an existing aligned route plan, accumulating into BF16 output without a global intermediate.Final W8 kernel performance (v46)
Measured CUDA: f37119c07d56. Test-only head
106dede4b7f5fixes down-reference post-dot scaling; CUDA, exact assertions and tolerances are unchanged. GB300 replay: all 9 original tests pass.CI: pre-commit/API/docs pass. Separate validation: all 23 executed non-VR200 jobs pass (22 required). W8: 9/9 on B200/GB200/GB300 x CUDA12.9/13.0 (54/54); dedicated multi-GPU/node suites are separate. GitHub CI is running (10 passed, 3 active). Appendix is historical.
S=stock SGLang Triton; C=candidate; O=old public export; W=source launcher. W/C is export parity. S runtime:
5407ec1a7dfee227a408702addcc15007ec7f126. CUDA SHA256 C:237a9b44549fa82fb3f7833ee51df76dd08d0927fc22674bc62c77c2182429b1; O:9be0cee09c84bf5dfd41fd65c5fe705761d614172dff6e147dd1e4a9584573ea(separate JIT, same flags/ABI).Same FP8 codes/scales, weights/routes to BF16. S times GEMM1, SwiGLU, intermediate FP8 quantization, GEMM2, combine (5 kernels); C/O/W fuse all five (1 kernel). us sum GPU durations; CPU gaps excluded. Input quantization, backend route alignment, weight conversion, allocation, JIT, warmup and per-sample output resets are untimed. Rounding/accumulation differ.
GB300/SM103a, 1 GPU, strict cold-L2 CUPTI, no fallback. Per arm: 5 warmups, 30 samples/round, 5 paired rounds/shape; odd C,S,O,W; even W,O,S,C. Rows=30-sample medians; summary durations=5-row medians; ratios=median paired ratios. All 2400 samples audited.
Fixtures: N=2I=256; scale blocks128x128; routing balancedness0.8. A8/A128: M8/128,K7168,E257,k9,shared,scale2.5,seeds28101/28102. Q8/Q128: M8/128,K2048,E512,k10,no shared,scale1,seeds28103/28104. Q uses Qwen geometry; fixtures are not real-request correctness.
S: BM/BN/BK=64/128/128,group-M=32,4 warps,3 stages,no up/down TMA or separate down config. C/O/W BM=8; C:192 threads,3 stages; z=2 iff M<=8 and K>=4096, else 1; y=1; disjoint Z output columns.
Summary (us; ratios >1 favor C):
All 20 paired rounds:
All 20 S/C>1. Q128 O/C<1 in rounds 2,3 (median 1.001057); W/C median 0.995176: export overhead, no significance claim.
SUCCESS/exit 0: submit-to-terminal 93.530691s, physical 93.431010s, harness 71.991005s.
FP32 scales:
((a*b)*r)*sbecomes((a*r)*s)*b(a=activation,b=W2,r=gated route,s=routed scale), not bit-identical. Qwen validates z1; H7168/z2 has performance evidence only.Final Qwen correctness (v46)
4xGB300,TP4/EP1; Qwen/Qwen3-Next-80B-A3B-Instruct-FP8@c5f5f263bdd5cc134092897864e8905d8fe7b928. Fixed AlphaMoE router: stock Triton experts B vs final W8 C (independent W8 attribution).
Independent audit: B=C=1259/1314 (95.8143%), delta 0; 7 gains/7 losses. Unchanged gates pass: both>=0.95,C-B>=-0.005. Five-shot GSM8K questions/prompts/references and wire/SDK inputs match. Each side: 1314 successful HTTP responses; 0 failures/retries/invalid/empty answers. Empty answers count wrong; paired outputs/HTTP receipts retained.
Runtime and CUDA identities are above. Dataset SHA256:
3730d312f6e3440559ace48831e51066acaca737f6eabec99bccb9e4b3c39d14.Four TP ranks: actual router kernels in B/C;
kernel_alpha_moe_w8a8_up_downin C. Real requests: H2048/I128/E512/k10/BM8,z1. C kernel-M discrete sets (count;min,max): decode(49;8,512), eager(97;515,911), prefill(28;64,16384). Exact sets/token counts retained; graph M denotes real-request replay geometry. Integrity passes.API B/C: 62.166970/50.689964s; harness 1448.686871s; physical turnaround 1460.407082s. Correctness wall times, not serving throughput or kernel timing.
Separate GB300 sanitizer replays preserve real eager inputs/layout: M1/H2048/I128/E512/k10/BM8,scale1,rank0/layer0,z1. Synccheck: 0 errors (23.510655s physical); racecheck: 0 hazards/errors/warnings (25.081664s). Scope excludes z2, graphs and other ranks. Capture: 545.017671s physical, 536.014468s harness, 513.267672s startup, 3.942568s request; no new accuracy/performance run.
Retained pre-optimization measurements and review history
The following unchanged record describes the earlier implementation and its original validation revisions. Its references to the current head or final results apply to that historical record, not the new revision above.
Real SGLang E2E correctness: PASS on the Qwen3-Next FP8 workload below. Holding AlphaMoE routing fixed and replacing Triton expert compute with this W8A8 compute gives 1263/1314 → 1260/1314 correct (96.1187% → 95.8904%, −0.2283 percentage points). Both sides exceed the unchanged 95% threshold and the paired drop is within 0.5 percentage points. The separate stock-to-combined comparison is 1261/1314 → 1260/1314; it is reported separately from W8A8 attribution. The retained audit checks all real question/prompt pairs, HTTP responses, runtime shapes and actual kernels on all four TP ranks. Full per-repeat serving results below retain the combined-path C128 regression.
Current head:
6818acabd7600e1cc6d13331dfa73c86f86c4947. The measured W8A8 implementation atf12b7d47f14c94fae8d3d44cc01e16c6eddd26f8, exercised inc6407025a445d0d6c3bcfd28a7326456b17e4387, is unchanged in this head, including its bindings and registration. The later main merge is source-equivalence evidence; these remain the original real-model measurements, with their exact runtime revisions recorded below. The measured public revision was based on mainaeab8e964bb76675012a2bd2b72bec06502af48a. Older historical measurements remain labelled separately and do not validate the current implementation.New kernel measurements on 2026-09-13 UTC use the retained combined FlashInfer checkout
f5c95353d3723360c55c6af313afd09bbd8bfdac; this API implementation is unchanged in the current PR head. This is a new GPU performance measurement, separate from the retained real-model correctness run.W8A8 fused expert-compute GPU performance
Hardware: NVIDIA GB300; one GPU per measurement. Five paired rounds per shape, 30 cold-L2 samples per arm per round, strict CUPTI activity tracing without timing-backend fallback. Each round has 5 explicit warmup invocations per arm. Compilation, autotuning, fixture allocation and output reset are outside the recorded intervals.
GPU sum is the sum of correlated kernel durations. GPU span is the first-to-last correlated activity interval and includes inter-kernel gaps; it can therefore include delays between host submissions. They are different measurements and are reported separately. Speedup = baseline duration / candidate duration; below 1× is a regression. Summary durations are medians of five round medians, and summary speedups are medians of the five paired ratios.
No separate SM-clock observations were retained by this W8A8 runner; no fixed-clock or clock-normalized comparison is claimed.
These deterministic performance fixtures reuse the retained generators. They are not real-model accuracy evidence; the separate real-model SGLang E2E correctness section remains authoritative. Five rounds are repeated workload samples, not five independent service starts.
Baseline: stock SGLang Triton prequantized/prealigned MoE compute sequence, using the stock selected configuration. Candidate: one fused W8A8 expert-compute kernel. Both receive the same FP8 activation codes, scales, FP32 route weights and logical expert weights. Input quantization, route alignment and static weight-layout conversion are outside both boundaries. Baseline GEMM1, SwiGLU, intermediate FP8 quantization, GEMM2 and combination are all measured; the candidate fuses this work. Each implementation retains its own intermediate rounding and accumulation behavior. The stock routine retains its ordinary intermediate-allocation calls; CPU gaps affect span but not summed GPU durations.
Qwen geometry uses routed scaling 1.0 and no shared expert. E=257 geometry uses routed scaling 2.5 and one shared expert. FP8 weight blocks are 128×128; activation quantization groups are 128. Per-shape stock default/tuned configuration selected by the existing stock dispatcher is recorded below; no new tuning search is performed by this runner.
Every paired round
Each duration below is the median of 30 samples; all raw sample distributions and correlated activity identities are retained in the audit artifacts.
Stock configuration and activity counts
{"alignment_block_m": 64, "config": {"BLOCK_SIZE_K": 128, "BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 128, "GROUP_SIZE_M": 32, "num_stages": 3, "num_warps": 4}, "down_config": null, "down_tma": false, "plan_extent": 4864, "up_tma": false}.{"alignment_block_m": 64, "config": {"BLOCK_SIZE_K": 128, "BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 128, "GROUP_SIZE_M": 32, "num_stages": 3, "num_warps": 4}, "down_config": null, "down_tma": false, "plan_extent": 29376, "up_tma": false}.{"alignment_block_m": 64, "config": {"BLOCK_SIZE_K": 128, "BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 128, "GROUP_SIZE_M": 32, "num_stages": 3, "num_warps": 4}, "down_config": null, "down_tma": false, "plan_extent": 3520, "up_tma": false}.{"alignment_block_m": 64, "config": {"BLOCK_SIZE_K": 128, "BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 128, "GROUP_SIZE_M": 32, "num_stages": 3, "num_warps": 4}, "down_config": null, "down_tma": false, "plan_extent": 16320, "up_tma": false}.The two successful timing steps took 49.199 s and 50.666 s physical execution (99.865 s total); their first-start-to-final-end turnaround was 512.419 s, including the intervening additional-shape submission. These are separate from the microsecond GPU intervals above.
Public contract and review fixes
[M,K]with FP32 per-token group-128 scales; FP8 weights with FP32 128×128 block scales. Gate/up rows use the provided eight-row interleave.K % 128 == 0; combined gate/up widthN % 256 == 0;block_m >= 8and divisible by eight. These are API constraints, not claims that every admitted shape has model-accuracy coverage.sorted_token_ids.numel() >= expert_ids.numel() * block_m, including inactive capacity.0 * infinity. Every nonzero block retains the original scale calculation. The device source consequently is no longer described as byte-identical to the original export.fi_tracedescriptions are included. JIT uses only requested exact SM100a/SM103a targets; SM100a requires CUDA ≥12.8 and SM103a CUDA ≥12.9.Real-model evaluation scope and acceptance
The SGLang integration is sgl-project/sglang#34072. Its W8 path is explicitly limited to Qwen3-Next FP8, BF16 activations, TP4/EP1,
E512/H2048/I_local128/routed_top_k10/BM8, no expert all-to-all, and separate shared experts. It rejects unsupported routing/activation options, fulltorch.compile, and speculative decoding.The evaluation used the pinned real Qwen checkpoint, the canonical full GSM8K 5-shot chat workload, identical baseline/candidate prompts, temperature zero, and retained per-question answers. Candidate accuracy was required to remain ≥0.95 and the paired accuracy delta ≥−0.005. Triton, router-only with Triton expert compute, and the combined AlphaMoE backend were measured separately. Shapes came only from live requests after health, with eager/graph receipts and GPU kernel execution evidence. No hand-authored tensor case or microbenchmark was counted as model correctness.
Performance reporting requires every baseline/candidate workload repeat, ratio, TTFT/TPOT, and available memory metrics. Five workload repeats are not described as five independent server restarts.
Recorded real-model E2E results
The following reports identify the evaluated source and model revisions. They retain every serving repeat, failed gate and unavailable metric. Historical reports remain separately labeled below.
Qwen W8A8 attribution, holding the AlphaMoE router fixed (graph)
Model:
Qwen/Qwen3-Next-80B-A3B-Instruct-FP8atc5f5f263bdd5cc134092897864e8905d8fe7b928. SGLang:5407ec1a7dfee227a408702addcc15007ec7f126; FlashInfer:c6407025a445d0d6c3bcfd28a7326456b17e4387.Comparison: AlphaMoE router + Triton MoE → AlphaMoE router + W8A8 MoE. TP4 / EP1 / DP1, execution mode
graph, speculative decoding disabled, shared-expert fusion disabled on both sides.FP8 checkpoint with the fixed 0.95 acceptance threshold. The retained five-shot scorer preserves the historical evaluation protocol; current SGLang default registrations use sgl-eval instead.
moe_runner_backendtritonflashinfer_alphamoeattention_backendtritontritonprefill_attention_backendNoneNonedecode_attention_backendNoneNonekv_cache_dtypeautoautochunked_prefill_size1638416384mem_fraction_static0.70.7cuda_graph_config{'decode': {'backend': 'full', 'bs': [1, 2, 4, 8, 12, 16, 24, 32, 40, 48, 56, 64, 72, 80, 88, 96, 104, 112, 120, 128, 136, 144, 152, 160, 168, 176, 184, 192, 200, 208, 216, 224, 232, 240, 248, 256, 272, 288, 304, 320, 336, 352, 368, 384, 400, 416, 432, 448, 464, 480, 496, 512], 'full_prefill_max_req': None, 'full_prefill_prefix_chunk_tokens': None, 'max_bs': 512, 'tc_compiler': 'eager'}, 'prefill': {'backend': 'breakable', 'bs': [4, 8, 12, 16, 20, 24, 28, 32, 48, 64, 80, 96, 112, 128, 144, 160, 176, 192, 208, 224, 240, 256, 288, 320, 352, 384, 416, 448, 480, 512, 576, 640, 704, 768, 832, 896, 960, 1024, 1280, 1536, 1792, 2048, 2304, 2560, 2816, 3072, 3328, 3584, 3840, 4096, 4608, 5120, 5632, 6144, 6656, 7168, 7680, 8192, 8704, 9216, 9728, 10240, 10752, 11264, 11776, 12288, 12800, 13312, 13824, 14336, 14848, 15360, 15872, 16384], 'full_prefill_max_req': None, 'full_prefill_prefix_chunk_tokens': None, 'max_bs': 16384, 'tc_compiler': 'eager'}}{'decode': {'backend': 'full', 'bs': [1, 2, 4, 8, 12, 16, 24, 32, 40, 48, 56, 64, 72, 80, 88, 96, 104, 112, 120, 128, 136, 144, 152, 160, 168, 176, 184, 192, 200, 208, 216, 224, 232, 240, 248, 256, 272, 288, 304, 320, 336, 352, 368, 384, 400, 416, 432, 448, 464, 480, 496, 512], 'full_prefill_max_req': None, 'full_prefill_prefix_chunk_tokens': None, 'max_bs': 512, 'tc_compiler': 'eager'}, 'prefill': {'backend': 'breakable', 'bs': [4, 8, 12, 16, 20, 24, 28, 32, 48, 64, 80, 96, 112, 128, 144, 160, 176, 192, 208, 224, 240, 256, 288, 320, 352, 384, 416, 448, 480, 512, 576, 640, 704, 768, 832, 896, 960, 1024, 1280, 1536, 1792, 2048, 2304, 2560, 2816, 3072, 3328, 3584, 3840, 4096, 4608, 5120, 5632, 6144, 6656, 7168, 7680, 8192, 8704, 9216, 9728, 10240, 10752, 11264, 11776, 12288, 12800, 13312, 13824, 14336, 14848, 15360, 15872, 16384], 'full_prefill_max_req': None, 'full_prefill_prefix_chunk_tokens': None, 'max_bs': 16384, 'tc_compiler': 'eager'}}This comparison holds the AlphaMoE router fixed and isolates replacing the Triton MoE computation with AlphaMoE W8A8.
GSM8K uses the repository's retained five-shot chat scorer (
sglang.test.simple_eval_mixed_prefix_gsm8k.GSM8KEval), all 1,314 held-out examples (the first five of the 1,319-example split supply the examples), temperature 0, top-p 1, and a 2,048-token generation limit. Both variants receive the same prompts and references.Accuracy gate: PASS. Both accuracies must be at least 0.95, and candidate minus baseline must be ≥ −0.005. The candidate gained 6 questions and lost 9; a passing accuracy gate does not assert bitwise equality or zero accuracy loss.
End-to-end kernel verification: PASS. This additionally requires complete server/request evidence and actual GPU kernel traces for each AlphaMoE variant in this comparison.
Performance
Five fixed-workload repetitions run per server variant and are paired by repeat ID and seed. They are not five independent server launches. All repeats are included. Throughput speedup is candidate output tokens/s divided by baseline output tokens/s; values below 1 are regressions. TTFT, TPOT and request E2E columns are per-repeat medians in milliseconds; lower is better.
Workload: 1024 input / 512 output tokens, 1,024 requests per repeat.
Performance gate: PASS. Every concurrency must have median paired speedup ≥ 1, and at least one must win all five paired repetitions.
Memory and execution evidence
Memory values below are resident-device snapshots after health/model/graph startup or after the named phase. They include model, KV cache and allocator reservations; they are not peak measurements or isolated CUDA Graph allocations.
Post-health GSM8K dispatch/capture receipts observed execution modes:
decode_graph_replay,eager,prefill_graph_replay. Startup capture and fixed-workload performance requests are excluded from this coverage.Observed AlphaMoE runtime shapes below come from those request receipts. Kernel M is the submitted kernel geometry, or the registered capture geometry when a real request replayed that graph. Dispatch M includes graph padding; real tokens are the actual request tokens before that padding. Each column lists its observed values separately, not a Cartesian product of supported shapes. An inclusive range contains only consecutive values that were all observed; missing values are not inferred. The stock baseline has no AlphaMoE kernel shape records.
decode_graph_replayalphamoe_fused_routereageralphamoe_fused_routerprefill_graph_replayalphamoe_fused_routerdecode_graph_replayalphamoe_fp8_block_scale_aligned_moedecode_graph_replayalphamoe_fused_routereageralphamoe_fp8_block_scale_aligned_moeeageralphamoe_fused_routerprefill_graph_replayalphamoe_fp8_block_scale_aligned_moeprefill_graph_replayalphamoe_fused_routerGPU execution witness: PASS. After unprofiled measurements, one stored real GSM8K prompt was replayed with a 32-token limit under SGLang's GPU/CUPTI profiler. This request is excluded from accuracy and performance. Required actual CUDA kernel symbols in all four TP traces:
kernel_alpha_moe_fused_router,kernel_alpha_moe_w8a8_up_down.Resumed execution and physical turnaround
The Qwen campaign spans workload submissions: an earlier submission reached its time limit, and the combined router/W8A8 variant resumed with a restarted server. Completed GSM8K outputs and sealed serving measurements were retained; the resume executed the unfinished serving repetitions and GPU profiling. The audit revalidated the retained evidence against the same model, source, runtime and evaluation contract.
There are five workload repetitions per concurrency and variant, paired by repeat ID and seed. These are not five independent server processes or restarts. The contributing attempts below show exactly where the recorded phases ran; restarting a server does not add a repetition.
baselinerouter_onlycombinedcombinedPhysical timing covers the entire three-variant Qwen campaign, including startup, interrupted work and cleanup. It is shared across the Qwen comparison tables, not a separate cost for each pair.
GSM8K API runtime and the per-repeat serving request latencies above are the measured workload results. They are reported separately from physical turnaround; interrupted, unsealed measurements are excluded from the performance table.
Qwen combined router and W8A8 versus stock MoE (graph)
Model:
Qwen/Qwen3-Next-80B-A3B-Instruct-FP8atc5f5f263bdd5cc134092897864e8905d8fe7b928. SGLang:5407ec1a7dfee227a408702addcc15007ec7f126; FlashInfer:c6407025a445d0d6c3bcfd28a7326456b17e4387.Comparison: stock MoE backend → AlphaMoE router + W8A8 MoE. TP4 / EP1 / DP1, execution mode
graph, speculative decoding disabled, shared-expert fusion disabled on both sides.FP8 checkpoint with the fixed 0.95 acceptance threshold. The retained five-shot scorer preserves the historical evaluation protocol; current SGLang default registrations use sgl-eval instead.
moe_runner_backendtritonflashinfer_alphamoeattention_backendtritontritonprefill_attention_backendNoneNonedecode_attention_backendNoneNonekv_cache_dtypeautoautochunked_prefill_size1638416384mem_fraction_static0.70.7cuda_graph_config{'decode': {'backend': 'full', 'bs': [1, 2, 4, 8, 12, 16, 24, 32, 40, 48, 56, 64, 72, 80, 88, 96, 104, 112, 120, 128, 136, 144, 152, 160, 168, 176, 184, 192, 200, 208, 216, 224, 232, 240, 248, 256, 272, 288, 304, 320, 336, 352, 368, 384, 400, 416, 432, 448, 464, 480, 496, 512], 'full_prefill_max_req': None, 'full_prefill_prefix_chunk_tokens': None, 'max_bs': 512, 'tc_compiler': 'eager'}, 'prefill': {'backend': 'breakable', 'bs': [4, 8, 12, 16, 20, 24, 28, 32, 48, 64, 80, 96, 112, 128, 144, 160, 176, 192, 208, 224, 240, 256, 288, 320, 352, 384, 416, 448, 480, 512, 576, 640, 704, 768, 832, 896, 960, 1024, 1280, 1536, 1792, 2048, 2304, 2560, 2816, 3072, 3328, 3584, 3840, 4096, 4608, 5120, 5632, 6144, 6656, 7168, 7680, 8192, 8704, 9216, 9728, 10240, 10752, 11264, 11776, 12288, 12800, 13312, 13824, 14336, 14848, 15360, 15872, 16384], 'full_prefill_max_req': None, 'full_prefill_prefix_chunk_tokens': None, 'max_bs': 16384, 'tc_compiler': 'eager'}}{'decode': {'backend': 'full', 'bs': [1, 2, 4, 8, 12, 16, 24, 32, 40, 48, 56, 64, 72, 80, 88, 96, 104, 112, 120, 128, 136, 144, 152, 160, 168, 176, 184, 192, 200, 208, 216, 224, 232, 240, 248, 256, 272, 288, 304, 320, 336, 352, 368, 384, 400, 416, 432, 448, 464, 480, 496, 512], 'full_prefill_max_req': None, 'full_prefill_prefix_chunk_tokens': None, 'max_bs': 512, 'tc_compiler': 'eager'}, 'prefill': {'backend': 'breakable', 'bs': [4, 8, 12, 16, 20, 24, 28, 32, 48, 64, 80, 96, 112, 128, 144, 160, 176, 192, 208, 224, 240, 256, 288, 320, 352, 384, 416, 448, 480, 512, 576, 640, 704, 768, 832, 896, 960, 1024, 1280, 1536, 1792, 2048, 2304, 2560, 2816, 3072, 3328, 3584, 3840, 4096, 4608, 5120, 5632, 6144, 6656, 7168, 7680, 8192, 8704, 9216, 9728, 10240, 10752, 11264, 11776, 12288, 12800, 13312, 13824, 14336, 14848, 15360, 15872, 16384], 'full_prefill_max_req': None, 'full_prefill_prefix_chunk_tokens': None, 'max_bs': 16384, 'tc_compiler': 'eager'}}This comparison changes both routing and MoE computation; its accuracy or speedup cannot be attributed to W8A8 alone.
GSM8K uses the repository's retained five-shot chat scorer (
sglang.test.simple_eval_mixed_prefix_gsm8k.GSM8KEval), all 1,314 held-out examples (the first five of the 1,319-example split supply the examples), temperature 0, top-p 1, and a 2,048-token generation limit. Both variants receive the same prompts and references.Accuracy gate: PASS. Both accuracies must be at least 0.95, and candidate minus baseline must be ≥ −0.005. The candidate gained 7 questions and lost 8; a passing accuracy gate does not assert bitwise equality or zero accuracy loss.
End-to-end kernel verification: PASS. This additionally requires complete server/request evidence and actual GPU kernel traces for each AlphaMoE variant in this comparison.
Performance
Five fixed-workload repetitions run per server variant and are paired by repeat ID and seed. They are not five independent server launches. All repeats are included. Throughput speedup is candidate output tokens/s divided by baseline output tokens/s; values below 1 are regressions. TTFT, TPOT and request E2E columns are per-repeat medians in milliseconds; lower is better.
Workload: 1024 input / 512 output tokens, 1,024 requests per repeat.
Performance gate: FAIL. Every concurrency must have median paired speedup ≥ 1, and at least one must win all five paired repetitions.
Memory and execution evidence
Memory values below are resident-device snapshots after health/model/graph startup or after the named phase. They include model, KV cache and allocator reservations; they are not peak measurements or isolated CUDA Graph allocations.
Post-health GSM8K dispatch/capture receipts observed execution modes:
decode_graph_replay,eager,prefill_graph_replay. Startup capture and fixed-workload performance requests are excluded from this coverage.Observed AlphaMoE runtime shapes below come from those request receipts. Kernel M is the submitted kernel geometry, or the registered capture geometry when a real request replayed that graph. Dispatch M includes graph padding; real tokens are the actual request tokens before that padding. Each column lists its observed values separately, not a Cartesian product of supported shapes. An inclusive range contains only consecutive values that were all observed; missing values are not inferred. The stock baseline has no AlphaMoE kernel shape records.
decode_graph_replayalphamoe_fp8_block_scale_aligned_moedecode_graph_replayalphamoe_fused_routereageralphamoe_fp8_block_scale_aligned_moeeageralphamoe_fused_routerprefill_graph_replayalphamoe_fp8_block_scale_aligned_moeprefill_graph_replayalphamoe_fused_routerGPU execution witness: PASS. After unprofiled measurements, one stored real GSM8K prompt was replayed with a 32-token limit under SGLang's GPU/CUPTI profiler. This request is excluded from accuracy and performance. Required actual CUDA kernel symbols in all four TP traces:
kernel_alpha_moe_fused_router,kernel_alpha_moe_w8a8_up_down.Resumed execution and physical turnaround
The Qwen campaign spans workload submissions: an earlier submission reached its time limit, and the combined router/W8A8 variant resumed with a restarted server. Completed GSM8K outputs and sealed serving measurements were retained; the resume executed the unfinished serving repetitions and GPU profiling. The audit revalidated the retained evidence against the same model, source, runtime and evaluation contract.
There are five workload repetitions per concurrency and variant, paired by repeat ID and seed. These are not five independent server processes or restarts. The contributing attempts below show exactly where the recorded phases ran; restarting a server does not add a repetition.
baselinerouter_onlycombinedcombinedPhysical timing covers the entire three-variant Qwen campaign, including startup, interrupted work and cleanup. It is shared across the Qwen comparison tables, not a separate cost for each pair.
GSM8K API runtime and the per-repeat serving request latencies above are the measured workload results. They are reported separately from physical turnaround; interrupted, unsealed measurements are excluded from the performance table.
Historical E2E correctness and performance
The following tables are restored reports from 2026-08-08, not results reproduced during this delivery. The source revisions and environments differ from the current candidate. Full raw per-request/per-repeat artifacts have not been revalidated in this session; rounded values are preserved as reported. These historical tables do not validate the evaluated revisions reported above.
Historical model:
Qwen/Qwen3-Next-80B-A3B-Instruct-FP8atc5f5f263bdd5cc134092897864e8905d8fe7b928, SGLang v0.5.16, 4×GB300, TP4/EP1, Triton baseline versus combined #4339 router + #4287 W8A8. This does not isolate either kernel.The candidate reached the old
accuracy >= 0.95anddelta >= -0.005gates exactly. This is an observed drop and a boundary pass, not proof of identical outputs or no accuracy loss.Serving: 1,024 requests per workload, 1,024 input / 512 output tokens, three workload repeats per concurrency. Throughput columns are reported means; speedup is the reported mean of paired repeat ratios, which need not equal a ratio of rounded means.
TTFT was reported worse at all three concurrencies; exact TTFT/TPOT rows were not restored. Prefill CUDA Graph allocation was reported as 2.76 → 58.04 GB/GPU. This is a historical graph-allocation report, not a newly measured process-memory peak. Individual baseline/candidate repeat values were not restored, so a full repeat table cannot be reconstructed. These are three repeats within one server deployment per backend, not three independently restarted deployments.
The fixed geometry was
E512/H2048/I_local128/routed_top_k10/BM8, with shared experts separate. Dynamic M was not traced; the workload does not establish any particular M bucket.Historical API timing, separate from model E2E
Both sides were reported to pass the then-used fixed FP8 comparison threshold before CUPTI cold-L2 timing. These fixture dimensions were not observed production M values and are not model-correctness evidence.
The boundaries differ: AlphaMoE receives prequantized activations and a prealigned route plan; the stock SGLang API includes dynamic quantization, routing/alignment and additional launches. These ratios are not equal-boundary GEMM speedups and do not predict service speed.
🔍 Related Issues
Companion router #4339, NVFP4 compute #4340, and SGLang integration #34072.
🚀 Pull Request Checklist
Thank you for contributing to FlashInfer! Before we review your pull request, please make sure the following items are complete.
✅ Pre-commit Checks
pre-commitby runningpip install pre-commit(or used your preferred method).pre-commit install.pre-commit run --all-filesand fixed any reported issues.🧪 Tests
unittest, etc.).New verification on 2026-09-12: the repository pre-commit suite passed on the PR changes before the latest main-only merge, including applicable formatting, Ruff and mypy checks. Checks were executed in the compute environment. The subsequent merge of main aeab8e9 changes unrelated attention files; the reviewed AlphaMoE files remain identical. Full-tree hooks, GPU compilation and runtime/model checks are not claimed by that result. Recorded SGLang model results and their gate outcomes appear above; merge acceptance still requires the stated criteria. Existing kernel tests and trace/ABI checks are engineering coverage; they do not establish model accuracy.
🔬 Experimental Track
flashinfer/experimental/and/or an@flashinfer_experimental_api. Tracking issue: #tests/experimental/and were validated on the intended hardware; a runnable example is included.flashinfer/aot.py, and no experimental backend is reachable frombackend="auto"withoutFLASHINFER_ALLOW_EXPERIMENTAL_AUTO_BACKENDS=1. (Calling an@flashinfer_experimental_apior naming a backend explicitly is itself the opt-in and needs no environment variable.)Reviewer Notes
Please review the zero-activation scale branch, device/capacity checks, weight layout, and API/trace/AOT integration. Review readiness permits code review; it does not assert fresh model accuracy or merge readiness. Assess the recorded evaluation revisions and gate outcomes above against the merge criteria.
Full-scope mypy validation on 2026-09-12 passed after explicitly typing the shared trace dictionaries. The touched-file hook suite also passed. This resolves the earlier CI mypy failure without changing device code or numerical tolerances. PR Test run #34722609670, attempt 2 now succeeds on
f12b7d47f14c94fae8d3d44cc01e16c6eddd26f8, with all 10 jobs successful. Earlier authorization-skipped runs and attempt 1's runner-shutdown AOT failures remain historical. These CI checks supplement the real-model validation reported above.