Flash attention for GPUs like in maxtext#149

Merged

entrpn merged 35 commits intoAI-Hypercomputer:mainfrom

ksikiric:kris/flux-FA

Feb 21, 2025

Commits on Feb 18, 2025

add support for flux vae. ~ wip

jfacevedo-google
authored and
ksikiric
committed
test for flux vae both encoding and decoding.

jfacevedo-google
authored and
ksikiric
committed
add clip text encoder test

jfacevedo-google
authored and
ksikiric
committed
remove transformers inside maxdiffusion, add transformers dependency. Start creating generation code for flux.

jfacevedo-google
authored and
ksikiric
committed
add double block to flux

jfacevedo-google
authored and
ksikiric
committed
forward pass for single double block.

jfacevedo-google
authored and
ksikiric
committed
trying to use scan.

jfacevedo-google
authored and
ksikiric
committed
add single stream block

jfacevedo-google
authored and
ksikiric
committed
finish transformer

jfacevedo-google
authored and
ksikiric
committed
convert pt weights to flax and load transformer state.

jfacevedo-google
authored and
ksikiric
committed
apply fsdp sharding, do one forward pass in the transformer.

jfacevedo-google
authored and
ksikiric
committed
wip - generate fn

jfacevedo-google
authored and
ksikiric
committed
working loop, bad generation

jfacevedo-google
authored and
ksikiric
committed
e2e, encoder offloading.

jfacevedo-google
authored and
ksikiric
committed
support both dev and schnell loading. Images still incorrect.

jfacevedo-google
authored and
ksikiric
committed
flux schnell working

jfacevedo-google
authored and
ksikiric
committed
removed unused code.

jfacevedo-google
authored and
ksikiric
committed
support dev

jfacevedo-google
authored and
ksikiric
committed
add sentencepiece requirement

jfacevedo-google
authored and
ksikiric
committed
fix repeated double and single blocks.

jfacevedo-google
authored and
ksikiric
committed
optimized flash block sizes for trillium.

jfacevedo-google
authored and
ksikiric
committed
clean up code and lint

jfacevedo-google
authored and
ksikiric
committed
fix sdxl generate smoke tests.

jfacevedo-google
authored and
ksikiric
committed
fix rest of unit tests.

jfacevedo-google
authored and
ksikiric
committed
update readme and some dependencies.

entrpn
authored and
ksikiric
committed
remove unused dependencies.

entrpn
authored and
ksikiric
committed
initial lora implementation for flux

jfacevedo-google
authored and
ksikiric
committed
adding another format lora support.

jfacevedo-google
authored and
ksikiric
committed
Support other format loras. update readme. Run code_style.

jfacevedo-google
authored and
ksikiric
committed
ruff

jfacevedo-google
authored and
ksikiric
committed
fix typo in readme.

entrpn
authored and
ksikiric
committed
Added FA support for GPUs
ksikiric
committed
ruff and code_style
ksikiric
committed

Commits on Feb 20, 2025