Crate inferencelayer

Expand description

Cross-vendor LFM2 inference on wgpu + hand-written WGSL compute kernels.

One WGSL source compiles to Metal, Vulkan, DX12, GL and WebGPU, so this runs on any GPU — the “all GPU cards” requirement that candle’s Metal/CUDA-only kernels can’t meet. The architecture mirrors the webml-community/lfm2-webgpu-kernels demo (Q4_0 dequant-in-GEMV, fused short-conv, flash-decode attention, on-GPU argmax) that reaches ~1400 tok/s; the candle LFM2 backend is the byte-exact correctness oracle every kernel is verified against.

This module is the foundation: a headless compute GpuCtx + a naive GEMV kernel, both verified before the fused/quantised kernels land on top.

Re-exports§

pub use embed_engine::EmbedEngine;
pub use encoder::EncKernels;
pub use encoder::EncoderGpu;
pub use encoder_cpu::CpuEncoder;
pub use encoder_weights::Act;
pub use encoder_weights::EncArch;
pub use encoder_weights::EncBatch;
pub use encoder_weights::EncoderConfig;
pub use encoder_weights::MaskKind;
pub use encoder_weights::MlpKind;
pub use encoder_weights::NormKind;
pub use encoder_weights::PosKind;
pub use forward::BatchCol;
pub use forward::BatchPlan;
pub use forward::EngineOpts;
pub use forward::Lfm2Gpu;
pub use forward::MtpEngine;
pub use forward::SpecGrammar;
pub use forward::SpecPlan;
pub use forward::StageBatchOut;
pub use forward::StageOut;
pub use pooling::EmbedOut;
pub use pooling::Pooling;
pub use replica::FleetPlan;
pub use replica::Replica;
pub use replica::ReplicaPolicy;
pub use replica::max_replicas;
pub use replica::plan_replicas;
pub use sampling::prompt_lookup_drafts;
pub use server::Emission;
pub use server::FinishReason;
pub use server::KvPool;
pub use server::RadixIndex;
pub use server::RequestParams;
pub use server::SamplingParams;
pub use server::Scheduler;
pub use server::ServeStats;
pub use server::generate_once;
pub use shard::Fleet;
pub use shard::FleetRegistry;
pub use shard::LoadAck;
pub use shard::ShardClient;
pub use shard::ShardedPipeline;
pub use shard::StageInfo;
pub use shard::WorkerOptions;
pub use shard::accept_fleet;
pub use shard::decode_greedy_mtp;
pub use shard::decode_greedy_mtp_duo;
pub use shard::decode_greedy_mtp_multi;
pub use shard::plan_split;
pub use shard::run_worker;
pub use shard::validate_chain;
pub use shard::webrtc_chain;
pub use shard::webrtc_pair;
pub use shard_serve::PipelineServe;
pub use shard_serve::ReplicatedServe;
pub use shard_serve::StageStat;
pub use vocab::ByteVocab;
pub use vocab::build_vocab_blob;
pub use weights::Layer;
pub use weights::Lfm2Config;
pub use weights::ModelKind;
pub use weights::Op;
pub use weights::Weights;
pub use weights::detect_model_kind;

Modules§

