Repository navigation
[Attn] Remove SDPA and use Varlen attention in Flux - #5296
Merged
Merged
Conversation
meta-cla
Bot
added
the
CLA Signed
This was referenced Oct 11, 2026
acisseJZhong
added this pull request to stack #5297
October 11, 2026 23:13
acisseJZhong
changed the title
acisseJZhong
force-pushed
the
refactor_attn
branch
from
October 12, 2026 04:40
d2ddc8d to
b65fb58
Compare
acisseJZhong
removed this pull request from stack #5297
October 12, 2026 05:31
acisseJZhong
added this pull request to stack #5308
October 12, 2026 05:32
acisseJZhong
marked this pull request as ready for review
October 12, 2026 06:41
acisseJZhong
requested review from
IvanKobzarev,
SherlockNoMad,
aditvenk,
fegin,
sanketpurandare,
tianyu-l,
wconstab and
wwwjn
as code owners
October 12, 2026 06:41
tianyu-l
approved these changes
Oct 12, 2026
| ) | ||
|
|
||
|
|
||
| def create_flux_attention_metadata( |
Contributor
There was a problem hiding this comment.
do we need this as public API?
also I feel it doesn't belong to "sharding"
Contributor
Author
There was a problem hiding this comment.
moved to model.py and changed the name to **_get_attention_metadata** to align with main decoder.
| return out_THV.unflatten(0, (batch_size, -1)) | ||
|
|
||
|
|
||
| def varlen_attention( |
Contributor
There was a problem hiding this comment.
should this be private function
| return q_BLHK.flatten(0, 1), k_BLHK.flatten(0, 1), v_BLHV.flatten(0, 1) | ||
|
|
||
|
|
||
| def prepare_flux_attention_inputs( |
| cp_group = spmd_mesh_group(CP) | ||
| if cp_group is not None: | ||
| k_BLHK, v_BLHV = ( | ||
| spmd.redistribute( |
Contributor
There was a problem hiding this comment.
flux cp is not via transform -- seems fine
Contributor
Author
There was a problem hiding this comment.
yes we are explicitly doing cp preprocess here as right now flux only have kv all gather and varlen. cc @fegin
Contributor
There was a problem hiding this comment.
For consistency, we should do transform. But we can do this only if we have more CP supports for Flex.
Contributor
Author
There was a problem hiding this comment.
left a todo comment
fegin
approved these changes
Oct 12, 2026
| cp_group = spmd_mesh_group(CP) | ||
| if cp_group is not None: | ||
| k_BLHK, v_BLHV = ( | ||
| spmd.redistribute( |
Contributor
There was a problem hiding this comment.
For consistency, we should do transform. But we can do this only if we have more CP supports for Flex.
acisseJZhong
force-pushed
the
refactor_attn
branch
from
October 12, 2026 17:29
b65fb58 to
574043b
Compare
Delete the experimental SDPA backend and its flat-input adapter. Use the shared Varlen backend in graph test recipes and pass its metadata into bitwise tests. Update the CP codegen check for Flash Attention operators. Use a math-only attention reference for CPU infrastructure tests; keep checkpoint-comparison SDPA isolated under scripts. Validation: - 7 targeted tests and 4 subtests passed (1 recipe case deselected). - Real Llama3 debug model: 10 synthetic AdamW steps, seed 42 and deterministic algorithms; eager and minimal_fx_tracer replay have bitwise equal losses, grad_norms, and all parameter gradients. - Serialized Varlen graphs replay changed document boundaries exactly. - Checkpoint reference matches the removed backend in FP32/BF16 forward and backward. - Non-Pyrefly all-file pre-commit checks pass; new helpers/tests typecheck. Full recipe/FSDP validation is blocked by missing local PyTorch APIs (analyze_pipeline_activation_liveness and param_dtype_override_fn). Global Pyrefly fails on existing environment import/path errors. Already-disabled AutoParallel tests still need Varlen metadata input sharding before they can be enabled.
acisseJZhong
force-pushed
the
refactor_attn
branch
from
October 12, 2026 17:42
f368422 to
6fdad0f
Compare
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
3 participants
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.
This PR removes the core SDPA backend and moves Flux and GraphTrainer recipes to shared Varlen attention, reducing the number of attention implementations to maintain.
Flux keeps its existing
[B, L, D]activations and packs Q/K/V only at attention calls. Per-sample document boundaries are built once per forward, preserving bidirectional attention within each sample and isolation between samples. Under CP, K/V are gathered before packing.Tests
Full Flux-dev and Flux-schnell (11.89B parameters), 8 H100s, BF16, FSDP=8, CP=1, full activation checkpointing, 256x256 images. Each backend ran 20 training steps using randomly initialized weights, the CC12M test shard, and offline tokenizers.
Original SDPA (
3210be8297) vs Varlen (d2ddc8d2d5), both cuDNN. PyTorch 2.14.0+cu130 / cuDNN 9.24; seed 42, TF32 off, deterministic=False, CUDA graphs off.Numerics
Maximum differences across 20 steps:
Both loss traces pass
rtol=1e-3, atol=1e-5; Schnell's grad_norm exceeds this tolerance. These runs are not bitwise identical.Separate 10-step
flux-debugregression, seed 42 and strict determinism, CP=1/2: outputs, loss, grad_norm, gradients, and final weights match the previous Varlen implementation (0b34e4ea36) bitwise.Performance
Median of steps 11-20, including data loading, encoders, forward/backward, optimizer, and communication; CUDA-synchronized wall time, slowest rank.
Schnell's 0.07% difference is within the observed step-to-step variation (5-6 ms standard deviation).