MasterJH5574 opened a new pull request, #16456: URL: https://github.com/apache/tvm/pull/16456
This PR introduces two changes to the (paged) KV cache: The first is introducing RoPE mode to PagedKVCache. Right now there are two modes: normal/inline. In "normal" mode, RoPE will be applied to input Q/K/V data before appending the K/V data to cache. In "inline" mode, the input K/V data is directly appending to cache, and the RoPE will be on-the-fly applied inside attention kernel. The main purpose of introducing RoPE mode is to balance the need of on-the-fly RoPE (in cases like Mistral where positions can cahnge) and the attention kernel performance. The second is introducing a new interface `AttentionWithFusedQKV` to KV cache. This function takes the input QKV data that is fused along the head dimension. And the fused QKV will be split into separate Q/K/V internally (note: requiring external workspace passed in). We introduce this function since in practice we note that when RoPE mode is "normal," it offers better performance if we fuse the QKV split and RoPE application. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
