From 5cd5ad719be5993a9c86198cf238d0cc2aef04a6 Mon Sep 17 00:00:00 2001
From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
Date: Fri, 23 Oct 2026 12:57:01 -0700
Subject: [PATCH 1/6] Export the KV cache functionally and share MLA's fused
 down-projection ranges

Add get_kv_cache_export_spec / capture_kv_cache_export_state to the functional
quantized-weight export API, and a kv_cache_formats argument to
build_hf_quantization_config, so Megatron-Bridge can export k_scale / v_scale
and kv_cache_quant_algo the way export_mcore_gpt_to_hf does.

vLLM and TensorRT-LLM fuse MLA's q_down and kv_down projections into one GEMM
with one scale per tensor. Calibration now gives the Megatron-Core
linear_q_down_proj / linear_kv_down_proj pair one per-tensor input and weight
range, and SHARED_PATTERNS groups them for static NVFP4.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
---
 CHANGELOG.rst                                 |  2 +
 .../torch/export/quantized_weight_export.py   | 73 ++++++++++++++++++-
 modelopt/torch/quantization/model_calib.py    |  9 ++-
 .../torch/quantization/plugins/megatron.py    | 35 +++++++++
 .../torch/quantization/utils/shared_input.py  |  9 ++-
 .../export/test_unified_export_megatron.py    | 29 ++++++++
 .../quantization/plugins/test_megatron.py     | 34 +++++++++
 .../export/test_quantized_weight_state.py     | 71 ++++++++++++++++++
 8 files changed, 257 insertions(+), 5 deletions(-)

diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index 1b4c6cbe2be..5f4f697de66 100755
--- a/CHANGELOG.rst
+++ b/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.
@@ -42,6 +43,7 @@ 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.
 - The default PTQ recipes now leave multi-token-prediction (MTP) layers unquantized in Megatron vision-language models too (``language_model.mtp.*``), as they already did at the model root.
 - Fix GPT-OSS draft training with pre-normalization teacher hidden states by enabling the target model's final normalization.
 - Fix vanilla DSpark exports to use the Qwen3 DSpark architecture, head names and configuration expected by vLLM. Re-export existing training checkpoints to use this format.
diff --git a/modelopt/torch/export/quantized_weight_export.py b/modelopt/torch/export/quantized_weight_export.py
index 79a29e7c400..cd1673f0735 100644
--- a/modelopt/torch/export/quantized_weight_export.py
+++ b/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,6 +54,8 @@
 )
 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,
@@ -58,8 +63,10 @@
 
 __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,47 @@ 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"
