veylon_attention.py

"""
veylon_attention.py  —  Veylon Alpha 1
======================================
Block-tiled Native GQA Sliding Window Attention for JAX / Keras-3 / TPU v5e.

Why NOT lax.scan + dynamic_slice
---------------------------------
The previous implementation used:

    jax.lax.scan(step, None, (jnp.arange(S), q_scan))

with `dynamic_slice(k_pad, (0,0,t,0), (B,Hkv,W,D))` inside the body.

XLA/XLA-TPU has a known pathology with this pattern:
  - The scan body contains a dynamic gather (dynamic_slice on a traced int `t`).
  - XLA's while-loop lowering stages ALL window materializations into a single
    large buffer to enable pipeline prefetching.
  - On TPU this produces a hidden tensor  [S, B, Hkv, W, D]  which is exactly
    what you see in the OOM allocation log:  f32[1024,8,4,512,64].
  - jax.checkpoint does NOT protect against this because it is a compiler
    (HLO-level) allocation, not a JAX-level rematerialization artifact.

The fix: block-tiled computation with fully static tensor shapes
----------------------------------------------------------------
Instead of iterating over S individual tokens we iterate over
(S / BLK) blocks of queries.  For each block:

  - q_blk  : [B, Hq,  BLK,        D]   <- static shape
  - k_blk  : [B, Hkv, BLK + W - 1, D]  <- static shape
  - scores : [B, Hkv, G, BLK, BLK+W-1] <- static shape

XLA sees ONLY static shapes inside the map body.  There is no
gather-inside-loop pattern.  XLA can freely fuse, pipeline, and
tile these einsums onto TPU systolic arrays without hidden buffers.

lax.fori_loop with dynamic_update_slice
----------------------------------------
We use `jax.lax.fori_loop` (not `lax.map` or `lax.scan`) because:
  - `lax.map` returns [n_blocks, B, Hq, BLK, D] — XLA stages this entire
    stack in HBM before the final transpose+reshape.  On large batches or
    many blocks this causes the VRAM spike you see in the profiler.
  - `lax.scan` has the same problem (stacks carry outputs).
  - `fori_loop` carries a single pre-allocated [B, Hq, S_pad, D] output
    buffer and writes each block with dynamic_update_slice.  XLA sees ONE
    static-shape buffer (same footprint as the final output) throughout,
    eliminating the n_blocks-deep intermediate stack entirely.

Tensor size invariants
----------------------
NEVER created:
    [B, Hq,  S, S]               full attention score matrix
    [S, B, Hkv, W, D]            per-token window stack (old scan bug)
    [B, Hq, Hkv, S, D]           duplicated KV heads

Largest tensors inside map body (all STATIC shapes):
    k_blk, v_blk  : [B, Hkv, BLK+W-1, D]
    scores        : [B, Hkv, G, BLK, BLK+W-1]
    probs         : same

Memory complexity
-----------------
  k_pad, v_pad  : O(B × Hkv × (S + W) × D)   linear in S, scales with Hkv
  q             : O(B × Hq  × S × D)
  Per block     : O(B × Hkv × (BLK + W) × D)  constant w.r.t. S
  Output        : O(B × Hq  × S × D)
  TOTAL         : O(B × (Hkv + Hq) × S × D)   strictly linear in S

TPU-specific notes
------------------
  * All shapes in map body are concrete at trace time — XLA never inserts
    shape-dependent conditionals or recompiles.
  * BF16 inputs are cast to FP32 before matmul/softmax, then cast back.
  * BLK should be a multiple of 128 on TPU v5e for optimal systolic array
    utilization (default 128, tunable via `block_size` parameter).
  * The precomputed mask is a static bool array passed as a closed-over
    constant — XLA fuses it into the einsum kernel at no extra memory cost.
  * jax.checkpoint on the map body protects backward-pass activations
    (one block at a time, not one token at a time — much coarser remat).
"""

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
# ---------------------------------------------------------------------------

# TPU v5e: 128-wide systolic arrays → BLK=128 saturates MXU
TPU_BLOCK_SIZE: int = 128

# GPU: warp size=32, tensor cores tile 16×16 or 8×16 → BLK=64 is safe default
# cuDNN FlashAttention internally tiles at 64 or 128 depending on head dim
GPU_BLOCK_SIZE: int = 64

# ---------------------------------------------------------------------------
# Backend detection
# ---------------------------------------------------------------------------

def _detect_backend() -> str:
    """
    Detect the current JAX backend.
    Returns 'tpu', 'gpu', or 'cpu'.
    """
    try:
        backend = jax.default_backend().lower()
        if 'tpu' in backend:
            return 'tpu'
        elif 'gpu' in backend or 'cuda' in backend:
            return 'gpu'
        return 'cpu'
    except Exception:
        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
