Add static KV cache (TensorScatter) support for Gemma4

src/mobius/__main__.py

@@ -119,6 +119,24 @@ def _cmd_build(args: argparse.Namespace) -> None:
     )
     from mobius.tasks import CausalLMTask, ModelTask
 
+    def _resolve_static_cache_task(model_type: str) -> ModelTask:
+        """Create the correct static cache task for the given model type."""
+        if model_type == "gemma4":
+            from mobius.tasks._gemma4 import Gemma4Task
+
+            return Gemma4Task(
+                static_cache=True,
+                max_seq_len=args.max_seq_len,
+            )
+        if model_type == "gemma4_text":
+            from mobius.tasks._gemma4 import Gemma4TextCausalLMTask
+
+            return Gemma4TextCausalLMTask(
+                static_cache=True,
+                max_seq_len=args.max_seq_len,
+            )
+        return CausalLMTask(static_cache=True, max_seq_len=args.max_seq_len)
+
     # Validate --max-seq-len requires --static-cache
     if args.max_seq_len is not None and not args.static_cache:
         raise SystemExit("Error: --max-seq-len can only be used with --static-cache.")
@@ -161,7 +179,14 @@ def _cmd_build(args: argparse.Namespace) -> None:
     load_weights = not args.no_weights
     task: str | ModelTask | None = args.task
     if args.static_cache:
-        task = CausalLMTask(static_cache=True, max_seq_len=args.max_seq_len)
+        # Defer task creation — we need to know the model type first.
+        # Store parameters for later resolution.
+        static_cache_params = {
+            "static_cache": True,
+            "max_seq_len": args.max_seq_len,
+        }
+    else:
+        static_cache_params = None
     trust_remote_code = args.trust_remote_code
     output_dir = args.output_dir
     os.makedirs(output_dir, exist_ok=True)
@@ -218,7 +243,9 @@ def _cmd_build(args: argparse.Namespace) -> None:
         config = _config_from_hf(hf_config, parent_config=parent_config)
         if dtype_override is not None:
             config = dataclasses.replace(config, dtype=dtype_override)
-        if task is None:
+        if static_cache_params is not None:
+            task = _resolve_static_cache_task(model_type)
+        elif task is None:
             task = _default_task_for_model(model_type)
         module_class = registry.get(model_type)
         model_module = module_class(config)
@@ -233,8 +260,18 @@ def _cmd_build(args: argparse.Namespace) -> None:
                 state_dict = model_module.preprocess_weights(state_dict)
             pkg.apply_weights(state_dict)
     else:
