[AMD] add dedicated jit-kernel-benchmark-test-amd stage + register portable JIT benches (#30307)

.github/workflows/pr-test-amd-rocm720.yml

@@ -39,6 +39,7 @@ on:
           - sgl-kernel-unit-test-2-gpu-amd-rocm720
           - stage-a-test-1-gpu-small-amd-rocm720
           - jit-kernel-unit-test-amd-rocm720
+          - jit-kernel-benchmark-test-amd-rocm720
           - stage-b-test-1-gpu-small-amd-rocm720
           - stage-b-test-1-gpu-small-amd-nondeterministic-rocm720
           - stage-b-test-1-gpu-small-amd-mi35x-rocm720
@@ -394,6 +395,44 @@ jobs:
         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' || '' }}
 
+  jit-kernel-benchmark-test-amd-rocm720:
+    needs: [check-changes]
+    if: |
+      always() &&
+      (
+        (contains(format(',{0},', inputs.target_stage || inputs.target_stage_select), ',jit-kernel-benchmark-test-amd-rocm720,')) ||
+        (
+          !(inputs.target_stage || inputs.target_stage_select) &&
+          needs.check-changes.outputs.jit_kernel == 'true'
+        )
+      )
+    strategy:
+      fail-fast: false
+      matrix:
+        runner: [linux-mi325-1gpu-sglang]
+    runs-on: ${{matrix.runner}}
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v4
+        with:
+          ref: ${{ inputs.pr_head_sha || inputs.ref || github.sha }}
+
+      - name: Ensure VRAM is clear
+        run: bash scripts/ci/amd/ensure_vram_clear.sh rocm
+
+      - name: Start CI container
+        run: bash scripts/ci/amd/amd_ci_start_container.sh --rocm-version rocm720
+        env:
+          GITHUB_WORKSPACE: ${{ github.workspace }}
+
+      - name: Install dependencies
+        run: |
+          bash scripts/ci/amd/amd_ci_install_dependency.sh
+      - name: Run JIT kernel benchmarks
+        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-benchmark-test-amd ${{ inputs.continue_on_error && '--continue-on-error' || '' }}
+
   stage-b-test-1-gpu-small-amd-rocm720:
     needs: [check-changes]
     if: |
@@ -1322,6 +1361,7 @@ jobs:
 
         stage-a-test-1-gpu-small-amd-rocm720,
         jit-kernel-unit-test-amd-rocm720,
+        jit-kernel-benchmark-test-amd-rocm720,
         stage-b-test-1-gpu-small-amd-rocm720,
         stage-b-test-1-gpu-small-amd-nondeterministic-rocm720,
         stage-b-test-1-gpu-small-amd-mi35x-rocm720,

.github/workflows/pr-test-amd.yml

@@ -27,6 +27,7 @@ on:
           - sgl-kernel-unit-test-2-gpu-amd
           - stage-a-test-1-gpu-small-amd
           - jit-kernel-unit-test-amd
+          - jit-kernel-benchmark-test-amd
           - stage-b-test-1-gpu-small-amd
           - stage-b-test-1-gpu-small-amd-nondeterministic
           - stage-b-test-1-gpu-small-amd-mi35x
@@ -391,6 +392,43 @@ jobs:
         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' || '' }}
 
+  jit-kernel-benchmark-test-amd:
+    name: ${{ format('jit-kernel-benchmark-test-amd (linux-{0}-1gpu-sglang)', inputs.runner_arch || 'mi325') }}
+    needs: [check-changes, call-gate]
+    if: |
+      always() && !cancelled() &&
+      (
+        (contains(format(',{0},', inputs.target_stage || inputs.target_stage_select), ',jit-kernel-benchmark-test-amd,')) ||
+        (
+          !(inputs.target_stage || inputs.target_stage_select) &&
+          (needs.call-gate.result == 'success' || needs.call-gate.result == 'skipped') &&
+          needs.check-changes.outputs.jit_kernel == 'true'
+        )
+      )
+    runs-on: ${{ format('linux-{0}-1gpu-sglang', inputs.runner_arch || 'mi325') }}
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v4
+        with:
+          ref: ${{ inputs.pr_head_sha || inputs.ref || github.sha }}
+
+      - name: Ensure VRAM is clear
+        run: bash scripts/ci/amd/ensure_vram_clear.sh rocm
+
+      - name: Start CI container
+        run: bash scripts/ci/amd/amd_ci_start_container.sh
+        env:
+          GITHUB_WORKSPACE: ${{ github.workspace }}
+
+      - name: Install dependencies
+        run: |
+          bash scripts/ci/amd/amd_ci_install_dependency.sh
+
+      - name: Run JIT kernel benchmarks
+        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-benchmark-test-amd ${{ needs.check-changes.outputs.continue_on_error == 'true' && '--continue-on-error' || '' }}
+
   # =============================================== Wait Jobs for Sequential PR Execution ====================================================
   # These jobs poll GitHub API to wait for previous stages to complete.
   # For PR runs: wait jobs run and enforce sequential execution via polling.
@@ -1270,6 +1308,7 @@ jobs:
         wait-for-stage-a-amd,
         stage-a-test-1-gpu-small-amd,
         jit-kernel-unit-test-amd,
+        jit-kernel-benchmark-test-amd,
         wait-for-stage-b-amd,
         stage-b-test-1-gpu-small-amd,
         stage-b-test-1-gpu-small-amd-nondeterministic,

test/registered/jit/benchmark/bench_activation.py

@@ -10,11 +10,12 @@ from sglang.jit_kernel.activation import relu2 as relu2_jit
 from sglang.jit_kernel.activation import silu_and_mul as silu_and_mul_jit
 from sglang.jit_kernel.benchmark import marker
 from sglang.jit_kernel.benchmark.utils import create_random
-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, stage="base-b-kernel-benchmark", runner_config="1-gpu-large"
 )
