# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES # # Licensed under the Apache License, Version 2.0 (the "License"). # You may not use this file except in compliance with the License. # To view a copy of this license, visit http://www.apache.org/licenses/LICENSE-2.0 # # No warranties are given. The work is provided "AS IS", without warranty of any kind, express or implied. # # SPDX-License-Identifier: Apache-2.0 import importlib import inspect from contextlib import nullcontext from copy import deepcopy from dataclasses import dataclass import re from typing import Any import warnings import torch import torch.nn as nn from fouroversix import ( DataType, ModelQuantizationConfig, QuantizationConfig, QuantizedTensor, RoundStyle, ScaleRule, quantize_model, quantize_to_fp4, ) from fouroversix.quantize.quantized_tensor import from_blocked from utils.nvfp4_kernel import fp4_dequantize _FUSED_KV_DEQUANT_DISABLED = False _FUSED_KV_DEQUANT_WARNED = False QUANTIZATION_TYPE = { "weight": "weight", "activation": "activation", "kv": "kv", } DEFAULT_GENERATOR_FILTERED_MODULES = [ "text_embedding.0", "text_embedding.2", "patch_embedding", "time_projection.1", "time_embedding.0", "time_embedding.2", "head.head", "head.modulation", "re:.*norm_k$", "re:.*norm_q$", "re:.*norm1$", "re:.*norm2$", "re:.*norm3$" ] DEFAULT_REAL_SCORE_FILTERED_MODULES = list(DEFAULT_GENERATOR_FILTERED_MODULES) DEFAULT_FAKE_SCORE_FILTERED_MODULES = list(DEFAULT_GENERATOR_FILTERED_MODULES) DEFAULT_FILTERED_MODULES = list(DEFAULT_GENERATOR_FILTERED_MODULES) FILTER_PROFILE_ALIASES = { "generator": "generator", "student": "generator", "real_score": "real_score", "teacher": "real_score", "fake_score": "fake_score", "critic": "fake_score", } @dataclass class LongLiveQuantizationConfig(QuantizationConfig): type: str = "weight" def __post_init__(self) -> None: super().__post_init__() if not isinstance(self.type, str): raise TypeError("Quantization type must be a string.") if self.type not in QUANTIZATION_TYPE: allowed = ", ".join(QUANTIZATION_TYPE.keys()) raise ValueError(f"Unknown quantization type '{self.type}'. Expected one of: {allowed}.") self.type = QUANTIZATION_TYPE[self.type] def _resolve_modules_to_not_convert( model: nn.Module, filtered_modules: list[str] | None, ) -> list[str]: if not filtered_modules: return [] exact_names = set() regex_patterns = [] for pattern in filtered_modules: if not isinstance(pattern, str): raise TypeError("Each filtered module pattern must be a string.") if pattern.startswith("re:"): regex_patterns.append(re.compile(pattern[3:])) else: exact_names.add(pattern) resolved = [] for module_name, _ in model.named_modules(): if not module_name: continue if module_name in exact_names or any( regex.search(module_name) for regex in regex_patterns ): resolved.append(module_name) return sorted(set(resolved)) def _get_default_filtered_modules(filter_profile: str | None) -> list[str]: if filter_profile is None: return list(DEFAULT_FILTERED_MODULES) normalized = FILTER_PROFILE_ALIASES.get(filter_profile, filter_profile) if normalized == "generator": return list(DEFAULT_GENERATOR_FILTERED_MODULES) if normalized == "real_score": return list(DEFAULT_REAL_SCORE_FILTERED_MODULES) if normalized == "fake_score": return list(DEFAULT_FAKE_SCORE_FILTERED_MODULES) allowed = ", ".join(sorted(FILTER_PROFILE_ALIASES)) raise ValueError( f"Unknown filter_profile '{filter_profile}'. Expected one of: {allowed}.", ) def _warn_for_te_config_mismatch(model_quant_config: ModelQuantizationConfig) -> None: config_entries = [("default", model_quant_config)] module_overrides = getattr(model_quant_config, "module_config_overrides", None) or {} config_entries.extend(sorted(module_overrides.items())) mismatched_rules = [] for module_name, module_config in config_entries: if getattr(module_config, "dtype", DataType.nvfp4) != DataType.nvfp4: raise NotImplementedError( "TransformerEngine replacement currently only supports NVFP4." ) for attr_name in ( "scale_rule", "activation_scale_rule", "weight_scale_rule", "gradient_scale_rule", ): rule = getattr(module_config, attr_name, None) if rule is not None and rule != ScaleRule.static_6: mismatched_rules.append(f"{module_name}:{attr_name}={rule}") if mismatched_rules: preview = ", ".join(mismatched_rules[:8]) if len(mismatched_rules) > 8: preview += ", ..." warnings.warn( "TransformerEngine NVFP4 path maps to `NVFP4BlockScaling` and does not " "replicate FourOverSix non-`static_6` scale rules exactly. " f"Mismatched config entries: {preview}", stacklevel=3, ) def _build_te_recipe(module_config: Any, te_recipe_kwargs: dict[str, Any] | None = None): recipe_module = importlib.import_module("transformer_engine.common.recipe") NVFP4BlockScaling = recipe_module.NVFP4BlockScaling recipe_kwargs = { "disable_2d_quantization": not getattr(module_config, "weight_scale_2d", False), "disable_stochastic_rounding": ( getattr(module_config, "gradient_round_style", RoundStyle.nearest) != RoundStyle.stochastic ), # FourOverSix only uses RHT in specific training paths, so keep TE conservative # by default and let callers override via `te_recipe_kwargs`. "disable_rht": True, } if te_recipe_kwargs: recipe_kwargs.update(te_recipe_kwargs) return NVFP4BlockScaling(**recipe_kwargs) class TransformerEngineLinear(nn.Module): """A lightweight wrapper that routes a linear layer through TransformerEngine.""" def __init__( self, module: nn.Linear, module_name: str, module_config: Any, inference_only: bool = False, low_precision_weights: bool = False, te_recipe_kwargs: dict[str, Any] | None = None, te_module_kwargs: dict[str, Any] | None = None, ) -> None: super().__init__() try: te = importlib.import_module("transformer_engine.pytorch") except ImportError as exc: raise ImportError( "TransformerEngine is not installed, but `use_transformer_engine=True` " "was requested." ) from exc if module.weight.device.type != "cuda": raise ValueError( "TransformerEngine replacement requires CUDA modules. " f"Module `{module_name}` is on `{module.weight.device}`." ) self.module_name = module_name self.in_features = module.in_features self.out_features = module.out_features self.inference_only = inference_only self.low_precision_weights = low_precision_weights self._te = te self._recipe = _build_te_recipe( module_config=module_config, te_recipe_kwargs=te_recipe_kwargs, ) module_kwargs = dict(te_module_kwargs or {}) module_kwargs.setdefault("device", module.weight.device) module_kwargs.setdefault("params_dtype", module.weight.dtype) module_kwargs.setdefault("name", module_name) fp8_model_init_fn = getattr(te, "fp8_model_init", None) if self.low_precision_weights and fp8_model_init_fn is None: warnings.warn( "TransformerEngine low-precision parameter init requested, but " "`fp8_model_init` is unavailable. Falling back to regular TE parameter " "storage for this inference path.", stacklevel=2, ) self.low_precision_weights = False model_init_context = ( fp8_model_init_fn( enabled=True, recipe=self._recipe, preserve_high_precision_init_val=False, ) if self.low_precision_weights else nullcontext() ) with model_init_context: self.linear = te.Linear( module.in_features, module.out_features, bias=module.bias is not None, **module_kwargs, ) self._load_from_linear(module) if self.inference_only: self.linear.requires_grad_(False) self.train(False) else: self.linear.weight.requires_grad_(module.weight.requires_grad) if self.linear.bias is not None and module.bias is not None: self.linear.bias.requires_grad_(module.bias.requires_grad) self.train(module.training) def _copy_tensor_into_parameter( self, destination: torch.Tensor, source: torch.Tensor, ) -> None: source = source.detach().to(device=destination.device) try: destination.copy_(source) return except Exception: pass destination.copy_(source.to(dtype=destination.dtype)) def _load_from_linear(self, module: nn.Linear) -> None: with torch.no_grad(): try: self._copy_tensor_into_parameter(self.linear.weight, module.weight) if module.bias is not None and self.linear.bias is not None: self._copy_tensor_into_parameter(self.linear.bias, module.bias) return except Exception as copy_exc: state_dict = { "weight": module.weight.detach().to(device=self.linear.weight.device), } if module.bias is not None: state_dict["bias"] = module.bias.detach().to( device=self.linear.weight.device, ) incompatible_keys = self.linear.load_state_dict(state_dict, strict=False) missing_keys = [ key for key in getattr(incompatible_keys, "missing_keys", []) if key != "_extra_state" ] unexpected_keys = list(getattr(incompatible_keys, "unexpected_keys", [])) if missing_keys or unexpected_keys: raise RuntimeError( "Failed to load weights into TransformerEngine linear " f"`{self.module_name}`. missing_keys={missing_keys}, " f"unexpected_keys={unexpected_keys}" ) from copy_exc @property def weight(self) -> torch.Tensor: return self.linear.weight @property def bias(self) -> torch.Tensor | None: return self.linear.bias def _autocast_context(self): autocast_fn = getattr(self._te, "autocast", None) if autocast_fn is not None: return autocast_fn(enabled=True, recipe=self._recipe) fp8_autocast_fn = getattr(self._te, "fp8_autocast", None) if fp8_autocast_fn is None: raise AttributeError( "TransformerEngine does not expose `autocast` or `fp8_autocast`." ) return fp8_autocast_fn(enabled=True, fp8_recipe=self._recipe) def forward(self, input: torch.Tensor) -> torch.Tensor: with self._autocast_context(): return self.linear(input) def extra_repr(self) -> str: return ( f"in_features={self.in_features}, " f"out_features={self.out_features}, " f"bias={self.bias is not None}, " f"inference_only={self.inference_only}, " f"low_precision_weights={self.low_precision_weights}, " "backend=transformer_engine" ) def quantize_model_with_optional_te( model: nn.Module, model_quant_config: ModelQuantizationConfig, *, use_transformer_engine: bool = False, te_inference_only: bool = False, te_low_precision_weights: bool | None = None, te_recipe_kwargs: dict[str, Any] | None = None, te_module_kwargs: dict[str, Any] | None = None, te_fallback_to_fouroversix: bool = False, **kwargs, ) -> list[str]: """ Quantize a model with FourOverSix by default, or replace `nn.Linear` with TransformerEngine wrappers when `use_transformer_engine=True`. """ if not use_transformer_engine: quantize_model(model, model_quant_config, **kwargs) return [] if te_low_precision_weights is None: te_low_precision_weights = te_inference_only if kwargs: if te_fallback_to_fouroversix: warnings.warn( "Additional kwargs passed to `quantize_model_with_optional_te` will " "only be forwarded to the fallback FourOverSix pass after " f"TransformerEngine replacement: {sorted(kwargs)}", stacklevel=2, ) else: warnings.warn( "Additional kwargs passed to `quantize_model` are ignored in the " f"TransformerEngine path: {sorted(kwargs)}", stacklevel=2, ) _warn_for_te_config_mismatch(model_quant_config) replaced_modules = [] for module_name, module in list(model.named_modules()): if ( module_name == "" or module_name in model_quant_config.modules_to_not_convert or not isinstance(module, nn.Linear) ): continue model.set_submodule( module_name, TransformerEngineLinear( module=module, module_name=module_name, module_config=model_quant_config.get_module_config(module_name), inference_only=te_inference_only, low_precision_weights=te_low_precision_weights, te_recipe_kwargs=te_recipe_kwargs, te_module_kwargs=te_module_kwargs, ), ) replaced_modules.append(module_name) if te_fallback_to_fouroversix: quantize_model(model, model_quant_config, **kwargs) return replaced_modules def _tensor_nbytes(tensor: torch.Tensor | None) -> int: if tensor is None: return 0 return tensor.numel() * tensor.element_size() def _materialize_transformer_engine_weights_for_inference( model: nn.Module, target_device: torch.device | str | None = None, cache_transposed_weights: bool = False, ) -> tuple[list[str], int, int]: del cache_transposed_weights materialized_modules = [] master_weight_bytes = 0 quantized_weight_bytes = 0 for module_name, module in model.named_modules(): if not isinstance(module, TransformerEngineLinear): continue if target_device is not None: module.to(device=torch.device(target_device)) quantized_weight_bytes += _tensor_nbytes(module.weight) quantized_weight_bytes += _tensor_nbytes(module.bias) materialized_modules.append(module_name) return materialized_modules, master_weight_bytes, quantized_weight_bytes def _materialize_mixed_quantized_weights_for_inference( model: nn.Module, target_device: torch.device | str | None = None, cache_transposed_weights: bool = False, ) -> tuple[list[str], int, int]: te_modules, te_master_bytes, te_quantized_bytes = ( _materialize_transformer_engine_weights_for_inference( model, target_device=target_device, cache_transposed_weights=cache_transposed_weights, ) ) f46_modules, f46_master_bytes, f46_quantized_bytes = ( _materialize_quantized_weights_for_inference( model, target_device=target_device, cache_transposed_weights=cache_transposed_weights, ) ) return ( sorted(set(te_modules + f46_modules)), te_master_bytes + f46_master_bytes, te_quantized_bytes + f46_quantized_bytes, ) def _materialize_quantized_weights_for_inference( model: nn.Module, target_device: torch.device | str | None = None, cache_transposed_weights: bool = False, ) -> tuple[list[str], int, int]: """ Materialize quantized weights and drop master weights. Optionally cache an additional transposed quantized layout for training paths that still require dgrad after the master weight is deleted (e.g. NVFP4 + LoRA). This function expects modules replaced by `fouroversix.quantize_model`. """ materialized_modules = [] master_weight_bytes = 0 quantized_weight_bytes = 0 for module_name, module in model.named_modules(): if not hasattr(module, "parameters_to_quantize") or not hasattr( module, "get_quantized_parameters", ): continue parameters_to_quantize = getattr(module, "parameters_to_quantize", ()) if callable(parameters_to_quantize): parameters_to_quantize = parameters_to_quantize() if not parameters_to_quantize: continue did_materialize = False for parameter_name in parameters_to_quantize: parameter = getattr(module, parameter_name, None) if parameter is None: continue if isinstance(parameter, nn.Parameter): parameter_tensor = parameter.data elif isinstance(parameter, torch.Tensor): parameter_tensor = parameter else: continue master_weight_bytes += parameter_tensor.numel() * parameter_tensor.element_size() get_quantized_parameters = module.get_quantized_parameters if ( cache_transposed_weights and "include_transposed" in inspect.signature( get_quantized_parameters, ).parameters ): quantized_params = get_quantized_parameters( parameter_name, parameter_tensor, include_transposed=True, ) else: quantized_params = get_quantized_parameters( parameter_name, parameter_tensor, ) for quantized_name, quantized_tensor in quantized_params.items(): if not isinstance(quantized_tensor, torch.Tensor): continue existing = getattr(module, quantized_name, None) dst_dtype = ( existing.dtype if isinstance(existing, torch.Tensor) else quantized_tensor.dtype ) if target_device is not None: dst_device = torch.device(target_device) elif isinstance(existing, torch.Tensor): dst_device = existing.device else: dst_device = quantized_tensor.device quantized_tensor = quantized_tensor.to( device=dst_device, dtype=dst_dtype, ) setattr(module, quantized_name, quantized_tensor) quantized_weight_bytes += ( quantized_tensor.numel() * quantized_tensor.element_size() ) # Drop high-precision master weight once quantized weights are materialized. if isinstance(getattr(module, parameter_name, None), nn.Parameter): module.register_parameter(parameter_name, None) else: setattr(module, parameter_name, None) did_materialize = True if did_materialize: if hasattr(module, "_quantized_weight"): delattr(module, "_quantized_weight") if hasattr(module, "_quantized_weight_transposed"): delattr(module, "_quantized_weight_transposed") if hasattr(module, "_quantized_weights"): delattr(module, "_quantized_weights") if hasattr(module, "config") and hasattr(module.config, "keep_master_weights"): module.config.keep_master_weights = False materialized_modules.append(module_name) return materialized_modules, master_weight_bytes, quantized_weight_bytes def quantize_model_with_filter( model: nn.Module, quant_config: ModelQuantizationConfig | dict | None = None, filtered_modules: list[str] | None = None, filter_profile: str | None = None, use_default_filtered_modules: bool = False, cast_model_to_bf16: bool = True, materialize_for_inference: bool = False, materialize_target_device: torch.device | str | None = None, use_transformer_engine: bool = False, te_inference_only: bool = False, te_low_precision_weights: bool | None = None, te_recipe_kwargs: dict[str, Any] | None = None, te_module_kwargs: dict[str, Any] | None = None, te_fallback_to_fouroversix: bool = False, verbose: bool = True, **kwargs, ) -> tuple[nn.Module, list[str]]: """ Quantize model with FourOverSix and optionally skip selected modules. `filtered_modules` supports: - Exact module names, e.g. "head.head" - Regex patterns prefixed with "re:", e.g. "re:.*norm1$" `filter_profile` selects which built-in filtered module profile to use when `use_default_filtered_modules=True`. Supported values: "generator"/"student" and "real_score"/"teacher". """ if quant_config is None: model_quant_config = ModelQuantizationConfig() elif isinstance(quant_config, dict): model_quant_config = ModelQuantizationConfig(**quant_config) elif isinstance(quant_config, ModelQuantizationConfig): model_quant_config = deepcopy(quant_config) else: raise TypeError( "quant_config must be ModelQuantizationConfig, dict, or None.", ) patterns = list(filtered_modules or []) if use_default_filtered_modules: patterns = _get_default_filtered_modules(filter_profile) + patterns matched_modules = _resolve_modules_to_not_convert(model, patterns) modules_to_not_convert = set(model_quant_config.modules_to_not_convert or []) modules_to_not_convert.update(matched_modules) model_quant_config.modules_to_not_convert = sorted(modules_to_not_convert) if cast_model_to_bf16: model.to(torch.bfloat16) resolved_te_low_precision_weights = ( te_inference_only if te_low_precision_weights is None else te_low_precision_weights ) te_replaced_modules = quantize_model_with_optional_te( model, model_quant_config, use_transformer_engine=use_transformer_engine, te_inference_only=te_inference_only, te_low_precision_weights=resolved_te_low_precision_weights, te_recipe_kwargs=te_recipe_kwargs, te_module_kwargs=te_module_kwargs, te_fallback_to_fouroversix=te_fallback_to_fouroversix, **kwargs, ) if materialize_for_inference: materialize_fn = _materialize_quantized_weights_for_inference if use_transformer_engine and te_fallback_to_fouroversix: materialize_fn = _materialize_mixed_quantized_weights_for_inference elif use_transformer_engine: materialize_fn = _materialize_transformer_engine_weights_for_inference materialized_modules, master_bytes, quantized_bytes = materialize_fn( model, target_device=materialize_target_device, ) if verbose: print( "[quantize_model_with_filter] " f"materialized_modules={len(materialized_modules)}, " f"master_weight={master_bytes / (1024 ** 3):.3f} GiB, " f"quantized_weight={quantized_bytes / (1024 ** 3):.3f} GiB", ) if verbose: profile_label = filter_profile or "default" print( "[quantize_model_with_filter] " f"profile={profile_label}, " f"matched={len(matched_modules)}, " f"total_excluded={len(model_quant_config.modules_to_not_convert)}", ) if use_transformer_engine: print( "[quantize_model_with_filter] " f"transformer_engine_replaced={len(te_replaced_modules)}, " f"inference_only={te_inference_only}, " f"low_precision_weights={resolved_te_low_precision_weights}, " f"fallback_to_fouroversix={te_fallback_to_fouroversix}", ) return model, matched_modules def _dequantize_kv_cache_fused_cuda(kv_list, max_blocks, num_heads, block_token_size, dtype): global _FUSED_KV_DEQUANT_DISABLED, _FUSED_KV_DEQUANT_WARNED if _FUSED_KV_DEQUANT_DISABLED or max_blocks <= 0: return None first_qt = kv_list[0] if first_qt.values.device.type != "cuda": return None try: from utils.kernel.kv_dequant import dequantize_kv_cache_fp4 blocks = kv_list[:max_blocks] values = [qt.values for qt in blocks] scale_factors = [qt.scale_factors for qt in blocks] amax = [qt.amax for qt in blocks] return dequantize_kv_cache_fp4( values, scale_factors, amax, num_heads=num_heads, block_token_size=block_token_size, dtype=dtype, scale_rule=first_qt.scale_rule, ) except Exception as exc: # pragma: no cover - exercised only when extension is stale/missing _FUSED_KV_DEQUANT_DISABLED = True if not _FUSED_KV_DEQUANT_WARNED: warnings.warn( "Fused CUDA KV-cache dequantization is unavailable; falling back to " f"the Triton per-block path. Reason: {exc}", stacklevel=2, ) _FUSED_KV_DEQUANT_WARNED = True return None def dequantize_kv_cache(kv_list, max_blocks, num_heads, block_token_size, dtype, device): """ Dequantize list of QuantizedTensor to a contiguous bf16 tensor. kv_list[block_idx] -> QuantizedTensor(block_token_size * num_heads, 128) Returns: [1, max_blocks * block_token_size, num_heads, 128] """ fused_result = _dequantize_kv_cache_fused_cuda( kv_list, max_blocks, num_heads, block_token_size, dtype, ) if fused_result is not None: return fused_result total_tokens = max_blocks * block_token_size result = torch.zeros([1, total_tokens, num_heads, 128], dtype=dtype, device=device) for block_idx in range(max_blocks): t_start = block_idx * block_token_size t_end = t_start + block_token_size # deq = kv_list[block_idx].dequantize(dtype) # triton fp4_dequantize qt = kv_list[block_idx] padded_shape = qt.padded_shape scales_2d = from_blocked( qt.scale_factors, (padded_shape[0], padded_shape[1] // 16), ) global_scale = qt.amax / ( qt.scale_rule.max_allowed_e2m1_value() * qt.scale_rule.max_allowed_e4m3_value() ) deq = fp4_dequantize( kv_list[block_idx].values, scales_2d, global_scale, block_size=16, dtype=dtype, ) result[0, t_start:t_end, :, :] = deq.view(block_token_size, num_heads, 128) return result def clone_quantized_tensor(qt): """Clone a QuantizedTensor by cloning its internal tensors.""" return QuantizedTensor( values=qt.values.clone(), scale_factors=qt.scale_factors.clone(), amax=qt.amax.clone() if qt.amax is not None else None, dtype=qt.dtype, original_shape=qt.original_shape, scale_rule=qt.scale_rule, padded_shape=qt.padded_shape, ) def copy_quantized_into(slot: QuantizedTensor, src: QuantizedTensor) -> None: """In-place copy a QuantizedTensor's data into a pre-allocated slot. Keeps the slot's `values`/`scale_factors`/`amax` buffers persistent (their addresses don't change) so cudagraph allocator does not see them as fresh outputs that can be reused across step boundaries. Used by the quantized KV cache rolling/insert paths. """ slot.values.copy_(src.values) slot.scale_factors.copy_(src.scale_factors) if src.amax is not None and slot.amax is not None: slot.amax.copy_(src.amax) def k_smooth(k: torch.Tensor) -> torch.Tensor: return k - k.mean(dim=-1, keepdim=True) def quantize_kv(k: torch.Tensor, v: torch.Tensor) -> torch.Tensor: B, S, H, D = k.shape # B is always 1 # S is the number of tokens # H is the number of heads # D is the dimension of the key and value config = QuantizationConfig(scale_rule="mse", backend="cuda") # per head quantization for head in range(H): k_head = k[:, :, head, :] v_head = v[:, :, head, :] k_head = k_smooth(k_head) v_head = v_head k_head = quantize_to_fp4(k_head, config) v_head = quantize_to_fp4(v_head, config) k[:, :, head, :] = k_head v[:, :, head, :] = v_head return k, v