refactor: move prune_lm_head from global flag to CausalLMTask option

Remove the MOBIUS_PRUNE_LM_HEAD feature flag and replace it with a
prune_lm_head constructor parameter on CausalLMTask (and
HybridCausalLMTask), mirroring the stable-API pattern already used by
the static_cache option.

When prune_lm_head=True, the task inserts Gather(axis=1, index=-1)
followed by Unsqueeze(axis=1) on the logits output after calling
module(), changing the output shape from [B, S, vocab] to [B, 1, vocab].
ONNX Runtime's graph optimizer can push this Gather backward through the
LM head MatMul to avoid computing all-token logits during prefill.

Update tests to use CausalLMTask(prune_lm_head=True) instead of the
override_flags context manager. Also merges main branch changes
(tencent_q1_0_use_native_2bit and static_cache_bias flags).

Signed-off-by: Copilot <copilot@github.com>
Signed-off-by: GitHub <noreply@github.com>

Author: Copilot

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

@@ -112,6 +112,12 @@ class MyCausalLMModel(CausalLMModel):
         self.lm_head = Linear(config.hidden_size, config.vocab_size, bias=False)
 ```
 
+> If your text model **subclasses `TextModel`** (rather than `nn.Module`) and
+> overrides `__init__` with `nn.Module.__init__(self)` to swap in a custom
+> decoder layer, you must also set `self.config = config` in that subclass —
+> `TextModel.forward` reads `self.config`. See troubleshooting §7.
+
+
 #### Class metadata attributes
 
 Every registered model class should set two class-level attributes:
@@ -190,9 +196,28 @@ Also export from `src/mobius/models/__init__.py`.
 
 ### 6. Update `ArchitectureConfig.from_transformers` if needed
 
-If the model has unusual config fields, update `from_transformers()` in
-`_configs.py`.  Use safe defaults (1.0 for multipliers, None for optional
-features) so existing models are unaffected.
+If the model has unusual top-level config fields (vocab size, head counts,
+RoPE knobs, etc.), update `from_transformers()` in `src/mobius/_configs/_base.py`.
+Use safe defaults (1.0 for multipliers, None for optional features) so
+existing models are unaffected.
+
+**For audio- or vision-capable models** (i.e. models whose HF config has an
+`audio_config` or `vision_config` sub-object), prefer adding a per-model hook
+under `src/mobius/_configs/per_model/` rather than editing the central file:
+
+```python
+# src/mobius/_configs/per_model/_my_model_vision.py
+from mobius._configs._extractors import register_vision_hook
+
+@register_vision_hook("my_model_type")
+def _my_model_vision(config, parent_config, model_type, fields):
+    fields.update(hidden_size=..., num_attention_heads=..., ...)
+    return None  # contribute fields, defer VisionConfig instantiation
+```
+
+Then add the new module to `src/mobius/_configs/per_model/__init__.py` so its
+side-effect registration runs at import time. The dispatcher filters hooks by
+the declared model_type strings, so unrelated models never see your hook.
 
 ### 7. Write tests
 
@@ -343,9 +368,10 @@ config fields (`embedding_multiplier`, `attention_multiplier`,
 
 ### 4. Config fields not extracted
 
-**Symptom:** Model builds but multipliers default to 1.0.  Add extraction
-to `ArchitectureConfig.from_transformers()` in `_configs.py` with safe
-defaults.
+**Symptom:** Model builds but multipliers default to 1.0. Add extraction
+to `ArchitectureConfig.from_transformers()` in `src/mobius/_configs/_base.py`
+with safe defaults. For audio/vision-specific fields, register a per-model
+hook under `src/mobius/_configs/per_model/` instead — see step 6 above.
 
 ### 5. Debugging workflow for logit mismatches
 
@@ -376,6 +402,25 @@ differs between loads, suspect `_init_weights` corruption. Set
 deterministic, `_init_weights` is the culprit. Then compare specific
 parameters between the loaded model and the safetensors checkpoint.
 
+### 7. `AttributeError: '<X>TextModel' object has no attribute 'config'`
+
+**Symptom:** Building (or running build-graph / GQA rewrite-rule tests for)
+a model raises `AttributeError: '…TextModel' object has no attribute
+'config'` from inside `TextModel.forward` (e.g. `_gqa_local_window_size`).
+
+**Root cause:** The base `TextModel.__init__` sets `self.config = config`,
+and `TextModel.forward` relies on it (sliding-window detection, etc.). A
+`TextModel` **subclass** that overrides `__init__` with
+`nn.Module.__init__(self)` — instead of `super().__init__(config)` — to swap
+in a custom decoder layer must re-establish the contract by setting
+`self.config = config` itself. Forgetting it crashes only that model.
+
+**Fix:** Add `self.config = config` right after `nn.Module.__init__(self)`
+in the subclass `__init__`. Real examples: `Glm4TextModel` (`models/glm.py`),
+`_LoRATextModel` (`models/phi.py`). Do **not** paper over it with
+`getattr(self, "config", None)` in the base — that silently disables
+config-driven features for the offending subclass.
+
 > For additional troubleshooting (gated attention split ordering, DeltaNet
 > scaling, identity node folding, fp32 upcast patterns, multi-token prefill,
 > embedding table off-by-one), read

.agents/skills/attention-optimization/SKILL.md

@@ -23,8 +23,8 @@ Use this skill when:
 | Scenario | Recommended | Why |
 |----------|------------|-----|
 | Causal only | `attn_mask=None` + `is_causal=1` | Enables Flash (fastest for prefill) |
-| Padding (batch>1) | `nonpad_kv_seqlens` (best) or bool mask | `nonpad_kv_seqlens` enables Flash + shared buffer with no mask |
-| Sliding window (simple) | Bool mask | Equally precise as float, uses less memory |
+| Padding (batch>1) | `nonpad_kv_seqlens` (+ static cache) or bool mask | `nonpad_kv_seqlens` enables Flash with no mask; pair with `TensorScatter` static cache for decode (can't combine with `past_key`/`past_value` inputs) |
+| Sliding window (simple) | GQA `local_window_size` or bool mask | `local_window_size` keeps the fast GQA path; bool mask if you need ONNX Attention |
 | Complex (sliding+KV-shared+dual head_dim) | Float additive bias | Avoids mask construction bugs in multi-constraint patterns |
 | Custom pattern | Float additive bias | Arbitrary values |
 
@@ -35,7 +35,7 @@ Use this skill when:
 | Pattern | Recommended mask type |
 |---------|----------------------|
 | Simple causal-only | No mask — use `is_causal=1` (enables Flash) |
-| Sliding window (simple model) | Bool mask (precise, less memory) |
+| Sliding window (simple, GQA path) | GQA `local_window_size` (fast path); bool mask only if you must use ONNX Attention |
 | KV-shared layers | Float additive bias |
 | Mixed head_dim (e.g. Gemma4) | Float additive bias |
 | Padding + causal | `nonpad_kv_seqlens` or bool mask |
@@ -116,26 +116,67 @@ single-token decode (memory-bandwidth bound regardless of kernel):
 | Symmetric heads | `head_size == v_head_size` |
 | GPU | SM≥8.0 (Ampere or newer) |
 
-### `nonpad_kv_seqlens` — the best padding solution
+### `nonpad_kv_seqlens` — variable-length without an explicit mask
 
 ONNX Attention opset 24 adds `nonpad_kv_seqlens` input, which tells
 the kernel the actual (non-padded) KV sequence length per batch item.
 This enables Flash Attention with variable-length sequences **without
 providing an explicit mask** — the kernel applies causal masking
 internally using the sequence length info.
 
+> ⚠️ **Cannot be combined with the `past_key` / `past_value` inputs.** ORT
+> *rejects* `nonpad_kv_seqlens` when `past_key`/`past_value` are supplied:
+> *"nonpad_kv_seqlens should not be used together with past_key and
+> past_value inputs."* It is therefore **not** usable with the *growing*
+> (dynamic) cache mode beyond the prefill pass. It **is** used at every
+> decode step in the **static-cache** mode below, where the full cache is
+> passed in the `key`/`value` slots and `past_key`/`past_value` are unused.
+
+#### In-place KV for ONNX Attention via `TensorScatter` (static cache)
+
+The opset-24 ONNX `Attention` schema has **no `past_present_share_buffer`
+attribute**, so the naive dynamic mode does `present = concat(past, new)`
+every step (an O(N) copy of distinct `past`/`present` tensors). But ONNX
+`Attention` **can** still update a KV cache *in place* — by pairing it with
+the opset-24 **`TensorScatter`** op:
+
+1. Pre-allocate a fixed-size KV buffer (`StaticCacheState` in mobius).
+2. `TensorScatter` writes the new token(s) into the buffer at
+   `write_indices` (in place when the buffer is IO-bound to the same
+   device memory) — no growing concat.
+3. Pass the **full** scattered cache in the `key`/`value` slots (not
+   `past_key`/`past_value`), with `nonpad_kv_seqlens` giving the valid
+   length and `is_causal=1` for masking.
+
+So in-place KV is **not** GQA-exclusive. The contrib **GroupQueryAttention**
+op has a built-in shared buffer (`past_present_share_buffer`); ONNX
+`Attention` reaches the same effect explicitly with `TensorScatter` +
+static cache. GQA's residual decode edge comes mostly from its dedicated
+`seq==1` decode kernels (XQA / Flash-decode), not from buffer management
+alone.
+
 ```python
