Address comment

Signed-off-by: Hui Gao <huig@nvidia.com>

Author: HuiGao-NV

tensorrt_llm/_torch/attention_backend/sparse/dsa.py

@@ -309,7 +309,8 @@ def __post_init__(self):
             [self.max_num_sequences, self.kv_cache_manager.max_blocks_per_seq],
             cache_name="indexer_k_cache_block_offsets",
             dtype=torch.int32,
-            capture_graph=capture_graph)
+            capture_graph=capture_graph,
+        )
         self.host_indexer_k_cache_block_offsets = torch.zeros_like(
             self.indexer_k_cache_block_offsets,
             device='cpu',
@@ -319,96 +320,117 @@ def __post_init__(self):
         # For mla_rope_append_paged_kv_assign_q
         if not self.enable_context_mla_with_cached_kv:
             self.ctx_cached_token_indptr = self.get_empty(
-                self.cuda_graph_buffers, (self.max_num_requests + 1, ),
+                self.cuda_graph_buffers,
+                (self.max_num_requests + 1, ),
                 cache_name="ctx_cached_token_indptr",
                 dtype=torch.int64,
-                capture_graph=capture_graph)
+                capture_graph=capture_graph,
+            )
             self.host_ctx_cached_token_indptr = torch.zeros_like(
                 self.ctx_cached_token_indptr,
                 device='cpu',
                 pin_memory=True,
             )
-            self.ctx_kv_indptr = self.get_empty(self.cuda_graph_buffers,
-                                                (self.max_num_requests + 1, ),
-                                                cache_name="ctx_kv_indptr",
-                                                dtype=torch.int64,
-                                                capture_graph=capture_graph)
+            self.ctx_kv_indptr = self.get_empty(
+                self.cuda_graph_buffers,
+                (self.max_num_requests + 1, ),
+                cache_name="ctx_kv_indptr",
+                dtype=torch.int64,
+                capture_graph=capture_graph,
+            )
             self.host_ctx_kv_indptr = torch.zeros_like(
                 self.ctx_kv_indptr,
                 device='cpu',
                 pin_memory=True,
             )
         # New generation buffers for dsa
         self.gen_cached_token_indptr = self.get_empty(
-            self.cuda_graph_buffers, (self.max_num_requests + 1, ),
+            self.cuda_graph_buffers,
+            (self.max_num_requests + 1, ),
             cache_name="gen_cached_token_indptr",
             dtype=torch.int64,
-            capture_graph=capture_graph)
+            capture_graph=capture_graph,
+        )
         self.host_gen_cached_token_indptr = torch.zeros_like(
             self.gen_cached_token_indptr,
             device='cpu',
             pin_memory=True,
         )
-        self.gen_kv_indptr = self.get_empty(self.cuda_graph_buffers,
-                                            (self.max_num_requests + 1, ),
-                                            cache_name="gen_kv_indptr",
-                                            dtype=torch.int64,
-                                            capture_graph=capture_graph)
+        self.gen_kv_indptr = self.get_empty(
+            self.cuda_graph_buffers,
+            (self.max_num_requests + 1, ),
+            cache_name="gen_kv_indptr",
+            dtype=torch.int64,
+            capture_graph=capture_graph,
+        )
         self.host_gen_kv_indptr = torch.zeros_like(
             self.gen_kv_indptr,
             device='cpu',
             pin_memory=True,
         )
         # Indexer metadata
         # Separate slot mappings for non-interleaved layout (flat byte indices)
-        self.slot_mapping_fp8 = self.get_empty(self.cuda_graph_buffers,
-                                               (self.max_num_tokens, ),
-                                               cache_name="slot_mapping_fp8",
-                                               dtype=torch.int64,
-                                               capture_graph=capture_graph)
+        self.slot_mapping_fp8 = self.get_empty(
+            self.cuda_graph_buffers,
+            (self.max_num_tokens, ),
+            cache_name="slot_mapping_fp8",
+            dtype=torch.int64,
+            capture_graph=capture_graph,
+        )
         self.host_slot_mapping_fp8 = torch.zeros_like(
             self.slot_mapping_fp8,
             device='cpu',
             pin_memory=True,
         )
         self.slot_mapping_scale = self.get_empty(
-            self.cuda_graph_buffers, (self.max_num_tokens, ),
+            self.cuda_graph_buffers,
+            (self.max_num_tokens, ),
             cache_name="slot_mapping_scale",
             dtype=torch.int64,
-            capture_graph=capture_graph)
+            capture_graph=capture_graph,
+        )
         self.host_slot_mapping_scale = torch.zeros_like(
             self.slot_mapping_scale,
             device='cpu',
             pin_memory=True,
         )
         # Per-token request index buffer for topk_indices conversion
