docs: update skills with proven Gemma 4 patterns (MoE, Any-to-Any, KV sharing)

.github/skills/adding-a-new-model/SKILL.md

@@ -819,3 +819,150 @@ When adding a new model, use these files as canonical references:
 | **Minimal** — encoder subclass | `models/layoutlmv3.py` | Extends `BertModel`, only overrides `preprocess_weights()`. Same pattern for encoder-only models. |
 | **Moderate** — custom components | `models/gemma.py` | Adds custom attention (soft-capping), custom MLP (GeGLU), and custom normalization. Good example of component subclassing. |
 | **Complex** — multi-model architecture | `models/qwen3_tts.py` | 4-model TTS split with talker, code predictor, embedding, and speaker encoder sub-modules. Shows how to structure multi-model architectures. |
+
+## KV sharing across layers (num_kv_shared_layers)
+
+Some models (e.g. Gemma 4) reduce parameter count by having the last N
+decoder layers **borrow** Key and Value states from an earlier "source" layer
+of the same type instead of projecting their own K,V. This is controlled by
+`num_kv_shared_layers` in the HuggingFace config.
+
+### What it means
+
+```
+first_kv_shared_idx = num_hidden_layers - num_kv_shared_layers
+
+Layers [0 .. first_kv_shared_idx - 1]:  normal — own k_proj, v_proj, k_norm
+Layers [first_kv_shared_idx .. end]:     shared — NO k_proj/v_proj weights
+```
+
+Each shared layer reuses K,V from the **last non-shared layer of the same
+attention type** (e.g. sliding vs. full attention). Only Q is computed fresh.
+
+### Impact on the checkpoint
+
+Shared layers have **no `k_proj`, `v_proj`, `k_norm`** keys in the
+HuggingFace checkpoint. `preprocess_weights` must not assert these keys
+exist for shared-layer indices — they simply won't be present.
+
+```python
+def preprocess_weights(self, state_dict):
+    # shared layers have no k/v proj — remove them silently if accidentally present
+    first_shared = self.config.num_hidden_layers - self.config.num_kv_shared_layers
+    for i in range(first_shared, self.config.num_hidden_layers):
+        for suffix in ("k_proj.weight", "v_proj.weight", "k_norm.weight"):
+            state_dict.pop(f"model.layers.{i}.self_attn.{suffix}", None)
+    return super().preprocess_weights(state_dict)
+```
+
+### Attention module: is_kv_shared_layer flag
+
+The attention class detects at `__init__` time whether it is a shared layer:
+
+```python
+class Gemma4Attention(nn.Module):
+    def __init__(self, config, layer_idx, layer_types, first_kv_shared_idx, ...):
+        self.is_kv_shared_layer = layer_idx >= first_kv_shared_idx > 0
+        prev_layers = layer_types[:first_kv_shared_idx]
+
+        if self.is_kv_shared_layer:
+            # Index of the source layer whose K,V this layer borrows
+            self.kv_shared_layer_index = (
+                len(prev_layers) - 1 - prev_layers[::-1].index(layer_types[layer_idx])
+            )
+            self.store_full_length_kv = False
+        else:
+            self.kv_shared_layer_index = None
+            # True for the last non-shared layer of each type that has downstream
+            # KV-shared layers depending on it — it stores K,V for reuse.
+            self.store_full_length_kv = first_kv_shared_idx > 0 and (
+                layer_idx
+                == len(prev_layers) - 1 - prev_layers[::-1].index(layer_types[layer_idx])
+            )
+
+        # All layers have Q projection
+        self.q_proj = Linear(config.hidden_size, num_heads * head_dim)
+        self.q_norm = RMSNorm(head_dim)
+        self.o_proj = Linear(num_heads * head_dim, config.hidden_size)
+
+        # Only non-shared layers have K/V projections
+        if not self.is_kv_shared_layer:
+            self.k_proj = Linear(config.hidden_size, num_kv_heads * head_dim)
+            self.v_proj = Linear(config.hidden_size, num_kv_heads * head_dim)
+            self.k_norm = RMSNorm(head_dim)
+```
+
+### forward(): shared layers consume shared_kv_states dict
+
+Pass a mutable `shared_kv_states` dict through the forward call. Source
+layers populate it; shared layers read from it:
+
+```python
+def forward(self, op, hidden_states, ..., shared_kv_states, past_key_value):
+    # Q projection (all layers)
+    query_states = self.q_proj(op, hidden_states)
+    ...
+
+    if self.is_kv_shared_layer:
+        # Borrow K,V from source layer (already in shared_kv_states)
+        src_key, src_value = shared_kv_states[self.kv_shared_layer_index]
+        # Reshape from present_kv 4D [B, kv_heads, total_seq, head_dim]
+        # to Attention input 3D [B, total_seq, kv_heads * head_dim]
+        src_key = op.Transpose(src_key, perm=[0, 2, 1, 3])
+        key_states = op.Reshape(src_key, ...)
+        value_states = ...
+    else:
+        # Normal K/V projection + norm
+        key_states = self.k_proj(op, hidden_states)
+        value_states = self.v_proj(op, hidden_states)
+        ...
+
+    hidden_out, present_kv = _apply_attention(op, query_states, key_states, ...)
+
+    if self.store_full_length_kv:
+        # Store present_kv [B, kv_heads, total_seq, head_dim] for downstream shared layers
+        shared_kv_states[self.layer_idx] = (present_kv_key, present_kv_value)
+
+    return hidden_out, present_kv
+```
+
+### Text model: KV cache has only num_kv_layers entries
+
+KV-shared layers do **not** append to `present_key_values`. The output list
+has `num_hidden_layers - num_kv_shared_layers` entries, not `num_hidden_layers`:
+
+```python
+# In Gemma4TextModel.forward():
+shared_kv_states: dict = {}
+present_key_values = []
+
+# past_key_values has only num_kv_layers entries (no entry for KV-shared layers).
+# Expand it to a full per-layer list so we can zip cleanly over all layers.
+if past_key_values is not None:
+    kv_iter = iter(past_key_values)
+    past_kvs: list = [
+        None if layer.self_attn.is_kv_shared_layer else next(kv_iter)
+        for layer in self.layers
+    ]
+else:
+    past_kvs = [None] * len(self.layers)
+
+for i, (layer, layer_type, past_kv) in enumerate(
+    zip(self.layers, self.layer_types, past_kvs)
+):
+    hidden_states, present_kv = layer(
+        op,
+        hidden_states=hidden_states,
+        attention_bias=attention_bias_dict[layer_type],
+        position_embeddings=position_embeddings_dict[layer_type],
+        shared_kv_states=shared_kv_states,
+        past_key_value=past_kv,
+    )
+    # KV-shared layers borrow K,V — exclude from present_key_values so the
+    # output has exactly num_kv_layers (not num_hidden_layers) entries.
+    if not layer.self_attn.is_kv_shared_layer:
+        present_key_values.append(present_kv)
+```
+
+The task's KV cache inputs/outputs must use the correct count:
+`num_kv_layers = config.num_hidden_layers - config.num_kv_shared_layers`.

