[AMD] Dsv4/pr1 fix run time issue (#25898)

Co-authored-by: wunhuang <wunhuang@amd.com>
Co-authored-by: Thomas Wang <1am9trash@gmail.com>
Co-authored-by: Xinyi Song <86638975+RolaoDenthu@users.noreply.github.com>
Co-authored-by: HaiShaw <hixiao@gmail.com>
Co-authored-by: amd-danli103 <danli103@amd.com>
Co-authored-by: Lin, Soga <soga.lin@amd.com>
Co-authored-by: Raiden-Makoto <Raiden-Makoto@users.noreply.github.com>
Co-authored-by: Hubert Lu <55214931+hubertlu-tw@users.noreply.github.com>
Co-authored-by: yichiche@amd.com <jacky.cheng>
Co-authored-by: yctseng0211 <yctseng@amd.com>
Co-authored-by: Bingxu Chen <bingxche@amd.com>

sgl-kernel/python/sgl_kernel/top_k.py

@@ -80,6 +80,38 @@ def fast_topk_transform_fused(
     return dst_page_table
 
 
+def deepseek_v4_topk_transform_512(
+    scores: torch.Tensor,
+    seq_lens: torch.Tensor,
+    page_table: torch.Tensor,
+    page_indices: torch.Tensor,
+    page_size: int,
+    raw_indices: Optional[torch.Tensor] = None,
+) -> None:
+    """
+    Performs the DeepSeek-V4 indexer top-k selection and writes the paged
+    physical slot indices into ``page_indices``. Supports topk up to 1024.
+    Optionally also writes the row-relative raw token positions into
+    ``raw_indices`` for hisparse capture.
+
+    Args:
+        scores: float32 ``[B, max_seq_len]`` indexer logits, contiguous on dim 1.
+        seq_lens: int32 ``[B]``, true KV length per batch row.
+        page_table: int32 ``[B, num_pages]``, logical->physical page table,
+            contiguous on dim 1.
+        page_indices: int32 ``[B, topk]``, output buffer, contiguous. Filled
+            with paged physical slots; -1 for padding entries.
+        page_size: power-of-2 page size.
+        raw_indices: optional int32 ``[B, topk]``, contiguous. If provided,
+            filled with raw token positions within each row.
+    """
+    if raw_indices is not None:
+        assert raw_indices.dim() == 2
+    torch.ops.sgl_kernel.deepseek_v4_topk_transform_512(
+        scores, seq_lens, page_table, page_indices, page_size, raw_indices
+    )
+
+
 def fast_topk_transform_ragged_fused(
     score: torch.Tensor,
     lengths: torch.Tensor,