Adding support for BlockedKV attention in CasualLM models

[vaibverm]

Objective:

This PR introduces the KV blocking technique for CausalLM models where the K/V cache is read and processed block by block in the attention computation. Number of desired KV blocks are defined at model initialization in the "from_pretrained" call to export the ONNX with required number of KV blocks. As a result, the following changes are introduced:

Changes:

1. SoftMax needs to be changed from regular SoftMax to online SoftMax where the running maximum and cumulative denominators are tracked and updated once each block is processed to retain mathematical accuracy compared to regular SoftMax.
2. Changes to CTXGather and CTXGatherCB custom ops to read only 1 block worth of data in each cache gather/read.
3. Changes to read_only function in QEffDynamicCache to allow reading of a cache block by block rather than full K/V cache.
4. Generation of attention mask per block.
5. Changes to eager_attention_forward implementation in the llama model to allow BlockedKV attention and online SoftMax implementation.
6. Wrapping the num_kv_blocks variable inside qaic_config to keep consistent calling style.
7. A new PyTorch transform to pass the num_kv_blocks variable to QEffLlamaAttention block.
8. A new constant added for num_kv_blocks.
9. Added tests to switch the BlockedKV feature on and off.

Please review and feel free to suggest changes and tests.

[vbaddi]

Thanks @vaibverm
Could you please address the conflicts and run the lint/format?

[vaibverm]

Hi @vbaddi,
I have addressed the conflicts but some workflows need approval. Would you be able to approve those?

[quic-rishinr]

if we are comparing int do we need torch.where? can we use min()?

[vaibverm]

torch.min() requires both inputs to be tensors. I tried torch.min() before but using torch.min() here leads to ONNX export time error.

[quic-rishinr]

if we are not using the attention_mask do we need this condition?

[vaibverm]

This is a causal model, so we use masking. The reason I kept attention_mask instead of the causal_mask_block in the condition was because earlier I was using a common method for both eager and blockedKV attention. I would suggest using the original condition testing if the overhead is not too high for compatibility reasons with regular eager attention.

[quic-rishinr]

Can we add a warning if the arcitecture doesnt support blocked KV

[vaibverm]

I wanted to have a broader discussion on this. I implemented blockedKV attention for Qwen2.5_VL model as well and the norm there was to use environment variable to switch between different blocking techniques. Is that the norm we want to keep across QEff? If yes, then we will not really need this transform anymore although I think using the PyTorch transform was a cleaner way to switch between different blocking techniques and consistent with current transform usage in QEff.

[vaibverm]

@quic-rishinr - Would you suggest we should go the route of using environment variables? Or would you prefer PyTorch transforms like above to implement the blocking?

[quic-rishinr]

please add @pytest.mark.on_qaic on the test as both tests would be using the qaic cards