Merge branch 'main' into justinchu/fix-castlike-dtype

Author: justinchuby

.github/copilot-instructions.md

@@ -101,6 +101,13 @@ and wrapper modules for nesting. See the `weight-name-alignment` skill.
 - Use ONNX opset 23 `op.Attention` with `q_num_heads`/`kv_num_heads`
   attributes (not `num_heads`)
 
+### Protobuf prohibition
+
+- **Zero explicit protobuf operations in this repo**
+- Never use `onnx.helper`, `onnx.TensorProto`, `onnx.ModelProto`, or any protobuf construction APIs
+- Always use `onnx_ir` APIs (`ir.Graph`, `ir.Node`, `ir.Function`, `ir.Value`, `ir.Tensor`) — `onnxscript.ir` is a deprecated alias, do not use it
+- This applies to all code: models, components, tasks, tests, utilities, and function bodies
+
 ### Comments and documentation
 
 - **Add inline comments that explain the model architecture**: annotate

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

@@ -46,8 +46,8 @@ If the logits match HuggingFace, you only need the registry entry.
 Create `src/mobius/models/<model_name>.py`.  The minimal template:
 
 ```python
-# Copyright (c) ONNX Project Contributors
-# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
 
 from __future__ import annotations
 
@@ -304,16 +304,28 @@ new model is a significant addition.
 
 ## Checklist
 
-- [ ] Model file in `src/mobius/models/` with Apache-2.0 copyright header
+This is the **implementation** checklist — the steps needed to wire up a new
+model in the codebase.  For the full **definition-of-done** quality checklist
+(L1–L5 tests, ORT GenAI, Foundry Local, Olive quantization, multi-dtype,
+code review), see the
+[quality-checklist skill](../quality-checklist/SKILL.md).
+
+- [ ] Model file in `src/mobius/models/` with Microsoft MIT copyright header
 - [ ] Class has `default_task` and `category` attributes (if not standard text-generation)
 - [ ] Class has a descriptive docstring (first paragraph used in generated docs)
 - [ ] `preprocess_weights` handles any key mismatches
 - [ ] Registered in `_create_default_registry()`
 - [ ] Exported from `models/__init__.py`
 - [ ] Config extraction works (`ArchitectureConfig.from_transformers`)
 - [ ] Tiny config in `tests/_test_configs.py` (with `is_representative` flag)
-- [ ] Integration test model in `tests/integration_test.py` (if small checkpoint available)
+- [ ] L2 YAML test case in `testdata/cases/` with `test_model_id`
+- [ ] L3 synthetic parity passes (`tests/synthetic_parity_test.py -k "<model_type>"`)
+- [ ] Integration test model in `tests/integration_test.py` (real-weight integration suite, if small checkpoint available)
+- [ ] L4 golden file generated and committed (`testdata/golden/`)
+- [ ] L5 generation golden file generated and committed
+- [ ] ORT GenAI test added to `tests/ort_genai_test.py` (text-generation and VLM models)
 - [ ] CLI build works (`mobius build --model ...`)
+- [ ] Multi-dtype correctness verified (fp32, fp16, bf16)
 
 **Note:** Default optimizer passes (CSE, deduplicate initializers, identity
 elimination, remove unused nodes/opsets) are applied automatically by
@@ -819,3 +831,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
+```