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 ...").
dee9371769