atomic-ops 0.1.0__py3-none-any.whl

atomic_ops/__init__.py

@@ -0,0 +1,34 @@
+"""
+Atomic Ops — Fused Gated DeltaNet-2 kernels for TPU v5e (Pallas/JAX).
+Ported from NVlabs DeltaNet Triton kernels.
+"""
+from importlib.metadata import version as _version, PackageNotFoundError as _PkgNotFound
+
+from .configs import KernelConfig, KAGGLE_SMALL, KAGGLE_MEDIUM, KAGGLE_LARGE, DEFAULT_CONFIG
+from .utils import is_tpu_available, estimate_memory, get_recommended_config
+from .fallback import gdn2_forward, gdn2_forward_trainable
+from .gdn2_fwd import gdn2_pallas_forward
+from .gdn2_pipeline import gdn2_pallas_forward_trainable
+from .reference import gdn2_chunked_wy_reference, gdn2_token_serial_reference
+
+try:
+    __version__ = _version("atomic_ops")
+except _PkgNotFound:
+    __version__ = "0.0.0.dev0"
+
+__all__ = [
+    "KernelConfig",
+    "KAGGLE_SMALL",
+    "KAGGLE_MEDIUM",
+    "KAGGLE_LARGE",
+    "DEFAULT_CONFIG",
+    "is_tpu_available",
+    "estimate_memory",
+    "get_recommended_config",
+    "gdn2_forward",
+    "gdn2_forward_trainable",
+    "gdn2_pallas_forward",
+    "gdn2_pallas_forward_trainable",
+    "gdn2_chunked_wy_reference",
+    "gdn2_token_serial_reference",
+]

atomic_ops/configs.py

