benchmarks: add gemma4_flex_inference for unsloth/gemma-4-E2B-it

Extends the flex_attention + paged KV + CUDA graphs engine from Qwen3 and
Llama-3.2 to Gemma-4-E2B-it via a new standalone file that imports the
shared helpers (PagedKVCache, PageTable, Sequence, LoRA double-copy,
drift verification, flex_attention_compiled, _apply_rotary) from
qwen3_flex_inference.py. The Qwen3 / Llama path is not modified.

Gemma-4 diverges from Qwen3 / Llama in ways that cannot be folded into a
single hasattr guard:

- KV-sharing layers. E2B has 35 layers; the upper 20 lack k_proj / v_proj
  / k_norm / v_norm and consume the full prefix K/V produced by a store
  layer further up the stack. We allocate a sidecar dict of
  [max_batch, n_kv, max_seq, head_dim] buffers at fixed device addresses,
  populated by store layers during prefill and read by shared layers
  through eager SDPA (their layout does not match the paged cache's
  block-mask shape).
- Dual attention regimes. full_attention (head_dim=512, rope_theta=1e6)
  and sliding_attention (head_dim=256, sliding_window=512) coexist. We
  precompute both (cos, sin) pairs via
  Gemma4TextRotaryEmbedding(x, position_ids, layer_type) and dispatch on
  self.layer_type inside the patched attention forward.
- Per-layer input embeddings. The walker threads the
  [B, S, num_layers, hidden_size_per_layer_input] table from
  get_per_layer_inputs + project_per_layer_inputs through each layer's
  per_layer_input_gate / act_fn / mul / per_layer_projection /
  post_per_layer_input_norm path.
- Four norms per block with double residuals. Attn (input_layernorm,
  post_attention_layernorm) and MLP (pre_feedforward_layernorm,
  post_feedforward_layernorm), plus the per-layer-input residual and
  layer_scalar multiply.
- Final logit softcap. tanh(logits / 30.0) * 30.0 on the lm_head output.

Transformers>=5.5.0 is required for the gemma4 module. A _require_gemma4
guard at main() exits with a clear install hint when the module is
missing, so the workspace's Qwen3 / Llama path stays on the existing
transformers install.

The text-only path loads Gemma4ForConditionalGeneration, drops the
vision and audio towers, and moves the language_model into a
Gemma4ForCausalLM shell so LoRA, state-dict hashing, and the double-copy
refresh treat it like any other HF decoder model.

CLI mirrors qwen3_flex_inference.py: --model_name (default
unsloth/gemma-4-E2B-it), --lora_adapter, --load_in_4bit,
--capture_cudagraph, --verify_no_drift, --chat_template {auto,grpo,
native}, --fa4_prefill / --no-fa4_prefill, plus decode / prefill
kernel_options for Triton block tuning.

Smoke-tested on B200 (sm_100) with bf16, batch 2, max_new_tokens 16,
no-fa4_prefill + BLOCK_M=32 / BLOCK_N=32 (Gemma-4 head_dim=256 exceeds
FA4's 128-limit on sm_100): 52 tok/s cold, coherent completions.

scripts/benchmarks/README.md gains a paragraph covering the
transformers>=5.5 dependency, the head_dim=256 constraint, and the
recommended kernel_options for B200.

scripts/benchmarks/README.md

@@ -96,6 +96,42 @@ flex_attention + paged KV + CUDA graphs stack is identical). Pass
 `--chat_template native` to use Llama's shipped Instruct template instead
 of the Qwen3 GRPO template.
 
+`gemma4_flex_inference.py` extends the engine to `unsloth/gemma-4-E2B-it`.
+Gemma-4 is not a drop-in: its text backbone has KV-sharing layers
+(layers 15-34 consume full-sequence K/V produced by a store layer), two
+attention regimes (`full_attention` with `head_dim=512` / rope_theta=1e6
+and `sliding_attention` with `head_dim=256` / sliding_window=512),
+per-layer input embeddings, four norms per block with double residuals,
+and a final logit softcap. The new file keeps the shared helpers
+(`PagedKVCache`, `PageTable`, LoRA double-copy, drift verification)
+imported from `qwen3_flex_inference.py` and adds:
+
+- a KV-sharing sidecar dict, sized `[max_batch, n_kv, max_seq, head_dim]`
+  per store layer, populated at prefill and read by the paired shared
+  layers through eager SDPA (shared layers don't fit the paged-cache
+  block-mask shape);
+- dual RoPE precomputation — `rotary_emb(x, pos, layer_type)` called once
+  per unique layer type, indexed by `self.layer_type`;
+- a walker that threads `per_layer_inputs` from
+  `get_per_layer_inputs` + `project_per_layer_inputs` into each layer
+  and applies the `layer_scalar` multiply at layer end;
+- `tanh(logits / final_logit_softcapping) * final_logit_softcapping`
+  applied on the lm_head output.
+
+Requires `transformers>=5.5.0` for the `gemma4` module; if absent the
+script exits with a clear install hint. Gemma-4 head_dim=256 exceeds FA4
+on sm_100 (B200), so pass `--no-fa4_prefill` and small Triton blocks:
+
+```bash
+CUDA_VISIBLE_DEVICES=6 python scripts/benchmarks/gemma4_flex_inference.py \
+    --model_name unsloth/gemma-4-E2B-it \
+    --n_prompts 64 --max_new_tokens 512 --capture_cudagraph \
+    --no-fa4_prefill \
+    --prefill_kernel_options '{"FORCE_USE_FLEX_ATTENTION": true, "BLOCK_M": 32, "BLOCK_N": 32}' \
+    --decode_kernel_options '{"BLOCK_M": 16, "BLOCK_N": 16}' \
+    --stats_path logs/flex_gemma4_bf16.json
+```
+
 | GPU          | arch      | sm    | Auto FA4 | Triton flex_attention |
 |--------------|-----------|-------|----------|------------------------|
 | A100         | Ampere    | sm_80 | off (uses Triton) | Works |