name: PR Test - JIT Kernel

on:
  workflow_call:
    inputs:
      jit_kernel:
        required: true
        type: string
      sgl_kernel:
        required: true
        type: string
      runner_config:
        required: true
        type: string
      runs_on_map:
        required: true
        type: string
      git_ref:
        required: false
        type: string
        default: ''
      test_parallel_dispatch:
        required: false
        type: string
        default: 'false'
      skip_pr_test_health_check:
        required: false
        type: boolean
        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
  PR_TEST_BYPASS_MAINTENANCE_ON_MAIN: ${{ github.ref == 'refs/heads/main' && 'true' || 'false' }}
  SKIP_PR_TEST_HEALTH_CHECK: ${{ inputs.skip_pr_test_health_check == true && 'true' || 'false' }}

jobs:
  jit-kernel-unit-test:
    # Runs whenever call-jit-kernel-tests dispatches this workflow. That caller is the
    # single gate (PR jit_kernel changes, or scheduled/parallel-dispatch full runs), so
    # the sub-jobs no longer re-exclude schedule/parallel-dispatch here.
    runs-on: 1-gpu-h100
    timeout-minutes: 240
    steps:
      - uses: actions/checkout@v4
        with:
          ref: ${{ inputs.git_ref || github.sha }}

      - uses: ./.github/actions/check-pr-test-health

      - uses: ./.github/actions/check-maintenance

      - name: Cleanup
        if: inputs.sgl_kernel == 'true'
        run: |
          ls -alh sgl-kernel/dist || true
          rm -rf sgl-kernel/dist/* || true

      - 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-cuda13.0

      - 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 test/
          python3 run_suite.py --hw cuda --suite base-b-kernel-unit-test-1-gpu-large

  jit-kernel-multigpu-unit-test:
    # Runs whenever call-jit-kernel-tests dispatches this workflow. That caller is the
    # single gate (PR jit_kernel changes, or scheduled/parallel-dispatch full runs), so
    # the sub-jobs no longer re-exclude schedule/parallel-dispatch here.
    runs-on: 8-gpu-h200
    timeout-minutes: 240
    steps:
      - uses: actions/checkout@v4
        with:
          ref: ${{ inputs.git_ref || github.sha }}

      - uses: ./.github/actions/check-maintenance

      - name: Cleanup
        if: inputs.sgl_kernel == 'true'
        run: |
          ls -alh sgl-kernel/dist || true
          rm -rf sgl-kernel/dist/* || true

      - 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-cuda13.0

      - 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 multi-GPU test
        timeout-minutes: 45
        run: |
          cd test/
          python3 run_suite.py --hw cuda --suite base-b-kernel-unit-test-8-gpu-h200

  jit-kernel-benchmark-test:
    # Runs whenever call-jit-kernel-tests dispatches this workflow. That caller is the
    # single gate (PR jit_kernel changes, or scheduled/parallel-dispatch full runs), so
    # the sub-jobs no longer re-exclude schedule/parallel-dispatch here.
    runs-on: 1-gpu-h100
    timeout-minutes: 240
    steps:
      - uses: actions/checkout@v4
        with:
          ref: ${{ inputs.git_ref || github.sha }}

      - uses: ./.github/actions/check-pr-test-health

      - uses: ./.github/actions/check-maintenance

      - name: Cleanup
        if: inputs.sgl_kernel == 'true'
        run: |
          ls -alh sgl-kernel/dist || true
          rm -rf sgl-kernel/dist/* || true

      - 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-cuda13.0

      - 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 benchmark tests
        timeout-minutes: 45
        run: |
          cd test/
          python3 run_suite.py --hw cuda --suite base-b-kernel-benchmark-test-1-gpu-large

  jit-kernel-b200-test:
    # Runs whenever call-jit-kernel-tests dispatches this workflow. That caller is the
    # single gate (PR jit_kernel changes, or scheduled/parallel-dispatch full runs), so
    # the sub-jobs no longer re-exclude schedule/parallel-dispatch here.
    runs-on: ${{ fromJson(inputs.runs_on_map)[inputs.runner_config] }}
    timeout-minutes: 240
    steps:
      - uses: actions/checkout@v4
        with:
          ref: ${{ inputs.git_ref || github.sha }}

      - uses: ./.github/actions/check-pr-test-health

      - uses: ./.github/actions/check-maintenance

      - name: Cleanup
        if: inputs.sgl_kernel == 'true'
        run: |
          ls -alh sgl-kernel/dist || true
          rm -rf sgl-kernel/dist/* || true

      - 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-cuda13.0

      - 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 B200 diffusion test
        timeout-minutes: 30
        run: |
          cd test/
          python3 run_suite.py --hw cuda --suite base-b-kernel-unit-test-4-gpu-b200