Minor fix

Signed-off-by: yizhang-nv <187001205+yizhang-nv@users.noreply.github.com>

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

@@ -1506,11 +1506,12 @@ 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,