Skip to content

[None][fix] Simplify idle disagg KV transfer progress check #17324

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Tabrizian wants to merge 3 commits into
base: main
Choose a base branch
from
+55 −140
Open

[None][fix] Simplify idle disagg KV transfer progress check #17324

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
5a98771
[None][fix] Simplify idle disagg KV transfer progress check
Tabrizian Aug 5, 2026
4a0ca6c
[None][fix] Drop redundant GEN transfer poll from the idle progress c…
Tabrizian Aug 7, 2026
f64c82d
[None][fix] Warn when the scheduler fits nothing on the idle disagg path
Tabrizian Aug 8, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
93 changes: 25 additions & 68 deletions tensorrt_llm/_torch/pyexecutor/py_executor.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2568,8 +2568,8 @@ def _executor_loop_pp(self):

# Stage 0: first PP rank schedules requests and propagates the result to all other PP ranks.
(scheduled_batch, fitting_disagg_gen_init_requests,
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)
if self.dist.rank != 0:
# Retry until current rank can run first PP's schedule result.
self._pp_retry_until_can_schedule(scheduled_batch)
Expand All @@ -2593,14 +2593,10 @@ def _executor_loop_pp(self):
self._prepare_disagg_gen_init(
fitting_disagg_gen_init_requests)

all_gen_first = self.active_requests and all(
req.py_disaggregated_params
and req.py_disaggregated_params.schedule_style ==
DisaggScheduleStyle.GENERATION_FIRST
for req in self.active_requests)
self._check_disagg_transfer_progress_when_idle(
num_fitting_reqs, fitting_disagg_gen_init_requests,
wait_for_disagg_gen_transfer_progress, all_gen_first)
if num_fitting_reqs == 0:
logger.warning(
"num_fitting_reqs=0, may not have enough kvCache")
self._check_disagg_transfer_progress_when_idle()

self.num_scheduled_requests = scheduled_batch.batch_size

