diff --git a/tensorrt_llm/_torch/disaggregation/transceiver.py b/tensorrt_llm/_torch/disaggregation/transceiver.py index c8bea047f730..c53e241893bc 100644 --- a/tensorrt_llm/_torch/disaggregation/transceiver.py +++ b/tensorrt_llm/_torch/disaggregation/transceiver.py @@ -750,11 +750,13 @@ def _build_prefill_chunk( chunk_start_pos, chunk_end_pos = req.py_last_context_chunk tpb = self._kv_cache_manager.tokens_per_block - # Include any reused prefix in the first transferred chunk. + # Include a reused prefix in the first regular chunk only when it was + # not already transferred before the context forward. is_first_chunk = chunk_start_pos == req.prepopulated_prompt_len + extends_to_prefix = is_first_chunk and not req.py_kv_prefix_sent is_last_chunk = req.context_remaining_length == 0 # Defer partial blocks except at the prompt's final chunk. - chunk_start = 0 if is_first_chunk else chunk_start_pos // tpb + chunk_start = 0 if extends_to_prefix else chunk_start_pos // tpb chunk_end = (chunk_end_pos + tpb - 1) // tpb if is_last_chunk else chunk_end_pos // tpb # The final chunk is sent even when it contains no complete block. diff --git a/tensorrt_llm/_torch/pyexecutor/llm_request.py b/tensorrt_llm/_torch/pyexecutor/llm_request.py index 93fc43fa106a..d32735b856cb 100644 --- a/tensorrt_llm/_torch/pyexecutor/llm_request.py +++ b/tensorrt_llm/_torch/pyexecutor/llm_request.py @@ -993,6 +993,10 @@ def __init__( self.py_kv_transfer_timed_out = False # Prevent recreation after a send session drops peer registration. self.py_kv_send_session_retired = False + # Set when the reuse-hit prefix blocks were transferred before the + # first context forward, so the first prefill chunk must not send them + # a second time. + self.py_kv_prefix_sent = False # 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 ad23219f7ebf..322956d2ce8b 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -4415,6 +4415,8 @@ def _executor_loop(self): if hasattr(self.drafter, "guided_decoder"): self.guided_decoder.rollback_draft_tokens() + self._send_kv_cache_early(scheduled_batch.context_requests) + scheduled_batch_stats = ( self._collect_scheduled_batch_stats(scheduled_batch) if self.enable_iter_perf_stats else None) @@ -5221,6 +5223,7 @@ def _executor_loop_overlap(self): # and let the later iteration to update it. should_process_previous_batch = can_queue or not can_queue_this_rank if can_queue: + self._send_kv_cache_early(scheduled_batch.context_requests) # The generation requests that do not have batch_idx # need to be in front of the batch due to the assumptions @@ -7640,6 +7643,51 @@ def _recv_disagg_gen_cache(self, new_gen_reqs): return + @nvtx_range("_send_kv_cache_early") + def _send_kv_cache_early(self, + scheduled_requests: List[LlmRequest]) -> None: + """Send reused prefix blocks before the first context forward.""" + if (self.kv_cache_transceiver is None + or not self.kv_cache_transceiver.pipeline_transfer_enabled): + return + + has_offload_tier = (self.kv_cache_manager.can_evict + if self._is_kv_manager_v2 else + self.kv_cache_manager.blocks_in_secondary_pool > 0) + if has_offload_tier: + return + + tokens_per_block = self.kv_cache_manager.tokens_per_block + requests = [] + for req in scheduled_requests: + if (not req.is_context_only_request + or req.is_finished_due_to_cancellation): + continue + + if (not req.is_first_context_chunk + or req.prepopulated_prompt_len <= 0): + continue + + # prepopulated_prompt_len is not block-aligned in general: partial + # block reuse matches mid-block, and a full prefix hit reports + # prompt_len - 1 because at least one token must be recomputed. + # Round down so only fully populated blocks go out early; the block + # straddling the boundary is still being written by this forward and + # is covered by the first chunk, whose start floors to the same + # block. + prefix_end = (req.prepopulated_prompt_len // tokens_per_block * + tokens_per_block) + if prefix_end == 0: + continue + req.py_last_context_chunk = (0, prefix_end) + req.py_kv_prefix_sent = True + requests.append(req) + + if not requests: + return + + self._send_kv_async(requests) + @nvtx_range("_send_kv_async") def _send_kv_async(self, scheduled_requests: List[LlmRequest]): # Order matters: reaping before the connector registers its transfer diff --git a/tests/unittest/disaggregated/test_chunked_transfer.py b/tests/unittest/disaggregated/test_chunked_transfer.py index 7cd2bad8b94b..2594dcd405c6 100644 --- a/tests/unittest/disaggregated/test_chunked_transfer.py +++ b/tests/unittest/disaggregated/test_chunked_transfer.py @@ -779,6 +779,98 @@ def test_pipelined_transfer_allows_generation_only_request(): executor.sampler.validate_request.assert_called_once_with(request) +def test_send_kv_cache_early_only_sends_reused_prefixes(): + from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor + + executor = MagicMock() + executor.kv_cache_transceiver.pipeline_transfer_enabled = True + executor._is_kv_manager_v2 = False + executor.kv_cache_manager.blocks_in_secondary_pool = 0 + executor.kv_cache_manager.tokens_per_block = 32 + + def make_request( + rid, *, is_first_chunk, prepopulated, cancelled=False, last_chunk=(None, None) + ): + return SimpleNamespace( + py_request_id=rid, + is_context_only_request=True, + is_finished_due_to_cancellation=cancelled, + is_first_context_chunk=is_first_chunk, + prepopulated_prompt_len=prepopulated, + py_last_context_chunk=last_chunk, + py_kv_prefix_sent=False, + ) + + completed = make_request(1, is_first_chunk=False, prepopulated=0, last_chunk=(0, 64)) + first_chunk = make_request(2, is_first_chunk=True, prepopulated=0) + reused_prefix = make_request(3, is_first_chunk=True, prepopulated=128) + cancelled = make_request(4, is_first_chunk=True, prepopulated=128, cancelled=True) + + result = PyExecutor._send_kv_cache_early( + executor, [completed, first_chunk, reused_prefix, cancelled] + ) + + executor._send_kv_async.assert_called_once_with([reused_prefix]) + assert reused_prefix.py_last_context_chunk == (0, 128) + assert reused_prefix.py_kv_prefix_sent + assert not first_chunk.py_kv_prefix_sent + assert not cancelled.py_kv_prefix_sent + assert result is None + + # A partial-block or full-prefix hit reports an unaligned prepopulated + # length; only whole blocks may be shipped ahead of the forward. + executor._send_kv_async.reset_mock() + unaligned = make_request(5, is_first_chunk=True, prepopulated=3894) + below_one_block = make_request(6, is_first_chunk=True, prepopulated=31) + + PyExecutor._send_kv_cache_early(executor, [unaligned, below_one_block]) + + executor._send_kv_async.assert_called_once_with([unaligned]) + assert unaligned.py_last_context_chunk == (0, 3872) + assert below_one_block.py_last_context_chunk == (None, None) + assert not below_one_block.py_kv_prefix_sent + + +@pytest.mark.parametrize( + ("is_kv_manager_v2", "blocks_in_secondary_pool", "can_evict"), + [(False, 1, False), (True, 0, True)], +) +def test_send_kv_cache_early_skips_offload_tiers( + is_kv_manager_v2, blocks_in_secondary_pool, can_evict +): + from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor + + executor = MagicMock() + executor.kv_cache_transceiver.pipeline_transfer_enabled = True + executor._is_kv_manager_v2 = is_kv_manager_v2 + executor.kv_cache_manager.blocks_in_secondary_pool = blocks_in_secondary_pool + executor.kv_cache_manager.can_evict = can_evict + request = SimpleNamespace( + is_context_only_request=True, + is_finished_due_to_cancellation=False, + is_first_context_chunk=True, + prepopulated_prompt_len=128, + py_last_context_chunk=(None, None), + py_kv_prefix_sent=False, + ) + + PyExecutor._send_kv_cache_early(executor, [request]) + + executor._send_kv_async.assert_not_called() + assert request.py_last_context_chunk == (None, None) + assert not request.py_kv_prefix_sent + + +def test_send_kv_cache_early_requires_pipelined_transfer(): + from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor + + executor = MagicMock() + executor.kv_cache_transceiver.pipeline_transfer_enabled = False + + assert PyExecutor._send_kv_cache_early(executor, []) is None + executor._send_kv_async.assert_not_called() + + def test_pipelined_last_chunk_sends_and_finalizes(): """respond_and_send_async sends the built chunk and finalizes on the last chunk.""" from tensorrt_llm._torch.disaggregation.transceiver import KvCacheTransceiverV2 @@ -952,6 +1044,7 @@ def test_pipelined_multiple_chunks_use_real_builder_and_tx_session(): py_beam_width=1, py_kv_send_session_retired=False, prepopulated_prompt_len=0, + py_kv_prefix_sent=False, is_generation_only_request=lambda: False, set_kv_cache_transfer_start=lambda _ts: None, state=LlmRequestState.CONTEXT_INIT, @@ -1001,6 +1094,7 @@ def _build_prefill_chunk_tokens_for( resident_blocks=None, sliding_window_size=_REUSE_TOTAL_BLOCKS * _REUSE_TPB, source_block_ids=None, + prefix_sent=False, ): """Drive the real _build_prefill_chunk for one chunk, in token coordinates. @@ -1013,6 +1107,8 @@ def _build_prefill_chunk_tokens_for( out, so a full-attention group carries max_attention_window. V2 spells the same thing as None. Passing a window shorter than the prompt makes the group genuinely windowed. + ``prefix_sent`` models the executor having already shipped the reused prefix + early. """ from tensorrt_llm._torch.disaggregation.transceiver import KvCacheTransceiverV2 @@ -1039,6 +1135,7 @@ def _build_prefill_chunk_tokens_for( req.py_beam_width = 1 req.prompt_len = _REUSE_TOTAL_BLOCKS * _REUSE_TPB req.prepopulated_prompt_len = prepopulated_tokens + req.py_kv_prefix_sent = prefix_sent req.py_last_context_chunk = (chunk_start_pos, chunk_end_pos) req.context_remaining_length = req.prompt_len - chunk_end_pos @@ -1050,6 +1147,7 @@ def _build_prefill_chunk_for( chunk_start_block, chunk_end_block, resident_blocks=None, + prefix_sent=False, ): """Drive the real _build_prefill_chunk for one block-aligned chunk.""" return _build_prefill_chunk_tokens_for( @@ -1057,6 +1155,7 @@ def _build_prefill_chunk_for( chunk_start_pos=chunk_start_block * _REUSE_TPB, chunk_end_pos=chunk_end_block * _REUSE_TPB, resident_blocks=chunk_end_block if resident_blocks is None else resident_blocks, + prefix_sent=prefix_sent, ) @@ -1191,6 +1290,20 @@ def test_first_chunk_covers_ctx_prefix_reuse(): assert kv_slice.is_last_slice is False +def test_first_chunk_skips_prefix_already_sent_early(): + """An early prefix send owns blocks [0, 3), so slice 0 starts at its own block.""" + kv_slice = _build_prefill_chunk_for( + prepopulated_blocks=3, + chunk_start_block=3, + chunk_end_block=6, + resident_blocks=6, + prefix_sent=True, + ) + + assert kv_slice.token_range == _reuse_token_range(3, 6) + assert np.array_equal(kv_slice.block_ids_per_layer_groups[0], np.arange(3, 6, dtype=np.int64)) + + @pytest.mark.parametrize( "prepopulated_blocks,chunk_start_block,chunk_end_block,expected_start_block", [ diff --git a/tests/unittest/disaggregated/test_kv_transfer.py b/tests/unittest/disaggregated/test_kv_transfer.py index 81ced03755da..5e343f996a49 100644 --- a/tests/unittest/disaggregated/test_kv_transfer.py +++ b/tests/unittest/disaggregated/test_kv_transfer.py @@ -226,6 +226,7 @@ def _send_prefill_chunks( req.py_beam_width = 1 # Set explicitly so MagicMock does not bypass first-chunk detection. req.prepopulated_prompt_len = prepopulated_blocks * tokens_per_block + req.py_kv_prefix_sent = False first_block = min(prepopulated_blocks, total_blocks) if chunk_size_blocks is None or chunk_size_blocks >= total_blocks - first_block: @@ -291,6 +292,7 @@ def test_build_prefill_chunk_slices_chunk_window_from_whole_prompt(): req.prompt_len = prompt_blocks * tokens_per_block req.py_beam_width = 1 req.prepopulated_prompt_len = 0 + req.py_kv_prefix_sent = False for chunk_idx in range(2): chunk_start = chunk_idx * chunk_blocks @@ -348,6 +350,7 @@ def test_build_prefill_chunk_defers_partial_swa_chunk(source_block_ids): req.prompt_len = prompt_blocks * tokens_per_block req.py_beam_width = 1 req.prepopulated_prompt_len = 0 + req.py_kv_prefix_sent = False req.py_last_context_chunk = (11 * tokens_per_block, 13 * tokens_per_block) req.context_remaining_length = 3 * tokens_per_block req.is_generation_only_request.return_value = False