[None][fix] fix token_range_end add extra_kv_num_tokens (#14258)

Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com>

Author: chuangz0

tensorrt_llm/_torch/disaggregation/transceiver.py

@@ -181,7 +181,12 @@ def _create_kv_slice(
         )
 
         if token_range is None and req.prompt_len > 0:
-            token_range = TokenRange(start=0, end=req.prompt_len)
+            # Align with KV cache allocation (resize_context /
+            # _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)
 
         groups = []
         for idx, lg in enumerate(layer_groups):

tests/unittest/disaggregated/test_cache_reuse_adapter.py

@@ -14,13 +14,16 @@
 # limitations under the License.
 """Tests for CacheReuseAdapter, _create_kv_slice SWA trim, and Sender token-start derivation."""
 
+from types import SimpleNamespace
+
 import numpy as np
 import pytest
 
 from tensorrt_llm._torch.disaggregation.base.transfer import TokenRange
 from tensorrt_llm._torch.disaggregation.native.transfer import Sender
 from tensorrt_llm._torch.disaggregation.resource.cache_reuse import CacheReuseAdapter
 from tensorrt_llm._torch.disaggregation.resource.page import AttentionLayerGroup
+from tensorrt_llm._torch.disaggregation.transceiver import KvCacheTransceiverV2
 
 # ---------------------------------------------------------------------------
 # _align_kv_blocks: contract unchanged.
@@ -114,6 +117,91 @@ def test_start_eq_end_rejected(self):
             TokenRange(start=256, end=256)
 
 
+# ---------------------------------------------------------------------------
+# _create_kv_slice: default TokenRange spans prompt_len + num_extra_kv_tokens
+# so transferred KV matches what resize_context / _get_context_bytes allocate.
+# ---------------------------------------------------------------------------
+
+
+def _build_transceiver_for_kv_slice(num_extra_kv_tokens: int, prompt_len: int):
+    """Stub a KvCacheTransceiverV2 so _create_kv_slice runs without dist setup.
+
+    Wires only the attributes the method touches:
+      - reuse adapter: tokens_per_block, per-layer-group cached count, block ids
+      - page table:    layer groups
+      - cache manager: num_extra_kv_tokens (read in this code path)
+    """
+    tokens_per_block = 8
+    layer_group = AttentionLayerGroup(pool_group_idx=0, kv_head_num_per_rank=1)
+    total_blocks = (prompt_len + num_extra_kv_tokens + tokens_per_block - 1) // tokens_per_block
+    block_ids = np.arange(total_blocks, dtype=np.int64)
+
+    reuse_adapter = SimpleNamespace(
+        tokens_per_block=tokens_per_block,
+        get_cached_token_count_per_layer_group=lambda req, layer_groups: [0] * len(layer_groups),
+        get_block_ids=lambda req, idx, lg: block_ids,
+    )
+    page_table = SimpleNamespace(layer_groups=[layer_group])
+    cache_manager = SimpleNamespace(num_extra_kv_tokens=num_extra_kv_tokens)
+
+    transceiver = object.__new__(KvCacheTransceiverV2)
+    transceiver._reuse_adapter = reuse_adapter
+    transceiver._page_table = page_table
+    transceiver._kv_cache_manager = cache_manager
+
+    req = SimpleNamespace(
+        prompt_len=prompt_len,
+        py_request_id=0,
+        is_generation_only_request=lambda: False,
+    )
+    return transceiver, req
+
+
+class TestCreateKvSliceTokenRange:
+    """Default TokenRange built by _create_kv_slice must align with KV-cache allocation.
+
+    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.
+    """
+
+    def test_includes_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)
+
+        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,
+        )
+
+    def test_defaults_to_prompt_len_when_no_extra(self):
+        prompt_len = 17
+        transceiver, req = _build_transceiver_for_kv_slice(
+            num_extra_kv_tokens=0, prompt_len=prompt_len
+        )
+
+        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)
+
+    def test_respects_explicit_token_range(self):
+        prompt_len = 17
+        transceiver, req = _build_transceiver_for_kv_slice(
+            num_extra_kv_tokens=7, prompt_len=prompt_len
+        )
+        explicit = TokenRange(start=0, end=8)
+
+        kv_slice = transceiver._create_kv_slice(req, token_range=explicit)
+
+        assert kv_slice.token_range is explicit
+
+
 # ---------------------------------------------------------------------------
 # CacheReuseAdapter.get_cached_token_count_per_layer_group: SWA clamp.
 # ---------------------------------------------------------------------------