xla
[ragged-paged-attn] Use hidden states in kv cache and support any num_kv_head
#8851

Merged

[ragged-paged-attn] Use hidden states in kv cache and support any num_kv_head #8851

vanbasten23 merged 4 commits into pytorch:master from bythew3i:ragged-attn-v2

Refactor kv cache to use hidden states

07d1cfe7

Fix undefined variable

3a891781

Fix kv cache shape in tests

ddbe8056

vanbasten23 commented on 2025-03-18

Rename kv_model_dim to kv_hidden_size

add6be02

vanbasten23 approved these changes on 2025-03-19

vanbasten23 merged 4190fc0e into master 230 days ago

Reviewers

vanbasten23

Assignees

No one assigned

Labels

None yet

Milestone

No milestone