jax-ml/jax · error · ValueError

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

Error message

4-bit block scaled MMA only supports K-fastest operands, but A is M-fastest

What it means

4-bit block-scaled MMA (e.g. MXFP4/NVFP4) hardware only supports operands whose fastest-varying dimension is K. If the A operand is laid out M-fastest, mma raises this error.

Source

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

    a_fastest = mma_utils.Dim.K
    a_k_instr_strides = None
    a_m_group_stride = a_k_group_stride = a_desc_base = None
  (
      (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

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Transpose A so K is the fastest-varying dimension and update swizzle accordingly
2. Use an 8-bit scaled type if an M-fastest A layout is required

Example fix

# before
a = TensorMemRefView(buf, (m, k), dt, layout=col_major)  # M-fastest
tcgen05.mma(a, b, d, a_scale=asc, b_scale=bsc, scale_block=16)
# after
a = TensorMemRefView(buf, (m, k), dt, layout=row_major)  # K-fastest
tcgen05.mma(a, b, d, a_scale=asc, b_scale=bsc, scale_block=16)

assert a_fastest == mma_utils.Dim.K, '4-bit scaled MMA needs K-fastest A'

When it happens

Trigger: Calling mma() with is_scaled=True, 4-bit A, and a layout where a_fastest is Dim.M (e.g. a transposed or col-major A).

Common situations: Reusing 8-bit kernel layouts for MXFP4; transposing A for coalescing purposes which flips the fastest dim.

Related errors

• 4-bit block scaled MMA only supports K-fastest operands, but
• Expected B scales to have a M=128 layout, got {b_scale.layou
• 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/cda5f1d2418fbfa7. Report an issue: GitHub.