tq4_sdpa: address PR review feedback

- gemma4_31b cuda_source_transformations: pass mask_is_causal=True to
  tq4_sdpa (the 12th arg was missing). Enables the prefill per-tile
  causal block-skip (~halves causal-triangle work); no-op for decode.
- tq4_sdpa kernel + tests: drop model-specific (Gemma/Qwen) references;
  describe configs by shape (head_dim / GQA ratio) instead.
- tests: add split-K coverage for B=2 (test_splitk_batch2) and a
  non-contiguous query stride (test_splitk_noncontiguous_query).

Author: Gasoonjia

backends/cuda/tests/test_tq4_sdpa.py

@@ -252,7 +252,7 @@ def test_gqa_prefill(self):
                 self._run_test(1, H_q, H_kv, 64, 64, 128, is_causal=True)
 
     def test_gqa_8x_head_dim_256(self):
-        """GQA 8:1 with head_dim=256 — matches Qwen 3.5 MoE config."""
+        """GQA 8:1 with head_dim=256."""
         self._run_test(1, 16, 2, 1, 128, 256)
         L = 64
         mask = torch.tril(torch.ones(1, 1, L, L, dtype=torch.bool, device="cuda"))
@@ -374,8 +374,8 @@ def test_float_mask_rejected(self):
                 float_mask,
             )
 
-    def test_qwen35_moe_config(self):
-        """Qwen 3.5 MoE: head_dim=256, GQA 16:2, decode + prefill."""
+    def test_config_hd256_gqa_16_2(self):
+        """head_dim=256, GQA 16:2, decode + prefill."""
         self._run_test(1, 16, 2, 1, 256, 256)
         self._run_test(1, 16, 2, 128, 128, 256, is_causal=True)
 
@@ -443,9 +443,8 @@ def test_output_shape_and_dtype(self):
     # Every test above calls tq4_sdpa WITHOUT kv_len and WITHOUT
     # mask_is_causal, so they only exercise the kv_len=None fallback
     # (full-Lk loop) at short KV. The cases below drive the actual
-    # long-context paths used in production by the Gemma-4 31B global
-    # layers (head_dim=512, GQA 8:4) and Qwen 3.5 MoE (head_dim=256,
-    # GQA 16:2):
+    # long-context paths at two representative GQA shapes (head_dim=512
+    # GQA 8:4, and head_dim=256 GQA 16:2):
     #   * the on-device kv_len scalar that bounds the KV loop to the
     #     filled context (decode), and
     #   * the mask_is_causal per-tile causal block-skip (prefill).
@@ -466,7 +465,7 @@ def test_output_shape_and_dtype(self):
     # of a kv_len-long context) there are two ways to place the causal
     # triangle. PyTorch F.sdpa(is_causal=True) uses TOP-LEFT alignment
     # (query row i attends to keys [0, i]) -- wrong for a KV cache. This
-    # kernel and gemma4_31b/model.py::_build_masks use BOTTOM-RIGHT
+    # kernel (and a KV-cache decoder's mask builder) use BOTTOM-RIGHT
     # alignment: query row i is absolute position (kv_len - Lq + i) and
     # attends to keys [0, kv_len - Lq + i]. So the reference below builds
     # an explicit bottom-right mask (q_pos >= cache_pos) rather than
