github
/
ColossalAI
mirror of https://github.com/hpcaitech/ColossalAI
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

[Inference/Feat] Feat quant kvcache step2 (#5674)

pull/5663/head
傅剑寒 2024-04-30 11:26:36 +08:00 committed by GitHub
parent 8ccb6714e7
commit 808ee6e4ad
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
4 changed files with 208 additions and 71 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

116
extensions/csrc/funcs/cast_functor.h
View File

@ -9,6 +9,7 @@
#endif #endif
#include <assert.h> #include <assert.h>
#include <stdint.h>
#include <functional> #include <functional>
@ -175,6 +176,16 @@ COLOSSAL_CAST_FUNCTOR_SPECIALIZATION(uint8_t, uint16_t, DEVICE, STMTS_WRAPPER({
return res.x; return res.x;
})) }))
// half raw -> fp8
COLOSSAL_CAST_FUNCTOR_SPECIALIZATION(uint16_t, uint8_t, DEVICE, STMTS_WRAPPER({
__half_raw tmp;
tmp.x = val;
__nv_fp8_storage_t res =
__nv_cvt_halfraw_to_fp8(
tmp, __NV_SATFINITE, __NV_E5M2);
return static_cast<uint8_t>(res);
}))
// fp8x2 -> half2 raw // fp8x2 -> half2 raw
COLOSSAL_CAST_FUNCTOR_SPECIALIZATION(uint16_t, uint32_t, DEVICE, STMTS_WRAPPER({ COLOSSAL_CAST_FUNCTOR_SPECIALIZATION(uint16_t, uint32_t, DEVICE, STMTS_WRAPPER({
union { union {
@ -222,6 +233,15 @@ COLOSSAL_CAST_FUNCTOR_SPECIALIZATION(uint8_t, half, DEVICE, STMTS_WRAPPER({
return half(res); return half(res);
})) }))
// half -> fp8
COLOSSAL_CAST_FUNCTOR_SPECIALIZATION(half, uint8_t, DEVICE, STMTS_WRAPPER({
__half_raw tmp(val);
__nv_fp8_storage_t res =
__nv_cvt_halfraw_to_fp8(
tmp, __NV_SATFINITE, __NV_E5M2);
return static_cast<uint8_t>(res);
}))
// fp8x2 -> half2 // fp8x2 -> half2
COLOSSAL_CAST_FUNCTOR_SPECIALIZATION(uint16_t, half2, DEVICE, STMTS_WRAPPER({ COLOSSAL_CAST_FUNCTOR_SPECIALIZATION(uint16_t, half2, DEVICE, STMTS_WRAPPER({
__half2_raw res = __half2_raw res =
@ -230,6 +250,15 @@ COLOSSAL_CAST_FUNCTOR_SPECIALIZATION(uint16_t, half2, DEVICE, STMTS_WRAPPER({
return half2(res); return half2(res);
})) }))
// half2 -> fp8x2
COLOSSAL_CAST_FUNCTOR_SPECIALIZATION(half2, uint16_t, DEVICE, STMTS_WRAPPER({
__half2_raw tmp(val);
__nv_fp8x2_storage_t res =
__nv_cvt_halfraw2_to_fp8x2(
tmp, __NV_SATFINITE, __NV_E5M2);
return static_cast<uint16_t>(res);
}))
// fp8x4 -> half4 // fp8x4 -> half4
COLOSSAL_CAST_FUNCTOR_SPECIALIZATION( COLOSSAL_CAST_FUNCTOR_SPECIALIZATION(
uint32_t, dtype::half4, DEVICE, STMTS_WRAPPER({ uint32_t, dtype::half4, DEVICE, STMTS_WRAPPER({
@ -242,6 +271,20 @@ COLOSSAL_CAST_FUNCTOR_SPECIALIZATION(
return res; return res;
})) }))
// half4 -> fp8x4
COLOSSAL_CAST_FUNCTOR_SPECIALIZATION(
dtype::half4, uint32_t, DEVICE, STMTS_WRAPPER({
half2 x, y;
x = val.x;
y = val.y;
uint16_t lo, hi;
lo = CastFunctor<half2, uint16_t>()(x);
hi = CastFunctor<half2, uint16_t>()(y);
uint32_t res;
asm volatile("mov.b32 %0, {%1, %2};\n" : "=r"(res) : "h"(lo), "h"(hi));
return res;
}))
// fp8x8 -> half8 // fp8x8 -> half8
COLOSSAL_CAST_FUNCTOR_SPECIALIZATION( COLOSSAL_CAST_FUNCTOR_SPECIALIZATION(
uint2, dtype::half8, DEVICE, STMTS_WRAPPER({ uint2, dtype::half8, DEVICE, STMTS_WRAPPER({
@ -314,6 +357,14 @@ COLOSSAL_CAST_FUNCTOR_SPECIALIZATION(
return res; return res;
})) }))
// float -> fp8
COLOSSAL_CAST_FUNCTOR_SPECIALIZATION(float, uint8_t, DEVICE, STMTS_WRAPPER({
__nv_fp8_storage_t res =
__nv_cvt_float_to_fp8(
val, __NV_SATFINITE, __NV_E5M2);
return static_cast<uint8_t>(res);
}))
// fp8x2 -> float2 // fp8x2 -> float2
COLOSSAL_CAST_FUNCTOR_SPECIALIZATION( COLOSSAL_CAST_FUNCTOR_SPECIALIZATION(
uint16_t, float2, DEVICE, STMTS_WRAPPER({ uint16_t, float2, DEVICE, STMTS_WRAPPER({
@ -328,6 +379,28 @@ COLOSSAL_CAST_FUNCTOR_SPECIALIZATION(
return make_float2(lof, hif); return make_float2(lof, hif);
})) }))
// float2 -> fp8x2
COLOSSAL_CAST_FUNCTOR_SPECIALIZATION(
float2, uint16_t, DEVICE, STMTS_WRAPPER({
uint16_t tmp1 =
static_cast<uint16_t>(CastFunctor<float, uint8_t>()(val.x));
uint16_t tmp2 =
static_cast<uint16_t>(CastFunctor<float, uint8_t>()(val.y));
uint16_t res = (tmp1 << 8U) | tmp2;
return res;
}))
// float4 -> fp8x4
COLOSSAL_CAST_FUNCTOR_SPECIALIZATION(float4, uint32_t, DEVICE, STMTS_WRAPPER({
uint32_t a, b, c, d;
a = CastFunctor<float, uint8_t>()(val.x);
b = CastFunctor<float, uint8_t>()(val.y);
c = CastFunctor<float, uint8_t>()(val.z);
d = CastFunctor<float, uint8_t>()(val.w);
return (a << 24U) | (b << 16U) |
(c << 8U) | d;
}))
// fp8x4 -> float4_ // fp8x4 -> float4_
COLOSSAL_CAST_FUNCTOR_SPECIALIZATION( COLOSSAL_CAST_FUNCTOR_SPECIALIZATION(
uint32_t, dtype::float4_, DEVICE, STMTS_WRAPPER({ uint32_t, dtype::float4_, DEVICE, STMTS_WRAPPER({
@ -338,6 +411,14 @@ COLOSSAL_CAST_FUNCTOR_SPECIALIZATION(
return res; return res;
})) }))
// fp8x4 -> float4
COLOSSAL_CAST_FUNCTOR_SPECIALIZATION(
uint32_t, float4, DEVICE, STMTS_WRAPPER({
dtype::float4_ tmp = CastFunctor<uint32_t, dtype::float4_>()(val);
float4 res = make_float4(tmp.x.x, tmp.x.y, tmp.y.x, tmp.y.y);
return res;
}))
// fp8x8 -> float8_ // fp8x8 -> float8_
COLOSSAL_CAST_FUNCTOR_SPECIALIZATION( COLOSSAL_CAST_FUNCTOR_SPECIALIZATION(
uint2, dtype::float8_, DEVICE, STMTS_WRAPPER({ uint2, dtype::float8_, DEVICE, STMTS_WRAPPER({
@ -352,16 +433,6 @@ COLOSSAL_CAST_FUNCTOR_SPECIALIZATION(
return res; return res;
})) }))
// half -> fp8
COLOSSAL_CAST_FUNCTOR_SPECIALIZATION(uint16_t, uint8_t, DEVICE, STMTS_WRAPPER({
__half_raw tmp;
tmp.x = val;
__nv_fp8_storage_t res =
__nv_cvt_halfraw_to_fp8(
tmp, __NV_SATFINITE, __NV_E5M2);
return static_cast<uint8_t>(res);
}))
// bf16 -> fp8 // bf16 -> fp8
COLOSSAL_CAST_FUNCTOR_SPECIALIZATION(__nv_bfloat16, uint8_t, DEVICE, COLOSSAL_CAST_FUNCTOR_SPECIALIZATION(__nv_bfloat16, uint8_t, DEVICE,
STMTS_WRAPPER({ STMTS_WRAPPER({
@ -376,19 +447,24 @@ COLOSSAL_CAST_FUNCTOR_SPECIALIZATION(__nv_bfloat16, uint8_t, DEVICE,
#endif #endif
})) }))
// float -> fp8 // bf162 -> fp8x2
COLOSSAL_CAST_FUNCTOR_SPECIALIZATION(float, uint8_t, DEVICE, STMTS_WRAPPER({ COLOSSAL_CAST_FUNCTOR_SPECIALIZATION(
__nv_fp8_storage_t res = __nv_bfloat162, uint16_t, DEVICE, STMTS_WRAPPER({
__nv_cvt_float_to_fp8( uint16_t a =
val, __NV_SATFINITE, __NV_E5M2); static_cast<uint16_t>(CastFunctor<__nv_bfloat16, uint8_t>()(val.x));
return static_cast<uint8_t>(res); uint16_t b =
static_cast<uint16_t>(CastFunctor<__nv_bfloat16, uint8_t>()(val.y));
return (a << 8U) | b;
})) }))
// fp8x4 -> float4 // bf164 -> fp8x4
COLOSSAL_CAST_FUNCTOR_SPECIALIZATION( COLOSSAL_CAST_FUNCTOR_SPECIALIZATION(
uint32_t, float4, DEVICE, STMTS_WRAPPER({ dtype::bfloat164, uint32_t, DEVICE, STMTS_WRAPPER({
dtype::float4_ tmp = CastFunctor<uint32_t, dtype::float4_>()(val); uint32_t res;
float4 res = make_float4(tmp.x.x, tmp.x.y, tmp.y.x, tmp.y.y); uint16_t a, b;
a = CastFunctor<__nv_bfloat162, uint16_t>()(val.x);
b = CastFunctor<__nv_bfloat162, uint16_t>()(val.y);
asm volatile("mov.b32 %0, {%1, %2};\n" : "=r"(res) : "h"(a), "h"(b));
return res; return res;
})) }))

126
extensions/csrc/kernel/cuda/context_kv_cache_memcpy_kernel.cu
View File

@ -4,16 +4,17 @@
#include "utils/vec_copy.h" #include "utils/vec_copy.h"
#include "common/micros.h" #include "common/micros.h"
using colossalAI::cuda::utils::copy_vector;
using colossalAI::cuda::utils::get_vec_size; using colossalAI::cuda::utils::get_vec_size;
using colossalAI::cuda::utils::copy;
using colossalAI::funcs::CastFunctor;
template<typename scalar_t, bool Aligned, int VecSize> template<typename T, typename CacheT, bool Aligned, int VecSize>
__global__ void context_kv_cache_memcpy_kernel( __global__ void context_kv_cache_memcpy_kernel(
const scalar_t* __restrict__ key, const T* __restrict__ key,
const scalar_t* __restrict__ value, const T* __restrict__ value,
scalar_t* __restrict__ key_cache, CacheT* __restrict__ key_cache,
scalar_t* __restrict__ value_cache, CacheT* __restrict__ value_cache,
const int* __restrict__ sequence_lengths, const int* __restrict__ sequence_lengths,
const int* __restrict__ cu_seqlens, const int* __restrict__ cu_seqlens,
const int* __restrict__ block_tables, const int* __restrict__ block_tables,
@ -54,8 +55,8 @@ __global__ void context_kv_cache_memcpy_kernel(
+ head_id * block_size * head_dim + head_id * block_size * head_dim
+ block_offset * head_dim + head_offset; + block_offset * head_dim + head_offset;
copy_vector<scalar_t, VecSize>(key_cache + target_id, key + key_src_id); copy<T, CacheT, VecSize>(key + key_src_id, key_cache + target_id);
copy_vector<scalar_t, VecSize>(value_cache + target_id, value + value_src_id); copy<T, CacheT, VecSize>(value + value_src_id, value_cache + target_id);
} }
// tail process // tail process
@ -69,22 +70,22 @@ __global__ void context_kv_cache_memcpy_kernel(
+ head_id * block_size * head_dim + head_id * block_size * head_dim
+ block_offset * head_dim + head_offset; + block_offset * head_dim + head_offset;
key_cache[target_id] = key[key_src_id]; key_cache[target_id] = CastFunctor<T, CacheT>()(key[key_src_id]);
value_cache[target_id] = value[value_src_id]; value_cache[target_id] = CastFunctor<T, CacheT>()(value[value_src_id]);
} }
} }
} }
template<typename scalar_t> template<typename T, typename CacheT>
void apply_context_kv_cache_memcpy( void apply_context_kv_cache_memcpy(
at::Tensor& key, // [num_tokens, head_num, head_dim] torch::Tensor& key, // [num_tokens, head_num, head_dim]
at::Tensor& value, // [num_tokens, head_num, head_dim] torch::Tensor& value, // [num_tokens, head_num, head_dim]
at::Tensor& key_cache, // [num_blocks, head_num, block_size, head_dim] torch::Tensor& key_cache, // [num_blocks, head_num, block_size, head_dim]
at::Tensor& value_cache, // [num_blocks, head_num, block_size, head_dim] torch::Tensor& value_cache, // [num_blocks, head_num, block_size, head_dim]
at::Tensor& sequence_lengths, // [batch_size] torch::Tensor& sequence_lengths, // [batch_size]
at::Tensor& cu_seqlens, // [batch_size + 1] torch::Tensor& cu_seqlens, // [batch_size + 1]
at::Tensor& block_tables, // [batch_size, max_seq_len] torch::Tensor& block_tables, // [batch_size, max_seq_len]
int max_seq_len_in_batch) int max_seq_len_in_batch)
{ {
int num_tokens = key.size(0); int num_tokens = key.size(0);
@ -97,7 +98,7 @@ void apply_context_kv_cache_memcpy(
int64_t value_stride = value.stride(0); int64_t value_stride = value.stride(0);
int block_table_stride = block_tables.stride(0); int block_table_stride = block_tables.stride(0);
int vec_size = get_vec_size<scalar_t>(key); int vec_size = get_vec_size<T>(key);
bool aligned = true; bool aligned = true;
if (head_dim % vec_size != 0) { if (head_dim % vec_size != 0) {
@ -112,11 +113,11 @@ void apply_context_kv_cache_memcpy(
#define CONTEXT_KV_CACHE_MEMCOPY_KERNEL_LAUNCH(__aligned, __vec_size) \ #define CONTEXT_KV_CACHE_MEMCOPY_KERNEL_LAUNCH(__aligned, __vec_size) \
do { \ do { \
context_kv_cache_memcpy_kernel<scalar_t, __aligned, __vec_size><<<grid, block, 0, stream>>>( \ context_kv_cache_memcpy_kernel<T, CacheT, __aligned, __vec_size><<<grid, block, 0, stream>>>( \
key.data_ptr<scalar_t>(), \ reinterpret_cast<T*>(key.data_ptr()), \
value.data_ptr<scalar_t>(), \ reinterpret_cast<T*>(value.data_ptr()), \
key_cache.data_ptr<scalar_t>(), \ reinterpret_cast<CacheT*>(key_cache.data_ptr()), \
value_cache.data_ptr<scalar_t>(), \ reinterpret_cast<CacheT*>(value_cache.data_ptr()), \
sequence_lengths.data_ptr<int>(), \ sequence_lengths.data_ptr<int>(), \
cu_seqlens.data_ptr<int>(), \ cu_seqlens.data_ptr<int>(), \
block_tables.data_ptr<int>(), \ block_tables.data_ptr<int>(), \
@ -161,26 +162,63 @@ void apply_context_kv_cache_memcpy(
} }
void context_kv_cache_memcpy( void context_kv_cache_memcpy(
at::Tensor& key, // [num_tokens, head_num, head_dim] torch::Tensor& key, // [num_tokens, head_num, head_dim]
at::Tensor& value, // [num_tokens, head_num, head_dim] torch::Tensor& value, // [num_tokens, head_num, head_dim]
at::Tensor& key_cache, // [num_blocks, head_num, block_size, head_dim] torch::Tensor& key_cache, // [num_blocks, head_num, block_size, head_dim]
at::Tensor& value_cache, // [num_blocks, head_num, block_size, head_dim] torch::Tensor& value_cache, // [num_blocks, head_num, block_size, head_dim]
at::Tensor& sequence_lengths, // [batch_size] torch::Tensor& sequence_lengths, // [batch_size]
at::Tensor& cu_seqlens, // [batch_size + 1] torch::Tensor& cu_seqlens, // [batch_size + 1]
at::Tensor& block_tables, // [batch_size, max_seq_len] torch::Tensor& block_tables, // [batch_size, max_seq_len]
int max_seq_len_in_batch) int max_seq_len_in_batch)
{ {
DISPATCH_FLOAT_HALF_AND_BFLOAT(
key.scalar_type(), TORCH_CHECK(key.scalar_type() == at::ScalarType::Float || key.scalar_type() == at::ScalarType::Half || key.scalar_type() == at::ScalarType::BFloat16,
"context_kv_cache_memcpy", "Dtype of key should be float, half or bfloat16!");
apply_context_kv_cache_memcpy<scalar_t>( TORCH_CHECK(key_cache.scalar_type() == at::ScalarType::Byte || key_cache.scalar_type() == key.scalar_type(),
key, "Dtype of query and kvcache should be the same unless dtype of kvcache is fp8!");
value,
key_cache,
value_cache, #define _(T, CacheT) \
sequence_lengths, apply_context_kv_cache_memcpy<T, CacheT>( \
cu_seqlens, key, \
block_tables, value, \
max_seq_len_in_batch key_cache, \
);) value_cache, \
sequence_lengths, \
cu_seqlens, \
block_tables, \
max_seq_len_in_batch \
)
if(key_cache.scalar_type() == at::ScalarType::Byte)
{
switch (key.scalar_type())
{
case at::ScalarType::Float:
_(float, uint8_t);
break;
case at::ScalarType::Half:
_(half, uint8_t);
break;
case at::ScalarType::BFloat16:
_(__nv_bfloat16, uint8_t);
break;
}
}
else
{
switch (key.scalar_type())
{
case at::ScalarType::Float:
_(float, float);
break;
case at::ScalarType::Half:
_(half, half);
break;
case at::ScalarType::BFloat16:
_(__nv_bfloat16, __nv_bfloat16);
break;
}
}
#undef _
} }

2
extensions/csrc/kernel/cuda/flash_decoding_attention_kernel.cu
View File

@ -372,7 +372,7 @@ void flash_decoding_attention(
TORCH_CHECK(query.scalar_type() == at::ScalarType::Float || query.scalar_type() == at::ScalarType::Half || query.scalar_type() == at::ScalarType::BFloat16, TORCH_CHECK(query.scalar_type() == at::ScalarType::Float || query.scalar_type() == at::ScalarType::Half || query.scalar_type() == at::ScalarType::BFloat16,
"Dtype of query should be float, half or bfloat16!"); "Dtype of query should be float, half or bfloat16!");
TORCH_CHECK(key_cache.scalar_type() == at::ScalarType::Byte || key_cache.scalar_type() == key_cache.scalar_type(), TORCH_CHECK(key_cache.scalar_type() == at::ScalarType::Byte || key_cache.scalar_type() == query.scalar_type(),
"Dtype of query and kvcache should be the same unless dtype of kvcache is fp8!"); "Dtype of query and kvcache should be the same unless dtype of kvcache is fp8!");
if(key_cache.scalar_type() == at::ScalarType::Byte) if(key_cache.scalar_type() == at::ScalarType::Byte)

31
extensions/csrc/kernel/cuda/utils/vec_copy.h
View File

@ -11,10 +11,9 @@ namespace colossalAI {
namespace cuda { namespace cuda {
namespace utils { namespace utils {
template <typename T, int VecSize> template <typename T, int vec_size>
__device__ __inline__ void copy_vector(T *dst, const T *src) { __device__ __inline__ void copy_vector(T *dst, const T *src) {
using VT = typename common::VecTypeTrait<T, VecSize>::Type; using VT = typename common::VecTypeTrait<T, vec_size>::Type;
// Note(LiuYang): Here static_cast can't be used for cast between two pointer
*(reinterpret_cast<VT *>(dst)) = *(reinterpret_cast<const VT *>(src)); *(reinterpret_cast<VT *>(dst)) = *(reinterpret_cast<const VT *>(src));
} }
@ -33,9 +32,33 @@ __device__ __inline__ void copy_zero_vector(T *dst) {
*(reinterpret_cast<VT *>(dst)) = funcs::CastFunctor<float, VT>()(0.0f); *(reinterpret_cast<VT *>(dst)) = funcs::CastFunctor<float, VT>()(0.0f);
} }
template <typename SrcT, typename DstT, int vec_size>
__device__ __inline__ void copy(const SrcT *src, DstT *dst) {
using SrcVT = typename common::VecTypeTrait<SrcT, vec_size>::Type;
using DstVT = typename common::VecTypeTrait<DstT, vec_size>::Type;
// Note(LiuYang): Here static_cast can't be used for cast between two pointer
*(reinterpret_cast<DstVT *>(dst)) = funcs::CastFunctor<SrcVT, DstVT>()(
*(reinterpret_cast<const SrcVT *>(src)));
}
template <typename T, int vec_size>
__device__ __inline__ void copy<T, T, vec_size>(const T *src, T *dst) {
using VT = typename common::VecTypeTrait<T, vec_size>::Type;
*(reinterpret_cast<VT *>(dst)) = *(reinterpret_cast<const VT *>(src));
}
template <>
__device__ __inline__ void copy<float, float, 8>(const float *src, float *dst) {
// Since the maximum memory alignment length is 128 bits, we choose float4
// here.
*(reinterpret_cast<float4 *>(dst)) = *(reinterpret_cast<const float4 *>(src));
*(reinterpret_cast<float4 *>(dst + 4)) =
*(reinterpret_cast<const float4 *>(src + 4));
}
template <typename T> template <typename T>
int get_vec_size(const torch::Tensor &tensor) { int get_vec_size(const torch::Tensor &tensor) {
uint64_t address = reinterpret_cast<uint64_t>(tensor.data_ptr<T>()); uint64_t address = reinterpret_cast<uint64_t>(tensor.data_ptr());
const int max_aligned_size = 128; const int max_aligned_size = 128;
const int dtype_size = sizeof(T) * 8; const int dtype_size = sizeof(T) * 8;