Replace bare F.scaled_dot_product_attention call inside MagiAttention
with FastVideo's LocalAttention layer so the backend selection (SDPA /
FlashAttn / SLA / SageAttn) flows through the standard configurable
path. Also drops the manual GQA repeat_interleave: SDPABackend's
enable_gqa=True handles num_heads_q != num_heads_kv directly.
LocalAttention requires a forward_context, so add set_forward_context
wrapping at the two call sites:
* Production: MagiHumanDenoisingStage's per-step DiT calls now run
inside set_forward_context(current_timestep=t, attn_metadata=None),
matching the pattern used by the generic DenoisingStage and other
custom denoise stages (ltx2, stable_audio, longcat).
* Parity tests: test_magi_human_parity.py and
test_magi_human_pipeline_parity.py wrap their direct DiT calls in
the same context so LocalAttention's get_forward_context() doesn't
assert.
Parity tests stay bit-exact (all 7 magi-human parity tests pass with
0.0/0.0 diff vs upstream). Production E2E still produces a coherent
video matching the prompt; the mp4 output is no longer byte-identical
to the pre-refactor version because SDPA dispatches to a different
kernel when enable_gqa=True at production sequence length (~3000
tokens), with mean per-pixel drift ~3/255 (~1.3%) -- visually
indistinguishable, just a different bf16 quantization noise pattern.
The MagiAttention block stays as a custom nn.Module rather than being
fully replaced by LocalAttention because of the per-modality packed
linears (PackedExpertLinear), per-head sigmoid gating, and per-modality
RMSNorm pattern that the FV layer abstractions don't model. The bare
SDPA call inside it is now the only piece reused from FV.
tests/local_tests/ holds local-only parity, smoke, and component tests
organized by model family (e.g. tests/local_tests/ltx2/,
tests/local_tests/stable_audio/). They typically require a checked-out
reference repo and/or model weights and are skipped in CI; see each family's
README.md for setup. The top-level
tests/local_tests/README.md is the family index.
The CI-backed test suite still lives in fastvideo/tests/.
SolitaryThinker
e24bc12c59