[quantization] Add generation step to StaticGemma4Runtime by dvsav · Pull Request #849 · Samsung/TICO

dvsav · 2026-07-28T12:43:45Z

Related issue: #768
Related PR: #822

What

This PR implements the generation step (Step 8) for the StaticGemma4Runtime in the TICO project, completing the full 8-step static runtime verification flow for Gemma4 E2B. It adds a greedy generation method with side-by-side validation against HuggingFace's reference implementation.

Additionally, this PR fixes the KV cache write logic in prefill to only store valid tokens (excluding padding), and adds proper shared-KV cache pre-population in decode_one for consumer layers.

Why

The static runtime needed a complete generation loop to verify end-to-end functionality beyond single-step decode verification. This enables:

Full sequence generation validation (not just single-token decode)
Token-level accuracy tracking against HF reference
Early detection of error accumulation across multiple decode steps
Visual inspection of generated text for qualitative assessment

Key Design Decisions

Greedy generation loop: The generate_greedy method reuses prefill() and decode_one() internally, ensuring consistency with the verified single-step implementations.
Image token preservation: The method extracts input IDs from _raw_inputs["input_ids"] (which preserves original image tokens) rather than llm_input_ids (which has image tokens replaced with pad_token_id for NPU compute).
Token accuracy metric: Instead of exact sequence match (which is unrealistic for quantized vs. FP models), we report token-level accuracy percentage and first mismatch position.
Detokenized output: Both runtime and HF generated sequences are detokenized and printed for visual inspection, enabling qualitative assessment beyond numerical metrics.
KV cache fix: In prefill, only valid K/V positions (up to valid_length) are written to the cache, excluding padding positions that contain garbage from padding tokens.
Shared-KV pre-population: In decode_one, shared-KV consumer layers are pre-populated from layer_caches before store layers update the cache, ensuring correct KV access across decode steps.

Changes

tico/quantization/recipes/debug/static_gemma4_runtime.py (+190/-9 lines):

StaticGemma4Runtime.prefill: Fixed KV cache write to only store valid-length K/V (excluding padding garbage).
StaticGemma4Runtime.decode_one: Added shared-KV cache pre-population for consumer layers before the layer loop.
StaticGemma4Runtime.generate_greedy (new, ~49 lines): Greedy generation method that:
- Runs prefill to process prompt + image
- Extracts input IDs from _raw_inputs to preserve image tokens
- Iteratively calls decode_one and collects generated tokens
- Supports early stopping on EOS token
verify_step_generation (new, ~92 lines): Validation function that:
- Runs runtime greedy generation
- Runs HF reference generation
- Reports token accuracy (%) and first mismatch position
- Prints detokenized text from both paths
run_static_gemma4_runtime: Wired Step 8 to call verify_step_generation and print generated text.

Tests

$ python -m pytest test/quantization/recipes/integration/test_static_gemma4_runtime_helpers.py -v
================================================================================ test session starts ================================================================================
platform linux -- Python 3.10.12, pytest-8.4.0, pluggy-1.6.0 -- /home/d.savchenkov/myenv/bin/python
cachedir: .pytest_cache
rootdir: /home/d.savchenkov/TICO
configfile: pyproject.toml
plugins: anyio-4.12.0, mock-3.15.1, xdist-3.7.0, cov-6.2.1
collected 21 items                                                                                                                                                                  

