[AMD] Dsv4/pr2 compressor opt (#26208)

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/benchmark/bench_dsv4_norm_rope.py

@@ -0,0 +1,75 @@
+"""Benchmark for DeepSeek-V4 fused norm + RoPE kernels."""
+
+import itertools
+
+import sgl_kernel
+import torch
+import triton
+import triton.testing
+
+try:
+    from sglang.utils import is_in_ci
+
+    IS_CI = is_in_ci()
+except ImportError:
+    IS_CI = False
+
+batch_sizes = [1] if IS_CI else [1, 4, 16, 64, 256]
+num_heads_list = [8] if IS_CI else [8, 16, 64]
+head_dims = [192] if IS_CI else [128, 192]
+
+configs = list(itertools.product(batch_sizes, num_heads_list, head_dims))
+
+
+def torch_rmsnorm_rope(
+    q: torch.Tensor, freqs_cis: torch.Tensor, positions: torch.Tensor, eps: float
+) -> torch.Tensor:
+    """Naive PyTorch reference: RMSNorm + RoPE."""
+    rms = torch.sqrt(q.float().pow(2).mean(dim=-1, keepdim=True) + eps)
+    q_normed = (q.float() / rms).to(q.dtype)
+    return q_normed
+
+
+@triton.testing.perf_report(
+    triton.testing.Benchmark(
+        x_names=["batch_size", "num_heads", "head_dim"],
+        x_vals=configs,
+        line_arg="provider",
+        line_vals=["sglang", "torch"],
+        line_names=["SGL Kernel", "PyTorch"],
+        styles=[("green", "-"), ("red", "--")],
+        ylabel="µs (median)",
+        plot_name="dsv4-q-norm-rope-performance",
+        args={},
+    )
+)
+def benchmark_q_norm_rope(batch_size, num_heads, head_dim, provider):
+    torch.manual_seed(42)
+    eps = 1e-6
+    max_pos = 8192
+    rope_dim = 64
+
+    q_input = torch.randn(
+        batch_size, num_heads, head_dim, dtype=torch.bfloat16, device="cuda"
+    )
+    q_output = torch.empty_like(q_input)
+    freqs_cis = torch.randn(max_pos, rope_dim, dtype=torch.float32, device="cuda")
+    positions = torch.randint(
+        0, max_pos, (batch_size,), dtype=torch.int32, device="cuda"
+    )
+
+    if provider == "sglang":
+        fn = lambda: sgl_kernel.dsv4_fused_q_norm_rope(
+            q_input, freqs_cis, positions, eps, q_output
+        )
+    else:
+        fn = lambda: torch_rmsnorm_rope(q_input, freqs_cis, positions, eps)
+
+    ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
+        fn, quantiles=[0.5, 0.2, 0.8]
+    )
+    return 1000 * ms, 1000 * max_ms, 1000 * min_ms
+
+
+if __name__ == "__main__":
+    benchmark_q_norm_rope.run(print_data=True)

sgl-kernel/csrc/common_extension.cc

