[NPU]bugfix: fix for dsv3.2 and dsvl2 (#17007)

Co-authored-by: Hexq0210 <893781835@qq.com>
Co-authored-by: liupeng374 <782420244@qq.com>
Co-authored-by: cy <chenyang08056032@163.com>

python/sglang/srt/configs/model_config.py

@@ -432,6 +432,7 @@ class ModelConfig:
             self.attention_arch = AttentionArch.MLA
             self.kv_lora_rank = self.hf_text_config.kv_lora_rank
             self.qk_rope_head_dim = self.hf_text_config.qk_rope_head_dim
+            self.qk_nope_head_dim = self.hf_text_config.qk_nope_head_dim
         elif "KimiVLForConditionalGeneration" in self.hf_config.architectures:
             self.head_dim = 256
             self.attention_arch = AttentionArch.MLA

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

@@ -64,36 +64,44 @@ def forward_mha_prepare_npu(
     kv_a, _ = latent_cache.split([m.kv_lora_rank, m.qk_rope_head_dim], dim=-1)
     latent_cache = latent_cache.unsqueeze(1)
 
-    B, S = q.shape[0], 1
-    cos, sin = m.rotary_emb.get_cos_sin_cache(
-        positions, hidden_states.dtype, offsets=None
-    )
-    q_pe = torch_npu.npu_interleave_rope(
-        q_pe.reshape(B, -1, S, m.qk_rope_head_dim),
-        cos,
-        sin,
-    )
-    q_pe = q_pe.reshape(B, -1, m.qk_rope_head_dim)
+    if m.use_deepseek_yarn_rope:
+        B, S = q.shape[0], 1
+        cos, sin = m.rotary_emb.get_cos_sin_cache(
+            positions, hidden_states.dtype, offsets=None
+        )
+        q_pe = torch_npu.npu_interleave_rope(
+            q_pe.reshape(B, -1, S, m.qk_rope_head_dim),
+            cos,
+            sin,
+        )
+        q_pe = q_pe.reshape(B, -1, m.qk_rope_head_dim)
 
-    ckv_cache, k_rope_cache = forward_batch.token_to_kv_pool.get_kv_buffer(m.layer_id)
-    _, _, k_pe, kv_a = torch_npu.npu_kv_rmsnorm_rope_cache(
-        latent_cache.view(-1, 1, 1, m.kv_lora_rank + m.qk_rope_head_dim),  # bnsd
-        m.kv_a_layernorm.weight,
-        cos,
-        sin,
-        forward_batch.out_cache_loc.to(torch.int64),
-        k_rope_cache,
-        ckv_cache,
-        k_rope_scale=None,
-        c_kv_scale=None,
-        k_rope_offset=None,
-        c_kv_offset=None,
-        epsilon=m.kv_a_layernorm.variance_epsilon,
-        cache_mode="PA_NZ" if is_fia_nz() else "PA_BNSD",
-        is_output_kv=True,
-    )  # adapter NZ
+        ckv_cache, k_rope_cache = forward_batch.token_to_kv_pool.get_kv_buffer(
+            m.layer_id
+        )
+        _, _, k_pe, kv_a = torch_npu.npu_kv_rmsnorm_rope_cache(
+            latent_cache.view(-1, 1, 1, m.kv_lora_rank + m.qk_rope_head_dim),  # bnsd
+            m.kv_a_layernorm.weight,
+            cos,
+            sin,
+            forward_batch.out_cache_loc.to(torch.int64),
+            k_rope_cache,
+            ckv_cache,
+            k_rope_scale=None,
+            c_kv_scale=None,
+            k_rope_offset=None,
+            c_kv_offset=None,
+            epsilon=m.kv_a_layernorm.variance_epsilon,
+            cache_mode="PA_NZ" if is_fia_nz() else "PA_BNSD",
+            is_output_kv=True,
+        )  # adapter NZ
 
-    k_pe = k_pe.reshape(B, -1, m.qk_rope_head_dim)
+        k_pe = k_pe.reshape(B, -1, m.qk_rope_head_dim)
+    else:
+        kv_a = m.kv_a_layernorm(kv_a)
+        k_pe = latent_cache[:, :, m.kv_lora_rank :]
+        if m.rotary_emb is not None:
+            q_pe, k_pe = m.rotary_emb(positions, q_pe, k_pe)
 
     q[..., m.qk_nope_head_dim :] = q_pe
 
@@ -288,10 +296,30 @@ 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)
+        if m.alt_stream is not None:
+            mla_event = torch.npu.Event()
+            mla_event.record()
+            with torch.npu.stream(m.alt_stream):
+                # alt stream waits for the completion of the event on the main stream to ensure data dependency is complete
+                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:
             (
                 q_pe,
                 k_pe,
@@ -303,12 +331,11 @@ 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)
+            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)
     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(
@@ -320,17 +347,24 @@ def forward_dsa_prepare_npu(
 
         q_lora = q.clone()  # required for topk_indices
 
-        m.alt_stream.wait_stream(torch.npu.current_stream())
-        with torch.npu.stream(m.alt_stream):
+        q_event = None
+        if m.alt_stream is not None:
+            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)
+                # record q to ensure memory space will not be released
+                q.record_stream(m.alt_stream)
+                q_event = m.alt_stream.record_event()
+        else:
             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)
+        k_nope = m.kv_a_layernorm(k_nope)
+        # main stream waits for the completion of the event on the alt stream to ensure data dependency is complete
+        if q_event is not None:
+            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)
 

python/sglang/srt/models/deepseek_v2.py

@@ -1220,6 +1220,7 @@ class DeepseekV2AttentionMLA(nn.Module, DeepseekMHAForwardMixin):
                 self.rotary_emb.forward = self.rotary_emb.forward_native
         else:
             self.rotary_emb = None
+        self.use_deepseek_yarn_rope = rope_scaling is not None
 
         self.attn_mqa = RadixAttention(
             self.num_local_heads,

test/registered/ascend/llm_models/test_ascend_deepseek_v3_2_exp_w8a8.py

@@ -0,0 +1,29 @@
+import unittest
+
+from sglang.test.ascend.gsm8k_ascend_mixin import GSM8KAscendMixin
+from sglang.test.ci.ci_register import register_npu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_npu_ci(est_time=400, suite="nightly-16-npu-a3", nightly=True)
+
+
+class TestDeepSeekV3_2ExpW8A8(GSM8KAscendMixin, CustomTestCase):
+    model = "/root/.cache/modelscope/hub/models/DeepSeek-V3.2-Exp-W8A8"
+    accuracy = 0.51
+    other_args = [
+        "--trust-remote-code",
+        "--mem-fraction-static",
+        "0.9",
+        "--attention-backend",
+        "ascend",
+        "--disable-cuda-graph",
+        "--tp-size",
+        "16",
+        "--quantization",
+        "modelslim",
+        "--disable-radix-cache",
+    ]
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/ascend/vlm_models/test_ascend_deepseek_vl2.py

@@ -0,0 +1,18 @@
+import unittest
+
+from sglang.test.ascend.vlm_utils import TestVLMModels
+from sglang.test.ci.ci_register import register_npu_ci
+
+register_npu_ci(est_time=400, suite="nightly-4-npu-a3", nightly=True)
+
+
+class TestGemmaModels(TestVLMModels):
+    model = "/root/.cache/modelscope/hub/models/deepseek-ai/deepseek-vl2"
+    mmmu_accuracy = 0.2
+
+    def test_vlm_mmmu_benchmark(self):
+        self._run_vlm_mmmu_test()
+
+
+if __name__ == "__main__":
+    unittest.main()