# ---------------------------------------------------------------------------

def _gpu_flash_gqa_swa(
    q: jnp.ndarray,
    k: jnp.ndarray,
    v: jnp.ndarray,
    window_size: int,
    use_remat: bool = True,
) -> jnp.ndarray:
    """
    GPU-optimized GQA Sliding Window Attention.

    cuDNN FlashAttention does NOT reliably support SWA masking across all
    JAX/cuDNN versions. Instead we use:

      - cuDNN for the raw QK^T matmul + softmax + V aggregation
        (via jax.nn.dot_product_attention without masking)
        only when window_size >= S (full attention — no masking needed).

      - XLA block-tiled path with GPU-friendly block_size=64 for SWA
        (window_size < S). This avoids the cuDNN engine config error
        while still running fast on CUDA via XLA's GPU backend.

    Both paths use BF16 compute and avoid materializing [S,S] matrices.
    """
    B,  Hq,  S,  D  = map(int, q.shape)
    _B, 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))

    # ── Full attention (no window mask): use cuDNN ────────────────────────────
    if window_size >= S:
        # [B, S, H, D] layout for cuDNN
        q_s = q.transpose(0, 2, 1, 3)
        k_s = k.transpose(0, 2, 1, 3)
        v_s = v.transpose(0, 2, 1, 3)

        def _full_attn(q_, k_, v_):
            return jax.nn.dot_product_attention(
                q_, k_, v_,
                scale=scale,
                is_causal=True,
                implementation='cudnn',
            )

        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.
    # XLA compiles this to efficient CUDA matmuls with BF16 tensor cores.
    # GPU_BLOCK_SIZE=64 matches CUDA warp/tensor-core tiling.
    return _block_gqa_swa(
        q=q,
        k=k,
        v=v,
        window_size=window_size,
        block_size=GPU_BLOCK_SIZE,
        use_remat=use_remat,
    )


# ---------------------------------------------------------------------------
# GPU decode path: single-token GQA SWA for inference
# ---------------------------------------------------------------------------

def _gpu_decode_swa(
    q: jnp.ndarray,
    k: jnp.ndarray,
    v: jnp.ndarray,
) -> jnp.ndarray:
    """
    GPU decode path for single-token generation (S=1).

    q: [B, Hq,  1, D]
    k: [B, Hkv, W, D]
    v: [B, Hkv, W, D]

    Returns: [B, Hq, 1, D]
    """
    B, Hq, S, D  = map(int, q.shape)
    _, Hkv, W, _ = map(int, k.shape)
    G = Hq // Hkv
    scale = 1.0 / math.sqrt(float(D))

    # [B, 1, Hq, D] and [B, W, Hkv, D] for cuDNN
    q_s = q.transpose(0, 2, 1, 3)
    k_s = k.transpose(0, 2, 1, 3)
    v_s = v.transpose(0, 2, 1, 3)

    try:
        # S=1 decode: all W tokens are past, no masking needed
        # cuDNN handles this as a batched GEMV — very fast
        out = jax.nn.dot_product_attention(
            q_s, k_s, v_s,
            scale=scale,
            is_causal=False,
            implementation='cudnn',
        )
        return out.transpose(0, 2, 1, 3).astype(q.dtype)
    except Exception:
        pass

    # Fallback: native GQA einsum (always works)
    q_g    = q.reshape(B, Hkv, G, 1, D)
    scores = (
        jnp.einsum('bngqd,bnkd->bngqk',
                   q_g.astype(jnp.float32),
                   k.astype(jnp.float32))
        * scale
    )
    probs  = jax.nn.softmax(scores, axis=-1)
    out_g  = jnp.einsum('bngqk,bnkd->bngqd', probs, v.astype(jnp.float32))
    return out_g.reshape(B, Hq, 1, D).astype(q.dtype)


# ---------------------------------------------------------------------------
# Utility helpers
# ---------------------------------------------------------------------------

def _next_power_of_2(x: int) -> int:
    x = int(x)
    if x <= 1:
        return 1
    return 1 << (x - 1).bit_length()


