fix(moe): write fused-shared routing replay as routed-only [T, K] by feih-nv · Pull Request #4894 · flashinfer-ai/flashinfer

feih-nv · 2026-09-02T10:05:55Z

📌 Description

Fixes #5009 part 1. DeepSeek routingMainKernel wrote routing_replay_out with the packed index token * (K+S) + k. Replay is [T, K] (routed ids only), so token 1 starts at K+S instead of K. S == 0 hid it; host/C++ rejected S > 0 (#4239), so the combo was unavailable.

This PR indexes replay with token * K + k and drops the S > 0 rejects. Packed ids/weights stay [T, K+S].

🔍 Related Issues

[Bug] routing_replay_out out-of-bounds writes: fused-shared stride and missing dim0 lower bound #5009
fix(moe): bound routing_replay_out dim0 from below in both validators #5072
feat(moe): support shared expert fusion for trtllm-gen fp4 moe #4239

🚀 Pull Request Checklist

✅ Pre-commit Checks

I have installed pre-commit by running pip install pre-commit (or used your 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

Tests have been added or updated as needed.
All tests are passing (unittest, etc.).

CPU reject

dim1 == K+S
non-DeepSeek + S > 0

GPU (SM100) accept

FP8 / FP4 fused-shared replay, T ∈ {8,32}, S ∈ {1,2}
one FP8 CUDA-graph capture on oversized [32, K]

Reviewer Notes

With S > 0, routing_replay_out is still [T, K] int16: the same K routed expert ids as S = 0 on the same logits, ids in [0, E). The S fused-shared slots live only in the internal packed ids/weights ([T, K+S], ids E … E+S-1). They are not appended to replay. A [T, K+S] replay buffer is rejected (dim1 must equal top_k).

Summary by CodeRabbit

New Features
- Routing replay is now supported with fused shared experts for FP8 and FP4 MoE operations.
- Replay buffers retain a routed-only layout with width equal to top_k; shared-expert slots are not recorded.
- Support includes oversized buffers and CUDA graph execution.
Documentation
- Updated integration and testing guidance for routed-only replay behavior.
Tests
- Added coverage for replay contents, layouts, sentinel rows, and FP8/FP4 execution paths.

coderabbitai · 2026-09-02T10:06:04Z

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Advanced

Run ID: 00351ae2-3509-4514-96c7-17ba09aeb47f

📥 Commits

Reviewing files that changed from the base of the PR and between 8eb2823 and a65baec.

📒 Files selected for processing (5)

csrc/trtllm_fused_moe_kernel_launcher.cu
flashinfer/fused_moe/backends/prims_ts/fp4_op.py
flashinfer/fused_moe/core.py
tests/moe/test_trtllm_gen_fused_moe.py
tests/moe/test_trtllm_gen_routed_fused_moe.py

💤 Files with no reviewable changes (1)

flashinfer/fused_moe/backends/prims_ts/fp4_op.py

Included review availability: Your plan provides up to 8 included reviews per hour; 7 remain after this review.

📝 Walkthrough

Walkthrough

Routing replay now uses a routed-only [num_tokens, top_k] layout with fused shared experts. Validation requires sufficient row capacity and the top_k width. FP8, FP4, and CUDA-graph tests cover the behavior.

Changes

Routing replay support

Layer / File(s)	Summary
Replay contract and validation `flashinfer/fused_moe/core.py`, `csrc/trtllm_fused_moe_kernel_launcher.cu`, `flashinfer/fused_moe/backends/prims_ts/fp4_op.py`, `tests/moe/test_trtllm_gen_fused_moe.py`	Replay validation requires `dim1 == top_k` and `dim0 >= num_tokens`. FP8 and FP4 fused-shared paths accept routed-only replay output. Invalid widths remain rejected.
Routed-only kernel layout `csrc/fused_moe/trtllm_backend/trtllm_fused_moe_routing_deepseek.cu`	The routing kernel writes replay rows with `top_k` stride and excludes fused shared-expert slots.
Behavior coverage and documentation `tests/moe/test_trtllm_gen_routed_fused_moe.py`, `docs/design_docs/moe_routing_test_decomposition.md`, `docs/vllm_routing_replay_integration.md`	Tests cover FP8, FP4, fused shared experts, oversized CUDA-graph buffers, replay contents, and invalid widths. Documentation describes the routed-only layout and test coverage location.

Priority: ➖ Normal

Estimated code review effort: 3 (Moderate) | ~25 minutes

Change: Bug fix · Severity of issue fixed: Medium

Merge Risk: ⚪ Minimal · up to a65ba

Replay buffers now retain the routed-only layout while accepting valid oversized CUDA-graph allocations. No concrete merge-blocking issue remains.

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 58.33% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 24 functions across 5 files.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly identifies the primary change: writing fused-shared routing replay in the routed-only [T, K] layout.
Description check	✅ Passed	The description explains the bug, the indexing fix, validation changes, related issues, tests, and reviewer focus areas. It is mostly complete and matches the required template. The checked statement …
Linked Issues check	✅ Passed	The PR meets the coding requirements in issue `#5009`. In `routingMainKernel`, replay writes use `blockIdx.x * params.mTopK + laneIdx`. Packed IDs and weights keep the `K+S` layout. Host and C++ validat…
Out of Scope Changes check	✅ Passed	The changes stay within issue `#5009`. The kernel change fixes replay indexing. Validation changes enable the required DeepSeek fused-shared configuration and preserve invalid-width rejection. The tests…

Fix all pre-merge checks with AI

✨ Finishing Touches 💡 1

🛠️ Fix failing CI checks 💡

Create stacked PR
Commit on current branch

🧪 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.}

