//! Warp-per-row q4tp decode matvecs for the whole-token graph.
//!
//! `q4tp_matvec16nl` gives a 64-lane sub-block four rows, walks the row's
//! 16-byte groups two and a half at a time (5120 columns) and then pays a
//! six-barrier workgroup tree for every sixteen rows. On an RTX 3090 that
//! held Qwen3.8-27B's narrow projections at 37-52 % of the bus (GDN qkv+z
//! 349 GB/s, out_proj 357, FFN down 488). These kernels give a 32-lane
//! subgroup R rows (`CMF_MV_SG_R`, 4 or 2; gate+up `CMF_GU_SG_R`, 2 or 4),
//! two groups a lane a step with every load of the step issued before the
//! first FMA (clamped, unconditional), no barrier at all, and finish with
//! subgroup shuffles.
//!
//! Each row is the 16nl row to the bit: lane l holds the partials of the
//! 16nl lanes l (groups l, l+64, …) and l+32 (groups l+32, l+96, …) in two
//! accumulators, adds them (the tree's stride-32 step), and the shuffle
//! steps 16, 8, 4, 2, 1 take the partner's partial in the tree's own
//! order. Up to three matrices of one input per dispatch, and the fused
//! gate+up+SiLU of `q4tp_matvec16nl_gu`.
//!
//! `CMF_MV_SG=0` keeps the 16nl kernels.
use super::*;
pub(crate) struct MvSg {
pub(crate) mv: wgpu::ComputePipeline,
pub(crate) mv_l: wgpu::BindGroupLayout,
/// The column-split twin (`ks_src`) with its own auto layout (the
/// bindings are the same, the layouts are per pipeline); built only
/// for `CMF_MV_SG_KS=1` (measured worse, see `sg_prep`).
pub(crate) ks: Option<(wgpu::ComputePipeline, wgpu::BindGroupLayout)>,
/// The two-rows-a-subgroup kernel for the thin grids (`sg_prep`), built
/// for `CMF_MV_SG_THIN=1` only (measured worse).
pub(crate) thin: Option<(wgpu::ComputePipeline, wgpu::BindGroupLayout)>,
pub(crate) gu: wgpu::ComputePipeline,
pub(crate) gu_l: wgpu::BindGroupLayout,
/// Rows a subgroup owns.
pub(crate) rpw: usize,
/// Subgroups (warps) a workgroup of the plain matvec holds.
pub(crate) wpg: usize,
/// Rows a subgroup owns in the fused gate+up.
pub(crate) grpw: usize,
/// Never-written output for the unused matrix slots.
dummy_y: wgpu::Buffer,
}
fn sg32(c: &Ctx) -> bool {
c.device.features().contains(wgpu::Features::SUBGROUP)
&& c.adapter_info.subgroup_min_size == 32
&& c.adapter_info.subgroup_max_size == 32
}
fn shape(var_r: &str, r_default: usize) -> usize {
match std::env::var(var_r).as_deref() {
Ok("2") => 2,
Ok("4") => 4,
_ => r_default,
}
}
fn compile(c: &Ctx, label: &str, src: &str, entry: &str) -> Option<wgpu::ComputePipeline> {
let scope = c.device.push_error_scope(wgpu::ErrorFilter::Validation);
let module = c.device.create_shader_module(wgpu::ShaderModuleDescriptor {
label: Some(label),
source: wgpu::ShaderSource::Wgsl(src.into()),
});
let p = c
.device
.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
label: Some(entry),
layout: None,
module: &module,
entry_point: Some(entry),
compilation_options: crate::gpu_wgpu::pco(),
cache: c.pipeline_cache.as_ref(),
});
if let Some(e) = pollster::block_on(scope.pop()) {
tracing::warn!("{label} module rejected ({e}): q4tp decode keeps the 16nl kernels");
return None;
}
Some(p)
}
fn build(c: &Ctx) -> Option<MvSg> {
if std::env::var("CMF_MV_SG").as_deref() == Ok("0") || !sg32(c) {
return None;
}
// Rows a subgroup owns (R); a lane loads two 16-byte groups a step,
// every load of the step issued before its first FMA. The fused
// gate+up carries two weight rows per output row, so it has its own.
let rpw = shape("CMF_MV_SG_R", 4);
let grpw = shape("CMF_GU_SG_R", 2);
// Measured on Qwen3.8-27B / RTX 3090 plain decode (tok/s): R4U2 36.3,
// R2U2 33.0, R8U2 31.3, R2U4 32.2, R4U4 32.5, R1 24.0; only R4U2 and
// R2U2 are kept. Gate+up: R2U2 36.4, R4U2 36.1, R2U4 35.5, R1 33.1.
let (src, rpw) = match rpw {
2 => (MV_SG_R2U2, 2),
_ => (MV_SG_R4U2, 4),
};
// Warps a workgroup (`CMF_MV_SG_W`: 1, 2, 4 or 8; default 8). On a
// 188-SM RTX PRO 6000 a 5120-row matrix is 160 workgroups of eight
// warps — 28 SMs idle — and the obvious cure, two warps a workgroup
// (640 workgroups), MEASURED WORSE: the FFN down 87 → 100 us a layer,
// the o projection 31 → 35, q 48 → 52 (one and four warps the same).
// Eight warps of one workgroup share an SM's L1 for the activation
// vector; spread thin, every warp streams it from L2. Same rows,
// same bits whichever way: a warp's work does not depend on which
// workgroup holds it. Kept as the record and the knob.
let wpg: usize = match std::env::var("CMF_MV_SG_W").as_deref() {
Ok("1") => 1,
Ok("2") => 2,
Ok("4") => 4,
_ => 8,
};
let with_wpg = |src: &str| {
src.replace(
"@compute @workgroup_size(256)\nfn q4tp_mv_sg(",
&format!("@compute @workgroup_size({})\nfn q4tp_mv_sg(", 32 * wpg),
)
.replace("let per = 8u * 4u;", &format!("let per = {wpg}u * 4u;"))
.replace("let per = 8u * 2u;", &format!("let per = {wpg}u * 2u;"))
.replace("(blk * 8u + warp)", &format!("(blk * {wpg}u + warp)"))
};
let src = with_wpg(src);
debug_assert!(src.contains(&format!("@workgroup_size({})", 32 * wpg)));
let (gsrc, grpw) = match grpw {
4 => (GU_SG_R4U2, 4),
_ => (GU_SG_R2U2, 2),
};
let mv = compile(c, "cmf-mv-sg", &src, "q4tp_mv_sg")?;
// The column-split twin: the eight warps of a workgroup take the SAME
// four rows, each an eighth of the columns, and the eight partials of
// a row are added in a fixed order. Four rows a workgroup puts a
// 5120-row matrix on 1280 workgroups (6.8 an SM on the 188-SM RTX PRO
// 6000) where the row-split kernel has 160 and leaves 28 SMs idle.
// MEASURED WORSE on that card for every shape (`sg_prep`); a knob.
// Two rows a subgroup for the matrices whose four-row grid leaves SMs
// idle: a 5120-row projection is 160 workgroups of 32 rows on a
// 188-SM card, 320 of 16 rows this way. The same per-row arithmetic
// (two groups a lane a step, the same tree), so the same bits — and
// MEASURED WORSE on that card (RTX PRO 6000, 65B): the FFN down 87 ->
// 97 us a layer, the o projection 31 -> 35. Four rows a warp keep
// more loads in flight than the idle SMs cost. A knob: `CMF_MV_SG_THIN=1`.
let thin = (rpw == 4 && std::env::var("CMF_MV_SG_THIN").as_deref() == Ok("1"))
.then(|| compile(c, "cmf-mv-sg-thin", &with_wpg(MV_SG_R2U2), "q4tp_mv_sg"))
.flatten()
.map(|p| {
let l = p.get_bind_group_layout(0);
(p, l)
});
let ks = (rpw == 4 && std::env::var("CMF_MV_SG_KS").as_deref() == Ok("1"))
.then(|| compile(c, "cmf-mv-sg-ks", &ks_src(), "q4tp_mv_sg_ks"))
.flatten()
.map(|p| {
let l = p.get_bind_group_layout(0);
(p, l)
});
let gu = compile(c, "cmf-gu-sg", gsrc, "q4tp_gu_sg")?;
let dummy_y = c.device.create_buffer(&wgpu::BufferDescriptor {
label: Some("mv-sg-dummy"),
size: 16,
usage: wgpu::BufferUsages::STORAGE,
mapped_at_creation: false,
});
Some(MvSg {
mv_l: mv.get_bind_group_layout(0),
mv,
ks,
thin,
gu_l: gu.get_bind_group_layout(0),
gu,
rpw,
wpg,
grpw,
dummy_y,
})
}
/// `MV_SG_R4U2` with the columns split across the eight warps: `run_*`
/// get the row block from the workgroup alone, walk the warp's eighth of
/// the groups, and hand their four warp sums to shared memory; the entry
/// point adds the eight in warp order and writes the rows.
fn ks_src() -> String {
let mut s = MV_SG_R4U2.to_string();
// The warp's slice of the groups.
assert_eq!(s.matches(" let base = (blk * 8u + warp) * 4u;\n").count(), 3);
s = s.replace(
" let base = (blk * 8u + warp) * 4u;\n",
" let base = blk * 4u;\n let gq = (gpr + 7u) / 8u;\n let gend = min(gpr, (warp + 1u) * gq);\n",
);
// A dead row block must still reach the barrier: no early return, an
// empty slice instead.
let old_loop = " if (!l0) { return; }\n var g = lane;\n loop {\n if (g >= gpr) { break; }\n";
assert_eq!(s.matches(old_loop).count(), 3);
s = s.replace(
old_loop,
" var g = warp * gq + lane;\n if (!l0) { g = gend; }\n loop {\n if (g >= gend) { break; }\n",
);
assert_eq!(s.matches(" if (g + 32u < gpr) {\n").count(), 3);
s = s.replace(" if (g + 32u < gpr) {\n", " if (g + 32u < gend) {\n");
// The warp sums go to shared memory instead of the output.
for m in ["a", "b", "c"] {
let old = format!(
" if (lane == 0u) {{\n if (l0) {{ sg_y{m}[r0] = t0; }}\n if (l1) {{ sg_y{m}[r1] = t1; }}\n if (l2) {{ sg_y{m}[r2] = t2; }}\n if (l3) {{ sg_y{m}[r3] = t3; }}\n }}\n"
);
assert_eq!(s.matches(old.as_str()).count(), 1, "sum store of run_{m}");
s = s.replace(
old.as_str(),
" if (lane == 0u) {\n ks_part[warp * 4u + 0u] = t0;\n ks_part[warp * 4u + 1u] = t1;\n ks_part[warp * 4u + 2u] = t2;\n ks_part[warp * 4u + 3u] = t3;\n }\n",
);
}
s = s.replace("var<storage, read_write> sg_yc : array<f32>;\n", "var<storage, read_write> sg_yc : array<f32>;\nvar<workgroup> ks_part: array<f32, 32>;\n");
let old_entry = s.find("@compute @workgroup_size(256)\nfn q4tp_mv_sg(").expect("sg entry");
s.truncate(old_entry);
s.push_str(
"@compute @workgroup_size(256)\nfn q4tp_mv_sg_ks(@builtin(workgroup_id) wid: vec3<u32>,\n @builtin(local_invocation_index) lid: u32,\n @builtin(subgroup_invocation_id) lane: u32) {\n let warp = lid >> 5u;\n let ba = (sg_p.rows_a + 3u) / 4u;\n let bb = (sg_p.rows_b + 3u) / 4u;\n let wb = wid.x;\n var m = 0u;\n var blk = wb;\n if (wb < ba) {\n run_a(wb, warp, lane);\n } else if (wb < ba + bb) {\n m = 1u;\n blk = wb - ba;\n run_b(blk, warp, lane);\n } else {\n m = 2u;\n blk = wb - ba - bb;\n run_c(blk, warp, lane);\n }\n workgroupBarrier();\n if (lid < 4u) {\n var t = ks_part[lid];\n for (var w = 1u; w < 8u; w = w + 1u) { t = t + ks_part[w * 4u + lid]; }\n let r = blk * 4u + lid;\n if (m == 0u) { if (r < sg_p.rows_a) { sg_ya[r] = t; } }\n else if (m == 1u) { if (r < sg_p.rows_b) { sg_yb[r] = t; } }\n else { if (r < sg_p.rows_c) { sg_yc[r] = t; } }\n }\n}\n",
);
s
}
pub(crate) fn pipes(c: &Ctx) -> Option<&MvSg> {
c.dense_mv_pipes.get_or_init(|| build(c)).as_ref()
}
/// Up to three q4tp matvecs of one input `xs` (`mats`: weight, output,
/// rows) as ONE dispatch: (pipeline, bind group, workgroups). None when the
/// kernel is missing or the width is not whole groups.
pub(crate) fn sg_prep<'a>(
c: &'a Ctx,
xs: &wgpu::Buffer,
cols: usize,
mats: &[(&wgpu::Buffer, &wgpu::Buffer, usize)],
) -> Option<(&'a wgpu::ComputePipeline, wgpu::BindGroup, u32)> {
let p = pipes(c)?;
if cols % 32 != 0 || mats.is_empty() || mats.len() > 3 {
return None;
}
let r = |i: usize| mats.get(i).map_or(0, |m| m.2);
// Two knobs for the thin grids (fewer four-row workgroups than two an
// SM), both MEASURED WORSE on an RTX PRO 6000 and off by default: the
// two-row kernel (`CMF_MV_SG_THIN=1`, twice the workgroups, the same
// bits: down 87 -> 97 us a layer, o 31 -> 35) and the column-split one
// (`CMF_MV_SG_KS=1`: q 48 -> 72, o 31 -> 43, down 87 -> 97 — eight
// warps reading eight segments of the same four rows fragment the
// memory streams). Four rows a warp, eight warps a workgroup, keep the
// most bytes in flight; the 28 idle SMs cost less than either cure.
let per_rs = p.wpg * p.rpw;
let groups_rs: usize = (0..3).map(|i| r(i).div_ceil(per_rs)).sum();
let sm = c.sm_count.unwrap_or(48) as usize;
let thin_grid = groups_rs < 2 * sm;
let (pipe, layout, per) = match (p.ks.as_ref(), p.thin.as_ref()) {
(Some((ks, l)), _) if ks_forced() => (ks, l, 4),
(_, Some((t, l))) if thin_grid => (t, l, p.wpg * 2),
_ => (&p.mv, &p.mv_l, per_rs),
};
let groups: usize = (0..3).map(|i| r(i).div_ceil(per)).sum();
if groups == 0 || groups as u32 > MAX_WG {
return None;
}
let u = uniform_u32x8(
c,
[r(0) as u32, r(1) as u32, r(2) as u32, (cols / 32) as u32, 0, 0, 0, 0],
);
let w = |i: usize| mats.get(i).map_or(mats[0].0, |m| m.0);
let y = |i: usize| mats.get(i).map_or(&p.dummy_y, |m| m.1);
let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
label: Some("mv-sg"),
layout,
entries: &[
bind_buf(0, w(0)),
bind_buf(1, w(1)),
bind_buf(2, w(2)),
bind_buf(3, w(0)),
bind_buf(4, w(1)),
bind_buf(5, w(2)),
bind_buf(6, xs),
bind_buf(7, y(0)),
bind_buf(8, y(1)),
bind_buf(9, y(2)),
bind_buf(10, &u),
],
});
Some((pipe, bind, groups as u32))
}
/// `CMF_MV_SG_KS=1`: the column-split kernel for every shape (a measurement knob).
fn ks_forced() -> bool {
static F: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
*F.get_or_init(|| std::env::var("CMF_MV_SG_KS").as_deref() == Ok("1"))
}
/// gate+up+activation of `q4tp_matvec16nl_gu` (act code 0 SiLU, 1 exact
/// GELU; `lim` the swiglu limit, 0 for none) as ONE dispatch.
#[allow(clippy::too_many_arguments)]
pub(crate) fn sg_gu_prep<'a>(
c: &'a Ctx,
gate: &wgpu::Buffer,
up: &wgpu::Buffer,
xs: &wgpu::Buffer,
act: &wgpu::Buffer,
inter: usize,
cols: usize,
act_code: u32,
lim_bits: u32,
) -> Option<(&'a wgpu::ComputePipeline, wgpu::BindGroup, u32)> {
let p = pipes(c)?;
if cols % 32 != 0 || inter == 0 {
return None;
}
let groups = inter.div_ceil(8 * p.grpw);
if groups as u32 > MAX_WG {
return None;
}
let u = uniform_u32x8(
c,
[inter as u32, inter as u32, 0, (cols / 32) as u32, lim_bits, act_code, 0, 0],
);
let bind = c.device.create_bind_group(&wgpu::BindGroupDescriptor {
label: Some("gu-sg"),
layout: &p.gu_l,
entries: &[
bind_buf(0, gate),
bind_buf(1, up),
bind_buf(3, gate),
bind_buf(4, up),
bind_buf(6, xs),
bind_buf(7, act),
bind_buf(10, &u),
],
});
Some((&p.gu, bind, groups as u32))
}
pub(crate) const MV_SG_R2U2: &str = r#"
struct SgP { rows_a: u32, rows_b: u32, rows_c: u32, gpr: u32, lim: u32, act: u32, _p0: u32, _p1: u32 };
@group(0) @binding(0) var<storage, read> sg_wa : array<u32>;
@group(0) @binding(1) var<storage, read> sg_wb : array<u32>;
@group(0) @binding(2) var<storage, read> sg_wc : array<u32>;
@group(0) @binding(3) var<storage, read> sg_w4a : array<vec4<u32>>;
@group(0) @binding(4) var<storage, read> sg_w4b : array<vec4<u32>>;
@group(0) @binding(5) var<storage, read> sg_w4c : array<vec4<u32>>;
@group(0) @binding(6) var<storage, read> sg_x : array<vec4<f32>>;
@group(0) @binding(7) var<storage, read_write> sg_ya : array<f32>;
@group(0) @binding(8) var<storage, read_write> sg_yb : array<f32>;
@group(0) @binding(9) var<storage, read_write> sg_yc : array<f32>;
@group(0) @binding(10) var<uniform> sg_p : SgP;
fn sg_byte_a(off: u32) -> u32 { return (sg_wa[off >> 2u] >> ((off & 3u) * 8u)) & 0xFFu; }
fn sg_byte_b(off: u32) -> u32 { return (sg_wb[off >> 2u] >> ((off & 3u) * 8u)) & 0xFFu; }
fn sg_byte_c(off: u32) -> u32 { return (sg_wc[off >> 2u] >> ((off & 3u) * 8u)) & 0xFFu; }
fn sg_nib(w: u32, sh: u32) -> f32 {
return bitcast<f32>(((w >> sh) & 0xFu) | 0x4B000000u) - 8388616.0;
}
fn sg_dot8(w: u32, a: vec4<f32>, b: vec4<f32>) -> f32 {
return sg_nib(w, 0u) * a.x
+ sg_nib(w, 4u) * a.y
+ sg_nib(w, 8u) * a.z
+ sg_nib(w, 12u) * a.w
+ sg_nib(w, 16u) * b.x
+ sg_nib(w, 20u) * b.y
+ sg_nib(w, 24u) * b.z
+ sg_nib(w, 28u) * b.w;
}
// The 64-lane tree of `q4tp_matvec16nl` from stride 16 down: lane l takes
// lane l+s's partial, in the same order (own + partner).
fn sg_tree(v0: f32) -> f32 {
var v = v0;
v = v + subgroupShuffleDown(v, 16u);
v = v + subgroupShuffleDown(v, 8u);
v = v + subgroupShuffleDown(v, 4u);
v = v + subgroupShuffleDown(v, 2u);
v = v + subgroupShuffleDown(v, 1u);
return v;
}
fn sg_erf(x: f32) -> f32 {
let a = abs(x);
let t = 1.0 / (1.0 + 0.3275911 * a);
let y = 1.0
- (((((1.0614054 * t - 1.4531521) * t + 1.4214138) * t - 0.28449674) * t
+ 0.2548296)
* t)
* exp(-a * a);
return select(y, -y, x < 0.0);
}
fn sg_gelu_erf(x: f32) -> f32 {
return 0.5 * x * (1.0 + sg_erf(x * 0.70710678));
}
fn sg_gelu_tanh(x: f32) -> f32 {
return 0.5 * x * (1.0 + tanh(0.7978846 * (x + 0.044715 * x * x * x)));
}
fn run_a(blk: u32, warp: u32, lane: u32) {
let rows = sg_p.rows_a;
let gpr = sg_p.gpr;
let params_w = rows * gpr * 4u;
let codes_b = rows * gpr * 16u + rows * 4u;
let cstride = (gpr * 5u + 7u) / 8u;
let base = (blk * 8u + warp) * 2u;
let r0 = base + 0u;
let l0 = r0 < rows;
let q0 = select(base, r0, l0);
let cm0 = codes_b + q0 * cstride;
let pm0 = unpack2x16float(sg_wa[params_w + q0]);
var ma0 = 0.0;
var mb0 = 0.0;
let r1 = base + 1u;
let l1 = r1 < rows;
let q1 = select(base, r1, l1);
let cm1 = codes_b + q1 * cstride;
let pm1 = unpack2x16float(sg_wa[params_w + q1]);
var ma1 = 0.0;
var mb1 = 0.0;
if (!l0) { return; }
var g = lane;
loop {
if (g >= gpr) { break; }
let g0 = min(g + 0u, gpr - 1u);
let bit0 = g0 * 5u;
let cbo0 = bit0 >> 3u;
let sh0 = bit0 & 7u;
let xo0 = g0 * 8u;
let xa0 = sg_x[xo0 + 0u];
let xb0 = sg_x[xo0 + 1u];
let xc0 = sg_x[xo0 + 2u];
let xd0 = sg_x[xo0 + 3u];
let xe0 = sg_x[xo0 + 4u];
let xf0 = sg_x[xo0 + 5u];
let xg0 = sg_x[xo0 + 6u];
let xh0 = sg_x[xo0 + 7u];
let g1 = min(g + 32u, gpr - 1u);
let bit1 = g1 * 5u;
let cbo1 = bit1 >> 3u;
let sh1 = bit1 & 7u;
let xo1 = g1 * 8u;
let xa1 = sg_x[xo1 + 0u];
let xb1 = sg_x[xo1 + 1u];
let xc1 = sg_x[xo1 + 2u];
let xd1 = sg_x[xo1 + 3u];
let xe1 = sg_x[xo1 + 4u];
let xf1 = sg_x[xo1 + 5u];
let xg1 = sg_x[xo1 + 6u];
let xh1 = sg_x[xo1 + 7u];
var cvm00 = sg_byte_a(cm0 + cbo0);
if (sh0 > 3u) { cvm00 = cvm00 | (sg_byte_a(cm0 + cbo0 + 1u) << 8u); }
let vm00 = sg_w4a[q0 * gpr + g0];
var cvm01 = sg_byte_a(cm0 + cbo1);
if (sh1 > 3u) { cvm01 = cvm01 | (sg_byte_a(cm0 + cbo1 + 1u) << 8u); }
let vm01 = sg_w4a[q0 * gpr + g1];
var cvm10 = sg_byte_a(cm1 + cbo0);
if (sh0 > 3u) { cvm10 = cvm10 | (sg_byte_a(cm1 + cbo0 + 1u) << 8u); }
let vm10 = sg_w4a[q1 * gpr + g0];
var cvm11 = sg_byte_a(cm1 + cbo1);
if (sh1 > 3u) { cvm11 = cvm11 | (sg_byte_a(cm1 + cbo1 + 1u) << 8u); }
let vm11 = sg_w4a[q1 * gpr + g1];
ma0 = ma0 + exp2(pm0.x + f32((cvm00 >> sh0) & 31u) * pm0.y)
* (sg_dot8(vm00.x, xa0, xb0) + sg_dot8(vm00.y, xc0, xd0)
+ sg_dot8(vm00.z, xe0, xf0) + sg_dot8(vm00.w, xg0, xh0));
ma1 = ma1 + exp2(pm1.x + f32((cvm10 >> sh0) & 31u) * pm1.y)
* (sg_dot8(vm10.x, xa0, xb0) + sg_dot8(vm10.y, xc0, xd0)
+ sg_dot8(vm10.z, xe0, xf0) + sg_dot8(vm10.w, xg0, xh0));
if (g + 32u < gpr) {
mb0 = mb0 + exp2(pm0.x + f32((cvm01 >> sh1) & 31u) * pm0.y)
* (sg_dot8(vm01.x, xa1, xb1) + sg_dot8(vm01.y, xc1, xd1)
+ sg_dot8(vm01.z, xe1, xf1) + sg_dot8(vm01.w, xg1, xh1));
mb1 = mb1 + exp2(pm1.x + f32((cvm11 >> sh1) & 31u) * pm1.y)
* (sg_dot8(vm11.x, xa1, xb1) + sg_dot8(vm11.y, xc1, xd1)
+ sg_dot8(vm11.z, xe1, xf1) + sg_dot8(vm11.w, xg1, xh1));
}
g = g + 64u;
}
let t0 = sg_tree(ma0 + mb0);
let t1 = sg_tree(ma1 + mb1);
if (lane == 0u) {
if (l0) { sg_ya[r0] = t0; }
if (l1) { sg_ya[r1] = t1; }
}
}
fn run_b(blk: u32, warp: u32, lane: u32) {
let rows = sg_p.rows_b;
let gpr = sg_p.gpr;
let params_w = rows * gpr * 4u;
let codes_b = rows * gpr * 16u + rows * 4u;
let cstride = (gpr * 5u + 7u) / 8u;
let base = (blk * 8u + warp) * 2u;
let r0 = base + 0u;
let l0 = r0 < rows;
let q0 = select(base, r0, l0);
let cm0 = codes_b + q0 * cstride;
let pm0 = unpack2x16float(sg_wb[params_w + q0]);
var ma0 = 0.0;
var mb0 = 0.0;
let r1 = base + 1u;
let l1 = r1 < rows;
let q1 = select(base, r1, l1);
let cm1 = codes_b + q1 * cstride;
let pm1 = unpack2x16float(sg_wb[params_w + q1]);
var ma1 = 0.0;
var mb1 = 0.0;
if (!l0) { return; }
var g = lane;
loop {
if (g >= gpr) { break; }
let g0 = min(g + 0u, gpr - 1u);
let bit0 = g0 * 5u;
let cbo0 = bit0 >> 3u;
let sh0 = bit0 & 7u;
let xo0 = g0 * 8u;
let xa0 = sg_x[xo0 + 0u];
let xb0 = sg_x[xo0 + 1u];
let xc0 = sg_x[xo0 + 2u];
let xd0 = sg_x[xo0 + 3u];
let xe0 = sg_x[xo0 + 4u];
let xf0 = sg_x[xo0 + 5u];
let xg0 = sg_x[xo0 + 6u];
let xh0 = sg_x[xo0 + 7u];
let g1 = min(g + 32u, gpr - 1u);
let bit1 = g1 * 5u;
let cbo1 = bit1 >> 3u;
let sh1 = bit1 & 7u;
let xo1 = g1 * 8u;
let xa1 = sg_x[xo1 + 0u];
let xb1 = sg_x[xo1 + 1u];
let xc1 = sg_x[xo1 + 2u];
let xd1 = sg_x[xo1 + 3u];
let xe1 = sg_x[xo1 + 4u];
let xf1 = sg_x[xo1 + 5u];
let xg1 = sg_x[xo1 + 6u];
let xh1 = sg_x[xo1 + 7u];
var cvm00 = sg_byte_b(cm0 + cbo0);
if (sh0 > 3u) { cvm00 = cvm00 | (sg_byte_b(cm0 + cbo0 + 1u) << 8u); }
let vm00 = sg_w4b[q0 * gpr + g0];
var cvm01 = sg_byte_b(cm0 + cbo1);
if (sh1 > 3u) { cvm01 = cvm01 | (sg_byte_b(cm0 + cbo1 + 1u) << 8u); }
let vm01 = sg_w4b[q0 * gpr + g1];
var cvm10 = sg_byte_b(cm1 + cbo0);
if (sh0 > 3u) { cvm10 = cvm10 | (sg_byte_b(cm1 + cbo0 + 1u) << 8u); }
let vm10 = sg_w4b[q1 * gpr + g0];
var cvm11 = sg_byte_b(cm1 + cbo1);
if (sh1 > 3u) { cvm11 = cvm11 | (sg_byte_b(cm1 + cbo1 + 1u) << 8u); }
let vm11 = sg_w4b[q1 * gpr + g1];
ma0 = ma0 + exp2(pm0.x + f32((cvm00 >> sh0) & 31u) * pm0.y)
* (sg_dot8(vm00.x, xa0, xb0) + sg_dot8(vm00.y, xc0, xd0)
+ sg_dot8(vm00.z, xe0, xf0) + sg_dot8(vm00.w, xg0, xh0));
ma1 = ma1 + exp2(pm1.x + f32((cvm10 >> sh0) & 31u) * pm1.y)
* (sg_dot8(vm10.x, xa0, xb0) + sg_dot8(vm10.y, xc0, xd0)
+ sg_dot8(vm10.z, xe0, xf0) + sg_dot8(vm10.w, xg0, xh0));
if (g + 32u < gpr) {
mb0 = mb0 + exp2(pm0.x + f32((cvm01 >> sh1) & 31u) * pm0.y)
* (sg_dot8(vm01.x, xa1, xb1) + sg_dot8(vm01.y, xc1, xd1)
+ sg_dot8(vm01.z, xe1, xf1) + sg_dot8(vm01.w, xg1, xh1));
mb1 = mb1 + exp2(pm1.x + f32((cvm11 >> sh1) & 31u) * pm1.y)
* (sg_dot8(vm11.x, xa1, xb1) + sg_dot8(vm11.y, xc1, xd1)
+ sg_dot8(vm11.z, xe1, xf1) + sg_dot8(vm11.w, xg1, xh1));
}
g = g + 64u;
}
let t0 = sg_tree(ma0 + mb0);
let t1 = sg_tree(ma1 + mb1);
if (lane == 0u) {
if (l0) { sg_yb[r0] = t0; }
if (l1) { sg_yb[r1] = t1; }
}
}
fn run_c(blk: u32, warp: u32, lane: u32) {
let rows = sg_p.rows_c;
let gpr = sg_p.gpr;
let params_w = rows * gpr * 4u;
let codes_b = rows * gpr * 16u + rows * 4u;
let cstride = (gpr * 5u + 7u) / 8u;
let base = (blk * 8u + warp) * 2u;
let r0 = base + 0u;
let l0 = r0 < rows;
let q0 = select(base, r0, l0);
let cm0 = codes_b + q0 * cstride;
let pm0 = unpack2x16float(sg_wc[params_w + q0]);
var ma0 = 0.0;
var mb0 = 0.0;
let r1 = base + 1u;
let l1 = r1 < rows;
let q1 = select(base, r1, l1);
let cm1 = codes_b + q1 * cstride;
let pm1 = unpack2x16float(sg_wc[params_w + q1]);
var ma1 = 0.0;
var mb1 = 0.0;
if (!l0) { return; }
var g = lane;
loop {
if (g >= gpr) { break; }
let g0 = min(g + 0u, gpr - 1u);
let bit0 = g0 * 5u;
let cbo0 = bit0 >> 3u;
let sh0 = bit0 & 7u;
let xo0 = g0 * 8u;
let xa0 = sg_x[xo0 + 0u];
let xb0 = sg_x[xo0 + 1u];
let xc0 = sg_x[xo0 + 2u];
let xd0 = sg_x[xo0 + 3u];
let xe0 = sg_x[xo0 + 4u];
let xf0 = sg_x[xo0 + 5u];
let xg0 = sg_x[xo0 + 6u];
let xh0 = sg_x[xo0 + 7u];
let g1 = min(g + 32u, gpr - 1u);
let bit1 = g1 * 5u;
let cbo1 = bit1 >> 3u;
let sh1 = bit1 & 7u;
let xo1 = g1 * 8u;
let xa1 = sg_x[xo1 + 0u];
let xb1 = sg_x[xo1 + 1u];
let xc1 = sg_x[xo1 + 2u];
let xd1 = sg_x[xo1 + 3u];
let xe1 = sg_x[xo1 + 4u];
let xf1 = sg_x[xo1 + 5u];
let xg1 = sg_x[xo1 + 6u];
let xh1 = sg_x[xo1 + 7u];
var cvm00 = sg_byte_c(cm0 + cbo0);
if (sh0 > 3u) { cvm00 = cvm00 | (sg_byte_c(cm0 + cbo0 + 1u) << 8u); }
let vm00 = sg_w4c[q0 * gpr + g0];
var cvm01 = sg_byte_c(cm0 + cbo1);
if (sh1 > 3u) { cvm01 = cvm01 | (sg_byte_c(cm0 + cbo1 + 1u) << 8u); }
let vm01 = sg_w4c[q0 * gpr + g1];
var cvm10 = sg_byte_c(cm1 + cbo0);
if (sh0 > 3u) { cvm10 = cvm10 | (sg_byte_c(cm1 + cbo0 + 1u) << 8u); }
let vm10 = sg_w4c[q1 * gpr + g0];
var cvm11 = sg_byte_c(cm1 + cbo1);
if (sh1 > 3u) { cvm11 = cvm11 | (sg_byte_c(cm1 + cbo1 + 1u) << 8u); }
let vm11 = sg_w4c[q1 * gpr + g1];
ma0 = ma0 + exp2(pm0.x + f32((cvm00 >> sh0) & 31u) * pm0.y)
* (sg_dot8(vm00.x, xa0, xb0) + sg_dot8(vm00.y, xc0, xd0)
+ sg_dot8(vm00.z, xe0, xf0) + sg_dot8(vm00.w, xg0, xh0));
ma1 = ma1 + exp2(pm1.x + f32((cvm10 >> sh0) & 31u) * pm1.y)
* (sg_dot8(vm10.x, xa0, xb0) + sg_dot8(vm10.y, xc0, xd0)
+ sg_dot8(vm10.z, xe0, xf0) + sg_dot8(vm10.w, xg0, xh0));
if (g + 32u < gpr) {
mb0 = mb0 + exp2(pm0.x + f32((cvm01 >> sh1) & 31u) * pm0.y)
* (sg_dot8(vm01.x, xa1, xb1) + sg_dot8(vm01.y, xc1, xd1)
+ sg_dot8(vm01.z, xe1, xf1) + sg_dot8(vm01.w, xg1, xh1));
mb1 = mb1 + exp2(pm1.x + f32((cvm11 >> sh1) & 31u) * pm1.y)
* (sg_dot8(vm11.x, xa1, xb1) + sg_dot8(vm11.y, xc1, xd1)
+ sg_dot8(vm11.z, xe1, xf1) + sg_dot8(vm11.w, xg1, xh1));
}
g = g + 64u;
}
let t0 = sg_tree(ma0 + mb0);
let t1 = sg_tree(ma1 + mb1);
if (lane == 0u) {
if (l0) { sg_yc[r0] = t0; }
if (l1) { sg_yc[r1] = t1; }
}
}
@compute @workgroup_size(256)
fn q4tp_mv_sg(@builtin(workgroup_id) wid: vec3<u32>,
@builtin(local_invocation_index) lid: u32,
@builtin(subgroup_invocation_id) lane: u32) {
let per = 8u * 2u;
let warp = lid >> 5u;
let ba = (sg_p.rows_a + per - 1u) / per;
let bb = (sg_p.rows_b + per - 1u) / per;
let wb = wid.x;
if (wb < ba) {
run_a(wb, warp, lane);
} else if (wb < ba + bb) {
run_b(wb - ba, warp, lane);
} else {
run_c(wb - ba - bb, warp, lane);
}
}
"#;
pub(crate) const MV_SG_R4U2: &str = r#"
struct SgP { rows_a: u32, rows_b: u32, rows_c: u32, gpr: u32, lim: u32, act: u32, _p0: u32, _p1: u32 };
@group(0) @binding(0) var<storage, read> sg_wa : array<u32>;
@group(0) @binding(1) var<storage, read> sg_wb : array<u32>;
@group(0) @binding(2) var<storage, read> sg_wc : array<u32>;
@group(0) @binding(3) var<storage, read> sg_w4a : array<vec4<u32>>;
@group(0) @binding(4) var<storage, read> sg_w4b : array<vec4<u32>>;
@group(0) @binding(5) var<storage, read> sg_w4c : array<vec4<u32>>;
@group(0) @binding(6) var<storage, read> sg_x : array<vec4<f32>>;
@group(0) @binding(7) var<storage, read_write> sg_ya : array<f32>;
@group(0) @binding(8) var<storage, read_write> sg_yb : array<f32>;
@group(0) @binding(9) var<storage, read_write> sg_yc : array<f32>;
@group(0) @binding(10) var<uniform> sg_p : SgP;
fn sg_byte_a(off: u32) -> u32 { return (sg_wa[off >> 2u] >> ((off & 3u) * 8u)) & 0xFFu; }
fn sg_byte_b(off: u32) -> u32 { return (sg_wb[off >> 2u] >> ((off & 3u) * 8u)) & 0xFFu; }
fn sg_byte_c(off: u32) -> u32 { return (sg_wc[off >> 2u] >> ((off & 3u) * 8u)) & 0xFFu; }
fn sg_nib(w: u32, sh: u32) -> f32 {
return bitcast<f32>(((w >> sh) & 0xFu) | 0x4B000000u) - 8388616.0;
}
fn sg_dot8(w: u32, a: vec4<f32>, b: vec4<f32>) -> f32 {
return sg_nib(w, 0u) * a.x
+ sg_nib(w, 4u) * a.y
+ sg_nib(w, 8u) * a.z
+ sg_nib(w, 12u) * a.w
+ sg_nib(w, 16u) * b.x
+ sg_nib(w, 20u) * b.y
+ sg_nib(w, 24u) * b.z
+ sg_nib(w, 28u) * b.w;
}
// The 64-lane tree of `q4tp_matvec16nl` from stride 16 down: lane l takes
// lane l+s's partial, in the same order (own + partner).
fn sg_tree(v0: f32) -> f32 {
var v = v0;
v = v + subgroupShuffleDown(v, 16u);
v = v + subgroupShuffleDown(v, 8u);
v = v + subgroupShuffleDown(v, 4u);
v = v + subgroupShuffleDown(v, 2u);
v = v + subgroupShuffleDown(v, 1u);
return v;
}
fn sg_erf(x: f32) -> f32 {
let a = abs(x);
let t = 1.0 / (1.0 + 0.3275911 * a);
let y = 1.0
- (((((1.0614054 * t - 1.4531521) * t + 1.4214138) * t - 0.28449674) * t
+ 0.2548296)
* t)
* exp(-a * a);
return select(y, -y, x < 0.0);
}
fn sg_gelu_erf(x: f32) -> f32 {
return 0.5 * x * (1.0 + sg_erf(x * 0.70710678));
}
fn sg_gelu_tanh(x: f32) -> f32 {
return 0.5 * x * (1.0 + tanh(0.7978846 * (x + 0.044715 * x * x * x)));
}
fn run_a(blk: u32, warp: u32, lane: u32) {
let rows = sg_p.rows_a;
let gpr = sg_p.gpr;
let params_w = rows * gpr * 4u;
let codes_b = rows * gpr * 16u + rows * 4u;
let cstride = (gpr * 5u + 7u) / 8u;
let base = (blk * 8u + warp) * 4u;
let r0 = base + 0u;
let l0 = r0 < rows;
let q0 = select(base, r0, l0);
let cm0 = codes_b + q0 * cstride;
let pm0 = unpack2x16float(sg_wa[params_w + q0]);
var ma0 = 0.0;
var mb0 = 0.0;
let r1 = base + 1u;
let l1 = r1 < rows;
let q1 = select(base, r1, l1);
let cm1 = codes_b + q1 * cstride;
let pm1 = unpack2x16float(sg_wa[params_w + q1]);
var ma1 = 0.0;
var mb1 = 0.0;
let r2 = base + 2u;
let l2 = r2 < rows;
let q2 = select(base, r2, l2);
let cm2 = codes_b + q2 * cstride;
let pm2 = unpack2x16float(sg_wa[params_w + q2]);
var ma2 = 0.0;
var mb2 = 0.0;
let r3 = base + 3u;
let l3 = r3 < rows;
let q3 = select(base, r3, l3);
let cm3 = codes_b + q3 * cstride;
let pm3 = unpack2x16float(sg_wa[params_w + q3]);
var ma3 = 0.0;
var mb3 = 0.0;
if (!l0) { return; }
var g = lane;
loop {
if (g >= gpr) { break; }
let g0 = min(g + 0u, gpr - 1u);
let bit0 = g0 * 5u;
let cbo0 = bit0 >> 3u;
let sh0 = bit0 & 7u;
let xo0 = g0 * 8u;
let xa0 = sg_x[xo0 + 0u];
let xb0 = sg_x[xo0 + 1u];
let xc0 = sg_x[xo0 + 2u];
let xd0 = sg_x[xo0 + 3u];
let xe0 = sg_x[xo0 + 4u];
let xf0 = sg_x[xo0 + 5u];
let xg0 = sg_x[xo0 + 6u];
let xh0 = sg_x[xo0 + 7u];
let g1 = min(g + 32u, gpr - 1u);
let bit1 = g1 * 5u;
let cbo1 = bit1 >> 3u;
let sh1 = bit1 & 7u;
let xo1 = g1 * 8u;
let xa1 = sg_x[xo1 + 0u];
let xb1 = sg_x[xo1 + 1u];
let xc1 = sg_x[xo1 + 2u];
let xd1 = sg_x[xo1 + 3u];
let xe1 = sg_x[xo1 + 4u];
let xf1 = sg_x[xo1 + 5u];
let xg1 = sg_x[xo1 + 6u];
let xh1 = sg_x[xo1 + 7u];
var cvm00 = sg_byte_a(cm0 + cbo0);
if (sh0 > 3u) { cvm00 = cvm00 | (sg_byte_a(cm0 + cbo0 + 1u) << 8u); }
let vm00 = sg_w4a[q0 * gpr + g0];
var cvm01 = sg_byte_a(cm0 + cbo1);
if (sh1 > 3u) { cvm01 = cvm01 | (sg_byte_a(cm0 + cbo1 + 1u) << 8u); }
let vm01 = sg_w4a[q0 * gpr + g1];
var cvm10 = sg_byte_a(cm1 + cbo0);
if (sh0 > 3u) { cvm10 = cvm10 | (sg_byte_a(cm1 + cbo0 + 1u) << 8u); }
let vm10 = sg_w4a[q1 * gpr + g0];
var cvm11 = sg_byte_a(cm1 + cbo1);
if (sh1 > 3u) { cvm11 = cvm11 | (sg_byte_a(cm1 + cbo1 + 1u) << 8u); }
let vm11 = sg_w4a[q1 * gpr + g1];
var cvm20 = sg_byte_a(cm2 + cbo0);
if (sh0 > 3u) { cvm20 = cvm20 | (sg_byte_a(cm2 + cbo0 + 1u) << 8u); }
let vm20 = sg_w4a[q2 * gpr + g0];
var cvm21 = sg_byte_a(cm2 + cbo1);
if (sh1 > 3u) { cvm21 = cvm21 | (sg_byte_a(cm2 + cbo1 + 1u) << 8u); }
let vm21 = sg_w4a[q2 * gpr + g1];
var cvm30 = sg_byte_a(cm3 + cbo0);
if (sh0 > 3u) { cvm30 = cvm30 | (sg_byte_a(cm3 + cbo0 + 1u) << 8u); }
let vm30 = sg_w4a[q3 * gpr + g0];
var cvm31 = sg_byte_a(cm3 + cbo1);
if (sh1 > 3u) { cvm31 = cvm31 | (sg_byte_a(cm3 + cbo1 + 1u) << 8u); }
let vm31 = sg_w4a[q3 * gpr + g1];
ma0 = ma0 + exp2(pm0.x + f32((cvm00 >> sh0) & 31u) * pm0.y)
* (sg_dot8(vm00.x, xa0, xb0) + sg_dot8(vm00.y, xc0, xd0)
+ sg_dot8(vm00.z, xe0, xf0) + sg_dot8(vm00.w, xg0, xh0));
ma1 = ma1 + exp2(pm1.x + f32((cvm10 >> sh0) & 31u) * pm1.y)
* (sg_dot8(vm10.x, xa0, xb0) + sg_dot8(vm10.y, xc0, xd0)
+ sg_dot8(vm10.z, xe0, xf0) + sg_dot8(vm10.w, xg0, xh0));
ma2 = ma2 + exp2(pm2.x + f32((cvm20 >> sh0) & 31u) * pm2.y)
* (sg_dot8(vm20.x, xa0, xb0) + sg_dot8(vm20.y, xc0, xd0)
+ sg_dot8(vm20.z, xe0, xf0) + sg_dot8(vm20.w, xg0, xh0));
ma3 = ma3 + exp2(pm3.x + f32((cvm30 >> sh0) & 31u) * pm3.y)
* (sg_dot8(vm30.x, xa0, xb0) + sg_dot8(vm30.y, xc0, xd0)
+ sg_dot8(vm30.z, xe0, xf0) + sg_dot8(vm30.w, xg0, xh0));
if (g + 32u < gpr) {
mb0 = mb0 + exp2(pm0.x + f32((cvm01 >> sh1) & 31u) * pm0.y)
* (sg_dot8(vm01.x, xa1, xb1) + sg_dot8(vm01.y, xc1, xd1)
+ sg_dot8(vm01.z, xe1, xf1) + sg_dot8(vm01.w, xg1, xh1));
mb1 = mb1 + exp2(pm1.x + f32((cvm11 >> sh1) & 31u) * pm1.y)
* (sg_dot8(vm11.x, xa1, xb1) + sg_dot8(vm11.y, xc1, xd1)
+ sg_dot8(vm11.z, xe1, xf1) + sg_dot8(vm11.w, xg1, xh1));
mb2 = mb2 + exp2(pm2.x + f32((cvm21 >> sh1) & 31u) * pm2.y)
* (sg_dot8(vm21.x, xa1, xb1) + sg_dot8(vm21.y, xc1, xd1)
+ sg_dot8(vm21.z, xe1, xf1) + sg_dot8(vm21.w, xg1, xh1));
mb3 = mb3 + exp2(pm3.x + f32((cvm31 >> sh1) & 31u) * pm3.y)
* (sg_dot8(vm31.x, xa1, xb1) + sg_dot8(vm31.y, xc1, xd1)
+ sg_dot8(vm31.z, xe1, xf1) + sg_dot8(vm31.w, xg1, xh1));
}
g = g + 64u;
}
let t0 = sg_tree(ma0 + mb0);
let t1 = sg_tree(ma1 + mb1);
let t2 = sg_tree(ma2 + mb2);
let t3 = sg_tree(ma3 + mb3);
if (lane == 0u) {
if (l0) { sg_ya[r0] = t0; }
if (l1) { sg_ya[r1] = t1; }
if (l2) { sg_ya[r2] = t2; }
if (l3) { sg_ya[r3] = t3; }
}
}
fn run_b(blk: u32, warp: u32, lane: u32) {
let rows = sg_p.rows_b;
let gpr = sg_p.gpr;
let params_w = rows * gpr * 4u;
let codes_b = rows * gpr * 16u + rows * 4u;
let cstride = (gpr * 5u + 7u) / 8u;
let base = (blk * 8u + warp) * 4u;
let r0 = base + 0u;
let l0 = r0 < rows;
let q0 = select(base, r0, l0);
let cm0 = codes_b + q0 * cstride;
let pm0 = unpack2x16float(sg_wb[params_w + q0]);
var ma0 = 0.0;
var mb0 = 0.0;
let r1 = base + 1u;
let l1 = r1 < rows;
let q1 = select(base, r1, l1);
let cm1 = codes_b + q1 * cstride;
let pm1 = unpack2x16float(sg_wb[params_w + q1]);
var ma1 = 0.0;
var mb1 = 0.0;
let r2 = base + 2u;
let l2 = r2 < rows;
let q2 = select(base, r2, l2);
let cm2 = codes_b + q2 * cstride;
let pm2 = unpack2x16float(sg_wb[params_w + q2]);
var ma2 = 0.0;
var mb2 = 0.0;
let r3 = base + 3u;
let l3 = r3 < rows;
let q3 = select(base, r3, l3);
let cm3 = codes_b + q3 * cstride;
let pm3 = unpack2x16float(sg_wb[params_w + q3]);
var ma3 = 0.0;
var mb3 = 0.0;
if (!l0) { return; }
var g = lane;
loop {
if (g >= gpr) { break; }
let g0 = min(g + 0u, gpr - 1u);
let bit0 = g0 * 5u;
let cbo0 = bit0 >> 3u;
let sh0 = bit0 & 7u;
let xo0 = g0 * 8u;
let xa0 = sg_x[xo0 + 0u];
let xb0 = sg_x[xo0 + 1u];
let xc0 = sg_x[xo0 + 2u];
let xd0 = sg_x[xo0 + 3u];
let xe0 = sg_x[xo0 + 4u];
let xf0 = sg_x[xo0 + 5u];
let xg0 = sg_x[xo0 + 6u];
let xh0 = sg_x[xo0 + 7u];
let g1 = min(g + 32u, gpr - 1u);
let bit1 = g1 * 5u;
let cbo1 = bit1 >> 3u;
let sh1 = bit1 & 7u;
let xo1 = g1 * 8u;
let xa1 = sg_x[xo1 + 0u];
let xb1 = sg_x[xo1 + 1u];
let xc1 = sg_x[xo1 + 2u];
let xd1 = sg_x[xo1 + 3u];
let xe1 = sg_x[xo1 + 4u];
let xf1 = sg_x[xo1 + 5u];
let xg1 = sg_x[xo1 + 6u];
let xh1 = sg_x[xo1 + 7u];
var cvm00 = sg_byte_b(cm0 + cbo0);
if (sh0 > 3u) { cvm00 = cvm00 | (sg_byte_b(cm0 + cbo0 + 1u) << 8u); }
let vm00 = sg_w4b[q0 * gpr + g0];
var cvm01 = sg_byte_b(cm0 + cbo1);
if (sh1 > 3u) { cvm01 = cvm01 | (sg_byte_b(cm0 + cbo1 + 1u) << 8u); }
let vm01 = sg_w4b[q0 * gpr + g1];
var cvm10 = sg_byte_b(cm1 + cbo0);
if (sh0 > 3u) { cvm10 = cvm10 | (sg_byte_b(cm1 + cbo0 + 1u) << 8u); }
let vm10 = sg_w4b[q1 * gpr + g0];
var cvm11 = sg_byte_b(cm1 + cbo1);
if (sh1 > 3u) { cvm11 = cvm11 | (sg_byte_b(cm1 + cbo1 + 1u) << 8u); }
let vm11 = sg_w4b[q1 * gpr + g1];
var cvm20 = sg_byte_b(cm2 + cbo0);
if (sh0 > 3u) { cvm20 = cvm20 | (sg_byte_b(cm2 + cbo0 + 1u) << 8u); }
let vm20 = sg_w4b[q2 * gpr + g0];
var cvm21 = sg_byte_b(cm2 + cbo1);
if (sh1 > 3u) { cvm21 = cvm21 | (sg_byte_b(cm2 + cbo1 + 1u) << 8u); }
let vm21 = sg_w4b[q2 * gpr + g1];
var cvm30 = sg_byte_b(cm3 + cbo0);
if (sh0 > 3u) { cvm30 = cvm30 | (sg_byte_b(cm3 + cbo0 + 1u) << 8u); }
let vm30 = sg_w4b[q3 * gpr + g0];
var cvm31 = sg_byte_b(cm3 + cbo1);
if (sh1 > 3u) { cvm31 = cvm31 | (sg_byte_b(cm3 + cbo1 + 1u) << 8u); }
let vm31 = sg_w4b[q3 * gpr + g1];
ma0 = ma0 + exp2(pm0.x + f32((cvm00 >> sh0) & 31u) * pm0.y)
* (sg_dot8(vm00.x, xa0, xb0) + sg_dot8(vm00.y, xc0, xd0)
+ sg_dot8(vm00.z, xe0, xf0) + sg_dot8(vm00.w, xg0, xh0));
ma1 = ma1 + exp2(pm1.x + f32((cvm10 >> sh0) & 31u) * pm1.y)
* (sg_dot8(vm10.x, xa0, xb0) + sg_dot8(vm10.y, xc0, xd0)
+ sg_dot8(vm10.z, xe0, xf0) + sg_dot8(vm10.w, xg0, xh0));
ma2 = ma2 + exp2(pm2.x + f32((cvm20 >> sh0) & 31u) * pm2.y)
* (sg_dot8(vm20.x, xa0, xb0) + sg_dot8(vm20.y, xc0, xd0)
+ sg_dot8(vm20.z, xe0, xf0) + sg_dot8(vm20.w, xg0, xh0));
ma3 = ma3 + exp2(pm3.x + f32((cvm30 >> sh0) & 31u) * pm3.y)
* (sg_dot8(vm30.x, xa0, xb0) + sg_dot8(vm30.y, xc0, xd0)
+ sg_dot8(vm30.z, xe0, xf0) + sg_dot8(vm30.w, xg0, xh0));
if (g + 32u < gpr) {
mb0 = mb0 + exp2(pm0.x + f32((cvm01 >> sh1) & 31u) * pm0.y)
* (sg_dot8(vm01.x, xa1, xb1) + sg_dot8(vm01.y, xc1, xd1)
+ sg_dot8(vm01.z, xe1, xf1) + sg_dot8(vm01.w, xg1, xh1));
mb1 = mb1 + exp2(pm1.x + f32((cvm11 >> sh1) & 31u) * pm1.y)
* (sg_dot8(vm11.x, xa1, xb1) + sg_dot8(vm11.y, xc1, xd1)
+ sg_dot8(vm11.z, xe1, xf1) + sg_dot8(vm11.w, xg1, xh1));
mb2 = mb2 + exp2(pm2.x + f32((cvm21 >> sh1) & 31u) * pm2.y)
* (sg_dot8(vm21.x, xa1, xb1) + sg_dot8(vm21.y, xc1, xd1)
+ sg_dot8(vm21.z, xe1, xf1) + sg_dot8(vm21.w, xg1, xh1));
mb3 = mb3 + exp2(pm3.x + f32((cvm31 >> sh1) & 31u) * pm3.y)
* (sg_dot8(vm31.x, xa1, xb1) + sg_dot8(vm31.y, xc1, xd1)
+ sg_dot8(vm31.z, xe1, xf1) + sg_dot8(vm31.w, xg1, xh1));
}
g = g + 64u;
}
let t0 = sg_tree(ma0 + mb0);
let t1 = sg_tree(ma1 + mb1);
let t2 = sg_tree(ma2 + mb2);
let t3 = sg_tree(ma3 + mb3);
if (lane == 0u) {
if (l0) { sg_yb[r0] = t0; }
if (l1) { sg_yb[r1] = t1; }
if (l2) { sg_yb[r2] = t2; }
if (l3) { sg_yb[r3] = t3; }
}
}
fn run_c(blk: u32, warp: u32, lane: u32) {
let rows = sg_p.rows_c;
let gpr = sg_p.gpr;
let params_w = rows * gpr * 4u;
let codes_b = rows * gpr * 16u + rows * 4u;
let cstride = (gpr * 5u + 7u) / 8u;
let base = (blk * 8u + warp) * 4u;
let r0 = base + 0u;
let l0 = r0 < rows;
let q0 = select(base, r0, l0);
let cm0 = codes_b + q0 * cstride;
let pm0 = unpack2x16float(sg_wc[params_w + q0]);
var ma0 = 0.0;
var mb0 = 0.0;
let r1 = base + 1u;
let l1 = r1 < rows;
let q1 = select(base, r1, l1);
let cm1 = codes_b + q1 * cstride;
let pm1 = unpack2x16float(sg_wc[params_w + q1]);
var ma1 = 0.0;
var mb1 = 0.0;
let r2 = base + 2u;
let l2 = r2 < rows;
let q2 = select(base, r2, l2);
let cm2 = codes_b + q2 * cstride;
let pm2 = unpack2x16float(sg_wc[params_w + q2]);
var ma2 = 0.0;
var mb2 = 0.0;
let r3 = base + 3u;
let l3 = r3 < rows;
let q3 = select(base, r3, l3);
let cm3 = codes_b + q3 * cstride;
let pm3 = unpack2x16float(sg_wc[params_w + q3]);
var ma3 = 0.0;
var mb3 = 0.0;
if (!l0) { return; }
var g = lane;
loop {
if (g >= gpr) { break; }
let g0 = min(g + 0u, gpr - 1u);
let bit0 = g0 * 5u;
let cbo0 = bit0 >> 3u;
let sh0 = bit0 & 7u;
let xo0 = g0 * 8u;
let xa0 = sg_x[xo0 + 0u];
let xb0 = sg_x[xo0 + 1u];
let xc0 = sg_x[xo0 + 2u];
let xd0 = sg_x[xo0 + 3u];
let xe0 = sg_x[xo0 + 4u];
let xf0 = sg_x[xo0 + 5u];
let xg0 = sg_x[xo0 + 6u];
let xh0 = sg_x[xo0 + 7u];
let g1 = min(g + 32u, gpr - 1u);
let bit1 = g1 * 5u;
let cbo1 = bit1 >> 3u;
let sh1 = bit1 & 7u;
let xo1 = g1 * 8u;
let xa1 = sg_x[xo1 + 0u];
let xb1 = sg_x[xo1 + 1u];
let xc1 = sg_x[xo1 + 2u];
let xd1 = sg_x[xo1 + 3u];
let xe1 = sg_x[xo1 + 4u];
let xf1 = sg_x[xo1 + 5u];
let xg1 = sg_x[xo1 + 6u];
let xh1 = sg_x[xo1 + 7u];
var cvm00 = sg_byte_c(cm0 + cbo0);
if (sh0 > 3u) { cvm00 = cvm00 | (sg_byte_c(cm0 + cbo0 + 1u) << 8u); }
let vm00 = sg_w4c[q0 * gpr + g0];
var cvm01 = sg_byte_c(cm0 + cbo1);
if (sh1 > 3u) { cvm01 = cvm01 | (sg_byte_c(cm0 + cbo1 + 1u) << 8u); }
let vm01 = sg_w4c[q0 * gpr + g1];
var cvm10 = sg_byte_c(cm1 + cbo0);
if (sh0 > 3u) { cvm10 = cvm10 | (sg_byte_c(cm1 + cbo0 + 1u) << 8u); }
let vm10 = sg_w4c[q1 * gpr + g0];
var cvm11 = sg_byte_c(cm1 + cbo1);
if (sh1 > 3u) { cvm11 = cvm11 | (sg_byte_c(cm1 + cbo1 + 1u) << 8u); }
let vm11 = sg_w4c[q1 * gpr + g1];
var cvm20 = sg_byte_c(cm2 + cbo0);
if (sh0 > 3u) { cvm20 = cvm20 | (sg_byte_c(cm2 + cbo0 + 1u) << 8u); }
let vm20 = sg_w4c[q2 * gpr + g0];
var cvm21 = sg_byte_c(cm2 + cbo1);
if (sh1 > 3u) { cvm21 = cvm21 | (sg_byte_c(cm2 + cbo1 + 1u) << 8u); }
let vm21 = sg_w4c[q2 * gpr + g1];
var cvm30 = sg_byte_c(cm3 + cbo0);
if (sh0 > 3u) { cvm30 = cvm30 | (sg_byte_c(cm3 + cbo0 + 1u) << 8u); }
let vm30 = sg_w4c[q3 * gpr + g0];
var cvm31 = sg_byte_c(cm3 + cbo1);
if (sh1 > 3u) { cvm31 = cvm31 | (sg_byte_c(cm3 + cbo1 + 1u) << 8u); }
let vm31 = sg_w4c[q3 * gpr + g1];
ma0 = ma0 + exp2(pm0.x + f32((cvm00 >> sh0) & 31u) * pm0.y)
* (sg_dot8(vm00.x, xa0, xb0) + sg_dot8(vm00.y, xc0, xd0)
+ sg_dot8(vm00.z, xe0, xf0) + sg_dot8(vm00.w, xg0, xh0));
ma1 = ma1 + exp2(pm1.x + f32((cvm10 >> sh0) & 31u) * pm1.y)
* (sg_dot8(vm10.x, xa0, xb0) + sg_dot8(vm10.y, xc0, xd0)
+ sg_dot8(vm10.z, xe0, xf0) + sg_dot8(vm10.w, xg0, xh0));
ma2 = ma2 + exp2(pm2.x + f32((cvm20 >> sh0) & 31u) * pm2.y)
* (sg_dot8(vm20.x, xa0, xb0) + sg_dot8(vm20.y, xc0, xd0)
+ sg_dot8(vm20.z, xe0, xf0) + sg_dot8(vm20.w, xg0, xh0));
ma3 = ma3 + exp2(pm3.x + f32((cvm30 >> sh0) & 31u) * pm3.y)
* (sg_dot8(vm30.x, xa0, xb0) + sg_dot8(vm30.y, xc0, xd0)
+ sg_dot8(vm30.z, xe0, xf0) + sg_dot8(vm30.w, xg0, xh0));
if (g + 32u < gpr) {
mb0 = mb0 + exp2(pm0.x + f32((cvm01 >> sh1) & 31u) * pm0.y)
* (sg_dot8(vm01.x, xa1, xb1) + sg_dot8(vm01.y, xc1, xd1)
+ sg_dot8(vm01.z, xe1, xf1) + sg_dot8(vm01.w, xg1, xh1));
mb1 = mb1 + exp2(pm1.x + f32((cvm11 >> sh1) & 31u) * pm1.y)
* (sg_dot8(vm11.x, xa1, xb1) + sg_dot8(vm11.y, xc1, xd1)
+ sg_dot8(vm11.z, xe1, xf1) + sg_dot8(vm11.w, xg1, xh1));
mb2 = mb2 + exp2(pm2.x + f32((cvm21 >> sh1) & 31u) * pm2.y)
* (sg_dot8(vm21.x, xa1, xb1) + sg_dot8(vm21.y, xc1, xd1)
+ sg_dot8(vm21.z, xe1, xf1) + sg_dot8(vm21.w, xg1, xh1));
mb3 = mb3 + exp2(pm3.x + f32((cvm31 >> sh1) & 31u) * pm3.y)
* (sg_dot8(vm31.x, xa1, xb1) + sg_dot8(vm31.y, xc1, xd1)
+ sg_dot8(vm31.z, xe1, xf1) + sg_dot8(vm31.w, xg1, xh1));
}
g = g + 64u;
}
let t0 = sg_tree(ma0 + mb0);
let t1 = sg_tree(ma1 + mb1);
let t2 = sg_tree(ma2 + mb2);
let t3 = sg_tree(ma3 + mb3);
if (lane == 0u) {
if (l0) { sg_yc[r0] = t0; }
if (l1) { sg_yc[r1] = t1; }
if (l2) { sg_yc[r2] = t2; }
if (l3) { sg_yc[r3] = t3; }
}
}
@compute @workgroup_size(256)
fn q4tp_mv_sg(@builtin(workgroup_id) wid: vec3<u32>,
@builtin(local_invocation_index) lid: u32,
@builtin(subgroup_invocation_id) lane: u32) {
let per = 8u * 4u;
let warp = lid >> 5u;
let ba = (sg_p.rows_a + per - 1u) / per;
let bb = (sg_p.rows_b + per - 1u) / per;
let wb = wid.x;
if (wb < ba) {
run_a(wb, warp, lane);
} else if (wb < ba + bb) {
run_b(wb - ba, warp, lane);
} else {
run_c(wb - ba - bb, warp, lane);
}
}
"#;
pub(crate) const GU_SG_R2U2: &str = r#"
struct SgP { rows_a: u32, rows_b: u32, rows_c: u32, gpr: u32, lim: u32, act: u32, _p0: u32, _p1: u32 };
@group(0) @binding(0) var<storage, read> sg_wa : array<u32>;
@group(0) @binding(1) var<storage, read> sg_wb : array<u32>;
@group(0) @binding(2) var<storage, read> sg_wc : array<u32>;
@group(0) @binding(3) var<storage, read> sg_w4a : array<vec4<u32>>;
@group(0) @binding(4) var<storage, read> sg_w4b : array<vec4<u32>>;
@group(0) @binding(5) var<storage, read> sg_w4c : array<vec4<u32>>;
@group(0) @binding(6) var<storage, read> sg_x : array<vec4<f32>>;
@group(0) @binding(7) var<storage, read_write> sg_ya : array<f32>;
@group(0) @binding(8) var<storage, read_write> sg_yb : array<f32>;
@group(0) @binding(9) var<storage, read_write> sg_yc : array<f32>;
@group(0) @binding(10) var<uniform> sg_p : SgP;
fn sg_byte_a(off: u32) -> u32 { return (sg_wa[off >> 2u] >> ((off & 3u) * 8u)) & 0xFFu; }
fn sg_byte_b(off: u32) -> u32 { return (sg_wb[off >> 2u] >> ((off & 3u) * 8u)) & 0xFFu; }
fn sg_byte_c(off: u32) -> u32 { return (sg_wc[off >> 2u] >> ((off & 3u) * 8u)) & 0xFFu; }
fn sg_nib(w: u32, sh: u32) -> f32 {
return bitcast<f32>(((w >> sh) & 0xFu) | 0x4B000000u) - 8388616.0;
}
fn sg_dot8(w: u32, a: vec4<f32>, b: vec4<f32>) -> f32 {
return sg_nib(w, 0u) * a.x
+ sg_nib(w, 4u) * a.y
+ sg_nib(w, 8u) * a.z
+ sg_nib(w, 12u) * a.w
+ sg_nib(w, 16u) * b.x
+ sg_nib(w, 20u) * b.y
+ sg_nib(w, 24u) * b.z
+ sg_nib(w, 28u) * b.w;
}
// The 64-lane tree of `q4tp_matvec16nl` from stride 16 down: lane l takes
// lane l+s's partial, in the same order (own + partner).
fn sg_tree(v0: f32) -> f32 {
var v = v0;
v = v + subgroupShuffleDown(v, 16u);
v = v + subgroupShuffleDown(v, 8u);
v = v + subgroupShuffleDown(v, 4u);
v = v + subgroupShuffleDown(v, 2u);
v = v + subgroupShuffleDown(v, 1u);
return v;
}
fn sg_erf(x: f32) -> f32 {
let a = abs(x);
let t = 1.0 / (1.0 + 0.3275911 * a);
let y = 1.0
- (((((1.0614054 * t - 1.4531521) * t + 1.4214138) * t - 0.28449674) * t
+ 0.2548296)
* t)
* exp(-a * a);
return select(y, -y, x < 0.0);
}
fn sg_gelu_erf(x: f32) -> f32 {
return 0.5 * x * (1.0 + sg_erf(x * 0.70710678));
}
fn sg_gelu_tanh(x: f32) -> f32 {
return 0.5 * x * (1.0 + tanh(0.7978846 * (x + 0.044715 * x * x * x)));
}
fn run_gu(blk: u32, warp: u32, lane: u32) {
let rows = sg_p.rows_a;
let gpr = sg_p.gpr;
let params_w = rows * gpr * 4u;
let codes_b = rows * gpr * 16u + rows * 4u;
let cstride = (gpr * 5u + 7u) / 8u;
let base = (blk * 8u + warp) * 2u;
let r0 = base + 0u;
let l0 = r0 < rows;
let q0 = select(base, r0, l0);
let cg0 = codes_b + q0 * cstride;
let pg0 = unpack2x16float(sg_wa[params_w + q0]);
var ga0 = 0.0;
var gb0 = 0.0;
let cu0 = codes_b + q0 * cstride;
let pu0 = unpack2x16float(sg_wb[params_w + q0]);
var ua0 = 0.0;
var ub0 = 0.0;
let r1 = base + 1u;
let l1 = r1 < rows;
let q1 = select(base, r1, l1);
let cg1 = codes_b + q1 * cstride;
let pg1 = unpack2x16float(sg_wa[params_w + q1]);
var ga1 = 0.0;
var gb1 = 0.0;
let cu1 = codes_b + q1 * cstride;
let pu1 = unpack2x16float(sg_wb[params_w + q1]);
var ua1 = 0.0;
var ub1 = 0.0;
if (!l0) { return; }
var g = lane;
loop {
if (g >= gpr) { break; }
let g0 = min(g + 0u, gpr - 1u);
let bit0 = g0 * 5u;
let cbo0 = bit0 >> 3u;
let sh0 = bit0 & 7u;
let xo0 = g0 * 8u;
let xa0 = sg_x[xo0 + 0u];
let xb0 = sg_x[xo0 + 1u];
let xc0 = sg_x[xo0 + 2u];
let xd0 = sg_x[xo0 + 3u];
let xe0 = sg_x[xo0 + 4u];
let xf0 = sg_x[xo0 + 5u];
let xg0 = sg_x[xo0 + 6u];
let xh0 = sg_x[xo0 + 7u];
let g1 = min(g + 32u, gpr - 1u);
let bit1 = g1 * 5u;
let cbo1 = bit1 >> 3u;
let sh1 = bit1 & 7u;
let xo1 = g1 * 8u;
let xa1 = sg_x[xo1 + 0u];
let xb1 = sg_x[xo1 + 1u];
let xc1 = sg_x[xo1 + 2u];
let xd1 = sg_x[xo1 + 3u];
let xe1 = sg_x[xo1 + 4u];
let xf1 = sg_x[xo1 + 5u];
let xg1 = sg_x[xo1 + 6u];
let xh1 = sg_x[xo1 + 7u];
var cvg00 = sg_byte_a(cg0 + cbo0);
if (sh0 > 3u) { cvg00 = cvg00 | (sg_byte_a(cg0 + cbo0 + 1u) << 8u); }
let vg00 = sg_w4a[q0 * gpr + g0];
var cvg01 = sg_byte_a(cg0 + cbo1);
if (sh1 > 3u) { cvg01 = cvg01 | (sg_byte_a(cg0 + cbo1 + 1u) << 8u); }
let vg01 = sg_w4a[q0 * gpr + g1];
var cvg10 = sg_byte_a(cg1 + cbo0);
if (sh0 > 3u) { cvg10 = cvg10 | (sg_byte_a(cg1 + cbo0 + 1u) << 8u); }
let vg10 = sg_w4a[q1 * gpr + g0];
var cvg11 = sg_byte_a(cg1 + cbo1);
if (sh1 > 3u) { cvg11 = cvg11 | (sg_byte_a(cg1 + cbo1 + 1u) << 8u); }
let vg11 = sg_w4a[q1 * gpr + g1];
var cvu00 = sg_byte_b(cu0 + cbo0);
if (sh0 > 3u) { cvu00 = cvu00 | (sg_byte_b(cu0 + cbo0 + 1u) << 8u); }
let vu00 = sg_w4b[q0 * gpr + g0];
var cvu01 = sg_byte_b(cu0 + cbo1);
if (sh1 > 3u) { cvu01 = cvu01 | (sg_byte_b(cu0 + cbo1 + 1u) << 8u); }
let vu01 = sg_w4b[q0 * gpr + g1];
var cvu10 = sg_byte_b(cu1 + cbo0);
if (sh0 > 3u) { cvu10 = cvu10 | (sg_byte_b(cu1 + cbo0 + 1u) << 8u); }
let vu10 = sg_w4b[q1 * gpr + g0];
var cvu11 = sg_byte_b(cu1 + cbo1);
if (sh1 > 3u) { cvu11 = cvu11 | (sg_byte_b(cu1 + cbo1 + 1u) << 8u); }
let vu11 = sg_w4b[q1 * gpr + g1];
ga0 = ga0 + exp2(pg0.x + f32((cvg00 >> sh0) & 31u) * pg0.y)
* (sg_dot8(vg00.x, xa0, xb0) + sg_dot8(vg00.y, xc0, xd0)
+ sg_dot8(vg00.z, xe0, xf0) + sg_dot8(vg00.w, xg0, xh0));
ga1 = ga1 + exp2(pg1.x + f32((cvg10 >> sh0) & 31u) * pg1.y)
* (sg_dot8(vg10.x, xa0, xb0) + sg_dot8(vg10.y, xc0, xd0)
+ sg_dot8(vg10.z, xe0, xf0) + sg_dot8(vg10.w, xg0, xh0));
ua0 = ua0 + exp2(pu0.x + f32((cvu00 >> sh0) & 31u) * pu0.y)
* (sg_dot8(vu00.x, xa0, xb0) + sg_dot8(vu00.y, xc0, xd0)
+ sg_dot8(vu00.z, xe0, xf0) + sg_dot8(vu00.w, xg0, xh0));
ua1 = ua1 + exp2(pu1.x + f32((cvu10 >> sh0) & 31u) * pu1.y)
* (sg_dot8(vu10.x, xa0, xb0) + sg_dot8(vu10.y, xc0, xd0)
+ sg_dot8(vu10.z, xe0, xf0) + sg_dot8(vu10.w, xg0, xh0));
if (g + 32u < gpr) {
gb0 = gb0 + exp2(pg0.x + f32((cvg01 >> sh1) & 31u) * pg0.y)
* (sg_dot8(vg01.x, xa1, xb1) + sg_dot8(vg01.y, xc1, xd1)
+ sg_dot8(vg01.z, xe1, xf1) + sg_dot8(vg01.w, xg1, xh1));
gb1 = gb1 + exp2(pg1.x + f32((cvg11 >> sh1) & 31u) * pg1.y)
* (sg_dot8(vg11.x, xa1, xb1) + sg_dot8(vg11.y, xc1, xd1)
+ sg_dot8(vg11.z, xe1, xf1) + sg_dot8(vg11.w, xg1, xh1));
ub0 = ub0 + exp2(pu0.x + f32((cvu01 >> sh1) & 31u) * pu0.y)
* (sg_dot8(vu01.x, xa1, xb1) + sg_dot8(vu01.y, xc1, xd1)
+ sg_dot8(vu01.z, xe1, xf1) + sg_dot8(vu01.w, xg1, xh1));
ub1 = ub1 + exp2(pu1.x + f32((cvu11 >> sh1) & 31u) * pu1.y)
* (sg_dot8(vu11.x, xa1, xb1) + sg_dot8(vu11.y, xc1, xd1)
+ sg_dot8(vu11.z, xe1, xf1) + sg_dot8(vu11.w, xg1, xh1));
}
g = g + 64u;
}
var gg0 = sg_tree(ga0 + gb0);
var uu0 = sg_tree(ua0 + ub0);
var gg1 = sg_tree(ga1 + gb1);
var uu1 = sg_tree(ua1 + ub1);
if (lane == 0u) {
let lim = bitcast<f32>(sg_p.lim);
if (l0) {
if (lim > 0.0) { uu0 = clamp(uu0, -lim, lim); gg0 = min(gg0, lim); }
if (sg_p.act == 1u) { sg_ya[r0] = sg_gelu_erf(gg0) * uu0; }
else if (sg_p.act == 2u) { sg_ya[r0] = sg_gelu_tanh(gg0) * uu0; }
else { sg_ya[r0] = (gg0 / (1.0 + exp(-gg0))) * uu0; }
}
if (l1) {
if (lim > 0.0) { uu1 = clamp(uu1, -lim, lim); gg1 = min(gg1, lim); }
if (sg_p.act == 1u) { sg_ya[r1] = sg_gelu_erf(gg1) * uu1; }
else if (sg_p.act == 2u) { sg_ya[r1] = sg_gelu_tanh(gg1) * uu1; }
else { sg_ya[r1] = (gg1 / (1.0 + exp(-gg1))) * uu1; }
}
}
}
@compute @workgroup_size(256)
fn q4tp_gu_sg(@builtin(workgroup_id) wid: vec3<u32>,
@builtin(local_invocation_index) lid: u32,
@builtin(subgroup_invocation_id) lane: u32) {
run_gu(wid.x, lid >> 5u, lane);
}
"#;
pub(crate) const GU_SG_R4U2: &str = r#"
struct SgP { rows_a: u32, rows_b: u32, rows_c: u32, gpr: u32, lim: u32, act: u32, _p0: u32, _p1: u32 };
@group(0) @binding(0) var<storage, read> sg_wa : array<u32>;
@group(0) @binding(1) var<storage, read> sg_wb : array<u32>;
@group(0) @binding(2) var<storage, read> sg_wc : array<u32>;
@group(0) @binding(3) var<storage, read> sg_w4a : array<vec4<u32>>;
@group(0) @binding(4) var<storage, read> sg_w4b : array<vec4<u32>>;
@group(0) @binding(5) var<storage, read> sg_w4c : array<vec4<u32>>;
@group(0) @binding(6) var<storage, read> sg_x : array<vec4<f32>>;
@group(0) @binding(7) var<storage, read_write> sg_ya : array<f32>;
@group(0) @binding(8) var<storage, read_write> sg_yb : array<f32>;
@group(0) @binding(9) var<storage, read_write> sg_yc : array<f32>;
@group(0) @binding(10) var<uniform> sg_p : SgP;
fn sg_byte_a(off: u32) -> u32 { return (sg_wa[off >> 2u] >> ((off & 3u) * 8u)) & 0xFFu; }
fn sg_byte_b(off: u32) -> u32 { return (sg_wb[off >> 2u] >> ((off & 3u) * 8u)) & 0xFFu; }
fn sg_byte_c(off: u32) -> u32 { return (sg_wc[off >> 2u] >> ((off & 3u) * 8u)) & 0xFFu; }
fn sg_nib(w: u32, sh: u32) -> f32 {
return bitcast<f32>(((w >> sh) & 0xFu) | 0x4B000000u) - 8388616.0;
}
fn sg_dot8(w: u32, a: vec4<f32>, b: vec4<f32>) -> f32 {
return sg_nib(w, 0u) * a.x
+ sg_nib(w, 4u) * a.y
+ sg_nib(w, 8u) * a.z
+ sg_nib(w, 12u) * a.w
+ sg_nib(w, 16u) * b.x
+ sg_nib(w, 20u) * b.y
+ sg_nib(w, 24u) * b.z
+ sg_nib(w, 28u) * b.w;
}
// The 64-lane tree of `q4tp_matvec16nl` from stride 16 down: lane l takes
// lane l+s's partial, in the same order (own + partner).
fn sg_tree(v0: f32) -> f32 {
var v = v0;
v = v + subgroupShuffleDown(v, 16u);
v = v + subgroupShuffleDown(v, 8u);
v = v + subgroupShuffleDown(v, 4u);
v = v + subgroupShuffleDown(v, 2u);
v = v + subgroupShuffleDown(v, 1u);
return v;
}
fn sg_erf(x: f32) -> f32 {
let a = abs(x);
let t = 1.0 / (1.0 + 0.3275911 * a);
let y = 1.0
- (((((1.0614054 * t - 1.4531521) * t + 1.4214138) * t - 0.28449674) * t
+ 0.2548296)
* t)
* exp(-a * a);
return select(y, -y, x < 0.0);
}
fn sg_gelu_erf(x: f32) -> f32 {
return 0.5 * x * (1.0 + sg_erf(x * 0.70710678));
}
fn sg_gelu_tanh(x: f32) -> f32 {
return 0.5 * x * (1.0 + tanh(0.7978846 * (x + 0.044715 * x * x * x)));
}
fn run_gu(blk: u32, warp: u32, lane: u32) {
let rows = sg_p.rows_a;
let gpr = sg_p.gpr;
let params_w = rows * gpr * 4u;
let codes_b = rows * gpr * 16u + rows * 4u;
let cstride = (gpr * 5u + 7u) / 8u;
let base = (blk * 8u + warp) * 4u;
let r0 = base + 0u;
let l0 = r0 < rows;
let q0 = select(base, r0, l0);
let cg0 = codes_b + q0 * cstride;
let pg0 = unpack2x16float(sg_wa[params_w + q0]);
var ga0 = 0.0;
var gb0 = 0.0;
let cu0 = codes_b + q0 * cstride;
let pu0 = unpack2x16float(sg_wb[params_w + q0]);
var ua0 = 0.0;
var ub0 = 0.0;
let r1 = base + 1u;
let l1 = r1 < rows;
let q1 = select(base, r1, l1);
let cg1 = codes_b + q1 * cstride;
let pg1 = unpack2x16float(sg_wa[params_w + q1]);
var ga1 = 0.0;
var gb1 = 0.0;
let cu1 = codes_b + q1 * cstride;
let pu1 = unpack2x16float(sg_wb[params_w + q1]);
var ua1 = 0.0;
var ub1 = 0.0;
let r2 = base + 2u;
let l2 = r2 < rows;
let q2 = select(base, r2, l2);
let cg2 = codes_b + q2 * cstride;
let pg2 = unpack2x16float(sg_wa[params_w + q2]);
var ga2 = 0.0;
var gb2 = 0.0;
let cu2 = codes_b + q2 * cstride;
let pu2 = unpack2x16float(sg_wb[params_w + q2]);
var ua2 = 0.0;
var ub2 = 0.0;
let r3 = base + 3u;
let l3 = r3 < rows;
let q3 = select(base, r3, l3);
let cg3 = codes_b + q3 * cstride;
let pg3 = unpack2x16float(sg_wa[params_w + q3]);
var ga3 = 0.0;
var gb3 = 0.0;
let cu3 = codes_b + q3 * cstride;
let pu3 = unpack2x16float(sg_wb[params_w + q3]);
var ua3 = 0.0;
var ub3 = 0.0;
if (!l0) { return; }
var g = lane;
loop {
if (g >= gpr) { break; }
let g0 = min(g + 0u, gpr - 1u);
let bit0 = g0 * 5u;
let cbo0 = bit0 >> 3u;
let sh0 = bit0 & 7u;
let xo0 = g0 * 8u;
let xa0 = sg_x[xo0 + 0u];
let xb0 = sg_x[xo0 + 1u];
let xc0 = sg_x[xo0 + 2u];
let xd0 = sg_x[xo0 + 3u];
let xe0 = sg_x[xo0 + 4u];
let xf0 = sg_x[xo0 + 5u];
let xg0 = sg_x[xo0 + 6u];
let xh0 = sg_x[xo0 + 7u];
let g1 = min(g + 32u, gpr - 1u);
let bit1 = g1 * 5u;
let cbo1 = bit1 >> 3u;
let sh1 = bit1 & 7u;
let xo1 = g1 * 8u;
let xa1 = sg_x[xo1 + 0u];
let xb1 = sg_x[xo1 + 1u];
let xc1 = sg_x[xo1 + 2u];
let xd1 = sg_x[xo1 + 3u];
let xe1 = sg_x[xo1 + 4u];
let xf1 = sg_x[xo1 + 5u];
let xg1 = sg_x[xo1 + 6u];
let xh1 = sg_x[xo1 + 7u];
var cvg00 = sg_byte_a(cg0 + cbo0);
if (sh0 > 3u) { cvg00 = cvg00 | (sg_byte_a(cg0 + cbo0 + 1u) << 8u); }
let vg00 = sg_w4a[q0 * gpr + g0];
var cvg01 = sg_byte_a(cg0 + cbo1);
if (sh1 > 3u) { cvg01 = cvg01 | (sg_byte_a(cg0 + cbo1 + 1u) << 8u); }
let vg01 = sg_w4a[q0 * gpr + g1];
var cvg10 = sg_byte_a(cg1 + cbo0);
if (sh0 > 3u) { cvg10 = cvg10 | (sg_byte_a(cg1 + cbo0 + 1u) << 8u); }
let vg10 = sg_w4a[q1 * gpr + g0];
var cvg11 = sg_byte_a(cg1 + cbo1);
if (sh1 > 3u) { cvg11 = cvg11 | (sg_byte_a(cg1 + cbo1 + 1u) << 8u); }
let vg11 = sg_w4a[q1 * gpr + g1];
var cvg20 = sg_byte_a(cg2 + cbo0);
if (sh0 > 3u) { cvg20 = cvg20 | (sg_byte_a(cg2 + cbo0 + 1u) << 8u); }
let vg20 = sg_w4a[q2 * gpr + g0];
var cvg21 = sg_byte_a(cg2 + cbo1);
if (sh1 > 3u) { cvg21 = cvg21 | (sg_byte_a(cg2 + cbo1 + 1u) << 8u); }
let vg21 = sg_w4a[q2 * gpr + g1];
var cvg30 = sg_byte_a(cg3 + cbo0);
if (sh0 > 3u) { cvg30 = cvg30 | (sg_byte_a(cg3 + cbo0 + 1u) << 8u); }
let vg30 = sg_w4a[q3 * gpr + g0];
var cvg31 = sg_byte_a(cg3 + cbo1);
if (sh1 > 3u) { cvg31 = cvg31 | (sg_byte_a(cg3 + cbo1 + 1u) << 8u); }
let vg31 = sg_w4a[q3 * gpr + g1];
var cvu00 = sg_byte_b(cu0 + cbo0);
if (sh0 > 3u) { cvu00 = cvu00 | (sg_byte_b(cu0 + cbo0 + 1u) << 8u); }
let vu00 = sg_w4b[q0 * gpr + g0];
var cvu01 = sg_byte_b(cu0 + cbo1);
if (sh1 > 3u) { cvu01 = cvu01 | (sg_byte_b(cu0 + cbo1 + 1u) << 8u); }
let vu01 = sg_w4b[q0 * gpr + g1];
var cvu10 = sg_byte_b(cu1 + cbo0);
if (sh0 > 3u) { cvu10 = cvu10 | (sg_byte_b(cu1 + cbo0 + 1u) << 8u); }
let vu10 = sg_w4b[q1 * gpr + g0];
var cvu11 = sg_byte_b(cu1 + cbo1);
if (sh1 > 3u) { cvu11 = cvu11 | (sg_byte_b(cu1 + cbo1 + 1u) << 8u); }
let vu11 = sg_w4b[q1 * gpr + g1];
var cvu20 = sg_byte_b(cu2 + cbo0);
if (sh0 > 3u) { cvu20 = cvu20 | (sg_byte_b(cu2 + cbo0 + 1u) << 8u); }
let vu20 = sg_w4b[q2 * gpr + g0];
var cvu21 = sg_byte_b(cu2 + cbo1);
if (sh1 > 3u) { cvu21 = cvu21 | (sg_byte_b(cu2 + cbo1 + 1u) << 8u); }
let vu21 = sg_w4b[q2 * gpr + g1];
var cvu30 = sg_byte_b(cu3 + cbo0);
if (sh0 > 3u) { cvu30 = cvu30 | (sg_byte_b(cu3 + cbo0 + 1u) << 8u); }
let vu30 = sg_w4b[q3 * gpr + g0];
var cvu31 = sg_byte_b(cu3 + cbo1);
if (sh1 > 3u) { cvu31 = cvu31 | (sg_byte_b(cu3 + cbo1 + 1u) << 8u); }
let vu31 = sg_w4b[q3 * gpr + g1];
ga0 = ga0 + exp2(pg0.x + f32((cvg00 >> sh0) & 31u) * pg0.y)
* (sg_dot8(vg00.x, xa0, xb0) + sg_dot8(vg00.y, xc0, xd0)
+ sg_dot8(vg00.z, xe0, xf0) + sg_dot8(vg00.w, xg0, xh0));
ga1 = ga1 + exp2(pg1.x + f32((cvg10 >> sh0) & 31u) * pg1.y)
* (sg_dot8(vg10.x, xa0, xb0) + sg_dot8(vg10.y, xc0, xd0)
+ sg_dot8(vg10.z, xe0, xf0) + sg_dot8(vg10.w, xg0, xh0));
ga2 = ga2 + exp2(pg2.x + f32((cvg20 >> sh0) & 31u) * pg2.y)
* (sg_dot8(vg20.x, xa0, xb0) + sg_dot8(vg20.y, xc0, xd0)
+ sg_dot8(vg20.z, xe0, xf0) + sg_dot8(vg20.w, xg0, xh0));
ga3 = ga3 + exp2(pg3.x + f32((cvg30 >> sh0) & 31u) * pg3.y)
* (sg_dot8(vg30.x, xa0, xb0) + sg_dot8(vg30.y, xc0, xd0)
+ sg_dot8(vg30.z, xe0, xf0) + sg_dot8(vg30.w, xg0, xh0));
ua0 = ua0 + exp2(pu0.x + f32((cvu00 >> sh0) & 31u) * pu0.y)
* (sg_dot8(vu00.x, xa0, xb0) + sg_dot8(vu00.y, xc0, xd0)
+ sg_dot8(vu00.z, xe0, xf0) + sg_dot8(vu00.w, xg0, xh0));
ua1 = ua1 + exp2(pu1.x + f32((cvu10 >> sh0) & 31u) * pu1.y)
* (sg_dot8(vu10.x, xa0, xb0) + sg_dot8(vu10.y, xc0, xd0)
+ sg_dot8(vu10.z, xe0, xf0) + sg_dot8(vu10.w, xg0, xh0));
ua2 = ua2 + exp2(pu2.x + f32((cvu20 >> sh0) & 31u) * pu2.y)
* (sg_dot8(vu20.x, xa0, xb0) + sg_dot8(vu20.y, xc0, xd0)
+ sg_dot8(vu20.z, xe0, xf0) + sg_dot8(vu20.w, xg0, xh0));
ua3 = ua3 + exp2(pu3.x + f32((cvu30 >> sh0) & 31u) * pu3.y)
* (sg_dot8(vu30.x, xa0, xb0) + sg_dot8(vu30.y, xc0, xd0)
+ sg_dot8(vu30.z, xe0, xf0) + sg_dot8(vu30.w, xg0, xh0));
if (g + 32u < gpr) {
gb0 = gb0 + exp2(pg0.x + f32((cvg01 >> sh1) & 31u) * pg0.y)
* (sg_dot8(vg01.x, xa1, xb1) + sg_dot8(vg01.y, xc1, xd1)
+ sg_dot8(vg01.z, xe1, xf1) + sg_dot8(vg01.w, xg1, xh1));
gb1 = gb1 + exp2(pg1.x + f32((cvg11 >> sh1) & 31u) * pg1.y)
* (sg_dot8(vg11.x, xa1, xb1) + sg_dot8(vg11.y, xc1, xd1)
+ sg_dot8(vg11.z, xe1, xf1) + sg_dot8(vg11.w, xg1, xh1));
gb2 = gb2 + exp2(pg2.x + f32((cvg21 >> sh1) & 31u) * pg2.y)
* (sg_dot8(vg21.x, xa1, xb1) + sg_dot8(vg21.y, xc1, xd1)
+ sg_dot8(vg21.z, xe1, xf1) + sg_dot8(vg21.w, xg1, xh1));
gb3 = gb3 + exp2(pg3.x + f32((cvg31 >> sh1) & 31u) * pg3.y)
* (sg_dot8(vg31.x, xa1, xb1) + sg_dot8(vg31.y, xc1, xd1)
+ sg_dot8(vg31.z, xe1, xf1) + sg_dot8(vg31.w, xg1, xh1));
ub0 = ub0 + exp2(pu0.x + f32((cvu01 >> sh1) & 31u) * pu0.y)
* (sg_dot8(vu01.x, xa1, xb1) + sg_dot8(vu01.y, xc1, xd1)
+ sg_dot8(vu01.z, xe1, xf1) + sg_dot8(vu01.w, xg1, xh1));
ub1 = ub1 + exp2(pu1.x + f32((cvu11 >> sh1) & 31u) * pu1.y)
* (sg_dot8(vu11.x, xa1, xb1) + sg_dot8(vu11.y, xc1, xd1)
+ sg_dot8(vu11.z, xe1, xf1) + sg_dot8(vu11.w, xg1, xh1));
ub2 = ub2 + exp2(pu2.x + f32((cvu21 >> sh1) & 31u) * pu2.y)
* (sg_dot8(vu21.x, xa1, xb1) + sg_dot8(vu21.y, xc1, xd1)
+ sg_dot8(vu21.z, xe1, xf1) + sg_dot8(vu21.w, xg1, xh1));
ub3 = ub3 + exp2(pu3.x + f32((cvu31 >> sh1) & 31u) * pu3.y)
* (sg_dot8(vu31.x, xa1, xb1) + sg_dot8(vu31.y, xc1, xd1)
+ sg_dot8(vu31.z, xe1, xf1) + sg_dot8(vu31.w, xg1, xh1));
}
g = g + 64u;
}
var gg0 = sg_tree(ga0 + gb0);
var uu0 = sg_tree(ua0 + ub0);
var gg1 = sg_tree(ga1 + gb1);
var uu1 = sg_tree(ua1 + ub1);
var gg2 = sg_tree(ga2 + gb2);
var uu2 = sg_tree(ua2 + ub2);
var gg3 = sg_tree(ga3 + gb3);
var uu3 = sg_tree(ua3 + ub3);
if (lane == 0u) {
let lim = bitcast<f32>(sg_p.lim);
if (l0) {
if (lim > 0.0) { uu0 = clamp(uu0, -lim, lim); gg0 = min(gg0, lim); }
if (sg_p.act == 1u) { sg_ya[r0] = sg_gelu_erf(gg0) * uu0; }
else if (sg_p.act == 2u) { sg_ya[r0] = sg_gelu_tanh(gg0) * uu0; }
else { sg_ya[r0] = (gg0 / (1.0 + exp(-gg0))) * uu0; }
}
if (l1) {
if (lim > 0.0) { uu1 = clamp(uu1, -lim, lim); gg1 = min(gg1, lim); }
if (sg_p.act == 1u) { sg_ya[r1] = sg_gelu_erf(gg1) * uu1; }
else if (sg_p.act == 2u) { sg_ya[r1] = sg_gelu_tanh(gg1) * uu1; }
else { sg_ya[r1] = (gg1 / (1.0 + exp(-gg1))) * uu1; }
}
if (l2) {
if (lim > 0.0) { uu2 = clamp(uu2, -lim, lim); gg2 = min(gg2, lim); }
if (sg_p.act == 1u) { sg_ya[r2] = sg_gelu_erf(gg2) * uu2; }
else if (sg_p.act == 2u) { sg_ya[r2] = sg_gelu_tanh(gg2) * uu2; }
else { sg_ya[r2] = (gg2 / (1.0 + exp(-gg2))) * uu2; }
}
if (l3) {
if (lim > 0.0) { uu3 = clamp(uu3, -lim, lim); gg3 = min(gg3, lim); }
if (sg_p.act == 1u) { sg_ya[r3] = sg_gelu_erf(gg3) * uu3; }
else if (sg_p.act == 2u) { sg_ya[r3] = sg_gelu_tanh(gg3) * uu3; }
else { sg_ya[r3] = (gg3 / (1.0 + exp(-gg3))) * uu3; }
}
}
}
@compute @workgroup_size(256)
fn q4tp_gu_sg(@builtin(workgroup_id) wid: vec3<u32>,
@builtin(local_invocation_index) lid: u32,
@builtin(subgroup_invocation_id) lane: u32) {
run_gu(wid.x, lid >> 5u, lane);
}
"#;
/// `gdn_step_par` (decode) and `gdn_step_par_k` (batch) with their tree
/// reductions finished by subgroup shuffles: the stride-64 step through
/// workgroup memory, the stride-32 step and the 16..1 strides inside the
/// first subgroup, in the tree's own pair order — the same sums to the
/// bit with four barriers a reduction where the loop takes nine.
pub(crate) const GDN_SG_WGSL: &str = r#"
struct GdnP { nv: u32, dk: u32, dv: u32, kd: u32, rep: u32, cdim: u32, eps: f32, tok: u32 };
@group(0) @binding(0) var<storage, read> gd_cq : array<f32>;
@group(0) @binding(2) var<storage, read> gd_a : array<f32>;
@group(0) @binding(3) var<storage, read> gd_b : array<f32>;
@group(0) @binding(4) var<storage, read> gd_alog : array<f32>;
@group(0) @binding(5) var<storage, read> gd_dtb : array<f32>;
@group(0) @binding(7) var<storage, read_write> gd_S4 : array<vec4<f32>>;
@group(0) @binding(8) var<storage, read_write> gd_o4 : array<vec4<f32>>;
@group(0) @binding(9) var<uniform> gd_p : GdnP;
var<workgroup> sr1: array<f32, 128>;
var<workgroup> sr4: array<vec4<f32>, 128>;
var<workgroup> sbc: vec4<f32>;
fn gd_softplus(x: f32) -> f32 {
if (x > 20.0) { return x; }
return log(1.0 + exp(x));
}
// The 128-lane tree `for stride 64..1: r[t] += r[t + stride]` and its
// trailing barrier, value of r[0] to every lane.
fn sgr1(t: u32, v: f32) -> f32 {
sr1[t] = v;
workgroupBarrier();
if (t < 64u) { sr1[t] = sr1[t] + sr1[t + 64u]; }
workgroupBarrier();
if (t < 32u) {
var s = sr1[t] + sr1[t + 32u];
s = s + subgroupShuffleDown(s, 16u);
s = s + subgroupShuffleDown(s, 8u);
s = s + subgroupShuffleDown(s, 4u);
s = s + subgroupShuffleDown(s, 2u);
s = s + subgroupShuffleDown(s, 1u);
if (t == 0u) { sbc.x = s; }
}
workgroupBarrier();
let r = sbc.x;
workgroupBarrier();
return r;
}
fn sgr4(t: u32, v: vec4<f32>) -> vec4<f32> {
sr4[t] = v;
workgroupBarrier();
if (t < 64u) { sr4[t] = sr4[t] + sr4[t + 64u]; }
workgroupBarrier();
if (t < 32u) {
var s = sr4[t] + sr4[t + 32u];
s = s + subgroupShuffleDown(s, 16u);
s = s + subgroupShuffleDown(s, 8u);
s = s + subgroupShuffleDown(s, 4u);
s = s + subgroupShuffleDown(s, 2u);
s = s + subgroupShuffleDown(s, 1u);
if (t == 0u) { sbc = s; }
}
workgroupBarrier();
let r = sbc;
workgroupBarrier();
return r;
}
@compute @workgroup_size(128)
fn gdn_step_par_sg(@builtin(workgroup_id) wid: vec3<u32>,
@builtin(local_invocation_id) lid: vec3<u32>) {
let h = wid.x;
let dj4 = wid.y;
let t = lid.x;
let dk = gd_p.dk;
let dv = gd_p.dv;
if (h >= gd_p.nv || dj4 * 4u >= dv) { return; }
let ko = h / gd_p.rep;
let qs = ko * dk;
let ks = gd_p.kd + ko * dk;
let nq = sgr1(t, select(0.0, gd_cq[qs + t] * gd_cq[qs + t], t < dk));
let nkn = sgr1(t, select(0.0, gd_cq[ks + t] * gd_cq[ks + t], t < dk));
let invq = 1.0 / (sqrt(nq + 1e-6) * sqrt(f32(dk)));
let invk = 1.0 / sqrt(nkn + 1e-6);
let abo = gd_p.tok * gd_p.nv;
let g = exp(-exp(gd_alog[h]) * gd_softplus(gd_a[abo + h] + gd_dtb[h]));
let beta = 1.0 / (1.0 + exp(-gd_b[abo + h]));
let s4base = (h * dk * dv) >> 2u;
let dv4 = dv >> 2u;
let vto = 2u * gd_p.kd + h * dv + dj4 * 4u;
let vt = vec4<f32>(gd_cq[vto], gd_cq[vto + 1u], gd_cq[vto + 2u], gd_cq[vto + 3u]);
let kf_t = select(0.0, gd_cq[ks + t] * invk, t < dk);
let qf_t = select(0.0, gd_cq[qs + t] * invq, t < dk);
var kv4 = vec4<f32>(0.0);
if (t < dk) {
kv4 = gd_S4[s4base + t * dv4 + dj4] * kf_t;
}
let kv = sgr4(t, kv4);
let delta = (vt - g * kv) * beta;
var contrib = vec4<f32>(0.0);
if (t < dk) {
let idx = s4base + t * dv4 + dj4;
let cell = g * gd_S4[idx] + kf_t * delta;
gd_S4[idx] = cell;
contrib = qf_t * cell;
}
let o = sgr4(t, contrib);
if (t == 0u) {
let zo4 = (gd_p.tok * gd_p.nv * dv) >> 2u;
gd_o4[zo4 + h * dv4 + dj4] = o;
}
}
"#;
pub(crate) const GDNK_SG_WGSL: &str = r#"
struct GdKP {
nv: u32, dk: u32, dv: u32, kd: u32,
rep: u32, cdim: u32, eps: f32, kb: u32,
stride: u32, ring_els: u32, p0: u32, p1: u32,
};
@group(0) @binding(0) var<storage, read> gdk_cq : array<f32>;
@group(0) @binding(2) var<storage, read> gdk_a : array<f32>;
@group(0) @binding(3) var<storage, read> gdk_b : array<f32>;
@group(0) @binding(4) var<storage, read> gdk_alog : array<f32>;
@group(0) @binding(5) var<storage, read> gdk_dtb : array<f32>;
@group(0) @binding(7) var<storage, read_write> gdk_S4 : array<vec4<f32>>;
@group(0) @binding(8) var<storage, read_write> gdk_o4 : array<vec4<f32>>;
@group(0) @binding(9) var<uniform> gdk_p : GdKP;
@group(0) @binding(10) var<storage, read_write> gdk_snap4 : array<vec4<f32>>;
var<workgroup> sr1: array<f32, 128>;
var<workgroup> sr4: array<vec4<f32>, 128>;
var<workgroup> sbc: vec4<f32>;
fn gd_softplus(x: f32) -> f32 {
if (x > 20.0) { return x; }
return log(1.0 + exp(x));
}
fn sgr1(t: u32, v: f32) -> f32 {
sr1[t] = v;
workgroupBarrier();
if (t < 64u) { sr1[t] = sr1[t] + sr1[t + 64u]; }
workgroupBarrier();
if (t < 32u) {
var s = sr1[t] + sr1[t + 32u];
s = s + subgroupShuffleDown(s, 16u);
s = s + subgroupShuffleDown(s, 8u);
s = s + subgroupShuffleDown(s, 4u);
s = s + subgroupShuffleDown(s, 2u);
s = s + subgroupShuffleDown(s, 1u);
if (t == 0u) { sbc.x = s; }
}
workgroupBarrier();
let r = sbc.x;
workgroupBarrier();
return r;
}
fn sgr4(t: u32, v: vec4<f32>) -> vec4<f32> {
sr4[t] = v;
workgroupBarrier();
if (t < 64u) { sr4[t] = sr4[t] + sr4[t + 64u]; }
workgroupBarrier();
if (t < 32u) {
var s = sr4[t] + sr4[t + 32u];
s = s + subgroupShuffleDown(s, 16u);
s = s + subgroupShuffleDown(s, 8u);
s = s + subgroupShuffleDown(s, 4u);
s = s + subgroupShuffleDown(s, 2u);
s = s + subgroupShuffleDown(s, 1u);
if (t == 0u) { sbc = s; }
}
workgroupBarrier();
let r = sbc;
workgroupBarrier();
return r;
}
@compute @workgroup_size(128)
fn gdn_step_par_k_sg(@builtin(workgroup_id) wid: vec3<u32>,
@builtin(local_invocation_id) lid: vec3<u32>) {
let h = wid.x;
let dj4 = wid.y;
let t = lid.x;
let dk = gdk_p.dk;
let dv = gdk_p.dv;
if (h >= gdk_p.nv || dj4 * 4u >= dv) { return; }
let ko = h / gdk_p.rep;
let dv4 = dv >> 2u;
let s4base = (h * dk * dv) >> 2u;
for (var i = 0u; i < gdk_p.kb; i = i + 1u) {
let cq0 = i * gdk_p.cdim;
let qs = cq0 + ko * dk;
let ks = cq0 + gdk_p.kd + ko * dk;
let nq = sgr1(t, select(0.0, gdk_cq[qs + t] * gdk_cq[qs + t], t < dk));
let nkn = sgr1(t, select(0.0, gdk_cq[ks + t] * gdk_cq[ks + t], t < dk));
let invq = 1.0 / (sqrt(nq + 1e-6) * sqrt(f32(dk)));
let invk = 1.0 / sqrt(nkn + 1e-6);
let abo = i * gdk_p.nv;
let g = exp(-exp(gdk_alog[h]) * gd_softplus(gdk_a[abo + h] + gdk_dtb[h]));
let beta = 1.0 / (1.0 + exp(-gdk_b[abo + h]));
let vto = cq0 + 2u * gdk_p.kd + h * dv + dj4 * 4u;
let vt = vec4<f32>(gdk_cq[vto], gdk_cq[vto + 1u], gdk_cq[vto + 2u], gdk_cq[vto + 3u]);
let kf_t = select(0.0, gdk_cq[ks + t] * invk, t < dk);
let qf_t = select(0.0, gdk_cq[qs + t] * invq, t < dk);
var kv4 = vec4<f32>(0.0);
if (t < dk) {
kv4 = gdk_S4[s4base + t * dv4 + dj4] * kf_t;
}
let kv = sgr4(t, kv4);
let delta = (vt - g * kv) * beta;
var contrib = vec4<f32>(0.0);
if (t < dk) {
let idx = s4base + t * dv4 + dj4;
let cell = g * gdk_S4[idx] + kf_t * delta;
gdk_S4[idx] = cell;
if (gdk_p.stride != 0u) {
gdk_snap4[(i * gdk_p.stride + gdk_p.ring_els) / 4u + idx] = cell;
}
contrib = qf_t * cell;
}
let o = sgr4(t, contrib);
if (t == 0u) {
gdk_o4[(i * gdk_p.nv * dv) / 4u + h * dv4 + dj4] = o;
}
}
}
"#;
/// The subgroup-tree GDN step twins (`CMF_GDN_SG=0` keeps the loops).
pub(crate) struct GdnSg {
pub(crate) par: wgpu::ComputePipeline,
pub(crate) par_l: wgpu::BindGroupLayout,
pub(crate) park: wgpu::ComputePipeline,
pub(crate) park_l: wgpu::BindGroupLayout,
}
fn build_gdn(c: &Ctx) -> Option<GdnSg> {
if std::env::var("CMF_GDN_SG").as_deref() == Ok("0") || !sg32(c) {
return None;
}
let par = compile(c, "cmf-gdn-sg", GDN_SG_WGSL, "gdn_step_par_sg")?;
let park = compile(c, "cmf-gdnk-sg", GDNK_SG_WGSL, "gdn_step_par_k_sg")?;
Some(GdnSg {
par_l: par.get_bind_group_layout(0),
par,
park_l: park.get_bind_group_layout(0),
park,
})
}
pub(crate) fn gdn(c: &Ctx) -> Option<&GdnSg> {
c.gdn_sg_pipes.get_or_init(|| build_gdn(c)).as_ref()
}