ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

benchmarks: gemma4_flex_inference -- drop sidecar, link shared layers to store cache

Browse source 
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 ...").
This commit is contained in:
Daniel Han 2026-04-21 09:40:12 +00:00
commit dee9371769
1 changed files with 109 additions and 148 deletions
Download patch file Download diff file Expand all files Collapse all files

253
scripts/benchmarks/gemma4_flex_inference.py
View file

@ -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