xserv/crates/xserv-kernels/src/quantization.rs

use std::cell::RefCell;
use std::collections::HashMap;
use std::ffi::c_void;
use xserv_cuda::GpuBuffer;
use xserv_tensor::{DType, Tensor};

// ============================================================
// FFI: custom CUDA kernels
// ============================================================

unsafe extern "C" {
    fn launch_dequant_fp8e4m3_to_bf16(
        src: *const c_void,
        scales: *const c_void,
        dst: *mut c_void,
        num_experts: i32, rows: i32, cols: i32,
        stream: *mut c_void,
    );
    fn launch_quantize_bf16_to_fp8e4m3_rowwise(
        src: *const c_void,
        dst: *mut c_void,
        scales: *mut c_void,
        num_rows: i32, cols: i32,
        stream: *mut c_void,
    );
    fn launch_rowwise_scale_bf16(
        data: *mut c_void,
        scales: *const c_void,
        num_rows: i32, cols: i32,
        stream: *mut c_void,
    );
}

// ============================================================
// FFI: cuBLASLt
// ============================================================

type CublasLtHandle = *mut c_void;
type CublasLtMatmulDesc = *mut c_void;
type CublasLtMatrixLayout = *mut c_void;
type CublasLtMatmulPreference = *mut c_void;

#[repr(C)]
#[derive(Clone, Copy)]
struct CublasLtMatmulAlgo {
    data: [u64; 8],
}

#[repr(C)]
struct CublasLtMatmulHeuristicResult {
    algo: CublasLtMatmulAlgo,
    workspace_size: usize,
    state: i32,
    _reserved: [f32; 4],
}

unsafe extern "C" {
    fn cublasLtCreate(handle: *mut CublasLtHandle) -> i32;
    fn cublasLtDestroy(handle: CublasLtHandle) -> i32;
    fn cublasLtMatmulDescCreate(desc: *mut CublasLtMatmulDesc, compute_type: i32, scale_type: i32) -> i32;
    fn cublasLtMatmulDescDestroy(desc: CublasLtMatmulDesc) -> i32;
    fn cublasLtMatmulDescSetAttribute(desc: CublasLtMatmulDesc, attr: i32, buf: *const c_void, size: usize) -> i32;
    fn cublasLtMatrixLayoutCreate(layout: *mut CublasLtMatrixLayout, dtype: i32, rows: u64, cols: u64, ld: i64) -> i32;
    fn cublasLtMatrixLayoutDestroy(layout: CublasLtMatrixLayout) -> i32;
    fn cublasLtMatrixLayoutSetAttribute(layout: CublasLtMatrixLayout, attr: i32, buf: *const c_void, size: usize) -> i32;
    fn cublasLtMatmulPreferenceCreate(pref: *mut CublasLtMatmulPreference) -> i32;
    fn cublasLtMatmulPreferenceDestroy(pref: CublasLtMatmulPreference) -> i32;
    fn cublasLtMatmulPreferenceSetAttribute(pref: CublasLtMatmulPreference, attr: i32, buf: *const c_void, size: usize) -> i32;
    fn cublasLtMatmulAlgoGetHeuristic(
        handle: CublasLtHandle, desc: CublasLtMatmulDesc,
        a_layout: CublasLtMatrixLayout, b_layout: CublasLtMatrixLayout,
        c_layout: CublasLtMatrixLayout, d_layout: CublasLtMatrixLayout,
        pref: CublasLtMatmulPreference,
        requested: i32,
        results: *mut CublasLtMatmulHeuristicResult,
        found: *mut i32,
    ) -> i32;
    fn cublasLtMatmul(
        handle: CublasLtHandle, desc: CublasLtMatmulDesc,
        alpha: *const c_void,
        a: *const c_void, a_layout: CublasLtMatrixLayout,
        b: *const c_void, b_layout: CublasLtMatrixLayout,
        beta: *const c_void,
        c: *const c_void, c_layout: CublasLtMatrixLayout,
        d: *mut c_void, d_layout: CublasLtMatrixLayout,
        algo: *const CublasLtMatmulAlgo,
        workspace: *mut c_void, workspace_size: usize,
        stream: *mut c_void,
    ) -> i32;
}

