[npu]add chunk gdn kernel and unify ssm state layout for ascend gdn backend (#36187)

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,

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(