From 5ff84e8b4fc43c199a0aa790ce6ec910ba0deb82 Mon Sep 17 00:00:00 2001 From: Iman Tabrizian <10105175+tabrizian@users.noreply.github.com> Date: Tue, 7 Jul 2026 14:19:50 -0700 Subject: [PATCH 1/2] [None][chore] Remove unused token_range parameter from _create_kv_slice The token_range parameter of KvCacheTransceiverV2._create_kv_slice is never passed by any caller; every call site relies on the default computed range TokenRange(0, prompt_len + num_extra_kv_tokens). Drop the parameter and the test case that existed only to exercise the explicit-range passthrough. The KVSlice.token_range field and its default computation are unchanged. Signed-off-by: Iman Tabrizian <10105175+tabrizian@users.noreply.github.com> --- tensorrt_llm/_torch/disaggregation/transceiver.py | 4 ++-- .../disaggregated/test_cache_reuse_adapter.py | 11 ----------- 2 files changed, 2 insertions(+), 13 deletions(-) diff --git a/tensorrt_llm/_torch/disaggregation/transceiver.py b/tensorrt_llm/_torch/disaggregation/transceiver.py index b3854da74de4..8445c5260283 100644 --- a/tensorrt_llm/_torch/disaggregation/transceiver.py +++ b/tensorrt_llm/_torch/disaggregation/transceiver.py @@ -176,7 +176,6 @@ def __exit__(self, _exc_type, _exc_val, _exc_tb): def _create_kv_slice( self, req: LlmRequest, - token_range: Optional[TokenRange] = None, is_last_slice: bool = True, ) -> KVSlice: adapter = self._reuse_adapter @@ -191,7 +190,8 @@ def _create_kv_slice( else [0] * len(layer_groups) ) - if token_range is None and req.prompt_len > 0: + 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 diff --git a/tests/unittest/disaggregated/test_cache_reuse_adapter.py b/tests/unittest/disaggregated/test_cache_reuse_adapter.py index 08d030a707b8..b006a9b7a197 100644 --- a/tests/unittest/disaggregated/test_cache_reuse_adapter.py +++ b/tests/unittest/disaggregated/test_cache_reuse_adapter.py @@ -342,17 +342,6 @@ def test_defaults_to_prompt_len_when_no_extra(self): 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. From ce9b56b53ed34b9d026548f0fd67d9c275baf594 Mon Sep 17 00:00:00 2001 From: Iman Tabrizian <10105175+tabrizian@users.noreply.github.com> Date: Tue, 7 Jul 2026 14:22:37 -0700 Subject: [PATCH 2/2] [None][chore] Remove unused is_last_slice parameter from _create_kv_slice Like token_range, the is_last_slice parameter is never passed by any caller; every call site relies on the True default. Hardcode is_last_slice=True in the constructed KVSlice. The KVSlice.is_last_slice field and the native transfer protocol are unchanged. Signed-off-by: Iman Tabrizian <10105175+tabrizian@users.noreply.github.com> --- tensorrt_llm/_torch/disaggregation/transceiver.py | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/tensorrt_llm/_torch/disaggregation/transceiver.py b/tensorrt_llm/_torch/disaggregation/transceiver.py index 8445c5260283..0c7123587ee4 100644 --- a/tensorrt_llm/_torch/disaggregation/transceiver.py +++ b/tensorrt_llm/_torch/disaggregation/transceiver.py @@ -173,11 +173,7 @@ def __enter__(self): def __exit__(self, _exc_type, _exc_val, _exc_tb): self.shutdown() - def _create_kv_slice( - self, - req: LlmRequest, - is_last_slice: bool = True, - ) -> KVSlice: + def _create_kv_slice(self, req: LlmRequest) -> KVSlice: adapter = self._reuse_adapter tpb = adapter.tokens_per_block assert self._page_table is not None @@ -242,7 +238,7 @@ def _create_kv_slice( mamba_state_index = self._kv_cache_manager.mamba_cache_index[req.py_request_id] return KVSlice( - is_last_slice=is_last_slice, + is_last_slice=True, block_ids_per_layer_groups=groups, mamba_state_index=mamba_state_index, token_range=token_range,