Add MaxText mHC-lite Pallas/Mosaic TPU kernel and profiling benchmarks.

src/maxtext/configs/base.yml

@@ -1312,6 +1312,13 @@ 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. Default of 128 is
+# optimal for TPU v7 memory constraints; 256 is optimal for TPU v6.
+mhc_pallas_kernel_bwd_block_size: 128
 
 ################################## DeepSeek Engram ##################################
 # Indices of transformer layers where Engram are integrated; leave empty [] to disable.

src/maxtext/configs/types.py

@@ -1692,6 +1692,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):

src/maxtext/kernels/mhc/__init__.py

@@ -0,0 +1,26 @@
+# 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 MhcContext
+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",
+    "UnsupportedInputError",
+]

src/maxtext/kernels/mhc/api.py

@@ -0,0 +1,166 @@
+# 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
+
+
+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,
+    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,
+    permutations: jax.Array,
+    *,
+    rms_epsilon: float = 1e-5,
+    pre_mapping_epsilon: float = 1e-6,
+    implementation: Implementation | Sequence[Implementation] | None = None,
+    block_size: int = common.DEFAULT_BLOCK_SIZE,
+    bwd_block_size: int = common.DEFAULT_BWD_BLOCK_SIZE,
+    vmem_limit_bytes: int = common.DEFAULT_VMEM_LIMIT_BYTES,
+    interpret: bool = False,
+) -> 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)`.
+    norm_scale: RMSNorm scale parameter of shape `(streams * embedding,)`.
+    pre_alpha: Projection matrix for pre-gate of shape `(streams * embedding,
+      streams)`.
+    pre_bias: Bias vector for pre-gate of shape `(streams,)`.
+    pre_scale: Scalar scale parameter for pre-gate of shape `(1,)`.
+    post_alpha: Projection matrix for post-gate of shape `(streams * embedding,
+      streams)`.
+    post_bias: Bias vector for post-gate of shape `(streams,)`.
+    post_scale: Scalar scale parameter for post-gate of shape `(1,)`.
+    res_alpha: Projection matrix for residual mixing of shape `(streams *
+      embedding, num_permutations)`.
+    res_bias: Bias vector for residual mixing of shape `(num_permutations,)`.
+    res_scale: Scalar scale parameter for residual mixing of shape `(1,)`.
+    permutations: All permutation matrices of shape `(num_permutations, streams,
+      streams)`.
+    rms_epsilon: Small constant added to RMSNorm denominator for numerical
+      stability.
+    pre_mapping_epsilon: Small constant added to pre-gate output.
+    implementation: Preferred implementation (`"mosaic"` or `"mosaic_tpu"`).
+    block_size: Token-axis Pallas block size for the forward kernels.
+    bwd_block_size: Token-axis block size for backward kernels.
+    vmem_limit_bytes: Scoped VMEM limit passed to the Mosaic compiler.
+    interpret: Whether to run the Pallas calls in interpret mode.
+
+  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,
+      norm_scale,
+      pre_alpha,
+      pre_bias,
+      pre_scale,
+      post_alpha,
+      post_bias,
+      post_scale,
+      res_alpha,
+      res_bias,
+      res_scale,
+      permutations,
+      rms_epsilon=rms_epsilon,
+      pre_mapping_epsilon=pre_mapping_epsilon,
+      block_size=block_size,
+      bwd_block_size=bwd_block_size,
+      vmem_limit_bytes=vmem_limit_bytes,
+      interpret=interpret,
+  )
+  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,
+    *,
+    block_size: int = common.DEFAULT_BLOCK_SIZE,
+    bwd_block_size: int = common.DEFAULT_POST_BWD_BLOCK_SIZE,
+    bwd_feature_block_size: int | None = None,
+    vmem_limit_bytes: int = common.DEFAULT_VMEM_LIMIT_BYTES,
+    interpret: bool = False,
+) -> 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`.
+    block_size: Token-axis Pallas block size for the forward kernel.
+    bwd_block_size: Token-axis block size for backward kernel.
+    bwd_feature_block_size: Feature-axis block size for backward kernel.
+    vmem_limit_bytes: Scoped VMEM limit passed to the Mosaic compiler.
+    interpret: Whether to run the Pallas calls in interpret mode.
+
+  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}'")
+  if bwd_feature_block_size is None:
+    bwd_feature_block_size = min(common.DEFAULT_POST_BWD_FEATURE_BLOCK_SIZE, context.x.shape[-1])
+  kernel_context = (context.x, context.h_post, context.residual)
+  return mhc_kernels_fwd.post(
+      layer_output,
+      kernel_context,
+      block_size=block_size,
+      bwd_block_size=bwd_block_size,
+      bwd_feature_block_size=bwd_feature_block_size,
+      vmem_limit_bytes=vmem_limit_bytes,
+      interpret=interpret,
+  )