// cuBLASLt constants
const CUBLAS_COMPUTE_32F: i32 = 68;
const CUDA_R_32F: i32 = 0;
const CUDA_R_16BF: i32 = 14;
const CUDA_R_8F_E4M3: i32 = 28;

// MatmulDesc attributes
const CUBLASLT_MATMUL_DESC_A_SCALE_POINTER: i32 = 17;
const CUBLASLT_MATMUL_DESC_B_SCALE_POINTER: i32 = 18;
const CUBLASLT_MATMUL_DESC_A_SCALE_MODE: i32 = 31;
const CUBLASLT_MATMUL_DESC_B_SCALE_MODE: i32 = 32;

// MatrixLayout attributes
const CUBLASLT_MATRIX_LAYOUT_BATCH_COUNT: i32 = 5;
const CUBLASLT_MATRIX_LAYOUT_STRIDED_BATCH_OFFSET: i32 = 6;

// Scale modes
const CUBLASLT_MATMUL_MATRIX_SCALE_SCALAR: i32 = 0;
const CUBLASLT_MATMUL_MATRIX_SCALE_OUTER_VEC_32F: i32 = 3;

// MatmulPreference attributes
const CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES: i32 = 1;

const WORKSPACE_BYTES: usize = 32 * 1024 * 1024;

const CUBLASLT_MATMUL_DESC_TRANSA: i32 = 3;

/// A fully-prepared FP8 matmul plan for one (M, N, K) shape: the matmul
/// descriptor, the four matrix layouts, and the heuristically-chosen algo.
/// Built once per shape and reused across every expert and every forward
/// pass — the heuristic search and descriptor/layout creation are the
/// expensive parts, so doing them once instead of per-expert-per-layer is
/// the difference between FP8 being faster or slower than BF16.
#[derive(Clone, Copy)]
struct Fp8Plan {
    desc: CublasLtMatmulDesc,
    a_layout: CublasLtMatrixLayout,
    b_layout: CublasLtMatrixLayout,
    c_layout: CublasLtMatrixLayout,
    d_layout: CublasLtMatrixLayout,
    algo: CublasLtMatmulAlgo,
    workspace_size: usize,
}

struct CublasLtContext {
    handle: CublasLtHandle,
    workspace: GpuBuffer,
    /// Persistent device scalar holding 1.0, used as the A/B scale pointer
    /// placeholder. Allocated once instead of per-expert.
    one_buf: GpuBuffer,
    /// Cache of prepared matmul plans keyed by (M, N, K).
    plans: HashMap<(usize, usize, usize), Fp8Plan>,
}

impl CublasLtContext {
    fn new() -> Self {
        let mut handle = std::ptr::null_mut();
        let status = unsafe { cublasLtCreate(&mut handle) };
        assert_eq!(status, 0, "cublasLtCreate failed: {status}");
        let workspace = GpuBuffer::alloc(WORKSPACE_BYTES).expect("alloc cublasLt workspace");
        let mut one_buf = GpuBuffer::alloc(4).expect("alloc cublasLt fp8 scale");
        one_buf.copy_from_host(&1.0f32.to_le_bytes()).expect("init fp8 scale");
        Self { handle, workspace, one_buf, plans: HashMap::new() }
    }

    /// Get the cached plan for (m, n, k), building (and caching) it on first use.
    fn plan(&mut self, m: usize, n: usize, k: usize) -> Fp8Plan {
        if let Some(p) = self.plans.get(&(m, n, k)) {
            return *p;
        }
        let one_ptr = self.one_buf.as_ptr() as *const c_void;
        let plan = unsafe { build_fp8_plan(self.handle, one_ptr, m, n, k) };
        self.plans.insert((m, n, k), plan);
        plan
    }
}

impl Drop for CublasLtContext {
    fn drop(&mut self) {
        // Tear down cached plans before destroying the handle.
        for (_, p) in self.plans.drain() {
            unsafe {
                cublasLtMatrixLayoutDestroy(p.a_layout);
                cublasLtMatrixLayoutDestroy(p.b_layout);
                cublasLtMatrixLayoutDestroy(p.c_layout);
                cublasLtMatrixLayoutDestroy(p.d_layout);
                cublasLtMatmulDescDestroy(p.desc);
            }
        }
        if !self.handle.is_null() {
            unsafe { cublasLtDestroy(self.handle) };
        }
    }
}

