"Combination of WebGPU and WASM for the best speed. WebGPU for all ML model parts, WASM for CPU-heavy things, JS/TS for the rest." — the goal here.
onnxruntime-web is the pragmatic baseline for every engine in this repo, but its
WebGPU EP has real ceilings: no fusion (each op is its own dispatch), GPU↔CPU
syncs on unsupported/dynamic ops, and no int8/int4 kernels. For a hot model the
faster path is hand-written WGSL over GPU-resident tensors — fused kernels,
nothing leaves the GPU between ops. src/gpu/compute.js is that foundation,
started with Kokoro TTS as the target.
GpuContext (src/gpu/compute.js) — pass a GPUDevice (navigator.gpu in the
browser, dawn in Node). Tensors are GPU-resident
{ buf, rows, cols }; only download() copies back to CPU.
| Kernel | WGSL | Notes |
|---|---|---|
matmul(a, b, {bias, act}) |
tiled 16×16, shared-memory | fused bias + activation (none/gelu/tanh/relu) in one dispatch |
conv1d(x, w, {…}) |
one thread / (Cout, Lout) | direct; small + depthwise/grouped convs |
conv1dFast(x, wRows, …) |
implicit GEMM, register-blocked | groups=1 hot vocoder convs; no im2col materialization |
conv1dGemm(x, wRows, …) |
im2col + tiled GEMM | groups=1; kept for reference (materializes patches) |
convTranspose1d(x, w, {…}) |
one thread / (Cout, Lout), gather form | iSTFTNet upsampler + iSTFT overlap-add; groups |
gatherCols(x, idxMap) |
index kernel | length regulator (duration-expand text→mel) |
lstm(x, w, r, b, hid) |
1 workgroup/direction, hidden units as threads | bidirectional, ONNX iofc gates, timesteps in-kernel (H ≤ 256) |
layernorm(x, γ, β) |
1 workgroup/row, 64-lane reduce | row-wise |
adain(x, scale, shift) |
1 workgroup/channel, 64-lane reduce | instance-norm over time + style affine (StyleTTS2 decoder) |
softmax(x) |
row-wise, numerically stable | for attention |
add / mul / leakyRelu |
elementwise | residuals, gating, iSTFTNet activation |
transpose / sliceCols / setCols |
index kernels | multi-head attention plumbing |
Every kernel is parity-checked against a CPU reference on the real GPU:
npm run gpu:verify # ✓ all kernels: max|gpu-cpu| ≤ tol
npm run gpu:bench # GEMM GFLOP/s + a GPU-resident FFN blockLatest run (M5 Pro, dawn):
matmul+bias+gelu max|gpu-cpu| = 2.7e-5 GEMM 512³ 532 GFLOP/s
layernorm max|gpu-cpu| = 2.4e-7 GEMM 512×512×2048 927 GFLOP/s
softmax max|gpu-cpu| = 7.5e-9 resident FFN block (200×512, d_ff 2048)
add / mul exact = 1.25 ms/block, 1 readback total
The FFN block — layernorm(x + W2·gelu(W1·x + b1) + b2), the ALBERT-encoder core —
chains 4 kernels with every intermediate resident on the GPU and a single
readback. That's the fusion+residency win ORT's per-op path can't match; a
6-layer encoder is ~7.5 ms of FFN compute.
The first real Kokoro sub-network is ported and numerically matches the ONNX
model (src/gpu/albert.js, npm run gpu:albert). Kokoro's PL-BERT: vocab 178,
embed 128, hidden 768, FFN 2048, 12 weight-shared layers, 12 heads (head_dim
64), gelu_new, LN eps 1e-12. Embeddings (gather + sum + LN) run on CPU; the whole
transformer stack — QKV projections, 12-head scaled-dot-product attention, output
projection, FFN, residual LayerNorms — runs GPU-resident on the kernels above.
Verified against the real Kokoro weights + an onnxruntime reference (input_ids → ALBERT output), on the M5 Pro GPU:
ALBERT input (embeds + 128→768 map): rel 1.5e-7
ALBERT output (12 layers) : max 1.3e-5 rel 3.4e-6 ← exact, fp32
Reproduce: kokoro-extract-albert.py (trace ALBERT weights out of the ONNX — the
Linear weights are anonymous onnx::MatMul_* initializers, found via the named
bias each feeds) + kokoro-ref-albert.py (expose the ALBERT in/out tensors as
graph outputs, run ORT for ground truth) → npm run gpu:albert. The attention
scale 1/√64 is folded into the query projection so no extra kernel is needed.
Gotcha: loading the weight .bins in Node — readFileSync().buffer is a view into
a shared pool, so slicing it grabs neighbouring garbage → NaN. Copy the exact byte
range (Uint8Array.from(buf)).
The two hard primitives for the prosody/duration predictors and the iSTFTNet decoder, both verified against the real Kokoro weights on the M5 Pro:
| op | target node | shape | result |
|---|---|---|---|
| bidirectional LSTM | predictor/lstm |
inp 640, hid 256, bidir | rel 4.6e-7 (gpu:lstm) |
| ConvTranspose1d | generator/ups.0 |
512→256, L 94→940, K 20, stride 10 | rel 4.7e-7 (gpu:convt) |
- LSTM matches the ONNX op exactly: gate order
iofc, no peephole, one workgroup per direction, hidden units as threads, timesteps looped in-kernel withh/cin workgroup memory. All six Kokoro LSTMs are bidirectional/hidden-256, so this one kernel covers them. - ConvTranspose1d is the iSTFTNet upsampler and the iSTFT overlap-add. Kokoro's iSTFT is a Cos/Sin DFT matmul + ConvTranspose — both now verified — so the vocoder's spectral tail composes from existing kernels; no separate FFT needed (n_fft is small here).
Reproduce: kokoro-ref-lstm.py / kokoro-ref-convt.py (extract weights + expose
the node's input/output as graph outputs for an ORT reference) → npm run gpu:lstm
/ npm run gpu:convt.
- Metal
tanhoverflows. It computesexp(x)directly, so a large argument →Inf/Inf = NaN(CPUMath.tanhsaturates). The gelu kernel clamps the tanh argument to ±20 (already ±1 to f32 precision). Small test matrices hid this; K=512 accumulators (~±40) surfaced it. - Flag constants aren't global in Node.
GPUBufferUsage/GPUMapModeare ambient in the browser;scripts/gpu-globals.mjsregisters them onglobalThisso the same kernel code runs under dawn. That's what makes headless parity testing possible (Chrome has nonavigator.gpuin the automation env). - GEMM: naive 1 TFLOP/s → register-blocked 2.2 TFLOP/s → f16-storage ~2.7
TFLOP/s (~20% of the M5 Pro's fp32 peak). f16 storage is the last portable-WGSL
lever (~1.3–1.5×, parity rel 3e-4); beyond that needs
simdgroupmatrix units.
The kernels are composed into the model graph, verified stage-by-stage against ONNX:
| stage | status |
|---|---|
input_ids → ALBERT → bert_encoder → d_en (the whole text frontend) |
✅ wired + parity rel 4.2e-6 (npm run gpu:kokoro) |
| DurationEncoder (alt. bidir-LSTM / AdaLayerNorm-with-style) → durations | kernels ready (lstm, layernorm, matmul); topology TODO |
length regulate (gatherCols) + F0/N predictor |
kernels ready (gatherCols, lstm, adain, conv1d) |
iSTFTNet decoder — AdaIN resblocks + convTranspose1d upsampling |
kernels ready (adain, conv1dFast, convTranspose1d, leakyRelu) |
| NSF harmonic source (F0→sine via cumsum/phase) + iSTFT generator tail | needs a cumsum/scan kernel + faithful topology (weight-norm, ScatterND) — the last, most intricate mile |
Everything up to d_en runs on the GPU and matches ONNX. The remaining stages need
no new compute kernels except a cumsum/scan for the NSF source; the work is
faithful topology reconstruction + weight extraction (the generator alone is ~500
nodes: 51 Conv, 51 Sin, phase math), which is a multi-session integration, not a
compute problem. Given the perf verdict is a tie with kokoro-js, finishing the
audio tail is a completeness/bundle exercise, not a speed win — the raw-WebGPU
effort pays off on the ORT-blocked models instead.
Kokoro 82M (StyleTTS2 + iSTFTNet) is ~7 subnets. Done so far: the ALBERT text encoder (parity above) and conv1d (regular + depthwise). Remaining, each verifiable through the same harness before wiring:
- ✅ ALBERT text encoder — parity vs ONNX (
gpu:albert). - ✅ conv1d (+ dilation, groups) — prosody predictor & decoder.
- ✅ bidirectional LSTM — duration/prosody predictors + text encoder (
gpu:lstm). - ✅ ConvTranspose1d — iSTFTNet upsampler + iSTFT overlap-add (
gpu:convt). - ✅ AdaIN — instance-norm over time + style affine (
adain). - ✅ LeakyReLU — iSTFTNet activation (
leakyRelu). - ✅ length regulator — duration-expand via
gatherCols. - harmonic+noise source for the generator (sine gen + the STFT is matmul/conv).
- embedding gather — currently CPU (a lookup); move to a kernel if it matters.
Every compute primitive is done and parity-clean (19 kernels, gpu:verify); the
full op set runs in one submit (gpu:kokoro-forward, 271/274). The only remaining
piece for an audio-producing end-to-end is the harmonic source generator — the rest
is wiring. But the perf verdict is settled: raw WebGPU ties kokoro-js at ~10×,
and neither fused-conv nor shared-tile fp16 breaks the tie. The realistic path to a
win is end-to-end f16 storage — a large refactor whose payoff is uncertain when
kokoro-js already ships at 10×. So the recommendation stands: build raw WebGPU for
the models where ORT is blocked (Nemotron int4, Parakeet int8-collapse), not for
Kokoro speed.
Before hand-wiring the whole StyleTTS2 graph (a large multi-block build), profile
where the compute is (kokoro-profile.py, ORT CPU, 2.05 s audio):
| module | share | op | share | |
|---|---|---|---|---|
| decoder (iSTFTNet) | 89.7% | Conv | 64% | |
| bert (ALBERT) | 5.4% | Sin (source) | 9.7% | |
| predictor | 1.4% | STFT | 5.7% | |
| text_encoder | 0.7% | ConvTranspose | 5.0% |
So the ALBERT encoder we ported perfectly is ~5% of the work — the whole
question is the vocoder convs (~106 GFLOP of conv for 2 s audio, dominated by
resblock convs at [128, 9841] K11). npm run gpu:kokoro-cost measures raw-WebGPU
conv at that dominant shape on the M5 Pro:
| conv path | throughput | projected conv-only RTFx |
|---|---|---|
direct (conv1d) |
252 GFLOP/s | ~4.9× — loses to kokoro-js |
im2col + tiled GEMM (conv1dGemm) |
876 GFLOP/s | ~17× — beats kokoro-js (~10×) |
The direct conv1d/convTranspose1d kernels are kept for correctness/parity and
small/grouped convs; the hot vocoder convs route through conv1dGemm.
The GEMM is register-blocked (64×64 block, 4×4 micro-tile per thread): 927 →
2131 GFLOP/s, and the dominant conv via conv1dGemm 876 → 1759 GFLOP/s (that
one conv, in isolation, projects to ~34× RTFx).
Correction to an earlier claim in this doc: the "~16% of peak / hard WGSL wall" was measured against a wrong ~13 TFLOP theoretical peak. The right baseline is an optimized-Metal GEMM. MLX (Metal simdgroup matmul) fp32 = ~5.9 TFLOP/s square on this M5 Pro (4096³ 5.87; fp16 6.15). Against that, the shipped kernel was already at 58%, and the textbook levers do close the gap on large GEMMs:
| variant (4096³) | TFLOP/s | % of MLX |
|---|---|---|
| v1 scalar 2D-blocktile (64×64 / 4×4, was "shipped") | 3.40 | 58% |
v2 transposed As + vec4 FMA from shared |
3.14 | neutral (compiler already coalesced shared reads) |
v3 v2 + vec4 GMEM loads (array<vec4<f32>>) |
3.71 | 63% |
| v4 128×128 block / 8×8 micro-tile, 16 vec4 accumulators | 4.08 | 70% ✓ |
The 8×8/128×128 tile that spilled as fp32 (461 GFLOP/s) does not spill with
vec4 accumulators — that was the miss. All variants parity 0.0 vs CPU (+ fused
gelu/relu 2.4e-7). matmul() now shape-dispatches: M,N,K ≥ 256 & K%8==0 & N%4==0
→ v4, else v1. simdgroup matrix units (MLX/MPS) remain unreachable in portable WGSL
— that's the last ~30%, not the whole gap.
The caveat that actually matters for this repo: the 70% win is on large square
GEMM only. Real speech-model GEMMs are thin (M~200, K/N 512–1024) → v1==v3==v4
(launch/occupancy-bound, kernel variant irrelevant): 200×512×512 = 0.10 TFLOP,
512×1024×4096 = 2.13 (46% of MLX 4.58). MLX's edge on thin shapes is lower launch
overhead, not a better inner kernel. So for the speech models the lever is
reducing dispatch count / keeping data resident (fusion), not a faster GEMM. v4
pays off on genuinely large-GEMM models = the ORT-blocked encoders (Nemotron/Whisper
big FFN). f16 storage still 2× on ≥2048³ square but 0× on thin — dead end for these.
Bench harness: scripts/gemm-bench-{large,shapes}.mjs, gemm-verify-v4.mjs.
Lever ordering, now fully measured: register-blocking (done, 2.2 TFLOP/s) → f16
storage (2× only on ≥2048³ square, 0× on Kokoro/Parakeet/Nemotron shapes) →
simdgroup (blocked on a WGSL extension). raw-WebGPU ties ORT-WebGPU on these
models because both hit the same occupancy/latency wall on thin GEMMs and neither
can reach the matrix units. The one genuine raw-WebGPU differentiator left is a
custom int4/int8 dequant kernel — ORT's WebGPU EP has no int kernels at all, so
Nemotron (int4-only) can't run on ORT-WebGPU at all, whereas a hand-written
int4→f32 dequant + GEMM runs it on the GPU. That's about running where ORT can't,
not out-GEMMing ORT — and it's built and verified (below).
ONNX MatMulNBits (bits=4, block_size=32) is how Nemotron's 219 encoder matmuls are
stored: packed int4 weights [N, nblk, 16] + per-block f32 scales + packed int4
zero-points. ORT's WebGPU EP has no int kernel, so it falls back to WASM (or
can't run). matmulNBits reads the packed int4 + scales + zero-points directly and
dequantizes in-shader (Y = A @ dequant(B)ᵀ, dequant(n,k) = (q−zp)·scale).
Verified against a real Nemotron layer (K=4352, N=1024) and a CPU dequant reference:
parity rel 5.3e-7 (exact — the int4 unpack + block dequant matches ORT's scheme),
and it runs on WebGPU at 3.3 ms/call — a matmul ORT-WebGPU cannot execute at all.
scripts/nemotron-extract-int4.py pulls a MatMulNBits (weights/scales/zero-points)
out of the ONNX; gpu:verify has a self-contained synthetic check (rel 1.2e-7).
This is the honest endpoint of the raw-WebGPU investigation: on speed it ties ORT everywhere (thin GEMMs, no matrix-unit access), but on capability it does one thing ORT-WebGPU can't — run int4 on the GPU. That's the reason to finish a raw-WebGPU Nemotron: not "faster," but "runs in-browser on the GPU at all."
npm run gpu:kokoro-forward replays all 274 real compute ops (Conv /
ConvTranspose / MatMul / Gemm / LSTM, actual shapes from an ORT profile) back-to-
back in a single submit, timing submit→GPU-finish (excludes CPU alloc/record):
| RTFx | |
|---|---|
| raw-WebGPU, all compute ops (M5 Pro, dawn) | ~10× |
| ORT CPU, same ops | ~9× |
| kokoro-js (ORT WebGPU, browser) | ~10× |
Verdict: with correct, register-blocked kernels raw WebGPU matches kokoro-js (~10×) — it does not clearly beat it, and the two obvious levers don't move it:
- Fused conv (
conv1dFast, implicit GEMM, no im2col): parity-clean and, at the dominant shape, exactly as fast as im2col+GEMM (~1750 GFLOP/s) — the heavy convs were already compute-bound, so removing the patch-matrix materialization saves memory but not time. The full-forward stays ~10× because the aggregate is dominated by the many smaller / depthwise convs, which are intrinsically lower-intensity, not by the one flagship conv. - fp16 shared-tile, f32 storage: slower — 515 vs 1750 GFLOP/s. With f32 global buffers there's no memory-bandwidth win (global reads stay f32), just conversion overhead. This is the wrong way to do fp16.
- fp16 storage (f16 global buffers): a 2× win on large GEMMs — but 0× on
Kokoro.
array<f16>storage is reachable in portable WGSL (shader-f16) and the f16 kernels (matmulF16,conv1dFastF16) are parity-clean (rel 3.3e-4). Measured: square GEMM 2048³ 1694 → 3410 GFLOP/s = 2.0×. BUT the f16 path was wired end-to-end through the full Kokoro forward and it changed nothing (202.2 vs 201.8 ms, ~10× both). Why: Kokoro's ops have small M (Cout ≤ 512, mostly 128), so with 64-row blocks there are only ~2 rows of parallelism — the kernel is occupancy/latency-bound there, and f16's bandwidth + 2× ALU can't be cashed in (conv f32 1760 vs f16 1749 GFLOP/sat Cout=128). f16 only helps when M is large. So the ~14× projection was wrong — it assumed the microbench gain carries; it doesn't (microbench ≠ pipeline, again). - …and f16 doesn't help the "large-GEMM" models either. I then tested the
Parakeet-v3 encoder (the supposed large-GEMM win): also 0× (replay 99.7 ms f32
vs 101 ms f16). The crossover is sharper than "large M" — measured f16-vs-f32 at
real shapes:
188·1024·1024,188·1024·4096,1504·1024·1024,512·1024·4096all 1.0×; only2048·2048·2048hits 1.5×. f16 helps only big square GEMMs (≥~2048³). Every real speech-model matmul is "thin" — hidden 1024, sequence a few hundred — and stays occupancy/latency-bound, where halving memory + 2× ALU buys nothing. So f16 is not a lever for Kokoro, Parakeet, or Nemotron. The kernels are kept (correct + occasionally useful) but the honest verdict is: f16 storage is a dead end for these workloads on this WGSL kernel.
Two hard-won measurement lessons:
- Denormals cost ~2×. Replaying with uninitialized (garbage) buffers ran at 389 ms; zero-initialized, the same ops ran at 202 ms. Flush-to-zero / clean inputs matter enormously on Metal.
- Microbench ≠ pipeline. One hot conv at 34× told a rosier story than the full op set at ~10×. Always measure the aggregate.
So the load-bearing question — is raw-WebGPU Kokoro worth building over kokoro-js? — answers: not for raw speed alone today (it's a tie); it's worth it for a smaller/ORT-free bundle, or as the vehicle for the models where ORT is blocked (Nemotron int4, Parakeet int8-collapse). A fused-conv + fp16 pass is what would turn the tie into a win.
Approach: extract each layer's weights + a reference intermediate from the Kokoro
ONNX (via onnxruntime-node), port the layer to WGSL, gate on parity, then chain
GPU-resident. Only when the whole graph is resident + parity-clean do we benchmark
against the current kokoro-js (ORT) path — that's the number that decides whether
raw WebGPU is worth shipping for Kokoro.