Fix SDPA vmap with GQA/MQA shapes (n_heads != n_kv_heads)#3385
Open
Brooooooklyn wants to merge 1 commit into
Open
Fix SDPA vmap with GQA/MQA shapes (n_heads != n_kv_heads)#3385Brooooooklyn wants to merge 1 commit into
Brooooooklyn wants to merge 1 commit into
