[None][chore] Remove one-warp-per-token policy from MoE A2A kernels

[bobboli]

Summary

Drop the WarpPolicy (one-warp-per-token) code path from MoE A2A dispatch and combine kernels, keeping only BlockPolicy (one-block-per-token).

Motivation

BlockPolicy was already the default. WarpPolicy is only beneficial when each token has very little work — i.e. very small hidden sizes — which is not representative of modern MoE models. The "many CTAs" concern that motivates WarpPolicy only triggers at batch sizes well above production. Keeping a code path that wins only outside the target regime — and silently regresses small/medium batches via the combine kernel — is not worth the maintenance cost.

Supporting data

B300x8, EP=8, NVLink one-sided, blog18 repro (tests/microbenchmarks/bench_moe_comm.py, --perfect_router, median dispatch+combine µs across ranks).

Modern hidden sizes — BlockPolicy wins across the full batch sweep:

config	batch	block (µs)	warp (µs)	block speedup
h=7168, FP8_BLOCK (DeepSeek-V3)	4096	1388	1427	+2.8%
h=4096, FP8_BLOCK	4096	816	858	+5.2%
h=4096, NVFP4	4096	576	603	+4.8%
h=4096, NO_QUANT	4096	818	856	+4.7%

WarpPolicy only catches up in a narrow corner — small hidden + aggressive quant + large batch:

config	batch	block (µs)	warp (µs)	warp speedup
h=1024, NVFP4	2048	132	104	+20.9%
h=1024, NVFP4	4096	229	189	+17.7%
h=2880, MXFP4 (gpt_oss)	4096	271	237	+12.7%

Combine-side WarpPolicy is multiple × slower than BlockPolicy at small/medium batches due to grid sizing (e.g. h=7168 FP8 b=1: 23 µs block vs 208 µs warp; h=4096 FP8 b=64: 23 µs vs 149 µs). This drags total latency far below block until b≈1024 across all configs.

What changes

• Remove WarpPolicy struct, SWITCH_POLICY macro, and all grid_size_warp paths in dispatch / prepare-combine / combine launches.
• Remove one_block_per_token field from MoeA2ADispatchParams and MoeA2ACombineParams.
• Remove getEnvMoeA2AOneBlockPerToken() and the TLLM_MOE_A2A_ONE_BLOCK_PER_TOKEN env knob.

No behavior change for users not overriding the env var.

[coderabbitai]

📝 Walkthrough

Walkthrough

This PR removes the one_block_per_token threading policy option from MoE A2A dispatch and combine kernels. The change eliminates an environment variable, removes related parameter fields, simplifies kernel implementations to use a single BlockPolicy variant with fixed one-block-per-token grid sizing, and removes associated operator initialization code.

Changes

MoE A2A Threading Policy Removal

Layer / File(s)	Summary
Environment variable API removal cpp/tensorrt_llm/common/envUtils.cpp, cpp/tensorrt_llm/common/envUtils.h	Removed getEnvMoeA2AOneBlockPerToken() function and declaration. Replaced the removed declaration with a DEV TODO section declaring new environment getter functions for MoE A2A dispatch and combine kernel block sizes (getEnvMoeA2ADispatchBlockSize() and getEnvMoeA2ACombineBlockSize()).
Kernel parameter struct cleanup cpp/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.h	Removed the one_block_per_token boolean field from both MoeA2ADispatchParams and MoeA2ACombineParams structures, changing their public struct layouts.
Kernel implementation and launch simplification cpp/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.cu	Removed WarpPolicy struct definition and SWITCH_POLICY macro. Simplified moeA2ADispatchKernel shared-memory layout to use a single fixed tile. Reworked moe_a2a_dispatch_launch, moe_a2a_prepare_combine_launch, and moe_a2a_combine_launch to always use BlockPolicy with grid size equal to token count (one block per token), removing all conditional warp-based branching and grid computation logic.
Operator parameter initialization cleanup cpp/tensorrt_llm/thop/moeAlltoAllOp.cpp	Removed environment-driven initialization of params.one_block_per_token from both moeA2ADispatchOp and moeA2ACombineOp, allowing parameters to use their default values.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~20 minutes

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 30.77% which is insufficient. The required threshold is 80.00%.	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 summarizes the main change: removing the warp-per-token policy from MoE A2A kernels, which aligns with the core objective of simplifying the kernel logic.
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.
Description check	✅ Passed	The PR description comprehensively explains the issue, motivation, supporting benchmarks, and exact changes. It clearly states the problem (WarpPolicy code path not worthwhile), provides performance data demonstrating BlockPolicy superiority, and lists all removed components.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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

[coderabbitai]

Actionable comments posted: 1

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)

cpp/tensorrt_llm/common/envUtils.cpp (1)

549-558: ⚠️ Potential issue | 🟠 Major | ⚡ Quick win

Treat malformed block-size env vars as invalid instead of throwing.

getIntEnv() uses std::stoi, so a value like TLLM_MOE_A2A_DISPATCH_BLOCK_SIZE=foo throws before sanitizeBlockSize() runs. That breaks the new “default 256 if unset or invalid” contract and turns a bad debug knob into a hard failure.

Proposed fix