feih-nv · 2026-09-03T11:09:28Z

@flashinfer-bot run

feih-nv · 2026-09-03T11:09:38Z

/bot run tests/moe

flashinfer-bot · 2026-09-03T11:10:22Z

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

flashinfer-bot · 2026-09-03T15:55:01Z

[FAILED] Pipeline #66024680 — 14/17 executed test jobs passed

Compared with nightly #65814627 (different CI configuration).

Unit Tests

GPU	CUDA 12.9	CUDA 13.0	Other	Notes
B200	✅ Pass	✅ Pass	—	—
GB200	✅ Pass	✅ Pass	—	—
GB300	✅ Pass	✅ Pass	—	—
H100	✅ Pass	✅ Pass	—	—
RTX Pro 6000 Blackwell	❔ Unknown	❔ Unknown	—	Unknown: `script failed before producing a JUnit report` (2 jobs; CUDA 12.9, CUDA 13.0)
VR200 CU134	—	—	❔ Unknown	Not compared: `tests.moe.test_trtllm_gen_routing` (668 failures) Not compared: `tests.moe.test_unified_moe` (3 failures) Not compared: `tests.moe.test_unified_moe_mxfp4` (2 failures) … and 1 more

✅ 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	Other	Notes
B300 (multi-GPU)	✅ Pass	✅ Pass	—	—
GB200 (multi-node)	✅ Pass	✅ Pass	—	—
GB300 (multi-node)	✅ Pass	✅ Pass	—	—

Failure details

Could not compare

