katolikov/smolvlm2-vknn

SmolVLM2-2.2B-Instruct for VKNN — full-GPU, one file (fp16 + int4)

SmolVLM2-2.2B-Instruct compiled for the VKNN Vulkan inference engine: every tensor-compute op runs on the GPU (0 CPU fallback), and the whole model — vision tower, connector, token embeddings, and the 24-layer Llama decoder at both prefill and decode shapes — ships as one .vxm file with a shared, content-deduped weight pool. No separate embeddings file, no vision/text split.

Parameters	2.2B (vision 434M + decoder 1.81B, untied lm_head)
Files	smolvlm2-2.2b-fp16.vxm (4.5 GB, fp16) · smolvlm2-2.2b-i4.vxm (1.35 GB, int4 weights)
Precision	fp16 weights + fp16 GPU compute (fp32 host boundary); the i4 file quantizes the MatMul weights to int4 (group 128, activation-salient outlier columns kept fp16, per-layer error guard) and computes in fp16
Buckets	vision [1,3,384,384] · embed [1,128]/[1,1] · decoder prefill S=128 · decoder decode S=1, C=512 KV slots
Image tokens	81 per 384×384 tile (one tile, do_image_splitting=False)
Measured on a current flagship phone GPU	vision 0.98 s · prefill TTFT 0.85 s · 6–7.5 tok/s decode (fp16 file)
Device requirement	fp16: ~4.6 GB of GPU-addressable memory; int4: ~1.5 GB — the int4 file is the one to reach devices with tighter mapping budgets

Unlike the int8/CPU builds of this model floating around, these are GPU builds: the fp16 file is gated against an fp32 onnxruntime reference (greedy token streams match; vision embeddings cosine ~1.0) rather than eyeballed. The int4 file is a calibration-free weight-only requantization of the gated fp16 file (vknn_compile in.vxm out.vxm -Os --quant-samples 0 — the compiler requantizes all five buckets over the shared weight pool); its answers stay grounded in the image with slightly less detail than fp16 (weight-only int4's expected envelope), also at 0 CPU fallback.

Run it

# device binary (examples/llm/vlm.cpp in the VKNN repo) + host tokenizer front-end
python3 examples/llm/vlm_host.py --serial <SERIAL> \
    --tokenizer <this repo's tokenizer files> \
    --image photo.jpg --question "How can I improve this shot?"

Or in the VKNN Android demo app: install, open Library, download SmolVLM2 2.2B, then use the VLM tab as a camera coach.

Export recipe (how this was made)

The official ONNX export of SmolVLM2 is unusable for a static-shape GPU planner (com.microsoft contrib ops + data-dependent NonZero). This build uses a hand-written export (torch.onnx.export(dynamo=False), opset 17, eager attention) with four load-bearing tricks:

1. Bypass Idefics3VisionEmbeddings.forward. Its NaViT variable-resolution position ids (bucketize + boolean indexing) are the only source of NonZero/GatherND. For a full unpadded 384×384 tile they collapse to arange(num_patches) — verified bit-identical (maxdiff 0.0) against the original module. The vision graph is plain SigLIP: Conv patch-embed + constant position embedding + 27 encoder blocks + the pixel-shuffle connector.
2. The 4-D causal mask is a graph INPUT. transformers' masking utils return a 4-D attention_mask as-is, so the host builds the additive mask and no mask-construction subgraph is traced at all (also sidesteps a torch.jit.trace crash in sdpa_mask).
3. Mask fill is -1e4 — never finfo.min, never -65504. -3.4e38 overflows fp16 to -inf and 0 * inf = NaN poisons the softmax; -65504 overflows on the first addition. exp(-1e4) underflows to exactly 0 in fp16, so masked positions contribute nothing and no inf ever exists on the compute path.
4. The decoder is traced ONCE with a dynamic sequence axis (validated by onnxruntime at both S=128 and S=1), then compiled into prefill and decode buckets of one .vxm. present.* outputs carry only the NEW KV rows — with 32 MHA KV heads (192 KiB/token), returning the full concatenated cache would move ~200 MB per decoded token across the host boundary.

The prompt's 81 <image> token rows are replaced by the vision embeddings on the host between the embedding lookup and prefill (both GPU buckets of the same file) — this is what keeps data-dependent ops out of the graph entirely.

Preprocessing: resize to exactly 384×384 (aspect distortion is the processor's own single-tile behavior), x/127.5 - 1 (rescale 1/255, normalize mean=std=0.5), CHW fp32. rope_theta = 130000, <image> id 49190, <end_of_utterance> id 49279.

Validation

• vision-embedding bypass: bit-identical (maxdiff 0.0) vs the stock module
• ONNX graphs: 100% ai.onnx, zero data-dependent-shape ops, ORT vs eager PyTorch maxdiff ≤ 4e-4 at S∈{3, 128, 1}
• device: --support-report shows 0 CPU fallback for all 5 buckets (632+1+1+852+756 nodes, all vulkan); an 80-token greedy device stream matches the fp32 onnxruntime reference token-for-token (80/80) on a real photo; vision embeddings cosine 0.999996
• int4: all five buckets requantized over the shared pool (vision 868→234 MB, decoder 3423→916 MB of weights), 0 CPU fallback, grounded same-subject answers on the photo gate
• decode is bounded by the fp32 host KV round-trip (~200 MB/token at C=512 with 32 MHA KV heads), not by compute — an on-GPU KV cache is the engine's next lever
• found in bringup and fixed in the engine: a mobile GPU driver's GLSL tanh() returns NaN for |x| > ~88.7 (exp-overflow in its symmetric lowering); SigLIP's tanh-GELU feeds it ±230, so the engine clamps tanh arguments to ±10 (bit-identical for every finite-result argument)