diff --git a/.github/workflows/pr-test-npu.yml b/.github/workflows/pr-test-npu.yml index cca050111..9020447dd 100644 --- a/.github/workflows/pr-test-npu.yml +++ b/.github/workflows/pr-test-npu.yml @@ -38,9 +38,10 @@ jobs: CACHING_URL="cache-service.nginx-pypi-cache.svc.cluster.local" sed -Ei "s@(ports|archive).ubuntu.com@${CACHING_URL}:8081@g" /etc/apt/sources.list pip config set global.index-url http://${CACHING_URL}/pypi/simple - pip config set global.trusted-host ${CACHING_URL} + pip config set global.extra-index-url "https://pypi.tuna.tsinghua.edu.cn/simple https://mirrors.aliyun.com/pypi/simple/" + pip config set global.trusted-host "${CACHING_URL} pypi.tuna.tsinghua.edu.cn mirrors.aliyun.com" - bash scripts/ci/npu_ci_install_dependency.sh + bash scripts/ci/npu_ci_install_dependency.sh 910b # copy required file from our daily cache cp ~/.cache/modelscope/hub/datasets/otavia/ShareGPT_Vicuna_unfiltered/ShareGPT_V3_unfiltered_cleaned_split.json /tmp # copy download through proxy @@ -53,13 +54,20 @@ jobs: SGLANG_IS_IN_CI: true HF_ENDPOINT: https://hf-mirror.com TORCH_EXTENSIONS_DIR: /tmp/torch_extensions + PYTORCH_NPU_ALLOC_CONF: "expandable_segments:True" + STREAMS_PER_DEVICE: 32 run: | + export PATH="/usr/local/Ascend/8.3.RC1/compiler/bishengir/bin:${PATH}" cd test/srt python3 run_suite.py --suite per-commit-1-ascend-npu per-commit-2-ascend-npu: if: github.event_name != 'pull_request' || contains(github.event.pull_request.labels.*.name, 'run-ci') runs-on: linux-arm64-npu-2 + strategy: + fail-fast: false + matrix: + part: [0, 1, 2] container: image: swr.cn-southwest-2.myhuaweicloud.com/base_image/ascend-ci/cann:8.2.rc1-910b-ubuntu22.04-py3.11 steps: @@ -72,24 +80,28 @@ jobs: CACHING_URL="cache-service.nginx-pypi-cache.svc.cluster.local" sed -Ei "s@(ports|archive).ubuntu.com@${CACHING_URL}:8081@g" /etc/apt/sources.list pip config set global.index-url http://${CACHING_URL}/pypi/simple - pip config set global.trusted-host ${CACHING_URL} + pip config set global.extra-index-url "https://pypi.tuna.tsinghua.edu.cn/simple https://mirrors.aliyun.com/pypi/simple/" + pip config set global.trusted-host "${CACHING_URL} pypi.tuna.tsinghua.edu.cn mirrors.aliyun.com" - bash scripts/ci/npu_ci_install_dependency.sh + bash scripts/ci/npu_ci_install_dependency.sh 910b # copy required file from our daily cache cp ~/.cache/modelscope/hub/datasets/otavia/ShareGPT_Vicuna_unfiltered/ShareGPT_V3_unfiltered_cleaned_split.json /tmp # copy download through proxy curl -o /tmp/test.jsonl -L https://gh-proxy.test.osinfra.cn/https://raw.githubusercontent.com/openai/grade-school-math/master/grade_school_math/data/test.jsonl - name: Run test - timeout-minutes: 90 + timeout-minutes: 60 env: SGLANG_USE_MODELSCOPE: true SGLANG_IS_IN_CI: true HF_ENDPOINT: https://hf-mirror.com TORCH_EXTENSIONS_DIR: /tmp/torch_extensions + PYTORCH_NPU_ALLOC_CONF: "expandable_segments:True" + STREAMS_PER_DEVICE: 32 run: | + export PATH="/usr/local/Ascend/8.3.RC1/compiler/bishengir/bin:${PATH}" cd test/srt - python3 run_suite.py --suite per-commit-2-ascend-npu + python3 run_suite.py --suite per-commit-2-ascend-npu --auto-partition-id ${{ matrix.part }} --auto-partition-size 3 per-commit-4-ascend-npu: if: github.event_name != 'pull_request' || contains(github.event.pull_request.labels.*.name, 'run-ci') @@ -106,22 +118,26 @@ jobs: CACHING_URL="cache-service.nginx-pypi-cache.svc.cluster.local" sed -Ei "s@(ports|archive).ubuntu.com@${CACHING_URL}:8081@g" /etc/apt/sources.list pip config set global.index-url http://${CACHING_URL}/pypi/simple - pip config set global.trusted-host ${CACHING_URL} + pip config set global.extra-index-url "https://pypi.tuna.tsinghua.edu.cn/simple https://mirrors.aliyun.com/pypi/simple/" + pip config set global.trusted-host "${CACHING_URL} pypi.tuna.tsinghua.edu.cn mirrors.aliyun.com" - bash scripts/ci/npu_ci_install_dependency.sh + bash scripts/ci/npu_ci_install_dependency.sh 910b # copy required file from our daily cache cp ~/.cache/modelscope/hub/datasets/otavia/ShareGPT_Vicuna_unfiltered/ShareGPT_V3_unfiltered_cleaned_split.json /tmp # copy download through proxy curl -o /tmp/test.jsonl -L https://gh-proxy.test.osinfra.cn/https://raw.githubusercontent.com/openai/grade-school-math/master/grade_school_math/data/test.jsonl - name: Run test - timeout-minutes: 120 + timeout-minutes: 60 env: SGLANG_USE_MODELSCOPE: true SGLANG_IS_IN_CI: true HF_ENDPOINT: https://hf-mirror.com TORCH_EXTENSIONS_DIR: /tmp/torch_extensions + PYTORCH_NPU_ALLOC_CONF: "expandable_segments:True" + STREAMS_PER_DEVICE: 32 run: | + export PATH="/usr/local/Ascend/8.3.RC1/compiler/bishengir/bin:${PATH}" cd test/srt python3 run_suite.py --suite per-commit-4-ascend-npu --timeout-per-file 3600 @@ -140,21 +156,25 @@ jobs: CACHING_URL="cache-service.nginx-pypi-cache.svc.cluster.local" sed -Ei "s@(ports|archive).ubuntu.com@${CACHING_URL}:8081@g" /etc/apt/sources.list pip config set global.index-url http://${CACHING_URL}/pypi/simple - pip config set global.trusted-host ${CACHING_URL} + pip config set global.extra-index-url "https://pypi.tuna.tsinghua.edu.cn/simple https://mirrors.aliyun.com/pypi/simple/" + pip config set global.trusted-host "${CACHING_URL} pypi.tuna.tsinghua.edu.cn mirrors.aliyun.com" - bash scripts/ci/npu_ci_install_dependency.sh + bash scripts/ci/npu_ci_install_dependency.sh a3 # copy required file from our daily cache cp ~/.cache/modelscope/hub/datasets/otavia/ShareGPT_Vicuna_unfiltered/ShareGPT_V3_unfiltered_cleaned_split.json /tmp # copy download through proxy curl -o /tmp/test.jsonl -L https://gh-proxy.test.osinfra.cn/https://raw.githubusercontent.com/openai/grade-school-math/master/grade_school_math/data/test.jsonl - name: Run test - timeout-minutes: 90 + timeout-minutes: 60 env: SGLANG_USE_MODELSCOPE: true SGLANG_IS_IN_CI: true HF_ENDPOINT: https://hf-mirror.com TORCH_EXTENSIONS_DIR: /tmp/torch_extensions + PYTORCH_NPU_ALLOC_CONF: "expandable_segments:True" + STREAMS_PER_DEVICE: 32 run: | + export PATH="/usr/local/Ascend/8.3.RC1/compiler/bishengir/bin:${PATH}" cd test/srt - python3 run_suite.py --suite per-commit-16-ascend-a3 --timeout-per-file 5400 + python3 run_suite.py --suite per-commit-16-ascend-a3 --timeout-per-file 3600 diff --git a/.github/workflows/release-docker-npu-nightly.yml b/.github/workflows/release-docker-npu-nightly.yml index dff45f2ac..3a1dee3b0 100644 --- a/.github/workflows/release-docker-npu-nightly.yml +++ b/.github/workflows/release-docker-npu-nightly.yml @@ -73,6 +73,6 @@ jobs: push: ${{ github.repository == 'sgl-project/sglang' && github.event_name != 'pull_request' }} provenance: false build-args: | - SGLANG_KERNEL_NPU_TAG=20250913 + SGLANG_KERNEL_NPU_TAG=20250926 CANN_VERSION=${{ matrix.cann_version }} DEVICE_TYPE=${{ matrix.device_type }} diff --git a/.github/workflows/release-docker-npu.yml b/.github/workflows/release-docker-npu.yml index 8fa6a983e..1a6ce903b 100644 --- a/.github/workflows/release-docker-npu.yml +++ b/.github/workflows/release-docker-npu.yml @@ -69,6 +69,6 @@ jobs: push: ${{ github.repository == 'sgl-project/sglang' && github.event_name != 'pull_request' }} provenance: false build-args: | - SGLANG_KERNEL_NPU_TAG=20250913 + SGLANG_KERNEL_NPU_TAG=20250926 CANN_VERSION=${{ matrix.cann_version }} DEVICE_TYPE=${{ matrix.device_type }} diff --git a/docker/Dockerfile.npu b/docker/Dockerfile.npu index df9235607..210369de7 100644 --- a/docker/Dockerfile.npu +++ b/docker/Dockerfile.npu @@ -6,12 +6,13 @@ ARG PYTHON_VERSION=py3.11 FROM quay.io/ascend/cann:$CANN_VERSION-$DEVICE_TYPE-$OS-$PYTHON_VERSION # Update pip & apt sources +ARG DEVICE_TYPE ARG PIP_INDEX_URL="https://pypi.org/simple/" ARG APTMIRROR="" ARG MEMFABRIC_URL=https://sglang-ascend.obs.cn-east-3.myhuaweicloud.com/sglang/mf_adapter-1.0.0-cp311-cp311-linux_aarch64.whl ARG PYTORCH_VERSION=2.6.0 ARG TORCHVISION_VERSION=0.21.0 -ARG PTA_URL="https://gitee.com/ascend/pytorch/releases/download/v7.1.0.1-pytorch2.6.0/torch_npu-2.6.0.post1-cp311-cp311-manylinux_2_28_aarch64.whl" +ARG PTA_URL="https://sglang-ascend.obs.cn-east-3.myhuaweicloud.com/ops/torch_npu-2.6.0.post2%2Bgit95d6260-cp311-cp311-linux_aarch64.whl" ARG VLLM_TAG=v0.8.5 ARG TRITON_ASCEND_URL="https://sglang-ascend.obs.cn-east-3.myhuaweicloud.com/sglang/triton_ascend-3.2.0%2Bgitb0ea0850-cp311-cp311-linux_aarch64.whl" ARG BISHENG_URL="https://sglang-ascend.obs.cn-east-3.myhuaweicloud.com/sglang/Ascend-BiSheng-toolkit_aarch64.run" @@ -71,7 +72,7 @@ RUN git clone --depth 1 https://github.com/vllm-project/vllm.git --branch $VLLM_ # TODO: install from pypi released triton-ascend RUN pip install torch==$PYTORCH_VERSION torchvision==$TORCHVISION_VERSION --index-url https://download.pytorch.org/whl/cpu --no-cache-dir \ - && wget ${PTA_URL} && pip install "./torch_npu-2.6.0.post1-cp311-cp311-manylinux_2_28_aarch64.whl" --no-cache-dir \ + && wget ${PTA_URL} && pip install "./torch_npu-2.6.0.post2+git95d6260-cp311-cp311-linux_aarch64.whl" --no-cache-dir \ && python3 -m pip install --no-cache-dir attrs==24.2.0 numpy==1.26.4 scipy==1.13.1 decorator==5.1.1 psutil==6.0.0 pytest==8.3.2 pytest-xdist==3.6.1 pyyaml pybind11 \ && pip install ${TRITON_ASCEND_URL} --no-cache-dir @@ -92,6 +93,13 @@ RUN pip install wheel==0.45.1 && git clone --branch $SGLANG_KERNEL_NPU_TAG http && cd .. && rm -rf sgl-kernel-npu \ && cd "$(pip show deep-ep | awk '/^Location:/ {print $2}')" && ln -s deep_ep/deep_ep_cpp*.so +# Install CustomOps +RUN wget https://sglang-ascend.obs.cn-east-3.myhuaweicloud.com/ops/CANN-custom_ops-8.2.0.0-$DEVICE_TYPE-linux.aarch64.run && \ + chmod a+x ./CANN-custom_ops-8.2.0.0-$DEVICE_TYPE-linux.aarch64.run && \ + ./CANN-custom_ops-8.2.0.0-$DEVICE_TYPE-linux.aarch64.run --quiet --install-path=/usr/local/Ascend/ascend-toolkit/latest/opp && \ + wget https://sglang-ascend.obs.cn-east-3.myhuaweicloud.com/ops/custom_ops-1.0.$DEVICE_TYPE-cp311-cp311-linux_aarch64.whl && \ + pip install ./custom_ops-1.0.$DEVICE_TYPE-cp311-cp311-linux_aarch64.whl + # Install Bisheng RUN wget ${BISHENG_URL} && chmod a+x Ascend-BiSheng-toolkit_aarch64.run && ./Ascend-BiSheng-toolkit_aarch64.run --install && rm Ascend-BiSheng-toolkit_aarch64.run diff --git a/python/sglang/srt/layers/attention/ascend_backend.py b/python/sglang/srt/layers/attention/ascend_backend.py index 65490b017..bc118d6c5 100644 --- a/python/sglang/srt/layers/attention/ascend_backend.py +++ b/python/sglang/srt/layers/attention/ascend_backend.py @@ -356,6 +356,11 @@ class AscendAttnBackend(AttentionBackend): assert ( layer.qk_head_dim != layer.v_head_dim ), "FIA only supports qk_head_dim != v_head_dim" + num_token_padding = q.shape[0] + q, k, v = [ + data[: forward_batch.num_token_non_padded_cpu] for data in [q, k, v] + ] + q_nope, q_rope = q.split([layer.v_head_dim, self.qk_rope_head_dim], dim=-1) k_nope, k_rope = k.split([layer.v_head_dim, self.qk_rope_head_dim], dim=-1) @@ -375,6 +380,18 @@ class AscendAttnBackend(AttentionBackend): next_tokens=0, ) + attn_output = attn_output.reshape(-1, layer.tp_q_head_num, layer.v_head_dim) + if num_token_padding != forward_batch.num_token_non_padded_cpu: + attn_output = torch.cat( + [ + attn_output, + attn_output.new_zeros( + num_token_padding - attn_output.shape[0], + *attn_output.shape[1:], + ), + ], + dim=0, + ) return attn_output def forward_decode_graph( diff --git a/python/sglang/srt/mem_cache/allocator_ascend.py b/python/sglang/srt/mem_cache/allocator_ascend.py index 14fc1d1e3..2c606187a 100644 --- a/python/sglang/srt/mem_cache/allocator_ascend.py +++ b/python/sglang/srt/mem_cache/allocator_ascend.py @@ -119,7 +119,7 @@ class AscendPagedTokenToKVPoolAllocator(PagedTokenToKVPoolAllocator): assert len(torch.unique(out_indices)) == len(out_indices) self.free_pages = self.free_pages[num_new_pages_item:] - return out_indices + return out_indices.int() def alloc_decode( self, diff --git a/python/sglang/srt/mem_cache/common.py b/python/sglang/srt/mem_cache/common.py index 7dcd69410..979b697c3 100644 --- a/python/sglang/srt/mem_cache/common.py +++ b/python/sglang/srt/mem_cache/common.py @@ -347,11 +347,7 @@ def alloc_for_extend( else: # Paged allocation - build last_loc last_loc = [ - ( - t[-1:] - if len(t) > 0 - else torch.tensor([-1], device=batch.tree_cache.device) - ) + (t[-1:] if len(t) > 0 else torch.tensor([-1], device=batch.device)) for t in prefix_tensors ] out_cache_loc = alloc_paged_token_slots_extend( diff --git a/python/sglang/srt/model_executor/npu_graph_runner.py b/python/sglang/srt/model_executor/npu_graph_runner.py index 67a31c62f..db7dcd159 100644 --- a/python/sglang/srt/model_executor/npu_graph_runner.py +++ b/python/sglang/srt/model_executor/npu_graph_runner.py @@ -22,7 +22,7 @@ from typing import TYPE_CHECKING, Optional, Union import numpy as np import torch -from sglang.srt.configs.model_config import AttentionArch +from sglang.srt.configs.model_config import AttentionArch, is_deepseek_nsa from sglang.srt.model_executor.cuda_graph_runner import CudaGraphRunner logger = logging.getLogger(__name__) @@ -75,7 +75,7 @@ class NPUGraphRunner(CudaGraphRunner): self.positions[: self.raw_num_token].copy_(forward_batch.positions) # Replay - if self.model_runner.model_config.index_head_dim is None: + if not is_deepseek_nsa(self.model_runner.model_config.hf_config): seq_lens = forward_batch.seq_lens.cpu().tolist() + [0] * ( self.bs - self.raw_bs ) diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index 3100e1490..fb9cd4f6c 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -1357,6 +1357,7 @@ class DeepseekV2AttentionMLA(nn.Module): inner_state = self.mla_preprocess.forward( positions, hidden_states, forward_batch, zero_allocator ) + inner_state = (*inner_state, None) # add a position for topk_indices elif attn_forward_method == AttnForwardMethod.NPU_MLA_SPARSE: inner_state = self.forward_npu_sparse_prepare( positions, hidden_states, forward_batch, zero_allocator diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 043bec845..94179cf7f 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -628,6 +628,16 @@ class ServerArgs: self.chunked_prefill_size = 2048 if self.cuda_graph_max_bs is None: self.cuda_graph_max_bs = 8 + elif is_npu() and gpu_mem < 32 * 1024: + # Atlas A2B4 + # (chunked_prefill_size 32k, cuda_graph_max_bs 16 if tp < 4 else 64) + if self.chunked_prefill_size is None: + self.chunked_prefill_size = 32768 + if self.cuda_graph_max_bs is None: + if self.tp_size < 4: + self.cuda_graph_max_bs = 16 + else: + self.cuda_graph_max_bs = 64 elif gpu_mem < 35 * 1024: # A10, 4090, 5090 # (chunked_prefill_size 2k, cuda_graph_max_bs 16 if tp < 4 else 80) @@ -651,6 +661,16 @@ class ServerArgs: self.cuda_graph_max_bs = 32 else: self.cuda_graph_max_bs = 160 + elif is_npu() and gpu_mem < 64 * 1024: + # Atlas A2 and Atlas A3 + # (chunked_prefill_size 32k, cuda_graph_max_bs 64 if tp < 4 else 128) + if self.chunked_prefill_size is None: + self.chunked_prefill_size = 32768 + if self.cuda_graph_max_bs is None: + if self.tp_size < 4: + self.cuda_graph_max_bs = 64 + else: + self.cuda_graph_max_bs = 128 elif gpu_mem < 90 * 1024: # H100, A100 # (chunked_prefill_size 8k, cuda_graph_max_bs 256 if tp < 4 else 512) diff --git a/scripts/ci/npu_ci_install_dependency.sh b/scripts/ci/npu_ci_install_dependency.sh index 4246bb419..8782a3191 100755 --- a/scripts/ci/npu_ci_install_dependency.sh +++ b/scripts/ci/npu_ci_install_dependency.sh @@ -2,6 +2,7 @@ set -euo pipefail PIP_INSTALL="pip install --no-cache-dir" +DEVICE_TYPE=$1 # Install the required dependencies in CI. @@ -39,8 +40,8 @@ TORCHVISION_VERSION=0.21.0 ${PIP_INSTALL} torch==$PYTORCH_VERSION torchvision==$TORCHVISION_VERSION --index-url https://download.pytorch.org/whl/cpu PTA_VERSION="v7.1.0.1-pytorch2.6.0" -PTA_NAME="torch_npu-2.6.0.post1-cp311-cp311-manylinux_2_28_aarch64.whl" -PTA_URL="https://gitee.com/ascend/pytorch/releases/download/${PTA_VERSION}/${PTA_NAME}" +PTA_NAME="torch_npu-2.6.0.post2+git95d6260-cp311-cp311-linux_aarch64.whl" +PTA_URL="https://sglang-ascend.obs.cn-east-3.myhuaweicloud.com/ops/torch_npu-2.6.0.post2%2Bgit95d6260-cp311-cp311-linux_aarch64.whl" wget -O "${PTA_NAME}" "${PTA_URL}" && ${PIP_INSTALL} "./${PTA_NAME}" @@ -58,11 +59,20 @@ wget -O "${BISHENG_NAME}" "${BISHENG_URL}" && chmod a+x "${BISHENG_NAME}" && "./ ### Install sgl-kernel-npu -SGL_KERNEL_NPU_TAG="20250913" +SGL_KERNEL_NPU_TAG="20250926" git clone --depth 1 https://github.com/sgl-project/sgl-kernel-npu.git --branch ${SGL_KERNEL_NPU_TAG} +# pin wheel to 0.45.1 ref: https://github.com/pypa/wheel/issues/662 +pip install wheel==0.45.1 (cd sgl-kernel-npu && bash ./build.sh && pip install output/deep_ep*.whl output/sgl_kernel_npu*.whl && cd "$(pip show deep-ep | grep -E '^Location:' | awk '{print $2}')" && ln -s deep_ep/deep_ep_cpp*.so) +### Install CustomOps (TODO: to be removed once merged into sgl-kernel-npu) +wget https://sglang-ascend.obs.cn-east-3.myhuaweicloud.com/ops/CANN-custom_ops-8.2.0.0-$DEVICE_TYPE-linux.aarch64.run +chmod a+x ./CANN-custom_ops-8.2.0.0-$DEVICE_TYPE-linux.aarch64.run +./CANN-custom_ops-8.2.0.0-$DEVICE_TYPE-linux.aarch64.run --quiet --install-path=/usr/local/Ascend/ascend-toolkit/latest/opp +wget https://sglang-ascend.obs.cn-east-3.myhuaweicloud.com/ops/custom_ops-1.0.$DEVICE_TYPE-cp311-cp311-linux_aarch64.whl +pip install ./custom_ops-1.0.$DEVICE_TYPE-cp311-cp311-linux_aarch64.whl + ### Install SGLang rm -rf python/pyproject.toml && mv python/pyproject_other.toml python/pyproject.toml ${PIP_INSTALL} -v -e "python[srt_npu]" diff --git a/test/srt/ascend/test_ascend_deepep.py b/test/srt/ascend/test_ascend_deepep.py index de51e35b3..eb05f559e 100644 --- a/test/srt/ascend/test_ascend_deepep.py +++ b/test/srt/ascend/test_ascend_deepep.py @@ -38,13 +38,10 @@ class TestAscendDeepEP(CustomTestCase): "--quantization", "w8a8_int8", "--mem-fraction-static", - 0.9, - "--max-running-requests", - 32, + 0.8, "--disable-radix-cache", "--chunked-prefill-size", 32768, - "--disable-cuda-graph", "--tp-size", 16, "--dp-size", @@ -58,7 +55,7 @@ class TestAscendDeepEP(CustomTestCase): ] cls.extra_envs = { - "HCCL_BUFFSIZE": "500", + "HCCL_BUFFSIZE": "1000", "SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK": "32", "SGLANG_NPU_USE_MLAPO": "1", }