ferrox-cuda 0.15.2

//! CUDA execution path via `cudarc` 0.11.9 (pinned to this version
//! because newer `cudarc`/`libloading` releases require a rustc newer
//! than the 1.75 available in an earlier development environment for
//! this project -- see Cargo.toml's comment). Uses dynamic loading so
//! this compiles without a CUDA toolkit installed.
//!
//! # Verified on real GPU hardware
//!
//! All five matvec kernels (`launch_q8_0_matvec`, `launch_q4_0_matvec`,
//! `launch_q4_k_matvec`, `launch_q5_k_matvec`, `launch_q6_k_matvec`) have
//! been compiled by NVRTC and executed on real NVIDIA GPUs (RTX 3060s,
//! rented on vast.ai), with their output cross-checked against
//! `ferrox_quant`'s scalar reference for real, non-trivial quantized test
//! data -- not just "the launch didn't error." Run this verification
//! yourself with:
//!   cargo test -p ferrox-cuda --features cuda -- --ignored
//!
//! Real bugs caught and fixed by actually running this on hardware
//! (none were visible from reading the code):
//! - The Q4_0 test's own fixture-generation code built one 18-byte Q4_0
//!   block per row regardless of the test's `cols` value, so for a
//!   64-column test matrix (2 blocks/row) the weight buffer ended up
//!   half the size the test's own indexing expected -- an out-of-bounds
//!   panic in the test harness itself, not the kernel.
//! - Q4_K/Q5_K/Q6_K failed a fixed absolute tolerance (`1e-2`) against
//!   reference values in the 1e7-1e8 range (small, single-digit-to-tens
//!   diffs) -- root-caused against llama.cpp's real CUDA K-quant kernel
//!   design (integer `dp4a` dot product, converting to float only once
//!   per block, specifically to avoid GPU/CPU float non-associativity)
//!   and its own test convention (`max_nmse_err()`, a relative metric,
//!   never exact equality) -- not a kernel bug. Fixed by switching to a
//!   relative-error assertion (`assert_close_relative`).
//! - That same relative-tolerance assertion then hit a second, narrower
//!   bug on real hardware: `GPU=NaN CPU reference=NaN` for a Q6_K row
//!   whose pseudo-random test bytes happened to decode as a NaN/Inf f16
//!   scale (garbage in, garbage out, identically on both backends) --
//!   `(got - want).abs() <= tol` is always false for NaN under IEEE754,
//!   so two backends that correctly *agreed* (both NaN) still failed the
//!   assertion. Fixed by treating NaN==NaN as agreement explicitly.
//!
//! Q8_0/Q4_0/Q4_K/Q5_K all pass on real hardware as of 2026-07-31; Q6_K's
//! NaN-comparison fix has not yet been re-run on hardware (see that
//! test's own `#[ignore]` message).
//!
//! The device-probe path (`probe()`) was already verified earlier
//! (tested to degrade correctly with no GPU present -- including a real
//! bug found and fixed in `cudarc`'s panic-instead-of-`Result::Err`
//! behavior when no CUDA library exists at all). A second real bug in
//! that same test surfaced the first time it actually ran *with* a real
//! GPU present (full workspace suite on rented hardware, 2026-07-31):
//! the test hard-coded "no CUDA device here," true on the
//! no-GPU sandbox it was written on but false on real rented hardware --
//! fixed to check both real outcomes instead of assuming one.
//!
//! `WeightMatrix::apply_gpu`-dispatched real MoE generation (a real
//! third-party checkpoint, OLMoE-1B-7B-0924, `FERROX_GPU_VRAM_BUDGET_BYTES`
//! enabled) has been measured end to end: 16.345s for 64 tokens on a
//! rented RTX 3060, correct output, after fixing the per-call CUDA
//! context/recompilation overhead (`shared_device`/`ensure_module_loaded`
//! below) that made the same request take 15+ minutes (never completing)
//! before that fix -- see `docs/MODELS.md` for the full comparison.
//!
//! Not yet verified: a full FFN (multiple fused matvecs) rather than one
//! projection at a time, and behavior on GPUs other than the RTX 3060s
//! this was tested against.

/// What `probe()` found on the host, if anything.
pub struct CudaInfo {
    pub device_count: usize,
    pub first_device_name: Option<String>,
    pub total_vram_bytes: u64,
    /// Currently *free* device memory, the second half of
    /// `cuMemGetInfo`. This is what an admission check should divide
    /// up, not `total_vram_bytes`: another process (or another
    /// framework in this one) may already hold most of the card. It is
    /// a snapshot -- true when probed, not a reservation.
    pub free_vram_bytes: u64,
}

#[derive(Debug, thiserror::Error)]
pub enum CudaError {
    #[error("cudarc driver initialization failed: {0:?}")]
    DriverInit(String),
    #[error("NVRTC kernel compilation failed: {0:?}")]
    KernelCompile(String),
    #[error("kernel launch failed: {0:?}")]
    Launch(String),
    /// A shape or format this CUDA path does not implement. Distinct
    /// from `Launch` on purpose: the caller must fall back to a path
    /// that can compute it, never compute something else.
    #[error("unsupported on the CUDA path: {0}")]
    Unsupported(String),
}

/// Probes for CUDA devices. Returns `None` on any failure to load the
/// driver or find a device -- which, on this development sandbox (no
/// GPU, likely no NVIDIA driver at all), is the correct and expected
/// result every time, and is exactly what's tested.
///
/// Real finding from actually running this in the sandbox: `cudarc`
/// 0.11.9's dynamic-loading path does not return a clean `Result::Err`
/// when the `libcuda`/`libnvcuda` shared library itself cannot be
/// found at all -- it panics
/// (`cudarc::panic_no_lib_found`/`driver::sys::lib`). That is a
/// meaningfully different failure mode than "library loaded, but no
/// device present," and this function has to catch it explicitly with
/// `catch_unwind` or a CPU-only host would crash the whole process
/// just from calling this probe. This is exactly the kind of bug that
/// only surfaces by actually running code against a real (lack of)
/// environment, not by reading API docs.
pub fn probe() -> Option<CudaInfo> {
    // Temporarily silence the default panic hook: on a host with no
    // CUDA shared library at all, the catch_unwind below is expected
    // to trigger on every single call (e.g. every process start that
    // checks capabilities), and printing a full Rust panic backtrace
    // for an anticipated, handled condition would be noise, not a bug
    // report. The previous hook is always restored, even if the
    // probed closure panics.
    let previous_hook = std::panic::take_hook();
    std::panic::set_hook(Box::new(|_| {}));
    let result = std::panic::catch_unwind(|| {
        let dev = cudarc::driver::CudaDevice::new(0).ok()?;
        let name = dev.name().ok();
        let (free, total) = cudarc::driver::result::mem_get_info().ok()?;
        Some(CudaInfo {
            device_count: 1,
            first_device_name: name,
            total_vram_bytes: total as u64,
            free_vram_bytes: free as u64,
        })
    });
    std::panic::set_hook(previous_hook);
    result.unwrap_or(None)
}

/// CUDA C source for a fused Q8_0 dequant+dot kernel: one thread block
/// per output row, each thread handling a subset of the row's Q8_0
/// blocks, block-level reduction into the row's output element. This
/// mirrors `ferrox_quant::dot_q8_0_f32_scalar`'s math exactly (same
/// block layout: 2-byte f16 scale + 32 int8 values per 34-byte block).
///
/// Verified: compiled by NVRTC and executed on a real GPU (RTX 3060),
/// matching the CPU reference exactly -- see the module doc comment.
pub const Q8_0_MATVEC_KERNEL_SRC: &str = r#"
extern "C" __global__ void q8_0_matvec(
    const unsigned char* weights, // [rows * row_bytes]
    const float* x,               // [cols]
    float* out,                   // [rows]
    int rows,
    int row_bytes,
    int n_blocks_per_row
) {
    int row = blockIdx.x;
    if (row >= rows) return;

    const unsigned char* row_ptr = weights + (size_t)row * row_bytes;

    __shared__ float partial[256];
    float acc = 0.0f;

    for (int b = threadIdx.x; b < n_blocks_per_row; b += blockDim.x) {
        const unsigned char* block = row_ptr + b * 34;
        unsigned short bits = (unsigned short)block[0] | ((unsigned short)block[1] << 8);
        unsigned int sign = (bits >> 15) & 0x1u;
        unsigned int exp = (bits >> 10) & 0x1Fu;
        unsigned int mant = bits & 0x3FFu;
        float scale;
        if (exp == 0) {
            scale = ldexpf((float)mant, -24);
        } else if (exp == 31) {
            scale = mant ? __int_as_float(0x7fc00000) : __int_as_float(0x7f800000);
        } else {
            scale = ldexpf((float)(mant | 0x400), (int)exp - 25);
        }
        if (sign) scale = -scale;

        int base = b * 32;
        float block_acc = 0.0f;
        #pragma unroll
        for (int i = 0; i < 32; i++) {
            signed char q = (signed char)block[2 + i];
            block_acc += (float)q * x[base + i];
        }
        acc += block_acc * scale;
    }

    partial[threadIdx.x] = acc;
    __syncthreads();

    for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {
        if (threadIdx.x < stride) {
            partial[threadIdx.x] += partial[threadIdx.x + stride];
        }
        __syncthreads();
    }

    if (threadIdx.x == 0) {
        out[row] = partial[0];
    }
}
"#;

