David
/
Foundry
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity
Foundry/cuda_kernels/matvec_kquants.cu

633 lines
28 KiB
Plaintext
Raw Blame History

/* SPDX-License-Identifier: AGPL-3.0-or-later */
/* Copyright (C) 2026 David (Corey) Herrera */
/**
* @file matvec_kquants.cu
* @brief sm_37 matvecs for repacked Q4_K/Q5_K/Q6_K/Q8_0/Q2_0_g128 and dense F32/F16/BF16 rows.
*
* Written from a behaviour specification (docs/cleanroom/f_cuq_spec.md); no ggml or
* llama.cpp source was consulted.
*
* Every kernel is bit-exact by contract: the order and grouping of every rounding is
* fixed (see the comments on each body), so a batched kernel reproduces the single-column
* kernel per column. Single-column kernels are the batched bodies at NC = 1 with a
* 2048-float x chunk; batched ones use a 512-float chunk per column.
*
* @status None
* @issues Kept as-is on purpose: Q4_K is min-lifted while Q5_K is computed straight;
* batched K-quant/Q8_0 kernels address column c at c*512 (cols < 512 is out of
* bounds, never launched); partial trailing blocks are ignored; dense batched
* trees assume a power-of-two blockDim.x; an odd final row is computed twice
* (row 0 duplicated) and the copy discarded.
* @todo None
*/
#include "matvec_nx_staging.cuh"
/* Shared prefix (floats) of the single-column Q4_K kernel: one sum per 32 staged x. */
#define FC_KQ_XSUM_FLOATS 64
/* Per-column shared prefix of the batched Q4_K kernels: one sum per 32 staged x. */
#define FC_NX_XSUM_PER_COL 16
#define FC_KQ_THREADS 256
#define FC_KQ_SINGLE_CHUNK 2048
#define FC_KQ_SUPERBLOCK 256
#define FC_KQ_KIND_Q4K 0
#define FC_KQ_KIND_Q5K 1
#define FC_KQ_KIND_Q6K 2
#define FC_Q4K_BLOCK_BYTES 152
#define FC_Q5K_BLOCK_BYTES 184
#define FC_Q6K_BLOCK_BYTES 216
#define FC_Q80_BLOCK_BYTES 36
#define FC_Q2_0_BLOCK_BYTES 36
#define FC_DENSE_F32 0
#define FC_DENSE_F16 1
#define FC_DENSE_BF16 2
extern __shared__ __align__(16) float fc_kq_shared[];
/* ---------------------------------------------------------------------------------------
* Shared building blocks
* ------------------------------------------------------------------------------------- */
/* Warp sum, shuffle-down tree 16/8/4/2/1; only lane 0's result is meaningful. */
__device__ __forceinline__ float fc_kq_warp_sum(float v) {
v += __shfl_down(v, 16);
v += __shfl_down(v, 8);
v += __shfl_down(v, 4);
v += __shfl_down(v, 2);
v += __shfl_down(v, 1);
return v;
}
/* 2^e as a float; folds to a constant when e is a compile-time value. */
__device__ __forceinline__ float fc_kq_pow2(int e) {
return __uint_as_float((unsigned)(127 + e) << 23);
}
/* Exact small-integer decode: OR the masked bits into the mantissa of 2^23, then one FMA
* scales and re-biases. Saves the shift and the int->float convert per element. */
__device__ __forceinline__ float fc_kq_magic(unsigned word, unsigned mask, float mul, float add) {
return __fmaf_rn(__uint_as_float(0x4B000000u | (word & mask)), mul, add);
}
/* (word >> pos) & 15, minus bias; pos <= 12 keeps the field inside the mantissa. */
__device__ __forceinline__ float fc_kq_nibble(unsigned word, int pos, float bias) {
return fc_kq_magic(word, 0xFu << pos, fc_kq_pow2(-pos),
-(fc_kq_pow2(23 - pos) + bias));
}
/* 16 * ((word >> pos) & (mask >> pos)): Q5_K fifth bit or Q6_K high pair. */
__device__ __forceinline__ float fc_kq_high16(unsigned word, unsigned mask, int pos) {
return fc_kq_magic(word, mask, fc_kq_pow2(4 - pos), -fc_kq_pow2(27 - pos));
}
/* Byte k of an 8-byte vector as a float (one byte-select convert). */
__device__ __forceinline__ float fc_kq_byte(uint2 v, int k) {
const unsigned word = k < 4 ? v.x : v.y;
return (float)((word >> (8 * (k & 3))) & 0xFFu);
}
/* Single-column x staging: a plain read-only copy of one chunk. */
__device__ __forceinline__ void fc_kq_stage_single(float *sx, const float *x, int cbase,
int cend) {
for (int i = threadIdx.x; i < cend; i += FC_KQ_THREADS)
sx[i] = __ldg(x + cbase + i);
}
/* Stage one chunk of NC columns; the single-column form when CHUNK is the 2048 chunk. */
template <int NC, int CHUNK>
__device__ __forceinline__ void fc_kq_stage(float *sx, const float *x, int col0, int nc,
int cols, int cbase, int cend) {
if (CHUNK == FC_KQ_SINGLE_CHUNK)
fc_kq_stage_single(sx, x, cbase, cend);
else
FC_NX_STAGE_X(sx, x, col0, cols, cbase, cend, NC, nc, FC_NX_STAGE_KEEP);
}
/* Lane 0 writes column c of row r when the row exists and the column is live. */
template <int NC>
__device__ __forceinline__ void fc_kq_store_row(float *out, const float (&acc)[NC], int rows,
int row, bool active, int col0, int nc) {
#pragma unroll
for (int c = 0; c < NC; c++) {
const float sum = fc_kq_warp_sum(acc[c]);
if ((threadIdx.x & 31) == 0 && active && c < nc)
out[(size_t)(col0 + c) * rows + row] = sum;
}
}
/* Both rows of a pair, column by column (keeps one column's sums live at a time). */
template <int NC>
__device__ __forceinline__ void fc_kq_store_rowpair(float *out, const float (&acc0)[NC],
const float (&acc1)[NC], int rows, int row,
bool active0, bool active1, int col0,
int nc) {
#pragma unroll
for (int c = 0; c < NC; c++) {
const float sum0 = fc_kq_warp_sum(acc0[c]);
const float sum1 = fc_kq_warp_sum(acc1[c]);
if ((threadIdx.x & 31) == 0 && c < nc) {
if (active0)
out[(size_t)(col0 + c) * rows + row] = sum0;
if (active1)
out[(size_t)(col0 + c) * rows + row + 1] = sum1;
}
}
}
/* ---------------------------------------------------------------------------------------
* Q4_K: min-lifted. Per lane and row, P += d * sum(sc*q*x) per super-block, and separately
* M += (dmin*mn) * xsum over owned (super-block, sub-block) pairs; result W(P - M).
* ------------------------------------------------------------------------------------- */
/* Sub-block activation sums of the staged chunk. The XOR visit order breaks shared bank
* conflicts and fixes the summation order (result-relevant). Groups of 4 threads share
* the shuffles, and the live test is uniform per group. */
template <int NC, int CHUNK>
__device__ __forceinline__ void fc_q4k_xsums(float *xs, const float *sx, int cend) {
const unsigned groups_per_col = CHUNK / 32;
const int quarter = threadIdx.x & 3;
const unsigned live_groups = (unsigned)cend >> 5;
for (unsigned gi = threadIdx.x >> 2; gi < NC * groups_per_col; gi += FC_KQ_THREADS / 4) {
const unsigned group = gi & (groups_per_col - 1);
if (group >= live_groups)
continue;
const float *src = sx + CHUNK * (gi / groups_per_col) + 32 * group + 8 * quarter;
const unsigned swizzle = gi & 7;
float sum = 0.0f;
#pragma unroll
for (unsigned k = 0; k < 8; k++)
sum += src[k ^ swizzle];
sum += __shfl_down(sum, 2);
sum += __shfl_down(sum, 1);
if (quarter == 0)
xs[gi] = sum;
}
}
/* Min half for one chunk. Lane l owns the pairs whose global index is l mod 32; the
* rotation by cbase/32 keeps that rule for 512-float chunks (and is 0 for 2048 ones). */
template <int NC, int CHUNK>
__device__ __forceinline__ void fc_q4k_min_half(const unsigned char *b0, const unsigned char *b1,
const float *xs, int nsb, int cbase, int lane,
float (&m0)[NC], float (&m1)[NC]) {
const int pi0 = (lane - (cbase >> 5)) & 31;
const unsigned char *blk0 = b0 + (pi0 >> 3) * FC_Q4K_BLOCK_BYTES;
const unsigned char *blk1 = b1 + (pi0 >> 3) * FC_Q4K_BLOCK_BYTES;
const int sub = 136 + (pi0 & 7);
/* At most two trips (one when batched): an unroll ladder would only cost registers. */
#pragma unroll 1
for (int pi = pi0; pi < 8 * nsb; pi += 32) {
const float mn0 = __fmul_rn(__ldg((const float *)(blk0 + 148)), (float)__ldg(blk0 + sub));
const float mn1 = __fmul_rn(__ldg((const float *)(blk1 + 148)), (float)__ldg(blk1 + sub));
#pragma unroll
for (int c = 0; c < NC; c++) {
const float xsum = xs[c * (CHUNK / 32) + pi];
m0[c] = __fmaf_rn(mn0, xsum, m0[c]);
m1[c] = __fmaf_rn(mn1, xsum, m1[c]);
}
blk0 += 4 * FC_Q4K_BLOCK_BYTES;
blk1 += 4 * FC_Q4K_BLOCK_BYTES;
}
}
/* Scale half of one super-block: A = sum(F(sc)*F(q) * x), then P = fma(d, A, P). */
template <int NC, int CHUNK>
__device__ __forceinline__ void fc_q4k_superblock(const unsigned char *b0, const unsigned char *b1,
const float *sx, int lane,
float (&p0)[NC], float (&p1)[NC]) {
const unsigned q0 = __ldg((const unsigned *)b0 + lane);
const unsigned q1 = __ldg((const unsigned *)b1 + lane);
const uint2 sc0 = __ldg((const uint2 *)(b0 + 128));
const uint2 sc1 = __ldg((const uint2 *)(b1 + 128));
const float d0 = __ldg((const float *)(b0 + 144));
const float d1 = __ldg((const float *)(b1 + 144));
const unsigned q0_hi = q0 >> 16, q1_hi = q1 >> 16;
float a0[NC], a1[NC];
#pragma unroll
for (int c = 0; c < NC; c++)
a0[c] = a1[c] = 0.0f;
#pragma unroll
for (int j = 0; j < 4; j++) {
#pragma unroll
for (int h = 0; h < 2; h++) {
const int s = 2 * j + h, pos = 8 * (j & 1) + 4 * h;
const float t0 = fc_kq_byte(sc0, s) * fc_kq_nibble(j < 2 ? q0 : q0_hi, pos, 0.0f);
const float t1 = fc_kq_byte(sc1, s) * fc_kq_nibble(j < 2 ? q1 : q1_hi, pos, 0.0f);
#pragma unroll
for (int c = 0; c < NC; c++) {
const float xv = sx[c * CHUNK + 64 * j + 32 * h + lane];
a0[c] = __fmaf_rn(t0, xv, a0[c]);
a1[c] = __fmaf_rn(t1, xv, a1[c]);
}
}
}
#pragma unroll
for (int c = 0; c < NC; c++) {
p0[c] = __fmaf_rn(d0, a0[c], p0[c]);
p1[c] = __fmaf_rn(d1, a1[c], p1[c]);
}
}
/* ---------------------------------------------------------------------------------------
* Q5_K: straight form. t = ((d*sc)*q) - (dmin*mn), four roundings; A += t*x; P = P + A.
* Scale/min/(d, dmin) are broadcast plain global loads on purpose.
* ------------------------------------------------------------------------------------- */
template <int NC, int CHUNK>
__device__ __forceinline__ void fc_q5k_superblock(const unsigned char *b0, const unsigned char *b1,
const float *sx, int lane,
float (&p0)[NC], float (&p1)[NC]) {
const unsigned q0 = __ldg((const unsigned *)b0 + lane);
const unsigned q1 = __ldg((const unsigned *)b1 + lane);
const unsigned h0 = __ldg(b0 + 128 + lane);
const unsigned h1 = __ldg(b1 + 128 + lane);
const uint2 sc0 = *(const uint2 *)(b0 + 160), sc1 = *(const uint2 *)(b1 + 160);
const uint2 mn0 = *(const uint2 *)(b0 + 168), mn1 = *(const uint2 *)(b1 + 168);
const float2 dd0 = *(const float2 *)(b0 + 176), dd1 = *(const float2 *)(b1 + 176);
const unsigned q0_hi = q0 >> 16, q1_hi = q1 >> 16;
float a0[NC], a1[NC];
#pragma unroll
for (int c = 0; c < NC; c++)
a0[c] = a1[c] = 0.0f;
#pragma unroll
for (int j = 0; j < 4; j++) {
#pragma unroll
for (int h = 0; h < 2; h++) {
const int s = 2 * j + h, bit = 8 * j + 4 * h;
const int pos = bit <= 16 ? bit : bit - 16;
const float n0 = fc_kq_nibble(bit <= 16 ? q0 : q0_hi, pos, 0.0f) +
fc_kq_high16(h0, 1u << s, s);
const float n1 = fc_kq_nibble(bit <= 16 ? q1 : q1_hi, pos, 0.0f) +
fc_kq_high16(h1, 1u << s, s);
const float t0 = __fsub_rn(__fmul_rn(__fmul_rn(dd0.x, fc_kq_byte(sc0, s)), n0),
__fmul_rn(dd0.y, fc_kq_byte(mn0, s)));
const float t1 = __fsub_rn(__fmul_rn(__fmul_rn(dd1.x, fc_kq_byte(sc1, s)), n1),
__fmul_rn(dd1.y, fc_kq_byte(mn1, s)));
#pragma unroll
for (int c = 0; c < NC; c++) {
const float xv = sx[c * CHUNK + 64 * j + 32 * h + lane];
a0[c] = __fmaf_rn(t0, xv, a0[c]);
a1[c] = __fmaf_rn(t1, xv, a1[c]);
}
}
}
#pragma unroll
for (int c = 0; c < NC; c++) {
p0[c] = __fadd_rn(p0[c], a0[c]);
p1[c] = __fadd_rn(p1[c], a1[c]);
}
}
/* ---------------------------------------------------------------------------------------
* Q6_K: A += ((sc - 128) * (q - 32)) * x (exact product), then P = fma(d, A, P).
* ------------------------------------------------------------------------------------- */
template <int NC, int CHUNK>
__device__ __forceinline__ void fc_q6k_superblock(const unsigned char *b0, const unsigned char *b1,
const float *sx, int lane,
float (&p0)[NC], float (&p1)[NC]) {
const unsigned lo0 = __ldg((const unsigned *)b0 + lane);
const unsigned lo1 = __ldg((const unsigned *)b1 + lane);
const unsigned hi0 = __ldg((const unsigned short *)(b0 + 128) + lane);
const unsigned hi1 = __ldg((const unsigned short *)(b1 + 128) + lane);
const float d0 = __ldg((const float *)(b0 + 208));
const float d1 = __ldg((const float *)(b1 + 208));
const uint2 sc0 = __ldg((const uint2 *)(b0 + 192 + 8 * (lane >> 4)));
const uint2 sc1 = __ldg((const uint2 *)(b1 + 192 + 8 * (lane >> 4)));
float a0[NC], a1[NC];
#pragma unroll
for (int c = 0; c < NC; c++)
a0[c] = a1[c] = 0.0f;
#pragma unroll
for (int half = 0; half < 2; half++) {
const unsigned w0 = half ? lo0 >> 16 : lo0, w1 = half ? lo1 >> 16 : lo1;
#pragma unroll
for (int quad = 0; quad < 4; quad++) {
const int i = 4 * half + quad;
const int pos = 8 * (quad & 1) + 4 * (quad >> 1), hpos = 8 * half + 2 * quad;
const float fq0 = fc_kq_nibble(w0, pos, 32.0f) + fc_kq_high16(hi0, 3u << hpos, hpos);
const float fq1 = fc_kq_nibble(w1, pos, 32.0f) + fc_kq_high16(hi1, 3u << hpos, hpos);
const float t0 = __fmul_rn(fc_kq_byte(sc0, i) - 128.0f, fq0);
const float t1 = __fmul_rn(fc_kq_byte(sc1, i) - 128.0f, fq1);
#pragma unroll
for (int c = 0; c < NC; c++) {
const float xv = sx[c * CHUNK + 32 * i + lane];
a0[c] = __fmaf_rn(t0, xv, a0[c]);
a1[c] = __fmaf_rn(t1, xv, a1[c]);
}
}
}
#pragma unroll
for (int c = 0; c < NC; c++) {
p0[c] = __fmaf_rn(d0, a0[c], p0[c]);
p1[c] = __fmaf_rn(d1, a1[c], p1[c]);
}
}
/* ---------------------------------------------------------------------------------------
* K-quant row-pair driver: 8 warps x 2 rows, two independent chains per warp sharing one
* shared-memory x read (Kepler dual-issue). An out-of-range second row re-reads row 0.
* Warps with no live row still run every barrier, staging, x-sum and the warp sum.
* ------------------------------------------------------------------------------------- */
template <int KIND, int NC, int CHUNK>
__device__ __forceinline__ void fc_kq_superblock(const unsigned char *b0, const unsigned char *b1,
const float *sx, int lane,
float (&p0)[NC], float (&p1)[NC]) {
if (KIND == FC_KQ_KIND_Q4K)
fc_q4k_superblock<NC, CHUNK>(b0, b1, sx, lane, p0, p1);
else if (KIND == FC_KQ_KIND_Q5K)
fc_q5k_superblock<NC, CHUNK>(b0, b1, sx, lane, p0, p1);
else
fc_q6k_superblock<NC, CHUNK>(b0, b1, sx, lane, p0, p1);
}
template <int KIND, int NC, int CHUNK>
__device__ __forceinline__ void fc_kq_rowpair_body(const unsigned char *w, const float *x,
float *out, int rows, int cols,
int row_stride, int col0, int nc) {
const int block_bytes = KIND == FC_KQ_KIND_Q4K ? FC_Q4K_BLOCK_BYTES
: KIND == FC_KQ_KIND_Q5K ? FC_Q5K_BLOCK_BYTES : FC_Q6K_BLOCK_BYTES;
const int xsum_per_col = CHUNK == FC_KQ_SINGLE_CHUNK ? FC_KQ_XSUM_FLOATS : FC_NX_XSUM_PER_COL;
const int xsum_floats = KIND == FC_KQ_KIND_Q4K ? NC * xsum_per_col : 0;
float *xs = fc_kq_shared;
float *sx = fc_kq_shared + xsum_floats;
const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31;
const int r0 = 16 * blockIdx.x + 2 * warp;
const bool active0 = r0 < rows, active1 = r0 + 1 < rows;
const unsigned char *b0 = w + (size_t)r0 * (size_t)row_stride;
const unsigned char *b1 = active1 ? b0 + (size_t)row_stride : b0;
float p0[NC], p1[NC], m0[NC], m1[NC];
#pragma unroll
for (int c = 0; c < NC; c++)
p0[c] = p1[c] = m0[c] = m1[c] = 0.0f;
for (int cbase = 0; cbase < cols; cbase += CHUNK) {
const int cend = min(CHUNK, cols - cbase);
__syncthreads();
fc_kq_stage<NC, CHUNK>(sx, x, col0, nc, cols, cbase, cend);
__syncthreads();
if (KIND == FC_KQ_KIND_Q4K) {
fc_q4k_xsums<NC, CHUNK>(xs, sx, cend);
__syncthreads();
}
if (!active0)
continue;
const int nsb = cend >> 8;
if (KIND == FC_KQ_KIND_Q4K)
fc_q4k_min_half<NC, CHUNK>(b0, b1, xs, nsb, cbase, lane, m0, m1);
const float *xb = sx;
for (int sb = 0; sb < nsb; sb++) {
fc_kq_superblock<KIND, NC, CHUNK>(b0, b1, xb, lane, p0, p1);
b0 += block_bytes;
b1 += block_bytes;
xb += FC_KQ_SUPERBLOCK;
}
}
if (KIND == FC_KQ_KIND_Q4K) {
#pragma unroll
for (int c = 0; c < NC; c++) {
p0[c] = __fsub_rn(p0[c], m0[c]);
p1[c] = __fsub_rn(p1[c], m1[c]);
}
}
fc_kq_store_rowpair<NC>(out, p0, p1, rows, r0, active0, active1, col0, nc);
}
/* ---------------------------------------------------------------------------------------
* Q8_0: warp per row, 8 rows per block. t = d * q (own rounding), P = fma(t, x, P).
* ------------------------------------------------------------------------------------- */
template <int NC, int CHUNK>
__device__ __forceinline__ void fc_q80_body(const unsigned char *w, const float *x, float *out,
int rows, int cols, int row_stride, int col0,
int nc) {
float *sx = fc_kq_shared;
const int lane = threadIdx.x & 31;
const int row = 8 * blockIdx.x + (threadIdx.x >> 5);
const bool active = row < rows;
const unsigned char *blk = w + (long long)row * row_stride;
float p[NC];
#pragma unroll
for (int c = 0; c < NC; c++)
p[c] = 0.0f;
for (int cbase = 0; cbase < cols; cbase += CHUNK) {
const int cend = min(CHUNK, cols - cbase);
__syncthreads();
fc_kq_stage<NC, CHUNK>(sx, x, col0, nc, cols, cbase, cend);
__syncthreads();
if (!active)
continue;
const int nblk = cend >> 5;
for (int b = 0; b < nblk; b++) {
const float t = __fmul_rn(__ldg((const float *)blk),
(float)__ldg((const signed char *)blk + 4 + lane));
#pragma unroll
for (int c = 0; c < NC; c++)
p[c] = __fmaf_rn(t, sx[c * CHUNK + 32 * b + lane], p[c]);
blk += FC_Q80_BLOCK_BYTES;
}
}
fc_kq_store_row<NC>(out, p, rows, row, active, col0, nc);
}
/* ---------------------------------------------------------------------------------------
* Q2_0_g128: activation-quantised integer dot per 32-element sub-block.
* ------------------------------------------------------------------------------------- */
/* Round a value >= 0 (or +Inf) to IEEE binary16 with round-to-nearest-even and widen it
* back exactly. Overflow (>= 65520) gives +Inf. Below the binary16 normal range, adding
* 0.5 rounds to a multiple of 2^-24 (the subnormal step) in one IEEE add, ties to even. */
__device__ __forceinline__ float fc_q2_0_round_via_half(float v) {
unsigned bits = __float_as_uint(v);
if (bits >= 0x477FF000u)
return __uint_as_float(0x7F800000u);
if (bits < 0x38800000u)
return __fsub_rn(__fadd_rn(v, 0.5f), 0.5f);
bits += 0x0FFFu + ((bits >> 13) & 1u);
return __uint_as_float(bits & ~0x1FFFu);
}
/* Integer dot of four 2-bit codes (value code-1) with q, wrapping mod 2^32. */
__device__ __forceinline__ int fc_q2_0_code_dot(unsigned codes, const int (&q)[4]) {
unsigned acc = 0;
#pragma unroll
for (int k = 0; k < 4; k++)
acc += (unsigned)((int)((codes >> (2 * k)) & 3u) - 1) * (unsigned)q[k];
return (int)acc;
}
/* Quantise this lane's four activations against its 8-lane sub-block's absmax/127.
* Returns the binary16-rounded activation scale; q receives the rounded integers. */
__device__ __forceinline__ float fc_q2_0_quantise(float4 xq, int (&q)[4]) {
const float xk[4] = { xq.x, xq.y, xq.z, xq.w };
float amax = 0.0f;
#pragma unroll
for (int k = 0; k < 4; k++)
amax = fmaxf(amax, fabsf(xk[k]));
amax = fmaxf(amax, __shfl_xor(amax, 1));
amax = fmaxf(amax, __shfl_xor(amax, 2));
amax = fmaxf(amax, __shfl_xor(amax, 4));
const float d_pre = __fmul_rn(amax, 1.0f / 127.0f);
const float inv = d_pre != 0.0f ? __fdividef(1.0f, d_pre) : 0.0f;
#pragma unroll
for (int k = 0; k < 4; k++)
q[k] = (int)roundf(__fmul_rn(xk[k], inv));
return fc_q2_0_round_via_half(d_pre);
}
__device__ __forceinline__ void fc_q2_0_body(const unsigned char *w, const float *x, float *out,
int rows, int cols, int row_stride) {
float *sx = fc_kq_shared;
const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31;
const int r0 = 16 * blockIdx.x + 2 * warp;
const bool active0 = r0 < rows, active1 = r0 + 1 < rows;
const unsigned char *b0 = w + (long long)r0 * row_stride;
const unsigned char *b1 = active1 ? b0 + row_stride : b0;
float p0[1] = { 0.0f }, p1[1] = { 0.0f };
for (int cbase = 0; cbase < cols; cbase += FC_KQ_SINGLE_CHUNK) {
const int cend = min(FC_KQ_SINGLE_CHUNK, cols - cbase);
__syncthreads();
fc_kq_stage_single(sx, x, cbase, cend);
__syncthreads();
if (!active0)
continue;
const int nblk = cend >> 7;
for (int b = 0; b < nblk; b++) {
int q[4];
const float d_act = fc_q2_0_quantise(((const float4 *)sx)[32 * b + lane], q);
const int dot0 = fc_q2_0_code_dot(__ldg(b0 + 4 + lane), q);
const int dot1 = fc_q2_0_code_dot(__ldg(b1 + 4 + lane), q);
const float d0 = __ldg((const float *)b0), d1 = __ldg((const float *)b1);
p0[0] = __fadd_rn(p0[0], __fmul_rn(d0, __fmul_rn(d_act, (float)dot0)));
p1[0] = __fadd_rn(p1[0], __fmul_rn(d1, __fmul_rn(d_act, (float)dot1)));
b0 += FC_Q2_0_BLOCK_BYTES;
b1 += FC_Q2_0_BLOCK_BYTES;
}
}
fc_kq_store_row<1>(out, p0, rows, r0, active0, 0, 1);
fc_kq_store_row<1>(out, p1, rows, r0 + 1, active1, 0, 1);
}
/* ---------------------------------------------------------------------------------------
* Dense F32/F16/BF16: block per row, strided FMA chain per thread, shared halving tree.
* Plain global loads (no staging). T = blockDim.x must be a power of two for a full sum.
* ------------------------------------------------------------------------------------- */
template <int WTYPE>
__device__ __forceinline__ float fc_dense_weight(const unsigned char *row, int i) {
if (WTYPE == FC_DENSE_F32)
return ((const float *)row)[i];
const unsigned short bits = ((const unsigned short *)row)[i];
if (WTYPE == FC_DENSE_BF16)
return __uint_as_float((unsigned)bits << 16);
float f;
asm("cvt.f32.f16 %0, %1;" : "=f"(f) : "h"(bits));
return f;
}
template <int WTYPE, int NC>
__device__ __forceinline__ void fc_dense_body(const unsigned char *w, const float *x, float *out,
int rows, int cols, int row_stride, int col0,
int nc) {
const int row = blockIdx.x;
if (row >= rows)
return;
const int tid = threadIdx.x, threads = blockDim.x;
const unsigned char *wrow = w + (size_t)row * (size_t)row_stride;
float *sd = fc_kq_shared;
float p[NC];
#pragma unroll
for (int c = 0; c < NC; c++)
p[c] = 0.0f;
for (int i = tid; i < cols; i += threads) {
const float wv = fc_dense_weight<WTYPE>(wrow, i);
#pragma unroll
for (int c = 0; c < NC; c++) {
if (c < nc)
p[c] = __fmaf_rn(wv, x[(size_t)(col0 + c) * cols + i], p[c]);
}
}
#pragma unroll
for (int c = 0; c < NC; c++)
sd[c * threads + tid] = p[c];
__syncthreads();
for (int s = threads / 2; s > 0; s >>= 1) {
if (tid < s) {
#pragma unroll
for (int c = 0; c < NC; c++)
sd[c * threads + tid] += sd[c * threads + tid + s];
}
__syncthreads();
}
if (tid != 0)
return;
for (int c = 0; c < nc; c++)
out[(size_t)(col0 + c) * rows + row] = sd[c * threads];
}
/* ---------------------------------------------------------------------------------------
* Entry points (unmangled). Parameter order is ABI: w, x, out, rows, cols, row_stride[, n].
* ------------------------------------------------------------------------------------- */
/* Uniform <NC, CHUNK> names per weight type, so one macro can emit each family. */
#define FC_KQ_BODY_ALIAS(NAME, BODY) \
template <int NC, int CHUNK> \
__device__ __forceinline__ void NAME(const unsigned char *w, const float *x, \
float *out, int rows, int cols, int row_stride, \
int col0, int nc) { \
BODY(w, x, out, rows, cols, row_stride, col0, nc); \
}
FC_KQ_BODY_ALIAS(fc_q4k_body, (fc_kq_rowpair_body<FC_KQ_KIND_Q4K, NC, CHUNK>))
FC_KQ_BODY_ALIAS(fc_q5k_body, (fc_kq_rowpair_body<FC_KQ_KIND_Q5K, NC, CHUNK>))
FC_KQ_BODY_ALIAS(fc_q6k_body, (fc_kq_rowpair_body<FC_KQ_KIND_Q6K, NC, CHUNK>))
FC_KQ_BODY_ALIAS(fc_f32_body, (fc_dense_body<FC_DENSE_F32, NC>))
FC_KQ_BODY_ALIAS(fc_f16_body, (fc_dense_body<FC_DENSE_F16, NC>))
FC_KQ_BODY_ALIAS(fc_bf16_body, (fc_dense_body<FC_DENSE_BF16, NC>))
#define FC_KQ_BATCHED(NAME, BODY, NC) \
extern "C" __global__ void NAME(const unsigned char *w, const float *x, float *out, \
int rows, int cols, int row_stride, int n) { \
const int col0 = (NC) * blockIdx.y; \
BODY<NC, FC_NX_CHUNK>(w, x, out, rows, cols, row_stride, col0, min((NC), n - col0)); \
}
/* Single-column kernel plus the 2-, 4- and 8-wide batched tiles of one weight type. */
#define FC_KQ_FAMILY(PREFIX, BODY) \
extern "C" __global__ void PREFIX(const unsigned char *w, const float *x, float *out, \
int rows, int cols, int row_stride) { \
BODY<1, FC_KQ_SINGLE_CHUNK>(w, x, out, rows, cols, row_stride, 0, 1); \
} \
FC_KQ_BATCHED(PREFIX##_nx2, BODY, 2) \
FC_KQ_BATCHED(PREFIX##_nx, BODY, 4) \
FC_KQ_BATCHED(PREFIX##_nx8, BODY, 8)
FC_KQ_FAMILY(fc_matvec_q4k, fc_q4k_body)
FC_KQ_FAMILY(fc_matvec_q5k, fc_q5k_body)
FC_KQ_FAMILY(fc_matvec_q6k, fc_q6k_body)
FC_KQ_FAMILY(fc_matvec_q80, fc_q80_body)
FC_KQ_FAMILY(fc_matvec_f32, fc_f32_body)
FC_KQ_FAMILY(fc_matvec_f16, fc_f16_body)
FC_KQ_FAMILY(fc_matvec_bf16, fc_bf16_body)
extern "C" __global__ void fc_matvec_q2_0(const unsigned char *w, const float *x, float *out,
int rows, int cols, int row_stride) {
fc_q2_0_body(w, x, out, rows, cols, row_stride);
}
Reference in New Issue View Git Blame Copy Permalink