[Feature]: Move query quantization to attention layer for all backends supporting query quantization.

[jmkuebler]

🚀 The feature, motivation and pitch

#24914 moved the query quantization out of the FlashAttn backend into attention/layer such that it can be fused by torch compile AND included in the CUDA graph.

#26534 moved it for Triton and FlashInfer.

This should be done for all backends.

Alternatives

No response

Additional context

No response

Before submitting a new issue...

• Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the documentation page, which can answer lots of frequently asked questions.

[adabeyta]

Sounds like a great idea, I'll take a look into moving query quantization to attention layer.

[jmkuebler]

Great thanks @adabeyta ! Note that for FA we did it in #24914

I just didn't get a chance to look into other backends yet (even to check which ones support quantized queries even).

So you help would be great to look into the other backends.

[aditya-kolsur]

@ProExpertProg I would like to contribute to this, could it be assigned to me? For which backends is it yet to be moved?

[AniruddhS24]

@ProExpertProg I would also like to contribute to this. @aditya-kolsur perhaps we can work on it together?

[ProExpertProg]

@aditya-kolsur @AniruddhS24 sorry for the late response. This has been moved for FA, Triton, and FlashInfer. Can you take a look at other backends to see if any of them quantize query and add support for them? You can divide the work as I assume there are still multiple backends to add support for.

[vnadathur]

@ProExpertProg are you looking to add quantized queries support to other attn backends? mayb that should be its own issue. I don't believe anything other than FA, FI, and Tritonattn currently support quantized queries. I think treeattn is the only one I can see getting qq support. what do u think?