From 176dbcb85d3b7737564e4911947035cc0af65f0a Mon Sep 17 00:00:00 2001 From: zhaozx-cn <59479021+zhaozx-cn@users.noreply.github.com> Date: Mon, 21 Sep 2026 09:51:19 +0800 Subject: [PATCH] [npu]add chunk gdn kernel and unify ssm state layout for ascend gdn backend (#36187) --- .../npu/attention/ascend_gdn_backend.py | 70 +++++++++++++++---- .../ascend_hybrid_linear_attn_backend.py | 7 +- .../HiCache/test_npu_hicache_mha.py | 2 - .../HiCache/test_npu_hicache_mla.py | 2 - 4 files changed, 59 insertions(+), 22 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 c8928c06e..ac1989e7e 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,6 +5,8 @@ from sgl_kernel_npu.fla.fused_gdn_gating import ( fused_gdn_gating_kernel_without_sigmoid, fused_gdn_gating_npu, ) +from sgl_kernel_npu.fla.l2norm import l2norm_fwd +from sgl_kernel_npu.fla.utils import prepare_chunk_indices from sglang.srt.hardware_backend.npu.attention.ascend_hybrid_linear_attn_backend import ( AscendMambaAttnBackendBase, @@ -69,6 +71,10 @@ class AscendGDNAttnBackend(AscendMambaAttnBackendBase): ) else: self.ssm_state_indices = cache_indices + if forward_mode.is_prefill(): + self.actual_seq_lengths = torch.diff( + self.forward_metadata.query_start_loc + ).to(torch.int32) def init_forward_metadata_out_graph( self, @@ -297,21 +303,21 @@ class AscendGDNAttnBackend(AscendMambaAttnBackendBase): ) actual_seq_len = query.shape[0] - query = query.view(1, actual_seq_len, layer.num_q_heads, layer.head_q_dim) - key = key.view(1, actual_seq_len, layer.num_k_heads, layer.head_k_dim) - value = value.view(1, actual_seq_len, layer.num_v_heads, layer.head_v_dim) - + query = query.view( + actual_seq_len, layer.num_q_heads, layer.head_q_dim + ).contiguous() + key = key.view( + actual_seq_len, layer.num_k_heads, layer.head_k_dim + ).contiguous() + value = value.view( + actual_seq_len, layer.num_v_heads, layer.head_v_dim + ).contiguous() g, beta = fused_gdn_gating(layer.A_log, a, b, layer.dt_bias) - core_attn_out, last_recurrent_state, h = self.kernel_dispatcher.extend( - q=query, - k=key, - v=value, - g=g, - beta=beta, - ssm_states=ssm_states, - cache_indices=cache_indices, - query_start_loc=query_start_loc, + + core_attn_out, last_recurrent_state, h = self.chunk_gdn( + query, key, value, g, beta, ssm_states[cache_indices], query_start_loc ) + if last_recurrent_state is not None: last_recurrent_state = last_recurrent_state.to( ssm_states.dtype, copy=False @@ -324,6 +330,44 @@ class AscendGDNAttnBackend(AscendMambaAttnBackendBase): return core_attn_out + def chunk_gdn( + self, + q, + k, + v, + g, + beta, + initial_state, + cu_seqlens, + ): + q = l2norm_fwd(q) # [T, Nk, Dk] + k = l2norm_fwd(k) # [T, Nk, Dk] + g = g.to(torch.float32) # [T, Nv] + scale = k.shape[-1] ** -0.5 + + chunk_size = 64 + chunk_indices = ( + prepare_chunk_indices(cu_seqlens, chunk_size) + if cu_seqlens is not None + else None + ) + h = k.new_empty( + len(chunk_indices), v.shape[-2], v.shape[-1], k.shape[-1] + ) # [NT, NH, DV, Dk] + + core_attn_out, last_recurrent_state = torch.ops.npu.chunk_gated_delta_rule( + q, + k, + v, + beta=beta, + initial_state=initial_state, + actual_seq_lengths=self.actual_seq_lengths, + scale=scale, + chunk_state=h, + g=g, + ) + return core_attn_out, last_recurrent_state, h + def fused_recurrent_gated_delta_rule_update( self, mix_qkv: torch.Tensor, diff --git a/python/sglang/srt/hardware_backend/npu/attention/ascend_hybrid_linear_attn_backend.py b/python/sglang/srt/hardware_backend/npu/attention/ascend_hybrid_linear_attn_backend.py index 3e60cd9ed..749e2b54e 100644 --- a/python/sglang/srt/hardware_backend/npu/attention/ascend_hybrid_linear_attn_backend.py +++ b/python/sglang/srt/hardware_backend/npu/attention/ascend_hybrid_linear_attn_backend.py @@ -290,11 +290,8 @@ class AscendHybridLinearAttnBackend(HybridLinearAttnBackend): track_mask = mamba_steps_to_track >= 0 # Track conv state from the verify-time window before rolling back # the working slot; NPU does not keep per-step conv intermediates. - track_indices = mamba_track_indices[track_mask] - if track_indices.numel() > 0: - conv_states[:, track_indices] = conv_states[ - :, dst_indices_tensor[track_mask] - ] + src_slots = torch.where(track_mask, dst_indices_tensor, mamba_track_indices) + conv_states[:, mamba_track_indices] = conv_states[:, src_slots] if dst_indices_tensor.numel() > 0: conv_state_rollback( diff --git a/test/registered/npu/basic_function/HiCache/test_npu_hicache_mha.py b/test/registered/npu/basic_function/HiCache/test_npu_hicache_mha.py index a7e5f05e9..18b878793 100644 --- a/test/registered/npu/basic_function/HiCache/test_npu_hicache_mha.py +++ b/test/registered/npu/basic_function/HiCache/test_npu_hicache_mha.py @@ -38,8 +38,6 @@ class TestAscendMhaHicache(CustomTestCase): "--attention-backend", "ascend", "--enable-hierarchical-cache", - "--hicache-ratio", - 1.2, ] def test_a_gsm8k(self): diff --git a/test/registered/npu/basic_function/HiCache/test_npu_hicache_mla.py b/test/registered/npu/basic_function/HiCache/test_npu_hicache_mla.py index 659c9dab8..e7d8675b0 100644 --- a/test/registered/npu/basic_function/HiCache/test_npu_hicache_mla.py +++ b/test/registered/npu/basic_function/HiCache/test_npu_hicache_mla.py @@ -40,8 +40,6 @@ class TestAscendMlaHicache(CustomTestCase): "--tp-size", 4, "--enable-hierarchical-cache", - "--hicache-ratio", - 1.2, ] def test_a_gsm8k(self):