Skip to content

v1.2.1 — Gemma 4 E2B decode CUDA-graph fix

Choose a tag to compare

@ai-hpc ai-hpc released this 21 Jun 00:50
· 9 commits to main since this release
d3ce26f

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_dyn ignored cache_head_dim — derived the KV
    position stride/offset from head_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_embedding missing 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