From 0a512804f11970323995785cc43b9f13e64477ff Mon Sep 17 00:00:00 2001
From: Asim Kadav <19524+asimkadav@users.noreply.github.com>
Date: Fri, 2 Oct 2026 20:15:45 -0700
Subject: [PATCH 1/2] Add experimental softmax reparameterization for output
 heads

Signed-off-by: Asim Kadav <19524+asimkadav@users.noreply.github.com>
---
 CHANGELOG.rst                                 |   2 +
 .../softmax_reparameterization/README.md      | 100 +++++++++++
 .../softmax_reparameterization/__init__.py    |  18 ++
 .../softmax_reparameterization/example.py     |  58 +++++++
 .../requirements.txt                          |   1 +
 .../softmax_reparameterization/search.py      | 157 ++++++++++++++++++
 .../test_softmax_reparameterization.py        |  43 +++++
 .../test_softmax_reparameterization.py        |  97 +++++++++++
 8 files changed, 476 insertions(+)
 create mode 100644 experimental/softmax_reparameterization/README.md
 create mode 100644 experimental/softmax_reparameterization/__init__.py
 create mode 100644 experimental/softmax_reparameterization/example.py
 create mode 100644 experimental/softmax_reparameterization/requirements.txt
 create mode 100644 experimental/softmax_reparameterization/search.py
 create mode 100644 tests/gpu/torch/quantization/test_softmax_reparameterization.py
 create mode 100644 tests/unit/torch/quantization/test_softmax_reparameterization.py

diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index adcaec96729..77d80da1d1e 100755
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -13,6 +13,8 @@ Changelog
 
 *Quantization*
 
+- Add an experimental validation-KL search over softmax-equivalent linear output heads before weight-only quantization. See ``experimental/softmax_reparameterization`` for usage and current limitations.
+
 - Backfill checkpoint aliases for nine more published NVFP4 releases, so each is reachable from its source model's hub path: ``zai-org/GLM-5.1`` and ``GLM-5.2``, ``MiniMaxAI/MiniMax-M2.5`` and ``MiniMax-M3``, ``deepseek-ai/DeepSeek-V3.1`` and ``DeepSeek-V3.2``, ``Qwen/Qwen3-235B-A22B-Instruct-2507`` and ``-Thinking-2507``, and ``Qwen/Qwen3.6-27B``. Each imports an existing general or architecture recipe wholesale rather than copying its body.
 - Add composed Hugging Face AutoQuantize recipes that run fixed PTQ or weight AutoQuantize before
   a separate KV-cache AutoQuantize stage, with independent resumable checkpoints for the weight and
