{"record":{"id":"8ae6bf1b22ba01b0","repo":"xai-org/x-algorithm","slug":"kv-seq-len-must-be-a-multiple-of-block-kv","errorCode":null,"errorMessage":"{kv_seq_len=} must be a multiple of {block_kv=}","messagePattern":"(.+?) must be a multiple of (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"phoenix/xrex/pallas/ranker_attention_fa3.py","lineNumber":113,"sourceCode":"        raise ValueError(f\"Expected {k.shape=} to be {kv_shape} (inferred from q)\")\n    if v.shape != kv_shape:\n        raise ValueError(f\"Expected {v.shape=} to be {kv_shape} (inferred from q)\")\n    if (dtype := q.dtype) != k.dtype or dtype != v.dtype:\n        raise ValueError(\n            f\"q, k, and v should all have the same dtype, got: {q.dtype}, {k.dtype}, {v.dtype}\"\n        )\n    if num_q_heads % num_kv_heads:\n        raise ValueError(f\"{num_q_heads=} must be divisible by and {num_kv_heads=}\")\n    q_heads_per_kv_head = num_q_heads // num_kv_heads\n    if head_dim % 64:\n        raise ValueError(f\"{head_dim=} must be divisible by 64\")\n    if jnp.dtype(dtype) not in map(jnp.dtype, [jnp.float16, jnp.bfloat16]):\n        raise NotImplementedError(f\"Only f16 and bf16 are supported, got dtype: {dtype}\")\n\n    max_concurrent_steps = min(config.max_concurrent_steps, kv_seq_len // config.block_kv)\n    block_q, block_kv = config.block_q, config.block_kv\n    if kv_seq_len % block_kv:\n        raise ValueError(f\"{kv_seq_len=} must be a multiple of {block_kv=}\")\n\n    def kernel(q_ref, k_ref, v_ref, bound_ref, out_ref, lse_ref, scoped):\n        batch = lax.axis_index(\"batch\")\n        q_head = lax.axis_index(\"heads\")\n        q_seq = lax.axis_index(\"q_seq\")\n        smem_buffers, buffer_barriers, consumed_barriers, schedule_barrier = scoped\n        wg_idx = lax.axis_index(\"wg\")\n        qo_smem2, k_smem, v_smem, lse_smem2 = smem_buffers\n        k_barriers, v_barriers, q_barriers = buffer_barriers\n        k_consumed_barriers, v_consumed_barriers = consumed_barriers\n        history_lower_bound = plgpu.load(bound_ref, (batch, 0))\n        history_upper_bound = plgpu.load(bound_ref, (batch, 1))\n        candidate_lower_bound = plgpu.load(bound_ref, (batch, 2))\n        candidate_upper_bound = plgpu.load(bound_ref, (batch, 3))\n\n        def perform_schedule_barrier():\n            plgpu.barrier_arrive(schedule_barrier)\n            plgpu.barrier_wait(schedule_barrier)","sourceCodeStart":95,"sourceCodeEnd":131,"githubUrl":"https://github.com/xai-org/x-algorithm/blob/24c60942c5c5fdad3a6addffb4c6e6d2f228f04f/phoenix/xrex/pallas/ranker_attention_fa3.py#L95-L131","documentation":"The kv sequence length must be a multiple of the configured block_kv so the kernel can tile the kv loop without a remainder tile (no ragged-tile handling is implemented).","triggerScenarios":"kv_seq_len=1000 with block_kv=128, or any (kv_seq_len % block_kv) != 0; also seq lens changed by chunked prefill or cache truncation.","commonSituations":"Arbitrary seq lens from tokenized batches; changing block_kv in TuningConfig to a size that no longer divides the padded context length.","solutions":["Pad kv sequence to a multiple of block_kv (with a bound/mask so padding is ignored)","Or choose block_kv that divides kv_seq_len (e.g. 64 divides more lengths)","Keep padding consistent for q as well if needed"],"exampleFix":"# before\nattn = attention(q, k, v, config=cfg)  # kv_seq_len=1000, block_kv=128\n# after\npad = (-kv_seq_len) % cfg.block_kv\nk, v = jnp.pad(k, ((0,0),(0,pad),(0,0),(0,0))), jnp.pad(v, ((0,0),(0,pad),(0,0),(0,0)))\nattn = attention(q, k, v, config=cfg, bound=bound)","handlingStrategy":"validation","validationCode":"pad = (-kv_seq_len) % config.block_kv\nif pad: k, v = pad_kv(k, pad), pad_kv(v, pad)  # and pass bound to mask padding","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Pad all sequence lengths to the kernel block lcm in the data pipeline","Write a pad_to_block(seq, block) utility reused for q and kv"],"tags":["jax","pallas","attention","shape-validation"],"backgroundTag":"sequence-length-not-aligned","analyzedSha":"24c60942c5c5fdad3a6addffb4c6e6d2f228f04f","analyzedAt":"2026-08-28T11:40:14.686Z","schemaVersion":2},"datasetVersion":"2026-08-28T16:17:29.566Z"}