feat(cute_dsl/moe): deterministic balanced autotune profile inputs

[leejnau]

Summary

Reduces fi-vs-trt parity distance for CuteDSL NVFP4 MoE autotune selections by giving the autotuner a deterministic, balanced approx-max-load distribution for token_selected_experts during profiling, plus seeded random tensor initializers for the other dynamic inputs. Brings fi's autotune-profile machinery into structural alignment with trt-llm's inputs_pre_hook mechanism (all 6 of trt-llm's CuteDSL TuningConfigs set this hook).

Three commits on the branch:

• 8faa2359 — Path A+C: seed all 4 tensor_initializers + port TuningConfig.inputs_pre_hook + CuteDslMoEInputsHelper (the load-bearing change)
• a009ca18 — small parity-cleanup: align use_cold_l2_cache=True with the rest of the codebase / trt-llm (provably no-op for DeepSeek-V3, value is forward-looking; flagged separately so reviewers don't conflate it with the A+C perf claim)
• 819bd648 — pin the CuteDslMoEInputsHelper input-layout contract with a small unit test

Empirical validation

10-run 30-cell sweep on NGC rc14, B200, DeepSeek-V3 (hidden=7168, intermediate=2048, num_experts=256, top_k=8):

	baseline	with this patch	delta
Overall mean |Δ%| from trt parity	8.48%	6.45%	−2.03pp
Small (N=1..256)	10.50%	7.44%	−3.07pp
Mid (N=512..2048)	6.54%	5.80%	−0.74pp
Large (N=4096..16384)	4.35%	4.13%	−0.21pp

Cells closer to parity: 19 / further: 7 / neutral: 19. Cells now within 1% of trt parity: 19 of 45 (was ~12 of 45 baseline).

Trade-offs (deliberately retained for parity gain)

• N=4096 EP=1 regresses from +1.24% to +8.89% (single cell; deterministic seed lands on a sub-optimal pick at this bucket).
• N=2048 EP=16 still bimodal under the seeded distribution; mean +19.87%.

Test plan

• pytest tests/moe/test_cute_dsl_fused_moe.py::TestInputsHelperContract — new unit test pins CuteDslMoEInputsHelper.inputs_pre_hook's input-layout contract (replaces inputs[2], passes through rest, deterministic across instances).
• pytest tests/moe/test_cute_dsl_fused_moe.py — existing CuteDSL MoE tests pass (no behavior change on the fast path).
• 10-run 30-cell parity sweep on B200 rc14 — all 450 cells parity_ok=True; numerical equivalence preserved.

Summary by CodeRabbit

• New Features

  • Optional input-preprocessing callback for autotuning profiles, preserved across derived configs
  • Deterministic MoE profiling inputs via a new helper and seeded initialization for reproducible tuning
  • Autotune runner enables a cold L2 cache option to improve MoE tuning realism

• Tests

  • Added CPU tests validating the input-preprocessing hook replaces the expected tensor, preserves others, and is deterministic

[coderabbitai]

Note

Reviews paused

It looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the reviews.auto_review.auto_pause_after_reviewed_commits setting.

Use the following commands to manage reviews:

• @coderabbitai resume to resume automatic reviews.
• @coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

• ▶️ Resume reviews
• 🔍 Trigger review

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: cc9f7496-03f1-4b87-9980-f0d10737c30c

📥 Commits

Reviewing files that changed from the base of the PR and between 5b000bf and 640e32e.

📒 Files selected for processing (1)

• flashinfer/fused_moe/cute_dsl/tuner.py

🚧 Files skipped from review as they are similar to previous changes (1)

• flashinfer/fused_moe/cute_dsl/tuner.py

---

📝 Walkthrough

Walkthrough

Adds an optional inputs_pre_hook to TuningConfig, implements CuteDslMoEInputsHelper to deterministically replace token_selected_experts for MoE autotuning, wires the helper and seeded input generation into the CuteDSL runner, and adds CPU tests validating the hook's contract and determinism.

Changes

MoE Autotuner Inputs Determinism

Layer / File(s)	Summary
Autotuner Hook Contract flashinfer/autotuner.py	TuningConfig gains optional inputs_pre_hook; _apply_tuning_overrides preserves it; AutoTuner.choose_one invokes the hook on synthesized tensors before per-tactic profiling.
CuteDslMoEInputsHelper Implementation flashinfer/fused_moe/cute_dsl/_inputs_helper.py	New CuteDslMoEInputsHelper computes balanced per-local-expert token counts, builds a deterministic (num_tokens, top_k) token_selected_experts via seeded per-expert randperm, and exposes inputs_pre_hook that replaces only that tensor and forwards others unchanged.
Runner Integration with Seeded Initialization flashinfer/fused_moe/cute_dsl/tuner.py	Runner imports and instantiates CuteDslMoEInputsHelper, seeds autotune tensor generators with torch.Generator.manual_seed(515), sets inputs_pre_hook=self._inputs_helper.inputs_pre_hook on TuningConfig, and enables use_cold_l2_cache=True.
Contract Validation Tests tests/moe/test_cute_dsl_fused_moe.py	Adds TestInputsHelperContract that builds synthetic inputs matching wrapper layout, asserts index-2 token_selected_experts is replaced with a fresh tensor preserving shape/dtype while other inputs pass through by identity, and verifies determinism across helper instances.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~25 minutes

Possibly related PRs

• flashinfer-ai/flashinfer#3025: Related edits to CuteDSL MoE autotuning inputs and token_selected_experts generation.
• flashinfer-ai/flashinfer#2958: Modifies _apply_tuning_overrides and related tuning-config override logic touched here.
• flashinfer-ai/flashinfer#3126: Changes to AutoTuner.choose_one profiling/input handling related to this hook integration.

Suggested reviewers

• aleozlx
• yzh119
• IwakuraRein
• samuellees
• jiahanc
• bkryu
• jimmyzho
• sricketts

🐰 I seed the tokens, one-two-three,
Reproducible hops for each expert tree,
A hook that swaps just one small slice,
Deterministic profiling — neat and precise,
Hop, tune, repeat — reproducible delight!

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 75.00% 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 'feat(cute_dsl/moe): deterministic balanced autotune profile inputs' accurately reflects the main changes, which add deterministic and balanced profiling inputs for CuteDSL MoE autotuning.
Description check	✅ Passed	The pull request description comprehensively covers the changes, objectives, empirical validation, trade-offs, and test plan, following the repository template structure.
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.

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

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

Tip

💬 Introducing Slack Agent: The best way for teams to turn conversations into code.

Slack Agent is built on CodeRabbit's deep understanding of your code, so your team can collaborate across the entire SDLC without losing context.

• Generate code and open pull requests
• Plan features and break down work
• Investigate incidents and troubleshoot customer tickets together
• Automate recurring tasks and respond to alerts with triggers
• Summarize progress and report instantly

Built for teams:

• Shared memory across your entire org—no repeating context
• Per-thread sandboxes to safely plan and execute work
• Governance built-in—scoped access, auditability, and budget controls

One agent for your entire SDLC. Right inside Slack.

👉 Get started

---

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 and usage tips.}

[gemini-code-assist]

Code Review

This pull request introduces a mechanism for deterministic and realistic input distributions during MoE autotuning by adding an "inputs_pre_hook" to the "TuningConfig". A new "CuteDslMoEInputsHelper" class implements balanced expert assignment via rejection sampling, replacing the previous random initialization. The "CuteDslMoETuner" is updated to utilize this helper and seeded tensor initializers. Review feedback identifies three issues: the "inputs_pre_hook" is not propagated during tuning overrides, a logic error in the expert assignment loop allows assignments when token counts are zero, and a multi-GPU indexing bug exists in the prioritization logic.

[gemini-code-assist]

If num_tokens_per_expert[j] is 0, the current logic will still assign expert j to one token because the inner loop over selection_order_j is entered and the break condition num_tokens_j <= 0 is only checked after the first assignment. Adding a check at the beginning of the expert loop prevents this and avoids unnecessary computation.

Suggested change

	for j, num_tokens_j in enumerate(num_tokens_per_expert):
	selection_order_j = selection_orders[j].tolist()
	for j, num_tokens_j in enumerate(num_tokens_per_expert):
	if num_tokens_j <= 0:
	continue
	selection_order_j = selection_orders[j].tolist()

[leejnau]

This is actually done deliberately to match the TRT-LLM behavior:

https://github.com/NVIDIA/TensorRT-LLM/blob/main/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py#L152-L160

[gemini-code-assist]

The limit calculation uses the local expert index j, but it should use the global expert index j + self.local_expert_offset to correctly identify tokens at risk of not hitting top_k experts across the entire MoE layer. Without this, the prioritization logic is effectively disabled for all ranks except the first one (where offset=0), breaking parity with the intended algorithm on multi-GPU setups.

Suggested change

	limit = self.top_k - (self.num_experts - j)
	limit = self.top_k - (self.num_experts - (j + self.local_expert_offset))

[leejnau]

This is actually done deliberately to match the TRT-LLM behavior:

https://github.com/NVIDIA/TensorRT-LLM/blob/main/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py#L146-L147

[qiching]

assume when num_tokens_on_curr_rank < num_local_experts, the base of divmod(num_tokens_on_curr_rank, num_local_experts) is 0, and the remaining experts should be assigned 0 tokens. For DSv3 with EP=16 and num_tokens=1: average=0.5, extra≈0.83, num_tokens_on_curr_rank=2, divmod(2,16)=(0,2) so experts 0 and 1 should each receive 1 token, while experts 2–15 should each receive 0 tokens; however, the current logic causes experts 2–15 to incorrectly receive 1 additional token each.

it deviates.

so how do you think?

[leejnau]

This is the same issue Gemini flagged earlier in this file, and your example reinforces this. However, TRT-LLM has identical behavior at https://github.com/NVIDIA/TensorRT-LLM/blob/main/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py#L151-L159

Since the goal is a faithful port, ensuring FlashInfer's autotune sees the same profile inputs TRT-LLM does, fixing here would deviate. We could instead fix it here and file an issue to TRT-LLM.

[qiching]

i agree with Gemini and i think here we should use global index:

global_j = j + self.local_expert_offset
limit = self.top_k - (self.num_experts - global_j)

[leejnau]

This was deliberately done for parity with TRT-LLM. TRT-LLM uses local j identically at https://github.com/NVIDIA/TensorRT-LLM/blob/main/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py#L146-L147, and the goal is port alignment. Switching to global j would deviate from that. We could instead fix it here and file an issue to TRT-LLM. (This would also apply to Gemini's other comment.)

[qiching]

Agree to merge. how about add a comment with the issue link to these two bug reports in the code to avoid repeated flagging?

[leejnau]

Agree to merge. how about add a comment with the issue link to these two bug reports in the code to avoid repeated flagging?

• Filed both as a single upstream tracking issue: [Bug]: [CuteDSL MoE] generate_token_selected_experts: ghost-token assignment to zero-allocation experts + dead-on-multi-rank prioritization NVIDIA/TensorRT-LLM#14146
• Added inline comments at both sites pointing to the TRT-LLM issue (commit 04b28b3) so future readers don't have to re-flag.

[qiching]

LGTM now! thanks your work

[nv-yunzheq]

/bot run

[flashinfer-bot]

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

[nv-yunzheq]

Do you want to use different seed value for tensor initialization?

[leejnau]

The shared seed is ok because the four lambdas draw from different shapes/dtypes/distributions, so the outputs are uncorrelated in practice despite sharing a seed.

[nv-yunzheq]

/bot run

[flashinfer-bot]

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

[nv-yunzheq]

/bot run

[flashinfer-bot]

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

[nv-yunzheq]

Approved as the unit test passed