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
from functools import partial
from typing import Optional

import jax
import jax.numpy as jnp


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


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

        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)

    # ── 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':
        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}")

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