+        )
+    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 +902,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 +944,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)
diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py
index 5e7e62c4189..2edbce2b55c 100644
--- a/modelopt/torch/quantization/model_calib.py
+++ b/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]] = []
@@ -155,6 +156,12 @@ 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.
     """
+    try:
+        from .plugins.megatron import _sync_mla_fused_down_proj_amax
+    except ImportError:
+        pass
+    else:
+        _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.
diff --git a/modelopt/torch/quantization/plugins/megatron.py b/modelopt/torch/quantization/plugins/megatron.py
index a96e7fd2a33..1bc5c94c214 100644
--- a/modelopt/torch/quantization/plugins/megatron.py
+++ b/modelopt/torch/quantization/plugins/megatron.py
@@ -92,6 +92,41 @@
 __all__ = []
 
 
+@torch.no_grad()
+def _sync_mla_fused_down_proj_amax(model: torch.nn.Module) -> None:
+    """Give 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 _check_nvfp4_static_tp_supported(model: torch.nn.Module) -> None:
     """Raise if using NVFP4-static weight quantization with TP>1.
 
diff --git a/modelopt/torch/quantization/utils/shared_input.py b/modelopt/torch/quantization/utils/shared_input.py
index d80de2fcf9d..d0977f9484c 100644
--- a/modelopt/torch/quantization/utils/shared_input.py
+++ b/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
@@ -57,6 +58,7 @@
 # immediate parent so MoE ``w1``/``w3`` groups are per-expert.
 SHARED_PATTERNS = (
     r"(?:(.*)\.)?(?:q_proj|k_proj|v_proj)",
+    r"(?:(.*)\.)?(?:linear_q_down_proj|linear_kv_down_proj)",
     r"(?:(.*)\.)?(?:gate_proj|up_proj)",
     r"(?:(.*)\.)?(?:w1|w3)",
 )
@@ -65,6 +67,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)",
+    r"(?:(.*)\.)?(?:linear_q_down_proj|linear_kv_down_proj)",
     r"(?:(.*)\.)?(?:gate_proj|up_proj)",
     r"(?:(.*)\.)?(?:experts|local_experts)\.\d+\.(?:w1|w3)",
 )
diff --git a/tests/gpu_megatron/torch/export/test_unified_export_megatron.py b/tests/gpu_megatron/torch/export/test_unified_export_megatron.py
index 043970ef42a..935f33e2952 100644
--- a/tests/gpu_megatron/torch/export/test_unified_export_megatron.py
+++ b/tests/gpu_megatron/torch/export/test_unified_export_megatron.py
@@ -14,6 +14,7 @@
 # limitations under the License.
 
 import json
+import re
 from copy import deepcopy
 from functools import partial
 from pathlib import Path
@@ -26,6 +27,7 @@
 from _test_utils.torch.export.unified_checkpoint import (
     assert_exported_checkpoint_matches,
     assert_per_expert_experts_complete,
+    load_safetensors_dir,
 )
 from _test_utils.torch.megatron.models import get_mcore_gpt_model, get_mcore_hybrid_model
 from _test_utils.torch.megatron.utils import get_forward
@@ -49,6 +51,11 @@
 from modelopt.torch.export import KV_CACHE_FP8, export_mcore_gpt_to_hf, import_mcore_gpt_from_hf
 from modelopt.torch.export.plugins.mcore_common import all_mcore_hf_export_mapping
 from modelopt.torch.export.quant_format import GGML_FORMATS
+from modelopt.torch.export.quantized_weight_export import (
+    build_hf_quantization_config,
+    capture_kv_cache_export_state,
+    get_kv_cache_export_spec,
+)
 from modelopt.torch.export.unified_export_megatron import GPTModelExporter
 from modelopt.torch.quantization.config import QuantizerAttributeConfig
 from modelopt.torch.quantization.nn import TensorQuantizer
@@ -501,6 +508,26 @@ def _verify_exported_metadata(export_dir: Path, model_type, quant_config, extra_
         assert generation_config["top_p"] == 0.5
 
 
+def _verify_functional_kv_cache_export(model, export_dir: Path):
+    """The functional KV-cache export writes the same scales and config as this exporter."""
+    attentions = [m for m in model.modules() if hasattr(m, "k_bmm_quantizer")]
+    states = [capture_kv_cache_export_state(attention) for attention in attentions]
+    exported = load_safetensors_dir(export_dir)
+    for kind in ("k_scale", "v_scale"):
+        names = sorted(
+            (name for name in exported if name.endswith(kind)),
+            key=lambda name: int(re.search(r"layers\.(\d+)\.", name).group(1)),
+        )
+        assert len(names) == len(states)
+        for name, state in zip(names, states):
+            torch.testing.assert_close(exported[name], getattr(state, kind).to(exported[name]))
+    config = build_hf_quantization_config({}, [get_kv_cache_export_spec(a) for a in attentions])
+    quantization_config = json.loads((export_dir / "config.json").read_text())[
+        "quantization_config"
+    ]
+    assert config["kv_cache_scheme"] == quantization_config["kv_cache_scheme"]
+
+
 def _test_unified_export_megatron(
     tmp_path,
     model_type,
@@ -661,6 +688,8 @@ def _test_unified_export_megatron(
 
     if quant_config:
         _verify_model_quant_config(tmp_export_dir, quant_config, kv_cache_quant_cfg)
+    if kv_cache_quant_cfg and rank == 0:
+        _verify_functional_kv_cache_export(model, tmp_export_dir)
 
     if rank == 0:
         _verify_exported_metadata(tmp_export_dir, model_type, quant_config, extra_module)
diff --git a/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py b/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py
index 9e2a7be8a0b..87e2265102c 100644
--- a/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py
+++ b/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py
@@ -2646,3 +2646,37 @@ def unrecognised(self, prefix="", sharded_offsets=(), metadata=None):
     )
     def test_extra_state_has_data(self, entry, expected):
         assert _output_layer_extra_state_has_data(entry) is expected
+
+
+@pytest.mark.parametrize(
+    "quant_cfg", ["FP8_DEFAULT_CFG", "NVFP4_DEFAULT_CFG", "NVFP4_W4A4_WEIGHT_MSE_FP8_SWEEP_CFG"]
+)
+def test_mla_fused_down_projections_share_ranges(distributed_setup_size_1, quant_cfg):
+    """vLLM and TRT-LLM fuse MLA's q_down and kv_down into one GEMM with one scale per tensor."""
+    initialize_for_megatron(tensor_model_parallel_size=1, seed=SEED)
+    model = get_mcore_gpt_model(
+        hidden_size=128,
+        num_attention_heads=2,
+        vocab_size=128,
+        max_sequence_length=32,
+        transformer_impl="transformer_engine",
+        multi_latent_attention=True,
+    ).cuda()
+    attentions = [layer.self_attention for layer in model.decoder.layers]
+    with torch.no_grad():
+        for attention in attentions:
+            attention.linear_kv_down_proj.weight.mul_(4)  # give the two weights different ranges
+
+    model = mtq.quantize(model, copy.deepcopy(getattr(mtq, quant_cfg)), get_forward(model))
+
+    for attention in attentions:
+        q_down, kv_down = attention.linear_q_down_proj, attention.linear_kv_down_proj
+        assert torch.equal(q_down.input_quantizer.amax, kv_down.input_quantizer.amax)
+        if q_down.weight_quantizer.is_nvfp4_static:
+            assert torch.equal(
+                q_down.weight_quantizer.global_amax, kv_down.weight_quantizer.global_amax
+            )
+        else:
+            assert torch.equal(q_down.weight_quantizer.amax, kv_down.weight_quantizer.amax)
+    # Clean up since this is not a spawned process
+    destroy_model_parallel()
diff --git a/tests/unit/torch/export/test_quantized_weight_state.py b/tests/unit/torch/export/test_quantized_weight_state.py
index 682bc0eb229..0cda6bb396b 100644
--- a/tests/unit/torch/export/test_quantized_weight_state.py
+++ b/tests/unit/torch/export/test_quantized_weight_state.py
@@ -24,8 +24,10 @@
 from modelopt.torch.export.convert_hf_config import convert_hf_quant_config_format
 from modelopt.torch.export.quantized_weight_export import (
     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,
@@ -696,3 +698,72 @@ def test_mixed_config_groups_moe_experts_by_projection_family():
         "model.layers.0.mlp.experts.down_proj",
     }
     assert "lm_head" in config["ignore"]