@@ -498,12 +497,12 @@ def _run_long_kv_test(
         causal mask) must confine attention to ``[0, kv_len)``.
 
         ``causal=True`` builds a bottom-right-aligned mask (the Lq queries
-        are the last Lq positions of a kv_len-long context), mirroring the
-        production ``q_pos >= cache_pos`` mask in gemma4_31b/model.py
-        ``_build_masks`` and the kernel's ``(kv_len - Lq) + seq_pos`` block
-        bound. We deliberately do NOT use ``F.sdpa(is_causal=True)`` for the
-        reference: PyTorch aligns is_causal top-left when L_q < L_kv, while
-        this kernel (and the model) align bottom-right.
+        are the last Lq positions of a kv_len-long context), mirroring a
+        KV-cache decoder's ``q_pos >= cache_pos`` mask and the kernel's
+        ``(kv_len - Lq) + seq_pos`` block bound. We deliberately do NOT use
+        ``F.sdpa(is_causal=True)`` for the reference: PyTorch aligns
+        is_causal top-left when L_q < L_kv, while this kernel (and such a
+        decoder) align bottom-right.
         """
         torch.manual_seed(seed)
         centroids, boundaries, rotation = _make_codebook_and_rotation(D)
@@ -592,6 +591,7 @@ def _run_splitk_vs_fused_test(
         Lq,
         kv_len,
         buffer_len,
+        B=1,
         seed=42,
     ):
         """Verify split-K output matches fused kernel output for same inputs.