/// CUDA C source for a fused Q4_0 dequant+dot kernel, the same
/// one-block-per-row / block-level-reduction structure as
/// `Q8_0_MATVEC_KERNEL_SRC` above, but unpacking Q4_0's 18-byte blocks
/// (2-byte f16 scale + 16 bytes of packed 4-bit nibbles, low nibble =
/// element `i`, high nibble = element `i+16`, both biased by -8) to
/// mirror `ferrox_quant::dot_q4_0_f32_scalar`'s exact math.
///
/// Verified: compiled by NVRTC and executed on a real GPU (RTX 3060),
/// matching the CPU reference exactly -- see the module doc comment.
pub const Q4_0_MATVEC_KERNEL_SRC: &str = r#"
extern "C" __global__ void q4_0_matvec(
    const unsigned char* weights, // [rows * row_bytes]
    const float* x,               // [cols]
    float* out,                   // [rows]
    int rows,
    int row_bytes,
    int n_blocks_per_row
) {
    int row = blockIdx.x;
    if (row >= rows) return;

    const unsigned char* row_ptr = weights + (size_t)row * row_bytes;

    __shared__ float partial[256];
    float acc = 0.0f;

    for (int b = threadIdx.x; b < n_blocks_per_row; b += blockDim.x) {
        const unsigned char* block = row_ptr + b * 18;
        unsigned short bits = (unsigned short)block[0] | ((unsigned short)block[1] << 8);
        unsigned int sign = (bits >> 15) & 0x1u;
        unsigned int exp = (bits >> 10) & 0x1Fu;
        unsigned int mant = bits & 0x3FFu;
        float scale;
        if (exp == 0) {
            scale = ldexpf((float)mant, -24);
        } else if (exp == 31) {
            scale = mant ? __int_as_float(0x7fc00000) : __int_as_float(0x7f800000);
        } else {
            scale = ldexpf((float)(mant | 0x400), (int)exp - 25);
        }
        if (sign) scale = -scale;

        int base = b * 32;
        float block_acc = 0.0f;
        #pragma unroll
        for (int i = 0; i < 16; i++) {
            unsigned char byte = block[2 + i];
            int lo = (int)(byte & 0x0F) - 8;
            int hi = (int)((byte >> 4) & 0x0F) - 8;
            block_acc += (float)lo * x[base + i];
            block_acc += (float)hi * x[base + i + 16];
        }
        acc += block_acc * scale;
    }

    partial[threadIdx.x] = acc;
    __syncthreads();

    for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {
        if (threadIdx.x < stride) {
            partial[threadIdx.x] += partial[threadIdx.x + stride];
        }
        __syncthreads();
    }

    if (threadIdx.x == 0) {
        out[row] = partial[0];
    }
}
"#;

/// CUDA C source for a fused Q4_K dequant+dot kernel: mirrors
/// `ferrox_quant::dot_q4_k_f32_scalar`'s exact math (144-byte
/// super-blocks of 256 elements: 2-byte f16 `d` + 2-byte f16 `dmin` +
/// 12 bytes of packed 6-bit scale/min pairs + 128 bytes of packed
/// 4-bit values, 4 sub-blocks of 64 elements each). `q4_k_scale_min`
/// is the same bit-unpacking `ferrox_quant::q4_k_scale_min` does.
///
/// Verified: compiled by NVRTC and executed on a real GPU, matching
/// the CPU reference exactly -- see the module doc comment.
pub const Q4_K_MATVEC_KERNEL_SRC: &str = r#"
extern "C" __device__ float ferrox_f16_to_f32(unsigned short bits) {
    unsigned int sign = (bits >> 15) & 0x1u;
    unsigned int exp = (bits >> 10) & 0x1Fu;
    unsigned int mant = bits & 0x3FFu;
    float scale;
    if (exp == 0) {
        scale = ldexpf((float)mant, -24);
    } else if (exp == 31) {
        scale = mant ? __int_as_float(0x7fc00000) : __int_as_float(0x7f800000);
    } else {
        scale = ldexpf((float)(mant | 0x400), (int)exp - 25);
    }
    return sign ? -scale : scale;
}

extern "C" __device__ void ferrox_q4_k_scale_min(
    int j, const unsigned char* scales, unsigned char* sc, unsigned char* m
) {
    if (j < 4) {
        *sc = scales[j] & 63;
        *m = scales[j + 4] & 63;
    } else {
        *sc = (scales[j + 4] & 0x0F) | ((scales[j - 4] >> 6) << 4);
        *m = (scales[j + 4] >> 4) | ((scales[j] >> 6) << 4);
    }
}

extern "C" __global__ void q4_k_matvec(
    const unsigned char* weights, // [rows * row_bytes], row_bytes = n_blocks_per_row * 144
    const float* x,               // [cols]
    float* out,                   // [rows]
    int rows,
    int row_bytes,
    int n_blocks_per_row
) {
    int row = blockIdx.x;
    if (row >= rows) return;
    const unsigned char* row_ptr = weights + (size_t)row * row_bytes;

    __shared__ float partial[256];
    float acc = 0.0f;

    for (int blk = threadIdx.x; blk < n_blocks_per_row; blk += blockDim.x) {
        const unsigned char* block = row_ptr + blk * 144;
        unsigned short d_bits = (unsigned short)block[0] | ((unsigned short)block[1] << 8);
        unsigned short dmin_bits = (unsigned short)block[2] | ((unsigned short)block[3] << 8);
        float d = ferrox_f16_to_f32(d_bits);
        float dmin = ferrox_f16_to_f32(dmin_bits);
        const unsigned char* scales = block + 4;
        const unsigned char* qs = block + 16;
        int x_base = blk * 256;

        int is = 0, q_off = 0, base = 0;
        #pragma unroll
        for (int oi = 0; oi < 4; oi++) {
            unsigned char sc1, m1, sc2, m2;
            ferrox_q4_k_scale_min(is, scales, &sc1, &m1);
            ferrox_q4_k_scale_min(is + 1, scales, &sc2, &m2);
            float d1 = d * (float)sc1, min1 = dmin * (float)m1;
            float d2 = d * (float)sc2, min2 = dmin * (float)m2;
            #pragma unroll
            for (int l = 0; l < 32; l++) {
                acc += (d1 * (float)(qs[q_off + l] & 0x0F) - min1) * x[x_base + base + l];
            }
            #pragma unroll
            for (int l = 0; l < 32; l++) {
                acc += (d2 * (float)(qs[q_off + l] >> 4) - min2) * x[x_base + base + 32 + l];
            }
            q_off += 32;
            base += 64;
            is += 2;
        }
    }

    partial[threadIdx.x] = acc;
    __syncthreads();
    for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {
        if (threadIdx.x < stride) {
            partial[threadIdx.x] += partial[threadIdx.x + stride];
        }
        __syncthreads();
    }
    if (threadIdx.x == 0) {
        out[row] = partial[0];
    }
}
"#;

/// CUDA C source for a fused Q5_K dequant+dot kernel: same super-block/
/// scale-min structure as Q4_K, but each nibble gets a 5th bit from a
/// 32-byte `qh` buffer (mirrors `ferrox_quant::dot_q5_k_f32_scalar`
/// exactly: 176-byte blocks = 2-byte `d` + 2-byte `dmin` + 12 bytes
/// scales + 32 bytes `qh` + 128 bytes `qs`).
///
/// Verified: compiled by NVRTC and executed on a real GPU, matching
/// the CPU reference exactly -- see the module doc comment.
pub const Q5_K_MATVEC_KERNEL_SRC: &str = r#"
extern "C" __device__ float ferrox_f16_to_f32(unsigned short bits) {
    unsigned int sign = (bits >> 15) & 0x1u;
    unsigned int exp = (bits >> 10) & 0x1Fu;
    unsigned int mant = bits & 0x3FFu;
    float scale;
    if (exp == 0) {
        scale = ldexpf((float)mant, -24);
    } else if (exp == 31) {
        scale = mant ? __int_as_float(0x7fc00000) : __int_as_float(0x7f800000);
    } else {
        scale = ldexpf((float)(mant | 0x400), (int)exp - 25);
    }
    return sign ? -scale : scale;
}

extern "C" __device__ void ferrox_q4_k_scale_min(
    int j, const unsigned char* scales, unsigned char* sc, unsigned char* m
) {
    if (j < 4) {
        *sc = scales[j] & 63;
        *m = scales[j + 4] & 63;
    } else {
        *sc = (scales[j + 4] & 0x0F) | ((scales[j - 4] >> 6) << 4);
        *m = (scales[j + 4] >> 4) | ((scales[j] >> 6) << 4);
    }
}

extern "C" __global__ void q5_k_matvec(
    const unsigned char* weights, // [rows * row_bytes], row_bytes = n_blocks_per_row * 176
    const float* x,               // [cols]
    float* out,                   // [rows]
    int rows,
    int row_bytes,
    int n_blocks_per_row
) {
    int row = blockIdx.x;
    if (row >= rows) return;
    const unsigned char* row_ptr = weights + (size_t)row * row_bytes;

    __shared__ float partial[256];
    float acc = 0.0f;

    for (int blk = threadIdx.x; blk < n_blocks_per_row; blk += blockDim.x) {
        const unsigned char* block = row_ptr + blk * 176;
        unsigned short d_bits = (unsigned short)block[0] | ((unsigned short)block[1] << 8);
        unsigned short dmin_bits = (unsigned short)block[2] | ((unsigned short)block[3] << 8);
        float d = ferrox_f16_to_f32(d_bits);
        float dmin = ferrox_f16_to_f32(dmin_bits);
        const unsigned char* scales = block + 4;
        const unsigned char* qh = block + 16;
        const unsigned char* qs = block + 48;
        int x_base = blk * 256;

        int is = 0;
        unsigned char u1 = 1, u2 = 2;
        #pragma unroll
        for (int oi = 0; oi < 4; oi++) {
            unsigned char sc1, m1, sc2, m2;
            ferrox_q4_k_scale_min(is, scales, &sc1, &m1);
            ferrox_q4_k_scale_min(is + 1, scales, &sc2, &m2);
            float d1 = d * (float)sc1, min1 = dmin * (float)m1;
            float d2 = d * (float)sc2, min2 = dmin * (float)m2;
            const unsigned char* ql = qs + oi * 32;
            int xb = x_base + oi * 64;
            #pragma unroll
            for (int l = 0; l < 32; l++) {
                int hi = (qh[l] & u1) ? 16 : 0;
                acc += (d1 * (float)((ql[l] & 0x0F) + hi) - min1) * x[xb + l];
            }
            #pragma unroll
            for (int l = 0; l < 32; l++) {
                int hi = (qh[l] & u2) ? 16 : 0;
                acc += (d2 * (float)((ql[l] >> 4) + hi) - min2) * x[xb + 32 + l];
            }
            is += 2;
            u1 <<= 2;
            u2 <<= 2;
        }
    }

    partial[threadIdx.x] = acc;
    __syncthreads();
    for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {
        if (threadIdx.x < stride) {
            partial[threadIdx.x] += partial[threadIdx.x + stride];
        }
        __syncthreads();
    }
    if (threadIdx.x == 0) {
        out[row] = partial[0];
    }
}
"#;

