Repository navigation
Conversation
| def _compute_attention(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, cu_seqlens, max_seqlen): | ||
| """Run the flash attention kernel. q/k/v are [total_tokens, heads, dim].""" | ||
| kwargs: dict = {"causal": True} | ||
| cu_seqlens_k = cu_seqlens.clone() |
![cursor[bot]](https://avatars.githubusercontent.com/in/1210556?s=60&v=4){kind=small}
| cu_seqlens=cu_seqlens, | ||
| chunk_indices=chunk_indices, | ||
| chunk_size=64, | ||
| ) |
There was a problem hiding this comment.
Delta-rule backward unpack arity mismatch
High Severity
chunk_gated_delta_rule_bwd is unpacked into eight values, but FLA's kernel returns six (dq, dk, dv, db, dg, dh0). Backward of the new fullgraph chunk_gated_delta_rule path raises ValueError and training cannot step.
Reviewed by Cursor Bugbot for commit 9e87b40. Configure here.
There was a problem hiding this comment.
There was a problem hiding this comment.
FlashAttention 3 passes both sequence-length tensors into FlashAttnVarlenFunc.apply. With fullgraph=True, passing the same tensor object for Q and K causes Dynamo to fail with
autograd.Function.apply: duplicate tensor input. The clone preserves identical values while giving K a distinct tensor identity. Without it, the one-node run failed at the first full-
attention layer; with it, all 32 layers captured and forward/backward completed.
There was a problem hiding this comment.
ok, figured it was along those lines, though I don't know why it would fail. not very important
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes and found 1 potential issue.
There are 2 total unresolved issues (including 1 from previous review).
❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.
Reviewed by Cursor Bugbot for commit 4214f8d. Configure here.
| key = key.contiguous() | ||
| value = value.contiguous() | ||
| decay = decay.contiguous() | ||
| beta = beta.contiguous() |
There was a problem hiding this comment.
Backward skips contiguous value tensors
Medium Severity
Forward now contigifies value and beta for the FLA kernels, but setup_context still saves the original custom-op inputs. Backward can therefore still receive non-contiguous value/beta under fullgraph, where torch.split views often stay non-contiguous, so chunk_gated_delta_rule_bwd can crash or return wrong gradients after forward succeeds.
Additional Locations (1)
Reviewed by Cursor Bugbot for commit 4214f8d. Configure here.
| query = query.contiguous() | ||
| key = key.contiguous() | ||
| value = value.contiguous() | ||
| decay = decay.contiguous() | ||
| beta = beta.contiguous() |
There was a problem hiding this comment.
Oh, yep, classic issue. Semi-surprised the kernel didn't error on non-contiguity; know some of the FLA ops guard on this.
On Blackwell (SM100) `attn = "auto"` resolves to FA4, whose public varlen wrapper calls `torch._guards.active_fake_mode()` via `is_fake_mode()`. Dynamo refuses to trace that function, so `fullgraph=True` fails in the full-attention layers (and without fullgraph the model silently splits around every attention call, which prevents CUDA graph replay). Wrap FA4's forward/backward kernels in `torch.library` custom ops with fake impls and a registered autograd, following the pattern already used for the GDN ops. Only `FlashAttention._compute_attention` is switched over; other FA4 call sites (afmoe, gpt_oss, Qwen3.5 vision) are unchanged. The wrapper relies on FA4's private `_flash_attn_fwd` / `_flash_attn_bwd`, so an FA4 bump may need a signature check. Tested on 8x B200 with configs/debug/fake/qwen3_5_9b_fullgraph_sft.toml: eager ~1.00 s/step vs reduce-overhead fullgraph ~0.635 s/step (1.57x). Outputs and gradients are bit-identical to stock FA4 (new unit test). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Note
Medium Risk
Touches shared flash-attention varlen metadata and introduces new custom autograd ops on the Qwen3.5 training path; behavior is guarded by tests but errors would affect numerics or compile stability.
Overview
Enables
torch.compile(fullgraph=True)on Qwen3.5 layers by replacing direct FLA calls in the linear-attention stack withtorch.librarycustom ops (causal_conv1d,chunk_gated_delta_rule) that wire FLA forward/backward and fake kernels for Dynamo.Qwen3_5GatedDeltaNetuses those wrappers on the non–context-parallel path and swapsFusedRMSNormGatedfor a localQwen3_5RMSNormGated; CP still uses eager FLA conv and the original chunk op.Flash attention now passes a cloned
cu_seqlens_kinto varlen flash (FA3/FA4) so compile does not reuse the same tensor for q and k metadata.Adds a debug fake-data SFT config for Qwen3.5-9B with
fullgraph = true, README note, and GPU tests for layer fullgraph forward/backward and gated RMS norm parity with FLA.Reviewed by Cursor Bugbot for commit 4214f8d. Bugbot is set up for automated code reviews on this repo. Configure here.