test/quantization/recipes/integration/test_static_gemma4_runtime_helpers.py::TestNormalizeValidTokenMask::test_batched_input PASSED                                           [  4%]
test/quantization/recipes/integration/test_static_gemma4_runtime_helpers.py::TestNormalizeValidTokenMask::test_shape_mismatch_raises PASSED                                   [  9%]
test/quantization/recipes/integration/test_static_gemma4_runtime_helpers.py::TestNormalizeValidTokenMask::test_with_attention_mask PASSED                                     [ 14%]
test/quantization/recipes/integration/test_static_gemma4_runtime_helpers.py::TestNormalizeValidTokenMask::test_without_attention_mask_no_pad_token_id PASSED                  [ 19%]
test/quantization/recipes/integration/test_static_gemma4_runtime_helpers.py::TestNormalizeValidTokenMask::test_without_attention_mask_uses_pad_token_id PASSED                [ 23%]
test/quantization/recipes/integration/test_static_gemma4_runtime_helpers.py::TestValidatePaddingLayout::test_batched_right_padding_invalid PASSED                             [ 28%]
test/quantization/recipes/integration/test_static_gemma4_runtime_helpers.py::TestValidatePaddingLayout::test_batched_right_padding_valid PASSED                               [ 33%]
test/quantization/recipes/integration/test_static_gemma4_runtime_helpers.py::TestValidatePaddingLayout::test_right_padding_invalid PASSED                                     [ 38%]
test/quantization/recipes/integration/test_static_gemma4_runtime_helpers.py::TestValidatePaddingLayout::test_right_padding_no_padding PASSED                                  [ 42%]
test/quantization/recipes/integration/test_static_gemma4_runtime_helpers.py::TestValidatePaddingLayout::test_right_padding_valid PASSED                                       [ 47%]
test/quantization/recipes/integration/test_static_gemma4_runtime_helpers.py::TestValidatePaddingLayout::test_unsupported_padding_side_raises PASSED                           [ 52%]
test/quantization/recipes/integration/test_static_gemma4_runtime_helpers.py::TestAllocateEmptyCache::test_all_full_attention_layers PASSED                                    [ 57%]
test/quantization/recipes/integration/test_static_gemma4_runtime_helpers.py::TestAllocateEmptyCache::test_all_sliding_layers PASSED                                           [ 61%]
test/quantization/recipes/integration/test_static_gemma4_runtime_helpers.py::TestAllocateEmptyCache::test_no_global_head_dim_falls_back PASSED                                [ 66%]
test/quantization/recipes/integration/test_static_gemma4_runtime_helpers.py::TestAllocateEmptyCache::test_per_layer_type_head_dim PASSED                                      [ 71%]
test/quantization/recipes/integration/test_static_gemma4_runtime_helpers.py::TestBuildDecodeMasksAndRope::test_full_attention_mask_shape PASSED                               [ 76%]
test/quantization/recipes/integration/test_static_gemma4_runtime_helpers.py::TestBuildDecodeMasksAndRope::test_full_attention_mask_values PASSED                              [ 80%]
test/quantization/recipes/integration/test_static_gemma4_runtime_helpers.py::TestBuildDecodeMasksAndRope::test_rope_computed_at_past_len PASSED                               [ 85%]
test/quantization/recipes/integration/test_static_gemma4_runtime_helpers.py::TestBuildDecodeMasksAndRope::test_sliding_window_mask_boundary_early_decode PASSED               [ 90%]
test/quantization/recipes/integration/test_static_gemma4_runtime_helpers.py::TestBuildDecodeMasksAndRope::test_sliding_window_mask_boundary_late_decode PASSED                [ 95%]
test/quantization/recipes/integration/test_static_gemma4_runtime_helpers.py::TestBuildDecodeMasksAndRope::test_sliding_window_mask_first_step PASSED                          [100%]

========================================================================== 21 passed, 2 warnings in 4.86s ===========================================================================

Example Script

$ python -m pdb ./tico/quantization/examples/inspector.py \
    --mode static-gemma4-runtime \
    --config ./tico/quantization/examples/configs/static_gemma4_runtime.yaml

[run_static_gemma4_runtime] Loading model: google/gemma-4-e2b-it
Warning: You are sending unauthenticated requests to the HF Hub. Please set a HF_TOKEN to enable higher rate limits and faster downloads.
Loading weights: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1951/1951 [00:00<00:00, 4661.76it/s]
[run_static_gemma4_runtime] Creating StaticGemma4Runtime ...
[run_static_gemma4_runtime] Step 1: verify build_static_inputs
[verify_step_build_static_inputs] All checks passed.
[run_static_gemma4_runtime] Step 8: verify generation
[generate_greedy] Next token: tensor([106])
[generate_greedy] Next token: tensor([106])
[generate_greedy] Next token: tensor([106])
[generate_greedy] Next token: tensor([106])
[generate_greedy] Next token: tensor([106])
[generate_greedy] Next token: tensor([106])
[generate_greedy] Next token: tensor([106])
[generate_greedy] Next token: tensor([106])
[generate_greedy] Next token: tensor([106])
[generate_greedy] Next token: tensor([106])
[generate_greedy] Next token: tensor([106])
[generate_greedy] Next token: tensor([106])
[generate_greedy] Next token: tensor([106])
[generate_greedy] Next token: tensor([106])
[generate_greedy] Next token: tensor([106])
[generate_greedy] Next token: tensor([106])
[verify_step_generation] Runtime generated 278 tokens
[verify_step_generation] HF generated 278 tokens
[verify_step_generation] Token accuracy: 100.00% (278/278)
[verify_step_generation] All tokens match!
[verify_step_generation] Runtime text: <|image><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><image|>Describe the image.<turn|><turn|><turn|><turn|><turn|><turn|><turn|><turn|><turn|><turn|><turn|><turn|><turn|><turn|><turn|><turn|>
[verify_step_generation] HF text: <|image><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><|image|><image|>Describe the image.<turn|><turn|><turn|><turn|><turn|><turn|><turn|><turn|><turn|><turn|><turn|><turn|><turn|><turn|><turn|><turn|>
[verify_step_generation] All checks passed.
[run_static_gemma4_runtime] Done.

