Export the KV cache in the functional export API and share MLA's fused down-projection ranges

CHANGELOG.rst

@@ -8,6 +8,7 @@ Changelog
 
 *Quantization*
 
+- The functional quantized-weight export API now exports the KV cache: ``capture_kv_cache_export_state`` returns an attention module's ``k_scale`` / ``v_scale``, and ``build_hf_quantization_config`` takes the KV-cache formats from ``get_kv_cache_export_spec``.
 - On Linux x86-64 with Python 3.11 and later, the ONNX extra now selects ONNX Runtime 1.30.x and CUDA 13 packages by default. GPU calibration requires a compatible CUDA 13 runtime and driver; Python 3.10 and other platforms retain their existing dependency selections.
 - Add imatrix calibration (``algorithm: {method: imatrix}``) for the GGML formats with a weighted search: each weight is encoded once, its search weighted by the square root of its inputs' per-column ``E[x^2]`` (llama.cpp's imatrix statistic), and HF fused-MoE experts get one statistic per expert. A weight that already holds final codes is left alone, so ``[{method: gptq}, {method: imatrix}]`` runs GPTQ where it applies and imatrix on the rest; the built-in IQ1 and IQ2 recipes now do this, and new ``<format>_experts_only`` recipes apply imatrix to the routed MoE experts alone.
 - GPTQ now also quantizes HuggingFace fused MoE experts, each with its own Hessian from the tokens routed to it; they were previously skipped. The built-in IQ1 and IQ2 recipes therefore run GPTQ on the experts too, and large MoEs need layerwise calibration to hold one layer's expert Hessians at a time.
@@ -44,6 +45,8 @@ Changelog
 
 **Bug Fixes**
 
+- Calibrating a Megatron-Core MLA model now gives ``linear_q_down_proj`` and ``linear_kv_down_proj`` one per-tensor input and weight range, as vLLM and TensorRT-LLM fuse them into one GEMM with one scale per tensor. Re-calibrate existing MLA checkpoints to get the shared ranges.
+- ``export_mcore_gpt_to_hf`` and ``capture_kv_cache_export_state`` now write NVFP4 KV-cache ``k_scale`` / ``v_scale`` as amax / 448, like the HF exporter and as vLLM and TensorRT-LLM expect. They wrote amax / (6 * 448), which overflows in the attention kernels. Re-export Megatron checkpoints with an NVFP4 KV cache.
 - Preserve accumulated histogram counts when computing an entropy calibration threshold with NumPy histograms.
 - ``mtq.quantize`` and ``mtq.calibrate`` now give every module back the training mode it had. They used to leave every module in the mode of the model passed in, so a module in eval mode inside a model in training mode came back in training mode.
 - Diffusers ``export_hf_checkpoint`` now honors ``save_modelopt_state=False`` (the default) when HuggingFace checkpointing is enabled. Pass ``save_modelopt_state=True`` to retain the ModelOpt state file.

modelopt/torch/export/quant_utils.py

@@ -315,10 +315,13 @@ def get_kv_cache_scaling_factor(
     ):
         return [None, None]
 
-    scaling_factors = [
-        get_scaling_factor(getattr(self_attention_module, quantizer))
-        for quantizer in ("k_bmm_quantizer", "v_bmm_quantizer")
-    ]
+    # Both FP8 and NVFP4 cache scales are amax / 448, as the HF exporter writes them and
+    # vLLM / TensorRT-LLM read them.
+    scaling_factors = []
+    for quantizer_name in ("k_bmm_quantizer", "v_bmm_quantizer"):
+        quantizer = getattr(self_attention_module, quantizer_name)
+        amax = quantizer.export_amax() if quantizer.is_enabled else None
+        scaling_factors.append(None if amax is None else amax.float() / 448.0)
 
     # For FP8, we recommend default KV-cache scaling factor to be 1. The
     # asymmetric format applies this only to K; V remains NVFP4.

modelopt/torch/export/quantized_weight_export.py