+
+
+_FP8_KV = {"num_bits": (4, 3)}
+_NVFP4_KV = {"num_bits": (2, 1), "block_sizes": {-1: 16, "type": "dynamic", "scale_bits": (4, 3)}}
+
+
+def _kv_attention(k_cfg, v_cfg, k_amax=896.0, v_amax=224.0) -> nn.Module:
+    attention = nn.Module()
+    for name, cfg, amax in (("k_bmm_quantizer", k_cfg, k_amax), ("v_bmm_quantizer", v_cfg, v_amax)):
+        quantizer = TensorQuantizer(QuantizerAttributeConfig(**cfg))
+        quantizer.amax = torch.tensor(amax)
+        setattr(attention, name, quantizer)
+    return attention
+
+
+@pytest.mark.parametrize(
+    ("k_cfg", "v_cfg", "quantization_format", "k_scale", "v_scale"),
+    [
+        # FP8 scales are amax / 448, raised to at least 1.0 like the other exporters do.
+        (_FP8_KV, _FP8_KV, "FP8", 2.0, 1.0),
+        (_NVFP4_KV, _NVFP4_KV, "NVFP4", 896 / (6 * 448), 224 / (6 * 448)),
+        (_FP8_KV, _NVFP4_KV, "FP8_K_NVFP4_V", 2.0, 224 / (6 * 448)),
+    ],
+)
+def test_kv_cache_capture(k_cfg, v_cfg, quantization_format, k_scale, v_scale):
+    attention = _kv_attention(k_cfg, v_cfg)
+
+    state = capture_kv_cache_export_state(attention)
+
+    assert get_kv_cache_export_spec(attention) == state.quantization_format == quantization_format
+    torch.testing.assert_close(state.k_scale.squeeze(), torch.tensor(k_scale))
+    torch.testing.assert_close(state.v_scale.squeeze(), torch.tensor(v_scale))
+    assert not state.k_scale.requires_grad and state.k_scale.device.type == "cpu"
+    if quantization_format == "FP8":
+        unclamped = capture_kv_cache_export_state(attention, clamp_fp8_scales=False)
+        assert unclamped.v_scale.item() == 0.5
+
+
+def test_kv_cache_spec_rejects_what_it_cannot_export():
+    attention = _kv_attention(_FP8_KV, _FP8_KV)
+    attention.k_bmm_quantizer.disable()
+    attention.v_bmm_quantizer.disable()
+    assert get_kv_cache_export_spec(attention) is None
+    assert capture_kv_cache_export_state(attention) is None
+    assert get_kv_cache_export_spec(nn.Linear(2, 2)) is None
+
+    attention.k_bmm_quantizer.enable()
+    with pytest.raises(NotImplementedError, match="both or neither"):
+        get_kv_cache_export_spec(attention)
+
+    attention = _kv_attention({"num_bits": 8}, {"num_bits": 8})
+    with pytest.raises(NotImplementedError, match="'INT8' KV cache"):
+        get_kv_cache_export_spec(attention)
+
+
+@pytest.mark.parametrize(
+    "weight_spec", [None, weight_export._QuantizedWeightExportSpec("fp8", 0, True)]
+)
+def test_hf_config_carries_the_kv_cache_format(weight_spec):
+    named_specs = {"model.layers.0.proj.weight": weight_spec}
+
+    config = build_hf_quantization_config(named_specs, kv_cache_formats=["FP8", None, "FP8"])
+
+    assert config["kv_cache_scheme"] == {"dynamic": False, "num_bits": 8, "type": "float"}
+    assert "kv_cache_scheme" not in build_hf_quantization_config(named_specs)
+    config = build_hf_quantization_config(named_specs, kv_cache_formats=["FP8_K_NVFP4_V"])
+    assert config["kv_cache_quant_algo"] == "FP8_K_NVFP4_V"
+    with pytest.raises(NotImplementedError, match="mixed KV-cache formats"):
+        build_hf_quantization_config(named_specs, kv_cache_formats=["FP8", "NVFP4"])

