397B v7x RL trainer: all GDN, MoE and engine speedups (measured + pending), ordered by impact

src/maxtext/configs/custom_mesh_and_rule/cp-as-ep.yml

@@ -63,6 +63,9 @@ logical_axis_rules: [
                       # MoE Weights
                       ['exp', ['context', 'expert']],
                       ['embed_moe', ['fsdp']],
+                      # wo's dims: FSDP on its mlp dim (see RoutedMoE.wo_kernel_axes)
+                      ['mlp_moe_out', ['fsdp']],
+                      ['embed_moe_out', []],
                       # ==========================================
                       # Standard MLP / Dense Layers / Model Structure
                       # ==========================================

src/maxtext/configs/types.py

@@ -1523,12 +1523,13 @@ class Qwen3Next(BaseModel):
       "highest",
       description=(
           "Precision of the f32 matmuls that compose the GDN recurrent state across sequence-sharded context"
-          " parallelism (kernels/gdn/gdn_bwd/cp_gdn.py). On TPU 'highest' runs 6 bf16 passes and 'high' runs 3."
-          " Only the XLA composition math changes; the Pallas kernels keep their own precision. It applies only"
-          " with use_gdn_kernel=True and sequence-sharded GDN context parallelism: gdn_cp_mode='seq', or 'auto',"
-          " which shards the sequence unless the GDN context parallelism is 2 and divides the number of key"
-          " heads. It has no effect with use_gdn_kernel=False (that path's context parallelism always uses"
-          " HIGHEST) or with head-sharded context parallelism."
+          " parallelism (kernels/gdn/gdn_bwd/cp_gdn.py). On TPU 'highest' runs 6 bf16 passes and 'high' runs 3;"
+          " with 'high' the rank-local forward transition product and backward state gradient run as Pallas scans"
+          " (kernels/gdn/gdn_bwd/dh_scan.py) at the same 3 passes. Only the composition math changes; the Pallas GDN"
+          " kernels keep their own precision. It applies only with use_gdn_kernel=True and sequence-sharded GDN"
+          " context parallelism: gdn_cp_mode='seq', or 'auto', which shards the sequence unless the GDN context"
+          " parallelism is 2 and divides the number of key heads. It has no effect with use_gdn_kernel=False (that"
+          " path's context parallelism always uses HIGHEST) or with head-sharded context parallelism."
       ),
   )
 

src/maxtext/kernels/gdn/compute_gdn.py