/// Build an FP8 matmul plan for one (m, n, k) shape. See `batched_gemm_fp8`
/// for the row-major → cuBLASLt col-major layout mapping (transA=T, transB=N,
/// m_lt=N, n_lt=M, k_lt=K). The B-scale pointer is initialised to `one_ptr`
/// and overwritten per-expert at call time.
unsafe fn build_fp8_plan(
    handle: CublasLtHandle,
    one_ptr: *const c_void,
    m: usize,
    n: usize,
    k: usize,
) -> Fp8Plan {
    let m_lt = n as u64;
    let n_lt = m as u64;
    let k_lt = k as u64;

    let mut desc: CublasLtMatmulDesc = std::ptr::null_mut();
    cublasLtMatmulDescCreate(&mut desc, CUBLAS_COMPUTE_32F, CUDA_R_32F);

    // transA=T (required for FP8 on Blackwell)
    let trans_a: i32 = 1;
    cublasLtMatmulDescSetAttribute(desc, CUBLASLT_MATMUL_DESC_TRANSA, &trans_a as *const i32 as _, 4);
    let ptr_sz = std::mem::size_of::<*const c_void>();
    cublasLtMatmulDescSetAttribute(desc, CUBLASLT_MATMUL_DESC_A_SCALE_POINTER, &one_ptr as *const _ as _, ptr_sz);
    cublasLtMatmulDescSetAttribute(desc, CUBLASLT_MATMUL_DESC_B_SCALE_POINTER, &one_ptr as *const _ as _, ptr_sz);

    // "A" layout (weights, transposed): physical (K, N) col-major, ld=K
    let mut a_layout: CublasLtMatrixLayout = std::ptr::null_mut();
    cublasLtMatrixLayoutCreate(&mut a_layout, CUDA_R_8F_E4M3, k_lt, m_lt, k as i64);
    // "B" layout (activations): physical (K, M) col-major, ld=K
    let mut b_layout: CublasLtMatrixLayout = std::ptr::null_mut();
    cublasLtMatrixLayoutCreate(&mut b_layout, CUDA_R_8F_E4M3, k_lt, n_lt, k as i64);
    // "C"/"D" layout (output): physical (N, M) col-major, ld=N
    let mut c_layout: CublasLtMatrixLayout = std::ptr::null_mut();
    cublasLtMatrixLayoutCreate(&mut c_layout, CUDA_R_16BF, m_lt, n_lt, m_lt as i64);
    let mut d_layout: CublasLtMatrixLayout = std::ptr::null_mut();
    cublasLtMatrixLayoutCreate(&mut d_layout, CUDA_R_16BF, m_lt, n_lt, m_lt as i64);

    let mut pref: CublasLtMatmulPreference = std::ptr::null_mut();
    cublasLtMatmulPreferenceCreate(&mut pref);
    let ws_bytes = WORKSPACE_BYTES as u64;
    cublasLtMatmulPreferenceSetAttribute(pref, CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES, &ws_bytes as *const u64 as _, 8);

    let mut heuristic = std::mem::zeroed::<CublasLtMatmulHeuristicResult>();
    let mut found: i32 = 0;
    let status = cublasLtMatmulAlgoGetHeuristic(
        handle, desc, a_layout, b_layout, c_layout, d_layout,
        pref, 1, &mut heuristic, &mut found,
    );
    assert!(status == 0 && found > 0,
        "cublasLtMatmulAlgoGetHeuristic failed for FP8 GEMM (m={m}, n={n}, k={k}): status={status}, found={found}");
    cublasLtMatmulPreferenceDestroy(pref);

    Fp8Plan {
        desc, a_layout, b_layout, c_layout, d_layout,
        algo: heuristic.algo,
        workspace_size: heuristic.workspace_size,
    }
}

thread_local! {
    static CUBLASLT_CTX: RefCell<CublasLtContext> = RefCell::new(CublasLtContext::new());
}

// ============================================================
// Public API
// ============================================================

