usp: differentiable operators for training (Ulysses/Ring/local attention)#1
Open
zhihengy wants to merge 3 commits into
Open
usp: differentiable operators for training (Ulysses/Ring/local attention)#1zhihengy wants to merge 3 commits into
zhihengy wants to merge 3 commits into
Conversation
ft_c.all_to_all_single has no registered autograd kernel, so backward silently stops there — fine for inference (no_grad) but breaks training (q/k/v projection grads come out None). Wrap it in an autograd Function; an even-split all-to-all is an involution, so its adjoint is the same collective on the upstream gradient. Restores the differentiability the FastVideo source had. Transparent to rollout. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
ring_attn called _templated_ring_attention (forward template) directly; the KV ring rotation isn't autograd-aware, so dK/dV never get the reverse-ring gradient pass that torch keeps in _templated_ring_attention_backward (normally wired by torch's context_parallel autograd.Function). dQ survived (Q isn't rotated) but to_k/to_v weight grads were ~0.59 rel error. Add _RingFlashAttention autograd Function (fwd=forward template, bwd=backward template, op=aten flash) and route ring_attn through it when grad is enabled; inference path unchanged. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
The jit FA kernels (sgl_kernel::fwd) have no autograd registration, so backward through USPAttention's local attention silently produces no q/k/v grads (fa4 on Blackwell has backward; fa3 on Hopper does not). Route the grad-enabled path through torch SDPA; inference is unchanged. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What these fix
All three follow the same pattern: a training-only branch guarded by
torch.is_grad_enabled()/requires_grad; the inference path is byte-identical.642026e31eUlysses all-to-all autograd —ft_c.all_to_all_singlehas no autograd kernel, so backward silently stops there (q/k/v projection grads come out None, with no error). An even-split all-to-all is an involution, so the adjoint is the same collective. Restores the differentiability the FastVideo source had.a928d729c5differentiable ring attention —ring_attncalled_templated_ring_attention(forward template) directly; the KV ring rotation is not autograd-aware, so dK/dV missed the reverse-ring pass (~0.59 rel error). Adds_RingFlashAttention(fwd/bwd templates, aten flash op); inference path unchanged.a606a3d0ffdifferentiable local attention — the jit FA kernels are forward-only on Hopper (fa3 viasgl_kernel::fwdhas no autograd registration; fa4 on Blackwell has backward). Grad-enabled path routes through torch SDPA.Invasiveness
2 files, ~120 insertions / 4 deletions. No API changes, no new deps. Validated by miles' SP parity suite before the dependency was removed (forward bitwise vs full-sequence reference, 6 bands incl. ring × grad-ckpt; details in radixark/miles_diffusion#21).
If upstreaming to sgl-project: main has since rewritten the USP all-to-all (sgl-project#24366, sgl-project#27143), so patch 1's wrap point needs re-adaptation there.
🤖 Generated with Claude Code