[TRTLLM-15178][fix] Pad an empty attention-DP scheduled batch so the fleet can make forward progress

[xwang233]

Description

Under attention DP, _can_queue vetoes the forward pass on every rank when any one rank's scheduled batch is empty:

tp_batch_sizes = self.dist.tp_allgather(scheduled_batch.batch_size)
can_queue = 0 not in tp_batch_sizes

_pad_attention_dp_dummy_request is meant to make an empty scheduled batch unreachable, but it guarantees a different invariant — every rank has at least one active request — and it necessarily runs before _schedule(), so the capacity scheduler's verdict cannot inform it. A rank whose only active request does not fit the free KV cache therefore counts as one active request, receives no padding dummy, and schedules an empty batch.

A rank-local cache shortage then stalls the whole fleet:

1. the empty rank vetoes can_queue everywhere, so peers never execute the context chunks they did schedule — _update_request_states, the only caller of move_to_next_context_chunk(), runs under if can_queue:;
2. no chunk completes, so no KV cache is released anywhere;
3. the starved rank stays starved, and the state re-arms every iteration.

Nothing on that path raises or times out, so the stall is silent: clients see zero errors and zero progress. On a disaggregated context server, where the starved rank is waiting on its own in-flight cache transfers to release the blocks it needs, the state can persist indefinitely.

Fix

PyExecutor._pad_empty_attention_dp_batch(), called immediately after _schedule() — the condition it repairs does not exist until the capacity scheduler has returned its verdict. A rank left with an empty batch appends a generation dummy to the scheduled batch.

Design points:

• Rank-local. _schedule() performs a tp_allgather inside _balance_adp_requests, so re-running scheduling on only the empty ranks would desynchronize that collective. The dummy is appended to scheduled_batch.generation_requests instead — pure rank-local state.
• A generation dummy, not the context dummy the pre-schedule path would pick: a rank that is empty precisely because it is short of KV cache must not be asked for max_num_tokens worth of it. A generation dummy costs 1 + max_total_draft_tokens tokens, and peers already run generation dummies alongside real context chunks.
• No new failure modes. Every allocation failure (OutOfPagesError, NoFreeSlotsError, insufficient capacity) degrades to today's behaviour rather than propagating: all ranks have yet to agree on can_queue, so a rank-local raise here would strand the peers in the collectives that follow.
• Bounded. No-op unless attention DP is enabled and the scheduled batch is empty; no-op when the rank has no active requests (the existing path covers that); no-op at max_num_active_requests, so it cannot trip the expected_num_active_requests assert in _pad_attention_dp_dummy_request; refuses to double-allocate the singleton ATTENTION_DP_DUMMY_REQUEST_ID; honors the benchmark-disagg fill gate, which suppresses forwards on purpose.
• KV cache manager V2. The dummy joins the batch after scheduling, so the V2 scheduler never grew its capacity; _revert_gen_alloc skips it, since reverting would shrink a real allocation.
• Scope. The pipeline-parallel loop does not call _prepare_and_schedule_batch and is unaffected.

Test Coverage

tests/unittest/_torch/executor/test_py_executor.py — nine new CPU-only cases covering the padding path: a rank with an unfittable active request gets a generation dummy in its scheduled batch; no-ops for non-empty batches, attention DP disabled, no active requests, an already-live dummy, and the active-request cap; graceful degradation on insufficient capacity and on either allocation error; and exclusion of the dummy from _revert_gen_alloc.

Validated on hardware, 4-node disaggregated attention-DP context servers:

• GB200 — the stall reproduced deterministically (3/3) on the unpatched build; with the fix the same workload runs to completion.
• GB300 — A/B against the same base commit and container, the fix compiled into the engine as the only variable: before → stalls (fleet frozen at 166 requests fetched), after → completes (2,526 requests). The empty-batch condition still fired 15 times on the fixed engine and was survived every time — the trigger is not avoided, it is made recoverable.

PR Checklist

• PR title uses the format [JIRA/NVBUG/None][type] Summary
• Description explains the change
• Test coverage added
• Documentation not applicable (internal scheduling fix, no user-facing surface)

Dev Engineer Review

• PyExecutor now pads rank-local empty scheduled batches with a generation dummy under attention data parallelism.
• The change avoids additional scheduling collectives and prevents forward-progress stalls.
• Dummy allocation failures do not raise.
• Cleanup remains active after padding.
• KV-cache allocation rollback excludes scheduler-bypassing dummies.
• The implementation respects active-request limits, existing dummy state, benchmark-disaggregated fill gates, allocation failures, and KV-cache manager V2 behavior.
• Post-schedule dummy allocations roll back when queue admission fails.
• No public or exported declarations changed.
• No configuration or test-list files changed.

QA Engineer Review

