Replace torch topk with custom topk in context phase

Signed-off-by: yuhangh <58161490+heyuhhh@users.noreply.github.com>

Author: heyuhhh

tensorrt_llm/_torch/attention_backend/sparse/kernel.py

@@ -672,7 +672,8 @@ def rocket_batch_to_flatten_kernel(
             token_mask = token_offsets < prefix_budget
 
             # Load from prefix_indices
-            prefix_indices = valid_idx_in_selected * num_kv_heads * prefix_budget + head_idx * prefix_budget + token_offsets
+            flattened_idx = valid_idx_in_selected * num_kv_heads + head_idx
+            prefix_indices = flattened_idx * prefix_budget + token_offsets
             prefix_values = tl.load(prefix_indices_ptr + prefix_indices,
                                     mask=token_mask,
                                     other=0)
@@ -717,26 +718,29 @@ def triton_rocket_batch_to_flatten(
         prefix_indices: torch.Tensor, input_lens: torch.Tensor,
         valid_seq_indices: torch.Tensor, output_offsets: torch.Tensor,
         batch_size: int, total_output_tokens: int, window_size: int,
-        prompt_budget: int) -> tuple[torch.Tensor, torch.Tensor]:
+        prompt_budget: int,
+        num_kv_heads: int) -> tuple[torch.Tensor, torch.Tensor]:
     """
     Flatten indices considering both valid and invalid batches.
     For valid sequences, combines prefix_indices with dynamically computed window indices.
     For invalid sequences, generates sequential indices.
 
     Args:
-        prefix_indices: Selected prefix indices [valid_batch_size, num_kv_heads, prefix_budget]
+        prefix_indices: Selected prefix indices [valid_batch_size * num_kv_heads, prefix_budget]
         input_lens: Lengths for all sequences [batch_size]
         valid_seq_indices: Valid sequence indices [valid_batch_size]
         output_offsets: Offset for each batch [batch_size + 1]
         batch_size: Number of batches
         total_output_tokens: Total number of output tokens
         window_size: Size of sliding window at the end
         prompt_budget: Total number of tokens for valid sequences (prefix_budget + window_size)
+        num_kv_heads: Number of KV heads
 
     Returns:
         sparse_indices: Flattened sparse indices [num_kv_heads, total_output_tokens]
     """
-    valid_batch_size, num_kv_heads, prefix_budget = prefix_indices.shape
+    total_tasks, prefix_budget = prefix_indices.shape
+    valid_batch_size = total_tasks // num_kv_heads
 
     # Create output tensor
     sparse_indices = torch.empty((num_kv_heads, total_output_tokens),

tensorrt_llm/_torch/attention_backend/sparse/rocket.py

@@ -75,10 +75,24 @@ def __post_init__(self):
                                              dtype=torch.int32)
 
         # Context length of RocketKV key for each valid sequence
-        self.k_context_lens = torch.empty(
-            self.max_num_sequences,
-            device='cpu',
+        self.k_context_lens_cuda = self.get_empty(
+            self.cuda_graph_buffers,
+            (self.max_num_sequences, ),
             dtype=torch.int32,
+            cache_name="k_context_lens_cuda",
+            capture_graph=capture_graph,
+        )
+        self.k_context_lens = torch.zeros_like(self.k_context_lens_cuda,
+                                               device='cpu',
+                                               dtype=torch.int32)
+
+        # Start index of RocketKV key for each valid sequence
+        self.k_context_start_cuda = self.get_empty(
+            None,
+            (self.max_num_sequences, ),
+            dtype=torch.int32,
+            cache_name="k_context_start_cuda",
+            capture_graph=capture_graph,
         )
 
         # Cumulative context lengths for each sequence
@@ -231,6 +245,8 @@ def prepare(self):
         # Only consider sequences that are long enough for sparse kv indices prediction in context phase
         self.k_context_lens[:valid_batch_size] = self.prompt_lens_cpu[
             valid_seq_indices] - self.window_size
+        self.k_context_lens_cuda[:valid_batch_size].copy_(
+            self.k_context_lens[:valid_batch_size], non_blocking=True)
 
         sparse_counts_ctx = torch.zeros(self.num_contexts,
                                         dtype=torch.int32,
@@ -399,12 +415,32 @@ def sparse_kv_predict(
                 padding=self.kernel_size // 2,
                 stride=1)
 
-            selected_prefix_indices = scores.topk(
-                self.prompt_budget - self.window_size,
-                dim=-1).indices.sort().values.to(torch.int32)
+            # Use indexer topk prefill to select topk prefix indices
+            total_tasks = metadata.valid_batch_size * self.num_kv_heads
+
+            selected_prefix_indices = torch.empty(
+                (total_tasks, self.prompt_budget - self.window_size),
+                device=qkv_input.device,
+                dtype=torch.int32)
+
+            scores = scores.view(total_tasks, -1)
+
+            row_starts = metadata.k_context_start_cuda[:metadata.
+                                                       valid_batch_size].repeat_interleave(
+                                                           self.num_kv_heads)
+            row_ends = metadata.k_context_lens_cuda[:metadata.
+                                                    valid_batch_size].repeat_interleave(
+                                                        self.num_kv_heads)
+            torch.ops.trtllm.indexer_topk_prefill(
+                scores, row_starts, row_ends, selected_prefix_indices,
+                self.prompt_budget - self.window_size)
+
+            # Sort selected prefix indices to keep topk indices in ascending order
+            selected_prefix_indices = torch.sort(selected_prefix_indices,
+                                                 dim=-1).values
         else:
             selected_prefix_indices = torch.empty(
-                (0, self.num_kv_heads, self.prompt_budget - self.window_size),
+                (0, self.prompt_budget - self.window_size),
                 device=qkv_input.device,
                 dtype=torch.int32)
 
@@ -416,7 +452,7 @@ def sparse_kv_predict(
             selected_prefix_indices, metadata.prompt_lens_cuda,
             metadata.valid_seq_indices_cuda, sparse_kv_offsets,
             metadata.num_contexts, metadata.total_sparse_ctx_indices,
-            self.window_size, self.prompt_budget)
+            self.window_size, self.prompt_budget, self.num_kv_heads)
 
         # Update KT cache
         kt_cache_tensor = metadata.kv_cache_manager.get_kt_buffers(