[Kernel] RPA: do the q token/kv-head swap on the 4D view to avoid a VMEM round trip - #3393

[Kernel] RPA: do the q token/kv-head swap on the 4D view to avoid a VMEM round trip#3393

shungcp wants to merge 1 commit into

vllm-project:mainvllm-project/tpu-inference:mainfrom

shungcp:rpa-q-layoutshungcp/tpu-inference:rpa-q-layoutCopy head branch name to clipboard