• Added CPU-only test coverage in tests/unittest/_torch/executor/test_py_executor.py for:
  • Empty scheduled-batch padding.
  • Generation-dummy insertion.
  • Padding eligibility and no-op conditions.
  • KV-cache capacity limits.
  • Dummy allocation failures.
  • Active-request limits.
  • Existing dummy requests.
  • Fleet-wide rollback and commit behavior.
  • Exclusion of dummies during generation-allocation rollback.
  • One-model MTP executor behavior with attention-DP padding disabled.
• These are test-code changes outside tests/integration/test_lists/.
• No corresponding test-db/ or qa/ coverage entry was identified.
• Verdict: needs follow-up.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: abedff2e-a116-4dc0-879c-4190c4a57dbb

📥 Commits

Reviewing files that changed from the base of the PR and between c5442d9 and 0723a36.

📒 Files selected for processing (1)

• tests/unittest/_torch/executor/test_py_executor.py

---

Walkthrough

Attention-DP scheduling now pads eligible rank-local empty batches with generation dummies. Allocation failures do not raise. Dummy resources integrate with cleanup and rollback rules. Tests cover eligibility, limits, failures, admission, cleanup, and rollback behavior.

Changes

Attention-DP padding

Layer / File(s)	Summary
Padding flow and allocation cleanup tensorrt_llm/_torch/pyexecutor/py_executor.py	The scheduler invokes _pad_empty_attention_dp_batch after scheduling. The method checks eligibility, allocates generation dummies, handles allocation failures, and appends successful dummies.
Dummy allocation rollback tensorrt_llm/_torch/pyexecutor/py_executor.py	Generation-allocation rollback skips scheduler-bypassing padding dummies. Tentative dummy resources remain reversible when queue admission fails.
Padding behavior validation tests/unittest/_torch/executor/test_py_executor.py	Tests cover eligibility, capacity, allocation errors, active-request limits, duplicate dummies, fleet-wide admission, cleanup, rollback exclusion, and MTP isolation.

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

Sequence Diagram(s)

sequenceDiagram
  participant PyExecutor
  participant Scheduler
  participant KVCache
  participant QueueAdmission
  PyExecutor->>Scheduler: schedule active requests
  Scheduler-->>PyExecutor: return rank-local batch
  PyExecutor->>KVCache: allocate generation dummy for empty batch
  KVCache-->>PyExecutor: return allocation or error
  PyExecutor->>QueueAdmission: submit scheduled batch
  QueueAdmission-->>PyExecutor: accept or reject batch
  PyExecutor->>KVCache: roll back tentative dummy resources on rejection

Loading

Possibly related PRs

• NVIDIA/TensorRT-LLM#16834: Modifies related Attention-DP dummy-padding logic and active-request handling.

Suggested reviewers: schetlur-nv, cascade812, pcastonguay

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title follows the required format and clearly identifies the attention-DP empty-batch padding fix.
Description check	✅ Passed	The description explains the issue, solution, design constraints, tests, hardware validation, and checklist status.
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.

✨ 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: 1

🧹 Nitpick comments (1)

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

1740-1755: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Annotate the new test helpers.

Annotate _run_pad_empty and _unfittable_rank parameters and return values.

As per coding guidelines, “Annotate every function.”

🤖 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 `@tests/unittest/_torch/executor/test_py_executor.py` around lines 1740 - 1755,
Add type annotations to the parameters and return values of the test helpers
_run_pad_empty and _unfittable_rank. Use the existing types for stub and
scheduled_batch, and annotate _unfittable_rank with the tuple type matching the
returned executor stub and ScheduledRequests instance.

Source: Coding guidelines

🤖 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 `@tensorrt_llm/_torch/pyexecutor/py_executor.py`:
- Around line 6088-6092: Update the post-schedule dummy handling in
py_executor.py around _pending_adp_dummy_request so every padding dummy is
registered for failed-queue cleanup, independent of
_enable_dsv4_adp_dummy_fixes, and ensure failed queueing releases it and removes
it from active requests. Add a regression test in
tests/unittest/_torch/executor/test_py_executor.py covering
_enable_dsv4_adp_dummy_fixes=False and verifying cleanup after queue failure.

---

Nitpick comments:
In `@tests/unittest/_torch/executor/test_py_executor.py`:
- Around line 1740-1755: Add type annotations to the parameters and return
values of the test helpers _run_pad_empty and _unfittable_rank. Use the existing
types for stub and scheduled_batch, and annotate _unfittable_rank with the tuple
type matching the returned executor stub and ScheduledRequests instance.

🪄 Autofix

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: ab43b2fd-b84f-4416-aa82-6088df05c550

📥 Commits

Reviewing files that changed from the base of the PR and between 2a40e27 and 92b1b82.

📒 Files selected for processing (2)

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

[xwang233]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #64425 [ run ] triggered by Bot. Commit: 308e9c5 Link to invocation

[tensorrt-cicd]

PR_Github #64425 [ run ] completed with state FAILURE. Commit: 308e9c5
/LLM/main/L0_MergeRequest_PR pipeline #52306 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

[Tabrizian]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #64503 [ run ] triggered by Bot. Commit: 0723a36 Link to invocation