From 9a96e075124e8bf4e151579e33e7ea0e5c3f78bd Mon Sep 17 00:00:00 2001 From: Shiyang Chen Date: Mon, 12 Oct 2026 11:51:13 -0700 Subject: [PATCH 1/6] Add GLM-5.3-Flash indexer scorer keyword arguments to vLLM fake-quant serving The k-pool indexer of GLM-5.3-Flash gets indexer_scorer_kwargs_quantizer: enabled with backend scorer_kwargs, its backend_extra_args become keyword arguments of that layer's DeepGEMM scorer kernel, e.g. the numerics options of a DeepGEMM build. register_indexer_scorer_kwargs_quantizer lets plugins add quantizers that turn their own configs into such arguments, without patching vLLM. The plugin checks early that the DeepGEMM build declares them in MQA_LOGITS_NUMERICS (stock DeepGEMM takes logits_dtype), widens BF16 scores to FP32 and maps NaN scores to -inf before vLLM's top-k, whose prefill kernel returns out-of-range indices on NaN. The scorer gets only tensors, so it cannot tell which layer calls it. The new vllm_layer_scope module keeps a process-wide weak registry from layer names to the converted modules that own them. The GLM-5.3-Flash indexer registers itself under its K cache's prefix, the name its op gets, and is in scope while the op runs. The op wrap goes under vLLM's eager_break_during_capture, as breakable CUDA graph replays call the undecorated op. The k-pool K cache writers now find their indexer the same way, instead of matching the cache tensor against every indexer; while indexer_k_quantizer is on, a write outside the indexer op or to another layer's cache raises. A GPU test runs the indexer op in a FULL CUDA graph and in vLLM's breakable CUDA graph: after the captured inputs change, replays must re-quantize the K-cache writes and score with each layer's scorer quantizers. The layer-scope tests need neither vLLM nor a GPU and are unit tests. Signed-off-by: Shiyang Chen --- CHANGELOG.rst | 1 + examples/vllm_serve/README.md | 45 ++ modelopt/torch/quantization/model_quant.py | 5 + .../torch/quantization/plugins/__init__.py | 1 + .../quantization/plugins/vllm_indexer.py | 380 ++++++++-- .../quantization/plugins/vllm_layer_scope.py | 122 ++++ .../quantization/test_vllm_dynamic_modules.py | 130 +++- .../quantization/test_vllm_indexer_quant.py | 653 +++++++++++++++++- .../plugins/test_vllm_layer_scope.py | 148 ++++ .../torch/quantization/test_quantize_cpu.py | 13 +- 10 files changed, 1393 insertions(+), 105 deletions(-) create mode 100644 modelopt/torch/quantization/plugins/vllm_layer_scope.py create mode 100644 tests/unit/torch/quantization/plugins/test_vllm_layer_scope.py diff --git a/CHANGELOG.rst b/CHANGELOG.rst index f3a20bc762b..e63217e7402 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -36,6 +36,7 @@ Changelog - Add a reusable local-Hessian NVFP4 PTQ recipe and the quantization recipe used for ``nvidia/Qwen3.8-27B-NVFP4``. - Add experimental dynamic FP8 fake quantization of GatedDeltaNet chunk-boundary states and WY activations for training through the standard ``quant_cfg`` interface. The fused GDN path requires ``fla-core==0.5.1`` and chunk size 64; state emulation requires SM89 or newer. - Add fake quantization of the sparse-attention indexer key cache and query for DeepSeek-V4 (vLLM and Megatron-Core) and GLM-5.3-Flash (vLLM) through the new ``indexer_k_quantizer`` and ``indexer_q_quantizer``. Enable them by importing the ``configs/ptq/units/indexer_k_nvfp4`` and ``configs/ptq/units/indexer_q_nvfp4`` units (NVFP4 with the global scale fixed to 1) into a recipe. +- Add ``indexer_scorer_kwargs_quantizer`` to pass keyword arguments, such as the numerics options of a DeepGEMM build, to the DeepGEMM indexer scorer kernel of GLM-5.3-Flash per layer in vLLM fake-quant serving; see `examples/vllm_serve/README.md `_. - Add the ``configs/ptq/units/kv_nvfp4_mla`` recipe unit for fake quantization of the MLA KV cache (DeepSeek-V3, GLM-5.3-Flash, ...) in vLLM fake-quant serving: NVFP4 for the latent and FP8 for the RoPE key. See `examples/vllm_serve/README.md `_ for an example recipe. - Add fake quantization of MoE expert-parallel communication to vLLM fake-quant serving through the new ``dispatch_quantizer`` and ``combine_quantizer`` on vLLM's routed experts, for MoE layers built on vLLM's fused-MoE layer and any all-to-all backend, and for the DeepGEMM Mega-MoE kernel of DeepSeek-V4 and Kimi-K3 (``--moe-backend deep_gemm_mega_moe``). Enable them by importing the ``configs/ptq/units/moe_dispatch_nvfp4`` and ``configs/ptq/units/moe_combine_nvfp4`` units into a recipe; see ``examples/vllm_serve/README.md``. - vLLM fake-quant serving now runs on pre-quantized checkpoints such as FP8 when the recipe leaves those layers unquantized (for example a KV-cache-only recipe), and on MLA models with an FP8 KV cache; both previously failed during quantization. diff --git a/examples/vllm_serve/README.md b/examples/vllm_serve/README.md index 3e821367e2e..4913aac4c6b 100644 --- a/examples/vllm_serve/README.md +++ b/examples/vllm_serve/README.md @@ -315,6 +315,51 @@ Notes: - vLLM quantizes the DeepSeek-V4 indexer key and query without a Hadamard rotation and the GLM-5.3-Flash ones after one, so the fake quantization applies in that basis. +### Indexer scorer kernel arguments + +GLM-5.3-Flash's indexer scores each key as `k_scale * sum_h W[h] * ReLU(q_h . k)` in DeepGEMM, from +the FP8 query and key and the effective head weights `W` (with the query scale folded in). +`indexer_scorer_kwargs_quantizer` passes keyword arguments to that scorer kernel, per layer, e.g. the +numerics options of a DeepGEMM build. It stays disabled unless a recipe enables it with +`backend: scorer_kwargs`; its `backend_extra_args` are the arguments. For example, this recipe sums +the heads in BF16: + +```yaml +# modelopt-schema: modelopt.recipe.config.ModelOptPTQRecipe +imports: + base_disable_all: configs/ptq/units/base_disable_all + +metadata: + description: BF16 head sum in the indexer scorer. +quantize: + algorithm: max + quant_cfg: + - $import: base_disable_all + - quantizer_name: '*indexer_scorer_kwargs_quantizer' + cfg: + backend: scorer_kwargs + backend_extra_args: {logits_dtype: bfloat16} +``` + +Plugins can add quantizers that turn their own configs into such arguments with +`register_indexer_scorer_kwargs_quantizer` of `modelopt.torch.quantization.plugins.vllm_indexer`. +Stock DeepGEMM takes `logits_dtype` (on SM100 GPUs); other arguments need a DeepGEMM build that takes +them and declares them, with their values, in `deep_gemm.MQA_LOGITS_NUMERICS`, such as a DeepGEMM +fork. ModelOpt ships none. vLLM imports an installed `deep_gemm` before its own copy, so install +such a build into the serving environment and serve: + +```bash +pip install --no-deps --no-build-isolation +RECIPE_PATH=.yaml \ + python vllm_serve_fakequant.py -tp 8 --host 0.0.0.0 --port 8000 +``` + +Notes: + +- Scores that come out NaN, e.g. from overflowing reduced-precision arithmetic in the scorer, are set + to `-inf` before top-k, as vLLM's prefill top-k mishandles NaN. +- DeepSeek-V4 rejects the scorer quantizers. + ## Fake-quantize MoE expert-parallel communication To study compressed expert-parallel (EP) communication, fake-quantize what the MoE layers send diff --git a/modelopt/torch/quantization/model_quant.py b/modelopt/torch/quantization/model_quant.py index ca49e92284f..12ca0359d7d 100644 --- a/modelopt/torch/quantization/model_quant.py +++ b/modelopt/torch/quantization/model_quant.py @@ -237,6 +237,11 @@ def _check_weight_quantization_took_effect(model: nn.Module, config: QuantizeCon _PLUGIN_QUANTIZERS = { "indexer_q_quantizer": (("*Indexer",), "sparse-attention indexer", "indexer_*_nvfp4"), "indexer_k_quantizer": (("*Indexer",), "sparse-attention indexer", "indexer_*_nvfp4"), + "indexer_scorer_kwargs_quantizer": ( + ("*Indexer",), + "sparse-attention indexer", + "indexer_*_nvfp4", + ), "dispatch_quantizer": (("*Experts*", "*FusedMoE"), "MoE experts module", "moe_*_nvfp4"), "combine_quantizer": (("*Experts*", "*FusedMoE"), "MoE experts module", "moe_*_nvfp4"), } diff --git a/modelopt/torch/quantization/plugins/__init__.py b/modelopt/torch/quantization/plugins/__init__.py index c1c50be74f4..161dd4bc320 100644 --- a/modelopt/torch/quantization/plugins/__init__.py +++ b/modelopt/torch/quantization/plugins/__init__.py @@ -72,6 +72,7 @@ with import_plugin("vllm"): from .vllm import * from .vllm_indexer import * + from .vllm_layer_scope import * from .vllm_mega_moe import * with import_plugin("trl"): diff --git a/modelopt/torch/quantization/plugins/vllm_indexer.py b/modelopt/torch/quantization/plugins/vllm_indexer.py index cc18b2bfea8..db43111f3c3 100644 --- a/modelopt/torch/quantization/plugins/vllm_indexer.py +++ b/modelopt/torch/quantization/plugins/vllm_indexer.py @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Fake quantization of the sparse-attention indexer query and K cache in vLLM. +"""Fake quantization of the sparse-attention indexer query, K cache and scorer in vLLM. Covers the CSA indexer of DeepSeek-V4 and the k-pool indexer of GLM-5.3-Flash. @@ -23,6 +23,12 @@ current step, the query) are dequantized, fake-quantized and quantized again with the kernels' scale rule. The names avoid the ``*[kv]_bmm_quantizer`` globs, so the KV-cache presets leave them disabled. + +GLM-5.3-Flash's DeepGEMM scorer computes ``score[q, k] = k_scale[k] * sum_h W[q, h] * +ReLU(q_h . k)``. ``indexer_scorer_kwargs_quantizer`` passes keyword arguments to that kernel, per +layer, e.g. the numerics options of a DeepGEMM build that declares them in +``MQA_LOGITS_NUMERICS``. :func:`register_indexer_scorer_kwargs_quantizer` adds quantizers that turn +their own configs into such arguments. They start disabled. """ import functools @@ -35,10 +41,18 @@ import torch from vllm.forward_context import get_forward_context -from ..nn import QuantModule, QuantModuleRegistry, TensorQuantizer +from ..config import QuantizerAttributeConfig +from ..nn import ( + QuantModule, + QuantModuleRegistry, + TensorQuantizer, + is_registered_quant_backend, + register_quant_backend, +) from .vllm import create_parallel_state +from .vllm_layer_scope import register_layer_owner, wrap_layer_callee, wrap_layer_op -__all__ = [] +__all__ = ["register_indexer_scorer_kwargs_quantizer"] # NotImplementedError: some vLLM model packages reject unsupported platforms at import # (e.g. glm5next on XPU). @@ -47,9 +61,15 @@ except (ImportError, NotImplementedError): VllmDeepseekV4Indexer = None +# Older vLLM releases have no DeepGEMM wrappers; GLM-5.3-Flash's indexer needs them. +try: + from vllm.utils import deep_gemm as vllm_deep_gemm +except ImportError: + vllm_deep_gemm = None + -def _import_glm5next_indexer() -> tuple[type | None, ModuleType | None]: - """Return GLM-5.3-Flash's ``Indexer`` and the module its indexer op binds as ``kpool_ops``.""" +def _import_glm5next_indexer() -> tuple[type | None, ModuleType | None, ModuleType | None]: + """Return GLM-5.3-Flash's ``Indexer``, the module of its indexer op and the op's ``kpool_ops``.""" for indexer_path, op_path in ( # vLLM main: the op module is platform-dispatched, follow it to the bound variant. ("vllm.models.glm5next.common.attention", "vllm.models.glm5next.sparse_indexer"), @@ -66,11 +86,11 @@ def _import_glm5next_indexer() -> tuple[type | None, ModuleType | None]: kpool_ops = op_module.kpool_ops except (ImportError, AttributeError, NotImplementedError): continue - return indexer_cls, kpool_ops - return None, None + return indexer_cls, op_module, kpool_ops + return None, None, None -VllmGlm5NextIndexer, _glm5next_kpool_ops = _import_glm5next_indexer() +VllmGlm5NextIndexer, _glm5next_op_module, _glm5next_kpool_ops = _import_glm5next_indexer() _INDEXER_FP8_MAX = 448.0 _INDEXER_K_SCALE_BYTES = 4 @@ -156,12 +176,80 @@ def _get_arg(args: tuple, kwargs: dict, name: str, pos: int): return kwargs[name] if name in kwargs else args[pos] +# The built-in quantizer whose ``backend_extra_args`` are the scorer kernel's keyword arguments. +_SCORER_KWARGS_QUANTIZER = "indexer_scorer_kwargs_quantizer" +_SCORER_KWARGS_BACKEND = "scorer_kwargs" +# Scorer keyword arguments of stock DeepGEMM; a build declares the others it accepts, with their +# values, in ``deep_gemm.MQA_LOGITS_NUMERICS``. +_STOCK_SCORER_KWARGS = ("logits_dtype",) + + +def _passed_scorer_kwargs(quantizer: TensorQuantizer) -> dict: + """``indexer_scorer_kwargs_quantizer``: its ``backend_extra_args``, passed on as they are.""" + if quantizer.backend != _SCORER_KWARGS_BACKEND: + raise ValueError( + f"{_SCORER_KWARGS_QUANTIZER} needs backend '{_SCORER_KWARGS_BACKEND}', " + f"got {quantizer.backend!r}." + ) + return dict(quantizer.backend_extra_args or {}) + + +# The quantizers whose configs become keyword arguments of the GLM-5.3-Flash indexer's DeepGEMM +# scorer kernel, with the functions that turn an enabled one into those arguments. +_SCORER_KWARGS_QUANTIZERS: dict[str, Callable[[TensorQuantizer], dict]] = { + _SCORER_KWARGS_QUANTIZER: _passed_scorer_kwargs +} + + +def register_indexer_scorer_kwargs_quantizer( + name: str, scorer_kwargs: Callable[[TensorQuantizer], dict] +) -> None: + """Add the quantizer ``name`` to the vLLM indexers that are converted from now on. + + It starts disabled. Enabled, ``scorer_kwargs(quantizer)`` turns its config into keyword + arguments of the GLM-5.3-Flash indexer's DeepGEMM scorer kernel, which every scorer call of + that layer then gets while the quantizer quantizes; it raises for a config it cannot express. + Other indexers reject the quantizer. The built-in ``indexer_scorer_kwargs_quantizer`` passes its + ``backend_extra_args`` (``backend: scorer_kwargs``) on as they are. + """ + _SCORER_KWARGS_QUANTIZERS[name] = scorer_kwargs + + +def _scorer_kwargs_backend(inputs: torch.Tensor, quantizer: TensorQuantizer) -> torch.Tensor: + raise RuntimeError( + f"backend '{_SCORER_KWARGS_BACKEND}' passes keyword arguments to the DeepGEMM indexer " + "scorer kernel and never runs on a tensor." + ) + + +if not is_registered_quant_backend(_SCORER_KWARGS_BACKEND): + register_quant_backend(_SCORER_KWARGS_BACKEND, _scorer_kwargs_backend) + + +def _scorer_logits_dtype(value) -> torch.dtype: + """The scorer's ``logits_dtype``, also from its name in a config (``bfloat16``).""" + dtype = getattr(torch, value, None) if isinstance(value, str) else value + if dtype not in (torch.float32, torch.bfloat16): + raise ValueError(f"The scorer's logits_dtype is float32 or bfloat16, got {value!r}.") + return dtype + + +def _widen_logits(logits: torch.Tensor) -> torch.Tensor: + """FP32 copy of ``logits`` with the same strides, which vLLM's top-k kernels read explicitly.""" + widened = torch.empty_strided( + logits.shape, logits.stride(), dtype=torch.float32, device=logits.device + ) + return widened.copy_(logits) + + class _QuantVLLMIndexerBase(QuantModule): - """Owner of ``indexer_q_quantizer`` and ``indexer_k_quantizer`` for one vLLM indexer layout.""" + """Owner of the indexer quantizers for one vLLM indexer layout.""" def _setup(self): self.indexer_q_quantizer = TensorQuantizer() self.indexer_k_quantizer = TensorQuantizer() + for name in _SCORER_KWARGS_QUANTIZERS: # opt-in, unlike the q and K quantizers + setattr(self, name, TensorQuantizer(QuantizerAttributeConfig(enable=False))) self.parallel_state = create_parallel_state() def forward(self, *args, **kwargs): @@ -171,6 +259,27 @@ def forward(self, *args, **kwargs): # would otherwise try to convert it a second time (inconsistent MRO). return super().forward(*args, **kwargs) + def modelopt_post_restore(self, prefix: str = ""): + """Also reject restored scorer quantizers that the layout cannot apply.""" + super().modelopt_post_restore(prefix) + self._validate_scorer_quantizers() + + def _enabled_scorer_quantizers(self) -> list[str]: + # A quantizer registered after the conversion is missing. + return [ + name + for name in _SCORER_KWARGS_QUANTIZERS + if hasattr(self, name) and getattr(self, name).is_enabled + ] + + def _validate_scorer_quantizers(self) -> None: + """Reject enabled scorer quantizers; the layouts that implement them override this.""" + if enabled := self._enabled_scorer_quantizers(): + raise NotImplementedError( + f"{', '.join(enabled)}: the indexer scorer quantizers are only implemented for the " + "GLM-5.3-Flash k-pool indexer." + ) + class _QuantVLLMDeepseekV4Indexer(_QuantVLLMIndexerBase): """DeepSeek-V4 indexer: fused kernels FP8-quantize the query and cache the compressed key. @@ -187,6 +296,7 @@ def _setup(self): self._indexer_weights_pos = _native_positional_index(self, "forward", "indexer_weights") def forward(self, *args, **kwargs): + self._validate_scorer_quantizers() q, q_scale, weights = super().forward(*args, **kwargs) if self.indexer_k_quantizer.is_enabled: self._check_fp8_indexer() @@ -238,19 +348,11 @@ def _requantize_written_keys(self, positions: torch.Tensor) -> None: ) -# GLM-5.3-Flash indexers converted in this process, looked up by the cache tensor a kernel writes. +# GLM-5.3-Flash indexers converted in this process, for the check that the K cache is written inside +# their indexer op. _glm5next_indexers: weakref.WeakSet = weakref.WeakSet() -def _glm5next_quantizer_for(kv_cache: torch.Tensor) -> TensorQuantizer | None: - """The enabled ``indexer_k_quantizer`` of the indexer that owns ``kv_cache``, if any.""" - for indexer in _glm5next_indexers: - quantizer = indexer.indexer_k_quantizer - if quantizer.is_enabled and indexer.k_cache.kv_cache.data_ptr() == kv_cache.data_ptr(): - return quantizer - return None - - def _kpool_prefill_written(arguments: dict) -> tuple[torch.Tensor, torch.Tensor]: """Slots written by ``kpool_compress_and_write_cache``: ``loc`` where valid and unmasked.""" loc = arguments["loc"] @@ -271,34 +373,119 @@ def _kpool_decode_written(arguments: dict) -> tuple[torch.Tensor, torch.Tensor]: return slots, valid -def _wrap_kpool_cache_writer(kpool_ops: ModuleType, name: str, written_slots: Callable) -> None: - """Wrap ``kpool_ops.`` so the pools it wrote are re-quantized after every call. +def _requantize_written_pools(written_slots: Callable) -> Callable: + """A K-cache writer call of an indexer: the kernel, then a re-quantization of what it wrote.""" - The wrap is permanent: the indexer op is a breakable-cudagraph eager break that vLLM replays - without re-entering ``Indexer.forward``, so a per-forward patch would miss every replayed step. - """ - original = getattr(kpool_ops, name) - if getattr(original, "_modelopt_indexer_k_wrapped", False): - return - signature = inspect.signature(original) - - @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() + def call(indexer, original: Callable, bound: inspect.BoundArguments): kv_cache = bound.arguments["kv_cache"] - quantizer = _glm5next_quantizer_for(kv_cache) - if quantizer is not None: - slots, valid = written_slots(bound.arguments) - _requantize_fp8_indexer_k_cache(kv_cache, slots, valid, quantizer) + if kv_cache.data_ptr() != indexer.k_cache.kv_cache.data_ptr(): + raise RuntimeError( + f"The indexer op of {indexer.k_cache.prefix} writes the K cache of another layer." + ) + out = original(*bound.args, **bound.kwargs) + slots, valid = written_slots(bound.arguments) + _requantize_fp8_indexer_k_cache(kv_cache, slots, valid, indexer.indexer_k_quantizer) return out - wrapper._modelopt_indexer_k_wrapped = True # type: ignore[attr-defined] - setattr(kpool_ops, name, wrapper) + return call + + +def _reject_k_cache_write_outside_op() -> None: + """Reject a K-cache write outside the indexer op, which alone names the layer of the cache.""" + if any(indexer.indexer_k_quantizer.is_enabled for indexer in _glm5next_indexers): + raise RuntimeError( + "vLLM wrote the GLM-5.3-Flash indexer K cache outside the indexer op, so " + "indexer_k_quantizer cannot tell whose cache it is." + ) + + +# GLM-5.3-Flash's K-cache writers, with the slots that a call writes. +_KPOOL_CACHE_WRITERS = { + "kpool_compress_and_write_cache": _kpool_prefill_written, + "kpool_decode_update_and_maybe_write_cache_batched": _kpool_decode_written, +} +# GLM-5.3-Flash's indexer op, which writes the K cache and scores and selects the keys. +_GLM5NEXT_INDEXER_OP = "sparse_attn_indexer_kpool" + + +def _wrap_deep_gemm_scorer(deep_gemm_utils: ModuleType, paged: bool) -> None: + """Wrap vLLM's DeepGEMM scorer in ``deep_gemm_utils`` to apply the scorer quantizers. + + ``paged``: the decode scorer, else the prefill one. The indexer op imports the scorer from the + module on every call. + """ + wrap_layer_callee( + deep_gemm_utils, + "fp8_fp4_paged_mqa_logits" if paged else "fp8_fp4_mqa_logits", + lambda indexer, original, bound: indexer._score(deep_gemm_utils, original, bound, paged), + applies=lambda indexer: bool(indexer._enabled_scorer_quantizers()), + ) + + +def _scorer_deep_gemm(deep_gemm_utils: ModuleType, scorer_kwargs: dict) -> ModuleType: + """The DeepGEMM module vLLM loaded, checked to accept the scorer keyword arguments here. + + Stock DeepGEMM takes ``logits_dtype``; a build declares the others it accepts, with their + values, in ``MQA_LOGITS_NUMERICS``. The scorer only runs for long sequences, so this fails + early instead. + """ + if hasattr(deep_gemm_utils, "_lazy_init"): # like vLLM's wrappers: JIT cache directory, PDL + deep_gemm_utils._lazy_init() + deep_gemm = deep_gemm_utils._import_deep_gemm() + if deep_gemm is None: + raise RuntimeError("The indexer scorer quantizers need DeepGEMM, which vLLM did not find.") + declared = getattr(deep_gemm, "MQA_LOGITS_NUMERICS", None) or {} + undeclared = { + key: value + for key, value in scorer_kwargs.items() + if key not in _STOCK_SCORER_KWARGS and value not in declared.get(key, ()) + } + if undeclared: + raise RuntimeError( + f"deep_gemm {getattr(deep_gemm, '__version__', '')} from " + f"{getattr(deep_gemm, '__file__', '?')} does not declare the indexer scorer keyword " + f"arguments {undeclared} in MQA_LOGITS_NUMERICS. Serve with a DeepGEMM build that " + "takes them." + ) + return deep_gemm + + +def _call_deep_gemm_scorer( + deep_gemm_utils: ModuleType, + arguments: dict, + weights: torch.Tensor, + scorer_kwargs: dict, + paged: bool, +) -> torch.Tensor: + """Call DeepGEMM's scorer with ``scorer_kwargs``, which vLLM's wrappers do not pass on.""" + deep_gemm = _scorer_deep_gemm(deep_gemm_utils, scorer_kwargs) + kwargs = {"clean_logits": arguments["clean_logits"], **scorer_kwargs} + if not paged: + return deep_gemm.fp8_fp4_mqa_logits( + arguments["q"], + arguments["kv"], + weights, + arguments["cu_seqlen_ks"], + arguments["cu_seqlen_ke"], + **kwargs, + ) + # Like vLLM's wrapper: DeepGEMM needs a unit last stride, and .contiguous() keeps a size-1 + # dim's stride. + block_tables = arguments["block_tables"] + if block_tables.dim() >= 2 and block_tables.stride(-1) != 1: + block_tables = block_tables.clone(memory_format=torch.contiguous_format) + if arguments.get("indices") is not None: + kwargs["indices"] = arguments["indices"] + return deep_gemm.fp8_fp4_paged_mqa_logits( + arguments["q"], + arguments["kv_cache"], + weights, + arguments["context_lens"], + block_tables, + arguments["schedule_metadata"], + arguments["max_model_len"], + **kwargs, + ) # GLM-5.3-Flash's fused Hadamard + FP8 quantization of the indexer query. @@ -308,12 +495,19 @@ def wrapper(*args, **kwargs): class _QuantVLLMGlm5NextIndexer(_QuantVLLMIndexerBase): """GLM-5.3-Flash indexer: kernels Hadamard-rotate and FP8-quantize the query and pooled keys. - The kernel entry points that write the indexer K cache are wrapped process-wide and re-quantize - the pools they wrote, on prefill and on decode pool completion. The query kernel is swapped for - a re-quantizing wrapper while ``forward`` runs; CUDA graph capture records the wrapper. + The query kernel is swapped for a re-quantizing wrapper while ``forward`` runs; CUDA graph + capture records the wrapper. The indexer op writes the K cache and scores the keys of the layer + that it names, so the op and the kernel entry points it calls are wrapped process-wide: the op + makes its layer's indexer the layer owner (see :mod:`.vllm_layer_scope`), the K-cache writers + re-quantize the pools they wrote, on prefill and on decode pool completion, and vLLM's DeepGEMM + scorer entry points apply the scorer quantizers. """ kpool_ops: ModuleType | None = _glm5next_kpool_ops + # The module whose ``sparse_attn_indexer_kpool`` the indexer op calls. + op_module: ModuleType | None = _glm5next_op_module + # The op imports its DeepGEMM scorer kernels from this module. + deep_gemm_utils: ModuleType | None = vllm_deep_gemm # The native forward looks up its query kernel in this module. indexer_module: ModuleType | None = ( inspect.getmodule(VllmGlm5NextIndexer) if VllmGlm5NextIndexer is not None else None @@ -321,19 +515,36 @@ class _QuantVLLMGlm5NextIndexer(_QuantVLLMIndexerBase): def _setup(self): super()._setup() - assert self.kpool_ops is not None # imported together with the registered indexer class + # Imported together with the registered indexer class. + assert self.kpool_ops is not None and self.op_module is not None + assert self.deep_gemm_utils is not None _glm5next_indexers.add(self) + # The op names the layer by its K cache. + register_layer_owner(self.k_cache.prefix, self) # Every layer calls these; only the first call wraps. - _wrap_kpool_cache_writer( - self.kpool_ops, "kpool_compress_and_write_cache", _kpool_prefill_written - ) - _wrap_kpool_cache_writer( - self.kpool_ops, - "kpool_decode_update_and_maybe_write_cache_batched", - _kpool_decode_written, + wrap_layer_op( + self.op_module, + _GLM5NEXT_INDEXER_OP, + "k_cache_prefix", + decorator=getattr(self.op_module, "eager_break_during_capture", None), ) + for name, written_slots in _KPOOL_CACHE_WRITERS.items(): + wrap_layer_callee( + self.kpool_ops, + name, + _requantize_written_pools(written_slots), + applies=lambda indexer: indexer.indexer_k_quantizer.is_enabled, + outside=_reject_k_cache_write_outside_op, + ) + # Prefill scores keys gathered from the K cache into one buffer, per query range. + _wrap_deep_gemm_scorer(self.deep_gemm_utils, paged=False) + # Decode scores keys read straight from the paged K cache. + _wrap_deep_gemm_scorer(self.deep_gemm_utils, paged=True) def forward(self, *args, **kwargs): + # The scorer only runs for sequences longer than index_topk, which calibration and warmup + # may never reach: fail here rather than on the first long request. + self._validate_scorer_quantizers() if not self.indexer_q_quantizer.is_enabled: return super().forward(*args, **kwargs) module = self.indexer_module @@ -357,6 +568,67 @@ def _quantize_query( q_fp8, q_scale = _fake_quantize_fp8_rows(q_fp8.float() * q_scale, self.indexer_q_quantizer) return q_fp8, q_scale[:, None] + def _validate_scorer_quantizers(self) -> None: + """Reject what the scorer cannot apply, also while calibrating. + + That includes a DeepGEMM that does not take the enabled quantizers' keyword arguments. + """ + if scorer_kwargs := self._scorer_kwargs(quantizing=False): + assert self.deep_gemm_utils is not None # checked in _setup + _scorer_deep_gemm(self.deep_gemm_utils, scorer_kwargs) + + def _scorer_kwargs(self, quantizing: bool = True) -> dict: + """The scorer keyword arguments of the enabled scorer-kwargs quantizers. + + ``quantizing``: only of those that quantize now; calibration runs the stock kernel. + """ + kwargs: dict = {} + for name in self._enabled_scorer_quantizers(): + quantizer = getattr(self, name) + if not isinstance(quantizer, TensorQuantizer): + raise ValueError(f"{name} takes a single format, not a list of formats.") + if quantizing and not quantizer._if_quant: + continue + for key, value in _SCORER_KWARGS_QUANTIZERS[name](quantizer).items(): + if key in kwargs and kwargs[key] != value: + raise ValueError( + f"{name} sets the scorer's {key} to {value!r}, another quantizer to " + f"{kwargs[key]!r}." + ) + kwargs[key] = value + if "logits_dtype" in kwargs: + kwargs["logits_dtype"] = _scorer_logits_dtype(kwargs["logits_dtype"]) + return kwargs + + def _score( + self, + deep_gemm_utils: ModuleType, + original: Callable, + bound: inspect.BoundArguments, + paged: bool, + ) -> torch.Tensor: + """Run the call ``bound`` of vLLM's scorer ``original`` with the scorer quantizers. + + ``weights`` is the effective W: the query scale and the model's normalization are folded in. + """ + arguments = bound.arguments + scorer_kwargs = self._scorer_kwargs() + weights = arguments["weights"] + if scorer_kwargs.get("logits_dtype") == torch.bfloat16: # DeepGEMM's BF16 scorer + weights = weights.to(torch.bfloat16) + if scorer_kwargs: + logits = _call_deep_gemm_scorer( + deep_gemm_utils, arguments, weights, scorer_kwargs, paged + ) + else: + logits = original(*bound.args, **bound.kwargs) + if logits.dtype != torch.float32: + logits = _widen_logits(logits) + # vLLM's prefill top-k returns out-of-range indices for rows with NaN scores, which the + # sparse attention then reads out of bounds; it handles -inf and +inf. An overflowing FP16 + # accumulation or rounding turns scores into NaN (e.g. +inf - inf in the head sum). + return logits.nan_to_num_(nan=float("-inf"), posinf=float("inf"), neginf=float("-inf")) + if VllmDeepseekV4Indexer is not None: QuantModuleRegistry.register({VllmDeepseekV4Indexer: "vllm_DeepseekV4Indexer"})( diff --git a/modelopt/torch/quantization/plugins/vllm_layer_scope.py b/modelopt/torch/quantization/plugins/vllm_layer_scope.py new file mode 100644 index 00000000000..171c3c41800 --- /dev/null +++ b/modelopt/torch/quantization/plugins/vllm_layer_scope.py @@ -0,0 +1,122 @@ +# 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. + +"""Reach a converted module from the functions that vLLM's layer ops call. + +vLLM's layer ops get the name of their layer, not the layer: its ``torch.ops.vllm`` custom ops +cannot take a module argument, and other ops, such as GLM-5.3-Flash's indexer op, follow the same +convention. The kernels they call get only tensors. :func:`register_layer_owner` maps a layer name +to the converted module that acts for the layer, :func:`wrap_layer_op` makes that module the +:func:`layer_owner` while an op that names the layer runs, and :func:`wrap_layer_callee` hands it +the calls of a function that such an op calls. + +The registry and the wraps are process-wide, for one converted model per process: another model +with the same layer names would get its owners. The wraps are permanent. A breakable CUDA graph +replays an op that ``eager_break_during_capture`` decorates by calling the undecorated function, so +the scope goes under that decorator. An op inside a captured graph has its scope while the graph is +captured, so the owner must launch only work that the graph can replay. +""" + +import contextvars +import functools +import inspect +import weakref +from collections.abc import Callable +from types import ModuleType +from typing import Any + +__all__ = ["layer_owner", "register_layer_owner", "wrap_layer_callee", "wrap_layer_op"] + +# Layer name -> the converted module that acts for the layer. Weak: it keeps no module alive. +_owners: weakref.WeakValueDictionary = weakref.WeakValueDictionary() +# Inside an op of a layer that no converted module owns. +_NO_OWNER = object() +# The owner of the layer whose op runs: None outside a wrapped op. +_layer_owner: contextvars.ContextVar = contextvars.ContextVar("_layer_owner", default=None) + + +def register_layer_owner(name: str, owner: Any) -> None: + """Make ``owner`` act for the layer ``name``, as the layer's ops name it; the latest wins.""" + _owners[name] = owner + + +def layer_owner() -> Any: + """The owner of the layer whose op runs, or None.""" + owner = _layer_owner.get() + return None if owner is _NO_OWNER else owner + + +def wrap_layer_op( + op_module: ModuleType, name: str, layer_arg: str, decorator: Callable | None = None +) -> None: + """Wrap the op ``op_module.`` to put the owner of the layer it names in scope. + + ``layer_arg`` is the op's argument with the layer name. If ``decorator`` decorates the op, the + wrap goes under it and ``decorator`` decorates the wrapper again. + """ + original = getattr(op_module, name) + if getattr(original, "_modelopt_layer_op", False): + return + inner = getattr(original, "__wrapped__", original) if decorator is not None else original + position = list(inspect.signature(inner).parameters).index(layer_arg) + + @functools.wraps(inner) + def wrapper(*args, **kwargs): + layer = kwargs[layer_arg] if layer_arg in kwargs else args[position] + # A str or vLLM's LayerName. + token = _layer_owner.set(_owners.get(getattr(layer, "value", layer), _NO_OWNER)) + try: + return inner(*args, **kwargs) + finally: + _layer_owner.reset(token) + + if inner is not original: + wrapper = decorator(wrapper) # type: ignore[misc] + wrapper._modelopt_layer_op = True # type: ignore[attr-defined] + setattr(op_module, name, wrapper) + + +def wrap_layer_callee( + module: ModuleType, + name: str, + call: Callable[[Any, Callable, inspect.BoundArguments], Any], + applies: Callable[[Any], bool], + outside: Callable[[], None] | None = None, +) -> None: + """Wrap ``module.``, a function that ops call, to hand its calls to the layer's owner. + + Inside an op that :func:`wrap_layer_op` wrapped, a call for whose owner ``applies(owner)`` + holds runs ``call(owner, original, bound)`` with the bound arguments of the call; other calls + run the original function. ``outside()`` runs before a call outside such an op, e.g. to reject + it. + """ + original = getattr(module, name) + if getattr(original, "_modelopt_layer_callee", False): + return + signature = inspect.signature(original) + + @functools.wraps(original) + def wrapper(*args, **kwargs): + owner = _layer_owner.get() + if owner is None and outside is not None: + outside() + if owner is None or owner is _NO_OWNER or not applies(owner): + return original(*args, **kwargs) + bound = signature.bind(*args, **kwargs) + bound.apply_defaults() + return call(owner, original, bound) + + wrapper._modelopt_layer_callee = True # type: ignore[attr-defined] + setattr(module, name, wrapper) diff --git a/tests/gpu_vllm/torch/quantization/test_vllm_dynamic_modules.py b/tests/gpu_vllm/torch/quantization/test_vllm_dynamic_modules.py index 8286272bb6b..9cc32bdcb27 100644 --- a/tests/gpu_vllm/torch/quantization/test_vllm_dynamic_modules.py +++ b/tests/gpu_vllm/torch/quantization/test_vllm_dynamic_modules.py @@ -28,6 +28,8 @@ from __future__ import annotations +import copy +import functools import gc import importlib.util import inspect @@ -48,6 +50,7 @@ from vllm import LLM, ModelRegistry, SamplingParams from vllm.distributed import cleanup_dist_env_and_memory from vllm.inputs import TokensPrompt +from vllm.utils.deep_gemm import _import_deep_gemm from vllm.utils.import_utils import has_deep_gemm import modelopt.torch.quantization as mtq @@ -56,6 +59,7 @@ from modelopt.torch.quantization.conversion import set_quantizer_by_cfg from modelopt.torch.quantization.nn import SequentialQuantizer, TensorQuantizer from modelopt.torch.quantization.plugins import vllm as vllm_plugin +from modelopt.torch.quantization.plugins import vllm_indexer from modelopt.torch.quantization.plugins.vllm import ( _ATTENTION_TYPES, VllmMLAAttention, @@ -66,7 +70,10 @@ configure_vllm_nvfp4_attention_quantizers, disable_compilation, ) -from modelopt.torch.quantization.plugins.vllm_indexer import _QuantVLLMIndexerBase +from modelopt.torch.quantization.plugins.vllm_indexer import ( + _QuantVLLMGlm5NextIndexer, + _QuantVLLMIndexerBase, +) def _load_example_module(name: str): @@ -788,7 +795,7 @@ def _missing(module, name, slots): } -def _boot_llm(model_dir, max_model_len=64, **extra): +def _boot_llm(model_dir, max_model_len=64, max_num_seqs=1, **extra): """Construct a vLLM engine on a tiny model. MoE fixtures override with ``moe_backend="triton"`` (pins the Triton @@ -801,7 +808,7 @@ def _boot_llm(model_dir, max_model_len=64, **extra): enforce_eager=True, gpu_memory_utilization=0.2, max_model_len=max_model_len, - max_num_seqs=1, + max_num_seqs=max_num_seqs, dtype="bfloat16", skip_tokenizer_init=True, **extra, @@ -887,13 +894,23 @@ def tiny_deepseek_llm(tmp_path_factory): "algorithm": "max", } +# GLM-5.3-Flash scores the keys only for sequences longer than index_topk (2048 tokens); two +# requests share a prefill chunk. +_GLM5_NEXT = ( + "Glm5NextForCausalLM", + create_tiny_glm5_next_config_dir, + { + "load_format": "dummy", + "max_model_len": 2304, + "max_num_batched_tokens": 2304, + "max_num_seqs": 2, + }, +) # model -> (architecture vLLM must know, tiny checkpoint builder, extra LLM kwargs) _SPARSE_ATTN_MODELS = { - "glm5_next": ( - "Glm5NextForCausalLM", - create_tiny_glm5_next_config_dir, - {"load_format": "dummy"}, - ), + "glm5_next": _GLM5_NEXT, + # In-kernel scorer numerics, from a DeepGEMM build with them on the Python path. + "glm5_next_numerics": _GLM5_NEXT, # DeepSeek-V4 computes the indexer query only when top-k selection is needed, i.e. beyond # compress_ratio * index_topk = 4 * 1024 tokens. "deepseek_v4": ( @@ -914,6 +931,14 @@ def tiny_sparse_attn_llm(request, tmp_path_factory): pytest.skip("vLLM's sparse-attention indexer needs DeepGEMM") if torch.cuda.get_device_capability()[0] not in (9, 10): pytest.skip("vLLM's sparse-attention indexer backends need Hopper or Blackwell") + if request.param == "glm5_next_numerics": + if not getattr(_import_deep_gemm(), "MQA_LOGITS_NUMERICS", None): + pytest.skip( + "vLLM's DeepGEMM has no scorer numerics: install a build with them or put one " + "first on PYTHONPATH" + ) + if torch.cuda.get_device_capability()[0] != 10: + pytest.skip("the in-kernel scorer numerics need an SM100 GPU") llm = _boot_llm(build(tmp_path_factory.mktemp(request.param)), **extra) try: yield llm @@ -944,21 +969,37 @@ def _indexer_cache_rows(kv_cache, mask): return values, scales.contiguous().view(torch.int32).squeeze(-1) -def _calibrate_and_clip_indexer_k(self): +def _glm5_next_scorer_cfg(in_kernel): + """GLM-5.3-Flash scorer kernel arguments: in-kernel numerics, or DeepGEMM's BF16 head reduction. + + Both run inside DeepGEMM's scorer, which applies them on SM100 GPUs only; elsewhere, none. + """ + if in_kernel: # BF16 post-ReLU products and FP16 Q.K accumulation + scorer_kwargs = {"mma_accum": "fp16", "post_relu": "bf16"} + elif torch.cuda.get_device_capability()[0] == 10: # stock DeepGEMM has it too + scorer_kwargs = {"logits_dtype": "bfloat16"} + else: + return [] + cfg = {"backend": "scorer_kwargs", "backend_extra_args": scorer_kwargs} + return [{"quantizer_name": "*indexer_scorer_kwargs_quantizer", "cfg": cfg}] + + +def _calibrate_and_clip_indexer_k(self, cfg): """Run on the worker: calibrate FP8 indexer q and K quantizers, then clip K to 1/8 of its amax. Calibration goes through real scheduled prefills: the fused indexers write their cache only when ``attn_metadata`` is set, which a dummy run does not do. The second prompt fills the - context, so that DeepSeek-V4 selects top-k and computes its query. + context, so that DeepSeek-V4 selects top-k and computes its query. From then on, the calls of + the GLM-5.3-Flash scorer quantizers are recorded. """ model = self.get_model() lengths = (40, self.model_config.max_model_len - 8) batches = [{"input_ids": torch.randint(1, 100, (1, n))} for n in lengths] forward_loop = _load_example_module("vllm_ptq_utils").calibrate_fun(batches, self) with disable_compilation(model): - mtq.quantize(model, _INDEXER_FP8_CFG, forward_loop=forward_loop) + mtq.quantize(model, cfg, forward_loop=forward_loop) - amaxes, self.indexer_k_snapshots = {"k": {}, "q": {}}, {} + amaxes, self.indexer_k_snapshots, self.indexer_scorer_calls = {"k": {}, "q": {}}, {}, [] for name, module in model.named_modules(): if isinstance(module, _QuantVLLMIndexerBase): for kind in amaxes: @@ -967,9 +1008,38 @@ def _calibrate_and_clip_indexer_k(self): if amaxes["k"][name]: module.indexer_k_quantizer.amax = module.indexer_k_quantizer.amax / 8 self.indexer_k_snapshots[name] = _cache_row_signature(module.k_cache.kv_cache) + if isinstance(module, _QuantVLLMGlm5NextIndexer): + module._score = functools.partial( + _record_scorer_call, module._score, name, self.indexer_scorer_calls + ) return amaxes +def _record_scorer_call(score, name, calls, deep_gemm_utils, original, bound, paged): + """Spy on ``_QuantVLLMGlm5NextIndexer._score``: the layer, the path, what top-k reads and the + in-kernel numerics passed to DeepGEMM.""" + passed = {} + call_scorer = vllm_indexer._call_deep_gemm_scorer + + def spy(deep_gemm_utils, arguments, weights, scorer_kwargs, paged): + passed.update(scorer_kwargs) + return call_scorer(deep_gemm_utils, arguments, weights, scorer_kwargs, paged) + + vllm_indexer._call_deep_gemm_scorer = spy + try: + logits = score(deep_gemm_utils, original, bound, paged) + finally: + vllm_indexer._call_deep_gemm_scorer = call_scorer + numerics = {key: passed[key] for key in ("mma_accum", "post_relu") if key in passed} + calls.append((name, paged, str(logits.dtype), bool(logits.isnan().any()), numerics)) + return logits + + +def _indexer_scorer_calls(self): + """Run on the worker: the scorer calls recorded since the calibration.""" + return self.indexer_scorer_calls + + def _indexer_k_rows_written(self): """Run on the worker: rows the indexer kernels wrote since the snapshot, max over the clip.""" torch.cuda.synchronize() @@ -1172,17 +1242,32 @@ def test_tiny_deepseek_mla_quantize(tiny_deepseek_llm): @pytest.mark.timeout(600) # engine boot and the DeepGEMM JIT dominate -def test_tiny_sparse_attn_indexer_quantize(tiny_sparse_attn_llm): - """The indexer query is fake-quantized and the K cache the kernels read holds QDQ keys.""" - amaxes = tiny_sparse_attn_llm.collective_rpc(_calibrate_and_clip_indexer_k)[0] +def test_tiny_sparse_attn_indexer_quantize(tiny_sparse_attn_llm, request): + """The indexer query is fake-quantized and the K cache the kernels read holds QDQ keys. + + On GLM-5.3-Flash on SM100 the scorer quantizers run too: a prompt longer than index_topk reaches + the DeepGEMM scorer in prefill, sharing its chunk with a short prompt, and in decode. With + ``glm5_next_numerics``, DeepGEMM's scorer kernel rounds the post-ReLU products to BF16 and + accumulates Q.K in FP16, which needs a DeepGEMM build with these numerics on the Python path. + """ + model = request.node.callspec.params["tiny_sparse_attn_llm"] + glm, in_kernel = model.startswith("glm5_next"), model == "glm5_next_numerics" + cfg = copy.deepcopy(_INDEXER_FP8_CFG) + scorer_cfg = _glm5_next_scorer_cfg(in_kernel) if glm else [] + cfg["quant_cfg"] += scorer_cfg + amaxes = tiny_sparse_attn_llm.collective_rpc(_calibrate_and_clip_indexer_k, args=(cfg,))[0] assert amaxes["k"], "no indexer was converted" for kind_amaxes in amaxes.values(): # calibrated: the quantizer saw the kernels' tensors assert all(a is not None and 0 < a < float("inf") for a in kind_amaxes.values()), amaxes # Prefill plus decode steps that complete further pools / compression groups. - prompts = [TokensPrompt(prompt_token_ids=list(range(1 + i, 41 + i))) for i in range(2)] + lengths = (2100, 40) if glm else (40, 40) + prompts = [ + TokensPrompt(prompt_token_ids=list(range(1 + i, n + 1 + i))) for i, n in enumerate(lengths) + ] params = SamplingParams(max_tokens=12, ignore_eos=True, temperature=0.0, detokenize=False) - tiny_sparse_attn_llm.generate(prompts, params) + outputs = tiny_sparse_attn_llm.generate(prompts, params) + assert [len(output.outputs[0].token_ids) for output in outputs] == [12, 12] for name, (rows, max_over_clip) in tiny_sparse_attn_llm.collective_rpc(_indexer_k_rows_written)[ 0 @@ -1191,6 +1276,17 @@ def test_tiny_sparse_attn_indexer_quantize(tiny_sparse_attn_llm): # Re-storing a clipped key in the FP8 cache rounds it up by at most 2**-4. assert max_over_clip <= 1.07, (name, rows, max_over_clip) + if scorer_cfg: + calls = tiny_sparse_attn_llm.collective_rpc(_indexer_scorer_calls)[0] + # Every indexer scored prefill and decode rows; vLLM's top-k reads FP32 without NaN. + assert {(name, paged) for name, paged, *_ in calls} == { + (name, paged) for name in amaxes["k"] for paged in (False, True) + }, calls + assert all(dtype == "torch.float32" and not nan for _, _, dtype, nan, _ in calls), calls + # The scorer passed the recipe's in-kernel numerics to DeepGEMM. + numerics = {"mma_accum": "fp16", "post_relu": "bf16"} if in_kernel else {} + assert all(kernel == numerics for *_, kernel in calls), calls + def _quantize_kv_preset_and_summarize(self): """Worker RPC: ``KV_QUANT_CFG=NVFP4_KV_CFG`` as FakeQuantWorker resolves it, fold, one more diff --git a/tests/gpu_vllm/torch/quantization/test_vllm_indexer_quant.py b/tests/gpu_vllm/torch/quantization/test_vllm_indexer_quant.py index e7f95811337..b60bddaf26c 100644 --- a/tests/gpu_vllm/torch/quantization/test_vllm_indexer_quant.py +++ b/tests/gpu_vllm/torch/quantization/test_vllm_indexer_quant.py @@ -13,14 +13,16 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Fake quantization of the sparse-attention indexer query and K cache in the vLLM plugin. +"""Fake quantization of the sparse-attention indexer query, K cache and scorer in the vLLM plugin. Covers the FP8 cache read-back/re-quantization shared by the fused indexer layouts, the query -re-quantization and the wiring of each layout adapter on stand-in modules, without booting an -``LLM`` (see ``test_vllm_dynamic_modules.py`` for the end-to-end DeepSeek-V4 and GLM-5.3-Flash -runs). +re-quantization, the scorer quantizers and the wiring of each layout adapter on stand-in modules, +without booting an ``LLM`` (see ``test_vllm_dynamic_modules.py`` for the end-to-end DeepSeek-V4 and +GLM-5.3-Flash runs). """ +import functools +import inspect import sys import weakref from types import SimpleNamespace @@ -30,8 +32,9 @@ import modelopt.torch.quantization as mtq from modelopt.torch.quantization.config import QuantizerAttributeConfig -from modelopt.torch.quantization.nn import QuantModuleRegistry, TensorQuantizer -from modelopt.torch.quantization.plugins import vllm_indexer +from modelopt.torch.quantization.conversion import set_quantizer_by_cfg +from modelopt.torch.quantization.nn import QuantModuleRegistry, SequentialQuantizer, TensorQuantizer +from modelopt.torch.quantization.plugins import vllm_indexer, vllm_layer_scope from modelopt.torch.utils.distributed import ParallelState HEAD_DIM = 128 @@ -312,9 +315,9 @@ def fwht128_quant_fp8(q): class _NativeGlm5NextIndexer(torch.nn.Module): - def __init__(self, kv_cache=None): + def __init__(self, kv_cache=None, prefix="layers.0.indexer.k_cache"): super().__init__() - self.k_cache = SimpleNamespace(kv_cache=kv_cache) + self.k_cache = SimpleNamespace(kv_cache=kv_cache, prefix=prefix) def forward(self, q): # Like vLLM's forward, look the query kernel up in the module globals on every call. @@ -326,11 +329,18 @@ class _TestGlm5NextIndexer(vllm_indexer._QuantVLLMGlm5NextIndexer, _NativeGlm5Ne def test_glm5next_kpool_cache_writers_requantize_written_pools( - parallel_state_stub, written_cache, monkeypatch + parallel_state_stub, written_cache, scorer_stand_ins, monkeypatch ): + """In the indexer op of a layer with the quantizer, the writers re-quantize what they wrote.""" kv_cache, slot_mapping = written_cache before = kv_cache.clone() - kernel_calls = [] + kernel_calls, requantized = [], [] + requantize = vllm_indexer._requantize_fp8_indexer_k_cache + monkeypatch.setattr( + vllm_indexer, + "_requantize_fp8_indexer_k_cache", + lambda kv_cache, *args: requantized.append(kv_cache) or requantize(kv_cache, *args), + ) def kpool_compress_and_write_cache( kv_cache, slot_k, slot_score, ape, loc, pool_size, head_dim=128, write_mask=None, **kwargs @@ -359,44 +369,58 @@ def kpool_decode_update_and_maybe_write_cache_batched( ), ) monkeypatch.setattr(_TestGlm5NextIndexer, "kpool_ops", kpool_ops) - monkeypatch.setattr(vllm_indexer, "_glm5next_indexers", weakref.WeakSet()) indexer = _TestGlm5NextIndexer.convert(_NativeGlm5NextIndexer(kv_cache)) # A second layer's conversion must not wrap the writers again. - other = _TestGlm5NextIndexer.convert(_NativeGlm5NextIndexer(torch.zeros_like(kv_cache))) + other_cache = torch.zeros_like(kv_cache) + other = _TestGlm5NextIndexer.convert( + _NativeGlm5NextIndexer(other_cache, prefix="layers.1.indexer.k_cache") + ) other.indexer_k_quantizer.disable() assert kpool_ops.kpool_compress_and_write_cache.__wrapped__ is kpool_compress_and_write_cache + op = scorer_stand_ins.op_module.sparse_attn_indexer_kpool + + def write(layer, writer, *args, **kwargs): + """``writer(*args, **kwargs)`` in the indexer op of ``layer``, like vLLM's op.""" + return op(None, f"layers.{layer}.indexer.k_cache", None, lambda: writer(*args, **kwargs)) - # No enabled quantizer: the kernel runs and the wrapper returns before looking up the cache. + # No enabled quantizer: the kernel runs, also outside the indexer op, and nothing is + # re-quantized. 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) + prefill = (kv_cache, None, None, None, slot_mapping, 4) + write(0, kpool_ops.kpool_compress_and_write_cache, *prefill) + kpool_ops.kpool_compress_and_write_cache(*prefill) + assert not requantized quantizer = indexer.indexer_k_quantizer = _fp8_quantizer(amax=2.0) - # Prefill: pools at ``loc`` masked by ``write_mask``; an unrelated cache is left alone. - other_cache = torch.zeros_like(kv_cache) - kpool_ops.kpool_compress_and_write_cache(other_cache, None, None, None, slot_mapping, 4) - assert torch.equal(other_cache, torch.zeros_like(kv_cache)) + # A write outside the indexer op names no layer, and one to another layer's cache is misrouted: + # both fail before the kernel runs. A layer without the quantizer re-quantizes nothing. + with pytest.raises(RuntimeError, match="outside the indexer op"): + kpool_ops.kpool_compress_and_write_cache(*prefill) + prefill_other = (other_cache, None, None, None, slot_mapping, 4) + with pytest.raises(RuntimeError, match="writes the K cache of another layer"): + write(0, kpool_ops.kpool_compress_and_write_cache, *prefill_other) + write(1, kpool_ops.kpool_compress_and_write_cache, *prefill_other) + assert not requantized + + # Prefill: pools at ``loc`` masked by ``write_mask``. write_mask = torch.tensor([True, True, False, True, True, True, True, True]) - kpool_ops.kpool_compress_and_write_cache( - kv_cache, None, None, None, slot_mapping, 4, write_mask=write_mask - ) + write(0, kpool_ops.kpool_compress_and_write_cache, *prefill, write_mask=write_mask) _assert_requantized(before, kv_cache, torch.tensor([9, 18, 30, 3]), quantizer) # Decode: ``[num_requests, next_n]`` slots, written only where a pool completes. before = kv_cache.clone() dec_slots = torch.tensor([[17, 9], [30, -1]], dtype=torch.int32) dec_pos = torch.tensor([[7, 8], [11, 15]], dtype=torch.int32) - kpool_ops.kpool_decode_update_and_maybe_write_cache_batched( - kv_cache, None, None, None, None, None, dec_slots, dec_pos, 4 + write( + 0, + kpool_ops.kpool_decode_update_and_maybe_write_cache_batched, + *(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", "prefill", "decode"] - assert getattr(kpool_ops.kpool_compress_and_write_cache, "_modelopt_indexer_k_wrapped", False) + assert kernel_calls == ["prefill"] * 4 + ["decode"] + assert len(requantized) == 2 and all(cache is kv_cache for cache in requantized) + assert getattr(kpool_ops.kpool_compress_and_write_cache, "_modelopt_layer_callee", False) def test_glm5next_indexer_requantizes_query(parallel_state_stub, monkeypatch): @@ -406,7 +430,6 @@ def test_glm5next_indexer_requantizes_query(parallel_state_stub, monkeypatch): ) monkeypatch.setattr(_TestGlm5NextIndexer, "kpool_ops", kpool_ops) monkeypatch.setattr(_TestGlm5NextIndexer, "indexer_module", sys.modules[__name__]) - monkeypatch.setattr(vllm_indexer, "_glm5next_indexers", weakref.WeakSet()) indexer = _TestGlm5NextIndexer.convert(_NativeGlm5NextIndexer()) indexer.indexer_k_quantizer.disable() torch.manual_seed(0) @@ -429,3 +452,571 @@ def test_glm5next_indexer_requantizes_query(parallel_state_stub, monkeypatch): monkeypatch.setattr(_TestGlm5NextIndexer, "indexer_module", SimpleNamespace()) with pytest.raises(NotImplementedError, match="fwht128_quant_fp8"): indexer(q) + + +# Indexer scorer quantizers (GLM-5.3-Flash). DeepGEMM and the indexer op are stand-ins: the +# scorer numerics need a DeepGEMM build that provides them, and the op a running engine. + +NUM_HEADS = 32 +NUMERICS = { + "version": 1, + "mma_accum": ("fp32", "fp16"), + "post_relu": ("none", "fp32", "bf16", "fp16", "fp16_dynamic", "mxint8"), + "reduce": ("fp32", "fp16"), +} +requires_cuda = pytest.mark.skipif( + not torch.cuda.is_available(), reason="the scorer quantizers run on the GPU" +) + + +class _FakeDeepGemm: + """DeepGEMM's scorers and vLLM's wrappers of them, returning the preset ``scores``. + + Keys outside each row's ``valid`` range hold garbage, as DeepGEMM leaves them with + ``clean_logits=False``. Every call records its weights, keyword arguments and logits. + """ + + def __init__(self): + self.calls = [] + self.scores = self.valid = None + self.module = SimpleNamespace( + MQA_LOGITS_NUMERICS=NUMERICS, + fp8_fp4_mqa_logits=self._mqa_logits, + fp8_fp4_paged_mqa_logits=self._paged_mqa_logits, + ) + self.utils = SimpleNamespace( + fp8_fp4_mqa_logits=self._vllm_mqa_logits, + fp8_fp4_paged_mqa_logits=self._vllm_paged_mqa_logits, + _import_deep_gemm=lambda: self.module, + _lazy_init=lambda: None, + ) + + def _logits(self, weights, block_table=None, **kwargs): + rows, keys = self.scores.shape + dtype = kwargs.get("logits_dtype", torch.float32) + garbage = torch.where(torch.arange(keys, device="cuda") % 2 == 0, float("nan"), 1e30) + # A view of a buffer with padded rows, like DeepGEMM's. + logits = torch.empty(rows, keys + 9, dtype=dtype, device="cuda")[:, :keys] + logits.copy_(torch.where(self.valid, self.scores, garbage)) + self.calls.append( + SimpleNamespace( + weights=weights, kwargs=kwargs, block_table=block_table, logits=logits.clone() + ) + ) + return logits + + def _mqa_logits(self, q, kv, weights, cu_seq_len_k_start, cu_seq_len_k_end, **kwargs): + return self._logits(weights, **kwargs) + + def _paged_mqa_logits( + self, q, kv_cache, weights, context_lens, block_table, schedule_meta, max_context_len, **kw + ): + return self._logits(weights, block_table=block_table, **kw) + + def _vllm_mqa_logits(self, q, kv, weights, cu_seqlen_ks, cu_seqlen_ke, clean_logits): + return self._mqa_logits( + q, kv, weights, cu_seqlen_ks, cu_seqlen_ke, clean_logits=clean_logits + ) + + def _vllm_paged_mqa_logits( + self, + q, + kv_cache, + weights, + context_lens, + block_tables, + schedule_metadata, + max_model_len, + clean_logits, + indices=None, + ): + kwargs = {} if indices is None else {"indices": indices} + return self._paged_mqa_logits( + q, + kv_cache, + weights, + context_lens, + block_tables, + schedule_metadata, + max_model_len, + clean_logits=clean_logits, + **kwargs, + ) + + +def _nan_to_minus_inf(logits): + """The scorer maps NaN, which vLLM's prefill top-k mishandles, to -inf; +-inf stay.""" + return logits.nan_to_num(nan=float("-inf"), posinf=float("inf"), neginf=float("-inf")) + + +def _eager_break(fn): + """Stand-in for vLLM's ``eager_break_during_capture``: graph replays call ``fn`` itself.""" + + @functools.wraps(fn) + def wrapper(*args, **kwargs): + return fn(*args, **kwargs) + + return wrapper + + +def _indexer_op(hidden_states, k_cache_prefix, kv_cache, run_scorer): + """Stand-in for GLM-5.3-Flash's indexer op, which looks its DeepGEMM scorer up per call.""" + return run_scorer() + + +@pytest.fixture(autouse=True) +def scorer_stand_ins(monkeypatch): + """Converted GLM-5.3-Flash indexers wrap these instead of the installed vLLM's modules. + + Each test starts with no converted indexers and an empty registry of layer owners. + """ + monkeypatch.setattr(vllm_indexer, "_glm5next_indexers", weakref.WeakSet()) + monkeypatch.setattr(vllm_layer_scope, "_owners", weakref.WeakValueDictionary()) + fake = _FakeDeepGemm() + op_module = SimpleNamespace( + sparse_attn_indexer_kpool=_eager_break(_indexer_op), + eager_break_during_capture=_eager_break, + ) + monkeypatch.setattr(_TestGlm5NextIndexer, "op_module", op_module) + monkeypatch.setattr(_TestGlm5NextIndexer, "deep_gemm_utils", fake.utils) + return SimpleNamespace(fake=fake, op_module=op_module) + + +@pytest.fixture +def glm(parallel_state_stub, scorer_stand_ins, monkeypatch): + """Converts GLM-5.3-Flash indexers (q and K quantizers off) and runs their indexer op.""" + kpool_ops = SimpleNamespace( + kpool_compress_and_write_cache=lambda kv_cache, loc: None, + kpool_decode_update_and_maybe_write_cache_batched=lambda kv_cache, slot_mapping: None, + ) + monkeypatch.setattr(_TestGlm5NextIndexer, "kpool_ops", kpool_ops) + + def convert(layer=0): + native = _NativeGlm5NextIndexer(prefix=f"layers.{layer}.indexer.k_cache") + indexer = _TestGlm5NextIndexer.convert(native) + indexer.indexer_q_quantizer.disable() + indexer.indexer_k_quantizer.disable() + return indexer + + def score(run_scorer, layer=0): + op = scorer_stand_ins.op_module.sparse_attn_indexer_kpool + return op(None, f"layers.{layer}.indexer.k_cache", None, run_scorer) + + return SimpleNamespace(fake=scorer_stand_ins.fake, convert=convert, score=score) + + +def _scorer_call(fake, paged, num_heads=NUM_HEADS, q_scale=None, block_tables=None, indices=None): + """Valid keys and the scorer call of the indexer op, on prefill or decode rows.""" + if paged: # rows score keys [0, context length) of their request's pages + context_lens = torch.tensor([[33, 64], [1, 100]], dtype=torch.int32, device="cuda") + start, end, num_keys = torch.zeros_like(context_lens).reshape(-1), context_lens, 128 + else: # rows of several requests score their ranges [ks, ke) of one gathered key buffer + start = torch.tensor([0, 0, 37, 37, 5], dtype=torch.int32, device="cuda") + end, num_keys = torch.tensor([70, 31, 133, 37, 140], dtype=torch.int32, device="cuda"), 140 + end = end.reshape(-1) + col = torch.arange(num_keys, device="cuda") + fake.valid = (col >= start[:, None]) & (col < end[:, None]) + fake.scores = torch.randn(fake.valid.shape, device="cuda") * 3 + weights = torch.randn(len(start), num_heads, device="cuda") + q = (None, q_scale) + utils = fake.utils + if paged: + if block_tables is None: + block_tables = torch.zeros(2, 4, dtype=torch.int32, device="cuda") + + def run_scorer(): + return utils.fp8_fp4_paged_mqa_logits( + q, None, weights, context_lens, block_tables, None, num_keys, False, indices + ) + else: + + def run_scorer(): + return utils.fp8_fp4_mqa_logits(q, None, weights, start, end, clean_logits=False) + + return SimpleNamespace(run=run_scorer, weights=weights, start=start, end=end) + + +def _kwargs_quantizer(**scorer_kwargs): + """``indexer_scorer_kwargs_quantizer`` that passes ``scorer_kwargs`` to the scorer kernel.""" + return TensorQuantizer( + QuantizerAttributeConfig(backend="scorer_kwargs", backend_extra_args=scorer_kwargs) + ) + + +def test_scorer_kwargs_quantizer_passes_its_backend_extra_args(glm): + indexer = glm.convert() + scorer_kwargs = {"mma_accum": "fp16", "post_relu": "bf16"} + cfg = {"backend": "scorer_kwargs", "backend_extra_args": scorer_kwargs} + set_quantizer_by_cfg( + indexer, [{"quantizer_name": "*indexer_scorer_kwargs_quantizer", "cfg": cfg}] + ) + assert indexer._scorer_kwargs() == scorer_kwargs + + indexer.indexer_scorer_kwargs_quantizer.disable_quant() # calibration runs the stock kernel + assert indexer._scorer_kwargs() == {} + + # Stock DeepGEMM's own logits_dtype, by its name in a config. + indexer.indexer_scorer_kwargs_quantizer = _kwargs_quantizer(logits_dtype="bfloat16") + assert indexer._scorer_kwargs() == {"logits_dtype": torch.bfloat16} + + +def test_registered_quantizers_turn_their_config_into_scorer_kwargs(glm, monkeypatch): + """E.g. a plugin's quantizer for numerics options of a DeepGEMM build.""" + + def post_relu(quantizer): + if quantizer.num_bits != (8, 7): + raise ValueError("BF16 only") + return {"post_relu": "bf16"} + + registered = dict(vllm_indexer._SCORER_KWARGS_QUANTIZERS) + monkeypatch.setattr(vllm_indexer, "_SCORER_KWARGS_QUANTIZERS", registered) + vllm_indexer.register_indexer_scorer_kwargs_quantizer("indexer_test_quantizer", post_relu) + indexer = glm.convert() + assert not indexer.indexer_test_quantizer.is_enabled + indexer.indexer_test_quantizer = TensorQuantizer( + QuantizerAttributeConfig(num_bits=(8, 7), backend="scorer_kwargs") + ) + indexer.indexer_scorer_kwargs_quantizer = _kwargs_quantizer(mma_accum="fp16") + assert indexer._scorer_kwargs() == {"mma_accum": "fp16", "post_relu": "bf16"} + + indexer.indexer_scorer_kwargs_quantizer = _kwargs_quantizer(post_relu="fp16") + with pytest.raises(ValueError, match=r"indexer_test_quantizer sets the scorer's post_relu to"): + indexer.modelopt_post_restore() + indexer.indexer_scorer_kwargs_quantizer.disable() + indexer.indexer_test_quantizer = TensorQuantizer( + QuantizerAttributeConfig(num_bits=(5, 10), backend="scorer_kwargs") + ) + with pytest.raises(ValueError, match="BF16 only"): + indexer.modelopt_post_restore() + + +@pytest.mark.parametrize( + ("quantizer", "match"), + [ + ( + lambda: TensorQuantizer(QuantizerAttributeConfig(num_bits=8)), + "needs backend 'scorer_kwargs', got None", + ), + (lambda: _kwargs_quantizer(logits_dtype="float16"), "float32 or bfloat16"), + ( + lambda: SequentialQuantizer(_kwargs_quantizer(), _kwargs_quantizer()), + "a single format", + ), + ], + ids=["backend", "logits_dtype", "list"], +) +def test_scorer_kwargs_quantizer_rejects_unsupported_config(glm, quantizer, match): + indexer = glm.convert() + indexer.indexer_scorer_kwargs_quantizer = quantizer() + with pytest.raises(ValueError, match=match): + indexer.modelopt_post_restore() + + +def test_glm5next_forward_validates_scorer_quantizers(glm): + """Bad configs and a DeepGEMM without the arguments fail in any forward, also in calibration. + + The scorer itself only runs for sequences longer than index_topk. + """ + indexer = glm.convert() + q = torch.randn(2, HEAD_DIM) + quantizer = indexer.indexer_scorer_kwargs_quantizer = _kwargs_quantizer(post_relu="mxint8") + quantizer.disable_quant() # calibrating: the scorer would run the stock kernel + indexer(q) + del glm.fake.module.MQA_LOGITS_NUMERICS # a DeepGEMM without the post_relu argument + with pytest.raises(RuntimeError, match=r"\{'post_relu': 'mxint8'\} in MQA_LOGITS_NUMERICS"): + indexer(q) + indexer.indexer_scorer_kwargs_quantizer = _kwargs_quantizer(logits_dtype="bfloat16") + indexer(q) # stock DeepGEMM's own argument + glm.fake.module.MQA_LOGITS_NUMERICS = {**NUMERICS, "post_relu": ("none",)} + indexer.indexer_scorer_kwargs_quantizer = _kwargs_quantizer(post_relu="mxint8") + with pytest.raises(RuntimeError, match=r"\{'post_relu': 'mxint8'\}"): + indexer(q) + indexer.indexer_scorer_kwargs_quantizer = _kwargs_quantizer(logits_dtype="float16") + with pytest.raises(ValueError, match="float32 or bfloat16"): + indexer(q) + + +@requires_cuda +@pytest.mark.parametrize("paged", [False, True], ids=["prefill", "decode"]) +def test_scorer_passes_through_without_enabled_quantizer(glm, paged): + indexer = glm.convert() + call = _scorer_call(glm.fake, paged) + logits = glm.score(call.run) + assert glm.fake.calls[-1].weights is call.weights + assert set(glm.fake.calls[-1].kwargs) <= {"clean_logits", "indices"} + assert torch.equal(logits.view(torch.int32), glm.fake.calls[-1].logits.view(torch.int32)) + + # Outside an indexer op there is no calling layer, whatever its quantizers. + indexer.indexer_scorer_kwargs_quantizer = _kwargs_quantizer(mma_accum="fp16") + logits = call.run() + assert set(glm.fake.calls[-1].kwargs) <= {"clean_logits", "indices"} + assert torch.equal(logits.view(torch.int32), glm.fake.calls[-1].logits.view(torch.int32)) + + +@requires_cuda +@pytest.mark.parametrize("paged", [False, True], ids=["prefill", "decode"]) +def test_bf16_reduction_widens_scores_to_fp32(glm, paged): + """DeepGEMM's BF16 head reduction takes BF16 weights and writes BF16 scores.""" + indexer = glm.convert() + indexer.indexer_scorer_kwargs_quantizer = _kwargs_quantizer(logits_dtype="bfloat16") + del glm.fake.module.MQA_LOGITS_NUMERICS # stock DeepGEMM has the BF16 reduction + call = _scorer_call(glm.fake, paged) + logits = glm.score(call.run) + + kernel = glm.fake.calls[-1] + assert kernel.kwargs["logits_dtype"] == torch.bfloat16 + assert torch.equal(kernel.weights, call.weights.to(torch.bfloat16)) + # vLLM's top-k reads FP32 scores with the strides of the kernel's output. + assert logits.dtype == torch.float32 + assert logits.shape == kernel.logits.shape and logits.stride() == (logits.shape[1] + 9, 1) + assert torch.equal(logits, _nan_to_minus_inf(kernel.logits.float())) + + +@requires_cuda +@pytest.mark.parametrize("paged", [False, True], ids=["prefill", "decode"]) +def test_scorer_maps_nan_scores_to_minus_inf(glm, paged): + """vLLM's prefill top-k returns out-of-range indices for rows with NaN; it handles +-inf. + + E.g. FP16 accumulation in the scorer overflows to +-inf, and the head sum turns those into NaN. + """ + indexer = glm.convert() + indexer.indexer_scorer_kwargs_quantizer = _kwargs_quantizer(mma_accum="fp16") + call = _scorer_call(glm.fake, paged) + special = torch.tensor([float("nan"), float("inf"), float("-inf")], device="cuda") + col = torch.arange(glm.fake.scores.shape[1], device="cuda") + picked = glm.fake.valid & (col % 7 == 0) + glm.fake.scores = torch.where(picked, special[col % 3], glm.fake.scores) + logits = glm.score(call.run) + + assert glm.fake.calls[-1].kwargs["mma_accum"] == "fp16" + assert not logits.isnan().any() + assert (logits[picked & glm.fake.scores.isnan()] == float("-inf")).all() + assert torch.equal(logits, _nan_to_minus_inf(glm.fake.calls[-1].logits)) # nothing else + + +@requires_cuda +def test_indexer_op_routes_scorer_to_calling_layer(glm, scorer_stand_ins): + first, second = glm.convert(0), glm.convert(1) + op = scorer_stand_ins.op_module.sparse_attn_indexer_kpool + # Wrapped once, under the breakable-cudagraph decorator, whose replays call op.__wrapped__. + assert op.__wrapped__.__wrapped__ is _indexer_op + first.indexer_scorer_kwargs_quantizer = _kwargs_quantizer(mma_accum="fp16") + callers = [] + call = _scorer_call(glm.fake, paged=False) + + def run_scorer(): + callers.append(vllm_layer_scope.layer_owner()) + return call.run() + + for layer, quantized in ((0, True), (1, False)): + op(None, f"layers.{layer}.indexer.k_cache", None, run_scorer) + assert ("mma_accum" in glm.fake.calls[-1].kwargs) == quantized + # The layer name wrapped like vLLM's ``LayerName``. + op(None, SimpleNamespace(value="layers.0.indexer.k_cache"), None, run_scorer) + # The op of a layer that no indexer owns, e.g. a draft model's: the scorer runs unchanged. + op(None, "draft.layers.0.indexer.k_cache", None, run_scorer) + assert "mma_accum" not in glm.fake.calls[-1].kwargs + assert callers == [first, second, first, None] + assert vllm_layer_scope.layer_owner() is None + assert second._enabled_scorer_quantizers() == [] + + +@requires_cuda +@pytest.mark.timeout(600) # the first FP8 fake quantization on the GPU may build its extension +@pytest.mark.parametrize("graph", ["full", "breakable"]) +def test_indexer_op_in_cuda_graph_replays(glm, scorer_stand_ins, written_cache, monkeypatch, graph): + """Replays of captured indexer ops re-quantize the K-cache writes and score with each layer's + scorer quantizers. + + A FULL CUDA graph captures the op whole: the wraps run at capture, and the replays repeat the + GPU work they launched, which must therefore be capturable. vLLM's breakable CUDA graph runs the + op eagerly at capture and at every replay, by calling the function that + ``eager_break_during_capture`` decorated: the scope goes under that decorator, or the replays + would write the K cache outside the op. + """ + op_module = scorer_stand_ins.op_module + if graph == "breakable": + breakable = pytest.importorskip("vllm.compilation.breakable_cudagraph") + monkeypatch.setattr(breakable, "is_breakable_cudagraph_enabled", lambda: True) + op_module.eager_break_during_capture = breakable.eager_break_during_capture + op_module.sparse_attn_indexer_kpool = breakable.eager_break_during_capture(_indexer_op) + source, loc = written_cache[0].cuda(), written_cache[1].cuda() + kpool_ops = _TestGlm5NextIndexer.kpool_ops + + def kpool_compress_and_write_cache(kv_cache, loc): # writes the new K rows of ``source`` + kv_cache.copy_(source) + + monkeypatch.setattr(kpool_ops, "kpool_compress_and_write_cache", kpool_compress_and_write_cache) + indexers = [glm.convert(layer) for layer in (0, 1)] + caches = [torch.zeros_like(source) for _ in indexers] + for indexer, kv_cache in zip(indexers, caches): + indexer.k_cache.kv_cache = kv_cache + quantizer = indexers[0].indexer_k_quantizer = _fp8_quantizer(amax=2.0).cuda() + indexers[0].indexer_scorer_kwargs_quantizer = _kwargs_quantizer(mma_accum="fp16") + call = _scorer_call(glm.fake, paged=False) + + def run(kv_cache): # what vLLM's indexer op does for its layer + kpool_ops.kpool_compress_and_write_cache(kv_cache, loc) + return call.run() + + def forward(): + return [ + op_module.sparse_attn_indexer_kpool( + None, f"layers.{layer}.indexer.k_cache", None, functools.partial(run, kv_cache) + ) + for layer, kv_cache in enumerate(caches) + ] + + forward() # warm up outside the capture + if graph == "full": + cuda_graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(cuda_graph): + logits = forward() + replay = cuda_graph.replay + else: + capture = breakable.BreakableCUDAGraphCapture() + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream), capture: + forward() + torch.cuda.current_stream().wait_stream(stream) + assert capture.num_eager_breaks == len(caches) + replay = capture.replay + + # New K rows and scores in the captured input buffers, as for the next step of the model. + new_source = written_cache[0].clone() + _write_fp8_indexer_cache( + torch.randn(len(loc), HEAD_DIM) * 5, new_source, written_cache[1].tolist() + ) + source.copy_(new_source) + glm.fake.scores.copy_(torch.randn_like(glm.fake.scores) * 3) + num_calls = len(glm.fake.calls) + replay() + torch.cuda.synchronize() + + _assert_requantized(source, caches[0], loc, quantizer) + assert torch.equal(caches[1], source) # a layer without the K quantizer + valid = glm.fake.valid + if graph == "full": # the replay recomputed the captured outputs from the new scores + for layer_logits in logits: + assert torch.equal(layer_logits[valid], glm.fake.scores[valid]) + # Only the scorer quantizers' layer maps the NaN of the unwritten keys to -inf. + assert not logits[0].isnan().any() and logits[1].isnan().any() + else: # the replay ran the scorer again, with each layer's arguments + calls = glm.fake.calls[num_calls:] + assert [call.kwargs.get("mma_accum") for call in calls] == ["fp16", None] + for replayed in calls: + assert torch.equal(replayed.logits[valid], glm.fake.scores[valid]) + + +@requires_cuda +def test_scorer_kwargs_reach_deep_gemm(glm): + indexer = glm.convert() + numerics = {"mma_accum": "fp16", "post_relu": "mxint8"} + indexer.indexer_scorer_kwargs_quantizer = _kwargs_quantizer(**numerics) + + glm.score(_scorer_call(glm.fake, paged=False).run) + assert glm.fake.calls[-1].kwargs == {"clean_logits": False, **numerics} + + # Like vLLM's wrapper: pass the request indices, and a block table with a unit last stride. + block_tables = torch.arange(8, dtype=torch.int32, device="cuda").reshape(4, 2).t() + indices = torch.arange(4, dtype=torch.int32, device="cuda") + glm.score(_scorer_call(glm.fake, True, block_tables=block_tables, indices=indices).run) + kernel = glm.fake.calls[-1] + assert kernel.kwargs == {"clean_logits": False, "indices": indices, **numerics} + assert kernel.block_table.stride(-1) == 1 and torch.equal(kernel.block_table, block_tables) + + +@requires_cuda +def test_scorer_rejects_undeclared_kwargs(glm): + indexer = glm.convert() + indexer.indexer_scorer_kwargs_quantizer = _kwargs_quantizer(post_relu="mxint8") + run = _scorer_call(glm.fake, paged=False).run + del glm.fake.module.MQA_LOGITS_NUMERICS # a DeepGEMM without the post_relu argument + with pytest.raises(RuntimeError, match="does not declare the indexer scorer keyword arguments"): + glm.score(run) + glm.fake.module.MQA_LOGITS_NUMERICS = {**NUMERICS, "post_relu": ("none",)} + with pytest.raises(RuntimeError, match=r"\{'post_relu': 'mxint8'\}"): + glm.score(run) + + +@requires_cuda +def test_scorer_wraps_vllm_deep_gemm(glm, monkeypatch): + """The wrappers bind vLLM's own scorer signatures and call the DeepGEMM it imports.""" + deep_gemm_utils = pytest.importorskip("vllm.utils.deep_gemm") + paged_scorer = getattr(deep_gemm_utils, "fp8_fp4_paged_mqa_logits", None) + if paged_scorer is None or "indices" not in inspect.signature(paged_scorer).parameters: + pytest.skip("needs vLLM's fp8_fp4 scorer wrappers with indices (newer vLLM, e.g. 0.30)") + for name in ("fp8_fp4_mqa_logits", "fp8_fp4_paged_mqa_logits"): + monkeypatch.setattr(deep_gemm_utils, name, getattr(deep_gemm_utils, name)) # restore + monkeypatch.setattr(deep_gemm_utils, "_import_deep_gemm", lambda: glm.fake.module) + monkeypatch.setattr(deep_gemm_utils, "_lazy_init", lambda: None) + monkeypatch.setattr(_TestGlm5NextIndexer, "deep_gemm_utils", deep_gemm_utils) + indexer = glm.convert() + indexer.indexer_scorer_kwargs_quantizer = _kwargs_quantizer(mma_accum="fp16") + glm.fake.utils = deep_gemm_utils # the indexer op calls vLLM's wrappers + + indices = torch.arange(4, dtype=torch.int32, device="cuda") + for paged in (False, True): + glm.score(_scorer_call(glm.fake, paged, indices=indices if paged else None).run) + expected = {"clean_logits": False, "mma_accum": "fp16"} + assert glm.fake.calls[-1].kwargs == {**expected, **({"indices": indices} if paged else {})} + + +@requires_cuda +@pytest.mark.timeout(600) # the JIT compiles the scorer +def test_in_kernel_numerics_in_the_real_scorer(glm, monkeypatch): + """BF16 post-ReLU products in the scorer of the DeepGEMM that vLLM imports, if it has them. + + Runs with a DeepGEMM build with the scorer numerics first on the Python path. + """ + deep_gemm_utils = pytest.importorskip("vllm.utils.deep_gemm") + if not getattr(deep_gemm_utils._import_deep_gemm(), "MQA_LOGITS_NUMERICS", None): + pytest.skip("vLLM's DeepGEMM has no scorer numerics (MQA_LOGITS_NUMERICS)") + if torch.cuda.get_device_capability()[0] != 10: + pytest.skip("the in-kernel scorer numerics need an SM100 GPU") + for name in ("fp8_fp4_mqa_logits", "fp8_fp4_paged_mqa_logits"): + monkeypatch.setattr(deep_gemm_utils, name, getattr(deep_gemm_utils, name)) # restore + monkeypatch.setattr(_TestGlm5NextIndexer, "deep_gemm_utils", deep_gemm_utils) + indexer = glm.convert() + indexer.indexer_scorer_kwargs_quantizer = _kwargs_quantizer(post_relu="bf16") + + # Small integer codes keep the dots exact and up to 1152, which BF16 rounds to 8 significant + # bits; signed powers of two as weights keep the head sum exact as well. + torch.manual_seed(0) + num_q, num_kv = 64, 256 + q = torch.randint(1, 4, (num_q, NUM_HEADS, HEAD_DIM), device="cuda").float() + q[:, ::4] *= -1 # heads with negative dots, which the ReLU zeroes + k = torch.randint(1, 4, (num_kv, HEAD_DIM), device="cuda").float() + sign = torch.randint(0, 2, (num_q, NUM_HEADS), device="cuda") * 2 - 1 + weights = (sign * 2.0 ** torch.randint(-2, 2, (num_q, NUM_HEADS), device="cuda")).float() + ks = torch.zeros(num_q, dtype=torch.int32, device="cuda") + ke = torch.full_like(ks, num_kv) + dots = torch.einsum("mhd,nd->mhn", q.double(), k.double()).relu() + + def scores(layer): + def run_scorer(): # the indexer op calls vLLM's wrapper + kv = (k.to(torch.float8_e4m3fn), torch.ones(num_kv, device="cuda")) + q_fp8 = (q.to(torch.float8_e4m3fn), None) + return deep_gemm_utils.fp8_fp4_mqa_logits(q_fp8, kv, weights, ks, ke, False) + + return glm.score(run_scorer, layer)[:, :num_kv].clone() + + def reference(post_relu): + return torch.einsum("mh,mhn->mn", weights.double(), post_relu(dots)).float() + + stock = scores(layer=1) # a layer without scorer quantizers + assert torch.equal(stock, reference(lambda x: x)) + rounded = scores(layer=0) + assert torch.equal(rounded, reference(lambda x: x.bfloat16().double())) + assert not torch.equal(rounded, stock) + + +@pytest.mark.parametrize("quantizer_name", list(vllm_indexer._SCORER_KWARGS_QUANTIZERS)) +def test_deepseek_v4_indexer_rejects_scorer_quantizers(parallel_state_stub, quantizer_name): + indexer = _TestV4Indexer.convert(_NativeV4Indexer()) + indexer.indexer_k_quantizer.disable() + indexer.indexer_q_quantizer.disable() + setattr(indexer, quantizer_name, _fp8_quantizer(amax=2.0)) + with pytest.raises(NotImplementedError, match=quantizer_name): + indexer(None, None, None, None, None, None) diff --git a/tests/unit/torch/quantization/plugins/test_vllm_layer_scope.py b/tests/unit/torch/quantization/plugins/test_vllm_layer_scope.py new file mode 100644 index 00000000000..3a190053afc --- /dev/null +++ b/tests/unit/torch/quantization/plugins/test_vllm_layer_scope.py @@ -0,0 +1,148 @@ +# 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. + +"""The owner of a vLLM layer, in scope while an op that names the layer runs.""" + +import functools +import gc +import weakref +from types import SimpleNamespace + +import pytest + +from modelopt.torch.quantization.plugins import vllm_layer_scope + + +class _Owner: + pass + + +def _decorate(fn): + """Stand-in for vLLM's ``eager_break_during_capture``.""" + + @functools.wraps(fn) + def decorated(*args, **kwargs): + return fn(*args, **kwargs) + + return decorated + + +@pytest.fixture +def ops(monkeypatch): + """An op that names its layer and the kernel it calls, with an empty registry of owners.""" + monkeypatch.setattr(vllm_layer_scope, "_owners", weakref.WeakValueDictionary()) + module = SimpleNamespace() + + def kernel(x, scale=2): + return x * scale + + def op(x, layer_name, run=None): + return (run or module.kernel)(x) + + module.kernel, module.op = kernel, _decorate(op) + return SimpleNamespace(module=module, kernel=kernel, op=op) + + +def test_callee_runs_for_the_owner_of_the_layer_whose_op_runs(ops): + owner = _Owner() + vllm_layer_scope.register_layer_owner("layer.0", owner) + calls, outside = [], [] + + def call(called_owner, original, bound): + calls.append((called_owner, dict(bound.arguments))) + return original(*bound.args, **bound.kwargs) + 1 + + vllm_layer_scope.wrap_layer_op(ops.module, "op", "layer_name", decorator=_decorate) + vllm_layer_scope.wrap_layer_callee( + ops.module, "kernel", call, applies=lambda o: o.active, outside=lambda: outside.append(1) + ) + owner.active = True + assert ops.module.op(3, "layer.0") == 7 # bound with the default scale + assert calls == [(owner, {"x": 3, "scale": 2})] + owner.active = False + assert ops.module.op(3, "layer.0") == 6 + assert ops.module.op(3, SimpleNamespace(value="layer.1")) == 6 # a layer without an owner + assert ops.module.kernel(3) == 6 # outside an op + assert outside == [1] and len(calls) == 1 + # A later registration of the name replaces the owner. + other = _Owner() + other.active = True + vllm_layer_scope.register_layer_owner("layer.0", other) + assert ops.module.op(3, "layer.0") == 7 and calls[-1][0] is other + assert vllm_layer_scope.layer_owner() is None + + +def test_scope_nests_and_ends_with_the_op(ops): + owner = _Owner() + vllm_layer_scope.register_layer_owner("layer.0", owner) + vllm_layer_scope.wrap_layer_op(ops.module, "op", "layer_name") + seen = [] + + def record(x): + seen.append(vllm_layer_scope.layer_owner()) + + def inner(x): + record(x) + return ops.module.op(x, "layer.1", run=record) + + ops.module.op(1, layer_name="layer.0", run=inner) + assert seen == [owner, None] + + def fail(x): + raise ValueError(x) + + with pytest.raises(ValueError): + ops.module.op(1, "layer.0", run=fail) + assert vllm_layer_scope.layer_owner() is None + + +def test_op_wrap_goes_under_its_decorator_once(ops): + owner = _Owner() + vllm_layer_scope.register_layer_owner("layer.0", owner) + for _ in range(2): + vllm_layer_scope.wrap_layer_op(ops.module, "op", "layer_name", decorator=_decorate) + vllm_layer_scope.wrap_layer_callee( + ops.module, "kernel", lambda *args: None, applies=lambda o: True + ) + assert ops.module.op.__wrapped__.__wrapped__ is ops.op + assert ops.module.kernel.__wrapped__ is ops.kernel + # A breakable-cudagraph replay calls the op's __wrapped__, which must still set the scope. + seen = [] + + def record(x): + seen.append(vllm_layer_scope.layer_owner()) + + ops.module.op.__wrapped__(1, "layer.0", run=record) + assert seen == [owner] + # Without the decorator, the wrap goes around the decorated op. + module = SimpleNamespace(op=_decorate(ops.op), kernel=ops.kernel) + vllm_layer_scope.wrap_layer_op(module, "op", "layer_name") + assert module.op.__wrapped__.__wrapped__ is ops.op + + +def test_owner_is_held_weakly(ops): + owner = _Owner() + vllm_layer_scope.register_layer_owner("layer.0", owner) + vllm_layer_scope.wrap_layer_op(ops.module, "op", "layer_name") + seen = [] + + def run(x): + seen.append(vllm_layer_scope.layer_owner() is not None) + + ops.module.op(1, "layer.0", run=run) + del owner + gc.collect() + ops.module.op(1, "layer.0", run=run) + assert seen == [True, False] diff --git a/tests/unit/torch/quantization/test_quantize_cpu.py b/tests/unit/torch/quantization/test_quantize_cpu.py index c4fdd4fcc27..59cc782c55c 100644 --- a/tests/unit/torch/quantization/test_quantize_cpu.py +++ b/tests/unit/torch/quantization/test_quantize_cpu.py @@ -729,15 +729,22 @@ def test_indexer_k_check_covers_all_ranks(skip_on_windows): spawn_multiprocess_job(2, _test_indexer_check_covers_all_ranks, backend="gloo") -def test_indexer_q_patterns_need_the_query_quantizer(): +@pytest.mark.parametrize( + "quantizer_name", + [ + "indexer_q_quantizer", + "indexer_scorer_kwargs_quantizer", + ], +) +def test_indexer_patterns_need_their_own_quantizer(quantizer_name): """Each indexer quantizer is checked on its own: a converted key quantizer does not count.""" config = copy.deepcopy(INDEXER_K_ONLY_CFG) config["quant_cfg"].append( - {"quantizer_name": "*indexer_q_quantizer", "cfg": {"num_bits": 8, "axis": None}} + {"quantizer_name": f"*{quantizer_name}", "cfg": {"num_bits": 8, "axis": None}} ) mtq.register(original_cls=_ToyIndexer, quantized_cls=_QuantToyIndexer) try: - with pytest.raises(RuntimeError, match=r"enables indexer_q_quantizer, but no"): + with pytest.raises(RuntimeError, match=rf"enables {quantizer_name}, but no"): mtq.quantize(_ToyIndexerModel(), config, lambda m: m(torch.randn(2, 16))) finally: mtq.unregister(_ToyIndexer) From 1f979e8907353c2aad7ce9f7acdc11c69d9ce3ba Mon Sep 17 00:00:00 2001 From: Shiyang Chen Date: Mon, 12 Oct 2026 22:06:35 -0700 Subject: [PATCH 2/6] Replace the layer scope's callee wrapper with a plain function patch wrap_layer_callee took the owner, a filter and a hook for calls outside an op, which tied its API to the indexer's uses. The layer scope now only tells which layer's op runs (current_layer) and the module that owns that layer (layer_owner), and patch_function patches a function once, calling a wrapper with the original. The GLM-5.3-Flash indexer's K-cache writers and DeepGEMM scorers become plain wrappers that check the scope themselves: outside the indexer op, current_layer() is None. Signed-off-by: Shiyang Chen --- .../quantization/plugins/vllm_indexer.py | 129 ++++++++++-------- .../quantization/plugins/vllm_layer_scope.py | 88 ++++++------ .../quantization/test_vllm_dynamic_modules.py | 4 +- .../quantization/test_vllm_indexer_quant.py | 7 +- .../plugins/test_vllm_layer_scope.py | 77 ++++++----- 5 files changed, 165 insertions(+), 140 deletions(-) diff --git a/modelopt/torch/quantization/plugins/vllm_indexer.py b/modelopt/torch/quantization/plugins/vllm_indexer.py index db43111f3c3..0e697048126 100644 --- a/modelopt/torch/quantization/plugins/vllm_indexer.py +++ b/modelopt/torch/quantization/plugins/vllm_indexer.py @@ -50,7 +50,13 @@ register_quant_backend, ) from .vllm import create_parallel_state -from .vllm_layer_scope import register_layer_owner, wrap_layer_callee, wrap_layer_op +from .vllm_layer_scope import ( + current_layer, + layer_owner, + patch_function, + register_layer_owner, + wrap_layer_op, +) __all__ = ["register_indexer_scorer_kwargs_quantizer"] @@ -373,55 +379,75 @@ def _kpool_decode_written(arguments: dict) -> tuple[torch.Tensor, torch.Tensor]: return slots, valid -def _requantize_written_pools(written_slots: Callable) -> Callable: - """A K-cache writer call of an indexer: the kernel, then a re-quantization of what it wrote.""" +# Once per function: the patched kernels run on every step. +_signature = functools.cache(inspect.signature) - def call(indexer, original: Callable, bound: inspect.BoundArguments): - kv_cache = bound.arguments["kv_cache"] - if kv_cache.data_ptr() != indexer.k_cache.kv_cache.data_ptr(): - raise RuntimeError( - f"The indexer op of {indexer.k_cache.prefix} writes the K cache of another layer." - ) - out = original(*bound.args, **bound.kwargs) - slots, valid = written_slots(bound.arguments) - _requantize_fp8_indexer_k_cache(kv_cache, slots, valid, indexer.indexer_k_quantizer) - return out - return call +def _bind(fn: Callable, args: tuple, kwargs: dict) -> inspect.BoundArguments: + """The call ``fn(*args, **kwargs)`` with its arguments by name, defaults included.""" + bound = _signature(fn).bind(*args, **kwargs) + bound.apply_defaults() + return bound -def _reject_k_cache_write_outside_op() -> None: - """Reject a K-cache write outside the indexer op, which alone names the layer of the cache.""" - if any(indexer.indexer_k_quantizer.is_enabled for indexer in _glm5next_indexers): +def _write_k_cache(written_slots: Callable, original: Callable, *args, **kwargs): + """A call of a K-cache writer; ``written_slots`` gives the slots that the call writes. + + In the indexer op of a layer with ``indexer_k_quantizer`` on: the kernel, then a re-quantization + of the pools it wrote. + """ + if current_layer() is None: + # Only the indexer op names the layer of the cache. + if any(indexer.indexer_k_quantizer.is_enabled for indexer in _glm5next_indexers): + raise RuntimeError( + "vLLM wrote the GLM-5.3-Flash indexer K cache outside the indexer op, so " + "indexer_k_quantizer cannot tell whose cache it is." + ) + return original(*args, **kwargs) + indexer = layer_owner() + if indexer is None or not indexer.indexer_k_quantizer.is_enabled: + return original(*args, **kwargs) + bound = _bind(original, args, kwargs) + kv_cache = bound.arguments["kv_cache"] + if kv_cache.data_ptr() != indexer.k_cache.kv_cache.data_ptr(): raise RuntimeError( - "vLLM wrote the GLM-5.3-Flash indexer K cache outside the indexer op, so " - "indexer_k_quantizer cannot tell whose cache it is." + f"The indexer op of {indexer.k_cache.prefix} writes the K cache of another layer." ) + out = original(*args, **kwargs) + slots, valid = written_slots(bound.arguments) + _requantize_fp8_indexer_k_cache(kv_cache, slots, valid, indexer.indexer_k_quantizer) + return out + + +def _run_scorer(paged: bool, original: Callable, *args, **kwargs) -> torch.Tensor: + """A call of vLLM's DeepGEMM scorer: with the scorer quantizers of the layer whose op runs. + + ``paged``: the decode scorer, else the prefill one. + """ + indexer = layer_owner() + if indexer is None or not indexer._enabled_scorer_quantizers(): + return original(*args, **kwargs) + return indexer._score(original, _bind(original, args, kwargs), paged) -# GLM-5.3-Flash's K-cache writers, with the slots that a call writes. +# GLM-5.3-Flash's K-cache writers, patched to re-quantize the slots that a call writes. _KPOOL_CACHE_WRITERS = { - "kpool_compress_and_write_cache": _kpool_prefill_written, - "kpool_decode_update_and_maybe_write_cache_batched": _kpool_decode_written, + "kpool_compress_and_write_cache": functools.partial(_write_k_cache, _kpool_prefill_written), + "kpool_decode_update_and_maybe_write_cache_batched": functools.partial( + _write_k_cache, _kpool_decode_written + ), +} +# vLLM's DeepGEMM scorers, patched to apply the scorer quantizers. The indexer op imports them from +# vLLM's module on every call. Prefill scores keys gathered from the K cache into one buffer, per +# query range; decode scores keys read straight from the paged K cache. +_DEEP_GEMM_SCORERS = { + "fp8_fp4_mqa_logits": functools.partial(_run_scorer, False), + "fp8_fp4_paged_mqa_logits": functools.partial(_run_scorer, True), } # GLM-5.3-Flash's indexer op, which writes the K cache and scores and selects the keys. _GLM5NEXT_INDEXER_OP = "sparse_attn_indexer_kpool" -def _wrap_deep_gemm_scorer(deep_gemm_utils: ModuleType, paged: bool) -> None: - """Wrap vLLM's DeepGEMM scorer in ``deep_gemm_utils`` to apply the scorer quantizers. - - ``paged``: the decode scorer, else the prefill one. The indexer op imports the scorer from the - module on every call. - """ - wrap_layer_callee( - deep_gemm_utils, - "fp8_fp4_paged_mqa_logits" if paged else "fp8_fp4_mqa_logits", - lambda indexer, original, bound: indexer._score(deep_gemm_utils, original, bound, paged), - applies=lambda indexer: bool(indexer._enabled_scorer_quantizers()), - ) - - def _scorer_deep_gemm(deep_gemm_utils: ModuleType, scorer_kwargs: dict) -> ModuleType: """The DeepGEMM module vLLM loaded, checked to accept the scorer keyword arguments here. @@ -498,9 +524,9 @@ class _QuantVLLMGlm5NextIndexer(_QuantVLLMIndexerBase): The query kernel is swapped for a re-quantizing wrapper while ``forward`` runs; CUDA graph capture records the wrapper. The indexer op writes the K cache and scores the keys of the layer that it names, so the op and the kernel entry points it calls are wrapped process-wide: the op - makes its layer's indexer the layer owner (see :mod:`.vllm_layer_scope`), the K-cache writers - re-quantize the pools they wrote, on prefill and on decode pool completion, and vLLM's DeepGEMM - scorer entry points apply the scorer quantizers. + makes its layer the current layer, whose indexer is the layer owner (see + :mod:`.vllm_layer_scope`), the K-cache writers re-quantize the pools they wrote, on prefill and + on decode pool completion, and vLLM's DeepGEMM scorer entry points apply the scorer quantizers. """ kpool_ops: ModuleType | None = _glm5next_kpool_ops @@ -528,18 +554,10 @@ def _setup(self): "k_cache_prefix", decorator=getattr(self.op_module, "eager_break_during_capture", None), ) - for name, written_slots in _KPOOL_CACHE_WRITERS.items(): - wrap_layer_callee( - self.kpool_ops, - name, - _requantize_written_pools(written_slots), - applies=lambda indexer: indexer.indexer_k_quantizer.is_enabled, - outside=_reject_k_cache_write_outside_op, - ) - # Prefill scores keys gathered from the K cache into one buffer, per query range. - _wrap_deep_gemm_scorer(self.deep_gemm_utils, paged=False) - # Decode scores keys read straight from the paged K cache. - _wrap_deep_gemm_scorer(self.deep_gemm_utils, paged=True) + for name, wrapper in _KPOOL_CACHE_WRITERS.items(): + patch_function(self.kpool_ops, name, wrapper) + for name, wrapper in _DEEP_GEMM_SCORERS.items(): + patch_function(self.deep_gemm_utils, name, wrapper) def forward(self, *args, **kwargs): # The scorer only runs for sequences longer than index_topk, which calibration and warmup @@ -601,16 +619,13 @@ def _scorer_kwargs(self, quantizing: bool = True) -> dict: return kwargs def _score( - self, - deep_gemm_utils: ModuleType, - original: Callable, - bound: inspect.BoundArguments, - paged: bool, + self, original: Callable, bound: inspect.BoundArguments, paged: bool ) -> torch.Tensor: """Run the call ``bound`` of vLLM's scorer ``original`` with the scorer quantizers. ``weights`` is the effective W: the query scale and the model's normalization are folded in. """ + assert self.deep_gemm_utils is not None # checked in _setup arguments = bound.arguments scorer_kwargs = self._scorer_kwargs() weights = arguments["weights"] @@ -618,7 +633,7 @@ def _score( weights = weights.to(torch.bfloat16) if scorer_kwargs: logits = _call_deep_gemm_scorer( - deep_gemm_utils, arguments, weights, scorer_kwargs, paged + self.deep_gemm_utils, arguments, weights, scorer_kwargs, paged ) else: logits = original(*bound.args, **bound.kwargs) diff --git a/modelopt/torch/quantization/plugins/vllm_layer_scope.py b/modelopt/torch/quantization/plugins/vllm_layer_scope.py index 171c3c41800..8b7edde6c6d 100644 --- a/modelopt/torch/quantization/plugins/vllm_layer_scope.py +++ b/modelopt/torch/quantization/plugins/vllm_layer_scope.py @@ -17,13 +17,14 @@ vLLM's layer ops get the name of their layer, not the layer: its ``torch.ops.vllm`` custom ops cannot take a module argument, and other ops, such as GLM-5.3-Flash's indexer op, follow the same -convention. The kernels they call get only tensors. :func:`register_layer_owner` maps a layer name -to the converted module that acts for the layer, :func:`wrap_layer_op` makes that module the -:func:`layer_owner` while an op that names the layer runs, and :func:`wrap_layer_callee` hands it -the calls of a function that such an op calls. - -The registry and the wraps are process-wide, for one converted model per process: another model -with the same layer names would get its owners. The wraps are permanent. A breakable CUDA graph +convention. The kernels they call get only tensors. :func:`wrap_layer_op` makes the layer that an op +names the :func:`current_layer` while the op runs, :func:`register_layer_owner` maps a layer name to +the converted module that acts for the layer, and :func:`layer_owner` gives that module for the +current layer. A function that the op calls, patched with :func:`patch_function`, can thus act for +the converted module. + +The registry and the patches are process-wide, for one converted model per process: another model +with the same layer names would get its owners. The patches are permanent. A breakable CUDA graph replays an op that ``eager_break_during_capture`` decorates by calling the undecorated function, so the scope goes under that decorator. An op inside a captured graph has its scope while the graph is captured, so the owner must launch only work that the graph can replay. @@ -37,14 +38,18 @@ from types import ModuleType from typing import Any -__all__ = ["layer_owner", "register_layer_owner", "wrap_layer_callee", "wrap_layer_op"] +__all__ = [ + "current_layer", + "layer_owner", + "patch_function", + "register_layer_owner", + "wrap_layer_op", +] # Layer name -> the converted module that acts for the layer. Weak: it keeps no module alive. _owners: weakref.WeakValueDictionary = weakref.WeakValueDictionary() -# Inside an op of a layer that no converted module owns. -_NO_OWNER = object() -# The owner of the layer whose op runs: None outside a wrapped op. -_layer_owner: contextvars.ContextVar = contextvars.ContextVar("_layer_owner", default=None) +# The name of the layer whose op runs: None outside a wrapped op. +_current_layer: contextvars.ContextVar = contextvars.ContextVar("_current_layer", default=None) def register_layer_owner(name: str, owner: Any) -> None: @@ -52,19 +57,24 @@ def register_layer_owner(name: str, owner: Any) -> None: _owners[name] = owner +def current_layer() -> str | None: + """The name of the layer whose op runs, or None outside a wrapped op.""" + return _current_layer.get() + + def layer_owner() -> Any: - """The owner of the layer whose op runs, or None.""" - owner = _layer_owner.get() - return None if owner is _NO_OWNER else owner + """The owner of the layer whose op runs, or None outside a wrapped op or without an owner.""" + name = _current_layer.get() + return None if name is None else _owners.get(name) def wrap_layer_op( op_module: ModuleType, name: str, layer_arg: str, decorator: Callable | None = None ) -> None: - """Wrap the op ``op_module.`` to put the owner of the layer it names in scope. + """Wrap the op ``op_module.`` to make the layer it names the current layer. ``layer_arg`` is the op's argument with the layer name. If ``decorator`` decorates the op, the - wrap goes under it and ``decorator`` decorates the wrapper again. + wrap goes under it and ``decorator`` decorates the wrapper again. A second wrap does nothing. """ original = getattr(op_module, name) if getattr(original, "_modelopt_layer_op", False): @@ -76,11 +86,11 @@ def wrap_layer_op( def wrapper(*args, **kwargs): layer = kwargs[layer_arg] if layer_arg in kwargs else args[position] # A str or vLLM's LayerName. - token = _layer_owner.set(_owners.get(getattr(layer, "value", layer), _NO_OWNER)) + token = _current_layer.set(getattr(layer, "value", layer)) try: return inner(*args, **kwargs) finally: - _layer_owner.reset(token) + _current_layer.reset(token) if inner is not original: wrapper = decorator(wrapper) # type: ignore[misc] @@ -88,35 +98,21 @@ def wrapper(*args, **kwargs): setattr(op_module, name, wrapper) -def wrap_layer_callee( - module: ModuleType, - name: str, - call: Callable[[Any, Callable, inspect.BoundArguments], Any], - applies: Callable[[Any], bool], - outside: Callable[[], None] | None = None, -) -> None: - """Wrap ``module.``, a function that ops call, to hand its calls to the layer's owner. +def patch_function(module: ModuleType, name: str, wrapper: Callable) -> None: + """Replace ``module.`` so that its calls run ``wrapper(original, *args, **kwargs)``. - Inside an op that :func:`wrap_layer_op` wrapped, a call for whose owner ``applies(owner)`` - holds runs ``call(owner, original, bound)`` with the bound arguments of the call; other calls - run the original function. ``outside()`` runs before a call outside such an op, e.g. to reject - it. + A second patch with the same ``wrapper`` does nothing; another wrapper wraps the patched one. """ original = getattr(module, name) - if getattr(original, "_modelopt_layer_callee", False): - return - signature = inspect.signature(original) + patched_fn = original + while patched_fn is not None: + if getattr(patched_fn, "_modelopt_patched_by", None) is wrapper: + return + patched_fn = getattr(patched_fn, "__wrapped__", None) @functools.wraps(original) - def wrapper(*args, **kwargs): - owner = _layer_owner.get() - if owner is None and outside is not None: - outside() - if owner is None or owner is _NO_OWNER or not applies(owner): - return original(*args, **kwargs) - bound = signature.bind(*args, **kwargs) - bound.apply_defaults() - return call(owner, original, bound) - - wrapper._modelopt_layer_callee = True # type: ignore[attr-defined] - setattr(module, name, wrapper) + def patched(*args, **kwargs): + return wrapper(original, *args, **kwargs) + + patched._modelopt_patched_by = wrapper # type: ignore[attr-defined] + setattr(module, name, patched) diff --git a/tests/gpu_vllm/torch/quantization/test_vllm_dynamic_modules.py b/tests/gpu_vllm/torch/quantization/test_vllm_dynamic_modules.py index 9cc32bdcb27..3f0f01c6798 100644 --- a/tests/gpu_vllm/torch/quantization/test_vllm_dynamic_modules.py +++ b/tests/gpu_vllm/torch/quantization/test_vllm_dynamic_modules.py @@ -1015,7 +1015,7 @@ def _calibrate_and_clip_indexer_k(self, cfg): return amaxes -def _record_scorer_call(score, name, calls, deep_gemm_utils, original, bound, paged): +def _record_scorer_call(score, name, calls, original, bound, paged): """Spy on ``_QuantVLLMGlm5NextIndexer._score``: the layer, the path, what top-k reads and the in-kernel numerics passed to DeepGEMM.""" passed = {} @@ -1027,7 +1027,7 @@ def spy(deep_gemm_utils, arguments, weights, scorer_kwargs, paged): vllm_indexer._call_deep_gemm_scorer = spy try: - logits = score(deep_gemm_utils, original, bound, paged) + logits = score(original, bound, paged) finally: vllm_indexer._call_deep_gemm_scorer = call_scorer numerics = {key: passed[key] for key in ("mma_accum", "post_relu") if key in passed} diff --git a/tests/gpu_vllm/torch/quantization/test_vllm_indexer_quant.py b/tests/gpu_vllm/torch/quantization/test_vllm_indexer_quant.py index b60bddaf26c..916c57aa26e 100644 --- a/tests/gpu_vllm/torch/quantization/test_vllm_indexer_quant.py +++ b/tests/gpu_vllm/torch/quantization/test_vllm_indexer_quant.py @@ -394,13 +394,15 @@ def write(layer, writer, *args, **kwargs): quantizer = indexer.indexer_k_quantizer = _fp8_quantizer(amax=2.0) # A write outside the indexer op names no layer, and one to another layer's cache is misrouted: - # both fail before the kernel runs. A layer without the quantizer re-quantizes nothing. + # both fail before the kernel runs. A layer without the quantizer re-quantizes nothing, and + # neither does the op of a layer that no indexer owns, e.g. a draft model's. with pytest.raises(RuntimeError, match="outside the indexer op"): kpool_ops.kpool_compress_and_write_cache(*prefill) prefill_other = (other_cache, None, None, None, slot_mapping, 4) with pytest.raises(RuntimeError, match="writes the K cache of another layer"): write(0, kpool_ops.kpool_compress_and_write_cache, *prefill_other) write(1, kpool_ops.kpool_compress_and_write_cache, *prefill_other) + write("draft", kpool_ops.kpool_compress_and_write_cache, *prefill_other) assert not requantized # Prefill: pools at ``loc`` masked by ``write_mask``. @@ -418,9 +420,8 @@ def write(layer, writer, *args, **kwargs): *(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"] * 4 + ["decode"] + assert kernel_calls == ["prefill"] * 5 + ["decode"] assert len(requantized) == 2 and all(cache is kv_cache for cache in requantized) - assert getattr(kpool_ops.kpool_compress_and_write_cache, "_modelopt_layer_callee", False) def test_glm5next_indexer_requantizes_query(parallel_state_stub, monkeypatch): diff --git a/tests/unit/torch/quantization/plugins/test_vllm_layer_scope.py b/tests/unit/torch/quantization/plugins/test_vllm_layer_scope.py index 3a190053afc..2678ad9ebba 100644 --- a/tests/unit/torch/quantization/plugins/test_vllm_layer_scope.py +++ b/tests/unit/torch/quantization/plugins/test_vllm_layer_scope.py @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""The owner of a vLLM layer, in scope while an op that names the layer runs.""" +"""The layer whose op runs and its owner, seen from the functions that the op calls.""" import functools import gc @@ -55,33 +55,32 @@ def op(x, layer_name, run=None): return SimpleNamespace(module=module, kernel=kernel, op=op) -def test_callee_runs_for_the_owner_of_the_layer_whose_op_runs(ops): +def _scope(): + return vllm_layer_scope.current_layer(), vllm_layer_scope.layer_owner() + + +def test_patched_callee_sees_the_layer_whose_op_runs_and_its_owner(ops): owner = _Owner() vllm_layer_scope.register_layer_owner("layer.0", owner) - calls, outside = [], [] + calls = [] - def call(called_owner, original, bound): - calls.append((called_owner, dict(bound.arguments))) - return original(*bound.args, **bound.kwargs) + 1 + def scaled_up(original, x, scale=2): + calls.append((*_scope(), x, scale)) + return original(x, scale) + 1 vllm_layer_scope.wrap_layer_op(ops.module, "op", "layer_name", decorator=_decorate) - vllm_layer_scope.wrap_layer_callee( - ops.module, "kernel", call, applies=lambda o: o.active, outside=lambda: outside.append(1) - ) - owner.active = True - assert ops.module.op(3, "layer.0") == 7 # bound with the default scale - assert calls == [(owner, {"x": 3, "scale": 2})] - owner.active = False - assert ops.module.op(3, "layer.0") == 6 - assert ops.module.op(3, SimpleNamespace(value="layer.1")) == 6 # a layer without an owner - assert ops.module.kernel(3) == 6 # outside an op - assert outside == [1] and len(calls) == 1 + vllm_layer_scope.patch_function(ops.module, "kernel", scaled_up) + assert ops.module.op(3, "layer.0") == 7 + # A layer without an owner, named by vLLM's LayerName. + assert ops.module.op(3, SimpleNamespace(value="layer.1")) == 7 + assert ops.module.kernel(3, scale=1) == 4 # outside an op + assert calls == [("layer.0", owner, 3, 2), ("layer.1", None, 3, 2), (None, None, 3, 1)] # A later registration of the name replaces the owner. other = _Owner() - other.active = True vllm_layer_scope.register_layer_owner("layer.0", other) - assert ops.module.op(3, "layer.0") == 7 and calls[-1][0] is other - assert vllm_layer_scope.layer_owner() is None + ops.module.op(3, "layer.0") + assert calls[-1][:2] == ("layer.0", other) + assert _scope() == (None, None) def test_scope_nests_and_ends_with_the_op(ops): @@ -91,21 +90,21 @@ def test_scope_nests_and_ends_with_the_op(ops): seen = [] def record(x): - seen.append(vllm_layer_scope.layer_owner()) + seen.append(_scope()) def inner(x): record(x) return ops.module.op(x, "layer.1", run=record) ops.module.op(1, layer_name="layer.0", run=inner) - assert seen == [owner, None] + assert seen == [("layer.0", owner), ("layer.1", None)] def fail(x): raise ValueError(x) with pytest.raises(ValueError): ops.module.op(1, "layer.0", run=fail) - assert vllm_layer_scope.layer_owner() is None + assert _scope() == (None, None) def test_op_wrap_goes_under_its_decorator_once(ops): @@ -113,25 +112,39 @@ def test_op_wrap_goes_under_its_decorator_once(ops): vllm_layer_scope.register_layer_owner("layer.0", owner) for _ in range(2): vllm_layer_scope.wrap_layer_op(ops.module, "op", "layer_name", decorator=_decorate) - vllm_layer_scope.wrap_layer_callee( - ops.module, "kernel", lambda *args: None, applies=lambda o: True - ) assert ops.module.op.__wrapped__.__wrapped__ is ops.op - assert ops.module.kernel.__wrapped__ is ops.kernel # A breakable-cudagraph replay calls the op's __wrapped__, which must still set the scope. seen = [] def record(x): - seen.append(vllm_layer_scope.layer_owner()) + seen.append(_scope()) ops.module.op.__wrapped__(1, "layer.0", run=record) - assert seen == [owner] + assert seen == [("layer.0", owner)] # Without the decorator, the wrap goes around the decorated op. - module = SimpleNamespace(op=_decorate(ops.op), kernel=ops.kernel) + module = SimpleNamespace(op=_decorate(ops.op)) vllm_layer_scope.wrap_layer_op(module, "op", "layer_name") assert module.op.__wrapped__.__wrapped__ is ops.op +def test_patch_function_applies_each_wrapper_once(ops): + def plus_one(original, *args, **kwargs): + return original(*args, **kwargs) + 1 + + def doubled(original, *args, **kwargs): + return original(*args, **kwargs) * 2 + + for _ in range(2): + vllm_layer_scope.patch_function(ops.module, "kernel", plus_one) + assert ops.module.kernel.__wrapped__ is ops.kernel + assert ops.module.kernel(3) == 7 + # Another wrapper wraps the patched function; patching with either again changes nothing. + vllm_layer_scope.patch_function(ops.module, "kernel", doubled) + vllm_layer_scope.patch_function(ops.module, "kernel", plus_one) + vllm_layer_scope.patch_function(ops.module, "kernel", doubled) + assert ops.module.kernel(3) == 14 + + def test_owner_is_held_weakly(ops): owner = _Owner() vllm_layer_scope.register_layer_owner("layer.0", owner) @@ -139,10 +152,10 @@ def test_owner_is_held_weakly(ops): seen = [] def run(x): - seen.append(vllm_layer_scope.layer_owner() is not None) + seen.append((vllm_layer_scope.current_layer(), vllm_layer_scope.layer_owner() is not None)) ops.module.op(1, "layer.0", run=run) del owner gc.collect() ops.module.op(1, "layer.0", run=run) - assert seen == [True, False] + assert seen == [("layer.0", True), ("layer.0", False)] From da61765a8220f2bf2374a190217277454c36dc63 Mon Sep 17 00:00:00 2001 From: Shiyang Chen Date: Tue, 13 Oct 2026 08:25:59 -0700 Subject: [PATCH 3/6] Point the scorer quantizer's unsupported-model error at its README section The table of plugin-only quantizers named the indexer_*_nvfp4 recipe units for indexer_scorer_kwargs_quantizer, which do not cover it. Its last column now says where the supported models are listed, and for the scorer quantizer that is its section of examples/vllm_serve/README.md. That section's serve command now runs the launcher by its path from the repository root, as the other examples do. Signed-off-by: Shiyang Chen --- examples/vllm_serve/README.md | 2 +- modelopt/torch/quantization/model_quant.py | 18 ++++++++++-------- .../torch/quantization/test_quantize_cpu.py | 10 +++++----- 3 files changed, 16 insertions(+), 14 deletions(-) diff --git a/examples/vllm_serve/README.md b/examples/vllm_serve/README.md index 4913aac4c6b..f6880b37858 100644 --- a/examples/vllm_serve/README.md +++ b/examples/vllm_serve/README.md @@ -351,7 +351,7 @@ such a build into the serving environment and serve: ```bash pip install --no-deps --no-build-isolation RECIPE_PATH=.yaml \ - python vllm_serve_fakequant.py -tp 8 --host 0.0.0.0 --port 8000 + python3 examples/vllm_serve/vllm_serve_fakequant.py -tp 8 --host 0.0.0.0 --port 8000 ``` Notes: diff --git a/modelopt/torch/quantization/model_quant.py b/modelopt/torch/quantization/model_quant.py index 12ca0359d7d..5067138b786 100644 --- a/modelopt/torch/quantization/model_quant.py +++ b/modelopt/torch/quantization/model_quant.py @@ -232,18 +232,20 @@ def _check_weight_quantization_took_effect(model: nn.Module, config: QuantizeCon ) +_INDEXER_UNITS = "the configs/ptq/units/indexer_*_nvfp4 recipe units" +_MOE_COMM_UNITS = "the configs/ptq/units/moe_*_nvfp4 recipe units" # Quantizers that only framework plugins add: (class-name patterns of the modules they belong on, -# what those modules are, recipe units listing the supported models). +# what those modules are, where the supported models are listed). _PLUGIN_QUANTIZERS = { - "indexer_q_quantizer": (("*Indexer",), "sparse-attention indexer", "indexer_*_nvfp4"), - "indexer_k_quantizer": (("*Indexer",), "sparse-attention indexer", "indexer_*_nvfp4"), + "indexer_q_quantizer": (("*Indexer",), "sparse-attention indexer", _INDEXER_UNITS), + "indexer_k_quantizer": (("*Indexer",), "sparse-attention indexer", _INDEXER_UNITS), "indexer_scorer_kwargs_quantizer": ( ("*Indexer",), "sparse-attention indexer", - "indexer_*_nvfp4", + 'the "Indexer scorer kernel arguments" section of examples/vllm_serve/README.md', ), - "dispatch_quantizer": (("*Experts*", "*FusedMoE"), "MoE experts module", "moe_*_nvfp4"), - "combine_quantizer": (("*Experts*", "*FusedMoE"), "MoE experts module", "moe_*_nvfp4"), + "dispatch_quantizer": (("*Experts*", "*FusedMoE"), "MoE experts module", _MOE_COMM_UNITS), + "combine_quantizer": (("*Experts*", "*FusedMoE"), "MoE experts module", _MOE_COMM_UNITS), } @@ -261,7 +263,7 @@ def _check_plugin_quantization_took_effect(model: nn.Module, config: QuantizeCon some waiting in calibration for the others. """ last_entry_per_pattern = {entry.quantizer_name: entry for entry in config.quant_cfg} - for quantizer_name, (owner_patterns, owner_kind, units) in _PLUGIN_QUANTIZERS.items(): + for quantizer_name, (owner_patterns, owner_kind, supported) in _PLUGIN_QUANTIZERS.items(): if not any( entry.enable and quantizer_name in pattern for pattern, entry in last_entry_per_pattern.items() @@ -295,7 +297,7 @@ def _check_plugin_quantization_took_effect(model: nn.Module, config: QuantizeCon f"The quantization config enables {quantizer_name}, but no {owner_kind} of this " f"model ({', '.join(unsupported)}) has one: the ModelOpt plugins do not support " "this model, its runtime backend or the installed vLLM / Megatron-Core version " - f"(supported models: see the configs/ptq/units/{units} recipe units)." + f"(supported models: see {supported})." ) if any(has for _, has, _ in per_rank) and not any(on for _, _, on in per_rank): raise RuntimeError( diff --git a/tests/unit/torch/quantization/test_quantize_cpu.py b/tests/unit/torch/quantization/test_quantize_cpu.py index 59cc782c55c..8feae97973d 100644 --- a/tests/unit/torch/quantization/test_quantize_cpu.py +++ b/tests/unit/torch/quantization/test_quantize_cpu.py @@ -730,13 +730,13 @@ def test_indexer_k_check_covers_all_ranks(skip_on_windows): @pytest.mark.parametrize( - "quantizer_name", + ("quantizer_name", "supported"), [ - "indexer_q_quantizer", - "indexer_scorer_kwargs_quantizer", + ("indexer_q_quantizer", r"configs/ptq/units/indexer_\*_nvfp4"), + ("indexer_scorer_kwargs_quantizer", r"scorer kernel arguments.*examples/vllm_serve"), ], ) -def test_indexer_patterns_need_their_own_quantizer(quantizer_name): +def test_indexer_patterns_need_their_own_quantizer(quantizer_name, supported): """Each indexer quantizer is checked on its own: a converted key quantizer does not count.""" config = copy.deepcopy(INDEXER_K_ONLY_CFG) config["quant_cfg"].append( @@ -744,7 +744,7 @@ def test_indexer_patterns_need_their_own_quantizer(quantizer_name): ) mtq.register(original_cls=_ToyIndexer, quantized_cls=_QuantToyIndexer) try: - with pytest.raises(RuntimeError, match=rf"enables {quantizer_name}, but no"): + with pytest.raises(RuntimeError, match=rf"enables {quantizer_name}, but no.*{supported}"): mtq.quantize(_ToyIndexerModel(), config, lambda m: m(torch.randn(2, 16))) finally: mtq.unregister(_ToyIndexer) From 0b4930cf77e6dd2cad55ec04880efc1f823c2619 Mon Sep 17 00:00:00 2001 From: Shiyang Chen Date: Tue, 13 Oct 2026 10:51:14 -0700 Subject: [PATCH 4/6] Test that a quantizer which only carries a config survives serving and a restore The vLLM indexers' scorer quantizers are enabled TensorQuantizers that never see a tensor, so they have no amax. A CPU test now takes such a quantizer, whose backend fails if it is called, through calibration, the quantization summary and fold_weight, as vLLM fake-quant serving does, and through a modelopt_state restore. It stays enabled and quantizing, gets no amax, and keeps its backend and backend_extra_args. Signed-off-by: Shiyang Chen --- .../torch/quantization/test_quantize_cpu.py | 46 +++++++++++++++++++ 1 file changed, 46 insertions(+) diff --git a/tests/unit/torch/quantization/test_quantize_cpu.py b/tests/unit/torch/quantization/test_quantize_cpu.py index 8feae97973d..dc6c66345b5 100644 --- a/tests/unit/torch/quantization/test_quantize_cpu.py +++ b/tests/unit/torch/quantization/test_quantize_cpu.py @@ -40,6 +40,8 @@ from modelopt.torch.quantization.nn.modules.tensor_quantizer import ( SequentialQuantizer, TensorQuantizer, + register_quant_backend, + unregister_quant_backend, ) from modelopt.torch.utils.distributed import ParallelState @@ -750,6 +752,50 @@ def test_indexer_patterns_need_their_own_quantizer(quantizer_name, supported): mtq.unregister(_ToyIndexer) +class _QuantToyScorerIndexer(_QuantToyIndexer): + def _setup(self): + super()._setup() + # Like the vLLM indexers' scorer quantizers: it carries a config and never sees a tensor. + self.indexer_scorer_kwargs_quantizer = TensorQuantizer() + + +def _never_on_a_tensor(inputs, quantizer): + raise AssertionError("a quantizer that only carries a config ran on a tensor") + + +def test_config_only_quantizer_survives_calibration_fold_summary_and_restore(): + """An enabled quantizer that never sees a tensor gets no amax and keeps its config.""" + config = copy.deepcopy(INDEXER_K_ONLY_CFG) + config["quant_cfg"] += [ + {"quantizer_name": "*weight_quantizer", "cfg": {"num_bits": 8, "axis": 0}}, + { + "quantizer_name": "*indexer_scorer_kwargs_quantizer", + "cfg": {"backend": "config_only", "backend_extra_args": {"mma_accum": "fp16"}}, + }, + ] + register_quant_backend("config_only", _never_on_a_tensor) + mtq.register(original_cls=_ToyIndexer, quantized_cls=_QuantToyScorerIndexer) + try: + model = mtq.quantize(_ToyIndexerModel(), config, lambda m: m(torch.randn(2, 16))) + state, weights = mto.modelopt_state(model), copy.deepcopy(model.state_dict()) + # The steps of vLLM fake-quant serving after quantization. + mtq.print_quant_summary(model) + mtq.fold_weight(model) + restored = mto.restore_from_modelopt_state(_ToyIndexerModel(), state) + restored.load_state_dict(weights) + finally: + mtq.unregister(_ToyIndexer) + unregister_quant_backend("config_only") + for quantizer in ( + model.indexer.indexer_scorer_kwargs_quantizer, + restored.indexer.indexer_scorer_kwargs_quantizer, + ): + assert quantizer.is_enabled and quantizer._if_quant + assert quantizer.backend == "config_only" + assert quantizer.backend_extra_args == {"mma_accum": "fp16"} + assert not hasattr(quantizer, "_amax") + + class _ToyMegaMoEExperts(torch.nn.Module): """Stands in for MoE experts no plugin converts, like vLLM's fused Mega-MoE kernels.""" From f65bdbef8bc465dc94ad5a9d7dba4e2411c04626 Mon Sep 17 00:00:00 2001 From: Shiyang Chen Date: Tue, 13 Oct 2026 13:46:11 -0700 Subject: [PATCH 5/6] Rename layer_owner to current_layer_owner In the layer scope, a layer is its name and its owner is the converted module registered under it. current_layer() gives the name of the layer whose op runs; current_layer_owner() now says by its name that it gives that layer's owner. Signed-off-by: Shiyang Chen --- modelopt/torch/quantization/plugins/vllm_indexer.py | 6 +++--- .../torch/quantization/plugins/vllm_layer_scope.py | 10 +++++----- .../torch/quantization/test_vllm_indexer_quant.py | 4 ++-- .../quantization/plugins/test_vllm_layer_scope.py | 6 ++++-- 4 files changed, 14 insertions(+), 12 deletions(-) diff --git a/modelopt/torch/quantization/plugins/vllm_indexer.py b/modelopt/torch/quantization/plugins/vllm_indexer.py index 0e697048126..07406d668ca 100644 --- a/modelopt/torch/quantization/plugins/vllm_indexer.py +++ b/modelopt/torch/quantization/plugins/vllm_indexer.py @@ -52,7 +52,7 @@ from .vllm import create_parallel_state from .vllm_layer_scope import ( current_layer, - layer_owner, + current_layer_owner, patch_function, register_layer_owner, wrap_layer_op, @@ -404,7 +404,7 @@ def _write_k_cache(written_slots: Callable, original: Callable, *args, **kwargs) "indexer_k_quantizer cannot tell whose cache it is." ) return original(*args, **kwargs) - indexer = layer_owner() + indexer = current_layer_owner() if indexer is None or not indexer.indexer_k_quantizer.is_enabled: return original(*args, **kwargs) bound = _bind(original, args, kwargs) @@ -424,7 +424,7 @@ def _run_scorer(paged: bool, original: Callable, *args, **kwargs) -> torch.Tenso ``paged``: the decode scorer, else the prefill one. """ - indexer = layer_owner() + indexer = current_layer_owner() if indexer is None or not indexer._enabled_scorer_quantizers(): return original(*args, **kwargs) return indexer._score(original, _bind(original, args, kwargs), paged) diff --git a/modelopt/torch/quantization/plugins/vllm_layer_scope.py b/modelopt/torch/quantization/plugins/vllm_layer_scope.py index 8b7edde6c6d..6092561340b 100644 --- a/modelopt/torch/quantization/plugins/vllm_layer_scope.py +++ b/modelopt/torch/quantization/plugins/vllm_layer_scope.py @@ -19,9 +19,9 @@ cannot take a module argument, and other ops, such as GLM-5.3-Flash's indexer op, follow the same convention. The kernels they call get only tensors. :func:`wrap_layer_op` makes the layer that an op names the :func:`current_layer` while the op runs, :func:`register_layer_owner` maps a layer name to -the converted module that acts for the layer, and :func:`layer_owner` gives that module for the -current layer. A function that the op calls, patched with :func:`patch_function`, can thus act for -the converted module. +the converted module that acts for the layer, and :func:`current_layer_owner` gives that module for +the current layer. A function that the op calls, patched with :func:`patch_function`, can thus act +for the converted module. The registry and the patches are process-wide, for one converted model per process: another model with the same layer names would get its owners. The patches are permanent. A breakable CUDA graph @@ -40,7 +40,7 @@ __all__ = [ "current_layer", - "layer_owner", + "current_layer_owner", "patch_function", "register_layer_owner", "wrap_layer_op", @@ -62,7 +62,7 @@ def current_layer() -> str | None: return _current_layer.get() -def layer_owner() -> Any: +def current_layer_owner() -> Any: """The owner of the layer whose op runs, or None outside a wrapped op or without an owner.""" name = _current_layer.get() return None if name is None else _owners.get(name) diff --git a/tests/gpu_vllm/torch/quantization/test_vllm_indexer_quant.py b/tests/gpu_vllm/torch/quantization/test_vllm_indexer_quant.py index 916c57aa26e..b7cfbe833fa 100644 --- a/tests/gpu_vllm/torch/quantization/test_vllm_indexer_quant.py +++ b/tests/gpu_vllm/torch/quantization/test_vllm_indexer_quant.py @@ -806,7 +806,7 @@ def test_indexer_op_routes_scorer_to_calling_layer(glm, scorer_stand_ins): call = _scorer_call(glm.fake, paged=False) def run_scorer(): - callers.append(vllm_layer_scope.layer_owner()) + callers.append(vllm_layer_scope.current_layer_owner()) return call.run() for layer, quantized in ((0, True), (1, False)): @@ -818,7 +818,7 @@ def run_scorer(): op(None, "draft.layers.0.indexer.k_cache", None, run_scorer) assert "mma_accum" not in glm.fake.calls[-1].kwargs assert callers == [first, second, first, None] - assert vllm_layer_scope.layer_owner() is None + assert vllm_layer_scope.current_layer_owner() is None assert second._enabled_scorer_quantizers() == [] diff --git a/tests/unit/torch/quantization/plugins/test_vllm_layer_scope.py b/tests/unit/torch/quantization/plugins/test_vllm_layer_scope.py index 2678ad9ebba..2681397b053 100644 --- a/tests/unit/torch/quantization/plugins/test_vllm_layer_scope.py +++ b/tests/unit/torch/quantization/plugins/test_vllm_layer_scope.py @@ -56,7 +56,7 @@ def op(x, layer_name, run=None): def _scope(): - return vllm_layer_scope.current_layer(), vllm_layer_scope.layer_owner() + return vllm_layer_scope.current_layer(), vllm_layer_scope.current_layer_owner() def test_patched_callee_sees_the_layer_whose_op_runs_and_its_owner(ops): @@ -152,7 +152,9 @@ def test_owner_is_held_weakly(ops): seen = [] def run(x): - seen.append((vllm_layer_scope.current_layer(), vllm_layer_scope.layer_owner() is not None)) + seen.append( + (vllm_layer_scope.current_layer(), vllm_layer_scope.current_layer_owner() is not None) + ) ops.module.op(1, "layer.0", run=run) del owner From d5b423893280508c436078a20ef60ee4a1a15460 Mon Sep 17 00:00:00 2001 From: Shiyang Chen Date: Tue, 13 Oct 2026 17:30:26 -0700 Subject: [PATCH 6/6] Check quantizers registered for the indexer scorer in mtq.quantize mtq.quantize raises when a config enables indexer_scorer_kwargs_quantizer but no indexer has it, so a recipe for an unsupported model or vLLM version does not silently keep the stock scorer. Quantizers that plugins add with register_indexer_scorer_kwargs_quantizer were not checked; registering one now adds it to the same check. Signed-off-by: Shiyang Chen --- .../quantization/plugins/vllm_indexer.py | 6 +++- .../quantization/test_vllm_indexer_quant.py | 31 +++++++++++++++++++ 2 files changed, 36 insertions(+), 1 deletion(-) diff --git a/modelopt/torch/quantization/plugins/vllm_indexer.py b/modelopt/torch/quantization/plugins/vllm_indexer.py index 07406d668ca..53e982435f4 100644 --- a/modelopt/torch/quantization/plugins/vllm_indexer.py +++ b/modelopt/torch/quantization/plugins/vllm_indexer.py @@ -216,9 +216,13 @@ def register_indexer_scorer_kwargs_quantizer( arguments of the GLM-5.3-Flash indexer's DeepGEMM scorer kernel, which every scorer call of that layer then gets while the quantizer quantizes; it raises for a config it cannot express. Other indexers reject the quantizer. The built-in ``indexer_scorer_kwargs_quantizer`` passes its - ``backend_extra_args`` (``backend: scorer_kwargs``) on as they are. + ``backend_extra_args`` (``backend: scorer_kwargs``) on as they are. Like for the built-in one, + ``mtq.quantize`` raises when a config enables the quantizer but no indexer has it. """ + from ..model_quant import _PLUGIN_QUANTIZERS + _SCORER_KWARGS_QUANTIZERS[name] = scorer_kwargs + _PLUGIN_QUANTIZERS[name] = _PLUGIN_QUANTIZERS[_SCORER_KWARGS_QUANTIZER] def _scorer_kwargs_backend(inputs: torch.Tensor, quantizer: TensorQuantizer) -> torch.Tensor: diff --git a/tests/gpu_vllm/torch/quantization/test_vllm_indexer_quant.py b/tests/gpu_vllm/torch/quantization/test_vllm_indexer_quant.py index b7cfbe833fa..f323ef8a509 100644 --- a/tests/gpu_vllm/torch/quantization/test_vllm_indexer_quant.py +++ b/tests/gpu_vllm/torch/quantization/test_vllm_indexer_quant.py @@ -31,6 +31,7 @@ import torch import modelopt.torch.quantization as mtq +from modelopt.torch.quantization import model_quant from modelopt.torch.quantization.config import QuantizerAttributeConfig from modelopt.torch.quantization.conversion import set_quantizer_by_cfg from modelopt.torch.quantization.nn import QuantModuleRegistry, SequentialQuantizer, TensorQuantizer @@ -671,6 +672,7 @@ def post_relu(quantizer): registered = dict(vllm_indexer._SCORER_KWARGS_QUANTIZERS) monkeypatch.setattr(vllm_indexer, "_SCORER_KWARGS_QUANTIZERS", registered) + monkeypatch.setattr(model_quant, "_PLUGIN_QUANTIZERS", dict(model_quant._PLUGIN_QUANTIZERS)) vllm_indexer.register_indexer_scorer_kwargs_quantizer("indexer_test_quantizer", post_relu) indexer = glm.convert() assert not indexer.indexer_test_quantizer.is_enabled @@ -691,6 +693,35 @@ def post_relu(quantizer): indexer.modelopt_post_restore() +class _UnconvertedIndexer(torch.nn.Module): + """An indexer that no plugin converts.""" + + def forward(self, x): + return x + + +def test_quantize_rejects_registered_quantizers_that_quantize_nothing(monkeypatch): + """Like the built-in scorer quantizer, a registered one must not silently quantize nothing.""" + monkeypatch.setattr( + vllm_indexer, "_SCORER_KWARGS_QUANTIZERS", dict(vllm_indexer._SCORER_KWARGS_QUANTIZERS) + ) + monkeypatch.setattr(model_quant, "_PLUGIN_QUANTIZERS", dict(model_quant._PLUGIN_QUANTIZERS)) + vllm_indexer.register_indexer_scorer_kwargs_quantizer("indexer_test_quantizer", lambda q: {}) + config = { + "quant_cfg": [ + {"quantizer_name": "*", "enable": False}, + {"quantizer_name": "*indexer_test_quantizer", "cfg": {"num_bits": 8, "axis": None}}, + ], + "algorithm": "max", + } + model = torch.nn.Sequential(torch.nn.Linear(4, 4), _UnconvertedIndexer()) + with pytest.raises( + RuntimeError, + match=r"indexer_test_quantizer.*\(_UnconvertedIndexer\).*scorer kernel arguments", + ): + mtq.quantize(model, config, lambda m: m(torch.randn(2, 4))) + + @pytest.mark.parametrize( ("quantizer", "match"), [