Fix error from rebase

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

Author: HuiGao-NV

tensorrt_llm/_torch/attention_backend/sparse/dsa.py

@@ -287,6 +287,7 @@ def __post_init__(self):
         capture_graph = torch.cuda.is_current_stream_capturing()
 
         self.indexer_k_cache_block_offsets = self.get_empty(
+            self.cuda_graph_buffers,
             [self.max_num_sequences, self.kv_cache_manager.max_blocks_per_seq],
             cache_name="indexer_k_cache_block_offsets",
             dtype=torch.int32,
@@ -300,7 +301,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 = self.get_empty(
-                (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)
@@ -309,7 +310,8 @@ def __post_init__(self):
                 device='cpu',
                 pin_memory=True,
             )
-            self.ctx_kv_indptr = self.get_empty((self.max_num_requests + 1, ),
+            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)
@@ -320,7 +322,7 @@ def __post_init__(self):
             )
         # New generation buffers for dsa
         self.gen_cached_token_indptr = self.get_empty(
-            (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)
@@ -329,7 +331,8 @@ def __post_init__(self):
             device='cpu',
             pin_memory=True,
         )
-        self.gen_kv_indptr = self.get_empty((self.max_num_requests + 1, ),
+        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)
@@ -340,7 +343,8 @@ def __post_init__(self):
         )
         # Indexer metadata
         # Separate slot mappings for non-interleaved layout (flat byte indices)
-        self.slot_mapping_fp8 = self.get_empty((self.max_num_tokens, ),
+        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)
@@ -350,7 +354,7 @@ def __post_init__(self):
             pin_memory=True,
         )
         self.slot_mapping_scale = self.get_empty(
-            (self.max_num_tokens, ),
+            self.cuda_graph_buffers, (self.max_num_tokens, ),
             cache_name="slot_mapping_scale",
             dtype=torch.int64,
             capture_graph=capture_graph)
@@ -360,26 +364,30 @@ def __post_init__(self):
             pin_memory=True,
         )
         # Per-token request index buffer for topk_indices conversion
-        self.req_idx_per_token = self.get_empty((self.max_num_tokens, ),
+        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)
         self.scheduler_metadata_buffer = self.get_empty(
-            (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.max_num_tokens, ),
+        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.max_num_tokens, ),
+        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)