[None][fix] Simplify idle disagg KV transfer progress check - #17324
[None][fix] Simplify idle disagg KV transfer progress check#17324Tabrizian wants to merge 3 commits into
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: |
|
PR_Github #65116 [ run ] completed with state
|
f64c82d to
267b3d8
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 #65678 [ run ] triggered by Bot. Commit: |
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 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 `@tests/unittest/_torch/executor/test_py_executor.py`:
- Around line 958-991: Add the explicit -> None return annotation to
test_polls_context_transfers_without_blocking,
test_does_not_repeat_gen_status_polled_by_loop_head, and
test_idle_poll_enters_no_extra_collective, without changing their behavior.
- Around line 958-1006: Add the explicit -> None return annotation to the three
new test methods: test_polls_context_transfers_without_blocking,
test_does_not_repeat_gen_status_polled_by_loop_head, and
test_idle_poll_enters_no_extra_collective. Leave their test logic unchanged.
🪄 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: b768574c-97f4-45a8-ba32-040f0a3d2fc2
📒 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
🚧 Files skipped from review as they are similar to previous changes (2)
- tests/unittest/_torch/executor/test_benchmark_disagg.py
- tensorrt_llm/_torch/pyexecutor/py_executor.py
| def test_polls_context_transfers_without_blocking(self): | ||
| executor = object.__new__(PyExecutor) | ||
| executor.dist = Mock(tp_size=1) | ||
| executor._check_disagg_gen_cache_transfer_status = Mock() | ||
| executor._check_disagg_ctx_cache_transfer_status = Mock() | ||
|
|
||
| PyExecutor._check_disagg_transfer_progress_when_idle( | ||
| executor, | ||
| num_fitting_reqs=0, | ||
| fitting_disagg_gen_init_requests=[], | ||
| wait_for_disagg_gen_transfer_progress=True, | ||
| all_gen_first=False, | ||
| ) | ||
| PyExecutor._check_disagg_transfer_progress_when_idle(executor) | ||
|
|
||
| executor._check_disagg_gen_cache_transfer_status.assert_called_once_with(1) | ||
| executor._check_disagg_ctx_cache_transfer_status.assert_not_called() | ||
| executor._check_disagg_ctx_cache_transfer_status.assert_called_once_with(0) | ||
|
|
||
| def test_peer_rank_enters_bounded_progress_poll(self): | ||
| def test_does_not_repeat_gen_status_polled_by_loop_head(self): | ||
| """The loop head already polls GEN status every iteration.""" | ||
| executor = object.__new__(PyExecutor) | ||
| executor.dist = Mock(tp_size=1, cp_size=4, world_size=4) | ||
| executor.dist.allreduce.return_value = 1 | ||
| executor.dist = Mock(tp_size=1) | ||
| executor._check_disagg_gen_cache_transfer_status = Mock() | ||
| executor._check_disagg_ctx_cache_transfer_status = Mock() | ||
|
|
||
| PyExecutor._check_disagg_transfer_progress_when_idle( | ||
| executor, | ||
| num_fitting_reqs=1, | ||
| fitting_disagg_gen_init_requests=[], | ||
| wait_for_disagg_gen_transfer_progress=True, | ||
| all_gen_first=False, | ||
| ) | ||
| PyExecutor._check_disagg_transfer_progress_when_idle(executor) | ||
|
|
||
| executor._check_disagg_gen_cache_transfer_status.assert_called_once_with(1) | ||
| executor._check_disagg_ctx_cache_transfer_status.assert_not_called() | ||
| executor.dist.allreduce.assert_called_once_with(0, op=ReduceOp.MAX) | ||
| executor._check_disagg_gen_cache_transfer_status.assert_not_called() | ||
|
|
||
| def test_falls_back_to_context_transfer_when_not_generation_blocked(self): | ||
| def test_idle_poll_enters_no_extra_collective(self): | ||
| """The context poll is rank-symmetric, so no gating collective is needed.""" | ||
| executor = object.__new__(PyExecutor) | ||
| executor.dist = Mock(tp_size=1) | ||
| executor.dist = Mock(tp_size=4, cp_size=4, world_size=16) | ||
| executor._check_disagg_gen_cache_transfer_status = Mock() | ||
| executor._check_disagg_ctx_cache_transfer_status = Mock() | ||
|
|
||
| PyExecutor._check_disagg_transfer_progress_when_idle( | ||
| executor, | ||
| num_fitting_reqs=0, | ||
| fitting_disagg_gen_init_requests=[], | ||
| wait_for_disagg_gen_transfer_progress=False, | ||
| all_gen_first=False, | ||
| ) | ||
| PyExecutor._check_disagg_transfer_progress_when_idle(executor) | ||
|
|
||
| executor._check_disagg_ctx_cache_transfer_status.assert_called_once_with(1) | ||
| executor._check_disagg_gen_cache_transfer_status.assert_not_called() | ||
| executor.dist.allreduce.assert_not_called() | ||
| executor.dist.tp_allreduce.assert_not_called() | ||
| executor.dist.tp_cp_allgather.assert_not_called() | ||
| executor._check_disagg_ctx_cache_transfer_status.assert_called_once_with(0) |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Add return annotations to the changed test methods.
test_polls_context_transfers_without_blocking, test_does_not_repeat_gen_status_polled_by_loop_head, and test_idle_poll_enters_no_extra_collective do not declare -> None.
Proposed fix
- def test_polls_context_transfers_without_blocking(self):
+ def test_polls_context_transfers_without_blocking(self) -> None:
...
- def test_does_not_repeat_gen_status_polled_by_loop_head(self):
+ def test_does_not_repeat_gen_status_polled_by_loop_head(self) -> None:
...
- def test_idle_poll_enters_no_extra_collective(self):
+ def test_idle_poll_enters_no_extra_collective(self) -> None:As per coding guidelines: “Annotate every function.”
📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| def test_polls_context_transfers_without_blocking(self): | |
| executor = object.__new__(PyExecutor) | |
| executor.dist = Mock(tp_size=1) | |
| executor._check_disagg_gen_cache_transfer_status = Mock() | |
| executor._check_disagg_ctx_cache_transfer_status = Mock() | |
| PyExecutor._check_disagg_transfer_progress_when_idle( | |
| executor, | |
| num_fitting_reqs=0, | |
| fitting_disagg_gen_init_requests=[], | |
| wait_for_disagg_gen_transfer_progress=True, | |
| all_gen_first=False, | |
| ) | |
| PyExecutor._check_disagg_transfer_progress_when_idle(executor) | |
| executor._check_disagg_gen_cache_transfer_status.assert_called_once_with(1) | |
| executor._check_disagg_ctx_cache_transfer_status.assert_not_called() | |
| executor._check_disagg_ctx_cache_transfer_status.assert_called_once_with(0) | |
| def test_peer_rank_enters_bounded_progress_poll(self): | |
| def test_does_not_repeat_gen_status_polled_by_loop_head(self): | |
| """The loop head already polls GEN status every iteration.""" | |
| executor = object.__new__(PyExecutor) | |
| executor.dist = Mock(tp_size=1, cp_size=4, world_size=4) | |
| executor.dist.allreduce.return_value = 1 | |
| executor.dist = Mock(tp_size=1) | |
| executor._check_disagg_gen_cache_transfer_status = Mock() | |
| executor._check_disagg_ctx_cache_transfer_status = Mock() | |
| PyExecutor._check_disagg_transfer_progress_when_idle( | |
| executor, | |
| num_fitting_reqs=1, | |
| fitting_disagg_gen_init_requests=[], | |
| wait_for_disagg_gen_transfer_progress=True, | |
| all_gen_first=False, | |
| ) | |
| PyExecutor._check_disagg_transfer_progress_when_idle(executor) | |
| executor._check_disagg_gen_cache_transfer_status.assert_called_once_with(1) | |
| executor._check_disagg_ctx_cache_transfer_status.assert_not_called() | |
| executor.dist.allreduce.assert_called_once_with(0, op=ReduceOp.MAX) | |
| executor._check_disagg_gen_cache_transfer_status.assert_not_called() | |
| def test_falls_back_to_context_transfer_when_not_generation_blocked(self): | |
| def test_idle_poll_enters_no_extra_collective(self): | |
| """The context poll is rank-symmetric, so no gating collective is needed.""" | |
| executor = object.__new__(PyExecutor) | |
| executor.dist = Mock(tp_size=1) | |
| executor.dist = Mock(tp_size=4, cp_size=4, world_size=16) | |
| executor._check_disagg_gen_cache_transfer_status = Mock() | |
| executor._check_disagg_ctx_cache_transfer_status = Mock() | |
| PyExecutor._check_disagg_transfer_progress_when_idle( | |
| executor, | |
| num_fitting_reqs=0, | |
| fitting_disagg_gen_init_requests=[], | |
| wait_for_disagg_gen_transfer_progress=False, | |
| all_gen_first=False, | |
| ) | |
| PyExecutor._check_disagg_transfer_progress_when_idle(executor) | |
| executor._check_disagg_ctx_cache_transfer_status.assert_called_once_with(1) | |
| executor._check_disagg_gen_cache_transfer_status.assert_not_called() | |
| executor.dist.allreduce.assert_not_called() | |
| executor.dist.tp_allreduce.assert_not_called() | |
| executor.dist.tp_cp_allgather.assert_not_called() | |
| executor._check_disagg_ctx_cache_transfer_status.assert_called_once_with(0) | |
| def test_polls_context_transfers_without_blocking(self) -> None: | |
| executor = object.__new__(PyExecutor) | |
| executor.dist = Mock(tp_size=1) | |
| executor._check_disagg_gen_cache_transfer_status = Mock() | |
| executor._check_disagg_ctx_cache_transfer_status = Mock() | |
| PyExecutor._check_disagg_transfer_progress_when_idle(executor) | |
| executor._check_disagg_ctx_cache_transfer_status.assert_called_once_with(0) | |
| def test_does_not_repeat_gen_status_polled_by_loop_head(self) -> None: | |
| """The loop head already polls GEN status every iteration.""" | |
| executor = object.__new__(PyExecutor) | |
| executor.dist = Mock(tp_size=1) | |
| executor._check_disagg_gen_cache_transfer_status = Mock() | |
| executor._check_disagg_ctx_cache_transfer_status = Mock() | |
| PyExecutor._check_disagg_transfer_progress_when_idle(executor) | |
| executor._check_disagg_gen_cache_transfer_status.assert_not_called() | |
| def test_idle_poll_enters_no_extra_collective(self) -> None: | |
| """The context poll is rank-symmetric, so no gating collective is needed.""" | |
| executor = object.__new__(PyExecutor) | |
| executor.dist = Mock(tp_size=4, cp_size=4, world_size=16) | |
| executor._check_disagg_gen_cache_transfer_status = Mock() | |
| executor._check_disagg_ctx_cache_transfer_status = Mock() | |
| PyExecutor._check_disagg_transfer_progress_when_idle(executor) | |
| executor.dist.allreduce.assert_not_called() | |
| executor.dist.tp_allreduce.assert_not_called() | |
| executor.dist.tp_cp_allgather.assert_not_called() | |
| executor._check_disagg_ctx_cache_transfer_status.assert_called_once_with(0) |
🤖 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 958 - 991,
Add the explicit -> None return annotation to
test_polls_context_transfers_without_blocking,
test_does_not_repeat_gen_status_polled_by_loop_head, and
test_idle_poll_enters_no_extra_collective, without changing their behavior.
Source: Coding guidelines
| def test_polls_context_transfers_without_blocking(self): | ||
| executor = object.__new__(PyExecutor) | ||
| executor.dist = Mock(tp_size=1) | ||
| executor._check_disagg_gen_cache_transfer_status = Mock() | ||
| executor._check_disagg_ctx_cache_transfer_status = Mock() | ||
|
|
||
| PyExecutor._check_disagg_transfer_progress_when_idle( | ||
| executor, | ||
| num_fitting_reqs=0, | ||
| fitting_disagg_gen_init_requests=[], | ||
| wait_for_disagg_gen_transfer_progress=True, | ||
| all_gen_first=False, | ||
| ) | ||
| PyExecutor._check_disagg_transfer_progress_when_idle(executor) | ||
|
|
||
| executor._check_disagg_gen_cache_transfer_status.assert_called_once_with(1) | ||
| executor._check_disagg_ctx_cache_transfer_status.assert_not_called() | ||
| executor._check_disagg_ctx_cache_transfer_status.assert_called_once_with(0) | ||
|
|
||
| def test_peer_rank_enters_bounded_progress_poll(self): | ||
| def test_does_not_repeat_gen_status_polled_by_loop_head(self): | ||
| """The loop head already polls GEN status every iteration.""" | ||
| executor = object.__new__(PyExecutor) | ||
| executor.dist = Mock(tp_size=1, cp_size=4, world_size=4) | ||
| executor.dist.allreduce.return_value = 1 | ||
| executor.dist = Mock(tp_size=1) | ||
| executor._check_disagg_gen_cache_transfer_status = Mock() | ||
| executor._check_disagg_ctx_cache_transfer_status = Mock() | ||
|
|
||
| PyExecutor._check_disagg_transfer_progress_when_idle( | ||
| executor, | ||
| num_fitting_reqs=1, | ||
| fitting_disagg_gen_init_requests=[], | ||
| wait_for_disagg_gen_transfer_progress=True, | ||
| all_gen_first=False, | ||
| ) | ||
| PyExecutor._check_disagg_transfer_progress_when_idle(executor) | ||
|
|
||
| executor._check_disagg_gen_cache_transfer_status.assert_called_once_with(1) | ||
| executor._check_disagg_ctx_cache_transfer_status.assert_not_called() | ||
| executor.dist.allreduce.assert_called_once_with(0, op=ReduceOp.MAX) | ||
| executor._check_disagg_gen_cache_transfer_status.assert_not_called() | ||
|
|
||
| def test_falls_back_to_context_transfer_when_not_generation_blocked(self): | ||
| def test_idle_poll_enters_no_extra_collective(self): | ||
| """The context poll is rank-symmetric, so no gating collective is needed.""" | ||
| executor = object.__new__(PyExecutor) | ||
| executor.dist = Mock(tp_size=1) | ||
| executor.dist = Mock(tp_size=4, cp_size=4, world_size=16) | ||
| executor._check_disagg_gen_cache_transfer_status = Mock() | ||
| executor._check_disagg_ctx_cache_transfer_status = Mock() | ||
|
|
||
| PyExecutor._check_disagg_transfer_progress_when_idle( | ||
| executor, | ||
| num_fitting_reqs=0, | ||
| fitting_disagg_gen_init_requests=[], | ||
| wait_for_disagg_gen_transfer_progress=False, | ||
| all_gen_first=False, | ||
| ) | ||
| PyExecutor._check_disagg_transfer_progress_when_idle(executor) | ||
|
|
||
| executor._check_disagg_ctx_cache_transfer_status.assert_called_once_with(1) | ||
| executor._check_disagg_gen_cache_transfer_status.assert_not_called() | ||
| executor.dist.allreduce.assert_not_called() | ||
| executor.dist.tp_allreduce.assert_not_called() | ||
| executor.dist.tp_cp_allgather.assert_not_called() | ||
| executor._check_disagg_ctx_cache_transfer_status.assert_called_once_with(0) | ||
|
|
||
| def test_gen_only_no_context_benchmark_polls_context_when_idle( | ||
| self, monkeypatch: pytest.MonkeyPatch | ||
| ) -> None: | ||
| monkeypatch.setenv("TRTLLM_DISAGG_BENCHMARK_GEN_ONLY", "1") | ||
| executor = object.__new__(PyExecutor) | ||
| executor.dist = Mock(tp_size=4, cp_size=1, world_size=4) | ||
| executor.dist.allreduce.return_value = 0 | ||
| executor.dist.tp_allreduce.return_value = 1 | ||
| executor._check_disagg_gen_cache_transfer_status = Mock() | ||
| executor._check_disagg_ctx_cache_transfer_status = Mock() | ||
|
|
||
| PyExecutor._check_disagg_transfer_progress_when_idle( | ||
| executor, | ||
| num_fitting_reqs=0, | ||
| fitting_disagg_gen_init_requests=[], | ||
| wait_for_disagg_gen_transfer_progress=True, | ||
| all_gen_first=False, | ||
| is_idle=True, | ||
| ) | ||
| PyExecutor._check_disagg_transfer_progress_when_idle(executor) | ||
|
|
||
| executor.dist.allreduce.assert_called_once_with(0, op=ReduceOp.MAX) | ||
| executor.dist.tp_allreduce.assert_called_once_with(1, op=ReduceOp.MAX) | ||
| executor._check_disagg_gen_cache_transfer_status.assert_not_called() | ||
| executor._check_disagg_ctx_cache_transfer_status.assert_called_once_with(1) | ||
| executor.dist.allreduce.assert_not_called() | ||
| executor.dist.tp_allreduce.assert_not_called() | ||
| executor._check_disagg_ctx_cache_transfer_status.assert_called_once_with(0) |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
fd -t f . tests/integration/test_lists
rg -n -C2 \
'test_py_executor|test_polls_context_transfers_without_blocking|test_does_not_repeat_gen_status_polled_by_loop_head|test_idle_poll_enters_no_extra_collective|test_gen_only_no_context_benchmark_polls_context_when_idle' \
tests/integration/test_lists/test-db tests/integration/test_lists/qa || trueRepository: NVIDIA/TensorRT-LLM
Length of output: 4823
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '%s\n' '--- changed files ---'
git diff --stat
git diff --name-status
printf '%s\n' '--- relevant test diff ---'
git diff -- tests/unittest/_torch/executor/test_py_executor.py
printf '%s\n' '--- repository guidance ---'
if [ -f CODING_GUIDELINES.md ]; then
sed -n '1,240p' CODING_GUIDELINES.md
fi
printf '%s\n' '--- test-list references ---'
rg -n -i -C2 \
'test_py_executor|tests/unittest/_torch/executor|unittest.*executor|_torch/executor' \
tests/integration/test_lists || trueRepository: NVIDIA/TensorRT-LLM
Length of output: 21943
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '%s\n' '--- repository state ---'
git status --short
git diff --numstat
printf '%s\n' '--- target test definitions ---'
sed -n '900,1030p' tests/unittest/_torch/executor/test_py_executor.py
printf '%s\n' '--- all relevant list entries ---'
rg -n -i \
'unittest/_torch/executor($|/)|test_py_executor|tests/unittest/_torch/executor' \
tests/integration/test_lists/test-db tests/integration/test_lists/qa || true
printf '%s\n' '--- QA list format and unit-test entries ---'
for file in tests/integration/test_lists/qa/*; do
if rg -q -i 'unittest/' "$file"; then
printf '%s\n' "--- $file"
rg -n -i -C1 'unittest/' "$file" | head -80
fi
done
printf '%s\n' '--- AST annotation check ---'
python3 - <<'PY'
import ast
from pathlib import Path
path = Path("tests/unittest/_torch/executor/test_py_executor.py")
tree = ast.parse(path.read_text())
names = {
"test_polls_context_transfers_without_blocking",
"test_does_not_repeat_gen_status_polled_by_loop_head",
"test_idle_poll_enters_no_extra_collective",
"test_gen_only_no_context_benchmark_polls_context_when_idle",
"_run_sync_idle_progress_rank",
"test_sync_transfer_skips_idle_progress_collectives",
}
for node in ast.walk(tree):
if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) and node.name in names:
args = [*node.args.posonlyargs, *node.args.args, *node.args.kwonlyargs]
missing = [arg.arg for arg in args if arg.annotation is None]
print(
f"{node.name}: return_annotation={node.returns is not None}; "
f"unannotated_parameters={missing}"
)
PYRepository: NVIDIA/TensorRT-LLM
Length of output: 8818
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '%s\n' '--- repository state ---'
git status --short
git diff --numstat
printf '%s\n' '--- target test definitions ---'
sed -n '900,1030p' tests/unittest/_torch/executor/test_py_executor.py
printf '%s\n' '--- all relevant list entries ---'
rg -n -i \
'unittest/_torch/executor($|/)|test_py_executor|tests/unittest/_torch/executor' \
tests/integration/test_lists/test-db tests/integration/test_lists/qa || true
printf '%s\n' '--- QA list format and unit-test entries ---'
for file in tests/integration/test_lists/qa/*; do
if rg -q -i 'unittest/' "$file"; then
printf '%s\n' "--- $file"
rg -n -i -C1 'unittest/' "$file" | head -80
fi
done
printf '%s\n' '--- AST annotation check ---'
python3 - <<'PY'
import ast
from pathlib import Path
path = Path("tests/unittest/_torch/executor/test_py_executor.py")
tree = ast.parse(path.read_text())
names = {
"test_polls_context_transfers_without_blocking",
"test_does_not_repeat_gen_status_polled_by_loop_head",
"test_idle_poll_enters_no_extra_collective",
"test_gen_only_no_context_benchmark_polls_context_when_idle",
"_run_sync_idle_progress_rank",
"test_sync_transfer_skips_idle_progress_collectives",
}
for node in ast.walk(tree):
if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) and node.name in names:
args = [*node.args.posonlyargs, *node.args.kwonlyargs, *node.args.args]
missing = [arg.arg for arg in args if arg.annotation is None]
print(
f"{node.name}: return_annotation={node.returns is not None}; "
f"unannotated_parameters={missing}"
)
PYRepository: NVIDIA/TensorRT-LLM
Length of output: 8818
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '%s\n' '--- test-list documentation ---'
sed -n '1,260p' tests/integration/test_lists/test-db/README.md
sed -n '1,260p' tests/integration/test_lists/qa/README.md
printf '%s\n' '--- nearby QA unit-test conventions ---'
rg -n -i -C2 \
'unittest/_torch|unittest/' \
tests/integration/test_lists/qa/llm_*.txt tests/integration/test_lists/qa/llm_*.yml | head -240Repository: NVIDIA/TensorRT-LLM
Length of output: 8151
Add -> None to the three new test methods: test_polls_context_transfers_without_blocking, test_does_not_repeat_gen_status_polled_by_loop_head, and test_idle_poll_enters_no_extra_collective.
Test coverage summary: sufficient. CI covers test_py_executor.py through the recursive unittest/_torch/executor entries in l0_cpu.yml, l0_dgx_b300.yml, l0_gb300_multi_gpus.yml, and l0_h100.yml. No QA entry is required for these unit tests.
🤖 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 958 - 1006,
Add the explicit -> None return annotation to the three new test methods:
test_polls_context_transfers_without_blocking,
test_does_not_repeat_gen_status_polled_by_loop_head, and
test_idle_poll_enters_no_extra_collective. Leave their test logic unchanged.
Source: Path instructions
|
PR_Github #65678 [ run ] completed with state
|
267b3d8 to
0c043a5
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #65728 [ run ] triggered by Bot. Commit: |
|
PR_Github #65728 [ run ] completed with state
|
|
Hi @Tabrizian , thanks for the effort. Could you clarify the purpose of this PR? Is it for code refactoring or performance optimization? This part of the code has frequently introduced issues in the past. If this PR is mainly a refactor, I'd suggest being cautious with the changes. If it's intended to improve performance, I'd recommend validating the benefits with actual measurements and minimizing the scope of modifications where possible. |
|
@Shixiaowei02 Posting performance numbers for the PR |
|
/bot run --disable-fail-fast |
`_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>
…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>
0c043a5 to
1f0504b
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #66184 [ run ] triggered by Bot. Commit: |
|
PR_Github #66185 [ run ] triggered by Bot. Commit: |
| # sync) and reaps any already-completed transfers without | ||
| # blocking on un-finished ones. | ||
| self._check_disagg_ctx_cache_transfer_status(0) | ||
| self._check_disagg_ctx_cache_transfer_status(0) |
There was a problem hiding this comment.
With the bounded wait gone in both directions, nothing paces this loop when the context worker has queued requests it cannot allocate blocks for, which is the state the warning below names.
| wait_for_disagg_gen_transfer_progress, | ||
| all_gen_first, | ||
| is_idle=scheduled_batch.batch_size == 0) | ||
| if num_fitting_reqs == 0: |
There was a problem hiding this comment.
The warning is back, but without the all_gen_first and fitting_disagg_gen_init_requests conjuncts it had before, so a generation server that is admitting requests normally will now log it every iteration.
| num_fitting_reqs, wait_for_disagg_gen_transfer_progress | ||
| ) = self._pp_schedule_and_propagate(microbatch_id) | ||
| num_fitting_reqs, | ||
| _) = self._pp_schedule_and_propagate(microbatch_id) |
There was a problem hiding this comment.
This still unpacks four values while #17245 changes _pp_schedule_and_propagate to return three, and since the fourth is now a bare underscore the conflicting hunk reads as cosmetic.
|
PR_Github #66184 [ run ] completed with state |
|
PR_Github #66185 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #66271 [ run ] triggered by Bot. Commit: |
|
/bot run --disable-fail-fast |
|
PR_Github #66304 [ run ] triggered by Bot. Commit: |
|
PR_Github #66271 [ run ] completed with state |
|
PR_Github #66304 [ run ] completed with state
|
5 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.