[ElmoPA]

Large MoT batches spend most computation on attention, especially repeated recomputation from nested checkpointing and 512-query chunks. Adds full-query fused efficient attention and an optional FlashAttention 4 backend with BF16 Q/K/V and FP32 outputs, parameters, remaining model and losses. All valid keys remain available to every query; image tokens, architecture, data sampling and objectives are unchanged.

New continuous and discrete MoT recipes default to Flash4 BF16 attention, unchunked queries and no block activation checkpointing, with TF32 disabled. Explicit native and unchunked-FP32 profiles remain available. Flash4 requires batch one, Hopper-or-newer CUDA hardware and the separately verified Torch 2.9.1 / FlashAttention 4.0.0b35 beta runtime. The ordinary shared project dependencies are unchanged; unsupported environments fail explicitly.

Validation:

• 132 model/training-policy tests, including real CUDA padding and gradient checks, at integration commit 01961df; Ruff passed.
• Eight native optimizer updates per each of three simulation recipes, scheduled validation and strict loading of both immutable step-4/8 checkpoints; maximum 16-time continuous and four-time discrete objective/prediction checks cover both embodiments. All nine checkpoint files remain retained and SHA-verified.
• Nine configuration/profile checks at cb95574 verify default selection without an attention override; Ruff passed.

On one paired, warmed 16-timestamp H100 example, forward/backward decreased from 106.01 s to 1.17 s (90.7×). This excludes loading, optimizer, validation and saving; peak allocated memory increased from 7.8 to 30.4 GiB. BF16 gradients differed by approximately 0.045% in global relative L2 in that example. Short checks do not establish unchanged convergence or final quality.

[ElmoPA]

Warning

This pull request is not mergeable via GitHub because a downstack PR is open. Once all requirements are satisfied, merge this PR as a stack on Graphite.
Learn more

• Add paired modality ablations for joint image forecast quality #265 : 2 dependent PRs (#276 , #320 )
• Train joint image, proprio and action MoT with smaller vector experts #264
• Add modular pixel video training and resilient log mirroring #262
• Optimize MoT attention and enable Flash4 BF16 defaults #261 👈 (View in Graphite)
• Add continuous-time MoT world model and offline data pipeline #260
• main

This stack of pull requests is managed by Graphite. Learn more about stacking.