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>
372 lines
14 KiB
Text
372 lines
14 KiB
Text
// Channel-major activations, with each row padded to four floats. The
|
|
// convolution GEMM consumes patches [input * taps, time] and resident
|
|
// weights [output, input * taps]. A 32x32x32 tile uses 8 KiB shared memory.
|
|
// Vulkan guarantees at least 65535 groups per dimension. Split large flat
|
|
// dispatches over two dimensions instead of imposing an activation limit.
|
|
let flat = {
|
|
workgroup_size: 128
|
|
index: fn() { return (group_id().y * 65535 + group_id().x) * 128 + local_id().x }
|
|
}
|
|
|
|
{
|
|
kokoro_gemm: {
|
|
workgroup_size: 128
|
|
a: shader.buffer(vec4f)
|
|
b: shader.buffer(vec4f)
|
|
bias: shader.buffer(f32)
|
|
dst: shader.buffer_rw(vec4f)
|
|
m: shader.uniform(u32)
|
|
n: shader.uniform(u32)
|
|
k: shader.uniform(u32)
|
|
astride: shader.uniform(u32)
|
|
bstride: shader.uniform(u32)
|
|
has_bias: shader.uniform(u32)
|
|
aa: shader.workgroup(f32, 1024)
|
|
bb: shader.workgroup(vec4f, 256)
|
|
main: fn() {
|
|
let t = local_id().x
|
|
let row = t / 8
|
|
let col = t % 8
|
|
let mr = group_id().y * 32
|
|
let nc = group_id().x * 8
|
|
var acc0 = vec4f(0.0)
|
|
var acc1 = vec4f(0.0)
|
|
var base = u32(0)
|
|
while base < self.k {
|
|
for half in 0..2 {
|
|
let r = row + half * 16
|
|
var av = vec4f(0.0)
|
|
var bv = vec4f(0.0)
|
|
if mr + r < self.m && base + col * 4 < self.k {
|
|
av = self.a[(mr + r) * self.astride / 4 + base / 4 + col]
|
|
}
|
|
if base + r < self.k && (nc + col) * 4 < self.n {
|
|
bv = self.b[(base + r) * self.bstride / 4 + nc + col]
|
|
}
|
|
for j in 0..4 { self.aa[r * 32 + col * 4 + j] = av[j] }
|
|
self.bb[r * 8 + col] = bv
|
|
}
|
|
barrier()
|
|
for j in 0..32 {
|
|
let bv = self.bb[j * 8 + col]
|
|
acc0 += self.aa[row * 32 + j] * bv
|
|
acc1 += self.aa[(row + 16) * 32 + j] * bv
|
|
}
|
|
barrier()
|
|
base += 32
|
|
}
|
|
if (nc + col) * 4 < self.n {
|
|
if mr + row < self.m {
|
|
var bv = 0.0
|
|
if self.has_bias != 0 { bv = self.bias[mr + row] }
|
|
self.dst[(mr + row) * self.bstride / 4 + nc + col] = acc0 + bv
|
|
}
|
|
if mr + row + 16 < self.m {
|
|
var bv = 0.0
|
|
if self.has_bias != 0 { bv = self.bias[mr + row + 16] }
|
|
self.dst[(mr + row + 16) * self.bstride / 4 + nc + col] = acc1 + bv
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// Four adjacent time positions per invocation. Stride-one interiors use
|
|
// one aligned vector load when possible, including dilated convolutions.
|
|
kokoro_im2col: flat{
|
|
x: shader.buffer(f32)
|
|
x4: shader.buffer(vec4f)
|
|
dst: shader.buffer_rw(vec4f)
|
|
ic: shader.uniform(u32)
|
|
iw: shader.uniform(u32)
|
|
xs: shader.uniform(u32)
|
|
ow: shader.uniform(u32)
|
|
os: shader.uniform(u32)
|
|
kw: shader.uniform(u32)
|
|
stride: shader.uniform(u32)
|
|
pad: shader.uniform(u32)
|
|
dilation: shader.uniform(u32)
|
|
main: fn() {
|
|
let i = self.index()
|
|
let time = (i % (self.os / 4)) * 4
|
|
let taprow = i / (self.os / 4)
|
|
if taprow < self.ic * self.kw {
|
|
let channel = taprow / self.kw
|
|
let tap = taprow % self.kw
|
|
var v = vec4f(0.0)
|
|
let start = i32(time * self.stride + tap * self.dilation) - i32(self.pad)
|
|
if self.stride == 1 && start >= 0 && start + 3 < i32(self.iw) && start % 4 == 0 && time + 3 < self.ow {
|
|
v = self.x4[(channel * self.xs + u32(start)) / 4]
|
|
} else {
|
|
for j in 0..4 {
|
|
let pos = start + i32(j * self.stride)
|
|
if time + j < self.ow && pos >= 0 && pos < i32(self.iw) {
|
|
v[j] = self.x[channel * self.xs + u32(pos)]
|
|
}
|
|
}
|
|
}
|
|
self.dst[i] = v
|
|
}
|
|
}
|
|
}
|
|
|
|
// Fold GEMM's [output * tap, input_time] into the upsampled sequence.
|
|
// Each output has one writer and visits only contributing taps.
|
|
kokoro_overlap: flat{
|
|
x: shader.buffer(f32)
|
|
bias: shader.buffer(f32)
|
|
dst: shader.buffer_rw(vec4f)
|
|
channels: shader.uniform(u32)
|
|
iw: shader.uniform(u32)
|
|
xs: shader.uniform(u32)
|
|
ow: shader.uniform(u32)
|
|
os: shader.uniform(u32)
|
|
kw: shader.uniform(u32)
|
|
stride: shader.uniform(u32)
|
|
pad: shader.uniform(u32)
|
|
has_bias: shader.uniform(u32)
|
|
main: fn() {
|
|
let i = self.index()
|
|
let channel = i / (self.os / 4)
|
|
let time = (i % (self.os / 4)) * 4
|
|
if channel < self.channels {
|
|
var value = vec4f(0.0)
|
|
if self.has_bias != 0 { value = vec4f(self.bias[channel]) }
|
|
for j in 0..4 {
|
|
let pos = time + j + self.pad
|
|
var tap = pos % self.stride
|
|
while tap < self.kw {
|
|
if time + j < self.ow && pos >= tap && (pos - tap) / self.stride < self.iw {
|
|
value[j] += self.x[(channel * self.kw + tap) * self.xs + (pos - tap) / self.stride]
|
|
}
|
|
tap += self.stride
|
|
}
|
|
}
|
|
self.dst[i] = value
|
|
}
|
|
}
|
|
}
|
|
|
|
// op: copy/scale, scaled add, leaky rectifier, Snake, or GELU.
|
|
kokoro_element: flat{
|
|
x: shader.buffer(vec4f)
|
|
b: shader.buffer(vec4f)
|
|
alpha: shader.buffer(f32)
|
|
dst: shader.buffer_rw(vec4f)
|
|
count: shader.uniform(u32)
|
|
stride: shader.uniform(u32)
|
|
op: shader.uniform(u32)
|
|
scale: shader.uniform(f32)
|
|
main: fn() {
|
|
let i = self.index()
|
|
if i < self.count / 4 {
|
|
let v = self.x[i]
|
|
var out = v * self.scale
|
|
if self.op == 1 { out = (v + self.b[i]) * self.scale }
|
|
if self.op == 2 { out = max(v, vec4f(0.0)) + min(v, vec4f(0.0)) * self.scale }
|
|
if self.op == 3 {
|
|
let a = self.alpha[i * 4 / self.stride]
|
|
let s = sin(v * a)
|
|
out = v + (1.0 / a) * s * s
|
|
}
|
|
if self.op == 4 {
|
|
out = 0.5 * v * (1.0 + tanh(0.7978845608 * (v + 0.044715 * v * v * v)))
|
|
}
|
|
self.dst[i] = out
|
|
}
|
|
}
|
|
}
|
|
|
|
// Two-pass variance over the logical row, excluding its padding. Affine
|
|
// mode 1 is per column; mode 2 is AdaIN's [gamma | beta] per channel.
|
|
// Fuse the following activation so each conditioned norm writes once.
|
|
kokoro_norm: {
|
|
workgroup_size: 128
|
|
x: shader.buffer(vec4f)
|
|
affine: shader.buffer(f32)
|
|
beta: shader.buffer(f32)
|
|
alpha: shader.buffer(f32)
|
|
dst: shader.buffer_rw(vec4f)
|
|
cols: shader.uniform(u32)
|
|
stride: shader.uniform(u32)
|
|
rows: shader.uniform(u32)
|
|
mode: shader.uniform(u32)
|
|
activation: shader.uniform(u32)
|
|
eps: shader.uniform(f32)
|
|
slope: shader.uniform(f32)
|
|
partial: shader.workgroup(f32, 128)
|
|
total: fn(v) {
|
|
let sum = subgroup_add(v)
|
|
if subgroup_lane() == 0 { self.partial[subgroup_id()] = sum }
|
|
barrier()
|
|
var result = 0.0
|
|
for j in 0..(self.workgroup_size + subgroup_size() - 1) / subgroup_size() {
|
|
result += self.partial[j]
|
|
}
|
|
barrier()
|
|
return result
|
|
}
|
|
main: fn() {
|
|
let row = group_id().x
|
|
let lane = local_id().x
|
|
let offset = row * self.stride / 4
|
|
var col = lane * 4
|
|
var sum = 0.0
|
|
while col < self.cols {
|
|
let v = self.x[offset + col / 4]
|
|
for j in 0..4 { if col + j < self.cols { sum += v[j] } }
|
|
col += self.workgroup_size * 4
|
|
}
|
|
let mean = self.total(sum) / float(self.cols)
|
|
col = lane * 4
|
|
var variance = 0.0
|
|
while col < self.cols {
|
|
let delta = self.x[offset + col / 4] - mean
|
|
for j in 0..4 { if col + j < self.cols { variance += delta[j] * delta[j] } }
|
|
col += self.workgroup_size * 4
|
|
}
|
|
let inv = 1.0 / sqrt(self.total(variance) / float(self.cols) + self.eps)
|
|
col = lane * 4
|
|
while col < self.cols {
|
|
var value = (self.x[offset + col / 4] - mean) * inv
|
|
if self.mode == 1 {
|
|
for j in 0..4 {
|
|
if col + j < self.cols { value[j] = value[j] * self.affine[col + j] + self.beta[col + j] }
|
|
}
|
|
}
|
|
if self.mode == 2 { value = value * (1.0 + self.affine[row]) + self.affine[self.rows + row] }
|
|
if self.activation == 1 { value = max(value, vec4f(0.0)) + min(value, vec4f(0.0)) * self.slope }
|
|
if self.activation == 2 {
|
|
let a = self.alpha[row]
|
|
let s = sin(value * a)
|
|
value += (1.0 / a) * s * s
|
|
}
|
|
for j in 0..4 { if col + j >= self.cols { value[j] = 0.0 } }
|
|
self.dst[offset + col / 4] = value
|
|
col += self.workgroup_size * 4
|
|
}
|
|
}
|
|
}
|
|
|
|
// op 0 transpose, 1 nearest repeat, 2 left reflection padding.
|
|
kokoro_layout: flat{
|
|
x: shader.buffer(f32)
|
|
dst: shader.buffer_rw(vec4f)
|
|
rows: shader.uniform(u32)
|
|
cols: shader.uniform(u32)
|
|
xs: shader.uniform(u32)
|
|
os: shader.uniform(u32)
|
|
ow: shader.uniform(u32)
|
|
op: shader.uniform(u32)
|
|
factor: shader.uniform(u32)
|
|
main: fn() {
|
|
let i = self.index()
|
|
let row = i / (self.os / 4)
|
|
let col = (i % (self.os / 4)) * 4
|
|
if row < self.rows {
|
|
var value = vec4f(0.0)
|
|
for j in 0..4 {
|
|
if col + j < self.ow {
|
|
if self.op == 0 { value[j] = self.x[(col + j) * self.xs + row] }
|
|
if self.op == 1 { value[j] = self.x[row * self.xs + (col + j) / self.factor] }
|
|
if self.op == 2 {
|
|
var src = i32(col + j) - i32(self.factor)
|
|
if src < 0 { src = -src }
|
|
value[j] = self.x[row * self.xs + u32(src)]
|
|
}
|
|
}
|
|
}
|
|
self.dst[i] = value
|
|
}
|
|
}
|
|
}
|
|
|
|
kokoro_depthwise: flat{
|
|
x: shader.buffer(f32)
|
|
w: shader.buffer(f32)
|
|
bias: shader.buffer(f32)
|
|
dst: shader.buffer_rw(vec4f)
|
|
channels: shader.uniform(u32)
|
|
iw: shader.uniform(u32)
|
|
xs: shader.uniform(u32)
|
|
ow: shader.uniform(u32)
|
|
os: shader.uniform(u32)
|
|
kw: shader.uniform(u32)
|
|
stride: shader.uniform(u32)
|
|
pad: shader.uniform(u32)
|
|
main: fn() {
|
|
let i = self.index()
|
|
let channel = i / (self.os / 4)
|
|
let time = (i % (self.os / 4)) * 4
|
|
if channel < self.channels {
|
|
var value = vec4f(0.0)
|
|
for tap in 0..self.kw {
|
|
for j in 0..4 {
|
|
let src = i32(time + j + self.pad) - i32(tap)
|
|
if time + j < self.ow && src >= 0 && src % i32(self.stride) == 0 && src / i32(self.stride) < i32(self.iw) {
|
|
value[j] += self.x[channel * self.xs + u32(src) / self.stride] * self.w[channel * self.kw + tap]
|
|
}
|
|
}
|
|
}
|
|
self.dst[i] = value + self.bias[channel]
|
|
}
|
|
}
|
|
}
|
|
|
|
// One workgroup owns one recurrent direction for the entire sequence.
|
|
// Input projections are a preceding GEMM. Hidden state stays in shared
|
|
// memory and cell state in registers; there is no dispatch per frame.
|
|
// Kokoro's recurrent directions have 256 hidden units.
|
|
kokoro_lstm: {
|
|
workgroup_size: 256
|
|
projected: shader.buffer(f32)
|
|
w: shader.buffer(vec4f)
|
|
bias: shader.buffer(f32)
|
|
dst: shader.buffer_rw(f32)
|
|
hidden: shader.uniform(u32)
|
|
time: shader.uniform(u32)
|
|
stride: shader.uniform(u32)
|
|
ws: shader.uniform(u32)
|
|
reverse: shader.uniform(u32)
|
|
state: shader.workgroup(f32, 256)
|
|
main: fn() {
|
|
let lane = local_id().x
|
|
self.state[lane] = 0.0
|
|
var cell = 0.0
|
|
barrier()
|
|
for step in 0..self.time {
|
|
var time = step
|
|
if self.reverse != 0 { time = self.time - 1 - step }
|
|
var gates = vec4f(0.0)
|
|
if lane < self.hidden {
|
|
for gate in 0..4 {
|
|
let row = gate * self.hidden + lane
|
|
gates[gate] = self.projected[row * self.stride + time] + self.bias[row]
|
|
}
|
|
var c = u32(0)
|
|
while c < self.hidden {
|
|
let h = vec4f(self.state[c], self.state[c + 1], self.state[c + 2], self.state[c + 3])
|
|
for gate in 0..4 {
|
|
let row = gate * self.hidden + lane
|
|
gates[gate] += dot(self.w[(row * self.ws + c) / 4], h)
|
|
}
|
|
c += 4
|
|
}
|
|
}
|
|
// No lane may overwrite the previous state until every
|
|
// lane has completed the recurrent matrix product.
|
|
barrier()
|
|
if lane < self.hidden {
|
|
let input = 1.0 / (1.0 + exp(-gates.x))
|
|
let forget = 1.0 / (1.0 + exp(-gates.y))
|
|
let update = tanh(gates.z)
|
|
let output = 1.0 / (1.0 + exp(-gates.w))
|
|
cell = forget * cell + input * update
|
|
let h = output * tanh(cell)
|
|
self.state[lane] = h
|
|
self.dst[lane * self.stride + time] = h
|
|
}
|
|
barrier()
|
|
}
|
|
}
|
|
}
|
|
}
|