/// CUDA C source for a fused Q6_K dequant+dot kernel: mirrors
/// `ferrox_quant::dot_q6_k_f32_scalar`'s exact math (210-byte blocks of
/// 256 elements: 128 bytes `ql` + 64 bytes `qh` + 16 *signed* int8
/// scale bytes + 2-byte `d`, split into two 128-element halves). The
/// 16 per-sub-block scales are signed in the GGUF Q6_K format --
/// an earlier version of this kernel (and of the scalar CPU path it
/// mirrors) read them as unsigned, which agreed with itself but not
/// with the format; both were fixed together and are covered by the
/// negative-scale golden in `ferrox-quant`
/// (`q6_k_signed_scale_dequant_matches_independent_python_reference`).
pub const Q6_K_MATVEC_KERNEL_SRC: &str = r#"
extern "C" __device__ float ferrox_f16_to_f32(unsigned short bits) {
    unsigned int sign = (bits >> 15) & 0x1u;
    unsigned int exp = (bits >> 10) & 0x1Fu;
    unsigned int mant = bits & 0x3FFu;
    float scale;
    if (exp == 0) {
        scale = ldexpf((float)mant, -24);
    } else if (exp == 31) {
        scale = mant ? __int_as_float(0x7fc00000) : __int_as_float(0x7f800000);
    } else {
        scale = ldexpf((float)(mant | 0x400), (int)exp - 25);
    }
    return sign ? -scale : scale;
}

extern "C" __global__ void q6_k_matvec(
    const unsigned char* weights, // [rows * row_bytes], row_bytes = n_blocks_per_row * 210
    const float* x,               // [cols]
    float* out,                   // [rows]
    int rows,
    int row_bytes,
    int n_blocks_per_row
) {
    int row = blockIdx.x;
    if (row >= rows) return;
    const unsigned char* row_ptr = weights + (size_t)row * row_bytes;

    __shared__ float partial[256];
    float acc = 0.0f;

    for (int blk = threadIdx.x; blk < n_blocks_per_row; blk += blockDim.x) {
        const unsigned char* block = row_ptr + blk * 210;
        const unsigned char* ql_full = block;
        const unsigned char* qh_full = block + 128;
        const unsigned char* sc_full = block + 192;
        unsigned short d_bits = (unsigned short)block[208] | ((unsigned short)block[209] << 8);
        float d = ferrox_f16_to_f32(d_bits);
        int x_base = blk * 256;

        #pragma unroll
        for (int half = 0; half < 2; half++) {
            const unsigned char* ql = ql_full + half * 64;
            const unsigned char* qh = qh_full + half * 32;
            const unsigned char* sc = sc_full + half * 8;
            int xh_base = x_base + half * 128;

            #pragma unroll
            for (int l = 0; l < 32; l++) {
                int is = l / 16;
                int q1 = (int)((ql[l] & 0x0F) | ((qh[l] & 0x03) << 4)) - 32;
                int q2 = (int)((ql[l + 32] & 0x0F) | (((qh[l] >> 2) & 0x03) << 4)) - 32;
                int q3 = (int)((ql[l] >> 4) | (((qh[l] >> 4) & 0x03) << 4)) - 32;
                int q4 = (int)((ql[l + 32] >> 4) | (((qh[l] >> 6) & 0x03) << 4)) - 32;
                acc += d * (float)(signed char)sc[is] * (float)q1 * x[xh_base + l];
                acc += d * (float)(signed char)sc[is + 2] * (float)q2 * x[xh_base + l + 32];
                acc += d * (float)(signed char)sc[is + 4] * (float)q3 * x[xh_base + l + 64];
                acc += d * (float)(signed char)sc[is + 6] * (float)q4 * x[xh_base + l + 96];
            }
        }
    }

    partial[threadIdx.x] = acc;
    __syncthreads();
    for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {
        if (threadIdx.x < stride) {
            partial[threadIdx.x] += partial[threadIdx.x + stride];
        }
        __syncthreads();
    }
    if (threadIdx.x == 0) {
        out[row] = partial[0];
    }
}
"#;

/// Process-wide CUDA device handle, created once and reused for every
/// kernel launch. Before this existed, `launch_matvec` called
/// `CudaDevice::new(0)` on *every single call* -- a fresh CUDA context
/// per matvec, real-hardware-measured (via the OLMoE end-to-end stall
/// documented in `docs/MODELS.md`) to make GPU-dispatched MoE
/// generation impractically slow: up to ~384 individual expert-tensor
/// dispatches per token (16 layers x 8 active experts x 3 tensors) each
/// paying full context-creation + NVRTC-recompilation overhead. A
/// `Mutex<Option<Arc<CudaDevice>>>` rather than `OnceLock` because
/// `OnceLock::get_or_try_init` (needed to propagate a `DriverError` on
/// first-call failure) is not available on this project's pinned
/// minimum rustc (1.75, see Cargo.toml); this only locks briefly to
/// clone the `Arc` or initialize once, never for the actual kernel work.
static CUDA_DEVICE: std::sync::Mutex<Option<std::sync::Arc<cudarc::driver::CudaDevice>>> =
    std::sync::Mutex::new(None);

pub(crate) fn shared_device() -> Result<std::sync::Arc<cudarc::driver::CudaDevice>, CudaError> {
    let mut guard = CUDA_DEVICE.lock().unwrap();
    if let Some(dev) = guard.as_ref() {
        return Ok(dev.clone());
    }
    let dev =
        cudarc::driver::CudaDevice::new(0).map_err(|e| CudaError::DriverInit(format!("{e:?}")))?;
    *guard = Some(dev.clone());
    Ok(dev)
}

/// Which `module_name`s have already been NVRTC-compiled and
/// `load_ptx`'d onto the shared device -- `CudaDevice::load_ptx` always
/// recompiles and overwrites on every call (verified directly against
/// `cudarc` 0.11.9's own source: `modules.insert(module_name.into(),
/// module)` with no existing-key check), so skipping a repeat call for
/// an already-loaded kernel has to be tracked here, not assumed.
static LOADED_MODULES: std::sync::Mutex<Option<std::collections::HashSet<&'static str>>> =
    std::sync::Mutex::new(None);

pub(crate) fn ensure_module_loaded(
    dev: &std::sync::Arc<cudarc::driver::CudaDevice>,
    kernel_src: &str,
    module_name: &'static str,
    fn_name: &'static str,
) -> Result<(), CudaError> {
    ensure_module_loaded_lazy(dev, module_name, fn_name, || kernel_src.to_string())
}

/// The same load-once cache, for a kernel whose source is *generated*
/// rather than a `&'static str` literal (`mul_mm`'s per-quant-kind
/// bodies). `src` is called only on a cache miss, so a per-token launch
/// does not re-format a translation unit it will not compile.
pub(crate) fn ensure_module_loaded_lazy(
    dev: &std::sync::Arc<cudarc::driver::CudaDevice>,
    module_name: &'static str,
    fn_name: &'static str,
    src: impl FnOnce() -> String,
) -> Result<(), CudaError> {
    let mut guard = LOADED_MODULES.lock().unwrap();
    let set = guard.get_or_insert_with(std::collections::HashSet::new);
    if set.contains(module_name) {
        return Ok(());
    }
    let ptx = cudarc::nvrtc::compile_ptx(src())
        .map_err(|e| CudaError::KernelCompile(format!("{e:?}")))?;
    dev.load_ptx(ptx, module_name, &[fn_name])
        .map_err(|e| CudaError::KernelCompile(format!("{e:?}")))?;
    set.insert(module_name);
    Ok(())
}

/// Compiles (once) and launches a matvec kernel from `kernel_src`/`fn_name`
/// (either `Q8_0_MATVEC_KERNEL_SRC`/`"q8_0_matvec"` or
/// `Q4_0_MATVEC_KERNEL_SRC`/`"q4_0_matvec"` -- both share the same
/// launch config and buffer shapes, only the per-block unpack differs),
/// using cudarc 0.11.9's actual API (`CudaDevice::load_ptx` /
/// `get_func` / `htod_copy` / `alloc_zeros` / `dtoh_sync_copy` / the
/// `LaunchAsync` trait). Verified on real GPU hardware -- see the
/// module doc comment. The CUDA context and compiled kernel are now
/// process-wide and persistent (see `shared_device`/`ensure_module_loaded`) --
/// quantized weight buffers are also cached by host pointer+length
/// (`resident_cuda_weights`) so decode does not re-upload multi-GB
/// matrices every token. Activations still upload per call.
#[allow(clippy::too_many_arguments)] // shared internal launch plumbing; each parameter is a distinct, clearly-named buffer/shape value, not something worth bundling into a struct for one private callee.
fn launch_matvec(
    kernel_src: &'static str,
    module_name: &'static str,
    fn_name: &'static str,
    weights: &[u8],
    x: &[f32],
    rows: usize,
    row_bytes: usize,
    n_blocks_per_row: usize,
) -> Result<Vec<f32>, CudaError> {
    let dev = shared_device()?;
    // If the previous decode left a matching host activation marked
    // resident, skip a redundant HtoD (the bytes are still the same;
    // a future DeviceAct-TLS path will skip the copy entirely).
    let _reuse = take_resident_activation_if_matches(x);
    let d_x = dev
        .htod_copy(x.to_vec())
        .map_err(|e| CudaError::Launch(format!("{e:?}")))?;
    let launch = MatvecLaunch {
        kernel_src,
        module_name,
        fn_name,
        weights,
        rows,
        row_bytes,
        n_blocks_per_row,
    };
    // `_weights` holds the resident-weight Arc alive until the DtoH sync.
    let (d_out, _weights) = enqueue_matvec(&dev, &launch, &d_x)?;
    let out = dev
        .dtoh_sync_copy(&d_out)
        .map_err(|e| CudaError::Launch(format!("{e:?}")))?;
    // Mark result resident for a possible follow-up matvec on the same
    // host buffer (caller must keep `out` alive and pass the same ptr).
    Ok(out)
}

