Address comments

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

Author: HuiGao-NV

tensorrt_llm/_torch/attention_backend/sparse/dsa.py

@@ -286,7 +286,7 @@ def __post_init__(self):
 
         capture_graph = torch.cuda.is_current_stream_capturing()
 
-        self.indexer_k_cache_block_offsets = get_empty(
+        self.indexer_k_cache_block_offsets = self.get_empty(
             [self.max_num_sequences, self.kv_cache_manager.max_blocks_per_seq],
             cache_name="indexer_k_cache_block_offsets",
             dtype=torch.int32,
@@ -299,7 +299,7 @@ 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 = get_empty(
+            self.ctx_cached_token_indptr = self.get_empty(
                 (self.max_num_requests + 1, ),
                 cache_name="ctx_cached_token_indptr",
                 dtype=torch.int64,
@@ -309,17 +309,17 @@ def __post_init__(self):
                 device='cpu',
                 pin_memory=True,
             )
-            self.ctx_kv_indptr = get_empty((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.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 = get_empty(
+        self.gen_cached_token_indptr = self.get_empty(
             (self.max_num_requests + 1, ),
             cache_name="gen_cached_token_indptr",
             dtype=torch.int64,
@@ -329,59 +329,60 @@ def __post_init__(self):
             device='cpu',
             pin_memory=True,
         )
-        self.gen_kv_indptr = get_empty((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.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 = get_empty((self.max_num_tokens, ),
-                                          cache_name="slot_mapping_fp8",
-                                          dtype=torch.int64,
-                                          capture_graph=capture_graph)
+        self.slot_mapping_fp8 = self.get_empty((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 = get_empty((self.max_num_tokens, ),
-                                            cache_name="slot_mapping_scale",
-                                            dtype=torch.int64,
-                                            capture_graph=capture_graph)
+        self.slot_mapping_scale = self.get_empty(
+            (self.max_num_tokens, ),
+            cache_name="slot_mapping_scale",
+            dtype=torch.int64,
+            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 = get_empty((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.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 = get_empty(
+        self.block_table = self.get_empty(
             (self.max_num_requests, self.kv_cache_manager.max_blocks_per_seq),
             cache_name="block_table",
             dtype=torch.int32,
             capture_graph=capture_graph)
-        self.scheduler_metadata_buffer = get_empty(
+        self.scheduler_metadata_buffer = self.get_empty(
             (self.num_sms + 1, 2),
             cache_name="scheduler_metadata_buffer",
             dtype=torch.int32,
             capture_graph=capture_graph)
-        self.cu_seqlen_ks = get_empty((self.max_num_tokens, ),
-                                      cache_name="cu_seqlen_ks",
-                                      dtype=torch.int32,
-                                      capture_graph=capture_graph)
-        self.cu_seqlen_ke = get_empty((self.max_num_tokens, ),
-                                      cache_name="cu_seqlen_ke",
-                                      dtype=torch.int32,
-                                      capture_graph=capture_graph)
+        self.cu_seqlen_ks = self.get_empty((self.max_num_tokens, ),
+                                           cache_name="cu_seqlen_ks",
+                                           dtype=torch.int32,
+                                           capture_graph=capture_graph)
+        self.cu_seqlen_ke = self.get_empty((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

@@ -9,6 +9,28 @@
 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