refactor(moe): enforce build() for all runners and refine CUTLASS autotuning in unified MoE API by feih-nv · Pull Request #4376 · flashinfer-ai/flashinfer

feih-nv · 2026-08-06T08:03:25Z

📌 Description

Enforce an explicit check_support() -> build() -> execute lifecycle for all registered unified MoE runners, and refine CUTLASS compound-tactic autotuning.

Runner lifecycle

Centralize check_support() → build() → execute enforcement in MoERunner.
Make builds idempotent and defer backend initialization until support validation succeeds.
Keep shape-dependent TRTLLM and b12x inner runners lazy.
Update direct-runner callers and add lifecycle regression tests.

CUTLASS staged autotuning

Add AutoTuner.rank_tactics() for ranked stage candidates.
Retain the top two tactics for GEMM1 and GEMM2, then profile their four combinations end to end.
Cache complete ranked shortlists in-process while preserving the existing persistent winner format.
Rebuild shortlists from persisted winners when tuning resumes.
Select the correct dynamic optimization profile before ranking.

This reduces the search cost from O(n1 × n2) to O(n1 + n2 + k²) while improving compound tactic selection.

🔍 Related Issues

Follow up tasks for CUTLASS BF16/W4A16 runners merged in feat(moe): add unified CUTLASS BF16 and W4A16 runners #4328.

🚀 Pull Request Checklist

✅ Pre-commit Checks

I have installed pre-commit by running pip install pre-commit (or used my preferred method).
I have installed the hooks with pre-commit install.
I have run the hooks manually with pre-commit run --all-files and fixed any reported issues.

🧪 Tests

Focused validation:

496 passed
31 skipped
All changed autotuner and unified MoE test files passed

Reviewer Notes

MoELayer now validates each runner with check_support() before calling its idempotent build() method. Direct-runner users must perform the same lifecycle explicitly.
TRTLLM and b12x inner runners remain lazily constructed because their configuration depends on runtime shapes.
CUTLASS staged ranking is currently orchestrated by get_valid_tactics(); a future autotuner abstraction could model multi-stage tuning declaratively.

coderabbitai · 2026-08-06T08:03:38Z

Caution

Review failed

An error occurred during the review process. Please try again later.

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests

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.

❤️ Share

_{Comment @coderabbitai help to get the list of available commands.}

Retain the fastest tactics from each GEMM stage so the outer tuner can evaluate a bounded set of end-to-end combinations instead of assuming independent winners compose optimally. Developed with AI assistance.

Keep complete process-local tactic rankings so warm and persisted winner caches retain compound refinement, and repair the CuTe DSL build lifecycle configuration access.

feih-nv · 2026-08-10T09:34:48Z

/bot run /tests/moe

flashinfer-bot · 2026-08-10T09:36:11Z

GitLab MR !1175 has been created, and the CI pipeline #61933729 is currently running. I'll report back once the pipeline job completes.

flashinfer-bot · 2026-08-11T16:05:49Z

[FAILED] Pipeline #61933729 — 9/18 executed test jobs passed

Compared with nightly #61930115.

Unit Tests

GPU	CUDA 12.9	CUDA 13.0	Notes
5090	⚠️ Infra	⚠️ Infra	Infrastructure: `CI infrastructure failure` (2 jobs; CUDA 12.9, CUDA 13.0)
B300	❌ New	🟡 Old	New: `tests.comm.test_ulysses_topology` (1 failure; CUDA 12.9) Old: `tests.attention.test_cute_dsl_hca_dsv4` (10 failures; CUDA 12.9, CUDA 13.0) Test timeout: `1 test file timed out: tests/moe/test_trtllm_gen_fused_moe.py` (1 job; CUDA 13.0)
GB200	🟡 Old	🟡 Old	Old: `tests.attention.test_cute_dsl_hca_dsv4` (10 failures; CUDA 12.9, CUDA 13.0)
GB300	🟡 Old	🟡 Old	Old: `tests.attention.test_cute_dsl_hca_dsv4` (10 failures; CUDA 12.9, CUDA 13.0)
H100	✅ Pass	⚠️ Infra	Infrastructure: `test infrastructure interrupted the job` (1 job; CUDA 13.0)
RTX Pro 6000 Blackwell	✅ Pass	✅ Pass	—