def apply_rope(
    x: jnp.ndarray,
    cos: jnp.ndarray,
    sin: jnp.ndarray,
    offset: int = 0,
) -> jnp.ndarray:
    """
    Apply Rotary Position Encoding to [..., S, D].
    cos/sin tables are pre-built for D//2 (half the head dim).
    Uses dynamic_slice — safe under jit and XLA outside of attention body.
    """
    if x.ndim < 2:
        raise ValueError(f"apply_rope: expected ≥2-D input, got shape {x.shape}")
    d = int(x.shape[-1])
    if d % 2 != 0:
        raise ValueError(f"apply_rope: head dim must be even, got {d}")
    half = d // 2
    seq_len = int(x.shape[-2])
    if offset + seq_len > int(cos.shape[0]):
        raise ValueError(
            f"apply_rope: RoPE table too small — "
            f"offset={offset}, seq_len={seq_len}, table_size={cos.shape[0]}"
        )
    cos_s = jax.lax.dynamic_slice(cos, (offset, 0), (seq_len, half)).astype(x.dtype)
    sin_s = jax.lax.dynamic_slice(sin, (offset, 0), (seq_len, half)).astype(x.dtype)
    while cos_s.ndim < x.ndim - 1:
        cos_s = cos_s[None]
        sin_s = sin_s[None]
    x1, x2 = x[..., :half], x[..., half:]
    return jnp.concatenate(
        [x1 * cos_s - x2 * sin_s, x1 * sin_s + x2 * cos_s],
        axis=-1,
    )


# ---------------------------------------------------------------------------
# Core kernel: block-tiled native GQA SWA
# ---------------------------------------------------------------------------

