diff --git a/python/sglang/kernels/aot/CMakeLists.txt b/python/sglang/kernels/aot/CMakeLists.txt index d456c2b29..0d22047a3 100644 --- a/python/sglang/kernels/aot/CMakeLists.txt +++ b/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) diff --git a/python/sglang/kernels/aot/Dockerfile b/python/sglang/kernels/aot/Dockerfile index 92c2626ee..f8f2853c3 100644 --- a/python/sglang/kernels/aot/Dockerfile +++ b/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} diff --git a/python/sglang/kernels/aot/README.md b/python/sglang/kernels/aot/README.md index 0227c4b17..f961324a1 100644 --- a/python/sglang/kernels/aot/README.md +++ b/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 diff --git a/python/sglang/kernels/aot/rename_wheels.sh b/python/sglang/kernels/aot/rename_wheels.sh index 550dfc68b..76a513679 100755 --- a/python/sglang/kernels/aot/rename_wheels.sh +++ b/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 diff --git a/python/sglang/kernels/aot/tests/test_flash_attn_sparse.py b/python/sglang/kernels/aot/tests/test_flash_attn_sparse.py index f8d344ef3..1126894a4 100644 --- a/python/sglang/kernels/aot/tests/test_flash_attn_sparse.py +++ b/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