xai-org/x-algorithm · critical · NotImplementedError

Causal attention not supported in the backwards pass yet.

Error message

Causal attention not supported in the backwards pass yet.

What it means

The backward (VJP) kernel of the FA3-style attention only implements the non-causal path; if TuningConfig.causal is True the backward pass raises NotImplementedError instead of producing wrong gradients.

Source

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

    return out, (q, k, v, out, lse)


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(

View on GitHub (pinned to 24c60942c5)

Solutions

Set causal=False in the config for training (use an explicit mask/bound if masking is needed)
Or use ranker_attention.py's mha backward or mha_reference for gradient computation
Keep causal attention for inference only with this kernel

Example fix

# before
cfg = TuningConfig(..., causal=True)
loss = loss_fn(attention(q, k, v, config=cfg))
grads = jax.grad(loss_fn)(params)
# after
cfg_train = replace(cfg, causal=False)
loss = loss_fn(attention(q, k, v, config=cfg_train))
grads = jax.grad(loss_fn)(params)

Defensive patterns

Strategy: fallback

Validate before calling

if config.causal:
    assert not requires_grad_path, "FA3 backward does not support causal; use non-causal config or reference impl"

Try / catch

try:
    grads = jax.grad(loss)(params)
except NotImplementedError:
    grads = jax.grad(reference_loss)(params)  # mha_reference fallback

Prevention

Gate training entry points on a causal-and-backward compatibility check
Use causal FA3 attention only in inference code paths

When it happens

Trigger: Creating attention with config.causal=True and then differentiating (jax.grad, loss.backward via vjp) through it, e.g. training a causal ranker model.

Common situations: Fine-tuning/training with causal masking; switching from inference-only usage (forward works fine with causal) to training.

Related errors

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