Make GDN support non-continuous B/A Tensor input to fix the accuracy regression of Qwen3.5-27B (#22312)

Signed-off-by: cs-cat <118669451+cs-cat@users.noreply.github.com>

python/sglang/srt/layers/attention/fla/fused_gdn_gating.py

@@ -16,6 +16,8 @@ def fused_gdn_gating_kernel(
     b,
     dt_bias,
     seq_len,
+    stride_a,
+    stride_b,
     NUM_HEADS: tl.constexpr,
     beta: tl.constexpr,
     threshold: tl.constexpr,
@@ -26,8 +28,8 @@ def fused_gdn_gating_kernel(
     off = i_b * seq_len * NUM_HEADS + i_s * NUM_HEADS + head_off
     mask = head_off < NUM_HEADS
     blk_A_log = tl.load(A_log + head_off, mask=mask)
-    blk_a = tl.load(a + off, mask=mask)
-    blk_b = tl.load(b + off, mask=mask)
+    blk_a = tl.load(a + i_b * stride_a + head_off, mask=mask)
+    blk_b = tl.load(b + i_b * stride_b + head_off, mask=mask)
     blk_bias = tl.load(dt_bias + head_off, mask=mask)
     x = blk_a.to(tl.float32) + blk_bias.to(tl.float32)
     softplus_x = tl.where(
@@ -49,6 +51,8 @@ def fused_gdn_gating(
 ) -> Tuple[torch.Tensor, torch.Tensor]:
     batch, num_heads = a.shape
     seq_len = 1
+    stride_a = a.stride(0)
+    stride_b = b.stride(0)
     grid = (batch, seq_len, triton.cdiv(num_heads, 8))
     g = torch.empty(1, batch, num_heads, dtype=torch.float32, device=a.device)
     beta_output = torch.empty(1, batch, num_heads, dtype=torch.float32, device=b.device)
@@ -60,6 +64,8 @@ def fused_gdn_gating(
         b,
         dt_bias,
         seq_len,
+        stride_a,
+        stride_b,
         num_heads,
         beta,
         threshold,

python/sglang/srt/layers/attention/fla/fused_sigmoid_gating_recurrent.py

@@ -30,6 +30,7 @@ def fused_sigmoid_gating_delta_rule_update_kernel(
     # ================================================
     scale,
     T,
+    stride_a,
     stride_q,
     stride_k,
     stride_v,
@@ -81,10 +82,10 @@ def fused_sigmoid_gating_delta_rule_update_kernel(
     # Gating computation pointers
     p_A_log = A_log + i_hv
     if IS_KDA:
-        p_a = a + (bos * HV + i_hv) * K + o_k
+        p_a = a + bos * stride_a + i_hv * K + o_k
         p_dt_bias = dt_bias + i_hv * K + o_k
     else:
-        p_a = a + bos * HV + i_hv
+        p_a = a + bos * stride_a + i_hv
         p_dt_bias = dt_bias + i_hv
 
     mask_k = o_k < K
@@ -220,10 +221,7 @@ def fused_sigmoid_gating_delta_rule_update_kernel(
         p_v += stride_v
         p_b += stride_b
         p_o += HV * V
-        if IS_KDA:
-            p_a += HV * K
-        else:
-            p_a += HV
+        p_a += stride_a
 
     # Store final state back to h0_source with bounds checking
     if not DISABLE_STATE_UPDATE:
@@ -278,6 +276,10 @@ def fused_sigmoid_gating_delta_rule_update(
     stride_k = k.stride()[1]
     stride_v = v.stride()[1]
     stride_b = b.stride()[-2]
+    # Both paths (KDA/GDN) advance p_a once per token, so use the token-axis stride.
+    # For 2D a ([T, ...]) this is stride(0); for 3D a ([B, T, ...]) this is stride(1).
+    # Using stride()[-2] covers GDN [T, HV] and KDA layouts ([T, HV*K] / [B, T, HV*K]).
+    stride_a = a.stride()[-2]
     HV = v.shape[2]
     N = B if cu_seqlens is None else len(cu_seqlens) - 1
     BK, BV = triton.next_power_of_2(K), min(triton.next_power_of_2(V), 32)
@@ -327,6 +329,7 @@ def fused_sigmoid_gating_delta_rule_update(
         stride_retrieve_parent_token_token=stride_retrieve_parent_token_token,
         scale=scale,
         T=T,
+        stride_a=stride_a,
         stride_q=stride_q,
         stride_k=stride_k,
         stride_v=stride_v,