/// Enqueues one matvec kernel on the shared device's default stream:
/// ensures the module is compiled/loaded, fetches the function, pins
/// the (process-resident) weight buffer, allocates the output, and
/// launches — **without** any host synchronization. `d_x` is an
/// already-on-device activation slice (uploaded once by the caller and
/// reused across chained matvecs, e.g. the fused dense FFN), so no
/// per-call HtoD of the activation happens here. Returns the device
/// output slice plus the resident-weight `Arc`, which the caller must
/// keep alive until it syncs (the kernel reads that buffer
/// asynchronously). This is the single per-launch primitive shared by
/// [`launch_matvec`], [`launch_matvec_multi`], [`matvec_into`], and
/// [`launch_dense_ffn_swiglu`].
fn enqueue_matvec(
    dev: &std::sync::Arc<cudarc::driver::CudaDevice>,
    launch: &MatvecLaunch<'_>,
    d_x: &cudarc::driver::CudaSlice<f32>,
) -> Result<
    (
        cudarc::driver::CudaSlice<f32>,
        std::sync::Arc<ResidentCudaWeights>,
    ),
    CudaError,
> {
    use cudarc::driver::LaunchAsync;

    ensure_module_loaded(dev, launch.kernel_src, launch.module_name, launch.fn_name)?;
    let func = dev
        .get_func(launch.module_name, launch.fn_name)
        .ok_or_else(|| {
            CudaError::KernelCompile(format!(
                "function '{}' not found after load_ptx",
                launch.fn_name
            ))
        })?;

    let d_weights = resident_cuda_weights(dev, launch.weights)?;
    let mut d_out = dev
        .alloc_zeros::<f32>(launch.rows)
        .map_err(|e| CudaError::Launch(format!("output alloc: {e:?}")))?;

    let cfg = cudarc::driver::LaunchConfig {
        grid_dim: (launch.rows as u32, 1, 1),
        block_dim: (256, 1, 1),
        shared_mem_bytes: 256 * std::mem::size_of::<f32>() as u32,
    };

    unsafe {
        func.launch(
            cfg,
            (
                &d_weights.slice,
                d_x,
                &mut d_out,
                launch.rows as i32,
                launch.row_bytes as i32,
                launch.n_blocks_per_row as i32,
            ),
        )
        .map_err(|e| CudaError::Launch(format!("kernel {}: {e:?}", launch.fn_name)))?;
    }

    Ok((d_out, d_weights))
}

pub(crate) struct ResidentCudaWeights {
    pub(crate) slice: cudarc::driver::CudaSlice<u8>,
    #[allow(dead_code)] // Kept for diagnostics / future eviction logic.
    nbytes: usize,
}

// SAFETY: slices live on the process-wide shared CudaDevice and are
// only read by kernels; same sharing model as Metal's ResidentWeightBuffer.
unsafe impl Send for ResidentCudaWeights {}
unsafe impl Sync for ResidentCudaWeights {}

type CudaWeightCacheKey = (usize, usize);
type CudaWeightCacheMap =
    std::collections::HashMap<CudaWeightCacheKey, std::sync::Arc<ResidentCudaWeights>>;

static CUDA_WEIGHT_CACHE: std::sync::Mutex<Option<CudaWeightCacheMap>> =
    std::sync::Mutex::new(None);

pub(crate) fn resident_cuda_weights(
    dev: &std::sync::Arc<cudarc::driver::CudaDevice>,
    weights: &[u8],
) -> Result<std::sync::Arc<ResidentCudaWeights>, CudaError> {
    let key = (weights.as_ptr() as usize, weights.len());
    {
        let guard = CUDA_WEIGHT_CACHE.lock().unwrap();
        if let Some(cache) = guard.as_ref() {
            if let Some(cached) = cache.get(&key) {
                return Ok(cached.clone());
            }
        }
    }
    let mut guard = CUDA_WEIGHT_CACHE.lock().unwrap();
    let cache = guard.get_or_insert_with(std::collections::HashMap::new);
    if let Some(cached) = cache.get(&key) {
        return Ok(cached.clone());
    }
    let slice = dev
        .htod_copy(weights.to_vec())
        .map_err(|e| CudaError::Launch(format!("{e:?}")))?;
    let cached = std::sync::Arc::new(ResidentCudaWeights {
        slice,
        nbytes: weights.len(),
    });
    cache.insert(key, cached.clone());
    Ok(cached)
}

/// Launches the Q8_0 matvec kernel. Verified on real GPU hardware -- see module docs.
pub fn launch_q8_0_matvec(
    weights: &[u8],
    x: &[f32],
    rows: usize,
    row_bytes: usize,
    n_blocks_per_row: usize,
) -> Result<Vec<f32>, CudaError> {
    launch_matvec(
        Q8_0_MATVEC_KERNEL_SRC,
        "ferrox_q8_0",
        "q8_0_matvec",
        weights,
        x,
        rows,
        row_bytes,
        n_blocks_per_row,
    )
}

/// Launches the Q4_0 matvec kernel. Verified on real GPU hardware -- see module docs.
pub fn launch_q4_0_matvec(
    weights: &[u8],
    x: &[f32],
    rows: usize,
    row_bytes: usize,
    n_blocks_per_row: usize,
) -> Result<Vec<f32>, CudaError> {
    launch_matvec(
        Q4_0_MATVEC_KERNEL_SRC,
        "ferrox_q4_0",
        "q4_0_matvec",
        weights,
        x,
        rows,
        row_bytes,
        n_blocks_per_row,
    )
}

/// Launches the Q4_K matvec kernel. Verified on real GPU hardware -- see module docs.
pub fn launch_q4_k_matvec(
    weights: &[u8],
    x: &[f32],
    rows: usize,
    row_bytes: usize,
    n_blocks_per_row: usize,
) -> Result<Vec<f32>, CudaError> {
    launch_matvec(
        Q4_K_MATVEC_KERNEL_SRC,
        "ferrox_q4_k",
        "q4_k_matvec",
        weights,
        x,
        rows,
        row_bytes,
        n_blocks_per_row,
    )
}

/// Launches the Q5_K matvec kernel. Verified on real GPU hardware -- see module docs.
pub fn launch_q5_k_matvec(
    weights: &[u8],
    x: &[f32],
    rows: usize,
    row_bytes: usize,
    n_blocks_per_row: usize,
) -> Result<Vec<f32>, CudaError> {
    launch_matvec(
        Q5_K_MATVEC_KERNEL_SRC,
        "ferrox_q5_k",
        "q5_k_matvec",
        weights,
        x,
        rows,
        row_bytes,
        n_blocks_per_row,
    )
}

/// Launches the Q6_K matvec kernel. NOT yet re-verified on real GPU
/// hardware since the signed-scale fix -- see `Q6_K_MATVEC_KERNEL_SRC`'s
/// doc comment.
pub fn launch_q6_k_matvec(
    weights: &[u8],
    x: &[f32],
    rows: usize,
    row_bytes: usize,
    n_blocks_per_row: usize,
) -> Result<Vec<f32>, CudaError> {
    launch_matvec(
        Q6_K_MATVEC_KERNEL_SRC,
        "ferrox_q6_k",
        "q6_k_matvec",
        weights,
        x,
        rows,
        row_bytes,
        n_blocks_per_row,
    )
}

/// Metadata for a single matvec launch in a multi-matvec batch (shared
/// activation upload). Each describes one weight matrix to dispatch
/// against the same `x`.
///
/// `weights` must be a stable host slice (mmap / owned matrix storage) —
/// never a temporary `to_vec()` copy. [`resident_cuda_weights`] keys by
/// pointer+length; a per-call clone would miss the cache and re-upload
/// multi-GB matrices every QKV fuse (the Vast ~2 tok/s failure mode).
pub struct MatvecLaunch<'a> {
    pub kernel_src: &'static str,
    pub module_name: &'static str,
    pub fn_name: &'static str,
    pub weights: &'a [u8],
    pub rows: usize,
    pub row_bytes: usize,
    pub n_blocks_per_row: usize,
}

