Restore deleted capabilities, completed and integrated properly

[AshishKumar4]

Fourth PR of the series — stacked on #12. Restores the two real capabilities Phase 1/2 deleted, rebuilt properly on the new foundation instead of reviving the broken originals.

Fused / flash attention (was EfficientAttention + --flash_attention)

One helper — scaled_dot_product_attention — now backs every attention module in the library:

• default (None): flax reference attention, runs everywhere, honors force_fp32_for_softmax
• 'xla' / 'cudnn': jax.nn.dot_product_attention, which dispatches to the fused cudnn flash kernel on supported GPUs
• 'tpu': the pallas TPU flash kernel with the 1/sqrt(d) scale passed explicitly — the deleted implementation passed none (pallas defaults sm_scale=1.0), inflating logits 8× at head_dim 64, and used a different param tree so its checkpoints were unusable on the normal path

The projections never change with the implementation, so checkpoints are interchangeable across hardware — train on TPU, sample on GPU, switch kernels freely. attention_impl is exposed on every model (DiT family, MM-DiT joint attention, UViT/UDiT, UNet) and as --attention_impl. A parity test pins reference-vs-fused agreement on identical params (verified to 6e-7 at the kernel level, 1e-4 at model level).

Video modeling (was unet_3d.py)

The old UNet3D was never wired into any registry/training path and subclassed diffusers flax blocks that upstream deleted. The replacement is VideoDiT — factorized spatial-temporal attention built entirely from the shared backbone: per-frame PatchSequenceEmbed (hilbert/zigzag work), spatial + temporal ModulatedBlock per layer, RoPE on the genuine 1D time axis, shared output head. Registered as video_dit and integrated end to end — including a fix for DiffusionInputConfig silently dropping the time axis for video shapes, which would have prevented any video model from initializing through the trainer.

Tests: forward shapes, temporal-mixing behavior check, registry construction, and a 5D end-to-end training run through the real GeneralDiffusionTrainer (the samplers were already made 5D-safe in Phase 1, so the whole video path now works: data → trainer → model → sampler).

UNet3D — restored properly (third commit)

The convolutional video UNet is back, rebuilt from our own UNet blocks instead of the deleted diffusers ones — which is what makes it actually valuable: the spatial module names mirror Unet exactly, so inflate_unet_params drops a trained 2D checkpoint straight into the 3D tree. The temporal attention blocks (RoPE over the frame axis, at every resolution level) are zero-initialized, so an inflated model reproduces the image model frame-by-frame and training only has to learn motion — the AnimateDiff recipe applied to our own pretrained models (e.g. the TRC-trained 128px text-conditional UNet). Verified by test: an inflated UNet3D matches the 2D Unet per-frame to 1e-5.

Registered as unet_3d; both video archs are wired into training.py. Suite: 67 passed.

Not restored, deliberately

favor_fastattn.py (Performer): its role — linear-cost token mixing — is covered by the S5 SSM mixer in ModulatedBlock, which is this repo's own research direction. Reviving 722 lines of gin-dependent vendored code would add a second linear-mixer path with no consumer. If linear attention specifically is ever wanted, the right shape is a third mixer= option in ModulatedBlock, not the old file.

Verification

pytest -m "not network": 65 passed (4 new tests). Suite remains xfail-free.

🤖 Generated with Claude Code

[AshishKumar4]

Landed on main as part of the full revamp stack (main fast-forwarded to bc9d06b). GitHub reports no new commits between this head and main, so every commit here is merged; closing as superseded by the stack merge rather than leaving it open.