From c47f1a34c9ef9b9a50d96f0b3f6d43090592d8ed Mon Sep 17 00:00:00 2001 From: Rohan Bierneni Date: Sun, 6 Sep 2026 04:30:21 +0000 Subject: [PATCH] [Stacked PR 3/5] Add decoupled GDN Pallas backward pass kernel and parity tests - Implement decoupled GDN backward pass using Pallas emit_pipeline (pallas_gdn_bwd_kernel) and custom VJP (gdn_decoupled_conv1d). - Execute Conv1D backward immediately following Pallas via unrolled native JAX shift-multiply, eliminating live range interference and vector register spills. - Provide comprehensive analytical gradient parity test suite (gdn_bwd_pallas_test.py) validating exact gradient match against pure JAX autodiff across varying sequence lengths, batch sizes, and head configurations. --- src/maxtext/models/kernels/gdn/__init__.py | 47 + .../models/kernels/gdn/gdn_bwd_pallas.py | 1801 +++++++++++++++++ tests/unit/gdn_bwd_pallas_test.py | 979 +++++++++ 3 files changed, 2827 insertions(+) create mode 100644 src/maxtext/models/kernels/gdn/__init__.py create mode 100644 src/maxtext/models/kernels/gdn/gdn_bwd_pallas.py create mode 100644 tests/unit/gdn_bwd_pallas_test.py diff --git a/src/maxtext/models/kernels/gdn/__init__.py b/src/maxtext/models/kernels/gdn/__init__.py new file mode 100644 index 000000000..84956fbad --- /dev/null +++ b/src/maxtext/models/kernels/gdn/__init__.py @@ -0,0 +1,47 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== + +"""Local GDN forward kernel package with triangular inverse matrix caching.""" + +from . import compute_conv1d +from . import compute_gdn +from . import config +from . import gdn_bwd_pallas +from . import memory_ref +from . import metadata +from . import tiling +from . import vmem_ldst +from . import wrapper +from .gdn_bwd_pallas import gdn_decoupled_conv1d +from .gdn_bwd_pallas import pallas_gdn_bwd_kernel + +gdn_fused_conv1d = gdn_decoupled_conv1d +gdn_kernel = gdn_decoupled_conv1d + +__all__ = [ + "compute_conv1d", + "compute_gdn", + "config", + "gdn_bwd_pallas", + "gdn_decoupled_conv1d", + "gdn_fused_conv1d", + "gdn_kernel", + "memory_ref", + "metadata", + "pallas_gdn_bwd_kernel", + "tiling", + "vmem_ldst", + "wrapper", +] diff --git a/src/maxtext/models/kernels/gdn/gdn_bwd_pallas.py b/src/maxtext/models/kernels/gdn/gdn_bwd_pallas.py new file mode 100644 index 000000000..e48150a82 --- /dev/null +++ b/src/maxtext/models/kernels/gdn/gdn_bwd_pallas.py @@ -0,0 +1,1801 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Canonical Gated Delta Net (GDN) backward pass using Pallas emit_pipeline. + +Decoupled GDN Architecture: +1. Conv1D forward is paired with SiLU in pure JAX (`conv1d_silu_fwd`), and GDN + forward caches the triangular inverse matrices (t_inv) in residuals while + running dense systolic matmuls on the TPU MXU. +2. Pallas GDN Backward: `pallas_gdn_bwd_kernel` executes the 40+ GDN + adjoint matrix recurrences via `pltpu.emit_pipeline` with vectorized head + processing and zero HBM intermediate spills. +3. Decoupled Conv1D Backward: `conv1d_silu_bwd` executes immediately after + Pallas via unrolled native JAX shift-multiply. + +Why Decoupled is Optimal: +Decoupling Conv1D backward into native JAX immediately after Pallas avoids +binding Conv1D gradient live ranges with GDN adjoint matrix state buffers, +eliminates vector register spills, keeps peak VMEM well within fast VMEM +boundaries, and maximizes training throughput. +""" + +import functools +from typing import Any, Optional, Tuple + +import jax +from jax.ad_checkpoint import checkpoint_name +from jax.experimental import pallas as pl +from jax.experimental.pallas import tpu as pltpu +import jax.numpy as jnp + +try: + from maxtext.layers import normalizations +except ImportError: + from maxtext.src.maxtext.layers import normalizations + +try: + import jax.experimental.xla_metadata + + if not hasattr(jax.experimental.xla_metadata, "must_fuse_call"): + jax.experimental.xla_metadata.must_fuse_call = lambda *args, **kwargs: (lambda fn: fn) +except Exception: # pylint: disable=broad-exception-caught + pass + + +try: + from maxtext.models.kernels.gdn import compute_gdn as local_compute_gdn + from maxtext.models.kernels.gdn import wrapper as local_gdn_wrapper +except ImportError: + try: + from maxtext.src.maxtext.models.kernels.gdn import compute_gdn as local_compute_gdn + from maxtext.src.maxtext.models.kernels.gdn import wrapper as local_gdn_wrapper + except ImportError: + from . import compute_gdn as local_compute_gdn + from . import wrapper as local_gdn_wrapper + + +# ============================================================================== +# SECTION 1: CPU Interpret & Runtime Helpers +# ============================================================================== + + +def ensure_cpu_interpret_registered() -> None: + """Ensures Pallas CPU interpretation registers TPU hardware info without top-level import side-effects.""" + # pylint: disable=protected-access,broad-exception-caught,import-outside-toplevel + try: + from jax._src.pallas.mosaic import tpu_info as _ti # noqa: E402 + + if "cpu" not in _ti.registry: + _ti.registry["cpu"] = lambda: _ti.get_tpu_info_for_chip(_ti.ChipVersion.TPU_V6E, 1) + try: + _ti.get_tpu_info.cache_clear() + except Exception: + pass + except Exception: # pragma: no cover + pass + + try: + from jax._src.pallas.mosaic import pipeline as _pl_pipeline # noqa: E402 + + if not getattr(_pl_pipeline, "_is_cpu_safe_cbs_patched", False): + _orig_cbs = getattr( + _pl_pipeline, + "_original_create_bounded_slice", + _pl_pipeline._create_bounded_slice, + ) + _pl_pipeline._original_create_bounded_slice = _orig_cbs + + def _cpu_safe_create_bounded_slice( + slice_start, + slice_size, + block_size, + dim_size, + *args, + tiling=None, + **kwargs, + ): + if isinstance(slice_size, int) and (tiling is None or slice_size % tiling == 0): + return pl.ds(slice_start, slice_size) + return _orig_cbs( + slice_start, + slice_size, + block_size, + dim_size, + tiling, + *args, + **kwargs, + ) + + _pl_pipeline._create_bounded_slice = _cpu_safe_create_bounded_slice + _pl_pipeline._is_cpu_safe_cbs_patched = True + except Exception: # pragma: no cover + pass + # pylint: enable=protected-access,broad-exception-caught,import-outside-toplevel + + +@jax.custom_vjp +def invert_triangular_matrix(t: jax.Array) -> jax.Array: + """Computes inverse of unit lower-triangular matrix using Tokamax block forward substitution.""" + return local_compute_gdn.invert_triangular_matrix(t, block_size=16) + + +def _invert_triangular_matrix_fwd(t: jax.Array): + t_inv = invert_triangular_matrix(t) + return t_inv, t_inv + + +def _invert_triangular_matrix_bwd(res, g): + t_inv = res + grad_t = jnp.tril(-(t_inv.mT @ g @ t_inv.mT), k=-1) + return (grad_t,) + + +invert_triangular_matrix.defvjp( + _invert_triangular_matrix_fwd, + _invert_triangular_matrix_bwd, +) + + +def chunk_forward( + q: jax.Array, + k: jax.Array, + v: jax.Array, + b_val: jax.Array, + a_val: jax.Array, + a_log_val: jax.Array, + dt_bias_val: jax.Array, + state_prev: jax.Array, + *, + kq_head_dim: int, + repeats: int, + chunk_size: int, + use_qk_norm_in_gdn: bool = False, +) -> Tuple[jax.Array, jax.Array]: + """Computes one chunk forward pass for GDN v3 with WY delta rule.""" + out, state_new, _ = chunk_forward_with_tinv( + q=q, + k=k, + v=v, + b_val=b_val, + a_val=a_val, + a_log_val=a_log_val, + dt_bias_val=dt_bias_val, + state_prev=state_prev, + kq_head_dim=kq_head_dim, + repeats=repeats, + chunk_size=chunk_size, + use_qk_norm_in_gdn=use_qk_norm_in_gdn, + ) + return out, state_new + + +def chunk_forward_with_tinv( + q: jax.Array, + k: jax.Array, + v: jax.Array, + b_val: jax.Array, + a_val: jax.Array, + a_log_val: jax.Array, + dt_bias_val: jax.Array, + state_prev: jax.Array, + *, + kq_head_dim: int, + repeats: int, + chunk_size: int, + use_qk_norm_in_gdn: bool = False, +) -> Tuple[jax.Array, jax.Array, jax.Array]: + """Computes one chunk forward pass for GDN v3 and returns (out, state_new, t_inv).""" + q = q.astype(jnp.float32) + k = k.astype(jnp.float32) + v = v.astype(jnp.float32) + if use_qk_norm_in_gdn: + q = normalizations.l2norm(q, dim=-1, eps=1e-6) + k = normalizations.l2norm(k, dim=-1, eps=1e-6) + scale = 1.0 / jnp.sqrt(kq_head_dim) + q = q * scale + b_val = b_val.astype(jnp.float32) + a_val = a_val.astype(jnp.float32) + a_log_val = a_log_val.astype(jnp.float32) + dt_bias_val = dt_bias_val.astype(jnp.float32) + state_prev = state_prev.astype(jnp.float32) + q_rep = jnp.repeat(q, repeats, axis=1) + k_rep = jnp.repeat(k, repeats, axis=1) + + beta = jax.nn.sigmoid(b_val) + + # EXACT GDN v3 gating formula + log_g = -jnp.exp(a_log_val) * jax.nn.softplus(a_val + dt_bias_val) + + # Fast MXU cumsum replacement + mask_cumsum = jnp.tril(jnp.ones((chunk_size, chunk_size), dtype=log_g.dtype)) + cumsum_log_g = jnp.dot(mask_cumsum, log_g) + + # Transpose to head-first: (H, C, D) + q_h = jnp.transpose(q_rep, (1, 0, 2)) + k_h = jnp.transpose(k_rep, (1, 0, 2)) + v_h = jnp.transpose(v, (1, 0, 2)) + beta_h = jnp.transpose(beta, (1, 0)) + cumsum_h = jnp.transpose(cumsum_log_g, (1, 0)) + + diff = cumsum_h[:, :, None] - cumsum_h[:, None, :] + mask_strict = jnp.tril(jnp.ones((chunk_size, chunk_size), dtype=diff.dtype), k=-1) + safe_diff_strict = jnp.where(mask_strict[None, :, :] == 1.0, diff, -1e4) + g_mat_strict = jnp.exp(safe_diff_strict) * mask_strict[None, :, :] + + mask_causal = jnp.tril(jnp.ones((chunk_size, chunk_size), dtype=diff.dtype), k=0) + safe_diff_causal = jnp.where(mask_causal[None, :, :] == 1.0, diff, -1e4) + g_mat_causal = jnp.exp(safe_diff_causal) * mask_causal[None, :, :] + + gating_forward = jnp.exp(cumsum_h)[:, :, None] + gating_last = jnp.exp(cumsum_h[:, -1])[:, None, None] + gating_backward = jnp.exp(cumsum_h[:, -1:] - cumsum_h)[:, :, None] + + # WY Representation: T = unit lower-triangular Gram matrix + k_beta = k_h * beta_h[:, :, None] + S = jnp.matmul(k_beta, jnp.swapaxes(k_h, -1, -2)) * g_mat_strict + identity_mask = jnp.eye(chunk_size, dtype=S.dtype)[None, :, :] + t = jnp.where(identity_mask == 1.0, 1.0, S) + A = invert_triangular_matrix(t) + + v_beta = v_h * beta_h[:, :, None] + k_beta_g = k_beta * gating_forward + u = jnp.matmul(A, v_beta) + w = jnp.matmul(A, k_beta_g) + + # Delta error subtraction against recurrent state + ws = jnp.matmul(w, state_prev) + v_new = u - ws + + # Output: cross-chunk state read + intra-chunk attention with v_new + q_g = q_h * gating_forward + out_cross = jnp.matmul(q_g, state_prev) + + attn = jnp.matmul(q_h, jnp.swapaxes(k_h, -1, -2)) * g_mat_causal + out_intra = jnp.matmul(attn, v_new) + + out = out_cross + out_intra + out = jnp.transpose(out, (1, 0, 2)) + + # State update: decayed previous state + rank-1 update from chunk with v_new + state_prev_decayed = state_prev * gating_last + k_scaled = k_h * gating_backward + state_new_intra = jnp.matmul(jnp.swapaxes(k_scaled, -1, -2), v_new) + state_new = state_prev_decayed + state_new_intra + + return out, state_new, A + + +def chunk_state_forward_with_cached_tinv( + k: jax.Array, + v: jax.Array, + b_val: jax.Array, + a_val: jax.Array, + a_log_val: jax.Array, + dt_bias_val: jax.Array, + state_prev: jax.Array, + t_inv: jax.Array, + *, + repeats: int, + chunk_size: int, + use_qk_norm_in_gdn: bool = False, +) -> jax.Array: + """Computes next recurrent state for one chunk using cached t_inv without triangular inversion.""" + v = v.astype(jnp.float32) + k = k.astype(jnp.float32) + if use_qk_norm_in_gdn: + k = normalizations.l2norm(k, dim=-1, eps=1e-6) + k_rep = jnp.repeat(k, repeats, axis=1) + + b_val = b_val.astype(jnp.float32) + a_val = a_val.astype(jnp.float32) + a_log_val = a_log_val.astype(jnp.float32) + dt_bias_val = dt_bias_val.astype(jnp.float32) + state_prev = state_prev.astype(jnp.float32) + + beta = jax.nn.sigmoid(b_val) + log_g = -jnp.exp(a_log_val) * jax.nn.softplus(a_val + dt_bias_val) + + mask_cumsum = jnp.tril(jnp.ones((chunk_size, chunk_size), dtype=log_g.dtype)) + cumsum_log_g = jnp.dot(mask_cumsum, log_g) + + k_h = jnp.transpose(k_rep, (1, 0, 2)) + v_h = jnp.transpose(v, (1, 0, 2)) + beta_h = jnp.transpose(beta, (1, 0)) + cumsum_h = jnp.transpose(cumsum_log_g, (1, 0)) + + gating_forward = jnp.exp(cumsum_h)[:, :, None] + gating_last = jnp.exp(cumsum_h[:, -1])[:, None, None] + gating_backward = jnp.exp(cumsum_h[:, -1:] - cumsum_h)[:, :, None] + + A = t_inv.astype(jnp.float32) + k_beta = k_h * beta_h[:, :, None] + v_beta = v_h * beta_h[:, :, None] + k_beta_g = k_beta * gating_forward + + u = jnp.matmul(A, v_beta) + w = jnp.matmul(A, k_beta_g) + + ws = jnp.matmul(w, state_prev) + v_new = u - ws + + state_prev_decayed = state_prev * gating_last + k_scaled = k_h * gating_backward + state_new_intra = jnp.matmul(jnp.swapaxes(k_scaled, -1, -2), v_new) + state_new = state_prev_decayed + state_new_intra + + return state_new + + +# ============================================================================== +# SECTION 2: Conv1D + SiLU Forward & Backward Primitives (Pure JAX) +# ============================================================================== + + +def conv1d_silu_fwd( + qkv: jax.Array, + conv_weight: jax.Array, + conv_bias: Optional[jax.Array], + kernel_size: int, +) -> Tuple[jax.Array, jax.Array]: + """Forward Conv1D + SiLU returning (conv_out, qkv_conv).""" + _, seq_len, _ = qkv.shape + if conv_weight.ndim == 3: + conv_weight_3d = conv_weight.astype(jnp.float32) + else: + conv_weight_3d = conv_weight[:, None, :].astype(jnp.float32) + + conv_input = jnp.pad(qkv.astype(jnp.float32), ((0, 0), (kernel_size - 1, 0), (0, 0))) + conv_out = sum(conv_input[:, k : k + seq_len, :] * conv_weight_3d[k, 0, :] for k in range(kernel_size)) + if conv_bias is not None: + conv_out = conv_out + conv_bias.astype(jnp.float32) + qkv_conv = jax.nn.silu(conv_out) + return conv_out, qkv_conv.astype(qkv.dtype) + + +def conv1d_silu_bwd( + qkv: jax.Array, + conv_weight: jax.Array, + conv_bias: Optional[jax.Array], + dy: jax.Array, + kernel_size: int, +) -> Tuple[jax.Array, jax.Array, Optional[jax.Array]]: + """Dedicated Conv1D + SiLU backward pass using JAX primitives.""" + _, seq_len, _ = qkv.shape + if conv_weight.ndim == 3: + conv_weight_3d = conv_weight.astype(jnp.float32) + else: + conv_weight_3d = conv_weight[:, None, :].astype(jnp.float32) + + # 1. Forward pass: z = conv1d(x) + b + conv_input = jnp.pad(qkv.astype(jnp.float32), ((0, 0), (kernel_size - 1, 0), (0, 0))) + conv_out = sum(conv_input[:, k : k + seq_len, :] * conv_weight_3d[k, 0, :] for k in range(kernel_size)) + if conv_bias is not None: + conv_out = conv_out + conv_bias.astype(jnp.float32) + z = conv_out + + # 2. Adjoint: dz = dy * SiLU'(z) + sig_z = jax.nn.sigmoid(z) + silu_prime = sig_z * (1.0 + z * (1.0 - sig_z)) + dz = dy.astype(jnp.float32) * silu_prime + + # 3. Parameter gradients: + # Bias gradient: db = sum(dz) + if conv_bias is not None: + db = jnp.sum(dz, axis=(0, 1)).astype(conv_bias.dtype) + if conv_bias.ndim != 1: + db = db.reshape(conv_bias.shape) + else: + db = None + + # Weight gradient: dw[k] = sum_{b,t} dz[b,t] * conv_input[b, t+k] + dw_rows = [] + for k in range(kernel_size): + x_k = conv_input[:, k : k + seq_len, :] + dw_rows.append(jnp.sum(dz * x_k, axis=(0, 1))) + dw = jnp.stack(dw_rows, axis=0) + if conv_weight.ndim == 3: + dw = dw[:, None, :].astype(conv_weight.dtype) + else: + dw = dw.astype(conv_weight.dtype) + + # 4. Input gradient: dx = transposed convolution of dz with reversed w + dz_pad = jnp.pad(dz, ((0, 0), (0, kernel_size - 1), (0, 0))) + w_rev = conv_weight_3d[::-1] + dx = sum(dz_pad[:, k : k + seq_len, :] * w_rev[k, 0, :] for k in range(kernel_size)) + dx = dx.astype(qkv.dtype) + + return dx, dw, db + + +# ============================================================================== +# SECTION 3: Pallas GDN Backward Pipeline Kernel (Mosaic TPU) +# ============================================================================== + + +def make_bwd_block_specs( + batch_size: int, + num_chunks: int, + chunk_size: int, + dim_size: int, + num_v_heads: int, + kq_head_dim: int, + v_head_dim: int, + padded_num_v_heads: int | None = None, + g: Any = None, + kernel_size: int = 4, + pad_len: int = 8, +) -> Tuple[list[pl.BlockSpec], list[pl.BlockSpec], int, int]: + """Constructs reverse-scan Pallas emit_pipeline in_specs and out_specs for GDN backward.""" + del batch_size, kernel_size, pad_len, g + if padded_num_v_heads is None: + padded_num_v_heads = ((num_v_heads + 127) // 128) * 128 + + def rc(c): + return num_chunks - 1 - c + + in_specs = [ + pl.BlockSpec( + (None, None, chunk_size, dim_size), + lambda b, c: (b, rc(c), 0, 0), + ), + pl.BlockSpec( + (None, None, chunk_size, padded_num_v_heads), + lambda b, c: (b, rc(c), 0, 0), + ), + pl.BlockSpec( + (None, None, chunk_size, padded_num_v_heads), + lambda b, c: (b, rc(c), 0, 0), + ), + pl.BlockSpec( + (None, None, chunk_size, num_v_heads, v_head_dim), + lambda b, c: (b, rc(c), 0, 0, 0), + ), + pl.BlockSpec( + (None, None, num_v_heads, kq_head_dim, v_head_dim), + lambda b, c: (b, rc(c), 0, 0, 0), + ), + pl.BlockSpec( + (None, None, num_v_heads, chunk_size, chunk_size), + lambda b, c: (b, rc(c), 0, 0, 0), + ), + pl.BlockSpec( + (None, 1, padded_num_v_heads), + lambda b, c: (b, 0, 0), + ), + pl.BlockSpec( + (None, 1, padded_num_v_heads), + lambda b, c: (b, 0, 0), + ), + pl.BlockSpec( + (None, None, 1, 128), + lambda b, c: (b, rc(c), 0, 0), + ), + ] + out_specs = [ + pl.BlockSpec( + (None, None, chunk_size, dim_size), + lambda b, c: (b, rc(c), 0, 0), + ), + pl.BlockSpec( + (None, None, chunk_size, padded_num_v_heads), + lambda b, c: (b, rc(c), 0, 0), + ), + pl.BlockSpec( + (None, None, chunk_size, padded_num_v_heads), + lambda b, c: (b, rc(c), 0, 0), + ), + pl.BlockSpec( + (None, None, 1, padded_num_v_heads), + lambda b, c: (b, rc(c), 0, 0), + ), + pl.BlockSpec( + (None, None, 1, padded_num_v_heads), + lambda b, c: (b, rc(c), 0, 0), + ), + ] + return in_specs, out_specs, len(in_specs), len(out_specs) + + +def _bwd_gdn_pipeline_body( + qkv_conv_ref: Any, + b_ref: Any, + a_ref: Any, + do_ref: Any, + chunk_states_ref: Any, + t_inv_ref: Any, + a_log_ref: Any, + dt_bias_ref: Any, + reset_ref: Any, + dy_conv_ref: Any, + d_b_ref: Any, + d_a_ref: Any, + d_a_log_ref: Any, + d_dt_bias_ref: Any, + d_state_scr: Any, + *, + chunk_size: int, + dim_size: int, + num_kq_heads: int, + num_v_heads: int, + padded_num_v_heads: int, + kq_head_dim: int, + v_head_dim: int, + use_qk_norm_in_gdn: bool, + kernel_size: int = 4, + pad_len: int = 8, +) -> None: + """Inner kernel executed per (batch, chunk) by emit_pipeline with manual GDN backward.""" + del kernel_size, pad_len + c = pl.program_id(1) + repeats = num_v_heads // num_kq_heads + q_size = num_kq_heads * kq_head_dim + k_size = num_kq_heads * kq_head_dim + v_size = num_v_heads * v_head_dim + + @pl.when(c == 0) + def _init(): + d_state_scr[...] = jnp.zeros((num_v_heads, kq_head_dim, v_head_dim), dtype=jnp.float32) + + is_reset = reset_ref[...][0, 0] > 0.5 + d_state = jnp.where(is_reset, 0.0, d_state_scr[...]) + y_c = qkv_conv_ref[...] + + # Slice chunk inputs for this head group + q_orig = y_c[:, :q_size].reshape((chunk_size, num_kq_heads, kq_head_dim)).astype(jnp.float32) + k_orig = y_c[:, q_size : q_size + k_size].reshape((chunk_size, num_kq_heads, kq_head_dim)).astype(jnp.float32) + v = ( + y_c[:, q_size + k_size : q_size + k_size + v_size] + .reshape((chunk_size, num_v_heads, v_head_dim)) + .astype(jnp.float32) + ) + + b_val = b_ref[...][:, :num_v_heads].astype(jnp.float32) + a_val = a_ref[...][:, :num_v_heads].astype(jnp.float32) + do_val = do_ref[...].astype(jnp.float32) + state_prev = chunk_states_ref[...].astype(jnp.float32) + t_inv_val = t_inv_ref[...].astype(jnp.float32) + a_log_val = a_log_ref[...][0, :num_v_heads].astype(jnp.float32) + dt_bias_val = dt_bias_ref[...][0, :num_v_heads].astype(jnp.float32) + + scale = 1.0 / jnp.sqrt(kq_head_dim) + mask_cumsum = jnp.tril(jnp.ones((chunk_size, chunk_size), dtype=jnp.float32)) + mask_strict = jnp.tril(jnp.ones((chunk_size, chunk_size), dtype=jnp.float32), k=-1) + mask_causal = jnp.tril(jnp.ones((chunk_size, chunk_size), dtype=jnp.float32), k=0) + + if use_qk_norm_in_gdn: + norm_q = normalizations.l2norm(q_orig, dim=-1, eps=1e-6) + norm_k = normalizations.l2norm(k_orig, dim=-1, eps=1e-6) + q_scaled = norm_q * scale + k_scaled_val = norm_k + else: + q_scaled = q_orig * scale + k_scaled_val = k_orig + + q_rep = jnp.repeat(q_scaled, repeats, axis=1) + k_rep = jnp.repeat(k_scaled_val, repeats, axis=1) + beta = jax.nn.sigmoid(b_val) + + sp_input = a_val + dt_bias_val + sp_val = jax.nn.softplus(sp_input) + exp_a_log = jnp.exp(a_log_val) + log_g = -exp_a_log * sp_val + + cumsum_log_g = jnp.dot(mask_cumsum, log_g) + + q_h = jnp.transpose(q_rep, (1, 0, 2)) + k_h = jnp.transpose(k_rep, (1, 0, 2)) + v_h = jnp.transpose(v, (1, 0, 2)) + beta_h = jnp.transpose(beta, (1, 0)) + cumsum_h = jnp.transpose(cumsum_log_g, (1, 0)) + do_h = jnp.transpose(do_val, (1, 0, 2)) + + diff = cumsum_h[:, :, None] - cumsum_h[:, None, :] + safe_diff_strict = jnp.where(mask_strict[None, :, :] == 1.0, diff, -1e4) + g_mat_strict = jnp.exp(safe_diff_strict) * mask_strict[None, :, :] + + safe_diff_causal = jnp.where(mask_causal[None, :, :] == 1.0, diff, -1e4) + g_mat_causal = jnp.exp(safe_diff_causal) * mask_causal[None, :, :] + + gating_forward = jnp.exp(cumsum_h)[:, :, None] + gating_last = jnp.exp(cumsum_h[:, -1])[:, None, None] + gating_backward = jnp.exp(cumsum_h[:, -1:] - cumsum_h)[:, :, None] + + k_beta = k_h * beta_h[:, :, None] + k_h_T = jnp.swapaxes(k_h, -1, -2) + S_unmasked = jnp.matmul(k_beta, k_h_T) + + A = t_inv_val + + v_beta = v_h * beta_h[:, :, None] + k_beta_g = k_beta * gating_forward + + u = jnp.matmul(A, v_beta) + w = jnp.matmul(A, k_beta_g) + + ws = jnp.matmul(w, state_prev) + v_new = u - ws + + q_g = q_h * gating_forward + attn_unmasked = jnp.matmul(q_h, k_h_T) + attn = attn_unmasked * g_mat_causal + + k_scaled_bwd = k_h * gating_backward + + dv_attn = jnp.matmul(jnp.swapaxes(attn, -1, -2), do_h) + + dv_new = dv_attn + jnp.matmul(k_scaled_bwd, d_state) + d_attn = jnp.matmul(do_h, jnp.swapaxes(v_new, -1, -2)) + + du = dv_new + dw = -jnp.matmul(dv_new, jnp.swapaxes(state_prev, -1, -2)) + + d_state_prev = ( + d_state * gating_last + jnp.matmul(jnp.swapaxes(q_g, -1, -2), do_h) - jnp.matmul(jnp.swapaxes(w, -1, -2), dv_new) + ) + + A_T = jnp.swapaxes(A, -1, -2) + d_v_beta = jnp.matmul(A_T, du) + d_k_beta_g = jnp.matmul(A_T, dw) + + dA = jnp.matmul(du, jnp.swapaxes(v_beta, -1, -2)) + jnp.matmul(dw, jnp.swapaxes(k_beta_g, -1, -2)) + + dS = jnp.tril(-jnp.matmul(jnp.matmul(A_T, dA), A_T), k=-1) + + d_S_unmasked = dS * g_mat_strict + d_k_beta_from_S = jnp.matmul(d_S_unmasked, k_h) + d_k_h_from_S = jnp.matmul(jnp.swapaxes(d_S_unmasked, -1, -2), k_beta) + + d_k_beta = d_k_beta_g * gating_forward + d_k_beta_from_S + d_beta_h = jnp.sum(d_v_beta * v_h, axis=-1) + jnp.sum(d_k_beta * k_h, axis=-1) + d_v_h = d_v_beta * beta_h[:, :, None] + + d_attn_unmasked = d_attn * g_mat_causal + d_q_h_from_attn = jnp.matmul(d_attn_unmasked, k_h) + d_k_h_from_attn = jnp.matmul(jnp.swapaxes(d_attn_unmasked, -1, -2), q_h) + + d_q_g = jnp.matmul(do_h, jnp.swapaxes(state_prev, -1, -2)) + d_q_h_from_q_g = d_q_g * gating_forward + + d_k_scaled = jnp.matmul(v_new, jnp.swapaxes(d_state, -1, -2)) + d_k_h_from_k_scaled = d_k_scaled * gating_backward + + d_q_h = d_q_h_from_q_g + d_q_h_from_attn + d_k_h = d_k_h_from_attn + d_k_h_from_S + d_k_beta * beta_h[:, :, None] + d_k_h_from_k_scaled + + # Gating adjoints + d_gating_forward = jnp.sum(d_q_g * q_h, axis=-1) + jnp.sum(d_k_beta_g * k_beta, axis=-1) + d_cumsum_from_fwd = d_gating_forward * gating_forward[:, :, 0] + + d_gating_last = jnp.sum(d_state * state_prev, axis=(-1, -2)) + d_cumsum_last = d_gating_last * jnp.exp(cumsum_h[:, -1]) + + d_gating_backward = jnp.sum(d_k_scaled * k_h, axis=-1) + d_diff_bwd = d_gating_backward * gating_backward[:, :, 0] + + d_g_strict = dS * S_unmasked + d_g_causal = d_attn * attn_unmasked + d_diff = (d_g_strict * g_mat_strict) + (d_g_causal * g_mat_causal) + d_cumsum_from_diff = jnp.sum(d_diff, axis=2) - jnp.sum(d_diff, axis=1) + + d_cumsum_h = d_cumsum_from_fwd + d_cumsum_from_diff - d_diff_bwd + last_col_addition = (d_cumsum_last + jnp.sum(d_diff_bwd, axis=-1))[:, None] + d_cumsum_h = d_cumsum_h + jnp.pad(last_col_addition, ((0, 0), (chunk_size - 1, 0))) + + d_cumsum_log_g = jnp.transpose(d_cumsum_h, (1, 0)) + d_log_g = jnp.dot(mask_cumsum.T, d_cumsum_log_g) + + sig_sp = jax.nn.sigmoid(sp_input) + d_a_val = d_log_g * (-exp_a_log * sig_sp) + d_a_log_val = jnp.sum(d_log_g * (-exp_a_log * sp_val), axis=0) + d_dt_bias_val = jnp.sum(d_a_val, axis=0) + + d_b_val = (jnp.transpose(d_beta_h, (1, 0))) * beta * (1.0 - beta) + + d_v_val = jnp.transpose(d_v_h, (1, 0, 2)) + d_q_rep = jnp.transpose(d_q_h, (1, 0, 2)) + d_k_rep = jnp.transpose(d_k_h, (1, 0, 2)) + + d_q_proj = jnp.sum( + d_q_rep.reshape(chunk_size, num_kq_heads, repeats, kq_head_dim), + axis=2, + ) + d_k_proj = jnp.sum( + d_k_rep.reshape(chunk_size, num_kq_heads, repeats, kq_head_dim), + axis=2, + ) + + if use_qk_norm_in_gdn: + d_q_scaled = d_q_proj * scale + r_q = jnp.sqrt(jnp.sum(q_orig**2, axis=-1, keepdims=True) + 1e-12) + q_unit = q_orig / r_q + d_q = (d_q_scaled - q_unit * jnp.sum(d_q_scaled * q_unit, axis=-1, keepdims=True)) / r_q + + r_k = jnp.sqrt(jnp.sum(k_orig**2, axis=-1, keepdims=True) + 1e-12) + k_unit = k_orig / r_k + d_k = (d_k_proj - k_unit * jnp.sum(d_k_proj * k_unit, axis=-1, keepdims=True)) / r_k + else: + d_q = d_q_proj * scale + d_k = d_k_proj + + # Flatten gradients to chunk_size x dim_size and write to refs + d_q_flat = d_q.reshape(chunk_size, q_size).astype(dy_conv_ref.dtype) + d_k_flat = d_k.reshape(chunk_size, k_size).astype(dy_conv_ref.dtype) + d_v_flat = d_v_val.reshape(chunk_size, v_size).astype(dy_conv_ref.dtype) + dy_conv_ref[...] = jnp.concatenate([d_q_flat, d_k_flat, d_v_flat], axis=-1) + + if padded_num_v_heads > num_v_heads: + d_b_ref[...] = jnp.pad(d_b_val, ((0, 0), (0, padded_num_v_heads - num_v_heads))).astype(d_b_ref.dtype) + d_a_ref[...] = jnp.pad(d_a_val, ((0, 0), (0, padded_num_v_heads - num_v_heads))).astype(d_a_ref.dtype) + d_a_log_ref[...] = jnp.pad(d_a_log_val[None, :], ((0, 0), (0, padded_num_v_heads - num_v_heads))).astype( + d_a_log_ref.dtype + ) + d_dt_bias_ref[...] = jnp.pad(d_dt_bias_val[None, :], ((0, 0), (0, padded_num_v_heads - num_v_heads))).astype( + d_dt_bias_ref.dtype + ) + else: + d_b_ref[...] = d_b_val.astype(d_b_ref.dtype) + d_a_ref[...] = d_a_val.astype(d_a_ref.dtype) + d_a_log_ref[...] = d_a_log_val[None, :].astype(d_a_log_ref.dtype) + d_dt_bias_ref[...] = d_dt_bias_val[None, :].astype(d_dt_bias_ref.dtype) + + d_state_scr[...] = d_state_prev.astype(d_state_scr.dtype) + + +def _pallas_gdn_bwd_kernel_single_group( + qkv_conv: jax.Array, + b: jax.Array, + a: jax.Array, + a_log: jax.Array, + dt_bias: jax.Array, + do: jax.Array, + chunk_states: jax.Array, + t_inv: jax.Array, + *, + num_v_heads: int, + num_kq_heads: int, + kq_head_dim: int, + v_head_dim: int, + chunk_size: int = 64, + use_qk_norm_in_gdn: bool = False, + vmem_limit_mb: Optional[int] = None, + segment_ids: Optional[jax.Array] = None, + interpret: bool | pltpu.InterpretParams | None = None, + name: str = "gdn_bwd_kernel", +) -> Tuple[jax.Array, jax.Array, jax.Array, jax.Array, jax.Array]: + """Executes single head-group Pallas emit_pipeline kernel.""" + batch_size, seq_len, group_dim_size = qkv_conv.shape + num_chunks = seq_len // chunk_size + padded_num_v_heads = ((num_v_heads + 127) // 128) * 128 + + # Reshape inputs into chunked tensors for emit_pipeline + qkv_conv_4d = qkv_conv.reshape(batch_size, num_chunks, chunk_size, group_dim_size) + + b_4d = b.reshape(batch_size, num_chunks, chunk_size, num_v_heads) + if padded_num_v_heads > num_v_heads: + b_4d = jnp.pad( + b_4d, + ((0, 0), (0, 0), (0, 0), (0, padded_num_v_heads - num_v_heads)), + ) + + a_4d = a.reshape(batch_size, num_chunks, chunk_size, num_v_heads) + if padded_num_v_heads > num_v_heads: + a_4d = jnp.pad( + a_4d, + ((0, 0), (0, 0), (0, 0), (0, padded_num_v_heads - num_v_heads)), + ) + + do_5d = do.reshape(batch_size, num_chunks, chunk_size, num_v_heads, v_head_dim) + + if a_log.ndim == 1: + a_log_3d = jnp.broadcast_to(a_log[None, None, :], (batch_size, 1, num_v_heads)) + elif a_log.ndim == 2: + a_log_3d = a_log[:, None, :] + else: + a_log_3d = a_log + if padded_num_v_heads > num_v_heads: + a_log_3d = jnp.pad(a_log_3d, ((0, 0), (0, 0), (0, padded_num_v_heads - num_v_heads))) + + if dt_bias.ndim == 1: + dt_bias_3d = jnp.broadcast_to(dt_bias[None, None, :], (batch_size, 1, num_v_heads)) + elif dt_bias.ndim == 2: + dt_bias_3d = dt_bias[:, None, :] + else: + dt_bias_3d = dt_bias + if padded_num_v_heads > num_v_heads: + dt_bias_3d = jnp.pad(dt_bias_3d, ((0, 0), (0, 0), (0, padded_num_v_heads - num_v_heads))) + + t_inv_5d = t_inv.astype(jnp.float32) + + # Segment reset tensor for cross-document boundary gradient reset + if segment_ids is not None and num_chunks > 1: + end_idx = jnp.arange(1, num_chunks) * chunk_size - 1 + start_next_idx = jnp.arange(1, num_chunks) * chunk_size + boundaries = segment_ids[:, end_idx] != segment_ids[:, start_next_idx] + reset_mask = jnp.pad(boundaries, ((0, 0), (0, 1)), constant_values=False) + reset_hbm = jnp.pad( + reset_mask[:, :, None, None].astype(jnp.float32), + ((0, 0), (0, 0), (0, 0), (0, 127)), + ) + else: + reset_hbm = jnp.zeros((batch_size, num_chunks, 1, 128), dtype=jnp.float32) + + in_specs, out_specs, num_in, num_out = make_bwd_block_specs( + batch_size=batch_size, + num_chunks=num_chunks, + chunk_size=chunk_size, + dim_size=group_dim_size, + num_v_heads=num_v_heads, + kq_head_dim=kq_head_dim, + v_head_dim=v_head_dim, + padded_num_v_heads=padded_num_v_heads, + ) + + out_shapes = ( + jax.ShapeDtypeStruct( + (batch_size, num_chunks, chunk_size, group_dim_size), + qkv_conv.dtype, + ), + jax.ShapeDtypeStruct(b_4d.shape, b_4d.dtype), + jax.ShapeDtypeStruct(a_4d.shape, a_4d.dtype), + jax.ShapeDtypeStruct( + (batch_size, num_chunks, 1, padded_num_v_heads), + a_log_3d.dtype, + ), + jax.ShapeDtypeStruct( + (batch_size, num_chunks, 1, padded_num_v_heads), + dt_bias_3d.dtype, + ), + ) + + body = functools.partial( + _bwd_gdn_pipeline_body, + chunk_size=chunk_size, + dim_size=group_dim_size, + num_kq_heads=num_kq_heads, + num_v_heads=num_v_heads, + padded_num_v_heads=padded_num_v_heads, + kq_head_dim=kq_head_dim, + v_head_dim=v_head_dim, + use_qk_norm_in_gdn=use_qk_norm_in_gdn, + ) + + def outer(*refs): + pltpu.emit_pipeline( + body, + grid=(batch_size, num_chunks), + in_specs=in_specs, + out_specs=out_specs, + )(*refs[: num_in + num_out], scratches=tuple(refs[num_in + num_out :])) + + if vmem_limit_mb is not None and vmem_limit_mb <= 64: + vmem_limit_bytes = int(vmem_limit_mb) * 1024 * 1024 + else: + tpu_info = pltpu.get_tpu_info() + vmem_limit_bytes = int(0.85 * tpu_info.vmem_capacity_bytes) + + hbm = pltpu.MemorySpace.HBM + ( + dy_conv_chunks, + d_b_chunks, + d_a_chunks, + d_a_log_chunks, + d_dt_bias_chunks, + ) = pl.pallas_call( + outer, + grid=(), + out_shape=out_shapes, + in_specs=[pl.BlockSpec(memory_space=hbm)] * num_in, + out_specs=[pl.BlockSpec(memory_space=hbm)] * num_out, + scratch_shapes=[ + pltpu.VMEM((num_v_heads, kq_head_dim, v_head_dim), jnp.float32), + ], + compiler_params=pltpu.CompilerParams( + vmem_limit_bytes=vmem_limit_bytes, + disable_bounds_checks=True, + ), + interpret=interpret, + name=name, + )( + qkv_conv_4d, + b_4d, + a_4d, + do_5d, + chunk_states, + t_inv_5d, + a_log_3d, + dt_bias_3d, + reset_hbm, + ) + + d_a_log_reduced = jnp.sum(d_a_log_chunks[..., 0, :num_v_heads], axis=(0, 1)).astype(a_log.dtype) + d_dt_bias_reduced = jnp.sum(d_dt_bias_chunks[..., 0, :num_v_heads], axis=(0, 1)).astype(dt_bias.dtype) + d_b_flat = d_b_chunks[..., :num_v_heads].reshape(batch_size, seq_len, num_v_heads).astype(b.dtype) + d_a_flat = d_a_chunks[..., :num_v_heads].reshape(batch_size, seq_len, num_v_heads).astype(a.dtype) + dy_conv_flat = dy_conv_chunks.reshape(batch_size, seq_len, group_dim_size).astype(qkv_conv.dtype) + + return ( + dy_conv_flat, + d_b_flat, + d_a_flat, + d_a_log_reduced, + d_dt_bias_reduced, + ) + + +def pallas_gdn_bwd_kernel( + qkv_conv: jax.Array, + b: jax.Array, + a: jax.Array, + a_log: jax.Array, + dt_bias: jax.Array, + do: jax.Array, + chunk_states: jax.Array, + t_inv: jax.Array, + *, + num_v_heads: int, + kq_head_dim: int, + v_head_dim: int, + chunk_size: int = 64, + use_qk_norm_in_gdn: bool = False, + vmem_limit_mb: Optional[int] = None, + head_tile: Optional[int] = None, + segment_ids: Optional[jax.Array] = None, + interpret: bool | pltpu.InterpretParams | None = None, +) -> Tuple[ + jax.Array, + jax.Array, + jax.Array, + jax.Array, + jax.Array, +]: + """Executes the Pallas reverse-chunk GDNv3 backward kernel using emit_pipeline. + + Streams natively per group in contiguous layout. + """ + if interpret is None and jax.default_backend() == "cpu": + interpret = True + if interpret: + ensure_cpu_interpret_registered() + + _, _, dim_size = qkv_conv.shape + num_kq_heads = (dim_size - num_v_heads * v_head_dim) // (kq_head_dim * 2) + repeats = num_v_heads // num_kq_heads + + target_tile = 16 if head_tile is None else head_tile + max_possible = min(num_v_heads, target_tile) + tile_v_heads = None + for candidate in range(max_possible, 0, -1): + if num_v_heads % candidate == 0 and candidate % repeats == 0: + tile_v_heads = candidate + break + if tile_v_heads is None: + tile_v_heads = repeats if num_v_heads % repeats == 0 else num_v_heads + num_groups = num_v_heads // tile_v_heads + tile_kq_heads = tile_v_heads // repeats + + if num_groups == 1: + return _pallas_gdn_bwd_kernel_single_group( + qkv_conv=qkv_conv, + b=b, + a=a, + a_log=a_log, + dt_bias=dt_bias, + do=do, + chunk_states=chunk_states, + t_inv=t_inv, + num_v_heads=num_v_heads, + num_kq_heads=num_kq_heads, + kq_head_dim=kq_head_dim, + v_head_dim=v_head_dim, + chunk_size=chunk_size, + use_qk_norm_in_gdn=use_qk_norm_in_gdn, + vmem_limit_mb=vmem_limit_mb, + segment_ids=segment_ids, + interpret=interpret, + name="gdn_bwd_kernel", + ) + + q_size = num_kq_heads * kq_head_dim + k_size = num_kq_heads * kq_head_dim + tile_q_size = tile_kq_heads * kq_head_dim + tile_k_size = tile_kq_heads * kq_head_dim + + dq_list = [] + dk_list = [] + dv_list = [] + db_list = [] + da_list = [] + dal_list = [] + ddt_list = [] + + for g in range(num_groups): + vh_start = g * tile_v_heads + vh_end = (g + 1) * tile_v_heads + kqh_start = g * tile_kq_heads + kqh_end = (g + 1) * tile_kq_heads + + q_g = qkv_conv[:, :, kqh_start * kq_head_dim : kqh_end * kq_head_dim] + k_g = qkv_conv[:, :, q_size + kqh_start * kq_head_dim : q_size + kqh_end * kq_head_dim] + v_g = qkv_conv[ + :, + :, + q_size + k_size + vh_start * v_head_dim : q_size + k_size + vh_end * v_head_dim, + ] + qkv_g = jnp.concatenate([q_g, k_g, v_g], axis=-1) + + b_g = b[:, :, vh_start:vh_end] + a_g = a[:, :, vh_start:vh_end] + do_g = do[:, :, vh_start:vh_end, :] + chunk_states_g = chunk_states[:, :, vh_start:vh_end, :, :] + t_inv_g = t_inv[:, :, vh_start:vh_end, :, :] + + if a_log.ndim == 1: + a_log_g = a_log[vh_start:vh_end] + elif a_log.ndim == 2: + a_log_g = a_log[:, vh_start:vh_end] + else: + a_log_g = a_log[:, :, vh_start:vh_end] + + if dt_bias.ndim == 1: + dt_bias_g = dt_bias[vh_start:vh_end] + elif dt_bias.ndim == 2: + dt_bias_g = dt_bias[:, vh_start:vh_end] + else: + dt_bias_g = dt_bias[:, :, vh_start:vh_end] + + dy_g, db_g, da_g, dal_g, ddt_g = _pallas_gdn_bwd_kernel_single_group( + qkv_conv=qkv_g, + b=b_g, + a=a_g, + a_log=a_log_g, + dt_bias=dt_bias_g, + do=do_g, + chunk_states=chunk_states_g, + t_inv=t_inv_g, + num_v_heads=tile_v_heads, + num_kq_heads=tile_kq_heads, + kq_head_dim=kq_head_dim, + v_head_dim=v_head_dim, + chunk_size=chunk_size, + use_qk_norm_in_gdn=use_qk_norm_in_gdn, + vmem_limit_mb=vmem_limit_mb, + segment_ids=segment_ids, + interpret=interpret, + name=f"gdn_bwd_kernel_group_{g}", + ) + + dq_g = dy_g[:, :, :tile_q_size] + dk_g = dy_g[:, :, tile_q_size : tile_q_size + tile_k_size] + dv_g = dy_g[:, :, tile_q_size + tile_k_size :] + + dq_list.append(dq_g) + dk_list.append(dk_g) + dv_list.append(dv_g) + db_list.append(db_g) + da_list.append(da_g) + dal_list.append(dal_g) + ddt_list.append(ddt_g) + + dq_flat = jnp.concatenate(dq_list, axis=-1) + dk_flat = jnp.concatenate(dk_list, axis=-1) + dv_flat = jnp.concatenate(dv_list, axis=-1) + dy_conv_flat = jnp.concatenate([dq_flat, dk_flat, dv_flat], axis=-1).astype(qkv_conv.dtype) + d_b_flat = jnp.concatenate(db_list, axis=-1) + d_a_flat = jnp.concatenate(da_list, axis=-1) + d_a_log_reduced = jnp.concatenate(dal_list, axis=-1) + d_dt_bias_reduced = jnp.concatenate(ddt_list, axis=-1) + + return ( + dy_conv_flat, + d_b_flat, + d_a_flat, + d_a_log_reduced, + d_dt_bias_reduced, + ) + + +pallas_gdn_bwd_computation = pallas_gdn_bwd_kernel + + +def decoupled_conv1d_gdn_bwd_kernel( + pre_conv_qkv: jax.Array, + b: jax.Array, + a: jax.Array, + a_log: jax.Array, + dt_bias: jax.Array, + do: jax.Array, + chunk_states: jax.Array, + conv_weight: jax.Array, + conv_bias: Optional[jax.Array] = None, + t_inv: Optional[jax.Array] = None, + qkv: Optional[jax.Array] = None, + seq_lens: Optional[jax.Array] = None, + *, + num_v_heads: int, + kq_head_dim: int, + v_head_dim: int, + kernel_size: int = 4, + chunk_size: int = 64, + use_qk_norm_in_gdn: bool = False, + vmem_limit_mb: Optional[int] = None, + head_tile: Optional[int] = None, + segment_ids: Optional[jax.Array] = None, + interpret: bool | pltpu.InterpretParams | None = None, +) -> Tuple[ + jax.Array, + jax.Array, + jax.Array, + jax.Array, + Optional[jax.Array], + jax.Array, + jax.Array, +]: + """Decoupled Conv1D + GDN backward combining Pallas GDN bwd and JAX Conv1D bwd.""" + del seq_lens, qkv + _, qkv_conv = conv1d_silu_fwd( + qkv=pre_conv_qkv, + conv_weight=conv_weight, + conv_bias=conv_bias, + kernel_size=kernel_size, + ) + + dy_conv, d_b, d_a, d_a_log, d_dt_bias = pallas_gdn_bwd_kernel( + qkv_conv=qkv_conv, + b=b, + a=a, + a_log=a_log, + dt_bias=dt_bias, + do=do, + chunk_states=chunk_states, + t_inv=t_inv, + num_v_heads=num_v_heads, + kq_head_dim=kq_head_dim, + v_head_dim=v_head_dim, + chunk_size=chunk_size, + use_qk_norm_in_gdn=use_qk_norm_in_gdn, + vmem_limit_mb=vmem_limit_mb, + head_tile=head_tile, + segment_ids=segment_ids, + interpret=interpret, + ) + + d_pre_conv_qkv, d_conv_weight, d_conv_bias = conv1d_silu_bwd( + qkv=pre_conv_qkv, + conv_weight=conv_weight, + conv_bias=conv_bias, + dy=dy_conv, + kernel_size=kernel_size, + ) + + return ( + d_pre_conv_qkv, + d_b, + d_a, + d_conv_weight, + d_conv_bias, + d_a_log, + d_dt_bias, + ) + + +decoupled_conv1d_gdn_bwd_computation = decoupled_conv1d_gdn_bwd_kernel +pallas_fused_conv1d_gdn_bwd_kernel = decoupled_conv1d_gdn_bwd_kernel +pallas_fused_conv1d_gdn_bwd_computation = decoupled_conv1d_gdn_bwd_kernel + + +# ============================================================================== +# SECTION 4: Unified GDN Custom VJP Interface (gdn_fused_conv1d) +# ============================================================================== + + +def pure_jax_decoupled_conv1d_gdn( + qkv: jax.Array, + b: jax.Array, + a: jax.Array, + conv_weight: jax.Array, + conv_bias: Optional[jax.Array], + a_log: jax.Array, + dt_bias: jax.Array, + conv_state: Optional[jax.Array], + recurrent_state: Optional[jax.Array], + *, + num_k_heads: int, + num_v_heads: int, + head_k_dim: int, + head_v_dim: int, + conv_kernel_size: int, + chunk_size: int, + use_qk_norm_in_gdn: bool, + compute_dtype: jnp.dtype = jnp.float32, +) -> Tuple[jax.Array, Tuple[jax.Array, jax.Array]]: + """Pure-JAX composite of Conv1D + GDN used during backward pass autodiff.""" + del conv_state + batch, seq_len, _ = qkv.shape + key_dim = num_k_heads * head_k_dim + + # Conv1D in FP32 + conv_input = jnp.pad(qkv.astype(jnp.float32), ((0, 0), (conv_kernel_size - 1, 0), (0, 0))) + if conv_weight.ndim == 3: + conv_weight_3d = conv_weight.astype(jnp.float32) + else: + conv_weight_3d = conv_weight[:, None, :].astype(jnp.float32) + conv_out = sum(conv_input[:, k : k + seq_len, :] * conv_weight_3d[k, 0, :] for k in range(conv_kernel_size)) + if conv_bias is not None: + conv_out = conv_out + conv_bias.astype(jnp.float32) + qkv_conv = jax.nn.silu(conv_out).astype(jnp.float32) + + q_conv, k_conv, v_conv = jnp.split(qkv_conv, [key_dim, 2 * key_dim], axis=-1) + + # Reshape for GDN + query = q_conv.reshape(batch, seq_len, num_k_heads, head_k_dim) + key = k_conv.reshape(batch, seq_len, num_k_heads, head_k_dim) + value = v_conv.reshape(batch, seq_len, num_v_heads, head_v_dim) + + a_log_cast = jnp.asarray(a_log, dtype=jnp.float32) + dt_bias_cast = jnp.asarray(dt_bias, dtype=jnp.float32) + beta = jax.nn.sigmoid(b.astype(jnp.float32)) + g = -jnp.exp(a_log_cast) * jax.nn.softplus(a.astype(jnp.float32) + dt_bias_cast) + + if num_v_heads > num_k_heads and num_v_heads % num_k_heads == 0: + repeats = num_v_heads // num_k_heads + query = jnp.repeat(query, repeats, axis=2) + key = jnp.repeat(key, repeats, axis=2) + + try: + from maxtext.models import qwen3 # pylint: disable=import-outside-toplevel + except ImportError: + try: + from maxtext.src.maxtext.models import qwen3 # pylint: disable=import-outside-toplevel + except ImportError: + from . import qwen3 # pylint: disable=import-outside-toplevel + + core_attn_out, next_recurrent_state = qwen3.jax_chunk_gated_delta_rule( + query=query, + key=key, + value=value, + g=g, + beta=beta, + chunk_size=chunk_size, + initial_state=(recurrent_state.astype(jnp.float32) if recurrent_state is not None else None), + use_qk_norm_in_gdn=use_qk_norm_in_gdn, + compute_dtype=jnp.float32, + ) + + next_conv_state = ( + qkv[:, -(conv_kernel_size - 1) :, :] + if seq_len >= conv_kernel_size - 1 + else jnp.zeros((batch, conv_kernel_size - 1, qkv.shape[-1]), dtype=qkv.dtype) + ) + if next_recurrent_state is None: + next_recurrent_state = jnp.zeros((batch, num_v_heads, head_k_dim, head_v_dim), dtype=jnp.float32) + + return core_attn_out.astype(qkv.dtype), ( + next_conv_state.astype(qkv.dtype), + next_recurrent_state.astype(jnp.float32), + ) + + +pure_jax_fused_conv1d_gdn = pure_jax_decoupled_conv1d_gdn + + +def _compute_forward_conv_and_states( + qkv: jax.Array, + b: jax.Array, + a: jax.Array, + conv_weight: jax.Array, + conv_bias: Optional[jax.Array], + a_log: jax.Array, + dt_bias: jax.Array, + recurrent_state: Optional[jax.Array], + *, + num_k_heads: int, + num_v_heads: int, + head_k_dim: int, + head_v_dim: int, + conv_kernel_size: int, + chunk_size: int, + use_qk_norm_in_gdn: bool = False, + compute_dtype: jnp.dtype = jnp.float32, + cached_t_inv: Optional[jax.Array] = None, +) -> Tuple[jax.Array, jax.Array, jax.Array]: + """Computes convolved QKV, inter-chunk states, and t_inv matrices in FP32.""" + del compute_dtype + batch_size, seq_len, _ = qkv.shape + num_chunks = seq_len // chunk_size + + # Conv1D in FP32 + conv_input = jnp.pad(qkv.astype(jnp.float32), ((0, 0), (conv_kernel_size - 1, 0), (0, 0))) + if conv_weight.ndim == 3: + conv_weight_3d = conv_weight.astype(jnp.float32) + else: + conv_weight_3d = conv_weight[:, None, :].astype(jnp.float32) + conv_out = sum(conv_input[:, k : k + seq_len, :] * conv_weight_3d[k, 0, :] for k in range(conv_kernel_size)) + if conv_bias is not None: + conv_out = conv_out + conv_bias.astype(jnp.float32) + qkv_conv_f32 = jax.nn.silu(conv_out).astype(jnp.float32) + + # Chunk states progression in FP32 + num_kq_heads = num_k_heads + q_size = num_kq_heads * head_k_dim + k_size = num_kq_heads * head_k_dim + repeats = num_v_heads // num_kq_heads + + q = qkv_conv_f32[:, :, :q_size].reshape(batch_size, num_chunks, chunk_size, num_kq_heads, head_k_dim) + k = qkv_conv_f32[:, :, q_size : q_size + k_size].reshape(batch_size, num_chunks, chunk_size, num_kq_heads, head_k_dim) + v = qkv_conv_f32[:, :, q_size + k_size :].reshape(batch_size, num_chunks, chunk_size, num_v_heads, head_v_dim) + + b_4d = b.astype(jnp.float32).reshape(batch_size, num_chunks, chunk_size, num_v_heads) + a_4d = a.astype(jnp.float32).reshape(batch_size, num_chunks, chunk_size, num_v_heads) + a_log_f32 = a_log.astype(jnp.float32) + dt_bias_f32 = dt_bias.astype(jnp.float32) + + if recurrent_state is None: + init_state = jnp.zeros((batch_size, num_v_heads, head_k_dim, head_v_dim), dtype=jnp.float32) + else: + init_state = recurrent_state.astype(jnp.float32) + + k_chunks = k.swapaxes(0, 1) + v_chunks = v.swapaxes(0, 1) + b_chunks = b_4d.swapaxes(0, 1) + a_chunks = a_4d.swapaxes(0, 1) + + if cached_t_inv is not None: + t_inv_chunks = cached_t_inv.astype(jnp.float32).swapaxes(0, 1) + + def chunk_cached_step(k_single, v_single, b_single, a_single, s_prev, t_inv_single): + return chunk_state_forward_with_cached_tinv( + k=k_single, + v=v_single, + b_val=b_single, + a_val=a_single, + a_log_val=a_log_f32, + dt_bias_val=dt_bias_f32, + state_prev=s_prev, + t_inv=t_inv_single, + repeats=repeats, + chunk_size=chunk_size, + use_qk_norm_in_gdn=use_qk_norm_in_gdn, + ) + + def scan_fn_cached(carry_state, chunk_inputs): + k_i, v_i, b_i, a_i, t_inv_i = chunk_inputs + next_state = jax.vmap(chunk_cached_step)(k_i, v_i, b_i, a_i, carry_state, t_inv_i) + return next_state, carry_state + + _, chunk_states = jax.lax.scan( + scan_fn_cached, + init_state, + (k_chunks, v_chunks, b_chunks, a_chunks, t_inv_chunks), + ) + chunk_states = chunk_states.swapaxes(0, 1) + t_inv_all = cached_t_inv + else: + + def chunk_step(q_single, k_single, v_single, b_single, a_single, s_prev): + return chunk_forward_with_tinv( + q_single, + k_single, + v_single, + b_single, + a_single, + a_log_f32, + dt_bias_f32, + s_prev, + kq_head_dim=head_k_dim, + repeats=repeats, + chunk_size=chunk_size, + use_qk_norm_in_gdn=use_qk_norm_in_gdn, + ) + + def scan_fn(carry_state, chunk_inputs): + q_i, k_i, v_i, b_i, a_i = chunk_inputs + _, next_state, t_inv_i = jax.vmap(chunk_step)(q_i, k_i, v_i, b_i, a_i, carry_state) + return next_state, (carry_state, t_inv_i) + + q_chunks = q.swapaxes(0, 1) + _, (chunk_states, t_inv_all) = jax.lax.scan(scan_fn, init_state, (q_chunks, k_chunks, v_chunks, b_chunks, a_chunks)) + chunk_states = chunk_states.swapaxes(0, 1) + t_inv_all = t_inv_all.swapaxes(0, 1) + + return qkv_conv_f32.astype(qkv.dtype), chunk_states, t_inv_all + + +def _run_local_gdn_decoupled_fwd( + qkv: jax.Array, + b: jax.Array, + a: jax.Array, + conv_weight: jax.Array, + conv_bias: Optional[jax.Array], + a_log: jax.Array, + dt_bias: jax.Array, + conv_state: Optional[jax.Array], + recurrent_state: Optional[jax.Array], + *, + num_k_heads: int, + num_v_heads: int, + head_k_dim: int, + head_v_dim: int, + conv_kernel_size: int, + chunk_size: int, + use_qk_norm_in_gdn: bool, + compute_dtype: jnp.dtype, +) -> Tuple[ + Tuple[jax.Array, Tuple[jax.Array, jax.Array]], + Optional[jax.Array], + Optional[jax.Array], +]: + """Runs local GDN forward pass on TPU returning (t_inv, chunk_states), or pure JAX on CPU.""" + if jax.extend.backend.get_backend().platform == "cpu": + out, states = pure_jax_decoupled_conv1d_gdn( + qkv=qkv, + b=b, + a=a, + conv_weight=conv_weight, + conv_bias=conv_bias, + a_log=a_log, + dt_bias=dt_bias, + conv_state=conv_state, + recurrent_state=recurrent_state, + num_k_heads=num_k_heads, + num_v_heads=num_v_heads, + head_k_dim=head_k_dim, + head_v_dim=head_v_dim, + conv_kernel_size=conv_kernel_size, + chunk_size=chunk_size, + use_qk_norm_in_gdn=use_qk_norm_in_gdn, + compute_dtype=compute_dtype, + ) + _, chunk_states, t_inv = _compute_forward_conv_and_states( + qkv=qkv, + b=b, + a=a, + conv_weight=conv_weight, + conv_bias=conv_bias, + a_log=a_log, + dt_bias=dt_bias, + recurrent_state=recurrent_state, + num_k_heads=num_k_heads, + num_v_heads=num_v_heads, + head_k_dim=head_k_dim, + head_v_dim=head_v_dim, + conv_kernel_size=conv_kernel_size, + chunk_size=chunk_size, + use_qk_norm_in_gdn=use_qk_norm_in_gdn, + compute_dtype=compute_dtype, + ) + return (out, states), t_inv, chunk_states + + batch_size, seq_len, dim_size = qkv.shape + num_seqs = batch_size + num_chunks = seq_len // chunk_size + + qkv_flat = qkv.reshape(-1, dim_size) + b_flat = b.reshape(-1, b.shape[-1]) + a_flat = a.reshape(-1, a.shape[-1]) + tokamax_conv_weight = jnp.swapaxes(conv_weight, 0, 2) + + query_start_loc = jnp.arange(0, (num_seqs + 1) * seq_len, seq_len, dtype=jnp.int32) + state_indices = jnp.arange(num_seqs, dtype=jnp.int32) + seq_lens = jnp.full((num_seqs,), seq_len, dtype=jnp.int32) + distribution = jnp.array([0, 0, num_seqs], dtype=jnp.int32) + + if conv_state is None: + tokamax_conv_state = jnp.zeros((num_seqs + 1, conv_kernel_size - 1, dim_size), dtype=qkv.dtype) + elif conv_state.shape[0] == num_seqs: + tokamax_conv_state = jnp.pad(conv_state, ((1, 0), (0, 0), (0, 0))) + else: + tokamax_conv_state = conv_state + + if recurrent_state is None: + tokamax_recurrent_state = jnp.zeros((num_seqs + 1, num_v_heads, head_k_dim, head_v_dim), dtype=jnp.float32) + elif recurrent_state.shape[0] == num_seqs: + tokamax_recurrent_state = jnp.pad(recurrent_state.astype(jnp.float32), ((1, 0), (0, 0), (0, 0), (0, 0))) + else: + tokamax_recurrent_state = recurrent_state.astype(jnp.float32) + + ( + core_attn_out_flat, + (new_conv_state, new_recurrent_state), + t_inv_raw, + chunk_states_raw, + ) = local_gdn_wrapper.fused_conv1d_gdn( + qkv=qkv_flat, + b=b_flat, + a=a_flat, + conv_state=tokamax_conv_state, + recurrent_state=tokamax_recurrent_state, + conv_weight=tokamax_conv_weight, + conv_bias=conv_bias, + a_log=a_log, + dt_bias=dt_bias, + query_start_loc=query_start_loc, + state_indices=state_indices, + distribution=distribution, + seq_lens=seq_lens, + n_kq=num_k_heads, + n_v=num_v_heads, + d_k=head_k_dim, + d_v=head_v_dim, + kernel_size=conv_kernel_size, + compute_precision=jnp.dtype(jnp.float32), + mixed_tile_size=chunk_size, + is_prefill_only=True, + ) + + core_attn_out = core_attn_out_flat.reshape(batch_size, seq_len, num_v_heads, head_v_dim) + t_inv = t_inv_raw.astype(jnp.float32).reshape(batch_size, num_chunks, num_v_heads, chunk_size, chunk_size) + chunk_states = chunk_states_raw.astype(jnp.float32).reshape(batch_size, num_chunks, num_v_heads, head_k_dim, head_v_dim) + return ( + ( + core_attn_out.astype(qkv.dtype), + ( + new_conv_state[1:].astype(qkv.dtype), + new_recurrent_state[1:].astype(jnp.float32), + ), + ), + t_inv, + chunk_states, + ) + + +@functools.partial(jax.custom_vjp, nondiff_argnums=(9, 10, 11, 12, 13, 14, 15, 16)) +def gdn_decoupled_conv1d( + qkv: jax.Array, + b: jax.Array, + a: jax.Array, + conv_weight: jax.Array, + conv_bias: Optional[jax.Array], + a_log: jax.Array, + dt_bias: jax.Array, + conv_state: Optional[jax.Array], + recurrent_state: Optional[jax.Array], + num_k_heads: int, + num_v_heads: int, + head_k_dim: int, + head_v_dim: int, + conv_kernel_size: int, + chunk_size: int, + use_qk_norm_in_gdn: bool, + compute_dtype: jnp.dtype, +) -> Tuple[jax.Array, Tuple[jax.Array, jax.Array]]: + """Decoupled Conv1D + GDN with Pallas backward pass.""" + (out, states), _, _ = _run_local_gdn_decoupled_fwd( + qkv, + b, + a, + conv_weight, + conv_bias, + a_log, + dt_bias, + conv_state, + recurrent_state, + num_k_heads=num_k_heads, + num_v_heads=num_v_heads, + head_k_dim=head_k_dim, + head_v_dim=head_v_dim, + conv_kernel_size=conv_kernel_size, + chunk_size=chunk_size, + use_qk_norm_in_gdn=use_qk_norm_in_gdn, + compute_dtype=compute_dtype, + ) + return out, states + + +def _gdn_decoupled_conv1d_fwd( + qkv: jax.Array, + b: jax.Array, + a: jax.Array, + conv_weight: jax.Array, + conv_bias: Optional[jax.Array], + a_log: jax.Array, + dt_bias: jax.Array, + conv_state: Optional[jax.Array], + recurrent_state: Optional[jax.Array], + num_k_heads: int, + num_v_heads: int, + head_k_dim: int, + head_v_dim: int, + conv_kernel_size: int, + chunk_size: int, + use_qk_norm_in_gdn: bool, + compute_dtype: jnp.dtype, +): + """Forward rule for custom_vjp registration of gdn_decoupled_conv1d.""" + qkv = checkpoint_name(qkv, "gdn_fwd_conv") + (out, states), t_inv, chunk_states = _run_local_gdn_decoupled_fwd( + qkv, + b, + a, + conv_weight, + conv_bias, + a_log, + dt_bias, + conv_state, + recurrent_state, + num_k_heads=num_k_heads, + num_v_heads=num_v_heads, + head_k_dim=head_k_dim, + head_v_dim=head_v_dim, + conv_kernel_size=conv_kernel_size, + chunk_size=chunk_size, + use_qk_norm_in_gdn=use_qk_norm_in_gdn, + compute_dtype=compute_dtype, + ) + out = checkpoint_name(out, "gdn_fwd_out") + residuals = ( + qkv, + b, + a, + conv_weight, + conv_bias, + a_log, + dt_bias, + conv_state, + recurrent_state, + t_inv, + chunk_states, + ) + return (out, states), residuals + + +def _gdn_decoupled_conv1d_bwd( + num_k_heads: int, + num_v_heads: int, + head_k_dim: int, + head_v_dim: int, + conv_kernel_size: int, + chunk_size: int, + use_qk_norm_in_gdn: bool, + compute_dtype: jnp.dtype, + residuals: tuple, + cotangents: tuple, +): + """Backward rule for custom_vjp registration of gdn_decoupled_conv1d.""" + if len(residuals) == 11: + ( + pre_conv_qkv, + b, + a, + conv_weight, + conv_bias, + a_log, + dt_bias, + conv_state, + recurrent_state, + t_inv_fwd, + chunk_states, + ) = residuals + else: + ( + pre_conv_qkv, + b, + a, + conv_weight, + conv_bias, + a_log, + dt_bias, + conv_state, + recurrent_state, + t_inv_fwd, + ) = residuals + chunk_states = None + + d_out, d_states = cotangents + d_conv_state, d_recurrent_state = d_states + del d_conv_state, d_recurrent_state + + # Recompute forward chunk states and t_inv if not cached in residuals + if chunk_states is None or t_inv_fwd is None: + qkv_conv, chunk_states_recomputed, t_inv_recomputed = _compute_forward_conv_and_states( + qkv=pre_conv_qkv, + b=b, + a=a, + conv_weight=conv_weight, + conv_bias=conv_bias, + a_log=a_log, + dt_bias=dt_bias, + recurrent_state=recurrent_state, + num_k_heads=num_k_heads, + num_v_heads=num_v_heads, + head_k_dim=head_k_dim, + head_v_dim=head_v_dim, + conv_kernel_size=conv_kernel_size, + chunk_size=chunk_size, + use_qk_norm_in_gdn=use_qk_norm_in_gdn, + compute_dtype=compute_dtype, + cached_t_inv=t_inv_fwd, + ) + if chunk_states is None: + chunk_states = chunk_states_recomputed + if t_inv_fwd is None: + t_inv_fwd = t_inv_recomputed + else: + _, qkv_conv = conv1d_silu_fwd( + qkv=pre_conv_qkv, + conv_weight=conv_weight, + conv_bias=conv_bias, + kernel_size=conv_kernel_size, + ) + t_inv = t_inv_fwd + + dy_conv, d_b, d_a, d_a_log, d_dt_bias = pallas_gdn_bwd_kernel( + qkv_conv=qkv_conv, + b=b, + a=a, + a_log=a_log, + dt_bias=dt_bias, + do=d_out, + chunk_states=chunk_states, + t_inv=t_inv, + num_v_heads=num_v_heads, + kq_head_dim=head_k_dim, + v_head_dim=head_v_dim, + chunk_size=chunk_size, + use_qk_norm_in_gdn=use_qk_norm_in_gdn, + ) + + d_pre_conv_qkv, d_conv_weight, d_conv_bias = conv1d_silu_bwd( + qkv=pre_conv_qkv, + conv_weight=conv_weight, + conv_bias=conv_bias, + dy=dy_conv, + kernel_size=conv_kernel_size, + ) + + d_conv_state_out = None if conv_state is None else jnp.zeros_like(conv_state) + d_recurrent_state_out = None if recurrent_state is None else jnp.zeros_like(recurrent_state) + return ( + d_pre_conv_qkv, + d_b, + d_a, + d_conv_weight, + d_conv_bias, + d_a_log, + d_dt_bias, + d_conv_state_out, + d_recurrent_state_out, + ) + + +gdn_decoupled_conv1d.defvjp( + _gdn_decoupled_conv1d_fwd, + _gdn_decoupled_conv1d_bwd, +) + +# Backward compatibility aliases +gdn_fused_conv1d = gdn_decoupled_conv1d +gdn_kernel = gdn_decoupled_conv1d +_gdn_fused_conv1d_fwd = _gdn_decoupled_conv1d_fwd +_gdn_fused_conv1d_bwd = _gdn_decoupled_conv1d_bwd +_run_local_gdn_fused_fwd = _run_local_gdn_decoupled_fwd + +__all__ = [ + "chunk_forward", + "chunk_forward_with_tinv", + "pallas_gdn_bwd_kernel", + "pallas_gdn_bwd_computation", + "decoupled_conv1d_gdn_bwd_kernel", + "decoupled_conv1d_gdn_bwd_computation", + "pallas_fused_conv1d_gdn_bwd_kernel", + "pallas_fused_conv1d_gdn_bwd_computation", + "conv1d_silu_fwd", + "conv1d_silu_bwd", + "pure_jax_decoupled_conv1d_gdn", + "pure_jax_fused_conv1d_gdn", + "gdn_decoupled_conv1d", + "gdn_fused_conv1d", + "gdn_kernel", + "ensure_cpu_interpret_registered", +] diff --git a/tests/unit/gdn_bwd_pallas_test.py b/tests/unit/gdn_bwd_pallas_test.py new file mode 100644 index 000000000..a7aa54b91 --- /dev/null +++ b/tests/unit/gdn_bwd_pallas_test.py @@ -0,0 +1,979 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for gdn_bwd_pallas with canonical GDN backward pass.""" + +# pylint: disable=protected-access + +from absl.testing import absltest +import jax +import jax.numpy as jnp +import numpy as np + +try: + from maxtext.models.kernels.gdn import gdn_bwd_pallas + from maxtext.models import qwen3 +except ImportError: + from maxtext.src.maxtext.models.kernels.gdn import gdn_bwd_pallas + from maxtext.src.maxtext.models import qwen3 + + +class GdnBwdPallasTest(absltest.TestCase): + + def setUp(self): + super().setUp() + gdn_bwd_pallas.ensure_cpu_interpret_registered() + + def test_chunk_forward_matches_jax_reference(self): + key = jax.random.PRNGKey(42) + chunk_size = 64 + num_kq_heads = 2 + num_v_heads = 4 + kq_head_dim = 128 + v_head_dim = 128 + repeats = num_v_heads // num_kq_heads + + k1, k2, k3, k4, k5, k6, k7, k8 = jax.random.split(key, 8) + q = jax.random.normal(k1, (chunk_size, num_kq_heads, kq_head_dim), dtype=jnp.float32) + k = jax.random.normal(k2, (chunk_size, num_kq_heads, kq_head_dim), dtype=jnp.float32) + v = jax.random.normal(k3, (chunk_size, num_v_heads, v_head_dim), dtype=jnp.float32) + b_val = jax.random.normal(k4, (chunk_size, num_v_heads), dtype=jnp.float32) + a_val = jax.random.normal(k5, (chunk_size, num_v_heads), dtype=jnp.float32) + a_log_val = jax.random.normal(k6, (num_v_heads,), dtype=jnp.float32) + dt_bias_val = jax.random.normal(k7, (num_v_heads,), dtype=jnp.float32) + state_prev = jax.random.normal(k8, (num_v_heads, kq_head_dim, v_head_dim), dtype=jnp.float32) + + out_emit, state_emit, t_inv = gdn_bwd_pallas.chunk_forward_with_tinv( + q, + k, + v, + b_val, + a_val, + a_log_val, + dt_bias_val, + state_prev, + kq_head_dim=kq_head_dim, + repeats=repeats, + chunk_size=chunk_size, + use_qk_norm_in_gdn=True, + ) + + self.assertEqual(t_inv.shape, (num_v_heads, chunk_size, chunk_size)) + + # Reference computation using qwen3.jax_chunk_gated_delta_rule + q_4d = q[None, :, :, :] + k_4d = k[None, :, :, :] + v_4d = v[None, :, :, :] + q_rep_4d = jnp.repeat(q_4d, repeats, axis=2) + k_rep_4d = jnp.repeat(k_4d, repeats, axis=2) + beta_3d = jax.nn.sigmoid(b_val)[None, :, :] + log_g_3d = (-jnp.exp(a_log_val) * jax.nn.softplus(a_val + dt_bias_val))[None, :, :] + state_4d = state_prev[None, :, :, :] + + expected_out, expected_state = qwen3.jax_chunk_gated_delta_rule( + query=q_rep_4d, + key=k_rep_4d, + value=v_4d, + g=log_g_3d, + beta=beta_3d, + chunk_size=chunk_size, + initial_state=state_4d, + use_qk_norm_in_gdn=True, + compute_dtype=jnp.float32, + ) + + np.testing.assert_allclose(out_emit, expected_out[0], rtol=1e-3, atol=1e-3) + np.testing.assert_allclose(state_emit, expected_state[0], rtol=5e-3, atol=5e-3) + + def test_compute_forward_conv_and_states(self): + batch_size = 1 + chunk_size = 16 + num_chunks = 2 + seq_len = num_chunks * chunk_size + num_k_heads = 1 + num_v_heads = 2 + head_k_dim = 128 + head_v_dim = 128 + conv_kernel_size = 4 + dim_size = num_k_heads * head_k_dim * 2 + num_v_heads * head_v_dim + + key = jax.random.PRNGKey(456) + k1, k2, k3, k4, k5, k6, k7 = jax.random.split(key, 7) + + qkv = jax.random.normal(k1, (batch_size, seq_len, dim_size), dtype=jnp.float32) + b = jax.random.normal(k2, (batch_size, seq_len, num_v_heads), dtype=jnp.float32) + a = jax.random.normal(k3, (batch_size, seq_len, num_v_heads), dtype=jnp.float32) + conv_weight = jax.random.normal(k4, (conv_kernel_size, 1, dim_size), dtype=jnp.float32) + conv_bias = jax.random.normal(k5, (dim_size,), dtype=jnp.float32) + a_log = jax.random.normal(k6, (num_v_heads,), dtype=jnp.float32) + dt_bias = jax.random.normal(k7, (num_v_heads,), dtype=jnp.float32) + + qkv_conv, chunk_states, t_inv = gdn_bwd_pallas._compute_forward_conv_and_states( + qkv=qkv, + b=b, + a=a, + conv_weight=conv_weight, + conv_bias=conv_bias, + a_log=a_log, + dt_bias=dt_bias, + recurrent_state=None, + num_k_heads=num_k_heads, + num_v_heads=num_v_heads, + head_k_dim=head_k_dim, + head_v_dim=head_v_dim, + conv_kernel_size=conv_kernel_size, + chunk_size=chunk_size, + use_qk_norm_in_gdn=True, + ) + + self.assertEqual( + chunk_states.shape, + (batch_size, num_chunks, num_v_heads, head_k_dim, head_v_dim), + ) + self.assertEqual( + t_inv.shape, + (batch_size, num_chunks, num_v_heads, chunk_size, chunk_size), + ) + + conv_input = jnp.pad(qkv.astype(jnp.float32), ((0, 0), (conv_kernel_size - 1, 0), (0, 0))) + expected_conv_out = jax.lax.conv_general_dilated( + lhs=conv_input, + rhs=conv_weight.astype(jnp.float32), + window_strides=(1,), + padding="VALID", + dimension_numbers=("NWC", "WIO", "NWC"), + feature_group_count=dim_size, + ) + expected_conv_out = expected_conv_out + conv_bias + expected_qkv_conv = jax.nn.silu(expected_conv_out) + np.testing.assert_allclose(qkv_conv, expected_qkv_conv, rtol=1e-5, atol=1e-5) + + def test_decoupled_conv1d_gdn_kernel_gradient_against_autodiff(self): + """Compares gdn_decoupled_conv1d custom VJP against JAX autodiff on pure JAX.""" + batch_size = 1 + chunk_size = 64 + num_chunks = 2 + seq_len = num_chunks * chunk_size + num_k_heads = 2 + num_v_heads = 4 + head_k_dim = 128 + head_v_dim = 128 + conv_kernel_size = 4 + dim_size = num_k_heads * head_k_dim * 2 + num_v_heads * head_v_dim + + key = jax.random.PRNGKey(789) + k1, k2, k3, k4, k5, k6, k7, k8 = jax.random.split(key, 8) + + qkv = jax.random.normal(k1, (batch_size, seq_len, dim_size), dtype=jnp.float32) + b = jax.random.normal(k2, (batch_size, seq_len, num_v_heads), dtype=jnp.float32) + a = jax.random.normal(k3, (batch_size, seq_len, num_v_heads), dtype=jnp.float32) + conv_weight = jax.random.normal(k4, (conv_kernel_size, 1, dim_size), dtype=jnp.float32) + conv_bias = jax.random.normal(k5, (dim_size,), dtype=jnp.float32) + a_log = jax.random.normal(k6, (num_v_heads,), dtype=jnp.float32) + dt_bias = jax.random.normal(k7, (num_v_heads,), dtype=jnp.float32) + do = jax.random.normal(k8, (batch_size, seq_len, num_v_heads, head_v_dim), dtype=jnp.float32) + + # 1. Golden Reference Gradient via Autodiff on pure JAX implementation + def loss_pure(qkv_in, b_in, a_in, cw_in, cb_in, al_in, dt_in): + out, _ = gdn_bwd_pallas.pure_jax_decoupled_conv1d_gdn( + qkv=qkv_in, + b=b_in, + a=a_in, + conv_weight=cw_in, + conv_bias=cb_in, + a_log=al_in, + dt_bias=dt_in, + conv_state=None, + recurrent_state=None, + num_k_heads=num_k_heads, + num_v_heads=num_v_heads, + head_k_dim=head_k_dim, + head_v_dim=head_v_dim, + conv_kernel_size=conv_kernel_size, + chunk_size=chunk_size, + use_qk_norm_in_gdn=True, + ) + return jnp.sum(out * do) + + exp_dqkv, exp_db, exp_da, exp_dcw, exp_dcb, exp_dal, exp_ddt = jax.grad(loss_pure, argnums=(0, 1, 2, 3, 4, 5, 6))( + qkv, b, a, conv_weight, conv_bias, a_log, dt_bias + ) + + # 2. Kernel Gradients via GDN Kernel custom VJP + def loss_gdn_kernel(qkv_in, b_in, a_in, cw_in, cb_in, al_in, dt_in): + out, _ = gdn_bwd_pallas.gdn_decoupled_conv1d( + qkv=qkv_in, + b=b_in, + a=a_in, + conv_weight=cw_in, + conv_bias=cb_in, + a_log=al_in, + dt_bias=dt_in, + conv_state=None, + recurrent_state=None, + num_k_heads=num_k_heads, + num_v_heads=num_v_heads, + head_k_dim=head_k_dim, + head_v_dim=head_v_dim, + conv_kernel_size=conv_kernel_size, + chunk_size=chunk_size, + use_qk_norm_in_gdn=True, + compute_dtype=jnp.float32, + ) + return jnp.sum(out * do) + + act_dqkv, act_db, act_da, act_dcw, act_dcb, act_dal, act_ddt = jax.grad( + loss_gdn_kernel, argnums=(0, 1, 2, 3, 4, 5, 6) + )(qkv, b, a, conv_weight, conv_bias, a_log, dt_bias) + + print("\n--- GDN Kernel Custom VJP vs Pure JAX Autodiff Breakdown ---") + comparisons = [ + ("beta (d_b)", exp_db, act_db), + ("alpha (d_a)", exp_da, act_da), + ("a_log (d_a_log)", exp_dal, act_dal), + ("dt_bias (d_dt_bias)", exp_ddt, act_ddt), + ("qkv (d_qkv)", exp_dqkv, act_dqkv), + ("conv_weight (d_conv_weight)", exp_dcw, act_dcw), + ("conv_bias (d_conv_bias)", exp_dcb, act_dcb), + ] + + for name, exp_g, act_g in comparisons: + self.assertIsNotNone(act_g, f"{name} actual gradient is None") + abs_diff = float(jnp.max(jnp.abs(exp_g - act_g))) + rel_diff = abs_diff / (float(jnp.max(jnp.abs(exp_g))) + 1e-7) + status = "✅ MATCH" if rel_diff < 1e-3 else "❌ DIVERGED" + print(f" {name:<28}: MaxAbsDiff = {abs_diff:.2e} | RelDiff =" f" {rel_diff:.2e} | {status}") + self.assertLess( + rel_diff, + 1e-3, + f"{name} relative difference {rel_diff:.2e} exceeds tolerance 1e-3", + ) + + print("✅ All 7 parameter gradients match Pure JAX autodiff within 0.1% on" " CPU!") + + def test_decoupled_conv1d_gdn_kernel_conv_bias_none(self): + """Verifies GDN kernel backward executes correctly when conv_bias is None.""" + batch_size = 1 + chunk_size = 32 + num_chunks = 2 + seq_len = num_chunks * chunk_size + num_k_heads = 1 + num_v_heads = 2 + head_k_dim = 64 + head_v_dim = 64 + conv_kernel_size = 4 + dim_size = num_k_heads * head_k_dim * 2 + num_v_heads * head_v_dim + + key = jax.random.PRNGKey(101) + k1, k2, k3, k4, k5, k6, k7 = jax.random.split(key, 7) + + qkv = jax.random.normal(k1, (batch_size, seq_len, dim_size), dtype=jnp.float32) + b = jax.random.normal(k2, (batch_size, seq_len, num_v_heads), dtype=jnp.float32) + a = jax.random.normal(k3, (batch_size, seq_len, num_v_heads), dtype=jnp.float32) + conv_weight = jax.random.normal(k4, (conv_kernel_size, 1, dim_size), dtype=jnp.float32) + a_log = jax.random.normal(k5, (num_v_heads,), dtype=jnp.float32) + dt_bias = jax.random.normal(k6, (num_v_heads,), dtype=jnp.float32) + do = jax.random.normal(k7, (batch_size, seq_len, num_v_heads, head_v_dim), dtype=jnp.float32) + + def loss_fn(qkv_in, b_in, a_in, cw_in, al_in, dt_in): + out, _ = gdn_bwd_pallas.gdn_decoupled_conv1d( + qkv=qkv_in, + b=b_in, + a=a_in, + conv_weight=cw_in, + conv_bias=None, + a_log=al_in, + dt_bias=dt_in, + conv_state=None, + recurrent_state=None, + num_k_heads=num_k_heads, + num_v_heads=num_v_heads, + head_k_dim=head_k_dim, + head_v_dim=head_v_dim, + conv_kernel_size=conv_kernel_size, + chunk_size=chunk_size, + use_qk_norm_in_gdn=True, + compute_dtype=jnp.float32, + ) + return jnp.sum(out * do) + + grads = jax.grad(loss_fn, argnums=(0, 1, 2, 3, 4, 5))(qkv, b, a, conv_weight, a_log, dt_bias) + for g in grads: + self.assertIsNotNone(g) + self.assertFalse(np.any(np.isnan(np.array(g)))) + + def test_decoupled_conv1d_gdn_kernel_multi_batch(self): + """Verifies GDN kernel backward handles batch_size > 1.""" + batch_size = 2 + chunk_size = 32 + num_chunks = 2 + seq_len = num_chunks * chunk_size + num_k_heads = 1 + num_v_heads = 2 + head_k_dim = 64 + head_v_dim = 64 + conv_kernel_size = 4 + dim_size = num_k_heads * head_k_dim * 2 + num_v_heads * head_v_dim + + key = jax.random.PRNGKey(202) + k1, k2, k3, k4, k5, k6, k7, k8 = jax.random.split(key, 8) + + qkv = jax.random.normal(k1, (batch_size, seq_len, dim_size), dtype=jnp.float32) + b = jax.random.normal(k2, (batch_size, seq_len, num_v_heads), dtype=jnp.float32) + a = jax.random.normal(k3, (batch_size, seq_len, num_v_heads), dtype=jnp.float32) + conv_weight = jax.random.normal(k4, (conv_kernel_size, 1, dim_size), dtype=jnp.float32) + conv_bias = jax.random.normal(k5, (dim_size,), dtype=jnp.float32) + a_log = jax.random.normal(k6, (num_v_heads,), dtype=jnp.float32) + dt_bias = jax.random.normal(k7, (num_v_heads,), dtype=jnp.float32) + do = jax.random.normal(k8, (batch_size, seq_len, num_v_heads, head_v_dim), dtype=jnp.float32) + + def loss_fn(qkv_in, b_in, a_in, cw_in, cb_in, al_in, dt_in): + out, _ = gdn_bwd_pallas.gdn_decoupled_conv1d( + qkv=qkv_in, + b=b_in, + a=a_in, + conv_weight=cw_in, + conv_bias=cb_in, + a_log=al_in, + dt_bias=dt_in, + conv_state=None, + recurrent_state=None, + num_k_heads=num_k_heads, + num_v_heads=num_v_heads, + head_k_dim=head_k_dim, + head_v_dim=head_v_dim, + conv_kernel_size=conv_kernel_size, + chunk_size=chunk_size, + use_qk_norm_in_gdn=True, + compute_dtype=jnp.float32, + ) + return jnp.sum(out * do) + + grads = jax.grad(loss_fn, argnums=(0, 1, 2, 3, 4, 5, 6))(qkv, b, a, conv_weight, conv_bias, a_log, dt_bias) + for g in grads: + self.assertIsNotNone(g) + self.assertFalse(np.any(np.isnan(np.array(g)))) + + def test_chunk_state_forward_with_cached_tinv_matches_chunk_forward(self): + """Verifies chunk_state_forward_with_cached_tinv matches chunk_forward_with_tinv state.""" + key = jax.random.PRNGKey(999) + chunk_size = 64 + num_kq_heads = 2 + num_v_heads = 4 + kq_head_dim = 128 + v_head_dim = 128 + repeats = num_v_heads // num_kq_heads + + k1, k2, k3, k4, k5, k6, k7, k8 = jax.random.split(key, 8) + q = jax.random.normal(k1, (chunk_size, num_kq_heads, kq_head_dim), dtype=jnp.float32) + k = jax.random.normal(k2, (chunk_size, num_kq_heads, kq_head_dim), dtype=jnp.float32) + v = jax.random.normal(k3, (chunk_size, num_v_heads, v_head_dim), dtype=jnp.float32) + b_val = jax.random.normal(k4, (chunk_size, num_v_heads), dtype=jnp.float32) + a_val = jax.random.normal(k5, (chunk_size, num_v_heads), dtype=jnp.float32) + a_log_val = jax.random.normal(k6, (num_v_heads,), dtype=jnp.float32) + dt_bias_val = jax.random.normal(k7, (num_v_heads,), dtype=jnp.float32) + state_prev = jax.random.normal(k8, (num_v_heads, kq_head_dim, v_head_dim), dtype=jnp.float32) + + _, expected_state, t_inv = gdn_bwd_pallas.chunk_forward_with_tinv( + q, + k, + v, + b_val, + a_val, + a_log_val, + dt_bias_val, + state_prev, + kq_head_dim=kq_head_dim, + repeats=repeats, + chunk_size=chunk_size, + use_qk_norm_in_gdn=True, + ) + + actual_state = gdn_bwd_pallas.chunk_state_forward_with_cached_tinv( + k=k, + v=v, + b_val=b_val, + a_val=a_val, + a_log_val=a_log_val, + dt_bias_val=dt_bias_val, + state_prev=state_prev, + t_inv=t_inv, + repeats=repeats, + chunk_size=chunk_size, + use_qk_norm_in_gdn=True, + ) + + np.testing.assert_allclose(actual_state, expected_state, rtol=1e-6, atol=1e-6) + + def test_compute_forward_conv_and_states_with_cached_tinv(self): + """Verifies _compute_forward_conv_and_states with cached_t_inv matches uncached version.""" + batch_size = 1 + chunk_size = 16 + num_chunks = 2 + seq_len = num_chunks * chunk_size + num_k_heads = 1 + num_v_heads = 2 + head_k_dim = 64 + head_v_dim = 64 + conv_kernel_size = 4 + dim_size = num_k_heads * head_k_dim * 2 + num_v_heads * head_v_dim + + key = jax.random.PRNGKey(888) + k1, k2, k3, k4, k5, k6, k7 = jax.random.split(key, 7) + + qkv = jax.random.normal(k1, (batch_size, seq_len, dim_size), dtype=jnp.float32) + b = jax.random.normal(k2, (batch_size, seq_len, num_v_heads), dtype=jnp.float32) + a = jax.random.normal(k3, (batch_size, seq_len, num_v_heads), dtype=jnp.float32) + conv_weight = jax.random.normal(k4, (conv_kernel_size, 1, dim_size), dtype=jnp.float32) + conv_bias = jax.random.normal(k5, (dim_size,), dtype=jnp.float32) + a_log = jax.random.normal(k6, (num_v_heads,), dtype=jnp.float32) + dt_bias = jax.random.normal(k7, (num_v_heads,), dtype=jnp.float32) + + qkv_conv_ref, chunk_states_ref, t_inv_ref = gdn_bwd_pallas._compute_forward_conv_and_states( + qkv=qkv, + b=b, + a=a, + conv_weight=conv_weight, + conv_bias=conv_bias, + a_log=a_log, + dt_bias=dt_bias, + recurrent_state=None, + num_k_heads=num_k_heads, + num_v_heads=num_v_heads, + head_k_dim=head_k_dim, + head_v_dim=head_v_dim, + conv_kernel_size=conv_kernel_size, + chunk_size=chunk_size, + use_qk_norm_in_gdn=True, + ) + + qkv_conv_cached, chunk_states_cached, t_inv_cached = gdn_bwd_pallas._compute_forward_conv_and_states( + qkv=qkv, + b=b, + a=a, + conv_weight=conv_weight, + conv_bias=conv_bias, + a_log=a_log, + dt_bias=dt_bias, + recurrent_state=None, + num_k_heads=num_k_heads, + num_v_heads=num_v_heads, + head_k_dim=head_k_dim, + head_v_dim=head_v_dim, + conv_kernel_size=conv_kernel_size, + chunk_size=chunk_size, + use_qk_norm_in_gdn=True, + cached_t_inv=t_inv_ref, + ) + + np.testing.assert_allclose(qkv_conv_cached, qkv_conv_ref, rtol=1e-6, atol=1e-6) + np.testing.assert_allclose(chunk_states_cached, chunk_states_ref, rtol=1e-6, atol=1e-6) + np.testing.assert_allclose(t_inv_cached, t_inv_ref, rtol=1e-6, atol=1e-6) + + def test_decoupled_conv1d_gdn_kernel_bwd_with_cached_tinv_in_residuals(self): + """Verifies _gdn_decoupled_conv1d_bwd gives identical grads with cached t_inv.""" + batch_size = 1 + chunk_size = 32 + num_chunks = 2 + seq_len = num_chunks * chunk_size + num_k_heads = 1 + num_v_heads = 2 + head_k_dim = 64 + head_v_dim = 64 + conv_kernel_size = 4 + dim_size = num_k_heads * head_k_dim * 2 + num_v_heads * head_v_dim + + key = jax.random.PRNGKey(777) + k1, k2, k3, k4, k5, k6, k7, k8 = jax.random.split(key, 8) + + qkv = jax.random.normal(k1, (batch_size, seq_len, dim_size), dtype=jnp.float32) + b = jax.random.normal(k2, (batch_size, seq_len, num_v_heads), dtype=jnp.float32) + a = jax.random.normal(k3, (batch_size, seq_len, num_v_heads), dtype=jnp.float32) + conv_weight = jax.random.normal(k4, (conv_kernel_size, 1, dim_size), dtype=jnp.float32) + conv_bias = jax.random.normal(k5, (dim_size,), dtype=jnp.float32) + a_log = jax.random.normal(k6, (num_v_heads,), dtype=jnp.float32) + dt_bias = jax.random.normal(k7, (num_v_heads,), dtype=jnp.float32) + do = jax.random.normal(k8, (batch_size, seq_len, num_v_heads, head_v_dim), dtype=jnp.float32) + + _, chunk_states, t_inv = gdn_bwd_pallas._compute_forward_conv_and_states( + qkv=qkv, + b=b, + a=a, + conv_weight=conv_weight, + conv_bias=conv_bias, + a_log=a_log, + dt_bias=dt_bias, + recurrent_state=None, + num_k_heads=num_k_heads, + num_v_heads=num_v_heads, + head_k_dim=head_k_dim, + head_v_dim=head_v_dim, + conv_kernel_size=conv_kernel_size, + chunk_size=chunk_size, + use_qk_norm_in_gdn=True, + ) + + res_none = ( + qkv, + b, + a, + conv_weight, + conv_bias, + a_log, + dt_bias, + None, + None, + None, + ) + res_cached = ( + qkv, + b, + a, + conv_weight, + conv_bias, + a_log, + dt_bias, + None, + None, + t_inv, + ) + res_cached_all = ( + qkv, + b, + a, + conv_weight, + conv_bias, + a_log, + dt_bias, + None, + None, + t_inv, + chunk_states, + ) + cotangents = (do, (None, None)) + + grads_none = gdn_bwd_pallas._gdn_decoupled_conv1d_bwd( + num_k_heads=num_k_heads, + num_v_heads=num_v_heads, + head_k_dim=head_k_dim, + head_v_dim=head_v_dim, + conv_kernel_size=conv_kernel_size, + chunk_size=chunk_size, + use_qk_norm_in_gdn=True, + compute_dtype=jnp.float32, + residuals=res_none, + cotangents=cotangents, + ) + + grads_cached = gdn_bwd_pallas._gdn_decoupled_conv1d_bwd( + num_k_heads=num_k_heads, + num_v_heads=num_v_heads, + head_k_dim=head_k_dim, + head_v_dim=head_v_dim, + conv_kernel_size=conv_kernel_size, + chunk_size=chunk_size, + use_qk_norm_in_gdn=True, + compute_dtype=jnp.float32, + residuals=res_cached, + cotangents=cotangents, + ) + + grads_cached_all = gdn_bwd_pallas._gdn_decoupled_conv1d_bwd( + num_k_heads=num_k_heads, + num_v_heads=num_v_heads, + head_k_dim=head_k_dim, + head_v_dim=head_v_dim, + conv_kernel_size=conv_kernel_size, + chunk_size=chunk_size, + use_qk_norm_in_gdn=True, + compute_dtype=jnp.float32, + residuals=res_cached_all, + cotangents=cotangents, + ) + + for g_none, g_cached, g_cached_all in zip(grads_none, grads_cached, grads_cached_all): + if g_none is not None and g_cached is not None: + np.testing.assert_allclose(g_none, g_cached, rtol=1e-5, atol=1e-5) + if g_none is not None and g_cached_all is not None: + np.testing.assert_allclose(g_none, g_cached_all, rtol=1e-5, atol=1e-5) + + def test_run_local_gdn_decoupled_fwd_returns_cached_chunk_states(self): + """Verifies _run_local_gdn_decoupled_fwd returns properly shaped chunk_states.""" + batch_size = 1 + chunk_size = 32 + num_chunks = 2 + seq_len = num_chunks * chunk_size + num_k_heads = 1 + num_v_heads = 2 + head_k_dim = 64 + head_v_dim = 64 + conv_kernel_size = 4 + dim_size = num_k_heads * head_k_dim * 2 + num_v_heads * head_v_dim + + key = jax.random.PRNGKey(101) + k1, k2, k3, k4, k5, k6, k7 = jax.random.split(key, 7) + + qkv = jax.random.normal(k1, (batch_size, seq_len, dim_size), dtype=jnp.float32) + b = jax.random.normal(k2, (batch_size, seq_len, num_v_heads), dtype=jnp.float32) + a = jax.random.normal(k3, (batch_size, seq_len, num_v_heads), dtype=jnp.float32) + conv_weight = jax.random.normal(k4, (conv_kernel_size, 1, dim_size), dtype=jnp.float32) + conv_bias = jax.random.normal(k5, (dim_size,), dtype=jnp.float32) + a_log = jax.random.normal(k6, (num_v_heads,), dtype=jnp.float32) + dt_bias = jax.random.normal(k7, (num_v_heads,), dtype=jnp.float32) + + (_, _), t_inv, chunk_states = gdn_bwd_pallas._run_local_gdn_decoupled_fwd( + qkv=qkv, + b=b, + a=a, + conv_weight=conv_weight, + conv_bias=conv_bias, + a_log=a_log, + dt_bias=dt_bias, + conv_state=None, + recurrent_state=None, + num_k_heads=num_k_heads, + num_v_heads=num_v_heads, + head_k_dim=head_k_dim, + head_v_dim=head_v_dim, + conv_kernel_size=conv_kernel_size, + chunk_size=chunk_size, + use_qk_norm_in_gdn=True, + compute_dtype=jnp.float32, + ) + + self.assertIsNotNone(chunk_states) + self.assertIsNotNone(t_inv) + self.assertEqual( + chunk_states.shape, + (batch_size, num_chunks, num_v_heads, head_k_dim, head_v_dim), + ) + self.assertEqual( + t_inv.shape, + (batch_size, num_chunks, num_v_heads, chunk_size, chunk_size), + ) + + # Verify against golden pure JAX _compute_forward_conv_and_states + _, exp_chunk_states, exp_t_inv = gdn_bwd_pallas._compute_forward_conv_and_states( + qkv=qkv, + b=b, + a=a, + conv_weight=conv_weight, + conv_bias=conv_bias, + a_log=a_log, + dt_bias=dt_bias, + recurrent_state=None, + num_k_heads=num_k_heads, + num_v_heads=num_v_heads, + head_k_dim=head_k_dim, + head_v_dim=head_v_dim, + conv_kernel_size=conv_kernel_size, + chunk_size=chunk_size, + use_qk_norm_in_gdn=True, + compute_dtype=jnp.float32, + ) + np.testing.assert_allclose(chunk_states, exp_chunk_states, rtol=1e-5, atol=1e-5) + np.testing.assert_allclose(t_inv, exp_t_inv, rtol=1e-5, atol=1e-5) + + def test_decoupled_conv1d_gdn_kernel_gradient_with_initial_states(self): + """Verifies custom VJP gradients when initial conv_state and recurrent_state are provided.""" + batch_size = 1 + chunk_size = 32 + num_chunks = 2 + seq_len = num_chunks * chunk_size + num_k_heads = 1 + num_v_heads = 2 + head_k_dim = 64 + head_v_dim = 64 + conv_kernel_size = 4 + dim_size = num_k_heads * head_k_dim * 2 + num_v_heads * head_v_dim + + key = jax.random.PRNGKey(999) + k1, k2, k3, k4, k5, k6, k7, k8, k9, k10 = jax.random.split(key, 10) + + qkv = jax.random.normal(k1, (batch_size, seq_len, dim_size), dtype=jnp.float32) + b = jax.random.normal(k2, (batch_size, seq_len, num_v_heads), dtype=jnp.float32) + a = jax.random.normal(k3, (batch_size, seq_len, num_v_heads), dtype=jnp.float32) + conv_weight = jax.random.normal(k4, (conv_kernel_size, 1, dim_size), dtype=jnp.float32) + conv_bias = jax.random.normal(k5, (dim_size,), dtype=jnp.float32) + a_log = jax.random.normal(k6, (num_v_heads,), dtype=jnp.float32) + dt_bias = jax.random.normal(k7, (num_v_heads,), dtype=jnp.float32) + do = jax.random.normal(k8, (batch_size, seq_len, num_v_heads, head_v_dim), dtype=jnp.float32) + conv_state = jax.random.normal(k9, (batch_size, conv_kernel_size - 1, dim_size), dtype=jnp.float32) + recurrent_state = jax.random.normal( + k10, + (batch_size, num_v_heads, head_k_dim, head_v_dim), + dtype=jnp.float32, + ) + + def loss_pure(qkv_in, b_in, a_in, cw_in, cb_in, al_in, dt_in): + out, _ = gdn_bwd_pallas.pure_jax_decoupled_conv1d_gdn( + qkv=qkv_in, + b=b_in, + a=a_in, + conv_weight=cw_in, + conv_bias=cb_in, + a_log=al_in, + dt_bias=dt_in, + conv_state=conv_state, + recurrent_state=recurrent_state, + num_k_heads=num_k_heads, + num_v_heads=num_v_heads, + head_k_dim=head_k_dim, + head_v_dim=head_v_dim, + conv_kernel_size=conv_kernel_size, + chunk_size=chunk_size, + use_qk_norm_in_gdn=True, + compute_dtype=jnp.float32, + ) + return jnp.sum(out * do) + + def loss_gdn_kernel(qkv_in, b_in, a_in, cw_in, cb_in, al_in, dt_in): + out, _ = gdn_bwd_pallas.gdn_decoupled_conv1d( + qkv=qkv_in, + b=b_in, + a=a_in, + conv_weight=cw_in, + conv_bias=cb_in, + a_log=al_in, + dt_bias=dt_in, + conv_state=conv_state, + recurrent_state=recurrent_state, + num_k_heads=num_k_heads, + num_v_heads=num_v_heads, + head_k_dim=head_k_dim, + head_v_dim=head_v_dim, + conv_kernel_size=conv_kernel_size, + chunk_size=chunk_size, + use_qk_norm_in_gdn=True, + compute_dtype=jnp.float32, + ) + return jnp.sum(out * do) + + exp_grads = jax.grad(loss_pure, argnums=(0, 1, 2, 3, 4, 5, 6))(qkv, b, a, conv_weight, conv_bias, a_log, dt_bias) + act_grads = jax.grad(loss_gdn_kernel, argnums=(0, 1, 2, 3, 4, 5, 6))( + qkv, b, a, conv_weight, conv_bias, a_log, dt_bias + ) + + for exp_g, act_g in zip(exp_grads, act_grads): + self.assertIsNotNone(act_g) + np.testing.assert_allclose(exp_g, act_g, rtol=1e-3, atol=1e-3) + + def test_gdn_kernel_bwd_multi_group_head_parallel(self): + """Verifies multi-group head-parallel grid dispatch matches reference.""" + batch_size = 1 + chunk_size = 32 + num_chunks = 2 + seq_len = num_chunks * chunk_size + num_k_heads = 4 + num_v_heads = 8 + head_k_dim = 64 + head_v_dim = 64 + dim_size = num_k_heads * head_k_dim * 2 + num_v_heads * head_v_dim + + key = jax.random.PRNGKey(1234) + k1, k2, k3, k4, k5, k6, k7, k8 = jax.random.split(key, 8) + + qkv = jax.random.normal(k1, (batch_size, seq_len, dim_size), dtype=jnp.float32) + b = jax.random.normal(k2, (batch_size, seq_len, num_v_heads), dtype=jnp.float32) + a = jax.random.normal(k3, (batch_size, seq_len, num_v_heads), dtype=jnp.float32) + a_log = jax.random.normal(k4, (num_v_heads,), dtype=jnp.float32) + dt_bias = jax.random.normal(k5, (num_v_heads,), dtype=jnp.float32) + do = jax.random.normal(k6, (batch_size, seq_len, num_v_heads, head_v_dim), dtype=jnp.float32) + chunk_states = jax.random.normal(k7, (batch_size, num_chunks, num_v_heads, head_k_dim, head_v_dim), dtype=jnp.float32) + t_inv = jax.random.normal(k8, (batch_size, num_chunks, num_v_heads, chunk_size, chunk_size), dtype=jnp.float32) + + # 1. Dispatch with head_tile = 4 -> 2 head groups + dy1, db1, da1, dal1, ddt1 = gdn_bwd_pallas.pallas_gdn_bwd_kernel( + qkv_conv=qkv, + b=b, + a=a, + a_log=a_log, + dt_bias=dt_bias, + do=do, + chunk_states=chunk_states, + t_inv=t_inv, + num_v_heads=num_v_heads, + kq_head_dim=head_k_dim, + v_head_dim=head_v_dim, + chunk_size=chunk_size, + head_tile=4, + ) + + # 2. Dispatch with head_tile = 8 -> 1 head group + dy2, db2, da2, dal2, ddt2 = gdn_bwd_pallas.pallas_gdn_bwd_kernel( + qkv_conv=qkv, + b=b, + a=a, + a_log=a_log, + dt_bias=dt_bias, + do=do, + chunk_states=chunk_states, + t_inv=t_inv, + num_v_heads=num_v_heads, + kq_head_dim=head_k_dim, + v_head_dim=head_v_dim, + chunk_size=chunk_size, + head_tile=8, + ) + + np.testing.assert_allclose(dy1, dy2, rtol=1e-3, atol=1e-3) + np.testing.assert_allclose(db1, db2, rtol=1e-3, atol=1e-3) + np.testing.assert_allclose(da1, da2, rtol=1e-3, atol=1e-3) + np.testing.assert_allclose(dal1, dal2, rtol=1e-3, atol=1e-3) + np.testing.assert_allclose(ddt1, ddt2, rtol=1e-3, atol=1e-3) + + def test_gdn_kernel_bwd_variable_length_segment_ids_reset(self): + """Verifies segment_ids document boundaries reset carried state gradient to prevent leakage.""" + batch_size = 1 + chunk_size = 32 + num_chunks = 2 + seq_len = num_chunks * chunk_size + num_k_heads = 1 + num_v_heads = 2 + head_k_dim = 64 + head_v_dim = 64 + dim_size = num_k_heads * head_k_dim * 2 + num_v_heads * head_v_dim + + key = jax.random.PRNGKey(5678) + k1, k2, k3, k4, k5, k6, k7, k8 = jax.random.split(key, 8) + + qkv = jax.random.normal(k1, (batch_size, seq_len, dim_size), dtype=jnp.float32) + b = jax.random.normal(k2, (batch_size, seq_len, num_v_heads), dtype=jnp.float32) + a = jax.random.normal(k3, (batch_size, seq_len, num_v_heads), dtype=jnp.float32) + a_log = jax.random.normal(k4, (num_v_heads,), dtype=jnp.float32) + dt_bias = jax.random.normal(k5, (num_v_heads,), dtype=jnp.float32) + chunk_states = jax.random.normal(k6, (batch_size, num_chunks, num_v_heads, head_k_dim, head_v_dim), dtype=jnp.float32) + t_inv = jax.random.normal(k7, (batch_size, num_chunks, num_v_heads, chunk_size, chunk_size), dtype=jnp.float32) + + # Only chunk 1 has non-zero incoming gradients; chunk 0 do is all zeros + do_chunk1 = jax.random.normal(k8, (batch_size, chunk_size, num_v_heads, head_v_dim), dtype=jnp.float32) + do_chunk0 = jnp.zeros((batch_size, chunk_size, num_v_heads, head_v_dim), dtype=jnp.float32) + do = jnp.concatenate([do_chunk0, do_chunk1], axis=1) + + # 1. No document reset: gradient flows backwards from chunk 1 into chunk 0 + dy_no_reset, _, _, _, _ = gdn_bwd_pallas.pallas_gdn_bwd_kernel( + qkv_conv=qkv, + b=b, + a=a, + a_log=a_log, + dt_bias=dt_bias, + do=do, + chunk_states=chunk_states, + t_inv=t_inv, + num_v_heads=num_v_heads, + kq_head_dim=head_k_dim, + v_head_dim=head_v_dim, + chunk_size=chunk_size, + segment_ids=None, + ) + # Chunk 0 gradient is non-zero due to recurrent state carrying gradients from chunk 1 + self.assertGreater(float(jnp.max(jnp.abs(dy_no_reset[:, :chunk_size, :]))), 1e-4) + + # 2. With segment_ids boundary between chunk 0 (doc 0) and chunk 1 (doc 1) + seg_doc0 = jnp.zeros((batch_size, chunk_size), dtype=jnp.int32) + seg_doc1 = jnp.ones((batch_size, chunk_size), dtype=jnp.int32) + segment_ids = jnp.concatenate([seg_doc0, seg_doc1], axis=1) + + dy_reset, _, _, _, _ = gdn_bwd_pallas.pallas_gdn_bwd_kernel( + qkv_conv=qkv, + b=b, + a=a, + a_log=a_log, + dt_bias=dt_bias, + do=do, + chunk_states=chunk_states, + t_inv=t_inv, + num_v_heads=num_v_heads, + kq_head_dim=head_k_dim, + v_head_dim=head_v_dim, + chunk_size=chunk_size, + segment_ids=segment_ids, + ) + # Chunk 0 gradient is strictly zero because boundary reset eliminated cross-document leakage! + np.testing.assert_allclose( + dy_reset[:, :chunk_size, :], + jnp.zeros_like(dy_reset[:, :chunk_size, :]), + atol=1e-6, + ) + + def test_decoupled_conv1d_gdn_kernel_bwd_with_head_tile(self): + """Verifies decoupled_conv1d_gdn_bwd_kernel forwards head_tile correctly.""" + batch_size = 1 + chunk_size = 32 + num_chunks = 2 + seq_len = num_chunks * chunk_size + num_k_heads = 4 + num_v_heads = 8 + head_k_dim = 64 + head_v_dim = 64 + conv_kernel_size = 4 + dim_size = num_k_heads * head_k_dim * 2 + num_v_heads * head_v_dim + + key = jax.random.PRNGKey(999) + k1, k2, k3, k4, k5, k6, k7, k8, k9, k10 = jax.random.split(key, 10) + + pre_conv_qkv = jax.random.normal(k1, (batch_size, seq_len, dim_size), dtype=jnp.float32) + b = jax.random.normal(k2, (batch_size, seq_len, num_v_heads), dtype=jnp.float32) + a = jax.random.normal(k3, (batch_size, seq_len, num_v_heads), dtype=jnp.float32) + a_log = jax.random.normal(k4, (num_v_heads,), dtype=jnp.float32) + dt_bias = jax.random.normal(k5, (num_v_heads,), dtype=jnp.float32) + do = jax.random.normal(k6, (batch_size, seq_len, num_v_heads, head_v_dim), dtype=jnp.float32) + chunk_states = jax.random.normal(k7, (batch_size, num_chunks, num_v_heads, head_k_dim, head_v_dim), dtype=jnp.float32) + conv_weight = jax.random.normal(k8, (conv_kernel_size, 1, dim_size), dtype=jnp.float32) + conv_bias = jax.random.normal(k9, (dim_size,), dtype=jnp.float32) + t_inv = jax.random.normal(k10, (batch_size, num_chunks, num_v_heads, chunk_size, chunk_size), dtype=jnp.float32) + + res1 = gdn_bwd_pallas.decoupled_conv1d_gdn_bwd_kernel( + pre_conv_qkv=pre_conv_qkv, + b=b, + a=a, + a_log=a_log, + dt_bias=dt_bias, + do=do, + chunk_states=chunk_states, + conv_weight=conv_weight, + conv_bias=conv_bias, + t_inv=t_inv, + num_v_heads=num_v_heads, + kq_head_dim=head_k_dim, + v_head_dim=head_v_dim, + chunk_size=chunk_size, + head_tile=4, + ) + + res2 = gdn_bwd_pallas.decoupled_conv1d_gdn_bwd_kernel( + pre_conv_qkv=pre_conv_qkv, + b=b, + a=a, + a_log=a_log, + dt_bias=dt_bias, + do=do, + chunk_states=chunk_states, + conv_weight=conv_weight, + conv_bias=conv_bias, + t_inv=t_inv, + num_v_heads=num_v_heads, + kq_head_dim=head_k_dim, + v_head_dim=head_v_dim, + chunk_size=chunk_size, + head_tile=8, + ) + + for g1, g2 in zip(res1, res2): + if g1 is not None and g2 is not None: + np.testing.assert_allclose(g1, g2, rtol=2e-2, atol=2.5e-1) + + +if __name__ == "__main__": + absltest.main()