File size: 5,425 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 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 | #include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <c10/cuda/CUDAException.h>
#include <torch/torch.h>
#include <cmath>
using BF = c10::BFloat16;
// The positive RMS scale cancels in quantization except for quant_eps.
// Keep it in the returned token norm; deliberately omit BF16 intermediates.
__global__ void rmsquant_kernel(uint8_t* packed, float* norms, const BF* x,
const BF* weight, const int32_t* permutation, const int8_t* signs,
const float* boundaries, float rms_eps, float quant_eps) {
constexpr int dim=2048;
__shared__ float values[dim], raw_sum[256], weighted_sum[256], bounds[15];
const int t=threadIdx.x, row=blockIdx.x;
float sx=0, sy=0;
for(int j=t;j<dim;j+=256){
int source=permutation[j];
float raw=float(x[row*dim+source]);
float weighted=raw*float(weight[source]);
values[j]=weighted*float(signs[j]);
sx=fmaf(raw,raw,sx);sy=fmaf(weighted,weighted,sy);
}
raw_sum[t]=sx;weighted_sum[t]=sy;
if(t<15)bounds[t]=boundaries[t];
__syncthreads();
for(int stride=128;stride>=32;stride>>=1){
if(t<stride){raw_sum[t]+=raw_sum[t+stride];weighted_sum[t]+=weighted_sum[t+stride];}
__syncthreads();
}
// Preserve the original reduction tree, but finish inside the first warp.
if(t<32){
float raw=raw_sum[t],weighted=weighted_sum[t];
#pragma unroll
for(int stride=16;stride;stride>>=1){
raw+=__shfl_down_sync(0xffffffff,raw,stride);
weighted+=__shfl_down_sync(0xffffffff,weighted,stride);
}
if(t==0){raw_sum[0]=raw;weighted_sum[0]=weighted;}
}
__syncthreads();
float scale=rsqrtf(raw_sum[0]/float(dim)+rms_eps);
float norm=sqrtf(weighted_sum[0]);
if(t==0)norms[row]=norm*scale;
float inv=1.f/(norm+quant_eps/scale);
float local[8];
#pragma unroll
for(int i=0;i<8;i++)local[i]=values[t+256*i]*inv;
// The first five Hadamard stages stay inside each warp.
#pragma unroll
for(int width=1;width<32;width<<=1){
#pragma unroll
for(int i=0;i<8;i++){
float other=__shfl_xor_sync(0xffffffff,local[i],width);
local[i]=(t&width)?other-local[i]:local[i]+other;
}
}
#pragma unroll
for(int i=0;i<8;i++)values[t+256*i]=local[i];
__syncthreads();
// Only three stages exchange values across warps.
for(int width=32;width<=128;width<<=1){
for(int b=t;b<dim/2;b+=256){
int group=b/width,offset=b-group*width;
int left=group*width*2+offset,right=left+width;
float a=values[left],v=values[right];values[left]=a+v;values[right]=a-v;
}
__syncthreads();
}
#pragma unroll
for(int i=0;i<8;i++)local[i]=values[t+256*i];
// Dimensions256/512/1024 are eight register values owned by one thread.
#pragma unroll
for(int width=1;width<8;width<<=1){
#pragma unroll
for(int b=0;b<4;b++){
int group=b/width,offset=b-group*width;
int left=group*width*2+offset,right=left+width;
float a=local[left],v=local[right];local[left]=a+v;local[right]=a-v;
}
}
// All reads finished before shared storage is reused for packing.
__syncthreads();
#pragma unroll
for(int i=0;i<8;i++)values[t+256*i]=local[i];
__syncthreads();
for(int j=t;j<dim/2;j+=256){
uint8_t codes[2];
#pragma unroll
for(int k=0;k<2;k++){
float v=values[2*j+k]*0.02209708691207961f;
int index=v>bounds[7]?8:0;
index+=v>bounds[index+3]?4:0;
index+=v>bounds[index+1]?2:0;
index+=v>bounds[index]?1:0;
codes[k]=index;
}
packed[row*(dim/2)+j]=codes[0]|(codes[1]<<4);
}
}
void rmsquant(torch::Tensor &packed, torch::Tensor &norms, torch::Tensor const &x,
torch::Tensor const &weight, torch::Tensor const &permutation,
torch::Tensor const &signs, torch::Tensor const &boundaries,
double rms_eps, double quant_eps) {
TORCH_CHECK(x.is_cuda() && x.is_contiguous() && x.scalar_type()==torch::kBFloat16
&& x.dim()==2 && x.size(1)==2048,"requires contiguous CUDA BF16 [rows,2048]");
for(auto* t:{&weight,&permutation,&signs,&boundaries})
TORCH_CHECK(t->device()==x.device() && t->is_contiguous() && t->dim()==1,"invalid constant device/layout");
TORCH_CHECK(weight.scalar_type()==torch::kBFloat16 && weight.numel()==2048,"invalid RMS weight");
TORCH_CHECK(permutation.scalar_type()==torch::kInt && permutation.numel()==2048,"invalid permutation");
TORCH_CHECK(signs.scalar_type()==torch::kChar && signs.numel()==2048,"invalid signs");
TORCH_CHECK(boundaries.scalar_type()==torch::kFloat && boundaries.numel()==15,"invalid boundaries");
TORCH_CHECK(packed.device()==x.device() && packed.is_contiguous() && packed.scalar_type()==torch::kUInt8
&& packed.dim()==2 && packed.size(0)==x.size(0) && packed.size(1)==1024,"invalid packed output");
TORCH_CHECK(norms.device()==x.device() && norms.is_contiguous() && norms.scalar_type()==torch::kFloat
&& norms.dim()==1 && norms.numel()==x.size(0),"invalid norm output");
TORCH_CHECK(std::isfinite(rms_eps) && rms_eps>0 && std::isfinite(quant_eps) && quant_eps>0,"epsilons must be finite and positive");
const at::cuda::OptionalCUDAGuard guard(device_of(x));
if(!x.size(0))return;
rmsquant_kernel<<<x.size(0),256,0,at::cuda::getCurrentCUDAStream()>>>(packed.data_ptr<uint8_t>(),norms.data_ptr<float>(),
x.data_ptr<BF>(),weight.data_ptr<BF>(),permutation.data_ptr<int32_t>(),signs.data_ptr<int8_t>(),
boundaries.data_ptr<float>(),rms_eps,quant_eps);
C10_CUDA_KERNEL_LAUNCH_CHECK();
}
|