From 850ce5d1b8c6940fa2e607136ef8fd77f5372ccf Mon Sep 17 00:00:00 2001 From: Daniel Mandragona Date: Fri, 14 Aug 2026 15:40:27 -0700 Subject: [PATCH] No public description Tested: PiperOrigin-RevId: 964926678 --- src/maxtext/configs/base.yml | 9 + src/maxtext/configs/types.py | 24 + src/maxtext/kernels/mhc/__init__.py | 38 ++ src/maxtext/kernels/mhc/api.py | 110 ++++ src/maxtext/kernels/mhc/common.py | 541 ++++++++++++++++++ src/maxtext/kernels/mhc/mhc_kernels_bwd.py | 467 ++++++++++++++++ src/maxtext/kernels/mhc/mhc_kernels_fwd.py | 328 +++++++++++ src/maxtext/layers/mhc.py | 113 ++-- tests/unit/mhc_test.py | 615 ++++++++++++++++++++- 9 files changed, 2192 insertions(+), 53 deletions(-) create mode 100644 src/maxtext/kernels/mhc/__init__.py create mode 100644 src/maxtext/kernels/mhc/api.py create mode 100644 src/maxtext/kernels/mhc/common.py create mode 100644 src/maxtext/kernels/mhc/mhc_kernels_bwd.py create mode 100644 src/maxtext/kernels/mhc/mhc_kernels_fwd.py diff --git a/src/maxtext/configs/base.yml b/src/maxtext/configs/base.yml index c88332f537..a4cf6877e3 100644 --- a/src/maxtext/configs/base.yml +++ b/src/maxtext/configs/base.yml @@ -1339,6 +1339,15 @@ sinkhorn_iterations: 20 # the expensive sinkhorn iterations, the downside to this approach is that # it is factorial in k. enable_mhc_lite: False +# Whether to use the Pallas TPU kernel implementation for mHC-lite when running on TPU. +use_mhc_pallas_kernel: False +# Block size for forward pass of MHC Pallas kernel. +mhc_pallas_kernel_fwd_block_size: 256 +# Block size for backward pass of MHC Pallas kernel. +# The backward pass is more memory intensive, so we were running into OOMs with +# TPU v7x if we go higher. For TPU v6 (Trillium), there is more VMEM, so 256 +# works and provides better results. +mhc_pallas_kernel_bwd_block_size: 128 ################################## DeepSeek Engram ################################## # Indices of transformer layers where Engram are integrated; leave empty [] to disable. diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index a86b36d96a..0c0197db99 100644 --- a/src/maxtext/configs/types.py +++ b/src/maxtext/configs/types.py @@ -1750,6 +1750,30 @@ class ManifoldConstrainedHyperConnections(BaseModel): "Practical only for a small mhc_expansion_rate (e.g., k=4)." ), ) + use_mhc_pallas_kernel: bool = Field( + False, + description=( + "Whether to use the Pallas TPU kernel implementation for" + " mHC-lite when running on TPU. Requires enable_mhc_lite=True." + ), + ) + mhc_pallas_kernel_fwd_block_size: int = Field( + 256, + description="Block size for forward pass of MHC Pallas kernel.", + ) + mhc_pallas_kernel_bwd_block_size: int = Field( + 128, + description=( + "Block size for backward pass of MHC Pallas kernel. Default of 128 is" + " optimal for TPU v7 memory constraints; 256 is optimal for TPU v6." + ), + ) + + @model_validator(mode="after") + def validate_mhc_kernel(self) -> "ManifoldConstrainedHyperConnections": + if self.use_mhc_pallas_kernel and not self.enable_mhc_lite: + raise ValueError("use_mhc_pallas_kernel=True requires enable_mhc_lite=True.") + return self class DilocoParams(BaseModel): diff --git a/src/maxtext/kernels/mhc/__init__.py b/src/maxtext/kernels/mhc/__init__.py new file mode 100644 index 0000000000..d07f40fa93 --- /dev/null +++ b/src/maxtext/kernels/mhc/__init__.py @@ -0,0 +1,38 @@ +# 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 +# +# https://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. +"""MaxText mHC-lite Pallas kernel package.""" + +from maxtext.kernels.mhc.api import MhcCoeffGradients +from maxtext.kernels.mhc.api import MhcCoeffOutputs +from maxtext.kernels.mhc.api import MhcCoeffParams +from maxtext.kernels.mhc.api import MhcContext +from maxtext.kernels.mhc.api import MhcDims +from maxtext.kernels.mhc.api import MhcKernelConfig +from maxtext.kernels.mhc.api import MhcWeights +from maxtext.kernels.mhc.api import post +from maxtext.kernels.mhc.api import pre +from maxtext.kernels.mhc.common import UnsupportedInputError + +__all__ = [ + "pre", + "post", + "MhcContext", + "MhcWeights", + "MhcKernelConfig", + "MhcDims", + "MhcCoeffParams", + "MhcCoeffOutputs", + "MhcCoeffGradients", + "UnsupportedInputError", +] diff --git a/src/maxtext/kernels/mhc/api.py b/src/maxtext/kernels/mhc/api.py new file mode 100644 index 0000000000..b01e44dacf --- /dev/null +++ b/src/maxtext/kernels/mhc/api.py @@ -0,0 +1,110 @@ +# 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 +# +# https://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. +"""Public API entrypoints for mHC-lite Pallas TPU kernel.""" + +from typing import Literal, Sequence +import jax +from maxtext.kernels.mhc import common +from maxtext.kernels.mhc import mhc_kernels_fwd + +type Implementation = Literal["mosaic", "mosaic_tpu", "xla"] +MhcContext = common.MHCContext +MhcWeights = common.MhcWeights +MhcKernelConfig = common.MhcKernelConfig +MhcDims = common.MhcDims +MhcCoeffParams = common.MhcCoeffParams +MhcCoeffOutputs = common.MhcCoeffOutputs +MhcCoeffGradients = common.MhcCoeffGradients + + +def _validate_implementation( + implementation: Implementation | Sequence[Implementation] | None, +) -> None: + """Validates that the requested implementation is supported.""" + if implementation is None: + return + valid = ("mosaic", "mosaic_tpu", "xla") + if isinstance(implementation, str): + if implementation not in valid: + raise ValueError(f"Unsupported implementation: '{implementation}'") + return + if not any(imp in valid for imp in implementation): + raise ValueError(f"Unsupported implementation: {implementation}") + + +def pre( + x: jax.Array, + weights: common.MhcWeights, + permutations: jax.Array, + *, + config: common.MhcKernelConfig = common.MhcKernelConfig(), + implementation: Implementation | Sequence[Implementation] | None = None, +) -> tuple[jax.Array, MhcContext]: + """Computes the branch input and opaque context for an mHC-wrapped branch. + + Uses the Pallas TPU kernel when running on TPU and the shape/dtype + contract is supported. + + Args: + x: Input streams of shape `(batch, sequence, streams, embedding)`. + weights: Structured `MhcWeights` container with all layer parameters. + permutations: All permutation matrices of shape `(num_permutations, streams, streams)`. + config: Structured `MhcKernelConfig` tuning and compiler configuration. + implementation: Preferred implementation (`"mosaic"` or `"mosaic_tpu"`). + + Returns: + A tuple `(layer_input, context)` where `layer_input` feeds the wrapped + model branch, and `context` is passed unchanged to `post`. + """ + permutations = jax.lax.stop_gradient(permutations) + _validate_implementation(implementation) + layer_input, kernel_context = mhc_kernels_fwd.pre( + x, + weights, + permutations, + config=config, + ) + x_context, h_post, residual = kernel_context + return layer_input, MhcContext( + x=x_context, + h_post=h_post, + residual=residual, + implementation="mosaic", + ) + + +def post( + layer_output: jax.Array, + context: MhcContext, + *, + config: common.MhcKernelConfig = common.MhcKernelConfig(), +) -> jax.Array: + """Runs the post-gate and residual stream mixing. + + Args: + layer_output: Output from the wrapped branch of shape `(batch, sequence, embedding)`. + context: Opaque `MhcContext` returned by `pre`. + config: Structured `MhcKernelConfig` tuning and compiler configuration. + + Returns: + Mixed output streams of shape `(batch, sequence, streams, embedding)`. + """ + if context.implementation not in ("mosaic", "mosaic_tpu"): + raise ValueError(f"Unsupported implementation in MhcContext: '{context.implementation}'") + kernel_context = (context.x, context.h_post, context.residual) + return mhc_kernels_fwd.post( + layer_output, + kernel_context, + config=config, + ) diff --git a/src/maxtext/kernels/mhc/common.py b/src/maxtext/kernels/mhc/common.py new file mode 100644 index 0000000000..98bd466689 --- /dev/null +++ b/src/maxtext/kernels/mhc/common.py @@ -0,0 +1,541 @@ +# 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 +# +# https://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. +"""Shared zero-cost abstractions, block math, and tiling for mHC-lite Pallas kernels.""" + +import dataclasses +import jax +from jax.experimental import pallas as pl +from jax.experimental.pallas import tpu as pltpu +import jax.numpy as jnp + +DEFAULT_BLOCK_SIZE = 128 +DEFAULT_BWD_BLOCK_SIZE = 32 +DEFAULT_POST_BWD_BLOCK_SIZE = 32 +DEFAULT_POST_BWD_FEATURE_BLOCK_SIZE = 1024 +DEFAULT_VMEM_LIMIT_BYTES = 128 * 1024 * 1024 +PARALLEL_DIMENSION_SEMANTICS = (pltpu.PARALLEL,) +SEQUENTIAL_DIMENSION_SEMANTICS = (pltpu.ARBITRARY,) +SEQUENTIAL_2D_DIMENSION_SEMANTICS = (pltpu.ARBITRARY, pltpu.ARBITRARY) +# Kernel-level context tuple: `(x, h_post, residual)`. +type KernelContext = tuple[jax.Array, jax.Array, jax.Array] + + +class UnsupportedInputError(ValueError): + """Known Mosaic shape, dtype, or tiling restriction.""" + + +@jax.tree_util.register_dataclass +@dataclasses.dataclass(frozen=True) +class MHCContext: + """Opaque token-local context passed from `pre` to `post`.""" + + x: jax.Array + h_post: jax.Array + residual: jax.Array + implementation: str = dataclasses.field(metadata={"static": True}) + + +@dataclasses.dataclass(frozen=True) +class MhcKernelConfig: + """Compiler and tiling configuration for mHC Pallas kernels.""" + + block_size: int = DEFAULT_BLOCK_SIZE + bwd_block_size: int = DEFAULT_BWD_BLOCK_SIZE + bwd_feature_block_size: int = DEFAULT_POST_BWD_FEATURE_BLOCK_SIZE + vmem_limit_bytes: int = DEFAULT_VMEM_LIMIT_BYTES + rms_epsilon: float = 1e-5 + pre_mapping_epsilon: float = 1e-6 + interpret: bool = False + + +@dataclasses.dataclass(frozen=True) +class MhcDims: + """Static dimension and cost descriptor for mHC-lite.""" + + tokens: int + streams: int + embedding: int + num_permutations: int = 24 + + @property + def flattened_size(self) -> int: + return self.streams * self.embedding + + @property + def pre_slice(self) -> slice: + return slice(0, self.streams) + + @property + def post_slice(self) -> slice: + return slice(self.streams, 2 * self.streams) + + @property + def res_slice(self) -> slice: + return slice(2 * self.streams, 2 * self.streams + self.num_permutations) + + @property + def phi_cols(self) -> int: + return 2 * self.streams + self.num_permutations + + def coeff_fwd_cost(self) -> pl.CostEstimate: + """Estimates compute and memory cost for the forward coefficient kernel. + + Mathematical Derivations: + - FLOPs: + 1. Fused Projection GEMM: `flattened_x (T, k*d) @ phi (k*d, 2*k+P)` + [einsum: `tf,fc->tc` where `f = k*d`, `c = 2*k + P`] + = `2 * T * (k*d) * (2*k + P)` FLOPs (2 FLOPs per multiply-accumulate). + 2. Permutation GEMM: `softmax_weights (T, P) @ permutations (P, k*k)` + [einsum: `tp,pij->tij` where `i, j` are streams `k, k`] + = `2 * T * P * k*k` FLOPs. + Total FLOPs = `2 * tokens * flattened_size * phi_cols + 2 * tokens * num_permutations * streams^2`. + - Transcendentals: + Sigmoid gating for pre/post gates and softmax over permutation logits + = `T * (k + k + P) = tokens * (2 * streams + num_permutations) = tokens * phi_cols`. + - Bytes Accessed: + Read input `x` (bfloat16: `T * k * d * 2` bytes) + read `phi` (float32: `(k*d) * phi_cols * 4` bytes) + = `tokens * streams * embedding * 2 + flattened_size * phi_cols * 4` bytes. + + Returns: + pl.CostEstimate with estimated FLOPs, transcendentals, and bytes accessed. + """ + flops = int( + 2 * self.tokens * self.flattened_size * self.phi_cols + + 2 * self.tokens * self.num_permutations * self.streams * self.streams + ) + transcendentals = int(self.tokens * (2 * self.streams + self.num_permutations)) + bytes_accessed = int(self.tokens * self.streams * self.embedding * 2 + self.flattened_size * self.phi_cols * 4) + return pl.CostEstimate(flops=flops, transcendentals=transcendentals, bytes_accessed=bytes_accessed) + + def pre_apply_fwd_cost(self) -> pl.CostEstimate: + """Estimates compute and memory cost for the forward pre-apply gating kernel. + + Mathematical Derivations: + - FLOPs: + Stream reduction `sum_s (h_pre[t, s] * x[t, s, d])` [einsum: `ts,tsd->td`] across `k` streams + for each of the `T * d` output elements (vector-matrix product `(1, k) @ (k, d) -> (1, d)` per token). + = `2 * T * k * d = 2 * tokens * streams * embedding` FLOPs. + - Transcendentals: + 0 (linear scaling and summation). + - Bytes Accessed: + Read input `x` (bfloat16: `T * k * d * 2` bytes) + write `layer_input` (bfloat16: `T * d * 2` bytes) + = `tokens * streams * embedding * 2 + tokens * embedding * 2` bytes. + + Returns: + pl.CostEstimate with estimated FLOPs, transcendentals, and bytes accessed. + """ + flops = int(2 * self.tokens * self.streams * self.embedding) + bytes_accessed = int(self.tokens * self.streams * self.embedding * 2 + self.tokens * self.embedding * 2) + return pl.CostEstimate(flops=flops, transcendentals=0, bytes_accessed=bytes_accessed) + + def post_apply_fwd_cost(self) -> pl.CostEstimate: + """Estimates compute and memory cost for the forward post-apply mixing kernel. + + Mathematical Derivations: + - FLOPs: + 1. Residual mixing contraction `sum_s_in (residual[t, s_in, s_out] * x[t, s_in, d])` + [einsum: `tkj,tkd->tjd`] (matrix product `(k, k) @ (k, d) -> (k, d)` per token): + `2 * k` FLOPs per output element over `T * k * d` elements = `2 * T * k^2 * d` FLOPs. + 2. Post-gating and accumulation `h_post[t, s] * layer_output[t, d] + residual_mix[t, s, d]` + [einsum: `ts,td->tsd`]: + `2 * T * k * d` FLOPs (1 multiply + 1 add per element). + Total FLOPs = `2 * tokens * streams^2 * embedding + 2 * tokens * streams * embedding`. + - Transcendentals: + 0. + - Bytes Accessed: + Read `x` (bfloat16: `T * k * d * 2` bytes) + write output (bfloat16: `T * k * d * 2` bytes) + + read `layer_output` and gating/residual context (`T * d * 4` bytes) + = `2 * tokens * streams * embedding * 2 + tokens * embedding * 4` bytes. + + Returns: + pl.CostEstimate with estimated FLOPs, transcendentals, and bytes accessed. + """ + flops = int( + 2 * self.tokens * self.streams * self.streams * self.embedding + 2 * self.tokens * self.streams * self.embedding + ) + bytes_accessed = int(2 * self.tokens * self.streams * self.embedding * 2 + self.tokens * self.embedding * 4) + return pl.CostEstimate(flops=flops, transcendentals=0, bytes_accessed=bytes_accessed) + + def pre_apply_bwd_cost(self) -> pl.CostEstimate: + """Estimates compute and memory cost for the backward pre-apply kernel. + + Mathematical Derivations: + - FLOPs: + 1. Activation cotangent `d_x = h_pre * d_layer_input` [einsum: `ts,td->tsd`] (multiply) + + accumulation with `d_x_acc` (add) = `2 * T * k * d` FLOPs. + 2. Gate cotangent `d_h_pre = sum_d (x * d_layer_input)` [einsum: `tsd,td->ts`] (multiply + reduce-add) + = `2 * T * k * d` FLOPs. + Total FLOPs = `4 * tokens * streams * embedding`. + - Transcendentals: + 0. + - Bytes Accessed: + Read `x` (`T * k * d * 2`) + read `d_x_acc` (`T * k * d * 2`) + write `d_x_acc_out` (`T * k * d * 2`) + + read `d_layer_input` (`T * d * 2`) + = `3 * tokens * streams * embedding * 2 + tokens * embedding * 2` bytes. + + Returns: + pl.CostEstimate with estimated FLOPs, transcendentals, and bytes accessed. + """ + flops = int(4 * self.tokens * self.streams * self.embedding) + bytes_accessed = int(3 * self.tokens * self.streams * self.embedding * 2 + self.tokens * self.embedding * 2) + return pl.CostEstimate(flops=flops, transcendentals=0, bytes_accessed=bytes_accessed) + + def coeff_bwd_cost(self) -> pl.CostEstimate: + """Estimates compute and memory cost for the backward coefficient kernel. + + Mathematical Derivations: + - FLOPs: + In-kernel forward recomputation plus VJP of `mhc_coeffs` (evaluating matrix multiplications + for activation cotangents `d_x` [einsum: `tc,fc->tf`], parameter gradients `d_phi` [einsum: `tf,tc->fc`], + and permutation gradients [einsum: `tij,pij->tp`]), scaling forward GEMMs by 2x: + Total FLOPs = `2 * (2 * tokens * flattened_size * phi_cols + 2 * tokens * num_permutations * streams^2)`. + - Transcendentals: + Gating and softmax evaluations across tokens + = `T * (k + k + P) = tokens * (2 * streams + num_permutations)`. + - Bytes Accessed: + Read `x` (`T * k * d * 2`) + read `d_x_acc` (`T * k * d * 2`) + write `d_x` (`T * k * d * 2`) + + read `phi` and write `d_phi` (`2 * (k*d) * phi_cols * 4`) + = `3 * tokens * streams * embedding * 2 + 2 * flattened_size * phi_cols * 4` bytes. + + Returns: + pl.CostEstimate with estimated FLOPs, transcendentals, and bytes accessed. + """ + flops = int( + 2 + * ( + 2 * self.tokens * self.flattened_size * self.phi_cols + + 2 * self.tokens * self.num_permutations * self.streams * self.streams + ) + ) + transcendentals = int(self.tokens * (2 * self.streams + self.num_permutations)) + bytes_accessed = int( + 3 * self.tokens * self.streams * self.embedding * 2 + 2 * self.flattened_size * self.phi_cols * 4 + ) + return pl.CostEstimate(flops=flops, transcendentals=transcendentals, bytes_accessed=bytes_accessed) + + def post_apply_bwd_cost(self) -> pl.CostEstimate: + """Estimates compute and memory cost for the backward post-apply kernel. + + Mathematical Derivations: + - FLOPs: + Backward VJP of post-apply evaluates transposed residual contraction `d_x = residual^T @ d_output` + [einsum: `tkj,tjd->tkd`] (`2 * T * k^2 * d`), gate reduction `d_layer_output = sum_s (h_post * d_output)` + [einsum: `ts,tsd->td`] (`2 * T * k * d`), and feature reductions `d_h_post` [einsum: `td,tsd->ts`] + (`2 * T * k * d`) and `d_residual` [einsum: `tjd,tkd->tjk`] (`2 * T * k^2 * d`): + Total FLOPs = `4 * tokens * streams^2 * embedding + 4 * tokens * streams * embedding` + = `2 * (2 * tokens * streams^2 * embedding + 2 * tokens * streams * embedding)`. + - Transcendentals: + 0. + - Bytes Accessed: + Read `x` (`T * k * d * 2`) + read `d_output` (`T * k * d * 2`) + write `d_x` (`T * k * d * 2`) + + read/write `layer_output` and `d_layer_output` (`2 * tokens * embedding * 4`) + = `3 * tokens * streams * embedding * 2 + 2 * tokens * embedding * 4` bytes. + + Returns: + pl.CostEstimate with estimated FLOPs, transcendentals, and bytes accessed. + """ + flops = int( + 2 + * ( + 2 * self.tokens * self.streams * self.streams * self.embedding + + 2 * self.tokens * self.streams * self.embedding + ) + ) + bytes_accessed = int(3 * self.tokens * self.streams * self.embedding * 2 + 2 * self.tokens * self.embedding * 4) + return pl.CostEstimate(flops=flops, transcendentals=0, bytes_accessed=bytes_accessed) + + +@jax.tree_util.register_dataclass +@dataclasses.dataclass(frozen=True) +class MhcWeights: + """Structured layer weights container for mHC-lite.""" + + norm_scale: jax.Array + pre_alpha: jax.Array + pre_bias: jax.Array + pre_scale: jax.Array + post_alpha: jax.Array + post_bias: jax.Array + post_scale: jax.Array + res_alpha: jax.Array + res_bias: jax.Array + res_scale: jax.Array + + def fold_norm_scale(self) -> jax.Array: + return fold_norm_scale(self.norm_scale, self.pre_alpha, self.post_alpha, self.res_alpha) + + def to_coeff_params(self) -> "MhcCoeffParams": + return MhcCoeffParams( + phi=self.fold_norm_scale(), + pre_scale=self.pre_scale, + pre_bias=self.pre_bias, + post_scale=self.post_scale, + post_bias=self.post_bias, + res_scale=self.res_scale, + res_bias=self.res_bias, + ) + + +@jax.tree_util.register_dataclass +@dataclasses.dataclass(frozen=True) +class MhcCoeffParams: + """Parameters required by the coefficient kernel.""" + + phi: jax.Array + pre_scale: jax.Array + pre_bias: jax.Array + post_scale: jax.Array + post_bias: jax.Array + res_scale: jax.Array + res_bias: jax.Array + + +@jax.tree_util.register_dataclass +@dataclasses.dataclass(frozen=True) +class MhcCoeffOutputs: + """Outputs generated by the coefficient kernel.""" + + h_pre: jax.Array + h_post: jax.Array + residual: jax.Array + + +@jax.tree_util.register_dataclass +@dataclasses.dataclass(frozen=True) +class MhcCoeffGradients: + """Backward gradients matching MhcCoeffParams.""" + + phi: jax.Array + pre_scale: jax.Array + pre_bias: jax.Array + post_scale: jax.Array + post_bias: jax.Array + res_scale: jax.Array + res_bias: jax.Array + + +def whole(shape: tuple[int, ...]) -> pl.BlockSpec: + """Returns a full-array BlockSpec for values that stay VMEM-resident.""" + return pl.BlockSpec(shape, lambda _: tuple(0 for _ in shape)) + + +def token_block_spec(shape: tuple[int, ...], block_size: int) -> pl.BlockSpec: + """Returns a BlockSpec tiling the leading token dimension.""" + block_shape = (block_size,) + shape[1:] + return pl.BlockSpec(block_shape, lambda i: (i,) + tuple(0 for _ in shape[1:])) + + +def feature_tiled_block_spec( + shape: tuple[int, ...], + block_size: int, + feature_block_size: int, + tiled_feature: bool = True, +) -> pl.BlockSpec: + """Returns a 2D BlockSpec over (token_idx, feature_idx) grid.""" + if tiled_feature: + block_shape = (block_size,) + shape[1:-1] + (feature_block_size,) + return pl.BlockSpec( + block_shape, + lambda token, feature: (token,) + tuple(0 for _ in shape[1:-1]) + (feature,), + ) + block_shape = (block_size,) + shape[1:] + return pl.BlockSpec( + block_shape, + lambda token, feature: (token,) + tuple(0 for _ in shape[1:]), + ) + + +def fold_norm_scale(norm_scale, pre_alpha, post_alpha, res_alpha) -> jax.Array: + """Folds the RMSNorm channel scale into the three projections.""" + alpha = jnp.concatenate((pre_alpha, post_alpha, res_alpha), axis=-1) + return norm_scale.astype(jnp.float32)[:, None] * alpha.astype(jnp.float32) + + +def fused_project_and_norm(flattened_x: jax.Array, phi: jax.Array, rms_epsilon: float) -> jax.Array: + """Computes linear projection fused with RMSNorm on the fly.""" + projected = jnp.dot(flattened_x, phi.astype(jnp.bfloat16), preferred_element_type=jnp.float32) + flattened_f32 = flattened_x.astype(jnp.float32) + mean_square = jnp.mean(flattened_f32 * flattened_f32, axis=-1, keepdims=True) + return projected * jax.lax.rsqrt(mean_square + rms_epsilon) + + +def compute_sigmoid_gate( + logits: jax.Array, + scale: jax.Array, + bias: jax.Array, + *, + multiplier: float = 1.0, + epsilon: float = 0.0, +) -> jax.Array: + """Computes scaled/biased sigmoid gating.""" + gate = jax.nn.sigmoid(scale.astype(jnp.float32) * logits + bias.astype(jnp.float32)) + if multiplier != 1.0: + gate = multiplier * gate + if epsilon != 0.0: + gate = gate + epsilon + return gate + + +def compute_residual_permutations( + logits: jax.Array, + scale: jax.Array, + bias: jax.Array, + permutations: jax.Array, + dims: MhcDims, +) -> jax.Array: + """Computes softmax permutation weighting and contraction.""" + weights = jax.nn.softmax( + scale.astype(jnp.float32) * logits + bias.astype(jnp.float32), + axis=-1, + ) + return jnp.dot( + weights, + permutations.reshape(dims.num_permutations, dims.streams * dims.streams).astype(jnp.float32), + ).reshape(dims.tokens, dims.streams, dims.streams) + + +def mhc_coeffs( + x: jax.Array, + coeff_params: MhcCoeffParams, + permutations: jax.Array, + *, + rms_epsilon: float, + pre_mapping_epsilon: float, +) -> MhcCoeffOutputs: + """Computes all mHC-lite coefficients without materializing normalized x.""" + tokens, streams, embedding = x.shape + dims = MhcDims(tokens=tokens, streams=streams, embedding=embedding, num_permutations=permutations.shape[0]) + flattened = x.reshape(tokens, dims.flattened_size) + projected = fused_project_and_norm(flattened, coeff_params.phi, rms_epsilon) + + h_pre = compute_sigmoid_gate( + projected[:, dims.pre_slice], + coeff_params.pre_scale, + coeff_params.pre_bias, + multiplier=1.0, + epsilon=pre_mapping_epsilon, + ) + h_post = compute_sigmoid_gate( + projected[:, dims.post_slice], + coeff_params.post_scale, + coeff_params.post_bias, + multiplier=2.0, + epsilon=0.0, + ) + residual = compute_residual_permutations( + projected[:, dims.res_slice], + coeff_params.res_scale, + coeff_params.res_bias, + permutations, + dims, + ) + return MhcCoeffOutputs(h_pre=h_pre, h_post=h_post, residual=residual) + + +def pre_apply(x: jax.Array, h_pre: jax.Array) -> jax.Array: + """Collapses the stream dimension before the wrapped model branch.""" + h_pre_f32 = h_pre.astype(jnp.float32) + return jnp.sum(h_pre_f32[:, :, None] * x.astype(jnp.float32), axis=1).astype(jnp.bfloat16) + + +def post_apply( + x: jax.Array, + layer_output: jax.Array, + h_post: jax.Array, + residual: jax.Array, +) -> jax.Array: + """Broadcasts the branch output and applies the residual stream mixing.""" + residual_mix = jnp.einsum( + "tkj,tkd->tjd", + residual.astype(jnp.bfloat16), + x, + preferred_element_type=jnp.float32, + ) + post_mix = h_post.astype(jnp.float32)[:, :, None] * layer_output.astype(jnp.float32)[:, None, :] + return (residual_mix + post_mix).astype(jnp.bfloat16) + + +def post_apply_bwd_pointwise( + d_output: jax.Array, + h_post: jax.Array, + residual: jax.Array, +) -> tuple[jax.Array, jax.Array]: + """Computes non-reduced pointwise gradients for post application.""" + d_output_f32 = d_output.astype(jnp.float32) + d_x = jnp.einsum( + "tkj,tjd->tkd", + residual.astype(jnp.bfloat16), + d_output_f32, + preferred_element_type=jnp.float32, + ) + d_layer_output = jnp.sum(h_post[:, :, None] * d_output_f32, axis=1) + return d_x, d_layer_output + + +def post_apply_bwd_reductions( + d_output: jax.Array, + layer_output: jax.Array, + x: jax.Array, +) -> tuple[jax.Array, jax.Array]: + """Computes feature-reduced gradients for post gating and residuals.""" + d_output_f32 = d_output.astype(jnp.float32) + d_h_post = jnp.sum(layer_output[:, None, :] * d_output_f32, axis=-1) + d_residual = jnp.einsum( + "tjd,tkd->tjk", + d_output_f32, + x, + preferred_element_type=jnp.float32, + ).transpose(0, 2, 1) + return d_h_post, d_residual + + +def validate_token_block_size(tokens: int, block_size: int, *, name: str) -> None: + """Validates a token-axis Pallas block size.""" + if block_size < 8 or block_size % 8: + raise UnsupportedInputError(f"{name} must be a positive multiple of 8; got {block_size}.") + if tokens % block_size: + raise UnsupportedInputError(f"The per-device token count ({tokens}) must be divisible by {name} ({block_size}).") + + +def validate_feature_block_size(embedding: int, block_size: int) -> None: + """Validates the feature tile used by the post-application backward.""" + if block_size < 128 or block_size % 128: + raise UnsupportedInputError(f"bwd_feature_block_size must be a positive multiple of 128; got {block_size}.") + if embedding % block_size: + raise UnsupportedInputError( + f"The embedding dimension ({embedding}) must be divisible by bwd_feature_block_size ({block_size})." + ) + + +def validate_inputs( + x: jax.Array, + block_size: int, + permutations_shape: tuple[int, ...] | None = None, + *, + block_size_name: str = "block_size", +) -> None: + """Validates the shape, dtype, and forward token block constraints.""" + if x.dtype != jnp.bfloat16: + raise UnsupportedInputError(f"The mHC Pallas kernel requires bfloat16 activations; got {x.dtype}.") + if x.ndim != 4: + raise UnsupportedInputError(f"Expected x to have shape (batch, sequence, streams, embedding); got {x.shape}.") + batch, sequence, streams, embedding = x.shape + if streams != 4 or (permutations_shape is not None and permutations_shape != (24, 4, 4)): + raise UnsupportedInputError( + "The optimized mHC Pallas kernel currently supports mHC-lite with" + f" expansion rate 4 only; got x.shape={x.shape} and permutations.shape={permutations_shape}." + ) + if embedding % 128: + raise UnsupportedInputError(f"The embedding dimension must be divisible by 128; got {embedding}.") + validate_token_block_size(batch * sequence, block_size, name=block_size_name) diff --git a/src/maxtext/kernels/mhc/mhc_kernels_bwd.py b/src/maxtext/kernels/mhc/mhc_kernels_bwd.py new file mode 100644 index 0000000000..c387b2609c --- /dev/null +++ b/src/maxtext/kernels/mhc/mhc_kernels_bwd.py @@ -0,0 +1,467 @@ +# 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 +# +# https://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. +"""Low-level Pallas backward kernels and custom VJP rules for mHC-lite.""" + +import jax +from jax.experimental import pallas as pl +from jax.experimental.pallas import tpu as pltpu +import jax.numpy as jnp +from maxtext.kernels.mhc import common + + +def _post_apply_bwd( + x: jax.Array, + layer_output: jax.Array, + h_post: jax.Array, + residual: jax.Array, + d_output: jax.Array, + config: common.MhcKernelConfig, +) -> tuple[jax.Array, jax.Array, jax.Array, jax.Array]: + """Builds the feature-tiled Pallas call for the post-branch backward pass.""" + tokens, streams, embedding = x.shape + feature_block_size = min(embedding, config.bwd_feature_block_size) + feature_blocks = embedding // feature_block_size + dims = common.MhcDims(tokens=tokens, streams=streams, embedding=embedding) + + def kernel( + x_ref, + layer_output_ref, + h_post_ref, + residual_ref, + d_output_ref, + d_x_ref, + d_layer_output_ref, + d_h_post_ref, + d_residual_ref, + ): + feature_block = pl.program_id(1) + + d_x, d_layer_output = common.post_apply_bwd_pointwise( + d_output_ref[...], + h_post_ref[...], + residual_ref[...], + ) + d_x_ref[...] = d_x.astype(d_x_ref.dtype) + d_layer_output_ref[...] = d_layer_output.astype(d_layer_output_ref.dtype) + + d_h_post, d_residual = common.post_apply_bwd_reductions( + d_output_ref[...], + layer_output_ref[...], + x_ref[...], + ) + + @pl.when(feature_block == 0) + def initialize_reductions(): + d_h_post_ref[...] = jnp.zeros_like(d_h_post_ref) + d_residual_ref[...] = jnp.zeros_like(d_residual_ref) + + d_h_post_ref[...] += d_h_post + d_residual_ref[...] += d_residual + + @pl.when(feature_block == feature_blocks - 1) + def round_d_residual(): + d_residual_ref[...] = d_residual_ref[...].astype(jnp.bfloat16).astype(d_residual_ref.dtype) + + d_x, d_layer_output, d_h_post, d_residual = pl.pallas_call( + kernel, + out_shape=( + jax.ShapeDtypeStruct((tokens, streams, embedding), x.dtype), + jax.ShapeDtypeStruct((tokens, embedding), layer_output.dtype), + jax.ShapeDtypeStruct((tokens, streams), h_post.dtype), + jax.ShapeDtypeStruct((tokens, streams, streams), residual.dtype), + ), + grid=(tokens // config.bwd_block_size, feature_blocks), + in_specs=( + common.feature_tiled_block_spec( + (tokens, streams, embedding), + config.bwd_block_size, + feature_block_size, + tiled_feature=True, + ), + common.feature_tiled_block_spec( + (tokens, embedding), + config.bwd_block_size, + feature_block_size, + tiled_feature=True, + ), + common.feature_tiled_block_spec( + (tokens, streams), + config.bwd_block_size, + feature_block_size, + tiled_feature=False, + ), + common.feature_tiled_block_spec( + (tokens, streams, streams), + config.bwd_block_size, + feature_block_size, + tiled_feature=False, + ), + common.feature_tiled_block_spec( + (tokens, streams, embedding), + config.bwd_block_size, + feature_block_size, + tiled_feature=True, + ), + ), + out_specs=( + common.feature_tiled_block_spec( + (tokens, streams, embedding), + config.bwd_block_size, + feature_block_size, + tiled_feature=True, + ), + common.feature_tiled_block_spec( + (tokens, embedding), + config.bwd_block_size, + feature_block_size, + tiled_feature=True, + ), + common.feature_tiled_block_spec( + (tokens, streams), + config.bwd_block_size, + feature_block_size, + tiled_feature=False, + ), + common.feature_tiled_block_spec( + (tokens, streams, streams), + config.bwd_block_size, + feature_block_size, + tiled_feature=False, + ), + ), + cost_estimate=dims.post_apply_bwd_cost(), + compiler_params=pltpu.CompilerParams( + vmem_limit_bytes=config.vmem_limit_bytes, + dimension_semantics=common.SEQUENTIAL_2D_DIMENSION_SEMANTICS, + ), + interpret=config.interpret, + )(x, layer_output, h_post, residual, d_output) + return d_x, d_layer_output, d_h_post, d_residual + + +def _pre_apply_bwd( + x: jax.Array, + h_pre: jax.Array, + d_layer_input: jax.Array, + d_x_acc: jax.Array, + config: common.MhcKernelConfig, +) -> tuple[jax.Array, jax.Array]: + """Builds the Pallas call for the pre-branch backward pass.""" + tokens, streams, embedding = x.shape + dims = common.MhcDims(tokens=tokens, streams=streams, embedding=embedding) + + def kernel(x_ref, h_pre_ref, d_layer_input_ref, d_x_acc_ref, d_x_ref, d_h_pre_ref): + _, vjp = jax.vjp(common.pre_apply, x_ref[...], h_pre_ref[...]) + d_x, d_h_pre = vjp(d_layer_input_ref[...]) + d_x_ref[...] = (d_x.astype(jnp.float32) + d_x_acc_ref[...].astype(jnp.float32)).astype(d_x_ref.dtype) + d_h_pre_ref[...] = d_h_pre + + d_x_acc_out, d_h_pre = pl.pallas_call( + kernel, + out_shape=( + jax.ShapeDtypeStruct((tokens, streams, embedding), x.dtype), + jax.ShapeDtypeStruct((tokens, streams), h_pre.dtype), + ), + grid=(tokens // config.bwd_block_size,), + in_specs=( + common.token_block_spec((tokens, streams, embedding), config.bwd_block_size), + common.token_block_spec((tokens, streams), config.bwd_block_size), + common.token_block_spec((tokens, embedding), config.bwd_block_size), + common.token_block_spec((tokens, streams, embedding), config.bwd_block_size), + ), + out_specs=( + common.token_block_spec((tokens, streams, embedding), config.bwd_block_size), + common.token_block_spec((tokens, streams), config.bwd_block_size), + ), + cost_estimate=dims.pre_apply_bwd_cost(), + compiler_params=pltpu.CompilerParams( + vmem_limit_bytes=config.vmem_limit_bytes, + dimension_semantics=common.PARALLEL_DIMENSION_SEMANTICS, + ), + interpret=config.interpret, + )(x, h_pre, d_layer_input, d_x_acc) + return d_x_acc_out, d_h_pre + + +def _coeff_bwd( + x: jax.Array, + coeff_params: common.MhcCoeffParams, + permutations: jax.Array, + d_outputs: common.MhcCoeffOutputs, + d_x_acc: jax.Array, + config: common.MhcKernelConfig, +) -> tuple[jax.Array, common.MhcCoeffGradients]: + """Builds the Pallas call for coefficients and parameter gradients.""" + tokens, streams, embedding = x.shape + dims = common.MhcDims( + tokens=tokens, + streams=streams, + embedding=embedding, + num_permutations=permutations.shape[0], + ) + + def kernel( + x_ref, + phi_ref, + pre_scale_ref, + pre_bias_ref, + post_scale_ref, + post_bias_ref, + res_scale_ref, + res_bias_ref, + permutations_ref, + d_h_pre_ref, + d_h_post_ref, + d_residual_ref, + d_x_acc_ref, + d_x_ref, + d_phi_ref, + d_pre_scale_ref, + d_pre_bias_ref, + d_post_scale_ref, + d_post_bias_ref, + d_res_scale_ref, + d_res_bias_ref, + ): + program_id = pl.program_id(0) + perms = permutations_ref[...] + + def mhc_coeffs_fn(x_val, params_val): + return common.mhc_coeffs( + x_val, + params_val, + perms, + rms_epsilon=config.rms_epsilon, + pre_mapping_epsilon=config.pre_mapping_epsilon, + ) + + param_refs = common.MhcCoeffParams( + phi=phi_ref, + pre_scale=pre_scale_ref, + pre_bias=pre_bias_ref, + post_scale=post_scale_ref, + post_bias=post_bias_ref, + res_scale=res_scale_ref, + res_bias=res_bias_ref, + ) + params_in = jax.tree.map(lambda ref: ref[...], param_refs) + cotangent_refs = common.MhcCoeffOutputs( + h_pre=d_h_pre_ref, + h_post=d_h_post_ref, + residual=d_residual_ref, + ) + cotangents = jax.tree.map(lambda ref: ref[...], cotangent_refs) + + _, vjp = jax.vjp(mhc_coeffs_fn, x_ref[...], params_in) + d_x, d_params = vjp(cotangents) + + d_x_ref[...] = (d_x.astype(jnp.float32) + d_x_acc_ref[...].astype(jnp.float32)).astype(d_x_ref.dtype) + + d_param_refs = common.MhcCoeffParams( + phi=d_phi_ref, + pre_scale=d_pre_scale_ref, + pre_bias=d_pre_bias_ref, + post_scale=d_post_scale_ref, + post_bias=d_post_bias_ref, + res_scale=d_res_scale_ref, + res_bias=d_res_bias_ref, + ) + + @pl.when(program_id == 0) + def initialize_reductions(): + def _zero(ref): + ref[...] = jnp.zeros_like(ref) + + jax.tree.map(_zero, d_param_refs) + + def _accumulate(ref, val): + ref[...] += val.astype(jnp.float32) + + jax.tree.map(_accumulate, d_param_refs, d_params) + + param_specs = jax.tree.map(lambda p: common.whole(p.shape), coeff_params) + param_out_shapes = jax.tree.map(lambda p: jax.ShapeDtypeStruct(p.shape, jnp.float32), coeff_params) + output_specs = jax.tree.map( + lambda out: common.token_block_spec(out.shape, config.bwd_block_size), + d_outputs, + ) + d_x, *d_param_grads = pl.pallas_call( + kernel, + out_shape=( + jax.ShapeDtypeStruct((tokens, streams, embedding), x.dtype), + *jax.tree.leaves(param_out_shapes), + ), + grid=(tokens // config.bwd_block_size,), + in_specs=( + common.token_block_spec((tokens, streams, embedding), config.bwd_block_size), + *jax.tree.leaves(param_specs), + common.whole(permutations.shape), + *jax.tree.leaves(output_specs), + common.token_block_spec((tokens, streams, embedding), config.bwd_block_size), + ), + out_specs=( + common.token_block_spec((tokens, streams, embedding), config.bwd_block_size), + *jax.tree.leaves(param_specs), + ), + cost_estimate=dims.coeff_bwd_cost(), + compiler_params=pltpu.CompilerParams( + vmem_limit_bytes=config.vmem_limit_bytes, + dimension_semantics=common.SEQUENTIAL_DIMENSION_SEMANTICS, + ), + interpret=config.interpret, + )( + x, + *jax.tree.leaves(coeff_params), + permutations, + *jax.tree.leaves(d_outputs), + d_x_acc, + ) + d_coeff_grads = common.MhcCoeffGradients(*d_param_grads) + return d_x, d_coeff_grads + + +def pre_bwd( + residuals: tuple[jax.Array, jax.Array], + cotangents: tuple[jax.Array, jax.Array, jax.Array, jax.Array], + x: jax.Array, + weights: common.MhcWeights, + permutations: jax.Array, + config: common.MhcKernelConfig, +) -> tuple[jax.Array, common.MhcWeights]: + """Computes pre-branch gradients with in-kernel input-gradient accumulation.""" + phi, h_pre = residuals + d_layer_input, d_x_acc, d_h_post, d_residual = cotangents + batch, sequence, streams, embedding = x.shape + tokens = batch * sequence + + x_flat = x.reshape(tokens, streams, embedding) + d_x_acc_flat = d_x_acc.reshape(tokens, streams, embedding) + d_h_post_flat = d_h_post.reshape(tokens, streams) + d_residual_flat = d_residual.reshape(tokens, streams, streams) + d_layer_input_flat = d_layer_input.reshape(tokens, embedding) + + d_x_acc_flat, d_h_pre = _pre_apply_bwd( + x_flat, + h_pre, + d_layer_input_flat, + d_x_acc_flat, + config=config, + ) + coeff_params = common.MhcCoeffParams( + phi=phi, + pre_scale=weights.pre_scale, + pre_bias=weights.pre_bias, + post_scale=weights.post_scale, + post_bias=weights.post_bias, + res_scale=weights.res_scale, + res_bias=weights.res_bias, + ) + d_outputs = common.MhcCoeffOutputs( + h_pre=d_h_pre, + h_post=d_h_post_flat, + residual=d_residual_flat, + ) + d_x, d_coeff_grads = _coeff_bwd( + x_flat, + coeff_params, + permutations, + d_outputs, + d_x_acc_flat, + config=config, + ) + _, phi_vjp = jax.vjp( + common.fold_norm_scale, + weights.norm_scale, + weights.pre_alpha, + weights.post_alpha, + weights.res_alpha, + ) + d_norm_scale, d_pre_alpha, d_post_alpha, d_res_alpha = phi_vjp(d_coeff_grads.phi) + d_weights = common.MhcWeights( + norm_scale=d_norm_scale, + pre_alpha=d_pre_alpha, + pre_bias=d_coeff_grads.pre_bias.astype(weights.pre_bias.dtype), + pre_scale=d_coeff_grads.pre_scale.astype(weights.pre_scale.dtype), + post_alpha=d_post_alpha, + post_bias=d_coeff_grads.post_bias.astype(weights.post_bias.dtype), + post_scale=d_coeff_grads.post_scale.astype(weights.post_scale.dtype), + res_alpha=d_res_alpha, + res_bias=d_coeff_grads.res_bias.astype(weights.res_bias.dtype), + res_scale=d_coeff_grads.res_scale.astype(weights.res_scale.dtype), + ) + return d_x.reshape(batch, sequence, streams, embedding), d_weights + + +def post_bwd( + cotangent: jax.Array, + layer_output: jax.Array, + x: jax.Array, + h_post: jax.Array, + residual: jax.Array, + config: common.MhcKernelConfig, +) -> tuple[jax.Array, jax.Array, jax.Array, jax.Array]: + """Computes post-branch gradients.""" + batch, sequence, streams, embedding = x.shape + tokens = batch * sequence + d_x, d_layer_output, d_h_post, d_residual = _post_apply_bwd( + x.reshape(tokens, streams, embedding), + layer_output.reshape(tokens, embedding), + h_post.reshape(tokens, streams), + residual.reshape(tokens, streams, streams), + cotangent.reshape(tokens, streams, embedding), + config=config, + ) + return ( + d_layer_output.reshape(batch, sequence, embedding), + d_x.reshape(batch, sequence, streams, embedding), + d_h_post.reshape(batch, sequence, streams), + d_residual.reshape(batch, sequence, streams, streams), + ) + + +def pre_op_bwd( + config: common.MhcKernelConfig, + permutations: jax.Array, + residuals: tuple[tuple[jax.Array, jax.Array], tuple[jax.Array, common.MhcWeights]], + cotangents: tuple[jax.Array, common.KernelContext], +) -> tuple[jax.Array, common.MhcWeights]: + """Custom-VJP backward rule for the low-level pre-branch entry point.""" + saved, (x, weights) = residuals + d_layer_input, (d_x, d_h_post, d_residual) = cotangents + return pre_bwd( + saved, + (d_layer_input, d_x, d_h_post, d_residual), + x, + weights, + permutations, + config=config, + ) + + +def post_op_bwd( + config: common.MhcKernelConfig, + saved: tuple[jax.Array, jax.Array, jax.Array, jax.Array], + d_output: jax.Array, +) -> tuple[jax.Array, jax.Array, jax.Array, jax.Array]: + """Custom-VJP backward rule for the low-level post-branch entry point.""" + layer_output, x, h_post, residual = saved + d_layer_output, d_x, d_h_post, d_residual = post_bwd( + d_output, + layer_output, + x, + h_post, + residual, + config=config, + ) + return d_layer_output, d_x, d_h_post, d_residual diff --git a/src/maxtext/kernels/mhc/mhc_kernels_fwd.py b/src/maxtext/kernels/mhc/mhc_kernels_fwd.py new file mode 100644 index 0000000000..f8243c4ed8 --- /dev/null +++ b/src/maxtext/kernels/mhc/mhc_kernels_fwd.py @@ -0,0 +1,328 @@ +# 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 +# +# https://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. +"""Low-level Pallas forward kernels and custom VJP functions for mHC-lite.""" + +import functools +import jax +from jax.experimental import pallas as pl +from jax.experimental.pallas import tpu as pltpu +import jax.numpy as jnp +from maxtext.kernels.mhc import common +from maxtext.kernels.mhc import mhc_kernels_bwd + + +def _coeff_fwd( + x: jax.Array, + coeff_params: common.MhcCoeffParams, + permutations: jax.Array, + config: common.MhcKernelConfig, +) -> common.MhcCoeffOutputs: + """Builds the Pallas call that computes the shared mHC coefficients.""" + tokens, streams, embedding = x.shape + dims = common.MhcDims( + tokens=tokens, + streams=streams, + embedding=embedding, + num_permutations=permutations.shape[0], + ) + + def kernel( + x_ref, + phi_ref, + pre_scale_ref, + pre_bias_ref, + post_scale_ref, + post_bias_ref, + res_scale_ref, + res_bias_ref, + permutations_ref, + h_pre_ref, + h_post_ref, + residual_ref, + ): + param_refs = common.MhcCoeffParams( + phi=phi_ref, + pre_scale=pre_scale_ref, + pre_bias=pre_bias_ref, + post_scale=post_scale_ref, + post_bias=post_bias_ref, + res_scale=res_scale_ref, + res_bias=res_bias_ref, + ) + params = jax.tree.map(lambda ref: ref[...], param_refs) + outputs = common.mhc_coeffs( + x_ref[...], + params, + permutations_ref[...], + rms_epsilon=config.rms_epsilon, + pre_mapping_epsilon=config.pre_mapping_epsilon, + ) + output_refs = common.MhcCoeffOutputs( + h_pre=h_pre_ref, + h_post=h_post_ref, + residual=residual_ref, + ) + + def _write_output(ref, val): + ref[...] = val + + jax.tree.map(_write_output, output_refs, outputs) + + param_specs = jax.tree.map(lambda p: common.whole(p.shape), coeff_params) + h_pre, h_post, residual = pl.pallas_call( + kernel, + out_shape=( + jax.ShapeDtypeStruct((tokens, streams), jnp.float32), + jax.ShapeDtypeStruct((tokens, streams), jnp.float32), + jax.ShapeDtypeStruct((tokens, streams, streams), jnp.float32), + ), + grid=(tokens // config.block_size,), + in_specs=( + common.token_block_spec((tokens, streams, embedding), config.block_size), + *jax.tree.leaves(param_specs), + common.whole(permutations.shape), + ), + out_specs=( + common.token_block_spec((tokens, streams), config.block_size), + common.token_block_spec((tokens, streams), config.block_size), + common.token_block_spec((tokens, streams, streams), config.block_size), + ), + cost_estimate=dims.coeff_fwd_cost(), + compiler_params=pltpu.CompilerParams( + vmem_limit_bytes=config.vmem_limit_bytes, + dimension_semantics=common.PARALLEL_DIMENSION_SEMANTICS, + ), + interpret=config.interpret, + )( + x, + *jax.tree.leaves(coeff_params), + permutations, + ) + return common.MhcCoeffOutputs(h_pre=h_pre, h_post=h_post, residual=residual) + + +def _pre_apply_fwd( + x: jax.Array, + h_pre: jax.Array, + config: common.MhcKernelConfig, +) -> jax.Array: + """Builds the Pallas call for the pre-branch forward pass.""" + tokens, streams, embedding = x.shape + dims = common.MhcDims(tokens=tokens, streams=streams, embedding=embedding) + + def kernel(x_ref, h_pre_ref, output_ref): + output_ref[...] = common.pre_apply(x_ref[...], h_pre_ref[...]) + + return pl.pallas_call( + kernel, + out_shape=jax.ShapeDtypeStruct((tokens, embedding), jnp.bfloat16), + grid=(tokens // config.block_size,), + in_specs=( + common.token_block_spec((tokens, streams, embedding), config.block_size), + common.token_block_spec((tokens, streams), config.block_size), + ), + out_specs=common.token_block_spec((tokens, embedding), config.block_size), + cost_estimate=dims.pre_apply_fwd_cost(), + compiler_params=pltpu.CompilerParams( + vmem_limit_bytes=config.vmem_limit_bytes, + dimension_semantics=common.PARALLEL_DIMENSION_SEMANTICS, + ), + interpret=config.interpret, + )(x, h_pre) + + +def _post_apply_fwd( + x: jax.Array, + layer_output: jax.Array, + h_post: jax.Array, + residual: jax.Array, + config: common.MhcKernelConfig, +) -> jax.Array: + """Builds the Pallas call for the post-branch forward pass.""" + tokens, streams, embedding = x.shape + dims = common.MhcDims(tokens=tokens, streams=streams, embedding=embedding) + + def kernel(x_ref, layer_output_ref, h_post_ref, residual_ref, output_ref): + output_ref[...] = common.post_apply( + x_ref[...], + layer_output_ref[...], + h_post_ref[...], + residual_ref[...], + ) + + return pl.pallas_call( + kernel, + out_shape=jax.ShapeDtypeStruct((tokens, streams, embedding), jnp.bfloat16), + grid=(tokens // config.block_size,), + in_specs=( + common.token_block_spec((tokens, streams, embedding), config.block_size), + common.token_block_spec((tokens, embedding), config.block_size), + common.token_block_spec((tokens, streams), config.block_size), + common.token_block_spec((tokens, streams, streams), config.block_size), + ), + out_specs=common.token_block_spec((tokens, streams, embedding), config.block_size), + cost_estimate=dims.post_apply_fwd_cost(), + compiler_params=pltpu.CompilerParams( + vmem_limit_bytes=config.vmem_limit_bytes, + dimension_semantics=common.PARALLEL_DIMENSION_SEMANTICS, + ), + interpret=config.interpret, + )(x, layer_output, h_post, residual) + + +def pre_fwd( + x: jax.Array, + weights: common.MhcWeights, + permutations: jax.Array, + config: common.MhcKernelConfig, +) -> tuple[tuple[jax.Array, common.KernelContext], tuple[jax.Array, jax.Array]]: + """Runs coefficient and pre-application forward kernels.""" + common.validate_inputs(x, config.block_size, permutations.shape) + batch, sequence, streams, embedding = x.shape + tokens = batch * sequence + x_flat = x.reshape(tokens, streams, embedding) + + coeff_params = weights.to_coeff_params() + outputs = _coeff_fwd(x_flat, coeff_params, permutations, config) + layer_input = _pre_apply_fwd(x_flat, outputs.h_pre, config) + + context: common.KernelContext = ( + x, + outputs.h_post.reshape(batch, sequence, streams), + outputs.residual.reshape(batch, sequence, streams, streams), + ) + primals_out = (layer_input.reshape(batch, sequence, embedding), context) + saved_residuals = (coeff_params.phi, outputs.h_pre) + return primals_out, saved_residuals + + +def post_fwd( + layer_output: jax.Array, + x: jax.Array, + h_post: jax.Array, + residual: jax.Array, + config: common.MhcKernelConfig, +) -> jax.Array: + """Runs the fused post-gate and residual-mixing forward kernel.""" + common.validate_inputs(x, config.block_size) + batch, sequence, streams, embedding = x.shape + tokens = batch * sequence + output = _post_apply_fwd( + x.reshape(tokens, streams, embedding), + layer_output.reshape(tokens, embedding), + h_post.reshape(tokens, streams), + residual.reshape(tokens, streams, streams), + config, + ) + return output.reshape(batch, sequence, streams, embedding) + + +@functools.partial(jax.custom_vjp, nondiff_argnums=(0, 2)) +def _pre_op( + config: common.MhcKernelConfig, + x: jax.Array, + permutations: jax.Array, + weights: common.MhcWeights, +) -> tuple[jax.Array, common.KernelContext]: + """Differentiable pre-branch mHC operation.""" + (layer_input, context), _ = pre_fwd(x, weights, permutations, config) + return layer_input, context + + +def _pre_op_fwd( + config: common.MhcKernelConfig, + x: jax.Array, + permutations: jax.Array, + weights: common.MhcWeights, +): + """Custom-VJP forward rule for the pre-branch operation.""" + primals_out, saved = pre_fwd(x, weights, permutations, config) + return primals_out, (saved, (x, weights)) + + +_pre_op.defvjp(_pre_op_fwd, mhc_kernels_bwd.pre_op_bwd) + + +@functools.partial(jax.custom_vjp, nondiff_argnums=(0,)) +def _post_op( + config: common.MhcKernelConfig, + layer_output: jax.Array, + x: jax.Array, + h_post: jax.Array, + residual: jax.Array, +) -> jax.Array: + """Differentiable post-branch mHC operation.""" + return post_fwd(layer_output, x, h_post, residual, config) + + +def _post_op_fwd( + config: common.MhcKernelConfig, + layer_output: jax.Array, + x: jax.Array, + h_post: jax.Array, + residual: jax.Array, +): + """Custom-VJP forward rule for the post-branch operation.""" + output = post_fwd(layer_output, x, h_post, residual, config) + return output, (layer_output, x, h_post, residual) + + +_post_op.defvjp(_post_op_fwd, mhc_kernels_bwd.post_op_bwd) + + +def pre( + x: jax.Array, + weights: common.MhcWeights, + permutations: jax.Array, + config: common.MhcKernelConfig = common.MhcKernelConfig(), +) -> tuple[jax.Array, common.KernelContext]: + """Runs the coefficient and pre-application kernels. + + Args: + x: Input streams with shape `[batch, sequence, streams, embedding]`. + weights: Structured weights container. + permutations: Permutation matrices with shape `[streams!, streams, + streams]`. + config: Kernel tuning and compiler configuration. + + Returns: + A pair containing the branch input and opaque kernel context. + """ + common.validate_inputs(x, config.block_size, permutations.shape) + common.validate_token_block_size(x.shape[0] * x.shape[1], config.bwd_block_size, name="bwd_block_size") + return _pre_op(config, x, permutations, weights) + + +def post( + layer_output: jax.Array, + context: common.KernelContext, + config: common.MhcKernelConfig = common.MhcKernelConfig(), +) -> jax.Array: + """Runs the fused post-gate and residual-mixing kernel. + + Args: + layer_output: Wrapped branch output with shape `[batch, sequence, + embedding]`. + context: Opaque context returned by `pre`. + config: Kernel tuning and compiler configuration. + + Returns: + Mixed output streams with shape `[batch, sequence, streams, embedding]`. + """ + x, h_post, residual = context + feature_block_size = min(x.shape[-1], config.bwd_feature_block_size) + common.validate_inputs(x, config.block_size) + common.validate_token_block_size(x.shape[0] * x.shape[1], config.bwd_block_size, name="bwd_block_size") + common.validate_feature_block_size(x.shape[-1], feature_block_size) + return _post_op(config, layer_output, x, h_post, residual) diff --git a/src/maxtext/layers/mhc.py b/src/maxtext/layers/mhc.py index d41cf5967e..817def4dc0 100644 --- a/src/maxtext/layers/mhc.py +++ b/src/maxtext/layers/mhc.py @@ -24,8 +24,9 @@ from jax.sharding import Mesh from maxtext.common.common_types import Array, Config from maxtext.common.common_types import HyperConnectionType -from maxtext.layers.initializers import default_bias_init, default_scalar_init, nd_dense_init, variable_to_logically_partitioned +from maxtext.kernels.mhc import api as mhc_kernel from maxtext.layers import nnx_wrappers +from maxtext.layers.initializers import default_bias_init, default_scalar_init, nd_dense_init, variable_to_logically_partitioned from maxtext.layers.normalizations import RMSNorm @@ -103,6 +104,9 @@ def __init__( self.weight_dtype = self.config.weight_dtype self.matmul_precision = jax.lax.Precision(self.config.matmul_precision) + if getattr(self.config, "use_mhc_pallas_kernel", False) and not self.config.enable_mhc_lite: + raise ValueError("use_mhc_pallas_kernel=True requires enable_mhc_lite=True.") + # Norm layer self.mhc_norm = RMSNorm( num_features=self.k * self.dim, @@ -188,6 +192,21 @@ def __init__( out_sharding=(None,), ) + def _get_mhc_weights(self) -> mhc_kernel.MhcWeights: + """Collects layer parameters into a structured MhcWeights PyTree.""" + return mhc_kernel.MhcWeights( + norm_scale=jnp.asarray(self.mhc_norm.scale[...], self.dtype), + pre_alpha=jnp.asarray(self.pre_alpha[...], self.dtype), + pre_bias=jnp.asarray(self.pre_beta[...], self.dtype), + pre_scale=jnp.asarray(self.pre_alpha_scale[...], self.dtype), + post_alpha=jnp.asarray(self.post_alpha[...], self.dtype), + post_bias=jnp.asarray(self.post_beta[...], self.dtype), + post_scale=jnp.asarray(self.post_alpha_scale[...], self.dtype), + res_alpha=jnp.asarray(self.res_alpha[...], self.dtype), + res_bias=jnp.asarray(self.res_beta[...], self.dtype), + res_scale=jnp.asarray(self.res_alpha_scale[...], self.dtype), + ) + def res_mapping(self, h_res: Array): """Helper function for residual mapping after matmul.""" # In MaxText, we match weight precision to activations before Matmul @@ -246,36 +265,54 @@ def __call__( # x shape: [batch, seq, expansion_rate, emb] b, s, k, d = x.shape - with jax.named_scope("mhc_norm"): - # 1. Flatten the tensor, and RMS normalization - norm_x = self.mhc_norm(jnp.reshape(x, (b, s, k * d))) - - # Fused Projections - pre_alpha = jnp.asarray(self.pre_alpha[...], self.dtype) - post_alpha = jnp.asarray(self.post_alpha[...], self.dtype) - res_alpha = jnp.asarray(self.res_alpha[...], self.dtype) - - alpha_concat = jnp.concatenate([pre_alpha, post_alpha, res_alpha], axis=-1) - - # MatMul on normalized input - h_concat = jnp.einsum("bsm,mn -> bsn", norm_x, alpha_concat, precision=self.matmul_precision) - - h_pre = h_concat[..., : self.k] - h_post = h_concat[..., self.k : 2 * self.k] - h_res = h_concat[..., 2 * self.k :] - - # 2. Pre mapping - pre_mapping = self.mapping( - h_pre, - self.pre_alpha_scale[...], - self.pre_beta[...], - 1.0, - eps=1e-6, - ) - # Moving away from einsum seems to allow XLA to perform better fusions - # https://github.com/AI-Hypercomputer/maxtext/pull/4664#discussion_r3677899970 - # bskd, bsk -> bsd - layer_input = jnp.sum(x * jnp.expand_dims(pre_mapping, axis=3), axis=2) + h_post = None + h_res = None + context = None + use_kernel = self.config.enable_mhc_lite and getattr(self.config, "use_mhc_pallas_kernel", False) + if use_kernel: + fwd_block_size = getattr(self.config, "mhc_pallas_kernel_fwd_block_size", 256) + bwd_block_size = getattr(self.config, "mhc_pallas_kernel_bwd_block_size", 128) + kernel_config = mhc_kernel.MhcKernelConfig( + block_size=fwd_block_size, + bwd_block_size=bwd_block_size, + rms_epsilon=self.config.normalization_layer_epsilon, + ) + weights = self._get_mhc_weights() + layer_input, context = mhc_kernel.pre( + x, + weights, + jnp.asarray(self.permutation_matrices, self.dtype), + config=kernel_config, + ) + else: + with jax.named_scope("mhc_norm"): + # 1. Flatten the tensor, and RMS normalization + norm_x = self.mhc_norm(jnp.reshape(x, (b, s, k * d))) + + # Fused Projections + pre_alpha = jnp.asarray(self.pre_alpha[...], self.dtype) + post_alpha = jnp.asarray(self.post_alpha[...], self.dtype) + res_alpha = jnp.asarray(self.res_alpha[...], self.dtype) + + alpha_concat = jnp.concatenate([pre_alpha, post_alpha, res_alpha], axis=-1) + + # MatMul on normalized input + h_concat = jnp.einsum("bsm,mn -> bsn", norm_x, alpha_concat, precision=self.matmul_precision) + h_pre = h_concat[..., : self.k] + h_post = h_concat[..., self.k : 2 * self.k] + h_res = h_concat[..., 2 * self.k :] + + # 2. Pre mapping + pre_mapping = self.mapping( + h_pre, + self.pre_alpha_scale[...], + self.pre_beta[...], + 1.0, + eps=1e-6, + ) + # Moving away from einsum seems to allow XLA to perform better fusions + # bskd, bsk -> bsd + layer_input = jnp.sum(x * jnp.expand_dims(pre_mapping, axis=3), axis=2) # 3. Pre-norm layer_input = norm_fn(layer_input) @@ -293,6 +330,20 @@ def __call__( else: raise ValueError(f"Unsupported type: {mhc_type}") + if use_kernel: + fwd_block_size = getattr(self.config, "mhc_pallas_kernel_fwd_block_size", 256) + bwd_block_size = getattr(self.config, "mhc_pallas_kernel_bwd_block_size", 128) + kernel_config = mhc_kernel.MhcKernelConfig( + block_size=fwd_block_size, + bwd_block_size=bwd_block_size, + ) + output = mhc_kernel.post( + layer_out, + context, + config=kernel_config, + ) + return output, metadata + # 5. Post mapping post_mapping = self.mapping( h_post, diff --git a/tests/unit/mhc_test.py b/tests/unit/mhc_test.py index fa2de233cf..921004ed78 100644 --- a/tests/unit/mhc_test.py +++ b/tests/unit/mhc_test.py @@ -14,7 +14,12 @@ """Test for DeepSeek Manifold-Constrained Hyper Connections (mHC).""" +import dataclasses +import itertools +import math import unittest +from unittest import mock +from absl.testing import absltest from absl.testing import parameterized from flax import nnx from flax.linen import partitioning as nn_partitioning @@ -23,13 +28,15 @@ from jax.sharding import Mesh from maxtext.common.common_types import HyperConnectionType from maxtext.configs import pyconfig +from maxtext.kernels import mhc as mhc_kernel +from maxtext.kernels.mhc import common as mhc_kernel_common from maxtext.layers import attention_mla, linears, mhc, moe from maxtext.layers.initializers import nd_dense_init from maxtext.layers.normalizations import RMSNorm from maxtext.utils import maxtext_utils +from tests.utils.test_helpers import get_test_config_path import numpy as np import pytest -from tests.utils.test_helpers import get_test_config_path class TestExpandReduce(unittest.TestCase): @@ -88,28 +95,57 @@ def test_doubly_stochastic_property(self): class TestMHC(parameterized.TestCase): """Test for MHC module""" - def _setup_mhc(self, rate, enable_mhc_lite=False): + def _setup_mhc( + self, + rate, + enable_mhc_lite=False, + use_mhc_pallas_kernel=False, + mhc_pallas_kernel_fwd_block_size=None, + mhc_pallas_kernel_bwd_block_size=None, + dim=16, + sequence_length=7, + per_device_batch_size=None, + dtype=None, + ): """Sets up the common configurations and modules for MHC testing.""" - self.dim = 16 + self.dim = dim + if per_device_batch_size is None: + per_device_batch_size = jax.device_count() + kwargs = { + "run_name": f"test_mhc_k{rate}", + "enable_checkpointing": False, + "model_name": "deepseek-custom", + "per_device_batch_size": per_device_batch_size, + "max_target_length": sequence_length, + "max_prefill_predict_length": sequence_length, + "attention": "dot_product", + "attention_type": "mla", + "routed_bias": True, + "routed_bias_update_rate": 0.01, + "load_balance_loss_weight": 0.02, + # override + "override_model_config": True, + "base_emb_dim": self.dim, + "mhc_expansion_rate": rate, + "enable_mhc_lite": enable_mhc_lite, + "use_mhc_pallas_kernel": use_mhc_pallas_kernel, + "decoder_block": "deepseek", + "num_experts": 4, + "num_experts_per_tok": 2, + "base_moe_mlp_dim": self.dim * 4, + "base_mlp_dim": self.dim * 4, + "engram_layers": [], + } + if mhc_pallas_kernel_fwd_block_size is not None: + kwargs["mhc_pallas_kernel_fwd_block_size"] = mhc_pallas_kernel_fwd_block_size + if mhc_pallas_kernel_bwd_block_size is not None: + kwargs["mhc_pallas_kernel_bwd_block_size"] = mhc_pallas_kernel_bwd_block_size + if dtype is not None: + kwargs["dtype"] = dtype + kwargs["weight_dtype"] = dtype self.config = pyconfig.initialize( [None, get_test_config_path()], - run_name=f"test_mhc_k{rate}", - enable_checkpointing=False, - model_name="deepseek-custom", - per_device_batch_size=jax.device_count(), - max_target_length=7, - max_prefill_predict_length=7, - attention="dot_product", - routed_bias_update_rate=0.01, - load_balance_loss_weight=0.02, - # override - override_model_config=True, - base_emb_dim=self.dim, - mhc_expansion_rate=rate, - enable_mhc_lite=enable_mhc_lite, - num_experts=4, - num_experts_per_tok=2, - engram_layers=[], + **kwargs, ) devices_array = maxtext_utils.create_device_mesh(self.config) self.mesh = Mesh(devices_array, self.config.mesh_axes) @@ -123,6 +159,7 @@ def _setup_mhc(self, rate, enable_mhc_lite=False): self.config.mhc_expansion_rate, self.config.emb_dim, ), + dtype=self.config.dtype, ) self.pre_norm = RMSNorm( @@ -227,7 +264,14 @@ def test_attention_layer_output_shape(self, rate): ) b, s, k, d = self.x.shape - output, metadata = module(self.pre_norm, layer, x=self.x, mhc_type=HyperConnectionType.ATTENTION) + positions = jnp.broadcast_to(jnp.arange(s)[None, :], (b, s)) + output, metadata = module( + self.pre_norm, + layer, + x=self.x, + mhc_type=HyperConnectionType.ATTENTION, + inputs_positions=positions, + ) self.assertDictEqual(metadata, {}) self.assertEqual(output.shape, (b, s, k, d)) @@ -299,6 +343,8 @@ def test_feature_flag_gates_lite(self): max_target_length=7, max_prefill_predict_length=7, attention="dot_product", + attention_type="mla", + routed_bias=True, routed_bias_update_rate=0.01, load_balance_loss_weight=0.02, # override @@ -306,8 +352,11 @@ def test_feature_flag_gates_lite(self): base_emb_dim=self.dim, mhc_expansion_rate=4, enable_mhc_lite=False, + decoder_block="deepseek", num_experts=4, num_experts_per_tok=2, + base_moe_mlp_dim=self.dim * 4, + base_mlp_dim=self.dim * 4, engram_layers=[], ) devices_array = maxtext_utils.create_device_mesh(self.config) @@ -324,6 +373,528 @@ def test_feature_flag_gates_lite(self): # Permutation matrices shouldn't be defined self.assertFalse(hasattr(module, "permutation_matrices")) + @parameterized.named_parameters( + ("KernelEnabled", True), + ("KernelDisabled", False), + ) + def test_use_mhc_pallas_kernel_dispatch(self, use_mhc_pallas_kernel): + """Verify that use_mhc_pallas_kernel flag controls kernel vs pure JAX dispatch.""" + self._setup_mhc( + 4, + enable_mhc_lite=True, + use_mhc_pallas_kernel=use_mhc_pallas_kernel, + dim=128, + sequence_length=256, + per_device_batch_size=1, + dtype="bfloat16", + ) + with nn_partitioning.axis_rules(self.config.logical_axis_rules): + module = mhc.ManifoldConstrainedHyperConnections(self.config, self.dim, self.mesh, self.rngs) + layer = linears.MlpBlock( + config=self.config, + mesh=self.mesh, + in_features=self.config.emb_dim, + intermediate_dim=self.config.moe_mlp_dim, + activations=self.config.mlp_activations, + intermediate_dropout_rate=self.config.dropout_rate, + dtype=self.config.dtype, + weight_dtype=self.config.weight_dtype, + model_mode=self.config.model_call_mode, + rngs=self.rngs, + ) + + real_pre = mhc.mhc_kernel.pre + real_post = mhc.mhc_kernel.post + + def fake_pre(*args, **kwargs): + config = kwargs.get("config", mhc.mhc_kernel.MhcKernelConfig()) + config = dataclasses.replace(config, interpret=True) + kwargs["config"] = config + return real_pre(*args, **kwargs) + + def fake_post(*args, **kwargs): + config = kwargs.get("config", mhc.mhc_kernel.MhcKernelConfig()) + config = dataclasses.replace(config, interpret=True) + kwargs["config"] = config + return real_post(*args, **kwargs) + + with ( + mock.patch.object(mhc.mhc_kernel, "pre", side_effect=fake_pre) as mock_pre, + mock.patch.object(mhc.mhc_kernel, "post", side_effect=fake_post) as mock_post, + ): + output, _ = module( + self.pre_norm, + layer, + x=self.x, + mhc_type=HyperConnectionType.MLP_DENSE, + ) + + if use_mhc_pallas_kernel: + mock_pre.assert_called_once() + mock_post.assert_called_once() + _, kwargs_pre = mock_pre.call_args + config_pre = kwargs_pre.get("config") + self.assertIsNotNone(config_pre) + self.assertEqual(config_pre.block_size, 256) + self.assertEqual(config_pre.bwd_block_size, 128) + + _, kwargs_post = mock_post.call_args + config_post = kwargs_post.get("config") + self.assertIsNotNone(config_post) + self.assertEqual(config_post.block_size, 256) + self.assertEqual(config_post.bwd_block_size, 128) + else: + mock_pre.assert_not_called() + mock_post.assert_not_called() + + self.assertEqual(output.shape, self.x.shape) + + def test_use_mhc_pallas_kernel_custom_block_size(self): + """Verify that custom block sizes are passed to the kernel.""" + self._setup_mhc( + 4, + enable_mhc_lite=True, + use_mhc_pallas_kernel=True, + mhc_pallas_kernel_fwd_block_size=128, + mhc_pallas_kernel_bwd_block_size=64, + dim=128, + sequence_length=128, + per_device_batch_size=1, + dtype="bfloat16", + ) + with nn_partitioning.axis_rules(self.config.logical_axis_rules): + module = mhc.ManifoldConstrainedHyperConnections(self.config, self.dim, self.mesh, self.rngs) + layer = linears.MlpBlock( + config=self.config, + mesh=self.mesh, + in_features=self.config.emb_dim, + intermediate_dim=self.config.moe_mlp_dim, + activations=self.config.mlp_activations, + intermediate_dropout_rate=self.config.dropout_rate, + dtype=self.config.dtype, + weight_dtype=self.config.weight_dtype, + model_mode=self.config.model_call_mode, + rngs=self.rngs, + ) + + real_pre = mhc.mhc_kernel.pre + real_post = mhc.mhc_kernel.post + + def fake_pre(*args, **kwargs): + config = kwargs.get("config", mhc.mhc_kernel.MhcKernelConfig()) + config = dataclasses.replace(config, interpret=True) + kwargs["config"] = config + return real_pre(*args, **kwargs) + + def fake_post(*args, **kwargs): + config = kwargs.get("config", mhc.mhc_kernel.MhcKernelConfig()) + config = dataclasses.replace(config, interpret=True) + kwargs["config"] = config + return real_post(*args, **kwargs) + + with ( + mock.patch.object(mhc.mhc_kernel, "pre", side_effect=fake_pre) as mock_pre, + mock.patch.object(mhc.mhc_kernel, "post", side_effect=fake_post) as mock_post, + ): + output, _ = module( + self.pre_norm, + layer, + x=self.x, + mhc_type=HyperConnectionType.MLP_DENSE, + ) + + mock_pre.assert_called_once() + mock_post.assert_called_once() + _, kwargs_pre = mock_pre.call_args + config_pre = kwargs_pre.get("config") + self.assertIsNotNone(config_pre) + self.assertEqual(config_pre.block_size, 128) + self.assertEqual(config_pre.bwd_block_size, 64) + + _, kwargs_post = mock_post.call_args + config_post = kwargs_post.get("config") + self.assertIsNotNone(config_post) + self.assertEqual(config_post.block_size, 128) + self.assertEqual(config_post.bwd_block_size, 64) + + self.assertEqual(output.shape, self.x.shape) + + def test_use_mhc_pallas_kernel_requires_enable_mhc_lite(self): + """Verify that use_mhc_pallas_kernel=True requires enable_mhc_lite=True.""" + # Test via pyconfig initialization + with self.assertRaises(ValueError): + self._setup_mhc( + 4, + enable_mhc_lite=False, + use_mhc_pallas_kernel=True, + ) + + # Test via direct layer initialization with mock config + self._setup_mhc(4) + mock_config = mock.MagicMock() + mock_config.use_mhc_pallas_kernel = True + mock_config.enable_mhc_lite = False + mock_config.dtype = jnp.bfloat16 + mock_config.weight_dtype = jnp.bfloat16 + mock_config.matmul_precision = "default" + mock_config.mhc_expansion_rate = 4 + with self.assertRaises(ValueError): + mhc.ManifoldConstrainedHyperConnections(mock_config, 16, self.mesh, self.rngs) + + def test_layer_vjp_parity_kernel_vs_baseline(self): + """Verify that ManifoldConstrainedHyperConnections with kernel matches baseline under VJP.""" + self._setup_mhc( + 4, + enable_mhc_lite=True, + use_mhc_pallas_kernel=False, + dim=128, + sequence_length=128, + per_device_batch_size=1, + dtype="bfloat16", + ) + with nn_partitioning.axis_rules(self.config.logical_axis_rules): + module_baseline = mhc.ManifoldConstrainedHyperConnections(self.config, self.dim, self.mesh, self.rngs) + + def layer_fn(inputs): + return inputs * 2.0 + + def forward_baseline(x): + out, _ = module_baseline(self.pre_norm, layer_fn, x=x, mhc_type=HyperConnectionType.MLP_DENSE) + return out + + (out_base, vjp_base) = jax.vjp(forward_baseline, self.x) + cotangent = jax.random.normal(jax.random.PRNGKey(123), self.x.shape, dtype=self.x.dtype) + (dx_base,) = vjp_base(cotangent) + + self._setup_mhc( + 4, + enable_mhc_lite=True, + use_mhc_pallas_kernel=True, + mhc_pallas_kernel_fwd_block_size=128, + mhc_pallas_kernel_bwd_block_size=64, + dim=128, + sequence_length=128, + per_device_batch_size=1, + dtype="bfloat16", + ) + module_kernel = mhc.ManifoldConstrainedHyperConnections(self.config, self.dim, self.mesh, self.rngs) + + real_pre = mhc.mhc_kernel.pre + real_post = mhc.mhc_kernel.post + + def fake_pre(*args, **kwargs): + cfg = kwargs.get("config", mhc.mhc_kernel.MhcKernelConfig()) + cfg = dataclasses.replace(cfg, interpret=True) + kwargs["config"] = cfg + return real_pre(*args, **kwargs) + + def fake_post(*args, **kwargs): + cfg = kwargs.get("config", mhc.mhc_kernel.MhcKernelConfig()) + cfg = dataclasses.replace(cfg, interpret=True) + kwargs["config"] = cfg + return real_post(*args, **kwargs) + + with ( + mock.patch.object(mhc.mhc_kernel, "pre", side_effect=fake_pre), + mock.patch.object(mhc.mhc_kernel, "post", side_effect=fake_post), + ): + + def forward_kernel(x): + out, _ = module_kernel(self.pre_norm, layer_fn, x=x, mhc_type=HyperConnectionType.MLP_DENSE) + return out + + (out_kern, vjp_kern) = jax.vjp(forward_kernel, self.x) + (dx_kern,) = vjp_kern(cotangent) + + np.testing.assert_allclose(out_kern, out_base, rtol=5e-2, atol=5e-2) + np.testing.assert_allclose( + np.asarray(dx_kern, np.float32), + np.asarray(dx_base, np.float32), + rtol=5e-2, + atol=5e-2, + ) + + +def _get_permutation_matrices(k: int) -> jax.Array: + """Generates all permutation matrices for k streams.""" + perms = jnp.array(list(itertools.permutations(range(k)))) + return jnp.eye(k, dtype=jnp.float32)[perms] + + +def _make_kernel_inputs(batch=2, sequence=64, streams=4, embedding=256, seed=0): + """Generates synthetic inputs and parameters for testing mHC kernels.""" + key = jax.random.PRNGKey(seed) + keys = jax.random.split(key, 12) + x = jax.random.normal(keys[0], (batch, sequence, streams, embedding), dtype=jnp.bfloat16) + norm_scale = jax.random.normal(keys[1], (streams * embedding,), dtype=jnp.bfloat16) + pre_alpha = jax.random.normal(keys[2], (streams * embedding, streams), dtype=jnp.bfloat16) * 0.1 + pre_bias = jax.random.normal(keys[3], (streams,), dtype=jnp.bfloat16) * 0.1 + pre_scale = jnp.array([1.0], dtype=jnp.bfloat16) + post_alpha = jax.random.normal(keys[5], (streams * embedding, streams), dtype=jnp.bfloat16) * 0.1 + post_bias = jax.random.normal(keys[6], (streams,), dtype=jnp.bfloat16) * 0.1 + post_scale = jnp.array([1.0], dtype=jnp.bfloat16) + num_perms = math.factorial(streams) + res_alpha = jax.random.normal(keys[8], (streams * embedding, num_perms), dtype=jnp.bfloat16) * 0.1 + res_bias = jax.random.normal(keys[9], (num_perms,), dtype=jnp.bfloat16) * 0.1 + res_scale = jnp.array([1.0], dtype=jnp.bfloat16) + permutations = _get_permutation_matrices(streams) + cotangent = jax.random.normal(keys[11], (batch, sequence, streams, embedding), dtype=jnp.bfloat16) + weights = mhc_kernel_common.MhcWeights( + norm_scale=norm_scale, + pre_alpha=pre_alpha, + pre_bias=pre_bias, + pre_scale=pre_scale, + post_alpha=post_alpha, + post_bias=post_bias, + post_scale=post_scale, + res_alpha=res_alpha, + res_bias=res_bias, + res_scale=res_scale, + ) + return x, weights, permutations, cotangent + + +def _run_pipeline_reference(x, weights: mhc_kernel_common.MhcWeights, permutations): + """Runs the native JAX/XLA reference implementation of the mHC pipeline.""" + batch, sequence, streams, embedding = x.shape + tokens = batch * sequence + flattened_size = streams * embedding + permutation_count = permutations.shape[0] + x_flat = x.reshape(tokens, streams, embedding) + flattened_f32 = x_flat.reshape(tokens, flattened_size).astype(jnp.float32) + normalized = ( + flattened_f32 + * jax.lax.rsqrt(jnp.mean(flattened_f32 * flattened_f32, axis=-1, keepdims=True) + 1e-5) + * weights.norm_scale.astype(jnp.float32) + ).astype(x.dtype) + + h_pre = ( + jax.nn.sigmoid( + weights.pre_scale.astype(jnp.float32) + * jnp.dot( + normalized, + weights.pre_alpha, + preferred_element_type=jnp.float32, + ) + + weights.pre_bias.astype(jnp.float32) + ) + + 1e-6 + ) + layer_input = jnp.sum( + h_pre[:, :, None] * x_flat.astype(jnp.float32), + axis=1, + ).astype(x.dtype) + + h_post = 2.0 * jax.nn.sigmoid( + weights.post_scale.astype(jnp.float32) + * jnp.dot( + normalized, + weights.post_alpha, + preferred_element_type=jnp.float32, + ) + + weights.post_bias.astype(jnp.float32) + ) + weights_res = jax.nn.softmax( + weights.res_scale.astype(jnp.float32) + * jnp.dot( + normalized, + weights.res_alpha, + preferred_element_type=jnp.float32, + ) + + weights.res_bias.astype(jnp.float32), + axis=-1, + ) + residual = jnp.dot( + weights_res, + permutations.reshape(permutation_count, streams * streams).astype(jnp.float32), + ).reshape(tokens, streams, streams) + + residual_mix = jnp.einsum( + "tkj,tkd->tjd", + residual.astype(x.dtype), + x_flat, + preferred_element_type=jnp.float32, + ) + post_mix = h_post.astype(jnp.float32)[:, :, None] * layer_input.astype(jnp.float32)[:, None, :] + return (residual_mix + post_mix).astype(x.dtype).reshape(x.shape) + + +def _run_pipeline_api( + x, + weights: mhc_kernel_common.MhcWeights, + permutations, + implementation=None, + config: mhc_kernel_common.MhcKernelConfig | None = None, + interpret=True, +): + """Runs the mHC Pallas kernel API pipeline.""" + if config is None: + config = mhc_kernel.MhcKernelConfig(interpret=interpret) + else: + config = dataclasses.replace(config, interpret=interpret) + layer_input, context = mhc_kernel.pre( + x, + weights, + permutations, + config=config, + implementation=implementation, + ) + return mhc_kernel.post(layer_input, context, config=config) + + +class TestMhcKernelsFwd(parameterized.TestCase): + """Unit tests for MaxText mHC-lite Pallas forward kernel.""" + + def test_doubly_stochastic(self): + x, weights, permutations, _ = _make_kernel_inputs(batch=1, sequence=128, streams=4, embedding=128) + config = mhc_kernel.MhcKernelConfig(interpret=True) + _, context = mhc_kernel.pre( + x, + weights, + permutations, + config=config, + ) + row_sums = jnp.sum(context.residual, axis=-1) + col_sums = jnp.sum(context.residual, axis=-2) + np.testing.assert_allclose(row_sums, np.ones_like(row_sums), rtol=1e-3, atol=1e-3) + np.testing.assert_allclose(col_sums, np.ones_like(col_sums), rtol=1e-3, atol=1e-3) + + @parameterized.named_parameters( + ("mosaic", "mosaic"), + ("auto", None), + ) + def test_forward_parity(self, implementation): + x, weights, permutations, _ = _make_kernel_inputs(batch=2, sequence=64, streams=4, embedding=256) + expected = _run_pipeline_reference(x, weights, permutations) + actual = _run_pipeline_api( + x, + weights, + permutations, + implementation=implementation, + interpret=True, + ) + np.testing.assert_allclose(actual, expected, rtol=5e-2, atol=5e-2) + + def test_unsupported_shape_raises_error(self): + x, weights, permutations, _ = _make_kernel_inputs(batch=1, sequence=16, streams=2, embedding=128) + config = mhc_kernel.MhcKernelConfig(interpret=True) + with self.assertRaises(mhc_kernel_common.UnsupportedInputError): + mhc_kernel.pre( + x, + weights, + permutations, + config=config, + ) + + +class TestMhcKernelsBwd(parameterized.TestCase): + """Unit tests for MaxText mHC-lite Pallas backward kernel.""" + + def test_forward_and_backward_vjp_parity(self): + x, weights, permutations, cotangent = _make_kernel_inputs(batch=2, sequence=64, streams=4, embedding=256) + expected_out, expected_vjp_fn = jax.vjp( + lambda x_, w_: _run_pipeline_reference(x_, w_, permutations), + x, + weights, + ) + expected_dx, expected_dw = expected_vjp_fn(cotangent) + + actual_out, actual_vjp_fn = jax.vjp( + lambda x_, w_: _run_pipeline_api(x_, w_, permutations, implementation=None, interpret=True), + x, + weights, + ) + actual_dx, actual_dw = actual_vjp_fn(cotangent) + + np.testing.assert_allclose(actual_out, expected_out, rtol=5e-2, atol=5e-2) + + actual_grads = (actual_dx,) + tuple(jax.tree_util.tree_leaves(actual_dw)) + expected_grads = (expected_dx,) + tuple(jax.tree_util.tree_leaves(expected_dw)) + self.assertEqual(len(actual_grads), len(expected_grads)) + for i, (actual_g, expected_g) in enumerate(zip(actual_grads, expected_grads)): + tol = 0.05 if actual_g.size == 1 else 0.02 + scale = max(float(np.max(np.abs(np.asarray(expected_g, np.float32)))), 1e-7) + np.testing.assert_allclose( + np.asarray(actual_g, np.float32), + np.asarray(expected_g, np.float32), + rtol=0.0, + atol=tol * scale, + err_msg=f"Gradient leaf {i} mismatch", + ) + + def test_forward_and_backward_vjp_parity_feature_tiled(self): + x, weights, permutations, cotangent = _make_kernel_inputs(batch=2, sequence=64, streams=4, embedding=256) + config = mhc_kernel.MhcKernelConfig(bwd_feature_block_size=128) + expected_out, expected_vjp_fn = jax.vjp( + lambda x_, w_: _run_pipeline_reference(x_, w_, permutations), + x, + weights, + ) + expected_dx, expected_dw = expected_vjp_fn(cotangent) + + actual_out, actual_vjp_fn = jax.vjp( + lambda x_, w_: _run_pipeline_api(x_, w_, permutations, config=config, interpret=True), + x, + weights, + ) + actual_dx, actual_dw = actual_vjp_fn(cotangent) + + np.testing.assert_allclose(actual_out, expected_out, rtol=5e-2, atol=5e-2) + + actual_grads = (actual_dx,) + tuple(jax.tree_util.tree_leaves(actual_dw)) + expected_grads = (expected_dx,) + tuple(jax.tree_util.tree_leaves(expected_dw)) + self.assertEqual(len(actual_grads), len(expected_grads)) + for i, (actual_g, expected_g) in enumerate(zip(actual_grads, expected_grads)): + tol = 0.05 if actual_g.size == 1 else 0.02 + scale = max(float(np.max(np.abs(np.asarray(expected_g, np.float32)))), 1e-7) + np.testing.assert_allclose( + np.asarray(actual_g, np.float32), + np.asarray(expected_g, np.float32), + rtol=0.0, + atol=tol * scale, + err_msg=f"Feature-tiled gradient leaf {i} mismatch", + ) + + +class TestMhcCostEstimates(unittest.TestCase): + """Unit tests for analytical CostEstimate computations on MhcDims.""" + + def setUp(self): + self.dims = mhc_kernel_common.MhcDims(tokens=256, streams=4, embedding=512, num_permutations=24) + + def test_dims_properties(self): + self.assertEqual(self.dims.flattened_size, 4 * 512) + self.assertEqual(self.dims.phi_cols, 2 * 4 + 24) + self.assertEqual(self.dims.pre_slice, slice(0, 4)) + self.assertEqual(self.dims.post_slice, slice(4, 8)) + self.assertEqual(self.dims.res_slice, slice(8, 32)) + + def test_cost_estimates_non_zero_and_valid(self): + costs = [ + ("coeff_fwd", self.dims.coeff_fwd_cost()), + ("pre_apply_fwd", self.dims.pre_apply_fwd_cost()), + ("post_apply_fwd", self.dims.post_apply_fwd_cost()), + ("pre_apply_bwd", self.dims.pre_apply_bwd_cost()), + ("coeff_bwd", self.dims.coeff_bwd_cost()), + ("post_apply_bwd", self.dims.post_apply_bwd_cost()), + ] + for name, cost in costs: + self.assertIsInstance(cost.flops, int, msg=f"{name} flops must be int") + self.assertGreater(cost.flops, 0, msg=f"{name} flops must be positive") + self.assertIsInstance(cost.bytes_accessed, int, msg=f"{name} bytes_accessed must be int") + self.assertGreater(cost.bytes_accessed, 0, msg=f"{name} bytes_accessed must be positive") + self.assertIsInstance(cost.transcendentals, int, msg=f"{name} transcendentals must be int") + self.assertGreaterEqual(cost.transcendentals, 0, msg=f"{name} transcendentals must be non-negative") + + # coeff kernels have transcendentals (sigmoid + softmax) + self.assertGreater(self.dims.coeff_fwd_cost().transcendentals, 0) + self.assertGreater(self.dims.coeff_bwd_cost().transcendentals, 0) + # pre and post apply kernels have 0 transcendentals + self.assertEqual(self.dims.pre_apply_fwd_cost().transcendentals, 0) + self.assertEqual(self.dims.post_apply_fwd_cost().transcendentals, 0) + self.assertEqual(self.dims.pre_apply_bwd_cost().transcendentals, 0) + self.assertEqual(self.dims.post_apply_bwd_cost().transcendentals, 0) + if __name__ == "__main__": - unittest.main() + absltest.main()