The sidecar design in the previous cut stored K/V at
[max_batch, n_kv, max_seq, D] layout, but the flex_attention block mask
is built for the paged cache's [1, H, n_pages*page_size, D] layout.
Shared layers running through that mismatch either needed a parallel
block mask (expensive to build per call, per layer) or had to fall back
to SDPA, which breaks the single-CUDA-graph capture story and silently
dropped the sliding-window mask on shared sliding layers.
This commit drops the sidecar entirely. Shared layers now reference the
store layer's `PagedKVCache` directly:
- `patch_gemma4_attention_forwards` allocates a cache on every
non-shared layer, then walks shared layers and points
`shared._paged_cache = store._paged_cache`, plus stashes the store
attention module itself on `shared._store_attn`.
- The store layer keeps its post-rotary `k`, `v` on
`self._last_k_val`, `self._last_v_val` before its own paged update
so shared successors can read the same packed prefill tensors.
- Shared-layer forward reads `_last_k_val` / `_last_v_val` on prefill
(q_len > 1) and `_paged_cache.k_cache` / `.v_cache` on decode. The
block_mask dispatched by `self.layer_type` works for both regimes
uniformly -- one block mask builder, one kernel compile per regime,
one CUDA graph per batch-size bucket.
Also fixes the 4-bit path: `AutoModelForCausalLM.from_pretrained` on
`unsloth/gemma-4-E2B-it-unsloth-bnb-4bit` resolves to
`Gemma4ForConditionalGeneration`, so `.model.embed_tokens` does not
exist. The loader now detects the multimodal wrapper, drops the vision
and audio towers, and moves the language_model into a
`Gemma4ForCausalLM` shell -- mirroring the bf16 path.
Benchmarks on a single B200 (sm_100), CUDA_VISIBLE_DEVICES=2, Gemma-4
E2B-it, n_prompts=64 n_rounds=5 max_new_tokens=512 max_batch_size=64
capture_cudagraph no-fa4_prefill BLOCK_M=32 BLOCK_N=32 (prefill) /
BLOCK_M=16 BLOCK_N=16 (decode):
| Config | Peak GB | Median tok/s | Best tok/s |
|-----------------|---------|--------------|------------|
| bf16 | 14.2 | 2794 | 2797 |
| bf16 + LoRA r32 | 23.0 | 2798 | 2801 |
| 4bit + LoRA r32 | 13.0 | 1659 | 1821 |
Drift verification (10 perturb+refresh cycles, noise_scale=0.01):
`base_bit_identical = true`, `inference_deterministic = true`.
Sample completions are coherent math reasoning ("Let the isosceles
trapezoid be $ABCD$ with bases ...").
The first cut routed every non-shared layer through one causal block mask
and relied on SDPA with is_causal=True for shared layers. That gives the
right semantics for Gemma-4's full_attention layers but silently drops the
sliding_attention window, so sliding layers attend far beyond their
512-token window as soon as the prefix grows.
This commit builds a block mask per attention regime (full_attention is
pure causal; sliding_attention is causal AND q_pos - kv_pos < window) and
passes both into each attention call as a dict, letting the patched
forward select by self.layer_type. For the shared-KV sidecar path, the
SDPA call now receives an explicit attn_mask composed the same way, so
sliding shared layers also respect the window. Strict less-than
comparison matches Unsloth's flex-attention convention for GPT-OSS.
`_causal_blockmask_with_window` and `_prefill_blockmask_with_window` are
local to this file; the shared helpers in `flex_paged_attention.py` stay
untouched. `FlexGemma4Inference` now caches both logical decode masks,
slices both per-row in `_decode_block_mask`, and runs the PageTable's
logical->physical conversion on each before passing them down.
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.
Daniel Han