From 4e7209caa81694f679855720cd9433632a1ee26a Mon Sep 17 00:00:00 2001 From: zhaozx-cn <59479021+zhaozx-cn@users.noreply.github.com> Date: Wed, 5 Aug 2026 22:22:49 +0800 Subject: [PATCH] [NPU] Add causal conv1d (#28267) --- .../npu/attention/ascend_gdn_backend.py | 100 +++++++++--------- 1 file changed, 49 insertions(+), 51 deletions(-) diff --git a/python/sglang/srt/hardware_backend/npu/attention/ascend_gdn_backend.py b/python/sglang/srt/hardware_backend/npu/attention/ascend_gdn_backend.py index 7e63f3052..cc9510ff9 100644 --- a/python/sglang/srt/hardware_backend/npu/attention/ascend_gdn_backend.py +++ b/python/sglang/srt/hardware_backend/npu/attention/ascend_gdn_backend.py @@ -5,11 +5,6 @@ from sgl_kernel_npu.fla.fused_gdn_gating import ( fused_gdn_gating_kernel_without_sigmoid, fused_gdn_gating_npu, ) -from sgl_kernel_npu.mamba.causal_conv1d import ( - causal_conv1d_fn_npu, - causal_conv1d_update_npu, - causal_conv1d_update_v2, -) from sglang.srt.hardware_backend.npu.attention.ascend_hybrid_linear_attn_backend import ( AscendMambaAttnBackendBase, @@ -26,8 +21,6 @@ from sglang.srt.model_executor.model_runner import ModelRunner from sglang.srt.speculative.eagle_info import EagleDraftInput, EagleVerifyInput fused_gdn_gating = fused_gdn_gating_npu -causal_conv1d_fn = causal_conv1d_fn_npu -causal_conv1d_update = causal_conv1d_update_npu class AscendGDNAttnBackend(AscendMambaAttnBackendBase): @@ -109,6 +102,13 @@ class AscendGDNAttnBackend(AscendMambaAttnBackendBase): self._prepare_mamba_track_metadata(forward_batch) self.graph_mode = False + def _get_conv_weights_t(self, layer: RadixLinearAttention) -> torch.Tensor: + w = getattr(layer, "_conv_weights_t", None) + if w is None: + w = layer.conv_weights.transpose(0, 1).contiguous() + layer._conv_weights_t = w + return w + def forward_decode( self, layer: RadixLinearAttention, @@ -125,16 +125,17 @@ class AscendGDNAttnBackend(AscendMambaAttnBackendBase): cache_indices = self.forward_metadata.mamba_cache_indices assert isinstance(mixed_qkv, torch.Tensor) - conv_states_tmp = conv_states.transpose(1, 2).clone() - mixed_qkv = causal_conv1d_update( + mixed_qkv = torch.ops.npu.causal_conv1d( mixed_qkv, - conv_states_tmp, - layer.conv_weights, - layer.bias, - layer.activation, - conv_state_indices=cache_indices, + self._get_conv_weights_t(layer), + conv_states=conv_states, + bias=layer.bias, + query_start_loc=query_start_loc, + cache_indices=cache_indices, + activation_mode=1, + pad_slot_id=-1, + run_mode=1, ) - conv_states[:] = conv_states_tmp.transpose(1, 2) query, key, value = torch.split( mixed_qkv, @@ -219,44 +220,41 @@ class AscendGDNAttnBackend(AscendMambaAttnBackendBase): dtype=torch.int32, device=mixed_qkv.device, ) - mixed_qkv = causal_conv1d_update_v2( - x=mixed_qkv.view(batch_size, draft_token_num, -1).contiguous(), - conv_state=conv_states.contiguous(), - weight=layer.conv_weights.transpose(0, 1).contiguous(), - bias=layer.bias, - activation=layer.activation, - conv_state_indices=cache_indices, - num_accepted_tokens=num_accepted_tokens, - pad_slot_id=-1, - validate_data=False, - ).view(seq_len, -1) - else: - mixed_qkv = mixed_qkv.transpose(0, 1) - if forward_metadata.has_mamba_track_mask: - mixed_qkv_to_track = mixed_qkv[ - :, forward_metadata.track_conv_indices - ].transpose(0, 1) - conv_states.transpose(1, 2)[ - forward_metadata.conv_states_mask_indices - ] = mixed_qkv_to_track - kernel_size = layer.conv_weights.shape[-1] - conv_states_for_prefill = conv_states[:, -(kernel_size - 1) :, :] - conv_states_tmp = conv_states_for_prefill.transpose(1, 2).contiguous() - - mixed_qkv = causal_conv1d_fn( + mixed_qkv = torch.ops.npu.causal_conv1d( mixed_qkv, - layer.conv_weights, - layer.bias, - activation=layer.activation, - conv_states=conv_states_tmp, - has_initial_state=has_initial_states, - cache_indices=cache_indices, + self._get_conv_weights_t(layer), + conv_states=conv_states, + bias=layer.bias, query_start_loc=query_start_loc, - seq_lens_cpu=forward_batch.extend_seq_lens_cpu, - ).transpose(0, 1)[:seq_len] - conv_states[:, -(kernel_size - 1) :, :] = conv_states_tmp.transpose( - 1, 2 - ).contiguous() + cache_indices=cache_indices, + num_accepted_tokens=num_accepted_tokens, + activation_mode=1, + pad_slot_id=-1, + run_mode=1, + ) + else: + if forward_metadata.has_mamba_track_mask: + mixed_qkv_to_track = mixed_qkv[forward_metadata.track_conv_indices] + conv_states[forward_metadata.conv_states_mask_indices] = ( + mixed_qkv_to_track + ) + kernel_size = layer.conv_weights.shape[-1] + conv_states_for_prefill = conv_states[ + :, -(kernel_size - 1) :, : + ].contiguous() + mixed_qkv = torch.ops.npu.causal_conv1d( + mixed_qkv, + self._get_conv_weights_t(layer), + conv_states=conv_states_for_prefill, + bias=layer.bias, + query_start_loc=query_start_loc, + cache_indices=cache_indices, + has_initial_state=has_initial_states, + activation_mode=1, + pad_slot_id=-1, + run_mode=0, + ) + conv_states[:, -(kernel_size - 1) :, :] = conv_states_for_prefill if is_target_verify: g, beta = fused_gdn_gating_kernel_without_sigmoid( layer.A_log, a, b, layer.dt_bias