# the state monoid: combine two partial (m, l, acc) summaries exactly
def combine(a, b):
    m1, l1, acc1 = a
    m2, l2, acc2 = b
    m = jnp.maximum(m1, m2)
    a1, a2 = jnp.exp(m1 - m), jnp.exp(m2 - m)
    return (m, l1 * a1 + l2 * a2, acc1 * a1[..., None] + acc2 * a2[..., None])

def summarize(s_block, v_block):
    m = jnp.max(s_block, axis=-1)
    p = jnp.exp(s_block - m[..., None])
    return (m, jnp.sum(p, axis=-1), p @ v_block)

def streaming_attention(q, k, v, block=64):
    s = q @ k.T
    n = s.shape[-1]
    state = summarize(s[:, :block], v[:block])
    for j in range(block, n, block):
        state = combine(state, summarize(s[:, j:j+block], v[j:j+block]))
    m, l, acc = state
    return acc / l[..., None]

import random

def reduce_bracketed(states, order):
    states = list(states)
    for idx in order:
        merged = combine(states[idx], states[idx + 1])
        states[idx:idx + 2] = [merged]
    return states[0]

s = q @ k.T
blocks = [summarize(s[:, j:j+64], v[j:j+64]) for j in range(0, 512, 64)]
base = reduce_bracketed(blocks, [0] * (len(blocks) - 1))  # left fold
for trial in range(20):
    order = [random.randrange(n) for n in range(len(blocks) - 1, 0, -1)]
    m, l, acc = reduce_bracketed(blocks, order)
    check(f"bracketing {trial}", acc / l[..., None], base[2] / base[1][..., None], tol=1e-5)
print("combine is associative under 20 random bracketings")

# one query block per grid step; the kv axis streams inside the kernel
def flash_kernel(q_ref, k_ref, v_ref, o_ref, *, block_kv):
    q = q_ref[...].astype(jnp.float32)
    n_kv = k_ref.shape[0]

    def step(j, state):
        m, l, acc = state
        kb = jax.lax.dynamic_slice_in_dim(k_ref[...], j * block_kv, block_kv).astype(jnp.float32)
        vb = jax.lax.dynamic_slice_in_dim(v_ref[...], j * block_kv, block_kv).astype(jnp.float32)
        s = q @ kb.T
        m_new = jnp.maximum(m, jnp.max(s, axis=-1))
        alpha = jnp.exp(m - m_new)
        p = jnp.exp(s - m_new[:, None])
        l_new = l * alpha + jnp.sum(p, axis=-1)
        acc_new = acc * alpha[:, None] + p @ vb
        return m_new, l_new, acc_new

    m0 = jnp.full((q.shape[0],), -jnp.inf, jnp.float32)
    l0 = jnp.zeros((q.shape[0],), jnp.float32)
    acc0 = jnp.zeros((q.shape[0], v_ref.shape[1]), jnp.float32)
    m, l, acc = jax.lax.fori_loop(0, n_kv // block_kv, step, (m0, l0, acc0))
    o_ref[...] = (acc / l[:, None]).astype(o_ref.dtype)

# reference math and the custom_vjp pair; forward saves only (O, m, l) style residuals
def attention_ref(q, k, v):
    return jax.nn.softmax(q @ k.T, axis=-1) @ v

@jax.custom_vjp
def attention(q, k, v):
    return attention_ref(q, k, v)

def fwd(q, k, v):
    s = q @ k.T
    m = jnp.max(s, axis=-1, keepdims=True)
    p = jnp.exp(s - m)
    l = jnp.sum(p, axis=-1, keepdims=True)
    o = (p / l) @ v
    return o, (q, k, v, o, m, l)

def bwd(res, do):
    q, k, v, o, m, l = res
    # recompute P blockwise in the real kernel; whole-array here for clarity
    p = jnp.exp(q @ k.T - m) / l
    dv = p.T @ do
    d = jnp.sum(do * o, axis=-1, keepdims=True)   # the identity that kills the Jacobian
    dp = do @ v.T
    ds = p * (dp - d)
    return ds @ k, ds.T @ q, dv

attention.defvjp(fwd, bwd)

q = jax.random.normal(jax.random.key(0), (128, 64))
k = jax.random.normal(jax.random.key(1), (256, 64))
v = jax.random.normal(jax.random.key(2), (256, 64))

def loss_custom(q, k, v): return jnp.sum(jnp.tanh(attention(q, k, v)))
def loss_ref(q, k, v): return jnp.sum(jnp.tanh(attention_ref(q, k, v)))

g1 = jax.grad(loss_custom, argnums=(0, 1, 2))(q, k, v)
g2 = jax.grad(loss_ref, argnums=(0, 1, 2))(q, k, v)
for name, a, b in zip("qkv", g1, g2):
    check(f"d{name}", a, b, tol=1e-5)

import functools

def causal_flash_kernel(q_ref, k_ref, v_ref, o_ref, *, block_q, block_kv):
    qi = pl.program_id(0)
    q = q_ref[...].astype(jnp.float32)
    row0 = qi * block_q

    def step(j, state):
        m, l, acc = state
        kb = jax.lax.dynamic_slice_in_dim(k_ref[...], j * block_kv, block_kv).astype(jnp.float32)
        vb = jax.lax.dynamic_slice_in_dim(v_ref[...], j * block_kv, block_kv).astype(jnp.float32)
        s = q @ kb.T
        cols = j * block_kv + jax.lax.broadcasted_iota(jnp.int32, s.shape, 1)
        rows = row0 + jax.lax.broadcasted_iota(jnp.int32, s.shape, 0)
        s = jnp.where(cols <= rows, s, -jnp.inf)      # edge block: mask inside
        m_new = jnp.maximum(m, jnp.max(s, axis=-1))
        alpha = jnp.exp(m - m_new)
        p = jnp.exp(s - m_new[:, None])
        l_new = l * alpha + jnp.sum(p, axis=-1)
        acc_new = acc * alpha[:, None] + p @ vb
        return m_new, l_new, acc_new

    # the causal structure IS the loop bound: blocks past the diagonal never run
    n_live = (row0 + block_q + block_kv - 1) // block_kv
    m0 = jnp.full((block_q,), -jnp.inf, jnp.float32)
    l0 = jnp.zeros((block_q,), jnp.float32)
    acc0 = jnp.zeros((block_q, v_ref.shape[1]), jnp.float32)
    m, l, acc = jax.lax.fori_loop(0, n_live, step, (m0, l0, acc0))
    o_ref[...] = (acc / l[:, None]).astype(o_ref.dtype)