mirror of
https://github.com/huggingface/text-generation-inference.git
synced 2025-04-24 00:12:08 +00:00
093a27c528
Add support for GPTQ Marlin kernels GPTQ Marlin extends the Marlin kernels to support common GPTQ configurations: - bits: 4 or 8 - groupsize: -1, 32, 64, or 128 - desc_act: true/false Using the GPTQ Marlin kernels requires repacking the parameters in the Marlin quantizer format. The kernels were contributed by Neural Magic to VLLM. We vendor them here for convenience.
351 lines
11 KiB
Plaintext
351 lines
11 KiB
Plaintext
#include "gptq_marlin.cuh"
|
|
|
|
namespace gptq_marlin {
|
|
|
|
static constexpr int repack_stages = 8;
|
|
|
|
static constexpr int repack_threads = 256;
|
|
|
|
static constexpr int tile_k_size = tile_size;
|
|
static constexpr int tile_n_size = tile_k_size * 4;
|
|
|
|
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800
|
|
|
|
template <int const num_threads, int const num_bits, bool const has_perm>
|
|
__global__ void marlin_repack_kernel(
|
|
uint32_t const* __restrict__ b_q_weight_ptr,
|
|
uint32_t const* __restrict__ perm_ptr, uint32_t* __restrict__ out_ptr,
|
|
int size_k, int size_n) {}
|
|
|
|
} // namespace gptq_marlin
|
|
|
|
torch::Tensor gptq_marlin_repack(torch::Tensor& b_q_weight, torch::Tensor& perm,
|
|
int64_t size_k, int64_t size_n,
|
|
int64_t num_bits) {
|
|
TORCH_CHECK_NOT_IMPLEMENTED(
|
|
false, "marlin_repack_from_gptq(..) requires CUDA_ARCH >= 8.0");
|
|
return torch::empty({1, 1});
|
|
}
|
|
|
|
#else
|
|
|
|
template <int const num_threads, int const num_bits, bool const has_perm>
|
|
__global__ void marlin_repack_kernel(
|
|
uint32_t const* __restrict__ b_q_weight_ptr,
|
|
uint32_t const* __restrict__ perm_ptr, uint32_t* __restrict__ out_ptr,
|
|
int size_k, int size_n) {
|
|
constexpr int pack_factor = 32 / num_bits;
|
|
|
|
int k_tiles = size_k / tile_k_size;
|
|
int n_tiles = size_n / tile_n_size;
|
|
int block_k_tiles = div_ceil(k_tiles, gridDim.x);
|
|
|
|
int start_k_tile = blockIdx.x * block_k_tiles;
|
|
if (start_k_tile >= k_tiles) {
|
|
return;
|
|
}
|
|
|
|
int finish_k_tile = min(start_k_tile + block_k_tiles, k_tiles);
|
|
|
|
// Wait until the next thread tile has been loaded to shared memory.
|
|
auto wait_for_stage = [&]() {
|
|
// We only have `stages - 2` active fetches since we are double buffering
|
|
// and can only issue the next fetch when it is guaranteed that the previous
|
|
// shared memory load is fully complete (as it may otherwise be
|
|
// overwritten).
|
|
cp_async_wait<repack_stages - 2>();
|
|
__syncthreads();
|
|
};
|
|
|
|
extern __shared__ int4 sh[];
|
|
|
|
constexpr int perm_size = tile_k_size / 4;
|
|
|
|
int4* sh_perm_ptr = sh;
|
|
int4* sh_pipe_ptr = sh_perm_ptr;
|
|
if constexpr (has_perm) {
|
|
sh_pipe_ptr += perm_size;
|
|
}
|
|
|
|
constexpr int tile_ints = tile_k_size / pack_factor;
|
|
|
|
constexpr int stage_n_threads = tile_n_size / 4;
|
|
constexpr int stage_k_threads = has_perm ? tile_k_size : tile_ints;
|
|
constexpr int stage_size = stage_k_threads * stage_n_threads;
|
|
|
|
auto load_perm_to_shared = [&](int k_tile_id) {
|
|
int first_k_int4 = (k_tile_id * tile_k_size) / 4;
|
|
|
|
int4 const* perm_int4_ptr = reinterpret_cast<int4 const*>(perm_ptr);
|
|
|
|
if (threadIdx.x < perm_size) {
|
|
sh_perm_ptr[threadIdx.x] = perm_int4_ptr[first_k_int4 + threadIdx.x];
|
|
}
|
|
__syncthreads();
|
|
};
|
|
|
|
auto fetch_to_shared = [&](int pipe, int k_tile_id, int n_tile_id) {
|
|
if (n_tile_id >= n_tiles) {
|
|
cp_async_fence();
|
|
return;
|
|
}
|
|
|
|
int first_n = n_tile_id * tile_n_size;
|
|
|
|
int4* sh_ptr = sh_pipe_ptr + stage_size * pipe;
|
|
|
|
if constexpr (has_perm) {
|
|
if (threadIdx.x < stage_size) {
|
|
int k_id = threadIdx.x / stage_n_threads;
|
|
int n_id = threadIdx.x % stage_n_threads;
|
|
|
|
uint32_t const* sh_perm_int_ptr =
|
|
reinterpret_cast<uint32_t const*>(sh_perm_ptr);
|
|
|
|
int src_k = sh_perm_int_ptr[k_id];
|
|
int src_k_packed = src_k / pack_factor;
|
|
|
|
cp_async4(
|
|
&sh_ptr[k_id * stage_n_threads + n_id],
|
|
reinterpret_cast<int4 const*>(&(
|
|
b_q_weight_ptr[src_k_packed * size_n + first_n + (n_id * 4)])));
|
|
}
|
|
|
|
} else {
|
|
if (threadIdx.x < stage_size) {
|
|
int k_id = threadIdx.x / stage_n_threads;
|
|
int n_id = threadIdx.x % stage_n_threads;
|
|
|
|
int first_k = k_tile_id * tile_k_size;
|
|
int first_k_packed = first_k / pack_factor;
|
|
|
|
cp_async4(&sh_ptr[k_id * stage_n_threads + n_id],
|
|
reinterpret_cast<int4 const*>(
|
|
&(b_q_weight_ptr[(first_k_packed + k_id) * size_n +
|
|
first_n + (n_id * 4)])));
|
|
}
|
|
}
|
|
|
|
cp_async_fence();
|
|
};
|
|
|
|
auto repack_tile = [&](int pipe, int k_tile_id, int n_tile_id) {
|
|
if (n_tile_id >= n_tiles) {
|
|
return;
|
|
}
|
|
|
|
int warp_id = threadIdx.x / 32;
|
|
int th_id = threadIdx.x % 32;
|
|
|
|
if (warp_id >= 4) {
|
|
return;
|
|
}
|
|
|
|
int tc_col = th_id / 4;
|
|
int tc_row = (th_id % 4) * 2;
|
|
|
|
constexpr int tc_offsets[4] = {0, 1, 8, 9};
|
|
|
|
int cur_n = warp_id * 16 + tc_col;
|
|
|
|
constexpr int sh_stride = 64;
|
|
constexpr uint32_t mask = (1 << num_bits) - 1;
|
|
|
|
int4* sh_stage_ptr = sh_pipe_ptr + stage_size * pipe;
|
|
uint32_t* sh_stage_int_ptr = reinterpret_cast<uint32_t*>(sh_stage_ptr);
|
|
|
|
uint32_t* sh_perm_int_ptr = reinterpret_cast<uint32_t*>(sh_perm_ptr);
|
|
|
|
uint32_t vals[8];
|
|
|
|
if constexpr (has_perm) {
|
|
for (int i = 0; i < 4; i++) {
|
|
int k_idx = tc_row + tc_offsets[i];
|
|
|
|
uint32_t src_k = sh_perm_int_ptr[k_idx];
|
|
uint32_t src_k_pos = src_k % pack_factor;
|
|
|
|
uint32_t b1_val = sh_stage_int_ptr[k_idx * sh_stride + cur_n];
|
|
uint32_t b1_cur_val = (b1_val >> (src_k_pos * num_bits)) & mask;
|
|
|
|
uint32_t b2_val = sh_stage_int_ptr[k_idx * sh_stride + cur_n + 8];
|
|
uint32_t b2_cur_val = (b2_val >> (src_k_pos * num_bits)) & mask;
|
|
|
|
vals[i] = b1_cur_val;
|
|
vals[4 + i] = b2_cur_val;
|
|
}
|
|
|
|
} else {
|
|
uint32_t b1_vals[tile_ints];
|
|
uint32_t b2_vals[tile_ints];
|
|
|
|
#pragma unroll
|
|
for (int i = 0; i < tile_ints; i++) {
|
|
b1_vals[i] = sh_stage_int_ptr[cur_n + sh_stride * i];
|
|
b2_vals[i] = sh_stage_int_ptr[cur_n + 8 + sh_stride * i];
|
|
}
|
|
|
|
#pragma unroll
|
|
for (int i = 0; i < 4; i++) {
|
|
int cur_elem = tc_row + tc_offsets[i];
|
|
int cur_int = cur_elem / pack_factor;
|
|
int cur_pos = cur_elem % pack_factor;
|
|
|
|
vals[i] = (b1_vals[cur_int] >> (cur_pos * num_bits)) & mask;
|
|
vals[4 + i] = (b2_vals[cur_int] >> (cur_pos * num_bits)) & mask;
|
|
}
|
|
}
|
|
|
|
constexpr int tile_size = tile_k_size * tile_n_size / pack_factor;
|
|
int out_offset = (k_tile_id * n_tiles + n_tile_id) * tile_size;
|
|
|
|
// Result of:
|
|
// https://github.com/NVIDIA/FasterTransformer/blob/main/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h
|
|
if constexpr (num_bits == 4) {
|
|
constexpr int pack_idx[8] = {0, 2, 4, 6, 1, 3, 5, 7};
|
|
|
|
uint32_t res = 0;
|
|
#pragma unroll
|
|
for (int i = 0; i < 8; i++) {
|
|
res |= vals[pack_idx[i]] << (i * 4);
|
|
}
|
|
|
|
out_ptr[out_offset + th_id * 4 + warp_id] = res;
|
|
|
|
} else {
|
|
constexpr int pack_idx[4] = {0, 2, 1, 3};
|
|
|
|
uint32_t res1 = 0;
|
|
uint32_t res2 = 0;
|
|
#pragma unroll
|
|
for (int i = 0; i < 4; i++) {
|
|
res1 |= vals[pack_idx[i]] << (i * 8);
|
|
res2 |= vals[4 + pack_idx[i]] << (i * 8);
|
|
}
|
|
|
|
out_ptr[out_offset + th_id * 8 + (warp_id * 2) + 0] = res1;
|
|
out_ptr[out_offset + th_id * 8 + (warp_id * 2) + 1] = res2;
|
|
}
|
|
};
|
|
|
|
auto start_pipes = [&](int k_tile_id, int n_tile_id) {
|
|
#pragma unroll
|
|
for (int pipe = 0; pipe < repack_stages - 1; pipe++) {
|
|
fetch_to_shared(pipe, k_tile_id, n_tile_id + pipe);
|
|
}
|
|
|
|
wait_for_stage();
|
|
};
|
|
#pragma unroll
|
|
for (int k_tile_id = start_k_tile; k_tile_id < finish_k_tile; k_tile_id++) {
|
|
int n_tile_id = 0;
|
|
|
|
if constexpr (has_perm) {
|
|
load_perm_to_shared(k_tile_id);
|
|
}
|
|
|
|
start_pipes(k_tile_id, n_tile_id);
|
|
|
|
while (n_tile_id < n_tiles) {
|
|
#pragma unroll
|
|
for (int pipe = 0; pipe < repack_stages; pipe++) {
|
|
fetch_to_shared((pipe + repack_stages - 1) % repack_stages, k_tile_id,
|
|
n_tile_id + pipe + repack_stages - 1);
|
|
repack_tile(pipe, k_tile_id, n_tile_id + pipe);
|
|
wait_for_stage();
|
|
}
|
|
n_tile_id += repack_stages;
|
|
}
|
|
}
|
|
}
|
|
|
|
} // namespace gptq_marlin
|
|
|
|
#define CALL_IF(NUM_BITS, HAS_PERM) \
|
|
else if (num_bits == NUM_BITS && has_perm == HAS_PERM) { \
|
|
cudaFuncSetAttribute( \
|
|
gptq_marlin::marlin_repack_kernel<gptq_marlin::repack_threads, \
|
|
NUM_BITS, HAS_PERM>, \
|
|
cudaFuncAttributeMaxDynamicSharedMemorySize, max_shared_mem); \
|
|
gptq_marlin::marlin_repack_kernel<gptq_marlin::repack_threads, NUM_BITS, \
|
|
HAS_PERM> \
|
|
<<<blocks, gptq_marlin::repack_threads, max_shared_mem, stream>>>( \
|
|
b_q_weight_ptr, perm_ptr, out_ptr, size_k, size_n); \
|
|
}
|
|
|
|
torch::Tensor gptq_marlin_repack(torch::Tensor& b_q_weight, torch::Tensor& perm,
|
|
int64_t size_k, int64_t size_n,
|
|
int64_t num_bits) {
|
|
// Verify compatibility with marlin tile of 16x64
|
|
TORCH_CHECK(size_k % gptq_marlin::tile_k_size == 0, "size_k = ", size_k,
|
|
" is not divisible by tile_k_size = ", gptq_marlin::tile_k_size);
|
|
TORCH_CHECK(size_n % gptq_marlin::tile_n_size == 0, "size_n = ", size_n,
|
|
" is not divisible by tile_n_size = ", gptq_marlin::tile_n_size);
|
|
|
|
TORCH_CHECK(num_bits == 4 || num_bits == 8,
|
|
"num_bits must be 4 or 8. Got = ", num_bits);
|
|
int const pack_factor = 32 / num_bits;
|
|
|
|
// Verify B
|
|
TORCH_CHECK((size_k / pack_factor) == b_q_weight.size(0),
|
|
"Shape mismatch: b_q_weight.size(0) = ", b_q_weight.size(0),
|
|
", size_k = ", size_k, ", pack_factor = ", pack_factor);
|
|
TORCH_CHECK(b_q_weight.size(1) == size_n,
|
|
"b_q_weight.size(1) = ", b_q_weight.size(1),
|
|
" is not size_n = ", size_n);
|
|
|
|
// Verify device and strides
|
|
TORCH_CHECK(b_q_weight.device().is_cuda(), "b_q_weight is not on GPU");
|
|
TORCH_CHECK(b_q_weight.is_contiguous(), "b_q_weight is not contiguous");
|
|
TORCH_CHECK(b_q_weight.dtype() == at::kInt, "b_q_weight type is not kInt");
|
|
|
|
TORCH_CHECK(perm.device().is_cuda(), "perm is not on GPU");
|
|
TORCH_CHECK(perm.is_contiguous(), "perm is not contiguous");
|
|
TORCH_CHECK(perm.dtype() == at::kInt, "perm type is not at::kInt");
|
|
|
|
// Alloc buffers
|
|
const at::cuda::OptionalCUDAGuard device_guard(device_of(b_q_weight));
|
|
auto options = torch::TensorOptions()
|
|
.dtype(b_q_weight.dtype())
|
|
.device(b_q_weight.device());
|
|
torch::Tensor out =
|
|
torch::empty({size_k / gptq_marlin::tile_size,
|
|
size_n * gptq_marlin::tile_size / pack_factor},
|
|
options);
|
|
|
|
// Detect if there is act_order
|
|
bool has_perm = perm.size(0) != 0;
|
|
|
|
// Get ptrs
|
|
uint32_t const* b_q_weight_ptr =
|
|
reinterpret_cast<uint32_t const*>(b_q_weight.data_ptr());
|
|
uint32_t const* perm_ptr = reinterpret_cast<uint32_t const*>(perm.data_ptr());
|
|
uint32_t* out_ptr = reinterpret_cast<uint32_t*>(out.data_ptr());
|
|
|
|
// Get dev info
|
|
int dev = b_q_weight.get_device();
|
|
cudaStream_t stream = at::cuda::getCurrentCUDAStream(dev);
|
|
int blocks;
|
|
cudaDeviceGetAttribute(&blocks, cudaDevAttrMultiProcessorCount, dev);
|
|
|
|
int max_shared_mem = 0;
|
|
cudaDeviceGetAttribute(&max_shared_mem,
|
|
cudaDevAttrMaxSharedMemoryPerBlockOptin, dev);
|
|
TORCH_CHECK(max_shared_mem > 0);
|
|
|
|
if (false) {
|
|
}
|
|
CALL_IF(4, false)
|
|
CALL_IF(4, true)
|
|
CALL_IF(8, false)
|
|
CALL_IF(8, true)
|
|
else {
|
|
TORCH_CHECK(false, "Unsupported repack config: num_bits = ", num_bits,
|
|
", has_perm = ", has_perm);
|
|
}
|
|
|
|
return out;
|
|
}
|
|
|
|
#endif
|