Repository: jax-ml/jax
Pull request #29613 (2025-06-20 14:56:57)

[Mosaic GPU] Add support for 2CTA MMA with M=128 (64 per block)
[Mosaic GPU] Add support for 2CTA MMA with M=128 (64 per block)

This requires us to add yet another layout, due to how non-uniform the TensorCore
requirements are...



Patch:
diff --git a/jax/_src/pallas/mosaic_gpu/lowering.py b/jax/_src/pallas/mosaic_gpu/lowering.py
index b1c46f0a0..5e672f477 100644
--- a/jax/_src/pallas/mosaic_gpu/lowering.py
+++ b/jax/_src/pallas/mosaic_gpu/lowering.py
@@ -313,7 +313,9 @@ def _run_scoped_resource_estimator(
         packing = 4 // aval.dtype.itemsize
       else:
         packing = 1
-      layout = tcgen05._infer_tmem_layout(aval.shape, packing=packing)
+      layout = tcgen05._infer_tmem_layout(
+          aval.shape, aval.collective, packing=packing
+      )
       cols_used = layout.cols_in_shape(
           aval.shape, mgpu_utils.dtype_to_ir_type(aval.dtype)
       )
@@ -453,7 +455,9 @@ class ModuleContext:
     else:
       packing = 1
     if layout is None:
-      layout = tcgen05._infer_tmem_layout(struct.shape, packing=packing)
+      layout = tcgen05._infer_tmem_layout(
+          struct.shape, collective, packing=packing
+      )
     mlir_dtype = mgpu_utils.dtype_to_ir_type(struct.dtype)
     unpadded_cols_used = layout.cols_in_shape(struct.shape, mlir_dtype)
     cols_used = tcgen05._alloc_ncols(unpadded_cols_used, exact_cols)
diff --git a/jax/experimental/mosaic/gpu/core.py b/jax/experimental/mosaic/gpu/core.py
index 94ce12af2..dfcd07245 100644
--- a/jax/experimental/mosaic/gpu/core.py
+++ b/jax/experimental/mosaic/gpu/core.py
@@ -418,7 +418,7 @@ def _construct_smem_reftree(
         )
         if layout is None:
           layout = tcgen05._infer_tmem_layout(
-              shape, 1 if packing is None else packing
+              shape, collective, 1 if packing is None else packing
           )
         num_cols = layout.cols_in_shape(shape, utils.dtype_to_ir_type(dtype))
         tmem_allocs.append(_TMEMAlloc(addr_ref, num_cols, collective))
diff --git a/jax/experimental/mosaic/gpu/fragmented_array.py b/jax/experimental/mosaic/gpu/fragmented_array.py
index 753008daf..f16c61170 100644
--- a/jax/experimental/mosaic/gpu/fragmented_array.py
+++ b/jax/experimental/mosaic/gpu/fragmented_array.py
@@ -2483,7 +2483,7 @@ class FragmentedArray:
     if any(
         len(dim_shape) != 1 for dim_shape in tiled_nested_shape[-layout.tiled_tiling_rank :]
     ):
-      raise NotImplementedError("Memory and register tiling incompatible")
+      raise NotImplementedError("Memory and register tiling too complicated")
     tiled_shape = list(itertools.chain.from_iterable(tiled_nested_shape))
     elem_tiled_strides = list(itertools.chain.from_iterable(tiled_nested_strides))
     lane_shape = [
diff --git a/jax/experimental/mosaic/gpu/tcgen05.py b/jax/experimental/mosaic/gpu/tcgen05.py
index c1ccdc7e7..d7f05812a 100644
--- a/jax/experimental/mosaic/gpu/tcgen05.py
+++ b/jax/experimental/mosaic/gpu/tcgen05.py
@@ -142,7 +142,12 @@ def mma(
     element_type2 = a.dtype
     if m != 128:
       raise ValueError(f"Only M=128 is supported for MMA with A in TMEM, but got M={m}")
-    if a.layout != (expected_layout := _infer_tmem_layout(a.shape, packing=2)):
+    # Watch out: this layout must be consistent with D's layout (up to packing).
+    expected_packing = 32 // utils.bitwidth(element_type)
+    expected_layout = _infer_tmem_layout(
+        a.shape, collective, packing=expected_packing
+    )
+    if a.layout != expected_layout:
       raise ValueError(
           f"A layout mismatch: expected {expected_layout}, got {a.layout}"
       )
@@ -171,18 +176,21 @@ def mma(
       )
     n_lane_groups = 1
   elif m == 64:
+    # Watch out: this layout must be consistent with A's layout (up to packing).
     # 2CTA M=128 instruction uses a different TMEM layout than 1CTA M=64.
-    if collective:
-      raise NotImplementedError("M=64 with collective MMA is not supported yet")
-    if d.layout != (expected_d_layout := tmem_half_lane_layout(n * num_cta, packing=1)):
+    expected_d_layout = _infer_tmem_layout(d.shape, collective, packing=1)
+    if d.layout != expected_d_layout:
       raise ValueError(
           f"Accumulator layout mismatch: expected {expected_d_layout}, got {d.layout}"
       )
-    n_lane_groups = 2
-    # We can't split N into groups if we would partition it below the tile size.
-    # TODO: We only need to check this if N is the minormost dim in B.
-    if 8 * swizzle // utils.bitwidth(element_type) > n // n_lane_groups:
-      raise ValueError("Swizzle is too big for MMA with M=64. Try lowering it.")
+    if collective:
+      n_lane_groups = 1
+    else:
+      n_lane_groups = 2
+      # We can't split N into groups if we would partition it below the tile size.
+      # TODO: We only need to check this if N is the minormost dim in B.
+      if 8 * swizzle // utils.bitwidth(element_type) > n // n_lane_groups:
+        raise ValueError("Swizzle is too big for MMA with M=64. Try lowering it.")
   else:
     raise ValueError(f"Only M=128 and M=64 are supported for MMA, but got M={m}")
   f32 = ir.F32Type.get()
@@ -557,7 +565,7 @@ class TMEMLayout(fa.TiledLayout):
     return math.prod(shape) // TMEM_ROWS // self.vector_length
 
 
-def _infer_tmem_layout(shape: tuple[int, int], packing: int = 1) -> TMEMLayout:
+def _infer_tmem_layout(shape: tuple[int, int], collective: bool, packing: int) -> TMEMLayout:
   if len(shape) != 2:
     raise ValueError(f"TMEM can only represent 2D shapes, got {shape}")
   if packing > 8 or packing.bit_count() != 1:
@@ -570,12 +578,16 @@ def _infer_tmem_layout(shape: tuple[int, int], packing: int = 1) -> TMEMLayout:
   if shape[0] == TMEM_ROWS:
     return tmem_default_layout(packing)
   elif shape[0] == TMEM_ROWS // 2:
-    return tmem_half_lane_layout(shape[1], packing)
+    if collective:
+      return tmem_m64_collective_layout(shape[1], packing)
+    else:
+      return tmem_half_lane_layout(shape[1], packing)
   else:
     raise ValueError(f"Unsupported shape: {shape}")
 
 
 def tmem_default_layout(packing: int = 1):
+  """A TMEM layout used for 1CTA MMA with M=128 and 2CTA MMA with M=256."""
   if packing > 8 or packing.bit_count() != 1:
     raise ValueError(f"Packing must be <= 8 and a power of 2, got: {packing}")
   return TMEMLayout(
@@ -587,6 +599,7 @@ def tmem_default_layout(packing: int = 1):
 
 
 def tmem_half_lane_layout(columns, packing: int = 1):
+  """A TMEM layout used for 1CTA MMA with M=64."""
   if packing > 8 or packing.bit_count() != 1:
     raise ValueError(f"Packing must be <= 8 and a power of 2, got: {packing}")
   if columns % 16:
@@ -603,6 +616,38 @@ def tmem_half_lane_layout(columns, packing: int = 1):
   )
 
 
+def tmem_m64_collective_layout(columns, packing: int = 1):
+  """A TMEM layout used for 2CTA MMA with M=128."""
+  if packing > 8 or packing.bit_count() != 1:
+    raise ValueError(f"Packing must be <= 8 and a power of 2, got: {packing}")
+  if columns % 16:
+    raise ValueError(f"Columns must be a multiple of 16, got: {columns}")
+  return TMEMLayout(
+      fa.Tiling((
+          (TMEM_ROWS // 2, columns),
+          (fa.WARP_SIZE, columns // 2),
+          (packing,),
+      )),
+      warp_dims=(-4, -5,),
+      lane_dims=(-3,),
+      vector_dim=-1,
+  )
+
+
+def fa_m64_collective_layout(columns):
+  """The register layout for transfers to/from tmem_m64_collective_layout."""
+  if columns % 8:
+    raise ValueError(f"Columns must be a multiple of 8, got: {columns}")
+  return fa.TiledLayout(
+      fa.Tiling((
+          (TMEM_ROWS // 2, columns), (fa.WARP_SIZE, columns // 2), (8, 8), (2,)
+      )),
+      warp_dims=(-6, -7),
+      lane_dims=(-3, -2),
+      vector_dim=-1,
+  )
+
+
 @dataclasses.dataclass(frozen=True)
 class TMEMRef:
   address: ir.Value
@@ -639,7 +684,7 @@ class TMEMRef:
         raise ValueError(
             "collective argument must be provided when TMEM layout is inferred"
         )
-      layout = _infer_tmem_layout(shape, collective)
+      layout = _infer_tmem_layout(shape, collective, packing=1)
     else:
       layout.check_type(shape, dtype)
     # TODO: Do we have to do this??
@@ -701,6 +746,12 @@ class TMEMRef:
     elif self.layout == tmem_half_lane_layout(self.shape[1], packing=2):
       packing = 2
       default_fa_layout = fa.WGMMA_LAYOUT
+    elif self.layout == tmem_m64_collective_layout(self.shape[1], packing=1):
+      packing = 1
+      default_fa_layout = fa_m64_collective_layout(self.shape[1])
+    elif self.layout == tmem_m64_collective_layout(self.shape[1], packing=2):
+      packing = 2
+      default_fa_layout = fa_m64_collective_layout(self.shape[1])
     else:
       raise ValueError(f"TMEM layout {self.layout} is not supported")
     if layout is None:
@@ -724,6 +775,12 @@ class TMEMRef:
       assert raw_registers.shape[0] == 4
       registers = np.concatenate([raw_registers[:2], raw_registers[2:]], axis=1)
       registers = registers.T.reshape(regs_shape)
+    elif layout == fa_m64_collective_layout(self.shape[1]) and self.layout == tmem_m64_collective_layout(self.shape[1], packing=packing):
+      regs_shape = layout.registers_shape(self.shape)
+      # We take half the columns, because they are split over halves of TMEM.
+      registers = _load_32xcols(
+          self.address, self.shape[1] // 2, self.dtype, packing
+      ).reshape(regs_shape)
     else:
       raise ValueError(
           f"Loads from TMEM layout {self.layout} to register layout"
@@ -763,6 +820,8 @@ class TMEMRef:
       registers = value.registers.T.reshape(2, -1)
       registers = np.concatenate(np.split(registers, 2, axis=1), axis=0)
       _store_32xcols(self.address, registers, packing)
+    elif value.layout == fa_m64_collective_layout(self.shape[1]) and self.layout == tmem_m64_collective_layout(self.shape[1], packing=packing):
+      _store_32xcols(self.address, value.registers.reshape(4, -1), packing)
     else:
       raise ValueError(
           f"Storing from register layout {value.layout} to TMEM layout"


Test patch:
diff --git a/tests/mosaic/gpu_test.py b/tests/mosaic/gpu_test.py
index d1ea797de..5c4cc6171 100644
--- a/tests/mosaic/gpu_test.py
+++ b/tests/mosaic/gpu_test.py
@@ -1120,17 +1120,21 @@ class TCGen05Test(TestCase):
   @parameterized.product(
       jax_dtype_packing=[(jnp.float32, 1), (jnp.float16, 1), (jnp.float16, 2)],
       reg_tmem_layout_m=[
-          (tcgen05.LAYOUT, lambda _, p: tcgen05.tmem_default_layout(p), 128),
-          (fa.WGMMA_LAYOUT, tcgen05.tmem_half_lane_layout, 64),
+          (lambda _: tcgen05.LAYOUT, lambda _, p: tcgen05.tmem_default_layout(p), 128),
+          (lambda _: fa.WGMMA_LAYOUT, tcgen05.tmem_half_lane_layout, 64),
+          (tcgen05.fa_m64_collective_layout, tcgen05.tmem_m64_collective_layout, 64),
       ],
   )
   def test_load_store_tmem_swizzle(self, jax_dtype_packing, reg_tmem_layout_m):
     jax_dtype, packing = jax_dtype_packing
-    reg_layout, tmem_layout_f, m = reg_tmem_layout_m
+    reg_layout_f, tmem_layout_f, m = reg_tmem_layout_m
     swizzle = 128
     in_mlir_dtype = utils.dtype_to_ir_type(jax_dtype)
     swizzle_elems = swizzle // bytewidth(in_mlir_dtype)
     tiling = (8, swizzle_elems)
+    # SMEM store planning for FA is too complicated for too long N...
+    n = 128 if jax_dtype != jnp.float32 else 64
+    reg_layout = reg_layout_f(n)
 
     def kernel(ctx, input, output, scratch):
       smem, barrier, tmem = scratch
@@ -1151,11 +1155,11 @@ class TCGen05Test(TestCase):
       )
       ctx.await_async_copy(0)
 
-    x = self.prng.uniform(-1, 1, (m, 128)).astype(jax_dtype)
+    x = self.prng.uniform(-1, 1, (m, n)).astype(jax_dtype)
     scratch_shape = [
         jax.ShapeDtypeStruct(tile_shape(x.shape, tiling), jax_dtype),
         mgpu.TMABarrier(),
-        mgpu.TMEM(x.shape, jax_dtype, layout=tmem_layout_f(128, packing)),
+        mgpu.TMEM(x.shape, jax_dtype, layout=tmem_layout_f(n, packing)),
     ]
     y = mgpu.as_gpu_kernel(
         kernel, (1, 1, 1), (128, 1, 1), x, x, scratch_shape
@@ -1469,22 +1473,21 @@ class TCGen05Test(TestCase):
       rhs_transpose=(False, True),
       in_jax_dtype=(jnp.float16,),
       out_jax_dtype=(jnp.float32,),
-      m=(256,),  # TODO(apaszke): 64, 192, 256
+      m=(128, 256),  # TODO(apaszke): 192, 256
       n=(128, 256),  # TODO(apaszke): 192, other non-power-of-2
-      k_steps=(1, 2),
       swizzle=(32, 64, 128,),
   )
   def test_mma_collective(
       self,
       m,
       n,
-      k_steps,
       swizzle,
       lhs_transpose,
       rhs_transpose,
       in_jax_dtype,
       out_jax_dtype,
   ):
+    k_steps = 2  # Reducing to 1 can be helpful while debugging.
     if out_jax_dtype == jnp.float16 and in_jax_dtype != jnp.float16:
       raise self.skipTest("Only f16 input is supported for f16 output.")
 
@@ -1552,7 +1555,7 @@ class TCGen05Test(TestCase):
         jax.ShapeDtypeStruct(tile_shape(x_block_shape, tiling), in_jax_dtype),
         jax.ShapeDtypeStruct(tile_shape(y_block_shape, tiling), in_jax_dtype),
         mgpu.TMABarrier(3),
-        mgpu.TMEM((128, n), out_jax_dtype, collective=True),
+        mgpu.TMEM((m_block_tile, n), out_jax_dtype, collective=True),
     ]
     z = mgpu.as_gpu_kernel(
         kernel, (2, 1, 1), (128, 1, 1), (x, y), out_shape, scratch_shape, cluster=(2, 1, 1)