Skip Qwen3.8's QSA indexer loop during calibration when it keeps every token

CHANGELOG.rst

@@ -19,6 +19,7 @@ Changelog
 - Add ``hessian_input_dtype`` to the GPTQ and ``local_hessian`` algorithm configs. ``bfloat16`` accumulates the float32 Hessian with a tensor-core GEMM on bfloat16 activations, several times faster and without a float32 copy of each batch; the default ``float32`` keeps the existing float32 GEMM.
 - Add ``examples/hf_ptq/convert_ggml_to_int8.py`` to expand a checkpoint exported with a GGML recipe (IQ formats, ``q8_0``) into an INT8 weight-only compressed-tensors checkpoint. Stock vLLM serves it with its integer kernels at about half the BF16 weight memory, reproducing the GGML weights up to rounding their scales to the model dtype.
 - Speed up NVFP4 QAT and QAD of Hugging Face fused MoE experts: all experts of a projection are now fake-quantized in one kernel launch instead of one per expert, about 2x faster per MoE layer on Qwen3-30B-A3B with FSDP2.
+- ``hf_ptq.py`` calibrates Qwen3.8 (``qwen4_exp``) without running its QSA indexer's per-query Python loop when every visible key fits the indexer's token budget, with bit-identical results. On Qwen3.8-Flash-Next this saves about two hours of layerwise calibration.
 
 *Speculative Decoding*
 

modelopt/torch/models/README.md

@@ -105,5 +105,17 @@ only; subsystem logic stays in the subsystem.
 Keep architecture-specific loading and calibration helpers alongside the model's
 `specs.py` (for example, `nemotron_h/mtp.py`). The `hf.py` dispatcher loads these
 helpers on demand; examples call its loading context and calibration hook.
+
+A helper module registered in `hf.py`'s `_PLUGINS` defines any of these hooks, all
+optional, and lists exactly the ones it defines in `__all__`. The dispatcher rejects any
+other name there, so a misspelled hook fails instead of being skipped.
+
+| Hook | Called by | Default without it |
+|---|---|---|
+| `prepare_for_loading(checkpoint_path, trust_remote_code, *, model_class)` | `prepare_model_for_loading` | no context |
+| `has_mtp_weights(checkpoint_path)` | `checkpoint_has_mtp` | `False` |
+| `prefers_bundled_code(hf_config, checkpoint_path, config_kwargs)` | `prefers_bundled_code` | `False` |
+| `prepare_for_calibration(model)` | `prepare_model_for_calibration` | nothing |
+
 Do not import runtime helpers from a model package's `__init__.py`: registering
 specs must remain independent of optional Transformers and calibration imports.

modelopt/torch/models/__init__.py

@@ -93,6 +93,7 @@
     qwen3_moe,
     qwen3_next,
     qwen3_vl_moe,
+    qwen4_exp,
     step3p5,
     step3p7,
 )

modelopt/torch/models/hf.py

@@ -26,44 +26,58 @@
 ]
 
 # Keep modeling dependencies lazy: importing model specs must not import calibration code.
+# A plugin defines any of these hooks and lists exactly those in its __all__ (see README.md).
+_HOOKS = frozenset(
+    {"prepare_for_loading", "has_mtp_weights", "prefers_bundled_code", "prepare_for_calibration"}
+)
 _PLUGINS = {
     "nemotron_h": "nemotron_h.mtp",
     "nemotron_h_omni": "nemotron_h.mtp",
+    "qwen4_exp": "qwen4_exp.qsa_indexer",
+    "qwen4_exp_text": "qwen4_exp.qsa_indexer",
 }
 
 
 def _get_plugin(model_type):
     """Return the optional ModelOpt support plugin for a Hugging Face model type."""
     name = _PLUGINS.get(model_type)
-    return None if name is None else importlib.import_module(f".{name}", __package__)
+    if name is None:
+        return None
+    plugin = importlib.import_module(f".{name}", __package__)
+    # Hooks are optional, so a misspelled one would be skipped silently: refuse it instead.
+    unknown = set(getattr(plugin, "__all__", ())) - _HOOKS
+    if unknown:
+        raise ValueError(
+            f"{plugin.__name__} exports {sorted(unknown)}, not hooks of {sorted(_HOOKS)}"
+        )
+    return plugin
 
 
 def prepare_model_for_loading(
     model_type, checkpoint_path: str, trust_remote_code: bool, *, model_class=None
 ):
     """Prepare checkpoint-only modules, optionally on the loader's selected concrete class."""
