[NPU] optimization for dsv3.2 (#14572)

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_query=actual_seq_qlen.to(
-                actual_seq_lengths_kv=actual_seq_lengths_kv.to(q.device),
+                    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",

python/sglang/srt/hardware_backend/npu/modules/deepseek_v2_attention_mla_npu.py

@@ -273,6 +273,10 @@ def forward_dsa_prepare_npu(
                 m.qk_rope_head_dim,
                 m.quant_config,
             )
+        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,
@@ -284,28 +288,36 @@ def forward_dsa_prepare_npu(
             ) = 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,

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(

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

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

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,6 +986,27 @@ class Indexer(CustomOp):
         sin = sin.repeat(1, 2).view(-1, 1, 1, self.rope_head_dim)
 
         bs = x.shape[0]
+        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(
@@ -987,13 +1014,20 @@ class Indexer(CustomOp):
                 [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.ops.npu.npu_rotary_mul(q_pe, cos, sin).view(
+            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)
 
+        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)
         k_pe, k_nope = torch.split(
@@ -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)
+        if self.alt_stream is not None:
-        weights = self.weights_proj(x.float())[0].to(torch.bfloat16)
+            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

python/sglang/srt/layers/layernorm.py

@@ -343,19 +343,7 @@ class LayerNorm(CustomOp):
         self,
         x: torch.Tensor,
     ) -> torch.Tensor:
-        orig_dtype = x.dtype
+        return self.forward_native(x)
-        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)
 
     def forward_cpu(
         self,

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)
 

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
 

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 (

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,6 +993,13 @@ 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:
+                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,
@@ -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(