jax-ml/jax · error · ValueError

Columns must be a multiple of 16, got: {columns}

Error message

Columns must be a multiple of 16, got: {columns}

What it means

tmem_half_lane_layout requires the column count to be a multiple of 16 because the underlying TMEM tiling for the 64-row layout uses 16-column granularity; a non-multiple would leave a partial tile that the hardware layout cannot express.

Source

Thrown at jax/experimental/mosaic/gpu/tcgen05.py:1101

def tmem_default_layout(packing: int = 1) -> TMEMLayout:
  """A TMEM layout used for 1CTA MMA with M=128 and 2CTA MMA with M=256."""
  if packing.bit_count() != 1:
    raise ValueError(f"Packing must be a power of 2, got: {packing}")
  return TMEMLayout(
      fa.Tiling(((TMEM_ROWS, packing), (fa.WARP_SIZE, packing))),
      warp_dims=(-4,),
      lane_dims=(-2,),
      vector_dim=-1,
  )


def tmem_half_lane_layout(columns, packing: int = 1) -> TMEMLayout:
  """A TMEM layout used for 1CTA MMA with M=64."""
  if packing > (columns // 2) 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 // 2, columns // 2),
          (packing,),
      )),
      warp_dims=(-5,),
      lane_dims=(-4, -3),
      vector_dim=-1,
  )


def tmem_m64_collective_layout(columns: int, packing: int = 1) -> TMEMLayout:
  """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}")

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pad N (columns) up to the next multiple of 16
2. Choose tile sizes where N % 16 == 0 (e.g. 16, 32, 64)
3. If partial results are needed, load the full padded region and mask in registers

Example fix

# before
layout = tcgen05.tmem_half_lane_layout(columns=24)
# after
layout = tcgen05.tmem_half_lane_layout(columns=32)

assert columns % 16 == 0, f'columns must be multiple of 16, got {columns}'

def columns_ok(columns: int) -> bool:
    return columns % 16 == 0

When it happens

Trigger: Calling tmem_half_lane_layout(columns=10) or any columns % 16 != 0; indirectly from tcgen05.load/store or infer_tmem_layout on a 64-row TMEMRef whose second dimension isn't a multiple of 16.

Common situations: MMA with M=64 and an N dimension like 8, 24, or 40 (common with small tiles or f8 dtypes) resulting in a TMEM shape violating the 16-column alignment.

Related errors

• Cannot assign layout to async load with gather indices since
• Minor dimension of shape must be divisible by packing, got:
• Packing must be a power of 2, got: {packing}
• Cannot slice TMEM with multiple tiles along rows.
• TMEM layout {self.layout} is not supported

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/e8d02de0c1e7250b. Report an issue: GitHub.