From 0e5a66dca400081066970f2789de6dd9f4e3f8d6 Mon Sep 17 00:00:00 2001 From: Michael <13900043+michaelzhang-ai@users.noreply.github.com> Date: Wed, 17 Jun 2026 23:28:56 -0700 Subject: [PATCH] [AMD] Register 3 JIT kernel unit tests for AMD CI (#27837) Co-authored-by: michaelzhang-ai --- .github/workflows/pr-test-amd-rocm720.yml | 6 +++--- .github/workflows/pr-test-amd.yml | 6 +++--- test/registered/jit/test_clamp_position.py | 3 ++- test/registered/jit/test_resolve_future_token_ids.py | 3 ++- test/registered/jit/test_rmsnorm_hf.py | 3 ++- 5 files changed, 12 insertions(+), 9 deletions(-) diff --git a/.github/workflows/pr-test-amd-rocm720.yml b/.github/workflows/pr-test-amd-rocm720.yml index 6174ee6d1..dd1c7f12e 100644 --- a/.github/workflows/pr-test-amd-rocm720.yml +++ b/.github/workflows/pr-test-amd-rocm720.yml @@ -268,7 +268,7 @@ jobs: run: | bash scripts/ci/amd/amd_ci_install_dependency.sh - name: Run test - timeout-minutes: 14 + timeout-minutes: 30 run: | docker exec -w /sglang-checkout/sgl-kernel/tests ci_sglang python3 -m pytest test_moe_align.py docker exec -w /sglang-checkout/sgl-kernel/tests ci_sglang python3 -m pytest test_moe_topk_softmax.py @@ -314,7 +314,7 @@ jobs: run: | bash scripts/ci/amd/amd_ci_install_dependency.sh - name: Run test - timeout-minutes: 20 + timeout-minutes: 30 run: | bash scripts/ci/amd/amd_ci_exec.sh -w "/sglang-checkout/test" python3 run_suite.py --hw amd --suite sgl-kernel-unit-test-2-gpu-amd @@ -393,7 +393,7 @@ jobs: run: | bash scripts/ci/amd/amd_ci_install_dependency.sh - name: Run JIT kernel unit tests - timeout-minutes: 10 + timeout-minutes: 30 run: | bash scripts/ci/amd/amd_ci_exec.sh -w "/sglang-checkout/test" python3 run_suite.py --hw amd --suite jit-kernel-unit-test-amd ${{ inputs.continue_on_error && '--continue-on-error' || '' }} diff --git a/.github/workflows/pr-test-amd.yml b/.github/workflows/pr-test-amd.yml index 0163201ee..a08d00967 100644 --- a/.github/workflows/pr-test-amd.yml +++ b/.github/workflows/pr-test-amd.yml @@ -250,7 +250,7 @@ jobs: bash scripts/ci/amd/amd_ci_install_dependency.sh - name: Run test - timeout-minutes: 14 + timeout-minutes: 30 env: CONTINUE_ON_ERROR: ${{ needs.check-changes.outputs.continue_on_error }} run: | @@ -308,7 +308,7 @@ jobs: bash scripts/ci/amd/amd_ci_install_dependency.sh - name: Run test - timeout-minutes: 20 + timeout-minutes: 30 env: CONTINUE_ON_ERROR: ${{ needs.check-changes.outputs.continue_on_error }} run: | @@ -386,7 +386,7 @@ jobs: bash scripts/ci/amd/amd_ci_install_dependency.sh - name: Run JIT kernel unit tests - timeout-minutes: 10 + timeout-minutes: 30 run: | bash scripts/ci/amd/amd_ci_exec.sh -w "/sglang-checkout/test" python3 run_suite.py --hw amd --suite jit-kernel-unit-test-amd ${{ needs.check-changes.outputs.continue_on_error == 'true' && '--continue-on-error' || '' }} diff --git a/test/registered/jit/test_clamp_position.py b/test/registered/jit/test_clamp_position.py index 4e8cc1e4f..846d29017 100644 --- a/test/registered/jit/test_clamp_position.py +++ b/test/registered/jit/test_clamp_position.py @@ -4,10 +4,11 @@ import pytest import torch from sglang.jit_kernel.clamp_position import clamp_position_cuda -from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci register_cuda_ci(est_time=12, suite="base-b-kernel-unit-1-gpu-large") register_cuda_ci(est_time=120, suite="nightly-kernel-1-gpu", nightly=True) +register_amd_ci(est_time=12, suite="jit-kernel-unit-test-amd") def _reference_clamp_position(seq_lens): diff --git a/test/registered/jit/test_resolve_future_token_ids.py b/test/registered/jit/test_resolve_future_token_ids.py index 062625ce2..760e9fbd9 100644 --- a/test/registered/jit/test_resolve_future_token_ids.py +++ b/test/registered/jit/test_resolve_future_token_ids.py @@ -4,10 +4,11 @@ import pytest import torch from sglang.jit_kernel.resolve_future_token_ids import resolve_future_token_ids_cuda -from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci register_cuda_ci(est_time=9, suite="base-b-kernel-unit-1-gpu-large") register_cuda_ci(est_time=120, suite="nightly-kernel-1-gpu", nightly=True) +register_amd_ci(est_time=9, suite="jit-kernel-unit-test-amd") def _reference_resolve(input_ids, future_map): diff --git a/test/registered/jit/test_rmsnorm_hf.py b/test/registered/jit/test_rmsnorm_hf.py index 02162172f..c26cdf0d8 100644 --- a/test/registered/jit/test_rmsnorm_hf.py +++ b/test/registered/jit/test_rmsnorm_hf.py @@ -11,10 +11,11 @@ from sglang.jit_kernel.rmsnorm_hf import ( rmsnorm_hf, ) from sglang.jit_kernel.utils import get_ci_test_range -from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci register_cuda_ci(est_time=30, suite="base-b-kernel-unit-1-gpu-large") register_cuda_ci(est_time=120, suite="nightly-kernel-1-gpu", nightly=True) +register_amd_ci(est_time=30, suite="jit-kernel-unit-test-amd") EPS = 1e-5 DEVICE = "cuda"