web-xpu-ops
v0.3.0
Published
Web primitive layer implementations for xPU: one reference per op, one implementation per backend, and tests that hold them together.
Maintainers
Readme
web-xpu-ops
Web primitive layer implementations for xPU.
One reference per op, one implementation per backend and target, and tests that hold them together.
Mission
The browser is becoming a heterogeneous machine. WebGPU is here, WebNN is arriving, WASM SIMD is everywhere, and the silicon underneath is a CPU, a GPU, and increasingly an NPU that all want different code for the same operation.
Nothing today lets you write an op once and answer the question "is this fast on the thing it is actually running on?" — and that question is the whole job at this layer. This is where the answer should live.
Not as fast as hand-tuned CUDA. That is not the bar. The bar is: on the hardware in front of the user, how close to that machine's ceiling are we, and does anyone notice when it regresses?
What exists today
Thirty-two ops, WGSL only, verified against their references on a real GPU. That is
the count of directories under ops/ — the number the backends table states and
the number harness/distribution.test.ts asserts against the tree, so it is the
one that cannot go stale silently. Counting the rows of the table below gives
something else and always will: a few ops take a row per entry point
(matvecQ8, matmulQ8, ropeAxes), and permute and dequant_transpose are
described in the LLM-engine sections that produced them rather than here.
Who runs these. examples/ carries two image models checked against their
own implementations — Anima-3.8B and
Z-Image — and an LLM engine in llm/. The Anima port and
the engine are importable, as web-xpu-ops/models/anima and
web-xpu-ops/llm/engine — see "Models and engines" below. The audio ops have a
consumer too, and it is in another repository:
VoxShot builds MioTTS on them —
browser text-to-speech and zero-shot voice cloning, istft's "same" padding,
snake.beta, group_norm and conv_transpose among others. It is why those
ops have the conventions they do; the rows below name MioCodec because MioCodec
is what they were written against. There is no TTS example here, deliberately —
the port exists there and one copy of it is better than two.
Speed is unmeasured for all but one of them. The roofline each would be
reported against does not exist yet, and a number without one would be a
statement about this GPU rather than about the kernel — so the column says so
rather than being left blank. The exception is axpy, whose whole reason to
exist is replacing two dispatches with one, so it had to be measured against
the pair it replaces (#152); its row carries both the figure and the conditions
it was taken under.
| op | notes |
| --- | --- |
| matvec | GEMV; torch.mv convention, streaming rather than tiled. Speed unmeasured |
| matvecQ8 | W8A32 GEMV: matvec with the weight held as int8 instead of f32. Weight is [N, ceil(K/4)] u32, four codes packed per word least-significant byte first — the layout a little-endian host gets for free by viewing an Int8Array's buffer as Uint32Array; scale is [N], quantize's per-row absmax convention, applied once per row after the dot product rather than per term. packQ8 packs quantize's Int32Array codes into this layout; the two compose (packQ8(quantize(w).output, N, K) + quantize(w).scales) rather than this op quantizing on its own. Speed unmeasured |
| matvecQ4G128 | W4A32 GEMV: matvecQ8 at four bits, with one scale per group of 128 columns instead of one per row. Weight is [N, ceil(K/8)] u32, eight two's-complement nibbles per word least-significant nibble first — packQ8's byte order at half the width; scale is [N, ceil(K/128)], quantizeQ4G128's per-group absmax. Range is symmetric [-7, 7], scale = absmax/7, and the reciprocal is formed as 7/absmax rather than 1/f32(scale) — all three stated because the alternatives are live conventions and one of them (Q4_0's [-8, 7]) measured better on weight error while breaking argmax. Not "Q4_0 compatible": block 128 vs 32, [-7, 7] vs [-8, 7], f32 scale vs fp16. See "The q4 format" below. Speed unmeasured |
| rmsnorm | workgroup reduction; eps guards an all-zero row. An optional per-group weight of [G, D] for QK-norm — row n takes group n % G, so the grouped axis has to be the one just left of D ([B, S, H, Dh], not [B, H, S, Dh]); G = 1 is the single shared gamma. Reduction stays over D alone, matching F.rms_norm(x, (D,)) * w rather than torch.nn.RMSNorm((H, Dh)), which reduces over both. Speed unmeasured |
| layernorm | two workgroup reductions; biased variance (1/D) and eps inside the sqrt, as torch.nn.functional.layer_norm. Speed unmeasured |
| group_norm | torch.nn.functional.group_norm: statistics pooled over a group of channels, affine applied per channel. For [N, C, L] each of the N × G groups reduces (C/G) × L values, so this is not layernorm with a longer row — that gives one mean per channel, and the two disagree with no error and no change of shape. G = 1 normalises the whole sample, G = C is InstanceNorm. Biased variance and eps inside the sqrt, both measured against torch 2.10 rather than inherited from layernorm (the unbiased form is out by 1.6e-1, eps outside the root by 4.3e-6). C % G != 0 throws, as torch does. Speed unmeasured |
| softmax | max-subtracted, so real logits do not overflow exp |
| activation | relu2, silu, elu, tanh, gelu, gelu_tanh. elu's alpha is a scalar hyperparameter defaulting to 1.0, as torch.nn.ELU. GELU is two functions, not one: the default gelu is the exact erf form (torch.nn.functional.gelu's own default, approximate="none") and gelu_tanh is approximate="tanh" — they differ by up to 4.73e-4, at x = 2.699, so neither is a silent stand-in for the other. Speed unmeasured |
| snake | Two entry points, because the name covers two functions. kernel: x + sin²(α·x)/α with a learned per-channel α (arXiv:2006.08195), as Snake1d in DACVAE. beta: x + sin²(α·x)/β with both learned per channel, as BigVGAN's SnakeBeta and MioCodec's decoder — β = α recovers the first. A checkpoint's alpha tensor does not say which it belongs to, and the BigVGAN family stores logarithms while DACVAE stores values; neither kernel exponentiates, because doing so would be wrong for the other. The epsilon is upstream's, inside the reciprocal and at upstream's value, guarding the divisor — α in the first, β in the second. sin² is the square of the sine. Its own op rather than an activation kind, because α is a buffer and a channel stride rather than a scalar — the reason is written out in ops/snake/reference.ts. Speed unmeasured |
| elementwise | add, multiply, over two equally sized arrays. elementwiseRows (entry point rows) is the same two operations with the right-hand side broadcast along the last dimension — [S, D] ⊕ [D], which is a Linear's bias for add and AdaLN's per-channel scale for multiply. It is a separate function, and takes S and D explicitly, rather than elementwise inferring a broadcast from the lengths: a [3,3] and a [3] admit two different broadcasts (torch 2.10: + c is [[1,3,5],…], + c.unsqueeze(1) is [[1,2,3],…]) and nothing in the lengths says which. Alignment is from the right as in NumPy/PyTorch, so b.length === a.length is a mistake and throws — torch refuses [2,3] + [6] for the same reason. Speed unmeasured |
| axpy | out[i] = y[i] + a * x[i] with a scalar a — torch.add(input, other, alpha=), BLAS's saxpy by name. Exists because ops/elementwise is same-shape only, so a rectified-flow sampler's latent += dt * velocity otherwise costs a full-length buffer of copies of dt plus multiply-then-add. Two entry points: kernel writes a third buffer, inplace updates y through a single read_write binding (aliasing y into kernel's two bindings instead is not an error you get told about — the command buffer is invalidated at finish() and the readback is all zeros). It also rounds once: this device's compiler contracts y + a*x into an FMA, which is what torch.add does on CPU and CUDA, where multiply-then-add rounds the product first — at a = f32(0.1), x = 3, y = -f32(0.3) the two-dispatch path cancels to exactly 0 and this one gives -2^-27. Measured (RTX 5090, driver 610.57.04, Dawn via [email protected], Node v25.6.1, f32, GPU timestamps, median of 5 per session, three sessions, otherwise-idle GPU; ceiling from harness/roofline.ts in the same sessions, 1.69-1.72 TB/s). At N = 262,144 (one 16×128×128 latent): 6.5 µs against 12.8-16.9 µs for elementwise(multiply) + elementwise(add) — 2.0-2.6x, but at only 480-487 GB/s, 28% of the ceiling, so neither path is bandwidth-bound there and the win is doing half the work rather than doing it faster. At N = 1,048,576 it is: 7.7-9.0 µs against 17.1-18.4 µs — 2.0-2.4x, at 1.40-1.63 TB/s or 83-95% of the ceiling, which is what a kernel moving 12 bytes per element and nothing else should reach. That larger size was unmeasurable while writing this op — it aborted or hung three times running — and is measurable now because #161 fixed the cause (the harness let Dawn's GPU instance be collected out from under the device); nothing in this op changed |
| rope | rotary position embedding, with KV-cache offset and NTK / YaRN context scaling. Follows jquesnelle/yarn and transformers, which agree; YaRN's attention temperature is included. An optional precomputed angle table (ropeCache); past its end the angle is recomputed rather than wrapped. headOffset / headCount rotate a subset of the heads and copy the rest through — the axis Irodori-TTS's _apply_rotary_half uses (chunk(2, dim=-2)), not channel-wise partial rotary (rotaryDim), which is the other thing "half-RoPE" is used to mean and is not implemented. Positions are f32, not integers: Z-Image indexes tokens by grid coordinate, but MiniMax-H3's visual VAE normalises each axis to (-1, 1) and multiplies by 2π, so its positions are fractional — a rotation by a fractional angle is the same rotation, and one binding is better than two kernels differing by a type. Speed unmeasured |
| ropeAxes | multi-axis RoPE, in ops/rope beside the 1-D one: the head dim split into contiguous per-axis blocks (Z-Image's [32, 48, 48]), each rotated by that token's own position on that axis. Follows Z-Image's RopeEmbedder / apply_rotary_emb (Tongyi-MAI/Z-Image @ 26f23ed), which is what decides all three of its conventions: positions arrive as an explicit [N, axes] Int32Array (upstream's ids) rather than being derived from a patch grid; one shared thetaBase (256 there), with the exponent normalised by the axis's own channel count, not by the head dim; and pairing is adjacent channels 2i/2i+1 — the same convention rope uses, torch.view_as_complex upstream, and not HF Llama's rotate_half, so a Z-Image checkpoint needs no permuteRopeChannels and a Llama-style one needs the same permutation as for rope. Angles are computed rather than tabulated, so upstream's axes_lens is not a parameter and a negative position turns backwards instead of wrapping to the end of a table. An odd axis dim throws (upstream cannot express one either). No scaling, no head range, no cache — none of the three is what Z-Image asks for, and ropeCacheAxes waits on a measurement. Speed unmeasured |
| alibi | linear attention-score bias (arXiv:2108.12409); slopes follow the paper's own get_slopes, including the non-monotonic appended tail for head counts that are not a power of two. Bias is the paper's relative form m * (j - i), not BLOOM's m * j; masking is the caller's — and attention's mask is an additive bias of exactly this shape (maskShape: [1, H, L]), so the two compose by addition. Speed unmeasured |
| pope | Legendre polynomial position table (arXiv:2405.04585, Eq. 14); order is the position, argument sweeps [-1, 1). posOffset is required because the paper does not say whether positions start at 0 or 1. Speed unmeasured |
| quantize | per-row absmax to int8, symmetric [-127, 127] |
| dequantize | applies both the weight and the activation scale |
| matmul | GEMM; torch.mm convention, shared-memory tiling. Speed unmeasured |
| matmulQ8 | W8A32 GEMM: matmul with the right-hand operand held as an int8 weight instead of f32, matvecQ8's own [M, ceil(K/4)] u32 packed wire format read in-kernel — no separate dequant/transpose pass. Scale is [M], quantize's per-row absmax convention, applied once per output element. Speed unmeasured |
| matmulQ4G128 | W4A32 GEMM: matmulQ8 at four bits, reading matvecQ4G128's own [M, ceil(K/8)] u32 packed wire format and its [M, ceil(K/128)] per-group scales in-kernel — the prefill half of the q4 format, same tiling (TILE = 16) and same argument names as matmulQ8. No bias, deliberately: matmul and matmulQ8 have none either, and a fused one cannot be measured against anything until the plain form agrees with the reference. Speed unmeasured |
| transpose | turned through workgroup memory so both read and write stay consecutive |
| reduce | sum / max / min / mean along an axis |
| gather | row selection, as torch.index_select(table, 0, indices) — not torch.gather; an out-of-range index gathers zeros |
| scatter | indexed writes; colliding indices accumulate — see below |
| stft / istft | torch.stft / torch.istft conventions: centred, reflect padding, one-sided, unnormalised, periodic Hann; istft divides by the w² envelope — see below. istft also takes padding: "same", the Vocos / X-Codec-2 / MioCodec vocoder convention that crops (nFft - hop) / 2 per end so T frames give T * hop samples — not a torch mode, and not composable from one, because the samples center: false would return fail NOLA. Speed unmeasured |
| conv | conv1d, conv2d and conv3d, as torch.nn.functional.conv1d / conv2d / conv3d — a cross-correlation, so the kernel is not flipped (on none of the axes); stride / padding / dilation / groups / optional bias throughout. 2D is NCHW and 3D is NCDHW ([N, Cin, D, H, W], weight [Cout, Cin/groups, KD, KH, KW]), and the spatial arguments take number | [H, W] and number | [D, H, W] — PyTorch's tuple order, measured rather than assumed. 3D exists for MiniMax-H3's visual VAE, which compresses time as well as space; its temporal axis is also what makes the axis order checkable, since [D, H, W] read in any other order gives a differently-shaped tensor rather than a plausible one. padding is an integer count on all three: 'same' / 'valid' are not accepted, because 'same' with an even effective kernel pads asymmetrically in torch and one integer per axis cannot say that — unlike istft's "same", which is a genuinely different result rather than sugar. padding_mode is not here either: reflect and a causal temporal pad are what H3 uses, and neither is a convolution argument. Speed unmeasured; 2D and 3D are one thread per output element with no tiling |
| pad | One axis of a tensor, as torch.nn.functional.pad: constant (with a value), reflect, replicate. reflect does not repeat the edge element and replicate does — swapping the two gives a tensor of the right shape whose entire interior is correct, which is why the distinction is measured against torch rather than described. before and after are separate, because MiniMax-H3's temporal padding is causal: frames before the data and none after, so the frame at t cannot see t + 1. Takes one axis, viewed as [outer, L, inner], and a caller pads three axes by calling it three times — measured to give what torch's single multi-axis call gives, element for element, which is not obvious when a reflection reads neighbours that are themselves reflections. Exists because ops/conv deliberately has no padding_mode. circular is absent: nothing uses it, and an unused mode is an untested one. Speed unmeasured |
| conv_transpose | 1D only, as torch.nn.functional.conv_transpose1d — the decoder half of conv, and what a DAC-style codec upsamples with. Weight is [Cin, Cout/groups, K], the transpose of conv's layout; padding crops the output rather than extending the input; output_padding lengthens the trailing end only and takes no part in the sum. The kernel is not flipped, for the same reason conv's is not. weight_norm is an offline conversion, not a flag here. Speed unmeasured |
| upsample | nearestUpsample2d: nearest-neighbour 2D resample over [N, C, H, W], as F.interpolate(x, size=(outH, outW), mode='nearest') — the size= path, not scale_factor=. The two are not two spellings of one thing: at H = 3, scale_factor=1.6 maps the four output rows 0, 0, 1, 1 and size=(4, ...) maps them 0, 0, 1, 2 (both measured), and only the size path is decided by integers rather than by how a float scale rounds. Source index is floor(dst * f32(inSize / outSize)) computed in f32, torch's OpenCV INTER_NEAREST formula — at H = 14 -> 46, destination row 23 copies source row 6, where exact integer arithmetic says 7. align_corners is not a parameter, because torch raises for it in this mode rather than defaulting it; mode='nearest-exact' is a different function (floor((dst + 0.5) * scale)) and is not implemented. Downsampling throws. No weights and no arithmetic on the values — this is the nearest upsample -> conv half of a decoder that avoids conv_transpose's checkerboard, so a decoder wants both ops and neither substitutes for the other. Speed unmeasured |
| attention | unfused SDPA in two dispatches; torch.nn.functional.scaled_dot_product_attention convention — scale is 1/sqrt(D) from the query's head dim, and causal is upper-left aligned (queryOffset = S - L gives causal_lower_right). mask is torch's float attn_mask, added to the scores (-Infinity masks), not the boolean one — the additive form is what composes with alibi; keyPaddingBias converts a boolean mask, whose polarity is torch's attn_mask (true = attend) and therefore the reverse of nn.MultiheadAttention's key_padding_mask. Broadcast over batch, heads and query rows via maskShape, so [B,1,1,S] and [B,H,L,S] are both legal. causal with mask throws, as torch does. A fully masked row returns zeros, as aten::_safe_softmax does and plain torch.softmax does not. Speed unmeasured |
| ctc_decode | greedy only. Collapse repeats then drop blanks, as torch.unique_consecutive + a blank filter does; blank=0 as in torch.nn.CTCLoss. Lengths are written by the kernel, so nothing reads back |
| flash_attention | the same function as attention, one dispatch, tiled online softmax; the [B, H, L, S] score matrix is never allocated, which is tested by counting bound bytes and not only by the answer. 132n + 44 bytes at L = S = n, D = Dv = 8 against unfused 4n² + 68n + 40. Takes the same additive mask as attention, on the same terms. Speed unmeasured |
| mel | filterbank construction and its application, as two kernels. Defaults are torchaudio.transforms.MelSpectrogram: HTK mel scale, unnormalised triangles, power spectrum, and AmplitudeToDB(stype="power") for the log — base 10, scaled by 20/power, flooring its argument at 1e-10 rather than adding an epsilon. { scale: "slaney", norm: "slaney" } gives librosa.filters.mel's defaults instead; on the same audio the two differ by 200x, so neither is a default worth leaving unstated. No top_db — it needs a reduction over the whole spectrogram. Speed unmeasured |
| moe | MoE routing: router, dispatch, gather. Softmax before top-k with the k gates renormalised or not, as MixtralSparseMoeBlock and norm_topk_prob (no default: the Switch Transformer must not renormalise at k = 1); top-k ties go to the lower expert index, which torch.topk leaves undefined; capacity overflow drops by rank, then by token index, as GShard / Switch / fairseq top2gating, not by arrival; the gate is applied in gather and only there. Speed unmeasured |
| gqa | grouped-query and multi-query attention: one op, parameterised by kvHeads — kvHeads = 1 is MQA and kvHeads = H is attention unchanged. Contiguous groups (kvHead = h / (H / kvHeads)), as enable_gqa=True in torch; H % kvHeads != 0 throws rather than guessing a grouping. What it buys, in bytes: one Llama-3-8B decoder layer (B=1, S=8192, D=Dv=128, f32) caches 268,435,456 at kvHeads=32, 67,108,864 at 8, 8,388,608 at 1 — over 32 layers, 32→8 saves 6,442,450,944 bytes. kvCacheBytes() computes it. Takes the same additive mask as attention, with maskHeads counting query heads. Speed unmeasured |
scatter: colliding indices accumulate
Two slots naming the same target add. That is a decision, and it is the one thing about this op a caller has to know before using it.
The alternative usually offered is "last write wins", and on a GPU that is not a rule, it is undefined behaviour with a reassuring name: the order threads reach a slot is unspecified, so last means whatever the driver did that day. Callers would build on whichever answer their first device gave. Accumulation is the only rule that returns the same answer for every possible ordering, and it is what the things scatter is actually used for — gradient accumulation, MoE dispatch, bincount — want anyway. The kernel pays an atomic per write for it.
It matches torch.zeros(N, D).scatter_add_(1, index, src), deliberately not
scatter_, which PyTorch itself documents as non-deterministic on collision.
Three departures from PyTorch, spelled out in ops/scatter/reference.ts: self
is implicitly zero, out-of-range indices are dropped rather than raising, and
index is i32 because WGSL has no 64-bit integer.
What stays order-dependent is the last bit or two of a collided sum — f32 addition is not associative. Measured on this GPU at 75 collisions per slot: up to 4.1e-7 relative, about three f32 epsilons, which is the tolerance those tests are set from.
stft / istft: the conventions, and what they match
Window, hop, centring, padding and normalisation each have more than one
reasonable answer. Choosing silently means every caller has a 50% chance of a
subtly wrong waveform, so all of them are named here and every one follows
torch.stft / torch.istft, checked against torch 2.10 numerically rather
than read off the documentation:
| thing | here | matches |
| --- | --- | --- |
| centring | center = true: frame f centres on sample f * hop | torch.stft(center=True) |
| padding | reflect, without repeating the edge sample | pad_mode="reflect" |
| frames | 1 + floor(L / hop) centred | torch |
| sidedness | one-sided, floor(nFft / 2) + 1 bins | onesided=True |
| scaling | none forward, 1 / nFft inverse | normalized=False |
| hannWindow | periodic | torch.hann_window, scipy.signal.get_window("hann"), librosa — not np.hanning, which is symmetric |
| Nyquist bin | counted once, imaginary part dropped | torch.fft.irfft |
istft divides by the overlap-added w² envelope rather than assuming the
window is COLA. That inverts any window satisfying NOLA — the same least-squares
inverse torch computes — and it matters for the ordinary case: a periodic Hann at
50% overlap is COLA in w but not in w², since sin⁴θ + cos⁴θ runs
between 0.5 and 1. Skipping the division is wrong by up to 2x and still sounds
like audio. The reference refuses a window whose envelope drops below 1e-11,
which is torch's own threshold, bracketed by bisection against it.
Two deliberate departures, both explained in ops/stft/reference.ts: the layout
is frame-major [frames, bins] where torch is [bins, frames], because a
vocoder head emits one row per frame — MioCodec's istft_head is a
Linear[.., 1922] with 1922 = 2 * (1920 / 2 + 1) — and because it lets the
kernel write consecutively. And asking for more output samples than the frames
reach raises, where torch pads the tail with zeros and warns; a silent zero tail
is indistinguishable from silence in a vocoder output.
The kernels are a naive DFT per frame, not an FFT. 1920 is 2^7 * 15, so radix-2
does not apply, and this runs beside a transformer over a few hundred frames.
Speed is unmeasured.
LLM tokenizer
llm/tokenizer.ts is not a WGSL op — it never touches a GPU device — but it
lives beside the ops because the LLM engine (#98) that consumes those ops
needs text turned into token ids and back before any of them run. It is a
from-scratch TypeScript reimplementation of SentencePiece unigram
encode (Viterbi) / decode: no wasm, no vendored binary, no runtime dependency
on the sentencepiece package.
import { SentencePieceTokenizer } from "web-xpu-ops/llm/tokenizer";
// vocab is whatever llm/tools/export_tokenizer.py produced for your model —
// fetch it, import it as JSON, however your bundler prefers.
const tokenizer = new SentencePieceTokenizer(vocab);
const ids = tokenizer.encode("<|system|>You are Alibi.</s>");
tokenizer.decode(ids); // "<|system|>You are Alibi.</s>"Every normalizer/byte-fallback/special-token behavior is read out of a real
model rather than assumed — this matters because SentencePiece's own common
defaults (NFKC normalization, an implicit leading ▁) do not apply to
every model, including the one this was built against
(Sarashina2.2-1B-Instruct). llm/tools/export_tokenizer.py parses
tokenizer.model's protobuf directly and refuses to silently drop a
normalizer_spec it does not implement; llm/tokenizer.ts's module doc
records exactly what was verified and how. Ground truth for correctness is
Python sentencepiece's SentencePieceProcessor, never transformers —
transformers 5.3.0 silently converts a UNIGRAM model with no
tokenizer.json beside it into an approximate BPE tokenizer that does not
reproduce true unigram segmentation.
llm/tools/gen_fixtures.py bakes real-tokenizer encode/decode pairs (80+
cases: Japanese/English/mixed, code, emoji, whitespace patterns, and the
<|system|>/<|user|>/</s>-style chat-template boundaries the engine
emits) that llm/tokenizer.test.ts holds this implementation to exactly.
Install
npm install web-xpu-opsCompiled JavaScript and .d.ts ship beside the .wgsl kernels, so nothing at
the other end needs a TypeScript toolchain to consume this. ESM only, Node ≥ 20.
Using it
An op is two things, and they are imported separately because they run in different places.
The reference is plain TypeScript with no dependencies — no GPU, no node:
imports — so the same import works in a browser, in Node and in a test runner:
import { matmul } from "web-xpu-ops/ops/matmul";
const c = matmul({ a, b, M: 64, N: 64, K: 64 }); // Float32Array, [M, N]It is the definition of what the op means, not a fast path. It accumulates in f64 and is written to be read rather than to be quick, which is the whole reason it can be trusted as the thing a kernel is checked against. Reach for it to verify a result, not to produce one at speed.
The kernel is the WGSL, published at a path that mirrors the resolution
grammar (<entry>[.<target>][.<dtype>].wgsl):
import code from "web-xpu-ops/ops/matmul/wgsl/kernel.wgsl?raw";Turning a .wgsl file into a string is your bundler's job rather than this
package's — ?raw is Vite's spelling, webpack wants asset/source, and fetching
the file at runtime works too. What is promised here is only that the path is
stable and the file is present.
You supply the GPUDevice. This library has no opinion about how an application
gets one, and the Node-and-Dawn runner under harness/ is test infrastructure
rather than a runtime: it imports vitest, so it is deliberately not published.
What is published from it is the contract — web-xpu-ops/harness carries
Runner, ResidentDevice, params and the kernel-source registry, and
imports nothing at runtime.
Models and engines
Issue #224. The Anima-3.8B port and the Llama engines ship as subpaths, so a
page or a worker can import the same code the verify-* scripts hold to the
goldens rather than copy it.
| subpath | what it is |
| --- | --- |
| web-xpu-ops/models/anima | the resident DiT forward, the CPU forwards it is checked against, both tokenizers, the Qwen3-0.6B encoder (CPU and GPU), the adapter, the sampler, the VAE decoder, and createBrowserResidentDevice / createBrowserRunner over navigator.gpu |
| web-xpu-ops/models/anima/kernels | animaKernels(load) — builds the three shader tables from WGSL you fetched or bundled; ANIMA_KERNEL_FILES names the files |
| web-xpu-ops/models/anima/fetch-weights | the Cache API loader (Range required), kept apart so a host with its own storage does not inherit its policy |
| web-xpu-ops/llm/engine | LlamaEngine, LlamaEngineQ8, LlamaEngineQ8Resident, the IndexedDB weight cache, and registerKernelSources |
| web-xpu-ops/harness | Runner, ResidentDevice, params, opKernel, registerKernelSources — the contract, no Dawn |
Every dispatching module takes its WGSL as strings. Nothing here reads a file or fetches one; the host says where the text comes from, once:
import { LLM_KERNEL_SOURCES, LlamaEngineQ8Resident, registerKernelSources } from "web-xpu-ops/llm/engine";
import { ANIMA_KERNEL_FILES, animaKernels } from "web-xpu-ops/models/anima/kernels";
// The engines resolve `(op, entry)` through a registry — fill it from what
// your bundler inlined (LLM_KERNEL_SOURCES lists the fifteen it will ask for).
registerKernelSources({ rmsnorm: { kernel: rmsnormWgsl }, matvec: { q8: matvecQ8Wgsl /* … */ } });
// The Anima forwards take explicit tables instead — one loader, three tables.
const { dit, encoder, vae } = animaKernels((op, entry) => wgslText[`ops/${op}/wgsl/${entry}.wgsl`]);Where the .wgsl text comes from is the same story as above: a bundler's raw
loader, or a fetch of web-xpu-ops/ops/<op>/wgsl/<entry>.wgsl at startup.
What is not here: the weights. The converted checkpoints are redistributed
under conditions still being settled (issue #190); the loaders take a URL or a
directory the host provides. Speed for these paths is measured in
examples/anima/README.md and
examples/anima-web/README.md, with the
conditions each number was taken under.
Everything below this line is design, not code. It is written down so the shape is decided before there is enough built for the shape to be hard to change.
The reference is the point
ops/rmsnorm/
reference.ts what correct means — plain, slow, obviously right
wgsl/kernel.wgsl
wgsl.test.tsBackends multiply; correctness does not. A kernel is only ever measured against the reference, never against another kernel. That is what stops a second backend from quietly redefining an op to match whatever it happens to compute, and it is why the reference is deliberately the slowest expression of the maths — its job is to be obviously right.
Targets
The same WGSL is not the same speed everywhere, and sometimes it should not be
the same WGSL. Workgroup size, tiling, whether subgroup operations exist, whether
f16 exists at all — these differ per target, and a single kernel that is
adequate everywhere is usually good nowhere.
Planned axis:
ops/matmul/
reference.ts
wgsl/
kernel.wgsl portable, correct, unremarkable
kernel.apple.wgsl unified memory, wide subgroups
kernel.nvidia.wgsl
kernel.amd.wgsl
kernel.soc.wgsl tight power and bandwidth budgets
webnn/graph.ts
wasm/kernel.tsSelection resolves in order, first hit wins:
explicit override → target + dtype → target → dtype → portableThe resolution is implemented (harness/resolve.ts, harness/target.ts);
the per-target kernels are not — no op ships a variant yet, so every op resolves
to its portable kernel today.
An op's wgsl/ directory holds one or more entry points, and each may have
variants. The filename is the whole grammar — <entry>[.<target>][.<dtype>].wgsl:
kernel.wgsl the default entry point, portable
inverse.wgsl a second entry point, portable (ops/stft)
scores.wgsl one of two entry points, portable (ops/attention)
kernel.nvidia.wgsl a target variant of kernel
inverse.f16.wgsl a dtype variant of inverse
scores.apple.f16.wgsl bothEvery entry point needs its own portable kernel, and resolution never leaves
the entry point it was asked about — istft falling back to the forward
transform because the chain ran off the end would be a wrong answer that still
looks like a result.
Entry points are named rather than inferred because an op genuinely may need
several kernels that are not variants of each other: stft computes the inverse
transform with different arithmetic, and attention is two dispatches with a
buffer between them, deliberately split so layout: "auto" cannot drop bindings
an entry point does not reference.
A suffix that is not a known target or dtype is an error, so kernel.nvidai.wgsl
is rejected rather than left in the tree looking tuned. A bare nvidia.wgsl
is rejected too, as ambiguous: it is far likelier to be a mis-written variant
than an entry point that happens to be called "nvidia".
What that cannot catch is a misspelt entry point — inverze.wgsl reads as a new
one. The asymmetry is deliberate. A misspelt variant fails silently, because
resolution falls back to portable and everything still runs; a misspelt entry
point fails loudly the moment its test asks for it by name. Only the silent one
needs a guard.
Overrides are a first-class input, not an escape hatch. Anyone integrating this will eventually know something the library cannot — that their sequence length is always 1, that their weights are static, that they would rather have lower peak memory than higher throughput. Refusing that knowledge makes the library something to work around. An override that names a variant which does not exist raises; falling back to portable would hand the caller a kernel they did not ask for and never say so.
Target detection has to be honest about how little is knowable. adapter.info
gives a vendor and an architecture string and not much else, so detection is a
hint, and the override exists partly because the hint will sometimes be wrong.
It is allowed to answer "I don't know", and that answer is portable rather than
a guess: intel covers both an on-die iGPU and a discrete Arc card, and nothing
in adapter.info says which, so it resolves to no target at all. A fallback
adapter is treated the same way — the vendor string there names silicon the code
is not running on.
Nothing picks a target-specific kernel silently. describeAdapter(adapter)
reports what the device said, which target was detected and why, and resolve()
returns the rung that hit along with every candidate it tried. A wrong guess that
is invisible is worse than a portable kernel.
It is a call the caller makes, not something the runner does on the way past.
createRunner's job is to run a kernel and read the result back; which target a
device looks like is a question about the device, and the two do not have to be
answered at the same moment.
Adding a variant cannot skip the reference test. eachVariant(url, entry, …)
builds an op's test loop from its wgsl/ directory rather than from a list
someone has to remember to extend, and unguardedOps fails the suite for an
entry point that grows a variant no test iterates. The check is per entry point,
not per op: looping scores says nothing about context. A target-specific
kernel that is fast and wrong is the failure this axis exists to design against.
All of it lives outside harness/index.ts — import harness/variants.js,
harness/resolve.js and harness/target.js directly. index.ts is what every
op's test imports, so anything re-exported from it lands on the import graph of
every GPU test; keeping resolution off it leaves index.ts and suite.ts
byte-identical to what they were before any of this existed.
Dtypes
f32 is the floor and always works. Everything above it is conditional:
| dtype | availability |
| --- | --- |
| f32 | always |
| f16 | requires the shader-f16 feature |
| i32 / u32 | always |
| packed int8 / int4 / ternary | manual bit packing; no native type |
The f16 line is not theoretical. Measured during this project's own work:
- Apple M3 —
shader-f16present,q4f16selected, and still slower than real time for the workload under test. - Linux / NVIDIA with Chrome's Vulkan backend disabled — Dawn falls back to an
ANGLE compatibility adapter advertising exactly one feature. No
shader-f16at all.
So f16 is not a switch to flip for a free win. It is a per-target, per-op
question with an answer that has to be measured, and the availability check is
load-bearing rather than defensive.
Integer paths matter more here than in a datacentre library: quantized weights are how a model fits down a browser's throat at all.
GEMM and GEMV are separate ops
Not two paths through one kernel. Two kernels.
GEMV — one vector against a matrix — reads every weight exactly once and reuses nothing. It is bandwidth-bound, and the only thing that matters is how fast the weights can be pulled through. Tiling for reuse is wasted code, because there is no reuse to find.
GEMM has reuse proportional to the tile size and is compute-bound once the tiling is right. Everything that makes GEMM fast is irrelevant to GEMV, and the work of keeping both inside one kernel is spent on branches that make each of them worse.
Autoregressive decoding is GEMV-shaped: batch of one, every step. Prefill is GEMM-shaped. A library that only does one of them well is only good at half of inference.
The q4 format (W4A32, group-128) — issue #137
The second quantized weight format in this library, beside quantize's per-row
int8. matvecQ4G128 reads it.
| | value | why not the alternative |
| --- | --- | --- |
| bits | 4, 8 codes per u32, least-significant nibble first | packQ8's byte order at half the width, so a converter that already emits q8 changes constants, not code |
| range | [-7, 7], symmetric | [-8, 7] (Q4_0) clips one tail only, which turns quantization error into a systematic bias — see below |
| scale | absmax / 7, f32, one per 128 contiguous columns of a row, [N, ceil(K/128)] | per-row loses too much at 4 bits; g64 was not distinguishable from g128 by anything but bpw |
| reciprocal | 7 / absmax, formed in f64 | 1 / f32(scale) (llama.cpp) rounds differently at code boundaries |
| rounding | Math.round — ties toward +Infinity | matches ops/quantize and llm/tools/quant_common.py exactly |
This is not Q4_0 and does not claim to be. It is in the same family
(per-block symmetric absmax, no zero point), and differs in three ways that
change the numbers: block 128 rather than 32, [-7, 7] rather than [-8, 7],
f32 scale rather than fp16. The group axis is the one GPTQ/AWQ call
group_size, but those carry a zero point and this does not. Q4_K / Q4_1 —
asymmetric, with a per-block minimum — are not implemented.
What is measured here, and what is not
Measured in this repository by ops/matvec/q4.quality.test.ts — CPU only,
so it cannot pass vacuously for want of a GPU; npm run test:file
ops/matvec/q4.quality.test.ts reprints the table. N=256, K=2560, deterministic
LCG input, f64 reference arithmetic. These are synthetic matrices, not model
weights:
| matrix | weight RMS rel: q8-row / q4-row / q4-g128 | GEMV peak rel: q8-row / q4-row / q4-g128 | | --- | --- | --- | | iid gaussian | 7.97e-3 / 1.45e-1 / 1.14e-1 | 9.44e-3 / 2.08e-1 / 1.50e-1 | | + a 20x outlier column every 64 | 4.05e-2 / 3.83e-1 / 2.90e-1 | 3.22e-2 / 2.85e-1 / 2.33e-1 | | per-128-column magnitude bands (1x/10x/100x) | 1.32e-2 / 1.68e-1 / 1.14e-1 | 1.88e-2 / 2.25e-1 / 1.36e-1 | | …the same, restricted to the 1x columns | 7.13e-1 / 1.00e+0 / 1.15e-1 | — |
The last row is the whole argument for the group axis, and it is the reason the
three rows above it look unimpressive: a global RMS-relative figure is dominated
by whichever columns are largest, so it barely moves. Restricted to the columns
that are two orders of magnitude smaller than their row's peak, per-row q4
scores exactly 1.0 — every code in those columns rounds to zero and the band
is annihilated — while g128 is unaffected at 1.15e-1. Note that q8's per-row
scale loses those columns too (7.13e-1): the axis is doing the work here, not
the bit width. An iid matrix cannot show any of this, because there is nothing
for a per-group scale to adapt to, and neither can uniformly-spread outliers —
both are in the table so that "group-wise is better" is not read as
unconditional. Real weights are neither.
Wire size, measured from the buffers packQ4 actually produces (same test):
4.250 bpw for q4-g128 at any K that is a multiple of 128 — 4 bits of code
plus one f32 scale per 128 weights — against 4 + 32/K for per-row q4 and
8 + 32/K for q8. The group axis costs a quarter of a bit per weight.
A [2560, 2560] weight is 26,214,400 bytes as f32, 6,563,840 as q8, and
3,481,600 as q4-g128.
Does a 4B model fit?
Arithmetic, from the measured 4.250 bpw: 4e9 parameters at 4.25 bits is
2,125,000,000 bytes ≈ 1.98 GiB, against 4.02 GiB at q8 (8.012 bpw) and
16 GiB at f32. That is the number issue #149 exists for, and it is arithmetic
rather than a measurement — this repository has no 4B checkpoint, has converted
none, and therefore does not know what the non-Linear tensors (embeddings,
norms, biases) add on top, nor whether any single tensor exceeds
maxStorageBufferBindingSize on a given browser. Issue #112 records that
exceeding a limit is answered with zeros rather than an error, so "it fits
in total" is not the same claim as "it loads".
Not measured here: what this does to a real model. No logits comparison, no
greedy-trajectory comparison, no audio. Issue #137's decisions rest on
voxshot's measurements on MioTTS-0.6B (per-row q4 at 4.5e-1 peak-relative
logit error against q8's 4.8e-2; [-8, 7] flipping the argmax in 4 of 4 cases
despite the best weight RMS error of any configuration tried), and those are
voxshot's numbers, not this repository's. The harness that produced them
(spike/miotts/measure_q4.ts) does not use these kernels, so running it against
weights quantized by quantizeQ4G128 would isolate this implementation's own
error — that comparison has not been run.
Speed is unmeasured, for the kernel and for the format.
Performance, measured against the machine's own ceiling
Correctness tests answer does it work. At this layer the other half of the job is is it fast, and nothing catches a regression there today.
Not against a spec sheet. WebGPU exposes no clock, no compute-unit count, no memory bus width — there is nothing to compute a theoretical peak from, and a number derived from a marketing figure would be worse than no number, because it would look authoritative.
So the ceiling is measured on the same device, in the same session:
- Bandwidth roofline — a kernel that does nothing but stream memory. The best bytes-per-second that device will give anyone.
- Compute roofline — a kernel that does nothing but fused multiply-adds out of registers. The best FLOPs that device will give anyone.
Every op then reports what fraction of the relevant roofline it reached, and
which of the two it is bound by. rmsnorm hitting 78% of measured bandwidth is a
statement about the kernel. "412 GFLOP/s" is a statement about the GPU someone
happened to run it on.
This also makes CI results comparable across wildly different hardware, which is the only way a browser-targeting library can have meaningful performance CI at all. A percentage regresses visibly on a laptop and a workstation alike; an absolute number regresses invisibly on both.
Op roadmap
Three layers, by what they are rather than by what they compute.
primitive/ — the algebra
matmul (GEMM) ✅ · matvec (GEMV) ✅ · conv ✅ (1D) · conv_transpose ✅ (1D) ·
upsample ✅ (nearest, 2D) · add · mul · gather ✅ · scatter ✅ ·
transpose ✅ · reduce ✅
matmul (GEMM) ✅ · matvec (GEMV) ✅ · conv ✅ (1D, 2D) · conv_transpose ✅ (1D) ·
add · mul · gather ✅ · scatter ✅ · transpose ✅ · reduce ✅
Small, total, boring. Everything else is built from these, and they are where target-specific tuning pays off most.
kernel/ — one fused, named operation
rope ✅ · rmsnorm ✅ · layernorm ✅ · group_norm ✅ · softmax ✅ ·
activation ✅ · snake ✅ · elementwise ✅ · axpy ✅ · quantize ✅ · dequantize ✅ ·
attention ✅ · flash_attention ✅ · ctc_decode ✅ (greedy) · mel ✅ ·
stft / istft ✅
Fusion is the reason this layer exists rather than being composed from
primitive/ at call time. flash_attention is not matmul + softmax +
matmul; it is the one that never writes the score matrix to memory, which is
the entire point of it.
mel, stft and istft are here because speech pipelines need them and no ML
kernel library ships them — they are DSP, so everyone assumes someone else has
them. The inverse STFT in particular is the thing ONNX cannot express, being
unable to carry complex tensors.
attention/ — the variants that are their own problem
Position: RoPE ✅ · multi-axis RoPE (ropeAxes) ✅ · ALiBi ✅ · PoPE ✅ · YaRN ✅ · NTK scaling ✅ · rotary cache ✅ · half-RoPE (head range) ✅
Sharing: GQA ✅ · MQA ✅ (one op — gqa, parameterised by kvHeads)
Routing: MoE router ✅ · MoE dispatch ✅ · MoE gather ✅
Serving: paged KV cache · speculative decode
Separated from kernel/ because these are not variations in arithmetic, they are
variations in what memory gets touched. Paged KV cache and speculative decoding
are not faster attention; they are different answers to where the state lives.
block/ — deliberately not here yet
Transformer, conformer, U-Net, decoder blocks are compositions, not kernels. They belong to whoever is building a model, and a kernel library that ships them starts making architectural decisions on its users' behalf.
The exception is a block that only pays off fused — where crossing it as separate
dispatches costs more than the arithmetic does. If one of those turns up, it
belongs in kernel/ under its own name, not in a block/ directory that invites
everything else in behind it.
llm/: an inference engine built from these ops
llm/ is a config-driven llama-architecture (decoder-only, GQA, RoPE, SiLU
MLP, RMSNorm) forward pass, composed entirely from the ops above rather than a
new fused kernel — the "block/" exception the roadmap notes above rules out
does not apply, because a llama decoder is not a block that only pays off
fused; it is a sequence of the primitives and kernels this repository
already ships, and issue #98
tracks that specific composition, not a new op.
// Within the repository. A consumer imports the same engine from the
// package as `web-xpu-ops/llm/engine` (0.3.0, issue #224) and registers the
// WGSL it bundled first — see "Models and engines" above.
import { LlamaEngine, TINY_FIXTURE_CONFIG } from "./llm/index.js";
const engine = new LlamaEngine(config, weights, runner.run);
const prefillLogits = await engine.forward(promptTokens); // matmul path, N > 1
const [nextLogits] = await engine.forward([nextToken]); // matvec path, N === 1Per layer: rmsnorm → QKV projection (fused) → rope(Q, K) → gqa (+ KV
cache) → O projection → residual → rmsnorm → gate/up projection (fused) →
silu(gate) * up → down projection → residual, then one final
rmsnorm → lm_head.
What "config-driven" means in practice — llm/config.ts#LlamaConfig has
one field per dimension a llama checkpoint's config.json carries (layers,
hidden size, query/KV heads, head dim, FFN width, vocab, RoPE base, RMSNorm
eps, weight tying). Two example configs live there: TINY_FIXTURE_CONFIG
(2 layers, hidden 64, 4 query / 2 KV heads, FFN 128, vocab 256 — what the
fixture below actually runs) and SARASHINA_2_2_1B_CONFIG (24 layers, hidden
1792, 16 query / 8 KV heads, FFN 6272, vocab 102400, RoPE base 500000 — issue
#96's real target, documentation only: running the actual checkpoint needs a
weight-conversion tool, a later issue — the tokenizer side is covered by
llm/tokenizer.ts above).
KV cache and dispatch fusion. The cache is pre-allocated f32,
[kvHeads, maxSeqLen, headDim] per layer (llm/kv-cache.ts) — no growth
during generation, optimisation left for later per the "correctness first"
rule below. Q/K/V and gate/up are each fused into one projection weight
(llm/reshape.ts#concatRows / #splitConcatRows) purely to cut GPU dispatch
count: this repository's webgpu (Dawn-through-Node) binding is measurably
unable to sustain an unfused prefill-then-decode run's ~195 dispatches inside
one device's lifetime on some machines, and fusing Q/K/V and gate/up brings
that down to ~155 without changing a single number the engine computes — see
llm/reshape.ts and llm/engine.wgsl.test.ts for the measurement.
Correctness is a fixture, not an eyeball. llm/tools/gen_fixture.py
builds a tiny, randomly-initialised llama model with transformers itself
(rule 7 — a real HF forward pass, attn_implementation="eager", seeded), runs
an 8-token prefill and a 4-step greedy decode, and writes the weights and
logits to llm/fixtures/tiny.*. llm/engine.wgsl.test.ts runs the same
tokens through LlamaEngine on a real GPU and checks logits against the
fixture (measured: worst absolute diff 1.49e-7, worst relative 3.4e-4
across prefill and decode — tighter than any individual op's own tolerance,
because a tiny random model's logits do not accumulate error the way a
trained one's sharper distributions might) and greedy-decoded tokens for
exact equality. The fixture (weights + logits + a JSON manifest) is committed
rather than generated at test time — 436 KiB total (420 KiB weights, 12 KiB
logits, 4 KiB manifest), small enough that regenerating it on every npm test
would buy nothing; see llm/tools/gen_fixture.py for the reproduction steps.
One channel-ordering fact worth knowing if you are reading the weight-loading
code: HF Llama's RoPE (rotate_half) pairs channel i with channel
i + headDim/2, while ops/rope pairs adjacent channels 2i/2i+1 —
the same rotation, numbered differently within a head. llm/weights.ts#permuteRopeChannels
relabels a checkpoint's Q/K projection rows accordingly (exact, not
approximate — the derivation and a numeric proof are in
llm/weights.ts and llm/rope-permutation.test.ts), and
llm/tools/gen_fixture.py applies the same permutation before writing
tiny.weights.bin. ropeAxes pairs adjacent channels too — Z-Image's
torch.view_as_complex(x.reshape(*, -1, 2)) is the same convention — so the
rule is per checkpoint, not per op: rotate_half weights need the permutation
whichever of the two rotates them.
Scope. f32 weights (this section) and int8 weights (next section) both
exist now; the tokenizer and sampler are exported from llm/index.ts and
driven by the browser demo's generation loop (below), not called by the
engine's own forward. Correctness first, per the rule below — nothing here
has been tuned for speed.
llm/: weight converter and the int8 (W8A32) engine path
Issue #105 closes the
loop from "an engine that runs a tiny fixture" to "an engine that loads and
generates from a real checkpoint": llm/tools/convert_weights.py converts a
HF safetensors checkpoint (bf16) into per-row int8, and LlamaEngineQ8
(llm/engine-q8.ts) runs it.
Every number in this section was measured on one machine (rule 9): NVIDIA
GeForce RTX 5090, driver 610.57.04, Linux (Arch, kernel 7.1.5-arch1-2),
backend Dawn via [email protected] under Node v25.6.1, Python 3.14 / NumPy 2.4
for the converter; checkpoint Sarashina2.2-1B-alibi-v1 (bf16, 2.68 GiB).
Timing figures are single observations from that machine, not averages —
comparable only against a rerun under the same conditions.
// llm/tools/convert_weights.py --model-dir <hf checkpoint> --out-dir <out>
// writes <out>/manifest.json + weights.codes.bin + weights.scales.bin + weights.norms.bin
import { loadConvertedWeightsQ8 } from "./llm/real-model-weights.js";
import { LlamaEngineQ8 } from "./llm/engine-q8.js";
const { config, weights } = loadConvertedWeightsQ8("<out-dir>", /* maxSeqLen */ 4096);
const engine = new LlamaEngineQ8(config, weights, runner.run);
const prefillLogits = await engine.forward(promptTokens);Converted format. Every Linear weight (wq/wk/wv/wo/gate/up/down/
lm_head) and the embedding table become per-row absmax int8 codes ([N, K],
one signed byte per element — ops/quantize/reference.ts#quantize's own
convention) plus an f32 scale per row; wq/wk get permuteRopeChannels
applied before quantizing (permuting rows and quantizing per row commute, so
this is equivalent to permuting the codes afterward — the tiny fixture's
generator does it the other way around, on purpose, for a reason its own
module doc explains). Norm weights stay f32. Manifest entries are
{ name, kind: "quant" | "norm", shape, codesOffset, scaleOffset } or
{ name, kind: "norm", shape, offset } — the exact shape llm/weights-q8-io.ts#buildLlamaWeightsQ8
parses, shared by the tiny fixture's loader (fixture-q8.ts) and the real
checkpoint's (real-model-weights.ts) so the two formats cannot silently
drift apart. Converting Sarashina2.2-1B-alibi-v1 (2.68 GiB bf16) produced a
1.41 GiB int8 checkpoint (1,407,451,136 bytes codes + 2,711,552 bytes scales +
351,232 bytes norms) in about 8 seconds.
LlamaEngineQ8's resident memory. Only the packed matvecQ8 wire format
(ops/matvec/reference.ts#packQ8) is kept per projection after construction —
not also the unpacked codes the loader handed in, since packing does not
change a weight's size (repacking would roughly double memory for nothing).
The embedding table is the one exception kept in its original (unpacked)
form, decoupled from the loader's shared buffer via a copy
(weights-q8.ts#cloneQuantizedLinear) — without that copy, the loader's
manifest-parsing convenience (every weight a view into one buffer covering
the whole checkpoint) would keep the entire raw codes buffer resident (~1.4
GiB) for the engine's whole lifetime just to reach the ~183 MiB embedding
table. Measured on the real checkpoint: loading is ~220ms, LlamaEngineQ8
construction (packing every projection) is ~1.1s, and process RSS falls from
~2.96 GiB (while the loader's own buffers are still referenced) to ~1.56 GiB
once the caller drops that reference and a GC runs.
Decode vs. prefill. Decode (tokens === 1) dispatches matvecQ8 directly
against the resident packed weight. Prefill (tokens > 1) dequantizes the
needed projection into a transient f32 matrix and runs matmul — issue
#105's own stated scope ("プリフィルは当面「行スケールdequantしてf32 matmul」
でもよい"), which happens once per generation (the prompt), not once per
token, since greedyGenerate calls forward with more than one token exactly
once. A prefill kernel that reads packed int8 directly is explicit follow-up
work, not done here.
The embedding table's gather is CPU-side, not a GPU dispatch. embedTokens
is quantized like every other weight, but LlamaEngineQ8 dequantizes only the
rows a forward call's tokens actually name, on the CPU
(weights-q8.ts#gatherDequantRows), instead of dispatching runGather
against a fully-dequantized 733 MiB table for a call that reads at most
maxSeqLen of its rows.
Correctness: an int8-quantization-aware fixture. llm/tools/gen_fixture_q8.py
quantizes the tiny fixture's own weights per row, substitutes the
dequantized weights back into the transformers model, and re-runs the
same prefill/decode loop gen_fixture.py runs — so the reference
(llm/fixtures/tiny_q8.*) has the same quantization error baked in that
LlamaEngineQ8 produces, rather than being compared against an f32-exact
answer a genuinely quantized engine could never match. llm/engine-q8.wgsl.test.ts
checks LlamaEngineQ8 against it on a real GPU: worst observed absolute diff
1.49e-7, relative 7.22e-4 (well inside the rel 1e-2, abs 5e-3
tolerance, and close to LlamaEngineQ8's own f32 counterpart's numbers —
evidence the int8 path and the Python reference agree on what quantization
error to expect, not that either one is loose). llm/quantize-parity.test.ts
separately checks that llm/tools/quant_common.py (the quantizer
convert_weights.py and gen_fixture_q8.py both call) rounds ties exactly
the way ops/quantize/reference.ts#quantize's Math.round does
(floor(x) + (frac >= 0.5), not np.round's banker's rounding — which
disagrees at every .5 boundary — and not the tempting np.floor(x + 0.5),
whose addition double-rounds just below half-integers; see
quant_common.py's module doc) by spawning the Python script and diffing its
output directly, including both boundary cases.
Real-checkpoint status. Converting, loading (real-model-weights.ts,
checked in real-model-weights.test.ts), and constructing LlamaEngineQ8
from the real Sarashina2.2-1B-alibi-v1 checkpoint all succeed and are
verified. Live GPU generation from the real checkpoint does not run yet,
for two separate reasons, neither a numerics problem (the int8 path is
verified above, to a tight tolerance, on real GPU dispatches):
- On the machine this was built on, a Node+Dawn binding fragility is triggered by real-model-scale CPU-bound work (loading and packing a ~1.4 GiB checkpoint, on the order of a second) immediately preceding a GPU dispatch; see #107 for the isolated repro (an unrelated large allocation, and separately a pure CPU busy-loop with no allocation at all, both reproduce it — GPU contention and buffer count/size were ruled out).
- Independently of #107 and of the machine, the real vocabulary size breaks
the lmHead projection against WebGPU's own limits: decode dispatches one
workgroup per output row (102,400 > the default 65,535
maxComputeWorkgroupsPerDimension), and prefill dequantizes lmHead to a ~700 MiB f32 matrix that exceeds the 512 MiB buffer/binding cap the harness requests (and the 128 MiB browser default). Found in review, not yet hit at runtime only because #107 aborts earlier; tracked as #112, which blocks #106's real-model demo as well.
llm/engine-q8.real-model.test.ts runs this end to end (prefill + greedy
decode, tok/s per step) when a converted checkpoint and an encoded prompt are
supplied via environment variables, and skips (visibly, not as a silent pass)
otherwise; resolving #107 and #112 is what turns its assertions on. Live
generation, llama.cpp comparison, and tok/s are planned for
#106 (browser demo),
where WebGPU runs in a separate GPU process rather than sharing Node's — the
condition #107 depends on does not exist there (but #112 applies to the
browser all the same).
llm/: sampler and token-level constraints
llm/sampler.ts turns a next-token logits vector into a token id — greedy
(argmax) or temperature + top-p (nucleus) — and takes an optional
Constraint:
interface Constraint {
nextAllowed(prefixTokens: readonly number[]): ReadonlySet<number> | null;
}null means unconstrained; an empty set means no legal continuation, which
sampleNext treats as an error rather than guessing. The mask is applied to
the raw logits before either sampling mode runs, so a constrained draw can
never land outside the allowed set regardless of temperature or top-p.
There is no repetition penalty, on purpose. The usual per-seen-token
logit penalty was tried against Alibi's Japanese model and measured to
degrade output rather than de-loop it — see
technologies-moe/alibi-ai#3.
A caller that wants to discourage repetition should express it as a
Constraint, this module's own tool for narrowing the next token, rather
than as a penalty baked into the sampler.
llm/constraints/line-format.ts is one Constraint: a state machine for a
fixed line shape — literal text, an enum choice, more literal text, free
text (a forbidden-character set and a max length), then EOS. It exists for
schemas as small as policy: <enum>\ntopic: <short text>, where a full
GBNF/grammar engine is more machinery than the shape needs. It is
tokenizer-agnostic: it takes an injected TokenCodec
(encode / idToToken / vocabSize) instead of depending on any one
tokenizer —
const constraint = new LineFormatConstraint(codec, {
segments: [
{ kind: "literal", text: "policy: " },
{ kind: "enum", choices: ["allow", "deny", "review"] },
{ kind: "literal", text: "\ntopic: " },
{ kind: "freeText", forbiddenChars: ["\n"], maxLength: 80 },
],
eosTokenId,
});
const next = sampleNext(logits, generatedSoFar, { mode: "greedy" }, constraint);Enum choices are matched by tokenizing each candidate once and walking a token-id trie — the "candidate string, tokenized, forward-matched" approach — so choices must be token-prefix-free: no choice's tokenization may be a strict prefix of another's, or completion would be ambiguous. Such a spec is rejected at construction rather than silently misclassified during generation.
llm/index.ts re-exports the tokenizer, the sampler and this constraint
(issue #106 — both had landed without ever being re-exported from the
package's own entry point, a gap #106's own text called out). The decode
loop that calls sampleNext once per step is still not inside
LlamaEngine/LlamaEngineQ8 itself — a caller drives it, the way
examples/llm-demo/src/main.ts does below — since a constraint is a
generation-level policy, not something the engine's forward (one call, one
set of logits) has an opinion about.
Browser demo (examples/llm-demo/)
Issue #106: every
piece above — tokenizer, LlamaEngineQ8, sampler, LineFormatConstraint —
running together in a real browser tab over WebGPU, loading a real converted
checkpoint over HTTP. This is the first time this repository's llm/ code
has generated anything outside Node; PR #108 could not complete a live GPU
dispatch on the machine it was built on (a Node+Dawn binding fragility
triggered by real-model-scale CPU-bound work immediately before a dispatch —
#107), and moved that
gate here on the strength of one fact: a browser's WebGPU implementation runs
in a separate GPU process, so the "CPU-bound work in the same process right
before a dispatch" condition #107 isolated does not exist there.
Running it
npm run demo:build # esbuild: examples/llm-demo/src/main.ts -> dist/bundle.js
npm run demo:serve # node examples/llm-demo/server.mjs — Node standard library onlyThen open http://localhost:8770/examples/llm-demo/ in a WebGPU-capable
browser. demo:serve serves this repository (so llm/data/*.vocab.json and
the demo's own dist/bundle.js are reachable) and additionally maps
/weights/ onto a converted-checkpoint directory outside this repository —
convert_weights.py's output (manifest.json + weights.codes.bin +
weights.scales.bin + weights.norms.bin), by default
technologies-moe/alibi-ai's third_party/webgpu-weights/sarashina2.2-1b-alibi-v1-q8/
(override with ALIBI_WEIGHTS_DIR). No Range support — Content-Length is
always set, since the page's progress bar reads it while streaming the ~1.4
GiB weights.codes.bin.
Persistent weight cache (issue #121)
loadWeightsQ8FromUrl (llm/browser-weights.ts) caches the checkpoint into
IndexedDB by default — the demo above (and technologies-moe/alibi-ai's own
integration, this issue's parent #96) no longer re-fetches ~1.4 GiB on every
visit. No demo-side code changed to get this: caching is on unless the fifth
argument's enabled is false, and every one of the demo's existing calls
already passes fewer than five arguments.
- Versioning: the cache key includes a SHA-256 hash of
manifest.json's raw bytes (`llm/weight-cache.ts#s
