feat(kernel): int8 coopmat MMQ GEMM lowers to Vulkan/Metal/ROCm, not just CUDA
The MMQ GEMM was targets(cuda) and read as though SPIR-V could not carry it. Nothing in cubecl blocked it: RADV gfx1151 advertises M16 N16 K16 / A=i8 B=i8 C=i32 R=i32 at subgroup scope, exactly the shape the kernel asks for. The blocker was ours. A cooperative-matrix op is plane-scoped: the fragment's 256 elements are spread across every lane of the plane, so a block narrower than the plane leaves the absent lanes' elements unwritten. The block was a literal 32 -- the CUDA warp. RADV's plane is 64, so half the tile was never written. The failure said so precisely: rows 0,1,4,5,8,9,12,13 bit-exact and the rest zero -- absence, not corruption. The block now comes from plane_size_max and partitions the fragment by it. One kernel source, one island, no fork. The manual-MMA family stays CUDA/ROCm only, and that is a real SPIR-V limit: its cooperative matrix is an opaque type with no queryable per-lane layout, so the explicit-register path cannot be expressed. Gated on the CPU oracle (35/35) and, on RADV, WMMA hello exact at 256/256 and MMQ within 1.6e-7 of the oracle at three tile shapes. The quant.rs comments claimed comptime shapes gave "full unroll". The shipped SPIR-V carried one static OpSDot; cubecl does not unroll a comptime-bounded loop without #[unroll]. Adding it multiplied the OpSDot count as advertised and moved decode by +0.38% -- noise, on a bandwidth-bound path -- so the claim is dropped from the comments rather than the loop from the kernel.
This commit is contained in:
Hanzo Dev
@@ -27,6 +27,10 @@
|
||||
//! };
|
||||
//! ```
|
||||
//!
|
||||
//! An island sits in either position: bound (`let x = island! { ... };`) when the arms yield a value,
|
||||
//! or bare (`island! { ... };`) when they act on shared memory — a warp-cooperative tile op has no
|
||||
//! per-lane result to bind, so the bare form is the natural one for it.
|
||||
//!
|
||||
//! The scrutinee is found BY TYPE: the kernel must take one `#[comptime] <name>: Target` parameter
|
||||
//! (any name), which the launch layer fills from the runtime — so the branch site names no target and
|
||||
//! the selection stays out of the authoring surface. `default` is mandatory and defines the oracle:
|
||||
@@ -41,7 +45,7 @@ use syn::{
|
||||
parse_macro_input, parse_quote,
|
||||
punctuated::Punctuated,
|
||||
visit_mut::{self, VisitMut},
|
||||
Expr, FnArg, Ident, ItemFn, Pat, Token, Type,
|
||||
Expr, ExprMacro, FnArg, Ident, ItemFn, Pat, Stmt, Token, Type,
|
||||
};
|
||||
|
||||
/// The backends a kernel can name. Each maps 1:1 to a runtime the same source lowers to.
|
||||
@@ -148,6 +152,27 @@ struct IslandRewrite {
|
||||
}
|
||||
|
||||
impl VisitMut for IslandRewrite {
|
||||
/// `island! { ... }` at STATEMENT position — the form an island takes when its arms act on shared
|
||||
/// memory instead of yielding a value (a warp-cooperative tile op has no per-lane result to bind).
|
||||
/// syn parses a braced macro statement as `Stmt::Macro`, never `Stmt::Expr`, so `visit_expr_mut`
|
||||
/// alone would never see it. Re-seat it as the equivalent expression statement and rewrite that,
|
||||
/// preserving the original trailing semicolon (or its absence, when the island is the block's value).
|
||||
fn visit_stmt_mut(&mut self, stmt: &mut Stmt) {
|
||||
if let Stmt::Macro(sm) = stmt {
|
||||
if sm.mac.path.is_ident("island") {
|
||||
let mut expr = Expr::Macro(ExprMacro {
|
||||
attrs: sm.attrs.clone(),
|
||||
mac: sm.mac.clone(),
|
||||
});
|
||||
let semi = sm.semi_token;
|
||||
self.visit_expr_mut(&mut expr);
|
||||
*stmt = Stmt::Expr(expr, semi);
|
||||
return;
|
||||
}
|
||||
}
|
||||
visit_mut::visit_stmt_mut(self, stmt);
|
||||
}
|
||||
|
||||
fn visit_expr_mut(&mut self, expr: &mut Expr) {
|
||||
if let Expr::Macro(mac) = expr {
|
||||
if mac.mac.path.is_ident("island") {
|
||||
|
||||
@@ -48,4 +48,3 @@ required-features = ["vulkan"]
|
||||
[[bin]]
|
||||
name = "mmq-check"
|
||||
path = "src/mmq_check.rs"
|
||||
required-features = ["cuda"]
|
||||
|
||||
@@ -577,14 +577,19 @@ fn main() {
|
||||
let (rows, k) = (4096usize, 4096usize);
|
||||
let ctrl = 256usize; // small-K control: reorder noise ~ ctrl*eps, should be ~1e-6
|
||||
|
||||
// `matvec-check dump` dispatches exactly the engine-bound MoE block kernels at their evo-optimal
|
||||
// baked nt (Q4_K 64, Q6_K 32) so `CUBECL_DEBUG_SPIRV=<dir>` writes exactly one .spv per kernel --
|
||||
// the committed hanzo-ml Vulkan artifact. Build with `--features vulkan,spirv-dump`.
|
||||
// `matvec-check dump` dispatches exactly the engine-bound MoE block kernels for each LIVE model
|
||||
// shape at its evo-optimal power-of-2 nt, so `CUBECL_DEBUG_SPIRV=<dir>` writes one .spv per (shape).
|
||||
// The DSL source is ONE fn per quant; comptime lowers it to a fast specialized artifact per shape
|
||||
// (like llama's templated kernels). Disambiguate by LocalSize (nt): q4k 64 = gate/up, 8 = down.
|
||||
// Qwen3-30B-A3B Q4_K_M: ffn_gate/up_exps Q4_K [128,768,2048]; ffn_down_exps Q6_K/Q4_K [128,2048,768].
|
||||
// `matvec-check dump` dispatches the engine-bound kernels once per LIVE model shape at that
|
||||
// shape's baked nt, so `CUBECL_DEBUG_SPIRV=<dir>` writes one .spv per shape. The DSL source is ONE
|
||||
// fn per quant; comptime lowers it to a specialized artifact per shape (like llama's templated
|
||||
// kernels). Disambiguate the two Q4_K MoE shapes by LocalSize (= nt): 64 = gate/up, 32 = down.
|
||||
// Qwen3-30B-A3B Q4_K_M: ffn_gate/up_exps Q4_K [128,768,2048]; ffn_down_exps Q6_K [128,2048,768].
|
||||
// Build with `--features vulkan,spirv-dump`.
|
||||
//
|
||||
// A dumped .spv is NOT yet the committed artifact: CubeCL names the entry point after the kernel,
|
||||
// while the hanzo-ml loader creates every pipeline with entry point `main` (vulkan_backend.rs).
|
||||
// Installing a dump verbatim therefore binds a nonexistent entry point -- undefined behaviour, and
|
||||
// RADV faults rather than erroring. Rename the entry point when installing:
|
||||
// spirv-dis in.spv | sed -E 's/(OpEntryPoint GLCompute %[0-9]+ )"[^"]+"/\1"main"/' \
|
||||
// | spirv-as --target-env vulkan1.2 -o hanzo-ml/src/vulkan/spv/<name>.spv
|
||||
#[cfg(all(feature = "vulkan", feature = "spirv-dump"))]
|
||||
if std::env::args().any(|a| a == "dump") {
|
||||
use cubecl::wgpu::{WgpuDevice, WgpuRuntime};
|
||||
@@ -705,6 +710,9 @@ fn main() {
|
||||
let c = WgpuRuntime::client(&WgpuDevice::default());
|
||||
check::<WgpuRuntime>("METAL", &c, rows, k);
|
||||
check_q4k::<WgpuRuntime>("METAL", &c, rows, k);
|
||||
// Same Qwen3-30B-A3B-shaped MoE decode the Vulkan block gates on, so the DSL's expert-gather is
|
||||
// proven on Metal too and the two backends are compared on identical shapes.
|
||||
check_moe_q4k::<WgpuRuntime>("METAL", &c, 128, 768, 8, 2048);
|
||||
check_dp4a::<WgpuRuntime>("METAL", &c, rows, k, true);
|
||||
}
|
||||
#[cfg(feature = "cuda")]
|
||||
|
||||
+348
-138
@@ -6,13 +6,24 @@
|
||||
//! and lowered to the hardware int8 matrix cores, instead of hand-written four times (CUDA fast_mmq,
|
||||
//! Vulkan mul_mm_q8_mmq, ROCm qmmq, Metal simdgroup).
|
||||
//!
|
||||
//! Two DSL routes to the int8 matrix cores, both registered by CubeCL on CUDA sm>=70/80:
|
||||
//! Two DSL routes to the int8 matrix cores, and they do NOT have the same reach:
|
||||
//! - high-level WMMA (`cmma::Matrix`, i8->i32 at 16x16x16): auto fragment load/store, one tile at a
|
||||
//! time. Ergonomic; the fragment layout is opaque so the per-block dequant-scale is applied via a
|
||||
//! shared-memory round-trip in the epilogue.
|
||||
//! shared-memory round-trip in the epilogue. PORTABLE -- lowers to WMMA on CUDA/ROCm, to
|
||||
//! `OpCooperativeMatrixMulAddKHR` (signed-component operands) on SPIR-V, and to simdgroup matrix on
|
||||
//! Metal. This is the prefill path.
|
||||
//! - manual MMA (`cmma::MmaDefinition`, i8->i32 at m16n8k32): `mma.sync` with explicit registers.
|
||||
//! The fragment element<->lane map IS exposed (`position_of_nth`), so scales could be applied in
|
||||
//! register -- the higher-ceiling path. Included here as the proven-by-CubeCL int8 probe.
|
||||
//! register -- the higher ceiling, but CUDA/ROCm ONLY: SPIR-V's cooperative matrix is an opaque
|
||||
//! type with no queryable per-lane layout, so cubecl-spirv rejects the whole manual family
|
||||
//! (see `mma_hello_i8`). Kept as the tuned peak path, not the portable one.
|
||||
//!
|
||||
//! Portability is carried by ONE island per GEMM, not by a second kernel: the accelerated arm issues
|
||||
//! cmma, the `default` arm expresses the identical int8 contraction as scalar MACs. `default` is what
|
||||
//! the CPU runtime runs (cubecl-cpu rejects every CoopMma op), which is what makes a bit-exact CPU
|
||||
//! oracle possible at all -- it gates the tiling, the shared-memory epilogue and the f32 accumulation
|
||||
//! order. The cmma arm itself is equivalent by construction (int8 accumulation is exact) and is gated
|
||||
//! on-GPU against the same `mmq_q8_ref`.
|
||||
//!
|
||||
//! MMQ math (Q8_0 weight x q8_1 activation), the exact contraction llama does:
|
||||
//! out[m,n] = sum over 32-blocks kb of (xs[m,kb] * wd[n,kb]) * sum_{k in block} xq[m,k]*wq[n,k]
|
||||
@@ -28,7 +39,11 @@ use crate::prelude::*;
|
||||
// ================================================================================================
|
||||
|
||||
/// out[16,16] = A[16,16] @ B[16,16]^T, int8 inputs, i32 accumulate. Single warp, one WMMA tile.
|
||||
#[kernel(targets(cuda), unchecked)]
|
||||
///
|
||||
/// No island: the kernel IS the cmma probe, so there is no meaningful non-cmma arm to fall back to.
|
||||
/// It therefore has no CPU oracle (cubecl-cpu rejects every CoopMma op) and is proven on-GPU against
|
||||
/// [`hello_i8_ref`]. The MMQ GEMMs below carry the oracle; this only answers "does the op JIT and run".
|
||||
#[kernel(targets(cuda, rocm, vulkan, metal), unchecked)]
|
||||
pub fn wmma_hello_i8(a: &Array<i8>, b: &Array<i8>, out: &mut Array<i32>) {
|
||||
let ma = cmma::Matrix::<i8>::from_slice(
|
||||
cmma::MatrixIdent::A,
|
||||
@@ -66,7 +81,13 @@ pub fn wmma_hello_i8(a: &Array<i8>, b: &Array<i8>, out: &mut Array<i32>) {
|
||||
}
|
||||
|
||||
/// Run the WMMA hello on a runtime; returns the 256 i32 outputs.
|
||||
///
|
||||
/// The block is ONE PLANE wide, queried from the runtime -- not a literal 32. A cooperative-matrix op
|
||||
/// is plane-scoped: the fragment's 256 elements are distributed across every lane of the plane, so a
|
||||
/// block narrower than the plane leaves the absent lanes' elements unwritten. The plane is 32 lanes on
|
||||
/// CUDA but 64 on RADV, which is why a hardcoded 32 silently produces a half-written tile there.
|
||||
pub fn wmma_hello_i8_run<R: Runtime>(client: &ComputeClient<R>, a: &[i8], b: &[i8]) -> Vec<i32> {
|
||||
let plane = client.properties().hardware.plane_size_max;
|
||||
let ah = client.create_from_slice(i8::as_bytes(a));
|
||||
let bh = client.create_from_slice(i8::as_bytes(b));
|
||||
let oh = client.create_from_slice(i32::as_bytes(&vec![0i32; 256]));
|
||||
@@ -74,7 +95,7 @@ pub fn wmma_hello_i8_run<R: Runtime>(client: &ComputeClient<R>, a: &[i8], b: &[i
|
||||
wmma_hello_i8::launch_unchecked::<R>(
|
||||
client,
|
||||
Grid::Static(1, 1, 1),
|
||||
Block::new_1d(32),
|
||||
Block::new_1d(plane),
|
||||
ArrayArg::from_raw_parts(ah.clone(), 256),
|
||||
ArrayArg::from_raw_parts(bh.clone(), 256),
|
||||
ArrayArg::from_raw_parts(oh.clone(), 256),
|
||||
@@ -105,7 +126,16 @@ pub fn hello_i8_ref(a: &[i8], b: &[i8]) -> Vec<i32> {
|
||||
// ================================================================================================
|
||||
|
||||
/// out[16,8] = A[16,32] @ B[32,8] over int8, i32 accumulate, via one `mma.sync.m16n8k32.s8`.
|
||||
#[kernel(targets(cuda), unchecked)]
|
||||
///
|
||||
/// NOT portable to Vulkan, and not by omission: manual MMA needs the fragment element<->lane map, and
|
||||
/// cubecl-spirv rejects every op that exposes it -- `RowIndex`/`ColIndex` (what `position_of_nth`
|
||||
/// emits), `LoadMatrix`/`StoreMatrix`, `ExecuteManual`, `ExecuteScaled` all hit one
|
||||
/// `panic!("Manual register management not currently supported in SPIR-V")`
|
||||
/// (cubecl-spirv-0.10.0/src/cmma.rs:37-44). SPIR-V's CooperativeMatrixKHR type is opaque by design:
|
||||
/// it has no defined per-lane layout to query, so this is a spec-level gap, not a missing match arm.
|
||||
/// The cubecl-cpp backends do implement it (cuda/dialect.rs, hip/dialect.rs:171, metal/dialect.rs:1219),
|
||||
/// hence cuda + rocm. The portable prefill path is the high-level WMMA route below.
|
||||
#[kernel(targets(cuda, rocm), unchecked)]
|
||||
pub fn mma_hello_i8(
|
||||
a: &Array<i8>, // [16,32] row-major
|
||||
b: &Array<i8>, // [32,8] row-major
|
||||
@@ -233,8 +263,13 @@ pub fn mma_hello_ref(a: &[i8], b: &[i8]) -> Vec<i32> {
|
||||
// fragment, stores the i32 tile to shared memory, then applies xs*wd per element into an f32 tile.
|
||||
// ================================================================================================
|
||||
|
||||
/// int8 tensor-core MMQ GEMM. Grid = (N/16, M/16); block = one warp (32 lanes).
|
||||
#[kernel(targets(cuda), unchecked)]
|
||||
/// int8 tensor-core MMQ GEMM. Grid = (N/16, M/16); block = ONE PLANE.
|
||||
///
|
||||
/// `plane` is comptime because a cooperative-matrix op is plane-scoped: the 16x16 fragment's 256
|
||||
/// elements are spread over every lane of the plane, so the block must be exactly one plane wide and
|
||||
/// the scalar prologue/epilogue must partition 256 by that same width (8 each on a 32-lane CUDA warp,
|
||||
/// 4 each on a 64-lane RADV wave). Hardcoding 32 leaves the absent lanes' elements unwritten.
|
||||
#[kernel(targets(cuda, rocm, vulkan, metal, cpu), unchecked)]
|
||||
pub fn mmq_q8_wmma(
|
||||
xq: &Array<i8>,
|
||||
xs: &Array<f32>,
|
||||
@@ -244,81 +279,113 @@ pub fn mmq_q8_wmma(
|
||||
#[comptime] m: usize,
|
||||
#[comptime] n: usize,
|
||||
#[comptime] k: usize,
|
||||
#[comptime] plane: usize,
|
||||
#[comptime] target: Target,
|
||||
) {
|
||||
let nt = CUBE_POS_X as usize; // output column-tile (16 cols of N)
|
||||
let mt = CUBE_POS_Y as usize; // output row-tile (16 rows of M)
|
||||
let lane = UNIT_POS as usize; // 0..31
|
||||
let lane = UNIT_POS as usize; // 0..plane-1
|
||||
let kb_count = k / 32;
|
||||
let mrow0 = mt * 16;
|
||||
let ncol0 = nt * 16;
|
||||
let per = 256 / plane; // tile elements this lane owns in the scalar prologue/epilogue
|
||||
|
||||
let mut accf = SharedMemory::<f32>::new(256usize); // f32 output tile [16x16]
|
||||
let mut ci = SharedMemory::<i32>::new(256usize); // i32 fragment store scratch [16x16]
|
||||
let mut ci = SharedMemory::<i32>::new(256usize); // i32 tile: the 32-block's int8 contraction
|
||||
|
||||
// Zero the f32 accumulator (each lane owns 8 of 256).
|
||||
// Zero the f32 accumulator (each lane owns `per` of 256).
|
||||
#[unroll]
|
||||
for e in 0usize..8 {
|
||||
accf[lane * 8 + e] = 0.0f32;
|
||||
for e in 0usize..per {
|
||||
accf[lane * per + e] = 0.0f32;
|
||||
}
|
||||
sync_cube();
|
||||
|
||||
for kb in 0..kb_count {
|
||||
let k0 = kb * 32;
|
||||
// Accumulate int8 over the 32-wide block into one i32 fragment (2 WMMA K=16 steps).
|
||||
let c = cmma::Matrix::<i32>::from_value(
|
||||
cmma::MatrixIdent::Accumulator,
|
||||
16usize,
|
||||
16usize,
|
||||
16usize,
|
||||
cmma::MatrixLayout::Undefined,
|
||||
0i32,
|
||||
);
|
||||
let a0 = cmma::Matrix::<i8>::from_slice(
|
||||
cmma::MatrixIdent::A,
|
||||
16usize,
|
||||
16usize,
|
||||
16usize,
|
||||
cmma::MatrixLayout::RowMajor,
|
||||
&xq.slice(mrow0 * k + k0, m * k),
|
||||
k as u32,
|
||||
);
|
||||
let b0 = cmma::Matrix::<i8>::from_slice(
|
||||
cmma::MatrixIdent::B,
|
||||
16usize,
|
||||
16usize,
|
||||
16usize,
|
||||
cmma::MatrixLayout::ColMajor,
|
||||
&wq.slice(ncol0 * k + k0, n * k),
|
||||
k as u32,
|
||||
);
|
||||
cmma::execute::<i8, i8, i32, i32>(&a0, &b0, &c, &c);
|
||||
// The 32-wide K-block's int8 contraction into the i32 tile `ci`. int8 accumulation is exact,
|
||||
// so both arms are the SAME integer value -- only the idiom differs.
|
||||
island! {
|
||||
// Accelerated: the hardware int8 matrix core, 2 x cmma(16x16x16) i8->i32. Lowers to WMMA
|
||||
// (CUDA/ROCm), OpCooperativeMatrixMulAddKHR with signed-component operands (SPIR-V), and
|
||||
// simdgroup matrix (Metal). The fragment layout is opaque, so the tile is spilled to
|
||||
// shared memory for the scale epilogue below.
|
||||
cuda | rocm | vulkan | metal => {
|
||||
let c = cmma::Matrix::<i32>::from_value(
|
||||
cmma::MatrixIdent::Accumulator,
|
||||
16usize,
|
||||
16usize,
|
||||
16usize,
|
||||
cmma::MatrixLayout::Undefined,
|
||||
0i32,
|
||||
);
|
||||
let a0 = cmma::Matrix::<i8>::from_slice(
|
||||
cmma::MatrixIdent::A,
|
||||
16usize,
|
||||
16usize,
|
||||
16usize,
|
||||
cmma::MatrixLayout::RowMajor,
|
||||
&xq.slice(mrow0 * k + k0, m * k),
|
||||
k as u32,
|
||||
);
|
||||
let b0 = cmma::Matrix::<i8>::from_slice(
|
||||
cmma::MatrixIdent::B,
|
||||
16usize,
|
||||
16usize,
|
||||
16usize,
|
||||
cmma::MatrixLayout::ColMajor,
|
||||
&wq.slice(ncol0 * k + k0, n * k),
|
||||
k as u32,
|
||||
);
|
||||
cmma::execute::<i8, i8, i32, i32>(&a0, &b0, &c, &c);
|
||||
|
||||
let a1 = cmma::Matrix::<i8>::from_slice(
|
||||
cmma::MatrixIdent::A,
|
||||
16usize,
|
||||
16usize,
|
||||
16usize,
|
||||
cmma::MatrixLayout::RowMajor,
|
||||
&xq.slice(mrow0 * k + k0 + 16, m * k),
|
||||
k as u32,
|
||||
);
|
||||
let b1 = cmma::Matrix::<i8>::from_slice(
|
||||
cmma::MatrixIdent::B,
|
||||
16usize,
|
||||
16usize,
|
||||
16usize,
|
||||
cmma::MatrixLayout::ColMajor,
|
||||
&wq.slice(ncol0 * k + k0 + 16, n * k),
|
||||
k as u32,
|
||||
);
|
||||
cmma::execute::<i8, i8, i32, i32>(&a1, &b1, &c, &c);
|
||||
let a1 = cmma::Matrix::<i8>::from_slice(
|
||||
cmma::MatrixIdent::A,
|
||||
16usize,
|
||||
16usize,
|
||||
16usize,
|
||||
cmma::MatrixLayout::RowMajor,
|
||||
&xq.slice(mrow0 * k + k0 + 16, m * k),
|
||||
k as u32,
|
||||
);
|
||||
let b1 = cmma::Matrix::<i8>::from_slice(
|
||||
cmma::MatrixIdent::B,
|
||||
16usize,
|
||||
16usize,
|
||||
16usize,
|
||||
cmma::MatrixLayout::ColMajor,
|
||||
&wq.slice(ncol0 * k + k0 + 16, n * k),
|
||||
k as u32,
|
||||
);
|
||||
cmma::execute::<i8, i8, i32, i32>(&a1, &b1, &c, &c);
|
||||
|
||||
// Spill the i32 tile to shared memory, then apply the per-block f32 scale in the epilogue.
|
||||
cmma::store(&mut ci.to_slice_mut(), &c, 16, cmma::MatrixLayout::RowMajor);
|
||||
cmma::store(&mut ci.to_slice_mut(), &c, 16, cmma::MatrixLayout::RowMajor);
|
||||
}
|
||||
// NORMATIVE oracle: the identical i32 contraction as portable scalar MACs, each lane
|
||||
// owning 8 of the tile's 256 elements. This arm carries no cmma op, so it is the arm the
|
||||
// CPU runtime executes -- and thus the bit-exact gate for the whole kernel's structure
|
||||
// (tiling, shared-memory epilogue, f32 accumulation order).
|
||||
default => {
|
||||
#[unroll]
|
||||
for e in 0usize..per {
|
||||
let idx = lane * per + e;
|
||||
let mm = idx / 16;
|
||||
let nn = idx % 16;
|
||||
let mut isum = 0i32;
|
||||
for l in 0usize..32 {
|
||||
isum += i32::cast_from(xq[(mrow0 + mm) * k + k0 + l])
|
||||
* i32::cast_from(wq[(ncol0 + nn) * k + k0 + l]);
|
||||
}
|
||||
ci[idx] = isum;
|
||||
}
|
||||
}
|
||||
};
|
||||
sync_cube();
|
||||
|
||||
// Epilogue: apply the per-block f32 scale. Shared by both arms -- the scale is a property of
|
||||
// block-quantized GEMM, not of the idiom that produced the int8 sum.
|
||||
#[unroll]
|
||||
for e in 0usize..8 {
|
||||
let idx = lane * 8 + e;
|
||||
for e in 0usize..per {
|
||||
let idx = lane * per + e;
|
||||
let mm = idx / 16;
|
||||
let nn = idx % 16;
|
||||
let xsc = xs[(mrow0 + mm) * kb_count + kb];
|
||||
@@ -329,15 +396,16 @@ pub fn mmq_q8_wmma(
|
||||
}
|
||||
|
||||
#[unroll]
|
||||
for e in 0usize..8 {
|
||||
let idx = lane * 8 + e;
|
||||
for e in 0usize..per {
|
||||
let idx = lane * per + e;
|
||||
let mm = idx / 16;
|
||||
let nn = idx % 16;
|
||||
out[(mrow0 + mm) * n + (ncol0 + nn)] = accf[idx];
|
||||
}
|
||||
}
|
||||
|
||||
/// Host launch for the MMQ GEMM. `iters` amortized kernel-only bench dispatches; returns (out, ms).
|
||||
/// Host launch for the MMQ GEMM -- the PRODUCTION surface. Derives the island tag from the runtime and
|
||||
/// hides it: callers pass only data. `iters` amortized kernel-only bench dispatches; returns (out, ms).
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub fn mmq_q8_wmma_run<R: Runtime>(
|
||||
client: &ComputeClient<R>,
|
||||
@@ -350,6 +418,25 @@ pub fn mmq_q8_wmma_run<R: Runtime>(
|
||||
k: usize,
|
||||
iters: usize,
|
||||
) -> (Vec<f32>, f64) {
|
||||
mmq_q8_wmma_run_with(client, xq, xs, wq, wd, m, n, k, iters, Target::of(client))
|
||||
}
|
||||
|
||||
/// Host launch with an explicit island tag. The oracle test pins [`Target::Cpu`] to run the normative
|
||||
/// arm; production goes through [`mmq_q8_wmma_run`], which derives the tag from the runtime.
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub fn mmq_q8_wmma_run_with<R: Runtime>(
|
||||
client: &ComputeClient<R>,
|
||||
xq: &[i8],
|
||||
xs: &[f32],
|
||||
wq: &[i8],
|
||||
wd: &[f32],
|
||||
m: usize,
|
||||
n: usize,
|
||||
k: usize,
|
||||
iters: usize,
|
||||
target: Target,
|
||||
) -> (Vec<f32>, f64) {
|
||||
let plane = client.properties().hardware.plane_size_max;
|
||||
let xqh = client.create_from_slice(i8::as_bytes(xq));
|
||||
let xsh = client.create_from_slice(f32::as_bytes(xs));
|
||||
let wqh = client.create_from_slice(i8::as_bytes(wq));
|
||||
@@ -360,7 +447,7 @@ pub fn mmq_q8_wmma_run<R: Runtime>(
|
||||
mmq_q8_wmma::launch_unchecked::<R>(
|
||||
c,
|
||||
grid.clone(),
|
||||
Block::new_1d(32),
|
||||
Block::new_1d(plane),
|
||||
ArrayArg::from_raw_parts(xqh.clone(), xq.len()),
|
||||
ArrayArg::from_raw_parts(xsh.clone(), xs.len()),
|
||||
ArrayArg::from_raw_parts(wqh.clone(), wq.len()),
|
||||
@@ -369,6 +456,8 @@ pub fn mmq_q8_wmma_run<R: Runtime>(
|
||||
m,
|
||||
n,
|
||||
k,
|
||||
plane as usize,
|
||||
target,
|
||||
);
|
||||
};
|
||||
launch(client);
|
||||
@@ -423,8 +512,12 @@ pub fn mmq_q8_ref(
|
||||
// the opaque i32 fragment (cmma::store -> shared -> scale). Requires M%32==0, N%64==0, K%32==0.
|
||||
// ================================================================================================
|
||||
|
||||
/// Tiled int8 MMQ GEMM. Grid = (N/64, M/32); block = 256 threads (8 warps, 2x4 of 16x16 subtiles).
|
||||
#[kernel(targets(cuda), unchecked)]
|
||||
/// Tiled int8 MMQ GEMM. Grid = (N/64, M/32); block = 8 PLANES (2x4 of 16x16 subtiles, one per plane).
|
||||
///
|
||||
/// `plane` is comptime for the reason given on [`mmq_q8_wmma`]: one plane cooperatively owns one 16x16
|
||||
/// fragment, so the block is 8 planes wide -- 256 threads on a 32-lane CUDA warp, 512 on a 64-lane RADV
|
||||
/// wave -- and every scalar partition below divides its buffer by that width rather than by a literal.
|
||||
#[kernel(targets(cuda, rocm, vulkan, metal, cpu), unchecked)]
|
||||
pub fn mmq_q8_wmma_blk(
|
||||
xq: &Array<i8>,
|
||||
xs: &Array<f32>,
|
||||
@@ -433,97 +526,122 @@ pub fn mmq_q8_wmma_blk(
|
||||
out: &mut Array<f32>,
|
||||
#[comptime] n: usize,
|
||||
#[comptime] k: usize,
|
||||
#[comptime] plane: usize,
|
||||
#[comptime] target: Target,
|
||||
) {
|
||||
let tid = UNIT_POS as usize;
|
||||
let warp = tid / 32; // 0..7
|
||||
let lane = tid % 32; // 0..31
|
||||
let warp = tid / plane; // 0..7 (one plane per 16x16 subtile)
|
||||
let lane = tid % plane; // 0..plane-1
|
||||
let wm = warp / 4; // 0..1 (M-subtile)
|
||||
let wn = warp % 4; // 0..3 (N-subtile)
|
||||
let mrow0 = CUBE_POS_Y as usize * 32;
|
||||
let ncol0 = CUBE_POS_X as usize * 64;
|
||||
let kb_count = k / 32;
|
||||
let nthread = 8 * plane; // block width: 8 subtiles, one plane each
|
||||
let per = 256 / plane; // fragment elements this lane owns in the scalar epilogue
|
||||
|
||||
let mut sa = SharedMemory::<i8>::new(1024usize); // A tile [32 x 32]
|
||||
let mut sb = SharedMemory::<i8>::new(2048usize); // B tile [64 x 32]
|
||||
let mut ci = SharedMemory::<i32>::new(2048usize); // per-warp i32 fragment scratch [8 x 256]
|
||||
let mut accf = SharedMemory::<f32>::new(2048usize); // f32 output tile [32 x 64]
|
||||
|
||||
for e in 0usize..8 {
|
||||
accf[tid * 8 + e] = 0.0f32;
|
||||
for e in 0usize..(2048 / nthread) {
|
||||
accf[tid * (2048 / nthread) + e] = 0.0f32;
|
||||
}
|
||||
sync_cube();
|
||||
|
||||
for kb in 0..kb_count {
|
||||
let k0 = kb * 32;
|
||||
// Stage A[32x32] and B[64x32] int8 tiles into shared memory (coalesced, reused across warps).
|
||||
for i in 0usize..4 {
|
||||
let idx = tid + i * 256;
|
||||
for i in 0usize..(1024 / nthread) {
|
||||
let idx = tid + i * nthread;
|
||||
sa[idx] = xq[(mrow0 + idx / 32) * k + k0 + idx % 32];
|
||||
}
|
||||
for i in 0usize..8 {
|
||||
let idx = tid + i * 256;
|
||||
for i in 0usize..(2048 / nthread) {
|
||||
let idx = tid + i * nthread;
|
||||
sb[idx] = wq[(ncol0 + idx / 32) * k + k0 + idx % 32];
|
||||
}
|
||||
sync_cube();
|
||||
|
||||
// This warp's 16x16 subtile: accumulate the 32-block via 2 WMMA(K=16) from shared memory.
|
||||
let c = cmma::Matrix::<i32>::from_value(
|
||||
cmma::MatrixIdent::Accumulator,
|
||||
16usize,
|
||||
16usize,
|
||||
16usize,
|
||||
cmma::MatrixLayout::Undefined,
|
||||
0i32,
|
||||
);
|
||||
let a0 = cmma::Matrix::<i8>::from_slice(
|
||||
cmma::MatrixIdent::A,
|
||||
16usize,
|
||||
16usize,
|
||||
16usize,
|
||||
cmma::MatrixLayout::RowMajor,
|
||||
&sa.slice(wm * 512, 1024),
|
||||
32,
|
||||
);
|
||||
let b0 = cmma::Matrix::<i8>::from_slice(
|
||||
cmma::MatrixIdent::B,
|
||||
16usize,
|
||||
16usize,
|
||||
16usize,
|
||||
cmma::MatrixLayout::ColMajor,
|
||||
&sb.slice(wn * 512, 2048),
|
||||
32,
|
||||
);
|
||||
cmma::execute::<i8, i8, i32, i32>(&a0, &b0, &c, &c);
|
||||
let a1 = cmma::Matrix::<i8>::from_slice(
|
||||
cmma::MatrixIdent::A,
|
||||
16usize,
|
||||
16usize,
|
||||
16usize,
|
||||
cmma::MatrixLayout::RowMajor,
|
||||
&sa.slice(wm * 512 + 16, 1024),
|
||||
32,
|
||||
);
|
||||
let b1 = cmma::Matrix::<i8>::from_slice(
|
||||
cmma::MatrixIdent::B,
|
||||
16usize,
|
||||
16usize,
|
||||
16usize,
|
||||
cmma::MatrixLayout::ColMajor,
|
||||
&sb.slice(wn * 512 + 16, 2048),
|
||||
32,
|
||||
);
|
||||
cmma::execute::<i8, i8, i32, i32>(&a1, &b1, &c, &c);
|
||||
// This warp's 16x16 subtile: the staged 32-block's int8 contraction into `ci`. Both arms read
|
||||
// the SAME shared-memory A/B tiles, so operand staging is proven by the oracle too.
|
||||
island! {
|
||||
// Accelerated: 2 x cmma(16x16x16) i8->i32 straight out of shared memory.
|
||||
cuda | rocm | vulkan | metal => {
|
||||
let c = cmma::Matrix::<i32>::from_value(
|
||||
cmma::MatrixIdent::Accumulator,
|
||||
16usize,
|
||||
16usize,
|
||||
16usize,
|
||||
cmma::MatrixLayout::Undefined,
|
||||
0i32,
|
||||
);
|
||||
let a0 = cmma::Matrix::<i8>::from_slice(
|
||||
cmma::MatrixIdent::A,
|
||||
16usize,
|
||||
16usize,
|
||||
16usize,
|
||||
cmma::MatrixLayout::RowMajor,
|
||||
&sa.slice(wm * 512, 1024),
|
||||
32,
|
||||
);
|
||||
let b0 = cmma::Matrix::<i8>::from_slice(
|
||||
cmma::MatrixIdent::B,
|
||||
16usize,
|
||||
16usize,
|
||||
16usize,
|
||||
cmma::MatrixLayout::ColMajor,
|
||||
&sb.slice(wn * 512, 2048),
|
||||
32,
|
||||
);
|
||||
cmma::execute::<i8, i8, i32, i32>(&a0, &b0, &c, &c);
|
||||
let a1 = cmma::Matrix::<i8>::from_slice(
|
||||
cmma::MatrixIdent::A,
|
||||
16usize,
|
||||
16usize,
|
||||
16usize,
|
||||
cmma::MatrixLayout::RowMajor,
|
||||
&sa.slice(wm * 512 + 16, 1024),
|
||||
32,
|
||||
);
|
||||
let b1 = cmma::Matrix::<i8>::from_slice(
|
||||
cmma::MatrixIdent::B,
|
||||
16usize,
|
||||
16usize,
|
||||
16usize,
|
||||
cmma::MatrixLayout::ColMajor,
|
||||
&sb.slice(wn * 512 + 16, 2048),
|
||||
32,
|
||||
);
|
||||
cmma::execute::<i8, i8, i32, i32>(&a1, &b1, &c, &c);
|
||||
|
||||
// Spill this warp's i32 tile, apply the per-block f32 scale, accumulate into the f32 tile.
|
||||
cmma::store(
|
||||
&mut ci.slice_mut(warp * 256, warp * 256 + 256),
|
||||
&c,
|
||||
16,
|
||||
cmma::MatrixLayout::RowMajor,
|
||||
);
|
||||
cmma::store(
|
||||
&mut ci.slice_mut(warp * 256, warp * 256 + 256),
|
||||
&c,
|
||||
16,
|
||||
cmma::MatrixLayout::RowMajor,
|
||||
);
|
||||
}
|
||||
// NORMATIVE oracle: the identical i32 contraction as portable scalar MACs.
|
||||
default => {
|
||||
for e in 0usize..per {
|
||||
let p = lane * per + e;
|
||||
let smm = p / 16;
|
||||
let snn = p % 16;
|
||||
let mut isum = 0i32;
|
||||
for l in 0usize..32 {
|
||||
isum += i32::cast_from(sa[(wm * 16 + smm) * 32 + l])
|
||||
* i32::cast_from(sb[(wn * 16 + snn) * 32 + l]);
|
||||
}
|
||||
ci[warp * 256 + p] = isum;
|
||||
}
|
||||
}
|
||||
};
|
||||
sync_cube();
|
||||
for e in 0usize..8 {
|
||||
let p = lane * 8 + e;
|
||||
|
||||
// Apply the per-block f32 scale, accumulate into the f32 tile.
|
||||
for e in 0usize..per {
|
||||
let p = lane * per + e;
|
||||
let smm = p / 16;
|
||||
let snn = p % 16;
|
||||
let gmm = wm * 16 + smm;
|
||||
@@ -535,15 +653,16 @@ pub fn mmq_q8_wmma_blk(
|
||||
sync_cube();
|
||||
}
|
||||
|
||||
for e in 0usize..8 {
|
||||
let p = lane * 8 + e;
|
||||
for e in 0usize..per {
|
||||
let p = lane * per + e;
|
||||
let gmm = wm * 16 + p / 16;
|
||||
let gnn = wn * 16 + p % 16;
|
||||
out[(mrow0 + gmm) * n + (ncol0 + gnn)] = accf[gmm * 64 + gnn];
|
||||
}
|
||||
}
|
||||
|
||||
/// Host launch for the tiled MMQ GEMM. Returns (out, ms/dispatch over `iters`).
|
||||
/// Host launch for the tiled MMQ GEMM -- the PRODUCTION surface; derives the island tag from the
|
||||
/// runtime. Returns (out, ms/dispatch over `iters`).
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub fn mmq_q8_wmma_blk_run<R: Runtime>(
|
||||
client: &ComputeClient<R>,
|
||||
@@ -556,6 +675,24 @@ pub fn mmq_q8_wmma_blk_run<R: Runtime>(
|
||||
k: usize,
|
||||
iters: usize,
|
||||
) -> (Vec<f32>, f64) {
|
||||
mmq_q8_wmma_blk_run_with(client, xq, xs, wq, wd, m, n, k, iters, Target::of(client))
|
||||
}
|
||||
|
||||
/// Host launch for the tiled MMQ GEMM with an explicit island tag. See [`mmq_q8_wmma_run_with`].
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub fn mmq_q8_wmma_blk_run_with<R: Runtime>(
|
||||
client: &ComputeClient<R>,
|
||||
xq: &[i8],
|
||||
xs: &[f32],
|
||||
wq: &[i8],
|
||||
wd: &[f32],
|
||||
m: usize,
|
||||
n: usize,
|
||||
k: usize,
|
||||
iters: usize,
|
||||
target: Target,
|
||||
) -> (Vec<f32>, f64) {
|
||||
let plane = client.properties().hardware.plane_size_max;
|
||||
let xqh = client.create_from_slice(i8::as_bytes(xq));
|
||||
let xsh = client.create_from_slice(f32::as_bytes(xs));
|
||||
let wqh = client.create_from_slice(i8::as_bytes(wq));
|
||||
@@ -566,7 +703,7 @@ pub fn mmq_q8_wmma_blk_run<R: Runtime>(
|
||||
mmq_q8_wmma_blk::launch_unchecked::<R>(
|
||||
c,
|
||||
grid.clone(),
|
||||
Block::new_1d(256),
|
||||
Block::new_1d(8 * plane),
|
||||
ArrayArg::from_raw_parts(xqh.clone(), xq.len()),
|
||||
ArrayArg::from_raw_parts(xsh.clone(), xs.len()),
|
||||
ArrayArg::from_raw_parts(wqh.clone(), wq.len()),
|
||||
@@ -574,6 +711,8 @@ pub fn mmq_q8_wmma_blk_run<R: Runtime>(
|
||||
ArrayArg::from_raw_parts(oh.clone(), m * n),
|
||||
n,
|
||||
k,
|
||||
plane as usize,
|
||||
target,
|
||||
);
|
||||
};
|
||||
launch(client);
|
||||
@@ -616,3 +755,74 @@ pub fn gen_mmq(m: usize, n: usize, k: usize) -> (Vec<i8>, Vec<f32>, Vec<i8>, Vec
|
||||
.collect();
|
||||
(xq, xs, wq, wd)
|
||||
}
|
||||
|
||||
#[cfg(all(test, feature = "cpu"))]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use cubecl::cpu::{CpuDevice, CpuRuntime};
|
||||
|
||||
fn cpu_client() -> ComputeClient<CpuRuntime> {
|
||||
CpuRuntime::client(&CpuDevice)
|
||||
}
|
||||
|
||||
/// Scale-relative gate: max|delta| / max|ref|. A per-element relative error is meaningless here --
|
||||
/// a signed int8 sum cancels to near-zero often, which explodes the per-element denominator and
|
||||
/// reports a false failure. A real decode/accumulation bug moves this number; cancellation does not.
|
||||
fn rel_to_max(got: &[f32], want: &[f32]) -> f32 {
|
||||
let mut maxabs = 0f32;
|
||||
let mut refmax = 1e-9f32;
|
||||
for (g, w) in got.iter().zip(want) {
|
||||
maxabs = maxabs.max((g - w).abs());
|
||||
refmax = refmax.max(w.abs());
|
||||
}
|
||||
maxabs / refmax
|
||||
}
|
||||
|
||||
/// The oracle. The naive MMQ GEMM's normative arm, on the CPU runtime, IS `mmq_q8_ref` -- bit-exact,
|
||||
/// not merely close. It can be bit-exact because the kernel accumulates the f32 scale in the same
|
||||
/// order the reference does (one 32-block at a time, ascending), and the int8 inner sum is integer.
|
||||
#[test]
|
||||
fn mmq_default_arm_is_bit_exact_on_cpu() {
|
||||
let client = cpu_client();
|
||||
for (m, n, k) in [(16usize, 16usize, 64usize), (32, 16, 128), (16, 32, 256)] {
|
||||
let (xq, xs, wq, wd) = gen_mmq(m, n, k);
|
||||
let (got, _) = mmq_q8_wmma_run_with::<CpuRuntime>(
|
||||
&client, &xq, &xs, &wq, &wd, m, n, k, 1, Target::Cpu,
|
||||
);
|
||||
let want = mmq_q8_ref(&xq, &xs, &wq, &wd, m, n, k);
|
||||
let gb: Vec<u32> = got.iter().map(|x| x.to_bits()).collect();
|
||||
let wb: Vec<u32> = want.iter().map(|x| x.to_bits()).collect();
|
||||
assert_eq!(gb, wb, "MMQ {m}x{n}x{k}: default arm != mmq_q8_ref (bit-exact gate)");
|
||||
}
|
||||
}
|
||||
|
||||
/// The tiled GEMM's normative arm agrees with the same reference through the shared-memory A/B
|
||||
/// staging path. Gated scale-relative rather than bit-exact: this kernel accumulates each warp's
|
||||
/// subtile independently, so the f32 add order differs from the reference's single ascending sweep.
|
||||
#[test]
|
||||
fn mmq_blk_default_arm_matches_ref_on_cpu() {
|
||||
let client = cpu_client();
|
||||
let (m, n, k) = (32usize, 64usize, 128usize);
|
||||
let (xq, xs, wq, wd) = gen_mmq(m, n, k);
|
||||
let (got, _) = mmq_q8_wmma_blk_run_with::<CpuRuntime>(
|
||||
&client, &xq, &xs, &wq, &wd, m, n, k, 1, Target::Cpu,
|
||||
);
|
||||
let want = mmq_q8_ref(&xq, &xs, &wq, &wd, m, n, k);
|
||||
let rel = rel_to_max(&got, &want);
|
||||
assert!(rel < 1e-6, "tiled MMQ {m}x{n}x{k}: rel_to_max={rel:.3e} vs mmq_q8_ref");
|
||||
}
|
||||
|
||||
/// The production surface derives the island tag from the runtime, so the CPU runtime resolves to
|
||||
/// the normative arm with no caller naming a target.
|
||||
#[test]
|
||||
fn production_run_selects_oracle_on_cpu() {
|
||||
let client = cpu_client();
|
||||
let (m, n, k) = (16usize, 16usize, 64usize);
|
||||
let (xq, xs, wq, wd) = gen_mmq(m, n, k);
|
||||
let (got, _) = mmq_q8_wmma_run::<CpuRuntime>(&client, &xq, &xs, &wq, &wd, m, n, k, 1);
|
||||
let want = mmq_q8_ref(&xq, &xs, &wq, &wd, m, n, k);
|
||||
let gb: Vec<u32> = got.iter().map(|x| x.to_bits()).collect();
|
||||
let wb: Vec<u32> = want.iter().map(|x| x.to_bits()).collect();
|
||||
assert_eq!(gb, wb);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,13 +1,16 @@
|
||||
//! mmq-check: does the DSL emit a working int8 tensor-core MMQ GEMM on CUDA?
|
||||
//! mmq-check: does the DSL emit a working int8 tensor-core MMQ GEMM on THIS backend?
|
||||
//!
|
||||
//! Staged, cheapest-kill-first (run with --features cuda on the GPU box):
|
||||
//! Staged, cheapest-kill-first (`--features cuda` or `--features vulkan` on the GPU box):
|
||||
//! 1a high-level WMMA hello (i8 16x16x16 -> i32) exact-int gate
|
||||
//! 1b manual MMA hello (i8 m16n8k32 -> i32) exact-int gate
|
||||
//! 2 Q8_0 x q8_1 MMQ tile (one/few tiles, K-loop) bit-close vs the quantized CPU oracle
|
||||
//! 1b manual MMA hello (i8 m16n8k32 -> i32) exact-int gate, CUDA/ROCm only
|
||||
//! 2 Q8_0 x q8_1 MMQ tile (one/few tiles, K-loop) scale-relative vs the quantized oracle
|
||||
//! 3 full prefill GEMM (M=512 N=4096 K=4096) verify + kernel-only GFLOP/s & GB/s
|
||||
//!
|
||||
//! Stage 1b is skipped off CUDA/ROCm by construction, not by failure: SPIR-V has no per-lane fragment
|
||||
//! layout to query, so the manual-MMA family cannot lower there (see `hanzo_kernel::mmq::mma_hello_i8`).
|
||||
//!
|
||||
//! Each stage is isolated with catch_unwind so a JIT failure in one path still lets the rest report
|
||||
//! (a decisive negative -- "the DSL can't express X on sm_121" -- must survive to the log).
|
||||
//! (a decisive negative -- "the DSL can't express X here" -- must survive to the log).
|
||||
|
||||
use hanzo_kernel::mmq::*;
|
||||
use hanzo_kernel::prelude::*;
|
||||
@@ -47,17 +50,18 @@ fn rand_i8(n: usize, seed: u64) -> Vec<i8> {
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn main() {
|
||||
use cubecl::cuda::{CudaDevice, CudaRuntime};
|
||||
let client = CudaRuntime::client(&CudaDevice::default());
|
||||
/// The whole probe, over any runtime -- ONE staged gate, every backend. `name` labels the log.
|
||||
fn probe<R: Runtime>(name: &str, client: &ComputeClient<R>) {
|
||||
let client = &client.clone();
|
||||
let plane = client.properties().hardware.plane_size_max;
|
||||
println!("== hanzo-kernel MMQ probe (CUDA, plane={plane}) ==\n");
|
||||
let target = Target::of(client);
|
||||
println!("== hanzo-kernel MMQ probe ({name}, plane={plane}, target={target:?}) ==\n");
|
||||
|
||||
// ---- Stage 1a: high-level WMMA hello ----
|
||||
let r = catch_unwind(AssertUnwindSafe(|| {
|
||||
let a = rand_i8(256, 0x1234);
|
||||
let b = rand_i8(256, 0x9E37);
|
||||
let got = wmma_hello_i8_run::<CudaRuntime>(&client, &a, &b);
|
||||
let got = wmma_hello_i8_run::<R>(client, &a, &b);
|
||||
let want = hello_i8_ref(&a, &b);
|
||||
let bad = got.iter().zip(&want).filter(|(x, y)| x != y).count();
|
||||
println!(
|
||||
@@ -75,33 +79,37 @@ fn main() {
|
||||
println!("[1a] WMMA hello i8 16x16x16->i32 PANIC/JIT-FAIL ✗ (see stderr above)");
|
||||
}
|
||||
|
||||
// ---- Stage 1b: manual MMA hello ----
|
||||
let r = catch_unwind(AssertUnwindSafe(|| {
|
||||
let a = rand_i8(16 * 32, 0xABCD);
|
||||
let b = rand_i8(32 * 8, 0x5555);
|
||||
let got = mma_hello_i8_run::<CudaRuntime>(&client, &a, &b, plane);
|
||||
let want = mma_hello_ref(&a, &b);
|
||||
let bad = got.iter().zip(&want).filter(|(x, y)| x != y).count();
|
||||
println!(
|
||||
"[1b] manual MMA i8 m16n8k32->i32 {} ({} / 128 mismatched)",
|
||||
if bad == 0 {
|
||||
"EXACT ✓"
|
||||
} else {
|
||||
"MISMATCH ✗"
|
||||
},
|
||||
bad
|
||||
);
|
||||
bad == 0
|
||||
}));
|
||||
if r.is_err() {
|
||||
println!("[1b] manual MMA i8 m16n8k32->i32 PANIC/JIT-FAIL ✗ (see stderr above)");
|
||||
// ---- Stage 1b: manual MMA hello (CUDA/ROCm only -- SPIR-V has no queryable fragment layout) ----
|
||||
if matches!(target, Target::Cuda | Target::Rocm) {
|
||||
let r = catch_unwind(AssertUnwindSafe(|| {
|
||||
let a = rand_i8(16 * 32, 0xABCD);
|
||||
let b = rand_i8(32 * 8, 0x5555);
|
||||
let got = mma_hello_i8_run::<R>(client, &a, &b, plane);
|
||||
let want = mma_hello_ref(&a, &b);
|
||||
let bad = got.iter().zip(&want).filter(|(x, y)| x != y).count();
|
||||
println!(
|
||||
"[1b] manual MMA i8 m16n8k32->i32 {} ({} / 128 mismatched)",
|
||||
if bad == 0 {
|
||||
"EXACT ✓"
|
||||
} else {
|
||||
"MISMATCH ✗"
|
||||
},
|
||||
bad
|
||||
);
|
||||
bad == 0
|
||||
}));
|
||||
if r.is_err() {
|
||||
println!("[1b] manual MMA i8 m16n8k32->i32 PANIC/JIT-FAIL ✗ (see stderr above)");
|
||||
}
|
||||
} else {
|
||||
println!("[1b] manual MMA i8 m16n8k32->i32 N/A (no fragment layout on {target:?})");
|
||||
}
|
||||
|
||||
// ---- Stage 2: MMQ tile(s), bit-close vs the quantized oracle ----
|
||||
for (m, n, k) in [(16usize, 16usize, 64usize), (32, 64, 256), (64, 128, 512)] {
|
||||
let r = catch_unwind(AssertUnwindSafe(|| {
|
||||
let (xq, xs, wq, wd) = gen_mmq(m, n, k);
|
||||
let (got, _) = mmq_q8_wmma_run::<CudaRuntime>(&client, &xq, &xs, &wq, &wd, m, n, k, 1);
|
||||
let (got, _) = mmq_q8_wmma_run::<R>(client, &xq, &xs, &wq, &wd, m, n, k, 1);
|
||||
let want = mmq_q8_ref(&xq, &xs, &wq, &wd, m, n, k);
|
||||
let rel = max_rel(&want, &got);
|
||||
let (maxabs, relmax) = err_robust(&want, &got);
|
||||
@@ -125,7 +133,7 @@ fn main() {
|
||||
let mrows = 8usize;
|
||||
let want = mmq_q8_ref(&xq, &xs, &wq, &wd, mrows, n, k); // exact 8-row stripe oracle
|
||||
|
||||
let (g0, ms0) = mmq_q8_wmma_run::<CudaRuntime>(&client, &xq, &xs, &wq, &wd, m, n, k, 50);
|
||||
let (g0, ms0) = mmq_q8_wmma_run::<R>(client, &xq, &xs, &wq, &wd, m, n, k, 50);
|
||||
let (a0, r0) = err_robust(&want, &g0[..mrows * n]);
|
||||
println!(
|
||||
"\n[3 ] GEMM {m}x{n}x{k} (naive 1-warp/16x16 tile) verify rel_to_max={:.2e} max_abs={:.2e} {}",
|
||||
@@ -138,8 +146,7 @@ fn main() {
|
||||
wbytes / (ms0 * 1e6)
|
||||
);
|
||||
|
||||
let (g1, ms1) =
|
||||
mmq_q8_wmma_blk_run::<CudaRuntime>(&client, &xq, &xs, &wq, &wd, m, n, k, 50);
|
||||
let (g1, ms1) = mmq_q8_wmma_blk_run::<R>(client, &xq, &xs, &wq, &wd, m, n, k, 50);
|
||||
let (a1, r1) = err_robust(&want, &g1[..mrows * n]);
|
||||
println!(
|
||||
"[3 ] GEMM {m}x{n}x{k} (tiled 8-warp/32x64 tile) verify rel_to_max={:.2e} max_abs={:.2e} {}",
|
||||
@@ -154,3 +161,16 @@ fn main() {
|
||||
println!("[3 ] GEMM 512x4096x4096 PANIC/JIT-FAIL ✗");
|
||||
}
|
||||
}
|
||||
|
||||
fn main() {
|
||||
#[cfg(feature = "cuda")]
|
||||
{
|
||||
use cubecl::cuda::{CudaDevice, CudaRuntime};
|
||||
probe::<CudaRuntime>("CUDA", &CudaRuntime::client(&CudaDevice::default()));
|
||||
}
|
||||
#[cfg(feature = "vulkan")]
|
||||
{
|
||||
use cubecl::wgpu::{WgpuDevice, WgpuRuntime};
|
||||
probe::<WgpuRuntime>("VULKAN", &WgpuRuntime::client(&WgpuDevice::default()));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -413,7 +413,7 @@ pub fn moe_matvec_q4k_blk<F: Float>(
|
||||
x: &Array<F>,
|
||||
ids: &Array<u32>,
|
||||
out: &mut Array<F>,
|
||||
#[comptime] n: usize, // comptime shape: outrow/n & nsub/nt fold to magic-multiply + full unroll
|
||||
#[comptime] n: usize, // comptime shape: outrow/n & nsub/nt fold to magic-multiply divides
|
||||
#[comptime] k: usize, // (the DSL's job -- one source fn, one fast .spv per live shape, like llama)
|
||||
#[comptime] nt: usize, // threads per output element (k/32 a multiple of nt)
|
||||
) {
|
||||
@@ -1089,7 +1089,7 @@ pub fn moe_matvec_q6k_blk<F: Float>(
|
||||
x: &Array<F>,
|
||||
ids: &Array<u32>,
|
||||
out: &mut Array<F>,
|
||||
#[comptime] n: usize, // comptime shape (see moe_matvec_q4k_blk): fold + full unroll
|
||||
#[comptime] n: usize, // comptime shape (see moe_matvec_q4k_blk): folds to magic-multiply divides
|
||||
#[comptime] k: usize,
|
||||
#[comptime] nt: usize,
|
||||
) {
|
||||
|
||||
Reference in New Issue
Block a user