Merge origin/main into feat/paged-cache; fold --paged-cache into --features

Update from main (brings in the cargo-style --features build option) and
fold the paged / block-table KV cache toggle into it:

- Add 'paged-cache' to _BUILD_FEATURES so 'mobius build --features paged-cache'
  enables the paged KV cache, matching static-cache / fp8-kv-cache / text-only.
- Remove the standalone --paged-cache boolean flag (main dropped the other
  boolean feature flags); --page-size / --num-pages remain as tuning params.
- Reword paged validation errors and docs (README, cli_reference, CHANGELOG)
  to reference --features paged-cache.
- Add a cli_test covering --features paged-cache.

Conflicts resolved in __main__.py, tasks/_causal_lm.py (keep both paged params
and prune_lm_head), and CHANGELOG.md.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Signed-off-by: justinchuby <justinchuby@users.noreply.github.com>

Author: justinchuby

.gitignore

@@ -228,3 +228,5 @@ docs/feature-flags.md
 *.m4a
 *.wav
 .olive-cache/
+.scratch/
+uv.lock

CHANGELOG.md

@@ -7,11 +7,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
-### Paged / block-table KV cache export (`--paged-cache`)
+### Paged / block-table KV cache export (`--features paged-cache`)
 
 #### Added
 
-- `CausalLMTask(paged_cache=True)` and the `mobius build --paged-cache` CLI flag
+- `CausalLMTask(paged_cache=True)` and `mobius build --features paged-cache`
   export a **paged / block-table KV cache** (onnx-genai `docs/DESIGN.md` §39.4
   Option C — vLLM PagedAttention / SGLang RadixAttention layout). KV lives in a
   shared per-layer **page pool** `key_pool.{i}` / `value_pool.{i}`
@@ -20,31 +20,102 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   for each newly written token. Attention writes new K/V into the pool via
   `ScatterND`, assembles the sequence's pages contiguously via
   `Gather(pool, block_table)`, then runs the opset-24 `Attention` op with