ark_asr: ARK-ASR-3B — the audio half of the current Open ASR leaderboard leader (4.76% mean WER).
ast: AST — the Audio Spectrogram Transformer (MIT/ast-finetuned-audioset), a vanilla ViT over a log-mel spectrogram for audio classification. Pure Rust, CPU.
bench_guard: The quiet-machine guard + pins ledger shared by every perf harness (scoreboard, encoder_scoreboard, enc-kernel-lab).
chatterbox: Chatterbox (Resemble AI) zero-shot voice cloning — Phase 0 scaffolding.
conv2d: 2D convolution foundation for the Docling ports (RT-DETRv2 layout + TableFormer) — CPU.
cpu_gemm: Prepacked f32 GEMM for the encoder’s linear layers: C[M,N] = A[M,K] · W[N,K]ᵀ (+ bias).
cv_resize: cv2-exact resize kernels for the TableFormer prep chain — CPU, encoder-cpu-gated.
deepencoder: DeepSeek-OCR DeepEncoder — the optical-compression vision encoder, as a CPU reference.
deepencoder_gpu: DeepEncoder GPU (wgpu/WGSL) path — numerically ISO to the CPU reference in crate::deepencoder.
deltanet: Gated DeltaNet (linear attention) — the recurrent primitive of Qwen3.5/3.6-MoE hybrid layers (30 of 40 layers in Qwen3.6-35B-A3B), verbatim-mirrored from transformers modeling_qwen3_5_moe.py (torch_recurrent_gated_delta_rule + Qwen3_5MoeGatedDeltaNet + Qwen3_5MoeRMSNormGated + torch_causal_conv1d_update), which follows Yang et al., “Gated Delta Networks” (arXiv 2412.06464).
diarize: Speaker diarization + voiceprint embeddings — pyannote segmentation-3.0 + NVIDIA TitaNet-large, with sherpa-onnx’s exact pipeline around them. CPU, f32.
diffusion_gemma: DiffusionGemma (google/diffusiongemma-26B-A4B-it) — Google’s block-diffusion Gemma: an encoder-decoder where the ENCODER is a causal Gemma-4 prefill producing a KV cache and the DECODER iteratively denoises a fixed-length token “canvas” with BIDIRECTIONAL self-attention over [encoder KV cache | canvas]. This module is the CPU f32 implementation, parity-gated against transformers.DiffusionGemmaDecoderModel (eager, f32) via scripts/export_diffusion_gemma_ref.py + tests/diffusion_gemma_parity.rs.
diffusion_gemma_gpu: DiffusionGemma’s CANVAS forward on wgpu (Metal / Vulkan / DX12 / WebGPU). The denoising decoder forward is where generation time goes — max_denoising_steps (default 48) decoder passes per canvas vs ONE encoder pass — so this module puts the canvas forward GPU-resident and leaves the encoder prefill and the (O(T·V), trivially cheap) sampler on CPU. The CPU DgDecoder is the parity oracle (tests/diffusion_gemma_parity.rs gates GPU vs CPU, and CPU is itself gated against transformers eager).
embed_engine: EmbedEngine — the device-agnostic embedding entry point: GPU when a wgpu adapter exists, the native CPU encoder otherwise.
encoder: Encoder GPU kernels (WGSL) — the stateless full-sequence compute path for embedding models.
encoder_cpu: CPU encoder runtime (feature encoder-cpu, native-only): the no-GPU fallback AND the parity oracle every encoder GPU kernel is verified against.
encoder_weights: Encoder-model configuration: config.json → the declarative fields the encoder executor branches on.
forward: Full LFM2 forward in WGSL — correctness-first (naive kernels, f32, one dispatch per op), verified against the CPU crate::reference::RefModel oracle. Optimisation (Q4_0 dequant-GEMV, fused norm+matmul, flash-decode attention, on-GPU argmax, one-sync/token) layers on top once correct.
gguf: GGUF container reading + the BINARY-weight (“Q1”) pack — the Bonsai/BitNet class.
gguf_write: GGUF WRITING + the native quantizer family — artifact PRODUCTION with zero external tools.
gliner: GLiNER — zero-shot span NER on the engine’s DeBERTa backbone.
gliner2: gliner2 (fastino/gliner2-multi-v1) — zero-shot RELATION extraction on the engine.
glinerrelex: GLiNER-relex — joint single-pass NER + relation extraction (knowledgator/gliner-relex-large-v1.0).
glinker: GLinker — BiEncoder token-level entity linking (knowledgator/gliner-linker-large-v1.0).
grammar: General structured-output grammar: a schema → flat-table → byte-FSM pipeline for GPU-resident constrained decoding (Workstream G).
iq_tables: IQ codebook grids — GENERATED by scripts/make_iq_tables.py from gguf-python’s packed tables (mirrors llama.cpp ggml-quants, MIT). Do not edit by hand. Row layout: [n_entries][elems_per_entry] flattened; indices select whole rows.
kld: The quantization RULER: KL divergence between a reference model’s next-token distribution and a candidate arm’s, accumulated over a calibration corpus.
kokoro: Kokoro-82M (hexgrad/Kokoro-82M, StyleTTS2/ISTFTNet lineage) — phonemes + voice style → 24 kHz waveform. Pure Rust, CPU. The engine’s third TTS architecture and its most composite:
mimi: Mimi — Kyutai’s streaming neural audio codec (Moshi’s tokenizer): 24 kHz mono ↔ 12.5 Hz × 8 codebooks (1 semantic VQ + 7-level acoustic RVQ, 2048 bins each). CPU f32, ported line-for-line from the official moshi package and parity-gated frame-exact via tests/fixtures/export_mimi.py / tests/mimi_parity.rs.
mimi_gpu: Mimi on wgpu — the full streaming codec (SEANet encoder/decoder, both 8-layer transformers, down/upsample, split RVQ) as TWO static dispatch plans, one per direction. The CPU port (mimi.rs, itself bit-exact vs the reference) is the oracle; this port exists because the CPU codec’s frame time tracks box load (9 ms/frame quiet, 190 ms at loadavg 88) while the GPU path is load-immune — the serve pipeline’s actual bottleneck.
moshi_gpu: Moshi’s depth transformer (“depformer”) on wgpu — the last CPU-resident half of the LM step. The 8 codebook steps are SEQUENTIAL (each consumes the previous step’s sampled token), which forbids per-step readbacks: the whole cycle is a STATIC plan in one compute pass with GPU-side token feedback — per-step argmax writes into a 9-slot token buffer (slot 0 = the text token, copied from the temporal argmax; slots 1..=8 = the audio codebooks) and the next step’s embedding gather reads it. One submit for temporal + text argmax + depformer; ONE 36-byte readback per 80 ms frame.
moshi_lm: Moshi 7B (moshiko) — the RQ-Transformer LM of the full-duplex speech stack: a 32-layer d4096 temporal decoder over 17 token streams (1 text + 8 generated-audio + 8 user-audio), plus the 6-layer d1024 depth transformer (“depformer”) that emits the 8 audio codebooks of one 80 ms frame sequentially. CPU f32, ported line-for-line from the official moshi package (models/lm.py LMModel/LMGen), parity-gated per stage and on a teacher-forced greedy trace via tests/fixtures/export_moshi_lm.py / tests/moshi_lm_parity.rs.
parakeet: Parakeet-TDT ASR (mlx-community/parakeet-tdt-0.6b-v3) — FastConformer encoder + a token-and-duration transducer, ported operation-for-operation from parakeet-mlx (the package the Python parakeet-asr service runs). CPU, f32.
parakeet_gpu: Parakeet’s FastConformer encoder on wgpu (Metal / Vulkan / DX12 / WebGPU) — the path that lets the ASR beat the MLX-on-GPU reference instead of trailing it on the CPU.
pocket_tts: Pocket TTS — Kyutai’s 100M-parameter CPU text-to-speech (kyutai/pocket-tts), ported to the engine as a pure-Rust CPU path. Text + a voice → 24 kHz waveform at ~12.5 latents/s.
pocket_tts_gpu: Pocket TTS’s codec on wgpu — the decode half (latent → 1920 samples) as ONE static dispatch plan, replayed per frame into a single command encoder.
pooling: Pooling & embedding-output semantics shared by the GPU encoder and the CPU oracle/runtime.
qwen3tts: Qwen3-TTS-12Hz (Alibaba, Apache 2.0) — Phase 0 scaffolding.
qwen3tts_codec: Qwen3-TTS-Tokenizer-12Hz codec — P4: the DECODER half (16 codes/frame → 1920 samples @ 24 kHz).
qwen3tts_gpu: Qwen3-TTS talker on wgpu (Metal / Vulkan / DX12 / WebGPU) — the AR backbone’s step on the GPU, the path to real-time synthesis (P9b). The CPU crate::qwen3tts::Talker stack stays the parity oracle; this module is gated against it (p9_gpu_talker_matches_cpu).
qwen3tts_spk: Qwen3-TTS speaker encoder — P6: reference audio → the x-vector that conditions the talker.
realtime: Realtime input discipline for the duplex loop, extracted so it is testable in isolation and shared with the server binary.
reference: CPU f32 reference LFM2 forward: the exact spec the WGSL kernels are verified against, and itself checked against the transformers gold continuation. Mirrors the candle LFM2 backend op-for-op (ShortConv ring-buffer, GQA attention with per-head QK-RMSNorm + RoPE, SwiGLU MLP, tied lm_head).
reference_qwen35: CPU f32 reference for Qwen3.5 hybrid decode (the arch behind NuExtract-3, Qwen3_5ForConditionalGeneration): 24 gated-DeltaNet linear-attention layers + 8 full-attention layers (every 4th), SwiGLU MLP, tied lm_head. This is a pure-CPU decoder — no wgpu/Vulkan — and it is the spec the GPU kernels are checked against, itself validated op-for-op against the transformers gold continuation (see bin/qwen35-cpu).
replica: Fleet REPLICATION planning — how to allocate a heterogeneous worker pool between pipeline depth (split one model instance across stages, to fit a bigger model) and replication (run several independent pipelines, for more throughput).
rtdetr: RT-DETRv2 layout detector (docling-layout-heron) — CPU, pure Rust, stage-gated vs transformers.
rtdetr_gpu: RT-DETRv2 backbone + CCFM on wgpu — the measured answer to “win the layout benchmark”.
sampling: Wasm-safe CPU sampling: penalties, logit bias, truncation (top-k / top-p / min-p), and per-token logprobs — pure functions over a logits slice with no GPU and no server-feature dependency, so the browser/wasm build and the native server share exactly this code and every rule is unit-tested on synthetic logits.
server: Multi-user serving core: paged KV-block allocator, radix prefix cache, and the continuous-batching scheduler that drives crate::forward::BatchPlan.
shard: P2P pipeline-parallel sharding: pool VRAM across devices (Macs, V100s, anything wgpu reaches) by giving each device a contiguous layer range and streaming the residual stream between them.
shard_serve: Micro-batched pipeline-parallel serving: the continuous-batching scheduler wired onto the shard chain. This is what turns a p-stage pipeline from 1/p utilization (M=1: one token walks the chain while p−1 devices idle) into a full pipeline: requests are partitioned into micro-batch groups (one per stage) and group g runs on stage s while group g−1 runs on stage s+1 — every device busy every tick.
siglip_vision: VisionEmbed — the device-agnostic SigLIP vision entry point: GPU transformer layers when a wgpu adapter exists, the native CPU encoder otherwise. Either way, patch embedding and the MAP attention-pooling head run on the CPU encoder (shared, oracle-identical code), so ONLY the 12 transformer layers move to the GPU — which is exactly the stage the parity gate covers.
simd: Hand-vectorized f32 primitives, for the kernels where the compiler will not do it for us.
simd_math: Deterministic vectorized transcendentals for the CPU encoder: exp, erf (A&S 7.1.26 — the SAME approximation the GPU’s ACT_FNS uses), and the activation/softmax passes built on them.
speecht5: SpeechT5 TTS (microsoft/speecht5_tts) — text → mel spectrogram → waveform. Pure Rust, CPU.
tableformer: TableFormer (docling-models v2.3.0, accurate) — CPU, pure Rust, stage-gated vs torch.
turboquant: TurboQuant KV-cache quantization (Zandieh–Daliri–Hadian–Mirrokni, arXiv 2504.19874): calibration-free, data-oblivious vector quantization built from (1) a fixed random rotation Π (orthogonal — inner products are preserved in the rotated domain, so attention never un-rotates keys: it rotates the QUERY once and dots against quantized coordinates), and (2) a per-coordinate Lloyd–Max quantizer whose centroids are precomputed ONCE on the analytic high-dimensional coordinate density (≈ N(0, 1/d) for unit vectors) — no calibration data.
vision: Qwen3.5-VL vision tower — native-resolution ViT, on the CPU, in pure Rust.
vision_glm: GLM-OCR vision tower (GlmOcrForConditionalGeneration.model.visual, the GLM-4.6V family) — native-resolution ViT on the CPU, in pure Rust.
vision_glm_gpu: GLM-OCR vision tower on the GPU (wgpu/WGSL) — numerically ISO to the CPU tower in crate::vision_glm.
vision_gpu: The Qwen3.5-VL vision tower on wgpu.
vits: VITS TTS (facebook/mms-tts-eng — MMS English) — text → waveform, end-to-end. Pure Rust, CPU.
vocab: Donor-side token→text display support. In pipeline sharding every donor already RECEIVES the token ids it processes (BatchCol.token) and, on the last stage, the ids it GENERATES — but a donor has no tokenizer (the browser build drops it; onig_sys can’t target wasm). Decoding ids→text, however, needs neither the tokenizer nor onig: it only needs a compact id→bytes table plus concatenation. The coordinator (which HAS the tokenizer) builds that table once from tokenizer.json (build_vocab_blob) and ships it; each donor decodes with ByteVocab.
weights: Model weight loading: safetensors (BF16 or F32 on disk) → Q4/f32 GPU buffers, organised per layer. Two architectures share the loader: LFM2 (conv/attention hybrid) and Gemma-3 text (uniform sliding/full attention with sandwich norms, GeGLU, and the (1+w) RMSNorm convention — folded into the buffers at load so the kernels stay convention-free).
whisper: OpenAI Whisper — speech-to-text, CPU f32 (the parity oracle and the portable fallback; the wgpu path reuses the encoder’s kernels, since Whisper’s encoder is a VANILLA bidirectional transformer — plain scaled-dot attention + GELU MLP, no rel-pos, no conv module).
whisper_gpu: Whisper’s encoder on wgpu (Metal / Vulkan / DX12 / WebGPU). Whisper is a VANILLA bidirectional transformer, so this is simpler than the conformer: the FFN/projection GEMMs + layernorms reuse the encoder’s kernels, and the only per-block custom kernel is plain scaled-dot attention (no rel-pos term — see crate::parakeet_gpu for the harder conformer variant). The conv stem stays on the CPU (runs once); the N transformer blocks run GPU-resident in one submit. The CPU WhisperEncoder is the parity oracle.

