xai-org/x-algorithm · error · ValueError

Need to specify backward blocks.

Error message

Need to specify backward blocks.

What it means

The backward kernel requires explicit backward tile sizes; if TuningConfig was created without all four of block_q_dkv, block_kv_dkv, block_q_dq, block_kv_dq (has_backward_blocks is False), the VJP raises because no defaults are assumed.

Source

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

def _attention_bwd(
    config: TuningConfig,
    save_residuals: bool,
    bound,
    sm_scale: float,
    cap: float,
    cap_method: str,
    z_loss_weight: float,
    res,
    do,
):
    del save_residuals
    q, k, v, out, lse = res

    if config.causal:
        raise NotImplementedError("Causal attention not supported in the backwards pass yet.")

    if not config.has_backward_blocks:
        raise ValueError("Need to specify backward blocks.")

    assert config.block_q_dq is not None
    assert config.block_kv_dq is not None
    assert config.block_q_dkv is not None
    assert config.block_kv_dkv is not None

    batch_size, q_seq_len, num_q_heads, head_dim = q.shape
    _, kv_seq_len, num_kv_heads, _ = k.shape
    q_heads_per_kv_head = num_q_heads // num_kv_heads
    dtype = q.dtype
    compute_wgs = config.compute_wgs_bwd

    num_q_tiles, rem = divmod(q_seq_len, config.block_q_dq * compute_wgs)
    if rem:
        raise NotImplementedError(
            f"{q_seq_len=} must be a multiple of {config.block_q_dq=} * {compute_wgs=}"
        )

View on GitHub (pinned to 24c60942c5)

Solutions

1. Add all four backward block sizes to the config (each a multiple of 64), e.g. 128/64/128/64
2. Verify config.has_backward_blocks before differentiating

Example fix

# before
cfg = TuningConfig(block_q=128, block_kv=128, max_concurrent_steps=2)
# after
cfg = TuningConfig(block_q=128, block_kv=128, max_concurrent_steps=2,
                  block_q_dkv=128, block_kv_dkv=64, block_q_dq=128, block_kv_dq=64)

assert config.has_backward_blocks, "set block_q_dkv/block_kv_dkv/block_q_dq/block_kv_dq before training"

When it happens

Trigger: Building a TuningConfig with only forward fields and calling jax.grad / vjp through attention or sharded_mha.

Common situations: Reusing an inference-tuned config for training; defaults changed so backward blocks must now be supplied manually.

Related errors

• Backward block sizes (block_q_dkv, block_kv_dkv, block_q_dq,
• Causal attention not supported in the backwards pass yet.
• {q_seq_len=} must be a multiple of {config.block_q_dq=} * {c
• Invalid backward pass implementation: {backward_pass_impl}
• cap_method must be in [tanh, soft_sign]

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