Add static KV cache (TensorScatter) support for Gemma4 - #246
Conversation
Performance Comparison
|
Codecov Report❌ Patch coverage is 📢 Thoughts on this report? Let us know! |
🏗️ Architecture Diff
No architecture changes detected. ✅ Legend: ⚪ No change · 🔵 Minor (attrs/inits) · 🟡 Moderate (nodes added/removed) · 🔴 Major (interface changed) |
Code Review🔴 Critical — Sliding-window masking silently disabled in static cache mode
The static-cache path calls Attention with The PR description calls this a "Phase 1 limitation matching the generic DecoderLayer," but the dynamic-cache GQA path here does enforce sliding via Suggested fix: either (a) reject 🟡 Major — Encoder-skip heuristic in
|
|
Thanks for the thorough review @titaiwangms! All three findings addressed in commit f50b2b8: 🟡 Major — Fragile encoder-skip heuristicReplaced the path-prefix blacklist ( 🟢 Minor —
|
This comment was marked as resolved.
This comment was marked as resolved.
|
Update: The hybrid per-layer approach (commit 0ebf291) fully addresses the Critical finding: 🔴 Sliding-window — RESOLVEDInstead of disabling sliding window or just warning, we now use per-layer hybrid cache:
This ensures sliding window constraints are correctly enforced via GQA's native 🟡 Fragile validation — RESOLVEDWhitelist 🟢
|
0e0e65a to
db1c23a
Compare
|
Will wait after the standard paths land |
24279cc to
b95ef3f
Compare
![github-code-quality[bot]](https://avatars.githubusercontent.com/u/9919?s=60&v=4){kind=small}
Enable --static-cache for Gemma4 models by: - Adding StaticCacheState dispatch to Gemma4DecoderLayer.forward() and Gemma4TextAttention.forward(), matching the pattern from DecoderLayer - Creating _make_gemma4_static_cache_inputs() that handles dual head_dim (sliding=256, full=512) and KV-shared layers (no cache entries) - Updating Gemma4TextCausalLMTask to support static_cache=True with pre-allocated fixed-size cache buffers and TensorScatter ops - Adding Gemma4DecoderLayer to _validate_static_cache_support() and scoping validation to skip encoder sub-modules - Updating CLI to resolve the correct task class for Gemma4 when --static-cache is used (text model class for multimodal models) In static cache mode, attention_mask is None and the Attention op uses is_causal=1 + nonpad_kv_seqlen for masking. GQA is automatically disabled since it requires attention_mask. Tested with google/gemma-4-e2b-it: produces 30 TensorScatter ops (15 non-shared layers x 2), 35 Attention ops, 0 GQA ops. All 1186 existing tests pass. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
- Fix input ordering: write_indices/nonpad_kv_seqlen now come after all cache inputs, matching the standard _make_static_cache_inputs pattern - Add 5 Gemma4 static cache tests: graph build, dual head_dim shapes, TensorScatter op counts, KV-shared layer exclusion, input ordering - Refactor CLI: remove ~50 lines of duplicated build() logic, use build(module_class=...) to override model class for multimodal models instead of reimplementing the build pipeline - No unrelated test deletions found (Finding 4 N/A) Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
Add static_cache/max_seq_len params to Gemma4Task so multimodal models (gemma4 model_type) get TensorScatter on their decoder while vision, audio, and embedding models stay unchanged. This removes the need for the text-only model class override in the CLI — build() now works directly with Gemma4Model + Gemma4Task for both dynamic and static cache modes. CLI _resolve_static_cache_task now maps: - gemma4 → Gemma4Task(static_cache=True) - gemma4_text → Gemma4TextCausalLMTask(static_cache=True) - others → CausalLMTask(static_cache=True) Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
…atic detection Address review findings from @titaiwangms: 1. [Critical] Sliding window warning: _validate_static_cache_support() now warns when the model uses sliding_window, since the static cache path (is_causal=1) does not enforce window constraints. 2. [Major] Replace fragile encoder-prefix blacklist with whitelist: validation now checks isinstance against _supported tuple, which naturally skips vision/audio encoder layers (different classes). 3. [Minor] Detect static_cache_mode from past_key_values content (isinstance StaticCacheState) instead of overloading attention_mask=None. Keeps the type contract of forward() clean. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
Replace the sliding-window UserWarning with a hard ValueError. The ONNX Attention op does not have a local_window_size attribute (only com.microsoft.GroupQueryAttention supports it), so static cache cannot enforce window constraints. This would silently produce incorrect outputs for sequences longer than the window. Models with sliding_window > 0 must use dynamic cache (without --static-cache) for correct behavior. Tests updated: - Default Gemma4 static cache test config uses full_attention only - New test: verify ValueError raised for sliding_window > 0 Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
Replace the per-model static cache approach with per-layer dispatch: - Full-attention layers → StaticCacheState (TensorScatter + Attention) - Sliding-attention layers → dynamic (past_key, past_value) tuples with GQA and local_window_size for correct window enforcement - KV-shared layers → None (borrow from source layers) This is required because the ONNX Attention op does not support local_window_size — only com.microsoft.GroupQueryAttention does. Sliding layers must use GQA to enforce window constraints correctly. Changes: - _make_gemma4_static_cache_inputs: creates StaticCacheState for full-attention layers and dynamic tuples for sliding layers - _register_hybrid_cache_outputs: uses updated_key_cache.N for static layers and present.N.key for dynamic layers - Gemma4TextModel.forward: per-layer dispatch based on isinstance(past_kv, StaticCacheState) - Removed sliding-window rejection from _validate_static_cache_support - Both task classes (text-only + multimodal) provide attention_mask for sliding GQA layers in hybrid mode Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
Signed-off-by: Justin Chu <justinchu@microsoft.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
Signed-off-by: Justin Chu <justinchu@microsoft.com>
Signed-off-by: Justin Chu <justinchu@microsoft.com>
Three fixes for static cache export after main merge: 1. Define past_seq_len for static cache path (sliding layers need it for dynamic cache within hybrid static/dynamic scheme) 2. Set attention_mask=None for static cache (uses is_causal + nonpad_kv_seqlen instead of explicit mask) 3. Guard create_attention_bias: skip when attention_mask is None (prevents Shape node with None input that crashes constant folding) 4. Fix test: use _config instead of unbound config variable Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
Handle 3D static cache tensors in KV-shared layer attention: - Static cache sources produce 3D [B, max_seq, kv_hidden] (already in BSNH-like layout), skip Transpose - Dynamic cache sources produce 4D [B, kv_heads, seq, head_dim], need BNSH→BSNH transpose + flatten - Pass nonpad_kv_seqlen from static source to KV-shared Attention for correct Flash Attention dispatch - Restore o_proj call before return - Fix else: keyword placement after GQA shared_kv block Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
…KV-shared - Static cache Attention: use is_causal=0 (nonpad_kv_seqlen enforces bounds; is_causal=1 causes wrong upper-left alignment for decode) - KV-shared layers: always use dynamic _apply_attention path even when borrowing from static cache source (nonpad + is_causal=0 triggers ORT CUDA kernel issue for decode) - Non-shared static layers (4, 9, 14): prefill works, decode needs further investigation for KV-shared layers Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
Correct Gemma4 static-cache RoPE setup and stale causality coverage, then carry forward the ONNXScript 0.7.1 rewrite/test API migrations and current synthetic-parity compatibility fixes required by the unpinned CI environment. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
The latest main supports maskless external-cache Attention with built-in causality while retaining non-causal mode for additive-bias callers. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
d8fbf6d to
5f24f98
Compare
2 participants
Summary
Adds
--static-cachesupport for Gemma4 models, enabling pre-allocated KV cache buffers with TensorScatter ops (opset 24).Changes
src/mobius/models/gemma4.pyStaticCacheStatedispatch toGemma4DecoderLayer.forward()— extracts static cache frompast_key_valueand passes to attentionstatic_cacheparameter toGemma4TextAttention.forward()— threads through to_apply_attention()static_cache_modeflag inGemma4TextModel.forward()— skips GQA, fallback mask construction, and usesNoneattention bias whenattention_maskisNonesrc/mobius/tasks/_gemma4.py_make_gemma4_static_cache_inputs()— creates per-layer static cache with correct dual head_dim (sliding=256, full=512) and skips KV-shared layersGemma4TextCausalLMTaskwithstatic_cache/max_seq_lenparams and static vs dynamic cache branchingsrc/mobius/tasks/_causal_lm.pyGemma4DecoderLayerto_validate_static_cache_support()src/mobius/__main__.py--static-cacheCLI to resolve the correct task class for Gemma4 (uses text model type for multimodal models)Testing
mobius build --model google/gemma-4-e2b-it --dtype f16 --static-cache --max-seq-len 4096produces correct model with 30 TensorScatter + 35 Attention + 0 GQA opsPhase 1 Scope / Limitations
This is Phase 1 — full-attention and sliding-attention layers both work, but:
is_causal=1without sliding window constraint (standard causal attention). This matches the generic DecoderLayer static cache behavior.