def _block_gqa_swa(
    q: jnp.ndarray,
    k: jnp.ndarray,
    v: jnp.ndarray,
    window_size: int,
    block_size: int = TPU_BLOCK_SIZE,
    use_remat: bool = True,
) -> jnp.ndarray:
    """
    Block-tiled Native GQA Sliding Window Attention.

    This function is the replacement for the scan+dynamic_slice approach.
    All tensor shapes inside the map body are STATIC — XLA never sees a
    gather-inside-loop pattern, so the hidden [S,B,H,W,D] buffer cannot form.

    Parameters
    ----------
    q           : [B, Hq,  S, D]  — any float dtype (BF16 in production)
    k           : [B, Hkv, S, D]  — same dtype
    v           : [B, Hkv, S, D]  — same dtype
    window_size : causal window W; token t attends to [max(0, t-W+1), t]
    block_size  : query block size BLK (tune to 128 for TPU v5e)
    use_remat   : wrap map body in jax.checkpoint (recommended for training)

    Returns
    -------
    [B, Hq, S, D] — same dtype as q
    """
    if q.ndim != 4 or k.ndim != 4 or v.ndim != 4:
        raise ValueError(
            f"_block_gqa_swa: expected 4-D inputs, "
            f"got q={q.shape}  k={k.shape}  v={v.shape}"
        )

    B,  Hq,  S,  D   = map(int, q.shape)
    _B, Hkv, Sk, Dk  = map(int, k.shape)

    if _B != B:
        raise ValueError(f"Batch mismatch: q B={B}, k B={_B}")
    if Dk != D:
        raise ValueError(f"Head-dim mismatch: q D={D}, k D={Dk}")
    if Sk != S:
        raise ValueError(f"Sequence-length mismatch: q S={S}, k S={Sk}")
    if Hq % Hkv != 0:
        raise ValueError(f"Hq={Hq} must be divisible by Hkv={Hkv}")

    G   = Hq // Hkv                  # queries per KV head
    W   = int(window_size)
    BLK = int(block_size)
    if BLK <= 0:
        raise ValueError(f"block_size must be > 0, got {BLK}")

    scale = 1.0 / math.sqrt(float(D))

    # ── Pad sequence to a multiple of BLK ────────────────────────────────────
    n_blocks = (S + BLK - 1) // BLK
    S_pad    = n_blocks * BLK          # ≥ S, multiple of BLK

    # ── Pad q along sequence axis ─────────────────────────────────────────────
    # [B, Hq, S_pad, D]
    # The extra (S_pad - S) tokens are zero-padded and trimmed from output.
    q_pad = jnp.pad(q, ((0, 0), (0, 0), (0, S_pad - S), (0, 0)))

    # ── Pad k/v on the LEFT by (W-1) for causal alignment ────────────────────
    # [B, Hkv, S_pad + W - 1, D]
    #
    # After padding, for query block starting at global position b*BLK:
    #   kv slice starts at offset b*BLK in k_pad
    #   kv slice has static length  BLK + W - 1
    #   It covers original positions  [b*BLK - (W-1), b*BLK + BLK - 1]
    #   which, after clipping to ≥ 0, is exactly the causal window.
    kv_pad_len = S_pad + W - 1        # total padded kv length (static)
    lpad       = W - 1                # left zero-padding width

    k_pad = jnp.pad(k, ((0, 0), (0, 0), (lpad, S_pad - S), (0, 0)))  # [B,Hkv,kv_pad_len,D]
    v_pad = jnp.pad(v, ((0, 0), (0, 0), (lpad, S_pad - S), (0, 0)))

    # ── Static per-block mask ─────────────────────────────────────────────────
    # Compute a boolean mask of shape [BLK, BLK + W - 1].
    # Entry [q_local, k_local] is True (mask=attend) when:
    #   (a) k is not from left-padding   → k_local >= W - 1 - (something)
    # We cannot compute the absolute positions statically because they depend
    # on block index b.  So we record the RELATIVE offsets and apply the
    # offset inside the map body using only static arithmetic on local indices.
    #
    # q_abs  = b*BLK + q_local          (q_local in 0..BLK-1)
    # k_abs  = b*BLK + k_local - lpad   (k_local in 0..BLK+W-2)
    #
    # Attend iff:
    #   k_abs >= 0                (not left padding)
    #   k_abs <= q_abs            (causal)
    #   q_abs - k_abs < W         (within window)
    #
    # q_abs - k_abs  = q_local - k_local + lpad   (b cancels out!)
    # k_abs >= 0     ↔ k_local >= lpad - b*BLK   (depends on b → handle in body)
    # k_abs <= q_abs ↔ k_local - q_local <= lpad  (b cancels out! → STATIC)
    #
    # Only the "k_abs >= 0" condition depends on b (first block only).
    # We handle it cheaply with a dynamic mask inside the body.
    # Everything else is b-independent and can be PRECOMPUTED once.

    kv_len = BLK + W - 1   # static length of kv slice per block

    q_local  = jnp.arange(BLK,    dtype=jnp.int32)   # [BLK]
    kv_local = jnp.arange(kv_len, dtype=jnp.int32)   # [kv_len]

    # Relative offset: delta[q, k] = q_local[q] - kv_local[k] + lpad
    #   = q_abs - k_abs    (b-independent)
    delta = q_local[:, None] - kv_local[None, :] + lpad  # [BLK, kv_len]

    # Static masks (b-independent)
    static_future_mask  = delta < 0          # k is in the future (causal)  → mask
    static_window_mask  = delta >= W         # k is too far back (window)   → mask
    static_base_mask    = static_future_mask | static_window_mask  # [BLK, kv_len]

    # ── Map body ──────────────────────────────────────────────────────────────

    def process_block(b: jnp.ndarray) -> jnp.ndarray:
        """
        b : scalar int32 block index in [0, n_blocks)

        Tensor shapes inside this function are ALL STATIC:
            q_blk   : [B, Hq,  BLK,    D]
            k_blk   : [B, Hkv, kv_len, D]
            v_blk   : [B, Hkv, kv_len, D]
            scores  : [B, Hkv, G, BLK, kv_len]
            probs   : same

        XLA has NO gather-inside-loop here.  The dynamic_slice start index
        `b * BLK` is a scalar multiply — XLA lowers this to a simple pointer
        offset, not a buffer materialisation.
        """
        blk_start = b * BLK  # scalar traced int32

        # Static-shape slices — the KEY difference from the scan approach.
        # XLA sees shapes (B,Hq,BLK,D) and (B,Hkv,kv_len,D) as compile-time
        # constants.  It cannot stage all blocks simultaneously because
        # lax.map gives it one block at a time with no output accumulation.
        q_blk = jax.lax.dynamic_slice(q_pad, (0, 0, blk_start,        0), (B, Hq,  BLK,    D))
        k_blk = jax.lax.dynamic_slice(k_pad, (0, 0, blk_start,        0), (B, Hkv, kv_len, D))
        v_blk = jax.lax.dynamic_slice(v_pad, (0, 0, blk_start,        0), (B, Hkv, kv_len, D))

        # ── Native GQA reshape (zero-copy view) ──────────────────────────
        # [B, Hq, BLK, D] → [B, Hkv, G, BLK, D]
        q_g = q_blk.reshape(B, Hkv, G, BLK, D)

        # ── Dot-product scores ────────────────────────────────────────────
        # [B, Hkv, G, BLK, kv_len]   ← STATIC, FUSED by XLA
        scores = (
            jnp.einsum(
                "bngqd,bnkd->bngqk",
                q_g.astype(jnp.float32),
                k_blk.astype(jnp.float32),
            )
            * scale
        )

        # ── Masking ───────────────────────────────────────────────────────
        # (a) b-independent mask (precomputed, fused as constant)
        mask = static_base_mask  # [BLK, kv_len]

        # (b) Left-padding mask: k_abs < 0  ↔  kv_local < lpad - blk_start
        #     Only non-trivial for b=0 (the very first block).
        #     For all subsequent blocks, lpad - blk_start < 0, so no extra masking.
        #     We compute it dynamically but it is a single scalar comparison
        #     broadcast — XLA will constant-fold it for b > 0 at runtime.
        leftpad_cutoff = lpad - blk_start   # scalar int32 (may be negative)
        leftpad_mask   = kv_local[None, :] < leftpad_cutoff  # [1, kv_len]
        mask = mask | leftpad_mask  # [BLK, kv_len]

        scores = jnp.where(
            mask[None, None, None, :, :],   # [1,1,1,BLK,kv_len]
            jnp.full_like(scores, -1e30),
            scores,
        )

        # ── Softmax + value aggregation (all FP32) ────────────────────────
        probs  = jax.nn.softmax(scores, axis=-1)            # [B,Hkv,G,BLK,kv_len]
        out_g  = jnp.einsum(
            "bngqk,bnkd->bngqd",
            probs,
            v_blk.astype(jnp.float32),
        )  # [B, Hkv, G, BLK, D]

        # ── Reshape + cast back to input dtype ────────────────────────────
        return out_g.reshape(B, Hq, BLK, D).astype(q.dtype)  # [B, Hq, BLK, D]

    # ── Gradient checkpointing ────────────────────────────────────────────────
    map_fn = jax.checkpoint(process_block) if use_remat else process_block

    # ── fori_loop: write each block directly into pre-allocated output ────────
    # lax.map returns [n_blocks, B, Hq, BLK, D] — XLA stages the ENTIRE stack
    # in HBM before the transpose+reshape.  For large n_blocks / batch this
    # wastes memory.
    #
    # lax.fori_loop carries a single [B, Hq, S_pad, D] output buffer and uses
    # dynamic_update_slice to write each block in-place.  XLA sees one static-
    # shape buffer (same size as the final output) instead of n_blocks copies.
    out_init = jnp.zeros((B, Hq, S_pad, D), dtype=q.dtype)

    def _write_block(b, out_buf):
        blk_out = map_fn(b)                             # [B, Hq, BLK, D]
        return jax.lax.dynamic_update_slice(
            out_buf,
            blk_out,
            (0, 0, b * BLK, 0),
        )

    out = jax.lax.fori_loop(0, n_blocks, _write_block, out_init)
    # out: [B, Hq, S_pad, D] — trim padding
    return out[:, :, :S, :]


