cache.h 662 Bytes
Newer Older
1
2
#pragma once

Woosuk Kwon's avatar
Woosuk Kwon committed
3
4
#include <torch/extension.h>

5
6
7
8
#include <map>
#include <vector>

void swap_blocks(
Woosuk Kwon's avatar
Woosuk Kwon committed
9
10
11
12
  torch::Tensor& src,
  torch::Tensor& dst,
  const std::map<int64_t, int64_t>& block_mapping);

13
void copy_blocks(
14
15
  std::vector<torch::Tensor>& key_caches,
  std::vector<torch::Tensor>& value_caches,
16
17
  const std::map<int64_t, std::vector<int64_t>>& block_mapping);

Woosuk Kwon's avatar
Woosuk Kwon committed
18
19
20
21
22
void reshape_and_cache(
  torch::Tensor& key,
  torch::Tensor& value,
  torch::Tensor& key_cache,
  torch::Tensor& value_cache,
23
24
  torch::Tensor& slot_mapping,
  const std::string& kv_cache_dtype);
Woosuk Kwon's avatar
Woosuk Kwon committed
25

26
27
28
29
// Just for unittest
void convert_fp8_e5m2(
  torch::Tensor& src_cache,
  torch::Tensor& dst_cache);