-    plugin = _get_plugin(model_type)
-    if plugin is None:
+    hook = getattr(_get_plugin(model_type), "prepare_for_loading", None)
+    if hook is None:
         return nullcontext()
-    return plugin.prepare_for_loading(checkpoint_path, trust_remote_code, model_class=model_class)
+    return hook(checkpoint_path, trust_remote_code, model_class=model_class)
 
 
 def checkpoint_has_mtp(model_type, checkpoint_path: str) -> bool:
     """Use the model plugin to detect MTP tensors before selecting a supported loading path."""
-    plugin = _get_plugin(model_type)
-    return plugin is not None and plugin.has_mtp_weights(checkpoint_path)
+    hook = getattr(_get_plugin(model_type), "has_mtp_weights", None)
+    return hook is not None and hook(checkpoint_path)
 
 
 def prefers_bundled_code(model_type, hf_config, checkpoint_path: str, config_kwargs: dict) -> bool:
     """Use the model plugin to tell whether the checkpoint's bundled code must replace the built-in class."""
-    plugin = _get_plugin(model_type)
-    return plugin is not None and plugin.prefers_bundled_code(
-        hf_config, checkpoint_path, config_kwargs
-    )
+    hook = getattr(_get_plugin(model_type), "prefers_bundled_code", None)
+    return hook is not None and hook(hf_config, checkpoint_path, config_kwargs)
 
 
 def prepare_model_for_calibration(model) -> None:
     """Run the optional model-specific hook that augments the default calibration forward."""
     plugin = _get_plugin(getattr(getattr(model, "config", None), "model_type", None))
-    if plugin is not None:
-        plugin.prepare_for_calibration(model)
+    hook = getattr(plugin, "prepare_for_calibration", None)
+    if hook is not None:
+        hook(model)

modelopt/torch/models/qwen4_exp/__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.
+
+"""Per-model specs for HF model type ``qwen4_exp``."""
+
+from . import specs

modelopt/torch/models/qwen4_exp/qsa_indexer.py

@@ -0,0 +1,80 @@
+# 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.
+
+"""Calibration support for Qwen3.8's QSA indexer."""
+
+import torch
+
+from modelopt.torch.utils import warn_rank_0
+from modelopt.torch.utils.network import bind_forward_method
+
+__all__ = ["prepare_for_calibration"]
+
+
+def _forward(
+    self, hidden_states, position_embeddings, attention_mask, past_key_values, *args, **kwargs
+):
+    """Transformers' indexer, short-circuited when every visible key fits the token budget.
+
+    The reference loops in Python over every (sample, query) pair. A query keeps its
+    ``block_topk`` best complete key blocks plus the incomplete tail, so when the keys hold no
+    more blocks than that, it keeps every visible token: the selection is the visibility mask the
+    indexer was given, which this returns directly. A call with arguments this does not know,
+    from a newer Transformers, goes to the reference too.
+    """
+    if args or kwargs or attention_mask.shape[-1] // self.compress_ratio > self.block_topk:
+        return self._modelopt_reference_forward(
+            hidden_states, position_embeddings, attention_mask, past_key_values, *args, **kwargs
+        )
+    # The reference runs the projection on every forward, so a quantizer on it gets calibrated.
+    qk = self.index_qk_proj(hidden_states)
+    if past_key_values is not None:
+        # The reference appends this step's indexer keys to the cache before selecting.
+        batch_size, seq_length, _ = hidden_states.shape
+        token_k = qk[..., self.index_n_heads * self.index_head_dim :]
+        past_key_values.update_indexer(
+            token_k.reshape(batch_size, seq_length, -1, self.index_head_dim).squeeze(2),
+            self.layer_idx,
+        )
+    # The reference reads head slot 0 of the 4-D mask and returns one head slot.
+    mask = attention_mask[:, :1]
+    if attention_mask.is_floating_point():
+        min_dtype = torch.finfo(attention_mask.dtype).min
+        return torch.where(mask == 0, mask.new_zeros(()), min_dtype)
+    return mask.clone()
+
+
+def prepare_for_calibration(model) -> bool:
+    """Give each QSA indexer of ``model`` the fast path; return whether ``model`` has one.
+
+    Calibration sequences rarely exceed the indexer's token budget, so on Qwen3.8-Flash-Next
+    (2,048 tokens) this skips the loop for every ``--calib_seq`` up to 2,048, with bit-identical
+    results, and runs the reference otherwise.
+    """
+    found = False
+    for module in model.modules():
+        if type(module).__name__ != "Qwen4ExpTextQSAIndexer":
+            continue
+        found = True
+        if not hasattr(module, "_modelopt_reference_forward"):
+            # Keeps an Accelerate hook already on the indexer, such as device_map's, around the fast path.
+            bind_forward_method(module, _forward, "_modelopt_reference_forward")
+    if not found:
+        warn_rank_0(
+            "No Qwen4ExpTextQSAIndexer found in this Qwen3.8 model, so calibration runs "
+            "Transformers' QSA indexer loop, which can take hours. If Transformers renamed the "
+            "class, update modelopt/torch/models/qwen4_exp/qsa_indexer.py."
+        )
+    return found

