Add NVIDIA Nemotron 3.5 Lightning export support

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

@@ -476,6 +476,25 @@ config.norm_topk_prob                # Whether to normalize routing weights
 config.routed_scaling_factor         # Post-normalization scale
 ```
 
+Current Transformers exposes NemotronH layer types as
+`linear_attention` / `full_attention`; older configs use
+`mamba` / `attention`. Normalize both vocabularies to Mobius
+`mamba2` / `full_attention`, preserve `mlp` and `moe` distinctly, and reject
+unknown values. Never map `mlp` to `moe` in parity fixtures.
+
+### Reduced-precision routing
+
+ONNX has no implicit mixed-float type promotion. NemotronH routing computes in
+fp32, so keep the correction-bias initializer in fp32 and explicitly cast the
+gate weight to fp32. Cast each expert output up, multiply and accumulate all
+routed contributions in fp32, then cast the completed routed tensor back once.
+Graph construction alone may miss this; execute fp16 and bf16 MoE paths.
+
+Official Nemotron 3.5 checkpoints also contain auxiliary `mtp.*` tensors.
+The base `NemotronHForCausalLM` generation graph does not instantiate them and
+marks them unexpected. Filter only that prefix and prove weight alignment still
+populates every base-decoder initializer.
+
 ### com.microsoft.MoE compatibility
 
 **Not compatible with NemotronH.** Three blockers:

.github/workflows/gpu_l4_golden_parity.yml

@@ -56,6 +56,12 @@ jobs:
           path: ~/.cache/huggingface
           key: hf-gpu-${{ hashFiles('testdata/cases/**/*.yaml') }}
 
+      - name: Cache reduced Nemotron fixture
+        uses: actions/cache@v6
+        with:
+          path: ${{ runner.temp }}/mobius-nemotron-cache
+          key: nemotron-reduced-d468880b6ad3c6e0d21377ce7242adaea4cc884d-schema-v1
+
       - name: Install PyTorch (CUDA)
         run: pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu128
 
@@ -82,11 +88,12 @@ jobs:
         env:
           HF_TOKEN: ${{ secrets.HF_TOKEN }}
           MOBIUS_TEST_DEVICE: cuda
+          MOBIUS_NEMOTRON_REDUCED_CACHE: ${{ runner.temp }}/mobius-nemotron-cache/reduced-d468880b6ad3c6e0d21377ce7242adaea4cc884d-schema-v1.safetensors
         run: |
           AFFECTED='${{ inputs.affected_models }}'
           if [ -z "$AFFECTED" ] || [ "$AFFECTED" = "[]" ]; then
             echo "Running all L4 golden comparison tests"
-            pytest tests/e2e_golden_test.py \
+            pytest tests/e2e_golden_test.py tests/*_real_weight_test.py \
               -m golden \
               -v \
               --timeout=300 \
@@ -98,7 +105,7 @@ jobs:
             # Convert JSON array to comma-separated list for --models
             MODELS=$(echo "$AFFECTED" | python -c "import json, sys; print(','.join(json.load(sys.stdin)))")
             if [ -n "$MODELS" ]; then
-              pytest tests/e2e_golden_test.py \
+              pytest tests/e2e_golden_test.py tests/*_real_weight_test.py \
                 -m golden \
                 -v \
                 --models "$MODELS" \

.github/workflows/gpu_l5_generation_e2e.yml

@@ -56,6 +56,12 @@ jobs:
           path: ~/.cache/huggingface
           key: hf-gpu-${{ hashFiles('testdata/cases/**/*.yaml') }}
 
+      - name: Cache reduced Nemotron fixture
+        uses: actions/cache@v6
+        with:
+          path: ${{ runner.temp }}/mobius-nemotron-cache
+          key: nemotron-reduced-d468880b6ad3c6e0d21377ce7242adaea4cc884d-schema-v1
+
       - name: Install PyTorch (CUDA)
         run: pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu128
 
@@ -64,17 +70,19 @@ jobs:
           pip install -r requirements/ci/requirements.txt
           pip install soundfile librosa ml_dtypes flatbuffers numpy packaging protobuf sympy coloredlogs
           pip install -e '.[testing,transformers]'
+          pip install -r examples/olive/nemotron-3_5-lightning-30b/requirements.txt --index-url https://packagefeedproxy.microsoft.io/pypi/simple
           pip install --pre --extra-index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/ort-cuda-12-nightly/pypi/simple/ onnxruntime-gpu onnxruntime-genai-cuda
 
       - name: Run L5 generation E2E tests
         env:
           HF_TOKEN: ${{ secrets.HF_TOKEN }}
           MOBIUS_TEST_DEVICE: cuda
+          MOBIUS_NEMOTRON_REDUCED_CACHE: ${{ runner.temp }}/mobius-nemotron-cache/reduced-d468880b6ad3c6e0d21377ce7242adaea4cc884d-schema-v1.safetensors
         run: |
           AFFECTED='${{ inputs.affected_models }}'
           if [ -z "$AFFECTED" ] || [ "$AFFECTED" = "[]" ]; then
             echo "Running all L5 generation E2E tests"
-            pytest tests/e2e_golden_test.py \
+            pytest tests/e2e_golden_test.py tests/*_real_weight_test.py \
               -m generation \
               -v \
               --timeout=300 \
@@ -86,7 +94,7 @@ jobs:
             # Convert JSON array to comma-separated list for --models
             MODELS=$(echo "$AFFECTED" | python -c "import json, sys; print(','.join(json.load(sys.stdin)))")
             if [ -n "$MODELS" ]; then
-              pytest tests/e2e_golden_test.py \
+              pytest tests/e2e_golden_test.py tests/*_real_weight_test.py \
                 -m generation \
                 -v \
                 --models "$MODELS" \

examples/olive/nemotron-3_5-lightning-30b/.gitignore

@@ -0,0 +1,2 @@
+cache/
+output/

examples/olive/nemotron-3_5-lightning-30b/README.md

@@ -0,0 +1,173 @@
+# Nemotron 3.5 Lightning: BF16 checkpoint + Olive
+
+This is **Option A** for
+[`nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16`](https://huggingface.co/nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16):
+export the official BF16 checkpoint to supported FP16 ONNX, quantize the model
+with Olive, assemble a direct ONNX Runtime package, then run cached generation.
+
+Every Hub access is pinned to revision
+`d468880b6ad3c6e0d21377ce7242adaea4cc884d`.
+
+## Architecture and runtime contract
+
+The checkpoint is a real `nemotron_h` model, not an alias:
+
+- 52 base-decoder layers mixing Mamba2, sigmoid-routed MoE, and full GQA.
+- 128 routed experts with top-6 selection and one shared expert.
+- `mtp.*` contains 270 auxiliary multi-token-prediction tensors. This export
+  intentionally targets the base `NemotronHForCausalLM` decoder; its forward
+  graph does not instantiate MTP, and upstream marks those keys unexpected.
+  No base-decoder generation input, cache, logit, or weight depends on them.
+
+This model's graph mixes `conv_state` plus `ssm_state` with sparse
+full-attention key/value caches. Mobius emits every field the current
+`genai_config.json` schema can represent (semantic inputs, key/value and
+convolution templates, and global cache-slot count) and leaves runtime
+acceptance to ORT GenAI. The schema currently has no `ssm_state` template. The
+validated recipe uses direct ONNX Runtime generation through `inference.py`;
+downstream load/generation outcomes are informational and do not gate export.
+
+## Install
+
+From the repository root:
+
+```powershell
+python -m pip install -e ".[transformers,testing]" `
+  --index-url https://packagefeedproxy.microsoft.io/pypi/simple
+python -m pip install -r examples\olive\nemotron-3_5-lightning-30b\requirements.txt `
+  --index-url https://packagefeedproxy.microsoft.io/pypi/simple
+```
+
+Use an ONNX Runtime GPU build with CUDA 12 and cuDNN 9 for CUDA inference.
+
+## Full export, quantization, and smoke test
+
+```powershell
+cd examples\olive\nemotron-3_5-lightning-30b
+python optimize.py `
+  --source-dir output\f16\cuda `
+  --output-dir output\Q4_K_M\cuda `
+  --ep cuda `
+  --precision q4_k_m
+```
+
+The script performs four gated steps:
+
+1. Downloads the exact 14-shard BF16 checkpoint revision and exports FP16 ONNX.
+   BF16 execution is rejected explicitly because corrected reduced-real parity
+   reaches `0.8594` max logit error, above the `1e-2` reduced-precision gate.
+2. Emits standard ONNX cache operations and applies the grouped-RMSNorm CUDA
+   workaround. CUDA still executes supported compute nodes; portable cache
+   operations avoid provider-dependent fused decode drift.
+3. Runs Olive Q4 K-quant with an explicitly CPU-only target. It also suppresses
+   Olive 0.13's unrelated GPU-EP DLL auto-registration, so a missing TensorRT
+   installation cannot abort CPU weight-only quantization.
+4. Reloads the assembled quantized package and generates four cached tokens.
+
+To reuse an existing source package:
+
+```powershell
+python optimize.py --skip-export `
+  --source-dir output\f16\cuda `
+  --output-dir output\Q4_K_M\cuda `
+  --ep cuda
+```
+
+`olive_q4.json` records the equivalent pass and provider-isolated target. Use
+`optimize.py` rather than invoking the JSON directly when the installed ORT
+wheel bundles unconfigured providers; the script contains the verified Olive
+0.13 registration isolation.
+
+## Package layout
+
+```text
+output/
+├── f16/cuda/
+│   ├── model.onnx
+│   ├── model.onnx.data
+│   ├── config.json
+│   ├── generation_config.json
+│   ├── tokenizer.json
+│   ├── tokenizer_config.json
+│   └── source_manifest.json
+└── Q4_K_M/cuda/
+    ├── model.onnx
+    ├── model.onnx.data
+    ├── config.json
+    ├── generation_config.json
+    ├── tokenizer.json
+    ├── tokenizer_config.json
+    └── source_manifest.json
+```
+
+This recipe intentionally uses direct ONNX Runtime, while
+`mobius build --runtime ort-genai` remains allowed and emits the best current
+schema metadata for downstream testing.
+
+## Direct generation and profiling
+
+```powershell
+python inference.py `
+  --model-dir output\Q4_K_M\cuda `
+  --device cuda `
+  --prompt "What is 84 * 3 / 2?" `
+  --max-new-tokens 20 `
+  --profile
+```
+
+The script initializes every cache from the saved graph, processes the prompt
+token by token, carries Mamba and KV state independently, and fails if CUDA was
+requested but not registered.
+
+## Reduced real-checkpoint validation
+
+The full checkpoint is 65.8 GB and cannot execute on the validation host's
+8 GB RTX A1000. The reproducible reduced check range-downloads 236 MiB of real
+weights while retaining production dimensions:
+
+- checkpoint layer 0: complete Mamba2 block;
+- layer 1: router, shared expert, and four complete routed experts;
+- layer 5: complete full-attention block;
+- sliced real embedding and LM-head rows plus final norm.
+
+```powershell
+python validate_reduced_checkpoint.py
+```
+
+The fixture is stored persistently under `~/.cache/mobius/` by default. Each
+range request validates status, `Content-Range`, declared length, and payload
+length, with three bounded attempts (1s then 2s backoff). The cache metadata
+must match the pinned model, revision, and fixture schema; writes are atomic.
+GPU CI restores the same revision/schema-keyed cache for L4 and L5.
+
+The supported matrix is intentionally limited to FP32/CPU and FP16/CUDA.
+Reproduce the BF16 rejection evidence separately without creating a supported
+package or weakening the production guard:
+
+```powershell
+python validate_reduced_checkpoint.py --bf16-rejection-evidence
+```
+
+Validated results on ORT 1.28.0 / Olive 0.13.0:
+
+| Variant | Full-logit max abs | Generated IDs | Placement |
+|---|---:|---|---|
+| FP32 CPU | `9.54e-6` | `12, 13, 12, 12` | CPU |
+| FP16 CUDA | `<= 0.0078125` (prefill + every cached step) | `12, 13, 12, 12` | portable ONNX graph on CUDA |
+| BF16 CUDA | rejected (`0.8594`) | N/A | fails numerical gate |
+| Olive Q4 | quantized | `12, 13, 12, 12` | portable ONNX graph on CUDA |
+
+The reduced package's portable weighted graph quantizes 17 matrix
+multiplications to `com.microsoft::MatMulNBits` and reloads successfully for
+multi-token generation. Record size/compression from the produced package;
+it varies with external-data serialization and Olive version.
+
+## Evidence-based waivers
+
+- Full-checkpoint L4/L5 coherent-text generation: requires roughly 66 GB just
+  for checkpoint storage and substantially more than 8 GB accelerator memory.
+- Full 30B Olive run: the recipe and reduced production-dimension pass are
+  validated; completing all 2,944 expert subgraphs requires a large-memory
+  host.
+- Foundry Local was not available on this host. That downstream validation
+  remains informational and does not block Mobius export.