[https://nvbugs/6403920][fix] Move allreduce strategy test to smoke, use full-warp hidden_size by MrGeva · Pull Request #16389 · NVIDIA/TensorRT-LLM

MrGeva · 2026-07-14T20:40:43Z

Description

Move the AutoDeploy allreduce-strategy test from tests/unittest/auto_deploy/multigpu/custom_ops/test_ad_dist_strategies.py to tests/unittest/auto_deploy/multigpu/smoke/test_ad_allreduce_strategies.py (pure relocation, no logic change).
Override the test's tiny model hidden_size from 64 to 256, so the fused all-reduce/RMSNorm kernel's blockReduceSum always launches exactly one full warp (32 threads) instead of the partial-warp (8-thread) block.

No kernel code (allReduceFusionKernels.cu) is touched at all — the kernel is exactly as it is on main today.

Motivation

Isolates the smoke-test-in-new-location + full-warp-hidden_size combination from the kernel fix in #16111, to determine whether CI failures seen on #16111 (an unrelated cuBLAS GEMM error and, separately, an MPI/PMIx runtime stall, on three different physical nodes across attempts) reproduce even without any kernel changes.

Test Coverage

Same test_allreduce_strategies test, all 6 strategies, hidden_size=256, run via the smoke directory.

🤖 Generated with Claude Code

Summary by CodeRabbit

Tests
- Added multi-GPU smoke coverage for multiple all-reduce strategies.
- Added throughput benchmark validation across supported distributed execution strategies.
- Improved test reliability with dataset reuse, kernel prewarming, GPU availability checks, and execution timeouts.
- Streamlined graph-level tests to focus on verifying all-reduce and all-gather strategy propagation.

Signed-off-by: Eran Geva <19514940+MrGeva@users.noreply.github.com>

Diagnostic-only change: overrides the tiny test model's hidden_size from 64 (which launches an 8-thread, partial-warp block in the fused all-reduce RMSNorm kernel) to 256 (exactly one full warp), so the smoke test no longer exercises the partial-warp reduction path at all. Two prior CI runs of the base branch (agent/fix-ad-allreduce-ci) failed with an identical CUBLAS_STATUS_EXECUTION_FAILED error in an unrelated cuBLAS GEMM (torch_linear_simple), on the same node (nsc-svg-slurm-1-gpu-120), at a different allreduce strategy each time. This branch isolates whether that failure is specific to the partial-warp kernel path or is unrelated node/infra flakiness that also reproduces on the full-warp path. Signed-off-by: Eran Geva <19514940+MrGeva@users.noreply.github.com>

MrGeva · 2026-07-14T20:40:52Z

/bot run

coderabbitai · 2026-07-14T20:46:30Z

📝 Walkthrough

Walkthrough

The change separates multi-GPU throughput strategy testing into a new smoke module, retaining graph-level propagation tests in the custom-ops module. The smoke tests prepare shared data, prewarm FlashInfer kernels, apply strategy-specific options, and invoke the throughput benchmark.

Changes

All-reduce strategy test coverage

Layer / File(s)	Summary
Retain graph strategy checks `tests/unittest/auto_deploy/multigpu/custom_ops/test_ad_dist_strategies.py`	Removes benchmark and dataset utilities while retaining graph-level all-reduce and all-gather strategy propagation checks.
Build smoke test infrastructure `tests/unittest/auto_deploy/multigpu/smoke/test_ad_allreduce_strategies.py`	Adds timeout handling, FlashInfer JIT prewarming, synthetic dataset generation, and shared dataset reuse.
Run strategy benchmark matrix `tests/unittest/auto_deploy/multigpu/smoke/test_ad_allreduce_strategies.py`	Adds parameterized multi-GPU throughput runs for six all-reduce strategies with auto-deploy configuration and exit-status assertions.

Estimated code review effort: 4 (Complex) | ~45 minutes

Suggested reviewers: suyoggupta

Sequence Diagram(s)

sequenceDiagram
  participant Pytest
  participant DatasetPreparer
  participant CliRunner
  participant AutoDeploy
  participant ThroughputBenchmark
  Pytest->>DatasetPreparer: generate shared synthetic dataset
  DatasetPreparer-->>Pytest: return dataset contents
  Pytest->>CliRunner: invoke throughput for each strategy
  CliRunner->>AutoDeploy: apply sharding and compilation options
  AutoDeploy->>ThroughputBenchmark: run multi-GPU benchmark
  ThroughputBenchmark-->>Pytest: return exit status

🚥 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: relocating the allreduce strategy test and adjusting hidden_size for smoke testing.
Description check	✅ Passed	The description covers the change, motivation, and test coverage, but does not fill the PR Checklist section.

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests

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

coderabbitai

Actionable comments posted: 4

🤖 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 `@tests/unittest/auto_deploy/multigpu/smoke/test_ad_allreduce_strategies.py`:
- Around line 40-52: Add Python 3.10+ type annotations to every function and
pytest fixture in this test module, including timeout and timeout_handler, all
parameters, and return types such as Iterator[None], Path, str, and None where
applicable. Ensure annotations cover the additional functions and fixtures
referenced by the review without changing their behavior.
- Around line 34-37: Remove the custom TimeoutError declaration and use Python’s
built-in TimeoutError throughout the test module, preserving the existing
timeout handling behavior.
- Line 194: Move the GPU-count guard for test_allreduce_strategies ahead of
shared_dataset setup and the autouse prewarm_flashinfer_jit fixture, using an
earlier collection or fixture-level skip so CPU and single-GPU runners exit
before dataset preparation or FlashInfer warmup.
- Around line 76-125: Update the FlashInfer prewarm block around the lock_f
acquisition and warmup steps to close and unlock the file in a finally block
whenever fcntl.flock succeeds. Ensure cleanup runs for exceptions from page or
sampling warmup, torch.cuda.empty_cache(), and future steps, while preserving
the existing warmup behavior.

🪄 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: fa3d6d0f-b8ed-44d1-bdb3-4c24d1938a1d

📥 Commits

Reviewing files that changed from the base of the PR and between 725d9a1 and 8991f60.

📒 Files selected for processing (2)

tests/unittest/auto_deploy/multigpu/custom_ops/test_ad_dist_strategies.py
tests/unittest/auto_deploy/multigpu/smoke/test_ad_allreduce_strategies.py

💤 Files with no reviewable changes (1)

tests/unittest/auto_deploy/multigpu/custom_ops/test_ad_dist_strategies.py

tensorrt-cicd · 2026-07-14T20:54:19Z

PR_Github #59275 [ run ] triggered by Bot. Commit: 8991f60 Link to invocation

tensorrt-cicd · 2026-07-14T21:08:32Z

PR_Github #59275 [ run ] completed with state FAILURE. Commit: 8991f60
/LLM/main/L0_MergeRequest_PR pipeline #47764 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

…path test_allreduce_strategies moved from tests/unittest/auto_deploy/multigpu/custom_ops/test_ad_dist_strategies.py to tests/unittest/auto_deploy/multigpu/smoke/test_ad_allreduce_strategies.py in the prior commit, but 6 waives.txt entries still referenced the old path, failing the release-check's stale-waive validator. Signed-off-by: Eran Geva <19514940+MrGeva@users.noreply.github.com>

MrGeva · 2026-07-14T21:12:28Z

/bot run

tensorrt-cicd · 2026-07-14T21:19:08Z

PR_Github #59281 [ run ] triggered by Bot. Commit: 6e0742c Link to invocation

tensorrt-cicd · 2026-07-14T22:28:21Z

PR_Github #59281 [ run ] completed with state SUCCESS. Commit: 6e0742c
/LLM/main/L0_MergeRequest_PR pipeline #47768 completed with status: 'SUCCESS'

CI Report

Link to invocation

MrGeva · 2026-07-15T06:29:50Z

/bot run --stage-list "DGX_B200-4_GPUs-AutoDeploy-1" --disable-reuse-test

tensorrt-cicd · 2026-07-15T06:36:11Z

PR_Github #59378 [ run ] triggered by Bot. Commit: 6e0742c Link to invocation

ZhanruiSunCh

LGTM for infra part.

tensorrt-cicd · 2026-07-15T13:59:08Z

PR_Github #59448 [ run ] completed with state SUCCESS. Commit: 0d5ff18
/LLM/main/L0_MergeRequest_PR pipeline #47916 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

MrGeva · 2026-07-16T05:20:01Z

/bot run

MrGeva · 2026-07-16T06:59:53Z

/bot run

tensorrt-cicd · 2026-07-16T07:06:17Z

PR_Github #59650 [ run ] triggered by Bot. Commit: 3806efd Link to invocation

tensorrt-cicd · 2026-07-16T07:30:37Z

PR_Github #59650 [ run ] completed with state FAILURE. Commit: 3806efd
/LLM/main/L0_MergeRequest_PR pipeline #48085 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

MrGeva · 2026-07-16T08:11:04Z

/bot run

tensorrt-cicd · 2026-07-16T08:17:09Z

PR_Github #59672 [ run ] triggered by Bot. Commit: 3806efd Link to invocation

tensorrt-cicd · 2026-07-16T08:37:09Z

PR_Github #59672 [ run ] completed with state FAILURE. Commit: 3806efd
/LLM/main/L0_MergeRequest_PR pipeline #48105 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

MrGeva · 2026-07-16T09:38:52Z

/bot run

tensorrt-cicd · 2026-07-16T09:44:29Z

PR_Github #59688 [ run ] triggered by Bot. Commit: 061e69d Link to invocation

tensorrt-cicd · 2026-07-16T10:10:36Z

PR_Github #59688 [ run ] completed with state FAILURE. Commit: 061e69d
/LLM/main/L0_MergeRequest_PR pipeline #48123 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

MrGeva · 2026-07-16T11:25:53Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-16T11:32:06Z

PR_Github #59697 [ run ] triggered by Bot. Commit: 061e69d Link to invocation

BowenFu

LGTM — faithful relocation of the allreduce strategy test to the smoke dir; no test logic dropped, and the smoke dir is already registered so the test still runs.

autoweave-bot · 2026-07-16T11:45:40Z

This looks like a great optimization pass, nice work! 🚀😌
Explore here →

_{Maintainer? Turn off weaves from non-maintainers →}
_{Carefully crafted by Subweave · 🧶 used ~16k LLM tokens}

tensorrt-cicd · 2026-07-16T14:26:23Z

PR_Github #59697 [ run ] completed with state SUCCESS. Commit: 061e69d
/LLM/main/L0_MergeRequest_PR pipeline #48128 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

MrGeva · 2026-07-16T14:54:47Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-16T15:01:00Z

PR_Github #59731 [ run ] triggered by Bot. Commit: 72ac8bc Link to invocation

tensorrt-cicd · 2026-07-16T16:28:59Z

PR_Github #59731 [ run ] completed with state SUCCESS. Commit: 72ac8bc
/LLM/main/L0_MergeRequest_PR pipeline #48158 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

MrGeva · 2026-07-16T17:14:38Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-16T17:20:13Z

PR_Github #59753 [ run ] triggered by Bot. Commit: 72ac8bc Link to invocation

MrGeva added 2 commits July 14, 2026 13:35

moved test to smoke

2280cbf

Signed-off-by: Eran Geva <19514940+MrGeva@users.noreply.github.com>

MrGeva requested a review from a team as a code owner July 14, 2026 20:40

MrGeva requested a review from suyoggupta July 14, 2026 20:40

github-actions Bot assigned MrGeva Jul 14, 2026

coderabbitai Bot reviewed Jul 14, 2026

View reviewed changes

MrGeva requested review from a team as code owners July 14, 2026 21:12

MrGeva requested review from BowenFu, EmmaQiaoCh, ZhanruiSunCh, brnguyen2, fredricz-20070104 and yufeiwu-nv July 14, 2026 21:12

MrGeva changed the title ~~[None][test] Move allreduce strategy test to smoke, use full-warp hidden_size~~ [https://nvbugs/6403920][fix] Move allreduce strategy test to smoke, use full-warp hidden_size Jul 15, 2026

tcherckez-nvidia approved these changes Jul 15, 2026

View reviewed changes

chzblych approved these changes Jul 15, 2026

View reviewed changes

ZhanruiSunCh approved these changes Jul 15, 2026

View reviewed changes

Merge branch 'main' into agent/allreduce-smoke-move-fullwarp

ff4d550

yufeiwu-nv removed their request for review July 16, 2026 05:47

Merge branch 'main' into agent/allreduce-smoke-move-fullwarp

3806efd

Merge branch 'main' into agent/allreduce-smoke-move-fullwarp

061e69d

BowenFu approved these changes Jul 16, 2026

View reviewed changes

MrGeva closed this Jul 16, 2026

auto-merge was automatically disabled July 16, 2026 14:54
Pull request was closed

MrGeva reopened this Jul 16, 2026

Merge branch 'main' into agent/allreduce-smoke-move-fullwarp

72ac8bc

MrGeva enabled auto-merge (squash) July 16, 2026 14:54

brnguyen2 approved these changes Jul 16, 2026

View reviewed changes

Uh oh!

Conversation

MrGeva commented Jul 14, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Motivation

Test Coverage

Summary by CodeRabbit

Uh oh!

MrGeva commented Jul 14, 2026

Uh oh!

coderabbitai Bot commented Jul 14, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Walkthrough

Changes

Sequence Diagram(s)

Uh oh!

coderabbitai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

tensorrt-cicd commented Jul 14, 2026

Uh oh!

tensorrt-cicd commented Jul 14, 2026

Uh oh!

MrGeva commented Jul 14, 2026

Uh oh!

tensorrt-cicd commented Jul 14, 2026

Uh oh!

tensorrt-cicd commented Jul 14, 2026

Uh oh!

MrGeva commented Jul 15, 2026

Uh oh!

tensorrt-cicd commented Jul 15, 2026

Uh oh!

ZhanruiSunCh left a comment

Choose a reason for hiding this comment

Uh oh!

tensorrt-cicd commented Jul 15, 2026

Uh oh!

MrGeva commented Jul 16, 2026

Uh oh!

MrGeva commented Jul 16, 2026

Uh oh!

tensorrt-cicd commented Jul 16, 2026

Uh oh!

tensorrt-cicd commented Jul 16, 2026

Uh oh!

MrGeva commented Jul 16, 2026

Uh oh!

tensorrt-cicd commented Jul 16, 2026

Uh oh!

tensorrt-cicd commented Jul 16, 2026

Uh oh!

MrGeva commented Jul 16, 2026

Uh oh!

tensorrt-cicd commented Jul 16, 2026

Uh oh!

tensorrt-cicd commented Jul 16, 2026

Uh oh!

MrGeva commented Jul 16, 2026

Uh oh!

tensorrt-cicd commented Jul 16, 2026

Uh oh!

BowenFu left a comment

Choose a reason for hiding this comment

Uh oh!

autoweave-bot commented Jul 16, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

tensorrt-cicd commented Jul 16, 2026

Uh oh!

MrGeva commented Jul 16, 2026

Uh oh!

tensorrt-cicd commented Jul 16, 2026

Uh oh!

tensorrt-cicd commented Jul 16, 2026

MrGeva commented Jul 14, 2026 •

edited

Loading

coderabbitai Bot commented Jul 14, 2026 •

edited

Loading

autoweave-bot commented Jul 16, 2026 •

edited

Loading