+register_amd_ci(est_time=30, stage="jit-kernel-benchmark", runner_config="amd")
 
 
 @torch.compile

test/registered/jit/benchmark/bench_add_constant.py

@@ -8,11 +8,12 @@ from sglang.jit_kernel.benchmark.utils import (
     get_benchmark_range,
     run_benchmark_no_cudagraph,
 )
-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=15, stage="base-b-kernel-benchmark", runner_config="1-gpu-large"
 )
+register_amd_ci(est_time=15, stage="jit-kernel-benchmark", runner_config="amd")
 
 CONSTANT = 7
 SIZE_LIST = get_benchmark_range(

test/registered/jit/benchmark/bench_clamp_position.py

@@ -16,7 +16,7 @@ from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 register_cuda_ci(
     est_time=13, stage="base-b-kernel-benchmark", runner_config="1-gpu-large"
 )
-register_amd_ci(est_time=16, stage="jit-kernel-unit", runner_config="amd")
+register_amd_ci(est_time=16, stage="jit-kernel-benchmark", runner_config="amd")
 
 SIZE_LIST = get_benchmark_range(
     full_range=[2**n for n in range(4, 16)],

test/registered/jit/benchmark/bench_dsv3_router_gemm.py

@@ -16,11 +16,12 @@ from sglang.jit_kernel.benchmark import marker
 from sglang.jit_kernel.benchmark.utils import create_random
 from sglang.jit_kernel.dsv3_router_gemm import dsv3_router_gemm
 from sglang.jit_kernel.utils import get_jit_cuda_arch, is_hip_runtime
-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=5, stage="base-b-kernel-benchmark", runner_config="1-gpu-large"
 )
+register_amd_ci(est_time=5, stage="jit-kernel-benchmark", runner_config="amd")
 
 # sgl_kernel AOT kernel is specialized for hidden_dim=7168 only.
 SGL_KERNEL_HIDDEN_DIM = 7168

test/registered/jit/benchmark/bench_fused_eh_norm.py

@@ -4,11 +4,12 @@ import torch
 
 from sglang.jit_kernel.benchmark import marker
 from sglang.jit_kernel.fused_eh_norm import fused_eh_norm
-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=6, stage="base-b-kernel-benchmark", runner_config="1-gpu-large"
 )
+register_amd_ci(est_time=6, stage="jit-kernel-benchmark", runner_config="amd")
 
 EPS = 1e-6
 

test/registered/jit/benchmark/bench_hicache.py

@@ -31,11 +31,12 @@ from sglang.jit_kernel.hicache import (
     transfer_hicache_all_layer,
     transfer_hicache_one_layer,
 )
-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=29, stage="base-b-kernel-benchmark", runner_config="1-gpu-large"
 )
+register_amd_ci(est_time=29, stage="jit-kernel-benchmark", runner_config="amd")
 
 DISABLE_TORCH = os.environ.get("DISABLE_TORCH", "0") == "1"
 PAGE_SIZE = 1

test/registered/jit/benchmark/bench_hisparse.py

@@ -7,11 +7,12 @@ import triton.testing
 
 from sglang.jit_kernel.benchmark.utils import DEFAULT_DEVICE, DEFAULT_DTYPE
 from sglang.jit_kernel.hisparse import load_cache_to_device_buffer_mla
-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, stage="base-b-kernel-benchmark", runner_config="1-gpu-large"
 )
+register_amd_ci(est_time=12, stage="jit-kernel-benchmark", runner_config="amd")
 
 DEVICE = DEFAULT_DEVICE
 DTYPE = DEFAULT_DTYPE

test/registered/jit/benchmark/bench_mla_kv_pack_quantize_fp8.py

@@ -18,11 +18,12 @@ from sglang.jit_kernel.mla_kv_pack_quantize_fp8 import (
     mla_kv_pack_quantize_fp8 as hybrid_pack,
 )
 from sglang.jit_kernel.utils import is_arch_support_pdl
-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=15, stage="base-b-kernel-benchmark", runner_config="1-gpu-large"
 )
+register_amd_ci(est_time=15, stage="jit-kernel-benchmark", runner_config="amd")
 
 
 @triton.jit