tests.moe.test_trtllm_gen_routing — 668 failures on VR200 CU134
- flashinfer.utils.BackendSupportedError: trtllm_gen_routing does not support compute capability 107
tests.moe.test_unified_moe — 3 failures on VR200 CU134
- NotImplementedError: Custom swiglu_alpha/swiglu_beta/swiglu_limit are not supported by the Rubin (SM107) gather grouped GEMM kernel yet.
tests.moe.test_unified_moe_mxfp4 — 2 failures on VR200 CU134
- RuntimeError: MoELayer: none of the configured backends ['TrtllmFp4Config'] are usable on arch sm107 for this configuration. Registered unified runners: [CutlassBf16Config, Cutl…
tests.moe.test_unified_moe_fuzz — 1 failure on VR200 CU134
- Failed: trtllm_mxint4_routed mxint4_swiglu_Llama4_hot1_e256_L128o128_k1_t4095_h1024_i256_s6: 1/4193280 elems exceed tol (rtol=0.3 atol=223; max|diff|=271.8, ‖ref‖∞=3424) CONFIG…

Timeouts, infrastructure, or incomplete jobs

Unknown: script failed before producing a JUnit report — RTX Pro 6000 Blackwell / CUDA 12.9, RTX Pro 6000 Blackwell / CUDA 13.0
- Jobs: unit_test_rtx-pro-6000-blackwell: [cu130], unit_test_rtx-pro-6000-blackwell: [cu129]

@feih-nv

Two review points from @feih-nv, who filed flashinfer-ai#5009. `num_tokens` was `Optional[int] = None` guarded by `is not None`, so a future entry point that forgot to pass it would silently skip the bound. Make it a required parameter: forgetting it is now a TypeError at the call site rather than a missing check at runtime. All seven callers already passed it; the two that also pass `num_fused_shared_experts` now name both keywords, since the parameter order changed. Add `test_routing_replay_out_rejects_undersized_dim0`, lifted from @feih-nv's `b81f0946` on flashinfer-ai#4894 with permission — FP8 and FP4, CPU tensors, no GPU needed. It pins the rejection this PR adds, which had no test of its own. flashinfer-ai#4894 drops that commit once this lands so the two do not collide on the validator and its call sites. Co-Authored-By: feih <feih@nvidia.com> Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

feih-nv · 2026-09-10T12:00:01Z

/bot run tests/moe

feih-nv · 2026-09-10T12:00:52Z

@flashinfer-bot run

flashinfer-bot · 2026-09-10T12:01:08Z

GitLab MR !1402 has been updated with latest changes, and the CI pipeline #67171205 is currently running. I'll report back once the pipeline job completes.

flashinfer-bot · 2026-09-10T17:11:10Z

[FAILED] Pipeline #67171205 — 10/17 executed test jobs passed

Compared with nightly #66937827.

Unit Tests

GPU	CUDA 12.9	CUDA 13.0	Other	Notes
B200	🟡 Old	🟡 Old	—	Old: `tests.moe.test_trtllm_gen_moe_autotune_tactics` (123 failures; CUDA 12.9, CUDA 13.0)
GB200	🟡 Old	❌ New	—	New: `tests.moe.test_trtllm_gen_moe_autotune_tactics` (1 failure; CUDA 13.0) Old: `tests.moe.test_trtllm_gen_moe_autotune_tactics` (123 failures; CUDA 12.9, CUDA 13.0)
GB300	❔ Unknown	❌ New	—	New: `tests.moe.test_trtllm_gen_moe_autotune_tactics` (1 failure; CUDA 13.0) Old: `tests.moe.test_trtllm_gen_moe_autotune_tactics` (70 failures; CUDA 13.0) Not compared: `tests.moe.test_trtllm_gen_moe_autotune_tactics` (74 failures; CUDA 12.9)
H100	✅ Pass	✅ Pass	—	—
RTX Pro 6000 Blackwell	✅ Pass	✅ Pass	—	—
VR200 CU134	—	—	❔ Unknown	Unknown: `script failed before producing a JUnit report` (1 job)

✅ 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	Other	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.moe.test_trtllm_gen_moe_autotune_tactics — 2 failures on GB200 / CUDA 13.0, GB300 / CUDA 13.0
- AssertionError: [NvFP4xNvFP4] forced tactic was not dispatched — autotuner did not log a cache hit; check '_moe_profile_shapes' against the actual MoEInputs layout. assert ('fla…

Pre-existing failures

tests.moe.test_trtllm_gen_moe_autotune_tactics — 316 failures on B200 / CUDA 12.9, B200 / CUDA 13.0, GB200 / CUDA 12.9, GB200 / CUDA 13.0, GB300 / CUDA 13.0
- AssertionError: the forced regression tactic was not dispatched through the autotuner cache assert ('flashinfer::trtllm_fp4_block_scale_moe', 'MoERunner') in set() + where set()…

Could not compare

tests.moe.test_trtllm_gen_moe_autotune_tactics — 74 failures on GB300 / CUDA 12.9
- AssertionError: the forced regression tactic was not dispatched through the autotuner cache assert ('flashinfer::trtllm_fp4_block_scale_moe', 'MoERunner') in set() + where set()…

Timeouts, infrastructure, or incomplete jobs

Unknown: script failed before producing a JUnit report — VR200 CU134
- Jobs: unit_test_vr200_cu134

feih-nv · 2026-09-11T07:20:55Z

/bot run tests/moe

feih-nv · 2026-09-11T07:21:04Z

@flashinfer-bot run

flashinfer-bot · 2026-09-11T07:21:04Z

GitLab MR !1402 has been updated with latest changes, and the CI pipeline #67329403 is currently running. I'll report back once the pipeline job completes.

flashinfer-bot · 2026-09-11T14:25:08Z

[FAILED] Pipeline #67329403 — 10/19 executed test jobs passed

Compared with nightly #67144516 (different CI configuration).

Unit Tests

GPU	CUDA 12.9	CUDA 13.0	CUDA 13.4	Notes
B200	🟡 Old	❔ Unknown	—	New: `tests.moe.test_trtllm_gen_moe_autotune_tactics` (3 failures; CUDA 13.0) Old: `tests.moe.test_trtllm_gen_moe_autotune_tactics` (122 failures; CUDA 12.9, CUDA 13.0)
GB200	🟡 Old	❔ Unknown	—	New: `tests.moe.test_trtllm_gen_moe_autotune_tactics` (3 failures; CUDA 13.0) Old: `tests.moe.test_trtllm_gen_moe_autotune_tactics` (122 failures; CUDA 12.9, CUDA 13.0)
GB300	❔ Unknown	❔ Unknown	—	New: `tests.moe.test_trtllm_gen_moe_autotune_tactics` (7 failures; CUDA 12.9, CUDA 13.0) Old: `tests.moe.test_trtllm_gen_moe_autotune_tactics` (143 failures; CUDA 12.9, CUDA 13.0)
H100	✅ Pass	✅ Pass	—	—
RTX Pro 6000 Blackwell	⚠️ Infra	⚠️ Infra	—	Infrastructure: `CI infrastructure failure` (2 jobs; CUDA 12.9, CUDA 13.0)
VR200	—	—	❔ Unknown	Unknown: `script failed before producing a JUnit report` (1 job; CUDA 13.4)

✅ 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	CUDA 13.4	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.moe.test_trtllm_gen_moe_autotune_tactics — 13 failures on B200 / CUDA 13.0, GB200 / CUDA 13.0, GB300 / CUDA 12.9, GB300 / CUDA 13.0
- AssertionError: [MxFP4xMxFP8] forced tactic was not dispatched — autotuner did not log a cache hit; check '_moe_profile_shapes' against the actual MoEInputs layout. assert ('fla…

Pre-existing failures

tests.moe.test_trtllm_gen_moe_autotune_tactics — 387 failures on B200 / CUDA 12.9, B200 / CUDA 13.0, GB200 / CUDA 12.9, GB200 / CUDA 13.0, GB300 / CUDA 12.9, GB300 / CUDA 13.0
- AssertionError: the forced regression tactic was not dispatched through the autotuner cache assert ('flashinfer::trtllm_fp4_block_scale_moe', 'MoERunner') in set() + where set()…

Timeouts, infrastructure, or incomplete jobs

Infrastructure: CI infrastructure failure — RTX Pro 6000 Blackwell / CUDA 12.9, RTX Pro 6000 Blackwell / CUDA 13.0
- Jobs: unit_test_rtx-pro-6000-blackwell: [cu130], unit_test_rtx-pro-6000-blackwell: [cu129]
Unknown: script failed before producing a JUnit report — VR200 / CUDA 13.4
- Jobs: unit_test_vr200_cu134

DeepSeek routing packed ids at stride top_k + S, so replay used the same index and scrambled multi-token [T, K] buffers. Write replay at stride top_k, drop the S>0 host/C++ rejects, and cover FP8/FP4 plus a CUDA-graph capture of the oversized buffer. AI-assisted.

Deleting it read as done. The fused-shared replay coverage landed in the fused-MoE tests instead; the standalone routing harness still has no S>0 case, because it would first have to build fused-shared weight rows.

Drop the int64_t casts (packed idxTopK is 32-bit) and put dim0 before dim1 in the C++ replay validator comment. AI-assisted.

feih-nv · 2026-09-14T03:11:10Z

@flashinfer-bot run

feih-nv · 2026-09-14T03:11:41Z

/bot run tests/moe

flashinfer-bot · 2026-09-14T03:12:00Z

GitLab MR !1402 has been updated with latest changes, and the CI pipeline #67710173 is currently running. I'll report back once the pipeline job completes.

flashinfer-bot · 2026-09-14T05:15:25Z

[SUCCESS] Pipeline #67710173: 18/19 executed test jobs passed

flashinfer-bot added the op: moe label Sep 2, 2026

feih-nv marked this pull request as ready for review September 3, 2026 11:08

feih-nv requested review from Anerudhan, Aneureka, IwakuraRein, StudyingShao, aleozlx, bkryu, cyx-6, dhiraj113, jiahanc, jimmyzho, kahyunnam, nv-yunzheq, qiching, qsang-nv, saltyminty, samuellees, sricketts, yongwww, yyihuang and yzh119 as code owners September 3, 2026 11:08

flashinfer-bot added the run-ci label Sep 3, 2026

This was referenced Sep 7, 2026

[Bug] routing_replay_out out-of-bounds writes: fused-shared stride and missing dim0 lower bound #5009

Closed

fix(moe): bound routing_replay_out dim0 from below in both validators #5072

Merged

feih-nv force-pushed the fix/moe-fused-shared-routing-replay branch from 8a42766 to 3b44f9d Compare September 10, 2026 09:43

flashinfer-bot added run-ci and removed run-ci labels Sep 10, 2026

flashinfer-bot added run-ci and removed run-ci labels Sep 11, 2026

feih-nv added 3 commits September 13, 2026 19:08

docs(moe): keep the standalone-routing replay follow-up

68e3e97

Deleting it read as done. The fused-shared replay coverage landed in the fused-MoE tests instead; the standalone routing harness still has no S>0 case, because it would first have to build fused-shared weight rows.

fix(moe): name DeepSeek replay index like packed writes

a65baec

Drop the int64_t casts (packed idxTopK is 32-bit) and put dim0 before dim1 in the C++ replay validator comment. AI-assisted.

feih-nv force-pushed the fix/moe-fused-shared-routing-replay branch from 8eb2823 to a65baec Compare September 14, 2026 02:11

feih-nv requested review from Jingfan-Sun and YangXu1990uiuc as code owners September 14, 2026 02:11

flashinfer-bot added run-ci and removed run-ci labels Sep 14, 2026

aleozlx approved these changes Sep 16, 2026

View reviewed changes

aleozlx merged commit 7554a6e into flashinfer-ai:main Sep 16, 2026
26 of 27 checks passed

feih-nv deleted the fix/moe-fused-shared-routing-replay branch September 16, 2026 06:09

Uh oh!

Conversation

feih-nv commented Sep 2, 2026 • edited by coderabbitai Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

📌 Description

🔍 Related Issues

🚀 Pull Request Checklist

✅ Pre-commit Checks

🧪 Tests

CPU reject

GPU (SM100) accept

Reviewer Notes

Summary by CodeRabbit

Uh oh!

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

Uh oh!

Walkthrough

Changes

❌ Failed checks (1 warning)

Uh oh!

feih-nv commented Sep 3, 2026

Uh oh!

feih-nv commented Sep 3, 2026

Uh oh!

flashinfer-bot commented Sep 3, 2026

Uh oh!

flashinfer-bot commented Sep 3, 2026

Unit Tests

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

Could not compare

Timeouts, infrastructure, or incomplete jobs

Uh oh!

feih-nv commented Sep 10, 2026

Uh oh!

feih-nv commented Sep 10, 2026

Uh oh!

flashinfer-bot commented Sep 10, 2026

Uh oh!

flashinfer-bot commented Sep 10, 2026

Unit Tests

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

New relative to nightly (attribution uncertain)

Pre-existing failures

Could not compare

Timeouts, infrastructure, or incomplete jobs

Uh oh!

feih-nv commented Sep 11, 2026

Uh oh!

feih-nv commented Sep 11, 2026

Uh oh!

flashinfer-bot commented Sep 11, 2026

Uh oh!

flashinfer-bot commented Sep 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!

feih-nv commented Sep 14, 2026

Uh oh!

feih-nv commented Sep 14, 2026

Uh oh!

flashinfer-bot commented Sep 14, 2026

Uh oh!

flashinfer-bot commented Sep 14, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

feih-nv commented Sep 2, 2026 •

edited by coderabbitai Bot

Loading

coderabbitai Bot commented Sep 2, 2026 •

edited

Loading