-  `nonpad_kv_seqlen` (input #6) — identical op contract to `--static-cache`, but
-  over non-contiguous pages. Because sequences can list the *same* physical page
-  in their `block_table`, the same graph expresses RadixAttention shared-prefix
-  pages with no change. Paging uses only standard ONNX ops (no custom op).
-  New tuning flags `--page-size` (default 16) and `--num-pages` (dynamic when
-  omitted). Requires `DecoderLayer` / `MoEDecoderLayer` models; mutually
-  exclusive with `--static-cache`. Targets a single active sequence
-  (`batch == 1`); multi-sequence batching is a documented TODO.
+  `nonpad_kv_seqlen` (input #6) — identical op contract to `--features
+  static-cache`, but over non-contiguous pages. Because sequences can list the
+  *same* physical page in their `block_table`, the same graph expresses
+  RadixAttention shared-prefix pages with no change. Paging uses only standard
+  ONNX ops (no custom op). New tuning flags `--page-size` (default 16) and
+  `--num-pages` (dynamic when omitted). Requires `DecoderLayer` /
+  `MoEDecoderLayer` models; mutually exclusive with `--features static-cache`.
+  Targets a single active sequence (`batch == 1`); multi-sequence batching is a
+  documented TODO.
+
+### NVIDIA Cosmos 3 Edge vision-language model (`cosmos3_edge`)
 
----
+#### Added
+
+- Support for the **full `cosmos3_edge` vision-language model**
+  (`nvidia/Cosmos3-Edge`, `Cosmos3EdgeForConditionalGeneration`) as a 3-model
+  onnxruntime-genai split (`decoder` + `vision_encoder` + `embedding`):
+  - **decoder**: grouped-query-attention text reasoner with a **non-gated
+    squared-ReLU FFN** (`hidden_act="relu2"`, `up_proj → relu2 → down_proj`)
+    and 3D multimodal RoPE (`mrope_section=[24, 20, 20]`); takes
+    `inputs_embeds`.
+  - **vision_encoder**: SigLIP vision tower + a new
+    `Cosmos3EdgeMultiModalProjector` (pre-shuffle `LayerNorm` → 2×2
+    pixel-shuffle → `linear_fc1` → GELU → `linear_fc2`).
+  - **embedding**: token embedding + image-feature fusion at
+    `image_token_id=19`.
+  `preprocess_weights` routes the single HF checkpoint to the three
+  sub-models: `model.visual.*` / `model.projector.*` → vision (with SigLIP
+  `mlp.fc1/fc2` → `up_proj/down_proj`), `embed_tokens` → embedding, the
+  top-level text tower (`layers.*` / `norm` / `lm_head`) → decoder (renaming
+  `self_attn.to_{q,k,v,out}` → `{q,k,v,o}_proj`), and drops the
+  generator-tower `k_norm_und_for_gen` key-norm. Built via a new
+  `Cosmos3EdgeVLTask` (`cosmos3-edge-vl`). The decoder-only text reasoner
+  remains available as `cosmos3_edge_text`.
+- **L1 graph-build tested only.** NVIDIA does not publish modeling code for
+  `cosmos3_edge` (not in `transformers`, no remote-code module), so the exact
+  pixel-shuffle ordering and numerical parity are unverifiable; L4/L5 parity
+  is deferred. The `cosmos3_omni` variants (`Cosmos3-Nano`/`-Super`) are
+  two-tower diffusion world models tracked separately.
+
+### Cargo-style `--features` build option
+
+#### Added
+
+- `mobius build --features <a,b,...>` collects the build-mode toggles under a
+  single Rust/cargo-style option. Accepts a comma-separated list and may be
+  repeated (`--features fp8-kv-cache,static-cache` or `--features fp8-kv-cache
+  --features static-cache`). Available features: `static-cache`, `fp8-kv-cache`,
+  `prune-lm-head`, `text-only`. Unknown feature names are rejected with an error
+  listing the valid set.
+
+#### Changed
+
+- The boolean flags `--static-cache`, `--fp8-kv-cache`, and `--text-only` have
+  been **removed** in favor of the equivalent `--features` value. Companion
+  value args (`--max-seq-len`, `--kv-cache-scale-file`) are unchanged.
+
+### Final-token LM-head pruning (`--features prune-lm-head`)
+
+#### Added
+
+- `build(prune_lm_head=True)` and `mobius build --features prune-lm-head`
+  select the final hidden-state position before the LM-head projection, reducing
+  prefill logits from `[B, S, vocab]` to `[B, 1, vocab]`. This avoids computing
+  unused per-token logits for single-token autoregressive generation. Models
+  with custom forward paths that do not support pre-projection pruning fail
+  explicitly instead of silently producing an unoptimized graph.
+
+### FP8 (E4M3) KV-cache export (`--features fp8-kv-cache`)
+
+#### Added
+
+- `build(fp8_kv_cache=True)` and `mobius build --features fp8-kv-cache` retype the
+  fused `GroupQueryAttention` KV cache to `FLOAT8E4M3FN` (per-tensor E4M3) after
+  GQA fusion, adding `k_scale`/`v_scale` initializers and the
+  `k_quant_type`/`v_quant_type="PER_TENSOR"`, `kv_cache_bit_width=8` attributes.
+  Halves KV-cache memory at long context on ORT runtimes with the FP8 KV kernel
+  (SM89+). `--kv-cache-scale-file` supplies calibrated per-layer scales
+  (onnxruntime-genai format); without it all layers use a unit scale of 1.0.
+  Only graph-input or empty-placeholder caches are retyped — a non-empty
+  initializer cache is skipped with a warning.
 
-### Text-only export for multimodal Gemma 4 (`--text-only`)
+### Text-only export for multimodal Gemma 4 (`--features text-only`)
 
 #### Added
 
-- `build(text_only=True)` and the `mobius build --text-only` CLI flag export the
+- `build(text_only=True)` and `mobius build --features text-only` export the
   **text backbone** of a unified multimodal checkpoint as a standalone
   decoder-only LLM. For `gemma4_unified` (`google/gemma-4-12B`) this remaps the
   model type to its text sibling (`gemma4_unified_text`) and strips the
   vision/audio config so the decoder fuses to `GroupQueryAttention` on
   GQA-capable execution providers (CUDA/DML) instead of the float-bias
   `Attention` path forced by the multimodal bidirectional vision-block overlay.
-  `--text-only` is rejected with `--config` / `--component` and now also bypasses
-  diffusers autodetect so `build()` validation runs (a diffusers/unsupported repo
-  raises instead of silently exporting a pipeline).
+  The `text-only` feature is rejected with `--config` / `--component` and now
+  also bypasses diffusers autodetect so `build()` validation runs (a
+  diffusers/unsupported repo raises instead of silently exporting a pipeline).
 
 #### Changed
 

README.md

@@ -112,6 +112,15 @@ mobius build --model Qwen/Qwen-Image-2512 output_dir/
 mobius build --model openai/whisper-tiny output_dir/
 ```
 
+Build-mode toggles use the cargo-style `--features` option. Available features
+are `static-cache`, `fp8-kv-cache`, `paged-cache`, `prune-lm-head`, and
+`text-only`. Pass them as a comma-separated list or repeat the option:
+
+```sh
+mobius build --model meta-llama/Llama-3.2-1B output_dir/ \
+    --features static-cache,prune-lm-head --max-seq-len 2048
+```
+
 See the [CLI Reference](https://onnxruntime.github.io/mobius/cli_reference.html) for all subcommands and flags.
 
 ### Examples

docs/cli_reference.md

@@ -168,31 +168,67 @@ mobius build --model Qwen/Qwen2.5-0.5B output/ \
     --ep cuda --dtype f16 --runtime ort-genai
 ```
 
-### Static Cache (`--static-cache`)
+### Build Features (`--features`)
 
+Build-mode toggles are collected under a single cargo-style `--features`
+option. Pass a comma-separated list (and/or repeat the flag):
+
+```
+--features fp8-kv-cache,static-cache
+--features prune-lm-head
+--features text-only
 ```
---static-cache
+
+Available features:
+
+| Feature | Effect |
+|---------|--------|
+| `static-cache` | Pre-allocate fixed-size KV cache buffers using `TensorScatter` (pair with `--max-seq-len N`). Requires `DecoderLayer` / `MoEDecoderLayer` models. Cannot combine with `--task`. |
+| `fp8-kv-cache` | Store the `GroupQueryAttention` KV cache as `FLOAT8E4M3FN` (per-tensor E4M3), halving KV-cache memory. Requires a GQA build (e.g. `--ep cuda --dtype f16`) and an ORT runtime with the FP8 KV-cache kernel (SM89+). Pair with `--kv-cache-scale-file` for calibrated scales. |
+| `paged-cache` | Export a paged / block-table KV cache (vLLM PagedAttention / SGLang RadixAttention layout): a shared page pool plus `block_table` / `slot_mapping`, using only standard ONNX ops. Tune with `--page-size N` (default 16) and `--num-pages N` (dynamic when omitted). Requires `DecoderLayer` / `MoEDecoderLayer` models. Cannot combine with `--task` or the `static-cache` feature. |
+| `prune-lm-head` | Select the final hidden-state position before the LM-head projection and emit logits shaped `[B, 1, vocab]`. Supported by models using the base `CausalLMModel.forward()` path; unsupported custom forwards fail explicitly. Use only when the downstream workflow does not need per-token logits. |
+| `text-only` | Export the text backbone of a multimodal checkpoint as a standalone decoder-only LLM (see below). |
+
+The legacy boolean flags `--static-cache`, `--fp8-kv-cache`, and
+`--text-only` have been removed in favor of `--features`.
+
+```bash
+mobius build --model meta-llama/Llama-3.2-1B output/ \
+    --features static-cache --max-seq-len 2048
+
+mobius build --model Qwen/Qwen2.5-0.5B output/ \
+    --ep cuda --dtype f16 --features fp8-kv-cache
+
+mobius build --model meta-llama/Llama-3.2-1B output/ \
+    --features prune-lm-head
+```
+
+### Static Cache (`--features static-cache`)
+
+```
+--features static-cache
 --max-seq-len N
 ```
 
 Pre-allocate fixed-size KV cache buffers using TensorScatter. Useful when
 the maximum sequence length is known up front.
 
-- `--static-cache` enables static cache mode. Requires models using
+- `--features static-cache` enables static cache mode. Requires models using
   `DecoderLayer` or `MoEDecoderLayer`.
 - `--max-seq-len N` sets the maximum sequence length for static cache
-  buffers. Only valid with `--static-cache`. Defaults to
+  buffers. Only valid with static cache. Defaults to
   `max_position_embeddings` from the model config.
 
 Cannot be combined with `--task`.
 
 #### Example
 
 ```bash
-mobius build --model meta-llama/Llama-3.2-1B output/ --static-cache
+mobius build --model meta-llama/Llama-3.2-1B output/ --features static-cache
 
 # With explicit max sequence length
-mobius build --model meta-llama/Llama-3.2-1B output/ --static-cache --max-seq-len 2048
+mobius build --model meta-llama/Llama-3.2-1B output/ \
+    --features static-cache --max-seq-len 2048
 ```
 
 ### Other Flags
@@ -206,13 +242,13 @@ mobius build --model meta-llama/Llama-3.2-1B output/ --static-cache --max-seq-le
 | `--max-shard-size SIZE` | Maximum shard size for safetensors external data (e.g. `5GB`). Only used with `--external-data safetensors`. |
 | `--trust-remote-code` | Trust remote code when loading the HuggingFace model config. |
 | `--component NAME` | Build only one component from a diffusers pipeline (e.g. `--component vae_decoder`). |
-| `--text-only` | Export the text backbone of a multimodal checkpoint as a standalone decoder-only LLM. Strips vision/audio routing so the decoder uses `GroupQueryAttention` on GQA-capable EPs (build with `--ep cuda`/`dml`). Currently supported for `gemma4_unified` (`google/gemma-4-12B`). Not compatible with `--config` or `--component`. |
+| `--kv-cache-scale-file PATH` | Optional JSON file of calibrated per-layer FP8 KV-cache scales (onnxruntime-genai format). Only used with the `fp8-kv-cache` feature; without it all layers use a unit scale of 1.0. |
 
 #### Text-only example
 
 ```bash
 # Export gemma-4-12B's text backbone as a GQA decoder-only LLM
-mobius build --model google/gemma-4-12B output/ --text-only --ep cuda --dtype f16
+mobius build --model google/gemma-4-12B output/ --features text-only --ep cuda --dtype f16
 ```
 
 For a full ORT-GenAI text-only package (with `genai_config.json`), use

examples/gemma4_12b_text_ort_genai.py

@@ -56,7 +56,7 @@
 
 The equivalent raw-ONNX (no genai_config) build is::
 
-    mobius build --model google/gemma-4-12B --text-only --ep cuda --dtype f16 \
+    mobius build --model google/gemma-4-12B --features text-only --ep cuda --dtype f16 \
         out/gemma4_12b_text_onnx/
 """
 

pyproject.toml

@@ -14,8 +14,9 @@ requires-python = ">=3.10"
 license = "MIT"
 dependencies = [
     "huggingface_hub",
+    "ml_dtypes",
     "numpy>=1.24.0",
-    "onnx_ir>=0.1.0",
+    "onnx_ir>=0.2.1",
     "onnx-shape-inference>=0.3.1",
     "onnxscript>=0.7.1",
     "safetensors",
@@ -136,6 +137,7 @@ ignore = [
     "RUF031",
     "RUF052",
     "RUF067",
+    "RUF105",
     "SIM102",
     "SIM108",
     "SIM114",

requirements/lintrunner/requirements.txt

@@ -1,3 +1,3 @@
 lintrunner-adapters>=0.14.0
 # RUFF, RUFF-FIX
-ruff==0.15.14
+ruff==0.16.1

scripts/generate_dashboard.py

@@ -145,7 +145,7 @@ def l5_passes(self) -> bool:
 
     @property
     def confidence_level(self) -> int:
-        """Return the highest confidence level achieved (0-5)."""
+        """The highest confidence level achieved (0-5)."""
         if self.l5_passes:
             return 5
         if self.l4_passes:

scripts/generate_sortformer_golden.py

@@ -0,0 +1,106 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+"""Regenerate the Sortformer diarization golden reference used by the L4/L5 tests.
+
+This produces ``testdata/golden/speech/sortformer_diarization.npz`` by running
+the *real* NeMo Sortformer model through the NeMo toolkit (the ground-truth
+reference implementation). It must be run inside an environment that has
+``nemo_toolkit`` installed (it is **not** a mobius runtime dependency)::
+
+    python -m venv /tmp/nemo_ref_venv
+    source /tmp/nemo_ref_venv/bin/activate
+    pip install "nemo_toolkit[asr]"
+    python scripts/generate_sortformer_golden.py \
+        --model nvidia/diar_streaming_sortformer_4spk-v2.1 \
+        --revision fafaab5faa1617a0ca52d38dd3dc4bd636800d3d \
+        --out testdata/golden/speech/sortformer_diarization.npz
+
+The offline forward path is ``frontend_encoder`` (mel features -> embedding
+sequence) followed by ``forward_infer`` (embeddings -> per-frame speaker
+activity sigmoids). The committed ``.npz`` stores the mel input, the encoder
+embeddings, and the speaker probabilities, plus a ``meta`` JSON blob (model id,
+revision, NeMo version, dtype, seed) so the reference is self-describing and
+auditable.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+
+import numpy as np
+import torch
+
+# Deterministic mel-feature fixture (also recorded in metadata).
+_SEED = 0
+_T = 400  # mel frames; with 8x subsampling -> 50 output diarization frames.
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--model", default="nvidia/diar_streaming_sortformer_4spk-v2.1")
+    parser.add_argument(
+        "--revision",
+        default="fafaab5faa1617a0ca52d38dd3dc4bd636800d3d",
+        help="HuggingFace Hub commit SHA to pin the reference model.",
+    )
+    parser.add_argument(
+        "--out",
+        default="testdata/golden/speech/sortformer_diarization.npz",
+    )
+    args = parser.parse_args()
+
+    import nemo  # type: ignore[import-not-found]
+    from huggingface_hub import hf_hub_download
+    from nemo.collections.asr.models import (  # type: ignore[import-not-found]
+        SortformerEncLabelModel,
+    )
+
+    torch.manual_seed(_SEED)
+
+    nemo_path = hf_hub_download(
+        repo_id=args.model,
+        filename="diar_streaming_sortformer_4spk-v2.1.nemo",
+        revision=args.revision,
+    )
+    model = SortformerEncLabelModel.restore_from(nemo_path, map_location="cpu")
+    model.eval()
+    # Offline (non-streaming) forward path: full-context attention.
+    model.streaming_mode = False
+
+    feat_dim = int(model.cfg.encoder.feat_in)
+    mel = torch.randn(1, feat_dim, _T)
+    mel_len = torch.tensor([_T], dtype=torch.long)
+
+    with torch.no_grad():
+        emb_seq, emb_len = model.frontend_encoder(
+            processed_signal=mel, processed_signal_length=mel_len
+        )
+        preds = model.forward_infer(emb_seq, emb_len)
+
+    num_spks = int(preds.shape[-1])
+    meta = {
+        "model_id": args.model,
+        "revision": args.revision,
+        "nemo_version": nemo.__version__,
+        "dtype": "float32",
+        "seed": _SEED,
+        "feat_dim": feat_dim,
+        "input_frames": _T,
+        "num_spks": num_spks,
+    }
+
+    np.savez_compressed(
+        args.out,
+        mel=mel.numpy().astype(np.float32),
+        emb_seq=emb_seq.numpy().astype(np.float32),
+        emb_len=emb_len.numpy().astype(np.int64),
+        preds=preds.numpy().astype(np.float32),
+        meta=np.array(json.dumps(meta)),
+    )
+    print(f"saved {args.out}\n{json.dumps(meta, indent=2)}")
+
+
+if __name__ == "__main__":
+    main()