modelopt/torch/models/qwen4_exp/specs.py

@@ -0,0 +1,25 @@
+# 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.
+
+"""Qwen3.8 specs (HF model type ``qwen4_exp``)."""
+
+from ..specs import ModelSpec, register
+
+__all__: list[str] = []
+
+# Only the model type is recorded, so the dispatcher in ``hf.py`` can find this package's
+# calibration helper. The fused MoE layout is left undeclared: export and quantization handle
+# these experts generically, and declaring it would change their behavior for this model.
+register(ModelSpec(model_type="qwen4_exp", min_transformers_version="5.16"))

modelopt/torch/quantization/model_calib.py

@@ -2765,10 +2765,12 @@ def _encode_on_each_gpu(by_device, encode):
     # CUDA extension, happens once before the workers start.
     first, *by_device[gpus[0]] = by_device[gpus[0]]
     encode(*first)
+    inference = torch.is_inference_mode_enabled()
 
     def worker(device):
-        torch.cuda.set_device(device)  # per thread, like grad mode
-        with torch.no_grad():
+        # The device, grad mode and inference mode are per thread: take the caller's.
+        torch.cuda.set_device(device)
+        with torch.inference_mode() if inference else torch.no_grad():
             for item in by_device[device]:
                 encode(*item)
         torch.cuda.synchronize(device)

tests/gpu/torch/quantization/test_ggml_imatrix_gpu.py

@@ -15,9 +15,9 @@
 
 """GPU tests for imatrix calibration of the GGML formats."""
 
-import copy
 import time
 
+import pytest
 import torch
 from torch import nn
 
@@ -49,29 +49,34 @@ def forward(self, x):
         return x
 
 
-def test_imatrix_encodes_each_gpus_weights_at_once(need_2_gpus, monkeypatch):
-    """Weights on two GPUs are encoded concurrently, to the same codes as on one GPU."""
+@pytest.mark.parametrize("inference_mode", [False, True])
+def test_imatrix_encodes_each_gpus_weights_at_once(need_2_gpus, monkeypatch, inference_mode):
+    """Weights on two GPUs are encoded concurrently, to the same codes as on one GPU, also as
+    inference tensors under the caller's inference mode, which worker threads do not inherit."""
     spans = []
     encode = GGMLImatrixEncoder.encode
 
     def timed_encode(self, weight, quantizer, second_moments):
         start = time.perf_counter()
         decoded = encode(self, weight, quantizer, second_moments)
         torch.cuda.synchronize(weight.device)
-        time.sleep(0.2)  # long enough that encodes on different GPUs must overlap to save time
+        # Stands in for a long encode: encodes on two GPUs overlap only if they run at once.
+        time.sleep(0.2)
         spans.append((weight.device.index, start, time.perf_counter()))
         return decoded
 
     monkeypatch.setattr(GGMLImatrixEncoder, "encode", timed_encode)
-    torch.manual_seed(0)
-    reference = _Layers().cuda()
-    model = copy.deepcopy(reference)
-    model.layers[2:].to("cuda:1")
-    inputs = torch.randn(8, 128, 512, device="cuda")
+    with torch.inference_mode(inference_mode):
+        torch.manual_seed(0)
+        reference = _Layers().cuda()
+        torch.manual_seed(0)
+        model = _Layers().cuda()
+        model.layers[2:].to("cuda:1")
+        inputs = torch.randn(8, 128, 512, device="cuda")
 
-    mtq.quantize(reference, IQ2_XXS_IMATRIX, forward_loop=lambda m: m(inputs))
-    spans.clear()
-    mtq.quantize(model, IQ2_XXS_IMATRIX, forward_loop=lambda m: m(inputs))
+        mtq.quantize(reference, IQ2_XXS_IMATRIX, forward_loop=lambda m: m(inputs))
+        spans.clear()
+        mtq.quantize(model, IQ2_XXS_IMATRIX, forward_loop=lambda m: m(inputs))
 
     for expected, layer in zip(reference.layers, model.layers):
         assert torch.equal(layer.weight.cuda(0), expected.weight)