diff --git a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py index 7518c909a..181b29c8b 100644 --- a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py +++ b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py @@ -1003,7 +1003,7 @@ class Indexer(CustomOp): ) # [bs, 64, 64 + 64] q_pe = q_pe.view(bs, self.n_heads, 1, self.rope_head_dim) - q_pe = torch_npu.npu_interleave_rope(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) @@ -1017,7 +1017,7 @@ class Indexer(CustomOp): ) # [bs, 64 + 64] k_pe = k_pe.view(-1, 1, 1, self.rope_head_dim) - k_pe = torch_npu.npu_interleave_rope(k_pe, cos, sin).view( + k_pe = torch_npu.npu_rotary_mul(k_pe, cos, sin).view( bs, 1, self.rope_head_dim ) # [bs, 1, d] k = torch.cat([k_pe, k_nope.unsqueeze(1)], dim=-1) # [bs, 1, 128] @@ -1087,7 +1087,7 @@ 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] + weights = self.weights_proj(x.float())[0].to(torch.bfloat16) block_table = ( block_table[: actual_seq_lengths_q.size()[0]] if is_prefill else block_table ) diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index f849b7f48..31348b5be 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -3653,9 +3653,10 @@ class DeepseekV2ForCausalLM(nn.Module): self_attn.w_kc = bind_or_assign( self_attn.w_kc, w_kc.transpose(1, 2).contiguous().transpose(1, 2) ) - self_attn.w_vc = bind_or_assign( - self_attn.w_vc, w_vc.contiguous().transpose(1, 2) - ) + w_vc = w_vc.contiguous().transpose(1, 2) + if _is_npu: + w_vc = w_vc.contiguous() + self_attn.w_vc = bind_or_assign(self_attn.w_vc, w_vc) if ( hasattr(self_attn.kv_b_proj, "weight_scale") and self_attn.w_scale is None diff --git a/test/srt/ascend/test_ascend_deepep.py b/test/srt/ascend/test_ascend_deepep.py index d6a94890a..f8cf760d5 100644 --- a/test/srt/ascend/test_ascend_deepep.py +++ b/test/srt/ascend/test_ascend_deepep.py @@ -47,6 +47,8 @@ class TestAscendDeepEP(CustomTestCase): 1, "--ep-size", 16, + "--max-running-requests", + 24, "--moe-a2a-backend", "deepep", "--deepep-mode",