chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,469 @@
|
||||
// Copyright (c) Microsoft Corporation.
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
// DeepSpeed Team
|
||||
|
||||
#include <cassert>
|
||||
#include "conversion_utils.h"
|
||||
#include "ds_kernel_utils.h"
|
||||
#include "memory_access_utils.h"
|
||||
#include "quantization.h"
|
||||
#include "reduction_utils.h"
|
||||
|
||||
#pragma once
|
||||
|
||||
using rop = reduce::ROpType;
|
||||
|
||||
namespace quantize {
|
||||
constexpr int granularity = 16;
|
||||
constexpr int h_per_load = granularity / sizeof(__half);
|
||||
constexpr int h2_per_load = granularity / sizeof(__half2);
|
||||
constexpr int max_threads = 1024;
|
||||
|
||||
/*
|
||||
Class to hold the quantization parameters for a given tensor.
|
||||
Holds the implementation of the quantization operation.
|
||||
*/
|
||||
|
||||
template <Type qType, int numBits>
|
||||
class Params {
|
||||
public:
|
||||
/*
|
||||
Quantization implementation, supports
|
||||
1) 4 Bit
|
||||
2) 8 Bit
|
||||
3) Symmetric
|
||||
4) Asymmetric
|
||||
Function Arguments :
|
||||
val : The __half value to quantize.
|
||||
*/
|
||||
DS_D_INLINE int8_t quantize(__half val);
|
||||
|
||||
template <typename T>
|
||||
DS_D_INLINE T dequantize(int8_t val);
|
||||
|
||||
DS_D_INLINE void store(float* params, int group_index);
|
||||
|
||||
// Initialize from memory
|
||||
DS_D_INLINE Params(const float* params, int group_index);
|
||||
};
|
||||
|
||||
template <int numBits>
|
||||
class Params<Type::Symmetric, numBits> {
|
||||
public:
|
||||
float scale;
|
||||
|
||||
DS_D_INLINE Params(float max)
|
||||
{
|
||||
if (max == 0) {
|
||||
scale = 1.0;
|
||||
} else {
|
||||
scale = (1 << numBits) / (2 * max);
|
||||
}
|
||||
}
|
||||
|
||||
DS_D_INLINE int8_t quantize(__half val)
|
||||
{
|
||||
constexpr int32_t q_min = -(1 << (numBits - 1));
|
||||
constexpr int32_t q_max = (1 << (numBits - 1)) - 1;
|
||||
|
||||
float val_f = conversion::to<float>(val) * scale;
|
||||
int32_t data_i32 = conversion::to<int32_t>(val_f);
|
||||
data_i32 = min(max(data_i32, q_min), q_max);
|
||||
return (int8_t)data_i32;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
DS_D_INLINE T dequantize(int8_t val)
|
||||
{
|
||||
const float val_deq_f = conversion::to<float>(val) * scale;
|
||||
return conversion::to<T>(val_deq_f);
|
||||
}
|
||||
|
||||
DS_D_INLINE void store(float* params, int group_index)
|
||||
{
|
||||
const float store_scale = 1 / scale;
|
||||
mem_access::store_global<sizeof(float)>(params + group_index, &store_scale);
|
||||
}
|
||||
|
||||
DS_D_INLINE Params(const float* params, int group_index)
|
||||
{
|
||||
mem_access::load_global<sizeof(float)>(&scale, params + group_index);
|
||||
}
|
||||
};
|
||||
|
||||
template <int numBits>
|
||||
class Params<Type::Asymmetric, numBits> {
|
||||
public:
|
||||
float scale;
|
||||
float offset;
|
||||
|
||||
DS_D_INLINE Params(float max, float min)
|
||||
{
|
||||
if (max == min) {
|
||||
scale = 1.0;
|
||||
} else {
|
||||
scale = ((1 << numBits)) / (max - min);
|
||||
}
|
||||
offset = (max + min) / 2;
|
||||
}
|
||||
|
||||
DS_D_INLINE int8_t quantize(__half val)
|
||||
{
|
||||
constexpr int32_t q_min = -(1 << (numBits - 1));
|
||||
constexpr int32_t q_max = (1 << (numBits - 1)) - 1;
|
||||
|
||||
float val_f = (conversion::to<float>(val) - offset) * scale;
|
||||
int32_t data_i32 = conversion::to<int32_t>(val_f);
|
||||
data_i32 = min(max(data_i32, q_min), q_max);
|
||||
return (int8_t)data_i32;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
DS_D_INLINE T dequantize(int8_t val)
|
||||
{
|
||||
const float val_deq_f = ((conversion::to<float>(val)) * scale) + offset;
|
||||
return conversion::to<T>(val_deq_f);
|
||||
}
|
||||
|
||||
DS_D_INLINE void store(float* params, int group_index)
|
||||
{
|
||||
// Codegen should turn this into stg.64
|
||||
const float store_scale = 1 / scale;
|
||||
mem_access::store_global<sizeof(float)>(params + 2 * group_index, &store_scale);
|
||||
mem_access::store_global<sizeof(float)>(params + 2 * group_index + 1, &offset);
|
||||
}
|
||||
|
||||
DS_D_INLINE Params(const float* params, int group_index)
|
||||
{
|
||||
// Codegen should turn this into ldg.64
|
||||
mem_access::load_global<sizeof(float)>(&scale, params + 2 * group_index);
|
||||
mem_access::load_global<sizeof(float)>(&offset, params + 2 * group_index + 1);
|
||||
}
|
||||
};
|
||||
|
||||
/*
|
||||
Group stats tracks the necessary statistics about the quantized group
|
||||
to abstract the particulars for the main loop.
|
||||
*/
|
||||
template <Type qType>
|
||||
class GroupStats {
|
||||
public:
|
||||
DS_D_INLINE void update(__half2 val);
|
||||
|
||||
DS_D_INLINE void reduce(cg::thread_block& tb, cg::thread_block_tile<hw_warp_size>& warp);
|
||||
};
|
||||
|
||||
template <>
|
||||
class GroupStats<Type::Symmetric> {
|
||||
public:
|
||||
// Symmetric quantization only tracks the maximum absolute value
|
||||
__half2 cur_max;
|
||||
float max;
|
||||
|
||||
/*
|
||||
Technically, this would give bad results if there
|
||||
are 0 values to process since the reduction would
|
||||
give -inf instead of 0. We do not consider this
|
||||
to be a reasonable edge case.
|
||||
*/
|
||||
DS_D_INLINE GroupStats() { cur_max = reduce::init<rop::Max, __half2>(); }
|
||||
|
||||
/*
|
||||
Updated the running absmax used to calculate params.
|
||||
Function Arguments :
|
||||
val : The __half2 value to update the running min and max with.
|
||||
*/
|
||||
DS_D_INLINE void update(__half2 val)
|
||||
{
|
||||
cur_max = reduce::element<rop::Max>(cur_max, __habs2(val));
|
||||
}
|
||||
|
||||
/*
|
||||
Function to return calculated quantization params.
|
||||
Template Arguments :
|
||||
numBits - Number of bits in quantized element. int : 8 or 4
|
||||
Function Arguments :
|
||||
tb - Threadblock object. cg::thread_block
|
||||
warp - Warp object. cg::thread_block_tile<hw_warp_size>
|
||||
*/
|
||||
template <int numBits, int threads_per_group>
|
||||
DS_D_INLINE Params<Type::Symmetric, numBits> get_params(
|
||||
cg::thread_block& tb,
|
||||
cg::thread_block_tile<hw_warp_size>& warp)
|
||||
{
|
||||
const float2 partial_max = conversion::to<float2>(cur_max);
|
||||
float max = reduce::element<rop::Max>(partial_max.x, partial_max.y);
|
||||
|
||||
reduce::partitioned_block<rop::Max, threads_per_group>(tb, warp, max);
|
||||
Params<Type::Symmetric, numBits> params(max);
|
||||
|
||||
return params;
|
||||
}
|
||||
};
|
||||
|
||||
template <>
|
||||
class GroupStats<Type::Asymmetric> {
|
||||
public:
|
||||
__half2 cur_max;
|
||||
__half2 cur_min;
|
||||
|
||||
/*
|
||||
Initialize cur_max to -inf, cur_min to inf since
|
||||
we are doing a true range analysis.
|
||||
*/
|
||||
DS_D_INLINE GroupStats()
|
||||
{
|
||||
cur_max = reduce::init<rop::Max, __half2>();
|
||||
cur_min = reduce::init<rop::Min, __half2>();
|
||||
}
|
||||
|
||||
/*
|
||||
Updated the running min and max used to calculate params.
|
||||
Function Arguments :
|
||||
val : The __half2 value to update the running min and max with.
|
||||
*/
|
||||
DS_D_INLINE void update(__half2 val)
|
||||
{
|
||||
cur_max = reduce::element<rop::Max>(cur_max, val);
|
||||
cur_min = reduce::element<rop::Min>(cur_min, val);
|
||||
}
|
||||
|
||||
/*
|
||||
Function to return calculated quantization params.
|
||||
Template Arguments :
|
||||
numBits - Number of bits in quantized element. int : 8 or 4
|
||||
Function Arguments :
|
||||
tb - Threadblock object. cg::thread_block
|
||||
warp - Warp object. cg::thread_block_tile<hw_warp_size>
|
||||
*/
|
||||
template <int numBits, int threads_per_group>
|
||||
DS_D_INLINE Params<Type::Asymmetric, numBits> get_params(
|
||||
cg::thread_block& tb,
|
||||
cg::thread_block_tile<hw_warp_size>& warp)
|
||||
{
|
||||
const float2 partial_max = conversion::to<float2>(cur_max);
|
||||
float max = reduce::element<rop::Max>(partial_max.x, partial_max.y);
|
||||
|
||||
const float2 partial_min = conversion::to<float2>(cur_min);
|
||||
float min = reduce::element<rop::Min>(partial_min.x, partial_min.y);
|
||||
|
||||
reduce::partitioned_block<rop::Max, rop::Min, threads_per_group>(tb, warp, max, min);
|
||||
|
||||
Params<Type::Asymmetric, numBits> params(max, min);
|
||||
|
||||
return params;
|
||||
}
|
||||
};
|
||||
|
||||
/*
|
||||
Device function that quantizes 16 bytes of __half type input data.
|
||||
Template Arguments :
|
||||
numBits - Number of bits in quantized element. int : 8 or 4
|
||||
qType - Type of quantization to perform. Type::Symmetric or Type::Asymmetric
|
||||
Function Arguments :
|
||||
local_output - Pointer to local memory to store quantized data. int8_t*
|
||||
data - Pointer to input data. __half*
|
||||
Params - Parameters for quantization. Params<qType, numBits>
|
||||
*/
|
||||
template <int numBits, Type qType>
|
||||
DS_D_INLINE void _chunk(int8_t* local_output, const __half* data, Params<qType, numBits> q_params);
|
||||
|
||||
/*
|
||||
Device function that quantizes 16 bytes of __half2 type input data.
|
||||
Template Arguments :
|
||||
numBits - Number of bits in quantized element. int : 8 or 4
|
||||
qType - Type of quantization to perform. Type::Symmetric or Type::Asymmetric
|
||||
Function Arguments :
|
||||
local_output - Pointer to local memory to store quantized data. int8_t*
|
||||
data - Pointer to input data. __half2*
|
||||
Params - Parameters for quantization. Params<qType, numBits>
|
||||
*/
|
||||
template <int numBits, Type qType>
|
||||
DS_D_INLINE void _chunk(int8_t* local_output, const __half2* data, Params<qType, numBits> q_params);
|
||||
|
||||
/*
|
||||
Helper function to do serial reduction on register-file arrays.
|
||||
Template Arguments :
|
||||
qType - Type of quantization to perform. Type::Symmetric or Type::Asymmetric
|
||||
numChunks - Number of bits in quantized element. int : 8 or 4
|
||||
Function Arguments :
|
||||
local_buffer - Pointer memory with input half2 data to be quantized.
|
||||
*/
|
||||
template <Type qType, int numChunks>
|
||||
DS_D_INLINE GroupStats<qType> _local_serial_reduce(__half2* local_buffer);
|
||||
|
||||
/*
|
||||
The main loop of the kernel that quantizes array in local memory of __half2 type input data, when
|
||||
Quantization parameters are pre-computed.
|
||||
Template Arguments :
|
||||
qType - Type of quantization to perform. Type::Symmetric or Type::Asymmetric
|
||||
numBits - Number of bits in quantized element. int : 8 or 4
|
||||
numChunks - Number of chunks(16 bytes of Input data). int : 8 or 4
|
||||
Function Arguments :
|
||||
local_buffer - Pointer memory with input half2 data to be quantized.
|
||||
scales - Pointer to output scales.
|
||||
offsets - Pointer to output offsets.
|
||||
output_data - Pointer to output data.
|
||||
elems_per_group - Number of elements to quantize in a group.
|
||||
q_params - Quantization parameters.
|
||||
*/
|
||||
template <int numBits, Type qType, int numChunks, int threads_per_group, int max_threads>
|
||||
DS_D_INLINE void local_array(cg::thread_block& tb,
|
||||
cg::thread_block_tile<hw_warp_size>& warp,
|
||||
__half2* local_buffer,
|
||||
float* __restrict__ scales,
|
||||
float* __restrict__ offsets,
|
||||
int8_t* __restrict__ output_data,
|
||||
const int& elems_per_group,
|
||||
const int& groups,
|
||||
Params<qType, numBits> q_params);
|
||||
|
||||
/*
|
||||
The main loop of the kernel that quantizes array in local memory of __half2 type input data.
|
||||
This function computes quantization parameters for each group.
|
||||
Template Arguments :
|
||||
qType - Type of quantization to perform. Type::Symmetric or Type::Asymmetric
|
||||
numBits - Number of bits in quantized element. int : 8 or 4
|
||||
numChunks - Number of chunks(16 bytes of Input data). int : 8 or 4
|
||||
Function Arguments :
|
||||
local_buffer - Pointer memory with input half2 data to be quantized.
|
||||
scales - Pointer to output scales.
|
||||
offsets - Pointer to output offsets.
|
||||
output_data - Pointer to output data.
|
||||
elems_per_group - Number of elements to quantize in a group.
|
||||
*/
|
||||
template <Type qType, int numBits, int numChunks, int threads_per_group, int max_threads>
|
||||
__device__ void local_array(__half2* local_buffer,
|
||||
float* __restrict__ scales,
|
||||
float* __restrict__ offsets,
|
||||
int8_t* __restrict__ output_data,
|
||||
const int& elems_per_group,
|
||||
const int& groups);
|
||||
|
||||
template <int numBits, Type qType>
|
||||
DS_D_INLINE void _chunk(int8_t* local_output, const __half* data, Params<qType, numBits> q_params)
|
||||
{
|
||||
constexpr int32_t elems = 16 / sizeof(__half);
|
||||
constexpr int32_t num_elems_packed = 8 / numBits;
|
||||
|
||||
#pragma unroll
|
||||
for (int i = 0, oi = 0; i < elems; i += num_elems_packed, oi++) {
|
||||
if (num_elems_packed == 1) {
|
||||
// TODO(cmikeh2): refactor to use conversion utils
|
||||
local_output[i] = q_params.quantize(data[i]);
|
||||
} else if (num_elems_packed == 2) {
|
||||
int8_t data_i8_1 = q_params.quantize(data[i]);
|
||||
int8_t data_i8_2 = q_params.quantize(data[i + 1]);
|
||||
auto data_i8 = PackedInt4{data_i8_2, data_i8_1};
|
||||
local_output[oi] = *((int8_t*)(&data_i8));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <int numBits, Type qType>
|
||||
DS_D_INLINE void _chunk(int8_t* local_output, const __half2* data, Params<qType, numBits> q_params)
|
||||
{
|
||||
const __half* data_cast = reinterpret_cast<const __half*>(data);
|
||||
_chunk<numBits>(local_output, data_cast, q_params);
|
||||
}
|
||||
|
||||
template <Type qType, int numChunks>
|
||||
DS_D_INLINE GroupStats<qType> _local_serial_reduce(__half2* local_buffer)
|
||||
{
|
||||
GroupStats<qType> stats;
|
||||
#pragma unroll
|
||||
for (int i = 0; i < numChunks * h2_per_load; i++) { stats.update(local_buffer[i]); }
|
||||
|
||||
return stats;
|
||||
}
|
||||
|
||||
template <Type qType, int numBits, int numChunks, int threads_per_group, int max_threads>
|
||||
DS_D_INLINE void local_array(cg::thread_block& tb,
|
||||
cg::thread_block_tile<hw_warp_size>& warp,
|
||||
__half2* local_buffer,
|
||||
float* __restrict__ global_params,
|
||||
int8_t* __restrict__ output_data,
|
||||
const int& elems_per_group,
|
||||
const int& groups,
|
||||
Params<qType, numBits> q_params)
|
||||
{
|
||||
constexpr int num_ele_int8 = 8 / numBits;
|
||||
constexpr int num_int8_out = quantize::h_per_load / num_ele_int8;
|
||||
|
||||
// Indexing offsets
|
||||
const int block_num =
|
||||
(tb.group_index().x * max_threads / threads_per_group) + tb.thread_index().y;
|
||||
const int block_offset = block_num * elems_per_group;
|
||||
const int elem_offset = tb.thread_index().x * quantize::h_per_load;
|
||||
const int base_offset = (block_offset + elem_offset) / num_ele_int8;
|
||||
const int stride = tb.size() * quantize::h_per_load / num_ele_int8;
|
||||
|
||||
int8_t local_output[num_int8_out];
|
||||
|
||||
if (tb.thread_index().x == 0 && block_num < groups) {
|
||||
q_params.store(
|
||||
global_params,
|
||||
(tb.group_index().x * max_threads / threads_per_group) + tb.thread_index().y);
|
||||
}
|
||||
#pragma unroll
|
||||
for (int i = 0; i < numChunks; i++) {
|
||||
if (elem_offset + i * stride * num_ele_int8 < elems_per_group && block_num < groups) {
|
||||
quantize::_chunk<numBits, qType>(
|
||||
local_output, local_buffer + i * quantize::h2_per_load, q_params);
|
||||
mem_access::store_global<num_int8_out>(output_data + (base_offset + i * stride),
|
||||
local_output);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <Type qType, int numBits, int numChunks, int threads_per_group, int max_threads>
|
||||
DS_D_INLINE void local_array(cg::thread_block& tb,
|
||||
cg::thread_block_tile<hw_warp_size>& warp,
|
||||
__half* local_buffer,
|
||||
float* __restrict__ global_params,
|
||||
int8_t* __restrict__ output_data,
|
||||
const int& elems_per_group,
|
||||
const int& groups,
|
||||
Params<qType, numBits> q_params)
|
||||
{
|
||||
__half2* local_buffer_h2 = reinterpret_cast<__half2*>(local_buffer);
|
||||
|
||||
quantize::local_array<qType, numBits, numChunks, threads_per_group, max_threads>(
|
||||
tb, warp, local_buffer, global_params, output_data, elems_per_group, groups, q_params);
|
||||
}
|
||||
|
||||
template <Type qType,
|
||||
int numBits,
|
||||
int numChunks,
|
||||
int threads_per_group = max_threads,
|
||||
int max_threads = 256>
|
||||
__device__ void local_array(__half2* local_buffer,
|
||||
float* __restrict__ global_params,
|
||||
int8_t* __restrict__ output_data,
|
||||
const int& elems_per_group,
|
||||
const int& groups)
|
||||
{
|
||||
cg::thread_block tb = cg::this_thread_block();
|
||||
cg::thread_block_tile<hw_warp_size> warp = cg::tiled_partition<hw_warp_size>(tb);
|
||||
|
||||
auto group_stats = _local_serial_reduce<qType, numChunks>(local_buffer);
|
||||
auto params = group_stats.template get_params<numBits, threads_per_group>(tb, warp);
|
||||
|
||||
quantize::local_array<qType, numBits, numChunks, threads_per_group, max_threads>(
|
||||
tb, warp, local_buffer, global_params, output_data, elems_per_group, groups, params);
|
||||
}
|
||||
|
||||
template <Type qType, int numBits, int numChunks, int threads_per_group, int max_threads>
|
||||
__device__ void local_array(__half* local_buffer,
|
||||
float* __restrict__ global_params,
|
||||
int8_t* __restrict__ output_data,
|
||||
const int& elems_per_group,
|
||||
const int& groups)
|
||||
{
|
||||
__half2* local_buffer_h2 = reinterpret_cast<__half2*>(local_buffer);
|
||||
quantize::local_array<qType, numBits, numChunks, threads_per_group, max_threads>(
|
||||
local_buffer_h2, global_params, output_data, elems_per_group, groups);
|
||||
}
|
||||
|
||||
} // namespace quantize
|
||||
Reference in New Issue
Block a user