732 lines
33 KiB
Plaintext
732 lines
33 KiB
Plaintext
/* SPDX-License-Identifier: AGPL-3.0-or-later */
|
|
/* Copyright (C) 2026 David (Corey) Herrera */
|
|
/*
|
|
* fused_layer.cu — multi-segment fused per-layer launches for Qwen3.5-MoE
|
|
* (Foundry)
|
|
*
|
|
* Target: sm_37 (K80, Kepler)
|
|
* Compile: nvcc -ptx -arch=sm_37 fused_layer.cu -o fused_layer.ptx
|
|
*
|
|
* Decode on the K80 is LAUNCH-BOUND (~80 us per launch vs ~35 us of work), so
|
|
* these kernels collapse the per-layer launch storm without touching any
|
|
* single matvec's internal accumulation order (bit-exactness is the gate):
|
|
*
|
|
* fc_fused_multi — up to FC_FUSED_MAX_SEG independent matvec "segments"
|
|
* that all read the SAME x vector, in ONE launch. A segment is one weight
|
|
* matrix (any repacked quant type or F32) with its own output buffer and
|
|
* an optional fused epilogue. The grid is the concatenation of every
|
|
* segment's blocks (block_base per segment); each block resolves its
|
|
* segment from blockIdx.x and runs EXACTLY the body its stand-alone
|
|
* kernel ran:
|
|
* kind 0 (quant): WARP-PER-ROW-PAIR, 8 warps/block (16 rows/block),
|
|
* staged-x shared chunks, __shfl_down reduce per row —
|
|
* duplicated from matvec_q3k.cu / matvec_kquants.cu
|
|
* (identical fp op order per row; the row pair only adds an
|
|
* independent accumulator chain for Kepler dual-issue ILP).
|
|
* kind 1 (F32): one block per row, 256-thread strided partials +
|
|
* shared-mem tree — duplicated from fc_matvec_f32 (and
|
|
* fc_moe_shexp_dot_gate's dot, same shape).
|
|
* kind 2 (quant gate+up pair): one warp computes the gate AND up dots
|
|
* for its row pair (each dot keeps its own accumulator and the
|
|
* stand-alone per-row order), writes silu(g)*u — the
|
|
* fc_moe_gateup_q3k pattern for the shared expert.
|
|
* Epilogues (applied once to the finished dot by the writing thread —
|
|
* value-identical to the small kernels they replace):
|
|
* 0 none out[row] = v
|
|
* 1 sigmoid out[row] = 1/(1+exp(-v)) (fc_sigmoid)
|
|
* 2 decay gate out[row] = a[row]*softplus(v+dt[row])
|
|
* (fc_ssm_decay_gate)
|
|
* 3 residual add out[row] = v; hid[row] += v (fc_add)
|
|
* 4 moe combine bo = moe[row] + v*sgate[0]; out[row] = bo;
|
|
* hid[row] += bo (fc_moe_shexp_dot_gate scale +
|
|
* fc_add combine + fc_add resid)
|
|
*
|
|
* fc_attn_qknorm_rope — per-head QK-RMSNorm + partial NEOX RoPE for Q and K
|
|
* in ONE launch (was 4: qk_rms_norm x2 + rope x2). One block per head
|
|
* (Q heads first, then K heads); the norm body is fc_qk_rms_norm
|
|
* verbatim, then one barrier, then the fc_rope_neox_partial pair rotation
|
|
* verbatim. Heads are independent, so fusing preserves per-head math.
|
|
*
|
|
* REPACKED WEIGHT LAYOUTS: see matvec_q3k.cu / matvec_kquants.cu. The lane
|
|
* dots below are duplicated from those files (PTX modules are standalone) and
|
|
* keep the identical fp op order, so results stay bit-identical.
|
|
*
|
|
* PTX symbols emitted (for embedding into ../cuda/fc_cuda_ptx.h):
|
|
* FC_PTX_FUSED_LAYER -> kernels fc_fused_multi, fc_attn_qknorm_rope
|
|
*
|
|
* Kepler notes: pure fp32; __shfl_down is the non-_sync Kepler variant.
|
|
*/
|
|
|
|
extern "C" {
|
|
|
|
/* Repacked block sizes (bytes) — must match fc_repack_weight. */
|
|
#define FC_FU_Q3K_RB 152
|
|
#define FC_FU_Q4K_RB 152
|
|
#define FC_FU_Q5K_RB 184
|
|
#define FC_FU_Q6K_RB 216
|
|
#define FC_FU_Q80_RB 36
|
|
|
|
/* Magic decodes (duplicated from matvec_q3k.cu / matvec_kquants.cu): exact
|
|
* float (v - bias) from the masked field at bit p, AND + OR + one FMA — no
|
|
* shift, no cvt (sm_37 runs int-mul / int32->f32 cvt at 32/clk, shifts at
|
|
* 64/clk vs AND/OR 160 and FMUL/FMA 192). BIT16/BIT2X16 scale a 1-/2-bit
|
|
* field by 16 (the K-quant high-bit weight). */
|
|
#define FC_MAGIC_NIB(w, p, bias) \
|
|
__fmaf_rn(__int_as_float(0x4B000000u | ((w) & (0xFu << (p)))), \
|
|
1.0f / (float)(1u << (p)), \
|
|
-(float)((1u << (23 - (p))) + (bias)))
|
|
#define FC_MAGIC_BYTE(w, p, bias) \
|
|
__fmaf_rn(__int_as_float(0x4B000000u | ((w) & (0xFFu << (p)))), \
|
|
1.0f / (float)(1u << (p)), \
|
|
-(float)((1u << (23 - (p))) + (bias)))
|
|
#define FC_MAGIC_BIT16(w, p) \
|
|
__fmaf_rn(__int_as_float(0x4B000000u | ((w) & (1u << (p)))), \
|
|
16.0f / (float)(1u << (p)), -(float)(1u << (27 - (p))))
|
|
#define FC_MAGIC_BIT2X16(w, p) \
|
|
__fmaf_rn(__int_as_float(0x4B000000u | ((w) & (3u << (p)))), \
|
|
16.0f / (float)(1u << (p)), -(float)(1u << (27 - (p))))
|
|
|
|
/* Device-side segment descriptor — layout must match fc_fused_seg_dev in
|
|
* fc_cuda.c (6 x u64 then 8 x s32, 80 bytes, 8-byte aligned). */
|
|
typedef struct {
|
|
unsigned long long w; /* weight bytes (repacked quant or raw F32) */
|
|
unsigned long long w2; /* kind 2: the up weight; else 0 */
|
|
unsigned long long out; /* [rows] f32 output */
|
|
unsigned long long aux0; /* ep 2: dt[rows]; ep 4: moe_out[rows] */
|
|
unsigned long long aux1; /* ep 2: a[rows]; ep 3/4: hidden[rows] */
|
|
unsigned long long aux2; /* ep 4: sgate[1] */
|
|
int rows, cols, row_stride;
|
|
int kind; /* 0 quant, 1 f32, 2 quant gate+up */
|
|
int qtype; /* GGUF id: 8, 11, 12, 13, 14 (kind 0/2) */
|
|
int ep; /* epilogue code (see header comment) */
|
|
int block_base; /* first grid block of this segment */
|
|
int pad;
|
|
} fc_fused_seg_dev;
|
|
|
|
/* Byte-aligned fields decode fastest through Kepler's 128/clk 8-bit cvt
|
|
* (single I2F.F32.U8 with byte select — cheaper than the magic FMA, which
|
|
* stays for the sub-byte fields). */
|
|
#define FC_CVT_BYTE(w, k) ((float)(((w) >> (8 * (k))) & 0xFFu))
|
|
|
|
__device__ __forceinline__ float fc_fu_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;
|
|
}
|
|
|
|
/* ---- Lane accumulators, duplicated from matvec_q3k.cu / matvec_kquants.cu.
|
|
*
|
|
* PINNED FP OP ORDER: the stand-alone kernels compiled (nvcc 11.4, default
|
|
* -fmad=true) to specific contraction shapes, and inlining them into this
|
|
* multi-kernel let nvcc re-fold the caller-level `partial += dot` into the
|
|
* element chain (different summation grouping = different bits — caught by
|
|
* the argmax/logit gate). Every accumulation below is therefore written with
|
|
* explicit __fmaf_rn/__fmul_rn/__fsub_rn/__fadd_rn intrinsics that reproduce
|
|
* the stand-alone kernels' EXACT PTX sequences (the magic decodes and the
|
|
* exact small-integer FMULs produce bit-identical term values to the old
|
|
* cvt-based decode):
|
|
* Q3_K/Q6_K: 8x acc = fma(mul(fs,fq), xb, acc) from 0.0f, then
|
|
* partial = fma(d, acc, partial) (mul(fs,fq) = exact int
|
|
* product, bit-identical to the old cvt(scale*q))
|
|
* Q4_K/Q5_K: 8x { mul(d,sc), mul(.,nib), mul(dmin,mn), sub,
|
|
* acc = fma(sub, xb, acc) } from 0.0f, then
|
|
* partial = add(partial, acc)
|
|
* Q8_0: partial = fma(mul(d, (float)q), xb, partial)
|
|
* Each function takes the running partial and returns the updated one. ---- */
|
|
|
|
__device__ __forceinline__ float fc_fu_q3k_acc(const unsigned char *blk,
|
|
const float *xb, int lane,
|
|
float partial) {
|
|
unsigned int wv = __ldg((const unsigned int *)blk + lane);
|
|
float d_all = __ldg((const float *)(blk + 144));
|
|
uint2 scw = __ldg((const uint2 *)(blk + 128) + (lane >> 4));
|
|
unsigned int wh = wv >> 16;
|
|
|
|
float acc = 0.0f;
|
|
#pragma unroll
|
|
for (int i = 0; i < 8; i++) {
|
|
/* nibble at bit 4i: direct mask while it stays below the fp32
|
|
* exponent (p <= 16), else from the >>16 half */
|
|
unsigned int qsrc = (i <= 4) ? wv : wh;
|
|
int qp = (i <= 4) ? 4 * i : 4 * i - 16;
|
|
unsigned int ssrc = (i < 4) ? scw.x : scw.y;
|
|
float fq = FC_MAGIC_NIB(qsrc, qp, 4);
|
|
float fs = __fsub_rn(FC_CVT_BYTE(ssrc, i & 3), 32.0f);
|
|
acc = __fmaf_rn(__fmul_rn(fs, fq), xb[32 * i + lane], acc);
|
|
}
|
|
return __fmaf_rn(d_all, acc, partial);
|
|
}
|
|
|
|
__device__ __forceinline__ float fc_fu_q4k_acc(const unsigned char *blk,
|
|
const float *xb, int lane,
|
|
float partial) {
|
|
unsigned int qw = __ldg((const unsigned int *)blk + lane);
|
|
uint2 scw = __ldg((const uint2 *)(blk + 128));
|
|
uint2 mnw = __ldg((const uint2 *)(blk + 136));
|
|
float2 dd = __ldg((const float2 *)(blk + 144));
|
|
float d = dd.x, dmin = dd.y;
|
|
unsigned int qwh = qw >> 16;
|
|
float acc = 0.0f;
|
|
#pragma unroll
|
|
for (int j = 0; j < 4; j++) {
|
|
#pragma unroll
|
|
for (int h = 0; h < 2; h++) {
|
|
int si = 2 * j + h;
|
|
int nbit = 8 * j + 4 * h;
|
|
unsigned int qsrc = (nbit <= 16) ? qw : qwh;
|
|
int qp = (nbit <= 16) ? nbit : nbit - 16;
|
|
unsigned int ssrc = (si < 4) ? scw.x : scw.y;
|
|
unsigned int msrc = (si < 4) ? mnw.x : mnw.y;
|
|
float fn = FC_MAGIC_NIB(qsrc, qp, 0);
|
|
float sc = FC_CVT_BYTE(ssrc, si & 3);
|
|
float mn = FC_CVT_BYTE(msrc, si & 3);
|
|
float t = __fsub_rn(__fmul_rn(__fmul_rn(d, sc), fn),
|
|
__fmul_rn(dmin, mn));
|
|
acc = __fmaf_rn(t, xb[64 * j + 32 * h + lane], acc);
|
|
}
|
|
}
|
|
return __fadd_rn(partial, acc);
|
|
}
|
|
|
|
__device__ __forceinline__ float fc_fu_q5k_acc(const unsigned char *blk,
|
|
const float *xb, int lane,
|
|
float partial) {
|
|
unsigned int qw = __ldg((const unsigned int *)blk + lane);
|
|
unsigned int qh = __ldg(blk + 128 + lane);
|
|
uint2 scw = __ldg((const uint2 *)(blk + 160));
|
|
uint2 mnw = __ldg((const uint2 *)(blk + 168));
|
|
float2 dd = __ldg((const float2 *)(blk + 176));
|
|
float d = dd.x, dmin = dd.y;
|
|
unsigned int qwh = qw >> 16;
|
|
float acc = 0.0f;
|
|
#pragma unroll
|
|
for (int j = 0; j < 4; j++) {
|
|
#pragma unroll
|
|
for (int h = 0; h < 2; h++) {
|
|
int si = 2 * j + h;
|
|
int nbit = 8 * j + 4 * h;
|
|
unsigned int qsrc = (nbit <= 16) ? qw : qwh;
|
|
int qp = (nbit <= 16) ? nbit : nbit - 16;
|
|
unsigned int ssrc = (si < 4) ? scw.x : scw.y;
|
|
unsigned int msrc = (si < 4) ? mnw.x : mnw.y;
|
|
float fn = FC_MAGIC_NIB(qsrc, qp, 0);
|
|
float fhi = FC_MAGIC_BIT16(qh, si); /* 16 * high bit */
|
|
float sc = FC_CVT_BYTE(ssrc, si & 3);
|
|
float mn = FC_CVT_BYTE(msrc, si & 3);
|
|
float t = __fsub_rn(__fmul_rn(__fmul_rn(d, sc),
|
|
__fadd_rn(fn, fhi)),
|
|
__fmul_rn(dmin, mn));
|
|
acc = __fmaf_rn(t, xb[64 * j + 32 * h + lane], acc);
|
|
}
|
|
}
|
|
return __fadd_rn(partial, acc);
|
|
}
|
|
|
|
__device__ __forceinline__ float fc_fu_q6k_acc(const unsigned char *blk,
|
|
const float *xb, int lane,
|
|
float partial) {
|
|
unsigned int qlw = __ldg((const unsigned int *)blk + lane);
|
|
unsigned int qh2 = __ldg((const unsigned short *)(blk + 128) + lane);
|
|
float d = __ldg((const float *)(blk + 208));
|
|
uint2 scw = __ldg((const uint2 *)(blk + 192) + (lane >> 4));
|
|
unsigned int qlh = qlw >> 16;
|
|
float acc = 0.0f;
|
|
#pragma unroll
|
|
for (int half = 0; half < 2; half++) {
|
|
#pragma unroll
|
|
for (int quad = 0; quad < 4; quad++) {
|
|
int i = 4 * half + quad;
|
|
int nbit = 16 * half + 8 * (quad & 1) + 4 * ((quad >> 1) & 1);
|
|
int hbit = 8 * half + 2 * quad;
|
|
unsigned int qsrc = (nbit <= 16) ? qlw : qlh;
|
|
int qp = (nbit <= 16) ? nbit : nbit - 16;
|
|
unsigned int ssrc = (i < 4) ? scw.x : scw.y;
|
|
float fn = FC_MAGIC_NIB(qsrc, qp, 32); /* nib - 32 */
|
|
float fhb = FC_MAGIC_BIT2X16(qh2, hbit); /* 16 * hb */
|
|
float fs = __fsub_rn(FC_CVT_BYTE(ssrc, i & 3), 128.0f);
|
|
float fq = __fadd_rn(fn, fhb); /* q, exact integer */
|
|
acc = __fmaf_rn(__fmul_rn(fs, fq), xb[32 * i + lane], acc);
|
|
}
|
|
}
|
|
return __fmaf_rn(d, acc, partial);
|
|
}
|
|
|
|
/*
|
|
* Interleaved ROW-PAIR lane dot: the two rows' 8-element FMA chains advance
|
|
* in lockstep per element (two independent chains the scheduler can
|
|
* dual-issue, unlike two back-to-back dependent 8-FMA chains), sharing one
|
|
* shared-memory x load per element. Each row's per-lane chain keeps the
|
|
* exact i = 0..7 order of fc_q3k_sb_lane_dot -> bit-exact per row.
|
|
*/
|
|
__device__ __forceinline__ void fc_fu_q3k_acc2(const unsigned char *blk0,
|
|
const unsigned char *blk1,
|
|
const float *xb, int lane,
|
|
float *p0, float *p1) {
|
|
unsigned int wv0 = __ldg((const unsigned int *)blk0 + lane);
|
|
unsigned int wv1 = __ldg((const unsigned int *)blk1 + lane);
|
|
float d0 = __ldg((const float *)(blk0 + 144));
|
|
float d1 = __ldg((const float *)(blk1 + 144));
|
|
uint2 sw0 = __ldg((const uint2 *)(blk0 + 128) + (lane >> 4));
|
|
uint2 sw1 = __ldg((const uint2 *)(blk1 + 128) + (lane >> 4));
|
|
unsigned int wh0 = wv0 >> 16, wh1 = wv1 >> 16;
|
|
|
|
float a0 = 0.0f, a1 = 0.0f;
|
|
#pragma unroll
|
|
for (int i = 0; i < 8; i++) {
|
|
int qp = (i <= 4) ? 4 * i : 4 * i - 16;
|
|
unsigned int q0 = (i <= 4) ? wv0 : wh0;
|
|
unsigned int q1 = (i <= 4) ? wv1 : wh1;
|
|
unsigned int s0 = (i < 4) ? sw0.x : sw0.y;
|
|
unsigned int s1 = (i < 4) ? sw1.x : sw1.y;
|
|
float xv = xb[32 * i + lane];
|
|
float fq0 = FC_MAGIC_NIB(q0, qp, 4);
|
|
float fs0 = __fsub_rn(FC_CVT_BYTE(s0, i & 3), 32.0f);
|
|
float fq1 = FC_MAGIC_NIB(q1, qp, 4);
|
|
float fs1 = __fsub_rn(FC_CVT_BYTE(s1, i & 3), 32.0f);
|
|
a0 = __fmaf_rn(__fmul_rn(fs0, fq0), xv, a0);
|
|
a1 = __fmaf_rn(__fmul_rn(fs1, fq1), xv, a1);
|
|
}
|
|
*p0 = __fmaf_rn(d0, a0, *p0);
|
|
*p1 = __fmaf_rn(d1, a1, *p1);
|
|
}
|
|
|
|
|
|
/* Interleaved ROW-PAIR Q4_K lane dot: both rows' chains advance per element
|
|
* (independent chains for dual-issue), sharing one x load. Each row keeps
|
|
* the exact per-element op order of fc_q4k_sb_lane_dot -> bit-exact. */
|
|
__device__ __forceinline__ void fc_fu_q4k_acc2(const unsigned char *blk0,
|
|
const unsigned char *blk1,
|
|
const float *xb, int lane,
|
|
float *p0, float *p1) {
|
|
unsigned int qw0 = __ldg((const unsigned int *)blk0 + lane);
|
|
unsigned int qw1 = __ldg((const unsigned int *)blk1 + lane);
|
|
uint2 sw0 = __ldg((const uint2 *)(blk0 + 128));
|
|
uint2 sw1 = __ldg((const uint2 *)(blk1 + 128));
|
|
uint2 mw0 = __ldg((const uint2 *)(blk0 + 136));
|
|
uint2 mw1 = __ldg((const uint2 *)(blk1 + 136));
|
|
float2 dd0 = __ldg((const float2 *)(blk0 + 144));
|
|
float2 dd1 = __ldg((const float2 *)(blk1 + 144));
|
|
unsigned int qh0 = qw0 >> 16, qh1 = qw1 >> 16;
|
|
float a0 = 0.0f, a1 = 0.0f;
|
|
#pragma unroll
|
|
for (int j = 0; j < 4; j++) {
|
|
#pragma unroll
|
|
for (int h = 0; h < 2; h++) {
|
|
int si = 2 * j + h;
|
|
int nbit = 8 * j + 4 * h;
|
|
int qp = (nbit <= 16) ? nbit : nbit - 16;
|
|
unsigned int q0 = (nbit <= 16) ? qw0 : qh0;
|
|
unsigned int q1 = (nbit <= 16) ? qw1 : qh1;
|
|
unsigned int s0 = (si < 4) ? sw0.x : sw0.y;
|
|
unsigned int s1 = (si < 4) ? sw1.x : sw1.y;
|
|
unsigned int m0 = (si < 4) ? mw0.x : mw0.y;
|
|
unsigned int m1 = (si < 4) ? mw1.x : mw1.y;
|
|
float xv = xb[64 * j + 32 * h + lane];
|
|
float t0 = __fsub_rn(
|
|
__fmul_rn(__fmul_rn(dd0.x, FC_CVT_BYTE(s0, si & 3)),
|
|
FC_MAGIC_NIB(q0, qp, 0)),
|
|
__fmul_rn(dd0.y, FC_CVT_BYTE(m0, si & 3)));
|
|
float t1 = __fsub_rn(
|
|
__fmul_rn(__fmul_rn(dd1.x, FC_CVT_BYTE(s1, si & 3)),
|
|
FC_MAGIC_NIB(q1, qp, 0)),
|
|
__fmul_rn(dd1.y, FC_CVT_BYTE(m1, si & 3)));
|
|
a0 = __fmaf_rn(t0, xv, a0);
|
|
a1 = __fmaf_rn(t1, xv, a1);
|
|
}
|
|
}
|
|
*p0 = __fadd_rn(*p0, a0);
|
|
*p1 = __fadd_rn(*p1, a1);
|
|
}
|
|
|
|
|
|
/* Interleaved ROW-PAIR Q5_K lane dot (see fc_q4k_sb_lane_dot2). */
|
|
__device__ __forceinline__ void fc_fu_q5k_acc2(const unsigned char *blk0,
|
|
const unsigned char *blk1,
|
|
const float *xb, int lane,
|
|
float *p0, float *p1) {
|
|
unsigned int qw0 = __ldg((const unsigned int *)blk0 + lane);
|
|
unsigned int qw1 = __ldg((const unsigned int *)blk1 + lane);
|
|
unsigned int qb0 = __ldg(blk0 + 128 + lane);
|
|
unsigned int qb1 = __ldg(blk1 + 128 + lane);
|
|
uint2 sw0 = __ldg((const uint2 *)(blk0 + 160));
|
|
uint2 sw1 = __ldg((const uint2 *)(blk1 + 160));
|
|
uint2 mw0 = __ldg((const uint2 *)(blk0 + 168));
|
|
uint2 mw1 = __ldg((const uint2 *)(blk1 + 168));
|
|
float2 dd0 = __ldg((const float2 *)(blk0 + 176));
|
|
float2 dd1 = __ldg((const float2 *)(blk1 + 176));
|
|
unsigned int qh0 = qw0 >> 16, qh1 = qw1 >> 16;
|
|
float a0 = 0.0f, a1 = 0.0f;
|
|
#pragma unroll
|
|
for (int j = 0; j < 4; j++) {
|
|
#pragma unroll
|
|
for (int h = 0; h < 2; h++) {
|
|
int si = 2 * j + h;
|
|
int nbit = 8 * j + 4 * h;
|
|
int qp = (nbit <= 16) ? nbit : nbit - 16;
|
|
unsigned int q0 = (nbit <= 16) ? qw0 : qh0;
|
|
unsigned int q1 = (nbit <= 16) ? qw1 : qh1;
|
|
unsigned int s0 = (si < 4) ? sw0.x : sw0.y;
|
|
unsigned int s1 = (si < 4) ? sw1.x : sw1.y;
|
|
unsigned int m0 = (si < 4) ? mw0.x : mw0.y;
|
|
unsigned int m1 = (si < 4) ? mw1.x : mw1.y;
|
|
float xv = xb[64 * j + 32 * h + lane];
|
|
float n0 = __fadd_rn(FC_MAGIC_NIB(q0, qp, 0),
|
|
FC_MAGIC_BIT16(qb0, si));
|
|
float n1 = __fadd_rn(FC_MAGIC_NIB(q1, qp, 0),
|
|
FC_MAGIC_BIT16(qb1, si));
|
|
float t0 = __fsub_rn(
|
|
__fmul_rn(__fmul_rn(dd0.x, FC_CVT_BYTE(s0, si & 3)), n0),
|
|
__fmul_rn(dd0.y, FC_CVT_BYTE(m0, si & 3)));
|
|
float t1 = __fsub_rn(
|
|
__fmul_rn(__fmul_rn(dd1.x, FC_CVT_BYTE(s1, si & 3)), n1),
|
|
__fmul_rn(dd1.y, FC_CVT_BYTE(m1, si & 3)));
|
|
a0 = __fmaf_rn(t0, xv, a0);
|
|
a1 = __fmaf_rn(t1, xv, a1);
|
|
}
|
|
}
|
|
*p0 = __fadd_rn(*p0, a0);
|
|
*p1 = __fadd_rn(*p1, a1);
|
|
}
|
|
|
|
|
|
/* Interleaved ROW-PAIR Q6_K lane dot (see fc_q4k_sb_lane_dot2). */
|
|
__device__ __forceinline__ void fc_fu_q6k_acc2(const unsigned char *blk0,
|
|
const unsigned char *blk1,
|
|
const float *xb, int lane,
|
|
float *p0, float *p1) {
|
|
unsigned int ql0 = __ldg((const unsigned int *)blk0 + lane);
|
|
unsigned int ql1 = __ldg((const unsigned int *)blk1 + lane);
|
|
unsigned int qh20 = __ldg((const unsigned short *)(blk0 + 128) + lane);
|
|
unsigned int qh21 = __ldg((const unsigned short *)(blk1 + 128) + lane);
|
|
float d0 = __ldg((const float *)(blk0 + 208));
|
|
float d1 = __ldg((const float *)(blk1 + 208));
|
|
uint2 sw0 = __ldg((const uint2 *)(blk0 + 192) + (lane >> 4));
|
|
uint2 sw1 = __ldg((const uint2 *)(blk1 + 192) + (lane >> 4));
|
|
unsigned int qlh0 = ql0 >> 16, qlh1 = ql1 >> 16;
|
|
float a0 = 0.0f, a1 = 0.0f;
|
|
#pragma unroll
|
|
for (int half = 0; half < 2; half++) {
|
|
#pragma unroll
|
|
for (int quad = 0; quad < 4; quad++) {
|
|
int i = 4 * half + quad;
|
|
int nbit = 16 * half + 8 * (quad & 1) + 4 * ((quad >> 1) & 1);
|
|
int hbit = 8 * half + 2 * quad;
|
|
int qp = (nbit <= 16) ? nbit : nbit - 16;
|
|
unsigned int q0 = (nbit <= 16) ? ql0 : qlh0;
|
|
unsigned int q1 = (nbit <= 16) ? ql1 : qlh1;
|
|
unsigned int s0 = (i < 4) ? sw0.x : sw0.y;
|
|
unsigned int s1 = (i < 4) ? sw1.x : sw1.y;
|
|
float xv = xb[32 * i + lane];
|
|
float fq0 = __fadd_rn(FC_MAGIC_NIB(q0, qp, 32),
|
|
FC_MAGIC_BIT2X16(qh20, hbit));
|
|
float fq1 = __fadd_rn(FC_MAGIC_NIB(q1, qp, 32),
|
|
FC_MAGIC_BIT2X16(qh21, hbit));
|
|
float fs0 = __fsub_rn(FC_CVT_BYTE(s0, i & 3), 128.0f);
|
|
float fs1 = __fsub_rn(FC_CVT_BYTE(s1, i & 3), 128.0f);
|
|
a0 = __fmaf_rn(__fmul_rn(fs0, fq0), xv, a0);
|
|
a1 = __fmaf_rn(__fmul_rn(fs1, fq1), xv, a1);
|
|
}
|
|
}
|
|
*p0 = __fmaf_rn(d0, a0, *p0);
|
|
*p1 = __fmaf_rn(d1, a1, *p1);
|
|
}
|
|
|
|
|
|
/* Super-block lane accumulate dispatch (256-elem super-block types). qtype
|
|
* is uniform per block, so the switch never diverges within a warp. */
|
|
__device__ __forceinline__ float fc_fu_sb_acc(int qtype, const unsigned char *blk,
|
|
const float *xb, int lane,
|
|
float partial) {
|
|
switch (qtype) {
|
|
case 11: return fc_fu_q3k_acc(blk, xb, lane, partial);
|
|
case 12: return fc_fu_q4k_acc(blk, xb, lane, partial);
|
|
case 13: return fc_fu_q5k_acc(blk, xb, lane, partial);
|
|
default: return fc_fu_q6k_acc(blk, xb, lane, partial); /* 14 */
|
|
}
|
|
}
|
|
|
|
/* Element-interleaved pair dispatch (two rows / two matrices per call —
|
|
* two independent accumulator chains for Kepler dual-issue). */
|
|
__device__ __forceinline__ void fc_fu_sb_acc2(int qtype,
|
|
const unsigned char *blk0,
|
|
const unsigned char *blk1,
|
|
const float *xb, int lane,
|
|
float *p0, float *p1) {
|
|
switch (qtype) {
|
|
case 11: fc_fu_q3k_acc2(blk0, blk1, xb, lane, p0, p1); break;
|
|
case 12: fc_fu_q4k_acc2(blk0, blk1, xb, lane, p0, p1); break;
|
|
case 13: fc_fu_q5k_acc2(blk0, blk1, xb, lane, p0, p1); break;
|
|
default: fc_fu_q6k_acc2(blk0, blk1, xb, lane, p0, p1); break; /* 14 */
|
|
}
|
|
}
|
|
|
|
__device__ __forceinline__ int fc_fu_sb_bytes(int qtype) {
|
|
switch (qtype) {
|
|
case 13: return FC_FU_Q5K_RB;
|
|
case 14: return FC_FU_Q6K_RB;
|
|
default: return FC_FU_Q3K_RB; /* 11, 12 (both 152) */
|
|
}
|
|
}
|
|
|
|
/* Fused epilogue — one store of the finished per-row value v. Each case is
|
|
* value-identical to the stand-alone kernel it replaces (see file header). */
|
|
__device__ __forceinline__ void fc_fu_store(const fc_fused_seg_dev *sg,
|
|
int row, float v) {
|
|
float *out = (float *)(size_t)sg->out;
|
|
int ep = sg->ep;
|
|
if (ep == 1) { /* sigmoid (fc_sigmoid) */
|
|
out[row] = 1.0f / (1.0f + expf(-v));
|
|
} else if (ep == 2) { /* decay gate (fc_ssm_decay_gate) */
|
|
const float *dt = (const float *)(size_t)sg->aux0;
|
|
const float *a = (const float *)(size_t)sg->aux1;
|
|
float t = v + dt[row];
|
|
float sp = (t > 20.0f) ? t : log1pf(expf(t));
|
|
out[row] = a[row] * sp;
|
|
} else if (ep == 3) { /* residual add (fc_add) */
|
|
float *hid = (float *)(size_t)sg->aux1;
|
|
out[row] = v;
|
|
hid[row] = hid[row] + v;
|
|
} else if (ep == 4) { /* moe combine + residual */
|
|
const float *moe = (const float *)(size_t)sg->aux0;
|
|
float *hid = (float *)(size_t)sg->aux1;
|
|
const float *sgt = (const float *)(size_t)sg->aux2;
|
|
/* Pinned mul-then-add: the stand-alone path was a mul kernel
|
|
* (shexp *= sgate) then an add kernel (moe_out + shexp) — keep the
|
|
* same two ops (no fma contraction). */
|
|
float bo = __fadd_rn(moe[row], __fmul_rn(v, sgt[0]));
|
|
out[row] = bo;
|
|
hid[row] = __fadd_rn(hid[row], bo);
|
|
} else {
|
|
out[row] = v;
|
|
}
|
|
}
|
|
|
|
/*
|
|
* Multi-segment fused matvec. All segments read the same x[.] (each segment
|
|
* uses its own cols). Launch: grid = sum of per-segment blocks (kind 0/2:
|
|
* ceil(rows/16), 16 rows per block; kind 1: rows), block = 256, shared =
|
|
* max over segments (kind 0/2: min(cols,2048)*4; kind 1: 1024).
|
|
*/
|
|
__global__ void fc_fused_multi(const fc_fused_seg_dev *segs, int nseg,
|
|
const float *x) {
|
|
extern __shared__ float sh[];
|
|
|
|
/* Resolve this block's segment (block_base ascending). */
|
|
int s = 0;
|
|
for (int t = 1; t < nseg; t++)
|
|
if (blockIdx.x >= segs[t].block_base)
|
|
s = t;
|
|
const fc_fused_seg_dev *sg = &segs[s];
|
|
const int rows = sg->rows;
|
|
const int cols = sg->cols;
|
|
const int rstride = sg->row_stride;
|
|
const int kind = sg->kind;
|
|
const int brel = blockIdx.x - sg->block_base;
|
|
|
|
if (kind == 1) {
|
|
/* F32 block-per-row — fc_matvec_f32 verbatim (256-thread strided
|
|
* partials + tree; also matches fc_moe_shexp_dot_gate's dot). */
|
|
int row = brel;
|
|
if (row >= rows)
|
|
return;
|
|
int tid = threadIdx.x;
|
|
int threads = blockDim.x;
|
|
const float *w_row = (const float *)
|
|
((const unsigned char *)(size_t)sg->w + (size_t)row * (size_t)rstride);
|
|
float partial = 0.0f;
|
|
for (int i = tid; i < cols; i += threads)
|
|
partial = __fmaf_rn(w_row[i], x[i], partial); /* fc_matvec_f32's
|
|
* fma chain */
|
|
sh[tid] = partial;
|
|
__syncthreads();
|
|
for (int t = threads / 2; t > 0; t >>= 1) {
|
|
if (tid < t)
|
|
sh[tid] += sh[tid + t];
|
|
__syncthreads();
|
|
}
|
|
if (tid == 0)
|
|
fc_fu_store(sg, row, sh[0]);
|
|
return;
|
|
}
|
|
|
|
/* Quant WARP-PER-ROW-PAIR (kind 0) / gate+up pair (kind 2) — staged-x
|
|
* chunk loop duplicated from matvec_q3k.cu / matvec_kquants.cu. Warp w
|
|
* handles rows brel*16 + 2w + {0,1} with independent accumulator chains
|
|
* (dual-issue ILP; staged-x LDS amortized across the pair); each row's
|
|
* per-lane chain and reduction are unchanged -> bit-exact per row. */
|
|
int warp = threadIdx.x >> 5;
|
|
int lane = threadIdx.x & 31;
|
|
int row0 = brel * 16 + warp * 2;
|
|
int active0 = row0 < rows;
|
|
int active1 = row0 + 1 < rows;
|
|
const int qtype = sg->qtype;
|
|
/* strength-reduced running block pointers (advanced per quant block) */
|
|
const unsigned char *w_row0 =
|
|
(const unsigned char *)(size_t)sg->w + (size_t)row0 * (size_t)rstride;
|
|
/* tail: clamp the pair onto row0, duplicate result discarded */
|
|
const unsigned char *w_row1 = active1 ? w_row0 + (size_t)rstride : w_row0;
|
|
const unsigned char *u_row0 =
|
|
(const unsigned char *)(size_t)sg->w2 + (size_t)row0 * (size_t)rstride;
|
|
const unsigned char *u_row1 = active1 ? u_row0 + (size_t)rstride : u_row0;
|
|
|
|
float pg0 = 0.0f, pu0 = 0.0f, pg1 = 0.0f, pu1 = 0.0f;
|
|
for (int cbase = 0; cbase < cols; cbase += 2048) {
|
|
int cend = cols - cbase < 2048 ? cols - cbase : 2048;
|
|
__syncthreads();
|
|
for (int i = threadIdx.x; i < cend; i += 256)
|
|
sh[i] = __ldg(x + cbase + i);
|
|
__syncthreads();
|
|
if (active0) {
|
|
if (qtype == 8) {
|
|
/* Q8_0: 32-elem blocks, lane l covers element l per block
|
|
* (fc_matvec_q80 verbatim per row). */
|
|
int nb = cend >> 5;
|
|
const float *xb = sh;
|
|
for (int b = 0; b < nb; b++) {
|
|
float d = __ldg((const float *)w_row0);
|
|
int q = (int)__ldg((const signed char *)(w_row0 + 4) + lane);
|
|
/* fc_matvec_q80's mul-then-fma, pinned. */
|
|
pg0 = __fmaf_rn(__fmul_rn(d, (float)q), xb[lane], pg0);
|
|
if (active1) {
|
|
float d1 = __ldg((const float *)w_row1);
|
|
int q1 = (int)__ldg((const signed char *)(w_row1 + 4) + lane);
|
|
pg1 = __fmaf_rn(__fmul_rn(d1, (float)q1), xb[lane], pg1);
|
|
}
|
|
if (kind == 2) {
|
|
float du = __ldg((const float *)u_row0);
|
|
int qu = (int)__ldg((const signed char *)(u_row0 + 4) + lane);
|
|
pu0 = __fmaf_rn(__fmul_rn(du, (float)qu), xb[lane], pu0);
|
|
if (active1) {
|
|
float du1 = __ldg((const float *)u_row1);
|
|
int qu1 = (int)__ldg((const signed char *)(u_row1 + 4) + lane);
|
|
pu1 = __fmaf_rn(__fmul_rn(du1, (float)qu1),
|
|
xb[lane], pu1);
|
|
}
|
|
}
|
|
w_row0 += FC_FU_Q80_RB;
|
|
w_row1 += FC_FU_Q80_RB;
|
|
if (kind == 2) {
|
|
u_row0 += FC_FU_Q80_RB;
|
|
u_row1 += FC_FU_Q80_RB;
|
|
}
|
|
xb += 32;
|
|
}
|
|
} else {
|
|
int bb = fc_fu_sb_bytes(qtype);
|
|
int nsb = cend >> 8;
|
|
const float *xb = sh;
|
|
for (int sb = 0; sb < nsb; sb++) {
|
|
fc_fu_sb_acc2(qtype, w_row0, w_row1, xb, lane, &pg0, &pg1);
|
|
if (kind == 2)
|
|
fc_fu_sb_acc2(qtype, u_row0, u_row1, xb, lane,
|
|
&pu0, &pu1);
|
|
w_row0 += bb;
|
|
w_row1 += bb;
|
|
if (kind == 2) {
|
|
u_row0 += bb;
|
|
u_row1 += bb;
|
|
}
|
|
xb += 256;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
pg0 = fc_fu_warp_sum(pg0);
|
|
pg1 = fc_fu_warp_sum(pg1);
|
|
if (kind == 2) {
|
|
pu0 = fc_fu_warp_sum(pu0);
|
|
pu1 = fc_fu_warp_sum(pu1);
|
|
}
|
|
if (active0 && lane == 0) {
|
|
/* kind 2: silu(gate)*up — identical to fc_moe_silu_mul / gateup. */
|
|
float v = (kind == 2) ? (pg0 / (1.0f + expf(-pg0))) * pu0 : pg0;
|
|
fc_fu_store(sg, row0, v);
|
|
}
|
|
if (active1 && lane == 0) {
|
|
float v = (kind == 2) ? (pg1 / (1.0f + expf(-pg1))) * pu1 : pg1;
|
|
fc_fu_store(sg, row0 + 1, v);
|
|
}
|
|
}
|
|
|
|
/*
|
|
* Fused per-head QK-RMSNorm + partial NEOX RoPE for Q and K (was 4 launches).
|
|
* Blocks [0, n_head) norm+rope Q heads with q_norm_w; blocks [n_head,
|
|
* n_head + n_head_kv) norm+rope K heads with k_norm_w. The norm body is
|
|
* fc_qk_rms_norm verbatim (threads = blockDim.x, tree over blockDim.x); the
|
|
* rope body is fc_rope_neox_partial verbatim (thread j < half rotates the
|
|
* (j, j+half) pair). One barrier between them orders the in-place update.
|
|
*
|
|
* Launch: grid = n_head + n_head_kv, block = min(256, head_dim), shared =
|
|
* blockDim.x * sizeof(float).
|
|
*/
|
|
__global__ void fc_attn_qknorm_rope(float *q, float *k, const float *q_norm_w,
|
|
const float *k_norm_w, int n_head,
|
|
int n_head_kv, int head_dim, float eps,
|
|
int rope_dim, int pos, float theta_base) {
|
|
extern __shared__ float sdata[];
|
|
|
|
int b = blockIdx.x;
|
|
if (b >= n_head + n_head_kv)
|
|
return;
|
|
float *h;
|
|
const float *w;
|
|
if (b < n_head) {
|
|
h = q + (size_t)b * head_dim;
|
|
w = q_norm_w;
|
|
} else {
|
|
h = k + (size_t)(b - n_head) * head_dim;
|
|
w = k_norm_w;
|
|
}
|
|
|
|
int tid = threadIdx.x;
|
|
int threads = blockDim.x;
|
|
|
|
/* --- fc_qk_rms_norm body (verbatim) --- */
|
|
float sum_sq = 0.0f;
|
|
for (int i = tid; i < head_dim; i += threads) {
|
|
float v = h[i];
|
|
sum_sq += v * v;
|
|
}
|
|
sdata[tid] = sum_sq;
|
|
__syncthreads();
|
|
for (int s = threads / 2; s > 0; s >>= 1) {
|
|
if (tid < s)
|
|
sdata[tid] += sdata[tid + s];
|
|
__syncthreads();
|
|
}
|
|
float rms = sqrtf(sdata[0] / (float)head_dim + eps);
|
|
float inv_rms = 1.0f / rms;
|
|
for (int i = tid; i < head_dim; i += threads) {
|
|
float v = h[i] * inv_rms;
|
|
if (w)
|
|
v *= w[i];
|
|
h[i] = v;
|
|
}
|
|
__syncthreads(); /* rope reads what other threads normed */
|
|
|
|
/* --- fc_rope_neox_partial body (verbatim) --- */
|
|
int rot = (rope_dim > 0 && rope_dim < head_dim) ? rope_dim : head_dim;
|
|
int half = rot / 2;
|
|
int j = tid;
|
|
if (j < half) {
|
|
float freq = 1.0f / powf(theta_base, (float)(2 * j) / (float)rot);
|
|
float angle = (float)pos * freq;
|
|
float cos_a = cosf(angle);
|
|
float sin_a = sinf(angle);
|
|
float x0 = h[j];
|
|
float x1 = h[j + half];
|
|
/* Pinned to the ORIGINAL rope kernel's SASS (ptxas contracts its
|
|
* PTX mul/mul/sub): h[j] = fma(x1, -sin, x0*cos),
|
|
* h[j+half] = fma(x1, cos, x0*sin) — one pre-rounded mul + one fused
|
|
* fma each, verified against nvdisasm of rope.ptx->cubin. */
|
|
h[j] = __fmaf_rn(x1, -sin_a, __fmul_rn(x0, cos_a));
|
|
h[j + half] = __fmaf_rn(x1, cos_a, __fmul_rn(x0, sin_a));
|
|
}
|
|
}
|
|
|
|
} /* extern "C" */
|