From 748d5f6593d164124a0152822950d2417899171c Mon Sep 17 00:00:00 2001
From: Suraj Kolla <surajkolla@google.com>
Date: Thu, 18 Jun 2026 23:36:18 +0000
Subject: [PATCH] integrate Tokamax GMM v2 in maxtext

---
 src/maxtext/configs/base.yml                  |    3 +
 src/maxtext/configs/types.py                  |   12 +
 src/maxtext/kernels/megablox/ops.py           |  111 +-
 .../pallas_mosaic_tpu_v2_gmm_kernel.py        | 1307 +++++++++++++++++
 .../pallas_mosaic_tpu_v2_tgmm_kernel.py       |  758 ++++++++++
 src/maxtext/layers/moe.py                     |    5 +-
 tests/integration/tokamax_test.py             |   31 +-
 tests/unit/configs_value_test.py              |   34 +
 8 files changed, 2232 insertions(+), 29 deletions(-)
 create mode 100644 src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_gmm_kernel.py
 create mode 100644 src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_tgmm_kernel.py

diff --git a/src/maxtext/configs/base.yml b/src/maxtext/configs/base.yml
index 84d3c03f57..cd9f4e5919 100644
--- a/src/maxtext/configs/base.yml
+++ b/src/maxtext/configs/base.yml
@@ -1211,6 +1211,9 @@ partial_rotary_factor: 1.0
 
 # Use tokamax library for gmm kernel implementation
 use_tokamax_gmm: false
+# Whether to use GMM v2 for MoE.
+# Requires use_tokamax_gmm: true. Currently incompatible with quantization.
+use_gmm_v2: false
 use_tokamax_splash: false
 # Setting this flag will use a non-pallas implementation.
 use_jax_splash: false
diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py
index e43f34f247..1428a38d69 100644
--- a/src/maxtext/configs/types.py
+++ b/src/maxtext/configs/types.py
@@ -610,6 +610,13 @@ class Attention(BaseModel):
       False,
       description="Whether to use the Tokamax library for GMM kernel implementation.",
   )
+  use_gmm_v2: bool = Field(
+      False,
+      description=(
+          "Whether to use GMM v2 (with bf16 activations and weights) for MoE."
+          " Requires use_tokamax_gmm: true. Currently incompatible with quantization."
+      ),
+  )
   ragged_block_size: int = Field(256, description="Block size for ragged attention.")
   enable_padding_causal_mask: bool = Field(True, description="Temporary flag for TE padding.")
   use_tokamax_splash: bool = Field(False, description="Whether to use tokamax splash attention.")
@@ -3189,6 +3196,11 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de
     if self.share_kv_projections and self.attention_type == "mla":
       raise ValueError("`share_kv_projections` is not compatible with `attention_type='mla'`.")
 
+    if self.use_gmm_v2 and (self.quantization or self.use_qwix_quantization):
+      raise ValueError("Quantization with GMM v2 is not supported yet.")
+    if self.use_gmm_v2 and not self.use_tokamax_gmm:
+      raise ValueError("GMM v2 requires `use_tokamax_gmm=true`.")
+
     for val in self.compress_ratios:
       if val != 0 and val < 4:
         raise ValueError(f"compress_ratio must be 0 (disabled) or >= 4, got {val}")
diff --git a/src/maxtext/kernels/megablox/ops.py b/src/maxtext/kernels/megablox/ops.py
index 21045f18c6..b7487c7368 100644
--- a/src/maxtext/kernels/megablox/ops.py
+++ b/src/maxtext/kernels/megablox/ops.py
@@ -22,12 +22,20 @@
 import jax
 import jax.numpy as jnp
 from maxtext.kernels.megablox import backend
+from maxtext.kernels.megablox import pallas_mosaic_tpu_v2_gmm_kernel as gmm_v2
+from maxtext.kernels.megablox import pallas_mosaic_tpu_v2_tgmm_kernel as tgmm_v2
 from maxtext.layers import quantizations
 import qwix
 import qwix.pallas as qpl
 import tokamax
 
 