/// Dequantize a 3D FP8 E4M3 tensor to BF16 using per-expert FP32 scales.
///
/// src:    [num_experts, rows, cols] FP8E4M3, contiguous, GPU
/// scales: [num_experts] F32, contiguous, GPU
///
/// Returns: [num_experts, rows, cols] BF16
pub fn dequant_fp8_to_bf16(src: &Tensor, scales: &Tensor) -> Tensor {
    assert_eq!(src.ndim(), 3, "dequant_fp8_to_bf16: src must be 3D");
    assert_eq!(src.dtype(), DType::FP8E4M3);
    assert!(src.is_contiguous());
    assert_eq!(scales.ndim(), 1);
    assert_eq!(scales.dtype(), DType::F32);
    assert!(scales.is_contiguous());

    let num_experts = src.shape()[0];
    let rows = src.shape()[1];
    let cols = src.shape()[2];
    assert_eq!(scales.shape()[0], num_experts);

    let out = Tensor::empty(&[num_experts, rows, cols], DType::BF16, src.device());

    unsafe {
        launch_dequant_fp8e4m3_to_bf16(
            src.data_ptr() as *const c_void,
            scales.data_ptr() as *const c_void,
            out.data_ptr() as *mut c_void,
            num_experts as i32, rows as i32, cols as i32,
            std::ptr::null_mut(),
        );
    }

    out
}

/// Dynamically quantize a contiguous BF16 tensor to FP8 E4M3 with per-row scales.
///
/// src: [num_rows, cols] or [batch, rows, cols] BF16, contiguous, GPU
/// Treats the tensor as 2D (flattens leading dims into num_rows).
///
/// Returns: (fp8_data [same shape] FP8E4M3, scales [total_rows] F32)
pub fn quantize_bf16_to_fp8_rowwise(src: &Tensor) -> (Tensor, Tensor) {
    assert_eq!(src.dtype(), DType::BF16);
    assert!(src.is_contiguous());
    assert!(src.ndim() >= 2);

    let cols = src.shape()[src.ndim() - 1];
    let num_rows: usize = src.shape()[..src.ndim() - 1].iter().product();

    let fp8_out = Tensor::empty(src.shape(), DType::FP8E4M3, src.device());
    let scales = Tensor::empty(&[num_rows], DType::F32, src.device());

    unsafe {
        launch_quantize_bf16_to_fp8e4m3_rowwise(
            src.data_ptr() as *const c_void,
            fp8_out.data_ptr() as *mut c_void,
            scales.data_ptr() as *mut c_void,
            num_rows as i32, cols as i32,
            std::ptr::null_mut(),
        );
    }

    (fp8_out, scales)
}