test/registered/jit/benchmark/bench_ngram_compute_decode.py

@@ -11,11 +11,12 @@ from sglang.jit_kernel.ngram_embedding import (
     compute_n_gram_ids,
     compute_n_gram_ids_decode,
 )
-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=15, stage="base-b-kernel-benchmark", runner_config="1-gpu-large"
 )
+register_amd_ci(est_time=15, stage="jit-kernel-benchmark", runner_config="amd")
 
 NE_N = 8
 NE_K = 2

test/registered/jit/benchmark/bench_ngram_update_token_table.py

@@ -11,11 +11,12 @@ from sglang.jit_kernel.ngram_embedding import (
     update_token_table,
     update_token_table_decode,
 )
-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=15, stage="base-b-kernel-benchmark", runner_config="1-gpu-large"
 )
+register_amd_ci(est_time=15, stage="jit-kernel-benchmark", runner_config="amd")
 
 MAX_CONTEXT_LEN = 4096
 BATCH_SIZE_LIST = get_benchmark_range(

test/registered/jit/benchmark/bench_online_c128_mtp.py

@@ -16,11 +16,12 @@ from sglang.jit_kernel.benchmark.utils import (
     run_benchmark_no_cudagraph,
 )
 from sglang.jit_kernel.dsv4.online_c128_mtp import _jit_online_c128_mtp_module
-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=10, stage="base-b-kernel-benchmark", runner_config="1-gpu-large"
 )
+register_amd_ci(est_time=10, stage="jit-kernel-benchmark", runner_config="amd")
 
 HEAD_DIM = 512
 STATE_DIM = HEAD_DIM * 3

test/registered/jit/benchmark/bench_resolve_future_token_ids.py

@@ -16,7 +16,7 @@ from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 register_cuda_ci(
     est_time=10, stage="base-b-kernel-benchmark", runner_config="1-gpu-large"
 )
-register_amd_ci(est_time=10, stage="jit-kernel-unit", runner_config="amd")
+register_amd_ci(est_time=10, stage="jit-kernel-benchmark", runner_config="amd")
 
 SIZE_LIST = get_benchmark_range(
     full_range=[2**n for n in range(4, 16)],  # 16 … 32K elements

test/registered/jit/benchmark/bench_store_cache.py

@@ -7,11 +7,12 @@ from sglang.jit_kernel.benchmark.utils import (
     create_random,
 )
 from sglang.jit_kernel.kvcache import store_cache
-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, stage="base-b-kernel-benchmark", runner_config="1-gpu-large"
 )
+register_amd_ci(est_time=9, stage="jit-kernel-benchmark", runner_config="amd")
 
 
 @torch.compile()

test/registered/jit/benchmark/diffusion/bench_qwen_image_modulation.py

@@ -9,12 +9,13 @@ from sglang.jit_kernel.diffusion.triton.scale_shift import (
     fuse_layernorm_scale_shift_gate_select01_kernel,
     fuse_residual_layernorm_scale_shift_gate_select01_kernel,
 )
-from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 from sglang.utils import is_in_ci
 
 register_cuda_ci(
     est_time=13, stage="base-b-kernel-benchmark", runner_config="1-gpu-large"
 )
+register_amd_ci(est_time=13, stage="jit-kernel-benchmark", runner_config="amd")
 
 if is_in_ci():
     B_RANGE, S_RANGE, D_RANGE = [1], [128], [3072]

test/registered/jit/benchmark/minimax/bench_minimax_decode_topk.py

@@ -15,11 +15,12 @@ from sglang.srt.layers.attention.minimax_sparse_ops.decode.flash_with_topk_idx i
     _topk_index_merge_kernel,
     _topk_index_partial_kernel,
 )
-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=8, stage="base-b-kernel-benchmark", runner_config="1-gpu-large"
 )
+register_amd_ci(est_time=8, stage="jit-kernel-benchmark", runner_config="amd")
 
 BLOCK_SIZE = 128
 TOPK = 16

test/registered/jit/benchmark/minimax/bench_minimax_store_kv_index.py

@@ -5,11 +5,12 @@ import torch
 
 from sglang.jit_kernel.benchmark import marker
 from sglang.jit_kernel.minimax_store_kv_index import store_kv_index
-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=6, stage="base-b-kernel-benchmark", runner_config="1-gpu-large"
 )
+register_amd_ci(est_time=6, stage="jit-kernel-benchmark", runner_config="amd")
 
 HEAD_DIM = 128
 NUM_KV_HEADS = 1

test/run_suite.py

@@ -44,6 +44,7 @@ PER_COMMIT_SUITES = {
         "stage-b-test-1-gpu-large-amd",
         "stage-b-test-2-gpu-large-amd",
         "jit-kernel-unit-test-amd",
+        "jit-kernel-benchmark-test-amd",
         "sgl-kernel-unit-test-2-gpu-amd",
         "stage-c-test-4-gpu-amd",
         "stage-c-test-large-8-gpu-amd",