add patch for gptq and peft
This commit is contained in:
John Smith
GitHub
144
GPTQ-for-LLaMa/autograd_4bit.py
Normal file
144
GPTQ-for-LLaMa/autograd_4bit.py
Normal file
@@ -0,0 +1,144 @@
|
||||
import quant
|
||||
import torch
|
||||
import numpy as np
|
||||
import torch.nn as nn
|
||||
|
||||
|
||||
def matmul4bit(x, qweight, scales, zeros):
|
||||
"""
|
||||
input x: (n, m)
|
||||
qweight: (j, k)
|
||||
where m == j*8
|
||||
|
||||
perform x @ qweight
|
||||
|
||||
return y:
|
||||
"""
|
||||
assert qweight.shape[0] * 8 == x.shape[-1]
|
||||
outshape = tuple(list(x.shape[:-1]) + [qweight.shape[1]])
|
||||
x = x.reshape(-1, x.shape[-1])
|
||||
assert x.shape[0] % 256 == 0
|
||||
y = torch.zeros((x.shape[0], qweight.shape[-1]), dtype=torch.float32, device=x.device)
|
||||
dtype = x.dtype
|
||||
x = x.float()
|
||||
quant.quant_cuda.vecquant4matmul(x, qweight, y, scales, zeros)
|
||||
y = y.to(dtype)
|
||||
return y.reshape(outshape)
|
||||
|
||||
|
||||
def matmul4bit_transpose(x, qweight, scales, zeros):
|
||||
"""
|
||||
input x: (n, m)
|
||||
qweight: (j, k)
|
||||
where m == k
|
||||
|
||||
perform qweight @ x.T
|
||||
|
||||
return y:
|
||||
"""
|
||||
assert qweight.shape[1] == x.shape[-1]
|
||||
outshape = tuple(list(x.shape[:-1]) + [qweight.shape[0] * 8])
|
||||
x = x.reshape(-1, x.shape[-1])
|
||||
assert x.shape[0] % 256 == 0
|
||||
y = torch.zeros((qweight.shape[0] * 8, x.shape[0]), dtype=torch.float32, device=x.device)
|
||||
dtype = x.dtype
|
||||
x = x.float()
|
||||
quant.quant_cuda.vecquant4transposematmul(x, qweight, y, scales, zeros)
|
||||
y = y.to(dtype)
|
||||
return y.reshape(outshape)
|
||||
|
||||
|
||||
class AutogradMatmul4bit(torch.autograd.Function):
|
||||
|
||||
@staticmethod
|
||||
def forward(ctx, x, qweight, scales, zeros):
|
||||
ctx.save_for_backward(x, qweight, scales, zeros)
|
||||
output = matmul4bit(x, qweight, scales, zeros).clone()
|
||||
return output # equals to torch.matmul(x, qweight)
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, grad_output):
|
||||
x, qweight, scales, zeros = ctx.saved_tensors
|
||||
# print(grad_output.shape, A.shape, B.shape)
|
||||
|
||||
# compute x @ qweight.T = (qweight @ x.T).T = f(x, qweight).T
|
||||
grad1 = matmul4bit_transpose(grad_output, qweight, scales, zeros)
|
||||
grad2 = torch.matmul(x.transpose(-1, -2), grad_output)
|
||||
|
||||
return grad1, grad2, None, None
|
||||
|
||||
|
||||
# Assumes layer is perfectly divisible into 256 * 256 blocks
|
||||
class Autograd4bitQuantLinear(nn.Module):
|
||||
|
||||
def __init__(self, infeatures, outfeatures):
|
||||
super().__init__()
|
||||
bits = 4
|
||||
self.in_features = infeatures
|
||||
self.out_features = outfeatures
|
||||
self.bits = bits
|
||||
self.register_buffer('zeros', torch.empty((outfeatures, 1)))
|
||||
self.register_buffer('scales', torch.empty((outfeatures, 1)))
|
||||
self.register_buffer('bias', torch.empty(outfeatures))
|
||||
self.register_buffer(
|
||||
'qweight', torch.empty((infeatures // 256 * (bits * 8), outfeatures), dtype=torch.int)
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
out = AutogradMatmul4bit.apply(x, self.qweight, self.scales, self.zeros)
|
||||
out += self.bias
|
||||
return out
|
||||
|
||||
|
||||
def make_quant_for_4bit_autograd(module, names, name=''):
|
||||
if isinstance(module, Autograd4bitQuantLinear):
|
||||
return
|
||||
for attr in dir(module):
|
||||
tmp = getattr(module, attr)
|
||||
name1 = name + '.' + attr if name != '' else attr
|
||||
if name1 in names:
|
||||
setattr(
|
||||
module, attr, Autograd4bitQuantLinear(tmp.in_features, tmp.out_features)
|
||||
)
|
||||
for name1, child in module.named_children():
|
||||
make_quant_for_4bit_autograd(child, names, name + '.' + name1 if name != '' else name1)
|
||||
|
||||
|
||||
def load_llama_model_4bit_low_ram(config_path, model_path):
|
||||
import transformers
|
||||
import accelerate
|
||||
from transformers import LLaMAConfig, LLaMAForCausalLM, LLaMATokenizer
|
||||
from modelutils import find_layers
|
||||
|
||||
print("Loading Model ...")
|
||||
t0 = time.time()
|
||||
|
||||
with accelerate.init_empty_weights():
|
||||
config = LLaMAConfig.from_pretrained(config_path)
|
||||
def noop(*args, **kwargs):
|
||||
pass
|
||||
torch.nn.init.kaiming_uniform_ = noop
|
||||
torch.nn.init.uniform_ = noop
|
||||
torch.nn.init.normal_ = noop
|
||||
torch.set_default_dtype(torch.half)
|
||||
transformers.modeling_utils._init_weights = False
|
||||
torch.set_default_dtype(torch.half)
|
||||
model = LLaMAForCausalLM(config)
|
||||
torch.set_default_dtype(torch.float)
|
||||
model = model.eval()
|
||||
layers = find_layers(model)
|
||||
for name in ['lm_head']:
|
||||
if name in layers:
|
||||
del layers[name]
|
||||
make_quant_for_4bit_autograd(model, layers)
|
||||
model = accelerate.load_checkpoint_and_dispatch(model=model, checkpoint=model_path, device_map='auto')
|
||||
model.cuda()
|
||||
model.seqlen = 2048
|
||||
|
||||
tokenizer = LLaMATokenizer.from_pretrained(config_path)
|
||||
tokenizer.truncation_side = 'left'
|
||||
|
||||
print(f"Loaded the model in {(time.time()-t0):.2f} seconds.")
|
||||
|
||||
return model, tokenizer
|
||||
|
||||
76
GPTQ-for-LLaMa/quant_cuda.cpp
Normal file
76
GPTQ-for-LLaMa/quant_cuda.cpp
Normal file
@@ -0,0 +1,76 @@
|
||||
#include <torch/all.h>
|
||||
#include <torch/python.h>
|
||||
#include <c10/cuda/CUDAGuard.h>
|
||||
|
||||
void vecquant2matmul_cuda(
|
||||
torch::Tensor vec, torch::Tensor mat, torch::Tensor mul,
|
||||
torch::Tensor scales, torch::Tensor zeros
|
||||
);
|
||||
|
||||
void vecquant2matmul(
|
||||
torch::Tensor vec, torch::Tensor mat, torch::Tensor mul,
|
||||
torch::Tensor scales, torch::Tensor zeros
|
||||
) {
|
||||
const at::cuda::OptionalCUDAGuard device_guard(device_of(vec));
|
||||
vecquant2matmul_cuda(vec, mat, mul, scales, zeros);
|
||||
}
|
||||
|
||||
void vecquant3matmul_cuda(
|
||||
torch::Tensor vec, torch::Tensor mat, torch::Tensor mul,
|
||||
torch::Tensor scales, torch::Tensor zeros
|
||||
);
|
||||
|
||||
void vecquant3matmul(
|
||||
torch::Tensor vec, torch::Tensor mat, torch::Tensor mul,
|
||||
torch::Tensor scales, torch::Tensor zeros
|
||||
) {
|
||||
const at::cuda::OptionalCUDAGuard device_guard(device_of(vec));
|
||||
vecquant3matmul_cuda(vec, mat, mul, scales, zeros);
|
||||
}
|
||||
|
||||
void vecquant4matmul_cuda(
|
||||
torch::Tensor vec, torch::Tensor mat, torch::Tensor mul,
|
||||
torch::Tensor scales, torch::Tensor zeros
|
||||
);
|
||||
|
||||
void vecquant4matmul(
|
||||
torch::Tensor vec, torch::Tensor mat, torch::Tensor mul,
|
||||
torch::Tensor scales, torch::Tensor zeros
|
||||
) {
|
||||
const at::cuda::OptionalCUDAGuard device_guard(device_of(vec));
|
||||
vecquant4matmul_cuda(vec, mat, mul, scales, zeros);
|
||||
}
|
||||
|
||||
void vecquant8matmul_cuda(
|
||||
torch::Tensor vec, torch::Tensor mat, torch::Tensor mul,
|
||||
torch::Tensor scales, torch::Tensor zeros
|
||||
);
|
||||
|
||||
void vecquant8matmul(
|
||||
torch::Tensor vec, torch::Tensor mat, torch::Tensor mul,
|
||||
torch::Tensor scales, torch::Tensor zeros
|
||||
) {
|
||||
const at::cuda::OptionalCUDAGuard device_guard(device_of(vec));
|
||||
vecquant8matmul_cuda(vec, mat, mul, scales, zeros);
|
||||
}
|
||||
|
||||
void vecquant4transposematmul_cuda(
|
||||
torch::Tensor vec, torch::Tensor mat, torch::Tensor mul,
|
||||
torch::Tensor scales, torch::Tensor zeros
|
||||
);
|
||||
|
||||
void vecquant4transposematmul(
|
||||
torch::Tensor vec, torch::Tensor mat, torch::Tensor mul,
|
||||
torch::Tensor scales, torch::Tensor zeros
|
||||
) {
|
||||
const at::cuda::OptionalCUDAGuard device_guard(device_of(vec));
|
||||
vecquant4transposematmul_cuda(vec, mat, mul, scales, zeros);
|
||||
}
|
||||
|
||||
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
|
||||
m.def("vecquant2matmul", &vecquant2matmul, "Vector 2-bit Quantized Matrix Multiplication (CUDA)");
|
||||
m.def("vecquant3matmul", &vecquant3matmul, "Vector 3-bit Quantized Matrix Multiplication (CUDA)");
|
||||
m.def("vecquant4matmul", &vecquant4matmul, "Vector 4-bit Quantized Matrix Multiplication (CUDA)");
|
||||
m.def("vecquant8matmul", &vecquant8matmul, "Vector 8-bit Quantized Matrix Multiplication (CUDA)");
|
||||
m.def("vecquant4transposematmul", &vecquant4transposematmul, "Vector 4-bit Transpose Quantized Matrix Multiplication (CUDA)");
|
||||
}
|
||||
480
GPTQ-for-LLaMa/quant_cuda_kernel.cu
Normal file
480
GPTQ-for-LLaMa/quant_cuda_kernel.cu
Normal file
@@ -0,0 +1,480 @@
|
||||
#include <torch/all.h>
|
||||
#include <torch/python.h>
|
||||
#include <cuda.h>
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
template <typename scalar_t>
|
||||
__global__ void VecQuant2MatMulKernel(
|
||||
const scalar_t* __restrict__ vec,
|
||||
const int* __restrict__ mat,
|
||||
scalar_t* __restrict__ mul,
|
||||
const scalar_t* __restrict__ scales,
|
||||
const scalar_t* __restrict__ zeros,
|
||||
int batch,
|
||||
int vec_height,
|
||||
int height,
|
||||
int width
|
||||
);
|
||||
|
||||
template <typename scalar_t>
|
||||
__global__ void VecQuant3MatMulKernel(
|
||||
const scalar_t* __restrict__ vec,
|
||||
const int* __restrict__ mat,
|
||||
scalar_t* __restrict__ mul,
|
||||
const scalar_t* __restrict__ scales,
|
||||
const scalar_t* __restrict__ zeros,
|
||||
int batch,
|
||||
int vec_height,
|
||||
int height,
|
||||
int width
|
||||
);
|
||||
|
||||
template <typename scalar_t>
|
||||
__global__ void VecQuant4MatMulKernel(
|
||||
const scalar_t* __restrict__ vec,
|
||||
const int* __restrict__ mat,
|
||||
scalar_t* __restrict__ mul,
|
||||
const scalar_t* __restrict__ scales,
|
||||
const scalar_t* __restrict__ zeros,
|
||||
int batch,
|
||||
int vec_height,
|
||||
int height,
|
||||
int width
|
||||
);
|
||||
|
||||
template <typename scalar_t>
|
||||
__global__ void VecQuant8MatMulKernel(
|
||||
const scalar_t* __restrict__ vec,
|
||||
const int* __restrict__ mat,
|
||||
scalar_t* __restrict__ mul,
|
||||
const scalar_t* __restrict__ scales,
|
||||
const scalar_t* __restrict__ zeros,
|
||||
int batch,
|
||||
int vec_height,
|
||||
int height,
|
||||
int width
|
||||
);
|
||||
|
||||
const int BLOCKWIDTH = 256;
|
||||
const int BLOCKHEIGHT2 = 16;
|
||||
const int BLOCKHEIGHT3 = 24;
|
||||
const int BLOCKHEIGHT4 = 32;
|
||||
const int BLOCKHEIGHT8 = 64;
|
||||
|
||||
__device__ inline unsigned int as_unsigned(int i) {
|
||||
return *reinterpret_cast<unsigned int*>(&i);
|
||||
}
|
||||
|
||||
void vecquant2matmul_cuda(
|
||||
torch::Tensor vec,
|
||||
torch::Tensor mat,
|
||||
torch::Tensor mul,
|
||||
torch::Tensor scales,
|
||||
torch::Tensor zeros
|
||||
) {
|
||||
int batch = vec.size(0);
|
||||
int vec_height = vec.size(1);
|
||||
int height = mat.size(0);
|
||||
int width = mat.size(1);
|
||||
|
||||
dim3 blocks(
|
||||
(height + BLOCKHEIGHT2 - 1) / BLOCKHEIGHT2,
|
||||
(width + BLOCKWIDTH - 1) / BLOCKWIDTH,
|
||||
batch
|
||||
);
|
||||
dim3 threads(BLOCKWIDTH);
|
||||
|
||||
AT_DISPATCH_FLOATING_TYPES(
|
||||
vec.type(), "vecquant2matmul_cuda", ([&] {
|
||||
VecQuant2MatMulKernel<<<blocks, threads>>>(
|
||||
vec.data<scalar_t>(), mat.data<int>(), mul.data<scalar_t>(),
|
||||
scales.data<scalar_t>(), zeros.data<scalar_t>(),
|
||||
batch, vec_height, height, width
|
||||
);
|
||||
})
|
||||
);
|
||||
}
|
||||
|
||||
template <typename scalar_t>
|
||||
__global__ void VecQuant2MatMulKernel(
|
||||
const scalar_t* __restrict__ vec,
|
||||
const int* __restrict__ mat,
|
||||
scalar_t* __restrict__ mul,
|
||||
const scalar_t* __restrict__ scales,
|
||||
const scalar_t* __restrict__ zeros,
|
||||
int batch,
|
||||
int vec_height,
|
||||
int height,
|
||||
int width
|
||||
) {
|
||||
int b = blockIdx.z;
|
||||
int h = BLOCKHEIGHT2 * blockIdx.x;
|
||||
int w = BLOCKWIDTH * blockIdx.y + threadIdx.x;
|
||||
|
||||
__shared__ scalar_t blockvec[BLOCKWIDTH];
|
||||
blockvec[threadIdx.x] = vec[b * vec_height + (h / BLOCKHEIGHT2) * BLOCKWIDTH + threadIdx.x];
|
||||
__syncthreads();
|
||||
|
||||
scalar_t scale = scales[w];
|
||||
scalar_t zero = zeros[w];
|
||||
|
||||
scalar_t res = 0;
|
||||
int i = width * h + w;
|
||||
int k = 0;
|
||||
|
||||
unsigned int tmp;
|
||||
|
||||
while (k < BLOCKWIDTH) {
|
||||
tmp = as_unsigned(mat[i]);
|
||||
res += (scale * scalar_t((tmp >> 0) & 0x3) - zero) * blockvec[k + 0];
|
||||
res += (scale * scalar_t((tmp >> 2) & 0x3) - zero) * blockvec[k + 1];
|
||||
res += (scale * scalar_t((tmp >> 4) & 0x3) - zero) * blockvec[k + 2];
|
||||
res += (scale * scalar_t((tmp >> 6) & 0x3) - zero) * blockvec[k + 3];
|
||||
res += (scale * scalar_t((tmp >> 8) & 0x3) - zero) * blockvec[k + 4];
|
||||
res += (scale * scalar_t((tmp >> 10) & 0x3) - zero) * blockvec[k + 5];
|
||||
res += (scale * scalar_t((tmp >> 12) & 0x3) - zero) * blockvec[k + 6];
|
||||
res += (scale * scalar_t((tmp >> 14) & 0x3) - zero) * blockvec[k + 7];
|
||||
res += (scale * scalar_t((tmp >> 16) & 0x3) - zero) * blockvec[k + 8];
|
||||
res += (scale * scalar_t((tmp >> 18) & 0x3) - zero) * blockvec[k + 9];
|
||||
res += (scale * scalar_t((tmp >> 20) & 0x3) - zero) * blockvec[k + 10];
|
||||
res += (scale * scalar_t((tmp >> 22) & 0x3) - zero) * blockvec[k + 11];
|
||||
res += (scale * scalar_t((tmp >> 24) & 0x3) - zero) * blockvec[k + 12];
|
||||
res += (scale * scalar_t((tmp >> 26) & 0x3) - zero) * blockvec[k + 13];
|
||||
res += (scale * scalar_t((tmp >> 28) & 0x3) - zero) * blockvec[k + 14];
|
||||
res += (scale * scalar_t((tmp >> 30) & 0x3) - zero) * blockvec[k + 15];
|
||||
i += width;
|
||||
k += 16;
|
||||
}
|
||||
|
||||
atomicAdd(&mul[b * width + w], res);
|
||||
}
|
||||
|
||||
void vecquant3matmul_cuda(
|
||||
torch::Tensor vec,
|
||||
torch::Tensor mat,
|
||||
torch::Tensor mul,
|
||||
torch::Tensor scales,
|
||||
torch::Tensor zeros
|
||||
) {
|
||||
int batch = vec.size(0);
|
||||
int vec_height = vec.size(1);
|
||||
int height = mat.size(0);
|
||||
int width = mat.size(1);
|
||||
|
||||
dim3 blocks(
|
||||
(height + BLOCKHEIGHT3 - 1) / BLOCKHEIGHT3,
|
||||
(width + BLOCKWIDTH - 1) / BLOCKWIDTH,
|
||||
batch
|
||||
);
|
||||
dim3 threads(BLOCKWIDTH);
|
||||
|
||||
AT_DISPATCH_FLOATING_TYPES(
|
||||
vec.type(), "vecquant3matmul_cuda", ([&] {
|
||||
VecQuant3MatMulKernel<<<blocks, threads>>>(
|
||||
vec.data<scalar_t>(), mat.data<int>(), mul.data<scalar_t>(),
|
||||
scales.data<scalar_t>(), zeros.data<scalar_t>(),
|
||||
batch, vec_height, height, width
|
||||
);
|
||||
})
|
||||
);
|
||||
}
|
||||
|
||||
template <typename scalar_t>
|
||||
__global__ void VecQuant3MatMulKernel(
|
||||
const scalar_t* __restrict__ vec,
|
||||
const int* __restrict__ mat,
|
||||
scalar_t* __restrict__ mul,
|
||||
const scalar_t* __restrict__ scales,
|
||||
const scalar_t* __restrict__ zeros,
|
||||
int batch,
|
||||
int vec_height,
|
||||
int height,
|
||||
int width
|
||||
) {
|
||||
int b = blockIdx.z;
|
||||
int h = BLOCKHEIGHT3 * blockIdx.x;
|
||||
int w = BLOCKWIDTH * blockIdx.y + threadIdx.x;
|
||||
|
||||
__shared__ scalar_t blockvec[BLOCKWIDTH];
|
||||
blockvec[threadIdx.x] = vec[b * vec_height + (h / BLOCKHEIGHT3) * BLOCKWIDTH + threadIdx.x];
|
||||
__syncthreads();
|
||||
|
||||
scalar_t scale = scales[w];
|
||||
scalar_t zero = zeros[w];
|
||||
|
||||
scalar_t res = 0;
|
||||
int i = width * h + w;
|
||||
int k = 0;
|
||||
|
||||
unsigned int tmp1;
|
||||
unsigned int tmp2;
|
||||
unsigned int tmp;
|
||||
|
||||
while (k < BLOCKWIDTH) {
|
||||
tmp1 = as_unsigned(mat[i]);
|
||||
res += (scale * scalar_t((tmp1 >> 0) & 0x7) - zero) * blockvec[k + 0];
|
||||
res += (scale * scalar_t((tmp1 >> 3) & 0x7) - zero) * blockvec[k + 1];
|
||||
res += (scale * scalar_t((tmp1 >> 6) & 0x7) - zero) * blockvec[k + 2];
|
||||
res += (scale * scalar_t((tmp1 >> 9) & 0x7) - zero) * blockvec[k + 3];
|
||||
res += (scale * scalar_t((tmp1 >> 12) & 0x7) - zero) * blockvec[k + 4];
|
||||
res += (scale * scalar_t((tmp1 >> 15) & 0x7) - zero) * blockvec[k + 5];
|
||||
res += (scale * scalar_t((tmp1 >> 18) & 0x7) - zero) * blockvec[k + 6];
|
||||
res += (scale * scalar_t((tmp1 >> 21) & 0x7) - zero) * blockvec[k + 7];
|
||||
res += (scale * scalar_t((tmp1 >> 24) & 0x7) - zero) * blockvec[k + 8];
|
||||
res += (scale * scalar_t((tmp1 >> 27) & 0x7) - zero) * blockvec[k + 9];
|
||||
i += width;
|
||||
tmp2 = as_unsigned(mat[i]);
|
||||
tmp = (tmp1 >> 30) | ((tmp2 << 2) & 0x4);
|
||||
tmp2 >>= 1;
|
||||
res += (scale * scalar_t(tmp) - zero) * blockvec[k + 10];
|
||||
k += 11;
|
||||
res += (scale * scalar_t((tmp2 >> 0) & 0x7) - zero) * blockvec[k + 0];
|
||||
res += (scale * scalar_t((tmp2 >> 3) & 0x7) - zero) * blockvec[k + 1];
|
||||
res += (scale * scalar_t((tmp2 >> 6) & 0x7) - zero) * blockvec[k + 2];
|
||||
res += (scale * scalar_t((tmp2 >> 9) & 0x7) - zero) * blockvec[k + 3];
|
||||
res += (scale * scalar_t((tmp2 >> 12) & 0x7) - zero) * blockvec[k + 4];
|
||||
res += (scale * scalar_t((tmp2 >> 15) & 0x7) - zero) * blockvec[k + 5];
|
||||
res += (scale * scalar_t((tmp2 >> 18) & 0x7) - zero) * blockvec[k + 6];
|
||||
res += (scale * scalar_t((tmp2 >> 21) & 0x7) - zero) * blockvec[k + 7];
|
||||
res += (scale * scalar_t((tmp2 >> 24) & 0x7) - zero) * blockvec[k + 8];
|
||||
res += (scale * scalar_t((tmp2 >> 27) & 0x7) - zero) * blockvec[k + 9];
|
||||
i += width;
|
||||
tmp1 = as_unsigned(mat[i]);
|
||||
tmp = (tmp2 >> 30) | ((tmp1 << 1) & 0x6);
|
||||
tmp1 >>= 2;
|
||||
res += (scale * scalar_t(tmp) - zero) * blockvec[k + 10];
|
||||
k += 11;
|
||||
res += (scale * scalar_t((tmp1 >> 0) & 0x7) - zero) * blockvec[k + 0];
|
||||
res += (scale * scalar_t((tmp1 >> 3) & 0x7) - zero) * blockvec[k + 1];
|
||||
res += (scale * scalar_t((tmp1 >> 6) & 0x7) - zero) * blockvec[k + 2];
|
||||
res += (scale * scalar_t((tmp1 >> 9) & 0x7) - zero) * blockvec[k + 3];
|
||||
res += (scale * scalar_t((tmp1 >> 12) & 0x7) - zero) * blockvec[k + 4];
|
||||
res += (scale * scalar_t((tmp1 >> 15) & 0x7) - zero) * blockvec[k + 5];
|
||||
res += (scale * scalar_t((tmp1 >> 18) & 0x7) - zero) * blockvec[k + 6];
|
||||
res += (scale * scalar_t((tmp1 >> 21) & 0x7) - zero) * blockvec[k + 7];
|
||||
res += (scale * scalar_t((tmp1 >> 24) & 0x7) - zero) * blockvec[k + 8];
|
||||
res += (scale * scalar_t((tmp1 >> 27) & 0x7) - zero) * blockvec[k + 9];
|
||||
i += width;
|
||||
k += 10;
|
||||
}
|
||||
|
||||
atomicAdd(&mul[b * width + w], res);
|
||||
}
|
||||
|
||||
void vecquant4matmul_cuda(
|
||||
torch::Tensor vec,
|
||||
torch::Tensor mat,
|
||||
torch::Tensor mul,
|
||||
torch::Tensor scales,
|
||||
torch::Tensor zeros
|
||||
) {
|
||||
int batch = vec.size(0);
|
||||
int vec_height = vec.size(1);
|
||||
int height = mat.size(0);
|
||||
int width = mat.size(1);
|
||||
|
||||
dim3 blocks(
|
||||
(height + BLOCKHEIGHT4 - 1) / BLOCKHEIGHT4,
|
||||
(width + BLOCKWIDTH - 1) / BLOCKWIDTH,
|
||||
batch
|
||||
);
|
||||
dim3 threads(BLOCKWIDTH);
|
||||
|
||||
AT_DISPATCH_FLOATING_TYPES(
|
||||
vec.type(), "vecquant4matmul_cuda", ([&] {
|
||||
VecQuant4MatMulKernel<<<blocks, threads>>>(
|
||||
vec.data<scalar_t>(), mat.data<int>(), mul.data<scalar_t>(),
|
||||
scales.data<scalar_t>(), zeros.data<scalar_t>(),
|
||||
batch, vec_height, height, width
|
||||
);
|
||||
})
|
||||
);
|
||||
}
|
||||
|
||||
template <typename scalar_t>
|
||||
__global__ void VecQuant4MatMulKernel(
|
||||
const scalar_t* __restrict__ vec,
|
||||
const int* __restrict__ mat,
|
||||
scalar_t* __restrict__ mul,
|
||||
const scalar_t* __restrict__ scales,
|
||||
const scalar_t* __restrict__ zeros,
|
||||
int batch,
|
||||
int vec_height,
|
||||
int height,
|
||||
int width
|
||||
) {
|
||||
int b = blockIdx.z;
|
||||
int h = BLOCKHEIGHT4 * blockIdx.x;
|
||||
int w = BLOCKWIDTH * blockIdx.y + threadIdx.x;
|
||||
|
||||
__shared__ scalar_t blockvec[BLOCKWIDTH];
|
||||
blockvec[threadIdx.x] = vec[b * vec_height + (h / BLOCKHEIGHT4) * BLOCKWIDTH + threadIdx.x];
|
||||
__syncthreads();
|
||||
|
||||
scalar_t scale = scales[w];
|
||||
scalar_t zero = zeros[w];
|
||||
|
||||
scalar_t res = 0;
|
||||
int i = width * h + w;
|
||||
int k = 0;
|
||||
|
||||
unsigned int tmp;
|
||||
|
||||
while (k < BLOCKWIDTH) {
|
||||
tmp = as_unsigned(mat[i]);
|
||||
res += (scale * scalar_t((tmp >> 0) & 0xF) - zero) * blockvec[k + 0];
|
||||
res += (scale * scalar_t((tmp >> 4) & 0xF) - zero) * blockvec[k + 1];
|
||||
res += (scale * scalar_t((tmp >> 8) & 0xF) - zero) * blockvec[k + 2];
|
||||
res += (scale * scalar_t((tmp >> 12) & 0xF) - zero) * blockvec[k + 3];
|
||||
res += (scale * scalar_t((tmp >> 16) & 0xF) - zero) * blockvec[k + 4];
|
||||
res += (scale * scalar_t((tmp >> 20) & 0xF) - zero) * blockvec[k + 5];
|
||||
res += (scale * scalar_t((tmp >> 24) & 0xF) - zero) * blockvec[k + 6];
|
||||
res += (scale * scalar_t((tmp >> 28) & 0xF) - zero) * blockvec[k + 7];
|
||||
i += width;
|
||||
k += 8;
|
||||
}
|
||||
|
||||
atomicAdd(&mul[b * width + w], res);
|
||||
}
|
||||
|
||||
void vecquant8matmul_cuda(
|
||||
torch::Tensor vec,
|
||||
torch::Tensor mat,
|
||||
torch::Tensor mul,
|
||||
torch::Tensor scales,
|
||||
torch::Tensor zeros
|
||||
) {
|
||||
int batch = vec.size(0);
|
||||
int vec_height = vec.size(1);
|
||||
int height = mat.size(0);
|
||||
int width = mat.size(1);
|
||||
|
||||
dim3 blocks(
|
||||
(height + BLOCKHEIGHT8 - 1) / BLOCKHEIGHT8,
|
||||
(width + BLOCKWIDTH - 1) / BLOCKWIDTH,
|
||||
batch
|
||||
);
|
||||
dim3 threads(BLOCKWIDTH);
|
||||
|
||||
AT_DISPATCH_FLOATING_TYPES(
|
||||
vec.type(), "vecquant8matmul_cuda", ([&] {
|
||||
VecQuant8MatMulKernel<<<blocks, threads>>>(
|
||||
vec.data<scalar_t>(), mat.data<int>(), mul.data<scalar_t>(),
|
||||
scales.data<scalar_t>(), zeros.data<scalar_t>(),
|
||||
batch, vec_height, height, width
|
||||
);
|
||||
})
|
||||
);
|
||||
}
|
||||
|
||||
template <typename scalar_t>
|
||||
__global__ void VecQuant8MatMulKernel(
|
||||
const scalar_t* __restrict__ vec,
|
||||
const int* __restrict__ mat,
|
||||
scalar_t* __restrict__ mul,
|
||||
const scalar_t* __restrict__ scales,
|
||||
const scalar_t* __restrict__ zeros,
|
||||
int batch,
|
||||
int vec_height,
|
||||
int height,
|
||||
int width
|
||||
) {
|
||||
int b = blockIdx.z;
|
||||
int h = BLOCKHEIGHT8 * blockIdx.x;
|
||||
int w = BLOCKWIDTH * blockIdx.y + threadIdx.x;
|
||||
|
||||
__shared__ scalar_t blockvec[BLOCKWIDTH];
|
||||
blockvec[threadIdx.x] = vec[b * vec_height + (h / BLOCKHEIGHT8) * BLOCKWIDTH + threadIdx.x];
|
||||
__syncthreads();
|
||||
|
||||
scalar_t scale = scales[w];
|
||||
scalar_t zero = zeros[w];
|
||||
|
||||
scalar_t res = 0;
|
||||
int i = width * h + w;
|
||||
int k = 0;
|
||||
|
||||
unsigned int tmp;
|
||||
|
||||
while (k < BLOCKWIDTH) {
|
||||
tmp = as_unsigned(mat[i]);
|
||||
res += (scale * scalar_t((tmp >> 0) & 0xFF) - zero) * blockvec[k + 0];
|
||||
res += (scale * scalar_t((tmp >> 8) & 0xFF) - zero) * blockvec[k + 1];
|
||||
res += (scale * scalar_t((tmp >> 16) & 0xFF) - zero) * blockvec[k + 2];
|
||||
res += (scale * scalar_t((tmp >> 24) & 0xFF) - zero) * blockvec[k + 3];
|
||||
i += width;
|
||||
k += 4;
|
||||
}
|
||||
|
||||
atomicAdd(&mul[b * width + w], res);
|
||||
}
|
||||
|
||||
template <typename scalar_t>
|
||||
__global__ void VecQuant4TransposeMatMulKernel(
|
||||
const scalar_t* __restrict__ vec,
|
||||
const int* __restrict__ mat,
|
||||
scalar_t* __restrict__ mul,
|
||||
const scalar_t* __restrict__ scales,
|
||||
const scalar_t* __restrict__ zeros,
|
||||
int batch,
|
||||
int vec_height,
|
||||
int height,
|
||||
int width
|
||||
) {
|
||||
int b = blockIdx.z;
|
||||
int h = BLOCKHEIGHT4 * blockIdx.x + threadIdx.x / 8;
|
||||
unsigned int shift = (unsigned int)((threadIdx.x % 8) * 4);
|
||||
int w = BLOCKWIDTH * blockIdx.y;
|
||||
|
||||
int n_rows = 8 * BLOCKHEIGHT4 * blockIdx.x + threadIdx.x;
|
||||
int n_cols = b;
|
||||
|
||||
__shared__ scalar_t blockvec[BLOCKWIDTH];
|
||||
blockvec[threadIdx.x] = vec[n_cols * vec_height + w + threadIdx.x];
|
||||
__syncthreads();
|
||||
|
||||
scalar_t res = 0;
|
||||
int i = width * h + w;
|
||||
int k = 0;
|
||||
int j = w;
|
||||
unsigned int tmp;
|
||||
while (k < BLOCKWIDTH) {
|
||||
tmp = as_unsigned(mat[i]);
|
||||
res += (scales[j] * scalar_t((tmp >> shift) & 0xF) - zeros[j]) * blockvec[k];
|
||||
i += 1;
|
||||
j += 1;
|
||||
k += 1;
|
||||
}
|
||||
|
||||
atomicAdd(&mul[n_cols * height * 8 + n_rows], res);
|
||||
}
|
||||
|
||||
void vecquant4transposematmul_cuda(
|
||||
torch::Tensor vec,
|
||||
torch::Tensor mat,
|
||||
torch::Tensor mul,
|
||||
torch::Tensor scales,
|
||||
torch::Tensor zeros
|
||||
) {
|
||||
int batch = vec.size(0);
|
||||
int vec_height = vec.size(1);
|
||||
int height = mat.size(0);
|
||||
int width = mat.size(1);
|
||||
|
||||
dim3 blocks(
|
||||
(height + BLOCKHEIGHT4 - 1) / BLOCKHEIGHT4,
|
||||
(width + BLOCKWIDTH - 1) / BLOCKWIDTH,
|
||||
batch
|
||||
);
|
||||
dim3 threads(BLOCKWIDTH);
|
||||
|
||||
AT_DISPATCH_FLOATING_TYPES(
|
||||
vec.type(), "vecquant4transposematmul_cuda", ([&] {
|
||||
VecQuant4TransposeMatMulKernel<<<blocks, threads>>>(
|
||||
vec.data<scalar_t>(), mat.data<int>(), mul.data<scalar_t>(),
|
||||
scales.data<scalar_t>(), zeros.data<scalar_t>(),
|
||||
batch, vec_height, height, width
|
||||
);
|
||||
})
|
||||
);
|
||||
}
|
||||
Reference in New Issue
Block a user