@@ -606,7 +606,6 @@ def _run_splitk_vs_fused_test(
         boundaries = boundaries.cuda()
         rotation = rotation.cuda()
 
-        B = 1
         k = torch.randn(B, H_kv, buffer_len, D, dtype=torch.bfloat16, device="cuda")
         v = torch.randn(B, H_kv, buffer_len, D, dtype=torch.bfloat16, device="cuda")
         # Add garbage tail to ensure split-K respects kv_len bound
@@ -671,22 +670,85 @@ def _run_splitk_vs_fused_test(
             cos,
             0.99,
             f"Split-K vs Fused cosine {cos:.5f} < 0.99 "
-            f"(H_q={H_q} H_kv={H_kv} D={D} kv_len={kv_len})",
+            f"(B={B} H_q={H_q} H_kv={H_kv} D={D} kv_len={kv_len})",
+        )
+
+    def test_splitk_batch2(self):
+        """Split-K decode (Lq=1) with batch size B=2.
+
+        Exercises the per-batch indexing in the split-K and reduce kernels
+        (b = pid_bh // H_grid). Split-K output must match the fused-kernel
+        path for the same inputs."""
+        self._run_splitk_vs_fused_test(
+            H_q=16, H_kv=2, D=256, Lq=1, kv_len=512, buffer_len=1024, B=2
+        )
+
+    def test_splitk_noncontiguous_query(self):
+        """Split-K decode (Lq=1, B=2) with a non-contiguous query.
+
+        The host wrapper rotates Q (Q @ Pi^T) before launching the kernel,
+        so a strided query must yield the same result as its contiguous
+        copy. Builds a query whose last-dim stride is 2 by slicing a padded
+        buffer, then checks it matches the contiguous query."""
+        H_q, H_kv, D, kv_len, B = 16, 2, 256, 512, 2
+        torch.manual_seed(42)
+        centroids, boundaries, rotation = _make_codebook_and_rotation(D)
+        centroids = centroids.cuda()
+        boundaries = boundaries.cuda()
+        rotation = rotation.cuda()
+
+        k = torch.randn(B, H_kv, kv_len, D, dtype=torch.bfloat16, device="cuda")
+        v = torch.randn(B, H_kv, kv_len, D, dtype=torch.bfloat16, device="cuda")
+        k_packed, k_norms = _compress(k, boundaries, rotation)
+        v_packed, v_norms = _compress(v, boundaries, rotation)
+
+        q = torch.randn(B, H_q, 1, D, dtype=torch.bfloat16, device="cuda")
+        # Non-contiguous alias with identical values (last-dim stride 2).
+        q_pad = torch.empty(B, H_q, 1, D, 2, dtype=torch.bfloat16, device="cuda")
+        q_pad[..., 0] = q
+        q_nc = q_pad[..., 0]
+        self.assertFalse(q_nc.is_contiguous(), "query should be non-contiguous")
+
+        kv_len_t = torch.tensor([kv_len], dtype=torch.int32, device="cuda")
+
+        def _run(query):
+            return self.tq4_sdpa(
+                query,
+                k_packed,
+                k_norms,
+                v_packed,
+                v_norms,
+                centroids,
+                rotation,
+                attn_mask=None,
+                is_causal=False,
+                scale=None,
+                kv_len=kv_len_t,
+                mask_is_causal=False,
+            )
+
+        out_contig = _run(q)
+        out_nc = _run(q_nc)
+
+        self.assertFalse(torch.isnan(out_nc).any(), "NaN in non-contiguous output")
+        cos = _cosine_sim(out_nc, out_contig)
+        self.assertGreater(
+            cos, 0.999, f"non-contiguous vs contiguous query cosine {cos:.5f}"
         )
 
-    def test_kv_len_clamp_decode_gemma_global(self):
-        """Decode (Lq=1) kv_len clamp at Gemma-4 31B global-layer shape
-        (head_dim=512, GQA 8:4). N=8192 leaves a 24k garbage tail in a 32k
-        buffer (clamp guard); N=32768 fills the buffer (full 32k loop)."""
+    def test_kv_len_clamp_decode_hd512_gqa_8_4(self):
+        """Decode (Lq=1) kv_len clamp at a head_dim=512, GQA 8:4 shape.
+        N=8192 leaves a 24k garbage tail in a 32k buffer (clamp guard);
+        N=32768 fills the buffer (full 32k loop)."""
         for N in (8192, 32768):
             with self.subTest(N=N):
                 self._run_long_kv_test(
                     H_q=8, H_kv=4, D=512, Lq=1, kv_len=N, buffer_len=32768
                 )
 
-    def test_kv_len_clamp_decode_gemma_global_splitk(self):
-        """Split-K decode (Lq=1) at Gemma global shape with long KV.
-        Verifies split-K output matches BOTH (a) fp32 reference over first
+    def test_kv_len_clamp_decode_hd512_gqa_8_4_splitk(self):
+        """Split-K decode (Lq=1) at a head_dim=512, GQA 8:4 shape with long
+        KV. Verifies split-K output matches BOTH (a) fp32 reference over first
         kv_len positions AND (b) existing fused-kernel output (byte-identical
         within fp tolerance). Uses garbage tail as negative control."""
         for N in (8192, 32768):
@@ -707,18 +769,18 @@ def test_kv_len_clamp_decode_gemma_global_splitk(self):
                     H_q=8, H_kv=4, D=512, Lq=1, kv_len=N, buffer_len=32768
                 )
 
-    def test_kv_len_clamp_decode_qwen(self):
-        """Decode (Lq=1) kv_len clamp at Qwen 3.5 MoE shape
-        (head_dim=256, GQA 16:2)."""
+    def test_kv_len_clamp_decode_hd256_gqa_16_2(self):
+        """Decode (Lq=1) kv_len clamp at a head_dim=256, GQA 16:2 shape."""
         for N in (8192, 32768):
             with self.subTest(N=N):
                 self._run_long_kv_test(
                     H_q=16, H_kv=2, D=256, Lq=1, kv_len=N, buffer_len=32768
                 )
 
-    def test_kv_len_clamp_decode_qwen_splitk(self):
-        """Split-K decode (Lq=1) at Qwen shape with long KV.
-        Verifies split-K output matches BOTH fp32 reference AND fused kernel."""
+    def test_kv_len_clamp_decode_hd256_gqa_16_2_splitk(self):
+        """Split-K decode (Lq=1) at a head_dim=256, GQA 16:2 shape with long
+        KV. Verifies split-K output matches BOTH fp32 reference AND fused
+        kernel."""
         for N in (8192, 32768):
             with self.subTest(N=N):
                 _ = self._run_long_kv_test(
@@ -734,12 +796,12 @@ def test_kv_len_clamp_decode_qwen_splitk(self):
                     H_q=16, H_kv=2, D=256, Lq=1, kv_len=N, buffer_len=32768
                 )
 
-    def test_mask_is_causal_prefill_gemma_global(self):
-        """Chunked prefill (Lq>1) with mask_is_causal at Gemma global shape.
-        The Lq queries are the last Lq of a kv_len-long context; the
-        per-tile causal block-skip plus bottom-right mask must match the
-        fp32 causal reference over the first kv_len positions. A garbage
-        tail beyond kv_len also exercises the clamp."""
+    def test_mask_is_causal_prefill_hd512_gqa_8_4(self):
+        """Chunked prefill (Lq>1) with mask_is_causal at a head_dim=512,
+        GQA 8:4 shape. The Lq queries are the last Lq of a kv_len-long
+        context; the per-tile causal block-skip plus bottom-right mask must
+        match the fp32 causal reference over the first kv_len positions. A
+        garbage tail beyond kv_len also exercises the clamp."""
         for Lq, kv_len, buf in ((256, 4096, 8192), (2048, 8192, 16384)):
             with self.subTest(Lq=Lq, kv_len=kv_len):
                 self._run_long_kv_test(
@@ -752,8 +814,9 @@ def test_mask_is_causal_prefill_gemma_global(self):
                     causal=True,
                 )
 
-    def test_mask_is_causal_prefill_qwen(self):
-        """Chunked prefill (Lq>1) with mask_is_causal at Qwen shape."""
+    def test_mask_is_causal_prefill_hd256_gqa_16_2(self):
+        """Chunked prefill (Lq>1) with mask_is_causal at a head_dim=256,
+        GQA 16:2 shape."""
         for Lq, kv_len, buf in ((256, 4096, 8192), (2048, 8192, 16384)):
             with self.subTest(Lq=Lq, kv_len=kv_len):
                 self._run_long_kv_test(
@@ -766,11 +829,11 @@ def test_mask_is_causal_prefill_qwen(self):
                     causal=True,
                 )
 
-    def test_kv_len_none_fallback_qwen(self):
+    def test_kv_len_none_fallback_hd256_gqa_16_2(self):
         """Regression: the kv_len=None fallback (HAS_KV_LEN False, full-Lk
-        loop) that the Qwen path relies on still matches the fp32 reference.
-        This guards the original behavior the kv_len feature must preserve
-        for callers that pass neither kv_len nor mask_is_causal."""
+        loop) still matches the fp32 reference. This guards the original
+        behavior the kv_len feature must preserve for callers that pass
+        neither kv_len nor mask_is_causal."""
         self._run_long_kv_test(
             H_q=16,
             H_kv=2,
@@ -787,9 +850,9 @@ def test_kv_len_none_fallback_qwen(self):
         "128k case is heavy for the 24GB CI runner; set TQ4_RUN_128K=1 to run",
     )
     def test_kv_len_clamp_128k(self):
-        """Full 131072-entry buffer (Qwen shape). (a) kv_len=8192 with a
-        ~123k garbage tail — the clamp keeps decode O(context) and never
-        touches the tail; (b) kv_len=131072 — correctness at true 128k
+        """Full 131072-entry buffer (head_dim=256, GQA 16:2). (a) kv_len=8192
+        with a ~123k garbage tail — the clamp keeps decode O(context) and
+        never touches the tail; (b) kv_len=131072 — correctness at true 128k
         scale. Gated behind TQ4_RUN_128K because the fp32 reference for (b)
         needs >~6GB and CI runs on a 24GB A10G."""
         self._run_long_kv_test(

backends/cuda/triton/kernels/tq4_sdpa.py

@@ -175,8 +175,8 @@ def _tq4_sdpa_fwd_kernel_body(
     # chunk it is chunk_end. This makes the global-layer attention O(context)
     # rather than O(max_seq_len) — the empty tail of the cache is never touched.
     # kv_len is read from a GPU scalar so the bound updates across CUDA-graph
-    # replays (decode is graph-captured). When not provided (HAS_KV_LEN False,
-    # e.g. qwen) it falls back to Lk, preserving the original behavior exactly.
+    # replays (decode is graph-captured). When not provided (HAS_KV_LEN False)
+    # it falls back to Lk, preserving the original behavior exactly.
     if HAS_KV_LEN:
         kv_len = tl.load(KV_LEN_ptr)
     else:
@@ -191,8 +191,8 @@ def _tq4_sdpa_fwd_kernel_body(
     # prefill analogue of the kv_len decode clamp and ~halves the causal-triangle
     # work. For decode (Lq=1, max(seq_pos)=0) this evaluates to kv_len, so decode
     # is byte-identical. Applied only when MASK_IS_CAUSAL (the caller guarantees a
-    # causal mask, e.g. Gemma's full-attention layers); otherwise the full kv_len
-    # bound is kept, which is safe for an arbitrary mask.
+    # causal mask); otherwise the full kv_len bound is kept, which is safe for an
+    # arbitrary mask.
     loop_end = kv_len
     if MASK_IS_CAUSAL:
         max_q_pos = (kv_len - Lq) + tl.max(seq_pos)
@@ -714,8 +714,9 @@ def tq4_sdpa(
         is_causal: apply causal masking (requires L_Q == L_KV)
         scale: softmax scale applied to ``Q @ K^T``. Defaults to
             ``1/sqrt(HEAD_DIM)`` when ``None``. Models that handle their
-            own normalization (e.g. Gemma 4 with QK-norm uses ``1.0``)
-            should pass an explicit value.
+            own normalization (e.g. QK-norm models that fold the
+            ``1/sqrt(d)`` factor into the trained weights, which then use
+            ``1.0``) should pass an explicit value.
         kv_len: Optional GPU int scalar = number of valid (filled) KV
             positions. When provided, the inner KV loop is bounded to
             ``kv_len`` instead of the full pre-allocated ``L_KV``, making
@@ -876,7 +877,7 @@ def tq4_sdpa(
 # Split-K decode kernel (flash-decoding) for TQ4
 # ==============================================================================
 # When L_q == 1 with GQA, the standard kernel launches only
-# ceil(num_groups / BLOCK_M) * B * H_kv CTAs (e.g. 4 for gemma global 8:4).
+# ceil(num_groups / BLOCK_M) * B * H_kv CTAs (e.g. 4 for a B=1, 8:4 GQA shape).
 # Split-K partitions the KV sequence across many CTAs for better occupancy,
 # then reduces partial results in a second kernel.
 
@@ -1201,7 +1202,7 @@ def _launch_tq4_decode_splitk(
         H_grid = H_KV
         # Size BLOCK_M to the actually-packed rows (L_q * num_groups; for decode
         # L_q=1 so == num_groups), rounded up to the bf16 tensor-core MMA minimum
-        # of 16 -- instead of a fixed 64. Gemma4 global has num_groups=8, so the
+        # of 16 -- instead of a fixed 64. For example with num_groups=8, the
         # old BLOCK_M=64 left 56/64 M-rows idle (the QK/PV MMAs still computed all
         # 64 rows) AND made acc[BLOCK_M, HEAD_DIM] fp32 = 64*512*4 = 128 KB/CTA,
         # which blows past the register file and spills to local memory. Matching

examples/models/gemma4_31b/cuda_source_transformations.py

@@ -100,9 +100,10 @@ def _turboquant_attention_forward(
         self.kv_cache.centroids,
         self.kv_cache.rotation,
         attn_mask,
-        False,  # is_causal — attn_mask already encodes causal masking
+        False,  # is_causal: attn_mask already encodes causal masking
         self.scaling,
         kv_len,
+        True,  # mask_is_causal: Gemma full-attention mask is standard causal
     )
 
     y = y.transpose(1, 2).contiguous().view(B, T, self.n_heads * self.head_dim)