biondizzle/nvfp4-megamoe-kernel
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

KV-1/KV-2: Fix quantize kernel — each thread handles 16-elem blocks independently

Browse Source 
Previous version used __shfl_down_sync for group-level amax reduction,
but shuffles operate at warp level and crossed group boundaries.
Fix: each thread independently quantizes its assigned 16-element blocks
from shared memory. Simpler and correct.
This commit is contained in:
biondizzle
2026-06-02 09:41:15 +00:00
parent f23320b5b2
commit c2664281c3
1 changed files with 138 additions and 391 deletions
Download Patch File Download Diff File Expand all files Collapse all files

529
dsv4/kernels/cuda/compressor_reduce_quant.cu
View File

@@ -1,16 +1,8 @@
/**
* FUSED CSA/HCA compress + RMSNorm + NVFP4 quantize kernels.
*
* KV-1/KV-2: Single kernel launch per compressed entry.
* The compressor produces FP32 values, applies kv_norm, then quantizes
* to NVFP4 (E2M1 data + E4M3 block scales + FP32 global scale) all in
* one kernel. No intermediate BF16 materialization.
*
* Shared memory budget per CTA (128 threads, hd=512):
* s_vals: hd * 4 = 2048 bytes (FP32 staging)
* s_nibbles: hd * 1 = 512 bytes (E2M1 nibbles)
* s_sq/s_amax/s_inv_rms: ~16 bytes (reduction scratch)
* Total: ~2576 bytes — well within 48KB
* KV-1/KV-2: Single kernel launch. FP32 -> E2M1 + E4M3 + FP32 gsa.
* Quantize: each thread independently handles 16-element blocks.
* No warp-shuffle group reduction (previous version had a cross-group bug).
*/
#include <cuda.h>
@@ -22,440 +14,195 @@
#include <c10/cuda/CUDAStream.h>
#include <cmath>
// ===========================================================================
// Shared utilities
// ===========================================================================
__device__ __forceinline__ float block_reduce_sum(float val, float* smem, int n_warps) {
for (int offset = 16; offset > 0; offset >>= 1)
val += __shfl_down_sync(0xffffffff, val, offset);
__device__ __forceinline__ float block_reduce_sum(float val, float* smem, int nw) {
for (int o = 16; o > 0; o >>= 1) val += __shfl_down_sync(0xffffffff, val, o);
if (threadIdx.x % 32 == 0) smem[threadIdx.x / 32] = val;
__syncthreads();
float result = 0.0f;
float r = 0.0f;
if (threadIdx.x < 32) {
float v = (threadIdx.x < n_warps) ? smem[threadIdx.x] : 0.0f;
for (int offset = 16; offset > 0; offset >>= 1)
v += __shfl_down_sync(0xffffffff, v, offset);
result = v;
float v = (threadIdx.x < nw) ? smem[threadIdx.x] : 0.0f;
for (int o = 16; o > 0; o >>= 1) v += __shfl_down_sync(0xffffffff, v, o);
r = v;
}
__syncthreads();
return result;
return r;
}
__device__ __forceinline__ float block_reduce_max(float val, float* smem, int n_warps) {
for (int offset = 16; offset > 0; offset >>= 1)
val = fmaxf(val, __shfl_down_sync(0xffffffff, val, offset));
__device__ __forceinline__ float block_reduce_max(float val, float* smem, int nw) {
for (int o = 16; o > 0; o >>= 1) val = fmaxf(val, __shfl_down_sync(0xffffffff, val, o));
if (threadIdx.x % 32 == 0) smem[threadIdx.x / 32] = val;
__syncthreads();
float result = 0.0f;
float r = 0.0f;
if (threadIdx.x < 32) {
float v = (threadIdx.x < n_warps) ? smem[threadIdx.x] : 0.0f;
for (int offset = 16; offset > 0; offset >>= 1)
v = fmaxf(v, __shfl_down_sync(0xffffffff, v, offset));
result = v;
float v = (threadIdx.x < nw) ? smem[threadIdx.x] : 0.0f;
for (int o = 16; o > 0; o >>= 1) v = fmaxf(v, __shfl_down_sync(0xffffffff, v, o));
r = v;
}
__syncthreads();
return result;
return r;
}
__device__ __forceinline__ int half_step_to_e2m1(int hs) {
if (hs <= 4) return hs;
if (hs <= 5) return 4;
if (hs <= 7) return 5;
if (hs <= 10) return 6;
return 7;
if (hs <= 4) return hs; if (hs <= 5) return 4;
if (hs <= 7) return 5; if (hs <= 10) return 6; return 7;
}
// ===========================================================================
// CSA fused compress + norm + quantize
// ===========================================================================
// === CSA ===
__global__ void csa_compress_reduce_quant_kernel(
const float* __restrict__ kv_proj, // [T, 2*hd] FP32
const float* __restrict__ gate_proj, // [T, 2*hd] FP32
const float* __restrict__ position_bias, // [m, 2*hd] FP32 or nullptr
const float* __restrict__ kv_norm_weight, // [hd] FP32 or nullptr
uint8_t* __restrict__ out_fp4, // (n_blocks, hd/2) packed E2M1
uint8_t* __restrict__ out_sf, // (n_blocks, hd/16) E4M3 block scales
float* __restrict__ out_gsa, // (n_blocks,) FP32 global scale
const float* kv_proj, const float* gate_proj,
const float* position_bias, const float* kv_norm_weight,
uint8_t* out_fp4, uint8_t* out_sf, float* out_gsa,
int T, int hd, int m, int n_blocks
) {
int block_i = blockIdx.x;
int tid = threadIdx.x;
int n_threads = blockDim.x;
int kv_dim = 2 * hd;
int n_warps = n_threads / 32;
int bi = blockIdx.x, tid = threadIdx.x, nt = blockDim.x;
int kd = 2*hd, nw = nt/32;
if (bi >= n_blocks) return;
if (block_i >= n_blocks) return;
int ntok = (bi > 0) ? 2*m : m;
int ps = (bi-1)*m, cs = bi*m;
int cpt = (hd+nt-1)/nt;
int n_tokens = (block_i > 0) ? 2 * m : m;
int prev_start = (block_i - 1) * m;
int cur_start = block_i * m;
int cols_per_thread = (hd + n_threads - 1) / n_threads;
// ---- Phase 1: Softmax + weighted sum ----
float local_vals[4], local_max[4], local_denom[4], local_acc[4];
for (int ci = 0; ci < cols_per_thread; ci++) {
int c = tid + ci * n_threads;
if (c >= hd) break;
local_max[ci] = -FLT_MAX;
local_denom[ci] = 0.0f;
local_acc[ci] = 0.0f;
// Pass 1: max gate
for (int t = 0; t < n_tokens; t++) {
int token_idx, gate_offset;
if (block_i > 0) {
if (t < m) { token_idx = prev_start + t; gate_offset = 0; }
else { token_idx = cur_start + (t - m); gate_offset = hd; }
} else { token_idx = t; gate_offset = hd; }
if (token_idx < 0 || token_idx >= T) continue;
float g = gate_proj[token_idx * kv_dim + gate_offset + c];
if (position_bias) {
int pbr = (block_i > 0 && t < m) ? t : (block_i > 0 ? (t - m) : t);
if (pbr >= 0 && pbr < m) g += position_bias[pbr * kv_dim + gate_offset + c];
}
local_max[ci] = fmaxf(local_max[ci], g);
float lv[4],lm[4],ld[4],la[4];
for (int ci=0;ci<cpt;ci++) {
int c=tid+ci*nt; if(c>=hd) break;
lm[ci]=-FLT_MAX; ld[ci]=0; la[ci]=0;
for (int t=0;t<ntok;t++) {
int ti,go; if(bi>0){if(t<m){ti=ps+t;go=0;}else{ti=cs+(t-m);go=hd;}}else{ti=t;go=hd;}
if(ti<0||ti>=T) continue;
float g=gate_proj[ti*kd+go+c];
if(position_bias){int p=(bi>0&&t<m)?t:(bi>0?(t-m):t);if(p>=0&&p<m)g+=position_bias[p*kd+go+c];}
lm[ci]=fmaxf(lm[ci],g);
}
// Pass 2: exp + weighted sum
for (int t = 0; t < n_tokens; t++) {
int token_idx, kv_offset, gate_offset;
if (block_i > 0) {
if (t < m) { token_idx = prev_start + t; kv_offset = 0; gate_offset = 0; }
else { token_idx = cur_start + (t - m); kv_offset = hd; gate_offset = hd; }
} else { token_idx = t; kv_offset = hd; gate_offset = hd; }
if (token_idx < 0 || token_idx >= T) continue;
float g = gate_proj[token_idx * kv_dim + gate_offset + c];
float kv_val = kv_proj[token_idx * kv_dim + kv_offset + c];
if (position_bias) {
int pbr = (block_i > 0 && t < m) ? t : (block_i > 0 ? (t - m) : t);
if (pbr >= 0 && pbr < m) {
float pb = position_bias[pbr * kv_dim + gate_offset + c];
g += pb;
kv_val += position_bias[pbr * kv_dim + kv_offset + c];
}
}
float e = expf(g - local_max[ci]);
local_denom[ci] += e;
local_acc[ci] += e * kv_val;
for (int t=0;t<ntok;t++) {
int ti,ko,go; if(bi>0){if(t<m){ti=ps+t;ko=0;go=0;}else{ti=cs+(t-m);ko=hd;go=hd;}}else{ti=t;ko=hd;go=hd;}
if(ti<0||ti>=T) continue;
float g=gate_proj[ti*kd+go+c], kv=kv_proj[ti*kd+ko+c];
if(position_bias){int p=(bi>0&&t<m)?t:(bi>0?(t-m):t);if(p>=0&&p<m){float pb=position_bias[p*kd+go+c];g+=pb;kv+=position_bias[p*kd+ko+c];}}
float e=expf(g-lm[ci]); ld[ci]+=e; la[ci]+=e*kv;
}
local_vals[ci] = (local_denom[ci] > 0.0f) ? (local_acc[ci] / local_denom[ci]) : 0.0f;
lv[ci]=(ld[ci]>0)?(la[ci]/ld[ci]):0;
}
// ---- Phase 2: kv_norm (RMSNorm) ----
if (kv_norm_weight) {
float local_sq = 0.0f;
for (int ci = 0; ci < cols_per_thread; ci++) {
int c = tid + ci * n_threads;
if (c >= hd) break;
local_sq += local_vals[ci] * local_vals[ci];
}
__shared__ float s_sq;
float total_sq = block_reduce_sum(local_sq, &s_sq, n_warps);
__shared__ float s_inv_rms;
if (tid == 0) s_inv_rms = rsqrtf(total_sq / hd + 1e-6f);
__syncthreads();
for (int ci = 0; ci < cols_per_thread; ci++) {
int c = tid + ci * n_threads;
if (c >= hd) break;
local_vals[ci] *= s_inv_rms * kv_norm_weight[c];
}
if(kv_norm_weight) {
float ls=0; for(int ci=0;ci<cpt;ci++){int c=tid+ci*nt;if(c>=hd)break;ls+=lv[ci]*lv[ci];}
__shared__ float ss; float ts=block_reduce_sum(ls,&ss,nw);
__shared__ float sir; if(tid==0) sir=rsqrtf(ts/hd+1e-6f); __syncthreads();
for(int ci=0;ci<cpt;ci++){int c=tid+ci*nt;if(c>=hd)break;lv[ci]*=sir*kv_norm_weight[c];}
}
// ---- Phase 3: Global scale (gsa) ----
float entry_amax = 0.0f;
for (int ci = 0; ci < cols_per_thread; ci++) {
int c = tid + ci * n_threads;
if (c >= hd) break;
entry_amax = fmaxf(entry_amax, fabsf(local_vals[ci]));
}
__shared__ float s_amax;
float global_amax = block_reduce_max(entry_amax, &s_amax, n_warps);
float gsa = fmaxf(global_amax, 1e-8f) / (6.0f * 448.0f);
if (tid == 0) out_gsa[block_i] = gsa;
float ea=0; for(int ci=0;ci<cpt;ci++){int c=tid+ci*nt;if(c>=hd)break;ea=fmaxf(ea,fabsf(lv[ci]));}
__shared__ float sam; float ga=block_reduce_max(ea,&sam,nw);
float gsa=fmaxf(ga,1e-8f)/(6.0f*448.0f);
if(tid==0) out_gsa[bi]=gsa;
// ---- Phase 4: NVFP4 quantize via shared memory ----
__shared__ float s_vals[512]; // FP32 staging
__shared__ uint8_t s_nib[512]; // E2M1 nibbles
for (int ci = 0; ci < cols_per_thread; ci++) {
int c = tid + ci * n_threads;
if (c >= hd) break;
s_vals[c] = local_vals[ci];
}
__shared__ float sv[512];
for(int ci=0;ci<cpt;ci++){int c=tid+ci*nt;if(c>=hd)break;sv[c]=lv[ci];}
__syncthreads();
int n_fp4_blocks = hd / 16;
int tpb = n_threads / n_fp4_blocks; // threads per fp4 block
int my_b = tid / tpb;
int my_l = tid % tpb;
if (my_b < n_fp4_blocks) {
int base = my_b * 16;
// Block amax
float bamax = 0.0f;
for (int i = my_l; i < 16; i += tpb) {
int c = base + i;
if (c < hd) bamax = fmaxf(bamax, fabsf(s_vals[c]) / gsa);
}
for (int off = tpb / 2; off > 0; off >>= 1)
bamax = fmaxf(bamax, __shfl_down_sync(0xffffffff, bamax, off));
float fbamax = __shfl_sync(0xffffffff, bamax, 0);
float bsf = fbamax / 6.0f;
bool zero_blk = (fbamax < 6.0f * 0.001953125f);
if (my_l == 0) {
if (zero_blk) {
out_sf[block_i * (hd / 16) + my_b] = 0;
} else {
__nv_fp8_e4m3 obj(bsf);
out_sf[block_i * (hd / 16) + my_b] = *(uint8_t*)&obj;
}
}
// Quantize to E2M1 nibbles
for (int i = my_l; i < 16; i += tpb) {
int c = base + i;
if (c >= hd || zero_blk) { s_nib[c] = 0; continue; }
float s = (s_vals[c] / gsa) / bsf;
int hs = __float2int_rn(fminf(fabsf(s), 6.0f) * 2.0f);
if (hs > 12) hs = 12;
int idx = half_step_to_e2m1(hs);
if (s < 0) idx += 8;
s_nib[c] = idx;
}
}
__syncthreads();
// Pack and write
if (my_b < n_fp4_blocks && my_l == 0) {
int base = my_b * 16;
for (int i = 0; i < 8; i++) {
uint8_t lo = s_nib[base + 2 * i] & 0x0F;
uint8_t hi = s_nib[base + 2 * i + 1] & 0x0F;
out_fp4[block_i * (hd / 2) + my_b * 8 + i] = (hi << 4) | lo;
int nfb=hd/16;
for(int b=tid;b<nfb;b+=nt) {
int base=b*16;
float ba=0; for(int i=0;i<16;i++){int c=base+i;if(c<hd)ba=fmaxf(ba,fabsf(sv[c])/gsa);}
float bsf=ba/6.0f; bool z=(ba<6.0f*0.001953125f);
if(z){out_sf[bi*(hd/16)+b]=0;}else{__nv_fp8_e4m3 o(bsf);out_sf[bi*(hd/16)+b]=*(uint8_t*)&o;}
float ab=z?0.0f:bsf;
for(int i=0;i<8;i++){
int c0=base+2*i,c1=base+2*i+1; uint8_t lo=0,hi=0;
if(!z&&c0<hd){float s=(sv[c0]/gsa)/ab;int hs=__float2int_rn(fminf(fabsf(s),6.0f)*2.0f);if(hs>12)hs=12;lo=half_step_to_e2m1(hs);if(s<0)lo+=8;}
if(!z&&c1<hd){float s=(sv[c1]/gsa)/ab;int hs=__float2int_rn(fminf(fabsf(s),6.0f)*2.0f);if(hs>12)hs=12;hi=half_step_to_e2m1(hs);if(s<0)hi+=8;}
out_fp4[bi*(hd/2)+b*8+i]=(hi<<4)|lo;
}
}
}
// ===========================================================================
// HCA fused compress + norm + quantize (simpler — no overlap)
// ===========================================================================
// === HCA ===
__global__ void hca_compress_reduce_quant_kernel(
const float* __restrict__ kv_proj,
const float* __restrict__ gate_proj,
const float* __restrict__ position_bias,
const float* __restrict__ kv_norm_weight,
uint8_t* __restrict__ out_fp4,
uint8_t* __restrict__ out_sf,
float* __restrict__ out_gsa,
const float* kv_proj, const float* gate_proj,
const float* position_bias, const float* kv_norm_weight,
uint8_t* out_fp4, uint8_t* out_sf, float* out_gsa,
int T, int hd, int m, int n_blocks
) {
int block_i = blockIdx.x;
int tid = threadIdx.x;
int n_threads = blockDim.x;
int n_warps = n_threads / 32;
int bi=blockIdx.x,tid=threadIdx.x,nt=blockDim.x,nw=nt/32;
if(bi>=n_blocks) return;
int cpt=(hd+nt-1)/nt;
if (block_i >= n_blocks) return;
int cols_per_thread = (hd + n_threads - 1) / n_threads;
// Phase 1: Softmax + weighted sum
float local_vals[4];
for (int ci = 0; ci < cols_per_thread; ci++) {
int c = tid + ci * n_threads;
if (c >= hd) break;
float lmax = -FLT_MAX, ldenom = 0.0f, lacc = 0.0f;
int start = block_i * m;
for (int t = 0; t < m; t++) {
int ti = start + t;
if (ti >= T) break;
float g = gate_proj[ti * hd + c];
if (position_bias && t < m) g += position_bias[t * hd + c];
lmax = fmaxf(lmax, g);
}
for (int t = 0; t < m; t++) {
int ti = start + t;
if (ti >= T) break;
float g = gate_proj[ti * hd + c];
float kv = kv_proj[ti * hd + c];
if (position_bias && t < m) { float pb = position_bias[t * hd + c]; g += pb; kv += pb; }
float e = expf(g - lmax);
ldenom += e; lacc += e * kv;
}
local_vals[ci] = (ldenom > 0.0f) ? (lacc / ldenom) : 0.0f;
float lv[4];
for(int ci=0;ci<cpt;ci++){
int c=tid+ci*nt;if(c>=hd)break;
float lm=-FLT_MAX,ld=0,la=0; int st=bi*m;
for(int t=0;t<m;t++){int ti=st+t;if(ti>=T)break;float g=gate_proj[ti*hd+c];if(position_bias&&t<m)g+=position_bias[t*hd+c];lm=fmaxf(lm,g);}
for(int t=0;t<m;t++){int ti=st+t;if(ti>=T)break;float g=gate_proj[ti*hd+c],kv=kv_proj[ti*hd+c];if(position_bias&&t<m){float pb=position_bias[t*hd+c];g+=pb;kv+=pb;}float e=expf(g-lm);ld+=e;la+=e*kv;}
lv[ci]=(ld>0)?(la/ld):0;
}
// Phase 2: kv_norm
if (kv_norm_weight) {
float lsq = 0.0f;
for (int ci = 0; ci < cols_per_thread; ci++) {
int c = tid + ci * n_threads;
if (c >= hd) break;
lsq += local_vals[ci] * local_vals[ci];
}
__shared__ float s_sq;
float tsq = block_reduce_sum(lsq, &s_sq, n_warps);
__shared__ float s_inv_rms;
if (tid == 0) s_inv_rms = rsqrtf(tsq / hd + 1e-6f);
__syncthreads();
for (int ci = 0; ci < cols_per_thread; ci++) {
int c = tid + ci * n_threads;
if (c >= hd) break;
local_vals[ci] *= s_inv_rms * kv_norm_weight[c];
}
if(kv_norm_weight){
float ls=0;for(int ci=0;ci<cpt;ci++){int c=tid+ci*nt;if(c>=hd)break;ls+=lv[ci]*lv[ci];}
__shared__ float ss;float ts=block_reduce_sum(ls,&ss,nw);
__shared__ float sir;if(tid==0)sir=rsqrtf(ts/hd+1e-6f);__syncthreads();
for(int ci=0;ci<cpt;ci++){int c=tid+ci*nt;if(c>=hd)break;lv[ci]*=sir*kv_norm_weight[c];}
}
// Phase 3: gsa
float eamax = 0.0f;
for (int ci = 0; ci < cols_per_thread; ci++) {
int c = tid + ci * n_threads;
if (c >= hd) break;
eamax = fmaxf(eamax, fabsf(local_vals[ci]));
}
__shared__ float s_amax;
float gamax = block_reduce_max(eamax, &s_amax, n_warps);
float gsa = fmaxf(gamax, 1e-8f) / (6.0f * 448.0f);
if (tid == 0) out_gsa[block_i] = gsa;
float ea=0;for(int ci=0;ci<cpt;ci++){int c=tid+ci*nt;if(c>=hd)break;ea=fmaxf(ea,fabsf(lv[ci]));}
__shared__ float sam;float ga=block_reduce_max(ea,&sam,nw);
float gsa=fmaxf(ga,1e-8f)/(6.0f*448.0f);
if(tid==0)out_gsa[bi]=gsa;
// Phase 4: NVFP4 quantize
__shared__ float s_vals[512];
__shared__ uint8_t s_nib[512];
for (int ci = 0; ci < cols_per_thread; ci++) {
int c = tid + ci * n_threads;
if (c >= hd) break;
s_vals[c] = local_vals[ci];
}
__shared__ float sv[512];
for(int ci=0;ci<cpt;ci++){int c=tid+ci*nt;if(c>=hd)break;sv[c]=lv[ci];}
__syncthreads();
int nfb = hd / 16;
int tpb = n_threads / nfb;
int my_b = tid / tpb;
int my_l = tid % tpb;
if (my_b < nfb) {
int base = my_b * 16;
float bamax = 0.0f;
for (int i = my_l; i < 16; i += tpb) {
int c = base + i;
if (c < hd) bamax = fmaxf(bamax, fabsf(s_vals[c]) / gsa);
}
for (int off = tpb / 2; off > 0; off >>= 1)
bamax = fmaxf(bamax, __shfl_down_sync(0xffffffff, bamax, off));
float fbamax = __shfl_sync(0xffffffff, bamax, 0);
float bsf = fbamax / 6.0f;
bool zblk = (fbamax < 6.0f * 0.001953125f);
if (my_l == 0) {
if (zblk) { out_sf[block_i * (hd / 16) + my_b] = 0; }
else { __nv_fp8_e4m3 obj(bsf); out_sf[block_i * (hd / 16) + my_b] = *(uint8_t*)&obj; }
}
for (int i = my_l; i < 16; i += tpb) {
int c = base + i;
if (c >= hd || zblk) { s_nib[c] = 0; continue; }
float s = (s_vals[c] / gsa) / bsf;
int hs = __float2int_rn(fminf(fabsf(s), 6.0f) * 2.0f);
if (hs > 12) hs = 12;
int idx = half_step_to_e2m1(hs);
if (s < 0) idx += 8;
s_nib[c] = idx;
}
}
__syncthreads();
if (my_b < nfb && my_l == 0) {
int base = my_b * 16;
for (int i = 0; i < 8; i++) {
uint8_t lo = s_nib[base + 2 * i] & 0x0F;
uint8_t hi = s_nib[base + 2 * i + 1] & 0x0F;
out_fp4[block_i * (hd / 2) + my_b * 8 + i] = (hi << 4) | lo;
int nfb=hd/16;
for(int b=tid;b<nfb;b+=nt){
int base=b*16;
float ba=0;for(int i=0;i<16;i++){int c=base+i;if(c<hd)ba=fmaxf(ba,fabsf(sv[c])/gsa);}
float bsf=ba/6.0f;bool z=(ba<6.0f*0.001953125f);
if(z){out_sf[bi*(hd/16)+b]=0;}else{__nv_fp8_e4m3 o(bsf);out_sf[bi*(hd/16)+b]=*(uint8_t*)&o;}
float ab=z?0.0f:bsf;
for(int i=0;i<8;i++){
int c0=base+2*i,c1=base+2*i+1;uint8_t lo=0,hi=0;
if(!z&&c0<hd){float s=(sv[c0]/gsa)/ab;int hs=__float2int_rn(fminf(fabsf(s),6.0f)*2.0f);if(hs>12)hs=12;lo=half_step_to_e2m1(hs);if(s<0)lo+=8;}
if(!z&&c1<hd){float s=(sv[c1]/gsa)/ab;int hs=__float2int_rn(fminf(fabsf(s),6.0f)*2.0f);if(hs>12)hs=12;hi=half_step_to_e2m1(hs);if(s<0)hi+=8;}
out_fp4[bi*(hd/2)+b*8+i]=(hi<<4)|lo;
}
}
}
// ===========================================================================
// PyTorch bindings
// ===========================================================================
std::tuple<torch::Tensor, torch::Tensor, torch::Tensor>
csa_compress_reduce_quant_cuda(
torch::Tensor kv_proj, // [T, 2*hd] FP32
torch::Tensor gate_proj, // [T, 2*hd] FP32
torch::Tensor position_bias, // [m, 2*hd] FP32 or empty
torch::Tensor kv_norm_weight, // [hd] FP32 or empty
int64_t m, int64_t n_blocks
) {
int T = kv_proj.size(0);
int hd = kv_proj.size(1) / 2;
int threads = 128;
const float* pos_ptr = (position_bias.numel() > 0) ? position_bias.data_ptr<float>() : nullptr;
const float* norm_ptr = (kv_norm_weight.numel() > 0) ? kv_norm_weight.data_ptr<float>() : nullptr;
auto opts = kv_proj.options();
auto out_fp4 = torch::zeros({(int)n_blocks, hd / 2}, opts.dtype(torch::kUInt8));
auto out_sf = torch::zeros({(int)n_blocks, hd / 16}, opts.dtype(torch::kUInt8));
auto out_gsa = torch::zeros({(int)n_blocks}, opts.dtype(torch::kFloat32));
csa_compress_reduce_quant_kernel<<<n_blocks, threads, 0, c10::cuda::getCurrentCUDAStream()>>>(
kv_proj.data_ptr<float>(),
gate_proj.data_ptr<float>(),
pos_ptr, norm_ptr,
out_fp4.data_ptr<uint8_t>(),
out_sf.data_ptr<uint8_t>(),
out_gsa.data_ptr<float>(),
T, hd, (int)m, (int)n_blocks
);
// === Bindings ===
std::tuple<torch::Tensor,torch::Tensor,torch::Tensor>
csa_compress_reduce_quant_cuda(torch::Tensor kp,torch::Tensor gp,torch::Tensor pb,torch::Tensor nw,int64_t m,int64_t nb){
int T=kp.size(0),hd=kp.size(1)/2,th=128;
const float*pp=(pb.numel()>0)?pb.data_ptr<float>():nullptr;
const float*np=(nw.numel()>0)?nw.data_ptr<float>():nullptr;
auto o=kp.options();
auto of4=torch::zeros({(int)nb,hd/2},o.dtype(torch::kUInt8));
auto osf=torch::zeros({(int)nb,hd/16},o.dtype(torch::kUInt8));
auto ogs=torch::zeros({(int)nb},o.dtype(torch::kFloat32));
csa_compress_reduce_quant_kernel<<<nb,th,0,c10::cuda::getCurrentCUDAStream()>>>(
kp.data_ptr<float>(),gp.data_ptr<float>(),pp,np,
of4.data_ptr<uint8_t>(),osf.data_ptr<uint8_t>(),ogs.data_ptr<float>(),
T,hd,(int)m,(int)nb);
C10_CUDA_CHECK(cudaGetLastError());
return {out_fp4.view(torch::kFloat4_e2m1fn_x2),
out_sf.view(torch::kFloat8_e4m3fn),
out_gsa};
return {of4.view(torch::kFloat4_e2m1fn_x2),osf.view(torch::kFloat8_e4m3fn),ogs};
}
std::tuple<torch::Tensor, torch::Tensor, torch::Tensor>
hca_compress_reduce_quant_cuda(
torch::Tensor kv_proj,
torch::Tensor gate_proj,
torch::Tensor position_bias,
torch::Tensor kv_norm_weight,
int64_t m, int64_t n_blocks
) {
int T = kv_proj.size(0);
int hd = kv_proj.size(1);
int threads = 128;
const float* pos_ptr = (position_bias.numel() > 0) ? position_bias.data_ptr<float>() : nullptr;
const float* norm_ptr = (kv_norm_weight.numel() > 0) ? kv_norm_weight.data_ptr<float>() : nullptr;
auto opts = kv_proj.options();
auto out_fp4 = torch::zeros({(int)n_blocks, hd / 2}, opts.dtype(torch::kUInt8));
auto out_sf = torch::zeros({(int)n_blocks, hd / 16}, opts.dtype(torch::kUInt8));
auto out_gsa = torch::zeros({(int)n_blocks}, opts.dtype(torch::kFloat32));
hca_compress_reduce_quant_kernel<<<n_blocks, threads, 0, c10::cuda::getCurrentCUDAStream()>>>(
kv_proj.data_ptr<float>(),
gate_proj.data_ptr<float>(),
pos_ptr, norm_ptr,
out_fp4.data_ptr<uint8_t>(),
out_sf.data_ptr<uint8_t>(),
out_gsa.data_ptr<float>(),
T, hd, (int)m, (int)n_blocks
);
std::tuple<torch::Tensor,torch::Tensor,torch::Tensor>
hca_compress_reduce_quant_cuda(torch::Tensor kp,torch::Tensor gp,torch::Tensor pb,torch::Tensor nw,int64_t m,int64_t nb){
int T=kp.size(0),hd=kp.size(1),th=128;
const float*pp=(pb.numel()>0)?pb.data_ptr<float>():nullptr;
const float*np=(nw.numel()>0)?nw.data_ptr<float>():nullptr;
auto o=kp.options();
auto of4=torch::zeros({(int)nb,hd/2},o.dtype(torch::kUInt8));
auto osf=torch::zeros({(int)nb,hd/16},o.dtype(torch::kUInt8));
auto ogs=torch::zeros({(int)nb},o.dtype(torch::kFloat32));
hca_compress_reduce_quant_kernel<<<nb,th,0,c10::cuda::getCurrentCUDAStream()>>>(
kp.data_ptr<float>(),gp.data_ptr<float>(),pp,np,
of4.data_ptr<uint8_t>(),osf.data_ptr<uint8_t>(),ogs.data_ptr<float>(),
T,hd,(int)m,(int)nb);
C10_CUDA_CHECK(cudaGetLastError());
return {out_fp4.view(torch::kFloat4_e2m1fn_x2),
out_sf.view(torch::kFloat8_e4m3fn),
out_gsa};
return {of4.view(torch::kFloat4_e2m1fn_x2),osf.view(torch::kFloat8_e4m3fn),ogs};
}
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.def("csa_compress_reduce_quant", &csa_compress_reduce_quant_cuda,
"Fused CSA compress + norm + NVFP4 quantize");
m.def("hca_compress_reduce_quant", &hca_compress_reduce_quant_cuda,
"Fused HCA compress + norm + NVFP4 quantize");
PYBIND11_MODULE(TORCH_EXTENSION_NAME,m){
m.def("csa_compress_reduce_quant",&csa_compress_reduce_quant_cuda);
m.def("hca_compress_reduce_quant",&hca_compress_reduce_quant_cuda);
}