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

[jimmyzho]

…#4237)

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

(cherry picked from commit 8b6af1a)

📌 Description

🔍 Related Issues

🚀 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

• 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.

If you are unsure about how to set up pre-commit, see the pre-commit documentation.

🧪 Tests

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

Reviewer Notes

[gemini-code-assist]

Caution

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

[coderabbitai]

Important

Review skipped

Auto reviews are disabled on base/target branches other than the default branch.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro Plus

Run ID: ed3ff6a1-02d9-4172-9c45-a5761e158c4f

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

---

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