xai-org/x-algorithm · error · NotImplementedError

{q_seq_len=} must be a multiple of {block_q * 2=}

Error message

{q_seq_len=} must be a multiple of {block_q * 2=}

What it means

The forward kernel processes q in tiles of 2*block_q (two query tiles per warp-group iteration), so q_seq_len must be a multiple of block_q*2; remainder tiles are unimplemented (NotImplementedError).

Source

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

        if save_residuals:
            scratch[3] = plgpu.SMEM((compute_wgs, block_q), jnp.float32)

        pl.run_scoped(
            lambda *args: kernel(q_ref, k_ref, v_ref, bound_ref, out_ref, lse_ref, args),
            scratch,
            (
                plgpu.Barrier(num_barriers=max_concurrent_steps),
                plgpu.Barrier(num_barriers=max_concurrent_steps),
                plgpu.Barrier(num_barriers=compute_wgs),
            ),
            (plgpu.Barrier(num_arrivals=compute_wgs, num_barriers=max_concurrent_steps),) * 2,
            plgpu.Barrier(num_arrivals=compute_wgs),
            collective_axes="wg",
        )

    num_q_tiles, rem = divmod(q_seq_len, block_q * 2)
    if rem:
        raise NotImplementedError(f"{q_seq_len=} must be a multiple of {block_q * 2=}")

    out_shape = [q, None]
    if save_residuals:
        out_shape[1] = jax.ShapeDtypeStruct((batch_size, num_q_heads, q_seq_len), jnp.float32)
    out, lse = plgpu.kernel(
        entry,
        out_shape=out_shape,
        grid=(num_q_heads, num_q_tiles, batch_size),
        grid_names=("heads", "q_seq", "batch"),
        num_threads=3,
        thread_name="wg",
        compiler_params=plgpu.CompilerParams(approx_math=True),
    )(q, k, v, bound)

    if save_residuals:
        assert lse is not None
        return out, (lse,)

View on GitHub (pinned to 24c60942c5)

Solutions

1. Pad q_seq_len (and correspondingly k/v) up to a multiple of 2*block_q and mask via bound
2. Or reduce block_q so 2*block_q divides q_seq_len

Example fix

# before
attn = attention(q, k, v, config=cfg)  # q_seq_len=200, block_q=128
# after
cfg = replace(cfg, block_q=64)  # 2*64=128 divides... or pad:
q = jnp.pad(q, ((0,0),(0,56),(0,0),(0,0)))  # pad to 256

divisor = config.block_q * 2
assert q_seq_len % divisor == 0 or plan_padding, f"pad q_seq_len to a multiple of {divisor}"

When it happens

Trigger: q_seq_len not divisible by 2*block_q, e.g. seq 100 with block_q=128 gives zero full tiles with remainder, or seq 200 with block_q=128 leaves remainder 72.

Common situations: Short local-attention windows in ranker models; small batch debugging with short prompts; increasing block_q in tuning so it no longer divides the seq len.

Related errors

• q, k, and v should all be 4D, got: {q.ndim=}, {k.ndim=}, {v.
• {kv_seq_len=} must be a multiple of {block_kv=}
• 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/b0ac0d412a51ce0a. Report an issue: GitHub.