[Stacked PR 2/5] Import local Tokamax GDN forward kernel with custom remat for backward pass support by Rohan-Bierneni · Pull Request #5152 · AI-Hypercomputer/maxtext

Rohan-Bierneni · 2026-09-06T04:30:16Z

Stacked PR Chain

Stack	Status	Branch	Base	PR
PR 1	🔗 Base PR	`rbierneni-gdn-1-ci-hygiene`	`main`	#5151 - [Stacked PR 1/5] Fix upstream pyink formatting and GMM v2 compatibility for Qwen3.5
PR 2	🚀 This PR	`rbierneni-gdn-2-fwd-kernel`	`rbierneni-gdn-1-ci-hygiene`	#5152 - [Stacked PR 2/5] Import local Tokamax GDN forward kernel with custom remat for backward pass support
PR 3	⏳ Stacked on PR 2	`rbierneni-gdn-3-bwd-kernel`	`rbierneni-gdn-2-fwd-kernel`	#5153 - [Stacked PR 3/5] Add decoupled GDN Pallas backward pass kernel and parity tests
PR 4	⏳ Stacked on PR 3	`rbierneni-gdn-4-model-integration`	`rbierneni-gdn-3-bwd-kernel`	#5154 - [Stacked PR 4/5] Enable GDN backward pass kernel and hybrid precision in Qwen3 model
PR 5	⏳ Stacked on PR 4	`rbierneni-gdnv3-bwd`	`rbierneni-gdn-4-model-integration`	#5098 - [Stacked PR 5/5] Add GDN backward pass 8k and 64k latency and memory benchmark suite

Description

This is PR 2 of 5 in the stacked series enabling the Pallas Gated Delta Net (GDN) backward pass kernel in MaxText.

This PR imports a local copy of the Tokamax Pallas Mosaic TPU forward kernel under src/maxtext/models/kernels/gdn/ to enable backward pass support with custom rematerialization (remat) in the forward pass:

Local Tokamax Forward Kernel Import: Vendors the Pallas Mosaic TPU GDN forward kernel directly into MaxText, ensuring consistent TPU VMEM tiling and execution without external package dependency mismatches.
Custom Rematerialization (Remat) for Backward Pass Support: Implements custom rematerialization in the forward pass by caching intermediate chunk states and triangular inverse matrices ($T^{-1}$) in residuals. This avoids expensive recomputation of forward triangular solves during the backward pass (PR 3).
Hardware-Optimized Tiling & Memory Management:
- tiling.py & config.py: Computes optimal chunk sizes, block dimensions, and VMEM layout heuristics across TPU generations (v4, v5e, v5p, v6e).
- memory_ref.py & vmem_ldst.py: Manages structured VMEM load/store operations, double-buffering, and register allocations.

Files Changed

src/maxtext/models/kernels/gdn/compute_conv1d.py: In-VMEM Conv1D forward primitives.
src/maxtext/models/kernels/gdn/compute_gdn.py: Core GDN chunk recurrence and triangular matrix inversion math with $T^{-1}$ caching.
src/maxtext/models/kernels/gdn/config.py: GDN kernel configuration dataclasses and validation.
src/maxtext/models/kernels/gdn/memory_ref.py: TPU VMEM buffer reference and allocation abstraction.
src/maxtext/models/kernels/gdn/metadata.py: Kernel launch metadata and grid layout definitions.
src/maxtext/models/kernels/gdn/pallas_mosaic_tpu.py: Low-level Pallas Mosaic TPU invocation bindings.
src/maxtext/models/kernels/gdn/tiling.py: Chunk and sublane tiling heuristics.
src/maxtext/models/kernels/gdn/vmem_ldst.py: Direct VMEM memory load/store operations.
src/maxtext/models/kernels/gdn/wrapper.py: High-level Python functional interface for GDN forward execution with custom remat caching.

Tests

Verified with pre-commit run --files ... (all hooks passed).
Verified forward execution and residual caching.

Checklist

I have performed a self-review of my code.
I have necessary comments in my code, particularly in hard-to-understand areas.
I have run end-to-end tests and verified pre-commit linters pass.

gemini-code-assist

Code Review

This pull request introduces a Pallas Mosaic TPU kernel implementation for the Causal Conv1D Gated Delta Rule (GDN) in MaxText, adding modules for Conv1D/GDN computation, dynamic tiling, VMEM load/store logic, and configuration. The review feedback highlights several critical issues and improvement opportunities: replacing the non-existent jax.Ref type annotations with pl.Ref across multiple files, avoiding compile-time tracer conversion errors with distribution[0], replacing the memory-intensive fused_transpose_broadcast with x.swapaxes to prevent VMEM overflow, simplifying the start_idx calculation and correcting shape comments in compute_conv1d.py, utilizing jnp.cumsum instead of a manual loop in compute_gdn.py, and adding an assertion to ensure chunk sizes are multiples of the block size.