✅ Pass · 🟡 Old failure · ❌ New failure · ⏱ Test timeout · ⚠️ Infrastructure · ❔ Unknown or unclassified · — Not run

Multi-GPU and Multi-Node Tests — 6/6 passed

GPU	CUDA 12.9	CUDA 13.0	Notes
B300 (multi-GPU)	✅ Pass	✅ Pass	—
GB200 (multi-node)	✅ Pass	✅ Pass	—
GB300 (multi-node)	✅ Pass	✅ Pass	—

Failure details

New relative to nightly (attribution uncertain)

tests.comm.test_ulysses_topology — 1 failure on B300 / CUDA 12.9
- AssertionError: only ranks [0] reported within 120s (likely a rank left the collective sequence early or hung in teardown); results so far: {0: ('pg-error', 'The server socket h…

Pre-existing failures

tests.attention.test_cute_dsl_hca_dsv4 — 30 failures on B300 / CUDA 12.9, B300 / CUDA 13.0, GB200 / CUDA 12.9, GB200 / CUDA 13.0, GB300 / CUDA 12.9, GB300 / CUDA 13.0
- AttributeError: module 'cutlass.cute.nvgpu.cpasync' has no attribute 'CopyBulkTensor2DGather4G2SOp'. Did you mean: 'CopyBulkTensorTileG2SOp'?

Timeouts, infrastructure, or incomplete jobs

Infrastructure: CI infrastructure failure — 5090 / CUDA 12.9, 5090 / CUDA 13.0
- Jobs: unit_test_5090: [cu130], unit_test_5090: [cu129]
Infrastructure: test infrastructure interrupted the job — H100 / CUDA 13.0
- Jobs: unit_test_h100: [cu130]
Test timeout: 1 test file timed out: tests/moe/test_trtllm_gen_fused_moe.py — B300 / CUDA 13.0
- Jobs: unit_test_b300: [cu130]

aleozlx

code review: lgtm. minor comment posted

## 📌 Description Adds non-EP fused shared experts to the unified TRTLLM-generated MoE API for DeepSeekV3 `FromLogits` routing. The unified block-FP8 runner supports DeepSeekFp8 and MxFp8, while the unified FP4 runner supports NVFP4, MXFP4, and W4A16. The underlying kernels already append shared expert IDs `E + k` with weight `1.0`; this PR completes the unified API integration: - Adds `ExpertConfig.num_fused_shared_experts` and validates the shared-expert contract: DeepSeekV3 routing, `K + S <= 32`, `E + S <= 512`, and no expert parallelism. - Updates `TrtllmFp8BlockRunner` for DeepSeekFp8 and MxFp8 and `TrtllmFp4RoutedRunner` for NVFP4, MXFP4, and W4A16 to accept `E + S` physical expert rows and forward `S` to the kernel. - Allocates `K + S` caller-owned routing outputs for FP4 in-kernel routing while keeping `ExpertConfig.local_num_experts` routed-only; block-FP8 instead passes empty placeholders and uses kernel-internal routing storage. - Requires explicit backend opt-in so unsupported runners reject `S > 0` through the mandatory `check_support()` lifecycle. - Adds separate block-FP8 and FP4 trace templates for `S > 0`, preserving routed `E`, physical `E + S`, and shared `S` as distinct geometry. - Fixes unified MoE autotune cache identity so in-memory and persisted keys include tactic-relevant expert geometry, routing, activation, and quantization state. ### Why the prerequisites are needed **Autotune cache keying (`c5cbe0c38`, `b4eb94bd5`).** Expert geometry is stored outside the profiled tensor list, so configurations differing only in `S` could otherwise share a tactic even though enumeration depends on `top_k + S`. Persisted keys also omit `runner_hash`. These changes introduce a shared `_cache_key_extras()` for tactic-relevant geometry, routing, activation, and quantization state and remove all runner-specific `__hash__()` overrides so every unified MoE runner follows the same cache-key contract. CUTLASS and CuTeDSL override `_cache_key_extras()` and call `super()` to append architecture, PDL, and inner-runner state without dropping the common fields. **Trace artifact refresh (`ec50a94b8`).** The committed block-FP8 DeepSeek artifact had drifted after the SwiGLU OA inputs were added, and its rendered reference omitted `_fp8_moe_run_experts`. Refreshing it separately establishes a clean baseline for the shared-expert artifacts. **Runner lifecycle ([#4376](#4376 `check_support()` is mandatory before `build()`, and packing or execution requires a built runner. Unsupported `S > 0` configurations are therefore rejected centrally without duplicating guards in every `pack_inputs()`. ## 🚀 Pull Request Checklist ### ✅ Pre-commit Checks - [x] I have installed `pre-commit` by running `pip install pre-commit` (or used my preferred method). - [x] I have installed the hooks with `pre-commit install`. - [x] I have run the hooks manually with `pre-commit run --all-files` and fixed any reported issues. ## 🧪 Tests Coverage added: - Numerical parity with the legacy flat API for DeepSeekFp8 at `S=1`, MxFp8 at `S=2`, NVFP4 at `S=1`, MXFP4 at `S=2`, and W4A16 at `S=1`. - DeepSeekFp8 `S=1` CUDA-graph capture and replay. - Configuration boundaries, unsupported-backend rejection, and in-memory and persisted autotune cache-key coverage. - Focused block-FP8 and FP4 trace dispatch, naming, geometry validation, generated artifacts, and curated `S > 0` fuzzer cases. Validation completed: - SM100: `tests/trace/` plus unified MoE and cache-key coverage — 1341 passed, 179 skipped, 0 failed before the final test-pruning and lifecycle-alignment edits. - SM103 (B300): focused block-FP8 shared-expert selection — 102 passed, 3 skipped, plus the DeepSeekFp8 and MxFp8 curated fuzzer cases. - Focused FP4 shared-expert tests in the CUDA container — 8 passed, including NVFP4, MXFP4, and W4A16 parity; all three curated FP4 fuzzer cases passed. - Focused cache-key integration tests for the unified contract, CUTLASS, and CuTeDSL — 3 passed. ## Reviewer Notes and Follow-ups - **Pre-routed inputs and routing replay:** callers can append shared IDs `E + k` and weights `1.0` themselves. Adding them inside FlashInfer would require re-striding `[T, K]` routing arrays to `[T, K + S]` or carrying dual-stride semantics through the post-routing pipeline. The replay buffer remains routed-only `[T, K]`, while the shared-expert routing kernel currently indexes it with the internal `K + S` stride, so replay is disabled for `S > 0`. Follow-up: fix the flat block-FP8 and FP4 replay paths while preserving the routed-only output contract. - **Cache invalidation:** the persisted-key change causes existing on-disk autotune entries to miss once and re-tune. This is benign but user-visible. - **EP remains unsupported:** the kernel maps a shared global ID to a weight row using `global_id - local_expert_offset`, which is valid only when the rank holds the full routed expert set. - **Physical versus routed rows:** `prepare_weights(num_local_experts=...)` takes the physical row count `E + S`, while `ExpertConfig.local_num_experts` remains routed-only `E`. - **Mixed quantization:** the current runners require homogeneous expert rows and therefore do not yet cover checkpoints such as `dsv4-flash-nvfp4`, whose routed experts use NVFP4 while shared experts use block-FP8. An initial implementation can compose an NVFP4 routed launch with a block-FP8 shared-expert launch and sum their BF16 outputs; a true single-launch implementation requires upstream kernel ABI and cubin support for separate routed/shared weight pointers, scale layouts, and dtypes. - **Test oracle:** shared-expert numerical tests compare with the legacy flat APIs because the existing self-fused pre-routed oracle cannot represent the important `S=1` and `S=2` cases under its expert-count divisibility constraints. - **Tracing:** artifacts were regenerated directly because `tests/trace/example.py` currently aborts in `recurrent_kda` with `nvidia-cutlass-dsl` 4.5.2. Follow-ups are to make Trace Apply select the concrete dispatched template per call, register reference-helper dependencies transitively and execute emitted functions in standalone tests, then repair example generation and add CI that regenerates `fi_trace_out` from an empty directory and rejects drift.

…otuning in unified MoE API (flashinfer-ai#4376) ## 📌 Description Enforce an explicit `check_support() -> build() -> execute` lifecycle for all registered unified MoE runners, and refine CUTLASS compound-tactic autotuning. ### Runner lifecycle - Centralize check_support() → build() → execute enforcement in MoERunner. - Make builds idempotent and defer backend initialization until support validation succeeds. - Keep shape-dependent TRTLLM and b12x inner runners lazy. - Update direct-runner callers and add lifecycle regression tests. ### CUTLASS staged autotuning - Add AutoTuner.rank_tactics() for ranked stage candidates. - Retain the top two tactics for GEMM1 and GEMM2, then profile their four combinations end to end. - Cache complete ranked shortlists in-process while preserving the existing persistent winner format. - Rebuild shortlists from persisted winners when tuning resumes. - Select the correct dynamic optimization profile before ranking. This reduces the search cost from O(n1 × n2) to O(n1 + n2 + k²) while improving compound tactic selection. ## 🔍 Related Issues - Follow up tasks for CUTLASS BF16/W4A16 runners merged in flashinfer-ai#4328. ## 🚀 Pull Request Checklist ### ✅ Pre-commit Checks - [x] I have installed `pre-commit` by running `pip install pre-commit` (or used my preferred method). - [x] I have installed the hooks with `pre-commit install`. - [x] I have run the hooks manually with `pre-commit run --all-files` and fixed any reported issues. ## 🧪 Tests Focused validation: - 496 passed - 31 skipped - All changed autotuner and unified MoE test files passed ## Reviewer Notes - `MoELayer` now validates each runner with `check_support()` before calling its idempotent `build()` method. Direct-runner users must perform the same lifecycle explicitly. - TRTLLM and b12x inner runners remain lazily constructed because their configuration depends on runtime shapes. - CUTLASS staged ranking is currently orchestrated by get_valid_tactics(); a future autotuner abstraction could model multi-stage tuning declaratively.

## 📌 Description Adds non-EP fused shared experts to the unified TRTLLM-generated MoE API for DeepSeekV3 `FromLogits` routing. The unified block-FP8 runner supports DeepSeekFp8 and MxFp8, while the unified FP4 runner supports NVFP4, MXFP4, and W4A16. The underlying kernels already append shared expert IDs `E + k` with weight `1.0`; this PR completes the unified API integration: - Adds `ExpertConfig.num_fused_shared_experts` and validates the shared-expert contract: DeepSeekV3 routing, `K + S <= 32`, `E + S <= 512`, and no expert parallelism. - Updates `TrtllmFp8BlockRunner` for DeepSeekFp8 and MxFp8 and `TrtllmFp4RoutedRunner` for NVFP4, MXFP4, and W4A16 to accept `E + S` physical expert rows and forward `S` to the kernel. - Allocates `K + S` caller-owned routing outputs for FP4 in-kernel routing while keeping `ExpertConfig.local_num_experts` routed-only; block-FP8 instead passes empty placeholders and uses kernel-internal routing storage. - Requires explicit backend opt-in so unsupported runners reject `S > 0` through the mandatory `check_support()` lifecycle. - Adds separate block-FP8 and FP4 trace templates for `S > 0`, preserving routed `E`, physical `E + S`, and shared `S` as distinct geometry. - Fixes unified MoE autotune cache identity so in-memory and persisted keys include tactic-relevant expert geometry, routing, activation, and quantization state. ### Why the prerequisites are needed **Autotune cache keying (`c5cbe0c38`, `b4eb94bd5`).** Expert geometry is stored outside the profiled tensor list, so configurations differing only in `S` could otherwise share a tactic even though enumeration depends on `top_k + S`. Persisted keys also omit `runner_hash`. These changes introduce a shared `_cache_key_extras()` for tactic-relevant geometry, routing, activation, and quantization state and remove all runner-specific `__hash__()` overrides so every unified MoE runner follows the same cache-key contract. CUTLASS and CuTeDSL override `_cache_key_extras()` and call `super()` to append architecture, PDL, and inner-runner state without dropping the common fields. **Trace artifact refresh (`ec50a94b8`).** The committed block-FP8 DeepSeek artifact had drifted after the SwiGLU OA inputs were added, and its rendered reference omitted `_fp8_moe_run_experts`. Refreshing it separately establishes a clean baseline for the shared-expert artifacts. **Runner lifecycle ([flashinfer-ai#4376](flashinfer-ai#4376 `check_support()` is mandatory before `build()`, and packing or execution requires a built runner. Unsupported `S > 0` configurations are therefore rejected centrally without duplicating guards in every `pack_inputs()`. ## 🚀 Pull Request Checklist ### ✅ Pre-commit Checks - [x] I have installed `pre-commit` by running `pip install pre-commit` (or used my preferred method). - [x] I have installed the hooks with `pre-commit install`. - [x] I have run the hooks manually with `pre-commit run --all-files` and fixed any reported issues. ## 🧪 Tests Coverage added: - Numerical parity with the legacy flat API for DeepSeekFp8 at `S=1`, MxFp8 at `S=2`, NVFP4 at `S=1`, MXFP4 at `S=2`, and W4A16 at `S=1`. - DeepSeekFp8 `S=1` CUDA-graph capture and replay. - Configuration boundaries, unsupported-backend rejection, and in-memory and persisted autotune cache-key coverage. - Focused block-FP8 and FP4 trace dispatch, naming, geometry validation, generated artifacts, and curated `S > 0` fuzzer cases. Validation completed: - SM100: `tests/trace/` plus unified MoE and cache-key coverage — 1341 passed, 179 skipped, 0 failed before the final test-pruning and lifecycle-alignment edits. - SM103 (B300): focused block-FP8 shared-expert selection — 102 passed, 3 skipped, plus the DeepSeekFp8 and MxFp8 curated fuzzer cases. - Focused FP4 shared-expert tests in the CUDA container — 8 passed, including NVFP4, MXFP4, and W4A16 parity; all three curated FP4 fuzzer cases passed. - Focused cache-key integration tests for the unified contract, CUTLASS, and CuTeDSL — 3 passed. ## Reviewer Notes and Follow-ups - **Pre-routed inputs and routing replay:** callers can append shared IDs `E + k` and weights `1.0` themselves. Adding them inside FlashInfer would require re-striding `[T, K]` routing arrays to `[T, K + S]` or carrying dual-stride semantics through the post-routing pipeline. The replay buffer remains routed-only `[T, K]`, while the shared-expert routing kernel currently indexes it with the internal `K + S` stride, so replay is disabled for `S > 0`. Follow-up: fix the flat block-FP8 and FP4 replay paths while preserving the routed-only output contract. - **Cache invalidation:** the persisted-key change causes existing on-disk autotune entries to miss once and re-tune. This is benign but user-visible. - **EP remains unsupported:** the kernel maps a shared global ID to a weight row using `global_id - local_expert_offset`, which is valid only when the rank holds the full routed expert set. - **Physical versus routed rows:** `prepare_weights(num_local_experts=...)` takes the physical row count `E + S`, while `ExpertConfig.local_num_experts` remains routed-only `E`. - **Mixed quantization:** the current runners require homogeneous expert rows and therefore do not yet cover checkpoints such as `dsv4-flash-nvfp4`, whose routed experts use NVFP4 while shared experts use block-FP8. An initial implementation can compose an NVFP4 routed launch with a block-FP8 shared-expert launch and sum their BF16 outputs; a true single-launch implementation requires upstream kernel ABI and cubin support for separate routed/shared weight pointers, scale layouts, and dtypes. - **Test oracle:** shared-expert numerical tests compare with the legacy flat APIs because the existing self-fused pre-routed oracle cannot represent the important `S=1` and `S=2` cases under its expert-count divisibility constraints. - **Tracing:** artifacts were regenerated directly because `tests/trace/example.py` currently aborts in `recurrent_kda` with `nvidia-cutlass-dsl` 4.5.2. Follow-ups are to make Trace Apply select the concrete dispatched template per call, register reference-helper dependencies transitively and execute emitted functions in standalone tests, then repair example generation and add CI that regenerates `fi_trace_out` from an empty directory and rejects drift.

flashinfer-bot added the op: moe label Aug 6, 2026

feih-nv force-pushed the feih/unified-moe-runner-build-lifecycle branch from 2f77ea3 to fc61572 Compare August 6, 2026 09:54

feih-nv changed the title ~~Feih/unified moe runner build lifecycle~~ refactor(moe): migrate all runners to always call build() Aug 6, 2026

feih-nv force-pushed the feih/unified-moe-runner-build-lifecycle branch from fc61572 to f666f9c Compare August 7, 2026 03:33

feih-nv mentioned this pull request Aug 7, 2026

Fused shared experts support via the unified API #4326

Merged

3 tasks

feih-nv added 2 commits August 9, 2026 23:30

refactor(moe): migrate runners to build lifecycle

d8274b8

feat(autotuner): refine compound CUTLASS tactics

59dfdba

Retain the fastest tactics from each GEMM stage so the outer tuner can evaluate a bounded set of end-to-end combinations instead of assuming independent winners compose optimally. Developed with AI assistance.

feih-nv force-pushed the feih/unified-moe-runner-build-lifecycle branch from c9cb3bb to 59dfdba Compare August 10, 2026 06:42

feih-nv changed the title ~~refactor(moe): migrate all runners to always call build()~~ refactor(moe): enforce build() for all runners and refine CUTLASS autotuning in unified MoE API Aug 10, 2026

fix(moe): preserve staged tuning across cache hits

d62d240

Keep complete process-local tactic rankings so warm and persisted winner caches retain compound refinement, and repair the CuTe DSL build lifecycle configuration access.

feih-nv marked this pull request as ready for review August 10, 2026 09:20

feih-nv requested review from Aneureka, IwakuraRein, StudyingShao, aleozlx, bkryu, jiahanc, nv-yunzheq, qiching, samuellees and yzh119 as code owners August 10, 2026 09:20

aleozlx reviewed Aug 12, 2026

View reviewed changes

Comment thread flashinfer/fused_moe/runners.py

aleozlx reviewed Aug 12, 2026

View reviewed changes

aleozlx approved these changes Aug 12, 2026

View reviewed changes

feih-nv merged commit 0659712 into flashinfer-ai:main Aug 12, 2026
48 of 49 checks passed

feih-nv deleted the feih/unified-moe-runner-build-lifecycle branch August 12, 2026 06:09

Navigation Menu

Uh oh!

refactor(moe): enforce build() for all runners and refine CUTLASS autotuning in unified MoE API - #4376

refactor(moe): enforce build() for all runners and refine CUTLASS autotuning in unified MoE API#4376
feih-nv merged 3 commits into
flashinfer-ai:mainfrom
feih-nv:feih/unified-moe-runner-build-lifecycle

feih-nv commented Aug 6, 2026 •

edited

Loading

Uh oh!

coderabbitai Bot commented Aug 6, 2026 •

edited

Loading

Review failed

Uh oh!

feih-nv commented Aug 10, 2026

Uh oh!

flashinfer-bot commented Aug 10, 2026

Uh oh!

flashinfer-bot commented Aug 11, 2026

New relative to nightly (attribution uncertain)

Pre-existing failures

Timeouts, infrastructure, or incomplete jobs

Uh oh!

Uh oh!

aleozlx left a comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Uh oh!

Conversation

feih-nv commented Aug 6, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

📌 Description

Runner lifecycle

CUTLASS staged autotuning

🔍 Related Issues

🚀 Pull Request Checklist

✅ Pre-commit Checks

🧪 Tests

Reviewer Notes

Uh oh!

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

Uh oh!

Review failed

Uh oh!

feih-nv commented Aug 10, 2026

Uh oh!

flashinfer-bot commented Aug 10, 2026

Uh oh!

flashinfer-bot commented Aug 11, 2026

Unit Tests

Multi-GPU and Multi-Node Tests — 6/6 passed

New relative to nightly (attribution uncertain)

Pre-existing failures

Timeouts, infrastructure, or incomplete jobs

Uh oh!

Uh oh!

aleozlx left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

feih-nv commented Aug 6, 2026 •

edited

Loading

coderabbitai Bot commented Aug 6, 2026 •

edited

Loading