Module gpu_offload 

Optional GPU offload for the output projection, the single most expensive matmul in a decode step.

Why only this one matmul

Measured on an Intel integrated GPU against 14 CPU cores at SmolLM2-360M’s real shapes (kopitiam-gpu/tests/matmul_timing.rs):

decode  attn q/o     (1 tok)   cpu   443µs   gpu 1.88ms   0.24x   LOSS
decode  mlp gate/up  (1 tok)   cpu  1.28ms   gpu 3.20ms   0.40x   LOSS
output head, quantized + resident        cpu 26.12ms  gpu 8.55ms   3.06x  WIN

A decode step is a single row of activations — far too little arithmetic to pay for moving a weight matrix across the bus. Offloading the small projections would make chat 2–4x slower. The output head is different on two counts: it is vocab x hidden (49152 x 960 for SmolLM2-360M), so it costs more per token than every other decode matmul combined; and its weight can stay resident on the device, so only the activation row crosses the bus per token.

Why the weight has to stay quantized, not merely resident

As f32 that weight is 188 MB, against a 128 MB max_storage_buffer_binding_size on the maintainer’s adapter — it cannot be bound at all. Held as block-scaled int8 it is ~47 MB of quants plus ~6 MB of scales. So “quantized” is not an optimisation here, it is the difference between running on the GPU and not.

The tied-embedding wrinkle, which is also a CPU-side finding

Neither SmolLM2 release ships a separate output.weight; both tie it to token_embd.weight. [crate::weights::ModelWeights] loads token_embd through load_tensor_f32 (the embedding lookup needs real floats), and the tied fallback clones that — so the biggest matmul in the model has been running against a fully dequantized f32 weight, skipping the fused quantized kernel every other projection gets. This module reads the original quantized bytes back off the model instead.

Structs

GpuOutputHead

The output projection, resident on the GPU.

Functions

split_q8_0

Splits raw GGUF Q8_0 bytes into the separated (scales, quants) layout kopitiam-gpu wants.