jax-ml/jax · error · NotImplementedError

TPU version must be 4 or higher.

Error message

TPU version must be 4 or higher.

What it means

get_tuned_block_sizes looks up auto-tuned (num_kv_pages_per_blk, num_queries_per_blk) values for the TPU ragged paged attention kernel; the tuning tables only cover TPU v4 and later. On older TPUs (v2/v3) or when the TPU version is detected as < 4 it raises NotImplementedError.

Source

Thrown at jax/experimental/pallas/ops/tpu/ragged_paged_attention/tuned_block_sizes.py:1458

  if num_devices is not None:
    name += f'-{num_devices}'
  return name


def get_tuned_block_sizes(
    q_dtype,
    kv_dtype,
    num_q_heads_per_blk,
    num_kv_heads_per_blk,
    head_dim,
    page_size,
    max_num_batched_tokens,
    pages_per_seq,
) -> tuple[int, int]:
  """Look up for the best (num_kv_pages_per_blk, num_queries_per_blk) from auto-tuned table."""
  tpu_version = get_tpu_version()
  if tpu_version < 4:
    raise NotImplementedError('TPU version must be 4 or higher.')
  key = (
      q_dtype,
      kv_dtype,
      num_q_heads_per_blk,
      num_kv_heads_per_blk,
      head_dim,
      page_size,
      max_num_batched_tokens,
      pages_per_seq,
  )
  key = simplify_key(key)
  device_name = get_device_name()

  # Default block sizes.
  bkv, bq = (128, 32)
  if tpu_version == 4:
    # This default block size is not tuned, only make sure there's no
    # OOM in vmem

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Run on TPU v4 or newer (v4, v5e, v5p, v6e)
2. Explicitly pass num_kv_pages_per_block and num_queries_per_block if the API path allows bypassing the tuned table
3. Check jax.devices() to confirm you are actually attached to a TPU and its version
4. Update libtpu/JAX so TPU version detection is correct

import jax
devs = jax.devices()
assert all(d.platform == 'tpu' for d in devs), 'requires TPU'

def is_tpu_v4_plus() -> bool:
    try:
        from jax.experimental.pallas.ops.tpu.ragged_paged_attention.tuned_block_sizes import get_tpu_version
        return get_tpu_version() >= 4
    except Exception:
        return False

try:
    sizes = get_tuned_block_sizes(...)
except NotImplementedError:
    sizes = None  # fall back to default block sizes

When it happens

Trigger: Running ragged_paged_attention with automatic block-size tuning on TPU v2/v3 hardware, or on a runtime/SDK where get_tpu_version() misdetects or returns a low version.

Common situations: Deploying JAX code written for v4/v5p/v6e onto older Cloud TPU slices or an emulator; running on CPU/GPU while the pallas TPU path is still imported and invoked.

Related errors

• masked load_p
• run_scoped_p with collective axes is not supported
• Non-decrementing wait is not supported.
• Only take_along_axis-like gathers supported
• Unsupported gather

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