diff --git a/.github/workflows/release-docker-npu-nightly.yml b/.github/workflows/release-docker-npu-nightly.yml index 1778038ac..ef9bc0c06 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=20251206 + SGLANG_KERNEL_NPU_TAG=2025.12.31 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 c3b936986..850efbae0 100644 --- a/.github/workflows/release-docker-npu.yml +++ b/.github/workflows/release-docker-npu.yml @@ -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 }} diff --git a/python/sglang/srt/layers/rotary_embedding.py b/python/sglang/srt/layers/rotary_embedding.py index 480579e01..faa35d3c0 100644 --- a/python/sglang/srt/layers/rotary_embedding.py +++ b/python/sglang/srt/layers/rotary_embedding.py @@ -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: diff --git a/python/sglang/srt/layers/vocab_parallel_embedding.py b/python/sglang/srt/layers/vocab_parallel_embedding.py index f171a15f2..192970904 100644 --- a/python/sglang/srt/layers/vocab_parallel_embedding.py +++ b/python/sglang/srt/layers/vocab_parallel_embedding.py @@ -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, diff --git a/python/sglang/srt/models/llama.py b/python/sglang/srt/models/llama.py index 4176af28d..ffcf73d46 100644 --- a/python/sglang/srt/models/llama.py +++ b/python/sglang/srt/models/llama.py @@ -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 diff --git a/python/sglang/srt/models/qwen3.py b/python/sglang/srt/models/qwen3.py index 9220831f6..aad79ce23 100644 --- a/python/sglang/srt/models/qwen3.py +++ b/python/sglang/srt/models/qwen3.py @@ -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: diff --git a/python/sglang/srt/models/qwen3_moe.py b/python/sglang/srt/models/qwen3_moe.py index e11678a9e..3648b2cf9 100644 --- a/python/sglang/srt/models/qwen3_moe.py +++ b/python/sglang/srt/models/qwen3_moe.py @@ -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), ) diff --git a/scripts/ci/npu_ci_install_dependency.sh b/scripts/ci/npu_ci_install_dependency.sh index 6172db3b4..8cc8fbf39 100755 --- a/scripts/ci/npu_ci_install_dependency.sh +++ b/scripts/ci/npu_ci_install_dependency.sh @@ -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)