Update veylon_attention.py

veylon_attention.py

@@ -81,12 +81,23 @@ TPU-specific notes
 from __future__ import annotations
 
 import math
+import os
 from functools import partial
 from typing import Optional
 
 import jax
 import jax.numpy as jnp
 
+# ---------------------------------------------------------------------------
+# Optional Pallas/Triton GPU kernel (FlashAttention-style, I/O-aware)
+# ---------------------------------------------------------------------------
+try:
+    from jax.experimental import pallas as pl
+    from jax.experimental.pallas import triton as plgpu
+    _PALLAS_GPU_AVAILABLE = True
+except Exception:
+    _PALLAS_GPU_AVAILABLE = False
+
 
 # ---------------------------------------------------------------------------
 # Public tuning constants
@@ -119,6 +130,498 @@ def _detect_backend() -> str:
         return 'cpu'
 
 
+# ---------------------------------------------------------------------------
+# Pallas/Triton GPU kernel — I/O-aware FlashAttention-style GQA SWA
+# ---------------------------------------------------------------------------
+#
+# This is a from-scratch FlashAttention-2 style kernel:
+#   - Tiles Q into BLOCK_Q-sized blocks, K/V into BLOCK_K-sized blocks.
+#   - Grid = (batch, kv_head, num_q_blocks). Each program instance owns one
+#     Q block for one KV head (covering its G query-head siblings via GQA
+#     broadcast inside the kernel — K/V are NEVER duplicated in HBM).
+#   - Online softmax: running max `m`, running sum `l`, running weighted
+#     accumulator `acc` are carried across the K-block loop via
+#     jax.lax.fori_loop. The full [BLK_Q, S] or [BLK_Q, window] score matrix
+#     is NEVER materialized — only one [BLK_Q, BLOCK_K] tile lives in VMEM
+#     at a time. This is the actual "I/O-aware" property: HBM traffic is
+#     O(S) reads of Q/K/V blocks, not O(S^2) score matrix writes.
+#   - Sliding-window + causal masking is applied per K-block using the same
+#     relative-offset trick as the existing XLA kernel (b-independent delta),
+#     so only K-blocks that intersect [q_pos - W + 1, q_pos] are visited —
+#     blocks fully outside the window are skipped via the loop bounds, not
+#     just masked, which is where the real compute savings come from vs the
+#     existing XLA block-tiled kernel (which still computes+masks every
+#     block inside a fixed kv_len window).
+#
+# Custom VJP: backward recomputes scores per (Q-block, K-block) pair from
+# saved Q, K, V, O, m, l (NOT saved scores/probs — that's the whole point,
+# same trick as FlashAttention). This keeps backward memory O(S) instead of
+# O(S * window).
+# ---------------------------------------------------------------------------
+
+_PALLAS_BLOCK_Q = 64
+_PALLAS_BLOCK_K = 64
+
+
+def _gpu_compute_capability() -> Optional[tuple]:
+    """Returns (major, minor) compute capability of the current GPU, or None
+    if it can't be determined. Used to gate the Pallas/Triton path, which
+    JAX only supports on Ampere (SM 8.0) and newer — Turing (T4, SM 7.5) and
+    older will FAIL_PRECONDITION at Triton compile time, not at import time,
+    so we must check this explicitly before attempting the kernel."""
+    try:
+        dev = jax.devices('gpu')[0]
+        # jaxlib exposes this via device_kind (e.g. "Tesla T4", "NVIDIA A100")
+        # or via compute_capability on newer jaxlib versions.
+        cc = getattr(dev, 'compute_capability', None)
+        if cc is not None:
+            major, minor = str(cc).split('.')[:2]
+            return (int(major), int(minor))
+        return None
+    except Exception:
+        return None
+
+
+_GPU_COMPUTE_CAPABILITY = None  # lazily populated on first check
+
+
+def _pallas_supported(D: int, dtype) -> bool:
+    """Conservative gate: only use the Pallas path for configs we've reasoned
+    through (head_dim multiple of 16 for tensor-core alignment, fp16/bf16/fp32,
+    Ampere-or-newer GPU). Anything else falls back to the cuDNN/XLA path
+    automatically."""
+    global _GPU_COMPUTE_CAPABILITY
+    if os.environ.get('VEYLON_DISABLE_PALLAS_ATTN', '0') == '1':
+        return False
+    if not _PALLAS_GPU_AVAILABLE:
+        return False
+    if D % 16 != 0:
+        return False
+    if dtype not in (jnp.float16, jnp.bfloat16, jnp.float32):
+        return False
+    if _GPU_COMPUTE_CAPABILITY is None:
+        _GPU_COMPUTE_CAPABILITY = _gpu_compute_capability() or (0, 0)
+    if _GPU_COMPUTE_CAPABILITY < (8, 0):
+        # Triton (Pallas GPU backend) requires Ampere or newer. T4 (7.5),
+        # V100 (7.0), P100 (6.0) all fail here — this is a hard hardware
+        # limit, not a bug, so we skip Pallas entirely rather than let it
+        # crash through a full Triton compile attempt.
+        return False
+    return True
+
+
+def _fa_fwd_kernel(
+    q_ref, k_ref, v_ref,       # inputs, VMEM-resident blocks
+    o_ref, m_ref, l_ref,       # outputs
+    *,
+    window: int,
+    block_q: int,
+    block_k: int,
+    seq_len: int,
+    scale: float,
+):
+    """
+    Pallas kernel body — one program instance handles ONE (batch, kv_head,
+    q_block) triple, looping internally over the K-blocks that intersect
+    the causal + sliding-window range for this Q block.
+
+    Ref shapes (per-program, already sliced by BlockSpec / index_map):
+      q_ref : [block_q, D]      (single query head's slice — see note below)
+      k_ref : [seq_len, D]      (full K for this batch/kv_head; we slice
+                                  inside the loop via pl.load with dynamic
+                                  start so only ONE [block_k, D] tile is
+                                  actually resident in VMEM at a time)
+      v_ref : [seq_len, D]      (same as k_ref)
+      o_ref : [block_q, D]      (output accumulator, written once at end)
+      m_ref, l_ref : [block_q, 1]  (running softmax stats, scratch)
+    """
+    q_block_idx = pl.program_id(2)
+    q_start = q_block_idx * block_q
+
+    q = q_ref[...].astype(jnp.float32) * scale  # [block_q, D]
+
+    m_i = jnp.full((block_q, 1), -jnp.inf, dtype=jnp.float32)
+    l_i = jnp.zeros((block_q, 1), dtype=jnp.float32)
+    acc = jnp.zeros_like(q)
+
+    # Range of K-blocks that can possibly intersect this Q-block's
+    # causal+window range. Query positions in this block span
+    # [q_start, q_start + block_q - 1]. Each attends to
+    # [q_pos - window + 1, q_pos]. So the union over the block spans
+    # [q_start - window + 1, q_start + block_q - 1].
+    k_lo = jnp.maximum(0, q_start - window + 1)
+    k_hi = jnp.minimum(seq_len, q_start + block_q)  # exclusive, causal cap
+    first_k_block = k_lo // block_k
+    num_k_blocks = (k_hi - first_k_block * block_k + block_k - 1) // block_k
+    num_k_blocks = jnp.maximum(num_k_blocks, 1)
+
+    def body(i, carry):
+        m_i, l_i, acc = carry
+        k_start = (first_k_block + i) * block_k
+
+        k_blk = pl.load(
+            k_ref, (pl.dslice(k_start, block_k), slice(None))
+        ).astype(jnp.float32)  # [block_k, D]
+        v_blk = pl.load(
+            v_ref, (pl.dslice(k_start, block_k), slice(None))
+        ).astype(jnp.float32)  # [block_k, D]
+
+        scores = jnp.dot(
+            q, k_blk.T, preferred_element_type=jnp.float32
+        )  # [block_q, block_k]
+
+        q_pos = q_start + jax.lax.broadcasted_iota(jnp.int32, (block_q, block_k), 0)
+        k_pos = k_start + jax.lax.broadcasted_iota(jnp.int32, (block_q, block_k), 1)
+        causal_ok = k_pos <= q_pos
+        window_ok = (q_pos - k_pos) < window
+        bounds_ok = k_pos < seq_len
+        mask = causal_ok & window_ok & bounds_ok
+        scores = jnp.where(mask, scores, -jnp.inf)
+
+        m_ij = jnp.max(scores, axis=-1, keepdims=True)           # [block_q, 1]
+        m_new = jnp.maximum(m_i, m_ij)
+        # Guard against all-masked rows (m_new stays -inf) -> exp(0)=1 issue
+        m_new_safe = jnp.where(m_new == -jnp.inf, 0.0, m_new)
+
+        p = jnp.exp(scores - m_new_safe)                          # [block_q, block_k]
+        p = jnp.where(mask, p, 0.0)
+
+        alpha = jnp.exp(jnp.where(m_i == -jnp.inf, m_new_safe, m_i) - m_new_safe)
+        l_new = l_i * alpha + jnp.sum(p, axis=-1, keepdims=True)
+        acc_new = acc * alpha + jnp.dot(p, v_blk, preferred_element_type=jnp.float32)
+
+        return m_new, l_new, acc_new
+
+    m_i, l_i, acc = jax.lax.fori_loop(0, num_k_blocks, body, (m_i, l_i, acc))
+
+    l_safe = jnp.where(l_i > 0, l_i, 1.0)
+    out = acc / l_safe
+
+    o_ref[...] = out.astype(o_ref.dtype)
+    m_ref[...] = m_i
+    l_ref[...] = l_i
+
+
+def _pallas_fwd_single_head(q, k, v, window, block_q, block_k, scale):
+    """
+    Runs the Pallas forward kernel for ONE query head against its KV head.
+    q: [B, S, D]   k, v: [B, S, D]   (already the per-head slices)
+    Returns: out [B, S, D], m [B, S, 1], l [B, S, 1]  (m, l saved for bwd)
+    """
+    B, S, D = map(int, q.shape)
+    n_q_blocks = (S + block_q - 1) // block_q
+    S_pad = n_q_blocks * block_q
+
+    q_p = jnp.pad(q, ((0, 0), (0, S_pad - S), (0, 0)))
+    # K/V padded on the right only; kernel bounds-checks k_pos < seq_len so
+    # right-padding is safe (never read past the pad due to k_hi clamp), but
+    # we still pad to a multiple of block_k so pl.load's static block shape
+    # never reads out-of-bounds memory.
+    n_k_blocks_total = (S + block_k - 1) // block_k
+    S_pad_k = n_k_blocks_total * block_k
+    k_p = jnp.pad(k, ((0, 0), (0, S_pad_k - S), (0, 0)))
+    v_p = jnp.pad(v, ((0, 0), (0, S_pad_k - S), (0, 0)))
+
+    kernel = partial(
+        _fa_fwd_kernel,
+        window=window, block_q=block_q, block_k=block_k,
+        seq_len=S, scale=scale,
+    )
+
+    out, m, l = pl.pallas_call(
+        kernel,
+        grid=(B, 1, n_q_blocks),
+        in_specs=[
+            pl.BlockSpec((None, block_q, D), lambda b, h, i: (b, i, 0)),
+            pl.BlockSpec((None, S_pad_k, D), lambda b, h, i: (b, 0, 0)),
+            pl.BlockSpec((None, S_pad_k, D), lambda b, h, i: (b, 0, 0)),
+        ],
+        out_specs=[
+            pl.BlockSpec((None, block_q, D), lambda b, h, i: (b, i, 0)),
+            pl.BlockSpec((None, block_q, 1), lambda b, h, i: (b, i, 0)),
+            pl.BlockSpec((None, block_q, 1), lambda b, h, i: (b, i, 0)),
+        ],
+        out_shape=[
+            jax.ShapeDtypeStruct((B, S_pad, D), q.dtype),
+            jax.ShapeDtypeStruct((B, S_pad, 1), jnp.float32),
+            jax.ShapeDtypeStruct((B, S_pad, 1), jnp.float32),
+        ],
+    )(q_p, k_p, v_p)
+
+    return out[:, :S, :], m[:, :S, :], l[:, :S, :]
+
+
+def _fa_bwd_kernel(
+    q_ref, k_ref, v_ref, o_ref, do_ref, m_ref, l_ref,
+    dq_ref, dk_ref, dv_ref,
+    *,
+    window: int,
+    block_q: int,
+    block_k: int,
+    seq_len: int,
+    scale: float,
+):
+    """
+    Backward kernel — one program per (batch, k_block). Recomputes scores
+    for each intersecting Q-block on the fly (from saved Q, K, V, m, l) and
+    accumulates dK/dV. dQ is accumulated via a separate pass below since it
+    is indexed by q_block, not k_block (standard FlashAttention-2 backward
+    split to avoid atomic adds across programs).
+    """
+    k_block_idx = pl.program_id(2)
+    k_start = k_block_idx * block_k
+
+    k_blk = k_ref[...].astype(jnp.float32)   # [block_k, D]
+    v_blk = v_ref[...].astype(jnp.float32)   # [block_k, D]
+
+    dk_acc = jnp.zeros_like(k_blk)
+    dv_acc = jnp.zeros_like(v_blk)
+
+    # Q-blocks that can intersect this K-block: q_pos >= k_pos (causal) and
+    # q_pos - k_pos < window. q spans [k_start, seq_len-1] roughly, capped
+    # by window on the upper side: q_pos < k_start + block_k + window - 1.
+    q_lo = k_start
+    q_hi = jnp.minimum(seq_len, k_start + block_k + window - 1)
+    first_q_block = q_lo // block_q
+    num_q_blocks = (q_hi - first_q_block * block_q + block_q - 1) // block_q
+    num_q_blocks = jnp.maximum(num_q_blocks, 1)
+
+    def body(i, carry):
+        dk_acc, dv_acc = carry
+        q_start = (first_q_block + i) * block_q
+
+        q_blk = pl.load(q_ref, (pl.dslice(q_start, block_q), slice(None))).astype(jnp.float32) * scale
+        do_blk = pl.load(do_ref, (pl.dslice(q_start, block_q), slice(None))).astype(jnp.float32)
+        m_blk = pl.load(m_ref, (pl.dslice(q_start, block_q), slice(None))).astype(jnp.float32)
+        l_blk = pl.load(l_ref, (pl.dslice(q_start, block_q), slice(None))).astype(jnp.float32)
+        o_blk = pl.load(o_ref, (pl.dslice(q_start, block_q), slice(None))).astype(jnp.float32)
+
+        scores = jnp.dot(q_blk, k_blk.T, preferred_element_type=jnp.float32)
+
+        q_pos = q_start + jax.lax.broadcasted_iota(jnp.int32, (block_q, block_k), 0)
+        k_pos = k_start + jax.lax.broadcasted_iota(jnp.int32, (block_q, block_k), 1)
+        causal_ok = k_pos <= q_pos
+        window_ok = (q_pos - k_pos) < window
+        bounds_ok = k_pos < seq_len
+        mask = causal_ok & window_ok & bounds_ok
+
+        l_safe = jnp.where(l_blk > 0, l_blk, 1.0)
+        p = jnp.where(mask, jnp.exp(scores - m_blk), 0.0) / l_safe  # [block_q, block_k]
+
+        dv_acc = dv_acc + jnp.dot(p.T, do_blk, preferred_element_type=jnp.float32)
+
+        dp = jnp.dot(do_blk, v_blk.T, preferred_element_type=jnp.float32)  # [block_q, block_k]
+        Di = jnp.sum(do_blk * o_blk, axis=-1, keepdims=True)  # [block_q, 1]
+        dscores = p * (dp - Di)
+        dscores = jnp.where(mask, dscores, 0.0)
+
+        dk_acc = dk_acc + jnp.dot(dscores.T, q_blk, preferred_element_type=jnp.float32) * scale
+
+        return dk_acc, dv_acc
+
+    dk_acc, dv_acc = jax.lax.fori_loop(0, num_q_blocks, body, (dk_acc, dv_acc))
+
+    dk_ref[...] = dk_acc.astype(dk_ref.dtype)
+    dv_ref[...] = dv_acc.astype(dv_ref.dtype)
+
+
+def _fa_bwd_dq_kernel(
+    q_ref, k_ref, v_ref, o_ref, do_ref, m_ref, l_ref,
+    dq_ref,
+    *,
+    window: int,
+    block_q: int,
+    block_k: int,
+    seq_len: int,
+    scale: float,
+):
+    """Separate pass computing dQ, one program per (batch, q_block), looping
+    over intersecting K-blocks. Kept separate from the dK/dV kernel because
+    dQ is naturally indexed by q_block and dK/dV by k_block — fusing both
+    into one kernel would need cross-program atomics, which Pallas/Triton
+    doesn't support cleanly. Recomputation cost (~2x score matmuls total
+    across both passes) is the standard FlashAttention-2 backward tradeoff."""
+    q_block_idx = pl.program_id(2)
+    q_start = q_block_idx * block_q
+
+    q_blk = q_ref[...].astype(jnp.float32) * scale
+    do_blk = do_ref[...].astype(jnp.float32)
+    m_blk = m_ref[...].astype(jnp.float32)
+    l_blk = l_ref[...].astype(jnp.float32)
+    o_blk = o_ref[...].astype(jnp.float32)
+
+    dq_acc = jnp.zeros_like(q_blk)
+
+    k_lo = jnp.maximum(0, q_start - window + 1)
+    k_hi = jnp.minimum(seq_len, q_start + block_q)
+    first_k_block = k_lo // block_k
+    num_k_blocks = (k_hi - first_k_block * block_k + block_k - 1) // block_k
+    num_k_blocks = jnp.maximum(num_k_blocks, 1)
+
+    def body(i, dq_acc):
+        k_start = (first_k_block + i) * block_k
+        k_blk = pl.load(k_ref, (pl.dslice(k_start, block_k), slice(None))).astype(jnp.float32)
+        v_blk = pl.load(v_ref, (pl.dslice(k_start, block_k), slice(None))).astype(jnp.float32)
+
+        scores = jnp.dot(q_blk, k_blk.T, preferred_element_type=jnp.float32)
+
+        q_pos = q_start + jax.lax.broadcasted_iota(jnp.int32, (block_q, block_k), 0)
+        k_pos = k_start + jax.lax.broadcasted_iota(jnp.int32, (block_q, block_k), 1)
+        causal_ok = k_pos <= q_pos
+        window_ok = (q_pos - k_pos) < window
+        bounds_ok = k_pos < seq_len
+        mask = causal_ok & window_ok & bounds_ok
+
+        l_safe = jnp.where(l_blk > 0, l_blk, 1.0)
+        p = jnp.where(mask, jnp.exp(scores - m_blk), 0.0) / l_safe
+
+        dp = jnp.dot(do_blk, v_blk.T, preferred_element_type=jnp.float32)
+        Di = jnp.sum(do_blk * o_blk, axis=-1, keepdims=True)
+        dscores = p * (dp - Di)
+        dscores = jnp.where(mask, dscores, 0.0)
+
+        dq_acc = dq_acc + jnp.dot(dscores, k_blk, preferred_element_type=jnp.float32) * scale
+        return dq_acc
+
+    dq_acc = jax.lax.fori_loop(0, num_k_blocks, body, dq_acc)
+    dq_ref[...] = dq_acc.astype(dq_ref.dtype)
+
+
+def _pallas_bwd_single_head(q, k, v, o, do, m, l, window, block_q, block_k, scale):
+    """Runs both backward kernels (dK/dV and dQ) for one query/KV head pair."""
+    B, S, D = map(int, q.shape)
+    n_q_blocks = (S + block_q - 1) // block_q
+    n_k_blocks = (S + block_k - 1) // block_k
+    S_pad_q = n_q_blocks * block_q
+    S_pad_k = n_k_blocks * block_k
+
+    pad_q = lambda x, fill=0.0: jnp.pad(x, ((0, 0), (0, S_pad_q - S), (0, 0)), constant_values=fill)
+    pad_k = lambda x: jnp.pad(x, ((0, 0), (0, S_pad_k - S), (0, 0)))
+
+    q_p, o_p, do_p = pad_q(q), pad_q(o), pad_q(do)
+    m_p = jnp.pad(m, ((0, 0), (0, S_pad_q - S), (0, 0)), constant_values=jnp.inf)
+    l_p = jnp.pad(l, ((0, 0), (0, S_pad_q - S), (0, 0)), constant_values=1.0)
+    k_p, v_p = pad_k(k), pad_k(v)
+
+    dkdv_kernel = partial(
+        _fa_bwd_kernel, window=window, block_q=block_q, block_k=block_k,
+        seq_len=S, scale=scale,
+    )
+    dk, dv = pl.pallas_call(
+        dkdv_kernel,
+        grid=(B, 1, n_k_blocks),
+        in_specs=[
+            pl.BlockSpec((None, S_pad_q, D), lambda b, h, j: (b, 0, 0)),   # q (full, sliced inside)
+            pl.BlockSpec((None, block_k, D), lambda b, h, j: (b, j, 0)),    # k block
+            pl.BlockSpec((None, block_k, D), lambda b, h, j: (b, j, 0)),    # v block
+            pl.BlockSpec((None, S_pad_q, D), lambda b, h, j: (b, 0, 0)),   # o (full)
+            pl.BlockSpec((None, S_pad_q, D), lambda b, h, j: (b, 0, 0)),   # do (full)
+            pl.BlockSpec((None, S_pad_q, 1), lambda b, h, j: (b, 0, 0)),   # m (full)
+            pl.BlockSpec((None, S_pad_q, 1), lambda b, h, j: (b, 0, 0)),   # l (full)
+        ],
+        out_specs=[
+            pl.BlockSpec((None, block_k, D), lambda b, h, j: (b, j, 0)),
+            pl.BlockSpec((None, block_k, D), lambda b, h, j: (b, j, 0)),
+        ],
+        out_shape=[
+            jax.ShapeDtypeStruct((B, S_pad_k, D), k.dtype),
+            jax.ShapeDtypeStruct((B, S_pad_k, D), v.dtype),
+        ],
+    )(q_p, k_p, v_p, o_p, do_p, m_p, l_p)
+
+    dq_kernel = partial(
+        _fa_bwd_dq_kernel, window=window, block_q=block_q, block_k=block_k,
+        seq_len=S, scale=scale,
+    )
+    dq = pl.pallas_call(
+        dq_kernel,
+        grid=(B, 1, n_q_blocks),
+        in_specs=[
+            pl.BlockSpec((None, block_q, D), lambda b, h, i: (b, i, 0)),   # q block
+            pl.BlockSpec((None, S_pad_k, D), lambda b, h, i: (b, 0, 0)),   # k (full)
+            pl.BlockSpec((None, S_pad_k, D), lambda b, h, i: (b, 0, 0)),   # v (full)
+            pl.BlockSpec((None, block_q, D), lambda b, h, i: (b, i, 0)),   # o block
+            pl.BlockSpec((None, block_q, D), lambda b, h, i: (b, i, 0)),   # do block
+            pl.BlockSpec((None, block_q, 1), lambda b, h, i: (b, i, 0)),   # m block
+            pl.BlockSpec((None, block_q, 1), lambda b, h, i: (b, i, 0)),   # l block
+        ],
+        out_specs=pl.BlockSpec((None, block_q, D), lambda b, h, i: (b, i, 0)),
+        out_shape=jax.ShapeDtypeStruct((B, S_pad_q, D), q.dtype),
+    )(q_p, k_p, v_p, o_p, do_p, m_p, l_p)
+
+    return dq[:, :S, :], dk[:, :S, :], dv[:, :S, :]
+
+
+@partial(jax.custom_vjp, nondiff_argnums=(3, 4, 5, 6))
+def _pallas_gqa_swa_head(q, k, v, window, block_q, block_k, scale):
+    """Single (query-head, kv-head) FlashAttention call with custom VJP.
+    q, k, v: [B, S, D] for ONE head pair (GQA broadcast handled by caller)."""
+    out, _, _ = _pallas_fwd_single_head(q, k, v, window, block_q, block_k, scale)
+    return out
+
+
+def _pallas_gqa_swa_head_fwd(q, k, v, window, block_q, block_k, scale):
+    out, m, l = _pallas_fwd_single_head(q, k, v, window, block_q, block_k, scale)
+    return out, (q, k, v, out, m, l)
+
+
+def _pallas_gqa_swa_head_bwd(window, block_q, block_k, scale, residuals, dout):
+    q, k, v, out, m, l = residuals
+    dq, dk, dv = _pallas_bwd_single_head(
+        q, k, v, out, dout, m, l, window, block_q, block_k, scale
+    )
+    return dq, dk, dv
+
+
+_pallas_gqa_swa_head.defvjp(_pallas_gqa_swa_head_fwd, _pallas_gqa_swa_head_bwd)
+
+
+def _pallas_flash_gqa_swa(
+    q: jnp.ndarray,
+    k: jnp.ndarray,
+    v: jnp.ndarray,
+    window_size: int,
+    block_q: int = _PALLAS_BLOCK_Q,
+    block_k: int = _PALLAS_BLOCK_K,
+) -> jnp.ndarray:
+    """
+    I/O-aware FlashAttention-style GQA SWA, entry point for the Pallas path.
+
+    q: [B, Hq,  S, D]
+    k: [B, Hkv, S, D]
+    v: [B, Hkv, S, D]
+
+    GQA is handled by vmapping the single-head kernel over KV heads, and
+    within each KV head over its G query-head siblings — K/V are never
+    physically duplicated; only the (small) grid iterates over G.
+    """
+    B, Hq, S, D = map(int, q.shape)
+    _, Hkv, Sk, Dk = map(int, k.shape)
+    if Hq % Hkv != 0:
+        raise ValueError(f"Hq={Hq} must be divisible by Hkv={Hkv}")
+    G = Hq // Hkv
+    scale = 1.0 / math.sqrt(float(D))
+
+    # [B, Hkv, G, S, D]
+    q_g = q.reshape(B, Hkv, G, S, D)
+
+    # vmap over (Hkv, G): each call gets q[B,S,D] for one query head and the
+    # matching k/v[B,S,D] for its KV head (broadcast across G, no copy of
+    # the underlying K/V buffer beyond what vmap's batching rule does).
+    def per_kv_head(q_kv, k_h, v_h):
+        # q_kv: [G, B, S, D]  k_h, v_h: [B, S, D]
+        fn = lambda qh: _pallas_gqa_swa_head(qh, k_h, v_h, window_size, block_q, block_k, scale)
+        return jax.vmap(fn)(q_kv)  # [G, B, S, D]
+
+    q_g_t = q_g.transpose(1, 2, 0, 3, 4)  # [Hkv, G, B, S, D]
+    k_t = k.transpose(1, 0, 2, 3)          # [Hkv, B, S, D]
+    v_t = v.transpose(1, 0, 2, 3)
+
+    out = jax.vmap(per_kv_head)(q_g_t, k_t, v_t)  # [Hkv, G, B, S, D]
+    out = out.transpose(2, 0, 1, 3, 4).reshape(B, Hq, S, D)  # [B, Hq, S, D]
+    return out.astype(q.dtype)
+
+
 # ---------------------------------------------------------------------------
 # GPU path: JAX cuDNN FlashAttention via jax.nn.dot_product_attention
 # ---------------------------------------------------------------------------