/// FP8 batched GEMM via cuBLASLt (transA=T required on Blackwell).
///
/// Computes: C[b] = scale_a[b] * scale_b[b] * (A_fp8[b] @ B_fp8_T[b]^T)
/// effectively C[b] = A[b, M, K] @ W[b, K, N] but W is stored transposed
/// as [b, N, K] for cuBLASLt FP8 compatibility.
///
/// a_fp8:    [batch, M, K] FP8E4M3  (activations, quantized per-row)
/// a_scales: [batch * M] F32        (per-token activation scales, applied post-GEMM)
/// b_fp8_t:  [batch, N, K] FP8E4M3  (weights, TRANSPOSED for cuBLASLt)
/// b_scales: [batch] F32            (per-expert scalar weight scales, applied in-GEMM)
///
/// Returns:  [batch, M, N] BF16
pub fn batched_gemm_fp8(
    a_fp8: &Tensor,
    a_scales: &Tensor,
    b_fp8_t: &Tensor,
    b_scales: &Tensor,
) -> Tensor {
    assert_eq!(a_fp8.ndim(), 3);
    assert_eq!(b_fp8_t.ndim(), 3);
    assert_eq!(a_fp8.dtype(), DType::FP8E4M3);
    assert_eq!(b_fp8_t.dtype(), DType::FP8E4M3);
    assert!(a_fp8.is_contiguous() && b_fp8_t.is_contiguous());
    assert_eq!(a_fp8.shape()[0], b_fp8_t.shape()[0]);
    // b_fp8_t is [batch, N, K] transposed, so b_fp8_t.shape[2] == K == a_fp8.shape[2]
    assert_eq!(a_fp8.shape()[2], b_fp8_t.shape()[2]);

    let batch = a_fp8.shape()[0];
    let m = a_fp8.shape()[1]; // tokens
    let k = a_fp8.shape()[2]; // hidden
    let n = b_fp8_t.shape()[1]; // out_dim (from transposed weight)

    // a_scales: [batch * M] per-token activation scales (applied post-GEMM, per row).
    // b_scales: [batch] per-expert scalar weight scales (applied in-GEMM via B-scale ptr).
    assert_eq!(a_scales.shape()[0], batch * m);
    assert_eq!(b_scales.shape()[0], batch);

    let c = Tensor::empty(&[batch, m, n], DType::BF16, a_fp8.device());

    // Strides (in bytes) for one expert slice
    let stride_a = m * k;     // FP8: 1 byte per elem
    let stride_b = n * k;     // FP8: 1 byte per elem (transposed: [N, K])
    let stride_c = m * n * 2; // BF16: 2 bytes per elem

    CUBLASLT_CTX.with(|cell| {
        let mut ctx = cell.borrow_mut();
        let handle = ctx.handle;
        let ws_ptr = ctx.workspace.as_ptr() as *mut c_void;
        // Build (or fetch) the cached plan for this shape — heuristic search and
        // descriptor/layout creation happen once per (m, n, k), not per-expert.
        let plan = ctx.plan(m, n, k);

        // alpha=1, beta=0. Per-expert weight scale is supplied via the cuBLASLt
        // B-scale pointer (device, scalar): cuBLASLt computes in the FP32 epilogue
        //   D = (1.0 * A_fp8) @ (b_scale[e] * B_fp8)^T = b_scale[e] * (A_fp8 @ B_fp8^T)
        // Per-token activation scale (a_scale) is applied post-GEMM (per row).
        let alpha: f32 = 1.0;
        let beta: f32 = 0.0;
        let ptr_sz = std::mem::size_of::<*const c_void>();

        for e in 0..batch {
            let a_ptr = unsafe { (a_fp8.data_ptr() as *const u8).add(e * stride_a) as *const c_void };
            let b_ptr = unsafe { (b_fp8_t.data_ptr() as *const u8).add(e * stride_b) as *const c_void };
            let c_ptr = unsafe { (c.data_ptr() as *mut u8).add(e * stride_c) as *mut c_void };
            // Device pointer to this expert's scalar weight scale (FP32, 4 bytes).
            let b_scale_ptr = unsafe { (b_scales.data_ptr() as *const u8).add(e * 4) as *const c_void };

            unsafe {
                cublasLtMatmulDescSetAttribute(
                    plan.desc, CUBLASLT_MATMUL_DESC_B_SCALE_POINTER,
                    &b_scale_ptr as *const _ as _, ptr_sz,
                );
                let status = cublasLtMatmul(
                    handle, plan.desc,
                    &alpha as *const f32 as _,
                    b_ptr,    // cuBLASLt "A" = weights
                    plan.a_layout,
                    a_ptr,    // cuBLASLt "B" = activations
                    plan.b_layout,
                    &beta as *const f32 as _,
                    c_ptr,    // C (unused with beta=0)
                    plan.c_layout,
                    c_ptr,    // D = output
                    plan.d_layout,
                    &plan.algo,
                    ws_ptr,
                    plan.workspace_size,
                    std::ptr::null_mut(),
                );
                assert_eq!(status, 0, "cublasLtMatmul FP8 failed for expert {e}: status={status}");
            }
        }
    });

    // Post-GEMM: multiply each row of c by its activation scale.
    // c is [batch, M, N] BF16. a_scales is [batch * M] F32.
    // This recovers the per-token scale that was divided out during quantization.
    let total_rows = (batch * m) as i32;
    let cols = n as i32;
    unsafe {
        launch_rowwise_scale_bf16(
            c.data_ptr() as *mut c_void,
            a_scales.data_ptr() as *const c_void,
            total_rows, cols,
            std::ptr::null_mut(),
        );
    }

    c
}