diff --git a/.github/workflows/pr-test-amd-rocm720.yml b/.github/workflows/pr-test-amd-rocm720.yml index 871aee1f1..0626d8451 100644 --- a/.github/workflows/pr-test-amd-rocm720.yml +++ b/.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, diff --git a/.github/workflows/pr-test-amd.yml b/.github/workflows/pr-test-amd.yml index 3ed01e7c4..b2162bdde 100644 --- a/.github/workflows/pr-test-amd.yml +++ b/.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, diff --git a/test/registered/jit/benchmark/bench_activation.py b/test/registered/jit/benchmark/bench_activation.py index a46c9a3e9..773b8581f 100644 --- a/test/registered/jit/benchmark/bench_activation.py +++ b/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 diff --git a/test/registered/jit/benchmark/bench_add_constant.py b/test/registered/jit/benchmark/bench_add_constant.py index 97d0b75df..f888533c3 100644 --- a/test/registered/jit/benchmark/bench_add_constant.py +++ b/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( diff --git a/test/registered/jit/benchmark/bench_clamp_position.py b/test/registered/jit/benchmark/bench_clamp_position.py index c2d74ef5f..cf3ca3dfe 100644 --- a/test/registered/jit/benchmark/bench_clamp_position.py +++ b/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)], diff --git a/test/registered/jit/benchmark/bench_dsv3_router_gemm.py b/test/registered/jit/benchmark/bench_dsv3_router_gemm.py index a6df11bdd..91018638c 100644 --- a/test/registered/jit/benchmark/bench_dsv3_router_gemm.py +++ b/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 diff --git a/test/registered/jit/benchmark/bench_fused_eh_norm.py b/test/registered/jit/benchmark/bench_fused_eh_norm.py index 9744c0750..e000efc28 100644 --- a/test/registered/jit/benchmark/bench_fused_eh_norm.py +++ b/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 diff --git a/test/registered/jit/benchmark/bench_hicache.py b/test/registered/jit/benchmark/bench_hicache.py index 6811d8b96..10f261ac3 100644 --- a/test/registered/jit/benchmark/bench_hicache.py +++ b/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 diff --git a/test/registered/jit/benchmark/bench_hisparse.py b/test/registered/jit/benchmark/bench_hisparse.py index 18fecab23..6d99729d0 100644 --- a/test/registered/jit/benchmark/bench_hisparse.py +++ b/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 diff --git a/test/registered/jit/benchmark/bench_mla_kv_pack_quantize_fp8.py b/test/registered/jit/benchmark/bench_mla_kv_pack_quantize_fp8.py index 606a1ec56..aa3f4080d 100644 --- a/test/registered/jit/benchmark/bench_mla_kv_pack_quantize_fp8.py +++ b/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 diff --git a/test/registered/jit/benchmark/bench_ngram_compute_decode.py b/test/registered/jit/benchmark/bench_ngram_compute_decode.py index 902d8df09..8801ab0a6 100644 --- a/test/registered/jit/benchmark/bench_ngram_compute_decode.py +++ b/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 diff --git a/test/registered/jit/benchmark/bench_ngram_update_token_table.py b/test/registered/jit/benchmark/bench_ngram_update_token_table.py index e0ab181a9..324f1c1c5 100644 --- a/test/registered/jit/benchmark/bench_ngram_update_token_table.py +++ b/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( diff --git a/test/registered/jit/benchmark/bench_online_c128_mtp.py b/test/registered/jit/benchmark/bench_online_c128_mtp.py index fdf2c7eb5..80f722f98 100644 --- a/test/registered/jit/benchmark/bench_online_c128_mtp.py +++ b/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 diff --git a/test/registered/jit/benchmark/bench_resolve_future_token_ids.py b/test/registered/jit/benchmark/bench_resolve_future_token_ids.py index c3c48a54f..8a88b5fb4 100644 --- a/test/registered/jit/benchmark/bench_resolve_future_token_ids.py +++ b/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 diff --git a/test/registered/jit/benchmark/bench_store_cache.py b/test/registered/jit/benchmark/bench_store_cache.py index 3e2eb80dd..4df15acee 100644 --- a/test/registered/jit/benchmark/bench_store_cache.py +++ b/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() diff --git a/test/registered/jit/benchmark/diffusion/bench_qwen_image_modulation.py b/test/registered/jit/benchmark/diffusion/bench_qwen_image_modulation.py index e614624a6..f155a679f 100644 --- a/test/registered/jit/benchmark/diffusion/bench_qwen_image_modulation.py +++ b/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] diff --git a/test/registered/jit/benchmark/minimax/bench_minimax_decode_topk.py b/test/registered/jit/benchmark/minimax/bench_minimax_decode_topk.py index 069b24d48..ec98c465c 100644 --- a/test/registered/jit/benchmark/minimax/bench_minimax_decode_topk.py +++ b/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 diff --git a/test/registered/jit/benchmark/minimax/bench_minimax_store_kv_index.py b/test/registered/jit/benchmark/minimax/bench_minimax_store_kv_index.py index d5f5a6032..5989147b1 100644 --- a/test/registered/jit/benchmark/minimax/bench_minimax_store_kv_index.py +++ b/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 diff --git a/test/run_suite.py b/test/run_suite.py index 4619bd73e..0fb50ff52 100644 --- a/test/run_suite.py +++ b/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",