[None][chore] Remove the disagg transfer admission controller - #17245
[None][chore] Remove the disagg transfer admission controller#17245Tabrizian wants to merge 2 commits into
Conversation
|
/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 (2)
🚧 Files skipped from review as they are similar to previous changes (1)
WalkthroughChangesDisaggregated transfer admission removal
Estimated code review effort: 3 (Moderate) | ~20 minutes Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
Remove the FCFS admission gate for disaggregated generation KV transfers from the PyTorch executor, along with the plumbing that existed only to serve it: - `DisaggTransferAdmissionController` / `DisaggTransferAdmissionResult` - `_get_disagg_transfer_admission_controller`, `_apply_disagg_transfer_admission` - `_revert_deferred_disagg_gen_init_alloc`, `_uses_kv_manager_v2`, `_revert_ctx_alloc` - The `wait_for_disagg_gen_transfer_progress` flag threaded through `SerializableSchedulerOutput` and the PP schedule broadcast, and the generation-side branch of `_check_disagg_transfer_progress_when_idle` it gated (`_sync_disagg_gen_status_entry`) The `max_tokens_in_buffer` / `kv_transfer_poll_interval_ms` cache transceiver config fields and the C++ `DisaggTransferAdmissionController` used by the TensorRT backend are left untouched. Signed-off-by: Iman Tabrizian <10105175+tabrizian@users.noreply.github.com>
2e982e4 to
8bcdc18
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #63843 [ run ] triggered by Bot. Commit: |
![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
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
tests/unittest/_torch/executor/test_py_executor.py (1)
474-485: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winPin the transfer mode for this test.
_check_disagg_transfer_progress_when_idle()returns before the expected poll whenTRTLLM_DISABLE_KV_CACHE_TRANSFER_OVERLAP=1orTRTLLM_DISAGG_BENCHMARK_GEN_ONLY=1. This test inherits both variables. A configured test environment can make the assertion at Line 487 fail without a product regression.Use
monkeypatchto set both variables to"0".Proposed fix
- def test_falls_back_to_context_transfer_when_idle(self): + def test_falls_back_to_context_transfer_when_idle(self, monkeypatch): + monkeypatch.setenv("TRTLLM_DISABLE_KV_CACHE_TRANSFER_OVERLAP", "0") + monkeypatch.setenv("TRTLLM_DISAGG_BENCHMARK_GEN_ONLY", "0") executor = object.__new__(PyExecutor)🤖 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 474 - 485, Update test_falls_back_to_context_transfer_when_idle to accept monkeypatch and set TRTLLM_DISABLE_KV_CACHE_TRANSFER_OVERLAP and TRTLLM_DISAGG_BENCHMARK_GEN_ONLY to "0" before invoking PyExecutor._check_disagg_transfer_progress_when_idle, ensuring the test always exercises the intended transfer path.tensorrt_llm/_torch/pyexecutor/py_executor.py (1)
2454-2460: 🗄️ Data Integrity & Integration | 🟠 Major | 🏗️ Heavy liftReconcile or roll back allocations from local PP scheduling.
KVCacheV2Scheduler.schedule_request()allocates generation KV and disaggregated-init KV. The non-rank-0 call discards its result, while rollback and disaggregated preparation use only the propagated schedule. If local capacity selects or evicts a different request, local KV and request state diverge. Use the propagated schedule for mutations, or roll back every local-only allocation. Add a PP test with unequal local KV capacity.🤖 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 `@tensorrt_llm/_torch/pyexecutor/py_executor.py` around lines 2454 - 2460, The local scheduler invocation around self.scheduler.schedule_request must not leave allocations or request-state mutations that differ from the propagated schedule. Update the PP scheduling flow to apply mutations only from the propagated schedule, or explicitly roll back all KV allocations and disaggregated-init state produced by non-rank-0 local scheduling; add coverage for unequal per-rank KV capacity to verify no divergence.
🧹 Nitpick comments (1)
tensorrt_llm/_torch/pyexecutor/py_executor.py (1)
2365-2368: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDeclare the three-value return type.
_pp_schedule_and_propagate()now exposes a changed scheduling contract. Add-> tuple[ScheduledRequests, list[LlmRequest], int]so static checks detect stale four-value callers.Proposed fix
- def _pp_schedule_and_propagate(self, microbatch_id: int): + def _pp_schedule_and_propagate( + self, microbatch_id: int + ) -> tuple[ScheduledRequests, list[LlmRequest], int]: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 `@tensorrt_llm/_torch/pyexecutor/py_executor.py` around lines 2365 - 2368, Annotate _pp_schedule_and_propagate with the explicit return type tuple[ScheduledRequests, list[LlmRequest], int], matching its three returned values and the updated scheduling contract.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 3311-3315: Update the synchronous branch of the context transfer
progress flow around _uses_async_disagg_gen_transfer() so it still polls and
drains completions from context respond_and_send_async() transfers without
entering the unsafe progress collective. Ensure completed sends release their
pinned KV blocks and add coverage for an in-flight context transfer with no
fitting request while KV-cache transfer overlap is disabled.
---
Outside diff comments:
In `@tensorrt_llm/_torch/pyexecutor/py_executor.py`:
- Around line 2454-2460: The local scheduler invocation around
self.scheduler.schedule_request must not leave allocations or request-state
mutations that differ from the propagated schedule. Update the PP scheduling
flow to apply mutations only from the propagated schedule, or explicitly roll
back all KV allocations and disaggregated-init state produced by non-rank-0
local scheduling; add coverage for unequal per-rank KV capacity to verify no
divergence.
In `@tests/unittest/_torch/executor/test_py_executor.py`:
- Around line 474-485: Update test_falls_back_to_context_transfer_when_idle to
accept monkeypatch and set TRTLLM_DISABLE_KV_CACHE_TRANSFER_OVERLAP and
TRTLLM_DISAGG_BENCHMARK_GEN_ONLY to "0" before invoking
PyExecutor._check_disagg_transfer_progress_when_idle, ensuring the test always
exercises the intended transfer path.
---
Nitpick comments:
In `@tensorrt_llm/_torch/pyexecutor/py_executor.py`:
- Around line 2365-2368: Annotate _pp_schedule_and_propagate with the explicit
return type tuple[ScheduledRequests, list[LlmRequest], int], matching its three
returned values and the updated scheduling contract.
🪄 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: aeb2cf6e-5399-44a0-b3a6-cb352167d8c4
📒 Files selected for processing (5)
tensorrt_llm/_torch/pyexecutor/py_executor.pytensorrt_llm/_torch/pyexecutor/scheduler/scheduler.pytests/unittest/_torch/executor/test_benchmark_disagg.pytests/unittest/_torch/executor/test_py_executor.pytests/unittest/_torch/executor/test_scheduler_serializable_output.py
💤 Files with no reviewable changes (2)
- tests/unittest/_torch/executor/test_scheduler_serializable_output.py
- tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py
| # A synchronous GEN receive is rank-local and blocking. One rank can | ||
| # still be receiving while another is idle, so entering either the | ||
| # generation or context progress collective here is unsafe. | ||
| # still be receiving while another is idle, so entering the context | ||
| # progress collective here is unsafe. | ||
| if not self._uses_async_disagg_gen_transfer(): | ||
| return |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/usr/bin/env bash
set -euo pipefail
rg -n -C 10 \
'TRTLLM_DISABLE_KV_CACHE_TRANSFER_OVERLAP|respond_and_send_async|check_context_transfer_status|_check_disagg_transfer_progress_when_idle|_check_kv_transfer_timeout' \
tensorrt_llm/_torch/pyexecutor
rg -n -C 8 --glob 'test_*.py' \
'sync.*context|context.*sync|_check_disagg_transfer_progress_when_idle' \
tests/unittest/_torch/executorRepository: NVIDIA/TensorRT-LLM
Length of output: 50375
🏁 Script executed:
#!/usr/bin/env bash
set -euo pipefail
printf '%s\n' '--- idle progress helper ---'
sed -n '3248,3370p' tensorrt_llm/_torch/pyexecutor/py_executor.py
printf '%s\n' '--- transfer start and context completion cleanup ---'
sed -n '5928,6075p' tensorrt_llm/_torch/pyexecutor/py_executor.py
printf '%s\n' '--- transfer manager completion API ---'
rg -n -C 8 \
'class .*Transfer|def (start_transfer|end_transfer|requests_in_transfer|has_any_inflight_requests)|end_transfer\(' \
tensorrt_llm/_torch/pyexecutor
printf '%s\n' '--- focused tests ---'
sed -n '430,650p' tests/unittest/_torch/executor/test_py_executor.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 36104
🏁 Script executed:
#!/usr/bin/env bash
set -euo pipefail
printf '%s\n' '--- _send_kv_async call sites and surrounding loop ---'
rg -n -C 12 'self\._send_kv_async\(' tensorrt_llm/_torch/pyexecutor/py_executor.py
sed -n '3860,3990p' tensorrt_llm/_torch/pyexecutor/py_executor.py
printf '%s\n' '--- synchronous receive and context-send ordering ---'
rg -n -C 18 \
'def _recv_disagg_gen_cache|request_and_receive_sync|def _send_kv_async|_check_disagg_ctx_cache_transfer_status\(0\)' \
tensorrt_llm/_torch/pyexecutor/py_executor.py
printf '%s\n' '--- context-transfer tests and async-send tests ---'
rg -n -C 12 \
'_send_kv_async|respond_and_send_async|check_context_transfer_status|in.flight|inflight|num_fitting_reqs' \
tests/unittest/_torch/executorRepository: NVIDIA/TensorRT-LLM
Length of output: 50376
🏁 Script executed:
#!/usr/bin/env bash
set -euo pipefail
printf '%s\n' '--- non-overlap loop control around scheduling and cleanup ---'
sed -n '3780,3978p' tensorrt_llm/_torch/pyexecutor/py_executor.py
printf '%s\n' '--- idle helper call sites ---'
rg -n -C 20 \
'_check_disagg_transfer_progress_when_idle\(' \
tensorrt_llm/_torch/pyexecutor/py_executor.py
printf '%s\n' '--- AST call-site verifier ---'
python3 - <<'PY'
import ast
from pathlib import Path
path = Path("tensorrt_llm/_torch/pyexecutor/py_executor.py")
tree = ast.parse(path.read_text())
wanted = {
"_send_kv_async",
"_check_disagg_ctx_cache_transfer_status",
"_check_disagg_transfer_progress_when_idle",
}
for node in ast.walk(tree):
if not isinstance(node, ast.Call):
continue
func = node.func
if not isinstance(func, ast.Attribute) or func.attr not in wanted:
continue
parents = []
# Recompute ancestors for this node with a compact recursive walk.
def walk(current, stack):
if current is node:
return stack
for child in ast.iter_child_nodes(current):
found = walk(child, stack + [current])
if found is not None:
return found
return None
ancestors = walk(tree, [])
functions = [
f"{type(parent).__name__}:{getattr(parent, 'name', '')}"
for parent in ancestors
if isinstance(parent, (ast.FunctionDef, ast.AsyncFunctionDef))
]
controls = [
type(parent).__name__
for parent in ancestors
if isinstance(parent, (ast.If, ast.For, ast.While, ast.Try))
]
print(
f"line={node.lineno} call={func.attr} "
f"function={functions[-1] if functions else '<module>'} "
f"controls={','.join(controls)}"
)
PYRepository: NVIDIA/TensorRT-LLM
Length of output: 19793
Drain context transfers in synchronous mode. When TRTLLM_DISABLE_KV_CACHE_TRANSFER_OVERLAP=1, context sends still use respond_and_send_async(), but the normal no-overlap loop skips _send_kv_async() and its completion poll when can_queue is false. Line 3314 also returns before polling. Completed sends can therefore keep KV blocks pinned and stall scheduling. Add a rank-safe context completion path and a test with an in-flight context transfer and no fitting request.
🤖 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 `@tensorrt_llm/_torch/pyexecutor/py_executor.py` around lines 3311 - 3315,
Update the synchronous branch of the context transfer progress flow around
_uses_async_disagg_gen_transfer() so it still polls and drains completions from
context respond_and_send_async() transfers without entering the unsafe progress
collective. Ensure completed sends release their pinned KV blocks and add
coverage for an in-flight context transfer with no fitting request while
KV-cache transfer overlap is disabled.
|
PR_Github #63843 [ run ] completed with state
|
cascade812
left a comment
cascade812
left a comment
There was a problem hiding this comment.
The removal itself looks clean to me.
Could you explain on why it's safe to drop the gate? Specifically:
- Without the gate, all scheduler-fitting gen INIT requests now start their KV transfers immediately, so concurrent transfers can exceed
max_tokens_in_buffer. Is that backpressure now handled in the transceiver layer itself (e.g., falling back to sync/unbuffered transfers), or was the original overflow concern found to be a non-issue in practice? - Or was the gate actively causing problems?
|
@cascade812 Thanks for the review. Yes, the cache transceiver already handles this and does not send multiple requests if it doesn't have enough cache transceiver buffer. The extra syncs added some overhead and complicated the request scheduling logic. |
|
/bot run --disable-fail-fast |
CI build #51784 failed Release-Check because pre-commit modified two files: - `yapf` rewrapped the `_sync_gen_only_benchmark_has_insufficient_kv` signature in `py_executor.py` (a yapf-owned file per the pre-commit `common-files` list). - `ruff` removed the now-unused `SerializableSchedulerOutput` import from `test_py_executor.py`; that symbol moved to `test_scheduler_serializable_output.py` with the test that uses it. Apply both so the branch is formatted as the hooks expect. Signed-off-by: Iman Tabrizian <10105175+tabrizian@users.noreply.github.com>
|
/bot run --disable-fail-fast |
|
PR_Github #64483 [ run ] triggered by Bot. Commit: |
|
PR_Github #64484 [ run ] triggered by Bot. Commit: |
|
PR_Github/17245-65c3ccb #64483 was force-killed by a newer pipeline run. |
3 participants
Description
Removes the FCFS admission gate for disaggregated generation KV transfers from
the PyTorch executor (
DisaggTransferAdmissionController), introduced inb6eacd1 (
[TRTLLM-12721][fix] Bound V2 context transfer polling, #15356).Removed along with it, because each became dead once the gate was gone:
DisaggTransferAdmissionController/DisaggTransferAdmissionResult_get_disagg_transfer_admission_controller,_apply_disagg_transfer_admission_revert_deferred_disagg_gen_init_alloc,_uses_kv_manager_v2— only evercalled to undo KV growth for gate-deferred requests
_revert_ctx_alloc— no remaining callerswait_for_disagg_gen_transfer_progressflag threaded throughSerializableSchedulerOutputand the PP schedule broadcast. With the gategone it is permanently
False, which made the generation-side branch of_check_disagg_transfer_progress_when_idle(and its_sync_disagg_gen_status_entryallreduce) unreachable, so both are removed.One behavioral note beyond dead-code deletion: in the PP loop, non-rank-0 ranks
still run
scheduler.schedule_request(...)for its request-state side effects,but no longer reconcile their local disagg candidates against rank 0's admitted
set. That reconciliation only had an effect when the gate deferred requests, so
it is a no-op now.
Deliberately left untouched:
DisaggTransferAdmissionController(
cpp/include/tensorrt_llm/batch_manager/disaggTransferAdmissionController.hand its
trtGptModelInflightBatchingwiring) used by the TensorRT backend.max_tokens_in_buffer/kv_transfer_poll_interval_msfields onCacheTransceiverConfig, so this PR carries no public API change.Test Coverage
This PR only removes code, so coverage is the existing suites minus the tests
for the removed paths:
tests/unittest/_torch/executor/test_py_executor.py— removedTestDisaggTransferAdmissionController(8 tests) andTestDisaggTransferAdmissionPP(2 tests). InTestDisaggTransferIdleProgress,the two tests covering the generation-side branch are removed; the two
context-side tests are kept and updated for the new signature, so the
remaining idle-progress behavior stays covered.
tests/unittest/_torch/executor/test_scheduler_serializable_output.py—SerializableSchedulerOutputpickle round-trip still covered, minus theremoved flag.
(
tests/integration/defs/accuracy/test_disaggregated_serving.py) exercise theend-to-end path this code sat on.
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
wait_for_disagg_gen_transfer_progressfrom scheduler serialization.QA Engineer Review
tests/unittest/_torch/executor/test_py_executor.pytests/unittest/_torch/executor/test_scheduler_serializable_output.pytests/unittest/_torch/executor/test_benchmark_disagg.py