Skip to content

Commit 3cceaa3

Browse files
[Bugfix] Fix Qwen3/DSV3/DSV3.2 model support (sgl-project#11510)
1 parent b0d20cd commit 3cceaa3

12 files changed

Lines changed: 102 additions & 33 deletions

File tree

.github/workflows/pr-test-npu.yml

Lines changed: 33 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -38,9 +38,10 @@ jobs:
3838
CACHING_URL="cache-service.nginx-pypi-cache.svc.cluster.local"
3939
sed -Ei "s@(ports|archive).ubuntu.com@${CACHING_URL}:8081@g" /etc/apt/sources.list
4040
pip config set global.index-url http://${CACHING_URL}/pypi/simple
41-
pip config set global.trusted-host ${CACHING_URL}
41+
pip config set global.extra-index-url "https://pypi.tuna.tsinghua.edu.cn/simple https://mirrors.aliyun.com/pypi/simple/"
42+
pip config set global.trusted-host "${CACHING_URL} pypi.tuna.tsinghua.edu.cn mirrors.aliyun.com"
4243
43-
bash scripts/ci/npu_ci_install_dependency.sh
44+
bash scripts/ci/npu_ci_install_dependency.sh 910b
4445
# copy required file from our daily cache
4546
cp ~/.cache/modelscope/hub/datasets/otavia/ShareGPT_Vicuna_unfiltered/ShareGPT_V3_unfiltered_cleaned_split.json /tmp
4647
# copy download through proxy
@@ -53,13 +54,20 @@ jobs:
5354
SGLANG_IS_IN_CI: true
5455
HF_ENDPOINT: https://hf-mirror.com
5556
TORCH_EXTENSIONS_DIR: /tmp/torch_extensions
57+
PYTORCH_NPU_ALLOC_CONF: "expandable_segments:True"
58+
STREAMS_PER_DEVICE: 32
5659
run: |
60+
export PATH="/usr/local/Ascend/8.3.RC1/compiler/bishengir/bin:${PATH}"
5761
cd test/srt
5862
python3 run_suite.py --suite per-commit-1-ascend-npu
5963
6064
per-commit-2-ascend-npu:
6165
if: github.event_name != 'pull_request' || contains(github.event.pull_request.labels.*.name, 'run-ci')
6266
runs-on: linux-arm64-npu-2
67+
strategy:
68+
fail-fast: false
69+
matrix:
70+
part: [0, 1, 2]
6371
container:
6472
image: swr.cn-southwest-2.myhuaweicloud.com/base_image/ascend-ci/cann:8.2.rc1-910b-ubuntu22.04-py3.11
6573
steps:
@@ -72,24 +80,28 @@ jobs:
7280
CACHING_URL="cache-service.nginx-pypi-cache.svc.cluster.local"
7381
sed -Ei "s@(ports|archive).ubuntu.com@${CACHING_URL}:8081@g" /etc/apt/sources.list
7482
pip config set global.index-url http://${CACHING_URL}/pypi/simple
75-
pip config set global.trusted-host ${CACHING_URL}
83+
pip config set global.extra-index-url "https://pypi.tuna.tsinghua.edu.cn/simple https://mirrors.aliyun.com/pypi/simple/"
84+
pip config set global.trusted-host "${CACHING_URL} pypi.tuna.tsinghua.edu.cn mirrors.aliyun.com"
7685
77-
bash scripts/ci/npu_ci_install_dependency.sh
86+
bash scripts/ci/npu_ci_install_dependency.sh 910b
7887
# copy required file from our daily cache
7988
cp ~/.cache/modelscope/hub/datasets/otavia/ShareGPT_Vicuna_unfiltered/ShareGPT_V3_unfiltered_cleaned_split.json /tmp
8089
# copy download through proxy
8190
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
8291
8392
- name: Run test
84-
timeout-minutes: 90
93+
timeout-minutes: 60
8594
env:
8695
SGLANG_USE_MODELSCOPE: true
8796
SGLANG_IS_IN_CI: true
8897
HF_ENDPOINT: https://hf-mirror.com
8998
TORCH_EXTENSIONS_DIR: /tmp/torch_extensions
99+
PYTORCH_NPU_ALLOC_CONF: "expandable_segments:True"
100+
STREAMS_PER_DEVICE: 32
90101
run: |
102+
export PATH="/usr/local/Ascend/8.3.RC1/compiler/bishengir/bin:${PATH}"
91103
cd test/srt
92-
python3 run_suite.py --suite per-commit-2-ascend-npu
104+
python3 run_suite.py --suite per-commit-2-ascend-npu --auto-partition-id ${{ matrix.part }} --auto-partition-size 3
93105
94106
per-commit-4-ascend-npu:
95107
if: github.event_name != 'pull_request' || contains(github.event.pull_request.labels.*.name, 'run-ci')
@@ -106,22 +118,26 @@ jobs:
106118
CACHING_URL="cache-service.nginx-pypi-cache.svc.cluster.local"
107119
sed -Ei "s@(ports|archive).ubuntu.com@${CACHING_URL}:8081@g" /etc/apt/sources.list
108120
pip config set global.index-url http://${CACHING_URL}/pypi/simple
109-
pip config set global.trusted-host ${CACHING_URL}
121+
pip config set global.extra-index-url "https://pypi.tuna.tsinghua.edu.cn/simple https://mirrors.aliyun.com/pypi/simple/"
122+
pip config set global.trusted-host "${CACHING_URL} pypi.tuna.tsinghua.edu.cn mirrors.aliyun.com"
110123
111-
bash scripts/ci/npu_ci_install_dependency.sh
124+
bash scripts/ci/npu_ci_install_dependency.sh 910b
112125
# copy required file from our daily cache
113126
cp ~/.cache/modelscope/hub/datasets/otavia/ShareGPT_Vicuna_unfiltered/ShareGPT_V3_unfiltered_cleaned_split.json /tmp
114127
# copy download through proxy
115128
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
116129
117130
- name: Run test
118-
timeout-minutes: 120
131+
timeout-minutes: 60
119132
env:
120133
SGLANG_USE_MODELSCOPE: true
121134
SGLANG_IS_IN_CI: true
122135
HF_ENDPOINT: https://hf-mirror.com
123136
TORCH_EXTENSIONS_DIR: /tmp/torch_extensions
137+
PYTORCH_NPU_ALLOC_CONF: "expandable_segments:True"
138+
STREAMS_PER_DEVICE: 32
124139
run: |
140+
export PATH="/usr/local/Ascend/8.3.RC1/compiler/bishengir/bin:${PATH}"
125141
cd test/srt
126142
python3 run_suite.py --suite per-commit-4-ascend-npu --timeout-per-file 3600
127143
@@ -140,21 +156,25 @@ jobs:
140156
CACHING_URL="cache-service.nginx-pypi-cache.svc.cluster.local"
141157
sed -Ei "s@(ports|archive).ubuntu.com@${CACHING_URL}:8081@g" /etc/apt/sources.list
142158
pip config set global.index-url http://${CACHING_URL}/pypi/simple
143-
pip config set global.trusted-host ${CACHING_URL}
159+
pip config set global.extra-index-url "https://pypi.tuna.tsinghua.edu.cn/simple https://mirrors.aliyun.com/pypi/simple/"
160+
pip config set global.trusted-host "${CACHING_URL} pypi.tuna.tsinghua.edu.cn mirrors.aliyun.com"
144161
145-
bash scripts/ci/npu_ci_install_dependency.sh
162+
bash scripts/ci/npu_ci_install_dependency.sh a3
146163
# copy required file from our daily cache
147164
cp ~/.cache/modelscope/hub/datasets/otavia/ShareGPT_Vicuna_unfiltered/ShareGPT_V3_unfiltered_cleaned_split.json /tmp
148165
# copy download through proxy
149166
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
150167
151168
- name: Run test
152-
timeout-minutes: 90
169+
timeout-minutes: 60
153170
env:
154171
SGLANG_USE_MODELSCOPE: true
155172
SGLANG_IS_IN_CI: true
156173
HF_ENDPOINT: https://hf-mirror.com
157174
TORCH_EXTENSIONS_DIR: /tmp/torch_extensions
175+
PYTORCH_NPU_ALLOC_CONF: "expandable_segments:True"
176+
STREAMS_PER_DEVICE: 32
158177
run: |
178+
export PATH="/usr/local/Ascend/8.3.RC1/compiler/bishengir/bin:${PATH}"
159179
cd test/srt
160-
python3 run_suite.py --suite per-commit-16-ascend-a3 --timeout-per-file 5400
180+
python3 run_suite.py --suite per-commit-16-ascend-a3 --timeout-per-file 3600

.github/workflows/release-docker-npu-nightly.yml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -73,6 +73,6 @@ jobs:
7373
push: ${{ github.repository == 'sgl-project/sglang' && github.event_name != 'pull_request' }}
7474
provenance: false
7575
build-args: |
76-
SGLANG_KERNEL_NPU_TAG=20250913
76+
SGLANG_KERNEL_NPU_TAG=20250926
7777
CANN_VERSION=${{ matrix.cann_version }}
7878
DEVICE_TYPE=${{ matrix.device_type }}

.github/workflows/release-docker-npu.yml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -69,6 +69,6 @@ jobs:
6969
push: ${{ github.repository == 'sgl-project/sglang' && github.event_name != 'pull_request' }}
7070
provenance: false
7171
build-args: |
72-
SGLANG_KERNEL_NPU_TAG=20250913
72+
SGLANG_KERNEL_NPU_TAG=20250926
7373
CANN_VERSION=${{ matrix.cann_version }}
7474
DEVICE_TYPE=${{ matrix.device_type }}

docker/Dockerfile.npu

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -6,12 +6,13 @@ ARG PYTHON_VERSION=py3.11
66
FROM quay.io/ascend/cann:$CANN_VERSION-$DEVICE_TYPE-$OS-$PYTHON_VERSION
77

88
# Update pip & apt sources
9+
ARG DEVICE_TYPE
910
ARG PIP_INDEX_URL="https://pypi.org/simple/"
1011
ARG APTMIRROR=""
1112
ARG MEMFABRIC_URL=https://sglang-ascend.obs.cn-east-3.myhuaweicloud.com/sglang/mf_adapter-1.0.0-cp311-cp311-linux_aarch64.whl
1213
ARG PYTORCH_VERSION=2.6.0
1314
ARG TORCHVISION_VERSION=0.21.0
14-
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"
15+
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"
1516
ARG VLLM_TAG=v0.8.5
1617
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"
1718
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_
7172

7273
# TODO: install from pypi released triton-ascend
7374
RUN pip install torch==$PYTORCH_VERSION torchvision==$TORCHVISION_VERSION --index-url https://download.pytorch.org/whl/cpu --no-cache-dir \
74-
&& wget ${PTA_URL} && pip install "./torch_npu-2.6.0.post1-cp311-cp311-manylinux_2_28_aarch64.whl" --no-cache-dir \
75+
&& wget ${PTA_URL} && pip install "./torch_npu-2.6.0.post2+git95d6260-cp311-cp311-linux_aarch64.whl" --no-cache-dir \
7576
&& 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 \
7677
&& pip install ${TRITON_ASCEND_URL} --no-cache-dir
7778

@@ -92,6 +93,13 @@ RUN pip install wheel==0.45.1 && git clone --branch $SGLANG_KERNEL_NPU_TAG http
9293
&& cd .. && rm -rf sgl-kernel-npu \
9394
&& cd "$(pip show deep-ep | awk '/^Location:/ {print $2}')" && ln -s deep_ep/deep_ep_cpp*.so
9495

96+
# Install CustomOps
97+
RUN wget https://sglang-ascend.obs.cn-east-3.myhuaweicloud.com/ops/CANN-custom_ops-8.2.0.0-$DEVICE_TYPE-linux.aarch64.run && \
98+
chmod a+x ./CANN-custom_ops-8.2.0.0-$DEVICE_TYPE-linux.aarch64.run && \
99+
./CANN-custom_ops-8.2.0.0-$DEVICE_TYPE-linux.aarch64.run --quiet --install-path=/usr/local/Ascend/ascend-toolkit/latest/opp && \
100+
wget https://sglang-ascend.obs.cn-east-3.myhuaweicloud.com/ops/custom_ops-1.0.$DEVICE_TYPE-cp311-cp311-linux_aarch64.whl && \
101+
pip install ./custom_ops-1.0.$DEVICE_TYPE-cp311-cp311-linux_aarch64.whl
102+
95103
# Install Bisheng
96104
RUN wget ${BISHENG_URL} && chmod a+x Ascend-BiSheng-toolkit_aarch64.run && ./Ascend-BiSheng-toolkit_aarch64.run --install && rm Ascend-BiSheng-toolkit_aarch64.run
97105

python/sglang/srt/layers/attention/ascend_backend.py

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -356,6 +356,11 @@ def forward_extend(
356356
assert (
357357
layer.qk_head_dim != layer.v_head_dim
358358
), "FIA only supports qk_head_dim != v_head_dim"
359+
num_token_padding = q.shape[0]
360+
q, k, v = [
361+
data[: forward_batch.num_token_non_padded_cpu] for data in [q, k, v]
362+
]
363+
359364
q_nope, q_rope = q.split([layer.v_head_dim, self.qk_rope_head_dim], dim=-1)
360365
k_nope, k_rope = k.split([layer.v_head_dim, self.qk_rope_head_dim], dim=-1)
361366

@@ -375,6 +380,18 @@ def forward_extend(
375380
next_tokens=0,
376381
)
377382

383+
attn_output = attn_output.reshape(-1, layer.tp_q_head_num, layer.v_head_dim)
384+
if num_token_padding != forward_batch.num_token_non_padded_cpu:
385+
attn_output = torch.cat(
386+
[
387+
attn_output,
388+
attn_output.new_zeros(
389+
num_token_padding - attn_output.shape[0],
390+
*attn_output.shape[1:],
391+
),
392+
],
393+
dim=0,
394+
)
378395
return attn_output
379396

380397
def forward_decode_graph(

python/sglang/srt/mem_cache/allocator_ascend.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -119,7 +119,7 @@ def alloc_extend(
119119
assert len(torch.unique(out_indices)) == len(out_indices)
120120

121121
self.free_pages = self.free_pages[num_new_pages_item:]
122-
return out_indices
122+
return out_indices.int()
123123

124124
def alloc_decode(
125125
self,

python/sglang/srt/mem_cache/common.py

Lines changed: 1 addition & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -347,11 +347,7 @@ def alloc_for_extend(
347347
else:
348348
# Paged allocation - build last_loc
349349
last_loc = [
350-
(
351-
t[-1:]
352-
if len(t) > 0
353-
else torch.tensor([-1], device=batch.tree_cache.device)
354-
)
350+
(t[-1:] if len(t) > 0 else torch.tensor([-1], device=batch.device))
355351
for t in prefix_tensors
356352
]
357353
out_cache_loc = alloc_paged_token_slots_extend(

python/sglang/srt/model_executor/npu_graph_runner.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,7 @@
2222
import numpy as np
2323
import torch
2424

25-
from sglang.srt.configs.model_config import AttentionArch
25+
from sglang.srt.configs.model_config import AttentionArch, is_deepseek_nsa
2626
from sglang.srt.model_executor.cuda_graph_runner import CudaGraphRunner
2727

2828
logger = logging.getLogger(__name__)
@@ -75,7 +75,7 @@ def replay(
7575
self.positions[: self.raw_num_token].copy_(forward_batch.positions)
7676

7777
# Replay
78-
if self.model_runner.model_config.index_head_dim is None:
78+
if not is_deepseek_nsa(self.model_runner.model_config.hf_config):
7979
seq_lens = forward_batch.seq_lens.cpu().tolist() + [0] * (
8080
self.bs - self.raw_bs
8181
)

python/sglang/srt/models/deepseek_v2.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1357,6 +1357,7 @@ def forward_prepare(
13571357
inner_state = self.mla_preprocess.forward(
13581358
positions, hidden_states, forward_batch, zero_allocator
13591359
)
1360+
inner_state = (*inner_state, None) # add a position for topk_indices
13601361
elif attn_forward_method == AttnForwardMethod.NPU_MLA_SPARSE:
13611362
inner_state = self.forward_npu_sparse_prepare(
13621363
positions, hidden_states, forward_batch, zero_allocator

python/sglang/srt/server_args.py

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -628,6 +628,16 @@ def _handle_gpu_memory_settings(self, gpu_mem):
628628
self.chunked_prefill_size = 2048
629629
if self.cuda_graph_max_bs is None:
630630
self.cuda_graph_max_bs = 8
631+
elif is_npu() and gpu_mem < 32 * 1024:
632+
# Atlas A2B4
633+
# (chunked_prefill_size 32k, cuda_graph_max_bs 16 if tp < 4 else 64)
634+
if self.chunked_prefill_size is None:
635+
self.chunked_prefill_size = 32768
636+
if self.cuda_graph_max_bs is None:
637+
if self.tp_size < 4:
638+
self.cuda_graph_max_bs = 16
639+
else:
640+
self.cuda_graph_max_bs = 64
631641
elif gpu_mem < 35 * 1024:
632642
# A10, 4090, 5090
633643
# (chunked_prefill_size 2k, cuda_graph_max_bs 16 if tp < 4 else 80)
@@ -651,6 +661,16 @@ def _handle_gpu_memory_settings(self, gpu_mem):
651661
self.cuda_graph_max_bs = 32
652662
else:
653663
self.cuda_graph_max_bs = 160
664+
elif is_npu() and gpu_mem < 64 * 1024:
665+
# Atlas A2 and Atlas A3
666+
# (chunked_prefill_size 32k, cuda_graph_max_bs 64 if tp < 4 else 128)
667+
if self.chunked_prefill_size is None:
668+
self.chunked_prefill_size = 32768
669+
if self.cuda_graph_max_bs is None:
670+
if self.tp_size < 4:
671+
self.cuda_graph_max_bs = 64
672+
else:
673+
self.cuda_graph_max_bs = 128
654674
elif gpu_mem < 90 * 1024:
655675
# H100, A100
656676
# (chunked_prefill_size 8k, cuda_graph_max_bs 256 if tp < 4 else 512)

0 commit comments

Comments
 (0)