jax-ml/jax · error · ValueError

float64 dot requires M>=16, N>=8, K>=16 per warp tile (Trito

Error message

float64 dot requires M>=16, N>=8, K>=16 per warp tile (Triton MMAv2 m8n8k4 layout); got {', '.join(errors)}

What it means

Triton's fp64 matrix multiply (MMAv2 m8n8k4 PTX path) imposes minimum per-warp tile dimensions of M>=16, N>=8, K>=16. The Pallas lowering validates this for float64 dots and raises ValueError listing which dimensions are too small.

Source

Thrown at jax/_src/pallas/triton/lowering.py:2424

  m, k = a_type.shape
  _, n = b_type.shape
  if a_type.element_type == ir.F64Type.get():
    # Triton's MMAv2 fp64 path uses the m8n8k4 PTX instruction but aggregates
    # it with NumRegisters={m:2, n:1, k:4}, producing an effective m16n8k16
    # per-warp tile.  Blocks smaller than these minimums cause repM/repN/repK
    # to round to zero, corrupting the ValueTable and segfaulting the compiler.
    #   M >= 16  (2 × instrM=8)
    #   N >=  8  (1 × instrN=8)
    #   K >= 16  (4 × instrK=4)
    errors = []
    if m < 16:
      errors.append(f"M={m} < 16")
    if n < 8:
      errors.append(f"N={n} < 8")
    if k < 16:
      errors.append(f"K={k} < 16")
    if errors:
      raise ValueError(
          f"float64 dot requires M>=16, N>=8, K>=16 per warp tile "
          f"(Triton MMAv2 m8n8k4 layout); got {', '.join(errors)}"
      )

  if a_type.element_type != b_type.element_type:
    raise ValueError(
        "a and b must have the same element type, but got:"
        f" {a_type.element_type} and {b_type.element_type}"
    )

  assert acc_dtype is not None
  acc = _zeros(ir.RankedTensorType.get([m, n], _dtype_to_ir_type(acc_dtype)))

  if precision in (
      lax.DotAlgorithmPreset.BF16_BF16_F32_X3,
      lax.DotAlgorithmPreset.BF16_BF16_F32_X6,
      lax.DotAlgorithmPreset.BF16_BF16_F32_X9,
  ):

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Increase the dot tile so M>=16, N>=8, K>=16 (e.g. use 16x16x16 minimum f64 tiles)
Switch inputs to float32 if full double precision isn't required
Pad the K dimension to a multiple >=16 if the natural K is smaller

Example fix

# before
BLOCK_M, BLOCK_N, BLOCK_K = 8, 8, 8  # f64 kernel
acc = pl.dot(a, b, out_dtype=jnp.float64)

# after
BLOCK_M, BLOCK_N, BLOCK_K = 16, 16, 16
acc = pl.dot(a, b, out_dtype=jnp.float64)

Defensive patterns

Strategy: validation

Validate before calling

def f64_tile_ok(m, n, k):
    return m >= 16 and n >= 8 and k >= 16
assert f64_tile_ok(BLOCK_M, BLOCK_N, BLOCK_K)

Type guard

def valid_f64_tile(m: int, n: int, k: int) -> bool:
    return m >= 16 and n >= 8 and k >= 16

Prevention

Default f64 kernels to 16x16x16 tiles or larger
Consider float32 unless double precision is strictly required

When it happens

Trigger: A pallas kernel doing an f64 dot with warp tile shapes like M=8, N=4, or K=8, e.g. tl.dot on float64 blocks of shape (8, 8).

Common situations: Tuning block sizes down for memory savings in a double-precision kernel; reusing f32 block sizes (which allow smaller tiles) in an f64 kernel; running on GPUs where fp64 needs the MMAv2 layout.

Related errors

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