From 6c8f62c3a7b80e341c615f4a72f1f55084c8a299 Mon Sep 17 00:00:00 2001
From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
Date: Fri, 23 Oct 2026 13:21:16 -0700
Subject: [PATCH 2/6] Make the new MLA and KV-cache tests fail loudly

Tear down Megatron in a finally block so one failing MLA case does not fail the
next ones on re-initialization, and require the KV export check to find
quantized attention modules.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
---
 .../export/test_unified_export_megatron.py    |  1 +
 .../quantization/plugins/test_megatron.py     | 49 ++++++++++---------
 2 files changed, 26 insertions(+), 24 deletions(-)

diff --git a/tests/gpu_megatron/torch/export/test_unified_export_megatron.py b/tests/gpu_megatron/torch/export/test_unified_export_megatron.py
index 935f33e2952..c70e5bb432f 100644
--- a/tests/gpu_megatron/torch/export/test_unified_export_megatron.py
+++ b/tests/gpu_megatron/torch/export/test_unified_export_megatron.py
@@ -512,6 +512,7 @@ def _verify_functional_kv_cache_export(model, export_dir: Path):
     """The functional KV-cache export writes the same scales and config as this exporter."""
     attentions = [m for m in model.modules() if hasattr(m, "k_bmm_quantizer")]
     states = [capture_kv_cache_export_state(attention) for attention in attentions]
