Running an LLM on VKNN

Running an LLM on VKNN

VKNN runs Qwen2.5-Coder-0.5B — a qwen2-architecture autoregressive decoder (~0.5B / 494M parameters, 24 decoder layers, apache-2.0) — entirely on the GPU with zero CPU compute fallbacks. Every tensor-compute op in the graph (embedding gather, all projections and MLP matmuls, RMSNorm, RoPE, GQA attention, SwiGLU, softmax, residual adds, the KV-cache concats, and the final logits matmul) executes on the Vulkan backend. The only host code is the BPE tokenizer, the greedy/sampling argmax on the logits readback (I/O, not a compute op), and the token loop itself.

This guide covers the path end to end: export → convert (with the 0-fallback gate) → the GPU op work that made it possible → the host-driven decode loop → the terminal chat app → validation → precision.

1. Export the ONNX (with a KV cache)

Export a with-past decoder — one that takes and returns a KV cache — with a plain opset, external weights beside the .onnx:

optimum-cli export onnx \
  --model Qwen/Qwen2.5-Coder-0.5B \
  --task text-generation-with-past \
  qwen-onnx/

This produces a model.onnx plus its external-data weight file (*.onnx_data) and the tokenizer files. The decoder graph has:

• Inputs: input_ids, attention_mask, position_ids, and past_key_values.{0..23}.key / past_key_values.{0..23}.value (24 layers).
• Outputs: logits and present.{0..23}.key / present.{0..23}.value.

Keep the external-data weight file beside the .onnx — VKNN's importer resolves it relative to the .onnx directory. A missing weight file loads as zeros.

Exporting with a plain opset (no ORT contrib fusions) decomposes RMSNorm, RoPE, and the GQA repeat_kv into ops VKNN already runs on the GPU, so no fused contrib op (com.microsoft.SimplifiedLayerNormalization, com.microsoft.RotaryEmbedding) needs a native kernel.

2. Convert to .vxm and prove 0 fallbacks

Compile the ONNX to a .vxm, declaring the dynamic axes (batch / sequence / past length), and write a support report — the 0-CPU-fallback oracle:

vknn_compile qwen-onnx/model.onnx qwen.vxm --fp16 \
  --shape input_ids=1x1 \
  --shape position_ids=1x1 \
  --shape attention_mask=1xC1 \
  --shape past_key_values.0.key=1xHxCxD \
  --shape past_key_values.0.value=1xHxCxD \
  # ... one --shape per past_key_values.<l>.key/value, all 24 layers ...
  --support-report report.json

python3 tools/check_model_support.py --engine-report report.json

Here C is the fixed past-context length, C1 = C + 1 (the past slots plus the one new token), H is the number of key/value heads, and D is the head dimension. The compiled model reports these back — examples/chat.cpp reads kv_heads, C, and head_dim from past_key_values.0.key at load, so the decode loop stays model-agnostic.

The full decoder compiles to 100% Vulkan: every node in report.json reads backend:"vulkan", and the summary has no "cpu" / "none" key. To make any residual gap a hard failure instead of a silent CPU landing, run with Config::allowCpuFallback = false, an empty Config::fallback, and island-folding off (Hint::GpuIslandFold = Off).

Two shape paths both work:

• Native dynamic-shape path — declare the axes with --shape (or Config::inputShapes on an ONNX-loaded session) and the importer resolves the dynamic attention shape chains at plan time.
• Static prebake path — a single fixed-max-context (fixed-C) plan, which is what the decode loop uses (one bucket, zero-key fast path every step). A prefill shape and a decode shape can be combined into one multi-bucket .vxm with --bucket.

3. The GPU op work that made 0-fallback possible

Closing the last gaps to a fully-GPU decoder took a few targeted additions to the Vulkan path:

• GPU IsNaN and And kernels. These were the only ONNX ops the exported decoder used that VKNN lacked — IsNaN as a float→bool guard, And as a boolean mask AND. Adding them (with their shape rules) also let the importer fully resolve the dynamic attention shape chains, because the shape-arithmetic subgraphs that feed the mask now const-fold to concrete shapes.
• A BOOL-target Cast gate. An int64→BOOL Cast now runs on the GPU (vk_gates.cpp), so the mask-construction casts stay on-device instead of falling back.

