xai-org/x-algorithm · error · ValueError

q, k, and v should all be 4D, got: {q.ndim=}, {k.ndim=}, {v.

Error message

q, k, and v should all be 4D, got: {q.ndim=}, {k.ndim=}, {v.ndim=}

What it means

_attention_forward requires q, k, v in BHSD layout (batch, seq, heads, head_dim) — exactly 4 dimensions each. This is the entry check before any shape inference happens.

Source

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

    @property
    def has_backward_blocks(self) -> bool:
        return self.block_q_dkv is not None


def _attention_forward(
    q,
    k,
    v,
    config: TuningConfig,
    save_residuals: bool = False,
    bound=None,
    sm_scale: float = 1.0,
    cap: float = -1.0,
    cap_method: str = "tanh",
):
    if q.ndim != 4 or k.ndim != 4 or v.ndim != 4:
        raise ValueError(f"q, k, and v should all be 4D, got: {q.ndim=}, {k.ndim=}, {v.ndim=}")
    batch_size, q_seq_len, num_q_heads, head_dim = q.shape
    _, kv_seq_len, num_kv_heads, _ = k.shape
    kv_shape = (batch_size, kv_seq_len, num_kv_heads, head_dim)
    if k.shape != kv_shape:
        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}")

View on GitHub (pinned to 24c60942c5)

Solutions

1. Reshape to 4D (batch, seq_len, num_heads, head_dim)
2. If you have (B, H, S, D), transpose axes 1 and 2
3. Add a head dimension of 1 for single-head inputs

Example fix

# before
out = attention(q, k, v)  # q is (B, H, S, D)
# after
q, k, v = (x.transpose(0, 2, 1, 3) for x in (q, k, v))
out = attention(q, k, v)

null

def is_bshd4(x) -> bool:
    return hasattr(x, "ndim") and x.ndim == 4

When it happens

Trigger: Passing 3D inputs (B, S, D without a head axis) or inputs still in (B, H, S, D) when the function expects (B, S, H, D), or accidentally batched extra leading dims.

Common situations: Layout confusion between BHSD and BSHD conventions across jax attention implementations; feeding outputs of einops.rearrange with the wrong pattern; passing single-example unbatched tensors.

Understand the failure class

Background: Tensor shape mismatch errors ("must have shape", "expected shape ... got ..."): when tensor dimensions disagree with what an op or layer was told to expect — this error's family across 6 libraries.

Related errors

• {kv_seq_len=} must be a multiple of {block_kv=}
• {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/fb166782622a96e6. Report an issue: GitHub.