@@ -0,0 +1,136 @@
+"""
+Central configuration, numerical-safety helpers, and shape validation
+shared across the forward/backward Pallas kernels.
+"""
+from __future__ import annotations
+
+import dataclasses as dc
+import os
+
+import jax
+import jax.numpy as jnp
+
+_HIGHEST = jax.lax.Precision.HIGHEST
+
+
+@dc.dataclass(frozen=True)
+class KernelConfig:
+    bt: int = 256
+    bc: int = 128
+    mb: int = 16
+    clip: float = 1e4
+    wy_eps: float = 0.0
+
+    @property
+    def n_sub(self) -> int:
+        return self.bt // self.bc
+
+    @property
+    def n_micro(self) -> int:
+        return self.bc // self.mb
+
+    def __post_init__(self):
+        if self.bt % self.bc != 0:
+            raise ValueError(f"bt={self.bt} must be divisible by bc={self.bc}")
+        if self.bt != 2 * self.bc:
+            raise ValueError(
+                f"bt={self.bt} must equal 2*bc (top-level WY-solve split "
+                f"supports only the 2-block case); got bc={self.bc}. "
+                f"Vary `mb` instead of `bc` to change solver granularity -- "
+                f"bc/mb do not affect numerical accuracy of the solve, only "
+                f"its speed (see grid_bt_bc_condition_diag.py Part 3)."
+            )
+        if self.bc % self.mb != 0:
+            raise ValueError(f"bc={self.bc} must be divisible by mb={self.mb}")
+        if not (0.0 <= self.wy_eps < 1.0):
+            raise ValueError(f"wy_eps={self.wy_eps} must be in [0, 1)")
+
+
+KAGGLE_SMALL = KernelConfig(bt=128, bc=64, mb=16, clip=1e4, wy_eps=1e-3)
+KAGGLE_MEDIUM = KernelConfig(bt=256, bc=128, mb=16, clip=1e4, wy_eps=1e-3)
+KAGGLE_LARGE = KernelConfig(bt=256, bc=128, mb=16, clip=5e3, wy_eps=1e-3)
+DEFAULT_CONFIG = KAGGLE_MEDIUM
+
+
+def sanitize(x, config: KernelConfig = DEFAULT_CONFIG):
+    """Standard clip + nan_to_num defense used at every kernel boundary."""
+    c = config.clip
+    return jnp.nan_to_num(jnp.clip(x, -c, c), nan=0.0, posinf=c, neginf=-c)
+
+
+def sanitize_h0(h0, config: KernelConfig = DEFAULT_CONFIG):
+    return sanitize(h0, config)
+
+
+def clip_acc(x, config: KernelConfig = DEFAULT_CONFIG):
+    """Same as sanitize; separate name kept for call-site clarity in
+    read-modify-write accumulation loops (see gdn2_bwd.py, Kernel B4)."""
+    return sanitize(x, config)
+
+
+def _reshape_to_chunks(t, bsz, n_chunks, H, D, bt):
+    t = t.reshape(bsz, n_chunks, bt, H, D)
+    return jnp.moveaxis(t, (1, 3), (2, 1))
+
+
+def _reshape_from_chunks(t, bsz, n_chunks, bt, H, D):
+    t2 = jnp.moveaxis(t, (1, 2, 3), (3, 1, 2))
+    return t2.reshape(bsz, n_chunks * bt, H, D)
+
+
+_GDN2_FWD_DIAG = os.environ.get("GDN2_FWD_DIAG", "0") == "1"
+_LARGE_THRESHOLD = 1e6  # suspiciously large but still-finite trigger level
+
+
+def _stage_diag(tag: str, x):
+    """No-op unless GDN2_FWD_DIAG=1. Diagnostic only -- never changes x."""
+    if not _GDN2_FWD_DIAG:
+        return x
+
+    finite_mask = jnp.isfinite(x)
+    all_finite = jnp.all(finite_mask)
+    n_nonfinite = jnp.sum(jnp.logical_not(finite_mask))
+    safe_x = jnp.where(finite_mask, x, 0.0)
+    max_abs = jnp.max(jnp.abs(safe_x))
+
+    def _report_nonfinite():
+        jax.debug.print(
+            "[GDN2-FWD-DIAG] non-finite at " + tag + ": n_nonfinite={n} max_abs_finite={m:.3e}",
+            n=n_nonfinite, m=max_abs,
+        )
+
+    def _report_large():
+        jax.debug.print(
+            "[GDN2-FWD-DIAG] suspiciously large (still finite) at " + tag + ": max_abs={m:.3e}",
+            m=max_abs,
+        )
+
+    jax.lax.cond(
+        jnp.logical_not(all_finite),
+        _report_nonfinite,
+        lambda: jax.lax.cond(max_abs > _LARGE_THRESHOLD, _report_large, lambda: None),
+    )
+    return x
+
+
+def validate_inputs(q, k, v, w, b, g, scale, h0, config: KernelConfig):
+    """Shape/dtype sanity checks shared by forward and backward entry points."""
+    if q.ndim != 4:
+        raise ValueError(f"q must be (batch, seq_len, heads, d_head); got shape {q.shape}")
+    bsz, L, H, D = q.shape
+    if D != 128:
+        raise ValueError(f"Kernels assume d_head=128 (MXU tile); got D={D}.")
+    if L % config.bt != 0:
+        raise ValueError(f"seq_len={L} must be divisible by config.bt={config.bt}.")
+
+    for name, t in (("k", k), ("v", v), ("w", w), ("b", b), ("g", g)):
+        if t.shape != q.shape:
+            raise ValueError(f"{name}.shape={t.shape} must match q.shape={q.shape}")
+
+    if h0 is not None:
+        expected_h0 = (bsz, H, D, D)
+        if h0.shape != expected_h0:
+            raise ValueError(f"h0.shape={h0.shape} must be {expected_h0}")
+
+    n_chunks = L // config.bt
+    return bsz, L, H, D, n_chunks

atomic_ops/fallback.py

@@ -0,0 +1,22 @@
+from __future__ import annotations
+import jax.numpy as jnp
+from .utils import is_tpu_available
+from .gdn2_fwd import gdn2_pallas_forward as _pallas_fwd
+from .gdn2_pipeline import gdn2_pallas_forward_trainable as _pallas_trainable
+from .reference import gdn2_chunked_wy_reference
+
+
+def gdn2_forward(q, k, v, w, b, g, scale, h0=None, config=None):
+    from .configs import DEFAULT_CONFIG
+    config = config or DEFAULT_CONFIG
+    if is_tpu_available() and q.shape[-1] == 128:
+        return _pallas_fwd(q, k, v, w, b, g, scale, h0=h0, config=config)
+    return gdn2_chunked_wy_reference(q, k, v, g, b, w, scale, chunk_size=config.bt, h0=h0)
+
+
+def gdn2_forward_trainable(q, k, v, w, b, g, scale, h0=None, config=None):
+    from .configs import DEFAULT_CONFIG
+    config = config or DEFAULT_CONFIG
+    if is_tpu_available() and q.shape[-1] == 128:
+        return _pallas_trainable(q, k, v, w, b, g, scale, h0=h0, config=config)
+    return gdn2_chunked_wy_reference(q, k, v, g, b, w, scale, chunk_size=config.bt, h0=h0)

