gemma export debt: tied-embedding dedup (#260), true-dynamic with_past default (#248), shared-KV trace fidelity (#194) - #290
Merged
Conversation
The shared-cache variants (SharedPositionalKVCache, PaddedSharedPositionalKVCache, OwnerReadOnlyKVCache) still read K/V back through the delegate's heap buffer (writeAt/sliceView, plus padHeadDim's raw copyToFloatArray + ctx.fromData), bypassing ctx.ops. Under tracing (ctx.isRecording, embedConstants=true) that baked an all-zero KV buffer as stablehlo.constant and disconnected the computed k_proj/v_proj — the same SKaiNET#763 bug fixed for PositionalKVCache in #193 — for any model with kvSharedLayers > 0. Mirror the #193 fix: - SharedPositionalKVCache / PaddedSharedPositionalKVCache get an ctx.isRecording branch that wires K/V functionally via ops.concat history (the padded write + slice-back read is an identity on the layer's own data, so no pad op is needed on the traced path). - OwnerReadOnlyKVCache returns the owner PositionalKVCache's traced K/V history (tracedKeys/tracedValues now internal) so followers attend over the owner's computed projections. Verified by a new KvCacheTraceFidelityTest case tracing a 3-block sandwich-norm Gemma with kvSharedLayers = 1: fails on the old code (zero KV constants), passes with the fix; no zero-KV stablehlo.constant and all projection dot_generals present. Closes #194 Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…he embedding twice (#260) When the checkpoint carries no explicit output.weight (tied embeddings, as in FunctionGemma-270M), Gemma4WeightLoader's fallback re-read the token_embd bytes into a SECOND tensor. Two tensors -> two BufferHandles -> the compiled export externalized the 262153x640 tied weight as two independent util.globals + two archive blobs: 77% of the weight archive (832 MiB bf16 where ~512 MiB would do) and ~640 MiB extra per in-memory copy during export. Both GGUF paths (sequential + streaming) now retain the loaded token_embd tensor and alias it (same Tensor, same BufferHandle) into output.weight. The trace then sees ONE weight, and the export emits one @t0 global consumed by both the embedding gather and the lm-head transpose/matmul. Verified with the real FunctionGemma-270M Q5_K_M checkpoint (FunctionGemmaExportTest): - exactly one 262153x640 util.global in gemma-gen.mlir (new assertion) - bf16 gemma.safetensors: 536,692,404 bytes (~512 MiB), down from 872,248,326 (~832 MiB) - deduped gemma-gen.mlir still iree-compiles (llvm-cpu vmfb). Closes #260 Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…the SENTINEL_PAST=7919 hack (#248) exportWithPast's dynamic KV graph was traced at a concrete sentinel prime (7919) and then regex-relaxed (`x7919x`/`x7920x` -> `x?x`) in the emitted StableHLO text — the most fragile of the MLIR-text rewrites tracked in #248 (a magic prime that must never collide with a real dim or SSA id). Engine 0.38+ traces a real `Dim.DYNAMIC` extent through concat/reshape and emits dynamic_broadcast_in_dim, previously opt-in via GEMMA_TRUE_DYNAMIC=1. Flip the default: true-dynamic tracing is now the standard path and needs no post-emit text rewrite. GEMMA_SENTINEL_PAST=1 rolls back to the legacy sentinel trace (relaxSeqDimToDynamic retained only for that rollback, marked LEGACY for removal with the remaining #248 items); the env var is passed through the exportFunctionGemma task. Verified with the real FunctionGemma-270M Q5_K_M checkpoint: - FunctionGemmaWithPastMlirDumpTest green — default emit carries `1x1x?x256` caches on inputs AND returns, zero sentinel dims leaked - emitted gemma-with-past.mlir compiles to a host vmfb (iree-compile --iree-input-type=stablehlo --iree-hal-target-backends=llvm-cpu) - FunctionGemmaWithPastCpuTest (two-graph prefill+decode loop vs the llama.cpp/board oracle) green. Refs #248 Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This was referenced Aug 11, 2026
Merged
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
1 participant
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.
Three items of gemma export debt, one commit each, all verified against the real FunctionGemma-270M Q5_K_M checkpoint.
T2.3 — shared-KV cache trace fidelity (#194)
SharedPositionalKVCache,PaddedSharedPositionalKVCacheandOwnerReadOnlyKVCacheread K/V back through the delegate's heap buffer (writeAt/sliceView, pluspadHeadDim's rawcopyToFloatArray+ctx.fromData), bypassingctx.ops. Under tracing that baked an all-zero KV buffer asstablehlo.constantand disconnected the computedk_proj/v_proj— the same SKaiNET#763 bug #193 fixed forPositionalKVCache— for any model withkvSharedLayers > 0.Fix mirrors #193: the write-through variants get a
ctx.isRecordingbranch that wires K/V functionally viaops.concathistory;OwnerReadOnlyKVCachereturns the owner's traced history (tracedKeys/tracedValues, nowinternal).T2.1 — tied embedding exported twice (#260)
Gemma4WeightLoader's tied-output fallback re-read thetoken_embdbytes into a SECOND tensor foroutput.weight→ twoBufferHandles → the compiled export externalized the 262153x640 tied weight as two independentutil.globals + two archive blobs (77% of the archive). Both GGUF paths (sequential + streaming) now alias the already-loaded embedding tensor intooutput.weight; the trace sees ONE weight and the export emits one@t0global consumed by both the embedding gather and the lm-head transpose/matmul.T2.2 — retire the SENTINEL_PAST=7919 hack (#248, item "Dynamic KV dims")
exportWithPasttraced the dynamic KV graph at a sentinel prime (7919) and regex-relaxedx7919x/x7920x→x?xin the emitted text. Engine 0.38+ traces a realDim.DYNAMICextent (previously opt-in viaGEMMA_TRUE_DYNAMIC=1). True-dynamic is now the default;GEMMA_SENTINEL_PAST=1rolls back to the legacy sentinel trace (relaxSeqDimToDynamicretained only for that rollback, marked LEGACY).Verification
KvCacheTraceFidelityTest.tracedKvSharedDecoderKeepsComputedKV(3-block sandwich Gemma,kvSharedLayers=1)dense<0.0>KV constants, all projectiondot_generals presentKVCacheVariantsTest(eager shared-cache semantics)FunctionGemmaExportTestw/ real GGUF: exactly ONE262153x640util.globalingemma-gen.mlir(new assertion)gemma.safetensorssizegemma-gen.mlir→iree-compilellvm-cpu vmfbFunctionGemmaInt8QuantTest(int8 archive window re-based to the deduped ~512 MiB bf16 → ~257 MiB int8)FunctionGemmaWithPastMlirDumpTest(default path now true-dynamic)1x1x?x256caches on inputs AND returns, zero sentinel dims leakedgemma-with-past.mlir→iree-compilellvm-cpu vmfbFunctionGemmaWithPastCpuTest(two-graph prefill+decode loop vs llama.cpp/board oracle):transformer-core/:llm-inference:gemma/:llm-runtime:kgemmafulljvmTest+apiCheck(GEMMA_GGUF set)Closes #194, closes #260. #248 stays open — the bf16/int8 text rewrites and the
forwardWithPast/refsForitems remain.🤖 Generated with Claude Code