[quantization] Add prefill verification step to StaticGemma4Runtime

[dvsav]

Related issue: #768
Related PR: #822

What

Implements the prefill verification step (Step 6) in the StaticGemma4Runtime for Gemma4 E2B. The runtime's prefill method now correctly handles PLE (Per-Layer Embeddings), per-layer-type KV cache head dimensions, and shared-KV bookkeeping. A new verify_step_prefill function performs side-by-side validation of runtime prefill logits against the HuggingFace FP reference, and Step 6 is wired into the run_static_gemma4_runtime runner.

Why

The static runtime's prefill path had four correctness gaps that prevented it from producing valid logits:

1. Missing PLE: When hidden_size_per_layer_input > 0, each decoder layer requires a per-layer input tensor. The runtime was not computing or passing these, causing a runtime error.
2. KV cache head_dim mismatch: Gemma4 uses different head dimensions per layer type (sliding = head_dim, full_attention = global_head_dim). The cache allocator used a single head_dim for all layers, causing a tensor shape mismatch.
3. Missing shared-KV: Some Gemma4 layers share K/V states with earlier layers of the same layer_type. The runtime did not maintain shared_kv_states or pass shared_key_value to consumer layers, causing a runtime error.
4. pixel_values dtype mismatch: pixel_values was not cast to the model's dtype, causing a dtype mismatch in the vision tower.

Without these fixes, the prefill path could not execute, blocking all downstream steps (decode, generation).

Key Design Decisions

• Per-layer-type KV cache allocation: _allocate_empty_cache now inspects layer_types[layer_idx] and selects head_dim from config.head_dim (sliding) or config.global_head_dim (full_attention). It also handles the general case where num_global_key_value_heads may differ, though for E2B this defaults to num_key_value_heads.
• PLE computed once, sliced per layer: get_per_layer_inputs and project_per_layer_inputs are called once before the layer loop to produce a (B, S, L, P) tensor. Each layer receives its slice per_layer_inputs[:, :, layer_idx, :].
• Shared-KV via a dict keyed by layer_type: A shared_kv_states: dict[str, tuple[Tensor, Tensor]] is maintained during the layer loop. Store layers (with store_full_length_kv=True) write their (new_k, new_v). Consumer layers (with is_kv_shared_layer=True) read from the dict and receive shared_key_value.
• Cache writes via copy_: For prefill, new_k/new_v cover the entire max_seq, so the full tensor is copied into the pre-allocated fixed cache via copy_() — no positional slicing is needed.
• final_logit_softcapping applied in verification, not in the adapter: The Gemma4LMHeadExportAdapter does not apply tanh(logits / sc) * sc. The verify_step_prefill function applies it post-hoc to match HF's output, keeping the adapter clean for NPU export.
• PEIR metric for comparison: Since the runtime uses quantized weights and the reference is FP, exact equality is not expected. PEIR (Peak-Error-to-Interval Ratio) and diff statistics are used to assess numerical closeness. PEIR is computed as compute_peir(ref_f, rt_f) — i.e., base = FP reference, target = quantized runtime — so the interval denominator is independent of quantization error.
• Cached raw processor output: build_static_inputs caches the raw (unpadded) processor output in batch["_raw_inputs"] so that verify_step_prefill can reuse it without re-running the processor (which involves expensive image preprocessing). A fallback re-runs the processor if the cache is not available.

Changes

tico/quantization/recipes/debug/static_gemma4_runtime.py (+236/-19):

1. __init__: Store self.text_model reference (from wrapped_model.language_model.wrapped) for PLE computation.
2. _allocate_empty_cache: Per-layer-type head_dim and num_kv_heads selection based on layer_types[layer_idx].
3. build_static_inputs: Cache raw processor output in batch["_raw_inputs"]; return type updated to dict[str, Any].
4. prefill:
   • Cast pixel_values to self.model.dtype.
   • Compute PLE via get_per_layer_inputs + project_per_layer_inputs when hidden_size_per_layer_input > 0.
   • Maintain shared_kv_states dict; pass per_layer_input and shared_key_value to each layer.
   • Handle both 3-tuple (non-shared) and single-tensor (shared) layer outputs.
   • Write full-sequence K/V into cache via copy_(); store full-length K/V from store layers into shared_kv_states.
