[PR] [KVCache] PagedKVCache refactor, FlashInfer JIT and MLA integration [tvm]

["MasterJH5574 (via GitHub)" <gi...@apache.org>]

MasterJH5574 opened a new pull request, #17674:
URL: https://github.com/apache/tvm/pull/17674

   This PR consists of the following parts:
   
   * We reorganized `paged_kv_cache.cc` by moving some of the utilities to `attn_utils.h`.
   
   * To integrate with the JIT kernel compilation in the latest FlashInfer project, while still being able to support attention kernels written with TIR, we introduced `AttnBackendFunc` in `attn_backend.h`, which exposes attention interfaces (e.g., `MHA`, `MLA`) to PagedKVCache. We subclass `AttnBackendFunc` and implement FlashInfer backends and TIR backends respectively.
   
   * With `AttnBackendFunc`, we refactored the PagedKVCache constructor. The new constructor is not backward compatible, and will break the existing compiled model libraries.
   
   * For both TIR and FlashInfer attention implementations, now we require an explicit attention softmax scale factor `sm_scale` to be passed in. Previously, it has an inlined `sm_scale` of `head_dim ** -0.5`. Due to the recent LLM inference techniques such as MLA weight absorption in DeepSeek models, the inlined `sm_scale` causes confusion and inconvenience. To keep attention interface standard and clear, we now require the explicit passing of `sm_scale`.
   
   * We refactored the existing GPU unit tests of the PagedKVCache, by updating from numpy to PyTorch for std calculation. This significantly reduces the test case run time.


-- 
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: commits-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

Re: [PR] [KVCache] PagedKVCache refactor, FlashInfer JIT and MLA integration [tvm]

["Hzfengsy (via GitHub)" <gi...@apache.org>]

Hzfengsy merged PR #17674:
URL: https://github.com/apache/tvm/pull/17674


-- 
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: commits-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org