diff --git a/backends/cuda/tests/test_tq4_sdpa.py b/backends/cuda/tests/test_tq4_sdpa.py index f4cc1d770ef..f9543b1ff18 100644 --- a/backends/cuda/tests/test_tq4_sdpa.py +++ b/backends/cuda/tests/test_tq4_sdpa.py @@ -20,7 +20,6 @@ import numpy as np import torch import torch.nn.functional as F - from executorch.backends.cuda.cuda_backend import CudaBackend from executorch.backends.cuda.cuda_partitioner import CudaPartitioner from executorch.backends.cuda.triton.kernels.tq4_sdpa import tq4_sdpa @@ -253,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")) @@ -375,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) @@ -444,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). @@ -467,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 @@ -499,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) @@ -584,31 +582,226 @@ def _run_long_kv_test( ) return cos - 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 _run_splitk_vs_fused_test( + self, + *, + H_q, + H_kv, + D, + Lq, + kv_len, + buffer_len, + B=1, + seed=42, + ): + """Verify split-K output matches fused kernel output for same inputs. + + Runs tq4_sdpa twice: once with kv_len (triggers split-K for Lq=1, kv_len>=256), + and once without kv_len (forces fused kernel path). Both outputs must match + within fp tolerance, proving split-K computes the same result. + """ + torch.manual_seed(seed) + centroids, boundaries, rotation = _make_codebook_and_rotation(D) + centroids = centroids.cuda() + boundaries = boundaries.cuda() + rotation = rotation.cuda() + + 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 + if buffer_len > kv_len: + g = buffer_len - kv_len + k[:, :, kv_len:, :] = ( + torch.randn(B, H_kv, g, D, dtype=torch.bfloat16, device="cuda") * 1000.0 + ) + v[:, :, kv_len:, :] = ( + torch.randn(B, H_kv, g, D, dtype=torch.bfloat16, device="cuda") * 1000.0 + ) + + q = torch.randn(B, H_q, Lq, D, dtype=torch.bfloat16, device="cuda") + + k_packed, k_norms = _compress(k, boundaries, rotation) + v_packed, v_norms = _compress(v, boundaries, rotation) + + # Split-K path: with kv_len (triggers split-K for Lq=1, kv_len>=256) + kv_len_t = torch.tensor([kv_len], dtype=torch.int32, device="cuda") + out_splitk = self.tq4_sdpa( + q, + 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, + ) + + # Fused kernel path: without kv_len (forces fused kernel) + # But we need to slice the buffer to kv_len to avoid garbage + k_packed_sliced = k_packed[:, :, :kv_len, :] + k_norms_sliced = k_norms[:, :, :kv_len, :] + v_packed_sliced = v_packed[:, :, :kv_len, :] + v_norms_sliced = v_norms[:, :, :kv_len, :] + + out_fused = self.tq4_sdpa( + q, + k_packed_sliced, + k_norms_sliced, + v_packed_sliced, + v_norms_sliced, + centroids, + rotation, + attn_mask=None, + is_causal=False, + scale=None, + kv_len=None, + mask_is_causal=False, + ) + + # Both outputs must match (split-K computes same result as fused) + self.assertFalse(torch.isnan(out_splitk).any(), "NaN in split-K output") + self.assertFalse(torch.isnan(out_fused).any(), "NaN in fused output") + cos = _cosine_sim(out_splitk, out_fused) + self.assertGreater( + cos, + 0.99, + f"Split-K vs Fused cosine {cos:.5f} < 0.99 " + 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_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_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_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): + with self.subTest(N=N): + # Run with split-K (kv_len >= 256 triggers split-K) + _ = self._run_long_kv_test( + H_q=8, + H_kv=4, + D=512, + Lq=1, + kv_len=N, + buffer_len=32768, + min_cosine=0.99, + ) + # Also verify split-K matches fused kernel by running without kv_len + # (which forces fused kernel path) and comparing outputs + self._run_splitk_vs_fused_test( + H_q=8, H_kv=4, D=512, Lq=1, kv_len=N, buffer_len=32768 + ) + + 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_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_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( + H_q=16, + H_kv=2, + D=256, + Lq=1, + kv_len=N, + buffer_len=32768, + min_cosine=0.99, + ) + self._run_splitk_vs_fused_test( + H_q=16, H_kv=2, D=256, Lq=1, kv_len=N, buffer_len=32768 + ) + + 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( @@ -621,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( @@ -635,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, @@ -656,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( diff --git a/backends/cuda/triton/kernels/tq4_sdpa.py b/backends/cuda/triton/kernels/tq4_sdpa.py index e1576f7e446..10f02c7fa3c 100644 --- a/backends/cuda/triton/kernels/tq4_sdpa.py +++ b/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 @@ -798,37 +799,509 @@ def tq4_sdpa( block_m = 32 if total_ctas_m64 < 4 * 84 else 64 pack_gqa = _should_pack_gqa(N_Q, num_groups, block_m) - _launch_tq4_kernel( + # Split-K decode dispatch: L_q == 1 AND kv_len >= threshold (256) + # Uses flash-decoding to partition KV across many CTAs for better occupancy + # Dispatch is static (based on buffer size N_KV, not runtime kv_len value) + # to be export/AOTI traceable. The kernel still uses kv_len on-device + # via tl.load for bounds checking (CUDA-graph safe). + _SPLITK_LKV_THRESHOLD = 256 + use_splitk = ( + N_Q == 1 + and HAS_KV_LEN + and kv_len_t is not None + and N_KV >= _SPLITK_LKV_THRESHOLD + ) + + if use_splitk: + _launch_tq4_decode_splitk( + q_rot, + k_packed, + k_n, + v_packed, + v_n, + lut_hi, + lut_lo, + mask_ptr, + out_rot, + kv_len_t, + B, + H_Q, + H_KV, + N_Q, + N_KV, + D, + sm_scale, + HAS_MASK, + HAS_KV_LEN, + stride_mb, + stride_mq, + stride_mk, + num_groups, + pack_gqa, + ) + else: + _launch_tq4_kernel( + q_rot, + k_packed, + k_n, + v_packed, + v_n, + lut_hi, + lut_lo, + mask_ptr, + out_rot, + kv_len_t, + B, + H_Q, + H_KV, + N_Q, + N_KV, + D, + sm_scale, + HAS_MASK, + HAS_KV_LEN, + MASK_IS_CAUSAL, + stride_mb, + stride_mq, + stride_mk, + is_causal, + num_groups, + pack_gqa, + ) + + # Post-rotate: convert from rotated space back to original space + return torch.matmul(out_rot, rotation.to(query.dtype)) + + +# ============================================================================== +# 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 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. + + +@triton.autotune( + configs=[ + triton.Config({"BLOCK_N": 32}, num_warps=2, num_stages=1), + triton.Config({"BLOCK_N": 32}, num_warps=4, num_stages=1), + triton.Config({"BLOCK_N": 64}, num_warps=2, num_stages=1), + triton.Config({"BLOCK_N": 64}, num_warps=4, num_stages=1), + triton.Config({"BLOCK_N": 64}, num_warps=4, num_stages=2), + triton.Config({"BLOCK_N": 128}, num_warps=4, num_stages=1), + triton.Config({"BLOCK_N": 128}, num_warps=4, num_stages=2), + triton.Config({"BLOCK_N": 128}, num_warps=4, num_stages=3), + triton.Config({"BLOCK_N": 128}, num_warps=8, num_stages=2), + triton.Config({"BLOCK_N": 256}, num_warps=4, num_stages=2), + triton.Config({"BLOCK_N": 256}, num_warps=8, num_stages=2), + ], + key=["Lk", "HEAD_DIM", "NUM_GROUPS", "HAS_MASK", "PACK_GQA"], +) +@triton.jit +def _tq4_sdpa_decode_splitk_kernel( + Q_ptr, + KP_ptr, + KN_ptr, + VP_ptr, + VN_ptr, + LUT_hi_ptr, + LUT_lo_ptr, + O_partial_ptr, + M_partial_ptr, + L_partial_ptr, + Mask_ptr, + KV_LEN_ptr, + B, + H_grid, + Lk, + stride_qb, + stride_qh, + stride_qm, + stride_qd, + stride_kpb, + stride_kph, + stride_kpn, + stride_kpd, + stride_knb, + stride_knh, + stride_knn, + stride_vpb, + stride_vph, + stride_vpn, + stride_vpd, + stride_vnb, + stride_vnh, + stride_vnn, + stride_op_s, + stride_op_b, + stride_op_h, + stride_op_d, + stride_mp_s, + stride_mp_b, + stride_mp_h, + stride_lp_s, + stride_lp_b, + stride_lp_h, + stride_mb, + stride_mq, + stride_mk, + sm_scale: tl.float32, + chunk_size, + HAS_MASK: tl.constexpr, + HAS_KV_LEN: tl.constexpr, + BLOCK_N: tl.constexpr, + HEAD_DIM: tl.constexpr, + HALF_D: tl.constexpr, + NUM_GROUPS: tl.constexpr, + PACK_GQA: tl.constexpr, + BLOCK_M: tl.constexpr, +): + split_id = tl.program_id(axis=0) + pid_bh = tl.program_id(axis=1) + b = pid_bh // H_grid + h_grid = pid_bh % H_grid + + # Compute KV chunk bounds for this split + start_n = split_id * chunk_size + if HAS_KV_LEN: + kv_len = tl.load(KV_LEN_ptr) + else: + kv_len = Lk + end_n = tl.minimum(start_n + chunk_size, kv_len) + + offs_d = tl.arange(0, HEAD_DIM) + offs_d_half = tl.arange(0, HALF_D) + offs_m = tl.arange(0, BLOCK_M) + + if PACK_GQA: + # Pack GQA: multiple Q heads folded into M dimension + seq_pos = offs_m // NUM_GROUPS + h_within = offs_m % NUM_GROUPS + h_q_rows = h_grid * NUM_GROUPS + h_within + h_kv = h_grid + row_valid = seq_pos < 1 # Lq=1 for decode + q_ptrs = Q_ptr + ( + b * stride_qb + + h_q_rows[:, None] * stride_qh + + seq_pos[:, None] * stride_qm + + offs_d[None, :] * stride_qd + ) + else: + seq_pos = offs_m + h_kv = h_grid // NUM_GROUPS + row_valid = offs_m < 1 + q_ptrs = Q_ptr + ( + b * stride_qb + + h_grid * stride_qh + + offs_m[:, None] * stride_qm + + offs_d[None, :] * stride_qd + ) + + q = tl.load(q_ptrs, mask=row_valid[:, None], other=0.0).to(tl.bfloat16) + + # Online softmax state + m_i = tl.full([BLOCK_M], -float("inf"), dtype=tl.float32) + l_i = tl.zeros([BLOCK_M], dtype=tl.float32) + acc = tl.zeros([BLOCK_M, HEAD_DIM], dtype=tl.float32) + + # Prescale for exp2-based softmax + qk_scale = sm_scale * 1.44269504 + + # TQ4 K/V base pointers (uniform: single KV head) + kp_base = KP_ptr + b * stride_kpb + h_kv * stride_kph + kn_base = KN_ptr + b * stride_knb + h_kv * stride_knh + vp_base = VP_ptr + b * stride_vpb + h_kv * stride_vph + vn_base = VN_ptr + b * stride_vnb + h_kv * stride_vnh + + offs_n_init = tl.arange(0, BLOCK_N) + + for tile_start in tl.range(start_n, end_n, BLOCK_N): + offs_n = tile_start + offs_n_init + kv_valid = offs_n < end_n + + # -- K decompression (LUT, no norm multiply on [N,D] tile) -- + kp_ptrs = ( + kp_base + offs_n[:, None] * stride_kpn + offs_d_half[None, :] * stride_kpd + ) + k_packed_data = tl.load(kp_ptrs, mask=kv_valid[:, None], other=0).to(tl.int32) + k = tl.join( + tl.load(LUT_hi_ptr + k_packed_data), + tl.load(LUT_lo_ptr + k_packed_data), + ).reshape(BLOCK_N, HEAD_DIM) + + # Q @ K^T with norm factored out + kn = tl.load(kn_base + offs_n * stride_knn, mask=kv_valid, other=0.0) + qk = (tl.dot(q, tl.trans(k)) * qk_scale * kn[None, :]).to(tl.float32) + + if HAS_MASK: + mask_ptrs = Mask_ptr + ( + b * stride_mb + + seq_pos[:, None] * stride_mq + + offs_n[None, :] * stride_mk + ) + mn_mask = row_valid[:, None] & kv_valid[None, :] + mask_block = tl.load(mask_ptrs, mask=mn_mask, other=False) + qk = tl.where(mask_block, qk, float("-inf")) + + qk = tl.where(kv_valid[None, :], qk, float("-inf")) + + # NaN-safe online softmax (exp2) + m_ij = tl.max(qk, 1) + m_new = tl.maximum(m_i, m_ij) + safe_alpha = tl.where(m_new > -float("inf"), m_i - m_new, 0.0) + alpha = tl.math.exp2(safe_alpha) + safe_p = tl.where( + m_new[:, None] > -float("inf"), qk - m_new[:, None], -float("inf") + ) + p = tl.math.exp2(safe_p) + l_ij = tl.sum(p, 1) + + # -- V decompression (LUT, norm factored into P) -- + vp_ptrs = ( + vp_base + offs_n[:, None] * stride_vpn + offs_d_half[None, :] * stride_vpd + ) + v_packed_data = tl.load(vp_ptrs, mask=kv_valid[:, None], other=0).to(tl.int32) + v = tl.join( + tl.load(LUT_hi_ptr + v_packed_data), + tl.load(LUT_lo_ptr + v_packed_data), + ).reshape(BLOCK_N, HEAD_DIM) + + # P @ (C*n) = (P*n) @ C — factor norm into P instead of V + vn = tl.load(vn_base + offs_n * stride_vnn, mask=kv_valid, other=0.0) + p_scaled = (p * vn[None, :]).to(tl.bfloat16) + acc = (acc * alpha[:, None] + tl.dot(p_scaled, v)).to(tl.float32) + l_i = (l_i * alpha + l_ij).to(tl.float32) + m_i = m_new + + # Store partial results + if PACK_GQA: + h_q_all = h_grid * NUM_GROUPS + h_within + o_ptrs = O_partial_ptr + ( + split_id * stride_op_s + + b * stride_op_b + + h_q_all[:, None] * stride_op_h + + offs_d[None, :] * stride_op_d + ) + m_ptrs = M_partial_ptr + ( + split_id * stride_mp_s + b * stride_mp_b + h_q_all * stride_mp_h + ) + l_ptrs = L_partial_ptr + ( + split_id * stride_lp_s + b * stride_lp_b + h_q_all * stride_lp_h + ) + tl.store(o_ptrs, acc, mask=row_valid[:, None]) + tl.store(m_ptrs, m_i, mask=row_valid) + tl.store(l_ptrs, l_i, mask=row_valid) + else: + o_ptrs = O_partial_ptr + ( + split_id * stride_op_s + + b * stride_op_b + + h_grid * stride_op_h + + offs_d * stride_op_d + ) + m_ptrs = M_partial_ptr + ( + split_id * stride_mp_s + b * stride_mp_b + h_grid * stride_mp_h + ) + l_ptrs = L_partial_ptr + ( + split_id * stride_lp_s + b * stride_lp_b + h_grid * stride_lp_h + ) + tl.store(o_ptrs, acc, mask=row_valid[:, None]) + tl.store(m_ptrs, m_i, mask=row_valid) + tl.store(l_ptrs, l_i, mask=row_valid) + + +@triton.jit +def _tq4_sdpa_decode_reduce_kernel( + O_partial_ptr, + M_partial_ptr, + L_partial_ptr, + O_ptr, + num_splits, + stride_op_s, + stride_op_b, + stride_op_h, + stride_op_d, + stride_mp_s, + stride_mp_b, + stride_mp_h, + stride_lp_s, + stride_lp_b, + stride_lp_h, + stride_ob, + stride_oh, + stride_om, + stride_od, + HEAD_DIM: tl.constexpr, +): + pid = tl.program_id(axis=0) + offs_d = tl.arange(0, HEAD_DIM) + + # Find global max across splits + m_global = tl.full([1], -float("inf"), dtype=tl.float32) + for s in tl.range(0, num_splits): + m_ptr = M_partial_ptr + s * stride_mp_s + pid * stride_mp_h + m_s = tl.load(m_ptr) + m_global = tl.maximum(m_global, m_s) + + # Rescale and sum partials + acc = tl.zeros([HEAD_DIM], dtype=tl.float32) + l_global = tl.zeros([1], dtype=tl.float32) + + for s in tl.range(0, num_splits): + m_ptr = M_partial_ptr + s * stride_mp_s + pid * stride_mp_h + l_ptr = L_partial_ptr + s * stride_lp_s + pid * stride_lp_h + o_ptrs = O_partial_ptr + ( + s * stride_op_s + pid * stride_op_h + offs_d * stride_op_d + ) + + m_s = tl.load(m_ptr) + l_s = tl.load(l_ptr) + o_s = tl.load(o_ptrs) + + # Rescale by exp(m_s - m_global) + alpha = tl.where(m_global > -float("inf"), tl.math.exp2(m_s - m_global), 0.0) + acc += o_s * alpha + l_global += l_s * alpha + + inv_l = tl.where(l_global > 0, 1.0 / l_global, 0.0) + acc = acc * inv_l + + o_out_ptrs = O_ptr + pid * stride_oh + offs_d * stride_od + tl.store(o_out_ptrs, acc.to(tl.bfloat16)) + + +def _launch_tq4_decode_splitk( + q_rot, + k_packed, + k_norms, + v_packed, + v_norms, + lut_hi, + lut_lo, + mask_ptr, + out_rot, + kv_len_ptr, + B, + H_Q, + H_KV, + L_Q, + L_KV, + D, + sm_scale, + HAS_MASK, + HAS_KV_LEN, + stride_mb, + stride_mq, + stride_mk, + num_groups, + pack_gqa, +): + HALF_D = D // 2 + + if pack_gqa: + 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. 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 + # BLOCK_M to num_groups removes both the wasted MMA rows and the spills. + _packed_m = L_Q * num_groups + BLOCK_M = max(16, 1 << (_packed_m - 1).bit_length()) + else: + H_grid = H_Q + BLOCK_M = 32 + + # Static num_splits for CUDA-graph compatibility + # Derive from buffer size, not runtime kv_len value + num_splits = min(max(triton.cdiv(L_KV, 256), 1), 128) + chunk_size = triton.cdiv(L_KV, num_splits) + + # Allocate partial result buffers + O_partial = torch.empty( + (num_splits, B, H_Q, D), device=q_rot.device, dtype=torch.float32 + ) + M_partial = torch.full( + (num_splits, B, H_Q), -float("inf"), device=q_rot.device, dtype=torch.float32 + ) + L_partial = torch.zeros( + (num_splits, B, H_Q), device=q_rot.device, dtype=torch.float32 + ) + + stride_op_s, stride_op_b, stride_op_h, stride_op_d = O_partial.stride() + stride_mp_s, stride_mp_b, stride_mp_h = M_partial.stride() + stride_lp_s, stride_lp_b, stride_lp_h = L_partial.stride() + + grid = (num_splits, B * H_grid) + wrap_triton(_tq4_sdpa_decode_splitk_kernel)[grid]( q_rot, k_packed, - k_n, + k_norms, v_packed, - v_n, + v_norms, lut_hi, lut_lo, - mask_ptr, - out_rot, - kv_len_t, + O_partial, + M_partial, + L_partial, + mask_ptr if HAS_MASK else 0, + kv_len_ptr if HAS_KV_LEN else 0, B, - H_Q, - H_KV, - N_Q, - N_KV, - D, - sm_scale, - HAS_MASK, - HAS_KV_LEN, - MASK_IS_CAUSAL, + H_grid, + L_KV, + *q_rot.stride(), + *k_packed.stride(), + *k_norms.stride(), + *v_packed.stride(), + *v_norms.stride(), + stride_op_s, + stride_op_b, + stride_op_h, + stride_op_d, + stride_mp_s, + stride_mp_b, + stride_mp_h, + stride_lp_s, + stride_lp_b, + stride_lp_h, stride_mb, stride_mq, stride_mk, - is_causal, - num_groups, - pack_gqa, + sm_scale, + chunk_size, + HAS_MASK=HAS_MASK, + HAS_KV_LEN=HAS_KV_LEN, + HEAD_DIM=D, + HALF_D=HALF_D, + NUM_GROUPS=num_groups, + PACK_GQA=pack_gqa, + BLOCK_M=BLOCK_M, ) - # Post-rotate: convert from rotated space back to original space - return torch.matmul(out_rot, rotation.to(query.dtype)) + # Reduce partials + grid_reduce = (B * H_Q,) + wrap_triton(_tq4_sdpa_decode_reduce_kernel)[grid_reduce]( + O_partial, + M_partial, + L_partial, + out_rot, + num_splits, + stride_op_s, + stride_op_b, + stride_op_h, + stride_op_d, + stride_mp_s, + stride_mp_b, + stride_mp_h, + stride_lp_s, + stride_lp_b, + stride_lp_h, + *out_rot.stride(), + HEAD_DIM=D, + num_warps=4, + num_stages=1, + ) @tq4_sdpa.register_fake diff --git a/examples/models/gemma4_31b/cuda_source_transformations.py b/examples/models/gemma4_31b/cuda_source_transformations.py index c16498d02b7..666d0c44e9d 100644 --- a/examples/models/gemma4_31b/cuda_source_transformations.py +++ b/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)