fix: resolve latent-MoE Nemotron-H shared and latent projection precision from checkpoint

docs/perf-model/limitations.md

@@ -44,6 +44,26 @@ which contributes to the decoder context but not the encoder embeddings.
 For example, a 70×70 image with patch size 14 and merge size 2 gives 36 encoder
 patches, 9 embeddings, and 12 context tokens.
 
+## Nemotron-H latent-MoE projection precision
+
+The SDK reads the precision of the shared-expert and `fc1`/`fc2` latent
+projections of latent-MoE Nemotron-H models from the checkpoint's per-layer
+metadata (`quantized_layers`, `config_groups`, exclusions). The result depends
+on the checkpoint and not on the model name. A projection that a
+MIXED_PRECISION listing leaves out is BF16, the rule vLLM and SGLang apply.
+Metadata that is not a MIXED_PRECISION listing leaves it on the global mode. On
+non-Blackwell systems an NVFP4 projection runs weight-only (`nvfp4_wo`), as the
+global NVFP4 mode does.
+
+If MoE layers give a projection different precisions, the SDK prices each
+precision for its own count of layers. Loading fails only when one module path
+has conflicting or unsupported algorithms. The MTP layers use the decoder
+layers' precision. An explicit `gemm_quant_mode` prices all four projections
+with one mode.
+
+Metadata without per-layer information keeps the global mode for all four
+projections.
+
 ## Other scoped boundaries
 
 - [Support](support-matrix.md) distinguishes engine queries from Replay and accuracy.

python/aisimulate/src/aisimulate_core/sdk/models/helpers.py

@@ -786,46 +786,34 @@ def _is_routing_expert_target(target: object) -> bool:
     return ".experts." in target_name or target_name.endswith(".experts") or "routing_expert" in target_name
 
 
+def _iter_mixed_precision_layer_algos(raw_config: dict):
+    """Yield target, normalized algorithm, and whether the target is a listed layer."""
+    quant = _get_language_quantization_config(raw_config) or {}
+    hf_config = raw_config.get("hf_quant_config")
+    hf_quant = hf_config.get("quantization", {}) if isinstance(hf_config, dict) else {}
+    for source in (hf_quant, quant):
+        if not isinstance(source, dict):
+            continue
+        layers = source.get("quantized_layers")
+        if not isinstance(layers, dict):
+            continue
+        for target, metadata in layers.items():
+            yield target, _mixed_precision_layer_algo(metadata), True
+    groups = quant.get("config_groups") if isinstance(quant, dict) else None
+    if isinstance(groups, dict):
+        for group in groups.values():
+            if isinstance(group, dict):
+                for target in group.get("targets") or []:
+                    yield target, _infer_mixed_precision_group_algo(group), False
+
+
 def _collect_mixed_precision_layer_algos(raw_config: dict) -> tuple[set[str], set[str]]:
     """Collect ModelOpt mixed precision algorithms by SDK GEMM/MoE category."""
     gemm_algos: set[str] = set()
     moe_algos: set[str] = set()
-
-    def add_target(target: object, algo_value: object) -> None:
-        algo = _normalize_mixed_precision_layer_algo(algo_value)
-        if algo is None:
-            return
-        if _is_routing_expert_target(target):
-            moe_algos.add(algo)
-        else:
-            gemm_algos.add(algo)
-
-    def add_quantized_layers(quantized_layers: object) -> None:
-        if not isinstance(quantized_layers, dict):
-            return
-        for target, metadata in quantized_layers.items():
-            add_target(target, _mixed_precision_layer_algo(metadata))
-
-    hf_quant_config = raw_config.get("hf_quant_config")
-    if isinstance(hf_quant_config, dict):
-        hf_quant_section = hf_quant_config.get("quantization")
-        if isinstance(hf_quant_section, dict):
-            add_quantized_layers(hf_quant_section.get("quantized_layers"))
-
-    quant_cfg = _get_language_quantization_config(raw_config)
-    if isinstance(quant_cfg, dict):
-        add_quantized_layers(quant_cfg.get("quantized_layers"))
-        config_groups = quant_cfg.get("config_groups")
-        if isinstance(config_groups, dict):
-            for group in config_groups.values():
-                if not isinstance(group, dict):
-                    continue
-                algo = _infer_mixed_precision_group_algo(group)
-                if algo is None:
-                    continue
-                for target in group.get("targets") or []:
-                    add_target(target, algo)
-
+    for target, algo, _ in _iter_mixed_precision_layer_algos(raw_config):
+        if algo is not None:
+            (moe_algos if _is_routing_expert_target(target) else gemm_algos).add(algo)
     return gemm_algos, moe_algos
 
 

