jax-ml/jax · error · NotImplementedError

{bkv_compute=} should be a multiple of {NUM_LANES}

Error message

{bkv_compute=} should be a multiple of {NUM_LANES}

What it means

Inside the forward attention kernel body, bkv_compute (the KV compute block size used for tiled softmax) must be divisible by NUM_LANES=8 so the running max can be tiled across lanes. If not, the kernel raises NotImplementedError.

Source

Thrown at jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_kernel.py:803

        # the kv_index program_id does not correspond to the actual coordinates
        # of the KV data. Make sure to use the 'unshrunk' index (coming from the
        # data_next array) when computing the mask.
        k_offset=global_kv_index * bkv + kv_compute_index * bkv_compute,
        bq=bq,
        mask_function=mask_function,
    )

    qk = apply_mask_and_soft_cap()
    assert not isinstance(qk, tuple)

    m_curr = qk.max(axis=-1)[:, None]
    assert m_curr.shape == (bq, 1)
    m_next = jnp.maximum(m_prev, m_curr)
    assert m_next.shape == (bq, NUM_LANES)

    bkv_repeats, rem = divmod(bkv_compute, NUM_LANES)
    if rem != 0:
      raise NotImplementedError(
          f"{bkv_compute=} should be a multiple of {NUM_LANES}"
      )

    s_curr = jnp.exp(qk - jnp.tile(m_next, (1, bkv_repeats)))
    assert s_curr.shape == (bq, bkv_compute)

    l_curr = jax.lax.broadcast_in_dim(s_curr.sum(axis=-1), l_prev.shape, (0,))
    assert l_curr.shape == (bq, NUM_LANES)

    alpha = jnp.exp(m_prev - m_next)
    l_next = l_curr + alpha * l_prev
    m_scratch_ref[...], l_scratch_ref[...] = m_next, l_next

    sv_dims = NN_DIM_NUMBERS if v_layout == HEAD_DIM_MINOR else NT_DIM_NUMBERS
    if v_layout == HEAD_DIM_MINOR:
      v = v_ref[slice_k, :]
    else:
      v = v_ref[:, slice_k]

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Set block_kv_compute to a multiple of 8 (typically block_kv_compute == block_kv, both multiples of 8)
2. If unsure, pass block_sizes=None and let the library choose default block sizes
3. Upgrade/downgrade JAX so the BlockSizes defaults match your tuning

Example fix

// before
block_sizes=BlockSizes(block_q=256, block_kv=256, block_kv_compute=124)
// after
block_sizes=BlockSizes(block_q=256, block_kv=256, block_kv_compute=128)

assert block_sizes.block_kv_compute % 8 == 0 or block_sizes.block_kv_compute is None

def compute_block_ok(bs) -> bool:
    return bs.block_kv_compute is None or bs.block_kv_compute % 8 == 0

When it happens

Trigger: Calling splash attention with block_sizes.block_kv_compute (bkv_compute) not divisible by 8, e.g. block_kv_compute=100 while block_kv=128.

Common situations: Manually splitting block_kv and block_kv_compute for memory tiling on TPU v4/v5e; using defaults from an older JAX version that did not enforce this.

Related errors

• block_kv must be a multiple of {NUM_LANES}
• block_q must be a multiple of {NUM_LANES}
• {bkv=} must be a multiple of {bkv_compute=}.
• {bkv_compute=} must be a multiple of {NUM_LANES}.
• {bq=} should not be greater than {q_seq_len=}

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