{"record":{"id":"a12295fa81f93101","repo":"xai-org/x-algorithm","slug":"need-to-specify-backward-blocks","errorCode":null,"errorMessage":"Need to specify backward blocks.","messagePattern":"Need to specify backward blocks\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"phoenix/xrex/pallas/ranker_attention_fa3.py","lineNumber":503,"sourceCode":"def _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(\n            f\"{q_seq_len=} must be a multiple of {config.block_q_dq=} * {compute_wgs=}\"\n        )\n","sourceCodeStart":485,"sourceCodeEnd":521,"githubUrl":"https://github.com/xai-org/x-algorithm/blob/24c60942c5c5fdad3a6addffb4c6e6d2f228f04f/phoenix/xrex/pallas/ranker_attention_fa3.py#L485-L521","documentation":"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.","triggerScenarios":"Building a TuningConfig with only forward fields and calling jax.grad / vjp through attention or sharded_mha.","commonSituations":"Reusing an inference-tuned config for training; defaults changed so backward blocks must now be supplied manually.","solutions":["Add all four backward block sizes to the config (each a multiple of 64), e.g. 128/64/128/64","Verify config.has_backward_blocks before differentiating"],"exampleFix":"# before\ncfg = TuningConfig(block_q=128, block_kv=128, max_concurrent_steps=2)\n# after\ncfg = TuningConfig(block_q=128, block_kv=128, max_concurrent_steps=2,\n                  block_q_dkv=128, block_kv_dkv=64, block_q_dq=128, block_kv_dq=64)","handlingStrategy":"validation","validationCode":"assert config.has_backward_blocks, \"set block_q_dkv/block_kv_dkv/block_q_dq/block_kv_dq before training\"","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Maintain separate train and inference TuningConfig factories","Assert has_backward_blocks in the training script startup"],"tags":["jax","pallas","attention","backward-pass","tuning-config"],"backgroundTag":"missing-config-for-backward-pass","analyzedSha":"24c60942c5c5fdad3a6addffb4c6e6d2f228f04f","analyzedAt":"2026-08-28T11:40:14.686Z","schemaVersion":2},"datasetVersion":"2026-08-28T16:17:29.566Z"}