python/aisimulate/src/aisimulate_core/sdk/models/nemotron_h.py

@@ -3,10 +3,101 @@
 
 from __future__ import annotations
 
+import re
+from collections import Counter
+from fnmatch import fnmatchcase
+
 import aisimulate_core.sdk.operations as ops
 from aisimulate_core.sdk import common
 from aisimulate_core.sdk.models.base import BaseModel, register_model
-from aisimulate_core.sdk.models.helpers import mtp_scale_factor
+from aisimulate_core.sdk.models.helpers import (
+    _iter_mixed_precision_layer_algos,
+    modelopt_mixed_precision_targets,
+    mtp_scale_factor,
+    quant_exclude_patterns,
+)
+from aisimulate_core.sdk.perf_database import is_blackwell_system
+
+# Decoder-layer prefixes: HF ``model``, native ``backbone``, and their language_model wrappers.
+_LAYER_PREFIXES = (
+    "model",
+    "backbone",
+    "language_model.model",
+    "language_model.backbone",
+    "model.language_model",
+)
+
+
+def _projection_matches(path: str, target: str) -> bool:
+    if target == "Linear":
+        return True
+    return bool(re.match(target[3:], path)) if target.startswith("re:") else fnmatchcase(path, target)
+
+
+def _projection_modes(raw: dict, pattern: str, default: common.GEMMQuantMode, system: str | None = None) -> dict:
+    """Resolve the shared-expert and latent projections from checkpoint metadata.
+
+    Returns ``{projection: {mode: number of MoE layers}}``. Layers with different
+    precisions get one operation each, so per-layer precision is kept exactly.
+    Conflicting or unsupported algorithms for one module path are rejected.
+    Group targets apply only without a ``quantized_layers`` listing. Exclusions win.
+
+    Metadata without per-layer information leaves every projection on ``default``.
+
+    A projection that a listing leaves out is BF16 only for a targets-only
+    ModelOpt MIXED_PRECISION export, the contract vLLM and SGLang implement
+    (see ``modelopt_mixed_precision_targets``). Any other metadata keeps
+    ``default`` for an unlisted projection instead of guessing a mode.
+    """
+    declarations = []
+    exact = {}
+    group_targets = []
+    layer_listing = False
+    exclusions = quant_exclude_patterns(raw)
+
+    def add(target, algo):
+        if target == "Linear" or target.startswith("re:") or any(char in target for char in "*?["):
+            declarations.append((target, algo))
+        else:
+            exact.setdefault(target, set()).add(algo)
+
+    for target, algo, listed in _iter_mixed_precision_layer_algos(raw):
+        layer_listing |= listed
+        if listed:
+            add(target, algo)
+        else:
+            group_targets.append((target, algo))
+    # A layer listing is authoritative: group targets only apply without one.
+    for target, algo in [] if layer_listing else group_targets:
+        add(target, algo)
+    # Non-Blackwell systems have no native FP4 tensor cores: NVFP4 runs weight-only, as
+    # resolve_nvfp4_for_system does for the global mode. Without a system, keep native NVFP4.
+    nvfp4 = common.GEMMQuantMode.nvfp4_wo if system and not is_blackwell_system(system) else common.GEMMQuantMode.nvfp4
+    modes = {
+        "fp8": common.GEMMQuantMode.fp8 if default == common.GEMMQuantMode.fp8 else common.GEMMQuantMode.fp8_static,
+        "mxfp8": common.GEMMQuantMode.fp8_block,
+        "nvfp4": nvfp4,
+        "w4a16_nvfp4": common.GEMMQuantMode.w4a16_nvfp4,
+    }
+    targets_only = layer_listing and modelopt_mixed_precision_targets(raw) is not None
+    unlisted = common.GEMMQuantMode.bfloat16 if targets_only else default
+    resolved = {}
+    for suffix in ("shared_experts.up_proj", "shared_experts.down_proj", "fc1_latent_proj", "fc2_latent_proj"):
+        values = []
+        for layer, kind in enumerate(pattern):
+            if kind != "E":
+                continue
+            paths = [f"{prefix}.layers.{layer}.mixer.{suffix}" for prefix in _LAYER_PREFIXES]
+            if any(_projection_matches(path, target) for path in paths for target in exclusions):
+                values.append(common.GEMMQuantMode.bfloat16)
+                continue
+            algos = {algo for target, algo in declarations if any(_projection_matches(path, target) for path in paths)}
+            algos.update(algo for path in paths for algo in exact.get(path, ()))
+            if len(algos) > 1 or any(algo not in modes for algo in algos):
+                raise ValueError(f"Unsupported mixed precision for Nemotron-H projection {suffix}: {algos}")
+            values.append(modes[next(iter(algos))] if algos else unlisted)
+        resolved[suffix] = dict(Counter(values))
+    return resolved
 
 
 @register_model("NEMOTRONH")
