{"record":{"id":"b4f87df82e68777a","repo":"jax-ml/jax","slug":"block-kv-must-be-a-multiple-of-num-lanes","errorCode":null,"errorMessage":"block_kv must be a multiple of {NUM_LANES}","messagePattern":"block_kv 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":668,"sourceCode":"      )\n      q_sequence = q_sequence_ref[:1, :]  # [1, bq]\n      q_sequence = jnp.broadcast_to(q_sequence, (k_slice.size, bq))\n\n    assert q_sequence.shape == k_sequence.shape\n    computed_mask = mask_function(q_sequence, k_sequence)\n    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","sourceCodeStart":650,"sourceCodeEnd":686,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_kernel.py#L650-L686","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Set block_sizes.block_kv to a multiple of 8 (e.g. 64, 128, 256)","If you don't need segment-based masking, drop the segment_ids argument","Use tpu_attention.DEFAULT_MASK_AND_BIAS_BLOCK_SIZES-style defaults or the library's recommended block size helper"],"exampleFix":"// before\nblock_sizes=BlockSizes(block_q=128, block_kv=100)\n// after\nblock_sizes=BlockSizes(block_q=128, block_kv=128)","handlingStrategy":"validation","validationCode":"NUM_LANES = 8\nassert block_sizes.block_kv % NUM_LANES == 0, 'block_kv must be multiple of 8 when using segment_ids'","typeGuard":"def valid_block_sizes(bs) -> bool:\n    return bs.block_kv % 8 == 0 and bs.block_q % 8 == 0","tryCatchPattern":null,"preventionTips":["Keep all block sizes multiples of 8 (NUM_LANES)","Centralize block-size config in one validated dataclass","When using segment_ids, always double-check lane alignment"],"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"}