andodeki/makepad
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions
makepad/libs/ai/vulkan/kernels/llm.splash

2239 lines
87 KiB
Text
Raw Permalink Blame History

// The LLM graph kernels (makepad-ai-llm's Vulkan executor): one kernel per
// graph op the qwen35 decode and prefill graphs run. Every stride is in
// ELEMENTS of the buffer's type; every tensor is passed by the address of
// its first element. Quantised weights are read as raw u32 words in their
// GGUF block encoding (q4_K/q5_K: 144/176 bytes, q6_K: 210, q8_0: 34 per 32).
// Quantised weight formats as the kernels read them: `value_at(bo, l)`
// decodes value `l` (0..256) of the 256-value unit at byte `bo` (q6_K and
// q8_0 in their repacked layouts, see vulkan_exec.rs `repack`).
let deq = {
w: shader.buffer(u32)
byte: fn(o) {
return (self.w[o >> 2] >> ((o & 3) * 8)) & 255
}
// The 4 bytes at byte `o` (any alignment) as one little-endian word.
word: fn(o) {
let i = o >> 2
let sh = (o & 3) * 8
let lo = self.w[i]
if sh == 0 {
return lo
}
return (lo >> sh) | (self.w[i + 1] << (32 - sh))
}
half: fn(o) {
let word = self.w[o >> 2]
return unpack_half2((word >> ((o & 2) * 8)) & 65535).x
}
scale_k4: fn(bo, si) {
if si < 4 {
return self.byte(bo + 4 + si) & 63
}
return (self.byte(bo + 4 + si + 4) & 15) | ((self.byte(bo + 4 + si - 4) >> 6) << 4)
}
min_k4: fn(bo, si) {
if si < 4 {
return self.byte(bo + 4 + si + 4) & 63
}
return (self.byte(bo + 4 + si + 4) >> 4) | ((self.byte(bo + 4 + si) >> 6) << 4)
}
signed: fn(b) {
if b > 127 {
return float(b) - 256.0
}
return float(b)
}
}
let deq_q6k = deq{
unit_bytes: 224
value_at: fn(bo, l) {
let d = self.half(bo + 208)
let n = l >> 7
let r = l & 127
let group = r >> 5
let lo = r & 31
let ql = self.byte(bo + n * 64 + (group & 1) * 32 + lo)
let qh = self.byte(bo + 128 + n * 32 + lo)
let q = ((ql >> (((group >> 1) & 1) * 4)) & 15) | (((qh >> (group * 2)) & 3) << 4)
let sc = self.signed(self.byte(bo + 192 + n * 8 + lo / 16 + 2 * group))
return d * sc * (float(q) - 32.0)
}
code_word: fn(bo, j, w) {
let n = j >> 2
let grp = j & 3
let ql = (self.word(bo + n * 64 + (grp & 1) * 32 + w * 4) >> ((grp >> 1) * 4)) & u32(252645135)
let qh = (self.word(bo + 128 + n * 32 + w * 4) >> (grp * 2)) & u32(50529027)
return ql | (qh << 4)
}
meta: fn(bo, j) {
let n = j >> 2
let grp = j & 3
let d = self.half(bo + 208)
let sc0 = self.signed(self.byte(bo + 192 + n * 8 + 2 * grp))
let sc1 = self.signed(self.byte(bo + 192 + n * 8 + 1 + 2 * grp))
return vec4(d * sc0, d * sc1, 0.0, 32.0)
}
}
let deq_q5k = deq{
unit_bytes: 176
value_at: fn(bo, l) {
let j64 = l >> 6
let rem = l & 63
let si = 2 * j64 + (rem >> 5)
let q = self.byte(bo + 48 + 32 * j64 + (rem & 31))
var nib = float(q & 15)
if rem >= 32 {
nib = float(q >> 4)
}
let hi = float((self.byte(bo + 16 + (rem & 31)) >> si) & 1) * 16.0
return self.half(bo) * float(self.scale_k4(bo, si)) * (nib + hi) - self.half(bo + 2) * float(self.min_k4(bo, si))
}
// Codes of int8-block j (0..8) of the unit, word w (0..8): 4 values
// as unsigned bytes; meta: (scale words 0-3, scale words 4-7, min
// coefficient, code offset).
code_word: fn(bo, j, w) {
let q = self.word(bo + 48 + (j >> 1) * 32 + w * 4)
let h = self.word(bo + 16 + w * 4)
return ((q >> ((j & 1) * 4)) & u32(252645135)) | (((h >> j) & u32(16843009)) << 4)
}
meta: fn(bo, j) {
let dm = unpack_half2(self.word(bo))
let sc = float(self.scale_k4(bo, j))
let mn = float(self.min_k4(bo, j))
return vec4(dm.x * sc, dm.x * sc, dm.y * mn, 0.0)
}
}
let deq_q4k = deq{
unit_bytes: 144
value_at: fn(bo, l) {
let j64 = l >> 6
let rem = l & 63
let si = 2 * j64 + (rem >> 5)
let q = self.byte(bo + 16 + 32 * j64 + (rem & 31))
var nib = float(q & 15)
if rem >= 32 {
nib = float(q >> 4)
}
return self.half(bo) * float(self.scale_k4(bo, si)) * nib - self.half(bo + 2) * float(self.min_k4(bo, si))
}
code_word: fn(bo, j, w) {
let q = self.word(bo + 16 + (j >> 1) * 32 + w * 4)
return (q >> ((j & 1) * 4)) & u32(252645135)
}
meta: fn(bo, j) {
let dm = unpack_half2(self.word(bo))
let sc = float(self.scale_k4(bo, j))
let mn = float(self.min_k4(bo, j))
return vec4(dm.x * sc, dm.x * sc, dm.y * mn, 0.0)
}
}
let deq_q80 = deq{
unit_bytes: 272
value_at: fn(bo, l) {
let blk = l >> 5
let d = unpack_half2(self.w[(bo >> 2) + (blk >> 1)])
var dv = d.x
if (blk & 1) == 1 {
dv = d.y
}
return dv * self.signed(self.byte(bo + 16 + l))
}
code_word: fn(bo, j, w) {
return self.word(bo + 16 + 32 * j + w * 4) ^ u32(2155905152)
}
meta: fn(bo, j) {
let d = self.half(bo + 2 * j)
return vec4(d, d, 0.0, 128.0)
}
}
// Quantised rows: dst[r, :] = dequant(src[rows[r], :]) — the token
// embedding lookup. One invocation per value.
let rows_main = fn() {
let i = global_id().x
let r = group_id().y
if i < self.ne0 {
let row = u32(self.rows[r])
let bo = row * self.w_row + (i / 256) * self.unit_bytes
self.dst[r * self.dst_row + i] = self.value_at(bo, i % 256)
}
}
// Quantised matrix x matrix for prefill (many activation columns) on
// int8 activations (`quantize_q8`): a 32-row by 32-column tile per
// workgroup, k in steps of one 32-value block. The loaders put the block's
// weight codes (format `code_word`) and per-row scales (`meta`) and the
// columns' biased int8 activations in shared memory; each invocation then
// accumulates 4 rows x 2 columns with packed 4x8-bit dot products.
let gemm_main = fn() {
let t = local_id().x
let tx = t % 16
let ty = t / 16
let row0 = group_id().x * 32
let col0 = group_id().y * 32
var acc0 = vec4(0.0, 0.0, 0.0, 0.0)
var acc1 = vec4(0.0, 0.0, 0.0, 0.0)
let ones = u32(16843009)
for kb in 0..self.k / 32 {
let unit = kb / 8
let j = kb % 8
// Weights: 32 rows x 8 words, two words per invocation.
let lr = t / 4
let lw = (t % 4) * 2
let row = min(row0 + lr, self.n - 1)
let bo = row * self.w_row + unit * self.unit_bytes
self.wq[lr * 8 + lw] = self.code_word(bo, j, lw)
self.wq[lr * 8 + lw + 1] = self.code_word(bo, j, lw + 1)
if (t % 4) == 0 {
self.wm[lr] = self.meta(bo, j)
}
// Activations: 32 columns x 8 words, two per invocation.
let col = min(col0 + lr, self.m - 1)
let xw = (col * self.k + kb * 32) / 4
self.xq_s[lr * 8 + lw] = self.xq[xw + lw]
self.xq_s[lr * 8 + lw + 1] = self.xq[xw + lw + 1]
if (t % 4) == 0 {
self.xd_s[lr] = self.xs[(col * self.k + kb * 32) / 32].x
}
barrier()
for c in 0..2 {
let cc = tx * 2 + c
// sum q over each half of the column's block.
var sq0 = 0
var sq1 = 0
for w in 0..4 {
sq0 += i32(dot4_u8(self.xq_s[cc * 8 + w], ones)) - 512
sq1 += i32(dot4_u8(self.xq_s[cc * 8 + 4 + w], ones)) - 512
}
let xd = self.xd_s[cc]
var part = vec4(0.0, 0.0, 0.0, 0.0)
for r in 0..4 {
let rr = ty * 4 + r
var d0 = 0
var d1 = 0
var c0 = 0
var c1 = 0
for w in 0..4 {
let a0 = self.wq[rr * 8 + w]
let a1 = self.wq[rr * 8 + 4 + w]
d0 += i32(dot4_u8(a0, self.xq_s[cc * 8 + w]))
d1 += i32(dot4_u8(a1, self.xq_s[cc * 8 + 4 + w]))
c0 += i32(dot4_u8(a0, ones))
c1 += i32(dot4_u8(a1, ones))
}
let mt = self.wm[rr]
// sum(code * q) per half, less the code offset, less the min.
let v0 = float(d0 - 128 * c0) - mt.w * float(sq0)
let v1 = float(d1 - 128 * c1) - mt.w * float(sq1)
part[r] = xd * (mt.x * v0 + mt.y * v1 - mt.z * float(sq0 + sq1))
}
if c == 0 {
acc0 += part
} else {
acc1 += part
}
}
barrier()
}
for c in 0..2 {
let col = col0 + tx * 2 + c
var acc = acc0
if c == 1 {
acc = acc1
}
if col < self.m {
for r in 0..4 {
let row = row0 + ty * 4 + r
if row < self.n {
self.dst[col * self.d_col + row] = acc[r]
}
}
}
}
}
// Prefill register tiles. Each invocation owns four rows by four columns;
// named vectors keep the hot loop out of dynamically indexed local arrays.
// Codes are decoded only into workgroup memory, never a second weight buffer.
// Stage two (or four) 32-value blocks per barrier pair. Weight/activation
// sums are computed by the loaders once, then reused by the whole tile.
let gemm_register = deq{
workgroup_size: 256
tile_n: 64
tile_k: 64
format: 5
unit_bytes: 176
w4: shader.buffer(vec4u)
xq4: shader.buffer(vec4u)
xs: shader.buffer(vec2f)
dst: shader.buffer_rw(f32)
k: shader.uniform(u32)
n: shader.uniform(u32)
m: shader.uniform(u32)
w_row: shader.uniform(u32)
d_col: shader.uniform(u32)
code16: fn(bo, j, h) {
let at = bo / 16
let mask = vec4u(u32(252645135))
if self.format == 6 {
let half = j >> 2
let g = j & 3
let lo = self.w4[at + half * 4 + (g & 1) * 2 + h]
let hi = self.w4[at + 8 + half * 2 + h]
return ((lo >> vec4u((g >> 1) * 4)) & mask)
| (((hi >> vec4u(g * 2)) & vec4u(u32(50529027))) << vec4u(4))
}
if self.format == 8 {
return self.w4[at + 1 + j * 2 + h] ^ vec4u(u32(2155905152))
}
var qs = u32(1)
if self.format == 5 { qs = u32(3) }
let lo = self.w4[at + qs + (j >> 1) * 2 + h]
var codes = (lo >> vec4u((j & 1) * 4)) & mask
if self.format == 5 {
let hi = self.w4[at + 1 + h]
codes = codes | (((hi >> vec4u(j)) & vec4u(u32(16843009))) << vec4u(4))
}
return codes
}
dot16: fn(a, b) {
return i32(dot4_u8(a.x, b.x) + dot4_u8(a.y, b.y)
+ dot4_u8(a.z, b.z) + dot4_u8(a.w, b.w))
}
sum16: fn(a) {
return self.dot16(a, vec4u(u32(16843009)))
}
dots: fn(a, b0, b1, b2, b3) {
return vec4i(self.dot16(a, b0), self.dot16(a, b1), self.dot16(a, b2), self.dot16(a, b3))
}
store: fn(r, c, v) {
if r < self.n {
if c < self.m { self.dst[c * self.d_col + r] = v.x }
if c + 1 < self.m { self.dst[(c + 1) * self.d_col + r] = v.y }
if c + 2 < self.m { self.dst[(c + 2) * self.d_col + r] = v.z }
if c + 3 < self.m { self.dst[(c + 3) * self.d_col + r] = v.w }
}
}
}
// Four rows by eight columns per invocation, 128 invocations per 64x64
// tile, staging 16 values of k per step. Q4/Q5/Q8 combine
// a whole 32-value integer dot before applying its scale and minimum.
// Metadata and bias corrections are prepared once per row/column block.
let gemm_wide = gemm_register{
workgroup_size: 128
tile_k: 32
wide_meta: fn(bo, j) {
if self.format == 6 {
let s = self.w[bo / 4 + 48 + j / 2] >> ((j & 1) * 16)
return vec4(self.signed(s & 255), self.signed((s >> 8) & 255), self.half(bo + 208), 32.0)
}
if self.format == 8 {
return vec4(self.half(bo + j * 2), 0.0, 0.0, 128.0)
}
// One aligned header load supplies both half scales and all eight
// packed sub-block scale/min pairs; no byte-at-a-time loads.
let header = self.w4[bo / 16]
let dm = unpack_half2(header.x)
let shift = (j & 3) * 8
var sc = (header.y >> shift) & 63
var mn = (header.z >> shift) & 63
if j >= 4 {
let extra = header.w >> shift
sc = (extra & 15) | (((header.y >> (shift + 6)) & 3) << 4)
mn = ((extra >> 4) & 15) | (((header.z >> (shift + 6)) & 3) << 4)
}
return vec4(dm.x * float(sc), dm.y * float(mn), 0.0, 0.0)
}
load_wide: fn(base, rn, cm) {
let lane = local_id().x
let blocks = u32(max(self.tile_k / 32, 1))
let entries = 64 * blocks
var i = lane
while i < entries * 2 {
let index = i % entries
let kb = base / 32 + index % blocks
if i < entries {
let row = min(rn + index / blocks, self.n - 1)
let bo = row * self.w_row + (kb / 8) * self.unit_bytes
var mt = self.wide_meta(bo, kb % 8)
if self.tile_k == 16 {
let h = (base / 16) & 1
let a = self.code16(bo, kb % 8, h)
self.wq[index] = a
if self.format == 6 {
var sc = mt.x
if h == 1 { sc = mt.y }
mt = vec4(mt.z * sc, 0.0, 0.0, 32.0)
}
mt.z = float(128 * self.sum16(a))
} else {
let a = self.code16(bo, kb % 8, u32(0))
let b = self.code16(bo, kb % 8, u32(1))
self.wq[index * 2] = a
self.wq[index * 2 + 1] = b
let ca = 128 * self.sum16(a)
let cb = 128 * self.sum16(b)
if self.format == 6 { self.wc[index] = vec2i(ca, cb) }
else { mt.z = float(ca + cb) }
}
self.wm[index] = mt
} else {
let col = min(cm + index / blocks, self.m - 1)
let at = col * (self.k / 32) + kb
let xs = self.xs[at]
if self.tile_k == 16 {
let a = self.xq4[at * 2 + ((base / 16) & 1)]
self.xq_s[index] = a
self.xm[index] = vec4(xs.x, float(self.sum16(a) - 2048), 0.0, 0.0)
} else {
let a = self.xq4[at * 2]
let b = self.xq4[at * 2 + 1]
self.xq_s[index * 2] = a
self.xq_s[index * 2 + 1] = b
if self.format == 6 {
let sa = float(self.sum16(a) - 2048)
self.xm[index] = vec4(xs.x, sa, xs.y - sa, 0.0)
} else {
// The quantiser has already computed the signed sum.
self.xm[index] = vec4(xs.x, xs.y, 0.0, 0.0)
}
}
}
i += self.workgroup_size
}
}
half_block: fn(at, dot, xd, sx) {
let mt = self.wm[at]
var v = vec4(float(dot.x), float(dot.y), float(dot.z), float(dot.w)) - mt.z
if self.format == 6 || self.format == 8 { v -= mt.w * sx }
return xd * (mt.x * v - mt.y * sx)
}
whole_block: fn(at, lo, hi, xd, sx0, sx1) {
let mt = self.wm[at]
if self.format == 6 {
let correction = self.wc[at]
let a = lo - vec4i(correction.x)
let b = hi - vec4i(correction.y)
let v0 = vec4(float(a.x), float(a.y), float(a.z), float(a.w)) - 32.0 * sx0
let v1 = vec4(float(b.x), float(b.y), float(b.z), float(b.w)) - 32.0 * sx1
// Both signed-scale products and their sum are exact in f32.
return xd * (mt.z * (mt.x * v0 + mt.y * v1))
}
return self.half_block(at, lo + hi, xd, sx0)
}
main: fn() {
let lane = local_id().x
let rn = group_id().x * 64
let cm = group_id().y * 64
let r = (lane / 8) * 4
let c = (lane % 8) * 8
let blocks = u32(max(self.tile_k / 32, 1))
let halves = u32(min(self.tile_k / 16, 2))
var c00 = vec4f(0)
var c01 = vec4f(0)
var c10 = vec4f(0)
var c11 = vec4f(0)
var c20 = vec4f(0)
var c21 = vec4f(0)
var c30 = vec4f(0)
var c31 = vec4f(0)
var err00 = vec4f(0)
var err01 = vec4f(0)
var err10 = vec4f(0)
var err11 = vec4f(0)
var err20 = vec4f(0)
var err21 = vec4f(0)
var err30 = vec4f(0)
var err31 = vec4f(0)
for base in 0..self.k / self.tile_k {
self.load_wide(base * self.tile_k, rn, cm)
barrier()
for d in 0..blocks {
let at = c * blocks + d
let x0 = self.xm[at]
let x1 = self.xm[at + blocks]
let x2 = self.xm[at + 2 * blocks]
let x3 = self.xm[at + 3 * blocks]
let x4 = self.xm[at + 4 * blocks]
let x5 = self.xm[at + 5 * blocks]
let x6 = self.xm[at + 6 * blocks]
let x7 = self.xm[at + 7 * blocks]
let xd0 = vec4(x0.x, x1.x, x2.x, x3.x)
let xd1 = vec4(x4.x, x5.x, x6.x, x7.x)
let sx00 = vec4(x0.y, x1.y, x2.y, x3.y)
let sx01 = vec4(x4.y, x5.y, x6.y, x7.y)
let sx10 = vec4(x0.z, x1.z, x2.z, x3.z)
let sx11 = vec4(x4.z, x5.z, x6.z, x7.z)
let b0 = self.xq_s[at * halves]
let b1 = self.xq_s[(at + blocks) * halves]
let b2 = self.xq_s[(at + 2 * blocks) * halves]
let b3 = self.xq_s[(at + 3 * blocks) * halves]
let b4 = self.xq_s[(at + 4 * blocks) * halves]
let b5 = self.xq_s[(at + 5 * blocks) * halves]
let b6 = self.xq_s[(at + 6 * blocks) * halves]
let b7 = self.xq_s[(at + 7 * blocks) * halves]
let wr = r * blocks + d
let a0 = self.wq[wr * halves]
let a1 = self.wq[(wr + blocks) * halves]
let a2 = self.wq[(wr + 2 * blocks) * halves]
let a3 = self.wq[(wr + 3 * blocks) * halves]
let lo00 = self.dots(a0, b0, b1, b2, b3)
let lo01 = self.dots(a0, b4, b5, b6, b7)
let lo10 = self.dots(a1, b0, b1, b2, b3)
let lo11 = self.dots(a1, b4, b5, b6, b7)
let lo20 = self.dots(a2, b0, b1, b2, b3)
let lo21 = self.dots(a2, b4, b5, b6, b7)
let lo30 = self.dots(a3, b0, b1, b2, b3)
let lo31 = self.dots(a3, b4, b5, b6, b7)
var p00 = vec4f(0)
var p01 = vec4f(0)
var p10 = vec4f(0)
var p11 = vec4f(0)
var p20 = vec4f(0)
var p21 = vec4f(0)
var p30 = vec4f(0)
var p31 = vec4f(0)
if self.tile_k == 16 {
p00 = self.half_block(wr, lo00, xd0, sx00)
p01 = self.half_block(wr, lo01, xd1, sx01)
p10 = self.half_block(wr + blocks, lo10, xd0, sx00)
p11 = self.half_block(wr + blocks, lo11, xd1, sx01)
p20 = self.half_block(wr + 2 * blocks, lo20, xd0, sx00)
p21 = self.half_block(wr + 2 * blocks, lo21, xd1, sx01)
p30 = self.half_block(wr + 3 * blocks, lo30, xd0, sx00)
p31 = self.half_block(wr + 3 * blocks, lo31, xd1, sx01)
} else {
let e0 = self.xq_s[at * 2 + 1]
let e1 = self.xq_s[(at + blocks) * 2 + 1]
let e2 = self.xq_s[(at + 2 * blocks) * 2 + 1]
let e3 = self.xq_s[(at + 3 * blocks) * 2 + 1]
let e4 = self.xq_s[(at + 4 * blocks) * 2 + 1]
let e5 = self.xq_s[(at + 5 * blocks) * 2 + 1]
let e6 = self.xq_s[(at + 6 * blocks) * 2 + 1]
let e7 = self.xq_s[(at + 7 * blocks) * 2 + 1]
let h0 = self.wq[wr * 2 + 1]
p00 = self.whole_block(wr, lo00, self.dots(h0, e0, e1, e2, e3), xd0, sx00, sx10)
p01 = self.whole_block(wr, lo01, self.dots(h0, e4, e5, e6, e7), xd1, sx01, sx11)
let h1 = self.wq[(wr + blocks) * 2 + 1]
p10 = self.whole_block(wr + blocks, lo10, self.dots(h1, e0, e1, e2, e3), xd0, sx00, sx10)
p11 = self.whole_block(wr + blocks, lo11, self.dots(h1, e4, e5, e6, e7), xd1, sx01, sx11)
let h2 = self.wq[(wr + 2 * blocks) * 2 + 1]
p20 = self.whole_block(wr + 2 * blocks, lo20, self.dots(h2, e0, e1, e2, e3), xd0, sx00, sx10)
p21 = self.whole_block(wr + 2 * blocks, lo21, self.dots(h2, e4, e5, e6, e7), xd1, sx01, sx11)
let h3 = self.wq[(wr + 3 * blocks) * 2 + 1]
p30 = self.whole_block(wr + 3 * blocks, lo30, self.dots(h3, e0, e1, e2, e3), xd0, sx00, sx10)
p31 = self.whole_block(wr + 3 * blocks, lo31, self.dots(h3, e4, e5, e6, e7), xd1, sx01, sx11)
}
if self.format == 6 {
// Retain the Q6 accumulation residual across blocks.
let y00 = p00 + err00
let y01 = p01 + err01
let y10 = p10 + err10
let y11 = p11 + err11
let y20 = p20 + err20
let y21 = p21 + err21
let y30 = p30 + err30
let y31 = p31 + err31
let t00 = c00 + y00
let t01 = c01 + y01
let t10 = c10 + y10
let t11 = c11 + y11
let t20 = c20 + y20
let t21 = c21 + y21
let t30 = c30 + y30
let t31 = c31 + y31
err00 = fma(vec4f(-1), t00, c00) + y00
err01 = fma(vec4f(-1), t01, c01) + y01
err10 = fma(vec4f(-1), t10, c10) + y10
err11 = fma(vec4f(-1), t11, c11) + y11
err20 = fma(vec4f(-1), t20, c20) + y20
err21 = fma(vec4f(-1), t21, c21) + y21
err30 = fma(vec4f(-1), t30, c30) + y30
err31 = fma(vec4f(-1), t31, c31) + y31
c00 = t00
c01 = t01
c10 = t10
c11 = t11
c20 = t20
c21 = t21
c30 = t30
c31 = t31
} else {
c00 += p00
c01 += p01
c10 += p10
c11 += p11
c20 += p20
c21 += p21
c30 += p30
c31 += p31
}
}
barrier()
}
self.store(rn + r, cm + c, c00)
self.store(rn + r, cm + c + 4, c01)
self.store(rn + r + 1, cm + c, c10)
self.store(rn + r + 1, cm + c + 4, c11)
self.store(rn + r + 2, cm + c, c20)
self.store(rn + r + 2, cm + c + 4, c21)
self.store(rn + r + 3, cm + c, c30)
self.store(rn + r + 3, cm + c + 4, c31)
}
}
let gemm_wide_16 = gemm_wide{
tile_k: 16
wq: shader.workgroup(vec4u, 64)
wm: shader.workgroup(vec4f, 64)
wc: shader.workgroup(vec2i, 64)
xq_s: shader.workgroup(vec4u, 64)
xm: shader.workgroup(vec4f, 64)
}
// Quantised matrix x vector on 8-bit activations. `quantize_q8` turns
// each 32 activations into int8 q = round(x / d), d = max|x| / 127, stored
// biased (q + 128, four per word) with (d, sum q) per block; the gemv then
// runs packed 4x8-bit unsigned dot products (dot4_u8) on the weights' own
// integer codes and applies the scales once per block. Integer work is
// what phone GPUs are slow at; this keeps it to a few ops per 4 weights.
//
// dst[c, row] = sum_k w[row, k] * x[c, k] for the m activation columns. A
// 32-lane group computes 4 consecutive rows; each lane takes 8 values of
// every 256-value unit (positions per format) and a butterfly sums lanes.
let gemv = {
workgroup_size: 128
w: shader.buffer(u32)
w4: shader.buffer(vec4u)
xq: shader.buffer(u32)
xs: shader.buffer(vec2f)
dst: shader.buffer_rw(f32)
k: shader.uniform(u32)
n: shader.uniform(u32)
m: shader.uniform(u32)
w_row: shader.uniform(u32)
d_col: shader.uniform(u32)
main: fn() {
let lane = local_id().x % 32
let row0 = (group_id().x * (self.workgroup_size / 32) + local_id().x / 32) * 4
if row0 < self.n {
let units = self.k / 256
for col in 0..self.m {
var acc = array(0.0, 0.0, 0.0, 0.0)
for u in 0..units {
// This unit's activations: word and block indices.
let xw = (col * self.k + u * 256) / 4
let xb = (col * self.k + u * 256) / 32
for r in 0..4 {
let row = min(row0 + r, self.n - 1)
acc[r] += self.unit_dot(row * self.w_row + u * self.unit_bytes, xw, xb, lane)
}
}
for r in 0..4 {
let total = self.lane_sum(acc[r])
if lane == 0 && row0 + r < self.n {
self.dst[col * self.d_col + row0 + r] = total
}
}
}
}
}
// Byte `o` of the weight tensor.
byte: fn(o) {
return (self.w[o >> 2] >> ((o & 3) * 8)) & 255
}
// The 4 bytes at byte `o` (any alignment) as one little-endian word.
word: fn(o) {
let i = o >> 2
let sh = (o & 3) * 8
let lo = self.w[i]
if sh == 0 {
return lo
}
return (lo >> sh) | (self.w[i + 1] << (32 - sh))
}
// sum(a_i * q_i) for 4 unsigned codes `a` against 4 biased
// activations `x` (q + 128 each), as an integer.
dot_codes: fn(a, x) {
return i32(dot4_u8(a, x)) - 128 * i32(dot4_u8(a, u32(16843009)))
}
// sum(q_i) of 4 biased activations.
sum_x: fn(x) {
return i32(dot4_u8(x, u32(16843009))) - 512
}
// Byte `j` (0..12) of the three words s0 s1 s2.
pick: fn(s0, s1, s2, j) {
var v = s2
if j < 4 {
v = s0
} else if j < 8 {
v = s1
}
return (v >> ((j & 3) * 8)) & 255
}
sc_k4: fn(s0, s1, s2, sb) {
if sb < 4 {
return self.pick(s0, s1, s2, sb) & 63
}
return (self.pick(s0, s1, s2, sb + 4) & 15) | ((self.pick(s0, s1, s2, sb - 4) >> 6) << 4)
}
mn_k4: fn(s0, s1, s2, sb) {
if sb < 4 {
return self.pick(s0, s1, s2, sb + 4) & 63
}
return (self.pick(s0, s1, s2, sb + 4) >> 4) | ((self.pick(s0, s1, s2, sb) >> 6) << 4)
}
lane_sum: fn(v) {
var s = v
s += subgroup_shuffle_xor(s, 16)
s += subgroup_shuffle_xor(s, 8)
s += subgroup_shuffle_xor(s, 4)
s += subgroup_shuffle_xor(s, 2)
s += subgroup_shuffle_xor(s, 1)
return s
}
}
// q4_K / q5_K: f16 d and dmin, 12 bytes of 6-bit scales and mins, (q5_K:
// 32 bytes of fifth bits), 128 bytes of nibbles; every part 16-aligned.
// Eight lanes take one 256-value unit, a lane 16 nibble bytes (one 16-byte
// load): 16 low-nibble values of scale group 2j and 16 high-nibble values
// of group 2j + 1, j = lane / 2; four units run side by side per row.
let gemv_k = {
workgroup_size: 256
w4: shader.buffer(vec4u)
xq4: shader.buffer(vec4u)
xs: shader.buffer(vec2f)
xs4: shader.buffer(vec4f)
dst: shader.buffer_rw(f32)
k: shader.uniform(u32)
n: shader.uniform(u32)
m: shader.uniform(u32)
w_row: shader.uniform(u32)
d_col: shader.uniform(u32)
// Rows per 32-lane group, 1 or 4 (4 share the activation loads; the
// accumulators are a vec4, never a dynamically indexed array, which
// phone compilers put in memory).
rpg: 1
main: fn() {
let lane = local_id().x % 32
let row0 = (group_id().x * (self.workgroup_size / 32) + local_id().x / 32) * self.rpg
if row0 < self.n {
let units = self.k / 256
let slot = lane >> 3
let part = lane & 7
let r1 = min(row0 + 1, self.n - 1)
let r2 = min(row0 + 2, self.n - 1)
let r3 = min(row0 + 3, self.n - 1)
for col in 0..self.m {
var acc = vec4(0.0, 0.0, 0.0, 0.0)
var u = slot
while u < units {
let xbase = col * self.k + u * 256 + (part >> 1) * 64 + (part & 1) * 16
let x_lo = self.xq4[xbase >> 4]
let x_hi = self.xq4[(xbase + 32) >> 4]
let sc = self.xs4[(col * self.k + u * 256) / 64 + (part >> 1)]
let ub = u * self.unit_bytes
acc.x += self.unit_part(row0 * self.w_row + ub, part, x_lo, x_hi, sc)
if self.rpg == 4 {
acc.y += self.unit_part(r1 * self.w_row + ub, part, x_lo, x_hi, sc)
acc.z += self.unit_part(r2 * self.w_row + ub, part, x_lo, x_hi, sc)
acc.w += self.unit_part(r3 * self.w_row + ub, part, x_lo, x_hi, sc)
}
u += 4
}
let t0 = self.lane_sum(acc.x)
if lane == 0 {
self.dst[col * self.d_col + row0] = t0
}
if self.rpg == 4 {
let t1 = self.lane_sum(acc.y)
let t2 = self.lane_sum(acc.z)
let t3 = self.lane_sum(acc.w)
if lane == 0 {
if row0 + 1 < self.n {
self.dst[col * self.d_col + row0 + 1] = t1
}
if row0 + 2 < self.n {
self.dst[col * self.d_col + row0 + 2] = t2
}
if row0 + 3 < self.n {
self.dst[col * self.d_col + row0 + 3] = t3
}
}
}
}
}
}
unit_part: fn(bo, part, x_lo, x_hi, sc) {
let head = self.w4[bo >> 4]
let dm = unpack_half2(head.x)
let q = self.w4[(bo + self.qs_at) / 16 + part]
let sb = (part >> 1) * 2
var lo = vec4u(q.x & u32(252645135), q.y & u32(252645135), q.z & u32(252645135), q.w & u32(252645135))
var hi = vec4u((q.x >> 4) & u32(252645135), (q.y >> 4) & u32(252645135), (q.z >> 4) & u32(252645135), (q.w >> 4) & u32(252645135))
if self.high_bit == 1 {
let h = self.w4[(bo + 16) / 16 + (part & 1)]
lo = lo | (((h >> vec4u(sb, sb, sb, sb)) & vec4u(u32(16843009), u32(16843009), u32(16843009), u32(16843009))) << vec4u(4, 4, 4, 4))
hi = hi | (((h >> vec4u(sb + 1, sb + 1, sb + 1, sb + 1)) & vec4u(u32(16843009), u32(16843009), u32(16843009), u32(16843009))) << vec4u(4, 4, 4, 4))
}
let sc_lo = float(self.sc_k4(head.y, head.z, head.w, sb))
let m_lo = float(self.mn_k4(head.y, head.z, head.w, sb))
let sc_hi = float(self.sc_k4(head.y, head.z, head.w, sb + 1))
let m_hi = float(self.mn_k4(head.y, head.z, head.w, sb + 1))
let d_lo = self.dot16(lo, x_lo)
let d_hi = self.dot16(hi, x_hi)
let s_lo = self.sum16(x_lo)
let s_hi = self.sum16(x_hi)
return sc.x * (dm.x * sc_lo * float(d_lo) - dm.y * m_lo * float(s_lo))
+ sc.z * (dm.x * sc_hi * float(d_hi) - dm.y * m_hi * float(s_hi))
}
// sum(a_i * q_i) over 16 unsigned codes against 16 biased activations.
dot16: fn(a, x) {
let ones = u32(16843009)
let raw = dot4_u8(a.x, x.x) + dot4_u8(a.y, x.y) + dot4_u8(a.z, x.z) + dot4_u8(a.w, x.w)
let asum = dot4_u8(a.x, ones) + dot4_u8(a.y, ones) + dot4_u8(a.z, ones) + dot4_u8(a.w, ones)
return i32(raw) - 128 * i32(asum)
}
// sum(q_i) of 16 biased activations.
sum16: fn(x) {
let ones = u32(16843009)
return i32(dot4_u8(x.x, ones) + dot4_u8(x.y, ones) + dot4_u8(x.z, ones) + dot4_u8(x.w, ones)) - 2048
}
pick: fn(s0, s1, s2, j) {
var v = s2
if j < 4 {
v = s0
} else if j < 8 {
v = s1
}
return (v >> ((j & 3) * 8)) & 255
}
sc_k4: fn(s0, s1, s2, sb) {
if sb < 4 {
return self.pick(s0, s1, s2, sb) & 63
}
return (self.pick(s0, s1, s2, sb + 4) & 15) | ((self.pick(s0, s1, s2, sb - 4) >> 6) << 4)
}
mn_k4: fn(s0, s1, s2, sb) {
if sb < 4 {
return self.pick(s0, s1, s2, sb + 4) & 63
}
return (self.pick(s0, s1, s2, sb + 4) >> 4) | ((self.pick(s0, s1, s2, sb) >> 6) << 4)
}
lane_sum: fn(v) {
var s = v
s += subgroup_shuffle_xor(s, 16)
s += subgroup_shuffle_xor(s, 8)
s += subgroup_shuffle_xor(s, 4)
s += subgroup_shuffle_xor(s, 2)
s += subgroup_shuffle_xor(s, 1)
return s
}
}
// Opt-in decode GEMVs. Each lane still loads 16-byte weight payloads;
// activation loads and corrections are outside the explicit row calls.
// All accumulators are registers, never dynamically indexed local arrays.
// Source-level counts per 16 nibble bytes (32 codes), excluding addresses,
// scale extraction and float accumulation: Q4 = 28 integer ops + 10
// dot4, Q5 = 56 integer ops + 10 dot4, Q6 = 56 integer ops + 10 dot4.
// Q4/Q5/Q6 additionally need 8 dot4 + 8 integer ops for activation sums
// shared across the rows. Q8 per 32 code bytes = 25 integer ops + 16 dot4, with
// its activation sum read from quantize_q8's metadata. Conversions omitted.
// These are algorithmic counts, not claims about a driver's instruction ISA.
let gemv_decode = gemv_k{
format: 5
unit_bytes: 176
row_lanes: 32
rpg: 1
groups_x: shader.uniform(u32)
partial: shader.workgroup(vec4f, 16)
main: fn() {
let t = local_id().x
let lane = t % self.row_lanes
let group = group_id().x + group_id().y * self.groups_x
let row0 = (group * (self.workgroup_size / self.row_lanes) + t / self.row_lanes) * self.rpg
let row1 = min(row0 + 1, self.n - 1)
let row2 = min(row0 + 2, self.n - 1)
let row3 = min(row0 + 3, self.n - 1)
let slot = lane >> 3
let part = lane & 7
for col in 0..self.m {
var acc = vec4(0.0, 0.0, 0.0, 0.0)
if row0 < self.n {
var u = slot
while u < self.k / 256 {
let xb = col * self.k + u * 256
var lo_at = xb + (part >> 1) * 64 + (part & 1) * 16
var hi_at = lo_at + 32
if self.format == 6 {
lo_at = xb + (part >> 2) * 128 + ((part >> 1) & 1) * 32 + (part & 1) * 16
hi_at = lo_at + 64
}
if self.format == 8 {
lo_at = xb + part * 32
hi_at = lo_at + 16
}
let xlo = self.xq4[lo_at >> 4]
let xhi = self.xq4[hi_at >> 4]
// Keep the original per-lane correction order. Splitting
// a full-block sum between lanes changes decode rounding.
var sc = vec4(0.0, 0.0, 0.0, 0.0)
if self.format == 4 || self.format == 5 {
let meta = self.xs4[xb / 64 + (part >> 1)]
sc = vec4(meta.x, float(self.sum16(xlo)), meta.z, float(self.sum16(xhi)))
} else if self.format == 6 {
sc = vec4(self.xs[lo_at / 32].x, float(self.sum16(xlo)), self.xs[hi_at / 32].x, float(self.sum16(xhi)))
} else {
let slo = self.xs[lo_at / 32]
sc = vec4(slo.x, slo.y, 0.0, 0.0)
}
let ub = u * self.unit_bytes
acc.x += self.decode_part(row0 * self.w_row + ub, part, xlo, xhi, sc)
if self.rpg >= 2 {
acc.y += self.decode_part(row1 * self.w_row + ub, part, xlo, xhi, sc)
}
if self.rpg == 4 {
acc.z += self.decode_part(row2 * self.w_row + ub, part, xlo, xhi, sc)
acc.w += self.decode_part(row3 * self.w_row + ub, part, xlo, xhi, sc)
}
u += self.row_lanes / 8
}
}
// Wide row groups use one subgroup reduction when the hardware
// width matches. Smaller subgroups use shared partials;
// every invocation reaches both barriers, including tail rows.
var total = acc
let sg = subgroup_size()
if self.row_lanes == 32 {
// Preserve the current kernel's tree, independently of
// the hardware subgroup width; this path needs no barrier.
total += subgroup_shuffle_xor(total, 16)
total += subgroup_shuffle_xor(total, 8)
total += subgroup_shuffle_xor(total, 4)
total += subgroup_shuffle_xor(total, 2)
total += subgroup_shuffle_xor(total, 1)
} else if sg == self.row_lanes {
total = subgroup_add(acc)
} else if sg < self.row_lanes {
let sub = subgroup_add(acc)
if subgroup_lane() == 0 {
self.partial[t / sg] = sub
}
barrier()
if lane == 0 {
let first = t / sg
total = self.partial[first]
for i in 1..self.row_lanes / sg {
total += self.partial[first + i]
}
}
barrier()
} else {
if self.row_lanes == 64 {
total += subgroup_shuffle_xor(total, 32)
}
total += subgroup_shuffle_xor(total, 16)
total += subgroup_shuffle_xor(total, 8)
total += subgroup_shuffle_xor(total, 4)
total += subgroup_shuffle_xor(total, 2)
total += subgroup_shuffle_xor(total, 1)
}
if lane == 0 && row0 < self.n {
self.dst[col * self.d_col + row0] = total.x
if self.rpg >= 2 && row0 + 1 < self.n {
self.dst[col * self.d_col + row0 + 1] = total.y
}
if self.rpg == 4 {
if row0 + 2 < self.n {
self.dst[col * self.d_col + row0 + 2] = total.z
}
if row0 + 3 < self.n {
self.dst[col * self.d_col + row0 + 3] = total.w
}
}
}
}
}
// Four codes sum to <=252 in each byte: no inter-byte carry. Bias
// correction uses one dot4 instead of four (5 dot4 per 16 codes).
dot16small: fn(a, x) {
let raw = self.dot16raw(a, x)
let sum = dot4_u8(a.x + a.y + a.z + a.w, u32(16843009))
return i32(raw) - (i32(sum) << 7)
}
dot16raw: fn(a, x) {
return dot4_u8(a.x, x.x) + dot4_u8(a.y, x.y) + dot4_u8(a.z, x.z) + dot4_u8(a.w, x.w)
}
decode_part: fn(bo, part, xlo, xhi, sc) {
let b16 = bo >> 4
let m4 = vec4u(u32(252645135), u32(252645135), u32(252645135), u32(252645135))
if self.format == 8 {
let ds = self.w4[b16]
var dw = ds.w
if part < 2 {
dw = ds.x
} else if part < 4 {
dw = ds.y
} else if part < 6 {
dw = ds.z
}
let d = unpack_half2(dw >> ((part & 1) * 16)).x
let flip = vec4u(u32(2155905152), u32(2155905152), u32(2155905152), u32(2155905152))
let a = self.w4[b16 + 1 + part * 2] ^ flip
let b = self.w4[b16 + 2 + part * 2] ^ flip
let ones = vec4u(u32(16843009), u32(16843009), u32(16843009), u32(16843009))
let raw = self.dot16raw(a, xlo) + self.dot16raw(b, xhi)
let wsum = self.dot16raw(a, ones) + self.dot16raw(b, ones)
let v = i32(raw) - ((i32(wsum) + i32(sc.y)) << 7)
return d * sc.x * float(v)
}
if self.format == 6 {
let half = part >> 2
let g = (part >> 1) & 1
let q = self.w4[b16 + half * 4 + (part & 3)]
let high = self.w4[b16 + 8 + half * 2 + (part & 1)]
let sh = g * 2
let h = high >> vec4u(sh, sh, sh, sh)
let m6 = vec4u(u32(808464432), u32(808464432), u32(808464432), u32(808464432))
let lo = (q & m4) | ((h << vec4u(4, 4, 4, 4)) & m6)
let hi = ((q >> vec4u(4, 4, 4, 4)) & m4) | (h & m6)
let scales = self.w4[b16 + 12]
var sl = scales.x
var shigh = scales.y
if half == 1 {
sl = scales.z
shigh = scales.w
}
let shift = 24 - (part & 3) * 8
let dl = float(i32(sl << shift) >> 24)
let dh = float(i32(shigh << shift) >> 24)
let d = unpack_half2(self.w4[b16 + 13].x & u32(65535)).x
let vl = self.dot16small(lo, xlo) - 32 * i32(sc.y)
let vh = self.dot16small(hi, xhi) - 32 * i32(sc.w)
return d * (dl * sc.x * float(vl) + dh * sc.z * float(vh))
}
let head = self.w4[b16]
let dm = unpack_half2(head.x)
var at = b16 + 1 + part
if self.format == 5 {
at += 2
}
let q = self.w4[at]
var lo = q & m4
var hi = (q >> vec4u(4, 4, 4, 4)) & m4
if self.format == 5 {
let high = self.w4[b16 + 1 + (part & 1)]
let sb = part & 6
let h = high >> vec4u(sb, sb, sb, sb)
let m5 = vec4u(u32(269488144), u32(269488144), u32(269488144), u32(269488144))
lo = lo | ((h << vec4u(4, 4, 4, 4)) & m5)
hi = hi | ((h << vec4u(3, 3, 3, 3)) & m5)
}
// Decode the adjacent sub-block scales/mins together as packed
// bytes, before extracting the pair and applying each scale once.
var scales = head.y & u32(1061109567)
var mins = head.z & u32(1061109567)
if part >= 4 {
scales = (head.w & u32(252645135)) | ((head.y >> 2) & u32(808464432))
mins = ((head.w >> 4) & u32(252645135)) | ((head.z >> 2) & u32(808464432))
}
let shift = (part & 2) * 8
let sp = scales >> shift
let mp = mins >> shift
let sl = float(sp & 255)
let sh = float((sp >> 8) & 255)
let ml = float(mp & 255)
let mh = float((mp >> 8) & 255)
let vl = self.dot16small(lo, xlo)
let vh = self.dot16small(hi, xhi)
return sc.x * (dm.x * sl * float(vl) - dm.y * ml * sc.y)
+ sc.z * (dm.x * sh * float(vh) - dm.y * mh * sc.w)
}
}
// Gated delta rule over the tokens of one sequence. 32 lanes own one
// state column (sv/32 elements each); the state is stored transposed
// (a column's sv values contiguous), q/k heads map to value heads by
// modulo, the gate decays the state before the k dot, beta/gate are
// indexed flat [G, H, T, S]. Writes the attention output then the
// final state behind it, or one state plane per token with checkpoints.
let gdn = {
workgroup_size: 128
// A pipeline constant: ordinary decode keeps only the final store.
checkpoints: 0
q: shader.buffer(f32)
k: shader.buffer(f32)
v: shader.buffer(f32)
g: shader.buffer(f32)
beta: shader.buffer(f32)
state: shader.buffer_rw(f32)
dst: shader.buffer_rw(f32)
sv: shader.uniform(u32)
h: shader.uniform(u32)
n_tokens: shader.uniform(u32)
n_seqs: shader.uniform(u32)
gate_g: shader.uniform(u32)
q_heads: shader.uniform(u32)
k_heads: shader.uniform(u32)
q_h: shader.uniform(u32)
q_t: shader.uniform(u32)
q_s: shader.uniform(u32)
k_h: shader.uniform(u32)
k_t: shader.uniform(u32)
k_s: shader.uniform(u32)
v_h: shader.uniform(u32)
v_t: shader.uniform(u32)
v_s: shader.uniform(u32)
main: fn() {
let lane = local_id().x % 32
let col = group_id().x * 4 + local_id().x / 32
let head = group_id().y
let seq = group_id().z
if col < self.sv {
let per = self.sv / 32
let iq = head % self.q_heads
let ik = head % self.k_heads
let scale = inverseSqrt(float(self.sv))
let sbase = self.state_base(seq, head, col, lane, per)
var ls = array(0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0)
for j in 0..per {
ls[j] = self.state[sbase + j]
}
var qo = self.q_start(seq, iq)
var ko = self.k_start(seq, ik)
var vo = self.v_start(seq, head)
var bo = seq * self.n_tokens * self.h + head
var go = (seq * self.n_tokens * self.h + head) * self.gate_g
for t in 0..self.n_tokens {
var s_k = 0.0
for j in 0..per {
let si = lane * per + j
var decay = 0.0
if self.gate_g == 1 {
decay = exp(self.g[go])
} else {
decay = exp(self.g[go + si])
}
ls[j] = ls[j] * decay
s_k += ls[j] * self.k[ko + si]
}
s_k = self.lane_sum(s_k)
let d = (self.v[vo + col] - s_k) * self.beta[bo]
var y = 0.0
for j in 0..per {
let si = lane * per + j
ls[j] = ls[j] + self.k[ko + si] * d
y += ls[j] * self.q[qo + si]
}
y = self.lane_sum(y)
if lane == 0 {
self.dst[((seq * self.n_tokens + t) * self.h + head) * self.sv + col] = y * scale
}
if self.checkpoints == 1 {
let out = self.n_seqs * self.n_tokens * self.h * self.sv
+ t * self.n_seqs * self.h * self.sv * self.sv + sbase
for j in 0..per {
self.dst[out + j] = ls[j]
}
}
qo += self.q_t
ko += self.k_t
vo += self.v_t
bo += self.h
go += self.h * self.gate_g
}
if self.checkpoints == 0 {
for j in 0..per {
self.write_state(sbase, j, ls[j])
}
}
}
}
q_start: fn(seq, head) { return seq * self.q_s + head * self.q_h }
k_start: fn(seq, head) { return seq * self.k_s + head * self.k_h }
v_start: fn(seq, head) { return seq * self.v_s + head * self.v_h }
state_base: fn(seq, head, col, lane, per) {
return ((seq * self.h + head) * self.sv + col) * self.sv + lane * per
}
write_state: fn(base, j, v) {
let out = self.n_seqs * self.n_tokens * self.h * self.sv + base
self.dst[out + j] = v
}
// The sum over the 32 lanes of this lane's group (any subgroup size
// >= 32: the butterfly stays inside the 32-aligned group).
lane_sum: fn(v) {
var s = v
s += subgroup_shuffle_xor(s, 16)
s += subgroup_shuffle_xor(s, 8)
s += subgroup_shuffle_xor(s, 4)
s += subgroup_shuffle_xor(s, 2)
s += subgroup_shuffle_xor(s, 1)
return s
}
}
// The same for sv == 128 (four state values per lane) with every q/k
// offset and stride a multiple of 4: each lane's state slice lives in
// one vec4 register and q/k/state move as 16-byte loads (a dynamically
// indexed local array goes to memory on some GPUs).
let gdn_v4 = gdn{
q4: shader.buffer(vec4f)
k4: shader.buffer(vec4f)
state4: shader.buffer_rw(vec4f)
dst4: shader.buffer_rw(vec4f)
main: fn() {
let lane = local_id().x % 32
let col = group_id().x * 4 + local_id().x / 32
let head = group_id().y
let seq = group_id().z
let iq = head % self.q_heads
let ik = head % self.k_heads
let scale = inverseSqrt(float(self.sv))
let si = self.state_base(seq, head, col, lane, u32(4)) / 4
var ls = self.state4[si]
var qo = self.q_start(seq, iq) / 4 + lane
var ko = self.k_start(seq, ik) / 4 + lane
var vo = self.v_start(seq, head)
var bo = seq * self.n_tokens * self.h + head
var go = (seq * self.n_tokens * self.h + head) * self.gate_g
for t in 0..self.n_tokens {
if self.gate_g == 1 {
ls = ls * exp(self.g[go])
} else {
let gi = go + lane * 4
ls = ls * vec4(exp(self.g[gi]), exp(self.g[gi + 1]), exp(self.g[gi + 2]), exp(self.g[gi + 3]))
}
let kv = self.k4[ko]
let s_k = self.lane_sum(dot(ls, kv))
let d = (self.v[vo + col] - s_k) * self.beta[bo]
ls = ls + kv * d
let y = self.lane_sum(dot(ls, self.q4[qo]))
if lane == 0 {
self.dst[((seq * self.n_tokens + t) * self.h + head) * self.sv + col] = y * scale
}
if self.checkpoints == 1 {
let out = (self.n_seqs * self.n_tokens * self.h * self.sv) / 4
+ t * self.n_seqs * self.h * self.sv * 32
+ ((seq * self.h + head) * self.sv + col) * 32 + lane
self.dst4[out] = ls
}
qo += self.q_t / 4
ko += self.k_t / 4
vo += self.v_t
bo += self.h
go += self.h * self.gate_g
}
if self.checkpoints == 0 {
self.write_state4(si, ls)
}
}
write_state4: fn(base, v) {
let out = (self.n_seqs * self.n_tokens * self.h * self.sv) / 4 + base
self.dst4[out] = v
}
}
// Every invocation owns its state slice, so the selected row can be read
// and written in place. Single-sequence address functions also keep the
// push constants within the 128-byte limit without extra parameter buffers.
let gdn_state = gdn{
rows: shader.buffer(i32)
state_row: shader.uniform(u32)
q_start: fn(seq, head) { return head * self.q_h }
k_start: fn(seq, head) { return head * self.k_h }
v_start: fn(seq, head) { return head * self.v_h }
state_base: fn(seq, head, col, lane, per) {
return u32(self.rows[0]) * self.state_row + (head * self.sv + col) * self.sv + lane * per
}
write_state: fn(base, j, v) { self.state[base + j] = v }
}
let gdn_v4_state = gdn_v4{
rows: shader.buffer(i32)
state_row: shader.uniform(u32)
q_start: gdn_state.q_start
k_start: gdn_state.k_start
v_start: gdn_state.v_start
state_base: gdn_state.state_base
write_state4: fn(base, v) { self.state4[base] = v }
}
{
// ---- copies --------------------------------------------------------
// dst = src, element by element in src order (ggml cpy/cont/dup): the
// flat index walks src's shape and lands at the same flat index of
// dst's shape. 4D, strided on both sides.
copy_f32: {
workgroup_size: 256
src: shader.buffer(f32)
dst: shader.buffer_rw(f32)
n: shader.uniform(u32)
groups_x: shader.uniform(u32)
s_ne0: shader.uniform(u32)
s_ne1: shader.uniform(u32)
s_ne2: shader.uniform(u32)
s_nb0: shader.uniform(u32)
s_nb1: shader.uniform(u32)
s_nb2: shader.uniform(u32)
s_nb3: shader.uniform(u32)
d_ne0: shader.uniform(u32)
d_ne1: shader.uniform(u32)
d_ne2: shader.uniform(u32)
d_nb0: shader.uniform(u32)
d_nb1: shader.uniform(u32)
d_nb2: shader.uniform(u32)
d_nb3: shader.uniform(u32)
main: fn() {
let g = group_id()
let i = (g.y * self.groups_x + g.x) * self.workgroup_size + local_id().x
if i < self.n {
let s0 = i % self.s_ne0
let s1 = (i / self.s_ne0) % self.s_ne1
let s2 = (i / (self.s_ne0 * self.s_ne1)) % self.s_ne2
let s3 = i / (self.s_ne0 * self.s_ne1 * self.s_ne2)
let d0 = i % self.d_ne0
let d1 = (i / self.d_ne0) % self.d_ne1
let d2 = (i / (self.d_ne0 * self.d_ne1)) % self.d_ne2
let d3 = i / (self.d_ne0 * self.d_ne1 * self.d_ne2)
self.dst[d0 * self.d_nb0 + d1 * self.d_nb1 + d2 * self.d_nb2 + d3 * self.d_nb3] =
self.src[s0 * self.s_nb0 + s1 * self.s_nb1 + s2 * self.s_nb2 + s3 * self.s_nb3]
}
}
}
// ---- rows ----------------------------------------------------------
// dst[r, :] = src[rows[r], :] (f32 rows).
get_rows_f32: {
workgroup_size: 256
src: shader.buffer(f32)
rows: shader.buffer(i32)
dst: shader.buffer_rw(f32)
ne0: shader.uniform(u32)
src_row: shader.uniform(u32)
dst_row: shader.uniform(u32)
main: fn() {
let i = global_id().x
let r = group_id().y
if i < self.ne0 {
let row = u32(self.rows[r])
self.dst[r * self.dst_row + i] = self.src[row * self.src_row + i]
}
}
}
// set_rows: dst[rows[r], :] = src[r, :], f32 rows into f32 (or f16 in
// the variant below).
set_rows_f32: {
workgroup_size: 256
src: shader.buffer(f32)
rows: shader.buffer(i32)
dst: shader.buffer_rw(f32)
ne0: shader.uniform(u32)
src_row: shader.uniform(u32)
dst_row: shader.uniform(u32)
main: fn() {
let i = global_id().x
let r = group_id().y
if i < self.ne0 {
let row = u32(self.rows[r])
self.dst[row * self.dst_row + i] = self.src[r * self.src_row + i]
}
}
}
set_rows_f16: {
workgroup_size: 256
src: shader.buffer(f32)
rows: shader.buffer(i32)
dst: shader.buffer_rw(f16)
ne0: shader.uniform(u32)
src_row: shader.uniform(u32)
dst_row: shader.uniform(u32)
main: fn() {
let i = global_id().x
let r = group_id().y
if i < self.ne0 {
let row = u32(self.rows[r])
self.dst[row * self.dst_row + i] = f16(self.src[r * self.src_row + i])
}
}
}
// ---- norms ---------------------------------------------------------
// RMS norm (l2 == 0): x / sqrt(mean(x^2) + eps); L2 norm (l2 == 1):
// x / sqrt(max(sum(x^2), eps)). One workgroup per row; rows along
// dim 0, the higher dims strided.
norm: {
workgroup_size: 256
x: shader.buffer(f32)
dst: shader.buffer_rw(f32)
// With w_mode 1 the result is also multiplied by the row weight w
// (a fused RmsNorm + Mul).
w: shader.buffer(f32)
w_mode: shader.uniform(u32)
ne0: shader.uniform(u32)
ne1: shader.uniform(u32)
ne2: shader.uniform(u32)
x_nb1: shader.uniform(u32)
x_nb2: shader.uniform(u32)
x_nb3: shader.uniform(u32)
d_nb1: shader.uniform(u32)
d_nb2: shader.uniform(u32)
d_nb3: shader.uniform(u32)
eps: shader.uniform(f32)
l2: shader.uniform(u32)
partial: shader.workgroup(f32, 32)
main: fn() {
let b = group_id().x
let i1 = b % self.ne1
let i2 = (b / self.ne1) % self.ne2
let i3 = b / (self.ne1 * self.ne2)
let xo = i1 * self.x_nb1 + i2 * self.x_nb2 + i3 * self.x_nb3
let dout = i1 * self.d_nb1 + i2 * self.d_nb2 + i3 * self.d_nb3
let t = local_id().x
var sum = 0.0
var c = t
while c < self.ne0 {
let v = self.x[xo + c]
sum += v * v
c += self.workgroup_size
}
sum = self.block_sum(sum)
var scale = 0.0
if self.l2 == 1 {
scale = 1.0 / sqrt(max(sum, self.eps))
} else {
scale = inverseSqrt(sum / float(self.ne0) + self.eps)
}
c = t
while c < self.ne0 {
var v = self.x[xo + c] * scale
if self.w_mode == 1 {
v = v * self.w[c]
}
self.dst[dout + c] = v
c += self.workgroup_size
}
}
// The workgroup's total of `v`, in every invocation.
block_sum: fn(v) {
let s = subgroup_add(v)
if subgroup_lane() == 0 {
self.partial[subgroup_id()] = s
}
barrier()
var total = 0.0
for k in 0..self.workgroup_size / subgroup_size() {
total += self.partial[k]
}
barrier()
return total
}
}
// ---- rope ----------------------------------------------------------
// Multi-section rope (MROPE / IMROPE): split-half pairs (x[ic],
// x[ic + n_dims/2]) rotated by a theta whose position comes from one of
// four position streams by section. One workgroup per (i1, i2, i3).
rope_multi: {
workgroup_size: 64
src: shader.buffer(f32)
pos: shader.buffer(i32)
dst: shader.buffer_rw(f32)
ne0: shader.uniform(u32)
ne2: shader.uniform(u32)
imrope: shader.uniform(u32)
n_dims: shader.uniform(u32)
sect_0: shader.uniform(u32)
sect_1: shader.uniform(u32)
sect_2: shader.uniform(u32)
sect_3: shader.uniform(u32)
freq_base: shader.uniform(f32)
freq_scale: shader.uniform(f32)
ext_factor: shader.uniform(f32)
attn_factor: shader.uniform(f32)
corr0: shader.uniform(f32)
corr1: shader.uniform(f32)
s_nb0: shader.uniform(u32)
s_nb1: shader.uniform(u32)
s_nb2: shader.uniform(u32)
s_nb3: shader.uniform(u32)
d_nb0: shader.uniform(u32)
d_nb1: shader.uniform(u32)
d_nb2: shader.uniform(u32)
d_nb3: shader.uniform(u32)
main: fn() {
let g = group_id()
let sbase = g.x * self.s_nb1 + g.y * self.s_nb2 + g.z * self.s_nb3
let dbase = g.x * self.d_nb1 + g.y * self.d_nb2 + g.z * self.d_nb3
let i2 = g.y
let half = self.n_dims / 2
var i0 = 2 * local_id().x
while i0 < self.ne0 {
if i0 < self.n_dims {
let ic = i0 / 2
let sect_dims = self.sect_0 + self.sect_1 + self.sect_2 + self.sect_3
let sector = ic % sect_dims
var stream = 3
if self.imrope == 1 {
if sector % 3 == 1 && sector < 3 * self.sect_1 {
stream = 1
} else if sector % 3 == 2 && sector < 3 * self.sect_2 {
stream = 2
} else if sector % 3 == 0 && sector < 3 * self.sect_0 {
stream = 0
}
} else {
if sector < self.sect_0 {
stream = 0
} else if sector < self.sect_0 + self.sect_1 {
stream = 1
} else if sector < self.sect_0 + self.sect_1 + self.sect_2 {
stream = 2
}
}
let theta_base = float(self.pos[i2 + self.ne2 * u32(stream)])
let theta_extrap = theta_base * pow(self.freq_base, -float(i0) / float(self.n_dims))
var mscale = self.attn_factor
var theta = self.freq_scale * theta_extrap
if self.ext_factor != 0.0 {
let y = (float(i0 / 2) - self.corr0) / max(0.001, self.corr1 - self.corr0)
let ramp = (1.0 - min(1.0, max(0.0, y))) * self.ext_factor
theta = theta * (1.0 - ramp) + theta_extrap * ramp
mscale *= 1.0 + 0.1 * log(1.0 / self.freq_scale)
}
let cos_t = cos(theta) * mscale
let sin_t = sin(theta) * mscale
let x0 = self.src[sbase + ic * self.s_nb0]
let x1 = self.src[sbase + (ic + half) * self.s_nb0]
self.dst[dbase + ic * self.d_nb0] = x0 * cos_t - x1 * sin_t
self.dst[dbase + (ic + half) * self.d_nb0] = x0 * sin_t + x1 * cos_t
} else {
self.dst[dbase + i0 * self.d_nb0] = self.src[sbase + i0 * self.s_nb0]
self.dst[dbase + (i0 + 1) * self.d_nb0] = self.src[sbase + (i0 + 1) * self.s_nb0]
}
i0 += 2 * self.workgroup_size
}
}
}
// ---- elementwise ---------------------------------------------------
// ggml unary ops by op code (abs 0 ... softplus 15), contiguous.
unary: {
workgroup_size: 256
x: shader.buffer(f32)
dst: shader.buffer_rw(f32)
// With gated 1 the result is multiplied by g (a fused Unary + Mul).
g: shader.buffer(f32)
n: shader.uniform(u32)
groups_x: shader.uniform(u32)
op: shader.uniform(u32)
gated: shader.uniform(u32)
g_width: shader.uniform(u32)
// Optional row-broadcast add before the activation.
bias: shader.buffer(f32)
add_width: shader.uniform(u32)
main: fn() {
let gr = group_id()
let i = (gr.y * self.groups_x + gr.x) * self.workgroup_size + local_id().x
if i < self.n {
var x = self.x[i]
if self.add_width > 0 {
x = x + self.bias[i % self.add_width]
}
var v = self.apply(self.op, x)
if self.gated == 1 {
var gi = i
if self.g_width > 0 {
gi = i % self.g_width
}
v = v * self.g[gi]
}
self.dst[i] = v
}
}
apply: fn(op, x) {
if op == 0 { return abs(x) }
if op == 1 { return sign(x) }
if op == 2 { return -x }
if op == 3 { return step(0.0, x) * min(1.0, abs(sign(x))) }
if op == 4 { return tanh(x) }
if op == 5 {
if x > 0.0 { return x }
return exp(x) - 1.0
}
if op == 6 { return max(x, 0.0) }
if op == 7 { return 1.0 / (1.0 + exp(-x)) }
if op == 8 { return 0.5 * x * (1.0 + tanh(0.7978845608 * (x + 0.044715 * x * x * x))) }
if op == 9 { return x / (1.0 + exp(-1.702 * x)) }
if op == 10 { return x / (1.0 + exp(-x)) }
if op == 11 { return x * min(1.0, max(0.0, (x + 3.0) / 6.0)) }
if op == 12 { return min(1.0, max(0.0, (x + 3.0) / 6.0)) }
if op == 13 { return exp(x) }
if op == 14 { return exp(x) - 1.0 }
if op == 15 {
if x > 20.0 { return x }
return log(1.0 + exp(x))
}
return x
}
}
// dst = act(a) * b (ggml glu with a separate gate tensor), contiguous.
glu: {
workgroup_size: 256
a: shader.buffer(f32)
b: shader.buffer(f32)
dst: shader.buffer_rw(f32)
n: shader.uniform(u32)
groups_x: shader.uniform(u32)
op: shader.uniform(u32)
main: fn() {
let g = group_id()
let i = (g.y * self.groups_x + g.x) * self.workgroup_size + local_id().x
if i < self.n {
let x = self.a[i]
var act = x / (1.0 + exp(-x))
if self.op == 0 {
act = max(x, 0.0)
} else if self.op == 1 {
act = 0.5 * x * (1.0 + tanh(0.7978845608 * (x + 0.044715 * x * x * x)))
} else if self.op == 5 {
act = x / (1.0 + exp(-1.702 * x))
}
self.dst[i] = act * self.b[i]
}
}
}
// dst = a op b with ggml's modulo broadcast of b (add 0, sub 1, mul 2,
// div 3). One invocation per dst element of dim 0, rows on y, the
// higher dims flattened on z.
binary: {
workgroup_size: 256
a: shader.buffer(f32)
b: shader.buffer(f32)
dst: shader.buffer_rw(f32)
op: shader.uniform(u32)
ne0: shader.uniform(u32)
ne1: shader.uniform(u32)
ne2: shader.uniform(u32)
b_ne0: shader.uniform(u32)
b_ne1: shader.uniform(u32)
b_ne2: shader.uniform(u32)
b_ne3: shader.uniform(u32)
a_nb0: shader.uniform(u32)
a_nb1: shader.uniform(u32)
a_nb2: shader.uniform(u32)
a_nb3: shader.uniform(u32)
b_nb0: shader.uniform(u32)
b_nb1: shader.uniform(u32)
b_nb2: shader.uniform(u32)
b_nb3: shader.uniform(u32)
d_nb0: shader.uniform(u32)
d_nb1: shader.uniform(u32)
d_nb2: shader.uniform(u32)
d_nb3: shader.uniform(u32)
main: fn() {
let i0 = global_id().x
let i1 = group_id().y
let i23 = group_id().z
if i0 < self.ne0 {
let i2 = i23 % self.ne2
let i3 = i23 / self.ne2
let av = self.a[i0 * self.a_nb0 + i1 * self.a_nb1 + i2 * self.a_nb2 + i3 * self.a_nb3]
let bv = self.b[(i0 % self.b_ne0) * self.b_nb0 + (i1 % self.b_ne1) * self.b_nb1
+ (i2 % self.b_ne2) * self.b_nb2 + (i3 % self.b_ne3) * self.b_nb3]
var r = av + bv
if self.op == 1 {
r = av - bv
} else if self.op == 2 {
r = av * bv
} else if self.op == 3 {
r = av / bv
}
self.dst[i0 * self.d_nb0 + i1 * self.d_nb1 + i2 * self.d_nb2 + i3 * self.d_nb3] = r
}
}
}
// ---- recurrent (qwen35 Gated DeltaNet layers) ------------------------
// Gather + concat + the final state window + scatter. One invocation
// owns a channel; it reads the old prefix in order before overwriting
// earlier entries of that same channel. The full convolution input is
// still materialized for the convolution, for decode and prefill alike.
conv_state: {
workgroup_size: 128
state: shader.buffer_rw(f32)
rows: shader.buffer(i32)
x: shader.buffer(f32)
dst: shader.buffer_rw(f32)
prefix: shader.uniform(u32)
channels: shader.uniform(u32)
tokens: shader.uniform(u32)
state_row: shader.uniform(u32)
x_t: shader.uniform(u32)
x_c: shader.uniform(u32)
main: fn() {
let c = global_id().x
if c < self.channels {
let sb = u32(self.rows[0]) * self.state_row + c * self.prefix
let db = c * (self.prefix + self.tokens)
for k in 0..self.prefix {
let v = self.state[sb + k]
self.dst[db + k] = v
if k >= self.tokens {
self.state[sb + k - self.tokens] = v
}
}
for t in 0..self.tokens {
let v = self.x[c * self.x_c + t * self.x_t]
let k = self.prefix + t
self.dst[db + k] = v
if k >= self.tokens {
self.state[sb + k - self.tokens] = v
}
}
}
}
}
// Short causal convolution (ggml ssm_conv): src [d_conv - 1 + T, C, S]
// with time along dim 0, w [d_conv, C], dst [C, T, S];
// dst[c, t, s] = sum_k src[t + k, c, s] * w[k, c]. One invocation per
// (channel, token, seq).
ssm_conv: {
workgroup_size: 128
src: shader.buffer(f32)
w: shader.buffer(f32)
dst: shader.buffer_rw(f32)
d_conv: shader.uniform(u32)
d_inner: shader.uniform(u32)
silu: shader.uniform(u32)
s_nb0: shader.uniform(u32)
s_nb1: shader.uniform(u32)
s_nb2: shader.uniform(u32)
w_nb1: shader.uniform(u32)
d_nb0: shader.uniform(u32)
d_nb1: shader.uniform(u32)
d_nb2: shader.uniform(u32)
main: fn() {
let c = global_id().x
let t = group_id().y
let s = group_id().z
if c < self.d_inner {
let base = s * self.s_nb2 + c * self.s_nb1
var sum = 0.0
for k in 0..self.d_conv {
sum += self.src[base + (t + k) * self.s_nb0] * self.w[c * self.w_nb1 + k]
}
if self.silu == 1 {
sum = sum / (1.0 + exp(-sum))
}
self.dst[s * self.d_nb2 + t * self.d_nb1 + c * self.d_nb0] = sum
}
}
}
// Gated delta rule (see `gdn`).
gated_delta_net: gdn{}
gated_delta_net_checkpoints: gdn{checkpoints: 1}
gated_delta_net_v4: gdn_v4{}
gated_delta_net_v4_checkpoints: gdn_v4{checkpoints: 1}
gated_delta_net_state: gdn_state{}
gated_delta_net_v4_state: gdn_v4_state{}
// ---- attention -----------------------------------------------------
// Attention for a few queries over the f16 KV cache: one workgroup per
// (head, query). Scores for every key in parallel into shared memory
// (in chunks of `chunk` keys), an online softmax across the chunks, and
// the output accumulated one dimension per invocation.
// q f32 [D, nq, H], k/v f16 [D, kc, Hkv], mask f16 [kc, nq], dst f32
// [Dv, H, nq].
attn_decode: {
workgroup_size: 128
chunk: 2048
q: shader.buffer(f32)
k: shader.buffer(f16)
v: shader.buffer(f16)
mask: shader.buffer(f16)
dst: shader.buffer_rw(f32)
d: shader.uniform(u32)
dv: shader.uniform(u32)
kc: shader.uniform(u32)
h: shader.uniform(u32)
hkv: shader.uniform(u32)
scale: shader.uniform(f32)
q_nb1: shader.uniform(u32)
q_nb2: shader.uniform(u32)
k_nb1: shader.uniform(u32)
k_nb2: shader.uniform(u32)
v_nb1: shader.uniform(u32)
v_nb2: shader.uniform(u32)
m_nb1: shader.uniform(u32)
d_nb1: shader.uniform(u32)
d_nb2: shader.uniform(u32)
qs: shader.workgroup(f32, 256)
scores: shader.workgroup(f32, 2048)
partial: shader.workgroup(f32, 32)
main: fn() {
let head = group_id().x
let iq = group_id().y
let t = local_id().x
let kvh = head / (self.h / self.hkv)
let qo = head * self.q_nb2 + iq * self.q_nb1
var i = t
while i < self.d {
self.qs[i] = self.q[qo + i]
i += self.workgroup_size
}
barrier()
// Each invocation keeps two output dimensions (Dv <= 256).
var acc0 = 0.0
var acc1 = 0.0
var m_run = -3.0e38
var l_run = 0.0
var base = u32(0)
while base < self.kc {
let n = min(self.chunk, self.kc - base)
var mx = -3.0e38
var j = t
while j < n {
let key = base + j
let mval = float(self.mask[iq * self.m_nb1 + key])
var s = -3.0e38
if mval > -65500.0 {
let ko = kvh * self.k_nb2 + key * self.k_nb1
var dot = 0.0
for c in 0..self.d {
dot += self.qs[c] * float(self.k[ko + c])
}
s = dot * self.scale + mval
}
self.scores[j] = s
mx = max(mx, s)
j += self.workgroup_size
}
mx = self.block_max(mx)
let m_new = max(m_run, mx)
let corr = exp(m_run - m_new)
var sum = 0.0
j = t
while j < n {
let p = exp(self.scores[j] - m_new)
self.scores[j] = p
sum += p
j += self.workgroup_size
}
sum = self.block_sum(sum)
barrier()
acc0 = acc0 * corr
acc1 = acc1 * corr
let c0 = t
let c1 = t + self.workgroup_size
for jj in 0..n {
let p = self.scores[jj]
let vo = kvh * self.v_nb2 + (base + jj) * self.v_nb1
if c0 < self.dv {
acc0 += p * float(self.v[vo + c0])
}
if c1 < self.dv {
acc1 += p * float(self.v[vo + c1])
}
}
l_run = l_run * corr + sum
m_run = m_new
barrier()
base += self.chunk
}
var inv = 0.0
if l_run > 0.0 {
inv = 1.0 / l_run
}
let out = iq * self.d_nb2 + head * self.d_nb1
if t < self.dv {
self.dst[out + t] = acc0 * inv
}
if t + self.workgroup_size < self.dv {
self.dst[out + t + self.workgroup_size] = acc1 * inv
}
}
block_sum: fn(v) {
let s = subgroup_add(v)
if subgroup_lane() == 0 {
self.partial[subgroup_id()] = s
}
barrier()
var total = 0.0
for k in 0..self.workgroup_size / subgroup_size() {
total += self.partial[k]
}
barrier()
return total
}
block_max: fn(v) {
let s = subgroup_max(v)
if subgroup_lane() == 0 {
self.partial[subgroup_id()] = s
}
barrier()
var total = -3.0e38
for k in 0..self.workgroup_size / subgroup_size() {
total = max(total, self.partial[k])
}
barrier()
return total
}
}
// ---- quantised matrix x vector (see `gemv` above) -------------------
// Activations to biased int8 per 32 (+ d and sum q per block): one
// invocation per block.
quantize_q8: {
workgroup_size: 64
x: shader.buffer(vec4f)
xq: shader.buffer_rw(u32)
xs: shader.buffer_rw(vec2f)
blocks: shader.uniform(u32)
k: shader.uniform(u32)
x_col: shader.uniform(u32)
main: fn() {
let b = global_id().x
if b < self.blocks {
let per_col = self.k / 32
let col = b / per_col
let at = (col * self.x_col + (b % per_col) * 32) / 4
var amax = 0.0
for i in 0..8 {
let v = abs(self.x[at + i])
amax = max(amax, max(max(v.x, v.y), max(v.z, v.w)))
}
let d = amax / 127.0
var inv = 0.0
if amax > 0.0 {
inv = 1.0 / d
}
var sum = 0
for i in 0..8 {
let v = self.x[at + i] * inv
let q = sign(v) * floor(abs(v) + 0.5)
sum += i32(q.x + q.y + q.z + q.w)
let u = q + 128.0
self.xq[b * 8 + i] = u32(u.x) | (u32(u.y) << 8) | (u32(u.z) << 16) | (u32(u.w) << 24)
}
self.xs[b] = vec2(d, float(sum))
}
}
}
gemv_q5k: gemv_k{
unit_bytes: 176
qs_at: 48
high_bit: 1
}
gemv_q4k: gemv_k{
unit_bytes: 144
qs_at: 16
high_bit: 0
}
// q6_K, repacked at upload into 224-byte blocks (the 210 bytes, then
// padding) so every part is 16-aligned: ql at 0, qh at 128, the int8
// scales at 192, f16 d at 208. Eight lanes per unit; lane p takes 16
// consecutive l of half n = p / 4: from one 16-byte ql load the low
// nibbles (group g = (p / 2) % 2) and high nibbles (group g + 2), the
// two high bits each from one 16-byte qh load. Codes are q + 32.
gemv_q6k: gemv_k{
unit_bytes: 224
main: fn() {
let lane = local_id().x % 32
let row = group_id().x * (self.workgroup_size / 32) + local_id().x / 32
if row < self.n {
let units = self.k / 256
let slot = lane >> 3
let part = lane & 7
let half = part >> 2
let g = (part >> 1) & 1
let l = (part & 1) * 16
for col in 0..self.m {
var acc = 0.0
var u = slot
while u < units {
let ub = col * self.k + u * 256 + half * 128
let x_lo = self.xq4[(ub + g * 32 + l) >> 4]
let x_hi = self.xq4[(ub + (g + 2) * 32 + l) >> 4]
let xd_lo = self.xs[(ub + g * 32) / 32].x
let xd_hi = self.xs[(ub + (g + 2) * 32) / 32].x
acc += self.unit_part(row * self.w_row + u * 224, part, x_lo, x_hi, vec4(xd_lo, 0.0, xd_hi, 0.0))
u += 4
}
let total = self.lane_sum(acc)
if lane == 0 {
self.dst[col * self.d_col + row] = total
}
}
}
}
unit_part: fn(bo, part, x_lo, x_hi, xd) {
let b16 = bo >> 4
let half = part >> 2
let g = (part >> 1) & 1
let ql = self.w4[b16 + half * 4 + ((part >> 1) & 1) * 2 + (part & 1)]
let qh = self.w4[b16 + 8 + half * 2 + (part & 1)]
let tail = self.w4[b16 + 12]
let d = unpack_half2(self.w4[b16 + 13].x & u32(65535)).x
let m4 = vec4u(u32(252645135), u32(252645135), u32(252645135), u32(252645135))
let m2 = vec4u(u32(50529027), u32(50529027), u32(50529027), u32(50529027))
let sh = g * 2
let lo = (ql & m4) | (((qh >> vec4u(sh, sh, sh, sh)) & m2) << vec4u(4, 4, 4, 4))
let hi = ((ql >> vec4u(4, 4, 4, 4)) & m4) | (((qh >> vec4u(sh + 4, sh + 4, sh + 4, sh + 4)) & m2) << vec4u(4, 4, 4, 4))
let sub = part & 1
let sc_lo = self.sbyte(tail, half * 8 + sub + 2 * g)
let sc_hi = self.sbyte(tail, half * 8 + sub + 2 * (g + 2))
let v_lo = self.dot16(lo, x_lo) - 32 * self.sum16(x_lo)
let v_hi = self.dot16(hi, x_hi) - 32 * self.sum16(x_hi)
return d * (sc_lo * xd.x * float(v_lo) + sc_hi * xd.z * float(v_hi))
}
// Signed byte `j` of a 16-byte value, as a float.
sbyte: fn(v, j) {
var w = v.w
if j < 4 {
w = v.x
} else if j < 8 {
w = v.y
} else if j < 12 {
w = v.z
}
let b = (w >> ((j & 3) * 8)) & 255
if b > 127 {
return float(b) - 256.0
}
return float(b)
}
}
// q8_0, repacked at upload: per 8 blocks (one 256-value unit) the eight
// f16 d (16 bytes), then the eight blocks' 32 int8 each. Lane p of the
// unit's eight takes block p: two 16-byte loads; sign bits flipped make
// the weights biased unsigned codes too.
gemv_q80: gemv_k{
unit_bytes: 272
main: fn() {
let lane = local_id().x % 32
let row = group_id().x * (self.workgroup_size / 32) + local_id().x / 32
if row < self.n {
let units = self.k / 256
let slot = lane >> 3
let part = lane & 7
for col in 0..self.m {
var acc = 0.0
var u = slot
while u < units {
let xb = col * self.k + u * 256 + part * 32
let x_a = self.xq4[xb >> 4]
let x_b = self.xq4[(xb >> 4) + 1]
let xd = self.xs[xb / 32].x
acc += self.unit_part(row * self.w_row + u * 272, part, x_a, x_b, vec4(xd, 0.0, 0.0, 0.0))
u += 4
}
let total = self.lane_sum(acc)
if lane == 0 {
self.dst[col * self.d_col + row] = total
}
}
}
}
unit_part: fn(bo, part, x_a, x_b, xd) {
let b16 = bo >> 4
let ds = self.w4[b16]
var dw = ds.w
if part < 2 {
dw = ds.x
} else if part < 4 {
dw = ds.y
} else if part < 6 {
dw = ds.z
}
let d = unpack_half2(dw).x
var d2 = d
if (part & 1) == 1 {
d2 = unpack_half2(dw).y
}
let flip = vec4u(u32(2155905152), u32(2155905152), u32(2155905152), u32(2155905152))
let wa = self.w4[b16 + 1 + part * 2] ^ flip
let wb = self.w4[b16 + 2 + part * 2] ^ flip
// sum (W - 128) q = sum W q' - 128 sum W - 128 sum q (q' = q + 128)
let ones = u32(16843009)
let raw = self.dot16raw(wa, x_a) + self.dot16raw(wb, x_b)
let wsum = self.dot16raw(wa, vec4u(ones, ones, ones, ones)) + self.dot16raw(wb, vec4u(ones, ones, ones, ones))
let v = i32(raw) - 128 * i32(wsum) - 128 * (self.sum16(x_a) + self.sum16(x_b))
return d2 * xd.x * float(v)
}
dot16raw: fn(a, x) {
return dot4_u8(a.x, x.x) + dot4_u8(a.y, x.y) + dot4_u8(a.z, x.z) + dot4_u8(a.w, x.w)
}
}
// word: 1 row/32 lanes, 256 threads. Per 16 payload bytes: Q4/Q5
// 18 dot4 including x sums, Q6 18, Q8 8; integer counts above.
gemv_q4k_word: gemv_decode{format: 4 unit_bytes: 144}
gemv_q5k_word: gemv_decode{}
gemv_q6k_word: gemv_decode{format: 6 unit_bytes: 224}
gemv_q80_word: gemv_decode{format: 8 unit_bytes: 272}
// rows2: 2 rows/32 lanes. Same weight ops; x loads halved and Q4/Q5/Q6
// amortises to 14 dot4/16 payload bytes (10 weight + 8/2 activation).
gemv_q4k_rows2: gemv_decode{format: 4 unit_bytes: 144 rpg: 2}
gemv_q5k_rows2: gemv_decode{rpg: 2}
gemv_q6k_rows2: gemv_decode{format: 6 unit_bytes: 224 rpg: 2}
gemv_q80_rows2: gemv_decode{format: 8 unit_bytes: 272 rpg: 2}
// rows4: 4 rows/32 lanes. Same weight ops; x loads quartered and Q4/Q5/Q6
// amortises to 12 dot4/16 payload bytes (10 weight + 8/4 activation).
gemv_q4k_rows4: gemv_decode{format: 4 unit_bytes: 144 rpg: 4}
gemv_q5k_rows4: gemv_decode{rpg: 4}
gemv_q6k_rows4: gemv_decode{format: 6 unit_bytes: 224 rpg: 4}
gemv_q80_rows4: gemv_decode{format: 8 unit_bytes: 272 rpg: 4}
// wide: 1 row/32 lanes, 512 threads; word's 18/18/18/8 dot4 counts
// per 16 payload bytes, twice as many independent rows per workgroup.
gemv_q4k_wide: gemv_decode{format: 4 unit_bytes: 144 workgroup_size: 512}
gemv_q5k_wide: gemv_decode{workgroup_size: 512}
gemv_q6k_wide: gemv_decode{format: 6 unit_bytes: 224 workgroup_size: 512}
gemv_q80_wide: gemv_decode{format: 8 unit_bytes: 272 workgroup_size: 512}
// lanes64: word's 18/18/18/8 dot4 per 16 payload bytes, eight units
// in flight per row; one subgroup_add per row at subgroup size 64.
gemv_q4k_lanes64: gemv_decode{format: 4 unit_bytes: 144 row_lanes: 64}
gemv_q5k_lanes64: gemv_decode{row_lanes: 64}
gemv_q6k_lanes64: gemv_decode{format: 6 unit_bytes: 224 row_lanes: 64}
gemv_q80_lanes64: gemv_decode{format: 8 unit_bytes: 272 row_lanes: 64}
// lanes128: word's 18/18/18/8 dot4 per 16 payload bytes, sixteen
// units in flight; one subgroup_add per row at subgroup size 128.
gemv_q4k_lanes128: gemv_decode{format: 4 unit_bytes: 144 row_lanes: 128}
gemv_q5k_lanes128: gemv_decode{row_lanes: 128}
gemv_q6k_lanes128: gemv_decode{format: 6 unit_bytes: 224 row_lanes: 128}
gemv_q80_lanes128: gemv_decode{format: 8 unit_bytes: 272 row_lanes: 128}
// ---- quantised row lookups and prefill matmuls (see `deq`) -------
// 64x64 output tile, 4x8 results per invocation, k steps of 16 (on the
// phone faster than 4x4 results, k 32 or 64).
gemm_q4k_r8k16: gemm_wide_16{format: 4 unit_bytes: 144}
gemm_q5k_r8k16: gemm_wide_16{}
gemm_q6k_r8k16: gemm_wide_16{format: 6 unit_bytes: 224}
gemm_q80_r8k16: gemm_wide_16{format: 8 unit_bytes: 272}
get_rows_q6k: deq_q6k{
workgroup_size: 256
rows: shader.buffer(i32)
dst: shader.buffer_rw(f32)
ne0: shader.uniform(u32)
w_row: shader.uniform(u32)
dst_row: shader.uniform(u32)
main: rows_main
}
gemm_q6k: deq_q6k{
workgroup_size: 128
xq: shader.buffer(u32)
xs: shader.buffer(vec2f)
dst: shader.buffer_rw(f32)
k: shader.uniform(u32)
n: shader.uniform(u32)
m: shader.uniform(u32)
w_row: shader.uniform(u32)
d_col: shader.uniform(u32)
wq: shader.workgroup(u32, 256)
wm: shader.workgroup(vec4f, 32)
xq_s: shader.workgroup(u32, 256)
xd_s: shader.workgroup(f32, 32)
main: gemm_main
}
get_rows_q5k: deq_q5k{
workgroup_size: 256
rows: shader.buffer(i32)
dst: shader.buffer_rw(f32)
ne0: shader.uniform(u32)
w_row: shader.uniform(u32)
dst_row: shader.uniform(u32)
main: rows_main
}
gemm_q5k: deq_q5k{
workgroup_size: 128
xq: shader.buffer(u32)
xs: shader.buffer(vec2f)
dst: shader.buffer_rw(f32)
k: shader.uniform(u32)
n: shader.uniform(u32)
m: shader.uniform(u32)
w_row: shader.uniform(u32)
d_col: shader.uniform(u32)
wq: shader.workgroup(u32, 256)
wm: shader.workgroup(vec4f, 32)
xq_s: shader.workgroup(u32, 256)
xd_s: shader.workgroup(f32, 32)
main: gemm_main
}
get_rows_q4k: deq_q4k{
workgroup_size: 256
rows: shader.buffer(i32)
dst: shader.buffer_rw(f32)
ne0: shader.uniform(u32)
w_row: shader.uniform(u32)
dst_row: shader.uniform(u32)
main: rows_main
}
gemm_q4k: deq_q4k{
workgroup_size: 128
xq: shader.buffer(u32)
xs: shader.buffer(vec2f)
dst: shader.buffer_rw(f32)
k: shader.uniform(u32)
n: shader.uniform(u32)
m: shader.uniform(u32)
w_row: shader.uniform(u32)
d_col: shader.uniform(u32)
wq: shader.workgroup(u32, 256)
wm: shader.workgroup(vec4f, 32)
xq_s: shader.workgroup(u32, 256)
xd_s: shader.workgroup(f32, 32)
main: gemm_main
}
get_rows_q80: deq_q80{
workgroup_size: 256
rows: shader.buffer(i32)
dst: shader.buffer_rw(f32)
ne0: shader.uniform(u32)
w_row: shader.uniform(u32)
dst_row: shader.uniform(u32)
main: rows_main
}
gemm_q80: deq_q80{
workgroup_size: 128
xq: shader.buffer(u32)
xs: shader.buffer(vec2f)
dst: shader.buffer_rw(f32)
k: shader.uniform(u32)
n: shader.uniform(u32)
m: shader.uniform(u32)
w_row: shader.uniform(u32)
d_col: shader.uniform(u32)
wq: shader.workgroup(u32, 256)
wm: shader.workgroup(vec4f, 32)
xq_s: shader.workgroup(u32, 256)
xd_s: shader.workgroup(f32, 32)
main: gemm_main
}
}
Reference in a new issue View git blame Copy permalink