Two correctness fixes in the Vulkan path keep the decoder numerically exact:

• Decode Int64 constant operands to fp32. The causal-mask construction uses Int64 constants; decoding them to fp32 when initializing the flat-path float buffers keeps the mask values correct through the fp16/fp32 compute path.
• Keep the Gather embedding index in fp32. The token-embedding Gather passes its indices as a float buffer; keeping the index in fp32 means a token id beyond the fp16 integer range (the vocab is 151936 wide) is not truncated when it selects the embedding row.

4. The host-driven decode loop

The decode loop is host-driven around a single fixed-max-context (fixed-C) plan — one static bucket that serves every step, so run() takes the zero-key fast path each time. The KV cache is not a graph re-plan; it lives in the past_key_values buffers:

• The past key/value buffers are the KV cache, retained across steps and turns.
• Stale slots are masked via the attention_mask input: a length-C+1 mask marks the p valid past slots plus the current token; the causality is an additive mask, so a masked slot contributes zero.
• position_ids carries the absolute position p of the current token across the whole conversation.
• After each step, the new token's key/value (at present sequence index C) is copied into cache slot p (clamped to C-1 past the compiled context).

Each step feeds one token (input_ids shape [1,1]) at absolute position p, the resident [1, kv_heads, C, head_dim] cache as the past inputs, and the mask over the valid slots. This reproduces the reference HuggingFace greedy stream token for token. Prefill runs each prompt token through the same step (the last prompt token's logits give the first generated token); decode repeats it until an EOS token or the token budget.

Tokenization (BPE) is the one host dependency; every graph compute op runs on the GPU, and the only unavoidable host touch inside the loop is the final logits readback for argmax — which is I/O, not compute.

5. The terminal chat app

Two pieces make an interactive chat:

• examples/chat.cpp → the device binary vknn_chat. It reads whitespace-separated prompt token ids on stdin (one turn per line), runs prefill + decode on the VKNN session, and streams generated token ids to stdout (one integer per line, flushed), then an END sentinel per turn. It keeps the conversation KV cache and absolute position across turns. Sampling is greedy at --temp 0, otherwise temperature + top-k + top-p.

  vknn_chat model.vxm [--backend vulkan|cpu] [--precision low|normal|high] [--fp32-tensors CSV]
            [--max-tokens N] [--temp T] [--top-k K] [--top-p P] [--eos ID] [--seed S]
  

• chat_host.py → the host front-end. It tokenizes with a HuggingFace AutoTokenizer, drives the device binary over adb (feeding prompt ids on stdin, reading generated ids from stdout), detokenizes and streams the completion to the terminal with a typewriter effect and a live tokens/s counter, and holds the REPL.

examples/chat.cpp is a device runner; add chat to the _vknn_examples list in CMakeLists.txt and ./build.sh --android builds vknn_chat.

Usage

# 1. Compile a fixed-C decode plan (see step 2) and build the runner.
#    (add `chat` to _vknn_examples in CMakeLists.txt first)
./build.sh --android --convert
./build.sh --android

# 2. Push the runner and model to the device.
SERIAL=<device-serial> ; DDIR=/data/local/tmp/vknn/qwen
adb -s $SERIAL shell "mkdir -p $DDIR"
adb -s $SERIAL push build-android/vknn_chat qwen.vxm $DDIR/

# 3. Run the host front-end (tokenizer + REPL + streaming display).
python3 chat_host.py --serial $SERIAL --ddir $DDIR --model qwen.vxm \
    --tokenizer qwen-onnx --precision low --max-tokens 128

chat_host.py opens a you> prompt, sends each line's token ids to the device, and streams the detokenized completion back. The KV cache persists across turns inside vknn_chat.

Example session

The tokenizer ships Qwen's ChatML template, so the base completion model can be driven as a chat. Asking questions — all generation runs on the GPU:

user>  Write a Python function to check if a number is prime.
model> def is_prime(n):
           if n <= 1:
               return False
           if n <= 3:
               return True
           if n % 2 == 0 or n % 3 == 0:
               return False
           i = 5
           while i * i <= n:
               if n % i == 0 or n % (i + 2) == 0:
                   return False
               i += 6
           return True

Or raw completion — feed the prefix directly instead of a chat turn:

prompt> In Python, a list comprehension is
model>  a concise way to create a new list by iterating over an existing list and applying a function to each element.

Greedy decoding on this 0.5B base model repeats on open-ended prompts; pass --temp 0.7 --top-p 0.9 (or --top-k 40) to chat_host.py / vknn_chat for more varied replies.

6. Validation

Correctness is validated against the HuggingFace transformers reference:

• Prefill logits argmax matches HF — a prefill forward on a reference prompt produces the same top token as the HF model.
• The greedy token sequence matches HF greedy — generating N tokens from a fixed prompt yields an identical greedy token stream token-for-token.

Any divergence is bisected by precision (see below) with --layer-dump / tools/compare_layers.py.

7. Precision

Every kernel accumulates in fp32 regardless of storage tier — the Config::precision tier governs only storage. The decoder runs correctly with fp16 weights (--fp16 at compile). A recommended progression:

1. Validate at --precision high first (full fp32 storage), with 0 fallbacks, and confirm the logits and greedy tokens match HF before touching precision.
2. Drop to --precision normal with a targeted --fp32-tensors list pinning the fragile tail — the final model/norm output, the wide lm_head logits (151936-wide), and the RoPE frequency tables — so an fp16 store never flips two near-tied top tokens under greedy argmax. --precision high stores activations fp32 throughout.

fp16 stores round to nearest even and saturate a finite result to ±65504 rather than overflowing to ±inf. Attention probabilities live in [0,1] and are well-conditioned, so the fp16 tier is usually fine for them; the wide-reduction tensors (the norm output and the logits) are the ones worth pinning fp32.

Performance

On-device latency at --precision low (fp16 activations), measured over the fixed-context plans with --repeat 20. Wall latency comes from non-profiled --timing runs (per-forward bind + segments + collect); GPU is the profiler's kernel total (Config::profile), which serializes each op behind a barrier and so is read only for the per-op breakdown, not for wall latency. TTFT is one prefill forward over a 32-token prompt; TPOT is one decode step at a 64-token context (tokens/s = 1000 / TPOT_median_ms). Both runs report 0 CPU fallbacks — every op is Vulkan.

Device	TTFT — prefill 32 tok (wall median / min / p90 · GPU)	TPOT — decode step (wall median / min / p90 · GPU)	tokens/s (median / best)
Mobile arm64 Vulkan GPU A	340 / 237 / 370 ms · 294 ms	147 / 133 / 153 ms · 140 ms	6.8 / 7.5
Mobile arm64 Vulkan GPU B (newer)	159 / 156 / 161 ms · 154 ms	115 / 105 / 117 ms · 119 ms	8.7 / 9.5

Both forwards are ~95% MatMul. The decode step is weight-memory-bandwidth-bound — one token streams all ~494M fp16 weights (~0.99 GB) through the GPU — so TPOT tracks memory throughput, not arithmetic. The single largest op is the final lm_head matmul (151936-wide), ~12 ms per decode step on GPU A. RMSNorm, softmax, RoPE (decomposed to slice/concat/mul), and the mask ops together are under 1 ms of a ~140 ms step. Wall runs a few ms above the GPU kernel total (host submit, I/O, barrier bubbles), and warms up under sustained load — the min column is the cool-device figure.

The fp16-activation path is numerically sound at the norm. RMSNorm is a native fused op — the lower_rmsnorm pass folds the exported mean/rsqrt/mul chain into one RMSNorm node (OpType::RMSNorm) that accumulates its variance reduction in fp32 regardless of the fp16 storage tier. That keeps the --precision low path faithful through the normalization, so the prefill-logits argmax matches the HuggingFace reference without pinning the norm to fp32. Long greedy streams remain precision-sensitive at the wide lm_head argmax (see Precision).