jax-ml/jax · error · NotImplementedError

block_q must be a multiple of {NUM_LANES}

Error message

block_q must be a multiple of {NUM_LANES}

What it means

When segment_ids are used and KV values are packed per-lane (not k_in_lanes), the query segment ids must be tiled in groups of NUM_LANES=8, so block_q must be divisible by 8. Otherwise the kernel cannot align query ids with KV lanes and raises NotImplementedError.

Source

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

    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

  if masks:
    mask = functools.reduce(jnp.logical_and, masks)
    qk = cap_logits(qk)
    qk = jnp.where(mask, qk, mask_value)
  else:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Set block_sizes.block_q to a multiple of 8 (e.g. 64, 128, 256)
2. Pad the query sequence length so a lane-aligned block_q fits
3. Drop segment_ids if segment masking is not required

Example fix

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

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

def valid_bq(bq: int) -> bool: return bq % 8 == 0

When it happens

Trigger: Calling splash attention with segment_ids and block_sizes.block_q % 8 != 0 (e.g. block_q=50) when the kernel takes the non-k_in_lanes branch for segment id comparison.

Common situations: Tuning block_q to exactly match a small query sequence length (e.g. seq_len=48 heads config with block_q=48); migrating configs between splash attention versions with different layout requirements.

Related errors

• block_kv 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/98ad3464f43d538a. Report an issue: GitHub.