Skip to content

Conversation

@mengshyu
Copy link
Contributor

This PR introduces KV Cache support for the CPU runtime:

  • Implementation of KV Cache TIR for CPU-based processing.
  • Updates to the relevant runtime components to integrate KV Cache.

@mengshyu mengshyu force-pushed the 0129-kvcache-cpu branch 4 times, most recently from 8e516a2 to 69b54e7 Compare January 29, 2025 20:00
This PR introduces KV Cache support for the CPU runtime:
* Implementation of KV Cache TIR for CPU-based processing.
* Updates to the relevant runtime components to integrate KV Cache.

Co-authored-by: HMZ <mzhuang@pllab.cs.nthu.edu.tw>
Co-authored-by: ShouChenChiu <scchiu@pllab.cs.nthu.edu.tw>
@mengshyu mengshyu changed the title Add KV Cache for CPU Runtime [KVCache] Add KV Cache for CPU Runtime Jan 30, 2025
@mengshyu
Copy link
Contributor Author

Hi @MasterJH5574 , could you help to review this PR? Thanks. cc. @tqchen

@MasterJH5574
Copy link
Contributor

Thank you @mengshyu! I wonder if we can add some unit tests? We can have a file that is similar to https://github.com/apache/tvm/blob/main/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_tir.py which focuses on the CPU backend.

Copy link
Contributor

@MasterJH5574 MasterJH5574 left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM, thank you @mengshyu!

@MasterJH5574 MasterJH5574 merged commit d93f4ad into apache:main Feb 2, 2025
14 checks passed
ShiboXing pushed a commit to ShiboXing/tvm that referenced this pull request Aug 10, 2025
* [KVCache] Add KV Cache for CPU Runtime

This PR introduces KV Cache support for the CPU runtime:
* Implementation of KV Cache TIR for CPU-based processing.
* Updates to the relevant runtime components to integrate KV Cache.

Co-authored-by: HMZ <mzhuang@pllab.cs.nthu.edu.tw>
Co-authored-by: ShouChenChiu <scchiu@pllab.cs.nthu.edu.tw>

* Add unit test

---------

Co-authored-by: HMZ <mzhuang@pllab.cs.nthu.edu.tw>
Co-authored-by: ShouChenChiu <scchiu@pllab.cs.nthu.edu.tw>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants