[None][chore] Remove unused token_range and is_last_slice parameters from _create_kv_slice (NVIDIA#16083)

Signed-off-by: Iman Tabrizian <10105175+tabrizian@users.noreply.github.com>

Author: Tabrizian

tensorrt_llm/_torch/disaggregation/transceiver.py

@@ -167,12 +167,7 @@ def __enter__(self):
     def __exit__(self, _exc_type, _exc_val, _exc_tb):
         self.shutdown()
 
-    def _create_kv_slice(
-        self,
-        req: LlmRequest,
-        token_range: Optional[TokenRange] = None,
-        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
@@ -185,7 +180,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
@@ -236,7 +232,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,

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.