xai-org/x-algorithm · error · ValueError

{kv_seq_len=} must be a multiple of {block_kv=}

Error message

{kv_seq_len=} must be a multiple of {block_kv=}

What it means

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).

Source

Thrown at phoenix/xrex/pallas/ranker_attention_fa3.py:113

        raise ValueError(f"Expected {k.shape=} to be {kv_shape} (inferred from q)")
    if v.shape != kv_shape:
        raise ValueError(f"Expected {v.shape=} to be {kv_shape} (inferred from q)")
    if (dtype := q.dtype) != k.dtype or dtype != v.dtype:
        raise ValueError(
            f"q, k, and v should all have the same dtype, got: {q.dtype}, {k.dtype}, {v.dtype}"
        )
    if num_q_heads % num_kv_heads:
        raise ValueError(f"{num_q_heads=} must be divisible by and {num_kv_heads=}")
    q_heads_per_kv_head = num_q_heads // num_kv_heads
    if head_dim % 64:
        raise ValueError(f"{head_dim=} must be divisible by 64")
    if jnp.dtype(dtype) not in map(jnp.dtype, [jnp.float16, jnp.bfloat16]):
        raise NotImplementedError(f"Only f16 and bf16 are supported, got dtype: {dtype}")

    max_concurrent_steps = min(config.max_concurrent_steps, kv_seq_len // config.block_kv)
    block_q, block_kv = config.block_q, config.block_kv
    if kv_seq_len % block_kv:
        raise ValueError(f"{kv_seq_len=} must be a multiple of {block_kv=}")

    def kernel(q_ref, k_ref, v_ref, bound_ref, out_ref, lse_ref, scoped):
        batch = lax.axis_index("batch")
        q_head = lax.axis_index("heads")
        q_seq = lax.axis_index("q_seq")
        smem_buffers, buffer_barriers, consumed_barriers, schedule_barrier = scoped
        wg_idx = lax.axis_index("wg")
        qo_smem2, k_smem, v_smem, lse_smem2 = smem_buffers
        k_barriers, v_barriers, q_barriers = buffer_barriers
        k_consumed_barriers, v_consumed_barriers = consumed_barriers
        history_lower_bound = plgpu.load(bound_ref, (batch, 0))
        history_upper_bound = plgpu.load(bound_ref, (batch, 1))
        candidate_lower_bound = plgpu.load(bound_ref, (batch, 2))
        candidate_upper_bound = plgpu.load(bound_ref, (batch, 3))

        def perform_schedule_barrier():
            plgpu.barrier_arrive(schedule_barrier)
            plgpu.barrier_wait(schedule_barrier)

View on GitHub (pinned to 24c60942c5)

Solutions

1. Pad kv sequence to a multiple of block_kv (with a bound/mask so padding is ignored)
2. Or choose block_kv that divides kv_seq_len (e.g. 64 divides more lengths)
3. Keep padding consistent for q as well if needed

Example fix

# before
attn = attention(q, k, v, config=cfg)  # kv_seq_len=1000, block_kv=128
# after
pad = (-kv_seq_len) % cfg.block_kv
k, v = jnp.pad(k, ((0,0),(0,pad),(0,0),(0,0))), jnp.pad(v, ((0,0),(0,pad),(0,0),(0,0)))
attn = attention(q, k, v, config=cfg, bound=bound)

pad = (-kv_seq_len) % config.block_kv
if pad: k, v = pad_kv(k, pad), pad_kv(v, pad)  # and pass bound to mask padding

When it happens

Trigger: 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.

Common situations: Arbitrary seq lens from tokenized batches; changing block_kv in TuningConfig to a size that no longer divides the padded context length.

Related errors

• q, k, and v should all be 4D, got: {q.ndim=}, {k.ndim=}, {v.
• {q_seq_len=} must be a multiple of {block_q * 2=}
• Expected bound to have shape (4,), got {bound_arr.shape}
• {q_seq_len=} must be a multiple of {config.block_q_dq=} * {c
• Invalid backward pass implementation: {backward_pass_impl}

AI-assisted analysis of xai-org/x-algorithm@24c60942c5 (2026-08-28). Data as JSON: /api/errors/8ae6bf1b22ba01b0. Report an issue: GitHub.