Update AOT kernels for Torch 2.13 (#33841)

python/sglang/kernels/aot/CMakeLists.txt

@@ -65,8 +65,8 @@ FetchContent_Populate(repo-fmt)
 # Triton kernel
 FetchContent_Declare(
     repo-triton
-    URL      https://${GITHUB_ARTIFACTORY}/triton-lang/triton/archive/v3.6.0.tar.gz
-    URL_HASH SHA256=be270ed11ca5a8fbd9d7941c5bbe9a23a9f6e2ffd372c8398346928bee464774
+    URL      https://${GITHUB_ARTIFACTORY}/triton-lang/triton/archive/v3.7.1.tar.gz
+    URL_HASH SHA256=7d998625d1035ac496d06a81117727647169422ee67b8889609f06fd5367c491
 )
 FetchContent_Populate(repo-triton)
 

python/sglang/kernels/aot/Dockerfile

@@ -97,10 +97,9 @@ RUN set -eux; \
 RUN --mount=type=cache,id=sgl-kernel-pip,target=/root/.cache/pip \
     set -eux; \
     case "${CUDA_VERSION}" in \
-      13.0) TORCH_VER=2.11.0; CU_TAG=cu130 ;; \
-      12.9) TORCH_VER=2.11.0; CU_TAG=cu129 ;; \
-      12.8) TORCH_VER=2.11.0; CU_TAG=cu128 ;; \
-      *)    TORCH_VER=2.11.0; CU_TAG=cu126 ;; \
+      13.0) TORCH_VER=2.13.0; CU_TAG=cu130 ;; \
+      12.9) TORCH_VER=2.13.0; CU_TAG=cu129 ;; \
+      *)    TORCH_VER=2.13.0; CU_TAG=cu126 ;; \
     esac; \
     ${PYTHON_ROOT_PATH}/bin/pip install torch==${TORCH_VER} --index-url ${PYTORCH_INDEX_BASE}/${CU_TAG}; \
     ${PYTHON_ROOT_PATH}/bin/pip install ninja setuptools==75.0.0 wheel==0.41.0 numpy uv scikit-build-core --index-url ${PIP_DEFAULT_INDEX}

python/sglang/kernels/aot/README.md

@@ -12,7 +12,7 @@
 `sglang-kernel` provides optimized compute primitives for LLM inference engines, enabling efficient inference for large language models and vision-language models through custom kernel operations. The source tree lives under the `python/sglang/kernels/aot/` directory and the Python import path remains `sgl_kernel`.
 
 ## Installation
-Requires torch == 2.11.0
+Requires torch == 2.13.0
 
 ```bash
 # Latest version

python/sglang/kernels/aot/rename_wheels.sh

@@ -1,5 +1,5 @@
 #!/usr/bin/env bash
-# Align CUDA wheel filenames (+cu124/+cu128/+cu129/+cu130) with internal METADATA Version and
+# Align CUDA wheel filenames (+cu124/+cu129/+cu130) with internal METADATA Version and
 # WHEEL tags after build (fixes pip "inconsistent version" when only the .whl name changed).
 # Unpack → patch WHEEL/METADATA → wheel pack (RECORD regenerated; no hand-editing).
 set -ex
@@ -9,8 +9,6 @@ WHEEL_DIR="dist"
 detect_cuda_suffix() {
     if ls /usr/local/ 2>/dev/null | grep -q "12.4"; then
         echo "+cu124"
-    elif ls /usr/local/ 2>/dev/null | grep -q "12.8"; then
-        echo "+cu128"
     elif ls /usr/local/ 2>/dev/null | grep -q "12.9"; then
         echo "+cu129"
     elif ls /usr/local/ 2>/dev/null | grep -q "13.0"; then

python/sglang/kernels/aot/tests/test_flash_attn_sparse.py

@@ -1,6 +1,6 @@
 import math
 import sys
-from typing import List, Optional, Tuple
+from typing import List, Optional
 
 import pytest
 import torch
@@ -370,27 +370,23 @@ def test_convert_vertical_slash_indexes_mergehead(causal):
         )
     )
 
-    # Manually create expected outputs for this input
-    # For demonstration, assume:
-    # - batch=1, head=2, num_rows=2, nnz_v=2, nnz_s=2
-    # Fill these expected tensors according to your kernel's behavior
-
-    expected_column_index = torch.tensor(
-        [[[[1, 0], [1, 3]], [[-1079459945, -1077788999], [-1080050043, -1104625879]]]],
-        dtype=torch.int32,
-        device="cuda",
-    )
+    # column_index is torch.empty-backed; only entries before column_count are valid.
+    expected_column_count = torch.zeros((1, 2, 2), dtype=torch.int32, device="cuda")
+    expected_column_index = [[[[], []], [[], []]]]
 
     if not causal:
-        # If non-causal mode output is different, update these values
-        expected_column_index = torch.tensor(
-            [[[[1, 0], [1, 3]], [[2, -1077788999], [2, -1104625879]]]],
-            dtype=torch.int32,
-            device="cuda",
+        expected_column_count = torch.tensor(
+            [[[1, 2], [1, 1]]], dtype=torch.int32, device="cuda"
         )
+        expected_column_index = [[[[1], [1, 3]], [[2], [2]]]]
 
-    # Assert that outputs match expectations
-    assert torch.equal(column_index, expected_column_index)
+    assert torch.equal(column_count, expected_column_count)
+    for batch_idx, batch_expected in enumerate(expected_column_index):
+        for head_idx, head_expected in enumerate(batch_expected):
+            for row_idx, expected_values in enumerate(head_expected):
+                count = int(column_count[batch_idx, head_idx, row_idx].item())
+                actual = column_index[batch_idx, head_idx, row_idx, :count].tolist()
+                assert actual == expected_values
 
 
 # skip cause use fa2 for test