[https://nvbugs/6093911][fix] Fix disagg gen-only benchmark hang under ADP router imbalance

[chienchunhung]

Summary by CodeRabbit

• Bug Fixes

  • Improved benchmark disaggregation fill gate to determine readiness based on executor state rather than request counts, ensuring more accurate synchronization across distributed ranks.
  • Enhanced diagnostic logging to report transfer status and request state information.

• Tests

  • Updated and expanded test coverage for state-based readiness validation, including new regression tests for distributed training scenarios.

Summary

Fixes four ADP/NIXL regressions PR #12208 introduced when it removed the blocking fill loop, and corrects a long-latent KV oversubscription in the Kimi-2 8k1k perf-sanity test config that #12208's fill-phase changes exposed.

Production fixes (2 commits)

PR #12208 replaced the blocking fill loop with a non-blocking gate to fix an MPI hard deadlock. The replacement was correct for MPI but introduced four ADP/NIXL-side regressions, all addressed here:

• State-based fill gate — count-based predicate is unsatisfiable under ADP per-rank rebalancing. Replaced with (num_fetched ≥ queue_size) ∧ (no active req in INIT/TRANS) ∧ (no transceiver pending). Fires uniformly across ranks.
• ADP router hard cap at max_num_active_requests (= per-rank max_batch_size) on both DefaultADPRouter and KVCacheAwareADPRouter — composes with upstream's fair_share_multiplier loose cap. Fixes nvbug 6071070; without it, bulk admission can route past max_batch_size and leave excess requests stuck in INIT.
• Fail-fast suppression during fill — PR [None][fix] return an explicit error if the requests can't be schedul… #12206's fail-fast was tripping on transient bulk-fetch state (some requests in INIT while others transfer). Resumes after the gate opens.
• tp_size-per-iter admission throttle in _pop_from_waiting_queue — restores the pre-[TRTLLM-11492][fix] Fix benchmark disagg deadlock by eliminating blocking fill loop #12208 conservative admission rate without reintroducing the blocking pattern. Scoped to non-warmup benchmark fill.

Test config fix (1 commit): Kimi-2 GEN KV fraction 0.6 → 0.78

Memory math (B200/GB200, 184 GiB HBM)

	0.6 (upstream)	0.78 (this PR)
KV pool	61.4 GiB	79.8 GiB
KV tokens / rank	1,875,008	2,437,504
Required at bs=256 × max_seq_len=9256	2,369,536	(same)
Sequences fit / 256	203 ❌	263 ✅
Headroom for transients	~40 GiB	~22 GiB

At 0.6, only 203 of 256 max-batch sequences fit per rank — 53 are perpetually paused, costing ~21 % of decode throughput across the entire run. 0.78 is the floor that fits the workload while preserving the ~22 GiB needed for NIXL receive buffers, NCCL workspace, EPLB layer churn, and CUDA caching-allocator fragmentation.

Why the test was passing before PR #12208

Same YAML, same oversubscription — but timing differed:

• Pre-[TRTLLM-11492][fix] Fix benchmark disagg deadlock by eliminating blocking fill loop #12208: count-based fetch loop completed in ~26 s; decode then ran for ~85–88 min under the chronic preemption; total wall landed just inside the test's 5400 s DEFAULT_TIMEOUT with 2–5 min of cluster-variance-dependent margin.
• Post-[TRTLLM-11492][fix] Fix benchmark disagg deadlock by eliminating blocking fill loop #12208 (and this PR): state-based fill flag remains active for ~84 min (tracks CTX prefill completion, which is the dominant cost when concurrency = capacity). The flag adds only ~50 s of all_gather overhead, but it ate the cluster-variance margin that had been hiding the oversubscription.

The test was always running on luck. Bumping to 0.78 makes it do what it claims (bs=256 steady-state) instead of depending on the right side of a 2–5 min margin.

Why this YAML change is bundled into this PR

The bump is a one-line test-config correctness fix that stands on its own — it would have been correct at any point in the test's lifetime. It belongs here because (a) the gate fixes are necessary but not sufficient for the Kimi-2 CI to reach green, and (b) the two changes share the same failure narrative.

Test coverage

Production fixes are covered by tests/unittest/_torch/executor/test_benchmark_disagg.py (10 test classes, 28 tests, no GPU required):

Production fix	Coverage
State-based fill gate	TestFillCompleteStateBased (4 tests, parameterized over the predicate's truth table), TestFillCompleteADP (3 tests, cross-rank allgather), TestFillCompleteADPRouterImbalance (2 tests, per-rank rebalancing edge cases), TestCheckBenchmarkDisaggGate (3 tests, wait-loop behavior), TestCanForwardGating (parameterized initial state)
ADP router hard cap	TestADPRouterPerRankCap (3 tests, parameterized over both DefaultADPRouter and KVCacheAwareADPRouter: cap applied at max, no rank exceeds, no-op below max)
Fail-fast suppression during fill	TestFailFastSuppressedDuringFill (3 tests: suppressed during fill, resumes after, full (fill_active × is_warmup) matrix)
tp_size-per-iter admission throttle	TestBenchmarkFillAdmissionFlowControl (2 tests, parameterized over tp_size ∈ {1, 4, 32}: cap applied during fill, full capacity used after)
Adjacent invariants	TestPadAttentionDpDummyBenchmarkDisagg (4 tests, dummy-request skip during fill), TestPrepareAndScheduleBatchNoBlock (1 test, fetch called once per iter)
End-to-end lifecycle	TestFillPhaseEndToEnd::test_full_lifecycle (admission-cap → KV-transfer-complete → gate-opens → normal admission, exercising both production code paths together)

The Kimi-2 perf-sanity test has been verified separately: https://nvbugspro.nvidia.com/bug/6093911?commentNumber=9

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)

• 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.

[chienchunhung]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #45018 [ run ] triggered by Bot. Commit: 16533e3 Link to invocation

[tensorrt-cicd]

PR_Github #45018 [ run ] completed with state SUCCESS. Commit: 16533e3
/LLM/main/L0_MergeRequest_PR pipeline #35331 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

Link to invocation

[chienchunhung]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #45236 [ run ] triggered by Bot. Commit: 16533e3 Link to invocation

[tensorrt-cicd]

PR_Github #45236 [ run ] completed with state SUCCESS. Commit: 16533e3
/LLM/main/L0_MergeRequest_PR pipeline #35497 completed with status: 'SUCCESS'

CI Report

Link to invocation

[coderabbitai]

📝 Walkthrough

Walkthrough

Refactors the benchmark-disagg fill gate readiness logic from counting generation requests to evaluating executor state, including active request states and KV cache transfer completion. In ADP mode, synchronizes readiness across TP ranks using allgather. Updates corresponding test suite to validate state-based predicate behavior and adds regression coverage for rank skew scenarios.

Changes

Cohort / File(s)	Summary
Core Predicate Logic tensorrt_llm/_torch/pyexecutor/py_executor.py	Refactored _is_benchmark_disagg_fill_complete from count-based thresholds to state-based evaluation. Now checks: (1) sufficient fetched requests, (2) active requests past INIT/TRANS_IN_PROGRESS states, (3) KV cache transceiver completion status. In ADP mode, synchronizes conditions (B) and (C) across TP ranks via tp_allgather. Updated logging to report rank-local state counts and transfer status.
Test Suite Updates tests/unittest/_torch/executor/test_benchmark_disagg.py	Comprehensively refactored tests to validate state-based predicate instead of count-based logic. Removed tests depending on old batch construction semantics. Updated ADP test coverage to verify tp_allgather behavior. Simplified dummy padding tests. Added regression suite for nvbug 6071070 validating predicate blocking under ADP router skew and INIT overflow scenarios.

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~45 minutes

Possibly related PRs

• [TRTLLM-11492][fix] Fix benchmark disagg deadlock by eliminating blocking fill loop #12208: Directly related—both PRs modify benchmark-disagg gating in _is_benchmark_disagg_fill_complete and corresponding test suite, introducing state-based gate evaluation and cross-rank synchronization in ADP mode.

Suggested reviewers

• pcastonguay

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 21.74% 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 correctly identifies the bug fix (disagg gen-only benchmark hang) and the root cause context (ADP router imbalance), directly matching the PR's primary objective.
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	PR description comprehensively explains the problem, solution, test coverage, and production fixes with detailed context on regressions and memory math.

_{✏️ 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

---

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

[coderabbitai]

🧹 Nitpick comments (1)

tests/unittest/_torch/executor/test_benchmark_disagg.py (1)

233-233: Optional: Linter warning about multiplication sign character.

Static analysis flags the × character as ambiguous (RUF002). While the mathematical notation is clearer for readers, you can silence this by using ASCII x or escaping if strict linting is required.

-    (e.g. 31×256 + 1×255 = 8191 < 8192). The new state-based predicate
+    (e.g. 31*256 + 1*255 = 8191 < 8192). The new state-based predicate

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@tests/unittest/_torch/executor/test_benchmark_disagg.py` at line 233, Replace
the non-ASCII multiplication character in the test string (e.g. "31×256 + 1×255
= 8191 < 8192") with an ASCII 'x' or escape it to satisfy the linter; update the
literal in tests/unittest/_torch/executor/test_benchmark_disagg.py so
occurrences of '×' become 'x' (or an escaped sequence) to silence RUF002 while
preserving the intended text.

🤖 Prompt for all review comments with AI agents

Verify each finding against the current code and only fix it if needed.

Nitpick comments:
In `@tests/unittest/_torch/executor/test_benchmark_disagg.py`:
- Line 233: Replace the non-ASCII multiplication character in the test string
(e.g. "31×256 + 1×255 = 8191 < 8192") with an ASCII 'x' or escape it to satisfy
the linter; update the literal in
tests/unittest/_torch/executor/test_benchmark_disagg.py so occurrences of '×'
become 'x' (or an escaped sequence) to silence RUF002 while preserving the
intended text.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: c7e428f7-465d-4269-9384-c8905b912cbc

📥 Commits

Reviewing files that changed from the base of the PR and between a492e5b and f98e2fd.

📒 Files selected for processing (2)

• tensorrt_llm/_torch/pyexecutor/py_executor.py
• tests/unittest/_torch/executor/test_benchmark_disagg.py

[chienchunhung]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #46198 [ run ] triggered by Bot. Commit: a2c6107 Link to invocation

[tensorrt-cicd]

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

[pcastonguay]

Some tests seem to fail due to your changes. For example:

[2026-04-30T01:31:33.474Z] FAILED unittest/DGX_B200-PyTorch-2/unittest/_torch/executor/test_benchmark_disagg.py::TestADPRouterPerRankCap::test_cap_not_applied_when_below_max[KVCacheAwareADPRouter] - assert 4 == 2

Also the release-check job is failing. Could you take a look?

[chienchunhung]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #46460 [ run ] triggered by Bot. Commit: d408754 Link to invocation

[chienchunhung]

Some tests seem to fail due to your changes. For example:

[2026-04-30T01:31:33.474Z] FAILED unittest/DGX_B200-PyTorch-2/unittest/_torch/executor/test_benchmark_disagg.py::TestADPRouterPerRankCap::test_cap_not_applied_when_below_max[KVCacheAwareADPRouter] - assert 4 == 2

Some of the new tests introduced by this PR failed after rebasing to the latest HEAD; updated the tests.

Also the release-check job is failing. Could you take a look?

Yup over the past few days I have been working closely with QA on the Kimi2 test.

[tensorrt-cicd]

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

[chienchunhung]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #48643 [ run ] triggered by Bot. Commit: 0562b7c Link to invocation

[chienchunhung]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #48652 [ run ] triggered by Bot. Commit: 3ce8cc3 Link to invocation

[tensorrt-cicd]

PR_Github #48652 [ run ] completed with state SUCCESS. Commit: 3ce8cc3
/LLM/main/L0_MergeRequest_PR pipeline #38432 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

[chienchunhung]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #48965 [ run ] triggered by Bot. Commit: 3ce8cc3 Link to invocation

[tensorrt-cicd]

PR_Github #48965 [ run ] completed with state FAILURE. Commit: 3ce8cc3
/LLM/main/L0_MergeRequest_PR pipeline #38710 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

[chienchunhung]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #48977 [ run ] triggered by Bot. Commit: 3ce8cc3 Link to invocation

[tensorrt-cicd]

PR_Github #48977 [ run ] completed with state FAILURE. Commit: 3ce8cc3
/LLM/main/L0_MergeRequest_PR pipeline #38722 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

[chienchunhung]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #49233 [ run ] triggered by Bot. Commit: 1e81601 Link to invocation

[tensorrt-cicd]

PR_Github #49233 [ run ] completed with state ABORTED. Commit: 1e81601

Link to invocation

[chienchunhung]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #49463 [ run ] triggered by Bot. Commit: 1e81601 Link to invocation

[chienchunhung]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #49486 [ run ] triggered by Bot. Commit: 1e81601 Link to invocation

[tensorrt-cicd]

PR_Github #49463 [ run ] completed with state ABORTED. Commit: 1e81601

Link to invocation

[tensorrt-cicd]

PR_Github #49486 [ run ] completed with state SUCCESS. Commit: 1e81601
/LLM/main/L0_MergeRequest_PR pipeline #39126 completed with status: 'SUCCESS'

CI Report

Link to invocation