/* 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);
}