@@ -37,6 +37,9 @@
 from ..quantization.nn import SequentialQuantizer, TensorQuantizer
 from .convert_hf_config import convert_hf_quant_config_format
 from .quant_format import (
+    KV_CACHE_FP8,
+    KV_CACHE_FP8_K_NVFP4_V,
+    KV_CACHE_NVFP4,
     QUANTIZATION_FP8,
     QUANTIZATION_FP8_PB_WO,
     QUANTIZATION_FP8_PC_PT,
@@ -51,15 +54,19 @@
 )
 from .quant_utils import (
     _get_quantization_from_quantizers,
+    get_kv_cache_dtype,
+    get_kv_cache_scaling_factor,
     get_scaling_factor,
     process_layer_quant_config,
     to_quantized_weight,
 )
 
 __all__ = [
     "build_hf_quantization_config",
+    "capture_kv_cache_export_state",
     "capture_quantized_weight_export_state",
     "export_quantized_weight_tensors",
+    "get_kv_cache_export_spec",
     "get_quantized_weight_export_spec",
     "merge_quantized_weight_export_states",
     "restore_quantized_weight_export_state",
@@ -83,6 +90,7 @@
         QUANTIZATION_MXFP8,
     }
 )
+_KV_CACHE_EXPORT_FORMATS = {KV_CACHE_FP8, KV_CACHE_FP8_K_NVFP4_V, KV_CACHE_NVFP4}
 _STATIC_INPUT_SCALE_FORMATS = {
     QUANTIZATION_FP8,
     QUANTIZATION_NVFP4,
@@ -126,6 +134,15 @@ class _QuantizedWeightExportSpec:
     input_quantizer_enabled: bool
 
 
+@dataclass(frozen=True)
+class _KVCacheExportState:
+    """Detached K and V cache scales of one attention module."""
+
+    quantization_format: str
+    k_scale: torch.Tensor
+    v_scale: torch.Tensor
+
+
 @dataclass(frozen=True)
 class _ExportStateTensorMetadata:
     name: str
@@ -497,6 +514,56 @@ def get_quantized_weight_export_spec(
     )
 
 
+def get_kv_cache_export_spec(attention: nn.Module) -> str | None:
+    """Return the KV-cache format of one attention module, or ``None`` when it is not quantized."""
+    quantizers = [getattr(attention, name, None) for name in ("k_bmm_quantizer", "v_bmm_quantizer")]
+    enabled = [quantizer is not None and quantizer.is_enabled for quantizer in quantizers]
+    if not any(enabled):
+        return None
+    if not all(enabled):
+        raise _UnsupportedQuantizedWeightExportFormatError(
+            "Functional export requires both or neither of the K and V cache quantizers"
+        )
+    # The captured state carries no K/V biases.
+    if any(
+        getattr(quantizer, attr, None) is not None
+        for quantizer in quantizers
+        for attr in ("bias", "bias_value")
+    ):
+        raise _UnsupportedQuantizedWeightExportFormatError(
+            "Functional export does not support affine KV caches"
+        )
+    quantization_format = get_kv_cache_dtype(attention)
+    if quantization_format not in _KV_CACHE_EXPORT_FORMATS:
+        raise _UnsupportedQuantizedWeightExportFormatError(
+            f"Functional export does not support the {quantization_format!r} KV cache"
+        )
+    return quantization_format
+
+
+def capture_kv_cache_export_state(
+    attention: nn.Module,
+    *,
+    clamp_fp8_scales: bool = True,
+    cpu: bool = True,
+) -> _KVCacheExportState | None:
+    """Capture the ``k_scale`` / ``v_scale`` of one attention module, or ``None`` when not quantized.
+
+    ``clamp_fp8_scales`` raises FP8 cache scales to at least 1.0, as the other exporters do.
+    """
+    quantization_format = get_kv_cache_export_spec(attention)
+    if quantization_format is None:
+        return None
+    k_scale, v_scale = get_kv_cache_scaling_factor(attention, clamp_fp8_scales=clamp_fp8_scales)
+    if k_scale is None or v_scale is None:
+        raise RuntimeError("Missing calibrated KV-cache scales")
+    if cpu:
+        k_scale, v_scale = k_scale.cpu(), v_scale.cpu()
+    return _KVCacheExportState(
+        quantization_format, k_scale.detach().clone(), v_scale.detach().clone()
+    )
+
+
 def split_quantized_weight_export_state(
     state: _QuantizedWeightExportState,
 ) -> tuple[object, tuple[torch.Tensor, ...]]:
@@ -844,8 +911,19 @@ def _quantized_layer_name(weight_name: str) -> str:
 def build_hf_quantization_config(
     named_states: Mapping[str, _QuantizedWeightExportState | _QuantizedWeightExportSpec | None]
     | Iterable[tuple[str, _QuantizedWeightExportState | _QuantizedWeightExportSpec | None]],
+    kv_cache_formats: Iterable[str | None] = (),
 ) -> dict[str, Any]:
-    """Build canonical HF configuration, rejecting unrepresentable activation settings."""
+    """Build canonical HF configuration, rejecting unrepresentable activation settings.
+
+    ``kv_cache_formats`` holds :func:`get_kv_cache_export_spec` of the attention modules; they must
+    agree, and ``None`` (an unquantized cache) is allowed alongside one format.
+    """
+    kv_cache_format_set = {fmt for fmt in kv_cache_formats if fmt is not None}
+    if len(kv_cache_format_set) > 1:
+        raise NotImplementedError(
+            "Cannot build canonical HF configuration for mixed KV-cache formats "
+            f"{sorted(kv_cache_format_set)}"
+        )
     layers: dict[str, tuple[str, int] | None] = {}
     for weight_name, state in dict(named_states).items():
         if state is not None:
@@ -875,4 +953,6 @@ def build_hf_quantization_config(
         "producer": {"name": "modelopt", "version": __version__},
         "quantization": process_layer_quant_config(layer_config),
     }
+    if kv_cache_format_set:
+        config["quantization"]["kv_cache_quant_algo"] = kv_cache_format_set.pop()
     return convert_hf_quant_config_format(config)

modelopt/torch/quantization/conversion.py

@@ -45,7 +45,7 @@
     TensorQuantizer,
 )
 from .utils import is_quantized, is_quantized_linear