src/maxtext/kernels/mhc/common.py

@@ -0,0 +1,151 @@
+# 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 block math and tiling constants for the MaxText 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})
+
+
+def whole(shape):
+  """Returns a full-array BlockSpec for values that stay VMEM-resident."""
+  return pl.BlockSpec(shape, lambda _: tuple(0 for _ in shape))
+
+
+def fold_norm_scale(norm_scale, pre_alpha, post_alpha, res_alpha):
+  """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 mhc_coeffs(
+    x,
+    phi,
+    pre_scale,
+    pre_bias,
+    post_scale,
+    post_bias,
+    res_scale,
+    res_bias,
+    permutations,
+    *,
+    rms_epsilon,
+    pre_mapping_epsilon,
+):
+  """Computes all mHC-lite coefficients without materializing normalized x."""
+  tokens, streams, embedding = x.shape
+  flattened = x.reshape(tokens, streams * embedding)
+  projected = jnp.dot(flattened, phi.astype(jnp.bfloat16), preferred_element_type=jnp.float32)
+
+  flattened_f32 = flattened.astype(jnp.float32)
+  mean_square = jnp.mean(flattened_f32 * flattened_f32, axis=-1, keepdims=True)
+  projected = projected * jax.lax.rsqrt(mean_square + rms_epsilon)
+
+  pre_logits = projected[:, :streams]
+  post_logits = projected[:, streams : 2 * streams]
+  res_logits = projected[:, 2 * streams :]
+  h_pre = jax.nn.sigmoid(pre_scale.astype(jnp.float32) * pre_logits + pre_bias.astype(jnp.float32)) + pre_mapping_epsilon
+  h_post = 2.0 * jax.nn.sigmoid(post_scale.astype(jnp.float32) * post_logits + post_bias.astype(jnp.float32))
+  weights = jax.nn.softmax(
+      res_scale.astype(jnp.float32) * res_logits + res_bias.astype(jnp.float32),
+      axis=-1,
+  )
+  permutation_count = permutations.shape[0]
+  residual = jnp.dot(
+      weights,
+      permutations.reshape(permutation_count, streams * streams).astype(jnp.float32),
+  ).reshape(tokens, streams, streams)
+  return h_pre, h_post, residual
+
+
+def pre_apply(x, h_pre):
+  """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, layer_output, h_post, residual):
+  """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 validate_token_block_size(tokens, block_size, *, name):
+  """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" f" {name} ({block_size}).")
+
+
+def validate_feature_block_size(embedding, block_size):
+  """Validates the feature tile used by the post-application backward."""
+  if block_size < 128 or block_size % 128:
+    raise UnsupportedInputError("bwd_feature_block_size must be a positive multiple of 128; got" f" {block_size}.")
+  if embedding % block_size:
+    raise UnsupportedInputError(
+        f"The embedding dimension ({embedding}) must be divisible by " f"bwd_feature_block_size ({block_size})."
+    )
+
+
+def validate_inputs(x, block_size, permutations_shape=None, *, block_size_name="block_size"):
+  """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("Expected x to have shape (batch, sequence, streams, embedding); got" f" {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"
+        " expansion rate 4 only; got"
+        f" 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)