5. verify_step_prefill (new function): Runs runtime prefill, applies final_logit_softcapping, re-runs HF Gemma4ForConditionalGeneration.forward on unpadded inputs (using cached raw processor output), and compares logits via PEIR / mean|diff| / max|diff| / argmax match.
6. run_static_gemma4_runtime: Wired Step 6 (verify_step_prefill) into the runner; updated docstring to reflect that prefill is now implemented.

test/quantization/recipes/integration/test_static_gemma4_runtime_helpers.py (+165):

New TestAllocateEmptyCache test class with 4 test cases:

• test_per_layer_type_head_dim: Mixed sliding/full layers — verifies correct head_dim per layer type.
• test_all_sliding_layers: All sliding → all caches use head_dim.
• test_all_full_attention_layers: All full_attention → all caches use global_head_dim.
• test_no_global_head_dim_falls_back: global_head_dim=None → all layers use head_dim.

Tests use SimpleNamespace mocks to test _allocate_empty_cache in isolation without requiring a full model.

Tests

All 15 existing helper tests in test/quantization/recipes/integration/test_static_gemma4_runtime_helpers.py continue to pass (11 original + 4 new TestAllocateEmptyCache 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-9.0.3, pluggy-1.6.0 -- /home/d.savchenkov/myenv/bin/python
cachedir: .pytest_cache
rootdir: /home/d.savchenkov/TICO
configfile: pyproject.toml
plugins: anyio-4.13.0
collected 15 items                                                                          

test/quantization/recipes/integration/test_static_gemma4_runtime_helpers.py::TestNormalizeValidTokenMask::test_batched_input PASSED [  6%]
test/quantization/recipes/integration/test_static_gemma4_runtime_helpers.py::TestNormalizeValidTokenMask::test_shape_mismatch_raises PASSED [ 13%]
test/quantization/recipes/integration/test_static_gemma4_runtime_helpers.py::TestNormalizeValidTokenMask::test_with_attention_mask PASSED [ 20%]
test/quantization/recipes/integration/test_static_gemma4_runtime_helpers.py::TestNormalizeValidTokenMask::test_without_attention_mask_no_pad_token_id PASSED [ 26%]
test/quantization/recipes/integration/test_static_gemma4_runtime_helpers.py::TestNormalizeValidTokenMask::test_without_attention_mask_uses_pad_token_id PASSED [ 33%]
test/quantization/recipes/integration/test_static_gemma4_runtime_helpers.py::TestValidatePaddingLayout::test_batched_right_padding_invalid PASSED [ 40%]
test/quantization/recipes/integration/test_static_gemma4_runtime_helpers.py::TestValidatePaddingLayout::test_batched_right_padding_valid PASSED [ 46%]
test/quantization/recipes/integration/test_static_gemma4_runtime_helpers.py::TestValidatePaddingLayout::test_right_padding_invalid PASSED [ 53%]
test/quantization/recipes/integration/test_static_gemma4_runtime_helpers.py::TestValidatePaddingLayout::test_right_padding_no_padding PASSED [ 60%]
test/quantization/recipes/integration/test_static_gemma4_runtime_helpers.py::TestValidatePaddingLayout::test_right_padding_valid PASSED [ 66%]
test/quantization/recipes/integration/test_static_gemma4_runtime_helpers.py::TestValidatePaddingLayout::test_unsupported_padding_side_raises PASSED [ 73%]
test/quantization/recipes/integration/test_static_gemma4_runtime_helpers.py::TestAllocateEmptyCache::test_all_full_attention_layers PASSED [ 80%]
test/quantization/recipes/integration/test_static_gemma4_runtime_helpers.py::TestAllocateEmptyCache::test_all_sliding_layers PASSED [ 86%]
test/quantization/recipes/integration/test_static_gemma4_runtime_helpers.py::TestAllocateEmptyCache::test_no_global_head_dim_falls_back PASSED [ 93%]
test/quantization/recipes/integration/test_static_gemma4_runtime_helpers.py::TestAllocateEmptyCache::test_per_layer_type_head_dim PASSED [100%]

============================== 15 passed, 2 warnings in 7.26s ===============================

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, 4241.73it/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 2: verify token_embedding
[verify_step_token_embedding] All checks passed.
[run_static_gemma4_runtime] Step 3: verify vision_prefill
[verify_step_vision_prefill] PEIR = 0.714285 %
[verify_step_vision_prefill] All checks passed.
[run_static_gemma4_runtime] Step 4: verify mm_fusion
[verify_step_mm_fusion] All checks passed.
[run_static_gemma4_runtime] Step 5: verify masks_and_rope
[verify_step_masks_and_rope] RoPE exact match for 2 layer types.
[verify_step_masks_and_rope] Mask patterns match for 2 layer types.
[verify_step_masks_and_rope] All checks passed.
[run_static_gemma4_runtime] Step 6: verify prefill
[verify_step_prefill] Runtime logits shape: (1, 262144)
[verify_step_prefill] Reference logits shape: (1, 262144)
[verify_step_prefill] mean|diff| = 0.05616927
[verify_step_prefill]  max|diff| = 0.50000000
[verify_step_prefill] PEIR       = 2.285714 %
[verify_step_prefill] Runtime argmax token: 106
[verify_step_prefill] Reference argmax token: 106
[verify_step_prefill] Argmax tokens MATCH.
[verify_step_prefill] All checks passed.
[run_static_gemma4_runtime] SKIP: decode (not implemented in this PR)
[run_static_gemma4_runtime] SKIP: generation (not implemented in this PR)
[run_static_gemma4_runtime] Done.

[Torrero]

LGTM

[mhs4670go]

The reference forward computes logits for every input position and slices the final position afterward. With the default max_seq=2048 and a 262,144-token vocabulary, this can allocate 1 GiB in BF16 or 2 GiB in FP32 just for the logits tensor. Please pass logits_to_keep=1 so the LM head runs only for the final token. use_cache=False should also be set because this verification does not consume the reference KV cache.

Suggested change

	ref_out = runtime.model(
	input_ids=raw_input_ids,
	pixel_values=pixel_values,
	image_position_ids=image_position_ids,
	attention_mask=ref_attention_mask,
	return_dict=True,
	)
	ref_logits = ref_out.logits[:, -1, :] # (1, vocab_size)
	ref_out = runtime.model(
	input_ids=raw_input_ids,
	pixel_values=pixel_values,
	image_position_ids=image_position_ids,
	attention_mask=ref_attention_mask,
	logits_to_keep=1,
	use_cache=False,
	return_dict=True,
	)
	ref_logits = ref_out.logits[:, -1, :]

[mhs4670go]

I think the runtime might as well does this instead. How about like this?

1. Introduce a shared helper.

def _apply_final_logit_softcapping(
    self, logits: torch.Tensor
) -> torch.Tensor:
    """Apply Gemma4 final logit softcapping."""
    softcap = getattr(
        self.text_config,
        "final_logit_softcapping",
        None,
    )
    if softcap is None:
        return logits

    softcap = float(softcap)
    return torch.tanh(logits / softcap) * softcap

2. Call the helper at the end of the prefill/decode function.

# def prefill
# ..

logits = self.lm_head(
    hidden_states[:, self.past_len - 1 : self.past_len, :]
)
logits = logits[:, -1, :]
return self._apply_final_logit_softcapping(logits)

# def decode_one
# ..

self.past_len += 1
logits = self.lm_head(hidden_states)[:, -1, :]
return self._apply_final_logit_softcapping(logits)

3. Remove the same logic from the verifier.