diff --git a/crates/metaltile-codegen/src/emit.rs b/crates/metaltile-codegen/src/emit.rs index c7514f82..76ca23c6 100644 --- a/crates/metaltile-codegen/src/emit.rs +++ b/crates/metaltile-codegen/src/emit.rs @@ -176,6 +176,7 @@ pub fn render_swift_wrappers(kernels: &[Kernel]) -> String { ); for k in kernels { emit_swift_wrapper(&mut out, k); + emit_swift_wrapper_threadgroups(&mut out, k); if k.wants_indirect_variant { emit_swift_wrapper_indirect(&mut out, k); } @@ -242,6 +243,66 @@ fn emit_swift_wrapper(out: &mut String, k: &Kernel) { writeln!(out, " }}\n").ok(); } +/// `dispatchThreadgroups` variant of `emit_swift_wrapper`. IDENTICAL bindings, +/// but `gridSize` is in THREADGROUPS and dispatch uses +/// `dispatchThreadgroups(_:threadsPerThreadgroup:)`. REQUIRED for coop_tile / +/// simdgroup_matrix kernels: `dispatchThreads` puts Metal in non-uniform- +/// threadgroup mode, which silently produces WRONG results for cooperative- +/// matrix ops even when the grid is a multiple of the threadgroup size. Plain +/// simd_sum/elementwise kernels are fine on either path. +fn emit_swift_wrapper_threadgroups(out: &mut String, k: &Kernel) { + use std::fmt::Write as _; + let fn_name = format!("{}_threadgroups", swift_safe_name(&k.name)); + writeln!( + out, + " /// dispatchThreadgroups variant of `{}` (gridSize in THREADGROUPS).", + k.name + ) + .ok(); + writeln!(out, " /// Use for coop_tile / simdgroup-matrix kernels.").ok(); + writeln!(out, " public static func {fn_name}(").ok(); + for p in &k.params { + let label = swift_safe_name(&p.name); + writeln!(out, " {label}: MTLBuffer, {label}Offset: Int = 0,").ok(); + } + for c in &k.constexprs { + let label = swift_safe_name(c.name.name()); + let swift_ty = swift_scalar_type(dtype_suffix(c.dtype)); + writeln!(out, " {label}: {swift_ty},").ok(); + } + writeln!(out, " gridSize: MTLSize,").ok(); + writeln!(out, " threadgroupSize: MTLSize,").ok(); + writeln!(out, " on commandBuffer: MTLCommandBuffer").ok(); + writeln!(out, " ) {{").ok(); + writeln!(out, " let pso = PSOCache.shared.pipelineState(for: \"{}\")", k.name).ok(); + writeln!( + out, + " guard let enc = commandBuffer.makeComputeCommandEncoder() else {{ return }}" + ) + .ok(); + writeln!(out, " enc.setComputePipelineState(pso)").ok(); + let mut slot = 0usize; + for p in &k.params { + let label = swift_safe_name(&p.name); + writeln!(out, " enc.setBuffer({label}, offset: {label}Offset, index: {slot})").ok(); + slot += 1; + } + for c in &k.constexprs { + let label = swift_safe_name(c.name.name()); + let len = swift_scalar_size(dtype_suffix(c.dtype)); + writeln!(out, " var {label}_v = {label}").ok(); + writeln!(out, " enc.setBytes(&{label}_v, length: {len}, index: {slot})").ok(); + slot += 1; + } + writeln!( + out, + " enc.dispatchThreadgroups(gridSize, threadsPerThreadgroup: threadgroupSize)" + ) + .ok(); + writeln!(out, " enc.endEncoding()").ok(); + writeln!(out, " }}\n").ok(); +} + /// Indirect-dispatch variant of `emit_swift_wrapper`. Same buffer + /// constexpr bindings, same PSO (the underlying kernel is unchanged), /// but the dispatch shape comes from an `MTLBuffer` carrying @@ -490,6 +551,7 @@ pub fn dtype_suffix(dt: DType) -> &'static str { DType::U32 => "u32", DType::I8 => "i8", DType::U8 => "u8", + DType::U16 => "u16", DType::I64 => "i64", DType::U64 => "u64", DType::I4 => "i4", @@ -590,4 +652,48 @@ mod tests { assert!(swift.contains("func mt_add_f32(")); assert!(!swift.contains("_indirect(")); } + + /// The `dispatchThreadgroups` wrapper (REQUIRED for coop_tile / + /// simdgroup-matrix kernels) is a distinct emit path: a `_threadgroups` + /// function taking `gridSize` + `threadgroupSize` in THREADGROUPS and + /// calling `dispatchThreadgroups(_:threadsPerThreadgroup:)`. Pins its + /// signature + buffer/constexpr binding + the dispatch call so a codegen + /// change to it surfaces as a reviewable diff (per docs/testing.md, a new + /// emit path lands with a fixture exercising it). + #[test] + fn emits_threadgroups_wrapper_with_bindings_and_dispatch() { + let mut k = dummy_kernel("mt_moe_bgemm_mma"); + // A second buffer + a constexpr so both binding loops are exercised. + k.params.push(Param { + name: "x".into(), + dtype: DType::F32, + shape: Shape::scalar(), + is_output: false, + kind: ParamKind::Tensor, + }); + k.constexprs.push(metaltile_core::ir::ConstExprDecl { + name: metaltile_core::constexpr::ConstExpr::new("k_in"), + dtype: DType::U32, + value: None, + }); + let swift = render_swift_wrappers(&[k]); + + // The threadgroups variant exists alongside the direct wrapper. + assert!(swift.contains("func mt_moe_bgemm_mma(")); + assert!(swift.contains("public static func mt_moe_bgemm_mma_threadgroups(")); + // gridSize is in THREADGROUPS and dispatch uses dispatchThreadgroups. + assert!(swift.contains("gridSize: MTLSize,")); + assert!(swift.contains("threadgroupSize: MTLSize,")); + assert!( + swift + .contains("dispatchThreadgroups(gridSize, threadsPerThreadgroup: threadgroupSize)") + ); + // Buffers bind by slot; the constexpr is set as inline bytes. + assert!(swift.contains("enc.setBuffer(out, offset: outOffset, index: 0)")); + assert!(swift.contains("enc.setBuffer(x, offset: xOffset, index: 1)")); + assert!(swift.contains("var k_in_v = k_in")); + assert!(swift.contains("enc.setBytes(&k_in_v, length: 4, index: 2)")); + // PSO lookup uses the base kernel name (no `_threadgroups` PSO exists). + assert!(swift.contains("pipelineState(for: \"mt_moe_bgemm_mma\")")); + } } diff --git a/crates/metaltile-codegen/src/passes/algebraic_simplify.rs b/crates/metaltile-codegen/src/passes/algebraic_simplify.rs index db00da5a..5af2d462 100644 --- a/crates/metaltile-codegen/src/passes/algebraic_simplify.rs +++ b/crates/metaltile-codegen/src/passes/algebraic_simplify.rs @@ -59,9 +59,8 @@ //! Tools", 2nd ed., §8.4–8.5. Classic treatment of algebraic identities and //! reduction in strength. -use std::collections::{BTreeMap, BTreeSet}; - use metaltile_core::ir::{BinOpKind, Block, BlockId, Kernel, Op, UnaryOpKind, ValueId}; +use rustc_hash::{FxHashMap, FxHashSet}; use super::remap; use crate::error::{Error, Result}; @@ -93,11 +92,11 @@ impl super::Pass for AlgebraicSimplifyPass { // --------------------------------------------------------------------------- /// Resolve transitive replacement chains: {v2→v1, v1→v0} becomes {v2→v0, v1→v0}. -fn resolve_transitive(map: &BTreeMap) -> BTreeMap { - let mut resolved = BTreeMap::new(); +fn resolve_transitive(map: &FxHashMap) -> FxHashMap { + let mut resolved = FxHashMap::default(); for (&key, &val) in map.iter() { let mut terminal = val; - let mut visited = BTreeSet::new(); + let mut visited = FxHashSet::default(); visited.insert(key); while let Some(&next) = map.get(&terminal) { if !visited.insert(terminal) { @@ -122,10 +121,11 @@ fn simplify_block_once(block: &mut Block) -> bool { let n = block.ops.len(); let mut const_overwrites: Vec<(usize, i64)> = Vec::new(); let mut op_replacements: Vec<(usize, Op)> = Vec::new(); - let mut vid_replacements: BTreeMap = BTreeMap::new(); + let mut vid_replacements: FxHashMap = + FxHashMap::with_capacity_and_hasher(block.ops.len(), Default::default()); // Build a map for peephole lookups. - let vid_to_op_pos: BTreeMap = block + let vid_to_op_pos: FxHashMap = block .results .iter() .enumerate() @@ -180,7 +180,7 @@ fn simplify_block_once(block: &mut Block) -> bool { // Remove dead ops whose results were redirected via ReplaceWithVid. // Without this, the same pattern re-matches on the next iteration, // producing the same replacement and causing an infinite fixpoint loop. - let dead_vids: BTreeSet = vid_replacements.keys().copied().collect(); + let dead_vids: FxHashSet = vid_replacements.keys().copied().collect(); if !dead_vids.is_empty() { let mut new_ops = Vec::new(); let mut new_results = Vec::new(); @@ -217,7 +217,7 @@ fn try_simplify( op: &Op, pos: usize, block: &Block, - vid_to_pos: &BTreeMap, + vid_to_pos: &FxHashMap, ) -> Option { match op { // ---- BinOp patterns ---- @@ -292,7 +292,7 @@ fn simplify_binop( rhs: ValueId, _pos: usize, block: &Block, - vid_to_pos: &BTreeMap, + vid_to_pos: &FxHashMap, ) -> Option { let lv = find_const_in_block(block, lhs); let rv = find_const_in_block(block, rhs); @@ -383,7 +383,7 @@ fn simplify_select( on_true: ValueId, on_false: ValueId, block: &Block, - vid_to_pos: &BTreeMap, + vid_to_pos: &FxHashMap, ) -> Option { // Select(true, a, b) → a if let Some(1) = find_const_in_block(block, cond) { @@ -427,7 +427,7 @@ fn find_const_in_block(block: &Block, vid: ValueId) -> Option { fn get_defining_op<'a>( vid: ValueId, block: &'a Block, - vid_to_pos: &BTreeMap, + vid_to_pos: &FxHashMap, ) -> Option<(usize, &'a Op)> { let &pos = vid_to_pos.get(&vid)?; Some((pos, &block.ops[pos])) @@ -437,7 +437,7 @@ fn get_defining_op<'a>( fn get_neg_arg( vid: ValueId, block: &Block, - vid_to_pos: &BTreeMap, + vid_to_pos: &FxHashMap, ) -> Option { let (_pos, op) = get_defining_op(vid, block, vid_to_pos)?; if let Op::UnaryOp { op: UnaryOpKind::Neg, value } = op { Some(*value) } else { None } @@ -447,7 +447,7 @@ fn get_neg_arg( fn get_not_arg( vid: ValueId, block: &Block, - vid_to_pos: &BTreeMap, + vid_to_pos: &FxHashMap, ) -> Option { let (_pos, op) = get_defining_op(vid, block, vid_to_pos)?; match op { @@ -476,7 +476,7 @@ fn get_not_arg( } /// Remap all ValueId references in an op. -fn remap_values_in_op(op: &mut Op, map: &BTreeMap) { +fn remap_values_in_op(op: &mut Op, map: &FxHashMap) { op.for_each_value_id_mut(&mut |v| { if let Some(&nv) = map.get(v) { *v = nv; diff --git a/crates/metaltile-codegen/src/passes/block_util.rs b/crates/metaltile-codegen/src/passes/block_util.rs index fb24abe0..fcadd8fc 100644 --- a/crates/metaltile-codegen/src/passes/block_util.rs +++ b/crates/metaltile-codegen/src/passes/block_util.rs @@ -6,9 +6,8 @@ //! within a [`Block`]. Used by Unroll, LICM, DSE, IfConversion, ValueSink, and //! any pass that rewrites the op/results arrays of a block. -use std::collections::BTreeSet; - use metaltile_core::ir::{Block, Op, ValueId}; +use rustc_hash::FxHashSet; /// Remove ops at the given indices from a block. /// @@ -18,7 +17,7 @@ pub fn remove_ops(block: &mut Block, indices: &[usize]) { if indices.is_empty() { return; } - let skip: BTreeSet = indices.iter().copied().collect(); + let skip: FxHashSet = indices.iter().copied().collect(); let old_ops = std::mem::take(&mut block.ops); let old_results = std::mem::take(&mut block.results); let mut new_ops = Vec::with_capacity(old_ops.len().saturating_sub(indices.len())); diff --git a/crates/metaltile-codegen/src/passes/copy_prop.rs b/crates/metaltile-codegen/src/passes/copy_prop.rs index d0de03f9..9a3a70a4 100644 --- a/crates/metaltile-codegen/src/passes/copy_prop.rs +++ b/crates/metaltile-codegen/src/passes/copy_prop.rs @@ -31,12 +31,11 @@ //! ACM TOPLAS 13(2):181–210. Sparse conditional constant propagation framework //! that subsumes copy propagation. -use std::collections::{BTreeMap, BTreeSet}; - use metaltile_core::{ dtype::DType, ir::{Block, BlockId, Kernel, Op, ValueId}, }; +use rustc_hash::{FxHashMap, FxHashSet}; use super::remap; use crate::error::{Error, Result}; @@ -64,11 +63,11 @@ impl super::Pass for CopyPropPass { } /// Resolve transitive replacement chains: {v2→v1, v1→v0} becomes {v2→v0, v1→v0}. -fn resolve_transitive(map: &BTreeMap) -> BTreeMap { - let mut resolved = BTreeMap::new(); +fn resolve_transitive(map: &FxHashMap) -> FxHashMap { + let mut resolved = FxHashMap::default(); for (&key, &val) in map.iter() { let mut terminal = val; - let mut visited = BTreeSet::new(); + let mut visited = FxHashSet::default(); visited.insert(key); while let Some(&next) = map.get(&terminal) { if !visited.insert(terminal) { @@ -91,7 +90,8 @@ fn copy_prop_block_fixpoint(block: &mut Block) { fn copy_prop_block_once(block: &mut Block) -> bool { let n = block.ops.len(); - let mut vid_replacements: BTreeMap = BTreeMap::new(); + let mut vid_replacements: FxHashMap = + FxHashMap::with_capacity_and_hasher(block.ops.len(), Default::default()); for i in 0..n { let op = &block.ops[i]; @@ -117,7 +117,7 @@ fn copy_prop_block_once(block: &mut Block) -> bool { // Remove dead ops whose results were redirected via identity propagation. // Without this, the same identity pattern re-matches on the next iteration, // producing the same replacement and causing an infinite fixpoint loop. - let dead_vids: BTreeSet = vid_replacements.keys().copied().collect(); + let dead_vids: FxHashSet = vid_replacements.keys().copied().collect(); if !dead_vids.is_empty() { let mut new_ops = Vec::new(); let mut new_results = Vec::new(); diff --git a/crates/metaltile-codegen/src/passes/licm.rs b/crates/metaltile-codegen/src/passes/licm.rs index 4099b363..2d38f2bf 100644 --- a/crates/metaltile-codegen/src/passes/licm.rs +++ b/crates/metaltile-codegen/src/passes/licm.rs @@ -38,7 +38,7 @@ use std::collections::{BTreeMap, BTreeSet}; use metaltile_core::ir::{Block, BlockId, Kernel, Op, ParamKind, ValueId}; -use rustc_hash::FxHashMap; +use rustc_hash::{FxHashMap, FxHashSet}; use super::remap; use crate::error::{Error, Result}; @@ -58,7 +58,7 @@ impl super::Pass for LicmPass { .collect(); // Build a definition map: ValueId -> BlockId where it's defined. - let mut def_block: BTreeMap = BTreeMap::new(); + let mut def_block: FxHashMap = FxHashMap::default(); for vid in kernel.body.results.iter().flatten() { def_block.insert(*vid, kernel.body.id); } @@ -101,7 +101,7 @@ impl super::Pass for LicmPass { fn licm_block( block: &mut Block, blocks: &mut FxHashMap, - def_block: &BTreeMap, + def_block: &FxHashMap, read_only: &BTreeSet, ) { let n = block.ops.len(); @@ -126,7 +126,7 @@ fn licm_block( // Build the initial invariant set: ValueIds defined before position `i` // in the parent block, plus any from ancestor blocks. - let mut invariant: BTreeSet = BTreeSet::new(); + let mut invariant: FxHashSet = FxHashSet::default(); for j in 0..i { if let Some(Some(vid)) = block.results.get(j) { invariant.insert(*vid); @@ -264,7 +264,7 @@ fn licm_block( /// Remove ops at given indices from a block. Indices must be sorted ascending. fn remove_ops_from_block(block: &mut Block, indices: &[usize]) { - let skip: BTreeSet = indices.iter().copied().collect(); + let skip: FxHashSet = indices.iter().copied().collect(); let old_ops = std::mem::take(&mut block.ops); let old_results = std::mem::take(&mut block.results); let mut new_ops = Vec::new(); diff --git a/crates/metaltile-codegen/src/passes/remap.rs b/crates/metaltile-codegen/src/passes/remap.rs index e90ef7a3..a9e91245 100644 --- a/crates/metaltile-codegen/src/passes/remap.rs +++ b/crates/metaltile-codegen/src/passes/remap.rs @@ -30,11 +30,27 @@ use smallvec::SmallVec; // remap_value_ids — mutate ValueId references in an Op // --------------------------------------------------------------------------- +/// `ValueId → ValueId` lookup abstraction so [`remap_value_ids`] accepts either +/// a `BTreeMap` or an `FxHashMap`. Callers pick the map type for performance; +/// the remap itself is a pure per-reference substitution, so iteration order is +/// irrelevant and either map yields identical output. +pub trait VidRemap { + fn remap(&self, v: ValueId) -> Option; +} + +impl VidRemap for BTreeMap { + fn remap(&self, v: ValueId) -> Option { self.get(&v).copied() } +} + +impl VidRemap for std::collections::HashMap { + fn remap(&self, v: ValueId) -> Option { self.get(&v).copied() } +} + /// Remap all `ValueId` references in `op` according to `map`. /// References not present in `map` are left unchanged. -pub fn remap_value_ids(op: &mut Op, map: &BTreeMap) { +pub fn remap_value_ids(op: &mut Op, map: &M) { op.for_each_value_id_mut(&mut |v| { - if let Some(&nv) = map.get(v) { + if let Some(nv) = map.remap(*v) { *v = nv; } }); diff --git a/crates/metaltile-core/src/dsl/dtype.rs b/crates/metaltile-core/src/dsl/dtype.rs index 3a07dcbb..4b73d835 100644 --- a/crates/metaltile-core/src/dsl/dtype.rs +++ b/crates/metaltile-core/src/dsl/dtype.rs @@ -23,6 +23,9 @@ pub enum DType { I4, /// 8-bit unsigned integer U8, + /// 16-bit unsigned integer (for reading misaligned GGUF quant blocks, e.g. + /// IQ2_XXS qs at byte offset 2 — u16-aligned, vs slow u8-recombine). + U16, /// 32-bit unsigned integer U32, /// 64-bit unsigned integer @@ -38,7 +41,7 @@ impl DType { pub const fn size_bytes(self) -> usize { match self { DType::F32 | DType::I32 | DType::U32 => 4, - DType::F16 | DType::BF16 => 2, + DType::F16 | DType::BF16 | DType::U16 => 2, DType::I8 | DType::U8 => 1, DType::I4 => 1, // packed, but addressable as 1 byte DType::U64 | DType::I64 => 8, @@ -46,15 +49,38 @@ impl DType { } } + // is_float / is_integer use an exhaustive `match` (not `matches!`) on + // purpose: adding a DType variant then fails to compile here until it's + // classified, so the two can never silently disagree or miss a new type. /// Whether this is a floating-point type. - pub const fn is_float(self) -> bool { matches!(self, DType::F32 | DType::F16 | DType::BF16) } + pub const fn is_float(self) -> bool { + match self { + DType::F32 | DType::F16 | DType::BF16 => true, + DType::I32 + | DType::I8 + | DType::I4 + | DType::U8 + | DType::U16 + | DType::U32 + | DType::U64 + | DType::I64 + | DType::Bool => false, + } + } /// Whether this is an integer type. pub const fn is_integer(self) -> bool { - matches!( - self, - DType::I32 | DType::I8 | DType::I4 | DType::U8 | DType::U32 | DType::U64 | DType::I64 - ) + match self { + DType::I32 + | DType::I8 + | DType::I4 + | DType::U8 + | DType::U16 + | DType::U32 + | DType::U64 + | DType::I64 => true, + DType::F32 | DType::F16 | DType::BF16 | DType::Bool => false, + } } /// Metal Shading Language name for this type. @@ -67,6 +93,7 @@ impl DType { DType::I8 => "char", DType::I4 => "char", // packed char DType::U8 => "uchar", + DType::U16 => "ushort", DType::U32 => "uint", DType::U64 => "ulong", DType::I64 => "long", @@ -84,6 +111,7 @@ impl DType { DType::I8 => "i8", DType::I4 => "i8", // stored as i8 DType::U8 => "u8", + DType::U16 => "u16", DType::U32 => "u32", DType::U64 => "u64", DType::I64 => "i64", @@ -107,6 +135,7 @@ impl DType { DType::I8 => "i8", DType::I4 => "i4", DType::U8 => "u8", + DType::U16 => "u16", DType::U32 => "u32", DType::U64 => "u64", DType::I64 => "i64", @@ -127,6 +156,7 @@ impl FromStr for DType { "i8" => Ok(DType::I8), "i4" => Ok(DType::I4), "u8" => Ok(DType::U8), + "u16" => Ok(DType::U16), "u32" => Ok(DType::U32), "u64" => Ok(DType::U64), "i64" => Ok(DType::I64), @@ -148,6 +178,7 @@ mod tests { DType::I8, DType::I4, DType::U8, + DType::U16, DType::U32, DType::U64, DType::I64, diff --git a/crates/metaltile-macros/src/kernel/sig.rs b/crates/metaltile-macros/src/kernel/sig.rs index 0fef3a1c..7db7284d 100644 --- a/crates/metaltile-macros/src/kernel/sig.rs +++ b/crates/metaltile-macros/src/kernel/sig.rs @@ -192,6 +192,7 @@ fn parse_dtype_generic( "bf16" => quote! { DType::BF16 }, "i32" => quote! { DType::I32 }, "u32" => quote! { DType::U32 }, + "u16" | "ushort" => quote! { DType::U16 }, "i8" => quote! { DType::I8 }, "u8" => quote! { DType::U8 }, "bool" => quote! { DType::Bool }, diff --git a/crates/metaltile-std/src/ffai/axpy_scalar_inplace.rs b/crates/metaltile-std/src/ffai/axpy_scalar_inplace.rs new file mode 100644 index 00000000..a037bd47 --- /dev/null +++ b/crates/metaltile-std/src/ffai/axpy_scalar_inplace.rs @@ -0,0 +1,62 @@ +//! Copyright 2026 0xClandestine, Ekryski, TheTom, Ambisphaeric +//! SPDX-License-Identifier: Apache-2.0 +//! In-place scaled-add — `accum[i] += scalar * src[i]`. +//! +//! Fuses the MoE per-expert epilogue (broadcast-mul + element-add into +//! the running accumulator) into a single dispatch. Saves 1 buffer +//! allocation (the broadcast scalar tensor) + 1 kernel dispatch (the +//! mul) per expert per layer. At top-K=6 × 43 layers = 258 saves per +//! token on DSv4-Flash. + +use metaltile::kernel; + +#[kernel] +pub fn ffai_axpy_scalar_inplace(src: Tensor, mut accum: Tensor, #[constexpr] scalar: f32) { + let i = tid; + let s = load(src[i]).cast::(); + let a = load(accum[i]).cast::(); + store(accum[i], a + scalar * s); +} + +pub mod kernel_tests { + use metaltile::{test::*, test_kernel}; + + use super::ffai_axpy_scalar_inplace; + use crate::utils::{pack_f32, unpack_f32}; + + fn setup(n: usize, dt: DType) -> TestSetup { + let scalar = 0.42f32; + let src: Vec = (0..n).map(|i| (i as f32 * 0.013 - 0.4).sin() * 1.2).collect(); + let accum_in: Vec = (0..n).map(|i| (i as f32 * 0.017 + 0.1).cos() * 0.8).collect(); + let src_dt = unpack_f32(&pack_f32(&src, dt), dt); + let accum_dt = unpack_f32(&pack_f32(&accum_in, dt), dt); + let expected: Vec = + accum_dt.iter().zip(&src_dt).map(|(a, s)| a + scalar * s).collect(); + TestSetup::new(ffai_axpy_scalar_inplace::kernel_ir_for(dt)) + .input(TestBuffer::from_vec("src", pack_f32(&src, dt), dt)) + .input(TestBuffer::from_vec("accum", pack_f32(&accum_in, dt), dt)) + .constexpr("scalar", scalar) + .expect(TestBuffer::from_vec("accum", pack_f32(&expected, dt), dt)) + .grid_1d(n, 256) + } + + #[test_kernel(dtypes = [f32, f16, bf16], tol = [1e-5, 5e-3, 5e-2])] + fn test_axpy_scalar_decode(dt: DType) -> TestSetup { setup(4096, dt) } +} + +pub mod kernel_benches { + use metaltile::{bench, test::*}; + + use super::ffai_axpy_scalar_inplace; + + #[bench(name = "ffai/axpy_scalar_inplace", dtypes = [f32, f16, bf16])] + fn bench_axpy_scalar(dt: DType) -> BenchSetup { + let n = 4096usize; + BenchSetup::new(ffai_axpy_scalar_inplace::kernel_ir_for(dt)) + .buffer(BenchBuffer::random("src", n, dt)) + .buffer(BenchBuffer::random("accum", n, dt).output()) + .constexpr("scalar", 0.5f32) + .grid_1d(n, 256) + .bytes_moved((n * dt.size_bytes() * 2) as u64) + } +} diff --git a/crates/metaltile-std/src/ffai/dsv4_compressor_pool.rs b/crates/metaltile-std/src/ffai/dsv4_compressor_pool.rs new file mode 100644 index 00000000..dfc072a1 --- /dev/null +++ b/crates/metaltile-std/src/ffai/dsv4_compressor_pool.rs @@ -0,0 +1,188 @@ +//! Copyright 2026 0xClandestine, Ekryski, TheTom, Ambisphaeric +//! SPDX-License-Identifier: Apache-2.0 +//! DSv4 CSA / HCA compressor — softmax-gated weighted pool + APE. +//! +//! Per-layer KV compressor used on CSA (`compress_ratio=4`) and HCA +//! (`compress_ratio=128`) layers. CSA uses **overlap** pooling +//! (`2 * ratio = 8` raw tokens per compressed entry); HCA uses +//! **non-overlap** pooling. Both reduce a pool window of post-`wkv`- +//! projected hidden states into one compressed KV entry per stride. +//! +//! ```text +//! for each pool window of `pool_len` raw KV entries: +//! g[pool_len] = softmax(gate_proj @ raw) # softmax-gated weights +//! out = sum_w g[w] * (kv_proj @ raw[w] + ape[w]) +//! ``` +//! +//! Inputs: +//! - `raw_kv: [pool_len, head_dim]` — already-`wkv`-projected KV +//! stream for the pool window (host slices the cache). +//! - `gate: [pool_len]` — pre-softmax gate weights +//! (`gate_proj @ x` per raw token, gathered into the pool). +//! - `ape: [pool_len, head_dim]` — learned absolute-position +//! embedding added per pool slot. +//! +//! Output: +//! - `compressed: [head_dim]` — one compressed KV entry. +//! +//! Pool semantics — overlap (CSA, m=4): pool_len=8 raw entries per +//! one compressed entry, stride=4 between successive compressed +//! entries. Non-overlap (HCA, m=128): pool_len=128, stride=128. +//! Wrappers in Ops.swift slice `raw_kv` to the right pool window +//! per decode step; this kernel just reduces. +//! +//! The compressor's RoPE base differs per layer-type — `rope_theta=10K` +//! for CSA, `compress_rope_theta=160K` for HCA — and is applied BEFORE +//! this kernel by the upstream wrapper (rotating the last +//! `qk_rope_head_dim=64` of each raw KV row). +//! +//! ## Dispatch +//! +//! 1D grid over output `head_dim` elements. Each thread computes one +//! output position. **Single-pass online-softmax** — Flash-style +//! running (max, sum, weighted-acc) in one loop over the pool. Avoids +//! the codegen-CSE hazard with sequential `for _w` loops sharing a +//! reused `gate[_w]` load (DSL collapses them and references variables +//! out of scope across loop boundaries). + +use metaltile::kernel; + +#[kernel] +pub fn ffai_dsv4_compressor_pool( + raw_kv: Tensor, + gate: Tensor, + ape: Tensor, + mut compressed: Tensor, + #[constexpr] head_dim: u32, + #[constexpr] pool_len: u32, +) { + let d = tid; + if d < head_dim { + // Online-softmax single pass: maintain running (max, sum) of + // gate exponentials and a running weighted accumulator of + // (raw + ape) at this thread's output dim. Rescale running + // state by `exp(old_max - new_max)` when a larger gate is + // seen — identical numerics to the two-pass form. + let mut g_max = neg_infinity(); + let mut g_sum = 0.0f32; + let mut acc = 0.0f32; + for _w in range(0u32, pool_len, 1u32) { + let g = load(gate[_w]); + let new_max = select(g > g_max, g, g_max); + let factor = exp(g_max - new_max); + let weight = exp(g - new_max); + let raw_val = load(raw_kv[_w * head_dim + d]).cast::(); + let ape_val = load(ape[_w * head_dim + d]).cast::(); + g_sum = g_sum * factor + weight; + acc = acc * factor + weight * (raw_val + ape_val); + g_max = new_max; + } + // Store the f32 result directly: the DSL narrows f32→T implicitly + // at the Store site (no explicit `.cast::()` needed). + store(compressed[d], acc / g_sum); + } +} + +#[allow(clippy::needless_range_loop)] +pub mod kernel_tests { + use metaltile::{test::*, test_kernel}; + + use super::ffai_dsv4_compressor_pool; + use crate::utils::{pack_f32, unpack_f32}; + + fn cpu_reference( + raw_kv: &[f32], + gate: &[f32], + ape: &[f32], + pool_len: usize, + head_dim: usize, + ) -> Vec { + let mut soft = vec![0f32; pool_len]; + let g_max = gate.iter().cloned().fold(f32::NEG_INFINITY, f32::max); + let mut g_sum = 0f32; + for w in 0..pool_len { + soft[w] = (gate[w] - g_max).exp(); + g_sum += soft[w]; + } + for w in 0..pool_len { + soft[w] /= g_sum; + } + let mut out = vec![0f32; head_dim]; + for d in 0..head_dim { + let mut acc = 0f32; + for w in 0..pool_len { + acc += soft[w] * (raw_kv[w * head_dim + d] + ape[w * head_dim + d]); + } + out[d] = acc; + } + out + } + + fn setup(pool_len: usize, head_dim: usize, dt: DType) -> TestSetup { + let raw_kv: Vec = + (0..pool_len * head_dim).map(|i| (i as f32 * 0.013 - 1.7).sin() * 1.2).collect(); + let ape: Vec = + (0..pool_len * head_dim).map(|i| (i as f32 * 0.021 - 0.4).cos() * 0.6).collect(); + let gate: Vec = (0..pool_len).map(|w| (w as f32 - 4.0) * 0.3).collect(); + let raw_dt = unpack_f32(&pack_f32(&raw_kv, dt), dt); + let ape_dt = unpack_f32(&pack_f32(&ape, dt), dt); + let expected = cpu_reference(&raw_dt, &gate, &ape_dt, pool_len, head_dim); + TestSetup::new(ffai_dsv4_compressor_pool::kernel_ir_for(dt)) + .input(TestBuffer::from_vec("raw_kv", pack_f32(&raw_kv, dt), dt)) + .input(TestBuffer::from_vec("gate", pack_f32(&gate, DType::F32), DType::F32)) + .input(TestBuffer::from_vec("ape", pack_f32(&ape, dt), dt)) + .input(TestBuffer::zeros("compressed", head_dim, dt)) + .constexpr("head_dim", head_dim as u32) + .constexpr("pool_len", pool_len as u32) + .expect(TestBuffer::from_vec("compressed", pack_f32(&expected, dt), dt)) + .grid_1d(head_dim, 256) + } + + /// CSA overlap-pool — 8 raw tokens, DSv4 head_dim=512. + #[test_kernel(dtypes = [f32, f16, bf16], tol = [1e-4, 5e-3, 5e-2])] + fn test_compressor_pool_csa(dt: DType) -> TestSetup { setup(8, 512, dt) } + + /// HCA non-overlap-pool — 128 raw tokens, DSv4 head_dim=512. + #[test_kernel(dtypes = [f32, f16, bf16], tol = [1e-4, 5e-3, 5e-2])] + fn test_compressor_pool_hca(dt: DType) -> TestSetup { setup(128, 512, dt) } +} + +pub mod kernel_benches { + use metaltile::{bench, test::*}; + + use super::ffai_dsv4_compressor_pool; + + #[bench(name = "ffai/dsv4_compressor_pool_csa", dtypes = [f32, f16, bf16])] + fn bench_csa(dt: DType) -> BenchSetup { + let (pool, head_dim) = (8usize, 512usize); + BenchSetup::new(ffai_dsv4_compressor_pool::kernel_ir_for(dt)) + .buffer(BenchBuffer::random("raw_kv", pool * head_dim, dt)) + .buffer(BenchBuffer::random("gate", pool, DType::F32)) + .buffer(BenchBuffer::random("ape", pool * head_dim, dt)) + .buffer(BenchBuffer::zeros("compressed", head_dim, dt).output()) + .constexpr("head_dim", head_dim as u32) + .constexpr("pool_len", pool as u32) + .grid_1d(head_dim, 256) + .bytes_moved( + ((2 * pool * head_dim + pool) * dt.size_bytes() + head_dim * dt.size_bytes()) + as u64, + ) + } + + #[bench(name = "ffai/dsv4_compressor_pool_hca", dtypes = [f32, f16, bf16])] + fn bench_hca(dt: DType) -> BenchSetup { + let (pool, head_dim) = (128usize, 512usize); + BenchSetup::new(ffai_dsv4_compressor_pool::kernel_ir_for(dt)) + .buffer(BenchBuffer::random("raw_kv", pool * head_dim, dt)) + .buffer(BenchBuffer::random("gate", pool, DType::F32)) + .buffer(BenchBuffer::random("ape", pool * head_dim, dt)) + .buffer(BenchBuffer::zeros("compressed", head_dim, dt).output()) + .constexpr("head_dim", head_dim as u32) + .constexpr("pool_len", pool as u32) + .grid_1d(head_dim, 256) + .bytes_moved( + ((2 * pool * head_dim + pool) * dt.size_bytes() + head_dim * dt.size_bytes()) + as u64, + ) + } +} diff --git a/crates/metaltile-std/src/ffai/dsv4_csa_sdpa_decode.rs b/crates/metaltile-std/src/ffai/dsv4_csa_sdpa_decode.rs new file mode 100644 index 00000000..a732b607 --- /dev/null +++ b/crates/metaltile-std/src/ffai/dsv4_csa_sdpa_decode.rs @@ -0,0 +1,450 @@ +//! Copyright 2026 0xClandestine, Ekryski, TheTom, Ambisphaeric +//! SPDX-License-Identifier: Apache-2.0 +//! DSv4 CSA sparse-gather SDPA decode for `head_dim == 512`. +//! +//! Single-token attention over a **selected subset** of KV cache +//! positions. Clone of [`crate::ffai::sdpa_decode_d512`] with the +//! dense `for _t in 0..n_kv` inner loop replaced by an index-gather +//! over a caller-supplied `selected_indices[n_selected]` buffer. +//! +//! The caller builds `selected_indices` by unioning: +//! - the top-512 indices returned by the Lightning Indexer +//! (`ffai_dsv4_indexer_topk_block`) +//! - the trailing `sliding_window` (DSv4 default: 128) most-recent +//! cache positions +//! +//! …and de-duplicating + sorting on the host. The kernel itself is +//! oblivious to the union logic — it just gathers `(K, V)` at each +//! listed index and runs the same online-softmax flash-attention. +//! +//! ## DISPATCH INVARIANTS +//! +//! Identical to `sdpa_decode_d512`: +//! - **TPG = 512** (16 SG × 32 lanes), 16 head-dim elements per lane. +//! - **`head_dim == 512`**, hardcoded. +//! - **Grid: 1 TG per q_head** — `grid = (nQHeads * 512, 1, 1)`, +//! `tg = (512, 1, 1)`. +//! - **`nQHeads % nKVHeads == 0`** (GQA integer fan-out). +//! - **`n_selected <= kv_stride`** — every selected index must be a +//! valid cache slot. +//! +//! ## Why a separate kernel +//! +//! Could fold sparse / dense as a runtime-selector on a `gather` +//! function constant, but Apple's MSL spec collapses constant-folded +//! branches at PSO compile time, so the dispatch hot path is the +//! same. Keeping the two as parallel kernels preserves the cleanest +//! emitted MSL for each path (no `if gather` per inner-loop iter). + +use metaltile::kernel; + +#[kernel] +pub fn ffai_dsv4_csa_sdpa_decode( + q: Tensor, + k: Tensor, + v: Tensor, + selected_indices: Tensor, + out: Tensor, + #[constexpr] head_dim: u32, + #[constexpr] n_selected: u32, + #[constexpr] kv_stride: u32, + #[constexpr] heads_per_group: u32, + #[constexpr] scale: f32, +) { + let q_head = tgid_x; + let kv_head = q_head / heads_per_group; + let sg = simd_id; + let lane = simd_lane; + let ns = n_simd; + threadgroup_alloc("tg_max", 32); + threadgroup_alloc("tg_sum", 32); + threadgroup_alloc("tg_out0", 1056); + threadgroup_alloc("tg_out1", 1056); + threadgroup_alloc("tg_out2", 1056); + threadgroup_alloc("tg_out3", 1056); + let q_off = q_head * head_dim; + let kv_head_base = kv_head * kv_stride * head_dim; + let d0 = lane * 16u32; + let q0 = load(q[q_off + d0]).cast::() * scale; + let q1 = load(q[q_off + d0 + 1u32]).cast::() * scale; + let q2 = load(q[q_off + d0 + 2u32]).cast::() * scale; + let q3 = load(q[q_off + d0 + 3u32]).cast::() * scale; + let q4 = load(q[q_off + d0 + 4u32]).cast::() * scale; + let q5 = load(q[q_off + d0 + 5u32]).cast::() * scale; + let q6 = load(q[q_off + d0 + 6u32]).cast::() * scale; + let q7 = load(q[q_off + d0 + 7u32]).cast::() * scale; + let q8 = load(q[q_off + d0 + 8u32]).cast::() * scale; + let q9 = load(q[q_off + d0 + 9u32]).cast::() * scale; + let q10 = load(q[q_off + d0 + 10u32]).cast::() * scale; + let q11 = load(q[q_off + d0 + 11u32]).cast::() * scale; + let q12 = load(q[q_off + d0 + 12u32]).cast::() * scale; + let q13 = load(q[q_off + d0 + 13u32]).cast::() * scale; + let q14 = load(q[q_off + d0 + 14u32]).cast::() * scale; + let q15 = load(q[q_off + d0 + 15u32]).cast::() * scale; + let mut run_max = neg_infinity(); + let mut run_sum = 0.0f32; + let mut o0 = 0.0f32; + let mut o1 = 0.0f32; + let mut o2 = 0.0f32; + let mut o3 = 0.0f32; + let mut o4 = 0.0f32; + let mut o5 = 0.0f32; + let mut o6 = 0.0f32; + let mut o7 = 0.0f32; + let mut o8 = 0.0f32; + let mut o9 = 0.0f32; + let mut o10 = 0.0f32; + let mut o11 = 0.0f32; + let mut o12 = 0.0f32; + let mut o13 = 0.0f32; + let mut o14 = 0.0f32; + let mut o15 = 0.0f32; + // Inner loop walks the gather list instead of [0, n_kv). + for _s in range(sg, n_selected, ns) { + let t_idx = load(selected_indices[_s]); + let base = kv_head_base + t_idx * head_dim; + let kv0 = base + d0; + let k0 = load(k[kv0]).cast::(); + let k1 = load(k[kv0 + 1u32]).cast::(); + let k2 = load(k[kv0 + 2u32]).cast::(); + let k3 = load(k[kv0 + 3u32]).cast::(); + let k4 = load(k[kv0 + 4u32]).cast::(); + let k5 = load(k[kv0 + 5u32]).cast::(); + let k6 = load(k[kv0 + 6u32]).cast::(); + let k7 = load(k[kv0 + 7u32]).cast::(); + let k8 = load(k[kv0 + 8u32]).cast::(); + let k9 = load(k[kv0 + 9u32]).cast::(); + let k10 = load(k[kv0 + 10u32]).cast::(); + let k11 = load(k[kv0 + 11u32]).cast::(); + let k12 = load(k[kv0 + 12u32]).cast::(); + let k13 = load(k[kv0 + 13u32]).cast::(); + let k14 = load(k[kv0 + 14u32]).cast::(); + let k15 = load(k[kv0 + 15u32]).cast::(); + let partial = q0 * k0 + + q1 * k1 + + q2 * k2 + + q3 * k3 + + q4 * k4 + + q5 * k5 + + q6 * k6 + + q7 * k7 + + q8 * k8 + + q9 * k9 + + q10 * k10 + + q11 * k11 + + q12 * k12 + + q13 * k13 + + q14 * k14 + + q15 * k15; + let score = simd_sum(partial); + let new_max = select(score > run_max, score, run_max); + let factor = exp(run_max - new_max); + let weight = exp(score - new_max); + run_sum = run_sum * factor + weight; + run_max = new_max; + let v0 = load(v[kv0]).cast::(); + let v1 = load(v[kv0 + 1u32]).cast::(); + let v2 = load(v[kv0 + 2u32]).cast::(); + let v3 = load(v[kv0 + 3u32]).cast::(); + let v4 = load(v[kv0 + 4u32]).cast::(); + let v5 = load(v[kv0 + 5u32]).cast::(); + let v6 = load(v[kv0 + 6u32]).cast::(); + let v7 = load(v[kv0 + 7u32]).cast::(); + let v8 = load(v[kv0 + 8u32]).cast::(); + let v9 = load(v[kv0 + 9u32]).cast::(); + let v10 = load(v[kv0 + 10u32]).cast::(); + let v11 = load(v[kv0 + 11u32]).cast::(); + let v12 = load(v[kv0 + 12u32]).cast::(); + let v13 = load(v[kv0 + 13u32]).cast::(); + let v14 = load(v[kv0 + 14u32]).cast::(); + let v15 = load(v[kv0 + 15u32]).cast::(); + o0 = o0 * factor + weight * v0; + o1 = o1 * factor + weight * v1; + o2 = o2 * factor + weight * v2; + o3 = o3 * factor + weight * v3; + o4 = o4 * factor + weight * v4; + o5 = o5 * factor + weight * v5; + o6 = o6 * factor + weight * v6; + o7 = o7 * factor + weight * v7; + o8 = o8 * factor + weight * v8; + o9 = o9 * factor + weight * v9; + o10 = o10 * factor + weight * v10; + o11 = o11 * factor + weight * v11; + o12 = o12 * factor + weight * v12; + o13 = o13 * factor + weight * v13; + o14 = o14 * factor + weight * v14; + o15 = o15 * factor + weight * v15; + } + if lane == 0 { + threadgroup_store("tg_max", sg, run_max); + threadgroup_store("tg_sum", sg, run_sum); + } + threadgroup_barrier(); + if sg == 0 { + let g_max_in = select(lane < ns, threadgroup_load("tg_max", lane), neg_infinity()); + let g_max = simd_max(g_max_in); + let g_sum_in = + select(lane < ns, threadgroup_load("tg_sum", lane) * exp(g_max_in - g_max), 0.0f32); + let g_sum = simd_sum(g_sum_in); + if lane == 0 { + threadgroup_store("tg_max", 0, g_max); + threadgroup_store("tg_sum", 0, g_sum); + } + } + threadgroup_barrier(); + let g_max = threadgroup_load("tg_max", 0); + let g_sum = threadgroup_load("tg_sum", 0); + let rescale = select(g_sum > 0.0f32, exp(run_max - g_max) / g_sum, 0.0f32); + let stride = ns + 1u32; + let idx = lane * stride + sg; + threadgroup_store("tg_out0", idx, o0 * rescale); + threadgroup_store("tg_out1", idx, o1 * rescale); + threadgroup_store("tg_out2", idx, o2 * rescale); + threadgroup_store("tg_out3", idx, o3 * rescale); + threadgroup_barrier(); + if sg == 0 { + let mut so0 = 0.0f32; + let mut so1 = 0.0f32; + let mut so2 = 0.0f32; + let mut so3 = 0.0f32; + for _g in range(0u32, ns, 1u32) { + let ri = lane * stride + _g; + so0 = so0 + threadgroup_load("tg_out0", ri); + so1 = so1 + threadgroup_load("tg_out1", ri); + so2 = so2 + threadgroup_load("tg_out2", ri); + so3 = so3 + threadgroup_load("tg_out3", ri); + } + let out_off = q_off + d0; + store(out[out_off], so0.cast::()); + store(out[out_off + 1u32], so1.cast::()); + store(out[out_off + 2u32], so2.cast::()); + store(out[out_off + 3u32], so3.cast::()); + } + threadgroup_barrier(); + threadgroup_store("tg_out0", idx, o4 * rescale); + threadgroup_store("tg_out1", idx, o5 * rescale); + threadgroup_store("tg_out2", idx, o6 * rescale); + threadgroup_store("tg_out3", idx, o7 * rescale); + threadgroup_barrier(); + if sg == 0 { + let mut so4 = 0.0f32; + let mut so5 = 0.0f32; + let mut so6 = 0.0f32; + let mut so7 = 0.0f32; + for _g in range(0u32, ns, 1u32) { + let ri = lane * stride + _g; + so4 = so4 + threadgroup_load("tg_out0", ri); + so5 = so5 + threadgroup_load("tg_out1", ri); + so6 = so6 + threadgroup_load("tg_out2", ri); + so7 = so7 + threadgroup_load("tg_out3", ri); + } + let out_off = q_off + d0; + store(out[out_off + 4u32], so4.cast::()); + store(out[out_off + 5u32], so5.cast::()); + store(out[out_off + 6u32], so6.cast::()); + store(out[out_off + 7u32], so7.cast::()); + } + threadgroup_barrier(); + threadgroup_store("tg_out0", idx, o8 * rescale); + threadgroup_store("tg_out1", idx, o9 * rescale); + threadgroup_store("tg_out2", idx, o10 * rescale); + threadgroup_store("tg_out3", idx, o11 * rescale); + threadgroup_barrier(); + if sg == 0 { + let mut so8 = 0.0f32; + let mut so9 = 0.0f32; + let mut so10 = 0.0f32; + let mut so11 = 0.0f32; + for _g in range(0u32, ns, 1u32) { + let ri = lane * stride + _g; + so8 = so8 + threadgroup_load("tg_out0", ri); + so9 = so9 + threadgroup_load("tg_out1", ri); + so10 = so10 + threadgroup_load("tg_out2", ri); + so11 = so11 + threadgroup_load("tg_out3", ri); + } + let out_off = q_off + d0; + store(out[out_off + 8u32], so8.cast::()); + store(out[out_off + 9u32], so9.cast::()); + store(out[out_off + 10u32], so10.cast::()); + store(out[out_off + 11u32], so11.cast::()); + } + threadgroup_barrier(); + threadgroup_store("tg_out0", idx, o12 * rescale); + threadgroup_store("tg_out1", idx, o13 * rescale); + threadgroup_store("tg_out2", idx, o14 * rescale); + threadgroup_store("tg_out3", idx, o15 * rescale); + threadgroup_barrier(); + if sg == 0 { + let mut so12 = 0.0f32; + let mut so13 = 0.0f32; + let mut so14 = 0.0f32; + let mut so15 = 0.0f32; + for _g in range(0u32, ns, 1u32) { + let ri = lane * stride + _g; + so12 = so12 + threadgroup_load("tg_out0", ri); + so13 = so13 + threadgroup_load("tg_out1", ri); + so14 = so14 + threadgroup_load("tg_out2", ri); + so15 = so15 + threadgroup_load("tg_out3", ri); + } + let out_off = q_off + d0; + store(out[out_off + 12u32], so12.cast::()); + store(out[out_off + 13u32], so13.cast::()); + store(out[out_off + 14u32], so14.cast::()); + store(out[out_off + 15u32], so15.cast::()); + } +} + +#[cfg(test)] +mod tests { + use metaltile_codegen::msl::MslGenerator; + use metaltile_core::ir::KernelMode; + + use super::ffai_dsv4_csa_sdpa_decode; + use crate::bench_types::DType; + + fn msl_for(dt: DType) -> String { + let mut k = ffai_dsv4_csa_sdpa_decode::kernel_ir_for(dt); + k.mode = KernelMode::Reduction; + MslGenerator::default().generate(&k).expect("ffai_dsv4_csa_sdpa_decode codegen succeeds") + } + + #[test] + fn codegen_produces_nonempty_msl_for_all_float_dtypes() { + for dt in [DType::F32, DType::F16, DType::BF16] { + let src = msl_for(dt); + assert!(!src.trim().is_empty(), "MSL for {dt:?} should not be empty"); + assert!( + src.contains("kernel void ffai_dsv4_csa_sdpa_decode"), + "MSL for {dt:?} should declare ffai_dsv4_csa_sdpa_decode:\n{src}", + ); + } + } +} + +pub mod kernel_tests { + use metaltile::{test::*, test_kernel}; + + use super::ffai_dsv4_csa_sdpa_decode; + use crate::utils::{pack_f32, unpack_f32}; + + /// Sparse-gather oracle: dense SDPA but only over the listed cache + /// positions. The online-softmax flash form converges to the same + /// result. + #[allow(clippy::too_many_arguments)] + fn naive_sparse_sdpa( + q: &[f32], + k: &[f32], + v: &[f32], + selected: &[u32], + n_q_heads: usize, + n_kv_heads: usize, + head_dim: usize, + kv_stride: usize, + scale: f32, + ) -> Vec { + let gqa = n_q_heads / n_kv_heads; + let n_sel = selected.len(); + let mut out = vec![0.0f32; n_q_heads * head_dim]; + for qh in 0..n_q_heads { + let kvh = qh / gqa; + let q_off = qh * head_dim; + let kv_slab = kvh * kv_stride * head_dim; + let mut scores = vec![0.0f32; n_sel]; + for (s, &idx) in selected.iter().enumerate() { + let k_off = kv_slab + (idx as usize) * head_dim; + let mut dot = 0.0f32; + for d in 0..head_dim { + dot += q[q_off + d] * k[k_off + d]; + } + scores[s] = dot * scale; + } + let m = scores.iter().copied().fold(f32::NEG_INFINITY, f32::max); + let mut sum = 0.0f32; + for sc in scores.iter_mut() { + *sc = (*sc - m).exp(); + sum += *sc; + } + let inv = if sum > 0.0 { 1.0 / sum } else { 0.0 }; + for d in 0..head_dim { + let mut acc = 0.0f32; + for (s, &idx) in selected.iter().enumerate() { + let v_off = kv_slab + (idx as usize) * head_dim + d; + acc += scores[s] * inv * v[v_off]; + } + out[q_off + d] = acc; + } + } + out + } + + fn ramp(n: usize, step: f32, start: f32) -> Vec { + (0..n).map(|i| ((start + i as f32 * step) % 2.0) - 1.0).collect() + } + + #[test_kernel(dtypes = [f32, f16, bf16], tol = [1e-3, 3e-3, 1.5e-2])] + fn test_dsv4_csa_sdpa_decode(dt: DType) -> TestSetup { + let (n_q_heads, n_kv_heads, head_dim) = (8usize, 4usize, 512usize); + let (n_selected, kv_stride) = (32usize, 128usize); + let heads_per_group = n_q_heads / n_kv_heads; + let scale = 1.0f32 / (head_dim as f32).sqrt(); + let q = unpack_f32(&pack_f32(&ramp(n_q_heads * head_dim, 0.013, -0.4), dt), dt); + let k = + unpack_f32(&pack_f32(&ramp(n_kv_heads * kv_stride * head_dim, 0.011, -0.5), dt), dt); + let v = + unpack_f32(&pack_f32(&ramp(n_kv_heads * kv_stride * head_dim, 0.007, -0.3), dt), dt); + // Pick a non-trivial gather: every 4th position from + // [0, kv_stride), so the kernel walks 32 indices spanning the + // full cache stride — exercises the per-thread index load + // path under the GQA fan-out. + let selected: Vec = (0..n_selected).map(|i| (i as u32) * 4).collect(); + let expected = naive_sparse_sdpa( + &q, &k, &v, &selected, n_q_heads, n_kv_heads, head_dim, kv_stride, scale, + ); + let sel_bytes: Vec = selected.iter().flat_map(|i| i.to_le_bytes()).collect(); + TestSetup::new(ffai_dsv4_csa_sdpa_decode::kernel_ir_for(dt)) + .mode(KernelMode::Reduction) + .input(TestBuffer::from_vec("q", pack_f32(&q, dt), dt)) + .input(TestBuffer::from_vec("k", pack_f32(&k, dt), dt)) + .input(TestBuffer::from_vec("v", pack_f32(&v, dt), dt)) + .input(TestBuffer::from_vec("selected_indices", sel_bytes, DType::U32)) + .input(TestBuffer::zeros("out", n_q_heads * head_dim, dt)) + .constexpr("head_dim", head_dim as u32) + .constexpr("n_selected", n_selected as u32) + .constexpr("kv_stride", kv_stride as u32) + .constexpr("heads_per_group", heads_per_group as u32) + .constexpr("scale", scale) + .expect(TestBuffer::from_vec("out", pack_f32(&expected, dt), dt)) + .grid_3d(n_q_heads as u32, 1, 1, [512, 1, 1]) + } +} + +pub mod kernel_benches { + use metaltile::{bench, test::*}; + + use super::ffai_dsv4_csa_sdpa_decode; + + #[bench(name = "ffai/dsv4_csa_sdpa_decode", dtypes = [f32, f16, bf16])] + fn bench_csa(dt: DType) -> BenchSetup { + let (n_q_heads, n_kv_heads, head_dim) = (32usize, 8usize, 512usize); + let (n_selected, kv_stride) = (640usize, 4096usize); // 512 top-k + 128 sliding window + let heads_per_group = n_q_heads / n_kv_heads; + let scale = 1.0f32 / (head_dim as f32).sqrt(); + let bytes = (2 * n_q_heads * head_dim + 2 * n_kv_heads * kv_stride * head_dim) + * dt.size_bytes() + + n_selected * 4; + BenchSetup::new(ffai_dsv4_csa_sdpa_decode::kernel_ir_for(dt)) + .mode(KernelMode::Reduction) + .buffer(BenchBuffer::random("q", n_q_heads * head_dim, dt)) + .buffer(BenchBuffer::random("k", n_kv_heads * kv_stride * head_dim, dt)) + .buffer(BenchBuffer::random("v", n_kv_heads * kv_stride * head_dim, dt)) + .buffer(BenchBuffer::random("selected_indices", n_selected, DType::U32)) + .buffer(BenchBuffer::zeros("out", n_q_heads * head_dim, dt).output()) + .constexpr("head_dim", head_dim as u32) + .constexpr("n_selected", n_selected as u32) + .constexpr("kv_stride", kv_stride as u32) + .constexpr("heads_per_group", heads_per_group as u32) + .constexpr("scale", scale) + .grid_3d(n_q_heads as u32, 1, 1, [512, 1, 1]) + .bytes_moved(bytes as u64) + } +} diff --git a/crates/metaltile-std/src/ffai/dsv4_fp8_block_dequant.rs b/crates/metaltile-std/src/ffai/dsv4_fp8_block_dequant.rs new file mode 100644 index 00000000..b848d931 --- /dev/null +++ b/crates/metaltile-std/src/ffai/dsv4_fp8_block_dequant.rs @@ -0,0 +1,238 @@ +//! Copyright 2026 0xClandestine, Ekryski, TheTom, Ambisphaeric +//! SPDX-License-Identifier: Apache-2.0 +//! Block-FP8 e4m3 dequant — DSv4 attention + router weight format. +//! +//! DSv4 attention + router weights ship as 1-byte-per-weight FP8 +//! e4m3 storage with **per-(128×128)-block fp32 scales**. +//! +//! ```text +//! weight [M, N] u8 — fp8 e4m3 bytes +//! scales [M/128, N/128] f32 — per-(128×128)-block scale +//! +//! out[i, j] = fp8_lut[weight[i, j]] * scales[i/128, j/128] +//! ``` +//! +//! FP8 e4m3: 1 sign + 4 exp + 3 mantissa, bias=7, max=±448, finite +//! everywhere except `0x7F` / `0xFF` (NaN). Bit-format conversion to +//! fp32 is ~6-8 ALU ops with denormal handling; on Apple GPUs (no +//! native FP8 type) a **256-entry LUT** (512 bytes total, host- +//! precomputed) is the proven faster path — confirmed by the +//! turboquant_plus M5 Max / M2 Pro analysis showing constant-cache +//! lookups cost 14% / 25% of decode time vs the arithmetic path. +//! +//! ## GPU-resident split (loader produces these from raw weights) +//! +//! 1. `weight_bytes [M * N]` u8 — raw fp8 e4m3 bytes +//! 2. `scales [M/128 * N/128]` f32 — block scales +//! 3. `fp8_lut [256]` f32 — `e4m3_to_fp32(byte)` +//! table uploaded once at runtime init +//! +//! Output dtype follows T (fp32 / fp16 / bf16). The dequant stays in +//! f32 (scale × LUT product) before the implicit narrowing at store. +//! +//! ## Apple GPU shape +//! +//! 1 thread per output value, `m_dim` constexpr lets the kernel +//! address the (i/128, j/128) scale tile without per-thread shape +//! math. The LUT fits in constant cache; the scale gather is once +//! per 128×128 = 16384 outputs, so cache-multicast is effectively +//! free. +//! +//! Bench-stage matters: for matrix-vector decode (FFAI hot path), +//! the dequant-then-matmul fallback (which this kernel is the +//! dequant half of) ships fp16 weights to a downstream gemv. The +//! fused `ffai_dsv4_fp8_block_gemv` lands as a follow-up — Apple +//! has no native FP8 multiply, but a dequant-on-the-fly variant +//! that interleaves the LUT lookup with the gemv accumulator can +//! avoid materialising the full fp16 matrix. + +use metaltile::kernel; + +// Bare `#[kernel]` — mixed-dtype param set (concrete u8/f32 + generic T). +#[kernel] +pub fn ffai_dsv4_fp8_block_dequant( + weight_bytes: Tensor, + scales: Tensor, + fp8_lut: Tensor, + out: Tensor, + #[constexpr] m_dim: u32, + #[constexpr] n_dim: u32, +) { + let i = tid; + let total = m_dim * n_dim; + if i < total { + let row = i / n_dim; + let col = i - row * n_dim; // i % n_dim + // 128×128 block-scale grid. `cols_per_block = n_dim / 128` + // is the number of column-blocks per row-block; reconstruct + // by `i / 16384`-style integer division on the 1D flat + // `[M/128, N/128]` scale buffer. + let block_row = row / 128u32; + let block_col = col / 128u32; + let n_block_cols = n_dim / 128u32; + let s = load(scales[block_row * n_block_cols + block_col]); + + let byte = load(weight_bytes[i]).cast::(); + let mag = load(fp8_lut[byte]); + + // Store the f32 result directly: the DSL narrows f32→T implicitly + // at the Store site (no explicit `.cast::()` needed). + store(out[i], s * mag); + } +} + +pub mod kernel_tests { + use metaltile::{test::*, test_kernel}; + + use super::ffai_dsv4_fp8_block_dequant; + use crate::utils::pack_f32; + + /// Build the FP8 e4m3 byte → fp32 LUT. e4m3 layout: + /// `[s (1)] [e (4)] [m (3)]` with `bias = 7`; subnormals + NaN + /// per the IEEE-754-ish FP8 spec used by DeepSeek-V3. + fn build_fp8_lut() -> [f32; 256] { + let mut lut = [0.0f32; 256]; + for byte in 0..256_u32 { + let sign = (byte >> 7) & 1; + let exp = (byte >> 3) & 0xf; + let mantissa = byte & 0x7; + let sign_mult = if sign == 0 { 1.0 } else { -1.0 }; + let value: f32 = if exp == 0xf && mantissa == 0x7 { + // 0x7F / 0xFF = NaN. Encode as f32 NaN so the + // downstream NaN filter (if any) sees it. + f32::NAN + } else if exp == 0 { + // Subnormal: mantissa / 2^9 (no implicit leading 1). + sign_mult * (mantissa as f32) * (1.0 / 64.0) * (1.0 / 8.0) + } else { + // Normal: (1 + m/8) * 2^(e - 7). + let m_frac = 1.0 + (mantissa as f32) / 8.0; + let exp_scale = (2.0_f32).powf((exp as f32) - 7.0); + sign_mult * m_frac * exp_scale + }; + lut[byte as usize] = value; + } + lut + } + + /// Reference quantizer: per (128×128) block compute the matching + /// scale (`max_abs / 448`), encode each value to the nearest + /// FP8 e4m3 byte via the LUT. Round-trip dequant should be + /// quant-error-only, not pattern noise. + fn quantize_block_fp8(values: &[f32], m: usize, n: usize) -> (Vec, Vec) { + assert_eq!(m % 128, 0, "block-FP8 needs M divisible by 128"); + assert_eq!(n % 128, 0, "block-FP8 needs N divisible by 128"); + let lut = build_fp8_lut(); + let block_rows = m / 128; + let block_cols = n / 128; + let mut bytes = vec![0u8; m * n]; + let mut scales = vec![0f32; block_rows * block_cols]; + for br in 0..block_rows { + for bc in 0..block_cols { + // amax over the 128×128 block. + let mut amax = 0.0f32; + for di in 0..128 { + for dj in 0..128 { + let v = values[(br * 128 + di) * n + bc * 128 + dj].abs(); + if v > amax { + amax = v; + } + } + } + let scale = if amax > 0.0 { amax / 448.0 } else { 1.0 }; + scales[br * block_cols + bc] = scale; + let inv = 1.0 / scale; + for di in 0..128 { + for dj in 0..128 { + let normalised = values[(br * 128 + di) * n + bc * 128 + dj] * inv; + // Nearest byte search. 256 bytes — brute force + // is fine for the test fixture (production + // quantization is offline anyway). + let mut best = 0u8; + let mut best_err = f32::INFINITY; + for byte in 0..256_u32 { + let mag = lut[byte as usize]; + if !mag.is_finite() { + continue; + } + let err = (normalised - mag).abs(); + if err < best_err { + best_err = err; + best = byte as u8; + } + } + bytes[(br * 128 + di) * n + bc * 128 + dj] = best; + } + } + } + } + (bytes, scales) + } + + fn cpu_dequant(bytes: &[u8], scales: &[f32], m: usize, n: usize) -> Vec { + let lut = build_fp8_lut(); + let block_cols = n / 128; + let mut out = vec![0f32; m * n]; + for i in 0..m { + for j in 0..n { + let br = i / 128; + let bc = j / 128; + let scale = scales[br * block_cols + bc]; + let mag = lut[bytes[i * n + j] as usize]; + out[i * n + j] = scale * mag; + } + } + out + } + + fn setup(m: usize, n: usize, dt: DType) -> TestSetup { + let values: Vec = (0..m * n).map(|i| (i as f32 * 0.00031 - 1.7).sin() * 5.0).collect(); + let (bytes, scales) = quantize_block_fp8(&values, m, n); + let lut = build_fp8_lut(); + let dequantized = cpu_dequant(&bytes, &scales, m, n); + TestSetup::new(ffai_dsv4_fp8_block_dequant::kernel_ir_for(dt)) + .input(TestBuffer::from_vec("weight_bytes", bytes, DType::U8)) + .input(TestBuffer::from_vec("scales", pack_f32(&scales, DType::F32), DType::F32)) + .input(TestBuffer::from_vec("fp8_lut", pack_f32(&lut, DType::F32), DType::F32)) + .input(TestBuffer::zeros("out", m * n, dt)) + .constexpr("m_dim", m as u32) + .constexpr("n_dim", n as u32) + .expect(TestBuffer::from_vec("out", pack_f32(&dequantized, dt), dt)) + .grid_1d(m * n, 256) + } + + /// Single-block (128×128) round-trip — the smallest valid shape. + #[test_kernel(dtypes = [f32, f16, bf16], tol = [1e-3, 5e-2, 2e-1])] + fn test_fp8_single_block(dt: DType) -> TestSetup { setup(128, 128, dt) } + + /// 4-block grid (256×256) — exercises the per-block-row/col + /// stride math. + #[test_kernel(dtypes = [f32, f16, bf16], tol = [1e-3, 5e-2, 2e-1])] + fn test_fp8_2x2_blocks(dt: DType) -> TestSetup { setup(256, 256, dt) } +} + +pub mod kernel_benches { + use metaltile::{bench, test::*}; + + use super::ffai_dsv4_fp8_block_dequant; + + #[bench(name = "ffai/dsv4_fp8_block_dequant", dtypes = [f32, f16, bf16])] + fn bench_fp8(dt: DType) -> BenchSetup { + // DSv4-Flash attention shape: wq_b is 1024 → 64*512 = 32768 + // (fused QK_nope + QK_rope across heads). Use 4096 × 4096 as + // a representative slab — same shape as wq_a. + let (m, n) = (4096usize, 4096usize); + let block_rows = m / 128; + let block_cols = n / 128; + BenchSetup::new(ffai_dsv4_fp8_block_dequant::kernel_ir_for(dt)) + .buffer(BenchBuffer::random("weight_bytes", m * n, DType::U8)) + .buffer(BenchBuffer::random("scales", block_rows * block_cols, DType::F32)) + .buffer(BenchBuffer::random("fp8_lut", 256, DType::F32)) + .buffer(BenchBuffer::zeros("out", m * n, dt).output()) + .constexpr("m_dim", m as u32) + .constexpr("n_dim", n as u32) + .grid_1d(m * n, 256) + // weights 1 B/elt + scale 4 B/(128*128 elts) + LUT 1 KB + output T + .bytes_moved((m * n + block_rows * block_cols * 4 + 1024 + m * n * dt.size_bytes()) as u64) + } +} diff --git a/crates/metaltile-std/src/ffai/dsv4_indexer_score.rs b/crates/metaltile-std/src/ffai/dsv4_indexer_score.rs new file mode 100644 index 00000000..d778aa1d --- /dev/null +++ b/crates/metaltile-std/src/ffai/dsv4_indexer_score.rs @@ -0,0 +1,157 @@ +//! Copyright 2026 0xClandestine, Ekryski, TheTom, Ambisphaeric +//! SPDX-License-Identifier: Apache-2.0 +//! DSv4 Lightning Indexer — per-position aggregate score. +//! +//! For each KV cache position `t`, compute one aggregate score the +//! top-k selector then sorts on: +//! +//! ```text +//! score[t] = sum_h w[h] * ReLU( q_idx[h] · k_idx[t, h] ) +//! ``` +//! +//! Where: +//! - `q_idx [n_heads, d_idx]` — projected query (one token). +//! - `k_idx [n_kv, n_heads, d_idx]` — projected key cache. +//! - `w [n_heads]` — per-head learnable scalar. +//! - `score [n_kv]` — output (f32; the top-k pass +//! downstream needs full mantissa for stable bitonic ordering). +//! +//! DSv4 production shape: `n_heads = 64`, `d_idx = 128`. The +//! indexer's job is to pick 512 cache positions per CSA / HCA layer +//! per decode step; this kernel produces the aggregate score, and a +//! follow-up `ffai_dsv4_indexer_topk` does the bitonic top-512. +//! +//! ## Dispatch +//! +//! 1 thread per cache position. Grid: `n_kv` threads in 1D. Each +//! thread walks `n_heads × d_idx = 8192` MADs on a hot 8 KB per-head +//! row of `k_idx`. Memory-bound (n_kv × 8 KB reads per call), so the +//! arithmetic cost is hidden under DRAM streaming on M5 Max. +//! +//! A fused per-simdgroup variant (1 SG per cache position, lanes +//! split d_idx, simd_sum across the dot) is the perf follow-up; this +//! straightforward version lands first with a GPU correctness test. + +use metaltile::kernel; + +#[kernel] +pub fn ffai_dsv4_indexer_score( + q_idx: Tensor, + k_idx: Tensor, + w: Tensor, + mut score: Tensor, + #[constexpr] n_heads: u32, + #[constexpr] d_idx: u32, + #[constexpr] n_kv: u32, +) { + let t = tid; + if t < n_kv { + let mut total = 0.0f32; + let kv_base = t * n_heads * d_idx; + for _h in range(0u32, n_heads, 1u32) { + let q_base = _h * d_idx; + let k_base = kv_base + _h * d_idx; + let mut dot = 0.0f32; + for _d in range(0u32, d_idx, 1u32) { + let q_val = load(q_idx[q_base + _d]).cast::(); + let k_val = load(k_idx[k_base + _d]).cast::(); + dot = dot + q_val * k_val; + } + // ReLU-clamp the per-head dot, then weight + accumulate. + let clamped = select(dot > 0.0f32, dot, 0.0f32); + let head_w = load(w[_h]); + total = total + head_w * clamped; + } + store(score[t], total); + } +} + +#[allow(clippy::needless_range_loop)] +pub mod kernel_tests { + use metaltile::{test::*, test_kernel}; + + use super::ffai_dsv4_indexer_score; + use crate::utils::{pack_f32, unpack_f32}; + + fn cpu_reference( + q_idx: &[f32], + k_idx: &[f32], + w: &[f32], + n_heads: usize, + d_idx: usize, + n_kv: usize, + ) -> Vec { + let mut out = vec![0f32; n_kv]; + for (t, slot) in out.iter_mut().enumerate() { + let mut total = 0f32; + for h in 0..n_heads { + let q_base = h * d_idx; + let k_base = t * n_heads * d_idx + h * d_idx; + let mut dot = 0f32; + for d in 0..d_idx { + dot += q_idx[q_base + d] * k_idx[k_base + d]; + } + let clamped = if dot > 0.0 { dot } else { 0.0 }; + total += w[h] * clamped; + } + *slot = total; + } + out + } + + fn setup(n_heads: usize, d_idx: usize, n_kv: usize, dt: DType) -> TestSetup { + let q_idx: Vec = + (0..n_heads * d_idx).map(|i| (i as f32 * 0.013 - 0.4).sin() * 0.8).collect(); + let k_idx: Vec = + (0..n_kv * n_heads * d_idx).map(|i| (i as f32 * 0.0073 + 0.2).cos() * 0.7).collect(); + // Mix of positive and negative head weights to exercise both + // sign branches of the per-head sum (a uniformly-positive w + // wouldn't catch a sign-flip bug in the per-head accumulator). + let w: Vec = (0..n_heads).map(|h| (h as f32 - n_heads as f32 / 2.0) * 0.05).collect(); + let q_dt = unpack_f32(&pack_f32(&q_idx, dt), dt); + let k_dt = unpack_f32(&pack_f32(&k_idx, dt), dt); + let expected = cpu_reference(&q_dt, &k_dt, &w, n_heads, d_idx, n_kv); + TestSetup::new(ffai_dsv4_indexer_score::kernel_ir_for(dt)) + .input(TestBuffer::from_vec("q_idx", pack_f32(&q_idx, dt), dt)) + .input(TestBuffer::from_vec("k_idx", pack_f32(&k_idx, dt), dt)) + .input(TestBuffer::from_vec("w", pack_f32(&w, DType::F32), DType::F32)) + .input(TestBuffer::zeros("score", n_kv, DType::F32)) + .constexpr("n_heads", n_heads as u32) + .constexpr("d_idx", d_idx as u32) + .constexpr("n_kv", n_kv as u32) + .expect(TestBuffer::from_vec("score", pack_f32(&expected, DType::F32), DType::F32)) + .grid_1d(n_kv, 256) + } + + /// Small shape sanity — 8 heads × 32 dim × 64 cache positions. + #[test_kernel(dtypes = [f32, f16, bf16], tol = [1e-4, 5e-2, 2e-1])] + fn test_indexer_score_small(dt: DType) -> TestSetup { setup(8, 32, 64, dt) } + + /// DSv4 production shape — 64 heads × 128 dim × 256 cache + /// positions (small `n_kv` to keep the test runtime bounded). + #[test_kernel(dtypes = [f32, f16, bf16], tol = [1e-3, 1e-1, 4e-1])] + fn test_indexer_score_dsv4(dt: DType) -> TestSetup { setup(64, 128, 256, dt) } +} + +pub mod kernel_benches { + use metaltile::{bench, test::*}; + + use super::ffai_dsv4_indexer_score; + + #[bench(name = "ffai/dsv4_indexer_score", dtypes = [f32, f16, bf16])] + fn bench_indexer(dt: DType) -> BenchSetup { + let (n_heads, d_idx, n_kv) = (64usize, 128usize, 4096usize); + let bytes = + (n_heads * d_idx + n_kv * n_heads * d_idx) * dt.size_bytes() + n_heads * 4 + n_kv * 4; + BenchSetup::new(ffai_dsv4_indexer_score::kernel_ir_for(dt)) + .buffer(BenchBuffer::random("q_idx", n_heads * d_idx, dt)) + .buffer(BenchBuffer::random("k_idx", n_kv * n_heads * d_idx, dt)) + .buffer(BenchBuffer::random("w", n_heads, DType::F32)) + .buffer(BenchBuffer::zeros("score", n_kv, DType::F32).output()) + .constexpr("n_heads", n_heads as u32) + .constexpr("d_idx", d_idx as u32) + .constexpr("n_kv", n_kv as u32) + .grid_1d(n_kv, 256) + .bytes_moved(bytes as u64) + } +} diff --git a/crates/metaltile-std/src/ffai/dsv4_indexer_topk.rs b/crates/metaltile-std/src/ffai/dsv4_indexer_topk.rs new file mode 100644 index 00000000..179f7bfd --- /dev/null +++ b/crates/metaltile-std/src/ffai/dsv4_indexer_topk.rs @@ -0,0 +1,171 @@ +//! Copyright 2026 0xClandestine, Ekryski, TheTom, Ambisphaeric +//! SPDX-License-Identifier: Apache-2.0 +//! DSv4 Lightning Indexer — top-k index selection over per-position +//! aggregate scores. +//! +//! Sits downstream of [`ffai_dsv4_indexer_score`]. Takes the +//! `score[n_kv]` produced by the indexer score kernel and returns the +//! indices of the K largest entries (DSv4 production: `K = 512`). The +//! returned indices feed CSA's sparse-gather SDPA inner loop. +//! +//! ## Single-block bitonic top-k +//! +//! For `n_kv <= 1024` the entire score array fits in one threadgroup's +//! shared memory. We do a parallel bitonic sort **descending** over +//! `(score, original_index)` pairs and emit the first `k` indices. +//! Same Batcher-pattern compare-and-swap as +//! [`crate::mlx::sort::mt_sort`] but with parallel storage for an +//! `u32` index tag so the original cache-position survives the swaps. +//! +//! Scores below `n_kv` are padded with `-INFINITY` so the sort +//! relegates them to the back regardless of K — caller does NOT need +//! a power-of-two `n_kv`, only an upper bound of 1024. +//! +//! ## DISPATCH INVARIANTS +//! +//! - **TPG = 256** (each thread owns 4 sort slots → 1024 total). +//! - **Grid: 1 threadgroup** — single-block sort. For `n_kv > 1024` +//! the multi-block-merge variant is the follow-up; the cleanest +//! path is `mt_sort` per chunk + a `mt_merge`-style co-rank merge +//! that keeps the index tag through the merge. +//! - `n_kv <= 1024` enforced at the call site. +//! - `k <= 1024` enforced at the call site. + +use metaltile::kernel; + +#[kernel] +pub fn ffai_dsv4_indexer_topk_block( + score: Tensor, + mut out_indices: Tensor, + #[constexpr] n_kv: u32, + #[constexpr] k: u32, +) { + let t = tid; + // Parallel TG buffers — sort the (score, index) pair in lockstep so + // each compare-and-swap moves both halves together. + threadgroup_alloc("tg_scores", 1024, f32); + threadgroup_alloc("tg_idx", 1024, u32); + + // Load 4 slots per thread. Beyond `n_kv` pads with `-INFINITY` so + // the descending sort sinks them to the tail and the first `k` + // slots are guaranteed to be real cache positions (when `k <= + // n_kv`). + let neg_inf = neg_infinity(); + for _e in range(0u32, 4u32, 1u32) { + let gi = t * 4u32 + _e; + let valid = gi < n_kv; + let raw_score = select(valid, load(score[gi]), neg_inf); + threadgroup_store("tg_scores", gi, raw_score); + threadgroup_store("tg_idx", gi, gi); + } + threadgroup_barrier(); + + // Bitonic sort DESCENDING. 10 outer stages (log2(1024)). The + // mt_sort kernel's barrier discipline (`if flip >= 7`) is reused: + // strides ≤ 64 stay within one simdgroup so the implicit lane + // ordering is enough; strides > 64 need a TG barrier. + for _k_stage in range(1u32, 11u32, 1u32) { + for _jb in range(0u32, _k_stage, 1u32) { + let flip = _k_stage - _jb - 1u32; + if flip >= 7u32 { + threadgroup_barrier(); + } + for _e in range(0u32, 4u32, 1u32) { + let gi = t * 4u32 + _e; + let partner = gi ^ (1u32 << flip); + if gi < partner { + let a_score = threadgroup_load("tg_scores", gi); + let b_score = threadgroup_load("tg_scores", partner); + let a_idx = threadgroup_load("tg_idx", gi); + let b_idx = threadgroup_load("tg_idx", partner); + let dir = (gi >> _k_stage) & 1u32; + // Descending: dir=0 keeps larger first → swap if `a < b`. + // dir=1 flips the sub-block (bitonic property). + let want_swap = select(dir == 0u32, a_score < b_score, a_score > b_score); + threadgroup_store("tg_scores", gi, select(want_swap, b_score, a_score)); + threadgroup_store("tg_scores", partner, select(want_swap, a_score, b_score)); + threadgroup_store("tg_idx", gi, select(want_swap, b_idx, a_idx)); + threadgroup_store("tg_idx", partner, select(want_swap, a_idx, b_idx)); + } + } + } + } + threadgroup_barrier(); + + // Emit the first `k` indices. Threads whose slots fall past `k` + // skip — the rest of `tg_idx` is sorted but unused. + for _e in range(0u32, 4u32, 1u32) { + let gi = t * 4u32 + _e; + if gi < k { + store(out_indices[gi], threadgroup_load("tg_idx", gi)); + } + } +} + +pub mod kernel_tests { + use metaltile::{test::*, test_kernel}; + + use super::ffai_dsv4_indexer_topk_block; + use crate::utils::pack_f32; + + /// CPU reference: argsort scores descending, take first `k` + /// original indices. Ties broken by ascending index (stable). + fn cpu_topk_indices(scores: &[f32], n_kv: usize, k: usize) -> Vec { + let mut paired: Vec<(f32, u32)> = + scores[..n_kv].iter().enumerate().map(|(i, &s)| (s, i as u32)).collect(); + paired.sort_by(|a, b| { + b.0.partial_cmp(&a.0).unwrap_or(std::cmp::Ordering::Equal).then(a.1.cmp(&b.1)) + }); + paired.iter().take(k).map(|(_, i)| *i).collect() + } + + fn setup(n_kv: usize, k: usize) -> TestSetup { + // Generate distinct scores so the top-k set is well-defined + // (no ties at the K-th boundary that could swap places between + // CPU and GPU stable-sort traversals). + let scores: Vec = (0..n_kv).map(|i| (i as f32 * 0.0379 - 1.7).sin() * 2.3).collect(); + let expected = cpu_topk_indices(&scores, n_kv, k); + // Pack u32 expected as raw little-endian bytes for the test + // framework — same path mxfp4_dequant uses for u32 outputs. + let expected_bytes: Vec = expected.iter().flat_map(|i| i.to_le_bytes()).collect(); + TestSetup::new(ffai_dsv4_indexer_topk_block::kernel_ir()) + .mode(KernelMode::Reduction) + .input(TestBuffer::from_vec("score", pack_f32(&scores, DType::F32), DType::F32)) + .input(TestBuffer::zeros("out_indices", k, DType::U32)) + .constexpr("n_kv", n_kv as u32) + .constexpr("k", k as u32) + .expect(TestBuffer::from_vec("out_indices", expected_bytes, DType::U32)) + .grid_3d(1, 1, 1, [256, 1, 1]) + } + + /// Small shape — 64 cache positions, top-8 — sanity check. + #[test_kernel(dtypes = [f32], tol = 0.0)] + fn test_topk_small(_dt: DType) -> TestSetup { setup(64, 8) } + + /// Mid shape — 256 positions, top-32. + #[test_kernel(dtypes = [f32], tol = 0.0)] + fn test_topk_mid(_dt: DType) -> TestSetup { setup(256, 32) } + + /// Full-block — 1024 positions, top-512 (DSv4 production `K`). + #[test_kernel(dtypes = [f32], tol = 0.0)] + fn test_topk_dsv4(_dt: DType) -> TestSetup { setup(1024, 512) } +} + +pub mod kernel_benches { + use metaltile::{bench, test::*}; + + use super::ffai_dsv4_indexer_topk_block; + + #[bench(name = "ffai/dsv4_indexer_topk_block", dtypes = [f32])] + fn bench_topk(_dt: DType) -> BenchSetup { + let (n_kv, k) = (1024usize, 512usize); + BenchSetup::new(ffai_dsv4_indexer_topk_block::kernel_ir()) + .mode(KernelMode::Reduction) + .buffer(BenchBuffer::random("score", n_kv, DType::F32)) + .buffer(BenchBuffer::zeros("out_indices", k, DType::U32).output()) + .constexpr("n_kv", n_kv as u32) + .constexpr("k", k as u32) + .grid_3d(1, 1, 1, [256, 1, 1]) + .bytes_moved((n_kv * 4 + k * 4) as u64) + } +} diff --git a/crates/metaltile-std/src/ffai/dsv4_mhc.rs b/crates/metaltile-std/src/ffai/dsv4_mhc.rs new file mode 100644 index 00000000..e9c7c070 --- /dev/null +++ b/crates/metaltile-std/src/ffai/dsv4_mhc.rs @@ -0,0 +1,247 @@ +//! Copyright 2026 0xClandestine, Ekryski, TheTom, Ambisphaeric +//! SPDX-License-Identifier: Apache-2.0 +//! DSv4 Manifold-Constrained Hyper-Connections (mHC) — runtime kernels. +//! +//! mHC carries a per-token `n_hc=4`-channel residual state +//! `H[hidden, n_hc, n_tokens]`. At every sub-block boundary the state +//! is collapsed to a single-channel input via per-token `pre[n_hc]` +//! weights, the sub-block runs, then the output is expanded back into +//! the 4 channels via per-token `post[n_hc]` weights + a per-token +//! `comb[n_hc, n_hc]` (Sinkhorn-normalized) residual remix: +//! +//! ```text +//! collapse: x[d, t] = sum_c pre[t, c] * H[d, c, t] +//! expand: H_new[d, dst, t] = block_out[d, t] * post[t, dst] +//! + sum_src comb[t, dst, src] * residual_H[d, src, t] +//! ``` +//! +//! `pre`, `post`, and `comb` are DYNAMIC per-token tensors produced +//! by [`crate::ffai::dsv4_mhc_sinkhorn_split`] from the `hc_*_fn @ +//! flatten(H)` 24-mix output. They are NOT stored model weights. +//! +//! ## Dispatch +//! +//! 1D grid over `hidden_dim`. Each thread loops over `n_tokens` and +//! the small `n_hc=4` inner — fine for decode (n_tokens=1) and OK +//! for short prefill chunks. Long-prefill rewrites would split the +//! grid over `(d, t)`. + +use metaltile::kernel; + +/// mHC collapse — `x[d, t] = sum_c pre[t, c] * H[d, c, t]`. +#[kernel] +pub fn ffai_dsv4_mhc_collapse( + state: Tensor, + pre: Tensor, + mut out: Tensor, + #[constexpr] hidden_dim: u32, + #[constexpr] n_hc: u32, + #[constexpr] n_tokens: u32, +) { + let d = tid; + if d < hidden_dim { + for _t in range(0u32, n_tokens, 1u32) { + let state_token_base = _t * n_hc * hidden_dim; + let pre_token_base = _t * n_hc; + let mut acc = 0.0f32; + for _c in range(0u32, n_hc, 1u32) { + let w = load(pre[pre_token_base + _c]); + let h = load(state[state_token_base + _c * hidden_dim + d]).cast::(); + acc = acc + w * h; + } + store(out[_t * hidden_dim + d], acc); + } + } +} + +/// mHC expand — channel-wise residual remix: +/// `H_new[d, dst, t] = block_out[d, t] * post[t, dst] +/// + sum_src comb[t, dst, src] * residual_H[d, src, t]` +/// +/// Reads from `residual_state` and writes to `state` — caller passes +/// the per-channel pre-update residual in `residual_state` and an +/// allocated output buffer in `state`. `comb` layout matches the +/// split kernel output: `comb[t * n_hc * n_hc + dst * n_hc + src]`. +#[kernel] +pub fn ffai_dsv4_mhc_expand( + block_out: Tensor, + post: Tensor, + comb: Tensor, + residual_state: Tensor, + mut state: Tensor, + #[constexpr] hidden_dim: u32, + #[constexpr] n_hc: u32, + #[constexpr] n_tokens: u32, +) { + let d = tid; + if d < hidden_dim { + for _t in range(0u32, n_tokens, 1u32) { + let state_token_base = _t * n_hc * hidden_dim; + let post_token_base = _t * n_hc; + let comb_token_base = _t * n_hc * n_hc; + let block_out_val = load(block_out[_t * hidden_dim + d]).cast::(); + for _dst in range(0u32, n_hc, 1u32) { + let post_w = load(post[post_token_base + _dst]); + let mut acc = block_out_val * post_w; + for _src in range(0u32, n_hc, 1u32) { + let comb_w = load(comb[comb_token_base + _dst * n_hc + _src]); + let resid = load(residual_state[state_token_base + _src * hidden_dim + d]) + .cast::(); + acc = acc + comb_w * resid; + } + store(state[state_token_base + _dst * hidden_dim + d], acc); + } + } + } +} + +pub mod kernel_tests { + use metaltile::{test::*, test_kernel}; + + use super::{ffai_dsv4_mhc_collapse, ffai_dsv4_mhc_expand}; + use crate::utils::{pack_f32, unpack_f32}; + + fn cpu_collapse( + state: &[f32], + pre: &[f32], + hidden_dim: usize, + n_hc: usize, + n_tokens: usize, + ) -> Vec { + let mut out = vec![0f32; n_tokens * hidden_dim]; + for t in 0..n_tokens { + for d in 0..hidden_dim { + let mut acc = 0f32; + for c in 0..n_hc { + acc += pre[t * n_hc + c] * state[t * n_hc * hidden_dim + c * hidden_dim + d]; + } + out[t * hidden_dim + d] = acc; + } + } + out + } + + fn setup_collapse(hidden_dim: usize, n_hc: usize, n_tokens: usize, dt: DType) -> TestSetup { + let state: Vec = (0..n_tokens * n_hc * hidden_dim) + .map(|i| (i as f32 * 0.011 - 0.7).sin() * 1.4) + .collect(); + let pre: Vec = (0..n_tokens * n_hc).map(|i| (i as f32 - 4.0) * 0.2 + 0.5).collect(); + let state_dt = unpack_f32(&pack_f32(&state, dt), dt); + let expected = cpu_collapse(&state_dt, &pre, hidden_dim, n_hc, n_tokens); + TestSetup::new(ffai_dsv4_mhc_collapse::kernel_ir_for(dt)) + .input(TestBuffer::from_vec("state", pack_f32(&state, dt), dt)) + .input(TestBuffer::from_vec("pre", pack_f32(&pre, DType::F32), DType::F32)) + .input(TestBuffer::zeros("out", n_tokens * hidden_dim, dt)) + .constexpr("hidden_dim", hidden_dim as u32) + .constexpr("n_hc", n_hc as u32) + .constexpr("n_tokens", n_tokens as u32) + .expect(TestBuffer::from_vec("out", pack_f32(&expected, dt), dt)) + .grid_1d(hidden_dim, 256) + } + + #[test_kernel(dtypes = [f32, f16, bf16], tol = [1e-5, 5e-3, 5e-2])] + fn test_mhc_collapse_decode(dt: DType) -> TestSetup { setup_collapse(4096, 4, 1, dt) } + + #[test_kernel(dtypes = [f32, f16, bf16], tol = [1e-5, 5e-3, 5e-2])] + fn test_mhc_collapse_batch(dt: DType) -> TestSetup { setup_collapse(512, 4, 4, dt) } + + fn cpu_expand( + block_out: &[f32], + post: &[f32], + comb: &[f32], + residual: &[f32], + hidden_dim: usize, + n_hc: usize, + n_tokens: usize, + ) -> Vec { + let mut out = vec![0f32; n_tokens * n_hc * hidden_dim]; + for t in 0..n_tokens { + for d in 0..hidden_dim { + let bo = block_out[t * hidden_dim + d]; + for dst in 0..n_hc { + let mut acc = bo * post[t * n_hc + dst]; + for src in 0..n_hc { + acc += comb[t * n_hc * n_hc + dst * n_hc + src] + * residual[t * n_hc * hidden_dim + src * hidden_dim + d]; + } + out[t * n_hc * hidden_dim + dst * hidden_dim + d] = acc; + } + } + } + out + } + + fn setup_expand(hidden_dim: usize, n_hc: usize, n_tokens: usize, dt: DType) -> TestSetup { + let block_out: Vec = + (0..n_tokens * hidden_dim).map(|i| (i as f32 * 0.017 - 1.2).sin() * 0.7).collect(); + let post: Vec = (0..n_tokens * n_hc).map(|i| (i as f32 - 2.0) * 0.3 + 0.4).collect(); + let comb: Vec = + (0..n_tokens * n_hc * n_hc).map(|i| (i as f32 * 0.05 - 0.1).cos() * 0.25).collect(); + let residual: Vec = (0..n_tokens * n_hc * hidden_dim) + .map(|i| (i as f32 * 0.0083 - 0.1).cos() * 0.9) + .collect(); + let block_out_dt = unpack_f32(&pack_f32(&block_out, dt), dt); + let residual_dt = unpack_f32(&pack_f32(&residual, dt), dt); + let expected = + cpu_expand(&block_out_dt, &post, &comb, &residual_dt, hidden_dim, n_hc, n_tokens); + TestSetup::new(ffai_dsv4_mhc_expand::kernel_ir_for(dt)) + .input(TestBuffer::from_vec("block_out", pack_f32(&block_out, dt), dt)) + .input(TestBuffer::from_vec("post", pack_f32(&post, DType::F32), DType::F32)) + .input(TestBuffer::from_vec("comb", pack_f32(&comb, DType::F32), DType::F32)) + .input(TestBuffer::from_vec("residual_state", pack_f32(&residual, dt), dt)) + .input(TestBuffer::zeros("state", n_tokens * n_hc * hidden_dim, dt)) + .constexpr("hidden_dim", hidden_dim as u32) + .constexpr("n_hc", n_hc as u32) + .constexpr("n_tokens", n_tokens as u32) + .expect(TestBuffer::from_vec("state", pack_f32(&expected, dt), dt)) + .grid_1d(hidden_dim, 256) + } + + #[test_kernel(dtypes = [f32, f16, bf16], tol = [1e-5, 5e-3, 5e-2])] + fn test_mhc_expand_decode(dt: DType) -> TestSetup { setup_expand(4096, 4, 1, dt) } + + #[test_kernel(dtypes = [f32, f16, bf16], tol = [1e-5, 5e-3, 5e-2])] + fn test_mhc_expand_batch(dt: DType) -> TestSetup { setup_expand(512, 4, 4, dt) } +} + +pub mod kernel_benches { + use metaltile::{bench, test::*}; + + use super::{ffai_dsv4_mhc_collapse, ffai_dsv4_mhc_expand}; + + #[bench(name = "ffai/dsv4_mhc_collapse", dtypes = [f32, f16, bf16])] + fn bench_collapse(dt: DType) -> BenchSetup { + let (hidden_dim, n_hc, n_tokens) = (4096usize, 4usize, 1usize); + BenchSetup::new(ffai_dsv4_mhc_collapse::kernel_ir_for(dt)) + .buffer(BenchBuffer::random("state", n_tokens * n_hc * hidden_dim, dt)) + .buffer(BenchBuffer::random("pre", n_tokens * n_hc, DType::F32)) + .buffer(BenchBuffer::zeros("out", n_tokens * hidden_dim, dt).output()) + .constexpr("hidden_dim", hidden_dim as u32) + .constexpr("n_hc", n_hc as u32) + .constexpr("n_tokens", n_tokens as u32) + .grid_1d(hidden_dim, 256) + .bytes_moved( + ((n_tokens * n_hc * hidden_dim + n_tokens * n_hc) * dt.size_bytes() + + n_tokens * hidden_dim * dt.size_bytes()) as u64, + ) + } + + #[bench(name = "ffai/dsv4_mhc_expand", dtypes = [f32, f16, bf16])] + fn bench_expand(dt: DType) -> BenchSetup { + let (hidden_dim, n_hc, n_tokens) = (4096usize, 4usize, 1usize); + BenchSetup::new(ffai_dsv4_mhc_expand::kernel_ir_for(dt)) + .buffer(BenchBuffer::random("block_out", n_tokens * hidden_dim, dt)) + .buffer(BenchBuffer::random("post", n_tokens * n_hc, DType::F32)) + .buffer(BenchBuffer::random("comb", n_tokens * n_hc * n_hc, DType::F32)) + .buffer(BenchBuffer::random("residual_state", n_tokens * n_hc * hidden_dim, dt)) + .buffer(BenchBuffer::zeros("state", n_tokens * n_hc * hidden_dim, dt).output()) + .constexpr("hidden_dim", hidden_dim as u32) + .constexpr("n_hc", n_hc as u32) + .constexpr("n_tokens", n_tokens as u32) + .grid_1d(hidden_dim, 256) + .bytes_moved( + ((2 * n_tokens * n_hc * hidden_dim + n_tokens * hidden_dim) * dt.size_bytes() + + n_tokens * (n_hc + n_hc * n_hc) * 4) as u64, + ) + } +} diff --git a/crates/metaltile-std/src/ffai/dsv4_mhc_sinkhorn_split.rs b/crates/metaltile-std/src/ffai/dsv4_mhc_sinkhorn_split.rs new file mode 100644 index 00000000..02f3cfac --- /dev/null +++ b/crates/metaltile-std/src/ffai/dsv4_mhc_sinkhorn_split.rs @@ -0,0 +1,353 @@ +//! Copyright 2026 0xClandestine, Ekryski, TheTom, Ambisphaeric +//! SPDX-License-Identifier: Apache-2.0 +//! DSv4 mHC dynamic-mix split — converts the `hc_fn @ flat` 24-mix +//! output into the three per-token control tensors `(pre, post, comb)` +//! the mHC collapse / expand kernels consume. +//! +//! ```text +//! mixes [24, n_tokens] = hc_fn @ flat (computed upstream) +//! scale [3] = (pre_scale, post_scale, comb_scale) +//! base [24] = per-slot bias +//! +//! pre[t, c] = sigmoid(mixes[t, c] * scale[0] + base[c]) + eps (c ∈ [0, 4)) +//! post[t, c] = 2 * sigmoid(mixes[t, 4+c] * scale[1] + base[4+c]) (c ∈ [0, 4)) +//! comb_raw[t, dst, src] = mixes[t, 8 + dst*4 + src] * scale[2] + base[8+...] +//! comb_softmax[t, dst, src] = softmax over src(comb_raw) + eps +//! comb[t, dst, src] = Sinkhorn(comb_softmax, iters) +//! ``` +//! +//! Sinkhorn-Knopp normalization: alternate row-normalize (`dst` axis) +//! and col-normalize (`src` axis) for `sinkhorn_iters` iterations. +//! Each normalize divides each slice by `sum + eps`. +//! +//! ## Dispatch +//! +//! 1 thread per token. `n_hc=4` is hardcoded — the 4×4 comb matrix +//! lives in 16 per-thread registers across the Sinkhorn loop. For +//! decode (n_tokens=1) the kernel is single-thread; the work fits in +//! a single dispatch even at prefill scale because per-token work is +//! O(n_hc² · iters) = 16 · iters ≈ <100 FMAs. + +use metaltile::kernel; + +#[kernel] +pub fn ffai_dsv4_mhc_sinkhorn_split( + mixes: Tensor, + scale: Tensor, + base: Tensor, + mut pre: Tensor, + mut post: Tensor, + mut comb: Tensor, + #[constexpr] n_tokens: u32, + #[constexpr] eps: f32, + #[constexpr] sinkhorn_iters: u32, +) { + let t = tid; + if t < n_tokens { + let pre_scale = load(scale[0]); + let post_scale = load(scale[1]); + let comb_scale = load(scale[2]); + + // ── Pre: 4 × (sigmoid + eps) ── + let mix_base = t * 24u32; + let p0 = load(mixes[mix_base]).cast::() * pre_scale + load(base[0]); + let p1 = load(mixes[mix_base + 1u32]).cast::() * pre_scale + load(base[1]); + let p2 = load(mixes[mix_base + 2u32]).cast::() * pre_scale + load(base[2]); + let p3 = load(mixes[mix_base + 3u32]).cast::() * pre_scale + load(base[3]); + let pre_t = t * 4u32; + store(pre[pre_t], 1.0f32 / (1.0f32 + exp(0.0f32 - p0)) + eps); + store(pre[pre_t + 1u32], 1.0f32 / (1.0f32 + exp(0.0f32 - p1)) + eps); + store(pre[pre_t + 2u32], 1.0f32 / (1.0f32 + exp(0.0f32 - p2)) + eps); + store(pre[pre_t + 3u32], 1.0f32 / (1.0f32 + exp(0.0f32 - p3)) + eps); + + // ── Post: 4 × (2 · sigmoid) ── + let q0 = load(mixes[mix_base + 4u32]).cast::() * post_scale + load(base[4]); + let q1 = load(mixes[mix_base + 5u32]).cast::() * post_scale + load(base[5]); + let q2 = load(mixes[mix_base + 6u32]).cast::() * post_scale + load(base[6]); + let q3 = load(mixes[mix_base + 7u32]).cast::() * post_scale + load(base[7]); + let post_t = t * 4u32; + store(post[post_t], 2.0f32 / (1.0f32 + exp(0.0f32 - q0))); + store(post[post_t + 1u32], 2.0f32 / (1.0f32 + exp(0.0f32 - q1))); + store(post[post_t + 2u32], 2.0f32 / (1.0f32 + exp(0.0f32 - q2))); + store(post[post_t + 3u32], 2.0f32 / (1.0f32 + exp(0.0f32 - q3))); + + // ── Comb 4×4: scale+bias, softmax-over-src, Sinkhorn ── + // + // Layout: comb[dst, src]. mixes[8 + dst*4 + src] supplies the + // raw value. base[8 + ...] is per-slot bias. Store result into + // `comb[t * 16 + dst*4 + src]`. + let r00 = load(mixes[mix_base + 8u32]).cast::() * comb_scale + load(base[8]); + let r01 = load(mixes[mix_base + 9u32]).cast::() * comb_scale + load(base[9]); + let r02 = load(mixes[mix_base + 10u32]).cast::() * comb_scale + load(base[10]); + let r03 = load(mixes[mix_base + 11u32]).cast::() * comb_scale + load(base[11]); + let r10 = load(mixes[mix_base + 12u32]).cast::() * comb_scale + load(base[12]); + let r11 = load(mixes[mix_base + 13u32]).cast::() * comb_scale + load(base[13]); + let r12 = load(mixes[mix_base + 14u32]).cast::() * comb_scale + load(base[14]); + let r13 = load(mixes[mix_base + 15u32]).cast::() * comb_scale + load(base[15]); + let r20 = load(mixes[mix_base + 16u32]).cast::() * comb_scale + load(base[16]); + let r21 = load(mixes[mix_base + 17u32]).cast::() * comb_scale + load(base[17]); + let r22 = load(mixes[mix_base + 18u32]).cast::() * comb_scale + load(base[18]); + let r23 = load(mixes[mix_base + 19u32]).cast::() * comb_scale + load(base[19]); + let r30 = load(mixes[mix_base + 20u32]).cast::() * comb_scale + load(base[20]); + let r31 = load(mixes[mix_base + 21u32]).cast::() * comb_scale + load(base[21]); + let r32 = load(mixes[mix_base + 22u32]).cast::() * comb_scale + load(base[22]); + let r33 = load(mixes[mix_base + 23u32]).cast::() * comb_scale + load(base[23]); + + // Softmax over src axis (per-row, dst fixed). Numerically stable. + let m0 = select(r00 > r01, r00, r01); + let m0 = select(m0 > r02, m0, r02); + let m0 = select(m0 > r03, m0, r03); + let e00 = exp(r00 - m0); + let e01 = exp(r01 - m0); + let e02 = exp(r02 - m0); + let e03 = exp(r03 - m0); + let s0 = e00 + e01 + e02 + e03; + let mut c00 = e00 / s0 + eps; + let mut c01 = e01 / s0 + eps; + let mut c02 = e02 / s0 + eps; + let mut c03 = e03 / s0 + eps; + + let m1 = select(r10 > r11, r10, r11); + let m1 = select(m1 > r12, m1, r12); + let m1 = select(m1 > r13, m1, r13); + let e10 = exp(r10 - m1); + let e11 = exp(r11 - m1); + let e12 = exp(r12 - m1); + let e13 = exp(r13 - m1); + let s1 = e10 + e11 + e12 + e13; + let mut c10 = e10 / s1 + eps; + let mut c11 = e11 / s1 + eps; + let mut c12 = e12 / s1 + eps; + let mut c13 = e13 / s1 + eps; + + let m2 = select(r20 > r21, r20, r21); + let m2 = select(m2 > r22, m2, r22); + let m2 = select(m2 > r23, m2, r23); + let e20 = exp(r20 - m2); + let e21 = exp(r21 - m2); + let e22 = exp(r22 - m2); + let e23 = exp(r23 - m2); + let s2 = e20 + e21 + e22 + e23; + let mut c20 = e20 / s2 + eps; + let mut c21 = e21 / s2 + eps; + let mut c22 = e22 / s2 + eps; + let mut c23 = e23 / s2 + eps; + + let m3 = select(r30 > r31, r30, r31); + let m3 = select(m3 > r32, m3, r32); + let m3 = select(m3 > r33, m3, r33); + let e30 = exp(r30 - m3); + let e31 = exp(r31 - m3); + let e32 = exp(r32 - m3); + let e33 = exp(r33 - m3); + let s3 = e30 + e31 + e32 + e33; + let mut c30 = e30 / s3 + eps; + let mut c31 = e31 / s3 + eps; + let mut c32 = e32 / s3 + eps; + let mut c33 = e33 / s3 + eps; + + // Sinkhorn-Knopp: alternate col-normalize and row-normalize. + // (The post-softmax matrix is already row-stochastic, so the + // first step is col-normalize.) Iterations chained via the + // 16 c__ live values — no array, no TG mem. + for _iter in range(0u32, sinkhorn_iters, 1u32) { + // Col normalize (each col divided by col sum + eps). + let cs0 = c00 + c10 + c20 + c30 + eps; + let cs1 = c01 + c11 + c21 + c31 + eps; + let cs2 = c02 + c12 + c22 + c32 + eps; + let cs3 = c03 + c13 + c23 + c33 + eps; + c00 = c00 / cs0; + c10 = c10 / cs0; + c20 = c20 / cs0; + c30 = c30 / cs0; + c01 = c01 / cs1; + c11 = c11 / cs1; + c21 = c21 / cs1; + c31 = c31 / cs1; + c02 = c02 / cs2; + c12 = c12 / cs2; + c22 = c22 / cs2; + c32 = c32 / cs2; + c03 = c03 / cs3; + c13 = c13 / cs3; + c23 = c23 / cs3; + c33 = c33 / cs3; + // Row normalize. + let rs0 = c00 + c01 + c02 + c03 + eps; + let rs1 = c10 + c11 + c12 + c13 + eps; + let rs2 = c20 + c21 + c22 + c23 + eps; + let rs3 = c30 + c31 + c32 + c33 + eps; + c00 = c00 / rs0; + c01 = c01 / rs0; + c02 = c02 / rs0; + c03 = c03 / rs0; + c10 = c10 / rs1; + c11 = c11 / rs1; + c12 = c12 / rs1; + c13 = c13 / rs1; + c20 = c20 / rs2; + c21 = c21 / rs2; + c22 = c22 / rs2; + c23 = c23 / rs2; + c30 = c30 / rs3; + c31 = c31 / rs3; + c32 = c32 / rs3; + c33 = c33 / rs3; + } + + // Write out [dst, src] layout: comb[t*16 + dst*4 + src]. + let comb_t = t * 16u32; + store(comb[comb_t], c00); + store(comb[comb_t + 1u32], c01); + store(comb[comb_t + 2u32], c02); + store(comb[comb_t + 3u32], c03); + store(comb[comb_t + 4u32], c10); + store(comb[comb_t + 5u32], c11); + store(comb[comb_t + 6u32], c12); + store(comb[comb_t + 7u32], c13); + store(comb[comb_t + 8u32], c20); + store(comb[comb_t + 9u32], c21); + store(comb[comb_t + 10u32], c22); + store(comb[comb_t + 11u32], c23); + store(comb[comb_t + 12u32], c30); + store(comb[comb_t + 13u32], c31); + store(comb[comb_t + 14u32], c32); + store(comb[comb_t + 15u32], c33); + } +} + +#[allow(clippy::needless_range_loop)] +pub mod kernel_tests { + use metaltile::{test::*, test_kernel}; + + use super::ffai_dsv4_mhc_sinkhorn_split; + use crate::utils::{pack_f32, unpack_f32}; + + #[allow(clippy::too_many_arguments)] + fn cpu_reference( + mixes: &[f32], + scale: &[f32], + base: &[f32], + n_tokens: usize, + eps: f32, + iters: usize, + ) -> (Vec, Vec, Vec) { + let mut pre = vec![0f32; n_tokens * 4]; + let mut post = vec![0f32; n_tokens * 4]; + let mut comb = vec![0f32; n_tokens * 16]; + for t in 0..n_tokens { + // pre + for c in 0..4 { + let raw = mixes[t * 24 + c] * scale[0] + base[c]; + pre[t * 4 + c] = 1.0 / (1.0 + (-raw).exp()) + eps; + } + // post + for c in 0..4 { + let raw = mixes[t * 24 + 4 + c] * scale[1] + base[4 + c]; + post[t * 4 + c] = 2.0 / (1.0 + (-raw).exp()); + } + // comb + let mut c = [[0f32; 4]; 4]; + for dst in 0..4 { + let mut row = [0f32; 4]; + for src in 0..4 { + row[src] = + mixes[t * 24 + 8 + dst * 4 + src] * scale[2] + base[8 + dst * 4 + src]; + } + let m = row.iter().cloned().fold(f32::NEG_INFINITY, f32::max); + let exps: [f32; 4] = [ + (row[0] - m).exp(), + (row[1] - m).exp(), + (row[2] - m).exp(), + (row[3] - m).exp(), + ]; + let s = exps.iter().sum::(); + for src in 0..4 { + c[dst][src] = exps[src] / s + eps; + } + } + for _ in 0..iters { + // Col normalize + for src in 0..4 { + let cs = c[0][src] + c[1][src] + c[2][src] + c[3][src] + eps; + for dst in 0..4 { + c[dst][src] /= cs; + } + } + // Row normalize + for dst in 0..4 { + let rs = c[dst][0] + c[dst][1] + c[dst][2] + c[dst][3] + eps; + for src in 0..4 { + c[dst][src] /= rs; + } + } + } + for dst in 0..4 { + for src in 0..4 { + comb[t * 16 + dst * 4 + src] = c[dst][src]; + } + } + } + (pre, post, comb) + } + + fn setup(n_tokens: usize, iters: usize, dt: DType) -> TestSetup { + let mixes: Vec = + (0..n_tokens * 24).map(|i| (i as f32 * 0.013 - 0.4).sin() * 0.8).collect(); + let scale: Vec = vec![0.5, 0.3, 0.7]; + let base: Vec = (0..24).map(|i| (i as f32 - 12.0) * 0.05).collect(); + let mixes_dt = unpack_f32(&pack_f32(&mixes, dt), dt); + let eps = 1e-6_f32; + let (pre_exp, post_exp, comb_exp) = + cpu_reference(&mixes_dt, &scale, &base, n_tokens, eps, iters); + TestSetup::new(ffai_dsv4_mhc_sinkhorn_split::kernel_ir_for(dt)) + .input(TestBuffer::from_vec("mixes", pack_f32(&mixes, dt), dt)) + .input(TestBuffer::from_vec("scale", pack_f32(&scale, DType::F32), DType::F32)) + .input(TestBuffer::from_vec("base", pack_f32(&base, DType::F32), DType::F32)) + .input(TestBuffer::zeros("pre", n_tokens * 4, DType::F32)) + .input(TestBuffer::zeros("post", n_tokens * 4, DType::F32)) + .input(TestBuffer::zeros("comb", n_tokens * 16, DType::F32)) + .constexpr("n_tokens", n_tokens as u32) + .constexpr("eps", eps) + .constexpr("sinkhorn_iters", iters as u32) + .expect(TestBuffer::from_vec("pre", pack_f32(&pre_exp, DType::F32), DType::F32)) + .expect(TestBuffer::from_vec("post", pack_f32(&post_exp, DType::F32), DType::F32)) + .expect(TestBuffer::from_vec("comb", pack_f32(&comb_exp, DType::F32), DType::F32)) + .grid_1d(n_tokens, 64) + } + + /// Single-token decode. + #[test_kernel(dtypes = [f32, f16, bf16], tol = [1e-5, 1e-3, 2e-2])] + fn test_mhc_sinkhorn_split_decode(dt: DType) -> TestSetup { setup(1, 1, dt) } + + /// Small batch + multi-iter Sinkhorn. + #[test_kernel(dtypes = [f32, f16, bf16], tol = [1e-5, 1e-3, 2e-2])] + fn test_mhc_sinkhorn_split_batch(dt: DType) -> TestSetup { setup(8, 3, dt) } +} + +pub mod kernel_benches { + use metaltile::{bench, test::*}; + + use super::ffai_dsv4_mhc_sinkhorn_split; + + #[bench(name = "ffai/dsv4_mhc_sinkhorn_split", dtypes = [f32, f16, bf16])] + fn bench_split(dt: DType) -> BenchSetup { + // 2 mHC blocks per layer × 43 layers = 86 dispatches/token. At + // decode (n_tokens=1) the per-dispatch work is trivial; bench + // shape covers a prefill chunk of 256 tokens. + let n_tokens = 256usize; + let iters = 1usize; + BenchSetup::new(ffai_dsv4_mhc_sinkhorn_split::kernel_ir_for(dt)) + .buffer(BenchBuffer::random("mixes", n_tokens * 24, dt)) + .buffer(BenchBuffer::random("scale", 3, DType::F32)) + .buffer(BenchBuffer::random("base", 24, DType::F32)) + .buffer(BenchBuffer::zeros("pre", n_tokens * 4, DType::F32).output()) + .buffer(BenchBuffer::zeros("post", n_tokens * 4, DType::F32).output()) + .buffer(BenchBuffer::zeros("comb", n_tokens * 16, DType::F32).output()) + .constexpr("n_tokens", n_tokens as u32) + .constexpr("eps", 1e-6_f32) + .constexpr("sinkhorn_iters", iters as u32) + .grid_1d(n_tokens, 64) + .bytes_moved((n_tokens * 24 * dt.size_bytes() + 27 * 4 + n_tokens * 24 * 4) as u64) + } +} diff --git a/crates/metaltile-std/src/ffai/dsv4_mxfp4_dequant.rs b/crates/metaltile-std/src/ffai/dsv4_mxfp4_dequant.rs new file mode 100644 index 00000000..6a470696 --- /dev/null +++ b/crates/metaltile-std/src/ffai/dsv4_mxfp4_dequant.rs @@ -0,0 +1,218 @@ +//! Copyright 2026 0xClandestine, Ekryski, TheTom, Ambisphaeric +//! SPDX-License-Identifier: Apache-2.0 +//! MXFP4 (OCP FP4 e2m1) block dequant. +//! +//! DSv4 MoE expert weights ship as native MXFP4 in the safetensors +//! checkpoint (separate code path from the GGUF IQ2_XXS recompression). +//! MXFP4 is the OCP Microscaling spec format: +//! +//! ```text +//! struct block_mxfp4 { +//! uint8_t e; // E8M0 unsigned exponent (block super-scale) +//! uint8_t qs[16]; // 32 packed 4-bit values, low nibble first +//! }; // 17 bytes per 32 values (BPW = 4.25) +//! ``` +//! +//! Each 4-bit code maps to one of 16 signed magnitudes (1 sign + 2 +//! exp + 1 mantissa per OCP MX v1.0): +//! +//! ```text +//! code | magnitude +//! ------+---------- +//! 0000 | +0 +//! 0001 | +0.5 +//! 0010 | +1.0 +//! 0011 | +1.5 +//! 0100 | +2.0 +//! 0101 | +3.0 +//! 0110 | +4.0 +//! 0111 | +6.0 +//! 1000 | -0 +//! 1001 | -0.5 +//! 1010 | -1.0 +//! 1011 | -1.5 +//! 1100 | -2.0 +//! 1101 | -3.0 +//! 1110 | -4.0 +//! 1111 | -6.0 +//! ``` +//! +//! No NaN in the value table — only the E8M0 scale carries a NaN +//! sentinel (`0xFF`). The decoder downstream must skip a NaN-scaled +//! block; this kernel does NOT special-case that. +//! +//! Dequant: `out = scale_f32 * mxfp4_lut[code]`. +//! +//! ## GPU-resident split (loader produces these from raw blocks) +//! +//! 1. `qs_packed [n_blocks * 4]` — `u32`, 16 packed-byte payload re-laid +//! as 4 LE u32 words per block (8 nibbles per word). +//! 2. `scales [n_blocks]` — `f32`, host-converted from E8M0 +//! (`scale = float_from_bits(uint32_t(e) << 23)` — Apple bit-cast +//! is unavailable in the DSL, so the loader does this once at +//! parse time). +//! 3. `lut [16]` — `f32`, the OCP-spec value table +//! uploaded once at runtime init. Shared across all dequant calls. +//! +//! ## Apple GPU shape notes +//! +//! - 16-entry LUT in a `Tensor` (read via `load(lut[code])`). +//! Apple's L1 multicast collapses the gather across lanes that +//! land on the same code. +//! - 1 thread per output value. Adjacent threads in a simdgroup +//! read consecutive nibbles within the same u32 word → coalesced +//! loads. +//! - This kernel is the dequant primitive; a fused-with-routing +//! MoE-GEMV variant lands in `dsv4_mxfp4_moe_gemv.rs` as a +//! follow-up (dequant→fp16→GEMM materialization is the perf wall +//! for MoE matmul, so the fused path skips the intermediate). + +use metaltile::kernel; + +// Bare `#[kernel]` — mixed-dtype param set (concrete u32/f32 + generic T). +#[kernel] +pub fn ffai_dsv4_mxfp4_dequant( + qs_packed: Tensor, + scales: Tensor, + lut: Tensor, + out: Tensor, + #[constexpr] n_values: u32, +) { + let i = tid; + if i < n_values { + let block = i / 32u32; + let in_block = i - block * 32u32; // 0..31 + // 32 nibbles per block, packed 8-per-u32 → word index ∈ [0..4), low-nibble-first. + let word_idx = in_block / 8u32; + let nibble_in_word = in_block & 7u32; + let word = load(qs_packed[block * 4u32 + word_idx]); + let code = (word >> (nibble_in_word * 4u32)) & 0xfu32; + let mag = load(lut[code]); + let s = load(scales[block]); + // Store the f32 result directly: the DSL narrows f32→T implicitly + // at the Store site (no explicit `.cast::()` needed). + store(out[i], s * mag); + } +} + +pub mod kernel_tests { + use metaltile::{test::*, test_kernel}; + + use super::ffai_dsv4_mxfp4_dequant; + use crate::utils::pack_f32; + + /// The canonical OCP-spec MXFP4 value table. Order matches the + /// 4-bit code's binary interpretation (low 4 bits of nibble). + const MXFP4_LUT: [f32; 16] = + [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0]; + + /// Reference quantizer: assign each input value to the nearest + /// MXFP4 magnitude, picking a per-block E8M0 scale that maximises + /// the representable range without saturating. Mirrors the + /// reference OCP encoder closely enough for round-trip dequant + /// invariance. + fn quantize_mxfp4(values: &[f32]) -> (Vec, Vec) { + assert_eq!(values.len() % 32, 0, "MXFP4 needs multiple-of-32 values"); + let n_blocks = values.len() / 32; + let mut qs_packed = Vec::with_capacity(n_blocks * 4); + let mut scales = Vec::with_capacity(n_blocks); + // Max representable magnitude in MXFP4 is 6.0; scale picks the + // exponent that maps `max_abs(block)` into that. + for block in values.chunks_exact(32) { + let amax = block.iter().map(|v| v.abs()).fold(0.0_f32, f32::max); + // E8M0 = unsigned 8-bit raw exponent, value 0xFF reserved + // as NaN sentinel. We compute the matching exponent in + // fp32 terms: `scale = 2^(e - 127)`, so `e = floor(log2(amax / 6)) + 127`. + let target = if amax > 0.0 { amax / 6.0 } else { 1.0 }; + let raw_exp = (target.log2().floor() as i32 + 127).clamp(0, 254); + let scale = (raw_exp - 127) as f32; + let scale_f32 = (2.0_f32).powf(scale); + scales.push(scale_f32); + let inv = if scale_f32 > 0.0 { 1.0 / scale_f32 } else { 0.0 }; + let mut nibbles = [0u8; 32]; + for (i, &v) in block.iter().enumerate() { + let normalised = v * inv; + // Nearest-magnitude search over the 16-entry table. + let mut best = 0u8; + let mut best_err = f32::INFINITY; + for (code, &mag) in MXFP4_LUT.iter().enumerate() { + let err = (normalised - mag).abs(); + if err < best_err { + best_err = err; + best = code as u8; + } + } + nibbles[i] = best; + } + // Repack 32 nibbles into 4 LE u32 words. + for w in 0..4 { + let mut word = 0u32; + for n in 0..8 { + word |= (nibbles[w * 8 + n] as u32) << (n * 4); + } + qs_packed.push(word); + } + } + (qs_packed, scales) + } + + fn cpu_dequant(qs_packed: &[u32], scales: &[f32]) -> Vec { + let n_blocks = scales.len(); + let mut out = Vec::with_capacity(n_blocks * 32); + for b in 0..n_blocks { + let s = scales[b]; + for i in 0..32_usize { + let word = qs_packed[b * 4 + i / 8]; + let nibble = (word >> ((i & 7) * 4)) & 0xf; + out.push(s * MXFP4_LUT[nibble as usize]); + } + } + out + } + + fn setup(n_blocks: usize, dt: DType) -> TestSetup { + let n = n_blocks * 32; + let values: Vec = (0..n).map(|i| (i as f32 * 0.0073 - 0.42).sin() * 2.7).collect(); + let (qs_packed, scales) = quantize_mxfp4(&values); + let dequantized = cpu_dequant(&qs_packed, &scales); + // Pack u32 vec as little-endian bytes for the test framework. + let qs_bytes: Vec = qs_packed.iter().flat_map(|w| w.to_le_bytes()).collect(); + TestSetup::new(ffai_dsv4_mxfp4_dequant::kernel_ir_for(dt)) + .input(TestBuffer::from_vec("qs_packed", qs_bytes, DType::U32)) + .input(TestBuffer::from_vec("scales", pack_f32(&scales, DType::F32), DType::F32)) + .input(TestBuffer::from_vec("lut", pack_f32(&MXFP4_LUT, DType::F32), DType::F32)) + .input(TestBuffer::zeros("out", n, dt)) + .constexpr("n_values", n as u32) + .expect(TestBuffer::from_vec("out", pack_f32(&dequantized, dt), dt)) + .grid_1d(n, 256) + } + + #[test_kernel(dtypes = [f32, f16, bf16], tol = [1e-4, 5e-3, 5e-2])] + fn test_mxfp4_single_block(dt: DType) -> TestSetup { setup(1, dt) } + + #[test_kernel(dtypes = [f32, f16, bf16], tol = [1e-4, 5e-3, 5e-2])] + fn test_mxfp4_many_blocks(dt: DType) -> TestSetup { setup(64, dt) } +} + +pub mod kernel_benches { + use metaltile::{bench, test::*}; + + use super::ffai_dsv4_mxfp4_dequant; + + #[bench(name = "ffai/dsv4_mxfp4_dequant", dtypes = [f32, f16, bf16])] + fn bench_mxfp4(dt: DType) -> BenchSetup { + // Representative MoE expert slab — 4096 × 2048 (DSv4 Flash + // expert intermediate=2048, hidden=4096). + let n = 4096 * 2048usize; + let n_blocks = n / 32; + BenchSetup::new(ffai_dsv4_mxfp4_dequant::kernel_ir_for(dt)) + .buffer(BenchBuffer::random("qs_packed", n_blocks * 4, DType::U32)) + .buffer(BenchBuffer::random("scales", n_blocks, DType::F32)) + .buffer(BenchBuffer::random("lut", 16, DType::F32)) + .buffer(BenchBuffer::zeros("out", n, dt).output()) + .constexpr("n_values", n as u32) + .grid_1d(n, 256) + // qs_packed 16 B + scale 4 B per block + lut 64 B once + output T. + .bytes_moved(((n_blocks * 20 + 64) + n * dt.size_bytes()) as u64) + } +} diff --git a/crates/metaltile-std/src/ffai/dsv4_partial_rope.rs b/crates/metaltile-std/src/ffai/dsv4_partial_rope.rs new file mode 100644 index 00000000..b39dec36 --- /dev/null +++ b/crates/metaltile-std/src/ffai/dsv4_partial_rope.rs @@ -0,0 +1,206 @@ +//! Copyright 2026 0xClandestine, Ekryski, TheTom, Ambisphaeric +//! SPDX-License-Identifier: Apache-2.0 +//! DSv4 partial RoPE — rotates only the tail `n_rot` dims of each +//! head, leaving the leading `n_nope = head_dim - n_rot` dims +//! untouched. +//! +//! DSv4 head_dim=512 splits as nope=448 + rope=64. Q and K are +//! computed full-rank, then this kernel rotates only the last 64 +//! dims of each head using the split-pair convention +//! `(dim_i, dim_i + n_rot/2)` for `dim_i ∈ [n_nope, n_nope + n_rot/2)`. +//! +//! ## Forward / inverse +//! +//! Forward rotation: `(x, y) → (x·cos θ − y·sin θ, x·sin θ + y·cos θ)` +//! Inverse rotation: `(x, y) → (x·cos θ + y·sin θ, −x·sin θ + y·cos θ)` +//! +//! The inverse is needed on the attention output before the grouped +//! O-LoRA — since K and V share the same tensor in MQA mode and only +//! K is logically rotated, the V-side contribution to the output +//! carries an extra rotation that has to be undone. Toggle via the +//! `inverse_flag` constexpr (`0` = forward, `1` = inverse). +//! +//! ## Dispatch +//! +//! Grid3D `[n_heads, half_rot, 1]` (one thread per rotation pair). +//! Operates in-place when `out == qk` — only the rope tail dims are +//! written. Caller is responsible for copying / passing through the +//! nope dims. + +use metaltile::kernel; + +#[kernel] +pub fn ffai_dsv4_partial_rope( + qk: Tensor, + out: Tensor, + #[constexpr] head_dim: u32, + #[constexpr] n_nope: u32, + #[constexpr] half_rot: u32, + #[constexpr] position: u32, + #[constexpr] theta_base: f32, + #[constexpr] inverse_flag: u32, + // YaRN (DSv4 compressed layers). Full layers pass freq_scale=1, + // ext_factor=0 → theta = pos·inv_freq (unchanged). The DSv4 reference + // cancels the YaRN magnitude scale (mscale=1), so it is omitted. + #[constexpr] freq_scale: f32, + #[constexpr] ext_factor: f32, + #[constexpr] corr_low: f32, + #[constexpr] corr_high: f32, +) { + let head = program_id::<0>(); + let pair_idx = program_id::<1>(); + let pair_f = pair_idx.cast::(); + let n_rot_f = (2u32 * half_rot).cast::(); + let inv_freq = exp2(0.0f32 - 2.0f32 * pair_f * log2(theta_base) / n_rot_f); + let pos_f = position.cast::(); + let theta_extrap = pos_f * inv_freq; + let theta_interp = freq_scale * theta_extrap; + // YaRN ramp: ramp = 1 - clamp((pair - corr_low)/(corr_high-corr_low), 0, 1). + let denom = select((corr_high - corr_low) > 0.001f32, corr_high - corr_low, 0.001f32); + let y_raw = (pair_f - corr_low) / denom; + let y_lo = select(y_raw > 0.0f32, y_raw, 0.0f32); + let y_cl = select(y_lo < 1.0f32, y_lo, 1.0f32); + let ramp_mix = (1.0f32 - y_cl) * ext_factor; + let theta_raw = theta_interp * (1.0f32 - ramp_mix) + theta_extrap * ramp_mix; + // Inverse rotation: flip the sign of theta (cos stays, sin flips). + let theta_signed = select(inverse_flag == 0u32, theta_raw, 0.0f32 - theta_raw); + let cos_t = cos(theta_signed); + let sin_t = sin(theta_signed); + let head_base = head * head_dim; + // DSv4 partial RoPE rotates ADJACENT + // pairs (tail[2p], tail[2p+1]) — GPT-J/interleaved convention, NOT the + // NEOX split-half (tail[p], tail[p+half_rot]). Pairing mismatch scrambles + // the q/k rope dims → wrong attention scores. inv_freq per pair index p + // is identical either way; only the dim layout differs. + let dim_lo = head_base + n_nope + 2u32 * pair_idx; + let dim_hi = head_base + n_nope + 2u32 * pair_idx + 1u32; + let x_lo = load(qk[dim_lo]).cast::(); + let x_hi = load(qk[dim_hi]).cast::(); + let o_lo = x_lo * cos_t - x_hi * sin_t; + let o_hi = x_lo * sin_t + x_hi * cos_t; + store(out[dim_lo], o_lo.cast::()); + store(out[dim_hi], o_hi.cast::()); +} + +pub mod kernel_tests { + use metaltile::{test::*, test_kernel}; + + use super::ffai_dsv4_partial_rope; + use crate::utils::{pack_f32, unpack_f32}; + + #[allow(clippy::too_many_arguments)] + fn cpu_reference( + qk: &[f32], + n_heads: usize, + head_dim: usize, + n_nope: usize, + half_rot: usize, + position: u32, + theta_base: f32, + inverse: bool, + ) -> Vec { + let mut out = qk.to_vec(); + for head in 0..n_heads { + for p in 0..half_rot { + let inv_freq = + (-(p as f32) * 2.0 * theta_base.ln() / (2.0 * half_rot as f32)).exp(); + let theta_raw = position as f32 * inv_freq; + let theta = if inverse { -theta_raw } else { theta_raw }; + let c = theta.cos(); + let s = theta.sin(); + let lo = head * head_dim + n_nope + 2 * p; + let hi = head * head_dim + n_nope + 2 * p + 1; + let x_lo = qk[lo]; + let x_hi = qk[hi]; + out[lo] = x_lo * c - x_hi * s; + out[hi] = x_lo * s + x_hi * c; + } + } + out + } + + #[allow(clippy::too_many_arguments)] + fn setup( + n_heads: usize, + head_dim: usize, + n_nope: usize, + half_rot: usize, + position: u32, + theta_base: f32, + inverse: bool, + dt: DType, + ) -> TestSetup { + let qk: Vec = + (0..n_heads * head_dim).map(|i| (i as f32 * 0.011 - 0.4).sin() * 1.2).collect(); + let qk_dt = unpack_f32(&pack_f32(&qk, dt), dt); + let expected = cpu_reference( + &qk_dt, n_heads, head_dim, n_nope, half_rot, position, theta_base, inverse, + ); + let inv_flag: u32 = if inverse { 1 } else { 0 }; + // `out` buffer is initialized to a copy of `qk` so the + // untouched nope dims pass through — kernel only writes the + // rope pairs (the in-place pattern callers rely on). + // freq_scale=1, ext_factor=0 → YaRN disabled (full-layer path). + TestSetup::new(ffai_dsv4_partial_rope::kernel_ir_for(dt)) + .mode(KernelMode::Grid3D) + .input(TestBuffer::from_vec("qk", pack_f32(&qk, dt), dt)) + .input(TestBuffer::from_vec("out", pack_f32(&qk, dt), dt)) + .constexpr("head_dim", head_dim as u32) + .constexpr("n_nope", n_nope as u32) + .constexpr("half_rot", half_rot as u32) + .constexpr("position", position) + .constexpr("theta_base", theta_base) + .constexpr("inverse_flag", inv_flag) + .constexpr("freq_scale", 1.0f32) + .constexpr("ext_factor", 0.0f32) + .constexpr("corr_low", 0.0f32) + .constexpr("corr_high", 0.0f32) + .expect(TestBuffer::from_vec("out", pack_f32(&expected, dt), dt)) + .grid_3d(n_heads as u32, half_rot as u32, 1, [1, 1, 1]) + } + + /// DSv4 production shape, forward, pos=64. + #[test_kernel(dtypes = [f32, f16, bf16], tol = [1e-4, 5e-3, 5e-2])] + fn test_partial_rope_dsv4_forward(dt: DType) -> TestSetup { + setup(64, 512, 448, 32, 64, 10_000.0, false, dt) + } + + /// DSv4 production shape, inverse, pos=64. + #[test_kernel(dtypes = [f32, f16, bf16], tol = [1e-4, 5e-3, 5e-2])] + fn test_partial_rope_dsv4_inverse(dt: DType) -> TestSetup { + setup(64, 512, 448, 32, 64, 10_000.0, true, dt) + } + + /// HCA compressed-stream theta (160K base). + #[test_kernel(dtypes = [f32, f16, bf16], tol = [1e-4, 5e-3, 5e-2])] + fn test_partial_rope_hca(dt: DType) -> TestSetup { + setup(64, 512, 448, 32, 16, 160_000.0, false, dt) + } +} + +pub mod kernel_benches { + use metaltile::{bench, test::*}; + + use super::ffai_dsv4_partial_rope; + + #[bench(name = "ffai/dsv4_partial_rope", dtypes = [f32, f16, bf16])] + fn bench_rope(dt: DType) -> BenchSetup { + let (n_heads, head_dim, n_nope, half_rot) = (64usize, 512usize, 448usize, 32usize); + BenchSetup::new(ffai_dsv4_partial_rope::kernel_ir_for(dt)) + .mode(KernelMode::Grid3D) + .buffer(BenchBuffer::random("qk", n_heads * head_dim, dt)) + .buffer(BenchBuffer::zeros("out", n_heads * head_dim, dt).output()) + .constexpr("head_dim", head_dim as u32) + .constexpr("n_nope", n_nope as u32) + .constexpr("half_rot", half_rot as u32) + .constexpr("position", 64u32) + .constexpr("theta_base", 10_000.0_f32) + .constexpr("inverse_flag", 0u32) + .constexpr("freq_scale", 1.0f32) + .constexpr("ext_factor", 0.0f32) + .constexpr("corr_low", 0.0f32) + .constexpr("corr_high", 0.0f32) + .grid_3d(n_heads as u32, half_rot as u32, 1, [1, 1, 1]) + .bytes_moved((4 * n_heads * half_rot * dt.size_bytes()) as u64) + } +} diff --git a/crates/metaltile-std/src/ffai/dsv4_partial_rope_rows.rs b/crates/metaltile-std/src/ffai/dsv4_partial_rope_rows.rs new file mode 100644 index 00000000..a8b9b649 --- /dev/null +++ b/crates/metaltile-std/src/ffai/dsv4_partial_rope_rows.rs @@ -0,0 +1,88 @@ +//! Copyright 2026 0xClandestine, Ekryski, TheTom, Ambisphaeric +//! SPDX-License-Identifier: Apache-2.0 +//! BATCHED DSv4 partial RoPE — the per-token ffai_dsv4_partial_rope applied +//! to N tokens in ONE dispatch (grid z = token). Prefill ropes token t at +//! absolute position `base_position + t`; the old prefill looped this N times +//! (3N tiny dispatches/layer for q + kv + inverse — the #1 warm-prefill cost). +//! Identical math to ffai_dsv4_partial_rope (adjacent GPT-J pairs, YaRN). +//! +//! qk/out are [n_tokens, n_heads, head_dim]. Grid3D [n_heads, half_rot, +//! n_tokens], one thread per (head, rot-pair, token). + +use metaltile::kernel; + +#[kernel] +#[allow(clippy::too_many_arguments)] +pub fn ffai_dsv4_partial_rope_rows( + qk: Tensor, + out: Tensor, + #[constexpr] head_dim: u32, + #[constexpr] n_nope: u32, + #[constexpr] half_rot: u32, + #[constexpr] n_heads: u32, + #[constexpr] base_position: u32, + #[constexpr] theta_base: f32, + #[constexpr] inverse_flag: u32, + #[constexpr] freq_scale: f32, + #[constexpr] ext_factor: f32, + #[constexpr] corr_low: f32, + #[constexpr] corr_high: f32, +) { + let head = program_id::<0>(); + let pair_idx = program_id::<1>(); + let token = program_id::<2>(); + let pair_f = pair_idx.cast::(); + let n_rot_f = (2u32 * half_rot).cast::(); + let inv_freq = exp2(0.0f32 - 2.0f32 * pair_f * log2(theta_base) / n_rot_f); + let pos_f = (base_position + token).cast::(); + let theta_extrap = pos_f * inv_freq; + let theta_interp = freq_scale * theta_extrap; + let denom = select((corr_high - corr_low) > 0.001f32, corr_high - corr_low, 0.001f32); + let y_raw = (pair_f - corr_low) / denom; + let y_lo = select(y_raw > 0.0f32, y_raw, 0.0f32); + let y_cl = select(y_lo < 1.0f32, y_lo, 1.0f32); + let ramp_mix = (1.0f32 - y_cl) * ext_factor; + let theta_raw = theta_interp * (1.0f32 - ramp_mix) + theta_extrap * ramp_mix; + let theta_signed = select(inverse_flag == 0u32, theta_raw, 0.0f32 - theta_raw); + let cos_t = cos(theta_signed); + let sin_t = sin(theta_signed); + let tok_base = token * n_heads * head_dim; + let head_base = tok_base + head * head_dim; + let dim_lo = head_base + n_nope + 2u32 * pair_idx; + let dim_hi = head_base + n_nope + 2u32 * pair_idx + 1u32; + let x_lo = load(qk[dim_lo]).cast::(); + let x_hi = load(qk[dim_hi]).cast::(); + let o_lo = x_lo * cos_t - x_hi * sin_t; + let o_hi = x_lo * sin_t + x_hi * cos_t; + store(out[dim_lo], o_lo.cast::()); + store(out[dim_hi], o_hi.cast::()); +} + +pub mod kernel_benches { + use metaltile::{bench, test::*}; + + use super::ffai_dsv4_partial_rope_rows; + + #[bench(name = "ffai/dsv4_partial_rope_rows", dtypes = [f32, f16, bf16])] + fn bench_rope_rows(dt: DType) -> BenchSetup { + let (n_heads, head_dim, n_nope, half_rot, n_tokens) = + (64usize, 512usize, 448usize, 32usize, 256usize); + BenchSetup::new(ffai_dsv4_partial_rope_rows::kernel_ir_for(dt)) + .mode(KernelMode::Grid3D) + .buffer(BenchBuffer::random("qk", n_tokens * n_heads * head_dim, dt)) + .buffer(BenchBuffer::zeros("out", n_tokens * n_heads * head_dim, dt).output()) + .constexpr("head_dim", head_dim as u32) + .constexpr("n_nope", n_nope as u32) + .constexpr("half_rot", half_rot as u32) + .constexpr("n_heads", n_heads as u32) + .constexpr("base_position", 0u32) + .constexpr("theta_base", 10_000.0_f32) + .constexpr("inverse_flag", 0u32) + .constexpr("freq_scale", 1.0f32) + .constexpr("ext_factor", 0.0f32) + .constexpr("corr_low", 0.0f32) + .constexpr("corr_high", 0.0f32) + .grid_3d(n_heads as u32, half_rot as u32, n_tokens as u32, [1, 1, 1]) + .bytes_moved((4 * n_tokens * n_heads * half_rot * dt.size_bytes()) as u64) + } +} diff --git a/crates/metaltile-std/src/ffai/dsv4_router_topk.rs b/crates/metaltile-std/src/ffai/dsv4_router_topk.rs new file mode 100644 index 00000000..6969ca70 --- /dev/null +++ b/crates/metaltile-std/src/ffai/dsv4_router_topk.rs @@ -0,0 +1,136 @@ +//! Copyright 2026 0xClandestine, Ekryski, TheTom, Ambisphaeric +//! SPDX-License-Identifier: Apache-2.0 +//! DeepSeek V4 router top-K selection — GPU-side, so the per-layer MoE +//! routing never round-trips to the CPU. +//! +//! DSv4 selects the top-K experts by the **biased** sqrtsoftplus score +//! and weights them by the **unbiased** score, renormalised to sum to 1. +//! This kernel emits the sum-to-1 weights only. `routed_scaling_factor` +//! is applied by the CALLER (`gpuRoutedFfnTail`) AFTER this renorm — +//! `final_weight_k = scaling * unbiased_k / Σ unbiased` — so it does NOT +//! cancel. Keeping it out of the kernel lets the kernel take just the two +//! score vectors: +//! +//! ```text +//! chosen = argmax-k(score_biased) (indices) +//! weight_i = score_unbiased[chosen_i] / Σ_j∈chosen score_unbiased[chosen_j] +//! ``` +//! +//! Mirrors the CPU selection in `forwardFfnSubblock` (top-K by biased, +//! weights from unbiased, sum-to-1). Structure adapted from +//! `mt_moe_router_topk` (masked argmax passes, one simdgroup/row). +//! +//! Dispatch: `[n_rows, 1, 1] × [32, 1, 1]` (pins tpg=32; one simdgroup +//! per token row — decode is n_rows=1). + +use metaltile::kernel; + +#[kernel] +pub fn mt_dsv4_router_topk( + score_biased: Tensor, + score_unbiased: Tensor, + mut indices_out: Tensor, + mut weights_out: Tensor, + #[constexpr] n_experts: u32, + #[constexpr] k: u32, +) { + let row = tgid_x; + let lane = tid; + let row_base = row * n_experts; + threadgroup_alloc("tg_chosen_idx", 32u32); + + // ── k masked-argmax passes over score_biased ───────────────────── + for it in range(0u32, k, 1u32) { + let mut best_val = neg_infinity(); + let mut best_idx = 0u32; + let n_per_lane = (n_experts + 31u32) / 32u32; + for r in range(0u32, n_per_lane, 1u32) { + let j = r * 32u32 + lane; + if j < n_experts { + let v = load(score_biased[row_base + j]).cast::(); + let mut chosen_mask = 0u32; + for p in range(0u32, it, 1u32) { + let cp = threadgroup_load("tg_chosen_idx", p); + chosen_mask = chosen_mask | select(j == cp, 1u32, 0u32); + } + let candidate = select(chosen_mask > 0u32, neg_infinity(), v); + let better = candidate > best_val; + best_val = select(better, candidate, best_val); + best_idx = select(better, j, best_idx); + } + } + let global_best_val = simd_max(best_val); + let i_have = best_val == global_best_val; + let my_idx_or_max = select(i_have, best_idx, 4294967295u32); + let global_best_idx = simd_min(my_idx_or_max); + if lane == 0u32 { + threadgroup_store("tg_chosen_idx", it, global_best_idx); + } + simdgroup_barrier_mem_none(); + } + + // ── weights = unbiased[chosen] / Σ unbiased[chosen] ────────────── + // tg_chosen_idx is f32-typed scratch; cast to u32 for the subscript. + let my_idx_f = select(lane < k, threadgroup_load("tg_chosen_idx", lane), 0.0f32); + let my_idx = my_idx_f.cast::(); + let my_unbiased = + select(lane < k, load(score_unbiased[row_base + my_idx]).cast::(), 0.0f32); + let sum_chosen = simd_sum(my_unbiased); + let weight = my_unbiased / sum_chosen; + if lane < k { + let out_base = row * k + lane; + store(indices_out[out_base], my_idx); + store(weights_out[out_base], weight.cast::()); + } +} + +/// out[i] = table[idx[i]] — gather a u32 table by u32 indices. Used to +/// remap raw routed expert ids into resident-pool packed slot ids on the +/// GPU (so the gather dispatch's `expert_ids` never touches the CPU). +#[kernel] +pub fn mt_remap_u32( + table: Tensor, + idx: Tensor, + mut out: Tensor, + #[constexpr] n: u32, +) { + let i = tid; + if i < n { + let e = load(idx[i]); + store(out[i], load(table[e])); + } +} + +pub mod kernel_benches { + use metaltile::{bench, test::*}; + + use super::{mt_dsv4_router_topk, mt_remap_u32}; + + #[bench(name = "ffai/moe/dsv4_router_topk", dtypes = [f32, f16, bf16])] + fn bench_dsv4_router_topk(dt: DType) -> BenchSetup { + let n_experts = 256usize; + let k = 6usize; + BenchSetup::new(mt_dsv4_router_topk::kernel_ir_for(dt)) + .mode(KernelMode::Reduction) + .buffer(BenchBuffer::random("score_biased", n_experts, dt)) + .buffer(BenchBuffer::random("score_unbiased", n_experts, dt)) + .buffer(BenchBuffer::zeros("indices_out", k, DType::U32).output()) + .buffer(BenchBuffer::zeros("weights_out", k, dt).output()) + .constexpr("n_experts", n_experts as u32) + .constexpr("k", k as u32) + .grid_3d(1, 1, 1, [32, 1, 1]) + .bytes_moved((2 * n_experts * dt.size_bytes()) as u64) + } + + #[bench(name = "ffai/moe/remap_u32", dtypes = [f32])] + fn bench_remap_u32(_dt: DType) -> BenchSetup { + let n = 6usize; + BenchSetup::new(mt_remap_u32::kernel_ir()) + .buffer(BenchBuffer::zeros("table", 256, DType::U32)) + .buffer(BenchBuffer::zeros("idx", n, DType::U32)) + .buffer(BenchBuffer::zeros("out", n, DType::U32).output()) + .constexpr("n", n as u32) + .grid_1d(n, 32) + .bytes_moved((n * 4 * 2) as u64) + } +} diff --git a/crates/metaltile-std/src/ffai/dsv4_swiglu_limit.rs b/crates/metaltile-std/src/ffai/dsv4_swiglu_limit.rs new file mode 100644 index 00000000..eaea2513 --- /dev/null +++ b/crates/metaltile-std/src/ffai/dsv4_swiglu_limit.rs @@ -0,0 +1,92 @@ +//! Copyright 2026 0xClandestine, Ekryski, TheTom, Ambisphaeric +//! SPDX-License-Identifier: Apache-2.0 +//! DSv4 SwiGLU with `swiglu_limit` clamp — `silu(min(gate, limit)) * +//! clip(up, -limit, +limit)`. DeepSeek-V4 trains the gated MLP / MoE +//! experts with this clamp (swiglu_limit=10); unclamped `silu(gate)*up` +//! overflows fp16 in the deep, high-magnitude layers and makes the +//! batched-prefill vs sequential-decode paths diverge (large values +//! round differently across kernels — the clamp pins them identical). +//! +//! Drop-in replacement for `mt_swiglu` on the DSv4 gate/up → inner path. + +use metaltile::kernel; + +#[kernel] +pub fn ffai_dsv4_swiglu_limit( + gate: Tensor, + up: Tensor, + out: Tensor, + #[constexpr] limit: f32, +) { + let idx = tid; + let g = load(gate[idx]).cast::(); + let u = load(up[idx]).cast::(); + // silu(min(gate, limit)) + let g_lim = select(g > limit, limit, g); + let s = g_lim / (1.0f32 + exp(0.0f32 - g_lim)); + // clip(up, -limit, +limit) + let neg = 0.0f32 - limit; + let u_lo = select(u < neg, neg, u); + let u_lim = select(u_lo > limit, limit, u_lo); + store(out[idx], (s * u_lim).cast::()); +} + +pub mod kernel_tests { + use metaltile::{test::*, test_kernel}; + + use super::ffai_dsv4_swiglu_limit; + use crate::utils::{pack_f32, unpack_f32}; + + fn setup(n: usize, dt: DType) -> TestSetup { + let limit = 10.0f32; + // Span values that exercise both clamp branches (|x| > 10). + let gate: Vec = (0..n).map(|i| (i % 41) as f32 * 0.8 - 16.0).collect(); + let up: Vec = (0..n).map(|i| (i % 37) as f32 * 0.9 - 16.0).collect(); + let g_dt = unpack_f32(&pack_f32(&gate, dt), dt); + let u_dt = unpack_f32(&pack_f32(&up, dt), dt); + let expected: Vec = g_dt + .iter() + .zip(&u_dt) + .map(|(&g, &u)| { + let gl = g.min(limit); + let s = gl / (1.0 + (-gl).exp()); + let ul = u.clamp(-limit, limit); + s * ul + }) + .collect(); + TestSetup::new(ffai_dsv4_swiglu_limit::kernel_ir_for(dt)) + .input(TestBuffer::from_vec("gate", pack_f32(&gate, dt), dt)) + .input(TestBuffer::from_vec("up", pack_f32(&up, dt), dt)) + .input(TestBuffer::zeros("out", n, dt)) + .constexpr("limit", limit) + .expect(TestBuffer::from_vec("out", pack_f32(&expected, dt), dt)) + .grid_1d(n, 256) + } + + // f16/bf16 tol mirrors the sister gated-product kernel + // `moe_down_swiglu_accum` (also `silu(gate)*up`). The clamp pins the + // product to |·| ≤ limit² (=100), and `silu` routes through GPU + // fast-`exp`, so the f16 output round lands ~15 ulp off the libm CPU + // oracle (err 7.8e-3 on M5; wider on older fast-math). f32 stays tight + // at 1e-4 — it proves the logic; the half tols absorb fast-math+round. + #[test_kernel(dtypes = [f32, f16, bf16], tol = [1e-4, 5e-2, 2e-1])] + fn test_ffai_dsv4_swiglu_limit(dt: DType) -> TestSetup { setup(1024, dt) } +} + +pub mod kernel_benches { + use metaltile::{bench, test::*}; + + use super::ffai_dsv4_swiglu_limit; + + #[bench(name = "ffai/swiglu/dsv4_swiglu_limit", dtypes = [f32, f16, bf16])] + fn bench_dsv4_swiglu_limit(dt: DType) -> BenchSetup { + let n = 1024 * 1024usize; + BenchSetup::new(ffai_dsv4_swiglu_limit::kernel_ir_for(dt)) + .buffer(BenchBuffer::random("gate", n, dt)) + .buffer(BenchBuffer::random("up", n, dt)) + .buffer(BenchBuffer::zeros("out", n, dt).output()) + .constexpr("limit", 10.0f32) + .grid_1d(n, 256) + .bytes_moved((3 * n * dt.size_bytes()) as u64) + } +} diff --git a/crates/metaltile-std/src/ffai/gate_up_swiglu_fused.rs b/crates/metaltile-std/src/ffai/gate_up_swiglu_fused.rs new file mode 100644 index 00000000..9956f521 --- /dev/null +++ b/crates/metaltile-std/src/ffai/gate_up_swiglu_fused.rs @@ -0,0 +1,99 @@ +//! Copyright 2026 0xClandestine, Ekryski, TheTom, Ambisphaeric +//! SPDX-License-Identifier: Apache-2.0 +//! Fused per-expert (gate gemv + up gemv + SwiGLU): `inner[r] = +//! silu(gate_w[r] · x) * (up_w[r] · x)`. Single dispatch replaces +//! gate_gemv → up_gemv → swiglu (3 dispatches). +//! +//! For MoE decode where each routed expert needs gate/up before +//! down, this fuses the 3-dispatch prelude into one. Saves 2 +//! commands per expert × top-K per layer. +//! +//! Geometry: one threadgroup per inner row (`intermediate = mat +//! shape[0]`). Each TG does TWO `strided_reduce_dot` + `reduce_sum` +//! passes (one for gate, one for up), then a per-row silu+mul + +//! store of inner. + +use metaltile::kernel; + +#[kernel] +pub fn ffai_gate_up_swiglu_fused( + gate_w: Tensor, + up_w: Tensor, + x: Tensor, + mut inner: Tensor, + #[constexpr] k: u32, +) { + let row = program_id::<0>(); + let rs = row * k; + let re = rs + k; + let g_dot = strided_reduce_dot(gate_w, x, rs, rs, re); + let g_full = reduce_sum(g_dot); + let u_dot = strided_reduce_dot(up_w, x, rs, rs, re); + let u_full = reduce_sum(u_dot); + // Match the unfused gate-gemv → silu → mul chain's precision: + // gate-gemv stores its result as T (f16/bf16 round-trip) before + // mt_swiglu's silu reads it. Round g + u back to T then to f32 so + // accumulated drift in long decode runs matches greedy on the + // unfused path exactly. + let g_t = g_full.cast::(); + let u_t = u_full.cast::(); + let g = g_t.cast::(); + let u = u_t.cast::(); + store(inner[row], silu(g) * u); +} + +pub mod kernel_tests { + use metaltile::{test::*, test_kernel}; + + use super::ffai_gate_up_swiglu_fused; + use crate::utils::{pack_f32, unpack_f32}; + + fn setup(m: usize, k: usize, dt: DType) -> TestSetup { + let gate: Vec = (0..m * k).map(|i| ((i % 17) as f32 - 8.0) * 0.01).collect(); + let up: Vec = (0..m * k).map(|i| ((i % 13) as f32 - 6.0) * 0.02).collect(); + let x: Vec = (0..k).map(|j| ((j % 11) as f32 - 5.0) * 0.03).collect(); + let g_dt = unpack_f32(&pack_f32(&gate, dt), dt); + let u_dt = unpack_f32(&pack_f32(&up, dt), dt); + let x_dt = unpack_f32(&pack_f32(&x, dt), dt); + let expected: Vec = (0..m) + .map(|r| { + let g_val: f32 = (0..k).map(|j| g_dt[r * k + j] * x_dt[j]).sum(); + let u_val: f32 = (0..k).map(|j| u_dt[r * k + j] * x_dt[j]).sum(); + let sig = g_val / (1.0 + (-g_val).exp()); + sig * u_val + }) + .collect(); + TestSetup::new(ffai_gate_up_swiglu_fused::kernel_ir_for(dt)) + .mode(KernelMode::Reduction) + .input(TestBuffer::from_vec("gate_w", pack_f32(&gate, dt), dt)) + .input(TestBuffer::from_vec("up_w", pack_f32(&up, dt), dt)) + .input(TestBuffer::from_vec("x", pack_f32(&x, dt), dt)) + .input(TestBuffer::zeros("inner", m, dt)) + .constexpr("k", k as u32) + .expect(TestBuffer::from_vec("inner", pack_f32(&expected, dt), dt)) + .grid_3d(m as u32, 1, 1, [256, 1, 1]) + } + + #[test_kernel(dtypes = [f32, f16, bf16], tol = [1e-3, 1e-1, 1.0])] + fn test_gus_small(dt: DType) -> TestSetup { setup(16, 256, dt) } +} + +pub mod kernel_benches { + use metaltile::{bench, test::*}; + + use super::ffai_gate_up_swiglu_fused; + + #[bench(name = "ffai/gate_up_swiglu_fused", dtypes = [f32, f16, bf16])] + fn bench_gus(dt: DType) -> BenchSetup { + let (m, k) = (2048usize, 4096usize); + BenchSetup::new(ffai_gate_up_swiglu_fused::kernel_ir_for(dt)) + .mode(KernelMode::Reduction) + .buffer(BenchBuffer::random("gate_w", m * k, dt)) + .buffer(BenchBuffer::random("up_w", m * k, dt)) + .buffer(BenchBuffer::random("x", k, dt)) + .buffer(BenchBuffer::zeros("inner", m, dt).output()) + .constexpr("k", k as u32) + .grid_3d(m as u32, 1, 1, [256, 1, 1]) + .bytes_moved((2 * m * k * dt.size_bytes() + k * dt.size_bytes()) as u64) + } +} diff --git a/crates/metaltile-std/src/ffai/gemm_q8.rs b/crates/metaltile-std/src/ffai/gemm_q8.rs new file mode 100644 index 00000000..292e4d4c --- /dev/null +++ b/crates/metaltile-std/src/ffai/gemm_q8.rs @@ -0,0 +1,207 @@ +//! Copyright 2026 0xClandestine, Ekryski, TheTom, Ambisphaeric +//! SPDX-License-Identifier: Apache-2.0 +//! Q8_0 multi-row GEMM — `out[r, :] = dequant(weight) · input[r, :]` for a +//! block of `n_rows` rows in one dispatch. The prefill counterpart of +//! `ffai_gemv_q8`: tiles the output into 32×32 blocks and stages a +//! `[32,16]` dequantized-weight tile + `[32,16]` input tile in threadgroup +//! memory, so each Q8 weight is read+dequanted ONCE and reused across all +//! 32 rows of the tile (the amortization that makes prefill fast — a +//! single-row gemv re-streams the weight per token). +//! +//! Weight is the resident Q8 split (`qs` int8 packed 4/u32 + per-32-block +//! `d` scale), laid out as `[out_dim, in_dim]` (row-major over values). +//! Mirrors `ffai_gemm`'s geometry exactly so the dispatch wrapper is +//! identical apart from the two extra weight buffers. +//! +//! ## DISPATCH INVARIANTS (same as ffai_gemm) +//! - TPG = 1024 (32×32). Grid: (out_dim/32) × (n_rows/32) threadgroups. +//! - `in_dim % 16 == 0` (K-tile) AND `in_dim % 32 == 0` (Q8 block). +//! - Row/col edges handled in-kernel (clamp loads to 0, skip OOB stores). + +use metaltile::kernel; + +#[kernel] +pub fn ffai_gemm_q8( + qs: Tensor, + d_f32: Tensor, + input: Tensor, + out: Tensor, + #[constexpr] in_dim: u32, + #[constexpr] out_dim: u32, + #[constexpr] n_rows: u32, +) { + let tid = simd_id * 32u32 + simd_lane; + let lr = tid / 32u32; // output row within tile (0..31) + let lo = tid % 32u32; // output col within tile (0..31) + threadgroup_alloc("gq8_w", 512); + threadgroup_alloc("gq8_x", 512); + let mut acc = 0.0f32; + for k0 in range(0u32, in_dim, 16u32) { + // threads 0..511 dequant one weight element each into the W tile. + if tid < 512u32 { + let s = tid; + let w_col = tgid_x * 32u32 + s / 16u32; + let w_valid = w_col < out_dim; + let w_col_safe = select(w_valid, w_col, 0u32); + // value index in the [out_dim, in_dim] weight, then Q8 unpack. + let vidx = w_col_safe * in_dim + k0 + s % 16u32; + let block = vidx / 32u32; + let lane = vidx & 31u32; + let word = load(qs[block * 8u32 + lane / 4u32]); + let by = (word >> ((lane & 3u32) * 8u32)) & 0xffu32; + let qf = by.cast::() - select(by > 127u32, 256.0f32, 0.0f32); + let w_raw = load(d_f32[block]) * qf; + threadgroup_store("gq8_w", s, select(w_valid, w_raw, 0.0f32)); + } + // threads 512..1023 load one input element each into the X tile. + if tid >= 512u32 { + let s = tid - 512u32; + let x_row = tgid_y * 32u32 + s / 16u32; + let x_valid = x_row < n_rows; + let x_row_safe = select(x_valid, x_row, 0u32); + let x_raw = load(input[x_row_safe * in_dim + k0 + s % 16u32]).cast::(); + threadgroup_store("gq8_x", s, select(x_valid, x_raw, 0.0f32)); + } + threadgroup_barrier(); + for k in range(0u32, 16u32, 1u32) { + let w = threadgroup_load("gq8_w", lo * 16u32 + k); + let x = threadgroup_load("gq8_x", lr * 16u32 + k); + acc = acc + w * x; + } + threadgroup_barrier(); + } + let r = tgid_y * 32u32 + lr; + let o = tgid_x * 32u32 + lo; + if r < n_rows { + if o < out_dim { + store(out[r * out_dim + o], acc.cast::()); + } + } +} + +/// GROUPED Q8 GEMM — the amortized fix for prefill O-LoRA-A. Same tiled +/// 32×32 / threadgroup-staged structure as `ffai_gemm_q8` (weight read+ +/// dequanted ONCE, reused across 32 rows), but the input is GROUPED: output +/// column `o` belongs to group `g = o / rows_per_group`, and group g reads +/// input columns `[g*in_dim, (g+1)*in_dim)` of a row that is `n_groups*in_dim` +/// wide. Replaces `ffai_grouped_gemv_q8_rows` (per-token gemv, ~47 ms/layer, +/// the #1 prefill attention hotspot) with O-LoRA-B-class amortized speed. +/// A 32-wide output tile lies entirely within one group (rows_per_group % +/// 32 == 0), so `g` is uniform per threadgroup → no per-thread divergence. +/// weight `[out_dim, in_dim]`; input `[n_rows, n_groups*in_dim]`; +/// out `[n_rows, out_dim]`. Grid + TPG identical to ffai_gemm_q8. +#[kernel] +#[allow(clippy::too_many_arguments)] +pub fn ffai_grouped_gemm_q8( + qs: Tensor, + d_f32: Tensor, + input: Tensor, + out: Tensor, + #[constexpr] in_dim: u32, + #[constexpr] out_dim: u32, + #[constexpr] n_rows: u32, + #[constexpr] n_groups: u32, + #[constexpr] rows_per_group: u32, +) { + let tid = simd_id * 32u32 + simd_lane; + let lr = tid / 32u32; + let lo = tid % 32u32; + // Group is uniform per threadgroup (tile cols all within one group). + let g = (tgid_x * 32u32) / rows_per_group; + let row_in_stride = n_groups * in_dim; + let in_col_off = g * in_dim; + threadgroup_alloc("gq8_w", 512); + threadgroup_alloc("gq8_x", 512); + let mut acc = 0.0f32; + for k0 in range(0u32, in_dim, 16u32) { + if tid < 512u32 { + let s = tid; + let w_col = tgid_x * 32u32 + s / 16u32; + let w_valid = w_col < out_dim; + let w_col_safe = select(w_valid, w_col, 0u32); + let vidx = w_col_safe * in_dim + k0 + s % 16u32; + let block = vidx / 32u32; + let lane = vidx & 31u32; + let word = load(qs[block * 8u32 + lane / 4u32]); + let by = (word >> ((lane & 3u32) * 8u32)) & 0xffu32; + let qf = by.cast::() - select(by > 127u32, 256.0f32, 0.0f32); + let w_raw = load(d_f32[block]) * qf; + threadgroup_store("gq8_w", s, select(w_valid, w_raw, 0.0f32)); + } + if tid >= 512u32 { + let s = tid - 512u32; + let x_row = tgid_y * 32u32 + s / 16u32; + let x_valid = x_row < n_rows; + let x_row_safe = select(x_valid, x_row, 0u32); + // GROUPED input read: row is n_groups*in_dim wide; this group's + // slice starts at g*in_dim. + let x_raw = + load(input[x_row_safe * row_in_stride + in_col_off + k0 + s % 16u32]).cast::(); + threadgroup_store("gq8_x", s, select(x_valid, x_raw, 0.0f32)); + } + threadgroup_barrier(); + for k in range(0u32, 16u32, 1u32) { + let w = threadgroup_load("gq8_w", lo * 16u32 + k); + let x = threadgroup_load("gq8_x", lr * 16u32 + k); + acc = acc + w * x; + } + threadgroup_barrier(); + } + let r = tgid_y * 32u32 + lr; + let o = tgid_x * 32u32 + lo; + if r < n_rows { + if o < out_dim { + store(out[r * out_dim + o], acc.cast::()); + } + } +} + +pub mod kernel_benches { + use metaltile::{bench, test::*}; + + use super::{ffai_gemm_q8, ffai_grouped_gemm_q8}; + + #[bench(name = "ffai/gemm/gemm_q8", dtypes = [f32, f16, bf16])] + fn bench_gemm_q8(dt: DType) -> BenchSetup { + let in_dim = 4096usize; + let out_dim = 2048usize; + let n_rows = 256usize; + let n_blocks = out_dim * in_dim / 32; + BenchSetup::new(ffai_gemm_q8::kernel_ir_for(dt)) + .mode(KernelMode::Reduction) + .buffer(BenchBuffer::random("qs", n_blocks * 8, DType::U32)) + .buffer(BenchBuffer::random("d_f32", n_blocks, DType::F32)) + .buffer(BenchBuffer::random("input", n_rows * in_dim, dt)) + .buffer(BenchBuffer::zeros("out", n_rows * out_dim, dt).output()) + .constexpr("in_dim", in_dim as u32) + .constexpr("out_dim", out_dim as u32) + .constexpr("n_rows", n_rows as u32) + .grid_3d((out_dim as u32).div_ceil(32), (n_rows as u32).div_ceil(32), 1, [1024, 1, 1]) + .bytes_moved((n_blocks * 36 + n_rows * in_dim * dt.size_bytes()) as u64) + } + + #[bench(name = "ffai/gemm/grouped_gemm_q8", dtypes = [f32, f16, bf16])] + fn bench_grouped_gemm_q8(dt: DType) -> BenchSetup { + // O-LoRA-A shape: 8 groups, per-group in_dim=4096 (nHeads*headDim/8), + // out_dim=8192 (8*1024), n_rows=256 tokens. + let in_dim = 4096usize; + let out_dim = 8192usize; + let n_rows = 256usize; + let n_groups = 8usize; + let rows_per_group = out_dim / n_groups; // 1024 + let n_blocks = out_dim * in_dim / 32; + BenchSetup::new(ffai_grouped_gemm_q8::kernel_ir_for(dt)) + .mode(KernelMode::Reduction) + .buffer(BenchBuffer::random("qs", n_blocks * 8, DType::U32)) + .buffer(BenchBuffer::random("d_f32", n_blocks, DType::F32)) + .buffer(BenchBuffer::random("input", n_rows * n_groups * in_dim, dt)) + .buffer(BenchBuffer::zeros("out", n_rows * out_dim, dt).output()) + .constexpr("in_dim", in_dim as u32) + .constexpr("out_dim", out_dim as u32) + .constexpr("n_rows", n_rows as u32) + .constexpr("n_groups", n_groups as u32) + .constexpr("rows_per_group", rows_per_group as u32) + .grid_3d((out_dim as u32).div_ceil(32), (n_rows as u32).div_ceil(32), 1, [1024, 1, 1]) + .bytes_moved((n_blocks * 36 + n_rows * n_groups * in_dim * dt.size_bytes()) as u64) + } +} diff --git a/crates/metaltile-std/src/ffai/gemm_q8_mpp.rs b/crates/metaltile-std/src/ffai/gemm_q8_mpp.rs new file mode 100644 index 00000000..f5234c69 --- /dev/null +++ b/crates/metaltile-std/src/ffai/gemm_q8_mpp.rs @@ -0,0 +1,260 @@ +//! Copyright 2026 0xClandestine, Ekryski, TheTom, Ambisphaeric +//! SPDX-License-Identifier: Apache-2.0 +//! Dense Q8_0 GEMM via cooperative-tensor MMA — the MMA replacement for the +//! SCALAR `ffai_gemm_q8` (which does `acc += w*x` on the ALUs, ~12× slower +//! than a simdgroup-MMA Q8 matmul). Used for the attention Q/KV/q_a/q_b +//! projections and the O-LoRA-B / shared-expert dense Q8 matmuls in prefill. +//! +//! Same 64×64×32 coop_tile geometry as `moe_bgemm_iq2xxs_bm64` (4 simdgroups, +//! 2×2 warp grid, 128 threads/tg) but with the Q8_0 dequant and NO expert +//! routing — every output row is the same dense weight, so the bm64 +//! sub-run/expert-boundary machinery collapses to a straight K-loop. +//! +//! Weight `[out_dim, k_in]` Q8_0 (qs int8 packed 4/u32 + per-32-block `d`). +//! x `[n_rows, k_in]`; out `[n_rows, out_dim]`. out[r,o]=Σ_k W[o,k]·x[r,k]. +//! Name contains `_mpp_` so FFAI's PSOCache LIVE-COMPILES it (the offline +//! metallib lowers cooperative-tensor MMA incorrectly — see PSOCache.isMppKernel). +//! +//! grid (threadgroups) = [ceil(out_dim/64), ceil(n_rows/64), 1], tg [128,1,1]. + +use metaltile::kernel; + +#[kernel] +pub fn ffai_gemm_q8_mpp( + x: Tensor, + qs: Tensor, + d_f32: Tensor, + mut out: Tensor, + #[constexpr] n_rows: u32, + #[constexpr] out_dim: u32, + #[constexpr] k_in: u32, +) { + let n_tile_base = tgid_x * 64u32; // output-feature tile (N dim) + let m_tile_base = tgid_y * 64u32; // token tile (M dim) + let sg = simd_group_id(); + let lane_in_tg = sg * 32u32 + simd_lane; + let sg_m_base = (sg / 2u32) * 32u32; + let sg_n_base = (sg & 1u32) * 32u32; + let x_m_row = lane_in_tg / 2u32; + let x_k_base = (lane_in_tg & 1u32) * 16u32; + threadgroup_alloc("Xs", 2048, coop_stage(T)); // 64 × 32 + threadgroup_alloc("Ws", 2048, coop_stage(T)); // 64 × 32 + threadgroup_alloc("OutScratch", 4096, f32); // 4 SG × 32 × 32 + coop_tile_setup( + "gemm", + 32, + 32, + 32, + coop_stage(T), + "accumulate", + "simdgroup", + f32, + false, + true, + false, + ); + coop_tile_zero("gemm"); + for kb in range(0u32, k_in, 32u32) { + // Stage X[m_tile_base..+64, kb..kb+32] → Xs. 128 lanes × 16. + let gr_x = m_tile_base + x_m_row; + let in_run_x = gr_x < n_rows; + let safe_gr_x = select(in_run_x, gr_x, 0u32); + let x_dev_base = safe_gr_x * k_in + kb + x_k_base; + let x_ws_base = x_m_row * 32u32 + x_k_base; + for _i in range(0u32, 16u32, 1u32) { + let xv = load(x[x_dev_base + _i]).cast::(); + threadgroup_store("Xs", x_ws_base + _i, select(in_run_x, xv, 0.0f32)); + } + // Dequant W[n_tile_base..+64, kb..kb+32] → Ws via Q8_0. 128 lanes × 16. + for _i in range(0u32, 16u32, 1u32) { + let flat = lane_in_tg * 16u32 + _i; + let w_row = flat / 32u32; // 0..63 (output feature within tile) + let k_local = flat & 31u32; // 0..31 (BK) + let global_col = n_tile_base + w_row; + let k = kb + k_local; + let vidx = global_col * k_in + k; + let block = vidx / 32u32; + let lane = vidx & 31u32; + let word = load(qs[block * 8u32 + lane / 4u32]); + let by = (word >> ((lane & 3u32) * 8u32)) & 0xffu32; + let qf = by.cast::() - select(by > 127u32, 256.0f32, 0.0f32); + let w = (load(d_f32[block]) * qf).cast::().cast::(); + threadgroup_store("Ws", w_row * 32u32 + k_local, w); + } + threadgroup_barrier(); + coop_tile_load_a("gemm", "Xs", true, coop_stage(T), 32, 32, sg_m_base * 32u32); + coop_tile_load_b("gemm", "Ws", true, coop_stage(T), 32, 32, sg_n_base * 32u32); + coop_tile_run("gemm"); + threadgroup_barrier(); + } + coop_tile_store_c("gemm", "OutScratch", true, f32, 32, 32, sg * 1024u32); + threadgroup_barrier(); + for _e in range(0u32, 32u32, 1u32) { + let flat = lane_in_tg * 32u32 + _e; + let mr = flat / 64u32; + let nc = flat & 63u32; + let gr = m_tile_base + mr; + let gc = n_tile_base + nc; + let in_run = (gr < n_rows) & (gc < out_dim); + if in_run { + let src_sg = (mr / 32u32) * 2u32 + nc / 32u32; + let v = threadgroup_load( + "OutScratch", + src_sg * 1024u32 + (mr & 31u32) * 32u32 + (nc & 31u32), + ); + store(out[gr * out_dim + gc], v.cast::()); + } + } +} + +/// GROUPED dense Q8 GEMM via cooperative-tensor MMA — the MMA version of +/// `ffai_grouped_gemm_q8` (which is scalar). For prefill O-LoRA-A: output +/// column o belongs to group g = o/rows_per_group, and group g reads input +/// columns [g*k_in, (g+1)*k_in) of an (n_groups*k_in)-wide activation row. +/// A 64-wide output tile is uniform-group (rows_per_group % 64 == 0), so the +/// group offset is computed once per threadgroup. Same 64×64×32 coop_tile +/// MMA as ffai_gemm_q8_mpp. Live-compiled (name has _mpp_). +#[kernel] +#[allow(clippy::too_many_arguments)] +pub fn ffai_grouped_gemm_q8_mpp( + x: Tensor, + qs: Tensor, + d_f32: Tensor, + mut out: Tensor, + #[constexpr] n_rows: u32, + #[constexpr] out_dim: u32, + #[constexpr] k_in: u32, + #[constexpr] n_groups: u32, + #[constexpr] rows_per_group: u32, +) { + let n_tile_base = tgid_x * 64u32; + let m_tile_base = tgid_y * 64u32; + let sg = simd_group_id(); + let lane_in_tg = sg * 32u32 + simd_lane; + let sg_m_base = (sg / 2u32) * 32u32; + let sg_n_base = (sg & 1u32) * 32u32; + let x_m_row = lane_in_tg / 2u32; + let x_k_base = (lane_in_tg & 1u32) * 16u32; + // Group is uniform per threadgroup (64-wide feature tile within one group). + let g = n_tile_base / rows_per_group; + let row_in_stride = n_groups * k_in; + let in_col_off = g * k_in; + threadgroup_alloc("Xs", 2048, coop_stage(T)); + threadgroup_alloc("Ws", 2048, coop_stage(T)); + threadgroup_alloc("OutScratch", 4096, f32); + coop_tile_setup( + "gemm", + 32, + 32, + 32, + coop_stage(T), + "accumulate", + "simdgroup", + f32, + false, + true, + false, + ); + coop_tile_zero("gemm"); + for kb in range(0u32, k_in, 32u32) { + let gr_x = m_tile_base + x_m_row; + let in_run_x = gr_x < n_rows; + let safe_gr_x = select(in_run_x, gr_x, 0u32); + // GROUPED input read: row is n_groups*k_in wide; this group's slice + // starts at g*k_in. + let x_dev_base = safe_gr_x * row_in_stride + in_col_off + kb + x_k_base; + let x_ws_base = x_m_row * 32u32 + x_k_base; + for _i in range(0u32, 16u32, 1u32) { + let xv = load(x[x_dev_base + _i]).cast::(); + threadgroup_store("Xs", x_ws_base + _i, select(in_run_x, xv, 0.0f32)); + } + for _i in range(0u32, 16u32, 1u32) { + let flat = lane_in_tg * 16u32 + _i; + let w_row = flat / 32u32; + let k_local = flat & 31u32; + let global_col = n_tile_base + w_row; + let k = kb + k_local; + let vidx = global_col * k_in + k; + let block = vidx / 32u32; + let lane = vidx & 31u32; + let word = load(qs[block * 8u32 + lane / 4u32]); + let by = (word >> ((lane & 3u32) * 8u32)) & 0xffu32; + let qf = by.cast::() - select(by > 127u32, 256.0f32, 0.0f32); + let w = (load(d_f32[block]) * qf).cast::().cast::(); + threadgroup_store("Ws", w_row * 32u32 + k_local, w); + } + threadgroup_barrier(); + coop_tile_load_a("gemm", "Xs", true, coop_stage(T), 32, 32, sg_m_base * 32u32); + coop_tile_load_b("gemm", "Ws", true, coop_stage(T), 32, 32, sg_n_base * 32u32); + coop_tile_run("gemm"); + threadgroup_barrier(); + } + coop_tile_store_c("gemm", "OutScratch", true, f32, 32, 32, sg * 1024u32); + threadgroup_barrier(); + for _e in range(0u32, 32u32, 1u32) { + let flat = lane_in_tg * 32u32 + _e; + let mr = flat / 64u32; + let nc = flat & 63u32; + let gr = m_tile_base + mr; + let gc = n_tile_base + nc; + let in_run = (gr < n_rows) & (gc < out_dim); + if in_run { + let src_sg = (mr / 32u32) * 2u32 + nc / 32u32; + let v = threadgroup_load( + "OutScratch", + src_sg * 1024u32 + (mr & 31u32) * 32u32 + (nc & 31u32), + ); + store(out[gr * out_dim + gc], v.cast::()); + } + } +} + +pub mod kernel_benches { + use metaltile::{bench, test::*}; + + use super::{ffai_gemm_q8_mpp, ffai_grouped_gemm_q8_mpp}; + + // q_a-like shape: in=4096, out=1024, 256 tokens. + #[bench(name = "ffai/gemm/gemm_q8_mpp", dtypes = [f32, f16, bf16])] + fn bench_gemm_q8_mpp(dt: DType) -> BenchSetup { + let n_rows = 256usize; + let out_dim = 1024usize; + let k_in = 4096usize; + let n_blocks = out_dim * k_in / 32; + BenchSetup::new(ffai_gemm_q8_mpp::kernel_ir_for(dt)) + .mode(KernelMode::Reduction) + .buffer(BenchBuffer::random("x", n_rows * k_in, dt)) + .buffer(BenchBuffer::random("qs", n_blocks * 8, DType::U32)) + .buffer(BenchBuffer::random("d_f32", n_blocks, DType::F32)) + .buffer(BenchBuffer::zeros("out", n_rows * out_dim, dt).output()) + .constexpr("n_rows", n_rows as u32) + .constexpr("out_dim", out_dim as u32) + .constexpr("k_in", k_in as u32) + .grid_3d((out_dim as u32).div_ceil(64), (n_rows as u32).div_ceil(64), 1, [128, 1, 1]) + .bytes_moved((n_blocks * 36 + n_rows * k_in * dt.size_bytes()) as u64) + } + + // O-LoRA-A shape: 8 groups, per-group in=4096, out=8192, 256 tokens. + #[bench(name = "ffai/gemm/grouped_gemm_q8_mpp", dtypes = [f32, f16, bf16])] + fn bench_grouped_gemm_q8_mpp(dt: DType) -> BenchSetup { + let n_rows = 256usize; + let out_dim = 8192usize; + let k_in = 4096usize; + let n_groups = 8usize; + let rows_per_group = out_dim / n_groups; + let n_blocks = out_dim * k_in / 32; + BenchSetup::new(ffai_grouped_gemm_q8_mpp::kernel_ir_for(dt)) + .mode(KernelMode::Reduction) + .buffer(BenchBuffer::random("x", n_rows * n_groups * k_in, dt)) + .buffer(BenchBuffer::random("qs", n_blocks * 8, DType::U32)) + .buffer(BenchBuffer::random("d_f32", n_blocks, DType::F32)) + .buffer(BenchBuffer::zeros("out", n_rows * out_dim, dt).output()) + .constexpr("n_rows", n_rows as u32) + .constexpr("out_dim", out_dim as u32) + .constexpr("k_in", k_in as u32) + .constexpr("n_groups", n_groups as u32) + .constexpr("rows_per_group", rows_per_group as u32) + .grid_3d((out_dim as u32).div_ceil(64), (n_rows as u32).div_ceil(64), 1, [128, 1, 1]) + .bytes_moved((n_blocks * 36 + n_rows * n_groups * k_in * dt.size_bytes()) as u64) + } +} diff --git a/crates/metaltile-std/src/ffai/gemv_axpy_inplace.rs b/crates/metaltile-std/src/ffai/gemv_axpy_inplace.rs new file mode 100644 index 00000000..49f2444c --- /dev/null +++ b/crates/metaltile-std/src/ffai/gemv_axpy_inplace.rs @@ -0,0 +1,102 @@ +//! Copyright 2026 0xClandestine, Ekryski, TheTom, Ambisphaeric +//! SPDX-License-Identifier: Apache-2.0 +//! Generic GEMV with weighted in-place axpy: `accum[r] += weight * (mat[r] · vec)`. +//! +//! Replaces the 3-kernel `out = gemv(mat, vec); scaled = out * weight_T; +//! accum += scaled` chain with a single dispatch. Aimed at the MoE +//! routed-expert combine step where weight is the normalised top-K +//! scalar and accum is moeAccum. +//! +//! Same reduction geometry as `mt_gemv` (one threadgroup per row, +//! simdgroup reduce over k). The only differences: +//! - `accum` instead of `out` (load + add + store, not just store) +//! - extra `weight: f32` constexpr scalar applied before the add +//! +//! `weight` is set via `setBytes` in the generated Swift binding +//! (constexpr scalars become runtime args at dispatch time per the +//! macro's emit path), so distinct dispatches with different +//! weights see their own value. + +use metaltile::kernel; + +#[kernel] +pub fn ffai_gemv_axpy_inplace( + mat: Tensor, + vec: Tensor, + mut accum: Tensor, + #[constexpr] k: u32, + #[constexpr] weight: f32, +) { + let row = program_id::<0>(); + let rs = row * k; + let re = rs + k; + let dot = strided_reduce_dot(mat, vec, rs, rs, re); + let result = reduce_sum(dot); + let prev = load(accum[row]).cast::(); + store(accum[row], prev + weight * result); +} + +pub mod kernel_tests { + use metaltile::{test::*, test_kernel}; + + use super::ffai_gemv_axpy_inplace; + use crate::utils::{pack_f32, unpack_f32}; + + fn setup(m: usize, k: usize, weight: f32, dt: DType) -> TestSetup { + let mat: Vec = (0..m * k).map(|i| ((i % 17) as f32 - 8.0) * 0.01).collect(); + let vec: Vec = (0..k).map(|j| ((j % 13) as f32 - 6.0) * 0.02).collect(); + let accum_in: Vec = (0..m).map(|r| (r as f32 % 5.0) * 0.1 - 0.2).collect(); + let mat_dt = unpack_f32(&pack_f32(&mat, dt), dt); + let vec_dt = unpack_f32(&pack_f32(&vec, dt), dt); + let accum_dt = unpack_f32(&pack_f32(&accum_in, dt), dt); + let expected: Vec = (0..m) + .map(|r| { + let dot: f32 = (0..k).map(|j| mat_dt[r * k + j] * vec_dt[j]).sum(); + accum_dt[r] + weight * dot + }) + .collect(); + TestSetup::new(ffai_gemv_axpy_inplace::kernel_ir_for(dt)) + .mode(KernelMode::Reduction) + .input(TestBuffer::from_vec("mat", pack_f32(&mat, dt), dt)) + .input(TestBuffer::from_vec("vec", pack_f32(&vec, dt), dt)) + .input(TestBuffer::from_vec("accum", pack_f32(&accum_in, dt), dt)) + .constexpr("k", k as u32) + .constexpr("weight", weight) + .expect(TestBuffer::from_vec("accum", pack_f32(&expected, dt), dt)) + .grid_3d(m as u32, 1, 1, [256, 1, 1]) + } + + #[test_kernel(dtypes = [f32, f16, bf16], tol = [1e-3, 1e-1, 1.0])] + fn test_axpy_small(dt: DType) -> TestSetup { setup(16, 256, 0.42, dt) } + + // f32 tol 5e-3 (vs 1e-3 on the k=256 tests): k=2048 here, and the + // GPU does a simdgroup tree-reduce while the CPU oracle sums + // sequentially. Over 2048 mixed-sign terms the f32 reorder gap grows + // with k (≈3.4e-3 on M5) — a reduction-order artifact, not a logic + // error. The f16/bf16 tols were already widened for this shape. + #[test_kernel(dtypes = [f32, f16, bf16], tol = [5e-3, 2e-1, 2.0])] + fn test_axpy_decode_shape(dt: DType) -> TestSetup { setup(4096, 2048, 0.167, dt) } + + #[test_kernel(dtypes = [f32, f16, bf16], tol = [1e-3, 1e-1, 1.0])] + fn test_axpy_neg_weight(dt: DType) -> TestSetup { setup(16, 256, -0.3, dt) } +} + +pub mod kernel_benches { + use metaltile::{bench, test::*}; + + use super::ffai_gemv_axpy_inplace; + + #[bench(name = "ffai/gemv_axpy_inplace", dtypes = [f32, f16, bf16])] + fn bench_gemv_axpy(dt: DType) -> BenchSetup { + let (m, k) = (4096usize, 2048usize); + BenchSetup::new(ffai_gemv_axpy_inplace::kernel_ir_for(dt)) + .mode(KernelMode::Reduction) + .buffer(BenchBuffer::random("mat", m * k, dt)) + .buffer(BenchBuffer::random("vec", k, dt)) + .buffer(BenchBuffer::random("accum", m, dt).output()) + .constexpr("k", k as u32) + .constexpr("weight", 0.167f32) + .grid_3d(m as u32, 1, 1, [256, 1, 1]) + .bytes_moved((m * k * dt.size_bytes() + 2 * m * dt.size_bytes()) as u64) + } +} diff --git a/crates/metaltile-std/src/ffai/gemv_q8.rs b/crates/metaltile-std/src/ffai/gemv_q8.rs new file mode 100644 index 00000000..7232d998 --- /dev/null +++ b/crates/metaltile-std/src/ffai/gemv_q8.rs @@ -0,0 +1,299 @@ +//! Copyright 2026 0xClandestine, Ekryski, TheTom, Ambisphaeric +//! SPDX-License-Identifier: Apache-2.0 +//! Q8_0 inline-dequant GEMV — `out[r] = Σ_k dequant(W[r,k]) · x[k]`, +//! reading the Q8_0 weight straight from its split buffers so the dense +//! attention/shared-expert projections stay 1 byte/weight resident +//! instead of being pre-expanded to f16 (2 bytes). The DSv4 attention +//! block is bandwidth-bound on these projections (q_b / output_a / +//! output_b are Q8_0 on disk, ~100M weights/layer); halving their bytes +//! roughly halves the attn GPU time. +//! +//! ## Q8_0 block (32 values) +//! d (f16 scale) + 32 int8 quants; value[i] = d · q_i8[i] +//! +//! ## Split inputs (loader produces these once, resident) +//! qs [m_out * (k_in/32) * 8] u32 — 32 int8/block packed as 8 LE u32 +//! d [m_out * (k_in/32)] f32 — per-block scale (fp16→f32) +//! x [k_in] T +//! out [m_out] T +//! +//! Dispatch (Reduction): grid (threadgroups) = [m_out, 1, 1], tg=[32,1,1]. + +use metaltile::kernel; + +#[kernel] +pub fn ffai_gemv_q8( + qs: Tensor, + d_f32: Tensor, + x: Tensor, + mut out: Tensor, + #[constexpr] k_in: u32, + #[constexpr] m_out: u32, +) { + let row = tgid_x; + let lane = tid; + let bpr = k_in / 32u32; + let qs_base = row * bpr * 8u32; + let d_base = row * bpr; + let mut acc = 0.0f32; + for b in range(lane, bpr, 32u32) { + let d = load(d_f32[d_base + b]); + let x_base = b * 32u32; + for w in range(0u32, 8u32, 1u32) { + let packed = load(qs[qs_base + b * 8u32 + w]); + for i in range(0u32, 4u32, 1u32) { + let by = (packed >> (i * 8u32)) & 0xffu32; + // sign-extend the byte to int8 range in the float domain + // (avoids ambiguous integer `select` in MSL). + let qf = by.cast::() - select(by > 127u32, 256.0f32, 0.0f32); + let val = d * qf; + acc = acc + val * load(x[x_base + w * 4u32 + i]).cast::(); + } + } + } + let total = simd_sum(acc); + if lane == 0u32 { + store(out[row], total.cast::()); + } +} + +/// Grouped Q8_0 gemv — `out[r] = Σ_k dequant(W[r,k]) · x[(r/rows_per_group)*k_in + k]`. +/// Each contiguous block of `rows_per_group` output rows reads its own +/// `k_in`-slice of `x`. Fuses the DSv4 grouped O-LoRA (8 groups × a +/// [1024,4096] Q8 slice, each on a different 4096-slice of the attention +/// output) into a SINGLE dispatch instead of 8. +#[kernel] +pub fn ffai_grouped_gemv_q8( + qs: Tensor, + d_f32: Tensor, + x: Tensor, + mut out: Tensor, + #[constexpr] k_in: u32, + #[constexpr] m_out: u32, + #[constexpr] rows_per_group: u32, +) { + let row = tgid_x; + let lane = tid; + let bpr = k_in / 32u32; + let qs_base = row * bpr * 8u32; + let d_base = row * bpr; + let x_base = (row / rows_per_group) * k_in; + let mut acc = 0.0f32; + for b in range(lane, bpr, 32u32) { + let d = load(d_f32[d_base + b]); + let x_blk = x_base + b * 32u32; + for w in range(0u32, 8u32, 1u32) { + let packed = load(qs[qs_base + b * 8u32 + w]); + for i in range(0u32, 4u32, 1u32) { + let by = (packed >> (i * 8u32)) & 0xffu32; + let qf = by.cast::() - select(by > 127u32, 256.0f32, 0.0f32); + acc = acc + d * qf * load(x[x_blk + w * 4u32 + i]).cast::(); + } + } + } + let total = simd_sum(acc); + if lane == 0u32 { + store(out[row], total.cast::()); + } +} + +/// BATCHED grouped Q8_0 gemv — ffai_grouped_gemv_q8 over `n_tokens` rows in +/// ONE dispatch (grid z/y = token). Prefill O-LoRA looped the per-token +/// grouped gemv N times; this folds it. x is [n_tokens, n_groups*k_in], +/// out is [n_tokens, m_out]; n_groups = m_out/rows_per_group. +/// Grid (Reduction): [m_out, n_tokens, 1], tg=[32,1,1]. +#[kernel] +#[allow(clippy::too_many_arguments)] +pub fn ffai_grouped_gemv_q8_rows( + qs: Tensor, + d_f32: Tensor, + x: Tensor, + mut out: Tensor, + #[constexpr] k_in: u32, + #[constexpr] m_out: u32, + #[constexpr] rows_per_group: u32, +) { + let row = tgid_x; + let token = tgid_y; + let lane = tid; + let bpr = k_in / 32u32; + let qs_base = row * bpr * 8u32; + let d_base = row * bpr; + let n_groups = m_out / rows_per_group; + let x_base = token * n_groups * k_in + (row / rows_per_group) * k_in; + let mut acc = 0.0f32; + for b in range(lane, bpr, 32u32) { + let d = load(d_f32[d_base + b]); + let x_blk = x_base + b * 32u32; + for w in range(0u32, 8u32, 1u32) { + let packed = load(qs[qs_base + b * 8u32 + w]); + for i in range(0u32, 4u32, 1u32) { + let by = (packed >> (i * 8u32)) & 0xffu32; + let qf = by.cast::() - select(by > 127u32, 256.0f32, 0.0f32); + acc = acc + d * qf * load(x[x_blk + w * 4u32 + i]).cast::(); + } + } + } + let total = simd_sum(acc); + if lane == 0u32 { + store(out[token * m_out + row], total.cast::()); + } +} + +/// TOKEN-TILED grouped Q8 gemv — the amortized fix for the prefill O-LoRA-A +/// hotspot. `ffai_grouped_gemv_q8_rows` re-reads each weight row from DRAM +/// once PER TOKEN (no amortization); at N=512 that's the single biggest +/// op in the attention block (~47 ms/layer). Here each threadgroup owns one +/// output row and a TILE of `tokens_per_tile` tokens: the Q8 weight block +/// (d + 8 packed = 32 int8) is loaded ONCE and applied to all T tokens, so +/// the weight DRAM traffic drops T-fold. T accumulators in a register stack. +/// grid (threadgroups) = [m_out, ceil(n_tokens/T), 1], threadgroup [32,1,1]. +#[kernel] +#[allow(clippy::too_many_arguments)] +pub fn ffai_grouped_gemv_q8_rows_tiled( + qs: Tensor, + d_f32: Tensor, + x: Tensor, + mut out: Tensor, + #[constexpr] k_in: u32, + #[constexpr] m_out: u32, + #[constexpr] rows_per_group: u32, + #[constexpr] n_tokens: u32, +) { + // Tile size is a compile-time LITERAL (stack_alloc can't size on a + // constexpr param — codegen emits the array decl before the constant + // is bound). T=8 tokens/tile → 8-fold weight-DRAM amortization. + let tokens_per_tile = 8u32; + let row = tgid_x; + let ttile = tgid_y; + let lane = tid; + let tok0 = ttile * tokens_per_tile; + let bpr = k_in / 32u32; + let qs_base = row * bpr * 8u32; + let d_base = row * bpr; + let n_groups = m_out / rows_per_group; + let group_off = (row / rows_per_group) * k_in; + + stack_alloc("acc", 8, "f32"); + for t in range(0u32, tokens_per_tile, 1u32) { + stack_store("acc", t, 0.0f32); + } + for b in range(lane, bpr, 32u32) { + let d = load(d_f32[d_base + b]); + let blk = b * 32u32; + for w in range(0u32, 8u32, 1u32) { + let packed = load(qs[qs_base + b * 8u32 + w]); // 4 int8 weights, read ONCE + for i in range(0u32, 4u32, 1u32) { + let by = (packed >> (i * 8u32)) & 0xffu32; + let qf = by.cast::() - select(by > 127u32, 256.0f32, 0.0f32); + let wv = d * qf; // weight value, reused across the T tokens + let kk = blk + w * 4u32 + i; + for t in range(0u32, tokens_per_tile, 1u32) { + let tok = tok0 + t; + if tok < n_tokens { + let xb = tok * n_groups * k_in + group_off + kk; + let prev = stack_load("acc", t); + stack_store("acc", t, prev + wv * load(x[xb]).cast::()); + } + } + } + } + } + for t in range(0u32, tokens_per_tile, 1u32) { + let tok = tok0 + t; + let total = simd_sum(stack_load("acc", t)); + if (lane == 0u32) & (tok < n_tokens) { + store(out[tok * m_out + row], total.cast::()); + } + } +} + +pub mod kernel_benches { + use metaltile::{bench, test::*}; + + use super::{ + ffai_gemv_q8, + ffai_grouped_gemv_q8, + ffai_grouped_gemv_q8_rows, + ffai_grouped_gemv_q8_rows_tiled, + }; + + #[bench(name = "ffai/gemv/q8", dtypes = [f32, f16, bf16])] + fn bench_gemv_q8(dt: DType) -> BenchSetup { + let m_out = 4096usize; + let k_in = 8192usize; + let bpr = k_in / 32; + BenchSetup::new(ffai_gemv_q8::kernel_ir_for(dt)) + .mode(KernelMode::Reduction) + .buffer(BenchBuffer::random("qs", m_out * bpr * 8, DType::U32)) + .buffer(BenchBuffer::random("d_f32", m_out * bpr, DType::F32)) + .buffer(BenchBuffer::random("x", k_in, dt)) + .buffer(BenchBuffer::zeros("out", m_out, dt).output()) + .constexpr("k_in", k_in as u32) + .constexpr("m_out", m_out as u32) + .grid_3d(m_out as u32, 1, 1, [32, 1, 1]) + .bytes_moved((m_out * bpr * 36 + k_in * dt.size_bytes()) as u64) + } + + #[bench(name = "ffai/gemv/grouped_q8", dtypes = [f32, f16, bf16])] + fn bench_grouped_gemv_q8(dt: DType) -> BenchSetup { + let m_out = 8192usize; + let k_in = 4096usize; + let bpr = k_in / 32; + BenchSetup::new(ffai_grouped_gemv_q8::kernel_ir_for(dt)) + .mode(KernelMode::Reduction) + .buffer(BenchBuffer::random("qs", m_out * bpr * 8, DType::U32)) + .buffer(BenchBuffer::random("d_f32", m_out * bpr, DType::F32)) + .buffer(BenchBuffer::random("x", 8 * k_in, dt)) + .buffer(BenchBuffer::zeros("out", m_out, dt).output()) + .constexpr("k_in", k_in as u32) + .constexpr("m_out", m_out as u32) + .constexpr("rows_per_group", 1024u32) + .grid_3d(m_out as u32, 1, 1, [32, 1, 1]) + .bytes_moved((m_out * bpr * 36 + 8 * k_in * dt.size_bytes()) as u64) + } + + #[bench(name = "ffai/gemv/grouped_q8_rows", dtypes = [f32, f16, bf16])] + fn bench_grouped_gemv_q8_rows(dt: DType) -> BenchSetup { + let m_out = 8192usize; + let k_in = 4096usize; + let n_tokens = 256usize; + let n_groups = m_out / 1024; + let bpr = k_in / 32; + BenchSetup::new(ffai_grouped_gemv_q8_rows::kernel_ir_for(dt)) + .mode(KernelMode::Reduction) + .buffer(BenchBuffer::random("qs", m_out * bpr * 8, DType::U32)) + .buffer(BenchBuffer::random("d_f32", m_out * bpr, DType::F32)) + .buffer(BenchBuffer::random("x", n_tokens * n_groups * k_in, dt)) + .buffer(BenchBuffer::zeros("out", n_tokens * m_out, dt).output()) + .constexpr("k_in", k_in as u32) + .constexpr("m_out", m_out as u32) + .constexpr("rows_per_group", 1024u32) + .grid_3d(m_out as u32, n_tokens as u32, 1, [32, 1, 1]) + .bytes_moved((m_out * bpr * 36 + n_tokens * n_groups * k_in * dt.size_bytes()) as u64) + } + + #[bench(name = "ffai/gemv/grouped_q8_rows_tiled", dtypes = [f32, f16, bf16])] + fn bench_grouped_gemv_q8_rows_tiled(dt: DType) -> BenchSetup { + let m_out = 8192usize; + let k_in = 4096usize; + let n_tokens = 256usize; + let tokens_per_tile = 8usize; + let n_groups = m_out / 1024; + let bpr = k_in / 32; + BenchSetup::new(ffai_grouped_gemv_q8_rows_tiled::kernel_ir_for(dt)) + .mode(KernelMode::Reduction) + .buffer(BenchBuffer::random("qs", m_out * bpr * 8, DType::U32)) + .buffer(BenchBuffer::random("d_f32", m_out * bpr, DType::F32)) + .buffer(BenchBuffer::random("x", n_tokens * n_groups * k_in, dt)) + .buffer(BenchBuffer::zeros("out", n_tokens * m_out, dt).output()) + .constexpr("k_in", k_in as u32) + .constexpr("m_out", m_out as u32) + .constexpr("rows_per_group", 1024u32) + .constexpr("n_tokens", n_tokens as u32) + .grid_3d(m_out as u32, (n_tokens as u32).div_ceil(tokens_per_tile as u32), 1, [ + 32, 1, 1, + ]) + .bytes_moved((m_out * bpr * 36 + n_tokens * n_groups * k_in * dt.size_bytes()) as u64) + } +} diff --git a/crates/metaltile-std/src/ffai/gguf_dequant_iq2_xxs.rs b/crates/metaltile-std/src/ffai/gguf_dequant_iq2_xxs.rs new file mode 100644 index 00000000..a334edad --- /dev/null +++ b/crates/metaltile-std/src/ffai/gguf_dequant_iq2_xxs.rs @@ -0,0 +1,209 @@ +//! Copyright 2026 0xClandestine, Ekryski, TheTom, Ambisphaeric +//! SPDX-License-Identifier: Apache-2.0 +//! GGUF IQ2_XXS block dequant — i-quant 2-bit-per-weight with codebook lookup. +//! +//! Follows the canonical `dequantize_row_iq2_xxs` reference algorithm. +//! +//! **WIP scaffold.** The kernel ABI + dispatch geometry + reference +//! quantizer are laid down so this file integrates with the rest of +//! the GGUF dequant pipeline today; the in-kernel codebook lookup is +//! stubbed pending the canonical `iq2xxs_grid` table (256 × 8-byte +//! signed-octet entries) — see TODO blocks below. +//! +//! ## On-disk block layout +//! +//! ```text +//! struct block_iq2_xxs { +//! uint16_t d; // fp16 super-scale (2 bytes) +//! uint16_t qs[32]; // 64 bytes — 8 groups of 32 outputs each +//! }; // 66 bytes per 256 values (BPW = 2.0625) +//! ``` +//! +//! Per group of 32 outputs (8 groups per block): +//! +//! ```text +//! // qs[4*g .. 4*g + 4] = 2 uint16 = 1 uint32 of grid-index payload +//! // qs[4*g + 4 .. 4*g + 8] = 2 uint16 = 1 uint32 of sign+scale payload +//! aux32_idx = u32 of qs[4*g + 0..4*g + 2] +//! aux32_sgn = u32 of qs[4*g + 2..4*g + 4] +//! scale_4bit = aux32_sgn >> 28 // upper nibble of sign u32 +//! db = d * (0.5 + scale_4bit) * 0.25 // per-group composite scale +//! for j in 0..4: +//! grid_idx = (aux32_idx >> (8*j)) & 0xff // byte j → 256-entry LUT key +//! grid_row = iq2xxs_grid[grid_idx] // 8 signed octets (i8x8) +//! sign_bits = (aux32_sgn >> (7*j)) & 0x7f // 7-bit sign vector +//! // Bit-popcount-parity expansion: the high bit is implicit XOR of +//! // the low 7 bits. Reconstruct as ±1 per octet. +//! for k in 0..8: +//! sign = sign_bits has bit k set ? -1 : +1 +//! out[32*g + 8*j + k] = db * sign * grid_row[k] +//! ``` +//! +//! ## GPU-resident split (the loader produces these from the packed block) +//! +//! 1. `qs_u32 [n_blocks * 16]` — `u32`, the 64 bytes of `qs[32]` +//! re-laid as 16 little-endian u32 words. Each 32-element group +//! consumes 2 u32 (indices + sign+scale). +//! 2. `d_f32 [n_blocks]` — `f32`, host-converted from fp16. +//! 3. `grid [256 * 8]` — `i8`, the canonical `iq2xxs_grid` +//! table from ggml-quants.c, repacked one signed octet per byte +//! (8 bytes per row, 256 rows). Shared across all dequant calls — +//! upload once at runtime init. +//! +//! ## Dispatch +//! +//! 1D grid over output values, but inner work is per-group-of-32 so +//! threads cooperate via the same (aux32_idx, aux32_sgn) pair. +//! Simplest shape: 1 thread per output, gather the relevant grid row +//! per output. Wasteful (8× redundant grid loads vs the per-group +//! pattern) but correctness-first; a per-group tile follow-up cuts +//! that. +//! +//! ## ABI +//! +//! ```text +//! qs_u32 [n_blocks * 16] u32 — packed grid-index + sign payloads +//! d_f32 [n_blocks] f32 — per-block super-scale +//! grid [2048] i8 — iq2xxs_grid as 256 × 8 signed octets +//! out [n_values] T — dequantized output +//! n_values u32 (constexpr) — total output count = n_blocks * 256 +//! ``` + +use metaltile::kernel; + +// Bare `#[kernel]` — see Q8_0 sibling for why; mixed concrete + +// generic param dtype set doesn't fit the legacy `bench(...)` shape. +// +// The dequant uses **two** runtime lookup tables from the +// canonical IQ2_XXS reference algorithm (Apache-2.0 here): +// +// 1. `grid: [256 × 8] u8` — the `iq2xxs_grid` octet table. +// Each row encodes 8 small-integer magnitudes (observed values: +// `{0x08, 0x19, 0x2b}` = `{8, 25, 43}` — the absolute octet +// magnitudes that the sign byte modulates). +// 2. `signs: [128] u8` — the `ksigns_iq2xs` sign byte +// table. Each entry maps a 7-bit sign-field index to an 8-bit +// mask; bit `l` of the mask is 1 → that octet flips sign. +// +// Both are uploaded once at runtime init and shared across all +// dequant calls. +#[kernel] +pub fn ffai_gguf_dequant_iq2_xxs( + qs_u32: Tensor, + d_f32: Tensor, + grid: Tensor, + signs: Tensor, + out: Tensor, + #[constexpr] n_values: u32, +) { + let i = tid; + if i < n_values { + let block = i / 256u32; + let in_block = i - block * 256u32; + let group = in_block / 32u32; // 0..7 — which of the 8 groups + let in_group = in_block & 31u32; // 0..31 — position inside group + let octet_within_index = in_group / 8u32; // 0..3 — which of 4 grid indices + let lane_in_octet = in_group & 7u32; // 0..7 — which octet byte + + // Two u32s per group: indices (qs_u32[block*16 + group*2]) and + // sign+scale payload (qs_u32[block*16 + group*2 + 1]). + let aux_idx = load(qs_u32[block * 16u32 + group * 2u32]); + let aux_sgn = load(qs_u32[block * 16u32 + group * 2u32 + 1u32]); + + let scale_4bit = aux_sgn >> 28u32; + // db = d * (0.5 + scale_4bit) * 0.25 — composite per-group scale. + let scale_factor = (scale_4bit.cast::().cast::() + 0.5) * 0.25; + let db = load(d_f32[block]) * scale_factor; + + // Grid lookup: byte `octet_within_index` of `aux_idx` is the + // 256-entry LUT key. + let grid_key = (aux_idx >> (octet_within_index * 8u32)) & 0xffu32; + let grid_row_base = grid_key * 8u32; + let octet_u = load(grid[grid_row_base + lane_in_octet]).cast::(); + let octet = octet_u.cast::().cast::(); + + // Sign reconstruction: 7-bit sign-index lives at bit + // `7 * octet_within_index` in aux_sgn; look it up in the + // `ksigns_iq2xs` table to get the 8-bit sign mask. Bit `l` + // of that mask says whether octet `l` flips sign. + let sign_idx = (aux_sgn >> (octet_within_index * 7u32)) & 0x7fu32; + let sign_mask = load(signs[sign_idx]).cast::(); + let lane_bit = sign_mask & (1u32 << lane_in_octet); + let sign = select(lane_bit != 0u32, -1.0f32, 1.0f32); + + // Store the f32 result directly: the DSL narrows f32→T implicitly + // at the Store site (an explicit `.cast::()` would emit a + // spurious same-type MSL cast). + store(out[i], db * sign * octet); + } +} + +pub mod kernel_tests { + //! **WIP**: end-to-end correctness ignored pending the canonical + //! `iq2xxs_grid` table. The kernel itself compiles + dispatches + //! cleanly; the codegen-shape invariants below catch ABI / IR + //! regressions even without the grid table. + + use metaltile::test::*; + + use super::ffai_gguf_dequant_iq2_xxs; + use crate::utils::pack_f32; + + /// MSL emission smoke test — confirms the kernel codegens for all + /// float output dtypes without a real input set. Catches IR / + /// macro-expansion regressions; correctness is gated separately + /// once `iq2xxs_grid` lands. + #[test] + fn codegen_iq2_xxs_smoke() { + for dt in [DType::F32, DType::F16, DType::BF16] { + let ir = ffai_gguf_dequant_iq2_xxs::kernel_ir_for(dt); + assert!(!ir.body.ops.is_empty(), "kernel body emitted no ops for {dt:?}"); + assert!(ir.params.iter().any(|p| p.name == "qs_u32"), "missing qs_u32 param"); + assert!(ir.params.iter().any(|p| p.name == "grid"), "missing grid param"); + } + } + + /// TODO end-to-end correctness vs the canonical reference. The + /// grid + ksigns tables ship on the FFAI loader side. Once they're + /// wired through Ops.swift this fixture should round-trip a + /// quantized + dequantized block + /// and assert within bf16 tolerance. + #[allow(dead_code)] + fn _placeholder_setup(n_blocks: usize, dt: DType) -> TestSetup { + let n = n_blocks * 256; + TestSetup::new(ffai_gguf_dequant_iq2_xxs::kernel_ir_for(dt)) + .input(TestBuffer::zeros("qs_u32", n_blocks * 16, DType::U32)) + .input(TestBuffer::from_vec( + "d_f32", + pack_f32(&vec![1.0; n_blocks], DType::F32), + DType::F32, + )) + .input(TestBuffer::zeros("grid", 2048, DType::U8)) + .input(TestBuffer::zeros("signs", 128, DType::U8)) + .input(TestBuffer::zeros("out", n, dt)) + .constexpr("n_values", n as u32) + .expect(TestBuffer::zeros("out", n, dt)) + .grid_1d(n, 256) + } +} + +pub mod kernel_benches { + use metaltile::{bench, test::*}; + + use super::ffai_gguf_dequant_iq2_xxs; + + #[bench(name = "ffai/gguf_dequant_iq2_xxs", dtypes = [f32, f16, bf16])] + fn bench_iq2_xxs(dt: DType) -> BenchSetup { + let n = 4096 * 4096usize; + let n_blocks = n / 256; + BenchSetup::new(ffai_gguf_dequant_iq2_xxs::kernel_ir_for(dt)) + .buffer(BenchBuffer::random("qs_u32", n_blocks * 16, DType::U32)) + .buffer(BenchBuffer::random("d_f32", n_blocks, DType::F32)) + .buffer(BenchBuffer::random("grid", 2048, DType::U8)) + .buffer(BenchBuffer::random("signs", 128, DType::U8)) + .buffer(BenchBuffer::zeros("out", n, dt).output()) + .constexpr("n_values", n as u32) + .grid_1d(n, 256) + .bytes_moved(((n_blocks * 64 + n_blocks * 4 + 2048 + 128) + n * dt.size_bytes()) as u64) + } +} diff --git a/crates/metaltile-std/src/ffai/gguf_dequant_iq2_xxs_raw.rs b/crates/metaltile-std/src/ffai/gguf_dequant_iq2_xxs_raw.rs new file mode 100644 index 00000000..9b916c6f --- /dev/null +++ b/crates/metaltile-std/src/ffai/gguf_dequant_iq2_xxs_raw.rs @@ -0,0 +1,104 @@ +//! Copyright 2026 0xClandestine, Ekryski, TheTom, Ambisphaeric +//! SPDX-License-Identifier: Apache-2.0 +//! GGUF IQ2_XXS dequant — raw-bytes variant. +//! +//! Same algorithm as `ffai_gguf_dequant_iq2_xxs` but reads `qs` from +//! the on-disk 66-byte block layout directly, skipping the CPU +//! preprocess that split each block into a separate `qs_u32` buffer. +//! +//! The block layout is `[u16 d (fp16); u8 qs[64]]` — 66 bytes per 256 +//! values. `d_f32` is still passed as a pre-staged buffer because the +//! DSL has no `bit_cast` intrinsic for the in-kernel fp16 +//! → f32 conversion; staging that single 32K-element vector on CPU +//! is ~30 ms per token vs the ~470 ms the per-block `qs` memcpy was +//! costing, so the asymmetric split is net positive. + +use metaltile::kernel; + +#[kernel] +pub fn ffai_gguf_dequant_iq2_xxs_raw( + raw_bytes: Tensor, + d_f32: Tensor, + grid: Tensor, + signs: Tensor, + out: Tensor, + #[constexpr] n_values: u32, +) { + let i = tid; + if i < n_values { + let block = i / 256u32; + let in_block = i - block * 256u32; + let group = in_block / 32u32; + let in_group = in_block & 31u32; + let octet_within_index = in_group / 8u32; + let lane_in_octet = in_group & 7u32; + + // Block base = block * 66, then skip the 2-byte fp16 d header + // → qs region starts at +2. Each group occupies 8 bytes (one + // `aux_idx` u32 + one `aux_sgn` u32). + let qs_base = block * 66u32 + 2u32 + group * 8u32; + let i0 = load(raw_bytes[qs_base]).cast::(); + let i1 = load(raw_bytes[qs_base + 1u32]).cast::(); + let i2 = load(raw_bytes[qs_base + 2u32]).cast::(); + let i3 = load(raw_bytes[qs_base + 3u32]).cast::(); + let aux_idx = i0 | (i1 << 8u32) | (i2 << 16u32) | (i3 << 24u32); + let s0 = load(raw_bytes[qs_base + 4u32]).cast::(); + let s1 = load(raw_bytes[qs_base + 5u32]).cast::(); + let s2 = load(raw_bytes[qs_base + 6u32]).cast::(); + let s3 = load(raw_bytes[qs_base + 7u32]).cast::(); + let aux_sgn = s0 | (s1 << 8u32) | (s2 << 16u32) | (s3 << 24u32); + + let scale_4bit = aux_sgn >> 28u32; + let scale_factor = (scale_4bit.cast::().cast::() + 0.5) * 0.25; + let db = load(d_f32[block]) * scale_factor; + + let grid_key = (aux_idx >> (octet_within_index * 8u32)) & 0xffu32; + let grid_row_base = grid_key * 8u32; + let octet_u = load(grid[grid_row_base + lane_in_octet]).cast::(); + let octet = octet_u.cast::().cast::(); + + let sign_idx = (aux_sgn >> (octet_within_index * 7u32)) & 0x7fu32; + let sign_mask = load(signs[sign_idx]).cast::(); + let lane_bit = sign_mask & (1u32 << lane_in_octet); + let sign = select(lane_bit != 0u32, -1.0f32, 1.0f32); + + store(out[i], db * sign * octet); + } +} + +#[cfg(test)] +pub mod kernel_tests { + use metaltile::test::*; + + use super::ffai_gguf_dequant_iq2_xxs_raw; + + #[test] + fn codegen_iq2_xxs_raw_smoke() { + for dt in [DType::F32, DType::F16, DType::BF16] { + let ir = ffai_gguf_dequant_iq2_xxs_raw::kernel_ir_for(dt); + assert!(!ir.body.ops.is_empty(), "kernel body emitted no ops for {dt:?}"); + assert!(ir.params.iter().any(|p| p.name == "raw_bytes"), "missing raw_bytes param"); + } + } +} + +pub mod kernel_benches { + use metaltile::{bench, test::*}; + + use super::ffai_gguf_dequant_iq2_xxs_raw; + + #[bench(name = "ffai/gguf_dequant_iq2_xxs_raw", dtypes = [f32, f16, bf16])] + fn bench_iq2_xxs_raw(dt: DType) -> BenchSetup { + let n = 4096 * 4096usize; + let n_blocks = n / 256; + BenchSetup::new(ffai_gguf_dequant_iq2_xxs_raw::kernel_ir_for(dt)) + .buffer(BenchBuffer::random("raw_bytes", n_blocks * 66, DType::U8)) + .buffer(BenchBuffer::random("d_f32", n_blocks, DType::F32)) + .buffer(BenchBuffer::random("grid", 2048, DType::U8)) + .buffer(BenchBuffer::random("signs", 128, DType::U8)) + .buffer(BenchBuffer::zeros("out", n, dt).output()) + .constexpr("n_values", n as u32) + .grid_1d(n, 256) + .bytes_moved(((n_blocks * 66 + n_blocks * 4 + 2048 + 128) + n * dt.size_bytes()) as u64) + } +} diff --git a/crates/metaltile-std/src/ffai/gguf_dequant_q2_k.rs b/crates/metaltile-std/src/ffai/gguf_dequant_q2_k.rs new file mode 100644 index 00000000..695aa7b1 --- /dev/null +++ b/crates/metaltile-std/src/ffai/gguf_dequant_q2_k.rs @@ -0,0 +1,266 @@ +//! Copyright 2026 0xClandestine, Ekryski, TheTom, Ambisphaeric +//! SPDX-License-Identifier: Apache-2.0 +//! GGUF Q2_K block dequant — k-quant 2-bit-per-weight with two-level scales. +//! +//! Follows the canonical `dequantize_row_q2_K` reference algorithm. +//! +//! ## On-disk block layout (decomposed CPU-side at load time) +//! +//! ```text +//! struct block_q2_K { +//! uint8_t scales[16]; // 16 bytes — low 4 bits = scale, high 4 bits = min +//! uint8_t qs[64]; // 64 bytes — 2-bit-packed quants, 4 vals per byte +//! uint16_t d; // 2 bytes — fp16 super-scale for scales +//! uint16_t dmin; // 2 bytes — fp16 super-scale for mins +//! }; // 84 bytes per 256 values (BPW = 2.625) +//! ``` +//! +//! Per output value `i ∈ [0, 256)`: +//! +//! ```text +//! sub = i / 16 // 0..15, picks the (4-bit scale, 4-bit min) pair +//! in_sub = i & 15 // 0..15 inside the sub-block +//! scale_byte = scales[sub] +//! scale_4bit = scale_byte & 0xf +//! min_4bit = (scale_byte >> 4) & 0xf +//! qs_byte = qs[i / 4] +//! shift = (i & 3) * 2 +//! q_2bit = (qs_byte >> shift) & 0x3 +//! out[i] = d * scale_4bit * q_2bit - dmin * min_4bit +//! ``` +//! +//! ## GPU-resident split (the loader produces these from the packed block) +//! +//! 1. `qs_packed [n_blocks * 16]` — `u32`, the 64 packed-quant bytes +//! per block re-laid as 16 u32 words. `qs_packed[block*16 + j]` +//! carries 16 two-bit quants in the lower / upper bytes of each +//! u32. Output index `i ∈ [0, 256)` → `u32 j = i / 16`, then a +//! `(i % 16) * 2`-bit shift on the byte that holds it. +//! 2. `scales [n_blocks * 16]` — `u8`, the raw scale/min byte +//! pairs (low nibble = scale, high nibble = min) — kept packed +//! because both nibbles are used per dequant. +//! 3. `d_f32 [n_blocks]` — `f32`, host-converted from fp16. +//! 4. `dmin_f32 [n_blocks]` — `f32`, host-converted from fp16. +//! +//! ## Dispatch +//! +//! 1D grid: one thread per *output value*. ~6 reads (1 qs_packed + 1 +//! scales + 1 each of d_f32 / dmin_f32, scales cache-multicast across +//! 16 lanes that share a sub-block) and ~4 arithmetic ops per output — +//! cleanly bandwidth-bound on Apple9. + +use metaltile::kernel; + +// Bare `#[kernel]` — see Q8_0 sibling for why; mixed concrete + +// generic param dtype set doesn't fit the legacy `bench(...)` shape. +#[kernel] +pub fn ffai_gguf_dequant_q2_k( + qs_packed: Tensor, + scales: Tensor, + d_f32: Tensor, + dmin_f32: Tensor, + out: Tensor, + #[constexpr] n_values: u32, +) { + let i = tid; + if i < n_values { + let block = i / 256u32; + let in_block = i - block * 256u32; + // Canonical Q2_K block layout: the 256 + // values are NOT 4-consecutive-per-byte. They split as 2 halves of + // 128; each half is 4 j-groups of 32; each j-group is two runs of 16 + // values that index 16 CONSECUTIVE qs bytes at a SHARED 2-bit shift + // (shift = j*2). The naive in_block/4 mapping was wrong. + let half = in_block / 128u32; // 0..1 → qs byte base half*32 + let yh = in_block - half * 128u32; // 0..127 + let jg = yh / 32u32; // 0..3 → shift = jg*2 + let yg = yh - jg * 32u32; // 0..31 + let sub_half = yg / 16u32; // 0..1 + let l = yg - sub_half * 16u32; // 0..15 → byte within the 16-run + let shift = jg * 2u32; + let q_byte = half * 32u32 + sub_half * 16u32 + l; // 0..63 + let scale_idx = half * 8u32 + jg * 2u32 + sub_half; // 0..15 + let word_idx = q_byte / 4u32; + let byte_in_word = q_byte & 3u32; + let word = load(qs_packed[block * 16u32 + word_idx]); + let qs_byte = (word >> (byte_in_word * 8u32)) & 0xffu32; + let q_2bit = (qs_byte >> shift) & 0x3u32; + + let scale_byte = load(scales[block * 16u32 + scale_idx]).cast::(); + let scale_4bit = scale_byte & 0xfu32; + let min_4bit = (scale_byte >> 4u32) & 0xfu32; + + let d = load(d_f32[block]); + let dmin = load(dmin_f32[block]); + + let scaled = + d * (scale_4bit.cast::().cast::()) * (q_2bit.cast::().cast::()); + let offset = dmin * (min_4bit.cast::().cast::()); + // Store the f32 result directly: the DSL narrows f32→T implicitly + // at the Store site (an explicit `.cast::()` would emit a + // spurious same-type MSL cast). + store(out[i], scaled - offset); + } +} + +pub mod kernel_tests { + use metaltile::{test::*, test_kernel}; + + use super::ffai_gguf_dequant_q2_k; + use crate::utils::pack_f32; + + /// Canonical Q2_K output-index → (qs byte 0..63, 2-bit shift). + /// Mirrors the GPU kernel's mapping exactly (llama.cpp + /// `dequantize_row_q2_K` order) so the quantizer, the CPU reference, + /// and the kernel all read/write the same byte for a given output — + /// the canonical scale index is simply `i / 16`. + fn q2_k_qpos(i: usize) -> (usize, u32) { + let half = i / 128; // 0..1 → qs byte base half*32 + let yh = i % 128; // 0..127 + let jg = yh / 32; // 0..3 → shift = jg*2 + let yg = yh % 32; // 0..31 + let sub_half = yg / 16; // 0..1 + let l = yg % 16; // 0..15 → byte within the 16-run + (half * 32 + sub_half * 16 + l, (jg * 2) as u32) + } + + /// Quantize a slice into the Q2_K GPU-resident split. + /// + /// This is a *non-trained* quantizer (per-sub-block min/max with + /// uniform 2-bit bucketing) sufficient to drive the kernel's + /// correctness test — `quantize_row_q2_K` in ggml-quants.c uses a + /// search-and-fit procedure tuned for perplexity that we don't + /// need to replicate for a dequant test. + fn quantize_q2_k(values: &[f32]) -> (Vec, Vec, Vec, Vec) { + assert_eq!(values.len() % 256, 0, "Q2_K needs multiple-of-256 values"); + let n_blocks = values.len() / 256; + let mut qs_packed = Vec::with_capacity(n_blocks * 16); + let mut scales = Vec::with_capacity(n_blocks * 16); + let mut d_f32 = Vec::with_capacity(n_blocks); + let mut dmin_f32 = Vec::with_capacity(n_blocks); + + for block in values.chunks_exact(256) { + // Per-sub-block (16 sub-blocks of 16 values each) compute + // mn, scale; quantize to 2 bits. + let mut sub_scales = [0u8; 16]; + let mut sub_mins = [0u8; 16]; + let mut qs_bytes = [0u8; 64]; + for s in 0..16 { + let sub = &block[s * 16..(s + 1) * 16]; + let mn = sub.iter().cloned().fold(f32::INFINITY, f32::min); + let mx = sub.iter().cloned().fold(f32::NEG_INFINITY, f32::max); + let scale = ((mx - mn) / 3.0).max(1e-30); + let mn_q = (-mn / scale).round().clamp(0.0, 15.0) as u8; + let scale_q = (scale * 15.0).round().clamp(0.0, 15.0) as u8; + // Reconstruct the encoded `d` / `dmin` super-scales so + // dequant exactly recovers the per-sub-block scale. + sub_scales[s] = scale_q; + sub_mins[s] = mn_q; + // Pack two 2-bit quants per nibble: out = round((x + min) / scale). + let recon_scale = (scale_q as f32) / 15.0; + let recon_min = -(mn_q as f32) * recon_scale; + for (i, &v) in sub.iter().enumerate() { + let q = + ((v - recon_min) / recon_scale.max(1e-30)).round().clamp(0.0, 3.0) as u8; + // Output index s*16+i (scale s governs the 16 outputs + // [s*16, s*16+16)); write the quant at its canonical + // byte/shift so the kernel recovers it. + let (q_byte, shift) = q2_k_qpos(s * 16 + i); + qs_bytes[q_byte] |= q << shift; + } + } + // Combine the per-sub-block 4-bit scale + 4-bit min into + // the on-disk scales[16] layout. + for s in 0..16 { + scales.push((sub_mins[s] << 4) | sub_scales[s]); + } + // Repack `qs_bytes` (64 bytes) into 16 u32 LE words. + for w in 0..16 { + let bs = w * 4; + qs_packed.push( + (qs_bytes[bs] as u32) + | ((qs_bytes[bs + 1] as u32) << 8) + | ((qs_bytes[bs + 2] as u32) << 16) + | ((qs_bytes[bs + 3] as u32) << 24), + ); + } + // The two super-scales encode `d = 1/15` and `dmin = 1` + // after the reconstruction above. Materialize as fp32 the + // same way the GGUF loader would after fp16-converting. + d_f32.push(half::f16::from_f32(1.0 / 15.0).to_f32()); + dmin_f32.push(half::f16::from_f32(1.0).to_f32()); + } + (qs_packed, scales, d_f32, dmin_f32) + } + + /// CPU reference. Mirrors the GPU kernel exactly so any kernel-side + /// divergence shows up as a tolerance miss. + fn cpu_dequant(qs_packed: &[u32], scales: &[u8], d_f32: &[f32], dmin_f32: &[f32]) -> Vec { + let n_blocks = d_f32.len(); + let mut out = Vec::with_capacity(n_blocks * 256); + for b in 0..n_blocks { + let d = d_f32[b]; + let dmin = dmin_f32[b]; + for i in 0..256_usize { + let sub = i / 16; // canonical scale index == i / 16 + let (q_byte, shift) = q2_k_qpos(i); + let word = qs_packed[b * 16 + q_byte / 4]; + let byte_in_word = q_byte & 3; + let qs_byte = (word >> (byte_in_word * 8)) & 0xff; + let q_2bit = (qs_byte >> shift) & 0x3; + let scale_byte = scales[b * 16 + sub] as u32; + let scale_4bit = scale_byte & 0xf; + let min_4bit = (scale_byte >> 4) & 0xf; + out.push(d * (scale_4bit as f32) * (q_2bit as f32) - dmin * (min_4bit as f32)); + } + } + out + } + + fn setup(n_blocks: usize, dt: DType) -> TestSetup { + let n = n_blocks * 256; + let values: Vec = (0..n).map(|i| (i as f32 * 0.007 - 0.5).sin() * 1.5).collect(); + let (qs_packed, scales, d_f32, dmin_f32) = quantize_q2_k(&values); + let dequantized = cpu_dequant(&qs_packed, &scales, &d_f32, &dmin_f32); + // Pack u32 vec as little-endian bytes for the test framework. + let qs_bytes: Vec = qs_packed.iter().flat_map(|w| w.to_le_bytes()).collect(); + TestSetup::new(ffai_gguf_dequant_q2_k::kernel_ir_for(dt)) + .input(TestBuffer::from_vec("qs_packed", qs_bytes, DType::U32)) + .input(TestBuffer::from_vec("scales", scales, DType::U8)) + .input(TestBuffer::from_vec("d_f32", pack_f32(&d_f32, DType::F32), DType::F32)) + .input(TestBuffer::from_vec("dmin_f32", pack_f32(&dmin_f32, DType::F32), DType::F32)) + .input(TestBuffer::zeros("out", n, dt)) + .constexpr("n_values", n as u32) + .expect(TestBuffer::from_vec("out", pack_f32(&dequantized, dt), dt)) + .grid_1d(n, 256) + } + + #[test_kernel(dtypes = [f32, f16, bf16], tol = [1e-4, 5e-3, 5e-2])] + fn test_gguf_q2_k_single_block(dt: DType) -> TestSetup { setup(1, dt) } + + #[test_kernel(dtypes = [f32, f16, bf16], tol = [1e-4, 5e-3, 5e-2])] + fn test_gguf_q2_k_many_blocks(dt: DType) -> TestSetup { setup(8, dt) } +} + +pub mod kernel_benches { + use metaltile::{bench, test::*}; + + use super::ffai_gguf_dequant_q2_k; + + #[bench(name = "ffai/gguf_dequant_q2_k", dtypes = [f32, f16, bf16])] + fn bench_q2_k(dt: DType) -> BenchSetup { + // Representative MoE-expert down-proj slab — 4096 × 4096. + let n = 4096 * 4096usize; + let n_blocks = n / 256; + BenchSetup::new(ffai_gguf_dequant_q2_k::kernel_ir_for(dt)) + .buffer(BenchBuffer::random("qs_packed", n_blocks * 16, DType::U32)) + .buffer(BenchBuffer::random("scales", n_blocks * 16, DType::U8)) + .buffer(BenchBuffer::random("d_f32", n_blocks, DType::F32)) + .buffer(BenchBuffer::random("dmin_f32", n_blocks, DType::F32)) + .buffer(BenchBuffer::zeros("out", n, dt).output()) + .constexpr("n_values", n as u32) + .grid_1d(n, 256) + // qs_packed 64 B + scales 16 B + 2*4 B per block + output T + .bytes_moved(((n_blocks * (64 + 16 + 8)) + n * dt.size_bytes()) as u64) + } +} diff --git a/crates/metaltile-std/src/ffai/gguf_dequant_q8_0.rs b/crates/metaltile-std/src/ffai/gguf_dequant_q8_0.rs new file mode 100644 index 00000000..9ca4b8dc --- /dev/null +++ b/crates/metaltile-std/src/ffai/gguf_dequant_q8_0.rs @@ -0,0 +1,166 @@ +//! Copyright 2026 0xClandestine, Ekryski, TheTom, Ambisphaeric +//! SPDX-License-Identifier: Apache-2.0 +//! GGUF Q8_0 block dequant — `out[i] = qs[i] * d[i/32]`. +//! +//! Follows the canonical `dequantize_row_q8_0` reference algorithm. +//! +//! ## On-disk block layout (decomposed CPU-side at load time) +//! +//! ```text +//! struct block_q8_0 { +//! uint16_t d; // fp16 super-scale (2 bytes) +//! int8_t qs[32]; // 32 quantized int8 values (32 bytes) +//! }; // 34 bytes per block +//! ``` +//! +//! GGUF blocks are tightly packed: block N starts at byte `34*N`. To +//! avoid in-kernel bit-cast / fp16-bit-reconstruction (which the +//! metaltile DSL doesn't expose today), the GGUF loader splits each +//! block into two GPU-resident tensors at parse time: +//! +//! 1. `qs_signed [n_blocks * 32]` — `u8` view of the original int8 +//! quants (signed reconstruction happens via the +//! `select(q >= 128, q - 256, q)` trick inside the kernel; no +//! arithmetic-shift / bit-cast intrinsic needed). +//! 2. `scales [n_blocks]` — `f32`, the fp16 super-scale +//! converted to f32 by the host loader. +//! +//! That single conversion pass is `O(n_blocks)` and runs once at load. +//! The hot-path kernel below does ~0 setup work per output value. +//! +//! ## Dispatch +//! +//! 1D grid: one thread per *output value*. Thread `tid` computes +//! `block = tid / 32`, `lane = tid % 32`. Reads +//! `q = qs_signed[tid]` and `d = scales[block]`. Adjacent lanes share +//! the same scale cache line — Apple's L1 multicast hides the gather. +//! +//! ## ABI +//! +//! ```text +//! qs_signed [n_values] u8 — the 32 int8 quants per block, packed +//! back-to-back (sign-reconstructed at use). +//! scales [n_blocks] f32 — host-extracted block super-scales. +//! out [n_values] T — dequantized output. +//! n_values u32 (constexpr) — total output count = n_blocks * 32. +//! ``` + +use metaltile::kernel; + +// Bare `#[kernel]` — kernel mixes concrete-dtype packed-byte inputs +// with a generic `Tensor` output, which doesn't fit the legacy +// `bench(...)` registration's `GenericEmpty` dispatch shape. The new +// declarative `#[bench]` attribute on the `kernel_benches::bench_q8_0` +// fn below registers this kernel for `tile bench` without the legacy +// path. +#[kernel] +pub fn ffai_gguf_dequant_q8_0( + qs_signed: Tensor, + scales: Tensor, + out: Tensor, + #[constexpr] n_values: u32, +) { + let i = tid; + // Bounds guard — caller may dispatch a power-of-two grid above + // `n_values` for alignment reasons. + if i < n_values { + let block = i / 32u32; + let q_u = load(qs_signed[i]).cast::(); + // Sign-extend the u8 to a signed f32 without a bit_cast: the + // u8 value `q_u >= 128` represents a negative int8, value + // `q_u - 256`. The `select` collapses to a `csel` in MSL. + let q_signed = select(q_u >= 128u32, q_u - 256u32, q_u); + let q = q_signed.cast::().cast::(); + let d = load(scales[block]); + // Store the f32 result directly: the DSL narrows f32→T implicitly + // at the Store site, so omitting `.cast::()` avoids a spurious + // f32→f32 MSL cast (measured 8.3e-3 numerical drift) and keeps the + // Store eligible for any future vectorize-window grouping. + store(out[i], q * d); + } +} + +pub mod kernel_tests { + use metaltile::{test::*, test_kernel}; + + use super::ffai_gguf_dequant_q8_0; + use crate::utils::pack_f32; + + /// Reference quantizer: mirrors `quantize_row_q8_0` in + /// ggml-quants.c. Returns the two GPU-resident tensors the kernel + /// expects: `(qs_signed, scales)`. + fn quantize_q8_0(values: &[f32]) -> (Vec, Vec) { + assert_eq!(values.len() % 32, 0, "Q8_0 needs multiple-of-32 values"); + let n_blocks = values.len() / 32; + let mut qs = Vec::with_capacity(n_blocks * 32); + let mut scales = Vec::with_capacity(n_blocks); + for block in values.chunks_exact(32) { + let amax = block.iter().map(|v| v.abs()).fold(0.0_f32, f32::max); + let d = if amax > 0.0 { amax / 127.0 } else { 1.0 }; + // Match the precision the on-disk fp16 super-scale gives + // us: round-trip through fp16 so the test's expected + // values match what the loader hands the kernel. + let d_f16 = half::f16::from_f32(d).to_f32(); + scales.push(d_f16); + let inv_d = if d_f16 > 0.0 { 1.0 / d_f16 } else { 0.0 }; + for &v in block { + let q = (v * inv_d).round().clamp(-128.0, 127.0) as i8; + qs.push(q as u8); + } + } + (qs, scales) + } + + fn cpu_dequant(qs: &[u8], scales: &[f32]) -> Vec { + let n_blocks = scales.len(); + let mut out = Vec::with_capacity(n_blocks * 32); + for b in 0..n_blocks { + let d = scales[b]; + for i in 0..32 { + let q = qs[b * 32 + i] as i8; + out.push(q as f32 * d); + } + } + out + } + + fn setup(n_blocks: usize, dt: DType) -> TestSetup { + let n = n_blocks * 32; + let values: Vec = (0..n).map(|i| (i as f32 * 0.013 - 0.4).sin() * 3.5).collect(); + let (qs, scales) = quantize_q8_0(&values); + let dequantized = cpu_dequant(&qs, &scales); + TestSetup::new(ffai_gguf_dequant_q8_0::kernel_ir_for(dt)) + .input(TestBuffer::from_vec("qs_signed", qs, DType::U8)) + .input(TestBuffer::from_vec("scales", pack_f32(&scales, DType::F32), DType::F32)) + .input(TestBuffer::zeros("out", n, dt)) + .constexpr("n_values", n as u32) + .expect(TestBuffer::from_vec("out", pack_f32(&dequantized, dt), dt)) + .grid_1d(n, 256) + } + + #[test_kernel(dtypes = [f32, f16, bf16], tol = [1e-4, 5e-3, 5e-2])] + fn test_gguf_q8_0_single_block(dt: DType) -> TestSetup { setup(1, dt) } + + #[test_kernel(dtypes = [f32, f16, bf16], tol = [1e-4, 5e-3, 5e-2])] + fn test_gguf_q8_0_many_blocks(dt: DType) -> TestSetup { setup(64, dt) } +} + +pub mod kernel_benches { + use metaltile::{bench, test::*}; + + use super::ffai_gguf_dequant_q8_0; + + #[bench(name = "ffai/gguf_dequant_q8_0", dtypes = [f32, f16, bf16])] + fn bench_q8_0(dt: DType) -> BenchSetup { + // 4096 × 4096 attn projection slab. + let n = 4096 * 4096usize; + let n_blocks = n / 32; + BenchSetup::new(ffai_gguf_dequant_q8_0::kernel_ir_for(dt)) + .buffer(BenchBuffer::random("qs_signed", n, DType::U8)) + .buffer(BenchBuffer::random("scales", n_blocks, DType::F32)) + .buffer(BenchBuffer::zeros("out", n, dt).output()) + .constexpr("n_values", n as u32) + .grid_1d(n, 256) + .bytes_moved(((n + n_blocks * 4) + n * dt.size_bytes()) as u64) + } +} diff --git a/crates/metaltile-std/src/ffai/gguf_iq2_xxs_extract_qs.rs b/crates/metaltile-std/src/ffai/gguf_iq2_xxs_extract_qs.rs new file mode 100644 index 00000000..5ee6b4aa --- /dev/null +++ b/crates/metaltile-std/src/ffai/gguf_iq2_xxs_extract_qs.rs @@ -0,0 +1,87 @@ +//! Copyright 2026 0xClandestine, Ekryski, TheTom, Ambisphaeric +//! SPDX-License-Identifier: Apache-2.0 +//! GGUF IQ2_XXS qs extraction: raw-bytes → packed `qs_u32`. +//! +//! Replaces the per-block CPU `memcpy(64)` loop that was eating ~26 ms +//! per expert (~470 ms / token across 18 routed experts at top-K=6 × +//! 3 weight roles on DSv4-Flash). The extract kernel reads 4 raw +//! bytes from the on-disk 66-byte block layout, assembles one u32, +//! and writes to the packed `qs_u32` staging buffer the existing +//! `ffai_gguf_dequant_iq2_xxs` kernel expects. +//! +//! ## Block layout +//! +//! ```text +//! struct block_iq2_xxs { +//! uint16_t d; // bytes 0..2 — fp16 super-scale +//! uint8_t qs[64]; // bytes 2..66 — packed grid index + sign payloads +//! }; // 66 bytes per 256 values +//! ``` +//! +//! Each block contributes 16 × u32 words to the output stream: +//! +//! ```text +//! qs_u32[block * 16 + w] = LE-u32 of raw_bytes[block * 66 + 2 + w * 4 .. +4] +//! ``` +//! +//! ## Dispatch +//! +//! 1 thread per output u32. Grid: `n_blocks * 16` threads in 1D. +//! Each thread does 4 byte loads + 3 shifts + 3 ORs + 1 u32 store — +//! ~12 ops. Memory-bound on the 4 raw byte reads, but they're +//! contiguous so Apple GPU coalescing is friendly within a simdgroup. + +use metaltile::kernel; + +#[kernel] +pub fn ffai_gguf_iq2_xxs_extract_qs( + raw_bytes: Tensor, + mut qs_u32: Tensor, + #[constexpr] n_blocks: u32, +) { + let w = tid; + let n_words = n_blocks * 16u32; + if w < n_words { + let block = w / 16u32; + let word_in_block = w - block * 16u32; + let byte_offset = block * 66u32 + 2u32 + word_in_block * 4u32; + let b0 = load(raw_bytes[byte_offset]).cast::(); + let b1 = load(raw_bytes[byte_offset + 1u32]).cast::(); + let b2 = load(raw_bytes[byte_offset + 2u32]).cast::(); + let b3 = load(raw_bytes[byte_offset + 3u32]).cast::(); + let packed = b0 | (b1 << 8u32) | (b2 << 16u32) | (b3 << 24u32); + store(qs_u32[w], packed); + } +} + +#[cfg(test)] +pub mod kernel_tests { + + use super::ffai_gguf_iq2_xxs_extract_qs; + + #[test] + fn codegen_extract_qs_smoke() { + let ir = ffai_gguf_iq2_xxs_extract_qs::kernel_ir_for(); + assert!(!ir.body.ops.is_empty(), "kernel body emitted no ops"); + assert!(ir.params.iter().any(|p| p.name == "raw_bytes"), "missing raw_bytes"); + assert!(ir.params.iter().any(|p| p.name == "qs_u32"), "missing qs_u32"); + } +} + +pub mod kernel_benches { + use metaltile::{bench, test::*}; + + use super::ffai_gguf_iq2_xxs_extract_qs; + + #[bench(name = "ffai/gguf_iq2_xxs_extract_qs", dtypes = [u32])] + fn bench_extract_qs(_dt: DType) -> BenchSetup { + let n = 4096 * 4096usize; + let n_blocks = n / 256; + BenchSetup::new(ffai_gguf_iq2_xxs_extract_qs::kernel_ir_for()) + .buffer(BenchBuffer::random("raw_bytes", n_blocks * 66, DType::U8)) + .buffer(BenchBuffer::zeros("qs_u32", n_blocks * 16, DType::U32).output()) + .constexpr("n_blocks", n_blocks as u32) + .grid_1d(n_blocks * 16, 256) + .bytes_moved((n_blocks * 66 + n_blocks * 16 * 4) as u64) + } +} diff --git a/crates/metaltile-std/src/ffai/mod.rs b/crates/metaltile-std/src/ffai/mod.rs index e832e734..e0b3cb05 100644 --- a/crates/metaltile-std/src/ffai/mod.rs +++ b/crates/metaltile-std/src/ffai/mod.rs @@ -26,6 +26,7 @@ pub mod aura_flash_pass2; pub mod aura_flash_sdpa; pub mod aura_score; pub mod aura_value; +pub mod axpy_scalar_inplace; pub mod batched_4_block_scaled_qgemv; pub mod batched_4_block_scaled_qmm; pub mod batched_4_qgemv; @@ -50,10 +51,23 @@ pub mod dequant_gather_block_scaled; pub mod dequant_gemv; pub mod dequant_gemv_expert_indexed; pub mod dequant_gemv_expert_indexed_block_scaled; +pub mod dsv4_compressor_pool; +pub mod dsv4_csa_sdpa_decode; +pub mod dsv4_fp8_block_dequant; +pub mod dsv4_indexer_score; +pub mod dsv4_indexer_topk; +pub mod dsv4_mhc; +pub mod dsv4_mhc_sinkhorn_split; +pub mod dsv4_mxfp4_dequant; +pub mod dsv4_partial_rope; +pub mod dsv4_partial_rope_rows; +pub mod dsv4_router_topk; +pub mod dsv4_swiglu_limit; pub mod fishspeech_conv1d; pub mod fishspeech_conv1d_block_scaled; pub mod flash_block_scaled_sdpa; pub mod flash_quantized_sdpa; +pub mod gate_up_swiglu_fused; pub mod gated_delta; pub mod gated_delta_prep; pub mod gated_delta_prep_chunk; @@ -64,6 +78,15 @@ pub mod gated_rms_norm_qgemv; pub mod gated_rmsnorm; pub mod gather; pub mod gemm; +pub mod gemm_q8; +pub mod gemm_q8_mpp; +pub mod gemv_axpy_inplace; +pub mod gemv_q8; +pub mod gguf_dequant_iq2_xxs; +pub mod gguf_dequant_iq2_xxs_raw; +pub mod gguf_dequant_q2_k; +pub mod gguf_dequant_q8_0; +pub mod gguf_iq2_xxs_extract_qs; pub mod kokoro; pub mod kv_cache; pub mod kv_cache_update_many; @@ -73,7 +96,23 @@ pub mod logits_top_p; pub mod logits_topk; pub mod mel_spectrogram; pub mod moe; +pub mod moe_bgemm_iq2xxs_bm64; +pub mod moe_bgemm_iq2xxs_mpp; +pub mod moe_bgemm_iq2xxs_view; +pub mod moe_bgemm_iq2xxs_view_u16_bm64; +pub mod moe_bgemm_q2k_bm64; +pub mod moe_bgemm_q2k_mpp; +pub mod moe_bgemm_q2k_view; +pub mod moe_bgemm_q2k_view_u16_bm64; pub mod moe_down_swiglu_accum; +pub mod moe_down_weighted_sum_f16; +pub mod moe_gather_down_q2k; +pub mod moe_gather_gemv_iq2xxs; +pub mod moe_gemv_rows_iq2xxs; +pub mod moe_gemv_rows_q2k; +pub mod moe_gemv_rows_view_iq2xxs; +pub mod moe_gemv_ws_iq2xxs; +pub mod moe_gemv_ws_q2k; pub mod moe_mpp; pub mod moe_mpp_block_scaled; pub mod moe_mpp_bm64; @@ -84,6 +123,8 @@ pub mod moe_mpp_bm8_block_scaled; pub mod moe_mpp_bm8_int8; pub mod moe_mpp_int8; pub mod moe_mpp_shared; +pub mod moe_router_sqrtsoftplus; +pub mod mt_vector_add; pub mod patch_embed; pub mod patch_embed_block_scaled; pub mod patch_embed_mma; @@ -102,8 +143,10 @@ pub mod sdpa_bidirectional_d128_relpos; pub mod sdpa_decode; // unified: d64, d96, d128, d256, d512 pub mod sdpa_decode_2pass; pub mod sdpa_decode_batched; +pub mod sdpa_decode_d512_sink; pub mod sdpa_multi; pub mod sdpa_multi_d256; +pub mod sdpa_prefill_d512_sink; pub mod ssm; pub mod ssm_replay; pub mod vocoder; diff --git a/crates/metaltile-std/src/ffai/moe_bgemm_iq2xxs_bm64.rs b/crates/metaltile-std/src/ffai/moe_bgemm_iq2xxs_bm64.rs new file mode 100644 index 00000000..17b877d0 --- /dev/null +++ b/crates/metaltile-std/src/ffai/moe_bgemm_iq2xxs_bm64.rs @@ -0,0 +1,190 @@ +//! Copyright 2026 0xClandestine, Ekryski, TheTom, Ambisphaeric +//! SPDX-License-Identifier: Apache-2.0 +//! HIGH-THROUGHPUT amortized MoE IQ2_XXS grouped BGEMM — the bm64 tiling +//! (64×64×32 tiles, 4 simdgroups / 128 threads, 2×2 warp grid) applied to +//! the IQ2_XXS dequant. The existing moe_bgemm_iq2xxs_mpp uses a tiny +//! 16×32×16 / 1-simdgroup tile (~4-10 GB/s) that is occupancy- and +//! amortization-starved; this mirrors moe_mpp_bm64's fast structure (each +//! 64-row weight tile dequanted once, reused across 64 token rows; 4 SGs for +//! occupancy) to reach high amortized-GEMM throughput. +//! +//! Pool format (slot-major, same as moe_bgemm_iq2xxs_mpp): qs u32 +//! [n_experts*nblk*16], d_f32 [n_experts*nblk], nblk = n_out*k_in/256. +//! `indices[row]` = packed expert/slot. x [m_total,k_in], out [m_total,n_out]. + +use metaltile::kernel; + +#[kernel] +#[allow(clippy::too_many_arguments)] +pub fn ffai_moe_bgemm_iq2xxs_bm64( + x: Tensor, + qs: Tensor, + d_f32: Tensor, + grid: Tensor, + signs: Tensor, + indices: Tensor, + mut out: Tensor, + #[constexpr] m_total: u32, + #[constexpr] n_out: u32, + #[constexpr] k_in: u32, +) { + let n_tile_base = tgid_x * 64u32; + let m_tile_base = tgid_y * 64u32; + let sg = simd_group_id(); + let lane_in_tg = sg * 32u32 + simd_lane; + let sg_m_base = (sg / 2u32) * 32u32; + let sg_n_base = (sg & 1u32) * 32u32; + let nblk_per_expert = n_out * k_in / 256u32; + // X coop-load: 128 lanes × 16 contiguous K = 2048 = BM(64)×TG_LD(32). + let x_m_row = lane_in_tg / 2u32; + let x_k_base = (lane_in_tg & 1u32) * 16u32; + threadgroup_alloc("Xs", 2048, coop_stage(T)); // 64 × 32 + threadgroup_alloc("Ws", 2048, coop_stage(T)); // 64 × 32 + threadgroup_alloc("OutScratch", 4096, f32); // 4 SG × 32 × 32 + coop_tile_setup( + "gemm", + 32, + 32, + 32, + coop_stage(T), + "accumulate", + "simdgroup", + f32, + false, + true, + false, + ); + let mut sub_offset = 0u32; + for _sub_iter in range(0u32, 64u32, 1u32) { + let cur_row = m_tile_base + sub_offset; + let cur_in_range = (sub_offset < 64u32) & (cur_row < m_total); + let cur_expert = select(cur_in_range, load(indices[cur_row]), 4294967295u32); + let mut sub_end = 64u32; + let mut found = 0u32; + for _ii in range(0u32, 64u32, 1u32) { + let probe = sub_offset + 1u32 + _ii; + let probe_row = m_tile_base + probe; + let probe_in_range = (probe < 64u32) & (probe_row < m_total); + if probe_in_range & (found == 0u32) { + let e = load(indices[probe_row]); + if e != cur_expert { + sub_end = probe; + found = 1u32; + } + } + if (probe < 64u32) & (probe_row >= m_total) & (found == 0u32) { + sub_end = probe; + found = 1u32; + } + } + let cur_valid = (cur_expert != 4294967295u32) & (sub_offset < 64u32); + if cur_valid { + let qs_expert_base = cur_expert * nblk_per_expert * 16u32; + let d_expert_base = cur_expert * nblk_per_expert; + coop_tile_zero("gemm"); + for kb in range(0u32, k_in, 32u32) { + // Stage X[m_tile_base..+64, kb..kb+32] → Xs. 128 lanes × 16. + let gr_x = m_tile_base + x_m_row; + let in_run_x = (x_m_row >= sub_offset) & (x_m_row < sub_end) & (gr_x < m_total); + let safe_gr_x = select(in_run_x, gr_x, 0u32); + let x_dev_base = safe_gr_x * k_in + kb + x_k_base; + let x_ws_base = x_m_row * 32u32 + x_k_base; + for _i in range(0u32, 16u32, 1u32) { + let xv = load(x[x_dev_base + _i]).cast::(); + threadgroup_store("Xs", x_ws_base + _i, select(in_run_x, xv, 0.0f32)); + } + // Dequant W[expert, n_tile_base..+64, kb..kb+32] → Ws via the + // IQ2_XXS formula. 128 lanes × 16 = 2048 = BN(64)×BK(32). The + // 16 elements a lane fills all share ONE w_row + ONE 32-block + // (flat=lane*16.._+15 never crosses a 32-boundary; kb is a + // mult of 32), so aux_idx/aux_sgn/db/block are CONSTANT across + // them — hoist those global loads out of the inner loop (16→1, + // the block-at-a-time dequant). Only the octet/sign extraction + // varies per element. + let flat0 = lane_in_tg * 16u32; + let w_row = flat0 / 32u32; // 0..63 (BN row) + let k_local0 = flat0 & 31u32; // 0 or 16 + let global_row = n_tile_base + w_row; + let vidx0 = global_row * k_in + kb + k_local0; + let block = vidx0 / 256u32; + let group = (vidx0 & 255u32) / 32u32; + let aux_idx = load(qs[qs_expert_base + block * 16u32 + group * 2u32]); + let aux_sgn = load(qs[qs_expert_base + block * 16u32 + group * 2u32 + 1u32]); + let scale_4bit = aux_sgn >> 28u32; + let db = load(d_f32[d_expert_base + block]) + * ((scale_4bit.cast::().cast::() + 0.5) * 0.25); + let ws_row_base = w_row * 32u32; + for _i in range(0u32, 16u32, 1u32) { + let k_local = k_local0 + _i; // 0..15 or 16..31 + let octet_within_index = (k_local & 31u32) / 8u32; + let lane_in_octet = k_local & 7u32; + let grid_key = (aux_idx >> (octet_within_index * 8u32)) & 0xffu32; + let octet = load(grid[grid_key * 8u32 + lane_in_octet]) + .cast::() + .cast::() + .cast::(); + let sign_idx = (aux_sgn >> (octet_within_index * 7u32)) & 0x7fu32; + let sign_mask = load(signs[sign_idx]).cast::(); + let lane_bit = sign_mask & (1u32 << lane_in_octet); + let sign = select(lane_bit != 0u32, -1.0f32, 1.0f32); + let w = (db * sign * octet).cast::().cast::(); + threadgroup_store("Ws", ws_row_base + k_local, w); + } + threadgroup_barrier(); + coop_tile_load_a("gemm", "Xs", true, coop_stage(T), 32, 32, sg_m_base * 32u32); + coop_tile_load_b("gemm", "Ws", true, coop_stage(T), 32, 32, sg_n_base * 32u32); + coop_tile_run("gemm"); + threadgroup_barrier(); + } + coop_tile_store_c("gemm", "OutScratch", true, f32, 32, 32, sg * 1024u32); + threadgroup_barrier(); + for _e in range(0u32, 32u32, 1u32) { + let flat = lane_in_tg * 32u32 + _e; + let mr = flat / 64u32; + let nc = flat & 63u32; + let gr = m_tile_base + mr; + let gc = n_tile_base + nc; + let in_run = (mr >= sub_offset) & (mr < sub_end) & (gr < m_total) & (gc < n_out); + if in_run { + let src_sg = (mr / 32u32) * 2u32 + nc / 32u32; + let v = threadgroup_load( + "OutScratch", + src_sg * 1024u32 + (mr & 31u32) * 32u32 + (nc & 31u32), + ); + store(out[gr * n_out + gc], v.cast::()); + } + } + threadgroup_barrier(); + } + sub_offset = sub_end; + } +} + +pub mod kernel_benches { + use metaltile::{bench, test::*}; + + use super::ffai_moe_bgemm_iq2xxs_bm64; + + #[bench(name = "ffai/moe/bgemm_iq2xxs_bm64", dtypes = [f32, f16, bf16])] + fn bench_bgemm_iq2xxs_bm64(dt: DType) -> BenchSetup { + let n_experts = 4usize; + let k_in = 4096usize; + let n_out = 2048usize; + let t_rows = 256usize; + let nblk = n_out * k_in / 256; + BenchSetup::new(ffai_moe_bgemm_iq2xxs_bm64::kernel_ir_for(dt)) + .mode(KernelMode::Reduction) + .buffer(BenchBuffer::random("x", t_rows * k_in, dt)) + .buffer(BenchBuffer::random("qs", n_experts * nblk * 16, DType::U32)) + .buffer(BenchBuffer::random("d_f32", n_experts * nblk, DType::F32)) + .buffer(BenchBuffer::random("grid", 2048, DType::U8)) + .buffer(BenchBuffer::random("signs", 128, DType::U8)) + .buffer(BenchBuffer::zeros("indices", t_rows, DType::U32)) + .buffer(BenchBuffer::zeros("out", t_rows * n_out, dt).output()) + .constexpr("m_total", t_rows as u32) + .constexpr("n_out", n_out as u32) + .constexpr("k_in", k_in as u32) + .grid_3d(n_out as u32 / 64, (t_rows as u32).div_ceil(64), 1, [128, 1, 1]) + .bytes_moved((n_experts * nblk * 64 + t_rows * k_in * dt.size_bytes()) as u64) + } +} diff --git a/crates/metaltile-std/src/ffai/moe_bgemm_iq2xxs_mpp.rs b/crates/metaltile-std/src/ffai/moe_bgemm_iq2xxs_mpp.rs new file mode 100644 index 00000000..2b6502a1 --- /dev/null +++ b/crates/metaltile-std/src/ffai/moe_bgemm_iq2xxs_mpp.rs @@ -0,0 +1,181 @@ +//! Copyright 2026 0xClandestine, Ekryski, TheTom, Ambisphaeric +//! SPDX-License-Identifier: Apache-2.0 +//! MPP MoE IQ2_XXS grouped BGEMM — the prefill counterpart of +//! `moe_gather_gemv_iq2xxs`. Multi-token (M>1) grouped matmul: rows are +//! pre-permuted by expert (`indices[row]` = packed expert slot), and each +//! expert's weight is read+dequanted ONCE and reused across all the tokens +//! routed to it — the amortization that makes MoE prefill fast. +//! +//! Structure is identical to `moe_mpp::mt_moe_gather_qmm_mma_int4_bm16_mpp` +//! (BM=16 / BN=32 / BK=16 coop-tile MMA, contiguous-expert sub-run walk); +//! ONLY the weight-staging loop differs — int4 nibble unpack → IQ2_XXS +//! grid+sign dequant (the same formula as the gemv, see +//! `gguf_dequant_iq2_xxs`). +//! +//! Layout: `qs [n_experts * nblk * 16] u32`, `d_f32 [n_experts * nblk]`, +//! where `nblk = n_out * k_in / 256`; weight value (row r, col k) lives at +//! flat index `r*k_in + k` → block `idx/256`, etc. `x [m_total, k_in]`, +//! `out [m_total, n_out]`. `grid [2048] u8`, `signs [128] u8`. Caller +//! supplies the by-expert-sorted rows + `indices`. + +use metaltile::kernel; + +#[kernel] +#[allow(clippy::too_many_arguments)] +pub fn ffai_moe_gather_bgemm_iq2xxs_mpp( + x: Tensor, + qs: Tensor, + d_f32: Tensor, + grid: Tensor, + signs: Tensor, + indices: Tensor, + mut out: Tensor, + #[constexpr] m_total: u32, + #[constexpr] n_out: u32, + #[constexpr] k_in: u32, +) { + let n_tile_base = tgid_x * 32u32; + let m_tile_base = tgid_y * 16u32; + let lane = simd_lane; + let nblk_per_expert = n_out * k_in / 256u32; + threadgroup_alloc("xs", 256, coop_stage(T)); // 16 × 16 + threadgroup_alloc("ws", 512, coop_stage(T)); // 32 × 16 + threadgroup_alloc("out_scratch", 512, f32); // 16 × 32 + coop_tile_setup( + "gemm", + 16, + 32, + 16, + coop_stage(T), + "accumulate", + "simdgroup", + f32, + false, + true, + false, + ); + let mut sub_offset = 0u32; + for _sub_iter in range(0u32, 16u32, 1u32) { + let cur_row = m_tile_base + sub_offset; + let cur_in_range = (sub_offset < 16u32) & (cur_row < m_total); + let cur_expert = select(cur_in_range, load(indices[cur_row]), 4294967295u32); + let mut sub_end = 16u32; + let mut found = 0u32; + for _ii in range(0u32, 16u32, 1u32) { + let probe = sub_offset + 1u32 + _ii; + let probe_row = m_tile_base + probe; + let probe_in_range = (probe < 16u32) & (probe_row < m_total); + if probe_in_range & (found == 0u32) { + let e = load(indices[probe_row]); + if e != cur_expert { + sub_end = probe; + found = 1u32; + } + } + if (probe < 16u32) & (probe_row >= m_total) & (found == 0u32) { + sub_end = probe; + found = 1u32; + } + } + let cur_valid = (cur_expert != 4294967295u32) & (sub_offset < 16u32); + if cur_valid { + let qs_expert_base = cur_expert * nblk_per_expert * 16u32; + let d_expert_base = cur_expert * nblk_per_expert; + coop_tile_zero("gemm"); + for kb in range(0u32, k_in, 16u32) { + // Stage X[m_tile_base..+16, kb..kb+16] → xs. 32 lanes × 8. + for _e in range(0u32, 8u32, 1u32) { + let flat = lane * 8u32 + _e; + let mr = flat / 16u32; + let kc = flat % 16u32; + let gr = m_tile_base + mr; + let in_run = (mr >= sub_offset) & (mr < sub_end) & (gr < m_total); + let safe_g = select(in_run, gr, 0u32); + let xv = load(x[safe_g * k_in + kb + kc]).cast::(); + threadgroup_store("xs", mr * 16u32 + kc, select(in_run, xv, 0.0f32)); + } + // Dequant W[expert, n_tile_base..+32, kb..kb+16] → ws via the + // IQ2_XXS grid+sign formula. 32 lanes = 32 BN rows; each lane + // dequants its row's 16 K-tile values. + let w_row = lane; // 0..31 + let global_row = n_tile_base + w_row; + for _kc in range(0u32, 16u32, 1u32) { + let k = kb + _kc; + let vidx = global_row * k_in + k; + let block = vidx / 256u32; + let in_block = vidx & 255u32; + let group = in_block / 32u32; + let in_group = in_block & 31u32; + let octet_within_index = in_group / 8u32; + let lane_in_octet = in_group & 7u32; + let aux_idx = load(qs[qs_expert_base + block * 16u32 + group * 2u32]); + let aux_sgn = load(qs[qs_expert_base + block * 16u32 + group * 2u32 + 1u32]); + let scale_4bit = aux_sgn >> 28u32; + let db = load(d_f32[d_expert_base + block]) + * ((scale_4bit.cast::().cast::() + 0.5) * 0.25); + let grid_key = (aux_idx >> (octet_within_index * 8u32)) & 0xffu32; + let octet = load(grid[grid_key * 8u32 + lane_in_octet]) + .cast::() + .cast::() + .cast::(); + let sign_idx = (aux_sgn >> (octet_within_index * 7u32)) & 0x7fu32; + let sign_mask = load(signs[sign_idx]).cast::(); + let lane_bit = sign_mask & (1u32 << lane_in_octet); + let sign = select(lane_bit != 0u32, -1.0f32, 1.0f32); + let w = (db * sign * octet).cast::().cast::(); + threadgroup_store("ws", w_row * 16u32 + _kc, w); + } + threadgroup_barrier(); + coop_tile_load_a("gemm", "xs", true, coop_stage(T), 16, 16); + coop_tile_load_b("gemm", "ws", true, coop_stage(T), 16, 32); + coop_tile_run("gemm"); + threadgroup_barrier(); + } + coop_tile_store_c("gemm", "out_scratch", true, f32, 32, 16); + threadgroup_barrier(); + for _e in range(0u32, 16u32, 1u32) { + let flat = lane * 16u32 + _e; + let mr = flat / 32u32; + let nc = flat % 32u32; + let gr = m_tile_base + mr; + let gc = n_tile_base + nc; + let in_run = (mr >= sub_offset) & (mr < sub_end) & (gr < m_total) & (gc < n_out); + if in_run { + let v = threadgroup_load("out_scratch", mr * 32u32 + nc); + store(out[gr * n_out + gc], v.cast::()); + } + } + threadgroup_barrier(); + } + sub_offset = sub_end; + } +} + +pub mod kernel_benches { + use metaltile::{bench, test::*}; + + use super::ffai_moe_gather_bgemm_iq2xxs_mpp; + + #[bench(name = "ffai/moe/gather_bgemm_iq2xxs_mpp", dtypes = [f32, f16, bf16])] + fn bench_bgemm_iq2xxs_mpp(dt: DType) -> BenchSetup { + let n_experts = 4usize; + let k_in = 4096usize; + let n_out = 2048usize; + let t_rows = 256usize; + let nblk = n_out * k_in / 256; + BenchSetup::new(ffai_moe_gather_bgemm_iq2xxs_mpp::kernel_ir_for(dt)) + .mode(KernelMode::Reduction) + .buffer(BenchBuffer::random("x", t_rows * k_in, dt)) + .buffer(BenchBuffer::random("qs", n_experts * nblk * 16, DType::U32)) + .buffer(BenchBuffer::random("d_f32", n_experts * nblk, DType::F32)) + .buffer(BenchBuffer::random("grid", 2048, DType::U8)) + .buffer(BenchBuffer::random("signs", 128, DType::U8)) + .buffer(BenchBuffer::zeros("indices", t_rows, DType::U32)) + .buffer(BenchBuffer::zeros("out", t_rows * n_out, dt).output()) + .constexpr("m_total", t_rows as u32) + .constexpr("n_out", n_out as u32) + .constexpr("k_in", k_in as u32) + .grid_3d(n_out as u32 / 32, (t_rows as u32).div_ceil(16), 1, [32, 1, 1]) + .bytes_moved((n_experts * nblk * 64 + t_rows * k_in * dt.size_bytes()) as u64) + } +} diff --git a/crates/metaltile-std/src/ffai/moe_bgemm_iq2xxs_view.rs b/crates/metaltile-std/src/ffai/moe_bgemm_iq2xxs_view.rs new file mode 100644 index 00000000..a267144b --- /dev/null +++ b/crates/metaltile-std/src/ffai/moe_bgemm_iq2xxs_view.rs @@ -0,0 +1,204 @@ +//! Copyright 2026 0xClandestine, Ekryski, TheTom, Ambisphaeric +//! SPDX-License-Identifier: Apache-2.0 +//! ZERO-COPY MoE IQ2_XXS grouped BGEMM — reads raw IQ2_XXS blocks straight +//! from a no-copy mmap VIEW buffer (see GGUFModelViews), +//! instead of the CPU-repacked `qs`/`d_f32` pool. Eliminates +//! the per-layer mmap→pool repack that made prefill CPU-bound. +//! +//! Identical MMA structure to `moe_bgemm_iq2xxs_mpp` (BM=16/BN=32/BK=16 +//! coop-tile, by-expert-sorted rows, contiguous-expert sub-run walk). ONLY +//! the weight-staging read differs: the IQ2_XXS block (66 bytes = 1 fp16 d +//! + 32 u16 qs) is read in place from the view. +//! +//! The SAME view buffer is bound twice — as `view_u16` (for the qs u16 +//! pairs that recombine into the aux32 words) and `view_f16` (for the +//! block's fp16 super-scale d). A block's bytes start at u16 index +//! `tensor_u16_off + expert*expert_u16_stride + block*33` +//! (66 bytes = 33 u16); d is u16[base], qs[g] is u16[base+1 + g*4 ..]. +//! `indices[row]` is the EXPERT id (not a packed pool slot). + +use metaltile::kernel; + +#[kernel] +#[allow(clippy::too_many_arguments)] +pub fn ffai_moe_bgemm_iq2xxs_view( + x: Tensor, + view_u8: Tensor, + grid: Tensor, + signs: Tensor, + indices: Tensor, + mut out: Tensor, + #[constexpr] m_total: u32, + #[constexpr] n_out: u32, + #[constexpr] k_in: u32, + #[constexpr] tensor_byte_off: u32, + #[constexpr] expert_byte_stride: u32, +) { + let n_tile_base = tgid_x * 32u32; + let m_tile_base = tgid_y * 16u32; + let lane = simd_lane; + threadgroup_alloc("xs", 256, coop_stage(T)); // 16 × 16 + threadgroup_alloc("ws", 512, coop_stage(T)); // 32 × 16 + threadgroup_alloc("out_scratch", 512, f32); // 16 × 32 + coop_tile_setup( + "gemm", + 16, + 32, + 16, + coop_stage(T), + "accumulate", + "simdgroup", + f32, + false, + true, + false, + ); + let mut sub_offset = 0u32; + for _sub_iter in range(0u32, 16u32, 1u32) { + let cur_row = m_tile_base + sub_offset; + let cur_in_range = (sub_offset < 16u32) & (cur_row < m_total); + let cur_expert = select(cur_in_range, load(indices[cur_row]), 4294967295u32); + let mut sub_end = 16u32; + let mut found = 0u32; + for _ii in range(0u32, 16u32, 1u32) { + let probe = sub_offset + 1u32 + _ii; + let probe_row = m_tile_base + probe; + let probe_in_range = (probe < 16u32) & (probe_row < m_total); + if probe_in_range & (found == 0u32) { + let e = load(indices[probe_row]); + if e != cur_expert { + sub_end = probe; + found = 1u32; + } + } + if (probe < 16u32) & (probe_row >= m_total) & (found == 0u32) { + sub_end = probe; + found = 1u32; + } + } + let cur_valid = (cur_expert != 4294967295u32) & (sub_offset < 16u32); + if cur_valid { + let expert_byte_base = tensor_byte_off + cur_expert * expert_byte_stride; + coop_tile_zero("gemm"); + for kb in range(0u32, k_in, 16u32) { + // Stage X[m_tile_base..+16, kb..kb+16] → xs. 32 lanes × 8. + for _e in range(0u32, 8u32, 1u32) { + let flat = lane * 8u32 + _e; + let mr = flat / 16u32; + let kc = flat % 16u32; + let gr = m_tile_base + mr; + let in_run = (mr >= sub_offset) & (mr < sub_end) & (gr < m_total); + let safe_g = select(in_run, gr, 0u32); + let xv = load(x[safe_g * k_in + kb + kc]).cast::(); + threadgroup_store("xs", mr * 16u32 + kc, select(in_run, xv, 0.0f32)); + } + // Dequant W[expert, n_tile_base..+32, kb..kb+16] → ws, + // reading raw IQ2_XXS blocks straight from the mmap view. + let w_row = lane; // 0..31 + let global_row = n_tile_base + w_row; + for _kc in range(0u32, 16u32, 1u32) { + let k = kb + _kc; + let vidx = global_row * k_in + k; + let block = vidx / 256u32; + let in_block = vidx & 255u32; + let group = in_block / 32u32; + let in_group = in_block & 31u32; + let octet_within_index = in_group / 8u32; + let lane_in_octet = in_group & 7u32; + // Block base in BYTES (66 bytes/block: 2-byte fp16 d + 64-byte qs). + let blk_byte = expert_byte_base + block * 66u32; + // d = block's leading fp16, decoded from 2 raw bytes (no + // buffer aliasing, no CPU extract). IEEE binary16: s|eeeee|mmmmmmmmmm. + let d_lo = load(view_u8[blk_byte]).cast::(); + let d_hi = load(view_u8[blk_byte + 1u32]).cast::(); + let d_bits = d_lo | (d_hi << 8u32); + let d_sign = select((d_bits & 0x8000u32) != 0u32, 0.0f32 - 1.0f32, 1.0f32); + let d_exp = (d_bits >> 10u32) & 0x1fu32; + let d_mant = (d_bits & 0x3ffu32).cast::().cast::(); + // normal: (1 + m/1024)·2^(e-15); subnormal (e==0): m·2^-24. + let d_norm = (1.0f32 + d_mant / 1024.0f32) + * exp2(d_exp.cast::().cast::() - 15.0f32); + let d_sub = d_mant * exp2(0.0f32 - 24.0f32); + let dval = d_sign * select(d_exp == 0u32, d_sub, d_norm); + // qs starts at byte blk_byte+2; group g uses 8 bytes + // (aux_idx u32 + aux_sgn u32), little-endian recombine from u8. + let q0 = blk_byte + 2u32 + group * 8u32; + let aux_idx = load(view_u8[q0]).cast::() + | (load(view_u8[q0 + 1u32]).cast::() << 8u32) + | (load(view_u8[q0 + 2u32]).cast::() << 16u32) + | (load(view_u8[q0 + 3u32]).cast::() << 24u32); + let aux_sgn = load(view_u8[q0 + 4u32]).cast::() + | (load(view_u8[q0 + 5u32]).cast::() << 8u32) + | (load(view_u8[q0 + 6u32]).cast::() << 16u32) + | (load(view_u8[q0 + 7u32]).cast::() << 24u32); + let scale_4bit = aux_sgn >> 28u32; + let db = dval * ((scale_4bit.cast::().cast::() + 0.5) * 0.25); + let grid_key = (aux_idx >> (octet_within_index * 8u32)) & 0xffu32; + let octet = load(grid[grid_key * 8u32 + lane_in_octet]) + .cast::() + .cast::() + .cast::(); + let sign_idx = (aux_sgn >> (octet_within_index * 7u32)) & 0x7fu32; + let sign_mask = load(signs[sign_idx]).cast::(); + let lane_bit = sign_mask & (1u32 << lane_in_octet); + let sign = select(lane_bit != 0u32, -1.0f32, 1.0f32); + let w = (db * sign * octet).cast::().cast::(); + threadgroup_store("ws", w_row * 16u32 + _kc, w); + } + threadgroup_barrier(); + coop_tile_load_a("gemm", "xs", true, coop_stage(T), 16, 16); + coop_tile_load_b("gemm", "ws", true, coop_stage(T), 16, 32); + coop_tile_run("gemm"); + threadgroup_barrier(); + } + coop_tile_store_c("gemm", "out_scratch", true, f32, 32, 16); + threadgroup_barrier(); + for _e in range(0u32, 16u32, 1u32) { + let flat = lane * 16u32 + _e; + let mr = flat / 32u32; + let nc = flat % 32u32; + let gr = m_tile_base + mr; + let gc = n_tile_base + nc; + let in_run = (mr >= sub_offset) & (mr < sub_end) & (gr < m_total) & (gc < n_out); + if in_run { + let v = threadgroup_load("out_scratch", mr * 32u32 + nc); + store(out[gr * n_out + gc], v.cast::()); + } + } + threadgroup_barrier(); + } + sub_offset = sub_end; + } +} + +pub mod kernel_benches { + use metaltile::{bench, test::*}; + + use super::ffai_moe_bgemm_iq2xxs_view; + + #[bench(name = "ffai/moe/bgemm_iq2xxs_view", dtypes = [f32, f16, bf16])] + fn bench_bgemm_iq2xxs_view(dt: DType) -> BenchSetup { + let n_experts = 4usize; + let k_in = 4096usize; + let n_out = 2048usize; + let t_rows = 256usize; + let nblk = n_out * k_in / 256; + // view holds n_experts × nblk IQ2 blocks of 66 bytes each. + let view_bytes = n_experts * nblk * 66; + BenchSetup::new(ffai_moe_bgemm_iq2xxs_view::kernel_ir_for(dt)) + .mode(KernelMode::Reduction) + .buffer(BenchBuffer::random("x", t_rows * k_in, dt)) + .buffer(BenchBuffer::random("view_u8", view_bytes, DType::U8)) + .buffer(BenchBuffer::random("grid", 2048, DType::U8)) + .buffer(BenchBuffer::random("signs", 128, DType::U8)) + .buffer(BenchBuffer::zeros("indices", t_rows, DType::U32)) + .buffer(BenchBuffer::zeros("out", t_rows * n_out, dt).output()) + .constexpr("m_total", t_rows as u32) + .constexpr("n_out", n_out as u32) + .constexpr("k_in", k_in as u32) + .constexpr("tensor_byte_off", 0u32) + .constexpr("expert_byte_stride", (nblk * 66) as u32) + .grid_3d(n_out as u32 / 32, (t_rows as u32).div_ceil(16), 1, [32, 1, 1]) + .bytes_moved((n_experts * nblk * 66 + t_rows * k_in * dt.size_bytes()) as u64) + } +} diff --git a/crates/metaltile-std/src/ffai/moe_bgemm_iq2xxs_view_u16_bm64.rs b/crates/metaltile-std/src/ffai/moe_bgemm_iq2xxs_view_u16_bm64.rs new file mode 100644 index 00000000..07adb09f --- /dev/null +++ b/crates/metaltile-std/src/ffai/moe_bgemm_iq2xxs_view_u16_bm64.rs @@ -0,0 +1,194 @@ +//! Copyright 2026 0xClandestine, Ekryski, TheTom, Ambisphaeric +//! SPDX-License-Identifier: Apache-2.0 +//! POOL-FREE amortized MoE IQ2_XXS bgemm — bm64 MMA tiling reading RAW +//! resident GGUF blocks directly via aligned u16 (DType::U16), eliminating +//! the per-layer CPU repack into a split pool. Combines: +//! - moe_bgemm_iq2xxs_bm64's 64×64×32 coop_tile MMA + dequant-hoist, and +//! - the u16 raw-block read (114 GB/s, vs the u8-recombine view's 3.5). +//! +//! `indices[row]` = GLOBAL expert id (not a pool slot). Weights read from +//! the resident mmap view at `tensor_byte_off + expert*expert_byte_stride`, +//! block stride 66 bytes (f16 d + u16[32] qs). This is the no-repack +//! architecture: read the resident weights in-place, no second copy. +//! x [m_total,k_in], out [m_total,n_out]. Live-compiled (name has bgemm). + +use metaltile::kernel; + +#[kernel] +#[allow(clippy::too_many_arguments)] +pub fn ffai_moe_bgemm_iq2xxs_view_u16_bm64( + x: Tensor, + view_u16: Tensor, + view_f16: Tensor, + grid: Tensor, + signs: Tensor, + indices: Tensor, + mut out: Tensor, + #[constexpr] m_total: u32, + #[constexpr] n_out: u32, + #[constexpr] k_in: u32, + #[constexpr] tensor_byte_off: u32, + #[constexpr] expert_byte_stride: u32, +) { + let n_tile_base = tgid_x * 64u32; + let m_tile_base = tgid_y * 64u32; + let sg = simd_group_id(); + let lane_in_tg = sg * 32u32 + simd_lane; + let sg_m_base = (sg / 2u32) * 32u32; + let sg_n_base = (sg & 1u32) * 32u32; + let x_m_row = lane_in_tg / 2u32; + let x_k_base = (lane_in_tg & 1u32) * 16u32; + threadgroup_alloc("Xs", 2048, coop_stage(T)); + threadgroup_alloc("Ws", 2048, coop_stage(T)); + threadgroup_alloc("OutScratch", 4096, f32); + coop_tile_setup( + "gemm", + 32, + 32, + 32, + coop_stage(T), + "accumulate", + "simdgroup", + f32, + false, + true, + false, + ); + let mut sub_offset = 0u32; + for _sub_iter in range(0u32, 64u32, 1u32) { + let cur_row = m_tile_base + sub_offset; + let cur_in_range = (sub_offset < 64u32) & (cur_row < m_total); + let cur_expert = select(cur_in_range, load(indices[cur_row]), 4294967295u32); + let mut sub_end = 64u32; + let mut found = 0u32; + for _ii in range(0u32, 64u32, 1u32) { + let probe = sub_offset + 1u32 + _ii; + let probe_row = m_tile_base + probe; + let probe_in_range = (probe < 64u32) & (probe_row < m_total); + if probe_in_range & (found == 0u32) { + let e = load(indices[probe_row]); + if e != cur_expert { + sub_end = probe; + found = 1u32; + } + } + if (probe < 64u32) & (probe_row >= m_total) & (found == 0u32) { + sub_end = probe; + found = 1u32; + } + } + let cur_valid = (cur_expert != 4294967295u32) & (sub_offset < 64u32); + if cur_valid { + // u16 base of this expert's raw blocks in the resident view. + let expert_u16_base = (tensor_byte_off + cur_expert * expert_byte_stride) / 2u32; + coop_tile_zero("gemm"); + for kb in range(0u32, k_in, 32u32) { + let gr_x = m_tile_base + x_m_row; + let in_run_x = (x_m_row >= sub_offset) & (x_m_row < sub_end) & (gr_x < m_total); + let safe_gr_x = select(in_run_x, gr_x, 0u32); + let x_dev_base = safe_gr_x * k_in + kb + x_k_base; + let x_ws_base = x_m_row * 32u32 + x_k_base; + for _i in range(0u32, 16u32, 1u32) { + let xv = load(x[x_dev_base + _i]).cast::(); + threadgroup_store("Xs", x_ws_base + _i, select(in_run_x, xv, 0.0f32)); + } + // Dequant W → Ws from the RAW block via u16. 16 elems/lane + // share one w_row + 32-block → hoist d/aux (read once). + let flat0 = lane_in_tg * 16u32; + let w_row = flat0 / 32u32; + let k_local0 = flat0 & 31u32; + let global_row = n_tile_base + w_row; + let vidx0 = global_row * k_in + kb + k_local0; + let block = vidx0 / 256u32; + let group = (vidx0 & 255u32) / 32u32; + // raw block at u16 offset expert_u16_base + block*33 (66 bytes). + let blk_u16 = expert_u16_base + block * 33u32; + // d = native f16 at the block start (byte = blk_u16*2). Read via + // the f16 view (Metal `half`) — exact, vs a hand-rolled decode. + let dval = load(view_f16[blk_u16]).cast::(); + // aux for this group at u16 offset blk_u16 + 1 + group*4. + let q = blk_u16 + 1u32 + group * 4u32; + // ushort→uint is a zero-extend (static_cast), so the combine is exact. + let aux_idx = load(view_u16[q]).cast::() + | (load(view_u16[q + 1u32]).cast::() << 16u32); + let aux_sgn = load(view_u16[q + 2u32]).cast::() + | (load(view_u16[q + 3u32]).cast::() << 16u32); + let scale_4bit = aux_sgn >> 28u32; + let db = dval * ((scale_4bit.cast::().cast::() + 0.5) * 0.25); + let ws_row_base = w_row * 32u32; + for _i in range(0u32, 16u32, 1u32) { + let k_local = k_local0 + _i; + let octet_within_index = (k_local & 31u32) / 8u32; + let lane_in_octet = k_local & 7u32; + let grid_key = (aux_idx >> (octet_within_index * 8u32)) & 0xffu32; + let octet = load(grid[grid_key * 8u32 + lane_in_octet]) + .cast::() + .cast::() + .cast::(); + let sign_idx = (aux_sgn >> (octet_within_index * 7u32)) & 0x7fu32; + let sign_mask = load(signs[sign_idx]).cast::(); + let lane_bit = sign_mask & (1u32 << lane_in_octet); + let sign = select(lane_bit != 0u32, -1.0f32, 1.0f32); + let w = (db * sign * octet).cast::().cast::(); + threadgroup_store("Ws", ws_row_base + k_local, w); + } + threadgroup_barrier(); + coop_tile_load_a("gemm", "Xs", true, coop_stage(T), 32, 32, sg_m_base * 32u32); + coop_tile_load_b("gemm", "Ws", true, coop_stage(T), 32, 32, sg_n_base * 32u32); + coop_tile_run("gemm"); + threadgroup_barrier(); + } + coop_tile_store_c("gemm", "OutScratch", true, f32, 32, 32, sg * 1024u32); + threadgroup_barrier(); + for _e in range(0u32, 32u32, 1u32) { + let flat = lane_in_tg * 32u32 + _e; + let mr = flat / 64u32; + let nc = flat & 63u32; + let gr = m_tile_base + mr; + let gc = n_tile_base + nc; + let in_run = (mr >= sub_offset) & (mr < sub_end) & (gr < m_total) & (gc < n_out); + if in_run { + let src_sg = (mr / 32u32) * 2u32 + nc / 32u32; + let v = threadgroup_load( + "OutScratch", + src_sg * 1024u32 + (mr & 31u32) * 32u32 + (nc & 31u32), + ); + store(out[gr * n_out + gc], v.cast::()); + } + } + threadgroup_barrier(); + } + sub_offset = sub_end; + } +} + +pub mod kernel_benches { + use metaltile::{bench, test::*}; + + use super::ffai_moe_bgemm_iq2xxs_view_u16_bm64; + + #[bench(name = "ffai/moe/bgemm_iq2xxs_view_u16_bm64", dtypes = [f32, f16, bf16])] + fn bench(dt: DType) -> BenchSetup { + let n_experts = 4usize; + let k_in = 4096usize; + let n_out = 2048usize; + let t_rows = 256usize; + let nblk = n_out * k_in / 256; + BenchSetup::new(ffai_moe_bgemm_iq2xxs_view_u16_bm64::kernel_ir_for(dt)) + .mode(KernelMode::Reduction) + .buffer(BenchBuffer::random("x", t_rows * k_in, dt)) + .buffer(BenchBuffer::random("view_u16", n_experts * nblk * 33, DType::U16)) + .buffer(BenchBuffer::random("view_f16", n_experts * nblk * 33, DType::F16)) + .buffer(BenchBuffer::random("grid", 2048, DType::U8)) + .buffer(BenchBuffer::random("signs", 128, DType::U8)) + .buffer(BenchBuffer::zeros("indices", t_rows, DType::U32)) + .buffer(BenchBuffer::zeros("out", t_rows * n_out, dt).output()) + .constexpr("m_total", t_rows as u32) + .constexpr("n_out", n_out as u32) + .constexpr("k_in", k_in as u32) + .constexpr("tensor_byte_off", 0u32) + .constexpr("expert_byte_stride", (nblk * 66) as u32) + .grid_3d(n_out as u32 / 64, (t_rows as u32).div_ceil(64), 1, [128, 1, 1]) + .bytes_moved((n_experts * nblk * 64 + t_rows * k_in * dt.size_bytes()) as u64) + } +} diff --git a/crates/metaltile-std/src/ffai/moe_bgemm_q2k_bm64.rs b/crates/metaltile-std/src/ffai/moe_bgemm_q2k_bm64.rs new file mode 100644 index 00000000..09e1e50d --- /dev/null +++ b/crates/metaltile-std/src/ffai/moe_bgemm_q2k_bm64.rs @@ -0,0 +1,187 @@ +//! Copyright 2026 0xClandestine, Ekryski, TheTom, Ambisphaeric +//! SPDX-License-Identifier: Apache-2.0 +//! HIGH-THROUGHPUT amortized MoE Q2_K grouped BGEMM (down) — bm64 tiling +//! (64×64×32, 4 simdgroups) with the canonical Q2_K dequant. Q2_K twin of +//! moe_bgemm_iq2xxs_bm64; replaces the slow 16×32 moe_bgemm_q2k_mpp. +//! Pool format: qs u32 [n_experts*nblk*16], scales u8 [n_experts*nblk*16], +//! d_f32/dmin_f32 [n_experts*nblk], nblk = n_out*k_in/256. + +use metaltile::kernel; + +#[kernel] +#[allow(clippy::too_many_arguments)] +pub fn ffai_moe_bgemm_q2k_bm64( + x: Tensor, + qs: Tensor, + scales: Tensor, + d_f32: Tensor, + dmin_f32: Tensor, + indices: Tensor, + mut out: Tensor, + #[constexpr] m_total: u32, + #[constexpr] n_out: u32, + #[constexpr] k_in: u32, +) { + let n_tile_base = tgid_x * 64u32; + let m_tile_base = tgid_y * 64u32; + let sg = simd_group_id(); + let lane_in_tg = sg * 32u32 + simd_lane; + let sg_m_base = (sg / 2u32) * 32u32; + let sg_n_base = (sg & 1u32) * 32u32; + let nblk_per_expert = n_out * k_in / 256u32; + let x_m_row = lane_in_tg / 2u32; + let x_k_base = (lane_in_tg & 1u32) * 16u32; + threadgroup_alloc("Xs", 2048, coop_stage(T)); + threadgroup_alloc("Ws", 2048, coop_stage(T)); + threadgroup_alloc("OutScratch", 4096, f32); + coop_tile_setup( + "gemm", + 32, + 32, + 32, + coop_stage(T), + "accumulate", + "simdgroup", + f32, + false, + true, + false, + ); + let mut sub_offset = 0u32; + for _sub_iter in range(0u32, 64u32, 1u32) { + let cur_row = m_tile_base + sub_offset; + let cur_in_range = (sub_offset < 64u32) & (cur_row < m_total); + let cur_expert = select(cur_in_range, load(indices[cur_row]), 4294967295u32); + let mut sub_end = 64u32; + let mut found = 0u32; + for _ii in range(0u32, 64u32, 1u32) { + let probe = sub_offset + 1u32 + _ii; + let probe_row = m_tile_base + probe; + let probe_in_range = (probe < 64u32) & (probe_row < m_total); + if probe_in_range & (found == 0u32) { + let e = load(indices[probe_row]); + if e != cur_expert { + sub_end = probe; + found = 1u32; + } + } + if (probe < 64u32) & (probe_row >= m_total) & (found == 0u32) { + sub_end = probe; + found = 1u32; + } + } + let cur_valid = (cur_expert != 4294967295u32) & (sub_offset < 64u32); + if cur_valid { + let qs_expert_base = cur_expert * nblk_per_expert * 16u32; + let sc_expert_base = cur_expert * nblk_per_expert * 16u32; + let blk_expert_base = cur_expert * nblk_per_expert; + coop_tile_zero("gemm"); + for kb in range(0u32, k_in, 32u32) { + let gr_x = m_tile_base + x_m_row; + let in_run_x = (x_m_row >= sub_offset) & (x_m_row < sub_end) & (gr_x < m_total); + let safe_gr_x = select(in_run_x, gr_x, 0u32); + let x_dev_base = safe_gr_x * k_in + kb + x_k_base; + let x_ws_base = x_m_row * 32u32 + x_k_base; + for _i in range(0u32, 16u32, 1u32) { + let xv = load(x[x_dev_base + _i]).cast::(); + threadgroup_store("Xs", x_ws_base + _i, select(in_run_x, xv, 0.0f32)); + } + // Dequant Q2_K W → Ws (canonical layout). 128 lanes × 16. The + // 16 elements share ONE w_row + ONE 256-block, so block/d/dmin + // are CONSTANT — hoist those global loads out (d_f32/dmin_f32 + // 16→1 each). qs-word + scales still vary per element. + let flat0 = lane_in_tg * 16u32; + let w_row = flat0 / 32u32; + let k_local0 = flat0 & 31u32; + let global_row = n_tile_base + w_row; + let block = (global_row * k_in + kb + k_local0) / 256u32; + let d = load(d_f32[blk_expert_base + block]); + let dmin = load(dmin_f32[blk_expert_base + block]); + let ws_row_base = w_row * 32u32; + for _i in range(0u32, 16u32, 1u32) { + let k_local = k_local0 + _i; + let k = kb + k_local; + let vidx = global_row * k_in + k; + let in_block = vidx & 255u32; + let half = in_block / 128u32; + let yh = in_block - half * 128u32; + let jg = yh / 32u32; + let yg = yh - jg * 32u32; + let sub_half = yg / 16u32; + let l = yg - sub_half * 16u32; + let shift = jg * 2u32; + let q_byte = half * 32u32 + sub_half * 16u32 + l; + let sub = half * 8u32 + jg * 2u32 + sub_half; + let word_idx = q_byte / 4u32; + let byte_in_word = q_byte & 3u32; + let word = load(qs[qs_expert_base + block * 16u32 + word_idx]); + let qs_byte = (word >> (byte_in_word * 8u32)) & 0xffu32; + let q_2bit = (qs_byte >> shift) & 0x3u32; + let scale_byte = + load(scales[sc_expert_base + block * 16u32 + sub]).cast::(); + let scale_4bit = scale_byte & 0xfu32; + let min_4bit = (scale_byte >> 4u32) & 0xfu32; + let wq = d + * scale_4bit.cast::().cast::() + * q_2bit.cast::().cast::() + - dmin * min_4bit.cast::().cast::(); + threadgroup_store("Ws", ws_row_base + k_local, wq.cast::().cast::()); + } + threadgroup_barrier(); + coop_tile_load_a("gemm", "Xs", true, coop_stage(T), 32, 32, sg_m_base * 32u32); + coop_tile_load_b("gemm", "Ws", true, coop_stage(T), 32, 32, sg_n_base * 32u32); + coop_tile_run("gemm"); + threadgroup_barrier(); + } + coop_tile_store_c("gemm", "OutScratch", true, f32, 32, 32, sg * 1024u32); + threadgroup_barrier(); + for _e in range(0u32, 32u32, 1u32) { + let flat = lane_in_tg * 32u32 + _e; + let mr = flat / 64u32; + let nc = flat & 63u32; + let gr = m_tile_base + mr; + let gc = n_tile_base + nc; + let in_run = (mr >= sub_offset) & (mr < sub_end) & (gr < m_total) & (gc < n_out); + if in_run { + let src_sg = (mr / 32u32) * 2u32 + nc / 32u32; + let v = threadgroup_load( + "OutScratch", + src_sg * 1024u32 + (mr & 31u32) * 32u32 + (nc & 31u32), + ); + store(out[gr * n_out + gc], v.cast::()); + } + } + threadgroup_barrier(); + } + sub_offset = sub_end; + } +} + +pub mod kernel_benches { + use metaltile::{bench, test::*}; + + use super::ffai_moe_bgemm_q2k_bm64; + + #[bench(name = "ffai/moe/bgemm_q2k_bm64", dtypes = [f32, f16, bf16])] + fn bench_bgemm_q2k_bm64(dt: DType) -> BenchSetup { + let n_experts = 4usize; + let k_in = 2048usize; + let n_out = 4096usize; + let t_rows = 256usize; + let nblk = n_out * k_in / 256; + BenchSetup::new(ffai_moe_bgemm_q2k_bm64::kernel_ir_for(dt)) + .mode(KernelMode::Reduction) + .buffer(BenchBuffer::random("x", t_rows * k_in, dt)) + .buffer(BenchBuffer::random("qs", n_experts * nblk * 16, DType::U32)) + .buffer(BenchBuffer::random("scales", n_experts * nblk * 16, DType::U8)) + .buffer(BenchBuffer::random("d_f32", n_experts * nblk, DType::F32)) + .buffer(BenchBuffer::random("dmin_f32", n_experts * nblk, DType::F32)) + .buffer(BenchBuffer::zeros("indices", t_rows, DType::U32)) + .buffer(BenchBuffer::zeros("out", t_rows * n_out, dt).output()) + .constexpr("m_total", t_rows as u32) + .constexpr("n_out", n_out as u32) + .constexpr("k_in", k_in as u32) + .grid_3d(n_out as u32 / 64, (t_rows as u32).div_ceil(64), 1, [128, 1, 1]) + .bytes_moved((n_experts * nblk * 84 + t_rows * k_in * dt.size_bytes()) as u64) + } +} diff --git a/crates/metaltile-std/src/ffai/moe_bgemm_q2k_mpp.rs b/crates/metaltile-std/src/ffai/moe_bgemm_q2k_mpp.rs new file mode 100644 index 00000000..adfd796e --- /dev/null +++ b/crates/metaltile-std/src/ffai/moe_bgemm_q2k_mpp.rs @@ -0,0 +1,179 @@ +//! Copyright 2026 0xClandestine, Ekryski, TheTom, Ambisphaeric +//! SPDX-License-Identifier: Apache-2.0 +//! MPP MoE Q2_K grouped BGEMM — prefill counterpart of +//! `moe_gather_down_q2k` (the down projection). Multi-token grouped matmul; +//! rows pre-permuted by expert, each expert's Q2_K weight read+dequanted +//! once and reused across its tokens via the MPP coop-tile MMA. Structure +//! identical to `moe_bgemm_iq2xxs_mpp`; only the weight dequant differs +//! (Q2_K: 4-bit scale/min sub-blocks × 2-bit quants — see +//! `gguf_dequant_q2_k`). +//! +//! Layout: `qs [n_experts*nblk*16] u32` (64 qs bytes/block as 16 u32), +//! `scales [n_experts*nblk*16] u8`, `d_f32`/`dmin_f32 [n_experts*nblk]`, +//! `nblk = n_out*k_in/256`. `x [m_total,k_in]`, `out [m_total,n_out]`. + +use metaltile::kernel; + +#[kernel] +#[allow(clippy::too_many_arguments)] +pub fn ffai_moe_gather_bgemm_q2k_mpp( + x: Tensor, + qs: Tensor, + scales: Tensor, + d_f32: Tensor, + dmin_f32: Tensor, + indices: Tensor, + mut out: Tensor, + #[constexpr] m_total: u32, + #[constexpr] n_out: u32, + #[constexpr] k_in: u32, +) { + let n_tile_base = tgid_x * 32u32; + let m_tile_base = tgid_y * 16u32; + let lane = simd_lane; + let nblk_per_expert = n_out * k_in / 256u32; + threadgroup_alloc("xs", 256, coop_stage(T)); + threadgroup_alloc("ws", 512, coop_stage(T)); + threadgroup_alloc("out_scratch", 512, f32); + coop_tile_setup( + "gemm", + 16, + 32, + 16, + coop_stage(T), + "accumulate", + "simdgroup", + f32, + false, + true, + false, + ); + let mut sub_offset = 0u32; + for _sub_iter in range(0u32, 16u32, 1u32) { + let cur_row = m_tile_base + sub_offset; + let cur_in_range = (sub_offset < 16u32) & (cur_row < m_total); + let cur_expert = select(cur_in_range, load(indices[cur_row]), 4294967295u32); + let mut sub_end = 16u32; + let mut found = 0u32; + for _ii in range(0u32, 16u32, 1u32) { + let probe = sub_offset + 1u32 + _ii; + let probe_row = m_tile_base + probe; + let probe_in_range = (probe < 16u32) & (probe_row < m_total); + if probe_in_range & (found == 0u32) { + let e = load(indices[probe_row]); + if e != cur_expert { + sub_end = probe; + found = 1u32; + } + } + if (probe < 16u32) & (probe_row >= m_total) & (found == 0u32) { + sub_end = probe; + found = 1u32; + } + } + let cur_valid = (cur_expert != 4294967295u32) & (sub_offset < 16u32); + if cur_valid { + let qs_expert_base = cur_expert * nblk_per_expert * 16u32; + let sc_expert_base = cur_expert * nblk_per_expert * 16u32; + let blk_expert_base = cur_expert * nblk_per_expert; + coop_tile_zero("gemm"); + for kb in range(0u32, k_in, 16u32) { + for _e in range(0u32, 8u32, 1u32) { + let flat = lane * 8u32 + _e; + let mr = flat / 16u32; + let kc = flat % 16u32; + let gr = m_tile_base + mr; + let in_run = (mr >= sub_offset) & (mr < sub_end) & (gr < m_total); + let safe_g = select(in_run, gr, 0u32); + let xv = load(x[safe_g * k_in + kb + kc]).cast::(); + threadgroup_store("xs", mr * 16u32 + kc, select(in_run, xv, 0.0f32)); + } + // Dequant Q2_K W[expert, n_tile_base+lane, kb..kb+16] → ws. + let w_row = lane; + let global_row = n_tile_base + w_row; + for _kc in range(0u32, 16u32, 1u32) { + let k = kb + _kc; + let vidx = global_row * k_in + k; + let block = vidx / 256u32; + let in_block = vidx & 255u32; + // Canonical Q2_K layout (see gguf_dequant_q2_k.rs). + let half = in_block / 128u32; + let yh = in_block - half * 128u32; + let jg = yh / 32u32; + let yg = yh - jg * 32u32; + let sub_half = yg / 16u32; + let l = yg - sub_half * 16u32; + let shift = jg * 2u32; + let q_byte = half * 32u32 + sub_half * 16u32 + l; + let sub = half * 8u32 + jg * 2u32 + sub_half; + let word_idx = q_byte / 4u32; + let byte_in_word = q_byte & 3u32; + let word = load(qs[qs_expert_base + block * 16u32 + word_idx]); + let qs_byte = (word >> (byte_in_word * 8u32)) & 0xffu32; + let q_2bit = (qs_byte >> shift) & 0x3u32; + let scale_byte = + load(scales[sc_expert_base + block * 16u32 + sub]).cast::(); + let scale_4bit = scale_byte & 0xfu32; + let min_4bit = (scale_byte >> 4u32) & 0xfu32; + let d = load(d_f32[blk_expert_base + block]); + let dmin = load(dmin_f32[blk_expert_base + block]); + let wq = d + * scale_4bit.cast::().cast::() + * q_2bit.cast::().cast::() + - dmin * min_4bit.cast::().cast::(); + threadgroup_store("ws", w_row * 16u32 + _kc, wq.cast::().cast::()); + } + threadgroup_barrier(); + coop_tile_load_a("gemm", "xs", true, coop_stage(T), 16, 16); + coop_tile_load_b("gemm", "ws", true, coop_stage(T), 16, 32); + coop_tile_run("gemm"); + threadgroup_barrier(); + } + coop_tile_store_c("gemm", "out_scratch", true, f32, 32, 16); + threadgroup_barrier(); + for _e in range(0u32, 16u32, 1u32) { + let flat = lane * 16u32 + _e; + let mr = flat / 32u32; + let nc = flat % 32u32; + let gr = m_tile_base + mr; + let gc = n_tile_base + nc; + let in_run = (mr >= sub_offset) & (mr < sub_end) & (gr < m_total) & (gc < n_out); + if in_run { + let v = threadgroup_load("out_scratch", mr * 32u32 + nc); + store(out[gr * n_out + gc], v.cast::()); + } + } + threadgroup_barrier(); + } + sub_offset = sub_end; + } +} + +pub mod kernel_benches { + use metaltile::{bench, test::*}; + + use super::ffai_moe_gather_bgemm_q2k_mpp; + + #[bench(name = "ffai/moe/gather_bgemm_q2k_mpp", dtypes = [f32, f16, bf16])] + fn bench_bgemm_q2k_mpp(dt: DType) -> BenchSetup { + let n_experts = 4usize; + let k_in = 2048usize; + let n_out = 4096usize; + let t_rows = 256usize; + let nblk = n_out * k_in / 256; + BenchSetup::new(ffai_moe_gather_bgemm_q2k_mpp::kernel_ir_for(dt)) + .mode(KernelMode::Reduction) + .buffer(BenchBuffer::random("x", t_rows * k_in, dt)) + .buffer(BenchBuffer::random("qs", n_experts * nblk * 16, DType::U32)) + .buffer(BenchBuffer::random("scales", n_experts * nblk * 16, DType::U8)) + .buffer(BenchBuffer::random("d_f32", n_experts * nblk, DType::F32)) + .buffer(BenchBuffer::random("dmin_f32", n_experts * nblk, DType::F32)) + .buffer(BenchBuffer::zeros("indices", t_rows, DType::U32)) + .buffer(BenchBuffer::zeros("out", t_rows * n_out, dt).output()) + .constexpr("m_total", t_rows as u32) + .constexpr("n_out", n_out as u32) + .constexpr("k_in", k_in as u32) + .grid_3d(n_out as u32 / 32, (t_rows as u32).div_ceil(16), 1, [32, 1, 1]) + .bytes_moved((n_experts * nblk * 84 + t_rows * k_in * dt.size_bytes()) as u64) + } +} diff --git a/crates/metaltile-std/src/ffai/moe_bgemm_q2k_view.rs b/crates/metaltile-std/src/ffai/moe_bgemm_q2k_view.rs new file mode 100644 index 00000000..44f78ba8 --- /dev/null +++ b/crates/metaltile-std/src/ffai/moe_bgemm_q2k_view.rs @@ -0,0 +1,194 @@ +//! Copyright 2026 0xClandestine, Ekryski, TheTom, Ambisphaeric +//! SPDX-License-Identifier: Apache-2.0 +//! ZERO-COPY MoE Q2_K grouped BGEMM — reads raw 84-byte Q2_K blocks straight +//! from a no-copy mmap VIEW buffer (the down-projection counterpart of +//! `moe_bgemm_iq2xxs_view`), instead of the CPU-repacked qs/scales/d/dmin +//! pool. Eliminates the per-layer mmap→pool repack for the down experts. +//! +//! Identical MMA structure to `moe_bgemm_q2k_mpp`; ONLY the weight read +//! differs. Raw Q2_K block (84 bytes, see gguf_dequant_q2_k + the gather +//! repack): scales[0..16] u8, qs[16..80] (64 bytes), d fp16[80..82], dmin +//! fp16[82..84]. A block's bytes start at +//! `tensor_byte_off + expert*expert_byte_stride + block*84`. +//! `indices[row]` is the EXPERT id. d/dmin are decoded inline from their raw +//! 2 bytes (IEEE binary16), so no buffer aliasing and no CPU extract. + +use metaltile::kernel; + +#[kernel] +#[allow(clippy::too_many_arguments)] +pub fn ffai_moe_bgemm_q2k_view( + x: Tensor, + view_u8: Tensor, + indices: Tensor, + mut out: Tensor, + #[constexpr] m_total: u32, + #[constexpr] n_out: u32, + #[constexpr] k_in: u32, + #[constexpr] tensor_byte_off: u32, + #[constexpr] expert_byte_stride: u32, +) { + let n_tile_base = tgid_x * 32u32; + let m_tile_base = tgid_y * 16u32; + let lane = simd_lane; + threadgroup_alloc("xs", 256, coop_stage(T)); + threadgroup_alloc("ws", 512, coop_stage(T)); + threadgroup_alloc("out_scratch", 512, f32); + coop_tile_setup( + "gemm", + 16, + 32, + 16, + coop_stage(T), + "accumulate", + "simdgroup", + f32, + false, + true, + false, + ); + let mut sub_offset = 0u32; + for _sub_iter in range(0u32, 16u32, 1u32) { + let cur_row = m_tile_base + sub_offset; + let cur_in_range = (sub_offset < 16u32) & (cur_row < m_total); + let cur_expert = select(cur_in_range, load(indices[cur_row]), 4294967295u32); + let mut sub_end = 16u32; + let mut found = 0u32; + for _ii in range(0u32, 16u32, 1u32) { + let probe = sub_offset + 1u32 + _ii; + let probe_row = m_tile_base + probe; + let probe_in_range = (probe < 16u32) & (probe_row < m_total); + if probe_in_range & (found == 0u32) { + let e = load(indices[probe_row]); + if e != cur_expert { + sub_end = probe; + found = 1u32; + } + } + if (probe < 16u32) & (probe_row >= m_total) & (found == 0u32) { + sub_end = probe; + found = 1u32; + } + } + let cur_valid = (cur_expert != 4294967295u32) & (sub_offset < 16u32); + if cur_valid { + let expert_byte_base = tensor_byte_off + cur_expert * expert_byte_stride; + coop_tile_zero("gemm"); + for kb in range(0u32, k_in, 16u32) { + for _e in range(0u32, 8u32, 1u32) { + let flat = lane * 8u32 + _e; + let mr = flat / 16u32; + let kc = flat % 16u32; + let gr = m_tile_base + mr; + let in_run = (mr >= sub_offset) & (mr < sub_end) & (gr < m_total); + let safe_g = select(in_run, gr, 0u32); + let xv = load(x[safe_g * k_in + kb + kc]).cast::(); + threadgroup_store("xs", mr * 16u32 + kc, select(in_run, xv, 0.0f32)); + } + // Dequant Q2_K W[expert, n_tile_base+lane, kb..kb+16] → ws, + // reading raw 84-byte Q2_K blocks straight from the mmap view. + let w_row = lane; + let global_row = n_tile_base + w_row; + for _kc in range(0u32, 16u32, 1u32) { + let k = kb + _kc; + let vidx = global_row * k_in + k; + let block = vidx / 256u32; + let in_block = vidx & 255u32; + // Canonical Q2_K layout (see gguf_dequant_q2_k.rs). + let half = in_block / 128u32; + let yh = in_block - half * 128u32; + let jg = yh / 32u32; + let yg = yh - jg * 32u32; + let sub_half = yg / 16u32; + let l = yg - sub_half * 16u32; + let shift = jg * 2u32; + let q_byte = half * 32u32 + sub_half * 16u32 + l; + let sub = half * 8u32 + jg * 2u32 + sub_half; + let blk_byte = expert_byte_base + block * 84u32; + // scales[0..16], qs[16..80], d fp16[80..82], dmin fp16[82..84]. + let scale_byte = load(view_u8[blk_byte + sub]).cast::(); + let scale_4bit = scale_byte & 0xfu32; + let min_4bit = (scale_byte >> 4u32) & 0xfu32; + let qs_byte = load(view_u8[blk_byte + 16u32 + q_byte]).cast::(); + let q_2bit = (qs_byte >> shift) & 0x3u32; + // d, dmin decoded inline from raw fp16 (IEEE binary16). + let d_lo = load(view_u8[blk_byte + 80u32]).cast::(); + let d_hi = load(view_u8[blk_byte + 81u32]).cast::(); + let d_bits = d_lo | (d_hi << 8u32); + let d_sign = select((d_bits & 0x8000u32) != 0u32, 0.0f32 - 1.0f32, 1.0f32); + let d_exp = (d_bits >> 10u32) & 0x1fu32; + let d_mant = (d_bits & 0x3ffu32).cast::().cast::(); + let d_norm = (1.0f32 + d_mant / 1024.0f32) + * exp2(d_exp.cast::().cast::() - 15.0f32); + let d_sub = d_mant * exp2(0.0f32 - 24.0f32); + let d = d_sign * select(d_exp == 0u32, d_sub, d_norm); + let m_lo = load(view_u8[blk_byte + 82u32]).cast::(); + let m_hi = load(view_u8[blk_byte + 83u32]).cast::(); + let m_bits = m_lo | (m_hi << 8u32); + let m_sign = select((m_bits & 0x8000u32) != 0u32, 0.0f32 - 1.0f32, 1.0f32); + let m_exp = (m_bits >> 10u32) & 0x1fu32; + let m_mant = (m_bits & 0x3ffu32).cast::().cast::(); + let m_norm = (1.0f32 + m_mant / 1024.0f32) + * exp2(m_exp.cast::().cast::() - 15.0f32); + let m_sub = m_mant * exp2(0.0f32 - 24.0f32); + let dmin = m_sign * select(m_exp == 0u32, m_sub, m_norm); + let wq = d + * scale_4bit.cast::().cast::() + * q_2bit.cast::().cast::() + - dmin * min_4bit.cast::().cast::(); + threadgroup_store("ws", w_row * 16u32 + _kc, wq.cast::().cast::()); + } + threadgroup_barrier(); + coop_tile_load_a("gemm", "xs", true, coop_stage(T), 16, 16); + coop_tile_load_b("gemm", "ws", true, coop_stage(T), 16, 32); + coop_tile_run("gemm"); + threadgroup_barrier(); + } + coop_tile_store_c("gemm", "out_scratch", true, f32, 32, 16); + threadgroup_barrier(); + for _e in range(0u32, 16u32, 1u32) { + let flat = lane * 16u32 + _e; + let mr = flat / 32u32; + let nc = flat % 32u32; + let gr = m_tile_base + mr; + let gc = n_tile_base + nc; + let in_run = (mr >= sub_offset) & (mr < sub_end) & (gr < m_total) & (gc < n_out); + if in_run { + let v = threadgroup_load("out_scratch", mr * 32u32 + nc); + store(out[gr * n_out + gc], v.cast::()); + } + } + threadgroup_barrier(); + } + sub_offset = sub_end; + } +} + +pub mod kernel_benches { + use metaltile::{bench, test::*}; + + use super::ffai_moe_bgemm_q2k_view; + + #[bench(name = "ffai/moe/bgemm_q2k_view", dtypes = [f32, f16, bf16])] + fn bench_bgemm_q2k_view(dt: DType) -> BenchSetup { + let n_experts = 4usize; + let k_in = 2048usize; + let n_out = 4096usize; + let t_rows = 256usize; + let nblk = n_out * k_in / 256; + let view_bytes = n_experts * nblk * 84; + BenchSetup::new(ffai_moe_bgemm_q2k_view::kernel_ir_for(dt)) + .mode(KernelMode::Reduction) + .buffer(BenchBuffer::random("x", t_rows * k_in, dt)) + .buffer(BenchBuffer::random("view_u8", view_bytes, DType::U8)) + .buffer(BenchBuffer::zeros("indices", t_rows, DType::U32)) + .buffer(BenchBuffer::zeros("out", t_rows * n_out, dt).output()) + .constexpr("m_total", t_rows as u32) + .constexpr("n_out", n_out as u32) + .constexpr("k_in", k_in as u32) + .constexpr("tensor_byte_off", 0u32) + .constexpr("expert_byte_stride", (nblk * 84) as u32) + .grid_3d(n_out as u32 / 32, (t_rows as u32).div_ceil(16), 1, [32, 1, 1]) + .bytes_moved((view_bytes + t_rows * k_in * dt.size_bytes()) as u64) + } +} diff --git a/crates/metaltile-std/src/ffai/moe_bgemm_q2k_view_u16_bm64.rs b/crates/metaltile-std/src/ffai/moe_bgemm_q2k_view_u16_bm64.rs new file mode 100644 index 00000000..a6c41824 --- /dev/null +++ b/crates/metaltile-std/src/ffai/moe_bgemm_q2k_view_u16_bm64.rs @@ -0,0 +1,195 @@ +//! Copyright 2026 0xClandestine, Ekryski, TheTom, Ambisphaeric +//! SPDX-License-Identifier: Apache-2.0 +//! POOL-FREE amortized MoE Q2_K bgemm — bm64 64×64×32 coop_tile MMA reading RAW +//! resident GGUF Q2_K blocks directly (no CPU deinterleave pool). The down- +//! projection twin of `moe_bgemm_iq2xxs_view_u16_bm64` (which won for gate/up); +//! the existing `moe_bgemm_q2k_view` is the OLD mpp tiling (slower than bm64). +//! +//! Q2_K block = 84 bytes = 42 u16: scales[0..16] u8 (4-bit d-scale low nibble + +//! 4-bit dmin-scale high nibble per 16-value sub-block), qs[16..80] (64 bytes, +//! 2-bit quants), d f16 @80, dmin f16 @82. `indices[row]` = slot/expert id; a +//! block's bytes start at `tensor_byte_off + expert*expert_byte_stride + +//! block*84`. d/dmin read via the f16 view (native half — exact), scales/qs via +//! the u16 view. x [m_total,k_in], out [m_total,n_out]. Live-compiled (bgemm). + +use metaltile::kernel; + +#[kernel] +#[allow(clippy::too_many_arguments)] +pub fn ffai_moe_bgemm_q2k_view_u16_bm64( + x: Tensor, + view_u16: Tensor, + view_f16: Tensor, + indices: Tensor, + mut out: Tensor, + #[constexpr] m_total: u32, + #[constexpr] n_out: u32, + #[constexpr] k_in: u32, + #[constexpr] tensor_byte_off: u32, + #[constexpr] expert_byte_stride: u32, +) { + let n_tile_base = tgid_x * 64u32; + let m_tile_base = tgid_y * 64u32; + let sg = simd_group_id(); + let lane_in_tg = sg * 32u32 + simd_lane; + let sg_m_base = (sg / 2u32) * 32u32; + let sg_n_base = (sg & 1u32) * 32u32; + let x_m_row = lane_in_tg / 2u32; + let x_k_base = (lane_in_tg & 1u32) * 16u32; + threadgroup_alloc("Xs", 2048, coop_stage(T)); + threadgroup_alloc("Ws", 2048, coop_stage(T)); + threadgroup_alloc("OutScratch", 4096, f32); + coop_tile_setup( + "gemm", + 32, + 32, + 32, + coop_stage(T), + "accumulate", + "simdgroup", + f32, + false, + true, + false, + ); + let mut sub_offset = 0u32; + for _sub_iter in range(0u32, 64u32, 1u32) { + let cur_row = m_tile_base + sub_offset; + let cur_in_range = (sub_offset < 64u32) & (cur_row < m_total); + let cur_expert = select(cur_in_range, load(indices[cur_row]), 4294967295u32); + let mut sub_end = 64u32; + let mut found = 0u32; + for _ii in range(0u32, 64u32, 1u32) { + let probe = sub_offset + 1u32 + _ii; + let probe_row = m_tile_base + probe; + let probe_in_range = (probe < 64u32) & (probe_row < m_total); + if probe_in_range & (found == 0u32) { + let e = load(indices[probe_row]); + if e != cur_expert { + sub_end = probe; + found = 1u32; + } + } + if (probe < 64u32) & (probe_row >= m_total) & (found == 0u32) { + sub_end = probe; + found = 1u32; + } + } + let cur_valid = (cur_expert != 4294967295u32) & (sub_offset < 64u32); + if cur_valid { + // u16 base of this expert's raw Q2_K blocks (84 bytes = 42 u16 each). + let expert_u16_base = (tensor_byte_off + cur_expert * expert_byte_stride) / 2u32; + coop_tile_zero("gemm"); + for kb in range(0u32, k_in, 32u32) { + let gr_x = m_tile_base + x_m_row; + let in_run_x = (x_m_row >= sub_offset) & (x_m_row < sub_end) & (gr_x < m_total); + let safe_gr_x = select(in_run_x, gr_x, 0u32); + let x_dev_base = safe_gr_x * k_in + kb + x_k_base; + let x_ws_base = x_m_row * 32u32 + x_k_base; + for _i in range(0u32, 16u32, 1u32) { + let xv = load(x[x_dev_base + _i]).cast::(); + threadgroup_store("Xs", x_ws_base + _i, select(in_run_x, xv, 0.0f32)); + } + // Dequant Q2_K W → Ws. 16 elems/lane share one w_row + 256-block + // → hoist d/dmin (read once via f16 view at block bytes 80/82 = + // u16 40/41). scales/qs vary per element. + let flat0 = lane_in_tg * 16u32; + let w_row = flat0 / 32u32; + let k_local0 = flat0 & 31u32; + let global_row = n_tile_base + w_row; + let block = (global_row * k_in + kb + k_local0) / 256u32; + let blk_u16 = expert_u16_base + block * 42u32; + let d = load(view_f16[blk_u16 + 40u32]).cast::(); + let dmin = load(view_f16[blk_u16 + 41u32]).cast::(); + let ws_row_base = w_row * 32u32; + for _i in range(0u32, 16u32, 1u32) { + let k_local = k_local0 + _i; + let vidx = global_row * k_in + kb + k_local; + let in_block = vidx & 255u32; + let half = in_block / 128u32; + let yh = in_block - half * 128u32; + let jg = yh / 32u32; + let yg = yh - jg * 32u32; + let sub_half = yg / 16u32; + let l = yg - sub_half * 16u32; + let shift = jg * 2u32; + let q_byte = half * 32u32 + sub_half * 16u32 + l; + let sub = half * 8u32 + jg * 2u32 + sub_half; + // qs word (u32) covering this byte: block bytes 16+word_idx*4, + // i.e. u16 offset blk_u16 + 8 + word_idx*2 (qs starts at byte 16). + let word_idx = q_byte / 4u32; + let byte_in_word = q_byte & 3u32; + let qw = blk_u16 + 8u32 + word_idx * 2u32; + let word = load(view_u16[qw]).cast::() + | (load(view_u16[qw + 1u32]).cast::() << 16u32); + let qs_byte = (word >> (byte_in_word * 8u32)) & 0xffu32; + let q_2bit = (qs_byte >> shift) & 0x3u32; + // scale byte `sub` is at block byte `sub` = u16 blk_u16 + sub/2. + let sc_word = load(view_u16[blk_u16 + sub / 2u32]).cast::(); + let scale_byte = (sc_word >> ((sub & 1u32) * 8u32)) & 0xffu32; + let scale_4bit = scale_byte & 0xfu32; + let min_4bit = (scale_byte >> 4u32) & 0xfu32; + let wq = d + * scale_4bit.cast::().cast::() + * q_2bit.cast::().cast::() + - dmin * min_4bit.cast::().cast::(); + threadgroup_store("Ws", ws_row_base + k_local, wq.cast::().cast::()); + } + threadgroup_barrier(); + coop_tile_load_a("gemm", "Xs", true, coop_stage(T), 32, 32, sg_m_base * 32u32); + coop_tile_load_b("gemm", "Ws", true, coop_stage(T), 32, 32, sg_n_base * 32u32); + coop_tile_run("gemm"); + threadgroup_barrier(); + } + coop_tile_store_c("gemm", "OutScratch", true, f32, 32, 32, sg * 1024u32); + threadgroup_barrier(); + for _e in range(0u32, 32u32, 1u32) { + let flat = lane_in_tg * 32u32 + _e; + let mr = flat / 64u32; + let nc = flat & 63u32; + let gr = m_tile_base + mr; + let gc = n_tile_base + nc; + let in_run = (mr >= sub_offset) & (mr < sub_end) & (gr < m_total) & (gc < n_out); + if in_run { + let src_sg = (mr / 32u32) * 2u32 + nc / 32u32; + let v = threadgroup_load( + "OutScratch", + src_sg * 1024u32 + (mr & 31u32) * 32u32 + (nc & 31u32), + ); + store(out[gr * n_out + gc], v.cast::()); + } + } + threadgroup_barrier(); + } + sub_offset = sub_end; + } +} + +pub mod kernel_benches { + use metaltile::{bench, test::*}; + + use super::ffai_moe_bgemm_q2k_view_u16_bm64; + + #[bench(name = "ffai/moe/bgemm_q2k_view_u16_bm64", dtypes = [f32, f16, bf16])] + fn bench(dt: DType) -> BenchSetup { + let n_experts = 4usize; + let k_in = 2048usize; + let n_out = 4096usize; + let t_rows = 256usize; + let nblk = n_out * k_in / 256; + BenchSetup::new(ffai_moe_bgemm_q2k_view_u16_bm64::kernel_ir_for(dt)) + .mode(KernelMode::Reduction) + .buffer(BenchBuffer::random("x", t_rows * k_in, dt)) + .buffer(BenchBuffer::random("view_u16", n_experts * nblk * 42, DType::U16)) + .buffer(BenchBuffer::random("view_f16", n_experts * nblk * 42, DType::F16)) + .buffer(BenchBuffer::zeros("indices", t_rows, DType::U32)) + .buffer(BenchBuffer::zeros("out", t_rows * n_out, dt).output()) + .constexpr("m_total", t_rows as u32) + .constexpr("n_out", n_out as u32) + .constexpr("k_in", k_in as u32) + .constexpr("tensor_byte_off", 0u32) + .constexpr("expert_byte_stride", (nblk * 84) as u32) + .grid_3d(n_out as u32 / 64, (t_rows as u32).div_ceil(64), 1, [128, 1, 1]) + .bytes_moved((n_experts * nblk * 84 + t_rows * k_in * dt.size_bytes()) as u64) + } +} diff --git a/crates/metaltile-std/src/ffai/moe_down_weighted_sum_f16.rs b/crates/metaltile-std/src/ffai/moe_down_weighted_sum_f16.rs new file mode 100644 index 00000000..17bf7dbe --- /dev/null +++ b/crates/metaltile-std/src/ffai/moe_down_weighted_sum_f16.rs @@ -0,0 +1,359 @@ +//! Copyright 2026 0xClandestine, Ekryski, TheTom, Ambisphaeric +//! SPDX-License-Identifier: Apache-2.0 +//! Fused multi-expert (top-K=6) down gemv + weighted accumulate. +//! +//! For routed experts with PRE-DEQUANTED f16 down weights and +//! PRE-COMPUTED inner activations (post-swiglu), this kernel does the +//! full 6-way weighted sum into moeAccum in ONE dispatch: +//! +//! ```text +//! moeAccum[r] += Σ_{k=0..5} w[k] * (down_w_k[r] · inner_k) +//! ``` +//! +//! Replaces 18 dispatches per layer (6 gemv + 6 mul + 6 add) with 1. +//! Saves ~17 × 100 µs encoding overhead per layer × 43 layers ≈ +//! 72 ms per token — the bulk of the dispatch-count overhead. +//! +//! Geometry: one threadgroup per output row of moeAccum. Each TG +//! sequentially walks the 6 routed slots, stages that slot's inner +//! activation in threadgroup memory, accumulates that slot's dot +//! product in a per-thread f32 register, then manually reduces the +//! six dot products via simd/threadgroup scratch. +//! This mirrors the barrier discipline used by +//! `moe_down_swiglu_accum_int4_chain8` and avoids chaining multiple +//! reduction helpers in one kernel body. + +use metaltile::kernel; + +#[kernel] +#[allow(clippy::too_many_arguments)] +pub fn ffai_moe_down_weighted_sum_6( + down_0: Tensor, + inner_0: Tensor, + down_1: Tensor, + inner_1: Tensor, + down_2: Tensor, + inner_2: Tensor, + down_3: Tensor, + inner_3: Tensor, + down_4: Tensor, + inner_4: Tensor, + down_5: Tensor, + inner_5: Tensor, + weights: Tensor, + mut accum: Tensor, + #[constexpr] k: u32, +) { + // DSv4 Flash uses moeIntermediate=2048. Keep the alloc literal fixed + // so the generated MSL has static TG memory; Swift preconditions the + // production wrapper to k <= 2048. + threadgroup_alloc("tg_inner", 2048, "f32"); + threadgroup_alloc("sgs", 8, "f32"); + threadgroup_alloc("dots", 6, "f32"); + + let row = program_id::<0>(); + let row_base = row * k; + let w0 = load(weights[0u32]); + let w1 = load(weights[1u32]); + let w2 = load(weights[2u32]); + let w3 = load(weights[3u32]); + let w4 = load(weights[4u32]); + let w5 = load(weights[5u32]); + let iters = (k + lsize - 1u32) / lsize; + let mut acc0 = 0.0f32; + let mut acc1 = 0.0f32; + let mut acc2 = 0.0f32; + let mut acc3 = 0.0f32; + let mut acc4 = 0.0f32; + let mut acc5 = 0.0f32; + + for s_iter in range(0u32, iters, 1u32) { + let d = s_iter * lsize + tid; + if d < k { + threadgroup_store("tg_inner", d, load(inner_0[d]).cast::()); + } + } + threadgroup_barrier(); + for s_iter in range(0u32, iters, 1u32) { + let d = s_iter * lsize + tid; + if d < k { + acc0 = + acc0 + load(down_0[row_base + d]).cast::() * threadgroup_load("tg_inner", d); + } + } + threadgroup_barrier(); + + for s_iter in range(0u32, iters, 1u32) { + let d = s_iter * lsize + tid; + if d < k { + threadgroup_store("tg_inner", d, load(inner_1[d]).cast::()); + } + } + threadgroup_barrier(); + for s_iter in range(0u32, iters, 1u32) { + let d = s_iter * lsize + tid; + if d < k { + acc1 = + acc1 + load(down_1[row_base + d]).cast::() * threadgroup_load("tg_inner", d); + } + } + threadgroup_barrier(); + + for s_iter in range(0u32, iters, 1u32) { + let d = s_iter * lsize + tid; + if d < k { + threadgroup_store("tg_inner", d, load(inner_2[d]).cast::()); + } + } + threadgroup_barrier(); + for s_iter in range(0u32, iters, 1u32) { + let d = s_iter * lsize + tid; + if d < k { + acc2 = + acc2 + load(down_2[row_base + d]).cast::() * threadgroup_load("tg_inner", d); + } + } + threadgroup_barrier(); + + for s_iter in range(0u32, iters, 1u32) { + let d = s_iter * lsize + tid; + if d < k { + threadgroup_store("tg_inner", d, load(inner_3[d]).cast::()); + } + } + threadgroup_barrier(); + for s_iter in range(0u32, iters, 1u32) { + let d = s_iter * lsize + tid; + if d < k { + acc3 = + acc3 + load(down_3[row_base + d]).cast::() * threadgroup_load("tg_inner", d); + } + } + threadgroup_barrier(); + + for s_iter in range(0u32, iters, 1u32) { + let d = s_iter * lsize + tid; + if d < k { + threadgroup_store("tg_inner", d, load(inner_4[d]).cast::()); + } + } + threadgroup_barrier(); + for s_iter in range(0u32, iters, 1u32) { + let d = s_iter * lsize + tid; + if d < k { + acc4 = + acc4 + load(down_4[row_base + d]).cast::() * threadgroup_load("tg_inner", d); + } + } + threadgroup_barrier(); + + for s_iter in range(0u32, iters, 1u32) { + let d = s_iter * lsize + tid; + if d < k { + threadgroup_store("tg_inner", d, load(inner_5[d]).cast::()); + } + } + threadgroup_barrier(); + for s_iter in range(0u32, iters, 1u32) { + let d = s_iter * lsize + tid; + if d < k { + acc5 = + acc5 + load(down_5[row_base + d]).cast::() * threadgroup_load("tg_inner", d); + } + } + + let sg0 = simd_sum(acc0); + if simd_lane == 0u32 { + threadgroup_store("sgs", simd_id, sg0); + } + threadgroup_barrier(); + if simd_id == 0u32 { + let v = select(simd_lane < n_simd, threadgroup_load("sgs", simd_lane), 0.0f32); + let total = simd_sum(v); + if simd_lane == 0u32 { + threadgroup_store("dots", 0u32, total); + } + } + threadgroup_barrier(); + + let sg1 = simd_sum(acc1); + if simd_lane == 0u32 { + threadgroup_store("sgs", simd_id, sg1); + } + threadgroup_barrier(); + if simd_id == 0u32 { + let v = select(simd_lane < n_simd, threadgroup_load("sgs", simd_lane), 0.0f32); + let total = simd_sum(v); + if simd_lane == 0u32 { + threadgroup_store("dots", 1u32, total); + } + } + threadgroup_barrier(); + + let sg2 = simd_sum(acc2); + if simd_lane == 0u32 { + threadgroup_store("sgs", simd_id, sg2); + } + threadgroup_barrier(); + if simd_id == 0u32 { + let v = select(simd_lane < n_simd, threadgroup_load("sgs", simd_lane), 0.0f32); + let total = simd_sum(v); + if simd_lane == 0u32 { + threadgroup_store("dots", 2u32, total); + } + } + threadgroup_barrier(); + + let sg3 = simd_sum(acc3); + if simd_lane == 0u32 { + threadgroup_store("sgs", simd_id, sg3); + } + threadgroup_barrier(); + if simd_id == 0u32 { + let v = select(simd_lane < n_simd, threadgroup_load("sgs", simd_lane), 0.0f32); + let total = simd_sum(v); + if simd_lane == 0u32 { + threadgroup_store("dots", 3u32, total); + } + } + threadgroup_barrier(); + + let sg4 = simd_sum(acc4); + if simd_lane == 0u32 { + threadgroup_store("sgs", simd_id, sg4); + } + threadgroup_barrier(); + if simd_id == 0u32 { + let v = select(simd_lane < n_simd, threadgroup_load("sgs", simd_lane), 0.0f32); + let total = simd_sum(v); + if simd_lane == 0u32 { + threadgroup_store("dots", 4u32, total); + } + } + threadgroup_barrier(); + + let sg5 = simd_sum(acc5); + if simd_lane == 0u32 { + threadgroup_store("sgs", simd_id, sg5); + } + threadgroup_barrier(); + if simd_id == 0u32 { + let v = select(simd_lane < n_simd, threadgroup_load("sgs", simd_lane), 0.0f32); + let total = simd_sum(v); + if simd_lane == 0u32 { + threadgroup_store("dots", 5u32, total); + } + } + threadgroup_barrier(); + + let prev = load(accum[row]).cast::(); + let d0 = threadgroup_load("dots", 0u32).cast::().cast::(); + let c0 = (w0 * d0).cast::().cast::(); + let a0 = (prev + c0).cast::().cast::(); + let d1 = threadgroup_load("dots", 1u32).cast::().cast::(); + let c1 = (w1 * d1).cast::().cast::(); + let a1 = (a0 + c1).cast::().cast::(); + let d2 = threadgroup_load("dots", 2u32).cast::().cast::(); + let c2 = (w2 * d2).cast::().cast::(); + let a2 = (a1 + c2).cast::().cast::(); + let d3 = threadgroup_load("dots", 3u32).cast::().cast::(); + let c3 = (w3 * d3).cast::().cast::(); + let a3 = (a2 + c3).cast::().cast::(); + let d4 = threadgroup_load("dots", 4u32).cast::().cast::(); + let c4 = (w4 * d4).cast::().cast::(); + let a4 = (a3 + c4).cast::().cast::(); + let d5 = threadgroup_load("dots", 5u32).cast::().cast::(); + let c5 = (w5 * d5).cast::().cast::(); + let a5 = (a4 + c5).cast::().cast::(); + store(accum[row], a5); +} + +pub mod kernel_tests { + use metaltile::{test::*, test_kernel}; + + use super::ffai_moe_down_weighted_sum_6; + use crate::utils::{pack_f32, unpack_f32}; + + fn setup(m: usize, k: usize, dt: DType) -> TestSetup { + let weights = [0.20f32, 0.18, 0.17, 0.16, 0.15, 0.14]; + let mut dws: Vec> = Vec::new(); + let mut inns: Vec> = Vec::new(); + for slot in 0..6 { + let dw: Vec = + (0..m * k).map(|i| ((i % 17) as f32 - 8.0) * 0.01 * (slot as f32 + 1.0)).collect(); + let inn: Vec = (0..k).map(|j| ((j % 13) as f32 - 6.0) * 0.02).collect(); + dws.push(dw); + inns.push(inn); + } + let accum_in: Vec = (0..m).map(|r| (r as f32 % 5.0) * 0.1).collect(); + let dws_dt: Vec> = dws.iter().map(|d| unpack_f32(&pack_f32(d, dt), dt)).collect(); + let inns_dt: Vec> = + inns.iter().map(|i| unpack_f32(&pack_f32(i, dt), dt)).collect(); + let accum_dt = unpack_f32(&pack_f32(&accum_in, dt), dt); + let expected: Vec = (0..m) + .map(|r| { + let mut s = accum_dt[r]; + for slot in 0..6 { + let dot: f32 = (0..k).map(|j| dws_dt[slot][r * k + j] * inns_dt[slot][j]).sum(); + s += weights[slot] * dot; + } + s + }) + .collect(); + TestSetup::new(ffai_moe_down_weighted_sum_6::kernel_ir_for(dt)) + .mode(KernelMode::Reduction) + .input(TestBuffer::from_vec("down_0", pack_f32(&dws[0], dt), dt)) + .input(TestBuffer::from_vec("inner_0", pack_f32(&inns[0], dt), dt)) + .input(TestBuffer::from_vec("down_1", pack_f32(&dws[1], dt), dt)) + .input(TestBuffer::from_vec("inner_1", pack_f32(&inns[1], dt), dt)) + .input(TestBuffer::from_vec("down_2", pack_f32(&dws[2], dt), dt)) + .input(TestBuffer::from_vec("inner_2", pack_f32(&inns[2], dt), dt)) + .input(TestBuffer::from_vec("down_3", pack_f32(&dws[3], dt), dt)) + .input(TestBuffer::from_vec("inner_3", pack_f32(&inns[3], dt), dt)) + .input(TestBuffer::from_vec("down_4", pack_f32(&dws[4], dt), dt)) + .input(TestBuffer::from_vec("inner_4", pack_f32(&inns[4], dt), dt)) + .input(TestBuffer::from_vec("down_5", pack_f32(&dws[5], dt), dt)) + .input(TestBuffer::from_vec("inner_5", pack_f32(&inns[5], dt), dt)) + .input(TestBuffer::from_vec("weights", pack_f32(&weights, DType::F32), DType::F32)) + .input(TestBuffer::from_vec("accum", pack_f32(&accum_in, dt), dt)) + .constexpr("k", k as u32) + .expect(TestBuffer::from_vec("accum", pack_f32(&expected, dt), dt)) + .grid_3d(m as u32, 1, 1, [256, 1, 1]) + } + + #[test_kernel(dtypes = [f32, f16, bf16], tol = [1e-3, 2e-1, 2.0])] + fn test_mds_small(dt: DType) -> TestSetup { setup(16, 256, dt) } +} + +pub mod kernel_benches { + use metaltile::{bench, test::*}; + + use super::ffai_moe_down_weighted_sum_6; + + #[bench(name = "ffai/moe_down_weighted_sum_6", dtypes = [f32, f16, bf16])] + fn bench_mds(dt: DType) -> BenchSetup { + let (m, k) = (4096usize, 2048usize); + BenchSetup::new(ffai_moe_down_weighted_sum_6::kernel_ir_for(dt)) + .mode(KernelMode::Reduction) + .buffer(BenchBuffer::random("down_0", m * k, dt)) + .buffer(BenchBuffer::random("inner_0", k, dt)) + .buffer(BenchBuffer::random("down_1", m * k, dt)) + .buffer(BenchBuffer::random("inner_1", k, dt)) + .buffer(BenchBuffer::random("down_2", m * k, dt)) + .buffer(BenchBuffer::random("inner_2", k, dt)) + .buffer(BenchBuffer::random("down_3", m * k, dt)) + .buffer(BenchBuffer::random("inner_3", k, dt)) + .buffer(BenchBuffer::random("down_4", m * k, dt)) + .buffer(BenchBuffer::random("inner_4", k, dt)) + .buffer(BenchBuffer::random("down_5", m * k, dt)) + .buffer(BenchBuffer::random("inner_5", k, dt)) + .buffer(BenchBuffer::random("weights", 6, DType::F32)) + .buffer(BenchBuffer::random("accum", m, dt).output()) + .constexpr("k", k as u32) + .grid_3d(m as u32, 1, 1, [256, 1, 1]) + .bytes_moved( + (6 * m * k * dt.size_bytes() + 6 * k * dt.size_bytes() + m * dt.size_bytes()) + as u64, + ) + } +} diff --git a/crates/metaltile-std/src/ffai/moe_gather_down_q2k.rs b/crates/metaltile-std/src/ffai/moe_gather_down_q2k.rs new file mode 100644 index 00000000..2f6c8d31 --- /dev/null +++ b/crates/metaltile-std/src/ffai/moe_gather_down_q2k.rs @@ -0,0 +1,134 @@ +//! Copyright 2026 0xClandestine, Ekryski, TheTom, Ambisphaeric +//! SPDX-License-Identifier: Apache-2.0 +//! Multi-expert Q2_K gather down-projection + weighted sum — inline +//! dequant of all `n_slots` routed experts' down weights, each applied +//! to that expert's own SwiGLU inner activation, combined by the router +//! weights into the routed MoE output. ONE dispatch replaces +//! `n_slots` × {dequant, gemv} + the `n_slots`-way weighted accumulate. +//! +//! out[m] = Σ_slot weights[slot] · ( Σ_k dequant(downW_slot[m, k]) · inner_slot[k] ) +//! +//! ## Inputs (Q2_K split format — produced by the FFAI loader) +//! +//! ```text +//! inners_all [n_slots * k_in] T — per-slot SwiGLU inner +//! qs_all [n_slots * nblk_per_expert * 16] u32 — 64 qs bytes / block as 16 LE u32 +//! scales_all [n_slots * nblk_per_expert * 16] u8 — 16 packed (4-bit scale|4-bit min) / block +//! d_all [n_slots * nblk_per_expert] f32 — per-block super-scale +//! dmin_all [n_slots * nblk_per_expert] f32 — per-block min super-scale +//! weights [n_slots] f32 — router combine weights +//! out [m_out] T — routed MoE output +//! k_in / m_out / n_slots (constexpr) +//! ``` +//! +//! `nblk_per_expert = m_out * (k_in / 256)`. Output row `m` of an expert +//! occupies blocks `[m * blocks_per_row .. +blocks_per_row)`. +//! +//! ## Dispatch (Reduction mode) +//! +//! grid (threadgroups) = `[m_out, 1, 1]`, threadgroup = `[32, 1, 1]`. +//! `tgid_x = m`, `tid = lane`. The 32-lane simdgroup strides the k axis +//! across all slots, folding the router weight into each partial; one +//! `simd_sum` then lane-0 store. Q2_K dequant math is identical to +//! `gguf_dequant_q2_k::ffai_gguf_dequant_q2_k`. + +use metaltile::kernel; + +#[kernel] +pub fn ffai_moe_gather_down_q2k( + inners_all: Tensor, + qs_all: Tensor, + scales_all: Tensor, + d_all: Tensor, + dmin_all: Tensor, + expert_ids: Tensor, + weights: Tensor, + mut out: Tensor, + #[constexpr] k_in: u32, + #[constexpr] m_out: u32, + #[constexpr] n_slots: u32, +) { + let m = tgid_x; + let lane = tid; + let blocks_per_row = k_in / 256u32; + let nblk_per_expert = m_out * blocks_per_row; + + let mut acc = 0.0f32; + for slot in range(0u32, n_slots, 1u32) { + let w_slot = load(weights[slot]); + // qs_all/scales_all/d_all/dmin_all hold ALL experts (resident); + // routed expert for this slot is expert_ids[slot]. + let expert = load(expert_ids[slot]); + let qs_row_base = (expert * nblk_per_expert + m * blocks_per_row) * 16u32; + let sc_row_base = (expert * nblk_per_expert + m * blocks_per_row) * 16u32; + let blk_row_base = expert * nblk_per_expert + m * blocks_per_row; + let inner_base = slot * k_in; + for k_iter in range(lane, k_in, 32u32) { + let b = k_iter / 256u32; // block within row + let in_block = k_iter & 255u32; // 0..255 + // Canonical Q2_K layout (see gguf_dequant_q2_k.rs): 256 values = + // 2 halves of 128; each half = 4 j-groups of 32; each j-group is + // two runs of 16 values over 16 CONSECUTIVE qs bytes at a shared + // shift (j*2). NOT the naive 4-consecutive-per-byte mapping. + let half = in_block / 128u32; + let yh = in_block - half * 128u32; + let jg = yh / 32u32; + let yg = yh - jg * 32u32; + let sub_half = yg / 16u32; + let l = yg - sub_half * 16u32; + let shift = jg * 2u32; + let q_byte = half * 32u32 + sub_half * 16u32 + l; + let sub = half * 8u32 + jg * 2u32 + sub_half; // scale index 0..15 + let word_idx = q_byte / 4u32; + let byte_in_word = q_byte & 3u32; + let word = load(qs_all[qs_row_base + b * 16u32 + word_idx]); + let qs_byte = (word >> (byte_in_word * 8u32)) & 0xffu32; + let q_2bit = (qs_byte >> shift) & 0x3u32; + let scale_byte = load(scales_all[sc_row_base + b * 16u32 + sub]).cast::(); + let scale_4bit = scale_byte & 0xfu32; + let min_4bit = (scale_byte >> 4u32) & 0xfu32; + let d = load(d_all[blk_row_base + b]); + let dmin = load(dmin_all[blk_row_base + b]); + let wq = + d * scale_4bit.cast::().cast::() * q_2bit.cast::().cast::() + - dmin * min_4bit.cast::().cast::(); + let wq_t = wq.cast::().cast::(); + let xv = load(inners_all[inner_base + k_iter]).cast::(); + acc = acc + w_slot * wq_t * xv; + } + } + let total = simd_sum(acc); + if lane == 0u32 { + store(out[m], total.cast::()); + } +} + +pub mod kernel_benches { + use metaltile::{bench, test::*}; + + use super::ffai_moe_gather_down_q2k; + + // n_slots=6; production down dims (m_out=4096, k_in=2048). + #[bench(name = "ffai/moe/gather_down_q2k", dtypes = [f32, f16, bf16])] + fn bench_gather_down_q2k(dt: DType) -> BenchSetup { + let n_slots = 6usize; + let m_out = 4096usize; + let k_in = 2048usize; + let nblk = m_out * (k_in / 256); + BenchSetup::new(ffai_moe_gather_down_q2k::kernel_ir_for(dt)) + .mode(KernelMode::Reduction) + .buffer(BenchBuffer::random("inners_all", n_slots * k_in, dt)) + .buffer(BenchBuffer::random("qs_all", n_slots * nblk * 16, DType::U32)) + .buffer(BenchBuffer::random("scales_all", n_slots * nblk * 16, DType::U8)) + .buffer(BenchBuffer::random("d_all", n_slots * nblk, DType::F32)) + .buffer(BenchBuffer::random("dmin_all", n_slots * nblk, DType::F32)) + .buffer(BenchBuffer::zeros("expert_ids", n_slots, DType::U32)) + .buffer(BenchBuffer::random("weights", n_slots, DType::F32)) + .buffer(BenchBuffer::zeros("out", m_out, dt).output()) + .constexpr("k_in", k_in as u32) + .constexpr("m_out", m_out as u32) + .constexpr("n_slots", n_slots as u32) + .grid_3d(m_out as u32, 1, 1, [32, 1, 1]) + .bytes_moved((n_slots * nblk * 84 + n_slots * k_in * dt.size_bytes()) as u64) + } +} diff --git a/crates/metaltile-std/src/ffai/moe_gather_gemv_iq2xxs.rs b/crates/metaltile-std/src/ffai/moe_gather_gemv_iq2xxs.rs new file mode 100644 index 00000000..77d6fcfd --- /dev/null +++ b/crates/metaltile-std/src/ffai/moe_gather_gemv_iq2xxs.rs @@ -0,0 +1,155 @@ +//! Copyright 2026 0xClandestine, Ekryski, TheTom, Ambisphaeric +//! SPDX-License-Identifier: Apache-2.0 +//! Multi-expert IQ2_XXS gather GEMV — inline dequant + matmul-vector for +//! all `n_slots` (=6) routed experts of one MoE role (gate / up) in ONE +//! dispatch. +//! +//! ## Why +//! +//! The DSv4 decode FFN ran one `dequant-to-f16` kernel + one `gemv` per +//! routed expert (6 experts × {gate,up} = 12 dispatches/layer × 43 = +//! ~516 cmd buffers just for gate+up), each dominated by Apple's +//! ~0.13 ms per-kernel dispatch floor rather than real compute. This +//! kernel folds (a) the IQ2_XXS dequant, (b) the gemv, and (c) the +//! 6-expert gather into a SINGLE dispatch per role: `out[slot, m]` for +//! `slot in 0..n_slots`, `m in 0..m_out`. No intermediate f16 weight +//! buffer is ever materialised — the quant bytes are read straight from +//! the (resident) split buffers. +//! +//! ## Inputs (split format — produced by the FFAI loader) +//! +//! ```text +//! x [k_in] T — shared activation +//! qs_all [n_slots * nblk_per_expert * 16] u32 — slot-major, the 64 +//! bytes of qs[32] per +//! block re-laid as 16 LE u32 +//! d_all [n_slots * nblk_per_expert] f32 — per-block super-scale +//! grid [2048] u8 — iq2xxs_grid, 256×8 octets +//! signs [128] u8 — ksigns_iq2xs sign masks +//! out [n_slots * m_out] T — result +//! k_in u32 (constexpr) — input dim (multiple of 256) +//! m_out u32 (constexpr) — output rows per expert +//! ``` +//! +//! `nblk_per_expert = m_out * (k_in / 256)`. Output row `m` of an expert +//! occupies blocks `[m * blocks_per_row .. +blocks_per_row)` where +//! `blocks_per_row = k_in / 256`. +//! +//! ## Dispatch (Reduction mode) +//! +//! grid (threadgroups) = `[m_out, n_slots, 1]`, threadgroup = `[32, 1, 1]`. +//! `tgid_x = m`, `tgid_y = slot`, `tid = lane` (0..31). The 32 lanes of +//! the single simdgroup stride over the row's groups-of-32, each lane +//! accumulating a partial dot product; `simd_sum` folds them and lane 0 +//! stores. The dequant math is identical to +//! `gguf_dequant_iq2_xxs::ffai_gguf_dequant_iq2_xxs` (the proven, +//! production reference) — see that file for the per-element derivation. + +use metaltile::kernel; + +#[kernel] +pub fn ffai_moe_gather_gemv_iq2xxs( + x: Tensor, + qs_all: Tensor, + d_all: Tensor, + expert_ids: Tensor, + grid: Tensor, + signs: Tensor, + mut out: Tensor, + #[constexpr] k_in: u32, + #[constexpr] m_out: u32, +) { + let m = tgid_x; + let slot = tgid_y; + let lane = tid; + + let blocks_per_row = k_in / 256u32; + let nblk_per_expert = m_out * blocks_per_row; + // `qs_all` / `d_all` hold ALL experts (resident, pre-split); the + // routed expert for this slot is `expert_ids[slot]`. + let expert = load(expert_ids[slot]); + let qs_row_base = (expert * nblk_per_expert + m * blocks_per_row) * 16u32; + let d_row_base = expert * nblk_per_expert + m * blocks_per_row; + + // Each block = 8 groups of 32 values. Lanes stride over groups. + let total_groups = blocks_per_row * 8u32; + let mut acc = 0.0f32; + for grp in range(lane, total_groups, 32u32) { + let b = grp / 8u32; // block within the row + let group = grp & 7u32; // group within the block (0..7) + let aux_idx = load(qs_all[qs_row_base + b * 16u32 + group * 2u32]); + let aux_sgn = load(qs_all[qs_row_base + b * 16u32 + group * 2u32 + 1u32]); + let scale_4bit = aux_sgn >> 28u32; + let db = + load(d_all[d_row_base + b]) * ((scale_4bit.cast::().cast::() + 0.5) * 0.25); + // x indices covered by this 32-value group. + let x_grp = b * 256u32 + group * 32u32; + for j in range(0u32, 4u32, 1u32) { + let grid_key = (aux_idx >> (j * 8u32)) & 0xffu32; + let grid_row_base = grid_key * 8u32; + let sign_idx = (aux_sgn >> (j * 7u32)) & 0x7fu32; + let sign_mask = load(signs[sign_idx]).cast::(); + for l in range(0u32, 8u32, 1u32) { + let octet = load(grid[grid_row_base + l]).cast::().cast::().cast::(); + let lane_bit = sign_mask & (1u32 << l); + let sign = select(lane_bit != 0u32, -1.0f32, 1.0f32); + // Round the dequanted weight through the activation dtype + // `T` before the multiply — the unfused path materialises + // an f16 weight buffer first, so matching that rounding + // keeps the fused path bit-identical (greedy-stable). + let w = (db * sign * octet).cast::().cast::(); + let xv = load(x[x_grp + j * 8u32 + l]).cast::(); + acc = acc + w * xv; + } + } + } + let total = simd_sum(acc); + if lane == 0u32 { + store(out[slot * m_out + m], total.cast::()); + } +} + +#[cfg(test)] +pub mod kernel_tests { + use metaltile::test::*; + + use super::ffai_moe_gather_gemv_iq2xxs; + + #[test] + fn codegen_gather_gemv_iq2xxs_smoke() { + for dt in [DType::F32, DType::F16, DType::BF16] { + let ir = ffai_moe_gather_gemv_iq2xxs::kernel_ir_for(dt); + assert!(!ir.body.ops.is_empty(), "no ops for {dt:?}"); + assert!(ir.params.iter().any(|p| p.name == "qs_all"), "missing qs_all"); + assert!(ir.params.iter().any(|p| p.name == "grid"), "missing grid"); + } + } +} + +pub mod kernel_benches { + use metaltile::{bench, test::*}; + + use super::ffai_moe_gather_gemv_iq2xxs; + + // n_slots=6 routed experts; production gate/up dims (m_out=2048, k_in=4096). + #[bench(name = "ffai/moe/gather_gemv_iq2xxs", dtypes = [f32, f16, bf16])] + fn bench_gather_gemv_iq2xxs(dt: DType) -> BenchSetup { + let n_slots = 6usize; + let m_out = 2048usize; + let k_in = 4096usize; + let nblk = m_out * (k_in / 256); + BenchSetup::new(ffai_moe_gather_gemv_iq2xxs::kernel_ir_for(dt)) + .mode(KernelMode::Reduction) + .buffer(BenchBuffer::random("x", k_in, dt)) + .buffer(BenchBuffer::random("qs_all", n_slots * nblk * 16, DType::U32)) + .buffer(BenchBuffer::random("d_all", n_slots * nblk, DType::F32)) + .buffer(BenchBuffer::zeros("expert_ids", n_slots, DType::U32)) + .buffer(BenchBuffer::random("grid", 2048, DType::U8)) + .buffer(BenchBuffer::random("signs", 128, DType::U8)) + .buffer(BenchBuffer::zeros("out", n_slots * m_out, dt).output()) + .constexpr("k_in", k_in as u32) + .constexpr("m_out", m_out as u32) + .grid_3d(m_out as u32, n_slots as u32, 1, [32, 1, 1]) + .bytes_moved((n_slots * nblk * 64 + n_slots * nblk * 4 + k_in * dt.size_bytes()) as u64) + } +} diff --git a/crates/metaltile-std/src/ffai/moe_gemv_rows_iq2xxs.rs b/crates/metaltile-std/src/ffai/moe_gemv_rows_iq2xxs.rs new file mode 100644 index 00000000..c84e35c8 --- /dev/null +++ b/crates/metaltile-std/src/ffai/moe_gemv_rows_iq2xxs.rs @@ -0,0 +1,109 @@ +//! Copyright 2026 0xClandestine, Ekryski, TheTom, Ambisphaeric +//! SPDX-License-Identifier: Apache-2.0 +//! Prefill MoE IQ2_XXS GEMV-over-rows — the fast decode gemv +//! (ffai_moe_gather_gemv_iq2xxs, ~270 GB/s) applied to a whole batch of +//! M = N*topK (token,expert) rows in ONE dispatch, instead of the +//! coop-tile MMA bgemm (ffai_moe_bgemm_iq2xxs_mpp, ~4-10 GB/s — 15-70x +//! slower). The bgemm's MMA staging + barriers dominate at these quant +//! shapes; the direct simd_sum dot-product the gemv uses is far faster. +//! +//! Each threadgroup (32 lanes) computes one output element +//! out[row, m] = dot(W[expert(row), m, :], x[row, :]), reading the +//! resident split qs/d pool. Rows are pre-permuted by expert (so x[row] +//! is the gathered activation for that (token,expert) pair) and +//! `expert_ids[row]` is the row's expert. NO amortization of the weight +//! dequant across rows of the same expert (each row re-dequants), but the +//! per-op throughput is so much higher than the MMA path that it still +//! wins; a weight-stationary version is the next step if this is the +//! bottleneck. Dequant math is identical to gather_gemv_iq2xxs. +//! +//! grid (threadgroups) = [m_out, m_total, 1], threadgroup = [32,1,1]. + +use metaltile::kernel; + +#[kernel] +pub fn ffai_moe_gemv_rows_iq2xxs( + x: Tensor, + qs_all: Tensor, + d_all: Tensor, + expert_ids: Tensor, + grid: Tensor, + signs: Tensor, + mut out: Tensor, + #[constexpr] k_in: u32, + #[constexpr] m_out: u32, + #[constexpr] m_total: u32, +) { + let m = tgid_x; // output row (0..m_out) + let row = tgid_y; // (token,expert) pair (0..m_total) + let lane = tid; + // grid y-extent is exactly m_total, so row is always in range. + + let blocks_per_row = k_in / 256u32; + let nblk_per_expert = m_out * blocks_per_row; + let expert = load(expert_ids[row]); + let qs_row_base = (expert * nblk_per_expert + m * blocks_per_row) * 16u32; + let d_row_base = expert * nblk_per_expert + m * blocks_per_row; + let x_base = row * k_in; + + let total_groups = blocks_per_row * 8u32; + let mut acc = 0.0f32; + for grp in range(lane, total_groups, 32u32) { + let b = grp / 8u32; + let group = grp & 7u32; + let aux_idx = load(qs_all[qs_row_base + b * 16u32 + group * 2u32]); + let aux_sgn = load(qs_all[qs_row_base + b * 16u32 + group * 2u32 + 1u32]); + let scale_4bit = aux_sgn >> 28u32; + let db = + load(d_all[d_row_base + b]) * ((scale_4bit.cast::().cast::() + 0.5) * 0.25); + let x_grp = b * 256u32 + group * 32u32; + for j in range(0u32, 4u32, 1u32) { + let grid_key = (aux_idx >> (j * 8u32)) & 0xffu32; + let grid_row_base = grid_key * 8u32; + let sign_idx = (aux_sgn >> (j * 7u32)) & 0x7fu32; + let sign_mask = load(signs[sign_idx]).cast::(); + for l in range(0u32, 8u32, 1u32) { + let octet = load(grid[grid_row_base + l]).cast::().cast::().cast::(); + let lane_bit = sign_mask & (1u32 << l); + let sign = select(lane_bit != 0u32, -1.0f32, 1.0f32); + let w = (db * sign * octet).cast::().cast::(); + let xv = load(x[x_base + x_grp + j * 8u32 + l]).cast::(); + acc = acc + w * xv; + } + } + } + let total = simd_sum(acc); + if lane == 0u32 { + store(out[row * m_out + m], total.cast::()); + } +} + +pub mod kernel_benches { + use metaltile::{bench, test::*}; + + use super::ffai_moe_gemv_rows_iq2xxs; + + // M=256 rows, production gate/up dims (m_out=2048, k_in=4096). + #[bench(name = "ffai/moe/gemv_rows_iq2xxs", dtypes = [f32, f16, bf16])] + fn bench_gemv_rows_iq2xxs(dt: DType) -> BenchSetup { + let m_total = 256usize; + let n_experts = 8usize; + let m_out = 2048usize; + let k_in = 4096usize; + let nblk = m_out * (k_in / 256); + BenchSetup::new(ffai_moe_gemv_rows_iq2xxs::kernel_ir_for(dt)) + .mode(KernelMode::Reduction) + .buffer(BenchBuffer::random("x", m_total * k_in, dt)) + .buffer(BenchBuffer::random("qs_all", n_experts * nblk * 16, DType::U32)) + .buffer(BenchBuffer::random("d_all", n_experts * nblk, DType::F32)) + .buffer(BenchBuffer::zeros("expert_ids", m_total, DType::U32)) + .buffer(BenchBuffer::random("grid", 2048, DType::U8)) + .buffer(BenchBuffer::random("signs", 128, DType::U8)) + .buffer(BenchBuffer::zeros("out", m_total * m_out, dt).output()) + .constexpr("k_in", k_in as u32) + .constexpr("m_out", m_out as u32) + .constexpr("m_total", m_total as u32) + .grid_3d(m_out as u32, m_total as u32, 1, [32, 1, 1]) + .bytes_moved((m_total * nblk * 64 + m_total * k_in * dt.size_bytes()) as u64) + } +} diff --git a/crates/metaltile-std/src/ffai/moe_gemv_rows_q2k.rs b/crates/metaltile-std/src/ffai/moe_gemv_rows_q2k.rs new file mode 100644 index 00000000..d5dc0211 --- /dev/null +++ b/crates/metaltile-std/src/ffai/moe_gemv_rows_q2k.rs @@ -0,0 +1,108 @@ +//! Copyright 2026 0xClandestine, Ekryski, TheTom, Ambisphaeric +//! SPDX-License-Identifier: Apache-2.0 +//! Prefill MoE Q2_K GEMV-over-rows (down projection) — the Q2_K twin of +//! ffai_moe_gemv_rows_iq2xxs. Replaces the slow coop-tile bgemm +//! (gather_bgemm_q2k_mpp ~10 GB/s) with the fast decode-style direct +//! simd_sum dot-product applied to a whole batch of M=(token,expert) rows +//! in ONE dispatch. Reads the resident split pool (qs u32 / scales u8 / +//! d f32 / dmin f32). Canonical Q2_K layout + dequant identical to +//! gather_bgemm_q2k_mpp; x is indexed per row, expert_ids[row] per row. +//! +//! grid (threadgroups) = [m_out, m_total, 1], threadgroup = [32,1,1]. +//! out[row, m] = dot(W_down[expert(row), m, :], x[row, :]). + +use metaltile::kernel; + +#[kernel] +#[allow(clippy::too_many_arguments)] +pub fn ffai_moe_gemv_rows_q2k( + x: Tensor, + qs: Tensor, + scales: Tensor, + d_f32: Tensor, + dmin_f32: Tensor, + expert_ids: Tensor, + mut out: Tensor, + #[constexpr] k_in: u32, + #[constexpr] m_out: u32, + #[constexpr] m_total: u32, +) { + let m = tgid_x; // output row (0..m_out) + let row = tgid_y; // (token,expert) pair (0..m_total) + let lane = tid; + + let blocks_per_row = k_in / 256u32; + let nblk_per_expert = m_out * blocks_per_row; + let expert = load(expert_ids[row]); + let qs_expert_base = (expert * nblk_per_expert) * 16u32; + let sc_expert_base = (expert * nblk_per_expert) * 16u32; + let blk_expert_base = expert * nblk_per_expert; + let x_base = row * k_in; + + let mut acc = 0.0f32; + // Each lane strides over the row's k values, dequant + multiply-accumulate. + for k in range(lane, k_in, 32u32) { + let vidx = m * k_in + k; + let block = vidx / 256u32; + let in_block = vidx & 255u32; + // Canonical Q2_K layout (see gguf_dequant_q2_k.rs / moe_bgemm_q2k_mpp). + let half = in_block / 128u32; + let yh = in_block - half * 128u32; + let jg = yh / 32u32; + let yg = yh - jg * 32u32; + let sub_half = yg / 16u32; + let l = yg - sub_half * 16u32; + let shift = jg * 2u32; + let q_byte = half * 32u32 + sub_half * 16u32 + l; + let sub = half * 8u32 + jg * 2u32 + sub_half; + let word_idx = q_byte / 4u32; + let byte_in_word = q_byte & 3u32; + let word = load(qs[qs_expert_base + block * 16u32 + word_idx]); + let qs_byte = (word >> (byte_in_word * 8u32)) & 0xffu32; + let q_2bit = (qs_byte >> shift) & 0x3u32; + let scale_byte = load(scales[sc_expert_base + block * 16u32 + sub]).cast::(); + let scale_4bit = scale_byte & 0xfu32; + let min_4bit = (scale_byte >> 4u32) & 0xfu32; + let d = load(d_f32[blk_expert_base + block]); + let dmin = load(dmin_f32[blk_expert_base + block]); + let wq = d * scale_4bit.cast::().cast::() * q_2bit.cast::().cast::() + - dmin * min_4bit.cast::().cast::(); + let w = wq.cast::().cast::(); + let xv = load(x[x_base + k]).cast::(); + acc = acc + w * xv; + } + let total = simd_sum(acc); + if lane == 0u32 { + store(out[row * m_out + m], total.cast::()); + } +} + +pub mod kernel_benches { + use metaltile::{bench, test::*}; + + use super::ffai_moe_gemv_rows_q2k; + + // M=256 rows, production down dims (m_out=4096 hidden, k_in=2048 intermediate). + #[bench(name = "ffai/moe/gemv_rows_q2k", dtypes = [f32, f16, bf16])] + fn bench_gemv_rows_q2k(dt: DType) -> BenchSetup { + let m_total = 256usize; + let n_experts = 8usize; + let m_out = 4096usize; + let k_in = 2048usize; + let nblk = m_out * (k_in / 256); + BenchSetup::new(ffai_moe_gemv_rows_q2k::kernel_ir_for(dt)) + .mode(KernelMode::Reduction) + .buffer(BenchBuffer::random("x", m_total * k_in, dt)) + .buffer(BenchBuffer::random("qs", n_experts * nblk * 16, DType::U32)) + .buffer(BenchBuffer::random("scales", n_experts * nblk * 16, DType::U8)) + .buffer(BenchBuffer::random("d_f32", n_experts * nblk, DType::F32)) + .buffer(BenchBuffer::random("dmin_f32", n_experts * nblk, DType::F32)) + .buffer(BenchBuffer::zeros("expert_ids", m_total, DType::U32)) + .buffer(BenchBuffer::zeros("out", m_total * m_out, dt).output()) + .constexpr("k_in", k_in as u32) + .constexpr("m_out", m_out as u32) + .constexpr("m_total", m_total as u32) + .grid_3d(m_out as u32, m_total as u32, 1, [32, 1, 1]) + .bytes_moved((m_total * nblk * 84 + m_total * k_in * dt.size_bytes()) as u64) + } +} diff --git a/crates/metaltile-std/src/ffai/moe_gemv_rows_view_iq2xxs.rs b/crates/metaltile-std/src/ffai/moe_gemv_rows_view_iq2xxs.rs new file mode 100644 index 00000000..afdc3979 --- /dev/null +++ b/crates/metaltile-std/src/ffai/moe_gemv_rows_view_iq2xxs.rs @@ -0,0 +1,227 @@ +//! Copyright 2026 0xClandestine, Ekryski, TheTom, Ambisphaeric +//! SPDX-License-Identifier: Apache-2.0 +//! SINGLE-COPY fast prefill MoE IQ2_XXS: the gemv-over-rows kernel +//! (moe_gemv_rows_iq2xxs, ~100 GB/s) reading raw 66-byte IQ2_XXS blocks +//! straight from the resident no-copy mmap VIEW (like moe_bgemm_iq2xxs_view) +//! instead of a repacked pool. This is the combination we want: +//! ONE resident copy of the weights (no 70GB repack pool) + the fast direct +//! simd_sum dot-product (no slow coop-tile MMA). Eliminates both the +//! double-memory and the per-chunk repack/re-read. +//! +//! out[row, m] = dot(W[expert(row), m, :], x[row, :]). Block bytes read +//! inline: d via fp16 decode (2 bytes), aux_idx/aux_sgn via u8-combine. +//! grid (threadgroups) = [m_out, m_total, 1], threadgroup = [32,1,1]. + +use metaltile::kernel; + +#[kernel] +#[allow(clippy::too_many_arguments)] +pub fn ffai_moe_gemv_rows_view_iq2xxs( + x: Tensor, + view_u8: Tensor, + grid: Tensor, + signs: Tensor, + expert_ids: Tensor, + mut out: Tensor, + #[constexpr] k_in: u32, + #[constexpr] m_out: u32, + #[constexpr] m_total: u32, + #[constexpr] tensor_byte_off: u32, + #[constexpr] expert_byte_stride: u32, +) { + let m = tgid_x; // output row (0..m_out) + let row = tgid_y; // (token,expert) pair (0..m_total) + let lane = tid; + + let blocks_per_row = k_in / 256u32; + let expert = load(expert_ids[row]); + let expert_byte_base = tensor_byte_off + expert * expert_byte_stride; + // Block index base for output row m within the expert (row-major [m_out,k_in]). + let row_block0 = m * blocks_per_row; + let x_base = row * k_in; + + let total_groups = blocks_per_row * 8u32; + let mut acc = 0.0f32; + for grp in range(lane, total_groups, 32u32) { + let b = grp / 8u32; // block within the row + let group = grp & 7u32; // group within the block (0..7) + let blk_byte = expert_byte_base + (row_block0 + b) * 66u32; + // d: leading fp16 (2 bytes), decoded inline (IEEE binary16). + let d_lo = load(view_u8[blk_byte]).cast::(); + let d_hi = load(view_u8[blk_byte + 1u32]).cast::(); + let d_bits = d_lo | (d_hi << 8u32); + let d_sign = select((d_bits & 0x8000u32) != 0u32, 0.0f32 - 1.0f32, 1.0f32); + let d_exp = (d_bits >> 10u32) & 0x1fu32; + let d_mant = (d_bits & 0x3ffu32).cast::().cast::(); + let d_norm = + (1.0f32 + d_mant / 1024.0f32) * exp2(d_exp.cast::().cast::() - 15.0f32); + let d_sub = d_mant * exp2(0.0f32 - 24.0f32); + let dval = d_sign * select(d_exp == 0u32, d_sub, d_norm); + // group g: aux_idx u32 + aux_sgn u32 at blk_byte+2+g*8, LE u8-combine. + let q0 = blk_byte + 2u32 + group * 8u32; + let aux_idx = load(view_u8[q0]).cast::() + | (load(view_u8[q0 + 1u32]).cast::() << 8u32) + | (load(view_u8[q0 + 2u32]).cast::() << 16u32) + | (load(view_u8[q0 + 3u32]).cast::() << 24u32); + let aux_sgn = load(view_u8[q0 + 4u32]).cast::() + | (load(view_u8[q0 + 5u32]).cast::() << 8u32) + | (load(view_u8[q0 + 6u32]).cast::() << 16u32) + | (load(view_u8[q0 + 7u32]).cast::() << 24u32); + let scale_4bit = aux_sgn >> 28u32; + let db = dval * ((scale_4bit.cast::().cast::() + 0.5) * 0.25); + let x_grp = b * 256u32 + group * 32u32; + for j in range(0u32, 4u32, 1u32) { + let grid_key = (aux_idx >> (j * 8u32)) & 0xffu32; + let grid_row_base = grid_key * 8u32; + let sign_idx = (aux_sgn >> (j * 7u32)) & 0x7fu32; + let sign_mask = load(signs[sign_idx]).cast::(); + for l in range(0u32, 8u32, 1u32) { + let octet = load(grid[grid_row_base + l]).cast::().cast::().cast::(); + let lane_bit = sign_mask & (1u32 << l); + let sign = select(lane_bit != 0u32, -1.0f32, 1.0f32); + let w = (db * sign * octet).cast::().cast::(); + let xv = load(x[x_base + x_grp + j * 8u32 + l]).cast::(); + acc = acc + w * xv; + } + } + } + let total = simd_sum(acc); + if lane == 0u32 { + store(out[row * m_out + m], total.cast::()); + } +} + +/// U16 variant — reads the raw IQ2_XXS blocks via ALIGNED u16 loads (the +/// block stride 66 and the qs offset 2 are both even, so every d/aux read is +/// u16-aligned) instead of the u8-recombine. 5 u16 loads/group vs 10 u8 → +/// tests whether the new DType::U16 lets the zero-copy view approach the pool +/// gemv's ~100 GB/s (the u8 path was ~3.5). If so, pool-elimination is viable. +#[kernel] +#[allow(clippy::too_many_arguments)] +pub fn ffai_moe_gemv_rows_view_u16_iq2xxs( + x: Tensor, + view_u16: Tensor, + grid: Tensor, + signs: Tensor, + expert_ids: Tensor, + mut out: Tensor, + #[constexpr] k_in: u32, + #[constexpr] m_out: u32, + #[constexpr] m_total: u32, + #[constexpr] tensor_byte_off: u32, + #[constexpr] expert_byte_stride: u32, +) { + let m = tgid_x; + let row = tgid_y; + let lane = tid; + let blocks_per_row = k_in / 256u32; + let expert = load(expert_ids[row]); + let expert_byte_base = tensor_byte_off + expert * expert_byte_stride; + let row_block0 = m * blocks_per_row; + let x_base = row * k_in; + let total_groups = blocks_per_row * 8u32; + let mut acc = 0.0f32; + for grp in range(lane, total_groups, 32u32) { + let b = grp / 8u32; + let group = grp & 7u32; + // u16 index of the block start (blk_byte is even → /2 exact). + let blk_u16 = (expert_byte_base + (row_block0 + b) * 66u32) / 2u32; + let d_bits = load(view_u16[blk_u16]).cast::(); + let d_sign = select((d_bits & 0x8000u32) != 0u32, 0.0f32 - 1.0f32, 1.0f32); + let d_exp = (d_bits >> 10u32) & 0x1fu32; + let d_mant = (d_bits & 0x3ffu32).cast::().cast::(); + let d_norm = + (1.0f32 + d_mant / 1024.0f32) * exp2(d_exp.cast::().cast::() - 15.0f32); + let d_sub = d_mant * exp2(0.0f32 - 24.0f32); + let dval = d_sign * select(d_exp == 0u32, d_sub, d_norm); + // group g: aux_idx + aux_sgn (2 u32 = 4 u16) at u16 offset blk+1+g*4. + let q = blk_u16 + 1u32 + group * 4u32; + let aux_idx = + load(view_u16[q]).cast::() | (load(view_u16[q + 1u32]).cast::() << 16u32); + let aux_sgn = load(view_u16[q + 2u32]).cast::() + | (load(view_u16[q + 3u32]).cast::() << 16u32); + let scale_4bit = aux_sgn >> 28u32; + let db = dval * ((scale_4bit.cast::().cast::() + 0.5) * 0.25); + let x_grp = b * 256u32 + group * 32u32; + for j in range(0u32, 4u32, 1u32) { + let grid_key = (aux_idx >> (j * 8u32)) & 0xffu32; + let grid_row_base = grid_key * 8u32; + let sign_idx = (aux_sgn >> (j * 7u32)) & 0x7fu32; + let sign_mask = load(signs[sign_idx]).cast::(); + for l in range(0u32, 8u32, 1u32) { + let octet = load(grid[grid_row_base + l]).cast::().cast::().cast::(); + let lane_bit = sign_mask & (1u32 << l); + let sign = select(lane_bit != 0u32, -1.0f32, 1.0f32); + let w = (db * sign * octet).cast::().cast::(); + let xv = load(x[x_base + x_grp + j * 8u32 + l]).cast::(); + acc = acc + w * xv; + } + } + } + let total = simd_sum(acc); + if lane == 0u32 { + store(out[row * m_out + m], total.cast::()); + } +} + +pub mod kernel_benches { + use metaltile::{bench, test::*}; + + use super::ffai_moe_gemv_rows_view_u16_iq2xxs; + + #[bench(name = "ffai/moe/gemv_rows_view_u16_iq2xxs", dtypes = [f32, f16, bf16])] + fn bench_gemv_rows_view_u16_iq2xxs(dt: DType) -> BenchSetup { + let m_total = 256usize; + let n_experts = 8usize; + let m_out = 2048usize; + let k_in = 4096usize; + let nblk = m_out * (k_in / 256); + let view_bytes = n_experts * nblk * 66; + BenchSetup::new(ffai_moe_gemv_rows_view_u16_iq2xxs::kernel_ir_for(dt)) + .mode(KernelMode::Reduction) + .buffer(BenchBuffer::random("x", m_total * k_in, dt)) + .buffer(BenchBuffer::random("view_u16", view_bytes / 2, DType::U16)) + .buffer(BenchBuffer::random("grid", 2048, DType::U8)) + .buffer(BenchBuffer::random("signs", 128, DType::U8)) + .buffer(BenchBuffer::zeros("expert_ids", m_total, DType::U32)) + .buffer(BenchBuffer::zeros("out", m_total * m_out, dt).output()) + .constexpr("k_in", k_in as u32) + .constexpr("m_out", m_out as u32) + .constexpr("m_total", m_total as u32) + .constexpr("tensor_byte_off", 0u32) + .constexpr("expert_byte_stride", (nblk * 66) as u32) + .grid_3d(m_out as u32, m_total as u32, 1, [32, 1, 1]) + .bytes_moved((m_total * nblk * 64 + m_total * k_in * dt.size_bytes()) as u64) + } +} + +pub mod kernel_benches_old { + use metaltile::{bench, test::*}; + + use super::ffai_moe_gemv_rows_view_iq2xxs; + + #[bench(name = "ffai/moe/gemv_rows_view_iq2xxs", dtypes = [f32, f16, bf16])] + fn bench_gemv_rows_view_iq2xxs(dt: DType) -> BenchSetup { + let m_total = 256usize; + let n_experts = 8usize; + let m_out = 2048usize; + let k_in = 4096usize; + let nblk = m_out * (k_in / 256); + let view_bytes = n_experts * nblk * 66; + BenchSetup::new(ffai_moe_gemv_rows_view_iq2xxs::kernel_ir_for(dt)) + .mode(KernelMode::Reduction) + .buffer(BenchBuffer::random("x", m_total * k_in, dt)) + .buffer(BenchBuffer::random("view_u8", view_bytes, DType::U8)) + .buffer(BenchBuffer::random("grid", 2048, DType::U8)) + .buffer(BenchBuffer::random("signs", 128, DType::U8)) + .buffer(BenchBuffer::zeros("expert_ids", m_total, DType::U32)) + .buffer(BenchBuffer::zeros("out", m_total * m_out, dt).output()) + .constexpr("k_in", k_in as u32) + .constexpr("m_out", m_out as u32) + .constexpr("m_total", m_total as u32) + .constexpr("tensor_byte_off", 0u32) + .constexpr("expert_byte_stride", (nblk * 66) as u32) + .grid_3d(m_out as u32, m_total as u32, 1, [32, 1, 1]) + .bytes_moved((view_bytes + m_total * k_in * dt.size_bytes()) as u64) + } +} diff --git a/crates/metaltile-std/src/ffai/moe_gemv_ws_iq2xxs.rs b/crates/metaltile-std/src/ffai/moe_gemv_ws_iq2xxs.rs new file mode 100644 index 00000000..5f147fd8 --- /dev/null +++ b/crates/metaltile-std/src/ffai/moe_gemv_ws_iq2xxs.rs @@ -0,0 +1,153 @@ +//! Copyright 2026 0xClandestine, Ekryski, TheTom, Ambisphaeric +//! SPDX-License-Identifier: Apache-2.0 +//! Prefill MoE IQ2_XXS WEIGHT-STATIONARY gemv — the amortized fix. +//! +//! The decode `gather_gemv_iq2xxs` runs at ~270 GB/s and `gemv_rows` +//! at ~105 GB/s, but the coop_tile MMA `bgemm_*_bm64` only hits ~21 GB/s +//! (MMA staging + barriers dominate the cheap IQ2 dequant). gemv_rows is +//! fast PER-OP but re-dequants each expert's weight row for EVERY token +//! row (no amortization → loses at large M). This kernel keeps gemv's +//! direct simd_sum dot-product but DEQUANTS each expert's weight row +//! ONCE into threadgroup memory and reuses it across all R rows of that +//! expert in the tile — fast per-op AND amortized. +//! +//! Each threadgroup (32 lanes) owns one output column `m` and a tile of +//! `rows_per_tile` consecutive (token,expert) rows. Rows are pre-permuted +//! by expert (contiguous per expert), so a tile is usually single-expert; +//! the weight row W[expert,m,:] is dequanted into `Ws` once and reused. +//! On an expert boundary inside the tile (rare — ~1 per expert segment) +//! it re-dequants. The expert-change test is threadgroup-uniform +//! (`expert_ids[row]` is the same for all lanes), so the barrier inside +//! the loop is safe. Dequant math identical to gemv_rows / gather_gemv. +//! +//! grid (threadgroups) = [m_out, ceil(m_total/rows_per_tile), 1], +//! threadgroup = [32,1,1]. + +use metaltile::kernel; + +#[kernel] +#[allow(clippy::too_many_arguments)] +pub fn ffai_moe_gemv_ws_iq2xxs( + x: Tensor, + qs_all: Tensor, + d_all: Tensor, + expert_ids: Tensor, + grid: Tensor, + signs: Tensor, + mut out: Tensor, + #[constexpr] k_in: u32, + #[constexpr] m_out: u32, + #[constexpr] m_total: u32, + #[constexpr] rows_per_tile: u32, +) { + let m = tgid_x; + let tile = tgid_y; + let lane = tid; + let row0 = tile * rows_per_tile; + + let blocks_per_row = k_in / 256u32; + let nblk_per_expert = m_out * blocks_per_row; + let total_groups = blocks_per_row * 8u32; + + // Ws holds the dequanted weight row W[cur_expert, m, 0..k_in] (f32 in + // threadgroup so the reuse dot is exact). 4096 = max k_in (gate/up); + // down (k_in=2048) uses the low half. + threadgroup_alloc("Ws", 4096, f32); + + let mut cur_expert = 4294967295u32; + for r in range(0u32, rows_per_tile, 1u32) { + let row = row0 + r; + if row < m_total { + let expert = load(expert_ids[row]); + // (Re)dequant W[expert, m, :] into Ws on an expert change. The + // test is threadgroup-uniform → the barriers are safe. A row-dot + // only READS Ws, so no per-row barrier is needed — only fence Ws + // before overwriting it on a (rare) expert change. + if expert != cur_expert { + if cur_expert != 4294967295u32 { + threadgroup_barrier(); // prior rows' reads complete before overwrite + } + let qs_row_base = (expert * nblk_per_expert + m * blocks_per_row) * 16u32; + let d_row_base = expert * nblk_per_expert + m * blocks_per_row; + for grp in range(lane, total_groups, 32u32) { + let b = grp / 8u32; + let group = grp & 7u32; + let aux_idx = load(qs_all[qs_row_base + b * 16u32 + group * 2u32]); + let aux_sgn = load(qs_all[qs_row_base + b * 16u32 + group * 2u32 + 1u32]); + let scale_4bit = aux_sgn >> 28u32; + let db = load(d_all[d_row_base + b]) + * ((scale_4bit.cast::().cast::() + 0.5) * 0.25); + let x_grp = b * 256u32 + group * 32u32; + for j in range(0u32, 4u32, 1u32) { + let grid_key = (aux_idx >> (j * 8u32)) & 0xffu32; + let grid_row_base = grid_key * 8u32; + let sign_idx = (aux_sgn >> (j * 7u32)) & 0x7fu32; + let sign_mask = load(signs[sign_idx]).cast::(); + for l in range(0u32, 8u32, 1u32) { + let octet = load(grid[grid_row_base + l]) + .cast::() + .cast::() + .cast::(); + let lane_bit = sign_mask & (1u32 << l); + let sign = select(lane_bit != 0u32, -1.0f32, 1.0f32); + let w = db * sign * octet; + threadgroup_store("Ws", x_grp + j * 8u32 + l, w); + } + } + } + threadgroup_barrier(); + cur_expert = expert; + } + // Dot Ws . x[row] over this lane's k-stride, then simd-reduce. + let x_base = row * k_in; + let mut acc = 0.0f32; + for grp in range(lane, total_groups, 32u32) { + let b = grp / 8u32; + let group = grp & 7u32; + let x_grp = b * 256u32 + group * 32u32; + for e in range(0u32, 32u32, 1u32) { + let k = x_grp + e; + let w = threadgroup_load("Ws", k); + let xv = load(x[x_base + k]).cast::(); + acc = acc + w * xv.cast::(); + } + } + let total = simd_sum(acc); + if lane == 0u32 { + store(out[row * m_out + m], total.cast::()); + } + } + } +} + +pub mod kernel_benches { + use metaltile::{bench, test::*}; + + use super::ffai_moe_gemv_ws_iq2xxs; + + // M=256 rows, production gate/up dims, 8 rows/tile. + #[bench(name = "ffai/moe/gemv_ws_iq2xxs", dtypes = [f32, f16, bf16])] + fn bench_gemv_ws_iq2xxs(dt: DType) -> BenchSetup { + let m_total = 256usize; + let n_experts = 8usize; + let m_out = 2048usize; + let k_in = 4096usize; + let rows_per_tile = 8usize; + let nblk = m_out * (k_in / 256); + BenchSetup::new(ffai_moe_gemv_ws_iq2xxs::kernel_ir_for(dt)) + .mode(KernelMode::Reduction) + .buffer(BenchBuffer::random("x", m_total * k_in, dt)) + .buffer(BenchBuffer::random("qs_all", n_experts * nblk * 16, DType::U32)) + .buffer(BenchBuffer::random("d_all", n_experts * nblk, DType::F32)) + .buffer(BenchBuffer::zeros("expert_ids", m_total, DType::U32)) + .buffer(BenchBuffer::random("grid", 2048, DType::U8)) + .buffer(BenchBuffer::random("signs", 128, DType::U8)) + .buffer(BenchBuffer::zeros("out", m_total * m_out, dt).output()) + .constexpr("k_in", k_in as u32) + .constexpr("m_out", m_out as u32) + .constexpr("m_total", m_total as u32) + .constexpr("rows_per_tile", rows_per_tile as u32) + .grid_3d(m_out as u32, (m_total as u32).div_ceil(rows_per_tile as u32), 1, [32, 1, 1]) + .bytes_moved((m_total * nblk * 64 + m_total * k_in * dt.size_bytes()) as u64) + } +} diff --git a/crates/metaltile-std/src/ffai/moe_gemv_ws_q2k.rs b/crates/metaltile-std/src/ffai/moe_gemv_ws_q2k.rs new file mode 100644 index 00000000..8919a007 --- /dev/null +++ b/crates/metaltile-std/src/ffai/moe_gemv_ws_q2k.rs @@ -0,0 +1,134 @@ +//! Copyright 2026 0xClandestine, Ekryski, TheTom, Ambisphaeric +//! SPDX-License-Identifier: Apache-2.0 +//! Prefill MoE Q2_K WEIGHT-STATIONARY gemv (down projection) — the Q2_K +//! twin of ffai_moe_gemv_ws_iq2xxs. Dequants each expert's weight row +//! W_down[expert,m,:] ONCE into threadgroup memory and reuses it across +//! all rows of that expert in the tile (amortized like bm64 but at gemv +//! speed). Rows are pre-permuted by expert (contiguous), so a tile is +//! usually single-expert; on an expert boundary it re-dequants. The +//! expert-change test is threadgroup-uniform → the barrier is safe. +//! Canonical Q2_K layout + dequant identical to gemv_rows_q2k. +//! +//! grid (threadgroups) = [m_out, ceil(m_total/rows_per_tile), 1], +//! threadgroup = [32,1,1]. + +use metaltile::kernel; + +#[kernel] +#[allow(clippy::too_many_arguments)] +pub fn ffai_moe_gemv_ws_q2k( + x: Tensor, + qs: Tensor, + scales: Tensor, + d_f32: Tensor, + dmin_f32: Tensor, + expert_ids: Tensor, + mut out: Tensor, + #[constexpr] k_in: u32, + #[constexpr] m_out: u32, + #[constexpr] m_total: u32, + #[constexpr] rows_per_tile: u32, +) { + let m = tgid_x; + let tile = tgid_y; + let lane = tid; + let row0 = tile * rows_per_tile; + + let blocks_per_row = k_in / 256u32; + let nblk_per_expert = m_out * blocks_per_row; + + // Ws holds the dequanted weight row W[cur_expert, m, 0..k_in] (f32). + // 4096 ≥ k_in (down k_in = intermediate = 2048). + threadgroup_alloc("Ws", 4096, f32); + + let mut cur_expert = 4294967295u32; + for r in range(0u32, rows_per_tile, 1u32) { + let row = row0 + r; + if row < m_total { + let expert = load(expert_ids[row]); + if expert != cur_expert { + if cur_expert != 4294967295u32 { + threadgroup_barrier(); + } + let qs_expert_base = (expert * nblk_per_expert) * 16u32; + let sc_expert_base = (expert * nblk_per_expert) * 16u32; + let blk_expert_base = expert * nblk_per_expert; + for k in range(lane, k_in, 32u32) { + let vidx = m * k_in + k; + let block = vidx / 256u32; + let in_block = vidx & 255u32; + let half = in_block / 128u32; + let yh = in_block - half * 128u32; + let jg = yh / 32u32; + let yg = yh - jg * 32u32; + let sub_half = yg / 16u32; + let l = yg - sub_half * 16u32; + let shift = jg * 2u32; + let q_byte = half * 32u32 + sub_half * 16u32 + l; + let sub = half * 8u32 + jg * 2u32 + sub_half; + let word_idx = q_byte / 4u32; + let byte_in_word = q_byte & 3u32; + let word = load(qs[qs_expert_base + block * 16u32 + word_idx]); + let qs_byte = (word >> (byte_in_word * 8u32)) & 0xffu32; + let q_2bit = (qs_byte >> shift) & 0x3u32; + let scale_byte = + load(scales[sc_expert_base + block * 16u32 + sub]).cast::(); + let scale_4bit = scale_byte & 0xfu32; + let min_4bit = (scale_byte >> 4u32) & 0xfu32; + let d = load(d_f32[blk_expert_base + block]); + let dmin = load(dmin_f32[blk_expert_base + block]); + let wq = d + * scale_4bit.cast::().cast::() + * q_2bit.cast::().cast::() + - dmin * min_4bit.cast::().cast::(); + threadgroup_store("Ws", k, wq); + } + threadgroup_barrier(); + cur_expert = expert; + } + let x_base = row * k_in; + let mut acc = 0.0f32; + for k in range(lane, k_in, 32u32) { + let w = threadgroup_load("Ws", k); + let xv = load(x[x_base + k]).cast::(); + acc = acc + w * xv.cast::(); + } + let total = simd_sum(acc); + if lane == 0u32 { + store(out[row * m_out + m], total.cast::()); + } + } + } +} + +pub mod kernel_benches { + use metaltile::{bench, test::*}; + + use super::ffai_moe_gemv_ws_q2k; + + // M=256 rows, production down dims (m_out=4096 hidden, k_in=2048). + #[bench(name = "ffai/moe/gemv_ws_q2k", dtypes = [f32, f16, bf16])] + fn bench_gemv_ws_q2k(dt: DType) -> BenchSetup { + let m_total = 256usize; + let n_experts = 8usize; + let m_out = 4096usize; + let k_in = 2048usize; + let rows_per_tile = 8usize; + let nblk = m_out * (k_in / 256); + BenchSetup::new(ffai_moe_gemv_ws_q2k::kernel_ir_for(dt)) + .mode(KernelMode::Reduction) + .buffer(BenchBuffer::random("x", m_total * k_in, dt)) + .buffer(BenchBuffer::random("qs", n_experts * nblk * 16, DType::U32)) + .buffer(BenchBuffer::random("scales", n_experts * nblk * 16, DType::U8)) + .buffer(BenchBuffer::random("d_f32", n_experts * nblk, DType::F32)) + .buffer(BenchBuffer::random("dmin_f32", n_experts * nblk, DType::F32)) + .buffer(BenchBuffer::zeros("expert_ids", m_total, DType::U32)) + .buffer(BenchBuffer::zeros("out", m_total * m_out, dt).output()) + .constexpr("k_in", k_in as u32) + .constexpr("m_out", m_out as u32) + .constexpr("m_total", m_total as u32) + .constexpr("rows_per_tile", rows_per_tile as u32) + .grid_3d(m_out as u32, (m_total as u32).div_ceil(rows_per_tile as u32), 1, [32, 1, 1]) + .bytes_moved((m_total * nblk * 84 + m_total * k_in * dt.size_bytes()) as u64) + } +} diff --git a/crates/metaltile-std/src/ffai/moe_router_sqrtsoftplus.rs b/crates/metaltile-std/src/ffai/moe_router_sqrtsoftplus.rs new file mode 100644 index 00000000..9bf4467d --- /dev/null +++ b/crates/metaltile-std/src/ffai/moe_router_sqrtsoftplus.rs @@ -0,0 +1,156 @@ +//! Copyright 2026 0xClandestine, Ekryski, TheTom, Ambisphaeric +//! SPDX-License-Identifier: Apache-2.0 +//! MoE router — sqrt(softplus(·)) + bias-correction scoring. +//! +//! The DeepSeek V4 (Flash + Pro) routing pattern, replacing V3.2's +//! sigmoid+bias gate per the V4 paper's `scoring_func: sqrtsoftplus`. +//! +//! ```text +//! score_unbiased[e] = sqrt(softplus(logits[e])) // routing signal +//! score_biased[e] = score_unbiased[e] + bias[e] // selection signal +//! ``` +//! +//! Top-k selection uses `score_biased`; the gather weight downstream +//! is `score_unbiased * routed_scaling_factor` (the bias is for +//! aux-loss-free load balancing, not weight magnitude — V3.2 / V4 +//! `noaux_tc` mechanism). Both outputs ship side-by-side so the +//! downstream top-k + normalize + scale chain consumes whichever it +//! needs without re-running the scoring math. +//! +//! Compared to the sister `ffai_moe_router_sigmoid_bias`: +//! - sigmoid+bias: `s = sigmoid(x)`; bounded in (0, 1). +//! - sqrtsoftplus: `s = sqrt(log(1 + exp(x)))`; unbounded, larger +//! dynamic range — paired with the bias-correction for selection. +//! +//! Pitfall: softplus(x) overflows in fp32 for large x. We use the +//! numerically-stable form `softplus(x) = max(x, 0) + log(1 + exp(-|x|))` +//! to avoid `exp` overflow on positive tails. +//! +//! ## ABI +//! +//! ```text +//! logits [n_experts] f32 — pre-scoring router output (`W_router · x`) +//! bias [n_experts] f32 — per-expert routing-bias (noaux_tc) +//! score_unbiased [n_experts] f32 — out: `sqrt(softplus(logits))` +//! score_biased [n_experts] f32 — out: `score_unbiased + bias` +//! ``` +//! +//! Grid is 1D elementwise — one thread per expert. Modal n_experts +//! across the production checkpoints is 288 (V4-Flash, ~256K-ctx +//! Pareto point); the kernel scales fine to V4-Pro's larger counts. + +use metaltile::kernel; + +// Bare `#[kernel]` — non-generic, all-f32 kernel doesn't fit the +// legacy `bench(...)` shape; declarative `#[bench]` below registers +// for `tile bench`. +#[kernel] +pub fn ffai_moe_router_sqrtsoftplus( + logits: Tensor, + bias: Tensor, + mut score_unbiased: Tensor, + mut score_biased: Tensor, +) { + let idx = tid; + let logit_val = load(logits[idx]); + // Numerically stable softplus: `max(x, 0) + log(1 + exp(-|x|))`. + // For x >> 0: ≈ x + log(1 + tiny) ≈ x + // For x << 0: ≈ 0 + log(1 + e^x) ≈ e^x + // + // Variable names deliberately avoid the `score` prefix — the DSL's + // codegen identifier mangler matches argument prefixes, so a local + // `score` collides with the `score_unbiased` / `score_biased` Tensor + // params and the binding is silently elided. + // Use the DSL free-function forms `log` / `exp` — the method-call + // forms `.ln()` / `.exp()` don't always emit their bindings when + // chained inside larger expressions (see e.g. + // `gated_delta_prep.rs` which uses the function form for the + // same `log(exp(x) + 1)` softplus pattern). + let abs_logit = select(logit_val >= 0.0f32, logit_val, -logit_val); + let neg_abs = 0.0f32 - abs_logit; + let exp_neg = exp(neg_abs); + let one_plus = 1.0f32 + exp_neg; + let log_term = log(one_plus); + let pos_part = select(logit_val >= 0.0f32, logit_val, 0.0f32); + let softplus_val = pos_part + log_term; + let routing_signal = sqrt(softplus_val); + let bias_val = load(bias[idx]); + store(score_unbiased[idx], routing_signal); + store(score_biased[idx], routing_signal + bias_val); +} + +pub mod kernel_tests { + use metaltile::{test::*, test_kernel}; + + use super::ffai_moe_router_sqrtsoftplus; + use crate::utils::pack_f32; + + /// CPU reference. Mirrors the GPU kernel for tight-tolerance check. + fn cpu_reference(n_experts: usize) -> (Vec, Vec, Vec, Vec) { + let dt = DType::F32; + // Logits cover the full numerically-interesting range: large + // positive (overflow tail of naive `exp(x)`), large negative + // (underflow tail of `log(1 + exp(x))` for naive `softplus`), + // and the dense-around-zero region where routing actually + // discriminates. + let logits: Vec = + (0..n_experts).map(|i| (i as f32 - n_experts as f32 / 2.0) * 0.25).collect(); + let bias: Vec = (0..n_experts).map(|i| (i % 13) as f32 * 0.02 - 0.13).collect(); + // Stable softplus matches the kernel exactly. + let score_unbiased: Vec = logits + .iter() + .map(|&x| { + let ax = x.abs(); + let pos = if x >= 0.0 { x } else { 0.0 }; + (pos + (1.0 + (-ax).exp()).ln()).sqrt() + }) + .collect(); + let score_biased: Vec = score_unbiased.iter().zip(&bias).map(|(s, b)| s + b).collect(); + let _ = dt; // dtype implied by the test_kernel decl + (logits, bias, score_unbiased, score_biased) + } + + fn setup(n_experts: usize) -> TestSetup { + let dt = DType::F32; + let (logits, bias, score_unbiased, score_biased) = cpu_reference(n_experts); + TestSetup::new(ffai_moe_router_sqrtsoftplus::kernel_ir()) + .input(TestBuffer::from_vec("logits", pack_f32(&logits, dt), dt)) + .input(TestBuffer::from_vec("bias", pack_f32(&bias, dt), dt)) + .input(TestBuffer::zeros("score_unbiased", n_experts, dt)) + .input(TestBuffer::zeros("score_biased", n_experts, dt)) + .expect(TestBuffer::from_vec("score_unbiased", pack_f32(&score_unbiased, dt), dt)) + .expect(TestBuffer::from_vec("score_biased", pack_f32(&score_biased, dt), dt)) + .grid_1d(n_experts, 64) + } + + // tol 2e-4: sqrtsoftplus routes through exp + sqrt, so the GPU↔CPU-oracle + // gap is dominated by transcendental fast-math and varies by GPU family + // (~5e-7 on M5, ~6e-5 on older Apple GPUs). 2e-4 holds across hardware + // while staying far tighter than a logic error would ever land. + /// V4-Flash production shape — 288 experts. + #[test_kernel(dtypes = [f32], tol = [2e-4])] + fn test_router_sqrtsoftplus_v4_flash(_dt: DType) -> TestSetup { setup(288) } + + /// V4-Pro / future-larger shape sanity check. + #[test_kernel(dtypes = [f32], tol = [2e-4])] + fn test_router_sqrtsoftplus_large(_dt: DType) -> TestSetup { setup(512) } +} + +pub mod kernel_benches { + use metaltile::{bench, test::*}; + + use super::ffai_moe_router_sqrtsoftplus; + + #[bench(name = "ffai/moe_router_sqrtsoftplus", dtypes = [f32])] + fn bench_router(_dt: DType) -> BenchSetup { + let dt = DType::F32; + let n_experts = 288usize; + BenchSetup::new(ffai_moe_router_sqrtsoftplus::kernel_ir()) + .buffer(BenchBuffer::random("logits", n_experts, dt)) + .buffer(BenchBuffer::random("bias", n_experts, dt)) + .buffer(BenchBuffer::zeros("score_unbiased", n_experts, dt).output()) + .buffer(BenchBuffer::zeros("score_biased", n_experts, dt).output()) + .grid_1d(n_experts, 64) + .bytes_moved((4 * n_experts * dt.size_bytes()) as u64) + } +} diff --git a/crates/metaltile-std/src/ffai/mt_vector_add.rs b/crates/metaltile-std/src/ffai/mt_vector_add.rs new file mode 100644 index 00000000..f94f6b05 --- /dev/null +++ b/crates/metaltile-std/src/ffai/mt_vector_add.rs @@ -0,0 +1,62 @@ +//! Copyright 2026 0xClandestine, Ekryski, TheTom, Ambisphaeric +//! SPDX-License-Identifier: Apache-2.0 +//! Generic elementwise vector add — `out[i] = a[i] + b[i]`. +//! +//! Trivial kernel that didn't already exist in the binary-op surface +//! (which is shape-spec-tied to MLX's `Binary` class registration). +//! Lands in `ffai/` so the bare-`#[kernel]` form emits cleanly into +//! `MetalTileKernels.swift` for FFAI callers. + +use metaltile::kernel; + +#[kernel] +pub fn ffai_vector_add(a: Tensor, b: Tensor, mut out: Tensor) { + let i = tid; + let av = load(a[i]).cast::(); + let bv = load(b[i]).cast::(); + store(out[i], av + bv); +} + +pub mod kernel_tests { + use metaltile::{test::*, test_kernel}; + + use super::ffai_vector_add; + use crate::utils::{pack_f32, unpack_f32}; + + fn setup(n: usize, dt: DType) -> TestSetup { + let a: Vec = (0..n).map(|i| (i as f32 * 0.013 - 0.4).sin() * 1.2).collect(); + let b: Vec = (0..n).map(|i| (i as f32 * 0.017 + 0.1).cos() * 0.8).collect(); + let a_dt = unpack_f32(&pack_f32(&a, dt), dt); + let b_dt = unpack_f32(&pack_f32(&b, dt), dt); + let expected: Vec = a_dt.iter().zip(&b_dt).map(|(x, y)| x + y).collect(); + TestSetup::new(ffai_vector_add::kernel_ir_for(dt)) + .input(TestBuffer::from_vec("a", pack_f32(&a, dt), dt)) + .input(TestBuffer::from_vec("b", pack_f32(&b, dt), dt)) + .input(TestBuffer::zeros("out", n, dt)) + .expect(TestBuffer::from_vec("out", pack_f32(&expected, dt), dt)) + .grid_1d(n, 256) + } + + #[test_kernel(dtypes = [f32, f16, bf16], tol = [1e-5, 5e-3, 5e-2])] + fn test_vector_add_decode(dt: DType) -> TestSetup { setup(4096, dt) } + + #[test_kernel(dtypes = [f32, f16, bf16], tol = [1e-5, 5e-3, 5e-2])] + fn test_vector_add_long(dt: DType) -> TestSetup { setup(16384, dt) } +} + +pub mod kernel_benches { + use metaltile::{bench, test::*}; + + use super::ffai_vector_add; + + #[bench(name = "ffai/vector_add", dtypes = [f32, f16, bf16])] + fn bench_add(dt: DType) -> BenchSetup { + let n = 4096usize; + BenchSetup::new(ffai_vector_add::kernel_ir_for(dt)) + .buffer(BenchBuffer::random("a", n, dt)) + .buffer(BenchBuffer::random("b", n, dt)) + .buffer(BenchBuffer::zeros("out", n, dt).output()) + .grid_1d(n, 256) + .bytes_moved((3 * n * dt.size_bytes()) as u64) + } +} diff --git a/crates/metaltile-std/src/ffai/sdpa_decode_d512_sink.rs b/crates/metaltile-std/src/ffai/sdpa_decode_d512_sink.rs new file mode 100644 index 00000000..4a972b25 --- /dev/null +++ b/crates/metaltile-std/src/ffai/sdpa_decode_d512_sink.rs @@ -0,0 +1,437 @@ +//! Copyright 2026 0xClandestine, Ekryski, TheTom, Ambisphaeric +//! SPDX-License-Identifier: Apache-2.0 +//! Single-token SDPA decode for `head_dim == 512` with **GPT-OSS-style +//! attention sink** — a per-head learnable scalar that participates in +//! the softmax denominator but contributes no value. +//! +//! ```text +//! logits[t] = (Q · Kₜ) * scale +//! M = max(max_t logits[t], sink_logit[q_head]) +//! denom = sum_t exp(logits[t] - M) + exp(sink_logit[q_head] - M) +//! out[d] = sum_t (exp(logits[t] - M) / denom) · V[t, d] +//! ``` +//! +//! Used by DeepSeek V4 HCA (hierarchical-coarse-attention) dense +//! layers — the model's `attn_sink` parameter is a `[n_heads]` fp32 +//! tensor learned alongside Q/K/V/O. +//! +//! Clone of [`crate::ffai::sdpa_decode_d512`] with the sink fold +//! applied in the cross-simdgroup reduction (where `g_max` and `g_sum` +//! are finalised). All other dispatch invariants — TPG=512, 16 dims +//! per lane, 4-phase output reduction — are identical. + +use metaltile::kernel; + +#[kernel] +pub fn ffai_sdpa_decode_d512_sink( + q: Tensor, + k: Tensor, + v: Tensor, + sink_logit: Tensor, + out: Tensor, + #[constexpr] head_dim: u32, + #[constexpr] n_kv: u32, + #[constexpr] kv_stride: u32, + #[constexpr] heads_per_group: u32, + #[constexpr] scale: f32, +) { + let q_head = tgid_x; + let kv_head = q_head / heads_per_group; + let sg = simd_id; + let lane = simd_lane; + let ns = n_simd; + threadgroup_alloc("tg_max", 32); + threadgroup_alloc("tg_sum", 32); + threadgroup_alloc("tg_out0", 1056); + threadgroup_alloc("tg_out1", 1056); + threadgroup_alloc("tg_out2", 1056); + threadgroup_alloc("tg_out3", 1056); + let q_off = q_head * head_dim; + let kv_head_base = kv_head * kv_stride * head_dim; + let d0 = lane * 16u32; + let q0 = load(q[q_off + d0]).cast::() * scale; + let q1 = load(q[q_off + d0 + 1u32]).cast::() * scale; + let q2 = load(q[q_off + d0 + 2u32]).cast::() * scale; + let q3 = load(q[q_off + d0 + 3u32]).cast::() * scale; + let q4 = load(q[q_off + d0 + 4u32]).cast::() * scale; + let q5 = load(q[q_off + d0 + 5u32]).cast::() * scale; + let q6 = load(q[q_off + d0 + 6u32]).cast::() * scale; + let q7 = load(q[q_off + d0 + 7u32]).cast::() * scale; + let q8 = load(q[q_off + d0 + 8u32]).cast::() * scale; + let q9 = load(q[q_off + d0 + 9u32]).cast::() * scale; + let q10 = load(q[q_off + d0 + 10u32]).cast::() * scale; + let q11 = load(q[q_off + d0 + 11u32]).cast::() * scale; + let q12 = load(q[q_off + d0 + 12u32]).cast::() * scale; + let q13 = load(q[q_off + d0 + 13u32]).cast::() * scale; + let q14 = load(q[q_off + d0 + 14u32]).cast::() * scale; + let q15 = load(q[q_off + d0 + 15u32]).cast::() * scale; + let mut run_max = neg_infinity(); + let mut run_sum = 0.0f32; + let mut o0 = 0.0f32; + let mut o1 = 0.0f32; + let mut o2 = 0.0f32; + let mut o3 = 0.0f32; + let mut o4 = 0.0f32; + let mut o5 = 0.0f32; + let mut o6 = 0.0f32; + let mut o7 = 0.0f32; + let mut o8 = 0.0f32; + let mut o9 = 0.0f32; + let mut o10 = 0.0f32; + let mut o11 = 0.0f32; + let mut o12 = 0.0f32; + let mut o13 = 0.0f32; + let mut o14 = 0.0f32; + let mut o15 = 0.0f32; + for _t in range(sg, n_kv, ns) { + let base = kv_head_base + _t * head_dim; + let kv0 = base + d0; + let k0 = load(k[kv0]).cast::(); + let k1 = load(k[kv0 + 1u32]).cast::(); + let k2 = load(k[kv0 + 2u32]).cast::(); + let k3 = load(k[kv0 + 3u32]).cast::(); + let k4 = load(k[kv0 + 4u32]).cast::(); + let k5 = load(k[kv0 + 5u32]).cast::(); + let k6 = load(k[kv0 + 6u32]).cast::(); + let k7 = load(k[kv0 + 7u32]).cast::(); + let k8 = load(k[kv0 + 8u32]).cast::(); + let k9 = load(k[kv0 + 9u32]).cast::(); + let k10 = load(k[kv0 + 10u32]).cast::(); + let k11 = load(k[kv0 + 11u32]).cast::(); + let k12 = load(k[kv0 + 12u32]).cast::(); + let k13 = load(k[kv0 + 13u32]).cast::(); + let k14 = load(k[kv0 + 14u32]).cast::(); + let k15 = load(k[kv0 + 15u32]).cast::(); + let partial = q0 * k0 + + q1 * k1 + + q2 * k2 + + q3 * k3 + + q4 * k4 + + q5 * k5 + + q6 * k6 + + q7 * k7 + + q8 * k8 + + q9 * k9 + + q10 * k10 + + q11 * k11 + + q12 * k12 + + q13 * k13 + + q14 * k14 + + q15 * k15; + let score = simd_sum(partial); + let new_max = select(score > run_max, score, run_max); + let factor = exp(run_max - new_max); + let weight = exp(score - new_max); + run_sum = run_sum * factor + weight; + run_max = new_max; + let v0 = load(v[kv0]).cast::(); + let v1 = load(v[kv0 + 1u32]).cast::(); + let v2 = load(v[kv0 + 2u32]).cast::(); + let v3 = load(v[kv0 + 3u32]).cast::(); + let v4 = load(v[kv0 + 4u32]).cast::(); + let v5 = load(v[kv0 + 5u32]).cast::(); + let v6 = load(v[kv0 + 6u32]).cast::(); + let v7 = load(v[kv0 + 7u32]).cast::(); + let v8 = load(v[kv0 + 8u32]).cast::(); + let v9 = load(v[kv0 + 9u32]).cast::(); + let v10 = load(v[kv0 + 10u32]).cast::(); + let v11 = load(v[kv0 + 11u32]).cast::(); + let v12 = load(v[kv0 + 12u32]).cast::(); + let v13 = load(v[kv0 + 13u32]).cast::(); + let v14 = load(v[kv0 + 14u32]).cast::(); + let v15 = load(v[kv0 + 15u32]).cast::(); + o0 = o0 * factor + weight * v0; + o1 = o1 * factor + weight * v1; + o2 = o2 * factor + weight * v2; + o3 = o3 * factor + weight * v3; + o4 = o4 * factor + weight * v4; + o5 = o5 * factor + weight * v5; + o6 = o6 * factor + weight * v6; + o7 = o7 * factor + weight * v7; + o8 = o8 * factor + weight * v8; + o9 = o9 * factor + weight * v9; + o10 = o10 * factor + weight * v10; + o11 = o11 * factor + weight * v11; + o12 = o12 * factor + weight * v12; + o13 = o13 * factor + weight * v13; + o14 = o14 * factor + weight * v14; + o15 = o15 * factor + weight * v15; + } + if lane == 0 { + threadgroup_store("tg_max", sg, run_max); + threadgroup_store("tg_sum", sg, run_sum); + } + threadgroup_barrier(); + if sg == 0 { + let g_max_in = select(lane < ns, threadgroup_load("tg_max", lane), neg_infinity()); + let g_max = simd_max(g_max_in); + let g_sum_in = + select(lane < ns, threadgroup_load("tg_sum", lane) * exp(g_max_in - g_max), 0.0f32); + let g_sum = simd_sum(g_sum_in); + if lane == 0 { + threadgroup_store("tg_max", 0, g_max); + threadgroup_store("tg_sum", 0, g_sum); + } + } + threadgroup_barrier(); + let g_max0 = threadgroup_load("tg_max", 0); + let g_sum0 = threadgroup_load("tg_sum", 0); + // Attention-sink fold: extend the softmax denominator by one + // "virtual" slot at score = `sink_logit[q_head]`. No value + // contribution, so the output accumulators rescale unchanged by + // the new max + sum. + let sink = load(sink_logit[q_head]); + let g_max = select(sink > g_max0, sink, g_max0); + let g_sum = g_sum0 * exp(g_max0 - g_max) + exp(sink - g_max); + let rescale = select(g_sum > 0.0f32, exp(run_max - g_max) / g_sum, 0.0f32); + let stride = ns + 1u32; + let idx = lane * stride + sg; + threadgroup_store("tg_out0", idx, o0 * rescale); + threadgroup_store("tg_out1", idx, o1 * rescale); + threadgroup_store("tg_out2", idx, o2 * rescale); + threadgroup_store("tg_out3", idx, o3 * rescale); + threadgroup_barrier(); + if sg == 0 { + let mut so0 = 0.0f32; + let mut so1 = 0.0f32; + let mut so2 = 0.0f32; + let mut so3 = 0.0f32; + for _g in range(0u32, ns, 1u32) { + let ri = lane * stride + _g; + so0 = so0 + threadgroup_load("tg_out0", ri); + so1 = so1 + threadgroup_load("tg_out1", ri); + so2 = so2 + threadgroup_load("tg_out2", ri); + so3 = so3 + threadgroup_load("tg_out3", ri); + } + let out_off = q_off + d0; + store(out[out_off], so0.cast::()); + store(out[out_off + 1u32], so1.cast::()); + store(out[out_off + 2u32], so2.cast::()); + store(out[out_off + 3u32], so3.cast::()); + } + threadgroup_barrier(); + threadgroup_store("tg_out0", idx, o4 * rescale); + threadgroup_store("tg_out1", idx, o5 * rescale); + threadgroup_store("tg_out2", idx, o6 * rescale); + threadgroup_store("tg_out3", idx, o7 * rescale); + threadgroup_barrier(); + if sg == 0 { + let mut so4 = 0.0f32; + let mut so5 = 0.0f32; + let mut so6 = 0.0f32; + let mut so7 = 0.0f32; + for _g in range(0u32, ns, 1u32) { + let ri = lane * stride + _g; + so4 = so4 + threadgroup_load("tg_out0", ri); + so5 = so5 + threadgroup_load("tg_out1", ri); + so6 = so6 + threadgroup_load("tg_out2", ri); + so7 = so7 + threadgroup_load("tg_out3", ri); + } + let out_off = q_off + d0; + store(out[out_off + 4u32], so4.cast::()); + store(out[out_off + 5u32], so5.cast::()); + store(out[out_off + 6u32], so6.cast::()); + store(out[out_off + 7u32], so7.cast::()); + } + threadgroup_barrier(); + threadgroup_store("tg_out0", idx, o8 * rescale); + threadgroup_store("tg_out1", idx, o9 * rescale); + threadgroup_store("tg_out2", idx, o10 * rescale); + threadgroup_store("tg_out3", idx, o11 * rescale); + threadgroup_barrier(); + if sg == 0 { + let mut so8 = 0.0f32; + let mut so9 = 0.0f32; + let mut so10 = 0.0f32; + let mut so11 = 0.0f32; + for _g in range(0u32, ns, 1u32) { + let ri = lane * stride + _g; + so8 = so8 + threadgroup_load("tg_out0", ri); + so9 = so9 + threadgroup_load("tg_out1", ri); + so10 = so10 + threadgroup_load("tg_out2", ri); + so11 = so11 + threadgroup_load("tg_out3", ri); + } + let out_off = q_off + d0; + store(out[out_off + 8u32], so8.cast::()); + store(out[out_off + 9u32], so9.cast::()); + store(out[out_off + 10u32], so10.cast::()); + store(out[out_off + 11u32], so11.cast::()); + } + threadgroup_barrier(); + threadgroup_store("tg_out0", idx, o12 * rescale); + threadgroup_store("tg_out1", idx, o13 * rescale); + threadgroup_store("tg_out2", idx, o14 * rescale); + threadgroup_store("tg_out3", idx, o15 * rescale); + threadgroup_barrier(); + if sg == 0 { + let mut so12 = 0.0f32; + let mut so13 = 0.0f32; + let mut so14 = 0.0f32; + let mut so15 = 0.0f32; + for _g in range(0u32, ns, 1u32) { + let ri = lane * stride + _g; + so12 = so12 + threadgroup_load("tg_out0", ri); + so13 = so13 + threadgroup_load("tg_out1", ri); + so14 = so14 + threadgroup_load("tg_out2", ri); + so15 = so15 + threadgroup_load("tg_out3", ri); + } + let out_off = q_off + d0; + store(out[out_off + 12u32], so12.cast::()); + store(out[out_off + 13u32], so13.cast::()); + store(out[out_off + 14u32], so14.cast::()); + store(out[out_off + 15u32], so15.cast::()); + } +} + +#[cfg(test)] +mod tests { + use metaltile_codegen::msl::MslGenerator; + use metaltile_core::ir::KernelMode; + + use super::ffai_sdpa_decode_d512_sink; + use crate::bench_types::DType; + + fn msl_for(dt: DType) -> String { + let mut k = ffai_sdpa_decode_d512_sink::kernel_ir_for(dt); + k.mode = KernelMode::Reduction; + MslGenerator::default().generate(&k).expect("ffai_sdpa_decode_d512_sink codegen succeeds") + } + + #[test] + fn codegen_produces_nonempty_msl_for_all_float_dtypes() { + for dt in [DType::F32, DType::F16, DType::BF16] { + let src = msl_for(dt); + assert!(!src.trim().is_empty(), "MSL for {dt:?} should not be empty"); + assert!( + src.contains("kernel void ffai_sdpa_decode_d512_sink"), + "MSL for {dt:?} should declare ffai_sdpa_decode_d512_sink:\n{src}", + ); + } + } +} + +#[allow(clippy::needless_range_loop)] +pub mod kernel_tests { + use metaltile::{test::*, test_kernel}; + + use super::ffai_sdpa_decode_d512_sink; + use crate::utils::{pack_f32, unpack_f32}; + + /// Dense softmax-attention oracle with a per-head `sink_logit` + /// scalar folded into the denominator (no value contribution). + #[allow(clippy::too_many_arguments)] + fn naive_sdpa_sink( + q: &[f32], + k: &[f32], + v: &[f32], + sink: &[f32], + n_q_heads: usize, + n_kv_heads: usize, + head_dim: usize, + n_kv: usize, + kv_stride: usize, + scale: f32, + ) -> Vec { + let gqa = n_q_heads / n_kv_heads; + let mut out = vec![0.0f32; n_q_heads * head_dim]; + for qh in 0..n_q_heads { + let kvh = qh / gqa; + let q_off = qh * head_dim; + let kv_slab = kvh * kv_stride * head_dim; + let mut scores = vec![0.0f32; n_kv]; + for (t, score) in scores.iter_mut().enumerate() { + let k_off = kv_slab + t * head_dim; + let mut dot = 0.0f32; + for d in 0..head_dim { + dot += q[q_off + d] * k[k_off + d]; + } + *score = dot * scale; + } + let mut m = scores.iter().copied().fold(f32::NEG_INFINITY, f32::max); + m = m.max(sink[qh]); + let mut sum = 0.0f32; + for s in scores.iter_mut() { + *s = (*s - m).exp(); + sum += *s; + } + sum += (sink[qh] - m).exp(); + let inv = if sum > 0.0 { 1.0 / sum } else { 0.0 }; + for d in 0..head_dim { + let mut acc = 0.0f32; + for (t, s) in scores.iter().enumerate() { + acc += *s * inv * v[kv_slab + t * head_dim + d]; + } + out[q_off + d] = acc; + } + } + out + } + + fn ramp(n: usize, step: f32, start: f32) -> Vec { + (0..n).map(|i| ((start + i as f32 * step) % 2.0) - 1.0).collect() + } + + #[test_kernel(dtypes = [f32, f16, bf16], tol = [1e-3, 3e-3, 1.5e-2])] + fn test_ffai_sdpa_decode_d512_sink(dt: DType) -> TestSetup { + let (n_q_heads, n_kv_heads, head_dim) = (8usize, 4usize, 512usize); + let (n_kv, kv_stride) = (64usize, 64usize); + let heads_per_group = n_q_heads / n_kv_heads; + let scale = 1.0f32 / (head_dim as f32).sqrt(); + let q = unpack_f32(&pack_f32(&ramp(n_q_heads * head_dim, 0.013, -0.4), dt), dt); + let k = + unpack_f32(&pack_f32(&ramp(n_kv_heads * kv_stride * head_dim, 0.011, -0.5), dt), dt); + let v = + unpack_f32(&pack_f32(&ramp(n_kv_heads * kv_stride * head_dim, 0.007, -0.3), dt), dt); + // Per-head sink — non-trivial scalars; some larger than expected + // max-score, some smaller, so the test exercises both branches + // of the sink-fold max selection. + let sink: Vec = (0..n_q_heads).map(|h| (h as f32 - 3.0) * 0.4).collect(); + let expected = naive_sdpa_sink( + &q, &k, &v, &sink, n_q_heads, n_kv_heads, head_dim, n_kv, kv_stride, scale, + ); + TestSetup::new(ffai_sdpa_decode_d512_sink::kernel_ir_for(dt)) + .mode(KernelMode::Reduction) + .input(TestBuffer::from_vec("q", pack_f32(&q, dt), dt)) + .input(TestBuffer::from_vec("k", pack_f32(&k, dt), dt)) + .input(TestBuffer::from_vec("v", pack_f32(&v, dt), dt)) + .input(TestBuffer::from_vec("sink_logit", pack_f32(&sink, DType::F32), DType::F32)) + .input(TestBuffer::zeros("out", n_q_heads * head_dim, dt)) + .constexpr("head_dim", head_dim as u32) + .constexpr("n_kv", n_kv as u32) + .constexpr("kv_stride", kv_stride as u32) + .constexpr("heads_per_group", heads_per_group as u32) + .constexpr("scale", scale) + .expect(TestBuffer::from_vec("out", pack_f32(&expected, dt), dt)) + .grid_3d(n_q_heads as u32, 1, 1, [512, 1, 1]) + } +} + +pub mod kernel_benches { + use metaltile::{bench, test::*}; + + use super::ffai_sdpa_decode_d512_sink; + + #[bench(name = "ffai/sdpa_decode_d512_sink", dtypes = [f32, f16, bf16])] + fn bench_sdpa_decode_d512_sink(dt: DType) -> BenchSetup { + let (n_q_heads, n_kv_heads, head_dim) = (32usize, 8usize, 512usize); + let (n_kv, kv_stride) = (4096usize, 4096usize); + let heads_per_group = n_q_heads / n_kv_heads; + let scale = 1.0f32 / (head_dim as f32).sqrt(); + let bytes = (2 * n_q_heads * head_dim + 2 * n_kv_heads * n_kv * head_dim) * dt.size_bytes() + + n_q_heads * 4; + BenchSetup::new(ffai_sdpa_decode_d512_sink::kernel_ir_for(dt)) + .mode(KernelMode::Reduction) + .buffer(BenchBuffer::random("q", n_q_heads * head_dim, dt)) + .buffer(BenchBuffer::random("k", n_kv_heads * kv_stride * head_dim, dt)) + .buffer(BenchBuffer::random("v", n_kv_heads * kv_stride * head_dim, dt)) + .buffer(BenchBuffer::random("sink_logit", n_q_heads, DType::F32)) + .buffer(BenchBuffer::zeros("out", n_q_heads * head_dim, dt).output()) + .constexpr("head_dim", head_dim as u32) + .constexpr("n_kv", n_kv as u32) + .constexpr("kv_stride", kv_stride as u32) + .constexpr("heads_per_group", heads_per_group as u32) + .constexpr("scale", scale) + .grid_3d(n_q_heads as u32, 1, 1, [512, 1, 1]) + .bytes_moved(bytes as u64) + } +} diff --git a/crates/metaltile-std/src/ffai/sdpa_prefill_d512_sink.rs b/crates/metaltile-std/src/ffai/sdpa_prefill_d512_sink.rs new file mode 100644 index 00000000..8dc89396 --- /dev/null +++ b/crates/metaltile-std/src/ffai/sdpa_prefill_d512_sink.rs @@ -0,0 +1,308 @@ +//! Copyright 2026 0xClandestine, Ekryski, TheTom, Ambisphaeric +//! SPDX-License-Identifier: Apache-2.0 +//! Multi-query causal sliding-window SDPA for head_dim=512 with attention +//! sink — the PREFILL counterpart of `ffai_sdpa_decode_d512_sink`. Each +//! (q_head, q_pos) threadgroup runs the same flash online-softmax over the +//! KV cache, but bounded causally: query at absolute position +//! `p = kv_base + q_pos` attends KV `[max(0, p+1-window) .. p]`. For the +//! DSv4 full-attn (compress_ratio=0) layers, `window=128`; pass a huge +//! window for full causal. MQA (n_kv_heads via heads_per_group). +//! +//! q/out are `[n_query, n_q_heads, head_dim]`; k/v are `[kv_stride, +//! head_dim]` per kv-head (the resident/prefill KV cache, absolute pos). + +use metaltile::kernel; + +#[kernel] +#[allow(clippy::too_many_arguments)] +pub fn ffai_sdpa_prefill_d512_sink( + q: Tensor, + k: Tensor, + v: Tensor, + sink_logit: Tensor, + out: Tensor, + #[constexpr] head_dim: u32, + #[constexpr] n_q_heads: u32, + #[constexpr] kv_stride: u32, + #[constexpr] heads_per_group: u32, + #[constexpr] window: u32, + #[constexpr] kv_base: u32, + #[constexpr] scale: f32, +) { + let q_head = tgid_x; + let q_pos = tgid_y; + let kv_head = q_head / heads_per_group; + let sg = simd_id; + let lane = simd_lane; + let ns = n_simd; + threadgroup_alloc("tg_max", 32); + threadgroup_alloc("tg_sum", 32); + threadgroup_alloc("tg_out0", 1056); + threadgroup_alloc("tg_out1", 1056); + threadgroup_alloc("tg_out2", 1056); + threadgroup_alloc("tg_out3", 1056); + let q_off = (q_pos * n_q_heads + q_head) * head_dim; + let kv_head_base = kv_head * kv_stride * head_dim; + let d0 = lane * 16u32; + let q0 = load(q[q_off + d0]).cast::() * scale; + let q1 = load(q[q_off + d0 + 1u32]).cast::() * scale; + let q2 = load(q[q_off + d0 + 2u32]).cast::() * scale; + let q3 = load(q[q_off + d0 + 3u32]).cast::() * scale; + let q4 = load(q[q_off + d0 + 4u32]).cast::() * scale; + let q5 = load(q[q_off + d0 + 5u32]).cast::() * scale; + let q6 = load(q[q_off + d0 + 6u32]).cast::() * scale; + let q7 = load(q[q_off + d0 + 7u32]).cast::() * scale; + let q8 = load(q[q_off + d0 + 8u32]).cast::() * scale; + let q9 = load(q[q_off + d0 + 9u32]).cast::() * scale; + let q10 = load(q[q_off + d0 + 10u32]).cast::() * scale; + let q11 = load(q[q_off + d0 + 11u32]).cast::() * scale; + let q12 = load(q[q_off + d0 + 12u32]).cast::() * scale; + let q13 = load(q[q_off + d0 + 13u32]).cast::() * scale; + let q14 = load(q[q_off + d0 + 14u32]).cast::() * scale; + let q15 = load(q[q_off + d0 + 15u32]).cast::() * scale; + let mut run_max = neg_infinity(); + let mut run_sum = 0.0f32; + let mut o0 = 0.0f32; + let mut o1 = 0.0f32; + let mut o2 = 0.0f32; + let mut o3 = 0.0f32; + let mut o4 = 0.0f32; + let mut o5 = 0.0f32; + let mut o6 = 0.0f32; + let mut o7 = 0.0f32; + let mut o8 = 0.0f32; + let mut o9 = 0.0f32; + let mut o10 = 0.0f32; + let mut o11 = 0.0f32; + let mut o12 = 0.0f32; + let mut o13 = 0.0f32; + let mut o14 = 0.0f32; + let mut o15 = 0.0f32; + // Causal sliding-window KV range for this query's absolute position. + let p1 = kv_base + q_pos + 1u32; + let lo = select(p1 > window, p1 - window, 0u32); + for _t in range(lo + sg, p1, ns) { + let base = kv_head_base + _t * head_dim; + let kv0 = base + d0; + let k0 = load(k[kv0]).cast::(); + let k1 = load(k[kv0 + 1u32]).cast::(); + let k2 = load(k[kv0 + 2u32]).cast::(); + let k3 = load(k[kv0 + 3u32]).cast::(); + let k4 = load(k[kv0 + 4u32]).cast::(); + let k5 = load(k[kv0 + 5u32]).cast::(); + let k6 = load(k[kv0 + 6u32]).cast::(); + let k7 = load(k[kv0 + 7u32]).cast::(); + let k8 = load(k[kv0 + 8u32]).cast::(); + let k9 = load(k[kv0 + 9u32]).cast::(); + let k10 = load(k[kv0 + 10u32]).cast::(); + let k11 = load(k[kv0 + 11u32]).cast::(); + let k12 = load(k[kv0 + 12u32]).cast::(); + let k13 = load(k[kv0 + 13u32]).cast::(); + let k14 = load(k[kv0 + 14u32]).cast::(); + let k15 = load(k[kv0 + 15u32]).cast::(); + let partial = q0 * k0 + + q1 * k1 + + q2 * k2 + + q3 * k3 + + q4 * k4 + + q5 * k5 + + q6 * k6 + + q7 * k7 + + q8 * k8 + + q9 * k9 + + q10 * k10 + + q11 * k11 + + q12 * k12 + + q13 * k13 + + q14 * k14 + + q15 * k15; + let score = simd_sum(partial); + let new_max = select(score > run_max, score, run_max); + let factor = exp(run_max - new_max); + let weight = exp(score - new_max); + run_sum = run_sum * factor + weight; + run_max = new_max; + let v0 = load(v[kv0]).cast::(); + let v1 = load(v[kv0 + 1u32]).cast::(); + let v2 = load(v[kv0 + 2u32]).cast::(); + let v3 = load(v[kv0 + 3u32]).cast::(); + let v4 = load(v[kv0 + 4u32]).cast::(); + let v5 = load(v[kv0 + 5u32]).cast::(); + let v6 = load(v[kv0 + 6u32]).cast::(); + let v7 = load(v[kv0 + 7u32]).cast::(); + let v8 = load(v[kv0 + 8u32]).cast::(); + let v9 = load(v[kv0 + 9u32]).cast::(); + let v10 = load(v[kv0 + 10u32]).cast::(); + let v11 = load(v[kv0 + 11u32]).cast::(); + let v12 = load(v[kv0 + 12u32]).cast::(); + let v13 = load(v[kv0 + 13u32]).cast::(); + let v14 = load(v[kv0 + 14u32]).cast::(); + let v15 = load(v[kv0 + 15u32]).cast::(); + o0 = o0 * factor + weight * v0; + o1 = o1 * factor + weight * v1; + o2 = o2 * factor + weight * v2; + o3 = o3 * factor + weight * v3; + o4 = o4 * factor + weight * v4; + o5 = o5 * factor + weight * v5; + o6 = o6 * factor + weight * v6; + o7 = o7 * factor + weight * v7; + o8 = o8 * factor + weight * v8; + o9 = o9 * factor + weight * v9; + o10 = o10 * factor + weight * v10; + o11 = o11 * factor + weight * v11; + o12 = o12 * factor + weight * v12; + o13 = o13 * factor + weight * v13; + o14 = o14 * factor + weight * v14; + o15 = o15 * factor + weight * v15; + } + if lane == 0 { + threadgroup_store("tg_max", sg, run_max); + threadgroup_store("tg_sum", sg, run_sum); + } + threadgroup_barrier(); + if sg == 0 { + let g_max_in = select(lane < ns, threadgroup_load("tg_max", lane), neg_infinity()); + let g_max = simd_max(g_max_in); + let g_sum_in = + select(lane < ns, threadgroup_load("tg_sum", lane) * exp(g_max_in - g_max), 0.0f32); + let g_sum = simd_sum(g_sum_in); + if lane == 0 { + threadgroup_store("tg_max", 0, g_max); + threadgroup_store("tg_sum", 0, g_sum); + } + } + threadgroup_barrier(); + let g_max0 = threadgroup_load("tg_max", 0); + let g_sum0 = threadgroup_load("tg_sum", 0); + let sink = load(sink_logit[q_head]); + let g_max = select(sink > g_max0, sink, g_max0); + let g_sum = g_sum0 * exp(g_max0 - g_max) + exp(sink - g_max); + let rescale = select(g_sum > 0.0f32, exp(run_max - g_max) / g_sum, 0.0f32); + let stride = ns + 1u32; + let idx = lane * stride + sg; + threadgroup_store("tg_out0", idx, o0 * rescale); + threadgroup_store("tg_out1", idx, o1 * rescale); + threadgroup_store("tg_out2", idx, o2 * rescale); + threadgroup_store("tg_out3", idx, o3 * rescale); + threadgroup_barrier(); + if sg == 0 { + let mut so0 = 0.0f32; + let mut so1 = 0.0f32; + let mut so2 = 0.0f32; + let mut so3 = 0.0f32; + for _g in range(0u32, ns, 1u32) { + let ri = lane * stride + _g; + so0 = so0 + threadgroup_load("tg_out0", ri); + so1 = so1 + threadgroup_load("tg_out1", ri); + so2 = so2 + threadgroup_load("tg_out2", ri); + so3 = so3 + threadgroup_load("tg_out3", ri); + } + let out_off = q_off + d0; + store(out[out_off], so0.cast::()); + store(out[out_off + 1u32], so1.cast::()); + store(out[out_off + 2u32], so2.cast::()); + store(out[out_off + 3u32], so3.cast::()); + } + threadgroup_barrier(); + threadgroup_store("tg_out0", idx, o4 * rescale); + threadgroup_store("tg_out1", idx, o5 * rescale); + threadgroup_store("tg_out2", idx, o6 * rescale); + threadgroup_store("tg_out3", idx, o7 * rescale); + threadgroup_barrier(); + if sg == 0 { + let mut so4 = 0.0f32; + let mut so5 = 0.0f32; + let mut so6 = 0.0f32; + let mut so7 = 0.0f32; + for _g in range(0u32, ns, 1u32) { + let ri = lane * stride + _g; + so4 = so4 + threadgroup_load("tg_out0", ri); + so5 = so5 + threadgroup_load("tg_out1", ri); + so6 = so6 + threadgroup_load("tg_out2", ri); + so7 = so7 + threadgroup_load("tg_out3", ri); + } + let out_off = q_off + d0; + store(out[out_off + 4u32], so4.cast::()); + store(out[out_off + 5u32], so5.cast::()); + store(out[out_off + 6u32], so6.cast::()); + store(out[out_off + 7u32], so7.cast::()); + } + threadgroup_barrier(); + threadgroup_store("tg_out0", idx, o8 * rescale); + threadgroup_store("tg_out1", idx, o9 * rescale); + threadgroup_store("tg_out2", idx, o10 * rescale); + threadgroup_store("tg_out3", idx, o11 * rescale); + threadgroup_barrier(); + if sg == 0 { + let mut so8 = 0.0f32; + let mut so9 = 0.0f32; + let mut so10 = 0.0f32; + let mut so11 = 0.0f32; + for _g in range(0u32, ns, 1u32) { + let ri = lane * stride + _g; + so8 = so8 + threadgroup_load("tg_out0", ri); + so9 = so9 + threadgroup_load("tg_out1", ri); + so10 = so10 + threadgroup_load("tg_out2", ri); + so11 = so11 + threadgroup_load("tg_out3", ri); + } + let out_off = q_off + d0; + store(out[out_off + 8u32], so8.cast::()); + store(out[out_off + 9u32], so9.cast::()); + store(out[out_off + 10u32], so10.cast::()); + store(out[out_off + 11u32], so11.cast::()); + } + threadgroup_barrier(); + threadgroup_store("tg_out0", idx, o12 * rescale); + threadgroup_store("tg_out1", idx, o13 * rescale); + threadgroup_store("tg_out2", idx, o14 * rescale); + threadgroup_store("tg_out3", idx, o15 * rescale); + threadgroup_barrier(); + if sg == 0 { + let mut so12 = 0.0f32; + let mut so13 = 0.0f32; + let mut so14 = 0.0f32; + let mut so15 = 0.0f32; + for _g in range(0u32, ns, 1u32) { + let ri = lane * stride + _g; + so12 = so12 + threadgroup_load("tg_out0", ri); + so13 = so13 + threadgroup_load("tg_out1", ri); + so14 = so14 + threadgroup_load("tg_out2", ri); + so15 = so15 + threadgroup_load("tg_out3", ri); + } + let out_off = q_off + d0; + store(out[out_off + 12u32], so12.cast::()); + store(out[out_off + 13u32], so13.cast::()); + store(out[out_off + 14u32], so14.cast::()); + store(out[out_off + 15u32], so15.cast::()); + } +} + +pub mod kernel_benches { + use metaltile::{bench, test::*}; + + use super::ffai_sdpa_prefill_d512_sink; + + #[bench(name = "ffai/sdpa/sdpa_prefill_d512_sink", dtypes = [f32, f16, bf16])] + fn bench_sdpa_prefill_d512_sink(dt: DType) -> BenchSetup { + let hd = 512usize; + let n_q = 64usize; + let n_query = 256usize; + let kv = 256usize; + BenchSetup::new(ffai_sdpa_prefill_d512_sink::kernel_ir_for(dt)) + .mode(KernelMode::Reduction) + .buffer(BenchBuffer::random("q", n_query * n_q * hd, dt)) + .buffer(BenchBuffer::random("k", kv * hd, dt)) + .buffer(BenchBuffer::random("v", kv * hd, dt)) + .buffer(BenchBuffer::random("sink_logit", n_q, DType::F32)) + .buffer(BenchBuffer::zeros("out", n_query * n_q * hd, dt).output()) + .constexpr("head_dim", hd as u32) + .constexpr("n_q_heads", n_q as u32) + .constexpr("kv_stride", kv as u32) + .constexpr("heads_per_group", n_q as u32) + .constexpr("window", 128u32) + .constexpr("kv_base", 0u32) + .constexpr("scale", 0.044194174f32) + .grid_3d(n_q as u32, n_query as u32, 1, [32, 1, 1]) + .bytes_moved(((kv * hd * 2 + n_query * n_q * hd) * dt.size_bytes()) as u64) + } +} diff --git a/crates/metaltile-std/tests/bm64_vs_gemvrows_iq2xxs.rs b/crates/metaltile-std/tests/bm64_vs_gemvrows_iq2xxs.rs new file mode 100644 index 00000000..7c36f6b9 --- /dev/null +++ b/crates/metaltile-std/tests/bm64_vs_gemvrows_iq2xxs.rs @@ -0,0 +1,113 @@ +//! Copyright 2026 TheTom +//! SPDX-License-Identifier: Apache-2.0 +//! Direct comparison: ffai_moe_bgemm_iq2xxs_bm64 vs ffai_moe_gemv_rows_iq2xxs +//! on IDENTICAL pool/x/indices. Both claim to compute gateP[row,m] = +//! W[expert(row),m,:]·x[row,:]; the prefill shows them disagreeing. This +//! reproduces it in isolation. cosine should be ~1.0. +#![cfg(target_os = "macos")] + +mod common; + +use std::collections::BTreeMap; + +use common::{Dt, gpu_lock, pack_bytes, pack_u32_bytes, unpack_bytes}; +use metaltile_core::ir::KernelMode; +use metaltile_runtime::Context; +use metaltile_std::ffai::{ + moe_bgemm_iq2xxs_bm64::ffai_moe_bgemm_iq2xxs_bm64, + moe_gemv_rows_iq2xxs::ffai_moe_gemv_rows_iq2xxs, +}; + +fn xorshift(s: &mut u32) -> u32 { + let mut x = *s; + x ^= x << 13; + x ^= x >> 17; + x ^= x << 5; + *s = x; + x +} + +#[test] +fn bm64_matches_gemvrows() { + let _g = gpu_lock(); + let probe = Context::new().expect("Context::new"); + if probe.chip_family().is_none_or(|lvl| lvl < 10) { + eprintln!("skip"); + return; + } + drop(probe); + + let n_experts = 64usize; + let k_in = 4096usize; + let n_out = 2048usize; + let t_rows = 30usize; + let nblk = n_out * k_in / 256; + // Prefill N=5 case: M=30 rows, each a DISTINCT expert (sorted), in one + // partial 64-tile — many size-1 sub-runs + sentinel padding to 64. + let indices: Vec = (0..t_rows as u32).collect(); + let mut st = 0xA11CE5u32; + let qs: Vec = (0..n_experts * nblk * 16).map(|_| xorshift(&mut st)).collect(); + let d: Vec = + (0..n_experts * nblk).map(|_| (xorshift(&mut st) % 1000) as f32 * 0.0005 + 0.001).collect(); + let grid: Vec = (0..2048).map(|i| ((i * 7) % 48) as u8).collect(); + let signs: Vec = (0..128).map(|i| (i * 2) as u8).collect(); + let x: Vec = + (0..t_rows * k_in).map(|_| ((xorshift(&mut st) % 2000) as f32 / 1000.0) - 1.0).collect(); + + let ctx = Context::new().unwrap(); + + // bm64 + let mut b: BTreeMap> = BTreeMap::new(); + b.insert("x".into(), pack_bytes(&x, Dt::F32)); + b.insert("qs".into(), pack_u32_bytes(&qs)); + b.insert("d_f32".into(), pack_bytes(&d, Dt::F32)); + b.insert("grid".into(), grid.clone()); + b.insert("signs".into(), signs.clone()); + b.insert("indices".into(), pack_u32_bytes(&indices)); + b.insert("out".into(), pack_bytes(&vec![0.0f32; t_rows * n_out], Dt::F32)); + b.insert("m_total".into(), (t_rows as u32).to_le_bytes().to_vec()); + b.insert("n_out".into(), (n_out as u32).to_le_bytes().to_vec()); + b.insert("k_in".into(), (k_in as u32).to_le_bytes().to_vec()); + let mut kb = ffai_moe_bgemm_iq2xxs_bm64::kernel_ir_for(Dt::F32.to_dtype()); + kb.mode = KernelMode::Reduction; + let rb = ctx + .dispatch_with_grid(&kb, &b, &BTreeMap::new(), [n_out / 64, t_rows.div_ceil(64), 1], [ + 128, 1, 1, + ]) + .unwrap(); + let got_bm = unpack_bytes(rb.outputs.get("out").unwrap(), Dt::F32); + + // gemv-rows (qs_all/d_all naming, m_out/m_total) + let mut g: BTreeMap> = BTreeMap::new(); + g.insert("x".into(), pack_bytes(&x, Dt::F32)); + g.insert("qs_all".into(), pack_u32_bytes(&qs)); + g.insert("d_all".into(), pack_bytes(&d, Dt::F32)); + g.insert("expert_ids".into(), pack_u32_bytes(&indices)); + g.insert("grid".into(), grid.clone()); + g.insert("signs".into(), signs.clone()); + g.insert("out".into(), pack_bytes(&vec![0.0f32; t_rows * n_out], Dt::F32)); + g.insert("k_in".into(), (k_in as u32).to_le_bytes().to_vec()); + g.insert("m_out".into(), (n_out as u32).to_le_bytes().to_vec()); + g.insert("m_total".into(), (t_rows as u32).to_le_bytes().to_vec()); + let mut kg = ffai_moe_gemv_rows_iq2xxs::kernel_ir_for(Dt::F32.to_dtype()); + kg.mode = KernelMode::Reduction; + let rg = + ctx.dispatch_with_grid(&kg, &g, &BTreeMap::new(), [n_out, t_rows, 1], [32, 1, 1]).unwrap(); + let got_gv = unpack_bytes(rg.outputs.get("out").unwrap(), Dt::F32); + + let mut dot = 0.0f64; + let mut na = 0.0f64; + let mut nb = 0.0f64; + let mut maxd = 0.0f32; + for (a, bb) in got_gv.iter().zip(&got_bm) { + dot += (*a as f64) * (*bb as f64); + na += (*a as f64).powi(2); + nb += (*bb as f64).powi(2); + maxd = maxd.max((a - bb).abs()); + } + let cos = dot / (na.sqrt() * nb.sqrt() + 1e-12); + eprintln!("gemvrows[0..6]={:?}", &got_gv[..6]); + eprintln!("bm64[0..6] ={:?}", &got_bm[..6]); + eprintln!("cos={cos:.6} maxAbsDiff={maxd:.4}"); + assert!(cos >= 0.999, "bm64 vs gemv-rows DISAGREE: cosine {cos:.6}"); +} diff --git a/crates/metaltile-std/tests/dsv4_partial_rope_rows_correctness.rs b/crates/metaltile-std/tests/dsv4_partial_rope_rows_correctness.rs new file mode 100644 index 00000000..8d4965e5 --- /dev/null +++ b/crates/metaltile-std/tests/dsv4_partial_rope_rows_correctness.rs @@ -0,0 +1,93 @@ +//! Copyright 2026 TheTom +//! SPDX-License-Identifier: Apache-2.0 +//! Batched partial RoPE (ffai_dsv4_partial_rope_rows) must equal the per-token +//! reference (token t roped at position t). NO model load. +#![cfg(target_os = "macos")] + +mod common; + +use std::collections::BTreeMap; + +use common::{Dt, gpu_lock, pack_bytes, unpack_bytes}; +use metaltile_core::ir::KernelMode; +use metaltile_runtime::Context; +use metaltile_std::ffai::dsv4_partial_rope_rows::ffai_dsv4_partial_rope_rows; + +#[test] +fn partial_rope_rows_matches_per_token_oracle() { + let _g = gpu_lock(); + let probe = Context::new().expect("Context::new"); + if probe.chip_family().is_none_or(|lvl| lvl < 10) { + eprintln!("skip"); + return; + } + drop(probe); + + let n_heads = 8usize; + let head_dim = 512usize; + let n_nope = 448usize; + let half_rot = 32usize; + let n_tokens = 40usize; + let theta_base = 10_000.0f32; + + let qk: Vec = (0..n_tokens * n_heads * head_dim) + .map(|i| ((i as f32) * 0.011 - 0.4).sin() * 1.2) + .collect(); + + // Oracle: token t roped at position t (adjacent GPT-J pairs, forward). + let mut want = qk.clone(); + for t in 0..n_tokens { + for h in 0..n_heads { + for p in 0..half_rot { + let inv_freq = + (-(p as f32) * 2.0 * theta_base.ln() / (2.0 * half_rot as f32)).exp(); + let theta = t as f32 * inv_freq; + let c = theta.cos(); + let s = theta.sin(); + let base = t * n_heads * head_dim + h * head_dim + n_nope + 2 * p; + let x_lo = qk[base]; + let x_hi = qk[base + 1]; + want[base] = x_lo * c - x_hi * s; + want[base + 1] = x_lo * s + x_hi * c; + } + } + } + + let mut buffers: BTreeMap> = BTreeMap::new(); + buffers.insert("qk".into(), pack_bytes(&qk, Dt::F32)); + buffers.insert("out".into(), pack_bytes(&qk, Dt::F32)); // pass-through nope dims + buffers.insert("head_dim".into(), (head_dim as u32).to_le_bytes().to_vec()); + buffers.insert("n_nope".into(), (n_nope as u32).to_le_bytes().to_vec()); + buffers.insert("half_rot".into(), (half_rot as u32).to_le_bytes().to_vec()); + buffers.insert("n_heads".into(), (n_heads as u32).to_le_bytes().to_vec()); + buffers.insert("base_position".into(), 0u32.to_le_bytes().to_vec()); + buffers.insert("theta_base".into(), theta_base.to_le_bytes().to_vec()); + buffers.insert("inverse_flag".into(), 0u32.to_le_bytes().to_vec()); + buffers.insert("freq_scale".into(), 1.0f32.to_le_bytes().to_vec()); + buffers.insert("ext_factor".into(), 0.0f32.to_le_bytes().to_vec()); + buffers.insert("corr_low".into(), 0.0f32.to_le_bytes().to_vec()); + buffers.insert("corr_high".into(), 0.0f32.to_le_bytes().to_vec()); + + let ctx = Context::new().unwrap(); + let mut k = ffai_dsv4_partial_rope_rows::kernel_ir_for(Dt::F32.to_dtype()); + k.mode = KernelMode::Grid3D; + let r = ctx + .dispatch_with_grid(&k, &buffers, &BTreeMap::new(), [n_heads, half_rot, n_tokens], [ + 1, 1, 1, + ]) + .unwrap(); + let got = unpack_bytes(r.outputs.get("out").unwrap(), Dt::F32); + + let mut maxd = 0.0f32; + let mut nan = 0; + for (a, b) in want.iter().zip(&got) { + if !b.is_finite() { + nan += 1; + continue; + } + maxd = maxd.max((a - b).abs()); + } + eprintln!("nan={nan} maxAbsDiff={maxd:.6}"); + assert_eq!(nan, 0); + assert!(maxd < 1e-3, "maxAbsDiff {maxd} too large — batched rope diverges from per-token"); +} diff --git a/crates/metaltile-std/tests/dsv4_router_topk_correctness.rs b/crates/metaltile-std/tests/dsv4_router_topk_correctness.rs new file mode 100644 index 00000000..f0e24c59 --- /dev/null +++ b/crates/metaltile-std/tests/dsv4_router_topk_correctness.rs @@ -0,0 +1,83 @@ +//! Copyright 2026 TheTom +//! SPDX-License-Identifier: Apache-2.0 +//! GPU correctness for `ffai::mt_dsv4_router_topk` — top-K by biased +//! score, weights = unbiased[chosen] renormalised to sum 1. +#![cfg(target_os = "macos")] + +mod common; + +use std::collections::BTreeMap; + +use common::{Dt, gpu_lock, pack_bytes, unpack_bytes, unpack_u32_bytes}; +use metaltile_core::ir::KernelMode; +use metaltile_runtime::Context; +use metaltile_std::ffai::dsv4_router_topk::{mt_dsv4_router_topk, mt_remap_u32}; + +#[test] +fn dsv4_router_topk_f32() { + let _g = gpu_lock(); + let n = 256usize; + let k = 6usize; + // Deterministic distinct biased scores so the top-K is unambiguous. + let biased: Vec = (0..n).map(|i| ((i * 37 + 11) % 251) as f32 * 0.1).collect(); + let unbiased: Vec = (0..n).map(|i| ((i * 53 + 7) % 199) as f32 * 0.01 + 0.05).collect(); + + // CPU reference. + let mut order: Vec = (0..n).collect(); + order.sort_by(|&a, &b| biased[b].partial_cmp(&biased[a]).unwrap()); + let chosen: Vec = order.iter().take(k).copied().collect(); + let wsum: f32 = chosen.iter().map(|&e| unbiased[e]).sum(); + let want_w: Vec = chosen.iter().map(|&e| unbiased[e] / wsum).collect(); + + let mut buffers: BTreeMap> = BTreeMap::new(); + buffers.insert("score_biased".into(), pack_bytes(&biased, Dt::F32)); + buffers.insert("score_unbiased".into(), pack_bytes(&unbiased, Dt::F32)); + buffers.insert("indices_out".into(), vec![0u8; k * 4]); + buffers.insert("weights_out".into(), pack_bytes(&vec![0.0f32; k], Dt::F32)); + buffers.insert("n_experts".into(), (n as u32).to_le_bytes().to_vec()); + buffers.insert("k".into(), (k as u32).to_le_bytes().to_vec()); + + let ctx = Context::new().expect("ctx"); + let mut kernel = mt_dsv4_router_topk::kernel_ir_for(Dt::F32.to_dtype()); + kernel.mode = KernelMode::Reduction; + let result = ctx + .dispatch_with_grid(&kernel, &buffers, &BTreeMap::new(), [1, 1, 1], [32, 1, 1]) + .expect("dispatch"); + let got_idx = unpack_u32_bytes(result.outputs.get("indices_out").expect("idx")); + let got_w = unpack_bytes(result.outputs.get("weights_out").expect("w"), Dt::F32); + + assert_eq!(got_idx.iter().map(|&x| x as usize).collect::>(), chosen, "indices"); + for (i, (g, w)) in got_w.iter().zip(want_w.iter()).enumerate() { + assert!((g - w).abs() < 1e-4, "weight {i}: got={g} want={w}"); + } +} + +/// `mt_remap_u32`: out[i] = table[idx[i]] — a plain u32 gather over `n` +/// elements. Non-generic Grid3D kernel (one thread per output), so it +/// dispatches via `kernel_ir()` and a [n,1,1] grid with [1,1,1] tg. +#[test] +fn remap_u32_matches_cpu() { + let _g = gpu_lock(); + let n = 6usize; + let table_len = 256usize; + let table: Vec = (0..table_len).map(|e| ((e * 37 + 11) % table_len) as u32).collect(); + let idx: Vec = vec![3, 200, 0, 255, 128, 64]; + let want: Vec = idx.iter().map(|&e| table[e as usize]).collect(); + + let to_bytes = |v: &[u32]| -> Vec { v.iter().flat_map(|x| x.to_le_bytes()).collect() }; + let mut buffers: BTreeMap> = BTreeMap::new(); + buffers.insert("table".into(), to_bytes(&table)); + buffers.insert("idx".into(), to_bytes(&idx)); + buffers.insert("out".into(), vec![0u8; n * 4]); + buffers.insert("n".into(), (n as u32).to_le_bytes().to_vec()); + + let ctx = Context::new().expect("ctx"); + let kernel = mt_remap_u32::kernel_ir(); + let result = ctx + .dispatch_with_grid(&kernel, &buffers, &BTreeMap::new(), [n, 1, 1], [1, 1, 1]) + .expect("dispatch"); + let got = unpack_u32_bytes(result.outputs.get("out").expect("out")); + + eprintln!("want={want:?} got={got:?}"); + assert_eq!(got, want, "remap_u32 gather mismatch"); +} diff --git a/crates/metaltile-std/tests/gemm_q8_correctness.rs b/crates/metaltile-std/tests/gemm_q8_correctness.rs new file mode 100644 index 00000000..713e01b1 --- /dev/null +++ b/crates/metaltile-std/tests/gemm_q8_correctness.rs @@ -0,0 +1,99 @@ +//! Copyright 2026 TheTom +//! SPDX-License-Identifier: Apache-2.0 +//! GPU correctness for `ffai::ffai_gemm_q8` — the multi-row Q8_0 tiled GEMM +//! used by DSv4 prefill. Oracle: triple-loop with the same Q8 dequant. +#![cfg(target_os = "macos")] + +mod common; + +use std::collections::BTreeMap; + +use common::{Dt, gpu_lock, pack_bytes, pack_u32_bytes, unpack_bytes}; +use metaltile_core::ir::KernelMode; +use metaltile_runtime::Context; +use metaltile_std::ffai::gemm_q8::ffai_gemm_q8; + +fn xorshift(s: &mut u32) -> u32 { + let mut x = *s; + x ^= x << 13; + x ^= x >> 17; + x ^= x << 5; + *s = x; + x +} + +fn run_case(dt: Dt, in_dim: usize, out_dim: usize, n_rows: usize, tol: f32) { + let _g = gpu_lock(); + let n_blocks = out_dim * in_dim / 32; + let mut st = 0x515E_2026u32; + let qs: Vec = (0..n_blocks * 8).map(|_| xorshift(&mut st)).collect(); + let d: Vec = + (0..n_blocks).map(|_| (xorshift(&mut st) % 1000) as f32 * 0.0002 + 0.001).collect(); + let mut input: Vec = + (0..n_rows * in_dim).map(|_| ((xorshift(&mut st) % 2000) as f32 / 1000.0) - 1.0).collect(); + // Round inputs to the kernel's dtype so the f32 oracle sees the same + // values the kernel loads (f16 input rounding, amplified by cancellation + // over in_dim terms, otherwise dominates the diff). Matches ffai_gemm's test. + if matches!(dt, Dt::F16) { + input = unpack_bytes(&pack_bytes(&input, Dt::F16), Dt::F16); + } + + // Dequant weight to a dense [out_dim, in_dim] f32 for the oracle. + let deq = |vidx: usize| -> f32 { + let block = vidx / 32; + let lane = vidx % 32; + let word = qs[block * 8 + lane / 4]; + let by = ((word >> ((lane % 4) * 8)) & 0xff) as i32; + let q = if by > 127 { by - 256 } else { by }; + d[block] * q as f32 + }; + let mut want = vec![0.0f32; n_rows * out_dim]; + for r in 0..n_rows { + for o in 0..out_dim { + let mut acc = 0.0f32; + for k in 0..in_dim { + acc += deq(o * in_dim + k) * input[r * in_dim + k]; + } + want[r * out_dim + o] = acc; + } + } + + let mut buffers: BTreeMap> = BTreeMap::new(); + buffers.insert("qs".into(), pack_u32_bytes(&qs)); + buffers.insert("d_f32".into(), pack_bytes(&d, Dt::F32)); + buffers.insert("input".into(), pack_bytes(&input, dt)); + buffers.insert("out".into(), pack_bytes(&vec![0.0f32; n_rows * out_dim], dt)); + buffers.insert("in_dim".into(), (in_dim as u32).to_le_bytes().to_vec()); + buffers.insert("out_dim".into(), (out_dim as u32).to_le_bytes().to_vec()); + buffers.insert("n_rows".into(), (n_rows as u32).to_le_bytes().to_vec()); + + let ctx = Context::new().expect("ctx"); + let mut kernel = ffai_gemm_q8::kernel_ir_for(dt.to_dtype()); + kernel.mode = KernelMode::Reduction; + let gx = (out_dim as u32).div_ceil(32); + let gy = (n_rows as u32).div_ceil(32); + let result = ctx + .dispatch_with_grid(&kernel, &buffers, &BTreeMap::new(), [gx as usize, gy as usize, 1], [ + 1024, 1, 1, + ]) + .expect("dispatch"); + let got = unpack_bytes(result.outputs.get("out").expect("out"), dt); + + for i in 0..(n_rows * out_dim) { + let w = want[i]; + let denom = w.abs().max(1.0); + assert!((got[i] - w).abs() / denom < tol, "i={i}: got={} want={}", got[i], w); + } +} + +#[test] +fn gemm_q8_f32() { + run_case(Dt::F32, 256, 64, 32, 1e-3); // aligned + run_case(Dt::F32, 512, 96, 40, 1e-3); // edge: out_dim/n_rows not mult of 32 +} + +#[test] +fn gemm_q8_f16() { + run_case(Dt::F16, 256, 64, 32, 2e-2); + run_case(Dt::F16, 512, 96, 40, 2e-2); +} diff --git a/crates/metaltile-std/tests/gemm_q8_mpp_correctness.rs b/crates/metaltile-std/tests/gemm_q8_mpp_correctness.rs new file mode 100644 index 00000000..4f28844a --- /dev/null +++ b/crates/metaltile-std/tests/gemm_q8_mpp_correctness.rs @@ -0,0 +1,172 @@ +//! Copyright 2026 TheTom +//! SPDX-License-Identifier: Apache-2.0 +//! GPU correctness for `ffai::ffai_gemm_q8_mpp` and `ffai_grouped_gemm_q8_mpp` +//! — the cooperative-tensor MMA Q8_0 GEMMs (dense + grouped). Oracle: triple +//! loop with the SAME Q8 dequant as `gemm_q8_correctness`; the MMA kernels +//! differ only in the matmul structure (64×64×32 coop_tile vs scalar). Both +//! live-compile (name has `_mpp_`). out[r,o] = Σ_k dequant(W[o,k]) · x[r,k]. +#![cfg(target_os = "macos")] + +mod common; + +use std::collections::BTreeMap; + +use common::{Dt, gpu_lock, pack_bytes, pack_u32_bytes, unpack_bytes}; +use metaltile_core::ir::KernelMode; +use metaltile_runtime::Context; +use metaltile_std::ffai::gemm_q8_mpp::{ffai_gemm_q8_mpp, ffai_grouped_gemm_q8_mpp}; + +fn xorshift(s: &mut u32) -> u32 { + let mut x = *s; + x ^= x << 13; + x ^= x >> 17; + x ^= x << 5; + *s = x; + x +} + +/// Q8 weight dequant: weight is [out_dim, k_in] row-major over values. +fn deq(qs: &[u32], d: &[f32], vidx: usize) -> f32 { + let block = vidx / 32; + let lane = vidx % 32; + let word = qs[block * 8 + lane / 4]; + let by = ((word >> ((lane % 4) * 8)) & 0xff) as i32; + let q = if by > 127 { by - 256 } else { by }; + d[block] * q as f32 +} + +fn need_apple10() -> bool { + let probe = Context::new().expect("Context::new"); + let ok = probe.chip_family().is_some_and(|lvl| lvl >= 10); + if !ok { + eprintln!("skip: needs Apple10+ GPU for coop_tile MMA"); + } + ok +} + +#[test] +fn gemm_q8_mpp_matches_oracle() { + let _g = gpu_lock(); + if !need_apple10() { + return; + } + // Aligned to the 64×64×32 tile. + let n_rows = 128usize; + let out_dim = 128usize; + let k_in = 128usize; + let n_blocks = out_dim * k_in / 32; + + let mut st = 0x4A3D_2026u32; + let qs: Vec = (0..n_blocks * 8).map(|_| xorshift(&mut st)).collect(); + let d: Vec = + (0..n_blocks).map(|_| (xorshift(&mut st) % 1000) as f32 * 0.0002 + 0.001).collect(); + let x: Vec = + (0..n_rows * k_in).map(|_| ((xorshift(&mut st) % 2000) as f32 / 1000.0) - 1.0).collect(); + + let mut want = vec![0.0f32; n_rows * out_dim]; + for r in 0..n_rows { + for o in 0..out_dim { + let mut acc = 0.0f32; + for k in 0..k_in { + acc += deq(&qs, &d, o * k_in + k) * x[r * k_in + k]; + } + want[r * out_dim + o] = acc; + } + } + + let mut buffers: BTreeMap> = BTreeMap::new(); + buffers.insert("x".into(), pack_bytes(&x, Dt::F32)); + buffers.insert("qs".into(), pack_u32_bytes(&qs)); + buffers.insert("d_f32".into(), pack_bytes(&d, Dt::F32)); + buffers.insert("out".into(), pack_bytes(&vec![0.0f32; n_rows * out_dim], Dt::F32)); + buffers.insert("n_rows".into(), (n_rows as u32).to_le_bytes().to_vec()); + buffers.insert("out_dim".into(), (out_dim as u32).to_le_bytes().to_vec()); + buffers.insert("k_in".into(), (k_in as u32).to_le_bytes().to_vec()); + + let ctx = Context::new().unwrap(); + let mut kernel = ffai_gemm_q8_mpp::kernel_ir_for(Dt::F32.to_dtype()); + kernel.mode = KernelMode::Reduction; + let gx = (out_dim as u32).div_ceil(64) as usize; + let gy = (n_rows as u32).div_ceil(64) as usize; + let r = ctx + .dispatch_with_grid(&kernel, &buffers, &BTreeMap::new(), [gx, gy, 1], [128, 1, 1]) + .unwrap(); + let got = unpack_bytes(r.outputs.get("out").unwrap(), Dt::F32); + + let mut maxd = 0.0f32; + for i in 0..(n_rows * out_dim) { + assert!(got[i].is_finite(), "non-finite at {i}"); + let denom = want[i].abs().max(1.0); + maxd = maxd.max((got[i] - want[i]).abs() / denom); + } + eprintln!("gemm_q8_mpp maxRelDiff={maxd:.6}"); + assert!(maxd < 1e-3, "gemm_q8_mpp diverges from oracle: maxRelDiff={maxd}"); +} + +#[test] +fn grouped_gemm_q8_mpp_matches_oracle() { + let _g = gpu_lock(); + if !need_apple10() { + return; + } + // 2 groups, rows_per_group=64 (64-wide feature tile is uniform-group). + let n_rows = 64usize; + let k_in = 64usize; + let n_groups = 2usize; + let rows_per_group = 64usize; + let out_dim = n_groups * rows_per_group; // 128 + let n_blocks = out_dim * k_in / 32; + + let mut st = 0x5B7C_2026u32; + let qs: Vec = (0..n_blocks * 8).map(|_| xorshift(&mut st)).collect(); + let d: Vec = + (0..n_blocks).map(|_| (xorshift(&mut st) % 1000) as f32 * 0.0002 + 0.001).collect(); + // Activation row is n_groups*k_in wide. + let x: Vec = (0..n_rows * n_groups * k_in) + .map(|_| ((xorshift(&mut st) % 2000) as f32 / 1000.0) - 1.0) + .collect(); + + let row_in_stride = n_groups * k_in; + let mut want = vec![0.0f32; n_rows * out_dim]; + for r in 0..n_rows { + for o in 0..out_dim { + let g = o / rows_per_group; + let in_col_off = g * k_in; + let mut acc = 0.0f32; + for k in 0..k_in { + acc += deq(&qs, &d, o * k_in + k) * x[r * row_in_stride + in_col_off + k]; + } + want[r * out_dim + o] = acc; + } + } + + let mut buffers: BTreeMap> = BTreeMap::new(); + buffers.insert("x".into(), pack_bytes(&x, Dt::F32)); + buffers.insert("qs".into(), pack_u32_bytes(&qs)); + buffers.insert("d_f32".into(), pack_bytes(&d, Dt::F32)); + buffers.insert("out".into(), pack_bytes(&vec![0.0f32; n_rows * out_dim], Dt::F32)); + buffers.insert("n_rows".into(), (n_rows as u32).to_le_bytes().to_vec()); + buffers.insert("out_dim".into(), (out_dim as u32).to_le_bytes().to_vec()); + buffers.insert("k_in".into(), (k_in as u32).to_le_bytes().to_vec()); + buffers.insert("n_groups".into(), (n_groups as u32).to_le_bytes().to_vec()); + buffers.insert("rows_per_group".into(), (rows_per_group as u32).to_le_bytes().to_vec()); + + let ctx = Context::new().unwrap(); + let mut kernel = ffai_grouped_gemm_q8_mpp::kernel_ir_for(Dt::F32.to_dtype()); + kernel.mode = KernelMode::Reduction; + let gx = (out_dim as u32).div_ceil(64) as usize; + let gy = (n_rows as u32).div_ceil(64) as usize; + let r = ctx + .dispatch_with_grid(&kernel, &buffers, &BTreeMap::new(), [gx, gy, 1], [128, 1, 1]) + .unwrap(); + let got = unpack_bytes(r.outputs.get("out").unwrap(), Dt::F32); + + let mut maxd = 0.0f32; + for i in 0..(n_rows * out_dim) { + assert!(got[i].is_finite(), "non-finite at {i}"); + let denom = want[i].abs().max(1.0); + maxd = maxd.max((got[i] - want[i]).abs() / denom); + } + eprintln!("grouped_gemm_q8_mpp maxRelDiff={maxd:.6}"); + assert!(maxd < 1e-3, "grouped_gemm_q8_mpp diverges from oracle: maxRelDiff={maxd}"); +} diff --git a/crates/metaltile-std/tests/gemv_q8_correctness.rs b/crates/metaltile-std/tests/gemv_q8_correctness.rs new file mode 100644 index 00000000..188c8604 --- /dev/null +++ b/crates/metaltile-std/tests/gemv_q8_correctness.rs @@ -0,0 +1,138 @@ +//! Copyright 2026 TheTom +//! SPDX-License-Identifier: Apache-2.0 +//! GPU correctness for `ffai::ffai_gemv_q8` — Q8_0 inline-dequant gemv +//! vs a CPU reference using the same dequant (value = d * int8). +#![cfg(target_os = "macos")] + +mod common; + +use std::collections::BTreeMap; + +use common::{Dt, gpu_lock, pack_bytes, pack_u32_bytes, unpack_bytes}; +use metaltile_core::ir::KernelMode; +use metaltile_runtime::Context; +use metaltile_std::ffai::gemv_q8::ffai_gemv_q8; + +fn xorshift(s: &mut u32) -> u32 { + let mut x = *s; + x ^= x << 13; + x ^= x >> 17; + x ^= x << 5; + *s = x; + x +} + +fn run_case(dt: Dt, k_in: usize, m_out: usize, tol: f32) { + let _g = gpu_lock(); + let bpr = k_in / 32; + let mut st = 0xC0FFEEu32; + // qs: m_out*bpr*8 u32 (random int8 payloads). d: per-block scale. x: input. + let qs: Vec = (0..m_out * bpr * 8).map(|_| xorshift(&mut st)).collect(); + let d: Vec = + (0..m_out * bpr).map(|_| (xorshift(&mut st) % 1000) as f32 * 0.0001 + 0.001).collect(); + let x: Vec = + (0..k_in).map(|_| ((xorshift(&mut st) % 2000) as f32 / 1000.0) - 1.0).collect(); + + // CPU reference. + let mut want = vec![0.0f32; m_out]; + for r in 0..m_out { + let mut acc = 0.0f32; + for b in 0..bpr { + let dv = d[r * bpr + b]; + for w in 0..8 { + let packed = qs[r * bpr * 8 + b * 8 + w]; + for i in 0..4 { + let by = ((packed >> (i * 8)) & 0xff) as i32; + let q = if by > 127 { by - 256 } else { by }; + acc += dv * q as f32 * x[b * 32 + w * 4 + i]; + } + } + } + want[r] = acc; + } + + let mut buffers: BTreeMap> = BTreeMap::new(); + buffers.insert("qs".into(), pack_u32_bytes(&qs)); + buffers.insert("d_f32".into(), pack_bytes(&d, Dt::F32)); + buffers.insert("x".into(), pack_bytes(&x, dt)); + buffers.insert("out".into(), pack_bytes(&vec![0.0f32; m_out], dt)); + buffers.insert("k_in".into(), (k_in as u32).to_le_bytes().to_vec()); + buffers.insert("m_out".into(), (m_out as u32).to_le_bytes().to_vec()); + + let ctx = Context::new().expect("ctx"); + let mut kernel = ffai_gemv_q8::kernel_ir_for(dt.to_dtype()); + kernel.mode = KernelMode::Reduction; + let result = ctx + .dispatch_with_grid(&kernel, &buffers, &BTreeMap::new(), [m_out, 1, 1], [32, 1, 1]) + .expect("dispatch"); + let got = unpack_bytes(result.outputs.get("out").expect("out"), dt); + + for (r, (g, w)) in got.iter().zip(want.iter()).enumerate() { + let denom = w.abs().max(1.0); + assert!((g - w).abs() / denom < tol, "dt={dt:?} r={r}: got={g} want={w}"); + } +} + +#[test] +fn gemv_q8_f32() { + run_case(Dt::F32, 1024, 64, 1e-3); + run_case(Dt::F32, 8192, 32, 1e-3); +} + +#[test] +fn gemv_q8_f16() { + run_case(Dt::F16, 1024, 64, 3e-2); + run_case(Dt::F16, 8192, 32, 4e-2); +} + +#[test] +fn grouped_gemv_q8_f32() { + use metaltile_std::ffai::gemv_q8::ffai_grouped_gemv_q8; + let _g = gpu_lock(); + let k_in = 4096usize; + let rows_per_group = 1024usize; + let n_groups = 8usize; + let m_out = rows_per_group * n_groups; + let bpr = k_in / 32; + let mut st = 0xABCDu32; + let qs: Vec = (0..m_out * bpr * 8).map(|_| xorshift(&mut st)).collect(); + let d: Vec = + (0..m_out * bpr).map(|_| (xorshift(&mut st) % 1000) as f32 * 0.0001 + 0.001).collect(); + let x: Vec = + (0..n_groups * k_in).map(|_| ((xorshift(&mut st) % 2000) as f32 / 1000.0) - 1.0).collect(); + let mut want = vec![0.0f32; m_out]; + for r in 0..m_out { + let xb = (r / rows_per_group) * k_in; + let mut acc = 0.0f32; + for b in 0..bpr { + let dv = d[r * bpr + b]; + for w in 0..8 { + let packed = qs[r * bpr * 8 + b * 8 + w]; + for i in 0..4 { + let by = ((packed >> (i * 8)) & 0xff) as i32; + let q = if by > 127 { by - 256 } else { by }; + acc += dv * q as f32 * x[xb + b * 32 + w * 4 + i]; + } + } + } + want[r] = acc; + } + let mut buffers: BTreeMap> = BTreeMap::new(); + buffers.insert("qs".into(), pack_u32_bytes(&qs)); + buffers.insert("d_f32".into(), pack_bytes(&d, Dt::F32)); + buffers.insert("x".into(), pack_bytes(&x, Dt::F32)); + buffers.insert("out".into(), pack_bytes(&vec![0.0f32; m_out], Dt::F32)); + buffers.insert("k_in".into(), (k_in as u32).to_le_bytes().to_vec()); + buffers.insert("m_out".into(), (m_out as u32).to_le_bytes().to_vec()); + buffers.insert("rows_per_group".into(), (rows_per_group as u32).to_le_bytes().to_vec()); + let ctx = Context::new().expect("ctx"); + let mut kernel = ffai_grouped_gemv_q8::kernel_ir_for(Dt::F32.to_dtype()); + kernel.mode = KernelMode::Reduction; + let result = ctx + .dispatch_with_grid(&kernel, &buffers, &BTreeMap::new(), [m_out, 1, 1], [32, 1, 1]) + .expect("dispatch"); + let got = unpack_bytes(result.outputs.get("out").expect("out"), Dt::F32); + for (r, (g, w)) in got.iter().zip(want.iter()).enumerate() { + assert!((g - w).abs() / w.abs().max(1.0) < 1e-3, "r={r}: got={g} want={w}"); + } +} diff --git a/crates/metaltile-std/tests/grouped_gemm_q8_correctness.rs b/crates/metaltile-std/tests/grouped_gemm_q8_correctness.rs new file mode 100644 index 00000000..fb490e48 --- /dev/null +++ b/crates/metaltile-std/tests/grouped_gemm_q8_correctness.rs @@ -0,0 +1,118 @@ +//! Copyright 2026 TheTom +//! SPDX-License-Identifier: Apache-2.0 +//! GPU correctness for `ffai::ffai_grouped_gemm_q8` — the GROUPED multi-row +//! Q8_0 tiled GEMM (O-LoRA-A prefill). Oracle: triple loop with the SAME Q8 +//! dequant as `gemm_q8_correctness`, but output column `o` belongs to group +//! `g = o / rows_per_group` and reads the `g`-th `in_dim`-slice of an +//! `n_groups*in_dim`-wide activation row. +#![cfg(target_os = "macos")] + +mod common; + +use std::collections::BTreeMap; + +use common::{Dt, gpu_lock, pack_bytes, pack_u32_bytes, unpack_bytes}; +use metaltile_core::ir::KernelMode; +use metaltile_runtime::Context; +use metaltile_std::ffai::gemm_q8::ffai_grouped_gemm_q8; + +fn xorshift(s: &mut u32) -> u32 { + let mut x = *s; + x ^= x << 13; + x ^= x >> 17; + x ^= x << 5; + *s = x; + x +} + +#[allow(clippy::too_many_arguments)] +fn run_case( + dt: Dt, + in_dim: usize, + out_dim: usize, + n_rows: usize, + n_groups: usize, + rows_per_group: usize, + tol: f32, +) { + let _g = gpu_lock(); + assert_eq!(out_dim, n_groups * rows_per_group, "out_dim must = n_groups*rows_per_group"); + // A 32-wide output tile must lie entirely within one group. + assert_eq!(rows_per_group % 32, 0, "rows_per_group must be a multiple of 32"); + + let n_blocks = out_dim * in_dim / 32; + let mut st = 0x6C1E_2026u32; + let qs: Vec = (0..n_blocks * 8).map(|_| xorshift(&mut st)).collect(); + let d: Vec = + (0..n_blocks).map(|_| (xorshift(&mut st) % 1000) as f32 * 0.0002 + 0.001).collect(); + // Activation row is n_groups*in_dim wide. + let mut input: Vec = (0..n_rows * n_groups * in_dim) + .map(|_| ((xorshift(&mut st) % 2000) as f32 / 1000.0) - 1.0) + .collect(); + if matches!(dt, Dt::F16) { + input = unpack_bytes(&pack_bytes(&input, Dt::F16), Dt::F16); + } + + // Q8 weight [out_dim, in_dim]. + let deq = |vidx: usize| -> f32 { + let block = vidx / 32; + let lane = vidx % 32; + let word = qs[block * 8 + lane / 4]; + let by = ((word >> ((lane % 4) * 8)) & 0xff) as i32; + let q = if by > 127 { by - 256 } else { by }; + d[block] * q as f32 + }; + let row_in_stride = n_groups * in_dim; + let mut want = vec![0.0f32; n_rows * out_dim]; + for r in 0..n_rows { + for o in 0..out_dim { + let g = o / rows_per_group; + let in_col_off = g * in_dim; + let mut acc = 0.0f32; + for k in 0..in_dim { + acc += deq(o * in_dim + k) * input[r * row_in_stride + in_col_off + k]; + } + want[r * out_dim + o] = acc; + } + } + + let mut buffers: BTreeMap> = BTreeMap::new(); + buffers.insert("qs".into(), pack_u32_bytes(&qs)); + buffers.insert("d_f32".into(), pack_bytes(&d, Dt::F32)); + buffers.insert("input".into(), pack_bytes(&input, dt)); + buffers.insert("out".into(), pack_bytes(&vec![0.0f32; n_rows * out_dim], dt)); + buffers.insert("in_dim".into(), (in_dim as u32).to_le_bytes().to_vec()); + buffers.insert("out_dim".into(), (out_dim as u32).to_le_bytes().to_vec()); + buffers.insert("n_rows".into(), (n_rows as u32).to_le_bytes().to_vec()); + buffers.insert("n_groups".into(), (n_groups as u32).to_le_bytes().to_vec()); + buffers.insert("rows_per_group".into(), (rows_per_group as u32).to_le_bytes().to_vec()); + + let ctx = Context::new().expect("ctx"); + let mut kernel = ffai_grouped_gemm_q8::kernel_ir_for(dt.to_dtype()); + kernel.mode = KernelMode::Reduction; + let gx = (out_dim as u32).div_ceil(32); + let gy = (n_rows as u32).div_ceil(32); + let result = ctx + .dispatch_with_grid(&kernel, &buffers, &BTreeMap::new(), [gx as usize, gy as usize, 1], [ + 1024, 1, 1, + ]) + .expect("dispatch"); + let got = unpack_bytes(result.outputs.get("out").expect("out"), dt); + + for i in 0..(n_rows * out_dim) { + let w = want[i]; + let denom = w.abs().max(1.0); + assert!((got[i] - w).abs() / denom < tol, "i={i}: got={} want={}", got[i], w); + } +} + +#[test] +fn grouped_gemm_q8_f32() { + // 2 groups, each in_dim=256, out_dim=64 (rows_per_group=32), 32 rows. + run_case(Dt::F32, 256, 64, 32, 2, 32, 1e-3); + // 4 groups, edge n_rows not a multiple of 32. + run_case(Dt::F32, 128, 128, 40, 4, 32, 1e-3); +} + +#[test] +fn grouped_gemm_q8_f16() { run_case(Dt::F16, 256, 64, 32, 2, 32, 2e-2); } diff --git a/crates/metaltile-std/tests/grouped_gemv_q8_rows_correctness.rs b/crates/metaltile-std/tests/grouped_gemv_q8_rows_correctness.rs new file mode 100644 index 00000000..2162a321 --- /dev/null +++ b/crates/metaltile-std/tests/grouped_gemv_q8_rows_correctness.rs @@ -0,0 +1,103 @@ +//! Copyright 2026 TheTom +//! SPDX-License-Identifier: Apache-2.0 +//! Batched grouped Q8 gemv (ffai_grouped_gemv_q8_rows) must equal the +//! per-token single kernel (ffai_grouped_gemv_q8) row by row. NO model load. +#![cfg(target_os = "macos")] + +mod common; + +use std::collections::BTreeMap; + +use common::{Dt, gpu_lock, pack_bytes, pack_u32_bytes, unpack_bytes}; +use metaltile_core::ir::KernelMode; +use metaltile_runtime::Context; +use metaltile_std::ffai::gemv_q8::{ffai_grouped_gemv_q8, ffai_grouped_gemv_q8_rows}; + +fn xorshift(s: &mut u32) -> u32 { + let mut x = *s; + x ^= x << 13; + x ^= x >> 17; + x ^= x << 5; + *s = x; + x +} + +#[test] +fn grouped_gemv_q8_rows_matches_single() { + let _g = gpu_lock(); + let probe = Context::new().expect("Context::new"); + if probe.chip_family().is_none_or(|lvl| lvl < 10) { + eprintln!("skip"); + return; + } + drop(probe); + + let m_out = 256usize; // n_groups*rows_per_group + let rows_per_group = 64usize; // n_groups = 4 + let n_groups = m_out / rows_per_group; + let k_in = 128usize; + let bpr = k_in / 32; + let n_tokens = 6usize; + + let mut st = 0x0DEFACE0u32; + let qs: Vec = (0..m_out * bpr * 8).map(|_| xorshift(&mut st)).collect(); + let d: Vec = (0..m_out * bpr) + .map(|_| ((xorshift(&mut st) % 1000) as f32 / 1000.0 - 0.5) * 0.05) + .collect(); + let x: Vec = (0..n_tokens * n_groups * k_in) + .map(|_| ((xorshift(&mut st) % 2000) as f32 / 1000.0) - 1.0) + .collect(); + + let ctx = Context::new().unwrap(); + let consts = |kk: usize, mm: usize, rpg: usize| -> BTreeMap> { + let mut c = BTreeMap::new(); + c.insert("k_in".into(), (kk as u32).to_le_bytes().to_vec()); + c.insert("m_out".into(), (mm as u32).to_le_bytes().to_vec()); + c.insert("rows_per_group".into(), (rpg as u32).to_le_bytes().to_vec()); + c + }; + + // Reference: single kernel per token. + let mut want = vec![0.0f32; n_tokens * m_out]; + let mut ks = ffai_grouped_gemv_q8::kernel_ir_for(Dt::F32.to_dtype()); + ks.mode = KernelMode::Reduction; + for t in 0..n_tokens { + let xt = &x[t * n_groups * k_in..(t + 1) * n_groups * k_in]; + let mut b: BTreeMap> = consts(k_in, m_out, rows_per_group); + b.insert("qs".into(), pack_u32_bytes(&qs)); + b.insert("d_f32".into(), pack_bytes(&d, Dt::F32)); + b.insert("x".into(), pack_bytes(xt, Dt::F32)); + b.insert("out".into(), pack_bytes(&vec![0.0f32; m_out], Dt::F32)); + let r = + ctx.dispatch_with_grid(&ks, &b, &BTreeMap::new(), [m_out, 1, 1], [32, 1, 1]).unwrap(); + let o = unpack_bytes(r.outputs.get("out").unwrap(), Dt::F32); + want[t * m_out..(t + 1) * m_out].copy_from_slice(&o); + } + + // Batched kernel. + let mut bb: BTreeMap> = consts(k_in, m_out, rows_per_group); + bb.insert("qs".into(), pack_u32_bytes(&qs)); + bb.insert("d_f32".into(), pack_bytes(&d, Dt::F32)); + bb.insert("x".into(), pack_bytes(&x, Dt::F32)); + bb.insert("out".into(), pack_bytes(&vec![0.0f32; n_tokens * m_out], Dt::F32)); + let mut kr = ffai_grouped_gemv_q8_rows::kernel_ir_for(Dt::F32.to_dtype()); + kr.mode = KernelMode::Reduction; + let r = ctx + .dispatch_with_grid(&kr, &bb, &BTreeMap::new(), [m_out, n_tokens, 1], [32, 1, 1]) + .unwrap(); + let got = unpack_bytes(r.outputs.get("out").unwrap(), Dt::F32); + + let mut maxd = 0.0f32; + let mut nan = 0; + for (a, b) in want.iter().zip(&got) { + if !b.is_finite() { + nan += 1; + continue; + } + maxd = maxd.max((a - b).abs()); + } + eprintln!("want[0..4]={:?} got[0..4]={:?}", &want[..4], &got[..4]); + eprintln!("nan={nan} maxAbsDiff={maxd:.6}"); + assert_eq!(nan, 0); + assert!(maxd < 1e-3, "maxAbsDiff {maxd} — batched grouped gemv diverges from single"); +} diff --git a/crates/metaltile-std/tests/grouped_gemv_q8_rows_tiled_correctness.rs b/crates/metaltile-std/tests/grouped_gemv_q8_rows_tiled_correctness.rs new file mode 100644 index 00000000..a96a0b60 --- /dev/null +++ b/crates/metaltile-std/tests/grouped_gemv_q8_rows_tiled_correctness.rs @@ -0,0 +1,105 @@ +//! Copyright 2026 TheTom +//! SPDX-License-Identifier: Apache-2.0 +//! GPU correctness for `ffai::ffai_grouped_gemv_q8_rows_tiled` — the +//! token-TILED grouped Q8 gemv (8-fold weight-DRAM amortization). It must +//! equal the proven `ffai_grouped_gemv_q8_rows` row-by-row: same grouped Q8 +//! dequant and dot product, only the per-token weight reuse differs. NO model +//! load. (Mirrors `grouped_gemv_q8_rows_correctness.rs`, which validates +//! `_rows` against the per-token single kernel.) +#![cfg(target_os = "macos")] + +mod common; + +use std::collections::BTreeMap; + +use common::{Dt, gpu_lock, pack_bytes, pack_u32_bytes, unpack_bytes}; +use metaltile_core::ir::KernelMode; +use metaltile_runtime::Context; +use metaltile_std::ffai::gemv_q8::{ffai_grouped_gemv_q8_rows, ffai_grouped_gemv_q8_rows_tiled}; + +fn xorshift(s: &mut u32) -> u32 { + let mut x = *s; + x ^= x << 13; + x ^= x >> 17; + x ^= x << 5; + *s = x; + x +} + +#[test] +fn grouped_gemv_q8_rows_tiled_matches_rows() { + let _g = gpu_lock(); + let probe = Context::new().expect("Context::new"); + if probe.chip_family().is_none_or(|lvl| lvl < 10) { + eprintln!("skip: needs Apple10+ GPU"); + return; + } + drop(probe); + + let m_out = 256usize; // n_groups*rows_per_group + let rows_per_group = 64usize; // n_groups = 4 + let n_groups = m_out / rows_per_group; + let k_in = 128usize; + let bpr = k_in / 32; + // n_tokens deliberately NOT a multiple of tokens_per_tile (8) so the + // final tile is partial — exercises the `tok < n_tokens` guard. + let n_tokens = 19usize; + + let mut st = 0x71ED_0011u32; + let qs: Vec = (0..m_out * bpr * 8).map(|_| xorshift(&mut st)).collect(); + let d: Vec = (0..m_out * bpr) + .map(|_| ((xorshift(&mut st) % 1000) as f32 / 1000.0 - 0.5) * 0.05) + .collect(); + let x: Vec = (0..n_tokens * n_groups * k_in) + .map(|_| ((xorshift(&mut st) % 2000) as f32 / 1000.0) - 1.0) + .collect(); + + let ctx = Context::new().unwrap(); + + // Reference: the proven batched `_rows` kernel. + let mut bref: BTreeMap> = BTreeMap::new(); + bref.insert("qs".into(), pack_u32_bytes(&qs)); + bref.insert("d_f32".into(), pack_bytes(&d, Dt::F32)); + bref.insert("x".into(), pack_bytes(&x, Dt::F32)); + bref.insert("out".into(), pack_bytes(&vec![0.0f32; n_tokens * m_out], Dt::F32)); + bref.insert("k_in".into(), (k_in as u32).to_le_bytes().to_vec()); + bref.insert("m_out".into(), (m_out as u32).to_le_bytes().to_vec()); + bref.insert("rows_per_group".into(), (rows_per_group as u32).to_le_bytes().to_vec()); + let mut kr = ffai_grouped_gemv_q8_rows::kernel_ir_for(Dt::F32.to_dtype()); + kr.mode = KernelMode::Reduction; + let rr = ctx + .dispatch_with_grid(&kr, &bref, &BTreeMap::new(), [m_out, n_tokens, 1], [32, 1, 1]) + .unwrap(); + let want = unpack_bytes(rr.outputs.get("out").unwrap(), Dt::F32); + + // Tiled kernel: grid y = ceil(n_tokens / 8). + let mut bt: BTreeMap> = BTreeMap::new(); + bt.insert("qs".into(), pack_u32_bytes(&qs)); + bt.insert("d_f32".into(), pack_bytes(&d, Dt::F32)); + bt.insert("x".into(), pack_bytes(&x, Dt::F32)); + bt.insert("out".into(), pack_bytes(&vec![0.0f32; n_tokens * m_out], Dt::F32)); + bt.insert("k_in".into(), (k_in as u32).to_le_bytes().to_vec()); + bt.insert("m_out".into(), (m_out as u32).to_le_bytes().to_vec()); + bt.insert("rows_per_group".into(), (rows_per_group as u32).to_le_bytes().to_vec()); + bt.insert("n_tokens".into(), (n_tokens as u32).to_le_bytes().to_vec()); + let mut kt = ffai_grouped_gemv_q8_rows_tiled::kernel_ir_for(Dt::F32.to_dtype()); + kt.mode = KernelMode::Reduction; + let gy = n_tokens.div_ceil(8); + let rt = + ctx.dispatch_with_grid(&kt, &bt, &BTreeMap::new(), [m_out, gy, 1], [32, 1, 1]).unwrap(); + let got = unpack_bytes(rt.outputs.get("out").unwrap(), Dt::F32); + + let mut maxd = 0.0f32; + let mut nan = 0; + for (a, b) in want.iter().zip(&got) { + if !b.is_finite() { + nan += 1; + continue; + } + maxd = maxd.max((a - b).abs()); + } + eprintln!("want[0..4]={:?} got[0..4]={:?}", &want[..4], &got[..4]); + eprintln!("nan={nan} maxAbsDiff={maxd:.6}"); + assert_eq!(nan, 0, "non-finite output"); + assert!(maxd < 1e-3, "maxAbsDiff {maxd} — tiled grouped gemv diverges from _rows"); +} diff --git a/crates/metaltile-std/tests/moe_bgemm_iq2xxs_bm64_correctness.rs b/crates/metaltile-std/tests/moe_bgemm_iq2xxs_bm64_correctness.rs new file mode 100644 index 00000000..84cdc54b --- /dev/null +++ b/crates/metaltile-std/tests/moe_bgemm_iq2xxs_bm64_correctness.rs @@ -0,0 +1,112 @@ +//! Copyright 2026 TheTom +//! SPDX-License-Identifier: Apache-2.0 +//! GPU correctness for the high-throughput bm64 IQ2_XXS BGEMM — must match +//! the proven 16×32 pool kernel (ffai_moe_gather_bgemm_iq2xxs_mpp) on +//! identical weights/x/indices (same dequant, only the tile geometry differs). +//! cosine ≥ 0.999. NO 86GB model load. +#![cfg(target_os = "macos")] + +mod common; + +use std::collections::BTreeMap; + +use common::{Dt, gpu_lock, pack_bytes, pack_u32_bytes, unpack_bytes}; +use metaltile_core::ir::KernelMode; +use metaltile_runtime::Context; +use metaltile_std::ffai::{ + moe_bgemm_iq2xxs_bm64::ffai_moe_bgemm_iq2xxs_bm64, + moe_bgemm_iq2xxs_mpp::ffai_moe_gather_bgemm_iq2xxs_mpp, +}; + +fn xorshift(s: &mut u32) -> u32 { + let mut x = *s; + x ^= x << 13; + x ^= x >> 17; + x ^= x << 5; + *s = x; + x +} + +#[test] +fn bgemm_iq2xxs_bm64_matches_pool_kernel() { + let _g = gpu_lock(); + let probe = Context::new().expect("Context::new"); + if probe.chip_family().is_none_or(|lvl| lvl < 10) { + eprintln!("skip: needs Apple10+ GPU"); + return; + } + drop(probe); + + let n_experts = 64usize; // MANY small expert runs (2 rows each) — like + let k_in = 256usize; // real prefill routing (~topK rows/expert), + let n_out = 128usize; // many sub-runs per 64-row m-tile. + let t_rows = 128usize; + let nblk = n_out * k_in / 256; + + let indices: Vec = (0..t_rows).map(|r| (r / (t_rows / n_experts)) as u32).collect(); + let mut st = 0x13572468u32; + let qs: Vec = (0..n_experts * nblk * 16).map(|_| xorshift(&mut st)).collect(); + let d: Vec = + (0..n_experts * nblk).map(|_| (xorshift(&mut st) % 1000) as f32 * 0.0005 + 0.001).collect(); + let grid: Vec = (0..2048).map(|i| ((i * 7) % 48) as u8).collect(); + let signs: Vec = (0..128).map(|i| (i * 2) as u8).collect(); + let x: Vec = + (0..t_rows * k_in).map(|_| ((xorshift(&mut st) % 2000) as f32 / 1000.0) - 1.0).collect(); + + let ctx = Context::new().unwrap(); + + // Reference: proven 16×32 pool kernel. + let mut pool: BTreeMap> = BTreeMap::new(); + pool.insert("x".into(), pack_bytes(&x, Dt::F32)); + pool.insert("qs".into(), pack_u32_bytes(&qs)); + pool.insert("d_f32".into(), pack_bytes(&d, Dt::F32)); + pool.insert("grid".into(), grid.clone()); + pool.insert("signs".into(), signs.clone()); + pool.insert("indices".into(), pack_u32_bytes(&indices)); + pool.insert("out".into(), pack_bytes(&vec![0.0f32; t_rows * n_out], Dt::F32)); + pool.insert("m_total".into(), (t_rows as u32).to_le_bytes().to_vec()); + pool.insert("n_out".into(), (n_out as u32).to_le_bytes().to_vec()); + pool.insert("k_in".into(), (k_in as u32).to_le_bytes().to_vec()); + let mut kp = ffai_moe_gather_bgemm_iq2xxs_mpp::kernel_ir_for(Dt::F32.to_dtype()); + kp.mode = KernelMode::Reduction; + let rp = ctx + .dispatch_with_grid(&kp, &pool, &BTreeMap::new(), [n_out / 32, t_rows.div_ceil(16), 1], [ + 32, 1, 1, + ]) + .unwrap(); + let want = unpack_bytes(rp.outputs.get("out").unwrap(), Dt::F32); + + // bm64 kernel (same buffers; 128 threads, n_out/64 × M/64 grid). + let mut b: BTreeMap> = pool.clone(); + b.insert("out".into(), pack_bytes(&vec![0.0f32; t_rows * n_out], Dt::F32)); + let mut kb = ffai_moe_bgemm_iq2xxs_bm64::kernel_ir_for(Dt::F32.to_dtype()); + kb.mode = KernelMode::Reduction; + let rb = ctx + .dispatch_with_grid(&kb, &b, &BTreeMap::new(), [n_out / 64, t_rows.div_ceil(64), 1], [ + 128, 1, 1, + ]) + .unwrap(); + let got = unpack_bytes(rb.outputs.get("out").unwrap(), Dt::F32); + + let mut dot = 0.0f64; + let mut na = 0.0f64; + let mut nb = 0.0f64; + let mut nan = 0; + let mut maxd = 0.0f32; + for (a, bb) in want.iter().zip(&got) { + if !a.is_finite() || !bb.is_finite() { + nan += 1; + continue; + } + dot += (*a as f64) * (*bb as f64); + na += (*a as f64).powi(2); + nb += (*bb as f64).powi(2); + maxd = maxd.max((a - bb).abs()); + } + let cos = dot / (na.sqrt() * nb.sqrt() + 1e-12); + eprintln!("pool[0..6]={:?}", &want[..6]); + eprintln!("bm64[0..6]={:?}", &got[..6]); + eprintln!("nan={nan} cos={cos:.6} maxAbsDiff={maxd:.6}"); + assert_eq!(nan, 0, "non-finite output"); + assert!(cos >= 0.999, "cosine {cos:.6} < 0.999"); +} diff --git a/crates/metaltile-std/tests/moe_bgemm_iq2xxs_mpp_correctness.rs b/crates/metaltile-std/tests/moe_bgemm_iq2xxs_mpp_correctness.rs new file mode 100644 index 00000000..97548253 --- /dev/null +++ b/crates/metaltile-std/tests/moe_bgemm_iq2xxs_mpp_correctness.rs @@ -0,0 +1,129 @@ +//! Copyright 2026 TheTom +//! SPDX-License-Identifier: Apache-2.0 +//! GPU correctness for `ffai::ffai_moe_gather_bgemm_iq2xxs_mpp` — the +//! prefill IQ2_XXS grouped BGEMM. Oracle: per-row IQ2_XXS dequant gemv +//! (same formula as ffai_gguf_dequant_iq2_xxs). Cosine ≥ 0.99 (MMA +//! accumulation order differs from the scalar oracle). +#![cfg(target_os = "macos")] + +mod common; + +use std::collections::BTreeMap; + +use common::{Dt, gpu_lock, pack_bytes, pack_u32_bytes, unpack_bytes}; +use metaltile_core::ir::KernelMode; +use metaltile_runtime::Context; +use metaltile_std::ffai::moe_bgemm_iq2xxs_mpp::ffai_moe_gather_bgemm_iq2xxs_mpp; + +fn xorshift(s: &mut u32) -> u32 { + let mut x = *s; + x ^= x << 13; + x ^= x >> 17; + x ^= x << 5; + *s = x; + x +} + +#[test] +fn bgemm_iq2xxs_mpp_matches_gemv_oracle() { + let _g = gpu_lock(); + let probe = Context::new().expect("Context::new"); + let family = probe.chip_family(); + if family.is_none_or(|lvl| lvl < 10) { + eprintln!("skip bgemm_iq2xxs_mpp: needs Apple10+ GPU (chip_family={family:?})"); + return; + } + drop(probe); + + let n_experts = 4usize; + let k_in = 256usize; // 1 IQ2 block per row + let n_out = 64usize; // BN=32 → 2 tiles + let t_rows = 64usize; // BM=16 → 4 tiles + let nblk_per_expert = n_out * k_in / 256; + + let indices: Vec = (0..t_rows).map(|r| (r / (t_rows / n_experts)) as u32).collect(); + + let mut st = 0x1F2D_3C4Bu32; + let qs: Vec = (0..n_experts * nblk_per_expert * 16).map(|_| xorshift(&mut st)).collect(); + let d: Vec = (0..n_experts * nblk_per_expert) + .map(|_| (xorshift(&mut st) % 1000) as f32 * 0.0005 + 0.001) + .collect(); + let grid: Vec = (0..2048).map(|i| ((i * 7) % 48) as u8).collect(); + let signs: Vec = (0..128).map(|i| (i * 2) as u8).collect(); + let x: Vec = + (0..t_rows * k_in).map(|_| ((xorshift(&mut st) % 2000) as f32 / 1000.0) - 1.0).collect(); + + // Oracle: per-row IQ2_XXS dequant gemv. + let deq = |e: usize, o: usize, k: usize| -> f32 { + let vidx = o * k_in + k; + let block = vidx / 256; + let in_block = vidx % 256; + let group = in_block / 32; + let in_group = in_block % 32; + let owi = in_group / 8; + let lio = in_group % 8; + let qb = e * nblk_per_expert * 16 + block * 16 + group * 2; + let aux_idx = qs[qb]; + let aux_sgn = qs[qb + 1]; + let s4 = aux_sgn >> 28; + let db = d[e * nblk_per_expert + block] * ((s4 as f32 + 0.5) * 0.25); + let key = ((aux_idx >> (owi * 8)) & 0xff) as usize; + let octet = grid[key * 8 + lio] as f32; + let sign_mask = signs[((aux_sgn >> (owi * 7)) & 0x7f) as usize] as u32; + let sign = if (sign_mask & (1 << lio)) != 0 { -1.0 } else { 1.0 }; + db * sign * octet + }; + let mut want = vec![0.0f32; t_rows * n_out]; + for r in 0..t_rows { + let e = indices[r] as usize; + for o in 0..n_out { + let mut acc = 0.0f32; + for k in 0..k_in { + acc += deq(e, o, k) * x[r * k_in + k]; + } + want[r * n_out + o] = acc; + } + } + + let mut buffers: BTreeMap> = BTreeMap::new(); + buffers.insert("x".into(), pack_bytes(&x, Dt::F32)); + buffers.insert("qs".into(), pack_u32_bytes(&qs)); + buffers.insert("d_f32".into(), pack_bytes(&d, Dt::F32)); + buffers.insert("grid".into(), grid.clone()); + buffers.insert("signs".into(), signs.clone()); + buffers.insert("indices".into(), pack_u32_bytes(&indices)); + buffers.insert("out".into(), pack_bytes(&vec![0.0f32; t_rows * n_out], Dt::F32)); + buffers.insert("m_total".into(), (t_rows as u32).to_le_bytes().to_vec()); + buffers.insert("n_out".into(), (n_out as u32).to_le_bytes().to_vec()); + buffers.insert("k_in".into(), (k_in as u32).to_le_bytes().to_vec()); + + let ctx = Context::new().unwrap(); + let mut k = ffai_moe_gather_bgemm_iq2xxs_mpp::kernel_ir_for(Dt::F32.to_dtype()); + k.mode = KernelMode::Reduction; + let r = ctx + .dispatch_with_grid(&k, &buffers, &BTreeMap::new(), [n_out / 32, t_rows.div_ceil(16), 1], [ + 32, 1, 1, + ]) + .unwrap(); + let got = unpack_bytes(r.outputs.get("out").unwrap(), Dt::F32); + + let mut dot = 0.0f64; + let mut na = 0.0f64; + let mut nb = 0.0f64; + let mut nan = 0; + for (a, b) in want.iter().zip(&got) { + if !a.is_finite() || !b.is_finite() { + nan += 1; + continue; + } + dot += (*a as f64) * (*b as f64); + na += (*a as f64).powi(2); + nb += (*b as f64).powi(2); + } + let cos = dot / (na.sqrt() * nb.sqrt() + 1e-12); + eprintln!("want[0..6]={:?}", &want[..6]); + eprintln!("got[0..6]={:?}", &got[..6]); + eprintln!("nan={nan} cos={cos:.6}"); + assert_eq!(nan, 0, "non-finite output"); + assert!(cos >= 0.99, "cosine {cos:.6} < 0.99"); +} diff --git a/crates/metaltile-std/tests/moe_bgemm_iq2xxs_view_correctness.rs b/crates/metaltile-std/tests/moe_bgemm_iq2xxs_view_correctness.rs new file mode 100644 index 00000000..a3256e1a --- /dev/null +++ b/crates/metaltile-std/tests/moe_bgemm_iq2xxs_view_correctness.rs @@ -0,0 +1,290 @@ +//! Copyright 2026 TheTom +//! SPDX-License-Identifier: Apache-2.0 +//! GPU correctness for `ffai::ffai_moe_bgemm_iq2xxs_view` — the ZERO-COPY +//! prefill IQ2_XXS grouped BGEMM that reads raw 66-byte IQ2_XXS blocks +//! straight from a no-copy mmap VIEW buffer (vs the repacked qs/d_f32 pool). +//! Same oracle as the pool kernel (per-row IQ2_XXS dequant gemv), but the +//! super-scale d round-trips through fp16 (the view stores the raw 2-byte +//! fp16 d, which the kernel decodes inline). Cosine ≥ 0.99. +//! +//! This is the memory-safe validation of the view byte math — NO 86 GB model +//! load. If this passes, the zero-copy resident-expert path is byte-correct +//! and the only remaining work is leak-safe mmap residency. +#![cfg(target_os = "macos")] + +mod common; + +use std::collections::BTreeMap; + +use common::{Dt, gpu_lock, pack_bytes, pack_u32_bytes, unpack_bytes}; +use half::f16; +use metaltile_core::ir::KernelMode; +use metaltile_runtime::Context; +use metaltile_std::ffai::moe_bgemm_iq2xxs_view::ffai_moe_bgemm_iq2xxs_view; + +fn xorshift(s: &mut u32) -> u32 { + let mut x = *s; + x ^= x << 13; + x ^= x >> 17; + x ^= x << 5; + *s = x; + x +} + +#[test] +fn bgemm_iq2xxs_view_matches_gemv_oracle() { + let _g = gpu_lock(); + let probe = Context::new().expect("Context::new"); + let family = probe.chip_family(); + if family.is_none_or(|lvl| lvl < 10) { + eprintln!("skip bgemm_iq2xxs_view: needs Apple10+ GPU (chip_family={family:?})"); + return; + } + drop(probe); + + let n_experts = 4usize; + let k_in = 256usize; // 1 IQ2 block per row + let n_out = 64usize; // BN=32 → 2 tiles + let t_rows = 64usize; // BM=16 → 4 tiles + let nblk_per_expert = n_out * k_in / 256; + let block_bytes = 66usize; // 2-byte fp16 d + 64-byte qs (8 groups × 8 bytes) + let expert_byte_stride = nblk_per_expert * block_bytes; + + let indices: Vec = (0..t_rows).map(|r| (r / (t_rows / n_experts)) as u32).collect(); + + let mut st = 0x1F2D_3C4Bu32; + let qs: Vec = (0..n_experts * nblk_per_expert * 16).map(|_| xorshift(&mut st)).collect(); + // d values rounded THROUGH fp16 so the oracle matches the kernel's inline + // fp16 decode exactly (the view stores raw 2-byte fp16). + let d_f16: Vec = (0..n_experts * nblk_per_expert) + .map(|_| f16::from_f32((xorshift(&mut st) % 1000) as f32 * 0.0005 + 0.001)) + .collect(); + let d: Vec = d_f16.iter().map(|h| h.to_f32()).collect(); + let grid: Vec = (0..2048).map(|i| ((i * 7) % 48) as u8).collect(); + let signs: Vec = (0..128).map(|i| (i * 2) as u8).collect(); + let x: Vec = + (0..t_rows * k_in).map(|_| ((xorshift(&mut st) % 2000) as f32 / 1000.0) - 1.0).collect(); + + // Pack the raw IQ2_XXS view: per expert, nblk blocks of 66 bytes. + // block: [d_lo, d_hi, then 8 groups of (aux_idx u32 LE, aux_sgn u32 LE)]. + let mut view = vec![0u8; n_experts * expert_byte_stride]; + for e in 0..n_experts { + for b in 0..nblk_per_expert { + let base = e * expert_byte_stride + b * block_bytes; + let db = d_f16[e * nblk_per_expert + b].to_bits(); + view[base] = (db & 0xff) as u8; + view[base + 1] = (db >> 8) as u8; + for grp in 0..8 { + let qb = e * nblk_per_expert * 16 + b * 16 + grp * 2; + let aux_idx = qs[qb]; + let aux_sgn = qs[qb + 1]; + let g0 = base + 2 + grp * 8; + view[g0..g0 + 4].copy_from_slice(&aux_idx.to_le_bytes()); + view[g0 + 4..g0 + 8].copy_from_slice(&aux_sgn.to_le_bytes()); + } + } + } + + // Oracle: per-row IQ2_XXS dequant gemv (d already fp16-rounded). + let deq = |e: usize, o: usize, k: usize| -> f32 { + let vidx = o * k_in + k; + let block = vidx / 256; + let in_block = vidx % 256; + let group = in_block / 32; + let in_group = in_block % 32; + let owi = in_group / 8; + let lio = in_group % 8; + let qb = e * nblk_per_expert * 16 + block * 16 + group * 2; + let aux_idx = qs[qb]; + let aux_sgn = qs[qb + 1]; + let s4 = aux_sgn >> 28; + let db = d[e * nblk_per_expert + block] * ((s4 as f32 + 0.5) * 0.25); + let key = ((aux_idx >> (owi * 8)) & 0xff) as usize; + let octet = grid[key * 8 + lio] as f32; + let sign_mask = signs[((aux_sgn >> (owi * 7)) & 0x7f) as usize] as u32; + let sign = if (sign_mask & (1 << lio)) != 0 { -1.0 } else { 1.0 }; + db * sign * octet + }; + let mut want = vec![0.0f32; t_rows * n_out]; + for r in 0..t_rows { + let e = indices[r] as usize; + for o in 0..n_out { + let mut acc = 0.0f32; + for k in 0..k_in { + acc += deq(e, o, k) * x[r * k_in + k]; + } + want[r * n_out + o] = acc; + } + } + + let mut buffers: BTreeMap> = BTreeMap::new(); + buffers.insert("x".into(), pack_bytes(&x, Dt::F32)); + buffers.insert("view_u8".into(), view); + buffers.insert("grid".into(), grid.clone()); + buffers.insert("signs".into(), signs.clone()); + buffers.insert("indices".into(), pack_u32_bytes(&indices)); + buffers.insert("out".into(), pack_bytes(&vec![0.0f32; t_rows * n_out], Dt::F32)); + buffers.insert("m_total".into(), (t_rows as u32).to_le_bytes().to_vec()); + buffers.insert("n_out".into(), (n_out as u32).to_le_bytes().to_vec()); + buffers.insert("k_in".into(), (k_in as u32).to_le_bytes().to_vec()); + buffers.insert("tensor_byte_off".into(), 0u32.to_le_bytes().to_vec()); + buffers.insert("expert_byte_stride".into(), (expert_byte_stride as u32).to_le_bytes().to_vec()); + + let ctx = Context::new().unwrap(); + let mut k = ffai_moe_bgemm_iq2xxs_view::kernel_ir_for(Dt::F32.to_dtype()); + k.mode = KernelMode::Reduction; + let r = ctx + .dispatch_with_grid(&k, &buffers, &BTreeMap::new(), [n_out / 32, t_rows.div_ceil(16), 1], [ + 32, 1, 1, + ]) + .unwrap(); + let got = unpack_bytes(r.outputs.get("out").unwrap(), Dt::F32); + + let mut dot = 0.0f64; + let mut na = 0.0f64; + let mut nb = 0.0f64; + let mut nan = 0; + for (a, b) in want.iter().zip(&got) { + if !a.is_finite() || !b.is_finite() { + nan += 1; + continue; + } + dot += (*a as f64) * (*b as f64); + na += (*a as f64).powi(2); + nb += (*b as f64).powi(2); + } + let cos = dot / (na.sqrt() * nb.sqrt() + 1e-12); + eprintln!("want[0..6]={:?}", &want[..6]); + eprintln!("got[0..6]={:?}", &got[..6]); + eprintln!("nan={nan} cos={cos:.6}"); + assert_eq!(nan, 0, "non-finite output"); + assert!(cos >= 0.99, "cosine {cos:.6} < 0.99"); +} + +/// Production-shape variant: k_in=4096 (16 blocks/row, exercises block +/// crossing), n_out=2048, and a NONZERO tensor_byte_off (simulates a tensor +/// that does not start at byte 0 of its no-copy view window). Catches +/// block-indexing / base-offset bugs the k_in=256, offset=0 case can't. +#[test] +fn bgemm_iq2xxs_view_prod_dims_nonzero_offset() { + let _g = gpu_lock(); + let probe = Context::new().expect("Context::new"); + if probe.chip_family().is_none_or(|lvl| lvl < 10) { + eprintln!("skip: needs Apple10+ GPU"); + return; + } + drop(probe); + + let n_experts = 3usize; + let k_in = 4096usize; // 16 IQ2 blocks per row + let n_out = 64usize; // keep t_rows*n_out small + let t_rows = 48usize; + let nblk_per_expert = n_out * k_in / 256; + let block_bytes = 66usize; + let expert_byte_stride = nblk_per_expert * block_bytes; + // Nonzero base: pretend the tensor starts 12,345,678 bytes into the view. + let tensor_byte_off = 12_345_678usize; + + let indices: Vec = (0..t_rows).map(|r| (r / (t_rows / n_experts)) as u32).collect(); + let mut st = 0x55AA_1234u32; + let qs: Vec = (0..n_experts * nblk_per_expert * 16).map(|_| xorshift(&mut st)).collect(); + let d_f16: Vec = (0..n_experts * nblk_per_expert) + .map(|_| f16::from_f32((xorshift(&mut st) % 1000) as f32 * 0.0005 + 0.001)) + .collect(); + let d: Vec = d_f16.iter().map(|h| h.to_f32()).collect(); + let grid: Vec = (0..2048).map(|i| ((i * 7) % 48) as u8).collect(); + let signs: Vec = (0..128).map(|i| (i * 2) as u8).collect(); + let x: Vec = + (0..t_rows * k_in).map(|_| ((xorshift(&mut st) % 2000) as f32 / 1000.0) - 1.0).collect(); + + // view buffer = tensor_byte_off padding + n_experts blocks. + let mut view = vec![0u8; tensor_byte_off + n_experts * expert_byte_stride]; + for e in 0..n_experts { + for b in 0..nblk_per_expert { + let base = tensor_byte_off + e * expert_byte_stride + b * block_bytes; + let db = d_f16[e * nblk_per_expert + b].to_bits(); + view[base] = (db & 0xff) as u8; + view[base + 1] = (db >> 8) as u8; + for grp in 0..8 { + let qb = e * nblk_per_expert * 16 + b * 16 + grp * 2; + let g0 = base + 2 + grp * 8; + view[g0..g0 + 4].copy_from_slice(&qs[qb].to_le_bytes()); + view[g0 + 4..g0 + 8].copy_from_slice(&qs[qb + 1].to_le_bytes()); + } + } + } + + let deq = |e: usize, o: usize, k: usize| -> f32 { + let vidx = o * k_in + k; + let block = vidx / 256; + let in_block = vidx % 256; + let group = in_block / 32; + let in_group = in_block % 32; + let owi = in_group / 8; + let lio = in_group % 8; + let qb = e * nblk_per_expert * 16 + block * 16 + group * 2; + let aux_idx = qs[qb]; + let aux_sgn = qs[qb + 1]; + let s4 = aux_sgn >> 28; + let db = d[e * nblk_per_expert + block] * ((s4 as f32 + 0.5) * 0.25); + let key = ((aux_idx >> (owi * 8)) & 0xff) as usize; + let octet = grid[key * 8 + lio] as f32; + let sign_mask = signs[((aux_sgn >> (owi * 7)) & 0x7f) as usize] as u32; + let sign = if (sign_mask & (1 << lio)) != 0 { -1.0 } else { 1.0 }; + db * sign * octet + }; + let mut want = vec![0.0f32; t_rows * n_out]; + for r in 0..t_rows { + let e = indices[r] as usize; + for o in 0..n_out { + let mut acc = 0.0f32; + for k in 0..k_in { + acc += deq(e, o, k) * x[r * k_in + k]; + } + want[r * n_out + o] = acc; + } + } + + let mut buffers: BTreeMap> = BTreeMap::new(); + buffers.insert("x".into(), pack_bytes(&x, Dt::F32)); + buffers.insert("view_u8".into(), view); + buffers.insert("grid".into(), grid.clone()); + buffers.insert("signs".into(), signs.clone()); + buffers.insert("indices".into(), pack_u32_bytes(&indices)); + buffers.insert("out".into(), pack_bytes(&vec![0.0f32; t_rows * n_out], Dt::F32)); + buffers.insert("m_total".into(), (t_rows as u32).to_le_bytes().to_vec()); + buffers.insert("n_out".into(), (n_out as u32).to_le_bytes().to_vec()); + buffers.insert("k_in".into(), (k_in as u32).to_le_bytes().to_vec()); + buffers.insert("tensor_byte_off".into(), (tensor_byte_off as u32).to_le_bytes().to_vec()); + buffers.insert("expert_byte_stride".into(), (expert_byte_stride as u32).to_le_bytes().to_vec()); + + let ctx = Context::new().unwrap(); + let mut k = ffai_moe_bgemm_iq2xxs_view::kernel_ir_for(Dt::F32.to_dtype()); + k.mode = KernelMode::Reduction; + let r = ctx + .dispatch_with_grid(&k, &buffers, &BTreeMap::new(), [n_out / 32, t_rows.div_ceil(16), 1], [ + 32, 1, 1, + ]) + .unwrap(); + let got = unpack_bytes(r.outputs.get("out").unwrap(), Dt::F32); + + let mut dot = 0.0f64; + let mut na = 0.0f64; + let mut nb = 0.0f64; + let mut nan = 0; + for (a, b) in want.iter().zip(&got) { + if !a.is_finite() || !b.is_finite() { + nan += 1; + continue; + } + dot += (*a as f64) * (*b as f64); + na += (*a as f64).powi(2); + nb += (*b as f64).powi(2); + } + let cos = dot / (na.sqrt() * nb.sqrt() + 1e-12); + eprintln!("prod want[0..6]={:?}", &want[..6]); + eprintln!("prod got[0..6]={:?}", &got[..6]); + eprintln!("prod nan={nan} cos={cos:.6}"); + assert_eq!(nan, 0, "non-finite output"); + assert!(cos >= 0.99, "cosine {cos:.6} < 0.99 (prod-dims/nonzero-offset)"); +} diff --git a/crates/metaltile-std/tests/moe_bgemm_q2k_bm64_correctness.rs b/crates/metaltile-std/tests/moe_bgemm_q2k_bm64_correctness.rs new file mode 100644 index 00000000..c5630d21 --- /dev/null +++ b/crates/metaltile-std/tests/moe_bgemm_q2k_bm64_correctness.rs @@ -0,0 +1,107 @@ +//! Copyright 2026 TheTom +//! SPDX-License-Identifier: Apache-2.0 +//! bm64 Q2_K BGEMM must match the proven 16×32 pool kernel +//! (ffai_moe_gather_bgemm_q2k_mpp). cosine ≥ 0.999. NO 86GB model load. +#![cfg(target_os = "macos")] + +mod common; + +use std::collections::BTreeMap; + +use common::{Dt, gpu_lock, pack_bytes, pack_u32_bytes, unpack_bytes}; +use metaltile_core::ir::KernelMode; +use metaltile_runtime::Context; +use metaltile_std::ffai::{ + moe_bgemm_q2k_bm64::ffai_moe_bgemm_q2k_bm64, + moe_bgemm_q2k_mpp::ffai_moe_gather_bgemm_q2k_mpp, +}; + +fn xorshift(s: &mut u32) -> u32 { + let mut x = *s; + x ^= x << 13; + x ^= x >> 17; + x ^= x << 5; + *s = x; + x +} + +#[test] +fn bgemm_q2k_bm64_matches_pool_kernel() { + let _g = gpu_lock(); + let probe = Context::new().expect("Context::new"); + if probe.chip_family().is_none_or(|lvl| lvl < 10) { + eprintln!("skip: needs Apple10+ GPU"); + return; + } + drop(probe); + + let n_experts = 4usize; + let k_in = 2048usize; // production down intermediate dim + let n_out = 4096usize; // production down hidden dim + let t_rows = 128usize; + let nblk = n_out * k_in / 256; + + let indices: Vec = (0..t_rows).map(|r| (r / (t_rows / n_experts)) as u32).collect(); + let mut st = 0x2BAD_F00Du32; + let qs: Vec = (0..n_experts * nblk * 16).map(|_| xorshift(&mut st)).collect(); + let scales: Vec = + (0..n_experts * nblk * 16).map(|_| (xorshift(&mut st) & 0xff) as u8).collect(); + let d: Vec = + (0..n_experts * nblk).map(|_| (xorshift(&mut st) % 1000) as f32 * 0.0003 + 0.001).collect(); + let dmin: Vec = + (0..n_experts * nblk).map(|_| (xorshift(&mut st) % 1000) as f32 * 0.0003 + 0.001).collect(); + let x: Vec = + (0..t_rows * k_in).map(|_| ((xorshift(&mut st) % 2000) as f32 / 1000.0) - 1.0).collect(); + + let ctx = Context::new().unwrap(); + let mut pool: BTreeMap> = BTreeMap::new(); + pool.insert("x".into(), pack_bytes(&x, Dt::F32)); + pool.insert("qs".into(), pack_u32_bytes(&qs)); + pool.insert("scales".into(), scales.clone()); + pool.insert("d_f32".into(), pack_bytes(&d, Dt::F32)); + pool.insert("dmin_f32".into(), pack_bytes(&dmin, Dt::F32)); + pool.insert("indices".into(), pack_u32_bytes(&indices)); + pool.insert("out".into(), pack_bytes(&vec![0.0f32; t_rows * n_out], Dt::F32)); + pool.insert("m_total".into(), (t_rows as u32).to_le_bytes().to_vec()); + pool.insert("n_out".into(), (n_out as u32).to_le_bytes().to_vec()); + pool.insert("k_in".into(), (k_in as u32).to_le_bytes().to_vec()); + let mut kp = ffai_moe_gather_bgemm_q2k_mpp::kernel_ir_for(Dt::F32.to_dtype()); + kp.mode = KernelMode::Reduction; + let rp = ctx + .dispatch_with_grid(&kp, &pool, &BTreeMap::new(), [n_out / 32, t_rows.div_ceil(16), 1], [ + 32, 1, 1, + ]) + .unwrap(); + let want = unpack_bytes(rp.outputs.get("out").unwrap(), Dt::F32); + + let mut b = pool.clone(); + b.insert("out".into(), pack_bytes(&vec![0.0f32; t_rows * n_out], Dt::F32)); + let mut kb = ffai_moe_bgemm_q2k_bm64::kernel_ir_for(Dt::F32.to_dtype()); + kb.mode = KernelMode::Reduction; + let rb = ctx + .dispatch_with_grid(&kb, &b, &BTreeMap::new(), [n_out / 64, t_rows.div_ceil(64), 1], [ + 128, 1, 1, + ]) + .unwrap(); + let got = unpack_bytes(rb.outputs.get("out").unwrap(), Dt::F32); + + let mut dot = 0.0f64; + let mut na = 0.0f64; + let mut nb = 0.0f64; + let mut nan = 0; + let mut maxd = 0.0f32; + for (a, bb) in want.iter().zip(&got) { + if !a.is_finite() || !bb.is_finite() { + nan += 1; + continue; + } + dot += (*a as f64) * (*bb as f64); + na += (*a as f64).powi(2); + nb += (*bb as f64).powi(2); + maxd = maxd.max((a - bb).abs()); + } + let cos = dot / (na.sqrt() * nb.sqrt() + 1e-12); + eprintln!("nan={nan} cos={cos:.6} maxAbsDiff={maxd:.6}"); + assert_eq!(nan, 0, "non-finite output"); + assert!(cos >= 0.999, "cosine {cos:.6} < 0.999"); +} diff --git a/crates/metaltile-std/tests/moe_bgemm_q2k_mpp_correctness.rs b/crates/metaltile-std/tests/moe_bgemm_q2k_mpp_correctness.rs new file mode 100644 index 00000000..d02ad8ed --- /dev/null +++ b/crates/metaltile-std/tests/moe_bgemm_q2k_mpp_correctness.rs @@ -0,0 +1,121 @@ +//! Copyright 2026 TheTom +//! SPDX-License-Identifier: Apache-2.0 +//! GPU correctness for `ffai::ffai_moe_gather_bgemm_q2k_mpp` — prefill Q2_K +//! grouped BGEMM. Oracle: per-row Q2_K dequant gemv. Cosine ≥ 0.99. +#![cfg(target_os = "macos")] + +mod common; + +use std::collections::BTreeMap; + +use common::{Dt, gpu_lock, pack_bytes, pack_u32_bytes, unpack_bytes}; +use metaltile_core::ir::KernelMode; +use metaltile_runtime::Context; +use metaltile_std::ffai::moe_bgemm_q2k_mpp::ffai_moe_gather_bgemm_q2k_mpp; + +fn xorshift(s: &mut u32) -> u32 { + let mut x = *s; + x ^= x << 13; + x ^= x >> 17; + x ^= x << 5; + *s = x; + x +} + +#[test] +fn bgemm_q2k_mpp_matches_gemv_oracle() { + let _g = gpu_lock(); + let probe = Context::new().expect("Context::new"); + if probe.chip_family().is_none_or(|lvl| lvl < 10) { + eprintln!("skip bgemm_q2k_mpp: needs Apple10+ GPU"); + return; + } + drop(probe); + + let n_experts = 4usize; + let k_in = 256usize; + let n_out = 64usize; + let t_rows = 64usize; + let nblk = n_out * k_in / 256; + + let indices: Vec = (0..t_rows).map(|r| (r / (t_rows / n_experts)) as u32).collect(); + let mut st = 0x2C0F_FEE2u32; + let qs: Vec = (0..n_experts * nblk * 16).map(|_| xorshift(&mut st)).collect(); + let scales: Vec = + (0..n_experts * nblk * 16).map(|_| (xorshift(&mut st) & 0xff) as u8).collect(); + let d: Vec = + (0..n_experts * nblk).map(|_| (xorshift(&mut st) % 1000) as f32 * 0.0003 + 0.001).collect(); + let dmin: Vec = + (0..n_experts * nblk).map(|_| (xorshift(&mut st) % 1000) as f32 * 0.0003 + 0.001).collect(); + let x: Vec = + (0..t_rows * k_in).map(|_| ((xorshift(&mut st) % 2000) as f32 / 1000.0) - 1.0).collect(); + + let deq = |e: usize, o: usize, k: usize| -> f32 { + let vidx = o * k_in + k; + let block = vidx / 256; + let in_block = vidx % 256; + let sub = in_block / 16; + let q_byte_idx = in_block / 4; + let word_idx = q_byte_idx / 4; + let byte_in_word = q_byte_idx % 4; + let word = qs[e * nblk * 16 + block * 16 + word_idx]; + let qs_byte = (word >> (byte_in_word * 8)) & 0xff; + let shift = (in_block % 4) * 2; + let q2 = (qs_byte >> shift) & 0x3; + let scale_byte = scales[e * nblk * 16 + block * 16 + sub] as u32; + let s4 = scale_byte & 0xf; + let m4 = (scale_byte >> 4) & 0xf; + d[e * nblk + block] * s4 as f32 * q2 as f32 - dmin[e * nblk + block] * m4 as f32 + }; + let mut want = vec![0.0f32; t_rows * n_out]; + for r in 0..t_rows { + let e = indices[r] as usize; + for o in 0..n_out { + let mut acc = 0.0f32; + for k in 0..k_in { + acc += deq(e, o, k) * x[r * k_in + k]; + } + want[r * n_out + o] = acc; + } + } + + let mut buffers: BTreeMap> = BTreeMap::new(); + buffers.insert("x".into(), pack_bytes(&x, Dt::F32)); + buffers.insert("qs".into(), pack_u32_bytes(&qs)); + buffers.insert("scales".into(), scales.clone()); + buffers.insert("d_f32".into(), pack_bytes(&d, Dt::F32)); + buffers.insert("dmin_f32".into(), pack_bytes(&dmin, Dt::F32)); + buffers.insert("indices".into(), pack_u32_bytes(&indices)); + buffers.insert("out".into(), pack_bytes(&vec![0.0f32; t_rows * n_out], Dt::F32)); + buffers.insert("m_total".into(), (t_rows as u32).to_le_bytes().to_vec()); + buffers.insert("n_out".into(), (n_out as u32).to_le_bytes().to_vec()); + buffers.insert("k_in".into(), (k_in as u32).to_le_bytes().to_vec()); + + let ctx = Context::new().unwrap(); + let mut k = ffai_moe_gather_bgemm_q2k_mpp::kernel_ir_for(Dt::F32.to_dtype()); + k.mode = KernelMode::Reduction; + let r = ctx + .dispatch_with_grid(&k, &buffers, &BTreeMap::new(), [n_out / 32, t_rows.div_ceil(16), 1], [ + 32, 1, 1, + ]) + .unwrap(); + let got = unpack_bytes(r.outputs.get("out").unwrap(), Dt::F32); + + let mut dot = 0.0f64; + let mut na = 0.0f64; + let mut nb = 0.0f64; + let mut nan = 0; + for (a, b) in want.iter().zip(&got) { + if !a.is_finite() || !b.is_finite() { + nan += 1; + continue; + } + dot += (*a as f64) * (*b as f64); + na += (*a as f64).powi(2); + nb += (*b as f64).powi(2); + } + let cos = dot / (na.sqrt() * nb.sqrt() + 1e-12); + eprintln!("nan={nan} cos={cos:.6}"); + assert_eq!(nan, 0); + assert!(cos >= 0.99, "cosine {cos:.6} < 0.99"); +} diff --git a/crates/metaltile-std/tests/moe_bgemm_q2k_view_correctness.rs b/crates/metaltile-std/tests/moe_bgemm_q2k_view_correctness.rs new file mode 100644 index 00000000..d2c8cf9d --- /dev/null +++ b/crates/metaltile-std/tests/moe_bgemm_q2k_view_correctness.rs @@ -0,0 +1,155 @@ +//! Copyright 2026 TheTom +//! SPDX-License-Identifier: Apache-2.0 +//! GPU correctness for `ffai::ffai_moe_bgemm_q2k_view` — the ZERO-COPY Q2_K +//! grouped BGEMM that reads raw 84-byte Q2_K blocks straight from a no-copy +//! mmap VIEW buffer. Rather than re-derive a scalar oracle for the canonical +//! Q2_K layout, this proves the view kernel produces the SAME output as the +//! PROVEN pool kernel (`ffai_moe_gather_bgemm_q2k_mpp`, the validated "Tokyo" +//! path) on identical logical weights — the view just reads the raw bytes the +//! pool would have been repacked from. cosine ≥ 0.999. NO 86 GB model load. +#![cfg(target_os = "macos")] + +mod common; + +use std::collections::BTreeMap; + +use common::{Dt, gpu_lock, pack_bytes, pack_u32_bytes, unpack_bytes}; +use half::f16; +use metaltile_core::ir::KernelMode; +use metaltile_runtime::Context; +use metaltile_std::ffai::{ + moe_bgemm_q2k_mpp::ffai_moe_gather_bgemm_q2k_mpp, + moe_bgemm_q2k_view::ffai_moe_bgemm_q2k_view, +}; + +fn xorshift(s: &mut u32) -> u32 { + let mut x = *s; + x ^= x << 13; + x ^= x >> 17; + x ^= x << 5; + *s = x; + x +} + +#[test] +fn bgemm_q2k_view_matches_pool_kernel() { + let _g = gpu_lock(); + let probe = Context::new().expect("Context::new"); + if probe.chip_family().is_none_or(|lvl| lvl < 10) { + eprintln!("skip bgemm_q2k_view: needs Apple10+ GPU"); + return; + } + drop(probe); + + let n_experts = 4usize; + let k_in = 256usize; + let n_out = 64usize; + let t_rows = 64usize; + let nblk = n_out * k_in / 256; + let block_bytes = 84usize; // scales[16] + qs[64] + d_f16[2] + dmin_f16[2] + let expert_byte_stride = nblk * block_bytes; + + let indices: Vec = (0..t_rows).map(|r| (r / (t_rows / n_experts)) as u32).collect(); + let mut st = 0x2C0F_FEE2u32; + let qs: Vec = (0..n_experts * nblk * 16).map(|_| xorshift(&mut st)).collect(); + let scales: Vec = + (0..n_experts * nblk * 16).map(|_| (xorshift(&mut st) & 0xff) as u8).collect(); + // d/dmin through fp16 (the view stores raw 2-byte fp16; the pool's d_f32 + // must be the same fp16-rounded value for an exact match). + let d_f16: Vec = (0..n_experts * nblk) + .map(|_| f16::from_f32((xorshift(&mut st) % 1000) as f32 * 0.0003 + 0.001)) + .collect(); + let dmin_f16: Vec = (0..n_experts * nblk) + .map(|_| f16::from_f32((xorshift(&mut st) % 1000) as f32 * 0.0003 + 0.001)) + .collect(); + let d: Vec = d_f16.iter().map(|h| h.to_f32()).collect(); + let dmin: Vec = dmin_f16.iter().map(|h| h.to_f32()).collect(); + let x: Vec = + (0..t_rows * k_in).map(|_| ((xorshift(&mut st) % 2000) as f32 / 1000.0) - 1.0).collect(); + + // ── Run the PROVEN pool kernel (split qs/scales/d/dmin arrays). ── + let mut pool: BTreeMap> = BTreeMap::new(); + pool.insert("x".into(), pack_bytes(&x, Dt::F32)); + pool.insert("qs".into(), pack_u32_bytes(&qs)); + pool.insert("scales".into(), scales.clone()); + pool.insert("d_f32".into(), pack_bytes(&d, Dt::F32)); + pool.insert("dmin_f32".into(), pack_bytes(&dmin, Dt::F32)); + pool.insert("indices".into(), pack_u32_bytes(&indices)); + pool.insert("out".into(), pack_bytes(&vec![0.0f32; t_rows * n_out], Dt::F32)); + pool.insert("m_total".into(), (t_rows as u32).to_le_bytes().to_vec()); + pool.insert("n_out".into(), (n_out as u32).to_le_bytes().to_vec()); + pool.insert("k_in".into(), (k_in as u32).to_le_bytes().to_vec()); + let ctx = Context::new().unwrap(); + let mut kp = ffai_moe_gather_bgemm_q2k_mpp::kernel_ir_for(Dt::F32.to_dtype()); + kp.mode = KernelMode::Reduction; + let rp = ctx + .dispatch_with_grid(&kp, &pool, &BTreeMap::new(), [n_out / 32, t_rows.div_ceil(16), 1], [ + 32, 1, 1, + ]) + .unwrap(); + let want = unpack_bytes(rp.outputs.get("out").unwrap(), Dt::F32); + + // ── Pack the SAME logical weights as raw 84-byte Q2_K blocks. ── + // scales[0..16], qs[16..80] (LE bytes of the 16 u32), d_f16[80..82], + // dmin_f16[82..84]. + let mut view = vec![0u8; n_experts * expert_byte_stride]; + for e in 0..n_experts { + for b in 0..nblk { + let base = e * expert_byte_stride + b * block_bytes; + let sc0 = e * nblk * 16 + b * 16; + view[base..base + 16].copy_from_slice(&scales[sc0..sc0 + 16]); + let qb0 = e * nblk * 16 + b * 16; + for w in 0..16 { + let o = base + 16 + w * 4; + view[o..o + 4].copy_from_slice(&qs[qb0 + w].to_le_bytes()); + } + let db = d_f16[e * nblk + b].to_bits(); + view[base + 80] = (db & 0xff) as u8; + view[base + 81] = (db >> 8) as u8; + let mb = dmin_f16[e * nblk + b].to_bits(); + view[base + 82] = (mb & 0xff) as u8; + view[base + 83] = (mb >> 8) as u8; + } + } + + let mut vb: BTreeMap> = BTreeMap::new(); + vb.insert("x".into(), pack_bytes(&x, Dt::F32)); + vb.insert("view_u8".into(), view); + vb.insert("indices".into(), pack_u32_bytes(&indices)); + vb.insert("out".into(), pack_bytes(&vec![0.0f32; t_rows * n_out], Dt::F32)); + vb.insert("m_total".into(), (t_rows as u32).to_le_bytes().to_vec()); + vb.insert("n_out".into(), (n_out as u32).to_le_bytes().to_vec()); + vb.insert("k_in".into(), (k_in as u32).to_le_bytes().to_vec()); + vb.insert("tensor_byte_off".into(), 0u32.to_le_bytes().to_vec()); + vb.insert("expert_byte_stride".into(), (expert_byte_stride as u32).to_le_bytes().to_vec()); + let mut kv = ffai_moe_bgemm_q2k_view::kernel_ir_for(Dt::F32.to_dtype()); + kv.mode = KernelMode::Reduction; + let rv = ctx + .dispatch_with_grid(&kv, &vb, &BTreeMap::new(), [n_out / 32, t_rows.div_ceil(16), 1], [ + 32, 1, 1, + ]) + .unwrap(); + let got = unpack_bytes(rv.outputs.get("out").unwrap(), Dt::F32); + + let mut dot = 0.0f64; + let mut na = 0.0f64; + let mut nb = 0.0f64; + let mut nan = 0; + let mut maxd = 0.0f32; + for (a, b) in want.iter().zip(&got) { + if !a.is_finite() || !b.is_finite() { + nan += 1; + continue; + } + dot += (*a as f64) * (*b as f64); + na += (*a as f64).powi(2); + nb += (*b as f64).powi(2); + maxd = maxd.max((a - b).abs()); + } + let cos = dot / (na.sqrt() * nb.sqrt() + 1e-12); + eprintln!("pool[0..6]={:?}", &want[..6]); + eprintln!("view[0..6]={:?}", &got[..6]); + eprintln!("nan={nan} cos={cos:.6} maxAbsDiff={maxd:.6}"); + assert_eq!(nan, 0, "non-finite output"); + assert!(cos >= 0.999, "cosine {cos:.6} < 0.999 (view kernel diverges from pool)"); +} diff --git a/crates/metaltile-std/tests/moe_bm64_ragged_correctness.rs b/crates/metaltile-std/tests/moe_bm64_ragged_correctness.rs new file mode 100644 index 00000000..df2f0b9c --- /dev/null +++ b/crates/metaltile-std/tests/moe_bm64_ragged_correctness.rs @@ -0,0 +1,318 @@ +//! Copyright 2026 TheTom +//! SPDX-License-Identifier: Apache-2.0 +//! Reproduction: bm64 amortized MoE GEMM vs gemv-rows (known-correct) on +//! RAGGED expert segments that straddle 64-row tile boundaries. +//! +//! The bm64 bench (`moe_bgemm_iq2xxs_bm64`) feeds `indices = zeros` → ALL +//! rows share expert 0 → the segment-straddle path (sub_offset/sub_end +//! probing inside a 64-row tile) is NEVER exercised. In-model the routed +//! rows are ragged (variable run-length per expert, boundaries not +//! 64-aligned) and bm64 diverges from gemv-rows on provably-identical +//! inputs. This test reproduces that WITHOUT the 86GB model: same pool, +//! same ragged indices, per-row diff to localize the bug. +//! +//! macOS-gated. + +#![cfg(target_os = "macos")] + +mod common; + +use std::collections::BTreeMap; + +use common::{Dt, gpu_lock, pack_bytes, pack_u32_bytes, unpack_bytes}; +use metaltile_core::ir::KernelMode; +use metaltile_runtime::Context; +use metaltile_std::ffai::{ + moe_bgemm_iq2xxs_bm64::ffai_moe_bgemm_iq2xxs_bm64, + moe_gemv_rows_iq2xxs::ffai_moe_gemv_rows_iq2xxs, +}; + +fn read_u32(p: &str) -> Vec { + let b = std::fs::read(p).unwrap(); + b.chunks_exact(4).map(|c| u32::from_le_bytes([c[0], c[1], c[2], c[3]])).collect() +} +fn read_f32(p: &str) -> Vec { + let b = std::fs::read(p).unwrap(); + b.chunks_exact(4).map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]])).collect() +} + +/// Replay the EXACT in-model dumped bytes (/tmp/bm64dump from FFAI +/// FFAI_DBG_DUMP=1). If bm64 diverges from gemv-rows here → data-triggered +/// kernel bug. If it matches → the in-model divergence is OUTSIDE the kernel +/// (buffer aliasing / cmd-buffer hazard / stale pool). +#[test] +#[ignore] // needs /tmp/bm64dump; run explicitly with --ignored +fn bm64_replay_real_dump() { + let _g = gpu_lock(); + let dir = "/tmp/bm64dump"; + let meta = std::fs::read_to_string(format!("{dir}/meta.txt")) + .expect("dump meta — run FFAI with FFAI_DBG_DUMP=1"); + eprintln!("[replay] {meta}"); + // M=30 hidden=4096 intermediate=2048 cap=64 nblk=32768 + let g = |k: &str| { + meta.split_whitespace() + .find_map(|t| t.strip_prefix(k).map(|v| v.parse::().unwrap())) + .unwrap() + }; + let m_total = g("M="); + let k_in = g("hidden="); + let n_out = g("intermediate="); + + let qs = read_u32(&format!("{dir}/qs.u32")); + let d = read_f32(&format!("{dir}/d.f32")); + let grid = std::fs::read(format!("{dir}/grid.u8")).unwrap(); + let signs = std::fs::read(format!("{dir}/signs.u8")).unwrap(); + let idx = read_u32(&format!("{dir}/idx.u32")); + let x = read_f32(&format!("{dir}/x.f32")); + eprintln!("[replay] qs={} d={} idx={:?}", qs.len(), d.len(), idx); + + for dt in [Dt::F32, Dt::F16] { + let bm = run_bm64(&x, &qs, &d, &grid, &signs, &idx, m_total, n_out, k_in, dt); + let gv = run_gemv(&x, &qs, &d, &grid, &signs, &idx, m_total, n_out, k_in, dt); + let mut worst = 0.0f32; + let mut wr = 0; + let mut bad = 0; + for r in 0..m_total { + let mut rm = 0.0f32; + for c in 0..n_out { + let dd = (bm[r * n_out + c] - gv[r * n_out + c]).abs(); + if dd > rm { + rm = dd; + } + } + if rm > 1e-2 { + bad += 1; + } + if rm > worst { + worst = rm; + wr = r; + } + } + let sab: f32 = bm.iter().map(|v| v.abs()).sum(); + let sag: f32 = gv.iter().map(|v| v.abs()).sum(); + eprintln!( + "[replay] dt={dt:?} bad={bad}/{m_total} worst={worst:.4e} @row={wr} slot={} sumAbs bm={sab:.1} gv={sag:.1}", + idx[wr] + ); + // Compare metaltile-runtime bm64 output vs the FFAI-runtime bm64 output + // on the SAME bytes. If replay-bm == in-model-gemv but != in-model-bm, + // the FFAI runtime wrongly executes the identical kernel. + if let (Ok(imb), Ok(img)) = ( + std::fs::read(format!("{dir}/gpBm_inmodel.f32")), + std::fs::read(format!("{dir}/gpGv_inmodel.f32")), + ) { + let imb: Vec = + imb.chunks_exact(4).map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]])).collect(); + let img: Vec = + img.chunks_exact(4).map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]])).collect(); + let sab_im: f32 = imb.iter().map(|v| v.abs()).sum(); + let sag_im: f32 = img.iter().map(|v| v.abs()).sum(); + let dbm: f32 = bm.iter().zip(&imb).map(|(a, b)| (a - b).abs()).fold(0.0, f32::max); + let dgm: f32 = bm.iter().zip(&img).map(|(a, b)| (a - b).abs()).fold(0.0, f32::max); + eprintln!( + "[replay] in-model: bm={sab_im:.1} gv={sag_im:.1} | replayBM vs inModelBM worst={dbm:.3} | replayBM vs inModelGV worst={dgm:.3}" + ); + } + } +} + +/// Deterministic LCG so the test is reproducible without rand. +struct Lcg(u64); +impl Lcg { + fn next_u32(&mut self) -> u32 { + self.0 = self.0.wrapping_mul(6364136223846793005).wrapping_add(1442695040888963407); + (self.0 >> 32) as u32 + } + fn next_f32(&mut self) -> f32 { + // small symmetric values so dequant*x stays in a sane range + ((self.next_u32() % 2000) as f32 - 1000.0) * 0.001 + } +} + +#[allow(clippy::too_many_arguments)] +fn run_bm64( + x: &[f32], + qs: &[u32], + d: &[f32], + grid: &[u8], + signs: &[u8], + idx: &[u32], + m_total: usize, + n_out: usize, + k_in: usize, + dt: Dt, +) -> Vec { + let mut buffers: BTreeMap> = BTreeMap::new(); + buffers.insert("x".into(), pack_bytes(x, dt)); + buffers.insert("qs".into(), pack_u32_bytes(qs)); + buffers.insert("d_f32".into(), pack_bytes(d, Dt::F32)); + buffers.insert("grid".into(), grid.to_vec()); + buffers.insert("signs".into(), signs.to_vec()); + buffers.insert("indices".into(), pack_u32_bytes(idx)); + buffers.insert("out".into(), vec![0u8; m_total * n_out * dt.bytes()]); + // #[constexpr] params bind as constant BUFFERS (LE u32), not fn_consts. + buffers.insert("m_total".into(), (m_total as u32).to_le_bytes().to_vec()); + buffers.insert("n_out".into(), (n_out as u32).to_le_bytes().to_vec()); + buffers.insert("k_in".into(), (k_in as u32).to_le_bytes().to_vec()); + + let ctx = Context::new().expect("ctx"); + let mut k = ffai_moe_bgemm_iq2xxs_bm64::kernel_ir_for(dt.to_dtype()); + k.mode = KernelMode::Reduction; + let gx = n_out / 64; + let gy = m_total.div_ceil(64); + let r = ctx + .dispatch_with_grid(&k, &buffers, &BTreeMap::new(), [gx, gy, 1], [128, 1, 1]) + .expect("bm64 dispatch"); + unpack_bytes(r.outputs.get("out").unwrap(), dt) +} + +#[allow(clippy::too_many_arguments)] +fn run_gemv( + x: &[f32], + qs: &[u32], + d: &[f32], + grid: &[u8], + signs: &[u8], + idx: &[u32], + m_total: usize, + n_out: usize, + k_in: usize, + dt: Dt, +) -> Vec { + let mut buffers: BTreeMap> = BTreeMap::new(); + buffers.insert("x".into(), pack_bytes(x, dt)); + buffers.insert("qs_all".into(), pack_u32_bytes(qs)); + buffers.insert("d_all".into(), pack_bytes(d, Dt::F32)); + buffers.insert("expert_ids".into(), pack_u32_bytes(idx)); + buffers.insert("grid".into(), grid.to_vec()); + buffers.insert("signs".into(), signs.to_vec()); + buffers.insert("out".into(), vec![0u8; m_total * n_out * dt.bytes()]); + buffers.insert("k_in".into(), (k_in as u32).to_le_bytes().to_vec()); + buffers.insert("m_out".into(), (n_out as u32).to_le_bytes().to_vec()); + buffers.insert("m_total".into(), (m_total as u32).to_le_bytes().to_vec()); + + let ctx = Context::new().expect("ctx"); + let mut k = ffai_moe_gemv_rows_iq2xxs::kernel_ir_for(dt.to_dtype()); + k.mode = KernelMode::Reduction; + let r = ctx + .dispatch_with_grid(&k, &buffers, &BTreeMap::new(), [n_out, m_total, 1], [32, 1, 1]) + .expect("gemv dispatch"); + unpack_bytes(r.outputs.get("out").unwrap(), dt) +} + +/// Build a ragged expert assignment: runs of variable length crossing the +/// 64-row tile boundary. Returns indices[m_total]. +fn ragged_indices(m_total: usize, n_experts: u32, seed: u64) -> Vec { + let mut rng = Lcg(seed); + let mut idx = Vec::with_capacity(m_total); + let mut e = 0u32; + while idx.len() < m_total { + let run = 1 + (rng.next_u32() % 7) as usize; // runs 1..7 → straddle 64 + for _ in 0..run { + if idx.len() >= m_total { + break; + } + idx.push(e); + } + e = (e + 1) % n_experts; + } + idx +} + +#[test] +fn bm64_matches_gemv_rows_on_ragged_segments_f32() { run_case(Dt::F32); } + +#[test] +fn bm64_matches_gemv_rows_on_ragged_segments_f16() { run_case(Dt::F16); } + +fn run_case(dtype: Dt) { + let _g = gpu_lock(); + let n_experts = 6u32; + let k_in = 4096usize; // 16 blocks/row — PRODUCTION gate/up dim + let n_out = 2048usize; // 32 n-tiles of 64 — PRODUCTION intermediate + let m_total = 200usize; // 4 m-tiles, ragged straddles + let nblk = n_out * k_in / 256; + + let mut rng = Lcg(0xDEADBEEF); + let qs: Vec = (0..n_experts as usize * nblk * 16).map(|_| rng.next_u32()).collect(); + let d: Vec = (0..n_experts as usize * nblk) + .map(|_| (rng.next_u32() % 100) as f32 * 0.01 + 0.01) + .collect(); + let grid: Vec = (0..2048).map(|i| ((i * 7) % 5) as u8).collect(); + let signs: Vec = (0..128).map(|_| (rng.next_u32() & 0xff) as u8).collect(); + let x: Vec = (0..m_total * k_in).map(|_| rng.next_f32()).collect(); + let idx = ragged_indices(m_total, n_experts, 0x1234); + + // Log the segment structure so we can correlate divergence with runs. + let mut runs: Vec<(u32, usize, usize)> = vec![]; // (expert, start, len) + let mut s = 0; + while s < m_total { + let e = idx[s]; + let mut t = s; + while t < m_total && idx[t] == e { + t += 1; + } + runs.push((e, s, t - s)); + s = t; + } + eprintln!("[ragged] m_total={m_total} runs={runs:?}"); + let tile_straddles: Vec<_> = + runs.iter().filter(|(_, st, ln)| st / 64 != (st + ln - 1) / 64).collect(); + eprintln!("[ragged] runs straddling a 64-row tile boundary: {tile_straddles:?}"); + + let bm = run_bm64(&x, &qs, &d, &grid, &signs, &idx, m_total, n_out, k_in, dtype); + let gv = run_gemv(&x, &qs, &d, &grid, &signs, &idx, m_total, n_out, k_in, dtype); + // RELATIVE tolerance: bm64 (f16-staged MMA, f32 accum) vs gemv-rows + // (f32 dot) differ only by rounding; the values reach ±300 so an + // absolute tol flags pure rounding. Scale by the max magnitude. + let mag: f32 = gv.iter().map(|v| v.abs()).fold(0.0, f32::max).max(1.0); + let tol = mag + * match dtype { + Dt::F32 => 1e-4, + _ => 4e-3, + }; + let sab: f32 = bm.iter().map(|v| v.abs()).sum(); + let sag: f32 = gv.iter().map(|v| v.abs()).sum(); + eprintln!("[case] dtype={dtype:?} mag={mag:.1} tol={tol:.3e} sumAbs bm={sab:.2} gv={sag:.2}"); + assert!(sag > 1e-3, "gemv output is all-zero — test is vacuous (dispatch not computing)"); + + // Per-row diagnostics. + let mut worst = 0.0f32; + let mut worst_row = 0usize; + let mut bad_rows = 0usize; + for row in 0..m_total { + let mut row_max = 0.0f32; + for c in 0..n_out { + let dlt = (bm[row * n_out + c] - gv[row * n_out + c]).abs(); + if dlt > row_max { + row_max = dlt; + } + } + let in_tile = row / 64; + let pos_in_tile = row % 64; + if row_max > tol { + bad_rows += 1; + eprintln!( + "[diff] row={row:>3} expert={} tile={in_tile} pos={pos_in_tile:>2} row_max={row_max:.4e} bm[0]={:.4} gv[0]={:.4}", + idx[row], + bm[row * n_out], + gv[row * n_out] + ); + } + if row_max > worst { + worst = row_max; + worst_row = row; + } + } + eprintln!( + "[summary] bad_rows={bad_rows}/{m_total} worst={worst:.4e} @row={worst_row} (expert={}, tile={}, pos={})", + idx[worst_row], + worst_row / 64, + worst_row % 64 + ); + + assert!( + worst < tol, + "bm64 diverges from gemv-rows ({dtype:?}) on ragged segments: worst |diff|={worst:.4e} @row={worst_row}" + ); +} diff --git a/crates/metaltile-std/tests/moe_gather_down_q2k_correctness.rs b/crates/metaltile-std/tests/moe_gather_down_q2k_correctness.rs new file mode 100644 index 00000000..ec934a38 --- /dev/null +++ b/crates/metaltile-std/tests/moe_gather_down_q2k_correctness.rs @@ -0,0 +1,155 @@ +//! Copyright 2026 TheTom +//! SPDX-License-Identifier: Apache-2.0 +//! GPU correctness for `ffai::moe_gather_down_q2k` — fused 6-expert +//! Q2_K inline-dequant down-projection + router-weighted sum. Validates +//! against a CPU reference running the identical (production-proven) +//! Q2_K dequant formula. +//! +//! macOS-gated. Shared gpu_lock. + +#![cfg(target_os = "macos")] + +mod common; + +use std::collections::BTreeMap; + +use common::{Dt, gpu_lock, pack_bytes, pack_u32_bytes, unpack_bytes}; +use metaltile_core::ir::KernelMode; +use metaltile_runtime::Context; +use metaltile_std::ffai::moe_gather_down_q2k::ffai_moe_gather_down_q2k; + +const N_SLOTS: usize = 6; + +fn xorshift(state: &mut u32) -> u32 { + let mut x = *state; + x ^= x << 13; + x ^= x >> 17; + x ^= x << 5; + *state = x; + x +} + +fn frand(state: &mut u32) -> f32 { (xorshift(state) as f32 / u32::MAX as f32) * 2.0 - 1.0 } + +#[allow(clippy::too_many_arguments)] +fn reference( + inners_all: &[f32], + qs_all: &[u32], + scales_all: &[u8], + d_all: &[f32], + dmin_all: &[f32], + weights: &[f32], + k_in: usize, + m_out: usize, +) -> Vec { + let blocks_per_row = k_in / 256; + let nblk_per_expert = m_out * blocks_per_row; + let mut out = vec![0.0f32; m_out]; + #[allow(clippy::needless_range_loop)] + for m in 0..m_out { + let mut acc = 0.0f32; + for slot in 0..N_SLOTS { + let w_slot = weights[slot]; + let qs_row_base = (slot * nblk_per_expert + m * blocks_per_row) * 16; + let blk_row_base = slot * nblk_per_expert + m * blocks_per_row; + let inner_base = slot * k_in; + for k in 0..k_in { + let b = k / 256; + let in_block = k % 256; + let sub = in_block / 16; + let q_byte_idx = in_block / 4; + let word_idx = q_byte_idx / 4; + let byte_in_word = q_byte_idx % 4; + let word = qs_all[qs_row_base + b * 16 + word_idx]; + let qs_byte = (word >> (byte_in_word * 8)) & 0xff; + let shift = (in_block % 4) * 2; + let q_2bit = (qs_byte >> shift) & 0x3; + let scale_byte = scales_all[qs_row_base + b * 16 + sub] as u32; + let scale_4bit = scale_byte & 0xf; + let min_4bit = (scale_byte >> 4) & 0xf; + let d = d_all[blk_row_base + b]; + let dmin = dmin_all[blk_row_base + b]; + let wq = d * scale_4bit as f32 * q_2bit as f32 - dmin * min_4bit as f32; + acc += w_slot * wq * inners_all[inner_base + k]; + } + } + out[m] = acc; + } + out +} + +#[allow(clippy::too_many_arguments)] +fn run_gpu( + inners_all: &[f32], + qs_all: &[u32], + scales_all: &[u8], + d_all: &[f32], + dmin_all: &[f32], + weights: &[f32], + dt: Dt, + k_in: usize, + m_out: usize, +) -> Vec { + let mut buffers: BTreeMap> = BTreeMap::new(); + buffers.insert("inners_all".into(), pack_bytes(inners_all, dt)); + buffers.insert("qs_all".into(), pack_u32_bytes(qs_all)); + buffers.insert("scales_all".into(), scales_all.to_vec()); + buffers.insert("d_all".into(), pack_bytes(d_all, Dt::F32)); + buffers.insert("dmin_all".into(), pack_bytes(dmin_all, Dt::F32)); + let eids: Vec = (0..N_SLOTS as u32).collect(); + buffers.insert("expert_ids".into(), pack_u32_bytes(&eids)); + buffers.insert("weights".into(), pack_bytes(weights, Dt::F32)); + buffers.insert("out".into(), pack_bytes(&vec![0.0_f32; m_out], dt)); + buffers.insert("k_in".into(), (k_in as u32).to_le_bytes().to_vec()); + buffers.insert("m_out".into(), (m_out as u32).to_le_bytes().to_vec()); + buffers.insert("n_slots".into(), (N_SLOTS as u32).to_le_bytes().to_vec()); + + let ctx = Context::new().expect("Context::new on macOS"); + let mut kernel = ffai_moe_gather_down_q2k::kernel_ir_for(dt.to_dtype()); + kernel.mode = KernelMode::Reduction; + let result = ctx + .dispatch_with_grid(&kernel, &buffers, &BTreeMap::new(), [m_out, 1, 1], [32, 1, 1]) + .expect("gather_down_q2k dispatch"); + unpack_bytes(result.outputs.get("out").expect("out"), dt) +} + +fn run_case(dt: Dt, k_in: usize, m_out: usize, tol: f32) { + let _g = gpu_lock(); + let blocks_per_row = k_in / 256; + let nblk_per_expert = m_out * blocks_per_row; + + let mut st = 0x9e37_79b9u32; + let inners_all: Vec = (0..N_SLOTS * k_in).map(|_| frand(&mut st)).collect(); + let qs_all: Vec = (0..N_SLOTS * nblk_per_expert * 16).map(|_| xorshift(&mut st)).collect(); + let scales_all: Vec = + (0..N_SLOTS * nblk_per_expert * 16).map(|_| (xorshift(&mut st) & 0xff) as u8).collect(); + let d_all: Vec = + (0..N_SLOTS * nblk_per_expert).map(|_| frand(&mut st).abs() * 0.05 + 0.01).collect(); + let dmin_all: Vec = + (0..N_SLOTS * nblk_per_expert).map(|_| frand(&mut st).abs() * 0.05 + 0.01).collect(); + let weights: Vec = (0..N_SLOTS).map(|_| frand(&mut st).abs() + 0.1).collect(); + + let want = + reference(&inners_all, &qs_all, &scales_all, &d_all, &dmin_all, &weights, k_in, m_out); + let got = + run_gpu(&inners_all, &qs_all, &scales_all, &d_all, &dmin_all, &weights, dt, k_in, m_out); + + assert_eq!(got.len(), want.len()); + for (i, (g, w)) in got.iter().zip(want.iter()).enumerate() { + let denom = w.abs().max(1.0); + let rel = (g - w).abs() / denom; + assert!(rel < tol, "dt={dt:?} m={i}: got={g} want={w} rel={rel}"); + } +} + +#[test] +fn gather_down_q2k_f32() { + run_case(Dt::F32, 512, 4, 1e-3); + run_case(Dt::F32, 2048, 8, 1e-3); +} + +#[test] +fn gather_down_q2k_f16() { + run_case(Dt::F16, 512, 4, 2e-2); + run_case(Dt::F16, 2048, 8, 3e-2); +} diff --git a/crates/metaltile-std/tests/moe_gather_gemv_iq2xxs_correctness.rs b/crates/metaltile-std/tests/moe_gather_gemv_iq2xxs_correctness.rs new file mode 100644 index 00000000..c892adcc --- /dev/null +++ b/crates/metaltile-std/tests/moe_gather_gemv_iq2xxs_correctness.rs @@ -0,0 +1,165 @@ +//! Copyright 2026 TheTom +//! SPDX-License-Identifier: Apache-2.0 +//! GPU correctness for `ffai::moe_gather_gemv_iq2xxs` — the fused +//! 6-expert IQ2_XXS inline-dequant gather GEMV used by the DSv4 decode +//! FFN. Validates GPU output against a CPU reference that runs the +//! identical (production-proven) IQ2_XXS dequant formula, so a wrong +//! index/stride or a wrongly folded gather surfaces here instead of as +//! garbage decode in FFAI. +//! +//! macOS-gated. Shared gpu_lock. + +#![cfg(target_os = "macos")] + +mod common; + +use std::collections::BTreeMap; + +use common::{Dt, gpu_lock, pack_bytes, pack_u32_bytes, unpack_bytes}; +use metaltile_core::ir::KernelMode; +use metaltile_runtime::Context; +use metaltile_std::ffai::moe_gather_gemv_iq2xxs::ffai_moe_gather_gemv_iq2xxs; + +const N_SLOTS: usize = 6; + +/// Deterministic pseudo-random u32 stream (xorshift) — avoids a dev-dep +/// on `rand` and keeps the test reproducible. +fn xorshift(state: &mut u32) -> u32 { + let mut x = *state; + x ^= x << 13; + x ^= x >> 17; + x ^= x << 5; + *state = x; + x +} + +fn frand(state: &mut u32) -> f32 { + // Uniform-ish in [-1, 1). + (xorshift(state) as f32 / u32::MAX as f32) * 2.0 - 1.0 +} + +/// CPU reference: dequant the IQ2_XXS split buffers and dot each +/// expert/row against `x`. Mirrors +/// `ffai_gguf_dequant_iq2_xxs` exactly. +fn reference( + x: &[f32], + qs_all: &[u32], + d_all: &[f32], + grid: &[u8], + signs: &[u8], + k_in: usize, + m_out: usize, +) -> Vec { + let blocks_per_row = k_in / 256; + let nblk_per_expert = m_out * blocks_per_row; + let mut out = vec![0.0f32; N_SLOTS * m_out]; + for slot in 0..N_SLOTS { + for m in 0..m_out { + let qs_row_base = (slot * nblk_per_expert + m * blocks_per_row) * 16; + let d_row_base = slot * nblk_per_expert + m * blocks_per_row; + let mut acc = 0.0f32; + for b in 0..blocks_per_row { + for group in 0..8usize { + let aux_idx = qs_all[qs_row_base + b * 16 + group * 2]; + let aux_sgn = qs_all[qs_row_base + b * 16 + group * 2 + 1]; + let scale_4bit = aux_sgn >> 28; + let db = d_all[d_row_base + b] * ((scale_4bit as f32 + 0.5) * 0.25); + let x_grp = b * 256 + group * 32; + for j in 0..4usize { + let grid_key = ((aux_idx >> (j * 8)) & 0xff) as usize; + let sign_idx = ((aux_sgn >> (j * 7)) & 0x7f) as usize; + let sign_mask = signs[sign_idx] as u32; + for l in 0..8usize { + let octet = grid[grid_key * 8 + l] as f32; + let sign = if (sign_mask & (1 << l)) != 0 { -1.0 } else { 1.0 }; + let w = db * sign * octet; + acc += w * x[x_grp + j * 8 + l]; + } + } + } + } + out[slot * m_out + m] = acc; + } + } + out +} + +#[allow(clippy::too_many_arguments)] +fn run_gpu( + x: &[f32], + qs_all: &[u32], + d_all: &[f32], + grid: &[u8], + signs: &[u8], + dt: Dt, + k_in: usize, + m_out: usize, +) -> Vec { + let mut buffers: BTreeMap> = BTreeMap::new(); + buffers.insert("x".into(), pack_bytes(x, dt)); + buffers.insert("qs_all".into(), pack_u32_bytes(qs_all)); + buffers.insert("d_all".into(), pack_bytes(d_all, Dt::F32)); + // Identity expert ids → slot s reads expert s (reproduces the + // contiguous-slot reference). + let eids: Vec = (0..N_SLOTS as u32).collect(); + buffers.insert("expert_ids".into(), pack_u32_bytes(&eids)); + buffers.insert("grid".into(), grid.to_vec()); + buffers.insert("signs".into(), signs.to_vec()); + buffers.insert("out".into(), pack_bytes(&vec![0.0_f32; N_SLOTS * m_out], dt)); + buffers.insert("k_in".into(), (k_in as u32).to_le_bytes().to_vec()); + buffers.insert("m_out".into(), (m_out as u32).to_le_bytes().to_vec()); + + let ctx = Context::new().expect("Context::new on macOS"); + let mut kernel = ffai_moe_gather_gemv_iq2xxs::kernel_ir_for(dt.to_dtype()); + kernel.mode = KernelMode::Reduction; + + // grid (threadgroups) = [m_out, n_slots, 1], one 32-lane simdgroup each. + let result = ctx + .dispatch_with_grid(&kernel, &buffers, &BTreeMap::new(), [m_out, N_SLOTS, 1], [32, 1, 1]) + .expect("gather_gemv_iq2xxs dispatch"); + unpack_bytes(result.outputs.get("out").expect("out"), dt) +} + +fn run_case(dt: Dt, k_in: usize, m_out: usize, tol: f32) { + let _g = gpu_lock(); + let blocks_per_row = k_in / 256; + let nblk_per_expert = m_out * blocks_per_row; + + let mut st = 0x1234_5678u32; + let x: Vec = (0..k_in).map(|_| frand(&mut st)).collect(); + let qs_all: Vec = (0..N_SLOTS * nblk_per_expert * 16).map(|_| xorshift(&mut st)).collect(); + // Positive-ish scales so output magnitudes stay reasonable. + let d_all: Vec = + (0..N_SLOTS * nblk_per_expert).map(|_| frand(&mut st).abs() + 0.1).collect(); + // iq2xxs_grid octets are small magnitudes; emulate with bytes 0..47. + let grid: Vec = (0..2048).map(|i| ((i * 7) % 48) as u8).collect(); + let signs: Vec = (0..128).map(|i| (i * 2) as u8).collect(); + + let want = reference(&x, &qs_all, &d_all, &grid, &signs, k_in, m_out); + let got = run_gpu(&x, &qs_all, &d_all, &grid, &signs, dt, k_in, m_out); + + assert_eq!(got.len(), want.len()); + for (i, (g, w)) in got.iter().zip(want.iter()).enumerate() { + let denom = w.abs().max(1.0); + let rel = (g - w).abs() / denom; + assert!( + rel < tol, + "dt={dt:?} idx={i} (slot={}, m={}): got={g} want={w} rel={rel}", + i / m_out, + i % m_out + ); + } +} + +#[test] +fn gather_gemv_iq2xxs_f32() { + // k_in=512 → 2 blocks/row; k_in=4096 → 16 blocks/row (production gate/up). + run_case(Dt::F32, 512, 4, 1e-3); + run_case(Dt::F32, 4096, 8, 1e-3); +} + +#[test] +fn gather_gemv_iq2xxs_f16() { + run_case(Dt::F16, 512, 4, 2e-2); + run_case(Dt::F16, 4096, 8, 3e-2); +} diff --git a/crates/metaltile-std/tests/moe_gemv_rows_iq2xxs_correctness.rs b/crates/metaltile-std/tests/moe_gemv_rows_iq2xxs_correctness.rs new file mode 100644 index 00000000..0205cc7b --- /dev/null +++ b/crates/metaltile-std/tests/moe_gemv_rows_iq2xxs_correctness.rs @@ -0,0 +1,124 @@ +//! Copyright 2026 TheTom +//! SPDX-License-Identifier: Apache-2.0 +//! GPU correctness for `ffai::ffai_moe_gemv_rows_iq2xxs` — the fast +//! gemv-over-rows prefill MoE kernel (replaces the slow coop-tile bgemm). +//! Oracle: per-row IQ2_XXS dequant gemv with PER-ROW x and per-row expert. +//! Same dequant as the proven gather_gemv; only x is row-indexed. Cos ≥ 0.99. +#![cfg(target_os = "macos")] + +mod common; + +use std::collections::BTreeMap; + +use common::{Dt, gpu_lock, pack_bytes, pack_u32_bytes, unpack_bytes}; +use metaltile_core::ir::KernelMode; +use metaltile_runtime::Context; +use metaltile_std::ffai::moe_gemv_rows_iq2xxs::ffai_moe_gemv_rows_iq2xxs; + +fn xorshift(s: &mut u32) -> u32 { + let mut x = *s; + x ^= x << 13; + x ^= x >> 17; + x ^= x << 5; + *s = x; + x +} + +#[test] +fn gemv_rows_iq2xxs_matches_oracle() { + let _g = gpu_lock(); + let probe = Context::new().expect("Context::new"); + if probe.chip_family().is_none_or(|lvl| lvl < 10) { + eprintln!("skip: needs Apple10+ GPU"); + return; + } + drop(probe); + + let n_experts = 8usize; + let k_in = 512usize; // 2 blocks/row + let m_out = 64usize; + let m_total = 40usize; + let nblk = m_out * (k_in / 256); + + // Each row picks an expert (sparse, like real routing). + let mut st = 0x71AC_0011u32; + let expert_ids: Vec = (0..m_total).map(|_| xorshift(&mut st) % n_experts as u32).collect(); + let qs: Vec = (0..n_experts * nblk * 16).map(|_| xorshift(&mut st)).collect(); + let d: Vec = + (0..n_experts * nblk).map(|_| (xorshift(&mut st) % 1000) as f32 * 0.0005 + 0.001).collect(); + let grid: Vec = (0..2048).map(|i| ((i * 7) % 48) as u8).collect(); + let signs: Vec = (0..128).map(|i| (i * 2) as u8).collect(); + let x: Vec = + (0..m_total * k_in).map(|_| ((xorshift(&mut st) % 2000) as f32 / 1000.0) - 1.0).collect(); + + let deq = |e: usize, o: usize, k: usize| -> f32 { + let vidx = o * k_in + k; + let block = vidx / 256; + let in_block = vidx % 256; + let group = in_block / 32; + let in_group = in_block % 32; + let owi = in_group / 8; + let lio = in_group % 8; + let qb = e * nblk * 16 + block * 16 + group * 2; + let aux_idx = qs[qb]; + let aux_sgn = qs[qb + 1]; + let s4 = aux_sgn >> 28; + let db = d[e * nblk + block] * ((s4 as f32 + 0.5) * 0.25); + let key = ((aux_idx >> (owi * 8)) & 0xff) as usize; + let octet = grid[key * 8 + lio] as f32; + let sign_mask = signs[((aux_sgn >> (owi * 7)) & 0x7f) as usize] as u32; + let sign = if (sign_mask & (1 << lio)) != 0 { -1.0 } else { 1.0 }; + db * sign * octet + }; + let mut want = vec![0.0f32; m_total * m_out]; + for r in 0..m_total { + let e = expert_ids[r] as usize; + for o in 0..m_out { + let mut acc = 0.0f32; + for k in 0..k_in { + acc += deq(e, o, k) * x[r * k_in + k]; + } + want[r * m_out + o] = acc; + } + } + + let mut buffers: BTreeMap> = BTreeMap::new(); + buffers.insert("x".into(), pack_bytes(&x, Dt::F32)); + buffers.insert("qs_all".into(), pack_u32_bytes(&qs)); + buffers.insert("d_all".into(), pack_bytes(&d, Dt::F32)); + buffers.insert("expert_ids".into(), pack_u32_bytes(&expert_ids)); + buffers.insert("grid".into(), grid.clone()); + buffers.insert("signs".into(), signs.clone()); + buffers.insert("out".into(), pack_bytes(&vec![0.0f32; m_total * m_out], Dt::F32)); + buffers.insert("k_in".into(), (k_in as u32).to_le_bytes().to_vec()); + buffers.insert("m_out".into(), (m_out as u32).to_le_bytes().to_vec()); + buffers.insert("m_total".into(), (m_total as u32).to_le_bytes().to_vec()); + + let ctx = Context::new().unwrap(); + let mut k = ffai_moe_gemv_rows_iq2xxs::kernel_ir_for(Dt::F32.to_dtype()); + k.mode = KernelMode::Reduction; + let r = ctx + .dispatch_with_grid(&k, &buffers, &BTreeMap::new(), [m_out, m_total, 1], [32, 1, 1]) + .unwrap(); + let got = unpack_bytes(r.outputs.get("out").unwrap(), Dt::F32); + + let mut dot = 0.0f64; + let mut na = 0.0f64; + let mut nb = 0.0f64; + let mut nan = 0; + for (a, b) in want.iter().zip(&got) { + if !a.is_finite() || !b.is_finite() { + nan += 1; + continue; + } + dot += (*a as f64) * (*b as f64); + na += (*a as f64).powi(2); + nb += (*b as f64).powi(2); + } + let cos = dot / (na.sqrt() * nb.sqrt() + 1e-12); + eprintln!("want[0..6]={:?}", &want[..6]); + eprintln!("got[0..6]={:?}", &got[..6]); + eprintln!("nan={nan} cos={cos:.6}"); + assert_eq!(nan, 0, "non-finite output"); + assert!(cos >= 0.99, "cosine {cos:.6} < 0.99"); +} diff --git a/crates/metaltile-std/tests/moe_gemv_rows_q2k_correctness.rs b/crates/metaltile-std/tests/moe_gemv_rows_q2k_correctness.rs new file mode 100644 index 00000000..1f3079b6 --- /dev/null +++ b/crates/metaltile-std/tests/moe_gemv_rows_q2k_correctness.rs @@ -0,0 +1,123 @@ +//! Copyright 2026 TheTom +//! SPDX-License-Identifier: Apache-2.0 +//! GPU correctness for `ffai::ffai_moe_gemv_rows_q2k` — proves it matches the +//! PROVEN pool bgemm `ffai_moe_gather_bgemm_q2k_mpp` on identical weights + +//! per-row x (same canonical Q2_K dequant; only the GEMM structure differs). +//! cosine ≥ 0.999. NO 86GB model load. +#![cfg(target_os = "macos")] + +mod common; + +use std::collections::BTreeMap; + +use common::{Dt, gpu_lock, pack_bytes, pack_u32_bytes, unpack_bytes}; +use metaltile_core::ir::KernelMode; +use metaltile_runtime::Context; +use metaltile_std::ffai::{ + moe_bgemm_q2k_mpp::ffai_moe_gather_bgemm_q2k_mpp, + moe_gemv_rows_q2k::ffai_moe_gemv_rows_q2k, +}; + +fn xorshift(s: &mut u32) -> u32 { + let mut x = *s; + x ^= x << 13; + x ^= x >> 17; + x ^= x << 5; + *s = x; + x +} + +#[test] +fn gemv_rows_q2k_matches_pool_kernel() { + let _g = gpu_lock(); + let probe = Context::new().expect("Context::new"); + if probe.chip_family().is_none_or(|lvl| lvl < 10) { + eprintln!("skip: needs Apple10+ GPU"); + return; + } + drop(probe); + + let n_experts = 4usize; + let k_in = 256usize; + let m_out = 64usize; + let t_rows = 64usize; // == m_total + let nblk = m_out * (k_in / 256); + + // Identity routing (row r -> expert r/(t_rows/n_experts)) so the pool + // kernel (indices=expert) and gemv-rows (expert_ids=expert) read the + // same expert per row; per-row x is shared between both. + let indices: Vec = (0..t_rows).map(|r| (r / (t_rows / n_experts)) as u32).collect(); + let mut st = 0x9E37_79B9u32; + let qs: Vec = (0..n_experts * nblk * 16).map(|_| xorshift(&mut st)).collect(); + let scales: Vec = + (0..n_experts * nblk * 16).map(|_| (xorshift(&mut st) & 0xff) as u8).collect(); + let d: Vec = + (0..n_experts * nblk).map(|_| (xorshift(&mut st) % 1000) as f32 * 0.0003 + 0.001).collect(); + let dmin: Vec = + (0..n_experts * nblk).map(|_| (xorshift(&mut st) % 1000) as f32 * 0.0003 + 0.001).collect(); + let x: Vec = + (0..t_rows * k_in).map(|_| ((xorshift(&mut st) % 2000) as f32 / 1000.0) - 1.0).collect(); + + let ctx = Context::new().unwrap(); + + // Pool kernel reference. + let mut pool: BTreeMap> = BTreeMap::new(); + pool.insert("x".into(), pack_bytes(&x, Dt::F32)); + pool.insert("qs".into(), pack_u32_bytes(&qs)); + pool.insert("scales".into(), scales.clone()); + pool.insert("d_f32".into(), pack_bytes(&d, Dt::F32)); + pool.insert("dmin_f32".into(), pack_bytes(&dmin, Dt::F32)); + pool.insert("indices".into(), pack_u32_bytes(&indices)); + pool.insert("out".into(), pack_bytes(&vec![0.0f32; t_rows * m_out], Dt::F32)); + pool.insert("m_total".into(), (t_rows as u32).to_le_bytes().to_vec()); + pool.insert("n_out".into(), (m_out as u32).to_le_bytes().to_vec()); + pool.insert("k_in".into(), (k_in as u32).to_le_bytes().to_vec()); + let mut kp = ffai_moe_gather_bgemm_q2k_mpp::kernel_ir_for(Dt::F32.to_dtype()); + kp.mode = KernelMode::Reduction; + let rp = ctx + .dispatch_with_grid(&kp, &pool, &BTreeMap::new(), [m_out / 32, t_rows.div_ceil(16), 1], [ + 32, 1, 1, + ]) + .unwrap(); + let want = unpack_bytes(rp.outputs.get("out").unwrap(), Dt::F32); + + // gemv-rows kernel. + let mut gv: BTreeMap> = BTreeMap::new(); + gv.insert("x".into(), pack_bytes(&x, Dt::F32)); + gv.insert("qs".into(), pack_u32_bytes(&qs)); + gv.insert("scales".into(), scales.clone()); + gv.insert("d_f32".into(), pack_bytes(&d, Dt::F32)); + gv.insert("dmin_f32".into(), pack_bytes(&dmin, Dt::F32)); + gv.insert("expert_ids".into(), pack_u32_bytes(&indices)); + gv.insert("out".into(), pack_bytes(&vec![0.0f32; t_rows * m_out], Dt::F32)); + gv.insert("k_in".into(), (k_in as u32).to_le_bytes().to_vec()); + gv.insert("m_out".into(), (m_out as u32).to_le_bytes().to_vec()); + gv.insert("m_total".into(), (t_rows as u32).to_le_bytes().to_vec()); + let mut kv = ffai_moe_gemv_rows_q2k::kernel_ir_for(Dt::F32.to_dtype()); + kv.mode = KernelMode::Reduction; + let rv = + ctx.dispatch_with_grid(&kv, &gv, &BTreeMap::new(), [m_out, t_rows, 1], [32, 1, 1]).unwrap(); + let got = unpack_bytes(rv.outputs.get("out").unwrap(), Dt::F32); + + let mut dot = 0.0f64; + let mut na = 0.0f64; + let mut nb = 0.0f64; + let mut nan = 0; + let mut maxd = 0.0f32; + for (a, b) in want.iter().zip(&got) { + if !a.is_finite() || !b.is_finite() { + nan += 1; + continue; + } + dot += (*a as f64) * (*b as f64); + na += (*a as f64).powi(2); + nb += (*b as f64).powi(2); + maxd = maxd.max((a - b).abs()); + } + let cos = dot / (na.sqrt() * nb.sqrt() + 1e-12); + eprintln!("pool[0..6]={:?}", &want[..6]); + eprintln!("gemv[0..6]={:?}", &got[..6]); + eprintln!("nan={nan} cos={cos:.6} maxAbsDiff={maxd:.6}"); + assert_eq!(nan, 0, "non-finite output"); + assert!(cos >= 0.999, "cosine {cos:.6} < 0.999"); +} diff --git a/crates/metaltile-std/tests/moe_gemv_rows_view_iq2xxs_correctness.rs b/crates/metaltile-std/tests/moe_gemv_rows_view_iq2xxs_correctness.rs new file mode 100644 index 00000000..aec4a4e6 --- /dev/null +++ b/crates/metaltile-std/tests/moe_gemv_rows_view_iq2xxs_correctness.rs @@ -0,0 +1,226 @@ +//! Copyright 2026 TheTom +//! SPDX-License-Identifier: Apache-2.0 +//! GPU correctness for `ffai::ffai_moe_gemv_rows_view_iq2xxs` (u8-recombine +//! reads) and `ffai_moe_gemv_rows_view_u16_iq2xxs` (aligned u16 reads) — the +//! zero-copy gemv-over-rows MoE kernels that read raw 66-byte IQ2_XXS blocks +//! straight from a no-copy view buffer. Oracle: per-row IQ2_XXS dequant gemv +//! (same dequant as the proven `moe_gemv_rows_iq2xxs`), with the super-scale +//! `d` round-tripped through fp16 (the view stores raw 2-byte fp16 `d`). +//! Cosine ≥ 0.99. +#![cfg(target_os = "macos")] + +mod common; + +use std::collections::BTreeMap; + +use common::{Dt, gpu_lock, pack_bytes, pack_u32_bytes, unpack_bytes}; +use half::f16; +use metaltile_core::ir::KernelMode; +use metaltile_runtime::Context; +use metaltile_std::ffai::moe_gemv_rows_view_iq2xxs::{ + ffai_moe_gemv_rows_view_iq2xxs, + ffai_moe_gemv_rows_view_u16_iq2xxs, +}; + +fn xorshift(s: &mut u32) -> u32 { + let mut x = *s; + x ^= x << 13; + x ^= x >> 17; + x ^= x << 5; + *s = x; + x +} + +struct Case { + k_in: usize, + m_out: usize, + m_total: usize, + expert_byte_stride: usize, + expert_ids: Vec, + grid: Vec, + signs: Vec, + x: Vec, + view: Vec, + want: Vec, +} + +fn build_case() -> Case { + let n_experts = 8usize; + let k_in = 512usize; // 2 blocks/row + let m_out = 64usize; + let m_total = 40usize; + let nblk_per_expert = m_out * (k_in / 256); + let block_bytes = 66usize; + let expert_byte_stride = nblk_per_expert * block_bytes; + + let mut st = 0x71AC_0011u32; + let expert_ids: Vec = (0..m_total).map(|_| xorshift(&mut st) % n_experts as u32).collect(); + let qs: Vec = (0..n_experts * nblk_per_expert * 16).map(|_| xorshift(&mut st)).collect(); + // d rounded THROUGH fp16 so the oracle matches the kernel's inline decode. + let d_f16: Vec = (0..n_experts * nblk_per_expert) + .map(|_| f16::from_f32((xorshift(&mut st) % 1000) as f32 * 0.0005 + 0.001)) + .collect(); + let d: Vec = d_f16.iter().map(|h| h.to_f32()).collect(); + let grid: Vec = (0..2048).map(|i| ((i * 7) % 48) as u8).collect(); + let signs: Vec = (0..128).map(|i| (i * 2) as u8).collect(); + let x: Vec = + (0..m_total * k_in).map(|_| ((xorshift(&mut st) % 2000) as f32 / 1000.0) - 1.0).collect(); + + // Raw IQ2_XXS view: per expert, nblk blocks of 66 bytes: + // [d_lo, d_hi, then 8 groups of (aux_idx u32 LE, aux_sgn u32 LE)]. + let mut view = vec![0u8; n_experts * expert_byte_stride]; + for e in 0..n_experts { + for b in 0..nblk_per_expert { + let base = e * expert_byte_stride + b * block_bytes; + let db = d_f16[e * nblk_per_expert + b].to_bits(); + view[base] = (db & 0xff) as u8; + view[base + 1] = (db >> 8) as u8; + for grp in 0..8 { + let qb = e * nblk_per_expert * 16 + b * 16 + grp * 2; + let g0 = base + 2 + grp * 8; + view[g0..g0 + 4].copy_from_slice(&qs[qb].to_le_bytes()); + view[g0 + 4..g0 + 8].copy_from_slice(&qs[qb + 1].to_le_bytes()); + } + } + } + + // Oracle: per-row IQ2_XXS dequant gemv. + let deq = |e: usize, o: usize, k: usize| -> f32 { + let vidx = o * k_in + k; + let block = vidx / 256; + let in_block = vidx % 256; + let group = in_block / 32; + let in_group = in_block % 32; + let owi = in_group / 8; + let lio = in_group % 8; + let qb = e * nblk_per_expert * 16 + block * 16 + group * 2; + let aux_idx = qs[qb]; + let aux_sgn = qs[qb + 1]; + let s4 = aux_sgn >> 28; + let db = d[e * nblk_per_expert + block] * ((s4 as f32 + 0.5) * 0.25); + let key = ((aux_idx >> (owi * 8)) & 0xff) as usize; + let octet = grid[key * 8 + lio] as f32; + let sign_mask = signs[((aux_sgn >> (owi * 7)) & 0x7f) as usize] as u32; + let sign = if (sign_mask & (1 << lio)) != 0 { -1.0 } else { 1.0 }; + db * sign * octet + }; + let mut want = vec![0.0f32; m_total * m_out]; + for r in 0..m_total { + let e = expert_ids[r] as usize; + for o in 0..m_out { + let mut acc = 0.0f32; + for k in 0..k_in { + acc += deq(e, o, k) * x[r * k_in + k]; + } + want[r * m_out + o] = acc; + } + } + + Case { k_in, m_out, m_total, expert_byte_stride, expert_ids, grid, signs, x, view, want } +} + +fn cosine(want: &[f32], got: &[f32]) -> (f64, usize) { + let mut dot = 0.0f64; + let mut na = 0.0f64; + let mut nb = 0.0f64; + let mut nan = 0; + for (a, b) in want.iter().zip(got) { + if !a.is_finite() || !b.is_finite() { + nan += 1; + continue; + } + dot += (*a as f64) * (*b as f64); + na += (*a as f64).powi(2); + nb += (*b as f64).powi(2); + } + (dot / (na.sqrt() * nb.sqrt() + 1e-12), nan) +} + +fn skip_no_gpu() -> bool { + let probe = Context::new().expect("Context::new"); + let skip = probe.chip_family().is_none_or(|lvl| lvl < 10); + if skip { + eprintln!("skip: needs Apple10+ GPU"); + } + skip +} + +#[test] +fn gemv_rows_view_iq2xxs_u8_matches_oracle() { + let _g = gpu_lock(); + if skip_no_gpu() { + return; + } + let c = build_case(); + + let mut buffers: BTreeMap> = BTreeMap::new(); + buffers.insert("x".into(), pack_bytes(&c.x, Dt::F32)); + buffers.insert("view_u8".into(), c.view.clone()); + buffers.insert("grid".into(), c.grid.clone()); + buffers.insert("signs".into(), c.signs.clone()); + buffers.insert("expert_ids".into(), pack_u32_bytes(&c.expert_ids)); + buffers.insert("out".into(), pack_bytes(&vec![0.0f32; c.m_total * c.m_out], Dt::F32)); + buffers.insert("k_in".into(), (c.k_in as u32).to_le_bytes().to_vec()); + buffers.insert("m_out".into(), (c.m_out as u32).to_le_bytes().to_vec()); + buffers.insert("m_total".into(), (c.m_total as u32).to_le_bytes().to_vec()); + buffers.insert("tensor_byte_off".into(), 0u32.to_le_bytes().to_vec()); + buffers + .insert("expert_byte_stride".into(), (c.expert_byte_stride as u32).to_le_bytes().to_vec()); + + let ctx = Context::new().unwrap(); + let mut k = ffai_moe_gemv_rows_view_iq2xxs::kernel_ir_for(Dt::F32.to_dtype()); + k.mode = KernelMode::Reduction; + let r = ctx + .dispatch_with_grid(&k, &buffers, &BTreeMap::new(), [c.m_out, c.m_total, 1], [32, 1, 1]) + .unwrap(); + let got = unpack_bytes(r.outputs.get("out").unwrap(), Dt::F32); + + let (cos, nan) = cosine(&c.want, &got); + eprintln!( + "[u8] want[0..6]={:?} got[0..6]={:?} nan={nan} cos={cos:.6}", + &c.want[..6], + &got[..6] + ); + assert_eq!(nan, 0, "non-finite output"); + assert!(cos >= 0.99, "cosine {cos:.6} < 0.99"); +} + +#[test] +fn gemv_rows_view_u16_iq2xxs_matches_oracle() { + let _g = gpu_lock(); + if skip_no_gpu() { + return; + } + let c = build_case(); + // The u16 variant reads the SAME raw bytes via aligned u16 loads. + let mut buffers: BTreeMap> = BTreeMap::new(); + buffers.insert("x".into(), pack_bytes(&c.x, Dt::F32)); + buffers.insert("view_u16".into(), c.view.clone()); + buffers.insert("grid".into(), c.grid.clone()); + buffers.insert("signs".into(), c.signs.clone()); + buffers.insert("expert_ids".into(), pack_u32_bytes(&c.expert_ids)); + buffers.insert("out".into(), pack_bytes(&vec![0.0f32; c.m_total * c.m_out], Dt::F32)); + buffers.insert("k_in".into(), (c.k_in as u32).to_le_bytes().to_vec()); + buffers.insert("m_out".into(), (c.m_out as u32).to_le_bytes().to_vec()); + buffers.insert("m_total".into(), (c.m_total as u32).to_le_bytes().to_vec()); + buffers.insert("tensor_byte_off".into(), 0u32.to_le_bytes().to_vec()); + buffers + .insert("expert_byte_stride".into(), (c.expert_byte_stride as u32).to_le_bytes().to_vec()); + + let ctx = Context::new().unwrap(); + let mut k = ffai_moe_gemv_rows_view_u16_iq2xxs::kernel_ir_for(Dt::F32.to_dtype()); + k.mode = KernelMode::Reduction; + let r = ctx + .dispatch_with_grid(&k, &buffers, &BTreeMap::new(), [c.m_out, c.m_total, 1], [32, 1, 1]) + .unwrap(); + let got = unpack_bytes(r.outputs.get("out").unwrap(), Dt::F32); + + let (cos, nan) = cosine(&c.want, &got); + eprintln!( + "[u16] want[0..6]={:?} got[0..6]={:?} nan={nan} cos={cos:.6}", + &c.want[..6], + &got[..6] + ); + assert_eq!(nan, 0, "non-finite output"); + assert!(cos >= 0.99, "cosine {cos:.6} < 0.99"); +} diff --git a/crates/metaltile-std/tests/moe_gemv_ws_iq2xxs_correctness.rs b/crates/metaltile-std/tests/moe_gemv_ws_iq2xxs_correctness.rs new file mode 100644 index 00000000..5d62967e --- /dev/null +++ b/crates/metaltile-std/tests/moe_gemv_ws_iq2xxs_correctness.rs @@ -0,0 +1,132 @@ +//! Copyright 2026 TheTom +//! SPDX-License-Identifier: Apache-2.0 +//! GPU correctness for `ffai::ffai_moe_gemv_ws_iq2xxs` — the WEIGHT-STATIONARY +//! prefill MoE IQ2_XXS gemv (dequants each expert's weight row ONCE into +//! threadgroup memory, reused across the tile's rows). Oracle: per-row +//! IQ2_XXS dequant gemv from the SAME split pool (`qs_all`/`d_all`) as the +//! proven `moe_gemv_rows_iq2xxs`. The amortization is transparent to the math, +//! so this is an exact f32 comparison. Cosine ≥ 0.999. +#![cfg(target_os = "macos")] + +mod common; + +use std::collections::BTreeMap; + +use common::{Dt, gpu_lock, pack_bytes, pack_u32_bytes, unpack_bytes}; +use metaltile_core::ir::KernelMode; +use metaltile_runtime::Context; +use metaltile_std::ffai::moe_gemv_ws_iq2xxs::ffai_moe_gemv_ws_iq2xxs; + +fn xorshift(s: &mut u32) -> u32 { + let mut x = *s; + x ^= x << 13; + x ^= x >> 17; + x ^= x << 5; + *s = x; + x +} + +#[test] +fn gemv_ws_iq2xxs_matches_oracle() { + let _g = gpu_lock(); + let probe = Context::new().expect("Context::new"); + if probe.chip_family().is_none_or(|lvl| lvl < 10) { + eprintln!("skip: needs Apple10+ GPU"); + return; + } + drop(probe); + + let n_experts = 8usize; + let k_in = 512usize; // 2 blocks/row + let m_out = 64usize; + let m_total = 40usize; + let rows_per_tile = 8usize; + let nblk = m_out * (k_in / 256); // per-expert blocks + + // Rows pre-permuted by expert (contiguous per expert) — the WS layout. A + // tile is usually single-expert; expert boundaries inside a tile trigger a + // re-dequant (still correct, just slower). + let mut st = 0x71AC_0011u32; + let expert_ids: Vec = (0..m_total).map(|r| (r * n_experts / m_total) as u32).collect(); + let qs_all: Vec = (0..n_experts * nblk * 16).map(|_| xorshift(&mut st)).collect(); + let d_all: Vec = + (0..n_experts * nblk).map(|_| (xorshift(&mut st) % 1000) as f32 * 0.0005 + 0.001).collect(); + let grid: Vec = (0..2048).map(|i| ((i * 7) % 48) as u8).collect(); + let signs: Vec = (0..128).map(|i| (i * 2) as u8).collect(); + let x: Vec = + (0..m_total * k_in).map(|_| ((xorshift(&mut st) % 2000) as f32 / 1000.0) - 1.0).collect(); + + // Oracle: per-row IQ2_XXS dequant (split pool layout, same as gemv_rows). + let deq = |e: usize, o: usize, k: usize| -> f32 { + let vidx = o * k_in + k; + let block = vidx / 256; + let in_block = vidx % 256; + let group = in_block / 32; + let in_group = in_block % 32; + let owi = in_group / 8; + let lio = in_group % 8; + let qb = e * nblk * 16 + block * 16 + group * 2; + let aux_idx = qs_all[qb]; + let aux_sgn = qs_all[qb + 1]; + let s4 = aux_sgn >> 28; + let db = d_all[e * nblk + block] * ((s4 as f32 + 0.5) * 0.25); + let key = ((aux_idx >> (owi * 8)) & 0xff) as usize; + let octet = grid[key * 8 + lio] as f32; + let sign_mask = signs[((aux_sgn >> (owi * 7)) & 0x7f) as usize] as u32; + let sign = if (sign_mask & (1 << lio)) != 0 { -1.0 } else { 1.0 }; + db * sign * octet + }; + let mut want = vec![0.0f32; m_total * m_out]; + for r in 0..m_total { + let e = expert_ids[r] as usize; + for o in 0..m_out { + let mut acc = 0.0f32; + for k in 0..k_in { + acc += deq(e, o, k) * x[r * k_in + k]; + } + want[r * m_out + o] = acc; + } + } + + let mut buffers: BTreeMap> = BTreeMap::new(); + buffers.insert("x".into(), pack_bytes(&x, Dt::F32)); + buffers.insert("qs_all".into(), pack_u32_bytes(&qs_all)); + buffers.insert("d_all".into(), pack_bytes(&d_all, Dt::F32)); + buffers.insert("expert_ids".into(), pack_u32_bytes(&expert_ids)); + buffers.insert("grid".into(), grid.clone()); + buffers.insert("signs".into(), signs.clone()); + buffers.insert("out".into(), pack_bytes(&vec![0.0f32; m_total * m_out], Dt::F32)); + buffers.insert("k_in".into(), (k_in as u32).to_le_bytes().to_vec()); + buffers.insert("m_out".into(), (m_out as u32).to_le_bytes().to_vec()); + buffers.insert("m_total".into(), (m_total as u32).to_le_bytes().to_vec()); + buffers.insert("rows_per_tile".into(), (rows_per_tile as u32).to_le_bytes().to_vec()); + + let ctx = Context::new().unwrap(); + let mut k = ffai_moe_gemv_ws_iq2xxs::kernel_ir_for(Dt::F32.to_dtype()); + k.mode = KernelMode::Reduction; + let gy = m_total.div_ceil(rows_per_tile); + let r = + ctx.dispatch_with_grid(&k, &buffers, &BTreeMap::new(), [m_out, gy, 1], [32, 1, 1]).unwrap(); + let got = unpack_bytes(r.outputs.get("out").unwrap(), Dt::F32); + + let mut dot = 0.0f64; + let mut na = 0.0f64; + let mut nb = 0.0f64; + let mut nan = 0; + let mut maxd = 0.0f32; + for (a, b) in want.iter().zip(&got) { + if !a.is_finite() || !b.is_finite() { + nan += 1; + continue; + } + dot += (*a as f64) * (*b as f64); + na += (*a as f64).powi(2); + nb += (*b as f64).powi(2); + maxd = maxd.max((a - b).abs()); + } + let cos = dot / (na.sqrt() * nb.sqrt() + 1e-12); + eprintln!("want[0..6]={:?} got[0..6]={:?}", &want[..6], &got[..6]); + eprintln!("nan={nan} cos={cos:.6} maxAbsDiff={maxd:.6}"); + assert_eq!(nan, 0, "non-finite output"); + assert!(cos >= 0.999, "cosine {cos:.6} < 0.999"); +} diff --git a/crates/metaltile-std/tests/moe_gemv_ws_q2k_correctness.rs b/crates/metaltile-std/tests/moe_gemv_ws_q2k_correctness.rs new file mode 100644 index 00000000..22570af2 --- /dev/null +++ b/crates/metaltile-std/tests/moe_gemv_ws_q2k_correctness.rs @@ -0,0 +1,123 @@ +//! Copyright 2026 TheTom +//! SPDX-License-Identifier: Apache-2.0 +//! GPU correctness for `ffai::ffai_moe_gemv_ws_q2k` — the WEIGHT-STATIONARY +//! prefill MoE Q2_K gemv (down projection). It dequants each expert's weight +//! row ONCE into threadgroup memory and reuses it across the tile; the math is +//! identical to the proven `ffai_moe_gemv_rows_q2k` (same canonical Q2_K +//! dequant, same split pool, same per-row dot). Oracle = the gemv-rows kernel +//! on the SAME inputs. Exact-ish f32 agreement (cosine ≥ 0.999), NO model load. +#![cfg(target_os = "macos")] + +mod common; + +use std::collections::BTreeMap; + +use common::{Dt, gpu_lock, pack_bytes, pack_u32_bytes, unpack_bytes}; +use metaltile_core::ir::KernelMode; +use metaltile_runtime::Context; +use metaltile_std::ffai::{ + moe_gemv_rows_q2k::ffai_moe_gemv_rows_q2k, + moe_gemv_ws_q2k::ffai_moe_gemv_ws_q2k, +}; + +fn xorshift(s: &mut u32) -> u32 { + let mut x = *s; + x ^= x << 13; + x ^= x >> 17; + x ^= x << 5; + *s = x; + x +} + +#[test] +fn gemv_ws_q2k_matches_gemv_rows() { + let _g = gpu_lock(); + let probe = Context::new().expect("Context::new"); + if probe.chip_family().is_none_or(|lvl| lvl < 10) { + eprintln!("skip: needs Apple10+ GPU"); + return; + } + drop(probe); + + let n_experts = 4usize; + let k_in = 256usize; // 1 block/row + let m_out = 64usize; + let m_total = 40usize; + let rows_per_tile = 8usize; + let nblk = m_out * (k_in / 256); + + // Rows pre-permuted by expert (contiguous) — the WS layout. + let expert_ids: Vec = (0..m_total).map(|r| (r * n_experts / m_total) as u32).collect(); + let mut st = 0x9E37_79B9u32; + let qs: Vec = (0..n_experts * nblk * 16).map(|_| xorshift(&mut st)).collect(); + let scales: Vec = + (0..n_experts * nblk * 16).map(|_| (xorshift(&mut st) & 0xff) as u8).collect(); + let d: Vec = + (0..n_experts * nblk).map(|_| (xorshift(&mut st) % 1000) as f32 * 0.0003 + 0.001).collect(); + let dmin: Vec = + (0..n_experts * nblk).map(|_| (xorshift(&mut st) % 1000) as f32 * 0.0003 + 0.001).collect(); + let x: Vec = + (0..m_total * k_in).map(|_| ((xorshift(&mut st) % 2000) as f32 / 1000.0) - 1.0).collect(); + + let ctx = Context::new().unwrap(); + + // Reference: proven gemv-rows kernel. + let mut bref: BTreeMap> = BTreeMap::new(); + bref.insert("x".into(), pack_bytes(&x, Dt::F32)); + bref.insert("qs".into(), pack_u32_bytes(&qs)); + bref.insert("scales".into(), scales.clone()); + bref.insert("d_f32".into(), pack_bytes(&d, Dt::F32)); + bref.insert("dmin_f32".into(), pack_bytes(&dmin, Dt::F32)); + bref.insert("expert_ids".into(), pack_u32_bytes(&expert_ids)); + bref.insert("out".into(), pack_bytes(&vec![0.0f32; m_total * m_out], Dt::F32)); + bref.insert("k_in".into(), (k_in as u32).to_le_bytes().to_vec()); + bref.insert("m_out".into(), (m_out as u32).to_le_bytes().to_vec()); + bref.insert("m_total".into(), (m_total as u32).to_le_bytes().to_vec()); + let mut kr = ffai_moe_gemv_rows_q2k::kernel_ir_for(Dt::F32.to_dtype()); + kr.mode = KernelMode::Reduction; + let rr = ctx + .dispatch_with_grid(&kr, &bref, &BTreeMap::new(), [m_out, m_total, 1], [32, 1, 1]) + .unwrap(); + let want = unpack_bytes(rr.outputs.get("out").unwrap(), Dt::F32); + + // Weight-stationary kernel. + let mut bws: BTreeMap> = BTreeMap::new(); + bws.insert("x".into(), pack_bytes(&x, Dt::F32)); + bws.insert("qs".into(), pack_u32_bytes(&qs)); + bws.insert("scales".into(), scales.clone()); + bws.insert("d_f32".into(), pack_bytes(&d, Dt::F32)); + bws.insert("dmin_f32".into(), pack_bytes(&dmin, Dt::F32)); + bws.insert("expert_ids".into(), pack_u32_bytes(&expert_ids)); + bws.insert("out".into(), pack_bytes(&vec![0.0f32; m_total * m_out], Dt::F32)); + bws.insert("k_in".into(), (k_in as u32).to_le_bytes().to_vec()); + bws.insert("m_out".into(), (m_out as u32).to_le_bytes().to_vec()); + bws.insert("m_total".into(), (m_total as u32).to_le_bytes().to_vec()); + bws.insert("rows_per_tile".into(), (rows_per_tile as u32).to_le_bytes().to_vec()); + let mut kw = ffai_moe_gemv_ws_q2k::kernel_ir_for(Dt::F32.to_dtype()); + kw.mode = KernelMode::Reduction; + let gy = m_total.div_ceil(rows_per_tile); + let rw = + ctx.dispatch_with_grid(&kw, &bws, &BTreeMap::new(), [m_out, gy, 1], [32, 1, 1]).unwrap(); + let got = unpack_bytes(rw.outputs.get("out").unwrap(), Dt::F32); + + let mut dot = 0.0f64; + let mut na = 0.0f64; + let mut nb = 0.0f64; + let mut nan = 0; + let mut maxd = 0.0f32; + for (a, b) in want.iter().zip(&got) { + if !a.is_finite() || !b.is_finite() { + nan += 1; + continue; + } + dot += (*a as f64) * (*b as f64); + na += (*a as f64).powi(2); + nb += (*b as f64).powi(2); + maxd = maxd.max((a - b).abs()); + } + let cos = dot / (na.sqrt() * nb.sqrt() + 1e-12); + eprintln!("rows[0..6]={:?} ws[0..6]={:?}", &want[..6], &got[..6]); + eprintln!("nan={nan} cos={cos:.6} maxAbsDiff={maxd:.6}"); + assert_eq!(nan, 0, "non-finite output"); + assert!(cos >= 0.999, "cosine {cos:.6} < 0.999"); +} diff --git a/crates/metaltile-std/tests/moe_q2k_view_u16_correctness.rs b/crates/metaltile-std/tests/moe_q2k_view_u16_correctness.rs new file mode 100644 index 00000000..541f8cab --- /dev/null +++ b/crates/metaltile-std/tests/moe_q2k_view_u16_correctness.rs @@ -0,0 +1,141 @@ +//! Copyright 2026 TheTom +//! SPDX-License-Identifier: Apache-2.0 +//! Q2_K view-bm64 (raw 84-byte blocks) must match the validated pool q2k_bm64 +//! on the SAME source. Multi-expert grouped indices + absolute zeroRows check +//! (the view-vs-pool diff alone masks bugs shared by both). macOS, no model. +#![cfg(target_os = "macos")] +mod common; +use std::collections::BTreeMap; + +use common::{Dt, gpu_lock, pack_bytes, pack_u32_bytes, unpack_bytes}; +use metaltile_core::{dtype::DType, ir::KernelMode}; +use metaltile_runtime::Context; +use metaltile_std::ffai::{ + moe_bgemm_q2k_bm64::ffai_moe_bgemm_q2k_bm64, + moe_bgemm_q2k_view_u16_bm64::ffai_moe_bgemm_q2k_view_u16_bm64, +}; + +fn xs(s: &mut u32) -> u32 { + let mut x = *s; + x ^= x << 13; + x ^= x >> 17; + x ^= x << 5; + *s = x; + x +} + +#[test] +fn q2k_view_u16_bm64_matches_pool() { + let _g = gpu_lock(); + let ctx = Context::new().unwrap(); + if ctx.chip_family().is_none_or(|lvl| lvl < 10) { + eprintln!("skip: needs Apple10+"); + return; + } + let n_experts = 4usize; + let pool_cap = 8usize; // pool bigger than experts + let k_in = 2048usize; + let n_out = 4096usize; + let m_total = 72usize; + let nblk = n_out * k_in / 256; + + // RAW Q2_K blocks: 84 bytes/block = scales[16] + qs[64] + d f16 + dmin f16. + let mut st = 0x2BADF00Du32; + let mut raw = vec![0u8; pool_cap * nblk * 84]; + for e in 0..n_experts { + for b in 0..nblk { + let base = (e * nblk + b) * 84; + for j in 0..80 { + raw[base + j] = (xs(&mut st) & 0xff) as u8; + } // scales+qs + let d = half::f16::from_f32(((xs(&mut st) % 1000) as f32) * 0.0003 + 0.001).to_bits(); + let dm = half::f16::from_f32(((xs(&mut st) % 1000) as f32) * 0.0003 + 0.001).to_bits(); + raw[base + 80..base + 82].copy_from_slice(&d.to_le_bytes()); + raw[base + 82..base + 84].copy_from_slice(&dm.to_le_bytes()); + } + } + // Deinterleave → pool: qs u32[*16], scales u8[*16], d f32, dmin f32. + let mut qs = vec![0u32; pool_cap * nblk * 16]; + let mut scales = vec![0u8; pool_cap * nblk * 16]; + let mut d = vec![0f32; pool_cap * nblk]; + let mut dmin = vec![0f32; pool_cap * nblk]; + for e in 0..n_experts { + for b in 0..nblk { + let base = (e * nblk + b) * 84; + let blk = e * nblk + b; + for j in 0..16 { + scales[blk * 16 + j] = raw[base + j]; + } + for w in 0..16 { + let o = base + 16 + w * 4; + qs[blk * 16 + w] = u32::from_le_bytes([raw[o], raw[o + 1], raw[o + 2], raw[o + 3]]); + } + d[blk] = half::f16::from_le_bytes([raw[base + 80], raw[base + 81]]).to_f32(); + dmin[blk] = half::f16::from_le_bytes([raw[base + 82], raw[base + 83]]).to_f32(); + } + } + let x: Vec = + (0..m_total * k_in).map(|_| ((xs(&mut st) % 2000) as f32 / 1000.0) - 1.0).collect(); + let indices: Vec = (0..m_total).map(|r| (r * n_experts / m_total) as u32).collect(); + let dt = Dt::F32; + + let run = |buffers: BTreeMap>, k: metaltile_core::ir::Kernel| -> Vec { + let mut kk = k; + kk.mode = KernelMode::Reduction; + let r = ctx + .dispatch_with_grid( + &kk, + &buffers, + &BTreeMap::new(), + [n_out / 64, m_total.div_ceil(64), 1], + [128, 1, 1], + ) + .expect("dispatch"); + unpack_bytes(r.outputs.get("out").unwrap(), dt) + }; + let mut pb: BTreeMap> = BTreeMap::new(); + pb.insert("x".into(), pack_bytes(&x, dt)); + pb.insert("qs".into(), pack_u32_bytes(&qs)); + pb.insert("scales".into(), scales); + pb.insert("d_f32".into(), pack_bytes(&d, Dt::F32)); + pb.insert("dmin_f32".into(), pack_bytes(&dmin, Dt::F32)); + pb.insert("indices".into(), pack_u32_bytes(&indices)); + pb.insert("out".into(), vec![0u8; m_total * n_out * 4]); + pb.insert("m_total".into(), (m_total as u32).to_le_bytes().to_vec()); + pb.insert("n_out".into(), (n_out as u32).to_le_bytes().to_vec()); + pb.insert("k_in".into(), (k_in as u32).to_le_bytes().to_vec()); + let pool_out = run(pb, ffai_moe_bgemm_q2k_bm64::kernel_ir_for(DType::F32)); + + let raw_u16: Vec = raw.clone(); // bytes; view_u16/view_f16 reinterpret + let mut vb: BTreeMap> = BTreeMap::new(); + vb.insert("x".into(), pack_bytes(&x, dt)); + vb.insert("view_u16".into(), raw_u16.clone()); + vb.insert("view_f16".into(), raw_u16); + vb.insert("indices".into(), pack_u32_bytes(&indices)); + vb.insert("out".into(), vec![0u8; m_total * n_out * 4]); + vb.insert("m_total".into(), (m_total as u32).to_le_bytes().to_vec()); + vb.insert("n_out".into(), (n_out as u32).to_le_bytes().to_vec()); + vb.insert("k_in".into(), (k_in as u32).to_le_bytes().to_vec()); + vb.insert("tensor_byte_off".into(), 0u32.to_le_bytes().to_vec()); + vb.insert("expert_byte_stride".into(), ((nblk * 84) as u32).to_le_bytes().to_vec()); + let view_out = run(vb, ffai_moe_bgemm_q2k_view_u16_bm64::kernel_ir_for(DType::F32)); + + let mut worst = 0.0f32; + let mut wi = 0; + for i in 0..pool_out.len() { + let dd = (pool_out[i] - view_out[i]).abs(); + if dd > worst { + worst = dd; + wi = i; + } + } + let mag: f32 = pool_out.iter().map(|v| v.abs()).fold(0.0, f32::max).max(1.0); + let vz = + (0..m_total).filter(|&r| (0..n_out).all(|c| view_out[r * n_out + c].abs() < 1e-6)).count(); + eprintln!( + "[q2k-view] worst={worst:.4e} @i={wi} pool={:.4} view={:.4} mag={mag:.2} viewZeroRows={vz}/{m_total}", + pool_out[wi], view_out[wi] + ); + assert!(worst < mag * 2e-3, "q2k view-bm64 != pool: worst={worst:.4e} (mag {mag:.2})"); + assert_eq!(vz, 0, "q2k view left {vz}/{m_total} rows zero"); +} diff --git a/crates/metaltile-std/tests/moe_view_u16_correctness.rs b/crates/metaltile-std/tests/moe_view_u16_correctness.rs new file mode 100644 index 00000000..9687338f --- /dev/null +++ b/crates/metaltile-std/tests/moe_view_u16_correctness.rs @@ -0,0 +1,207 @@ +//! Copyright 2026 TheTom +//! SPDX-License-Identifier: Apache-2.0 +//! Localize the u16 view-bm64 dequant bug: run the POOL bm64 (validated) and +//! the view-u16 bm64 on the SAME IQ2_XXS source (raw 66-byte blocks → +//! deinterleaved split pool) with single-expert indices. They must match. +//! In-model the view path gives argmax 79872 (wrong) vs the pool's Tokyo; +//! this isolates whether it's the raw u16 read / d-decode / indexing. +//! macOS-gated. +#![cfg(target_os = "macos")] + +mod common; +use std::collections::BTreeMap; + +use common::{Dt, gpu_lock, pack_bytes, pack_u32_bytes, unpack_bytes}; +use metaltile_core::{dtype::DType, ir::KernelMode}; +use metaltile_runtime::Context; +use metaltile_std::ffai::{ + moe_bgemm_iq2xxs_bm64::ffai_moe_bgemm_iq2xxs_bm64, + moe_bgemm_iq2xxs_view_u16_bm64::ffai_moe_bgemm_iq2xxs_view_u16_bm64, +}; + +struct Lcg(u64); +impl Lcg { + fn u16v(&mut self) -> u16 { + self.0 = self.0.wrapping_mul(6364136223846793005).wrapping_add(1); + (self.0 >> 33) as u16 + } +} + +#[test] +fn view_u16_bm64_matches_pool_bm64() { + let _g = gpu_lock(); + // In-model condition (VIEW_RAGGED=1): M ragged (72 = 64+8, NOT a mult of 64), + // indices grouped across multiple experts (permuted-by-expert layout), and a + // pool sized LARGER than the experts actually filled (poolCap > n_experts) so + // some slots are uninitialized. The original case (M=256, single expert, + // n_experts==poolCap) passes — this exercises the ragged/multi-expert path. + let ragged = std::env::var("VIEW_RAGGED").is_ok(); + // VIEW_TINYRUNS: mimic in-model — M=72 across MANY experts (~1-2 rows each + // = lots of tiny runs in a 64-row tile), oversized pool. This is what the + // in-model prefill actually feeds (47 experts / 72 rows). + let tiny = std::env::var("VIEW_TINYRUNS").is_ok(); + let k_in = 4096usize; + let n_out = 2048usize; + let m_total = if ragged || tiny { 72usize } else { 256usize }; + let n_experts = if tiny { 47usize } else { 4usize }; + let pool_cap = if ragged || tiny { 64usize } else { n_experts }; + let nblk = n_out * k_in / 256; + // RAW blocks: per expert, nblk blocks of 33 u16 (1 d-f16-bits + 32 qs). + let mut rng = Lcg(0xBEEF); + let constant = std::env::var("VIEW_CONST").is_ok(); + // Buffers sized to pool_cap experts (tail slots beyond n_experts left zero — + // mimics the in-model oversized resident pool). Real data only in 0..n_experts. + let raw_u16: Vec = (0..pool_cap * nblk * 33) + .map(|i| { + let e = i / (nblk * 33); + if e >= n_experts { + return 0u16; + } + if i % 33 == 0 { + half::f16::from_f32(if constant { + 1.0 + } else { + ((rng.u16v() % 200) as f32 - 100.0) * 0.001 + }) + .to_bits() + } else if constant { + 0u16 + } else { + rng.u16v() + } + }) + .collect(); + // Deinterleave → split pool: qs u32[pool_cap*nblk*16], d f32[pool_cap*nblk]. + let mut qs = vec![0u32; pool_cap * nblk * 16]; + let mut d = vec![0f32; pool_cap * nblk]; + for e in 0..n_experts { + for b in 0..nblk { + let base = (e * nblk + b) * 33; + d[e * nblk + b] = half::f16::from_bits(raw_u16[base]).to_f32(); + for i in 0..16 { + // 32 qs u16 → 16 u32 + qs[(e * nblk + b) * 16 + i] = raw_u16[base + 1 + 2 * i] as u32 + | ((raw_u16[base + 1 + 2 * i + 1] as u32) << 16); + } + } + } + let x: Vec = (0..m_total * k_in).map(|i| ((i % 17) as f32 - 8.0) * 0.01).collect(); + let grid: Vec = (0..2048).map(|i| ((i * 7) % 5) as u8).collect(); + let signs: Vec = (0..128).map(|i| (i * 13 % 256) as u8).collect(); + let dt = Dt::F32; + // indices = GLOBAL expert id per row (== slot, pool is compacted). Ragged: + // grouped runs across n_experts (permuted-by-expert), so the last m-tile + // (rows 64..71) is a partial run — the exact in-model shape. + // group(r) = run id (0..n_experts), rows grouped by run (permuted-by-expert). + // VIEW_SCRAMBLE (default for tiny/ragged): map each run to a NON-MONOTONIC + // slot via a multiplicative permutation — mimics in-model rawG.slotOf, which + // assigns slots in routing order (not sorted). All slots stay in [0,n_experts) + // so they're all filled. This is the unreplicated in-model condition. + let scramble = std::env::var("VIEW_NOSCRAMBLE").is_err() && (tiny || ragged); + let indices: Vec = if std::env::var("VIEW_INMODEL").is_ok() { + // EXACT in-model L0 gate slotGIdx (72 rows, 47 experts, runs up to len 4). + [ + 16, 9, 9, 40, 40, 4, 21, 15, 15, 5, 0, 0, 42, 11, 1, 1, 29, 41, 41, 18, 14, 14, 23, 7, + 7, 17, 17, 39, 39, 39, 39, 43, 27, 31, 12, 10, 10, 24, 24, 13, 13, 19, 19, 19, 8, 35, + 28, 44, 25, 25, 26, 6, 20, 33, 33, 3, 3, 3, 36, 30, 30, 38, 45, 45, 32, 2, 2, 22, 37, + 37, 34, 46, + ] + .iter() + .map(|&x| x as u32) + .collect() + } else if let Ok(s) = std::env::var("VIEW_SINGLE_E") { + let e: u32 = s.parse().unwrap(); + vec![e; m_total] // all rows → ONE expert (isolate per-expert offset) + } else { + (0..m_total) + .map(|r| { + let g = r * n_experts / m_total; + let slot = if scramble { (g * 7) % n_experts } else { g }; + slot as u32 + }) + .collect() + }; + + // POOL bm64 + let ctx = Context::new().unwrap(); + let run = + |buffers: BTreeMap>, kernel_ir: metaltile_core::ir::Kernel| -> Vec { + let mut k = kernel_ir; + k.mode = KernelMode::Reduction; + let r = ctx + .dispatch_with_grid( + &k, + &buffers, + &BTreeMap::new(), + [n_out / 64, m_total.div_ceil(64), 1], + [128, 1, 1], + ) + .expect("dispatch"); + unpack_bytes(r.outputs.get("out").unwrap(), dt) + }; + let mut pb: BTreeMap> = BTreeMap::new(); + pb.insert("x".into(), pack_bytes(&x, dt)); + pb.insert("qs".into(), pack_u32_bytes(&qs)); + pb.insert("d_f32".into(), pack_bytes(&d, Dt::F32)); + pb.insert("grid".into(), grid.clone()); + pb.insert("signs".into(), signs.clone()); + pb.insert("indices".into(), pack_u32_bytes(&indices)); + pb.insert("out".into(), vec![0u8; m_total * n_out * 4]); + pb.insert("m_total".into(), (m_total as u32).to_le_bytes().to_vec()); + pb.insert("n_out".into(), (n_out as u32).to_le_bytes().to_vec()); + pb.insert("k_in".into(), (k_in as u32).to_le_bytes().to_vec()); + let pool_out = run(pb, ffai_moe_bgemm_iq2xxs_bm64::kernel_ir_for(DType::F32)); + + // VIEW-u16 bm64 + let raw_bytes: Vec = raw_u16.iter().flat_map(|v| v.to_le_bytes()).collect(); + let mut vb: BTreeMap> = BTreeMap::new(); + vb.insert("x".into(), pack_bytes(&x, dt)); + vb.insert("view_u16".into(), raw_bytes.clone()); + vb.insert("view_f16".into(), raw_bytes); + vb.insert("grid".into(), grid); + vb.insert("signs".into(), signs); + vb.insert("indices".into(), pack_u32_bytes(&indices)); + vb.insert("out".into(), vec![0u8; m_total * n_out * 4]); + vb.insert("m_total".into(), (m_total as u32).to_le_bytes().to_vec()); + vb.insert("n_out".into(), (n_out as u32).to_le_bytes().to_vec()); + vb.insert("k_in".into(), (k_in as u32).to_le_bytes().to_vec()); + vb.insert("tensor_byte_off".into(), 0u32.to_le_bytes().to_vec()); + vb.insert("expert_byte_stride".into(), ((nblk * 66) as u32).to_le_bytes().to_vec()); + let view_out = run(vb, ffai_moe_bgemm_iq2xxs_view_u16_bm64::kernel_ir_for(DType::F32)); + + let mut worst = 0.0f32; + let mut wi = 0; + for i in 0..pool_out.len() { + let dd = (pool_out[i] - view_out[i]).abs(); + if dd > worst { + worst = dd; + wi = i; + } + } + eprintln!( + "[view-u16] worst |pool-view|={worst:.4e} @i={wi} pool={:.4} view={:.4}; pool[0..4]={:?} view[0..4]={:?}", + pool_out[wi], + view_out[wi], + &pool_out[0..4], + &view_out[0..4] + ); + // ABSOLUTE row-coverage check: the view-vs-pool diff masks bugs SHARED by + // both kernels (e.g. a run-detection that drops rows). Count nonzero rows + // in EACH independently — both must cover all M rows. + let rowzero = |out: &[f32]| -> usize { + (0..m_total).filter(|&r| (0..n_out).all(|c| out[r * n_out + c].abs() < 1e-6)).count() + }; + let pz = rowzero(&pool_out); + let vz = rowzero(&view_out); + eprintln!("[view-u16] zeroRows pool={pz}/{m_total} view={vz}/{m_total}"); + let mag: f32 = pool_out.iter().map(|v| v.abs()).fold(0.0, f32::max).max(1.0); + assert!( + worst < mag * 1e-3, + "view-u16 bm64 diverges from pool bm64: worst={worst:.4e} (mag {mag:.2})" + ); + assert_eq!( + vz, 0, + "view-u16 bm64 left {vz}/{m_total} rows ZERO (run-detection drops rows — masked by view==pool)" + ); + assert_eq!(pz, 0, "pool bm64 left {pz}/{m_total} rows ZERO"); +} diff --git a/crates/metaltile-std/tests/sdpa_decode_d512_sink_nkv1.rs b/crates/metaltile-std/tests/sdpa_decode_d512_sink_nkv1.rs new file mode 100644 index 00000000..c0e634ca --- /dev/null +++ b/crates/metaltile-std/tests/sdpa_decode_d512_sink_nkv1.rs @@ -0,0 +1,91 @@ +//! Copyright 2026 TheTom +//! SPDX-License-Identifier: Apache-2.0 +//! Isolated check: ffai_sdpa_decode_d512_sink with n_kv=1 (single visible +//! KV — the first decode token), many q heads (MQA), per-head sink. The +//! in-source test only covers n_kv=64; DSv4 token 0 hits n_kv=1. +#![cfg(target_os = "macos")] + +mod common; + +use std::collections::BTreeMap; + +use common::{Dt, gpu_lock, pack_bytes, unpack_bytes}; +use metaltile_core::ir::KernelMode; +use metaltile_runtime::Context; +use metaltile_std::ffai::sdpa_decode_d512_sink::ffai_sdpa_decode_d512_sink; + +fn xorshift(s: &mut u32) -> u32 { + let mut x = *s; + x ^= x << 13; + x ^= x >> 17; + x ^= x << 5; + *s = x; + x +} +fn frand(s: &mut u32) -> f32 { (xorshift(s) as f32 / u32::MAX as f32) * 0.2 - 0.1 } + +#[test] +fn sdpa_decode_d512_sink_nkv1_64heads() { + let _g = gpu_lock(); + let dt = Dt::F32; + let hd = 512usize; + let n_q = 64usize; + let n_kv = 1usize; + let kv_stride = 128usize; // like the real sliding-window cache + let heads_per_group = n_q; // MQA: 1 kv head + let scale = 1.0f32 / (hd as f32).sqrt(); + let mut st = 0x5D9A_2026u32; + + let q: Vec = (0..n_q * hd).map(|_| frand(&mut st)).collect(); + let k: Vec = (0..kv_stride * hd).map(|_| frand(&mut st)).collect(); + let v: Vec = (0..kv_stride * hd).map(|_| frand(&mut st)).collect(); + let sink: Vec = (0..n_q).map(|_| frand(&mut st)).collect(); + + // Oracle. + let mut want = vec![0.0f32; n_q * hd]; + #[allow(clippy::needless_range_loop)] + for h in 0..n_q { + let qo = h * hd; + let mut s = 0.0f32; + for d in 0..hd { + s += q[qo + d] * scale * k[d]; + } // single key at row 0 + let m = s.max(sink[h]); + let denom = (s - m).exp() + (sink[h] - m).exp(); + let w = (s - m).exp() / denom; + for d in 0..hd { + want[qo + d] = w * v[d]; + } + } + + let mut b: BTreeMap> = BTreeMap::new(); + b.insert("q".into(), pack_bytes(&q, dt)); + b.insert("k".into(), pack_bytes(&k, dt)); + b.insert("v".into(), pack_bytes(&v, dt)); + b.insert("sink_logit".into(), pack_bytes(&sink, Dt::F32)); + b.insert("out".into(), pack_bytes(&vec![0.0f32; n_q * hd], dt)); + b.insert("head_dim".into(), (hd as u32).to_le_bytes().to_vec()); + b.insert("n_kv".into(), (n_kv as u32).to_le_bytes().to_vec()); + b.insert("kv_stride".into(), (kv_stride as u32).to_le_bytes().to_vec()); + b.insert("heads_per_group".into(), (heads_per_group as u32).to_le_bytes().to_vec()); + b.insert("scale".into(), scale.to_le_bytes().to_vec()); + + let ctx = Context::new().unwrap(); + let mut kern = ffai_sdpa_decode_d512_sink::kernel_ir_for(dt.to_dtype()); + kern.mode = KernelMode::Reduction; + let r = ctx.dispatch_with_grid(&kern, &b, &BTreeMap::new(), [n_q, 1, 1], [512, 1, 1]).unwrap(); + let got = unpack_bytes(r.outputs.get("out").unwrap(), dt); + + let mut max_rel = 0.0f32; + for i in 0..(n_q * hd) { + let denom = want[i].abs().max(0.05); + let rel = (got[i] - want[i]).abs() / denom; + if rel > max_rel { + max_rel = rel; + } + } + eprintln!("nkv1 64head max_rel={max_rel}"); + // Spot-check head 5 specifically (the FFAI case that zeroed). + eprintln!("head5: got={:?} want={:?}", &got[5 * hd..5 * hd + 4], &want[5 * hd..5 * hd + 4]); + assert!(max_rel < 1e-3, "max_rel={max_rel}"); +} diff --git a/crates/metaltile-std/tests/sdpa_prefill_d512_sink_correctness.rs b/crates/metaltile-std/tests/sdpa_prefill_d512_sink_correctness.rs new file mode 100644 index 00000000..4461f284 --- /dev/null +++ b/crates/metaltile-std/tests/sdpa_prefill_d512_sink_correctness.rs @@ -0,0 +1,120 @@ +//! Copyright 2026 TheTom +//! SPDX-License-Identifier: Apache-2.0 +//! GPU correctness for `ffai::ffai_sdpa_prefill_d512_sink` — multi-query +//! causal sliding-window SDPA (d512, attn sink, MQA). Oracle: per-(q_pos, +//! q_head) causal softmax with sink over the KV window. +#![cfg(target_os = "macos")] + +mod common; + +use std::collections::BTreeMap; + +use common::{Dt, gpu_lock, pack_bytes, unpack_bytes}; +use metaltile_core::ir::KernelMode; +use metaltile_runtime::Context; +use metaltile_std::ffai::sdpa_prefill_d512_sink::ffai_sdpa_prefill_d512_sink; + +fn xorshift(s: &mut u32) -> u32 { + let mut x = *s; + x ^= x << 13; + x ^= x >> 17; + x ^= x << 5; + *s = x; + x +} +fn frand(s: &mut u32) -> f32 { (xorshift(s) as f32 / u32::MAX as f32) * 0.2 - 0.1 } + +fn run_case(dt: Dt, n_query: usize, kv: usize, window: usize, tol: f32) { + run_case_nq(dt, 4, n_query, kv, window, tol); +} + +fn run_case_nq(dt: Dt, n_q: usize, n_query: usize, kv: usize, window: usize, tol: f32) { + let _g = gpu_lock(); + let hd = 512usize; + let heads_per_group = n_q; // MQA: 1 kv head + let scale = 1.0f32 / (hd as f32).sqrt(); + let mut st = 0x5D9A_2026u32; + + let q: Vec = (0..n_query * n_q * hd).map(|_| frand(&mut st)).collect(); + let k: Vec = (0..kv * hd).map(|_| frand(&mut st)).collect(); + let v: Vec = (0..kv * hd).map(|_| frand(&mut st)).collect(); + let sink: Vec = (0..n_q).map(|_| frand(&mut st)).collect(); + + // Oracle (kv_base = 0). + let mut want = vec![0.0f32; n_query * n_q * hd]; + for qp in 0..n_query { + let p1 = qp + 1; + let lo = p1.saturating_sub(window); + #[allow(clippy::needless_range_loop)] + for h in 0..n_q { + let qo = (qp * n_q + h) * hd; + let mut scores = vec![0.0f32; p1 - lo]; + for (i, t) in (lo..p1).enumerate() { + let mut s = 0.0f32; + for d in 0..hd { + s += q[qo + d] * scale * k[t * hd + d]; + } + scores[i] = s; + } + let mut m = sink[h]; + for &s in &scores { + if s > m { + m = s; + } + } + let mut denom = (sink[h] - m).exp(); + for &s in &scores { + denom += (s - m).exp(); + } + for d in 0..hd { + let mut acc = 0.0f32; + for (i, t) in (lo..p1).enumerate() { + acc += ((scores[i] - m).exp() / denom) * v[t * hd + d]; + } + want[qo + d] = acc; + } + } + } + + let mut b: BTreeMap> = BTreeMap::new(); + b.insert("q".into(), pack_bytes(&q, dt)); + b.insert("k".into(), pack_bytes(&k, dt)); + b.insert("v".into(), pack_bytes(&v, dt)); + b.insert("sink_logit".into(), pack_bytes(&sink, Dt::F32)); + b.insert("out".into(), pack_bytes(&vec![0.0f32; n_query * n_q * hd], dt)); + b.insert("head_dim".into(), (hd as u32).to_le_bytes().to_vec()); + b.insert("n_q_heads".into(), (n_q as u32).to_le_bytes().to_vec()); + b.insert("kv_stride".into(), (kv as u32).to_le_bytes().to_vec()); + b.insert("heads_per_group".into(), (heads_per_group as u32).to_le_bytes().to_vec()); + b.insert("window".into(), (window as u32).to_le_bytes().to_vec()); + b.insert("kv_base".into(), 0u32.to_le_bytes().to_vec()); + b.insert("scale".into(), scale.to_le_bytes().to_vec()); + + let ctx = Context::new().unwrap(); + let mut kern = ffai_sdpa_prefill_d512_sink::kernel_ir_for(dt.to_dtype()); + kern.mode = KernelMode::Reduction; + let r = + ctx.dispatch_with_grid(&kern, &b, &BTreeMap::new(), [n_q, n_query, 1], [32, 1, 1]).unwrap(); + let got = unpack_bytes(r.outputs.get("out").unwrap(), dt); + + for i in 0..(n_query * n_q * hd) { + let denom = want[i].abs().max(0.05); + assert!((got[i] - want[i]).abs() / denom < tol, "i={i}: got={} want={}", got[i], want[i]); + } +} + +#[test] +fn sdpa_prefill_d512_sink_f32() { + run_case(Dt::F32, 8, 8, 128, 1e-3); // full causal (window > kv) + run_case(Dt::F32, 40, 40, 16, 1e-3); // sliding window active +} + +#[test] +fn sdpa_prefill_d512_sink_single_token_64heads() { + // DSv4 decode-parity: 1 query, 1 KV, 64 q heads (MQA), per-head sink. + run_case_nq(Dt::F32, 64, 1, 1, 128, 1e-3); + run_case_nq(Dt::F32, 64, 4, 4, 128, 1e-3); +} + +#[test] +fn sdpa_prefill_d512_sink_f16() { run_case(Dt::F16, 8, 8, 128, 2e-2); } diff --git a/crates/metaltile/src/runner/gpu.rs b/crates/metaltile/src/runner/gpu.rs index 3379d22f..0b26c3da 100644 --- a/crates/metaltile/src/runner/gpu.rs +++ b/crates/metaltile/src/runner/gpu.rs @@ -560,7 +560,7 @@ pub fn buffer_typed(runner: &GpuRunner, vals: &[f32], dt: DType) -> GpuBuffer { DType::I8 => runner.buffer_bytes(&vals.iter().map(|&v| v as i8 as u8).collect::>()), DType::U8 => runner.buffer_bytes(&vals.iter().map(|&v| v as u8).collect::>()), DType::Bool => runner.buffer_f32(vals), - DType::I4 | DType::U64 | DType::I64 => { + DType::I4 | DType::U16 | DType::U64 | DType::I64 => { unimplemented!("buffer_typed: unsupported dtype {dt:?}") }, } @@ -598,7 +598,7 @@ pub fn read_typed(runner: &GpuRunner, buf: &GpuBuffer, n: usize, dt: DType) -> V bytes.iter().map(|&b| b as f32).collect() }, DType::Bool => runner.read_f32_slice(buf, n), - DType::I4 | DType::U64 | DType::I64 => { + DType::I4 | DType::U16 | DType::U64 | DType::I64 => { unimplemented!("read_typed: unsupported dtype {dt:?}") }, }