[JIT kernel] CI support jit kernel tests (#15939)

.github/workflows/pr-test.yml

@@ -40,6 +40,7 @@ jobs:
     outputs:
       main_package: ${{ steps.filter.outputs.main_package || steps.scheduled.outputs.main_package }}
       sgl_kernel: ${{ steps.filter.outputs.sgl_kernel || steps.scheduled.outputs.sgl_kernel }}
+      jit_kernel: ${{ steps.filter.outputs.jit_kernel || steps.scheduled.outputs.jit_kernel }}
       multimodal_gen: ${{ steps.filter.outputs.multimodal_gen || steps.scheduled.outputs.multimodal_gen }}
       max_parallel: ${{ steps.set-parallel.outputs.max_parallel }}
       b200_runner: ${{ steps.set-runner.outputs.b200_runner }}
@@ -63,6 +64,10 @@ jobs:
               - ".github/workflows/pr-test.yml"
             sgl_kernel:
               - "sgl-kernel/**"
+            jit_kernel:
+              - "python/sglang/jit_kernel/**"
+              - "python/*.toml"
+              - ".github/workflows/pr-test.yml"
             multimodal_gen:
               - "python/sglang/multimodal_gen/**"
               - "python/sglang/cli/**"
@@ -75,6 +80,7 @@ jobs:
         run: |
           echo "main_package=true" >> $GITHUB_OUTPUT
           echo "sgl_kernel=false" >> $GITHUB_OUTPUT
+          echo "jit_kernel=true" >> $GITHUB_OUTPUT
           echo "multimodal_gen=true" >> $GITHUB_OUTPUT
 
       - name: Set max-parallel based on high-priority label
@@ -124,6 +130,7 @@ jobs:
             echo "|-------------------|---------|"
             echo "| main_package      | ${{ steps.filter.outputs.main_package || steps.scheduled.outputs.main_package }} |"
             echo "| sgl_kernel        | ${{ steps.filter.outputs.sgl_kernel || steps.scheduled.outputs.sgl_kernel }} |"
+            echo "| jit_kernel        | ${{ steps.filter.outputs.jit_kernel || steps.scheduled.outputs.jit_kernel }} |"
             echo "| multimodal_gen    | ${{ steps.filter.outputs.multimodal_gen || steps.scheduled.outputs.multimodal_gen }} |"
             echo "| max_parallel      | ${{ steps.set-parallel.outputs.max_parallel }} |"
             echo "| b200_runner       | ${{ steps.set-runner.outputs.b200_runner }} |"
@@ -137,6 +144,7 @@ jobs:
     if: |
       needs.check-changes.outputs.main_package == 'true' ||
       needs.check-changes.outputs.sgl_kernel == 'true' ||
+      needs.check-changes.outputs.jit_kernel == 'true' ||
       needs.check-changes.outputs.multimodal_gen == 'true'
     uses: ./.github/workflows/pr-gate.yml
     secrets: inherit
@@ -392,6 +400,27 @@ jobs:
   #         cd sgl-kernel
   #         pytest tests/
 
+  # =============================================== jit-kernel ====================================================
+
+  jit-kernel-unit-test:
+    needs: [check-changes, call-gate]
+    if: needs.check-changes.outputs.jit_kernel == 'true'
+    runs-on: 1-gpu-runner
+    env:
+      RUNNER_LABELS: 1-gpu-runner
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Install dependencies
+        run: |
+          bash scripts/ci/ci_install_dependency.sh
+
+      - name: Run test
+        timeout-minutes: 30
+        run: |
+          cd python/sglang/jit_kernel
+          pytest tests/
+
   # =============================================== primary ====================================================
 
   stage-a-test-1:
@@ -1572,6 +1601,8 @@ jobs:
         sgl-kernel-benchmark-test,
         sgl-kernel-b200-test,
 
+        jit-kernel-unit-test,
+
         multimodal-gen-test-1-gpu,
         multimodal-gen-test-2-gpu,
 

python/sglang/jit_kernel/tests/test_per_tensor_quant_fp8.py

@@ -43,7 +43,7 @@ def torch_scaled_fp8_quant(tensor, inv_scale):
     "num_tokens,hidden_dim",
     list(itertools.product([128, 256, 512], [512, 2048, 4096])),
 )
-def test_per_tensor_quant_compare_implementations(
+def test_jit_per_tensor_quant_compare_implementations(
     num_tokens: int,
     hidden_dim: int,
 ):