jax-ml/jax · error · NotImplementedError

block_kv must be a multiple of {NUM_LANES}

Error message

block_kv must be a multiple of {NUM_LANES}

What it means

When using segment_ids with Splash Attention on TPU, key/value sequence block sizes must be laid out in SIMD lanes (NUM_LANES=8). If block_kv is not divisible by NUM_LANES, the kernel cannot tile the KV segment ids to compare against query ids, so it raises NotImplementedError.

Source

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

      )
      q_sequence = q_sequence_ref[:1, :]  # [1, bq]
      q_sequence = jnp.broadcast_to(q_sequence, (k_slice.size, bq))

    assert q_sequence.shape == k_sequence.shape
    computed_mask = mask_function(q_sequence, k_sequence)
    if computed_mask.dtype != jnp.dtype(jnp.bool_):
      raise ValueError(
          "Mask function must return a boolean-valued array, but got:"
          f" {computed_mask.dtype}"
      )
    masks.append(computed_mask)

  if q_segment_ids_ref is not None:
    if k_in_lanes:
      kv_ids = kv_segment_ids_ref[:1, k_slice]  # [1, k_slice]
      repeats, rem = divmod(kv_ids.shape[1], NUM_LANES)
      if rem:
        raise NotImplementedError(f"block_kv must be a multiple of {NUM_LANES}")
      q_ids = jnp.tile(q_segment_ids_ref[:], (1, repeats))  # [bq, bkv]
    else:
      assert bq == q_segment_ids_ref.shape[-1]
      repeats, rem = divmod(bq, NUM_LANES)
      if rem:
        raise NotImplementedError(f"block_q must be a multiple of {NUM_LANES}")
      kv_ids = jnp.tile(
          kv_segment_ids_ref[k_slice, :], (1, repeats)
      )  # [k_slice, bq]
      q_ids = q_segment_ids_ref[:1, :]  # [1, bq]
    masks.append(q_ids == kv_ids)

  def cap_logits(logits):
    if attn_logits_soft_cap is not None:
      logits = jnp.tanh(qk / attn_logits_soft_cap)
      return logits * attn_logits_soft_cap
    else:
      return logits

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Set block_sizes.block_kv to a multiple of 8 (e.g. 64, 128, 256)
2. If you don't need segment-based masking, drop the segment_ids argument
3. Use tpu_attention.DEFAULT_MASK_AND_BIAS_BLOCK_SIZES-style defaults or the library's recommended block size helper

Example fix

// before
block_sizes=BlockSizes(block_q=128, block_kv=100)
// after
block_sizes=BlockSizes(block_q=128, block_kv=128)

NUM_LANES = 8
assert block_sizes.block_kv % NUM_LANES == 0, 'block_kv must be multiple of 8 when using segment_ids'

def valid_block_sizes(bs) -> bool:
    return bs.block_kv % 8 == 0 and bs.block_q % 8 == 0

When it happens

Trigger: Calling splash attention with both segment_ids and block_sizes where block_sizes.block_kv % 8 != 0, e.g. block_kv=128 is fine but block_kv=100 or block_kv=12 raises this when k_in_lanes packing applies.

Common situations: Copying block_sizes tuned for GPU flash attention (e.g. block_kv=64 usually fine, but odd experimental sizes like 7 or 48*odd) into TPU splash attention; short KV sequences forcing small blocks.

Related errors

• block_q must be a multiple of {NUM_LANES}
• {bkv_compute=} should be a multiple of {NUM_LANES}
• {bkv=} must be a multiple of {bkv_compute=}.
• {bkv_compute=} must be a multiple of {NUM_LANES}.
• Invalid shape for q segment_ids: {segment_ids.q.shape}. Expe

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