# ---------------------------------------------------------------------------
# Public API
# ---------------------------------------------------------------------------

def flash_splash_attention(
    q: jnp.ndarray,
    k: jnp.ndarray,
    v: jnp.ndarray,
    window_size: int,
    backend: Optional[str] = None,
    use_gqa: bool = True,
    start_pos: int = 0,
    block_size: int = TPU_BLOCK_SIZE,
    use_remat: bool = True,
) -> jnp.ndarray:
    """
    Block-tiled GQA-native Sliding Window Attention — main entry point.

    Automatically dispatches to the optimal kernel for the current backend:

    ┌──────────┬──────────────────────────────────────────────────────────┐
    │ Backend  │ Kernel                                                   │
    ├──────────┼──────────────────────────────────────────────────────────┤
    │ TPU v5e  │ _block_gqa_swa — block-tiled lax.map, static shapes,    │
    │          │ BF16 matmul, gradient checkpoint per block               │
    ├──────────┼──────────────────────────────────────────────────────────┤
    │ GPU/CUDA │ _gpu_flash_gqa_swa — cuDNN FlashAttention v2/v3 via      │
    │          │ jax.nn.dot_product_attention, native GQA, SWA mask       │
    │          │ Falls back to JAX XLA SDPA then explicit einsum          │
    ├──────────┼──────────────────────────────────────────────────────────┤
    │ CPU      │ _block_gqa_swa — same as TPU (block_size=64 for cache)  │
    └──────────┴──────────────────────────────────────────────────────────┘

    Parameters
    ----------
    q, k, v      : [B, Hq, S, D] / [B, Hkv, S, D]
    window_size  : causal window W
    backend      : override auto-detection ('tpu', 'gpu', 'cpu', or None)
    use_gqa      : API-compat flag; GQA is always native here
    start_pos    : KV-cache offset for decode mode (training path: leave at 0)
    block_size   : query block size BLK (128 for TPU, 64 for GPU)
    use_remat    : gradient checkpointing (recommended for training)

    Returns
    -------
    [B, Hq, S, D] — same dtype as q
    """
    _ = use_gqa
    _ = start_pos

    # ── Backend detection ─────────────────────────────────────────────────────
    active_backend = (backend or _detect_backend()).lower()
    if 'tpu' in active_backend:
        active_backend = 'tpu'
    elif 'gpu' in active_backend or 'cuda' in active_backend:
        active_backend = 'gpu'
    else:
        active_backend = 'cpu'

    # ── 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,
            v=v,
            window_size=int(window_size),
            use_remat=use_remat,
        )
    else:
        # TPU and CPU both use block-tiled JAX kernel
        # GPU_BLOCK_SIZE is used for CPU (cache-friendly), TPU_BLOCK_SIZE for TPU
        blk = block_size if active_backend == 'tpu' else GPU_BLOCK_SIZE
        return _block_gqa_swa(
            q=q,
            k=k,
            v=v,
            window_size=int(window_size),
            block_size=int(blk),
            use_remat=use_remat,
        )
