genie-ai-runtime v1.1 → v1.3.1: Gemma 4 on a Jetson, from llama.cpp parity to beating it at depth

genie-ai-runtime is my from-scratch CUDA inference engine for the Jetson Orin Nano Super — the one I wrote about beating llama.cpp on Qwen3-4B by designing around the board’s unified memory. v1.1.0 added Gemma 4 E2B support and pulled Gemma’s decode up to llama.cpp parity and its prefill up 4.1×. The cycle didn’t stop there — through v1.3.1 the engine went from matching llama.cpp to beating it at every context depth on decode, with a separate TTFT win on short prompts. This post covers the whole arc.

The model support is the headline, but it’s not what I want to write about. What I want to write about is how the speedups happened — because every single one came from the same loop: profile, find the one bottleneck that matters, fix it, measure, repeat. No grand redesign. Just the profiler pointing at the next wall, over and over. This is the field journal of that cycle.

Table of contents

What Gemma 4 E2B actually is (and why it’s hard on-device)

The engine already ran Qwen3-4B. Gemma 4 E2B is a Gemma-3n-class architecture, and it is a genuinely different animal — a pile of non-standard pieces that each need their own forward-pass support before a single token is correct:

• Per-layer-type attention — layers alternate between local sliding attention (head_dim 256, a 512-token window) and global full attention (head_dim 512).
• Dual proportional RoPE — two rotary bases, θ = 1e4 on local layers and θ = 1e6 on global layers.
• Per-Layer Embeddings (PLE) — a per-layer embedding table the model mixes in at every layer. (I flagged PLE in my roadmap as one of the algorithmic tricks that bend the memory wall; this is me actually implementing it.)
• Trailing-layer KV sharing, GeGLU double-wide MLP, scaled input embeddings, weightless V-RMSNorm, and final-logit soft-capping.

None of that is optional — get the RoPE base wrong on one layer type, or miss the soft-cap, and the output silently diverges. The bar I hold the engine to is greedy-identical to llama.cpp: same model, same prompt, bit-identical tokens. v1.1 clears it on Orin across the usual prompts and on a 1043-token counting prompt that crosses the 512-token sliding window — both engines continue “231, 232, 233…” identically, which is the real test that the windowed attention and KV handling are correct past the window edge.

The architecture also pays off in memory, which is the whole game on an 8 GB board. Only 15 of the layers own a KV cache (the rest share), so sizing the KV pool to 15 instead of 35 layers is ~57% less KV memory — headroom that, on a board already running STT, TTS, and Home Assistant, is the difference between fitting and not.

The decode arc: 12 → 23.5 tok/s (llama.cpp parity)

Decode is memory-bandwidth-bound: one token per pass, streaming the weights. So the profiler’s answer was unsurprising — the time was in the dequant + matrix-vector paths — but the fix is all in matching llama.cpp’s int8 machinery:

