[https://nvbugs/6527301][fix] Do not block ranks with scheduled forward work on disagg gen KV-transfer progress - #17299
Conversation
…on disagg gen transfer progress Signed-off-by: Tingfeng Xian <289617005+nv-xtf@users.noreply.github.com>
…it constants with underscore Signed-off-by: Tingfeng Xian <289617005+nv-xtf@users.noreply.github.com>
WalkthroughPyExecutor now uses bitwise-OR status reductions to coordinate transfer progress and forward work across ranks. Idle polling uses non-blocking or bounded blocking behavior based on the combined status. Tests cover the updated call sites and polling paths. ChangesDisaggregated transfer progress
Estimated code review effort: 3 (Moderate) | ~20 minutes Sequence Diagram(s)sequenceDiagram
participant PyExecutor
participant torch.distributed
participant TransferProgress
PyExecutor->>torch.distributed: Reduce transfer and forward-work status with ReduceOp.BOR
torch.distributed-->>PyExecutor: Return combined status
PyExecutor->>TransferProgress: Poll non-blockingly when forward work exists
PyExecutor->>TransferProgress: Poll with bounded blocking when all ranks are idle
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.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tests/unittest/_torch/executor/test_py_executor.py (1)
644-711: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winAdd return annotations to the changed test methods.
Lines 644, 663, 685, and 711 define changed functions without
-> None. Add return annotations.Proposed fix
- def test_polls_generation_transfer_when_admission_blocked_and_idle(self): + def test_polls_generation_transfer_when_admission_blocked_and_idle(self) -> None: ... - def test_same_rank_progress_need_with_forward_work_reaps_nonblocking(self): + def test_same_rank_progress_need_with_forward_work_reaps_nonblocking(self) -> None: ... - def test_peer_backpressure_with_forward_work_reaps_nonblocking(self): + def test_peer_backpressure_with_forward_work_reaps_nonblocking(self) -> None: ... - def test_peer_backpressure_with_all_ranks_idle_enters_bounded_poll(self): + def test_peer_backpressure_with_all_ranks_idle_enters_bounded_poll(self) -> None:🤖 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 644 - 711, Add the explicit -> None return annotation to each changed test method: test_polls_generation_transfer_when_admission_blocked_and_idle, test_same_rank_progress_need_with_forward_work_reaps_nonblocking, test_peer_backpressure_with_forward_work_reaps_nonblocking, and test_peer_backpressure_with_all_ranks_idle_enters_bounded_poll. Do not alter their test logic.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.
Outside diff comments:
In `@tests/unittest/_torch/executor/test_py_executor.py`:
- Around line 644-711: Add the explicit -> None return annotation to each
changed test method:
test_polls_generation_transfer_when_admission_blocked_and_idle,
test_same_rank_progress_need_with_forward_work_reaps_nonblocking,
test_peer_backpressure_with_forward_work_reaps_nonblocking, and
test_peer_backpressure_with_all_ranks_idle_enters_bounded_poll. Do not alter
their test logic.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: fd4e15f2-b330-4a75-9dc2-3d5f359699e9
📒 Files selected for processing (3)
tensorrt_llm/_torch/pyexecutor/py_executor.pytests/unittest/_torch/executor/test_benchmark_disagg.pytests/unittest/_torch/executor/test_py_executor.py
|
/bot run --disable-fail-fast |
|
PR_Github #64001 [ run ] triggered by Bot. Commit: |
|
Under attention DP the rank signalling need has an empty batch by construction, so the queue-admission vote still blocks every rank, the throttle is gone but no work is unblocked. Does the repro sit outside that regime? |
|
PR_Github #64001 [ run ] completed with state
|
@Shixiaowei02 Fair challenge — you're right. Verified the chain: in the standard v1/v2 flow a NEED-signalling rank has an empty batch by construction, so _can_queue gates forward fleet-wide regardless of the poll — in that regime this patch only removes the pacing without unfreezing decode. I had missed the _can_queue layer. Where the patch genuinely helps is when the final batch diverges from num_fitting_reqs (one existing path: the MixedMamba disagg WAR) — there it removes an up-to-5s stall in front of a forward that would actually run. Whether it addresses the actual NVBug 6527301 issue needs re-validation. Converting this PR to draft until then. Note #17324 is also changing this logic (never waiting on KV transfer under async disagg, non-blocking reap only) — the work overlaps, so I'll coordinate with that PR before converging on a design. |
3 participants
Description
Problem
In the disaggregated GEN worker's scheduler loop,
_check_disagg_transfer_progress_when_idleamplifies local admission backpressure into a global stall: when any rank goes idle waiting for KV-transfer admission budget (in-flight transfers fillingmax_tokens_in_buffer, candidates deferred), the flag is broadcast viaallreduce(MAX)and every rank — including ranks that already scheduled a decode/context batch this iteration — enters the blockingcheck_gen_transfer_status(1)poll, up tokv_transfer_poll_interval_ms(default 5s) per iteration.Scheduling happens before the wait and there is no re-scheduling afterwards, so a transfer completed during the wait can never join the current batch — for ranks with scheduled work the block has zero benefit and simply freezes decode.
Solution
Extend the cross-rank sync from a single boolean to two status bits, combined in one
allreduce(BOR)(same collective count and scalar payload as the previous MAX):_DISAGG_GEN_NEED_TRANSFER_PROGRESS: this rank is idle, admitted nothing this iteration, and has candidates blocked by in-flight transfers;_DISAGG_GEN_HAS_FORWARD_WORK: this rank scheduled a non-empty batch (scheduled_batch.batch_size > 0, passed explicitly from the call sites).All ranks see the identical OR-combined global state and decide uniformly:
check(0), then proceed to forwardcheck(1)(preserves the anti-busy-spin semantics from #15356 / #15737)Two invariants are preserved: (1) all ranks branch on the same global state, so the consensus collectives inside
check_gen_transfer_statusalways receive uniformat_least_request_num(no deadlock risk); (2) the existing contract "any rank needs progress ⇒ a status check happens at this point" is unchanged, so error/cancellation detection timing is unaffected.Compatibility boundary
Behavior is unchanged for single-rank and lockstep-scheduled (non-attention-DP) configurations in the common scheduler flow: the batch is drawn from the fitting set, so
batch_size > 0impliesnum_fitting_reqs > 0and the two bits are mutually exclusive on one rank. One known exception: the MixedMamba disagg WAR in_scheduleoverwritesnum_fittingafter filtering context requests (without counting retained generation requests), making the bits coexist on one rank — that corner moves from the old blocking(1)to non-blocking(0), i.e. a rank that can run forward no longer blocks, which is the work-conserving behavior this fix intends.Out of scope (tracked separately under NVBug 6527301): transfer budget sizing, CTX prefill performance, and the 60s default transfer timeout. The analogous ctx-side branch is left untouched.
Test Coverage
tests/unittest/_torch/executor/test_py_executor.py(TestDisaggTransferIdleProgress, CPU-only):test_peer_backpressure_with_forward_work_reaps_nonblocking: the core scenario — peer backpressure + local scheduled batch ⇒ all ranks call(0)(replacestest_peer_rank_enters_bounded_progress_poll, which asserted the old all-ranks-block behavior);test_peer_backpressure_with_all_ranks_idle_enters_bounded_poll: globally idle ⇒ bounded blocking(1)is preserved;test_same_rank_progress_need_with_forward_work_reaps_nonblocking: NEED and WORK coexisting on one rank (reachable via the MixedMamba disagg WAR) ⇒(0);test_polls_generation_transfer_when_admission_blocked_and_idle: single-rank idle blocking semantics preserved;tests/unittest/_torch/executor/test_benchmark_disagg.py: updated the_prepare_and_schedule_batchargument assertion.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-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin 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.Dev Engineer Review
allreduce(BOR).check(0)polling when transfer progress is required.check(1)polling only when all ranks are idle.QA Engineer Review
tests/unittest/_torch/executor/test_py_executor.py.tests/unittest/_torch/executor/test_benchmark_disagg.py.tests/integration/test_lists/changes were included.