diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 77a8fffe1b8..01ebb68405c 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -13,6 +13,7 @@ Changelog - Add a calibration-free streaming Kimi-K3 converter and checkpoint-mirror recipe for NVFP4 routed experts with ``input_scale=1.0`` and 128x128 block-FP8 KDA/MLA attention weights. The converter operates shard-by-shard on the source checkpoint's packed MXFP4 experts instead of loading the 2.8T model through the in-memory ``hf_ptq.py`` path. - Add ``mtq.temporarily_fold_weights`` for repeated frozen-weight inference and ``mtq.preserve_quantizer_attributes_context`` for restoring temporary quantizer property and type changes. Temporary folding snapshots affected fake-quant weights on a configurable device and restores them with their quantizer state; retained pre-quant scales are inactive, while shared weights, shared quantizers, and ``SequentialQuantizer`` weights are unsupported. - Add the ``nvfp4_act_headroom`` calibration algorithm for NVFP4 **activation** global scales. Instead of setting the global scale from the largest per-block amax seen during calibration (plain ``max``, which leaves no room above it so any larger activation saturates), it anchors the scale to a low percentile of the per-block amax distribution, leaving the rest of the FP8 block-scale range as headroom: ``amax = max(rho * anchor, upper)``, where ``anchor`` and ``upper`` are the per-block amaxes at ``anchor_percentile`` (default 1) and ``upper_percentile`` (default 99.99; set to 100 to never clip calibration data), and ``rho`` (default 16384) is the headroom factor. Applies only to NVFP4 dynamic-block input quantizers; ``SequentialQuantizer`` activation quantizers raise. Weight scales are an orthogonal axis selected by a nested ``weight_scale_algorithm`` (``max`` by default, or ``mse`` / ``local_hessian``), so one recipe can combine a weight calibration with this activation policy in a single pass. Ships ``modelopt_recipes/general/ptq/nvfp4_act_headroom-kv_fp8_cast.yaml``, which mirrors ``nvfp4_default-kv_fp8_cast`` with only the calibration algorithm swapped and exports a standard NVFP4 checkpoint. +- Add layer-wise KV-cache AutoQuantize through ``mtq.auto_quantize`` with ``cost_model="kv_cache"``. It measures isolated full-vocabulary forward KL for caller-supplied K/V formats, solves a width-weighted additive storage-constrained recipe across eligible layers, preserves search-disabled layers in their existing format, exports the selected per-attention mapping in unified HF checkpoints, and writes a JSON sensitivity report alongside the checkpoint. A cast-mode FP8/NVFP4 recipe at 5.4 bits/scalar is included. - Add ``layerwise.export_dir``: layerwise calibration writes each decoder layer to its own quantized checkpoint shard as it finishes, so no separate ``export_hf_checkpoint()`` pass is needed and, with ``layerwise.checkpoint_dir``, an interrupted run resumes without redoing finished layers. Supports FP8 and NVFP4 on single-process models, resident or offloaded; other formats and placements raise ``NotImplementedError`` before calibration starts. *Megatron Framework (M-LM / M-Bridge)* diff --git a/examples/hf_ptq/README.md b/examples/hf_ptq/README.md index aaac7a5f04d..230ee474164 100755 --- a/examples/hf_ptq/README.md +++ b/examples/hf_ptq/README.md @@ -430,8 +430,39 @@ leaving the original recipe unchanged. For models without backprop support (e.g. Llama-4), use the `kl_div` scoring method — see the shipped `general/auto_quantize/nvfp4_fp8_kl_div_at_5p4bits` recipe. -KV cache is applied as a uniform post-step, not part of the per-layer search. An AutoQuantize recipe -falls back to `--kv_cache_qformat` (default `fp8_cast`) unless it sets an explicit `kv_cache` field. +Weight AutoQuantize recipes still apply KV cache as a uniform post-step and fall back to +`--kv_cache_qformat` (default `fp8_cast`) unless they set an explicit `kv_cache` field. + +KV-cache AutoQuantize recipes use the same `mtq.auto_quantize` API and set +`constraints.cost_model: kv_cache` with an `effective_bits` target. Their +`candidate_formats` are complete K/V cache configs whose config-level `effective_bits` includes +packed scale overhead. The width-weighted budget covers eligible layers; `disabled_layers` are +preserved and excluded. BF16 is used only as the isolated-KL reference, not as a solver choice. +The shipped recipe searches FP8-cast K/V (8.0 bits/scalar) and NVFP4-cast K/V +(4.5 bits/scalar) at 5.4 bits/scalar. It intentionally excludes FP8-K/NVFP4-V because the +companion vLLM implementation does not support that asymmetric per-layer format: + +```bash +python hf_ptq.py \ + --pyt_ckpt_path Qwen/Qwen3-1.7B \ + --recipe general/auto_quantize/kv_fp8_nvfp4_cast_kl_div_at_5p4bits \ + --auto_quantize_checkpoint /path/to/kv_autoquant.pth \ + --export_path /path/to/qwen3-1.7b-mixed-kv +``` + +Each candidate uses an explicit constant scale, avoiding an additional calibration pass while +keeping persistent K/V scales in the unified HF checkpoint. Unified export records the selected +formats in `kv_cache_quantized_layers` and writes +the JSON-safe sensitivity report to `kv_cache_auto_quantize_report.json`; +`--auto_quantize_checkpoint` stores the resumable raw search state. + +> [!NOTE] +> Layer-wise KV checkpoints require the companion +> [vLLM mixed-KV metadata consumer](https://github.com/vllm-project/vllm/pull/52813) or a later +> vLLM release containing it. The repository's currently pinned vLLM 0.26.0 does not consume +> `kv_cache_quantized_layers`, so these checkpoints are export-only in that stock environment. +> Do not deploy them with the pinned runtime. Full FP8 K/V and full NVFP4 K/V use existing vLLM +> kernels once the layer-wise metadata consumer is available. The one runtime flag is `--auto_quantize_checkpoint` — save/restore the search state to resume an interrupted search (skips re-scoring): @@ -468,6 +499,8 @@ mtq.calibrate(model, algorithm="max", forward_loop=calibrate_loop) ModelOpt enables quantization of LLMs across multiple GPU nodes using FSDP2 for distributed model sharding and calibration, exposed via the `--use_fsdp2` flag on the standard `hf_ptq.py` entry point. +> *KV-cache AutoQuantize recipes are not supported with `--use_fsdp2` and are rejected before model loading. Distributed KV sensitivity scoring, selection, and checkpoint writes must be synchronized before this combination can be enabled safely. Existing weight AutoQuantize recipes retain their previous experimental warning with FSDP2.* + ### Usage #### Slurm (recommended) diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py index 0e3c8035bf8..4e6e8cc8534 100755 --- a/examples/hf_ptq/hf_ptq.py +++ b/examples/hf_ptq/hf_ptq.py @@ -99,6 +99,28 @@ from modelopt.torch.utils.vlm_dataset_utils import get_vlm_dataset_dataloader RAND_SEED = 1234 +_FSDP2_KV_AUTOQUANT_ERROR = ( + "KV-cache AutoQuantize does not support --use_fsdp2 until distributed sensitivity scoring, " + "selection, and checkpoint writes are synchronized across ranks." +) +_FSDP2_AUTOQUANT_WARNING = ( + "AutoQuantize with --use_fsdp2 has not been validated end-to-end yet " + "(distributed calibration, sensitivity scoring, and recipe/checkpoint " + "synchronization across ranks); use at your own risk." +) + + +def _select_unpadded_logits(logits: torch.Tensor, batch: dict[str, Any]) -> torch.Tensor: + """Return logits only for token positions selected by ``attention_mask``.""" + attention_mask = batch.get("attention_mask") + if attention_mask is None: + return logits + if logits.shape[:-1] != attention_mask.shape: + raise ValueError( + "AutoQuant KL logits and attention_mask must have matching token dimensions; " + f"got {tuple(logits.shape[:-1])} and {tuple(attention_mask.shape)}." + ) + return logits[attention_mask.bool()] def _kv_cfg_uses_constant_amax(kv_quant_cfg: list[dict[str, Any]]) -> bool: @@ -251,7 +273,7 @@ def make_calib_dataloader( ) else: assert tokenizer is not None and isinstance( - tokenizer, (PreTrainedTokenizer, PreTrainedTokenizerFast) + tokenizer, PreTrainedTokenizer | PreTrainedTokenizerFast ), "The PreTrainedTokenizer must be set" # Labels are only needed for gradient-based auto_quantize include_labels = autoquant_gradient_recipe @@ -341,6 +363,36 @@ def _mtq_candidate_formats(formats) -> list[dict]: return quantization_formats +def _mtq_kv_candidate_formats(formats) -> list[tuple[dict, str]]: + """Translate format-agnostic KV candidates while preserving useful preset names.""" + candidates = [] + for idx, fmt in enumerate(formats): + quant_cfg = fmt.model_dump(exclude_none=True) + candidate_quantizers = quant_cfg.get("quant_cfg", []) + name = None + nvfp4_quantizers = type(fmt)(**KV_QUANT_CFG_CHOICES["nvfp4"]).model_dump(exclude_none=True)[ + "quant_cfg" + ] + fp8_quantizers = type(fmt)(**KV_QUANT_CFG_CHOICES["fp8"]).model_dump(exclude_none=True)[ + "quant_cfg" + ] + if len(fp8_quantizers) != 1: + raise RuntimeError("The FP8 KV preset must contain exactly one quantizer entry.") + fp8_k_quantizer = copy.deepcopy(fp8_quantizers[0]) + fp8_k_quantizer["quantizer_name"] = "*.k_bmm_quantizer" + if candidate_quantizers == [*nvfp4_quantizers, fp8_k_quantizer]: + name = "fp8_k_nvfp4_v" + for preset_name, preset in KV_QUANT_CFG_CHOICES.items(): + normalized_preset_quantizers = ( + type(fmt)(**preset).model_dump(exclude_none=True).get("quant_cfg", []) + ) + if normalized_preset_quantizers == candidate_quantizers: + name = preset_name + break + candidates.append((quant_cfg, name or f"KV_CACHE_FORMAT_{idx}")) + return candidates + + def _mtq_inputs_from_auto_quantize_config( aq_config, args: argparse.Namespace, fixed_quantize_config=None ) -> dict: @@ -352,6 +404,16 @@ def _mtq_inputs_from_auto_quantize_config( to ``--kv_cache_qformat`` when the recipe omits it. """ constraints = aq_config.constraints.model_dump(exclude_none=True) + is_kv_search = aq_config.constraints.cost_model == "kv_cache" + if is_kv_search: + return { + "search_domain": "kv_cache", + "constraints": constraints, + "quantization_formats": _mtq_kv_candidate_formats(aq_config.candidate_formats), + "disabled_layers": aq_config.disabled_layers, + "method": aq_config.auto_quantize_method, + "score_size": aq_config.score_size, + } # cost_excluded_layers (sibling of disabled_layers) maps to the mtq cost key: these layers are # kept out of the bit-budget denominator (cost_weight 0) — e.g. VL vision towers — distinct from # disabled_layers, which removes them from the search. @@ -383,6 +445,7 @@ def _mtq_inputs_from_auto_quantize_config( for search_space in aq_config.module_search_spaces ] return { + "search_domain": "weight", "constraints": constraints, "quantization_formats": quantization_formats, "fixed_quantization_config": fixed_quantization_config, @@ -416,16 +479,13 @@ def auto_quantize( "Auto Quantization is not supported for pipeline parallel size > 1" ) - if args.use_fsdp2: - warnings.warn( - "AutoQuantize with --use_fsdp2 has not been validated end-to-end yet " - "(distributed calibration, sensitivity scoring, and recipe/checkpoint " - "synchronization across ranks); use at your own risk." - ) - inputs = _mtq_inputs_from_auto_quantize_config( aq_config, args, fixed_quantize_config=fixed_quantize_config ) + if args.use_fsdp2: + if inputs["search_domain"] == "kv_cache": + raise NotImplementedError(_FSDP2_KV_AUTOQUANT_ERROR) + warnings.warn(_FSDP2_AUTOQUANT_WARNING) # base-model lm_head handling (mirrors the CLI helper) is_base_model = ( @@ -464,30 +524,45 @@ def forward_step(model, batch): output = model(**inputs_) if is_base_model: assert full_model is not None - return full_model.lm_head(output.last_hidden_state) - return output.logits + logits = full_model.lm_head(output.last_hidden_state) + else: + logits = output.logits + return _select_unpadded_logits(logits, batch) else: raise ValueError( f"Invalid auto_quantize method: {inputs['method']}. Must be 'gradient' or 'kl_div'" ) + auto_quantize_kwargs: dict[str, Any] = { + "constraints": inputs["constraints"], + "data_loader": calib_dataloader, + "forward_step": forward_step, + "quantization_formats": inputs["quantization_formats"], + "num_calib_steps": len(calib_dataloader), + "num_score_steps": min( + len(calib_dataloader), max(inputs["score_size"] // args.batch_size, 1) + ), + "verbose": True, + "disabled_layers": inputs["disabled_layers"], + "method": inputs["method"], + "checkpoint": args.auto_quantize_checkpoint, + } + if inputs["search_domain"] == "weight": + auto_quantize_kwargs.update( + { + "loss_func": loss_func, + "fixed_quantization_config": inputs["fixed_quantization_config"], + "module_search_spaces": inputs["module_search_spaces"], + } + ) + language_model, _ = mtq.auto_quantize( language_model, - constraints=inputs["constraints"], - data_loader=calib_dataloader, - forward_step=forward_step, - loss_func=loss_func, - quantization_formats=inputs["quantization_formats"], - fixed_quantization_config=inputs["fixed_quantization_config"], - module_search_spaces=inputs["module_search_spaces"], - num_calib_steps=len(calib_dataloader), - num_score_steps=min(len(calib_dataloader), max(inputs["score_size"] // args.batch_size, 1)), - verbose=True, - disabled_layers=inputs["disabled_layers"], - method=inputs["method"], - checkpoint=args.auto_quantize_checkpoint, + **auto_quantize_kwargs, ) + if inputs["search_domain"] == "kv_cache": + return language_model # KV cache quantization is uniform; applied after the LP search. kv_cache_quant_cfg = inputs["kv_cache_quant_cfg"] @@ -512,9 +587,22 @@ def _recipe_is_auto_quantize(recipe: str | None) -> bool: return recipe is not None and isinstance(load_recipe(recipe), ModelOptAutoQuantizeRecipe) +def _recipe_is_kv_auto_quantize(recipe: str | None) -> bool: + """True if ``recipe`` resolves to a KV AutoQuantize recipe (peeked before model load).""" + if recipe is None: + return False + loaded_recipe = load_recipe(recipe) + return ( + isinstance(loaded_recipe, ModelOptAutoQuantizeRecipe) + and loaded_recipe.auto_quantize.constraints.cost_model == "kv_cache" + ) + + def load_model(args: argparse.Namespace): # If low memory mode is enabled, we compress the model while loading the HF checkpoint. calibration_only = False + if args.use_fsdp2 and _recipe_is_kv_auto_quantize(args.recipe): + raise NotImplementedError(_FSDP2_KV_AUTOQUANT_ERROR) if args.use_fsdp2: hf_config = AutoConfig.from_pretrained( args.pyt_ckpt_path, trust_remote_code=args.trust_remote_code @@ -694,7 +782,7 @@ def sparsity_main( # Different calibration datasets are also available, e.g., "pile" and "wikipedia" # Please also check the docstring for the datasets available assert tokenizer is not None and isinstance( - tokenizer, (PreTrainedTokenizer, PreTrainedTokenizerFast) + tokenizer, PreTrainedTokenizer | PreTrainedTokenizerFast ), "The PreTrainedTokenizer must be set" calib_dataloader = get_dataset_dataloader( dataset_name=args.dataset, @@ -1187,7 +1275,7 @@ def quantize_main( if args.recipe is not None: print(f"Use recipe {args.recipe} for quantization") recipe = load_recipe(args.recipe) - if not isinstance(recipe, (ModelOptPTQRecipe, ModelOptAutoQuantizeRecipe)): + if not isinstance(recipe, ModelOptPTQRecipe | ModelOptAutoQuantizeRecipe): raise TypeError( f"Expected PTQ or AutoQuantize recipe, but got {type(recipe).__name__} " f"from {args.recipe}" @@ -1438,9 +1526,10 @@ def parse_args() -> argparse.Namespace: help=( "PTQ or AutoQuantize recipe YAML file or name without suffix (e.g. " "general/ptq/nvfp4_default-kv_fp8_cast, general/auto_quantize/nvfp4_fp8_at_4p8bits). " - "KV cache source depends on the recipe type: PTQ recipes bake KV cache into quant_cfg " - "and --kv_cache_qformat is ignored; AutoQuantize recipes fall back to --kv_cache_qformat " - "unless the recipe sets an explicit kv_cache field." + "KV cache behavior depends on the recipe type: PTQ recipes configure it in quant_cfg " + "and ignore --kv_cache_qformat; weight AutoQuantize recipes use their kv_cache setting " + "or fall back to --kv_cache_qformat; KV-cache AutoQuantize recipes select per-layer K/V " + "formats from candidate_formats and ignore --kv_cache_qformat." ), default=None, ) diff --git a/modelopt/recipe/config.py b/modelopt/recipe/config.py index 2ca8f4f462b..28de372adad 100644 --- a/modelopt/recipe/config.py +++ b/modelopt/recipe/config.py @@ -151,13 +151,18 @@ class AutoQuantizeConstraints(ModeloptBaseConfig): effective_bits: float = ModeloptField( default=4.8, - title="Effective bits per weight", - description="Average weight-storage bits target for the LP, in (0, 16].", + title="Effective bits", + description=( + "Average storage-bits target for the selected cost model, in (0, 16]. Defaults to 4.8." + ), ) - cost_model: Literal["weight", "active_moe"] = ModeloptField( + cost_model: Literal["weight", "active_moe", "kv_cache"] = ModeloptField( default="weight", title="Cost model", - description="'weight' counts all weights equally; 'active_moe' scales routed-expert weights.", + description=( + "'weight' counts all weights equally; 'active_moe' scales routed-expert weights; " + "'kv_cache' accounts for paired K/V-cache storage." + ), ) cost: AutoQuantizeCost | None = ModeloptField( default=None, @@ -172,6 +177,12 @@ def _validate_effective_bits(cls, v: float) -> float: raise ValueError(f"effective_bits must be in (0, 16], got {v}") return v + @model_validator(mode="after") + def _validate_cost_settings(self): + if self.cost_model == "kv_cache" and self.cost is not None: + raise ValueError("KV-cache AutoQuant does not accept weight cost settings.") + return self + class AutoQuantizeModuleSearchSpace(ModeloptBaseConfig): """Candidate formats selectable for modules matching one or more name patterns.""" @@ -272,6 +283,25 @@ def _has_search_space(self): "auto_quantize requires candidate_formats or at least one module_search_spaces " "entry. For uniform quantization, use a PTQ recipe instead." ) + if self.constraints.cost_model == "kv_cache": + if self.auto_quantize_method != "kl_div": + raise ValueError( + "KV-cache AutoQuant currently requires auto_quantize_method=kl_div." + ) + if self.module_search_spaces: + raise ValueError( + "KV-cache AutoQuant uses one candidate space for all eligible attention " + "layers; module_search_spaces is not supported." + ) + if self.kv_cache is not None: + raise ValueError( + "KV-cache AutoQuant candidate_formats replace the uniform kv_cache post-step." + ) + if self.cost_excluded_layers: + raise ValueError( + "KV-cache AutoQuant does not support cost_excluded_layers; use " + "disabled_layers to exclude non-KV-cache modules from the search." + ) return self diff --git a/modelopt/torch/export/convert_hf_config.py b/modelopt/torch/export/convert_hf_config.py index 45fa0c30f3b..534aefe3ba8 100644 --- a/modelopt/torch/export/convert_hf_config.py +++ b/modelopt/torch/export/convert_hf_config.py @@ -269,6 +269,14 @@ def convert_hf_quant_config_format(input_config: dict[str, Any]) -> dict[str, An if kv_cache_quant_algo: if kv_cache_quant_algo == "FP8": new_config["kv_cache_scheme"] = {"dynamic": False, "num_bits": 8, "type": "float"} + elif kv_cache_quant_algo == "MIXED_PRECISION": + new_config["kv_cache_quant_algo"] = kv_cache_quant_algo + new_config["kv_cache_quantized_layers"] = original_quantization_details.get( + "kv_cache_quantized_layers", {} + ) + new_config["kv_cache_schema_version"] = original_quantization_details.get( + "kv_cache_schema_version", 1 + ) else: # TODO: Handle other kv cache quantization algorithms new_config["kv_cache_scheme"] = kv_cache_quant_algo diff --git a/modelopt/torch/export/model_config.py b/modelopt/torch/export/model_config.py index ebafc1d9d7a..cc84275be41 100755 --- a/modelopt/torch/export/model_config.py +++ b/modelopt/torch/export/model_config.py @@ -50,6 +50,7 @@ FUSION_FREE_FORMATS = frozenset({QUANTIZATION_FP8, QUANTIZATION_NONE, QUANTIZATION_FP8_PB_REAL}) KV_CACHE_FP8 = "FP8" +KV_CACHE_FP8_K_NVFP4_V = "FP8_K_NVFP4_V" KV_CACHE_INT8 = "INT8" KV_CACHE_NVFP4 = "NVFP4" KV_CACHE_NVFP4_AFFINE = "NVFP4_AFFINE" diff --git a/modelopt/torch/export/model_utils.py b/modelopt/torch/export/model_utils.py index f27405ec83f..cf89ddc2e5b 100755 --- a/modelopt/torch/export/model_utils.py +++ b/modelopt/torch/export/model_utils.py @@ -108,7 +108,7 @@ def is_multimodal_model(model): config = model.config # Check for Nemotron-Parse encoder-decoder architecture - architectures = getattr(config, "architectures", []) + architectures = getattr(config, "architectures", None) or [] is_nemotron_parse = any("nemotronparse" in arch.lower() for arch in architectures) return ( @@ -137,12 +137,17 @@ def get_language_model_from_vl(model) -> list[nn.Module] | None: >>> # lineage[0] is vlm_model >>> # lineage[1] is vlm_model.language_model """ - # always prioritize model.model.langauge_model + candidates = [] if hasattr(model, "model") and hasattr(model.model, "language_model"): - return [model, model.model, model.model.language_model] - + candidates.append([model, model.model, model.model.language_model]) if hasattr(model, "language_model"): - return [model, model.language_model] + candidates.append([model, model.language_model]) + if len(candidates) > 1: + raise ValueError( + "Found multiple language-model roots; refusing to select one by traversal order." + ) + if candidates: + return candidates[0] # Pattern 3: For encoder-decoder VL models (e.g., Nemotron-Parse), the decoder is the language model. # Only match if the model is detected as multimodal to avoid matching non-VLM encoder-decoder diff --git a/modelopt/torch/export/quant_aware_conversion.py b/modelopt/torch/export/quant_aware_conversion.py index 2e32f123869..fa594f104c4 100644 --- a/modelopt/torch/export/quant_aware_conversion.py +++ b/modelopt/torch/export/quant_aware_conversion.py @@ -278,7 +278,7 @@ def _map(name: str) -> str: def revert_quant_config_names(quantization: dict, mapper) -> None: - """Revert ``exclude_modules`` / ``quantized_layers`` keys to hub names, in place. + """Revert layer-reference keys to hub names, in place. ``mapper`` is the callable from :func:`build_reverse_name_mapper` (a no-op when ``None``). Applies to the ModelOpt ``{"quantization": {...}}`` sub-dict before it is @@ -293,6 +293,11 @@ def revert_quant_config_names(quantization: dict, mapper) -> None: quantized_layers = quantization.get("quantized_layers") if isinstance(quantized_layers, dict) and quantized_layers: quantization["quantized_layers"] = {mapper(k): v for k, v in quantized_layers.items()} + kv_cache_quantized_layers = quantization.get("kv_cache_quantized_layers") + if isinstance(kv_cache_quantized_layers, dict) and kv_cache_quantized_layers: + quantization["kv_cache_quantized_layers"] = { + mapper(k): v for k, v in kv_cache_quantized_layers.items() + } def _assert_experts_pre_expanded( diff --git a/modelopt/torch/export/quant_utils.py b/modelopt/torch/export/quant_utils.py index c86af3aa9f5..2d8f94366a6 100755 --- a/modelopt/torch/export/quant_utils.py +++ b/modelopt/torch/export/quant_utils.py @@ -18,7 +18,6 @@ import logging from collections import defaultdict from collections.abc import Generator -from types import SimpleNamespace from typing import Any from warnings import warn @@ -50,6 +49,7 @@ from ..quantization.nn import NVFP4StaticQuantizer, SequentialQuantizer, TensorQuantizer from .model_config import ( KV_CACHE_FP8, + KV_CACHE_FP8_K_NVFP4_V, KV_CACHE_INT8, KV_CACHE_NVFP4, KV_CACHE_NVFP4_AFFINE, @@ -389,27 +389,36 @@ def get_kv_cache_scaling_factor(self_attention_module: nn.Module) -> list[torch. for quantizer in ("k_bmm_quantizer", "v_bmm_quantizer") ] - # For FP8, we recommend default kv cache scaling factor to be 1. - if get_kv_cache_dtype(self_attention_module) == KV_CACHE_FP8: - for i, factor in enumerate(scaling_factors): - if factor is None: - continue - if factor.item() > 0.5: - warn( - f"Warning: Large KV activation detected: {factor.item()}, " - "Quantized KV cache may lead to higher accuracy drop." - ) - scaling_factors[i] = torch.max( - factor, torch.tensor([1.0], dtype=torch.float, device=factor.device) + # For FP8, we recommend default KV-cache scaling factor to be 1. The + # asymmetric format applies this only to K; V remains NVFP4. + kv_cache_dtype = get_kv_cache_dtype(self_attention_module) + if kv_cache_dtype == KV_CACHE_FP8: + fp8_indices = range(len(scaling_factors)) + elif kv_cache_dtype == KV_CACHE_FP8_K_NVFP4_V: + fp8_indices = (0,) + else: + fp8_indices = () + for i in fp8_indices: + factor = scaling_factors[i] + if factor is None: + continue + if factor.item() > 0.5: + warn( + f"Warning: Large KV activation detected: {factor.item()}, " + "Quantized KV cache may lead to higher accuracy drop." ) + scaling_factors[i] = torch.max( + factor, torch.tensor([1.0], dtype=torch.float, device=factor.device) + ) return scaling_factors def get_kv_cache_dtype(modules: list[nn.Module] | nn.Module) -> str | None: """Returns the kv_cache dtype. - If num_bits of output_quantizer is (4, 3) then returns FP8; if it is 8, returns int8, - otherwise returns None. + K/V quantizers are inspected as a pair so FP8 K with NVFP4 V remains + distinguishable from uniform FP8 or NVFP4. The output quantizer is retained + as a fallback for the unified Megatron export path. Args: modules: The module or list of modules to inspect. @@ -424,6 +433,29 @@ def get_kv_cache_dtype(modules: list[nn.Module] | nn.Module) -> str | None: modules = [modules] for module in modules: + k_quantizer = getattr(module, "k_bmm_quantizer", None) + v_quantizer = getattr(module, "v_bmm_quantizer", None) + if ( + k_quantizer is not None + and v_quantizer is not None + and k_quantizer.is_enabled + and v_quantizer.is_enabled + ): + k_dtype = _compute_kv_cache_dtype( + [k_quantizer.num_bits], hasattr(k_quantizer, "_bias_value") + ) + v_dtype = _compute_kv_cache_dtype( + [v_quantizer.num_bits], hasattr(v_quantizer, "_bias_value") + ) + if k_dtype == KV_CACHE_FP8 and v_dtype == KV_CACHE_NVFP4: + return KV_CACHE_FP8_K_NVFP4_V + if k_dtype == v_dtype: + return k_dtype + raise NotImplementedError( + "Unsupported mixed K/V cache quantization pair: " + f"K uses {k_dtype}, while V uses {v_dtype}." + ) + # Case where the module has both k_bmm_quantizer and v_bmm_quantizer # Still check for output quantizer for the unified_megatron_export path for quantizer in ("k_bmm_quantizer", "v_bmm_quantizer", "output_quantizer"): @@ -1001,7 +1033,7 @@ def _postprocess_single_tensor( key: str, value: torch.Tensor, kv_cache_max_bound: float, - kv_cache_format: str | None, + kv_cache_format: str | dict[str, dict[str, str]] | None, is_modelopt_qlora: bool = False, ) -> tuple[str | None, torch.Tensor | None]: """Per-tensor subset of :func:`postprocess_state_dict`, for streaming export. @@ -1035,12 +1067,15 @@ def _postprocess_single_tensor( if key.endswith(old_suffix): prefix = key[: -len(old_suffix)] if "_amax" in key: - assert kv_cache_format in [KV_CACHE_FP8, KV_CACHE_NVFP4, KV_CACHE_NVFP4_AFFINE], ( - "Invalid KV cache quantization format." - ) + layer_quantization = _resolve_kv_cache_format_for_key(key, kv_cache_format) + assert layer_quantization in [ + KV_CACHE_FP8, + KV_CACHE_NVFP4, + KV_CACHE_NVFP4_AFFINE, + ], "Invalid KV cache quantization format." assert kv_cache_max_bound > 0, "Maxbound must be greater than zero." value = value.float() / kv_cache_max_bound - if kv_cache_format == KV_CACHE_FP8 and value.item() > 0.5: + if layer_quantization == KV_CACHE_FP8 and value.item() > 0.5: logger.warning( "Large KV activations detected. Quantized KV cache may lead to higher accuracy drop." ) @@ -1051,10 +1086,40 @@ def _postprocess_single_tensor( return None, None +def _resolve_kv_cache_format_for_key( + key: str, quantization: str | dict[str, dict[str, str]] | None +) -> str | None: + """Resolve uniform or per-layer metadata to the format for this K or V tensor.""" + if isinstance(quantization, dict): + matches = [ + (layer_name, layer_config.get("quant_algo")) + for layer_name, layer_config in quantization.items() + if key == layer_name or key.startswith(layer_name + ".") + ] + quantization = max(matches, key=lambda item: len(item[0]))[1] if matches else None + if quantization == KV_CACHE_FP8_K_NVFP4_V: + if key.endswith("k_bmm_quantizer._amax"): + return KV_CACHE_FP8 + if key.endswith("v_bmm_quantizer._amax"): + return KV_CACHE_NVFP4 + return None + return quantization + + +def _get_kv_cache_postprocess_config( + quantization_details: dict[str, Any], +) -> str | dict[str, dict[str, str]] | None: + """Return the uniform format or layer map consumed by both HF exporters.""" + kv_cache_format = quantization_details.get("kv_cache_quant_algo") + if kv_cache_format == "MIXED_PRECISION": + return quantization_details.get("kv_cache_quantized_layers", {}) + return kv_cache_format + + def postprocess_state_dict( state_dict: dict, maxbound: float, - quantization: str | None, + quantization: str | dict[str, dict[str, str]] | None, is_modelopt_qlora: bool = False, tied_map: "TiedWeightMap | None" = None, ) -> dict: @@ -1063,7 +1128,8 @@ def postprocess_state_dict( Args: state_dict: The full model state_dict. maxbound: The maximum bound value for the output quantizer. - quantization: The KV cache quantization format. + quantization: The uniform KV cache quantization format, or a per-attention-layer + ``{layer_name: {"quant_algo": ...}}`` mapping for mixed precision. is_modelopt_qlora: Whether the model is a modelopt-trained QLoRA model. tied_map: Optional :class:`TiedWeightMap`. When provided, tied-weight dedup is authoritative and name-based: a declared alias key whose canonical @@ -1101,15 +1167,18 @@ def _export_key(key: str) -> str: prefix = key[: -len(old_suffix)] if "_amax" in key: - assert quantization in [KV_CACHE_FP8, KV_CACHE_NVFP4, KV_CACHE_NVFP4_AFFINE], ( - "Invalid KV cache quantization format." - ) + layer_quantization = _resolve_kv_cache_format_for_key(key, quantization) + assert layer_quantization in [ + KV_CACHE_FP8, + KV_CACHE_NVFP4, + KV_CACHE_NVFP4_AFFINE, + ], "Invalid KV cache quantization format." assert maxbound > 0, "Maxbound must be greater than zero." value = value.float() / maxbound # Warn if scale exceeds threshold - if quantization == KV_CACHE_FP8 and value.item() > 0.5: + if layer_quantization == KV_CACHE_FP8 and value.item() > 0.5: logger.warning( "Large KV activations detected. Quantized KV cache may lead to higher accuracy drop." ) @@ -1601,7 +1670,7 @@ def get_quant_config( block_size = None # Create base config - quant_config = { + quant_config: dict[str, Any] = { "producer": { "name": "modelopt", "version": __version__, @@ -1619,7 +1688,9 @@ def get_quant_config( # It also holds awq_block_size information for applicable layers. layer_config_dict = {} - kv_cache_format = QUANTIZATION_NONE + kv_cache_formats: set[str] = set() + kv_cache_quantized_layers: dict[str, dict[str, str]] = {} + kv_cache_eligible_layers = 0 for name, module in dict(model.named_modules()).items(): # Check for standard quantizers or any quantizers from weight attributes weight_names = list(weight_attr_names(module)) @@ -1672,21 +1743,18 @@ def get_quant_config( layer_config_dict[name + ".quantization"] = quantization_format layer_config_dict[name + ".awq_block_size"] = block_size - not_enabled = SimpleNamespace(is_enabled=False) - # Find kv cache quant format - if ( - getattr(module, "k_bmm_quantizer", not_enabled).is_enabled - or getattr(module, "v_bmm_quantizer", not_enabled).is_enabled - or getattr(module, "output_quantizer", not_enabled).is_enabled - ): - module_kv_quant = get_kv_cache_dtype(module) - if kv_cache_format == QUANTIZATION_NONE: - kv_cache_format = module_kv_quant - else: - assert kv_cache_format == module_kv_quant, ( - "Do not support mixed precision kv cache quantization" - ) + has_kv_quantizers = all( + hasattr(module, quantizer_name) + for quantizer_name in ("k_bmm_quantizer", "v_bmm_quantizer") + ) + if has_kv_quantizers: + kv_cache_eligible_layers += 1 + if module.k_bmm_quantizer.is_enabled and module.v_bmm_quantizer.is_enabled: + module_kv_quant = get_kv_cache_dtype(module) + if module_kv_quant != QUANTIZATION_NONE: + kv_cache_formats.add(module_kv_quant) + kv_cache_quantized_layers[name] = {"quant_algo": module_kv_quant} # MoE routers/gates are intentionally kept in original precision. On transformers>=5.0 they # are not nn.Linear modules (e.g. TopKRouter), never receive a quantizer, and would otherwise @@ -1699,8 +1767,23 @@ def get_quant_config( # Process per layer quantization config dict quant_config["quantization"].update(process_layer_quant_config(layer_config_dict)) - if kv_cache_format is not None: - quant_config["quantization"]["kv_cache_quant_algo"] = kv_cache_format + all_kv_layers_quantized = ( + kv_cache_eligible_layers > 0 and len(kv_cache_quantized_layers) == kv_cache_eligible_layers + ) + if len(kv_cache_formats) == 1 and all_kv_layers_quantized: + quant_config["quantization"]["kv_cache_quant_algo"] = next(iter(kv_cache_formats)) + elif kv_cache_quantized_layers: + weight_quant_algo = quant_config["quantization"].get("quant_algo") + if weight_quant_algo not in (None, "MIXED_PRECISION"): + raise NotImplementedError( + "Mixed-precision KV-cache export with a uniform quantized-weight format is " + "not supported yet. Use BF16 weights or a mixed-weight AutoQuant recipe." + ) + quant_config["quantization"]["quant_algo"] = "MIXED_PRECISION" + quant_config["quantization"].setdefault("quantized_layers", {}) + quant_config["quantization"]["kv_cache_quant_algo"] = "MIXED_PRECISION" + quant_config["quantization"]["kv_cache_quantized_layers"] = kv_cache_quantized_layers + quant_config["quantization"]["kv_cache_schema_version"] = 1 return quant_config diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index d98211cc3fb..0bea45348b4 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -16,6 +16,7 @@ """Code that export quantized Hugging Face models for deployment.""" import contextlib +import copy import json import re import shutil @@ -103,6 +104,7 @@ revert_weight_conversion_quant_aware, ) from .quant_utils import ( + _get_kv_cache_postprocess_config, fuse_prequant_layernorm, fuse_prequant_to_linear, get_activation_scaling_factor, @@ -1013,12 +1015,13 @@ def _export_transformers_checkpoint( # We define kv cache scale as amax / 448 for both FP8 and NVFP4 KV cache quantization. kv_cache_max_bound = 448 - kv_cache_format = quant_config["quantization"]["kv_cache_quant_algo"] + quantization_details = quant_config["quantization"] + kv_cache_postprocess_config = _get_kv_cache_postprocess_config(quantization_details) quantized_state_dict = postprocess_state_dict( quantized_state_dict, kv_cache_max_bound, - kv_cache_format, + kv_cache_postprocess_config, is_modelopt_qlora, tied_map=tied_map, ) @@ -1492,6 +1495,7 @@ def _write_hf_export_config( model: nn.Module, hf_quant_config: dict | None, export_dir: Path, + name_mapper: Callable[[str], str] | None = None, ) -> None: """Write hf_quant_config.json (if quantized) and embed quantization_config into config.json.""" quantization_details = (hf_quant_config or {}).get("quantization", {}) @@ -1505,6 +1509,27 @@ def _write_hf_export_config( json.dump(hf_quant_config, file, indent=4) quantization_config = convert_hf_quant_config_format(hf_quant_config) + kv_autoquant_report = next( + ( + getattr(module, "_modelopt_kv_cache_auto_quantize_state") + for module in model.modules() + if hasattr(module, "_modelopt_kv_cache_auto_quantize_state") + ), + None, + ) + if kv_autoquant_report is not None: + kv_autoquant_report = copy.deepcopy(kv_autoquant_report) + if name_mapper is not None: + kv_autoquant_report["layers"] = { + name_mapper(name): value + for name, value in kv_autoquant_report["layers"].items() + } + signature_layers = kv_autoquant_report.get("search_signature", {}).get("layers", []) + for layer in signature_layers: + layer["name"] = name_mapper(layer["name"]) + with open(f"{export_dir}/kv_cache_auto_quantize_report.json", "w") as file: + json.dump(kv_autoquant_report, file, indent=4) + original_config = f"{export_dir}/config.json" with open(original_config) as file: config_data = json.load(file) @@ -1602,6 +1627,7 @@ def export_hf_checkpoint( ) if getattr(model, "hf_quantizer", None) is not None: model.hf_quantizer = None + name_mapper = None try: name_mapper = build_reverse_name_mapper(model) if name_mapper is not None and hf_quant_config: @@ -1611,7 +1637,7 @@ def export_hf_checkpoint( f"Quant-aware reverse weight conversion skipped ({exc}); exported tensor " "names may not match the original HF hub checkpoint." ) - _write_hf_export_config(model, hf_quant_config, export_dir) + _write_hf_export_config(model, hf_quant_config, export_dir, name_mapper) return post_state_dict, hf_quant_config = _export_transformers_checkpoint(model, dtype, **kwargs) @@ -1633,6 +1659,7 @@ def export_hf_checkpoint( # and fails). Best-effort and atomic: any failure (an op we cannot reverse yet, # transformers API drift, unexpected shapes) falls back to the in-memory names for BOTH # weights and config so they stay mutually consistent. + name_mapper = None try: name_mapper = build_reverse_name_mapper(model) export_state_dict = revert_weight_conversion_quant_aware(model, export_state_dict) @@ -1667,7 +1694,7 @@ def export_hf_checkpoint( finally: _unpatch_revert_weight_conversion(_patches) - _write_hf_export_config(model, hf_quant_config, export_dir) + _write_hf_export_config(model, hf_quant_config, export_dir, name_mapper) except Exception as e: warnings.warn( diff --git a/modelopt/torch/export/unified_export_hf_streaming.py b/modelopt/torch/export/unified_export_hf_streaming.py index a6493aac44c..a52670eaa89 100644 --- a/modelopt/torch/export/unified_export_hf_streaming.py +++ b/modelopt/torch/export/unified_export_hf_streaming.py @@ -32,7 +32,11 @@ from safetensors.torch import save_file from .quant_aware_conversion import build_reverse_name_mapper -from .quant_utils import _postprocess_single_tensor, get_quant_config +from .quant_utils import ( + _get_kv_cache_postprocess_config, + _postprocess_single_tensor, + get_quant_config, +) from .registry import ExportContext from .unified_export_hf import ( _add_mtp_exclusions, @@ -249,7 +253,7 @@ def _export_transformers_checkpoint_streaming( # --- Per-tensor constants --- kv_cache_max_bound = 448 - kv_cache_format = quant_config["quantization"]["kv_cache_quant_algo"] + kv_cache_format = _get_kv_cache_postprocess_config(quant_config["quantization"]) # --- Tied alias keys to skip --- # data_ptr() is unreliable for disk-offloaded weights, so we use _tied_weights_keys. diff --git a/modelopt/torch/quantization/_auto_quantize_cost.py b/modelopt/torch/quantization/_auto_quantize_cost.py index 297e8cd8dc3..3834e264500 100644 --- a/modelopt/torch/quantization/_auto_quantize_cost.py +++ b/modelopt/torch/quantization/_auto_quantize_cost.py @@ -31,6 +31,7 @@ EXCLUDED_MODULE_NAME_PATTERNS_KEY: Final = "excluded_module_name_patterns" COST_MODEL_WEIGHT: Final = "weight" COST_MODEL_ACTIVE_MOE: Final = "active_moe" +COST_MODEL_KV_CACHE: Final = "kv_cache" _ROUTED_MOE_EXPERT_NAME_RE = re.compile(r"(^|\.)experts(\.|$)") _ACTIVE_MOE_TOP_K_ATTRS = ( @@ -68,7 +69,7 @@ def _iter_model_configs(model: nn.Module): def _get_first_numeric_config_attr(config: Any, attr_names: tuple[str, ...]) -> float | None: for attr_name in attr_names: value = getattr(config, attr_name, None) - if isinstance(value, (int, float)) and not isinstance(value, bool): + if isinstance(value, int | float) and not isinstance(value, bool): return float(value) return None @@ -185,7 +186,7 @@ def normalize_cost_constraints( ) if not ( - isinstance(active_moe_expert_ratio, (int, float)) + isinstance(active_moe_expert_ratio, int | float) and not isinstance(active_moe_expert_ratio, bool) and 0.0 < active_moe_expert_ratio <= 1.0 ): @@ -206,9 +207,22 @@ def module_cost_weight( return base_weight +class KVCacheCostModel(AutoQuantizeCostModel): + """Account for the exact storage of paired K/V-cache candidates.""" + + name = COST_MODEL_KV_CACHE + supported_cost_keys = frozenset() + + @staticmethod + def candidate_cost(k_width: int, v_width: int, k_bits: float, v_bits: float) -> float: + """Return candidate storage in bits for one attention layer and token.""" + return k_width * k_bits + v_width * v_bits + + _COST_MODELS: Final = { COST_MODEL_WEIGHT: WeightCostModel(), COST_MODEL_ACTIVE_MOE: ActiveMoECostModel(), + COST_MODEL_KV_CACHE: KVCacheCostModel(), } diff --git a/modelopt/torch/quantization/kv_cache_auto_quant.py b/modelopt/torch/quantization/kv_cache_auto_quant.py new file mode 100644 index 00000000000..dc89a0a628c --- /dev/null +++ b/modelopt/torch/quantization/kv_cache_auto_quant.py @@ -0,0 +1,907 @@ +# 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. + +"""Layer-wise KV-cache AutoQuant using isolated forward KL sensitivity.""" + +from __future__ import annotations + +import fnmatch +import math +from contextlib import contextmanager +from typing import TYPE_CHECKING, Any, cast + +import torch +import torch.nn as nn +import torch.nn.functional as F +from tqdm import tqdm + +from modelopt.torch.opt.hparam import Hparam +from modelopt.torch.opt.searcher import LPS, BaseSearcher, SearchConfig, SearchStateDict +from modelopt.torch.utils import print_rank_0 + +from ._auto_quantize_cost import ( + COST_MODEL_KV_CACHE, + KVCacheCostModel, + get_auto_quantize_cost_model, + normalize_auto_quantize_constraints, +) +from .config import QuantizeConfig +from .conversion import set_quantizer_by_cfg +from .nn import TensorQuantizer + +__all__ = ["AutoQuantizeKVSearcher"] + +if TYPE_CHECKING: + from collections.abc import Callable + +_KV_QUANTIZER_ATTRS = ("k_bmm_quantizer", "v_bmm_quantizer") +_KV_AUTOQUANT_SCHEMA_VERSION = 1 +_KV_CANDIDATE_HOLDER_NAME = "layer" +_KV_CANDIDATE_NAMES = {f"{_KV_CANDIDATE_HOLDER_NAME}.{attr}" for attr in _KV_QUANTIZER_ATTRS} +_NON_KV_PROBE_NAMES = { + f"{_KV_CANDIDATE_HOLDER_NAME}.{name}" + for name in ( + "q_bmm_quantizer", + "p_bmm_quantizer", + "input_quantizer", + "output_quantizer", + "q_proj.input_quantizer", + "q_proj.weight_quantizer", + "q_proj.output_quantizer", + ) +} + + +def _disabled_quantizer() -> TensorQuantizer: + quantizer = TensorQuantizer() + quantizer.disable() + return quantizer + + +def _candidate_quantizers(config: QuantizeConfig) -> dict[str, TensorQuantizer]: + """Build a candidate with the same qualified K/V names used by a full model.""" + _validate_candidate_patterns(config) + root = nn.Module() + holder = nn.Module() + root.add_module(_KV_CANDIDATE_HOLDER_NAME, holder) + for attr in _KV_QUANTIZER_ATTRS: + setattr(holder, attr, _disabled_quantizer()) + set_quantizer_by_cfg(root, config.quant_cfg) + quantizers = {attr: getattr(holder, attr) for attr in _KV_QUANTIZER_ATTRS} + for attr, quantizer in quantizers.items(): + if not isinstance(quantizer, TensorQuantizer) or not quantizer.is_enabled: + raise ValueError( + f"KV-cache candidate must enable {attr}; got {type(quantizer).__name__}." + ) + return quantizers + + +def _validate_candidate_patterns(config: QuantizeConfig) -> None: + """Require every ordered config entry to match only a qualified K/V name.""" + matched_names: set[str] = set() + probe_names = _KV_CANDIDATE_NAMES | _NON_KV_PROBE_NAMES + for entry in config.quant_cfg: + if entry.parent_class is not None: + raise ValueError("KV-cache AutoQuant candidates do not support parent_class filters.") + matches = {name for name in probe_names if fnmatch.fnmatch(name, entry.quantizer_name)} + if not matches: + raise ValueError( + "KV-cache AutoQuant candidate pattern " + f"{entry.quantizer_name!r} does not match a supported qualified K/V quantizer." + ) + non_kv_matches = matches - _KV_CANDIDATE_NAMES + if non_kv_matches: + raise ValueError( + "KV-cache AutoQuant candidates may configure only k_bmm_quantizer and " + f"v_bmm_quantizer; pattern {entry.quantizer_name!r} also matches " + f"{sorted(non_kv_matches)}." + ) + matched_names.update(matches) + if matched_names != _KV_CANDIDATE_NAMES: + raise ValueError( + "KV-cache AutoQuant candidates must completely configure both " + "k_bmm_quantizer and v_bmm_quantizer." + ) + + +def _algorithm_method(config: QuantizeConfig) -> str | None: + algorithm = config.algorithm + if algorithm is None or isinstance(algorithm, str): + return algorithm + if isinstance(algorithm, dict): + return algorithm.get("method") + return getattr(algorithm, "method", None) + + +def _deployable_kv_bits(quantizer: TensorQuantizer) -> float: + """Return storage bits for the narrow K/V formats supported by unified export.""" + if quantizer.bias is not None: + raise ValueError("KV-cache AutoQuant does not support affine candidates yet.") + if quantizer.is_fp8: + return 8.0 + if quantizer.is_nvfp4_dynamic and quantizer.block_sizes.get(-1) == 16: + return 4.5 + raise ValueError( + "KV-cache AutoQuant candidates must use unified-export-compatible per-tensor FP8 " + "or block-16 dynamic NVFP4 quantizers." + ) + + +def _candidate_kv_bits(config: QuantizeConfig) -> tuple[float, float]: + quantizers = _candidate_quantizers(config) + return ( + _deployable_kv_bits(quantizers["k_bmm_quantizer"]), + _deployable_kv_bits(quantizers["v_bmm_quantizer"]), + ) + + +def _validate_deployable_candidate(config: QuantizeConfig) -> None: + quantizers = _candidate_quantizers(config) + k_quantizer = quantizers["k_bmm_quantizer"] + v_quantizer = quantizers["v_bmm_quantizer"] + k_bits = _deployable_kv_bits(k_quantizer) + v_bits = _deployable_kv_bits(v_quantizer) + if k_bits != v_bits and (k_bits, v_bits) != (8.0, 4.5): + raise ValueError( + "Unified export supports only uniform FP8, uniform NVFP4, or FP8-K/NVFP4-V " + "KV-cache AutoQuant candidates." + ) + + algorithm_method = _algorithm_method(config) + for attr, quantizer in quantizers.items(): + if quantizer._dynamic: + raise ValueError( + f"KV-cache AutoQuant candidate {attr} uses top-level dynamic quantization, " + "which does not retain a persistent export scale." + ) + will_calibrate = algorithm_method == "max" and not quantizer._use_constant_amax + if not hasattr(quantizer, "_amax") and not will_calibrate: + raise ValueError( + f"KV-cache AutoQuant candidate {attr} has no persistent export scale. " + "Use max calibration or constant_amax; dynamic and use_constant_amax-only " + "candidates cannot be exported." + ) + + assert config.effective_bits is not None + actual_effective_bits = (k_bits + v_bits) / 2.0 + if not math.isclose(config.effective_bits, actual_effective_bits, rel_tol=0.0, abs_tol=1e-12): + raise ValueError( + "KV-cache AutoQuant candidate effective_bits does not match its configured K/V " + f"storage cost: declared {config.effective_bits}, actual {actual_effective_bits}." + ) + + +def _validate_kv_only_config(config: QuantizeConfig) -> None: + if config.effective_bits is None: + raise ValueError( + "Each KV-cache AutoQuant candidate must declare config-level effective_bits." + ) + algorithm_method = _algorithm_method(config) + if algorithm_method != "max": + if algorithm_method is not None: + raise ValueError( + "KV-cache AutoQuant supports only non-structural calibration algorithms " + f"None and 'max'; got {algorithm_method!r}." + ) + _validate_deployable_candidate(config) + + +def _validate_search_inputs( + constraints: dict[str, Any], + quantization_formats: list[tuple[dict[str, Any], str]], + num_calib_steps: int, + num_score_steps: int, +) -> tuple[float, list[tuple[str, QuantizeConfig]]]: + """Validate a KV-cache search before the caller converts the model.""" + if ( + set(constraints) - {"effective_bits", "cost_model", "cost"} + or constraints.get("cost_model") != COST_MODEL_KV_CACHE + or "effective_bits" not in constraints + or constraints.get("cost") not in (None, {}) + ): + raise ValueError( + "KV-cache AutoQuant requires an effective_bits target with " + f"cost_model='kv_cache'; got {constraints}." + ) + target_bits = float(constraints["effective_bits"]) + if not (0 < target_bits <= 16): + raise ValueError(f"effective_bits must be in (0, 16], got {target_bits}.") + if num_calib_steps <= 0: + raise ValueError("num_calib_steps must be positive.") + if num_score_steps <= 0: + raise ValueError("num_score_steps must be positive.") + + candidates = [] + seen_names = set() + for raw_config, name in quantization_formats: + if name in seen_names: + raise ValueError(f"Duplicate KV-cache AutoQuant candidate name: {name!r}.") + config = QuantizeConfig(**raw_config) + _validate_kv_only_config(config) + candidates.append((name, config)) + seen_names.add(name) + if not candidates: + raise ValueError("KV-cache AutoQuant requires at least one candidate format.") + return target_bits, candidates + + +def _projection_width(module: nn.Module, side: str) -> int | None: + projection = getattr(module, f"{side}_proj", None) + out_features = getattr(projection, "out_features", None) + if isinstance(out_features, int) and out_features > 0: + return out_features + + config = getattr(module, "config", None) + num_kv_heads = getattr(config, "num_key_value_heads", None) + head_dim = getattr(module, "head_dim", None) or getattr(config, "head_dim", None) + if not isinstance(head_dim, int): + hidden_size = getattr(config, "hidden_size", None) + num_heads = getattr(config, "num_attention_heads", None) + if isinstance(hidden_size, int) and isinstance(num_heads, int) and num_heads > 0: + head_dim = hidden_size // num_heads + if isinstance(num_kv_heads, int) and isinstance(head_dim, int): + return num_kv_heads * head_dim + return None + + +def _kv_scalar_weight(module: nn.Module, name: str) -> int: + k_width = _projection_width(module, "k") + v_width = _projection_width(module, "v") + if k_width is None or v_width is None: + raise ValueError( + "Cannot determine exact KV width for eligible attention layer " + f"{name!r}. Expected k_proj/v_proj.out_features or config " + "num_key_value_heads plus head_dim." + ) + return k_width + v_width + + +def _validate_candidate_cost_geometry( + candidates: list[tuple[str, QuantizeConfig]], + layers: list[tuple[str, nn.Module, int]], +) -> None: + candidate_bits = [_candidate_kv_bits(config) for _, config in candidates] + if all(k_bits == v_bits for k_bits, v_bits in candidate_bits): + return + + unequal_width_layers = [] + for name, module, _ in layers: + k_width = _projection_width(module, "k") + v_width = _projection_width(module, "v") + if k_width != v_width: + unequal_width_layers.append(f"{name} (K={k_width}, V={v_width})") + if unequal_width_layers: + raise ValueError( + "KV-cache AutoQuant cannot cost asymmetric K/V candidates on layers with unequal " + "K/V widths: " + ", ".join(unequal_width_layers) + "." + ) + + +def _eligible_layers( + model: nn.Module, disabled_layers: list[str] | str | None +) -> list[tuple[str, nn.Module, int]]: + patterns = [disabled_layers] if isinstance(disabled_layers, str) else disabled_layers or [] + boundaries = [] + names_by_identity: dict[int, list[str]] = {} + for name, module in model.named_modules(remove_duplicate=False): + if not all(hasattr(module, attr) for attr in _KV_QUANTIZER_ATTRS): + continue + boundaries.append((name, module)) + names_by_identity.setdefault(id(module), []).append(name) + aliases = [names for names in names_by_identity.values() if len(names) > 1] + if aliases: + raise ValueError( + f"KV-cache attention boundaries are registered through aliases: {aliases}." + ) + + layers = [] + for name, module in boundaries: + if any(fnmatch.fnmatch(name, pattern) for pattern in patterns): + continue + layers.append((name, module, _kv_scalar_weight(module, name))) + if not layers: + raise ValueError("KV-cache AutoQuant found no eligible attention layers.") + return layers + + +def _apply_layer_quantizers(module: nn.Module, quantizers: dict[str, TensorQuantizer]) -> None: + for attr, quantizer in quantizers.items(): + setattr(module, attr, quantizer) + + +@contextmanager +def _freeze_existing_quantizers(model: nn.Module, candidate_quantizers: list[TensorQuantizer]): + """Freeze calibration without changing existing quantizers' execution mode.""" + candidate_ids = {id(quantizer) for quantizer in candidate_quantizers} + states = [] + for module in model.modules(): + if ( + not isinstance(module, TensorQuantizer) + or id(module) in candidate_ids + or not module.is_enabled + ): + continue + states.append((module, module._if_calib)) + module.disable_calib() + try: + yield + finally: + for quantizer, if_calib in states: + quantizer._if_calib = if_calib + + +def _get_logits( + forward_step: Callable[[nn.Module, Any], torch.Tensor], model: nn.Module, data: Any +) -> torch.Tensor: + logits = forward_step(model, data) + if not isinstance(logits, torch.Tensor): + raise TypeError("KV-cache AutoQuant forward_step must return a logits tensor.") + if logits.ndim < 2 or logits.shape[-1] == 0: + raise ValueError( + "KV-cache AutoQuant forward_step must return logits with a non-empty vocabulary " + "dimension." + ) + if not torch.isfinite(logits).all(): + raise ValueError("KV-cache AutoQuant encountered NaN or Inf logits.") + return logits + + +def _solve_additive_recipe( + layer_names: list[str], + layer_widths: list[tuple[int, int]], + candidate_names: list[str], + candidate_kv_bits: list[tuple[float, float]], + scores: list[list[float]], + target_bits: float, + verbose: bool, +) -> tuple[list[int], str]: + cost_model = get_auto_quantize_cost_model(COST_MODEL_KV_CACHE) + assert isinstance(cost_model, KVCacheCostModel) + denominator = float(sum(k_width + v_width for k_width, v_width in layer_widths)) + candidate_costs = [ + [ + cost_model.candidate_cost(k_width, v_width, k_bits, v_bits) / 16.0 + for k_bits, v_bits in candidate_kv_bits + ] + for k_width, v_width in layer_widths + ] + max_cost = denominator * target_bits / 16.0 + lps = LPS( + name="KVCacheAutoQuant", + constraints={"kv_cache_size_after_compression": max_cost}, + constraints_to_candidate_costs={"kv_cache_size_after_compression": candidate_costs}, + candidate_scores=scores, + objective_type="minimize", + verbose=verbose, + ) + selections, status = lps() + if status != "Optimal": + minimum_bits = sum(min(costs) for costs in candidate_costs) * 16.0 / denominator + raise ValueError( + f"KV-cache AutoQuant could not satisfy effective_bits={target_bits}; " + f"minimum achievable value is {minimum_bits:.4f}. Solver status: {status}." + ) + if len(selections) != len(layer_names): + raise RuntimeError( + "KV-cache AutoQuant solver returned an invalid selection count: " + f"{len(selections)} for {len(layer_names)} layers and candidates {candidate_names}." + ) + return selections, status + + +def _search_signature( + candidates: list[tuple[str, QuantizeConfig]], + layers: list[tuple[str, nn.Module, int]], + num_calib_steps: int, + num_score_steps: int, +) -> dict[str, Any]: + return { + "schema_version": _KV_AUTOQUANT_SCHEMA_VERSION, + "num_calib_steps": num_calib_steps, + "num_score_steps": num_score_steps, + "candidates": [ + { + "name": name, + "config": config.model_dump(mode="json", exclude_none=True), + } + for name, config in candidates + ], + "layers": [ + { + "name": name, + "k_width": _projection_width(module, "k"), + "v_width": _projection_width(module, "v"), + } + for name, module, _ in layers + ], + } + + +def _checkpoint_state_is_compatible(state: dict[str, Any], signature: dict[str, Any]) -> bool: + return state.get("search_signature") == signature + + +def _quantizer_state_dict( + candidate_quantizers: dict[str, dict[str, dict[str, TensorQuantizer]]], +) -> dict[str, dict[str, dict[str, dict[str, torch.Tensor]]]]: + return { + layer_name: { + candidate_name: { + attr: quantizer.state_dict() for attr, quantizer in layer_quantizers.items() + } + for candidate_name, layer_quantizers in layer_candidates.items() + } + for layer_name, layer_candidates in candidate_quantizers.items() + } + + +def _restore_quantizer_state_dict( + candidate_quantizers: dict[str, dict[str, dict[str, TensorQuantizer]]], + state: dict[str, dict[str, dict[str, dict[str, torch.Tensor]]]], +) -> None: + """Restore calibration buffers into config-created candidate quantizers.""" + for layer_name, layer_candidates in candidate_quantizers.items(): + for candidate_name, layer_quantizers in layer_candidates.items(): + for attr, quantizer in layer_quantizers.items(): + quantizer_state = state[layer_name][candidate_name][attr] + for key, value in quantizer_state.items(): + if "." not in key and key not in quantizer._buffers: + quantizer.register_buffer(key, torch.empty_like(value)) + quantizer.load_state_dict(quantizer_state) + + +def _validate_persistent_candidate_scales( + candidate_quantizers: dict[str, dict[str, dict[str, TensorQuantizer]]], +) -> None: + """Require every calibrated candidate scale to be persistent in its state dict.""" + for layer_name, layer_candidates in candidate_quantizers.items(): + for candidate_name, layer_quantizers in layer_candidates.items(): + for attr, quantizer in layer_quantizers.items(): + if "_amax" not in quantizer.state_dict(): + raise ValueError( + f"KV-cache AutoQuant candidate {candidate_name!r} for " + f"{layer_name!r}/{attr} has no persistent export scale after calibration." + ) + + +def _report_state(state: dict[str, Any]) -> dict[str, Any]: + """Return the JSON-safe search report, excluding calibration tensors.""" + return {key: value for key, value in state.items() if key != "quantizer_state"} + + +class QuantKVRecipeHparam(Hparam): + """One paired K/V format decision for an attention layer.""" + + def __init__( + self, + name: str, + module: nn.Module, + candidates: list[tuple[str, QuantizeConfig]], + ) -> None: + super().__init__(range(len(candidates)), original=0) + self.name = name + self.module = module + self.candidates = candidates + self.original_quantizers = {attr: getattr(module, attr) for attr in _KV_QUANTIZER_ATTRS} + self.reference_quantizers = {attr: _disabled_quantizer() for attr in _KV_QUANTIZER_ATTRS} + self.candidate_quantizers = { + index: _candidate_quantizers(config) for index, (_, config) in enumerate(candidates) + } + k_width = _projection_width(module, "k") + v_width = _projection_width(module, "v") + assert k_width is not None and v_width is not None + self.k_width = k_width + self.v_width = v_width + self.use_reference() + + @property + def active(self) -> int: + """Return the selected candidate index.""" + assert isinstance(self._active, int) + return self._active + + @active.setter + def active(self, value: int | None) -> None: + if value is None: + assert isinstance(self.original, int) + value = self.original + assert value in self.choices + self._active = value + _apply_layer_quantizers(self.module, self.candidate_quantizers[value]) + + def use_reference(self) -> None: + """Use BF16/no-quant K/V as the scoring reference, never as a solver choice.""" + _apply_layer_quantizers(self.module, self.reference_quantizers) + + def restore_original(self) -> None: + """Restore the K/V quantizer objects present before search.""" + _apply_layer_quantizers(self.module, self.original_quantizers) + + def candidate_name(self, index: int) -> str: + return self.candidates[index][0] + + def candidate_bits(self, index: int) -> tuple[float, float]: + return _candidate_kv_bits(self.candidates[index][1]) + + def candidate_cost(self, index: int, cost_model: KVCacheCostModel) -> float: + k_bits, v_bits = self.candidate_bits(index) + return cost_model.candidate_cost(self.k_width, self.v_width, k_bits, v_bits) + + +class AutoQuantizeKVSearcher(BaseSearcher): + """KV-cache AutoQuant backend using the shared search/checkpoint lifecycle.""" + + method_name = "kl_div" + + @property + def default_search_config(self) -> SearchConfig: + """Return KV-specific fields layered on the shared search configuration.""" + config = super().default_search_config + config.update( + { + "quantization_formats": [], + "forward_step": None, + "num_calib_steps": 512, + "num_score_steps": 128, + "disabled_layers": None, + } + ) + return config + + @property + def default_state_dict(self) -> SearchStateDict: + """Return the checkpointed KV search state.""" + return { + "schema_version": _KV_AUTOQUANT_SCHEMA_VERSION, + "method": self.method_name, + "cost_model": COST_MODEL_KV_CACHE, + "search_signature": None, + "calibration_complete": False, + "num_calib_steps": 0, + "score_reduction": "mean_per_scored_token", + "num_score_steps": 0, + "num_scored_tokens": 0, + "candidates": [], + "layers": {}, + "quantizer_state": {}, + "requested_constraints": {}, + "best": { + "recipe": {}, + "constraints": {}, + "score": float("inf"), + "is_satisfied": False, + "solver_status": None, + }, + } + + def sanitize_search_config(self, config: SearchConfig | None) -> SearchConfig: + """Validate the data inputs required by isolated forward-KL scoring.""" + config = super().sanitize_search_config(config) + if config["data_loader"] is None: + raise ValueError("data_loader must be provided for KV-cache AutoQuant.") + if config["forward_step"] is None: + raise ValueError("forward_step must be provided for KV-cache AutoQuant.") + return config + + def load_search_checkpoint(self) -> bool: + """Load compatible fields before validating the resolved KV search signature.""" + return super().load_search_checkpoint(strict=False) + + @property + def _candidate_quantizer_map( + self, + ) -> dict[str, dict[str, dict[str, TensorQuantizer]]]: + return { + hparam.name: { + hparam.candidate_name(index): quantizers + for index, quantizers in hparam.candidate_quantizers.items() + } + for hparam in self._hparams + } + + def restore_original_quantizers(self) -> None: + """Restore pre-search K/V objects after a failed search.""" + for hparam in getattr(self, "_hparams", []): + hparam.restore_original() + + def _calibrate_candidates(self) -> None: + from .model_quant import calibrate + + data_loader = self.config["data_loader"] + forward_step = self.config["forward_step"] + num_calib_steps = self.config["num_calib_steps"] + for candidate_index, (_, config) in enumerate(self._candidates): + for hparam in self._hparams: + hparam.active = candidate_index + + if config.algorithm is not None: + + def calibration_loop(calibration_model): + for step, data in enumerate(data_loader): + if step >= num_calib_steps: + break + _get_logits(forward_step, calibration_model, data) + + active_quantizers = [ + quantizer + for hparam in self._hparams + for quantizer in hparam.candidate_quantizers[candidate_index].values() + ] + calibration_proxy = nn.Module() + calibration_proxy.quantizers = nn.ModuleList(active_quantizers) + with _freeze_existing_quantizers(self.model, active_quantizers): + calibrate( + calibration_proxy, + algorithm=config.algorithm, + forward_loop=lambda _: calibration_loop(self.model), + ) + + for hparam in self._hparams: + hparam.use_reference() + + candidate_quantizers = self._candidate_quantizer_map + _validate_persistent_candidate_scales(candidate_quantizers) + self.quantizer_state = _quantizer_state_dict(candidate_quantizers) + self.calibration_complete = True + self.num_calib_steps = num_calib_steps + self.save_search_checkpoint(verbose=self.config["verbose"]) + + def before_search(self) -> None: + """Resolve attention decisions and calibrate or restore candidate scales.""" + super().before_search() + self.constraints = normalize_auto_quantize_constraints(self.model, self.constraints) + target_bits, self._candidates = _validate_search_inputs( + self.constraints, + self.config["quantization_formats"], + self.config["num_calib_steps"], + self.config["num_score_steps"], + ) + self._target_bits = target_bits + layers = _eligible_layers(self.model, self.config["disabled_layers"]) + _validate_candidate_cost_geometry(self._candidates, layers) + signature = _search_signature( + self._candidates, + layers, + self.config["num_calib_steps"], + self.config["num_score_steps"], + ) + if self.search_signature is not None and not _checkpoint_state_is_compatible( + self.state_dict(), signature + ): + raise ValueError( + "KV-cache AutoQuant checkpoint does not match the current candidates, scoring " + "setup, or eligible layers. Use a different checkpoint path." + ) + self.search_signature = signature + self._hparams = [ + QuantKVRecipeHparam(name, module, self._candidates) for name, module, _ in layers + ] + self._cost_model = get_auto_quantize_cost_model(COST_MODEL_KV_CACHE) + assert isinstance(self._cost_model, KVCacheCostModel) + self.candidates = [ + { + "name": name, + "effective_bits": config.effective_bits, + "k_bits": _candidate_kv_bits(config)[0], + "v_bits": _candidate_kv_bits(config)[1], + "config": config.model_dump(mode="python", exclude_none=True), + } + for name, config in self._candidates + ] + self.model.eval() + + if self.calibration_complete: + if not self.quantizer_state: + raise ValueError( + "KV-cache AutoQuant checkpoint is missing calibrated quantizer state. " + "Use a different checkpoint path." + ) + candidate_quantizers = self._candidate_quantizer_map + _restore_quantizer_state_dict(candidate_quantizers, self.quantizer_state) + _validate_persistent_candidate_scales(candidate_quantizers) + if self.config["verbose"]: + print_rank_0("KV-cache AutoQuant restored calibration from checkpoint.") + else: + self._calibrate_candidates() + + def _estimate_sensitivity_scores(self) -> None: + candidate_names = [name for name, _ in self._candidates] + score_sums: dict[str, dict[str, torch.Tensor | None]] = { + hparam.name: dict.fromkeys(candidate_names) for hparam in self._hparams + } + scored_tokens = 0 + scored_steps = 0 + iterator = tqdm( + self.config["data_loader"], + total=self.config["num_score_steps"], + desc="Estimating KV-cache KL sensitivity", + disable=not self.config["verbose"], + ) + for data in iterator: + if scored_steps >= self.config["num_score_steps"]: + break + logits_ref = _get_logits(self.config["forward_step"], self.model, data) + log_prob_ref = torch.log_softmax(logits_ref.float(), dim=-1) + scored_tokens += logits_ref.numel() // logits_ref.shape[-1] + + for hparam in self._hparams: + for candidate_index, candidate_name in enumerate(candidate_names): + hparam.active = candidate_index + logits_quant = _get_logits(self.config["forward_step"], self.model, data) + if logits_quant.shape != logits_ref.shape: + raise ValueError( + "KV-cache AutoQuant forward_step returned different reference and " + f"candidate logits shapes: {tuple(logits_ref.shape)} and " + f"{tuple(logits_quant.shape)}." + ) + score = F.kl_div( + torch.log_softmax(logits_quant.float(), dim=-1), + log_prob_ref, + reduction="sum", + log_target=True, + ) + previous_score = score_sums[hparam.name][candidate_name] + score_sums[hparam.name][candidate_name] = ( + score if previous_score is None else previous_score + score + ) + hparam.use_reference() + scored_steps += 1 + + if scored_steps == 0 or scored_tokens == 0: + raise ValueError("KV-cache AutoQuant data_loader produced no scoring batches.") + + self.layers = {} + for hparam in self._hparams: + scores = {} + for candidate_name in candidate_names: + score_sum = score_sums[hparam.name][candidate_name] + if score_sum is None: + raise RuntimeError( + "KV-cache AutoQuant did not collect a score for " + f"{hparam.name!r}/{candidate_name!r}." + ) + if not torch.isfinite(score_sum): + raise ValueError( + "KV-cache AutoQuant produced a non-finite KL score for " + f"{hparam.name!r}/{candidate_name!r}." + ) + scores[candidate_name] = float(score_sum.item()) / scored_tokens + self.layers[hparam.name] = { + "k_width": hparam.k_width, + "v_width": hparam.v_width, + "scores": scores, + } + self.num_score_steps = scored_steps + self.num_scored_tokens = scored_tokens + self.save_search_checkpoint(verbose=self.config["verbose"]) + + def _solve(self) -> None: + candidate_names = [name for name, _ in self._candidates] + candidate_kv_bits = [_candidate_kv_bits(config) for _, config in self._candidates] + layers = cast("dict[str, dict[str, Any]]", self.layers) + scores = [ + [layers[hparam.name]["scores"][name] for name in candidate_names] + for hparam in self._hparams + ] + selections, status = _solve_additive_recipe( + [hparam.name for hparam in self._hparams], + [(hparam.k_width, hparam.v_width) for hparam in self._hparams], + candidate_names, + candidate_kv_bits, + scores, + self._target_bits, + self.config["verbose"], + ) + denominator = float(sum(hparam.k_width + hparam.v_width for hparam in self._hparams)) + cost_model = self._cost_model + assert isinstance(cost_model, KVCacheCostModel) + total_cost = sum( + ( + hparam.candidate_cost(selected, cost_model) + for hparam, selected in zip(self._hparams, selections) + ), + start=0.0, + ) + achieved_bits = total_cost / denominator + selected_score = sum( + layer_scores[selected] for selected, layer_scores in zip(selections, scores) + ) + recipe = {} + for hparam, selected in zip(self._hparams, selections): + hparam.active = selected + selected_name = hparam.candidate_name(selected) + self.layers[hparam.name]["selected"] = selected_name + recipe[hparam.name] = selected_name + if self.config["verbose"]: + print_rank_0(f"KV-cache AutoQuant selected {selected_name} for {hparam.name}.") + self.requested_constraints = { + "effective_bits": self._target_bits, + "cost_model": COST_MODEL_KV_CACHE, + } + self.best = { + "recipe": recipe, + "constraints": { + "effective_bits": achieved_bits, + "cost_model": COST_MODEL_KV_CACHE, + }, + "score": selected_score, + "is_satisfied": achieved_bits <= self._target_bits + 1e-12, + "solver_status": status, + } + + @torch.inference_mode() + def run_search(self) -> None: + """Score candidates when needed, solve the budget, and apply the selection.""" + if not self.layers: + self._estimate_sensitivity_scores() + self._solve() + self.save_search_checkpoint(verbose=self.config["verbose"]) + + def after_search(self) -> None: + """Attach the JSON-safe sensitivity report used by unified export.""" + self.model._modelopt_kv_cache_auto_quantize_state = _report_state(self.state_dict()) + + +def _config_entry_dict(entry: Any) -> dict[str, Any]: + if hasattr(entry, "model_dump"): + return entry.model_dump(mode="python", exclude_none=True) + return dict(entry) + + +def get_kv_cache_auto_quantize_config( + search_state: dict[str, Any], + constraints: dict[str, Any] | None = None, + verbose: bool = False, +) -> dict[str, Any]: + """Build a flat K/V quantization config, optionally re-solving at a new target.""" + requested = constraints or search_state.get("requested_constraints") + if not isinstance(requested, dict): + raise ValueError("KV-cache AutoQuant search state has no requested constraints.") + target_bits = float(requested["effective_bits"]) + if requested.get("cost_model", COST_MODEL_KV_CACHE) != COST_MODEL_KV_CACHE: + raise ValueError("KV-cache search state can only be re-solved with cost_model='kv_cache'.") + + candidates = search_state["candidates"] + candidate_names = [candidate["name"] for candidate in candidates] + candidate_kv_bits = [ + (float(candidate["k_bits"]), float(candidate["v_bits"])) for candidate in candidates + ] + layers = search_state["layers"] + layer_names = list(layers) + selections, _ = _solve_additive_recipe( + layer_names, + [(layers[name]["k_width"], layers[name]["v_width"]) for name in layer_names], + candidate_names, + candidate_kv_bits, + [ + [layers[name]["scores"][candidate] for candidate in candidate_names] + for name in layer_names + ], + target_bits, + verbose, + ) + + quant_cfg: list[dict[str, Any]] = [] + for layer_name, selected in zip(layer_names, selections): + config = QuantizeConfig(**candidates[selected]["config"]) + for entry in config.quant_cfg: + entry_dict = _config_entry_dict(entry) + pattern = entry_dict["quantizer_name"] + for attr in _KV_QUANTIZER_ATTRS: + if fnmatch.fnmatch(f"{_KV_CANDIDATE_HOLDER_NAME}.{attr}", pattern): + resolved = dict(entry_dict) + resolved["quantizer_name"] = f"{layer_name}.{attr}" + quant_cfg.append(resolved) + return {"quant_cfg": quant_cfg, "algorithm": "max"} diff --git a/modelopt/torch/quantization/model_quant.py b/modelopt/torch/quantization/model_quant.py index 3f6040fd4ef..f14df6fb466 100644 --- a/modelopt/torch/quantization/model_quant.py +++ b/modelopt/torch/quantization/model_quant.py @@ -15,6 +15,7 @@ """User-facing quantization API.""" +import copy import fnmatch import inspect import os @@ -38,9 +39,12 @@ ) from modelopt.torch.utils import atomic_print +from ._auto_quantize_cost import COST_MODEL_KV_CACHE from .algorithms import AutoQuantizeGradientSearcher, AutoQuantizeKLDivSearcher, QuantRecipe from .algorithms import get_auto_quantize_config as _get_auto_quantize_config from .config import QuantizeAlgoCfgType +from .kv_cache_auto_quant import AutoQuantizeKVSearcher, get_kv_cache_auto_quantize_config +from .kv_cache_auto_quant import _validate_search_inputs as _validate_kv_cache_search_inputs from .mode import QuantizeModeRegistry, get_modelike_from_algo_cfg from .nn import QuantModule, SequentialQuantizer, TensorQuantizer from .utils import is_quantized @@ -272,7 +276,7 @@ def forward_loop(model) -> None: def auto_quantize( model: nn.Module, constraints: dict[str, Any] | None = None, - quantization_formats: list[dict[str, Any] | str] | None = None, + quantization_formats: list[dict[str, Any] | str | tuple[dict[str, Any], str]] | None = None, data_loader: Iterable | None = None, forward_step: Callable[[nn.Module, Any], Any | torch.Tensor] | None = None, loss_func: Callable[[Any, Any], torch.Tensor] | None = None, @@ -281,7 +285,7 @@ def auto_quantize( num_calib_steps: int = 512, num_score_steps: int = 128, verbose: bool = False, - method: str = "gradient", + method: str | None = None, checkpoint: str | None = None, module_search_spaces: list[dict[str, Any]] | None = None, fixed_quantization_config: dict[str, Any] | str | None = None, @@ -305,9 +309,11 @@ def auto_quantize( model: A pytorch model with quantizer modules. constraints: Constraints for the search. ``effective_bits`` specifies the effective number of bits for the quantized model and defaults to 4.8. ``cost_model`` selects the metric - used for the effective-bits constraint and currently supports ``"weight"`` (default) - and ``"active_moe"``. Additional cost-model parameters are provided through the nested - ``cost`` dict. + used for the effective-bits constraint and supports ``"weight"`` (default), + ``"active_moe"``, and ``"kv_cache"``. The KV-cache cost model dispatches to isolated + forward-KL scoring over paired K/V formats; BF16/no-quant is its scoring reference but + is never solver-selectable. Additional cost-model parameters are provided through the + nested ``cost`` dict. Here is an example for valid ``effective_bits`` argument: @@ -326,6 +332,9 @@ def auto_quantize( }, } + # For paired K/V-cache formats with exact K/V storage accounting + constraints = {"effective_bits": 5.4, "cost_model": "kv_cache"} + quantization_formats: A list of quantization format config dictionaries or string names to search for. Each config dictionary should be valid as a ``config`` argument in :meth:`quantize `. @@ -526,6 +535,99 @@ def _process_quantization_formats(formats, custom_name_prefix): processed.append((quant_cfg, name)) return processed + is_kv_search = constraints is not None and constraints.get("cost_model") == COST_MODEL_KV_CACHE + if is_kv_search: + if ( + torch.distributed.is_available() + and torch.distributed.is_initialized() + and torch.distributed.get_world_size() > 1 + ): + raise RuntimeError( + "KV-cache AutoQuant is single-process only; distributed scoring, selection, " + "and checkpoint writes are not synchronized." + ) + if method not in (None, "kl_div"): + raise ValueError("cost_model='kv_cache' requires method='kl_div'.") + if fixed_quantization_config is not None or module_search_spaces: + raise ValueError( + "KV-cache AutoQuant does not support fixed_quantization_config or " + "module_search_spaces." + ) + if loss_func is not None or forward_backward_step is not None: + raise ValueError( + "KV-cache AutoQuant uses forward KL and does not accept loss_func or " + "forward_backward_step." + ) + if not isinstance(quantization_formats, list) or not quantization_formats: + raise ValueError( + "cost_model='kv_cache' requires a non-empty quantization_formats list." + ) + if data_loader is None or forward_step is None: + raise ValueError( + "data_loader and forward_step must be provided for KV-cache AutoQuant." + ) + processed_kv_formats = [] + for index, candidate in enumerate(quantization_formats): + if isinstance(candidate, tuple): + raw_config, name = candidate + elif isinstance(candidate, str): + if not hasattr(mtq, candidate): + raise ValueError(f"Unknown KV-cache quantization format: {candidate!r}.") + raw_config, name = getattr(mtq, candidate), candidate + elif isinstance(candidate, dict): + raw_config = candidate + name = QuantRecipe.get_auto_name_for_config(candidate) or f"KV_CACHE_FORMAT_{index}" + else: + raise TypeError( + "KV-cache quantization formats must be config dictionaries, preset names, " + "or (config, name) tuples." + ) + if not isinstance(raw_config, dict): + raise TypeError("KV-cache AutoQuant formats must resolve to config dictionaries.") + if not isinstance(name, str) or not name: + raise ValueError("KV-cache AutoQuant candidate names must be non-empty strings.") + processed_kv_formats.append((raw_config, name)) + + assert constraints is not None + _validate_kv_cache_search_inputs( + constraints, + processed_kv_formats, + num_calib_steps, + num_score_steps, + ) + converted_for_search = not is_quantized(model) + conversion_snapshot = _snapshot_model_structure(model) if converted_for_search else [] + is_training = model.training + searcher = AutoQuantizeKVSearcher() + try: + if converted_for_search: + model = apply_mode(model, mode="auto_quantize", registry=QuantizeModeRegistry) + set_quantizer_by_cfg(model, [{"quantizer_name": "*", "enable": False}]) + search_config = { + "quantization_formats": processed_kv_formats, + "data_loader": data_loader, + "forward_step": forward_step, + "num_calib_steps": num_calib_steps, + "num_score_steps": num_score_steps, + "disabled_layers": disabled_layers, + "verbose": verbose, + "checkpoint": checkpoint, + } + searcher.search( + model, + cast("ConstraintsDict", constraints), + config=search_config, + ) + return model, searcher.state_dict() + except Exception: + searcher.restore_original_quantizers() + model.train(is_training) + if converted_for_search: + _restore_model_structure(conversion_snapshot) + raise + + method = method or "gradient" + if fixed_quantization_config is None and quantization_formats is None: quantization_formats = [mtq.NVFP4_AWQ_LITE_CFG, mtq.FP8_DEFAULT_CFG] elif fixed_quantization_config is not None and quantization_formats is None: @@ -659,6 +761,33 @@ def _process_quantization_formats(formats, custom_name_prefix): return model, searcher.state_dict() +def _snapshot_model_structure( + model: nn.Module, +) -> list[tuple[nn.Module, type[nn.Module], dict[str, Any]]]: + """Capture lightweight module metadata for failure-atomic fresh conversion.""" + return [ + ( + module, + type(module), + { + key: copy.copy(value) if isinstance(value, dict | list | set) else value + for key, value in module.__dict__.items() + }, + ) + for module in model.modules() + ] + + +def _restore_model_structure( + snapshot: list[tuple[nn.Module, type[nn.Module], dict[str, Any]]], +) -> None: + """Undo an in-place quantization conversion without copying model tensors.""" + for module, original_type, original_state in reversed(snapshot): + object.__setattr__(module, "__class__", original_type) + module.__dict__.clear() + module.__dict__.update(original_state) + + def get_auto_quantize_config(search_state, constraints=None, verbose=False): """Build a flat quant config from auto_quantize search_state. @@ -692,6 +821,8 @@ def get_auto_quantize_config(search_state, constraints=None, verbose=False): # fresh_model = load_model(...) # fresh_model = mtq.quantize(fresh_model, config, forward_loop=calibrate_loop) """ + if search_state.get("cost_model") == COST_MODEL_KV_CACHE: + return get_kv_cache_auto_quantize_config(search_state, constraints, verbose=verbose) return _get_auto_quantize_config(search_state, constraints, verbose=verbose) diff --git a/modelopt_recipes/general/auto_quantize/kv_fp8_nvfp4_cast_kl_div_at_5p4bits.yaml b/modelopt_recipes/general/auto_quantize/kv_fp8_nvfp4_cast_kl_div_at_5p4bits.yaml new file mode 100644 index 00000000000..28fb8262c08 --- /dev/null +++ b/modelopt_recipes/general/auto_quantize/kv_fp8_nvfp4_cast_kl_div_at_5p4bits.yaml @@ -0,0 +1,44 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +# Layer-wise KV-cache search over FP8-cast and NVFP4-cast at 5.4 bits/scalar. +# Isolated full-vocabulary forward KL is measured with every other eligible +# attention layer kept in BF16. + +# modelopt-schema: modelopt.recipe.config.ModelOptAutoQuantizeRecipe +imports: + base_disabled_layers: configs/auto_quantize/units/base_disabled_layers + fp8: configs/numerics/fp8 + nvfp4: configs/numerics/nvfp4 + +metadata: + recipe_type: auto_quantize + description: Layer-wise FP8-cast/NVFP4-cast KV-cache search at 5.4 bits using forward KL. + +auto_quantize: + constraints: + effective_bits: 5.4 + cost_model: kv_cache + + candidate_formats: + - quant_cfg: + - quantizer_name: "*[kv]_bmm_quantizer" + cfg: + $import: fp8 + constant_amax: 448.0 + algorithm: + effective_bits: 8.0 + - quant_cfg: + - quantizer_name: "*[kv]_bmm_quantizer" + cfg: + $import: nvfp4 + constant_amax: 448.0 + algorithm: + effective_bits: 4.5 + + auto_quantize_method: kl_div + score_size: 128 + + disabled_layers: + - $import: base_disabled_layers + - "*mtp*" diff --git a/tests/_test_utils/torch/transformers_models.py b/tests/_test_utils/torch/transformers_models.py index cf75e50e107..dc70c180057 100644 --- a/tests/_test_utils/torch/transformers_models.py +++ b/tests/_test_utils/torch/transformers_models.py @@ -220,6 +220,7 @@ def get_tiny_qwen3vl(**config_kwargs) -> PreTrainedModel: "head_dim": 8, "max_position_embeddings": 32, "vocab_size": 32, + "rope_scaling": {"rope_type": "default", "mrope_section": [1, 1, 2]}, } text_kwargs.update(config_kwargs) # Pass as dicts — transformers 5.3.0 Qwen3VLConfig.__init__ only handles diff --git a/tests/examples/hf_ptq/test_hf_ptq_args.py b/tests/examples/hf_ptq/test_hf_ptq_args.py index 6a2c36e4a7c..761d55a386f 100644 --- a/tests/examples/hf_ptq/test_hf_ptq_args.py +++ b/tests/examples/hf_ptq/test_hf_ptq_args.py @@ -13,6 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import copy import getpass import importlib import sys @@ -20,11 +21,14 @@ from types import SimpleNamespace import pytest +import torch import yaml +from _test_utils.torch.transformers_models import get_tiny_qwen3 from modelopt.recipe import load_recipe from modelopt.recipe.config import AutoQuantizeConfig, AutoQuantizeConstraints from modelopt.recipe.presets import QUANT_CFG_CHOICES +from modelopt.torch.quantization import tensor_quant from modelopt.torch.quantization.config import QuantizeConfig _EXAMPLES_DIR = Path(__file__).resolve().parents[3] / "examples" / "hf_ptq" @@ -55,6 +59,19 @@ def _parse_hf_ptq_args(monkeypatch, *args): return hf_ptq, parsed_args +def test_recipe_help_distinguishes_weight_and_kv_autoquant(monkeypatch, capsys): + hf_ptq = _import_hf_ptq(monkeypatch) + monkeypatch.setattr(sys, "argv", ["hf_ptq.py", "--help"]) + + with pytest.raises(SystemExit) as exc_info: + hf_ptq.parse_args() + + assert exc_info.value.code == 0 + help_text = " ".join(capsys.readouterr().out.split()) + assert "weight AutoQuantize recipes use their kv_cache setting" in help_text + assert "KV-cache AutoQuantize recipes select per-layer K/V formats" in help_text + + def test_autoquant_recipe_builds_mtq_inputs(monkeypatch): """The recipe path maps an AutoQuantizeConfig to the expected mtq.auto_quantize inputs.""" hf_ptq, args = _parse_hf_ptq_args( @@ -83,6 +100,156 @@ def test_autoquant_recipe_builds_mtq_inputs(monkeypatch): assert inputs["quantization_formats"][1] == QUANT_CFG_CHOICES["fp8"] +def test_kv_autoquant_recipe_builds_kv_search_inputs(monkeypatch): + hf_ptq, args = _parse_hf_ptq_args( + monkeypatch, "--pyt_ckpt_path", "dummy", "--kv_cache_qformat", "fp8_cast" + ) + aq = load_recipe("general/auto_quantize/kv_fp8_nvfp4_cast_kl_div_at_5p4bits").auto_quantize + inputs = hf_ptq._mtq_inputs_from_auto_quantize_config(aq, args) + + assert inputs["search_domain"] == "kv_cache" + assert inputs["constraints"] == {"effective_bits": 5.4, "cost_model": "kv_cache"} + assert inputs["method"] == "kl_div" + assert [config["effective_bits"] for config, _ in inputs["quantization_formats"]] == [ + 8.0, + 4.5, + ] + assert aq.cost_excluded_layers == [] + assert "*mtp*" in inputs["disabled_layers"] + assert "kv_cache_quant_cfg" not in inputs + + +def test_hf_ptq_kv_autoquant_invokes_public_api(monkeypatch): + """The HF entry point runs the real public KV AutoQuant path on an offline Qwen fixture.""" + hf_ptq = _import_hf_ptq(monkeypatch) + monkeypatch.setattr( + tensor_quant, + "dynamic_block_quantize_op", + lambda inputs, *_args, **_kwargs: torch.zeros_like(inputs), + ) + model = get_tiny_qwen3(num_hidden_layers=1) + aq = load_recipe("general/auto_quantize/kv_fp8_nvfp4_cast_kl_div_at_5p4bits").auto_quantize + args = SimpleNamespace( + calib_with_images=False, + inference_pipeline_parallel=1, + use_fsdp2=False, + kv_cache_qformat="none", + batch_size=1, + auto_quantize_checkpoint=None, + ) + data = [{"input_ids": torch.randint(0, model.config.vocab_size, (1, 8))}] + + hf_ptq.auto_quantize(args, model, data, aq, full_model=model) + + attention = model.model.layers[0].self_attn + assert attention.k_bmm_quantizer.num_bits == (2, 1) + assert attention.v_bmm_quantizer.num_bits == (2, 1) + assert attention.k_bmm_quantizer.amax == 448.0 + assert attention.v_bmm_quantizer.amax == 448.0 + + +def test_kv_autoquant_names_asymmetric_export_format(monkeypatch): + """The supported FP8-K/NVFP4-V candidate has a stable semantic name.""" + hf_ptq = _import_hf_ptq(monkeypatch) + mixed_config = copy.deepcopy(hf_ptq.KV_QUANT_CFG_CHOICES["nvfp4"]) + fp8_k_quantizer = copy.deepcopy(hf_ptq.KV_QUANT_CFG_CHOICES["fp8"]["quant_cfg"][0]) + fp8_k_quantizer["quantizer_name"] = "*.k_bmm_quantizer" + mixed_config["quant_cfg"].append(fp8_k_quantizer) + mixed_config["effective_bits"] = 6.25 + + candidates = hf_ptq._mtq_kv_candidate_formats([QuantizeConfig(**mixed_config)]) + + assert candidates[0][1] == "fp8_k_nvfp4_v" + + +def test_kv_autoquant_kl_excludes_padding_positions(monkeypatch): + hf_ptq = _import_hf_ptq(monkeypatch) + logits = torch.arange(2 * 4 * 3).reshape(2, 4, 3) + attention_mask = torch.tensor([[1, 1, 0, 0], [0, 1, 1, 0]]) + + selected = hf_ptq._select_unpadded_logits(logits, {"attention_mask": attention_mask}) + + assert torch.equal(selected, logits[attention_mask.bool()]) + + +def test_kv_autoquant_kl_rejects_misaligned_attention_mask(monkeypatch): + hf_ptq = _import_hf_ptq(monkeypatch) + + with pytest.raises(ValueError, match="matching token dimensions"): + hf_ptq._select_unpadded_logits(torch.zeros(2, 4, 3), {"attention_mask": torch.ones(2, 3)}) + + +def test_kv_autoquant_rejects_fsdp2(monkeypatch): + hf_ptq = _import_hf_ptq(monkeypatch) + monkeypatch.setattr( + hf_ptq, + "_mtq_inputs_from_auto_quantize_config", + lambda *_args, **_kwargs: {"search_domain": "kv_cache"}, + ) + args = SimpleNamespace( + calib_with_images=False, + inference_pipeline_parallel=1, + use_fsdp2=True, + ) + + with pytest.raises(NotImplementedError, match="KV-cache AutoQuantize does not support"): + hf_ptq.auto_quantize(args, torch.nn.Module(), [], SimpleNamespace()) + + +def test_weight_autoquant_retains_fsdp2_warning(monkeypatch): + hf_ptq = _import_hf_ptq(monkeypatch) + model = torch.nn.Module() + inputs = { + "search_domain": "weight", + "constraints": {"effective_bits": 8.0}, + "quantization_formats": [], + "fixed_quantization_config": None, + "module_search_spaces": [], + "disabled_layers": [], + "kv_cache_quant_cfg": None, + "method": "gradient", + "score_size": 1, + } + monkeypatch.setattr( + hf_ptq, "_mtq_inputs_from_auto_quantize_config", lambda *_args, **_kwargs: inputs + ) + monkeypatch.setattr( + hf_ptq.mtq, "auto_quantize", lambda search_model, **_kwargs: (search_model, {}) + ) + args = SimpleNamespace( + calib_with_images=False, + inference_pipeline_parallel=1, + use_fsdp2=True, + batch_size=1, + auto_quantize_checkpoint=None, + ) + + with pytest.warns(UserWarning, match="use at your own risk"): + assert hf_ptq.auto_quantize(args, model, [], SimpleNamespace()) is model + + +def test_fsdp2_kv_autoquant_rejected_before_model_load(monkeypatch): + hf_ptq = _import_hf_ptq(monkeypatch) + monkeypatch.setattr(hf_ptq, "_recipe_is_kv_auto_quantize", lambda _: True) + monkeypatch.setattr( + hf_ptq.AutoConfig, + "from_pretrained", + lambda *_args, **_kwargs: pytest.fail("The model config must not be loaded."), + ) + + with pytest.raises(NotImplementedError, match="KV-cache AutoQuantize does not support"): + hf_ptq.load_model(SimpleNamespace(use_fsdp2=True, recipe="autoquant")) + + +def test_fsdp2_preload_guard_distinguishes_weight_and_kv_autoquant(monkeypatch): + hf_ptq = _import_hf_ptq(monkeypatch) + + assert hf_ptq._recipe_is_kv_auto_quantize( + "general/auto_quantize/kv_fp8_nvfp4_cast_kl_div_at_5p4bits" + ) + assert not hf_ptq._recipe_is_kv_auto_quantize("general/auto_quantize/nvfp4_fp8_at_5p4bits") + + def test_autoquant_recipe_cost_excluded_layers_map_into_cost(monkeypatch): """Top-level cost_excluded_layers maps to the mtq constraints.cost.excluded_module_name_patterns key (distinct from disabled_layers), so a cost-exclusion recipe matches the nested mtq dict.""" diff --git a/tests/unit/recipe/test_loader.py b/tests/unit/recipe/test_loader.py index 3dfb9906a54..cacc6bcd520 100644 --- a/tests/unit/recipe/test_loader.py +++ b/tests/unit/recipe/test_loader.py @@ -28,6 +28,9 @@ import modelopt.torch.quantization.config as qcfg from modelopt.recipe.config import ( + AutoQuantizeConfig, + AutoQuantizeConstraints, + AutoQuantizeCost, ModelOptAutoQuantizeRecipe, ModelOptDFlashRecipe, ModelOptEagleRecipe, @@ -1768,6 +1771,22 @@ def test_load_recipe_autoquantize_minimal(tmp_path): assert aq.module_search_spaces == [] +def test_autoquantize_constraints_use_effective_bits_for_kv_cost_model(): + assert AutoQuantizeConstraints.model_fields["effective_bits"].default == 4.8 + assert AutoQuantizeConstraints().effective_bits == 4.8 + + constraints = AutoQuantizeConstraints(effective_bits=5.4, cost_model="kv_cache") + assert constraints.effective_bits == 5.4 + assert constraints.cost_model == "kv_cache" + + with pytest.raises(ValueError, match="does not accept weight cost settings"): + AutoQuantizeConstraints( + effective_bits=5.4, + cost_model="kv_cache", + cost=AutoQuantizeCost(active_moe_expert_ratio=0.5), + ) + + def test_load_recipe_autoquantize_active_moe_cost_roundtrip(tmp_path): """cost_model + cost.active_moe_expert_ratio parse and dump to the mtq constraints dict shape.""" recipe_file = tmp_path / "aq.yml" @@ -1907,6 +1926,7 @@ def test_load_recipe_autoquantize_fixed_baseline_requires_explicit_search(tmp_pa [ "general/auto_quantize/nvfp4_fp8_at_5p4bits", "general/auto_quantize/nvfp4_fp8_kl_div_at_5p4bits", + "general/auto_quantize/kv_fp8_nvfp4_cast_kl_div_at_5p4bits", "general/auto_quantize/nvfp4_mse_fp8_at_6p0bits", "general/auto_quantize/w4a8_awq_beta_fp8_at_6p0bits", "general/auto_quantize/w4a16_nvfp4_fp8_at_6p0bits-active_moe", @@ -1918,11 +1938,45 @@ def test_load_recipe_autoquantize_builtin_general(recipe_path): assert isinstance(recipe, ModelOptAutoQuantizeRecipe) assert len(recipe.auto_quantize.candidate_formats) >= 2 assert recipe.auto_quantize.auto_quantize_method in ("gradient", "kl_div") - # Both shared base units must be spliced in: the removed --auto_quantize_* CLI shim appended - # them unconditionally, so a general recipe is the migration target and must match it. Without - # cost_excluded_layers a VL/MTP model counts its vision tower in the effective-bits denominator. assert "*output_layer*" in recipe.auto_quantize.disabled_layers - assert recipe.auto_quantize.cost_excluded_layers == ["*visual*", "*mtp*", "*vision_tower*"] + if recipe.auto_quantize.constraints.cost_model == "kv_cache": + assert "*mtp*" in recipe.auto_quantize.disabled_layers + assert recipe.auto_quantize.cost_excluded_layers == [] + else: + assert recipe.auto_quantize.cost_excluded_layers == [ + "*visual*", + "*mtp*", + "*vision_tower*", + ] + + +def test_load_recipe_kv_autoquantize_contract(): + recipe = load_recipe("general/auto_quantize/kv_fp8_nvfp4_cast_kl_div_at_5p4bits") + aq = recipe.auto_quantize + + assert aq.constraints.effective_bits == 5.4 + assert aq.constraints.cost_model == "kv_cache" + assert aq.auto_quantize_method == "kl_div" + assert "*mtp*" in aq.disabled_layers + assert aq.cost_excluded_layers == [] + assert all(candidate.algorithm is None for candidate in aq.candidate_formats) + assert [fmt.effective_bits for fmt in aq.candidate_formats] == [8.0, 4.5] + for fmt in aq.candidate_formats: + for entry in fmt.quant_cfg: + assert entry.quantizer_name == "*[kv]_bmm_quantizer" + assert not entry.cfg.use_constant_amax + assert entry.cfg.constant_amax == 448.0 + assert fmt.algorithm is None + + +def test_kv_autoquantize_rejects_cost_excluded_layers(): + with pytest.raises(ValueError, match=r"cost_excluded_layers.*disabled_layers"): + AutoQuantizeConfig( + constraints=AutoQuantizeConstraints(effective_bits=8.0, cost_model="kv_cache"), + candidate_formats=[qcfg.QuantizeConfig(quant_cfg=[], effective_bits=8.0)], + auto_quantize_method="kl_div", + cost_excluded_layers=["*mtp*"], + ) def _all_shipped_ptq_recipe_paths(): diff --git a/tests/unit/torch/export/test_convert_hf_config.py b/tests/unit/torch/export/test_convert_hf_config.py new file mode 100644 index 00000000000..a0b8295b485 --- /dev/null +++ b/tests/unit/torch/export/test_convert_hf_config.py @@ -0,0 +1,83 @@ +# 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. + +import json + +import torch + +from modelopt.torch.export.convert_hf_config import convert_hf_quant_config_format +from modelopt.torch.export.unified_export_hf import _write_hf_export_config + + +def test_convert_mixed_kv_cache_config_preserves_layer_map(): + layer_map = { + "model.layers.0.self_attn": {"quant_algo": "FP8"}, + "model.layers.1.self_attn": {"quant_algo": "NVFP4"}, + } + converted = convert_hf_quant_config_format( + { + "producer": {"name": "modelopt", "version": "test"}, + "quantization": { + "quant_algo": "MIXED_PRECISION", + "quantized_layers": {}, + "kv_cache_quant_algo": "MIXED_PRECISION", + "kv_cache_quantized_layers": layer_map, + "kv_cache_schema_version": 1, + }, + } + ) + + assert converted["quant_method"] == "modelopt" + assert converted["quant_algo"] == "MIXED_PRECISION" + assert converted["config_groups"] == {} + assert converted["kv_cache_quant_algo"] == "MIXED_PRECISION" + assert converted["kv_cache_quantized_layers"] == layer_map + assert converted["kv_cache_schema_version"] == 1 + + +def test_write_hf_export_config_writes_mapped_kv_autoquant_report(tmp_path): + layer_name = "model.layers.0.self_attn" + model = torch.nn.Module() + model._modelopt_kv_cache_auto_quantize_state = { + "layers": {layer_name: {"selected": "fp8"}}, + "search_signature": {"layers": [{"name": layer_name}]}, + } + quant_config = { + "producer": {"name": "modelopt", "version": "test"}, + "quantization": { + "quant_algo": None, + "kv_cache_quant_algo": "MIXED_PRECISION", + "kv_cache_quantized_layers": {layer_name: {"quant_algo": "FP8"}}, + "kv_cache_schema_version": 1, + }, + } + (tmp_path / "config.json").write_text("{}") + + _write_hf_export_config( + model, + quant_config, + tmp_path, + name_mapper=lambda name: f"hub.{name}", + ) + + report = json.loads((tmp_path / "kv_cache_auto_quantize_report.json").read_text()) + assert report["layers"] == {f"hub.{layer_name}": {"selected": "fp8"}} + assert report["search_signature"]["layers"] == [{"name": f"hub.{layer_name}"}] + assert model._modelopt_kv_cache_auto_quantize_state["layers"] == { + layer_name: {"selected": "fp8"} + } + assert (tmp_path / "hf_quant_config.json").is_file() + exported_config = json.loads((tmp_path / "config.json").read_text()) + assert exported_config["quantization_config"]["kv_cache_quant_algo"] == "MIXED_PRECISION" diff --git a/tests/unit/torch/export/test_get_quantization.py b/tests/unit/torch/export/test_get_quantization.py index e7ca68d0b69..1202b1fcae4 100644 --- a/tests/unit/torch/export/test_get_quantization.py +++ b/tests/unit/torch/export/test_get_quantization.py @@ -27,12 +27,15 @@ import modelopt.torch.quantization as mtq from modelopt.torch.export.layer_utils import get_quantization_format from modelopt.torch.export.model_config import ( + KV_CACHE_FP8, + KV_CACHE_FP8_K_NVFP4_V, + KV_CACHE_NVFP4, QUANTIZATION_FP8, QUANTIZATION_NVFP4, QUANTIZATION_W4A8_AWQ, ) -from modelopt.torch.export.quant_utils import get_quant_config -from modelopt.torch.quantization.nn import NVFP4StaticQuantizer +from modelopt.torch.export.quant_utils import get_quant_config, postprocess_state_dict +from modelopt.torch.quantization.nn import NVFP4StaticQuantizer, TensorQuantizer @pytest.mark.parametrize( @@ -64,6 +67,161 @@ def test_nvfp4_static_quantizer_export(): assert quant_config["quantization"]["group_size"] == 16 +def test_projection_output_quantizers_are_not_exported_as_kv_cache(): + model = ToyModel() + config = { + "quant_cfg": [ + {"quantizer_name": "*", "enable": False}, + { + "quantizer_name": "*.weight_quantizer", + "cfg": {"num_bits": (4, 3), "axis": None}, + "enable": True, + }, + { + "quantizer_name": "*.input_quantizer", + "cfg": {"num_bits": (4, 3), "axis": None}, + "enable": True, + }, + { + "quantizer_name": "*.output_quantizer", + "cfg": {"num_bits": (4, 3), "axis": None}, + "enable": True, + }, + ], + "algorithm": "max", + } + mtq.quantize(model, config, lambda x: x(torch.randn(1, 4, 10))) + + quantization = get_quant_config(model)["quantization"] + + assert quantization["quant_algo"] == "FP8" + assert quantization["kv_cache_quant_algo"] is None + assert "kv_cache_quantized_layers" not in quantization + + +def test_mixed_kv_cache_quantization_exports_per_layer_map(): + class FakeAttention(torch.nn.Module): + def __init__(self): + super().__init__() + self.k_bmm_quantizer = TensorQuantizer() + self.v_bmm_quantizer = TensorQuantizer() + + model = torch.nn.Module() + model.attn0 = FakeAttention() + model.attn1 = FakeAttention() + model.attn2 = FakeAttention() + mtq.set_quantizer_by_cfg( + model.attn0, + [ + { + "quantizer_name": "*[kv]_bmm_quantizer", + "cfg": {"num_bits": (4, 3), "use_constant_amax": True}, + } + ], + ) + mtq.set_quantizer_by_cfg( + model.attn1, + [ + { + "quantizer_name": "*[kv]_bmm_quantizer", + "cfg": { + "num_bits": (2, 1), + "block_sizes": {-1: 16, "type": "dynamic", "scale_bits": (4, 3)}, + "use_constant_amax": True, + }, + } + ], + ) + mtq.set_quantizer_by_cfg( + model.attn2, + [ + { + "quantizer_name": "*k_bmm_quantizer", + "cfg": {"num_bits": (4, 3), "use_constant_amax": True}, + }, + { + "quantizer_name": "*v_bmm_quantizer", + "cfg": { + "num_bits": (2, 1), + "block_sizes": {-1: 16, "type": "dynamic", "scale_bits": (4, 3)}, + "use_constant_amax": True, + }, + }, + ], + ) + + quantization = get_quant_config(model)["quantization"] + assert quantization["quant_algo"] == "MIXED_PRECISION" + assert quantization["kv_cache_quant_algo"] == "MIXED_PRECISION" + assert quantization["quantized_layers"] == {} + assert quantization["kv_cache_quantized_layers"] == { + "attn0": {"quant_algo": "FP8"}, + "attn1": {"quant_algo": "NVFP4"}, + "attn2": {"quant_algo": "FP8_K_NVFP4_V"}, + } + + +def test_unsupported_asymmetric_kv_cache_pair_fails_export(): + class FakeAttention(torch.nn.Module): + def __init__(self): + super().__init__() + self.k_bmm_quantizer = TensorQuantizer() + self.v_bmm_quantizer = TensorQuantizer() + + model = FakeAttention() + mtq.set_quantizer_by_cfg( + model, + [ + { + "quantizer_name": "*k_bmm_quantizer", + "cfg": { + "num_bits": (2, 1), + "block_sizes": {-1: 16, "type": "dynamic", "scale_bits": (4, 3)}, + "use_constant_amax": True, + }, + }, + { + "quantizer_name": "*v_bmm_quantizer", + "cfg": {"num_bits": (4, 3), "use_constant_amax": True}, + }, + ], + ) + + with pytest.raises(NotImplementedError, match="Unsupported mixed K/V cache"): + get_quant_config(model) + + +def test_mixed_kv_cache_postprocess_uses_each_layers_format(): + state_dict = { + "attn0.k_bmm_quantizer._amax": torch.tensor([448.0]), + "attn0.v_bmm_quantizer._amax": torch.tensor([224.0]), + "attn1.k_bmm_quantizer._amax": torch.tensor([112.0]), + "attn1.v_bmm_quantizer._amax": torch.tensor([56.0]), + } + layer_formats = { + "attn0": {"quant_algo": KV_CACHE_FP8}, + "attn1": {"quant_algo": KV_CACHE_NVFP4}, + "attn2": {"quant_algo": KV_CACHE_FP8_K_NVFP4_V}, + } + state_dict.update( + { + "attn2.k_bmm_quantizer._amax": torch.tensor([448.0]), + "attn2.v_bmm_quantizer._amax": torch.tensor([112.0]), + } + ) + + processed = postprocess_state_dict(state_dict, 448.0, layer_formats) + + assert processed == { + "attn0.k_proj.k_scale": torch.tensor([1.0]), + "attn0.v_proj.v_scale": torch.tensor([0.5]), + "attn1.k_proj.k_scale": torch.tensor([0.25]), + "attn1.v_proj.v_scale": torch.tensor([0.125]), + "attn2.k_proj.k_scale": torch.tensor([1.0]), + "attn2.v_proj.v_scale": torch.tensor([0.25]), + } + + class _FakeTopKRouter(torch.nn.Module): """Mimics a transformers>=5.0 MoE router: owns a ``weight`` but is NOT an ``nn.Linear``. diff --git a/tests/unit/torch/export/test_offload_export.py b/tests/unit/torch/export/test_offload_export.py index 242a64f762b..09551c6bf60 100644 --- a/tests/unit/torch/export/test_offload_export.py +++ b/tests/unit/torch/export/test_offload_export.py @@ -36,9 +36,13 @@ ) import modelopt.torch.quantization as mtq -from modelopt.torch.export.model_config import KV_CACHE_FP8 +from modelopt.torch.export.model_config import KV_CACHE_FP8, KV_CACHE_FP8_K_NVFP4_V, KV_CACHE_NVFP4 from modelopt.torch.export.model_utils import TiedWeightMap -from modelopt.torch.export.quant_utils import _postprocess_single_tensor +from modelopt.torch.export.quant_utils import ( + _get_kv_cache_postprocess_config, + _postprocess_single_tensor, + _resolve_kv_cache_format_for_key, +) from modelopt.torch.export.unified_export_hf import _export_quantized_weight from modelopt.torch.export.unified_export_hf_streaming import ( _parse_shard_size, @@ -319,6 +323,53 @@ def test_postprocess_kv_scale_renamed_and_divided(): assert abs(val.item() - 0.5) < 1e-5 +@pytest.mark.parametrize( + ("layer_name", "quant_algo", "side", "resolved_format"), + [ + ("model.layers.0.self_attn", KV_CACHE_FP8_K_NVFP4_V, "k", KV_CACHE_FP8), + ("model.layers.0.self_attn", KV_CACHE_FP8_K_NVFP4_V, "v", KV_CACHE_NVFP4), + ("model.layers.1.self_attn", KV_CACHE_NVFP4, "k", KV_CACHE_NVFP4), + ], +) +def test_postprocess_resolves_mixed_kv_format_per_layer_and_side( + layer_name, quant_algo, side, resolved_format +): + quantization = { + "kv_cache_quant_algo": "MIXED_PRECISION", + "kv_cache_quantized_layers": {layer_name: {"quant_algo": quant_algo}}, + } + postprocess_config = _get_kv_cache_postprocess_config(quantization) + original_key = f"{layer_name}.{side}_bmm_quantizer._amax" + + assert _resolve_kv_cache_format_for_key(original_key, postprocess_config) == resolved_format + + key, val = _postprocess_single_tensor( + original_key, + torch.tensor(224.0), + 448.0, + postprocess_config, + ) + + assert key == f"{layer_name}.{side}_proj.{side}_scale" + assert val.item() == pytest.approx(0.5) + + +@pytest.mark.parametrize(("side", "resolved_format"), [("k", KV_CACHE_FP8), ("v", KV_CACHE_NVFP4)]) +def test_postprocess_resolves_uniform_asymmetric_kv_format(side, resolved_format): + original_key = f"model.layers.0.self_attn.{side}_bmm_quantizer._amax" + + assert _resolve_kv_cache_format_for_key(original_key, KV_CACHE_FP8_K_NVFP4_V) == resolved_format + key, val = _postprocess_single_tensor( + original_key, + torch.tensor(224.0), + 448.0, + KV_CACHE_FP8_K_NVFP4_V, + ) + + assert key == f"model.layers.0.self_attn.{side}_proj.{side}_scale" + assert val.item() == pytest.approx(0.5) + + def test_postprocess_scale_squeezed(): """3D scale tensors with shape[0]==1 are squeezed.""" t = torch.ones(1, 4, 4) diff --git a/tests/unit/torch/export/test_quant_aware_conversion.py b/tests/unit/torch/export/test_quant_aware_conversion.py index e3aa22bf3d8..06a1b88f222 100644 --- a/tests/unit/torch/export/test_quant_aware_conversion.py +++ b/tests/unit/torch/export/test_quant_aware_conversion.py @@ -516,6 +516,7 @@ def test_revert_quant_config_names_mapper(): "lm_head", ], "quantized_layers": {"model.layers.0.mlp.experts.0.w1": {"quant_algo": "NVFP4"}}, + "kv_cache_quantized_layers": {"model.layers.0.mlp.experts.0": {"quant_algo": "FP8"}}, } revert_quant_config_names(quant, mapper) assert quant["exclude_modules"] == [ @@ -524,6 +525,7 @@ def test_revert_quant_config_names_mapper(): "lm_head", ] assert "model.layers.0.block_sparse_moe.experts.0.w1" in quant["quantized_layers"] + assert "model.layers.0.block_sparse_moe.experts.0" in quant["kv_cache_quantized_layers"] # mapper(None) is a no-op q2 = {"exclude_modules": ["x*"]} revert_quant_config_names(q2, None) diff --git a/tests/unit/torch/export/test_unified_export_hf.py b/tests/unit/torch/export/test_unified_export_hf.py index 469ffd9227e..9de51052179 100644 --- a/tests/unit/torch/export/test_unified_export_hf.py +++ b/tests/unit/torch/export/test_unified_export_hf.py @@ -25,7 +25,11 @@ ) import modelopt.torch.quantization as mtq -from modelopt.torch.export.model_utils import TiedWeightMap +from modelopt.torch.export.model_utils import ( + TiedWeightMap, + get_language_model_from_vl, + is_multimodal_model, +) from modelopt.torch.export.quant_utils import ( fuse_prequant_layernorm, postprocess_state_dict, @@ -35,6 +39,25 @@ from modelopt.torch.quantization.nn import TensorQuantizer +def test_multimodal_detection_accepts_null_architectures(): + """Unified export treats absent architecture metadata as an empty list.""" + model = SimpleNamespace(config=SimpleNamespace(architectures=None)) + + assert not is_multimodal_model(model) + + +@pytest.mark.parametrize("aliased", [False, True]) +def test_language_model_extraction_rejects_competing_or_aliased_roots(aliased): + """Language-model extraction must not select ambiguous roots by traversal order.""" + model = torch.nn.Module() + model.model = torch.nn.Module() + model.model.language_model = torch.nn.Module() + model.language_model = model.model.language_model if aliased else torch.nn.Module() + + with pytest.raises(ValueError, match="multiple language-model roots"): + get_language_model_from_vl(model) + + @pytest.mark.parametrize( ("configured_dtype", "dtype", "expected_dtype", "warning_count"), [ diff --git a/tests/unit/torch/quantization/test_kv_cache_auto_quant.py b/tests/unit/torch/quantization/test_kv_cache_auto_quant.py new file mode 100644 index 00000000000..7df6c2764bd --- /dev/null +++ b/tests/unit/torch/quantization/test_kv_cache_auto_quant.py @@ -0,0 +1,811 @@ +# 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. + +import json + +import pytest +import torch +import torch.nn as nn +from _test_utils.torch.transformers_models import get_tiny_llama, get_tiny_qwen3, get_tiny_qwen3vl + +import modelopt.torch.quantization as mtq +from modelopt.torch.export.quant_utils import get_kv_cache_dtype, get_quant_config +from modelopt.torch.opt.searcher import BaseSearcher +from modelopt.torch.quantization import model_quant, tensor_quant +from modelopt.torch.quantization.config import QuantizeConfig +from modelopt.torch.quantization.kv_cache_auto_quant import ( + AutoQuantizeKVSearcher, + _candidate_quantizers, + _eligible_layers, + _kv_scalar_weight, + _solve_additive_recipe, + _validate_kv_only_config, +) +from modelopt.torch.quantization.nn import TensorQuantizer + + +@pytest.fixture +def nvfp4_fake_quant_stub(monkeypatch): + """Keep CPU search tests independent of the CUDA-only NVFP4 fake-quant kernel.""" + + monkeypatch.setattr( + tensor_quant, + "dynamic_block_quantize_op", + lambda inputs, *_args, **_kwargs: torch.zeros_like(inputs), + ) + + +def _quantizer_cfg(bits, *, constant_amax=None): + cfg = {"num_bits": bits} + if bits == (2, 1): + cfg["block_sizes"] = {-1: 16, "type": "dynamic", "scale_bits": (4, 3)} + if constant_amax is not None: + cfg["constant_amax"] = constant_amax + return cfg + + +def _kv_config(bits, effective_bits, *, algorithm="max", constant_amax=None): + return QuantizeConfig( + quant_cfg=[ + { + "quantizer_name": "*[kv]_bmm_quantizer", + "cfg": _quantizer_cfg(bits, constant_amax=constant_amax), + } + ], + algorithm=algorithm, + effective_bits=effective_bits, + ) + + +def _asymmetric_kv_config(): + return QuantizeConfig( + quant_cfg=[ + { + "quantizer_name": "*[kv]_bmm_quantizer", + "cfg": _quantizer_cfg((2, 1), constant_amax=1.0), + }, + { + "quantizer_name": "*.k_bmm_quantizer", + "cfg": _quantizer_cfg((4, 3), constant_amax=1.0), + }, + ], + algorithm=None, + effective_bits=6.25, + ) + + +def test_kv_candidate_requires_exact_bits_and_both_sides(): + _validate_kv_only_config(_kv_config((4, 3), 8.0)) + + with pytest.raises(ValueError, match="config-level effective_bits"): + _validate_kv_only_config( + QuantizeConfig( + quant_cfg=[ + { + "quantizer_name": "*[kv]_bmm_quantizer", + "cfg": {"num_bits": (4, 3), "use_constant_amax": True}, + } + ] + ) + ) + with pytest.raises(ValueError, match="completely configure both"): + _validate_kv_only_config( + QuantizeConfig( + quant_cfg=[ + { + "quantizer_name": "*k_bmm_quantizer", + "cfg": {"num_bits": (4, 3), "use_constant_amax": True}, + } + ], + effective_bits=8.0, + ) + ) + + +def test_kv_autoquant_uses_shared_searcher_lifecycle(): + assert issubclass(AutoQuantizeKVSearcher, BaseSearcher) + + +@pytest.mark.parametrize("algorithm", ["svdquant", {"method": "smoothquant"}, {"method": "mse"}]) +def test_kv_candidate_rejects_structural_or_unscoped_algorithms(algorithm): + config = _kv_config((4, 3), 8.0).model_copy(update={"algorithm": algorithm}) + + with pytest.raises(ValueError, match="only non-structural calibration algorithms"): + _validate_kv_only_config(config) + + +def test_kv_additive_solver_spends_fp8_on_more_sensitive_layer(): + selections, status = _solve_additive_recipe( + layer_names=["layer0", "layer1"], + layer_widths=[(128, 128), (128, 128)], + candidate_names=["fp8", "nvfp4"], + candidate_kv_bits=[(8.0, 8.0), (4.5, 4.5)], + scores=[[0.0, 10.0], [0.0, 1.0]], + target_bits=6.25, + verbose=False, + ) + + assert status == "Optimal" + assert selections == [0, 1] + + +def test_kv_scalar_weight_counts_k_and_v_widths(): + module = nn.Module() + module.k_proj = nn.Linear(32, 24, bias=False) + module.v_proj = nn.Linear(32, 16, bias=False) + + assert _kv_scalar_weight(module, "attention") == 40 + + +@pytest.mark.parametrize( + ("config", "expected_format"), + [ + (_kv_config((4, 3), 8.0), "FP8"), + (_kv_config((2, 1), 4.5), "NVFP4"), + (_kv_config((4, 3), 8.0, algorithm=None, constant_amax=1.0), "FP8"), + (_kv_config((2, 1), 4.5, algorithm=None, constant_amax=1.0), "NVFP4"), + ], +) +def test_kv_candidate_accepts_export_supported_persistent_formats(config, expected_format): + _validate_kv_only_config(config) + quantizers = _candidate_quantizers(config) + module = nn.Module() + for name, quantizer in quantizers.items(): + setattr(module, name, quantizer) + + assert get_kv_cache_dtype(module) == expected_format + + +@pytest.mark.parametrize( + ("config", "match"), + [ + (_kv_config((4, 3), 8.0, algorithm=None), "no persistent export scale"), + ( + QuantizeConfig( + quant_cfg=[ + { + "quantizer_name": "*[kv]_bmm_quantizer", + "cfg": {"num_bits": (4, 3), "use_constant_amax": True}, + } + ], + algorithm="max", + effective_bits=8.0, + ), + "no persistent export scale", + ), + (_kv_config(8, 8.0, algorithm=None, constant_amax=1.0), "per-tensor FP8"), + (_kv_config(4, 4.0, algorithm=None, constant_amax=1.0), "per-tensor FP8"), + (_kv_config((4, 3), 6.0), "does not match its configured K/V storage cost"), + ], +) +def test_kv_candidate_rejects_non_exportable_or_incorrect_cost(config, match): + with pytest.raises(ValueError, match=match): + _validate_kv_only_config(config) + + +def test_kv_candidate_rejects_top_level_dynamic_fp8(): + config = QuantizeConfig( + quant_cfg=[ + { + "quantizer_name": "*[kv]_bmm_quantizer", + "cfg": {"num_bits": (4, 3), "type": "dynamic"}, + } + ], + algorithm="max", + effective_bits=8.0, + ) + + with pytest.raises(ValueError, match="top-level dynamic"): + _validate_kv_only_config(config) + + +class _ToyKVAttention(nn.Module): + def __init__(self, width, gain): + super().__init__() + self.k_proj = nn.Linear(width, width, bias=False) + self.v_proj = nn.Linear(width, width, bias=False) + self.k_bmm_quantizer = nn.Identity() + self.v_bmm_quantizer = nn.Identity() + self.gain = gain + + def forward(self, x): + return x + self.gain * (self.k_bmm_quantizer(x) + self.v_bmm_quantizer(x)) + + +class _ToyKVModel(nn.Module): + def __init__(self, width=8): + super().__init__() + self.attn0 = _ToyKVAttention(width, gain=0.25) + self.attn1 = _ToyKVAttention(width, gain=2.0) + self.lm_head = nn.Linear(width, width, bias=False) + + def forward(self, x): + return self.lm_head(self.attn1(self.attn0(x))) + + +def test_kv_autoquant_rejects_missing_scale_after_calibration(monkeypatch): + model = _ToyKVModel() + original_quantizers = { + name: (module.k_bmm_quantizer, module.v_bmm_quantizer) + for name, module in (("attn0", model.attn0), ("attn1", model.attn1)) + } + monkeypatch.setattr(model_quant, "calibrate", lambda *_args, **_kwargs: None) + + with pytest.raises(ValueError, match="no persistent export scale after calibration"): + mtq.auto_quantize( + model, + {"effective_bits": 8.0, "cost_model": "kv_cache"}, + [(_kv_config((4, 3), 8.0).model_dump(), "fp8")], + [torch.randn(1, 2, 8)], + lambda search_model, batch: search_model(batch), + num_calib_steps=1, + num_score_steps=1, + ) + + assert model.training + for name, module in (("attn0", model.attn0), ("attn1", model.attn1)): + assert (module.k_bmm_quantizer, module.v_bmm_quantizer) == original_quantizers[name] + + +def test_kv_eligible_layers_supports_hybrid_attention_mixers_only(): + """Hybrid decoders include attention mixers but exclude nonattention mixers.""" + model = nn.Module() + model.layers = nn.ModuleList([nn.Module(), nn.Module()]) + model.layers[0].mixer = nn.Linear(8, 8, bias=False) + model.layers[1].mixer = _ToyKVAttention(8, gain=1.0) + + layers = _eligible_layers(model, disabled_layers=None) + + assert [(name, width) for name, _, width in layers] == [("layers.1.mixer", 16)] + + +def test_kv_eligible_layers_rejects_aliased_attention_boundary(): + """An attention object registered at multiple paths must not be selected by traversal order.""" + model = nn.Module() + attention = _ToyKVAttention(8, gain=1.0) + model.attention = attention + model.attention_alias = attention + + with pytest.raises(ValueError, match="registered through aliases"): + _eligible_layers(model, disabled_layers=None) + + +def test_kv_autoquant_scores_and_applies_one_format_per_layer(tmp_path, nvfp4_fake_quant_stub): + torch.manual_seed(123) + model = _ToyKVModel() + data = [torch.randn(2, 3, 8), torch.randn(2, 3, 8)] + candidates = [ + ( + { + "quant_cfg": [ + { + "quantizer_name": "*[kv]_bmm_quantizer", + "cfg": {"num_bits": (4, 3), "constant_amax": 1.0}, + } + ], + "algorithm": None, + "effective_bits": 8.0, + }, + "fp8", + ), + ( + { + "quant_cfg": [ + { + "quantizer_name": "*[kv]_bmm_quantizer", + "cfg": { + "num_bits": (2, 1), + "block_sizes": { + -1: 16, + "type": "dynamic", + "scale_bits": (4, 3), + }, + "constant_amax": 1.0, + }, + } + ], + "algorithm": None, + "effective_bits": 4.5, + }, + "nvfp4", + ), + ] + + model, state = mtq.auto_quantize( + model, + {"effective_bits": 6.25, "cost_model": "kv_cache"}, + candidates, + data, + lambda model, batch: model(batch), + num_calib_steps=2, + num_score_steps=2, + checkpoint=str(tmp_path / "kv_search.pth"), + ) + + assert state["best"]["constraints"]["effective_bits"] == pytest.approx(6.25) + assert state["best"]["is_satisfied"] + assert model.training + assert {layer["selected"] for layer in state["layers"].values()} == { + "fp8", + "nvfp4", + } + for layer_name, layer_state in state["layers"].items(): + layer = model.get_submodule(layer_name) + assert layer.k_bmm_quantizer.num_bits == layer.v_bmm_quantizer.num_bits + expected_bits = (4, 3) if layer_state["selected"] == "fp8" else (2, 1) + assert layer.k_bmm_quantizer.num_bits == expected_bits + + restored_model = _ToyKVModel().eval() + restored_model, restored_state = mtq.auto_quantize( + restored_model, + {"effective_bits": 6.25, "cost_model": "kv_cache"}, + candidates, + data, + lambda *_: pytest.fail("A compatible checkpoint must skip calibration and scoring."), + num_calib_steps=2, + num_score_steps=2, + checkpoint=str(tmp_path / "kv_search.pth"), + ) + + assert restored_state["best"] == state["best"] + assert restored_state["layers"] == state["layers"] + assert not restored_model.training + for layer_name, layer_state in restored_state["layers"].items(): + layer = restored_model.get_submodule(layer_name) + expected_bits = (4, 3) if layer_state["selected"] == "fp8" else (2, 1) + assert layer.k_bmm_quantizer.num_bits == expected_bits + + resolved_config = mtq.get_auto_quantize_config(state, {"effective_bits": 4.5}) + assert resolved_config["algorithm"] == "max" + assert {entry["quantizer_name"] for entry in resolved_config["quant_cfg"]} == { + "attn0.k_bmm_quantizer", + "attn0.v_bmm_quantizer", + "attn1.k_bmm_quantizer", + "attn1.v_bmm_quantizer", + } + + re_solved_model = _ToyKVModel() + _, re_solved_state = mtq.auto_quantize( + re_solved_model, + {"effective_bits": 4.5, "cost_model": "kv_cache"}, + candidates, + data, + lambda *_: pytest.fail("Changing only the budget must re-solve without re-scoring."), + num_calib_steps=2, + num_score_steps=2, + checkpoint=str(tmp_path / "kv_search.pth"), + ) + assert re_solved_state["best"]["constraints"]["effective_bits"] == pytest.approx(4.5) + assert {layer["selected"] for layer in re_solved_state["layers"].values()} == {"nvfp4"} + + +def test_kv_autoquant_honors_ordered_qualified_override_and_cost(nvfp4_fake_quant_stub): + model = _ToyKVModel() + candidate = (_asymmetric_kv_config().model_dump(exclude_none=True), "fp8_k_nvfp4_v") + + model, state = mtq.auto_quantize( + model, + {"effective_bits": 6.25, "cost_model": "kv_cache"}, + [candidate], + [torch.randn(1, 2, 8)], + lambda search_model, batch: search_model(batch), + num_calib_steps=1, + num_score_steps=1, + ) + + assert state["candidates"][0]["effective_bits"] == pytest.approx(6.25) + assert state["best"]["constraints"]["effective_bits"] == pytest.approx(6.25) + for layer_state in state["layers"].values(): + assert layer_state["selected"] == "fp8_k_nvfp4_v" + for layer in (model.attn0, model.attn1): + assert layer.k_bmm_quantizer.num_bits == (4, 3) + assert layer.v_bmm_quantizer.num_bits == (2, 1) + assert get_quant_config(model)["quantization"]["kv_cache_quant_algo"] == "FP8_K_NVFP4_V" + + +def test_kv_autoquant_rejects_asymmetric_candidate_for_unequal_kv_widths( + nvfp4_fake_quant_stub, +): + model = _ToyKVModel() + model.attn0.k_proj = nn.Linear(8, 12, bias=False) + model.attn0.v_proj = nn.Linear(8, 8, bias=False) + + with pytest.raises(ValueError, match=r"asymmetric K/V candidates.*unequal K/V widths"): + mtq.auto_quantize( + model, + {"effective_bits": 6.25, "cost_model": "kv_cache"}, + [(_asymmetric_kv_config().model_dump(exclude_none=True), "fp8_k_nvfp4_v")], + [torch.randn(1, 2, 8)], + lambda search_model, batch: search_model(batch), + num_calib_steps=1, + num_score_steps=1, + ) + + +def test_kv_autoquant_rejects_invalid_logits_and_restores_model_state(): + model = _ToyKVModel() + original_quantizers = { + name: (module.k_bmm_quantizer, module.v_bmm_quantizer) + for name, module in (("attn0", model.attn0), ("attn1", model.attn1)) + } + candidates = [ + ( + { + "quant_cfg": [ + { + "quantizer_name": "*[kv]_bmm_quantizer", + "cfg": { + "num_bits": (2, 1), + "block_sizes": { + -1: 16, + "type": "dynamic", + "scale_bits": (4, 3), + }, + "constant_amax": 1.0, + }, + } + ], + "algorithm": None, + "effective_bits": 4.5, + }, + "nvfp4", + ) + ] + + with pytest.raises(ValueError, match="non-empty vocabulary dimension"): + mtq.auto_quantize( + model, + {"effective_bits": 4.5, "cost_model": "kv_cache"}, + candidates, + [torch.randn(2, 3, 8)], + lambda *_: torch.ones(8), + num_calib_steps=1, + num_score_steps=1, + ) + + assert model.training + for name, module in (("attn0", model.attn0), ("attn1", model.attn1)): + assert (module.k_bmm_quantizer, module.v_bmm_quantizer) == original_quantizers[name] + + +def test_public_kv_autoquant_converts_hf_attention_and_searches(tmp_path, nvfp4_fake_quant_stub): + torch.manual_seed(123) + model = get_tiny_llama(num_hidden_layers=2) + data = [{"input_ids": torch.randint(0, model.config.vocab_size, (1, 8))} for _ in range(2)] + candidates = [ + ( + { + "quant_cfg": [ + { + "quantizer_name": "*[kv]_bmm_quantizer", + "cfg": {"num_bits": (4, 3)}, + }, + ], + "algorithm": "max", + "effective_bits": 8.0, + }, + "fp8", + ), + ( + { + "quant_cfg": [ + { + "quantizer_name": "*[kv]_bmm_quantizer", + "cfg": { + "num_bits": (2, 1), + "block_sizes": { + -1: 16, + "type": "dynamic", + "scale_bits": (4, 3), + }, + "constant_amax": 1.0, + }, + } + ], + "algorithm": None, + "effective_bits": 4.5, + }, + "nvfp4", + ), + ] + + model, state = mtq.auto_quantize( + model, + {"effective_bits": 6.25, "cost_model": "kv_cache"}, + candidates, + data, + lambda search_model, batch: search_model(**batch).logits, + num_calib_steps=2, + num_score_steps=2, + checkpoint=str(tmp_path / "hf_kv_search.pth"), + ) + + assert len(state["layers"]) == model.config.num_hidden_layers + assert state["best"]["constraints"]["effective_bits"] == pytest.approx(6.25) + assert all( + layer.self_attn.k_bmm_quantizer.is_enabled and layer.self_attn.v_bmm_quantizer.is_enabled + for layer in model.model.layers + ) + non_kv_quantizers = [ + quantizer + for name, quantizer in model.named_modules() + if isinstance(quantizer, TensorQuantizer) + and not name.endswith(("k_bmm_quantizer", "v_bmm_quantizer")) + ] + assert non_kv_quantizers + assert all(not quantizer.is_enabled for quantizer in non_kv_quantizers) + assert all( + layer.self_attn.q_proj.weight_quantizer.num_bits == 8 for layer in model.model.layers + ) + exported_quantization = get_quant_config(model)["quantization"] + assert exported_quantization["quantized_layers"] == {} + assert exported_quantization["kv_cache_quantized_layers"] + assert set(exported_quantization["kv_cache_quantized_layers"]) <= { + f"model.layers.{idx}.self_attn" for idx in range(model.config.num_hidden_layers) + } + + restored_model = get_tiny_llama(num_hidden_layers=2) + restored_model, restored_state = mtq.auto_quantize( + restored_model, + {"effective_bits": 6.25, "cost_model": "kv_cache"}, + candidates, + data, + lambda *_: pytest.fail("A compatible checkpoint must skip calibration and scoring."), + num_calib_steps=2, + num_score_steps=2, + checkpoint=str(tmp_path / "hf_kv_search.pth"), + ) + + assert restored_state["best"] == state["best"] + assert restored_state["layers"] == state["layers"] + assert any( + hasattr(layer.self_attn.k_bmm_quantizer, "_amax") + for layer in restored_model.model.layers + if layer.self_attn.k_bmm_quantizer.num_bits == (4, 3) + ) + + +@pytest.mark.parametrize( + ("model_factory", "expected_layer", "disabled_layers"), + [ + (get_tiny_qwen3, "model.layers.0.self_attn", None), + (get_tiny_qwen3vl, "model.language_model.layers.0.self_attn", "*visual*"), + ], +) +def test_public_kv_autoquant_selects_qwen_causal_attention_only( + model_factory, expected_layer, disabled_layers +): + """Plain and conditional Qwen models expose only causal attention to the KV search.""" + model = model_factory(num_hidden_layers=1) + text_config = getattr(model.config, "text_config", model.config) + data = [{"input_ids": torch.randint(0, text_config.vocab_size, (1, 8))}] + candidate = ( + _kv_config((4, 3), 8.0, algorithm=None, constant_amax=1.0).model_dump(), + "fp8", + ) + + model, state = mtq.auto_quantize( + model, + {"effective_bits": 8.0, "cost_model": "kv_cache"}, + [candidate], + data, + lambda search_model, batch: search_model(**batch).logits, + num_calib_steps=1, + num_score_steps=1, + disabled_layers=disabled_layers, + ) + + assert set(state["layers"]) == {expected_layer} + report = model._modelopt_kv_cache_auto_quantize_state + assert json.loads(json.dumps(report))["layers"][expected_layer]["selected"] == "fp8" + exported = get_quant_config(model)["quantization"] + if "kv_cache_quantized_layers" in exported: + assert set(exported["kv_cache_quantized_layers"]) == {expected_layer} + else: + assert exported["kv_cache_quant_algo"] == "FP8" + + +def test_public_kv_autoquant_validation_and_runtime_failures_are_atomic(): + model = get_tiny_llama(num_hidden_layers=1) + original_types = {name: type(module) for name, module in model.named_modules()} + invalid_candidate = _kv_config((4, 3), 8.0).model_dump() + invalid_candidate["algorithm"] = "svdquant" + + with pytest.raises(ValueError, match="only non-structural calibration algorithms"): + mtq.auto_quantize( + model, + {"effective_bits": 8.0, "cost_model": "kv_cache"}, + [invalid_candidate], + [], + lambda *_: pytest.fail("Validation must run before model conversion."), + num_calib_steps=1, + num_score_steps=1, + ) + + assert not hasattr(model, "_modelopt_state") + assert {name: type(module) for name, module in model.named_modules()} == original_types + + valid_candidate = _kv_config((4, 3), 8.0, algorithm=None, constant_amax=1.0).model_dump() + data = [{"input_ids": torch.randint(0, model.config.vocab_size, (1, 8))}] + with pytest.raises(ValueError, match="non-empty vocabulary dimension"): + mtq.auto_quantize( + model, + {"effective_bits": 8.0, "cost_model": "kv_cache"}, + [valid_candidate], + data, + lambda *_: torch.ones(8), + num_calib_steps=1, + num_score_steps=1, + ) + + assert not hasattr(model, "_modelopt_state") + assert {name: type(module) for name, module in model.named_modules()} == original_types + + +def test_public_kv_autoquant_rejects_unmatched_or_unexportable_candidates_before_conversion(): + model = get_tiny_llama(num_hidden_layers=1) + original_types = {name: type(module) for name, module in model.named_modules()} + unmatched = _kv_config((4, 3), 8.0).model_dump() + unmatched["quant_cfg"].append({"quantizer_name": "q_proj.*_quantizer", "cfg": {"num_bits": 2}}) + invalid_candidates = [ + (unmatched, "does not match a supported qualified K/V quantizer"), + (_kv_config((4, 3), 8.0, algorithm=None).model_dump(), "no persistent export scale"), + ( + _kv_config(8, 8.0, algorithm=None, constant_amax=1.0).model_dump(), + "per-tensor FP8", + ), + ] + + for candidate, match in invalid_candidates: + with pytest.raises(ValueError, match=match): + mtq.auto_quantize( + model, + {"effective_bits": 8.0, "cost_model": "kv_cache"}, + [candidate], + [], + lambda *_: pytest.fail("Validation must run before model conversion."), + num_calib_steps=1, + num_score_steps=1, + ) + assert not hasattr(model, "_modelopt_state") + assert {name: type(module) for name, module in model.named_modules()} == original_types + + +def test_public_kv_autoquant_rejects_distributed_execution_before_mutation(monkeypatch): + model = get_tiny_llama(num_hidden_layers=1) + original_types = {name: type(module) for name, module in model.named_modules()} + monkeypatch.setattr(torch.distributed, "is_available", lambda: True) + monkeypatch.setattr(torch.distributed, "is_initialized", lambda: True) + monkeypatch.setattr(torch.distributed, "get_world_size", lambda: 2) + + with pytest.raises(RuntimeError, match="single-process only"): + mtq.auto_quantize( + model, + {"effective_bits": 8.0, "cost_model": "kv_cache"}, + [], + [], + lambda *_: pytest.fail("Distributed validation must fail before search."), + ) + + assert not hasattr(model, "_modelopt_state") + assert {name: type(module) for name, module in model.named_modules()} == original_types + + +def test_public_kv_autoquant_preserves_fixed_layers_and_weight_quantizers( + monkeypatch, nvfp4_fake_quant_stub +): + torch.manual_seed(123) + model = get_tiny_llama(num_hidden_layers=2) + data = [{"input_ids": torch.randint(0, model.config.vocab_size, (1, 8))}] + fixed_kv_config = { + "quant_cfg": [ + { + "quantizer_name": "*[kv]_bmm_quantizer", + "cfg": {"num_bits": (4, 3), "constant_amax": 1.0}, + } + ], + "algorithm": None, + } + model = mtq.quantize(model, fixed_kv_config) + fixed_weight_quantizer = model.model.layers[0].self_attn.q_proj.weight_quantizer + fixed_weight_quantizer.enable() + fixed_weight_quantizer.amax = torch.tensor(1.0) + fixed_weight_quantizer.disable_quant() + fixed_weight_quantizer.disable_calib() + observed_fixed_states = [] + fixed_hook = fixed_weight_quantizer.register_forward_hook( + lambda module, _inputs, _output: observed_fixed_states.append( + (module.is_enabled, module._if_quant, module._if_calib) + ) + ) + fixed_qdq_quantizer = model.model.layers[1].self_attn.q_proj.weight_quantizer + fixed_qdq_quantizer.enable() + fixed_qdq_quantizer.amax = torch.tensor(1.0) + observed_qdq_states = [] + qdq_hook = fixed_qdq_quantizer.register_forward_hook( + lambda module, _inputs, _output: observed_qdq_states.append( + (module.is_enabled, module._if_quant, module._if_calib) + ) + ) + calibration_states = [] + real_calibrate = model_quant.calibrate + + def calibrate_with_state_check(*args, **kwargs): + calibration_states.append( + (fixed_weight_quantizer._if_quant, fixed_weight_quantizer._if_calib) + ) + result = real_calibrate(*args, **kwargs) + calibration_states.append( + (fixed_weight_quantizer._if_quant, fixed_weight_quantizer._if_calib) + ) + return result + + monkeypatch.setattr(model_quant, "calibrate", calibrate_with_state_check) + + try: + model, state = mtq.auto_quantize( + model, + {"effective_bits": 4.5, "cost_model": "kv_cache"}, + [ + ( + { + "quant_cfg": [ + { + "quantizer_name": "*[kv]_bmm_quantizer", + "cfg": { + "num_bits": (2, 1), + "block_sizes": { + -1: 16, + "type": "dynamic", + "scale_bits": (4, 3), + }, + }, + }, + ], + "algorithm": "max", + "effective_bits": 4.5, + }, + "nvfp4", + ) + ], + data, + lambda search_model, batch: search_model(**batch).logits, + num_calib_steps=1, + num_score_steps=1, + disabled_layers="model.layers.1.self_attn", + ) + finally: + fixed_hook.remove() + qdq_hook.remove() + + assert set(state["layers"]) == {"model.layers.0.self_attn"} + assert observed_fixed_states + assert all(state == (True, False, False) for state in observed_fixed_states) + assert calibration_states == [(False, False), (False, False)] + assert model.model.layers[0].self_attn.k_bmm_quantizer.num_bits == (2, 1) + assert model.model.layers[0].self_attn.q_proj.weight_quantizer.is_enabled + assert not fixed_weight_quantizer._if_quant + assert not fixed_weight_quantizer._if_calib + assert fixed_weight_quantizer.amax.item() == pytest.approx(1.0) + assert observed_qdq_states + assert all(quantizer_state == (True, True, False) for quantizer_state in observed_qdq_states) + assert fixed_qdq_quantizer._if_quant + assert not fixed_qdq_quantizer._if_calib + assert fixed_qdq_quantizer.amax.item() == pytest.approx(1.0) + fixed_attention = model.model.layers[1].self_attn + assert fixed_attention.k_bmm_quantizer.is_enabled + assert fixed_attention.v_bmm_quantizer.is_enabled + assert fixed_attention.k_bmm_quantizer.num_bits == (4, 3) + assert fixed_attention.v_bmm_quantizer.num_bits == (4, 3)