+DLHS_RAGGED_DOT_DIM_NUMS = jax.lax.RaggedDotDimensionNumbers(
+    dot_dimension_numbers=(([1], [2]), ([], [])),
+    lhs_ragged_dimensions=[0],
+    rhs_group_dimensions=[0],
+)
+
 DRHS_RAGGED_DOT_DIM_NUMS = jax.lax.RaggedDotDimensionNumbers(
     dot_dimension_numbers=(([0], [0]), ([], [])),
     lhs_ragged_dimensions=[0],
@@ -65,6 +73,7 @@ def gmm(
     # TODO(amandaliang): get rid of the qwix_rule in favor of Qwix's interception feature
     qwix_rule: qwix.QtRule | None = None,
     use_manual_quantization: bool = False,  # used in batchsplit
+    use_gmm_v2: bool = False,
 ):
   """Grouped matrix multiplication operation."""
   quantization_rule = None
@@ -84,7 +93,7 @@ def gmm(
       )
 
   gmm_fwd_bwd = lambda *args: _gmm_fwd(*args)[0]  # pylint: disable=C3001
-  gmm_fwd_bwd = jax.custom_vjp(gmm_fwd_bwd, nondiff_argnums=(3, 4, 7, 8, 9, 10, 11, 12, 13, 14))
+  gmm_fwd_bwd = jax.custom_vjp(gmm_fwd_bwd, nondiff_argnums=(3, 4, 7, 8, 9, 10, 11, 12, 13, 14, 15))
   gmm_fwd_bwd.defvjp(_gmm_fwd, functools.partial(_gmm_bwd, lhs.dtype, rhs.dtype))
   return gmm_fwd_bwd(
       lhs,
@@ -102,6 +111,7 @@ def gmm(
       use_manual_quantization,
       lhs_vma_axes,
       rhs_vma_axes,
+      use_gmm_v2,
   )
 
 
@@ -131,6 +141,7 @@ def _gmm_fwd(
     use_manual_quantization: bool = False,
     lhs_vma_axes: tuple = tuple(),
     rhs_vma_axes: tuple = tuple(),
+    use_gmm_v2: bool = False,
 ) -> tuple[
     jnp.ndarray,
     tuple[
@@ -178,6 +189,19 @@ def _gmm_fwd(
     if transpose_rhs:
       rhs = rhs.swapaxes(1, 2)
 
+  if use_gmm_v2:
+    out = gmm_v2.gmm_v2(
+        lhs=lhs,
+        rhs=rhs,
+        group_sizes=group_sizes,
+        tile_info=gmm_v2.TileSizes(
+            tile_m=tiling[0],
+            tile_k=tiling[1],
+            tile_n=tiling[2],
+        ),
+        preferred_element_type=preferred_element_type,
+    )
+  elif use_tokamax_backend:
     # manual_axis_type is for gmm with shard_map check_vma=True, needs tokamax > 0.0.12
     out_kwargs = {}
     if use_manual_quantization:
@@ -226,6 +250,7 @@ def _gmm_bwd(
     use_manual_quantization: bool,
     lhs_vma_axes: tuple,
     rhs_vma_axes: tuple,
+    use_gmm_v2: bool,
     residual: tuple[
         jnp.ndarray | qpl.QArray,
         jnp.ndarray | qpl.QArray,
@@ -274,10 +299,10 @@ def _gmm_bwd(
         channelwise_axes=[] if quantization_rule.disable_channelwise_axes else [1],
         calibration_method=quantization_rule.bwd_calibration_method,
     )
-  if use_tokamax_backend:
+  if use_tokamax_backend or use_gmm_v2:
     # Handle transpose_rhs manually
     dlhs_rhs = rhs
-    if not transpose_rhs:
+    if transpose_rhs:
       dlhs_rhs = dlhs_rhs.swapaxes(1, 2)
 
     # manual_axis_type is for gmm with shard_map check_vma=True, needs tokamax > 0.0.12
@@ -290,29 +315,63 @@ def _gmm_bwd(
           varying=frozenset(["expert"]), unreduced=frozenset(["data", "fsdp"])
       )
 
-    dlhs = tokamax.ragged_dot(
-        lhs=dlhs_dout,
-        rhs=dlhs_rhs,
-        group_sizes=group_sizes,
-        precision=jax.lax.Precision.DEFAULT,
-        preferred_element_type=lhs_dtype,
-        # `group_offset` is not yet supported
-        group_offset=None,
-        implementation="mosaic",
-        **dlhs_kwargs,
-    )
-    drhs = tokamax.ragged_dot_general(
-        lhs=lhs,
-        rhs=drhs_dout,
-        group_sizes=group_sizes,
-        ragged_dot_dimension_numbers=DRHS_RAGGED_DOT_DIM_NUMS,
-        precision=jax.lax.Precision.DEFAULT,
-        preferred_element_type=rhs_dtype,
-        # `group_offset` is not yet supported
-        group_offset=None,
-        implementation="mosaic",
-        **drhs_kwargs,
-    )
+    if use_gmm_v2:
+      dlhs = gmm_v2.gmm_v2(
+          lhs=dlhs_dout,
+          rhs=dlhs_rhs.swapaxes(1, 2),  # requires rhs to be [g, n, k]
+          group_sizes=group_sizes,
+          tile_info=gmm_v2.TileSizes(
+              tile_m=tiling[3],
+              tile_k=tiling[4],
+              tile_n=tiling[5],
+          ),
+          preferred_element_type=lhs_dtype,
+      )
+
+      # tgmm_v2_op requires lhs and rhs to have the same dtype.
+      if lhs.dtype != drhs_dout.dtype:
+        drhs_dout = drhs_dout.astype(lhs.dtype)
+
+      drhs = tgmm_v2.tgmm_v2(
+          lhs=lhs,
+          rhs=drhs_dout,
+          group_sizes=group_sizes,
+          num_actual_groups=num_actual_groups,
+          precision=jax.lax.Precision.DEFAULT,
+          preferred_element_type=rhs_dtype,
+          group_offset=group_offset,
+          tile_info=gmm_v2.TileSizes(
+              tile_m=tiling[6],
+              tile_k=tiling[7],
+              tile_n=tiling[8],
+          ),
+      )
+    else:
+      dlhs = tokamax.ragged_dot_general(
+          lhs=dlhs_dout,
+          rhs=dlhs_rhs,
+          group_sizes=group_sizes,
+          ragged_dot_dimension_numbers=DLHS_RAGGED_DOT_DIM_NUMS,
+          precision=jax.lax.Precision.DEFAULT,
+          preferred_element_type=lhs_dtype,
+          # `group_offset` is not yet supported
+          group_offset=None,
+          implementation="mosaic",
+          **dlhs_kwargs,
+      )
+
+      drhs = tokamax.ragged_dot_general(
+          lhs=lhs,
+          rhs=drhs_dout,
+          group_sizes=group_sizes,
+          ragged_dot_dimension_numbers=DRHS_RAGGED_DOT_DIM_NUMS,
+          precision=jax.lax.Precision.DEFAULT,
+          preferred_element_type=rhs_dtype,
+          # `group_offset` is not yet supported
+          group_offset=None,
+          implementation="mosaic",
+          **drhs_kwargs,
+      )
     if quantization_rule and quantization_rule.bwd_qtype and weight_gather_axes:
       # Scatter back in reverse order of gather
       for axis_name, axis_idx in reversed(weight_gather_axes):
diff --git a/src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_gmm_kernel.py b/src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_gmm_kernel.py
new file mode 100644
index 0000000000..df706c1db8
--- /dev/null
+++ b/src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_gmm_kernel.py
@@ -0,0 +1,1307 @@
+# Copyright 2026 DeepMind Technologies Limited. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+# Forked from:
+# https://github.com/openxla/tokamax/blob/3f332fcf85dcb87aab661d00228ed71a09b5fd56/
+# tokamax/_src/ops/ragged_dot/pallas_mosaic_tpu_v2_gmm_kernel.py
+"""GMM kernel implemented using Pallas."""
+
+from abc import ABC, abstractmethod
+import dataclasses
+import functools
+from typing import Any, Callable, Tuple
+
+import jax
+from jax import lax
+from jax.experimental import pallas as pl
+from jax.experimental.pallas import tpu as pltpu
+import jax.numpy as jnp
+
+# Util.
+
+
+def swigluoai(gate: jax.Array, up: jax.Array, *, alpha: float = 1.702, limit: float = 7.0) -> jax.Array:
+  """Activation used in some models such as GPT-OSS."""
+
+  gate = jnp.clip(gate, max=limit)
+  up = jnp.clip(up, min=-limit, max=limit)
+  glu = gate * jax.nn.sigmoid(alpha * gate)
+  return (up + 1.0) * glu
+
+
+def apply_act_fn(acc: jax.Array, fuse_act: str | None):
+  """Applies a fused activation function to the accumulator.
+
+  This function is used when an activation function is fused with the matrix
+  multiplication. The input accumulator `acc` is expected to contain
+  concatenated results for both the 'gate' and 'up' projections.
+
+  Args:
+    acc: The accumulator array, with the last dimension being 2 * tile_n.
+    fuse_act: The name of the activation function to apply. Supported values are
+      "silu", "gelu", and "swigluoai". If None, no activation is applied.
+
+  Returns:
+    The result of applying the activation function.
+
+  Raises:
+    NotImplementedError: If an unsupported `fuse_act` is provided.
+  """
+
+  if fuse_act is None:
+    return acc
+
+  acc_gate, acc_up = jnp.split(acc, 2, -1)
+  match fuse_act:
+    case "silu":
+      return jax.nn.silu(acc_gate) * acc_up
+    case "gelu":
+      return jax.nn.gelu(acc_gate) * acc_up
+    case "swigluoai":
+      return swigluoai(acc_gate, acc_up)
+    case _:
+      raise NotImplementedError(f"Unsupported activation function: {fuse_act}")
+
+
+def align_to(x, a):
+  return pl.cdiv(x, a) * a
+
+
+# Define data classes.
+
+
+class RhsRef(ABC):
+  """Abstract class that defines interfaces for rhs values."""
+
+  @abstractmethod
+  def get_weight(self) -> jax.Array:
+    ...
+
+  @abstractmethod
+  def get_scale(self) -> jax.Array:
+    ...
+
+  @abstractmethod
+  def get_bias(self) -> jax.Array:
+    ...
+
+
+@jax.tree_util.register_dataclass
+@dataclasses.dataclass(frozen=True)
+class WeightsRef(RhsRef):
+  """Dataclass for a single weights."""
+
+  weight: Any
+  scale: Any | None
+  bias: Any | None
+
+  def get_weight(self) -> jax.Array:
+    return self.weight[...]
+
+  def get_scale(self) -> jax.Array:
+    assert self.scale is not None
+    return self.scale[...]
+
+  def get_bias(self) -> jax.Array:
+    assert self.bias is not None
+    return self.bias[...]
+
+
+@jax.tree_util.register_dataclass
+@dataclasses.dataclass(frozen=True)
+class FusedWeightsRef(RhsRef):
+  """Dataclass for gate and up weights used in fused activation."""
+
+  gate: WeightsRef
+  up: WeightsRef
+
+  def get_weight(self) -> jax.Array:
+    w_gate = self.gate.get_weight()
+    w_up = self.up.get_weight()
+    return jnp.concatenate([w_gate, w_up], axis=-1)
+
+  def get_scale(self) -> jax.Array:
+    s_gate = self.gate.get_scale()
+    s_up = self.up.get_scale()
+    return jnp.concatenate([s_gate, s_up], axis=-1)
+
+  def get_bias(self) -> jax.Array:
+    b_gate = self.gate.get_bias()
+    b_up = self.up.get_bias()
+    return jnp.concatenate([b_gate, b_up], axis=-1)
+
+
+@jax.tree_util.register_dataclass
+@dataclasses.dataclass(frozen=True)
+class MetadataRef:
+  gm_id_to_group_id: jax.Array
+  gm_id_to_m_offset: jax.Array
+
+
+@dataclasses.dataclass(frozen=True)
+class TileSizes:
+  tile_m: int
+  tile_k: int
+  tile_n: int
+
+
+@dataclasses.dataclass(frozen=True)
+class Dimensions:
+  size_m: int
+  size_k: int
+  size_n: int
+  size_group: int
+  size_lhs_group: int
+  size_lhs_sublane: int
+
+
+@dataclasses.dataclass(frozen=True)
+class InputConfigs:
+  """Configuration parameters for input tensors."""
+
+  quant_dtype: jnp.dtype | None
+  quant_block_size: int | None
+  dtype: jnp.dtype
+  has_bias: bool = False
+  has_scale: bool = False
+
+  @property
+  def should_bitcast(self) -> bool:
+    bits = jax.dtypes.itemsize_bits(self.dtype)
+    return bits < 8
+
+  @property
+  def should_dequantize_before_matmul(self) -> bool:
+    if not self.has_scale:
+      return False
+    assert self.quant_block_size is not None
+    mxu_size = pltpu.get_tpu_info().mxu_column_size
+    return self.quant_block_size < mxu_size
+
+  @property
+  def should_dequantize_after_matmul(self) -> bool:
+    return self.has_scale and not self.should_dequantize_before_matmul
+
+
+@dataclasses.dataclass(frozen=True)
+class GmmConfigs:
+  """Full configuration details for GMM execution."""
+
+  tiles: TileSizes
+  dims: Dimensions
+  lhs_cfgs: InputConfigs
+  rhs_cfgs: InputConfigs
+  out_dtype: jnp.dtype
+  acc_dtype: jnp.dtype
+  zero_init: bool
+  fuse_act: str | None
+
+  @property
+  def num_quant_blocks_per_tile_k(self) -> int:
+    return pl.cdiv(self.tiles.tile_k, self.rhs_cfgs.quant_block_size)
+
+  @property
+  def out_size_n(self) -> int:
+    if self.fuse_act is None:
+      return self.dims.size_n
+    else:
+      return self.dims.size_n // 2
+
+
+TileFn = Callable[[Dimensions, InputConfigs, InputConfigs, int, str | None], TileSizes]
+
+
+class IndexMaps:
+  """Index maps for GMM kernel."""
+
+  def __init__(self, metadata_ref: MetadataRef, cfgs: GmmConfigs):
+    self.metadata_ref = metadata_ref
+    self.cfgs = cfgs
+
+  def lhs_index_map(self, _: jax.Array, gm_id: jax.Array, k_id: jax.Array):
+    m_start = self.metadata_ref.gm_id_to_m_offset[gm_id]
+    m_end = self.metadata_ref.gm_id_to_m_offset[gm_id + 1]
+
+    row_start = m_start // self.cfgs.dims.size_lhs_sublane
+    row_end = pl.cdiv(m_end, self.cfgs.dims.size_lhs_sublane)
+    row_size = row_end - row_start
+
+    return (pl.ds(row_start, row_size), 0, k_id)
+
+  def rhs_weight_index_map(self, n_id: jax.Array, gm_id: jax.Array, k_id: jax.Array):
+    group_id = self.metadata_ref.gm_id_to_group_id[gm_id]
+    return (group_id, k_id, n_id)
+
+  def rhs_bias_index_map(self, n_id: jax.Array, gm_id: jax.Array, _: jax.Array):
+    group_id = self.metadata_ref.gm_id_to_group_id[gm_id]
+    return (group_id, 0, n_id)
+
+  def rhs_scale_index_map(self, n_id: jax.Array, gm_id: jax.Array, k_id: jax.Array):
+    group_id = self.metadata_ref.gm_id_to_group_id[gm_id]
+    # Simply multiplying k_id by num_quant_blocks_per_tile_k will not work
+    # since a single quant block could be shared along multiple k tile.
+    k_row = k_id * self.cfgs.tiles.tile_k
+    b_row = k_row // self.cfgs.rhs_cfgs.quant_block_size
+    b_tile_id = b_row // self.cfgs.num_quant_blocks_per_tile_k
+    return (group_id, b_tile_id, 0, n_id)
+
+  def out_index_map(self, n_id: jax.Array, gm_id: jax.Array, _: jax.Array):
+    """Calculates index map for the output tensor."""
+    is_last_gm = gm_id == (pl.num_programs(1) - 1)
+    m_start = self.metadata_ref.gm_id_to_m_offset[gm_id]
+    m_end = self.metadata_ref.gm_id_to_m_offset[gm_id + 1]
+
+    row_start = m_start // self.cfgs.dims.size_lhs_sublane
+    capped_row_end = m_end // self.cfgs.dims.size_lhs_sublane
+    last_row_end = pl.cdiv(m_end, self.cfgs.dims.size_lhs_sublane)
+    row_end = jnp.where(is_last_gm, last_row_end, capped_row_end)
+    row_size = row_end - row_start
+
+    return (pl.ds(row_start, row_size), 0, n_id)
+
+
+def generate_block_specs(
+    metadata_ref: MetadataRef, cfgs: GmmConfigs
+) -> Tuple[Tuple[pl.BlockSpec, WeightsRef], pl.BlockSpec]:
+  """Generates block specs for the given lhs, rhs, and out refs."""
+
+  index_map = IndexMaps(metadata_ref, cfgs)
+  bounded_slice_gm = pl.BoundedSlice(cfgs.tiles.tile_m // cfgs.dims.size_lhs_sublane)
+
+  lhs_block_spec = pl.BlockSpec(
+      (bounded_slice_gm, cfgs.dims.size_lhs_sublane, cfgs.tiles.tile_k),
+      index_map.lhs_index_map,
+  )
+
+  tile_k_rhs = cfgs.tiles.tile_k
+  if cfgs.rhs_cfgs.should_bitcast:
+    packing = pl.cdiv(32, jax.dtypes.itemsize_bits(cfgs.rhs_cfgs.dtype))
+    tile_k_rhs //= packing
+
+  rhs_weight_spec = pl.BlockSpec(
+      (None, tile_k_rhs, cfgs.tiles.tile_n),
+      index_map.rhs_weight_index_map,
+      pipeline_mode=pl.Buffered(buffer_count=3),
+  )
+  rhs_scale_block_spec = rhs_bias_block_spec = None
+  if cfgs.rhs_cfgs.has_bias:
+    rhs_bias_block_spec = pl.BlockSpec(
+        (None, 1, cfgs.tiles.tile_n),
+        index_map.rhs_bias_index_map,
+    )
+  if cfgs.rhs_cfgs.has_scale:
+    rhs_scale_block_spec = pl.BlockSpec(
+        (None, cfgs.num_quant_blocks_per_tile_k, 1, cfgs.tiles.tile_n),
+        index_map.rhs_scale_index_map,
+    )
+
+  rhs_block_spec = WeightsRef(
+      weight=rhs_weight_spec,
+      scale=rhs_scale_block_spec,
+      bias=rhs_bias_block_spec,
+  )
+
+  out_block_spec = pl.BlockSpec(
+      (bounded_slice_gm, cfgs.dims.size_lhs_sublane, cfgs.tiles.tile_n),
+      index_map.out_index_map,
+  )
+
+  return (lhs_block_spec, rhs_block_spec), out_block_spec
+
+
+# Define kernels.
+
+
+def inner_kernel(
+    # In
+    tiled_lhs_ref: jax.Array,
+    # [tile_m // size_lhs_sublane, size_lhs_sublane, tile_k]
+    tiled_rhs_ref: RhsRef,  # [tile_k, tile_n]
+    # Out
+    tiled_out_ref: jax.Array,
+    # [tile_m // size_lhs_sublane, size_lhs_sublane, tile_n]
+    # Scratch
+    partial_out_ref: jax.Array,  # [size_lhs_sublane, tile_n]
+    acc_ref: jax.Array,  # [tile_m, tile_n]
+    metadata_ref: MetadataRef,
+    *,
+    cfgs: GmmConfigs,
+):
+  """Inner kernel invoked by emit_pipeline to perform matmul.
+
+  tiled_lhs_ref and tiled_out_ref points to rows [m_start:m_end] of lhs and out.
+  Additionally, m_start and m_end does not have to align with tile boundaries
+  [m_offset:m_offset+tile_m]. Therefore, rows [m_offset:m_start] and
+  [m_end:m_offset+tile_m] of tiled_lhs_ref and tiled_out_ref will contain
+  invalid data and needs to be masked out.
+
+  Args:
+    tiled_lhs_ref: Contains value lhs[m_start:m_end, k_start:k_end]
+    tiled_rhs_ref: Contains value rhs[g_id, k_start:k_end, n_start:n_end]. where
+      g_id is the group associated with lhs[m_start:m_end, :]
+    tiled_out_ref: Contains value out[m_start:m_end, n_start:n_end]
+    partial_out_ref: Contains last size_lhs_sublane rows of the previous output.
+      Will be initialized to zero if this is first tile for grid[n_id, :, :].
+    acc_ref: Reference to the accumulator.
+    metadata_ref: Reference to the metadata.
+    cfgs: GmmConfigs.
+  """
+
+  def _matmul(is_first_k_step: bool, is_last_k_step: bool):
+    tpu_info = pltpu.get_tpu_info()
+    mxu_size = tpu_info.mxu_column_size
+
+    # Step 1: Input pre-processing.
+    tiled_lhs = tiled_lhs_ref.reshape(-1, cfgs.tiles.tile_k)[...]
+    tiled_rhs = tiled_rhs_ref.get_weight()
+    # When rhs is packed (quantized dtype packed into uint32), unpack it
+    # back to the original dtype using pltpu.bitcast which operates on K
+    # axis. This expands the K dimension back to tile_k.
+    if cfgs.rhs_cfgs.should_bitcast:
+      tiled_rhs = pltpu.bitcast(tiled_rhs, cfgs.rhs_cfgs.dtype)
+    rhs_tile_n = tiled_rhs.shape[1]
+
+    # This should only be taken in the case where we don't requantize
+    # the scales and thus we need to dequantize inside VMEM to avoid small
+    # contracting dimensions
+    if cfgs.rhs_cfgs.should_dequantize_before_matmul:
+      rhs_qbs = cfgs.rhs_cfgs.quant_block_size
+      tiled_rhs_scale = tiled_rhs_ref.get_scale().astype(acc_ref.dtype)
+      num_blocks = cfgs.num_quant_blocks_per_tile_k
+      tiled_rhs_dequant = tiled_rhs.astype(acc_ref.dtype).reshape(num_blocks, rhs_qbs, rhs_tile_n)
+      tiled_rhs_dequant = tiled_rhs_dequant * tiled_rhs_scale
+      tiled_rhs = tiled_rhs_dequant.reshape(cfgs.tiles.tile_k, rhs_tile_n)
+
+    valid_k = cfgs.dims.size_k % cfgs.tiles.tile_k
+    if is_last_k_step and valid_k != 0:
+      mask_rhs = lax.broadcasted_iota(jnp.int32, tiled_rhs.shape, 0) < valid_k
+      tiled_rhs = jnp.where(mask_rhs, tiled_rhs, 0)
+
+    # Step 2: Matmul.
+    acc_list = []
+    if cfgs.lhs_cfgs.quant_dtype is None:
+      # Unquantized matmul path.
+      rhs_qbs = cfgs.rhs_cfgs.quant_block_size
+
+      for start_n in range(0, rhs_tile_n, mxu_size):
+        end_n = min(rhs_tile_n, start_n + mxu_size)
+        col_size = end_n - start_n
+
+        acc_n = jnp.zeros((cfgs.tiles.tile_m, col_size), dtype=acc_ref.dtype)
+        for b_id in range(cfgs.num_quant_blocks_per_tile_k):
+          start_k = b_id * rhs_qbs
+          end_k = start_k + rhs_qbs
+
+          block_acc = jnp.matmul(
+              tiled_lhs[:, start_k:end_k],
+              tiled_rhs[start_k:end_k, start_n:end_n],
+              preferred_element_type=jnp.float32,
+          ).astype(acc_ref.dtype)
+
+          if cfgs.rhs_cfgs.should_dequantize_after_matmul:
+            tiled_rhs_scale = tiled_rhs_ref.get_scale()
+            block_acc *= tiled_rhs_scale[b_id, :, start_n:end_n].astype(acc_ref.dtype)
+
+          acc_n += block_acc
+        acc_list.append(acc_n)
+    else:
+      # Quantized matmul path.
+      lhs_q_dtype = cfgs.lhs_cfgs.quant_dtype
+      q_block_size = cfgs.lhs_cfgs.quant_block_size
+
+      if jnp.issubdtype(lhs_q_dtype, jnp.floating):
+        dtype_max = float(jnp.finfo(lhs_q_dtype).max)
+        preferred_element_type = jnp.float32
+      else:
+        dtype_max = float(jnp.iinfo(lhs_q_dtype).max)
+        preferred_element_type = jnp.int32
+
+      # Without n outer loop, result of quantized matmul becomes available only
+      # at the last iteration of the loop. This means [tile_m, tile_n] value
+      # needs to be stored until the last iteration. By adding n outer loop,
+      # result of [tile_m, mxu_size] becomes available at the end of every k
+      # inner loop which can be used to pipeline subsequent VPU or VST ops with
+      # MXU ops for the next [tile_m, mxu_size].
+      for start_n in range(0, rhs_tile_n, mxu_size):
+        end_n = min(rhs_tile_n, start_n + mxu_size)
+        col_size = end_n - start_n
+
+        acc_n = jnp.zeros((cfgs.tiles.tile_m, col_size), dtype=acc_ref.dtype)
+        for start_k in range(0, cfgs.tiles.tile_k, q_block_size):
+          end_k = min(cfgs.tiles.tile_k, start_k + q_block_size)
+
+          block_lhs = tiled_lhs[:, start_k:end_k]
+          block_rhs = tiled_rhs[start_k:end_k, start_n:end_n]
+
+          # Perform lhs quantization. Note that for every block_lhs,
+          # same computation will be performed tiles_n//mxu_size times.
+          # But we can let compiler perform CSE and avoid recomputation.
+          block_abs_max = jnp.max(jnp.abs(block_lhs), axis=1, keepdims=True)
+          block_scale = block_abs_max / dtype_max
+
+          # If block_scale=0, it will cause division by zero and return either
+          # NaN or Inf. Since this can cause numeric issue when downcasting to
+          # quantized value, we convert them into 0.
+          block_scale_inv = jnp.where(block_scale == 0, 0, 1 / block_scale)
+          # Convert lhs into quantized dtype.
+          block_lhs_q = (block_lhs * block_scale_inv).astype(lhs_q_dtype)
+
+          # Unlike unquantized path, compiler may not perform implicit type
+          # conversion due to numeric concerns. As this can cause unsupported
+          # matmul error, explicit type conversion is performed.
+          if not tpu_info.is_matmul_supported(lhs_q_dtype, block_rhs.dtype):
+            block_rhs = block_rhs.astype(lhs_q_dtype)
+
+          block_acc = jnp.matmul(
+              block_lhs_q,
+              block_rhs,
+              preferred_element_type=preferred_element_type,
+          ).astype(acc_ref.dtype)
+
+          block_acc *= block_scale.astype(acc_ref.dtype)
+
+          # Apply rhs subchannel scale per quant block.
+          if cfgs.rhs_cfgs.should_dequantize_after_matmul:
+            b_id = start_k // cfgs.rhs_cfgs.quant_block_size
+            rhs_scale_slice = tiled_rhs_ref.get_scale()
+            block_acc *= rhs_scale_slice[b_id, :, start_n:end_n].astype(acc_ref.dtype)
+
+          acc_n += block_acc
+        acc_list.append(acc_n)
+    acc = jnp.concatenate(acc_list, axis=1)
+
+    # Step 3: Output post-processing.
+    if not is_first_k_step:
+      acc += acc_ref[...]
+
+    if is_last_k_step:
+      if cfgs.rhs_cfgs.has_bias:
+        tiled_rhs_bias = tiled_rhs_ref.get_bias()
+        acc += tiled_rhs_bias.astype(acc.dtype)
+
+      acc = apply_act_fn(acc, cfgs.fuse_act)
+
+      gm_id = pl.program_id(1)
+
+      # Mask out rows that does not belong to the current group.
+      m_start = metadata_ref.gm_id_to_m_offset[gm_id]
+      m_end = metadata_ref.gm_id_to_m_offset[gm_id + 1]
+      m_offset = m_start - m_start % cfgs.dims.size_lhs_sublane
+
+      m_start_local = m_start - m_offset
+      m_end_local = m_end - m_offset
+
+      iota = lax.broadcasted_iota(jnp.int32, acc.shape, 0)
+      mask = jnp.logical_and(m_start_local <= iota, iota < m_end_local)
+      acc_masked = jnp.where(mask, acc, 0).reshape(tiled_out_ref.shape)
+
+      # Write the final output to the output ref.
+      tiled_out_ref[...] = acc_masked.astype(tiled_out_ref.dtype)
+
+      # If this is the first tile for grid[n_id, :, :], we initialize the
+      # partial out to zeros. Otherwise, partial out from last tile of
+      # grid[n_id-1, :, :] can be used and cause numeric issues.
+      partial_out_zeros = jnp.zeros_like(partial_out_ref)
+
+      # Accumulate the partial output from the previous step.
+      tiled_out_ref[0] += jnp.where(gm_id == 0, partial_out_zeros, partial_out_ref[...])
+
+      # Consider following case where size_lhs_sublane = 4, number denotes group
+      # id and | denotes boundaries between sublanes:
+      # | 0 0 1 2 | 2 2 2 2 | 3 3 4 4 |
+      #
+      # Assuming group id of current step is 1, current step will not completely
+      # fill size_lhs_sublane rows and will be revisited at the next step. By
+      # storing the partial rows into the partial_out_ref, the next step can
+      # read them and accumulate to them.  Additionally, for group id of 2,
+      # since it completely fills the size_lhs_sublane rows, we need to zero out
+      # partial_out_ref to avoid numeric error for group 3.
+      last_row = m_end_local // cfgs.dims.size_lhs_sublane
+      partial_out_ref[...] = jnp.where(
+          m_end_local % cfgs.dims.size_lhs_sublane == 0,
+          partial_out_zeros,
+          tiled_out_ref[last_row],
+      )
+    else:
+      acc_ref[...] = acc
+
+  # Define matmul wrapper functions.
+  @jax.named_scope("matmul_first_last")
+  def matmul_first_last():
+    _matmul(is_first_k_step=True, is_last_k_step=True)
+
+  @jax.named_scope("matmul_first")
+  def matmul_first():
+    _matmul(is_first_k_step=True, is_last_k_step=False)
+
+  @jax.named_scope("matmul")
+  def matmul():
+    _matmul(is_first_k_step=False, is_last_k_step=False)
+
+  @jax.named_scope("matmul_last")
+  def matmul_last():
+    _matmul(is_first_k_step=False, is_last_k_step=True)
+
+  # Select and execute matmul function based on the current step.
+  num_k = pl.num_programs(2)
+  k_id = pl.program_id(2)
+
+  is_first_k_step = k_id == 0
+  is_last_k_step = k_id == (num_k - 1)
+
+  lax.cond(
+      is_first_k_step,
+      lambda: lax.cond(
+          is_last_k_step,
+          matmul_first_last,
+          matmul_first,
+      ),
+      lambda: lax.cond(
+          is_last_k_step,
+          matmul_last,
+          matmul,
+      ),
+  )
+
+
+def fill_metadata(
+    lhs_group_sizes_ref: jax.Array,  # int32[size_lhs_group]
+    group_offset_ref: jax.Array,  # int32[1]
+    metadata_ref: MetadataRef,
+    *,
+    cfgs: GmmConfigs,
+) -> jax.Array:
+  """Fills the metadata for the given lhs group sizes and group offset.
+
+  Iterates over the lhs group sizes and if the group id is valid, determines
+  the number of gm tiles that are needed to process the current group. Then,
+  it fills starting and ending offset (gm_id_to_m_offset), and the group id
+  (gm_id_to_group_id) for each gm tile.
+
+  Args:
+    lhs_group_sizes_ref: The group sizes of lhs.
+    group_offset_ref: Offset of the first group to process.
+    metadata_ref: Metadata that is used to determine the group id and m offsets
+      for each gmm tile.
+    cfgs: GmmConfigs.
+
+  Returns:
+      The number of gm tiles to process lhs with given group offset.
+  """
+
+  group_offset = group_offset_ref[0]
+  max_num_group = group_offset + cfgs.dims.size_group
+  metadata_ref.gm_id_to_m_offset[0] = 0
+
+  @jax.named_scope("inner_tm_loop")
+  def inner_tm_loop(tm_id, curr_m_offset, *, end_m_offset, group_id):
+    local_offset = curr_m_offset % cfgs.dims.size_lhs_sublane
+    tm_size = jnp.minimum(cfgs.tiles.tile_m - local_offset, end_m_offset - curr_m_offset)
+
+    metadata_ref.gm_id_to_group_id[tm_id] = group_id
+
+    next_m_offset = curr_m_offset + tm_size
+    metadata_ref.gm_id_to_m_offset[tm_id] = curr_m_offset
+    metadata_ref.gm_id_to_m_offset[tm_id + 1] = next_m_offset
+
+    return next_m_offset
+
+  @jax.named_scope("outer_group_loop")
+  def outer_group_loop(lhs_group_id, carry):
+    num_gm, start_m_offset = carry
+
+    group_id = lhs_group_id - group_offset
+    group_size = lhs_group_sizes_ref[lhs_group_id]
+    end_m_offset = start_m_offset + group_size
+
+    # Assume following arguments:
+    # - size_lhs_sublane & tile_m = 4
+    # - group_size = 3
+    # - start_m_offset = 7
+    #
+    # If we visualize it, it will look like this where:
+    # - |: denotes boundaries between sublanes
+    # - 0: denotes values for other groups
+    # - 1: denotes values for the current group
+    # | 0 0 0 0 | 0 0 0 1 | 1 1 0 0 |
+    #
+    # In this example, we see that we require processing 2 m tiles.
+    # But, performing a naive cdiv(group_size, tile_m) will return 1.
+    # Instead, adding local_offset will give us the correct value.
+    local_offset = start_m_offset % cfgs.dims.size_lhs_sublane
+    aligned_group_size = group_size + local_offset
+    curr_num_gm = pl.cdiv(aligned_group_size, cfgs.tiles.tile_m)
+
+    # We need to handle cases where we should not process the group.
+    # 1. Even if group_size is 0, if local_offset is not 0, cdiv will return 1.
+    # 2. If group comes before the group_offset, we should not process it.
+    should_process = jnp.logical_and(group_size > 0, group_id >= 0)
+    curr_num_gm = jnp.where(should_process, curr_num_gm, 0)
+    next_num_gm = num_gm + curr_num_gm
+
+    tm_loop_fn = functools.partial(
+        inner_tm_loop,
+        end_m_offset=end_m_offset,
+        group_id=group_id,
+    )
+    lax.fori_loop(num_gm, next_num_gm, tm_loop_fn, start_m_offset)
+
+    return next_num_gm, end_m_offset
+
+  num_gm, _ = lax.fori_loop(0, max_num_group, outer_group_loop, (0, 0))
+  return num_gm
+
+
+def zero_out_start(
+    out_ref: jax.Array,  # [size_m, size_n]
+    zero_ref: jax.Array,  # [tile_zero_m, num_lanes]
+    semaphore_ref: jax.Array,  # [1]
+    metadata_ref: MetadataRef,
+    num_gm: jax.Array,
+    *,
+    dims: Dimensions,
+):
+  """Zero out output rows that are not used in the computation."""
+
+  num_lanes = pltpu.get_tpu_info().num_lanes
+  assert num_lanes == zero_ref.shape[-1]
+  zero_ref[...] = jnp.zeros_like(zero_ref)
+
+  zero_dma = zero_ref.reshape(-1, dims.size_lhs_sublane, num_lanes)
+  out_dma = out_ref.reshape(-1, dims.size_lhs_sublane, out_ref.shape[-1])
+  row_size = zero_dma.shape[0]
+
+  compute_start = metadata_ref.gm_id_to_m_offset[0]
+  compute_end = metadata_ref.gm_id_to_m_offset[num_gm]
+
+  left_zero_start = 0
+  left_zero_end = compute_start // dims.size_lhs_sublane
+  left_zero_size = left_zero_end - left_zero_start
+  left_num_loops = pl.cdiv(left_zero_size, row_size)
+
+  right_zero_start = pl.cdiv(compute_end, dims.size_lhs_sublane)
+  right_zero_end = out_dma.shape[0]
+  right_zero_size = right_zero_end - right_zero_start
+  right_num_loops = pl.cdiv(right_zero_size, row_size)
+
+  def fill_zero(i, zero_size, *, start, end):
+    dma_start = start + i * row_size
+    dma_end = jnp.minimum(dma_start + row_size, end)
+    dma_size = dma_end - dma_start
+
+    # Static loop. Will be unrolled during compile time.
+    for n_start in range(0, out_dma.shape[-1], num_lanes):
+      n_end = n_start + num_lanes
+      pltpu.make_async_copy(
+          src_ref=zero_dma.at[pl.ds(0, dma_size)],
+          dst_ref=out_dma.at[pl.ds(dma_start, dma_size), :, n_start:n_end],
+          sem=semaphore_ref.at[0],
+      ).start(priority=1)
+
+    return zero_size + dma_size
+
+  @jax.named_scope("left_fill_zero")
+  def left_fill_zero(i, zero_size):
+    return fill_zero(i, zero_size, start=left_zero_start, end=left_zero_end)
+
+  @jax.named_scope("right_fill_zero")
+  def right_fill_zero(i, zero_size):
+    return fill_zero(i, zero_size, start=right_zero_start, end=right_zero_end)
+
+  zero_size = lax.fori_loop(0, left_num_loops, left_fill_zero, 0)
+  zero_size = lax.fori_loop(0, right_num_loops, right_fill_zero, zero_size)
+  return zero_size
+
+
+def zero_out_end(
+    out_ref: jax.Array,  # [size_m, size_n]
+    semaphore_ref: jax.Array,  # [1]
+    zero_size: jax.Array,
+    *,
+    dims: Dimensions,
+):
+  out_dma = out_ref.reshape(-1, dims.size_lhs_sublane, out_ref.shape[-1])
+  pltpu.make_async_copy(
+      src_ref=out_dma.at[pl.ds(0, zero_size)],
+      dst_ref=out_dma.at[pl.ds(0, zero_size)],
+      sem=semaphore_ref.at[0],
+  ).wait()
+
+
+def kernel_main(
+    # Scalar prefetch
+    lhs_group_sizes_ref: jax.Array,  # int32[size_lhs_group]
+    group_offset_ref: jax.Array,  # int32[1]
+    # In
+    lhs_ref: jax.Array,  # [size_m, size_k]
+    rhs_ref: WeightsRef,  # [size_group, size_k, size_n]
+    # Out
+    out_ref: jax.Array,  # [size_m, size_n]
+    # Scratch memory
+    partial_out_ref: jax.Array,  # [size_lhs_sublane, tile_n]
+    acc_ref: jax.Array,  # [tile_m, tile_n]
+    metadata_ref: MetadataRef,
+    zero_ref: jax.Array | None,  # [tile_zero_m, num_lanes]
+    semaphore_ref: jax.Array | None,  # [1]
+    *,
+    cfgs: GmmConfigs,
+):
+  """Entry point for GMM kernel.
+
+  Computes metadata to determine which rows of lhs needs processing and how
+  they will be tiled. And then, invoke inner kernel using metadata.
+
+  Uses the following notation:
+  - g: rhs group dimension
+  - m: Batch dimension
+  - gm: Batch tiling dimension. Aligned to size_lhs_sublane and has tile size
+    of tile_m. Skips over empty groups and accounts for revisited tiles.
+  - k: in dimension
+  - n: out dimension
+
+  Args:
+    lhs_group_sizes_ref: Reference to the group sizes of lhs.
+    group_offset_ref: Reference to the group offset.
+    lhs_ref: Reference to the lhs.
+    rhs_ref: Reference to the rhs.
+    out_ref: Reference to the out.
+    partial_out_ref: Reference to the partial output.
+    acc_ref: Reference to the accumulator.
+    metadata_ref: Reference to the metadata.
+    zero_ref: Scratch memory for storing zero values used in initialization.
+    semaphore_ref: Semaphore for zero initialization DMAs.
+    cfgs: GmmConfigs.
+  """
+
+  num_k = pl.cdiv(cfgs.dims.size_k, cfgs.tiles.tile_k)
+  num_n = pl.cdiv(cfgs.out_size_n, cfgs.tiles.tile_n)
+
+  # Pack along K (2nd minor dim) so that pltpu.bitcast can unpack inside the
+  # kernel.
+  # [G, K, N] -> [G, K//packing, N] uint32
+  if cfgs.rhs_cfgs.should_bitcast:
+    rhs_weight = rhs_ref.weight.bitcast(jnp.uint32)
+    rhs_ref = dataclasses.replace(rhs_ref, weight=rhs_weight)
+
+  # Fill metadata buffer and return number of group & m iterations.
+  num_gm = fill_metadata(
+      lhs_group_sizes_ref,
+      group_offset_ref,
+      metadata_ref,
+      cfgs=cfgs,
+  )
+
+  if cfgs.zero_init:
+    zero_size = zero_out_start(
+        out_ref,
+        zero_ref,
+        semaphore_ref,
+        metadata_ref,
+        num_gm,
+        dims=cfgs.dims,
+    )
+
+  (lhs_spec, rhs_spec), out_spec = generate_block_specs(metadata_ref, cfgs)
+
+  if cfgs.fuse_act is not None:
+    rhs_up_ref = jax.tree.map(lambda x: x.at[..., cfgs.out_size_n :], rhs_ref)
+    rhs_ref = FusedWeightsRef(gate=rhs_ref, up=rhs_up_ref)
+
+    rhs_spec = FusedWeightsRef(
+        gate=rhs_spec,
+        up=rhs_spec,
+    )
+
+  # Execute the inner kernel.
+  pipeline_fn = pltpu.emit_pipeline(
+      functools.partial(inner_kernel, cfgs=cfgs),
+      grid=(num_n, num_gm, num_k),
+      in_specs=(lhs_spec, rhs_spec),
+      out_specs=out_spec,
+  )
+
+  # Bounded slice requires second last dim to be aligned to the sublane size.
+  # rhs_ref uses static tiling thus reshape is not needed.
+  lhs_in = lhs_ref.reshape(-1, cfgs.dims.size_lhs_sublane, lhs_ref.shape[-1])
+  out_in = out_ref.reshape(-1, cfgs.dims.size_lhs_sublane, out_ref.shape[-1])
+  scratches = [partial_out_ref, acc_ref, metadata_ref]
+  pipeline_fn(lhs_in, rhs_ref, out_in, scratches=scratches)
+
+  if cfgs.zero_init:
+    zero_out_end(out_ref, semaphore_ref, zero_size, dims=cfgs.dims)
+
+
+def calculate_tiling(
+    dims: Dimensions,
+    lhs_cfgs: InputConfigs,
+    rhs_cfgs: InputConfigs,
+    vmem_limit_bytes: int,
+    fuse_act: str | None = None,
+) -> TileSizes:
+  """Calculate optimal tile sizes for GMM kernel."""
+
+  lhs_dtype = lhs_cfgs.quant_dtype or lhs_cfgs.dtype
+  rhs_dtype = rhs_cfgs.dtype
+
+  lhs_bits = jax.dtypes.itemsize_bits(lhs_dtype)
+  rhs_bits = jax.dtypes.itemsize_bits(rhs_dtype)
+
+  # When using bf16 for lhs and rhs, 128 is the largest tile_m value that is
+  # safe to use for most scenarios. But if lower bitwidth is used, we need
+  # to tweak tile_m to account for using faster hardware unit.
+  # TODO: Account for different TPU hardware specs.
+  bf16_bf16_tile_m = 128
+  lhs_mod = min(pl.cdiv(16, lhs_bits), 2)
+  rhs_mod = min(pl.cdiv(16, rhs_bits), 2)
+  tile_m = bf16_bf16_tile_m * lhs_mod // rhs_mod
+  tile_m = min(tile_m, dims.size_m)
+
+  # To avoid stalling MXU, we add some buffer room where tile_n cannot go
+  # smaller than 2x of mxu_column_size.
+  tile_n_limit = pltpu.get_tpu_info().mxu_column_size * 2
+  tile_n_limit = min(tile_n_limit, dims.size_n)
+
+  size_n_per_rhs = dims.size_n
+  fuse_act_factor = 1
+  if fuse_act is not None:
+    # When computing activation function, rhs is concatenated along dim n.
+    fuse_act_factor = 2
+    size_n_per_rhs //= fuse_act_factor
+    tile_n_limit //= fuse_act_factor
+
+  def _is_tile_k_quant_block_compatible(tk: int) -> bool:
+    if tk % rhs_cfgs.quant_block_size != 0 and rhs_cfgs.quant_block_size % tk != 0:
+      return False
+    return True
+
+  # Initialize tile_k and tile_n to their maximum valid values.
+  num_k_tiles = num_n_tiles = 1
+  num_lanes = pltpu.get_tpu_info().num_lanes
+  tile_k = align_to(dims.size_k, num_lanes)
+  tile_n = align_to(size_n_per_rhs, num_lanes)
+
+  def _gmm_vmem_estimate(tn: int, tk: int) -> int:
+    # 1. LHS tile (double-buffered)
+    lhs_tile_bytes = lhs_bits // 8
+    lhs_vmem = 2 * tile_m * tk * lhs_tile_bytes
+
+    # 2. RHS tile (triple-buffered, includes scale and bias if present)
+    # If fuse_act is enabled, we have both gate and up weights,
+    # so RHS memory is doubled.
+    rhs_weight_vmem = tk * tn * rhs_bits // 8
+    rhs_scale_vmem = 0
+    if rhs_cfgs.has_scale and rhs_cfgs.quant_block_size is not None:
+      num_quant_blocks_per_tile_k = pl.cdiv(tk, rhs_cfgs.quant_block_size)
+      rhs_scale_vmem = num_quant_blocks_per_tile_k * tn * 4
+    rhs_bias_vmem = 0
+    if rhs_cfgs.has_bias:
+      rhs_bias_vmem = tn * 4
+    rhs_vmem = fuse_act_factor * (3 * rhs_weight_vmem + 2 * rhs_scale_vmem + 2 * rhs_bias_vmem)
+
+    # 3. Accumulator
+    acc_cols = fuse_act_factor * tn
+    acc_dtype_bytes = 2 if lhs_cfgs.quant_dtype is not None else 4
+    acc_vmem = tile_m * acc_cols * acc_dtype_bytes
+
+    # 4. Output tile (double-buffered)
+    out_dtype_bytes = jax.dtypes.itemsize_bits(lhs_cfgs.dtype) // 8
+    out_vmem = 2 * tile_m * tn * out_dtype_bytes
+
+    return lhs_vmem + rhs_vmem + acc_vmem + out_vmem
+
+  # Multiple k tiles will introduce accumulation overhead. Thus, we first try
+  # to fit the tensors into vmem by only adjusting tile_n.
+
+  # Decrease tile_n until total memory fits in vmem limit.
+  while _gmm_vmem_estimate(tile_n, tile_k) > vmem_limit_bytes and tile_n > tile_n_limit:
+    num_n_tiles += 1
+    tile_n = align_to(size_n_per_rhs, num_n_tiles * num_lanes) // num_n_tiles
+
+  # If decreasing tile_n is no longer possible, we decrease tile_k instead.
+  if tile_n < tile_n_limit:
+    num_n_tiles -= 1
+    tile_n = align_to(size_n_per_rhs, num_n_tiles * num_lanes) // num_n_tiles
+
+    # Decrease tile_k until total memory fits in vmem limit and tile_k is valid.
+    while _gmm_vmem_estimate(tile_n, tile_k) > vmem_limit_bytes or not _is_tile_k_quant_block_compatible(tile_k):
+      num_k_tiles += 1
+      tile_k = align_to(dims.size_k, num_k_tiles * num_lanes) // num_k_tiles
+
+  if tile_n == 0 or tile_k == 0:
+    final_estimate = _gmm_vmem_estimate(tile_n, tile_k)
+    raise ValueError(
+        f"Could not find valid tile sizes for {dims=} and" f" {final_estimate=} (limit: {vmem_limit_bytes})."
+    )
+
+  return TileSizes(tile_m=tile_m, tile_k=tile_k, tile_n=tile_n)
+
+
+def validate_inputs(
+    lhs: jax.Array,
+    rhs: jax.Array,
+    rhs_scale: jax.Array | None,
+    rhs_bias: jax.Array | None,
+    group_sizes: jax.Array,
+    group_offset: jax.Array,
+    fuse_act: str | None = None,
+) -> Dimensions:
+  """Validates the inputs for the GMM kernel."""
+
+  size_m = lhs.shape[0]
+  size_group, size_k, size_n = rhs.shape
+  size_lhs_group = group_sizes.shape[0]
+
+  assert size_group <= size_lhs_group
+  assert lhs.shape == (size_m, size_k)
+  assert rhs.shape == (size_group, size_k, size_n)
+  if rhs_bias is not None:
+    assert rhs_bias.shape == (size_group, 1, size_n)
+  if rhs_scale is not None:
+    num_quant_blocks = rhs_scale.shape[1]
+    assert rhs_scale.shape == (size_group, num_quant_blocks, 1, size_n)
+    assert size_k % num_quant_blocks == 0
+
+  assert group_offset.shape == (1,)
+
+  size_lhs_sublane = pltpu.get_tpu_info().get_sublane_tiling(lhs.dtype)
+  size_lhs_sublane = min(size_lhs_sublane, size_m)
+  if fuse_act is not None:
+    num_lanes = pltpu.get_tpu_info().num_lanes
+    if size_n % (2 * num_lanes) != 0:
+      raise ValueError(
+          f"{size_n=} should be divisible by 2 * num_lanes when fuse_act is "
+          "enabled since we need to split n dimension for gate and up."
+      )
+
+  return Dimensions(
+      size_m=size_m,
+      size_k=size_k,
+      size_n=size_n,
+      size_group=size_group,
+      size_lhs_group=size_lhs_group,
+      size_lhs_sublane=size_lhs_sublane,
+  )
+
+
+def get_cost_estimate(cfgs: GmmConfigs):
+  """Returns the cost estimate for the GMM kernel."""
+
+  dims = cfgs.dims
+  lhs_dtype = cfgs.lhs_cfgs.quant_dtype or cfgs.lhs_cfgs.dtype
+  rhs_dtype = cfgs.rhs_cfgs.dtype
+
+  # We use bits for rhs since it could sub-byte dtype like int4.
+  rhs_bits = jax.dtypes.itemsize_bits(rhs_dtype)
+  fp32_bytes = jnp.dtype(jnp.float32).itemsize
+
+  # TODO: Add compute flops for quant, dequant, and bias.
+  flops = 2 * dims.size_m * dims.size_k * dims.size_n
+
+  lhs_bytes = dims.size_m * dims.size_k * lhs_dtype.itemsize
+
+  rhs_size = dims.size_group * dims.size_k * dims.size_n
+  rhs_bytes = rhs_size * rhs_bits // 8
+  if cfgs.rhs_cfgs.has_scale:
+    num_quant_blocks = pl.cdiv(dims.size_k, cfgs.rhs_cfgs.quant_block_size)
+    rhs_bytes += dims.size_group * num_quant_blocks * dims.size_n * fp32_bytes
+  if cfgs.rhs_cfgs.has_bias:
+    rhs_bytes += dims.size_group * dims.size_n * fp32_bytes
+
+  out_bytes = dims.size_m * cfgs.out_size_n * cfgs.out_dtype.itemsize
+
+  total_bytes = lhs_bytes + rhs_bytes + out_bytes
+
+  return pl.CostEstimate(
+      flops=flops,
+      bytes_accessed=total_bytes,
+      transcendentals=0,
+  )
+
+
+def get_scope_name(cfgs: GmmConfigs) -> str:
+  dims = cfgs.dims
+  tiles = cfgs.tiles
+  return (
+      f"gmm_v2-g_{dims.size_group}-m_{dims.size_m}-k_{dims.size_k}-act_{cfgs.fuse_act}"
+      f"-n_{dims.size_n}-tm_{tiles.tile_m}-tk_{tiles.tile_k}-tn_{tiles.tile_n}"
+  )
+
+
+def make_gmm_configs(
+    lhs: jax.Array,
+    rhs: jax.Array,
+    rhs_scale: jax.Array | None,
+    rhs_bias: jax.Array | None,
+    group_sizes: jax.Array,
+    group_offset: jax.Array,
+    *,
+    tile_info: TileSizes | TileFn,
+    vmem_limit_bytes: int | None,
+    out_dtype: jnp.dtype | None,
+    acc_dtype: jnp.dtype | None,
+    maybe_quantize_lhs: bool,
+    zero_initialize: bool,
+    fuse_act: str | None = None,
+):
+  """Fills the GMM config for the GMM kernel."""
+
+  dims = validate_inputs(lhs, rhs, rhs_scale, rhs_bias, group_sizes, group_offset, fuse_act)
+
+  if rhs_scale is not None:
+    has_scale = True
+    rhs_quant_dtype = rhs.dtype
+    num_blocks = rhs_scale.shape[1]
+    block_size = dims.size_k // num_blocks
+  else:
+    has_scale = False
+    rhs_quant_dtype = None
+    block_size = dims.size_k
+
+  rhs_cfgs = InputConfigs(
+      quant_dtype=rhs_quant_dtype,
+      quant_block_size=block_size,
+      dtype=rhs.dtype,
+      has_bias=rhs_bias is not None,
+      has_scale=has_scale,
+  )
+
+  lhs_q_dtype = None
+  if maybe_quantize_lhs and rhs_cfgs.should_dequantize_after_matmul:
+    # Choose lhs quantization dtype based on TPU hardware support.
+    is_rhs_float = jnp.issubdtype(rhs_quant_dtype, jnp.floating)
+    tpu_info = pltpu.get_tpu_info()
+    # Check if there is hardware compute support for rhs dtype group.
+    if tpu_info.fp8_ops_per_second > 0:
+      # Special handling for 4-bit integer rhs as it can be converted to fp8
+      # without a numeric issues. Note that this is not the case for 4-bit
+      # floating rhs as conversion to int8 will cause numeric issues.
+      is_rhs_4bits = jax.dtypes.itemsize_bits(rhs_quant_dtype) == 4
+      if is_rhs_float or is_rhs_4bits:
+        lhs_q_dtype = jnp.float8_e4m3fn.dtype
+    if tpu_info.int8_ops_per_second > 0:
+      if not is_rhs_float:
+        lhs_q_dtype = jnp.int8.dtype
+
+  lhs_cfgs = InputConfigs(
+      quant_dtype=lhs_q_dtype,
+      # Input quantization involves reading all elements in a block to compute
+      # scale value. Since this operation is very memory intensive, we use a
+      # block size that is small enough to minimize memory overhead but large
+      # enough to minimize compute overhead of quantization.
+      quant_block_size=512,
+      dtype=lhs.dtype,
+  )
+
+  if out_dtype is None:
+    out_dtype = lhs.dtype
+
+  if acc_dtype is None:
+    if lhs_cfgs.quant_dtype is None:
+      acc_dtype = jnp.float32.dtype
+    else:
+      # Input quantization requires elementwise ops which can put pressure on
+      # VPUs. Using faster bf16 hardware during accumulation can help offset the
+      # pressure.
+      acc_dtype = jnp.bfloat16.dtype
+
+  if isinstance(tile_info, TileSizes):
+    tiles = tile_info
+  else:
+    tiles = tile_info(dims, lhs_cfgs, rhs_cfgs, vmem_limit_bytes, fuse_act)
+
+  return GmmConfigs(
+      dims=dims,
+      tiles=tiles,
+      lhs_cfgs=lhs_cfgs,
+      rhs_cfgs=rhs_cfgs,
+      out_dtype=jnp.dtype(out_dtype),
+      acc_dtype=jnp.dtype(acc_dtype),
+      zero_init=zero_initialize,
+      fuse_act=fuse_act,
+  )
+
+
+def get_metadata(cfgs: GmmConfigs) -> dict[str, str | int | float]:
+  cfgs_dict = dataclasses.asdict(cfgs)
+  ret = {}
+  for path, val in jax.tree_util.tree_leaves_with_path(cfgs_dict):
+    key = jax.tree_util.keystr(path, simple=True, separator=".")
+    if not isinstance(val, str | int | float):
+      val = str(val)
+    ret[key] = val
+  return ret
+
+
+@jax.jit(
+    static_argnames=[
+        "tile_info",
+        "vmem_limit_bytes",
+        "precision",
+        "preferred_element_type",
+        "acc_dtype",
+        "maybe_quantize_lhs",
+        "zero_initialize",
+        "fuse_act",
+    ]
+)
+def gmm_v2(
+    lhs: jax.Array,  # [size_m, size_k]
+    rhs: jax.Array,  # [size_group, size_k, size_n]
+    group_sizes: jax.Array,  # int32[size_lhs_group]
+    rhs_scale: jax.Array | None = None,  # [size_group, num_blocks, 1, out_size]
+    rhs_bias: jax.Array | None = None,  # [size_group, 1, out_size]
+    group_offset: jax.Array | None = None,  # int32[1]
+    *,
+    tile_info: TileSizes | TileFn = calculate_tiling,
+    vmem_limit_bytes: int | None = None,
+    precision: jax.lax.Precision = jax.lax.Precision.DEFAULT,
+    preferred_element_type: jnp.dtype | None = None,
+    acc_dtype: jnp.dtype | None = None,
+    maybe_quantize_lhs: bool = True,
+    zero_initialize: bool = True,
+    fuse_act: str | None = None,
+) -> jax.Array:
+  """GMM kernel implemented with emit_pipeline.
+
+  Dynamically calculate offset lhs/out tiles to reduce redundant computations.
+  Additionally, it adjusts dma size based on number of valid rows and utilize
+  triple buffering on weights to better utilize memory.
+
+  Args:
+    lhs: lhs with shape [size_m, size_k].
+    rhs: rhs with shape [size_group, size_k, size_n].
+    group_sizes: The group sizes of lhs rows of shape [size_lhs_group,].
+    rhs_scale: The rhs scale of shape [size_group, num_blocks, 1, out_size].
+    rhs_bias: The rhs bias of shape [size_group, 1, out_size].
+    group_offset: Optional. The group offset of shape [1,].
+    tile_info: The tile sizes or tile function to use.
+    vmem_limit_bytes: Optional vmem limit in bytes.
+    precision: Unused. Exists for compatibility reasons.
+    preferred_element_type: Optional jnp.dtype for the output matrix.
+    acc_dtype: Optional jnp.dtype for the accumulator.
+    maybe_quantize_lhs: Quantize lhs if set to True and rhs is quantized.
+    zero_initialize: Whether to initialize unvisited output elements to zero.
+    fuse_act: Activation function to fuse with GMM, None if no fusion.
+
+  Returns:
+    Output of shape [size_m, size_n].
+  """
+
+  del precision
+
+  if group_offset is None:
+    group_offset = jnp.array([0], dtype=jnp.int32)
+  else:
+    if jnp.isscalar(group_offset):
+      group_offset = group_offset[None]
+
+  if vmem_limit_bytes is None:
+    vmem_limit_bytes = int(pltpu.get_tpu_info().vmem_capacity_bytes * 0.9)
+
+  cfgs = make_gmm_configs(
+      lhs,
+      rhs,
+      rhs_scale,
+      rhs_bias,
+      group_sizes,
+      group_offset,
+      tile_info=tile_info,
+      vmem_limit_bytes=vmem_limit_bytes,
+      out_dtype=preferred_element_type,
+      acc_dtype=acc_dtype,
+      maybe_quantize_lhs=maybe_quantize_lhs,
+      zero_initialize=zero_initialize,
+      fuse_act=fuse_act,
+  )
+  dims = cfgs.dims
+  tiles = cfgs.tiles
+
+  # Prepare block specs.
+  rhs_scale_spec = rhs_bias_spec = None
+  if rhs_scale is not None:
+    rhs_scale = rhs_scale.astype(jnp.float32)
+    rhs_scale_spec = pl.BlockSpec(memory_space=pltpu.HBM)
+  if rhs_bias is not None:
+    rhs_bias = rhs_bias.astype(jnp.float32)
+    rhs_bias_spec = pl.BlockSpec(memory_space=pltpu.HBM)
+
+  # Initialize scratch shapes.
+  max_num_gm = dims.size_group + pl.cdiv(dims.size_m, tiles.tile_m) - 1
+  acc_cols = 2 * tiles.tile_n if cfgs.fuse_act is not None else tiles.tile_n
+  scratch_shapes = [
+      # partial_out_ref
+      pltpu.VMEM((dims.size_lhs_sublane, tiles.tile_n), cfgs.out_dtype),
+      # acc_ref
+      pltpu.VMEM((tiles.tile_m, acc_cols), cfgs.acc_dtype),
+      # metadata_ref
+      MetadataRef(
+          gm_id_to_group_id=pltpu.SMEM((max_num_gm,), jnp.int32),
+          gm_id_to_m_offset=pltpu.SMEM((max_num_gm + 1,), jnp.int32),
+      ),
+  ]
+
+  num_lanes = pltpu.get_tpu_info().num_lanes
+  if cfgs.zero_init:
+    # TODO: Create better heuristics for determining this value.
+    target_zero_ref_bytes = 2 * 1024 * 1024
+
+    # Zero initialization is done by tiling size_m dim where each tile invokes
+    # zero initializing DMA for up-to tile_zero_m rows. This means larger
+    # tile_zero_m will result in fewer number of tiles and lead to smaller
+    # overhead. However, in order to invoke DMA call up-to tile_zero_m rows, we
+    # need to store equivalent sized memory in VMEM buffer for the duration of
+    # DMA. Storing [tile_zero_m, size_n] in buffer will trigger OOM if
+    # tile_zero_m is too large. Instead, if we set column size as num_lanes
+    # (which is smallest allowed column size for DMA) and reuse the buffer by
+    # size_n//num_lanes times in a single tile, we can significantly increase
+    # tile_zero_m without triggering OOM.
+    out_bytes = jnp.dtype(cfgs.out_dtype).itemsize
+    tile_zero_m = target_zero_ref_bytes // num_lanes // out_bytes
+    tile_zero_m = min(tile_zero_m, dims.size_m)
+
+    scratch_shapes += [
+        pltpu.VMEM((tile_zero_m, num_lanes), cfgs.out_dtype),
+        pltpu.SemaphoreType.DMA((1,)),
+    ]
+  else:
+    scratch_shapes += [None, None]
+
+  aligned_n = align_to(cfgs.out_size_n, num_lanes)
+  out_init = jax.ShapeDtypeStruct((dims.size_m, aligned_n), cfgs.out_dtype)
+  rhs_weights = WeightsRef(weight=rhs, scale=rhs_scale, bias=rhs_bias)
+
+  return pl.pallas_call(
+      functools.partial(kernel_main, cfgs=cfgs),
+      out_shape=out_init,
+      grid_spec=pltpu.PrefetchScalarGridSpec(
+          num_scalar_prefetch=2,
+          in_specs=[
+              pl.BlockSpec(memory_space=pltpu.HBM),
+              WeightsRef(
+                  weight=pl.BlockSpec(memory_space=pltpu.HBM),
+                  scale=rhs_scale_spec,
+                  bias=rhs_bias_spec,
+              ),
+          ],
+          out_specs=pl.BlockSpec(memory_space=pltpu.HBM),
+          scratch_shapes=scratch_shapes,
+      ),
+      compiler_params=pltpu.CompilerParams(
+          vmem_limit_bytes=vmem_limit_bytes,
+          disable_bounds_checks=True,
+      ),
+      name=get_scope_name(cfgs),
+      cost_estimate=get_cost_estimate(cfgs),
+      metadata=get_metadata(cfgs),
+  )(group_sizes, group_offset, lhs, rhs_weights)[:, : cfgs.out_size_n]
diff --git a/src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_tgmm_kernel.py b/src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_tgmm_kernel.py
new file mode 100644
index 0000000000..779358199a
--- /dev/null
+++ b/src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_tgmm_kernel.py
@@ -0,0 +1,758 @@
+# Copyright 2026 DeepMind Technologies Limited. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+# Forked from:
+# https://github.com/openxla/tokamax/blob/3f332fcf85dcb87aab661d00228ed71a09b5fd56/
+# tokamax/_src/ops/ragged_dot/pallas_mosaic_tpu_v2_tgmm_kernel.py
+"""TGMM kernel"""
+
+import dataclasses
+import functools
+from typing import Any, Callable, Tuple
+
+import jax
+from jax import lax
+from jax.experimental import pallas as pl
+from jax.experimental.pallas import tpu as pltpu
+import jax.numpy as jnp
+
+from maxtext.kernels.megablox import pallas_mosaic_tpu_v2_gmm_kernel as gmm_v2
+
+
+@jax.tree_util.register_dataclass
+@dataclasses.dataclass(frozen=True)
+class OperandRef:
+  """Bundles a kernel operand with its optional per-N scale.
+
+  Registered as a pytree so it can be passed as a single 'pl.pallas_call' /
+  'emit_pipeline' operand (and in_spec). When 'scale' is None it contributes no
+  pytree leaf, so the kernel signature stays fixed-arity regardless of whether a
+  scale is present; the kernel just reads 'rhs_ref.scale' (None or a ref).
+  """
+
+  value: Any
+  scale: Any | None = None
+
+
+TileTgmmFn = Callable[
+    [
+        gmm_v2.Dimensions,
+        gmm_v2.InputConfigs,
+        gmm_v2.InputConfigs,
+        int,
+        jnp.dtype,
+        jnp.dtype,
+        int,
+    ],
+    gmm_v2.TileSizes,
+]
+
+
+def get_scope_name(cfgs: gmm_v2.GmmConfigs) -> str:
+  dims = cfgs.dims
+  tiles = cfgs.tiles
+  return (
+      f"tgmm_v2-g_{dims.size_group}-m_{dims.size_m}-k_{dims.size_k}-act_{cfgs.fuse_act}"
+      f"-n_{dims.size_n}-tm_{tiles.tile_m}-tk_{tiles.tile_k}-tn_{tiles.tile_n}"
+  )
+
+
+def get_cost_estimate(cfgs: gmm_v2.GmmConfigs) -> pl.CostEstimate:
+  dims = cfgs.dims
+  flops = 2 * dims.size_m * dims.size_k * dims.size_n
+  lhs_bytes = dims.size_m * dims.size_k * cfgs.lhs_cfgs.dtype.itemsize
+  rhs_bytes = dims.size_m * dims.size_n * cfgs.rhs_cfgs.dtype.itemsize
+  out_bytes = dims.size_group * dims.size_k * dims.size_n * cfgs.out_dtype.itemsize
+  return pl.CostEstimate(
+      flops=flops,
+      bytes_accessed=lhs_bytes + rhs_bytes + out_bytes,
+      transcendentals=0,
+  )
+
+
+def calculate_tgmm_tiling(
+    dims: gmm_v2.Dimensions,
+    lhs_cfgs: gmm_v2.InputConfigs,
+    rhs_cfgs: gmm_v2.InputConfigs,
+    vmem_limit_bytes: int,
+    out_dtype: jnp.dtype,
+    acc_dtype: jnp.dtype,
+    target_zero_ref_bytes: int,
+) -> gmm_v2.TileSizes:
+  """Calculate optimal tile sizes for TGMM kernel."""
+  # In tgmm, we calculate lhs.T @ dout which doesn't require quantization.
+  # Since we use it in MOE, the m can be dynamic and small. So we don't
+  # want it to be too big. At the same time, because the mxu size is 256, the
+  # rhs is divided into 256x256 tiles. The lhs is divided to blocks of 256-wide
+  # (256 on the contracting dimension) rows. So any size less than 256 will
+  # have the same perf as using 256.
+  bf16_bf16_tile_m = 256
+  tile_m = min(bf16_bf16_tile_m, dims.size_m)
+  tile_m = max(tile_m, dims.size_lhs_sublane)
+
+  num_k_tiles = num_n_tiles = 1
+  num_lanes = pltpu.get_tpu_info().num_lanes
+  tile_n = gmm_v2.align_to(dims.size_n, num_lanes)
+  # To avoid stalling MXU, we add some buffer room where tile_n cannot go
+  # smaller than 2x of mxu_column_size.
+  tile_n_lower_bound = pltpu.get_tpu_info().mxu_column_size * 2
+  tile_n_lower_bound = min(tile_n_lower_bound, dims.size_n)
+  tile_k = gmm_v2.align_to(dims.size_k, num_lanes)
+
+  def within_vmem_limit(tile_m, tile_k, tile_n):
+    acc_bytes = jax.dtypes.itemsize_bits(acc_dtype) // 8
+    out_bytes = jax.dtypes.itemsize_bits(out_dtype) // 8
+    lhs_bytes = jax.dtypes.itemsize_bits(lhs_cfgs.dtype) // 8
+    rhs_bytes = jax.dtypes.itemsize_bits(rhs_cfgs.dtype) // 8
+    num_buffers = 2
+    # For lhs, we use (num_buffers+1). +1 is needed because we are doing
+    # lhs.T @ rhs, lhs cannot be fed directly into MXU and has to go through
+    # XLU's transpose. in order to reduce redundant XLU computation, instead
+    # of performing XLU's transpose every time lhs is pushed into XLU, it
+    # caches the transposed value into VMEM. this increases VMEM requirement.
+    budget = (
+        tile_k * tile_n * (acc_bytes + num_buffers * out_bytes)
+        + (num_buffers + 1) * (tile_m * tile_k * lhs_bytes)
+        + num_buffers * (tile_m * tile_n * rhs_bytes)
+        # Reserve VMEM for zero_ref. Use the upper bound target_zero_ref_bytes
+        # since the actual zero_ref size depends on out_dtype/size_k and is
+        # always <= this value.
+        + target_zero_ref_bytes
+    )
+    return budget <= vmem_limit_bytes
+
+  prev_tile_n = tile_n
+  while not within_vmem_limit(tile_m, tile_k, tile_n):
+    num_n_tiles += 1
+    # The reason why we do "tile_n * num_n_tiles must cover size_n." is
+    # tile_n must be a multiple of num_lanes and
+    # tile_n * num_n_tiles must cover size_n.
+    tile_n = gmm_v2.align_to(dims.size_n, num_n_tiles * num_lanes) // num_n_tiles
+    # If size_n is small and awkwardly sized (e.g., size_n=100, num_lanes=128),
+    # align_to(100, N*128) // N can get stuck at a constant value (128) as N
+    # grows. If that constant value is above the floor and budget still
+    # doesn't fit, the loop never terminates. That's why we need to check if
+    # "tile_n >= prev_tile_n".
+    if tile_n < tile_n_lower_bound or tile_n >= prev_tile_n:
+      break
+    prev_tile_n = tile_n
+
+  if tile_n >= tile_n_lower_bound and within_vmem_limit(tile_m, tile_k, tile_n):
+    return gmm_v2.TileSizes(tile_m=tile_m, tile_k=tile_k, tile_n=tile_n)
+
+  if tile_n < tile_n_lower_bound:
+    num_n_tiles -= 1
+    tile_n = gmm_v2.align_to(dims.size_n, num_n_tiles * num_lanes) // num_n_tiles
+
+  prev_tile_k = tile_k
+  while not within_vmem_limit(tile_m, tile_k, tile_n):
+    num_k_tiles += 1
+    tile_k = gmm_v2.align_to(dims.size_k, num_k_tiles * num_lanes) // num_k_tiles
+    if tile_k < num_lanes or tile_k >= prev_tile_k:
+      break
+    prev_tile_k = tile_k
+
+  if tile_k < num_lanes:
+    num_k_tiles -= 1
+    tile_k = gmm_v2.align_to(dims.size_k, num_k_tiles * num_lanes) // num_k_tiles
+
+  if not within_vmem_limit(tile_m, tile_k, tile_n):
+    raise ValueError(
+        f"Could not find valid tile sizes for tgmm. dims={dims},"
+        f" tiles=({tile_m},{tile_k},{tile_n}), vmem={vmem_limit_bytes}"
+    )
+  return gmm_v2.TileSizes(tile_m=tile_m, tile_k=tile_k, tile_n=tile_n)
+
+
+def make_tgmm_configs(
+    lhs: jax.Array,  # [m, k]
+    rhs: jax.Array,  # [m, n]
+    rhs_scale: jax.Array,  # [1, 1, n] (per-N scale)
+    group_sizes: jax.Array,
+    num_actual_groups: int,
+    *,
+    tile_info: gmm_v2.TileSizes | TileTgmmFn,
+    vmem_limit_bytes: int | None,
+    out_dtype: jnp.dtype,
+    acc_dtype: jnp.dtype | None,
+    target_zero_ref_bytes: int,
+):
+  """Fills the GMM config for the TGMM kernel."""
+  assert out_dtype, "out_dtype cannot be None"
+  assert lhs.shape[0] == rhs.shape[0], (
+      f"lhs and rhs m-dim mismatch: {lhs.shape[0]}!={rhs.shape[0]} {lhs.shape}" f" vs {rhs.shape}"
+  )
+  size_m, size_k = lhs.shape
+  _, size_n = rhs.shape
+  if rhs_scale is not None:
+    # rhs_scale.shape[0] is the number of quant blocks along the m (reduction)
+    # dimension. tgmm_v2 only implements per-N (per-output-channel) scaling,
+    # i.e. a single m-block: rhs_scale.shape == (1, 1, size_n). A leading dim
+    # > 1 means the caller wants sub-channel (per-m-block) quantization.
+    if rhs_scale.ndim == 3 and rhs_scale.shape[0] > 1:
+      raise NotImplementedError(
+          "tgmm_v2 only supports per-N rhs_scale with shape (1, 1, size_n);"
+          f" got {rhs_scale.shape}, which implies {rhs_scale.shape[0]}"
+          " sub-channel quant blocks along the m (reduction) dimension."
+          " Sub-channel quantization is not implemented."
+      )
+    assert rhs_scale.shape == (1, 1, size_n), (
+        f"expecting rhs_scale.shape to be (1, 1, size_n) but got" f" {rhs_scale.shape}"
+    )
+  # size_lhs_sublane is used in tgmm_inner_kernel to set the
+  # (m/size_lhs_sublane, size_lhs_sublane, ...) reshape tile used on the m-axis
+  # for both 'tiled_lhs_ref' and 'tiled_rhs_ref'.
+  size_lhs_sublane = pltpu.get_tpu_info().get_sublane_tiling(lhs.dtype)
+  size_lhs_sublane = min(size_lhs_sublane, size_m)
+  size_rhs_sublane = pltpu.get_tpu_info().get_sublane_tiling(rhs.dtype)
+  size_rhs_sublane = min(size_rhs_sublane, size_m)
+  assert size_lhs_sublane == size_rhs_sublane, (
+      f"size_lhs_sublane should be the same as size_rhs_sublane {lhs.dtype=}," f" {rhs.dtype=}"
+  )
+  dims = gmm_v2.Dimensions(
+      size_m=size_m,
+      size_k=size_k,
+      size_n=size_n,
+      size_group=num_actual_groups,  # weight.shape[0]
+      size_lhs_group=group_sizes.shape[0],
+      size_lhs_sublane=size_lhs_sublane,
+  )
+
+  rhs_quant_block_size_m = size_m
+  rhs_cfgs = gmm_v2.InputConfigs(
+      quant_dtype=None,
+      quant_block_size=rhs_quant_block_size_m,
+      dtype=rhs.dtype,
+      has_scale=(rhs_scale is not None),
+  )
+  lhs_cfgs = gmm_v2.InputConfigs(
+      quant_dtype=None,
+      quant_block_size=-1,
+      dtype=lhs.dtype,
+  )
+
+  fuse_act = None  # fuse_act has to be None in tgmm.
+  if acc_dtype is None:
+    acc_dtype = jnp.float32.dtype
+  if isinstance(tile_info, gmm_v2.TileSizes):
+    tiles = tile_info
+  else:
+    tiles = tile_info(
+        dims,
+        lhs_cfgs,
+        rhs_cfgs,
+        vmem_limit_bytes,
+        out_dtype,
+        acc_dtype,
+        target_zero_ref_bytes,
+    )
+
+  return gmm_v2.GmmConfigs(
+      dims=dims,
+      tiles=tiles,
+      lhs_cfgs=lhs_cfgs,
+      rhs_cfgs=rhs_cfgs,
+      out_dtype=jnp.dtype(out_dtype),
+      acc_dtype=jnp.dtype(acc_dtype),
+      # GMM's 'zero_init' zeros unvisited m-rows via DMA, which doesn't apply to
+      # tgmm's [num_groups, k, n] output. The actual zero-initialization for
+      # tgmm accumulation happens at the 'pallas_call' level.
+      zero_init=False,
+      fuse_act=fuse_act,
+  )
+
+
+def tgmm_inner_kernel(
+    tiled_lhs_ref: jax.Array,
+    # [tile_m // size_lhs_sublane, size_lhs_sublane, tile_k]
+    tiled_rhs_ref: OperandRef,
+    # .value: [tile_m // size_lhs_sublane, size_lhs_sublane, tile_n]
+    # .scale: [1, 1, tile_n] or None
+    tiled_out_ref: jax.Array,
+    acc_ref: jax.Array,
+    metadata_ref: gmm_v2.MetadataRef,
+    *,
+    cfgs: gmm_v2.GmmConfigs,
+):
+  """Inner kernel for TGMM computation.
+
+  This kernel performs the matrix multiplication for a single tile of the output
+  in the TGMM operation (lhs.T @ rhs). It handles masking for partial groups
+  and accumulation across different group-major tiles.
+
+  Args:
+    tiled_lhs_ref: Reference to the tiled LHS data.
+    tiled_rhs_ref: OperandRef bundling the tiled RHS data ('.value') and its
+      optional per-N scale ('.scale', None when there is no scale).
+    tiled_out_ref: Reference to the tiled output buffer [None, tile_k, tile_n].
+    acc_ref: Scratch memory for accumulation [tile_k, tile_n].
+    metadata_ref: Contains metadata like group offsets and group IDs.
+    cfgs: GmmConfigs object containing kernel configurations.
+  """
+  # NB: grid=(num_n, num_k, num_gm)
+  tiled_rhs_scale_ref = tiled_rhs_ref.scale
+
+  tiled_lhs_ref = tiled_lhs_ref.reshape(-1, tiled_lhs_ref.shape[-1])
+  tiled_rhs_ref = tiled_rhs_ref.value.reshape(-1, tiled_rhs_ref.value.shape[-1])
+  gm_id = pl.program_id(2)
+
+  def _matmul(is_new_group: bool, is_group_changing: bool):
+
+    # Mask out invalid rows in the LHS/RHS tiles.
+    # The DMA loads tiles aligned to sublane boundaries, but the actual group
+    # data may not start/end on those boundaries.
+    m_start = metadata_ref.gm_id_to_m_offset[gm_id]
+    m_end = metadata_ref.gm_id_to_m_offset[gm_id + 1]
+    m_offset = m_start - m_start % cfgs.dims.size_lhs_sublane
+    m_start_local = m_start - m_offset
+    m_end_local = m_end - m_offset
+    lhs_iota = lax.broadcasted_iota(jnp.int32, tiled_lhs_ref.shape, 0)
+    lhs_mask = jnp.logical_and(m_start_local <= lhs_iota, lhs_iota < m_end_local)
+    lhs_masked = jnp.where(lhs_mask, tiled_lhs_ref[...], 0)
+    # If there are no NaNs, masking both lhs and rhs shouldn't be necessary.
+    # But without masking both, we sometimes see the result contain NaNs so we
+    # decide to mask both to be safe.
+    rhs_iota = lax.broadcasted_iota(jnp.int32, tiled_rhs_ref.shape, 0)
+    rhs_mask = jnp.logical_and(m_start_local <= rhs_iota, rhs_iota < m_end_local)
+    rhs_masked = jnp.where(rhs_mask, tiled_rhs_ref[...], 0)
+
+    acc = jax.lax.dot_general(
+        lhs_masked,
+        rhs_masked,
+        (((0,), (0,)), ((), ())),
+        preferred_element_type=jnp.float32,
+    )
+
+    if not is_new_group:
+      acc += acc_ref[...]
+
+    if is_group_changing:
+      if cfgs.rhs_cfgs.has_scale:
+        scale_slice = tiled_rhs_scale_ref[0]
+        acc *= scale_slice
+      tiled_out_ref[...] = acc.astype(tiled_out_ref.dtype)
+    else:
+      acc_ref[...] = acc
+
+  @jax.named_scope("matmul_new_group_and_changing")
+  def matmul_new_group_and_changing():
+    _matmul(is_new_group=True, is_group_changing=True)
+
+  @jax.named_scope("matmul_new_group")
+  def matmul_new_group():
+    _matmul(is_new_group=True, is_group_changing=False)
+
+  @jax.named_scope("matmul")
+  def matmul():
+    _matmul(is_new_group=False, is_group_changing=False)
+
+  @jax.named_scope("matmul_group_changing")
+  def matmul_group_changing():
+    _matmul(is_new_group=False, is_group_changing=True)
+
+  prev_gm_id = jnp.where(gm_id > 0, gm_id - 1, 0)
+  is_first_gm = gm_id == 0
+  group_id_changed = metadata_ref.gm_id_to_group_id[gm_id] != metadata_ref.gm_id_to_group_id[prev_gm_id]
+  new_group = jnp.logical_or(is_first_gm, group_id_changed)
+
+  is_last_gm = gm_id == (pl.num_programs(2) - 1)
+  next_gm_id = jnp.where(is_last_gm, gm_id, gm_id + 1)
+  next_group_id = metadata_ref.gm_id_to_group_id[next_gm_id]
+  cur_group_id = metadata_ref.gm_id_to_group_id[gm_id]
+  group_is_changing = jnp.logical_or(is_last_gm, cur_group_id != next_group_id)
+
+  lax.cond(
+      new_group,
+      lambda: lax.cond(
+          group_is_changing,
+          # gm_id is the only one in its group =>
+          # group_size + local_offset ≤ tile_m.
+          matmul_new_group_and_changing,
+          # matmul_new_group: first gm_id of a multi-gm group =>
+          # group spans ≥ 2 gm_ids.
+          matmul_new_group,
+      ),
+      lambda: lax.cond(
+          group_is_changing,
+          # matmul_group_changing: last gm_id of a multi-gm group.
+          matmul_group_changing,
+          # matmul: middle gm_id => group spans ≥ 3 gm_ids =>
+          # group_size + local_offset > 2*tile_m.
+          matmul,
+      ),
+  )
+
+
+class TgmmIndexMaps:
+  """Index maps for TGMM kernel."""
+
+  def __init__(self, metadata_ref: gmm_v2.MetadataRef, cfgs: gmm_v2.GmmConfigs):
+    self.metadata_ref = metadata_ref
+    self.cfgs = cfgs
+
+  def lhs_index_map(self, n_id: jax.Array, k_id: jax.Array, gm_id: jax.Array):
+    m_start = self.metadata_ref.gm_id_to_m_offset[gm_id]
+    m_end = self.metadata_ref.gm_id_to_m_offset[gm_id + 1]
+
+    row_start = m_start // self.cfgs.dims.size_lhs_sublane
+    row_end = pl.cdiv(m_end, self.cfgs.dims.size_lhs_sublane)
+    row_size = row_end - row_start
+    return (pl.ds(row_start, row_size), 0, k_id)
+
+  def rhs_index_map(self, n_id: jax.Array, k_id: jax.Array, gm_id: jax.Array):
+    m_start = self.metadata_ref.gm_id_to_m_offset[gm_id]
+    m_end = self.metadata_ref.gm_id_to_m_offset[gm_id + 1]
+
+    row_start = m_start // self.cfgs.dims.size_lhs_sublane
+    row_end = pl.cdiv(m_end, self.cfgs.dims.size_lhs_sublane)
+    row_size = row_end - row_start
+    return (pl.ds(row_start, row_size), 0, n_id)
+
+  def rhs_scale_index_map(self, n_id: jax.Array, k_id: jax.Array, gm_id: jax.Array):
+    return (0, 0, n_id)
+
+  def out_index_map(self, n_id: jax.Array, k_id: jax.Array, gm_id: jax.Array):
+    group_id = self.metadata_ref.gm_id_to_group_id[gm_id]
+    return (group_id, k_id, n_id)
+
+
+def generate_tgmm_block_specs(
+    metadata_ref: gmm_v2.MetadataRef, cfgs: gmm_v2.GmmConfigs
+) -> Tuple[Tuple[pl.BlockSpec, OperandRef], pl.BlockSpec]:
+  """Generates block specs for the given lhs, rhs, and out refs."""
+  index_map = TgmmIndexMaps(metadata_ref, cfgs)
+  # NB: in tgmm, LHS is reshaped from (M, K) to (-1, size_lhs_sublane, K) so
+  # that DMA transfers are aligned to sublane boundaries. The first dimension
+  # after this reshape has size tile_m // size_lhs_sublane — i.e., the number of
+  # "sublane-rows" in a tile.
+  bounded_slice_gm = pl.BoundedSlice(cfgs.tiles.tile_m // cfgs.dims.size_lhs_sublane)
+  lhs_block_spec = pl.BlockSpec(
+      (bounded_slice_gm, cfgs.dims.size_lhs_sublane, cfgs.tiles.tile_k),
+      index_map.lhs_index_map,
+  )
+  rhs_block_spec = pl.BlockSpec(
+      (bounded_slice_gm, cfgs.dims.size_lhs_sublane, cfgs.tiles.tile_n),
+      index_map.rhs_index_map,
+  )
+  rhs_scale_block_spec = None
+  if cfgs.rhs_cfgs.has_scale:
+    rhs_scale_block_spec = pl.BlockSpec(
+        (1, 1, cfgs.tiles.tile_n),
+        index_map.rhs_scale_index_map,
+    )
+  rhs_spec = OperandRef(value=rhs_block_spec, scale=rhs_scale_block_spec)
+  out_block_spec = pl.BlockSpec(
+      (None, cfgs.tiles.tile_k, cfgs.tiles.tile_n),
+      index_map.out_index_map,
+  )
+
+  return (lhs_block_spec, rhs_spec), out_block_spec
+
+
+def zero_out_start(
+    lhs_group_sizes_ref,  # int32[size_lhs_group]
+    group_offset_ref,  # int32[1]
+    out_ref,  # [num_actual_groups, k, n]
+    zero_ref,  # [tile_zero_k, num_lanes]
+    semaphore_ref,  # [1]
+):
+  """If group_sizes[i]==0, kick off async DMAs to zero out drhs[i]."""
+  num_actual_groups, aligned_k, aligned_n = out_ref.shape
+  tile_zero_k = zero_ref.shape[0]
+  zero_ref = zero_ref.reshape(1, tile_zero_k, -1)
+  num_lanes = pltpu.get_tpu_info().num_lanes
+  assert aligned_n % num_lanes == 0
+
+  zero_ref[...] = jnp.zeros_like(zero_ref)
+
+  def fill_zero(local_group_id, should_copy):
+    should_copy_int = should_copy.astype(int)
+    for i in range(pl.cdiv(aligned_k, tile_zero_k)):
+      size_k_to_copy = min(tile_zero_k, aligned_k - i * tile_zero_k)
+      for j in range(aligned_n // num_lanes):
+        src = zero_ref.at[pl.ds(0, should_copy_int), pl.ds(0, size_k_to_copy)]
+        dst = out_ref.at[
+            pl.ds(local_group_id, should_copy_int),
+            pl.ds(i * tile_zero_k, size_k_to_copy),
+            pl.ds(j * num_lanes, num_lanes),
+        ]
+        pltpu.make_async_copy(
+            src_ref=src,
+            dst_ref=dst,
+            sem=semaphore_ref.at[0],
+        ).start(priority=1)
+    return 1
+
+  num_groups_to_zero = 0
+  group_offset = group_offset_ref[0]
+  for local_group_id in range(num_actual_groups):
+    global_group_id = local_group_id + group_offset
+    should_copy = lhs_group_sizes_ref[global_group_id] == 0
+    num_groups_to_zero += should_copy.astype(int)
+    fill_zero(local_group_id, should_copy)
+
+  return num_groups_to_zero
+
+
+def zero_out_end(
+    num_groups_to_zero,
+    out_ref,  # [num_actual_groups, k, n]
+    semaphore_ref,  # [1]
+):
+  """Drain the DMAs started by zero_out_start."""
+  dst = out_ref.at[pl.ds(0, num_groups_to_zero),]
+  src = dst
+  pltpu.make_async_copy(
+      src_ref=src,
+      dst_ref=dst,
+      sem=semaphore_ref.at[0],
+  ).wait()
+
+
+def tgmm_kernel_main(
+    lhs_group_sizes_ref,  # int32[size_lhs_group]
+    group_offset_ref,  # int32[1]
+    lhs_ref,  # [m, k]
+    rhs_ref,  # OperandRef: .value [m, n], .scale [1, 1, n] or None
+    out_ref,  # [num_actual_groups, k, n]
+    # Scratch memory
+    acc_ref: jax.Array,  # [tile_k, tile_n]
+    metadata_ref: gmm_v2.MetadataRef,
+    zero_ref: jax.Array,  # [tile_zero_k, num_lanes]
+    semaphore_ref: jax.Array,  # [1]
+    *,
+    cfgs,
+):
+  """Main kernel function for TGMM computation.
+
+  Args:
+    lhs_group_sizes_ref: Reference to the group sizes of lhs.
+    group_offset_ref: Reference to the group offset.
+    lhs_ref: Reference to the LHS array [m, k].
+    rhs_ref: OperandRef bundling the RHS array ('.value' [m, n]) and its
+      optional per-N scale ('.scale' [1, 1, n], None when there is no scale).
+    out_ref: Reference to the output array [num_groups, k, n].
+    acc_ref: Scratch memory reference for accumulation [tile_k, tile_n].
+    metadata_ref: Reference to the metadata structure.
+    zero_ref: Scratch buffer for zeroing empty groups' output.
+    semaphore_ref: DMA semaphore for the zeroing copies.
+    cfgs: GmmConfigs object containing kernel configurations.
+  """
+  num_groups_to_zero = zero_out_start(
+      lhs_group_sizes_ref,
+      group_offset_ref,
+      out_ref,
+      zero_ref,
+      semaphore_ref,
+  )
+
+  num_k = pl.cdiv(cfgs.dims.size_k, cfgs.tiles.tile_k)
+  num_n = pl.cdiv(cfgs.dims.size_n, cfgs.tiles.tile_n)
+  num_gm = gmm_v2.fill_metadata(
+      lhs_group_sizes_ref,
+      group_offset_ref,
+      metadata_ref,
+      cfgs=cfgs,
+  )
+
+  in_specs, out_specs = generate_tgmm_block_specs(metadata_ref, cfgs)
+  pipeline_fn = pltpu.emit_pipeline(
+      functools.partial(tgmm_inner_kernel, cfgs=cfgs),
+      grid=(num_n, num_k, num_gm),
+      in_specs=in_specs,
+      out_specs=out_specs,
+  )
+  lhs_in = lhs_ref.reshape(-1, cfgs.dims.size_lhs_sublane, lhs_ref.shape[-1])
+  rhs_value = rhs_ref.value
+  rhs_in = rhs_value.reshape(-1, cfgs.dims.size_lhs_sublane, rhs_value.shape[-1])
+  rhs_operand = OperandRef(value=rhs_in, scale=rhs_ref.scale)
+  scratches = [acc_ref, metadata_ref]
+
+  pipeline_fn(lhs_in, rhs_operand, out_ref, scratches=scratches)
+  zero_out_end(
+      num_groups_to_zero,
+      out_ref,
+      semaphore_ref,
+  )
+
+
+def validate_tgmm_inputs(
+    group_sizes: jax.Array,
+    num_actual_groups: int,
+    group_offset: jax.Array | None = None,
+) -> None:
+  """Validates inputs to 'tgmm_v2'.
+
+  Call this eagerly before invoking the kernel. It is not jit-safe because it
+  concretizes 'group_offset'.
+
+  Args:
+    group_sizes: The sizes of each group.
+    num_actual_groups: The number of actual groups.
+    group_offset: An optional offset for the group indices.
+  """
+  if group_offset is None:
+    group_offset = jnp.array([0], dtype=jnp.int32)
+  elif jnp.isscalar(group_offset):
+    assert group_offset.size == 1
+    if jnp.isscalar(group_offset):
+      group_offset = group_offset[None]
+  if group_sizes.size < group_offset[0] + num_actual_groups:
+    raise ValueError(
+        f"group_sizes.size ({group_sizes.size}) must be >= group_offset"
+        f" ({group_offset[0]}) + num_actual_groups ({num_actual_groups})"
+    )
+
+
+@jax.jit(
+    static_argnames=[
+        "num_actual_groups",
+        "tile_info",
+        "vmem_limit_bytes",
+        "precision",
+        "preferred_element_type",
+        "acc_dtype",
+    ],
+)
+def tgmm_v2(
+    lhs: jax.Array,  # [size_m, size_k]
+    rhs: jax.Array,  # [size_m, size_n]
+    group_sizes: jax.Array,
+    num_actual_groups: int,
+    rhs_scale: jax.Array | None = None,  # [1, 1, size_n] (per-N scale)
+    group_offset: jax.Array | None = None,
+    *,
+    tile_info: gmm_v2.TileSizes | TileTgmmFn = calculate_tgmm_tiling,
+    vmem_limit_bytes: int | None = None,
+    precision: jax.lax.Precision = jax.lax.Precision.DEFAULT,
+    preferred_element_type: jnp.dtype | None = None,
+    acc_dtype: jnp.dtype | None = None,
+):
+  """Computes a transposed grouped matrix multiplication.
+
+  This kernel computes
+  grad_rhs=lhs[sizes[i-1]:sizes[i], :].T @ rhs[sizes[i-1]:sizes[i], :], aka
+  grad_rhs = lhs.T @ grad.
+
+  Args:
+    lhs: The left-hand side array with shape [size_m, size_k].
+    rhs: The right-hand side array with shape [size_m, size_n].
+    group_sizes: The group sizes of lhs with shape [size_lhs_group].
+    num_actual_groups: The actual number of groups: weight.shape[0].
+    rhs_scale: The per-N scale of the rhs.
+    group_offset: An optional offset for the group indices.
+    tile_info: Specifies the tiling strategy. Can be a `TileSizes` object or a
+      function to calculate it.
+    vmem_limit_bytes: The VMEM limit in bytes for the kernel.
+    precision: Unused. Exists for compatibility reasons.
+    preferred_element_type: Optional jnp.dtype for the output matrix.
+    acc_dtype: Optional jnp.dtype for the accumulator.
+
+  Returns:
+    The result of the transposed grouped matrix multiplication, with shape
+    [num_actual_groups, size_k, size_n].
+  """
+  del precision
+  if group_offset is None:
+    group_offset = jnp.array([0], dtype=jnp.int32)
+  else:
+    if jnp.isscalar(group_offset):
+      group_offset = group_offset[None]
+  if vmem_limit_bytes is None:
+    vmem_limit_bytes = int(pltpu.get_tpu_info().vmem_capacity_bytes * 0.9)
+
+  # Target VMEM size for the zero-init scratch buffer (zero_ref). The actual
+  # allocation may be smaller after capping by size_k and rounding down to a
+  # sublane multiple, so this also serves as an upper bound used by
+  # calculate_tgmm_tiling when reserving VMEM for zero_ref.
+  target_zero_ref_bytes = 2 * 1024 * 1024
+
+  cfgs = make_tgmm_configs(
+      lhs,
+      rhs,
+      rhs_scale,
+      group_sizes,
+      num_actual_groups,
+      tile_info=tile_info,
+      vmem_limit_bytes=vmem_limit_bytes,
+      out_dtype=preferred_element_type,
+      acc_dtype=acc_dtype,
+      target_zero_ref_bytes=target_zero_ref_bytes,
+  )
+  dims = cfgs.dims
+  tiles = cfgs.tiles
+
+  num_lanes = pltpu.get_tpu_info().num_lanes
+  aligned_n = gmm_v2.align_to(dims.size_n, num_lanes)
+  # Pad K up to a tile_k multiple so (a) every k-tile written by the matmul
+  # stays in-bounds, and (b) the zero-init path can slice in sublane-aligned
+  # chunks. tile_k is num_lanes-aligned, which is also sublane-tile-aligned.
+  aligned_k = gmm_v2.align_to(dims.size_k, tiles.tile_k)
+  out_init = jax.ShapeDtypeStruct((num_actual_groups, aligned_k, aligned_n), cfgs.out_dtype)
+  max_num_gm = dims.size_group + pl.cdiv(dims.size_m, tiles.tile_m) - 1
+  scratch_shapes = [
+      # acc_ref
+      pltpu.VMEM((tiles.tile_k, tiles.tile_n), cfgs.acc_dtype),
+      # metadata_ref
+      gmm_v2.MetadataRef(
+          gm_id_to_group_id=pltpu.SMEM((max_num_gm,), jnp.int32),
+          gm_id_to_m_offset=pltpu.SMEM((max_num_gm + 1,), jnp.int32),
+      ),
+  ]
+
+  # Prepare zero initializing the drhs[i, :, :] where the group_size[i] is 0.
+  out_bytes = jnp.dtype(cfgs.out_dtype).itemsize
+  tile_zero_k = target_zero_ref_bytes // num_lanes // out_bytes
+  tile_zero_k = min(tile_zero_k, dims.size_k)
+  size_out_sublane = pltpu.get_tpu_info().get_sublane_tiling(cfgs.out_dtype)
+  tile_zero_k = (tile_zero_k // size_out_sublane) * size_out_sublane
+  assert tile_zero_k > 0
+  scratch_shapes += [
+      pltpu.VMEM((tile_zero_k, num_lanes), cfgs.out_dtype),
+      pltpu.SemaphoreType.DMA((1,)),
+  ]
+
+  if rhs_scale is not None:
+    rhs_scale = rhs_scale.astype(jnp.float32)
+    pad_n = aligned_n - dims.size_n
+    if pad_n > 0:
+      rhs_scale = jnp.pad(rhs_scale, ((0, 0), (0, 0), (0, pad_n)))
+  rhs = OperandRef(value=rhs, scale=rhs_scale)
+  hbm_spec = pl.BlockSpec(memory_space=pltpu.HBM)
+  in_specs = [
+      hbm_spec,  # lhs
+      # the tree.map build a
+      # OperandRef(value=hbm_spec, scale=None if scale is None else hbm_spec.
+      jax.tree.map(lambda _: hbm_spec, rhs),  # rhs
+  ]
+
+  return pl.pallas_call(
+      functools.partial(tgmm_kernel_main, cfgs=cfgs),
+      out_shape=out_init,
+      grid_spec=pltpu.PrefetchScalarGridSpec(
+          num_scalar_prefetch=2,
+          in_specs=in_specs,
+          out_specs=pl.BlockSpec(memory_space=pltpu.HBM),
+          scratch_shapes=scratch_shapes,
+      ),
+      compiler_params=pltpu.CompilerParams(
+          vmem_limit_bytes=vmem_limit_bytes,
+          disable_bounds_checks=True,
+      ),
+      name=get_scope_name(cfgs),
+      cost_estimate=get_cost_estimate(cfgs),
+      # the metadata here is for profiling, debugging, and cost modeling.
+      # It does not affect the kernel's computation.
+      metadata=gmm_v2.get_metadata(cfgs),
+  )(group_sizes, group_offset, lhs, rhs)[:, : dims.size_k, : dims.size_n]
diff --git a/src/maxtext/layers/moe.py b/src/maxtext/layers/moe.py
index 020956098c..7e69d0b938 100644
--- a/src/maxtext/layers/moe.py
+++ b/src/maxtext/layers/moe.py
@@ -1297,7 +1297,9 @@ def extract_vma(tensor):
       # We support three implementations for gmm - tokamax, older forked kernel, or jax.lax.ragged_dot
       # For quantized tokamax we call a forked version that supports our quantization recipes.
       if self.config.use_tokamax_gmm:
-        if self.config.quantization:  # tokamax (quantized)
+        # tokamax gmm v1 (quantized) or tokamax gmm v2 (unquantized)
+        # tokamax gmm v2 (quantized) not supported yet
+        if self.config.quantization or self.config.use_gmm_v2:
           output = mblx.gmm(
               lhs=inputs,
               rhs=kernel,
@@ -1312,6 +1314,7 @@ def extract_vma(tensor):
               weight_gather_axes=weight_gather_axes,
               lhs_vma_axes=lhs_vma_axes,
               rhs_vma_axes=rhs_vma_axes,
+              use_gmm_v2=self.config.use_gmm_v2,
           )
         else:  # tokamax (unquantized)
           output = tokamax.ragged_dot(
diff --git a/tests/integration/tokamax_test.py b/tests/integration/tokamax_test.py
index 4a270b8164..d2a4d291ad 100644
--- a/tests/integration/tokamax_test.py
+++ b/tests/integration/tokamax_test.py
@@ -34,14 +34,21 @@ class Train(parameterized.TestCase):
       {
           "testcase_name": "gmm bf16",
           "quantization": "",
+          "use_gmm_v2": False,
       },
       {
           "testcase_name": "gmm fp8",
           "quantization": "fp8_full",
+          "use_gmm_v2": False,
+      },
+      {
+          "testcase_name": "gmm v2 bf16",
+          "quantization": "",
+          "use_gmm_v2": True,
       },
   )
   @pytest.mark.tpu_only
-  def test_different_configs(self, quantization: str):
+  def test_different_configs(self, quantization: str, use_gmm_v2: bool):
     """Smoke train with small config."""
     test_tmpdir = os.environ.get("TEST_TMPDIR", gettempdir())
     outputs_dir = os.environ.get("TEST_UNDECLARED_OUTPUTS_DIR", test_tmpdir)
@@ -66,13 +73,33 @@ def test_different_configs(self, quantization: str):
         "sparse_matmul=True",
         "megablox=False",
         "use_tokamax_gmm=True",
+        f"use_gmm_v2={use_gmm_v2}",
+        # tile sizes
+        "wi_tile_fwd_batch_seq=128",
+        "wi_tile_fwd_embed_dim=128",
+        "wi_tile_fwd_mlp_dim=128",
+        "wi_tile_dlhs_batch_seq=128",
+        "wi_tile_dlhs_embed_dim=128",
+        "wi_tile_dlhs_mlp_dim=128",
+        "wi_tile_drhs_batch_seq=128",
+        "wi_tile_drhs_embed_dim=128",
+        "wi_tile_drhs_mlp_dim=128",
+        "wo_tile_fwd_batch_seq=128",
+        "wo_tile_fwd_embed_dim=128",
+        "wo_tile_fwd_mlp_dim=128",
+        "wo_tile_dlhs_batch_seq=128",
+        "wo_tile_dlhs_embed_dim=128",
+        "wo_tile_dlhs_mlp_dim=128",
+        "wo_tile_drhs_batch_seq=128",
+        "wo_tile_drhs_embed_dim=128",
+        "wo_tile_drhs_mlp_dim=128",
         # tokamax splash
         "max_target_length=128",
         "attention=flash",
         "use_tokamax_splash=True",
         # quantization
         f"quantization={quantization}",
-        "use_qwix_quantization=True",
+        f"use_qwix_quantization={quantization == 'fp8_full'}",
         "weight_quantization_calibration_method=fixed,-224,224",
         "act_quantization_calibration_method=fixed,-224,224",
         # train
diff --git a/tests/unit/configs_value_test.py b/tests/unit/configs_value_test.py
index 103329ee49..d28001f7c8 100644
--- a/tests/unit/configs_value_test.py
+++ b/tests/unit/configs_value_test.py
@@ -138,6 +138,40 @@ def test_initialize_pydantic_bad_keys(self):
           ]
       )
 
+  def test_gmm_v2_quantization_disallowed(self):
+    """Tests that use_gmm_v2=True with quantization enabled is disallowed."""
+    argv = [
+        "",
+        _BASE_CONFIG_PATH,
+        "run_name=test",
+        "use_gmm_v2=true",
+        "quantization=fp8_full",
+    ]
+    with self.assertRaises(pydantic.ValidationError):
+      pyconfig.initialize(argv)
+
+    argv_qwix = [
+        "",
+        _BASE_CONFIG_PATH,
+        "run_name=test",
+        "use_gmm_v2=true",
+        "use_qwix_quantization=true",
+    ]
+    with self.assertRaises(pydantic.ValidationError):
+      pyconfig.initialize(argv_qwix)
+
+  def test_gmm_v2_requires_tokamax_gmm(self):
+    """Tests that use_gmm_v2=True requires use_tokamax_gmm=True."""
+    argv = [
+        "",
+        _BASE_CONFIG_PATH,
+        "run_name=test",
+        "use_gmm_v2=true",
+        "use_tokamax_gmm=false",
+    ]
+    with self.assertRaises(pydantic.ValidationError):
+      pyconfig.initialize(argv)
+
 
 if __name__ == "__main__":
   unittest.main()