[quantization] Add token embedding, vision prefill, and mm_fusion verification steps in Gemma4 static runtime

[dvsav]

Related issue: #768
Related PR: #822

What

This PR adds three new side-by-side verification steps to the StaticGemma4Runtime for the Gemma4 E2B model: token_embedding, vision_prefill, and mm_fusion. Each step independently validates a runtime adapter's output against the corresponding HF transformers reference implementation, ensuring numerical equivalence before the static runtime is used for actual inference. It also refactors all verify_step_* functions to return their computed results, eliminating redundant recomputation across steps.

Why

The StaticGemma4Runtime is being built incrementally, one workflow step at a time. Before proceeding to prefill, decode, and generation, we need confidence that each upstream component (token embedding, vision tower, multimodal fusion) produces identical results to the HF reference. Side-by-side validation catches subtle bugs early — such as dtype mismatches, incorrect embedding scaling, or wrong fusion semantics — rather than debugging them through the full prefill pipeline. The refactoring to return results (instead of bool) eliminates redundant token_embedding and vision_prefill calls that were previously duplicated in the mm_fusion step.

Key Design Decisions

1. Return computed results, not bool: Since verification failures raise exceptions (assert_close, raise ValueError), returning True was meaningless. Each verify_step_* now returns its runtime-computed tensor (or batch dict), allowing downstream steps to reuse it without recomputation.

2. vision_prefill uses PEIR, not exact match, for HF comparison: The runtime's vision adapter wraps a quantized vision tower (with fake Q-DQ ops), so exact equality against the FP HF model is not expected. PEIR (Peak-Error-to-Interval Ratio) quantifies the deviation. An additional exact-match assertion against the quantized model's get_image_features ensures the adapter itself introduces no error beyond quantization.

3. mm_fusion compares fixed_slot_fuse vs HF masked_scatter: The runtime uses a static torch.cat-based fusion (fixed contiguous slot), while HF uses dynamic masked_scatter. The verification re-runs the processor to recover raw input_ids with image tokens, builds the image mask, and applies masked_scatter to confirm both paths produce identical fused embeddings.

4. pixel_values cast to model dtype: The HF processor outputs float32, but the quantized vision tower weights are BFloat16. The vision_prefill step casts pixel_values to model.dtype before passing to both runtime and HF reference sides.

5. mm_fusion accepts text_embeds and visual_embeds as parameters rather than recomputing them internally, following the "reuse results" principle.

Changes

• tico/quantization/recipes/debug/static_gemma4_runtime.py
  • verify_step_build_static_inputs: Changed return type from bool to dict[str, torch.Tensor], returns the batch dict.
  • Added verify_step_token_embedding: Validates Gemma4TokenEmbeddingExportAdapter output against HF Gemma4TextScaledWordEmbedding, including a sanity check that sqrt(hidden_size) embedding scale is applied. Returns text_embeds tensor.
  • Added verify_step_vision_prefill: Validates Gemma4VisionPrefillExportAdapter output against HF model.get_image_features().pooler_output (with PEIR metric) and quantized model's get_image_features (exact match). Casts pixel_values to model dtype. Returns visual_embeds tensor.
  • Added verify_step_mm_fusion: Validates Gemma4MMFusionExportAdapter (fixed_slot_fuse) against HF masked_scatter reference. Accepts text_embeds and visual_embeds as parameters. Returns fused_embeds tensor.
  • run_static_gemma4_runtime: Updated to chain return values — batch from Step 1, text_embeds from Step 2, visual_embeds from Step 3, passed to Step 4. Removed redundant build_static_inputs call.

Tests

• All 11 existing tests in test/quantization/recipes/integration/test_static_gemma4_runtime_helpers.py continue to pass (tests _normalize_valid_token_mask and _validate_padding_layout helper functions).
• Manual end-to-end verification via the inspector script confirms all 4 steps pass:
  • Step 1 (build_static_inputs): All checks passed.
  • Step 2 (token_embedding): All checks passed.
  • Step 3 (vision_prefill): PEIR = 7.142857e-03, all checks passed.
  • Step 4 (mm_fusion): All checks passed.

Example Script

$ python ./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, 4382.96it/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 = 7.142857e-03
[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] SKIP: prefill (not implemented in this PR)
[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]

There is mismatch in the description.
Actually, you test this assertion between runtime and quantized model not HF model.

[dvsav]

👍 Thank you for noticing! The docstring was updated to accurately describe the two-tier verification - PEIR (informational) against the HF FP model, and assert_close (hard assertion) against the quantized model - rather than claiming the function asserts exact match against HF.

[Torrero]

This verification does not use visual_start_idx and num_visual_tokens, this differs from runtime.mm_fusion.

[dvsav]

👍 Added an explicit check that the image token positions in raw_input_ids match the static layout (visual_start_idx and num_visual_tokens). Previously, the masked_scatter reference path implicitly assumed the image tokens were contiguous at visual_start_idx (validated only by Step 1's build_static_inputs). Without this check, a layout mismatch would produce a confusing "mm_fusion mismatch" assertion error instead of a clear diagnostic. The new check raises a descriptive ValueError if the start position or token count doesn't match the static layout.

[mhs4670go]

The first argument should be reference value: hf_visual_embeds.

[dvsav]

👍 Thanks for catching that! Fixed the compute_peir argument order in verify_step_vision_prefill. The compute_peir(base, target) function is asymmetric — base is the reference used for both the numerator and the interval denominator, while target is the value being compared (used only in the numerator). The original code had the arguments swapped (compute_peir(rt_visual_embeds, hf_visual_embeds)), which incorrectly computed the interval denominator from the quantized runtime output instead of the HF FP reference. The fix swaps them to compute_peir(hf_visual_embeds, rt_visual_embeds), ensuring the interval is computed from the FP reference as the metric's design intends.

[Torrero]

LGTM, thank you!

[mhs4670go]

LGTM