Optimize TPU Flash Attention (20x XLA compilation speed-up on 32k long context)#908

Merged

ds-hwang merged 1 commit intoapple:mainfrom ds-hwang:flsh_op

Jan 7, 2025

+257-122

Commits on Jan 7, 2025

Optimize TPU Flash Attention (20x XLA compilation speed-up on 32k long context)
ds-hwang
committed