[milesial]

Before this change, latent-MoE Nemotron-H models priced the shared-expert and fc1/fc2 latent projections with the global GEMM precision, usually FP8, even when the checkpoint gave them a different precision. After this change, the SDK reads the precision of these projections from the checkpoint metadata, for example NVFP4 shared experts and BF16 latent projections. In a simulated GB200 / vLLM 0.31.0 prediction at TP1, the predicted mean TTFT goes down by 4.84% and the mean TPOT goes up by 1.68%. These are simulator predictions, not hardware measurements.

Why and what changed

• sdk/models/nemotron_h.py reads the precision of the four projections from quantized_layers listings and from exclusions. It also reads config-group targets and Linear class targets. It accepts the model, backbone and language_model.* path prefixes.
• A projection that a targets-only MIXED_PRECISION listing leaves out is BF16. This is the vLLM and SGLang rule that helpers.py already cites (fix: model nvidia/GLM-5.3-NVFP4 ModelOpt MIXED_PRECISION targets-only quantization #610). For other metadata, an unlisted projection keeps the global precision.
• On a system that is not Blackwell, an NVFP4 projection resolves to weight-only nvfp4_wo, from ModelConfig.system. This is the existing policy for the global precision. Blackwell keeps native nvfp4.
• An explicit gemm_quant_mode and MoE models without a latent projection keep the current behavior.
• If the MoE layers give a projection different precisions, the SDK emits one GEMM per precision under the same op name, each counting its own layers. Same-name ops add up, so the total equals the per-layer cost. Conflicting or unsupported algorithms for one module path still fail.
• helpers.py: _collect_mixed_precision_layer_algos now uses the new _iter_mixed_precision_layer_algos generator, which calls _mixed_precision_layer_algo from fix: model nvidia/GLM-5.3-NVFP4 ModelOpt MIXED_PRECISION targets-only quantization #610. The result does not change.

Review map

• Risk level: medium.
• Start with: _projection_modes and NemotronHModel.create in sdk/models/nemotron_h.py.
• Public or serialized contract changed: no. The predicted precision changes only for latent-MoE Nemotron-H checkpoints whose metadata lists or excludes these projections.
• Compatibility or rollback concern: uniform checkpoints keep one op per projection with the same count as before. A checkpoint with uneven per-layer precision now gets one op per precision. For example, a public Super NVFP4 export that lists its projections only in part (shared up: FP8 in 39 layers, unlisted in 1) loads and prices each layer at its own precision. Metadata without per-layer information keeps the global precision. docs/perf-model/limitations.md describes this.

Evidence

• Tests: tests/unit/sdk/models/test_nemotron_mixed_precision.py has 38 tests on synthetic metadata. They cover uniform NVFP4 and mixed FP8/NVFP4 metadata through Task.build_model_config on Hopper and Blackwell, in prefill and decode. They also cover exclusions, explicit overrides and the packaged Super and Ultra configs. For a 40-layer partial listing, they check the per-precision layer counts, weight bytes against a per-layer reference, and latency against single-precision references. tests/unit/sdk/models and tests/unit/sdk/task_v2 pass. Ruff passes.
• Fast CI: pass at 089afa7 (rebased on main d9620e0, same patches).
• Full CI: pass at 089afa7 (63 pass, 7 skipped).
• CodeRabbit reviewed commit: none. CodeRabbit does not review draft PRs.
• Codex reviewed commit: none on GitHub.
• Negative or boundary cases: conflicting or unsupported algorithms on one module path fail to load. An explicit gemm_quant_mode prices all four projections with one mode. Metadata that is not MIXED_PRECISION keeps the global precision for an unlisted projection. Hopper resolves NVFP4 to nvfp4_wo, and the Hopper cases fail without that fix.
• Expected-value derivation (for exact numerical assertions): the tests assert precision modes only, not latencies.
• Before/after output, trace, benchmark, or golden diff: the comparison uses the public aisimulate predict CLI and the RustForwardPassPerfModel API. The target is simulated gb200 with vLLM 0.31.0 in strict op-level mode (SILICON, fallback deny). Base is main 17625dc. The head is the earlier PR head 6f28796. That checkpoint gives every MoE layer the same precision, so 517590e prices it the same way. Both use the same native build and identical table data. The model metadata comes from a pinned checkpoint that lists the shared experts as NVFP4 and excludes the latent projections. Not every public Super export has this layout.
  • Per-op precision: shared up and down change from FP8 to NVFP4. fc1 and fc2 change from FP8 to BF16. The dense projections do not change.
  • With the packaged tables only, strict mode fails in the same way before and after: MoE data missing. The packaged gb200 vLLM 0.31.0 tables have no MoE rows for this shape, so no strict latency exists.
  • With added MoE rows, strict mode succeeds at TP1 and TP2 in both variants. TP4 and TP8 still fail with MoE data missing.
  • Public CLI, TP1, 1024 input and 16 output tokens, concurrency 4, 8 requests, seeded MoE rows. Both variants exit 0, complete 8 of 8 requests, produce 128 output tokens, and report no errors.

Predicted metric	Base	Head	Change
Mean TTFT (ms)	95.27	90.66	-4.84%
Mean TPOT (ms)	6.136	6.239	+1.68%
Output tok/s (1 GPU)	341.7	347.4	+1.66%

• A separate run used an earlier unseeded MoE dataset from a different GPU. It gave -4.85% TTFT, +1.69% TPOT and +1.67% output tok/s. These similar changes describe the two runs only. They do not measure a seed effect or the accuracy on hardware.
• In strict mode at TP1 prefill, the count of operations priced from collected tables goes from 6 to 10. The count of operations estimated from other tables goes from 8 to 4.
• Public refreshed metadata: the published nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-NVFP4 (revision ff433f54) and nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B-NVFP4 (revision 02462641) configs were used as local inputs only. They are not part of this PR. With them, the previous head 28b18f6 fails to load Super ("different precision across layers", shared up FP8 in 39 layers and BF16 in 1). This head loads Super on Blackwell and Hopper with exact per-precision layer counts. Strict TP1 succeeds with the added MoE rows, and the public CLI completes 8 of 8 requests with 128 output tokens. Ultra gives one precision per projection and behaves the same on both heads. These Super numbers use the export's own uneven precision, so they do not compare with the pinned-checkpoint table above.

Modeling or data provenance

• The added MoE rows are not part of this PR. They are NVFP4 rows for hidden size 1024, intermediate size 2688, top-k 22 and 512 experts. They come from GB200 with vLLM 0.31.0, measured per rank on one GPU, for TP1 and TP2 only. The run added them to a copy of the table.
• Each variant ran once. The simulator is deterministic. This PR claims no host timing and no GPU timing, and the run is not an AgentX workload.
• A strict pass does not mean that the tables measure every operation:
  • Mamba conv1d and ssm use speed-of-light estimates (sol).
  • Attention uses vLLM 0.25.0 rows under the declared 0.31.0.
  • Dense fp8_static GEMMs are estimated from FP8 rows.
  • Norms and activations use formulas.

Tracking

• Closes: none.
• Related PRs or issues: fix(perfmodel): add Nemotron NVFP4 and FlashInfer Mamba support #608 changes the same nemotron_h.py and helpers.py sections. This PR takes only the projection precision part, with no Mamba controls. fix: model nvidia/GLM-5.3-NVFP4 ModelOpt MIXED_PRECISION targets-only quantization #610 is merged into this branch.

[coderabbitai]

Important

Draft PR not reviewed

Draft PRs are not automatically reviewed by default.

• Trigger a manual review

To automatically review draft PRs, update your CodeRabbit configuration:

reviews:
  auto_review:
    drafts: true

• Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

_{Comment @coderabbitai help to get the list of available commands.}

[milesial]

Verdict: changes requested.

I checked the PR premise against the upstream nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-NVFP4 metadata (config.json quantization_config and hf_quant_config.json, fetched from the HF hub today). It does not match the PR description, and the new resolver fails on it:

projection	FP8	NVFP4	unlisted (of 40 E layers)
shared_experts.up_proj	39	0	1
shared_experts.down_proj	37	1	2
fc1_latent_proj	3	0	37
fc2_latent_proj	1	0	39

• Shared experts are FP8 in nearly every layer, not NVFP4. The latent projections are mostly unlisted, not excluded (ignore/exclude_modules are empty).
• Because the listing is partial and layer-dependent, _projection_modes raises different precision across layers for every projection. Loading this checkpoint from the hub or a local download, which worked before this PR, now fails. See inline comments.
• The bundled model_configs/*Nemotron-3-Super*NVFP4* and *Ultra*NVFP4* files carry only quant_algo: NVFP4 with no quantized_layers or config_groups, so the default bundled path is a no-op. The unit test uses synthetic config_groups metadata and does not exercise either real shape.

I could not fetch the Ultra 550B NVFP4 upstream metadata (gated), so the Ultra numbers are unverified.

Optional (not required):

• _iter_mixed_precision_layer_algos: quant.get("config_groups") is no longer guarded by isinstance(quant, dict); a string-valued quantization_config would raise where the old code tolerated it.
• The # shortcut: comment style on the create branch is not used elsewhere in this module.
• helpers.py refactor itself is behavior-preserving as far as I can tell.

[milesial]

Ready from code review; no required changes found.

Reviewed 0c58048 against main at bf373a1. Traced metadata collection, exclusions and aliases, explicit-mode precedence, both operation pipelines, and the Rust-backed GEMM precision handoff. The new tests cover the principal success and rejection paths. The documented limits are clear: bundled metadata keeps the global mode, and layer-dependent precision requires an explicit override.

Host-cost assessment: resolution runs during construction; literal layer declarations are indexed, avoiding a full listing scan for every projection/layer. No added per-query resolution was found. This is a source assessment, not a measured performance result.

Validation: git diff --check passed and the checkout is clean. Fast CI passed; broader CI was still pending when inspected. Tests, builds, benchmarks, and the reported E2E results were not rerun or independently reproduced under this read-only review.

[milesial]

Superseded by #617, which now contains both commits of this PR unchanged.