[quantization] Add static decode step for StaticGemma4Runtime - #843
Conversation
| cache.past_k[:, :, : self.past_len, :], | ||
| cache.past_v[:, :, : self.past_len, :], |
There was a problem hiding this comment.
Please use a fixed past cache, for example (B, Hkv, max_seq - 1, D) as in StaticLlamaRuntime, compact valid prefill K/V into its prefix, and pass the full buffer on every decode invocation.
There was a problem hiding this comment.
Please use a fixed past cache, for example (B, Hkv, max_seq - 1, D) as in
StaticLlamaRuntime, compact valid prefill K/V into its prefix, and pass the full buffer on every decode invocation.
Thank you for the suggestion, but... I've investigated adopting the StaticLlamaRuntime pattern for StaticGemma4Runtime.decode_one, and found that it requires changes beyond the scope of this PR.
The Issue:
The Gemma4 text attention wrapper (quant_text_attention.py) performs a capacity check in _build_attention_mask that raises when k_len exceeds the causal_mask_template size:
if k_len > self.causal_mask_template.size(-1):
raise RuntimeError(
f"Key length exceeds causal mask capacity: k_len={k_len}, capacity={self.causal_mask_template.size(-1)}."
)When passing the full cache buffer (2048) and concatenating with the new token's K/V, we get k_len = 2049, which triggers this error.
Why StaticLlamaRuntime Works:
StaticLlamaRuntime uses max_seq - 1 cache capacity AND processes prefill via chunked append_prefill, so the cache never contains more than max_seq - 1 valid tokens. The decode concatenation produces k_len <= max_seq - 1, which passes the capacity check.
StaticGemma4Runtime Difference:
Gemma4 uses single-shot prefill that returns full max_seq K/V tensors. Passing this full buffer to decode and concatenating with the new token produces k_len = max_seq + 1, exceeding capacity.
Required Changes for StaticLlamaRuntime Pattern:
To support passing the full buffer while only using valid positions, we could do the following:
quant_text_attention.py: Addcache_valid_lengthparameter to_build_present_key_valueto slice past K/V before concatenationquant_text_decoder_layer.py: Forwardcache_valid_lengththrough the decoder layerexport_adapters.py: UpdateGemma4TextDecoderLayerDecodeExportAdapter.forwardto accept and forwardcache_valid_length
These changes touch core wrappers and, I think, require a separate, focused PR.
Current Approach:
For this PR, decode_one slices the cache to past_len before passing:
past_key_value=(
cache.past_k[:, :, :self.past_len, :],
cache.past_v[:, :, :self.past_len, :],
)This ensures k_len = past_len + 1 <= max_seq, passing the capacity check while the attention mask (which is sliced to match k_len by _normalize_attention_mask_shape) correctly indicates valid positions.
Future Work:
I can file a follow-up issue to add cache_valid_length support to the Gemma4 attention wrapper infrastructure, enabling the cleaner StaticLlamaRuntime pattern for StaticGemma4Runtime. What do you think?
There was a problem hiding this comment.
I think there is a misunderstanding about the proposed cache shape.
The suggestion was not to pass the current max_seq-sized cache buffer in full. The persistent past cache itself should be allocated with capacity max_seq - 1.
With max_seq = 2048, the decode shapes would therefore be:
past_key_value: (B, Hkv, 2047, D)
new_key_value: (B, Hkv, 1, D)
present KV: (B, Hkv, 2048, D)
Therefore, the concatenation does not produce k_len = 2049; it produces exactly k_len = max_seq, which satisfies the existing capacity check. No cache_valid_length argument or wrapper change is needed.
The single-shot prefill output may still have sequence length max_seq. The runtime should copy only the valid prompt K/V range into the prefix of the max_seq - 1 persistent cache, as StaticLlamaRuntime.prefill() already
does. StaticLlamaRuntime does not require chunked append-prefill for its initial prefill; append_prefill is a separate multi-turn path.
For decode, the fixed physical layout is:
[past cache: max_seq - 1 slots][current token: 1 slot]
Unused past-cache slots are masked. The decode mask should therefore expose [:past_len] plus the final slot [-1], rather than placing the current token at physical index past_len.
For a shared-KV store layer, the fixed-size shared K/V passed to consumer layers can similarly be constructed from the max_seq - 1 past cache and the returned one-token K/V delta.
Slicing by a runtime cache_valid_length inside the exported wrapper would make the K dimension dynamic again, which is what the fixed-shape runtime should avoid.
If this is still difficult to follow, we can merge this PR as-is, and I can submit a follow-up PR implementing the fixed-shape cache design. Alternatively, you can implement it directly in this PR.
There was a problem hiding this comment.
Ok, let me handle that in the next PR)
| # --- Step 7: decode validation --- | ||
| print("[run_static_gemma4_runtime] Step 7: verify decode") | ||
| rt_decode_logits = verify_step_decode( | ||
| runtime, batch, cfg.prompt, image, prefill_logits=rt_logits, num_decode_steps=1 |
There was a problem hiding this comment.
| runtime, batch, cfg.prompt, image, prefill_logits=rt_logits, num_decode_steps=1 | |
| runtime, batch, cfg.prompt, image, prefill_logits=rt_logits, num_decode_steps=cfg.verify_steps, |
Implements decode_one with PLE, shared-KV support, and verify_step_decode validation. Co-authored-by: Cline TICO-DCO-1.0-Signed-off-by: d.savchenkov <d.savchenkov@partner.samsung.com>
2 participants
Related issue: #768
Related PR: #822
What
This PR implements the decode step for the
StaticGemma4Runtimein the TICO project, enabling side-by-side validation of single-token autoregressive decoding against Hugging Face's reference implementation for Gemma4 E2B models.Why
The static-shape NPU inference flow requires separate prefill and decode paths. While prefill was previously implemented, decode was a skeleton placeholder. This PR adds the full decode implementation to:
Key Design Decisions
past_lenbefore passing to decode layersmax_seq-sized cache buffer contains stale padding data. Passing only valid tokens (0..past_len) ensures the attention wrapper's capacity check passes and avoids carrying garbage through attention computation.past_len + 1entriespast_len, the valid range is 0..past_len inclusive. Consumer layers need exactly this range to compute correct attention scores.verify_step_prefillreturn valueverify_step_decodefunction acceptsprefill_logitsas a parameter to avoid redundant prefill computation, reducing runtime while maintaining test isolation.DynamicCachesame_token_modefor multi-step decodesame_token_mode=Trueparameter forces both paths to use HF reference tokens, enabling isolated quantization error measurement.Changes
tico/quantization/recipes/debug/static_gemma4_runtime.pybuild_decode_masks_and_rope: Replaced placeholder with full implementation generating per-layer-type full/sliding decode masks and RoPE tensors at the current decode positiondecode_one: Added PLE (Per-Layer Embeddings) support, shared-KV bookkeeping, and proper KV cache slicing (past_leninstead of fullmax_seq)verify_step_decode:DynamicCachereference with PEIR metricssame_token_mode: bool = Falseparameter for controlled multi-step comparisonrun_static_gemma4_runtime: Integrated step 7 (decode verification) into the smoke flow, reusing prefill logitstest/quantization/recipes/integration/test_static_gemma4_runtime_helpers.pyTestBuildDecodeMasksAndRope: New test class with 6 tests for decode mask generation:test_full_attention_mask_shape: Verifies mask shape is(batch_size, 1, max_seq)test_full_attention_mask_values: Verifies positions0..past_lenare0.0(allowed)test_sliding_window_mask_boundary_early_decode: Verifies window covers[0, past_len]whenpast_len < sliding_windowtest_sliding_window_mask_boundary_late_decode: Verifies window covers[past_len-sw+1, past_len]whenpast_len >= sliding_windowtest_sliding_window_mask_first_step: Verifies atpast_len=0, only position 0 is allowedtest_rope_computed_at_past_len: Verifies RoPE is computed at the correct decode positionTests
Example Script