@@ -214,6 +214,21 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) {
   m.def("apply_shuffle_mul_sum(Tensor input, Tensor output, Tensor permutation, Tensor? factors) -> ()");
   m.impl("apply_shuffle_mul_sum", torch::kCUDA, &apply_shuffle_mul_sum);
 
+  // DeepSeek-V4 fused norm + rope
+  m.def(
+      "dsv4_fused_q_norm_rope(Tensor q_input, Tensor! q_output, Tensor freqs_cis, Tensor positions, float eps) -> ()");
+  m.impl("dsv4_fused_q_norm_rope", torch::kCUDA, &dsv4_fused_q_norm_rope);
+
+  m.def(
+      "dsv4_fused_k_norm_rope_flashmla(Tensor kv, Tensor kv_weight, Tensor freqs_cis, Tensor positions, "
+      "Tensor out_loc, Tensor! kvcache, float eps, int page_size) -> ()");
+  m.impl("dsv4_fused_k_norm_rope_flashmla", torch::kCUDA, &dsv4_fused_k_norm_rope_flashmla);
+
+  m.def(
+      "dsv4_fused_q_indexer_rope_hadamard_quant(Tensor q_input, Tensor! q_fp8, Tensor weight, "
+      "Tensor! weights_out, float weight_scale, Tensor freqs_cis, Tensor positions) -> ()");
+  m.impl("dsv4_fused_q_indexer_rope_hadamard_quant", torch::kCUDA, &dsv4_fused_q_indexer_rope_hadamard_quant);
+
   m.def(
       "fused_qk_norm_rope(Tensor! qkv, int num_heads_q, "
       "int num_heads_k, int num_heads_v, int head_dim, float eps, "

sgl-kernel/include/sgl_kernel_ops.h

@@ -368,6 +368,35 @@ void apply_shuffle_mul_sum(
     const torch::Tensor& permutation,
     const std::optional<torch::Tensor>& factors);
 
+/*
+ * From csrc/elementwise (DeepSeek-V4 norm + rope)
+ */
+void dsv4_fused_q_norm_rope(
+    const at::Tensor& q_input,
+    at::Tensor& q_output,
+    const at::Tensor& freqs_cis,
+    const at::Tensor& positions,
+    double eps);
+
+void dsv4_fused_k_norm_rope_flashmla(
+    const at::Tensor& kv,
+    const at::Tensor& kv_weight,
+    const at::Tensor& freqs_cis,
+    const at::Tensor& positions,
+    const at::Tensor& out_loc,
+    at::Tensor& kvcache,
+    double eps,
+    int64_t page_size);
+
+void dsv4_fused_q_indexer_rope_hadamard_quant(
+    const at::Tensor& q_input,
+    at::Tensor& q_fp8,
+    const at::Tensor& weight,
+    at::Tensor& weights_out,
+    double weight_scale,
+    const at::Tensor& freqs_cis,
+    const at::Tensor& positions);
+
 void fused_qk_norm_rope(
     torch::Tensor& qkv,
     int64_t num_heads_q,

sgl-kernel/python/sgl_kernel/__init__.py

@@ -36,6 +36,9 @@ else:
         concat_mla_absorb_q,
         concat_mla_k,
         copy_to_gpu_no_ce,
+        dsv4_fused_k_norm_rope_flashmla,
+        dsv4_fused_q_indexer_rope_hadamard_quant,
+        dsv4_fused_q_norm_rope,
         fused_add_rmsnorm,
         gelu_and_mul,
         gelu_tanh_and_mul,
@@ -125,6 +128,7 @@ else:
 
     if torch.version.hip is not None:
         from sgl_kernel.elementwise import gelu_quick
+        from sgl_kernel.top_k import deepseek_v4_topk_transform_512
 
     if hasattr(torch.version, "musa") and torch.version.musa is not None:
         from sgl_kernel.musa import (
@@ -152,6 +156,9 @@ else:
         "cutlass_mla_get_workspace_size",
         "dsv3_fused_a_gemm",
         "dsv3_router_gemm",
+        "dsv4_fused_k_norm_rope_flashmla",
+        "dsv4_fused_q_indexer_rope_hadamard_quant",
+        "dsv4_fused_q_norm_rope",
         "es_fp8_blockwise_scaled_grouped_mm",
         "es_sm100_mxfp8_blockscaled_grouped_mm",
         "es_sm100_mxfp8_blockscaled_grouped_quant",
@@ -205,6 +212,7 @@ else:
 
     if torch.version.hip is not None:
         _DEBUG_EXPORT_NAMES.append("gelu_quick")
+        _DEBUG_EXPORT_NAMES.append("deepseek_v4_topk_transform_512")
 
     for _name in _DEBUG_EXPORT_NAMES:
         if _name in globals():

sgl-kernel/tests/test_dsv4_norm_rope.py

@@ -0,0 +1,130 @@
+"""Tests for DeepSeek-V4 fused norm + RoPE kernels."""
+
+import math
+
+import pytest
+import sgl_kernel
+import torch
+
+
+def _ref_rmsnorm_self(x: torch.Tensor, eps: float) -> torch.Tensor:
+    """Reference: RMSNorm without weight (identity weight)."""
+    rms = torch.sqrt(x.float().pow(2).mean(dim=-1, keepdim=True) + eps)
+    return (x.float() / rms).to(x.dtype)
+
+
+def _ref_rope_interleaved(
+    x: torch.Tensor, freqs_cis: torch.Tensor, positions: torch.Tensor, rope_dim: int
+) -> torch.Tensor:
+    """Reference: apply RoPE to the last `rope_dim` elements (interleaved re/im)."""
+    out = x.clone()
+    B = x.size(0)
+    head_dim = x.size(-1)
+    nope_dim = head_dim - rope_dim
+
+    for b in range(B):
+        pos = positions[b].item()
+        freq = freqs_cis[pos]  # (rope_dim,) interleaved [re0, im0, re1, im1, ...]
+        rope_part = out[b, ..., nope_dim:].float()
+        # Reshape to pairs
+        pairs = rope_part.reshape(*rope_part.shape[:-1], rope_dim // 2, 2)
+        x_real = pairs[..., 0]
+        x_imag = pairs[..., 1]
+        freq_pairs = freq.reshape(rope_dim // 2, 2)
+        f_real = freq_pairs[:, 0]
+        f_imag = freq_pairs[:, 1]
+        rot_real = x_real * f_real - x_imag * f_imag
+        rot_imag = x_real * f_imag + x_imag * f_real
+        result = torch.stack([rot_real, rot_imag], dim=-1).reshape(rope_part.shape)
+        out[b, ..., nope_dim:] = result.to(x.dtype)
+    return out
+
+
+@pytest.mark.parametrize("batch_size", [1, 4, 16])
+@pytest.mark.parametrize("num_heads", [1, 8])
+@pytest.mark.parametrize("head_dim", [128, 192])
+def test_fused_q_norm_rope_correctness(batch_size, num_heads, head_dim):
+    """Test Q norm + rope against reference."""
+    torch.manual_seed(42)
+    rope_dim = 64
+    max_pos = 512
+    eps = 1e-6
+
+    q_input = torch.randn(
+        batch_size, num_heads, head_dim, dtype=torch.bfloat16, device="cuda"
+    )
+    freqs_cis = torch.randn(max_pos, rope_dim, dtype=torch.float32, device="cuda")
+    positions = torch.randint(
+        0, max_pos, (batch_size,), dtype=torch.int32, device="cuda"
+    )
+
+    q_output = sgl_kernel.dsv4_fused_q_norm_rope(q_input, freqs_cis, positions, eps)
+
+    # Reference
+    normed = _ref_rmsnorm_self(q_input, eps)
+    expected = _ref_rope_interleaved(normed, freqs_cis, positions, rope_dim)
+
+    torch.testing.assert_close(q_output.float(), expected.float(), rtol=1e-2, atol=1e-2)
+
+
+def test_fused_q_norm_rope_zero_batch():
+    """Empty batch should not crash."""
+    q_input = torch.empty(0, 8, 192, dtype=torch.bfloat16, device="cuda")
+    freqs_cis = torch.randn(512, 64, dtype=torch.float32, device="cuda")
+    positions = torch.empty(0, dtype=torch.int32, device="cuda")
+    q_output = sgl_kernel.dsv4_fused_q_norm_rope(q_input, freqs_cis, positions)
+    assert q_output.shape == q_input.shape
+
+
+def test_fused_q_norm_rope_preallocated_output():
+    """Test with pre-allocated output tensor."""
+    torch.manual_seed(42)
+    B, H, D = 4, 8, 192
+    q_input = torch.randn(B, H, D, dtype=torch.bfloat16, device="cuda")
+    freqs_cis = torch.randn(512, 64, dtype=torch.float32, device="cuda")
+    positions = torch.randint(0, 512, (B,), dtype=torch.int32, device="cuda")
+    q_output = torch.empty_like(q_input)
+
+    result = sgl_kernel.dsv4_fused_q_norm_rope(
+        q_input, freqs_cis, positions, q_output=q_output
+    )
+    assert result is q_output
+
+
+@pytest.mark.parametrize("batch_size", [1, 8])
+def test_fused_q_indexer_rope_hadamard_quant_runs(batch_size):
+    """Smoke test: kernel runs without errors and produces finite results."""
+    torch.manual_seed(42)
+    num_heads = 4
+    head_dim = 128
+    rope_dim = 64
+    max_pos = 256
+
+    q_input = torch.randn(
+        batch_size, num_heads, head_dim, dtype=torch.bfloat16, device="cuda"
+    )
+    q_fp8 = torch.empty(
+        batch_size, num_heads, head_dim, dtype=torch.uint8, device="cuda"
+    )
+    weight = torch.randn(batch_size, num_heads, dtype=torch.bfloat16, device="cuda")
+    weights_out = torch.empty(
+        batch_size, num_heads, 1, dtype=torch.float32, device="cuda"
+    )
+    freqs_cis = torch.randn(max_pos, rope_dim, dtype=torch.float32, device="cuda")
+    positions = torch.randint(
+        0, max_pos, (batch_size,), dtype=torch.int32, device="cuda"
+    )
+    weight_scale = 0.5
+
+    sgl_kernel.dsv4_fused_q_indexer_rope_hadamard_quant(
+        q_input, q_fp8, weight, weights_out, weight_scale, freqs_cis, positions
+    )
+
+    assert torch.isfinite(weights_out).all(), "weights_out contains non-finite values"
+    assert q_fp8.any(), "q_fp8 should not be all zeros"
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main([__file__, "-v"]))