FlexAttention?

[johnnynunez]

Is it compatible flexattention from pytorch 2.6.0?

[jainapurva]

@drisspg

[drisspg]

Can you add some more context here @johnnynunez

[johnnynunez]

Can you add some more context here @johnnynunez

I want to quantize the lerobot pizero model, which has FlexAttention. @drisspg
https://huggingface.co/blog/pi0

context:
In the future, we plan on extending this support to allow for quantized versions of attention or things like RadixAttention as well.

https://pytorch.org/blog/flexattention/

[drisspg]

So currently all of our quantization APIs target linear layers and are orthogonal to flex attention. Therefore, yes, flex attention should work. Flex-Tension currently doesn't support low precision inputs, however, that is planned - no ETA yet though

[johnnynunez]

So currently all of our quantization APIs target linear layers and are orthogonal to flex attention. Therefore, yes, flex attention should work. Flex-Tension currently doesn't support low precision inputs, however, that is planned - no ETA yet though

thanks! I'm going to try

[drisspg]

Let me know if anything comes up!

[moinnadeem]

Flex-Tension currently doesn't support low precision inputs, however, that is planned - no ETA yet though

@drisspg this is a good point -- what will happen with a low-precision input? will it get upcast to bf16 for the actual matmul? if so, are you basically seeing VRAM savings but no time savings?

[drisspg]

I have an example for doing this and @danielvegamyhre is starting to investigate and ultimately make this a well supported path.

For an fp8 mm the mm will be in low precision and in theory can utilize tensors cores w/ fp8 support on H100 and accumulated into high