diff --git a/src/maxtext/configs/base.yml b/src/maxtext/configs/base.yml
index 08d22bae24..16bc05bd95 100644
--- a/src/maxtext/configs/base.yml
+++ b/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.
diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py
index acaf519a3d..7fd14d14b3 100644
--- a/src/maxtext/configs/types.py
+++ b/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):
diff --git a/src/maxtext/kernels/mhc/__init__.py b/src/maxtext/kernels/mhc/__init__.py
new file mode 100644
index 0000000000..5f3afd3b5a
--- /dev/null
+++ b/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",
+]
diff --git a/src/maxtext/kernels/mhc/api.py b/src/maxtext/kernels/mhc/api.py
new file mode 100644
index 0000000000..174e0a8d0d
--- /dev/null
+++ b/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,
+  )
diff --git a/src/maxtext/kernels/mhc/common.py b/src/maxtext/kernels/mhc/common.py
new file mode 100644
index 0000000000..b71544cee7
--- /dev/null
+++ b/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)
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..103338ce8c
--- /dev/null
+++ b/src/maxtext/kernels/mhc/mhc_kernels_bwd.py
@@ -0,0 +1,540 @@
+# 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,
+    layer_output,
+    h_post,
+    residual,
+    d_output,
+    *,
+    block_size,
+    feature_block_size,
+    vmem_limit_bytes,
+    interpret,
+):
+  """Builds the feature-tiled Pallas call for the post-branch backward pass."""
+  tokens, streams, embedding = x.shape
+  feature_blocks = embedding // feature_block_size
+
+  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_output_f32 = d_output_ref[...].astype(jnp.float32)
+
+    d_x = jnp.einsum(
+        "tkj,tjd->tkd",
+        residual_ref[...].astype(jnp.bfloat16),
+        d_output_f32,
+        preferred_element_type=jnp.float32,
+    )
+    d_layer_output = jnp.sum(h_post_ref[...][:, :, None] * d_output_f32, axis=1)
+    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 = jnp.sum(layer_output_ref[...][:, None, :] * d_output_f32, axis=-1)
+    d_residual = jnp.einsum(
+        "tjd,tkd->tjk",
+        d_output_f32,
+        x_ref[...],
+        preferred_element_type=jnp.float32,
+    ).transpose(0, 2, 1)
+
+    @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)
+
+  return 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 // block_size, feature_blocks),
+      in_specs=(
+          pl.BlockSpec(
+              (block_size, streams, feature_block_size),
+              lambda token, feature: (token, 0, feature),
+          ),
+          pl.BlockSpec(
+              (block_size, feature_block_size),
+              lambda token, feature: (token, feature),
+          ),
+          pl.BlockSpec((block_size, streams), lambda token, feature: (token, 0)),
+          pl.BlockSpec(
+              (block_size, streams, streams),
+              lambda token, feature: (token, 0, 0),
+          ),
+          pl.BlockSpec(
+              (block_size, streams, feature_block_size),
+              lambda token, feature: (token, 0, feature),
+          ),
+      ),
+      out_specs=(
+          pl.BlockSpec(
+              (block_size, streams, feature_block_size),
+              lambda token, feature: (token, 0, feature),
+          ),
+          pl.BlockSpec(
+              (block_size, feature_block_size),
+              lambda token, feature: (token, feature),
+          ),
+          pl.BlockSpec((block_size, streams), lambda token, feature: (token, 0)),
+          pl.BlockSpec(
+              (block_size, streams, streams),
+              lambda token, feature: (token, 0, 0),
+          ),
+      ),
+      cost_estimate=pl.CostEstimate(
+          flops=int(2 * (2 * tokens * streams * streams * embedding + tokens * streams * embedding)),
+          transcendentals=0,
+          bytes_accessed=int(3 * tokens * streams * embedding * 2 + 2 * tokens * embedding * 4),
+      ),
+      compiler_params=pltpu.CompilerParams(
+          vmem_limit_bytes=vmem_limit_bytes,
+          dimension_semantics=common.SEQUENTIAL_2D_DIMENSION_SEMANTICS,
+      ),
+      interpret=interpret,
+  )(x, layer_output, h_post, residual, d_output)
+
+
+def _pre_apply_bwd(x, h_pre, d_layer_input, d_x_acc, *, block_size, vmem_limit_bytes, interpret):
+  """Builds the Pallas call for the pre-branch backward pass."""
+  tokens, streams, embedding = x.shape
+
+  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
+
+  return pl.pallas_call(
+      kernel,
+      out_shape=(
+          jax.ShapeDtypeStruct((tokens, streams, embedding), x.dtype),
+          jax.ShapeDtypeStruct((tokens, streams), h_pre.dtype),
+      ),
+      grid=(tokens // block_size,),
+      in_specs=(
+          pl.BlockSpec((block_size, streams, embedding), lambda i: (i, 0, 0)),
+          pl.BlockSpec((block_size, streams), lambda i: (i, 0)),
+          pl.BlockSpec((block_size, embedding), lambda i: (i, 0)),
+          pl.BlockSpec((block_size, streams, embedding), lambda i: (i, 0, 0)),
+      ),
+      out_specs=(
+          pl.BlockSpec((block_size, streams, embedding), lambda i: (i, 0, 0)),
+          pl.BlockSpec((block_size, streams), lambda i: (i, 0)),
+      ),
+      cost_estimate=pl.CostEstimate(
+          flops=int(4 * tokens * streams * embedding),
+          transcendentals=0,
+          bytes_accessed=int(3 * tokens * streams * embedding * 2 + tokens * embedding * 2),
+      ),
+      compiler_params=pltpu.CompilerParams(
+          vmem_limit_bytes=vmem_limit_bytes,
+          dimension_semantics=common.PARALLEL_DIMENSION_SEMANTICS,
+      ),
+      interpret=interpret,
+  )(x, h_pre, d_layer_input, d_x_acc)
+
+
+def _coeff_bwd(
+    x,
+    phi,
+    pre_scale,
+    pre_bias,
+    post_scale,
+    post_bias,
+    res_scale,
+    res_bias,
+    permutations,
+    d_h_pre,
+    d_h_post,
+    d_residual,
+    d_x_acc,
+    *,
+    block_size,
+    vmem_limit_bytes,
+    interpret,
+    rms_epsilon,
+    pre_mapping_epsilon,
+):
+  """Builds the Pallas call for coefficients and parameter gradients."""
+  tokens, streams, embedding = x.shape
+  flattened_size = streams * embedding
+  permutation_count = 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)
+    permutations_value = permutations_ref[...]
+
+    def mhc_coeffs_fn(
+        x_value,
+        phi_value,
+        pre_scale_value,
+        pre_bias_value,
+        post_scale_value,
+        post_bias_value,
+        res_scale_value,
+        res_bias_value,
+    ):
+      return common.mhc_coeffs(
+          x_value,
+          phi_value,
+          pre_scale_value,
+          pre_bias_value,
+          post_scale_value,
+          post_bias_value,
+          res_scale_value,
+          res_bias_value,
+          permutations_value,
+          rms_epsilon=rms_epsilon,
+          pre_mapping_epsilon=pre_mapping_epsilon,
+      )
+
+    _, vjp = jax.vjp(
+        mhc_coeffs_fn,
+        x_ref[...],
+        phi_ref[...],
+        pre_scale_ref[...],
+        pre_bias_ref[...],
+        post_scale_ref[...],
+        post_bias_ref[...],
+        res_scale_ref[...],
+        res_bias_ref[...],
+    )
+    (
+        d_x,
+        d_phi,
+        d_pre_scale,
+        d_pre_bias,
+        d_post_scale,
+        d_post_bias,
+        d_res_scale,
+        d_res_bias,
+    ) = vjp((d_h_pre_ref[...], d_h_post_ref[...], d_residual_ref[...]))
+    d_x_ref[...] = (d_x.astype(jnp.float32) + d_x_acc_ref[...].astype(jnp.float32)).astype(d_x_ref.dtype)
+
+    @pl.when(program_id == 0)
+    def initialize_reductions():
+      d_phi_ref[...] = jnp.zeros_like(d_phi_ref)
+      d_pre_scale_ref[...] = jnp.zeros_like(d_pre_scale_ref)
+      d_pre_bias_ref[...] = jnp.zeros_like(d_pre_bias_ref)
+      d_post_scale_ref[...] = jnp.zeros_like(d_post_scale_ref)
+      d_post_bias_ref[...] = jnp.zeros_like(d_post_bias_ref)
+      d_res_scale_ref[...] = jnp.zeros_like(d_res_scale_ref)
+      d_res_bias_ref[...] = jnp.zeros_like(d_res_bias_ref)
+
+    d_phi_ref[...] += d_phi.astype(jnp.float32)
+    d_pre_scale_ref[...] += d_pre_scale.astype(jnp.float32)
+    d_pre_bias_ref[...] += d_pre_bias.astype(jnp.float32)
+    d_post_scale_ref[...] += d_post_scale.astype(jnp.float32)
+    d_post_bias_ref[...] += d_post_bias.astype(jnp.float32)
+    d_res_scale_ref[...] += d_res_scale.astype(jnp.float32)
+    d_res_bias_ref[...] += d_res_bias.astype(jnp.float32)
+
+  return pl.pallas_call(
+      kernel,
+      out_shape=(
+          jax.ShapeDtypeStruct((tokens, streams, embedding), x.dtype),
+          jax.ShapeDtypeStruct((flattened_size, 2 * streams + permutation_count), jnp.float32),
+          jax.ShapeDtypeStruct((1,), jnp.float32),
+          jax.ShapeDtypeStruct((streams,), jnp.float32),
+          jax.ShapeDtypeStruct((1,), jnp.float32),
+          jax.ShapeDtypeStruct((streams,), jnp.float32),
+          jax.ShapeDtypeStruct((1,), jnp.float32),
+          jax.ShapeDtypeStruct((permutation_count,), jnp.float32),
+      ),
+      grid=(tokens // block_size,),
+      in_specs=(
+          pl.BlockSpec((block_size, streams, embedding), lambda i: (i, 0, 0)),
+          common.whole((flattened_size, 2 * streams + permutation_count)),
+          common.whole((1,)),
+          common.whole((streams,)),
+          common.whole((1,)),
+          common.whole((streams,)),
+          common.whole((1,)),
+          common.whole((permutation_count,)),
+          common.whole((permutation_count, streams, streams)),
+          pl.BlockSpec((block_size, streams), lambda i: (i, 0)),
+          pl.BlockSpec((block_size, streams), lambda i: (i, 0)),
+          pl.BlockSpec((block_size, streams, streams), lambda i: (i, 0, 0)),
+          pl.BlockSpec((block_size, streams, embedding), lambda i: (i, 0, 0)),
+      ),
+      out_specs=(
+          pl.BlockSpec((block_size, streams, embedding), lambda i: (i, 0, 0)),
+          common.whole((flattened_size, 2 * streams + permutation_count)),
+          common.whole((1,)),
+          common.whole((streams,)),
+          common.whole((1,)),
+          common.whole((streams,)),
+          common.whole((1,)),
+          common.whole((permutation_count,)),
+      ),
+      cost_estimate=pl.CostEstimate(
+          flops=int(
+              2
+              * (
+                  2 * tokens * flattened_size * (2 * streams + permutation_count)
+                  + 2 * tokens * permutation_count * streams * streams
+              )
+          ),
+          transcendentals=int(tokens * (streams + permutation_count)),
+          bytes_accessed=int(
+              3 * tokens * streams * embedding * 2 + 2 * flattened_size * (2 * streams + permutation_count) * 4
+          ),
+      ),
+      compiler_params=pltpu.CompilerParams(
+          vmem_limit_bytes=vmem_limit_bytes,
+          dimension_semantics=common.SEQUENTIAL_DIMENSION_SEMANTICS,
+      ),
+      interpret=interpret,
+  )(
+      x,
+      phi,
+      pre_scale,
+      pre_bias,
+      post_scale,
+      post_bias,
+      res_scale,
+      res_bias,
+      permutations,
+      d_h_pre,
+      d_h_post,
+      d_residual,
+      d_x_acc,
+  )
+
+
+def pre_bwd(
+    residuals,
+    cotangents,
+    x,
+    norm_scale,
+    pre_alpha,
+    pre_bias,
+    pre_scale,
+    post_alpha,
+    post_bias,
+    post_scale,
+    res_alpha,
+    res_bias,
+    res_scale,
+    permutations,
+    *,
+    rms_epsilon,
+    pre_mapping_epsilon,
+    block_size,
+    vmem_limit_bytes,
+    interpret,
+):
+  """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 = d_x_acc.reshape(tokens, streams, embedding)
+  d_h_post = d_h_post.reshape(tokens, streams)
+  d_residual = d_residual.reshape(tokens, streams, streams)
+  d_layer_input = d_layer_input.reshape(tokens, embedding)
+
+  d_x_acc, d_h_pre = _pre_apply_bwd(
+      x_flat,
+      h_pre,
+      d_layer_input,
+      d_x_acc,
+      block_size=block_size,
+      vmem_limit_bytes=vmem_limit_bytes,
+      interpret=interpret,
+  )
+  (
+      d_x,
+      d_phi,
+      d_pre_scale,
+      d_pre_bias,
+      d_post_scale,
+      d_post_bias,
+      d_res_scale,
+      d_res_bias,
+  ) = _coeff_bwd(
+      x_flat,
+      phi,
+      pre_scale,
+      pre_bias,
+      post_scale,
+      post_bias,
+      res_scale,
+      res_bias,
+      permutations,
+      d_h_pre,
+      d_h_post,
+      d_residual,
+      d_x_acc,
+      block_size=block_size,
+      vmem_limit_bytes=vmem_limit_bytes,
+      interpret=interpret,
+      rms_epsilon=rms_epsilon,
+      pre_mapping_epsilon=pre_mapping_epsilon,
+  )
+  _, phi_vjp = jax.vjp(common.fold_norm_scale, norm_scale, pre_alpha, post_alpha, res_alpha)
+  d_norm_scale, d_pre_alpha, d_post_alpha, d_res_alpha = phi_vjp(d_phi)
+  d_x = d_x.reshape(batch, sequence, streams, embedding)
+  return (
+      d_x,
+      d_norm_scale,
+      d_pre_alpha,
+      d_pre_bias.astype(pre_bias.dtype),
+      d_pre_scale.astype(pre_scale.dtype),
+      d_post_alpha,
+      d_post_bias.astype(post_bias.dtype),
+      d_post_scale.astype(post_scale.dtype),
+      d_res_alpha,
+      d_res_bias.astype(res_bias.dtype),
+      d_res_scale.astype(res_scale.dtype),
+      jnp.zeros_like(permutations),
+  )
+
+
+def post_bwd(
+    cotangent,
+    layer_output,
+    x,
+    h_post,
+    residual,
+    *,
+    block_size,
+    feature_block_size,
+    vmem_limit_bytes,
+    interpret,
+):
+  """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),
+      block_size=block_size,
+      feature_block_size=feature_block_size,
+      vmem_limit_bytes=vmem_limit_bytes,
+      interpret=interpret,
+  )
+  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(
+    _block_size,
+    bwd_block_size,
+    vmem_limit_bytes,
+    interpret,
+    rms_epsilon,
+    pre_mapping_epsilon,
+    residuals,
+    cotangents,
+):
+  """Custom-VJP backward rule for the low-level pre-branch entry point."""
+  saved, primals = 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),
+      *primals,
+      rms_epsilon=rms_epsilon,
+      pre_mapping_epsilon=pre_mapping_epsilon,
+      block_size=bwd_block_size,
+      vmem_limit_bytes=vmem_limit_bytes,
+      interpret=interpret,
+  )
+
+
+def post_op_bwd(
+    _block_size,
+    bwd_block_size,
+    bwd_feature_block_size,
+    vmem_limit_bytes,
+    interpret,
+    saved,
+    d_output,
+):
+  """Custom-VJP backward rule for the low-level post-branch entry point."""
+  x, layer_output, h_post, residual = saved
+  d_layer_output, d_x, d_h_post, d_residual = post_bwd(
+      d_output,
+      layer_output,
+      x,
+      h_post,
+      residual,
+      block_size=bwd_block_size,
+      feature_block_size=bwd_feature_block_size,
+      vmem_limit_bytes=vmem_limit_bytes,
+      interpret=interpret,
+  )
+  return d_x, d_layer_output, 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..f42c359889
--- /dev/null
+++ b/src/maxtext/kernels/mhc/mhc_kernels_fwd.py
@@ -0,0 +1,548 @@
+# 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
+
+type PreResiduals = tuple[jax.Array, jax.Array]
+
+
+def _coeff_fwd(
+    x,
+    phi,
+    pre_scale,
+    pre_bias,
+    post_scale,
+    post_bias,
+    res_scale,
+    res_bias,
+    permutations,
+    *,
+    block_size,
+    vmem_limit_bytes,
+    interpret,
+    rms_epsilon,
+    pre_mapping_epsilon,
+):
+  """Builds the Pallas call that computes the shared mHC coefficients."""
+  tokens, streams, embedding = x.shape
+  flattened_size = streams * embedding
+  permutation_count = 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,
+  ):
+    h_pre, h_post, residual = common.mhc_coeffs(
+        x_ref[...],
+        phi_ref[...],
+        pre_scale_ref[...],
+        pre_bias_ref[...],
+        post_scale_ref[...],
+        post_bias_ref[...],
+        res_scale_ref[...],
+        res_bias_ref[...],
+        permutations_ref[...],
+        rms_epsilon=rms_epsilon,
+        pre_mapping_epsilon=pre_mapping_epsilon,
+    )
+    h_pre_ref[...] = h_pre
+    h_post_ref[...] = h_post
+    residual_ref[...] = residual
+
+  cost = pl.CostEstimate(
+      flops=int(
+          2 * tokens * flattened_size * (2 * streams + permutation_count)
+          + 2 * tokens * permutation_count * streams * streams
+      ),
+      transcendentals=int(tokens * (streams + permutation_count)),
+      bytes_accessed=int(tokens * streams * embedding * 2 + flattened_size * (2 * streams + permutation_count) * 4),
+  )
+  return 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 // block_size,),
+      in_specs=(
+          pl.BlockSpec((block_size, streams, embedding), lambda i: (i, 0, 0)),
+          common.whole((flattened_size, 2 * streams + permutation_count)),
+          common.whole((1,)),
+          common.whole((streams,)),
+          common.whole((1,)),
+          common.whole((streams,)),
+          common.whole((1,)),
+          common.whole((permutation_count,)),
+          common.whole((permutation_count, streams, streams)),
+      ),
+      out_specs=(
+          pl.BlockSpec((block_size, streams), lambda i: (i, 0)),
+          pl.BlockSpec((block_size, streams), lambda i: (i, 0)),
+          pl.BlockSpec((block_size, streams, streams), lambda i: (i, 0, 0)),
+      ),
+      cost_estimate=cost,
+      compiler_params=pltpu.CompilerParams(
+          vmem_limit_bytes=vmem_limit_bytes,
+          dimension_semantics=common.PARALLEL_DIMENSION_SEMANTICS,
+      ),
+      interpret=interpret,
+  )(
+      x,
+      phi,
+      pre_scale,
+      pre_bias,
+      post_scale,
+      post_bias,
+      res_scale,
+      res_bias,
+      permutations,
+  )
+
+
+def _pre_apply_fwd(x, h_pre, *, block_size, vmem_limit_bytes, interpret):
+  tokens, streams, embedding = x.shape
+
+  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 // block_size,),
+      in_specs=(
+          pl.BlockSpec((block_size, streams, embedding), lambda i: (i, 0, 0)),
+          pl.BlockSpec((block_size, streams), lambda i: (i, 0)),
+      ),
+      out_specs=pl.BlockSpec((block_size, embedding), lambda i: (i, 0)),
+      cost_estimate=pl.CostEstimate(
+          flops=int(2 * tokens * streams * embedding),
+          transcendentals=0,
+          bytes_accessed=int(tokens * streams * embedding * 2 + tokens * embedding * 2),
+      ),
+      compiler_params=pltpu.CompilerParams(
+          vmem_limit_bytes=vmem_limit_bytes,
+          dimension_semantics=common.PARALLEL_DIMENSION_SEMANTICS,
+      ),
+      interpret=interpret,
+  )(x, h_pre)
+
+
+def _post_apply_fwd(
+    x,
+    layer_output,
+    h_post,
+    residual,
+    *,
+    block_size,
+    vmem_limit_bytes,
+    interpret,
+):
+  """Builds the Pallas call for the post-branch forward pass."""
+  tokens, streams, embedding = x.shape
+
+  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 // block_size,),
+      in_specs=(
+          pl.BlockSpec((block_size, streams, embedding), lambda i: (i, 0, 0)),
+          pl.BlockSpec((block_size, embedding), lambda i: (i, 0)),
+          pl.BlockSpec((block_size, streams), lambda i: (i, 0)),
+          pl.BlockSpec((block_size, streams, streams), lambda i: (i, 0, 0)),
+      ),
+      out_specs=pl.BlockSpec((block_size, streams, embedding), lambda i: (i, 0, 0)),
+      cost_estimate=pl.CostEstimate(
+          flops=int(2 * tokens * streams * streams * embedding + tokens * streams * embedding),
+          transcendentals=0,
+          bytes_accessed=int(2 * tokens * streams * embedding * 2 + tokens * embedding * 4),
+      ),
+      compiler_params=pltpu.CompilerParams(
+          vmem_limit_bytes=vmem_limit_bytes,
+          dimension_semantics=common.PARALLEL_DIMENSION_SEMANTICS,
+      ),
+      interpret=interpret,
+  )(x, layer_output, h_post, residual)
+
+
+def pre_fwd(
+    x,
+    norm_scale,
+    pre_alpha,
+    pre_bias,
+    pre_scale,
+    post_alpha,
+    post_bias,
+    post_scale,
+    res_alpha,
+    res_bias,
+    res_scale,
+    permutations,
+    *,
+    rms_epsilon,
+    pre_mapping_epsilon,
+    block_size,
+    vmem_limit_bytes,
+    interpret,
+) -> tuple[tuple[jax.Array, ...], PreResiduals]:
+  """Runs the coefficient and pre-application forward kernels."""
+  common.validate_inputs(x, block_size, permutations.shape)
+  batch, sequence, streams, embedding = x.shape
+  x_flat = x.reshape(batch * sequence, streams, embedding)
+  phi = common.fold_norm_scale(norm_scale, pre_alpha, post_alpha, res_alpha)
+  h_pre, h_post, residual = _coeff_fwd(
+      x_flat,
+      phi,
+      pre_scale,
+      pre_bias,
+      post_scale,
+      post_bias,
+      res_scale,
+      res_bias,
+      permutations,
+      block_size=block_size,
+      vmem_limit_bytes=vmem_limit_bytes,
+      interpret=interpret,
+      rms_epsilon=rms_epsilon,
+      pre_mapping_epsilon=pre_mapping_epsilon,
+  )
+  layer_input = _pre_apply_fwd(
+      x_flat,
+      h_pre,
+      block_size=block_size,
+      vmem_limit_bytes=vmem_limit_bytes,
+      interpret=interpret,
+  )
+  output = (
+      layer_input.reshape(batch, sequence, embedding),
+      x,
+      h_post.reshape(batch, sequence, streams),
+      residual.reshape(batch, sequence, streams, streams),
+  )
+  return output, (phi, h_pre)
+
+
+def post_fwd(
+    layer_output,
+    x,
+    h_post,
+    residual,
+    *,
+    block_size,
+    vmem_limit_bytes,
+    interpret,
+) -> jax.Array:
+  """Runs the fused post-gate and residual-mixing forward kernel."""
+  common.validate_inputs(x, 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),
+      block_size=block_size,
+      vmem_limit_bytes=vmem_limit_bytes,
+      interpret=interpret,
+  )
+  return output.reshape(batch, sequence, streams, embedding)
+
+
+@functools.partial(jax.custom_vjp, nondiff_argnums=(0, 1, 2, 3, 4, 5))
+def _pre_op(
+    block_size,
+    _bwd_block_size,
+    vmem_limit_bytes,
+    interpret,
+    rms_epsilon,
+    pre_mapping_epsilon,
+    x,
+    norm_scale,
+    pre_alpha,
+    pre_bias,
+    pre_scale,
+    post_alpha,
+    post_bias,
+    post_scale,
+    res_alpha,
+    res_bias,
+    res_scale,
+    permutations,
+):
+  """Differentiable pre-branch mHC operation."""
+  (layer_input, x_context, h_post, residual), _ = pre_fwd(
+      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,
+      vmem_limit_bytes=vmem_limit_bytes,
+      interpret=interpret,
+  )
+  return layer_input, (x_context, h_post, residual)
+
+
+def _pre_op_fwd(
+    block_size,
+    _bwd_block_size,
+    vmem_limit_bytes,
+    interpret,
+    rms_epsilon,
+    pre_mapping_epsilon,
+    x,
+    norm_scale,
+    pre_alpha,
+    pre_bias,
+    pre_scale,
+    post_alpha,
+    post_bias,
+    post_scale,
+    res_alpha,
+    res_bias,
+    res_scale,
+    permutations,
+):
+  """Custom-VJP forward rule for the pre-branch operation."""
+  primals = (
+      x,
+      norm_scale,
+      pre_alpha,
+      pre_bias,
+      pre_scale,
+      post_alpha,
+      post_bias,
+      post_scale,
+      res_alpha,
+      res_bias,
+      res_scale,
+      permutations,
+  )
+  (layer_input, x_context, h_post, residual), saved = pre_fwd(
+      *primals,
+      rms_epsilon=rms_epsilon,
+      pre_mapping_epsilon=pre_mapping_epsilon,
+      block_size=block_size,
+      vmem_limit_bytes=vmem_limit_bytes,
+      interpret=interpret,
+  )
+  output = (layer_input, (x_context, h_post, residual))
+  return output, (saved, primals)
+
+
+_pre_op.defvjp(_pre_op_fwd, mhc_kernels_bwd.pre_op_bwd)
+
+
+@functools.partial(jax.custom_vjp, nondiff_argnums=(0, 1, 2, 3, 4))
+def _post_op(
+    block_size,
+    _bwd_block_size,
+    _bwd_feature_block_size,
+    vmem_limit_bytes,
+    interpret,
+    x,
+    layer_output,
+    h_post,
+    residual,
+):
+  """Differentiable post-branch mHC operation."""
+  return post_fwd(
+      layer_output,
+      x,
+      h_post,
+      residual,
+      block_size=block_size,
+      vmem_limit_bytes=vmem_limit_bytes,
+      interpret=interpret,
+  )
+
+
+def _post_op_fwd(
+    block_size,
+    _bwd_block_size,
+    _bwd_feature_block_size,
+    vmem_limit_bytes,
+    interpret,
+    x,
+    layer_output,
+    h_post,
+    residual,
+):
+  """Custom-VJP forward rule for the post-branch operation."""
+  output = post_fwd(
+      layer_output,
+      x,
+      h_post,
+      residual,
+      block_size=block_size,
+      vmem_limit_bytes=vmem_limit_bytes,
+      interpret=interpret,
+  )
+  return output, (x, layer_output, h_post, residual)
+
+
+_post_op.defvjp(_post_op_fwd, mhc_kernels_bwd.post_op_bwd)
+
+
+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,
+    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, common.KernelContext]:
+  """Runs the coefficient and pre-application kernels.
+
+  Args:
+    x: Input streams with shape `[batch, sequence, streams, embedding]`.
+    norm_scale: RMSNorm scale with shape `[streams * embedding]`.
+    pre_alpha: Pre-gate projection with shape `[streams * embedding, streams]`.
+    pre_bias: Pre-gate bias with shape `[streams]`.
+    pre_scale: Pre-gate scalar scale with shape `[1]`.
+    post_alpha: Post-gate projection with shape `[streams * embedding,
+      streams]`.
+    post_bias: Post-gate bias with shape `[streams]`.
+    post_scale: Post-gate scalar scale with shape `[1]`.
+    res_alpha: Residual projection with shape `[streams * embedding, streams!]`.
+    res_bias: Residual bias with shape `[streams!]`.
+    res_scale: Residual scalar scale with shape `[1]`.
+    permutations: Permutation matrices with shape `[streams!, streams,
+      streams]`.
+    rms_epsilon: Epsilon used by RMSNorm.
+    pre_mapping_epsilon: Epsilon added to the pre-gate output.
+    block_size: Token-axis Pallas block size for the forward kernels.
+    bwd_block_size: Token-axis block size for the coefficient and
+      pre-application 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 pair containing the branch input and opaque kernel context.
+  """
+  common.validate_inputs(x, block_size, permutations.shape)
+  common.validate_token_block_size(x.shape[0] * x.shape[1], bwd_block_size, name="bwd_block_size")
+  return _pre_op(
+      block_size,
+      bwd_block_size,
+      vmem_limit_bytes,
+      interpret,
+      rms_epsilon,
+      pre_mapping_epsilon,
+      x,
+      norm_scale,
+      pre_alpha,
+      pre_bias,
+      pre_scale,
+      post_alpha,
+      post_bias,
+      post_scale,
+      res_alpha,
+      res_bias,
+      res_scale,
+      permutations,
+  )
+
+
+def post(
+    layer_output: jax.Array,
+    context: common.KernelContext,
+    *,
+    block_size: int = common.DEFAULT_BLOCK_SIZE,
+    bwd_block_size: int = common.DEFAULT_POST_BWD_BLOCK_SIZE,
+    bwd_feature_block_size: int = common.DEFAULT_POST_BWD_FEATURE_BLOCK_SIZE,
+    vmem_limit_bytes: int = common.DEFAULT_VMEM_LIMIT_BYTES,
+    interpret: bool = False,
+) -> 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`.
+    block_size: Token-axis Pallas block size for the forward kernel.
+    bwd_block_size: Token-axis block size for the post-application backward
+      kernel.
+    bwd_feature_block_size: Feature-axis block size for the post-application
+      backward kernel.
+    vmem_limit_bytes: Scoped VMEM limit passed to the Mosaic compiler.
+    interpret: Whether to run the Pallas call in interpret mode.
+
+  Returns:
+    Mixed output streams with shape `[batch, sequence, streams, embedding]`.
+  """
+  x, h_post, residual = context
+  bwd_feature_block_size = min(x.shape[-1], bwd_feature_block_size)
+  common.validate_inputs(x, block_size)
+  common.validate_token_block_size(x.shape[0] * x.shape[1], bwd_block_size, name="bwd_block_size")
+  common.validate_feature_block_size(x.shape[-1], bwd_feature_block_size)
+  return _post_op(
+      block_size,
+      bwd_block_size,
+      bwd_feature_block_size,
+      vmem_limit_bytes,
+      interpret,
+      x,
+      layer_output,
+      h_post,
+      residual,
+  )
diff --git a/src/maxtext/layers/mhc.py b/src/maxtext/layers/mhc.py
index d41cf5967e..05787f6971 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,
@@ -246,36 +250,59 @@ 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)
+      layer_input, context = mhc_kernel.pre(
+          x,
+          jnp.asarray(self.mhc_norm.scale[...], self.dtype),
+          jnp.asarray(self.pre_alpha[...], self.dtype),
+          jnp.asarray(self.pre_beta[...], self.dtype),
+          jnp.asarray(self.pre_alpha_scale[...], self.dtype),
+          jnp.asarray(self.post_alpha[...], self.dtype),
+          jnp.asarray(self.post_beta[...], self.dtype),
+          jnp.asarray(self.post_alpha_scale[...], self.dtype),
+          jnp.asarray(self.res_alpha[...], self.dtype),
+          jnp.asarray(self.res_beta[...], self.dtype),
+          jnp.asarray(self.res_alpha_scale[...], self.dtype),
+          jnp.asarray(self.permutation_matrices, self.dtype),
+          rms_epsilon=self.config.normalization_layer_epsilon,
+          block_size=fwd_block_size,
+          bwd_block_size=bwd_block_size,
+      )
+    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 +320,17 @@ 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)
+      output = mhc_kernel.post(
+          layer_out,
+          context,
+          block_size=fwd_block_size,
+          bwd_block_size=bwd_block_size,
+      )
+      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..6e29a3d368 100644
--- a/tests/unit/mhc_test.py
+++ b/tests/unit/mhc_test.py
@@ -14,7 +14,11 @@
 
 """Test for DeepSeek Manifold-Constrained Hyper Connections (mHC)."""
 
+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 +27,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 +94,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 +158,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 +263,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 +342,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 +351,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 +372,426 @@ 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):
+        kwargs["interpret"] = True
+        return real_pre(*args, **kwargs)
+
+      def fake_post(*args, **kwargs):
+        kwargs["interpret"] = True
+        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
+          self.assertEqual(kwargs_pre.get("block_size"), 256)
+          self.assertEqual(kwargs_pre.get("bwd_block_size"), 128)
+
+          _, kwargs_post = mock_post.call_args
+          self.assertEqual(kwargs_post.get("block_size"), 256)
+          self.assertEqual(kwargs_post.get("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):
+        kwargs["interpret"] = True
+        return real_pre(*args, **kwargs)
+
+      def fake_post(*args, **kwargs):
+        kwargs["interpret"] = True
+        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
+        self.assertEqual(kwargs_pre.get("block_size"), 128)
+        self.assertEqual(kwargs_pre.get("bwd_block_size"), 64)
+
+        _, kwargs_post = mock_post.call_args
+        self.assertEqual(kwargs_post.get("block_size"), 128)
+        self.assertEqual(kwargs_post.get("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 _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)
+  differentiable = (
+      x,
+      norm_scale,
+      pre_alpha,
+      pre_bias,
+      pre_scale,
+      post_alpha,
+      post_bias,
+      post_scale,
+      res_alpha,
+      res_bias,
+      res_scale,
+  )
+  return differentiable, permutations, cotangent
+
+
+def _run_pipeline_reference(primals, permutations):
+  """Runs the native JAX/XLA reference implementation of the mHC pipeline."""
+  (
+      x,
+      norm_scale,
+      pre_alpha,
+      pre_bias,
+      pre_scale,
+      post_alpha,
+      post_bias,
+      post_scale,
+      res_alpha,
+      res_bias,
+      res_scale,
+  ) = primals
+  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)
+      * norm_scale.astype(jnp.float32)
+  ).astype(x.dtype)
+
+  h_pre = (
+      jax.nn.sigmoid(
+          pre_scale.astype(jnp.float32)
+          * jnp.dot(
+              normalized,
+              pre_alpha,
+              preferred_element_type=jnp.float32,
+          )
+          + 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(
+      post_scale.astype(jnp.float32)
+      * jnp.dot(
+          normalized,
+          post_alpha,
+          preferred_element_type=jnp.float32,
+      )
+      + post_bias.astype(jnp.float32)
+  )
+  weights = jax.nn.softmax(
+      res_scale.astype(jnp.float32)
+      * jnp.dot(
+          normalized,
+          res_alpha,
+          preferred_element_type=jnp.float32,
+      )
+      + res_bias.astype(jnp.float32),
+      axis=-1,
+  )
+  residual = jnp.dot(
+      weights,
+      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(primals, permutations, implementation=None, interpret=True):
+  """Runs the mHC Pallas kernel API pipeline."""
+  (
+      x,
+      norm_scale,
+      pre_alpha,
+      pre_bias,
+      pre_scale,
+      post_alpha,
+      post_bias,
+      post_scale,
+      res_alpha,
+      res_bias,
+      res_scale,
+  ) = primals
+  layer_input, context = mhc_kernel.pre(
+      x,
+      norm_scale,
+      pre_alpha,
+      pre_bias,
+      pre_scale,
+      post_alpha,
+      post_bias,
+      post_scale,
+      res_alpha,
+      res_bias,
+      res_scale,
+      permutations,
+      implementation=implementation,
+      interpret=interpret,
+  )
+  return mhc_kernel.post(layer_input, context, interpret=interpret)
+
+
+class TestMhcKernelsFwd(parameterized.TestCase):
+  """Unit tests for MaxText mHC-lite Pallas forward kernel."""
+
+  def test_doubly_stochastic(self):
+    differentiable, permutations, _ = _make_kernel_inputs(batch=1, sequence=128, streams=4, embedding=128)
+    (
+        x,
+        norm_scale,
+        pre_alpha,
+        pre_bias,
+        pre_scale,
+        post_alpha,
+        post_bias,
+        post_scale,
+        res_alpha,
+        res_bias,
+        res_scale,
+    ) = differentiable
+    _, context = mhc_kernel.pre(
+        x,
+        norm_scale,
+        pre_alpha,
+        pre_bias,
+        pre_scale,
+        post_alpha,
+        post_bias,
+        post_scale,
+        res_alpha,
+        res_bias,
+        res_scale,
+        permutations,
+        interpret=True,
+    )
+    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):
+    differentiable, permutations, _ = _make_kernel_inputs(batch=2, sequence=64, streams=4, embedding=256)
+    expected = _run_pipeline_reference(differentiable, permutations)
+    actual = _run_pipeline_api(
+        differentiable,
+        permutations,
+        implementation=implementation,
+        interpret=True,
+    )
+    np.testing.assert_allclose(actual, expected, rtol=5e-2, atol=5e-2)
+
+  def test_unsupported_shape_raises_error(self):
+    differentiable, permutations, _ = _make_kernel_inputs(batch=1, sequence=16, streams=2, embedding=128)
+    (
+        x,
+        norm_scale,
+        pre_alpha,
+        pre_bias,
+        pre_scale,
+        post_alpha,
+        post_bias,
+        post_scale,
+        res_alpha,
+        res_bias,
+        res_scale,
+    ) = differentiable
+    with self.assertRaises(mhc_kernel_common.UnsupportedInputError):
+      mhc_kernel.pre(
+          x,
+          norm_scale,
+          pre_alpha,
+          pre_bias,
+          pre_scale,
+          post_alpha,
+          post_bias,
+          post_scale,
+          res_alpha,
+          res_bias,
+          res_scale,
+          permutations,
+          interpret=True,
+      )
+
+
+class TestMhcKernelsBwd(parameterized.TestCase):
+  """Unit tests for MaxText mHC-lite Pallas backward kernel."""
+
+  def test_forward_and_backward_vjp_parity(self):
+    differentiable, permutations, cotangent = _make_kernel_inputs(batch=2, sequence=64, streams=4, embedding=256)
+    expected_out, expected_vjp_fn = jax.vjp(
+        lambda *args: _run_pipeline_reference(args, permutations),
+        *differentiable,
+    )
+    expected_grads = expected_vjp_fn(cotangent)
+
+    actual_out, actual_vjp_fn = jax.vjp(
+        lambda *args: _run_pipeline_api(args, permutations, implementation=None, interpret=True),
+        *differentiable,
+    )
+    actual_grads = actual_vjp_fn(cotangent)
+
+    np.testing.assert_allclose(actual_out, expected_out, rtol=5e-2, atol=5e-2)
+    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",
+      )
+
 
 if __name__ == "__main__":
-  unittest.main()
+  absltest.main()