+    assert states and None not in states
     exported = load_safetensors_dir(export_dir)
     for kind in ("k_scale", "v_scale"):
         names = sorted(
diff --git a/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py b/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py
index 87e2265102c..ed85ea2a717 100644
--- a/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py
+++ b/tests/gpu_megatron/torch/quantization/plugins/test_megatron.py
@@ -2654,29 +2654,30 @@ def test_extra_state_has_data(self, entry, expected):
 def test_mla_fused_down_projections_share_ranges(distributed_setup_size_1, quant_cfg):
     """vLLM and TRT-LLM fuse MLA's q_down and kv_down into one GEMM with one scale per tensor."""
     initialize_for_megatron(tensor_model_parallel_size=1, seed=SEED)
-    model = get_mcore_gpt_model(
-        hidden_size=128,
-        num_attention_heads=2,
-        vocab_size=128,
-        max_sequence_length=32,
-        transformer_impl="transformer_engine",
-        multi_latent_attention=True,
-    ).cuda()
-    attentions = [layer.self_attention for layer in model.decoder.layers]
-    with torch.no_grad():
-        for attention in attentions:
-            attention.linear_kv_down_proj.weight.mul_(4)  # give the two weights different ranges
+    try:
+        model = get_mcore_gpt_model(
+            hidden_size=128,
+            num_attention_heads=2,
+            vocab_size=128,
+            max_sequence_length=32,
+            transformer_impl="transformer_engine",
+            multi_latent_attention=True,
+        ).cuda()
+        attentions = [layer.self_attention for layer in model.decoder.layers]
+        with torch.no_grad():  # give the two weights different ranges
+            for attention in attentions:
+                attention.linear_kv_down_proj.weight.mul_(4)
 
-    model = mtq.quantize(model, copy.deepcopy(getattr(mtq, quant_cfg)), get_forward(model))
+        model = mtq.quantize(model, copy.deepcopy(getattr(mtq, quant_cfg)), get_forward(model))
 
-    for attention in attentions:
-        q_down, kv_down = attention.linear_q_down_proj, attention.linear_kv_down_proj
-        assert torch.equal(q_down.input_quantizer.amax, kv_down.input_quantizer.amax)
-        if q_down.weight_quantizer.is_nvfp4_static:
-            assert torch.equal(
-                q_down.weight_quantizer.global_amax, kv_down.weight_quantizer.global_amax
-            )
-        else:
-            assert torch.equal(q_down.weight_quantizer.amax, kv_down.weight_quantizer.amax)
-    # Clean up since this is not a spawned process
-    destroy_model_parallel()
+        for attention in attentions:
+            q_down, kv_down = attention.linear_q_down_proj, attention.linear_kv_down_proj
+            assert torch.equal(q_down.input_quantizer.amax, kv_down.input_quantizer.amax)
+            if q_down.weight_quantizer.is_nvfp4_static:
+                assert torch.equal(
+                    q_down.weight_quantizer.global_amax, kv_down.weight_quantizer.global_amax
+                )
+            else:
+                assert torch.equal(q_down.weight_quantizer.amax, kv_down.weight_quantizer.amax)
+    finally:
+        destroy_model_parallel()  # not a spawned process

From 7b0a424926e5f566fd08e06e2127640715098aa3 Mon Sep 17 00:00:00 2001
From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
Date: Fri, 23 Oct 2026 13:29:01 -0700
Subject: [PATCH 3/6] Share MLA's down-projection ranges after the MSE weight
 search too

The MSE and local-Hessian weight searches refine each per-tensor weight range on
its own, which split the q_down / kv_down pair again; share it after the search
as well. The sync moves into model_calib (it only uses attribute names), with
CPU tests for the per-tensor, MSE, W4A8 and disabled-member cases.

Record the shared-group patterns in the quantize metadata so restore groups
modules as calibration did: checkpoints saved before the MLA pair was grouped
restore with each member's own range instead of tying both into one buffer.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
---
 modelopt/torch/quantization/conversion.py     | 24 +++--
 modelopt/torch/quantization/model_calib.py    | 44 +++++++--
 .../torch/quantization/plugins/megatron.py    | 35 -------
 .../torch/quantization/utils/shared_input.py  |  9 +-
 .../torch/quantization/test_shared_input.py   | 97 ++++++++++++++++++-
 5 files changed, 157 insertions(+), 52 deletions(-)

diff --git a/modelopt/torch/quantization/conversion.py b/modelopt/torch/quantization/conversion.py
index 78a6e03a16b..fbe373878b0 100644
--- a/modelopt/torch/quantization/conversion.py
+++ b/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,14 @@ def update_quantize_metadata(
     metadata["quantizer_state"] = quantizer_state(model)
     if shared_state_metadata := SharedWeightGlobalAmaxState.metadata(model):
         metadata["shared_quant_states"] = shared_state_metadata
+        # max / mse / local_hessian all carry ``shared_states`` (via _SharedStatesConfig); restore
+        # rebuilds the ties with the patterns in effect now.
+        metadata["shared_quant_patterns"] = SharedWeightGlobalAmaxState.resolve_patterns(
+            shared_states=getattr(config, "shared_states", None)
+        )
     else:
         metadata.pop("shared_quant_states", None)
+        metadata.pop("shared_quant_patterns", None)
 
 
 def _linear_attention_modules(model):
diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py
index 2edbce2b55c..915f631d0c7 100644
--- a/modelopt/torch/quantization/model_calib.py
+++ b/modelopt/torch/quantization/model_calib.py
@@ -147,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.
 
@@ -156,12 +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.
     """
-    try:
-        from .plugins.megatron import _sync_mla_fused_down_proj_amax
-    except ImportError:
-        pass
-    else:
-        _sync_mla_fused_down_proj_amax(model)
+    _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.
@@ -839,6 +869,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:
diff --git a/modelopt/torch/quantization/plugins/megatron.py b/modelopt/torch/quantization/plugins/megatron.py
index 1bc5c94c214..a96e7fd2a33 100644
--- a/modelopt/torch/quantization/plugins/megatron.py
+++ b/modelopt/torch/quantization/plugins/megatron.py
@@ -92,41 +92,6 @@
 __all__ = []
 
 
-@torch.no_grad()
-def _sync_mla_fused_down_proj_amax(model: torch.nn.Module) -> None:
-    """Give 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 _check_nvfp4_static_tp_supported(model: torch.nn.Module) -> None:
     """Raise if using NVFP4-static weight quantization with TP>1.
 
diff --git a/modelopt/torch/quantization/utils/shared_input.py b/modelopt/torch/quantization/utils/shared_input.py
index d0977f9484c..c13bbcdb6cc 100644
--- a/modelopt/torch/quantization/utils/shared_input.py
+++ b/modelopt/torch/quantization/utils/shared_input.py
@@ -53,12 +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)",
-    r"(?:(.*)\.)?(?:linear_q_down_proj|linear_kv_down_proj)",
+    MLA_DOWN_PROJ_PATTERN,
     r"(?:(.*)\.)?(?:gate_proj|up_proj)",
     r"(?:(.*)\.)?(?:w1|w3)",
 )
@@ -67,7 +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)",
-    r"(?:(.*)\.)?(?:linear_q_down_proj|linear_kv_down_proj)",
+    MLA_DOWN_PROJ_PATTERN,
     r"(?:(.*)\.)?(?:gate_proj|up_proj)",
     r"(?:(.*)\.)?(?:experts|local_experts)\.\d+\.(?:w1|w3)",
 )
diff --git a/tests/unit/torch/quantization/test_shared_input.py b/tests/unit/torch/quantization/test_shared_input.py
index 63933a819ff..81cb893cefa 100644
--- a/tests/unit/torch/quantization/test_shared_input.py
+++ b/tests/unit/torch/quantization/test_shared_input.py
@@ -27,8 +27,12 @@
 import modelopt.torch.opt as mto
 import modelopt.torch.quantization as mtq
 from modelopt.torch.quantization.config import MaxCalibConfig, QuantizerAttributeConfig
-from modelopt.torch.quantization.model_calib import max_calibrate
-from modelopt.torch.quantization.nn import NVFP4StaticQuantizer, TensorQuantizer
+from modelopt.torch.quantization.model_calib import _sync_mla_fused_down_proj_amax, max_calibrate
+from modelopt.torch.quantization.nn import (
+    NVFP4StaticQuantizer,
+    SequentialQuantizer,
+    TensorQuantizer,
+)
 from modelopt.torch.quantization.utils import (
     SHARED_PATTERNS,
     SharedWeightGlobalAmaxState,
@@ -507,3 +511,92 @@ def test_config_rejects_invalid_shared_states(self):
             MaxCalibConfig(shared_states={"weight_global_amax": {"pattern": [r".*"]}})
         with pytest.raises(ValueError, match="invalid regex"):
             MaxCalibConfig(shared_states={"weight_global_amax": {"patterns": ["("]}})
+
+
+class _DummyMLA(nn.Module):
+    """Megatron-Core MLA's q_down / kv_down projections, which deployment fuses into one GEMM."""
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.linear_q_down_proj = nn.Linear(32, 32, bias=False)
+        self.linear_kv_down_proj = nn.Linear(32, 32, bias=False)
+        with torch.no_grad():
+            self.linear_kv_down_proj.weight.mul_(4)  # give the two weights different ranges
+
+    def forward(self, x):
+        return self.linear_q_down_proj(x) + self.linear_kv_down_proj(x)
+
+
+_PER_TENSOR_INT8 = {"num_bits": 8, "axis": None}
+_STATIC_NVFP4 = {
+    "num_bits": (2, 1),
+    "block_sizes": {-1: NVFP4_BLOCK, "type": "static", "scale_bits": (4, 3)},
+}
+
+
+def _quantize_mla(weight_cfg, algorithm, input_cfg=None):
+    quant_cfg = [
+        {"enable": False, "quantizer_name": "*"},
+        {"cfg": weight_cfg, "quantizer_name": "*weight_quantizer"},
+    ]
+    if input_cfg is not None:
+        quant_cfg.append({"cfg": input_cfg, "quantizer_name": "*input_quantizer"})
+    mla = _DummyMLA()
+    mtq.quantize(
+        mla, {"quant_cfg": quant_cfg, "algorithm": algorithm}, lambda m: m(torch.randn(4, 32))
+    )
+    return mla
+
+
+@pytest.mark.parametrize("algorithm", ["max", {"method": "mse"}])
+def test_mla_down_projections_share_per_tensor_ranges(algorithm):
+    """The MSE weight search refines each range on its own, so the pair is shared again after it."""
+    mla = _quantize_mla(_PER_TENSOR_INT8, algorithm, input_cfg=_PER_TENSOR_INT8)
+    q_down, kv_down = mla.linear_q_down_proj, mla.linear_kv_down_proj
+
+    assert torch.equal(q_down.input_quantizer.amax, kv_down.input_quantizer.amax)
+    assert torch.equal(q_down.weight_quantizer.amax, kv_down.weight_quantizer.amax)
+    assert kv_down.weight_quantizer.amax > q_down.weight.abs().max()
+
+
+def test_mla_down_projections_share_the_fp8_stage_of_sequential_weight_quantizers():
+    mla = _DummyMLA()
+    mtq.replace_quant_module(mla)
+    first_stages = []
+    for scale, proj in enumerate((mla.linear_q_down_proj, mla.linear_kv_down_proj), start=1):
+        proj.weight_quantizer = SequentialQuantizer(TensorQuantizer(), TensorQuantizer())
+        proj.weight_quantizer[0].amax = torch.full((32, 1), float(scale))
+        proj.weight_quantizer[1].amax = torch.tensor(float(scale))
+        proj.input_quantizer.amax = torch.tensor(float(scale))
+        first_stages.append(proj.weight_quantizer[0].amax.clone())
+    mla.linear_kv_down_proj.input_quantizer.disable()
+
+    _sync_mla_fused_down_proj_amax(mla)
+
+    q_down, kv_down = mla.linear_q_down_proj, mla.linear_kv_down_proj
+    assert q_down.weight_quantizer[1].amax == kv_down.weight_quantizer[1].amax == 2.0
+    for proj, first_stage in zip((q_down, kv_down), first_stages):
+        assert torch.equal(proj.weight_quantizer[0].amax, first_stage)  # per-channel: left alone
+    assert q_down.input_quantizer.amax == 1.0  # the other side is not quantized
+
+
+def test_mla_down_projections_of_older_checkpoints_restore_ungrouped(tmp_path):
+    """A checkpoint saved before the MLA pair was grouped keeps each member's own range."""
+    mla = _quantize_mla(_STATIC_NVFP4, "max")
+    shared = mla.linear_q_down_proj.weight_quantizer.global_amax.clone()
+    assert torch.equal(shared, mla.linear_kv_down_proj.weight_quantizer.global_amax)
+    path = tmp_path / "model.pth"
+    mto.save(mla, path)
+
+    restored = _DummyMLA()
+    mto.restore(restored, path)
+    _single_global_amax_state(restored)
+
+    # Older checkpoints did not record the patterns; restore must not tie the pair.
+    checkpoint = torch.load(path, weights_only=False)
+    for _, mode_state in checkpoint["modelopt_state"]["modelopt_state_dict"]:
+        mode_state["metadata"].pop("shared_quant_patterns", None)
+    torch.save(checkpoint, path)
+    restored = _DummyMLA()
+    mto.restore(restored, path)
+    assert not hasattr(restored, "_shared_quant_states")

From 6172cf14b3b6bf0a368427bab927cf104c89db1c Mon Sep 17 00:00:00 2001
From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
Date: Fri, 23 Oct 2026 23:45:44 -0700
Subject: [PATCH 4/6] Record the grouping patterns calibration used, not the
 config's

update_quantize_metadata resolved the shared-state patterns from the config it
was given, but LoRA and AutoQuantize pass a default QuantizeConfig, which would
record the default patterns for a model calibrated with custom ones. Each shared
state now keeps the patterns that formed it, and the metadata records those.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
---
 modelopt/torch/quantization/conversion.py     |  8 ++---
 .../torch/quantization/utils/shared_input.py  | 11 ++++++
 .../torch/quantization/test_shared_input.py   | 34 ++++++++++++++++++-
 3 files changed, 47 insertions(+), 6 deletions(-)

diff --git a/modelopt/torch/quantization/conversion.py b/modelopt/torch/quantization/conversion.py
index fbe373878b0..d82f916b248 100644
--- a/modelopt/torch/quantization/conversion.py
+++ b/modelopt/torch/quantization/conversion.py
@@ -212,11 +212,9 @@ def update_quantize_metadata(
     metadata["quantizer_state"] = quantizer_state(model)
     if shared_state_metadata := SharedWeightGlobalAmaxState.metadata(model):
         metadata["shared_quant_states"] = shared_state_metadata
-        # max / mse / local_hessian all carry ``shared_states`` (via _SharedStatesConfig); restore
-        # rebuilds the ties with the patterns in effect now.
-        metadata["shared_quant_patterns"] = SharedWeightGlobalAmaxState.resolve_patterns(
-            shared_states=getattr(config, "shared_states", None)
-        )
+        # 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)
diff --git a/modelopt/torch/quantization/utils/shared_input.py b/modelopt/torch/quantization/utils/shared_input.py
index c13bbcdb6cc..1d7e1c2a139 100644
--- a/modelopt/torch/quantization/utils/shared_input.py
+++ b/modelopt/torch/quantization/utils/shared_input.py
@@ -109,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."""
@@ -241,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)
@@ -278,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``."""
diff --git a/tests/unit/torch/quantization/test_shared_input.py b/tests/unit/torch/quantization/test_shared_input.py
index 81cb893cefa..9ea01e64dbe 100644
--- a/tests/unit/torch/quantization/test_shared_input.py
+++ b/tests/unit/torch/quantization/test_shared_input.py
@@ -26,7 +26,12 @@
 
 import modelopt.torch.opt as mto
 import modelopt.torch.quantization as mtq
-from modelopt.torch.quantization.config import MaxCalibConfig, QuantizerAttributeConfig
+from modelopt.torch.quantization.config import (
+    MaxCalibConfig,
+    QuantizeConfig,
+    QuantizerAttributeConfig,
+)
+from modelopt.torch.quantization.conversion import update_quantize_metadata
 from modelopt.torch.quantization.model_calib import _sync_mla_fused_down_proj_amax, max_calibrate
 from modelopt.torch.quantization.nn import (
     NVFP4StaticQuantizer,
@@ -600,3 +605,30 @@ def test_mla_down_projections_of_older_checkpoints_restore_ungrouped(tmp_path):
     restored = _DummyMLA()
     mto.restore(restored, path)
     assert not hasattr(restored, "_shared_quant_states")
+
+
+def test_custom_grouping_patterns_survive_save_and_restore(tmp_path):
+    """Restore groups as calibration did, even where the metadata is written with a default config."""
+    q_and_k = r"(?:(.*)\.)?(?:q_proj|k_proj)"
+    cfg = {
+        "quant_cfg": [
+            {"enable": False, "quantizer_name": "*"},
+            {"cfg": _STATIC_NVFP4, "quantizer_name": "*weight_quantizer"},
+        ],
+        "algorithm": {
+            "method": "max",
+            "shared_states": {"weight_global_amax": {"patterns": [q_and_k]}},
+        },
+    }
+    attn = _DummyAttention()
+    mtq.quantize(attn, cfg, lambda m: m(torch.randn(2, 32)))
+    assert _single_global_amax_state(attn).members == (attn.q_proj, attn.k_proj)
+    metadata = {}
+    update_quantize_metadata(attn, QuantizeConfig(), metadata)  # as LoRA and AutoQuantize do
+    assert metadata["shared_quant_patterns"] == [q_and_k]
+
+    path = tmp_path / "model.pth"
+    mto.save(attn, path)
+    restored = _DummyAttention()
+    mto.restore(restored, path)
+    assert _single_global_amax_state(restored).members == (restored.q_proj, restored.k_proj)

From 56a55d50c8da0eab792de16f5ef0364572b4643e Mon Sep 17 00:00:00 2001
From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
Date: Mon, 26 Oct 2026 11:02:12 -0700
Subject: [PATCH 5/6] Write NVFP4 KV-cache scales as amax / 448 in the Megatron
 and functional exporters

vLLM and TensorRT-LLM read NVFP4 k_scale / v_scale as amax / 448, as the HF
exporter writes them; amax / (6 * 448) overflowed in the attention kernels.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
---
 CHANGELOG.rst                                         |  1 +
 modelopt/torch/export/quant_utils.py                  | 11 +++++++----
 .../unit/torch/export/test_quantized_weight_state.py  |  6 +++---
 3 files changed, 11 insertions(+), 7 deletions(-)

diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index 8cd35516d8c..844f96893d6 100755
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -44,6 +44,7 @@ Changelog
 
 **Bug Fixes**
 
+- ``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.
 - 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.
 - 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.
 - The default PTQ recipes now leave multi-token-prediction (MTP) layers unquantized in Megatron vision-language models too (``language_model.mtp.*``), as they already did at the model root.
diff --git a/modelopt/torch/export/quant_utils.py b/modelopt/torch/export/quant_utils.py
index 2e73258f9cd..5fb6f32ddea 100755
--- a/modelopt/torch/export/quant_utils.py
+++ b/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.
diff --git a/tests/unit/torch/export/test_quantized_weight_state.py b/tests/unit/torch/export/test_quantized_weight_state.py
index ecf9fb2c8e5..19edc21cab0 100644
--- a/tests/unit/torch/export/test_quantized_weight_state.py
+++ b/tests/unit/torch/export/test_quantized_weight_state.py
@@ -733,10 +733,10 @@ def _kv_attention(k_cfg, v_cfg, k_amax=896.0, v_amax=224.0) -> nn.Module:
 @pytest.mark.parametrize(
     ("k_cfg", "v_cfg", "quantization_format", "k_scale", "v_scale"),
     [
-        # FP8 scales are amax / 448, raised to at least 1.0 like the other exporters do.
+        # Scales are amax / 448; FP8 ones are raised to at least 1.0 like the other exporters do.
         (_FP8_KV, _FP8_KV, "FP8", 2.0, 1.0),
-        (_NVFP4_KV, _NVFP4_KV, "NVFP4", 896 / (6 * 448), 224 / (6 * 448)),
-        (_FP8_KV, _NVFP4_KV, "FP8_K_NVFP4_V", 2.0, 224 / (6 * 448)),
+        (_NVFP4_KV, _NVFP4_KV, "NVFP4", 2.0, 0.5),
+        (_FP8_KV, _NVFP4_KV, "FP8_K_NVFP4_V", 2.0, 0.5),
     ],
 )
 def test_kv_cache_capture(k_cfg, v_cfg, quantization_format, k_scale, v_scale):

From 6bd1d6d2c6c6eba3b6b369b8d59926053fca6eb4 Mon Sep 17 00:00:00 2001
From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
Date: Mon, 26 Oct 2026 12:55:33 -0700
Subject: [PATCH 6/6] Reject affine KV caches in the functional KV-cache export

capture_kv_cache_export_state carries only k_scale / v_scale, so an affine
FP8 cache (FP8_AFFINE_KV_CFG) would lose its calibrated K/V biases.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
---
 modelopt/torch/export/quantized_weight_export.py       | 9 +++++++++
 tests/unit/torch/export/test_quantized_weight_state.py | 5 +++++
 2 files changed, 14 insertions(+)

diff --git a/modelopt/torch/export/quantized_weight_export.py b/modelopt/torch/export/quantized_weight_export.py
index cd1673f0735..c3e70624d90 100644
--- a/modelopt/torch/export/quantized_weight_export.py
+++ b/modelopt/torch/export/quantized_weight_export.py
@@ -524,6 +524,15 @@ def get_kv_cache_export_spec(attention: nn.Module) -> str | None:
         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(
diff --git a/tests/unit/torch/export/test_quantized_weight_state.py b/tests/unit/torch/export/test_quantized_weight_state.py
index 19edc21cab0..73e92734954 100644
--- a/tests/unit/torch/export/test_quantized_weight_state.py
+++ b/tests/unit/torch/export/test_quantized_weight_state.py
@@ -769,6 +769,11 @@ def test_kv_cache_spec_rejects_what_it_cannot_export():
     with pytest.raises(NotImplementedError, match="'INT8' KV cache"):
         get_kv_cache_export_spec(attention)
 
+    affine_fp8 = {**_FP8_KV, "bias": {-2: None, -4: None, "type": "static"}}
+    attention = _kv_attention(affine_fp8, affine_fp8)
+    with pytest.raises(NotImplementedError, match="affine KV caches"):
+        get_kv_cache_export_spec(attention)
+
 
 @pytest.mark.parametrize(
     "weight_spec", [None, weight_export._QuantizedWeightExportSpec("fp8", 0, True)]