-        self.req_idx_per_token = self.get_empty(self.cuda_graph_buffers,
-                                                (self.max_num_tokens, ),
-                                                cache_name="req_idx_per_token",
-                                                dtype=torch.int32,
-                                                capture_graph=capture_graph)
+        self.req_idx_per_token = self.get_empty(
+            self.cuda_graph_buffers,
+            (self.max_num_tokens, ),
+            cache_name="req_idx_per_token",
+            dtype=torch.int32,
+            capture_graph=capture_graph,
+        )
         # Block table for topk_indices conversion (shared for context and generation)
         self.block_table = self.get_empty(
             self.cuda_graph_buffers,
             (self.max_num_requests, self.kv_cache_manager.max_blocks_per_seq),
             cache_name="block_table",
             dtype=torch.int32,
-            capture_graph=capture_graph)
+            capture_graph=capture_graph,
+        )
         self.scheduler_metadata_buffer = self.get_empty(
-            self.cuda_graph_buffers, (self.num_sms + 1, 2),
+            self.cuda_graph_buffers,
+            (self.num_sms + 1, 2),
             cache_name="scheduler_metadata_buffer",
             dtype=torch.int32,
-            capture_graph=capture_graph)
-        self.cu_seqlen_ks = self.get_empty(self.cuda_graph_buffers,
-                                           (self.max_num_tokens, ),
-                                           cache_name="cu_seqlen_ks",
-                                           dtype=torch.int32,
-                                           capture_graph=capture_graph)
-        self.cu_seqlen_ke = self.get_empty(self.cuda_graph_buffers,
-                                           (self.max_num_tokens, ),
-                                           cache_name="cu_seqlen_ke",
-                                           dtype=torch.int32,
-                                           capture_graph=capture_graph)
+            capture_graph=capture_graph,
+        )
+        self.cu_seqlen_ks = self.get_empty(
+            self.cuda_graph_buffers,
+            (self.max_num_tokens, ),
+            cache_name="cu_seqlen_ks",
+            dtype=torch.int32,
+            capture_graph=capture_graph,
+        )
+        self.cu_seqlen_ke = self.get_empty(
+            self.cuda_graph_buffers,
+            (self.max_num_tokens, ),
+            cache_name="cu_seqlen_ke",
+            dtype=torch.int32,
+            capture_graph=capture_graph,
+        )
 
     def prepare(self):
         super().prepare()

tensorrt_llm/_torch/memory_buffer_utils.py

@@ -1,6 +1,5 @@
 import contextlib
 import math
-from collections import OrderedDict
 from dataclasses import dataclass
 from typing import Optional
 
@@ -9,28 +8,6 @@
 from tensorrt_llm.logger import logger
 
 
-def get_smallest_key_greater_than(ordered_dict, target_value):
-    """
-    Return (k, ordered_dict[k]) where k is the smallest key with k >= target_value,
-    or (None, None) if not found.
-    """
-    min_key = min((k for k in ordered_dict.keys() if k >= target_value),
-                  default=None)
-    return (min_key, ordered_dict[min_key]) if min_key is not None else (None,
-                                                                         None)
-
-
-def get_biggest_key_smaller_than(ordered_dict, target_value):
-    """
-    Return (k, ordered_dict[k]) where k is the largest key with k < target_value,
-    or (None, None) if not found.
-    """
-    max_key = max((k for k in ordered_dict.keys() if k < target_value),
-                  default=None)
-    return (max_key, ordered_dict[max_key]) if max_key is not None else (None,
-                                                                         None)
-
-
 def get_size_in_byte(target_shape: list[int], target_dtype: torch.dtype):
     return math.prod(target_shape) * target_dtype.itemsize
 
@@ -57,7 +34,6 @@ class Buffers:
 
     def __init__(self):
         self.buffers: dict[str, list[BufferBlock]] = {}
-        self.managed_buffers = OrderedDict()
         self.max_buffer_concurrency = 0
 
     @staticmethod
@@ -74,48 +50,12 @@ def _view_as(buffer: torch.Tensor, target_shape: list[int],
         return buffer[:required_memory_size].view(target_dtype).view(
             target_shape)
 
-    def _get_managed_buffer(self, required_memory_size: int):
-        size, buffer = get_smallest_key_greater_than(self.managed_buffers,
-                                                     required_memory_size)
-
-        if size is not None and buffer is not None:
-            return buffer
-
-        size_1, buffer_1 = get_biggest_key_smaller_than(self.managed_buffers,
-                                                        required_memory_size)
-        if size_1 is not None and buffer is not None:
-            del self.managed_buffers[size_1]
-
-        new_buffer_tensor = None
-        try:
-            with torch.cuda.memory.use_mem_pool(get_shared_pool()):
-                new_buffer_tensor = torch.zeros((required_memory_size, ),
-                                                device='cuda',
-                                                dtype=torch.uint8)
-        except Exception as ex:
-            # Need to check if this is an OOM exception
-            logger.debug(
-                f"Exception happened to create tensor from given memory pool: {str(ex)}"
-            )
-            # if exception happens during allocating memory from shared pool, retry
-            # to allocate from default pool
-            new_buffer_tensor = torch.zeros((required_memory_size, ),
-                                            device='cuda',
-                                            dtype=torch.uint8)
-
-        self.managed_buffers[required_memory_size] = new_buffer_tensor
-
-        return new_buffer_tensor
-
     def get_buffer(self, tensor_shape: list[int], dtype: torch.dtype,
                    buffer_name: str, reserve_buffer: bool):
 
         # all buffers are allocated with 1 byte element size
         required_memory_size = math.prod(tensor_shape) * dtype.itemsize
 
-        if buffer_name is None or len(buffer_name) == 0:
-            return _get_managed_buffer(required_memory_size)
-
         candidate_blocks = self.buffers.get(buffer_name, [])
 
         # Find the best-fit available buffer.