Module encoder 

Encoder GPU kernels (WGSL) — the stateless full-sequence compute path for embedding models.

Distinct from the decoder kernels in forward.rs by design: encoders are single-shot over a known ragged batch (no KV pool, no block tables, no ring buffers), weights are f16/f32 (never Q4 — encoder models are small and precision-sensitive), and attention is a mask MODE (bidirectional / causal / local-window) over exact seq_starts ranges rather than paged causal decode. Every kernel here is parity-tested against the encoder_cpu oracle (tolerance-gated, never bitwise — see tests/encoder_kernels.rs).

Structs

EncKernels

The compiled encoder kernel set. Built once per device; the f16 GEMM twin exists only when the adapter reports SHADER_F16.

EncoderGpu

The GPU encoder: a prebuilt dispatch plan (bind groups + uniforms baked at load) over persistent scratch sized max_tokens, replayed per batch with only the token-count uniform fields and the ragged-batch buffers rewritten. One submit + one readback per encode.

Constants

MAX_VISION_BATCH

Max images per batched SigLIP vision replay (run_layers_gpu_batched). The plan scratch is sized MAX_VISION_BATCH · n_patches; callers chunk larger batches. 32 × 196 patches × 768 hidden ≈ 19 MB per scratch buffer — negligible VRAM, large dispatch-overhead amortization.

Functions

enc_gemm2_src

GEMM v2 — the same contract as [enc_gemm_src] (y[m,n] = act(x[m,k] · Wᵀ + bias), same bindings, same Meta), register-blocked.

enc_gemm3_f16a_src

GEMM v3-f16a — the v3 structure with F16 ARITHMETIC (the P7 contingency): x is converted to f16 ONCE at staging, the weight is stored f16 [k,n]-transposed, and the inner FMAs and accumulators are f16 — Apple’s shader ALUs run f16 at twice the f32 rate, which is the whole remaining edge MPSGraph holds at the small/mid shapes. The EPILOGUE is f32 (bias, activation, store), so downstream buffers keep their type.

enc_gemm3_sk_f16a_src

The f16-ARITHMETIC split-K partial: the same patch over enc_gemm3_f16a_src — its epilogue already stores f32(acc), so the partials stay f32 and the SAME reduce kernel serves both precisions.

enc_gemm3_sk_src

GEMM v3 SPLIT-K — the occupancy answer for the mid-band widths (n = 768/1024-class) at SMALL m, where no tile has enough output blocks to fill the GPU (m=64 × n=768 at (16,64) is 48 workgroups on a ~40-core machine). The reduction dimension supplies the missing parallelism: grid.z chunks each own a contiguous run of k-tiles and write RAW partials (y_part[c][m,n], no bias, no activation); [ENC_GEMM3_SK_REDUCE] then folds the chunks in fixed order and applies the epilogue. Chunk count is num_workgroups.z (one pipeline serves any split), the tile is the same vec4/[k,n]-transposed body as enc_gemm3_src, and a chunk whose window falls past k writes zeros — never stale scratch.

enc_gemm3_src

GEMM v3 — vec4 shared memory, on a PRE-TRANSPOSED weight.

enc_gemm4_f16_src

GEMM v4 — HARDWARE MATRIX UNITS (EXPERIMENTAL_COOPERATIVE_MATRIX).

enc_gemm4_f16w_src

v4-f16 over an f16-STORED [k,n] weight (the f16a upload layout) — w binds as array<f16>, no conversion at stage; x still converts once.

enc_gemm4_src

gelu_tanh_on_gpu

Run the GPU’s apply_act (tanh-GELU) over xs — the isolation harness for the overflow regression test. The activation is normally FUSED into the GEMM epilogue, which is exactly why it went unnoticed that it could return NaN: nothing ever exercised it on its own.

gemm2_tile

The row-tile height for an [m, n] output: the LARGEST tile (best register blocking) whose grid still fills the GPU, else the smallest.