wehub-resource-sync
617 lines
24 KiB
C++
617 lines
24 KiB
C++
// 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.
|
|
|
|
#pragma once
|
|
|
|
#include <cstdint>
|
|
#include "paddle/phi/backends/gpu/gpu_launch_config.h"
|
|
#include "paddle/phi/core/dense_tensor.h"
|
|
#include "paddle/phi/core/kernel_registry.h"
|
|
#include "paddle/phi/kernels/funcs/aligned_vector.h"
|
|
#include "paddle/phi/kernels/funcs/common_shape.h"
|
|
namespace phi {
|
|
|
|
template <typename IndexT>
|
|
__global__ void weight_permute_kernel_wint8(const int8_t* input_data_dev,
|
|
int8_t* output_data_dev,
|
|
IndexT numel,
|
|
IndexT total_k,
|
|
IndexT total_n) {
|
|
CUDA_KERNEL_LOOP_TYPE(linear_idx, numel, IndexT) {
|
|
IndexT k_id = linear_idx / total_n;
|
|
IndexT n_id = linear_idx % total_n;
|
|
IndexT k_mod_16 = k_id % 16;
|
|
|
|
constexpr int map[16] = {
|
|
0, 2, 4, 6, 8, 10, 12, 14, 1, 3, 5, 7, 9, 11, 13, 15};
|
|
IndexT permute_kk = map[k_mod_16] + k_id / 16 * 16;
|
|
|
|
IndexT permute_index = permute_kk % 64 + permute_kk / 64 * 128 +
|
|
64 * (n_id % 2) + total_k * 2 * (n_id / 2);
|
|
uint8_t shift_quant_weight = static_cast<uint8_t>(
|
|
static_cast<int32_t>(input_data_dev[linear_idx]) + 128);
|
|
output_data_dev[permute_index] =
|
|
*reinterpret_cast<int8_t*>(&shift_quant_weight);
|
|
}
|
|
}
|
|
|
|
template <typename IndexT>
|
|
__global__ void weight_permute_kernel_wint4(const int8_t* input_data_dev,
|
|
int8_t* output_data_dev,
|
|
IndexT numel,
|
|
IndexT total_k,
|
|
IndexT total_n) {
|
|
CUDA_KERNEL_LOOP_TYPE(linear_idx, numel, IndexT) {
|
|
IndexT k_id = linear_idx / total_n;
|
|
IndexT n_id = linear_idx % total_n;
|
|
// k_id is 8_bit index.
|
|
constexpr int map[16] = {
|
|
0, 4, 8, 12, 1, 5, 9, 13, 2, 6, 10, 14, 3, 7, 11, 15};
|
|
|
|
IndexT permute_kk = map[k_id % 16] + k_id / 16 * 16;
|
|
IndexT permute_index = permute_kk % 32 + permute_kk / 32 * 128 +
|
|
32 * (n_id % 4) + total_k * 2 * (n_id / 4);
|
|
int8_t shift_quant_weight = input_data_dev[linear_idx];
|
|
output_data_dev[permute_index] =
|
|
*reinterpret_cast<int8_t*>(&shift_quant_weight);
|
|
}
|
|
}
|
|
|
|
// convetr 0,1,2,3,4,5,6,7 4bit -> 0,2,4,6,1,3,5,7
|
|
__global__ void weight_interval_kernel_wint4(int8_t* output_data_dev,
|
|
int64_t numel) {
|
|
constexpr int value_per_interval_thread = 4;
|
|
int64_t linear_idx =
|
|
(static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x) *
|
|
value_per_interval_thread;
|
|
int64_t stride =
|
|
static_cast<int64_t>(blockDim.x) * gridDim.x * value_per_interval_thread;
|
|
|
|
for (; linear_idx < numel; linear_idx += stride) {
|
|
uint32_t value = *reinterpret_cast<uint32_t*>(output_data_dev + linear_idx);
|
|
uint32_t result = 0;
|
|
|
|
constexpr int map[8] = {0, 2, 4, 6, 1, 3, 5, 7};
|
|
|
|
for (int ii = 0; ii < 8; ii++) {
|
|
uint32_t tmp = value >> (map[ii] * 4);
|
|
tmp &= 0x0F;
|
|
tmp = (tmp + 8) & 0x0F;
|
|
tmp = tmp << (ii * 4);
|
|
result |= tmp;
|
|
}
|
|
|
|
*reinterpret_cast<uint32_t*>(output_data_dev + linear_idx) = result;
|
|
}
|
|
}
|
|
|
|
/*
|
|
For SM70 volta arch, weightonly int8 dequantize invoked in load global memory.
|
|
So it only need interleave in K-dimension
|
|
K_index: 0 1 2 3 -> 0 2 1 3
|
|
*/
|
|
template <typename IndexT>
|
|
__global__ void weight_interleave_add_bias_kernel_wint8(
|
|
const int8_t* input_data_dev,
|
|
int8_t* output_data_dev,
|
|
IndexT numel,
|
|
IndexT total_n) {
|
|
CUDA_KERNEL_LOOP_TYPE(linear_idx, numel, IndexT) {
|
|
IndexT k_id = linear_idx / total_n;
|
|
IndexT n_id = linear_idx % total_n;
|
|
constexpr int n_interleaved_factor = 4;
|
|
IndexT n_interleave_group_id = n_id / n_interleaved_factor;
|
|
IndexT n_interleave_id = n_id % n_interleaved_factor;
|
|
if (n_interleave_id == 1 || n_interleave_id == 2) {
|
|
/*
|
|
0001 xor 0011 -> 0010
|
|
0010 xor 0011 -> 0001
|
|
*/
|
|
n_interleave_id ^= 3;
|
|
}
|
|
const IndexT new_n_id =
|
|
n_interleave_group_id * n_interleaved_factor + n_interleave_id;
|
|
const IndexT interleave_idx = k_id * total_n + new_n_id;
|
|
|
|
uint8_t shift_quant_weight = static_cast<uint8_t>(
|
|
static_cast<int32_t>(input_data_dev[linear_idx]) + 128);
|
|
output_data_dev[interleave_idx] =
|
|
*reinterpret_cast<int8_t*>(&shift_quant_weight);
|
|
}
|
|
}
|
|
|
|
/*
|
|
For SM70 volta arch, weightonly int4 dequantize invoked in load global memory.
|
|
So it only need interleave in K-dimension
|
|
K_index: 0 1 2 3 4 5 6 7 -> 0 2 4 6 1 3 5 7
|
|
*/
|
|
template <typename IndexT>
|
|
__global__ void weight_interleave_add_bias_kernel_wint4(int8_t* input_data_dev,
|
|
int8_t* output_data_dev,
|
|
IndexT numel,
|
|
IndexT total_n) {
|
|
const IndexT num_registers = numel / 4;
|
|
uint32_t* packed_input = reinterpret_cast<uint32_t*>(input_data_dev);
|
|
uint32_t* packed_output = reinterpret_cast<uint32_t*>(output_data_dev);
|
|
CUDA_KERNEL_LOOP_TYPE(i, num_registers, IndexT) {
|
|
uint32_t current_pack = packed_input[i];
|
|
uint32_t transformed_pack = 0;
|
|
#pragma unroll
|
|
for (int idx = 0; idx < 8; ++idx) {
|
|
const int offset = idx / 4;
|
|
const int src = (idx % 4) * 2 + offset;
|
|
|
|
const int src_shift = src * 4;
|
|
const int dst_shift = idx * 4;
|
|
|
|
const uint32_t src_bits = ((current_pack >> src_shift) + 8) & 0xF;
|
|
transformed_pack |= (src_bits << dst_shift);
|
|
}
|
|
packed_output[i] = transformed_pack;
|
|
}
|
|
}
|
|
|
|
template <typename GPUContext, typename IndexT>
|
|
void weight_permute_gpu_impl(const GPUContext& dev_ctx,
|
|
int8_t* input_data,
|
|
int8_t* output_data,
|
|
const std::vector<int64_t>& shape,
|
|
const int32_t arch,
|
|
const std::string& algo) {
|
|
auto total_k = shape[0];
|
|
auto total_n = shape[1];
|
|
auto numel = total_k * total_n;
|
|
auto gpu_config = backends::gpu::GetGpuLaunchConfig1D(dev_ctx, numel, 1);
|
|
int grid_size = gpu_config.GetGridSize();
|
|
int block_size = gpu_config.GetBlockSize();
|
|
if ((arch == 100) || (arch == 90) || (arch == 89) || (arch == 86) ||
|
|
(arch == 80) || (arch == 75)) {
|
|
if (algo == "weight_only_int4") {
|
|
numel /= 2;
|
|
weight_permute_kernel_wint4<IndexT><<<grid_size, block_size>>>(
|
|
input_data, output_data, numel, total_k, total_n);
|
|
weight_interval_kernel_wint4<<<grid_size, block_size>>>(output_data,
|
|
numel);
|
|
} else {
|
|
weight_permute_kernel_wint8<IndexT><<<grid_size, block_size>>>(
|
|
input_data, output_data, numel, total_k, total_n);
|
|
}
|
|
} else if (arch == 70) {
|
|
if (algo == "weight_only_int4") {
|
|
weight_interleave_add_bias_kernel_wint4<IndexT>
|
|
<<<grid_size, block_size>>>(input_data, output_data, numel, total_n);
|
|
} else {
|
|
weight_interleave_add_bias_kernel_wint8<IndexT>
|
|
<<<grid_size, block_size>>>(input_data, output_data, numel, total_n);
|
|
}
|
|
}
|
|
}
|
|
|
|
template <typename GPUContext>
|
|
void weight_permute_gpu(const GPUContext& dev_ctx,
|
|
int8_t* input_data,
|
|
int8_t* output_data,
|
|
const std::vector<int64_t>& shape,
|
|
const int32_t arch,
|
|
const std::string& algo) {
|
|
int64_t numel = shape[0] * shape[1];
|
|
if (numel <= std::numeric_limits<int>::max()) {
|
|
weight_permute_gpu_impl<GPUContext, int>(
|
|
dev_ctx, input_data, output_data, shape, arch, algo);
|
|
} else {
|
|
weight_permute_gpu_impl<GPUContext, int64_t>(
|
|
dev_ctx, input_data, output_data, shape, arch, algo);
|
|
}
|
|
}
|
|
|
|
template <typename T, int VectorSize = 8, typename ScaleT>
|
|
__global__ void per_channel_quant_gpu(const T* weight_data,
|
|
int8_t* quanted_weight_data,
|
|
ScaleT* scale_data,
|
|
int total_k,
|
|
int64_t total_vec_n) {
|
|
int64_t n = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
|
|
if (n < total_vec_n) {
|
|
const int4* vec_weight_data_ptr =
|
|
reinterpret_cast<const int4*>(weight_data);
|
|
int2* vec_quanted_weight_data =
|
|
reinterpret_cast<int2*>(quanted_weight_data);
|
|
AlignedVector<float, VectorSize> abs_max;
|
|
#pragma unroll
|
|
for (int i = 0; i < VectorSize; ++i) {
|
|
abs_max[i] = static_cast<float>(0.0f);
|
|
}
|
|
#pragma unroll
|
|
for (int k = 0; k < total_k; ++k) {
|
|
int64_t linear_index = k * total_vec_n + n;
|
|
AlignedVector<T, VectorSize> weight;
|
|
*reinterpret_cast<int4*>(&weight) = vec_weight_data_ptr[linear_index];
|
|
#pragma unroll
|
|
for (int i = 0; i < VectorSize; ++i) {
|
|
abs_max[i] = fmaxf((abs_max[i]), fabsf((weight[i])));
|
|
}
|
|
}
|
|
AlignedVector<ScaleT, VectorSize> scale;
|
|
#pragma unroll
|
|
for (int i = 0; i < VectorSize; ++i) {
|
|
scale[i] = static_cast<ScaleT>(abs_max[i] / static_cast<float>(127.0f));
|
|
}
|
|
*reinterpret_cast<float4*>(scale_data + VectorSize * n) =
|
|
*reinterpret_cast<float4*>(&scale);
|
|
|
|
for (int k = 0; k < total_k; ++k) {
|
|
AlignedVector<int8_t, VectorSize> quanted_weight;
|
|
int64_t linear_index = k * total_vec_n + n;
|
|
AlignedVector<T, VectorSize> weight;
|
|
*reinterpret_cast<int4*>(&weight) =
|
|
*reinterpret_cast<const int4*>(vec_weight_data_ptr + linear_index);
|
|
#pragma unroll
|
|
for (int i = 0; i < VectorSize; ++i) {
|
|
float scaled_weight =
|
|
(static_cast<float>(weight[i]) / static_cast<float>(abs_max[i])) *
|
|
static_cast<float>(127.0);
|
|
int8_t clipped_weight = static_cast<int8_t>(
|
|
lroundf(fmaxf(-127.0f, fminf(127.0f, scaled_weight))));
|
|
quanted_weight[i] = clipped_weight;
|
|
}
|
|
*reinterpret_cast<int2*>(vec_quanted_weight_data + linear_index) =
|
|
*reinterpret_cast<int2*>(&quanted_weight);
|
|
}
|
|
}
|
|
}
|
|
|
|
template <typename T, int VectorSize = 8, typename ScaleT>
|
|
__global__ void per_channel_quant_gpu_int4_row_pack(const T* weight_data,
|
|
int8_t* quanted_weight_data,
|
|
ScaleT* scale_data,
|
|
int total_k,
|
|
int64_t total_vec_n) {
|
|
int64_t n = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
|
|
if (n < total_vec_n) {
|
|
const int4* vec_weight_data_ptr =
|
|
reinterpret_cast<const int4*>(weight_data);
|
|
int* vec_quanted_weight_data = reinterpret_cast<int*>(quanted_weight_data);
|
|
AlignedVector<float, VectorSize> abs_max;
|
|
#pragma unroll
|
|
for (int i = 0; i < VectorSize; ++i) {
|
|
abs_max[i] = static_cast<float>(0.0f);
|
|
}
|
|
#pragma unroll
|
|
for (int k = 0; k < total_k; ++k) {
|
|
int64_t linear_index = k * total_vec_n + n;
|
|
AlignedVector<T, VectorSize> weight;
|
|
*reinterpret_cast<int4*>(&weight) = vec_weight_data_ptr[linear_index];
|
|
#pragma unroll
|
|
for (int i = 0; i < VectorSize; ++i) {
|
|
abs_max[i] = fmaxf((abs_max[i]), fabsf((weight[i])));
|
|
}
|
|
}
|
|
AlignedVector<ScaleT, VectorSize> scale;
|
|
#pragma unroll
|
|
for (int i = 0; i < VectorSize; ++i) {
|
|
scale[i] = static_cast<ScaleT>(abs_max[i] / static_cast<float>(7.0f));
|
|
}
|
|
*reinterpret_cast<float4*>(scale_data + VectorSize * n) =
|
|
*reinterpret_cast<float4*>(&scale);
|
|
for (int k = 0; k < total_k; ++k) {
|
|
int64_t linear_index = k * total_vec_n + n;
|
|
AlignedVector<T, VectorSize> weight;
|
|
AlignedVector<int8_t, VectorSize / 2> quanted_weight;
|
|
*reinterpret_cast<int4*>(&weight) =
|
|
*reinterpret_cast<const int4*>(vec_weight_data_ptr + linear_index);
|
|
#pragma unroll
|
|
for (int i = 0; i < VectorSize / 2; ++i) {
|
|
int8_t packed_int4s = 0;
|
|
for (int pack = 0; pack < 2; ++pack) {
|
|
int vector_index = i * 2 + pack;
|
|
const float r_scale = 1 / static_cast<float>(scale[vector_index]);
|
|
const float weight_elt =
|
|
static_cast<float>(weight[vector_index]) * r_scale;
|
|
float scaled_weight = roundf(weight_elt);
|
|
int int_weight = static_cast<int>(scaled_weight);
|
|
#ifdef PADDLE_WITH_HIP
|
|
int8_t clipped_weight = max(-7, min(7, int_weight)) + 8;
|
|
#else
|
|
int8_t clipped_weight = max(-7, min(7, int_weight));
|
|
#endif
|
|
packed_int4s |= ((clipped_weight & 0x0F) << (4 * pack));
|
|
}
|
|
quanted_weight[i] = packed_int4s;
|
|
}
|
|
*reinterpret_cast<int*>(vec_quanted_weight_data + linear_index) =
|
|
*reinterpret_cast<int*>(&quanted_weight);
|
|
}
|
|
}
|
|
}
|
|
|
|
template <typename T, int VectorSize = 8, typename ScaleT>
|
|
__global__ void per_channel_quant_gpu_int4_col_pack(const T* weight_data,
|
|
int8_t* quanted_weight_data,
|
|
ScaleT* scale_data,
|
|
int total_k,
|
|
int64_t total_vec_n) {
|
|
int64_t n = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
|
|
if (n < total_vec_n) {
|
|
const int4* vec_weight_data_ptr =
|
|
reinterpret_cast<const int4*>(weight_data);
|
|
int2* vec_quanted_weight_data =
|
|
reinterpret_cast<int2*>(quanted_weight_data);
|
|
AlignedVector<float, VectorSize> abs_max;
|
|
#pragma unroll
|
|
for (int i = 0; i < VectorSize; ++i) {
|
|
abs_max[i] = static_cast<float>(0.0f);
|
|
}
|
|
#pragma unroll
|
|
for (int k = 0; k < total_k; ++k) {
|
|
int64_t linear_index = k * total_vec_n + n;
|
|
AlignedVector<T, VectorSize> weight;
|
|
*reinterpret_cast<int4*>(&weight) = vec_weight_data_ptr[linear_index];
|
|
#pragma unroll
|
|
for (int i = 0; i < VectorSize; ++i) {
|
|
abs_max[i] = fmaxf((abs_max[i]), static_cast<float>(fabsf(weight[i])));
|
|
}
|
|
}
|
|
AlignedVector<ScaleT, VectorSize> scale;
|
|
#pragma unroll
|
|
for (int i = 0; i < VectorSize; ++i) {
|
|
scale[i] = static_cast<ScaleT>(abs_max[i] / static_cast<float>(7.0f));
|
|
}
|
|
*reinterpret_cast<float4*>(scale_data + VectorSize * n) =
|
|
*reinterpret_cast<float4*>(&scale);
|
|
|
|
for (int k = 0; k < total_k / 2; ++k) {
|
|
AlignedVector<int8_t, VectorSize> quanted_weight;
|
|
for (int packed_idx = 0; packed_idx < 2; ++packed_idx) {
|
|
int64_t linear_index = (k * 2 + packed_idx) * total_vec_n + n;
|
|
AlignedVector<T, VectorSize> weight;
|
|
*reinterpret_cast<int4*>(&weight) =
|
|
*reinterpret_cast<const int4*>(vec_weight_data_ptr + linear_index);
|
|
#pragma unroll
|
|
for (int i = 0; i < VectorSize; ++i) {
|
|
const float weight_elt =
|
|
(static_cast<float>(weight[i]) / static_cast<float>(abs_max[i])) *
|
|
static_cast<float>(7.0);
|
|
const float scaled_weight = lroundf(weight_elt);
|
|
int int_weight = static_cast<int>(scaled_weight);
|
|
const int8_t clipped_weight = fmaxf(-7, fminf(7, int_weight));
|
|
quanted_weight[i] &= ~(0x0F << (4 * packed_idx));
|
|
quanted_weight[i] |= ((clipped_weight & 0x0F) << (4 * packed_idx));
|
|
}
|
|
}
|
|
int64_t linear_index_new = k * total_vec_n + n;
|
|
*reinterpret_cast<int2*>(vec_quanted_weight_data + linear_index_new) =
|
|
*reinterpret_cast<int2*>(&quanted_weight);
|
|
}
|
|
}
|
|
}
|
|
|
|
template <typename T, typename GPUContext, typename ScaleT>
|
|
void weight_quant_gpu(const GPUContext& dev_ctx,
|
|
const T* weight_data,
|
|
int8_t* quanted_weight_data,
|
|
ScaleT* scale_data,
|
|
const std::vector<int64_t>& shape,
|
|
const int32_t arch,
|
|
const std::string& algo) {
|
|
int64_t total_k = shape[0];
|
|
int64_t total_n = shape[1];
|
|
int64_t numel = total_k * total_n;
|
|
constexpr int kWarpSize = 32;
|
|
constexpr int kBlockSize = 64;
|
|
constexpr int kWarpNum = kBlockSize / kWarpSize;
|
|
constexpr int kVectorSize = 128 / sizeof(T) / 8;
|
|
PADDLE_ENFORCE_EQ(total_n % kVectorSize,
|
|
0,
|
|
common::errors::PreconditionNotMet(
|
|
"Currently, weight_quant_gpu kernel only support n "
|
|
"with multiple of %d, please use",
|
|
kVectorSize));
|
|
int64_t vec_total_n = total_n / kVectorSize;
|
|
int64_t kGridSize =
|
|
max((vec_total_n + kBlockSize - 1) / kBlockSize, int64_t(1));
|
|
if (algo == "weight_only_int4") {
|
|
#ifdef PADDLE_WITH_HIP
|
|
per_channel_quant_gpu_int4_row_pack<T, kVectorSize>
|
|
<<<kGridSize, kBlockSize>>>(
|
|
weight_data, quanted_weight_data, scale_data, total_k, vec_total_n);
|
|
#else
|
|
if ((arch == 100) || (arch == 90) || (arch == 89) || (arch == 86) ||
|
|
(arch == 80) || (arch == 75)) {
|
|
per_channel_quant_gpu_int4_col_pack<T, kVectorSize>
|
|
<<<kGridSize, kBlockSize>>>(weight_data,
|
|
quanted_weight_data,
|
|
scale_data,
|
|
total_k,
|
|
vec_total_n);
|
|
} else if ((arch == 70)) {
|
|
per_channel_quant_gpu_int4_row_pack<T, kVectorSize>
|
|
<<<kGridSize, kBlockSize>>>(weight_data,
|
|
quanted_weight_data,
|
|
scale_data,
|
|
total_k,
|
|
vec_total_n);
|
|
}
|
|
#endif
|
|
} else {
|
|
per_channel_quant_gpu<T, kVectorSize><<<kGridSize, kBlockSize>>>(
|
|
weight_data, quanted_weight_data, scale_data, total_k, vec_total_n);
|
|
}
|
|
}
|
|
|
|
template <typename IndexT>
|
|
__global__ void weight_permute_transpose_interleave_kernel_w4a8(
|
|
const int8_t* input_data_ptr,
|
|
int8_t* output_data_ptr,
|
|
IndexT original_k,
|
|
IndexT original_n) {
|
|
// every 2 k-direction 8bit , ie 4 k-direction 4bit,
|
|
// is packed to 2 int8, and assigned to a new new_index.
|
|
// so here / 4.
|
|
IndexT numel = original_k * original_n / 4;
|
|
CUDA_KERNEL_LOOP_TYPE(linear_idx, numel, IndexT) {
|
|
const IndexT k_group_id = linear_idx / original_n;
|
|
const IndexT n_id = linear_idx % original_n;
|
|
|
|
uint16_t res = 0;
|
|
for (int j = 0; j < 2; j++) {
|
|
const IndexT k_id = k_group_id * 2 + j;
|
|
uint16_t val = input_data_ptr[k_id * original_n + n_id];
|
|
val = val & 0xFF;
|
|
val = val << (j * 8);
|
|
res |= val;
|
|
}
|
|
|
|
constexpr int map[8] = {0, 2, 4, 6, 1, 3, 5, 7};
|
|
// remember output(in 16 bit granularity)'shape is
|
|
// [16, 4, original_k/64, original_n/4]
|
|
// index is :
|
|
// [k_group_id % 16, n_id % 4, k_group_id/16, n_id/4]
|
|
const IndexT new_index = map[k_group_id % 8] + k_group_id % 16 / 8 * 8 +
|
|
(n_id % 4) * 16 + k_group_id / 16 * (16 * 4) +
|
|
n_id / 4 * (original_k);
|
|
|
|
reinterpret_cast<uint16_t*>(output_data_ptr)[new_index] = res;
|
|
}
|
|
}
|
|
|
|
__global__ void w4a8_inplace_permute(uint32_t* output_data_ptr, int64_t numel) {
|
|
CUDA_KERNEL_LOOP_TYPE(linear_idx, numel, int64_t) {
|
|
const uint32_t value = output_data_ptr[linear_idx];
|
|
|
|
uint32_t res = 0;
|
|
|
|
const int map[8] = {0, 2, 4, 6, 1, 3, 5, 7};
|
|
for (int i = 0; i < 8; i++) {
|
|
uint32_t tmp = value >> (i * 4);
|
|
tmp = tmp & 0x0F;
|
|
tmp = tmp << (map[i] * 4);
|
|
res |= tmp;
|
|
}
|
|
output_data_ptr[linear_idx] = res;
|
|
}
|
|
}
|
|
|
|
template <typename GPUContext>
|
|
void weight_permute_gpu_w4a8(const GPUContext& dev_ctx,
|
|
const int8_t* input_data,
|
|
int8_t* output_data,
|
|
const std::vector<int64_t>& shape,
|
|
const int32_t arch,
|
|
const std::string& algo) {
|
|
auto original_k = shape[0] * 2;
|
|
auto original_n = shape[1];
|
|
auto original_numel = original_k * original_n;
|
|
auto gpu_config =
|
|
backends::gpu::GetGpuLaunchConfig1D(dev_ctx, original_numel, 1);
|
|
int grid_size = gpu_config.GetGridSize();
|
|
VLOG(2) << "weight_permute_gpu: original_k = " << original_k
|
|
<< "original_n = " << original_n << "grid size = " << grid_size;
|
|
if (arch > 70) {
|
|
if (algo == "w4a8") {
|
|
dim3 block_dim(128);
|
|
if (original_numel <= std::numeric_limits<int>::max()) {
|
|
weight_permute_transpose_interleave_kernel_w4a8<int>
|
|
<<<grid_size, block_dim>>>(
|
|
input_data, output_data, original_k, original_n);
|
|
} else {
|
|
weight_permute_transpose_interleave_kernel_w4a8<int64_t>
|
|
<<<grid_size, block_dim>>>(
|
|
input_data, output_data, original_k, original_n);
|
|
}
|
|
w4a8_inplace_permute<<<grid_size, block_dim>>>(
|
|
reinterpret_cast<uint32_t*>(output_data), original_numel / 8);
|
|
}
|
|
} else {
|
|
common::errors::Unimplemented(
|
|
"The algo %s support need arch > 70, but got algo = %d.", algo, arch);
|
|
}
|
|
}
|
|
|
|
template <typename IndexT>
|
|
__global__ void weight_permute_interleave_kernelw4afp8(const int8_t* input_data,
|
|
int8_t* output_data,
|
|
IndexT original_k,
|
|
IndexT original_n) {
|
|
IndexT numel = original_k * original_n / 4;
|
|
const IndexT pack_group_size = 64;
|
|
const IndexT thread_group_size = pack_group_size / 4; // 16
|
|
const IndexT thread_k_stride = original_k / 4;
|
|
CUDA_KERNEL_LOOP_TYPE(linear_idx, numel, IndexT) {
|
|
const IndexT n_id = linear_idx / thread_k_stride;
|
|
const IndexT k_id = linear_idx % thread_k_stride;
|
|
const IndexT k_group_idx = k_id / thread_group_size;
|
|
const IndexT k_idx_in_group = k_id % thread_group_size;
|
|
|
|
const int8_t* src = input_data +
|
|
k_group_idx * pack_group_size / 2 * original_n +
|
|
k_idx_in_group * original_n + n_id;
|
|
|
|
int8_t tmp0 = src[0];
|
|
int8_t tmp1 = src[pack_group_size / 4 * original_n];
|
|
|
|
int8_t tmp00 = (tmp0 & 0xF0) + 112;
|
|
int8_t tmp01 = ((tmp0 << 4) & 0xF0) + 112;
|
|
int8_t tmp10 = (tmp1 & 0xF0) + 112;
|
|
int8_t tmp11 = ((tmp1 << 4) & 0xF0) + 112;
|
|
|
|
uint8_t utmp00 = *(reinterpret_cast<uint8_t*>(&tmp00));
|
|
uint8_t utmp01 = *(reinterpret_cast<uint8_t*>(&tmp01));
|
|
uint8_t utmp10 = *(reinterpret_cast<uint8_t*>(&tmp10));
|
|
uint8_t utmp11 = *(reinterpret_cast<uint8_t*>(&tmp11));
|
|
|
|
int8_t dst0 = (utmp01 & 0xF0) | ((utmp11 & 0xF0) >> 4);
|
|
int8_t dst1 = (utmp00 & 0xF0) | ((utmp10 & 0xF0) >> 4);
|
|
|
|
int8_t* dst = output_data + n_id * original_k / 2 +
|
|
(k_group_idx * pack_group_size / 2) + k_idx_in_group * 2;
|
|
dst[0] = dst0;
|
|
dst[1] = dst1;
|
|
}
|
|
}
|
|
|
|
template <typename GPUContext>
|
|
void weight_permute_gpu_w4afp8(const GPUContext& dev_ctx,
|
|
const int8_t* input_data,
|
|
int8_t* output_data,
|
|
const std::vector<int64_t>& shape,
|
|
const int32_t arch,
|
|
const std::string& algo) {
|
|
auto original_k = shape[0] * 2;
|
|
auto original_n = shape[1];
|
|
auto original_numel = original_k * original_n;
|
|
auto gpu_config =
|
|
backends::gpu::GetGpuLaunchConfig1D(dev_ctx, original_numel, 1);
|
|
int grid_size = gpu_config.GetGridSize();
|
|
VLOG(2) << "weight_permute_gpu: original_k = " << original_k
|
|
<< "original_n = " << original_n << "grid size = " << grid_size;
|
|
if (arch > 70) {
|
|
if (algo == "w4afp8") {
|
|
dim3 block_dim(128);
|
|
if (original_numel <= std::numeric_limits<int>::max()) {
|
|
weight_permute_interleave_kernelw4afp8<int><<<grid_size, block_dim>>>(
|
|
input_data, output_data, original_k, original_n);
|
|
} else {
|
|
weight_permute_interleave_kernelw4afp8<int64_t>
|
|
<<<grid_size, block_dim>>>(
|
|
input_data, output_data, original_k, original_n);
|
|
}
|
|
}
|
|
} else {
|
|
common::errors::Unimplemented(
|
|
"The algo %s support need arch > 70, but got algo = %d.", algo, arch);
|
|
}
|
|
}
|
|
|
|
} // namespace phi
|