+namespace
+{
+
+int getSanitizedBlockSize(char const* name)
+{
+    try
+    {
+        return sanitizeBlockSize(getIntEnv(name));
+    }
+    catch (std::exception const&)
+    {
+        TLLM_LOG_WARNING("Invalid value for %s. Falling back to 256.", name);
+        return sanitizeBlockSize(std::nullopt);
+    }
+}
+
+} // namespace
+
 int getEnvMoeA2ADispatchBlockSize()
 {
-    static int const kBlock = sanitizeBlockSize(getIntEnv("TLLM_MOE_A2A_DISPATCH_BLOCK_SIZE"));
+    static int const kBlock = getSanitizedBlockSize("TLLM_MOE_A2A_DISPATCH_BLOCK_SIZE");
     return kBlock;
 }
 
 int getEnvMoeA2ACombineBlockSize()
 {
-    static int const kBlock = sanitizeBlockSize(getIntEnv("TLLM_MOE_A2A_COMBINE_BLOCK_SIZE"));
+    static int const kBlock = getSanitizedBlockSize("TLLM_MOE_A2A_COMBINE_BLOCK_SIZE");
     return kBlock;
 }

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@cpp/tensorrt_llm/common/envUtils.cpp` around lines 549 - 558, The
getEnvMoeA2ADispatchBlockSize and getEnvMoeA2ACombineBlockSize rely on getIntEnv
which calls std::stoi and will throw for malformed values (e.g. "foo"); change
these accessors to treat parse errors as "no value" so sanitizeBlockSize can
apply the default. Concretely, wrap the integer parsing used by getIntEnv (or
add a safe wrapper) to catch std::invalid_argument/std::out_of_range and return
an indicator of "unset/invalid" (e.g., optional/negative sentinel) instead of
letting the exception propagate, then call sanitizeBlockSize with that safe
result in both getEnvMoeA2ADispatchBlockSize and getEnvMoeA2ACombineBlockSize so
malformed env vars fall back to the default 256 behavior. Ensure you reference
getIntEnv, sanitizeBlockSize, getEnvMoeA2ADispatchBlockSize, and
getEnvMoeA2ACombineBlockSize when making the change.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@cpp/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.cu`:
- Around line 625-630: The launch block containing SWITCH_BOOL(...,
SWITCH_TOP_K(..., { ... launchWithPdlWhenEnabled("moeA2ADispatchKernel",
kernel_fn, ... ) }) is failing pre-commit formatting; run clang-format (or your
repo's formatter) on this region so the SWITCH_BOOL / SWITCH_TOP_K nesting, the
auto kernel_fn = moeA2ADispatchKernel<BlockPolicy, TOP_K, EPLB_STATS>; line, and
the launchWithPdlWhenEnabled(...) call are reformatted to match project style
and fix pre-commit checks.

---

Outside diff comments:
In `@cpp/tensorrt_llm/common/envUtils.cpp`:
- Around line 549-558: The getEnvMoeA2ADispatchBlockSize and
getEnvMoeA2ACombineBlockSize rely on getIntEnv which calls std::stoi and will
throw for malformed values (e.g. "foo"); change these accessors to treat parse
errors as "no value" so sanitizeBlockSize can apply the default. Concretely,
wrap the integer parsing used by getIntEnv (or add a safe wrapper) to catch
std::invalid_argument/std::out_of_range and return an indicator of
"unset/invalid" (e.g., optional/negative sentinel) instead of letting the
exception propagate, then call sanitizeBlockSize with that safe result in both
getEnvMoeA2ADispatchBlockSize and getEnvMoeA2ACombineBlockSize so malformed env
vars fall back to the default 256 behavior. Ensure you reference getIntEnv,
sanitizeBlockSize, getEnvMoeA2ADispatchBlockSize, and
getEnvMoeA2ACombineBlockSize when making the change.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 80651b15-7e08-4e09-bdf0-9b7dffc84353

📥 Commits

Reviewing files that changed from the base of the PR and between 526d7ee and 7839c5b.

📒 Files selected for processing (5)

• cpp/tensorrt_llm/common/envUtils.cpp
• cpp/tensorrt_llm/common/envUtils.h
• cpp/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.cu
• cpp/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.h
• cpp/tensorrt_llm/thop/moeAlltoAllOp.cpp

[bobboli]

Addressed CodeRabbit feedback in ec26ba0f76:

1. clang-format applied to moeAlltoAllKernels.cu launch block.
2. Wrapped sanitizeBlockSize(getIntEnv(...)) in getSanitizedBlockSizeFromEnv() helper that catches std::stoi exceptions, logs a warning, and falls back to default — malformed block-size env vars no longer crash.

[bobboli]

/bot run

[tensorrt-cicd]

PR_Github #50289 [ run ] triggered by Bot. Commit: ec26ba0 Link to invocation

[tensorrt-cicd]

PR_Github #50289 [ run ] completed with state FAILURE. Commit: ec26ba0
/LLM/main/L0_MergeRequest_PR pipeline #39819 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

[bobboli]

/bot run

[tensorrt-cicd]

PR_Github #50409 [ run ] triggered by Bot. Commit: ec26ba0 Link to invocation

[tensorrt-cicd]

PR_Github #50409 [ run ] completed with state SUCCESS. Commit: ec26ba0
/LLM/main/L0_MergeRequest_PR pipeline #39933 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

[bobboli]

/bot run

[tensorrt-cicd]

PR_Github #50447 [ run ] triggered by Bot. Commit: ec26ba0 Link to invocation

[tensorrt-cicd]

PR_Github #50447 [ run ] completed with state SUCCESS. Commit: ec26ba0
/LLM/main/L0_MergeRequest_PR pipeline #39966 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

[bobboli]

/bot run

[tensorrt-cicd]

PR_Github #50503 [ run ] triggered by Bot. Commit: ec26ba0 Link to invocation

[tensorrt-cicd]

PR_Github #50503 [ run ] completed with state SUCCESS. Commit: ec26ba0
/LLM/main/L0_MergeRequest_PR pipeline #40011 completed with status: 'SUCCESS'

CI Report

Link to invocation