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.
8fb0c2e2a7