jax-ml/jax · error · ValueError

Expected B scales to have a M=128 layout, got {b_scale.layou

Error message

Expected B scales to have a M=128 layout, got {b_scale.layout}

What it means

For block-scaled MMA with M=128, the B scale tensor must use the standard scales_layout(). This ValueError fires when any other layout is attached to b_scale in the m=128 path.

Source

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

          f"Unsupported element type for block scaling: {a_element_type}"
      )
    k_scales = k // scale_block
    if a_scale.shape != (TMEM_ROWS, k_scales):
      raise ValueError(
          f"A scale shape mismatch: expected ({TMEM_ROWS}, {k_scales}), got"
          f" {a_scale.shape}"
      )
    if a_scale.layout != scales_layout():
      raise ValueError(f"A scale layout {a_scale.layout} is not supported")
    if collective and m == 64:
      if b_scale.layout != b_scales_m64_collective_layout():
        raise ValueError(
            "Expected B scales to have a M=64 collective layout, got"
            f" {b_scale.layout}"
        )
    elif m == 128:
      if b_scale.layout != scales_layout():
        raise ValueError(
            f"Expected B scales to have a M=128 layout, got {b_scale.layout}"
        )
    else:
      raise AssertionError("Should not happen")
    if b_scale.shape[0] % 128 or b_scale.shape[0] < n * num_cta:
      raise ValueError(
          f"B scale shape[0] must be a multiple of 128 and >= N={n * num_cta},"
          f" got {b_scale.shape[0]}"
      )
    if b_scale.shape[1] != k_scales:
      raise ValueError(
          f"B scale shape mismatch: expected ({b_scale.shape[0]}, {k_scales}),"
          f" got {b_scale.shape}"
      )
  if is_sparse:
    sparse_group_elems = 8 if utils.bitwidth(a_element_type) == 4 else 4
    # Each sparse group has 2 entries.
    expected_meta_k = k // sparse_group_elems * 2

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use scales_layout() for b_scale when m=128
2. If you need collective M=64, keep b_scales_m64_collective_layout() and m=64 consistent

Example fix

# before
b_scale = TensorMemRefView(buf, shape, dt, layout=b_scales_m64_collective_layout())
# after (m=128)
b_scale = TensorMemRefView(buf, shape, dt, layout=scales_layout())

assert b_scale.layout == scales_layout() if m == 128 else True

When it happens

Trigger: Calling mma() with m=128 and b_scale.layout != scales_layout().

Common situations: Switching an M=64 collective kernel to M=128 without changing the B scale layout back; applying swizzled operand layouts to scales.

Related errors

• 4-bit block scaled MMA only supports K-fastest operands, but
• 4-bit block scaled MMA only supports K-fastest operands, but
• MMA with element type {elem_type_str} does not support block
• Scale element type mismatch: expected f8e8m0fnu or f8e4m3fn,
• Unsupported element type for block scaling: {a_element_type}

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