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

[NPU] use causal_conv1d_update_v2 for performance (#24595)

Browse Source 
Co-authored-by: iridiumine <iridiumine@users.noreply.github.com>
This commit is contained in:
iridiumine
2026-05-12 17:04:02 +08:00
committed by GitHub
co-authored by iridiumine
parent 1efe9e2164
commit 693f49793a
1 changed files with 14 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/hardware_backend/npu/attention/ascend_gdn_backend.py
+14 -14
View File
@@ -8,6 +8,7 @@ from sgl_kernel_npu.fla.fused_gdn_gating import (
from sgl_kernel_npu.mamba.causal_conv1d import ( from sgl_kernel_npu.mamba.causal_conv1d import (
causal_conv1d_fn_npu, causal_conv1d_fn_npu,
causal_conv1d_update_npu, causal_conv1d_update_npu,
causal_conv1d_update_v2,
) )
from sglang.srt.hardware_backend.npu.attention.ascend_hybrid_linear_attn_backend import ( from sglang.srt.hardware_backend.npu.attention.ascend_hybrid_linear_attn_backend import (
@@ -224,9 +225,7 @@ class AscendGDNAttnBackend(AscendMambaAttnBackendBase):
else: else:
has_initial_states = forward_batch.extend_prefix_lens > 0 has_initial_states = forward_batch.extend_prefix_lens > 0
if is_target_verify: if is_target_verify:
draft_token_num = forward_batch.spec_info.draft_token_num
num_token_padding = mixed_qkv.shape[0] num_token_padding = mixed_qkv.shape[0]
batch_size = cache_indices.shape[0]
if ( if (
not self.graph_mode not self.graph_mode
and forward_batch.num_token_non_padded_cpu != num_token_padding and forward_batch.num_token_non_padded_cpu != num_token_padding
@@ -236,23 +235,24 @@ class AscendGDNAttnBackend(AscendMambaAttnBackendBase):
b = b[: forward_batch.num_token_non_padded_cpu] b = b[: forward_batch.num_token_non_padded_cpu]
seq_len = forward_batch.num_token_non_padded_cpu seq_len = forward_batch.num_token_non_padded_cpu
mixed_qkv_reshaped = mixed_qkv.view(batch_size, draft_token_num, -1) batch_size = cache_indices.shape[0]
num_accept_tokens = torch.full( draft_token_num = forward_batch.spec_info.draft_token_num
num_accepted_tokens = torch.full(
(batch_size,), (batch_size,),
draft_token_num, draft_token_num,
dtype=torch.int32, dtype=torch.int32,
device=mixed_qkv.device, device=mixed_qkv.device,
) )
mixed_qkv = torch.ops.npu.causal_conv1d_update( mixed_qkv = causal_conv1d_update_v2(
mixed_qkv_reshaped, x=mixed_qkv.view(batch_size, draft_token_num, -1).contiguous(),
layer.conv_weights.transpose(0, 1).contiguous(), conv_state=conv_states.contiguous(),
conv_states, weight=layer.conv_weights.transpose(0, 1).contiguous(),
cache_indices, bias=layer.bias,
layer.bias, activation=layer.activation,
num_accept_tokens, conv_state_indices=cache_indices,
None, num_accepted_tokens=num_accepted_tokens,
layer.activation == "silu", pad_slot_id=-1,
self.pad_slot_id, validate_data=False,
).view(seq_len, -1) ).view(seq_len, -1)
else: else:
mixed_qkv = mixed_qkv.transpose(0, 1) mixed_qkv = mixed_qkv.transpose(0, 1)