def decode_swa(
    q: jnp.ndarray,
    k: jnp.ndarray,
    v: jnp.ndarray,
) -> jnp.ndarray:
    """
    Decode-time SWA for one-token generation.

    Dispatches to the optimal kernel for the current backend:
      - GPU/CUDA: cuDNN SDPA (batched GEMV, extremely fast for S=1)
      - TPU/CPU:  native GQA einsum (same as before)

    q: [B, Hq,  1, D]
    k: [B, Hkv, W, D]  (sliding window cache)
    v: [B, Hkv, W, D]

    Returns: [B, Hq, 1, D]
    """
    if q.ndim != 4 or k.ndim != 4 or v.ndim != 4:
        raise ValueError(
            f"decode_swa: expected 4-D tensors, got q={q.shape}, k={k.shape}, v={v.shape}"
        )

    B, Hq, S, D = map(int, q.shape)
    _, Hkv, W, Dk = map(int, k.shape)

    if S != 1:
        raise ValueError(f"decode_swa expects one token, got S={S}")
    if D != Dk:
        raise ValueError(f"Head-dim mismatch: q D={D}, k D={Dk}")
    if Hq % Hkv != 0:
        raise ValueError(f"Hq={Hq} must be divisible by Hkv={Hkv}")

    # ── Backend dispatch ──────────────────────────────────────────────────────
    active_backend = _detect_backend()

    if active_backend == 'gpu':
        return _gpu_decode_swa(q, k, v)

    # ── TPU/CPU: native GQA einsum (original implementation) ─────────────────
    G = Hq // Hkv
    scale = 1.0 / math.sqrt(float(D))

    q_g = q.reshape(B, Hkv, G, 1, D)

    scores = (
        jnp.einsum(
            "bngqd,bnkd->bngqk",
            q_g.astype(jnp.float32),
            k.astype(jnp.float32),
        )
        * scale
    )  # [B, Hkv, G, 1, W]

    probs = jax.nn.softmax(scores, axis=-1)

    out_g = jnp.einsum(
        "bngqk,bnkd->bngqd",
        probs,
        v.astype(jnp.float32),
    )

    return out_g.reshape(B, Hq, 1, D).astype(q.dtype)

def local_causal_attention(
    q: jnp.ndarray,
    k: jnp.ndarray,
    v: jnp.ndarray,
) -> jnp.ndarray:
    """
    Full causal attention with native GQA — no windowing.

    ⚠  Creates an O(S²) score tensor [B, Hkv, G, S, S].
       Use only for short sequences, unit tests, or reference baselines.

    q: [B, Hq, S, D]  |  k, v: [B, Hkv, S, D]
    returns: [B, Hq, S, D]
    """
    if q.ndim != 4 or k.ndim != 4 or v.ndim != 4:
        raise ValueError(
            f"local_causal_attention: expected 4-D tensors, "
            f"got q={q.shape}  k={k.shape}  v={v.shape}"
        )
    B, Hq, S, D = map(int, q.shape)
    _,  Hkv, K, _ = 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))
    q_g   = q.reshape(B, Hkv, G, S, D)
    scores = (
        jnp.einsum("bngsd,bnkd->bngsk", q_g.astype(jnp.float32), k.astype(jnp.float32))
        * scale
    )
    qi = jnp.arange(S, dtype=jnp.int32)[:, None]
    ki = jnp.arange(K, dtype=jnp.int32)[None, :]
    scores = jnp.where(ki > qi, -1e30, scores)
    probs  = jax.nn.softmax(scores, axis=-1)
    out_g  = jnp.einsum("bngsk,bnkd->bngsd", probs, v.astype(jnp.float32))
    return out_g.reshape(B, Hq, S, D).astype(q.dtype)


# ---------------------------------------------------------------------------
# Validation suite
# ---------------------------------------------------------------------------

