File size: 4,611 Bytes
b35ab1f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 | #include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <c10/cuda/CUDAException.h>
#include <torch/torch.h>
template<class T>
__global__ void packed_gemv(T* out, const uint8_t* x, const uint8_t* w,
const float* xn, const c10::BFloat16* wn, const int8_t* ac, const int8_t* wc,
const T* bias, int n, int packed_k, float as, float ws) {
const int lane = threadIdx.x & 31;
const int col = blockIdx.x * 4 + threadIdx.x / 32;
const int row = blockIdx.y;
__shared__ int a_codes[16], w_codes[16];
if (threadIdx.x < 16) { a_codes[threadIdx.x] = ac[threadIdx.x]; w_codes[threadIdx.x] = wc[threadIdx.x]; }
__syncthreads();
if (col >= n) return;
int sum = 0;
if (packed_k % 4 == 0 && (reinterpret_cast<uintptr_t>(x) & 3) == 0 && (reinterpret_cast<uintptr_t>(w) & 3) == 0) {
const uint32_t* xp = reinterpret_cast<const uint32_t*>(x + row * packed_k);
const uint32_t* wp = reinterpret_cast<const uint32_t*>(w + col * packed_k);
for (int k = lane; k < packed_k / 4; k += 32) {
const uint32_t av = xp[k], bv = wp[k];
uint32_t al = 0, ah = 0, bl = 0, bh = 0;
#pragma unroll
for (int j = 0; j < 4; ++j) {
al |= uint32_t(uint8_t(a_codes[(av >> (j * 8)) & 15])) << (j * 8);
ah |= uint32_t(uint8_t(a_codes[(av >> (j * 8 + 4)) & 15])) << (j * 8);
bl |= uint32_t(uint8_t(w_codes[(bv >> (j * 8)) & 15])) << (j * 8);
bh |= uint32_t(uint8_t(w_codes[(bv >> (j * 8 + 4)) & 15])) << (j * 8);
}
sum = __dp4a(static_cast<int>(al), static_cast<int>(bl), sum);
sum = __dp4a(static_cast<int>(ah), static_cast<int>(bh), sum);
}
} else {
for (int k = lane; k < packed_k; k += 32) {
const uint8_t a = x[row * packed_k + k];
const uint8_t b = w[col * packed_k + k];
sum += a_codes[a & 15] * w_codes[b & 15];
sum += a_codes[a >> 4] * w_codes[b >> 4];
}
}
#pragma unroll
for (int offset = 16; offset; offset >>= 1) sum += __shfl_down_sync(0xffffffff, sum, offset);
if (lane == 0) {
float value = static_cast<float>(sum);
value *= xn[row] * static_cast<float>(wn[col]) * (as * ws);
if (bias) value += static_cast<float>(bias[col]);
out[row * n + col] = static_cast<T>(value);
}
}
void orbitquant_gemv(torch::Tensor &out, torch::Tensor const &x, torch::Tensor const &w, torch::Tensor const &xn, torch::Tensor const &wn, torch::Tensor const &ac, torch::Tensor const &wc, torch::Tensor const &bias, bool has_bias, double activation_scale, double weight_scale) {
TORCH_CHECK(x.is_cuda() && x.scalar_type() == torch::kUInt8 && x.dim() == 2, "x must be rank-2 packed CUDA uint8");
const torch::Tensor* tensors[] = {&out, &w, &xn, &wn, &ac, &wc};
for (const auto* t : tensors) {
TORCH_CHECK(t->device() == x.device() && t->is_contiguous(), "tensors must be contiguous on the input CUDA device");
}
TORCH_CHECK(x.is_contiguous() && w.scalar_type() == torch::kUInt8, "packed inputs must be contiguous uint8");
TORCH_CHECK(x.size(0) >= 1 && x.size(0) <= 8 && x.size(1) > 0 && x.size(1) <= 8192, "supported rows 1..8 and packed K 1..8192");
TORCH_CHECK(wn.scalar_type() == torch::kBFloat16 && xn.scalar_type() == torch::kFloat, "norm dtypes must be BF16 weights and FP32 activations");
TORCH_CHECK(ac.scalar_type() == torch::kChar && wc.scalar_type() == torch::kChar && ac.numel() == 16 && wc.numel() == 16, "codebooks must contain 16 INT8 entries");
TORCH_CHECK(out.dim() == 2 && out.size(0) == x.size(0) && out.size(1) == wn.numel() && xn.numel() == x.size(0), "output and norm shapes mismatch");
TORCH_CHECK(w.numel() == wn.numel() * x.size(1), "packed weight size mismatch");
TORCH_CHECK(out.scalar_type() == torch::kBFloat16 || out.scalar_type() == torch::kHalf, "output must be BF16 or FP16");
if (has_bias) TORCH_CHECK(bias.device() == x.device() && bias.is_contiguous() && bias.scalar_type() == out.scalar_type() && bias.numel() == wn.numel(), "invalid bias");
if (wn.numel() == 0) return;
const at::cuda::OptionalCUDAGuard guard(device_of(x));
const dim3 grid((wn.numel() + 3) / 4, x.size(0));
const auto stream = at::cuda::getCurrentCUDAStream();
#define LAUNCH(T) packed_gemv<T><<<grid, 128, 0, stream>>>(out.data_ptr<T>(), x.data_ptr<uint8_t>(), w.data_ptr<uint8_t>(), xn.data_ptr<float>(), wn.data_ptr<c10::BFloat16>(), ac.data_ptr<int8_t>(), wc.data_ptr<int8_t>(), has_bias ? bias.data_ptr<T>() : nullptr, wn.numel(), x.size(1), activation_scale, weight_scale)
if (out.scalar_type() == torch::kBFloat16) { LAUNCH(c10::BFloat16); } else { LAUNCH(c10::Half); }
#undef LAUNCH
C10_CUDA_KERNEL_LAUNCH_CHECK();
}
|