fix(moe): pad trtllm-gen route map by one element to avoid OOB read

[syuoni]

📌 Description

The routed dynamic-batch batched-GEMM kernels (Bmm_*_dynB with -routeAct)
read one int32 past the end of ptrRouteMap, from the last batch-dim CTA, on
every launch. The read is speculative — its value is never consumed, so outputs
are always bitwise correct — but it faults with CUDA_ERROR_ILLEGAL_ADDRESS
whenever the allocation happens to end at a mapped-region boundary. That
allocator-placement dependence is what makes it surface as flaky MoE autotune
and inference crashes.

Root cause is kernel-side: an off-by-one clamp in the hoisted load-task
initializer (WarpGrpThreadIdx is clamped to the load-group size instead of
size − 1), so off-group threads compute row tileN of a tileN-row tile, i.e.
index (ctaIdxY + 1) * tileN. For the last CTA that is exactly one element past
the shape documented in KernelParamsDecl.h
([sum(divUpMul(N[bi], tileN) for bi in B)]) — which is what we allocate. It has
been reported to the kernel owners and is being fixed there.

This PR is the integration-side workaround: allocate the route map with one
extra element, so the already-shipped prebuilt cubins stay in bounds. The
overrun is provably always exactly one int32 from one CTA, so +1 is
sufficient by construction, not a heuristic. The pad slot's contents are
irrelevant since the value is never used.

Two allocation sites, both permuted_idx_to_token_idx (→ routeMap). Marked
WAR + TODO so the +1 can be dropped once regenerated cubins land.

🔍 Related Issues

Likely explains #3530 and #3168 (and possibly #4012, #2776) — all report
intermittent NVFP4 MoE autotune crashes, and #3168 additionally reports silent
garbage output, which is the expected signature of an OOB read that usually
lands on mapped memory.

🚀 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.).

No test added — the failure is allocator-placement dependent and does not
reproduce deterministically through the public API. Validation was done on
v0.6.15.post1 (B200, SM100, Kimi-K2.5 NVFP4 TP4):

Check	stock	route map +1
compute-sanitizer memcheck, M=2 bucket × 4 ranks	40 × Invalid __global__ read of size 4	0 violations (816 profiles/rank)
Guard-page probe (route map placed at the tail of a VMM mapping)	faults on every launch	clean
Full-model TP4 autotune crash loop	crashed by attempt 2 in 3/3 loops	6/6 clean

The guard-page probe is the load-bearing one: it removes all dependence on
allocator placement, so it is deterministic in both directions.

Summary by CodeRabbit

• Bug Fixes
  • Improved stability for fused mixture-of-experts routing operations by preventing out-of-bounds memory access in supported execution paths.
  • Applied the safeguard to both standard and FP4 block-scale routing workflows.

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro Plus

Run ID: 231372bc-fb41-4fad-98ea-79051b701424

📥 Commits

Reviewing files that changed from the base of the PR and between 7adc546 and 0b1f2ef.

📒 Files selected for processing (1)

• csrc/trtllm_fused_moe_kernel_launcher.cu

---

📝 Walkthrough

Walkthrough

The MoE routing workspace allocation for permuted_idx_to_token_idx increases by one int32 element in the standard and FP4 block-scale launchers, with comments documenting the temporary workaround.

Changes

Routing workspace padding

Layer / File(s)	Summary
Route-map allocation padding csrc/trtllm_fused_moe_kernel_launcher.cu	Standard and FP4 block-scale routing preparation allocate max_num_padded_tokens + 1 elements for permuted_idx_to_token_idx and document the temporary out-of-bounds-read workaround.

Estimated code review effort: 1 (Trivial) | ~5 minutes

Suggested reviewers: iwakurarein, aleozlx, jdebache, djns99, nv-yunzheq

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly and concisely describes the main change: padding the route map to avoid an out-of-bounds read.
Description check	✅ Passed	The description covers the change, rationale, related issues, checklist, and validation details, so it is mostly complete.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches 💡 1

🛠️ Fix failing CI checks 💡

• Fix failing CI checks

🧪 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

• X
• Mastodon
• Reddit
• LinkedIn

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

[samuellees]

/bot run tests/moe

[zhyncs]

Per our discussion, let's take this hot fix for now. The proper fix is expected to take longer to land, and we can replace this workaround once it's ready.

[flashinfer-bot]

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

[samuellees]

cc @rosenrodt for viz

[rosenrodt]

Thanks for the fix. Will notify here when cubin fix lands

[samuellees]

LGTM. For CI fails on GB300, it's about AllReduce and not relative with this PR. And this is not relative with 5090.