From f1f0cb514171c77af254baf2ce3f72258bb7ac67 Mon Sep 17 00:00:00 2001 From: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com> Date: Tue, 18 Aug 2026 13:39:52 -0700 Subject: [PATCH 1/4] [nvbugs/6621362][fix] Unblock disagg KV receivers on two idle/abort paths Two independent defects on the disaggregated KV transfer path, both found while investigating the conc512 stress case. Each is fixed at its own error site; neither is a workaround. 1. An idle generation worker with an in-flight receive routed its wait to the context path, which is a no-op for a receive-only worker: check_context_transfer_status() returns on `not _ever_had_send_session` *above* its poll, so it never sleeps. The worker spun the scheduler loop and starved the GIL from the transfer threads that alone complete the receive it was waiting on. An in-flight receive now also votes for the generation wait, which does reach its poll interval. Previously this was reachable only under admission-budget pressure. 2. Of the three failure exits in Sender._deliver_kv_to_agent, the `session is None` exit was the only one that did not send a FAILED last-slice result. A session can be deregistered (cancel_request or the context transfer timeout) while its slice is still queued, so the peer's RX task future stayed unresolved for the full kv_transfer_timeout_ms with its KV pages pinned. All three exits now share _abort_receiver_slice(). The shared helper uses the thread-local DEALER cache: it runs on _process_task_queue worker threads, while self._dealers is unsynchronized and documented as listener-thread-only (the success path already does this). It tolerates a send failure like its sibling _send_failed_result_to_receiver, since the local task is already failed and a dead peer must not become a second, unhandled failure. The reported gate is not fixed by this change and the waiver is left in place: the residual is a prefill capacity ceiling, not a defect. This case is the only ctxtp1/gentp1 stress entry at concurrency 512, and the 5% aiperf gate it fails was itself added after the bug's PASSED commit. Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com> --- .../_torch/disaggregation/native/transfer.py | 52 ++++++++++++------- tensorrt_llm/_torch/pyexecutor/py_executor.py | 17 ++++-- .../_torch/executor/test_py_executor.py | 33 ++++++++++++ tests/unittest/disaggregated/test_bounce.py | 39 ++++++++++++++ 4 files changed, 119 insertions(+), 22 deletions(-) diff --git a/tensorrt_llm/_torch/disaggregation/native/transfer.py b/tensorrt_llm/_torch/disaggregation/native/transfer.py index 672887b82ac6..22939612d8ac 100644 --- a/tensorrt_llm/_torch/disaggregation/native/transfer.py +++ b/tensorrt_llm/_torch/disaggregation/native/transfer.py @@ -537,6 +537,7 @@ def _deliver_kv_to_agent(self, write_meta: WriteMeta): assert write_meta.src_ptrs.size == write_meta.dst_ptrs.size == write_meta.sizes.size, ( f"WriteMeta ptr/size mismatch for unique_rid={write_meta.unique_rid}" ) + assert write_meta.slice_id is not None with self._sessions_lock: session = self._get_session(write_meta.unique_rid) @@ -546,8 +547,12 @@ def _deliver_kv_to_agent(self, write_meta: WriteMeta): ) logger.error(msg) write_meta.task.fail(RuntimeError(msg)) + # The session can be deregistered (cancel_request/ctx timeout) while + # this slice is still queued. Without a result frame the peer's RX + # task stays unresolved for the whole kv_transfer_timeout_ms with its + # KV pages pinned, so abort it here as the sibling exits below do. + self._abort_receiver_slice(write_meta) return - assert write_meta.slice_id is not None task = session.kv_tasks[write_meta.slice_id] timer = task._perf_timer if timer: @@ -573,15 +578,7 @@ def _deliver_kv_to_agent(self, write_meta: WriteMeta): task.fail( RuntimeError(f"session {write_meta.unique_rid} {status.value}, transfer aborted") ) - self._get_or_connect_dealer(write_meta.peer_endpoint).send( - _make_kv_result_msg( - self._instance_rank, - write_meta.unique_rid, - write_meta.slice_id, - True, # is_last_slice — ensures receiver resolves its task future - AgentResult.FAILED, - ) - ) + self._abort_receiver_slice(write_meta) return from .bounce import build_send_request, encode_result_tail @@ -603,15 +600,7 @@ def _deliver_kv_to_agent(self, write_meta: WriteMeta): f"{write_meta.unique_rid} slice={write_meta.slice_id}: {e}" ) task.fail(RuntimeError(f"build_send_request failed: {e}")) - self._get_or_connect_dealer(write_meta.peer_endpoint).send( - _make_kv_result_msg( - self._instance_rank, - write_meta.unique_rid, - write_meta.slice_id, - True, # is_last_slice — ensures receiver resolves its task future - AgentResult.FAILED, - ) - ) + self._abort_receiver_slice(write_meta) return if timer: timer.record_transfer_start(write_meta.peer_rank) @@ -1131,6 +1120,31 @@ def _respond_with_kv(self, _send_id: bytes, message: list[bytes]): task._perf_timer.record_push_start(trans_meta.peer_rank) self._enqueue(trans_meta) + def _abort_receiver_slice(self, write_meta: WriteMeta): + """Tell the peer this slice failed so it resolves its RX task future now. + + Called from _deliver_kv_to_agent on a _process_task_queue worker thread, + hence the thread-local DEALER cache: self._dealers is unsynchronized and + listener-thread-only. A send failure is swallowed like in + _send_failed_result_to_receiver — the local task is already failed and a + dead peer must not turn into a second, unhandled failure. + """ + try: + self._get_or_connect_thread_dealer(write_meta.peer_endpoint).send( + _make_kv_result_msg( + self._instance_rank, + write_meta.unique_rid, + write_meta.slice_id, + True, # is_last_slice — ensures receiver resolves its task future + AgentResult.FAILED, + ) + ) + except Exception as e: + logger.warning( + f"_deliver_kv_to_agent: failed to abort receiver slice for " + f"rid={write_meta.unique_rid} slice={write_meta.slice_id}: {e}" + ) + def _send_failed_result_to_receiver(self, info: RecvReqInfo): try: peer_ri = self._registrar.get_peer_rank_info(info.instance_name, info.instance_rank) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 19ab0cffc873..6b644bb9c719 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -3626,8 +3626,18 @@ def _check_disagg_transfer_progress_when_idle( and not self._is_disagg_gen_only_no_context_benchmark()): return - local_need_gen_check = (uses_async_gen_transfer and local_needs_progress - and wait_for_disagg_gen_transfer_progress) + # An in-flight receive must vote too, not just admission-budget + # pressure. The context wait is a no-op for a receive-only worker: + # check_context_transfer_status() returns on `not + # _ever_had_send_session` *above* its poll, so it never sleeps. Such a + # worker would spin the scheduler loop and starve the GIL from the + # transfer threads that alone finish the receive it is waiting on, + # whereas the generation wait does reach its poll interval. + local_need_gen_check = ( + uses_async_gen_transfer and local_needs_progress + and (wait_for_disagg_gen_transfer_progress + or any(req.is_disagg_generation_transmission_in_progress + for req in self.active_requests))) any_need_gen_check = self._sync_disagg_gen_status_entry( local_need_gen_check) @@ -3635,7 +3645,8 @@ def _check_disagg_transfer_progress_when_idle( if local_need_gen_check: logger.debug( "Waiting for generation KV cache transfer progress to " - "free disagg admission budget") + "complete an in-flight receive or free disagg admission " + "budget") self._check_disagg_gen_cache_transfer_status(1) return diff --git a/tests/unittest/_torch/executor/test_py_executor.py b/tests/unittest/_torch/executor/test_py_executor.py index 8bcc11a1f120..ff90745e1fe7 100644 --- a/tests/unittest/_torch/executor/test_py_executor.py +++ b/tests/unittest/_torch/executor/test_py_executor.py @@ -971,6 +971,7 @@ def test_gen_transfer_status_skips_sync_mode(self, monkeypatch): def test_polls_generation_transfer_when_admission_blocked(self): executor = object.__new__(PyExecutor) + executor.active_requests = [] executor.dist = Mock(tp_size=1) executor._check_disagg_gen_cache_transfer_status = Mock() executor._check_disagg_ctx_cache_transfer_status = Mock() @@ -988,6 +989,7 @@ def test_polls_generation_transfer_when_admission_blocked(self): def test_peer_rank_enters_bounded_progress_poll(self): executor = object.__new__(PyExecutor) + executor.active_requests = [] executor.dist = Mock(tp_size=1, cp_size=4, world_size=4) executor.dist.allreduce.return_value = 1 executor._check_disagg_gen_cache_transfer_status = Mock() @@ -1007,6 +1009,7 @@ def test_peer_rank_enters_bounded_progress_poll(self): def test_falls_back_to_context_transfer_when_not_generation_blocked(self): executor = object.__new__(PyExecutor) + executor.active_requests = [] executor.dist = Mock(tp_size=1) executor._check_disagg_gen_cache_transfer_status = Mock() executor._check_disagg_ctx_cache_transfer_status = Mock() @@ -1022,11 +1025,40 @@ def test_falls_back_to_context_transfer_when_not_generation_blocked(self): executor._check_disagg_ctx_cache_transfer_status.assert_called_once_with(1) executor._check_disagg_gen_cache_transfer_status.assert_not_called() + @pytest.mark.parametrize("in_flight", [True, False]) + def test_inflight_gen_receive_routes_idle_wait_to_generation(self, in_flight): + """An in-flight receive alone must select the generation wait. + + The context wait never sleeps for a receive-only worker, so routing + there would spin and starve the transfer threads finishing the receive. + """ + executor = object.__new__(PyExecutor) + executor.active_requests = [_make_disagg_transfer_request(1, 32, in_progress=in_flight)] + 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, + ) + + if in_flight: + executor._check_disagg_gen_cache_transfer_status.assert_called_once_with(1) + executor._check_disagg_ctx_cache_transfer_status.assert_not_called() + else: + executor._check_disagg_ctx_cache_transfer_status.assert_called_once_with(1) + executor._check_disagg_gen_cache_transfer_status.assert_not_called() + 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.active_requests = [] 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 @@ -1167,6 +1199,7 @@ def complete_or_error(req): def test_peer_cp_rank_enters_context_progress_poll(self): executor = object.__new__(PyExecutor) + executor.active_requests = [] 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] diff --git a/tests/unittest/disaggregated/test_bounce.py b/tests/unittest/disaggregated/test_bounce.py index 65cb12a60c21..0274a9fa0d15 100644 --- a/tests/unittest/disaggregated/test_bounce.py +++ b/tests/unittest/disaggregated/test_bounce.py @@ -19,6 +19,7 @@ """ import queue +import threading from types import SimpleNamespace import numpy as np @@ -220,6 +221,44 @@ def test_make_kv_result_msg_uses_binary_frame(result_name): assert tfr._AGENT_RESULT_BY_CODE[code] is result +def test_deliver_kv_aborts_receiver_when_session_is_gone(): + """A slice whose TxSession was deregistered must still FAIL the peer's slice. + + cancel_request()/ctx-timeout can drop the session while the slice is queued. + Returning without a result frame leaves the receiver's task future + unresolved for the whole kv_transfer_timeout_ms with its KV pages pinned. + """ + tfr = pytest.importorskip("tensorrt_llm._torch.disaggregation.native.transfer") + + sent = [] + failures = [] + write_meta = SimpleNamespace( + unique_rid=4242, + slice_id=3, + task=SimpleNamespace(fail=failures.append), + peer_endpoint="tcp://peer:1234", + src_ptrs=SimpleNamespace(size=1), + dst_ptrs=SimpleNamespace(size=1), + sizes=SimpleNamespace(size=1), + ) + + sender = object.__new__(tfr.Sender) + sender._sessions = {} # the real _get_session() then reports the session gone + sender._sessions_lock = threading.Lock() + sender._instance_rank = 1 + sender._shutdown = True # nothing real for __del__ to reap + sender._thread_local = threading.local() + sender._thread_local.dealers = {write_meta.peer_endpoint: SimpleNamespace(send=sent.append)} + + tfr.Sender._deliver_kv_to_agent(sender, write_meta) + + assert [type(e) for e in failures] == [RuntimeError] + assert len(sent) == 1, "receiver was never told the slice failed" + _, rid, slice_id, is_last, code, _ = tfr._KV_RESULT_PREFIX.unpack(sent[0][1]) + assert (rid, slice_id, is_last) == (4242, 3, True) + assert tfr._AGENT_RESULT_BY_CODE[code] is tfr.AgentResult.FAILED + + # --------------------------------------------------------------------------- # # fan-in safety gate — equal total//num_writers split only for uniform writers # --------------------------------------------------------------------------- # From 832a44fdc5be8089657495c40fcd06a609a01ef5 Mon Sep 17 00:00:00 2001 From: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com> Date: Tue, 18 Aug 2026 19:21:03 -0700 Subject: [PATCH 2/4] [nvbugs/6621362][fix] Size the qwen3.5 disagg transfer budget to its batch DisaggTransferAdmissionController derives its budget from max_tokens_in_buffer, but the two are different quantities. max_tokens_in_buffer is a *per-buffer* arena size -- cacheTransBuffer.cpp:371-374 computes preAllocBufferSize = transferBufferSize * (recvBufferCount + sendBufferCount), so the field sizes ONE buffer and concurrency is a separate multiplicand (1..3 by default). The controller instead divides it by tokens_per_block and spends the quotient as an aggregate pool across all in-flight requests. With this config's 16384 that pool is 16384/32 = 512 blocks, while one 8K-ISL request needs 8192/32 = 256, so exactly 2 of the 512 concurrent requests this test offers can have a generation transfer in flight. The remaining ones queue until the router gives up at req_timeout_secs=180, which is the observed failure: 2362 requests returning code=500 with a blank detail body. Size the budget to the batch the same file already declares (max_batch_size=128 * 8K ISL = 1048576), giving 1048576/32/256 = 128 concurrent transfers. The gate stays enforcing -- it now admits the declared batch instead of 2 -- and the capacity scheduler still applies the real KV limits before admission runs. Raising the value cannot over-allocate any staging arena on this path: tensorrt_llm/_torch/disaggregation/ has zero references to max_tokens_in_buffer, and Qwen3_5.get_preferred_transceiver_runtime() returns "PYTHON", so no C++ transfer buffer is sized from this field here. Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com> --- ...agg_config_ctxtp1_gentp1_qwen3_5_4b_fp8_tllm.yaml | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp1_gentp1_qwen3_5_4b_fp8_tllm.yaml b/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp1_gentp1_qwen3_5_4b_fp8_tllm.yaml index aaf757903d9e..96f3c9730c51 100644 --- a/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp1_gentp1_qwen3_5_4b_fp8_tllm.yaml +++ b/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp1_gentp1_qwen3_5_4b_fp8_tllm.yaml @@ -18,7 +18,13 @@ context_servers: print_iter_log: true cache_transceiver_config: backend: DEFAULT - max_tokens_in_buffer: 16384 + # DisaggTransferAdmissionController spends max_tokens_in_buffer / + # tokens_per_block as an aggregate budget for concurrent generation KV + # transfers, so an arena-sized value (~one sequence) admits ~2 of the 512 + # requests this test offers at 8K ISL and the rest die on the router's 180s + # timeout. Sized to the batch declared above (max_batch_size * 8K ISL); + # this model runs the Python transceiver, which allocates no such arena. + max_tokens_in_buffer: 1048576 generation_servers: num_instances: 1 tensor_parallel_size: 1 @@ -47,4 +53,6 @@ generation_servers: print_iter_log: true cache_transceiver_config: backend: DEFAULT - max_tokens_in_buffer: 16384 + # Kept in step with the context server above: the generation side is the one + # whose admission gate defers the transfers. + max_tokens_in_buffer: 1048576 From 6c90a7400a985ba3b542226dadbca44603d32e92 Mon Sep 17 00:00:00 2001 From: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com> Date: Tue, 18 Aug 2026 19:26:52 -0700 Subject: [PATCH 3/4] [nvbugs/6621362][fix] Charge the context KV transfer timeout to the transfer only The context-side deadline is stamped right after respond_and_send_async(), which only creates the send session. The KV write cannot start until the generation peer asks for the data, and the session stays in SessionStatus.INIT until every peer rank's request info arrives. At stress concurrency the peer is late for a structural reason -- generation max_batch_size=128 against concurrency=512 -- so tens of seconds of generation-slot queue wait are charged to the transfer. Measured on the qwen3.5 8K/512 stress run, bucketing the reported elapsed times: 627 of 1241 context timeouts fire in the 60-70s bucket, i.e. they trip the deadline the first time it is checked, having spent the entire budget before any bytes could move. Each one fails a request that never got to transfer. Rebase the deadline while the send session is still waiting for its peer, via a new context_transfer_is_waiting_for_peer() hook. It defaults to False on the base transceiver, so the C++ runtime is unaffected; the Python V2 transceiver reports the INIT boundary it already tracks for scheduling. Bound the rebase rather than resetting unconditionally. This deadline is the only path that ends a context transfer whose peer never asks -- check_context_transfer_status sees WaitResult.TIMEOUT but deliberately keeps the request in progress, the session stays in INIT so _collect_done never returns it, and _try_cancel_request declines while the request is still in requests_in_transfer -- so an unbounded rebase would pin its KV pages for the process lifetime. py_kv_transfer_peer_wait_start records the original stamp and is never rebased; once the total peer wait exceeds 3x kv_transfer_timeout_ms the deadline applies again. That ceiling covers the measured peer-wait spread (max 182.7s) and lands at the disaggregated router's own req_timeout_secs=180 default, past which no peer will ask. A transfer that has actually started still times out, and a peer that never asks still expires -- both covered by the accompanying tests, one of which fails if the ceiling is removed. Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com> --- .../_torch/disaggregation/transceiver.py | 5 ++ .../_torch/pyexecutor/kv_cache_transceiver.py | 11 +++ tensorrt_llm/_torch/pyexecutor/llm_request.py | 3 + tensorrt_llm/_torch/pyexecutor/py_executor.py | 41 ++++++++- .../_torch/executor/test_py_executor.py | 83 +++++++++++++++++++ 5 files changed, 142 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/disaggregation/transceiver.py b/tensorrt_llm/_torch/disaggregation/transceiver.py index 45bc4b5e5a56..a8e0517b4390 100644 --- a/tensorrt_llm/_torch/disaggregation/transceiver.py +++ b/tensorrt_llm/_torch/disaggregation/transceiver.py @@ -965,6 +965,11 @@ def _assert_disagg_history_declared(self, req: LlmRequest) -> None: f"_try_schedule_disagg_gen_init." ) + def context_transfer_is_waiting_for_peer(self, req: LlmRequest) -> bool: + # The send session stays in SessionStatus.INIT until every peer rank's + # request info has arrived; only then can any KV be written. + return not self._transfer_worker.has_all_peer_req_infos_for_send(get_unique_rid(req)) + def cancel_request(self, req: LlmRequest) -> bool: """Cancel the transfer for the given request. diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py index 65f613337fff..ddaf647dd8de 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py @@ -263,6 +263,17 @@ def cancel_request(self, req: LlmRequest): def supports_inflight_request_cancellation(self) -> bool: return False + def context_transfer_is_waiting_for_peer(self, req: LlmRequest) -> bool: + """Whether a context send is still waiting for its peer to ask for the data. + + respond_and_send_async() only creates the send session; the KV write + cannot start until the generation peer requests it. Runtimes that can + observe that boundary report it here so the transfer timeout measures + the transfer instead of the peer's queueing delay. Default False keeps + the timeout unchanged for runtimes that cannot distinguish the two. + """ + return False + def has_poisoned_transfer_buffer(self) -> bool: return False diff --git a/tensorrt_llm/_torch/pyexecutor/llm_request.py b/tensorrt_llm/_torch/pyexecutor/llm_request.py index 041c4933542a..08a4ed780b92 100644 --- a/tensorrt_llm/_torch/pyexecutor/llm_request.py +++ b/tensorrt_llm/_torch/pyexecutor/llm_request.py @@ -991,6 +991,9 @@ def __init__( self.is_cuda_graph_dummy = False self.py_kv_transfer_start_time = None self.py_kv_transfer_timed_out = False + # Set alongside py_kv_transfer_start_time for a context send and never + # rebased, so waiting for a peer that never asks stays bounded. + self.py_kv_transfer_peer_wait_start = None # Encoder-decoder runtime state. ``py_encoder_output`` holds the # packed encoder hidden states produced by the encoder iteration as diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 6b644bb9c719..611e641d8370 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -6724,6 +6724,36 @@ def _check_gen_cache_transfer_errors_consensus(self) -> None: requests=error_requests, charge_budget=False) + # How much of kv_transfer_timeout_ms a context send may spend waiting for + # its generation peer to ask for the data before the deadline stops being + # rebased. The peer arrives late for a structural reason (generation + # max_batch_size below the offered concurrency), so the wait must be + # tolerated -- but a peer that never asks has to stay reclaimable, since + # this deadline is the only path that ends such a transfer. 3x covers the + # measured peer-wait spread on the 8K/512 stress run (max 182.7s) and lands + # at the disaggregated router's own req_timeout_secs=180 default, past + # which no peer will ask. + _CTX_PEER_WAIT_TIMEOUT_MULTIPLIER = 3 + + def _context_transfer_peer_wait_is_within_ceiling(self, req: LlmRequest, + current_time: float, + timeout_ms: int) -> bool: + """Whether a context send may keep rebasing its transfer deadline. + + The clock is stamped when the send session is created, but the write + only starts once the peer requests the data, so charging the peer wait + to the transfer times out transfers that never got to run. Rebasing is + bounded by the ceiling above so a peer that never asks still expires. + """ + if not self.kv_cache_transceiver.context_transfer_is_waiting_for_peer( + req): + return False + if req.py_kv_transfer_peer_wait_start is None: + return False + peer_wait_ms = (current_time - + req.py_kv_transfer_peer_wait_start) * 1000 + return peer_wait_ms <= timeout_ms * self._CTX_PEER_WAIT_TIMEOUT_MULTIPLIER + @nvtx_range("_check_kv_transfer_timeout") def _check_kv_transfer_timeout(self): if not self.kv_cache_transceiver: @@ -6732,8 +6762,9 @@ def _check_kv_transfer_timeout(self): if timeout_ms is None: return + current_time = time.monotonic() + def flag_if_kv_transfer_timed_out(req: LlmRequest, type: str) -> None: - current_time = time.monotonic() if req.py_kv_transfer_start_time is None: return elapsed_time = (current_time - req.py_kv_transfer_start_time) * 1000 @@ -6748,6 +6779,10 @@ def flag_if_kv_transfer_timed_out(req: LlmRequest, type: str) -> None: req.py_kv_transfer_timed_out = True for req in self.async_transfer_manager.requests_in_transfer().values(): + if self._context_transfer_peer_wait_is_within_ceiling( + req, current_time, timeout_ms): + req.py_kv_transfer_start_time = current_time + continue flag_if_kv_transfer_timed_out(req, "context") for req in self.active_requests: @@ -7213,6 +7248,7 @@ def _prepare_disagg_gen_transmission_complete(self, scheduled_batch): req.decoding_iter = 1 req.py_decoding_iter = 1 req.py_kv_transfer_start_time = None + req.py_kv_transfer_peer_wait_start = None req.py_kv_transfer_timed_out = False first_gen_tokens = req.context_phase_params.first_gen_tokens ctx_draft_tokens = req.context_phase_params.draft_tokens @@ -7458,6 +7494,8 @@ def kv_connector_request_finished(req: LlmRequest): if self.kv_cache_transceiver.kv_transfer_timeout_ms is not None: req.py_kv_transfer_start_time = time.monotonic() + req.py_kv_transfer_peer_wait_start = ( + req.py_kv_transfer_start_time) if self.kv_connector_manager: if not self.disable_overlap_scheduler: @@ -7557,6 +7595,7 @@ def _check_disagg_ctx_cache_transfer_status(self, atLeastNum: int = 0): # cancellation is disabled: a queued transfer that can be # cancelled is immediately released from the async manager. request.py_kv_transfer_start_time = None + request.py_kv_transfer_peer_wait_start = None request.state = LlmRequestState.DISAGG_CONTEXT_COMPLETE self._end_transfer_and_maybe_terminate(request) diff --git a/tests/unittest/_torch/executor/test_py_executor.py b/tests/unittest/_torch/executor/test_py_executor.py index ff90745e1fe7..6736c2a46b36 100644 --- a/tests/unittest/_torch/executor/test_py_executor.py +++ b/tests/unittest/_torch/executor/test_py_executor.py @@ -1219,6 +1219,89 @@ def test_peer_cp_rank_enters_context_progress_poll(self): executor.dist.tp_cp_allgather.assert_called_once_with(0) +class TestContextKvTransferTimeoutClock: + """The context transfer deadline must measure the transfer, not the peer wait.""" + + TIMEOUT_MS = 60000 + + @classmethod + def _make_executor(cls, waiting_for_peer, elapsed_s, peer_wait_elapsed_s=None): + executor = object.__new__(PyExecutor) + started = time.monotonic() - elapsed_s + peer_wait_start = ( + started if peer_wait_elapsed_s is None else time.monotonic() - peer_wait_elapsed_s + ) + req = types.SimpleNamespace( + py_request_id=7, + py_kv_transfer_start_time=started, + py_kv_transfer_peer_wait_start=peer_wait_start, + py_kv_transfer_timed_out=False, + is_disagg_generation_transmission_in_progress=False, + ) + executor.kv_cache_transceiver = Mock( + kv_transfer_timeout_ms=cls.TIMEOUT_MS, + **{"context_transfer_is_waiting_for_peer.return_value": waiting_for_peer}, + ) + executor.async_transfer_manager = Mock(**{"requests_in_transfer.return_value": {7: req}}) + executor.active_requests = [] + executor._is_disagg_inflight_cancel_active = Mock(return_value=False) + return executor, req + + def test_peer_wait_does_not_trip_the_deadline(self): + executor, req = self._make_executor(waiting_for_peer=True, elapsed_s=120) + + PyExecutor._check_kv_transfer_timeout(executor) + + assert not req.py_kv_transfer_timed_out + # Rebased, so the next check starts from the moment the write can begin. + assert (time.monotonic() - req.py_kv_transfer_start_time) < 1 + + def test_stalled_transfer_still_trips_the_deadline(self): + """Rebasing must not become a way to never time out.""" + executor, req = self._make_executor(waiting_for_peer=False, elapsed_s=120) + + PyExecutor._check_kv_transfer_timeout(executor) + + assert req.py_kv_transfer_timed_out + + def test_transfer_inside_budget_is_not_flagged(self): + executor, req = self._make_executor(waiting_for_peer=False, elapsed_s=5) + + PyExecutor._check_kv_transfer_timeout(executor) + + assert not req.py_kv_transfer_timed_out + + def test_peer_that_never_asks_eventually_times_out(self): + """A peer that never sends REQUEST_DATA must stay reclaimable. + + This deadline is the only path that ends such a transfer, so an + unbounded rebase would pin its KV pages for the process lifetime. + """ + ceiling_ms = self.TIMEOUT_MS * PyExecutor._CTX_PEER_WAIT_TIMEOUT_MULTIPLIER + executor, req = self._make_executor( + waiting_for_peer=True, + elapsed_s=120, + peer_wait_elapsed_s=ceiling_ms / 1000 + 1, + ) + + PyExecutor._check_kv_transfer_timeout(executor) + + assert req.py_kv_transfer_timed_out + + def test_peer_wait_rebases_repeatedly_below_the_ceiling(self): + """Repeated checks while waiting must not accumulate toward the deadline.""" + executor, req = self._make_executor( + waiting_for_peer=True, elapsed_s=120, peer_wait_elapsed_s=30 + ) + + for _ in range(5): + PyExecutor._check_kv_transfer_timeout(executor) + assert not req.py_kv_transfer_timed_out + + # The peer-wait origin is never rebased, so the ceiling still applies. + assert (time.monotonic() - req.py_kv_transfer_peer_wait_start) >= 30 + + @pytest.mark.usefixtures("_clear_disagg_transfer_mode_env") class TestDisaggTransferAdmissionPP: def test_pp_schedule_applies_gate_before_serializing(self): From 79f3dfdf1ea7cd711f1b86ea07ef170870a34895 Mon Sep 17 00:00:00 2001 From: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com> Date: Tue, 18 Aug 2026 19:27:28 -0700 Subject: [PATCH 4/4] [nvbugs/6621362][fix] Unwaive the qwen3.5 conc512 disagg stress test The gpt_oss_120b_eagle_triton_stress entry filed under the same bug id is a different configuration and stays waived. Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com> --- .../_torch/disaggregation/native/transfer.py | 2 +- tensorrt_llm/_torch/pyexecutor/py_executor.py | 29 ++++++++++--------- tests/integration/test_lists/waives.txt | 1 - 3 files changed, 17 insertions(+), 15 deletions(-) diff --git a/tensorrt_llm/_torch/disaggregation/native/transfer.py b/tensorrt_llm/_torch/disaggregation/native/transfer.py index 22939612d8ac..c8fba0b6ac86 100644 --- a/tensorrt_llm/_torch/disaggregation/native/transfer.py +++ b/tensorrt_llm/_torch/disaggregation/native/transfer.py @@ -1141,7 +1141,7 @@ def _abort_receiver_slice(self, write_meta: WriteMeta): ) except Exception as e: logger.warning( - f"_deliver_kv_to_agent: failed to abort receiver slice for " + f"_abort_receiver_slice: failed to abort receiver slice for " f"rid={write_meta.unique_rid} slice={write_meta.slice_id}: {e}" ) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 611e641d8370..ca1602eb5ca7 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -6735,24 +6735,26 @@ def _check_gen_cache_transfer_errors_consensus(self) -> None: # which no peer will ask. _CTX_PEER_WAIT_TIMEOUT_MULTIPLIER = 3 - def _context_transfer_peer_wait_is_within_ceiling(self, req: LlmRequest, - current_time: float, - timeout_ms: int) -> bool: + def _context_transfer_peer_wait_is_within_ceiling( + self, req: LlmRequest, current_time: float, + ceiling_ms: float) -> bool: """Whether a context send may keep rebasing its transfer deadline. The clock is stamped when the send session is created, but the write only starts once the peer requests the data, so charging the peer wait - to the transfer times out transfers that never got to run. Rebasing is - bounded by the ceiling above so a peer that never asks still expires. + to the transfer times out transfers that never got to run. The local + stamps are checked first: the transceiver query walks per-request peer + bookkeeping under a lock, and a request past the ceiling cannot be + rebased regardless of what it reports. """ - if not self.kv_cache_transceiver.context_transfer_is_waiting_for_peer( - req): - return False if req.py_kv_transfer_peer_wait_start is None: return False peer_wait_ms = (current_time - req.py_kv_transfer_peer_wait_start) * 1000 - return peer_wait_ms <= timeout_ms * self._CTX_PEER_WAIT_TIMEOUT_MULTIPLIER + if peer_wait_ms > ceiling_ms: + return False + return self.kv_cache_transceiver.context_transfer_is_waiting_for_peer( + req) @nvtx_range("_check_kv_transfer_timeout") def _check_kv_transfer_timeout(self): @@ -6763,6 +6765,7 @@ def _check_kv_transfer_timeout(self): return current_time = time.monotonic() + peer_wait_ceiling_ms = timeout_ms * self._CTX_PEER_WAIT_TIMEOUT_MULTIPLIER def flag_if_kv_transfer_timed_out(req: LlmRequest, type: str) -> None: if req.py_kv_transfer_start_time is None: @@ -6780,7 +6783,7 @@ def flag_if_kv_transfer_timed_out(req: LlmRequest, type: str) -> None: for req in self.async_transfer_manager.requests_in_transfer().values(): if self._context_transfer_peer_wait_is_within_ceiling( - req, current_time, timeout_ms): + req, current_time, peer_wait_ceiling_ms): req.py_kv_transfer_start_time = current_time continue flag_if_kv_transfer_timed_out(req, "context") @@ -7493,9 +7496,9 @@ def kv_connector_request_finished(req: LlmRequest): self.kv_cache_transceiver.respond_and_send_async(req) if self.kv_cache_transceiver.kv_transfer_timeout_ms is not None: - req.py_kv_transfer_start_time = time.monotonic() - req.py_kv_transfer_peer_wait_start = ( - req.py_kv_transfer_start_time) + transfer_start = time.monotonic() + req.py_kv_transfer_start_time = transfer_start + req.py_kv_transfer_peer_wait_start = transfer_start if self.kv_connector_manager: if not self.disable_overlap_scheduler: diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index b8535429c074..df9f50f53aaa 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -108,7 +108,6 @@ disaggregated/test_disaggregated.py::test_disaggregated_genbs1[TinyLlama-1.1B-Ch disaggregated/test_disaggregated.py::test_disaggregated_qwen3_32b_fp8[Qwen3/Qwen3-32B-FP8] SKIP (https://nvbugs/6566734) disaggregated/test_disaggregated.py::test_disaggregated_stress_test[input8k-output1k-conc512-deepseek_r1_v2_fp4_stress] SKIP (https://nvbugs/6621358) disaggregated/test_disaggregated.py::test_disaggregated_stress_test[input8k-output1k-conc512-gpt_oss_120b_eagle_triton_stress] SKIP (https://nvbugs/6621362) -disaggregated/test_disaggregated.py::test_disaggregated_stress_test[input8k-output1k-conc512-qwen3_5_4b_fp8_stress] SKIP (https://nvbugs/6621362) disaggregated/test_workers.py::test_workers_conversation_router[TinyLlama-1.1B-Chat-v1.0] SKIP (https://nvbugs/6162322) disaggregated/test_workers.py::test_workers_kv_cache_aware_router_deepseek_v3_lite_bf16[DeepSeek-V3-Lite-bf16] SKIP (https://nvbugs/6162322) disaggregated/test_workers.py::test_workers_kv_cache_aware_router_eviction[TinyLlama-1.1B-Chat-v1.0] SKIP (https://nvbugs/6162322)