maxtext icon indicating copy to clipboard operation
maxtext copied to clipboard

Published 20 hours ago • verified publisher icon google

FlashAttention Support - TPUv3

Open maciek-pioro opened this issue 7 months ago • 0 comments

Is FlashAttention supported on TPUv3? The same config that works on TPUv4 fails on TPUv3 with the following error: jaxlib.xla_extension.XlaRuntimeError: INTERNAL: Mosaic failed to compile TPU kernel: Unsupported input data type in matrix multiplication. However, after setting attention from autoselected to dot_product, the error disappears.

maciek-pioro avatar Jul 19 '24 10:07 maciek-pioro