Expand Down Expand Up @@ -3514,63 +3510,27 @@ def _allgather_model_parallel_status(
return self.dist.tp_allgather(local_status)
return [local_status]

def _sync_disagg_gen_status_entry(self, local_need_check: bool) -> int:
if self._dist_size(self.dist, "world_size") > 1:
return self.dist.allreduce(int(local_need_check), op=ReduceOp.MAX)
return int(local_need_check)
def _check_disagg_transfer_progress_when_idle(self) -> None:
"""Reap completed context KV transfers so their blocks can be freed.

def _sync_disagg_ctx_status_entry(self, local_need_check: bool) -> int:
if self._dist_size(self.dist, "cp_size") > 1:
return int(any(self.dist.tp_cp_allgather(int(local_need_check))))
if self._dist_size(self.dist, "tp_size") > 1:
return self.dist.tp_allreduce(int(local_need_check),
op=ReduceOp.MAX)
return int(local_need_check)

def _check_disagg_transfer_progress_when_idle(
self, num_fitting_reqs: int,
fitting_disagg_gen_init_requests: List[LlmRequest],
wait_for_disagg_gen_transfer_progress: bool,
all_gen_first: bool) -> None:
local_need_check = (num_fitting_reqs == 0
and not fitting_disagg_gen_init_requests)
The poll is non-blocking and rank-symmetric: every rank enters it
unconditionally on every disagg iteration, so the consensus performed
inside the status call stays aligned without an extra collective here.
Ranks with nothing in flight simply reap nothing.

Generation transfers are deliberately not polled here: the loop head
already ran `_check_disagg_gen_transfer_status` this iteration, and any
receive started since then by `_prepare_disagg_gen_init` is polled by
`_recv_disagg_gen_cache` right after it is issued. A poll here would
only repeat the GEN status call and its consensus.
"""
# 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

local_need_gen_check = (local_need_check
and wait_for_disagg_gen_transfer_progress)

any_need_gen_check = self._sync_disagg_gen_status_entry(
local_need_gen_check)
if any_need_gen_check > 0:
if local_need_gen_check:
logger.debug(
"Waiting for generation KV cache transfer progress to "
"free disagg admission budget")
self._check_disagg_gen_cache_transfer_status(1)
return

any_need_check = self._sync_disagg_ctx_status_entry(local_need_check)
if any_need_check > 0:
if local_need_check and not all_gen_first:
logger.warning(
"num_fitting_reqs=0 and fitting_disagg_gen_init_requests is empty, may not have enough kvCache"
)
# Local conditions warrant a blocking wait for at least one
# in-flight transfer to complete so KV blocks can be freed.
self._check_disagg_ctx_cache_transfer_status(1)
else:
# Either (a) a peer rank needed the call but we didn't, or
# (b) all active requests are gen-first so we don't
# actively block. In both cases the non-blocking variant
# still runs the internal allgather (keeping all ranks in
# 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)

def _sync_gen_only_benchmark_has_insufficient_kv(
self, scheduler_fitting_disagg_gen_init_requests: List[LlmRequest],
Expand Down Expand Up @@ -3715,13 +3675,10 @@ def _prepare_and_schedule_batch(self):
# into the transfer window this iteration.
self._prepare_disagg_gen_init(admitted_disagg_gen_init_requests)

all_gen_first = self.active_requests and all(
req.py_disaggregated_params and req.py_disaggregated_params.
schedule_style == DisaggScheduleStyle.GENERATION_FIRST
for req in self.active_requests)
self._check_disagg_transfer_progress_when_idle(
num_fitting_reqs, admitted_disagg_gen_init_requests,
wait_for_disagg_gen_transfer_progress, all_gen_first)
if num_fitting_reqs == 0:
logger.warning(
"num_fitting_reqs=0, may not have enough kvCache")
self._check_disagg_transfer_progress_when_idle()

# In gen-only benchmark mode, all requests must fit in KV cache
# simultaneously. If some requests are stuck in INIT state and the
Expand Down
7 changes: 2 additions & 5 deletions tests/unittest/_torch/executor/test_benchmark_disagg.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1168,9 +1168,7 @@ def test_partial_transfer_admission_uses_only_admitted_requests(self):
assert result is not None
ex._apply_disagg_transfer_admission.assert_called_once_with(candidates)
ex._prepare_disagg_gen_init.assert_called_once_with([admitted_req])
ex._check_disagg_transfer_progress_when_idle.assert_called_once_with(
0, [admitted_req], False, False
)
ex._check_disagg_transfer_progress_when_idle.assert_called_once_with()
ex._handle_errors.assert_not_called()

def test_fill_with_no_init_requests_does_not_kill(self):
Expand Down Expand Up @@ -1203,8 +1201,7 @@ def test_transfer_admission_backpressure_does_not_kill(self, monkeypatch):
)
ex._apply_disagg_transfer_admission.assert_called_once_with([fitting_req])
ex._prepare_disagg_gen_init.assert_called_once_with([])
ex._check_disagg_gen_cache_transfer_status.assert_called_once_with(1)
ex._check_disagg_ctx_cache_transfer_status.assert_not_called()
ex._check_disagg_ctx_cache_transfer_status.assert_called_once_with(0)
ex._handle_errors.assert_not_called()

@pytest.mark.parametrize(
Expand Down
95 changes: 28 additions & 67 deletions tests/unittest/_torch/executor/test_py_executor.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,6 @@

import pytest

from tensorrt_llm._torch.distributed.communicator import ReduceOp
from tensorrt_llm._torch.pyexecutor.executor_request_queue import (
SHUTDOWN_REQUEST_ID,
RequestQueueItem,
Expand Down Expand Up @@ -636,57 +635,51 @@ def test_gen_transfer_status_skips_sync_mode(self, monkeypatch):

executor._check_disagg_gen_cache_transfer_status.assert_not_called()

def test_polls_generation_transfer_when_admission_blocked(self):
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=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)

def test_gen_only_no_context_benchmark_skips_idle_polls(self, monkeypatch):
monkeypatch.setenv("TRTLLM_DISAGG_BENCHMARK_GEN_ONLY", "1")
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=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_ctx_cache_transfer_status.assert_not_called()
executor._check_disagg_gen_cache_transfer_status.assert_not_called()

def test_sync_benchmark_skips_idle_transfer_collectives(self, monkeypatch):
Expand All @@ -697,13 +690,7 @@ def test_sync_benchmark_skips_idle_transfer_collectives(self, monkeypatch):
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.dist.allreduce.assert_not_called()
executor.dist.tp_allreduce.assert_not_called()
Expand All @@ -718,13 +705,7 @@ def test_sync_non_benchmark_skips_idle_transfer_collectives(self, monkeypatch):
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.dist.allreduce.assert_not_called()
executor.dist.tp_allreduce.assert_not_called()
Expand Down Expand Up @@ -813,26 +794,6 @@ def complete_or_error(req):
charge_budget=False,
)

def test_peer_cp_rank_enters_context_progress_poll(self):
executor = object.__new__(PyExecutor)
executor.dist = Mock(tp_size=1, cp_size=4, world_size=4)
executor.dist.allreduce.return_value = 0
executor.dist.tp_cp_allgather.return_value = [0, 1, 0, 0]
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=False,
all_gen_first=False,
)

executor._check_disagg_ctx_cache_transfer_status.assert_called_once_with(0)
executor._check_disagg_gen_cache_transfer_status.assert_not_called()
executor.dist.tp_cp_allgather.assert_called_once_with(0)


@pytest.mark.usefixtures("_clear_disagg_transfer_mode_env")
class TestDisaggTransferAdmissionPP:
Expand Down
Loading