Torrero · 2026-07-31T11:04:44Z

+    print(f"[verify_step_generation] Runtime text: {rt_text}")
+    print(f"[verify_step_generation] HF text: {hf_text}")
+
+    print("[verify_step_generation] All checks passed.")


This message will be printed even if there are some mismatches.

👍 Fixed.
Now the message accurately reflects the verification result:

"All checks passed." - printed only when all tokens match between runtime and HF generation

"Mismatches detected." - printed when token mismatches exist

Torrero · 2026-07-31T11:07:49Z

            dtype=hidden_states.dtype,
        )

+        valid_positions = (attention_masks["full_attention"] == 0.0).sum().item()


Where do you plan to use valid_positions? This variable is computed but it is never used.

👍 Removed the unused variable.

Torrero · 2026-07-31T11:25:16Z

+                # Write only valid K/V (excluding padding) into the cache.
+                # new_k/new_v from prefill have shape (B, kv_heads, max_seq, head_dim)
+                # but positions >= valid_length contain garbage from padding tokens.
+                valid_len = int(batch["valid_length"].item())


It would be better to move this statement outside the loop to avoid a GPU→CPU sync per layer.

👍 Moved valid_len = int(batch["valid_length"].item()) outside the prefill layer loop.

Implement greedy generation with side-by-side HF verification for Gemma4 E2B. Co-authored-by: Cline TICO-DCO-1.0-Signed-off-by: d.savchenkov <d.savchenkov@partner.samsung.com>

Torrero

LGTM

dvsav force-pushed the generate branch 3 times, most recently from 576b4f2 to acbdd28 Compare July 30, 2026 08:20

dvsav marked this pull request as ready for review July 30, 2026 08:41

dvsav requested a review from Torrero July 30, 2026 08:41

Torrero reviewed Jul 31, 2026

View reviewed changes

[quantization] Add generation step to StaticGemma4Runtime

09fe22d

Implement greedy generation with side-by-side HF verification for Gemma4 E2B. Co-authored-by: Cline TICO-DCO-1.0-Signed-off-by: d.savchenkov <d.savchenkov@partner.samsung.com>

dvsav force-pushed the generate branch from acbdd28 to 09fe22d Compare July 31, 2026 12:49

dvsav requested a review from Torrero July 31, 2026 12:55

Torrero approved these changes Jul 31, 2026

View reviewed changes

dvsav merged commit 5436e4c into Samsung:main Jul 31, 2026
7 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[quantization] Add generation step to StaticGemma4Runtime - #849

[quantization] Add generation step to StaticGemma4Runtime#849
dvsav merged 1 commit into
Samsung:mainfrom
dvsav:generate

dvsav commented Jul 28, 2026 •

edited

Loading

Uh oh!

Torrero Jul 31, 2026

Uh oh!

dvsav Jul 31, 2026

Uh oh!

Torrero Jul 31, 2026

Uh oh!

dvsav Jul 31, 2026

Uh oh!

Torrero Jul 31, 2026

Uh oh!

dvsav Jul 31, 2026

Uh oh!

Torrero left a comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

Conversation

dvsav commented Jul 28, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What

Why

Key Design Decisions

Changes

Tests

Example Script

Uh oh!

Torrero Jul 31, 2026

Choose a reason for hiding this comment

Uh oh!

dvsav Jul 31, 2026

Choose a reason for hiding this comment

Uh oh!

Torrero Jul 31, 2026

Choose a reason for hiding this comment

Uh oh!

dvsav Jul 31, 2026

Choose a reason for hiding this comment

Uh oh!

Torrero Jul 31, 2026

Choose a reason for hiding this comment

Uh oh!

dvsav Jul 31, 2026

Choose a reason for hiding this comment

Uh oh!

Torrero left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

dvsav commented Jul 28, 2026 •

edited

Loading