Repository navigation
lean: faster prompt reading and decode, GPU-resident sampling, more weight formats, transformers-exact repetition penalty, no WGSL overrides - #8
Merged
Conversation
…er token forward_prefill_suffix (used whenever part of the prompt is already in the KV cache: every follow-up turn, and any prompt that shares a few leading tokens with the previous one) ran the new tokens one by one through the decode path, about one decode step (20 ms on SmolLM2-1.7B) per token. For head_dim 64 and 128 it now goes through the batched prefill: tiled matmuls, RoPE and the K/V scatter offset by the resident prefix, and attention that reads the cache (attn_prefill_tiled.wgsl, CACHE build). Other head dims keep the per-token path. prefill_sweep --suffix times it and checks it against the full prefill.
…60M (Llama architecture)
…ns (cold and suffix prefill, decode, exactness hash)
…ken read at submit time The decode loop kept one step submitted ahead and asked for step t+1's map only after step t's had resolved. A browser that resolves buffer maps on a coarse schedule (Firefox, about every 100 ms) then delivered one token per resolution, 9 tokens/s on SmolLM2-1.7B against 48 in Chromium. Up to 8 steps are now in flight and each step's map is requested when it is submitted, so one resolution returns every step that already finished. The steps ahead are discarded when the loop stops, as the single step ahead was.
Firefox on Linux (RTX 3080, Vulkan) lost the device with 'Not enough memory left' while a 1 GB model was loading: queue.write_buffer data stays staged until the next submit and the browser cannot hold more than about a hundred MB of it. Writes are now at most 8 MB and a submit with no command buffers goes out whenever 8 MB are pending, so the staged data is released as the load goes. The same sizes on every device. Reproduced with a plain WebGPU page (many 1 MB buffers with a flush every 32 MB fail, every 8 MB pass).
…he report; phases and fresh-engine check
Tensors a quantised GGUF keeps at F32 (SmolLM2's Q4_0 file holds three ffn_down weights as Q4_1, which load as F32) ran prefill through the one-scalar-loop-per-output kernel: three layers took about 40 percent of a 128-token prefill on the M2 (511 ms of 1013 ms in the down projections). They now use the 64x64 register-blocked tile of the Q4_0 prefill kernel with F32 weight tiles (k multiple of 32).
Firefox resolves a buffer map only when a queue submission wakes its poll, so every read back (a sampled token's logits, a prefill's logits) waited about 100 ms, whatever the GPU time: 10 tokens/s with the chat page's sampler on an RTX 3080 against 60 with this change. While a map is pending an empty queue.submit goes out every timer tick; the wait ends as soon as the map resolves, so browsers that resolve on their own are not delayed.
On an RTX 3080 in Firefox the two depths decode at the same speed (220 tokens/s); fewer steps are wasted when generation stops.
Each Engine owns its device. Dropping it now calls device.destroy(), so freeing a Lean releases all its buffers at once instead of when the browser collects the wrapper objects. Adds a bench 'loads' phase (repeated load and free), run-doc note, ENGINE_BUILD browser-perf-03.
Each greedy step appends its argmax id to a GPU block buffer; the host maps one block of GREEDY_BLOCK ids at a time instead of one id per step, with two blocks in flight. Stopping inside a block rolls the KV length back to the step-by-step value. Tests cover stops at and around block boundaries.
…shared by 8 lanes Each key's dot product is spread over 8 adjacent threads reading contiguous vec4 slices instead of one thread walking a whole row. The tile max and sum use broadcast reads of workgroup memory instead of two barrier trees, and the V accumulation is unrolled by four keys. Adds the decode_ctx example (decode speed against context length, logit dump and compare).
The WGSL spec makes a local named dot shadow the builtin, so calling dot() after the declaration is rejected by Tint (Chromium) although naga accepts it.
…nd measured device ceilings
… and matmul kernels Q4_1 blocks (nibble * d + m) were dequantized to F32 at load, 4.8x the bytes read per step. The nibbles repack as for Q4_0 and d and m are stored interleaved in the scales buffer; decode (one and four rows per lane), short and long prefill and the naive kernel gain a Q41 override.
…device grants shader-f16 The adapter's shader-f16 feature is requested whenever it is offered; the F32 tiled kernels stay the path on every other device. LEAN_F16=0 keeps them on an f16 adapter. Adds an ignored test comparing f16 and F32 prefill logits and greedy tokens against each other and the reference fixtures.
accelerate is needed by transformers to load a GGUF for the fixture scripts.
… back per step Repetition penalty (a count per vocabulary id, applied per occurrence), top-k (tree of 1024-item bitonic sorts), temperature, softmax, top-p and the draw run in three kernels at the end of the decode step. The chosen id feeds the next step's embedding on the GPU, so nothing but the id comes back, and two steps stay in flight. The kernels use only multiplications and additions (no division, polynomial exp) and a counter-based generator keyed by seed and token position, and lean_core::resident_sampling holds the same computation on the CPU: the first token (chosen from the prefill logits) and the CPU backend use it, so every backend gives the same reply for the same seed. Ties at the k-th logit are cut at k, lower id first. The random stream differs from sample()'s, which still serves top_k 0 and 64 and up through the full-logits readback. This replaces the candidate readback of the previous commit (host-side penalty, top-p and draw over 64 candidates), which is removed. Bumps ENGINE_BUILD on every loading URL.
…cks of 8 measured no faster
The 64x64 tile uses a 16x8 workgroup in which each thread produces 8 rows x 4 columns (32 accumulator floats), instead of 16x16 threads with 4x4 blocks. A thread issues 12 shared-memory vec4 loads per 4 k for 32 outputs instead of 8 for 16, and the workgroup has half as many threads at each barrier. Staging loops over the tile with the workgroup as stride. All weight formats (Q4_0, Q8_0, Q4_1, Q5_0, Q4_K) and both the f32 and f16 kernels use the new layout; the per-output summation order is unchanged.
…, chat page results
… reduction per row NEON and wasm simd128 row kernel (dot_rows) shared by the one-row dot and a new matvec_q: the activation block is loaded once for four weight rows, each row accumulates block partial sums times the scale in four lanes, and the horizontal sum happens once per row instead of once per 16 weights. f32 activations as before. A tail of one to three rows takes the one-row instance of the same code, so every output has the same bits whichever group or thread computes it. The one-row linear and the 32-column team items call matvec_q; the scalar build keeps its per-row dots.
…ates cpu_lora_parity and lora_parity compare GPU logits against transformers (and PEFT) and against the CPU path at F32 tolerances, but did not set LEAN_F16=0, so on an adapter that grants shader-f16 they ran the f16 prefill matmul. Its rounding moved the answer logits by up to 5.8e-2 (LoRA) and 1.6e-2 (base model) and cpu_lora_parity failed at its 1e-2 clear_lora check. With the F32 kernels the GPU agrees with HF+PEFT to 2.1e-4 and with the CPU base model to 3.4e-5. The f16 path stays held to the F32 path by f16_parity.
…Processor Real next-token logits of SmolLM2-360M and Qwen2.5-0.5B (Q4_0 GGUFs, dequantized by transformers, f32 CPU) after a chat prompt plus a 48-token greedy continuation, with the processor applied over the whole input_ids at penalties 1.1 and 1.3. Stores the raw and processed logit of every distinct history id as f32 bits.
RepetitionPenaltyLogitsProcessor gathers the logits of input_ids, penalizes them and scatters them back, so an id is penalized once however often it occurs. lean applied the penalty once per occurrence on every path: the CPU sampler looped over the history, and the resident sampler (the GPU kernel and its CPU reference, used on every backend) looped count times. A token seen 20 times was divided by 1.1^20 instead of 1.1. The CPU sampler now penalizes the deduplicated history; the GPU kernel and sample_reference apply the penalty once to every id with a nonzero count. The counts buffer and everything after the penalty are unchanged, and greedy decoding never reaches this code. Tests: the CPU sampler and the GPU kernel against the transformers fixture (fixture_rep_penalty.json). The CPU sampler matches bit for bit; the kernel multiplies by 1 / penalty where transformers divides and stays within one unit in the last place.
…mpling, before and after
Every pipeline-overridable constant (HEAD_DIM, CACHE, Q8, Q41, Q50, K4, GATE_UP, BF16, WINDOW) is now a plain `const`, and engine.rs writes the per-pipeline value into the declaration before the shader module is created. No pipeline depends on the implementation evaluating override expressions at pipeline creation; attn_prefill_tiled's DV, RS and UPL, whose initializers referenced another override, were rejected by one WebGPU implementation and left the prefill pipelines invalid. Values are unchanged, so every kernel computes the same thing.
Pipeline creation at engine start runs inside validation and internal error scopes. If the implementation rejects any pipeline, Engine::new returns an error, so the backend ladder falls back to the CPU instead of running a forward pass with invalid pipelines.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Performance, correctness and portability work on lean, measured in Chromium, WebKit and Firefox on an Apple M2 and on an Android phone. Run docs in docs/runs/ give every number.
Performance
shader-f16(greedy tokens identical to the f32 path on every fixture).Formats
Correctness
Portability
overrideconstants: build-time values are written into the shader source before compilation, the same on every device. Safari 26.3 rejected overrides whose default depends on another override, which left the attention pipelines invalid and produced garbage.Output hashes are identical across Chromium, WebKit and Firefox for five models (greedy and sampled), and match the CPU path. Build tag 2026-10-18-lean-01 on every loading URL.