1428 lines
52 KiB
Plaintext
1428 lines
52 KiB
Plaintext
/* Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
|
|
|
|
Licensed under the Apache License, Version 2.0 (the "License");
|
|
you may not use this file except in compliance with the License.
|
|
You may obtain a copy of the License at
|
|
|
|
http://www.apache.org/licenses/LICENSE-2.0
|
|
|
|
Unless required by applicable law or agreed to in writing, software
|
|
distributed under the License is distributed on an "AS IS" BASIS,
|
|
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
See the License for the specific language governing permissions and
|
|
limitations under the License. */
|
|
|
|
#include "paddle/phi/kernels/funcs/weight_only_gemv.h"
|
|
|
|
#include <assert.h>
|
|
#include <stdint.h>
|
|
#include <cmath>
|
|
#include "paddle/phi/backends/gpu/gpu_context.h"
|
|
#include "paddle/phi/common/amp_type_traits.h"
|
|
#include "paddle/phi/common/datatype_traits.h"
|
|
#include "paddle/phi/core/kernel_registry.h"
|
|
|
|
namespace phi {
|
|
|
|
namespace {
|
|
|
|
#ifdef PADDLE_WITH_CUDA
|
|
constexpr int kWarpSize = 32;
|
|
constexpr int kPerBlockWarpNum = 8;
|
|
|
|
/////////////////////////////////////////////////////////////////////
|
|
|
|
template <typename T>
|
|
struct CUDA_HALF_2_TYPE_TARIS {};
|
|
|
|
template <>
|
|
struct CUDA_HALF_2_TYPE_TARIS<half> {
|
|
using type = half2;
|
|
};
|
|
|
|
#ifdef PADDLE_CUDA_BF16
|
|
template <>
|
|
struct CUDA_HALF_2_TYPE_TARIS<__nv_bfloat16> {
|
|
using type = __nv_bfloat162;
|
|
};
|
|
#endif
|
|
|
|
template <typename T>
|
|
__device__ inline void fast_cvt_4_packed_signed_i8s_to_2_half2s(
|
|
T halves[4], int8_t signed_chars[4]) {
|
|
assert(false);
|
|
}
|
|
|
|
// Specialization for fast cast from FP16 -> int8
|
|
template <>
|
|
__device__ inline void fast_cvt_4_packed_signed_i8s_to_2_half2s(
|
|
half halves[4], int8_t signed_chars[4]) {
|
|
#if CUDA_ARCH_FP16_SUPPORTED(__CUDA_ARCH__)
|
|
uint32_t* h = reinterpret_cast<uint32_t*>(halves);
|
|
uint32_t i8s = *reinterpret_cast<uint32_t*>(signed_chars);
|
|
|
|
static constexpr uint32_t mask_for_elt_01 = 0x5150;
|
|
static constexpr uint32_t mask_for_elt_23 = 0x5352;
|
|
static constexpr uint32_t start_byte_for_fp16 = 0x64646464;
|
|
asm volatile("prmt.b32 %0,%1,%2,%3;\n"
|
|
: "=r"(h[0])
|
|
: "r"(i8s), "n"(start_byte_for_fp16), "n"(mask_for_elt_01));
|
|
asm volatile("prmt.b32 %0,%1,%2,%3;\n"
|
|
: "=r"(h[1])
|
|
: "r"(i8s), "n"(start_byte_for_fp16), "n"(mask_for_elt_23));
|
|
|
|
static constexpr uint32_t I8s_TO_F16s_MAGIC_NUM = 0x64806480;
|
|
asm volatile("sub.f16x2 %0, %1, %2;\n"
|
|
: "=r"(h[0])
|
|
: "r"(h[0]), "r"(I8s_TO_F16s_MAGIC_NUM));
|
|
asm volatile("sub.f16x2 %0, %1, %2;\n"
|
|
: "=r"(h[1])
|
|
: "r"(h[1]), "r"(I8s_TO_F16s_MAGIC_NUM));
|
|
#endif
|
|
}
|
|
|
|
// Specialization for fast cast from BF16 -> int8
|
|
#ifdef PADDLE_CUDA_BF16
|
|
template <>
|
|
__device__ inline void fast_cvt_4_packed_signed_i8s_to_2_half2s(
|
|
__nv_bfloat16 halves[4], int8_t signed_chars[4]) {
|
|
#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800))
|
|
uint32_t* bf16_result_ptr = reinterpret_cast<uint32_t*>(halves);
|
|
uint32_t i8s = *reinterpret_cast<uint32_t*>(signed_chars);
|
|
|
|
static constexpr uint32_t fp32_base = 0x4B000000;
|
|
float fp32_intermediates[4];
|
|
|
|
// Construct FP32s, bfloat does not have enough mantissa for IADD trick
|
|
uint32_t* fp32_intermediates_casted =
|
|
reinterpret_cast<uint32_t*>(fp32_intermediates);
|
|
fp32_intermediates_casted[0] = __byte_perm(i8s, fp32_base, 0x7650);
|
|
fp32_intermediates_casted[1] = __byte_perm(i8s, fp32_base, 0x7651);
|
|
fp32_intermediates_casted[2] = __byte_perm(i8s, fp32_base, 0x7652);
|
|
fp32_intermediates_casted[3] = __byte_perm(i8s, fp32_base, 0x7653);
|
|
|
|
// Subtract out fp32_base + 128 to make the unsigned integer signed.
|
|
#pragma unroll
|
|
for (int ii = 0; ii < 4; ++ii) {
|
|
fp32_intermediates[ii] -= 8388736.f;
|
|
}
|
|
|
|
// Truncate the fp32 representation and pack up as bfloat16s.
|
|
#pragma unroll
|
|
for (int ii = 0; ii < 2; ++ii) {
|
|
bf16_result_ptr[ii] = __byte_perm(fp32_intermediates_casted[2 * ii + 0],
|
|
fp32_intermediates_casted[2 * ii + 1],
|
|
0x7632);
|
|
}
|
|
#else
|
|
// Disable this on architectures older than Ampere since they lack hardware
|
|
// for bf16 mma. If one wishes to use HMMA on older hardware, they should
|
|
// Convert directly to FP16 using FP16 converters.
|
|
assert(false);
|
|
#endif
|
|
}
|
|
#endif
|
|
|
|
/* Gelu Activation */
|
|
|
|
__forceinline__ __device__ float copysignf_pos(float a, float b) {
|
|
float r;
|
|
r = __int_as_float(__float_as_int(a) | (__float_as_int(b) & 0x80000000));
|
|
return r;
|
|
}
|
|
|
|
__inline__ __device__ float tanh_opt(float x) {
|
|
#if (__CUDA_ARCH__ >= 750 && CUDART_VERSION >= 11000)
|
|
float r;
|
|
asm("tanh.approx.f32 %0,%1; \n\t" : "=f"(r) : "f"(x));
|
|
return r;
|
|
#else
|
|
const float exp_val = -1.f * fabs(2 * x);
|
|
return copysignf_pos((1.0f - __expf(exp_val)) / (__expf(exp_val) + 1.0f), x);
|
|
#endif
|
|
}
|
|
|
|
template <typename T, bool EnableFastGelu>
|
|
struct GeluActivation {
|
|
using return_type = T;
|
|
static __device__ __forceinline__ T apply(const T& val) {
|
|
if (!EnableFastGelu) return val;
|
|
const float cdf =
|
|
0.5f * (1.0f + tanh_opt((0.7978845608028654f *
|
|
(val + 0.044715f * val * val * val))));
|
|
return val * cdf;
|
|
}
|
|
};
|
|
|
|
template <typename T>
|
|
struct ConvertFloatFunc {
|
|
ConvertFloatFunc() {}
|
|
static __device__ __forceinline__ float apply(const T& val) {
|
|
assert(false);
|
|
return 0.0f;
|
|
}
|
|
};
|
|
|
|
template <>
|
|
struct ConvertFloatFunc<half> {
|
|
static __device__ __forceinline__ float apply(const half& val) {
|
|
return __half2float(val);
|
|
}
|
|
};
|
|
|
|
#ifdef PADDLE_CUDA_BF16
|
|
template <>
|
|
struct ConvertFloatFunc<__nv_bfloat16> {
|
|
static __device__ __forceinline__ float apply(const __nv_bfloat16& val) {
|
|
return __bfloat162float(val);
|
|
}
|
|
};
|
|
#endif
|
|
|
|
template <typename T>
|
|
struct ConvertDstFunc_2 {
|
|
static __device__ __forceinline__ T apply(const float& val) { assert(false); }
|
|
};
|
|
|
|
template <typename T>
|
|
struct ConvertDstFunc {
|
|
static __device__ __forceinline__ T apply(const float& val) { assert(false); }
|
|
};
|
|
|
|
template <>
|
|
struct ConvertDstFunc<half> {
|
|
static __device__ __forceinline__ half apply(const float& val) {
|
|
return __float2half_rn(val);
|
|
}
|
|
};
|
|
|
|
template <>
|
|
struct ConvertDstFunc<half2> {
|
|
static __device__ __forceinline__ half2 apply(const float& val) {
|
|
return __float2half2_rn(val);
|
|
}
|
|
};
|
|
|
|
template <>
|
|
struct ConvertDstFunc_2<half2> {
|
|
static __device__ __forceinline__ half2 apply(const half& val) {
|
|
return __half2half2(val);
|
|
}
|
|
};
|
|
#ifdef PADDLE_CUDA_BF16
|
|
template <>
|
|
struct ConvertDstFunc<__nv_bfloat16> {
|
|
static __device__ __forceinline__ __nv_bfloat16 apply(const float& val) {
|
|
return __float2bfloat16_rn(val);
|
|
}
|
|
};
|
|
|
|
template <>
|
|
struct ConvertDstFunc<__nv_bfloat162> {
|
|
static __device__ __forceinline__ __nv_bfloat162 apply(const float& val) {
|
|
return __float2bfloat162_rn(val);
|
|
}
|
|
};
|
|
|
|
template <>
|
|
struct ConvertDstFunc_2<__nv_bfloat162> {
|
|
static __device__ __forceinline__ __nv_bfloat162
|
|
apply(const __nv_bfloat16& val) {
|
|
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800
|
|
__nv_bfloat162 val2;
|
|
val2.x = val;
|
|
val2.y = val;
|
|
return val2;
|
|
#else
|
|
return __bfloat162bfloat162(val);
|
|
#endif
|
|
}
|
|
};
|
|
#endif
|
|
|
|
template <typename T>
|
|
struct HalfMul {
|
|
static __device__ __forceinline__ T apply(const T& x, const T& y) {
|
|
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 530
|
|
return __hmul(x, y);
|
|
#else
|
|
float res = static_cast<float>(float16(x)) * static_cast<float>(float16(y));
|
|
return float16(res).to_half();
|
|
#endif
|
|
}
|
|
};
|
|
|
|
template <typename T>
|
|
struct HalfMulAdd {
|
|
static __device__ __forceinline__ T apply(const T& x,
|
|
const T& y,
|
|
const T& z) {
|
|
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 530
|
|
return __hfma2(x, y, z);
|
|
#else
|
|
assert(0 && "HalfMulAdd cuda version error");
|
|
#endif
|
|
}
|
|
};
|
|
|
|
#ifdef PADDLE_CUDA_BF16
|
|
template <>
|
|
struct HalfMul<__nv_bfloat16> {
|
|
static __device__ __forceinline__ __nv_bfloat16
|
|
apply(const __nv_bfloat16& x, const __nv_bfloat16& y) {
|
|
#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800))
|
|
return __hmul(x, y);
|
|
#else
|
|
return __float2bfloat16_rn(__bfloat162float(x) * __bfloat162float(y));
|
|
#endif
|
|
}
|
|
};
|
|
|
|
template <>
|
|
struct HalfMulAdd<__nv_bfloat162> {
|
|
static __device__ __forceinline__ __nv_bfloat162
|
|
apply(const __nv_bfloat162& x,
|
|
const __nv_bfloat162& y,
|
|
const __nv_bfloat162& z) {
|
|
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
|
|
return __hfma2(x, y, z);
|
|
#else
|
|
assert(0 && "HalfMulAdd cuda version error");
|
|
#endif
|
|
}
|
|
};
|
|
#endif
|
|
|
|
/*
|
|
Int8 Weightonly GEMV.
|
|
X: 1 x k
|
|
Weight(ColMajor): n x k
|
|
Each Warp Process: 1 x k matmul 1 x k
|
|
*/
|
|
template <typename T, bool Bias, bool Gelu>
|
|
__global__ void int8_weight_only_gemv(const T* input,
|
|
const int8_t* weight,
|
|
const T* scale_list,
|
|
const T* bias,
|
|
T* output,
|
|
const int k,
|
|
const int n) {
|
|
constexpr int kWarpSize = 32;
|
|
constexpr int kVecSize = 16;
|
|
T vec_input[kVecSize];
|
|
int8_t vec_weight[kVecSize];
|
|
T vec_weight_f16[kVecSize];
|
|
|
|
const int warp_id = threadIdx.x / kWarpSize;
|
|
const int lane_id = threadIdx.x % kWarpSize;
|
|
const int tile_id = blockIdx.x * blockDim.x / kWarpSize + warp_id;
|
|
const int row_id = tile_id * 2 + ((lane_id % 8) > 3 ? 1 : 0);
|
|
weight += tile_id * k * 2;
|
|
|
|
float v = 0.f, scale = static_cast<float>(scale_list[row_id]), v_bias;
|
|
|
|
if (Bias) {
|
|
v_bias = ConvertFloatFunc<T>::apply(bias[row_id]);
|
|
}
|
|
|
|
#pragma unroll
|
|
for (int i = lane_id * kVecSize; i < k * 2; i += kVecSize * kWarpSize) {
|
|
*reinterpret_cast<int4*>(vec_weight) =
|
|
*reinterpret_cast<const int4*>(weight + i); // NOLINT
|
|
*reinterpret_cast<float4*>(vec_input) = // NOLINT
|
|
*reinterpret_cast<const float4*>(input + i / 128 * 64 +
|
|
(i % 64)); // NOLINT
|
|
*reinterpret_cast<float4*>(vec_input + 8) = // NOLINT
|
|
*reinterpret_cast<const float4*>(input + i / 128 * 64 + (i % 64) +
|
|
8); // NOLINT
|
|
|
|
#pragma unroll
|
|
for (int p = 0; p < kVecSize; p += 4) {
|
|
fast_cvt_4_packed_signed_i8s_to_2_half2s<T>(vec_weight_f16 + p,
|
|
vec_weight + p);
|
|
}
|
|
#pragma unroll
|
|
for (int p = 0; p < kVecSize; ++p) {
|
|
v += ConvertFloatFunc<T>::apply(
|
|
HalfMul<T>::apply(vec_input[p], vec_weight_f16[p / 8 + (p % 8) * 2]));
|
|
}
|
|
}
|
|
// Do WarpReduceSum.
|
|
v += __shfl_xor_sync(0xffffffff, v, 16);
|
|
v += __shfl_xor_sync(0xffffffff, v, 8);
|
|
v += __shfl_xor_sync(0xffffffff, v, 2);
|
|
v += __shfl_xor_sync(0xffffffff, v, 1);
|
|
if (lane_id == 0 || lane_id == 4) {
|
|
if (Bias) {
|
|
output[row_id] = ConvertDstFunc<T>::apply(
|
|
GeluActivation<float, Gelu>::apply(v * scale + v_bias));
|
|
} else {
|
|
output[row_id] = ConvertDstFunc<T>::apply(
|
|
GeluActivation<float, Gelu>::apply(v * scale));
|
|
}
|
|
}
|
|
}
|
|
|
|
enum class WeightOnlyQuantType { Int4b, Int8b };
|
|
|
|
enum class WeightOnlyType { PerChannel, GroupWise };
|
|
|
|
template <WeightOnlyQuantType QType>
|
|
struct WeightLayoutDetails;
|
|
|
|
template <>
|
|
struct WeightLayoutDetails<WeightOnlyQuantType::Int4b> {
|
|
// Every four rows of the original weights are interleaved into a row with
|
|
// stride of 64, so if each thread processes 32 elements(for int4, we can use
|
|
// ldg.128 to load weights), then every group of two adjacent threads will
|
|
// alternately process four different row weights for example every 256
|
|
// consecutive int4 elements [256*i, 256*(i+1)-1] of row N under interleave
|
|
// layout, the first 64 are from [64*i, 64*(i+1)-1] of row 4N before
|
|
// interleaving, and the second 64 are from [64*i, 64*(i+1)-1] of row 4N+1
|
|
// before interleaving, and so on. So if each thread loads 32 int4 elements,
|
|
// then the elements of each 2 adjacent threads of each 8 consecutive threads
|
|
// will come from row 4N ~ 4N+3 respectively before interleaving.
|
|
static constexpr int kElemBits = 4;
|
|
static constexpr int kInterleave = 4;
|
|
static constexpr int kStride = 64;
|
|
|
|
// The index remapping here is to counteracts the effect of
|
|
// cutlass::permute_B_rows_for_mixed_gemm input 0 1 2 3 4 5 6 7 8 9 10 11 12
|
|
// 13 14 15 ... 31 weight 0 1 8 9 16 17 24 25 2 3 10 11 18 19 26 27 4 5 12 13
|
|
// 20 21 28 29 6 7 14 15 22 23 30 31
|
|
static constexpr int kShuffleSize = 32;
|
|
static constexpr int kShuffleBasicTile = 2;
|
|
static constexpr int kShuffleContinuous = 4;
|
|
static constexpr int kShuffleStrided = 4;
|
|
|
|
// The rearrangement here counteracts the effect of
|
|
// cutlass::add_bias_and_interleave_int4s_inplace Input int8 data layout
|
|
// [elt_7 elt_5 elt_3 elt_1 elt_6 elt_4 elt_2 elt_0] (each elt
|
|
// occupies 4 bits)
|
|
//
|
|
// Converted fp16 data layout
|
|
// [elt_7 elt_6 elt_5 elt_4 elt_3 elt_2 elt_1 elt_0] (each elt
|
|
// occupies 16 bits)
|
|
static constexpr int kConvertCount = 8;
|
|
// using Converter
|
|
// =
|
|
// cutlass::FastInterleavedAndBiasedNumericArrayConverter<cutlass::half_t,
|
|
// cutlass::uint4b_t, kConvertCount>;
|
|
|
|
// Each warp completes the internal reduce and writes the [Batch * NPerBlock *
|
|
// Interleave] results to the corresponding address in shared memory
|
|
template <int Num, int WarpSize>
|
|
__device__ __forceinline__ static void sync(float* res,
|
|
float (*sm)[Num * kInterleave]) {
|
|
#pragma unroll
|
|
for (int i = 0; i < Num; ++i) {
|
|
res[i] += __shfl_xor_sync(~0, res[i], 16);
|
|
res[i] += __shfl_xor_sync(~0, res[i], 8);
|
|
res[i] += __shfl_xor_sync(~0, res[i], 1);
|
|
}
|
|
__syncthreads();
|
|
int warp = threadIdx.x / WarpSize, lane = threadIdx.x % WarpSize;
|
|
if (lane == 0 || lane == 2 || lane == 4 || lane == 6) {
|
|
#pragma unroll
|
|
for (int i = 0; i < Num; ++i) {
|
|
sm[warp][i * kInterleave + lane / 2] = res[i];
|
|
}
|
|
}
|
|
__syncthreads();
|
|
}
|
|
};
|
|
|
|
template <>
|
|
struct WeightLayoutDetails<WeightOnlyQuantType::Int8b> {
|
|
// Every two rows of the original weights are interleaved into a row with
|
|
// stride of 64, so if each thread processes 16 elements(for int8, we can use
|
|
// ldg.128 to load weights), then every group of four adjacent threads will
|
|
// alternately process two different row weights for example every 128
|
|
// consecutive int8 elements [128*i, 128*(i+1)-1] of row N under interleave
|
|
// layout, the first 64 are from [64*i, 64*(i+1)-1] of row 2N before
|
|
// interleaving, and the last 64 are from [64*i, 64*(i+1)-1] of row 2N+1
|
|
// before interleaving. So if each thread loads 16 int8 elements, then the
|
|
// elements of the first four and last four threads of each 8 consecutive
|
|
// threads will come from row 2N and row 2N+1 respectively before
|
|
// interleaving.
|
|
static constexpr int kElemBits = 8;
|
|
static constexpr int kInterleave = 2;
|
|
static constexpr int kStride = 64;
|
|
|
|
// The index remapping here is to counteracts the effect of
|
|
// cutlass::permute_B_rows_for_mixed_gemm input 0 1 2 3 4 5 6 7 8 9 10 11 12
|
|
// 13 14 15 weight 0 1 8 9 2 3 10 11 4 5 12 13 6 7 14 15
|
|
static constexpr int kShuffleSize = 16;
|
|
static constexpr int kShuffleBasicTile = 2;
|
|
static constexpr int kShuffleContinuous = 2;
|
|
static constexpr int kShuffleStrided = 4;
|
|
|
|
// The rearrangement here counteracts the effect of
|
|
// cutlass::add_bias_and_interleave_int8s_inplace Input int8 data layout
|
|
// [elt_3 elt_1 elt_2 elt_0] (each elt occupies 8 bits)
|
|
//
|
|
// Converted fp16 data layout
|
|
// [elt_3 elt_2 elt_1 elt_0] (each elt occupies 16 bits)
|
|
static constexpr int kConvertCount = 4;
|
|
// using Converter =
|
|
// cutlass::FastInterleavedAndBiasedNumericArrayConverter<cutlass::half_t,
|
|
// uint8_t, kConvertCount>;
|
|
|
|
// Each warp completes the internal reduce and writes the [Batch * NPerBlock *
|
|
// Interleave] results to the corresponding address in shared memory
|
|
template <int Num, int WarpSize>
|
|
__device__ __forceinline__ static void sync(float* res,
|
|
float (*sm)[Num * kInterleave]) {
|
|
#pragma unroll
|
|
for (int i = 0; i < Num; ++i) {
|
|
res[i] += __shfl_xor_sync(~0, res[i], 16);
|
|
res[i] += __shfl_xor_sync(~0, res[i], 8);
|
|
res[i] += __shfl_xor_sync(~0, res[i], 2);
|
|
res[i] += __shfl_xor_sync(~0, res[i], 1);
|
|
}
|
|
__syncthreads();
|
|
int warp = threadIdx.x / WarpSize, lane = threadIdx.x % WarpSize;
|
|
if (lane == 0 || lane == 4) {
|
|
#pragma unroll
|
|
for (int i = 0; i < Num; ++i) {
|
|
sm[warp][i * kInterleave + lane / 4] = res[i];
|
|
}
|
|
}
|
|
__syncthreads();
|
|
}
|
|
};
|
|
|
|
template <WeightOnlyQuantType QType>
|
|
struct WeightOnlyKernelDetails {
|
|
using Layout = WeightLayoutDetails<QType>;
|
|
|
|
static constexpr int kElemBits = Layout::kElemBits;
|
|
static constexpr int kInterleave = Layout::kInterleave;
|
|
static constexpr int kStride = Layout::kStride;
|
|
|
|
static constexpr int kShuffleSize = Layout::kShuffleSize;
|
|
static constexpr int kShuffleBasicTile = Layout::kShuffleBasicTile;
|
|
static constexpr int kShuffleContinuous = Layout::kShuffleContinuous;
|
|
static constexpr int kShuffleStrided = Layout::kShuffleStrided;
|
|
|
|
// using Converter = typename Layout::Converter;
|
|
static constexpr int kConvertCount = Layout::kConvertCount;
|
|
|
|
// Use ldg128 load data from global memory
|
|
static constexpr int kAccessSize = 128;
|
|
using AccessType = uint4;
|
|
|
|
static constexpr int kElemsPerByte = 8 / kElemBits;
|
|
static constexpr int kElemsPerThread = kAccessSize / kElemBits;
|
|
static constexpr int kBytePerThread = kElemsPerThread / kElemsPerByte;
|
|
static constexpr int kThreadsNumPerTile = kStride / kElemsPerThread;
|
|
static constexpr int kThreadsNumPerInterleave =
|
|
kThreadsNumPerTile * kInterleave;
|
|
|
|
static constexpr int kConvertIters = kElemsPerThread / kConvertCount;
|
|
|
|
// Each thread loads 16(int8b)/32(int4b) quantized weight elements each time
|
|
// through ldg128 So more times of ldg128 are needed to load the same number
|
|
// of fp16 activation elements.
|
|
static constexpr int kActivationElemNumPerAccess =
|
|
kAccessSize / (sizeof(half) * 8);
|
|
static constexpr int kActivationAccessNum =
|
|
kElemsPerThread / kActivationElemNumPerAccess;
|
|
};
|
|
|
|
struct WeightOnlyPerChannel;
|
|
template <int GS>
|
|
struct WeightOnlyGroupWise;
|
|
|
|
template <typename WeightOnlyFlag>
|
|
struct WeightOnlyProperties;
|
|
|
|
template <>
|
|
struct WeightOnlyProperties<WeightOnlyPerChannel> {
|
|
static constexpr bool kIsFineGrained = false;
|
|
static constexpr int kGroupSize = 0;
|
|
};
|
|
|
|
template <int GS>
|
|
struct WeightOnlyProperties<WeightOnlyGroupWise<GS>> {
|
|
static constexpr bool kIsFineGrained = true;
|
|
static constexpr int kGroupSize = GS;
|
|
};
|
|
|
|
template <typename T,
|
|
WeightOnlyQuantType QType,
|
|
typename WeightOnlyFlag,
|
|
bool Zero,
|
|
int BlockSize>
|
|
struct WeightOnlyScaleLoader {
|
|
using Details = WeightOnlyKernelDetails<QType>;
|
|
static constexpr bool kIsFineGrained =
|
|
WeightOnlyProperties<WeightOnlyFlag>::kIsFineGrained;
|
|
static constexpr int kGroupSize =
|
|
WeightOnlyProperties<WeightOnlyFlag>::kGroupSize;
|
|
|
|
private:
|
|
const T* _scales;
|
|
const T* _zeros;
|
|
int _stride;
|
|
int _offset;
|
|
|
|
public:
|
|
__device__ __forceinline__ WeightOnlyScaleLoader(const T* scales,
|
|
const T* zeros,
|
|
int initial_offset,
|
|
int stride)
|
|
: _scales(scales), _zeros(zeros), _stride(stride) {
|
|
_scales += initial_offset;
|
|
if (Zero) {
|
|
_zeros += initial_offset;
|
|
}
|
|
// Calculate the k dimension index of the element processed by the current
|
|
// thread of layout before interleave Used to load scales and zeros in
|
|
// groupwise weight only quant
|
|
_offset =
|
|
threadIdx.x / Details::kThreadsNumPerInterleave * Details::kStride +
|
|
(threadIdx.x % Details::kThreadsNumPerTile) * Details::kElemsPerThread;
|
|
}
|
|
|
|
__device__ __forceinline__ void load(T* scale, T* zero, int nid) {
|
|
int offset = nid * Details::kInterleave;
|
|
|
|
// TODO(freeliuzc): cpplint has bug here
|
|
#ifndef WIN32
|
|
if constexpr (kIsFineGrained) {
|
|
#else
|
|
if (kIsFineGrained) {
|
|
#endif
|
|
offset += _offset / kGroupSize * _stride;
|
|
}
|
|
*scale = _scales[offset];
|
|
|
|
// TODO(freeliuzc): cpplint has bug here
|
|
#ifndef WIN32
|
|
if constexpr (Zero) {
|
|
#else
|
|
if (Zero) {
|
|
#endif
|
|
*zero = _zeros[offset];
|
|
} else {
|
|
*zero = static_cast<T>(0.f);
|
|
}
|
|
}
|
|
|
|
__device__ __forceinline__ void advance() {
|
|
_offset += BlockSize * Details::kElemsPerThread / Details::kInterleave;
|
|
}
|
|
|
|
__device__ __forceinline__ int offset() { return _offset; }
|
|
}; // NOLINT
|
|
|
|
template <typename T, WeightOnlyQuantType QType>
|
|
struct WeightOnlyConverter {};
|
|
|
|
template <>
|
|
struct WeightOnlyConverter<half, WeightOnlyQuantType::Int8b> {
|
|
static __device__ inline void convert(half halves[4],
|
|
int8_t signed_chars[4]) {
|
|
uint32_t* h = reinterpret_cast<uint32_t*>(halves);
|
|
uint32_t i8s = *reinterpret_cast<uint32_t*>(signed_chars);
|
|
|
|
static constexpr uint32_t mask_for_elt_01 = 0x5150;
|
|
static constexpr uint32_t mask_for_elt_23 = 0x5352;
|
|
static constexpr uint32_t start_byte_for_fp16 = 0x64646464;
|
|
asm volatile("prmt.b32 %0,%1,%2,%3;\n"
|
|
: "=r"(h[0])
|
|
: "r"(i8s), "n"(start_byte_for_fp16), "n"(mask_for_elt_01));
|
|
asm volatile("prmt.b32 %0,%1,%2,%3;\n"
|
|
: "=r"(h[1])
|
|
: "r"(i8s), "n"(start_byte_for_fp16), "n"(mask_for_elt_23));
|
|
|
|
static constexpr uint32_t I8s_TO_F16s_MAGIC_NUM = 0x64806480;
|
|
asm volatile("sub.f16x2 %0, %1, %2;\n"
|
|
: "=r"(h[0])
|
|
: "r"(h[0]), "r"(I8s_TO_F16s_MAGIC_NUM));
|
|
asm volatile("sub.f16x2 %0, %1, %2;\n"
|
|
: "=r"(h[1])
|
|
: "r"(h[1]), "r"(I8s_TO_F16s_MAGIC_NUM));
|
|
}
|
|
};
|
|
|
|
#ifdef PADDLE_CUDA_BF16
|
|
template <>
|
|
struct WeightOnlyConverter<__nv_bfloat16, WeightOnlyQuantType::Int8b> {
|
|
static __device__ inline void convert(__nv_bfloat16 halves[4],
|
|
int8_t signed_chars[4]) {
|
|
#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800))
|
|
uint32_t* bf16_result_ptr = reinterpret_cast<uint32_t*>(halves);
|
|
uint32_t i8s = *reinterpret_cast<uint32_t*>(signed_chars);
|
|
|
|
static constexpr uint32_t fp32_base = 0x4B000000;
|
|
float fp32_intermediates[4];
|
|
|
|
// Construct FP32s, bfloat does not have enough mantissa for IADD trick
|
|
uint32_t* fp32_intermediates_casted =
|
|
reinterpret_cast<uint32_t*>(fp32_intermediates);
|
|
fp32_intermediates_casted[0] = __byte_perm(i8s, fp32_base, 0x7650);
|
|
fp32_intermediates_casted[1] = __byte_perm(i8s, fp32_base, 0x7651);
|
|
fp32_intermediates_casted[2] = __byte_perm(i8s, fp32_base, 0x7652);
|
|
fp32_intermediates_casted[3] = __byte_perm(i8s, fp32_base, 0x7653);
|
|
|
|
// Subtract out fp32_base + 128 to make the unsigned integer signed.
|
|
#pragma unroll
|
|
for (int ii = 0; ii < 4; ++ii) {
|
|
fp32_intermediates[ii] -= 8388736.f;
|
|
}
|
|
|
|
// Truncate the fp32 representation and pack up as bfloat16s.
|
|
#pragma unroll
|
|
for (int ii = 0; ii < 2; ++ii) {
|
|
bf16_result_ptr[ii] = __byte_perm(fp32_intermediates_casted[2 * ii + 0],
|
|
fp32_intermediates_casted[2 * ii + 1],
|
|
0x7632);
|
|
}
|
|
#else
|
|
// Disable this on architectures older than Ampere since they lack hardware
|
|
// for bf16 mma. If one wishes to use HMMA on older hardware, they should
|
|
// Convert directly to FP16 using FP16 converters.
|
|
assert(false);
|
|
#endif
|
|
}
|
|
};
|
|
#endif
|
|
|
|
template <>
|
|
struct WeightOnlyConverter<half, WeightOnlyQuantType::Int4b> {
|
|
static __device__ inline void convert(half halves[8],
|
|
int8_t signed_chars[4]) {
|
|
uint32_t* h = reinterpret_cast<uint32_t*>(halves);
|
|
uint32_t i4s = *reinterpret_cast<uint32_t*>(signed_chars);
|
|
|
|
// First, we extract the i4s and construct an intermediate fp16 number.
|
|
static constexpr uint32_t immLut = (0xf0 & 0xcc) | 0xaa;
|
|
static constexpr uint32_t BOTTOM_MASK = 0x000f000f;
|
|
static constexpr uint32_t TOP_MASK = 0x00f000f0;
|
|
static constexpr uint32_t I4s_TO_F16s_MAGIC_NUM = 0x64006400;
|
|
|
|
// Note that the entire sequence only requires 1 shift instruction. This is
|
|
// thanks to the register packing format and the fact that we force our
|
|
// integers to be unsigned, and account for this in the fp16 subtractions.
|
|
// In addition, I exploit the fact that sub and fma have the same throughput
|
|
// in order to convert elt_23 and elt_67 to fp16 without having to shift
|
|
// them to the bottom bits before hand.
|
|
|
|
// Shift right by 8 to now consider elt_45 and elt_67. Issue first to hide
|
|
// RAW dependency if we issue immediately before required.
|
|
const uint32_t top_i4s = i4s >> 8;
|
|
// Extract elt_01 - (i4s & 0x000f000f) | 0x64006400
|
|
asm volatile(
|
|
"lop3.b32 %0, %1, %2, %3, %4;\n"
|
|
: "=r"(h[0])
|
|
: "r"(i4s), "n"(BOTTOM_MASK), "n"(I4s_TO_F16s_MAGIC_NUM), "n"(immLut));
|
|
// Extract elt_23 (i4s & 0x00f000f0) | 0x64006400
|
|
asm volatile(
|
|
"lop3.b32 %0, %1, %2, %3, %4;\n"
|
|
: "=r"(h[1])
|
|
: "r"(i4s), "n"(TOP_MASK), "n"(I4s_TO_F16s_MAGIC_NUM), "n"(immLut));
|
|
// Extract elt_45 (top_i4s & 0x000f000f) | 0x64006400
|
|
asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n"
|
|
: "=r"(h[2])
|
|
: "r"(top_i4s),
|
|
"n"(BOTTOM_MASK),
|
|
"n"(I4s_TO_F16s_MAGIC_NUM),
|
|
"n"(immLut));
|
|
// Extract elt_67 (top_i4s & 0x00f000f0) | 0x64006400
|
|
asm volatile(
|
|
"lop3.b32 %0, %1, %2, %3, %4;\n"
|
|
: "=r"(h[3])
|
|
: "r"(top_i4s), "n"(TOP_MASK), "n"(I4s_TO_F16s_MAGIC_NUM), "n"(immLut));
|
|
|
|
// I use inline PTX below because I am not sure if the compiler will emit
|
|
// float2half instructions if I use the half2 ctor. In this case, I chose
|
|
// performance reliability over code readability.
|
|
|
|
// This is the half2 {1032, 1032} represented as an integer.
|
|
static constexpr uint32_t FP16_TOP_MAGIC_NUM = 0x64086408;
|
|
// This is the half2 {1 / 16, 1 / 16} represented as an integer.
|
|
static constexpr uint32_t ONE_SIXTEENTH = 0x2c002c00;
|
|
// This is the half2 {-72, -72} represented as an integer.
|
|
static constexpr uint32_t NEG_72 = 0xd480d480;
|
|
// Finally, we construct the output numbers.
|
|
// Convert elt_01
|
|
asm volatile("sub.f16x2 %0, %1, %2;\n"
|
|
: "=r"(h[0])
|
|
: "r"(h[0]), "r"(FP16_TOP_MAGIC_NUM));
|
|
// Convert elt_23
|
|
asm volatile("fma.rn.f16x2 %0, %1, %2, %3;\n"
|
|
: "=r"(h[1])
|
|
: "r"(h[1]), "r"(ONE_SIXTEENTH), "r"(NEG_72));
|
|
// Convert elt_45
|
|
asm volatile("sub.f16x2 %0, %1, %2;\n"
|
|
: "=r"(h[2])
|
|
: "r"(h[2]), "r"(FP16_TOP_MAGIC_NUM));
|
|
// Convert elt_67
|
|
asm volatile("fma.rn.f16x2 %0, %1, %2, %3;\n"
|
|
: "=r"(h[3])
|
|
: "r"(h[3]), "r"(ONE_SIXTEENTH), "r"(NEG_72));
|
|
}
|
|
};
|
|
|
|
#ifdef PADDLE_CUDA_BF16
|
|
template <>
|
|
struct WeightOnlyConverter<__nv_bfloat16, WeightOnlyQuantType::Int4b> {
|
|
static __device__ inline void convert(__nv_bfloat16 halves[8],
|
|
int8_t signed_chars[4]) {
|
|
#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800))
|
|
uint32_t* h = reinterpret_cast<uint32_t*>(halves);
|
|
uint32_t const source_i4s = *reinterpret_cast<uint32_t*>(signed_chars);
|
|
|
|
static constexpr uint32_t immLut = (0xf0 & 0xcc) | 0xaa;
|
|
static constexpr uint32_t MASK = 0x000f000f;
|
|
static constexpr uint32_t I4s_TO_BF16s_MAGIC_NUM = 0x43004300;
|
|
|
|
// We don't have enough mantissa to remove as much shift overhead as FP16,
|
|
// so we must loop. No shift needed for first item.
|
|
uint32_t i4s = source_i4s;
|
|
asm volatile(
|
|
"lop3.b32 %0, %1, %2, %3, %4;\n"
|
|
: "=r"(h[0])
|
|
: "r"(i4s), "n"(MASK), "n"(I4s_TO_BF16s_MAGIC_NUM), "n"(immLut));
|
|
#pragma unroll
|
|
for (int ii = 1; ii < 4; ++ii) {
|
|
i4s >>= 4;
|
|
// (i4s & 0x000f000f) | 0x43004300
|
|
asm volatile(
|
|
"lop3.b32 %0, %1, %2, %3, %4;\n"
|
|
: "=r"(h[ii])
|
|
: "r"(i4s), "n"(MASK), "n"(I4s_TO_BF16s_MAGIC_NUM), "n"(immLut));
|
|
}
|
|
|
|
// This is the BF16 {-136, -136} represented as an integer.
|
|
static constexpr uint32_t BF16_BIAS = 0xC308C308;
|
|
static constexpr uint32_t BF16_ONE = 0x3F803F80;
|
|
|
|
// Finally, we construct the output numbers.
|
|
#pragma unroll
|
|
for (int ii = 0; ii < 4; ++ii) {
|
|
// Since this section is for Ampere+, we use bf16 fma to do the bias
|
|
// subtraction
|
|
asm("fma.rn.bf16x2 %0, %1, %2, %3;\n"
|
|
: "=r"(h[ii])
|
|
: "r"(h[ii]), "r"(BF16_ONE), "r"(BF16_BIAS));
|
|
}
|
|
|
|
#else
|
|
// Disable this on architectures older than Ampere since they lack hardware
|
|
// for bf16 mma. If one wishes to use HMMA on older hardware, they should
|
|
// Convert directly to FP16 using FP16 converters.
|
|
assert(false);
|
|
#endif
|
|
}
|
|
};
|
|
#endif
|
|
|
|
template <typename VecType, typename T0, typename T1>
|
|
__device__ __forceinline__ void load(T0* dst, T1* src, size_t offset = 0) {
|
|
*reinterpret_cast<VecType*>(dst) =
|
|
*(reinterpret_cast<const VecType*>(src) + offset);
|
|
}
|
|
|
|
template <typename T, WeightOnlyQuantType QType, typename Details>
|
|
struct WeightPostProcessor {
|
|
static __device__ __forceinline__ void run(T* weights_vec,
|
|
T* weights_f16,
|
|
T* scale,
|
|
T* zero,
|
|
int NPerBlock,
|
|
int idx) {}
|
|
};
|
|
|
|
template <typename T, typename Details>
|
|
struct WeightPostProcessor<T, WeightOnlyQuantType::Int4b, Details> {
|
|
static __device__ __forceinline__ void run(T* weights_vec,
|
|
T* weights_f16,
|
|
T* scale,
|
|
T* zero,
|
|
int NPerBlock,
|
|
int idx) {
|
|
using HALF_2_TYPE = typename CUDA_HALF_2_TYPE_TARIS<T>::type;
|
|
#pragma unroll
|
|
for (int i = 0; i < Details::kShuffleContinuous; ++i) {
|
|
#pragma unroll
|
|
for (int j = 0; j < Details::kShuffleStrided; ++j) {
|
|
// Dequantize the weights and arrange the shuffled elements back to
|
|
// the correct order in the register array
|
|
HALF_2_TYPE v = *reinterpret_cast<HALF_2_TYPE*>(
|
|
weights_vec + i * Details::kShuffleBasicTile +
|
|
j * Details::kShuffleContinuous * Details::kShuffleBasicTile);
|
|
v = HalfMulAdd<HALF_2_TYPE>::apply(
|
|
v,
|
|
ConvertDstFunc_2<HALF_2_TYPE>::apply(scale[idx]),
|
|
ConvertDstFunc_2<HALF_2_TYPE>::apply(zero[idx]));
|
|
weights_f16[(i * Details::kShuffleStrided * Details::kShuffleBasicTile +
|
|
j * Details::kShuffleBasicTile + 0) *
|
|
NPerBlock +
|
|
idx] = v.x;
|
|
weights_f16[(i * Details::kShuffleStrided * Details::kShuffleBasicTile +
|
|
j * Details::kShuffleBasicTile + 1) *
|
|
NPerBlock +
|
|
idx] = v.y;
|
|
}
|
|
}
|
|
}
|
|
};
|
|
|
|
template <typename T, typename Details>
|
|
struct WeightPostProcessor<T, WeightOnlyQuantType::Int8b, Details> {
|
|
static __device__ __forceinline__ void run(T* weights_vec,
|
|
T* weights_f16,
|
|
T* scale,
|
|
T* zero,
|
|
int NPerBlock,
|
|
int idx) {
|
|
#pragma unroll
|
|
for (int p = 0; p < 16; ++p) {
|
|
weights_f16[p * NPerBlock + idx] =
|
|
weights_vec[p / 8 + (p % 8) * 2] * scale[idx];
|
|
}
|
|
}
|
|
};
|
|
|
|
template <typename T,
|
|
WeightOnlyQuantType QType,
|
|
typename WeightOnlyFlag,
|
|
bool Gelu,
|
|
bool Zero,
|
|
bool Bias,
|
|
int NPerBlock,
|
|
int Batch,
|
|
int BlockSize>
|
|
__global__ void weight_only_batched_gemv_multi_warp(const T* in,
|
|
const int8_t* qweight,
|
|
const T* bias,
|
|
const T* scales,
|
|
const T* zeros,
|
|
T* out,
|
|
const int n,
|
|
const int k) {
|
|
static_assert(NPerBlock == 1 || (NPerBlock % 2 == 0),
|
|
"NPerBlock must be 1 or even in gemv multi warp kernel. ");
|
|
using Details = WeightOnlyKernelDetails<QType>;
|
|
|
|
// using Converter = typename Details::Converter;
|
|
using AccType = typename Details::AccessType;
|
|
using CvtSrcType = int8_t;
|
|
using CvtResType = T;
|
|
using ScaleLoader =
|
|
WeightOnlyScaleLoader<T, QType, WeightOnlyFlag, Zero, BlockSize>;
|
|
using WeightProcessor = WeightPostProcessor<T, QType, Details>;
|
|
|
|
extern __shared__ uint8_t shmem[];
|
|
constexpr int Interleave = Details::kInterleave;
|
|
constexpr int WarpSize = 32;
|
|
constexpr int Num = Batch * NPerBlock;
|
|
const int tid = threadIdx.x;
|
|
const int bid = blockIdx.x;
|
|
const int n_start_id = bid * NPerBlock * Interleave;
|
|
using HALF_2_TYPE = typename CUDA_HALF_2_TYPE_TARIS<T>::type;
|
|
// Calculate the n-dimensional index of the data processed by the current
|
|
// thread in the interleave tile
|
|
const int interleave_n_id = (tid / Details::kThreadsNumPerTile) % Interleave;
|
|
|
|
qweight += n_start_id * k / Details::kElemsPerByte;
|
|
ScaleLoader scale_loader(scales, zeros, n_start_id + interleave_n_id, n);
|
|
|
|
float(*sm)[Num * Interleave] =
|
|
reinterpret_cast<float(*)[Num * Interleave]>(shmem);
|
|
|
|
// In order to take advantage of hfma2, we use fp16 for accumulation within
|
|
// threads and fp32 for accumulation between threads.
|
|
T accumulator[Num];
|
|
for (int i = 0; i < Num; ++i) {
|
|
accumulator[i] = ConvertDstFunc<T>::apply(0.f);
|
|
}
|
|
|
|
// Iteration in k dimensions
|
|
for (int local_k = tid * Details::kElemsPerThread; local_k < k * Interleave;
|
|
local_k += BlockSize * Details::kElemsPerThread) {
|
|
T weights_f16[Details::kElemsPerThread * NPerBlock];
|
|
T scale[NPerBlock], zero[NPerBlock];
|
|
#pragma unroll
|
|
for (int idx = 0; idx < NPerBlock; ++idx) {
|
|
// Load quantized weight and scales/zeros
|
|
int8_t weights_quantized[Details::kBytePerThread];
|
|
load<AccType>(weights_quantized,
|
|
qweight + idx * Interleave * k / Details::kElemsPerByte +
|
|
local_k / Details::kElemsPerByte);
|
|
scale_loader.load(scale + idx, zero + idx, idx);
|
|
T weights_vec[Details::kElemsPerThread];
|
|
|
|
#pragma unroll
|
|
for (int i = 0; i < Details::kConvertIters; ++i) {
|
|
// Use cutlass::FastInterleavedAndBiasedNumericArrayConverter for I2F
|
|
// type conversion
|
|
WeightOnlyConverter<T, QType>::convert(
|
|
weights_vec + i * Details::kConvertCount,
|
|
weights_quantized +
|
|
i * Details::kConvertCount / Details::kElemsPerByte);
|
|
}
|
|
// Assign weight and apply scales.
|
|
// Currently not support zero.
|
|
WeightProcessor::run(
|
|
weights_vec, weights_f16, scale, zero, NPerBlock, idx);
|
|
}
|
|
#pragma unroll
|
|
for (int b = 0; b < Batch; ++b) {
|
|
T in_v[Details::kElemsPerThread];
|
|
#pragma unroll
|
|
for (int idx = 0; idx < Details::kActivationAccessNum; ++idx) {
|
|
load<AccType>(in_v + idx * Details::kActivationElemNumPerAccess,
|
|
in + b * k + scale_loader.offset() +
|
|
idx * Details::kActivationElemNumPerAccess);
|
|
}
|
|
// Perform vector inner product and accumulate
|
|
#ifndef WIN32
|
|
if constexpr (NPerBlock == 1) {
|
|
#else
|
|
if (NPerBlock == 1) {
|
|
#endif
|
|
HALF_2_TYPE v = ConvertDstFunc<HALF_2_TYPE>::apply(0.f);
|
|
#pragma unroll
|
|
for (int y = 0; y < Details::kElemsPerThread; y += 2) {
|
|
v = HalfMulAdd<HALF_2_TYPE>::apply(
|
|
*reinterpret_cast<HALF_2_TYPE*>(weights_f16 + y),
|
|
*reinterpret_cast<HALF_2_TYPE*>(in_v + y),
|
|
v);
|
|
}
|
|
accumulator[b] = accumulator[b] + ConvertDstFunc<T>::apply(v.x + v.y);
|
|
} else {
|
|
#pragma unroll
|
|
for (int x = 0; x < NPerBlock / 2; ++x) {
|
|
#pragma unroll
|
|
for (int y = 0; y < Details::kElemsPerThread; ++y) {
|
|
*reinterpret_cast<HALF_2_TYPE*>(accumulator + b * NPerBlock +
|
|
x * 2) =
|
|
HalfMulAdd<HALF_2_TYPE>::apply(
|
|
*reinterpret_cast<HALF_2_TYPE*>(weights_f16 +
|
|
y * NPerBlock + x * 2),
|
|
ConvertDstFunc_2<HALF_2_TYPE>::apply(in_v[y]),
|
|
*reinterpret_cast<HALF_2_TYPE*>(accumulator +
|
|
b * NPerBlock + x * 2));
|
|
}
|
|
}
|
|
}
|
|
}
|
|
scale_loader.advance();
|
|
}
|
|
float reses[Num];
|
|
#pragma unroll
|
|
for (int i = 0; i < Num; ++i) {
|
|
reses[i] = ConvertFloatFunc<T>::apply(accumulator[i]);
|
|
}
|
|
|
|
// Each warp completes the internal reduce and writes the [Batch * NPerBlock *
|
|
// Interleave] results to the corresponding address in shared memory
|
|
Details::Layout::sync<Num, WarpSize>(reses, sm);
|
|
|
|
// Each thread is responsible for the accumulation and store to global memory
|
|
// of one element
|
|
for (int i = tid; i < Num * Interleave; i += BlockSize) {
|
|
int nid = i % (NPerBlock * Interleave);
|
|
float v = 0.f;
|
|
for (int j = 0; j < BlockSize / WarpSize; ++j) {
|
|
v += sm[j][i];
|
|
}
|
|
float bias_v = 0.f;
|
|
#ifndef WIN32
|
|
if constexpr (Bias) {
|
|
#else
|
|
if (Bias) {
|
|
#endif
|
|
bias_v = ConvertFloatFunc<T>::apply(bias[n_start_id + nid]);
|
|
}
|
|
int b = i / NPerBlock / Interleave;
|
|
out[b * n + n_start_id + nid] = ConvertDstFunc<T>::apply(
|
|
GeluActivation<float, Gelu>::apply(v + bias_v));
|
|
}
|
|
}
|
|
#endif
|
|
|
|
template <typename T,
|
|
WeightOnlyQuantType QType,
|
|
typename WeightOnlyFlag,
|
|
int NPerBlock,
|
|
int Batch,
|
|
int BlockSize>
|
|
void select_activation_and_bias(const T* input,
|
|
const int8_t* weight,
|
|
const T* bias,
|
|
const T* scales,
|
|
const int m,
|
|
const int n,
|
|
const int k,
|
|
const std::string& act_method,
|
|
T* output,
|
|
cudaStream_t stream) {
|
|
#ifdef PADDLE_WITH_CUDA
|
|
static constexpr int kInterleave = WeightLayoutDetails<QType>::kInterleave;
|
|
dim3 grid(n / NPerBlock / kInterleave);
|
|
dim3 block(BlockSize);
|
|
int size = sizeof(float) * BlockSize / 32 * Batch * NPerBlock * kInterleave;
|
|
if (bias) {
|
|
if (act_method == "gelu") {
|
|
weight_only_batched_gemv_multi_warp<T,
|
|
QType,
|
|
WeightOnlyFlag,
|
|
true,
|
|
false,
|
|
true,
|
|
NPerBlock,
|
|
Batch,
|
|
BlockSize>
|
|
<<<grid, block, size, stream>>>(
|
|
input, weight, bias, scales, /*zeros*/ nullptr, output, n, k);
|
|
} else if (act_method == "None") {
|
|
weight_only_batched_gemv_multi_warp<T,
|
|
QType,
|
|
WeightOnlyFlag,
|
|
false,
|
|
false,
|
|
true,
|
|
NPerBlock,
|
|
Batch,
|
|
BlockSize>
|
|
<<<grid, block, size, stream>>>(
|
|
input, weight, bias, scales, /*zeros*/ nullptr, output, n, k);
|
|
} else {
|
|
PADDLE_THROW(
|
|
errors::InvalidArgument("Currently, weightonly GEMV act_method "
|
|
"only support `gelu`, `None`. "));
|
|
}
|
|
} else {
|
|
if (act_method == "gelu") {
|
|
weight_only_batched_gemv_multi_warp<T,
|
|
QType,
|
|
WeightOnlyFlag,
|
|
true,
|
|
false,
|
|
false,
|
|
NPerBlock,
|
|
Batch,
|
|
BlockSize>
|
|
<<<grid, block, size, stream>>>(
|
|
input, weight, bias, scales, /*zeros*/ nullptr, output, n, k);
|
|
} else if (act_method == "None") {
|
|
weight_only_batched_gemv_multi_warp<T,
|
|
QType,
|
|
WeightOnlyFlag,
|
|
false,
|
|
false,
|
|
false,
|
|
NPerBlock,
|
|
Batch,
|
|
BlockSize>
|
|
<<<grid, block, size, stream>>>(
|
|
input, weight, bias, scales, /*zeros*/ nullptr, output, n, k);
|
|
} else {
|
|
PADDLE_THROW(
|
|
errors::InvalidArgument("Currently, weightonly GEMV act_method "
|
|
"only support `gelu`, `None`. "));
|
|
}
|
|
}
|
|
#endif
|
|
}
|
|
|
|
template <typename T, typename WeightOnlyFlag>
|
|
void weight_only_batched_gemv_launcher(
|
|
const T* input,
|
|
const int8_t* weight,
|
|
const T* bias,
|
|
const T* scales,
|
|
int m,
|
|
int n,
|
|
int k,
|
|
const std::string& weight_only_quant_type,
|
|
const std::string& act_method,
|
|
T* output,
|
|
cudaStream_t stream) {
|
|
#ifdef PADDLE_WITH_CUDA
|
|
if (weight_only_quant_type == "int4") {
|
|
switch (m) {
|
|
case 1: {
|
|
select_activation_and_bias<T,
|
|
WeightOnlyQuantType::Int4b,
|
|
WeightOnlyFlag,
|
|
1,
|
|
1,
|
|
192>(
|
|
input, weight, bias, scales, m, n, k, act_method, output, stream);
|
|
break;
|
|
}
|
|
case 2: {
|
|
select_activation_and_bias<T,
|
|
WeightOnlyQuantType::Int4b,
|
|
WeightOnlyFlag,
|
|
2,
|
|
2,
|
|
128>(
|
|
input, weight, bias, scales, m, n, k, act_method, output, stream);
|
|
break;
|
|
}
|
|
case 3: {
|
|
select_activation_and_bias<T,
|
|
WeightOnlyQuantType::Int4b,
|
|
WeightOnlyFlag,
|
|
2,
|
|
3,
|
|
256>(
|
|
input, weight, bias, scales, m, n, k, act_method, output, stream);
|
|
break;
|
|
}
|
|
case 4: {
|
|
select_activation_and_bias<T,
|
|
WeightOnlyQuantType::Int4b,
|
|
WeightOnlyFlag,
|
|
4,
|
|
4,
|
|
256>(
|
|
input, weight, bias, scales, m, n, k, act_method, output, stream);
|
|
break;
|
|
}
|
|
default: {
|
|
throw std::runtime_error(
|
|
"Weight only cuda kernel only supported bs <= 4");
|
|
break;
|
|
}
|
|
}
|
|
} else if (weight_only_quant_type == "int8") {
|
|
switch (m) {
|
|
case 1: {
|
|
select_activation_and_bias<T,
|
|
WeightOnlyQuantType::Int8b,
|
|
WeightOnlyFlag,
|
|
2,
|
|
1,
|
|
256>(
|
|
input, weight, bias, scales, m, n, k, act_method, output, stream);
|
|
break;
|
|
}
|
|
case 2: {
|
|
select_activation_and_bias<T,
|
|
WeightOnlyQuantType::Int8b,
|
|
WeightOnlyFlag,
|
|
2,
|
|
2,
|
|
256>(
|
|
input, weight, bias, scales, m, n, k, act_method, output, stream);
|
|
break;
|
|
}
|
|
case 3: {
|
|
select_activation_and_bias<T,
|
|
WeightOnlyQuantType::Int8b,
|
|
WeightOnlyFlag,
|
|
2,
|
|
3,
|
|
256>(
|
|
input, weight, bias, scales, m, n, k, act_method, output, stream);
|
|
break;
|
|
}
|
|
case 4: {
|
|
select_activation_and_bias<T,
|
|
WeightOnlyQuantType::Int8b,
|
|
WeightOnlyFlag,
|
|
2,
|
|
4,
|
|
256>(
|
|
input, weight, bias, scales, m, n, k, act_method, output, stream);
|
|
break;
|
|
}
|
|
default: {
|
|
throw std::runtime_error(
|
|
"Weight only cuda kernel only supported bs <= 4");
|
|
break;
|
|
}
|
|
}
|
|
} else {
|
|
PADDLE_THROW(common::errors::InvalidArgument(
|
|
"WeightOnlyGemvKernel quant_type only support 'int4' or 'int8'."));
|
|
}
|
|
#endif
|
|
}
|
|
|
|
} // namespace
|
|
|
|
template <typename T, typename Context>
|
|
void WeightOnlyGemvWrapper(const Context& dev_ctx,
|
|
const T* input,
|
|
const int8_t* weight,
|
|
const T* bias,
|
|
const T* scales,
|
|
int m,
|
|
int n,
|
|
int k,
|
|
int group_size,
|
|
const std::string& weight_only_quant_type,
|
|
const std::string& weight_only_type,
|
|
const std::string& act_method,
|
|
T* output) {
|
|
using DataType = typename PDDataTypeTraits<T>::DataType;
|
|
if (weight_only_type == "per_channel") {
|
|
PADDLE_ENFORCE_EQ(group_size,
|
|
-1,
|
|
common::errors::InvalidArgument(
|
|
"group size must be -1 in per-channel mode."));
|
|
|
|
weight_only_batched_gemv_launcher<DataType, WeightOnlyPerChannel>(
|
|
reinterpret_cast<const DataType*>(input),
|
|
reinterpret_cast<const int8_t*>(weight),
|
|
reinterpret_cast<const DataType*>(bias),
|
|
reinterpret_cast<const DataType*>(scales),
|
|
m,
|
|
n,
|
|
k,
|
|
weight_only_quant_type,
|
|
act_method,
|
|
reinterpret_cast<DataType*>(output),
|
|
dev_ctx.stream());
|
|
} else if (weight_only_type == "group_wise") {
|
|
if (group_size == 64) {
|
|
weight_only_batched_gemv_launcher<DataType, WeightOnlyGroupWise<64>>(
|
|
reinterpret_cast<const DataType*>(input),
|
|
reinterpret_cast<const int8_t*>(weight),
|
|
reinterpret_cast<const DataType*>(bias),
|
|
reinterpret_cast<const DataType*>(scales),
|
|
m,
|
|
n,
|
|
k,
|
|
weight_only_quant_type,
|
|
act_method,
|
|
reinterpret_cast<DataType*>(output),
|
|
dev_ctx.stream());
|
|
} else if (group_size == 128) {
|
|
weight_only_batched_gemv_launcher<DataType, WeightOnlyGroupWise<128>>(
|
|
reinterpret_cast<const DataType*>(input),
|
|
reinterpret_cast<const int8_t*>(weight),
|
|
reinterpret_cast<const DataType*>(bias),
|
|
reinterpret_cast<const DataType*>(scales),
|
|
m,
|
|
n,
|
|
k,
|
|
weight_only_quant_type,
|
|
act_method,
|
|
reinterpret_cast<DataType*>(output),
|
|
dev_ctx.stream());
|
|
} else {
|
|
PADDLE_THROW(common::errors::InvalidArgument(
|
|
"WeightOnlyGemvKernel group_size only support 64 or 128."));
|
|
}
|
|
} else {
|
|
PADDLE_THROW(common::errors::InvalidArgument(
|
|
"WeightOnlyGemvKernel type only support "
|
|
"'per_channel' or 'group_wise'."));
|
|
}
|
|
}
|
|
|
|
template <>
|
|
void WeightOnlyGemvWrapper(const GPUContext& dev_ctx,
|
|
const float* input,
|
|
const int8_t* weight,
|
|
const float* bias,
|
|
const float* scales,
|
|
int m,
|
|
int n,
|
|
int k,
|
|
int group_size,
|
|
const std::string& weight_only_quant_type,
|
|
const std::string& weight_only_type,
|
|
const std::string& act_method,
|
|
float* output) {
|
|
PADDLE_THROW(common::errors::Unimplemented(
|
|
"WeightOnlyGemvKernel type only support 'float16' and 'bfloa16."
|
|
"Not support float32."));
|
|
}
|
|
|
|
template <typename T, typename Context>
|
|
void WeightOnlyGemvKernel(const Context& dev_ctx,
|
|
const DenseTensor& x,
|
|
const DenseTensor& weight,
|
|
const optional<DenseTensor>& bias,
|
|
const DenseTensor& weight_scale,
|
|
int group_size,
|
|
const std::string& weight_only_quant_type,
|
|
const std::string& weight_only_type,
|
|
const std::string& act_method,
|
|
DenseTensor* out) {
|
|
const T* x_data = x.data<T>();
|
|
const int8_t* weight_data = weight.data<int8_t>();
|
|
// Actually, we pass the weight datatype is uint8_t type.
|
|
const T* bias_data = bias ? bias.get().data<T>() : nullptr;
|
|
const T* weight_scale_data = weight_scale.data<T>();
|
|
T* out_data = dev_ctx.template Alloc<T>(out);
|
|
int64_t m = x.dims()[0];
|
|
// TODO(large-tensor): downstream functors may still use int
|
|
|
|
int64_t k = x.dims()[1];
|
|
// TODO(large-tensor): downstream functors may still use int
|
|
|
|
int64_t n = weight.dims()[0];
|
|
// TODO(large-tensor): downstream functors may still use int
|
|
|
|
WeightOnlyGemvWrapper<T>(dev_ctx,
|
|
x_data,
|
|
weight_data,
|
|
bias_data,
|
|
weight_scale_data,
|
|
m,
|
|
n,
|
|
k,
|
|
group_size,
|
|
weight_only_quant_type,
|
|
weight_only_type,
|
|
act_method,
|
|
out_data);
|
|
}
|
|
|
|
template void WeightOnlyGemvWrapper(const GPUContext& dev_ctx,
|
|
const float* input,
|
|
const int8_t* weight,
|
|
const float* bias,
|
|
const float* scales,
|
|
int m,
|
|
int n,
|
|
int k,
|
|
int group_size,
|
|
const std::string& weight_only_quant_type,
|
|
const std::string& weight_only_type,
|
|
const std::string& act_method,
|
|
float* output);
|
|
|
|
template void WeightOnlyGemvWrapper(const GPUContext& dev_ctx,
|
|
const phi::float16* input,
|
|
const int8_t* weight,
|
|
const phi::float16* bias,
|
|
const phi::float16* scales,
|
|
int m,
|
|
int n,
|
|
int k,
|
|
int group_size,
|
|
const std::string& weight_only_quant_type,
|
|
const std::string& weight_only_type,
|
|
const std::string& act_method,
|
|
phi::float16* output);
|
|
#ifdef PADDLE_CUDA_BF16
|
|
template void WeightOnlyGemvWrapper(const GPUContext& dev_ctx,
|
|
const phi::bfloat16* input,
|
|
const int8_t* weight,
|
|
const phi::bfloat16* bias,
|
|
const phi::bfloat16* scales,
|
|
int m,
|
|
int n,
|
|
int k,
|
|
int group_size,
|
|
const std::string& weight_only_quant_type,
|
|
const std::string& weight_only_type,
|
|
const std::string& act_method,
|
|
phi::bfloat16* output);
|
|
#endif
|
|
|
|
} // namespace phi
|