[quantization] Implement prefill masks and RoPE verification in StaticGemma4Runtime

[dvsav]

Related issue: #768
Related PR: #822

What

Implements real per-layer-type attention masks and RoPE (Rotary Position Embedding) generation in StaticGemma4Runtime, replacing the previous placeholder implementation. Adds Step 5 of the side-by-side verification flow (verify_step_masks_and_rope) that validates the runtime's masks and RoPE against Hugging Face transformers reference outputs.

Why

The static Gemma4 runtime requires CPU-owned, pre-computed attention masks and RoPE tensors that will be fed to NPU-executed decoder layers. Gemma4 E2B has a mixed layer-type architecture — a repeating pattern of 5 sliding-attention layers followed by 1 full-attention layer — each requiring distinct mask patterns and distinct RoPE configurations:

• full_attention: proportional RoPE (partial_rotary_factor=0.25, global_head_dim=512, rope_theta=1_000_000) with a standard causal mask.
• sliding_attention: default RoPE (head_dim=256, rope_theta=10_000) with a sliding-window causal mask (sliding_window=512).

The previous build_prefill_masks_and_rope returned shape-compatible zero/one placeholders. This PR replaces them with correct implementations and adds verification to ensure they exactly match HF's Gemma4TextRotaryEmbedding and create_masks_for_generate.

Key Design Decisions

1. Per-layer-type RoPE via HF module: The runtime calls the model's own Gemma4TextRotaryEmbedding module (self.model.model.language_model.rotary_emb) directly, passing the layer_type argument. This guarantees the runtime uses the exact same inv_freq buffers and attention_scaling that the HF model uses, avoiding reimplementation errors.

2. Key-only padding masking (matching HF convention): Both mask builders (_build_full_attention_mask and _build_sliding_window_attention_mask) apply padding only via key_valid (column masking). Query rows for padding positions are not zeroed out — they retain their causal/sliding-window pattern. This matches HF's create_masks_for_generate, which only controls which keys are visible, not which queries are active. Padding query outputs are simply discarded downstream. An earlier version incorrectly included query_valid, which truncated padding rows to all-False and caused a pattern mismatch.

3. Boolean pattern comparison for masks: HF may use different mask_value representations or return None depending on the attention implementation. The verification compares boolean attention patterns (allowed vs. forbidden) rather than exact float values, and temporarily forces _attn_implementation="eager" to obtain 4D float masks from HF's create_masks_for_generate.

4. Exact match for RoPE: Since the runtime and HF reference use the same module, RoPE comparison is an exact torch.testing.assert_close check.

Changes

All changes are in tico/quantization/recipes/debug/static_gemma4_runtime.py (293 insertions, 29 deletions):

• Added _build_full_attention_mask: Builds a standard causal mask with key-only padding masking. Returns (B, 1, S, S) additive bias.
• Added _build_sliding_window_attention_mask: Builds a sliding-window causal mask (kv_idx > q_idx - sliding_window) with key-only padding masking. Returns (B, 1, S, S) additive bias.
• Replaced placeholder build_prefill_masks_and_rope: Now computes position_ids = arange(seq_len), builds both mask types from the text config's sliding_window, and computes per-layer-type RoPE via the HF rotary_emb module. Returns (attention_masks, position_embeddings) dicts keyed by layer type.
• Added verify_step_masks_and_rope: Side-by-side verification comparing runtime RoPE (exact match) and masks (boolean pattern match) against HF reference. Returns (rt_masks, rt_rope) to avoid recomputation in downstream steps.
• Updated run_static_gemma4_runtime: Calls Step 5 after Step 4.

Tests

No new unit test files in this PR — verification is done via the example script's side-by-side comparison flow. All 5 verification steps pass:

[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] 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.

Example Script

python -m tico.quantization.recipes.debug.static_gemma4_runtime \
  --config tico/quantization/examples/configs/static_gemma4_runtime.yaml

[mhs4670go]

LGTM

[Torrero]

@dvsav
This comparison is unclear for me, in both cases you use HF's module rotary_emb with almost the same artificial dummy_hidden input. Could you please clarify your idea, or maybe I missed something.

[dvsav]

👍 I agree, but, I'll fix that in the following PR, if you don't mind.