[Barry-Delaney]

Dev Engineer Review

The Rubin locality-domain tuning configuration now constrains optional inplace output buffers using the permuted-token bucket. The output row count follows the bucket size. The output scale-factor count uses the bucket size, weight width, and scaling_vector_size. This should allow cache-key reuse across token counts in the same bucket.

The scale-factor inference assumes input 1 has the expected weight shape and that its second dimension represents the interleaved gate/linear width. Existing cache entries for this operation may not match the updated keys. The author reports that None outputs retain their existing keys and other runners are unchanged. Rubin hardware execution and performance results were not supplied. Review severity counts are unavailable.

QA Engineer Review

The test file adds coverage for scale-factor sizing, cache-key reuse across two token counts in the same bucket, absent outputs, and optimization-profile output dimensions. The cache-key and profile tests require the Rubin CuTe DSL and are skipped when it is unavailable. The supplied results do not report test execution.

The test file is not listed for these new tests in the inspected CI or manual-QA lists. The CI list includes a different test from the same file. The two locality-domain correctness tests are present in the test file, but their CI inclusion and execution were not established. Coverage verdict: needs follow-up.

Per-File QA Perspective

• tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py: Verify that the output scale-factor shape matches the shared locality-domain buffer for supported weight shapes and scaling-vector sizes. Verify cache-key reuse for token counts in the same bucket and compatibility with regenerated tuning-cache entries.
• tests/unittest/_torch/moe/kernels/parallel/test_cute_dsl_moe.py: Added helper sizing and Rubin tuning tests cover scale-factor count, cache keys, optional outputs, and profile shapes. The two new DSL-dependent tests are not listed in the inspected CI or manual-QA lists; the two locality-domain correctness tests also lack verified list coverage.

Description

Problem

On Rubin (SM107), the locality-domain NVFP4 MoE FC1 op
(trtllm::cute_dsl_nvfp4_gather_grouped_gemm_act_fusion_locality_domain_inplace_rubin)
tunes Sm107BlockScaledContiguousGatherGroupedGemmActFusionRunner with twelve inputs. Inputs 10 and 11
are the shared inplace output and output scale-factor buffers that both partitions write into.

The runner's TuningConfig buckets input 7 (permuted_idx_to_expanded_idx) and constrains inputs
0/2/5/6 to it. It says nothing about 10/11. AutoTuner._find_nearest_profile copies every dimension
that no spec mentions into the cache key verbatim, on the assumption that it is static. These two are
not static: the caller allocates fc1_out / fc1_out_sf from the actual permuted-token count
(TrtllmCutedslGroupedGemmNvfp4Impl._run_moe_nvfp4_locality_domain), so the raw count re-enters the
key and defeats the bucketing on input 7.

What goes wrong:

• A tuning sweep stores one entry per (bucket, output-buffer size) pair, not one per bucket.
• At inference, a lookup in the same bucket hits only when the buffer size happens to match one that
  was tuned. Otherwise the op falls back to the default tactic, and the tuned result goes unused.
• Duplicate entries for the same bucket are tuned independently, so they can disagree.

Fix

• Add ConstraintSpec(10, 0, infer_shape_max_num_permuted_tokens) and a ConstraintSpec for input
  11 driven by a new GatherGroupedGemmInputsHelper.infer_shape_locality_domain_output_sf_numel.
  Both dims become pure functions of input 7's bucket.

