From 568f432ab949e408a38bf9b399574aceb45963e3 Mon Sep 17 00:00:00 2001 From: Rohan Bierneni Date: Sun, 6 Sep 2026 04:30:13 +0000 Subject: [PATCH] [Stacked PR 2/5] Add local Pallas GDN forward kernel with T_inv matrix 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. --- .../models/kernels/gdn/compute_conv1d.py | 75 +++ src/maxtext/models/kernels/gdn/compute_gdn.py | 466 +++++++++++++ src/maxtext/models/kernels/gdn/config.py | 154 +++++ src/maxtext/models/kernels/gdn/memory_ref.py | 616 ++++++++++++++++++ src/maxtext/models/kernels/gdn/metadata.py | 144 ++++ .../models/kernels/gdn/pallas_mosaic_tpu.py | 106 +++ src/maxtext/models/kernels/gdn/tiling.py | 337 ++++++++++ src/maxtext/models/kernels/gdn/vmem_ldst.py | 263 ++++++++ src/maxtext/models/kernels/gdn/wrapper.py | 539 +++++++++++++++ 9 files changed, 2700 insertions(+) create mode 100644 src/maxtext/models/kernels/gdn/compute_conv1d.py create mode 100644 src/maxtext/models/kernels/gdn/compute_gdn.py create mode 100644 src/maxtext/models/kernels/gdn/config.py create mode 100644 src/maxtext/models/kernels/gdn/memory_ref.py create mode 100644 src/maxtext/models/kernels/gdn/metadata.py create mode 100644 src/maxtext/models/kernels/gdn/pallas_mosaic_tpu.py create mode 100644 src/maxtext/models/kernels/gdn/tiling.py create mode 100644 src/maxtext/models/kernels/gdn/vmem_ldst.py create mode 100644 src/maxtext/models/kernels/gdn/wrapper.py diff --git a/src/maxtext/models/kernels/gdn/compute_conv1d.py b/src/maxtext/models/kernels/gdn/compute_conv1d.py new file mode 100644 index 000000000..c57abf2ee --- /dev/null +++ b/src/maxtext/models/kernels/gdn/compute_conv1d.py @@ -0,0 +1,75 @@ +# 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. +# ============================================================================== + +"""In-VMEM causal depthwise Conv1D computation.""" + +import jax +import jax.numpy as jnp + +try: + from maxtext.models.kernels.gdn import config +except (ImportError, ModuleNotFoundError): + try: + from maxtext.src.maxtext.models.kernels.gdn import config + except (ImportError, ModuleNotFoundError): + from . import config + + +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] + conv_bias: jax.Array | None, # [dim_size] + cfg: config.GDNConfig, +) -> tuple[jax.Array, jax.Array]: + """Perform causal Conv1D. Returns Conv1D output and convolution states.""" + + assert lhs.ndim == 4 + + out_list = [] + + for c_idx in range(cfg.chunk_size): + out = jnp.zeros((cfg.seq_tile_size, 1, cfg.dim_size), jnp.float32) + + end_idx = c_idx + cfg.prev_kernel_size + start_idx = 1 + end_idx - cfg.kernel_size + for k in range(cfg.kernel_size): + lhs_curr = lhs[:, start_idx + k] + out += lhs_curr * conv_weight[k : k + 1] + + if conv_bias is not None: + out += conv_bias.reshape(1, 1, -1) + + out_list.append(out) + + # Last prev_kernel_size elements needs to be returned as conv_state. However, + # real_sizes may be smaller than chunk_size. Therefore, slicing last + # prev_kernel_size elements does not guarantee numeric correctness. Instead, + # kernel iterate each rows and perform masking to fetch correct values. + # NOTE: lhs[:, : prev_kernel_size] can be skipped since they were loaded from + # previous conv states. + new_conv_state = lhs[:, 1 : cfg.kernel_size] + real_sizes = real_sizes.reshape(-1, 1, 1, 1) + # NOTE: Even though for loop is invoked twice, since they are static loops, + # compiler will perform loop fusion. + for c_idx in range(2, cfg.chunk_size + 1): + row_end = c_idx + cfg.prev_kernel_size + new_conv_state = jnp.where( + c_idx == real_sizes, + lhs[:, c_idx:row_end], + new_conv_state, + ) + + return jnp.stack(out_list, axis=1), new_conv_state diff --git a/src/maxtext/models/kernels/gdn/compute_gdn.py b/src/maxtext/models/kernels/gdn/compute_gdn.py new file mode 100644 index 000000000..08356d237 --- /dev/null +++ b/src/maxtext/models/kernels/gdn/compute_gdn.py @@ -0,0 +1,466 @@ +# 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. +# ============================================================================== + +"""Core GDN forward computation with triangular inverse matrix caching.""" + +import jax +import jax.numpy as jnp + +try: + from maxtext.models.kernels.gdn import config +except (ImportError, ModuleNotFoundError): + try: + from maxtext.src.maxtext.models.kernels.gdn import config + except (ImportError, ModuleNotFoundError): + from . import config + + +def l2_norm(x: jax.Array, eps: float = 1e-6) -> jax.Array: + norm = jnp.sqrt(jnp.sum(x * x, axis=-1, keepdims=True, dtype=x.dtype) + eps) + return x / norm + + +def get_mask_dtype(dtype: jnp.dtype) -> jnp.dtype: + match jnp.dtype(dtype).itemsize: + case 4: + return jnp.int32 + case 2: + return jnp.int16 + case _: + raise ValueError(f"Unsupported dtype: {dtype}") + + +# NOTE: Fork of recurrent_scan_v2.py but applied various optimizations. +def invert_triangular_matrix(t: jax.Array, block_size: int = 16) -> jax.Array: + """Compute invert matrix of a given triauglar matrix.""" + + # NOTE: if chunk_size=1, compiler will perform DCE. + out_dtype = t.dtype + chunk = t.shape[-1] + block_size = min(block_size, chunk) + 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 + 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_blocks = [] + iota_r = jax.lax.broadcasted_iota(jnp.int32, t.shape, 1) + iota_c = jax.lax.broadcasted_iota(jnp.int32, t.shape, 2) + identity_mask = jnp.where(iota_r == iota_c, 1.0, 0.0) + for i in range(num_blocks): + start, end = i * block_size, (i + 1) * block_size + e_block = identity_mask[:, start:end, :] + + if i == 0: + target_b = e_block + else: + interaction_t = t[:, start:end, :start] + solved_x = jnp.concatenate(x_blocks, axis=1) + prev_sum = jax.lax.dot( + interaction_t, + solved_x, + dimension_numbers=(((2,), (1,)), ((0,), (0,))), + preferred_element_type=jnp.float32, + ) + target_b = e_block - prev_sum + + # NOTE: Utilize fp32 to minimize cost of sublane rolling. + local_t = t[:, start:end, start:end].astype(jnp.float32) + x_block = local_forward_sub(local_t, target_b) + x_blocks.append(x_block.astype(out_dtype)) + + return jnp.concatenate(x_blocks, axis=1) + + +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) + + +def chunked_gdn_per_seq( + q_large: jax.Array, # [num_kq_heads, chunk, kq_head_dim] + k_large: jax.Array, # [num_kq_heads, chunk, kq_head_dim] + v_large: jax.Array, # [num_v_heads, chunk, v_head_dim] + gating_log: jax.Array, # [1, 1, num_v_heads] + beta: jax.Array, # [1, 1, num_v_heads] + state_prev: jax.Array, # [num_v_heads, kq_head_dim, v_head_dim] + cfg: config.GDNConfig, +) -> tuple[jax.Array, jax.Array, jax.Array]: + """Perform chunked GDN over input [num_heads, chunk, head_dim].""" + + # NOTE: Repeat along non lane/sublane dim is free. + q_repeat = jnp.repeat(q_large, cfg.v_per_kq_head, axis=0) + k_repeat = jnp.repeat(k_large, cfg.v_per_kq_head, axis=0) + + # Compute cumulative sum of decay. + # [1, 1, num_v_heads] + 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) + + # [num_v_heads, chunk, 1] + g_cum_sum_log = fused_transpose_broadcast(g_cum_sum_log, src_dim=2, dst_dim=0) + g_cum_sum_log = g_cum_sum_log[: cfg.num_v_heads] + beta = fused_transpose_broadcast(beta, src_dim=2, dst_dim=0) + beta_large = beta[: cfg.num_v_heads] + + # [num_v_heads, 1, chunk] + g_cum_sum_log_t = fused_transpose_broadcast(g_cum_sum_log, src_dim=1, dst_dim=2) + # [num_v_heads, chunk, chunk] + g_cum_sum_diff_log = g_cum_sum_log - g_cum_sum_log_t + gating_map = jnp.exp(g_cum_sum_diff_log) + # [num_v_heads, chunk, 1] + gating_backward = jnp.exp(-g_cum_sum_diff_log[..., -1:]) + # [num_v_heads, chunk, 1] + gating_forward = jnp.exp(g_cum_sum_log) + # [num_v_heads, 1, 1] + gating_last = gating_forward[:, -1:] + + mask_dtype = get_mask_dtype(cfg.dtypes.compute) + iota_r = jax.lax.broadcasted_iota(mask_dtype, gating_map.shape, 1) + iota_c = jax.lax.broadcasted_iota(mask_dtype, gating_map.shape, 2) + identity_mask = iota_r == iota_c + strictly_lower_mask = iota_r > iota_c + lower_mask = iota_r >= iota_c + # [num_v_heads, chunk, chunk] + gating_map_masked = jnp.where(strictly_lower_mask, gating_map, 0) + + # [num_v_heads, chunk, kq_head_dim] + k_beta_repeat = k_repeat * beta_large + + # [num_v_heads, chunk, chunk] + beta_k_k_t = jax.lax.dot( + k_beta_repeat, + k_repeat, + dimension_numbers=(((2,), (2,)), ((0,), (0,))), + preferred_element_type=jnp.float32, + ).astype(cfg.dtypes.compute) + gating_beta_k_k_t = gating_map_masked * beta_k_k_t + t = jnp.where(identity_mask, 1, gating_beta_k_k_t) + + # [num_v_heads, chunk, chunk] + t_inv = invert_triangular_matrix(t) + + # [num_v_heads, chunk, v_head_dim] + v_beta_large = v_large * beta_large + # [num_v_heads, chunk, kv_head_dim] + k_beta_gating = k_beta_repeat * gating_forward + # NOTE: If v_head_dim < mxu size, concatenating them will help increase mxu + # utilization. Also, if v_head_dim is multiple of lane size, concat / split + # along lane dim is free - making this optimization strictly beneficial. + # [num_v_heads, chunk, v_head_dim + kq_head_dim] + merged_v_k = jnp.concat([v_beta_large, k_beta_gating], axis=-1) + merged_uw = jax.lax.dot( + t_inv, + merged_v_k, + dimension_numbers=(((2,), (1,)), ((0,), (0,))), + preferred_element_type=jnp.float32, + ).astype(cfg.dtypes.compute) + + # [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,))), + 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) + ws = ws.astype(cfg.dtypes.compute) + + # [num_v_heads, chunk, v_head_dim] + u_ws = u - ws + + # [num_v_heads, chunk, kq_head_dim] + k_repeat_gating = k_repeat * gating_backward + + # [num_v_heads, kq_head_dim, v_head_dim] + state_new = jax.lax.dot( + k_repeat_gating, + u_ws, + dimension_numbers=(((1,), (1,)), ((0,), (0,))), + preferred_element_type=jnp.float32, + ) + + # [num_v_heads, kq_head_dim, v_head_dim] + state_updated = state_prev * gating_last + state = state_updated + state_new + + # [num_kq_heads, chunk, chunk] + out_qk = jax.lax.dot( + q_large, + k_large, + dimension_numbers=(((2,), (2,)), ((0,), (0,))), + preferred_element_type=jnp.float32, + ).astype(cfg.dtypes.compute) + # NOTE: must perform repeat after matmul to reduce required compute. + # [num_v_heads, chunk, chunk] + out_qk = jnp.repeat(out_qk, cfg.v_per_kq_head, axis=0) + out_qk *= gating_map + out_qk = jnp.where(lower_mask, out_qk, 0) + + # [num_v_heads, chunk, v_head_dim] + out_new = jax.lax.dot( + out_qk, + u_ws, + dimension_numbers=(((2,), (1,)), ((0,), (0,))), + preferred_element_type=jnp.float32, + ) + out = out_updated + out_new + + return out, state, t_inv + + +def chunked_gdn( + real_sizes: jax.Array, + q_large: jax.Array, + k_large: jax.Array, + v_large: jax.Array, + b_large: jax.Array, + a_large: jax.Array, + state_prev: jax.Array, + 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].""" + + 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) + mask = iota < real_sizes.reshape(-1, 1, 1, 1).astype(mask_dtype) + + # [seqs, num_kq_heads, chunk, kq_head_dim] + q_large = jnp.where(mask, q_large.astype(cfg.dtypes.compute), 0) + k_large = jnp.where(mask, k_large.astype(cfg.dtypes.compute), 0) + # [seqs, num_v_heads, chunk, v_head_dim] + v_large = jnp.where(mask, v_large.astype(cfg.dtypes.compute), 0) + + b_large = b_large.astype(cfg.dtypes.compute) + a_large = a_large.astype(cfg.dtypes.compute) + + a_log = a_log.reshape(1, 1, 1, -1).astype(cfg.dtypes.compute) + dt_bias = dt_bias.reshape(1, 1, 1, -1).astype(cfg.dtypes.compute) + + # NOTE: Any element-wise computations should occur before repeat. + q_large = l2_norm(q_large) + q_scale = cfg.kq_head_dim**-0.5 + q_large *= q_scale + k_large = l2_norm(k_large) + + # [seqs, 1, chunk, num_v_heads] + beta = jax.nn.sigmoid(b_large) + gating_log = -jnp.exp(a_log) * jax.nn.softplus(a_large + dt_bias) + + beta = jnp.where(mask, beta, 0) + # NOTE: Masked gating_log will evaluate to jnp.exp(0)=1. gating (decay) must + # be masked to 1 since it signifies that strength of state from previous row + # will be 1 (i.e., no decay) if current row is invalid. + gating_log = jnp.where(mask, gating_log, 0) + + out_list = [] + state_list = [] + t_inv_list = [] + for idx in range(cfg.seq_tile_size): + out, state, t_inv = chunked_gdn_per_seq( + q_large[idx], + k_large[idx], + v_large[idx], + gating_log[idx], + beta[idx], + state_prev[idx], + cfg, + ) + out_list.append(out.swapaxes(0, 1)) + state_list.append(state) + t_inv_list.append(t_inv) + out = jnp.stack(out_list, axis=0) + state = jnp.stack(state_list, axis=0) + t_inv = jnp.stack(t_inv_list, axis=0) + return out, state, t_inv + + +def recurrent_gdn_per_seq( + q_compact: jax.Array, # [num_kq_heads, chunk, 1, kq_head_dim] + k_compact: jax.Array, # [num_kq_heads, chunk, 1, kq_head_dim] + k_compact_t: jax.Array, # [num_kq_heads, chunk, kq_head_dim, 1] + v_compact: jax.Array, # [num_v_heads, chunk, 1, v_head_dim] + gating_log: jax.Array, # [num_v_heads, chunk, 1, 1] + beta: jax.Array, # [num_v_heads, chunk, 1, 1] + state: jax.Array, # [num_v_heads, kq_head_dim, v_head_dim] + cfgs: config.GDNConfig, +) -> tuple[jax.Array, jax.Array]: + """Perform recurrent GDN over input [num_heads, chunk, 1, head_dim].""" + + out_list = [] + for c_idx in range(cfgs.chunk_size): + # [num_v_heads, 1, kq_head_dim] + q_curr = q_compact[:, c_idx] + q_curr = jnp.repeat(q_curr, cfgs.v_per_kq_head, axis=0) + k_curr = k_compact[:, c_idx] + k_curr = jnp.repeat(k_curr, cfgs.v_per_kq_head, axis=0) + + # [num_v_heads, 1, v_head_dim] + v_curr = v_compact[:, c_idx] + + # [num_v_heads, kq_head_dim, 1] + k_curr_t = k_compact_t[:, c_idx] + k_curr_t = jnp.repeat(k_curr_t, cfgs.v_per_kq_head, axis=0) + + # [num_v_heads, 1, 1] + beta_curr = beta[:, c_idx] + gating_curr = gating_log[:, c_idx] + + # [num_v_heads, kq_head_dim, v_head_dim] + state_updated = state * gating_curr + + # [num_v_heads, 1, v_head_dim] + v_updated = jax.lax.dot( + k_curr, + state_updated, + dimension_numbers=(((2,), (1,)), ((0,), (0,))), + preferred_element_type=jnp.float32, + ).astype(cfgs.dtypes.compute) + + # [num_v_heads, 1, v_head_dim] + v_diff = v_curr - v_updated + v_new = beta_curr * v_diff + + # [num_v_heads, kq_head_dim, v_head_dim] + # NOTE: Multiplication with k_curr_t needs to be deferred as much as + # possible as it expands the dimension size by kq_head_dim. + state_new = k_curr_t * v_new + # [num_v_heads, kq_head_dim, v_head_dim] + state = state_updated + state_new + + # [num_v_heads, 1, v_head_dim] + out = jax.lax.dot( + q_curr, + state, + dimension_numbers=(((2,), (1,)), ((0,), (0,))), + preferred_element_type=jnp.float32, + ).astype(cfgs.dtypes.compute) + + out_list.append(out[:, 0, :]) + + return jnp.stack(out_list, axis=0), state + + +def recurrent_gdn( + real_sizes: jax.Array, + q_compact: jax.Array, + k_compact: jax.Array, + v_compact: jax.Array, + b_compact: jax.Array, + a_compact: jax.Array, + state_prev: jax.Array, + a_log: jax.Array, + dt_bias: jax.Array, + cfg: config.GDNConfig, +) -> tuple[jax.Array, jax.Array, jax.Array]: + """Perform recurrent GDN over input [seq, num_heads, chunk, 1, head_dim].""" + + mask_dtype = get_mask_dtype(cfg.dtypes.compute) + iota = jax.lax.broadcasted_iota(mask_dtype, (cfg.seq_tile_size, 1, cfg.chunk_size, 1, 1), 2) + mask = iota < real_sizes.reshape(-1, 1, 1, 1, 1).astype(mask_dtype) + + # [seqs, num_kq_heads, chunk, 1, kq_head_dim] + q_compact = jnp.where(mask, q_compact.astype(cfg.dtypes.compute), 0) + k_compact = jnp.where(mask, k_compact.astype(cfg.dtypes.compute), 0) + # [seqs, num_v_heads, chunk, 1, v_head_dim] + v_compact = jnp.where(mask, v_compact.astype(cfg.dtypes.compute), 0) + + b_compact = b_compact.astype(cfg.dtypes.compute) + a_compact = a_compact.astype(cfg.dtypes.compute) + + a_log = a_log.reshape(1, 1, 1, 1, -1).astype(cfg.dtypes.compute) + dt_bias = dt_bias.reshape(1, 1, 1, 1, -1).astype(cfg.dtypes.compute) + + # [seqs, num_kq_heads, chunk, 1, kq_head_dim] + q_compact = l2_norm(q_compact) + q_scale = cfg.kq_head_dim**-0.5 + q_compact *= q_scale + k_compact = l2_norm(k_compact) + k_compact_t = fused_transpose_broadcast(k_compact, src_dim=4, dst_dim=3) + + beta = jax.nn.sigmoid(b_compact) + gating_log = -jnp.exp(a_log) * jax.nn.softplus(a_compact + dt_bias) + + beta = jnp.where(mask, beta, 0) + # NOTE: Masked gating_log will evaluate to jnp.exp(0)=1. gating (decay) must + # be masked to 1 since it signifies that strength of state from previous row + # will be 1 (i.e., no decay) if current row is invalid. + gating_log = jnp.where(mask, gating_log, 0) + gating_log = jnp.exp(gating_log) + + beta = fused_transpose_broadcast(beta, src_dim=4, dst_dim=1) + beta = beta[:, : cfg.num_v_heads] + gating_log = fused_transpose_broadcast(gating_log, src_dim=4, dst_dim=1) + gating_log = gating_log[:, : cfg.num_v_heads] + + out_list = [] + new_state_list = [] + + for idx in range(cfg.seq_tile_size): + out, state = recurrent_gdn_per_seq( + q_compact[idx], + k_compact[idx], + k_compact_t[idx], + v_compact[idx], + gating_log[idx], + beta[idx], + state_prev[idx], + cfg, + ) + out_list.append(out) + new_state_list.append(state) + + out = jnp.stack(out_list, axis=0) + new_recurrent_state = jnp.stack(new_state_list, axis=0) + t_inv = jnp.ones((cfg.seq_tile_size, cfg.num_v_heads, 1, 1), dtype=cfg.dtypes.compute) + + return out, new_recurrent_state, t_inv diff --git a/src/maxtext/models/kernels/gdn/config.py b/src/maxtext/models/kernels/gdn/config.py new file mode 100644 index 000000000..1332ec4d4 --- /dev/null +++ b/src/maxtext/models/kernels/gdn/config.py @@ -0,0 +1,154 @@ +# 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. +# ============================================================================== + +"""GDN Configuration dataclass defining GDN tiling, dtypes, and kernel configurations.""" + +import dataclasses +import enum +from typing import Any + +import jax +from jax.experimental import pallas as pl +from jax.experimental.pallas import tpu as pltpu +import jax.numpy as jnp + + +DEFAULT_VMEM_LIMIT_FACTOR: float = 0.90 + + +class GDNMode(enum.StrEnum): + """Execution mode for GDN kernel.""" + + BATCHED = enum.auto() + PER_SEQ = enum.auto() + + def get_seq_tile_size(self, tile_size: int) -> int: + if self == GDNMode.BATCHED: + return tile_size + return 1 + + def get_chunk_size(self, tile_size: int) -> int: + if self == GDNMode.BATCHED: + return 1 + return tile_size + + +@jax.tree_util.register_dataclass +@dataclasses.dataclass(frozen=True) +class Dtypes: + """Dtypes used across GDN operations.""" + + act_in: jnp.dtype + act_out: jnp.dtype + compute: jnp.dtype + recurrent_state: jnp.dtype + conv_state: jnp.dtype + + +def get_vmem_limit_bytes( + vmem_limit_factor: float = DEFAULT_VMEM_LIMIT_FACTOR, +) -> int: + """Returns the maximum allowable VMEM capacity budget in bytes.""" + tpu_info = pltpu.get_tpu_info() + return int(vmem_limit_factor * tpu_info.vmem_capacity_bytes) + + +@jax.tree_util.register_dataclass +@dataclasses.dataclass(frozen=True) +class GDNConfig: + """Configuration dataclass for GDN kernel.""" + + mode: GDNMode + dtypes: Dtypes + batch_size: int + dim_size: int + kernel_size: int + tile_size: int + num_kq_heads: int + num_v_heads: int + kq_head_dim: int + v_head_dim: int + num_buffers: int = 2 + + @property + def chunk_size(self) -> int: + return self.mode.get_chunk_size(self.tile_size) + + @property + def seq_tile_size(self) -> int: + return self.mode.get_seq_tile_size(self.tile_size) + + @property + def prev_kernel_size(self) -> int: + return self.kernel_size - 1 + + @property + def v_dim_size(self) -> int: + return self.num_v_heads * self.v_head_dim + + @property + def kq_dim_size(self) -> int: + return self.num_kq_heads * self.kq_head_dim + + @property + def v_per_kq_head(self) -> int: + return self.num_v_heads // self.num_kq_heads + + @property + def aligned_num_v_heads(self) -> int: + tpu_info = pltpu.get_tpu_info() + num_lanes = tpu_info.num_lanes + return pl.cdiv(self.num_v_heads, num_lanes) * num_lanes + + def get_kernel_name(self) -> str: + return f"fused_conv1d_gdn_{self.mode.value}_b{self.seq_tile_size}" f"_c{self.chunk_size}" + + def get_metadata(self) -> dict[str, str | int | float]: + cfgs_dict = dataclasses.asdict(self) + ret = {} + for path, val in jax.tree_util.tree_leaves_with_path(cfgs_dict): + key = jax.tree_util.keystr(path, simple=True, separator=".") + if not isinstance(val, str | int | float): + val = str(val) + ret[key] = val + return ret + + def get_out_shape(self) -> jax.ShapeDtypeStruct: + return jax.ShapeDtypeStruct( + (self.batch_size, self.num_v_heads, self.v_head_dim), + self.dtypes.act_out, + ) + + def get_scratch_shape_dict(self) -> dict[str, Any]: + """Returns dictionary of scratch shapes for Pallas pipeline.""" + conv_shape = (self.seq_tile_size, self.prev_kernel_size, 1, self.dim_size) + recurrent_shape = ( + self.seq_tile_size, + self.num_v_heads, + self.kq_head_dim, + self.v_head_dim, + ) + + carry_conv_scratch = carry_recurrent_scratch = None + # NOTE: Currently, batched mode only supports case where 1 seq = 1 tile. + # Therefore, inter tile carry is not needed. + if self.mode != GDNMode.BATCHED: + carry_conv_scratch = pltpu.VMEM(conv_shape, jnp.float32) + carry_recurrent_scratch = pltpu.VMEM(recurrent_shape, jnp.float32) + + return { + "carry_conv_scratch_ref": carry_conv_scratch, + "carry_recurrent_scratch_ref": carry_recurrent_scratch, + } diff --git a/src/maxtext/models/kernels/gdn/memory_ref.py b/src/maxtext/models/kernels/gdn/memory_ref.py new file mode 100644 index 000000000..eb6db2f28 --- /dev/null +++ b/src/maxtext/models/kernels/gdn/memory_ref.py @@ -0,0 +1,616 @@ +# 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. +# ============================================================================== + +"""Weight and state reference dataclasses for VMEM.""" + +import dataclasses +import functools +from typing import Any + +import jax +from jax.experimental import pallas as pl +from jax.experimental.pallas import tpu as pltpu +import jax.numpy as jnp + +try: + from maxtext.models.kernels.gdn import config +except (ImportError, ModuleNotFoundError): + try: + from maxtext.src.maxtext.models.kernels.gdn import config + except (ImportError, ModuleNotFoundError): + from . import config + + +def _flat_pos(shape: tuple[int, ...], indices: tuple[Any, ...]) -> Any: + """Row-major flat offset of `indices` into a logical array of `shape`.""" + strides = pl.strides_from_shape(shape) + assert len(strides) == len(indices) + + pos = 0 + for stride, idx in zip(strides, indices): + pos += stride * idx + return pos + + +@jax.tree_util.register_dataclass +@dataclasses.dataclass(frozen=True) +class ConvWeightsRef: + weight: Any + bias: Any | None = None + + +@jax.tree_util.register_dataclass +@dataclasses.dataclass(frozen=True) +class GDNWeightsRef: + a_log: Any + dt_bias: Any + + +@jax.tree_util.register_dataclass +@dataclasses.dataclass(frozen=True) +class WeightRefs: + conv: ConvWeightsRef + gdn: GDNWeightsRef + + +class FieldOffset: + """Descriptor returning the record field at ``data[pos + offset]``. + + Reads a single dynamically-indexed element rather than a slice, since JAX + can't slice a range with traced indices. Read-only: metadata is never + written. + """ + + def __init__(self, offset: int): + self.offset = offset + + def __get__(self, obj, objtype=None): + if obj is None: + return self + return obj.data[obj.pos + self.offset] + + +# Per-p_id metadata is an array of structs: each p_id's fields sit contiguously +# and FieldOffset(k) reads the k-th word of its struct. +# +# Packed struct: [r_base, packed_word]. +# Fields share packed_word to save SMEM: is_first_tile(0), is_last_tile(1), r_size(2..15), s_idx(16..31). +@dataclasses.dataclass(frozen=True) +class PackedPIdRecord: + """Packed struct [r_base, packed_word]; the four small fields bit-slice word. + + Each bit field masks after shifting, which also clears the sign bits that + ``>>`` extends on the signed int32 word. + """ + + STRUCT_SIZE = 2 + FIRST_TILE_SHIFT = 0 + LAST_TILE_SHIFT = 1 + R_SIZE_SHIFT = 2 + S_IDX_SHIFT = 16 + FLAG_MASK = 1 + R_SIZE_MASK = (1 << (S_IDX_SHIFT - R_SIZE_SHIFT)) - 1 + S_IDX_MASK = (1 << (32 - S_IDX_SHIFT)) - 1 + MAX_SEQS = S_IDX_MASK + 1 + + data: Any + pos: Any + r_base = FieldOffset(0) + word = FieldOffset(1) + + @property + def s_idx(self): + return (self.word >> self.S_IDX_SHIFT) & self.S_IDX_MASK + + @property + def r_size(self): + return (self.word >> self.R_SIZE_SHIFT) & self.R_SIZE_MASK + + @property + def is_first_tile(self): + return (self.word & self.FLAG_MASK) != 0 + + @property + def is_last_tile(self): + return ((self.word >> self.LAST_TILE_SHIFT) & self.FLAG_MASK) != 0 + + @classmethod + def pack( + cls, + s_idx: jax.Array, + r_size: jax.Array, + is_first_tile: jax.Array, + is_last_tile: jax.Array, + ) -> jax.Array: + """Packs s_idx, row size and two tile-state flags into one int32 word.""" + + s_idx = s_idx.reshape(-1).astype(jnp.int32) + r_size = r_size.reshape(-1).astype(jnp.int32) + is_first_tile = is_first_tile.reshape(-1).astype(jnp.int32) + is_last_tile = is_last_tile.reshape(-1).astype(jnp.int32) + word = s_idx << cls.S_IDX_SHIFT + word |= r_size << cls.R_SIZE_SHIFT + word |= is_last_tile << cls.LAST_TILE_SHIFT + word |= is_first_tile << cls.FIRST_TILE_SHIFT + return word + + +@jax.tree_util.register_dataclass +@dataclasses.dataclass(frozen=True) +class MetadataRef: + """Container for sequence and chunk metadata references.""" + + num_tiles: Any + # Array of structs holding every p_id's metadata + records: Any + s_idx_has_initial_state: Any + s_idx_to_state_indices: Any + shape: tuple[int, ...] = dataclasses.field(metadata={"static": True}) + + def get_record(self, p_id, idx) -> PackedPIdRecord: + """View of one p_id's metadata: .r_base / .s_idx / .r_size / .is_*_tile.""" + record_idx = _flat_pos(self.shape, (p_id, idx)) + return PackedPIdRecord(self.records, record_idx * PackedPIdRecord.STRUCT_SIZE) + + @classmethod + def create( # pyrefly: ignore[bad-override] + cls, + cfgs: config.GDNConfig, + num_tiles: jax.Array, + p_id_to_s_idx: jax.Array, + p_id_to_r_base: jax.Array, + p_id_to_r_size: jax.Array, + p_id_is_first_tile: jax.Array, + p_id_is_last_tile: jax.Array, + s_idx_has_initial_state: jax.Array, + s_idx_to_state_indices: jax.Array, + ): + """Constructs a MetadataRef instance from packing inputs.""" + # NOTE: First dim does not matter when it comes to calculating stride. + shape = (1, cfgs.seq_tile_size) + assert s_idx_has_initial_state.shape[0] <= PackedPIdRecord.MAX_SEQS, ( + f"Number of sequences ({s_idx_has_initial_state.shape[0]}) exceeds" + f" PackedPIdRecord limit ({PackedPIdRecord.MAX_SEQS})." + ) + assert cfgs.tile_size <= PackedPIdRecord.R_SIZE_MASK, ( + f"Tile size ({cfgs.tile_size}) exceeds PackedPIdRecord limit" f" ({PackedPIdRecord.R_SIZE_MASK})." + ) + + r_base = p_id_to_r_base.reshape(-1).astype(jnp.int32) + word = PackedPIdRecord.pack(p_id_to_s_idx, p_id_to_r_size, p_id_is_first_tile, p_id_is_last_tile) + fields = [r_base, word] + # Interleave fields into one array of structs: [rec0_f0, rec0_f1, ...]. + records = jnp.stack(fields, axis=-1).reshape(-1) + + return cls( + num_tiles=num_tiles, + records=records, + s_idx_has_initial_state=s_idx_has_initial_state, + s_idx_to_state_indices=s_idx_to_state_indices, + shape=shape, + ) + + def __len__(self) -> int: + return len(jax.tree_util.tree_leaves(self)) + + +@dataclasses.dataclass(frozen=True, kw_only=True) +class BaseBufferedRef(pltpu.BufferedRef): + """Base class for Pallas BufferedRef with attached GDN configuration.""" + + cfg: config.GDNConfig = dataclasses.field(metadata={"static": True}) + # NOTE: Despite being ref, metadata_ref should be set to static. This is + # because the memory will be allocated outside of kernel and metadata_ref + # merely points to the reference. + metadata_ref: MetadataRef = dataclasses.field(metadata={"static": True}) + + @classmethod + def create( # pyrefly: ignore[bad-override] + cls, + spec: pl.BlockSpec, + dtype_or_type: jax.Array, + buffer_type: pltpu.BufferType, + buffer_count: int, + use_lookahead: bool, + cfg: config.GDNConfig, + metadata_ref: MetadataRef, + ): + """Creates a BaseBufferedRef wrapping a standard Pallas BufferedRef.""" + standard_ref = pltpu.BufferedRef.create( + spec=spec, + dtype_or_type=dtype_or_type, + buffer_type=buffer_type, + buffer_count=buffer_count, + grid_rank=1, + use_lookahead=use_lookahead, + ) + return cls( + cfg=cfg, + metadata_ref=metadata_ref, + **{f.name: getattr(standard_ref, f.name) for f in dataclasses.fields(pltpu.BufferedRef)}, + ) + + +@jax.tree_util.register_dataclass +@dataclasses.dataclass(frozen=True, kw_only=True) +class InBufferedRef(BaseBufferedRef): + """Input double-buffered DMA reference for Pallas pipeline.""" + + def copy_in(self, src_ref: jax.Array, grid_indices: tuple[int | jax.Array]): + """Initiates asynchronous copy into current input VMEM slot.""" + assert self.sem_recvs is not None + assert self.window_ref is not None + slot = self.current_copy_in_slot + sem = self.sem_recvs.at[slot] + vmem_ref = self.window_ref.at[slot] + p_id = grid_indices[0] + + for idx in range(self.cfg.seq_tile_size): + record = self.metadata_ref.get_record(p_id, idx) + r_base = record.r_base + dma_size = record.r_size + pltpu.make_async_copy( + src_ref.at[pl.ds(r_base, dma_size)], + vmem_ref.at[idx, pl.ds(0, dma_size)], # pyrefly: ignore[missing-attribute] + sem, + ).start() + + def wait_in(self, src_ref: jax.Array, grid_indices: tuple[int | jax.Array]): + """Waits for asynchronous input copy to complete.""" + assert self.sem_recvs is not None + assert self.window_ref is not None + slot = self.current_wait_in_slot + sem = self.sem_recvs.at[slot] + vmem_ref = self.window_ref.at[slot] + p_id = grid_indices[0] + + dma_size = 0 + for idx in range(self.cfg.seq_tile_size): + dma_size += self.metadata_ref.get_record(p_id, idx).r_size + + pltpu.make_async_copy( + vmem_ref.at[0, pl.ds(0, dma_size)], # pyrefly: ignore[missing-attribute] + vmem_ref.at[0, pl.ds(0, dma_size)], # pyrefly: ignore[missing-attribute] + sem, + ).wait() + + +@jax.tree_util.register_dataclass +@dataclasses.dataclass(frozen=True, kw_only=True) +class OutBufferedRef(BaseBufferedRef): + """Output double-buffered DMA reference for Pallas pipeline.""" + + def copy_out(self, dst_ref: jax.Array, grid_indices: tuple[int | jax.Array]): + """Initiates asynchronous copy from current output VMEM slot.""" + assert self.sem_sends is not None + assert self.window_ref is not None + slot = self.current_copy_out_slot + sem = self.sem_sends.at[slot] + vmem_ref = self.window_ref.at[slot] + p_id = grid_indices[0] + + for idx in range(self.cfg.seq_tile_size): + record = self.metadata_ref.get_record(p_id, idx) + r_base = record.r_base + dma_size = record.r_size + pltpu.make_async_copy( + vmem_ref.at[idx, pl.ds(0, dma_size)], # pyrefly: ignore[missing-attribute] + dst_ref.at[pl.ds(r_base, dma_size)], + sem, + ).start() + + def wait_out(self, dst_ref: jax.Array, grid_indices: tuple[int | jax.Array]): + """Waits for asynchronous output copy to complete.""" + assert self.sem_sends is not None + assert self.window_ref is not None + slot = self.current_wait_out_slot + sem = self.sem_sends.at[slot] + vmem_ref = self.window_ref.at[slot] + p_id = grid_indices[0] + + dma_size = 0 + for idx in range(self.cfg.seq_tile_size): + dma_size += self.metadata_ref.get_record(p_id, idx).r_size + + pltpu.make_async_copy( + vmem_ref.at[0, pl.ds(0, dma_size)], # pyrefly: ignore[missing-attribute] + vmem_ref.at[0, pl.ds(0, dma_size)], # pyrefly: ignore[missing-attribute] + sem, + ).wait() + + +@jax.tree_util.register_dataclass +@dataclasses.dataclass(frozen=True, kw_only=True) +class TInvBufferedRef(BaseBufferedRef): + """DMA buffer for triangular inverse matrix caching (t_inv).""" + + def copy_out(self, dst_ref: jax.Array, grid_indices: tuple[int | jax.Array]): + assert self.sem_sends is not None + assert self.window_ref is not None + slot = self.current_copy_out_slot + sem = self.sem_sends.at[slot] + vmem_ref = self.window_ref.at[slot] + p_id = grid_indices[0] + + for idx in range(self.cfg.seq_tile_size): + pltpu.make_async_copy( + vmem_ref.at[idx], + dst_ref.at[p_id + idx], + sem, + ).start() + + def wait_out(self, dst_ref: jax.Array, grid_indices: tuple[int | jax.Array]): + assert self.sem_sends is not None + assert self.window_ref is not None + slot = self.current_wait_out_slot + sem = self.sem_sends.at[slot] + vmem_ref = self.window_ref.at[slot] + + for idx in range(self.cfg.seq_tile_size): + pltpu.make_async_copy( + vmem_ref.at[idx], + vmem_ref.at[idx], + sem, + ).wait() + + +@jax.tree_util.register_dataclass +@dataclasses.dataclass(frozen=True, kw_only=True) +class ChunkStatesBufferedRef(BaseBufferedRef): + """DMA buffer for caching intermediate recurrent chunk states.""" + + def copy_out(self, dst_ref: jax.Array, grid_indices: tuple[int | jax.Array]): + assert self.sem_sends is not None + assert self.window_ref is not None + slot = self.current_copy_out_slot + sem = self.sem_sends.at[slot] + vmem_ref = self.window_ref.at[slot] + p_id = grid_indices[0] + + for idx in range(self.cfg.seq_tile_size): + pltpu.make_async_copy( + vmem_ref.at[idx], + dst_ref.at[p_id + idx], + sem, + ).start() + + def wait_out(self, dst_ref: jax.Array, grid_indices: tuple[int | jax.Array]): + assert self.sem_sends is not None + assert self.window_ref is not None + slot = self.current_wait_out_slot + sem = self.sem_sends.at[slot] + vmem_ref = self.window_ref.at[slot] + + for idx in range(self.cfg.seq_tile_size): + pltpu.make_async_copy( + vmem_ref.at[idx], + vmem_ref.at[idx], + sem, + ).wait() + + +@jax.tree_util.register_dataclass +@dataclasses.dataclass(frozen=True, kw_only=True) +class StateBufferedRef(BaseBufferedRef): + """State buffered reference with conditional first/last tile DMA.""" + + def copy_in(self, src_ref: jax.Array, grid_indices: tuple[int | jax.Array]): + """Initiates asynchronous state copy into VMEM slot.""" + assert self.sem_recvs is not None + assert self.window_ref is not None + slot = self.current_copy_in_slot + sem = self.sem_recvs.at[slot] + vmem_ref = self.window_ref.at[slot] + p_id = grid_indices[0] + + for idx in range(self.cfg.seq_tile_size): + record = self.metadata_ref.get_record(p_id, idx) + is_first_tile = record.is_first_tile + s_idx = record.s_idx + state_idx = self.metadata_ref.s_idx_to_state_indices[s_idx] + has_initial_state = self.metadata_ref.s_idx_has_initial_state[s_idx] + should_read = jnp.logical_and(is_first_tile, has_initial_state) + dma_size = jnp.where(should_read, 1, 0) + + pltpu.make_async_copy( + src_ref.at[pl.ds(state_idx, dma_size)], + vmem_ref.at[pl.ds(idx, dma_size)], # pyrefly: ignore[missing-attribute] + sem, + ).start() + + def wait_in(self, src_ref: jax.Array, grid_indices: tuple[int | jax.Array]): + """Waits for asynchronous state input copy to complete.""" + assert self.sem_recvs is not None + assert self.window_ref is not None + slot = self.current_wait_in_slot + sem = self.sem_recvs.at[slot] + vmem_ref = self.window_ref.at[slot] + p_id = grid_indices[0] + + dma_size = 0 + for idx in range(self.cfg.seq_tile_size): + record = self.metadata_ref.get_record(p_id, idx) + is_first_tile = record.is_first_tile + s_idx = record.s_idx + has_initial_state = self.metadata_ref.s_idx_has_initial_state[s_idx] + should_read = jnp.logical_and(is_first_tile, has_initial_state) + dma_size += jnp.where(should_read, 1, 0) + + pltpu.make_async_copy( + vmem_ref.at[pl.ds(0, dma_size)], # pyrefly: ignore[missing-attribute] + vmem_ref.at[pl.ds(0, dma_size)], # pyrefly: ignore[missing-attribute] + sem, + ).wait() + + def copy_out(self, dst_ref: jax.Array, grid_indices: tuple[int | jax.Array]): + """Initiates asynchronous state copy from VMEM slot to HBM.""" + assert self.sem_sends is not None + assert self.window_ref is not None + slot = self.current_copy_out_slot + sem = self.sem_sends.at[slot] + vmem_ref = self.window_ref.at[slot] + p_id = grid_indices[0] + + for idx in range(self.cfg.seq_tile_size): + record = self.metadata_ref.get_record(p_id, idx) + is_last_tile = record.is_last_tile + s_idx = record.s_idx + state_idx = self.metadata_ref.s_idx_to_state_indices[s_idx] + dma_size = jnp.where(is_last_tile, 1, 0) + + pltpu.make_async_copy( + vmem_ref.at[pl.ds(idx, dma_size)], # pyrefly: ignore[missing-attribute] + dst_ref.at[pl.ds(state_idx, dma_size)], + sem, + ).start() + + def wait_out(self, dst_ref: jax.Array, grid_indices: tuple[int | jax.Array]): + """Waits for asynchronous state output copy to complete.""" + assert self.sem_sends is not None + assert self.window_ref is not None + slot = self.current_wait_out_slot + sem = self.sem_sends.at[slot] + vmem_ref = self.window_ref.at[slot] + p_id = grid_indices[0] + + dma_size = 0 + for idx in range(self.cfg.seq_tile_size): + is_last_tile = self.metadata_ref.get_record(p_id, idx).is_last_tile + dma_size += jnp.where(is_last_tile, 1, 0) + + pltpu.make_async_copy( + vmem_ref.at[pl.ds(0, dma_size)], # pyrefly: ignore[missing-attribute] + vmem_ref.at[pl.ds(0, dma_size)], # pyrefly: ignore[missing-attribute] + sem, + ).wait() + + +def create_allocs( + metadata_ref: MetadataRef, + qkv_ref: jax.Array, + b_ref: jax.Array, + a_ref: jax.Array, + out_ref: jax.Array, + conv_state_ref: jax.Array, + recurrent_state_ref: jax.Array, + cfg: config.GDNConfig, + t_inv_ref: jax.Array | None = None, + chunk_states_ref: jax.Array | None = None, +) -> tuple[Any, ...]: + """Creates all Pallas double-buffered allocations for GDN kernel.""" + qkv_shape = (cfg.seq_tile_size, cfg.chunk_size, 1, cfg.dim_size) + ba_shape = (cfg.seq_tile_size, cfg.chunk_size, 1, cfg.aligned_num_v_heads) + + out_shape = ( + cfg.seq_tile_size, + cfg.chunk_size, + cfg.num_v_heads, + cfg.v_head_dim, + ) + conv_shape = (cfg.seq_tile_size, cfg.prev_kernel_size, 1, cfg.dim_size) + recurrent_shape = ( + cfg.seq_tile_size, + cfg.num_v_heads, + cfg.kq_head_dim, + cfg.v_head_dim, + ) + + pipeline_mode = pl.Buffered(buffer_count=cfg.num_buffers, use_lookahead=False) + + block_spec_partial = functools.partial( + pl.BlockSpec, + memory_space=pltpu.VMEM, + index_map=lambda i: (i,), + pipeline_mode=pipeline_mode, + ) + + qkv_spec = block_spec_partial(block_shape=qkv_shape) + ba_spec = block_spec_partial(block_shape=ba_shape) + in_buffered_partial = functools.partial( + InBufferedRef.input, + buffer_count=pipeline_mode.buffer_count, + use_lookahead=pipeline_mode.use_lookahead, + cfg=cfg, + metadata_ref=metadata_ref, + ) + qkv_alloc = in_buffered_partial(spec=qkv_spec, dtype_or_type=qkv_ref) + b_alloc = in_buffered_partial(spec=ba_spec, dtype_or_type=b_ref) + a_alloc = in_buffered_partial(spec=ba_spec, dtype_or_type=a_ref) + + out_alloc = OutBufferedRef.output( + spec=block_spec_partial(block_shape=out_shape), + dtype_or_type=out_ref, + buffer_count=pipeline_mode.buffer_count, + use_lookahead=pipeline_mode.use_lookahead, + cfg=cfg, + metadata_ref=metadata_ref, + ) + + conv_spec = block_spec_partial(block_shape=conv_shape) + recurrent_spec = block_spec_partial(block_shape=recurrent_shape) + state_buffered_partial = functools.partial( + StateBufferedRef.input_output, + buffer_count=pipeline_mode.buffer_count, + use_lookahead=pipeline_mode.use_lookahead, + cfg=cfg, + metadata_ref=metadata_ref, + ) + conv_alloc = state_buffered_partial(spec=conv_spec, dtype_or_type=conv_state_ref) + recurrent_alloc = state_buffered_partial(spec=recurrent_spec, dtype_or_type=recurrent_state_ref) + + allocs = [ + qkv_alloc, + b_alloc, + a_alloc, + conv_alloc, + recurrent_alloc, + out_alloc, + ] + + if t_inv_ref is not None: + t_inv_shape = ( + cfg.seq_tile_size, + cfg.num_v_heads, + cfg.chunk_size, + cfg.chunk_size, + ) + t_inv_alloc = TInvBufferedRef.output( + spec=block_spec_partial(block_shape=t_inv_shape), + dtype_or_type=t_inv_ref, + buffer_count=pipeline_mode.buffer_count, + use_lookahead=pipeline_mode.use_lookahead, + cfg=cfg, + metadata_ref=metadata_ref, + ) + allocs.append(t_inv_alloc) + + if chunk_states_ref is not None: + chunk_states_shape = ( + cfg.seq_tile_size, + cfg.num_v_heads, + cfg.kq_head_dim, + cfg.v_head_dim, + ) + chunk_states_alloc = ChunkStatesBufferedRef.output( + spec=block_spec_partial(block_shape=chunk_states_shape), + dtype_or_type=chunk_states_ref, + buffer_count=pipeline_mode.buffer_count, + use_lookahead=pipeline_mode.use_lookahead, + cfg=cfg, + metadata_ref=metadata_ref, + ) + allocs.append(chunk_states_alloc) + + return tuple(allocs) diff --git a/src/maxtext/models/kernels/gdn/metadata.py b/src/maxtext/models/kernels/gdn/metadata.py new file mode 100644 index 000000000..bb9cb7af8 --- /dev/null +++ b/src/maxtext/models/kernels/gdn/metadata.py @@ -0,0 +1,144 @@ +# 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. +# ============================================================================== + +"""Metadata references for sequence mapping and grid distribution.""" + +import jax +from jax.experimental import pallas as pl +import jax.numpy as jnp + +try: + from maxtext.models.kernels.gdn import config + from maxtext.models.kernels.gdn import memory_ref +except (ImportError, ModuleNotFoundError): + try: + from maxtext.src.maxtext.models.kernels.gdn import config + from maxtext.src.maxtext.models.kernels.gdn import memory_ref + except (ImportError, ModuleNotFoundError): + from . import config + from . import memory_ref + + +def compute_batched_seq_metadata( + cfg: config.GDNConfig, + seq_lens: jax.Array, + query_start_loc: jax.Array, + state_indices: jax.Array, + end_seq: jax.Array, +) -> memory_ref.MetadataRef: + """Metadata for computing multiple sequences per tile.""" + + max_seqs = seq_lens.size + all_seqs = jnp.arange(max_seqs) + + # NOTE: Only supports use case where query_lens[i] = 1 where i < end_seq. + # This must be guaranteed by the function caller. + # TODO(b/534541682): Add error handling when above condition is not met. + query_lens = query_start_loc[1:] - query_start_loc[:-1] + is_valid_seqs = jnp.where(all_seqs < end_seq, True, False) + has_initial_state = (seq_lens - query_lens) > 0 + all_valid_seqs = jnp.where(is_valid_seqs, all_seqs, 0) + + return memory_ref.MetadataRef.create( + cfgs=cfg, + num_tiles=pl.cdiv(end_seq, cfg.tile_size), + p_id_to_s_idx=all_valid_seqs, + p_id_to_r_base=all_valid_seqs, + p_id_to_r_size=jnp.where(is_valid_seqs, 1, 0), + p_id_is_first_tile=is_valid_seqs, + p_id_is_last_tile=is_valid_seqs, + s_idx_has_initial_state=has_initial_state, + s_idx_to_state_indices=state_indices, + ) + + +def compute_per_seq_metadata( + cfg: config.GDNConfig, + seq_lens: jax.Array, + query_start_loc: jax.Array, + state_indices: jax.Array, + start_seq: jax.Array, + end_seq: jax.Array, +) -> memory_ref.MetadataRef: + """Metadata for computing single sequence per tile.""" + + max_seqs = seq_lens.size + max_tokens = cfg.batch_size + all_seqs = jnp.arange(max_seqs) + all_tokens = jnp.arange(max_tokens) + + # Shift to ensure first element is for start_seq. + query_start_loc = jnp.roll(query_start_loc, shift=-start_seq) + seq_lens = jnp.roll(seq_lens, shift=-start_seq) + state_indices = jnp.roll(state_indices, shift=-start_seq) + + query_lens = query_start_loc[1:] - query_start_loc[:-1] + # NOTE: query_lens is used for calculating num_tiles. Defensive programming + # that masks out all the other values (seq_lens, state_indices) are not needed + # since they will not be visited as long as num_tiles is correct. + num_seqs = end_seq - start_seq + query_lens = jnp.where(all_seqs < num_seqs, query_lens, 0) + + # Calculate number of tiles needed for each sequence. + s_idx_to_num_tiles = pl.cdiv(query_lens, cfg.chunk_size) + # Calculate starting p_id of each sequence. + s_idx_to_start_p_id = jnp.cumulative_sum(s_idx_to_num_tiles, include_initial=True) + # Map tile index to seq index. + # Consider following case: + # all_seqs = [0 1 2 3 4] + # s_idx_to_num_tiles = [1 2 3 0 1] + # jnp.repeat will return following results: + # p_id_to_s_idx = [0 1 1 2 2 2 4] + # This means p_id_to_s_idx[i] will point to its corresponding seq index. + + # NOTE: To make jnp.repeat jit compilable, we add total_repeat_length. This + # introduces padding to p_id_to_s_idx[i] where i >= num_tiles. Since the + # kernel only checks value up-to p_id_to_s_idx[num_tiles-1], padded value + # will not impact kernel execution. + p_id_to_s_idx = jnp.repeat(all_seqs, s_idx_to_num_tiles, total_repeat_length=max_tokens) + # Map program id (p_id) to tile id of a sequence. + p_id_to_t_id = all_tokens - s_idx_to_start_p_id[p_id_to_s_idx] + # Map tile index to starting row of its activation. + p_id_to_r_base = query_start_loc[p_id_to_s_idx] + p_id_to_t_id * cfg.chunk_size + # Calculate number of rows to calculate / fetch for each tile. + p_id_to_r_size = jnp.minimum( + query_start_loc[p_id_to_s_idx + 1] - p_id_to_r_base, + cfg.tile_size, + ) + + # Calculate predicate used for state DMA. State is read if program id (p_id) + # is the first tile of a sequence and the sequence had been computed before + # (chunked prefill, decode, etc). State is written if the program id is the + # last tile of a sequence. + has_initial_state = (seq_lens - query_lens) > 0 + p_id_is_first_tile = p_id_to_t_id == 0 + p_id_is_last_tile = p_id_to_t_id == (s_idx_to_num_tiles[p_id_to_s_idx] - 1) + + # NOTE: Since query_lens[i] = 0 where i >= num_seqs, s_idx_to_num_tiles[i] + # where i >= num_seqs will also be 0. Therefore, s_idx_to_num_tiles.sum() + # will contain number of tiles for valid sequence. + num_tiles = s_idx_to_num_tiles.sum() + + return memory_ref.MetadataRef.create( + cfgs=cfg, + num_tiles=num_tiles, + p_id_to_s_idx=p_id_to_s_idx, + p_id_to_r_base=p_id_to_r_base, + p_id_to_r_size=p_id_to_r_size, + p_id_is_first_tile=p_id_is_first_tile, + p_id_is_last_tile=p_id_is_last_tile, + s_idx_has_initial_state=has_initial_state, + s_idx_to_state_indices=state_indices, + ) diff --git a/src/maxtext/models/kernels/gdn/pallas_mosaic_tpu.py b/src/maxtext/models/kernels/gdn/pallas_mosaic_tpu.py new file mode 100644 index 000000000..52d781e82 --- /dev/null +++ b/src/maxtext/models/kernels/gdn/pallas_mosaic_tpu.py @@ -0,0 +1,106 @@ +# 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. +# ============================================================================== +"""Pallas Mosaic TPU kernel implementation for Causal Conv1D Gated Delta Rule.""" + +import dataclasses +from typing import Optional, override + +import jax +from jax.experimental.pallas import tpu as pltpu +import jax.numpy as jnp +from tokamax._src.ops.causal_conv1d_gated_delta_rule import base + +try: + from maxtext.models.kernels.gdn import config + from maxtext.models.kernels.gdn import wrapper +except (ImportError, ModuleNotFoundError): + try: + from maxtext.src.maxtext.models.kernels.gdn import config + from maxtext.src.maxtext.models.kernels.gdn import wrapper + except (ImportError, ModuleNotFoundError): + from . import config + from . import wrapper + +GDNConfig = config.GDNConfig + + +@dataclasses.dataclass(frozen=True, kw_only=True) +class PallasMosaicTpuCausalConv1dGatedDeltaRule(base.CausalConv1dGatedDeltaRule[GDNConfig]): + """Wrapper for the tokamax Op API for Pallas Mosaic TPU kernel.""" + + # pylint: disable=redefined-outer-name + def _fwd( + self, + qkv: jax.Array, + b: jax.Array, + a: jax.Array, + conv_state: jax.Array, + recurrent_state: jax.Array, + conv_weight: jax.Array, + conv_bias: Optional[jax.Array], + a_log: jax.Array, + dt_bias: jax.Array, + query_start_loc: jax.Array, + state_indices: jax.Array, + distribution: jax.Array, + seq_lens: jax.Array, + *, + n_kq: int, + n_v: int, + d_k: int, + d_v: int, + kernel_size: int, + zero_initialize_out: bool = True, + compute_precision: jnp.dtype = jnp.float32.dtype, + decode_tile_size: int = 4, + mixed_tile_size: int = 64, + config: GDNConfig | None = None, + return_residuals: bool = False, + ) -> tuple[tuple[tuple[jax.Array, jax.Array], jax.Array], None]: + """Forward execution rule for Causal Conv1D Gated Delta Rule.""" + del return_residuals + _ = config + out_act, states, *_ = wrapper.fused_conv1d_gdn( + qkv=qkv, + b=b, + a=a, + conv_state=conv_state, + recurrent_state=recurrent_state, + conv_weight=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=n_kq, + n_v=n_v, + d_k=d_k, + d_v=d_v, + kernel_size=kernel_size, + zero_initialize_out=zero_initialize_out, + compute_precision=compute_precision, + decode_tile_size=decode_tile_size, + mixed_tile_size=mixed_tile_size, + ) + return (states, out_act), None + + @override + def supported_on(self, device: jax.Device) -> bool: + try: + return device.platform == "tpu" and pltpu.get_tpu_info().generation >= 6 + except Exception: # pylint: disable=broad-exception-caught + return False diff --git a/src/maxtext/models/kernels/gdn/tiling.py b/src/maxtext/models/kernels/gdn/tiling.py new file mode 100644 index 000000000..b666a7c41 --- /dev/null +++ b/src/maxtext/models/kernels/gdn/tiling.py @@ -0,0 +1,337 @@ +# 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. +# ============================================================================== + +"""Dynamic tiling heuristics and VMEM memory estimation for Fused Conv1D-GDN.""" + +from jax.experimental import pallas as pl +import jax.numpy as jnp + +try: + from maxtext.models.kernels.gdn import config +except (ImportError, ModuleNotFoundError): + try: + from maxtext.src.maxtext.models.kernels.gdn import config + except (ImportError, ModuleNotFoundError): + from . import config + + +def align_to(x: int, alignment: int) -> int: + """Aligns an integer upward to the nearest multiple of alignment.""" + return pl.cdiv(x, alignment) * alignment + + +def get_vmem_estimate_bytes( + tile_b: int, + chunk_sz: int, + n_kq: int, + n_v: int, + d_k: int, + d_v: int, + kernel_size: int, + act_in_bytes: int, + act_out_bytes: int, + conv_state_bytes: int, + rec_state_bytes: int, + num_lanes: int, + conv_state_dim_size: int, + is_decode: bool = False, +) -> int: + """Estimates total on-chip VMEM footprint in bytes for a GDN tile.""" + aligned_num_v_heads = align_to(n_v, num_lanes) + aligned_d_k = align_to(d_k, num_lanes) + aligned_d_v = align_to(d_v, num_lanes) + dim_size = align_to(2 * n_kq * d_k + n_v * d_v, num_lanes) + aligned_out_dim = align_to(n_v * d_v, num_lanes) + + # 1. Double-buffered input activation buffers (QKV, A, B). + qkv_bytes = 2 * (tile_b * chunk_sz * dim_size * act_in_bytes) + b_bytes = 2 * (tile_b * chunk_sz * aligned_num_v_heads * act_in_bytes) + a_bytes = 2 * (tile_b * chunk_sz * aligned_num_v_heads * act_in_bytes) + + # 2. Double-buffered state cache buffers (convolution and recurrent states). + conv_state_buffer_bytes = 2 * (tile_b * max(0, kernel_size - 1) * conv_state_dim_size * conv_state_bytes) + recurrent_state_buffer_bytes = 2 * (tile_b * n_v * d_v * d_k * rec_state_bytes) + + # 3. Double-buffered output activation buffer. + out_bytes = 2 * (tile_b * chunk_sz * aligned_out_dim * act_out_bytes) + + # 4. Temporary scratch buffers (allocated in non-batched mode). + if is_decode: + scratch_conv_bytes = 0 + scratch_recurrent_bytes = 0 + else: + scratch_conv_bytes = tile_b * max(0, kernel_size - 1) * dim_size * conv_state_bytes + scratch_recurrent_bytes = tile_b * n_v * d_v * d_k * rec_state_bytes + + # 5. Static weight cache references in on-chip memory. + weights_bytes = ((kernel_size - 1) * dim_size * 4) + (dim_size * 4) + (aligned_num_v_heads * 8) + + # 6. Working memory for intra-chunk recurrence and projections. + intermediate_bytes = n_v * (5 * chunk_sz * chunk_sz + 3 * chunk_sz * (aligned_d_v + aligned_d_k)) * 4 + + return ( + qkv_bytes + + b_bytes + + a_bytes + + conv_state_buffer_bytes + + recurrent_state_buffer_bytes + + out_bytes + + scratch_conv_bytes + + scratch_recurrent_bytes + + weights_bytes + + intermediate_bytes + ) + + +def calculate_decode_tile_size( + batch_size: int, + n_kq: int, + n_v: int, + d_k: int, + d_v: int, + conv_state_dim_size: int, + act_in_dtype: jnp.dtype, + act_out_dtype: jnp.dtype, + conv_state_dtype: jnp.dtype, + recurrent_state_dtype: jnp.dtype, + num_lanes: int, + vmem_capacity_limit_bytes: int, + kernel_size: int = 4, +) -> int: + """Derives optimal batch tile size for decode execution. + + Searches candidate batch tile sizes within maximum VMEM capacity limits. + + Args: + batch_size: Total batch size of the active decode sequence. + n_kq: Number of key/query heads. + n_v: Number of value heads. + d_k: Key head dimension. + d_v: Value head dimension. + conv_state_dim_size: Feature dimension size for conv state. + act_in_dtype: Data type for input activations. + act_out_dtype: Data type for output activations. + conv_state_dtype: Data type for conv state cache. + recurrent_state_dtype: Data type for recurrent state matrix. + num_lanes: Number of lanes for TPU vector layout alignment. + vmem_capacity_limit_bytes: Maximum allowed VMEM capacity in bytes. + kernel_size: 1D convolution kernel window size. + + Returns: + Derived batch tile size fitting within VMEM capacity limits. + """ + # Return a minimum valid tile size of 1 for empty or zero-length batches. + if batch_size <= 0: + return 1 + + act_in_bytes = jnp.dtype(act_in_dtype).itemsize + act_out_bytes = jnp.dtype(act_out_dtype).itemsize + conv_state_bytes = jnp.dtype(conv_state_dtype).itemsize + rec_state_bytes = jnp.dtype(recurrent_state_dtype).itemsize + + # Balance vector compute density against on-chip VMEM capacity: + # - Cap tile size across batch size tiers to maximize vector lane compute + # density. + # - Floor at tile_b = 4 for small batches to ensure compute density. + # - When value head count is large (n_v >= 64), recurrent state working + # memory scales up, so cap max_decode_b to 4 to prevent on-chip memory + # overflow. + if n_v >= 64: + max_decode_b = 2 + elif batch_size <= 64: + max_decode_b = 4 + elif batch_size <= 128: + max_decode_b = 8 + elif batch_size <= 256: + max_decode_b = 16 + else: + max_decode_b = 32 + + decode_candidates = [c for c in (32, 16, 8, 4, 2, 1) if c <= batch_size and c <= max_decode_b] + decode_tile_size = decode_candidates[-1] + + for cand in decode_candidates: + vmem_est = get_vmem_estimate_bytes( + tile_b=cand, + chunk_sz=1, + n_kq=n_kq, + n_v=n_v, + d_k=d_k, + d_v=d_v, + kernel_size=kernel_size, + act_in_bytes=act_in_bytes, + act_out_bytes=act_out_bytes, + conv_state_bytes=conv_state_bytes, + rec_state_bytes=rec_state_bytes, + num_lanes=num_lanes, + conv_state_dim_size=conv_state_dim_size, + is_decode=True, + ) + if vmem_est <= vmem_capacity_limit_bytes: + return cand + + return decode_tile_size + + +def calculate_mixed_tile_size( + seq_len: int, + n_kq: int, + n_v: int, + d_k: int, + d_v: int, + conv_state_dim_size: int, + act_in_dtype: jnp.dtype, + act_out_dtype: jnp.dtype, + conv_state_dtype: jnp.dtype, + recurrent_state_dtype: jnp.dtype, + num_lanes: int, + vmem_capacity_limit_bytes: int, + kernel_size: int = 4, +) -> int: + """Derives optimal chunk tile size for prefill and mixed execution. + + Searches candidate tile sizes within maximum VMEM capacity limits. + + Args: + seq_len: Sequence length of the active prefill or mixed sequence. + n_kq: Number of key/query heads. + n_v: Number of value heads. + d_k: Key head dimension. + d_v: Value head dimension. + conv_state_dim_size: Feature dimension size for conv state. + act_in_dtype: Data type for input activations. + act_out_dtype: Data type for output activations. + conv_state_dtype: Data type for conv state cache. + recurrent_state_dtype: Data type for recurrent state matrix. + num_lanes: Number of lanes for TPU vector layout alignment. + vmem_capacity_limit_bytes: Maximum allowed VMEM capacity in bytes. + kernel_size: 1D convolution kernel window size. + + Returns: + Derived chunk tile size fitting within VMEM capacity limits. + """ + # Return a minimum valid chunk size of 1 for empty or zero-length sequences. + if seq_len <= 0: + return 1 + + act_in_bytes = jnp.dtype(act_in_dtype).itemsize + act_out_bytes = jnp.dtype(act_out_dtype).itemsize + conv_state_bytes = jnp.dtype(conv_state_dtype).itemsize + rec_state_bytes = jnp.dtype(recurrent_state_dtype).itemsize + + # Limit chunk size to C <= 128: above 128, intra-chunk triangular + # solve operations and vector register pressure outweigh systolic compute + # density gains. + # When value head count is large (n_v >= 64), intra-chunk intermediate + # memory scales up, so cap chunk search space to C <= 64 to avoid on-chip + # memory overflow. + max_chunk_cap = 64 if n_v >= 64 else 128 + prefill_candidates = [c for c in (128, 64, 32, 16, 8, 4, 2, 1) if c <= seq_len and c <= max_chunk_cap] + mixed_tile_size = prefill_candidates[-1] + for candidate in prefill_candidates: + vmem_est = get_vmem_estimate_bytes( + tile_b=1, + chunk_sz=candidate, + n_kq=n_kq, + n_v=n_v, + d_k=d_k, + d_v=d_v, + kernel_size=kernel_size, + act_in_bytes=act_in_bytes, + act_out_bytes=act_out_bytes, + conv_state_bytes=conv_state_bytes, + rec_state_bytes=rec_state_bytes, + num_lanes=num_lanes, + conv_state_dim_size=conv_state_dim_size, + is_decode=False, + ) + if vmem_est <= vmem_capacity_limit_bytes: + return candidate + + return mixed_tile_size + + +def get_tile_sizes( + batch_size: int, + num_seqs: int, + padded_batch_size: int, + n_kq: int, + n_v: int, + d_k: int, + d_v: int, + kernel_size: int, + conv_state_dim_size: int, + act_in_dtype: jnp.dtype, + act_out_dtype: jnp.dtype, + conv_state_dtype: jnp.dtype, + recurrent_state_dtype: jnp.dtype, + num_lanes: int, + decode_tile_size: int | None = None, + mixed_tile_size: int | None = None, +) -> tuple[int, int]: + """Derives optimal decode and mixed tile sizes fitting within VMEM limits.""" + vmem_capacity_limit_bytes = config.get_vmem_limit_bytes() + + if decode_tile_size is None or decode_tile_size <= 0: + decode_tile_size = calculate_decode_tile_size( + batch_size=padded_batch_size, + n_kq=n_kq, + n_v=n_v, + d_k=d_k, + d_v=d_v, + conv_state_dim_size=conv_state_dim_size, + act_in_dtype=act_in_dtype, + act_out_dtype=act_out_dtype, + conv_state_dtype=conv_state_dtype, + recurrent_state_dtype=recurrent_state_dtype, + num_lanes=num_lanes, + vmem_capacity_limit_bytes=vmem_capacity_limit_bytes, + kernel_size=kernel_size, + ) + + if mixed_tile_size is None or mixed_tile_size <= 0: + if batch_size <= num_seqs: + # When all sequences have length 1 (decode), size prefill chunks to 1. + effective_prefill_seq_len = 1 + else: + # Estimate maximum sequence length across uniform and mixed prefill + # batches. In an adversarial mixed batch of B tokens with N_seq sequences, + # the largest prefill sequence length is bounded by + # B - (N_seq - 1) = B - N_seq + 1. + effective_prefill_seq_len = max( + batch_size // max(1, num_seqs), + batch_size - num_seqs + 1, + ) + + mixed_tile_size = calculate_mixed_tile_size( + seq_len=effective_prefill_seq_len, + n_kq=n_kq, + n_v=n_v, + d_k=d_k, + d_v=d_v, + conv_state_dim_size=conv_state_dim_size, + act_in_dtype=act_in_dtype, + act_out_dtype=act_out_dtype, + conv_state_dtype=conv_state_dtype, + recurrent_state_dtype=recurrent_state_dtype, + num_lanes=num_lanes, + vmem_capacity_limit_bytes=vmem_capacity_limit_bytes, + kernel_size=kernel_size, + ) + + # Guarantee strictly positive tile sizes (>= 1) for Pallas grid compilation. + decode_tile_size = max(1, min(decode_tile_size, batch_size)) + mixed_tile_size = max(1, min(mixed_tile_size, batch_size)) + return decode_tile_size, mixed_tile_size diff --git a/src/maxtext/models/kernels/gdn/vmem_ldst.py b/src/maxtext/models/kernels/gdn/vmem_ldst.py new file mode 100644 index 000000000..0253c6344 --- /dev/null +++ b/src/maxtext/models/kernels/gdn/vmem_ldst.py @@ -0,0 +1,263 @@ +# 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. +# ============================================================================== + +"""VMEM load/store pre-processing logic.""" + +import jax +from jax.experimental.pallas import tpu as pltpu +import jax.numpy as jnp + +try: + from maxtext.models.kernels.gdn import config + from maxtext.models.kernels.gdn import memory_ref +except (ImportError, ModuleNotFoundError): + try: + from maxtext.src.maxtext.models.kernels.gdn import config + from maxtext.src.maxtext.models.kernels.gdn import memory_ref + except (ImportError, ModuleNotFoundError): + from . import config + from . import memory_ref + + +def load_as_qkv_large(qkv_vmem_ref: jax.Ref, cfgs: config.GDNConfig) -> tuple[jax.Array, jax.Array, jax.Array]: + """Split qkv and transpose by performing 1 load per chunk for large layout. + + Args: + qkv_vmem_ref: qkv reference in VMEM containing concatenated values of q, k, + and v of shape [seq_tile_size, chunk_size, 1, num_kq_heads * kq_head_dim * + 2 + num_v_heads * v_head_dim]. + cfgs: GDN configuration object. + + Returns: + q, k: [seq_tile_size, num_kq_heads, chunk_size, kq_head_dim] + v: [seq_tile_size, num_v_heads, chunk_size, v_head_dim] + """ + + num_lanes = pltpu.get_tpu_info().num_lanes + lanes_per_col = qkv_vmem_ref.shape[-1] // num_lanes + kq_lanes_per_head = cfgs.kq_head_dim // num_lanes + k_offset = cfgs.num_kq_heads * kq_lanes_per_head + + q_large_list = [] + k_large_list = [] + v_large_list = [] + + qkv_slot_flat_ref = qkv_vmem_ref.reshape(-1, num_lanes) # pyrefly: ignore[missing-attribute] + for kq_head in range(cfgs.num_kq_heads): + q_head_list = [] + k_head_list = [] + for lane in range(kq_lanes_per_head): + q_lane = kq_head * kq_lanes_per_head + lane + k_lane = k_offset + q_lane + + q_head_list.append(qkv_slot_flat_ref[q_lane::lanes_per_col]) + k_head_list.append(qkv_slot_flat_ref[k_lane::lanes_per_col]) + q_large_list.append(jnp.concat(q_head_list, axis=-1)) + k_large_list.append(jnp.concat(k_head_list, axis=-1)) + v_offset = kq_lanes_per_head * cfgs.num_kq_heads * 2 + v_lanes_per_head = cfgs.v_head_dim // num_lanes + for v_head in range(cfgs.num_v_heads): + v_head_list = [] + for lane in range(v_lanes_per_head): + v_lane = v_offset + v_head * v_lanes_per_head + lane + v_head_list.append(qkv_slot_flat_ref[v_lane::lanes_per_col]) + v_large_list.append(jnp.concat(v_head_list, axis=-1)) + + q_large = jnp.stack(q_large_list, axis=0) + k_large = jnp.stack(k_large_list, axis=0) + v_large = jnp.stack(v_large_list, axis=0) + + 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]: + """Split qkv and transpose by performing 1 load per head for compact layout. + + Args: + qkv_vmem_ref: qkv reference in VMEM containing concatenated values of q, k, + and v of shape [seq_tile_size, chunk_size, 1, num_kq_heads * kq_head_dim * + 2 + num_v_heads * v_head_dim]. + cfg: GDN configuration object. + + Returns: + q, k: [seq_tile_size, num_kq_heads, chunk_size, 1, kq_head_dim] + v: [seq_tile_size, num_v_heads, chunk_size, 1, v_head_dim] + """ + + k_offset = cfg.num_kq_heads * cfg.kq_head_dim + v_offset = cfg.num_kq_heads * 2 * cfg.kq_head_dim + + q_compact_list = [] + k_compact_list = [] + v_compact_list = [] + + for kq_head in range(cfg.num_kq_heads): + q_start = kq_head * cfg.kq_head_dim + q_end = q_start + cfg.kq_head_dim + k_start = k_offset + q_start + k_end = k_start + cfg.kq_head_dim + q_compact_list.append(qkv_vmem_ref[..., q_start:q_end]) + k_compact_list.append(qkv_vmem_ref[..., k_start:k_end]) + for v_head in range(cfg.num_v_heads): + v_start = v_offset + v_head * cfg.v_head_dim + v_end = v_start + cfg.v_head_dim + v_compact_list.append(qkv_vmem_ref[..., v_start:v_end]) + + q_compact = jnp.stack(q_compact_list, axis=1) + k_compact = jnp.stack(k_compact_list, axis=1) + v_compact = jnp.stack(v_compact_list, axis=1) + + return q_compact, k_compact, v_compact + + +def load_compact_to_large(vmem_ref: jax.Ref) -> jax.Array: + """Use strided load to convert compact to large layout without transpose.""" + + # NOTE: Only support 32-bits for now. + assert vmem_ref.dtype.itemsize == 4 + assert vmem_ref.shape[-2] == 1 + col_size = vmem_ref.shape[-1] + new_shape = vmem_ref.shape[:-2] + (col_size,) + tpu_info = pltpu.get_tpu_info() + num_lanes = tpu_info.num_lanes + + vreg_list = [] + vmem_ref = vmem_ref.reshape(-1, col_size) # pyrefly: ignore[missing-attribute] + for col_start in range(0, col_size, num_lanes): + col_end = min(col_start + num_lanes, col_size) + vreg = vmem_ref[..., col_start:col_end] + vreg_list.append(vreg) + return jnp.concat(vreg_list, axis=-1).reshape(new_shape) + + +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, +) -> tuple[jax.Array, jax.Array, jax.Array]: + """Load correct states from HBM or prior tile, and masks invalid states. + + Reference metadata to select the appropriate prior states. If `is_first_tile` + is True, it selects states read from HBM. If it is False, it selects + carry states from previous tile. If `has_initial_state` is False, states are + zero initialized. + + Args: + metadata_ref: Metadata reference containing grid and sequence mappings. + p_id: Current Pallas program ID. + conv_state_slot_ref: Convolution state read from HBM of shape + [seq_tile_size, prev_kernel_size, 1, dim_size]. + recurrent_slot_ref: Recurrent state read from HBM of shape [seq_tile_size, + num_v_heads, kq_head_dim, v_head_dim]. + carry_conv_scratch_ref: Optional inter-tile convolution carry of shape + [seq_tile_size, prev_kernel_size, 1, dim_size]. + carry_recurrent_scratch_ref: Optional inter-tile recurrent state carry of + shape [seq_tile_size, num_v_heads, kq_head_dim, v_head_dim]. + cfg: GDN configuration object. + + Returns: + real_sizes: Valid token count per sequence tile of shape [seq_tile_size]. + prev_conv_state: Selected convolution state of shape [seq_tile_size, + prev_kernel_size, 1, dim_size] in float32. + prev_recurrent_state: Selected recurrent state of shape [seq_tile_size, + num_v_heads, kq_head_dim, v_head_dim]. + """ + + real_sizes_list = [] + prev_conv_state_list = [] + prev_recurrent_state_list = [] + + for idx in range(cfg.seq_tile_size): + record = metadata_ref.get_record(p_id, idx) + s_idx = record.s_idx + real_sizes = record.r_size + is_first_tile = record.is_first_tile + has_initial_state = metadata_ref.s_idx_has_initial_state[s_idx] + + # NOTE: Conv1D mandates fp32 due to its usage of compact layout. + hbm_conv_state = conv_state_slot_ref[idx].astype(jnp.float32) + prev_conv_state = jnp.where(has_initial_state, hbm_conv_state, 0) + + if carry_conv_scratch_ref is not None: + prev_tile_conv = carry_conv_scratch_ref[idx] + prev_conv_state = jnp.where(is_first_tile, prev_conv_state, prev_tile_conv) + + hbm_recurrent_state = recurrent_slot_ref[idx] + prev_recurrent_state = jnp.where(has_initial_state, hbm_recurrent_state, 0) + + if carry_recurrent_scratch_ref is not None: + prev_tile_recurrent_scratch = carry_recurrent_scratch_ref[idx] + prev_recurrent_state = jnp.where(is_first_tile, prev_recurrent_state, prev_tile_recurrent_scratch) + + real_sizes_list.append(real_sizes) + prev_conv_state_list.append(prev_conv_state) + prev_recurrent_state_list.append(prev_recurrent_state) + + real_sizes = jnp.stack(real_sizes_list, axis=0) + prev_conv_state = jnp.stack(prev_conv_state_list, axis=0) + prev_recurrent_state = jnp.stack(prev_recurrent_state_list, axis=0) + + return real_sizes, prev_conv_state, prev_recurrent_state + + +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, +) -> tuple[jax.Array, jax.Array, jax.Array, jax.Array, jax.Array]: + """Load activations from VMEM as a compact layout.""" + + qkv_vmem_ref[...] = qkv_vreg + q_compact, k_compact, v_compact = load_as_qkv_compact(qkv_vmem_ref, cfgs) + b_compact = jnp.expand_dims(b_vmem_ref[...], axis=1) + a_compact = jnp.expand_dims(a_vmem_ref[...], axis=1) + return q_compact, k_compact, v_compact, b_compact, a_compact + + +def load_activation_as_large( + qkv_vreg: jax.Array, + qkv_vmem_ref: jax.Ref, + b_vmem_ref: jax.Ref, + a_vmem_ref: jax.Ref, + cfgs: config.GDNConfig, +) -> tuple[jax.Array, jax.Array, jax.Array, jax.Array, jax.Array]: + """Load activations from VMEM as a large layout.""" + + qkv_vmem_ref[...] = qkv_vreg + + q_large_list = [] + k_large_list = [] + v_large_list = [] + for idx in range(cfgs.seq_tile_size): + q_large, k_large, v_large = load_as_qkv_large(qkv_vmem_ref.at[idx], cfgs) + q_large_list.append(q_large) + k_large_list.append(k_large) + v_large_list.append(v_large) + + q_large = jnp.stack(q_large_list, axis=0) + k_large = jnp.stack(k_large_list, axis=0) + v_large = jnp.stack(v_large_list, axis=0) + b_large = load_compact_to_large(b_vmem_ref) + a_large = load_compact_to_large(a_vmem_ref) + b_large = jnp.expand_dims(b_large, axis=1) + a_large = jnp.expand_dims(a_large, axis=1) + + return q_large, k_large, v_large, b_large, a_large diff --git a/src/maxtext/models/kernels/gdn/wrapper.py b/src/maxtext/models/kernels/gdn/wrapper.py new file mode 100644 index 000000000..cfaacbee2 --- /dev/null +++ b/src/maxtext/models/kernels/gdn/wrapper.py @@ -0,0 +1,539 @@ +# 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. +# ============================================================================== + +"""Top-level Pallas kernel wrapper for fused Conv1D-GDN with triangular inverse caching.""" + +import functools + +import jax +from jax.experimental import pallas as pl +from jax.experimental.pallas import tpu as pltpu +import jax.numpy as jnp + +try: + from maxtext.models.kernels.gdn import compute_conv1d + from maxtext.models.kernels.gdn import compute_gdn + from maxtext.models.kernels.gdn import config + from maxtext.models.kernels.gdn import memory_ref + from maxtext.models.kernels.gdn import metadata + from maxtext.models.kernels.gdn import tiling + from maxtext.models.kernels.gdn import vmem_ldst +except (ImportError, ModuleNotFoundError): + try: + from maxtext.src.maxtext.models.kernels.gdn import compute_conv1d + from maxtext.src.maxtext.models.kernels.gdn import compute_gdn + from maxtext.src.maxtext.models.kernels.gdn import config + from maxtext.src.maxtext.models.kernels.gdn import memory_ref + from maxtext.src.maxtext.models.kernels.gdn import metadata + from maxtext.src.maxtext.models.kernels.gdn import tiling + from maxtext.src.maxtext.models.kernels.gdn import vmem_ldst + except (ImportError, ModuleNotFoundError): + from . import compute_conv1d + from . import compute_gdn + from . import config + from . import memory_ref + from . import metadata + from . import tiling + from . import vmem_ldst + + +def inner_kernel( + # Inputs. + qkv_slot_ref: jax.Ref, # [seq, chunk, 1, dim_size] + b_slot_ref: jax.Ref, # [seq, chunk, 1, num_v_heads] + a_slot_ref: jax.Ref, # [seq, chunk, 1, num_v_heads] + conv_state_slot_ref: jax.Ref, # [seq, prev_kernel_size, 1, dim_size] + recurrent_slot_ref: jax.Ref, # [seq, num_v_heads, kq_head, v_head] + # Outputs. + out_slot_ref: jax.Array, # [seq * chunk, num_v_heads, v_head] + t_inv_slot_ref: jax.Array, # [seq, num_v_heads, chunk, chunk] + *args, + cfg: config.GDNConfig, + **kwargs, +) -> None: + """Orchestrates computation of Conv1D and GDN for a single tile. + + This kernel acts as a facade adhering to strict separation of concerns. It + operates VMEM reference without knowledge on DMA logic. Furthermore, the + kernel invokes vmem_ldst to pre-processes data needed for compute and + invokes compute_conv1d and compute_gdn for actual compute. + """ + if cfg.mode == config.GDNMode.PER_SEQ: + chunk_states_slot_ref = args[0] + metadata_ref = args[1] + weights_ref = args[2] + carry_conv_scratch_ref = args[3] if len(args) > 3 else None + carry_recurrent_scratch_ref = args[4] if len(args) > 4 else None + else: + chunk_states_slot_ref = None + metadata_ref = args[0] + weights_ref = args[1] + carry_conv_scratch_ref = args[2] if len(args) > 2 else None + carry_recurrent_scratch_ref = args[3] if len(args) > 3 else None + + p_id = pl.program_id(0) + + # Prepare states. + real_sizes, prev_conv, prev_recurrent = vmem_ldst.load_and_select_states( + metadata_ref=metadata_ref, + p_id=p_id, + conv_state_slot_ref=conv_state_slot_ref, + recurrent_slot_ref=recurrent_slot_ref, + carry_conv_scratch_ref=carry_conv_scratch_ref, + carry_recurrent_scratch_ref=carry_recurrent_scratch_ref, + cfg=cfg, + ) + + # Step 1: Conv1D. + qkv_in_compact = qkv_slot_ref[...].astype(jnp.float32) + qkv_in_compact = jnp.concat([prev_conv, qkv_in_compact], axis=1) + + # Prepare conv1d weights. + conv_weight = weights_ref.conv.weight[...].astype(jnp.float32) + conv_bias = None + if weights_ref.conv.bias is not None: + conv_bias = weights_ref.conv.bias[...].astype(jnp.float32) + + qkv_out_compact, new_conv_state = compute_conv1d.causal_conv1d( + real_sizes=real_sizes, + lhs=qkv_in_compact, + conv_weight=conv_weight, + conv_bias=conv_bias, + cfg=cfg, + ) + + conv_state_slot_ref[...] = new_conv_state + if carry_conv_scratch_ref is not None: + carry_conv_scratch_ref[...] = new_conv_state + + # Apply activation function. + qkv_out_compact = jax.nn.silu(qkv_out_compact) + + # Step 2: GDN. + padding_size = cfg.aligned_num_v_heads - cfg.num_v_heads + a_log = jnp.pad(weights_ref.gdn.a_log[...], ((0, padding_size))) + dt_bias = jnp.pad(weights_ref.gdn.dt_bias[...], ((0, padding_size))) + + if cfg.chunk_size == 1: + q_compact, k_compact, v_compact, b_compact, a_compact = vmem_ldst.load_activation_as_compact( + qkv_vreg=qkv_out_compact, + qkv_vmem_ref=qkv_slot_ref, + b_vmem_ref=b_slot_ref, + a_vmem_ref=a_slot_ref, + cfgs=cfg, + ) + + out, new_recurrent_state, t_inv = compute_gdn.recurrent_gdn( + q_compact=q_compact, + k_compact=k_compact, + v_compact=v_compact, + b_compact=b_compact, + a_compact=a_compact, + state_prev=prev_recurrent, + a_log=a_log, + dt_bias=dt_bias, + cfg=cfg, + real_sizes=real_sizes, + ) + else: + q_large, k_large, v_large, b_large, a_large = vmem_ldst.load_activation_as_large( + qkv_vreg=qkv_out_compact, + qkv_vmem_ref=qkv_slot_ref, + b_vmem_ref=b_slot_ref, + a_vmem_ref=a_slot_ref, + cfgs=cfg, + ) + + out, new_recurrent_state, t_inv = compute_gdn.chunked_gdn( + q_large=q_large, + k_large=k_large, + v_large=v_large, + b_large=b_large, + a_large=a_large, + state_prev=prev_recurrent, + a_log=a_log, + dt_bias=dt_bias, + cfg=cfg, + real_sizes=real_sizes, + ) + + # Store output, recurrent, and t_inv to vmem. + out_slot_ref[...] = out.astype(out_slot_ref.dtype) + recurrent_slot_ref[...] = new_recurrent_state.astype(recurrent_slot_ref.dtype) + t_inv_slot_ref[...] = t_inv.astype(t_inv_slot_ref.dtype) + if chunk_states_slot_ref is not None: + chunk_states_slot_ref[...] = prev_recurrent.astype(chunk_states_slot_ref.dtype) + + if carry_recurrent_scratch_ref is not None: + carry_recurrent_scratch_ref[...] = new_recurrent_state + + +def outer_kernel( + # Inputs. + metadata_ref: memory_ref.MetadataRef, + qkv_ref: jax.Array, + b_ref: jax.Array, + a_ref: jax.Array, + conv_state_ref: jax.Array, + recurrent_state_ref: jax.Array, + _: jax.Array, + weights_ref: memory_ref.WeightRefs, + # Outputs. + out_ref: jax.Array, + conv_state_out_ref: jax.Array, + recurrent_state_out_ref: jax.Array, + t_inv_ref: jax.Array, + *args, + carry_conv_scratch_ref: jax.Array | None = None, + carry_recurrent_scratch_ref: jax.Array | None = None, + cfg: config.GDNConfig, + **kwargs, +) -> None: + """Setup memory allocations and emit pipeline for running inner_kernel.""" + del conv_state_out_ref, recurrent_state_out_ref + + chunk_states_ref = args[0] if (len(args) > 0 and cfg.mode == config.GDNMode.PER_SEQ) else None + + allocs = memory_ref.create_allocs( + metadata_ref=metadata_ref, + qkv_ref=qkv_ref, + b_ref=b_ref, + a_ref=a_ref, + out_ref=out_ref, + conv_state_ref=conv_state_ref, + recurrent_state_ref=recurrent_state_ref, + cfg=cfg, + t_inv_ref=t_inv_ref, + chunk_states_ref=chunk_states_ref, + ) + qkv_alloc = allocs[0] + b_alloc = allocs[1] + a_alloc = allocs[2] + conv_alloc = allocs[3] + recurrent_alloc = allocs[4] + out_alloc = allocs[5] + t_inv_alloc = allocs[6] + chunk_states_alloc = allocs[7] if len(allocs) > 7 else None + + num_tiles = metadata_ref.num_tiles[...] + + out_specs = [out_alloc.spec, t_inv_alloc.spec] + if chunk_states_alloc is not None: + out_specs.append(chunk_states_alloc.spec) + + pipeline_func = pltpu.emit_pipeline( + body=functools.partial( + inner_kernel, + cfg=cfg, + ), + grid=(num_tiles,), + in_specs=( + qkv_alloc.spec, + b_alloc.spec, + a_alloc.spec, + conv_alloc.spec, + recurrent_alloc.spec, + ), + out_specs=tuple(out_specs), + ) + + @pl.with_scoped(allocations=allocs) + def _run(allocations): + out_args = [out_ref, t_inv_ref] + if chunk_states_ref is not None: + out_args.append(chunk_states_ref) + pipeline_func( + qkv_ref, + b_ref, + a_ref, + conv_state_ref, + recurrent_state_ref, + *out_args, + scratches=( + metadata_ref, + weights_ref, + carry_conv_scratch_ref, + carry_recurrent_scratch_ref, + ), + allocations=allocations, + ) + + # pylint: disable=no-value-for-parameter + _run() + + +@jax.jit( + donate_argnames=("conv_state", "recurrent_state"), + static_argnames=( + "n_kq", + "n_v", + "d_k", + "d_v", + "kernel_size", + "decode_tile_size", + "mixed_tile_size", + "zero_initialize_out", + "compute_precision", + "is_prefill_only", + ), +) +def fused_conv1d_gdn( + qkv: jax.Array, # [batch_size, n_kq * d_k * 2 + n_v * d_v = dim_size] + b: jax.Array, # [batch_size, n_v] + a: jax.Array, # [batch_size, n_v] + conv_state: jax.Array, # [num_seqs + 1, kernel_size - 1, dim_size] + recurrent_state: jax.Array, # [num_seqs + 1, nv, dk, dv] + conv_weight: jax.Array, # [kernel_size - 1, dim_size] + conv_bias: jax.Array | None, # [dim_size] + a_log: jax.Array, # [n_v] + dt_bias: jax.Array, # [n_v] + query_start_loc: jax.Array, # [num_seqs + 1] + state_indices: jax.Array, # [num_seqs] + distribution: jax.Array, # [3] + seq_lens: jax.Array, # [num_seqs] + *, + n_kq: int, + n_v: int, + d_k: int, + d_v: int, + kernel_size: int, + zero_initialize_out: bool = True, + compute_precision: jnp.dtype = jnp.float32.dtype, + decode_tile_size: int | None = None, + mixed_tile_size: int | None = None, + is_prefill_only: bool = False, +) -> tuple[jax.Array, tuple[jax.Array, jax.Array], jax.Array, jax.Array]: + """Perform conv1d and gdn in a single fused kernel, returning (out, states, t_inv, chunk_states).""" + act_in_dtype = qkv.dtype + act_out_dtype = qkv.dtype + conv_out_dtype = conv_state.dtype + recurrent_out_dtype = recurrent_state.dtype + assert a.dtype == b.dtype == qkv.dtype == act_in_dtype + + qkv = qkv.astype(jnp.float32) + b = b.astype(jnp.float32) + a = a.astype(jnp.float32) + conv_state = conv_state.astype(jnp.float32) + + # Step 1: Validate inputs. + num_seqs = state_indices.size + batch_size, dim = qkv.shape + assert conv_weight.shape == (dim, 1, kernel_size) + if conv_bias is not None: + assert conv_bias.shape == (dim,) + assert query_start_loc.shape == (num_seqs + 1,) + assert state_indices.shape == (num_seqs,) + assert distribution.shape == (3,) + + num_lanes = pltpu.get_tpu_info().num_lanes + packing = 4 // act_in_dtype.itemsize + padded_batch_size = pl.cdiv(batch_size, packing) * packing + conv_state_dim_size = conv_state.shape[-1] + + decode_tile_size, mixed_tile_size = tiling.get_tile_sizes( + batch_size=batch_size, + num_seqs=num_seqs, + padded_batch_size=padded_batch_size, + n_kq=n_kq, + n_v=n_v, + d_k=d_k, + d_v=d_v, + kernel_size=kernel_size, + conv_state_dim_size=conv_state_dim_size, + act_in_dtype=act_in_dtype, + act_out_dtype=act_out_dtype, + conv_state_dtype=conv_state.dtype, + recurrent_state_dtype=recurrent_state.dtype, + num_lanes=num_lanes, + decode_tile_size=decode_tile_size, + mixed_tile_size=mixed_tile_size, + ) + + batch_padding_size = padded_batch_size - batch_size + aligned_num_v_heads = tiling.align_to(n_v, num_lanes) + num_v_padding_size = aligned_num_v_heads - n_v + qkv = jnp.pad(qkv, ((0, batch_padding_size), (0, 0))) + b = jnp.pad(b, ((0, batch_padding_size), (0, num_v_padding_size))) + a = jnp.pad(a, ((0, batch_padding_size), (0, num_v_padding_size))) + + qkv = qkv.reshape(padded_batch_size, 1, -1) + b = b.reshape(padded_batch_size, 1, -1) + a = a.reshape(padded_batch_size, 1, -1) + + # Step 3: States and weights pre-processing. + conv_state_shape = conv_state.shape + conv_state = conv_state.reshape(-1, kernel_size - 1, 1, dim) + conv_weight = conv_weight.swapaxes(0, 2).astype(jnp.float32) + conv_bias = conv_bias.astype(jnp.float32) if conv_bias is not None else None + + # Step 4: Wrap inputs for the kernel. + conv_weights = memory_ref.ConvWeightsRef(weight=conv_weight, bias=conv_bias) + gdn_weights = memory_ref.GDNWeightsRef(a_log=a_log, dt_bias=dt_bias) + weights = memory_ref.WeightRefs(conv=conv_weights, gdn=gdn_weights) + + # Step 5: Create specs. + smem_spec = pl.BlockSpec(memory_space=pltpu.SMEM) + vmem_spec = pl.BlockSpec(memory_space=pltpu.VMEM) + hbm_spec = pl.BlockSpec(memory_space=pltpu.HBM) + weights_spec = jax.tree.map(lambda _: vmem_spec, weights) + + def call_kernel( + in_conv_state: jax.Array, + in_recurrent_state: jax.Array, + in_act: jax.Array | None, + mode: config.GDNMode, + ) -> tuple[jax.Array, ...]: + if mode == config.GDNMode.BATCHED: + tile_size = decode_tile_size + else: + tile_size = mixed_tile_size + + cfg = config.GDNConfig( + mode=mode, + batch_size=padded_batch_size, + kernel_size=kernel_size, + tile_size=tile_size, + dim_size=dim, + num_kq_heads=n_kq, + num_v_heads=n_v, + kq_head_dim=d_k, + v_head_dim=d_v, + dtypes=config.Dtypes( + act_in=act_in_dtype, + act_out=act_out_dtype, + compute=compute_precision, + recurrent_state=in_recurrent_state.dtype, + conv_state=in_conv_state.dtype, + ), + ) + + if mode == config.GDNMode.BATCHED: + metadata_obj = metadata.compute_batched_seq_metadata( + cfg=cfg, + seq_lens=seq_lens, + query_start_loc=query_start_loc, + state_indices=state_indices, + end_seq=distribution[0], + ) + else: + metadata_obj = metadata.compute_per_seq_metadata( + cfg=cfg, + seq_lens=seq_lens, + query_start_loc=query_start_loc, + state_indices=state_indices, + start_seq=distribution[0], + end_seq=distribution[-1], + ) + + metadata_spec = jax.tree.map(lambda _: smem_spec, metadata_obj) + + in_out_spec = None + input_output_aliases = {len(metadata_obj) + 3: 1, len(metadata_obj) + 4: 2} + out_shape = cfg.get_out_shape() + + if in_act is None and zero_initialize_out: + in_act = jnp.zeros_like(out_shape) + if in_act is not None: + out_shape = in_act + in_out_spec = hbm_spec + input_output_aliases[len(metadata_obj) + 5] = 0 + + num_chunks = cfg.batch_size // cfg.chunk_size + t_inv_shape = jax.ShapeDtypeStruct( + (num_chunks, cfg.num_v_heads, cfg.chunk_size, cfg.chunk_size), + cfg.dtypes.compute, + ) + + if mode == config.GDNMode.PER_SEQ: + chunk_states_shape = jax.ShapeDtypeStruct( + (num_chunks, cfg.num_v_heads, cfg.kq_head_dim, cfg.v_head_dim), + cfg.dtypes.compute, + ) + out_shape_tuple = ( + out_shape, + in_conv_state, + in_recurrent_state, + t_inv_shape, + chunk_states_shape, + ) + out_specs_tuple = (hbm_spec, hbm_spec, hbm_spec, hbm_spec, hbm_spec) + else: + out_shape_tuple = ( + out_shape, + in_conv_state, + in_recurrent_state, + t_inv_shape, + ) + out_specs_tuple = (hbm_spec, hbm_spec, hbm_spec, hbm_spec) + + return pl.pallas_call( + functools.partial(outer_kernel, cfg=cfg), + out_shape=out_shape_tuple, + in_specs=( + metadata_spec, + hbm_spec, + hbm_spec, + hbm_spec, + hbm_spec, + hbm_spec, + in_out_spec, + weights_spec, + ), + out_specs=out_specs_tuple, + scratch_shapes=cfg.get_scratch_shape_dict(), + input_output_aliases=input_output_aliases, + compiler_params=pltpu.CompilerParams( + disable_bounds_checks=True, + vmem_limit_bytes=config.get_vmem_limit_bytes(), + ), + name=cfg.get_kernel_name(), + metadata=cfg.get_metadata(), + )( + metadata_obj, + qkv, + b, + a, + in_conv_state, + in_recurrent_state, + in_act, + weights, + ) + + if not is_prefill_only: + try: + if int(distribution[0]) == 0: + is_prefill_only = True + except (TypeError, ValueError, jax.errors.TracerIntegerConversionError): + pass + + if not is_prefill_only: + out_act, out_conv_state, out_recurrent_state, _ = call_kernel( + conv_state, recurrent_state, None, config.GDNMode.BATCHED + ) + else: + out_act = None + out_conv_state = conv_state + out_recurrent_state = recurrent_state + + out_act, out_conv_state, out_recurrent_state, t_inv, chunk_states = call_kernel( + out_conv_state, out_recurrent_state, out_act, config.GDNMode.PER_SEQ + ) + + out_act = out_act.reshape(padded_batch_size, -1)[:batch_size] + out_conv_state = out_conv_state.astype(conv_out_dtype) + out_conv_state = out_conv_state.reshape(conv_state_shape) + out_recurrent_state = out_recurrent_state.astype(recurrent_out_dtype) + + return out_act, (out_conv_state, out_recurrent_state), t_inv, chunk_states