diff --git a/python/sglang/srt/configs/model_config.py b/python/sglang/srt/configs/model_config.py index 9bfd16210..6bdb83747 100644 --- a/python/sglang/srt/configs/model_config.py +++ b/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 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 dc811fcca..740d11df7 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 @@ -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) diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index b6986977a..365fd6690 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/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, diff --git a/test/registered/ascend/llm_models/test_ascend_deepseek_v3_2_exp_w8a8.py b/test/registered/ascend/llm_models/test_ascend_deepseek_v3_2_exp_w8a8.py new file mode 100644 index 000000000..46ad904e7 --- /dev/null +++ b/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() diff --git a/test/registered/ascend/vlm_models/test_ascend_deepseek_vl2.py b/test/registered/ascend/vlm_models/test_ascend_deepseek_vl2.py new file mode 100644 index 000000000..9becccff5 --- /dev/null +++ b/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()