#include #include #include #include void swap_blocks( torch::Tensor& src, torch::Tensor& dst, const std::map& block_mapping); void copy_blocks( std::vector& key_caches, std::vector& value_caches, const std::map>& block_mapping); void reshape_and_cache( torch::Tensor& key, torch::Tensor& value, torch::Tensor& key_cache, torch::Tensor& value_cache, torch::Tensor& slot_mapping); void gather_cached_kv( torch::Tensor& key, torch::Tensor& value, torch::Tensor& key_cache, torch::Tensor& value_cache, torch::Tensor& slot_mapping); PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def( "swap_blocks", &swap_blocks, "Swap in (out) the cache blocks from src to dst"); m.def( "copy_blocks", ©_blocks, "Copy the cache blocks from src to dst"); m.def( "reshape_and_cache", &reshape_and_cache, "Reshape the key and value tensors and cache them"); m.def( "gather_cached_kv", &gather_cached_kv, "Gather key and value from the cache into contiguous KV tensors"); }