diff --git a/src/backend.rs b/src/backend.rs index cf06895f..7de333e2 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<()>; fn prequant_gemm_x(&self, _w: &GpuTensor, _n: usize, _x: u64, _x_stride: usize) -> anyhow::Result<()> { Ok(()) } diff --git a/src/cuda/kernels.cu b/src/cuda/kernels.cu index d5325603..f46873de 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]; @@ -829,8 +838,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); } @@ -989,6 +1007,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)); @@ -997,6 +1017,7 @@ __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; } // The mma with every accumulator starting at I2F_BIAS: d = I2F_BIAS + sum, so i2f_small's integer add comes free. @@ -1012,17 +1033,36 @@ __device__ __forceinline__ void mma_s8_16832_b(int* d, u32 a0, u32 a1, u32 a2, u // 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) { @@ -1166,6 +1206,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 ---- @@ -1866,7 +1907,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 \ @@ -1955,8 +1996,13 @@ __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; } } +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 + unsigned mk = warp_max_u(bk); + bi = warp_min_u(bk == mk ? bi : 0xffffffffu); +#else unsigned mk = __reduce_max_sync(0xffffffff, bk); bi = __reduce_min_sync(0xffffffff, bk == mk ? bi : 0xffffffffu); +#endif float bv = __uint_as_float(mk - 1); if ((bi & 31) == lane) taken |= 1u << (bi >> 5); if (lane == j) { id = bi; mine = bv; } @@ -2059,7 +2105,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) { @@ -2069,6 +2122,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 b067e214..e24652eb 100644 --- a/src/cuda/mod.rs +++ b/src/cuda/mod.rs @@ -38,6 +38,10 @@ 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, + /// The device's opt-in shared memory per block (96 KB on Ampere, 64 KB on Turing). + max_smem: u32, /// `argmax_rows` partials, AMAX_PARTS per row. am: CudaSlice<[i32; 2]>, /// Recorded graphs by `Backend::graph` key. @@ -181,20 +185,28 @@ 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}")])) .chain(["attn_decode_gqa4_128".to_string()]); 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); } @@ -202,7 +214,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(), fr: Default::default() }) + Ok(Gpu { ctx, side: stream.fork()?, stream, nb: AtomicUsize::new(usize::MAX), fns, q8, kv_q8, int8_mma, max_smem: max_optin, am, graphs: Default::default(), ring, done, lag: Default::default(), fr: Default::default() }) } fn f(&self, name: &str) -> &CudaFunction { &self.fns[name] } @@ -272,8 +284,8 @@ impl Gpu { let reuse = same && q8.1 == key; // The int8 path loads q8 rows 16 bytes at a time: their stride (cols * 11 / 8) is aligned only for cols % 128 == 0. // A SAME_X gemm inside a batch cannot quantize x (the quantize and the gemm would land on different streams): - // when the cached q8 rows are not x's, use the f16 kernel. - let int8 = Self::mmq(ty) && cols.is_multiple_of(128); + // when the cached q8 rows are not x's, use the f16 kernel. Turing has no int8 tensor-core mma. + let int8 = self.int8_mma && Self::mmq(ty) && cols.is_multiple_of(128); if !int8 || (same && !reuse && self.nb.load(Relaxed) < BESIDE) { if !same { q8.1 = None; } launch!(self, &kname(if sm { "gemms" } else { "gemm" }, ty), cfg(grid, 2 * GEMM_BM as u32, 0), @@ -298,7 +310,7 @@ impl Gpu { /// Quantize `x` for the int8 GEMM, unless the cache already holds the same rows; `None` when `ty` or `cols` /// selects the f16 path. Shared by `gemm_raw` and `prequant_gemm_x`. fn q8_x(&self, q8: &mut (CudaSlice, Option<[u64; 4]>), ty: GgmlType, cols: usize, n_x: usize, x: u64, x_stride: usize) -> anyhow::Result> { - if !Self::mmq(ty) || !cols.is_multiple_of(128) { return Ok(None); } + if !self.int8_mma || !Self::mmq(ty) || !cols.is_multiple_of(128) { return Ok(None); } let key = Some([x, x_stride as u64, n_x as u64, cols as u64]); if q8.1 == key { return Ok(Some(self.ptr(&q8.0))); } let need = n_x * (cols + cols / 8 + cols / 4); @@ -346,8 +358,19 @@ impl Gpu { if !(1..=4).contains(&nt) { bail!("gemvn: {nt} inputs (1 to 4)"); } let (ncols, k1) = Self::gemv_cols("gemv", w.ty, w.cols)?; let (mut c, rpw) = Self::gemv_cfg(w.ty, w.rows, w.cols, 1); - c.shared_mem_bytes = if xq == 0 { c.shared_mem_bytes * nt as u32 } else { 0 }; + let per = c.shared_mem_bytes; + c.shared_mem_bytes = if xq == 0 { per * nt as u32 } else { 0 }; c.grid_dim.2 = z; + // Turing has 64 KB of opt-in shared memory, less than gemvn's nt cached x rows need for wide matrices + // (4 x 17408 columns is 77 KB). Run one token at a time there instead of failing the launch. + if xq == 0 && c.shared_mem_bytes > self.max_smem { + c.shared_mem_bytes = per; + for t in 0..nt { + let (xt, yt) = (x + 4 * t as u64 * x_stride as u64, y + 4 * t as u64 * y_stride as u64); + launch!(self, &k1, c, w.ptr, xt, yt, w.rows as i32, ncols, rpw, xq, pair.0, pair.1, pair.2); + } + return Ok(()); + } // One token: gemv, whose kernel has no nt branches (gemvn's nt = 4 path costs it registers). if nt == 1 { launch!(self, &k1, c, w.ptr, x, y, w.rows as i32, ncols, rpw, xq, pair.0, pair.1, pair.2); @@ -367,8 +390,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, - // attn_decode_gqa4_128 one. - let cp = !self.kv_q8 && (hd == 128 || hd == 256); + // attn_decode_gqa4_128 one. 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 = !self.kv_q8 && (hd == 128 || (hd == 256 && self.int8_mma)); let g4 = cp && g == 4 && hd == 128; let (name, cp) = if g4 { ("attn_decode_gqa4_128".to_string(), 2 * 32 * hd * 2) } else { (format!("attn_decode_gqa_{hd}"), if cp { 2 * 2 * 32 * hd * 2 } else { 0 }) }; if gq_min != i32::MAX { launch!(self, &name, 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); } @@ -519,6 +543,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 prequant_gemm_x(&self, w: &GpuTensor, n: usize, x: u64, x_stride: usize) -> anyhow::Result<()> { let mut q8 = self.q8.lock().unwrap(); self.q8_x(&mut q8, w.ty, w.cols, n, x, x_stride)?; @@ -669,7 +695,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 b8311440..b28b2ea9 100644 --- a/src/vulkan/mod.rs +++ b/src/vulkan/mod.rs @@ -1060,6 +1060,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 prequant_gemm_x(&self, w: &GpuTensor, n: usize, x: u64, x_stride: usize) -> anyhow::Result<()> { let mut q8 = self.3.lock().unwrap(); diff --git a/tests/kernels.rs b/tests/kernels.rs index a5d8f2c4..e97f320b 100644 --- a/tests/kernels.rs +++ b/tests/kernels.rs @@ -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()); @@ -1180,7 +1180,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() }; } }