Skip to content

v1.2.0 — decode optimization cycle

Choose a tag to compare

@ai-hpc ai-hpc released this 20 Jun 13:01
· 17 commits to main since this release
2dd6a9c

A decode-focused optimization cycle on top of v1.1.0, for Gemma 4 E2B on
Jetson Orin Nano Super. No correctness change (greedy output unchanged) and
no change to the Qwen3 path — every optimization is env-gated (JLLM_*=0).

What's in it

  • perf(rope): precomputed cos/sin table — no per-decode-step trig.
  • perf(attn): warp-parallel Q·Kᵀ + block softmax in the decode attention kernel.
  • perf(gemv): dp4a triple QKV GEMV (Q4_K+Q4_K+Q6_K), one shared q8_1 activation.
  • perf(kernels): half2-vectorized vec_add / RMSNorm / GeGLU / SwiGLU / scale.
  • perf(engine): Gemma 4 CUDA-graph support — inert for E2B (20 shared-KV
    layers → graph correctly skipped); benefits only a non-shared-KV variant.
  • fix(engine): decode-graph skip messages now print once (were per-token).
  • chore(bench): scripts/bench_v11x.sh decode profiler.

Measured (Orin Nano Super, gemma-4-E2B-it-Q4_K_M, MAXN_SUPER)

Same-conditions decode A/B (v1.1.0 vs v1.2.0 binaries, identical prompt, 3
interleaved runs each, ≤0.3 tok/s variance), at three context depths:

Context depth (prompt len) v1.1.0 v1.2.0 Δ
shallow (~11 tok) 29.1 31.0 +6.5 %
~450 (421 tok) 22.1 24.4 +10.4 %
512 (511 tok) 21.1 23.8 ~+12 %

The decode gain grows with context depth — the warp-parallel decode
attention does more work per token as the KV cache grows, so it saves more
there. The 512-depth row is the apples-to-apples comparison with v1.1.0's
"23.5 tok/s at parity" methodology (v1.1.0 reads 21.1 here on
jetson_clocks-unlocked MAXN_SUPER). The gain comes from the attention kernel,
the dp4a QKV triple, and the half2 kernels; the RoPE table and the Gemma 4
graph path contribute negligibly for E2B.

Full Changelog: v1.1.0...v1.2.0