@@ -44,6 +135,14 @@ def create(cls, model_info: dict, model_config, backend_name: str) -> BaseModel:
             backend_name=backend_name,
         )
         # extra_params is NemotronHConfig with hybrid layer configuration.
+        # Non-latent MoE keeps the global GEMM mode until its calibration is qualified.
+        if not model_info["gemm_quant_mode_is_explicit"] and model_info["extra_params"].moe_latent_size:
+            model._projection_quant_modes = _projection_modes(
+                model_info.get("raw_config", {}),
+                model_info["extra_params"].hybrid_override_pattern,
+                model_config.gemm_quant_mode,
+                model_config.system,
+            )
         model.set_hybrid_config(model_info["extra_params"])
         return model
 
@@ -56,6 +155,7 @@ def __init__(self, topk: int, num_experts: int, moe_inter_size: int, *args, back
         self._num_experts = num_experts
         self._moe_inter_size = moe_inter_size
         self._hybrid_config: common.NemotronHConfig | None = None
+        self._projection_quant_modes: dict[str, dict[common.GEMMQuantMode, int]] = {}
         self._power_law_alpha = 1.01  # follow DeepSeek MoE
         # MTP (num_nextn_predict_layers > 0): scale generation iteration cost by
         # the small extra-layer factor (nextn+L)/L. Accepted-token progress is
@@ -84,6 +184,20 @@ def set_hybrid_config(self, hybrid_config: common.NemotronHConfig) -> None:
         self._build_context_ops()
         self._build_generation_ops()
 
+    def _projection_gemms(self, name: str, suffix: str, count, n: int, k: int, quant_mode, **kwargs) -> list:
+        """One GEMM per checkpoint precision of a projection, each priced for its own MoE layers.
+
+        Ops that share a name fold additively, so the buckets sum to the per-layer cost.
+        """
+        buckets = self._projection_quant_modes.get(suffix)
+        if not buckets:
+            return [ops.GEMM(name, count, n, k, quant_mode, **kwargs)]
+        total = sum(buckets.values())
+        return [
+            ops.GEMM(name, count if layers == total else count * layers / total, n, k, mode, **kwargs)
+            for mode, layers in buckets.items()
+        ]
+
     def _count_layer_types(self) -> dict[str, int]:
         """Count occurrences of each layer type in the pattern."""
         pattern = self._hybrid_config.hybrid_override_pattern
@@ -235,8 +349,9 @@ def _build_context_ops(self) -> None:
             # See TensorRT-LLM modeling_nemotron_h.py NemotronHMLP class
             self.context_ops.extend(
                 [
-                    ops.GEMM(
+                    *self._projection_gemms(
                         "context_shared_up_gemm",
+                        "shared_experts.up_proj",
                         count,
                         cfg.moe_shared_expert_intermediate_size // tp_size,
                         h,
@@ -249,8 +364,9 @@ def _build_context_ops(self) -> None:
                         cfg.moe_shared_expert_intermediate_size // tp_size,
                         0.8,
                     ),
-                    ops.GEMM(
+                    *self._projection_gemms(
                         "context_shared_down_gemm",
+                        "shared_experts.down_proj",
                         count,
                         h,
                         cfg.moe_shared_expert_intermediate_size // tp_size,
@@ -275,9 +391,10 @@ def _build_context_ops(self) -> None:
             # When latent-MoE is active, fc1_latent_proj (h -> latent) runs before
             # dispatch, so dispatch/experts/combine all communicate in latent dim.
             if cfg.moe_latent_size > 0:
-                self.context_ops.append(
-                    ops.GEMM(
+                self.context_ops.extend(
+                    self._projection_gemms(
                         "context_fc1_latent_proj_gemm",
+                        "fc1_latent_proj",
                         count,
                         cfg.moe_latent_size // tp_size,
                         h,
@@ -333,9 +450,10 @@ def _build_context_ops(self) -> None:
                 # fc2_latent_proj (latent -> h) runs after expert combine,
                 # before the post-MoE allreduce. Modeled as a row-parallel GEMM
                 # whose partial sums are aggregated by the following AllReduce.
-                self.context_ops.append(
-                    ops.GEMM(
+                self.context_ops.extend(
+                    self._projection_gemms(
                         "context_fc2_latent_proj_gemm",
+                        "fc2_latent_proj",
                         count,
                         h,
                         cfg.moe_latent_size // tp_size,
@@ -528,8 +646,9 @@ def _build_generation_ops(self) -> None:
             # See TensorRT-LLM modeling_nemotron_h.py NemotronHMLP class
             self.generation_ops.extend(
                 [
-                    ops.GEMM(
+                    *self._projection_gemms(
                         "generation_shared_up_gemm",
+                        "shared_experts.up_proj",
                         count,
                         cfg.moe_shared_expert_intermediate_size // tp_size,
                         h,
@@ -542,8 +661,9 @@ def _build_generation_ops(self) -> None:
                         cfg.moe_shared_expert_intermediate_size // tp_size,
                         0.8,
                     ),
-                    ops.GEMM(
+                    *self._projection_gemms(
                         "generation_shared_down_gemm",
+                        "shared_experts.down_proj",
                         count,
                         h,
                         cfg.moe_shared_expert_intermediate_size // tp_size,
@@ -568,9 +688,10 @@ def _build_generation_ops(self) -> None:
             # When latent-MoE is active, fc1_latent_proj (h -> latent) runs before
             # dispatch, so dispatch/experts/combine all communicate in latent dim.
             if cfg.moe_latent_size > 0:
-                self.generation_ops.append(
-                    ops.GEMM(
+                self.generation_ops.extend(
+                    self._projection_gemms(
                         "generation_fc1_latent_proj_gemm",
+                        "fc1_latent_proj",
                         count,
                         cfg.moe_latent_size // tp_size,
                         h,
@@ -626,9 +747,10 @@ def _build_generation_ops(self) -> None:
                 # fc2_latent_proj (latent -> h) runs after expert combine,
                 # before the post-MoE allreduce. Modeled as a row-parallel GEMM
                 # whose partial sums are aggregated by the following AllReduce.
-                self.generation_ops.append(
-                    ops.GEMM(
+                self.generation_ops.extend(
+                    self._projection_gemms(
                         "generation_fc2_latent_proj_gemm",
+                        "fc2_latent_proj",
                         count,
                         h,
                         cfg.moe_latent_size // tp_size,