v1.2.1 — Gemma 4 E2B decode CUDA-graph fix
Fixes the Gemma 4 E2B decode CUDA-graph path and enables it for the E2B
model (which uses KV sharing). The graph was added in v1.2.0 but excluded for
shared-KV models and carried three latent bugs — none ever shipped to users
(the path was guard-disabled for every real Gemma 4 model).
Bugs fixed
flash_attention_decode_dynignoredcache_head_dim— derived the KV
position stride/offset fromhead_dim, so Gemma sliding layers (head_dim
256, cache slot 512) read each KV slot at half stride.- PLE input buffer aliased scratch — the captured graph reused the same
scratch offsets for its per-layer buffers and overwrote the PLE values
mid-replay. Fixed with a dedicated persistent buffer. scale_embeddingmissing from the graph path (dominant bug) — the
forward ran on an embedding ~39× too small (Gemma ScaledWordEmbedding). Input
RMSNorm hides it from attention, but the residual stream and the per-layer
embedding block read the embedding directly → degeneration.
Also adds a device-position Q-only RoPE (rope_inplace_dyn) so the captured
graph can rotate Q for the trailing KV-sharing layers.
Performance (Orin Nano Super, gemma-4-E2B-it-Q4_K_M, MAXN_SUPER)
Clean equal-token decode A/B:
| Context depth | graph OFF | graph ON | Δ |
|---|---|---|---|
| shallow / typical | 30.7 | 32.4 | +5.5 % |
| 512-token depth | 23.9 | 24.0 | ~neutral |
The graph collapses ~448 per-token kernel launches into one cudaGraphLaunch,
recovering host-launch overhead that matters most when tokens are fast (shallow
context); at depth the per-token GPU work dominates and the win tapers to
neutral. Default-on; JLLM_DECODE_GRAPH=0 reverts. Greedy output matches the
per-step path (a 511-token prompt crossing the sliding window is byte-identical),
modulo the existing int8-GEMM-class borderline tie-breaks.
Full Changelog: v1.2.0...v1.2.1