@@ -51,18 +51,24 @@ def invert_triangular_matrix(
   num_blocks = chunk // block_size
 
   def local_forward_sub(t_mat: jax.Array, b_mat: jax.Array) -> jax.Array:
-    x_list = []
-    for i in range(block_size):
-      b_i = b_mat[:, i, :]
-      if i == 0:
-        x_i = b_i
+    # Right-looking substitution: once row j is final, subtract t[i, j] * x_j from every later row i.
+    # The rows stay in place as two sublane-aligned halves, so no step re-stacks the solved rows or
+    # moves a row of t from lanes to sublanes. Each row subtracts the same products as a left-looking
+    # solve, one at a time in column order.
+    half = block_size // 2
+    x_lo, x_hi = b_mat[:, :half, :], b_mat[:, half:, :]
+    t_lo, t_hi = t_mat[:, :half, :], t_mat[:, half:, :]
+    rows_lo = jax.lax.broadcasted_iota(jnp.int32, (1, half, 1), 1)
+    rows_hi = jax.lax.broadcasted_iota(jnp.int32, (1, block_size - half, 1), 1)
+    for j in range(block_size - 1):
+      if j < half:
+        x_j = x_lo[:, j : j + 1, :]
+        x_lo = jnp.where(rows_lo > j, x_lo - t_lo[:, :, j : j + 1] * x_j, x_lo)
+        x_hi = x_hi - t_hi[:, :, j : j + 1] * x_j
       else:
-        stacked_x = jnp.stack(x_list, axis=1)
-        all_prev_t = t_mat[:, i, :i]
-        prev_sum = jnp.sum(all_prev_t[..., None] * stacked_x, axis=1)
-        x_i = b_i - prev_sum
-      x_list.append(x_i)
-    return jnp.stack(x_list, axis=1)
+        x_j = x_hi[:, j - half : j - half + 1, :]
+        x_hi = jnp.where(rows_hi > j - half, x_hi - t_hi[:, :, j : j + 1] * x_j, x_hi)
+    return jnp.concatenate([x_lo, x_hi], axis=1)
 
   x_blocks = []
   iota_r = jax.lax.broadcasted_iota(jnp.int32, t.shape, 1)
@@ -114,8 +120,8 @@ def chunked_gdn_per_seq(
     cfg: config.GDNConfig,
     seg_c_col: jax.Array | None = None,  # [1, chunk, 1]
     seg_aux_col: jax.Array | None = None,  # [1, chunk, 1]
-) -> tuple[jax.Array, jax.Array, jax.Array]:
-  """Perform chunked GDN over input [num_heads, chunk, head_dim]."""
+) -> tuple[jax.Array | None, jax.Array, jax.Array]:
+  """Perform chunked GDN over input [num_heads, chunk, head_dim]; out is None in PER_SEQ_STATES mode."""
   # bf16 activations keep single-pass MXU matmuls even with an f32 compute dtype
   # (HIGHEST would triple the forward MXU cost); HIGHEST needs f32 activations and compute.
   dot_prec = (
@@ -236,23 +242,36 @@ def chunked_gdn_per_seq(
   # [num_v_heads, chunk, v_head_dim]
   u, w = jnp.split(merged_uw, [cfg.v_head_dim], axis=-1)
 
-  # [num_v_heads, chunk, kq_head_dim]
-  q_large_gating = q_repeat * gating_forward
-  # NOTE: Concatenate lhs with same rhs to leverage weight
-  # stationary architecture.
-  # [num_v_heads, 2 * chunk, kq_head_dim]
-  merged_w_q = jnp.concat([w, q_large_gating], axis=1)
-  # [num_v_heads, 2 * chunk, v_head_dim]
-  merged_ws_out_updated = jax.lax.dot(
-      merged_w_q,
-      state_prev,
-      dimension_numbers=(((2,), (1,)), ((0,), (0,))),
-      precision=dot_prec,
-      preferred_element_type=jnp.float32,
-  )
+  states_only = cfg.mode == config.GDNMode.PER_SEQ_STATES
+  if states_only:
+    # No output: only the w rows of the merged matmul below. Each row's dot product is the same,
+    # so ws is bit-identical.
+    # [num_v_heads, chunk, v_head_dim]
+    ws = jax.lax.dot(
+        w,
+        state_prev,
+        dimension_numbers=(((2,), (1,)), ((0,), (0,))),
+        precision=dot_prec,
+        preferred_element_type=jnp.float32,
+    )
+  else:
+    # [num_v_heads, chunk, kq_head_dim]
+    q_large_gating = q_repeat * gating_forward
+    # NOTE: Concatenate lhs with same rhs to leverage weight
+    # stationary architecture.
+    # [num_v_heads, 2 * chunk, kq_head_dim]
+    merged_w_q = jnp.concat([w, q_large_gating], axis=1)
+    # [num_v_heads, 2 * chunk, v_head_dim]
+    merged_ws_out_updated = jax.lax.dot(
+        merged_w_q,
+        state_prev,
+        dimension_numbers=(((2,), (1,)), ((0,), (0,))),
+        precision=dot_prec,
+        preferred_element_type=jnp.float32,
+    )
 
-  # NOTE: Splitting along non sublane/lane dim is free.
-  ws, out_updated = jnp.split(merged_ws_out_updated, 2, axis=1)
+    # NOTE: Splitting along non sublane/lane dim is free.
+    ws, out_updated = jnp.split(merged_ws_out_updated, 2, axis=1)
   ws = ws.astype(cfg.dtypes.compute)
 
   # [num_v_heads, chunk, v_head_dim]
@@ -273,6 +292,8 @@ def chunked_gdn_per_seq(
   # [num_v_heads, kq_head_dim, v_head_dim]
   state_updated = state_prev * gating_last
   state = state_updated + state_new
+  if states_only:
+    return None, state, t_inv
 
   # [num_kq_heads, chunk, chunk]
   out_qk = jax.lax.dot(
@@ -314,8 +335,8 @@ def chunked_gdn(
     a_log: jax.Array,
     dt_bias: jax.Array,
     cfg: config.GDNConfig,
-) -> tuple[jax.Array, jax.Array, jax.Array]:
-  """Perform chunked GDN over input [seq, num_heads, chunk, head_dim]."""
+) -> tuple[jax.Array | None, jax.Array, jax.Array]:
+  """Perform chunked GDN over input [seq, num_heads, chunk, head_dim]; out is None in PER_SEQ_STATES mode."""
 
   mask_dtype = get_mask_dtype(cfg.dtypes.compute)
   iota = jax.lax.broadcasted_iota(mask_dtype, (cfg.seq_tile_size, 1, cfg.chunk_size, 1), 2)
@@ -373,10 +394,11 @@ def chunked_gdn(
         seg_c_col=seg_c_per_seq[idx] if cfg.has_seg_ids else None,
         seg_aux_col=seg_aux_per_seq[idx] if cfg.has_seg_ids else None,
     )
-    out_list.append(out.swapaxes(0, 1))
+    if out is not None:
+      out_list.append(out.swapaxes(0, 1))
     state_list.append(state)
     t_inv_list.append(t_inv)
-  out = jnp.stack(out_list, axis=0)
+  out = jnp.stack(out_list, axis=0) if out_list else None
   state = jnp.stack(state_list, axis=0)
   t_inv = jnp.stack(t_inv_list, axis=0)
   return out, state, t_inv

src/maxtext/kernels/gdn/config.py

@@ -24,14 +24,17 @@
 from jax.experimental.pallas import tpu as pltpu
 import jax.numpy as jnp
 
-DEFAULT_VMEM_LIMIT_FACTOR: float = 0.90
+# The 64-v-head prefill kernel needs about 58 MiB of the 64 MiB VMEM on TPU7x.
+DEFAULT_VMEM_LIMIT_FACTOR: float = 0.95
 
 
 class GDNMode(enum.StrEnum):
   """Execution mode for GDN kernel."""
 
   BATCHED = enum.auto()
   PER_SEQ = enum.auto()
+  # PER_SEQ that only carries the conv/recurrent states and t_inv: no output, no chunk states.
+  PER_SEQ_STATES = enum.auto()
 
   def get_seq_tile_size(self, tile_size: int) -> int:
     if self == GDNMode.BATCHED:
@@ -82,6 +85,8 @@ class GDNConfig:
   num_buffers: int = 2
   has_seg_ids: bool = False
   use_qk_norm_in_gdn: bool = True
+  # Prefill-only tiles are whole chunks: qkv is DMA'd in act_in dtype and upcast in VMEM.
+  is_prefill_only: bool = False
 
   @property
   def chunk_size(self) -> int:
@@ -150,7 +155,13 @@ def get_scratch_shape_dict(self) -> dict[str, Any]:
       carry_conv_scratch = pltpu.VMEM(conv_shape, jnp.float32)
       carry_recurrent_scratch = pltpu.VMEM(recurrent_shape, jnp.float32)
 
+    # Compact fp32 copy of a prefill qkv chunk, read by Conv1D and the q/k/v loads.
+    qkv_scratch = None
+    if self.is_prefill_only:
+      qkv_scratch = pltpu.VMEM((self.seq_tile_size, self.chunk_size, 1, self.dim_size), jnp.float32)
+
     return {
         "carry_conv_scratch_ref": carry_conv_scratch,
         "carry_recurrent_scratch_ref": carry_recurrent_scratch,
+        "qkv_scratch_ref": qkv_scratch,
     }