-# Enables Flash + past_present_share_buffer for efficient KV cache
+# Dynamic (growing) mode — nonpad_kv_seqlens is PREFILL ONLY here,
+# because it cannot be combined with past_key/past_value:
 attn_out = op.Attention(
     query, key, value,
-    attn_mask=None,           # nullptr → Flash eligible
-    past_key=past_k,
-    past_value=past_v,
-    nonpad_kv_seqlens=seqlens_k,  # opset 24
+    attn_mask=None,                # nullptr → Flash eligible
+    nonpad_kv_seqlens=seqlens_k,   # opset 24, prefill pass only
+    q_num_heads=num_heads,
+    kv_num_heads=kv_heads,
+    is_causal=1,
+)
+
+# Static-cache mode — in-place KV at EVERY step (prefill + decode):
+updated_k = op.TensorScatter(key_cache, key, write_indices, axis=1)
+updated_v = op.TensorScatter(value_cache, value, write_indices, axis=1)
+attn_out = op.Attention(
+    query, updated_k, updated_v,   # full cache in key/value slots
+    None, None, None,              # no attn_mask, no past_key/past_value
+    nonpad_kv_seqlens,             # valid length per batch — used every step
     q_num_heads=num_heads,
     kv_num_heads=kv_heads,
     is_causal=1,
-    past_present_share_buffer=1,
 )
 ```
 
@@ -214,24 +255,61 @@ borrow K/V from a layer with different `head_size`, creating
 | Attention bias | ❌ Rejected | ✅ Supported |
 | Flash Attention | ✅ (no mask) | ✅ (no mask) |
 | XQA kernel | ✅ | ❌ |
-| KV cache management | Built-in (`past_present_share_buffer`) | Manual (separate past/present) |
-| Variable-length | Via `seqlens_k` | Via `nonpad_kv_seqlens` (opset 24) |
+| In-place KV buffer | ✅ built-in `past_present_share_buffer` | ✅ via `TensorScatter` + static cache (no growing concat) |
+| Sliding window | ✅ `local_window_size` attribute | Via float/bool bias only |
+| Variable-length | Via `seqlens_k` (works with past KV) | Via `nonpad_kv_seqlens` (with static cache, not the `past_key`/`past_value` inputs) |
 
 **Guideline:** Use Contrib GQA when you don't need attention bias
-(simple causal models). Use ONNX Attention when you need float bias
-(sliding window, KV-shared, padding).
+(simple causal models, sliding-window via `local_window_size`). Use ONNX
+Attention when you need a float bias (KV-shared, dual head_dim, or
+mixed/alternating per-layer windows that one global window can't express).
+
+### GQA sliding window via `local_window_size`
+
+GroupQueryAttention takes a `local_window_size` attribute that masks each
+query to the most recent `W` keys (positions `[i-W+1, i]`) — exactly
+matching HuggingFace `sliding_window=W`. This keeps a uniform-window model
+on the fast GQA path instead of forcing it onto ONNX `Attention` with a
+baked float window mask. In mobius this is wired in `TextModel.forward`
+from `config.sliding_window` (see `GQAContext.local_window_size`), guarded
+to uniformly-sliding models — mixed `layer_types` (Gemma2/3/4, gpt-oss)
+use custom per-layer `GQAContext`s instead.
+
+> Note: `local_window_size` only *masks* attention; it does not shrink the
+> physical KV buffer, so bounding memory still needs a circular/static
+> cache. Also, the post-hoc GQA rewrite (`RotaryAttentionToGQA`) cannot
+> recover a window from an already-baked float mask, so sliding windows
+> must be set on the **direct** GQA path (GQAContext), not via the rewrite.
+>
+> Decode trade-off: a non-default `local_window_size != -1` disqualifies the
+> dedicated `seq==1` **XQA** decode kernel (which requires `local_window==-1`,
+> see the CUDA cascade table above), so windowed GQA decode falls back to
+> Flash-decode/MEA. This is still on the fast GQA path and is the correct
+> trade for models that genuinely need the window; do not set it on
+> full-attention models.
 
 ## Key takeaways for model builders
 
 1. **Flash requires `attn_mask == nullptr`** — any explicit mask
    disables Flash. Use `is_causal=1` instead.
 2. **`nonpad_kv_seqlens`** enables Flash with variable-length sequences
-   without an explicit mask.
+   without an explicit mask. It cannot be combined with the
+   `past_key`/`past_value` inputs (so it is prefill-only in the *growing*
+   cache mode), but it is used at **every** step in the static-cache mode
+   (full cache in the `key`/`value` slots).
 3. **GQA contrib op rejects `attention_bias`** — use standard ONNX
    `Attention` if you need bias with GQA.
 4. **SM≥8.0** (Ampere+) required for Flash on all paths.
 5. **Float bias is safer** than bool mask for complex attention patterns.
 6. **MEA requires alignment** — `total_kv % 4 == 0` for bias tensors.
+7. **In-place KV is not GQA-exclusive.** GQA has a built-in shared buffer
+   (`past_present_share_buffer`); ONNX `Attention` (opset 24) has no such
+   attribute but reaches the same effect with **`TensorScatter` + a static
+   cache** (vs. the naive growing `concat(past, new)`). GQA's residual
+   decode edge is mostly its dedicated `seq==1` kernels (XQA/Flash-decode).
+8. **Sliding window on the fast path:** set GQA `local_window_size` (=
+   `config.sliding_window`) for uniform-window models instead of baking a
+   float window mask into ONNX `Attention`.
 
 ## Cross-references
 

.agents/skills/debugging-memcpy/SKILL.md

@@ -65,11 +65,12 @@ pkg = build(model_id, execution_provider='cuda', dtype='float16')
 tmpdir = tempfile.mkdtemp(prefix="memcpy_profile_")
 
 for name, model in pkg.items():
-    # IMPORTANT: lower opset to 23 — CUDA EP doesn't register many ops
-    # (including Reshape, Cast) at opset 24.  Without this, you'll see
-    # hundreds of false-positive Memcpy from ops that work fine at opset 23.
-    if model.opset_imports.get("", 0) > 23:
-        model.opset_imports[""] = 23
+    # NOTE: ORT ≤1.24.x didn't register CUDA kernels for some opset 24
+    # standard ops.  If profiling on older ORT, lower opset to 23 to
+    # avoid false-positive Memcpy nodes.  On current ORT this is
+    # unnecessary — the correct fix is updating ORT kernel registration.
+    # if model.opset_imports.get("", 0) > 23:
+    #     model.opset_imports[""] = 23
 
     path = os.path.join(tmpdir, f"{name}.onnx")
     ir.save(model, path, external_data=f"{name}.onnx.data")
@@ -120,14 +121,16 @@ This gives you the exact node names causing Memcpy.  Common categories:
 - **Equal/Cast on token IDs** — usually low-impact (small tensors)
 - **CumSum on INT64** — inherent, no GPU kernel
 
-### Critical: opset 24 false positives
+### Critical: opset 24 false positives (ORT ≤1.24.x)
 
-**Always lower opset to 23 before profiling.**  ORT CUDA EP (≤1.24.x)
-does not register kernels for many standard ops at opset 24, including
-`Reshape`, `Cast`, and others.  A Gemma4 decoder at opset 24 shows
-**280 Memcpy** nodes; at opset 23, it shows **4**.  The
-`ort_lower_opset_for_ep` flag in `_flags.py` handles this at runtime,
-but you must apply it manually when profiling with raw ORT sessions.
+**On ORT ≤1.24.x**, CUDA EP did not register kernels for many standard
+ops at opset 24, including `Reshape`, `Cast`, and others.  A Gemma4
+decoder at opset 24 showed **280 Memcpy** nodes; at opset 23, just **4**.
+The correct fix is to update ORT kernel registration for the missing
+opset versions — not to lower the model's opset.  The
+`ort_lower_opset_for_ep` flag in `_flags.py` is available as a
+workaround (disabled by default, opt-in via
+`MOBIUS_ORT_LOWER_OPSET_FOR_EP=1`).
 
 ## CPU-only op reference (ORT CUDA EP)
 
@@ -223,6 +226,11 @@ and the `GreaterOrEqual` for causality.
 with `is_causal=1`.  Not applicable to `GroupQueryAttention` (GQA handles
 masking internally via `local_window_size`).
 
+> **Note**: For complex attention patterns (e.g. dual head_dim, KV-shared
+> layers, mixed sliding/full attention), prefer **float additive bias** over
+> bool masks.  Float masks are batch-safe and work correctly with MEA
+> (Memory Efficient Attention) on CUDA.
+
 ### Pattern 3: Use GQA's built-in local_window_size
 
 **Problem**: Sliding-window attention requires an explicit mask (CumSum-based)
@@ -303,12 +311,14 @@ at opset 23.  Results after optimization:
 | Decoder | 4 | `input_ids` (input), 2× `Equal` (token masks), `Where` (bool mask) |
 | Embedding | 3 | `input_ids` (input), 2× `Equal` (token masks) |
 
-### Opset 24 trap
+### Opset 24 historical note (ORT ≤1.24.x)
 
-Without opset lowering, the decoder showed **280 Memcpy** nodes because
-CUDA EP doesn't register `Reshape`, `Cast`, and other standard ops at
-opset 24.  The `ort_lower_opset_for_ep` flag (enabled by default) fixes
-this at runtime.  Always lower to opset 23 before profiling.
+On ORT ≤1.24.x, the decoder showed **280 Memcpy** nodes because
+CUDA EP didn't register `Reshape`, `Cast`, and other standard ops at
+opset 24.  This has been fixed in newer ORT versions.  The
+`ort_lower_opset_for_ep` flag (disabled by default, opt-in via
+`MOBIUS_ORT_LOWER_OPSET_FOR_EP=1`) is available as a workaround for
+older ORT builds.
 
 ## Impact assessment
 

.agents/skills/debugging-multimodal/SKILL.md

@@ -254,9 +254,28 @@ ORT bug: microsoft/onnxruntime#28107
 
 ### Opset 24 kernel registration
 
-ORT ≤1.24.x CUDA/TRT EPs don't register kernels for opset 24.
-**Fix:** Use the `ort_lower_opset_for_ep` feature flag (enabled by
-default). See `src/mobius/_flags.py`.
+ORT ≤1.24.x CUDA/TRT EPs didn't register kernels for opset 24.
+This has been fixed in newer ORT versions.  The `ort_lower_opset_for_ep`
+feature flag is available as a workaround (disabled by default, opt-in
+via `MOBIUS_ORT_LOWER_OPSET_FOR_EP=1`).  See `src/mobius/_flags.py`.
+
+### Encoder input dtype alignment
+
+Encoder task inputs should be declared with `dtype=config.dtype` so
+entry tensors match the model compute dtype (float32/float16/bfloat16).
+In the current codebase, multimodal encoder task builders set encoder
+inputs directly to `config.dtype` (there is no `_cast_encoder_input()`
+helper in `src/mobius/tasks/_base.py`).
+
+### GQA for KV-shared layers
+
+Gemma4 KV-shared layers now emit `GroupQueryAttention` with empty K/V
+inputs (`kv_sequence_length=0`) and borrowed source-layer KV wired via
+`past_key`/`past_value`, avoiding extra Transpose/Reshape cache ops.
+
+Runtime support depends on ORT having KV-shared GQA support (tracked in
+microsoft/onnxruntime#28242; still upstreaming as of this writing). On
+ORT builds without that support, this path can fail at runtime.
 
 ### NaN for large head_dim (> 256)