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!