wehub-resource-sync
527 lines
25 KiB
Plaintext
527 lines
25 KiB
Plaintext
// Copyright (c) Microsoft Corporation.
|
|
// SPDX-License-Identifier: Apache-2.0
|
|
|
|
// DeepSpeed Team
|
|
|
|
#include <stdexcept>
|
|
#include "fp_context.h"
|
|
#include "fp_quantize.h"
|
|
#include "memory_access_utils.h"
|
|
#include "reduction_utils.h"
|
|
|
|
#include <cuda.h>
|
|
#include <stdint.h>
|
|
|
|
#include <cuda_fp16.h>
|
|
#include <curand_kernel.h>
|
|
|
|
#ifdef BF16_AVAILABLE
|
|
#include <cuda_bf16.h>
|
|
#endif
|
|
#include <cuda_runtime_api.h>
|
|
|
|
using ROp = reduce::ROpType;
|
|
|
|
namespace quantization {
|
|
|
|
constexpr int access_granularity = 16;
|
|
constexpr int quanitzed_access_granularity = 4;
|
|
constexpr int quanitzed_access_granularity_6bits = 2;
|
|
constexpr int threads = 256;
|
|
constexpr int warps = threads / 32;
|
|
|
|
} // namespace quantization
|
|
|
|
template <int _mantisa_bits, int q_mantisa_bits, int stochastic_rounding>
|
|
__device__ void round(uint32_t& mantisa, uint32_t& dst_exponent, curandStatePhilox4_32_10_t* state)
|
|
{
|
|
constexpr uint32_t mantisa_mask = (1U << (_mantisa_bits - q_mantisa_bits)) - 1;
|
|
uint32_t offset = stochastic_rounding ? (curand_poisson(state, 10) & mantisa_mask)
|
|
: 1U << (_mantisa_bits - q_mantisa_bits - 1);
|
|
mantisa += offset;
|
|
dst_exponent += (((mantisa & ~mantisa_mask) == (1U << _mantisa_bits)) ? 1 : 0);
|
|
}
|
|
|
|
template <int _mantisa_bits, int _exponent_bits, int q_mantisa_bits, int q_exponent_bits>
|
|
__device__ void clip(uint32_t& exponent, uint32_t& mantisa)
|
|
{
|
|
constexpr uint32_t max_exponent = (1 << (q_exponent_bits - 1)) + (1 << (_exponent_bits - 1));
|
|
constexpr uint32_t min_exponent =
|
|
(1 << (_exponent_bits - 1)) - ((1 << (q_exponent_bits - 1)) - 1);
|
|
if (exponent > max_exponent) {
|
|
exponent = max_exponent;
|
|
mantisa = (((uint32_t)-1) >> (32 - q_mantisa_bits)) << 1; //.11 .. 10
|
|
}
|
|
if (exponent < min_exponent) {
|
|
exponent = min_exponent;
|
|
mantisa = 0;
|
|
}
|
|
}
|
|
|
|
template <typename T,
|
|
int unroll,
|
|
int _mantisa_bits,
|
|
int _exponent_bits,
|
|
int total_q_bits = 8,
|
|
int q_mantisa_bits = 3,
|
|
int stochastic_rounding = 0>
|
|
__global__ void apply_quantization(T* val,
|
|
uint8_t* q_val,
|
|
int group_size,
|
|
std::pair<uint64_t, uint64_t> seed,
|
|
float q_range)
|
|
{
|
|
int tidx = threadIdx.x;
|
|
int wid = tidx >> 5;
|
|
int lane = tidx & 0x1f;
|
|
int gid = blockIdx.x * quantization::warps + wid;
|
|
|
|
constexpr int q_exponent_bits = total_q_bits - q_mantisa_bits - 1;
|
|
constexpr uint32_t _mantisa_mask = (1 << _mantisa_bits) - 1;
|
|
constexpr uint32_t _exponent_mask = ((1 << _exponent_bits) - 1) << _mantisa_bits;
|
|
constexpr uint32_t _sign_mask = 1U << (_mantisa_bits + _exponent_bits);
|
|
// CG helpers
|
|
cg::thread_block tb = cg::this_thread_block();
|
|
cg::thread_block_tile<hw_warp_size> warp = cg::tiled_partition<hw_warp_size>(tb);
|
|
|
|
constexpr uint32_t vector_size = quantization::access_granularity / sizeof(T);
|
|
constexpr uint32_t load_stride = vector_size * hw_warp_size;
|
|
constexpr uint32_t store_stride = (total_q_bits * vector_size / 8) * hw_warp_size;
|
|
const uint32_t thread_offset = lane * vector_size;
|
|
const uint32_t store_thread_offset = lane * (total_q_bits * vector_size / 8);
|
|
const uint32_t base_load_offset = gid * group_size + thread_offset;
|
|
const uint32_t base_store_offset =
|
|
gid * ((group_size * total_q_bits / 8) + 4) +
|
|
store_thread_offset; // 4-byte for saving the scale per group
|
|
const T* load_base_ptr = val + base_load_offset;
|
|
T tmp_buf[unroll * vector_size];
|
|
T cur_max;
|
|
reduce::init<ROp::Max>(&cur_max);
|
|
|
|
int idx = blockIdx.x * blockDim.x + threadIdx.x;
|
|
curandStatePhilox4_32_10_t state;
|
|
curand_init(seed.first, idx, seed.second, &state);
|
|
|
|
#pragma unroll
|
|
for (int i = 0; i < unroll; i++) {
|
|
if (i * load_stride + thread_offset < group_size) {
|
|
mem_access::load_global<quantization::access_granularity>(
|
|
&tmp_buf[vector_size * i], load_base_ptr + i * load_stride);
|
|
for (int j = 0; j < vector_size; j++)
|
|
cur_max = reduce::element<ROp::Max>(cur_max, __habs(tmp_buf[i * vector_size + j]));
|
|
}
|
|
}
|
|
reduce::_block<T, 1, ROp::Max>(tb, warp, &cur_max);
|
|
|
|
int mantisa_mask = ((1 << q_mantisa_bits) - 1);
|
|
mantisa_mask <<= (_mantisa_bits - q_mantisa_bits);
|
|
|
|
uint8_t* store_base_ptr = q_val + base_store_offset;
|
|
float scale = (float)q_range / conversion::to<float>(cur_max);
|
|
#pragma unroll
|
|
for (int i = 0; i < unroll; i++) {
|
|
if (i * load_stride + thread_offset < group_size) {
|
|
uint64_t q_buf = 0;
|
|
uint64_t q_buf1 = 0;
|
|
#pragma unroll
|
|
for (int j = 0; j < vector_size; j++) {
|
|
float val_f = conversion::to<float>(tmp_buf[i * vector_size + j]) * scale;
|
|
uint32_t* data = reinterpret_cast<uint32_t*>(&val_f);
|
|
uint32_t sign = (data[0] & _sign_mask) >> (_mantisa_bits + _exponent_bits);
|
|
uint32_t cur_exponent = (data[0] & _exponent_mask) >> _mantisa_bits;
|
|
uint32_t dst_mantisa = (data[0] & _mantisa_mask);
|
|
|
|
uint32_t dst_exponent = cur_exponent;
|
|
|
|
round<_mantisa_bits, q_mantisa_bits, stochastic_rounding>(
|
|
dst_mantisa, dst_exponent, &state);
|
|
if (cur_exponent != 0)
|
|
clip<_mantisa_bits, _exponent_bits, q_mantisa_bits, q_exponent_bits>(
|
|
dst_exponent, dst_mantisa);
|
|
|
|
dst_mantisa = (dst_mantisa & mantisa_mask) >> (_mantisa_bits - q_mantisa_bits);
|
|
|
|
if (dst_exponent != (1 << q_exponent_bits) - 1)
|
|
dst_exponent = (dst_exponent - ((1 << (_exponent_bits - 1)) - 1)) +
|
|
(1 << (q_exponent_bits - 1)) - 1;
|
|
if (total_q_bits == 8 || total_q_bits == 4 || total_q_bits == 6)
|
|
q_buf = q_buf |
|
|
((uint64_t)((uint8_t)(sign << (q_exponent_bits + q_mantisa_bits) |
|
|
(dst_exponent << q_mantisa_bits) | dst_mantisa))
|
|
<< j * total_q_bits);
|
|
else if (total_q_bits == 12) {
|
|
if (j < 5)
|
|
q_buf =
|
|
q_buf |
|
|
((uint64_t)((uint16_t)(sign << (q_exponent_bits + q_mantisa_bits) |
|
|
(dst_exponent << q_mantisa_bits) | dst_mantisa))
|
|
<< j * total_q_bits);
|
|
else
|
|
q_buf1 =
|
|
q_buf1 |
|
|
((uint64_t)((uint16_t)(sign << (q_exponent_bits + q_mantisa_bits) |
|
|
(dst_exponent << q_mantisa_bits) | dst_mantisa))
|
|
<< (j - 5) * total_q_bits);
|
|
}
|
|
}
|
|
if (total_q_bits == 12) {
|
|
uint64_t last_nibble_mask = 0xf;
|
|
last_nibble_mask = q_buf1 & last_nibble_mask;
|
|
q_buf = (last_nibble_mask << 60) | q_buf;
|
|
q_buf1 >>= 4;
|
|
}
|
|
uint8_t* int8_data = reinterpret_cast<uint8_t*>(&q_buf);
|
|
uint8_t* int8_data1 = reinterpret_cast<uint8_t*>(&q_buf1);
|
|
if (total_q_bits == 6) {
|
|
mem_access::store_global<quantization::quanitzed_access_granularity_6bits>(
|
|
store_base_ptr + i * store_stride, int8_data);
|
|
mem_access::store_global<quantization::quanitzed_access_granularity_6bits>(
|
|
store_base_ptr + i * store_stride +
|
|
quantization::quanitzed_access_granularity_6bits,
|
|
int8_data + quantization::quanitzed_access_granularity_6bits);
|
|
mem_access::store_global<quantization::quanitzed_access_granularity_6bits>(
|
|
store_base_ptr + i * store_stride +
|
|
quantization::quanitzed_access_granularity_6bits * 2,
|
|
int8_data + 2 * quantization::quanitzed_access_granularity_6bits);
|
|
} else {
|
|
mem_access::store_global<quantization::quanitzed_access_granularity>(
|
|
store_base_ptr + i * store_stride, int8_data);
|
|
|
|
if (total_q_bits > 4) {
|
|
mem_access::store_global<quantization::quanitzed_access_granularity>(
|
|
store_base_ptr + i * store_stride +
|
|
quantization::quanitzed_access_granularity,
|
|
int8_data + quantization::quanitzed_access_granularity);
|
|
if (total_q_bits == 12) {
|
|
mem_access::store_global<quantization::quanitzed_access_granularity>(
|
|
store_base_ptr + i * store_stride +
|
|
quantization::quanitzed_access_granularity * 2,
|
|
int8_data1);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
if (lane == 0) {
|
|
float q_scale = conversion::to<float>(cur_max) / (float)q_range;
|
|
uint8_t* scale_as_int8 = reinterpret_cast<uint8_t*>(&q_scale);
|
|
uint32_t scale_offset =
|
|
gid * ((group_size * total_q_bits / 8) + 4) + (group_size * total_q_bits / 8);
|
|
if (total_q_bits != 6)
|
|
mem_access::store_global<quantization::quanitzed_access_granularity>(
|
|
q_val + scale_offset, scale_as_int8);
|
|
else {
|
|
mem_access::store_global<quantization::quanitzed_access_granularity_6bits>(
|
|
q_val + scale_offset, scale_as_int8);
|
|
mem_access::store_global<quantization::quanitzed_access_granularity_6bits>(
|
|
q_val + scale_offset + quantization::quanitzed_access_granularity_6bits,
|
|
scale_as_int8 + quantization::quanitzed_access_granularity_6bits);
|
|
}
|
|
}
|
|
}
|
|
|
|
template <typename T,
|
|
int mantisa_bits,
|
|
int total_q_bits = 16,
|
|
int q_mantisa_bits = 3,
|
|
int q_exponent_bits = 4>
|
|
__global__ void apply_dequantization(uint8_t* val, T* q_val, int group_size, int total_num_elements)
|
|
{
|
|
constexpr uint32_t vector_size = quantization::access_granularity / sizeof(T);
|
|
int tidx = (blockIdx.x * blockDim.x + threadIdx.x) * vector_size;
|
|
|
|
constexpr int quantized_bits = q_mantisa_bits + q_exponent_bits + 1;
|
|
constexpr int q_exponent_bits = total_q_bits - mantisa_bits - 1;
|
|
constexpr uint16_t _mantisa_mask = (1 << q_mantisa_bits) - 1;
|
|
constexpr uint16_t _exponent_mask = ((1 << q_exponent_bits) - 1) << q_mantisa_bits;
|
|
constexpr uint16_t _sign_mask = 1U << (q_mantisa_bits + q_exponent_bits);
|
|
const uint32_t g_index = (tidx / group_size);
|
|
const uint32_t group_size_bytes = (group_size * quantized_bits / 8);
|
|
const uint8_t* load_base_ptr =
|
|
val + g_index * (group_size_bytes + 4) + (tidx % group_size) * quantized_bits / 8;
|
|
|
|
T* store_base_ptr = q_val + tidx;
|
|
float scale;
|
|
|
|
uint8_t* scale_as_int8 = reinterpret_cast<uint8_t*>(&scale);
|
|
if (quantized_bits == 6) {
|
|
mem_access::load_global<quantization::quanitzed_access_granularity>(
|
|
scale_as_int8, val + g_index * (group_size_bytes + 4) + group_size_bytes);
|
|
mem_access::load_global<quantization::quanitzed_access_granularity_6bits>(
|
|
scale_as_int8 + quantization::quanitzed_access_granularity_6bits,
|
|
val + g_index * (group_size_bytes + 4) + group_size_bytes +
|
|
quantization::quanitzed_access_granularity_6bits);
|
|
} else
|
|
mem_access::load_global<quantization::quanitzed_access_granularity>(
|
|
scale_as_int8, val + g_index * (group_size_bytes + 4) + group_size_bytes);
|
|
|
|
if (tidx < total_num_elements) {
|
|
uint64_t q_buf_in;
|
|
uint64_t q_buf_in1;
|
|
uint8_t* int8_data = reinterpret_cast<uint8_t*>(&q_buf_in);
|
|
uint8_t* int8_data1 = reinterpret_cast<uint8_t*>(&q_buf_in1);
|
|
if (quantized_bits == 6) {
|
|
mem_access::load_global<quantization::quanitzed_access_granularity_6bits>(
|
|
int8_data, load_base_ptr);
|
|
mem_access::load_global<quantization::quanitzed_access_granularity_6bits>(
|
|
int8_data + quantization::quanitzed_access_granularity_6bits,
|
|
load_base_ptr + quantization::quanitzed_access_granularity_6bits);
|
|
mem_access::load_global<quantization::quanitzed_access_granularity_6bits>(
|
|
int8_data + quantization::quanitzed_access_granularity_6bits * 2,
|
|
load_base_ptr + quantization::quanitzed_access_granularity_6bits * 2);
|
|
|
|
} else {
|
|
mem_access::load_global<quantization::quanitzed_access_granularity>(int8_data,
|
|
load_base_ptr);
|
|
if (quantized_bits > 4) {
|
|
mem_access::load_global<quantization::quanitzed_access_granularity>(
|
|
int8_data + quantization::quanitzed_access_granularity,
|
|
load_base_ptr + quantization::quanitzed_access_granularity);
|
|
if (quantized_bits == 12) {
|
|
mem_access::load_global<quantization::quanitzed_access_granularity>(
|
|
int8_data1, load_base_ptr + quantization::quanitzed_access_granularity * 2);
|
|
}
|
|
}
|
|
}
|
|
T store_buf[vector_size];
|
|
uint16_t* q_buf = reinterpret_cast<uint16_t*>(store_buf);
|
|
#pragma unroll
|
|
for (int j = 0; j < vector_size; j++) {
|
|
uint16_t new_data;
|
|
if (j < 5 || quantized_bits != 12) {
|
|
new_data = (uint16_t)(q_buf_in >> (j * quantized_bits));
|
|
} else {
|
|
if (j == 5) {
|
|
new_data = (uint16_t)(q_buf_in1);
|
|
new_data = (uint16_t)((new_data << 4) | (q_buf_in >> 60));
|
|
} else
|
|
new_data = (uint16_t)(q_buf_in1 >> ((j - 6) * quantized_bits + 8));
|
|
}
|
|
|
|
uint16_t sign = (new_data & _sign_mask) >> (q_mantisa_bits + q_exponent_bits);
|
|
uint16_t dst_exponent = (new_data & _exponent_mask) >> q_mantisa_bits;
|
|
uint16_t dst_mantisa = (new_data & _mantisa_mask);
|
|
|
|
if (dst_exponent != (1 << q_exponent_bits) - 1)
|
|
dst_exponent = (dst_exponent - ((1 << (q_exponent_bits - 1)) - 1)) +
|
|
(1 << (q_exponent_bits - 1)) - 1;
|
|
|
|
q_buf[j] =
|
|
((sign << (q_exponent_bits + mantisa_bits)) | (dst_exponent << mantisa_bits) |
|
|
(dst_mantisa << (mantisa_bits - q_mantisa_bits)));
|
|
float up_cast = conversion::to<float>(store_buf[j]);
|
|
store_buf[j] = conversion::to<T>(up_cast * scale);
|
|
}
|
|
mem_access::store_global<quantization::access_granularity>(store_base_ptr, store_buf);
|
|
}
|
|
}
|
|
|
|
#define LAUNCH_FOR_QUANTIZATION_UNROLL(COUNT) \
|
|
case COUNT: \
|
|
apply_quantization<T, \
|
|
COUNT, \
|
|
mantisa, \
|
|
exponent, \
|
|
CONST_Q_BITS, \
|
|
CONST_Q_MANTISA_BITS, \
|
|
CONST_STOCHASTIC_ROUNDING> \
|
|
<<<grid, block, 0, stream>>>(val, q_val, group_size, seed, q_range); \
|
|
break;
|
|
|
|
template <typename T, int mantisa, int exponent>
|
|
void launch_quantization(T* val,
|
|
uint8_t* q_val,
|
|
int num_groups,
|
|
int group_size,
|
|
cudaStream_t stream,
|
|
float q_range,
|
|
int q_bits,
|
|
int q_mantisa_bits,
|
|
int stochastic_rounding)
|
|
{
|
|
const dim3 grid((num_groups + quantization::warps - 1) / quantization::warps);
|
|
const dim3 block(quantization::threads);
|
|
|
|
std::pair<uint64_t, uint64_t> seed = FPContext::Instance().IncrementOffset(16);
|
|
|
|
constexpr int vals_per_unroll = hw_warp_size * quantization::access_granularity / sizeof(T);
|
|
|
|
const int copy_unroll = (group_size + vals_per_unroll - 1) / vals_per_unroll;
|
|
QUANT_SWITCH((q_bits - q_mantisa_bits - 1) * q_mantisa_bits + stochastic_rounding, [&] {
|
|
switch (copy_unroll) {
|
|
LAUNCH_FOR_QUANTIZATION_UNROLL(1)
|
|
LAUNCH_FOR_QUANTIZATION_UNROLL(2)
|
|
LAUNCH_FOR_QUANTIZATION_UNROLL(3)
|
|
LAUNCH_FOR_QUANTIZATION_UNROLL(4)
|
|
LAUNCH_FOR_QUANTIZATION_UNROLL(5)
|
|
LAUNCH_FOR_QUANTIZATION_UNROLL(6)
|
|
}
|
|
});
|
|
}
|
|
#define INSTANTIATE_LAUNCH_QUANTIZATION(T, mantisa, exponent) \
|
|
template void launch_quantization<T, mantisa, exponent>( \
|
|
T*, uint8_t*, int, int, cudaStream_t, float q_range, int, int, int);
|
|
// fp8(E4M3), nearest-rounding
|
|
#ifdef BF16_AVAILABLE
|
|
INSTANTIATE_LAUNCH_QUANTIZATION(__nv_bfloat16, 23, 8);
|
|
#endif
|
|
INSTANTIATE_LAUNCH_QUANTIZATION(__half, 23, 8);
|
|
|
|
template <typename T, int mantisa>
|
|
void launch_dequantization(uint8_t* val,
|
|
T* q_val,
|
|
int num_groups,
|
|
int group_size,
|
|
int q_mantisa_bits,
|
|
int q_exponent_bits,
|
|
cudaStream_t stream)
|
|
{
|
|
int blocks = ((num_groups * group_size) - 1) /
|
|
(quantization::threads * (quantization::access_granularity / sizeof(T))) +
|
|
1;
|
|
const dim3 grid(blocks);
|
|
const dim3 block(quantization::threads);
|
|
DEQUANT_SWITCH(q_mantisa_bits * q_exponent_bits, [&] {
|
|
apply_dequantization<T, mantisa, 16, CONST_Q_MANTISA_BITS, CONST_Q_EXPONENT_BITS>
|
|
<<<grid, block, 0, stream>>>(val, q_val, group_size, (num_groups * group_size));
|
|
});
|
|
}
|
|
#define INSTANTIATE_LAUNCH_DEQUANTIZATION(T, mantisa) \
|
|
template void launch_dequantization<T, mantisa>(uint8_t*, T*, int, int, int, int, cudaStream_t);
|
|
// fp8(E4M3)
|
|
#ifdef BF16_AVAILABLE
|
|
INSTANTIATE_LAUNCH_DEQUANTIZATION(__nv_bfloat16, 7);
|
|
#endif
|
|
INSTANTIATE_LAUNCH_DEQUANTIZATION(__half, 10);
|
|
|
|
template <typename T,
|
|
int mantisa_bits,
|
|
int total_q_bits = 16,
|
|
int q_mantisa_bits = 3,
|
|
int q_exponent_bits = 4>
|
|
__global__ void apply_selective_dequantization(uint8_t* val,
|
|
T* q_val,
|
|
int32_t* indexes,
|
|
int group_size,
|
|
int total_num_elements)
|
|
{
|
|
int index = indexes[blockIdx.x];
|
|
constexpr uint32_t vector_size = quantization::access_granularity / sizeof(T);
|
|
int tidx = (blockIdx.y * blockDim.x + threadIdx.x) * vector_size;
|
|
int input_index = index * total_num_elements + tidx;
|
|
constexpr int quantized_bits = q_mantisa_bits + q_exponent_bits + 1;
|
|
constexpr int q_exponent_bits = total_q_bits - mantisa_bits - 1;
|
|
constexpr uint16_t _mantisa_mask = (1 << q_mantisa_bits) - 1;
|
|
constexpr uint16_t _exponent_mask = ((1 << q_exponent_bits) - 1) << q_mantisa_bits;
|
|
constexpr uint16_t _sign_mask = 1U << (q_mantisa_bits + q_exponent_bits);
|
|
const uint32_t g_index = (input_index / group_size);
|
|
const uint32_t group_size_bytes = (group_size * quantized_bits / 8);
|
|
const uint8_t* load_base_ptr =
|
|
val + g_index * (group_size_bytes + 4) + (input_index % group_size) * quantized_bits / 8;
|
|
|
|
T* store_base_ptr = q_val + tidx + blockIdx.x * total_num_elements;
|
|
float scale;
|
|
|
|
uint8_t* scale_as_int8 = reinterpret_cast<uint8_t*>(&scale);
|
|
if (quantized_bits == 6) {
|
|
mem_access::load_global<quantization::quanitzed_access_granularity>(
|
|
scale_as_int8, val + g_index * (group_size_bytes + 4) + group_size_bytes);
|
|
mem_access::load_global<quantization::quanitzed_access_granularity_6bits>(
|
|
scale_as_int8 + quantization::quanitzed_access_granularity_6bits,
|
|
val + g_index * (group_size_bytes + 4) + group_size_bytes +
|
|
quantization::quanitzed_access_granularity_6bits);
|
|
} else
|
|
mem_access::load_global<quantization::quanitzed_access_granularity>(
|
|
scale_as_int8, val + g_index * (group_size_bytes + 4) + group_size_bytes);
|
|
|
|
if (tidx < total_num_elements) {
|
|
uint64_t q_buf_in;
|
|
uint64_t q_buf_in1;
|
|
uint8_t* int8_data = reinterpret_cast<uint8_t*>(&q_buf_in);
|
|
uint8_t* int8_data1 = reinterpret_cast<uint8_t*>(&q_buf_in1);
|
|
if (quantized_bits == 6) {
|
|
mem_access::load_global<quantization::quanitzed_access_granularity_6bits>(
|
|
int8_data, load_base_ptr);
|
|
mem_access::load_global<quantization::quanitzed_access_granularity_6bits>(
|
|
int8_data + quantization::quanitzed_access_granularity_6bits,
|
|
load_base_ptr + quantization::quanitzed_access_granularity_6bits);
|
|
mem_access::load_global<quantization::quanitzed_access_granularity_6bits>(
|
|
int8_data + quantization::quanitzed_access_granularity_6bits * 2,
|
|
load_base_ptr + quantization::quanitzed_access_granularity_6bits * 2);
|
|
} else {
|
|
mem_access::load_global<quantization::quanitzed_access_granularity>(int8_data,
|
|
load_base_ptr);
|
|
if (quantized_bits > 4) {
|
|
mem_access::load_global<quantization::quanitzed_access_granularity>(
|
|
int8_data + quantization::quanitzed_access_granularity,
|
|
load_base_ptr + quantization::quanitzed_access_granularity);
|
|
if (quantized_bits == 12) {
|
|
mem_access::load_global<quantization::quanitzed_access_granularity>(
|
|
int8_data1, load_base_ptr + quantization::quanitzed_access_granularity * 2);
|
|
}
|
|
}
|
|
}
|
|
T store_buf[vector_size];
|
|
uint16_t* q_buf = reinterpret_cast<uint16_t*>(store_buf);
|
|
#pragma unroll
|
|
for (int j = 0; j < vector_size; j++) {
|
|
uint16_t new_data;
|
|
if (j < 5 || quantized_bits != 12) {
|
|
new_data = (uint16_t)(q_buf_in >> (j * quantized_bits));
|
|
} else {
|
|
if (j == 5) {
|
|
new_data = (uint16_t)(q_buf_in1);
|
|
new_data = (uint16_t)((new_data << 4) | (q_buf_in >> 60));
|
|
} else
|
|
new_data = (uint16_t)(q_buf_in1 >> ((j - 6) * quantized_bits + 8));
|
|
}
|
|
|
|
uint16_t sign = (new_data & _sign_mask) >> (q_mantisa_bits + q_exponent_bits);
|
|
uint16_t dst_exponent = (new_data & _exponent_mask) >> q_mantisa_bits;
|
|
uint16_t dst_mantisa = (new_data & _mantisa_mask);
|
|
|
|
if (dst_exponent != (1 << q_exponent_bits) - 1)
|
|
dst_exponent = (dst_exponent - ((1 << (q_exponent_bits - 1)) - 1)) +
|
|
(1 << (q_exponent_bits - 1)) - 1;
|
|
|
|
q_buf[j] =
|
|
((sign << (q_exponent_bits + mantisa_bits)) | (dst_exponent << mantisa_bits) |
|
|
(dst_mantisa << (mantisa_bits - q_mantisa_bits)));
|
|
float up_cast = conversion::to<float>(store_buf[j]);
|
|
store_buf[j] = conversion::to<T>(up_cast * scale);
|
|
}
|
|
mem_access::store_global<quantization::access_granularity>(store_base_ptr, store_buf);
|
|
}
|
|
}
|
|
|
|
template <typename T, int mantisa>
|
|
void launch_selective_dequantization(uint8_t* val,
|
|
T* q_val,
|
|
int32_t* indexes,
|
|
int num_groups,
|
|
int group_size,
|
|
int num_indexes,
|
|
int q_mantisa_bits,
|
|
int q_exponent_bits,
|
|
cudaStream_t stream)
|
|
{
|
|
int total_elements_per_index = (num_groups / num_indexes) * group_size;
|
|
int blocks = (total_elements_per_index - 1) /
|
|
(quantization::threads * (quantization::access_granularity / sizeof(T))) +
|
|
1;
|
|
const dim3 grid(num_indexes, blocks);
|
|
const dim3 block(quantization::threads);
|
|
DEQUANT_SWITCH(q_mantisa_bits * q_exponent_bits, [&] {
|
|
apply_selective_dequantization<T, mantisa, 16, CONST_Q_MANTISA_BITS, CONST_Q_EXPONENT_BITS>
|
|
<<<grid, block, 0, stream>>>(val, q_val, indexes, group_size, total_elements_per_index);
|
|
});
|
|
}
|
|
#define INSTANTIATE_LAUNCH_SELECTIVE_DEQUANTIZATION(T, mantisa) \
|
|
template void launch_selective_dequantization<T, mantisa>( \
|
|
uint8_t*, T*, int32_t*, int, int, int, int, int, cudaStream_t);
|
|
// fp8(E4M3)
|
|
#ifdef BF16_AVAILABLE
|
|
INSTANTIATE_LAUNCH_SELECTIVE_DEQUANTIZATION(__nv_bfloat16, 7);
|
|
#endif
|
|
INSTANTIATE_LAUNCH_SELECTIVE_DEQUANTIZATION(__half, 10);
|