atomic_ops/gdn2_bwd.py

@@ -0,0 +1,342 @@
+"""
+Backward kernels: B1 (state) -> B2 (dAqk/dv) -> B3 (WY/dqkg) -> B4 (intra) -> B5 (reverse cumsum).
+"""
+from __future__ import annotations
+
+import jax
+import jax.numpy as jnp
+from jax.experimental import pallas as pl
+from jax.experimental.pallas import tpu as pltpu
+
+from .configs import (
+    KernelConfig, DEFAULT_CONFIG, sanitize, clip_acc,
+    _reshape_to_chunks as _r2c, _reshape_from_chunks as _r2f,
+)
+
+_HIGHEST = jax.lax.Precision.HIGHEST
+
+
+# ---------- B5 ----------
+def reverse_cumsum_bwd(dgc, chunk_size: int, config: KernelConfig = DEFAULT_CONFIG):
+    C = chunk_size
+    idx = jnp.arange(C)
+    triu_ones = (idx[:, None] <= idx[None, :]).astype(jnp.float32)
+    dg_raw = jnp.einsum("ij,...jd->...id", triu_ones, dgc.astype(jnp.float32), precision=_HIGHEST)
+    return sanitize(dg_raw, config)
+
+
+# ---------- B1 ----------
+def gdn2_dhu_backward(do, dv_partial, w_pseudo, qg, kg, gc_last, scale, dht=None,
+                       config: KernelConfig = DEFAULT_CONFIG):
+    bsz, H, n_chunks, BT, D = qg.shape
+    if dht is None:
+        dht = jnp.zeros((bsz, H, D, D), dtype=jnp.float32)
+    dht = sanitize(dht, config)
+
+    to_scan = tuple(jnp.moveaxis(x, 2, 0) for x in (do, dv_partial, w_pseudo, qg, kg, gc_last))
+
+    def step(dh_carry, inputs):
+        do_c, dvp_c, wp_c, qg_c, kg_c, gclast_c = inputs
+        decay_c = jnp.exp(gclast_c)[..., None]
+
+        dqh = scale * do_c
+        contrib_from_output = jnp.einsum("bhid,bhiv->bhdv", qg_c, dqh, precision=_HIGHEST)
+        contrib_from_state = dh_carry * decay_c
+
+        dv_write = jnp.einsum("bhid,bhdv->bhiv", kg_c, dh_carry, precision=_HIGHEST)
+        dv_new_c = dvp_c + dv_write
+        dv_new_c = sanitize(dv_new_c, config)
+
+        contrib_from_vnew = -jnp.einsum("bhjd,bhjv->bhdv", wp_c, dv_new_c, precision=_HIGHEST)
+
+        dh_pre_c = contrib_from_output + contrib_from_state + contrib_from_vnew
+        dh_pre_c = sanitize(dh_pre_c, config)
+        return dh_pre_c, (dh_pre_c, dv_new_c)
+
+    dh0, (dh_all_rev, dv_all_rev) = jax.lax.scan(step, dht, to_scan, reverse=True)
+    dh_all = jnp.moveaxis(dh_all_rev, 0, 2)
+    dv_all = jnp.moveaxis(dv_all_rev, 0, 2)
+    return dh_all, dh0, dv_all
+
+
+# ---------- B2 ----------
+def _kernel_b2_body(aqk_ref, vnew_ref, do_ref, daqk_ref, dvnew_ref, *, bt: int, config: KernelConfig):
+    Aqk = aqk_ref[0, 0, 0].astype(jnp.float32)
+    v_new = vnew_ref[0, 0, 0].astype(jnp.float32)
+    do = do_ref[0, 0, 0].astype(jnp.float32)
+
+    idx = jnp.arange(bt)
+    causal = (idx[:, None] >= idx[None, :]).astype(jnp.float32)
+
+    dAqk = jnp.dot(do, v_new.T, precision=_HIGHEST) * causal
+    dv_new = jnp.dot(Aqk.T, do, precision=_HIGHEST)
+
+    daqk_ref[0, 0, 0] = sanitize(dAqk, config)
+    dvnew_ref[0, 0, 0] = sanitize(dv_new, config)
+
+
+def dav_backward_pallas(Aqk, v_new, do, config: KernelConfig = DEFAULT_CONFIG):
+    bsz, H, n_chunks, _BT, D = v_new.shape
+    grid = (bsz, H, n_chunks)
+    aqk_spec = pl.BlockSpec((1, 1, 1, config.bt, config.bt), lambda i, h, c: (i, h, c, 0, 0))
+    io_spec = pl.BlockSpec((1, 1, 1, config.bt, D), lambda i, h, c: (i, h, c, 0, 0))
+
+    dAqk, dv_new = pl.pallas_call(
+        lambda *refs: _kernel_b2_body(*refs, bt=config.bt, config=config),
+        grid=grid,
+        in_specs=[aqk_spec, io_spec, io_spec],
+        out_specs=[aqk_spec, io_spec],
+        out_shape=[
+            jax.ShapeDtypeStruct(Aqk.shape, jnp.float32),
+            jax.ShapeDtypeStruct(v_new.shape, jnp.float32),
+        ],
+        compiler_params=pltpu.CompilerParams(vmem_limit_bytes=64 * 1024 * 1024),
+    )(Aqk, v_new, do)
+    return dAqk, dv_new
+
+
+# ---------- B3 ----------
+def _kernel_b3_body(q_ref, k_ref, b_ref, w_ref, v_ref, gc_ref, a_ref, akk_ref,
+                     hpre_ref, vnew_ref, do_ref, dv_ref, dhnext_ref,
+                     dq_ref, dk_ref, db_ref, dw_ref, dvraw_ref, dgc_ref, dakk_ref,
+                     *, scale: float, bt: int, wy_eps: float, config: KernelConfig):
+    q_c = q_ref[0, 0, 0].astype(jnp.float32)
+    k_c = k_ref[0, 0, 0].astype(jnp.float32)
+    b_c = b_ref[0, 0, 0].astype(jnp.float32)
+    w_c = w_ref[0, 0, 0].astype(jnp.float32)
+    v_c = v_ref[0, 0, 0].astype(jnp.float32)
+    gc = gc_ref[0, 0, 0].astype(jnp.float32)
+    A = a_ref[0, 0, 0].astype(jnp.float32)
+    h_pre = hpre_ref[0, 0, 0].astype(jnp.float32)
+    v_new = vnew_ref[0, 0, 0].astype(jnp.float32)
+    do = do_ref[0, 0, 0].astype(jnp.float32)
+    dv = dv_ref[0, 0, 0].astype(jnp.float32)
+    dh_next = dhnext_ref[0, 0, 0].astype(jnp.float32)
+
+    C = bt
+    gc_last = gc[C - 1]
+
+    kb_decayed = b_c * k_c * jnp.exp(gc)
+    kg = k_c * jnp.exp(gc_last[None, :] - gc)
+    qg = q_c * jnp.exp(gc)
+    wv = w_c * v_c
+
+    dqh_up = scale * do
+    dqg = jnp.dot(dqh_up, h_pre.T, precision=_HIGHEST)
+
+    dwh = -dv
+    dw_pseudo = jnp.dot(dwh, h_pre.T, precision=_HIGHEST)
+    du = dv
+
+    dkg = jnp.dot(v_new, dh_next.T, precision=_HIGHEST)
+
+    dA_from_w = jnp.dot(dw_pseudo, kb_decayed.T, precision=_HIGHEST)
+    dkb_decayed = jnp.dot(A.T, dw_pseudo, precision=_HIGHEST)
+
+    dA_from_u = jnp.dot(du, wv.T, precision=_HIGHEST)
+    dwv = jnp.dot(A.T, du, precision=_HIGHEST)
+
+    dA_total = dA_from_w + dA_from_u
+    dA_total = sanitize(dA_total, config)
+
+    idx = jnp.arange(C)
+    strict = (idx[:, None] > idx[None, :]).astype(jnp.float32)
+
+    tmp = jnp.dot(dA_total, A.T, precision=_HIGHEST)
+    tmp = sanitize(tmp, config)
+    dAkk_raw = -jnp.dot(A.T, tmp, precision=_HIGHEST)
+    dAkk_raw = dAkk_raw * (1.0 - wy_eps)
+    dAkk = dAkk_raw * strict
+
+    dk_from_kb = dkb_decayed * jnp.exp(gc) * b_c
+    db = dkb_decayed * jnp.exp(gc) * k_c
+    dgc_from_kb = dkb_decayed * kb_decayed
+
+    dx = dkg * kg
+    dk_from_kg = dkg * jnp.exp(gc_last[None, :] - gc)
+    dgc_from_kg = -dx
+    dgc_last_contrib = jnp.sum(dx, axis=0)
+
+    dq = dqg * jnp.exp(gc)
+    dgc_from_qg = dqg * qg
+
+    dw = dwv * v_c
+    dv_raw = dwv * w_c
+
+    dk = dk_from_kb + dk_from_kg
+    dgc = dgc_from_kb + dgc_from_qg + dgc_from_kg
+
+    decay_h_row = jnp.exp(gc_last)
+    dgc_last_from_decay = decay_h_row * jnp.sum(dh_next * h_pre, axis=-1)
+    dgc_last_total = dgc_last_contrib + dgc_last_from_decay
+
+    row_mask = (idx == (C - 1)).astype(jnp.float32)[:, None]
+    dgc = dgc + row_mask * dgc_last_total[None, :]
+
+    dq_ref[0, 0, 0] = sanitize(dq, config)
+    dk_ref[0, 0, 0] = sanitize(dk, config)
+    db_ref[0, 0, 0] = sanitize(db, config)
+    dw_ref[0, 0, 0] = sanitize(dw, config)
+    dvraw_ref[0, 0, 0] = sanitize(dv_raw, config)
+    dakk_ref[0, 0, 0] = sanitize(dAkk, config)
+    dgc_ref[0, 0, 0] = sanitize(dgc, config)
+
+
+def wy_dqkg_backward_pallas(q, k, b, w, v, gc, A, Akk, h_pre_all, v_new_all,
+                             do, dv, dh_next_all, scale, config: KernelConfig = DEFAULT_CONFIG):
+    bsz, H, n_chunks, _BT, D = q.shape
+    grid = (bsz, H, n_chunks)
+
+    io_spec = pl.BlockSpec((1, 1, 1, config.bt, D), lambda i, h, c: (i, h, c, 0, 0))
+    score_spec = pl.BlockSpec((1, 1, 1, config.bt, config.bt), lambda i, h, c: (i, h, c, 0, 0))
+    h_spec = pl.BlockSpec((1, 1, 1, D, D), lambda i, h, c: (i, h, c, 0, 0))
+
+    dq, dk, db, dw, dv_raw, dgc, dAkk = pl.pallas_call(
+        lambda *refs: _kernel_b3_body(*refs, scale=scale, bt=config.bt, wy_eps=config.wy_eps, config=config),
+        grid=grid,
+        in_specs=[io_spec, io_spec, io_spec, io_spec, io_spec, io_spec,
+                   score_spec, score_spec, h_spec, io_spec, io_spec, io_spec, h_spec],
+        out_specs=[io_spec, io_spec, io_spec, io_spec, io_spec, io_spec, score_spec],
+        out_shape=[
+            jax.ShapeDtypeStruct((bsz, H, n_chunks, config.bt, D), jnp.float32),
+            jax.ShapeDtypeStruct((bsz, H, n_chunks, config.bt, D), jnp.float32),
+            jax.ShapeDtypeStruct((bsz, H, n_chunks, config.bt, D), jnp.float32),
+            jax.ShapeDtypeStruct((bsz, H, n_chunks, config.bt, D), jnp.float32),
+            jax.ShapeDtypeStruct((bsz, H, n_chunks, config.bt, D), jnp.float32),
+            jax.ShapeDtypeStruct((bsz, H, n_chunks, config.bt, D), jnp.float32),
+            jax.ShapeDtypeStruct((bsz, H, n_chunks, config.bt, config.bt), jnp.float32),
+        ],
+        compiler_params=pltpu.CompilerParams(vmem_limit_bytes=100 * 1024 * 1024),
+    )(q, k, b, w, v, gc, A, Akk, h_pre_all, v_new_all, do, dv, dh_next_all)
+
+    return dict(dq=dq, dk=dk, db=db, dw=dw, dv_raw=dv_raw, dgc=dgc, dAkk=dAkk)
+
+
+# ---------- B4 ----------
+def _dL_pair_sum(dM, edecay, R):
+    tmp = dM[:, :, None] * edecay
+    tmp = tmp * R[None, :, :]
+    return jnp.sum(tmp, axis=1)
+
+
+def _dR_pair_sum(dM, edecay, L):
+    tmp = dM[:, :, None] * edecay
+    tmp = tmp * L[:, None, :]
+    return jnp.sum(tmp, axis=0)
+
+
+def _dgc_pair_sum(dM, edecay, L, R, clipmask):
+    weight = dM[:, :, None] * L[:, None, :] * R[None, :, :] * edecay * clipmask
+    dgc_i = jnp.sum(weight, axis=1)
+    dgc_j = -jnp.sum(weight, axis=0)
+    return dgc_i, dgc_j
+
+
+def _kernel_b4_body(q_ref, k_ref, b_ref, g_ref, daqk_ref, dakk_ref,
+                     dq_ref, dk_ref, db_ref, dgc_ref, *, scale: float, bt: int, bc: int, n_sub: int,
+                     config: KernelConfig):
+    q_full = q_ref[0, 0, 0].astype(jnp.float32)
+    k_full = k_ref[0, 0, 0].astype(jnp.float32)
+    b_full = b_ref[0, 0, 0].astype(jnp.float32)
+    g_raw = g_ref[0, 0, 0].astype(jnp.float32)
+    dAqk = daqk_ref[0, 0, 0].astype(jnp.float32)
+    dAkk = dakk_ref[0, 0, 0].astype(jnp.float32)
+
+    bt_idx = jnp.arange(bt)
+    tril_ones_bt = (bt_idx[:, None] >= bt_idx[None, :]).astype(jnp.float32)
+    gc = jnp.dot(tril_ones_bt, g_raw, precision=_HIGHEST)
+
+    bk_full = b_full * k_full
+
+    dq_ref[0, 0, 0] = jnp.zeros_like(q_full)
+    dk_ref[0, 0, 0] = jnp.zeros_like(k_full)
+    db_ref[0, 0, 0] = jnp.zeros_like(k_full)
+    dgc_ref[0, 0, 0] = jnp.zeros_like(g_raw)
+
+    for si in range(n_sub):
+        for sj in range(si + 1):
+            i0, i1 = si * bc, (si + 1) * bc
+            j0, j1 = sj * bc, (sj + 1) * bc
+
+            q_i = q_full[i0:i1]
+            k_i = k_full[i0:i1]
+            k_j = k_full[j0:j1]
+            b_i = b_full[i0:i1]
+            bk_i = bk_full[i0:i1]
+            gc_i = gc[i0:i1]
+            gc_j = gc[j0:j1]
+
+            dM_qk = dAqk[i0:i1, j0:j1]
+            dM_kk = dAkk[i0:i1, j0:j1]
+            if si == sj:
+                idx = jnp.arange(bc)
+                causal = (idx[:, None] >= idx[None, :]).astype(jnp.float32)
+                strict = (idx[:, None] > idx[None, :]).astype(jnp.float32)
+                dM_qk = dM_qk * causal
+                dM_kk = dM_kk * strict
+
+            decay_diff = gc_i[:, None, :] - gc_j[None, :, :]
+            clipmask = ((decay_diff >= -20.0) & (decay_diff <= 20.0)).astype(jnp.float32)
+            edecay = jnp.exp(jnp.clip(decay_diff, -20.0, 20.0))
+
+            L_qk = scale * q_i
+            R_qk = k_j
+            dL_qk = _dL_pair_sum(dM_qk, edecay, R_qk)
+            dR_qk = _dR_pair_sum(dM_qk, edecay, L_qk)
+            dgc_i_qk, dgc_j_qk = _dgc_pair_sum(dM_qk, edecay, L_qk, R_qk, clipmask)
+
+            L_kk = bk_i
+            R_kk = k_j
+            dL_kk = _dL_pair_sum(dM_kk, edecay, R_kk)
+            dR_kk = _dR_pair_sum(dM_kk, edecay, L_kk)
+            dgc_i_kk, dgc_j_kk = _dgc_pair_sum(dM_kk, edecay, L_kk, R_kk, clipmask)
+
+            dq_ref[0, 0, 0, i0:i1] = clip_acc(dq_ref[0, 0, 0, i0:i1] + dL_qk * scale, config)
+            db_ref[0, 0, 0, i0:i1] = clip_acc(db_ref[0, 0, 0, i0:i1] + dL_kk, config)
+            dk_ref[0, 0, 0, j0:j1] = clip_acc(dk_ref[0, 0, 0, j0:j1] + dR_qk + dR_kk, config)
+            dgc_ref[0, 0, 0, i0:i1] = clip_acc(dgc_ref[0, 0, 0, i0:i1] + dgc_i_qk + dgc_i_kk, config)
+            dgc_ref[0, 0, 0, j0:j1] = clip_acc(dgc_ref[0, 0, 0, j0:j1] + dgc_j_qk + dgc_j_kk, config)
+
+    dbk_final = db_ref[0, 0, 0]
+    dk_final = dk_ref[0, 0, 0] + dbk_final * b_full
+    db_final = dbk_final * k_full
+    dq_final = dq_ref[0, 0, 0]
+    dgc_final = dgc_ref[0, 0, 0]
+
+    dq_ref[0, 0, 0] = sanitize(dq_final, config)
+    dk_ref[0, 0, 0] = sanitize(dk_final, config)
+    db_ref[0, 0, 0] = sanitize(db_final, config)
+    dgc_ref[0, 0, 0] = sanitize(dgc_final, config)
+
+def intra_backward_pallas(dAqk, dAkk, q, k, b, g, scale, config: KernelConfig = DEFAULT_CONFIG, interpret: bool = False):
+    bsz, L, H, D = q.shape
+    n_chunks = L // config.bt
+
+    def reshape_in(t):
+        return _r2c(t, bsz, n_chunks, H, D, config.bt)
+
+    q_r, k_r, b_r, g_r = map(reshape_in, (q, k, b, g))
+
+    grid = (bsz, H, n_chunks)
+    io_spec = pl.BlockSpec((1, 1, 1, config.bt, D), lambda i, h, c: (i, h, c, 0, 0))
+    score_spec = pl.BlockSpec((1, 1, 1, config.bt, config.bt), lambda i, h, c: (i, h, c, 0, 0))
+
+    dq, dk, db, dgc = pl.pallas_call(
+        lambda *refs: _kernel_b4_body(
+            *refs, scale=scale, bt=config.bt, bc=config.bc, n_sub=config.n_sub,
+            config=config,
+        ),
+        grid=grid,
+        in_specs=[io_spec, io_spec, io_spec, io_spec, score_spec, score_spec],
+        out_specs=[io_spec, io_spec, io_spec, io_spec],
+        out_shape=[
+            jax.ShapeDtypeStruct((bsz, H, n_chunks, config.bt, D), jnp.float32),
+            jax.ShapeDtypeStruct((bsz, H, n_chunks, config.bt, D), jnp.float32),
+            jax.ShapeDtypeStruct((bsz, H, n_chunks, config.bt, D), jnp.float32),
+            jax.ShapeDtypeStruct((bsz, H, n_chunks, config.bt, D), jnp.float32),
+        ],
+        compiler_params=pltpu.CompilerParams(vmem_limit_bytes=150 * 1024 * 1024),
+        interpret=interpret,
+    )(q_r, k_r, b_r, g_r, dAqk, dAkk)
+
+    return dq, dk, db, dgc