Split pr-test.yml: extract sgl-kernel, jit-kernel, and multimodal-gen tests into separate workflow files (#21219)

.github/workflows/pr-test-jit-kernel.yml

@@ -0,0 +1,133 @@
+name: PR Test - JIT Kernel
+
+on:
+  workflow_call:
+    inputs:
+      jit_kernel:
+        required: true
+        type: string
+      pr_head_sha:
+        required: false
+        type: string
+        default: ''
+      git_ref:
+        required: false
+        type: string
+        default: ''
+      target_stage:
+        required: false
+        type: string
+        default: ''
+      test_parallel_dispatch:
+        required: false
+        type: string
+        default: 'false'
+
+# Workflow-level env is NOT inherited from the caller in reusable workflows (verified by CI test).
+# The github context (including github.event_name) IS inherited from the caller.
+env:
+  SGLANG_IS_IN_CI: true
+  SGLANG_CUDA_COREDUMP: "1"
+  SGLANG_JIT_DEEPGEMM_FAST_WARMUP: true
+  SGLANG_PR_TEST_BYPASS_MAINTENANCE_ON_MAIN: ${{ github.ref == 'refs/heads/main' && 'true' || 'false' }}
+
+jobs:
+  jit-kernel-unit-test:
+    if: |
+      github.event_name != 'schedule' &&
+      inputs.test_parallel_dispatch != 'true' &&
+      !inputs.target_stage
+    runs-on: 1-gpu-h100
+    timeout-minutes: 240
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          ref: ${{ inputs.pr_head_sha || inputs.git_ref || github.sha }}
+
+      - uses: ./.github/actions/check-maintenance
+        with:
+          github-token: ${{ github.token }}
+
+      - name: Install dependencies
+        timeout-minutes: 20
+        run: |
+          bash scripts/ci/cuda/ci_install_dependency.sh diffusion
+
+      - name: Run test
+        timeout-minutes: 30
+        run: |
+          cd python/sglang/jit_kernel
+          pytest tests/
+
+  jit-kernel-unit-test-nightly:
+    if: |
+      github.event_name == 'schedule' &&
+      inputs.jit_kernel == 'true'
+    runs-on: 1-gpu-h100
+    timeout-minutes: 240
+    env:
+      SGLANG_JIT_KERNEL_RUN_FULL_TESTS: "1"
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          ref: ${{ inputs.pr_head_sha || inputs.git_ref || github.sha }}
+
+      - uses: ./.github/actions/check-maintenance
+        with:
+          github-token: ${{ github.token }}
+
+      - name: Install dependencies
+        timeout-minutes: 20
+        run: |
+          bash scripts/ci/cuda/ci_install_dependency.sh
+
+      - name: Run full nightly test
+        timeout-minutes: 60
+        run: |
+          cd python/sglang/jit_kernel
+          pytest tests/
+
+  jit-kernel-benchmark-test:
+    if: |
+      github.event_name != 'schedule' &&
+      inputs.test_parallel_dispatch != 'true' &&
+      !inputs.target_stage
+    runs-on: 1-gpu-h100
+    timeout-minutes: 240
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          ref: ${{ inputs.pr_head_sha || inputs.git_ref || github.sha }}
+
+      - uses: ./.github/actions/check-maintenance
+        with:
+          github-token: ${{ github.token }}
+
+      - name: Install dependencies
+        timeout-minutes: 20
+        run: |
+          bash scripts/ci/cuda/ci_install_dependency.sh diffusion
+
+      - name: Run benchmark tests
+        timeout-minutes: 45
+        run: |
+          cd python/sglang/jit_kernel/benchmark
+          echo "Running jit-kernel benchmark tests in CI mode..."
+
+          failures=()
+
+          for bench_file in bench_*.py; do
+            echo "Testing $bench_file..."
+            if ! timeout 120 python3 "$bench_file"; then
+              failures+=("$bench_file")
+            fi
+            echo "Completed $bench_file"
+            echo "---"
+          done
+
+          if [ ${#failures[@]} -ne 0 ]; then
+            echo "The following benchmark tests failed: ${failures[*]}"
+            exit 1
+          fi
+
+          echo "All jit-kernel benchmark tests completed successfully!"

.github/workflows/pr-test-multimodal-gen.yml

@@ -0,0 +1,191 @@
+name: PR Test - Multimodal Gen
+
+on:
+  workflow_call:
+    inputs:
+      multimodal_gen:
+        required: true
+        type: string
+      sgl_kernel:
+        required: true
+        type: string
+      continue_on_error:
+        required: false
+        type: string
+        default: 'false'
+      pr_head_sha:
+        required: false
+        type: string
+        default: ''
+      git_ref:
+        required: false
+        type: string
+        default: ''
+      target_stage:
+        required: false
+        type: string
+        default: ''
+      test_parallel_dispatch:
+        required: false
+        type: string
+        default: 'false'
+      caller_needs_failure:
+        required: false
+        type: string
+        default: 'false'
+
+# Workflow-level env is NOT inherited from the caller in reusable workflows.
+# The github context (including github.event_name) IS inherited from the caller.
+env:
+  SGLANG_IS_IN_CI: true
+  SGLANG_CUDA_COREDUMP: "1"
+  SGLANG_PR_TEST_BYPASS_MAINTENANCE_ON_MAIN: ${{ github.ref == 'refs/heads/main' && 'true' || 'false' }}
+
+jobs:
+  multimodal-gen-test-1-gpu:
+    if: |
+      (inputs.target_stage == 'multimodal-gen-test-1-gpu') ||
+      (
+        !inputs.target_stage &&
+        ((github.event_name == 'schedule' || inputs.test_parallel_dispatch == 'true') || (inputs.caller_needs_failure != 'true' && !cancelled())) &&
+        inputs.multimodal_gen == 'true'
+      )
+    runs-on: 1-gpu-h100
+    timeout-minutes: 240
+    strategy:
+      fail-fast: false
+      matrix:
+        part: [0, 1]
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v4
+        with:
+          ref: ${{ inputs.pr_head_sha || inputs.git_ref || github.sha }}
+
+      - uses: ./.github/actions/check-maintenance
+        with:
+          github-token: ${{ github.token }}
+
+      - name: Download artifacts
+        if: inputs.sgl_kernel == 'true'
+        uses: actions/download-artifact@v4
+        with:
+          path: sgl-kernel/dist/
+          merge-multiple: true
+          pattern: wheel-python3.10-cuda12.9
+
+      - name: Install dependencies
+        timeout-minutes: 20
+        run: |
+          CUSTOM_BUILD_SGL_KERNEL=${{inputs.sgl_kernel}} bash scripts/ci/cuda/ci_install_dependency.sh diffusion
+      - name: Run diffusion server tests
+        timeout-minutes: 240
+        env:
+          RUNAI_STREAMER_MEMORY_LIMIT: 0
+          CONTINUE_ON_ERROR_FLAG: ${{ inputs.continue_on_error == 'true' && '--continue-on-error' || '' }}
+        run: |
+          cd python
+          python3 sglang/multimodal_gen/test/run_suite.py \
+            --suite 1-gpu \
+            --partition-id ${{ matrix.part }} \
+            --total-partitions 2 \
+            $CONTINUE_ON_ERROR_FLAG
+
+      - uses: ./.github/actions/upload-cuda-coredumps
+        if: always()
+        with:
+          artifact-suffix: ${{ matrix.part }}
+
+  multimodal-gen-test-2-gpu:
+    if: |
+      (inputs.target_stage == 'multimodal-gen-test-2-gpu') ||
+      (
+        !inputs.target_stage &&
+        ((github.event_name == 'schedule' || inputs.test_parallel_dispatch == 'true') || (inputs.caller_needs_failure != 'true' && !cancelled())) &&
+        inputs.multimodal_gen == 'true'
+      )
+    runs-on: 2-gpu-h100
+    timeout-minutes: 240
+    strategy:
+      fail-fast: false
+      matrix:
+        part: [0, 1]
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v4
+        with:
+          ref: ${{ inputs.pr_head_sha || inputs.git_ref || github.sha }}
+
+      - uses: ./.github/actions/check-maintenance
+        with:
+          github-token: ${{ github.token }}
+
+      - name: Download artifacts
+        if: inputs.sgl_kernel == 'true'
+        uses: actions/download-artifact@v4
+        with:
+          path: sgl-kernel/dist/
+          merge-multiple: true
+          pattern: wheel-python3.10-cuda12.9
+
+      - name: Install dependencies
+        timeout-minutes: 20
+        run: |
+          CUSTOM_BUILD_SGL_KERNEL=${{inputs.sgl_kernel}} bash scripts/ci/cuda/ci_install_dependency.sh diffusion
+
+      - name: Run diffusion server tests
+        timeout-minutes: 240
+        env:
+          RUNAI_STREAMER_MEMORY_LIMIT: 0
+          CONTINUE_ON_ERROR_FLAG: ${{ inputs.continue_on_error == 'true' && '--continue-on-error' || '' }}
+        run: |
+          cd python
+          python3 sglang/multimodal_gen/test/run_suite.py \
+            --suite 2-gpu \
+            --partition-id ${{ matrix.part }} \
+            --total-partitions 2 \
+            $CONTINUE_ON_ERROR_FLAG
+
+      - uses: ./.github/actions/upload-cuda-coredumps
+        if: always()
+        with:
+          artifact-suffix: ${{ matrix.part }}
+
+  multimodal-gen-unit-test:
+    if: |
+      (inputs.target_stage == 'multimodal-gen-unit-test') ||
+      (
+        !inputs.target_stage &&
+        ((github.event_name == 'schedule' || inputs.test_parallel_dispatch == 'true') || (inputs.caller_needs_failure != 'true' && !cancelled())) &&
+        inputs.multimodal_gen == 'true'
+      )
+    runs-on: 1-gpu-h100
+    timeout-minutes: 120
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v4
+        with:
+          ref: ${{ inputs.pr_head_sha || inputs.git_ref || github.sha }}
+
+      - uses: ./.github/actions/check-maintenance
+        with:
+          github-token: ${{ github.token }}
+
+      - name: Download artifacts
+        if: inputs.sgl_kernel == 'true'
+        uses: actions/download-artifact@v4
+        with:
+          path: sgl-kernel/dist/
+          merge-multiple: true
+          pattern: wheel-python3.10-cuda12.9
+
+      - name: Install dependencies
+        timeout-minutes: 20
+        run: |
+          CUSTOM_BUILD_SGL_KERNEL=${{inputs.sgl_kernel}} bash scripts/ci/cuda/ci_install_dependency.sh diffusion
+
+      - name: Run diffusion unit tests
+        timeout-minutes: 60
+        run: |
+          cd python
+          python3 sglang/multimodal_gen/test/run_suite.py --suite unit

.github/workflows/pr-test-sgl-kernel.yml

@@ -0,0 +1,209 @@
+name: PR Test - SGL Kernel
+
+on:
+  workflow_call:
+    inputs:
+      sgl_kernel:
+        required: true
+        type: string
+      b200_runner:
+        required: true
+        type: string
+      pr_head_sha:
+        required: false
+        type: string
+        default: ''
+      git_ref:
+        required: false
+        type: string
+        default: ''
+
+# Workflow-level env is NOT inherited from the caller in reusable workflows.
+# The github context (including github.event_name) IS inherited from the caller.
+env:
+  SGLANG_IS_IN_CI: true
+  SGLANG_CUDA_COREDUMP: "1"
+  SGLANG_PR_TEST_BYPASS_MAINTENANCE_ON_MAIN: ${{ github.ref == 'refs/heads/main' && 'true' || 'false' }}
+
+jobs:
+  sgl-kernel-unit-test:
+    runs-on: 1-gpu-h100
+    timeout-minutes: 240
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          ref: ${{ inputs.pr_head_sha || inputs.git_ref || github.sha }}
+
+      - uses: ./.github/actions/check-maintenance
+        with:
+          github-token: ${{ github.token }}
+
+      - name: Cleanup
+        run: |
+          ls -alh sgl-kernel/dist || true
+          rm -rf sgl-kernel/dist/* || true
+
+      - name: Download artifacts
+        uses: actions/download-artifact@v4
+        with:
+          path: sgl-kernel/dist/
+          merge-multiple: true
+          pattern: wheel-python3.10-cuda12.9
+
+      - name: Install dependencies
+        timeout-minutes: 20
+        run: |
+          CUSTOM_BUILD_SGL_KERNEL=${{inputs.sgl_kernel}} bash scripts/ci/cuda/ci_install_dependency.sh diffusion
+
+      - name: Run test
+        timeout-minutes: 30
+        run: |
+          cd sgl-kernel
+          pytest tests/
+
+  sgl-kernel-mla-test:
+    runs-on: 1-gpu-h100
+    timeout-minutes: 240
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          ref: ${{ inputs.pr_head_sha || inputs.git_ref || github.sha }}
+
+      - uses: ./.github/actions/check-maintenance
+        with:
+          github-token: ${{ github.token }}
+
+      - name: Cleanup
+        run: |
+          ls -alh sgl-kernel/dist || true
+          rm -rf sgl-kernel/dist/* || true
+
+      - name: Download artifacts
+        uses: actions/download-artifact@v4
+        with:
+          path: sgl-kernel/dist/
+          merge-multiple: true
+          pattern: wheel-python3.10-cuda12.9
+
+      - name: Install dependencies
+        timeout-minutes: 20
+        run: |
+          CUSTOM_BUILD_SGL_KERNEL=${{inputs.sgl_kernel}} bash scripts/ci/cuda/ci_install_dependency.sh
+
+      - name: Run test
+        timeout-minutes: 30
+        run: |
+          cd test/registered/mla
+          python3 test_mla_deepseek_v3.py
+
+  sgl-kernel-benchmark-test:
+    runs-on: 1-gpu-h100
+    timeout-minutes: 240
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          ref: ${{ inputs.pr_head_sha || inputs.git_ref || github.sha }}
+
+      - uses: ./.github/actions/check-maintenance
+        with:
+          github-token: ${{ github.token }}
+
+      - name: Cleanup
+        run: |
+          ls -alh sgl-kernel/dist || true
+          rm -rf sgl-kernel/dist/* || true
+
+      - name: Download artifacts
+        uses: actions/download-artifact@v4
+        with:
+          path: sgl-kernel/dist/
+          merge-multiple: true
+          pattern: wheel-python3.10-cuda12.9
+
+      - name: Install dependencies
+        timeout-minutes: 20
+        run: |
+          CUSTOM_BUILD_SGL_KERNEL=${{inputs.sgl_kernel}} bash scripts/ci/cuda/ci_install_dependency.sh
+
+      - name: Run benchmark tests
+        timeout-minutes: 45
+        run: |
+          cd sgl-kernel/benchmark
+          echo "Running sgl-kernel benchmark tests in CI mode..."
+
+          echo "CI environment variable: $CI"
+          echo "GITHUB_ACTIONS environment variable: $GITHUB_ACTIONS"
+
+          for bench_file in bench_*.py; do
+            echo "Testing $bench_file..."
+            timeout 60 python3 "$bench_file" || echo "Warning: $bench_file timed out or failed, continuing..."
+            echo "Completed $bench_file"
+            echo "---"
+          done
+
+          echo "All benchmark tests completed!"
+
+  sgl-kernel-b200-test:
+    runs-on: ${{ inputs.b200_runner }}
+    timeout-minutes: 240
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          ref: ${{ inputs.pr_head_sha || inputs.git_ref || github.sha }}
+
+      - uses: ./.github/actions/check-maintenance
+        with:
+          github-token: ${{ github.token }}
+
+      - name: Cleanup
+        run: |
+          ls -alh sgl-kernel/dist || true
+          rm -rf sgl-kernel/dist/* || true
+
+      - name: Download artifacts
+        uses: actions/download-artifact@v4
+        with:
+          path: sgl-kernel/dist/
+          merge-multiple: true
+          pattern: wheel-python3.10-cuda12.9
+
+      - name: Install dependencies
+        timeout-minutes: 20
+        run: |
+          CUSTOM_BUILD_SGL_KERNEL=${{inputs.sgl_kernel}} bash scripts/ci/cuda/ci_install_dependency.sh diffusion
+
+      - name: Run sgl-kernel unit tests on B200
+        timeout-minutes: 30
+        run: |
+          cd sgl-kernel
+          pytest tests/
+
+  # Adding a single CUDA13 smoke test to verify that the kernel builds and runs
+  # TODO: Add back this test when it can pass on CI
+  # cuda13-kernel-smoke-test:
+  #   if: inputs.sgl_kernel == 'true'
+  #   runs-on: x64-cu13-kernel-tests
+  #   steps:
+  #     - uses: actions/checkout@v4
+
+  #     - name: Cleanup
+  #       run: |
+  #         ls -alh sgl-kernel/dist || true
+  #         rm -rf sgl-kernel/dist/* || true
+
+  #     - name: Download CUDA 13.0 artifacts
+  #       uses: actions/download-artifact@v4
+  #       with:
+  #         path: sgl-kernel/dist/
+  #         merge-multiple: true
+  #         pattern: wheel-python3.10-cuda13.0
+
+  #     - name: Install dependencies
+  #       run: |
+  #         CUSTOM_BUILD_SGL_KERNEL=${{inputs.sgl_kernel}} bash scripts/ci/cuda/ci_install_dependency.sh
+
+  #     - name: Run kernel unit tests
+  #       timeout-minutes: 30
+  #       run: |
+  #         cd sgl-kernel
+  #         pytest tests/

.github/workflows/pr-test.yml

@@ -129,15 +129,18 @@ jobs:
               - "test/**"
             multimodal_gen:
               - ".github/workflows/pr-test.yml"
+              - ".github/workflows/pr-test-multimodal-gen.yml"
               - "python/pyproject.toml"
               - "python/sglang/multimodal_gen/**"
               - "python/sglang/jit_kernel/**"
               - "python/sglang/cli/**"
             jit_kernel:
               - ".github/workflows/pr-test.yml"
+              - ".github/workflows/pr-test-jit-kernel.yml"
               - "python/pyproject.toml"
               - "python/sglang/jit_kernel/**"
             sgl_kernel:
+              - ".github/workflows/pr-test-sgl-kernel.yml"
               - "sgl-kernel/**"
 
       # For /rerun-stage (workflow_dispatch with target_stage), dorny/paths-filter doesn't work
@@ -171,7 +174,7 @@ jobs:
           echo "..."
 
           # Check for sgl-kernel changes
-          if echo "$CHANGED_FILES" | grep -q "^sgl-kernel/"; then
+          if echo "$CHANGED_FILES" | grep -qE "^(sgl-kernel/|\.github/workflows/pr-test-sgl-kernel\.yml)"; then
             echo "sgl_kernel=true" >> $GITHUB_OUTPUT
             echo "Detected sgl-kernel changes"
           else
@@ -189,7 +192,7 @@ jobs:
           fi
 
           # Check for jit_kernel changes
-          if echo "$CHANGED_FILES" | grep -qE "^(python/sglang/jit_kernel/|python/pyproject\.toml|\.github/workflows/pr-test\.yml)"; then
+          if echo "$CHANGED_FILES" | grep -qE "^(python/sglang/jit_kernel/|python/pyproject\.toml|\.github/workflows/pr-test\.yml|\.github/workflows/pr-test-jit-kernel\.yml)"; then
             echo "jit_kernel=true" >> $GITHUB_OUTPUT
             echo "Detected jit_kernel changes"
           else
@@ -197,7 +200,7 @@ jobs:
           fi
 
           # Check for multimodal_gen changes
-          if echo "$CHANGED_FILES" | grep -qE "^(python/sglang/multimodal_gen/|python/sglang/cli/|python/pyproject\.toml|\.github/workflows/pr-test\.yml)"; then
+          if echo "$CHANGED_FILES" | grep -qE "^(python/sglang/multimodal_gen/|python/sglang/cli/|python/pyproject\.toml|\.github/workflows/pr-test\.yml|\.github/workflows/pr-test-multimodal-gen\.yml)"; then
             echo "multimodal_gen=true" >> $GITHUB_OUTPUT
             echo "Detected multimodal_gen changes"
           else
@@ -510,325 +513,34 @@ jobs:
           path: sgl-kernel/dist/*
           if-no-files-found: error
 
-  sgl-kernel-unit-test:
+  call-sgl-kernel-tests:
     needs: [check-changes, call-gate, sgl-kernel-build-wheels]
-    # Skip for scheduled runs and when target_stage is set
     if: |
       github.event_name != 'schedule' &&
       inputs.test_parallel_dispatch != true &&
       !inputs.target_stage &&
       needs.check-changes.outputs.sgl_kernel == 'true'
-    runs-on: 1-gpu-h100
+    uses: ./.github/workflows/pr-test-sgl-kernel.yml
-    timeout-minutes: 240
+    with:
-    steps:
+      sgl_kernel: ${{ needs.check-changes.outputs.sgl_kernel }}
-      - uses: actions/checkout@v4
+      b200_runner: ${{ needs.check-changes.outputs.b200_runner }}
-        with:
+      pr_head_sha: ${{ inputs.pr_head_sha || '' }}
-          ref: ${{ inputs.pr_head_sha || inputs.git_ref || github.sha }}
+      git_ref: ${{ inputs.git_ref || '' }}
-
+    secrets: inherit
-      - uses: ./.github/actions/check-maintenance
-        with:
-          github-token: ${{ github.token }}
-
-      - name: Cleanup
-        run: |
-          ls -alh sgl-kernel/dist || true
-          rm -rf sgl-kernel/dist/* || true
-
-      - name: Download artifacts
-        uses: actions/download-artifact@v4
-        with:
-          path: sgl-kernel/dist/
-          merge-multiple: true
-          pattern: wheel-python3.10-cuda12.9
-
-      - name: Install dependencies
-        timeout-minutes: 20
-        run: |
-          CUSTOM_BUILD_SGL_KERNEL=${{needs.check-changes.outputs.sgl_kernel}} bash scripts/ci/cuda/ci_install_dependency.sh diffusion
-
-      - name: Run test
-        timeout-minutes: 30
-        run: |
-          cd sgl-kernel
-          pytest tests/
-
-  sgl-kernel-mla-test:
-    needs: [check-changes, call-gate, sgl-kernel-build-wheels]
-    # Skip for scheduled runs and when target_stage is set
-    if: |
-      github.event_name != 'schedule' &&
-      inputs.test_parallel_dispatch != true &&
-      !inputs.target_stage &&
-      needs.check-changes.outputs.sgl_kernel == 'true'
-    runs-on: 1-gpu-h100
-    timeout-minutes: 240
-    steps:
-      - uses: actions/checkout@v4
-        with:
-          ref: ${{ inputs.pr_head_sha || inputs.git_ref || github.sha }}
-
-      - uses: ./.github/actions/check-maintenance
-        with:
-          github-token: ${{ github.token }}
-
-      - name: Cleanup
-        run: |
-          ls -alh sgl-kernel/dist || true
-          rm -rf sgl-kernel/dist/* || true
-
-      - name: Download artifacts
-        uses: actions/download-artifact@v4
-        with:
-          path: sgl-kernel/dist/
-          merge-multiple: true
-          pattern: wheel-python3.10-cuda12.9
-
-      - name: Install dependencies
-        timeout-minutes: 20
-        run: |
-          CUSTOM_BUILD_SGL_KERNEL=${{needs.check-changes.outputs.sgl_kernel}} bash scripts/ci/cuda/ci_install_dependency.sh
-
-      - name: Run test
-        timeout-minutes: 30
-        run: |
-          cd test/registered/mla
-          python3 test_mla_deepseek_v3.py
-
-  sgl-kernel-benchmark-test:
-    needs: [check-changes, call-gate, sgl-kernel-build-wheels]
-    # Skip for scheduled runs and when target_stage is set
-    if: |
-      github.event_name != 'schedule' &&
-      inputs.test_parallel_dispatch != true &&
-      !inputs.target_stage &&
-      needs.check-changes.outputs.sgl_kernel == 'true'
-    runs-on: 1-gpu-h100
-    timeout-minutes: 240
-    steps:
-      - uses: actions/checkout@v4
-        with:
-          ref: ${{ inputs.pr_head_sha || inputs.git_ref || github.sha }}
-
-      - uses: ./.github/actions/check-maintenance
-        with:
-          github-token: ${{ github.token }}
-
-      - name: Cleanup
-        run: |
-          ls -alh sgl-kernel/dist || true
-          rm -rf sgl-kernel/dist/* || true
-
-      - name: Download artifacts
-        uses: actions/download-artifact@v4
-        with:
-          path: sgl-kernel/dist/
-          merge-multiple: true
-          pattern: wheel-python3.10-cuda12.9
-
-      - name: Install dependencies
-        timeout-minutes: 20
-        run: |
-          CUSTOM_BUILD_SGL_KERNEL=${{needs.check-changes.outputs.sgl_kernel}} bash scripts/ci/cuda/ci_install_dependency.sh
-
-      - name: Run benchmark tests
-        timeout-minutes: 45
-        run: |
-          cd sgl-kernel/benchmark
-          echo "Running sgl-kernel benchmark tests in CI mode..."
-
-          echo "CI environment variable: $CI"
-          echo "GITHUB_ACTIONS environment variable: $GITHUB_ACTIONS"
-
-          for bench_file in bench_*.py; do
-            echo "Testing $bench_file..."
-            timeout 60 python3 "$bench_file" || echo "Warning: $bench_file timed out or failed, continuing..."
-            echo "Completed $bench_file"
-            echo "---"
-          done
-
-          echo "All benchmark tests completed!"
-
-  sgl-kernel-b200-test:
-    needs: [check-changes, sgl-kernel-build-wheels]
-    # Skip for scheduled runs and when target_stage is set
-    if: |
-      github.event_name != 'schedule' &&
-      inputs.test_parallel_dispatch != true &&
-      !inputs.target_stage &&
-      needs.check-changes.outputs.sgl_kernel == 'true'
-    runs-on: ${{ needs.check-changes.outputs.b200_runner }}
-    timeout-minutes: 240
-    steps:
-      - uses: actions/checkout@v4
-        with:
-          ref: ${{ inputs.pr_head_sha || inputs.git_ref || github.sha }}
-
-      - uses: ./.github/actions/check-maintenance
-        with:
-          github-token: ${{ github.token }}
-
-      - name: Cleanup
-        run: |
-          ls -alh sgl-kernel/dist || true
-          rm -rf sgl-kernel/dist/* || true
-
-      - name: Download artifacts
-        uses: actions/download-artifact@v4
-        with:
-          path: sgl-kernel/dist/
-          merge-multiple: true
-          pattern: wheel-python3.10-cuda12.9
-
-      - name: Install dependencies
-        timeout-minutes: 20
-        run: |
-          CUSTOM_BUILD_SGL_KERNEL=${{needs.check-changes.outputs.sgl_kernel}} bash scripts/ci/cuda/ci_install_dependency.sh diffusion
-
-      - name: Run sgl-kernel unit tests on B200
-        timeout-minutes: 30
-        run: |
-          cd sgl-kernel
-          pytest tests/
-
-  # Adding a single CUDA13 smoke test to verify that the kernel builds and runs
-  # TODO: Add back this test when it can pass on CI
-  # cuda13-kernel-smoke-test:
-  #   needs: [check-changes, sgl-kernel-build-wheels]
-  #   if: needs.check-changes.outputs.sgl_kernel == 'true'
-  #   runs-on: x64-cu13-kernel-tests
-  #   steps:
-  #     - uses: actions/checkout@v4
-
-  #     - name: Cleanup
-  #       run: |
-  #         ls -alh sgl-kernel/dist || true
-  #         rm -rf sgl-kernel/dist/* || true
-
-  #     - name: Download CUDA 13.0 artifacts
-  #       uses: actions/download-artifact@v4
-  #       with:
-  #         path: sgl-kernel/dist/
-  #         merge-multiple: true
-  #         pattern: wheel-python3.10-cuda13.0
-
-  #     - name: Install dependencies
-  #       run: |
-  #         CUSTOM_BUILD_SGL_KERNEL=${{needs.check-changes.outputs.sgl_kernel}} bash scripts/ci/cuda/ci_install_dependency.sh
-
-  #     - name: Run kernel unit tests
-  #       timeout-minutes: 30
-  #       run: |
-  #         cd sgl-kernel
-  #         pytest tests/
 
   # =============================================== jit-kernel ====================================================
 
-  jit-kernel-unit-test:
+  call-jit-kernel-tests:
     needs: [check-changes, call-gate]
-    # Skip for scheduled runs and when target_stage is set
+    if: needs.check-changes.outputs.jit_kernel == 'true'
-    if: |
+    uses: ./.github/workflows/pr-test-jit-kernel.yml
-      github.event_name != 'schedule' &&
+    with:
-      inputs.test_parallel_dispatch != true &&
+      jit_kernel: ${{ needs.check-changes.outputs.jit_kernel }}
-      !inputs.target_stage &&
+      pr_head_sha: ${{ inputs.pr_head_sha || '' }}
-      needs.check-changes.outputs.jit_kernel == 'true'
+      git_ref: ${{ inputs.git_ref || '' }}
-    runs-on: 1-gpu-h100
+      target_stage: ${{ inputs.target_stage || '' }}
-    timeout-minutes: 240
+      test_parallel_dispatch: ${{ inputs.test_parallel_dispatch == true && 'true' || 'false' }}
-    steps:
+    secrets: inherit
-      - uses: actions/checkout@v4
-        with:
-          ref: ${{ inputs.pr_head_sha || inputs.git_ref || github.sha }}
-
-      - uses: ./.github/actions/check-maintenance
-        with:
-          github-token: ${{ github.token }}
-
-      - name: Install dependencies
-        timeout-minutes: 20
-        run: |
-          bash scripts/ci/cuda/ci_install_dependency.sh diffusion
-
-      - name: Run test
-        timeout-minutes: 30
-        run: |
-          cd python/sglang/jit_kernel
-          pytest tests/
-
-  jit-kernel-unit-test-nightly:
-    needs: [check-changes]
-    if: |
-      github.event_name == 'schedule' &&
-      needs.check-changes.outputs.jit_kernel == 'true'
-    runs-on: 1-gpu-h100
-    timeout-minutes: 240
-    env:
-      SGLANG_JIT_KERNEL_RUN_FULL_TESTS: "1"
-    steps:
-      - uses: actions/checkout@v4
-        with:
-          ref: ${{ inputs.pr_head_sha || inputs.git_ref || github.sha }}
-
-      - uses: ./.github/actions/check-maintenance
-        with:
-          github-token: ${{ github.token }}
-
-      - name: Install dependencies
-        timeout-minutes: 20
-        run: |
-          bash scripts/ci/cuda/ci_install_dependency.sh
-
-      - name: Run full nightly test
-        timeout-minutes: 60
-        run: |
-          cd python/sglang/jit_kernel
-          pytest tests/
-
-  jit-kernel-benchmark-test:
-    needs: [check-changes, call-gate]
-    # Skip for scheduled runs and when target_stage is set
-    if: |
-      github.event_name != 'schedule' &&
-      inputs.test_parallel_dispatch != true &&
-      !inputs.target_stage &&
-      needs.check-changes.outputs.jit_kernel == 'true'
-    runs-on: 1-gpu-h100
-    timeout-minutes: 240
-    steps:
-      - uses: actions/checkout@v4
-        with:
-          ref: ${{ inputs.pr_head_sha || inputs.git_ref || github.sha }}
-
-      - uses: ./.github/actions/check-maintenance
-        with:
-          github-token: ${{ github.token }}
-
-      - name: Install dependencies
-        timeout-minutes: 20
-        run: |
-          bash scripts/ci/cuda/ci_install_dependency.sh diffusion
-
-      - name: Run benchmark tests
-        timeout-minutes: 45
-        run: |
-          cd python/sglang/jit_kernel/benchmark
-          echo "Running jit-kernel benchmark tests in CI mode..."
-
-          failures=()
-
-          for bench_file in bench_*.py; do
-            echo "Testing $bench_file..."
-            if ! timeout 120 python3 "$bench_file"; then
-              failures+=("$bench_file")
-            fi
-            echo "Completed $bench_file"
-            echo "---"
-          done
-
-          if [ ${#failures[@]} -ne 0 ]; then
-            echo "The following benchmark tests failed: ${failures[*]}"
-            exit 1
-          fi
-
-          echo "All jit-kernel benchmark tests completed successfully!"
 
   # =============================================== primary ====================================================
 
@@ -1164,165 +876,32 @@ jobs:
       - uses: ./.github/actions/upload-cuda-coredumps
         if: always()
 
-  multimodal-gen-test-1-gpu:
+  call-multimodal-gen-tests:
     needs: [check-changes, call-gate, sgl-kernel-build-wheels]
     if: |
       always() &&
+      !cancelled() &&
       (
-        (inputs.target_stage == 'multimodal-gen-test-1-gpu') ||
+        inputs.target_stage == 'multimodal-gen-test-1-gpu' ||
+        inputs.target_stage == 'multimodal-gen-test-2-gpu' ||
+        inputs.target_stage == 'multimodal-gen-unit-test' ||
         (
           !inputs.target_stage &&
           ((github.event_name == 'schedule' || inputs.test_parallel_dispatch == true) || (!failure() && !cancelled())) &&
           needs.check-changes.outputs.multimodal_gen == 'true'
         )
       )
-    runs-on: 1-gpu-h100
+    uses: ./.github/workflows/pr-test-multimodal-gen.yml
-    timeout-minutes: 240
+    with:
-    strategy:
+      multimodal_gen: ${{ needs.check-changes.outputs.multimodal_gen }}
-      fail-fast: false
+      sgl_kernel: ${{ needs.check-changes.outputs.sgl_kernel }}
-      matrix:
+      continue_on_error: ${{ needs.check-changes.outputs.continue_on_error }}
-        part: [0, 1]
+      pr_head_sha: ${{ inputs.pr_head_sha || '' }}
-    steps:
+      git_ref: ${{ inputs.git_ref || '' }}
-      - name: Checkout code
+      target_stage: ${{ inputs.target_stage || '' }}
-        uses: actions/checkout@v4
+      test_parallel_dispatch: ${{ inputs.test_parallel_dispatch == true && 'true' || 'false' }}
-        with:
+      caller_needs_failure: ${{ (needs.call-gate.result == 'failure' || needs.sgl-kernel-build-wheels.result == 'failure' || needs.check-changes.result == 'failure') && 'true' || 'false' }}
-          ref: ${{ inputs.pr_head_sha || inputs.git_ref || github.sha }}
+    secrets: inherit
-
-      - uses: ./.github/actions/check-maintenance
-        with:
-          github-token: ${{ github.token }}
-
-      - name: Download artifacts
-        if: needs.check-changes.outputs.sgl_kernel == 'true'
-        uses: actions/download-artifact@v4
-        with:
-          path: sgl-kernel/dist/
-          merge-multiple: true
-          pattern: wheel-python3.10-cuda12.9
-
-      - name: Install dependencies
-        timeout-minutes: 20
-        run: |
-          CUSTOM_BUILD_SGL_KERNEL=${{needs.check-changes.outputs.sgl_kernel}} bash scripts/ci/cuda/ci_install_dependency.sh diffusion
-      - name: Run diffusion server tests
-        timeout-minutes: 240
-        env:
-          RUNAI_STREAMER_MEMORY_LIMIT: 0
-          CONTINUE_ON_ERROR_FLAG: ${{ needs.check-changes.outputs.continue_on_error == 'true' && '--continue-on-error' || '' }}
-        run: |
-          cd python
-          python3 sglang/multimodal_gen/test/run_suite.py \
-            --suite 1-gpu \
-            --partition-id ${{ matrix.part }} \
-            --total-partitions 2 \
-            $CONTINUE_ON_ERROR_FLAG
-
-      - uses: ./.github/actions/upload-cuda-coredumps
-        if: always()
-        with:
-          artifact-suffix: ${{ matrix.part }}
-
-  multimodal-gen-test-2-gpu:
-    needs: [check-changes, call-gate, sgl-kernel-build-wheels]
-    if: |
-      always() &&
-      (
-        (inputs.target_stage == 'multimodal-gen-test-2-gpu') ||
-        (
-          !inputs.target_stage &&
-          ((github.event_name == 'schedule' || inputs.test_parallel_dispatch == true) || (!failure() && !cancelled())) &&
-          needs.check-changes.outputs.multimodal_gen == 'true'
-        )
-      )
-    runs-on: 2-gpu-h100
-    timeout-minutes: 240
-    strategy:
-      fail-fast: false
-      matrix:
-        part: [0, 1]
-    steps:
-      - name: Checkout code
-        uses: actions/checkout@v4
-        with:
-          ref: ${{ inputs.pr_head_sha || inputs.git_ref || github.sha }}
-
-      - uses: ./.github/actions/check-maintenance
-        with:
-          github-token: ${{ github.token }}
-
-      - name: Download artifacts
-        if: needs.check-changes.outputs.sgl_kernel == 'true'
-        uses: actions/download-artifact@v4
-        with:
-          path: sgl-kernel/dist/
-          merge-multiple: true
-          pattern: wheel-python3.10-cuda12.9
-
-      - name: Install dependencies
-        timeout-minutes: 20
-        run: |
-          CUSTOM_BUILD_SGL_KERNEL=${{needs.check-changes.outputs.sgl_kernel}} bash scripts/ci/cuda/ci_install_dependency.sh diffusion
-
-      - name: Run diffusion server tests
-        timeout-minutes: 240
-        env:
-          RUNAI_STREAMER_MEMORY_LIMIT: 0
-          CONTINUE_ON_ERROR_FLAG: ${{ needs.check-changes.outputs.continue_on_error == 'true' && '--continue-on-error' || '' }}
-        run: |
-          cd python
-          python3 sglang/multimodal_gen/test/run_suite.py \
-            --suite 2-gpu \
-            --partition-id ${{ matrix.part }} \
-            --total-partitions 2 \
-            $CONTINUE_ON_ERROR_FLAG
-
-      - uses: ./.github/actions/upload-cuda-coredumps
-        if: always()
-        with:
-          artifact-suffix: ${{ matrix.part }}
-
-  multimodal-gen-unit-test:
-    needs: [check-changes, call-gate, sgl-kernel-build-wheels]
-    if: |
-      always() &&
-      (
-        (inputs.target_stage == 'multimodal-gen-unit-test') ||
-        (
-          !inputs.target_stage &&
-          ((github.event_name == 'schedule' || inputs.test_parallel_dispatch == true) || (!failure() && !cancelled())) &&
-          needs.check-changes.outputs.multimodal_gen == 'true'
-        )
-      )
-    runs-on: 1-gpu-h100
-    timeout-minutes: 120
-    steps:
-      - name: Checkout code
-        uses: actions/checkout@v4
-        with:
-          ref: ${{ inputs.pr_head_sha || inputs.git_ref || github.sha }}
-
-      - uses: ./.github/actions/check-maintenance
-        with:
-          github-token: ${{ github.token }}
-
-      - name: Download artifacts
-        if: needs.check-changes.outputs.sgl_kernel == 'true'
-        uses: actions/download-artifact@v4
-        with:
-          path: sgl-kernel/dist/
-          merge-multiple: true
-          pattern: wheel-python3.10-cuda12.9
-
-      - name: Install dependencies
-        timeout-minutes: 20
-        run: |
-          CUSTOM_BUILD_SGL_KERNEL=${{needs.check-changes.outputs.sgl_kernel}} bash scripts/ci/cuda/ci_install_dependency.sh diffusion
-
-      - name: Run diffusion unit tests
-        timeout-minutes: 60
-        run: |
-          cd python
-          python3 sglang/multimodal_gen/test/run_suite.py --suite unit
 
   stage-c-test-4-gpu-h100:
     needs: [check-changes, call-gate, wait-for-stage-b]
@@ -1737,21 +1316,15 @@ jobs:
         check-changes,
 
         sgl-kernel-build-wheels,
-        sgl-kernel-unit-test,
+        sgl-kernel-build-wheels-arm,
-        sgl-kernel-mla-test,
+        call-sgl-kernel-tests,
-        sgl-kernel-benchmark-test,
-        sgl-kernel-b200-test,
 
         wait-for-stage-a,
         wait-for-stage-b,
 
-        jit-kernel-unit-test,
+        call-jit-kernel-tests,
-        jit-kernel-unit-test-nightly,
-        jit-kernel-benchmark-test,
 
-        multimodal-gen-unit-test,
+        call-multimodal-gen-tests,
-        multimodal-gen-test-1-gpu,
-        multimodal-gen-test-2-gpu,
 
         stage-a-test-1-gpu-small,
         stage-a-test-cpu,