if __name__ == "__main__":
    import sys

    PASS  = "\033[92m✓\033[0m"
    FAIL  = "\033[91m✗\033[0m"
    HDR   = "\033[1;94m"
    RST   = "\033[0m"
    failures = 0

    def section(t):  print(f"\n{HDR}{'─'*62}{RST}\n{HDR}  {t}{RST}\n{HDR}{'─'*62}{RST}")
    def ok(m):       print(f"  {PASS}  {m}")
    def fail(m):
        global failures; failures += 1
        print(f"  {FAIL}  {m}", file=sys.stderr)

    print(f"\n{HDR}{'═'*62}{RST}")
    print(f"{HDR}  Veylon Attention — Block-tiled GQA SWA — Validation{RST}")
    print(f"{HDR}{'═'*62}{RST}")

    # ── 1. Shape and dtype ───────────────────────────────────────────────────
    section("1 · Shape and dtype correctness")

    B, Hq, Hkv, S, D, W = 1, 8, 2, 128, 64, 32
    ks = jax.random.split(jax.random.PRNGKey(0), 3)
    q = jax.random.normal(ks[0], (B, Hq,  S, D), dtype=jnp.bfloat16)
    k = jax.random.normal(ks[1], (B, Hkv, S, D), dtype=jnp.bfloat16)
    v = jax.random.normal(ks[2], (B, Hkv, S, D), dtype=jnp.bfloat16)

    out = flash_splash_attention(q, k, v, window_size=W, block_size=32)

    if out.shape == (B, Hq, S, D):   ok(f"Output shape  : {out.shape}")
    else:                             fail(f"Shape wrong — expected {(B,Hq,S,D)}, got {out.shape}")
    if out.dtype == jnp.bfloat16:    ok(f"Output dtype  : {out.dtype}")
    else:                             fail(f"dtype wrong — expected bfloat16, got {out.dtype}")
    if not jnp.any(jnp.isnan(out)):  ok("No NaNs")
    else:                             fail("Output contains NaNs")

    # ── 2. Numerical agreement with reference SWA ────────────────────────────
    section("2 · Numerical agreement: block SWA ≈ reference SWA (W=S)")

    B_, Hq_, Hkv_, S_, D_ = 1, 4, 2, 48, 16
    ks = jax.random.split(jax.random.PRNGKey(1), 3)
    qn = jax.random.normal(ks[0], (B_, Hq_,  S_, D_), dtype=jnp.float32)
    kn = jax.random.normal(ks[1], (B_, Hkv_, S_, D_), dtype=jnp.float32)
    vn = jax.random.normal(ks[2], (B_, Hkv_, S_, D_), dtype=jnp.float32)

    for W_t in [4, 8, 16, S_]:
        out_blk  = flash_splash_attention(qn, kn, vn, window_size=W_t, block_size=16)
        out_ref  = local_causal_attention(qn, kn, vn) if W_t == S_ else None

        # Build reference SWA inline for each W_t
        G_ = Hq_ // Hkv_
        sc = 1.0 / math.sqrt(D_)
        qg = qn.reshape(B_, Hkv_, G_, S_, D_)
        sc_ref = jnp.einsum("bngsd,bnkd->bngsk", qg, kn) * sc
        qi = jnp.arange(S_)[:, None]; ki = jnp.arange(S_)[None, :]
        sc_ref = jnp.where((ki > qi) | (qi - ki >= W_t), -1e30, sc_ref)
        pr_ref = jax.nn.softmax(sc_ref, axis=-1)
        out_swa_ref = jnp.einsum("bngsk,bnkd->bngsd", pr_ref, vn).reshape(B_, Hq_, S_, D_)

        err = float(jnp.max(jnp.abs(out_blk - out_swa_ref)))
        if err < 1e-4:  ok(f"W={W_t:3d}  max|block - ref_swa| = {err:.2e}")
        else:           fail(f"W={W_t:3d}  MISMATCH: max err = {err:.2e}")

    # ── 3. Memory scaling ────────────────────────────────────────────────────
    section("3 · Memory scaling  (2× S → ~2× memory, not 4×)")
    print("  Shape + completion check at S = 256, 512, 1024, 2048")
    for S_t in [256, 512, 1024, 2048]:
        ks = jax.random.split(jax.random.PRNGKey(S_t), 3)
        qt = jax.random.normal(ks[0], (1, 8, S_t, 64), dtype=jnp.bfloat16)
        kt = jax.random.normal(ks[1], (1, 2, S_t, 64), dtype=jnp.bfloat16)
        vt = jax.random.normal(ks[2], (1, 2, S_t, 64), dtype=jnp.bfloat16)
        ot = flash_splash_attention(qt, kt, vt, window_size=128)
        if ot.shape == (1, 8, S_t, 64): ok(f"S={S_t:5d}  →  {ot.shape}")
        else:                            fail(f"S={S_t}  wrong shape {ot.shape}")

    # ── 4. Native GQA isolation ──────────────────────────────────────────────
    section("4 · Native GQA head isolation (no KV duplication)")
    B_, Hq_, Hkv_, S_, D_, W_ = 1, 4, 2, 32, 16, 16
    G_ = Hq_ // Hkv_
    ks  = jax.random.split(jax.random.PRNGKey(7), 3)
    qg  = jax.random.normal(ks[0], (B_, Hq_, S_, D_), dtype=jnp.float32)
    kgq = jax.random.normal(ks[1], (B_, Hkv_, S_, D_), dtype=jnp.float32)
    vgq = jax.random.normal(ks[2], (B_, Hkv_, S_, D_), dtype=jnp.float32)
    base  = flash_splash_attention(qg, kgq,                               vgq,                               window_size=W_, block_size=16)
    zero0 = flash_splash_attention(qg, kgq.at[:,0].set(0.), vgq.at[:,0].set(0.), window_size=W_, block_size=16)
    changed   = not jnp.allclose(base[:, :G_],  zero0[:, :G_],  atol=1e-4)
    unchanged = jnp.allclose(    base[:, G_:],  zero0[:, G_:],  atol=1e-4)
    ok(f"Q heads 0..{G_-1} changed when KV head 0 zeroed")   if changed   else fail("GQA dependency broken")
    ok(f"Q heads {G_}..{Hq_-1} unchanged (correct isolation)") if unchanged else fail("Cross-group contamination")

    # ── 5. Causality ─────────────────────────────────────────────────────────
    section("5 · Causality")
    B_, Hq_, Hkv_, S_, D_, W_ = 1, 4, 2, 24, 8, 8
    pivot = S_ // 2
    ks = jax.random.split(jax.random.PRNGKey(42), 3)
    qc = jax.random.normal(ks[0], (B_, Hq_, S_, D_), dtype=jnp.float32)
    kc = jax.random.normal(ks[1], (B_, Hkv_, S_, D_), dtype=jnp.float32)
    vc = jax.random.normal(ks[2], (B_, Hkv_, S_, D_), dtype=jnp.float32)
    out_base = flash_splash_attention(qc, kc, vc, window_size=W_, block_size=8)
    out_corr = flash_splash_attention(qc,
                                      kc.at[:,:,pivot:].set(999.),
                                      vc.at[:,:,pivot:].set(999.),
                                      window_size=W_, block_size=8)
    if jnp.allclose(out_base[:,:,:pivot], out_corr[:,:,:pivot], atol=1e-5):
        ok(f"Tokens 0..{pivot-1} unaffected by corruption of tokens {pivot}+")
    else:
        fail("Causality violated — future tokens leaked into past outputs")

    # ── 6. Window boundary ───────────────────────────────────────────────────
    section("6 · Window boundary  (no leakage past W tokens)")
    B_, Hq_, Hkv_, S_, D_, W_ = 1, 2, 1, 24, 8, 4
    ks = jax.random.split(jax.random.PRNGKey(11), 3)
    qw = jax.random.normal(ks[0], (B_, Hq_, S_, D_), dtype=jnp.float32)
    kw = jax.random.normal(ks[1], (B_, Hkv_, S_, D_), dtype=jnp.float32)
    vw = jax.random.normal(ks[2], (B_, Hkv_, S_, D_), dtype=jnp.float32)
    out_w  = flash_splash_attention(qw, kw, vw, window_size=W_, block_size=4)
    out_wm = flash_splash_attention(qw,
                                    kw.at[:,:,0].set(999.),
                                    vw.at[:,:,0].set(999.),
                                    window_size=W_, block_size=4)
    pivot_w = W_  # first token where token 0 is outside the window
    if jnp.allclose(out_w[:,:,pivot_w:], out_wm[:,:,pivot_w:], atol=1e-5):
        ok(f"Tokens {pivot_w}+ unaffected by modifying token 0  (W={W_} boundary)")
    else:
        fail(f"Window boundary violated — token 0 leaked into token {pivot_w}+")

    # ── 7. Non-power-of-2 sequence length ────────────────────────────────────
    section("7 · Non-power-of-2 sequence lengths  (S=100, 200, 500)")
    for S_t in [100, 200, 500]:
        ks = jax.random.split(jax.random.PRNGKey(S_t+1), 3)
        qt = jax.random.normal(ks[0], (1, 4, S_t, 16), dtype=jnp.float32)
        kt = jax.random.normal(ks[1], (1, 2, S_t, 16), dtype=jnp.float32)
        vt = jax.random.normal(ks[2], (1, 2, S_t, 16), dtype=jnp.float32)
        ot = flash_splash_attention(qt, kt, vt, window_size=32, block_size=32)
        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.")
    else:             print(f"  {FAIL}  {failures} test(s) failed.", file=sys.stderr)
    print(f"{HDR}{'═'*62}{RST}\n")
    sys.exit(failures)