metal-flash-attention: causal training ungated; fused RoPE + Attention for Flux1/Flux2, Z-Image and Krea

[bghira]

This pull request introduces support for Metal Flash Attention with fused RoPE (Rotary Position Embedding) in the Flux and Flux2 model attention implementations. The main change is the integration of the maybe_metal_flash_rope_attention function, which enables the use of a fast, fused attention backend when possible, falling back to the standard rotary embedding logic otherwise. Documentation in multiple languages is updated to clarify the behavior and statistics of the fused RoPE path, especially distinguishing between no-grad and gradient calls.

The most important changes are:

Model Code Integration:

• Integrated maybe_metal_flash_rope_attention into the attention call paths of FluxAttnProcessor3_0, FluxSingleAttnProcessor3_0, and both flux/transformer.py and flux2/transformer.py, allowing eligible calls to use Metal Flash Attention with fused RoPE for improved performance. If the fused backend is unavailable or ineligible, the code falls back to the standard RoPE application. (simpletuner/helpers/models/flux/attention.py [1] [2] [3] [4]; simpletuner/helpers/models/flux/transformer.py [5] [6]; simpletuner/helpers/models/flux2/transformer.py [7] [8] [9]

Documentation Updates:

• Updated documentation in English, Spanish, Hindi, Japanese, Portuguese, and Chinese to clarify that no-grad calls through the fused RoPE entry point are counted under rope_instream, while gradient calls are counted under rope_autograd. Also clarified that eligible no-grad calls avoid eager rotation passes and FP32 materializations, improving clarity for users and developers. (documentation/experimental/METAL_FLASH_ATTENTION.md [1] METAL_FLASH_ATTENTION.es.md [2] METAL_FLASH_ATTENTION.hi.md [3] METAL_FLASH_ATTENTION.ja.md [4] METAL_FLASH_ATTENTION.pt-BR.md [5] METAL_FLASH_ATTENTION.zh.md [6]

These changes collectively enable more efficient attention computation on supported hardware, improve code maintainability by centralizing the fused attention logic, and enhance documentation for both users and contributors.