[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/elementwise.py

@@ -324,6 +324,85 @@ if torch.version.hip is not None:
         return out
 
 
+def dsv4_fused_q_norm_rope(
+    q_input: torch.Tensor,
+    freqs_cis: torch.Tensor,
+    positions: torch.Tensor,
+    eps: float = 1e-6,
+    q_output: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """DeepSeek-V4 fused Q RMSNorm (no weight) + RoPE.
+
+    Parameters
+    ----------
+    q_input  : (B, num_q_heads, head_dim) bfloat16
+    freqs_cis: (max_pos, rope_dim) float32, re/im interleaved
+    positions: (B,) int32
+    eps      : RMSNorm epsilon
+    q_output : optional pre-allocated output tensor
+    """
+    if q_output is None:
+        q_output = torch.empty_like(q_input)
+    torch.ops.sgl_kernel.dsv4_fused_q_norm_rope.default(
+        q_input, q_output, freqs_cis, positions, eps
+    )
+    return q_output
+
+
+def dsv4_fused_k_norm_rope_flashmla(
+    kv: torch.Tensor,
+    kv_weight: torch.Tensor,
+    freqs_cis: torch.Tensor,
+    positions: torch.Tensor,
+    out_loc: torch.Tensor,
+    kvcache: torch.Tensor,
+    eps: float = 1e-6,
+    page_size: int = 1,
+) -> None:
+    """DeepSeek-V4 fused K RMSNorm + RoPE + FlashMLA FP8 store.
+
+    Parameters
+    ----------
+    kv       : (B, 512) bfloat16
+    kv_weight: (512,) bfloat16
+    freqs_cis: (max_pos, 64) float32
+    positions: (B,) int32
+    out_loc  : (B,) int32  cache slot ids
+    kvcache  : (npages, page_bytes) uint8
+    eps      : RMSNorm epsilon
+    page_size: page size (power of 2)
+    """
+    torch.ops.sgl_kernel.dsv4_fused_k_norm_rope_flashmla.default(
+        kv, kv_weight, freqs_cis, positions, out_loc, kvcache, eps, page_size
+    )
+
+
+def dsv4_fused_q_indexer_rope_hadamard_quant(
+    q_input: torch.Tensor,
+    q_fp8: torch.Tensor,
+    weight: torch.Tensor,
+    weights_out: torch.Tensor,
+    weight_scale: float,
+    freqs_cis: torch.Tensor,
+    positions: torch.Tensor,
+) -> None:
+    """DeepSeek-V4 fused Q indexer: RoPE + Hadamard + FP8 quant.
+
+    Parameters
+    ----------
+    q_input    : (B, num_heads, 128) bfloat16
+    q_fp8      : (B, num_heads, 128) fp8_e4m3 output
+    weight     : (B, num_heads) bfloat16
+    weights_out: (B, num_heads, 1) float32 output
+    weight_scale: scalar
+    freqs_cis  : (max_pos, 64) float32
+    positions  : (B,) int32
+    """
+    torch.ops.sgl_kernel.dsv4_fused_q_indexer_rope_hadamard_quant.default(
+        q_input, q_fp8, weight, weights_out, weight_scale, freqs_cis, positions
+    )
+
+
 def rotary_embedding(
     positions: torch.Tensor,
     query: torch.Tensor,

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,