From 71094cb1567cd8dec6814e23e55d2a0cd508812c Mon Sep 17 00:00:00 2001 From: PsyMan47 Date: Fri, 6 Nov 2026 11:06:19 +0100 Subject: [PATCH] cuda: run on Turing (sm_75) Turing has no cp.async, no __reduce_max/min_sync and no m16n8k16 f16 or m16n8k16/k32 int8 mma, so the kernels do not compile there. The f16 GEMM becomes two m16n8k8 steps on the same accumulator (a0/a1 and a2/a3 are the two k-halves of the k16 A fragment, b0/b1 the two halves of B). cp.async copies through registers; a predicated-off cp.async zero-fills its bytes, and the prefill attention relies on that for rows past the last key, so the fallback writes zeros too. The two __reduce_*_sync in topk_warp are warp shuffles. There is no int8 GEMM on Turing either: gemmi/gemmis are not built and gemm_raw takes the f16 path. The new Backend::int8_gemm lets the kernel tests build the reference the backend actually uses (Vulkan keeps its old behavior). sm_75 has 64 KB of shared memory a block. attn_prefill at head_dim 256 uses a 16-key tile instead of 32, and attn_decode_gqa at head_dim 256 skips its cp.async KV prefetch. The dynamic shared memory opt-in now follows the device limit minus the kernel's statics instead of a fixed 96 KB. --- src/backend.rs | 3 +++ src/cuda/kernels.cu | 61 ++++++++++++++++++++++++++++++++++++++++----- src/cuda/mod.rs | 28 ++++++++++++++++----- src/vulkan/mod.rs | 1 + tests/kernels.rs | 10 ++++---- 5 files changed, 86 insertions(+), 17 deletions(-) diff --git a/src/backend.rs b/src/backend.rs index bb188493..90424bab 100644 --- a/src/backend.rs +++ b/src/backend.rs @@ -142,6 +142,9 @@ pub trait Backend: Send { /// y = silu(y) * result (SwiGLU with the gate already in y), 2 adds to y; `| SAME_X` as below. `n_x`: rows of x /// the tiles read. fn gemm(&self, ty: GgmlType, rows: usize, cols: usize, tiles: u64, n_tiles: usize, x: u64, x_stride: usize, n_x: usize, xrow: u64, y: u64, y_stride: usize, act: i32) -> anyhow::Result<()>; + /// Whether `gemm` and `gemm_dense` run on int8 tensor cores for `ty` (the alternative is f16 products), so tests + /// and callers build their reference the way this backend will. + fn int8_gemm(&self, ty: GgmlType) -> bool; /// y[m] (op) W x[m] for positions m < n, x and y rows `x_stride` / `y_stride` floats apart. fn gemm_dense(&self, w: &GpuTensor, n: usize, x: u64, x_stride: usize, y: u64, y_stride: usize, act: i32) -> anyhow::Result<()>; /// y = W x (all rows of `w`). diff --git a/src/cuda/kernels.cu b/src/cuda/kernels.cu index 4871e142..8f0b5119 100644 --- a/src/cuda/kernels.cu +++ b/src/cuda/kernels.cu @@ -23,6 +23,15 @@ __device__ __forceinline__ float warp_max(float v) { for (int o = 16; o > 0; o >>= 1) v = fmaxf(v, __shfl_xor_sync(0xffffffff, v, o)); return v; } +// Unsigned integer forms of warp_max (__reduce_max_sync/__reduce_min_sync need sm_80; the shuffles run on sm_70+). +__device__ __forceinline__ unsigned warp_max_u(unsigned v) { + for (int o = 16; o > 0; o >>= 1) v = max(v, __shfl_xor_sync(0xffffffff, v, o)); + return v; +} +__device__ __forceinline__ unsigned warp_min_u(unsigned v) { + for (int o = 16; o > 0; o >>= 1) v = min(v, __shfl_xor_sync(0xffffffff, v, o)); + return v; +} // Sum over the whole block; every thread gets the result. blockDim.x must be a multiple of 32. __device__ float block_sum(float v) { __shared__ float s[32]; @@ -813,8 +822,17 @@ __device__ __forceinline__ void ldsm4t(u32* r, const u16* p) { : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3]) : "r"(a)); } __device__ __forceinline__ void mma16816(float* c, u32 a0, u32 a1, u32 a2, u32 a3, u32 b0, u32 b1) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 + // Turing (sm_75) has no m16n8k16 f16 mma: the same product is two m16n8k8 steps, k = 0..7 and k = 8..15 + // (the a0/a1, a2/a3 and b0, b1 fragments are those half-K fragments of the k16 layout). + asm volatile("mma.sync.aligned.m16n8k8.row.col.f32.f16.f16.f32 {%0,%1,%2,%3}, {%4,%5}, {%6}, {%0,%1,%2,%3};" + : "+f"(c[0]), "+f"(c[1]), "+f"(c[2]), "+f"(c[3]) : "r"(a0), "r"(a1), "r"(b0)); + asm volatile("mma.sync.aligned.m16n8k8.row.col.f32.f16.f16.f32 {%0,%1,%2,%3}, {%4,%5}, {%6}, {%0,%1,%2,%3};" + : "+f"(c[0]), "+f"(c[1]), "+f"(c[2]), "+f"(c[3]) : "r"(a2), "r"(a3), "r"(b1)); +#else asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 {%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%0,%1,%2,%3};" : "+f"(c[0]), "+f"(c[1]), "+f"(c[2]), "+f"(c[3]) : "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b0), "r"(b1)); +#endif } __device__ __forceinline__ u32 h2pack(float a, float b) { return (u32)f2h(a) | ((u32)f2h(b) << 16); } @@ -973,6 +991,8 @@ extern "C" __global__ void quant_q8(const float* x, int x_stride, u8* out, int c } } +// int8 tensor-core mma (sm_80+). On Turing the int8 path is not built: gemm_raw uses the f16 GEMM there. +#if !defined(__CUDA_ARCH__) || __CUDA_ARCH__ >= 800 __device__ __forceinline__ void mma_s8_16816(int* c, u32 a0, u32 a1, u32 b0) { asm volatile("mma.sync.aligned.m16n8k16.row.col.s32.s8.s8.s32 {%0,%1,%2,%3}, {%4,%5}, {%6}, {%0,%1,%2,%3};" : "+r"(c[0]), "+r"(c[1]), "+r"(c[2]), "+r"(c[3]) : "r"(a0), "r"(a1), "r"(b0)); @@ -981,22 +1001,42 @@ __device__ __forceinline__ void mma_s8_16832(int* c, u32 a0, u32 a1, u32 a2, u32 asm volatile("mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32 {%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%0,%1,%2,%3};" : "+r"(c[0]), "+r"(c[1]), "+r"(c[2]), "+r"(c[3]) : "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b0), "r"(b1)); } +#endif // int -> float for |c| < 2^22 without the quarter-rate I2F: 2^23 + 2^22 + c has c in its low mantissa bits. __device__ __forceinline__ float i2f_small(int c) { return __int_as_float(0x4B400000 + c) - 12582912.0f; } // Weight types whose scale (and min) is constant over 32 weights: one 32-deep mma per scale. template __host__ __device__ constexpr bool scale32() { return T == 42 || T == 2 || T == 3 || T == 6 || T == 7 || T == 8 || T == 12 || T == 13 || T == 16 || T == 19 || T == 18 || T == 21 || T == 23; } __device__ __forceinline__ void cp_async16(void* dst, const void* src, bool ok) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 + // No cp.async before sm_80: copy through registers. A predicated-off cp.async still zero-fills its 16 bytes, + // and the kernels rely on that for rows past the last one, so the fallback writes zeros too. + if (ok) *(uint4*)dst = *(const uint4*)src; else *(uint4*)dst = make_uint4(0, 0, 0, 0); +#else unsigned d = (unsigned)__cvta_generic_to_shared(dst); asm volatile("cp.async.cg.shared.global [%0], [%1], 16, %2;" :: "r"(d), "l"(src), "r"(ok ? 16 : 0)); +#endif } __device__ __forceinline__ void cp_async8(void* dst, const void* src, bool ok) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 + if (ok) *(uint2*)dst = *(const uint2*)src; else *(uint2*)dst = make_uint2(0, 0); +#else unsigned d = (unsigned)__cvta_generic_to_shared(dst); asm volatile("cp.async.ca.shared.global [%0], [%1], 8, %2;" :: "r"(d), "l"(src), "r"(ok ? 8 : 0)); +#endif +} +__device__ __forceinline__ void cp_async_commit() { +#if !defined(__CUDA_ARCH__) || __CUDA_ARCH__ >= 800 + asm volatile("cp.async.commit_group;"); +#endif +} +template __device__ __forceinline__ void cp_async_wait() { +#if !defined(__CUDA_ARCH__) || __CUDA_ARCH__ >= 800 + asm volatile("cp.async.wait_group %0;" :: "n"(N)); +#endif } -__device__ __forceinline__ void cp_async_commit() { asm volatile("cp.async.commit_group;"); } -template __device__ __forceinline__ void cp_async_wait() { asm volatile("cp.async.wait_group %0;" :: "n"(N)); } // SM: as in gemm_q (gemmis_*, the tiles of at most 32 positions, 16 x 16 outputs per warp). +#if !defined(__CUDA_ARCH__) || __CUDA_ARCH__ >= 800 template __device__ __forceinline__ void gemmi_q(const GemmTile* tiles, const u8* wd, int nd, const u8* __restrict__ x, const int* xrow, float* y, int y_stride, int rows, int n64, int act) { @@ -1128,6 +1168,7 @@ __device__ __forceinline__ void gemmi_q(const GemmTile* tiles, const u8* wd, int gemmi_q(tiles, wd, nd, x, xrow, y, y_stride, rows, n64, act); \ } ITYPES(DEF_GEMMI) +#endif // ---- elementwise / norm ---- @@ -1781,7 +1822,7 @@ __device__ void attn_decode_t(const float* q, const u16* kc, const u16* vc, floa float scale, int gq_min, int pz, int nt, const int* sel, int sst #define ATT_TOK int t = blockIdx.x % nt; q += (size_t)t * nh * hd; part += (size_t)t * pz; step += 2 * t; if (sel) sel += (size_t)t * sst; #define ATT_RUN(G) attn_decode_t(q, kc, vc, part, step, sel, nh, nkv, scale, blockIdx.x / nt) -#define ATT_RUN_CP(G) attn_decode_t(q, kc, vc, part, step, sel, nh, nkv, scale, blockIdx.x / nt) +#define ATT_RUN_CP(G) attn_decode_t= 800))>(q, kc, vc, part, step, sel, nh, nkv, scale, blockIdx.x / nt) #define DEF_ATT(HD) \ extern "C" __global__ void attn_decode_##HD(ATT_ARGS) { \ constexpr int PER = HD / 32; ATT_TOK \ @@ -1861,8 +1902,8 @@ __device__ void topk_warp(const float* p, int n, int k, float extra_scale, int l unsigned i = lane + 32 * r, key = __float_as_uint(v[r]) + 1; if (i < n && !(taken >> r & 1) && key > bk) { bk = key; bi = i; } } - unsigned mk = __reduce_max_sync(0xffffffff, bk); - bi = __reduce_min_sync(0xffffffff, bk == mk ? bi : 0xffffffffu); + unsigned mk = warp_max_u(bk); + bi = warp_min_u(bk == mk ? bi : 0xffffffffu); float bv = __uint_as_float(mk - 1); if ((bi & 31) == lane) taken |= 1u << (bi >> 5); if (lane == j) { id = bi; mine = bv; } @@ -1965,7 +2006,14 @@ extern "C" __global__ void moe_combine_seq(float* x, const float* ye, const floa // gate at +HD per head; out [T, nh * HD] = softmax(q K^T * scale) V * sigmoid(gate). Shared: Q, K and V tiles, all // row-major; mma fragments come from ldmatrix (.trans for V). #define FA_BQ 64 -#define FA_BK 32 +#define FA_BK_DEFAULT 32 +// Turing's 64 KB shared memory cannot hold the HD = 256 tile at 32 keys either: 16 fit (the kernel keeps the same +// 64 queries a block and four warps). +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 +#define FA_BK_FOR(HD) ((HD) == 256 ? 16 : FA_BK_DEFAULT) +#else +#define FA_BK_FOR(HD) FA_BK_DEFAULT +#endif // QSA mask (see qsa_select): whether query row t (mw words a row) sees key position `key`. The mask has rows for the // tile padding past T too (their output is dropped). __device__ __forceinline__ bool qsa_bit(const unsigned* qm, int mw, int t, int key) { @@ -1975,6 +2023,7 @@ template __device__ __forceinline__ void attn_prefill_t(const float* q, const u16* kc, const u16* vc, const float* qg, float* out, const int* step, int T, int nh, int nkv, float scale, const unsigned* qm, int mw) { constexpr int PQ = HD + 8; + constexpr int FA_BK = FA_BK_FOR(HD); extern __shared__ __align__(16) u16 fa_sm[]; u16 *qs = fa_sm, *ks = qs + FA_BQ * PQ, *vs = ks + FA_BK * PQ; // The last query tiles see the most keys: start them first so the short ones fill the tail. diff --git a/src/cuda/mod.rs b/src/cuda/mod.rs index c9b4b295..692227f6 100644 --- a/src/cuda/mod.rs +++ b/src/cuda/mod.rs @@ -38,6 +38,8 @@ pub struct Gpu { q8: std::sync::Mutex<(CudaSlice, Option<[u64; 4]>)>, /// The kernels were compiled for a q8 KV cache (`--kv q8`). pub kv_q8: bool, + /// The GPU has int8 tensor-core mma (sm_80+). Turing runs the prefill GEMM on the f16 path instead. + pub int8_mma: bool, /// `argmax_rows` partials, AMAX_PARTS per row. am: CudaSlice<[i32; 2]>, /// Recorded graphs by `Backend::graph` key. @@ -179,19 +181,27 @@ impl Gpu { eprintln!("cuda device {device}: {}", ctx.name().unwrap_or_default()); let stream = ctx.new_stream()?; let (maj, min) = ctx.compute_capability()?; + let int8_mma = maj >= 8; + let max_optin = ctx.attribute(cudarc::driver::sys::CUdevice_attribute::CU_DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_BLOCK_OPTIN) + .unwrap_or(48 << 10).max(48 << 10) as u32; let arch: &'static str = Box::leak(format!("sm_{maj}{min}").into_boxed_str()); let module = ctx.load_module(ptx(&source(kv_q8), arch)?).context("load cubin")?; let mut fns = HashMap::new(); let names = KERNELS.iter().map(|s| s.to_string()) .chain(QTYPES.iter().flat_map(|&t| ["dequant", "gemv", "gemvn", "moe_gemv", "gemm", "gemms"].map(|p| kname(p, t)))) .chain(QTYPES.iter().flat_map(|&t| ["gemv", "gemvn", "moe_gemv"].map(|p| kname(p, t) + "_p"))) - .chain(QTYPES.iter().filter(|&&t| !Self::float_ty(t)).flat_map(|&t| [kname("gemmi", t), kname("gemmis", t)])) + .chain(QTYPES.iter().filter(|&&t| int8_mma && !Self::float_ty(t)).flat_map(|&t| [kname("gemmi", t), kname("gemmis", t)])) .chain([64, 128, 256, 512].iter().flat_map(|h| [format!("attn_decode_{h}"), format!("attn_decode_gqa_{h}")])); for n in names { let f = module.load_function(&n).with_context(|| format!("kernel {n}"))?; // gemvn caches nt quantized inputs: 4 x 17408 columns (Qwen3.8-27B ffn_down) need 77 KB. if n.starts_with("attn_prefill") || n.starts_with("attn_decode_gqa") || n.starts_with("gemvn") { - f.set_attribute(cudarc::driver::sys::CUfunction_attribute::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, 96 << 10)?; + // Opt in to as much dynamic shared memory as the device has left after the kernel's statics (96 KB + // on Ampere and newer, 64 KB on Turing). + let static_smem = f.shared_size_bytes().unwrap_or(0).max(0) as u32; + let smem = max_optin.saturating_sub(static_smem).min(96 << 10) as i32; + f.set_attribute(cudarc::driver::sys::CUfunction_attribute::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, smem) + .with_context(|| format!("shared memory attribute for {n}"))?; } fns.insert(n, f); } @@ -199,7 +209,7 @@ impl Gpu { let am = stream.alloc_zeros(64 * AMAX_PARTS)?; let ring = std::sync::Mutex::new((unsafe { ctx.alloc_pinned::(RING)? }, 0, VecDeque::new())); let done = ctx.new_event(None)?; - Ok(Gpu { ctx, side: stream.fork()?, stream, nb: AtomicUsize::new(usize::MAX), fns, q8, kv_q8, am, graphs: Default::default(), ring, done, lag: Default::default() }) + Ok(Gpu { ctx, side: stream.fork()?, stream, nb: AtomicUsize::new(usize::MAX), fns, q8, kv_q8, int8_mma, am, graphs: Default::default(), ring, done, lag: Default::default() }) } fn f(&self, name: &str) -> &CudaFunction { &self.fns[name] } @@ -245,7 +255,7 @@ impl Gpu { let (same, act) = (act & SAME_X != 0, act & !SAME_X); let mut q8 = self.q8.lock().unwrap(); // The int8 path loads q8 rows 16 bytes at a time: their stride (cols * 11 / 8) is aligned only for cols % 128 == 0. - if !Self::mmq(ty) || cols % 128 != 0 { + if !self.int8_mma || !Self::mmq(ty) || cols % 128 != 0 { if !same { q8.1 = None; } launch!(self, &kname(if sm { "gemms" } else { "gemm" }, ty), cfg(grid, 2 * GEMM_BM as u32, 0), tiles, wd, nd as i32, x, x_stride as i32, xrow, y, y_stride as i32, rows as i32, (cols / 64) as i32, act); @@ -332,7 +342,9 @@ impl Gpu { // gq_min i32::MAX or 0 fixes the regime, so the kernel that would exit at once is not launched. if gq_min != 0 { launch!(self, &format!("attn_decode_{hd}"), cfg(((nh * nt) as u32, ATT_SPLITS_1 as u32), 256, (8 * hd * 4) as u32), q, kc, vc, part, step, nh as i32, nkv as i32, hd as i32, scale, gq_min, pz, nt as i32, sel, sst as i32); } // attn_decode_gqa with the f16 cache at head_dim 128 or 256 also holds 2 chunks of K and V for its cp.async. - let cp = if !self.kv_q8 && (hd == 128 || hd == 256) { 2 * 2 * 32 * hd * 2 } else { 0 }; + // Turing (sm_75) has 64 KB of shared memory a block: the head_dim 256 chunks do not fit, so it loads them + // through registers instead (kernels.cu ATT_RUN_CP). + let cp = if !self.kv_q8 && (hd == 128 || (hd == 256 && self.int8_mma)) { 2 * 2 * 32 * hd * 2 } else { 0 }; if gq_min != i32::MAX { launch!(self, &format!("attn_decode_gqa_{hd}"), cfg(((nh / g * nt) as u32, ATT_SPLITS as u32), 256, (8 * hd * 4 + cp) as u32), q, kc, vc, part, step, nh as i32, nkv as i32, hd as i32, scale, gq_min, pz, nt as i32, sel, sst as i32); } launch!(self, "attn_merge", LaunchConfig { grid_dim: (nh as u32, 1, nt as u32), block_dim: (256, 1, 1), shared_mem_bytes: 0 }, part, qg, out, hd as i32, ATT_SPLITS as i32, step, ATT_SPLITS_1 as i32, gq_min, xq); Ok(()) @@ -475,6 +487,8 @@ impl Backend for Gpu { self.gemm_raw(w.ty, w.rows, w.cols, 0, n.div_ceil(GEMM_BM), w.ptr, n, x, x_stride, n, 0, y, y_stride, act) } + fn int8_gemm(&self, ty: GgmlType) -> bool { self.int8_mma && Self::mmq(ty) } + fn gemv(&self, w: &GpuTensor, x: u64, y: u64) -> anyhow::Result<()> { let (ncols, kn) = Self::gemv_cols("gemv", w.ty, w.cols)?; let (c, rpw) = Self::gemv_cfg(w.ty, w.rows, w.cols, 1); @@ -617,7 +631,9 @@ impl Backend for Gpu { fn attn_prefill_m(&self, q: u64, kc: u64, vc: u64, qg: u64, out: u64, step: u64, t: usize, nh: usize, nkv: usize, hd: usize, scale: f32, qm: u64, mw: usize) -> anyhow::Result<()> { if !matches!(hd, 64 | 128 | 256) { bail!("attention prefill: head_dim {hd} unsupported (64, 128 or 256)"); } - let smem = (2 * (64 + 32 + 32) * (hd + 8)) as u32; + // Turing keeps 16 keys a tile for head_dim 256 (see FA_BK_FOR in kernels.cu). + let bk = if !self.int8_mma && hd == 256 { 16 } else { 32 }; + let smem = (2 * (64 + bk + bk) * (hd + 8)) as u32; launch!(self, &format!("attn_prefill_{hd}"), cfg((t.div_ceil(64) as u32, nh as u32), 128, smem), q, kc, vc, qg, out, step, t as i32, nh as i32, nkv as i32, scale, qm, mw as i32); Ok(()) diff --git a/src/vulkan/mod.rs b/src/vulkan/mod.rs index fa46a8a5..f39d8770 100644 --- a/src/vulkan/mod.rs +++ b/src/vulkan/mod.rs @@ -1040,6 +1040,7 @@ impl Backend for Vk { fn gemm_dense(&self, w: &GpuTensor, n: usize, x: u64, x_stride: usize, y: u64, y_stride: usize, act: i32) -> anyhow::Result<()> { self.gemm_raw(w.ty, w.rows, w.cols, 0, n.div_ceil(128), w.ptr, n, x, x_stride, n, 0, y, y_stride, act) } + fn int8_gemm(&self, ty: GgmlType) -> bool { crate::cuda::Gpu::mmq(ty) } fn gemv(&self, w: &GpuTensor, x: u64, y: u64) -> anyhow::Result<()> { self.gemvq(w, x, 0, 0, y, 0, 1) } fn gemvn(&self, w: &GpuTensor, x: u64, x_stride: usize, y: u64, y_stride: usize, nt: usize) -> anyhow::Result<()> { self.gemvq(w, x, x_stride, 0, y, y_stride, nt) diff --git a/tests/kernels.rs b/tests/kernels.rs index 8a37292c..83ec528d 100644 --- a/tests/kernels.rs +++ b/tests/kernels.rs @@ -2,7 +2,7 @@ mod common; use common::{down, up}; use rusting_strata::backend::{ATT_GQA_MIN, ATT_SPLITS, GEMM_BM, GpuTensor}; -use rusting_strata::{cuda::Gpu, gguf::Gguf, quant::GgmlType, weights::{Tensor, Weights}}; +use rusting_strata::{gguf::Gguf, quant::GgmlType, weights::{Tensor, Weights}}; fn xs(n: usize) -> Vec { (0..n).map(|i| ((i as f32 * 0.7548777).fract() - 0.5) * 2.0).collect() @@ -249,7 +249,7 @@ fn gemv_bandwidth() { fn f16r(v: f32) -> f32 { half::f16::from_f32(v).to_f32() } /// Tensor-core GEMM over a tile list matches per-row dot products, for every type, with a gathered X and partial -/// tiles. Types on the int8 path (`Gpu::mmq`) see X as int8 per 32 (as the GEMV path does), the others fp16 inputs. +/// tiles. Types on the int8 path (`Backend::int8_gemm`) see X as int8 per 32 (as the GEMV path does), the others fp16 inputs. #[test] fn gemm_matches_reference_for_every_type() { let Some(p) = common::model_path() else { return }; @@ -280,7 +280,7 @@ fn gemm_matches_reference_for_every_type() { t.row(r, &mut wr); for (i, &xr) in xrow.iter().enumerate() { let xv = &x[xr as usize * t.cols..][..t.cols]; - want[i * rows + r] = if Gpu::mmq(t.ty) { + want[i * rows + r] = if gpu.int8_gemm(t.ty) { wr.iter().zip(as_seen(t.ty, xv)).map(|(a, b)| a * b).sum() } else { wr.iter().zip(xv).map(|(a, b)| f16r(*a) * f16r(*b)).sum() @@ -291,7 +291,7 @@ fn gemm_matches_reference_for_every_type() { eprintln!("{:8} {name}: gemm rel err {e:.2e}", ty.name()); assert!(e <= 1e-4, "{} gemm ({name}): rel err {e}", ty.name()); // Vulkan without cooperative matrices sums the f16 products in another order than mma.sync. - if bk == "vulkan" && !Gpu::mmq(t.ty) && std::env::var_os("STRATA_VK_NO_COOP").is_some() { continue; } + if bk == "vulkan" && !gpu.int8_gemm(t.ty) && std::env::var_os("STRATA_VK_NO_COOP").is_some() { continue; } let f = first.entry(name.clone()).or_insert_with(|| got.clone()); let bits = got.iter().zip(f.iter()).filter(|(a, b)| a.to_bits() != b.to_bits()).count(); assert_eq!(bits, 0, "{} gemm ({name}) on {bk}: {bits} values differ from the first backend", ty.name()); @@ -1112,7 +1112,7 @@ fn gemv_partial_superblocks_and_q2_0() { for r in 0..rows { t.row(r, &mut wr); for (i, xv) in xm.chunks(cols).enumerate() { - want[i * rows + r] = if Gpu::mmq(ty) && cols % 128 == 0 { wr.iter().zip(as_seen(ty, xv)).map(|(a, b)| a * b).sum() } + want[i * rows + r] = if gpu.int8_gemm(ty) && cols % 128 == 0 { wr.iter().zip(as_seen(ty, xv)).map(|(a, b)| a * b).sum() } else { wr.iter().zip(xv).map(|(a, b)| f16r(*a) * f16r(*b)).sum() }; } }