From b49f8a6cf18ae3ffdda79f66eba8cc9193792b00 Mon Sep 17 00:00:00 2001 From: Chuang Zhu <111838961+chuangz0@users.noreply.github.com> Date: Wed, 1 Jul 2026 08:47:49 +0000 Subject: [PATCH 1/2] [None][fix] align KV slice token_range.end with transferred block count _create_kv_slice set the default token_range.end to req.prompt_len + num_extra_kv_tokens, while the block list is trimmed to total_blocks = ceil(prompt_len / tpb). The sender reconstructs total_blocks from token_range.end via ceil(end / tpb) to derive per-layer token starts, so the two block counts disagree whenever num_extra_kv_tokens pushes token_range.end across a block boundary. In current configurations the resulting offset is applied symmetrically to the src and dst token starts and cancels in _align_kv_blocks, so no transfer is corrupted today (dst_start_token is always None and num_extra_kv_tokens is smaller than tokens_per_block). It is nonetheless a latent violation of the KVSlice contract (total_blocks = ceil(token_range.end / tpb)) that would surface if dst_start_token becomes non-None, or if num_extra_kv_tokens exceeds tokens_per_block together with asymmetric prefix reuse. Set token_range.end to req.prompt_len so it matches the trimmed block list. Not transferring the extra speculative-decoding blocks follows the NVIDIA/TensorRT-LLM PR #14546 review discussion; the accuracy impact on speculative decoding is not yet confirmed. Update tests: flip test_includes_num_extra_kv_tokens to test_excludes_num_extra_kv_tokens and add a block-boundary regression. Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com> --- .../_torch/disaggregation/transceiver.py | 18 +++++---- .../disaggregated/test_cache_reuse_adapter.py | 38 +++++++++++++------ 2 files changed, 36 insertions(+), 20 deletions(-) diff --git a/tensorrt_llm/_torch/disaggregation/transceiver.py b/tensorrt_llm/_torch/disaggregation/transceiver.py index 5ff8b305596f..1a5d36705b34 100644 --- a/tensorrt_llm/_torch/disaggregation/transceiver.py +++ b/tensorrt_llm/_torch/disaggregation/transceiver.py @@ -182,12 +182,16 @@ def _create_kv_slice(self, req: LlmRequest) -> KVSlice: token_range = None if req.prompt_len > 0: - # Align with KV cache allocation (prepare_disagg_gen_init / - # _get_context_bytes), which reserves prompt_len + - # num_extra_kv_tokens slots for speculative decoding methods - # (e.g. EAGLE3) that consume extra KV positions per request. - num_extra_kv_tokens = getattr(self._kv_cache_manager, "num_extra_kv_tokens", 0) or 0 - token_range = TokenRange(start=0, end=req.prompt_len + num_extra_kv_tokens) + # end must match the trimmed block list below (ceil(prompt_len / tpb) + # blocks). num_extra_kv_tokens slots (speculative decoding) are not + # transferred. In the previously added support for ctx disabling + # speculative decoding while gen enables it, both sides currently + # use prompt_len as the transfer range, so the ranges stay + # consistent. + # TODO: the accuracy impact of not transferring num_extra_kv_tokens + # on MTP and other speculative decoding paths is currently unclear; + # revisit whether these extra KV slots need to be transferred. + token_range = TokenRange(start=0, end=req.prompt_len) groups = [] for idx, lg in enumerate(layer_groups): @@ -196,8 +200,6 @@ def _create_kv_slice(self, req: LlmRequest) -> KVSlice: continue block_ids = adapter.get_block_ids(req, idx, lg) # Limit to prompt_len blocks, matching C++ cacheFormatter behavior. - # Extra blocks from num_extra_kv_tokens (speculative decoding) have - # uninitialized KV data and must not be transferred. total_blocks = (req.prompt_len + tpb - 1) // tpb if block_ids.size > total_blocks: block_ids = block_ids[:total_blocks] diff --git a/tests/unittest/disaggregated/test_cache_reuse_adapter.py b/tests/unittest/disaggregated/test_cache_reuse_adapter.py index b006a9b7a197..c6828bac8684 100644 --- a/tests/unittest/disaggregated/test_cache_reuse_adapter.py +++ b/tests/unittest/disaggregated/test_cache_reuse_adapter.py @@ -269,8 +269,8 @@ def test_start_eq_end_rejected(self): # --------------------------------------------------------------------------- -# _create_kv_slice: default TokenRange spans prompt_len + num_extra_kv_tokens -# so transferred KV matches what resize_context / _get_context_bytes allocate. +# _create_kv_slice: default TokenRange spans prompt_len, matching the +# trimmed block list actually transferred. # --------------------------------------------------------------------------- @@ -310,15 +310,14 @@ def _build_transceiver_for_kv_slice(num_extra_kv_tokens: int, prompt_len: int): class TestCreateKvSliceTokenRange: - """Default TokenRange built by _create_kv_slice must align with KV-cache allocation. + """token_range.end must be prompt_len, matching the trimmed block list. - KV cache allocation in resize_context (V2) and prepare_resources (V1) reserves - prompt_len + num_extra_kv_tokens slots whenever speculative decoding (e.g. - EAGLE3, MTP) consumes extra KV positions per request. The transferred token - range must cover the same span, otherwise the receiver under-receives KV. + The sender reconstructs total_blocks from token_range.end (ceil(end / tpb)), + so end must stay at prompt_len -- not prompt_len + num_extra_kv_tokens -- + to match the blocks actually transferred. """ - def test_includes_num_extra_kv_tokens(self): + def test_excludes_num_extra_kv_tokens(self): prompt_len = 17 num_extra_kv_tokens = 7 transceiver, req = _build_transceiver_for_kv_slice(num_extra_kv_tokens, prompt_len) @@ -326,10 +325,25 @@ def test_includes_num_extra_kv_tokens(self): kv_slice = transceiver._create_kv_slice(req) assert kv_slice.token_range is not None - assert (kv_slice.token_range.start, kv_slice.token_range.end) == ( - 0, - prompt_len + num_extra_kv_tokens, - ) + assert (kv_slice.token_range.start, kv_slice.token_range.end) == (0, prompt_len) + + def test_extra_tokens_do_not_cross_block_boundary(self): + # Reconstructed total_blocks (ceil(end / tpb)) must match the blocks sent. + prompt_len = 16 + num_extra_kv_tokens = 7 + transceiver, req = _build_transceiver_for_kv_slice(num_extra_kv_tokens, prompt_len) + tpb = transceiver._reuse_adapter.tokens_per_block + + # Setup must actually exercise a boundary crossing: prompt_len ends on a + # block boundary and the extra tokens would otherwise add a block. + assert prompt_len % tpb == 0 + assert (prompt_len + num_extra_kv_tokens + tpb - 1) // tpb == prompt_len // tpb + 1 + + kv_slice = transceiver._create_kv_slice(req) + + end = kv_slice.token_range.end + transferred_blocks = kv_slice.block_ids_per_layer_groups[0].size + assert (end + tpb - 1) // tpb == transferred_blocks def test_defaults_to_prompt_len_when_no_extra(self): prompt_len = 17 From b308579a7ebd1c1a280f6bf7f04c1abe8b56f441 Mon Sep 17 00:00:00 2001 From: Chuang Zhu <111838961+chuangz0@users.noreply.github.com> Date: Wed, 15 Jul 2026 07:58:54 +0000 Subject: [PATCH 2/2] [None][fix] raise on dst/src block count mismatch instead of trimming Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com> --- .../_torch/disaggregation/native/transfer.py | 14 +++++--------- 1 file changed, 5 insertions(+), 9 deletions(-) diff --git a/tensorrt_llm/_torch/disaggregation/native/transfer.py b/tensorrt_llm/_torch/disaggregation/native/transfer.py index fe5105bc8652..c1bc6403b8fb 100644 --- a/tensorrt_llm/_torch/disaggregation/native/transfer.py +++ b/tensorrt_llm/_torch/disaggregation/native/transfer.py @@ -757,18 +757,14 @@ def _build_kv_write_meta(self, task: KVSendTask, req_info: RecvReqInfo) -> Write src_block_ids = src_block_ids_per_groups[self_lg] dst_block_ids = dst_block_ids_per_groups[peer_lg] - # Speculative decoding: generation may have one extra draft-token block. + # Both sides trim block lists to ceil(prompt_len / tpb) in + # _create_kv_slice, so dst must never exceed src. A smaller dst + # (generation prefix-cache reuse) is handled via dst_start below. block_diff = dst_block_ids.size - src_block_ids.size - if block_diff == 1: - logger.debug( - f"Trimming 1 extra dst block for draft tokens: " - f"src={src_block_ids.size}, dst={dst_block_ids.size}" - ) - dst_block_ids = dst_block_ids[:-1] - elif block_diff > 1: + if block_diff > 0: raise ValueError( f"src/dst block count mismatch: {src_block_ids.size} vs " - f"{dst_block_ids.size} (expected diff <= 1)" + f"{dst_block_ids.size} (dst must not exceed src)" ) tpb = extractor.page_table.tokens_per_block token_range = task._slice.token_range