release: merge fused MoE expert-combine + f32-direct quantize (gfx1151)

2026-06-25 22:09:33 +00:00
parent 41b0476653 c820d9209f
commit 5ff6df6b5a
3 changed files with 235 additions and 8 deletions
@@ -1012,6 +1012,10 @@ impl QTensor {
                // QStorage::indexed_moe_forward. HANZO_MOE_QMMQ_FALLBACK forces matvec for the A/B.
                let use_qmmq = t > 1 && std::env::var("HANZO_MOE_QMMQ_FALLBACK").is_err();
                let x_flat = match x_exp.dtype() {
+                    // qmmq quantizes f16/f32 activations natively, so keep the model's dtype and skip
+                    // the f32->f16 cast (a 16.7M-elem read+write per gate/up). Other dtypes (bf16 with
+                    // a symmetric expert type) still cast to f16.
+                    DType::F16 | DType::F32 if use_qmmq => x_exp.reshape((nrows, k))?.contiguous()?,
                    _ if use_qmmq => x_exp.reshape((nrows, k))?.to_dtype(DType::F16)?.contiguous()?,
                    DType::BF16 | DType::F16 => x_exp.reshape((nrows, k))?.contiguous()?,
                    _ => x_exp.reshape((nrows, k))?.to_dtype(DType::F16)?.contiguous()?,
@@ -1206,6 +1210,41 @@ fn rocm_moe_bank_cache(
    CACHE.get_or_init(|| std::sync::Mutex::new(std::collections::HashMap::new()))
 }

+/// FUSED MoE expert-combine: `out[i,j] = sum_e scores[i,e] * ys[i,e,j]`, reducing the per-expert
+/// outputs `ys` [t, topk, n] by the router weights `scores` [t, topk] into [t, n]. On ROCm this is
+/// ONE fused kernel (`RocmDevice::moe_combine`) instead of `ys.broadcast_mul(scores).sum(Minus2)`,
+/// which cast ys -> f32, wrote a [t,topk,n] f32 product temp, and ran an 8-wide strided reduce.
+/// Other backends keep the generic broadcast-mul + sum.
+pub fn moe_combine(ys: &Tensor, scores: &Tensor) -> Result<Tensor> {
+    let (t, topk, n) = ys.dims3()?;
+    #[cfg(feature = "rocm")]
+    if let Device::Rocm(dev) = ys.device() {
+        if std::env::var("HANZO_MOE_COMBINE_FALLBACK").is_ok() {
+            return ys.broadcast_mul(&scores.unsqueeze(D::Minus1)?)?.sum(D::Minus2);
+        }
+        let ys_c = ys.contiguous()?;
+        let scores_c = scores.to_dtype(DType::F32)?.contiguous()?;
+        let (ys_store, _) = ys_c.storage_and_layout();
+        let yr = match &*ys_store {
+            Storage::Rocm(r) => r,
+            _ => crate::bail!("moe_combine: ys not on rocm after contiguous()"),
+        };
+        let (sc_store, _) = scores_c.storage_and_layout();
+        let sr = match &*sc_store {
+            Storage::Rocm(r) => r,
+            _ => crate::bail!("moe_combine: scores not on rocm after contiguous()"),
+        };
+        let out = dev.moe_combine(yr, sr, t, topk, n)?;
+        return Ok(crate::tensor::from_storage(
+            Storage::Rocm(out),
+            (t, n),
+            crate::op::BackpropOp::none(),
+            false,
+        ));
+    }
+    ys.broadcast_mul(&scores.unsqueeze(D::Minus1)?)?.sum(D::Minus2)
+}
+
 thread_local! {
    static DEQUANTIZE_ALL: bool = {
        match std::env::var("DEQUANTIZE_ALL") {
@@ -1417,6 +1456,10 @@ impl QMatMul {
                // forces the matvec on prefill too (the before/after A/B + oracle-equivalence lever).
                let use_qmmq = t > 1 && std::env::var("HANZO_MOE_QMMQ_FALLBACK").is_err();
                let x_flat = match x_exp.dtype() {
+                    // qmmq quantizes f16/f32 activations natively, so keep the model's dtype and skip
+                    // the f32->f16 cast (a 16.7M-elem read+write per gate/up). Other dtypes (bf16 with
+                    // a symmetric expert type) still cast to f16.
+                    DType::F16 | DType::F32 if use_qmmq => x_exp.reshape((nrows, k))?.contiguous()?,
                    _ if use_qmmq => x_exp.reshape((nrows, k))?.to_dtype(DType::F16)?.contiguous()?,
                    DType::BF16 | DType::F16 => x_exp.reshape((nrows, k))?.contiguous()?,
                    _ => x_exp.reshape((nrows, k))?.to_dtype(DType::F16)?.contiguous()?,
@@ -1178,9 +1178,12 @@ impl RocmDevice {
        k: usize,
    ) -> Result<(RocmStorage, RocmStorage)> {
        use hanzo_rocm_kernels::kernel::{KernelSource, QuantKernel};
-        let x_ptr = match &x.slice {
-            RocmStorageSlice::F16(s) => s.as_ptr(),
-            _ => crate::bail!("quantize_q8: x must be f16"),
+        // Activations arrive f16 (dense prefill) or f32 (the MoE compute dtype); quantize each in its
+        // native dtype so the MoE glue need not pay an f32->f16 cast before this kernel.
+        let (kname, x_ptr) = match &x.slice {
+            RocmStorageSlice::F16(s) => ("quantize_q8", s.as_ptr()),
+            RocmStorageSlice::F32(s) => ("quantize_q8_f32", s.as_ptr()),
+            other => crate::bail!("quantize_q8: x must be f16 or f32, got {:?}", other.dtype()),
        };
        let nblk = k / 32;
        let xq = self.alloc::<u8>(m * k)?;
@@ -1197,7 +1200,7 @@ impl RocmDevice {
                self,
                QuantKernel::NAME,
                QuantKernel::CODE,
-                "quantize_q8",
+                kname,
                grid,
                block,
                &mut [
@@ -1259,9 +1262,13 @@ impl RocmDevice {
        k: usize,
    ) -> Result<(RocmStorage, RocmStorage, RocmStorage)> {
        use hanzo_rocm_kernels::kernel::{KernelSource, QuantKernel};
-        let x_ptr = match &x.slice {
-            RocmStorageSlice::F16(s) => s.as_ptr(),
-            _ => crate::bail!("quantize_q8_1: x must be f16"),
+        // f16 (dense prefill), f32 (MoE compute dtype), or bf16 (decode) -- quantize in the native
+        // dtype, no f32->f16 cast in the MoE glue.
+        let (kname, x_ptr) = match &x.slice {
+            RocmStorageSlice::F16(s) => ("quantize_q8_1", s.as_ptr()),
+            RocmStorageSlice::F32(s) => ("quantize_q8_1_f32", s.as_ptr()),
+            RocmStorageSlice::BF16(s) => ("quantize_q8_1_bf16", s.as_ptr()),
+            other => crate::bail!("quantize_q8_1: x must be f16/f32/bf16, got {:?}", other.dtype()),
        };
        let nblk = k / 32;
        let xq = self.alloc::<u8>(m * k)?;
@@ -1280,7 +1287,7 @@ impl RocmDevice {
                self,
                QuantKernel::NAME,
                QuantKernel::CODE,
-                "quantize_q8_1",
+                kname,
                grid,
                block,
                &mut [
@@ -1535,6 +1542,67 @@ impl RocmDevice {
        }
        Ok(RocmStorage { slice: RocmStorageSlice::F16(out), device: self.clone() })
    }
+
+    /// FUSED MoE expert-combine: `out[i,j] = sum_e scores[i,e] * ys[i,e,j]` in ONE launch. Replaces
+    /// the engine's `ys.broadcast_mul(scores).sum(Minus2)` (a 16.7M-elem ys->f32 cast + a 16.7M-elem
+    /// product temp + an 8-wide reduce at 2M blocks x 8 threads). `ys` is [ntok, topk, n] (f16/bf16/
+    /// f32, contiguous), `scores` is [ntok, topk] f32 contiguous; `out` is [ntok, n] in ys's dtype.
+    pub fn moe_combine(
+        &self,
+        ys: &RocmStorage,
+        scores: &RocmStorage,
+        ntok: usize,
+        topk: usize,
+        n: usize,
+    ) -> Result<RocmStorage> {
+        use hanzo_rocm_kernels::kernel::{KernelSource, QuantKernel};
+        use std::ffi::c_void;
+        let sc_ptr = match &scores.slice {
+            RocmStorageSlice::F32(s) => s.as_ptr(),
+            other => crate::bail!("moe_combine: scores must be f32, got {:?}", other.dtype()),
+        };
+        let total = ntok * n;
+        let nt = ntok as i32;
+        let tk = topk as i32;
+        let ni = n as i32;
+        let (grid, block) = launch_config(total);
+        macro_rules! launch_combine {
+            ($variant:ident, $rty:ty, $kernel:literal) => {{
+                let ys_ptr = match &ys.slice {
+                    RocmStorageSlice::$variant(s) => s.as_ptr(),
+                    _ => unreachable!(),
+                };
+                let out = self.alloc::<$rty>(total)?;
+                let out_ptr = out.as_ptr();
+                unsafe {
+                    launch_kernel(
+                        self,
+                        QuantKernel::NAME,
+                        QuantKernel::CODE,
+                        $kernel,
+                        grid,
+                        block,
+                        &mut [
+                            &nt as *const i32 as *mut c_void,
+                            &tk as *const i32 as *mut c_void,
+                            &ni as *const i32 as *mut c_void,
+                            (&ys_ptr) as *const *mut c_void as *mut c_void,
+                            (&sc_ptr) as *const *mut c_void as *mut c_void,
+                            (&out_ptr) as *const *mut c_void as *mut c_void,
+                        ],
+                    )?;
+                }
+                RocmStorage { slice: RocmStorageSlice::$variant(out), device: self.clone() }
+            }};
+        }
+        let out = match &ys.slice {
+            RocmStorageSlice::F16(_) => launch_combine!(F16, f16, "moe_combine_f16"),
+            RocmStorageSlice::BF16(_) => launch_combine!(BF16, bf16, "moe_combine_bf16"),
+            RocmStorageSlice::F32(_) => launch_combine!(F32, f32, "moe_combine_f32"),
+            other => crate::bail!("moe_combine: ys dtype {:?} unsupported", other.dtype()),
+        };
+        Ok(out)
+    }
 }

 unsafe fn launch_kernel(
@@ -139,6 +139,122 @@ extern "C" __global__ void quantize_q8_1_bf16(
    }
 }

+// f32 activation variant of quantize_q8 (symmetric): reads float instead of __half. MoE prefill
+// activations arrive in f32 (the model's compute dtype); quantizing f32 directly skips the f32->f16
+// to_dtype cast (a 16.7M-elem read+write per gate/up) the GEMM glue used to pay before this kernel.
+extern "C" __global__ void quantize_q8_f32(
+    const int M,
+    const int K,
+    const float* __restrict__ x,
+    int8_t* __restrict__ xq,
+    __half* __restrict__ xd
+) {
+    const int nblk = K >> 5;
+    const int wid = blockIdx.x * (blockDim.x >> 5) + (threadIdx.x >> 5);
+    if (wid >= M * nblk) {
+        return;
+    }
+    const int m = wid / nblk;
+    const int blk = wid % nblk;
+    const int lane = threadIdx.x & 31;
+    const size_t idx = (size_t)m * K + (size_t)blk * 32 + lane;
+    const float v = x[idx];
+    float a = fabsf(v);
+    #pragma unroll
+    for (int off = 16; off > 0; off >>= 1) {
+        a = fmaxf(a, __shfl_xor(a, off));
+    }
+    const float inv = (a > 0.0f) ? (127.0f / a) : 0.0f;
+    int q = (int)roundf(v * inv);
+    q = max(-127, min(127, q));
+    xq[idx] = (int8_t)q;
+    if (lane == 0) {
+        xd[(size_t)m * nblk + blk] = __float2half(a / 127.0f);
+    }
+}
+
+// f32 activation variant of quantize_q8_1 (asymmetric: also emits the per-block int8 sum xs).
+extern "C" __global__ void quantize_q8_1_f32(
+    const int M,
+    const int K,
+    const float* __restrict__ x,
+    int8_t* __restrict__ xq,
+    __half* __restrict__ xd,
+    int* __restrict__ xs
+) {
+    const int nblk = K >> 5;
+    const int wid = blockIdx.x * (blockDim.x >> 5) + (threadIdx.x >> 5);
+    if (wid >= M * nblk) {
+        return;
+    }
+    const int m = wid / nblk;
+    const int blk = wid % nblk;
+    const int lane = threadIdx.x & 31;
+    const size_t idx = (size_t)m * K + (size_t)blk * 32 + lane;
+    const float v = x[idx];
+    float a = fabsf(v);
+    #pragma unroll
+    for (int off = 16; off > 0; off >>= 1) {
+        a = fmaxf(a, __shfl_xor(a, off));
+    }
+    const float inv = (a > 0.0f) ? (127.0f / a) : 0.0f;
+    int q = (int)roundf(v * inv);
+    q = max(-127, min(127, q));
+    xq[idx] = (int8_t)q;
+    int qsum = q;
+    #pragma unroll
+    for (int off = 16; off > 0; off >>= 1) {
+        qsum += __shfl_xor(qsum, off);
+    }
+    if (lane == 0) {
+        xd[(size_t)m * nblk + blk] = __float2half(a / 127.0f);
+        xs[(size_t)m * nblk + blk] = qsum;
+    }
+}
+
+// Fused MoE expert-combine: out[i, j] = sum_{e<topk} scores[i, e] * ys[i, e, j].
+// Replaces the engine's broadcast_mul + sum(Minus2): that path cast ys -> f32 (a 16.7M-elem
+// cast), wrote a 16.7M-elem f32 product temp, then ran an 8-wide reduce at 2M blocks x 8 threads
+// (24/32 lanes idle, dst-strided). This reads ys ONCE (coalesced over j), accumulates topk in f32,
+// and writes out[i, j] -- one launch, one pass, no temp. NOT a qmmq_core variant (no matrix cores,
+// no weight decode) -- a standalone glue-reduction kernel.
+template <typename T>
+__device__ __forceinline__ void moe_combine_core(
+    const int ntok, const int topk, const int N,
+    const T* __restrict__ ys,         // [ntok, topk, N]
+    const float* __restrict__ scores, // [ntok, topk]
+    T* __restrict__ out               // [ntok, N]
+) {
+    const size_t total = (size_t)ntok * N;
+    const size_t stride = (size_t)gridDim.x * blockDim.x;
+    for (size_t gid = (size_t)blockIdx.x * blockDim.x + threadIdx.x; gid < total; gid += stride) {
+        const int i = (int)(gid / N);
+        const int j = (int)(gid - (size_t)i * N);
+        const float* __restrict__ sc = scores + (size_t)i * topk;
+        const T* __restrict__ yp = ys + ((size_t)i * topk) * N + j;
+        float acc = 0.0f;
+        for (int e = 0; e < topk; ++e) {
+            acc += sc[e] * (float)yp[(size_t)e * N];
+        }
+        out[gid] = (T)acc;
+    }
+}
+
+extern "C" __global__ void moe_combine_f16(
+    const int ntok, const int topk, const int N,
+    const __half* __restrict__ ys, const float* __restrict__ scores, __half* __restrict__ out
+) { moe_combine_core<__half>(ntok, topk, N, ys, scores, out); }
+
+extern "C" __global__ void moe_combine_bf16(
+    const int ntok, const int topk, const int N,
+    const hip_bfloat16* __restrict__ ys, const float* __restrict__ scores, hip_bfloat16* __restrict__ out
+) { moe_combine_core<hip_bfloat16>(ntok, topk, N, ys, scores, out); }
+
+extern "C" __global__ void moe_combine_f32(
+    const int ntok, const int topk, const int N,
+    const float* __restrict__ ys, const float* __restrict__ scores, float* __restrict__ out
+) { moe_combine_core<float>(ntok, topk, N, ys, scores, out); }
+
 // get_scale_min_k4 scale/min unpack -- DEFINED below (reused by the scalar Q4_K path + the unified
 // core); forward-declared here so the dp4a Q4_K kernel can call the SAME unpack (one source of
 // truth for the 6-bit scale layout).