131 lines
4.7 KiB
C++
131 lines
4.7 KiB
C++
#include "common.h"
|
|
#include "vec.h"
|
|
|
|
namespace {
|
|
|
|
template <typename scalar_t>
|
|
inline void copy_stub(scalar_t* __restrict__ dst, const scalar_t* __restrict__ src, int size) {
|
|
int d = 0;
|
|
#if defined(CPU_CAPABILITY_AVX512)
|
|
using Vec = at::vec::Vectorized<scalar_t>;
|
|
constexpr int kVecSize = Vec::size();
|
|
|
|
for (; d <= size - kVecSize; d += kVecSize) {
|
|
Vec data = Vec::loadu(src + d);
|
|
data.store(dst + d);
|
|
}
|
|
#endif
|
|
for (; d < size; ++d) {
|
|
dst[d] = src[d];
|
|
}
|
|
}
|
|
|
|
template <typename scalar_t, typename index_t>
|
|
void store_cache_kernel_impl(
|
|
const scalar_t* __restrict__ k,
|
|
const scalar_t* __restrict__ v,
|
|
scalar_t* __restrict__ k_cache,
|
|
scalar_t* __restrict__ v_cache,
|
|
const index_t* __restrict__ indices,
|
|
int64_t batch_size,
|
|
int64_t num_pages,
|
|
int64_t row_dim,
|
|
int64_t k_stride,
|
|
int64_t v_stride,
|
|
int64_t kc_stride,
|
|
int64_t vc_stride) {
|
|
at::parallel_for(0, batch_size, 0, [&](int64_t begin, int64_t end) {
|
|
for (int64_t bs = begin; bs < end; ++bs) {
|
|
const int64_t idx = static_cast<int64_t>(indices[bs]);
|
|
const scalar_t* k_ptr = k + bs * k_stride;
|
|
const scalar_t* v_ptr = v + bs * v_stride;
|
|
scalar_t* kc_ptr = k_cache + idx * kc_stride;
|
|
scalar_t* vc_ptr = v_cache + idx * vc_stride;
|
|
copy_stub(kc_ptr, k_ptr, row_dim);
|
|
copy_stub(vc_ptr, v_ptr, row_dim);
|
|
}
|
|
});
|
|
}
|
|
|
|
} // anonymous namespace
|
|
|
|
// check tensor last two dimensions are contiguous
|
|
#define CHECK_LAST2_DIM_CONTIGUOUS(x, ndim) \
|
|
do { \
|
|
const auto& _x = (x); \
|
|
const auto _ndim = _x.dim(); \
|
|
const auto _strides = _x.strides(); \
|
|
const auto _sizes = _x.sizes(); \
|
|
TORCH_CHECK(_ndim == ndim, #x " must have " #ndim " dimensions"); \
|
|
TORCH_CHECK( \
|
|
_ndim >= 2 && _strides[_ndim - 1] == 1 && _strides[_ndim - 2] == _sizes[_ndim - 1], \
|
|
#x " must be contiguous at the last two dimensions"); \
|
|
} while (0)
|
|
|
|
// [NB]: store_cache takes 3 dimension tensors,
|
|
// This is to avoid the overhead of creating a new TensorImpl
|
|
// from .view(-1, row_dim)
|
|
//
|
|
// k : [batch_size, num_heads, head_size] -> [batch_size, row_dim]
|
|
// v : [batch_size, num_heads, head_size] -> [batch_size, row_dim]
|
|
// k_cache : [num_pages, num_heads, head_size] -> [num_pages, row_dim]
|
|
// v_cache : [num_pages, num_heads, head_size] -> [num_pages, row_dim]
|
|
// indices : [batch_size]
|
|
//
|
|
void store_cache_cpu(
|
|
const at::Tensor& k,
|
|
const at::Tensor& v,
|
|
const at::Tensor& k_cache,
|
|
const at::Tensor& v_cache,
|
|
const at::Tensor& indices,
|
|
std::optional<int64_t> row_dim) {
|
|
CHECK_LAST2_DIM_CONTIGUOUS(k, 3);
|
|
CHECK_LAST2_DIM_CONTIGUOUS(v, 3);
|
|
CHECK_LAST2_DIM_CONTIGUOUS(k_cache, 3);
|
|
CHECK_LAST2_DIM_CONTIGUOUS(v_cache, 3);
|
|
CHECK_INPUT(indices);
|
|
|
|
int64_t batch_size = k.size(0);
|
|
int64_t num_heads = k.size(1);
|
|
int64_t head_size = k.size(2);
|
|
int64_t num_pages = k_cache.size(0);
|
|
int64_t row_dim_value = num_heads * head_size;
|
|
if (row_dim.has_value()) {
|
|
CHECK_EQ(row_dim.value(), row_dim_value);
|
|
}
|
|
CHECK_EQ(indices.size(0), batch_size);
|
|
|
|
// strides: batch dimension (dim 0) stride in elements
|
|
int64_t k_stride = k.stride(0);
|
|
int64_t v_stride = v.stride(0);
|
|
int64_t kc_stride = k_cache.stride(0);
|
|
int64_t vc_stride = v_cache.stride(0);
|
|
|
|
const auto dtype = k.scalar_type();
|
|
TORCH_CHECK(
|
|
dtype == v.scalar_type() && dtype == k_cache.scalar_type() && dtype == v_cache.scalar_type(),
|
|
"store_cache_cpu: input tensors must have the same dtype");
|
|
const auto index_dtype = indices.scalar_type();
|
|
TORCH_CHECK(index_dtype == at::kLong || index_dtype == at::kInt, "indices must be int64 or int32");
|
|
|
|
// dtype : [bfloat16, float16, uint8] for fp8 KV stored as uint8
|
|
// index_dtype : [int64, int32]
|
|
AT_DISPATCH_REDUCED_FLOATING_TYPES_AND(at::ScalarType::Byte, dtype, "store_cache_cpu", [&] {
|
|
AT_DISPATCH_INDEX_TYPES(index_dtype, "store_cache_cpu_index", [&] {
|
|
store_cache_kernel_impl<scalar_t, index_t>(
|
|
k.data_ptr<scalar_t>(),
|
|
v.data_ptr<scalar_t>(),
|
|
k_cache.data_ptr<scalar_t>(),
|
|
v_cache.data_ptr<scalar_t>(),
|
|
indices.data_ptr<index_t>(),
|
|
batch_size,
|
|
num_pages,
|
|
row_dim_value,
|
|
k_stride,
|
|
v_stride,
|
|
kc_stride,
|
|
vc_stride);
|
|
});
|
|
});
|
|
}
|