diff --git a/experimental/softmax_reparameterization/README.md b/experimental/softmax_reparameterization/README.md
new file mode 100644
index 00000000000..40e063a5fc8
--- /dev/null
+++ b/experimental/softmax_reparameterization/README.md
@@ -0,0 +1,100 @@
+# Softmax reparameterization for output-head quantization
+
+An experimental implementation of [Kadav et al., *Softmax Reparameterization for
+Output-Head Quantization*](https://arxiv.org/abs/2609.31291).
+
+For a vocabulary-by-hidden weight matrix `W`, each candidate is
+`W_t = W - t * W.mean(dim=0)`. Every vocabulary row receives the same vector
+subtraction, preserving the full-precision softmax distribution algebraically.
+Quantizing these equivalent heads produces different residuals. We fit the same
+ModelOpt quantizer independently for each candidate and select the coefficient
+with the lowest forward KL to the **original** head on validation states.
+
+## Usage
+
+Install ModelOpt from this checkout, then run the synthetic example:
+
+```bash
+python -m experimental.softmax_reparameterization.example --device cpu
+python -m experimental.softmax_reparameterization.example --device cuda --algorithm gptq
+```
+
+For an existing head, capture the final hidden states **entering the output
+projection**. Use articles disjoint across quantizer fitting, coefficient
+validation, and final test evaluation; exclude padding tokens. The caller must
+ensure these splits are disjoint. Each tensor has shape `[tokens, hidden_width]`
+and may reside on CPU to limit accelerator memory usage.
+
+```python
+import copy
+
+import torch
+
+import modelopt.torch.opt as mto
+import modelopt.torch.quantization as mtq
+from experimental.softmax_reparameterization import head_kl, search_head
+
+# source_head is an unquantized torch.nn.Linear immediately preceding softmax.
+config = copy.deepcopy(mtq.INT4_BLOCKWISE_WEIGHT_ONLY_CFG)
+result = search_head(source_head, fit_states, validation_states, config)
+print(result.coefficient, result.validation_kl)
+print("Test KL:", head_kl(source_head, result.model, test_states))
+mto.save(result.model, "selected_head.pt")
+restored = mto.restore(torch.nn.Sequential(copy.deepcopy(source_head)), "selected_head.pt")
+```
+
+The returned model is a one-layer `Sequential`, with the head at index `0`; quantizer name patterns refer
+to this container, not to the original language model. The default blockwise INT4
+preset applies to it without changing ModelOpt's default LM-head exclusions.
+To use GPTQ, set `config["algorithm"] = "gptq"` before search. This is ModelOpt's
+calibration implementation; its rounding and fitting conventions need not match
+the paper's quantizers.
+
+The ordered default grid is `[-2, -1, -0.5, 0, 0.5, 1, 1.5, 2, 2.5, 3, 4, 5, 6, 8]`.
+Zero is required; one provides the fixed mean-centering comparison. Exact ties
+keep the first candidate. The full validation curve is returned, and the selected
+calibrated candidate is retained without refitting. Including zero guarantees
+non-increasing **measured validation KL**, not improvement on unseen data.
+
+## Scope and numerical behavior
+
+- Supports plain, unsharded `torch.nn.Linear` heads, with optional unchanged bias,
+  on CPU or a single CUDA device. Inputs and weights use FP32, FP16, or BF16;
+  log-softmax and KL are computed in FP32 and accumulated in FP64.
+- Supports weight-only fake quantization with `max` or `gptq` calibration.
+  Activation quantization, composite quantizers, nonlinear logit transformations
+  (including soft-capping), and grouped coefficient search are outside this PR.
+- Every candidate has independent weights and quantizer state. The source head
+  and any tied input embedding remain untouched, including on failure. Using the
+  selected head in a tied model requires a separate output weight allocation and
+  disabling weight tying when saving/reloading the full model.
+- Equivalence holds in exact arithmetic. Casting shifted weights and low-precision
+  matrix products can introduce small differences even before quantization.
+- The search runs the head only and does not retrain or rerun the decoder. Peak
+  memory includes the source, candidate, selected head, and temporary folded-weight
+  copies; GPTQ additionally needs a hidden-width-squared Hessian.
+
+## Model support and deployment
+
+This first integration tests synthetic linear heads. It does not claim that the
+paper's model-level results have been reproduced with ModelOpt. Capture and
+full-model adapters are intentionally left to the caller, who must verify a
+shift-compatible logit path. Replacing a Hugging Face child module also needs a
+model-level ModelOpt conversion/export workflow; the standalone head container
+must not be mistaken for an export-ready language-model checkpoint.
+
+The shift is stored in weights and requires no extra inference operator for a
+linear-softmax head. This implementation evaluates fake-quantized heads and tests
+ModelOpt head save/restore. Packed export and TensorRT-LLM/vLLM/SGLang serving
+are **not validated** by this prototype; no latency benefit is claimed here.
+
+## Tests
+
+```bash
+pytest tests/unit/torch/quantization/test_softmax_reparameterization.py -q
+pytest tests/gpu/torch/quantization/test_softmax_reparameterization.py -q
+```
+
+CPU tests compare selection with independent ModelOpt candidates, check softmax
+equivalence, source/tied-embedding isolation, tie handling, invalid inputs, and
+save/restore. CUDA tests exercise BF16 heads with INT4 max and GPTQ calibration.
diff --git a/experimental/softmax_reparameterization/__init__.py b/experimental/softmax_reparameterization/__init__.py
new file mode 100644
index 00000000000..e9536cb56ca
--- /dev/null
+++ b/experimental/softmax_reparameterization/__init__.py
@@ -0,0 +1,18 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Experimental scalar search over equivalent linear-softmax output heads."""
+
+from .search import *
diff --git a/experimental/softmax_reparameterization/example.py b/experimental/softmax_reparameterization/example.py
new file mode 100644
index 00000000000..6d41295e2e7
--- /dev/null
+++ b/experimental/softmax_reparameterization/example.py
@@ -0,0 +1,58 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Run a reproducible, synthetic head-only search without downloading a model."""
+
+import argparse
+import copy
+import json
+
+import torch
+from torch import nn
+
+import modelopt.torch.quantization as mtq
+from experimental.softmax_reparameterization import head_kl, search_head
+
+
+def main():
+    """Compare selected and unshifted heads on disjoint synthetic test states."""
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--device", choices=["cpu", "cuda"], default="cpu")
+    parser.add_argument("--algorithm", choices=["max", "gptq"], default="max")
+    args = parser.parse_args()
+    torch.manual_seed(0)
+    head = nn.Linear(128, 512, bias=False).to(args.device).eval()
+    with torch.no_grad():
+        head.weight.add_(torch.randn(128, device=args.device) * 0.1)
+    fit, validation, test = torch.randn(3, 64, 128, device=args.device).unbind()
+    config = copy.deepcopy(mtq.INT4_BLOCKWISE_WEIGHT_ONLY_CFG)
+    config["algorithm"] = args.algorithm
+    result = search_head(head, fit, validation, config)
+    baseline = search_head(head, fit, validation, config, coefficients=(0,))
+    print(
+        json.dumps(
+            {
+                "coefficient": result.coefficient,
+                "validation_kl": result.validation_kl,
+                "test_kl_unshifted": head_kl(head, baseline.model, test),
+                "test_kl_selected": head_kl(head, result.model, test),
+            },
+            indent=2,
+        )
+    )
+
+
+if __name__ == "__main__":
+    main()
diff --git a/experimental/softmax_reparameterization/requirements.txt b/experimental/softmax_reparameterization/requirements.txt
new file mode 100644
index 00000000000..72a4e098730
--- /dev/null
+++ b/experimental/softmax_reparameterization/requirements.txt
@@ -0,0 +1 @@
+# No dependencies beyond base ModelOpt. Tests require pytest.
diff --git a/experimental/softmax_reparameterization/search.py b/experimental/softmax_reparameterization/search.py
new file mode 100644
index 00000000000..303b91cb412
--- /dev/null
+++ b/experimental/softmax_reparameterization/search.py
@@ -0,0 +1,157 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Validation-selected output-head shifts (https://arxiv.org/abs/2609.31291)."""
+
+import copy
+import math
+from collections.abc import Sequence
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+from torch import nn
+
+import modelopt.torch.quantization as mtq
+from modelopt.torch.quantization.nn import TensorQuantizer
+
+__all__ = ["DEFAULT_COEFFICIENTS", "SearchResult", "head_kl", "search_head"]
+
+DEFAULT_COEFFICIENTS = (-2, -1, -0.5, 0, 0.5, 1, 1.5, 2, 2.5, 3, 4, 5, 6, 8)
+
+
+@dataclass
+class SearchResult:
+    """Selected ModelOpt head container and token-averaged validation KL for every candidate."""
+
+    model: nn.Sequential
+    coefficient: float
+    validation_kl: tuple[tuple[float, float], ...]
+
+
+def _check_states(states: torch.Tensor, width: int) -> None:
+    if states.ndim != 2 or states.shape[0] == 0 or states.shape[1] != width:
+        raise ValueError("States must be a nonempty [tokens, hidden_width] tensor.")
+    if not states.is_floating_point() or not torch.isfinite(states).all():
+        raise ValueError("States must contain finite floating-point values.")
+
+
+@torch.no_grad()
+def head_kl(
+    source: nn.Linear,
+    candidate: nn.Module,
+    states: torch.Tensor,
+    batch_size: int = 128,
+) -> float:
+    """Mean forward KL from source to candidate on final hidden states.
+
+    Both readouts use the source weight dtype; log-softmax and KL use FP32.
+    States may reside on CPU and are moved one batch at a time. Use disjoint
+    test states here after freezing the coefficient selected by ``search_head``.
+    """
+    _check_states(states, source.in_features)
+    if batch_size <= 0:
+        raise ValueError("batch_size must be positive.")
+    total = torch.zeros((), device=source.weight.device, dtype=torch.float64)
+    for batch in states.split(batch_size):
+        batch = batch.to(device=source.weight.device, dtype=source.weight.dtype)
+        log_p = source(batch).float().log_softmax(-1)
+        log_q = candidate(batch).float().log_softmax(-1)
+        total += (log_p.exp() * (log_p - log_q)).sum(-1).double().sum()
+    value = (total / states.shape[0]).item()
+    if not math.isfinite(value):
+        raise ValueError("Non-finite KL; check the head, quantizer, and hidden states.")
+    return value
+
+
+@torch.no_grad()
+def search_head(
+    head: nn.Linear,
+    fit_states: torch.Tensor,
+    validation_states: torch.Tensor,
+    quant_config: dict[str, Any],
+    coefficients: Sequence[float] = DEFAULT_COEFFICIENTS,
+    batch_size: int = 128,
+) -> SearchResult:
+    """Fit shifted copies of a linear-softmax head and select by validation KL.
+
+    Each candidate subtracts ``t * head.weight.float().mean(0)`` from every
+    vocabulary row before independently fitting the same ModelOpt quantizer.
+    The original head (including a shared embedding parameter) is never mutated.
+
+    Args:
+        head: Plain, unquantized ``nn.Linear`` immediately preceding softmax.
+            Nonlinear logit paths and distributed/sharded heads are not supported.
+        fit_states: Final hidden states used exclusively for quantizer fitting.
+        validation_states: States from disjoint articles used only for selection.
+            The caller is responsible for enforcing article-level disjointness.
+        quant_config: ModelOpt weight-only fake-quant config for ``max`` or ``gptq``.
+            Patterns address a one-layer Sequential (e.g. ``0.weight_quantizer``).
+        coefficients: Finite, unique ordered candidates including zero; ties keep
+            the first candidate. Including one provides a fixed-centering control.
+        batch_size: Maximum number of token states per readout.
+
+    Returns:
+        The selected calibrated Sequential, coefficient, and validation curve.
+        Keep the container when using ModelOpt save/restore. Selection guarantees
+        no worse measured validation KL than t=0, not improvement on unseen data.
+    """
+    if type(head) is not nn.Linear:
+        raise TypeError("head must be a plain, unquantized nn.Linear.")
+    if head.weight.device.type not in ("cpu", "cuda"):
+        raise ValueError("Only materialized CPU or CUDA heads are supported.")
+    if head.weight.dtype not in (torch.float32, torch.float16, torch.bfloat16):
+        raise ValueError("Use FP32, FP16, or BF16 head weights.")
+    values = tuple(float(t) for t in coefficients)
+    if not values or not all(math.isfinite(t) for t in values):
+        raise ValueError("coefficients must be nonempty and finite.")
+    if 0.0 not in values or len(set(values)) != len(values):
+        raise ValueError("coefficients must be unique and include zero.")
+    if batch_size <= 0:
+        raise ValueError("batch_size must be positive.")
+    _check_states(fit_states, head.in_features)
+    _check_states(validation_states, head.in_features)
+    algorithm = quant_config.get("algorithm", "max")
+    method = algorithm.get("method") if isinstance(algorithm, dict) else algorithm
+    if method not in ("max", "gptq"):
+        raise ValueError("This prototype supports max and gptq calibration only.")
+
+    def fit(model):
+        for batch in fit_states.split(batch_size):
+            model(batch.to(device=head.weight.device, dtype=head.weight.dtype))
+
+    mean = head.weight.float().mean(0)
+    best_model, best_t, best_kl = None, 0.0, math.inf
+    curve = []
+    for coefficient in values:
+        candidate = nn.Sequential(copy.deepcopy(head)).eval().requires_grad_(False)
+        if coefficient != 0:
+            candidate[0].weight.copy_(head.weight.float() - coefficient * mean)
+        mtq.quantize(candidate, copy.deepcopy(quant_config), forward_loop=fit)
+        quantizer = candidate[0].weight_quantizer
+        if not isinstance(quantizer, TensorQuantizer) or not quantizer.is_enabled:
+            raise ValueError("Enable a single weight quantizer for the head.")
+        if not quantizer.fake_quant:
+            raise ValueError("Search requires fake quantization; pack only after selection.")
+        if candidate[0].input_quantizer.is_enabled or candidate[0].output_quantizer.is_enabled:
+            raise ValueError("Search supports weight-only quantization.")
+        # Avoid requantizing the full vocabulary matrix for each validation batch.
+        with mtq.temporarily_fold_weights(candidate):
+            kl = head_kl(head, candidate, validation_states, batch_size)
+        curve.append((coefficient, kl))
+        if kl < best_kl:
+            best_model, best_t, best_kl = candidate, coefficient, kl
+    assert best_model is not None
+    return SearchResult(best_model, best_t, tuple(curve))
diff --git a/tests/gpu/torch/quantization/test_softmax_reparameterization.py b/tests/gpu/torch/quantization/test_softmax_reparameterization.py
new file mode 100644
index 00000000000..8a52978810e
--- /dev/null
+++ b/tests/gpu/torch/quantization/test_softmax_reparameterization.py
@@ -0,0 +1,43 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""CUDA validation with ModelOpt INT4 max and GPTQ calibration."""
+
+import copy
+
+import pytest
+import torch
+from torch import nn
+
+import modelopt.torch.quantization as mtq
+from experimental.softmax_reparameterization import head_kl, search_head
+
+
+@pytest.mark.parametrize("algorithm", ["max", "gptq"])
+def test_cuda_head_search(algorithm):
+    torch.manual_seed(7)
+    head = nn.Linear(128, 256, bias=False, device="cuda", dtype=torch.bfloat16).eval()
+    original = head.weight.detach().clone()
+    fit, val, test = torch.randn(3, 64, 128, device="cuda", dtype=torch.bfloat16).unbind()
+    config = copy.deepcopy(mtq.INT4_BLOCKWISE_WEIGHT_ONLY_CFG)
+    config["algorithm"] = algorithm
+    result = search_head(head, fit, val, config, (0, 1, 4), batch_size=32)
+    curve = dict(result.validation_kl)
+    assert curve[result.coefficient] <= curve[0]
+    assert head_kl(head, result.model, val, 64) == pytest.approx(
+        curve[result.coefficient], abs=2e-5
+    )
+    assert head_kl(head, result.model, test) >= -1e-6
+    torch.testing.assert_close(head.weight, original, rtol=0, atol=0)
diff --git a/tests/unit/torch/quantization/test_softmax_reparameterization.py b/tests/unit/torch/quantization/test_softmax_reparameterization.py
new file mode 100644
index 00000000000..3ba46c67573
--- /dev/null
+++ b/tests/unit/torch/quantization/test_softmax_reparameterization.py
@@ -0,0 +1,97 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""End-to-end scalar selection using ModelOpt weight quantization."""
+
+import copy
+
+import pytest
+import torch
+from torch import nn
+
+import modelopt.torch.opt as mto
+import modelopt.torch.quantization as mtq
+from experimental.softmax_reparameterization import head_kl, search_head
+
+
+@pytest.mark.parametrize("bias", [False, True])
+def test_selection_matches_independent_candidates_and_preserves_source(bias, tmp_path):
+    torch.manual_seed(12)
+    head = nn.Linear(16, 32, bias=bias)
+    embedding = nn.Embedding(32, 16)
+    embedding.weight = head.weight
+    original = head.weight.detach().clone()
+    states = torch.randn(23, 16)
+    fit, val, test = states[:8], states[8:17], states[17:]
+    config = {
+        "quant_cfg": [
+            {"quantizer_name": "*", "enable": False},
+            {"quantizer_name": "*weight_quantizer", "cfg": {"num_bits": 4, "axis": 0}},
+        ],
+        "algorithm": "max",
+    }
+    original_config = copy.deepcopy(config)
+    grid = (0, 1, -1, 4)
+    result = search_head(head, fit, val, config, grid, batch_size=4)
+    independent = []
+    for t in grid:
+        shifted = copy.deepcopy(head).eval()
+        with torch.no_grad():
+            shifted.weight.copy_(original - t * original.mean(0))
+        torch.testing.assert_close(
+            shifted(test).softmax(-1), head(test).softmax(-1), atol=1e-6, rtol=1e-5
+        )
+        quantized = mtq.quantize(nn.Sequential(shifted), copy.deepcopy(config))
+        independent.append(head_kl(head, quantized, val, batch_size=9))
+    assert [kl for _, kl in result.validation_kl] == pytest.approx(independent, abs=1e-7)
+    assert result.coefficient == grid[independent.index(min(independent))]
+    assert head_kl(head, result.model, val) <= independent[0] + 1e-7
+    assert config == original_config
+    assert embedding.weight is head.weight
+    assert head.training
+    assert not hasattr(head, "weight_quantizer")
+    torch.testing.assert_close(head.weight, original, rtol=0, atol=0)
+    assert result.model[0].weight.data_ptr() != head.weight.data_ptr()
+    path = tmp_path / "head.pt"
+    mto.save(result.model, path)
+    restored = mto.restore(nn.Sequential(copy.deepcopy(head)), path)
+    torch.testing.assert_close(restored(test), result.model(test), rtol=0, atol=0)
+
+
+def test_ties_keep_first_and_invalid_candidates_fail():
+    head = nn.Linear(16, 4, bias=False)
+    with torch.no_grad():
+        head.weight.zero_()
+    states = torch.ones(3, 16)
+    config = {
+        "quant_cfg": [
+            {"quantizer_name": "*", "enable": False},
+            {"quantizer_name": "*weight_quantizer", "cfg": {"num_bits": 4, "axis": 0}},
+        ],
+        "algorithm": "max",
+    }
+    result = search_head(head, states, states.clone(), config, (1, 0, -1))
+    assert result.coefficient == 1
+    assert all(kl == 0 for _, kl in result.validation_kl)
+    for grid in ((1,), (0, 0), (0, float("nan"))):
+        with pytest.raises(ValueError, match="coefficients"):
+            search_head(head, states, states, config, grid)
+    with pytest.raises(ValueError, match="nonempty"):
+        search_head(head, states, states[:0], config)
+    config["quant_cfg"].append(
+        {"quantizer_name": "*input_quantizer", "cfg": {"num_bits": 8, "axis": None}}
+    )
+    with pytest.raises(ValueError, match="weight-only"):
+        search_head(head, states, states, config, (0,))

From e8bf828e2b9b4f20cbc784a84c1f1929be461060 Mon Sep 17 00:00:00 2001
From: Asim Kadav <19524+asimkadav@users.noreply.github.com>
Date: Fri, 2 Oct 2026 20:24:42 -0700
Subject: [PATCH 2/2] Give CUDA tests a unique module name for combined
 collection

Signed-off-by: Asim Kadav <19524+asimkadav@users.noreply.github.com>
---
 experimental/softmax_reparameterization/README.md               | 2 +-
 ...ameterization.py => test_softmax_reparameterization_cuda.py} | 0
 2 files changed, 1 insertion(+), 1 deletion(-)
 rename tests/gpu/torch/quantization/{test_softmax_reparameterization.py => test_softmax_reparameterization_cuda.py} (100%)

diff --git a/experimental/softmax_reparameterization/README.md b/experimental/softmax_reparameterization/README.md
index 40e063a5fc8..f6ceade3c16 100644
--- a/experimental/softmax_reparameterization/README.md
+++ b/experimental/softmax_reparameterization/README.md
@@ -92,7 +92,7 @@ are **not validated** by this prototype; no latency benefit is claimed here.
 
 ```bash
 pytest tests/unit/torch/quantization/test_softmax_reparameterization.py -q
-pytest tests/gpu/torch/quantization/test_softmax_reparameterization.py -q
+pytest tests/gpu/torch/quantization/test_softmax_reparameterization_cuda.py -q
 ```
 
 CPU tests compare selection with independent ModelOpt candidates, check softmax
diff --git a/tests/gpu/torch/quantization/test_softmax_reparameterization.py b/tests/gpu/torch/quantization/test_softmax_reparameterization_cuda.py
similarity index 100%
rename from tests/gpu/torch/quantization/test_softmax_reparameterization.py
rename to tests/gpu/torch/quantization/test_softmax_reparameterization_cuda.py