From 47a6dfd708b5587b7ba487f46150fb029b970634 Mon Sep 17 00:00:00 2001 From: Michael <13900043+michaelzhang-ai@users.noreply.github.com> Date: Wed, 8 Jul 2026 13:38:17 -0700 Subject: [PATCH] [AMD] Register 3 ROCm-portable JIT kernel tests for AMD CI (#30212) --- .../jit/test_dsv32_indexer_fusion.py | 3 +- test/registered/jit/test_rmsnorm.py | 39 ++++++++++-- test/registered/jit/test_rope.py | 62 ++++++++++++++++--- 3 files changed, 91 insertions(+), 13 deletions(-) diff --git a/test/registered/jit/test_dsv32_indexer_fusion.py b/test/registered/jit/test_dsv32_indexer_fusion.py index b01dce1a2..51210b605 100644 --- a/test/registered/jit/test_dsv32_indexer_fusion.py +++ b/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 diff --git a/test/registered/jit/test_rmsnorm.py b/test/registered/jit/test_rmsnorm.py index 8ec878f35..a7ba83ce4 100644 --- a/test/registered/jit/test_rmsnorm.py +++ b/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() - output_flashinfer = torch.empty_like(input) - flashinfer_rmsnorm(input_flashinfer, weight, output=output_flashinfer) + input_ref = input.clone() + output_ref = torch.empty_like(input) + 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]) diff --git a/test/registered/jit/test_rope.py b/test/registered/jit/test_rope.py index b70eaa1a9..170739712 100644 --- a/test/registered/jit/test_rope.py +++ b/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 - pass + """Pure-PyTorch RoPE reference (in place), used as the ROCm fallback.""" + 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)