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_a = tl.load(a + i_b * stride_a + head_off, mask=mask)
-    blk_b = tl.load(b + 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 += stride_a
-            p_a += HV * K
-        else:
-            p_a += HV
 
     # 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,

test/registered/attention/test_gdn_noncontiguous_stride.py

@@ -0,0 +1,255 @@
+"""
+Tests that fused_gdn_gating and fused_sigmoid_gating_delta_rule_update
+produce correct results when a/b inputs are non-contiguous,
+as happens with Qwen3.5-27B (v_per_group=3) via mixed_ba.split().
+"""
+
+import unittest
+
+import torch
+
+from sglang.srt.layers.attention.fla.fused_gdn_gating import fused_gdn_gating
+from sglang.srt.layers.attention.fla.fused_sigmoid_gating_recurrent import (
+    fused_sigmoid_gating_delta_rule_update,
+)
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=30, suite="stage-b-test-1-gpu-large")
+
+
+def _make_noncontiguous_ab(batch, num_heads, dtype=torch.bfloat16, device="cuda"):
+    """
+    Simulate Qwen3.5 fallback: mixed_ba.split([nv_tp, nv_tp], dim=-1).
+    Returns (b, a) as split views with stride(0) = 2 * num_heads.
+    Also returns contiguous copies for reference comparison.
+    """
+    mixed_ba = torch.randn(batch, 2 * num_heads, dtype=dtype, device=device)
+    b, a = mixed_ba.split([num_heads, num_heads], dim=-1)
+
+    # For batch=1, PyTorch may still report contiguous even when split keeps
+    # a widened leading stride. Validate stride semantics unconditionally.
+    if batch > 1:
+        assert not a.is_contiguous(), "a should be non-contiguous from split"
+        assert not b.is_contiguous(), "b should be non-contiguous from split"
+    assert a.stride(0) == 2 * num_heads
+    assert b.stride(0) == 2 * num_heads
+    return b, a, b.contiguous(), a.contiguous()
+
+
+@unittest.skipIf(not torch.cuda.is_available(), "Test requires CUDA")
+class TestFusedGdnGatingNonContiguous(unittest.TestCase):
+    """Test fused_gdn_gating with non-contiguous a/b."""
+
+    def _run_test(self, batch, num_heads):
+        A_log = torch.randn(num_heads, dtype=torch.float32, device="cuda")
+        dt_bias = torch.randn(num_heads, dtype=torch.bfloat16, device="cuda")
+
+        b, a, b_contig, a_contig = _make_noncontiguous_ab(batch, num_heads)
+
+        g_ref, beta_ref = fused_gdn_gating(A_log, a_contig, b_contig, dt_bias)
+        g_test, beta_test = fused_gdn_gating(A_log, a, b, dt_bias)
+
+        self.assertTrue(
+            torch.allclose(g_test, g_ref, rtol=0, atol=0),
+            f"g mismatch: max diff = {(g_test - g_ref).abs().max().item()}",
+        )
+        self.assertTrue(
+            torch.allclose(beta_test, beta_ref, rtol=0, atol=0),
+            f"beta mismatch: max diff = {(beta_test - beta_ref).abs().max().item()}",
+        )
+
+    def test_small(self):
+        self._run_test(batch=4, num_heads=8)
+
+    def test_qwen35_27b_tp1(self):
+        """Qwen3.5-27B TP=1: nv_tp=48."""
+        self._run_test(batch=16, num_heads=48)
+
+    def test_qwen35_27b_tp2(self):
+        """Qwen3.5-27B TP=2: nv_tp=24."""
+        self._run_test(batch=32, num_heads=24)
+
+    def test_single_batch(self):
+        self._run_test(batch=1, num_heads=48)
+
+
+@unittest.skipIf(not torch.cuda.is_available(), "Test requires CUDA")
+class TestFusedSigmoidGatingDeltaRuleUpdateNonContiguous(unittest.TestCase):
+    """Test fused_sigmoid_gating_delta_rule_update with non-contiguous a/b."""
+
+    def _run_test(self, batch, T, num_v_heads, head_k_dim, head_v_dim):
+        num_k_heads = num_v_heads  # simplification for GDN
+        HV = num_v_heads
+        K = head_k_dim
+        V = head_v_dim
+
+        A_log = torch.randn(HV, dtype=torch.float32, device="cuda")
+        dt_bias = torch.randn(HV, dtype=torch.bfloat16, device="cuda")
+
+        q = torch.randn(batch, T, num_k_heads, K, dtype=torch.bfloat16, device="cuda")
+        k = torch.randn(batch, T, num_k_heads, K, dtype=torch.bfloat16, device="cuda")
+        v = torch.randn(batch, T, HV, V, dtype=torch.bfloat16, device="cuda")
+
+        # Simulate non-contiguous a/b from split
+        mixed_ba = torch.randn(batch * T, 2 * HV, dtype=torch.bfloat16, device="cuda")
+        b_nc, a_nc = mixed_ba.split([HV, HV], dim=-1)
+        b_c, a_c = b_nc.contiguous(), a_nc.contiguous()
+
+        # Build cu_seqlens for varlen (one token per sequence)
+        cu_seqlens = torch.arange(0, batch * T + 1, T, dtype=torch.int32, device="cuda")
+
+        cache_len = batch + 4
+        ssm_states = torch.zeros(
+            cache_len, HV, K, V, dtype=torch.float32, device="cuda"
+        )
+        state_indices = torch.arange(batch, dtype=torch.int32, device="cuda")
+
+        # Reference: contiguous a/b
+        ssm_ref = ssm_states.clone()
+        out_ref = fused_sigmoid_gating_delta_rule_update(
+            A_log=A_log,
+            dt_bias=dt_bias,
+            q=q,
+            k=k,
+            v=v,
+            a=a_c,
+            b=b_c,
+            initial_state_source=ssm_ref,
+            initial_state_indices=state_indices,
+            cu_seqlens=cu_seqlens,
+            softplus_beta=1.0,
+            softplus_threshold=20.0,
+            is_kda=False,
+        )
+
+        # Test: non-contiguous a/b
+        ssm_test = ssm_states.clone()
+        out_test = fused_sigmoid_gating_delta_rule_update(
+            A_log=A_log,
+            dt_bias=dt_bias,
+            q=q,
+            k=k,
+            v=v,
+            a=a_nc,
+            b=b_nc,
+            initial_state_source=ssm_test,
+            initial_state_indices=state_indices,
+            cu_seqlens=cu_seqlens,
+            softplus_beta=1.0,
+            softplus_threshold=20.0,
+            is_kda=False,
+        )
+
+        max_out_diff = (out_test - out_ref).abs().max().item()
+        max_state_diff = (ssm_test - ssm_ref).abs().max().item()
+
+        self.assertTrue(
+            torch.allclose(out_test, out_ref, rtol=0, atol=0),
+            f"output mismatch: max diff = {max_out_diff}",
+        )
+        self.assertTrue(
+            torch.allclose(ssm_test, ssm_ref, rtol=0, atol=0),
+            f"state mismatch: max diff = {max_state_diff}",
+        )
+
+    def test_decode_single_token(self):
+        """Standard decode: T=1, batch>1."""
+        self._run_test(batch=4, T=1, num_v_heads=8, head_k_dim=64, head_v_dim=32)
+
+    def test_qwen35_decode(self):
+        """Qwen3.5-27B like config: HV=48."""
+        self._run_test(batch=8, T=1, num_v_heads=48, head_k_dim=128, head_v_dim=128)
+
+    def test_multi_token(self):
+        """target_verify style: T>1."""
+        self._run_test(batch=4, T=4, num_v_heads=8, head_k_dim=64, head_v_dim=32)
+
+
+@unittest.skipIf(not torch.cuda.is_available(), "Test requires CUDA")
+class TestFusedSigmoidGatingKDAStride(unittest.TestCase):
+    """Regression test: KDA path handles non-contiguous a/b after stride_a refactor."""
+
+    def test_kda_noncontiguous_matches_contiguous(self):
+        """KDA path should produce identical outputs/states for contiguous vs non-contiguous a/b."""
+        token_num = 4
+        num_heads = 8
+        head_dim = 128
+        HV = num_heads
+        K = head_dim
+
+        A_log = torch.randn(1, 1, HV, 1, dtype=torch.float32, device="cuda")
+        dt_bias = torch.randn(HV * K, dtype=torch.bfloat16, device="cuda")
+
+        mixed_a = torch.randn(
+            token_num, 2 * HV * K, dtype=torch.bfloat16, device="cuda"
+        )
+        a_nc, _ = mixed_a.split([HV * K, HV * K], dim=-1)
+        a_c = a_nc.contiguous()
+        self.assertFalse(a_nc.is_contiguous())
+
+        mixed_b = torch.randn(1, token_num, 2 * HV, dtype=torch.bfloat16, device="cuda")
+        b_nc, _ = mixed_b.split([HV, HV], dim=-1)
+        b_c = b_nc.contiguous()
+        self.assertFalse(b_nc.is_contiguous())
+
+        q = torch.randn(1, token_num, HV, K, dtype=torch.bfloat16, device="cuda")
+        k = torch.randn(1, token_num, HV, K, dtype=torch.bfloat16, device="cuda")
+        v = torch.randn(1, token_num, HV, K, dtype=torch.bfloat16, device="cuda")
+
+        cu_seqlens = torch.tensor([0, 1, 2, 3, 4], device="cuda", dtype=torch.int32)
+        cache_len = 64
+        ssm_states = torch.zeros(
+            cache_len, HV, K, K, dtype=torch.float32, device="cuda"
+        )
+        cache_indices = torch.tensor([0, 2, 5, 8], device="cuda", dtype=torch.int32)
+
+        # Reference: contiguous a/b
+        ssm_ref = ssm_states.clone()
+        out_ref = fused_sigmoid_gating_delta_rule_update(
+            A_log=A_log,
+            dt_bias=dt_bias,
+            q=q,
+            k=k,
+            v=v,
+            a=a_c,
+            b=b_c,
+            initial_state_source=ssm_ref,
+            initial_state_indices=cache_indices,
+            cu_seqlens=cu_seqlens,
+            use_qk_l2norm_in_kernel=True,
+            softplus_beta=1.0,
+            softplus_threshold=20.0,
+            is_kda=True,
+        )
+
+        # Test: non-contiguous a/b from split
+        ssm_test = ssm_states.clone()
+        out_test = fused_sigmoid_gating_delta_rule_update(
+            A_log=A_log,
+            dt_bias=dt_bias,
+            q=q,
+            k=k,
+            v=v,
+            a=a_nc,
+            b=b_nc,
+            initial_state_source=ssm_test,
+            initial_state_indices=cache_indices,
+            cu_seqlens=cu_seqlens,
+            use_qk_l2norm_in_kernel=True,
+            softplus_beta=1.0,
+            softplus_threshold=20.0,
+            is_kda=True,
+        )
+
+        self.assertTrue(
+            torch.allclose(out_test, out_ref, rtol=0, atol=0),
+            f"KDA output mismatch: max diff = {(out_test - out_ref).abs().max().item()}",
+        )
+        self.assertTrue(
+            torch.allclose(ssm_test, ssm_ref, rtol=0, atol=0),
+            f"KDA state mismatch: max diff = {(ssm_test - ssm_ref).abs().max().item()}",
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()