• dp4a int8 MMVQ dot-product paths for Q4_K and Q6_K (#104–#106) — do the quantized mat-vec in int8 with the dp4a instruction instead of dequant-to-fp.
• A warp-per-row PLE GEMV (#103) for Gemma’s per-layer embedding projection.
• A Q4_K uint32 fast path (#102) for the dequant inner loop.

Together: 12 → 23.5 tok/s, dead even with llama.cpp’s token-generation rate. At ~23.5 tok/s a short home-assistant reply lands in roughly a second — past the threshold where the device feels like it’s thinking with you rather than buffering. (That was the v1.1 milestone: parity. A later split-K rewrite pushed decode past llama.cpp at every depth — see the v1.3.0 update below.)

The prefill kernel cycle: 152 → 620 tok/s on a 1261-token prompt (4.1×)

This is the part I’m proud of, because it’s the profiler at its most honest. Prefill is compute-bound (all the prompt tokens at once), so this is where tensor cores and faithful int8 MMQ should win — and each PR below is one profiler reading, one fix, one measured delta on the same 1261-token Gemma prompt:

Two of those deserve more than a table row.

#122 — the missing operand load. genie’s MMQ kernel had the tensor-core math but not llama.cpp’s operand-staging (ldmatrix into the right register tiles for the mma instruction). That’s the unglamorous 80% of a fast MMQ kernel — feeding the tensor cores correctly — and vendoring llama.cpp’s mma.cuh tile<> machinery is what made the rest of the cycle possible. Lesson: “I have a tensor-core kernel” and “I have llama.cpp’s tensor-core kernel” are very different sentences.

#126 — the profiler overturning a prior conclusion. Earlier in the project I had a finding on record: tensor-core attention degrades Gemma’s quality. It was wrong. The profiler (and a careful numerics check) showed the culprit wasn’t fp16 inputs — it was fp16 accumulation losing precision in the attention softmax/score path. Switch to fp16-in, f32-accumulate via cuBLAS and tensor-core attention is both correct and the single biggest win of the cycle at +43%. A wrong conclusion sat in my notes for weeks; profiling the actual failure, not the assumed one, is what dislodged it.

Earlier in the cycle, an F32 query-tiled flash attention (#118) fixed a long-context collapse — prefill throughput had been falling off a cliff on long prompts — taking the 8192-token prefill from 63 → 126 tok/s (2×). Plus register-pressure and warp-tiling passes on the Q4_K MMQ (#119, #120).

The cycle kept going: v1.2 → v1.3.1 (parity → beating llama.cpp)

v1.1 matched llama.cpp on decode. The next four releases pushed past it — same loop, new walls.

v1.2.0 — a decode optimization cycle

With prefill on tensor cores, the profiler moved back to decode. A batch of memory-bound wins, none individually huge, compounding: a precomputed cos/sin RoPE table (no per-step trig), warp-parallel Q·Kᵀ + block softmax in the decode attention kernel, a dp4a triple QKV GEMV (Q4_K + Q4_K + Q6_K sharing one q8_1 activation), and half2-vectorized vec_add / RMSNorm / GeGLU / SwiGLU. The interesting part is where the gain landed:

The decode gain grows with context depth — the warp-parallel attention does more work per token as the KV cache grows, so it saves more there. That was the profiler already hinting at where the real decode-at-depth problem lived (v1.3.0, below).

v1.2.1 — the Gemma CUDA-graph path, and three bugs I’m glad I caught

Earlier I noted decode CUDA graphs were off for Gemma 4. v1.2.0 added a Gemma graph path but left it inert for E2B; v1.2.1 fixed it and turned it on. The graph collapses ~448 per-token kernel launches into one cudaGraphLaunch — pure host-overhead recovery, worth +5.5% at shallow depth (where tokens are fast and launch overhead dominates), tapering to neutral at depth (where GPU work dominates).

Getting there meant fixing three latent bugs the guard had been hiding — and one is a perfect cautionary tale:

• flash_attention_decode_dyn ignored cache_head_dim — it derived the KV stride from head_dim, so Gemma’s sliding layers (head_dim 256 in a 512-wide cache slot) read every KV slot at half stride.
• The PLE input buffer aliased scratch — the captured graph reused scratch offsets and overwrote the per-layer-embedding values mid-replay. Fixed with a dedicated persistent buffer.
• scale_embedding was missing from the graph path — the dominant bug. Gemma uses a ScaledWordEmbedding, so without the scale the forward ran on an embedding ~39× too small. Input RMSNorm hides that from the attention block — so attention looked fine — but the residual stream and the per-layer embedding read the raw embedding directly, and the model quietly degenerated. A bug invisible exactly where you’d look first is the kind only a byte-level greedy A/B catches.

v1.3.0 — split-K decode attention: now it beats llama.cpp

This is the one I’m proudest of in the whole arc, and it started with an ncu reading that was almost embarrassing. The decode-attention kernel launched grid = (n_heads) — 8 blocks on an 8-SM Orin: ~11% occupancy, 3% throughput. Each per-head block then walked the entire KV cache serially, so decode throughput fell off a cliff with depth — genie was 2.2× behind llama.cpp at 4096 tokens (8.6 vs 19.1 tok/s).

The fix is flash-decoding / split-K: split each head’s KV range across several blocks, each computing an online-softmax partial, then a reduce kernel combines them per head. Suddenly the 8-SM machine has enough blocks to fill, and the per-token work parallelizes instead of serializing:

genie now leads llama.cpp at every depth, and the at-depth collapse is gone. The split-K kernels are numerically exact vs the single-block path (standalone test: max abs diff 0.0), handle FP16 and INT8 KV plus Gemma’s sliding-window mask, and are wired into both the per-step and CUDA-graph decode paths. Shallow decode keeps the single-block kernel (below 64 tokens there isn’t enough KV to split).

v1.3.1 — the fastest kernel depends on the prompt length

The last wall was TTFT on short prompts. The cuBLAS tensor-core attention that won the v1.1 prefill cycle (#126) carries real per-call overhead — handle setup, ~70 batched-GEMM + softmax launches per layer, N×N score-buffer traffic — and that overhead dominates when the sequence is short. So v1.3.1 length-dispatches: below a ~768-token crossover it runs the fused F32 query-tiled kernel; above it, cuBLAS’s tensor-core GEMMs win. Best of both:

Typical 50–300-token interactive prompts now prefill 1.4–2× faster (straight TTFT); long context unchanged. Against llama.cpp’s real cold prompt-eval (llama-simple), genie’s short-prefill is 2–7× faster — llama pays a ~1–2 s cold warmup genie doesn’t. The v1.1 lesson (#126: cuBLAS TC attention is the prefill win) wasn’t wrong, just incomplete: it’s the win at length, and the wrong tool for a chat-sized prompt.

What I learned

• Profile the failure you have, not the one you assume. The fp16-accumulation bug (#126) is the whole lesson in miniature: a plausible-but-wrong conclusion (“TC attention hurts Gemma”) cost weeks until I profiled the actual precision loss instead of trusting the earlier verdict.
• Faithfulness to a reference kernel is a real, separable skill. The gap between genie’s MMQ and llama.cpp’s wasn’t the math — it was the operand load (#122). Re-deriving a fast kernel from scratch is much harder than understanding exactly what the reference does and matching it.
• Numerics on the edge have teeth. Gemma 4 forces FP16 KV — its V activations have ~10× RMS outliers that INT8 per-row quantization collapses, so the INT8 KV that’s fine for Qwen3 would wreck Gemma. And int8-activation GEMMs can shift greedy tie-breaks on borderline tokens (the same tradeoff llama.cpp makes); factual greedy outputs are unchanged, but you have to know that’s the boundary you’re trading on.
• A small grid starves a small GPU. The biggest win in the whole arc (split-K, v1.3.0) wasn’t a math change — it was occupancy. grid = (n_heads) is 8 blocks on an 8-SM Orin: a correct kernel running at 11% occupancy. On a small GPU, “is there enough parallel work to fill the machine?” is the first question, not the last.
• The fastest kernel is regime-dependent — so dispatch on it. cuBLAS tensor-core attention wins at long context and loses at chat length (v1.3.1); the single-block decode kernel wins when there’s too little KV to split (v1.3.0). The fix often isn’t a better kernel, it’s picking the right one per sequence length / depth.
• Graph capture is treacherous — but worth it. CUDA Graphs collapse hundreds of per-token launches into one (a real shallow-depth win), but Gemma 4’s KV sharing, sliding-window stride, and scaled embeddings made the captured path subtle: off for Gemma at v1.1, correct only after three capture-only bugs in v1.2.1 — including a missing scale_embedding that ran the model on an embedding ~39× too small, invisible behind RMSNorm until a byte-level A/B caught it.
• Gate everything; regress nothing. Every kernel in this cycle is env-gated (JLLM_*=0 reverts) and the Qwen3 path is byte-for-byte unchanged. New-model work should never be able to silently slow the model that already shipped.

I’ll keep it honest about where the gaps are now. Decode leads llama.cpp at every depth (1.1–1.6×) and the at-depth collapse is fixed. Short-prefill beats llama.cpp’s cold prompt-eval 2–7× (it pays a warmup genie doesn’t). The one place llama.cpp still wins is long-prompt prefill on a warm server, where its heavily-tuned MMQ is roughly 6× genie’s — that’s the next wall. A real result, not a finished one.

Why it matters

The point of all this is a single sentence: Gemma 4 E2B runs greedy-identical to llama.cpp on an 8 GB Jetson — decode now faster than llama.cpp at every depth (up to 1.58×), short-prompt TTFT 1.4–2× quicker than where the prefill cycle started, in ~57% less KV memory. That’s a current, capable on-device model for genie-claw, the private home agent this engine exists to serve — another model the household can run with nothing leaving the house.

And it’s the roadmap ethos in one release: the wins live at the seam between the algorithm (Gemma’s PLE, sliding-window attention, the int8 quant formats) and the kernel (MMQ operand loads, tensor-core accumulation, fused write-backs). You don’t find that seam by reading about it. You find it with a profiler, one wall at a time.