From c05d3afb5d8bc9b354cb924bf4e467ecb8ce424f Mon Sep 17 00:00:00 2001 From: ZhengdQin <46387172+ZhengdQin@users.noreply.github.com> Date: Fri, 12 Dec 2025 14:52:16 +0800 Subject: [PATCH] [NPU] optimization for dsv3.2 (#14572) --- .../npu/attention/ascend_backend.py | 8 ++- .../modules/deepseek_v2_attention_mla_npu.py | 52 +++++++++------ .../srt/hardware_backend/npu/moe/topk.py | 15 ++--- .../npu/quantization/modelslim.py | 8 ++- .../sglang/srt/hardware_backend/npu/utils.py | 8 +++ .../srt/layers/attention/nsa/nsa_indexer.py | 64 +++++++++++++++---- python/sglang/srt/layers/layernorm.py | 14 +--- python/sglang/srt/layers/linear.py | 5 ++ .../quantization/compressed_tensors/utils.py | 2 +- python/sglang/srt/models/deepseek_nextn.py | 2 +- python/sglang/srt/models/deepseek_v2.py | 29 +++++++-- 11 files changed, 140 insertions(+), 67 deletions(-) diff --git a/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py b/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py index 0e21e680c..8342e01f7 100644 --- a/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py +++ b/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py @@ -573,8 +573,12 @@ class AscendAttnBackend(AttentionBackend): key_rope=k_pe, sparse_indices=topk_indices, scale_value=layer.scaling, - actual_seq_lengths_query=actual_seq_qlen, - actual_seq_lengths_kv=actual_seq_lengths_kv.to(q.device), + actual_seq_lengths_query=actual_seq_qlen.to( + device=q_nope.device, dtype=torch.int32 + ), + actual_seq_lengths_kv=actual_seq_lengths_kv.to( + device=q_nope.device, dtype=torch.int32 + ), block_table=self.forward_metadata.block_tables, sparse_block_size=1, layout_query="TND", diff --git a/python/sglang/srt/hardware_backend/npu/modules/deepseek_v2_attention_mla_npu.py b/python/sglang/srt/hardware_backend/npu/modules/deepseek_v2_attention_mla_npu.py index 8b916d8de..58cd9194e 100644 --- a/python/sglang/srt/hardware_backend/npu/modules/deepseek_v2_attention_mla_npu.py +++ b/python/sglang/srt/hardware_backend/npu/modules/deepseek_v2_attention_mla_npu.py @@ -273,39 +273,51 @@ def forward_dsa_prepare_npu( m.qk_rope_head_dim, m.quant_config, ) - ( - q_pe, - k_pe, - q_nope_out, - k_nope, - forward_batch, - zero_allocator, - positions, - ) = m.mla_preprocess.forward( - positions, hidden_states, forward_batch, zero_allocator - ) - + mla_event = torch.npu.Event() + mla_event.record() + with torch.npu.stream(m.alt_stream): + torch.npu.current_stream().wait_event(mla_event) + ( + q_pe, + k_pe, + q_nope_out, + k_nope, + forward_batch, + zero_allocator, + positions, + ) = m.mla_preprocess.forward( + positions, hidden_states, forward_batch, zero_allocator + ) fused_qkv_a_proj_out = m.fused_qkv_a_proj_with_mqa(hidden_states)[0] q, _ = fused_qkv_a_proj_out.split( [m.q_lora_rank, m.kv_lora_rank + m.qk_rope_head_dim], dim=-1 ) q_lora = m.q_a_layernorm(q) + torch.npu.current_stream().wait_stream(m.alt_stream) else: fused_qkv_a_proj_out = m.fused_qkv_a_proj_with_mqa(hidden_states)[0] q, latent_cache = fused_qkv_a_proj_out.split( [m.q_lora_rank, m.kv_lora_rank + m.qk_rope_head_dim], dim=-1 ) - k_nope = latent_cache[..., : m.kv_lora_rank] + # overlap qk norm q = m.q_a_layernorm(q) - k_nope = m.kv_a_layernorm(k_nope) q_lora = q.clone() # required for topk_indices - k_nope = k_nope.unsqueeze(1) - q = m.q_b_proj(q)[0].view(-1, m.num_local_heads, m.qk_head_dim) + + m.alt_stream.wait_stream(torch.npu.current_stream()) + with torch.npu.stream(m.alt_stream): + q = m.q_b_proj(q_lora)[0].view(-1, m.num_local_heads, m.qk_head_dim) + q.record_stream(m.alt_stream) + q_event = m.alt_stream.record_event() + + k_nope, k_pe = latent_cache.unsqueeze(1).split( + [m.kv_lora_rank, m.qk_rope_head_dim], dim=-1 + ) + k_nope = m.kv_a_layernorm(k_nope).unsqueeze(1) + torch.npu.current_stream().wait_event(q_event) q_nope, q_pe = q.split([m.qk_nope_head_dim, m.qk_rope_head_dim], dim=-1) - k_pe = latent_cache[..., m.kv_lora_rank :].unsqueeze(1) q_nope_out = torch.bmm(q_nope.transpose(0, 1), m.w_kc) @@ -367,7 +379,11 @@ def forward_dsa_core_npu( device=attn_output.device, ) - if not forward_batch.forward_mode.is_decode(): + if ( + forward_batch.forward_mode.is_extend() + and not forward_batch.forward_mode.is_draft_extend(include_v2=True) + and not forward_batch.forward_mode.is_target_verify() + ): attn_output = attn_output.transpose(0, 1) torch.bmm( attn_output, diff --git a/python/sglang/srt/hardware_backend/npu/moe/topk.py b/python/sglang/srt/hardware_backend/npu/moe/topk.py index e00c33a26..813c12f6a 100644 --- a/python/sglang/srt/hardware_backend/npu/moe/topk.py +++ b/python/sglang/srt/hardware_backend/npu/moe/topk.py @@ -39,8 +39,7 @@ def fused_topk_npu( topk_weights = topk_weights.to(torch.float32) elif use_grouped_topk and correction_bias is not None: - routed_scaling_factor = topk_config.routed_scaling_factor or 1 - + # Force set routed_scaling_factor = 1 to optimize renormalize topk_weights, topk_ids, _ = torch.ops.npu.npu_moe_gating_top_k( router_logits.to(torch.float32), k=topk_config.top_k, @@ -50,18 +49,12 @@ def fused_topk_npu( group_select_mode=1, renorm=0, norm_type=1, - routed_scaling_factor=routed_scaling_factor, + routed_scaling_factor=( + 1 if renormalize else topk_config.routed_scaling_factor + ), eps=float(1e-20), ) - if renormalize: - topk_weights_sum = ( - topk_weights.sum(dim=-1, keepdim=True) - if topk_config.num_fused_shared_experts == 0 - else topk_weights[:, :-1].sum(dim=-1, keepdim=True) - ) - topk_weights = topk_weights / topk_weights_sum - else: topk_config.torch_native = True return select_experts( diff --git a/python/sglang/srt/hardware_backend/npu/quantization/modelslim.py b/python/sglang/srt/hardware_backend/npu/quantization/modelslim.py index aae786836..6bf9b3d34 100644 --- a/python/sglang/srt/hardware_backend/npu/quantization/modelslim.py +++ b/python/sglang/srt/hardware_backend/npu/quantization/modelslim.py @@ -159,8 +159,10 @@ class ModelSlimConfig(QuantizationConfig): proj_name, packed_modules_mapping_subset[proj_name][0] ) self.is_dynamic = ( - self.quant_description[prefix_in_quant_config + ".weight"] + self.quant_description.get(prefix_in_quant_config + ".weight", "") == "W8A8_DYNAMIC" + or self.quant_description.get("quant_method", "") + == "modelslim" # TODO: This path is for compress-tensor config，needs refactor @zhengdqin ) if self.is_layer_skipped(prefix, packed_modules_mapping_subset): return UnquantizedLinearMethod() @@ -199,7 +201,7 @@ class ModelSlimConfig(QuantizationConfig): is_skipped = None for shard_prefix in shard_prefixes: is_shard_skipped = ( - self.quant_description[shard_prefix + ".weight"] == "FLOAT" + self.quant_description.get(shard_prefix + ".weight", "") == "FLOAT" ) if is_skipped is None: @@ -211,7 +213,7 @@ class ModelSlimConfig(QuantizationConfig): "to have the same precision." ) else: - is_skipped = self.quant_description[prefix + ".weight"] == "FLOAT" + is_skipped = self.quant_description.get(prefix + ".weight", "") == "FLOAT" assert is_skipped is not None return is_skipped diff --git a/python/sglang/srt/hardware_backend/npu/utils.py b/python/sglang/srt/hardware_backend/npu/utils.py index 070f1b252..715af908e 100644 --- a/python/sglang/srt/hardware_backend/npu/utils.py +++ b/python/sglang/srt/hardware_backend/npu/utils.py @@ -13,6 +13,7 @@ if TYPE_CHECKING: logger = logging.getLogger(__name__) _is_npu = is_npu() +indexer_weight_stream = None class NPUACLFormat(IntEnum): @@ -110,3 +111,10 @@ def npu_format_cast( import torch_npu return torch_npu.npu_format_cast(tensor, acl_format.value) + + +def get_indexer_weight_stream(): + global indexer_weight_stream + if indexer_weight_stream is None: + indexer_weight_stream = torch.npu.Stream() + return indexer_weight_stream diff --git a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py index c89c38f5d..2b2ae7e52 100644 --- a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py +++ b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py @@ -10,12 +10,18 @@ from sglang.srt.custom_op import CustomOp from sglang.srt.layers.layernorm import LayerNorm from sglang.srt.utils import add_prefix, ceil_align, is_cuda, is_hip, is_npu +global _use_multi_stream + if is_cuda(): try: import deep_gemm except ImportError as e: deep_gemm = e +if is_npu(): + import custom_ops # noqa: F401 + import torch_npu + from sglang.srt.hardware_backend.npu.utils import get_indexer_weight_stream from sglang.srt.layers import deep_gemm_wrapper from sglang.srt.layers.attention.nsa.utils import ( @@ -980,19 +986,47 @@ class Indexer(CustomOp): sin = sin.repeat(1, 2).view(-1, 1, 1, self.rope_head_dim) bs = x.shape[0] - q = self.wq_b(q_lora)[0] # [bs, 1536] @ [1536, 64 * 128] = [bs, 64 * 128] - q = q.view(bs, self.n_heads, self.head_dim) # [bs, 64, 128] - q_pe, q_nope = torch.split( - q, - [self.rope_head_dim, self.head_dim - self.rope_head_dim], - dim=-1, - ) # [bs, 64, 64 + 64] + if self.alt_stream is not None: + self.alt_stream.wait_stream(torch.npu.current_stream()) + with torch.npu.stream(self.alt_stream): + q = self.wq_b(q_lora)[ + 0 + ] # [bs, 1536] @ [1536, 64 * 128] = [bs, 64 * 128] + wq_b_event = self.alt_stream.record_event() + q = q.view(bs, self.n_heads, self.head_dim) # [bs, 64, 128] + q_pe, q_nope = torch.split( + q, + [self.rope_head_dim, self.head_dim - self.rope_head_dim], + dim=-1, + ) # [bs, 64, 64 + 64] + q_pe = q_pe.view(bs, self.n_heads, 1, self.rope_head_dim) + q_pe = torch_npu.npu_rotary_mul(q_pe, cos, sin).view( + bs, self.n_heads, self.rope_head_dim + ) # [bs, n, d] + q = torch.cat([q_pe, q_nope], dim=-1) + q.record_stream(self.alt_stream) + q_rope_event = self.alt_stream.record_event() + else: + q = self.wq_b(q_lora)[0] # [bs, 1536] @ [1536, 64 * 128] = [bs, 64 * 128] + q = q.view(bs, self.n_heads, self.head_dim) # [bs, 64, 128] + q_pe, q_nope = torch.split( + q, + [self.rope_head_dim, self.head_dim - self.rope_head_dim], + dim=-1, + ) # [bs, 64, 64 + 64] + q_pe = q_pe.view(bs, self.n_heads, 1, self.rope_head_dim) + q_pe = torch_npu.npu_rotary_mul(q_pe, cos, sin).view( + bs, self.n_heads, self.rope_head_dim + ) # [bs, n, d] + q = torch.cat([q_pe, q_nope], dim=-1) - q_pe = q_pe.view(bs, self.n_heads, 1, self.rope_head_dim) - q_pe = torch.ops.npu.npu_rotary_mul(q_pe, cos, sin).view( - bs, self.n_heads, self.rope_head_dim - ) # [bs, n, d] - q = torch.cat([q_pe, q_nope], dim=-1) + indexer_weight_stream = get_indexer_weight_stream() + indexer_weight_stream.wait_stream(torch.npu.current_stream()) + with torch.npu.stream(indexer_weight_stream): + x = x.view(-1, self.hidden_size) + weights = self.weights_proj(x.float())[0].to(torch.bfloat16) + weights.record_stream(indexer_weight_stream) + weights_event = indexer_weight_stream.record_event() k_proj = self.wk(x)[0] # [b, s, 7168] @ [7168, 128] = [b, s, 128] k = self.k_norm(k_proj) @@ -1072,8 +1106,10 @@ class Indexer(CustomOp): past_key_states = forward_batch.token_to_kv_pool.get_index_k_buffer(layer_id) - x = x.view(-1, self.hidden_size) - weights = self.weights_proj(x.float())[0].to(torch.bfloat16) + if self.alt_stream is not None: + torch.npu.current_stream().wait_event(q_rope_event) + torch.npu.current_stream().wait_event(weights_event) + block_table = forward_batch.attn_backend.forward_metadata.block_tables if ( is_prefill diff --git a/python/sglang/srt/layers/layernorm.py b/python/sglang/srt/layers/layernorm.py index bf9098fc0..e543a63ee 100644 --- a/python/sglang/srt/layers/layernorm.py +++ b/python/sglang/srt/layers/layernorm.py @@ -343,19 +343,7 @@ class LayerNorm(CustomOp): self, x: torch.Tensor, ) -> torch.Tensor: - orig_dtype = x.dtype - x = x.to(self.dtype) - - mean = x.mean(dim=-1, keepdim=True) - variance = (x - mean).pow(2).mean(dim=-1, keepdim=True) - x = (x - mean) * torch.rsqrt(variance + self.variance_epsilon) - - if self.elementwise_affine: - x = x * self.weight.to(self.dtype) - if self.use_bias: - x = x + self.bias.to(self.dtype) - - return x.to(orig_dtype) + return self.forward_native(x) def forward_cpu( self, diff --git a/python/sglang/srt/layers/linear.py b/python/sglang/srt/layers/linear.py index b63680d02..e64fd5070 100644 --- a/python/sglang/srt/layers/linear.py +++ b/python/sglang/srt/layers/linear.py @@ -245,6 +245,11 @@ class ReplicatedLinear(LinearBase): else: raise ValueError(f"{loaded_weight} are not all equal") + if param.dtype == torch.int8 or loaded_weight.dtype == torch.int8: + assert ( + param.dtype == loaded_weight.dtype + ), "init para dtype and loaded weight dtype should be the same" + assert param.size() == loaded_weight.size() param.data.copy_(loaded_weight) diff --git a/python/sglang/srt/layers/quantization/compressed_tensors/utils.py b/python/sglang/srt/layers/quantization/compressed_tensors/utils.py index ddefa5ea3..ab01c94d5 100644 --- a/python/sglang/srt/layers/quantization/compressed_tensors/utils.py +++ b/python/sglang/srt/layers/quantization/compressed_tensors/utils.py @@ -79,7 +79,7 @@ def check_equal_or_regex_match(layer_name: str, targets: Iterable[str]) -> bool: if target starts with 're:' to any target in list. """ for target in targets: - if _is_equal_or_regex_match(layer_name, target): + if _is_equal_or_regex_match(layer_name, target, check_contains=True): return True return False diff --git a/python/sglang/srt/models/deepseek_nextn.py b/python/sglang/srt/models/deepseek_nextn.py index e1b341cef..2aa062e5d 100644 --- a/python/sglang/srt/models/deepseek_nextn.py +++ b/python/sglang/srt/models/deepseek_nextn.py @@ -97,7 +97,7 @@ class DeepseekModelNextN(nn.Module): self.eh_proj = nn.Linear(2 * config.hidden_size, config.hidden_size, bias=False) - self.alt_stream = torch.cuda.Stream() if _is_cuda else None + self.alt_stream = torch.cuda.Stream() if _is_cuda or _is_npu else None layer_name = "decoder" if _is_npu and ( diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index e210a83af..342462517 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -700,6 +700,7 @@ class DeepseekV2MoE(nn.Module): dict(tp_rank=0, tp_size=1) if get_moe_a2a_backend().is_deepep() or get_moe_a2a_backend().is_mooncake() + or get_moe_a2a_backend().is_ascend_fuseep() or should_use_flashinfer_cutlass_moe_fp4_allgather() else {} ), @@ -738,7 +739,11 @@ class DeepseekV2MoE(nn.Module): self.top_k = config.num_experts_per_tok - if get_moe_a2a_backend().is_deepep() or get_moe_a2a_backend().is_mooncake(): + if ( + get_moe_a2a_backend().is_deepep() + or get_moe_a2a_backend().is_mooncake() + or get_moe_a2a_backend().is_ascend_fuseep() + ): # TODO: we will support tp < ep in the future self.ep_size = get_moe_expert_parallel_world_size() self.num_experts = ( @@ -755,7 +760,9 @@ class DeepseekV2MoE(nn.Module): ) self._enable_a2a_moe = ( - get_moe_a2a_backend().is_deepep() or get_moe_a2a_backend().is_mooncake() + get_moe_a2a_backend().is_deepep() + or get_moe_a2a_backend().is_mooncake() + or get_moe_a2a_backend().is_ascend_fuseep() ) self._fuse_shared_experts_inside_sbo = SboFlags.fuse_shared_experts_inside_sbo() @@ -986,7 +993,14 @@ class DeepseekV2MoE(nn.Module): # router_logits: (num_tokens, n_experts) router_logits = self.gate(hidden_states, forward_batch=forward_batch) if not sbo_enabled_flag: - shared_output = self._forward_shared_experts(hidden_states) + if self.alt_stream is not None: + self.alt_stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(self.alt_stream): + shared_output = self._forward_shared_experts(hidden_states) + shared_output.record_stream(self.alt_stream) + shared_event = self.alt_stream.record_event() + else: + shared_output = self._forward_shared_experts(hidden_states) topk_output = self.topk( hidden_states, router_logits, @@ -1105,6 +1119,12 @@ class DeepseekV2MoE(nn.Module): topk_output=topk_output, ) + if ( + hidden_states.shape[0] > 0 + and not sbo_enabled_flag + and self.alt_stream is not None + ): + torch.cuda.current_stream().wait_event(shared_event) if shared_output is not None: x = shared_output if self.experts.should_fuse_routed_scaling_factor_in_topk: @@ -2991,7 +3011,8 @@ class DeepseekV2Model(nn.Module): else: self.embed_tokens = PPMissingLayer() - self.alt_stream = torch.cuda.Stream() if _is_cuda else None + self.alt_stream = torch.cuda.Stream() if _is_cuda or _is_npu else None + self.layers, self.start_layer, self.end_layer = make_layers( config.num_hidden_layers, lambda idx, prefix: DeepseekV2DecoderLayer(