From e4136a4503f269a63f2d1a72637379ffc19a77c3 Mon Sep 17 00:00:00 2001 From: "d.savchenkov" Date: Wed, 22 Jul 2026 11:17:13 +0300 Subject: [PATCH] [quantization] Implement prefill masks and RoPE verification in StaticGemma4Runtime Add real per-layer-type attention masks and RoPE to StaticGemma4Runtime. Co-authored-by: Cline TICO-DCO-1.0-Signed-off-by: d.savchenkov --- .../recipes/debug/static_gemma4_runtime.py | 322 ++++++++++++++++-- 1 file changed, 293 insertions(+), 29 deletions(-) diff --git a/tico/quantization/recipes/debug/static_gemma4_runtime.py b/tico/quantization/recipes/debug/static_gemma4_runtime.py index a2b2e08b..a81547b6 100644 --- a/tico/quantization/recipes/debug/static_gemma4_runtime.py +++ b/tico/quantization/recipes/debug/static_gemma4_runtime.py @@ -97,6 +97,90 @@ def _validate_padding_layout( raise ValueError("Right padding expected but not found") +def _build_full_attention_mask( + valid_token_mask: torch.Tensor, + device: torch.device, + dtype: torch.dtype, + mask_value: float, +) -> torch.Tensor: + """Build a standard causal attention mask with padding. + + Args: + valid_token_mask: Boolean mask of shape ``(B, S)`` — True for real + tokens, False for padding. + device: Target device. + dtype: Target dtype for the output mask. + mask_value: Value used to mask forbidden positions (typically + ``torch.finfo(dtype).min``). + + Returns: + Additive attention bias of shape ``(B, 1, S, S)``. ``0.0`` allows + attention, ``mask_value`` forbids it. + """ + batch_size, seq_len = valid_token_mask.shape + # Causal: kv_idx <= q_idx + causal = torch.tril( + torch.ones(seq_len, seq_len, device=device, dtype=torch.bool) + ) # (S, S) + causal = causal.unsqueeze(0).expand(batch_size, -1, -1) # (B, S, S) + + key_valid = valid_token_mask.to(device).unsqueeze(1).expand(-1, seq_len, -1) + + valid = causal & key_valid + mask = torch.zeros(batch_size, seq_len, seq_len, device=device, dtype=dtype) + mask = mask.masked_fill(~valid, mask_value) + return mask.unsqueeze(1) # (B, 1, S, S) + + +def _build_sliding_window_attention_mask( + valid_token_mask: torch.Tensor, + sliding_window: int, + device: torch.device, + dtype: torch.dtype, + mask_value: float, +) -> torch.Tensor: + """Build a sliding-window causal attention mask with padding. + + A token at position ``q`` can attend to position ``k`` if: + - ``k <= q`` (causal), AND + - ``k > q - sliding_window`` (within the sliding window) + + This matches HF's ``sliding_window_causal_mask_function`` which is + ``and_masks(sliding_window_overlay, causal_mask_function)`` where + ``sliding_window_overlay`` returns ``kv_idx > q_idx - sliding_window``. + + Args: + valid_token_mask: Boolean mask of shape ``(B, S)``. + sliding_window: Window size (number of past tokens visible, inclusive + of the current token). + device: Target device. + dtype: Target dtype. + mask_value: Value for masked positions. + + Returns: + Additive attention bias of shape ``(B, 1, S, S)``. + """ + batch_size, seq_len = valid_token_mask.shape + + q_idx = torch.arange(seq_len, device=device).view(seq_len, 1) # (S, 1) + kv_idx = torch.arange(seq_len, device=device).view(1, seq_len) # (1, S) + + # Causal: kv_idx <= q_idx + causal = kv_idx <= q_idx # (S, S) broadcast + # Sliding window: kv_idx > q_idx - sliding_window + window = kv_idx > (q_idx - sliding_window) # (S, S) broadcast + + valid_pattern = causal & window # (S, S) + valid_pattern = valid_pattern.unsqueeze(0).expand(batch_size, -1, -1) # (B, S, S) + + key_valid = valid_token_mask.to(device).unsqueeze(1).expand(-1, seq_len, -1) + + valid = valid_pattern & key_valid + mask = torch.zeros(batch_size, seq_len, seq_len, device=device, dtype=dtype) + mask = mask.masked_fill(~valid, mask_value) + return mask.unsqueeze(1) # (B, 1, S, S) + + # ============================================================================= # Data Classes # ============================================================================= @@ -144,6 +228,7 @@ def __init__( self.model = model.eval().to(device) self.processor = processor + self.layout = layout self.device = torch.device(device) self.config = model.config @@ -345,40 +430,75 @@ def build_prefill_masks_and_rope( ) -> tuple[dict[str, torch.Tensor], dict[str, tuple[torch.Tensor, torch.Tensor]]]: """Build CPU-owned static masks and RoPE tensors for prefill. - This skeleton returns shape-compatible placeholder tensors so downstream - runtime code and linters can type-check while the exact Gemma4 mask and - RoPE implementation is developed. The final implementation should - replace this method with full/sliding attention masks and layer-type - specific RoPE generated from the Gemma4 text configuration. + Produces two dictionaries keyed by layer type + (``"full_attention"``, ``"sliding_attention"``): + + - ``attention_masks``: 4D additive bias tensors of shape + ``(batch, 1, seq_len, seq_len)``. ``0.0`` allows attention, + ``mask_value`` forbids it. Full-attention layers get a standard + causal mask; sliding-attention layers get a sliding-window causal + mask. Both masks also zero out padding positions. + + - ``position_embeddings``: ``(cos, sin)`` pairs of shape + ``(batch, seq_len, head_dim)`` per layer type, computed via the HF + model's ``Gemma4TextRotaryEmbedding`` module. Full-attention layers + use proportional RoPE (partial_rotary_factor=0.25, global_head_dim); + sliding-attention layers use default RoPE (head_dim). """ batch_size, seq_len = input_ids.shape - runtime_dtype = torch.float32 - if attention_mask.is_floating_point(): - runtime_dtype = attention_mask.dtype - - full_mask = torch.zeros( - batch_size, - 1, - seq_len, - seq_len, - device=self.device, - dtype=runtime_dtype, - ) - head_dim = int( - getattr( - self.text_config, - "head_dim", - self.text_config.hidden_size // self.text_config.num_attention_heads, - ) + + # --- Determine runtime dtype from model weights --- + runtime_dtype = next(self.model.parameters()).dtype + + # --- Valid token mask (boolean) --- + valid_token_mask = attention_mask.to(self.device).bool() # (B, S) + + # --- Position IDs: arange(seq_len) for prefill (no past cache) --- + position_ids = torch.arange(seq_len, device=self.device).unsqueeze(0) # (1, S) + + # --- Attention masks --- + mask_value = torch.finfo(runtime_dtype).min + + # Full attention: standard causal + padding + full_mask = _build_full_attention_mask( + valid_token_mask, self.device, runtime_dtype, mask_value ) - cos = torch.ones( - batch_size, seq_len, head_dim, device=self.device, dtype=runtime_dtype + + # Sliding attention: sliding-window causal + padding + sliding_window = int(getattr(self.text_config, "sliding_window", 1024)) + sliding_mask = _build_sliding_window_attention_mask( + valid_token_mask, + sliding_window, + self.device, + runtime_dtype, + mask_value, ) - sin = torch.zeros_like(cos) layer_types = set(getattr(self.text_config, "layer_types", ["full_attention"])) - attention_masks = {layer_type: full_mask for layer_type in layer_types} - position_embeddings = {layer_type: (cos, sin) for layer_type in layer_types} + attention_masks = { + "full_attention": full_mask, + "sliding_attention": sliding_mask, + } + # If the model only has one layer type, keep only that entry + if not hasattr(self.text_config, "layer_types"): + attention_masks = {"full_attention": full_mask} + + # --- RoPE (per layer type) --- + # Use the HF model's Gemma4TextRotaryEmbedding module directly. + # It has per-layer-type inv_freq buffers and attention_scaling. + rotary_emb = self.model.model.language_model.rotary_emb + + # Dummy hidden states for shape/dtype inference (rotary_emb only uses + # x.device and x.dtype) + dummy_hidden = torch.zeros( + batch_size, seq_len, 1, device=self.device, dtype=runtime_dtype + ) + + position_embeddings: dict[str, tuple[torch.Tensor, torch.Tensor]] = {} + for layer_type in layer_types: + cos, sin = rotary_emb(dummy_hidden, position_ids, layer_type) + position_embeddings[layer_type] = (cos, sin) + return attention_masks, position_embeddings @torch.no_grad() @@ -848,6 +968,146 @@ def verify_step_mm_fusion( return rt_fused +@torch.no_grad() +def verify_step_masks_and_rope( + runtime: StaticGemma4Runtime, + batch: dict[str, torch.Tensor], +) -> tuple[dict[str, torch.Tensor], dict[str, tuple[torch.Tensor, torch.Tensor]]]: + """Side-by-side validation of ``build_prefill_masks_and_rope`` against HF. + + Validates two things: + + 1. **Attention masks**: The runtime's full-causal and sliding-window-causal + masks are compared against HF's ``create_masks_for_generate``. Since HF + may return ``None`` (for sdpa) or use a different ``mask_value``, we + compare boolean attention patterns (allowed vs. forbidden) rather than + exact float values. + + 2. **RoPE (cos, sin)**: The runtime's per-layer-type RoPE is compared + against HF's ``Gemma4TextRotaryEmbedding.forward`` — the same module the + runtime uses internally. This is an exact-match check. + + HF reference (modeling_gemma4.py L1688–1707): + ``causal_mask_mapping = {`` + `` "full_attention": create_causal_mask(...),`` + `` "sliding_attention": create_sliding_window_causal_mask(...),`` + ``}`` + ``position_embeddings[layer_type] = self.rotary_emb(hidden_states, position_ids, layer_type)`` + + Args: + runtime: The ``StaticGemma4Runtime`` instance. + batch: The batch dict returned by ``build_static_inputs``. + + Returns the ``(attention_masks, position_embeddings)`` tuple from + ``build_prefill_masks_and_rope``. + """ + import torch.testing + + from transformers.masking_utils import create_masks_for_generate + + llm_input_ids = batch["llm_input_ids"].to(runtime.device) + attention_mask = batch["attention_mask"].to(runtime.device) + max_seq = llm_input_ids.shape[1] + runtime_dtype = next(runtime.model.parameters()).dtype + + # --- Runtime side --- + rt_masks, rt_rope = runtime.build_prefill_masks_and_rope( + llm_input_ids, attention_mask + ) + + # --- HF reference: RoPE --- + # Use the HF model's rotary_emb directly (same module the runtime uses) + rotary_emb = runtime.model.model.language_model.rotary_emb + position_ids = torch.arange(max_seq, device=runtime.device).unsqueeze(0) + dummy_hidden = torch.zeros( + 1, max_seq, 1, device=runtime.device, dtype=runtime_dtype + ) + + layer_types = set(runtime.text_config.layer_types) + for layer_type in layer_types: + ref_cos, ref_sin = rotary_emb(dummy_hidden, position_ids, layer_type) + rt_cos, rt_sin = rt_rope[layer_type] + + torch.testing.assert_close( + rt_cos, + ref_cos, + msg=f"RoPE cos mismatch for layer_type={layer_type!r}", + ) + torch.testing.assert_close( + rt_sin, + ref_sin, + msg=f"RoPE sin mismatch for layer_type={layer_type!r}", + ) + + print( + f"[verify_step_masks_and_rope] RoPE exact match for {len(layer_types)} layer types." + ) + + # --- HF reference: masks --- + # Build dummy inputs_embeds for shape inference + hidden_size = int(runtime.text_config.hidden_size) + dummy_embeds = torch.zeros( + 1, max_seq, hidden_size, device=runtime.device, dtype=runtime_dtype + ) + + # create_masks_for_generate needs a 2D attention mask (1=valid, 0=padding) + attn_mask_2d = attention_mask.to(runtime.device).long() # (1, max_seq) + + # Temporarily force eager attention to get 4D float masks from HF + original_attn_impl = getattr(runtime.text_config, "_attn_implementation", "sdpa") + runtime.text_config._attn_implementation = "eager" + try: + ref_masks = create_masks_for_generate( + config=runtime.text_config, + inputs_embeds=dummy_embeds, + attention_mask=attn_mask_2d, + past_key_values=None, + position_ids=position_ids, + ) + finally: + runtime.text_config._attn_implementation = original_attn_impl + + # --- Compare masks as boolean patterns --- + # HF may return None for some layer types (e.g. if no masking needed). + # We compare the allowed/forbidden pattern, not exact float values. + for layer_type in layer_types: + rt_mask = rt_masks[layer_type] # (1, 1, S, S) + + if isinstance(ref_masks, dict): + ref_mask = ref_masks.get(layer_type, None) + else: + # Single mask (all layers same type) + ref_mask = ref_masks + + if ref_mask is None: + # HF returned None — means no masking (all positions allowed). + # This happens when there's no padding and sdpa handles causality + # internally. Compare against our mask's "allowed" pattern. + # For prefill with no padding, all causal positions should be 0.0. + print( + f"[verify_step_masks_and_rope] HF returned None mask for " + f"{layer_type!r}, skipping mask comparison (RoPE still verified)." + ) + continue + + # Convert to boolean: True = allowed (not mask_value), False = masked + mask_value = torch.finfo(runtime_dtype).min + rt_allowed = rt_mask.squeeze(1).squeeze(0) > (mask_value / 2) # (S, S) + ref_allowed = ref_mask.squeeze(1).squeeze(0) > (mask_value / 2) # (S, S) + + torch.testing.assert_close( + rt_allowed, + ref_allowed, + msg=f"Attention mask pattern mismatch for layer_type={layer_type!r}", + ) + + print( + f"[verify_step_masks_and_rope] Mask patterns match for {len(layer_types)} layer types." + ) + print("[verify_step_masks_and_rope] All checks passed.") + return rt_masks, rt_rope + + def run_static_gemma4_runtime(cfg: StaticGemma4RuntimeConfig) -> None: """Run the Gemma4 E2B static runtime smoke flow. @@ -905,7 +1165,11 @@ def run_static_gemma4_runtime(cfg: StaticGemma4RuntimeConfig) -> None: runtime, text_embeds, visual_embeds, cfg.prompt, image ) - # --- Steps 5+: prefill / decode / generation (not yet implemented) --- + # --- Step 5: masks_and_rope validation --- + print("[run_static_gemma4_runtime] Step 5: verify masks_and_rope") + masks, rope = verify_step_masks_and_rope(runtime, batch) + + # --- Steps 6+: prefill / decode / generation (not yet implemented) --- print("[run_static_gemma4_runtime] SKIP: prefill (not implemented in this PR)") print("[run_static_gemma4_runtime] SKIP: decode (not implemented in this PR)") print("[run_static_gemma4_runtime] SKIP: generation (not implemented in this PR)")