74faf62ef1
The AI hub serves local models on the LAN and now yields GPU memory to the UI (shrinks the chat context, then unloads; grows back after; never loads into the UI reserve), groups models by type, and reports each model's memory split. LLM decode on Vulkan with Splash kernels, one shared GPU device per process, vision on chat turns, an on-device text detector, music models and stems. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
373 lines
13 KiB
Text
373 lines
13 KiB
Text
// Whisper: row-major f32 activations, packed f16 or f32 weights, f32 sums.
|
|
// Weight rows are padded to eight elements. Addresses and strides refer to
|
|
// elements, never descriptors. No kernel assumes a 32-lane subgroup.
|
|
let matrix = {
|
|
workgroup_size: 128
|
|
half_weights: 1
|
|
x: shader.buffer(f32)
|
|
x4: shader.buffer(vec4f)
|
|
w4: shader.buffer(vec4u)
|
|
wf4: shader.buffer(vec4f)
|
|
bias: shader.buffer(f32)
|
|
dst: shader.buffer_rw(f32)
|
|
m: shader.uniform(u32)
|
|
n: shader.uniform(u32)
|
|
k: shader.uniform(u32)
|
|
ws: shader.uniform(u32)
|
|
has_bias: shader.uniform(u32)
|
|
load_x: fn(row, col) {
|
|
if row < self.m && col + 3 < self.k && self.k % 4 == 0 {
|
|
return self.x4[(row * self.k + col) / 4]
|
|
}
|
|
var v = vec4f(0.0)
|
|
if row < self.m {
|
|
for c in 0..4 {
|
|
if col + c < self.k { v[c] = self.x[row * self.k + col + c] }
|
|
}
|
|
}
|
|
return v
|
|
}
|
|
}
|
|
|
|
// A 32x32 output tile and 32-wide reduction tile, 8 KiB shared memory.
|
|
// Each invocation accumulates two rows of four columns. The global reads
|
|
// are vec4f activations and vec4u packed weights (eight f16 values/load).
|
|
let gemm = matrix{
|
|
ax: shader.workgroup(f32, 1024)
|
|
bw: shader.workgroup(f32, 1024)
|
|
main: fn() {
|
|
let t = local_id().x
|
|
let row = t / 8
|
|
let col = (t % 8) * 4
|
|
let mr = group_id().y * 32
|
|
let nr = group_id().x * 32
|
|
var acc0 = vec4f(0.0)
|
|
var acc1 = vec4f(0.0)
|
|
var base = u32(0)
|
|
while base < self.k {
|
|
let a0 = self.load_x(mr + row, base + col)
|
|
let a1 = self.load_x(mr + row + 16, base + col)
|
|
for c in 0..4 {
|
|
self.ax[row * 32 + col + c] = a0[c]
|
|
self.ax[(row + 16) * 32 + col + c] = a1[c]
|
|
}
|
|
let wr = t / 4
|
|
let wc = (t % 4) * 8
|
|
var lo = vec4f(0.0)
|
|
var hi = vec4f(0.0)
|
|
if nr + wr < self.n && base + wc < self.k {
|
|
let off = (nr + wr) * self.ws + base + wc
|
|
if self.half_weights == 1 {
|
|
let p = self.w4[off / 8]
|
|
lo = vec4f(unpack_half2(p.x), unpack_half2(p.y))
|
|
hi = vec4f(unpack_half2(p.z), unpack_half2(p.w))
|
|
} else {
|
|
lo = self.wf4[off / 4]
|
|
hi = self.wf4[off / 4 + 1]
|
|
}
|
|
}
|
|
for c in 0..4 {
|
|
self.bw[(wc + c) * 32 + wr] = lo[c]
|
|
self.bw[(wc + c + 4) * 32 + wr] = hi[c]
|
|
}
|
|
barrier()
|
|
for j in 0..32 {
|
|
let w = vec4f(self.bw[j * 32 + col], self.bw[j * 32 + col + 1],
|
|
self.bw[j * 32 + col + 2], self.bw[j * 32 + col + 3])
|
|
acc0 += self.ax[row * 32 + j] * w
|
|
acc1 += self.ax[(row + 16) * 32 + j] * w
|
|
}
|
|
barrier()
|
|
base += 32
|
|
}
|
|
for c in 0..4 {
|
|
if nr + col + c < self.n {
|
|
var b = 0.0
|
|
if self.has_bias != 0 { b = self.bias[nr + col + c] }
|
|
if mr + row < self.m {
|
|
self.dst[(mr + row) * self.n + nr + col + c] = acc0[c] + b
|
|
}
|
|
if mr + row + 16 < self.m {
|
|
self.dst[(mr + row + 16) * self.n + nr + col + c] = acc1[c] + b
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// Decode: four output rows per workgroup, eight weights per vector load.
|
|
// Reductions span all actual subgroups, including a single subgroup of 128.
|
|
let gemv = matrix{
|
|
partial: shader.workgroup(vec4f, 128)
|
|
main: fn() {
|
|
let t = local_id().x
|
|
let row = group_id().x * 4
|
|
let batch = group_id().y
|
|
var acc = vec4f(0.0)
|
|
var c = t * 8
|
|
while c < self.k {
|
|
let x0 = self.load_x(batch, c)
|
|
let x1 = self.load_x(batch, c + 4)
|
|
for r in 0..4 {
|
|
if row + r < self.n {
|
|
let off = (row + r) * self.ws + c
|
|
var lo = vec4f(0.0)
|
|
var hi = vec4f(0.0)
|
|
if self.half_weights == 1 {
|
|
let p = self.w4[off / 8]
|
|
lo = vec4f(unpack_half2(p.x), unpack_half2(p.y))
|
|
hi = vec4f(unpack_half2(p.z), unpack_half2(p.w))
|
|
} else {
|
|
lo = self.wf4[off / 4]
|
|
hi = self.wf4[off / 4 + 1]
|
|
}
|
|
acc[r] += dot(x0, lo) + dot(x1, hi)
|
|
}
|
|
}
|
|
c += self.workgroup_size * 8
|
|
}
|
|
let total = vec4f(subgroup_add(acc.x), subgroup_add(acc.y),
|
|
subgroup_add(acc.z), subgroup_add(acc.w))
|
|
if subgroup_lane() == 0 { self.partial[subgroup_id()] = total }
|
|
barrier()
|
|
if t == 0 {
|
|
var result = vec4f(0.0)
|
|
for s in 0..(self.workgroup_size + subgroup_size() - 1) / subgroup_size() {
|
|
result += self.partial[s]
|
|
}
|
|
for r in 0..4 {
|
|
if row + r < self.n {
|
|
var b = 0.0
|
|
if self.has_bias != 0 { b = self.bias[row + r] }
|
|
self.dst[batch * self.n + row + r] = result[r] + b
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
{
|
|
whisper_gemm_f16: gemm{}
|
|
whisper_gemm_f32: gemm{half_weights: 0}
|
|
whisper_gemv_f16: gemv{}
|
|
whisper_gemv_f32: gemv{half_weights: 0}
|
|
|
|
// op 0 add, 1 multiply, 2 GELU. b_len supports scalar, row and full
|
|
// elementwise broadcasts; callers validate the corresponding shapes.
|
|
whisper_elementwise: {
|
|
workgroup_size: 128
|
|
a: shader.buffer(f32)
|
|
b: shader.buffer(f32)
|
|
dst: shader.buffer_rw(f32)
|
|
n: shader.uniform(u32)
|
|
b_len: shader.uniform(u32)
|
|
op: shader.uniform(u32)
|
|
main: fn() {
|
|
let i = global_id().x
|
|
if i < self.n {
|
|
let x = self.a[i]
|
|
var out = x
|
|
if self.op == 0 {
|
|
out = x + self.b[i % self.b_len]
|
|
} else if self.op == 1 {
|
|
out = x * self.b[i % self.b_len]
|
|
} else {
|
|
out = 0.5 * x * (1.0 + tanh(0.7978845608 * (x + 0.044715 * x * x * x)))
|
|
}
|
|
self.dst[i] = out
|
|
}
|
|
}
|
|
}
|
|
|
|
// Two-pass variance matches the CPU formula (avoids cancellation in
|
|
// E[x*x] - E[x]*E[x]). Every lane participates in both reductions.
|
|
whisper_norm: {
|
|
workgroup_size: 128
|
|
x: shader.buffer(f32)
|
|
mul: shader.buffer(f32)
|
|
add: shader.buffer(f32)
|
|
dst: shader.buffer_rw(f32)
|
|
cols: shader.uniform(u32)
|
|
affine: shader.uniform(u32)
|
|
eps: shader.uniform(f32)
|
|
partial: shader.workgroup(f32, 128)
|
|
block_sum: fn(v) {
|
|
let s = subgroup_add(v)
|
|
if subgroup_lane() == 0 { self.partial[subgroup_id()] = s }
|
|
barrier()
|
|
var sum = 0.0
|
|
for j in 0..(self.workgroup_size + subgroup_size() - 1) / subgroup_size() {
|
|
sum += self.partial[j]
|
|
}
|
|
barrier()
|
|
return sum
|
|
}
|
|
main: fn() {
|
|
let t = local_id().x
|
|
let row = group_id().x * self.cols
|
|
var sum = 0.0
|
|
var c = t
|
|
while c < self.cols {
|
|
sum += self.x[row + c]
|
|
c += self.workgroup_size
|
|
}
|
|
let mean = self.block_sum(sum) / float(self.cols)
|
|
var variance = 0.0
|
|
c = t
|
|
while c < self.cols {
|
|
let delta = self.x[row + c] - mean
|
|
variance += delta * delta
|
|
c += self.workgroup_size
|
|
}
|
|
let inv = 1.0 / sqrt(self.block_sum(variance) / float(self.cols) + self.eps)
|
|
c = t
|
|
while c < self.cols {
|
|
var value = (self.x[row + c] - mean) * inv
|
|
if self.affine != 0 { value = value * self.mul[c] + self.add[c] }
|
|
self.dst[row + c] = value
|
|
c += self.workgroup_size
|
|
}
|
|
}
|
|
}
|
|
|
|
// Front-end convolution lowered to a GEMM. The first input is
|
|
// channels-first mel; subsequent convolution inputs are time-major.
|
|
whisper_im2col: {
|
|
workgroup_size: 128
|
|
x: shader.buffer(f32)
|
|
dst: shader.buffer_rw(f32)
|
|
ic: shader.uniform(u32)
|
|
iw: shader.uniform(u32)
|
|
ow: shader.uniform(u32)
|
|
kw: shader.uniform(u32)
|
|
stride: shader.uniform(u32)
|
|
pad: shader.uniform(u32)
|
|
time_major: shader.uniform(u32)
|
|
main: fn() {
|
|
let i = global_id().x
|
|
let width = self.ic * self.kw
|
|
if i < self.ow * width {
|
|
let t = i / width
|
|
let ci = (i % width) / self.kw
|
|
let kk = i % self.kw
|
|
let pos = i32(t * self.stride + kk) - i32(self.pad)
|
|
var v = 0.0
|
|
if pos >= 0 && pos < i32(self.iw) {
|
|
if self.time_major != 0 {
|
|
v = self.x[u32(pos) * self.ic + ci]
|
|
} else {
|
|
v = self.x[ci * self.iw + u32(pos)]
|
|
}
|
|
}
|
|
self.dst[i] = v
|
|
}
|
|
}
|
|
}
|
|
|
|
// Online attention: eight queries x 32 keys, 64 channels/head, <20 KiB
|
|
// shared memory. Q/K/V are read with 16-byte loads. Scores never leave
|
|
// the tile, so a 1500-frame encoder has no quadratic allocation.
|
|
// The 64-channel head is shared by every supported Whisper model.
|
|
whisper_attention: {
|
|
workgroup_size: 128
|
|
q: shader.buffer(vec4f)
|
|
k: shader.buffer(vec4f)
|
|
v: shader.buffer(vec4f)
|
|
dst: shader.buffer_rw(vec4f)
|
|
nq: shader.uniform(u32)
|
|
nk: shader.uniform(u32)
|
|
heads: shader.uniform(u32)
|
|
past: shader.uniform(u32)
|
|
causal: shader.uniform(u32)
|
|
scale: shader.uniform(f32)
|
|
qs: shader.workgroup(vec4f, 128)
|
|
ks: shader.workgroup(vec4f, 512)
|
|
vs: shader.workgroup(vec4f, 512)
|
|
scores: shader.workgroup(f32, 256)
|
|
maxima: shader.workgroup(f32, 8)
|
|
sums: shader.workgroup(f32, 8)
|
|
correction: shader.workgroup(f32, 8)
|
|
main: fn() {
|
|
let t = local_id().x
|
|
let head = group_id().x
|
|
let first = group_id().y * 8
|
|
let qi = t / 16
|
|
let dim = t % 16
|
|
let row_stride = self.heads * 16
|
|
var query = vec4f(0.0)
|
|
if first + qi < self.nq {
|
|
query = self.q[(first + qi) * row_stride + head * 16 + dim]
|
|
}
|
|
self.qs[t] = query
|
|
if t < 8 {
|
|
self.maxima[t] = -3.0e38
|
|
self.sums[t] = 0.0
|
|
}
|
|
barrier()
|
|
var acc = vec4f(0.0)
|
|
var base = u32(0)
|
|
while base < self.nk {
|
|
var p = t
|
|
while p < 512 {
|
|
let key = base + p / 16
|
|
let off = key * row_stride + head * 16 + p % 16
|
|
var key_vec = vec4f(0.0)
|
|
var val_vec = vec4f(0.0)
|
|
if key < self.nk {
|
|
key_vec = self.k[off]
|
|
val_vec = self.v[off]
|
|
}
|
|
self.ks[p] = key_vec
|
|
self.vs[p] = val_vec
|
|
p += self.workgroup_size
|
|
}
|
|
barrier()
|
|
// Each invocation owns two scores; no subgroup-size
|
|
// assumption or cross-subgroup shuffle is involved.
|
|
for pair in 0..2 {
|
|
let score_i = t + pair * 128
|
|
let qr = score_i / 32
|
|
let kr = score_i % 32
|
|
var score = -3.0e38
|
|
if first + qr < self.nq && base + kr < self.nk {
|
|
if self.causal == 0 || base + kr <= self.past + first + qr {
|
|
var product = 0.0
|
|
for d4 in 0..16 {
|
|
product += dot(self.qs[qr * 16 + d4], self.ks[kr * 16 + d4])
|
|
}
|
|
score = product * self.scale
|
|
}
|
|
}
|
|
self.scores[score_i] = score
|
|
}
|
|
barrier()
|
|
if dim == 0 {
|
|
var mx = self.maxima[qi]
|
|
for j in 0..32 { mx = max(mx, self.scores[qi * 32 + j]) }
|
|
let corr = exp(self.maxima[qi] - mx)
|
|
var sum = 0.0
|
|
for j in 0..32 {
|
|
let score = self.scores[qi * 32 + j]
|
|
var prob = 0.0
|
|
if score > -3.0e38 { prob = exp(score - mx) }
|
|
self.scores[qi * 32 + j] = prob
|
|
sum += prob
|
|
}
|
|
self.correction[qi] = corr
|
|
self.sums[qi] = self.sums[qi] * corr + sum
|
|
self.maxima[qi] = mx
|
|
}
|
|
barrier()
|
|
acc = acc * self.correction[qi]
|
|
for j in 0..32 {
|
|
acc += self.scores[qi * 32 + j] * self.vs[j * 16 + dim]
|
|
}
|
|
barrier()
|
|
base += 32
|
|
}
|
|
if first + qi < self.nq {
|
|
self.dst[(first + qi) * row_stride + head * 16 + dim] = acc / self.sums[qi]
|
|
}
|
|
}
|
|
}
|
|
}
|