@@ -162,24 +665,21 @@ def _gpu_flash_gqa_swa(
         k_s = k.transpose(0, 2, 1, 3)
         v_s = v.transpose(0, 2, 1, 3)
 
-        def _full_attn(q_s, k_s, v_s):
-            try:
-                out = jax.nn.dot_product_attention(
-                    q_s, k_s, v_s,
-                    scale=scale,
-                    is_causal=True,
-                    implementation='cudnn',
-                )
-                return out
-            except Exception:
-                # cuDNN unavailable: fall through to XLA path below
-                return None
+        def _full_attn(q_, k_, v_):
+            return jax.nn.dot_product_attention(
+                q_, k_, v_,
+                scale=scale,
+                is_causal=True,
+                implementation='cudnn',
+            )
 
-        result = _full_attn(q_s, k_s, v_s)
-        if result is not None:
-            if use_remat:
-                result = jax.checkpoint(lambda q, k, v: _full_attn(q, k, v))(q_s, k_s, v_s)
-            return result.transpose(0, 2, 1, 3).astype(q.dtype)
+        if use_remat:
+            # jax.checkpoint recomputes _full_attn on the backward pass; the
+            # function itself still runs exactly ONCE per forward pass.
+            result = jax.checkpoint(_full_attn)(q_s, k_s, v_s)
+        else:
+            result = _full_attn(q_s, k_s, v_s)
+        return result.transpose(0, 2, 1, 3).astype(q.dtype)
 
     # ── SWA path: XLA block-tiled kernel (GPU block_size=64) ─────────────────
     # This is the fast path for SWA on GPU.
@@ -573,6 +1073,34 @@ def flash_splash_attention(
 
     # ── Dispatch ──────────────────────────────────────────────────────────────
     if active_backend == 'gpu':
+        B, Hq, S, D = map(int, q.shape)
+        if _pallas_supported(D, q.dtype):
+            try:
+                return _pallas_flash_gqa_swa(
+                    q, k, v,
+                    window_size=int(window_size),
+                    block_q=min(_PALLAS_BLOCK_Q, S) if S < _PALLAS_BLOCK_Q else _PALLAS_BLOCK_Q,
+                    block_k=min(_PALLAS_BLOCK_K, S) if S < _PALLAS_BLOCK_K else _PALLAS_BLOCK_K,
+                )
+            except Exception as e:
+                # Any Pallas/Triton compile or runtime failure (unsupported
+                # GPU arch, block size mismatch, etc.) falls back silently to
+                # the proven cuDNN/XLA path below — training never crashes
+                # because of this optimization. Set
+                # VEYLON_DEBUG_PALLAS_ATTN=1 to see what actually failed.
+                if os.environ.get('VEYLON_DEBUG_PALLAS_ATTN', '0') == '1':
+                    print(f"[veylon_attention] Pallas path failed, falling back: "
+                          f"{type(e).__name__}: {e}")
+        # cuDNN fused attention only accepts fp16/bf16/fp8 — fp32 inputs must
+        # go through the plain-XLA fallback further down in
+        # _gpu_flash_gqa_swa rather than crashing on the cuDNN dtype check.
+        if q.dtype not in (jnp.float16, jnp.bfloat16):
+            return _block_gqa_swa(
+                q=q, k=k, v=v,
+                window_size=int(window_size),
+                block_size=int(GPU_BLOCK_SIZE),
+                use_remat=use_remat,
+            )
         return _gpu_flash_gqa_swa(
             q=q,
             k=k,
@@ -836,6 +1364,61 @@ if __name__ == "__main__":
         if ot.shape == (1, 4, S_t, 16): ok(f"S={S_t}  →  {ot.shape}")
         else:                            fail(f"S={S_t}  wrong shape {ot.shape}")
 
+    # ── 8. Pallas GPU kernel (forward correctness + gradient check) ─────────
+    section("8 · Pallas/Triton FlashAttention kernel (GPU only)")
+    if not _PALLAS_GPU_AVAILABLE:
+        print("  (skipped — Pallas not importable in this environment)")
+    elif _detect_backend() != 'gpu':
+        print("  (skipped — no GPU backend detected)")
+    elif not _pallas_supported(16, jnp.float16):
+        cc = _gpu_compute_capability()
+        if cc is not None and cc < (8, 0):
+            print(f"  (skipped — GPU compute capability {cc[0]}.{cc[1]} < 8.0; "
+                  f"Triton/Pallas requires Ampere or newer. cuDNN path handles "
+                  f"FlashAttention on this GPU instead.)")
+        else:
+            print("  (skipped — Pallas gated off for this config; "
+                  "set VEYLON_DEBUG_PALLAS_ATTN=1 for details)")
+    else:
+        B_, Hq_, Hkv_, S_, D_, W_ = 1, 4, 2, 130, 16, 24
+        ks = jax.random.split(jax.random.PRNGKey(99), 3)
+        qp = jax.random.normal(ks[0], (B_, Hq_,  S_, D_), dtype=jnp.float32) * 0.1
+        kp = jax.random.normal(ks[1], (B_, Hkv_, S_, D_), dtype=jnp.float32) * 0.1
+        vp = jax.random.normal(ks[2], (B_, Hkv_, S_, D_), dtype=jnp.float32) * 0.1
+
+        try:
+            out_pallas = _pallas_flash_gqa_swa(qp, kp, vp, window_size=W_, block_q=32, block_k=32)
+
+            # Reference via existing XLA block-tiled kernel
+            out_ref = _block_gqa_swa(qp, kp, vp, window_size=W_, block_size=32, use_remat=False)
+
+            err = float(jnp.max(jnp.abs(out_pallas - out_ref)))
+            if err < 1e-3:
+                ok(f"Forward matches XLA reference: max err = {err:.2e}")
+            else:
+                fail(f"Forward MISMATCH vs XLA reference: max err = {err:.2e}")
+
+            # Gradient check: compare d(sum(out))/d(q,k,v) against XLA reference
+            def loss_pallas(q, k, v):
+                return jnp.sum(_pallas_flash_gqa_swa(q, k, v, window_size=W_, block_q=32, block_k=32))
+
+            def loss_ref(q, k, v):
+                return jnp.sum(_block_gqa_swa(q, k, v, window_size=W_, block_size=32, use_remat=False))
+
+            gp = jax.grad(loss_pallas, argnums=(0, 1, 2))(qp, kp, vp)
+            gr = jax.grad(loss_ref, argnums=(0, 1, 2))(qp, kp, vp)
+
+            names = ['dQ', 'dK', 'dV']
+            for name, gp_i, gr_i in zip(names, gp, gr):
+                gerr = float(jnp.max(jnp.abs(gp_i - gr_i)))
+                if gerr < 1e-2:
+                    ok(f"{name} matches XLA autodiff: max err = {gerr:.2e}")
+                else:
+                    fail(f"{name} MISMATCH vs XLA autodiff: max err = {gerr:.2e}")
+
+        except Exception as e:
+            fail(f"Pallas kernel raised an exception: {type(e).__name__}: {e}")
+
     # ── Summary ──────────────────────────────────────────────────────────────
     print(f"\n{HDR}{'═'*62}{RST}")
     if failures == 0: print(f"  {PASS}  All tests passed.")