.github/skills/moe-models/SKILL.md

@@ -260,3 +260,124 @@ should be small enough for CI (~1-4B params).  The test pattern:
 3. Optionally test greedy generation (token-ID matching)
 
 See `tests/moe_integration_test.py` for the complete pattern.
+
+## Direct MoE op emission (com.microsoft.MoE)
+
+OnnxRuntime ships a fused `com.microsoft.MoE` contrib op (CUDA float32/fp16/bf16,
+CPU float32/fp16). For new model architectures, **emit it directly** — like
+`com.microsoft.GroupQueryAttention` — rather than relying on a rewrite rule.
+
+### When to use
+
+Use `com.microsoft.MoE` when:
+- The model uses top-k MoE routing (standard softmax gate → TopK)
+- All expert weights are the same shape (no dynamic expert counts)
+- The EP's `caps.supports_fused_moe` is `True`
+
+Fall back to the loop-over-experts path when `supports_fused_moe` is `False`
+(CPU EP without contrib ops, or EPs that don't support the custom op).
+
+### Gate output: full pre-topk router_probs
+
+The op takes the **full** `(num_tokens, num_experts)` probability tensor and
+performs top-k selection internally via the `k` attribute. The gate must
+produce the full softmax distribution — not already-selected top-k indices.
+
+```python
+# In your gate forward(), return shape [num_tokens, num_experts]
+router_probs = op.Softmax(op.MatMul(hidden_states, self.weight), axis=-1)
+```
+
+### Emission pattern (from Gemma 4 implementation)
+
+```python
+from mobius._build_context import ep_capabilities
+
+caps = ep_capabilities()
+if caps.supports_fused_moe:
+    moe_out = op.CastLike(
+        op.MoE(                        # type: ignore[attr-defined]
+            normed_hidden,             # [num_tokens, hidden_size]
+            router_probs,              # [num_tokens, num_experts] — full pre-topk
+            self.fc1_experts_weights,  # [E, inter_size, hidden_size]
+            self.fc2_experts_weights,  # [E, hidden_size, inter_size]
+            activation_type="silu",
+            k=self._top_k,
+            normalize_routing_weights=1,
+            _domain="com.microsoft",
+        ),
+        normed_hidden,  # CastLike: preserve bf16/fp16/fp32 — NOT hardcoded float32
+    )
+else:
+    moe_out = self._dispatch_moe_fallback(op, normed_hidden, router_probs)
+```
+
+**Critical: use `CastLike` after the MoE op.** The `com.microsoft.MoE` custom
+op has `type=None` on its output — ONNX type propagation cannot infer the
+output dtype. `op.CastLike(moe_output, target=input)` restores the correct
+dtype (bf16/fp16/fp32), which allows downstream ops to share scalar
+initializers and avoids hard-coded `Cast` to float32.
+
+### preprocess_weights: expert weight mapping
+
+HuggingFace Gemma 4 stores experts as a 3D tensor per projection:
+`layers.N.experts.gate_up_proj [E, 2*inter, H]`
+`layers.N.experts.down_proj    [E, H, inter]`
+
+Map these to the parameter names used by the ONNX MoE op:
+
+```python
+def preprocess_weights(self, state_dict):
+    for key in list(state_dict.keys()):
+        if ".experts.gate_up_proj" in key:
+            new_key = key.replace(".experts.gate_up_proj", ".fc1_experts_weights")
+            state_dict[new_key] = state_dict.pop(key)
+        elif ".experts.down_proj" in key:
+            new_key = key.replace(".experts.down_proj", ".fc2_experts_weights")
+            state_dict[new_key] = state_dict.pop(key)
+    return super().preprocess_weights(state_dict)
+```
+
+For models that store per-expert weights separately (one matrix per expert),
+stack them into 3D tensors in `preprocess_weights`:
+
+```python
+n = config.num_local_experts
+gate = torch.stack([state_dict.pop(f"experts.{i}.gate_proj.weight") for i in range(n)])
+down = torch.stack([state_dict.pop(f"experts.{i}.down_proj.weight") for i in range(n)])
+state_dict["fc1_experts_weights"] = gate   # [E, inter, hidden]
+state_dict["fc2_experts_weights"] = down   # [E, hidden, inter]
+```
+
+### EP capability check (matches GQA pattern)
+
+```python
+# In _execution_providers.py EpCapabilities:
+supports_fused_moe: bool = True  # set False for EPs without com.microsoft.MoE
+
+# In model forward():
+from mobius._execution_providers import ep_capabilities
+caps = ep_capabilities()
+if caps.supports_fused_moe:
+    # emit com.microsoft.MoE
+else:
+    # fallback loop
+```
+
+### Fallback: TopKGate + loop dispatch
+
+When `supports_fused_moe` is False, implement a static unroll:
+
+```python
+def _dispatch_moe_fallback(self, op, hidden, router_probs):
+    k_tensor = op.Constant(value_ints=[self._top_k])
+    top_weights, top_indices = op.TopK(router_probs, k_tensor, axis=-1)
+    top_weights = op.Softmax(top_weights, axis=-1)   # renormalize
+    output = op.CastLike(op.ConstantOfShape(op.Shape(hidden), value=0.0), hidden)
+    for e_idx in range(self._num_experts):
+        w1 = op.Squeeze(op.Gather(self.fc1_experts_weights, [e_idx], axis=0), [0])
+        w2 = op.Squeeze(op.Gather(self.fc2_experts_weights, [e_idx], axis=0), [0])
+        # expert output, gated by routing weight
+        ...
+    return output
+```

.github/skills/multimodal-models/SKILL.md

@@ -533,3 +533,87 @@ gathered = op.Gather(padded, indices, axis=0)
 # Where mask selects only real features; padding row is never used
 result = op.Where(image_mask, gathered, text_embeddings)
 ```
+
+## Any-to-Any 4-model task (vision+audio+text)
+
+Some models come in two tiers: small variants support vision **and** audio
+(Any-to-Any), while large variants support vision only (Image-Text-to-Text).
+Each tier uses a different number of exported ONNX models.
+
+### Tier split
+
+| Tier | Models | ONNX split | Example |
+|------|--------|-----------|---------|
+| Small Any-to-Any | E2B, E4B | 4 models: decoder + vision + **audio** + embedding | `google/gemma-4-E2B-it` |
+| Large Image-Text-to-Text | 26B-A4B, 31B | 3 models: decoder + vision + embedding | `google/gemma-4-26B-A4B-it` |
+
+The task class detects which tier to use from the config (e.g. whether
+`config.audio is not None` — `ArchitectureConfig.from_transformers` populates
+the `audio` field when the HuggingFace config contains an audio sub-config).
+
+### 4-model task structure
+
+```
+decoder          inputs_embeds [B, S, H] → logits + KV cache
+vision_encoder   pixel_values [B, 3, H, W] → image_features [num_image_tokens, H]
+audio_encoder    input_features [B, T, mel] → audio_features [num_audio_tokens, H]
+embedding        input_ids + image_features + audio_features → inputs_embeds [B, S, H]
+```
+
+Reference implementation: `Gemma4AnyToAnyTask` in
+`src/mobius/tasks/_gemma4.py`. This follows the same 4-model structural pattern as
+`Phi4MMMultiModalTask` in `src/mobius/tasks/_phi4mm_multimodal.py`
+(each modality is a separate ONNX model; embedding splices features at placeholder
+positions), though the exact I/O shapes differ per architecture.
+
+### Audio encoder wiring
+
+The audio encoder takes raw mel-spectrogram frames and outputs token-level
+features at the text hidden size:
+
+```python
+# In Gemma4AnyToAnyTask._build_audio():
+input_features = ir.Value(
+    name="input_features",
+    shape=ir.Shape([batch, time, input_size]),   # [B, T, 128]
+    type=ir.TensorType(config.dtype),
+)
+audio_features = audio_encoder(op, input_features)
+# audio_features: [B, T//4, text_hidden_size]
+```
+
+The audio encoder (`Gemma4AudioEncoder` / `_Gemma4AudioEncoderModel`) is
+its own `nn.Module` subgraph exported as the `"audio"` key in the
+`ModelPackage`.
+
+### Embedding model fuses all modalities
+
+The embedding model receives `input_ids`, `image_features`, and
+`audio_features` as separate inputs and splices them into the token
+embedding sequence at the placeholder positions:
+
+```python
+# In Gemma4AnyToAnyTask._build_embedding():
+inputs_embeds = embedding(
+    op,
+    input_ids=input_ids,            # [B, S]
+    image_features=image_features,  # [num_image_tokens, H]
+    audio_features=audio_features,  # [num_audio_tokens, H]
+)
+# returns inputs_embeds: [B, S, H]
+```
+
+### Task class tier detection
+
+```python
+class MyAnyToAnyTask(ModelTask):
+    def build(self, module, config):
+        models = {}
+        models["decoder"] = self._build_decoder(module.decoder, config)
+        models["vision"] = self._build_vision(module.vision_encoder, config)
+        models["embedding"] = self._build_embedding(module.embedding, config)
+        # Build audio encoder only when audio config is present
+        if config.audio is not None:
+            models["audio"] = self._build_audio(module.audio_encoder, config)
+        return ModelPackage(models, config=config)
+```