executorch 1.6.0__cp315-cp315-win_amd64.whl

executorch/backends/cuda/triton/kernels/sdpa.py

@@ -0,0 +1,2496 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+#
+# The GQA "pack GQA" optimization is adapted from FlashAttention
+# (Tri Dao, 2023-2025):
+#   https://github.com/Dao-AILab/flash-attention
+#   flash_attn/cute/pack_gqa.py — PackGQA class
+#   hopper/heuristics.h — should_pack_gqa() tile-utilization heuristic
+# Licensed under BSD-3-Clause.
+#
+# Pack GQA folds multiple Q heads that share the same KV head into the M
+# (sequence) dimension of a single tile, so K/V are loaded once per KV head
+# instead of once per Q head. The power-of-2 kernel autotunes whether to pack
+# along with its tile; the non-power-of-2 kernel keeps the tile-utilization
+# heuristic, which packs for short seqlen_q (e.g. decode) and not for long
+# seqlen_q (e.g. prefill).
+
+"""
+Triton SDPA Kernel for ExecuTorch CUDA Backend.
+
+This module provides a Triton-optimized implementation of scaled dot-product attention
+that can replace the default ATen/Edge SDPA operator during graph transformation to allow
+us export the model without decomposing the SDPA operator under libtorch free environment
+and have better performance.
+
+GQA support: when enable_gqa=True and H_q > H_kv, the kernel uses "pack GQA"
+(adapted from FlashAttention) to fold multiple Q heads sharing the same KV head
+into the M (sequence) dimension of a single tile. This avoids redundant K/V reads
+and improves tile utilization, especially during decode (seqlen_q=1).
+"""
+
+import functools
+import math
+from typing import Optional
+
+import torch
+import triton
+import triton.language as tl
+from executorch.backends.cuda.autotune.inputs import autotune_inputs
+from torch.library import triton_op, wrap_triton
+
+
+def _kv_len_scenarios(args) -> list[int]:
+    """KV lengths to autotune the kernels that read KV_LEN_ptr at: a short, a
+    middle and a full cache, never fewer than the queries. Autotuning would
+    otherwise see a zero-filled length and time empty sweeps."""
+    lk = int(args["Lk"])
+    lq = int(args.get("Lq", args.get("LQ", 1)))
+    return sorted({min(lk, max(lq, length)) for length in (1024, lk // 8, lk)})
+
+
+def _visible_mask_scenarios(args) -> list[bool]:
+    """An all-visible mask for each of ``_kv_len_scenarios``, paired by position.
+
+    The kernels skip KV blocks no row of a tile can see, so the all-False bool
+    mask autotuning would otherwise see times no attention at all and picks
+    whichever tile loads the mask fastest. All-visible times the full sweep
+    up to each KV length."""
+    return [True] * len(_kv_len_scenarios(args))
+
+
+def _is_power_of_2(n: int) -> bool:
+    """Check if n is a power of 2."""
+    return n > 0 and (n & (n - 1)) == 0
+
+
+# KV length at/above which split-K attention is used instead of the standard
+# kernel. The replacement pass still routes only L_q == 1 directly to
+# triton.sdpa_decode_splitk.
+_SPLITK_LKV_THRESHOLD = 256
+
+
+_TMA_PREFILL_LKV_THRESHOLD = 16384
+
+
+def _cuda_compile_target_is_sm90_or_newer() -> bool:
+    """Return whether the CUDA code being compiled targets SM90 or newer."""
+    if torch.version.hip is not None:
+        return False
+
+    # Match the target selected by AOTInductor rather than assuming that the
+    # export host is also the inference target. Multi-architecture PTE builds
+    # compile each native variant independently, so this value corresponds to
+    # the variant currently being lowered.
+    from torch._inductor.codegen.cuda.compile_utils import _nvcc_arch_as_compile_option
+
+    target = _nvcc_arch_as_compile_option().removesuffix("a")
+    return target.isdigit() and int(target) >= 90
+
+
+def _tma_prefill_config(
+    head_dim: int, query_len: int
+) -> Optional[tuple[int, int, int, int]]:
+    """Return a validated TMA config for common transformer head dimensions."""
+    if head_dim == 64:
+        return 64, 64, 3, 4
+    if head_dim == 128:
+        if query_len >= 1024:
+            return 128, 64, 3, 8
+        if query_len >= 512:
+            return 64, 64, 3, 4
+    return None
+
+
+# Decode split-K occupancy target. A sweep across both production attention
+# families showed that targeting 16/9 waves gives a good balance between split
+# kernel occupancy and reduction/empty-CTA overhead. Keep the ratio integral so
+# the launch policy is deterministic and does not depend on floating-point
+# rounding.
+_SPLITK_TARGET_WAVES_NUMERATOR = 16
+_SPLITK_TARGET_WAVES_DENOMINATOR = 9
+_SPLITK_MAX_SPLITS = 128
+
+# Do not create more static split CTAs than one per 128 elements in the KV
+# buffer. This permits 16 splits for the 2048-element sliding-window family.
+# A 40-state, >L2 working-set A/B showed gains at the production prompt lengths;
+# single-state measurements are cache-hot and are not representative here.
+# For the 128K global family with grid_y=4 on a 170-SM device, the occupancy
+# target selects 76 splits instead of reaching the merge-base buffer cap of 128.
+_SPLITK_BUFFER_ELEMENTS_PER_SPLIT = 128
+
+
+@functools.lru_cache(maxsize=None)
+def _device_sm_count(device_index: int) -> int:
+    """SM (multiprocessor) count of a CUDA device.
+
+    Read once per device from torch device properties and bake it into the
+    exported artifact before cuda-graph capture. The supported deployment flow
+    exports a native artifact per target GPU and may merge those variants into
+    one PTE; moving one native artifact to a same-architecture GPU with a
+    different SM count remains correct, but may not retain optimal scheduling.
+    """
+    return torch.cuda.get_device_properties(device_index).multi_processor_count
+
+
+def _decode_splitk_config(
+    L_kv: int, grid_y: int, device: torch.device
+) -> tuple[int, int]:
+    """Hardware-derived split count for the flash-decoding decode path.
+
+    The split kernel launches a (num_splits, grid_y) grid where
+    grid_y = B * H_kv. We size num_splits so the total launched CTAs
+    (num_splits * grid_y) target 16/9 waves of this device's SM array, capped by
+    cdiv(L_kv, 128) so a short KV buffer is not needlessly over-split and
+    clamped to >= 1.
+
+    num_splits depends only on host-side constants (L_kv buffer size, the
+    static grid_y, and the device SM count), so it is fixed before cuda-graph
+    capture and identical on every replay.
+    """
+    device_index = (
+        device.index if device.index is not None else torch.cuda.current_device()
+    )
+    n_sm = _device_sm_count(device_index)
+    # Cap so num_splits * grid_y ~= 16/9 * n_sm (>= 1). Use ceil division: a
+    # partial final wave is preferable to leaving an SM idle.
+    sm_cap = max(
+        triton.cdiv(
+            _SPLITK_TARGET_WAVES_NUMERATOR * n_sm,
+            _SPLITK_TARGET_WAVES_DENOMINATOR * max(grid_y, 1),
+        ),
+        1,
+    )
+    buffer_cap = max(triton.cdiv(L_kv, _SPLITK_BUFFER_ELEMENTS_PER_SPLIT), 1)
+    # Avoid letting very high-SM GPUs over-split long global-attention buffers:
+    # beyond 128 partitions the extra reduction and launch work outweighed the
+    # occupancy gain in CUDA-graph measurements. Lower-SM devices continue to
+    # use their hardware-derived cap.
+    num_splits = min(buffer_cap, sm_cap, _SPLITK_MAX_SPLITS)
+    return num_splits, triton.cdiv(L_kv, num_splits)
+
+
+def _next_power_of_2(x: int) -> int:
+    """Get the next power of 2 >= x, clamped to [16, 256].
+
+    Used for HEAD_DIM tiling where tile sizes below 16 waste warps
+    and head dims above 256 are unsupported.
+    """
+    if x <= 16:
+        return 16
+    if x <= 32:
+        return 32
+    if x <= 64:
+        return 64
+    if x <= 128:
+        return 128
+    return 256
+
+
+def _next_power_of_2_unclamped(x: int) -> int:
+    """Get the next power of 2 >= x (no clamping).
+
+    Used for GQA group-count tiling where num_groups can be small (1, 2, ...)
+    and should not be inflated to 16.
+    """
+    if x <= 0:
+        return 1
+    return 1 << (x - 1).bit_length()
+
+
+def _should_pack_gqa(L_q: int, num_groups: int, block_m: int) -> bool:
+    """Decide whether to use pack GQA based on tile utilization.
+
+    Pack GQA folds multiple Q heads into the M dimension so they share
+    the same K/V loads. This helps when seqlen_q is small relative to
+    BLOCK_M (e.g., decode with seqlen_q=1).
+
+    Heuristic from FlashAttention (hopper/heuristics.h, should_pack_gqa):
+    compare tile utilization with and without packing; pack if it
+    improves efficiency by >10%.
+
+    Reference: https://github.com/Dao-AILab/flash-attention/blob/main/hopper/heuristics.h
+    """
+    if num_groups <= 1:
+        return False
+
+    def round_up(a, b):
+        return ((a + b - 1) // b) * b
+
+    nopack_eff = L_q / round_up(L_q, block_m)
+    pack_eff = (L_q * num_groups) / round_up(L_q * num_groups, block_m)
+    return nopack_eff < 0.9 * pack_eff
+
+
+def _validate_qkv_shapes(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    enable_gqa: bool = False,
+) -> tuple[int, int, int, int, int, int, int]:
+    """
+    Validate dimensions and return shape info.
+    Args:
+        query: Query tensor [B, H_q, L_q, D]
+        key: Key tensor [B, H_kv, L_kv, D]
+        value: Value tensor [B, H_kv, L_kv, D]
+        enable_gqa: If True, H_q must be a multiple of H_kv (GQA/MQA).
+    Returns:
+        Tuple of (B, H_q, H_kv, L_q, L_kv, D_q, D_kv)
+    Raises:
+        RuntimeError: If dimensions are incompatible
+    """
+    B_q, H_q, L_q, D_q = query.shape
+    B_k, H_k, L_kv_k, D_k = key.shape
+    B_v, H_v, L_kv_v, D_v = value.shape
+    # Validate batch dimensions
+    if not (B_q == B_k == B_v):
+        raise RuntimeError(
+            f"Batch dimension must match; got B_q={B_q}, B_k={B_k}, B_v={B_v}."
+        )
+    # Validate head dimensions
+    if not (H_k == H_v):
+        raise RuntimeError(f"K and V head counts must match; got H_k={H_k}, H_v={H_v}.")
+    if enable_gqa:
+        if H_q % H_k != 0:
+            raise RuntimeError(
+                f"GQA requires H_q divisible by H_kv; got H_q={H_q}, H_kv={H_k}."
+            )
+    else:
+        if not (H_q == H_k):
+            raise RuntimeError(
+                f"Head counts must match (or use enable_gqa=True); "
+                f"got H_q={H_q}, H_k={H_k}."
+            )
+    # Head dimension must match
+    if not (D_q == D_k == D_v):
+        raise RuntimeError(
+            f"Head dimension must match across Q, K, V; got D_q={D_q}, D_k={D_k}, D_v={D_v}."
+        )
+    # Heads that are not a power of 2 must fit in one tile, and
+    # _next_power_of_2 caps the tile at 256.
+    if not _is_power_of_2(D_q) and _next_power_of_2(D_q) < D_q:
+        raise RuntimeError(
+            "Head dimensions that are not a power of 2 must be at most "
+            f"{_next_power_of_2(D_q)}; got D={D_q}."
+        )
+    # Key and Value sequence lengths must match
+    if L_kv_k != L_kv_v:
+        raise RuntimeError(
+            f"Key and Value must have the same sequence length; got L_k={L_kv_k}, L_v={L_kv_v}."
+        )
+    return B_q, H_q, H_k, L_q, L_kv_k, D_q, D_k
+
+
+# ==============================================================================
+# Non-power-of-2 HEAD_DIM kernel
+# ==============================================================================
+@triton.jit
+def _sdpa_fwd_kernel_non_pow2(
+    q_ptr,
+    k_ptr,
+    v_ptr,
+    o_ptr,
+    mask_ptr,
+    kv_len_ptr,
+    B,
+    H_grid,
+    LQ,
+    LK,
+    HEAD_DIM,
+    stride_qb,
+    stride_qh,
+    stride_ql,
+    stride_qd,
+    stride_kb,
+    stride_kh,
+    stride_kl,
+    stride_kd,
+    stride_vb,
+    stride_vh,
+    stride_vl,
+    stride_vd,
+    stride_ob,
+    stride_oh,
+    stride_ol,
+    stride_od,
+    stride_mb,
+    stride_mh,
+    stride_mlq,
+    stride_mlk,
+    scale,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_D: tl.constexpr,
+    HAS_MASK: tl.constexpr,
+    IS_CAUSAL: tl.constexpr,
+    HAS_KV_LEN: tl.constexpr,
+    NUM_GROUPS: tl.constexpr,
+    PACK_GQA: tl.constexpr,
+    MASK_IS_CAUSAL: tl.constexpr = False,
+):
+    """
+    SDPA forward kernel for non-power-of-2 HEAD_DIM.
+    Uses dynamic masking to handle arbitrary head dimensions.
+
+    PACK_GQA: when True, multiple Q heads sharing the same KV head are
+    folded into the M dimension. The grid iterates over H_kv heads and
+    each tile processes up to BLOCK_M rows from the packed (head, seq)
+    space. K/V are loaded once per KV head.
+    """
+    pid_m = tl.program_id(axis=0)
+    pid_bh = tl.program_id(axis=1)
+
+    b = pid_bh // H_grid
+    h_grid = pid_bh % H_grid
+
+    offs_packed = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    offs_d = tl.arange(0, BLOCK_D)
+    d_mask = offs_d < HEAD_DIM
+
+    if PACK_GQA:
+        seq_pos = offs_packed // NUM_GROUPS
+        h_within = offs_packed % NUM_GROUPS
+        h_q_rows = h_grid * NUM_GROUPS + h_within
+        h_kv = h_grid
+        row_valid = seq_pos < LQ
+        q_ptrs = (
+            q_ptr
+            + b * stride_qb
+            + h_q_rows[:, None] * stride_qh
+            + seq_pos[:, None] * stride_ql
+            + offs_d[None, :] * stride_qd
+        )
+    else:
+        seq_pos = offs_packed
+        h_kv = h_grid // NUM_GROUPS
+        row_valid = offs_packed < LQ
+        q_ptrs = (
+            q_ptr
+            + b * stride_qb
+            + h_grid * stride_qh
+            + offs_packed[:, None] * stride_ql
+            + offs_d[None, :] * stride_qd
+        )
+
+    q = tl.load(q_ptrs, mask=row_valid[:, None] & d_mask[None, :], other=0.0)
+
+    k_base = k_ptr + b * stride_kb + h_kv * stride_kh
+    v_base = v_ptr + b * stride_vb + h_kv * stride_vh
+
+    acc = tl.zeros((BLOCK_M, BLOCK_D), dtype=tl.float32)
+    m_i = tl.full((BLOCK_M,), -float("inf"), dtype=tl.float32)
+    l_i = tl.full((BLOCK_M,), 1.0, dtype=tl.float32)
+
+    qk_scale_log2 = scale * 1.4426950408889634
+
+    if HAS_MASK:
+        mask_b_base = mask_ptr + b * stride_mb
+
+    NEG_INF: tl.constexpr = float("-inf")
+
+    # Bound the KV loop to valid (filled) positions; see pow2 body for details.
+    if HAS_KV_LEN:
+        kv_len = tl.load(kv_len_ptr)
+    else:
+        kv_len = LK
+
+    for start_n in tl.range(0, kv_len, BLOCK_N, num_stages=2):
+        offs_n = start_n + tl.arange(0, BLOCK_N)
+        kv_col_mask = offs_n < kv_len
+
+        k_ptrs = k_base + (offs_n[:, None] * stride_kl + offs_d[None, :] * stride_kd)
+        k = tl.load(k_ptrs, mask=kv_col_mask[:, None] & d_mask[None, :], other=0.0)
+
+        qk = tl.dot(q, tl.trans(k))
+        qk = (qk * qk_scale_log2).to(tl.float32)
+
+        if IS_CAUSAL:
+            causal_mask = offs_n[None, :] > seq_pos[:, None]
+            qk = tl.where(causal_mask, tl.full(qk.shape, NEG_INF, dtype=tl.float32), qk)
+
+        if MASK_IS_CAUSAL:
+            causal_mask = offs_n[None, :] > (kv_len - LQ) + seq_pos[:, None]
+            qk = tl.where(causal_mask, tl.full(qk.shape, NEG_INF, dtype=tl.float32), qk)
+
+        if HAS_MASK:
+            m_ptrs = (
+                mask_b_base
+                + seq_pos[:, None] * stride_mlq
+                + offs_n[None, :] * stride_mlk
+            )
+            tile_valid = row_valid[:, None] & kv_col_mask[None, :]
+            keep = tl.load(m_ptrs, mask=tile_valid, other=False)
+            qk = tl.where(keep, qk, tl.full(qk.shape, NEG_INF, dtype=tl.float32))
+
+        qk = tl.where(
+            kv_col_mask[None, :], qk, tl.full(qk.shape, NEG_INF, dtype=tl.float32)
+        )
+
+        m_ij = tl.maximum(m_i, tl.max(qk, 1).to(tl.float32))
+        safe_diff = tl.where(
+            m_ij[:, None] > -float("inf"), qk - m_ij[:, None], -float("inf")
+        )
+        p = tl.math.exp2(safe_diff).to(tl.float32)
+        l_ij = tl.sum(p, 1).to(tl.float32)
+        safe_alpha_diff = tl.where(m_ij > -float("inf"), m_i - m_ij, 0.0)
+        alpha = tl.math.exp2(safe_alpha_diff).to(tl.float32)
+
+        acc = (acc * alpha[:, None]).to(tl.float32)
+
+        v_ptrs = v_base + (offs_n[:, None] * stride_vl + offs_d[None, :] * stride_vd)
+        v = tl.load(v_ptrs, mask=kv_col_mask[:, None] & d_mask[None, :], other=0.0)
+
+        acc = tl.dot(p.to(v.dtype), v, acc).to(tl.float32)
+
+        l_i = (l_i * alpha + l_ij).to(tl.float32)
+        m_i = m_ij
+
+    out = acc / l_i[:, None]
+
+    if PACK_GQA:
+        o_ptrs = (
+            o_ptr
+            + b * stride_ob
+            + h_q_rows[:, None] * stride_oh
+            + seq_pos[:, None] * stride_ol
+            + offs_d[None, :] * stride_od
+        )
+    else:
+        o_ptrs = (
+            o_ptr
+            + b * stride_ob
+            + h_grid * stride_oh
+            + offs_packed[:, None] * stride_ol
+            + offs_d[None, :] * stride_od
+        )
+    tl.store(o_ptrs, out.to(tl.bfloat16), mask=row_valid[:, None] & d_mask[None, :])
+
+
+# ==============================================================================
+# Power-of-2 HEAD_DIM kernels
+# ==============================================================================
+@triton.jit
+def _sdpa_fwd_kernel_body(  # noqa: C901
+    Q_ptr,
+    K_ptr,
+    V_ptr,
+    O_ptr,
+    Mask_ptr,
+    KV_LEN_ptr,
+    B,
+    H_grid,
+    Lq,
+    Lk,
+    stride_qb,
+    stride_qh,
+    stride_qm,
+    stride_qd,
+    stride_kb,
+    stride_kh,
+    stride_kn,
+    stride_kd,
+    stride_vb,
+    stride_vh,
+    stride_vn,
+    stride_vd,
+    stride_ob,
+    stride_oh,
+    stride_om,
+    stride_od,
+    stride_mb,
+    stride_mq,
+    stride_mk,
+    sm_scale: tl.float32,
+    HAS_MASK: tl.constexpr,
+    IS_CAUSAL: tl.constexpr,
+    HAS_KV_LEN: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    HEAD_DIM: tl.constexpr,
+    NUM_GROUPS: tl.constexpr,
+    PACK_GQA: tl.constexpr,
+    MASK_IS_CAUSAL: tl.constexpr = False,
+    KV_LEN_MULTIPLE: tl.constexpr = 1,
+):
+    """
+    Shared kernel body for SDPA forward pass.
+
+    PACK_GQA: when True, multiple Q heads sharing the same KV head are
+    folded into the M dimension (adapted from FlashAttention's pack_gqa).
+    The grid iterates over H_kv heads; each tile processes rows from the
+    packed (head, seq) space. K/V are loaded once per KV head, eliminating
+    redundant HBM reads across Q heads in a group.
+
+    When False, the grid iterates over H_q heads and each program handles
+    one Q head with simple h_kv = h_q // NUM_GROUPS remapping.
+
+    KV_LEN_MULTIPLE: kv_len is a multiple of this. When BLOCK_N divides it,
+    no tile reaches past kv_len, so the kernel skips the key bound check.
+    """
+    pid_m = tl.program_id(axis=0)
+    pid_bh = tl.program_id(axis=1)
+    b = pid_bh // H_grid
+    h_grid = pid_bh % H_grid
+
+    offs_packed = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    offs_d = tl.arange(0, HEAD_DIM)
+
+    if PACK_GQA:
+        # Decompose packed index: heads interleaved with positions
+        # [h0_pos0, h1_pos0, ..., h(G-1)_pos0, h0_pos1, h1_pos1, ...]
+        seq_pos = offs_packed // NUM_GROUPS
+        h_within = offs_packed % NUM_GROUPS
+        h_q_rows = h_grid * NUM_GROUPS + h_within  # [BLOCK_M] vector
+        h_kv = h_grid
+        row_valid = seq_pos < Lq
+
+        # Scattered Q load: each row may be a different Q head
+        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_packed
+        h_kv = h_grid // NUM_GROUPS
+        row_valid = offs_packed < Lq
+
+        # Uniform Q load: all rows are the same Q head
+        q_ptrs = Q_ptr + (
+            b * stride_qb
+            + h_grid * stride_qh
+            + offs_packed[:, None] * stride_qm
+            + offs_d[None, :] * stride_qd
+        )
+
+    q_mask = row_valid[:, None] & (offs_d[None, :] < HEAD_DIM)
+    q = tl.load(q_ptrs, mask=q_mask, other=0.0).to(tl.bfloat16)
+
+    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)
+
+    offs_n_init = tl.arange(0, BLOCK_N)
+
+    # Bound the KV loop to the number of valid (filled) positions instead of the
+    # full pre-allocated buffer Lk. For decode this is input_pos+1; for a prefill
+    # chunk it is chunk_end. This makes full-attention (global) layers 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) it
+    # falls back to Lk, preserving the original behavior exactly.
+    if HAS_KV_LEN:
+        kv_len = tl.load(KV_LEN_ptr)
+    else:
+        kv_len = Lk
+
+    # Window-aware early-exit. A KV block that is fully masked (sliding-window
+    # or causal) contributes nothing to the online softmax — every entry is
+    # -inf, so p=0 and m_i/l_i/acc are left unchanged. We detect such blocks up
+    # front and skip their K/V loads and both matmuls. This is exact: it only
+    # skips work the mask would have zeroed out anyway. At seq=2048 the 50
+    # sliding-window(1024) layers and the 10 causal layers each leave roughly
+    # half (or more) of their KV blocks fully masked, so this is a large cut to
+    # the dominant prefill cost. The skip condition is a CTA-wide reduction, so
+    # the branch is uniform and turns into a real skip (not predication).
+    if IS_CAUSAL:
+        max_seq_pos = tl.max(seq_pos)
+    if MASK_IS_CAUSAL:
+        max_kv_pos = (kv_len - Lq) + tl.max(seq_pos)
+
+    for start_n in tl.range(0, kv_len, BLOCK_N):
+        offs_n = start_n + offs_n_init
+
+        # Decide whether any row in this tile actually attends to this KV block.
+        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] & (offs_n[None, :] < kv_len)
+            mask_block = tl.load(mask_ptrs, mask=mn_mask, other=False)
+            block_active = tl.sum(mask_block.to(tl.int32)) > 0
+        elif IS_CAUSAL:
+            # Block is entirely in the future for every row -> skip.
+            block_active = start_n <= max_seq_pos
+        elif MASK_IS_CAUSAL:
+            block_active = start_n <= max_kv_pos
+        else:
+            block_active = True
+
+        if block_active:
+            # K load: uniform (single KV head, shared across Q heads in tile)
+            k_ptrs = K_ptr + (
+                b * stride_kb
+                + h_kv * stride_kh
+                + (offs_n[:, None] * stride_kn)
+                + (offs_d[None, :] * stride_kd)
+            )
+            k_mask = (offs_n[:, None] < kv_len) & (offs_d[None, :] < HEAD_DIM)
+            k = tl.load(k_ptrs, mask=k_mask, other=0.0).to(tl.bfloat16)
+
+            qk = (tl.dot(q, tl.trans(k)).to(tl.float32) * sm_scale).to(tl.float32)
+            # Keys past kv_len load as zeros, which would score 0 rather than -inf.
+            # A mask or a causal bound already drops them.
+            if (
+                not (HAS_MASK or IS_CAUSAL or MASK_IS_CAUSAL)
+                and KV_LEN_MULTIPLE % BLOCK_N != 0
+            ):
+                qk = tl.where(
+                    offs_n[None, :] < kv_len,
+                    qk,
+                    tl.full(qk.shape, -float("inf"), dtype=tl.float32),
+                )
+
+            if HAS_MASK:
+                qk = tl.where(
+                    mask_block, qk, tl.full(qk.shape, -float("inf"), dtype=tl.float32)
+                )
+
+            if IS_CAUSAL:
+                causal = offs_n[None, :] > seq_pos[:, None]
+                qk = tl.where(
+                    causal, tl.full(qk.shape, -float("inf"), dtype=tl.float32), qk
+                )
+
+            if MASK_IS_CAUSAL:
+                causal = offs_n[None, :] > (kv_len - Lq) + seq_pos[:, None]
+                qk = tl.where(
+                    causal, tl.full(qk.shape, -float("inf"), dtype=tl.float32), qk
+                )
+
+            m_ij = tl.maximum(m_i, tl.max(qk, axis=1).to(tl.float32))
+            safe_diff = tl.where(
+                m_ij[:, None] > -float("inf"), qk - m_ij[:, None], -float("inf")
+            )
+            p_f32 = tl.exp(safe_diff).to(tl.float32)
+            l_ij = tl.sum(p_f32, axis=1).to(tl.float32)
+            safe_alpha_diff = tl.where(m_ij > -float("inf"), m_i - m_ij, 0.0)
+            alpha = tl.exp(safe_alpha_diff).to(tl.float32)
+
+            # V load: uniform (single KV head)
+            v_ptrs = V_ptr + (
+                b * stride_vb
+                + h_kv * stride_vh
+                + (offs_n[:, None] * stride_vn)
+                + (offs_d[None, :] * stride_vd)
+            )
+            v_mask = (offs_n[:, None] < kv_len) & (offs_d[None, :] < HEAD_DIM)
+            v = tl.load(v_ptrs, mask=v_mask, other=0.0).to(tl.bfloat16)
+
+            p_bf16 = p_f32.to(tl.bfloat16)
+            acc = (acc * alpha[:, None] + tl.dot(p_bf16, v)).to(tl.float32)
+            l_i = (l_i * alpha + l_ij).to(tl.float32)
+            m_i = m_ij
+
+    inv_l_i = tl.where(l_i > 0, 1.0 / l_i, 0.0)
+    acc = acc * inv_l_i[:, None]
+
+    # O store: scattered when PACK_GQA, uniform otherwise
+    if PACK_GQA:
+        o_ptrs = O_ptr + (
+            b * stride_ob
+            + h_q_rows[:, None] * stride_oh
+            + seq_pos[:, None] * stride_om
+            + offs_d[None, :] * stride_od
+        )
+    else:
+        o_ptrs = O_ptr + (
+            b * stride_ob
+            + h_grid * stride_oh
+            + offs_packed[:, None] * stride_om
+            + offs_d[None, :] * stride_od
+        )
+    o_mask = row_valid[:, None] & (offs_d[None, :] < HEAD_DIM)
+    tl.store(o_ptrs, acc.to(tl.bfloat16), mask=o_mask)
+
+
+# Prefill / standard-path tile configs. ONE autotuned kernel spanning BLOCK_M in
+# {16..128}; `_sdpa_prefill_prune` drops configs whose fp32 accumulator
+# acc[BLOCK_M, HEAD_DIM] would spill registers for the runtime HEAD_DIM, so the
+# kernel is high-occupancy AND HEAD_DIM-agnostic (64/80/96/128/256/512). This
+# replaces the old fixed BLOCK_M=64 (m64) / BLOCK_M=32 (m32) wrappers + Python
+# CTA-count selector: at HEAD_DIM=512 the m64 path spilled acc[64,512] fp32
+# (128 KB/CTA -> ~280 reg spills -> ~30 TFLOP/s); the autotuner now picks a
+# non-spilling, well-pipelined tile per HEAD_DIM (e.g. BLOCK_M=32 at 512).
+_SDPA_PREFILL_CONFIGS = [
+    triton.Config({"BLOCK_M": 16, "BLOCK_N": 32}, num_warps=4, num_stages=2),
+    triton.Config({"BLOCK_M": 16, "BLOCK_N": 64}, num_warps=4, num_stages=3),
+    triton.Config({"BLOCK_M": 32, "BLOCK_N": 32}, num_warps=4, num_stages=2),
+    triton.Config({"BLOCK_M": 32, "BLOCK_N": 32}, num_warps=4, num_stages=3),
+    triton.Config({"BLOCK_M": 32, "BLOCK_N": 64}, num_warps=4, num_stages=2),
+    triton.Config({"BLOCK_M": 32, "BLOCK_N": 64}, num_warps=8, num_stages=3),
+    triton.Config({"BLOCK_M": 64, "BLOCK_N": 32}, num_warps=8, num_stages=2),
+    triton.Config({"BLOCK_M": 64, "BLOCK_N": 64}, num_warps=4, num_stages=2),
+    triton.Config({"BLOCK_M": 64, "BLOCK_N": 64}, num_warps=8, num_stages=3),
+    triton.Config({"BLOCK_M": 64, "BLOCK_N": 128}, num_warps=8, num_stages=3),
+    triton.Config({"BLOCK_M": 128, "BLOCK_N": 64}, num_warps=8, num_stages=3),
+    triton.Config({"BLOCK_M": 128, "BLOCK_N": 128}, num_warps=8, num_stages=3),
+]
+
+
+def _sdpa_prefill_prune(configs, nargs, **kwargs):
+    """Drop configs whose fp32 acc[BLOCK_M, HEAD_DIM] would spill registers.
+
+    Keeps ``BLOCK_M * HEAD_DIM <= 4096 * num_warps`` (the measured A100 no-spill
+    boundary: HEAD_DIM=512 -> BLOCK_M<=32 at 4 warps / <=64 at 8 warps;
+    HEAD_DIM=128 -> BLOCK_M<=128 at 4 warps). This guarantees a high-occupancy
+    pick for any HEAD_DIM and a non-empty result (the BLOCK_M=16 configs satisfy
+    the budget for every HEAD_DIM<=1024). SMEM-OOR tiles (large
+    BLOCK_N*HEAD_DIM*num_stages) are pruned by the autotuner at benchmark time.
+    """
+    head_dim = kwargs.get("HEAD_DIM")
+    if head_dim is None and nargs is not None:
+        head_dim = nargs.get("HEAD_DIM")
+    if head_dim is None:
+        return configs
+    kept = [c for c in configs if c.kwargs["BLOCK_M"] * head_dim <= 4096 * c.num_warps]
+    if not kept:
+        kept = [min(configs, key=lambda c: c.kwargs["BLOCK_M"] / c.num_warps)]
+    return kept
+
+
+# Pack GQA is tuned with the tile: which wins depends on the query count, the
+# group size, the tile and how many CTAs the device runs at once, which a
+# utilization rule cannot weigh.
+_SDPA_FWD_CONFIGS = [
+    triton.Config(
+        {**config.kwargs, "PACK_GQA": pack},
+        num_warps=config.num_warps,
+        num_stages=config.num_stages,
+    )
+    for config in _SDPA_PREFILL_CONFIGS
+    for pack in (False, True)
+]
+_SDPA_BLOCK_N_SIZES = {config.kwargs["BLOCK_N"] for config in _SDPA_FWD_CONFIGS}
+
+
+def _sdpa_fwd_prune(configs, nargs, **kwargs):
+    """``_sdpa_prefill_prune``, and without GQA only the unpacked tiles."""
+    num_groups = kwargs.get("NUM_GROUPS")
+    if num_groups is None and nargs is not None:
+        num_groups = nargs.get("NUM_GROUPS")
+    if num_groups is not None and num_groups <= 1:
+        configs = [c for c in configs if not c.kwargs["PACK_GQA"]]
+    return _sdpa_prefill_prune(configs, nargs, **kwargs)
+
+
+@autotune_inputs(KV_LEN_ptr=_kv_len_scenarios, Mask_ptr=_visible_mask_scenarios)
+@triton.autotune(
+    configs=_SDPA_FWD_CONFIGS,
+    key=[
+        "Lq",
+        "Lk",
+        "HEAD_DIM",
+        "HAS_MASK",
+        "IS_CAUSAL",
+        "MASK_IS_CAUSAL",
+        "NUM_GROUPS",
+    ],
+    prune_configs_by={"early_config_prune": _sdpa_fwd_prune},
+)
+@triton.jit
+def _sdpa_fwd_kernel(
+    Q_ptr,
+    K_ptr,
+    V_ptr,
+    O_ptr,
+    Mask_ptr,
+    KV_LEN_ptr,
+    B,
+    H_q,
+    Lq,
+    Lk,
+    stride_qb,
+    stride_qh,
+    stride_qm,
+    stride_qd,
+    stride_kb,
+    stride_kh,
+    stride_kn,
+    stride_kd,
+    stride_vb,
+    stride_vh,
+    stride_vn,
+    stride_vd,
+    stride_ob,
+    stride_oh,
+    stride_om,
+    stride_od,
+    stride_mb,
+    stride_mq,
+    stride_mk,
+    sm_scale: tl.float32,
+    HAS_MASK: tl.constexpr,
+    IS_CAUSAL: tl.constexpr,
+    HAS_KV_LEN: tl.constexpr,
+    HEAD_DIM: tl.constexpr,
+    NUM_GROUPS: tl.constexpr,
+    PACK_GQA: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    MASK_IS_CAUSAL: tl.constexpr = False,
+    KV_LEN_MULTIPLE: tl.constexpr = 1,
+):
+    # Packed, the grid walks the KV heads; unpacked, the query heads.
+    if PACK_GQA:
+        H_grid = H_q // NUM_GROUPS
+    else:
+        H_grid = H_q
+    _sdpa_fwd_kernel_body(
+        Q_ptr,
+        K_ptr,
+        V_ptr,
+        O_ptr,
+        Mask_ptr,
+        KV_LEN_ptr,
+        B,
+        H_grid,
+        Lq,
+        Lk,
+        stride_qb,
+        stride_qh,
+        stride_qm,
+        stride_qd,
+        stride_kb,
+        stride_kh,
+        stride_kn,
+        stride_kd,
+        stride_vb,
+        stride_vh,
+        stride_vn,
+        stride_vd,
+        stride_ob,
+        stride_oh,
+        stride_om,
+        stride_od,
+        stride_mb,
+        stride_mq,
+        stride_mk,
+        sm_scale,
+        HAS_MASK=HAS_MASK,
+        IS_CAUSAL=IS_CAUSAL,
+        HAS_KV_LEN=HAS_KV_LEN,
+        BLOCK_M=BLOCK_M,
+        BLOCK_N=BLOCK_N,
+        HEAD_DIM=HEAD_DIM,
+        NUM_GROUPS=NUM_GROUPS,
+        PACK_GQA=PACK_GQA,
+        MASK_IS_CAUSAL=MASK_IS_CAUSAL,
+        KV_LEN_MULTIPLE=KV_LEN_MULTIPLE,
+    )
+
+
+@triton.jit
+def _sdpa_prefill_tma_kernel(
+    Q_ptr,
+    K_ptr,
+    V_ptr,
+    O_ptr,
+    KV_LEN_ptr,
+    H_grid,
+    Lq,
+    Lk,
+    stride_qb,
+    stride_qh,
+    stride_qm,
+    stride_kb,
+    stride_kh,
+    stride_kn,
+    stride_vb,
+    stride_vh,
+    stride_vn,
+    stride_ob,
+    stride_oh,
+    stride_om,
+    stride_od,
+    sm_scale,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    HEAD_DIM: tl.constexpr,
+    NUM_GROUPS: tl.constexpr,
+    NUM_STAGES: tl.constexpr,
+):
+    """TMA global-causal prefill for SM90+ and a device-resident KV bound."""
+    pid_m = tl.program_id(0)
+    pid_bh = tl.program_id(1)
+    b = pid_bh // H_grid
+    h_q = pid_bh % H_grid
+    h_kv = h_q // NUM_GROUPS
+
+    q_desc = tl.make_tensor_descriptor(
+        Q_ptr + b * stride_qb + h_q * stride_qh,
+        shape=[Lq, HEAD_DIM],
+        strides=[stride_qm, 1],
+        block_shape=[BLOCK_M, HEAD_DIM],
+    )
+    k_desc = tl.make_tensor_descriptor(
+        K_ptr + b * stride_kb + h_kv * stride_kh,
+        shape=[Lk, HEAD_DIM],
+        strides=[stride_kn, 1],
+        block_shape=[BLOCK_N, HEAD_DIM],
+    )
+    v_desc = tl.make_tensor_descriptor(
+        V_ptr + b * stride_vb + h_kv * stride_vh,
+        shape=[Lk, HEAD_DIM],
+        strides=[stride_vn, 1],
+        block_shape=[BLOCK_N, HEAD_DIM],
+    )
+
+    q_start = pid_m * BLOCK_M
+    offs_m = q_start + tl.arange(0, BLOCK_M)
+    offs_n_base = tl.arange(0, BLOCK_N)
+    offs_d = tl.arange(0, HEAD_DIM)
+    q = tl.load_tensor_descriptor(q_desc, [q_start, 0])
+    kv_len = tl.minimum(tl.load(KV_LEN_ptr), Lk)
+    absolute_q = (kv_len - Lq) + offs_m
+    m_i = tl.full([BLOCK_M], -float("inf"), tl.float32)
+    l_i = tl.zeros([BLOCK_M], tl.float32)
+    acc = tl.zeros([BLOCK_M, HEAD_DIM], tl.float32)
+    scale = sm_scale.to(tl.float32)
+
+    for start_n in tl.range(0, kv_len, BLOCK_N, num_stages=NUM_STAGES):
+        offs_n = start_n + offs_n_base
+        k = tl.load_tensor_descriptor(k_desc, [start_n, 0])
+        qk = tl.dot(q, tl.trans(k)).to(tl.float32) * scale
+        valid = (offs_n[None, :] < kv_len) & (offs_n[None, :] <= absolute_q[:, None])
+        qk = tl.where(valid, qk, -float("inf"))
+        m_ij = tl.maximum(m_i, tl.max(qk, axis=1))
+        safe_m = tl.where(m_ij == -float("inf"), 0.0, m_ij)
+        alpha = tl.exp(m_i - safe_m)
+        p = tl.exp(qk - safe_m[:, None])
+        l_i = l_i * alpha + tl.sum(p, axis=1)
+        acc = acc * alpha[:, None]
+        v = tl.load_tensor_descriptor(v_desc, [start_n, 0])
+        acc = tl.dot(p.to(tl.bfloat16), v, acc).to(tl.float32)
+        m_i = m_ij
+
+    inv_l = tl.where(l_i > 0, 1.0 / l_i, 0.0)
+    out = acc * inv_l[:, None]
+    o_ptrs = (
+        O_ptr
+        + b * stride_ob
+        + h_q * stride_oh
+        + offs_m[:, None] * stride_om
+        + offs_d[None, :] * stride_od
+    )
+    tl.store(o_ptrs, out.to(tl.bfloat16), mask=offs_m[:, None] < Lq)
+
+
+def _validate_sdpa_inputs(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    dropout_p: float,
+    enable_gqa: bool,
+) -> None:
+    """Validate SDPA input tensors and unsupported feature flags."""
+    if not (query.is_cuda and key.is_cuda and value.is_cuda):
+        raise RuntimeError("Q, K, V must be CUDA tensors.")
+    if (
+        query.dtype != torch.bfloat16
+        or key.dtype != torch.bfloat16
+        or value.dtype != torch.bfloat16
+    ):
+        raise RuntimeError("Expected bfloat16 inputs")
+    if query.dim() != 4 or key.dim() != 4 or value.dim() != 4:
+        raise RuntimeError(
+            f"Expected 4D tensors shaped [B, H, L, D]; got "
+            f"query.dim()={query.dim()}, key.dim()={key.dim()}, "
+            f"value.dim()={value.dim()}."
+        )
+    if dropout_p != 0.0:
+        raise RuntimeError(
+            "dropout_p must be 0.0 (not supported in this implementation)."
+        )
+
+
+def _prepare_mask_params(
+    attn_mask: Optional[torch.Tensor],
+    B: int,
+    L_q: int,
+    L_kv: int,
+) -> tuple[bool, torch.Tensor, int, int, int]:
+    """Prepare attention mask parameters for kernel invocation."""
+    if attn_mask is None:
+        return False, 0, 0, 0, 0
+
+    if attn_mask.dtype != torch.bool:
+        raise RuntimeError("attn_mask must have dtype torch.bool")
+    if not attn_mask.is_cuda:
+        raise RuntimeError("attn_mask must be a CUDA tensor")
+    if attn_mask.shape[1] != 1:
+        raise RuntimeError(
+            f"attn_mask head dimension must be 1 (broadcast over heads); "
+            f"per-head masks are not supported. Got attn_mask.shape={attn_mask.shape}"
+        )
+    if (
+        attn_mask.shape[0] != B
+        or attn_mask.shape[2] != L_q
+        or attn_mask.shape[3] != L_kv
+    ):
+        raise RuntimeError(
+            f"attn_mask shape mismatch: expected [B={B}, 1, L_q={L_q}, L_kv={L_kv}], "
+            f"got {attn_mask.shape}"
+        )
+    return (
+        True,
+        attn_mask,
+        attn_mask.stride(0),
+        attn_mask.stride(2),
+        attn_mask.stride(3),
+    )
+
+
+def _launch_pow2_kernel(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    out: torch.Tensor,
+    B: int,
+    H_q: int,
+    H_kv: int,
+    L_q: int,
+    L_kv: int,
+    D: int,
+    sm_scale: float,
+    HAS_MASK: bool,
+    Mask_ptr: torch.Tensor,
+    stride_mb: int,
+    stride_mq: int,
+    stride_mk: int,
+    is_causal: bool,
+    num_groups: int,
+    kv_len_ptr: Optional[torch.Tensor] = None,
+    HAS_KV_LEN: bool = False,
+    mask_is_causal: bool = False,
+) -> None:
+    """Launch power-of-2 optimized SDPA kernel."""
+    stride_qb, stride_qh, stride_qm, stride_qd = query.stride()
+    stride_kb, stride_kh, stride_kn, stride_kd = key.stride()
+    stride_vb, stride_vh, stride_vn, stride_vd = value.stride()
+    stride_ob, stride_oh, stride_om, stride_od = out.stride()
+
+    # Hopper/Blackwell: TMA materially improves the long-context global
+    # causal path. Sliding-window attention keeps the existing kernel, as do
+    # pre-SM90 GPUs. The same TMA kernel is safe for short chunks and avoids a
+    # host sync on the device-resident kv_len scalar.
+    tma_config = _tma_prefill_config(D, L_q)
+    if (
+        HAS_KV_LEN
+        and mask_is_causal
+        and not HAS_MASK
+        and query.stride(-1) == 1
+        and key.stride(-1) == 1
+        and value.stride(-1) == 1
+        and tma_config is not None
+        and L_kv >= _TMA_PREFILL_LKV_THRESHOLD
+        and L_q > 4
+        and _cuda_compile_target_is_sm90_or_newer()
+    ):
+        block_m, block_n, tma_num_stages, tma_num_warps = tma_config
+        wrap_triton(_sdpa_prefill_tma_kernel)[(triton.cdiv(L_q, block_m), B * H_q)](
+            query,
+            key,
+            value,
+            out,
+            kv_len_ptr,
+            H_q,
+            L_q,
+            L_kv,
+            stride_qb,
+            stride_qh,
+            stride_qm,
+            stride_kb,
+            stride_kh,
+            stride_kn,
+            stride_vb,
+            stride_vh,
+            stride_vn,
+            stride_ob,
+            stride_oh,
+            stride_om,
+            stride_od,
+            sm_scale,
+            BLOCK_M=block_m,
+            BLOCK_N=block_n,
+            HEAD_DIM=D,
+            NUM_GROUPS=num_groups,
+            NUM_STAGES=tma_num_stages,
+            num_warps=tma_num_warps,
+            num_stages=tma_num_stages,
+        )
+        return
+
+    def grid(meta):
+        if meta["PACK_GQA"]:
+            return (triton.cdiv(L_q * num_groups, meta["BLOCK_M"]), B * H_kv)
+        return (triton.cdiv(L_q, meta["BLOCK_M"]), B * H_q)
+
+    # The largest tile size that divides L_kv. A kv_len or a symbolic length is
+    # only known at run time, so those always check bounds. A mask or a causal
+    # bound makes the check unneeded, so those calls keep 1 and share kernels.
+    kv_len_multiple = 1
+    if isinstance(L_kv, int) and not (HAS_KV_LEN or HAS_MASK or is_causal):
+        kv_len_multiple = max(
+            (n for n in _SDPA_BLOCK_N_SIZES if L_kv % n == 0), default=1
+        )
+
+    # Single autotuned kernel: the config set spans BLOCK_M in {16..128}, each
+    # with and without pack GQA, and `_sdpa_prefill_prune` keeps only
+    # non-spilling tiles for this HEAD_DIM, so the autotuner picks a
+    # high-occupancy tile (small BLOCK_M for large HEAD_DIM, larger BLOCK_M /
+    # more CTAs for small problems).
+    wrap_triton(_sdpa_fwd_kernel)[grid](
+        query,
+        key,
+        value,
+        out,
+        Mask_ptr if HAS_MASK else 0,
+        kv_len_ptr if HAS_KV_LEN else 0,
+        B,
+        H_q,
+        L_q,
+        L_kv,
+        stride_qb,
+        stride_qh,
+        stride_qm,
+        stride_qd,
+        stride_kb,
+        stride_kh,
+        stride_kn,
+        stride_kd,
+        stride_vb,
+        stride_vh,
+        stride_vn,
+        stride_vd,
+        stride_ob,
+        stride_oh,
+        stride_om,
+        stride_od,
+        stride_mb,
+        stride_mq,
+        stride_mk,
+        sm_scale,
+        HAS_MASK=HAS_MASK,
+        IS_CAUSAL=is_causal,
+        HAS_KV_LEN=HAS_KV_LEN,
+        HEAD_DIM=D,
+        NUM_GROUPS=num_groups,
+        MASK_IS_CAUSAL=mask_is_causal,
+        KV_LEN_MULTIPLE=kv_len_multiple,
+    )
+
+
+def _launch_non_pow2_kernel(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    out: torch.Tensor,
+    attn_mask: Optional[torch.Tensor],
+    B: int,
+    H_q: int,
+    H_kv: int,
+    L_q: int,
+    L_kv: int,
+    D: int,
+    sm_scale: float,
+    HAS_MASK: bool,
+    is_causal: bool,
+    num_groups: int,
+    pack_gqa: bool,
+    kv_len_ptr: Optional[torch.Tensor] = None,
+    HAS_KV_LEN: bool = False,
+    mask_is_causal: bool = False,
+) -> None:
+    """Launch non-power-of-2 SDPA kernel with dynamic HEAD_DIM masking."""
+    stride_qb, stride_qh, stride_qm, stride_qd = query.stride()
+    stride_kb, stride_kh, stride_kn, stride_kd = key.stride()
+    stride_vb, stride_vh, stride_vn, stride_vd = value.stride()
+    stride_ob, stride_oh, stride_om, stride_od = out.stride()
+
+    BLOCK_D = _next_power_of_2(D)
+    BLOCK_N = 64 if BLOCK_D >= 256 else 128
+    BLOCK_M = 32
+    num_warps = 4
+    num_stages = 2
+
+    if pack_gqa:
+        H_grid = H_kv
+        Lq_packed = L_q * num_groups
+    else:
+        H_grid = H_q
+        Lq_packed = L_q
+
+    if HAS_MASK:
+        mask_ptr = attn_mask
+        stride_mb_np2 = attn_mask.stride(0)
+        stride_mh_np2 = attn_mask.stride(1)
+        stride_mlq_np2 = attn_mask.stride(2)
+        stride_mlk_np2 = attn_mask.stride(3)
+    else:
+        mask_ptr = 0
+        stride_mb_np2 = stride_mh_np2 = stride_mlq_np2 = stride_mlk_np2 = 0
+
+    def grid_non_pow2(meta):
+        return (triton.cdiv(Lq_packed, meta["BLOCK_M"]), B * H_grid)
+
+    wrap_triton(_sdpa_fwd_kernel_non_pow2)[grid_non_pow2](
+        query,
+        key,
+        value,
+        out,
+        mask_ptr,
+        kv_len_ptr if HAS_KV_LEN else 0,
+        B,
+        H_grid,
+        L_q,
+        L_kv,
+        D,
+        stride_qb,
+        stride_qh,
+        stride_qm,
+        stride_qd,
+        stride_kb,
+        stride_kh,
+        stride_kn,
+        stride_kd,
+        stride_vb,
+        stride_vh,
+        stride_vn,
+        stride_vd,
+        stride_ob,
+        stride_oh,
+        stride_om,
+        stride_od,
+        stride_mb_np2,
+        stride_mh_np2,
+        stride_mlq_np2,
+        stride_mlk_np2,
+        sm_scale,
+        BLOCK_M=BLOCK_M,
+        BLOCK_N=BLOCK_N,
+        BLOCK_D=BLOCK_D,
+        HAS_MASK=HAS_MASK,
+        IS_CAUSAL=is_causal,
+        HAS_KV_LEN=HAS_KV_LEN,
+        NUM_GROUPS=num_groups,
+        PACK_GQA=pack_gqa,
+        MASK_IS_CAUSAL=mask_is_causal,
+        num_warps=num_warps,
+        num_stages=num_stages,
+    )
+
+
+@triton_op("triton::sdpa", mutates_args={})
+def sdpa(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    attn_mask: Optional[torch.Tensor] = None,
+    dropout_p: float = 0.0,
+    is_causal: bool = False,
+    scale: float = 0.0,
+    enable_gqa: bool = False,
+    kv_len: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """
+    Triton fused Scaled Dot-Product Attention with GQA pack optimization.
+
+    When enable_gqa=True and H_q > H_kv, this kernel automatically decides
+    whether to use "pack GQA" (folding Q heads into the M dimension so they
+    share K/V loads) based on a tile-utilization heuristic from FlashAttention.
+
+    Args:
+        query: Query tensor [B, H_q, L_q, D], dtype torch.bfloat16. A D that
+            is not a power of 2 must be at most 256.
+        key: Key tensor [B, H_kv, L_kv, D], dtype torch.bfloat16
+        value: Value tensor [B, H_kv, L_kv, D], dtype torch.bfloat16
+        attn_mask: Optional bool mask [B, 1, L_q, L_kv] (broadcast over heads)
+        dropout_p: must be 0.0
+        is_causal: apply causal masking. With ``kv_len``, non-square inputs use
+            bottom-right causal alignment without materializing a dense mask.
+        scale: attention scale (default: 1/sqrt(D))
+        enable_gqa: allow H_q != H_kv (GQA/MQA)
+        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 attention O(context) instead
+            of O(max_seq_len). It is read on-device (no host sync) so the bound
+            updates correctly under CUDA-graph replay (decode). For decode pass
+            ``input_pos + 1``; for a prefill chunk pass ``chunk_end``. When None
+            the loop runs over the full ``L_kv`` (original behavior). Supplying
+            it for an L_q==1 decode with a large buffer also routes through the
+            split-K flash-decoding kernel for occupancy.
+    Returns:
+        Output tensor [B, H_q, L_q, D], dtype torch.bfloat16
+    """
+    _validate_sdpa_inputs(query, key, value, dropout_p, enable_gqa)
+
+    B, H_q, H_kv, L_q, L_kv, D_q, _ = _validate_qkv_shapes(
+        query, key, value, enable_gqa
+    )
+    D = D_q
+    num_groups = H_q // H_kv
+
+    if is_causal and L_q != L_kv and kv_len is None:
+        raise RuntimeError(
+            f"Causal masking requires L_q == L_kv; got L_q={L_q}, L_kv={L_kv}. "
+            "For non-square causal attention, pass kv_len."
+        )
+
+    out = torch.empty((B, H_q, L_q, D), device=query.device, dtype=query.dtype)
+    sm_scale = 1.0 / math.sqrt(D) if scale == 0.0 else scale
+    HAS_MASK, Mask_ptr, stride_mb, stride_mq, stride_mk = _prepare_mask_params(
+        attn_mask, B, L_q, L_kv
+    )
+
+    # Optional length bound: device int32 scalar, clamped to the buffer size for
+    # OOB safety. Reshaped to [1] so the kernel can ``tl.load`` element 0. No
+    # ``.item()`` — keeps it CUDA-graph-safe (value updates on replay).
+    HAS_KV_LEN = kv_len is not None
+    if HAS_KV_LEN:
+        kv_len_t = torch.clamp(
+            kv_len.reshape(1).to(torch.int32), max=int(L_kv)
+        ).contiguous()
+    else:
+        kv_len_t = None
+
+    # A device-resident KV bound provides the absolute query offset needed for
+    # bottom-right causal masking. Keep an explicit mask, if present: the
+    # kernels compose it with the causal bound rather than silently dropping it.
+    mask_is_causal = is_causal and HAS_KV_LEN
+    kernel_is_causal = is_causal and not mask_is_causal
+
+    # Split-K dispatch with a kv_len bound, power-of-2 head dimension, and a
+    # large KV buffer. The legacy decode launcher remains isolated at L_q == 1;
+    # the generalized launcher handles only verifier-sized query blocks.
+    if HAS_KV_LEN and _is_power_of_2(D) and L_kv >= _SPLITK_LKV_THRESHOLD:
+        if L_q == 1:
+            _launch_decode_splitk(
+                query,
+                key,
+                value,
+                out,
+                B,
+                H_q,
+                H_kv,
+                L_kv,
+                D,
+                sm_scale,
+                HAS_MASK,
+                Mask_ptr,
+                stride_mb,
+                stride_mq,
+                stride_mk,
+                num_groups,
+                kv_len_t,
+                HAS_KV_LEN,
+            )
+            return out
+        if 2 <= L_q <= 4:
+            _launch_small_query_splitk(
+                query,
+                key,
+                value,
+                out,
+                B,
+                H_q,
+                H_kv,
+                L_q,
+                L_kv,
+                D,
+                sm_scale,
+                HAS_MASK,
+                Mask_ptr,
+                stride_mb,
+                stride_mq,
+                stride_mk,
+                num_groups,
+                kv_len_t,
+                HAS_KV_LEN,
+                mask_is_causal,
+            )
+            return out
+
+    if _is_power_of_2(D):
+        _launch_pow2_kernel(
+            query,
+            key,
+            value,
+            out,
+            B,
+            H_q,
+            H_kv,
+            L_q,
+            L_kv,
+            D,
+            sm_scale,
+            HAS_MASK,
+            Mask_ptr,
+            stride_mb,
+            stride_mq,
+            stride_mk,
+            kernel_is_causal,
+            num_groups,
+            kv_len_t,
+            HAS_KV_LEN,
+            mask_is_causal,
+        )
+    else:
+        # The non-pow2 kernel always uses BLOCK_M=32 and is not tuned over pack
+        # GQA, so a tile-utilization rule decides it.
+        pack_gqa = _should_pack_gqa(L_q, num_groups, 32)
+        _launch_non_pow2_kernel(
+            query,
+            key,
+            value,
+            out,
+            attn_mask,
+            B,
+            H_q,
+            H_kv,
+            L_q,
+            L_kv,
+            D,
+            sm_scale,
+            HAS_MASK,
+            kernel_is_causal,
+            num_groups,
+            pack_gqa,
+            kv_len_t,
+            HAS_KV_LEN,
+            mask_is_causal,
+        )
+
+    return out
+
+
+# Register the abstract/fake implementation for torch.export
+# This is critical to avoid accessing real tensor data during export
+@sdpa.register_fake
+def _sdpa_abstract(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    attn_mask: Optional[torch.Tensor] = None,
+    dropout_p: float = 0.0,
+    is_causal: bool = False,
+    scale: float = 0.0,
+    enable_gqa: bool = False,
+    kv_len: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """
+    Abstract/fake implementation for torch.export.
+    This just returns an empty tensor with the correct shape/dtype/device.
+    """
+    # Validate dtypes match
+    assert query.dtype == key.dtype == value.dtype, "Q, K, V must have the same dtype"
+    # Validate kqv's shape and get the output shape
+    B, H_q, _H_kv, L_q, _, D_q, _ = _validate_qkv_shapes(query, key, value, enable_gqa)
+
+    return torch.empty(B, H_q, L_q, D_q, dtype=query.dtype, device=query.device)
+
+
+# ==============================================================================
+# Split-K decode kernel (flash-decoding)
+# ==============================================================================
+# When L_q == 1 with GQA, the standard kernel launches only
+# ceil(num_groups / BLOCK_M) * B * H_kv CTAs (e.g. 2 for Qwen3.5 MoE).
+# Split-K partitions the KV sequence across many CTAs for better occupancy,
+# then reduces partial results in a second kernel.
+
+
+@autotune_inputs(KV_LEN_ptr=_kv_len_scenarios, Mask_ptr=_visible_mask_scenarios)
+@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"],
+)
+@triton.jit
+def _sdpa_decode_splitk_kernel(
+    Q_ptr,
+    K_ptr,
+    V_ptr,
+    O_partial_ptr,
+    M_partial_ptr,
+    L_partial_ptr,
+    Mask_ptr,
+    KV_LEN_ptr,
+    B,
+    H_kv,
+    Lk,
+    stride_qb,
+    stride_qh,
+    stride_qm,
+    stride_qd,
+    stride_kb,
+    stride_kh,
+    stride_kn,
+    stride_kd,
+    stride_vb,
+    stride_vh,
+    stride_vn,
+    stride_vd,
+    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,
+    NUM_GROUPS: tl.constexpr,
+    BLOCK_G: tl.constexpr,
+):
+    sm_scale = sm_scale.to(tl.float32)
+    split_id = tl.program_id(axis=0)
+    pid_bh = tl.program_id(axis=1)
+    b = pid_bh // H_kv
+    h_kv = pid_bh % H_kv
+
+    start_n = split_id * chunk_size
+    # Bound the decode KV sweep to the valid (filled) positions. Empty splits
+    # store m = -inf and l = acc = 0, which the reduce weights to zero. kv_len is
+    # read on-device (CUDA-graph safe); falls back to Lk when not provided.
+    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_g = tl.arange(0, BLOCK_G)
+    g_valid = offs_g < NUM_GROUPS
+    h_q_heads = h_kv * NUM_GROUPS + offs_g  # [BLOCK_G]
+
+    # Load Q for all heads in this group: [BLOCK_G, HEAD_DIM]
+    q_ptrs = Q_ptr + (
+        b * stride_qb
+        + h_q_heads[:, None] * stride_qh
+        + 0 * stride_qm
+        + offs_d[None, :] * stride_qd
+    )
+    q = tl.load(q_ptrs, mask=g_valid[:, None], other=0.0).to(tl.bfloat16)
+
+    m_i = tl.full([BLOCK_G], -float("inf"), dtype=tl.float32)
+    l_i = tl.zeros([BLOCK_G], dtype=tl.float32)
+    acc = tl.zeros([BLOCK_G, HEAD_DIM], dtype=tl.float32)
+
+    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
+        n_valid = offs_n < end_n
+
+        k_ptrs = K_ptr + (
+            b * stride_kb
+            + h_kv * stride_kh
+            + offs_n[:, None] * stride_kn
+            + offs_d[None, :] * stride_kd
+        )
+        k = tl.load(k_ptrs, mask=n_valid[:, None], other=0.0).to(tl.bfloat16)
+
+        # QK: [BLOCK_G, BLOCK_N]
+        qk = (tl.dot(q, tl.trans(k)).to(tl.float32) * sm_scale).to(tl.float32)
+
+        # Mask out-of-bounds KV positions
+        qk = tl.where(
+            n_valid[None, :],
+            qk,
+            tl.full(qk.shape, -float("inf"), dtype=tl.float32),
+        )
+
+        if HAS_MASK:
+            mask_ptrs = Mask_ptr + (
+                b * stride_mb + 0 * stride_mq + offs_n[None, :] * stride_mk
+            )
+            mask_block = tl.load(mask_ptrs, mask=n_valid[None, :], other=False)
+            qk = tl.where(
+                mask_block, qk, tl.full(qk.shape, -float("inf"), dtype=tl.float32)
+            )
+
+        m_ij = tl.maximum(m_i, tl.max(qk, axis=1).to(tl.float32))
+        safe_diff = tl.where(
+            m_ij[:, None] > -float("inf"), qk - m_ij[:, None], -float("inf")
+        )
+        p_f32 = tl.exp(safe_diff).to(tl.float32)
+        l_ij = tl.sum(p_f32, axis=1).to(tl.float32)
+        safe_alpha_diff = tl.where(m_ij > -float("inf"), m_i - m_ij, 0.0)
+        alpha = tl.exp(safe_alpha_diff).to(tl.float32)
+
+        v_ptrs = V_ptr + (
+            b * stride_vb
+            + h_kv * stride_vh
+            + offs_n[:, None] * stride_vn
+            + offs_d[None, :] * stride_vd
+        )
+        v = tl.load(v_ptrs, mask=n_valid[:, None], other=0.0).to(tl.bfloat16)
+
+        p_bf16 = p_f32.to(tl.bfloat16)
+        acc = (acc * alpha[:, None] + tl.dot(p_bf16, v)).to(tl.float32)
+        l_i = (l_i * alpha + l_ij).to(tl.float32)
+        m_i = m_ij
+
+    # Padded group lanes can alias later heads, so all three stores need g_valid.
+    h_q_all = h_kv * NUM_GROUPS + offs_g  # [BLOCK_G]
+    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
+    )
+    tl.store(o_ptrs, acc, mask=g_valid[:, None])
+
+    m_ptrs = M_partial_ptr + (
+        split_id * stride_mp_s + b * stride_mp_b + h_q_all * stride_mp_h
+    )
+    tl.store(m_ptrs, m_i, mask=g_valid)
+
+    ll_ptrs = L_partial_ptr + (
+        split_id * stride_lp_s + b * stride_lp_b + h_q_all * stride_lp_h
+    )
+    tl.store(ll_ptrs, l_i, mask=g_valid)
+
+
+@triton.jit
+def _sdpa_decode_reduce_kernel(
+    O_partial_ptr,
+    M_partial_ptr,
+    L_partial_ptr,
+    O_ptr,
+    num_splits,
+    stride_op_s,
+    stride_op_h,
+    stride_op_d,
+    stride_mp_s,
+    stride_mp_h,
+    stride_lp_s,
+    stride_lp_h,
+    stride_oh,
+    stride_od,
+    BLOCK_S: tl.constexpr,
+    HEAD_DIM: tl.constexpr,
+):
+    pid = tl.program_id(axis=0)
+    offs_s = tl.arange(0, BLOCK_S)
+    offs_d = tl.arange(0, HEAD_DIM)
+    s_valid = offs_s < num_splits
+
+    m_ptrs = M_partial_ptr + offs_s * stride_mp_s + pid * stride_mp_h
+    m = tl.load(m_ptrs, mask=s_valid, other=-float("inf"))
+    m_global = tl.max(m, axis=0)
+    alpha = tl.where(
+        s_valid & (m > -float("inf")) & (m_global > -float("inf")),
+        tl.exp(m - m_global),
+        0.0,
+    )
+
+    l_ptrs = L_partial_ptr + offs_s * stride_lp_s + pid * stride_lp_h
+    l_values = tl.load(l_ptrs, mask=s_valid, other=0.0)
+    o_ptrs = O_partial_ptr + (
+        offs_s[:, None] * stride_op_s
+        + pid * stride_op_h
+        + offs_d[None, :] * stride_op_d
+    )
+    o = tl.load(o_ptrs, mask=s_valid[:, None], other=0.0)
+
+    l_global = tl.sum(l_values * alpha, axis=0)
+    acc = tl.sum(o * alpha[:, None], axis=0)
+
+    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_decode_splitk(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    out: torch.Tensor,
+    B: int,
+    H_q: int,
+    H_kv: int,
+    L_kv: int,
+    D: int,
+    sm_scale: float,
+    HAS_MASK: bool,
+    Mask_ptr,
+    stride_mb: int,
+    stride_mq: int,
+    stride_mk: int,
+    num_groups: int,
+    kv_len_ptr: Optional[torch.Tensor] = None,
+    HAS_KV_LEN: bool = False,
+) -> None:
+    num_splits, chunk_size = _decode_splitk_config(L_kv, B * H_kv, query.device)
+
+    # Each valid (split, batch, query head) slot belongs to exactly one program,
+    # and empty KV ranges still reach the stores with m=-inf and l=acc=0.
+    O_partial = torch.empty(
+        (num_splits, B, H_q, D), device=query.device, dtype=torch.float32
+    )
+    M_partial = torch.empty(
+        (num_splits, B, H_q), device=query.device, dtype=torch.float32
+    )
+    L_partial = torch.empty(
+        (num_splits, B, H_q), device=query.device, dtype=torch.float32
+    )
+
+    stride_qb, stride_qh, stride_qm, stride_qd = query.stride()
+    stride_kb, stride_kh, stride_kn, stride_kd = key.stride()
+    stride_vb, stride_vh, stride_vn, stride_vd = value.stride()
+    _, stride_oh, _, stride_od = out.stride()
+    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_split = (num_splits, B * H_kv)
+    wrap_triton(_sdpa_decode_splitk_kernel)[grid_split](
+        query,
+        key,
+        value,
+        O_partial,
+        M_partial,
+        L_partial,
+        Mask_ptr if HAS_MASK else 0,
+        kv_len_ptr if HAS_KV_LEN else 0,
+        B,
+        H_kv,
+        L_kv,
+        stride_qb,
+        stride_qh,
+        stride_qm,
+        stride_qd,
+        stride_kb,
+        stride_kh,
+        stride_kn,
+        stride_kd,
+        stride_vb,
+        stride_vh,
+        stride_vn,
+        stride_vd,
+        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,
+        chunk_size,
+        HAS_MASK=HAS_MASK,
+        HAS_KV_LEN=HAS_KV_LEN,
+        HEAD_DIM=D,
+        NUM_GROUPS=num_groups,
+        BLOCK_G=_next_power_of_2_unclamped(num_groups),
+    )
+
+    grid_reduce = (B * H_q,)
+    wrap_triton(_sdpa_decode_reduce_kernel)[grid_reduce](
+        O_partial,
+        M_partial,
+        L_partial,
+        out,
+        num_splits,
+        stride_op_s,
+        stride_op_h,
+        stride_op_d,
+        stride_mp_s,
+        stride_mp_h,
+        stride_lp_s,
+        stride_lp_h,
+        stride_oh,
+        stride_od,
+        BLOCK_S=_next_power_of_2_unclamped(num_splits),
+        HEAD_DIM=D,
+        num_warps=4,
+        num_stages=1,
+    )
+
+
+@triton_op("triton::sdpa_decode_splitk", mutates_args={})
+def sdpa_decode_splitk(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    attn_mask: Optional[torch.Tensor] = None,
+    dropout_p: float = 0.0,
+    is_causal: bool = False,
+    scale: float = 0.0,
+    enable_gqa: bool = False,
+    phi: float = 5.0,
+    kv_len: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Split-K flash-decoding SDPA for L_q=1 (decode step).
+
+    Tracks a stable online-softmax maximum per split, then rescales partial
+    results by the global maximum during reduction.
+
+    Signature mirrors sdpa() for drop-in use with torch.cond dispatch.
+    enable_gqa is accepted but ignored — GQA is handled natively via
+    H_q // H_kv grouping; no packed-GQA tradeoff exists at L_q=1.
+
+    kv_len: optional GPU int scalar bounding the KV sweep to the valid
+    (filled) positions (O(context) instead of O(max_seq_len)). Read
+    on-device, CUDA-graph safe. When None, sweeps the full L_kv.
+
+    phi is deprecated, accepted for operator-schema compatibility, and ignored.
+    """
+    _validate_sdpa_inputs(query, key, value, dropout_p, enable_gqa)
+
+    B, H_q, L_q, D = query.shape
+    _, H_kv, L_kv, _ = key.shape
+
+    out = torch.empty((B, H_q, L_q, D), device=query.device, dtype=query.dtype)
+
+    # is_causal is a no-op at L_q=1 (single query can't attend to future
+    # positions), so we accept it silently for API compatibility with callers
+    # that always pass is_causal=True for decode.
+
+    # Validation — only check at runtime (concrete shapes), not during AOTI
+    # tracing where shapes are symbolic. torch.cond traces both branches with
+    # the same symbolic L_q, so L_q is not necessarily 1 during tracing.
+    if isinstance(L_q, int):
+        if L_q != 1:
+            raise RuntimeError(
+                f"sdpa_decode_splitk requires L_q == 1 (decode); got L_q={L_q}"
+            )
+        if H_q % H_kv != 0:
+            raise RuntimeError(
+                f"H_q must be divisible by H_kv; got H_q={H_q}, H_kv={H_kv}"
+            )
+        if not _is_power_of_2(D):
+            raise RuntimeError(
+                f"sdpa_decode_splitk requires power-of-2 head dim; got D={D}"
+            )
+
+    num_groups = H_q // H_kv
+    sm_scale = 1.0 / math.sqrt(D) if scale == 0.0 else scale
+    HAS_MASK, Mask_ptr, stride_mb, stride_mq, stride_mk = _prepare_mask_params(
+        attn_mask, B, L_q, L_kv
+    )
+
+    HAS_KV_LEN = kv_len is not None
+    if HAS_KV_LEN:
+        kv_len_t = torch.clamp(
+            kv_len.reshape(1).to(torch.int32), max=int(L_kv)
+        ).contiguous()
+    else:
+        kv_len_t = None
+
+    _launch_decode_splitk(
+        query,
+        key,
+        value,
+        out,
+        B,
+        H_q,
+        H_kv,
+        L_kv,
+        D,
+        sm_scale,
+        HAS_MASK,
+        Mask_ptr,
+        stride_mb,
+        stride_mq,
+        stride_mk,
+        num_groups,
+        kv_len_t,
+        HAS_KV_LEN,
+    )
+    return out
+
+
+@sdpa_decode_splitk.register_fake
+def _sdpa_decode_splitk_abstract(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    attn_mask: Optional[torch.Tensor] = None,
+    dropout_p: float = 0.0,
+    is_causal: bool = False,
+    scale: float = 0.0,
+    enable_gqa: bool = False,
+    phi: float = 5.0,
+    kv_len: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    assert query.dtype == key.dtype == value.dtype, "Q, K, V must have the same dtype"
+    B, H_q, L_q, D = query.shape
+    return torch.empty(B, H_q, L_q, D, dtype=query.dtype, device=query.device)
+
+
+# ==============================================================================
+# Split-K small-query kernel (flash-decoding)
+# ==============================================================================
+# With a small query block and GQA, the standard kernel launches only a few
+# CTAs. Split-K partitions the KV sequence across many CTAs for better
+# occupancy, then reduces per-query partial results in a second kernel.
+
+
+@autotune_inputs(KV_LEN_ptr=_kv_len_scenarios, Mask_ptr=_visible_mask_scenarios)
+@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", "LQ", "HAS_MASK", "MASK_IS_CAUSAL"],
+)
+@triton.jit
+def _sdpa_small_query_splitk_kernel(
+    Q_ptr,
+    K_ptr,
+    V_ptr,
+    O_partial_ptr,
+    M_partial_ptr,
+    L_partial_ptr,
+    Mask_ptr,
+    KV_LEN_ptr,
+    B,
+    H_kv,
+    Lk,
+    stride_qb,
+    stride_qh,
+    stride_qm,
+    stride_qd,
+    stride_kb,
+    stride_kh,
+    stride_kn,
+    stride_kd,
+    stride_vb,
+    stride_vh,
+    stride_vn,
+    stride_vd,
+    stride_op_s,
+    stride_op_b,
+    stride_op_h,
+    stride_op_m,
+    stride_op_d,
+    stride_mp_s,
+    stride_mp_b,
+    stride_mp_h,
+    stride_mp_m,
+    stride_lp_s,
+    stride_lp_b,
+    stride_lp_h,
+    stride_lp_m,
+    stride_mb,
+    stride_mq,
+    stride_mk,
+    sm_scale: tl.float32,
+    chunk_size,
+    HAS_MASK: tl.constexpr,
+    HAS_KV_LEN: tl.constexpr,
+    MASK_IS_CAUSAL: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    HEAD_DIM: tl.constexpr,
+    NUM_GROUPS: tl.constexpr,
+    LQ: tl.constexpr,
+    BLOCK_G: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+):
+    sm_scale = sm_scale.to(tl.float32)
+    split_id = tl.program_id(axis=0)
+    pid_bh = tl.program_id(axis=1)
+    b = pid_bh // H_kv
+    h_kv = pid_bh % H_kv
+
+    start_n = split_id * chunk_size
+    # Bound the decode KV sweep to the valid (filled) positions. Empty splits
+    # store m = -inf and l = acc = 0, which the reduce weights to zero. kv_len is
+    # read on-device (CUDA-graph safe); falls back to Lk when not provided.
+    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_m = tl.arange(0, BLOCK_M)
+    query_rows = offs_m // BLOCK_G
+    group_heads = offs_m % BLOCK_G
+    row_valid = (query_rows < LQ) & (group_heads < NUM_GROUPS)
+    h_q_heads = h_kv * NUM_GROUPS + group_heads
+
+    # Flatten query rows and GQA groups so each KV tile is shared by all
+    # small-query/head combinations handled by this CTA.
+    q_ptrs = Q_ptr + (
+        b * stride_qb
+        + h_q_heads[:, None] * stride_qh
+        + query_rows[:, None] * stride_qm
+        + offs_d[None, :] * stride_qd
+    )
+    q = tl.load(q_ptrs, mask=row_valid[:, None], other=0.0).to(tl.bfloat16)
+
+    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)
+
+    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
+        n_valid = offs_n < end_n
+
+        k_ptrs = K_ptr + (
+            b * stride_kb
+            + h_kv * stride_kh
+            + offs_n[:, None] * stride_kn
+            + offs_d[None, :] * stride_kd
+        )
+        k = tl.load(k_ptrs, mask=n_valid[:, None], other=0.0).to(tl.bfloat16)
+
+        # QK: [BLOCK_M, BLOCK_N]
+        qk = (tl.dot(q, tl.trans(k)).to(tl.float32) * sm_scale).to(tl.float32)
+
+        # Mask out-of-bounds KV positions
+        qk = tl.where(
+            n_valid[None, :],
+            qk,
+            tl.full(qk.shape, -float("inf"), dtype=tl.float32),
+        )
+
+        if HAS_MASK:
+            mask_ptrs = Mask_ptr + (
+                b * stride_mb
+                + query_rows[:, None] * stride_mq
+                + offs_n[None, :] * stride_mk
+            )
+            mask_block = tl.load(
+                mask_ptrs,
+                mask=row_valid[:, None] & n_valid[None, :],
+                other=False,
+            )
+            qk = tl.where(
+                mask_block, qk, tl.full(qk.shape, -float("inf"), dtype=tl.float32)
+            )
+
+        if MASK_IS_CAUSAL:
+            absolute_q = (kv_len - LQ) + query_rows
+            causal = offs_n[None, :] > absolute_q[:, None]
+            qk = tl.where(
+                causal, tl.full(qk.shape, -float("inf"), dtype=tl.float32), qk
+            )
+
+        m_ij = tl.maximum(m_i, tl.max(qk, axis=1).to(tl.float32))
+        safe_diff = tl.where(
+            m_ij[:, None] > -float("inf"), qk - m_ij[:, None], -float("inf")
+        )
+        p_f32 = tl.exp(safe_diff).to(tl.float32)
+        l_ij = tl.sum(p_f32, axis=1).to(tl.float32)
+        safe_alpha_diff = tl.where(m_ij > -float("inf"), m_i - m_ij, 0.0)
+        alpha = tl.exp(safe_alpha_diff).to(tl.float32)
+
+        v_ptrs = V_ptr + (
+            b * stride_vb
+            + h_kv * stride_vh
+            + offs_n[:, None] * stride_vn
+            + offs_d[None, :] * stride_vd
+        )
+        v = tl.load(v_ptrs, mask=n_valid[:, None], other=0.0).to(tl.bfloat16)
+
+        p_bf16 = p_f32.to(tl.bfloat16)
+        acc = (acc * alpha[:, None] + tl.dot(p_bf16, v)).to(tl.float32)
+        l_i = (l_i * alpha + l_ij).to(tl.float32)
+        m_i = m_ij
+
+    # Padded query/head rows can alias valid outputs, so all stores need row_valid.
+    o_ptrs = O_partial_ptr + (
+        split_id * stride_op_s
+        + b * stride_op_b
+        + h_q_heads[:, None] * stride_op_h
+        + query_rows[:, None] * stride_op_m
+        + offs_d[None, :] * stride_op_d
+    )
+    tl.store(o_ptrs, acc, mask=row_valid[:, None])
+
+    m_ptrs = M_partial_ptr + (
+        split_id * stride_mp_s
+        + b * stride_mp_b
+        + h_q_heads * stride_mp_h
+        + query_rows * stride_mp_m
+    )
+    tl.store(m_ptrs, m_i, mask=row_valid)
+
+    ll_ptrs = L_partial_ptr + (
+        split_id * stride_lp_s
+        + b * stride_lp_b
+        + h_q_heads * stride_lp_h
+        + query_rows * stride_lp_m
+    )
+    tl.store(ll_ptrs, l_i, mask=row_valid)
+
+
+@triton.jit
+def _sdpa_small_query_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_m,
+    stride_op_d,
+    stride_mp_s,
+    stride_mp_b,
+    stride_mp_h,
+    stride_mp_m,
+    stride_lp_s,
+    stride_lp_b,
+    stride_lp_h,
+    stride_lp_m,
+    stride_ob,
+    stride_oh,
+    stride_om,
+    stride_od,
+    H_Q,
+    BLOCK_S: tl.constexpr,
+    LQ: tl.constexpr,
+    HEAD_DIM: tl.constexpr,
+):
+    pid = tl.program_id(axis=0)
+    query_row = pid % LQ
+    h_q = (pid // LQ) % H_Q
+    b = pid // (LQ * H_Q)
+    offs_s = tl.arange(0, BLOCK_S)
+    offs_d = tl.arange(0, HEAD_DIM)
+    s_valid = offs_s < num_splits
+
+    m_ptrs = (
+        M_partial_ptr
+        + offs_s * stride_mp_s
+        + b * stride_mp_b
+        + h_q * stride_mp_h
+        + query_row * stride_mp_m
+    )
+    m = tl.load(m_ptrs, mask=s_valid, other=-float("inf"))
+    m_global = tl.max(m, axis=0)
+    alpha = tl.where(
+        s_valid & (m > -float("inf")) & (m_global > -float("inf")),
+        tl.exp(m - m_global),
+        0.0,
+    )
+
+    l_ptrs = (
+        L_partial_ptr
+        + offs_s * stride_lp_s
+        + b * stride_lp_b
+        + h_q * stride_lp_h
+        + query_row * stride_lp_m
+    )
+    l_values = tl.load(l_ptrs, mask=s_valid, other=0.0)
+    o_ptrs = O_partial_ptr + (
+        offs_s[:, None] * stride_op_s
+        + b * stride_op_b
+        + h_q * stride_op_h
+        + query_row * stride_op_m
+        + offs_d[None, :] * stride_op_d
+    )
+    o = tl.load(o_ptrs, mask=s_valid[:, None], other=0.0)
+
+    l_global = tl.sum(l_values * alpha, axis=0)
+    acc = tl.sum(o * alpha[:, None], axis=0)
+
+    inv_l = tl.where(l_global > 0, 1.0 / l_global, 0.0)
+    acc = acc * inv_l
+
+    o_out_ptrs = (
+        O_ptr
+        + b * stride_ob
+        + h_q * stride_oh
+        + query_row * stride_om
+        + offs_d * stride_od
+    )
+    tl.store(o_out_ptrs, acc.to(tl.bfloat16))
+
+
+def _launch_small_query_splitk(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    out: torch.Tensor,
+    B: int,
+    H_q: int,
+    H_kv: int,
+    L_q: int,
+    L_kv: int,
+    D: int,
+    sm_scale: float,
+    HAS_MASK: bool,
+    Mask_ptr,
+    stride_mb: int,
+    stride_mq: int,
+    stride_mk: int,
+    num_groups: int,
+    kv_len_ptr: Optional[torch.Tensor] = None,
+    HAS_KV_LEN: bool = False,
+    mask_is_causal: bool = False,
+) -> None:
+    num_splits, chunk_size = _decode_splitk_config(L_kv, B * H_kv, query.device)
+
+    # Each valid (split, batch, query head, query row) slot belongs to exactly
+    # one program, and empty KV ranges still store m=-inf and l=acc=0.
+    O_partial = torch.empty(
+        (num_splits, B, H_q, L_q, D), device=query.device, dtype=torch.float32
+    )
+    M_partial = torch.empty(
+        (num_splits, B, H_q, L_q), device=query.device, dtype=torch.float32
+    )
+    L_partial = torch.empty(
+        (num_splits, B, H_q, L_q), device=query.device, dtype=torch.float32
+    )
+
+    stride_qb, stride_qh, stride_qm, stride_qd = query.stride()
+    stride_kb, stride_kh, stride_kn, stride_kd = key.stride()
+    stride_vb, stride_vh, stride_vn, stride_vd = value.stride()
+    stride_ob, stride_oh, stride_om, stride_od = out.stride()
+    stride_op_s, stride_op_b, stride_op_h, stride_op_m, stride_op_d = O_partial.stride()
+    stride_mp_s, stride_mp_b, stride_mp_h, stride_mp_m = M_partial.stride()
+    stride_lp_s, stride_lp_b, stride_lp_h, stride_lp_m = L_partial.stride()
+
+    block_g = _next_power_of_2_unclamped(num_groups)
+    grid_split = (num_splits, B * H_kv)
+    wrap_triton(_sdpa_small_query_splitk_kernel)[grid_split](
+        query,
+        key,
+        value,
+        O_partial,
+        M_partial,
+        L_partial,
+        Mask_ptr if HAS_MASK else 0,
+        kv_len_ptr if HAS_KV_LEN else 0,
+        B,
+        H_kv,
+        L_kv,
+        stride_qb,
+        stride_qh,
+        stride_qm,
+        stride_qd,
+        stride_kb,
+        stride_kh,
+        stride_kn,
+        stride_kd,
+        stride_vb,
+        stride_vh,
+        stride_vn,
+        stride_vd,
+        stride_op_s,
+        stride_op_b,
+        stride_op_h,
+        stride_op_m,
+        stride_op_d,
+        stride_mp_s,
+        stride_mp_b,
+        stride_mp_h,
+        stride_mp_m,
+        stride_lp_s,
+        stride_lp_b,
+        stride_lp_h,
+        stride_lp_m,
+        stride_mb,
+        stride_mq,
+        stride_mk,
+        sm_scale,
+        chunk_size,
+        HAS_MASK=HAS_MASK,
+        HAS_KV_LEN=HAS_KV_LEN,
+        MASK_IS_CAUSAL=mask_is_causal,
+        HEAD_DIM=D,
+        NUM_GROUPS=num_groups,
+        LQ=L_q,
+        BLOCK_G=block_g,
+        BLOCK_M=_next_power_of_2_unclamped(L_q * block_g),
+    )
+
+    grid_reduce = (B * H_q * L_q,)
+    wrap_triton(_sdpa_small_query_reduce_kernel)[grid_reduce](
+        O_partial,
+        M_partial,
+        L_partial,
+        out,
+        num_splits,
+        stride_op_s,
+        stride_op_b,
+        stride_op_h,
+        stride_op_m,
+        stride_op_d,
+        stride_mp_s,
+        stride_mp_b,
+        stride_mp_h,
+        stride_mp_m,
+        stride_lp_s,
+        stride_lp_b,
+        stride_lp_h,
+        stride_lp_m,
+        stride_ob,
+        stride_oh,
+        stride_om,
+        stride_od,
+        H_q,
+        BLOCK_S=_next_power_of_2_unclamped(num_splits),
+        LQ=L_q,
+        HEAD_DIM=D,
+        num_warps=4,
+        num_stages=1,
+    )
+
+
+@triton_op("triton::sdpa_small_query_splitk", mutates_args={})
+def sdpa_small_query_splitk(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    attn_mask: Optional[torch.Tensor] = None,
+    dropout_p: float = 0.0,
+    is_causal: bool = False,
+    scale: float = 0.0,
+    enable_gqa: bool = False,
+    phi: float = 5.0,
+    kv_len: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Split-K flash-decoding SDPA for small query blocks (2 <= L_q <= 4).
+
+    Tracks a stable online-softmax maximum per split, then rescales partial
+    results by the global maximum during reduction.
+
+    Signature mirrors sdpa() for drop-in use with torch.cond dispatch.
+    enable_gqa is accepted but ignored — GQA is handled natively via
+    H_q // H_kv grouping; query rows and GQA groups share each KV tile.
+
+    kv_len: optional GPU int scalar bounding the KV sweep to the valid
+    (filled) positions (O(context) instead of O(max_seq_len)). Read
+    on-device, CUDA-graph safe. When None, sweeps the full L_kv.
+
+    phi is deprecated, accepted for operator-schema compatibility, and ignored.
+    """
+    _validate_sdpa_inputs(query, key, value, dropout_p, enable_gqa)
+
+    B, H_q, L_q, D = query.shape
+    _, H_kv, L_kv, _ = key.shape
+
+    out = torch.empty((B, H_q, L_q, D), device=query.device, dtype=query.dtype)
+
+    # Cached multi-query attention needs an explicit bottom-right-aligned mask.
+    if is_causal:
+        torch._check(
+            L_q == 1,
+            lambda: "sdpa_small_query_splitk requires an explicit mask when L_q > 1",
+        )
+
+    # Validation — only check at runtime (concrete shapes), not during AOTI
+    # tracing where shapes are symbolic. torch.cond traces both branches with
+    # the same symbolic L_q, so L_q is not necessarily concrete during tracing.
+    if isinstance(L_q, int):
+        if not 2 <= L_q <= 4:
+            raise RuntimeError(
+                f"sdpa_small_query_splitk requires 2 <= L_q <= 4; got L_q={L_q}"
+            )
+        if H_q % H_kv != 0:
+            raise RuntimeError(
+                f"H_q must be divisible by H_kv; got H_q={H_q}, H_kv={H_kv}"
+            )
+        if not _is_power_of_2(D):
+            raise RuntimeError(
+                f"sdpa_small_query_splitk requires power-of-2 head dim; got D={D}"
+            )
+
+    num_groups = H_q // H_kv
+    sm_scale = 1.0 / math.sqrt(D) if scale == 0.0 else scale
+    HAS_MASK, Mask_ptr, stride_mb, stride_mq, stride_mk = _prepare_mask_params(
+        attn_mask, B, L_q, L_kv
+    )
+
+    HAS_KV_LEN = kv_len is not None
+    if HAS_KV_LEN:
+        kv_len_t = torch.clamp(
+            kv_len.reshape(1).to(torch.int32), max=int(L_kv)
+        ).contiguous()
+    else:
+        kv_len_t = None
+
+    _launch_small_query_splitk(
+        query,
+        key,
+        value,
+        out,
+        B,
+        H_q,
+        H_kv,
+        L_q,
+        L_kv,
+        D,
+        sm_scale,
+        HAS_MASK,
+        Mask_ptr,
+        stride_mb,
+        stride_mq,
+        stride_mk,
+        num_groups,
+        kv_len_t,
+        HAS_KV_LEN,
+    )
+    return out
+
+
+@sdpa_small_query_splitk.register_fake
+def _sdpa_small_query_splitk_abstract(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    attn_mask: Optional[torch.Tensor] = None,
+    dropout_p: float = 0.0,
+    is_causal: bool = False,
+    scale: float = 0.0,
+    enable_gqa: bool = False,
+    phi: float = 5.0,
+    kv_len: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    assert query.dtype == key.dtype == value.dtype, "Q, K, V must have the same dtype"
+    B, H_q, L_q, D = query.shape
+    return torch.empty(B, H_q, L_q, D, dtype=query.dtype, device=query.device)