[None][fix] Simplify idle disagg KV transfer progress check - #17324
[None][fix] Simplify idle disagg KV transfer progress check#17324Tabrizian wants to merge 2 commits into
Conversation
`_check_disagg_transfer_progress_when_idle` gated its work behind two rank-collectives (`_sync_disagg_gen_status_entry` / `_sync_disagg_ctx_status_entry`) and then issued a blocking `atLeastNum=1` wait on whichever direction won the vote. The vote input was derived from purely local 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 an extra 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_status` and `_check_disagg_gen_cache_transfer_status` already 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 symmetric without 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. Removes the now-unused `_sync_disagg_gen_status_entry` and `_sync_disagg_ctx_status_entry` helpers and drops the per-iteration `all_gen_first` scan over `active_requests` at both call sites. Signed-off-by: Iman Tabrizian <10105175+tabrizian@users.noreply.github.com>
|
/bot run --disable-fail-fast |
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (3)
💤 Files with no reviewable changes (1)
🚧 Files skipped from review as they are similar to previous changes (2)
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 suppression, and context-transfer backpressure. ChangesDisaggregated transfer polling
Estimated code review effort: 3 (Moderate) | ~20 minutes 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.
🧹 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
|
|
…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>
|
/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. |
4 participants
Dev Engineer Review
PyExecutor.QA Engineer Review
tests/integration/test_lists/entries were modified.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.