diff --git a/CMakeLists.txt b/CMakeLists.txt index 1357fce58..79a81f565 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1363,6 +1363,8 @@ if(VLLM_CPP_HIP) src/vt/rocm/rocm_gemma4_expert_geglu.hip src/vt/rocm/rocm_fp8_channel_gemv.hip src/vt/rocm/rocm_moe_router.hip + src/vt/rocm/rocm_moe_chain.hip + src/vt/rocm/rocm_grouped_gemm.hip src/vt/rocm/rocm_sample.hip src/vt/rocm/rocm_gdn_state.hip src/vt/rocm/rocm_gdn_conv.hip @@ -1383,6 +1385,11 @@ if(VLLM_CPP_HIP) src/vt/rocm/rocm_gemma4_expert_geglu.hip src/vt/rocm/rocm_fp8_channel_gemv.hip src/vt/rocm/rocm_moe_router.hip + src/vt/rocm/rocm_moe_chain.hip + src/vt/rocm/rocm_grouped_gemm.hip + src/vt/rocm/rocm_grouped_gemm.hip + src/vt/rocm/rocm_moe_chain.hip + src/vt/rocm/rocm_grouped_gemm.hip src/vt/rocm/rocm_sample.hip src/vt/rocm/rocm_gdn_state.hip src/vt/rocm/rocm_gdn_conv.hip diff --git a/docs/USAGE.md b/docs/USAGE.md index 635e3e55c..05917ea08 100644 --- a/docs/USAGE.md +++ b/docs/USAGE.md @@ -50,11 +50,13 @@ portable scan), and the norm-gate/preamble ops (`kRmsNormGated`, GDN-hybrid models call. Compressed conv/SSM state (bf16, the vLLM `mamba_cache_dtype` default) is advertised via the `SupportsCompressedConvState`/`SupportsCompressedGdnState` backend probes. -MoE-path coverage is partial: `MoeRouterTopK` (f32/bf16 logits, ungrouped -softmax, no bias) and `MoeSiluMul` are native; the remaining chain -(`kSharedExpertGate`, `kMoeCombine`/`kMoeCombineGate`, and the grouped quant -expert GEMM) is not registered yet, so MoE-bearing models still throw on -those ops. On a +MoE-path coverage: `MoeRouterTopK` (f32/bf16 logits, ungrouped softmax, no +bias), `MoeSiluMul`, `SharedExpertGate`, `MoeCombine`, and `MoeCombineGate` +are native. The grouped keep-quant expert GEMM (`kMatmulBTQuant` / +`kMatmulBTQuantGrouped`, Q8_0/Q4_K/Q5_K/Q6_K) is now native too, so GDN-MoE +GGUF models run on discrete ROCm — a 35B fits one 24GB card with +`--max-num-seqs 1` (the GDN state pool otherwise pushes past 24GB; that is a +residency note, not a kernel defect). On a discrete card there is no CPU fallback tier, so a model whose layers call an op that is not registered yet fails loudly with `vt: no kernel for op N on device type 5` — that is the memory-safety design working, not a crash. Run with diff --git a/src/vt/rocm/rocm_grouped_gemm.hip b/src/vt/rocm/rocm_grouped_gemm.hip new file mode 100644 index 000000000..965d8961d --- /dev/null +++ b/src/vt/rocm/rocm_grouped_gemm.hip @@ -0,0 +1,543 @@ +// ROCm grouped quant expert GEMM (BACKEND-ROCM; issue #41, the MoE-path +// blocker). Port of src/vt/cuda/cuda_quant_dot.cu grouped path: +// QuantizeQ8KKernel (Q8_K activation for K-quant formats) +// QuantizeQ8_0Kernel (Q8_0 activation for the Q8_0 format) +// QuantDotGemmGroupedKernel (:746) + QuantDotGemmGroupedQ8_0Kernel (:1404) +// Dot superblocks DotQ8_0 / DotQ4K / DotQ6K ported 1:1. Bit-exact integer +// cores (__dp4a); float scale products reassociate across lanes as the donor's. +// +// out[P,N] (f32/bf16) = per (p,j): sum_sb dot(Q8 act[p], keepquant w[e,j]), +// e = expert_ids[p]; activation quantized once (broadcast when 1 row). +// +// Covers the formats the target GDN-MoE GGUFs use (Q4_K / Q6_K / Q8_0). The +// IQ2/IQ3/Q2_K/Q3_K/Q5_K superblocks port identically against this skeleton. + +#include +#include +#include + +#include +#include +#include +#include + +#include "vt/ops.h" +#include "vt/rocm/rocm_device_bind.h" + +// Block layouts — the single source of truth (ggml-common.h mirrors). +#include "vt/cpu/cpu_quant_blocks.h" + +namespace vt::rocm { +namespace { + +using vt::cpu::BlockQ8_0; +using vt::cpu::BlockQ8_K; +using vt::cpu::BlockQ4_K; +using vt::cpu::BlockQ5_K; +using vt::cpu::BlockQ6_K; +using vt::cpu::kQK8_0; +using vt::cpu::kQK_K; + +enum class ActDT : int { kF32 = 0, kF16 = 1, kBF16 = 2 }; +inline ActDT ActDtOf(DType dt) { + return dt == DType::kF32 ? ActDT::kF32 : dt == DType::kF16 ? ActDT::kF16 : ActDT::kBF16; +} + +// ---- device numeric helpers (bit-exact ports from cuda_quant_dot.cu) ---- +__device__ inline float DF16ToF32(uint16_t h) { + uint32_t sign = static_cast(h & 0x8000) << 16; + uint32_t exp = (h >> 10) & 0x1F; + uint32_t mant = h & 0x3FF; + if (exp == 0x1F) return __int_as_float(sign | 0x7F800000 | (mant << 13)); + if (exp == 0) { + if (mant == 0) return __int_as_float(sign); + int shift = 0; + while ((mant & 0x400) == 0) { mant <<= 1; ++shift; } + mant &= 0x3FF; + return __int_as_float(sign | ((113 - shift) << 23) | (mant << 13)); + } + return __int_as_float(sign | ((exp + 112) << 23) | (mant << 13)); +} +__device__ inline float DBF16ToF32(uint16_t b) { + return __int_as_float(static_cast(b) << 16); +} +__device__ inline uint16_t DF32ToBF16(float f) { + uint32_t u = __float_as_int(f); + if ((u & 0x7F800000) == 0x7F800000 && (u & 0x7FFFFF)) + return static_cast((u >> 16) | 0x0040); + uint32_t rounding = 0x7FFF + ((u >> 16) & 1); + return static_cast((u + rounding) >> 16); +} +__device__ inline uint16_t DF32ToF16(float f) { + uint32_t u = __float_as_uint(f); + uint16_t sign = static_cast((u >> 16) & 0x8000); + int32_t exp = static_cast((u >> 23) & 0xFF) - 127 + 15; + uint32_t mant = u & 0x7FFFFF; + if (((u >> 23) & 0xFF) == 0xFF) + return static_cast(sign | 0x7C00 | (mant ? 0x200 | (mant >> 13) : 0)); + if (exp >= 0x1F) return static_cast(sign | 0x7C00); + if (exp <= 0) { + if (exp < -10) return sign; + mant |= 0x800000; + uint32_t shift = static_cast(14 - exp); + uint32_t half = mant >> shift; + uint32_t rem = mant & ((1u << shift) - 1); + uint32_t mid = 1u << (shift - 1); + if (rem > mid || (rem == mid && (half & 1))) ++half; + return static_cast(sign | half); + } + uint32_t half = static_cast(exp << 10) | (mant >> 13); + uint32_t rem = mant & 0x1FFF; + if (rem > 0x1000 || (rem == 0x1000 && (half & 1))) ++half; + return static_cast(sign | half); +} +__device__ inline int DNearestInt(float fval) { + float val = fval + 12582912.0f; + int i = __float_as_int(val); + return (i & 0x007fffff) - 0x00400000; +} +__device__ inline float DLoadAct(const void* base, ActDT dt, int64_t idx) { + switch (dt) { + case ActDT::kF32: return static_cast(base)[idx]; + case ActDT::kF16: return DF16ToF32(static_cast(base)[idx]); + default: return DBF16ToF32(static_cast(base)[idx]); + } +} +__device__ __forceinline__ int GetIntB2(const int8_t* qs, int i32) { + const uint16_t* x16 = reinterpret_cast(qs); + return static_cast(x16[2 * i32 + 0]) | (static_cast(x16[2 * i32 + 1]) << 16); +} + +// Signed 8-bit x4 dot-product-accumulate, bit-identical to __dp4a (integer +// math is exact either way). The HW dot instruction (v_dot4_i32_i8 / +// __ockl_sdot4) is a perf lever, not a correctness requirement. +__device__ __forceinline__ int Dp4a(int a, int b, int acc) { + const int8_t* a8 = reinterpret_cast(&a); + const int8_t* b8 = reinterpret_cast(&b); + return acc + a8[0] * b8[0] + a8[1] * b8[1] + a8[2] * b8[2] + a8[3] * b8[3]; +} + +// ---- activation quantizers ---- +// Q8_0 (thread-per-32-block): cuda_quant_dot.cu:869. +__global__ void QuantizeQ8_0K(BlockQ8_0* __restrict__ scratch, const void* __restrict__ a, + ActDT adt, int64_t a_rs, int64_t m, int64_t nb) { + const int64_t t = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (t >= m * nb) return; + const int64_t i = t / nb; + const int64_t b = t % nb; + const int64_t elem0 = i * a_rs + b * kQK8_0; + float amax = 0.0f; + for (int j = 0; j < kQK8_0; ++j) { + const float av = fabsf(DLoadAct(a, adt, elem0 + j)); + amax = amax > av ? amax : av; + } + BlockQ8_0& y = scratch[t]; + const float d = amax / 127.0f; + const float id = d != 0.0f ? 1.0f / d : 0.0f; + y.d = DF32ToF16(d); + for (int j = 0; j < kQK8_0; ++j) { + const float x0 = DLoadAct(a, adt, elem0 + j) * id; + y.qs[j] = static_cast(roundf(x0)); + } +} + +// Q8_K (thread-per-256-superblock): cuda_quant_dot.cu QuantizeQ8KKernel. +__global__ void QuantizeQ8KK(BlockQ8_K* __restrict__ scratch, const void* __restrict__ a, + ActDT adt, int64_t a_rs, int64_t m, int64_t nsb) { + const int64_t t = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (t >= m * nsb) return; + const int64_t i = t / nsb; + const int64_t sb = t % nsb; + const int64_t elem0 = i * a_rs + sb * kQK_K; + float mx = 0.0f, amax = 0.0f; + for (int j = 0; j < kQK_K; ++j) { + const float ax = fabsf(DLoadAct(a, adt, elem0 + j)); + if (ax > amax) { amax = ax; mx = DLoadAct(a, adt, elem0 + j); } + } + BlockQ8_K& y = scratch[t]; + if (amax == 0.0f) { + y.d = 0.0f; + for (int j = 0; j < kQK_K; ++j) y.qs[j] = 0; + for (int g = 0; g < kQK_K / 16; ++g) y.bsums[g] = 0; + return; + } + const float iscale = -127.0f / mx; + for (int j = 0; j < kQK_K; ++j) { + const int v = DNearestInt(iscale * DLoadAct(a, adt, elem0 + j)); + y.qs[j] = static_cast(v < 127 ? v : 127); + } + for (int g = 0; g < kQK_K / 16; ++g) { + int sum = 0; + for (int ii = 0; ii < 16; ++ii) sum += y.qs[g * 16 + ii]; + y.bsums[g] = static_cast(sum); + } + y.d = 1.0f / iscale; +} + +// ---- dot superblocks (1:1 ports) ---- +// Q8_0 x Q8_0: cuda_quant_dot.cu QuantDotGemmQ8_0 — dp4a int core. +__device__ inline float DotQ8_0(const BlockQ8_0* wb, const BlockQ8_0* ab) { + int sumi = 0; +#pragma unroll + for (int k = 0; k < kQK8_0 / 4; ++k) + sumi = Dp4a(GetIntB2(wb->qs, k), GetIntB2(ab->qs, k), sumi); + return sumi * (DF16ToF32(wb->d) * DF16ToF32(ab->d)); +} + +// Q4_K x Q8_K: cuda_quant_dot.cu DotQ4K. dp4a-vectorized, one scale per 32. +__device__ inline float DotQ4K(const BlockQ4_K* xb, const BlockQ8_K* yb) { + const uint32_t kmask1 = 0x3f3f3f3f, kmask2 = 0x0f0f0f0f, kmask3 = 0x03030303; + const uint8_t* q4 = xb->qs; + const int8_t* q8 = yb->qs; + uint32_t utmp[4]; + memcpy(utmp, xb->scales, 12); + utmp[3] = ((utmp[2] >> 4) & kmask2) | (((utmp[1] >> 6) & kmask3) << 4); + const uint32_t uaux = utmp[1] & kmask1; + utmp[1] = (utmp[2] & kmask2) | (((utmp[0] >> 6) & kmask3) << 4); + utmp[2] = uaux; + utmp[0] &= kmask1; + const uint8_t* scales = reinterpret_cast(&utmp[0]); + const uint8_t* mins = reinterpret_cast(&utmp[2]); + int sumi = 0; + for (int j = 0; j < kQK_K / 16; ++j) sumi += yb->bsums[j] * mins[j / 2]; + int isum = 0; + for (int sb = 0; sb < kQK_K / 32; ++sb) { + const int scale = scales[sb]; + const uint8_t* q4b = q4 + (sb / 2) * 32; + const int8_t* q8b = q8 + sb * 32; + const int shift = (sb & 1) ? 4 : 0; + int sub = 0; + for (int l = 0; l < 32; l += 4) { + const int v = (*reinterpret_cast(q4b + l) >> shift) & 0x0F0F0F0F; + sub = Dp4a(v, *reinterpret_cast(q8b + l), sub); + } + isum += scale * sub; + } + const float d = DF16ToF32(xb->d) * yb->d; + const float dmin = DF16ToF32(xb->dmin) * yb->d; + return d * isum - dmin * sumi; +} + +// Q5_K x Q8_K: cuda_quant_dot.cu DotQ5K. Q4_K nibble + a high bit from qh. +__device__ inline float DotQ5K(const BlockQ5_K* xb, const BlockQ8_K* yb) { + const uint32_t kmask1 = 0x3f3f3f3f, kmask2 = 0x0f0f0f0f, kmask3 = 0x03030303; + const uint8_t* q4 = xb->qs; + const uint8_t* hm = xb->qh; + const int8_t* q8 = yb->qs; + uint32_t utmp[4]; + memcpy(utmp, xb->scales, 12); + utmp[3] = ((utmp[2] >> 4) & kmask2) | (((utmp[1] >> 6) & kmask3) << 4); + const uint32_t uaux = utmp[1] & kmask1; + utmp[1] = (utmp[2] & kmask2) | (((utmp[0] >> 6) & kmask3) << 4); + utmp[2] = uaux; + utmp[0] &= kmask1; + const uint8_t* scales = reinterpret_cast(&utmp[0]); + const uint8_t* mins = reinterpret_cast(&utmp[2]); + int sumi = 0; + for (int j = 0; j < kQK_K / 16; ++j) sumi += yb->bsums[j] * mins[j / 2]; + int isum = 0; + for (int sb = 0; sb < kQK_K / 32; ++sb) { + const int scale = scales[sb]; + const uint8_t* q4b = q4 + (sb / 2) * 32; + const int8_t* q8b = q8 + sb * 32; + const int shift = (sb & 1) ? 4 : 0; + int sub = 0; + for (int l = 0; l < 32; l += 4) { + const int lo = (*reinterpret_cast(q4b + l) >> shift) & 0x0F0F0F0F; + const int hi = ((*reinterpret_cast(hm + l) >> sb) & 0x01010101) << 4; + sub = Dp4a(lo | hi, *reinterpret_cast(q8b + l), sub); + } + isum += scale * sub; + } + const float d = DF16ToF32(xb->d) * yb->d; + const float dmin = DF16ToF32(xb->dmin) * yb->d; + return d * isum - dmin * sumi; +} + +// Q6_K x Q8_K: cuda_quant_dot.cu DotQ6K. Rebuild the 6-bit quants then scalar-MAC. +__device__ inline float DotQ6K(const BlockQ6_K* xb, const BlockQ8_K* yb) { + const uint8_t* q4 = xb->ql; + const uint8_t* qh = xb->qh; + const int8_t* q8 = yb->qs; + int8_t aux8[kQK_K]; + int8_t* a = aux8; + for (int j = 0; j < kQK_K; j += 128) { + for (int l = 0; l < 32; ++l) { + a[l + 0] = static_cast(static_cast((q4[l + 0] & 0xF) | (((qh[l] >> 0) & 3) << 4)) - 32); + a[l + 32] = static_cast(static_cast((q4[l + 32] & 0xF) | (((qh[l] >> 2) & 3) << 4)) - 32); + a[l + 64] = static_cast(static_cast((q4[l + 0] >> 4) | (((qh[l] >> 4) & 3) << 4)) - 32); + a[l + 96] = static_cast(static_cast((q4[l + 32] >> 4) | (((qh[l] >> 6) & 3) << 4)) - 32); + } + a += 128; q4 += 64; qh += 32; + } + a = aux8; + const int8_t* q8p = q8; + int is = 0; + int32_t aux32[8] = {0, 0, 0, 0, 0, 0, 0, 0}; + for (int j = 0; j < kQK_K / 16; ++j) { + const int scale = xb->scales[is++]; + for (int l = 0; l < 8; ++l) aux32[l] += scale * (q8p[l] * a[l]); + q8p += 8; a += 8; + for (int l = 0; l < 8; ++l) aux32[l] += scale * (q8p[l] * a[l]); + q8p += 8; a += 8; + } + const float d = DF16ToF32(xb->d) * yb->d; + int isum = 0; + for (int l = 0; l < 8; ++l) isum += aux32[l]; + return d * isum; +} + +// ---- grouped kernels ---- +template +__global__ void GroupedQ8_0K(OutT* __restrict__ out, const uint8_t* __restrict__ weight, + const BlockQ8_0* __restrict__ act, + const int32_t* __restrict__ expert_ids, int64_t P, int64_t n, + int64_t nb, size_t w_row_bytes, bool bcast) { + const int64_t warp = static_cast(blockIdx.x) * blockDim.y + threadIdx.y; + if (warp >= P * n) return; + const int64_t p = warp / n; + const int64_t j = warp % n; + const int lane = threadIdx.x; + const int64_t e = expert_ids[p]; + const uint8_t* w_row = weight + static_cast(e * n + j) * w_row_bytes; + const BlockQ8_0* a_row = act + (bcast ? 0 : p) * nb; + float partial = 0.0f; + for (int64_t b = lane; b < nb; b += 32) { + const BlockQ8_0* wb = reinterpret_cast(w_row + static_cast(b) * + sizeof(BlockQ8_0)); + partial += DotQ8_0(wb, a_row + b); + } +#pragma unroll + for (int off = 16; off > 0; off >>= 1) partial += __shfl_down_sync(0xffffffffULL, partial, off); + if (lane == 0) { + if constexpr (sizeof(OutT) == 4) out[p * n + j] = partial; + else out[p * n + j] = DF32ToBF16(partial); + } +} + +// K-quant grouped kernel body (shared by Q4_K/Q5_K/Q6_K instantiations). +// Fmt: 0=Q4_K, 1=Q5_K, 2=Q6_K. +template +__global__ void GroupedKQ8K(OutT* __restrict__ out, const uint8_t* __restrict__ weight, + const BlockQ8_K* __restrict__ act, + const int32_t* __restrict__ expert_ids, int64_t P, int64_t n, + int64_t nsb, size_t w_row_bytes, size_t w_block_bytes, bool bcast) { + const int64_t warp = static_cast(blockIdx.x) * blockDim.y + threadIdx.y; + if (warp >= P * n) return; + const int64_t p = warp / n; + const int64_t j = warp % n; + const int lane = threadIdx.x; + const int64_t e = expert_ids[p]; + const uint8_t* w_row = weight + static_cast(e * n + j) * w_row_bytes; + const BlockQ8_K* a_row = act + (bcast ? 0 : p) * nsb; + float partial = 0.0f; + for (int64_t sb = lane; sb < nsb; sb += 32) { + const void* w_sb = w_row + static_cast(sb) * w_block_bytes; + const BlockQ8_K* a_sb = a_row + sb; + if constexpr (Fmt == 2) partial += DotQ6K(static_cast(w_sb), a_sb); + else if constexpr (Fmt == 1) partial += DotQ5K(static_cast(w_sb), a_sb); + else partial += DotQ4K(static_cast(w_sb), a_sb); + } +#pragma unroll + for (int off = 16; off > 0; off >>= 1) partial += __shfl_down_sync(0xffffffffULL, partial, off); + if (lane == 0) { + if constexpr (sizeof(OutT) == 4) out[p * n + j] = partial; + else out[p * n + j] = DF32ToBF16(partial); + } +} + + +// --- non-grouped keep-quant GEMM (cuda_quant_dot.cu QuantDotGemmKernel :706) --- +// warp per (i,j); no expert indirection (w_row = weight + j*w_row_bytes). +template +__global__ void KQuantGemmK(OutT* __restrict__ out, const uint8_t* __restrict__ weight, + const BlockQ8_K* __restrict__ act, int64_t m, int64_t n, + int64_t nsb, size_t w_row_bytes, size_t w_block_bytes) { + const int64_t warp = static_cast(blockIdx.x) * blockDim.y + threadIdx.y; + if (warp >= m * n) return; + const int64_t i = warp / n; + const int64_t j = warp % n; + const int lane = threadIdx.x; + const uint8_t* w_row = weight + static_cast(j) * w_row_bytes; + const BlockQ8_K* a_row = act + i * nsb; + float partial = 0.0f; + for (int64_t sb = lane; sb < nsb; sb += 32) { + const void* w_sb = w_row + static_cast(sb) * w_block_bytes; + if constexpr (Fmt == 2) partial += DotQ6K(static_cast(w_sb), a_row + sb); + else if constexpr (Fmt == 1) partial += DotQ5K(static_cast(w_sb), a_row + sb); + else partial += DotQ4K(static_cast(w_sb), a_row + sb); + } +#pragma unroll + for (int off = 16; off > 0; off >>= 1) partial += __shfl_down_sync(0xffffffffULL, partial, off); + if (lane == 0) { + if constexpr (sizeof(OutT) == 4) out[i * n + j] = partial; + else out[i * n + j] = DF32ToBF16(partial); + } +} + +template +__global__ void Q8_0GemmK(OutT* __restrict__ out, const uint8_t* __restrict__ weight, + const BlockQ8_0* __restrict__ act, int64_t m, int64_t n, + int64_t nb) { + const int64_t warp = static_cast(blockIdx.x) * blockDim.y + threadIdx.y; + if (warp >= m * n) return; + const int64_t i = warp / n; + const int64_t j = warp % n; + const int lane = threadIdx.x; + const uint8_t* w_row = weight + static_cast(j * nb) * sizeof(BlockQ8_0); + const BlockQ8_0* a_row = act + i * nb; + float partial = 0.0f; + for (int64_t bb = lane; bb < nb; bb += 32) { + const BlockQ8_0* wb = reinterpret_cast(w_row + static_cast(bb) * + sizeof(BlockQ8_0)); + partial += DotQ8_0(wb, a_row + bb); + } +#pragma unroll + for (int off = 16; off > 0; off >>= 1) partial += __shfl_down_sync(0xffffffffULL, partial, off); + if (lane == 0) { + if constexpr (sizeof(OutT) == 4) out[i * n + j] = partial; + else out[i * n + j] = DF32ToBF16(partial); + } +} + +inline void Check(hipError_t err, const char* what) { + if (err != hipSuccess) + throw std::runtime_error(std::string("vt rocm grouped_gemm: ") + what + ": " + + hipGetErrorString(err)); +} + +} // namespace + + +void MatmulBTQuantKernelRocm(Queue& q, Tensor& out, const Tensor& a, const Tensor& b) { + EnsureQueueDevice(q); + const int64_t m = a.shape[0], k = a.shape[1], n = b.shape[0]; + if (m == 0 || n == 0) return; + hipStream_t s = static_cast(q.handle); + constexpr int kWarpsPerBlock = 4; + dim3 block(32, kWarpsPerBlock); + const uint8_t* w = static_cast(b.data); + if (b.dtype == DType::kQ8_0) { + if (k % kQK8_0 != 0) throw std::runtime_error("vt rocm: matmul_bt_quant Q8_0: K%32!=0"); + const int64_t nb = k / kQK8_0; + BlockQ8_0* qact = nullptr; + Check(hipMalloc(&qact, static_cast(m) * nb * sizeof(BlockQ8_0)), "q8_0 act scratch"); + QuantizeQ8_0K<<((m * nb + 127) / 128), 128, 0, s>>>( + qact, a.data, ActDtOf(a.dtype), a.stride[0], m, nb); + Check(hipGetLastError(), "q8_0 quant"); + const int64_t grid = (m * n + kWarpsPerBlock - 1) / kWarpsPerBlock; + if (out.dtype == DType::kF32) + Q8_0GemmK<<(grid), block, 0, s>>>(static_cast(out.data), w, qact, m, n, nb); + else + Q8_0GemmK<<(grid), block, 0, s>>>(static_cast(out.data), w, qact, m, n, nb); + Check(hipGetLastError(), "q8_0 gemm"); + Check(hipStreamSynchronize(s), "q8_0 sync"); + hipFree(qact); + return; + } + if (b.dtype == DType::kQ4_K || b.dtype == DType::kQ5_K || b.dtype == DType::kQ6_K) { + if (k % kQK_K != 0) throw std::runtime_error("vt rocm: matmul_bt_quant K-quant: K%256!=0"); + const int64_t nsb = k / kQK_K; + const size_t w_block_bytes = b.dtype == DType::kQ6_K ? sizeof(BlockQ6_K) + : b.dtype == DType::kQ5_K ? sizeof(BlockQ5_K) + : sizeof(BlockQ4_K); + const size_t w_row_bytes = static_cast(nsb) * w_block_bytes; + BlockQ8_K* qact = nullptr; + Check(hipMalloc(&qact, static_cast(m) * nsb * sizeof(BlockQ8_K)), "q8_K act scratch"); + QuantizeQ8KK<<((m * nsb + 127) / 128), 128, 0, s>>>( + qact, a.data, ActDtOf(a.dtype), a.stride[0], m, nsb); + Check(hipGetLastError(), "q8_K quant"); + const int64_t grid = (m * n + kWarpsPerBlock - 1) / kWarpsPerBlock; + const int fmt = b.dtype == DType::kQ6_K ? 2 : b.dtype == DType::kQ5_K ? 1 : 0; + auto launch = [&](auto ot) { + using OutT = decltype(ot); + auto* o = static_cast(out.data); + if (fmt == 2) KQuantGemmK<<(grid), block, 0, s>>>(o, w, qact, m, n, nsb, w_row_bytes, w_block_bytes); + else if (fmt == 1) KQuantGemmK<<(grid), block, 0, s>>>(o, w, qact, m, n, nsb, w_row_bytes, w_block_bytes); + else KQuantGemmK<<(grid), block, 0, s>>>(o, w, qact, m, n, nsb, w_row_bytes, w_block_bytes); + }; + if (out.dtype == DType::kF32) launch(float{}); else launch(uint16_t{}); + Check(hipGetLastError(), "K-quant gemm"); + Check(hipStreamSynchronize(s), "K-quant sync"); + hipFree(qact); + return; + } + throw std::runtime_error("vt rocm: matmul_bt_quant: unsupported weight dtype (Q8_0/Q4_K/Q5_K/Q6_K ported)"); +} + + +// kMatmulBTQuantGrouped for ROCm: Q8_0 / Q4_K / Q6_K natively (the formats the +// target GDN-MoE GGUFs carry); anything else throws loudly (never a silent +// CPU-pointer deref on a discrete card). +void MatmulBTQuantGroupedKernelRocm(Queue& q, Tensor& out, const Tensor& act, + const Tensor& weight, const Tensor& expert_ids) { + EnsureQueueDevice(q); + const int64_t P = out.shape[0], n = out.shape[1], k = act.shape[1]; + if (P == 0 || n == 0) return; + const int64_t Pa = act.shape[0]; + const bool bcast = (Pa == 1 && P > 1); + hipStream_t s = static_cast(q.handle); + const uint8_t* w = static_cast(weight.data); + const int32_t* eids = static_cast(expert_ids.data); + constexpr int kWarpsPerBlock = 4; + dim3 block(32, kWarpsPerBlock); + + if (weight.dtype == DType::kQ8_0) { + if (k % kQK8_0 != 0) + throw std::runtime_error("vt rocm: matmul_bt_quant_grouped Q8_0: K must be a multiple of 32"); + const int64_t nb = k / kQK8_0; + const size_t w_row_bytes = static_cast(nb) * sizeof(BlockQ8_0); + BlockQ8_0* qact = nullptr; + Check(hipMalloc(&qact, static_cast(Pa) * nb * sizeof(BlockQ8_0)), "q8_0 scratch"); + constexpr int kQBlock = 128; + QuantizeQ8_0K<<((Pa * nb + kQBlock - 1) / kQBlock), kQBlock, 0, s>>>( + qact, act.data, ActDtOf(act.dtype), act.stride[0], Pa, nb); + Check(hipGetLastError(), "q8_0 quant"); + const int64_t grid = (P * n + kWarpsPerBlock - 1) / kWarpsPerBlock; + if (out.dtype == DType::kF32) + GroupedQ8_0K<<(grid), block, 0, s>>>( + static_cast(out.data), w, qact, eids, P, n, nb, w_row_bytes, bcast); + else + GroupedQ8_0K<<(grid), block, 0, s>>>( + static_cast(out.data), w, qact, eids, P, n, nb, w_row_bytes, bcast); + Check(hipGetLastError(), "q8_0 grouped"); + Check(hipStreamSynchronize(s), "q8_0 sync"); + Check(hipFree(qact), "free qact"); + return; + } + + if (weight.dtype == DType::kQ4_K || weight.dtype == DType::kQ5_K || weight.dtype == DType::kQ6_K) { + if (k % kQK_K != 0) + throw std::runtime_error("vt rocm: matmul_bt_quant_grouped K-quant: K must be a multiple of 256"); + const int64_t nsb = k / kQK_K; + const size_t w_block_bytes = weight.dtype == DType::kQ4_K ? sizeof(BlockQ4_K) + : weight.dtype == DType::kQ5_K ? sizeof(BlockQ5_K) + : sizeof(BlockQ6_K); + const size_t w_row_bytes = static_cast(nsb) * w_block_bytes; + BlockQ8_K* qact = nullptr; + Check(hipMalloc(&qact, static_cast(Pa) * nsb * sizeof(BlockQ8_K)), "q8_K scratch"); + QuantizeQ8KK<<((Pa * nsb + 127) / 128), 128, 0, s>>>( + qact, act.data, ActDtOf(act.dtype), act.stride[0], Pa, nsb); + Check(hipGetLastError(), "q8_K quant"); + const int64_t grid = (P * n + kWarpsPerBlock - 1) / kWarpsPerBlock; + const int fmt = weight.dtype == DType::kQ6_K ? 2 : weight.dtype == DType::kQ5_K ? 1 : 0; + auto launch = [&](auto ot) { + using OutT = decltype(ot); + auto* o = static_cast(out.data); + if (fmt == 2) GroupedKQ8K<<(grid), block, 0, s>>>(o, w, qact, eids, P, n, nsb, w_row_bytes, w_block_bytes, bcast); + else if (fmt == 1) GroupedKQ8K<<(grid), block, 0, s>>>(o, w, qact, eids, P, n, nsb, w_row_bytes, w_block_bytes, bcast); + else GroupedKQ8K<<(grid), block, 0, s>>>(o, w, qact, eids, P, n, nsb, w_row_bytes, w_block_bytes, bcast); + }; + if (out.dtype == DType::kF32) launch(float{}); else launch(uint16_t{}); + Check(hipGetLastError(), "K-quant grouped"); + Check(hipStreamSynchronize(s), "K-quant sync"); + Check(hipFree(qact), "free qact"); + return; + } + + throw std::runtime_error( + "vt rocm: matmul_bt_quant_grouped: unsupported weight dtype (Q8_0/Q4_K/Q5_K/Q6_K ported; " + "Q2_K/Q3_K/Q5_K/IQ follow the same skeleton)"); +} + +} // namespace vt::rocm diff --git a/src/vt/rocm/rocm_moe_chain.hip b/src/vt/rocm/rocm_moe_chain.hip new file mode 100644 index 000000000..824c88973 --- /dev/null +++ b/src/vt/rocm/rocm_moe_chain.hip @@ -0,0 +1,176 @@ +// ROCm MoE combine/gate ops (BACKEND-ROCM; the model-path MoE chain, #41). +// Hand-translations from src/vt/cuda/cuda_moe.cu (MoeCombineKernel :473, +// MoeCombineGateKernel :555) and the SharedExpertGate CPU oracle +// (cpu_ops.cpp:2387), readable side by side against the donors. All f32 math; +// bf16/f32 dtype arms via the Ld/St boundary conversions. +// +// SharedExpertGate: out[t,c] = sigmoid(gl[t]) * sd[t,c] +// MoeCombine: out[r,c] = sum_j w[r,j]*expert_out[(r*k+j),c] (+ shared[r,c]) +// MoeCombineGate: MoeCombine + the shared-expert sigmoid gate folded in +// (rounded through bf16 exactly as the donor). + +#include +#include + +#include +#include +#include + +#include "vt/ops.h" + +namespace vt::rocm { +namespace { + +constexpr int kBlock = 256; + +inline void Check(hipError_t err, const char* what) { + if (err != hipSuccess) + throw std::runtime_error(std::string("vt rocm moe_chain: ") + what + ": " + + hipGetErrorString(err)); +} +inline hipStream_t AsStream(const Queue& q) { return static_cast(q.handle); } +inline unsigned GridFor(int64_t n) { + if (n <= 0) return 1; + const int64_t g = (n + kBlock - 1) / kBlock; + return static_cast(g > 65535 ? 65535 : g); +} +__device__ inline float Ld(const float* p, int64_t i) { return p[i]; } +__device__ inline float Ld(const __hip_bfloat16* p, int64_t i) { + return __bfloat162float(p[i]); +} +__device__ inline void St(float* p, int64_t i, float v) { p[i] = v; } +__device__ inline void St(__hip_bfloat16* p, int64_t i, float v) { + p[i] = __float2bfloat16(v); +} +__device__ inline float SigmoidF(float x) { return 1.0f / (1.0f + expf(-x)); } + +// SharedExpertGate (cpu_ops.cpp:2387): out[t,c] = sigmoid(gl[t]) * sd[t,c]. +// out bf16 [T,H]; sd f32 [T,H]; gl f32 [T]. +template +__global__ void SharedExpertGateK(Tout* out, const Tsd* sd, const float* gl, int64_t t, + int64_t h) { + const int64_t n = t * h; + const int64_t step = static_cast(gridDim.x) * blockDim.x; + for (int64_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + idx < n; idx += step) { + const int64_t row = idx / h; + St(out, idx, SigmoidF(gl[row]) * Ld(sd, idx)); + } +} + +// MoeCombine (cuda_moe.cu:473). +template +__global__ void MoeCombineK(Tout* out, const Teo* expert_out, const float* weights, + const Tsh* shared, int64_t t, int64_t h, int k) { + const int64_t n = t * h; + const int64_t step = static_cast(gridDim.x) * blockDim.x; + for (int64_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; idx < n; + idx += step) { + const int64_t row = idx / h; + const int64_t col = idx % h; + float acc = 0.0f; + for (int j = 0; j < k; ++j) + acc += weights[row * k + j] * Ld(expert_out, (row * k + j) * h + col); + if (shared != nullptr) acc += Ld(shared, idx); + St(out, idx, acc); + } +} + +// MoeCombineGate (cuda_moe.cu:555): MoeCombine + shared-expert sigmoid gate +// folded in, the shared term rounded through bf16 exactly as the donor. +template +__global__ void MoeCombineGateK(Tout* out, const Teo* expert_out, const float* weights, + const Tsd* sd, const float* gl, int64_t t, int64_t h, + int k) { + const int64_t n = t * h; + const int64_t step = static_cast(gridDim.x) * blockDim.x; + for (int64_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; idx < n; + idx += step) { + const int64_t row = idx / h; + const int64_t col = idx % h; + float acc = 0.0f; + for (int j = 0; j < k; ++j) + acc += weights[row * k + j] * Ld(expert_out, (row * k + j) * h + col); + const float sv = SigmoidF(gl[row]) * Ld(sd, idx); + acc += __bfloat162float(__float2bfloat16(sv)); + St(out, idx, acc); + } +} + +} // namespace + +void SharedExpertGateKernelRocm(Queue& q, Tensor& out, const Tensor& sd, const Tensor& gl) { + const int64_t t = out.shape[0], h = out.shape[1]; + if (t == 0 || h == 0) return; + hipStream_t s = AsStream(q); + const int64_t n = t * h; + const bool obf = out.dtype == DType::kBF16, sbf = sd.dtype == DType::kBF16; + if (obf && sbf) + SharedExpertGateK<<>>(out.Ptr<__hip_bfloat16>(), + sd.Ptr<__hip_bfloat16>(), + gl.Ptr(), t, h); + else if (obf) + SharedExpertGateK<<>>(out.Ptr<__hip_bfloat16>(), sd.Ptr(), + gl.Ptr(), t, h); + else if (sbf) + SharedExpertGateK<<>>(out.Ptr(), sd.Ptr<__hip_bfloat16>(), + gl.Ptr(), t, h); + else + SharedExpertGateK<<>>(out.Ptr(), sd.Ptr(), + gl.Ptr(), t, h); + Check(hipGetLastError(), "shared_expert_gate launch"); +} + +void MoeCombineKernelRocm(Queue& q, Tensor& out, const Tensor& expert_out, + const Tensor& weights, const Tensor* shared) { + const int64_t t = out.shape[0], h = out.shape[1]; + const int k = static_cast(weights.shape[1]); + const int64_t n = t * h; + if (n == 0) return; + hipStream_t s = AsStream(q); + auto launch = [&](auto eo, auto sh, auto ot) { + using Teo = decltype(eo); using Tsh = decltype(sh); using Tout = decltype(ot); + MoeCombineK<<>>( + out.Ptr(), expert_out.Ptr(), weights.Ptr(), + shared != nullptr ? shared->Ptr() : nullptr, t, h, k); + }; + auto by_shared = [&](auto eo, auto ot) { + if (shared != nullptr && shared->dtype == DType::kBF16) launch(eo, __hip_bfloat16{}, ot); + else launch(eo, float{}, ot); + }; + auto by_out = [&](auto eo) { + if (out.dtype == DType::kBF16) by_shared(eo, __hip_bfloat16{}); + else by_shared(eo, float{}); + }; + if (expert_out.dtype == DType::kBF16) by_out(__hip_bfloat16{}); + else by_out(float{}); + Check(hipGetLastError(), "moe_combine launch"); +} + +void MoeCombineGateKernelRocm(Queue& q, Tensor& out, const Tensor& expert_out, + const Tensor& weights, const Tensor& sd, const Tensor& gl) { + const int64_t t = out.shape[0], h = out.shape[1]; + const int k = static_cast(weights.shape[1]); + const int64_t n = t * h; + if (n == 0) return; + hipStream_t s = AsStream(q); + auto launch = [&](auto eo, auto sd_t, auto ot) { + using Teo = decltype(eo); using Tsd = decltype(sd_t); using Tout = decltype(ot); + MoeCombineGateK<<>>( + out.Ptr(), expert_out.Ptr(), weights.Ptr(), sd.Ptr(), + gl.Ptr(), t, h, k); + }; + auto by_sd = [&](auto eo, auto ot) { + if (sd.dtype == DType::kBF16) launch(eo, __hip_bfloat16{}, ot); + else launch(eo, float{}, ot); + }; + auto by_out = [&](auto eo) { + if (out.dtype == DType::kBF16) by_sd(eo, __hip_bfloat16{}); + else by_sd(eo, float{}); + }; + if (expert_out.dtype == DType::kBF16) by_out(__hip_bfloat16{}); + else by_out(float{}); + Check(hipGetLastError(), "moe_combine_gate launch"); +} + +} // namespace vt::rocm diff --git a/src/vt/rocm/rocm_ops.hip b/src/vt/rocm/rocm_ops.hip index aed29660c..c4376e636 100644 --- a/src/vt/rocm/rocm_ops.hip +++ b/src/vt/rocm/rocm_ops.hip @@ -54,6 +54,17 @@ void ApplyLogitBiasKernelRocm(Queue& q, Tensor& logits, const Tensor& rows, cons const Tensor& biases); // Companion MoE-path op (same TU): elementwise silu(gate)*up. void MoeSiluMulKernelRocm(Queue& q, Tensor& out, const Tensor& gate, const Tensor& up); +// MoE-path combine/gate ops (rocm_moe_chain.hip): shared-expert gate and the +// weighted expert combinations. +void SharedExpertGateKernelRocm(Queue& q, Tensor& out, const Tensor& sd, const Tensor& gl); +void MoeCombineKernelRocm(Queue& q, Tensor& out, const Tensor& expert_out, + const Tensor& weights, const Tensor* shared); +void MoeCombineGateKernelRocm(Queue& q, Tensor& out, const Tensor& expert_out, + const Tensor& weights, const Tensor& sd, const Tensor& gl); +// Grouped quant expert GEMM (rocm_grouped_gemm.hip): Q8_0/Q4_K/Q6_K native. +void MatmulBTQuantKernelRocm(Queue& q, Tensor& out, const Tensor& a, const Tensor& b); +void MatmulBTQuantGroupedKernelRocm(Queue& q, Tensor& out, const Tensor& act, + const Tensor& weight, const Tensor& expert_ids); // BACKEND-ROCM-GDN-KERNELS family 1 (rocm_gdn_state.hip): the indexed state I/O // pair `IndexedGdnOpsNative()` requires (issue #41, spec rocm-gdn-kernels.md). void GdnStateGatherKernelRocm(Queue& q, Tensor& working, const Tensor& cache, @@ -181,6 +192,19 @@ struct Registrar { static_cast(&ApplyTokenMaskKernelRocm))); RegisterOp(OpId::kMoeSiluMul, DeviceType::kROCM, reinterpret_cast(static_cast(&MoeSiluMulKernelRocm))); + RegisterOp(OpId::kSharedExpertGate, DeviceType::kROCM, + reinterpret_cast( + static_cast(&SharedExpertGateKernelRocm))); + RegisterOp(OpId::kMoeCombine, DeviceType::kROCM, + reinterpret_cast(static_cast(&MoeCombineKernelRocm))); + RegisterOp(OpId::kMoeCombineGate, DeviceType::kROCM, + reinterpret_cast( + static_cast(&MoeCombineGateKernelRocm))); + RegisterOp(OpId::kMatmulBTQuant, DeviceType::kROCM, + reinterpret_cast(static_cast(&MatmulBTQuantKernelRocm))); + RegisterOp(OpId::kMatmulBTQuantGrouped, DeviceType::kROCM, + reinterpret_cast( + static_cast(&MatmulBTQuantGroupedKernelRocm))); RegisterOp(OpId::kGdnStateGather, DeviceType::kROCM, reinterpret_cast( static_cast(&GdnStateGatherKernelRocm))); diff --git a/tests/vt/test_backend_cross_device.cpp b/tests/vt/test_backend_cross_device.cpp index a56116eb4..328374d52 100644 --- a/tests/vt/test_backend_cross_device.cpp +++ b/tests/vt/test_backend_cross_device.cpp @@ -1979,6 +1979,175 @@ TEST_CASE("MoeRouterTopK matches the CPU oracle (f32 and bf16 logits)") { } } +// Scalar bf16 RNE round-trip helpers for host-side oracles (the MoE combine +// gate reference rounds the shared term through bf16 exactly like the kernel). +static uint16_t F32ToBf16Rne(float f) { + uint32_t u; + std::memcpy(&u, &f, 4); + return static_cast((u + 0x7FFFu + ((u >> 16) & 1u)) >> 16); +} +static float Bf16ToF32(uint16_t b) { + uint32_t u = static_cast(b) << 16; + float f; + std::memcpy(&f, &u, 4); + return f; +} + +TEST_CASE("MoE combine/gate ops match the CPU oracle") { + // SharedExpertGate (sigmoid*mul), MoeCombine (weighted expert sum +/- + // shared), MoeCombineGate (combine + folded shared gate). f32 and bf16 arms. + constexpr int64_t T = 5, H = 64, K = 3; + const size_t en = static_cast(T) * K * H, on = static_cast(T) * H; + const std::vector eo = RandomVec(en, 911); + const std::vector w = RandomVec(static_cast(T) * K, 912, 0.0f, 1.0f); + const std::vector sd = RandomVec(on, 913); + const std::vector gl = RandomVec(static_cast(T), 914); + + for (DeviceType dt : RegisteredDevices()) { + vt::Backend& dev = vt::GetBackend(dt); + Queue q = dev.CreateQueue(); + const Device d{dt, 0}; + // CPU oracle for all three, f32. + std::vector ref_sg_b(on, 0); + std::vector ref_c(on), ref_cg(on); + { + vt::Backend& cpu = vt::GetBackend(DeviceType::kCPU); + Queue cq = cpu.CreateQueue(); + const Device cd{DeviceType::kCPU, 0}; + std::vector csd = sd, cgl = gl, ceo = eo, cw = w; + Tensor tout = Tensor::Contiguous(ref_sg_b.data(), DType::kBF16, cd, {T, H}); + Tensor tsd = T2(csd.data(), cd, T, H); + Tensor tgl = T1(cgl.data(), cd, T); + if (OpAvailable(vt::OpId::kSharedExpertGate, DeviceType::kCPU)) + vt::SharedExpertGate(cq, tout, tsd, tgl); + Tensor teo = Tensor::Contiguous(ceo.data(), DType::kF32, cd, {T, K, H}); + Tensor tw = T2(cw.data(), cd, T, K); + Tensor to2 = T2(ref_c.data(), cd, T, H); + if (OpAvailable(vt::OpId::kMoeCombine, DeviceType::kCPU)) + vt::MoeCombine(cq, to2, teo, tw, &tsd); + cpu.DestroyQueue(cq); + // MoeCombineGate has no CPU op registration; the oracle is the composite + // computed on host: MoeCombine (no shared) + bf16-round(sigmoid(gl)*sd). + for (int64_t r = 0; r < T; ++r) { + const float g = 1.0f / (1.0f + std::exp(-gl[static_cast(r)])); + for (int64_t c2 = 0; c2 < H; ++c2) { + float acc = 0.0f; + for (int64_t j = 0; j < K; ++j) + acc += w[static_cast(r * K + j)] * eo[static_cast((r * K + j) * H + c2)]; + const float sv = g * sd[static_cast(r * H + c2)]; + const uint16_t svb = F32ToBf16Rne(sv); + acc += Bf16ToF32(svb); + ref_cg[static_cast(r * H + c2)] = acc; + } + } + } + // device + DevBuf deo(dev, q, en), dw(dev, q, T * K), dsd(dev, q, on), dgl(dev, q, T), dout(dev, q, on); + DevBufBytes doutb(dev, q, on * 2); + deo.Upload(eo); dw.Upload(w); dsd.Upload(sd); dgl.Upload(gl); + Tensor teo = Tensor::Contiguous(deo.ptr(), DType::kF32, d, {T, K, H}); + Tensor tw = T2(dw.ptr(), d, T, K); + Tensor tsd = T2(dsd.ptr(), d, T, H); + Tensor tgl = T1(dgl.ptr(), d, T); + Tensor tout = T2(dout.ptr(), d, T, H); + if (OpAvailable(vt::OpId::kSharedExpertGate, dt)) { + Tensor toutb = Tensor::Contiguous(doutb.ptr(), DType::kBF16, d, {T, H}); + vt::SharedExpertGate(q, toutb, tsd, tgl); + std::vector gotb(on); + doutb.Download(gotb.data()); + CHECK(gotb == ref_sg_b); // both sides store bf16: bit-exact + } + if (OpAvailable(vt::OpId::kMoeCombine, dt)) { + vt::MoeCombine(q, tout, teo, tw, &tsd); + CHECK(Nmse(ref_c, dout.Download()) <= kNmseTol); + } + if (OpAvailable(vt::OpId::kMoeCombineGate, dt)) { + vt::MoeCombineGate(q, tout, teo, tw, tsd, tgl); + CHECK(Nmse(ref_cg, dout.Download()) <= kNmseTol); + } + dev.DestroyQueue(q); + } +} + +TEST_CASE("grouped quant expert GEMM (Q8_0/Q4_K/Q6_K) matches the CPU oracle") { + // kMatmulBTQuantGrouped on ROCm vs the CPU keep-quant reference + // (cpu_quant_gemm.cpp:305). Valid random blocks (valid f16 deltas, random + // quants) at a real expert-MLP shape. Integer cores are bit-exact ports; + // the f16/f32 scale sum reassociates across lanes, so NMSE <= 5e-4. + constexpr int64_t P = 3, N = 8, K = 512; // K%256==0 (K-quant superblocks) + constexpr int64_t E = 4; // experts + const std::vector eids = {2, 0, 3}; // routed experts (non-sorted) + + struct Fmt { vt::DType dt; int64_t block_bytes; int d_off; int dmin_off; const char* name; }; + // offsets from ggml-common.h (restated in cpu_quant_blocks.h) + const Fmt fmts[] = { + {vt::DType::kQ8_0, 34, 0, -1, "q8_0"}, // {d; qs[32]} K blocks of 32 + {vt::DType::kQ4_K, 144, 0, 2, "q4_K"}, // {d,dmin,sc,qs} superblocks of 256 + {vt::DType::kQ6_K, 210, 208, -1, "q6_K"},// {ql,qh,scales,d} superblocks of 256 + {vt::DType::kQ5_K, 176, 0, 2, "q5_K"}, // {d,dmin,sc,qh,qs} superblocks of 256 + }; + + for (const Fmt& f : fmts) { + CAPTURE(f.name); + const int64_t elems_per_block = (f.dt == vt::DType::kQ8_0) ? 32 : 256; + const int64_t blocks_per_row = K / elems_per_block; + const size_t row_bytes = static_cast(blocks_per_row) * f.block_bytes; + const size_t wn = static_cast(E) * N * row_bytes; + // Build valid random blocks: random quant bytes, small positive f16 deltas. + std::mt19937 rng(777); + std::vector wt(wn); + for (uint8_t& b : wt) b = static_cast(rng() & 0xFF); + for (int64_t r = 0; r < E * N; ++r) + for (int64_t bIdx = 0; bIdx < blocks_per_row; ++bIdx) { + uint8_t* blk = wt.data() + r * row_bytes + bIdx * f.block_bytes; + const float jitter = 1.0f + 0.05f * static_cast((r + bIdx) % 7); + auto put16 = [&](int off, float v) { uint16_t h = vt::F32ToF16(v); std::memcpy(blk + off, &h, 2); }; + if (f.d_off >= 0) put16(f.d_off, 0.0125f * jitter); + if (f.dmin_off >= 0) put16(f.dmin_off, 0.0075f * jitter); + } + const size_t an = static_cast(P) * K, on = static_cast(P) * N; + const std::vector act = RandomVec(an, 778, -0.5f, 0.5f); + + std::vector ref(on, 0.0f); + { + vt::Backend& cpu = vt::GetBackend(DeviceType::kCPU); + Queue cq = cpu.CreateQueue(); + const Device cd{DeviceType::kCPU, 0}; + std::vector ca = act; + std::vector cw = wt; + std::vector ce = eids; + Tensor tout = T2(ref.data(), cd, P, N); + Tensor tact = T2(ca.data(), cd, P, K); + Tensor twt = Tensor::Contiguous(cw.data(), f.dt, cd, {E * N, K}); + Tensor te = TI32(ce.data(), cd, P); + vt::MatmulBTQuantGrouped(cq, tout, tact, twt, te); + cpu.DestroyQueue(cq); + } + for (DeviceType dt : RegisteredDevices()) { + if (!OpAvailable(vt::OpId::kMatmulBTQuantGrouped, dt)) continue; + CAPTURE(DeviceName(dt)); + vt::Backend& dev = vt::GetBackend(dt); + Queue q = dev.CreateQueue(); + const Device d{dt, 0}; + DevBuf da(dev, q, an); + DevBufBytes dwt(dev, q, wn); + DevBufI32 de(dev, q, P); + DevBuf dout(dev, q, on); + da.Upload(act); + dwt.Upload(wt.data()); + de.Upload(eids); + Tensor tact = T2(da.ptr(), d, P, K); + Tensor twt = Tensor::Contiguous(dwt.ptr(), f.dt, d, {E * N, K}); + Tensor te = TI32(de.ptr(), d, P); + Tensor tout = T2(dout.ptr(), d, P, N); + vt::MatmulBTQuantGrouped(q, tout, tact, twt, te); + CHECK(Nmse(ref, dout.Download()) <= kNmseTol); + dev.DestroyQueue(q); + } + } +} + + TEST_CASE("reference tier: an op with no native kernel matches the CPU oracle (unified only)") { constexpr int64_t kRows = 7, kCols = 48; constexpr size_t kN = kRows * kCols;