Address review: skip disabled GLM-5.3-Flash kpool hooks, explain DSA-only freeze

Signed-off-by: Shiyang Chen <shiychen@nvidia.com>

CHANGELOG.rst

@@ -26,7 +26,7 @@ Changelog
 - Add support for quantizing and calibrating enabled operators outside the transformer layers, such as ``lm_head``, when using layerwise calibration.
 - Add an end-to-end BEVFormer ONNX PTQ example with temporal calibration data generation, INT8 and FP8 quantization, TensorRT engine building, and nuScenes accuracy evaluation. See `examples/onnx_ptq/bevformer/README.md <https://github.com/NVIDIA/Model-Optimizer/tree/main/examples/onnx_ptq/bevformer>`_ for details.
 - Add a reusable local-Hessian NVFP4 PTQ recipe and the quantization recipe used for ``nvidia/Qwen3.8-27B-NVFP4``.
-- Add fake quantization of the sparse-attention indexer key cache for Megatron-Core (``DSAIndexer``, ``CSAIndexer``) and vLLM (DeepSeek-V3.2, DeepSeek-V4, GLM-5.x, GLM-5.3-Flash) through the new ``indexer_k_quantizer``. Enable it by importing the ``configs/ptq/units/indexer_k_nvfp4`` unit (NVFP4 with its global scale fixed to 1) into a recipe; the KV-cache presets leave it disabled. Megatron-Core ``DSAIndexer`` models need ``dsa_indexer_rotate_activation=False``.
+- Add fake quantization of the sparse-attention indexer key cache for Megatron-Core (``DSAIndexer``, ``CSAIndexer``) and vLLM (DeepSeek-V3.2, DeepSeek-V4, GLM-5.x, GLM-5.3-Flash) through the new ``indexer_k_quantizer``. Enable it by importing the ``configs/ptq/units/indexer_k_nvfp4`` unit (NVFP4 with its global scale fixed to 1) into a recipe; Megatron-Core ``DSAIndexer`` models need ``dsa_indexer_rotate_activation=False``.
 
 *Speculative Decoding*
 

examples/megatron_bridge/distill.py

@@ -528,7 +528,9 @@ def main(args: argparse.Namespace, owns_the_run: bool = True):
     if DSAIndexer is not None and not (student_provider.dsa_indexer_loss_coeff or 0) > 0:
         # Without the indexer loss MCore runs the DSA indexer under no_grad, so its parameters never
         # get a gradient; left trainable, they fail the grad-buffer bucket reset of
-        # overlap_grad_reduce. Freeze them before DDP builds its buffers.
+        # overlap_grad_reduce. Freeze them before DDP builds its buffers. DeepSeek-V4's CSAIndexer
+        # does not hit this: MCore runs it with grad enabled and scales its loss by the coefficient,
+        # so its parameters still receive (zero) gradients.
         def _freeze_dsa_indexers_hook(model_chunks):
             for chunk in model_chunks:
                 for module in unwrap_model(chunk).modules():

modelopt/torch/quantization/plugins/vllm_indexer.py

@@ -342,6 +342,9 @@ def _wrap_kpool_cache_writer(kpool_ops: ModuleType, name: str, written_slots: Ca
     @functools.wraps(original)
     def wrapper(*args, **kwargs):
         out = original(*args, **kwargs)
+        # Conversion registers every indexer, also those whose quantizer is disabled.
+        if not any(indexer.indexer_k_quantizer.is_enabled for indexer in _glm5next_indexers):
+            return out
         bound = signature.bind(*args, **kwargs)
         bound.apply_defaults()
         kv_cache = bound.arguments["kv_cache"]

modelopt_recipes/configs/ptq/units/indexer_k_nvfp4.yaml

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# 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");

tests/gpu_megatron/torch/quantization/plugins/test_megatron.py

@@ -2118,6 +2118,7 @@ def test_indexer_k_quant_interleaved_rope(dist_workers_size_1):
 
 def _csa_indexer_model():
     """A standalone DeepSeek-V4 CSA indexer (compress ratio 4), as built by the dsv4_hybrid spec."""
+    # Imported here: CSA (older megatron-core) and Transformer Engine are optional; the test skips.
     from megatron.core.extensions.transformer_engine import TELinear, TENorm
     from megatron.core.models.common.embeddings import RotaryEmbedding
     from megatron.core.process_groups_config import ProcessGroupCollection

tests/gpu_vllm/torch/quantization/test_vllm_indexer_quant.py

@@ -20,6 +20,7 @@ of each layout adapter on stand-in modules, without booting an ``LLM`` (see
 ``test_vllm_dynamic_modules.py`` for the end-to-end DeepSeek-V3.2 run).
 """
 
+import weakref
 from types import SimpleNamespace
 
 import pytest
@@ -409,8 +410,19 @@ def test_glm5next_kpool_cache_writers_requantize_written_pools(
         ),
     )
     monkeypatch.setattr(_TestGlm5NextIndexer, "kpool_ops", kpool_ops)
+    monkeypatch.setattr(vllm_indexer, "_glm5next_indexers", weakref.WeakSet())
     indexer = _TestGlm5NextIndexer.convert(_NativeGlm5NextIndexer(kv_cache))
     vllm_indexer._install_kpool_cache_hooks(kpool_ops)  # a second indexer must not double-wrap
+
+    # No enabled quantizer: the kernel runs and the wrapper returns before looking up the cache.
+    indexer.indexer_k_quantizer.disable()
+    with monkeypatch.context() as m:
+        lookups = []
+        m.setattr(vllm_indexer, "_glm5next_quantizer_for", lookups.append)
+        kpool_ops.kpool_compress_and_write_cache(kv_cache, None, None, None, slot_mapping, 4)
+    assert not lookups
+    assert torch.equal(kv_cache, before)
+
     quantizer = indexer.indexer_k_quantizer = _fp8_quantizer(amax=2.0)
 
     # Prefill: pools at ``loc`` masked by ``write_mask``; an unrelated cache is left alone.
@@ -431,5 +443,5 @@ def test_glm5next_kpool_cache_writers_requantize_written_pools(
         kv_cache, None, None, None, None, None, dec_slots, dec_pos, 4
     )
     _assert_requantized(before, kv_cache, torch.tensor([17, 30]), quantizer)
-    assert kernel_calls == ["prefill", "prefill", "decode"]
+    assert kernel_calls == ["prefill", "prefill", "prefill", "decode"]
     assert getattr(kpool_ops.kpool_compress_and_write_cache, "_modelopt_indexer_k_wrapped", False)