jax-ml/jax · error · ValueError

4-bit block scaled MMA only supports K-fastest operands…

Error message

4-bit block scaled MMA only supports K-fastest operands, but B is N-fastest

What it means

4-bit block-scaled MMA computes scale groups along K, so the B operand must also be K-fastest. It fires when B's layout is N-fastest, after the corresponding A check passed.

Solutions

Store/transpose B so K is the fastest dimension
Adjust b_swizzle to a value compatible with the K-fastest layout

Example fix

# before
b = weights  # N-fastest
tcgen05.mma(a, b, d, a_scale=asc, b_scale=bsc, scale_block=16)
# after
b_t = transpose_to_k_fastest(weights)
tcgen05.mma(a, b_t, d, a_scale=asc, b_scale=bsc, scale_block=16)

Defensive patterns

Strategy: validation

Validate before calling

assert b_fastest == mma_utils.Dim.K, '4-bit scaled MMA needs K-fastest B'

Prevention

Pre-transpose weights to K-major at quantization time
Assert b_fastest for 4-bit scaled paths

When it happens

Trigger: Calling mma() with is_scaled=True, 4-bit operands, and b_fastest == Dim.N (B stored row-major over (N, K) transposed).

Common situations: Standard GEMM weight layout (N-major B) reused from an 8-bit kernel; missing transpose of the weights for NVFP4.

Related errors

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

Appendix: source

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

      (b_desc_base, b_k_instr_strides),
      (b_n_group_stride, b_k_group_stride),
      b_fastest,
  ) = mma_utils.create_descriptor(
      b,
      swizzle=b_swizzle,
      group_size=(k_group_elems, n_group_elems),
      logical_k_major=True,
      mma_bytewidth_k=64 if is_sparse else 32,
      split_const=True,
  )

  if is_scaled and utils.bitwidth(mma_a_element_type) == 4:
    if a_fastest != mma_utils.Dim.K:
      raise ValueError(
          "4-bit block scaled MMA only supports K-fastest operands, but A is M-fastest"
      )
    if b_fastest != mma_utils.Dim.K:
      raise ValueError(
          "4-bit block scaled MMA only supports K-fastest operands, but B is N-fastest"
      )
  if is_sparse:
    if b_swizzle == 32 and b_fastest == mma_utils.Dim.K:
      raise NotImplementedError(
          "B tiling too small. Increase swizzle or transpose the input."
      )

  # Step 4. Issue the instructions.
  true = arith.constant(ir.IntegerType.get_signless(1), 1)
  n_collective_group_elems = n_group_elems * num_cta
  n_col_groups = n_groups // n_lane_groups
  assert d.layout.base_tile_shape[0] % 4 == 0
  lanes_per_n_group = d.layout.base_tile_shape[0] // 4
  a_sparse_addr_base = a_sparse_metadata.address if is_sparse else None
  a_scale_addr_base = a_scale.address if is_scaled else None  # pyrefly: ignore[missing-attribute]
  b_scale_addr_base = b_scale.address if is_scaled else None  # pyrefly: ignore[missing-attribute]
  # B scales are padded when N is short, so it can't be derived from n_collective_group_elems.

View on GitHub (pinned to 1e1c6a8fc0)