chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

[NPU][Bugfix] Fix qwen3 error when enable-dp-lm-head (#16115)

Browse Source
This commit is contained in:
chenxu214
2026-01-08 15:15:43 +08:00
committed by GitHub
parent 3d51ae18a1
commit 7dd679cbb9
8 changed files with 52 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
.github/workflows/release-docker-npu-nightly.yml vendored
View File

@@ -73,6 +73,6 @@ jobs:
push: ${{ github.repository == 'sgl-project/sglang' && github.event_name != 'pull_request' }}
provenance: false
build-args: |
SGLANG_KERNEL_NPU_TAG=20251206
SGLANG_KERNEL_NPU_TAG=2025.12.31
CANN_VERSION=${{ matrix.cann_version }}
DEVICE_TYPE=${{ matrix.device_type }}

2
.github/workflows/release-docker-npu.yml vendored
View File

@@ -81,7 +81,7 @@ jobs:
push: ${{ github.repository == 'sgl-project/sglang' && github.event_name != 'pull_request' }}
provenance: false
build-args: |
SGLANG_KERNEL_NPU_TAG=20251206
SGLANG_KERNEL_NPU_TAG=2025.12.31
CANN_VERSION=${{ matrix.cann_version }}
DEVICE_TYPE=${{ matrix.device_type }}
SGLANG_TAG=${{ steps.version.outputs.version }}

4
python/sglang/srt/layers/rotary_embedding.py
View File

@@ -294,8 +294,8 @@ class RotaryEmbedding(MultiPlatformOp):
assert (
fused_set_kv_buffer_arg is None
), "fused_set_kv_buffer_arg is not supported for npu implementation"
rotary_mode = "half"
if query.dtype == torch.bfloat16 and self.cos_sin_cache.dtype == torch.float:
return self.forward_native(positions, query, key, offsets)
if self.is_neox_style:
rotary_mode = "half"
else:

4
python/sglang/srt/layers/vocab_parallel_embedding.py
View File

@@ -31,6 +31,7 @@ from sglang.srt.utils import (
cpu_has_amx_support,
get_compiler_backend,
is_cpu,
is_npu,
set_weight_attrs,
)
@@ -38,6 +39,7 @@ DEFAULT_VOCAB_PADDING_SIZE = 64
_is_cpu_amx_available = cpu_has_amx_support()
_is_cpu = is_cpu()
_is_npu = is_npu()
logger = logging.getLogger(__name__)
@@ -123,7 +125,7 @@ class VocabParallelEmbeddingShardIndices:
assert self.num_added_elements <= self.num_added_elements_padded
@torch.compile(dynamic=True, backend=get_compiler_backend())
@torch.compile(dynamic=True, backend=get_compiler_backend(), disable=_is_npu)
def get_masked_input_and_mask(
input_: torch.Tensor,
org_vocab_start_index: int,

41
python/sglang/srt/models/llama.py
View File

@@ -52,10 +52,14 @@ from sglang.srt.model_loader.weight_utils import (
maybe_remap_kv_scale_name,
)
from sglang.srt.server_args import get_global_server_args
from sglang.srt.utils import add_prefix, make_layers
from sglang.srt.utils import add_prefix, is_npu, make_layers
from sglang.utils import get_exception_traceback
logger = logging.getLogger(__name__)
_is_npu = is_npu()
if _is_npu:
from sgl_kernel_npu.norm.split_qkv_rmsnorm_rope import split_qkv_rmsnorm_rope
class LlamaMLP(nn.Module):
@@ -185,15 +189,44 @@ class LlamaAttention(nn.Module):
prefix=add_prefix("attn", prefix),
)
def forward_prepare_native(self, positions, hidden_states):
qkv, _ = self.qkv_proj(hidden_states)
q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
q, k = self.rotary_emb(positions, q, k)
return q, k, v
def forward_prepare_npu(self, positions, hidden_states, forward_batch):
qkv, _ = self.qkv_proj(hidden_states)
if self.attn.layer_id == forward_batch.token_to_kv_pool.start_layer:
self.rotary_emb.get_cos_sin_with_position(positions)
q, k, v = split_qkv_rmsnorm_rope(
qkv,
self.rotary_emb.position_sin,
self.rotary_emb.position_cos,
self.q_size,
self.kv_size,
self.head_dim,
)
return q, k, v
def forward(
self,
positions: torch.Tensor,
hidden_states: torch.Tensor,
forward_batch: ForwardBatch,
) -> torch.Tensor:
qkv, _ = self.qkv_proj(hidden_states)
q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
q, k = self.rotary_emb(positions, q, k)
if not _is_npu or not hasattr(self.rotary_emb, "get_cos_sin_with_position"):
q, k, v = self.forward_prepare_native(
positions=positions,
hidden_states=hidden_states,
)
else:
q, k, v = self.forward_prepare_npu(
positions=positions,
hidden_states=hidden_states,
forward_batch=forward_batch,
)
attn_output = self.attn(q, k, v, forward_batch)
output, _ = self.o_proj(attn_output)
return output

7
python/sglang/srt/models/qwen3.py
View File

@@ -161,12 +161,12 @@ class Qwen3Attention(nn.Module):
qkv,
self.rotary_emb.position_sin,
self.rotary_emb.position_cos,
self.q_norm.weight,
self.k_norm.weight,
self.q_size,
self.kv_size,
self.head_dim,
self.q_norm.variance_epsilon,
eps=self.q_norm.variance_epsilon,
q_weight=self.q_norm.weight,
k_weight=self.k_norm.weight,
q_bias=getattr(self.q_norm, "bias", None),
k_bias=getattr(self.k_norm, "bias", None),
)
@@ -372,6 +372,7 @@ class Qwen3ForCausalLM(nn.Module):
config.vocab_size,
config.hidden_size,
quant_config=quant_config,
use_attn_tp_group=get_global_server_args().enable_dp_lm_head,
prefix=add_prefix("lm_head", prefix),
)
else:

6
python/sglang/srt/models/qwen3_moe.py
View File

@@ -523,12 +523,12 @@ class Qwen3MoeAttention(nn.Module):
qkv,
self.rotary_emb.position_sin,
self.rotary_emb.position_cos,
self.q_norm.weight,
self.k_norm.weight,
self.q_size,
self.kv_size,
self.head_dim,
self.q_norm.variance_epsilon,
eps=self.q_norm.variance_epsilon,
q_weight=self.q_norm.weight,
k_weight=self.k_norm.weight,
q_bias=getattr(self.q_norm, "bias", None),
k_bias=getattr(self.k_norm, "bias", None),
)

2
scripts/ci/npu_ci_install_dependency.sh
View File

@@ -49,7 +49,7 @@ wget -O "${BISHENG_NAME}" "${BISHENG_URL}" && chmod a+x "${BISHENG_NAME}" && "./
### Install sgl-kernel-npu
SGL_KERNEL_NPU_TAG="20251206"
SGL_KERNEL_NPU_TAG="2025.12.31"
git clone --depth 1 https://github.com/sgl-project/sgl-kernel-npu.git --branch ${SGL_KERNEL_NPU_TAG}
(cd sgl-kernel-npu && bash ./build.sh && ${PIP_INSTALL} output/deep_ep*.whl output/sgl_kernel_npu*.whl && cd "$(python3 -m pip show deep-ep | grep -E '^Location:' | awk '{print $2}')" && ln -s deep_ep/deep_ep_cpp*.so)