2ff767b513
Co-authored-by: Gregory Shtrasberg <Gregory.Shtrasberg@amd.com> Co-authored-by: HaiShaw <hixiao@gmail.com> Co-authored-by: AdrianAbeyta <Adrian.Abeyta@amd.com> Co-authored-by: Matthew Wong <Matthew.Wong2@amd.com> Co-authored-by: root <root@gt-pla-u18-08.pla.dcgpu> Co-authored-by: mawong-amd <156021403+mawong-amd@users.noreply.github.com> Co-authored-by: ttbachyinsda <ttbachyinsda@outlook.com> Co-authored-by: guofangze <guofangze@kuaishou.com> Co-authored-by: Michael Goin <mgoin64@gmail.com> Co-authored-by: jacobthebanana <50071502+jacobthebanana@users.noreply.github.com> Co-authored-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
31 lines
681 B
C++
31 lines
681 B
C++
#pragma once
|
|
|
|
#include <torch/extension.h>
|
|
|
|
#include <map>
|
|
#include <vector>
|
|
|
|
void swap_blocks(
|
|
torch::Tensor& src,
|
|
torch::Tensor& dst,
|
|
const std::map<int64_t, int64_t>& block_mapping);
|
|
|
|
void copy_blocks(
|
|
std::vector<torch::Tensor>& key_caches,
|
|
std::vector<torch::Tensor>& value_caches,
|
|
const std::map<int64_t, std::vector<int64_t>>& block_mapping);
|
|
|
|
void reshape_and_cache(
|
|
torch::Tensor& key,
|
|
torch::Tensor& value,
|
|
torch::Tensor& key_cache,
|
|
torch::Tensor& value_cache,
|
|
torch::Tensor& slot_mapping,
|
|
const std::string& kv_cache_dtype,
|
|
const float kv_scale);
|
|
|
|
// Just for unittest
|
|
void convert_fp8(
|
|
torch::Tensor& src_cache,
|
|
torch::Tensor& dst_cache);
|