minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[NPU] bug fix: w_vc need contiguous for NPU batch_matmul_transpose ops (#13980)

Browse Source
This commit is contained in:
ZhengdQin
2025-12-03 19:35:18 +08:00
committed by GitHub
parent 96cc10834a
commit d122e32467
3 changed files with 9 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/attention/nsa/nsa_indexer.py
+3 -3
View File
@@ -1003,7 +1003,7 @@ class Indexer(CustomOp):
) # [bs, 64, 64 + 64] ) # [bs, 64, 64 + 64]
q_pe = q_pe.view(bs, self.n_heads, 1, self.rope_head_dim) 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, self.n_heads, self.rope_head_dim
) # [bs, n, d] ) # [bs, n, d]
q = torch.cat([q_pe, q_nope], dim=-1) q = torch.cat([q_pe, q_nope], dim=-1)
@@ -1017,7 +1017,7 @@ class Indexer(CustomOp):
) # [bs, 64 + 64] ) # [bs, 64 + 64]
k_pe = k_pe.view(-1, 1, 1, self.rope_head_dim) 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, self.rope_head_dim
) # [bs, 1, d] ) # [bs, 1, d]
k = torch.cat([k_pe, k_nope.unsqueeze(1)], dim=-1) # [bs, 1, 128] 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) past_key_states = forward_batch.token_to_kv_pool.get_index_k_buffer(layer_id)
x = x.view(-1, self.hidden_size) 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 = (
block_table[: actual_seq_lengths_q.size()[0]] if is_prefill else block_table block_table[: actual_seq_lengths_q.size()[0]] if is_prefill else block_table
) )
python/sglang/srt/models/deepseek_v2.py
+4 -3
View File
@@ -3653,9 +3653,10 @@ class DeepseekV2ForCausalLM(nn.Module):
self_attn.w_kc = bind_or_assign( self_attn.w_kc = bind_or_assign(
self_attn.w_kc, w_kc.transpose(1, 2).contiguous().transpose(1, 2) self_attn.w_kc, w_kc.transpose(1, 2).contiguous().transpose(1, 2)
) )
self_attn.w_vc = bind_or_assign( w_vc = w_vc.contiguous().transpose(1, 2)
self_attn.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 ( if (
hasattr(self_attn.kv_b_proj, "weight_scale") hasattr(self_attn.kv_b_proj, "weight_scale")
and self_attn.w_scale is None and self_attn.w_scale is None
test/srt/ascend/test_ascend_deepep.py
+2
View File
@@ -47,6 +47,8 @@ class TestAscendDeepEP(CustomTestCase):
1, 1,
"--ep-size", "--ep-size",
16, 16,
"--max-running-requests",
24,
"--moe-a2a-backend", "--moe-a2a-backend",
"deepep", "deepep",
"--deepep-mode", "--deepep-mode",