{"record":{"id":"98ad3464f43d538a","repo":"jax-ml/jax","slug":"block-q-must-be-a-multiple-of-num-lanes","errorCode":null,"errorMessage":"block_q must be a multiple of {NUM_LANES}","messagePattern":"block_q must be a multiple of (.+?)","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_kernel.py","lineNumber":674,"sourceCode":"    if computed_mask.dtype != jnp.dtype(jnp.bool_):\n      raise ValueError(\n          \"Mask function must return a boolean-valued array, but got:\"\n          f\" {computed_mask.dtype}\"\n      )\n    masks.append(computed_mask)\n\n  if q_segment_ids_ref is not None:\n    if k_in_lanes:\n      kv_ids = kv_segment_ids_ref[:1, k_slice]  # [1, k_slice]\n      repeats, rem = divmod(kv_ids.shape[1], NUM_LANES)\n      if rem:\n        raise NotImplementedError(f\"block_kv must be a multiple of {NUM_LANES}\")\n      q_ids = jnp.tile(q_segment_ids_ref[:], (1, repeats))  # [bq, bkv]\n    else:\n      assert bq == q_segment_ids_ref.shape[-1]\n      repeats, rem = divmod(bq, NUM_LANES)\n      if rem:\n        raise NotImplementedError(f\"block_q must be a multiple of {NUM_LANES}\")\n      kv_ids = jnp.tile(\n          kv_segment_ids_ref[k_slice, :], (1, repeats)\n      )  # [k_slice, bq]\n      q_ids = q_segment_ids_ref[:1, :]  # [1, bq]\n    masks.append(q_ids == kv_ids)\n\n  def cap_logits(logits):\n    if attn_logits_soft_cap is not None:\n      logits = jnp.tanh(qk / attn_logits_soft_cap)\n      return logits * attn_logits_soft_cap\n    else:\n      return logits\n\n  if masks:\n    mask = functools.reduce(jnp.logical_and, masks)\n    qk = cap_logits(qk)\n    qk = jnp.where(mask, qk, mask_value)\n  else:","sourceCodeStart":656,"sourceCodeEnd":692,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_kernel.py#L656-L692","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Set block_sizes.block_q to a multiple of 8 (e.g. 64, 128, 256)","Pad the query sequence length so a lane-aligned block_q fits","Drop segment_ids if segment masking is not required"],"exampleFix":"// before\nblock_sizes=BlockSizes(block_q=48, block_kv=128)\n// after\nblock_sizes=BlockSizes(block_q=64, block_kv=128)","handlingStrategy":"validation","validationCode":"assert block_sizes.block_q % 8 == 0, 'block_q must be multiple of 8 when using segment_ids'","typeGuard":"def valid_bq(bq: int) -> bool: return bq % 8 == 0","tryCatchPattern":null,"preventionTips":["Use multiples of 8 for every block dimension","Don't hand-tune block sizes to odd sequence lengths; pad instead"],"tags":["jax","pallas","tpu","splash-attention","segment-ids","block-size"],"backgroundTag":"block-size-constraint-violation","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}