-from .utils.shared_input import SharedWeightGlobalAmaxState
+from .utils.shared_input import MLA_DOWN_PROJ_PATTERN, SharedWeightGlobalAmaxState
 
 __all__ = [
     "preserve_quantizer_attributes_context",
@@ -115,11 +115,17 @@ def _restore_shared_quant_state_aliases(
     """Rebuild shared-state ties before checkpoint tensor values are loaded."""
     if not metadata.get("shared_quant_states"):
         return
-    # max / mse / local_hessian all carry ``shared_states`` (via _SharedStatesConfig) and use
-    # the same grouping; resolve the patterns that were in effect at save and rebuild the ties.
-    patterns = SharedWeightGlobalAmaxState.resolve_patterns(
-        shared_states=getattr(config, "shared_states", None)
-    )
+    patterns = metadata.get("shared_quant_patterns")
+    if patterns is None:
+        # Saved before the patterns were recorded, when the MLA down projections were not grouped:
+        # tying them now would load both members' ranges into one buffer.
+        patterns = [
+            pattern
+            for pattern in SharedWeightGlobalAmaxState.resolve_patterns(
+                shared_states=getattr(config, "shared_states", None)
+            )
+            if pattern != MLA_DOWN_PROJ_PATTERN
+        ]
     SharedWeightGlobalAmaxState.restore(model, patterns=patterns)
 
 
@@ -206,8 +212,12 @@ def update_quantize_metadata(
     metadata["quantizer_state"] = quantizer_state(model)
     if shared_state_metadata := SharedWeightGlobalAmaxState.metadata(model):
         metadata["shared_quant_states"] = shared_state_metadata
+        # Record the patterns calibration grouped with, so restore ties the same groups; ``config``
+        # need not be the calibration algorithm's (e.g. LoRA and AutoQuantize pass a default one).
+        metadata["shared_quant_patterns"] = SharedWeightGlobalAmaxState.grouping_patterns(model)
     else:
         metadata.pop("shared_quant_states", None)
+        metadata.pop("shared_quant_patterns", None)
 
 
 def _linear_attention_modules(model):

modelopt/torch/quantization/model_calib.py

@@ -100,9 +100,10 @@ def _is_calibrated_nvfp4_static(q) -> bool:
 
 
 def _collect_grouped_linears(model: nn.Module) -> list[list[nn.Module]]:
-    """Collect name-based sibling groups (Q/K/V, gate/up, w1/w3) of calibrated NVFP4-static linears."""
+    """Collect name-based sibling groups (Q/K/V, MLA down, gate/up, w1/w3) of calibrated NVFP4-static linears."""
     patterns: tuple[tuple[str, ...], ...] = (
         ("q_proj", "k_proj", "v_proj"),
+        ("linear_q_down_proj", "linear_kv_down_proj"),
         *list_all_possible("gate_up_pairs"),
     )
     groups: list[list[nn.Module]] = []
@@ -146,6 +147,41 @@ def _check_grouped_weight_global_amax_synced(model: nn.Module) -> None:
         )
 
 
+@torch.no_grad()
+def _sync_mla_fused_down_proj_amax(model: torch.nn.Module) -> None:
+    """Give Megatron-Core MLA's q_down / kv_down projections one per-tensor input and weight range each.
+
+    vLLM and TRT-LLM fuse them into one GEMM (``fused_qkv_a_proj``) with one scale per tensor.
+    Static NVFP4 weights share their global range through ``SHARED_PATTERNS`` instead.
+    """
+    for module in model.modules():
+        members = [
+            getattr(module, name, None) for name in ("linear_q_down_proj", "linear_kv_down_proj")
+        ]
+        if any(member is None for member in members):
+            continue
+        for quantizer_name in ("input_quantizer", "weight_quantizer"):
+            quantizers = []
+            for member in members:
+                quantizer = getattr(member, quantizer_name, None)
+                # W4A8 keeps its per-tensor FP8 range in the last quantizer.
+                if isinstance(quantizer, SequentialQuantizer):
+                    quantizer = quantizer[-1]
+                amax = getattr(quantizer, "_amax", None)
+                if (
+                    quantizer is None
+                    or not quantizer.is_enabled
+                    or amax is None
+                    or amax.numel() != 1
+                ):
+                    break
+                quantizers.append((quantizer, amax))
+            else:
+                shared = torch.stack([amax.float().reshape(()) for _, amax in quantizers]).max()
+                for quantizer, amax in quantizers:
+                    quantizer.amax = shared.to(amax).reshape(amax.shape)
+
+
 def _finalize_with_shared_state(model: nn.Module, weight_patterns: list[str]) -> None:
     """Finalize calibrated static quantizers and attached shared state.
 
@@ -155,6 +191,7 @@ def _finalize_with_shared_state(model: nn.Module, weight_patterns: list[str]) ->
     the default patterns, verifies the name groups were actually synced. Call once ``_amax``
     is final: single-process, or after the distributed amax sync.
     """
+    _sync_mla_fused_down_proj_amax(model)
     SharedWeightGlobalAmaxState.populate(model)
     promote_static_block_weight_quantizers(model)
     # Under the default patterns, verify the fusible name groups were actually synced.
@@ -833,6 +870,8 @@ def _mse_calibrate_weights(
 
                 pbar.update(1)
     pbar.close()
+    # The search refines each weight's range on its own; share the fused MLA pair's again.
+    _sync_mla_fused_down_proj_amax(model)
 
 
 class _LocalHessianAccumulator:

modelopt/torch/quantization/utils/shared_input.py

@@ -16,9 +16,10 @@
 """Shared quantization state for fusible sibling modules.
 
 Weight ``global_amax`` must be unified across modules that get **fused** at export
-(q/k/v -> qkv, gate/up -> gate_up) so they quantize with one per-tensor scale.
-:func:`find_shared_input_groups` discovers these groups by regex over module FQNs;
-``SHARED_PATTERNS`` covers the standard q/k/v, gate/up and per-expert w1/w3 names.
+(q/k/v -> qkv, gate/up -> gate_up, MLA q_down/kv_down -> fused_qkv_a_proj) so they quantize
+with one per-tensor scale. :func:`find_shared_input_groups` discovers these groups by regex over
+module FQNs; ``SHARED_PATTERNS`` covers the standard q/k/v, gate/up, per-expert w1/w3 and MCore
+MLA down-projection names.
 
 Discovery is name/pattern-based (not input-hook-based) on purpose: "shares an input
 tensor" is broader than "gets fused" — e.g. a ``shared_expert_gate`` reads the same
@@ -52,11 +53,15 @@
     "iter_shared_quant_states",
 ]
 
-# Groups export fuses: q/k/v -> qkv, gate/up (incl. Mixtral w1/w3) -> gate_up.
+# Megatron-Core MLA's q_down / kv_down projections, which deployment fuses into fused_qkv_a_proj.
+MLA_DOWN_PROJ_PATTERN = r"(?:(.*)\.)?(?:linear_q_down_proj|linear_kv_down_proj)"
+
+# Groups export fuses: q/k/v -> qkv, MLA q_down/kv_down, gate/up (incl. Mixtral w1/w3) -> gate_up.
 # Regexes are ``re.fullmatch``-ed against module FQNs; ``(?:(.*)\.)?`` captures the
 # immediate parent so MoE ``w1``/``w3`` groups are per-expert.
 SHARED_PATTERNS = (
     r"(?:(.*)\.)?(?:q_proj|k_proj|v_proj)",
+    MLA_DOWN_PROJ_PATTERN,
     r"(?:(.*)\.)?(?:gate_proj|up_proj)",
     r"(?:(.*)\.)?(?:w1|w3)",
 )
@@ -65,6 +70,7 @@
 # same ``experts``/``local_experts`` container. Kept unused until a state opts into it.
 SHARED_PATTERNS_ACROSS_EXPERTS = (
     r"(?:(.*)\.)?(?:q_proj|k_proj|v_proj)",
+    MLA_DOWN_PROJ_PATTERN,
     r"(?:(.*)\.)?(?:gate_proj|up_proj)",
     r"(?:(.*)\.)?(?:experts|local_experts)\.\d+\.(?:w1|w3)",
 )
@@ -103,6 +109,8 @@ def __init__(self) -> None:
         object.__setattr__(self, "_parent", None)
         object.__setattr__(self, "_members", ())
         object.__setattr__(self, "_hook_handles", [])
+        # Patterns that formed this group, recorded so restore groups modules the same way.
+        object.__setattr__(self, "patterns", ())
 
     def set_members(self, parent: nn.Module, members: Sequence[nn.Module]) -> None:
         """Set the owning parent and linked member modules."""
@@ -235,6 +243,7 @@ def attach(cls, model: nn.Module, patterns: Sequence[str] | None = None) -> int:
             target_quantizer_kind=cls.target_quantizer_kind,
         ):
             state = cls()
+            object.__setattr__(state, "patterns", tuple(patterns or ()))
             state.set_members(parent, members)
             created = _register_parent_shared_state(parent, state)
             n_created += int(created)
@@ -272,6 +281,14 @@ def restore(cls, model: nn.Module, patterns: Sequence[str] | None = None) -> Non
         for state in iter_shared_quant_states(model, cls):
             state.restore_from_members()
 
+    @classmethod
+    def grouping_patterns(cls, model: nn.Module) -> list[str]:
+        """Return the patterns that formed this state's groups in ``model``."""
+        patterns: list[str] = []
+        for state in iter_shared_quant_states(model, cls):
+            patterns += [pattern for pattern in state.patterns if pattern not in patterns]
+        return patterns
+
     @classmethod
     def metadata(cls, model: nn.Module) -> dict[str, bool]:
         """Return restore metadata for this state when present in ``model``."""