[https://nvbugs/6457853][fix] Allow trtllm-gen MoE autotuner when local_num_experts < top_k

[dongfengy]

Description

Fixes https://nvbugs/6457853 — gpt-oss-120b (128 experts, top_k=4) crashes at autotuner warmup when EP64 leaves 128/64 = 2 local experts per rank:

AssertionError: random_local requires local_num_experts >= top_k; got local_num_experts=2, top_k=4

To time MoE kernel tactics, the autotuner fabricates a dummy token→expert assignment instead of running the real router. Both dummy generators (random and balanced, per TRTLLM_GEN_MOE_AUTOTUNE_DUMMY_DISTRIBUTION) required top_k distinct LOCAL ids per row, which is impossible once EP > num_experts / top_k. The real inference path has no such limit: routing picks top_k distinct GLOBAL experts, of which at most local_num_experts live on any rank — so this was autotuner-only.

Fix: keep dummy rows production-shaped — fill min(top_k, local_num_experts) distinct local ids and pad the remaining slots with deterministic out-of-shard ids (remote slots do no local work). No in-row duplicates; zero behavior change when local_num_experts >= top_k.

Verified on main (b8604c4, GB200): the 1-GPU repro hits the exact assert on stock code and passes full autotune profiling with this fix; at TP64/EP64 on 16×4 GB200 NVL72, stock fails on all 64 ranks at warmup while the fixed run completes trtllm-bench cleanly (8×32k/8k requests, 48 min). The same A/B was also reproduced on 1.3.0rc18 and rc21.

Test Coverage

• New: tests/unittest/_torch/misc/test_autotuner.py::test_trtllm_gen_moe_dummy_topk_local_experts_less_than_topk[random|balanced] — fails before the fix, passes after; checks rows are production-shaped (distinct ids, all local experts present, pad ids valid).

PR Checklist

Please review the following before submitting your PR:

• PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.

• PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.

• Test cases are provided for new code paths (see test instructions)

• If PR introduces API changes, an appropriate PR label is added - either api-compatible or api-breaking. For api-breaking, include BREAKING in the PR title.

• Any new dependencies have been scanned for license and vulnerabilities

• CODEOWNERS updated if ownership changes

• Documentation updated as needed

• Update tava architecture diagram if there is a significant design change in PR.

• The reviewers assigned automatically/manually are appropriate for the PR.

• Please check this after reviewing the above items as appropriate for this PR.

🤖 Generated with Claude Code

[coderabbitai]

📝 Walkthrough

Walkthrough

Dummy top-k generation now tolerates expert pools smaller than top_k by allowing modular duplicates in balanced mode and falling back to that mode for random local generation. A parametrized autotuner test validates both distributions.

Changes

Dummy top-k generation

Layer / File(s)	Summary
Balanced modular top-k generation tensorrt_llm/_torch/custom_ops/trtllm_gen_custom_ops.py	Balanced dummy generation allows modular wraparound duplicates when the target expert pool is smaller than top_k.
Local distribution fallback and validation tensorrt_llm/_torch/custom_ops/trtllm_gen_custom_ops.py, tests/unittest/_torch/misc/test_autotuner.py	Random local generation falls back to balanced generation for undersized pools, with parametrized coverage validating shapes, dtypes, and expert bounds for both distributions.

Estimated code review effort: 2 (Simple) | ~10 minutes

Suggested reviewers: cascade812

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
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.
Title check	✅ Passed	The title matches the main change and follows the required ticket/type/summary format.
Description check	✅ Passed	The description includes the required Description, Test Coverage, and PR Checklist sections with clear details.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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

[dongfengy]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #61035 [ run ] triggered by Bot. Commit: 87dbbce Link to invocation

[tensorrt-cicd]

PR_Github #61035 [ run ] completed with state FAILURE. Commit: 87dbbce
/LLM/main/L0_MergeRequest_PR pipeline #49290 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[dongfengy]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #61080 [ run ] triggered by Bot. Commit: 87dbbce Link to invocation

[tensorrt-cicd]

PR_Github #61080 [ run ] completed with state SUCCESS. Commit: 87dbbce
/LLM/main/L0_MergeRequest_PR pipeline #49335 completed with status: 'SUCCESS'

CI Report

Link to invocation