{"record":{"id":"ab1df89221597407","repo":"xai-org/x-algorithm","slug":"causal-attention-not-supported-in-the-backwards-pa","errorCode":null,"errorMessage":"Causal attention not supported in the backwards pass yet.","messagePattern":"Causal attention not supported in the backwards pass yet\\.","errorType":"validation","errorClass":"NotImplementedError","httpStatus":null,"severity":"critical","filePath":"phoenix/xrex/pallas/ranker_attention_fa3.py","lineNumber":500,"sourceCode":"    return out, (q, k, v, out, lse)\n\n\ndef _attention_bwd(\n    config: TuningConfig,\n    save_residuals: bool,\n    bound,\n    sm_scale: float,\n    cap: float,\n    cap_method: str,\n    z_loss_weight: float,\n    res,\n    do,\n):\n    del save_residuals\n    q, k, v, out, lse = res\n\n    if config.causal:\n        raise NotImplementedError(\"Causal attention not supported in the backwards pass yet.\")\n\n    if not config.has_backward_blocks:\n        raise ValueError(\"Need to specify backward blocks.\")\n\n    assert config.block_q_dq is not None\n    assert config.block_kv_dq is not None\n    assert config.block_q_dkv is not None\n    assert config.block_kv_dkv is not None\n\n    batch_size, q_seq_len, num_q_heads, head_dim = q.shape\n    _, kv_seq_len, num_kv_heads, _ = k.shape\n    q_heads_per_kv_head = num_q_heads // num_kv_heads\n    dtype = q.dtype\n    compute_wgs = config.compute_wgs_bwd\n\n    num_q_tiles, rem = divmod(q_seq_len, config.block_q_dq * compute_wgs)\n    if rem:\n        raise NotImplementedError(","sourceCodeStart":482,"sourceCodeEnd":518,"githubUrl":"https://github.com/xai-org/x-algorithm/blob/24c60942c5c5fdad3a6addffb4c6e6d2f228f04f/phoenix/xrex/pallas/ranker_attention_fa3.py#L482-L518","documentation":"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.","triggerScenarios":"Creating attention with config.causal=True and then differentiating (jax.grad, loss.backward via vjp) through it, e.g. training a causal ranker model.","commonSituations":"Fine-tuning/training with causal masking; switching from inference-only usage (forward works fine with causal) to training.","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"],"exampleFix":"# before\ncfg = TuningConfig(..., causal=True)\nloss = loss_fn(attention(q, k, v, config=cfg))\ngrads = jax.grad(loss_fn)(params)\n# after\ncfg_train = replace(cfg, causal=False)\nloss = loss_fn(attention(q, k, v, config=cfg_train))\ngrads = jax.grad(loss_fn)(params)","handlingStrategy":"fallback","validationCode":"if config.causal:\n    assert not requires_grad_path, \"FA3 backward does not support causal; use non-causal config or reference impl\"","typeGuard":null,"tryCatchPattern":"try:\n    grads = jax.grad(loss)(params)\nexcept NotImplementedError:\n    grads = jax.grad(reference_loss)(params)  # mha_reference fallback","preventionTips":["Gate training entry points on a causal-and-backward compatibility check","Use causal FA3 attention only in inference code paths"],"tags":["jax","pallas","attention","backward-pass","causal"],"backgroundTag":"unsupported-backward-pass","analyzedSha":"24c60942c5c5fdad3a6addffb4c6e6d2f228f04f","analyzedAt":"2026-08-28T11:40:14.686Z","schemaVersion":2},"datasetVersion":"2026-08-28T16:17:29.566Z"}