jax-ml/jax · error · ValueError

B scale shape[0] must be a multiple of 128 and >= N={n * num

Error message

B scale shape[0] must be a multiple of 128 and >= N={n * num_cta}, got {b_scale.shape[0]}

What it means

The B scale tensor's leading dimension must be a multiple of 128 and at least n * num_cta to satisfy TMEM allocation granularity for block-scaled MMA. This error reports the offending b_scale.shape[0].

Source

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

          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
    if a_sparse_metadata.shape != (m, expected_meta_k):
      raise ValueError(
          f"A sparse metadata shape mismatch: expected {(m, expected_meta_k)},"
          f" got {a_sparse_metadata.shape}"
      )
    if a_sparse_metadata.dtype != ir.IntegerType.get_signless(2):

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Pad b_scale.shape[0] up to the next multiple of 128 that is >= n * num_cta
Double-check n and num_cta values passed to mma match the shape math

Example fix

# before
n, num_cta = 256, 2
b_scale = make_scales(shape=(128, k_scales))
# after
b_scale = make_scales(shape=(512, k_scales))  # multiple of 128 and >= 512

Defensive patterns

Strategy: validation

Validate before calling

rows = ((n * num_cta + 127) // 128) * 128
assert b_scale.shape[0] % 128 == 0 and b_scale.shape[0] >= n * num_cta, f'use {rows}'

Type guard

def valid_b_scale_rows(rows: int, n: int, num_cta: int) -> bool:
    return rows % 128 == 0 and rows >= n * num_cta

Prevention

Allocate B scales with padding to 128-row granularity
Pass num_cta consistently with the warp-specialization config

When it happens

Trigger: Passing b_scale whose shape[0] is not a multiple of 128 (e.g. 64, 100) or smaller than N (n * num_cta), e.g. N=256 with shape[0]=128.

Common situations: Allocating B scales sized exactly to N when N is not a multiple of 128; forgetting to scale up the N dimension for multi-CTA (num_cta=2) kernels.

Related errors

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