Fix SGLang third party test (#4652)

Resolve the conflicts when applying sglang patch.

Author: leonling-ll

.github/workflows/sglang-tests.yml

@@ -74,23 +74,27 @@ jobs:
           echo "REPORTS=$PWD/reports" >> $GITHUB_ENV
 
       - name: Install SGLang
-        id: install-sglang
+        id: install
         run: |
           git clone https://github.com/sgl-project/sglang.git
           cd sglang
-          git apply ../benchmarks/third_party/sglang/sglang.patch
-          pip install ./python[dev_xpu]
+          git apply ../benchmarks/third_party/sglang/sglang-fix.patch
+          pip install "./python[dev_xpu]"
 
-      # Install Pytorch and Triton after SGLANG to ensure that the correct versions are used
       - name: Setup PyTorch
         uses: ./.github/actions/setup-pytorch
 
       - name: Setup Triton
         uses: ./.github/actions/setup-triton
 
       - name: Run SGLANG tests
-        if: ${{ steps.install.outcome == 'success' && steps.install-sglang.outcome == 'success' && !cancelled() }}
+        if: ${{ steps.install.outcome == 'success' && !cancelled() }}
         run: |
-            pip install pytest pytest-xdist
-            cd sglang
-            pytest -vvv -n 4 test/srt/test_triton_attention_kernels.py
+          ./scripts/test-triton.sh --sglang --skip-pip-install --skip-pytorch-install
+
+      - name: Upload test report
+        if: ${{ steps.install.outcome == 'success' && !cancelled() }}
+        uses: actions/upload-artifact@v4
+        with:
+          name: test-reports
+          path: reports

benchmarks/third_party/sglang/sglang.patch renamed to benchmarks/third_party/sglang/sglang-fix.patch

@@ -1,39 +1,12 @@
 diff --git a/python/sglang/srt/utils.py b/python/sglang/srt/utils.py
-index 884e715f..14e5df33 100644
+index bc2affa1..8ef91e66 100644
 --- a/python/sglang/srt/utils.py
 +++ b/python/sglang/srt/utils.py
-@@ -77,12 +77,20 @@ from torch.func import functional_call
- from torch.library import Library
- from torch.profiler import ProfilerActivity, profile, record_function
- from torch.utils._contextlib import _DecoratorContextManager
--from triton.runtime.cache import (
--    FileCacheManager,
--    default_cache_dir,
--    default_dump_dir,
--    default_override_dir,
--)
-+try:
-+    from triton.runtime.cache import (
-+        FileCacheManager,
-+        default_cache_dir,
-+        default_dump_dir,
-+        default_override_dir,
-+    )
-+except ImportError:
-+    from triton.runtime.cache import FileCacheManager
-+    from triton.knobs import cache as tt_cache
-+
-+    default_cache_dir = lambda: tt_cache.dir
-+    default_dump_dir = lambda: tt_cache.dump_dir
-+    default_override_dir = lambda: tt_cache.override_dir
- 
- logger = logging.getLogger(__name__)
+@@ -228,6 +228,22 @@ def is_flashinfer_available():
+     return importlib.util.find_spec("flashinfer") is not None and is_cuda()
 
-@@ -156,6 +164,18 @@ def is_xpu() -> bool:
- def is_npu() -> bool:
-     return hasattr(torch, "npu") and torch.npu.is_available()
 
-+def infer_device():
++def auto_detect_device():
 +    """
 +    Infer the device type based on the current environment.
 +    """
@@ -43,23 +16,27 @@ index 884e715f..14e5df33 100644
 +        return "xpu"
 +    elif is_hpu():
 +        return "hpu"
++    elif is_npu():
++        return "npu"
 +    else:
 +        return "cpu"
- 
- def is_flashinfer_available():
-     """
++
++
+ _ENABLE_TORCH_INFERENCE_MODE = get_bool_env_var(
+     "SGLANG_ENABLE_TORCH_INFERENCE_MODE", "false"
+ )
 diff --git a/test/srt/test_triton_attention_kernels.py b/test/srt/test_triton_attention_kernels.py
-index 47eb16a9..9d6a0af0 100644
+index 47eb16a9..cce70fb9 100644
 --- a/test/srt/test_triton_attention_kernels.py
 +++ b/test/srt/test_triton_attention_kernels.py
 @@ -16,8 +16,11 @@ from sglang.srt.layers.attention.triton_ops.prefill_attention import (
      context_attention_fwd,
  )
  from sglang.test.test_utils import CustomTestCase
-+from sglang.srt.utils import infer_device
++from sglang.srt.utils import auto_detect_device
 
 
-+device = infer_device()
++device = auto_detect_device()
 +
  class TestTritonAttention(CustomTestCase):
 

scripts/test-triton.sh

@@ -20,6 +20,7 @@ TEST:
     --attention
     --instrumentation
     --inductor
+    --sglang
 
 OPTION:
     --unskip
@@ -57,6 +58,7 @@ TEST_BENCHMARK_GEMM=false
 TEST_BENCHMARK_ATTENTION=false
 TEST_INSTRUMENTATION=false
 TEST_INDUCTOR=false
+TEST_SGLANG=false
 VENV=false
 TRITON_TEST_REPORTS=false
 TRITON_TEST_WARNING_REPORTS=false
@@ -141,6 +143,11 @@ while (( $# != 0 )); do
       TEST_DEFAULT=false
       shift
       ;;
+    --sglang)
+      TEST_SGLANG=true
+      TEST_DEFAULT=false
+      shift
+      ;;
     --venv)
       VENV=true
       shift
@@ -470,6 +477,29 @@ run_inductor_tests() {
   grep AlbertForMaskedLM inductor_log.csv | grep -q ,pass,
 }
 
+run_sglang_tests() {
+  echo "***************************************************"
+  echo "******    Running SGLang Triton tests       ******"
+  echo "***************************************************"
+
+  if ! [ -d "./sglang" ]; then
+    git clone https://github.com/sgl-project/sglang.git
+  fi
+  cd sglang
+
+  if ! pip list | grep "sglang" ; then
+    git apply $TRITON_PROJ/benchmarks/third_party/sglang/sglang-fix.patch
+    pip install "./python[dev_xpu]"
+
+    # SGLang installation breaks the default PyTorch and Triton versions, so we need to reinstall them.
+    $SCRIPTS_DIR/install-pytorch.sh --force-reinstall
+    $SCRIPTS_DIR/compile-triton.sh --triton
+  fi
+
+  pip install pytest pytest-xdist
+  run_pytest_command -vvv -n ${PYTEST_MAX_PROCESSES:-4} test/srt/test_triton_attention_kernels.py
+}
+
 test_triton() {
   if [ "$TEST_UNIT" = true ]; then
     run_unit_tests
@@ -517,6 +547,9 @@ test_triton() {
   if [ "$TEST_INDUCTOR" == true ]; then
     run_inductor_tests
   fi
+  if [ "$TEST_SGLANG" == true ]; then
+    run_sglang_tests
+  fi
 }
 
 install_deps