Minor fix

Signed-off-by: yizhang-nv <[email protected]>

Author: yizhang-nv

tensorrt_llm/_torch/attention_backend/trtllm.py

@@ -726,7 +726,6 @@ def _post_init_with_buffers(self, buffers) -> None:
                 capture_graph=capture_graph,
             )
             self.host_kv_cache_block_offsets = self.kv_cache_manager.host_kv_cache_block_offsets
-            assert self.host_kv_cache_block_offsets.shape == self.kv_cache_block_offsets.shape, f"host_kv_cache_block_offsets and kv_cache_block_offsets should have the same shape, but got {self.host_kv_cache_block_offsets.shape} and {self.kv_cache_block_offsets.shape}"
             self.block_ids_per_seq = None
             self.kv_block_ids_per_seq = None
             if self.enable_flash_mla:

tensorrt_llm/_torch/pyexecutor/resource_manager.py

@@ -33,8 +33,6 @@
 from tensorrt_llm.runtime.kv_cache_manager_v2 import (LayerId, TokenIdExt,
                                                       _KVCache)
 from tensorrt_llm.runtime.kv_cache_manager_v2._config import DataRole
-from tensorrt_llm.runtime.kv_cache_manager_v2._copy_engine import \
-    copy_batch_block_offsets as copy_batch_block_offsets_nanobind
 from tensorrt_llm.runtime.kv_cache_manager_v2._utils import (exact_div,
                                                              typed_range)
 from tensorrt_llm.sampling_params import SamplingParams
@@ -1506,20 +1504,17 @@ def append_to_kv_heads_per_layer(num_kv_heads_per_layer: List[int],
 
         self.enable_block_reuse = kv_cache_config.enable_block_reuse
 
-        self.index_mapper = IndexMapper(max_batch_size, max_beam_width)
+        # Plus 1 for cuda graph dummy request
+        self.index_mapper = IndexMapper(max_batch_size + 1, max_beam_width)
 
         self.host_kv_cache_block_offsets = torch.empty(
             self.num_pools,
-            max_batch_size * max_beam_width,
+            (max_batch_size + 1) * max_beam_width,
             2,  # key and value
             self.max_blocks_per_seq,
             dtype=torch.int32,
             pin_memory=True,
             device='cpu')
-        import os
-
-        # V2 is using zero copy api for index copy. By default we use zero copy api.
-        self.index_copy_v1 = os.environ.get("INDEX_COPY_V1", "0") == "1"
 
     @property
     def blocks_in_primary_pool(self) -> int:
@@ -1734,8 +1729,6 @@ def add_dummy_requests(
     def free_resources(self, request: LlmRequest, pin_on_release: bool = False):
         kv_cache = self.kv_cache_map.pop(request.py_request_id)
         kv_cache.close()
-        if self.index_copy_v1:
-            return
         self.index_mapper.remove_sequence(request.py_request_id)
 
     def get_batch_cache_indices(
@@ -1918,17 +1911,6 @@ def update_resources(self,
     def copy_batch_block_offsets(self, dst_tensor: torch.Tensor,
                                  request_ids: List[int], beam_width: int,
                                  num_contexts: int, num_gen: int):
-        if self.index_copy_v1:
-            self.copy_batch_block_offsets_v1(dst_tensor, request_ids,
-                                             beam_width, num_contexts, num_gen)
-        else:
-            self.copy_batch_block_offsets_v2(dst_tensor, request_ids,
-                                             beam_width, num_contexts, num_gen)
-
-    @nvtx_range("copy_batch_block_offsets_v2")
-    def copy_batch_block_offsets_v2(self, dst_tensor: torch.Tensor,
-                                    request_ids: List[int], beam_width: int,
-                                    num_contexts: int, num_gen: int):
         assert beam_width == 1, "beam_width must be 1 for KVCacheManagerV2"
 
         assert num_contexts + num_gen == len(
@@ -1941,35 +1923,11 @@ def copy_batch_block_offsets_v2(self, dst_tensor: torch.Tensor,
             self.host_kv_cache_block_offsets, dst_tensor, copy_idx, True,
             torch.cuda.current_stream().cuda_stream)
 
-    def copy_batch_block_offsets_v1(self, dst_tensor: torch.Tensor,
-                                    request_ids: List[int], beam_width: int,
-                                    num_contexts: int, num_gen: int):
-        assert beam_width == 1, "beam_width must be 1 for KVCacheManager"
-
-        num_seqs = num_contexts + num_gen * beam_width
-
-        for offset, end in [(0, num_contexts), (num_contexts, num_seqs)]:
-            batch_cache_indices = []
-            for pool_idx in range(self.num_pools):
-                for req_id in request_ids[offset:end]:
-                    batch_cache_indices.append(
-                        self.kv_cache_map[req_id].get_page_indices(
-                            pool_idx, 0).buffer_info())
-            if len(batch_cache_indices) > 0:
-                copy_batch_block_offsets_nanobind(
-                    self.host_kv_cache_block_offsets, end - offset,
-                    batch_cache_indices, self.num_pools, offset)
-
-        dst_tensor[:, :num_seqs].copy_(
-            self.host_kv_cache_block_offsets[:, :num_seqs], non_blocking=True)
-
     def _create_kv_cache(self, request_id: int, lora_task_id: int,
                          input_tokens: Sequence[TokenIdExt]):
         assert request_id not in self.kv_cache_map, f"KV cache for request {request_id} already exists"
         kv_cache = self.impl.create_kv_cache(lora_task_id, input_tokens)
         self.kv_cache_map[request_id] = kv_cache
-        if self.index_copy_v1:
-            return kv_cache
         index = self.index_mapper.add_new_sequence(request_id)
         for i in range(self.max_beam_width):
             for pool_idx in range(self.num_pools):