kv_dispatch: apply gemma-4 per-layer epilogue (PLE + layer_scalar) in cached prefill/decode (fixes #177, stacked on #178)#181
Open
layer5one wants to merge 6 commits into
Conversation
Gemma 4 GGUFs describe heterogeneous attention (sliding vs global layers) via a per-layer head_count_kv array and *_swa twin keys. The flat mapping collapsed these to one number and dropped the rest, so gemma-4 models whose global layers ship no attn_v tensor (12B, 31B: attention_k_eq_v) indexed past the end of the V collection at inference time (ndarray OOB panic). Emit the flat HF-style keys the safetensors detect path already reads: num_key_value_heads + num_global_key_value_heads from the per-layer array, head_dim/global_head_dim from key_length_swa/key_length, dual rope bases, partial_rotary_factor, sliding_window, and attention_k_eq_v detected from the tensor inventory (fewer attn_v than blocks), as llama.cpp does. Verified: gemma-4-12B-it-qat-q4_0 GGUF converts and INFERs without panic (was: index out of bounds at arraytraits.rs:36).
The GGUF-to-HF key replacement table is llama-shaped: it maps ffn_norm to post_attention_layernorm (correct only for the llama two-norm layer) and has no entries for attn_q_norm, attn_k_norm, post_attention_norm, post_ffw_norm, or layer_output_scale. On gemma 2/3/4 GGUFs this silently drops five tensor kinds per layer (240 tensors on the gemma-4 12B: QK norms, pre/post-FFN norms, per-layer output scale) and writes the pre-FFN norm weights into the post-attention slot. Gemma-4 relies on QK-norm instead of 1/sqrt(d) attention scaling, so a vindex converted without these tensors produces unscaled-attention garbage even after the geometry fix. Add a gemma-specific replacement pre-pass, gated on gemma2/gemma3/ gemma4* architectures (gemma 1 is llama-layout and keeps the generic path), and thread the GGUF architecture string into tensor-key normalization via normalize_gguf_key_for_arch. layer_output_scale maps to layer_scalar with no .weight suffix, matching Gemma4Arch::layer_scalar_key. Unit tests cover the gemma mappings plus llama/gemma-1 regression guards.
larql_inference::encode_prompt exists precisely because gemma-4 ships a tokenizer.json whose post-processor does not add BOS even though the model requires it, but the LQL query paths never used it: INFER, EXPLAIN INFER, WALK, and EXPLAIN WALK all called tokenizer.encode directly, in six call-sites across both the vindex and dense-weight backends. On gemma-4 the silently missing BOS is the difference between prose and digit babble — the raw prompt "The capital of France is" top-1s to "1" instead of continuing the sentence — and it presents as a model/numerics bug rather than an input bug. Route all six call-sites through encode_prompt: - Dense Weight backend (INFER, EXPLAIN INFER dense): the loaded ModelWeights already carries the detected architecture — use it via a shared encode_dense_prompt helper. - Vindex backend (INFER, EXPLAIN INFER, WALK, EXPLAIN WALK): add larql_vindex::arch_from_vindex_config, exposing the existing build_arch_json -> detect_from_json reconstruction (previously only reachable through a full weight load) so tokenization can see arch.bos_token_id() without loading weights. Legacy vindexes with no recorded model_config keep the old tokenizer-only behaviour via the encode_vindex_prompt fallback. encode_prompt is a no-op for architectures whose bos_token_id() is None (everything except Gemma4Arch today) and is idempotent when the tokenizer already emitted BOS, so non-gemma-4 models tokenize byte- identically to before.
Every non-K legacy nibble format and Q6_K decoded (and, where larql is the producer, encoded) a private interleaved layout: element pairs (2j, 2j+1) from the low/high nibbles of byte j, with Q5/Q6 high bits taken from consecutive 1- or 2-bit fields. ggml's actual layouts are planar: the low nibbles of a block hold the first half of the elements and the high nibbles the second half (Q4_0/Q4_1/Q5_0/Q5_1: elements j and j+16 per byte; Q6_K: per 128-element half, ql[l]/ql[l+32] low nibbles hold elements l/l+32 and the high nibbles elements l+64/l+96, with qh[l] packing the two high bits of all four at shifts 0/2/4/6). Q4_K (and the other K-quants with explicit llama.cpp-cited layouts) were already correct, which is why Q4_K_M models worked while the gemma-4 12B QAT GGUF — Q4_0 projections plus a Q6_K token embedding — produced structured garbage: every 32-element run of every quantized tensor was internally permuted with wrong high bits, at the right magnitude. Round-trip tests passed because larql's own encoders packed the same wrong layout; nothing pinned the bytes to ggml ground truth. Fixed surfaces, all moved to the ggml layout together: - larql-models: dequantize_q4_0/q4_1/q5_0/q5_1 (legacy.rs), dequantize_q6_k + q6k_row_dot (scalar & NEON) + q6k_row_scaled_add (q6_k.rs, now sharing a single q6k_subblock_vals helper), quantize_q4_0 (quantize.rs). - larql-compute: quantize_q4_0 + quantize_q6_k + decode_q6k_superblock_into (q4_common.rs), q6k_matvec (rewritten on the shared helper), q6k_q8k_matvec scalar (q4k_q8k_dot.rs), and the q4_dot.c C kernel (NEON dot loses its vzip — planar pairs directly with sequential Q8; scalar matvec/vecmat re-indexed). The Q6_K×Q8_K NEON and hand-asm kernels still assume the old layout; the dispatcher routes all arches through the fixed scalar path until they are reworked and re-verified on ARM (TODO(q6k-planar)). New tests pin dequantize_q6_k, dequantize_q4_0, q6k_row_dot, and q6k_row_scaled_add to ground-truth bytes+values from the real gemma-4 GGUF as decoded by gguf-py (mirrors llama.cpp), so layout drift can no longer hide behind internal round-trip consistency. Existing tests that asserted the interleaved layout are updated to planar. MIGRATION: vindexes containing larql-quantized Q4_0/Q6_K blocks written before this change (Q4K-format vindexes quantize V projections and FFN slices to Q6_K; interleaved-q4 FFN stores use Q4_0) decode incorrectly with the fixed readers and need re-extraction.
Both keys are present in gemma GGUF metadata (gemma4.final_logit_softcapping,
{arch}.attention.layer_norm_rms_epsilon) but were dropped, leaving the
reconstructed config on detect defaults. The softcap is monotonic so it
never flips a top-1, but it shapes the softmax; the eps emit is generic
across RMSNorm architectures.
… cached prefill/decode The dispatch helpers ran only attention->FFN per layer, dropping apply_per_layer_embedding and apply_layer_scalar. Gemma-4 applies both on every layer, so the default KV-cached decode produced garbage while single-pass INFER (and --kv-cache none) were reference-exact. Mirrors kv_prefill_run/kv_decode_step_run in larql-kv/src/generation.rs across all four helpers (sync/async x prefill/decode). No-op for archs without PLE/layer_scalar. Verified on gemma-4-12B-it-qat-q4_0 vindex: cached generation now token-identical to --kv-cache none; INFER fingerprints unchanged (template -> The at 100.00%). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Fixes #177.
Stacked on #178 — this branch is #178 plus one commit. The new material here is only the last commit (
1af99eb,crates/larql-inference/src/kv_dispatch/helpers.rs, +71/−6). Merge #178 first and this becomes a one-file diff; alternatively merging this closes both.Same disclosure as on #178/#179: this was debugged and verified by a Claude (AI) agent running in my homelab; I'm relaying its firsthand-verified work.
Bug
On gemma-4 dense, single-pass
INFERis reference-exact (after #178) but the default KV-cached generation path (run/chat) emits garbage on identical input.Root cause: the dispatch helpers in
kv_dispatch/helpers.rsrun only attention → FFN per layer, dropping gemma-4's per-layer epilogue —apply_per_layer_embedding(PLE) +apply_layer_scalar— which gemma-4 applies on every layer. The non-dispatch reference path (kv_prefill_run/kv_decode_step_runinlarql-kv/src/generation.rs) applies both, which is why--kv-cache nonewas correct while the cached engines weren't. Other architectures have neither step (both are no-ops), so existing tinymodel parity tests never caught it.Fix
All four helpers (sync/async × prefill/decode) now mirror the reference path:
kv_prefill_via_dispatch): precomputes PLE inputs from the prompt token ids and applies PLE + layer_scalar after each layer's FFN.kv_prefill_from_hidden_via_dispatch(no token ids available): documented caveat — PLE is skipped, layer_scalar still applied. This entry point is only used by the Gemma-3 multimodal path today, which has no PLE, so behavior there is unchanged.Verification
Model: gemma-4-12B-it QAT q4_0, GGUF→vindex (artifact + scripts public at taykso/gemma-4-12b-it-qat-q4_0-vindex):
--kv-cache noneon identical templated input: token-identical ("The capital of France is Paris.").INFERfingerprints unchanged (chat-template prompt →'The'100.00%).Notes
layer_scalar+ PLE tensors, pinningStandardEngineoutput againstNoCacheEngine. That's the gap that let this slip past the existing parity tests. Happy to add it to this PR if wanted.larql-compute/src/attention/gpu.rs(run_attention_block_gpu) uses the olderapply_rope_partialand skips the per-layer effective-rope-base handling — harmless for gemma-4, wrong for rope-scaled models if that path is ever used with them. Flagging for whenever the GPU path gets attention.