diff --git a/scripts/benchmarks/gemma4_flex_inference.py b/scripts/benchmarks/gemma4_flex_inference.py index 6708cde860..337a544860 100644 --- a/scripts/benchmarks/gemma4_flex_inference.py +++ b/scripts/benchmarks/gemma4_flex_inference.py @@ -6,13 +6,13 @@ its text backbone diverges from Qwen3/Llama in ways that cannot be folded into a single `hasattr(self, "q_norm")` branch. The divergences, and how this file handles them: -1. KV-sharing layers. E2B has 35 layers; the upper 20 (layers 15-34) lack - `k_proj`, `v_proj`, `k_norm`, `v_norm` entirely and consume the full - prefix K/V produced by a "store" layer further up the stack. Paging - the KV for shared layers would cost more than it saves, so the shared - layers read from a pre-sized sidecar dict - (`FlexGemma4Inference.shared_kv_buffer`) whose tensors live at fixed - device addresses for CUDA graph safety. +1. KV-sharing layers. E2B has 35 layers; the upper 15 lack `k_proj`, + `v_proj`, `k_norm`, `v_norm` entirely and consume the K/V produced + by a "store" layer further up the stack. We link each shared layer's + `_paged_cache` to the store layer's `PagedKVCache` so flex_attention + reads the same pages that the store layer populated a few layers + earlier -- no sidecar, no SDPA fallback, one block mask per regime + works for every layer. 2. Dual attention types. Each layer is either `full_attention` (`head_dim=512`, `rope_theta=1e6`, `partial_rotary_factor=0.25`) or `sliding_attention` (`head_dim=256`, `rope_theta=10000`, @@ -159,8 +159,8 @@ def _require_gemma4(): def _apply_rotary_q(q, cos, sin): - """Rotary on Q alone; used when K comes from the shared sidecar and - already carries its rotary from the store layer.""" + """Rotary on Q alone; used on shared layers where K is read + pre-rotated from the store layer's paged cache.""" cos = cos.unsqueeze(1) sin = sin.unsqueeze(1) @@ -172,36 +172,33 @@ def _apply_rotary_q(q, cos, sin): return (q * cos) + (rotate_half(q) * sin) -def make_flex_gemma4_attention_forward( - page_table: PageTable, shared_kv_buffer: dict -): +def make_flex_gemma4_attention_forward(page_table: PageTable): """Return a new `forward` method for `Gemma4TextAttention` that routes - through flex_attention against either the paged cache (non-shared - layers) or the full-length shared KV sidecar. + through flex_attention against a paged KV cache. - Three layer kinds: - - shared (`self.is_kv_shared_layer == True`): no `k_proj`/`v_proj`/ - `k_norm`/`v_norm`. Read K/V from - `shared_kv_buffer[self.kv_shared_layer_index]`. - - store (`self.store_full_length_kv == True`): standard q/k/v - projection. After rotary, write the full-sequence K/V - into `shared_kv_buffer[self.layer_idx]` on prefill. - Also updates the paged cache for its own attention. - - plain (neither flag set): standard q/k/v + paged cache. + Two layer kinds: + - non-shared (`self.is_kv_shared_layer == False`): standard q/k/v + projection; writes new K/V into `self._paged_cache` + (which is the layer's own PagedKVCache). + - shared (`self.is_kv_shared_layer == True`): no `k_proj`/ + `v_proj`/`k_norm`/`v_norm`. Reads K/V directly from + the *store* layer's paged cache, which the patching + helper has already linked onto `self._paged_cache`. + No write -- the store layer populated the cache for + the same positions earlier in the walker, so the + shared layer just attends over those pages with the + identical block mask. - Shared layers attend over the prefix K/V only (populated at prefill). - Decode-time tokens generated by store layers are NOT written into the - sidecar -- this is the documented simplification from the plan; it - trades exactness during decode for a fixed-address sidecar that is - safe under CUDA graph capture. + Linking shared layers to the store layer's paged cache keeps one + block-mask + one KV layout across the whole stack and lets + flex_attention handle every layer uniformly (no sidecar, no SDPA + fallback, one CUDA graph capture). `position_embeddings` is a dict keyed by `layer_type`; we pick the right (cos, sin) pair before rotary. - Expects `self._paged_cache` to be set on non-shared layers (None on - shared layers). Shared layers still keep `self.q_proj`, `self.q_norm`. `self.v_proj` may be None when the config sets `attention_k_eq_v` - (global head dim with shared K=V); that branch reuses k_states. + (global head dim with shared K=V); that branch reuses k_raw. """ def forward( @@ -225,65 +222,26 @@ def make_flex_gemma4_attention_forward( q = self.q_norm(self.q_proj(hidden_states).view(hidden_shape)).transpose(1, 2) if getattr(self, "is_kv_shared_layer", False): - # Shared layer: read sidecar written by the paired store layer - # during prefill. Q still goes through rotary. The sidecar - # layout ([B, n_kv, max_seq, D]) does not match the paged - # cache's block-mask shape, so we route shared layers through - # the eager SDPA kernel instead of flex_attention. This is - # slower per-layer but keeps the sidecar design simple and - # CUDA-graph safe. + # Shared layer reads from the paired store layer's K/V. + # `PagedKVCache.update` returns different shapes for prefill + # vs decode: + # - prefill: the packed k_val/v_val [1, H, L_packed, D] + # - decode : the full paged pool k_cache/v_cache + # [1, H, n_pages*page_size, D] + # The prefill block_mask is sized for L_packed and the decode + # block_mask is sized for the paged pool, so we need to match + # the same shape here. The store layer stashes its + # post-rotary k/v as `_last_k_val` / `_last_v_val` during + # prefill; at decode time we read from its `_paged_cache` + # (the same buffer the shared layer was linked to at patch). q = _apply_rotary_q(q, cos, sin) - shared_k, shared_v = shared_kv_buffer[self.kv_shared_layer_index] - B = q.shape[0] - k = shared_k[:B] - v = shared_v[:B] - Hq = q.shape[1] - Hkv = k.shape[1] - if Hq != Hkv: - groups = Hq // Hkv - k = k.repeat_interleave(groups, dim = 1) - v = v.repeat_interleave(groups, dim = 1) - - # Build the attn_mask matching this layer's regime: - # - full_attention : pure causal. - # - sliding_attn : causal AND q_pos - kv_pos < window. - # For prefill (q_len > 1) the mask is a [q_len, kv_len] bool; - # for decode (q_len == 1) it becomes a [1, kv_len] row where - # the single q position is `input_pos` (passed in - # flex_input_pos) and kv_positions run 0..kv_len-1. - q_len = q.shape[-2] - kv_len = k.shape[-2] - window = getattr(self, "sliding_window", None) - if q_len > 1: - # Prefill: q_len == prefill_packed_len, kv_len should equal - # q_len in a clean run. Build per-batch positions. - q_pos = torch.arange(q_len, device = q.device) - kv_pos = torch.arange(kv_len, device = q.device) - attn_mask = q_pos[:, None] >= kv_pos[None, :] - if window is not None: - attn_mask = attn_mask & (q_pos[:, None] - kv_pos[None, :] < window) + store_attn = self._store_attn + if q.shape[-2] > 1: + k = store_attn._last_k_val + v = store_attn._last_v_val else: - # Decode: q_pos per batch comes from flex_input_pos. - q_pos = flex_input_pos.view(B, 1) # [B, 1] - kv_pos = torch.arange(kv_len, device = q.device)[None, :] - attn_mask = q_pos >= kv_pos - if window is not None: - attn_mask = attn_mask & (q_pos - kv_pos < window) - # SDPA expects mask shape [B, 1, 1, kv_len]. - attn_mask = attn_mask.unsqueeze(1).unsqueeze(1) - - attn_output = F.scaled_dot_product_attention( - q, - k, - v, - attn_mask = attn_mask, - scale = self.scaling, - ) - attn_output = ( - attn_output.transpose(1, 2).reshape(*input_shape, -1).contiguous() - ) - return self.o_proj(attn_output), None - + k = self._paged_cache.k_cache + v = self._paged_cache.v_cache else: k_raw = self.k_proj(hidden_states).view(hidden_shape) k = self.k_norm(k_raw).transpose(1, 2) @@ -297,15 +255,14 @@ def make_flex_gemma4_attention_forward( v = k_raw.transpose(1, 2) q, k = _apply_rotary(q, k, cos, sin) - # Prefill-only sidecar write for store layers. + # Store layers stash the post-rotary k/v so any shared + # successors can read the same packed prefill tensors. This + # assignment is a pointer rebind, not a copy; CUDA graph + # capture sees a stable attribute reference. Plain + # non-shared layers don't need this. if getattr(self, "store_full_length_kv", False): - is_prefill = q.shape[-2] > 1 - if is_prefill: - shared_k, shared_v = shared_kv_buffer[self.layer_idx] - B = k.shape[0] - S = k.shape[-2] - shared_k[:B, :, :S, :].copy_(k) - shared_v[:B, :, :S, :].copy_(v) + self._last_k_val = k + self._last_v_val = v if self._paged_cache is not None and flex_input_pos is not None: k, v = self._paged_cache.update(flex_input_pos, k, v, flex_batch_idx) @@ -329,26 +286,44 @@ def make_flex_gemma4_attention_forward( def patch_gemma4_attention_forwards( - model: torch.nn.Module, page_table: PageTable, shared_kv_buffer: dict + model: torch.nn.Module, page_table: PageTable ): - """Attach a PagedKVCache to every non-shared attention layer and swap - in the flex_attention forward above. Shared layers get `_paged_cache = - None` because their K/V comes from the sidecar. + """Attach a PagedKVCache to every non-shared attention layer, link + every shared attention layer to its store layer's cache, and swap in + the flex_attention forward above. + + Three passes: + 1. Allocate a PagedKVCache on each non-shared layer (the cache + shape depends on that layer's head_dim and num_kv_heads, which + vary across Gemma-4 layers). + 2. Walk shared layers and set + `shared._paged_cache = store._paged_cache`, where `store` is + `model.model.layers[shared.kv_shared_layer_index]`. The shared + layer's forward reads `k_cache` / `v_cache` directly; the store + layer's `update()` writes populate the same tensors. + 3. Bind the flex forward. """ - fwd = make_flex_gemma4_attention_forward(page_table, shared_kv_buffer) + fwd = make_flex_gemma4_attention_forward(page_table) for layer in model.model.layers: attn = layer.self_attn if getattr(attn, "is_kv_shared_layer", False): - attn._paged_cache = None - else: - n_kv = attn.k_proj.out_features // attn.head_dim - attn._paged_cache = PagedKVCache( - page_table, - n_heads = n_kv, - head_dim = attn.head_dim, - dtype = model.dtype, - ).to(model.device) - attn.forward = types.MethodType(fwd, attn) + continue + n_kv = attn.k_proj.out_features // attn.head_dim + attn._paged_cache = PagedKVCache( + page_table, + n_heads = n_kv, + head_dim = attn.head_dim, + dtype = model.dtype, + ).to(model.device) + for layer in model.model.layers: + attn = layer.self_attn + if not getattr(attn, "is_kv_shared_layer", False): + continue + store_attn = model.model.layers[attn.kv_shared_layer_index].self_attn + attn._paged_cache = store_attn._paged_cache + attn._store_attn = store_attn + for layer in model.model.layers: + layer.self_attn.forward = types.MethodType(fwd, layer.self_attn) # --- model forward walker -------------------------------------------------- @@ -496,38 +471,7 @@ class FlexGemma4Inference: device = self.device.type, ) - # Allocate shared-KV sidecar buffers. Only "store" layers get an - # entry; shared layers read by the store layer's `layer_idx`. - # Buffers live at fixed device addresses, so CUDA graph replay - # reads stable pointers -- safe because store layers only write - # during prefill (not inside captured decode graphs) and shared - # layers only read. - # - # Per-layer num_kv_heads is derived from `k_proj.out_features / - # head_dim` rather than `config.num_key_value_heads`, because - # global-attention layers may use `num_global_key_value_heads` - # under `attention_k_eq_v`. - self.shared_kv_buffer: dict = {} - for i, layer in enumerate(model.model.layers): - attn = layer.self_attn - if not getattr(attn, "store_full_length_kv", False): - continue - hd = attn.head_dim - n_kv = attn.k_proj.out_features // hd - K = torch.zeros( - max_batch_size, - n_kv, - max_seq_length, - hd, - dtype = model.dtype, - device = self.device, - ) - V = torch.zeros_like(K) - self.shared_kv_buffer[i] = (K, V) - - patch_gemma4_attention_forwards( - model, self.page_table, self.shared_kv_buffer - ) + patch_gemma4_attention_forwards(model, self.page_table) self.input_pos_buffer = torch.zeros( max_batch_size, dtype = torch.int32, device = self.device @@ -983,14 +927,31 @@ def main(): f"or drop --load_in_4bit for bf16." ) print(f"[flex-gemma4] loading 4-bit base: {bnb_model_name}") - model = AutoModelForCausalLM.from_pretrained( + # AutoModelForCausalLM resolves to `Gemma4ForConditionalGeneration` + # for Gemma-4 -- mirror the bf16 path and move the + # language_model into a ForCausalLM shell so downstream code can + # reach `model.model.layers` / `model.model.embed_tokens`. + full_model = AutoModelForCausalLM.from_pretrained( bnb_model_name, attn_implementation = "eager", device_map = "cuda:0", ) - if getattr(model.config, "tie_word_embeddings", False): - model.lm_head.weight = model.model.embed_tokens.weight + if hasattr(full_model.model, "language_model"): + lang_model = full_model.model.language_model + full_model.model.vision_tower = None + full_model.model.audio_tower = None + full_model.model.embed_vision = None + full_model.model.embed_audio = None + text_cfg = full_model.config.text_config + model = Gemma4ForCausalLM(text_cfg) + model.model = lang_model + model.lm_head.weight = lang_model.embed_tokens.weight + else: + model = full_model + if getattr(model.config, "tie_word_embeddings", False): + model.lm_head.weight = model.model.embed_tokens.weight model.eval() + del full_model if args.lora_adapter: from peft import PeftModel