[DeepseekV3.2][NSA][Indexer] Fix PAGED top-k transform for NSA indexer chunked execution on H200 (#14325)

python/sglang/srt/layers/attention/nsa/nsa_indexer.py

@@ -370,6 +370,8 @@ class Indexer(CustomOp):
         k_scale_list = []
         ks_list = []
         ke_list = []
+        # Token-to-batch mapping for PAGED chunk alignment
+        token_to_batch_idx: List[int] = []
 
         q_offset = 0
         k_offset = 0
@@ -401,6 +403,7 @@ class Indexer(CustomOp):
             ks_list.append(ks)
             ke_list.append(ke)
 
+            token_to_batch_idx.extend([i] * extend_seq_len)
             q_offset += extend_seq_len
             k_offset += seq_len
 
@@ -473,6 +476,13 @@ class Indexer(CustomOp):
             (token_nums, self.index_topk), -1, device=device, dtype=torch.int32
         )
 
+        # Only materialize batch index tensor when PAGED path needs it
+        token_to_batch_idx_tensor = None
+        if global_topk_offset is None:
+            token_to_batch_idx_tensor = torch.tensor(
+                token_to_batch_idx, dtype=torch.long, device=device
+            )
+
         start = 0
         while start < q_offset:
             end = min(start + max_rows, q_offset)
@@ -488,17 +498,28 @@ class Indexer(CustomOp):
 
             lengths_chunk = seq_lens_expanded[start:end]
 
-            topk_offset_chunk = (
-                global_topk_offset[start:end]
-                if global_topk_offset is not None
-                else None
-            )
+            # RAGGED: use global offset; PAGED: construct local cu_seqlens_q per chunk
+            if global_topk_offset is not None:
+                # RAGGED path
+                topk_offset_chunk = global_topk_offset[start:end]
+                cu_seqlens_q_chunk = None
+                batch_idx_chunk = None
+            else:
+                # PAGED path: treat each token as a length-1 sequence
+                topk_offset_chunk = None
+                B_chunk = logits_chunk.shape[0]
+                cu_seqlens_q_chunk = torch.ones(
+                    B_chunk, dtype=torch.int32, device=device
+                )
+                batch_idx_chunk = token_to_batch_idx_tensor[start:end]
 
             raw_topk_chunk = metadata.topk_transform(
                 logits_chunk,
                 self.index_topk,
                 ks=ks[start:end],
+                cu_seqlens_q=cu_seqlens_q_chunk,
                 ke_offset=lengths_chunk,
+                batch_idx_list=batch_idx_chunk,
                 topk_indices_offset_override=topk_offset_chunk,
             )
             topk_result[start:end] = raw_topk_chunk