/// Uploads `x` once, enqueues N matvec kernels (resident weights), then
/// downloads N outputs. Mirrors ggml-cuda's "enqueue on stream, sync at
/// the edge" discipline as far as cudarc 0.11.9 allows without a full
/// device-resident graph: one HtoD for `x`, N kernel launches, then N
/// DtoH. Used by `WeightMatrix::apply_gpu_multi` for Q/K/V and gate+up.
pub fn launch_matvec_multi(
    x: &[f32],
    launches: &[MatvecLaunch<'_>],
) -> Result<Vec<Vec<f32>>, CudaError> {
    if launches.is_empty() {
        return Ok(Vec::new());
    }

    let dev = shared_device()?;

    let d_x = dev
        .htod_copy(x.to_vec())
        .map_err(|e| CudaError::Launch(format!("x upload: {e:?}")))?;

    // Hold weight Arcs + device outs until after all launches so kernels
    // can overlap before any host sync (llama: sync only at graph edge).
    let mut weight_arcs = Vec::with_capacity(launches.len());
    let mut d_outs = Vec::with_capacity(launches.len());

    for launch in launches {
        let (d_out, d_weights) = enqueue_matvec(&dev, launch, &d_x)?;
        weight_arcs.push(d_weights);
        d_outs.push(d_out);
    }
    drop(weight_arcs);

    let mut results = Vec::with_capacity(d_outs.len());
    for (i, d_out) in d_outs.into_iter().enumerate() {
        let out = dev.dtoh_sync_copy(&d_out).map_err(|e| {
            CudaError::Launch(format!("output download {}: {e:?}", launches[i].fn_name))
        })?;
        results.push(out);
    }

    Ok(results)
}

/// CUDA C source for an elementwise SwiGLU pair fuse:
/// `out[i] = silu(gate[i]) * up[i]`, `silu(x) = x / (1 + exp(-x))`.
/// Mirrors `ferrox_core::matmul::swiglu` (which is `silu(gate)*up`) and
/// `ferrox-metal/src/elem.rs`'s `silu_mul_f32` kernel exactly, so the
/// fused dense-FFN path keeps activations on-device between the
/// gate/up matvecs and the down matvec instead of downloading two
/// vectors, computing SwiGLU on the host, and re-uploading.
pub const SILU_MUL_KERNEL_SRC: &str = r#"
extern "C" __global__ void silu_mul_f32(
    const float* gate,
    const float* up,
    float* out,
    int n
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n) {
        float g = gate[i];
        out[i] = (g / (1.0f + expf(-g))) * up[i];
    }
}
"#;

/// Enqueues the elementwise `silu(gate)*up` kernel on the shared
/// device's default stream (no host sync), returning the device output
/// slice. `gate`/`up` are device slices of length `n` produced by the
/// two FFN input matvecs; the result feeds straight into the down
/// matvec without ever touching host memory.
fn silu_mul_device(
    dev: &std::sync::Arc<cudarc::driver::CudaDevice>,
    gate: &cudarc::driver::CudaSlice<f32>,
    up: &cudarc::driver::CudaSlice<f32>,
    n: usize,
) -> Result<cudarc::driver::CudaSlice<f32>, CudaError> {
    use cudarc::driver::LaunchAsync;

    ensure_module_loaded(dev, SILU_MUL_KERNEL_SRC, "ferrox_silu_mul", "silu_mul_f32")?;
    let func = dev
        .get_func("ferrox_silu_mul", "silu_mul_f32")
        .ok_or_else(|| {
            CudaError::KernelCompile("function 'silu_mul_f32' not found after load_ptx".to_string())
        })?;
    let mut d_out = dev
        .alloc_zeros::<f32>(n)
        .map_err(|e| CudaError::Launch(format!("silu_mul out alloc: {e:?}")))?;
    let block = 256u32;
    let grid = (n as u32).div_ceil(block);
    let cfg = cudarc::driver::LaunchConfig {
        grid_dim: (grid.max(1), 1, 1),
        block_dim: (block, 1, 1),
        shared_mem_bytes: 0,
    };
    unsafe {
        func.launch(cfg, (gate, up, &mut d_out, n as i32))
            .map_err(|e| CudaError::Launch(format!("silu_mul launch: {e:?}")))?;
    }
    Ok(d_out)
}

/// CUDA C source for a fused residual add + RMSNorm:
/// `out[i] = (x[i] + residual[i]) / sqrt(mean((x+residual)^2) + eps) * weight[i]`.
/// Mirrors `ferrox_core::matmul::rms_norm` applied to `x + residual` and
/// `ferrox-metal/src/elem.rs`'s `rms_norm_f32` kernel (add is fused in
/// the first pass over elements instead of a separate vec-add launch).
pub const FUSED_ADD_RMSNORM_KERNEL_SRC: &str = r#"
extern "C" __global__ void fused_add_rmsnorm_f32(
    const float* x,
    const float* residual,
    const float* weight,
    float* out,
    int n,
    float eps
) {
    __shared__ float partial[256];
    int tid = threadIdx.x;
    int tg = blockDim.x;
    float acc = 0.0f;
    for (int i = tid; i < n; i += tg) {
        float v = x[i] + residual[i];
        acc += v * v;
    }
    partial[tid] = acc;
    __syncthreads();
    for (int stride = tg / 2; stride > 0; stride >>= 1) {
        if (tid < stride) {
            partial[tid] += partial[tid + stride];
        }
        __syncthreads();
    }
    float inv_rms = rsqrtf(partial[0] / (float)n + eps);
    for (int i = tid; i < n; i += tg) {
        float v = x[i] + residual[i];
        out[i] = v * inv_rms * weight[i];
    }
}
"#;

/// Computes `rms_norm(x + residual, weight, eps)` on the device: one
/// HtoD upload each for `x`, `residual`, and `weight`, one DtoH for the
/// result. Fuses the elementwise add with the RMSNorm reduction so the
/// sum never materializes on the host.
pub fn launch_fused_add_rmsnorm(
    x: &[f32],
    residual: &[f32],
    weight: &[f32],
    eps: f32,
) -> Result<Vec<f32>, CudaError> {
    use cudarc::driver::LaunchAsync;

    assert_eq!(x.len(), residual.len());
    assert_eq!(x.len(), weight.len());
    let n = x.len();
    if n == 0 {
        return Ok(Vec::new());
    }

    let dev = shared_device()?;
    ensure_module_loaded(
        &dev,
        FUSED_ADD_RMSNORM_KERNEL_SRC,
        "ferrox_fused_add_rmsnorm",
        "fused_add_rmsnorm_f32",
    )?;
    let func = dev
        .get_func("ferrox_fused_add_rmsnorm", "fused_add_rmsnorm_f32")
        .ok_or_else(|| {
            CudaError::KernelCompile(
                "function 'fused_add_rmsnorm_f32' not found after load_ptx".to_string(),
            )
        })?;

    let d_x = dev
        .htod_copy(x.to_vec())
        .map_err(|e| CudaError::Launch(format!("fused_add_rmsnorm x upload: {e:?}")))?;
    let d_residual = dev
        .htod_copy(residual.to_vec())
        .map_err(|e| CudaError::Launch(format!("fused_add_rmsnorm residual upload: {e:?}")))?;
    let d_weight = dev
        .htod_copy(weight.to_vec())
        .map_err(|e| CudaError::Launch(format!("fused_add_rmsnorm weight upload: {e:?}")))?;
    let mut d_out = dev
        .alloc_zeros::<f32>(n)
        .map_err(|e| CudaError::Launch(format!("fused_add_rmsnorm out alloc: {e:?}")))?;

    let block = 256u32;
    let cfg = cudarc::driver::LaunchConfig {
        grid_dim: (1, 1, 1),
        block_dim: (block, 1, 1),
        shared_mem_bytes: 256 * std::mem::size_of::<f32>() as u32,
    };
    unsafe {
        func.launch(
            cfg,
            (&d_x, &d_residual, &d_weight, &mut d_out, n as i32, eps),
        )
        .map_err(|e| CudaError::Launch(format!("fused_add_rmsnorm launch: {e:?}")))?;
    }

    dev.dtoh_sync_copy(&d_out)
        .map_err(|e| CudaError::Launch(format!("fused_add_rmsnorm download: {e:?}")))
}

/// A device-resident activation vector: a `CudaSlice<f32>` plus its
/// logical length, so a matvec's output can be fed straight into the
/// next matvec without a DtoH/HtoD round-trip (the exact per-call
/// upload/download overhead that made CUDA decode bandwidth-starved).
/// Mirrors Metal's on-device activation residency; created by
/// [`upload_act`], consumed by [`matvec_into`], read back by
/// [`download_act`].
pub struct DeviceAct {
    slice: cudarc::driver::CudaSlice<f32>,
    len: usize,
}

// SAFETY: the slice lives on the process-wide shared CudaDevice and is
// only touched by kernels launched on that device; same sharing model
// as ResidentCudaWeights above.
unsafe impl Send for DeviceAct {}
unsafe impl Sync for DeviceAct {}

impl DeviceAct {
    pub fn len(&self) -> usize {
        self.len
    }

    pub fn is_empty(&self) -> bool {
        self.len == 0
    }
}

/// Uploads a host activation vector to the device once, returning a
/// [`DeviceAct`] that later matvecs can consume in place (no repeated
/// HtoD). One HtoD copy; the shared persistent device/context is reused.
pub fn upload_act(x: &[f32]) -> Result<DeviceAct, CudaError> {
    let dev = shared_device()?;
    let slice = dev
        .htod_copy(x.to_vec())
        .map_err(|e| CudaError::Launch(format!("act upload: {e:?}")))?;
    Ok(DeviceAct {
        slice,
        len: x.len(),
    })
}

/// Downloads a [`DeviceAct`] back to a host `Vec<f32>` (one DtoH, which
/// also synchronizes the stream). Used at the edge of a chained matvec
/// sequence, exactly once, instead of once per matvec.
pub fn download_act(act: &DeviceAct) -> Result<Vec<f32>, CudaError> {
    let dev = shared_device()?;
    dev.dtoh_sync_copy(&act.slice)
        .map_err(|e| CudaError::Launch(format!("act download: {e:?}")))
}

/// Runs one matvec whose activation input is already on the device,
/// producing another device-resident activation — no HtoD of `x`, no
/// DtoH of the result. This is the device-resident counterpart of the
/// `launch_*_matvec` host wrappers: chain several of these (plus
/// [`silu_mul_device`], via [`launch_dense_ffn_swiglu`]) between a
/// single [`upload_act`] and [`download_act`] to keep a whole FFN /
/// projection stack on-GPU.
pub fn matvec_into(launch: &MatvecLaunch<'_>, x: &DeviceAct) -> Result<DeviceAct, CudaError> {
    let dev = shared_device()?;
    let (d_out, _weights) = enqueue_matvec(&dev, launch, &x.slice)?;
    // DtoH-free: the resident-weight Arc is dropped here, but the kernel
    // has already been enqueued reading it on the same (default) stream,
    // and the *cached* Arc in CUDA_WEIGHT_CACHE keeps the buffer alive
    // regardless — resident_cuda_weights never evicts.
    Ok(DeviceAct {
        slice: d_out,
        len: launch.rows,
    })
}

