{"record":{"id":"f617b54ed8a80d8e","repo":"xai-org/x-algorithm","slug":"q-seq-len-must-be-a-multiple-of-config-block-q","errorCode":null,"errorMessage":"{q_seq_len=} must be a multiple of {config.block_q_dq=} * {compute_wgs=}","messagePattern":"(.+?) must be a multiple of (.+?) \\* (.+?)","errorType":"validation","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"phoenix/xrex/pallas/ranker_attention_fa3.py","lineNumber":518,"sourceCode":"        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\n    num_kv_tiles, rem = divmod(kv_seq_len, config.block_kv_dkv * compute_wgs)\n    if rem:\n        raise NotImplementedError(\n            f\"{kv_seq_len=} must be a multiple of {config.block_kv_dkv=} * {compute_wgs=}\"\n        )\n\n    num_q_tiles_in_dkv, rem = divmod(q_seq_len, config.block_q_dkv)\n    if rem:\n        raise NotImplementedError(f\"{q_seq_len=} must be a multiple of {config.block_q_dkv=}\")\n\n    num_kv_tiles_in_dq, rem = divmod(kv_seq_len, config.block_kv_dq)\n    if rem:\n        raise NotImplementedError(f\"{kv_seq_len=} must be a multiple of {config.block_kv_dq=}\")\n\n    bound_arr = _normalize_bound(bound, batch_size, q_seq_len)","sourceCodeStart":500,"sourceCodeEnd":536,"githubUrl":"https://github.com/xai-org/x-algorithm/blob/24c60942c5c5fdad3a6addffb4c6e6d2f228f04f/phoenix/xrex/pallas/ranker_attention_fa3.py#L500-L536","documentation":"The backward dkv/dq kernels split work across compute_wgs warp groups, so q_seq_len must be divisible by config.block_q_dq * compute_wgs (= compute_wgs_bwd). Remainder sequences cannot be tiled and raise NotImplementedError.","triggerScenarios":"Differentiating with q_seq_len not a multiple of block_q_dq * compute_wgs_bwd, e.g. seq 1000 with block_q_dq=128 and compute_wgs=2 (divisor 256).","commonSituations":"Short fine-tuning sequences; adjusting compute_wgs_bwd or backward block sizes without re-checking sequence padding.","solutions":["Pad q (and kv consistently) to a multiple of block_q_dq * compute_wgs_bwd and mask with bound","Or pick block sizes whose product with compute_wgs divides q_seq_len (e.g. 64 divisors)"],"exampleFix":"# before\ngrads = jax.grad(loss)(params)  # q_seq_len=1000, block_q_dq=128, compute_wgs=2\n# after\n# pad q_seq_len to 1024 (multiple of 256) and mask with bound, or:\ncfg = replace(cfg, block_q_dq=64)  # 64*2=128 divides 1000? no -> pad to 1024\nq = jnp.pad(q, ((0,0),(0,24),(0,0),(0,0)))","handlingStrategy":"validation","validationCode":"divisor = config.block_q_dq * config.compute_wgs_bwd\nassert q_seq_len % divisor == 0, f\"pad q_seq_len to a multiple of {divisor}\"","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Compute backward divisors (block*_dq/dkv * compute_wgs_bwd) when choosing padding","Reuse one padded seq-len for forward and backward alignment checks"],"tags":["jax","pallas","attention","backward-pass","shape-validation"],"backgroundTag":"sequence-length-not-aligned","analyzedSha":"24c60942c5c5fdad3a6addffb4c6e6d2f228f04f","analyzedAt":"2026-08-28T11:40:14.686Z","schemaVersion":2},"datasetVersion":"2026-08-28T16:17:29.566Z"}