Structs§

ConvFusedF16: Fused conv pipeline (in_proj + norm + gating + depthwise conv).
GpuCtx: A headless wgpu compute context (device + queue) with buffer up/download helpers. No surface, so it initialises on any backend wgpu supports.
Q4LmHead: Tiled Q4 lm_head pipeline (8 rows/workgroup).
Q4LmHeadSparse: Sparse Q4 lm_head pipeline (see [q4_lmhead_sparse_src]); rows per workgroup = 8, matching the SPARSE_ARGS kernel’s indirect-workgroup computation.

Functions§

f16_lmhead_lcpp_src: gemv_f16_k_lcpp_src’s LM-head form, derived by the SAME mechanical transform q1_lmhead_lcpp_src applies to its base: the head runs at vocab-scale M, so its row grid exceeds the 32768 per-dimension dispatch cap and needs ROW BANDS folded into gy (gy = gy_rows·ncols, and the kernel decodes wid.y as (col, band)). The accumulate arm is dropped — the head always overwrites.
gemm_f16_src: gemm_q4_src’s f16-STORAGE twin: same tiles, weights read as plain f16 (the same buffer layout gemv_f16_k_lcpp_src binds — 8 halves per vec4<u32>, 4 per 32-element block).
gemm_q1_src: Tiled Q4 GEMM for PREFILL widths — the real weight-shared shape the wide plan needed (the KC16 tweaks measured single-digit GB/s; per-column lcpp re-reads weights ncols×). Classic blocked GEMM: BN=32 columns per tile share one weight stream (weights read ⌈ncols/32⌉× per step instead of ncols×), BM=64 rows per workgroup, BK=32 (exactly one Q4 block), the activation tile staged in 4 KiB of shared memory. 256 threads as a 16×16 grid, each owning a 4-row × 2-column register tile. Bind layout identical to the gemv family (drop-in for the gemv_k slot); dims.z = 1 accumulates into y (the out-proj epilogue). Pure baseline WGSL — no subgroup ops, portable to every adapter. gemm_q4_src’s BINARY-weight twin (Q1/Bonsai): the same BM=32 x BN=16 tile — the two-phase coalesced x staging, the once-per-WG weight staging, the fully-unrolled inner product — with the Q4 nibble decode swapped for the Q1 sign decode (one u32 of sign bits + one f16 scale per 128-weight superblock). This is what makes Q1 WIDE plans weight-SHARED: per step the weights stream ceil(k/16) times instead of k times, which is the whole prefill wall on a 27B (measured 65 tok/s per-column vs the sweep bound). Body is GENERATED like the Q4 original (“fully-unrolled”: dynamically-indexed private arrays land in DRAM-backed Naga scratch — the measured 4.2x lesson).
gemm_q1_src_chunk: Chunked staging variants: stage chunk 32-blocks per barrier round. The ablation on the V100S attributed 52% of the kernel to the staging skeleton (2 barriers per 32-block = 320 barriers/WG at n=5120); chunk=4 stages one whole 128-weight SUPERBLOCK per round (4x fewer barriers) and hoists the block scale to ONE multiply per superblock (it is mathematically per-superblock — decode was 38%, and a third of that was the per-group scale mul).
gemm_q1_xt_rp4_src
gemm_q1_xt_src: gemm_q1_src reading PRE-TRANSPOSED activations: global x layout = [n elems][ncols/4] vec4 (columns packed) — exactly the shared-tile layout, so staging is ONE coalesced copy + ONE barrier and the dynamic-[comp] transpose select-chains are gone. Ablation on the V100S attributed 52% of the kernel to the staging skeleton; this variant measured 3.83 -> 5.07 TFLOPS (1.32x, same-run, parity-exact — FP order identical, so batch results stay bitwise). Pair with transpose_cols_src on each producer output (~20 us per input vs ~150 ms saved per 64-col wide step on the 27B).
gemm_q4_splitk_reduce_src: The split-K fold: y[i] = (acc ? y[i] : 0) + sum_s part[s][i], slices in fixed ascending order — the same result every run, independent of workgroup scheduling.
gemm_q4_splitk_src: SPLIT-K twin of gemm_q4_src for the SKINNY (ncols <= 64) serving GEMMs: gy carries ncol_tiles x S — each workgroup covers one k-slice of the block walk and writes an f32 partial [S, ncols, m]; gemm_q4_splitk_reduce_src folds the slices in FIXED ascending order (deterministic) and applies the acc flag. Generated by string-transforming the production kernel, so the per-slice dequant + FMA order is byte-identical to gemm_q4_src — only the summation grouping changes (NOT bitwise vs the serial walk; opt-in at the plan level).
gemm_q4_src
gemm_q8_0n_src: gemm_q4_src’s native-Q8_0 twin: padded 9-word blocks (f16 d + 32 i8), the layout gemv_q8_0n_k_lcpp_src and the Q8 policy’s pack_q8_0n produce.
gemv_f16_k_lcpp_src
gemv_iq4_nl_k_lcpp_src: Native IQ4_NL GEMV — padded 5-word blocks, the 16-entry nonlinear codebook applied in-shader. Same lcpp discipline as the rest of the family (WG==subgroup==32, 4 rows/WG, subgroupAdd fold, 4 bindings). The codebook lookups go through ONE helper call site — naga’s MSL backend names reinterpret temps after operand expression handles, and repeated intrinsic calls on the same operand collide (the dp4a lesson).
gemv_iq4_xs_k_lcpp_src: Native IQ4_XS GEMV — 34-word superblocks as-is (136 bytes, naturally aligned): f16 d + packed 6-bit (−32-biased) sub-scales + 8 sub-blocks of codebook nibbles.
gemv_iq_grid_src: Generic IQ superblock pad: [f16 d | 2B pad | payload…] to a word-aligned stride. Used by every 2-byte-d-first IQ layout (XXS/XS/S families and IQ1_S). Value-lossless. Native GEMV for the grid-codebook IQ family (IQ1/IQ2/IQ3): the GRID rides binding 0 (the scales slot — same 5-binding shape as Q4_0), padded raw superblocks ride binding 1. One emitter, per-type inner body; every kernel keeps the lcpp shell (WG==subgroup==32, 4 rows/WG, subgroupAdd). ksigns is COMPUTED (i | parity<<7); grid rows unpack via ONE helper call site per operand shape (the naga-MSL name-collision rule).
gemv_q1_f16x_src: gemv_q1_k_lcpp_src reading PACKED-f16 activations: x bound as vec4<u32> (8 f16 per 16-byte load) and unpacked with unpack2x16float — HALF the x-load instructions per 32-block (4 vs 8). The decode GEMV ablation put x-loads at 41% of the kernel; measured same-run on all five Bonsai decode shapes: 1.30-1.50x (e.g. qkv 38.9 -> 26.9 us, w2 106 -> 82 us; 117 -> 170 GB/s weight bandwidth). Accumulation stays f32; pairing with pack_f16_src on each producer output costs ~3 us/input at decode widths.
gemv_q1_k_lcpp_src: gemv_q4_k_lcpp_src’s BINARY-weight twin (the Bonsai/BitNet “Q1” class: sign bits with per-128-weight f16 scales — 1.125 bits/weight on disk, 4.5 bytes per 32-weight block here vs Q4_0’s 18). Same lcpp discipline: WG = one 32-lane subgroup, 4 rows per workgroup, each lane owns k-blocks strided by 32, subgroupAdd(vec4) folds the lanes. The weight word’s bit i is x-element i’s sign (1 → +1, 0 → −1 — the convention the Q1_0 GGUF carries); decode is a select, no multiplies on the weight side. At the decode bandwidth wall the byte ratio IS the speed ratio — tests/q1_kernel.rs holds the parity gate and the bench.
gemv_q1_rp4_f16x_src: gemv_q1_f16x_src on the ROWPACK4 layout: the 4 rows’ sign words arrive as ONE vec4<u32> load and their scales as ONE vec4<f32> load — the per-lane serial load chain drops from 16 instructions per 32-block to 6 (with the packed-f16 x). Measured COLD (rotating weight regions, the regime production actually runs in): 1.55-1.66x over the shipped path, weight bandwidth 104 -> 219 GB/s (tests/gemv_cold_bench.rs).
gemv_q4_k_lcpp_src: llama.cpp-shaped GEMV (“lcpp”, lab-measured 3-7.4× over the tree/nbar shapes on V100/Vulkan at 1-16 columns): WG == subgroup == 32 lanes, FOUR rows per workgroup sharing every activation load, subgroupAdd reduction — zero barriers anywhere. Columns ride gy (weights re-read per column; still 761 GB/s effective at 16 columns — the wall-clock win over the weight-shared shapes is what matters). fp32 accumulation. Requires subgroup support.
gemv_q4_k_sg_solo_src: SOLO (ncols == 1) twin of the subgroup Q4 GEMV: same math, no KC column loop and no array<f32, KC> accumulator. That array is dynamically indexed, so Metal spills it to scratch and every acc[cc] += … costs a memory round trip — measured 1.22-1.26x on all three Moshi shapes (tests/moshi_kernel_bw.rs). The realtime decode plan is single-column; batch and prefill plans keep the KC kernel. Per-row reduction order is unchanged.
gemv_q4_k_sg_src: Subgroup GEMV (v3): one SUBGROUP per row, lanes stride blocks, subgroupAdd reduction — zero barriers, zero shared memory. Column loop guarded by ncols, so M=1 pays one column.
gemv_q4_k_src: Batched [gemv_q4_src]: grid (⌈M/NR⌉, k). Same NR=2 body; col = wid.y.
gemv_q4k_k_lcpp_src: gemv_q4_k_lcpp_src’s f16-STORAGE twin: weights held as plain f16, no scale table and no quantization beyond the f16 rounding itself. This is the engine’s REFERENCE weight arm — the one the quantization ruler (crate::kld) measures Q4/Q1 against, because checkpoints ship bf16 and f16 round-trips them at these magnitudes, so “f16 vs Q4” isolates what Q4 costs instead of comparing two lossy schemes to each other.
gemv_q5_0n_k_lcpp_src: Native Q5_0 GEMV — the legacy 32-block type llama-quantize FALLS BACK to when a K-quant cannot fit (ne0 % 256 != 0, the norm on small-hidden models like SmolLM2). Disk blocks are 22 bytes (unaligned); the loader pads to 24 bytes (6 words): word0 = f16 d (+pad), word1 = qh (the 5th bits), words 2-5 = 16 nibble bytes. w = d·((nib | bit<<4) − 16).
gemv_q5k_k_lcpp_src: Native Q5_K GEMV — gemv_q4k_k_lcpp_src plus the 5th-bit plane. 176-byte superblocks are already word-aligned (44 words): f16 d+dmin (word 0), 12 scale bytes (words 1-3, the SAME 6-bit pack as Q4_K), 32 qh bytes (words 4-11), 128 nibble bytes (words 12-43). Sub-block s takes bit s of qh[l]; w = d·sc·(nib | bit<<4) − dmin·mn.
gemv_q6k_k_lcpp_src: gemv_q4k_k_lcpp_src’s Q6_K sibling: native 6-bit superblocks decoded in-shader. The disk block is 210 bytes — NOT word-aligned — so the loader pads each block to 212 bytes (53 words); the pad is dead space (+1% bytes, lossless) and buys aligned vec4 nibble tricks instead of per-byte gymnastics. Layout per padded block: 128 ql bytes (words 0-31), 64 qh bytes (words 32-47), 16 signed scales (words 48-51), f16 d (word 52 low half).
gemv_q8_0n_k_lcpp_src: Native Q8_0 GEMV — llama.cpp’s 34-byte blocks (f16 d + 32 i8) padded to 36 bytes (9 words): word0 = d (+pad), words 1-8 = the signed bytes. w = d·q. (The engine’s OWN planar Q8 kernels are Moshi-scoped and use a different layout; this one reads the GGUF bytes as-is, so an artifact’s Q8_0 tensors serve without any repacking.)
gemv_q8_k_sg_solo_src: Q8_0 decode GEMV (subgroup): biased-u8 quants (q+128), TWO vec4<u32> per 32-block — the Moshi FFN band (4-bit destroys its gating matrices; see moshi_lm’s q4sim sweep). SOLO (ncols == 1) twin of gemv_q8_k_sg_src: identical math, but the KC column loop and its array<f32, KC> accumulator are gone. A dynamically indexed array does NOT live in registers on Metal — it spills to scratch, and every acc[cc] += … becomes a memory round trip. Stripping it measured 0.29 → 0.15 ms on 11264x4096 (the roofline variant of tests/gemv_roofline.rs runs at exactly this speed, which is how the cost was found). The realtime decode path is always single-column, so this is what it dispatches; the KC kernel stays for batch/prefill plans. Reduction order per row is unchanged.
gemv_q8_k_sg_src
gemv_q8_k_src: Portable (non-subgroup) Q8_0 decode GEMV — the fallback twin of gemv_q8_k_sg_src for adapters without guaranteed subgroups (GL-class, or OSFKB_NO_SUBGROUPS). 16 rows/wg × 16 software-lanes/row with a workgroup-tree reduction, matching gemv_q4_k_src’s tiling so it rides the same !sg dispatch (gemv_rows_per_wg = 16). Same bindings as the subgroup kernel; the FP reduction order differs, so it is a QUALITY band, not bitwise.
iq4_nl_pad_blocks: Pad raw 22-byte Q5_0 blocks to the 24-byte layout gemv_q5_0n_k_lcpp_src reads: word0 = d (+2 pad bytes), word1 = qh, words 2-5 = qs. Pure re-layout. IQ4_NL 18-byte blocks -> 20-byte (5-word) kernel layout: word0 = f16 d (+2B pad), words 1-4 = the 16 codebook-nibble bytes. Value-lossless rearrangement.
iq_grid_geom: (padded words, grid table, block bytes) for a grid-IQ ggml type — the loader’s one-stop map.
iq_pad_blocks
load_encoder_tokenizer: Load a checkpoint’s tokenizer.json SANITIZED for the ragged encoder: any baked-in padding is REMOVED and truncation is pinned to max_seq.
mlp_gate_f16_lcpp_src: [mlp_gate_q4_lcpp_src]’s f16-STORAGE twin: the SAME fused gate+up discipline (one WG = 32-lane subgroup, 4 rows, pre-normed activation from the lcpp split-rmsnorm, subgroupAdd fold), with both weight streams read as plain f16 instead of Q4 nibbles + scales.
mlp_gate_q1_f16x_src: [mlp_gate_q4_lcpp_src]’s BINARY-weight twin (the Q1 class): the SAME fused gate+up discipline (split-rmsnorm feeds pre-normed x; one WG = 32-lane subgroup, 4 rows), with the two weight streams decoded as signs instead of Q4 nibbles. Bindings are identical to the Q4 kernel (s1, b1, s3, b3, x, y, dims, epsm), so the plan swaps ONLY the pipeline. mlp_gate_q1_lcpp_src reading PACKED-f16 activations (same transform as gemv_q1_f16x_src): x bound as vec4<u32>, unpacked with unpack2x16float — half the x-load instructions, shared across BOTH weight streams. Derived from the f32 source by mechanical substitution so the two can never drift.
mlp_gate_q1_lcpp_src
mlp_gate_q1_rp4_f16x_src
mlp_gate_q8_0n_lcpp_src: mlp_gate_f16_lcpp_src’s native-Q8_0 twin: the fused gate+up MLP with BOTH weight streams read as padded 9-word Q8_0 blocks. In a native-GGUF model the gate/up pair is served Q8_0N even when the artifact stored Q5_0 — the upcast (d8 = d5, q8 = q5 − 16) is EXACT, and one fused kernel beats carrying a per-legacy-type MLP zoo. Same six bindings as the f16 twin (w1, w3, x, y, dims, epsm); the act(gate)·up epilogue is byte-identical to the Q4 kernel’s.
pack_f16_src: Pack an f32 vector into f16 pairs (pack2x16float) for gemv_q1_f16x_src. One thread per output word; widths are always even (multiples of 32). ~3 us at decode sizes.
q1_lmhead_lcpp_src: llama.cpp-shaped lm_head (WG=32=subgroup, 4 rows sharing activation loads, subgroupAdd, zero barriers; head dims layout (M, N, ncols, gy_rows), gy = ncols). Lab + serving measurements put this 2.4-6× over both the nbar and tiled shapes at 1-8 columns on V100/Vulkan — the small-width band both the MTP draft chain and the M=1 verify live in. [q4_lmhead_lcpp_src]’s BINARY-weight twin: the same vocab-scale row-band folding over the Q1 GEMV, so the head runs the Bonsai/BitNet sign weights (OSFKB_Q1_WEIGHTS). Same binding layout as the Q4 head (scales, bits, x, y, dims), so Q4LmHead’s geometry is unchanged.
q1_rp4_repack_src: Repack a Q1 weight into the ROWPACK4 layout gemv_q1_rp4_f16x_src reads: the sign words of each 4-row group as one vec4<u32> [m/4, n/32] and their scales as one vec4<f32> [m/4, n/128]. One-time GPU pass per matrix (~1 ms for the largest).
q4_0_to_q8_0n: Losslessly upcast raw Q4_0 blocks (18 B) to the padded Q8_0N layout: d8 = d4, q8 = nib − 8. Exact for the same reason as the Q5_0 upcast.
q5_0_pad_blocks
q5_0_to_q8_0n: Losslessly upcast raw Q5_0 blocks (22 B) to the padded Q8_0N layout (36 B): d8 = d5, q8 = (nib | bit<<4) − 16. Every representable Q5_0 value is exactly representable — the scale is copied bit-for-bit and the integer fits i8 — so this is a RE-ENCODING, not a requantization. Used when a concatenated site (q|k|v, gate|up) mixes legacy types.
q6k_pad_blocks: Pad raw Q6_K blocks (210 B) to the 212-byte (53-word) layout gemv_q6k_k_lcpp_src reads. Pure re-layout — the two pad bytes are never read.
q8_0_pad_blocks: Pad raw 34-byte Q8_0 blocks to the 36-byte layout gemv_q8_0n_k_lcpp_src reads: word0 = d (+2 pad bytes), words 1-8 = the 32 signed bytes. Pure re-layout.
q8_0n_lmhead_lcpp_src: gemv_q8_0n_k_lcpp_src’s LM-head form — the SAME mechanical band-fold transform f16_lmhead_lcpp_src applies to its base (vocab-scale row grids exceed the 32768 dispatch cap; accumulate arm dropped). Anchors asserted so the two sources can never drift.
test_head_kc4: Test hook: the KC4 tiled lm_head source (head-gate comparisons).
test_head_kc16: Test hook: the KC16 tiled lm_head source.
test_kernel_kc4: Test hook: the KC4 subgroup GEMV source (prefill-gate comparisons).
test_kernel_kc16: Test hook: the KC16 prefill GEMV source.
transpose_cols_src: Transpose [ncols, n] row-major f32 into the [n][ncols/4] column-packed vec4 layout gemm_q1_xt_src reads. Thread = one output vec4; tid ordered so global READS coalesce along e. Total traffic 2x the tile (~2.5 MB at 64x5120) — ~20 us, amortized across every GEMM consuming the same input.