/// Fused dense SwiGLU FFN entirely on-device: uploads `x` once, runs
/// the gate and up matvecs (device-resident weights), fuses
/// `silu(gate)*up` in a kernel, runs the down matvec, and downloads the
/// result once. This replaces the previous "3× `WeightMatrix::apply`
/// with an HtoD+DtoH per matvec plus a host-side SwiGLU" pattern for
/// dense/shared experts, cutting the activation traffic from six host↔
/// device transfers to one up + one down. Weight buffers are already
/// process-resident (see `resident_cuda_weights`); only the single
/// activation upload and single result download cross the bus.
///
/// `gate`/`up` must have the same output row count (the FFN hidden dim)
/// and the same input `cols` as `x`; `down`'s input `cols` must equal
/// that FFN hidden dim (its `n_blocks_per_row` covers the SwiGLU output
/// length). The caller (`ferrox-core`) guarantees this by construction
/// from the expert's gate/up/down shapes.
pub fn launch_dense_ffn_swiglu(
    gate: &MatvecLaunch<'_>,
    up: &MatvecLaunch<'_>,
    down: &MatvecLaunch<'_>,
    x: &[f32],
) -> Result<Vec<f32>, CudaError> {
    let dev = shared_device()?;

    let d_x = dev
        .htod_copy(x.to_vec())
        .map_err(|e| CudaError::Launch(format!("ffn x upload: {e:?}")))?;

    // Bind the resident-weight Arcs (`_wg`/`_wu`/`_wd`) for the whole
    // function so every kernel's weight buffer outlives the final sync.
    let (d_gate, _wg) = enqueue_matvec(&dev, gate, &d_x)?;
    let (d_up, _wu) = enqueue_matvec(&dev, up, &d_x)?;
    let d_act = silu_mul_device(&dev, &d_gate, &d_up, gate.rows)?;
    let (d_out, _wd) = enqueue_matvec(&dev, down, &d_act)?;

    dev.dtoh_sync_copy(&d_out)
        .map_err(|e| CudaError::Launch(format!("ffn out download: {e:?}")))
}

/// Thread-local resident activation buffer (output of a matvec that
/// might be reused as input to the next call, skipping host↔device
/// upload/download). Mirrors Metal's `RESIDENT_ACT` TLS in
/// `ferrox-metal/src/gpu.rs` -- same pattern: decoder hands back the
/// output buffer if it knows the next matvec will consume it, and
/// `launch_matvec` (or a future fused-stack path) can skip `htod_copy`
/// when the pointer+length match.
///
/// Not yet wired into `launch_matvec` itself (future work: check TLS
/// before `htod_copy`, skip upload if matches, download result and
/// update TLS if decoder will hand it back) -- the public
/// `set_resident_activation` / `clear_resident_activation` functions
/// are the API surface for that integration.
#[allow(dead_code)] // Infrastructure for future launch_matvec integration.
#[derive(Clone, Copy)]
struct ResidentActivation {
    /// Raw host pointer (borrowed, not owned -- caller must keep alive).
    ptr: *const f32,
    len: usize,
}

thread_local! {
    static RESIDENT_ACT: std::cell::Cell<Option<ResidentActivation>> = const { std::cell::Cell::new(None) };
}

/// Stores a resident activation buffer pointer for the current thread.
/// Used by dense-decode paths when the output of one matvec (e.g.
/// final_norm) is known to be the input of the next (e.g. lm_head).
/// Call this *after* downloading the result to host, passing the host
/// slice pointer+length. The next `launch_matvec` call can skip
/// `htod_copy` if it sees the same pointer.
pub fn set_resident_activation(x: &[f32]) {
    RESIDENT_ACT.set(Some(ResidentActivation {
        ptr: x.as_ptr(),
        len: x.len(),
    }));
}

/// Clears the resident activation TLS (use after a decode that doesn't
/// consume it, or at the start of a fresh generation, to avoid stale
/// pointer reuse).
pub fn clear_resident_activation() {
    RESIDENT_ACT.set(None);
}

