[AMD] Register 3 ROCm-portable JIT kernel tests for AMD CI (#30212)

test/registered/jit/test_dsv32_indexer_fusion.py

@@ -27,12 +27,13 @@ from sglang.jit_kernel.fused_store_index_cache import (
     fused_store_index_k_cache,
 )
 from sglang.srt.utils import is_hip
-from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 
 _is_hip = is_hip()
 
 register_cuda_ci(est_time=45, stage="base-b-kernel-unit", runner_config="1-gpu-large")
 register_cuda_ci(est_time=90, suite="nightly-kernel-1-gpu", nightly=True)
+register_amd_ci(est_time=45, suite="jit-kernel-unit-test-amd")
 
 HEAD_DIM = 128
 ROPE_DIM = 64

test/registered/jit/test_rmsnorm.py

@@ -5,10 +5,12 @@ import pytest
 import torch
 
 from sglang.jit_kernel.utils import get_ci_test_range
-from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.srt.utils import is_hip
+from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 
 register_cuda_ci(est_time=45, stage="base-b-kernel-unit", runner_config="1-gpu-large")
 register_cuda_ci(est_time=240, suite="nightly-kernel-1-gpu", nightly=True)
+register_amd_ci(est_time=45, suite="jit-kernel-unit-test-amd")
 
 
 EPS = 1e-6
@@ -40,6 +42,33 @@ def flashinfer_rmsnorm(
     rmsnorm(input, weight, out=output, eps=eps)
 
 
+def torch_rmsnorm(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    *,
+    output: torch.Tensor,
+    eps: float = EPS,
+) -> None:
+    x = input.float()
+    normed = x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + eps)
+    output.copy_((normed * weight.float()).to(output.dtype))
+
+
+def reference_rmsnorm(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    *,
+    output: torch.Tensor,
+    eps: float = EPS,
+) -> None:
+    # NVIDIA uses flashinfer (the bitwise reference); flashinfer is CUDA-only,
+    # so on ROCm fall back to the torch reference (matches flashinfer math).
+    if is_hip():
+        torch_rmsnorm(input, weight, output=output, eps=eps)
+    else:
+        flashinfer_rmsnorm(input, weight, output=output, eps=eps)
+
+
 BS_LIST = [2**n for n in range(0, 14)]
 BS_LIST += [x + 1 + i for i, x in enumerate(BS_LIST)]
 SUPPORTED_HIDDEN_SIZE_LIST = [
@@ -81,9 +110,9 @@ def test_rmsnorm(
     input = torch.randn(batch_size, hidden_size, device=DEVICE, dtype=dtype)
     weight = torch.randn(hidden_size, device=DEVICE, dtype=dtype)
 
-    input_flashinfer = input.clone()
+    input_ref = input.clone()
-    output_flashinfer = torch.empty_like(input)
+    output_ref = torch.empty_like(input)
-    flashinfer_rmsnorm(input_flashinfer, weight, output=output_flashinfer)
+    reference_rmsnorm(input_ref, weight, output=output_ref)
 
     if specify_out:
         output_sglang = torch.empty_like(input)
@@ -92,7 +121,7 @@ def test_rmsnorm(
         output_sglang = input.clone()
         sglang_jit_rmsnorm(output_sglang, weight, output=output_sglang)
 
-    torch.testing.assert_close(output_sglang, output_flashinfer, atol=1e-2, rtol=1e-2)
+    torch.testing.assert_close(output_sglang, output_ref, atol=1e-2, rtol=1e-2)
 
 
 @pytest.mark.parametrize("hidden_size", [64, 128, 256, 512, 8192, 8704, 16384])

test/registered/jit/test_rope.py

@@ -5,10 +5,12 @@ import torch
 import triton
 
 from sglang.jit_kernel.utils import get_ci_test_range
-from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.srt.utils import is_hip
+from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 
 register_cuda_ci(est_time=64, stage="base-b-kernel-unit", runner_config="1-gpu-large")
 register_cuda_ci(est_time=256, suite="nightly-kernel-1-gpu", nightly=True)
+register_amd_ci(est_time=64, suite="jit-kernel-unit-test-amd")
 
 DEVICE = "cuda"
 DTYPE = torch.bfloat16
@@ -78,6 +80,31 @@ def flashinfer_rope(
     )
 
 
+def _rope_rotate(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, is_neox: bool):
+    """Rotate the first ``rotary_dim`` channels of ``x`` in place.
+
+    ``x``: [nnz, num_heads, head_size]; ``cos``/``sin``: [nnz, rotary_dim // 2].
+    Matches flashinfer's ``apply_rope_with_cos_sin_cache_inplace`` convention:
+    NeoX splits the rotary block into halves; non-NeoX (GPT-J) uses interleaved
+    even/odd pairs. Channels beyond ``rotary_dim`` are left untouched.
+    """
+    rotary_dim = cos.shape[-1] * 2
+    xf = x[..., :rotary_dim].to(torch.float32)
+    cos = cos[:, None, :]  # [nnz, 1, rotary_dim // 2]
+    sin = sin[:, None, :]
+    if is_neox:
+        x1, x2 = xf[..., : rotary_dim // 2], xf[..., rotary_dim // 2 :]
+        out1 = x1 * cos - x2 * sin
+        out2 = x2 * cos + x1 * sin
+        rotated = torch.cat((out1, out2), dim=-1)
+    else:
+        x1, x2 = xf[..., 0::2], xf[..., 1::2]
+        out1 = x1 * cos - x2 * sin
+        out2 = x2 * cos + x1 * sin
+        rotated = torch.stack((out1, out2), dim=-1).flatten(-2)
+    x[..., :rotary_dim] = rotated.to(x.dtype)
+
+
 def torch_impl_rope(
     q: torch.Tensor,
     k: torch.Tensor,
@@ -85,8 +112,29 @@ def torch_impl_rope(
     positions: torch.Tensor,
     is_neox: bool,
 ) -> None:
-    # TODO: implement a pure-PyTorch reference for extra coverage
+    """Pure-PyTorch RoPE reference (in place), used as the ROCm fallback."""
-    pass
+    rotary_dim = cos_sin_cache.shape[-1]
+    half = rotary_dim // 2
+    gathered = cos_sin_cache[positions.long()]
+    cos, sin = gathered[:, :half], gathered[:, half:]
+    _rope_rotate(q, cos, sin, is_neox)
+    _rope_rotate(k, cos, sin, is_neox)
+
+
+def reference_rope(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    cos_sin_cache: torch.Tensor,
+    positions: torch.Tensor,
+    is_neox: bool,
+) -> None:
+    # NVIDIA uses flashinfer (the reference); flashinfer is CUDA-only, so on
+    # ROCm fall back to the torch reference (matches flashinfer's cos/sin-cache
+    # application semantics).
+    if is_hip():
+        torch_impl_rope(q, k, cos_sin_cache, positions, is_neox)
+    else:
+        flashinfer_rope(q, k, cos_sin_cache, positions, is_neox)
 
 
 # ---------------------------------------------------------------------------
@@ -132,7 +180,7 @@ def test_rope(
     q_fi, k_fi = q.clone(), k.clone()
     q_jit, k_jit = q.clone(), k.clone()
 
-    flashinfer_rope(q_fi, k_fi, cos_sin_cache, positions, is_neox)
+    reference_rope(q_fi, k_fi, cos_sin_cache, positions, is_neox)
     sglang_jit_rope(q_jit, k_jit, cos_sin_cache, positions, is_neox)
 
     atol = rtol = 1e-2
@@ -154,7 +202,7 @@ def test_rope_position_dtypes(dtype: torch.dtype) -> None:
     q_fi, k_fi = q.clone(), k.clone()
     q_jit, k_jit = q.clone(), k.clone()
 
-    flashinfer_rope(q_fi, k_fi, cos_sin_cache, positions.long(), is_neox)
+    reference_rope(q_fi, k_fi, cos_sin_cache, positions.long(), is_neox)
     sglang_jit_rope(q_jit, k_jit, cos_sin_cache, positions, is_neox)
 
     atol = rtol = 1e-2
@@ -180,7 +228,7 @@ def test_partial_rope(batch_size: int, is_neox: bool, rope_dim: int, head_dim: i
     q_jit, k_jit = q.clone(), k.clone()
     rope = ..., slice(rope_dim)  # NOTE: flashinfer by default apply to first rope_dim
 
-    flashinfer_rope(q_fi, k_fi, cos_sin_cache, positions.long(), is_neox)
+    reference_rope(q_fi, k_fi, cos_sin_cache, positions.long(), is_neox)
     sglang_jit_rope(q_jit[rope], k_jit[rope], cos_sin_cache, positions, is_neox)
 
     atol = rtol = 1e-2
@@ -223,7 +271,7 @@ def test_fused_rope_store(
 
     # --- reference: separate RoPE then manual scatter ---
     q_ref, k_ref = q.clone(), k.clone()
-    flashinfer_rope(q_ref, k_ref, cos_sin_cache, positions, is_neox)
+    reference_rope(q_ref, k_ref, cos_sin_cache, positions, is_neox)
     k_cache_ref[out_loc] = k_ref.view(batch_size, -1)
     v_cache_ref[out_loc] = v.view(batch_size, -1)