+        model_id_or_path = args.model
+        if static_cache_params is not None:
+            # Detect model type to resolve the correct static cache task.
+            import transformers
+
+            hf_config = transformers.AutoConfig.from_pretrained(
+                model_id_or_path, trust_remote_code=trust_remote_code
+            )
+            task = _resolve_static_cache_task(getattr(hf_config, "model_type", ""))
+
         pkg = build(
-            args.model,
+            model_id_or_path,
             task=task,
             dtype=dtype_override,
             load_weights=load_weights,

src/mobius/models/gemma4.py

@@ -54,7 +54,7 @@
 from mobius.models.gemma3_text import Gemma3TextScaledWordEmbedding
 
 if TYPE_CHECKING:
-    from mobius.components._attention import GQAContext
+    from mobius.components._attention import GQAContext, StaticCacheState
 
 
 # ---------------------------------------------------------------------------
@@ -874,6 +874,7 @@ def forward(
         shared_kv_states: dict | None = None,
         past_key_value: tuple | None = None,
         is_causal: int = 1,
+        static_cache: StaticCacheState | None = None,
     ):
         from mobius.components._attention import (
             GQAContext,
@@ -902,7 +903,7 @@ def forward(
 
         if self.is_kv_shared_layer:
             # KV-shared layers borrow K,V from a source layer (no own KV cache).
-            src_key, src_value = shared_kv_states[self.kv_shared_layer_index]
+            src_key, src_value = shared_kv_states[self.kv_shared_layer_index][:2]
 
             if use_gqa:
                 # GQA path for shared KV: pass empty K/V tensors and wire the
@@ -951,43 +952,27 @@ def forward(
                     **gqa_attrs,
                 )
             else:
-                # Fallback Attention path (non-GQA / CPU-style graphs).
-                #
-                # The shared K,V buffer is the source layer's present K/V, 4D
-                # BNSH [batch, kv_heads, kv_len, head_dim], and already contains
-                # the FULL sequence (with RoPE applied).  Transpose it to the 3D
-                # layout the Attention op expects:
-                #   [batch, kv_len, kv_heads * head_dim].
-                #
-                # Reshape to the STATIC ``kv_heads * head_dim`` hidden size (not a
-                # dynamic ``-1``): the source present-value comes out of an
-                # Attention op whose head_dim onnxruntime does not always
-                # propagate, so a ``-1`` here leaves the Attention op's value
-                # head size (and therefore this layer's attention output width)
-                # unknown.  onnxruntime then infers a mismatched o_proj input dim
-                # and rejects the graph at session creation with a MatMul
-                # "Incompatible dimensions" shape-inference error.  A concrete
-                # last dim lets it infer the attention output width correctly.
-                shared_kv_hidden = self.num_key_value_heads * self.head_dim
-                src_key = op.Transpose(src_key, perm=[0, 2, 1, 3])
-                src_key = op.Reshape(src_key, [0, 0, shared_kv_hidden])
-                src_value = op.Transpose(src_value, perm=[0, 2, 1, 3])
-                src_value = op.Reshape(src_value, [0, 0, shared_kv_hidden])
-
-                # IMPORTANT: pass is_causal=0 here, NOT the caller's is_causal.
-                #
-                # We feed the FULL shared sequence as key/value with no past, so
-                # q_len < kv_len during decode.  ``attention_bias`` from
-                # ``create_attention_bias`` already bakes in the complete
-                # bottom-right causal (+ sliding + padding) mask, so causality is
-                # fully handled by the bias.  If we ALSO set is_causal=1 the
-                # Attention op applies its own built-in causal mask on top — and
-                # for q_len < kv_len the two EPs disagree on its alignment (per
-                # the ONNX spec is_causal is UPPER-LEFT aligned: CUDA follows the
-                # spec and a single decode query attends only to kv[0], while the
-                # CPU EP bottom-right aligns).  That double-masking is what made
-                # gemma4 decode diverge on CUDA.  Relying solely on the float
-                # bias (is_causal=0) is correct and identical on CPU and CUDA.
+                # Fallback Attention path: transpose shared KV from BNSH to 3D.
+                # Source K/V shape depends on whether the source layer uses
+                # static or dynamic cache:
+                #   Dynamic: [B, kv_heads, total_seq, head_dim] (4D present)
+                #   Static:  [B, max_seq, kv_heads*head_dim]    (3D updated cache)
+                # Static cache sources are already 3D — skip reshape.
+                is_static_source = src_key.shape is not None and len(src_key.shape) == 3
+                if not is_static_source:
+                    shared_kv_hidden = self.num_key_value_heads * self.head_dim
+                    src_key = op.Transpose(src_key, perm=[0, 2, 1, 3])
+                    src_key = op.Reshape(src_key, [0, 0, shared_kv_hidden])
+                    src_value = op.Transpose(src_value, perm=[0, 2, 1, 3])
+                    src_value = op.Reshape(src_value, [0, 0, shared_kv_hidden])
+
+                # KV-shared layers always use the dynamic Attention path with
+                # mask (attention_bias). Even when the source layer uses static
+                # cache, the KV-shared layer's Attention uses the source's full
+                # cache as K/V with past_key=None (no own KV concat).
+                # The nonpad_kv_seqlen path is NOT used here because ORT's
+                # is_causal=0 + nonpad_kv_seqlen triggers a CUDA kernel issue
+                # for KV-shared decode (S_q=1, S_kv=max_seq).
                 attn_output, present_key, present_value = _apply_attention(
                     op,
                     query_states,
@@ -1071,6 +1056,7 @@ def forward(
                 shared_kv_states[self.layer_idx] = (
                     present_key,
                     present_value,
+                    None,  # no nonpad_kv_seqlen for GQA path
                 )
         else:
             # K projection + per-head K norm + optional RoPE
@@ -1123,14 +1109,19 @@ def forward(
                 num_key_value_heads=self.num_key_value_heads,
                 scale=self.scaling,
                 softcap=self.softcap,
+                static_cache=static_cache,
                 is_causal=is_causal,
             )
 
             # Source layers store K,V for downstream KV-shared layers.
+            # Include nonpad_kv_seqlen for static cache sources so
+            # KV-shared layers can pass it to the Attention op.
             if self.provides_shared_kv and shared_kv_states is not None:
+                nonpad = static_cache.nonpad_kv_seqlen if static_cache else None
                 shared_kv_states[self.layer_idx] = (
                     present_key,
                     present_value,
+                    nonpad,
                 )
 
         attn_output = self.o_proj(op, attn_output)
@@ -1332,9 +1323,19 @@ def forward(
         position_embeddings: tuple | None,
         shared_kv_states: dict,
         per_layer_input: ir.Value | None,
-        past_key_value: tuple | None,
+        past_key_value: tuple | StaticCacheState | None,
         is_causal: int = 1,
     ):
+        # Dispatch StaticCacheState: extract it from past_key_value so that
+        # the attention module receives it as a separate parameter.
+        from mobius.components._attention import StaticCacheState
+
+        if isinstance(past_key_value, StaticCacheState):
+            static_cache = past_key_value
+            past_key_value = None
+        else:
+            static_cache = None
+
         # Attention block: pre-norm -> attn -> post-norm -> residual
         residual = hidden_states
         hidden_states = self.input_layernorm(op, hidden_states)
@@ -1345,6 +1346,7 @@ def forward(
             position_embeddings=position_embeddings,
             shared_kv_states=shared_kv_states,
             past_key_value=past_key_value,
+            static_cache=static_cache,
             is_causal=is_causal,
         )
         hidden_states = self.post_attention_layernorm(op, attn_output)
@@ -1810,7 +1812,7 @@ def forward(
         self,
         op: OpBuilder,
         input_ids: ir.Value | None,
-        attention_mask: ir.Value,
+        attention_mask: ir.Value | None,
         position_ids: ir.Value,
         past_key_values: list | None = None,
         inputs_embeds: ir.Value | None = None,
@@ -1846,11 +1848,10 @@ def forward(
         # KV-shared layers fall back to standard Attention because they
         # borrow K,V from another layer (no own KV cache).
         from mobius._build_context import get_build_dtype
-        from mobius.components._attention import GQAContext
+        from mobius.components._attention import GQAContext, StaticCacheState
 
         caps = ep_capabilities()
         dtype = get_build_dtype()
-
         # Bidirectional vision-block overlay (Gemma4 larger models). When
         # active, contiguous vision-token blocks attend bidirectionally on
         # BOTH full and sliding layers. This cannot be expressed by the
@@ -1886,6 +1887,9 @@ def forward(
             )
         use_block_overlay = bidirectional and block_sequence_ids is not None
 
+        # GQA is available when attention_mask exists and the EP supports it.
+        # In hybrid mode, sliding layers use GQA while full-attention layers
+        # use the static Attention path with TensorScatter.
         use_gqa = (
             attention_mask is not None
             and dtype in caps.gqa_dtypes
@@ -1903,7 +1907,7 @@ def forward(
             # GQA references these caches directly; without the call the
             # parameters are never emitted into the graph.
             _ = self.rotary_emb_local(op, position_ids)
-            _ = self.rotary_emb_global(op, position_ids)
+            global_pos_emb = self.rotary_emb_global(op, position_ids)
 
             # seqlens_k[b] = sum(attention_mask[b]) - 1  (last valid KV idx)
             # total_seq_len = attention_mask.shape[1]     (past + current)
@@ -1949,6 +1953,12 @@ def forward(
                 "sliding_attention": None,
                 "full_attention": None,
             }
+            # In hybrid mode, static-cache full-attention layers need RoPE
+            # embeddings for the standard Attention path (not GQA).
+            if past_key_values is not None and any(
+                isinstance(kv, StaticCacheState) for kv in past_key_values if kv is not None
+            ):
+                position_embeddings_dict["full_attention"] = global_pos_emb
         else:
             position_embeddings_dict = {
                 "sliding_attention": self.rotary_emb_local(op, position_ids),
@@ -1959,7 +1969,12 @@ def forward(
         query_input = input_ids if input_ids is not None else hidden_states
         fallback_bias_dict: dict[str, ir.Value | None] = {}
         need_fallback = not use_gqa
-        if need_fallback:
+        if need_fallback and attention_mask is not None:
+            # All fallback layers use float additive bias masks encoding
+            # causal + sliding window + padding constraints. Float bias
+            # works with both unfused and MEA kernel paths on CUDA EP.
+            # When attention_mask is None (static cache), the Attention op uses
+            # nonpad_kv_seqlen to bound the externally managed cache.
             fallback_bias_dict = {
                 "sliding_attention": create_attention_bias(
                     op,
@@ -2003,14 +2018,24 @@ def forward(
         ):
             per_layer_input = per_layer_list[i] if per_layer_list is not None else None
 
-            # Per-layer decision: use GQA when available. KV-shared layers
-            # also use GQA (with empty K/V and shared past buffer).
-            if use_gqa:
+            # Per-layer cache/attention dispatch:
+            # - StaticCacheState → static path (TensorScatter + Attention)
+            # - Dynamic tuple → GQA path (with local_window_size)
+            # - None (no cache) → fallback Attention path
+            is_layer_static = isinstance(past_kv, StaticCacheState)
+
+            if is_layer_static:
+                attn_bias = None
+                pos_emb = position_embeddings_dict[layer_type]
+            elif use_gqa:
                 attn_bias = gqa_ctx_dict[layer_type]
                 pos_emb = None
-            else:
+            elif fallback_bias_dict:
                 attn_bias = fallback_bias_dict[layer_type]
                 pos_emb = fallback_pos_dict[layer_type]
+            else:
+                attn_bias = None
+                pos_emb = position_embeddings_dict.get(layer_type)
 
             hidden_states, present_kv = layer(
                 op,

src/mobius/tasks/_causal_lm.py

@@ -400,10 +400,15 @@ def _register_static_cache_outputs(
 def _validate_static_cache_support(module: nn.Module) -> None:
     """Check that the module's decoder layers support StaticCacheState.
 
-    Only :class:`DecoderLayer` and :class:`MoEDecoderLayer` have the
-    ``isinstance(StaticCacheState)`` dispatch in ``forward()``.  Custom
-    decoder layers will silently unpack the NamedTuple as a regular
-    ``(key, value)`` tuple, producing wrong results.
+    Only :class:`DecoderLayer`, :class:`MoEDecoderLayer`, and
+    :class:`Gemma4DecoderLayer` have the ``isinstance(StaticCacheState)``
+    dispatch in ``forward()``.  Custom decoder layers will silently
+    unpack the NamedTuple as a regular ``(key, value)`` tuple, producing
+    wrong results.
+
+    Also warns when the model uses sliding-window attention, since the
+    static cache path does not enforce window constraints (the Attention
+    op uses ``is_causal=1`` without ``local_window_size``).
 
     NOTE: The following models are NOT yet supported in static cache
     mode and will raise TypeError from this check:
@@ -427,24 +432,29 @@ def _validate_static_cache_support(module: nn.Module) -> None:
         TypeError: If any decoder layer is not a supported type.
     """
     from mobius.components._decoder import DecoderLayer
+    from mobius.models.gemma4 import Gemma4DecoderLayer
     from mobius.models.moe import MoEDecoderLayer
 
+    _supported = (DecoderLayer, MoEDecoderLayer, Gemma4DecoderLayer)
+
+    # Whitelist-based validation: only check layers that have self_attn/attn
+    # (decoder-like), and accept those that are in the supported tuple.
+    # This naturally skips vision/audio encoder layers since they use
+    # different classes (e.g. Gemma4VisionEncoderLayer).
     for name, child in module.named_modules():
         if not isinstance(child, nn.ModuleList):
             continue
         for i, layer in enumerate(child):
             if not isinstance(layer, nn.Module):
                 continue
-            # Check modules that look like decoder layers: they have an
-            # attention sub-module named either "self_attn" (standard) or
-            # "attn" (GPT-2 style).
             if not hasattr(layer, "self_attn") and not hasattr(layer, "attn"):
                 continue
-            if not isinstance(layer, (DecoderLayer, MoEDecoderLayer)):
+            if not isinstance(layer, _supported):
                 raise TypeError(
                     f"Static cache mode requires decoder layers that "
-                    f"inherit from DecoderLayer or MoEDecoderLayer, but "
-                    f"{name}[{i}] is {type(layer).__name__}. Either use a "
-                    f"compatible model or add StaticCacheState dispatch to "
+                    f"inherit from DecoderLayer, MoEDecoderLayer, or "
+                    f"Gemma4DecoderLayer, but {name}[{i}] is "
+                    f"{type(layer).__name__}. Either use a compatible "
+                    f"model or add StaticCacheState dispatch to "
                     f"{type(layer).__name__}.forward()."
                 )