[Runtime] Use preferred host memory (pinned memory) in KV cache

apache / tvm

Open deep learning compiler stack for cpu, gpu and specialized accelerators

https://tvm.apache.org/

Apache License 2.0

11.67k stars 3.45k forks source link

[Runtime] Use preferred host memory (pinned memory) in KV cache #17036

Closed MasterJH5574 closed 4 months ago

MasterJH5574 commented 4 months ago

This PR updates the PagedKVCache with the pinned memory support, which can reduce the copy overhead between CPU and GPU.

This PR also bumps FlashInfer version, which now supports

specifying kernels to build via cmake,
pinned memory as host memory.

We also update CMakeLists.txt and config.cmake to include the FlashInfer compile options. Prior to this PR, the kernels being built is hardcoded in FlashInfer header files.

tqchen commented 4 months ago

@tvm-bot rerun