gemini-code-assist · 2026-09-06T04:33:28Z

+    from . import memory_ref
+
+
+def load_as_qkv_large(qkv_vmem_ref: jax.Ref, cfgs: config.GDNConfig) -> tuple[jax.Array, jax.Array, jax.Array]:


jax.Ref is not a standard JAX type and does not exist in the public jax namespace. It should be pl.Ref (from jax.experimental.pallas). This will cause type errors or AttributeError at runtime/import time.

Suggested change

def load_as_qkv_large(qkv_vmem_ref: jax.Ref, cfgs: config.GDNConfig) -> tuple[jax.Array, jax.Array, jax.Array]:

def load_as_qkv_large(qkv_vmem_ref: pl.Ref, cfgs: config.GDNConfig) -> tuple[jax.Array, jax.Array, jax.Array]:

gemini-code-assist · 2026-09-06T04:33:28Z

+def load_activation_as_compact(
+    qkv_vreg: jax.Array,
+    qkv_vmem_ref: jax.Ref,
+    b_vmem_ref: jax.Ref,
+    a_vmem_ref: jax.Ref,
+    cfgs: config.GDNConfig,


jax.Ref is not a standard JAX type and does not exist in the public jax namespace. It should be pl.Ref (from jax.experimental.pallas). This will cause type errors or AttributeError at runtime/import time.

def load_activation_as_compact(\n qkv_vreg: jax.Array,\n qkv_vmem_ref: pl.Ref,\n b_vmem_ref: pl.Ref,\n a_vmem_ref: pl.Ref,\n cfgs: config.GDNConfig,

gemini-code-assist · 2026-09-06T04:33:28Z

+  if not is_prefill_only:
+    try:
+      if int(distribution[0]) == 0:
+        is_prefill_only = True
+    except (TypeError, ValueError, jax.errors.TracerIntegerConversionError):
+      pass


Since distribution is a tracer under JIT, int(distribution[0]) will always raise TracerIntegerConversionError at compile time, meaning this check is dead code and the except block will always catch it and do pass. If you want to conditionally run the batched kernel dynamically, you should use jax.lax.cond. Otherwise, remove this misleading block and rely on Pallas handling 0-sized grids gracefully.

Suggested change

if not is_prefill_only:

try:

if int(distribution[0]) == 0:

is_prefill_only = True

except (TypeError, ValueError, jax.errors.TracerIntegerConversionError):

pass

# Under JIT, distribution[0] is a tracer, so we cannot dynamically set is_prefill_only in Python control flow.\n # We rely on Pallas handling 0-sized grids gracefully.

gemini-code-assist · 2026-09-06T04:33:28Z

+def fused_transpose_broadcast(x: jax.Array, src_dim: int, dst_dim: int) -> jax.Array:
+  """Perform 1D transpose where results are broadcasted along src_dim."""
+  assert x.shape[dst_dim] == 1
+
+  dtype = x.dtype
+  mask_dtype = get_mask_dtype(dtype)
+  mask_shape = list(x.shape)
+  mask_size = mask_shape[src_dim]
+  mask_shape[dst_dim] = mask_size
+  src_mask = jax.lax.broadcasted_iota(mask_dtype, mask_shape, src_dim)
+  dst_mask = jax.lax.broadcasted_iota(mask_dtype, mask_shape, dst_dim)
+  mask = src_mask == dst_mask
+  return jnp.where(mask, x, 0).sum(axis=src_dim, keepdims=True, dtype=dtype)


The entire fused_transpose_broadcast function can be replaced with x.swapaxes(src_dim, dst_dim). The current implementation creates a 3D intermediate tensor of size O(N^3) (e.g., up to 8MB for 128x128x128), which is extremely dangerous for VMEM capacity on TPU and can easily cause VMEM overflow or compilation failures. Using swapaxes is standard, highly optimized, and memory-efficient.

def fused_transpose_broadcast(x: jax.Array, src_dim: int, dst_dim: int) -> jax.Array:\n \"\"\"Perform 1D transpose where results are broadcasted along src_dim.\"\"\"\n return x.swapaxes(src_dim, dst_dim)

gemini-code-assist · 2026-09-06T04:33:28Z

+  return q_large, k_large, v_large
+
+
+def load_as_qkv_compact(qkv_vmem_ref: jax.Ref, cfg: config.GDNConfig) -> tuple[jax.Array, jax.Array, jax.Array]:


jax.Ref is not a standard JAX type and does not exist in the public jax namespace. It should be pl.Ref (from jax.experimental.pallas). This will cause type errors or AttributeError at runtime/import time.

Suggested change

def load_as_qkv_compact(qkv_vmem_ref: jax.Ref, cfg: config.GDNConfig) -> tuple[jax.Array, jax.Array, jax.Array]:

def load_as_qkv_compact(qkv_vmem_ref: pl.Ref, cfg: config.GDNConfig) -> tuple[jax.Array, jax.Array, jax.Array]:

gemini-code-assist · 2026-09-06T04:33:29Z

+def load_and_select_states(
+    metadata_ref: memory_ref.MetadataRef,
+    p_id: jax.Array,
+    conv_state_slot_ref: jax.Ref,
+    recurrent_slot_ref: jax.Ref,
+    carry_conv_scratch_ref: jax.Ref | None,
+    carry_recurrent_scratch_ref: jax.Ref | None,
+    cfg: config.GDNConfig,


jax.Ref is not a standard JAX type and does not exist in the public jax namespace. It should be pl.Ref (from jax.experimental.pallas). This will cause type errors or AttributeError at runtime/import time.

def load_and_select_states(\n metadata_ref: memory_ref.MetadataRef,\n p_id: jax.Array,\n conv_state_slot_ref: pl.Ref,\n recurrent_slot_ref: pl.Ref,\n carry_conv_scratch_ref: pl.Ref | None,\n carry_recurrent_scratch_ref: pl.Ref | None,\n cfg: config.GDNConfig,

gemini-code-assist · 2026-09-06T04:33:29Z

+    end_idx = c_idx + cfg.prev_kernel_size
+    start_idx = 1 + end_idx - cfg.kernel_size


The calculation of start_idx can be simplified. Since cfg.prev_kernel_size = cfg.kernel_size - 1, the expression 1 + end_idx - cfg.kernel_size mathematically simplifies to c_idx. Simplifying this improves readability and removes the unused end_idx variable.

Suggested change

end_idx = c_idx + cfg.prev_kernel_size

start_idx = 1 + end_idx - cfg.kernel_size

start_idx = c_idx

gemini-code-assist · 2026-09-06T04:33:29Z

+def causal_conv1d(
+    real_sizes: jax.Array,  # [seq]
+    lhs: jax.Array,  # [seq, chunk, q, dim_size]
+    conv_weight: jax.Array,  # [prev_kernel_size, 1, dim_size]


The comment indicates that conv_weight has shape [prev_kernel_size, 1, dim_size], but it actually has shape [kernel_size, 1, dim_size]. Updating the comment prevents confusion.

Suggested change

conv_weight: jax.Array, # [prev_kernel_size, 1, dim_size]

conv_weight: jax.Array, # [kernel_size, 1, dim_size]

gemini-code-assist · 2026-09-06T04:33:29Z

+  g_cum_sum_list = [gating_log[:, :1]]
+  for row in range(1, cfg.chunk_size):
+    g_cum_sum_list.append(g_cum_sum_list[-1] + gating_log[:, row : row + 1])
+  # [1, chunk, num_v_heads]
+  g_cum_sum_log = jnp.concat(g_cum_sum_list, axis=1)


The Python loop and list concatenation can be replaced with jnp.cumsum(gating_log, axis=1). This is much cleaner, faster, and avoids creating many intermediate nodes in the Jaxpr from the Python loop, improving compile times and HLO optimization.

Suggested change

g_cum_sum_list = [gating_log[:, :1]]

for row in range(1, cfg.chunk_size):

g_cum_sum_list.append(g_cum_sum_list[-1] + gating_log[:, row : row + 1])

# [1, chunk, num_v_heads]

g_cum_sum_log = jnp.concat(g_cum_sum_list, axis=1)

# [1, chunk, num_v_heads]\n g_cum_sum_log = jnp.cumsum(gating_log, axis=1)

gemini-code-assist · 2026-09-06T04:33:29Z

+  out_dtype = t.dtype
+  chunk = t.shape[-1]
+  block_size = min(block_size, chunk)
+  num_blocks = chunk // block_size


Add an assertion to ensure that chunk is a multiple of block_size. If a non-power-of-two chunk size is used in the future, a mismatch would cause silent correctness bugs because num_blocks would discard the remainder.

assert chunk % block_size == 0, f"chunk size {chunk} must be a multiple of block_size {block_size}"\n num_blocks = chunk // block_size

…x caching - Integrate local Pallas Mosaic TPU forward kernel for Gated Delta Net (GDN). - Implement VMEM allocation strategies, tiling heuristics, and memory reference utilities optimized for TPU v4/v5/v6e architectures. - Add support for caching chunk states and triangular inverse matrices (t_inv) in forward residuals to enable decoupled backward execution without redundant forward recomputation.

codecov · 2026-09-06T07:14:14Z

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

Rohan-Bierneni requested review from NuojCheng, RissyRan, aireenmei, bvandermoon, gagika, gobbleturk, huytransformer, igorts-git, jiangjy1982, parambole, richjames0, shralex, shuningjin and xibinliu as code owners September 6, 2026 04:30

gemini-code-assist Bot reviewed Sep 6, 2026

View reviewed changes

Rohan-Bierneni changed the title ~~[Stacked PR 2/5] Add local Pallas GDN forward kernel with T_inv matrix caching~~ [Stacked PR 2/5] Import local Tokamax GDN forward kernel with custom remat for backward pass support Sep 6, 2026

Rohan-Bierneni force-pushed the rbierneni-gdn-2-fwd-kernel branch from 6a7e646 to 568f432 Compare September 6, 2026 07:00

Navigation Menu

Uh oh!

[Stacked PR 2/5] Import local Tokamax GDN forward kernel with custom remat for backward pass support - #5152

[Stacked PR 2/5] Import local Tokamax GDN forward kernel with custom remat for backward pass support#5152
Rohan-Bierneni wants to merge 1 commit into
rbierneni-gdn-1-ci-hygienefrom
rbierneni-gdn-2-fwd-kernel

Rohan-Bierneni commented Sep 6, 2026 •

edited

Loading

Uh oh!

gemini-code-assist Bot left a comment

Uh oh!

gemini-code-assist Bot Sep 6, 2026

Uh oh!

gemini-code-assist Bot Sep 6, 2026

Uh oh!

gemini-code-assist Bot Sep 6, 2026

Uh oh!

gemini-code-assist Bot Sep 6, 2026

Uh oh!

gemini-code-assist Bot Sep 6, 2026

Uh oh!

gemini-code-assist Bot Sep 6, 2026

Uh oh!

gemini-code-assist Bot Sep 6, 2026

Uh oh!

gemini-code-assist Bot Sep 6, 2026

Uh oh!

gemini-code-assist Bot Sep 6, 2026

Uh oh!

gemini-code-assist Bot Sep 6, 2026

Uh oh!

codecov Bot commented Sep 6, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

		from . import memory_ref


		def load_as_qkv_large(qkv_vmem_ref: jax.Ref, cfgs: config.GDNConfig) -> tuple[jax.Array, jax.Array, jax.Array]:

	def load_as_qkv_large(qkv_vmem_ref: jax.Ref, cfgs: config.GDNConfig) -> tuple[jax.Array, jax.Array, jax.Array]:
	def load_as_qkv_large(qkv_vmem_ref: pl.Ref, cfgs: config.GDNConfig) -> tuple[jax.Array, jax.Array, jax.Array]:

		return q_large, k_large, v_large


		def load_as_qkv_compact(qkv_vmem_ref: jax.Ref, cfg: config.GDNConfig) -> tuple[jax.Array, jax.Array, jax.Array]:

	def load_as_qkv_compact(qkv_vmem_ref: jax.Ref, cfg: config.GDNConfig) -> tuple[jax.Array, jax.Array, jax.Array]:
	def load_as_qkv_compact(qkv_vmem_ref: pl.Ref, cfg: config.GDNConfig) -> tuple[jax.Array, jax.Array, jax.Array]:

		end_idx = c_idx + cfg.prev_kernel_size
		start_idx = 1 + end_idx - cfg.kernel_size

	end_idx = c_idx + cfg.prev_kernel_size
	start_idx = 1 + end_idx - cfg.kernel_size
	start_idx = c_idx

	conv_weight: jax.Array, # [prev_kernel_size, 1, dim_size]
	conv_weight: jax.Array, # [kernel_size, 1, dim_size]

Uh oh!

Conversation

Rohan-Bierneni commented Sep 6, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Stacked PR Chain

Description

Files Changed

Tests

Checklist

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist Bot Sep 6, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Sep 6, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Sep 6, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Sep 6, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Sep 6, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Sep 6, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Sep 6, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Sep 6, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Sep 6, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Sep 6, 2026

Choose a reason for hiding this comment

Uh oh!

codecov Bot commented Sep 6, 2026

Codecov Report

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Rohan-Bierneni commented Sep 6, 2026 •

edited

Loading