[None][fix] Simplify idle disagg KV transfer progress check - #17324
Conversation
|
/bot run --disable-fail-fast |
|
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:
WalkthroughPyExecutor now uses parameterless, nonblocking idle polling for asynchronous context transfers. Pipeline and non-pipeline schedulers discard unused results. Tests cover symmetric polling, synchronous modes, benchmark behavior, and context-transfer backpressure. ChangesDisaggregated transfer polling
Estimated code review effort: 3 (Moderate) | ~20 minutes Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches 💡 1🛠️ Fix failing CI checks 💡
🧪 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.
🧹 Nitpick comments (1)
tests/unittest/_torch/executor/test_py_executor.py (1)
638-674: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd annotations to the new test methods.
The three changed test methods lack the required return annotations. Add
-> Noneto each method. Add a precise type formonkeypatchintest_gen_only_no_context_benchmark_skips_idle_polls.Proposed change
- def test_polls_both_transfer_directions_without_blocking(self): + def test_polls_both_transfer_directions_without_blocking(self) -> None: ... - def test_idle_poll_enters_no_extra_collective(self): + def test_idle_poll_enters_no_extra_collective(self) -> None: ... - def test_gen_only_no_context_benchmark_skips_idle_polls(self, monkeypatch): + def test_gen_only_no_context_benchmark_skips_idle_polls( + self, monkeypatch: pytest.MonkeyPatch) -> None: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 638 - 674, Update the three new test methods in the PyExecutor test class to include a None return annotation; additionally annotate the monkeypatch parameter in test_gen_only_no_context_benchmark_skips_idle_polls with the precise pytest monkeypatch fixture type used by the project.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.
Nitpick comments:
In `@tests/unittest/_torch/executor/test_py_executor.py`:
- Around line 638-674: Update the three new test methods in the PyExecutor test
class to include a None return annotation; additionally annotate the monkeypatch
parameter in test_gen_only_no_context_benchmark_skips_idle_polls with the
precise pytest monkeypatch fixture type used by the project.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: c8ae1ee8-a729-4826-b7b7-ca61c4c5920b
📒 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
|
PR_Github #64131 [ run ] triggered by Bot. Commit: |
…lectives dropped) 115b474 replaced the body of _check_disagg_transfer_progress_when_idle with an early return, on the reasoning that "removal is safe because transfer completion is still reaped at the other call sites in the executor loop". That reasoning is wrong. Every other reap site sits behind the attention-DP _can_queue gate, so a rank whose scheduled batch is empty stops reaping altogether. Combined with the ADP empty-scheduled-batch forward-progress veto (_can_queue vetoes the forward pass on every rank when any one rank's SCHEDULED batch is empty, which _pad_attention_dp_dummy_request does not prevent because it runs before the capacity scheduler), that converts a transient stall into a permanent, silent, fleet-wide hang. Measured, not argued. The GLM-5.2 AgentX CTX-only conc32 cell (Lyris GB200, 4 nodes, max_seq_len 512k, tp8 both roles) hangs 3/3 on this branch, frozen mid-warmup at returned={30,34,26}/37 with exactly one CTX "Observed timeout on context request". Restoring the two non-blocking reaps below -- identical image (923af1a sm100), identical config, this the only variable -- completed warmup and reached the measurement phase with 0 transfer timeouts and the watchdog silent (SLURM 2591522). The empty-batch trigger itself was observed during that run and survived. What this does and does not change: - The per-iteration votes stay gone. _sync_disagg_gen_status_entry (WORLD) and _sync_disagg_ctx_status_entry (TP/CP) are not reinstated, so the 99%-of-method cost 115b474 measured is not reintroduced, and neither is the blocking atLeastNum=1 wait. - Both calls are non-blocking (atLeastNum=0) and rank-uniform: every rank enters them unconditionally and each performs its own internal consensus, so no vote is needed and ranks cannot diverge. - This is NOT a fix for the ADP veto, which is a genuine upstream bug and needs _pad_empty_attention_dp_batch (xiaow, separately). This only removes the permanence that stubbing out the reap introduced. - The transfer admission controller is deliberately left as the passthrough 115b474 made it. Only the idle reap is restored here; the admission arm is still under A/B. Upstream equivalent: NVIDIA#17324 Pre-commit bypassed: this branch carries a pre-existing test-list validation failure (llm_function_core.txt references TestKimiK2, the file defines TestKimiK25) unrelated to this commit, which touches only py_executor.py. Formatting verified separately with the pinned yapf 0.43.0. Signed-off-by: Iman Tabrizian <10105175+tabrizian@users.noreply.github.com>
nv-xtf
left a comment
nv-xtf
left a comment
There was a problem hiding this comment.
The simplification makes sense.Two concerns:
- What paces the idle loop after removing (1)? INIT/TRANS requests remain active, so fetch does not wait and repeated (0) polls may spin hot. When an empty ADP rank already makes _can_queue() false fleet-wide, bounded waiting does not delay a forward that could run in that iteration.
- The loop head already performs gen(0) every iteration. Unless new receives were started during scheduling, this appears to run GEN status—and any required consensus—twice.
This directly overlaps with #17299 (NVBug 6527301; currently draft while we re-validate). Could we converge on a design that keeps non-blocking polling by default, but retains bounded waiting when the batch is not globally queueable and transfer progress can unblock it? The existing _can_queue() result could potentially be computed earlier and reused.
|
PR_Github #64131 [ run ] completed with state
|
|
|
/bot run --disable-fail-fast |
|
PR_Github #64479 [ run ] triggered by Bot. Commit: |
The event-driven direction sounds right as the final solution — +1 to that. One thing I'd like to double-check on "spinning adds no overhead": under ADP the fleet spins together, and the ranks with work don't idle — each iteration they schedule, grow KV capacity for the batch, fail _can_queue, then roll it back via _revert_gen_alloc, plus the per-iteration collectives. Wouldn't that churn at spin frequency be real overhead? |
|
In the KV starved state the loop still has work queued, so the request queue does not block and nothing waits, which is the state the deleted warning named. A short bounded wait on that branch, or keeping the warning, would cover it. |
|
PR_Github #64479 [ run ] completed with state
|
8d6b372 to
4a0ca6c
Compare
|
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 #64555 [ run ] triggered by Bot. Commit: |
…heck `_check_disagg_transfer_progress_when_idle` polled both directions, but the GEN poll was always a repeat of one that already ran earlier in the same iteration: - The loop head (`_executor_loop_pp` / `_prepare_and_schedule_batch`) calls `_check_disagg_gen_transfer_status`, which enters `_check_disagg_gen_cache_transfer_status(0)` unconditionally. - If scheduling started new receives, `_prepare_disagg_gen_init` -> `_recv_disagg_gen_cache` already polls GEN status right after issuing them. So in both cases the second call re-ran the GEN status query and its internal cross-rank consensus for nothing. Keep only the CTX poll here. The synchronous-transfer early return is unchanged: a synchronous GEN receive is rank-local and blocking, so one rank can still be receiving while another is idle, which makes entering the context progress collective unsafe. Signed-off-by: Iman Tabrizian <10105175+tabrizian@users.noreply.github.com>
The idle progress check no longer blocks on a transfer, and the request queue does not block either while INIT/TRANS requests are active, so nothing named the KV-starved state after the blocking branch was removed. Log it again. Signed-off-by: Iman Tabrizian <10105175+tabrizian@users.noreply.github.com>
Removing the blocking `atLeastNum=1` wait also removed the only pacing for the idle-blocked case. `_fetch_and_enqueue_requests` uses a zero timeout while any request is active, so once nothing fits, the loop re-ran the ADP allgather, both transfer-status gathers and a full schedule pass at full speed until a transfer completed, burning a core and contending with the transceiver's own progress threads. Sleep 1ms on iterations where `_can_queue` came back False. `can_queue` is already a fleet-wide consensus, so every rank pauses on the same iterations without adding a collective, and the loop head still runs each iteration so cancellations, control requests, shutdown and generation-transfer reaping keep being serviced. Signed-off-by: Iman Tabrizian <10105175+tabrizian@users.noreply.github.com>
The restored warning is broader than the one it replaced, which also required `not fitting_disagg_gen_init_requests` and suppressed the generation-first case. `num_fitting_reqs` does not count fitting gen-init requests -- the capacity scheduler returns those separately -- so the condition holds during normal generation ramp-up and while a context server drains in-flight transfers, making it fire every iteration in states that are not actually degraded. Drop it rather than re-adding the guards: the caller already has the information, and a log that cannot distinguish KV exhaustion from routine ramp-up is not worth the rate-limiting it would need. `num_fitting_reqs` becomes unused at both scheduling sites, so those revert to `_`. `_pp_schedule_and_propagate` still returns it as part of its serialized contract. Signed-off-by: Iman Tabrizian <10105175+tabrizian@users.noreply.github.com>
2f815d5 to
cd39ea8
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #68393 [ run ] triggered by Bot. Commit: |
|
PR_Github #68378 [ run ] completed with state |
|
PR_Github #68393 [ run ] completed with state |
… work The 1ms sleep sat between `_can_queue` and the rest of the iteration, so it delayed work that was already finished and waiting to be handed off: on a `not can_queue` pass the overlap loop still runs `_update_requests`, `_send_kv_async` and the first-token responses for the previous batch, and both loops still flush pending transfer responses and enter the synced transfer-timeout collective. None of that needed to wait a millisecond. Move the sleep to the end of the iteration, once that work has drained, and put it behind `_pace_idle_disagg_loop`. Sleeping there also means the pending-transfer check reads the state the pass left behind rather than the state it started with, so a transfer that landed during the pass no longer costs a sleep. Only pace when a KV transfer is what the loop is waiting on. Context sends are tracked by the transfer manager and generation receives by request state, so both directions are covered. Outside disaggregated serving, and on a rank with nothing in flight, the sleep is skipped entirely -- it never gates a collective, so ranks taking it on different iterations is safe. Signed-off-by: Iman Tabrizian <10105175+tabrizian@users.noreply.github.com>
The pacing added for the other two loops skipped `_executor_loop_pp`, which can reach the same state: a microbatch that cannot be queued leaves the loop re-running the schedule pass and its ring collectives at full speed. It is only the same state once the ring is empty. While microbatches are still in flight, `fetch_executed_batches` sets `must_get = not can_queue` and blocks on the response queue until one finishes, so the loop is already paced -- and sleeping there would delay the relay of batches that are finished and waiting to be handled. Pace only when nothing is queued and nothing is outstanding, and do it after stage 3.4 so a rebalance that the drained ring just unblocked still completes in the same iteration. Signed-off-by: Iman Tabrizian <10105175+tabrizian@users.noreply.github.com>
|
/bot run --disable-fail-fast |
|
PR_Github #68906 [ run ] triggered by Bot. Commit: |
|
PR_Github #68906 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #69260 [ run ] triggered by Bot. Commit: |
|
PR_Github #69260 [ run ] completed with state |
PR NVIDIA#17324 left verbose, redundant docstrings/comments on _check_disagg_transfer_progress_when_idle, _pace_idle_disagg_loop and _pp_ring_is_drained, and dropped an is_idle parameter that the body of _check_disagg_transfer_progress_when_idle still referenced (a NameError on the single-rank CTX idle-reap path). Restore the parameter and its call sites, and cut the comments down to what's non-obvious. Also drops the PP-ring test added alongside the pacing change, which only exercised a private helper directly. Signed-off-by: Iman Tabrizian <10105175+tabrizian@users.noreply.github.com>
Drops the _pace_idle_disagg_loop unit test coverage added by PR NVIDIA#17324. Signed-off-by: Iman Tabrizian <10105175+tabrizian@users.noreply.github.com>
The function is called every iteration regardless of idleness, and the multi-rank/async-transfer branch already polls context status unconditionally, so the single-rank branch's extra is_idle and has_any_inflight_requests gating was inconsistent with that. Poll unconditionally there too, matching how every other branch behaves. Also fixes test_sync_single_rank_ctx_reaps_idle_transfer, which called the old (pre-NVIDIA#17324) five-argument signature and would TypeError. Signed-off-by: Iman Tabrizian <10105175+tabrizian@users.noreply.github.com>
Stops one commit short of origin/main tip 5ae89d6 (NVIDIA#17324), which leaves an undefined is_idle reference in py_executor.py and fails the ruff-legacy hook. Signed-off-by: Pranav Shrestha <254760092+pranav-nvidia@users.noreply.github.com>
PR NVIDIA#17324 left verbose, redundant docstrings/comments on _check_disagg_transfer_progress_when_idle, _pace_idle_disagg_loop and _pp_ring_is_drained, and dropped an is_idle parameter that the body of _check_disagg_transfer_progress_when_idle still referenced (a NameError on the single-rank CTX idle-reap path). Restore the parameter and its call sites, and cut the comments down to what's non-obvious. Also drops the PP-ring test added alongside the pacing change, which only exercised a private helper directly. Signed-off-by: Iman Tabrizian <10105175+tabrizian@users.noreply.github.com>
Drops the _pace_idle_disagg_loop unit test coverage added by PR NVIDIA#17324. Signed-off-by: Iman Tabrizian <10105175+tabrizian@users.noreply.github.com>
The function is called every iteration regardless of idleness, and the multi-rank/async-transfer branch already polls context status unconditionally, so the single-rank branch's extra is_idle and has_any_inflight_requests gating was inconsistent with that. Poll unconditionally there too, matching how every other branch behaves. Also fixes test_sync_single_rank_ctx_reaps_idle_transfer, which called the old (pre-NVIDIA#17324) five-argument signature and would TypeError. Signed-off-by: Iman Tabrizian <10105175+tabrizian@users.noreply.github.com>
8 participants
Dev Engineer Review
PyExecutor.atLeastNum=0.QA Engineer Review
tests/integration/test_lists/entries were modified, so no CI or manual-QA test-list coverage is shown.Description
_check_disagg_transfer_progress_when_idlegated its work behind two rank-collectives(
_sync_disagg_gen_status_entry/_sync_disagg_ctx_status_entry) and then issued a blockingatLeastNum=1wait on whichever direction won the vote. The vote input was derived from purelylocal scheduler state (
num_fitting_reqs,fitting_disagg_gen_init_requests,wait_for_disagg_gen_transfer_progress,all_gen_first), so every disagg iteration paid for anextra allreduce or allgather just to decide whether to poll, and the winning branch could block the
executor loop on an unfinished transfer.
Both
_check_disagg_ctx_cache_transfer_statusand_check_disagg_gen_cache_transfer_statusalready perform their own internal cross-rank consensus and are safe to enter unconditionally with
atLeastNum=0. Entering both non-blocking polls on every iteration keeps all ranks symmetricwithout the extra collective, and reaps completed transfers so their KV blocks are freed just the
same. Ranks with nothing in flight simply reap nothing.
The synchronous-transfer early return is preserved: a synchronous GEN receive is rank-local and
blocking, so one rank can still be receiving while another is idle, which makes entering either
progress collective unsafe. This covers both
TRTLLM_DISABLE_KV_CACHE_TRANSFER_OVERLAP=1and thegen_only_no_contextbenchmark mode.Also removes the now-unused
_sync_disagg_gen_status_entryand_sync_disagg_ctx_status_entryhelpers, and drops the per-iteration
all_gen_firstscan overactive_requestsat both call sites(
_executor_loop_ppand_prepare_and_schedule_batch).No API change; the method is private to
PyExecutor.Test Coverage
Updated
tests/unittest/_torch/executor/test_py_executor.py::TestDisaggTransferIdleProgress:test_polls_both_transfer_directions_without_blocking— both directions are polled withatLeastNum=0.test_idle_poll_enters_no_extra_collective— withtp_size=4, cp_size=4, world_size=16, noallreduce/tp_allreduce/tp_cp_allgatheris entered.test_gen_only_no_context_benchmark_skips_idle_polls— new coverage for theTRTLLM_DISAGG_BENCHMARK_GEN_ONLY=1branch of the preserved guard.test_sync_benchmark_skips_idle_transfer_collectives/test_sync_non_benchmark_skips_idle_transfer_collectives— retained, confirming thesynchronous-transfer early return still suppresses every poll and collective.
Tests asserting the removed vote-then-block behavior
(
test_polls_generation_transfer_when_admission_blocked,test_peer_rank_enters_bounded_progress_poll,test_falls_back_to_context_transfer_when_not_generation_blocked,test_peer_cp_rank_enters_context_progress_poll) are removed or replaced.tests/unittest/_torch/executor/test_benchmark_disagg.pyupdated for the new call signature andfor the non-blocking poll pair under transfer-admission backpressure.
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.