Fix SDPA vmap with GQA/MQA shapes (n_heads != n_kv_heads)#3385
Open
Brooooooklyn wants to merge 1 commit intoml-explore:mainfrom
Open
Fix SDPA vmap with GQA/MQA shapes (n_heads != n_kv_heads)#3385Brooooooklyn wants to merge 1 commit intoml-explore:mainfrom
Brooooooklyn wants to merge 1 commit intoml-explore:mainfrom