/// Checks if a resident activation matches `x` (same pointer+length),
/// and if so, returns `true` and clears the TLS. `launch_matvec` can
/// then skip `htod_copy(x)` because the data is already on device from
/// the prior call.
///
/// Not yet called by `launch_matvec` -- future work to wire this check
/// in before the `htod_copy` inside that function.
fn take_resident_activation_if_matches(x: &[f32]) -> bool {
    RESIDENT_ACT
        .take()
        .is_some_and(|res| res.len == x.len() && res.ptr == x.as_ptr())
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn probe_degrades_cleanly_regardless_of_real_hardware_presence() {
        // This assertion used to hard-code "no CUDA device here," true
        // on the sandbox this was written on but false the first time
        // this actually ran on rented GPU hardware
        // (real RTX 3060, 2026-07-31) --
        // caught a real test bug, not a probe() bug: a CI-style test
        // that only passes in one specific environment isn't testing
        // the real invariant (probe() never panics and reports
        // something sane either way), it's testing "this machine has no
        // GPU," which isn't ferrox's to assert. Check both real outcomes
        // instead of assuming one.
        match probe() {
            None => {} // no device present -- also correct, nothing further to check
            Some(info) => {
                assert!(
                    info.device_count >= 1,
                    "reported a device but device_count=0"
                );
                assert!(
                    info.total_vram_bytes > 0,
                    "reported a device but total_vram_bytes=0"
                );
                assert!(
                    info.free_vram_bytes <= info.total_vram_bytes,
                    "free VRAM ({}) cannot exceed total ({})",
                    info.free_vram_bytes,
                    info.total_vram_bytes
                );
            }
        }
    }

    /// Builds `rows` real, non-trivial Q8_0-quantized rows (not all-
    /// zero weights, which would pass trivially even with a broken
    /// kernel) using `ferrox_quant::quantize_q8_0`, so the ignored GPU
    /// tests below check actual numerical agreement with the CPU
    /// reference, not just "the launch didn't error."
    fn real_q8_0_test_matrix(rows: usize, cols: usize) -> (Vec<u8>, Vec<f32>, Vec<f32>) {
        let mut weights = Vec::new();
        let mut all_rows_f32 = Vec::new();
        for r in 0..rows {
            let row: Vec<f32> = (0..cols)
                .map(|i| (((r * cols + i) as f32) * 0.037).sin())
                .collect();
            weights.extend(ferrox_quant::quantize_q8_0(&row));
            all_rows_f32.push(row);
        }
        let x: Vec<f32> = (0..cols).map(|i| ((i as f32) * 0.019).cos()).collect();
        let expected: Vec<f32> = all_rows_f32
            .iter()
            .map(|row| ferrox_quant::dot_q8_0_f32_scalar(&ferrox_quant::quantize_q8_0(row), &x))
            .collect();
        (weights, x, expected)
    }

    #[test]
    #[ignore = "requires real CUDA hardware -- verified passing on an RTX 3060 (vast.ai); run with --ignored on a CUDA-capable machine to re-verify"]
    fn launch_q8_0_matvec_matches_cpu_reference() {
        // Run manually on a machine with an actual CUDA device:
        //   cargo test -p ferrox-cuda --features cuda -- --ignored
        let rows = 4;
        let cols = 64;
        let row_bytes = (cols / ferrox_quant::Q8_0_BLOCK_ELEMS) * ferrox_quant::Q8_0_BLOCK_BYTES;
        let (weights, x, expected) = real_q8_0_test_matrix(rows, cols);

        let result = launch_q8_0_matvec(
            &weights,
            &x,
            rows,
            row_bytes,
            cols / ferrox_quant::Q8_0_BLOCK_ELEMS,
        )
        .expect("kernel launch must succeed on real CUDA hardware");
        assert_eq!(result.len(), expected.len());
        for (i, (got, want)) in result.iter().zip(expected.iter()).enumerate() {
            assert!(
                (got - want).abs() < 1e-2,
                "row {i}: GPU={got} CPU reference={want}"
            );
        }
    }

    #[test]
    #[ignore = "requires real CUDA hardware -- verified passing on an RTX 3060 (vast.ai); run with --ignored on a CUDA-capable machine to re-verify"]
    fn launch_q4_0_matvec_matches_cpu_reference() {
        // Run manually on a machine with an actual CUDA device:
        //   cargo test -p ferrox-cuda --features cuda -- --ignored
        let rows = 4;
        let cols = 64;
        let blocks_per_row = cols / ferrox_quant::Q4_0_BLOCK_ELEMS;
        let row_bytes = blocks_per_row * ferrox_quant::Q4_0_BLOCK_BYTES;

        // Regression note: an earlier version of this test built exactly
        // one 18-byte Q4_0 block per row regardless of `cols`, so for
        // cols=64 (2 blocks/row, row_bytes=36) the `weights` buffer ended
        // up 72 bytes instead of the 144 bytes `row_bytes_slice` below
        // actually indexes into -- caught by a real out-of-bounds panic
        // the first time this test ran on real GPU hardware. Fixed by
        // building `blocks_per_row` blocks per row, matching `cols`.
        let mut weights = Vec::new();
        for r in 0..rows {
            for b in 0..blocks_per_row {
                weights.extend_from_slice(
                    &half::f16::from_f32(0.05 + (r * blocks_per_row + b) as f32 * 0.01)
                        .to_le_bytes(),
                );
                for i in 0..16u8 {
                    let lo = (i + r as u8 + b as u8) % 16;
                    let hi = (15 - i + r as u8 + b as u8) % 16;
                    weights.push(lo | (hi << 4));
                }
            }
        }
        let x: Vec<f32> = (0..cols).map(|i| ((i as f32) * 0.09).sin()).collect();
        let expected: Vec<f32> = (0..rows)
            .map(|r| {
                let row_bytes_slice = &weights[r * row_bytes..(r + 1) * row_bytes];
                ferrox_quant::dot_q4_0_f32_scalar(row_bytes_slice, &x)
            })
            .collect();

        let result = launch_q4_0_matvec(
            &weights,
            &x,
            rows,
            row_bytes,
            cols / ferrox_quant::Q4_0_BLOCK_ELEMS,
        )
        .expect("kernel launch must succeed on real CUDA hardware");
        assert_eq!(result.len(), expected.len());
        for (i, (got, want)) in result.iter().zip(expected.iter()).enumerate() {
            assert!(
                (got - want).abs() < 1e-2,
                "row {i}: GPU={got} CPU reference={want}"
            );
        }
    }

    /// Deterministic pseudo-random byte generator for building real,
    /// non-trivial K-quant block bytes (no `quantize_qX_k` producer
    /// exists in `ferrox_quant` -- these formats are load-only, never
    /// produced by ferrox -- so tests build arbitrary-but-valid-shaped
    /// bytes directly, the same convention `ferrox-core::weight_matrix`'s
    /// and `ferrox-models::kimi_loader`'s own MXFP4 tests already use).
    fn pseudo_bytes(seed: u32, len: usize) -> Vec<u8> {
        let mut state = seed.wrapping_mul(2654435761).wrapping_add(1);
        (0..len)
            .map(|_| {
                state = state.wrapping_mul(1103515245).wrapping_add(12345);
                (state >> 16) as u8
            })
            .collect()
    }

    /// GPU-vs-CPU agreement check for the K-quant kernels below, using a
    /// *relative* error bound instead of the fixed absolute `1e-2` the
    /// Q8_0/Q4_0 tests above use. The K-quant kernels apply each block's
    /// scale/min *inside* a per-element `acc += (d1 * q - min1) * x[i]`
    /// accumulation (hundreds of float multiply-adds per row), unlike
    /// Q8_0/Q4_0's structure (accumulate an exact integer dot product
    /// first, multiply by scale once at the end) -- so K-quant results are
    /// exposed to per-element GPU-vs-CPU float rounding differences
    /// (nvcc/NVRTC contracts `a*b+c`-shaped expressions into a single-
    /// rounding `fma` instruction by default; plain Rust `f32` arithmetic
    /// does not auto-contract). Real hardware run (RTX 3060, 2026-07-31)
    /// measured relative errors around 1e-7 (e.g. diff=40 against a
    /// ~97,454,420-magnitude reference) -- consistent with float32
    /// machine epsilon times a handful of differently-ordered rounding
    /// steps, not a logic bug. This is exactly why llama.cpp's own real
    /// CUDA K-quant kernels (`vec_dot_q6_K_q8_1_impl_mmvq`,
    /// `ggml/src/ggml-cuda/vecdotq.cuh`) quantize activations to int8 too
    /// and reduce via an *integer* `dp4a` dot product, converting to
    /// float only once per block for the final scale multiply --
    /// deliberately avoiding this exact class of divergence -- and why
    /// llama.cpp's own backend-comparison tests
    /// (`tests/test-backend-ops.cpp`) check `max_nmse_err()` (a relative
    /// error metric), never exact/absolute equality, for quantized ops.
    fn assert_close_relative(got: f32, want: f32, row: usize) {
        // Real hardware run (RTX 3060, 2026-07-31) hit `GPU=NaN CPU
        // reference=NaN` for Q6_K row 1: `real_k_quant_test_matrix`'s
        // pseudo-random block bytes aren't pinned to safe values (unlike
        // some other tests here), so by chance a row's raw `d`/scale
        // bytes decoded as an f16 NaN/Inf bit pattern -- garbage in,
        // garbage out, identically on both backends. GPU and CPU *agree*
        // in that case (both NaN from the same degenerate input), which
        // is a real pass, not a failure to paper over; plain `(got -
        // want).abs() <= tol` fails to see that because IEEE754 NaN
        // comparisons are always false. Only a mismatched
        // NaN-vs-non-NaN pairing is a genuine disagreement.
        if want.is_nan() {
            assert!(
                got.is_nan(),
                "row {row}: CPU reference is NaN but GPU={got} is not"
            );
            return;
        }
        let tol = 1e-4 * want.abs().max(1.0);
        assert!(
            (got - want).abs() <= tol,
            "row {row}: GPU={got} CPU reference={want} (relative tolerance {tol})"
        );
    }

    /// Builds `rows` real (non-zero, non-trivial) blocks of `block_bytes`
    /// each for a K-quant format, and the matching `expected` output via
    /// `scalar_dot` (`ferrox_quant::dot_q{4,5,6}_k_f32_scalar` -- already
    /// independently verified elsewhere in this workspace), so the
    /// ignored GPU tests below check real numerical agreement with that
    /// trusted CPU reference, not just "the launch didn't error."
    fn real_k_quant_test_matrix(
        rows: usize,
        cols: usize,
        block_bytes: usize,
        scalar_dot: impl Fn(&[u8], &[f32]) -> f32,
    ) -> (Vec<u8>, Vec<f32>, Vec<f32>) {
        let n_blocks_per_row = cols / 256;
        let row_bytes = n_blocks_per_row * block_bytes;
        let mut weights = Vec::with_capacity(rows * row_bytes);
        for r in 0..rows {
            weights.extend(pseudo_bytes(r as u32 + 1, row_bytes));
        }
        let x: Vec<f32> = (0..cols).map(|i| ((i as f32) * 0.021).sin()).collect();
        let expected: Vec<f32> = (0..rows)
            .map(|r| scalar_dot(&weights[r * row_bytes..(r + 1) * row_bytes], &x))
            .collect();
        (weights, x, expected)
    }

    #[test]
    #[ignore = "requires real CUDA hardware -- verified passing on an RTX 3060 (vast.ai, 2026-07-31) with the relative-tolerance fix; run with --ignored on a CUDA-capable machine to re-verify"]
    fn launch_q4_k_matvec_matches_cpu_reference() {
        let rows = 3;
        let cols = 256; // 1 Q4_K super-block per row
        let (weights, x, expected) = real_k_quant_test_matrix(
            rows,
            cols,
            ferrox_quant::Q4_K_BLOCK_BYTES,
            ferrox_quant::dot_q4_k_f32_scalar,
        );

        let result = launch_q4_k_matvec(&weights, &x, rows, ferrox_quant::Q4_K_BLOCK_BYTES, 1)
            .expect("kernel launch must succeed on real CUDA hardware");
        assert_eq!(result.len(), expected.len());
        for (i, (got, want)) in result.iter().zip(expected.iter()).enumerate() {
            assert_close_relative(*got, *want, i);
        }
    }

    #[test]
    #[ignore = "requires real CUDA hardware -- verified passing on an RTX 3060 (vast.ai, 2026-07-31) with the relative-tolerance fix; run with --ignored on a CUDA-capable machine to re-verify"]
    fn launch_q5_k_matvec_matches_cpu_reference() {
        let rows = 3;
        let cols = 256; // 1 Q5_K super-block per row
        let (weights, x, expected) = real_k_quant_test_matrix(
            rows,
            cols,
            ferrox_quant::Q5_K_BLOCK_BYTES,
            ferrox_quant::dot_q5_k_f32_scalar,
        );

        let result = launch_q5_k_matvec(&weights, &x, rows, ferrox_quant::Q5_K_BLOCK_BYTES, 1)
            .expect("kernel launch must succeed on real CUDA hardware");
        assert_eq!(result.len(), expected.len());
        for (i, (got, want)) in result.iter().zip(expected.iter()).enumerate() {
            assert_close_relative(*got, *want, i);
        }
    }

    #[test]
    #[ignore = "requires real CUDA hardware -- 2026-07-31 real RTX 3060 run failed on row 1 with GPU=NaN CPU reference=NaN (assert_close_relative didn't treat NaN==NaN as agreement; fixed, but not yet re-verified on hardware); run with --ignored on a CUDA-capable machine to re-verify"]
    fn launch_q6_k_matvec_matches_cpu_reference() {
        let rows = 3;
        let cols = 512; // 2 Q6_K super-blocks per row
        let (weights, x, expected) = real_k_quant_test_matrix(
            rows,
            cols,
            ferrox_quant::Q6_K_BLOCK_BYTES,
            ferrox_quant::dot_q6_k_f32_scalar,
        );

        let result = launch_q6_k_matvec(&weights, &x, rows, ferrox_quant::Q6_K_BLOCK_BYTES * 2, 2)
            .expect("kernel launch must succeed on real CUDA hardware");
        assert_eq!(result.len(), expected.len());
        for (i, (got, want)) in result.iter().zip(expected.iter()).enumerate() {
            assert_close_relative(*got, *want, i);
        }
    }

    #[test]
    #[ignore = "requires real CUDA hardware -- verifies that launch_matvec_multi (shared x upload) matches N sequential single-matvec launches; run with --ignored on a CUDA-capable machine"]
    fn launch_matvec_multi_matches_sequential() {
        // Build three Q4_K matrices (different shapes) that share the same
        // activation dimension.
        let cols = 256;
        let rows_a = 2;
        let rows_b = 3;
        let rows_c = 4;
        let (weights_a, x, expected_a) = real_k_quant_test_matrix(
            rows_a,
            cols,
            ferrox_quant::Q4_K_BLOCK_BYTES,
            ferrox_quant::dot_q4_k_f32_scalar,
        );
        let (weights_b, _, expected_b) = real_k_quant_test_matrix(
            rows_b,
            cols,
            ferrox_quant::Q4_K_BLOCK_BYTES,
            ferrox_quant::dot_q4_k_f32_scalar,
        );
        let (weights_c, _, expected_c) = real_k_quant_test_matrix(
            rows_c,
            cols,
            ferrox_quant::Q4_K_BLOCK_BYTES,
            ferrox_quant::dot_q4_k_f32_scalar,
        );

        let launches = [
            MatvecLaunch {
                kernel_src: Q4_K_MATVEC_KERNEL_SRC,
                module_name: "ferrox_q4_k",
                fn_name: "q4_k_matvec",
                weights: weights_a.as_slice(),
                rows: rows_a,
                row_bytes: ferrox_quant::Q4_K_BLOCK_BYTES,
                n_blocks_per_row: 1,
            },
            MatvecLaunch {
                kernel_src: Q4_K_MATVEC_KERNEL_SRC,
                module_name: "ferrox_q4_k",
                fn_name: "q4_k_matvec",
                weights: weights_b.as_slice(),
                rows: rows_b,
                row_bytes: ferrox_quant::Q4_K_BLOCK_BYTES,
                n_blocks_per_row: 1,
            },
            MatvecLaunch {
                kernel_src: Q4_K_MATVEC_KERNEL_SRC,
                module_name: "ferrox_q4_k",
                fn_name: "q4_k_matvec",
                weights: weights_c.as_slice(),
                rows: rows_c,
                row_bytes: ferrox_quant::Q4_K_BLOCK_BYTES,
                n_blocks_per_row: 1,
            },
        ];

        let results = launch_matvec_multi(&x, &launches)
            .expect("multi-matvec must succeed on real CUDA hardware");
        assert_eq!(results.len(), 3);

        // Check first matrix output.
        assert_eq!(results[0].len(), expected_a.len());
        for (i, (got, want)) in results[0].iter().zip(expected_a.iter()).enumerate() {
            assert_close_relative(*got, *want, i);
        }

        // Check second matrix output.
        assert_eq!(results[1].len(), expected_b.len());
        for (i, (got, want)) in results[1].iter().zip(expected_b.iter()).enumerate() {
            assert_close_relative(*got, *want, i);
        }

        // Check third matrix output.
        assert_eq!(results[2].len(), expected_c.len());
        for (i, (got, want)) in results[2].iter().zip(expected_c.iter()).enumerate() {
            assert_close_relative(*got, *want, i);
        }
    }

    /// The fused on-device dense SwiGLU FFN
    /// (`launch_dense_ffn_swiglu`: one `x` upload → gate matvec → up
    /// matvec → `silu(gate)*up` kernel → down matvec → one download)
    /// must produce the same result as running the three matvecs
    /// separately on the CPU reference and doing SwiGLU on the host.
    /// This proves the device-resident activation chaining (no per-
    /// matvec DtoH/HtoD) and the `silu_mul_f32` kernel are correct, not
    /// just that the launches didn't error. Q8_0 is used because
    /// `ferrox_quant::quantize_q8_0` gives an exact producer/consumer
    /// pair, keeping this a real numerical-agreement check.
    #[test]
    #[ignore = "requires real CUDA hardware -- run with --ignored on a CUDA-capable machine to verify the fused dense-FFN activation-residency path"]
    fn launch_dense_ffn_swiglu_matches_sequential_cpu() {
        let hidden_dim = 64; // 2 Q8_0 blocks per row
        let ffn_dim = 96; // 3 Q8_0 blocks per row; gate/up output length

        let make_row = |cols: usize, seed: f32| -> Vec<f32> {
            (0..cols)
                .map(|i| (((i as f32) - (cols as f32) / 2.0) * 0.013 * seed).sin())
                .collect()
        };
        // Build a quantized [rows, cols] matrix plus keep each row's f32
        // so the CPU reference dots the *same* quantized bytes back.
        let build = |rows: usize, cols: usize, seed: f32| -> (Vec<u8>, usize) {
            let mut packed = Vec::new();
            for r in 0..rows {
                packed.extend(ferrox_quant::quantize_q8_0(&make_row(
                    cols,
                    seed + r as f32,
                )));
            }
            let row_bytes =
                (cols / ferrox_quant::Q8_0_BLOCK_ELEMS) * ferrox_quant::Q8_0_BLOCK_BYTES;
            (packed, row_bytes)
        };

        let (gate_w, gate_rb) = build(ffn_dim, hidden_dim, 1.0);
        let (up_w, up_rb) = build(ffn_dim, hidden_dim, 2.0);
        let (down_w, down_rb) = build(hidden_dim, ffn_dim, 3.0);
        let x = make_row(hidden_dim, 0.7);

        let blk = |cols: usize| cols / ferrox_quant::Q8_0_BLOCK_ELEMS;
        let gate = MatvecLaunch {
            kernel_src: Q8_0_MATVEC_KERNEL_SRC,
            module_name: "ferrox_q8_0",
            fn_name: "q8_0_matvec",
            weights: gate_w.as_slice(),
            rows: ffn_dim,
            row_bytes: gate_rb,
            n_blocks_per_row: blk(hidden_dim),
        };
        let up = MatvecLaunch {
            kernel_src: Q8_0_MATVEC_KERNEL_SRC,
            module_name: "ferrox_q8_0",
            fn_name: "q8_0_matvec",
            weights: up_w.as_slice(),
            rows: ffn_dim,
            row_bytes: up_rb,
            n_blocks_per_row: blk(hidden_dim),
        };
        let down = MatvecLaunch {
            kernel_src: Q8_0_MATVEC_KERNEL_SRC,
            module_name: "ferrox_q8_0",
            fn_name: "q8_0_matvec",
            weights: down_w.as_slice(),
            rows: hidden_dim,
            row_bytes: down_rb,
            n_blocks_per_row: blk(ffn_dim),
        };

        let gpu = launch_dense_ffn_swiglu(&gate, &up, &down, &x)
            .expect("fused FFN must launch on real CUDA hardware");

        // CPU reference: gate/up matvecs, host SwiGLU, down matvec.
        let cpu_matvec = |w: &[u8], row_bytes: usize, rows: usize, act: &[f32]| -> Vec<f32> {
            (0..rows)
                .map(|r| {
                    ferrox_quant::dot_q8_0_f32_scalar(&w[r * row_bytes..(r + 1) * row_bytes], act)
                })
                .collect::<Vec<f32>>()
        };
        let g = cpu_matvec(&gate_w, gate_rb, ffn_dim, &x);
        let u = cpu_matvec(&up_w, up_rb, ffn_dim, &x);
        let act: Vec<f32> = g
            .iter()
            .zip(u.iter())
            .map(|(gv, uv)| (gv / (1.0 + (-gv).exp())) * uv)
            .collect();
        let expected = cpu_matvec(&down_w, down_rb, hidden_dim, &act);

        assert_eq!(gpu.len(), expected.len());
        for (i, (got, want)) in gpu.iter().zip(expected.iter()).enumerate() {
            assert_close_relative(*got, *want, i);
        }
    }

    /// `upload_act` → `matvec_into` → `download_act` (device-resident
    /// activation, no per-matvec host round-trip) must agree with the
    /// host-wrapper `launch_q8_0_matvec` for the same weights/input.
    #[test]
    #[ignore = "requires real CUDA hardware -- run with --ignored to verify DeviceAct residency vs the host-wrapper matvec"]
    fn matvec_into_matches_host_wrapper() {
        let rows = 4;
        let cols = 64;
        let row_bytes = (cols / ferrox_quant::Q8_0_BLOCK_ELEMS) * ferrox_quant::Q8_0_BLOCK_BYTES;
        let (weights, x, _expected) = real_q8_0_test_matrix(rows, cols);

        let host = launch_q8_0_matvec(&weights, &x, rows, row_bytes, cols / 32)
            .expect("host-wrapper matvec must launch");

        let d_x = upload_act(&x).expect("upload");
        let launch = MatvecLaunch {
            kernel_src: Q8_0_MATVEC_KERNEL_SRC,
            module_name: "ferrox_q8_0",
            fn_name: "q8_0_matvec",
            weights: weights.as_slice(),
            rows,
            row_bytes,
            n_blocks_per_row: cols / 32,
        };
        let d_out = matvec_into(&launch, &d_x).expect("device matvec");
        let device = download_act(&d_out).expect("download");

        assert_eq!(host.len(), device.len());
        for (i, (h, d)) in host.iter().zip(device.iter()).enumerate() {
            assert_close_relative(*d, *h, i);
        }
    }

    /// Fused `rms_norm(x + residual, weight, eps)` on device must agree
    /// with the CPU reference (`ferrox_core::matmul::rms_norm` on the
    /// elementwise sum).
    #[test]
    #[ignore = "requires real CUDA hardware -- run with --ignored to verify fused_add_rmsnorm vs CPU rms_norm"]
    fn launch_fused_add_rmsnorm_matches_cpu_reference() {
        let n = 128;
        let x: Vec<f32> = (0..n).map(|i| ((i as f32) * 0.031).sin()).collect();
        let residual: Vec<f32> = (0..n).map(|i| ((i as f32) * 0.017).cos()).collect();
        let weight: Vec<f32> = (0..n).map(|i| 1.0 + ((i as f32) * 0.003).sin()).collect();
        let eps = 1e-5f32;

        let sum: Vec<f32> = x.iter().zip(residual.iter()).map(|(a, b)| a + b).collect();
        let mean_sq = sum.iter().map(|v| v * v).sum::<f32>() / n as f32;
        let scale = 1.0 / (mean_sq + eps).sqrt();
        let expected: Vec<f32> = sum
            .iter()
            .zip(weight.iter())
            .map(|(v, w)| v * scale * w)
            .collect();

        let gpu = launch_fused_add_rmsnorm(&x, &residual, &weight, eps)
            .expect("fused_add_rmsnorm must launch on CUDA hardware");

        assert_eq!(gpu.len(), expected.len());
        for (i, (got, want)) in gpu.iter().zip(expected.iter()).enumerate() {
            assert_close_relative(*got, *want, i);
        }
    }
}