[TRTLLM-14660][feat] Early reused cache transfer in Transceiver V2

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.

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

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

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,13 +1147,15 @@ 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(
         prepopulated_tokens=prepopulated_blocks * _REUSE_TPB,
         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",
     [

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