• The scale-factor element count is m * n // scaling_vector_size, with n = input_shapes[1][1]
  (the per-partition interleaved weight width). This matches:

  • the caller's allocation (m * full_interm // scaling_vector_size, where full_interm == n for
    two partitions);
  • the kernel-side assertion in forward (m * interm_size // scaling_vector_size * 2);
  • the per-row factor that _resize_locality_domain_outputs derives in the inputs pre-hook.

  The value has to be exact: the pre-hook infers the per-row scale count from the profiling tensor,
  so a wrong size would make every profiled tactic fail the kernel's shape assertion.

• scaling_vector_size is taken from the runner rather than hard-coded.

• Refresh the stale comment on IS_CUTLASS_DSL_RUBIN_AVAILABLE in cute_dsl_utils.py: the public nvidia-cutlass-dsl 4.8.0 packages already ship cutlass.utils.rubin_helpers, so the flag is set by the import probe and the new test runs wherever that package is installed. Comment-only change.

Scope and compatibility

• Single-tensor (non-locality-domain) path: passes output_tensor=None / output_sf_tensor=None.
  The autotuner records these as (0,), and both _optimization_profiles and _find_nearest_profile
  skip constraint specs on such inputs, so the cache keys on this path are unchanged.
• Other runners: already constrain their inplace outputs, so they are unchanged:
  • Sm100 runners;
  • the Sm107 NVFP4 finalize runner (output at index 5);
  • the Sm107 fused FC12 runner;
  • the BF16 Sm107 gather runner, whose inplace output at index 7 is constrained when present.
• Persisted caches: locality-domain FC1 cache keys change, because the output dims are now -1
  instead of the raw row count. Existing on-disk autotuner cache entries for this op no longer match.
  In inference mode the op falls back to the default tactic until the next tuning run regenerates
  them.
• Profiling buffers: profiling now allocates bucket-sized output buffers instead of reusing the
  caller's buffers. The inputs pre-hook resizes them to the generated routing as before.

Performance

Expected effect: the locality-domain FC1 op reuses one tuned tactic per bucket across token counts,
instead of falling back to the default tactic whenever the buffer size does not match a tuned one.
End-to-end numbers have not been measured for this PR.

Test Coverage

New test in tests/unittest/_torch/moe/kernels/parallel/test_cute_dsl_moe.py (skipped without the Rubin CuTe DSL):

• test_sm107_nvfp4_gather_act_fusion_cache_key_ignores_inplace_output_rows: two permuted-token counts in the same tuning bucket now produce the same cache key, and every generated profile sizes the scale-factor buffer (input 11) as m * n // 16, the size the kernel asserts.

Not yet run on Rubin hardware. The existing locality-domain correctness tests (test_nvfp4_gather_grouped_gemm_swiglu_locality_domain_rubin, test_moe_module_locality_domain_correctness_rubin) should be part of CI for this change.

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.

GitHub Bot Help

To see a list of available CI bot commands, please comment /bot help.

[coderabbitai]

No actionable comments were generated in the recent review.

Recent review info

Run configuration

• Source mode: Auto
• Repository: NVIDIA/TensorRT-LLM
• Branch: main
• File: .coderabbit.yaml
• Commit: 6d20096
• Layers: Repository YAML
• Review profile: CHILL
• Plan: Enterprise
• Run ID: 471bd087-3626-4914-a420-8070e66a3644

Commits

Reviewing files that changed from the base of the PR and between 6d20096 and 36a9291.

Files selected for processing: 2

• tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py
• tests/unittest/_torch/moe/kernels/parallel/test_cute_dsl_moe.py

Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 11 remain after this review.

Walkthrough

The gather grouped-GEMM tuning configuration now infers optional in-place output shapes from the permuted-token bucket. Rubin tests cover scale-factor sizing, autotuner cache keys, and optimization profile output shapes.

Changes

Rubin locality-domain output sizing

Layer / File(s)	Summary
Infer and constrain output shapes tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py	The helper infers the shared output scale-factor element count from the maximum permuted-token count, weight width, and scaling-vector size. The tuning configuration applies inferred row and scale-factor shapes to optional in-place outputs.
Validate tuning keys and profiles tests/unittest/_torch/moe/kernels/parallel/test_cute_dsl_moe.py	Rubin tests check inferred scale-factor sizes, cache keys for inputs in the same tuning bucket, absent-output key entries, and output shapes in generated optimization profiles.

Priority: Pending latest changes

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

Change: Bug fix

Suggested reviewers: juney-nvidia, leslie-fang25, yihuilu512

Merge Risk: Minimal · up to 36a92

Output shapes are grouped by token bucket. No merge-blocking defect was identified; Rubin hardware execution remains unverified.

Pre-merge checks: 4 · 1

Failed checks: 1 inconclusive

Check name	Status	Explanation	Resolution
Docstring Coverage		Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff with cyclomatic complexity of at least 2. Documentat…	Write docstrings for the functions missing them to satisfy the coverage threshold.

Passed checks: 4 passed

Check name	Status	Explanation
Title check		The title clearly identifies the fix and its main purpose: excluding Rubin inplace FC1 outputs from the CuteDSL MoE autotune key. It follows the required [None][fix] format.
Linked Issues check		Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check		Check skipped because no linked issues were found for this pull request.
Description check		The PR description follows the repository template. It includes a Description section with Problem, Fix, Scope and compatibility, and Performance subsections. It includes a Test Coverage section that …

Full details: Description check

Explanation

The PR description follows the repository template. It includes a Description section with Problem, Fix, Scope and compatibility, and Performance subsections. It includes a Test Coverage section that names the new test and says it has not been run on Rubin hardware. It includes the PR Checklist with the last item checked. The title line is not part of the provided body, so the title format cannot be verified here. The description explains what changes and why, and it names the affected paths and compatibility impact. The description does not state an NVBugs or JIRA ticket ID, and it does not state an API label decision. These are minor gaps and do not block the check.

Full details: Docstring Coverage

Explanation

Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff with cyclomatic complexity of at least 2. Documentation-only edits and deletions within surviving functions are included; unchanged diff context and fully deleted functions are excluded. Analyzed 3 functions across 1 files. (1 skipped: 1 too large.)

Resolution

Write docstrings for the functions missing them to satisfy the coverage threshold.

Finishing Touches: 💡 1

🛠️ Fix failing CI checks: 💡

• Commit to this branch
• Create a new PR

🧪 Generate unit tests (beta)

• Create a new PR

• Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

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