{"record":{"id":"546b29eab4d56ca8","repo":"jax-ml/jax","slug":"bkv-compute-should-be-a-multiple-of-num-lanes","errorCode":null,"errorMessage":"{bkv_compute=} should be a multiple of {NUM_LANES}","messagePattern":"(.+?) should 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":803,"sourceCode":"        # the kv_index program_id does not correspond to the actual coordinates\n        # of the KV data. Make sure to use the 'unshrunk' index (coming from the\n        # data_next array) when computing the mask.\n        k_offset=global_kv_index * bkv + kv_compute_index * bkv_compute,\n        bq=bq,\n        mask_function=mask_function,\n    )\n\n    qk = apply_mask_and_soft_cap()\n    assert not isinstance(qk, tuple)\n\n    m_curr = qk.max(axis=-1)[:, None]\n    assert m_curr.shape == (bq, 1)\n    m_next = jnp.maximum(m_prev, m_curr)\n    assert m_next.shape == (bq, NUM_LANES)\n\n    bkv_repeats, rem = divmod(bkv_compute, NUM_LANES)\n    if rem != 0:\n      raise NotImplementedError(\n          f\"{bkv_compute=} should be a multiple of {NUM_LANES}\"\n      )\n\n    s_curr = jnp.exp(qk - jnp.tile(m_next, (1, bkv_repeats)))\n    assert s_curr.shape == (bq, bkv_compute)\n\n    l_curr = jax.lax.broadcast_in_dim(s_curr.sum(axis=-1), l_prev.shape, (0,))\n    assert l_curr.shape == (bq, NUM_LANES)\n\n    alpha = jnp.exp(m_prev - m_next)\n    l_next = l_curr + alpha * l_prev\n    m_scratch_ref[...], l_scratch_ref[...] = m_next, l_next\n\n    sv_dims = NN_DIM_NUMBERS if v_layout == HEAD_DIM_MINOR else NT_DIM_NUMBERS\n    if v_layout == HEAD_DIM_MINOR:\n      v = v_ref[slice_k, :]\n    else:\n      v = v_ref[:, slice_k]","sourceCodeStart":785,"sourceCodeEnd":821,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_kernel.py#L785-L821","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Set block_kv_compute to a multiple of 8 (typically block_kv_compute == block_kv, both multiples of 8)","If unsure, pass block_sizes=None and let the library choose default block sizes","Upgrade/downgrade JAX so the BlockSizes defaults match your tuning"],"exampleFix":"// before\nblock_sizes=BlockSizes(block_q=256, block_kv=256, block_kv_compute=124)\n// after\nblock_sizes=BlockSizes(block_q=256, block_kv=256, block_kv_compute=128)","handlingStrategy":"validation","validationCode":"assert block_sizes.block_kv_compute % 8 == 0 or block_sizes.block_kv_compute is None","typeGuard":"def compute_block_ok(bs) -> bool:\n    return bs.block_kv_compute is None or bs.block_kv_compute % 8 == 0","tryCatchPattern":null,"preventionTips":["Default to block_kv_compute == block_kv","Validate all BlockSizes fields against NUM_LANES=8 in a single helper"],"tags":["jax","pallas","tpu","splash-attention","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"}