[TRTLLM-15178][fix] Pad an empty attention-DP scheduled batch so the fleet can make forward progress - #17379
Conversation
…n make forward progress
`_can_queue` vetoes the forward pass on every attention-DP 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 that 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: peers never execute the
context chunks they did schedule, because `_update_request_states` -- the only
caller of `move_to_next_context_chunk()` -- runs under `if can_queue:`. No chunk
completes, so no KV cache is released, so the starved rank stays starved.
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: after `_schedule()`, a rank left with an empty batch appends a generation
dummy to the scheduled batch. Rank-local by design -- `_schedule()` performs a
`tp_allgather` inside `_balance_adp_requests`, so re-scheduling on only the
empty ranks would desynchronize that collective. A generation dummy is used
rather than the context dummy the pre-schedule path would pick, because a rank
that is empty precisely because it is short of KV cache must not be asked for
`max_num_tokens` worth of it. Every allocation failure degrades to today's
behaviour rather than propagating, since all ranks have yet to agree on
`can_queue` and a rank-local raise would strand the peers in the collectives
that follow. The pipeline-parallel loop does not call
`_prepare_and_schedule_batch` and is unaffected.
Verified on 4-node disaggregated attention-DP context servers on GB200 and
GB300: the stall reproduced deterministically without the change, and with it
the empty-batch condition still occurs and the fleet drives through every
occurrence.
Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com>
|
Note Reviews pausedIt 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 Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
WalkthroughAttention-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, rollback, and MTP isolation. ChangesAttention-DP padding
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant PyExecutor
participant Scheduler
participant KVAllocator
participant QueueAdmission
PyExecutor->>Scheduler: schedule active requests
Scheduler-->>PyExecutor: return rank-local batch
PyExecutor->>KVAllocator: allocate generation dummy for empty batch
KVAllocator-->>PyExecutor: return allocation or error
PyExecutor->>QueueAdmission: submit scheduled batch
QueueAdmission-->>PyExecutor: accept or reject batch
PyExecutor->>KVAllocator: roll back tentative dummy resources on rejection
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
tests/unittest/_torch/executor/test_py_executor.py (1)
1740-1755: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueAnnotate the new test helpers.
Annotate
_run_pad_emptyand_unfittable_rankparameters 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
📒 Files selected for processing (2)
tensorrt_llm/_torch/pyexecutor/py_executor.pytests/unittest/_torch/executor/test_py_executor.py
xwang233
changed the title
…mments Address review feedback: the docstring and inline comments carried more detail than the code needs. Trim to the reasoning a reader must have -- why the padding runs after _schedule(), why it is rank-local, and why a generation dummy -- and leave the full analysis to the PR description. Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com>
… model The post-schedule dummy was only registered for rollback under _enable_dsv4_adp_dummy_fixes, so on other models a fleet-wide can_queue=False left it parked with its KV cache allocated until the next successful forward. _pending_adp_dummy_request is written in exactly two places -- the DeepSeek-V4 branch of _pad_attention_dp_dummy_request and the new padding path -- and every other flow leaves it None, where _finalize_adp_dummy_allocation already returns early. Dropping the model gate from that method is therefore a no-op for the existing paths and closes the gap for the new one. Rollback is needed because the usual teardown, which terminates every attention-DP dummy in _handle_responses, runs only under 'if can_queue:'. Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com>
|
/bot run --disable-fail-fast |
|
PR_Github #64425 [ run ] triggered by Bot. Commit: |
|
PR_Github #64425 [ run ] completed with state
|
… stub
`_pad_empty_attention_dp_batch` is called unconditionally from
`_prepare_and_schedule_batch` and reads `self.enable_attention_dp` on its
first line. `TestOneModelMTPDraftTokenScheduling` builds its executor with
`object.__new__(PyExecutor)` and never assigned that attribute, so the new
call raised:
AttributeError: 'PyExecutor' object has no attribute 'enable_attention_dp'
failing CPU-Generic-arm-1 in `unittest/_torch/executor`.
Set it to False on the stub, matching `test_benchmark_disagg.py`'s executor
stub, so the new padding path returns immediately for a test that does not
exercise attention DP.
Signed-off-by: Iman Tabrizian <10105175+tabrizian@users.noreply.github.com>
|
/bot run --disable-fail-fast |
|
PR_Github #64503 [ run ] triggered by Bot. Commit: |
|
PR_Github #64503 [ run ] completed with state
|
Signed-off-by: Iman Tabrizian <10105175+tabrizian@users.noreply.github.com>
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
|
/bot run --disable-fail-fast |
|
PR_Github #64649 [ run ] triggered by Bot. Commit: |
|
PR_Github #64649 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #64814 [ run ] triggered by Bot. Commit: |
|
PR_Github #64814 [ run ] completed with state
|
|
/bot run --disable-fail-fast --reuse-test |
…ed-batch-deadlock Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com> # Conflicts: # tests/unittest/_torch/executor/test_py_executor.py
|
/bot run --disable-fail-fast |
|
PR_Github #64822 [ run ] triggered by Bot. Commit: |
|
PR_Github #64822 [ run ] completed with state
|
3 participants
Description
Under attention DP,
_can_queuevetoes the forward pass on every rank when any one rank's scheduled batch is empty:_pad_attention_dp_dummy_requestis 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:
can_queueeverywhere, so peers never execute the context chunks they did schedule —_update_request_states, the only caller ofmove_to_next_context_chunk(), runs underif can_queue:;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:
_schedule()performs atp_allgatherinside_balance_adp_requests, so re-running scheduling on only the empty ranks would desynchronize that collective. The dummy is appended toscheduled_batch.generation_requestsinstead — pure rank-local state.max_num_tokensworth of it. A generation dummy costs1 + max_total_draft_tokenstokens, and peers already run generation dummies alongside real context chunks.OutOfPagesError,NoFreeSlotsError, insufficient capacity) degrades to today's behaviour rather than propagating: all ranks have yet to agree oncan_queue, so a rank-local raise here would strand the peers in the collectives that follow.max_num_active_requests, so it cannot trip theexpected_num_active_requestsassert in_pad_attention_dp_dummy_request; refuses to double-allocate the singletonATTENTION_DP_DUMMY_REQUEST_ID; honors the benchmark-disagg fill gate, which suppresses forwards on purpose._revert_gen_allocskips it, since reverting would shrink a real allocation._prepare_and_schedule_batchand 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:
PR Checklist
[JIRA/NVBUG/None][type] SummaryDev Engineer Review
PyExecutorpads rank-local empty scheduled batches with a generation dummy under attention data parallelism.QA Engineer Review
tests/unittest/_torch/executor/test_py_executor.pyfor:tests/integration/test_lists/.test-db/orqa/coverage entry was identified.