hanzo/ml
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge DSL f32-direct Q4_K decode matvec (hanzo-ml 0.11.94, hanzo-kernel 0.2.39)

Browse Source 
The DSL twin of mul_mat_vec_q4k_cm (plane_sum + coalesced reads) earns the
live Vulkan dense Q4_K decode path where it measures >= hand (gated nout>=4096,
DSL/hand 1.022x on the dominant FFN shape); small shapes keep hand. One DSL
Rust source replaces the hand fork on the shapes that dominate decode bytes.

# Conflicts:
#	hanzo-kernel/Cargo.toml
This commit is contained in:
Hanzo Dev
2026-07-22 16:19:42 -07:00
parent 2f0d540664 1611dc6e19
commit a29925ad9e
6 changed files with 375 additions and 72 deletions
Download Patch File Download Diff File Expand all files Collapse all files
hanzo-kernel/Cargo.toml
+1 -1
View File
@@ -1,6 +1,6 @@
[package]
name = "hanzo-kernel"
version = "0.2.38"
version = "0.2.39"
edition = "2021"
description = "Hanzo's first-party GPU kernel DSL: one Rust source, lowered to CUDA/ROCm/Vulkan/Metal."
license = "BSD-3-Clause"
hanzo-kernel/src/main.rs
+118 -1
View File
@@ -14,7 +14,7 @@ use hanzo_kernel::quant::{
moe_route, moe_route_ref, moe_route_run, gemv, gemv_run,
moe_matvec_q4k_dp4a_blk, moe_matvec_q4k_dp4a_blk_run, quant_act_q8_cpu,
moe_matvec_q6k_dp4a_blk, moe_matvec_q6k_dp4a_blk_run,
matvec_q4k_dp4a_blk, matvec_q4k_dp4a_blk_run, pack_q4k, QK8_0,
matvec_q4k_dp4a_blk, matvec_q4k_dp4a_blk_run, matvec_q4k_f32_blk_run, pack_q4k, QK8_0,
};
use std::time::Instant;
use hanzo_kernel::norm::rms_norm_run;
@@ -653,6 +653,20 @@ fn main() {
return;
}
// `matvec-check dumpf32`: emit ONLY the f32-direct Q4_K decode matvec .spv (matvec_q4k_f32_blk) at
// the schedule the cold sweep crowned (nt=64, nr=1). k and nout are runtime (meta[0] / out.len()),
// so this ONE .spv serves every dense Q4_K decode shape. CUBECL_DEBUG_SPIRV=<dir> writes the raw
// artifact; spv_to_ml.sh renames the entry point to `main` for ml's VulkanDevice::dispatch.
#[cfg(all(feature = "vulkan", feature = "spirv-dump"))]
if std::env::args().any(|a| a == "dumpf32") {
use cubecl::wgpu::{WgpuDevice, WgpuRuntime};
let c = WgpuRuntime::client(&WgpuDevice::default());
let (wqs, wsc, wd, wdm, x) = gen_q4k(2048, 2048);
let _ = matvec_q4k_f32_blk_run::<WgpuRuntime>(&c, &wqs, &wsc, &wd, &wdm, &x, 2048, 2048, 64, 1);
println!("dumpf32: matvec_q4k_f32_blk nt=64 nr=1 dispatched -- .spv emitted");
return;
}
// `matvec-check dump-flash`: emit ONLY the flash-attention .spv (isolated from the dp4a/mmq dumps
// above, which need extensions a software adapter lacks). The flash kernel's cmma island lowers to
// OpCooperativeMatrixMulAddKHR; the .spv is written by cubecl at codegen, BEFORE the driver validates
@@ -700,6 +714,109 @@ fn main() {
return;
}
// `matvec-check f32sweep`: cold GB/s sweep of the f32-direct Q4_K decode matvec (DSL
// matvec_q4k_f32_blk) over its {WG x NR} schedule at the live dense Q4_K decode shapes
// (Qwen3-1.7B: ffn_gate/up nout=6144, attn_q/o nout=2048, attn_k nout=1024; k=2048). Uses the
// bank-rotated, bit-exact-gated MatvecQ4kF32Eval so each timed dispatch reads a weight bank the
// 32 MB MALL has already evicted -- the real per-token decode regime, not the warm small-shape
// benches. Names the fastest (nt, nr) per shape (the width the engine dispatch selects by row
// count) and its cold weight-bandwidth -- the screen that decides whether the DSL twin is worth
// the in-engine A/B against the hand mul_mat_vec_q4k_cm.
#[cfg(feature = "vulkan")]
if std::env::args().any(|a| a == "f32sweep") {
use cubecl::wgpu::{WgpuDevice, WgpuRuntime};
use hanzo_kernel::quant::{matvec_q4k_f32_space, MatvecQ4kF32Eval};
use hanzo_kernel::tune::Evaluator;
let c = WgpuRuntime::client(&WgpuDevice::default());
println!("== cold f32-direct Q4_K matvec sweep (DSL matvec_q4k_f32_blk), bank-rotated >> 32 MB MALL ==");
let space = matvec_q4k_f32_space();
for (nout, k) in [(6144usize, 2048usize), (2048, 2048), (1024, 2048)] {
let wbytes = nout * (k / 256) * 144;
let nbanks = (64 * 1024 * 1024 / wbytes + 1).max(6); // footprint >= 64 MB (2x MALL)
let (wqs, wsc, wd, wdm, x) = gen_q4k(nout, k);
let eval = MatvecQ4kF32Eval::new(&c, &space, &wqs, &wsc, &wd, &wdm, &x, nout, k, nbanks, 6);
println!(
"-- {nout}x{k} ({} MB/bank x {nbanks} banks = {} MB cold footprint) --",
wbytes / (1024 * 1024),
nbanks * wbytes / (1024 * 1024)
);
let mut best = (f64::INFINITY, 0usize, 0usize, 0f64);
for cfg in space.enumerate() {
let nt = cfg.get(&space, "WG") as usize;
let nr = cfg.get(&space, "NR") as usize;
let ms = eval.measure(&cfg, nbanks); // iters = nbanks: one cold pass over all banks
if !ms.is_finite() {
println!(" nt={nt:<3} nr={nr} REJECT (diverged from Q4_K oracle)");
continue;
}
let gbps = wbytes as f64 / (ms * 1e6);
println!(" nt={nt:<3} nr={nr} {ms:.4} ms {gbps:.0} GB/s");
if ms < best.0 {
best = (ms, nt, nr, gbps);
}
}
println!(
" >> BEST f32 {nout}x{k}: nt={} nr={} {:.4} ms {:.0} GB/s (worst_rel {:.2e})",
best.1, best.2, best.0, best.3, eval.worst_rel()
);
// dp4a-block anchor at the same shape + same cold bank-rotation: the DSL twin of the OTHER
// shipped decode kernel. The hand mul_mat_vec_q4k_cm shipped +14.9% over the dp4a block
// in-engine, so the f32/dp4a ratio here calibrates the f32 twin against the hand CM without
// the microbench-vs-engine harness gap (both DSL kernels run through the same wgpu path).
use cubecl::server::Handle;
let packed = pack_q4k(&wqs, &wsc, &wd, &wdm);
let (xq, xs, xsum) = quant_act_q8_cpu(&x, 1, k);
let dbanks: Vec<Handle> =
(0..nbanks).map(|_| c.create_from_slice(u32::as_bytes(&packed))).collect();
let xqh = c.create_from_slice(u32::as_bytes(&xq));
let xsh = c.create_from_slice(f32::as_bytes(&xs));
let xsumh = c.create_from_slice(f32::as_bytes(&xsum));
let dmeta = c.create_from_slice(u32::as_bytes(&[k as u32]));
let doh = c.create_from_slice(f32::as_bytes(&vec![0.0f32; nout]));
let dp4a = |nt: usize, bank: &Handle| unsafe {
matvec_q4k_dp4a_blk::launch_unchecked::<f32, WgpuRuntime>(
&c,
Grid::Static(nout as u32, 1, 1),
Block::new_1d(nt as u32),
ArrayArg::from_raw_parts(bank.clone(), packed.len()),
ArrayArg::from_raw_parts(xqh.clone(), xq.len()),
ArrayArg::from_raw_parts(xsh.clone(), xs.len()),
ArrayArg::from_raw_parts(xsumh.clone(), xsum.len()),
ArrayArg::from_raw_parts(doh.clone(), nout),
ArrayArg::from_raw_parts(dmeta.clone(), 1),
nt,
);
};
let mut dbest = (f64::INFINITY, 0usize);
for nt in [32usize, 64, 128] {
for i in 0..(2 * nbanks) {
dp4a(nt, &dbanks[i % nbanks]);
}
let _ = c.read_one_unchecked(doh.clone());
let mut m = f64::INFINITY;
for _ in 0..6 {
let t = std::time::Instant::now();
for i in 0..nbanks {
dp4a(nt, &dbanks[i % nbanks]);
}
let _ = c.read_one_unchecked(doh.clone());
m = m.min(t.elapsed().as_secs_f64() * 1e3 / nbanks as f64);
}
let g = wbytes as f64 / (m * 1e6);
println!(" dp4a nt={nt:<3} {m:.4} ms {g:.0} GB/s");
if m < dbest.0 {
dbest = (m, nt);
}
}
let dg = wbytes as f64 / (dbest.0 * 1e6);
println!(
" >> dp4a BEST {nout}x{k}: nt={} {:.4} ms {:.0} GB/s || f32/dp4a = {:.2}x (needs >= 1.15x to match hand CM)",
dbest.1, dbest.0, dg, best.3 / dg
);
}
return;
}
println!("hanzo-kernel :: one #[device] matvec_q8 source, lowered per backend, gated bit-exact\n");
#[cfg(feature = "cpu")]
hanzo-kernel/src/quant.rs
+144 -67
View File
@@ -731,85 +731,160 @@ pub fn matvec_q4k_f32_blk<F: Float>(
wq: &Array<u32>, // packed Q4_K, 36 u32/superblock
x: &Array<F>, // activation, f32 [k]
out: &mut Array<F>,
meta: &Array<u32>, // [k]
meta: &Array<u32>, // [k, nout]
#[comptime] nt: usize, // threads cooperating on one row's k (the workgroup-width axis)
#[comptime] nr: usize, // output rows per workgroup (the row tile)
#[comptime] tgt: Target, // island scrutinee: Vulkan takes the coalesced plane_sum idiom
) {
let wgid = CUBE_POS as usize;
let t = UNIT_POS as usize;
let k = meta[0] as usize;
let nb = k / 256;
let nsub = k / 32;
let per = (nsub + nt - 1) / nt;
let row0 = wgid * nr;
let nout = out.len();
let mut acc = Array::<F>::new(nr);
#[unroll]
for n in 0..nr {
acc[n] = F::new(0.0);
}
for j in 0..per {
let sb = j * nt + t; // sub-block (0..nsub) this thread contracts on this pass
if sb < nsub {
let sup = sb / 8; // superblock
let jloc = sb % 8; // sub-block within the superblock (0..7)
let shift = ((jloc % 2) * 4) as u32; // low (even sub-block) or high (odd) nibble
let cwoff = 4 + (jloc / 2) * 8; // qs u32 base within a superblock for this sub-block
let xbase = sup * 256 + (jloc / 2) * 64 + (jloc % 2) * 32; // activation base of the 32 weights
let nblocks = k / 256;
let nout = meta[1] as usize; // nout from meta (not out.len()): the cubecl info buffer strips for ml dispatch
let row0 = CUBE_POS as usize * nr;
island! {
// Vulkan: the DSL twin of the hand mul_mat_vec_q4k_cm (nt = plane = 64). 16 threads cooperate
// on ONE super-block reading ADJACENT qs u32 words (thread itid reads word 4+itid, coalesced
// 64-byte burst) while ITS = nt/16 super-blocks stream in parallel; the per-row partial reduces
// with `plane_sum` (subgroupAdd) -- no shared memory, no tree barriers. Each thread owns a
// 16-weight sub-position (4 sub-blocks x 4 bytes) so the 16 threads cover all 8 sub-blocks and
// the 32 qs words (4+itid, 4+itid+16) exactly once. Requires nt = the hardware plane size.
vulkan => {
let its = nt / 16;
let itid = t & 15;
let ix = t / 16;
let hg = itid / 8; // 0/1 pair-group -> sub-blocks {0,1,4,5} or {2,3,6,7}
let qq = (itid & 7) * 4; // 0,4,..,28: first of this thread's 4 weights in each sub-block
let qw = itid; // qs u32 word (low group); high group at qw+16
let j0 = 2 * hg;
let j1 = j0 + 1;
let j2 = 4 + 2 * hg;
let j3 = j2 + 1;
let ya = hg * 64 + qq;
let yb = ya + 32;
let yc = 128 + hg * 64 + qq;
let yd = yc + 32;
let mut acc = Array::<F>::new(nr);
#[unroll]
for n in 0..nr {
acc[n] = F::new(0.0);
}
#[unroll]
for n in 0..nr {
let row = row0 + n;
if row < nout {
let bb = (row * nb + sup) * 36; // packed superblock base
let d = F::cast_from(f16lo_to_f32(wq[bb]));
let dmin = F::cast_from(f16lo_to_f32(wq[bb] >> 16));
let dj = d * F::cast_from(q4k_sc(wq, bb + 1, jloc));
let mj = dmin * F::cast_from(q4k_m(wq, bb + 1, jloc));
let cw = bb + cwoff;
let mut sdot = F::new(0.0);
let mut sx = F::new(0.0);
#[unroll]
for g in 0..8usize {
let word = (wq[cw + g] >> shift) & 0x0F0F0F0F;
let xb = xbase + g * 4;
let x0 = x[xb];
let x1 = x[xb + 1];
let x2 = x[xb + 2];
let x3 = x[xb + 3];
sdot += F::cast_from(word & 0xFF) * x0
+ F::cast_from((word >> 8) & 0xFF) * x1
+ F::cast_from((word >> 16) & 0xFF) * x2
+ F::cast_from((word >> 24) & 0xFF) * x3;
sx += x0 + x1 + x2 + x3;
let rowbase = row * nblocks * 36;
let mut blk = ix;
while blk < nblocks {
let b = rowbase + blk * 36;
let d = F::cast_from(f16lo_to_f32(wq[b]));
let dmin = F::cast_from(f16lo_to_f32(wq[b] >> 16));
let wlo = wq[b + 4 + qw]; // sub-blocks j0 (low nibble), j1 (high nibble)
let whi = wq[b + 4 + qw + 16]; // sub-blocks j2 (low), j3 (high)
let dj0 = d * F::cast_from(q4k_sc(wq, b + 1, j0));
let mj0 = dmin * F::cast_from(q4k_m(wq, b + 1, j0));
let dj1 = d * F::cast_from(q4k_sc(wq, b + 1, j1));
let mj1 = dmin * F::cast_from(q4k_m(wq, b + 1, j1));
let dj2 = d * F::cast_from(q4k_sc(wq, b + 1, j2));
let mj2 = dmin * F::cast_from(q4k_m(wq, b + 1, j2));
let dj3 = d * F::cast_from(q4k_sc(wq, b + 1, j3));
let mj3 = dmin * F::cast_from(q4k_m(wq, b + 1, j3));
let yblk = blk * 256;
#[unroll]
for l in 0..4usize {
let s8 = (l * 8) as u32;
acc[n] += (dj0 * F::cast_from((wlo >> s8) & 0x0F) - mj0) * x[yblk + ya + l];
acc[n] += (dj1 * F::cast_from((wlo >> (s8 + 4)) & 0x0F) - mj1) * x[yblk + yb + l];
acc[n] += (dj2 * F::cast_from((whi >> s8) & 0x0F) - mj2) * x[yblk + yc + l];
acc[n] += (dj3 * F::cast_from((whi >> (s8 + 4)) & 0x0F) - mj3) * x[yblk + yd + l];
}
blk += its;
}
}
}
// Plane reduction (one plane per row): sums all nt lanes' partials with no shared memory.
#[unroll]
for n in 0..nr {
let total = plane_sum(acc[n]);
if t == 0 {
let row = row0 + n;
if row < nout {
out[row] = total;
}
// Affine decode folded over the sub-block: sum (dj*q - mj)*x = dj*sum(q*x) - mj*sum(x).
acc[n] += dj * sdot - mj * sx;
}
}
}
}
// Per-row tree reduction of the nt partials through shared memory (nr sequential passes; both small).
let mut smem = SharedMemory::<F>::new(nt);
#[unroll]
for n in 0..nr {
smem[t] = acc[n];
sync_cube();
let mut stride = CUBE_DIM / 2;
while stride > 0 {
if UNIT_POS < stride {
let v = smem[(UNIT_POS + stride) as usize];
smem[t] += v;
// CPU oracle + non-Vulkan GPUs: `nt` threads split the row's sub-blocks (ceil+guard, so nt may
// exceed nsub), each accumulating its slice, then a shared-memory tree reduce. Works for any nt.
default => {
let nsub = k / 32;
let per = (nsub + nt - 1) / nt;
let mut acc = Array::<F>::new(nr);
#[unroll]
for n in 0..nr {
acc[n] = F::new(0.0);
}
sync_cube();
stride /= 2;
}
if t == 0 {
let row = row0 + n;
if row < nout {
out[row] = smem[0];
for j in 0..per {
let sb = j * nt + t;
if sb < nsub {
let sup = sb / 8;
let jloc = sb % 8;
let shift = ((jloc % 2) * 4) as u32;
let cwoff = 4 + (jloc / 2) * 8;
let xbase = sup * 256 + (jloc / 2) * 64 + (jloc % 2) * 32;
#[unroll]
for n in 0..nr {
let row = row0 + n;
if row < nout {
let bb = (row * nblocks + sup) * 36;
let d = F::cast_from(f16lo_to_f32(wq[bb]));
let dmin = F::cast_from(f16lo_to_f32(wq[bb] >> 16));
let dj = d * F::cast_from(q4k_sc(wq, bb + 1, jloc));
let mj = dmin * F::cast_from(q4k_m(wq, bb + 1, jloc));
let cw = bb + cwoff;
let mut sdot = F::new(0.0);
let mut sx = F::new(0.0);
#[unroll]
for g in 0..8usize {
let word = (wq[cw + g] >> shift) & 0x0F0F0F0F;
let xb = xbase + g * 4;
let x0 = x[xb];
let x1 = x[xb + 1];
let x2 = x[xb + 2];
let x3 = x[xb + 3];
sdot += F::cast_from(word & 0xFF) * x0
+ F::cast_from((word >> 8) & 0xFF) * x1
+ F::cast_from((word >> 16) & 0xFF) * x2
+ F::cast_from((word >> 24) & 0xFF) * x3;
sx += x0 + x1 + x2 + x3;
}
acc[n] += dj * sdot - mj * sx;
}
}
}
}
let mut smem = SharedMemory::<F>::new(nt);
#[unroll]
for n in 0..nr {
smem[t] = acc[n];
sync_cube();
let mut stride = CUBE_DIM / 2;
while stride > 0 {
if UNIT_POS < stride {
let v = smem[(UNIT_POS + stride) as usize];
smem[t] += v;
}
sync_cube();
stride /= 2;
}
if t == 0 {
let row = row0 + n;
if row < nout {
out[row] = smem[0];
}
}
sync_cube();
}
}
sync_cube();
}
}
@@ -823,7 +898,7 @@ pub fn matvec_q4k_f32_blk_run<R: Runtime>(
let packed = pack_q4k(wqs, wsc, wd, wdm);
let wh = client.create_from_slice(u32::as_bytes(&packed));
let xh = client.create_from_slice(f32::as_bytes(x));
let meta = client.create_from_slice(u32::as_bytes(&[k as u32]));
let meta = client.create_from_slice(u32::as_bytes(&[k as u32, nout as u32]));
let oh = client.create_from_slice(f32::as_bytes(&vec![0.0f32; nout]));
let grid = (nout as u32).div_ceil(nr as u32); // one workgroup per nr output rows
unsafe {
@@ -832,9 +907,10 @@ pub fn matvec_q4k_f32_blk_run<R: Runtime>(
ArrayArg::from_raw_parts(wh.clone(), packed.len()),
ArrayArg::from_raw_parts(xh.clone(), x.len()),
ArrayArg::from_raw_parts(oh.clone(), nout),
ArrayArg::from_raw_parts(meta.clone(), 1),
ArrayArg::from_raw_parts(meta.clone(), 2),
nt,
nr,
Target::of(client),
);
}
f32::from_bytes(&client.read_one_unchecked(oh)).to_vec()
@@ -911,7 +987,7 @@ impl<'a, R: Runtime> MatvecQ4kF32Eval<'a, R> {
let banks: Vec<Handle> =
(0..nbanks.max(1)).map(|_| client.create_from_slice(u32::as_bytes(&packed))).collect();
let xh = client.create_from_slice(f32::as_bytes(x));
let meta = client.create_from_slice(u32::as_bytes(&[k as u32]));
let meta = client.create_from_slice(u32::as_bytes(&[k as u32, rows as u32]));
let outh = client.create_from_slice(f32::as_bytes(&vec![0.0f32; rows]));
let oracle = matvec_q4k_ref(wqs, wsc, wd, wdm, x, rows, k);
let maxref = oracle.iter().fold(0f32, |a, &v| a.max(v.abs())).max(1e-30);
@@ -953,9 +1029,10 @@ impl<'a, R: Runtime> MatvecQ4kF32Eval<'a, R> {
ArrayArg::from_raw_parts(bank.clone(), self.packed_len),
ArrayArg::from_raw_parts(self.xh.clone(), self.x_len),
ArrayArg::from_raw_parts(self.outh.clone(), self.rows),
ArrayArg::from_raw_parts(self.meta.clone(), 1),
ArrayArg::from_raw_parts(self.meta.clone(), 2),
nt,
nr,
Target::of(self.client),
);
}
}
hanzo-ml/Cargo.toml
+1 -1
View File
@@ -1,6 +1,6 @@
[package]
name = "hanzo-ml"
version = "0.11.93"
version = "0.11.94"
edition = "2021"
description = "Fast multi-backend tensor & ML framework for Rust (CPU/CUDA/Metal/Vulkan/ROCm) with quantization — the compute core of the Hanzo stack."
repository = "https://github.com/hanzoai/ml"
hanzo-ml/src/vulkan/spv/matvec_q4k_f32_blk.spv
BIN
View File
Binary file not shown.
hanzo-ml/src/vulkan_backend.rs
+111 -2
View File
@@ -83,6 +83,11 @@ fn kernel_spv(name: &str) -> Result<&'static [u8]> {
"mul_mat_vec_q4k" => spv!("mul_mat_vec_q4k"),
"mul_mat_vec_q4k_sg" => spv!("mul_mat_vec_q4k_sg"),
"mul_mat_vec_q4k_cm" => spv!("mul_mat_vec_q4k_cm"),
// DSL twin of mul_mat_vec_q4k_cm (hanzo-kernel quant::matvec_q4k_f32_blk, nt=64 nr=1): the
// f32-direct cooperative Q4_K decode matvec, lowered from the DSL. Bindings wq(0) x(1) out(2)
// meta(3)=[k]; k and nout are runtime so ONE .spv serves every dense shape. VK_Q4K_DSL routes
// dense decode through it for the in-engine A/B against the hand mul_mat_vec_q4k_cm.
"matvec_q4k_f32_blk" => include_bytes!("vulkan/spv/matvec_q4k_f32_blk.spv"),
"mul_mat_vec_q5k" => spv!("mul_mat_vec_q5k"),
"mul_mat_vec_q5k_sg" => spv!("mul_mat_vec_q5k_sg"),
"mul_mat_vec_q6k" => spv!("mul_mat_vec_q6k"),
@@ -1294,8 +1299,9 @@ impl VulkanDevice {
crate::bail!("matvec_q4k_gpu: x count {} < k {k}", x.count);
}
let out = self.alloc_f32(nout)?;
// DEFAULT Q4_K decode path where the device advertises subgroup arithmetic: the coalesced
// multi-thread f32 matvec (mul_mat_vec_q4k_cm). TPB=16 threads cooperate on one super-block,
// Resident-MoE-bank Q4_K decode (woff != 0) + the VK_Q4K_DSL_OFF regression fallback + the
// bit-exact oracle for the DSL twin above: the coalesced multi-thread f32 matvec
// (mul_mat_vec_q4k_cm). TPB=16 threads cooperate on one super-block,
// adjacent threads reading adjacent qs u32 words (a coalesced 64-byte burst); Q4K_CM_ROWS rows
// per workgroup amortise the shared activation load; the per-row partial reduces with a subgroup
// add. The memory-BW lever for decode on this bandwidth-bound UMA APU -- the Q4_K analogue of
@@ -1303,6 +1309,30 @@ impl VulkanDevice {
// quantize dispatch drops. Decode is bit-identical to BlockQ4K::to_float; unpackHalf2x16 decodes
// the f16 d/dmin so subnormal scales stay exact. VK_Q4K_CM_OFF reverts to the dp4a block kernel
// for the A/B. Q4K_CM_ROWS MUST equal the shader's `NR` constant.
// Large dense Q4_K decode: the DSL f32-direct coalesced plane_sum matvec (hanzo-kernel
// matvec_q4k_f32_blk), the DSL twin of mul_mat_vec_q4k_cm lowered from ONE Rust source. Its
// Vulkan island arm mirrors the hand kernel's structure -- 16 threads per super-block reading
// ADJACENT qs words (coalesced burst), ITS super-blocks in parallel, plane_sum (subgroupAdd)
// reduction, no shared-mem tree. Reads x as f32 (no quantize) with k and nout in a meta buffer
// (cubecl has no push constants); grid = nout workgroups (nr=1 baked).
// The DSL twin earns the live path only where it MEASURES >= hand (campaign migration
// criterion): cold in-engine A/B on evo gfx1151 -- DSL/hand 1.022x at the dominant 6144-row
// shape (FFN gate/up, the bulk of the Q4_K bytes) but 0.947x @ 2048 and 0.871x @ 1024, where the
// plane_sum reduction's fixed cost outweighs the tiny matvec. So route large dense shapes to the
// DSL and keep hand for small ones -- the (device,shape) autotuner is the eventual owner of this
// split; Q4K_DSL_MIN_ROWS is the measured crossover until it lands. VK_Q4K_DSL_OFF forces hand.
// Dense only; the resident-MoE-bank case (woff != 0) stays on hand.
const Q4K_DSL_MIN_ROWS: usize = 4096;
if woff == 0
&& nout >= Q4K_DSL_MIN_ROWS
&& self.inner.subgroup_matvec
&& std::env::var_os("VK_Q4K_DSL_OFF").is_none()
{
let meta = self.upload_u32(&[k as u32, nout as u32])?;
let bufs = [wq.buffer, x.buffer, out.buffer, meta.buffer];
self.dispatch_out("matvec_q4k_f32_blk", &bufs, 2, &[], (nout as u32, 1, 1))?;
return Ok(out);
}
if self.inner.subgroup_matvec && std::env::var_os("VK_Q4K_CM_OFF").is_none() {
const Q4K_CM_ROWS: u32 = 2;
let push = push_u32(&[nout as u32, k as u32, woff as u32]);
@@ -8275,6 +8305,85 @@ mod dsl_dispatch_proof {
}
}
/// Cold in-engine A/B: the DSL f32-direct twin (`matvec_q4k_f32_blk`, routed by VK_Q4K_DSL) vs the
/// shipped hand kernel (`mul_mat_vec_q4k_cm`), BOTH through the production `dev.matvec_q4k` dispatch
/// at the live dense Qwen3-1.7B decode shapes, bank-rotated so each timed dispatch reads weight the
/// 32 MB MALL has evicted. Both are gated bit-exact against the same CPU oracle. `dev.matvec_q4k`
/// uploads x + reads back per call, an overhead identical for both arms, so the DSL/hand ratio is
/// the honest in-engine sign (the precise cold kernel GB/s is the matvec-check f32sweep + dp4a
/// anchor). HANZO_Q4K_AB=1 arms it; run in a quiet window only (wall-clock timing).
#[test]
fn q4k_dsl_vs_hand_cm_ab() {
if std::env::var_os("HANZO_Q4K_AB").is_none() {
eprintln!("[q4k-ab] set HANZO_Q4K_AB=1 to run the cold A/B (skipped)");
return;
}
use crate::quantized::k_quants::{BlockQ4K, GgmlType};
let dev = match VulkanDevice::new(0) {
Ok(d) => d,
Err(e) => { eprintln!("[q4k-ab] no vulkan device ({e}); skipping"); return; }
};
eprintln!("== cold in-engine A/B: DSL matvec_q4k_f32_blk vs hand mul_mat_vec_q4k_cm (dev.matvec_q4k) ==");
for &(nout, k) in &[(6144usize, 2048usize), (2048, 2048), (1024, 2048)] {
let nb = k / 256;
let gen_w = |row: usize, i: usize| (((row * 11 + i * 5) % 900) as f32) / 450.0 - 1.0;
let mut blocks: Vec<BlockQ4K> =
(0..nout * nb).map(|_| unsafe { std::mem::zeroed() }).collect();
let mut wdeq = vec![0f32; nout * k];
for r in 0..nout {
let rowf: Vec<f32> = (0..k).map(|i| gen_w(r, i)).collect();
BlockQ4K::from_float(&rowf, &mut blocks[r * nb..(r + 1) * nb]);
BlockQ4K::to_float(&blocks[r * nb..(r + 1) * nb], &mut wdeq[r * k..(r + 1) * k]);
}
let u32s: &[u32] = unsafe {
std::slice::from_raw_parts(blocks.as_ptr() as *const u32, blocks.len() * 36)
};
let wbytes = nout * (k / 256) * 144;
let nbanks = (64 * 1024 * 1024 / wbytes + 1).max(6);
let banks: Vec<_> = (0..nbanks).map(|_| dev.upload_u32(u32s).unwrap()).collect();
let x: Vec<f32> = (0..k).map(|i| (((i * 19 + 3) % 700) as f32) / 350.0 - 1.0).collect();
let mut cpu = vec![0f32; nout];
for (r, c) in cpu.iter_mut().enumerate() {
*c = (0..k).map(|i| wdeq[r * k + i] * x[i]).sum();
}
let maxref = cpu.iter().fold(0f32, |m, &v| m.max(v.abs())).max(1e-30);
let rel_of = |got: &[f32]| {
got.iter().zip(&cpu).map(|(a, b)| (a - b).abs()).fold(0f32, f32::max) / maxref
};
let meas = |dsl: bool| -> (f64, f32) {
if dsl {
unsafe { std::env::remove_var("VK_Q4K_DSL_OFF") };
} else {
unsafe { std::env::set_var("VK_Q4K_DSL_OFF", "1") };
}
let rel = rel_of(&dev.matvec_q4k(&banks[0], &x, nout, k).unwrap());
for i in 0..(2 * nbanks) {
let _ = dev.matvec_q4k(&banks[i % nbanks], &x, nout, k).unwrap();
}
let mut best = f64::INFINITY;
for _ in 0..8 {
let t = std::time::Instant::now();
for i in 0..nbanks {
let _ = dev.matvec_q4k(&banks[i % nbanks], &x, nout, k).unwrap();
}
best = best.min(t.elapsed().as_secs_f64() * 1e3 / nbanks as f64);
}
(best, rel)
};
let (hand_ms, hand_rel) = meas(false);
let (dsl_ms, dsl_rel) = meas(true);
unsafe { std::env::remove_var("VK_Q4K_DSL_OFF") };
let hand_g = wbytes as f64 / (hand_ms * 1e6);
let dsl_g = wbytes as f64 / (dsl_ms * 1e6);
eprintln!(
" {nout}x{k}: hand CM {hand_ms:.4}ms {hand_g:.0} GB/s (rel {hand_rel:.1e}) | DSL {dsl_ms:.4}ms {dsl_g:.0} GB/s (rel {dsl_rel:.1e}) | DSL/hand {:.3}x",
dsl_g / hand_g
);
assert!(hand_rel < 1e-3, "hand CM diverged from oracle: {hand_rel:.2e}");
assert!(dsl_rel < 1e-3, "DSL f32 diverged from oracle: {dsl_rel:.2e}");
}
}
// dp4a twin of the Q6_K proof: same split bank and shape, dispatched through
// `moe_matvec_blk_dp4a_gpu` with `with_xsum = false` (Q6_K derives its half-block activation sums
// in-register, so xsum is not a binding). This gates the GLUE — bank order + the no-xsum binding