@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / d0105b84f0

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 7, 16:30 UTC

add mixed precision support for modelopt (#35047)

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

sychen52 committed Feb 26, 2026 at 13:56 UTC d0105b84f00fadc18d9e7859d3e76887b7f5772c

4 files changed +250 -83

vllm/config/model.py

@@ -883,6 +883,7 @@ class ModelConfig:

883 "modelopt",

884 "modelopt_fp4",

885 "modelopt_mxfp8",

886 + "modelopt_mixed",

887 "petit_nvfp4",

888 # Ensure heavy backends are probed last to avoid unnecessary

889 # imports during override detection (e.g., MXFP4 imports Triton)

vllm/model_executor/layers/quantization/__init__.py

+8 -1

@@ -18,6 +18,7 @@ QuantizationMethods = Literal[

18 "modelopt",

19 "modelopt_fp4",

20 "modelopt_mxfp8",

21 + "modelopt_mixed",

22 "gguf",

23 "gptq_marlin",

24 "awq_marlin",

@@ -120,7 +121,12 @@ def get_quantization_config(quantization: str) -> type[QuantizationConfig]:

121 from .gptq import GPTQConfig

122 from .gptq_marlin import GPTQMarlinConfig

123 from .inc import INCConfig

123 - from .modelopt import ModelOptFp8Config, ModelOptMxFp8Config, ModelOptNvFp4Config

124 + from .modelopt import (

125 + ModelOptFp8Config,

126 + ModelOptMixedPrecisionConfig,

127 + ModelOptMxFp8Config,

128 + ModelOptNvFp4Config,

129 + )

130 from .moe_wna16 import MoeWNA16Config

131 from .mxfp4 import Mxfp4Config

132 from .petit import PetitNvFp4Config

@@ -135,6 +141,7 @@ def get_quantization_config(quantization: str) -> type[QuantizationConfig]:

141 "modelopt": ModelOptFp8Config,

142 "modelopt_fp4": ModelOptNvFp4Config,

143 "modelopt_mxfp8": ModelOptMxFp8Config,

144 + "modelopt_mixed": ModelOptMixedPrecisionConfig,

145 "gguf": GGUFConfig,

146 "gptq_marlin": GPTQMarlinConfig,

147 "awq_marlin": AWQMarlinConfig,

vllm/model_executor/layers/quantization/modelopt.py

+228 -79

@@ -114,6 +114,8 @@ QUANT_ALGOS = [

114 "NVFP4",

115 # MXFP8

116 "MXFP8",

117 + # MIXED_PRECISION,

118 + "MIXED_PRECISION",

119 ]

120 KV_CACHE_QUANT_ALGOS = ["FP8"]

121

@@ -235,6 +237,26 @@ class ModelOptQuantConfigBase(QuantizationConfig):

237

238 self.exclude_modules = hf_to_vllm_mapper.apply_list(new_exclude_modules)

239

240 + @staticmethod

241 + def _extract_modelopt_quant_algo(

242 + hf_quant_cfg: dict[str, Any] | None,

243 + ) -> str | None:

244 + """Extract upper-cased quant_algo from a modelopt config.

245 +

246 + Returns the quant_algo string (upper-cased), or None if the config

247 + is not a modelopt config.

248 + """

249 + if hf_quant_cfg is None:

250 + return None

251 + if hf_quant_cfg.get("quant_method", "").lower() != "modelopt":

252 + return None

253 + if "quantization" in hf_quant_cfg:

254 + quant_config = hf_quant_cfg["quantization"]

255 + if isinstance(quant_config, dict):

256 + return str(quant_config.get("quant_algo", "")).upper()

257 + return None

258 + return str(hf_quant_cfg.get("quant_algo", "")).upper()

259 +

260 @staticmethod

261 def get_config_filenames() -> list[str]:

262 return ["hf_quant_config.json"]

@@ -272,10 +294,20 @@ class ModelOptQuantConfigBase(QuantizationConfig):

294 # "exclude_modules" is the key in the legacy hf_quant_config.json

295 exclude_modules = quant_config.get("exclude_modules", [])

296 else:

275 - # Compressed-tensors style format:

297 + # Compressed-tensors style format (config.json quantization_config):

298 # {"quant_algo": "...", "quant_method": "modelopt"}

299 quant_method = config.get("quant_algo")

278 - kv_cache_quant_method = config.get("kv_cache_quant_algo")

300 +

301 + # "kv_cache_scheme" (a dict) instead of "kv_cache_quant_algo" (a string).

302 + kv_cache_scheme = config.get("kv_cache_scheme")

303 + if isinstance(kv_cache_scheme, dict) and (

304 + kv_cache_scheme.get("type") == "float"

305 + and kv_cache_scheme.get("num_bits") == 8

306 + ):

307 + kv_cache_quant_method = "FP8"

308 + else:

309 + kv_cache_quant_method = None

310 +

311 # "ignore" is the key in config.json

312 exclude_modules = config.get("ignore", [])

313 group_size_raw = config.get("group_size")

@@ -379,32 +411,9 @@ class ModelOptFp8Config(ModelOptQuantConfigBase):

411 def override_quantization_method(

412 cls, hf_quant_cfg, user_quant

413 ) -> QuantizationMethods | None:

382 - """Detect if this ModelOpt config should be used based on

383 - quantization config."""

384 -

385 - if hf_quant_cfg is None:

386 - return None

387 -

388 - # Use the community standard 'quant_method'

389 - quant_method = hf_quant_cfg.get("quant_method", "").lower()

390 -

391 - # Only proceed if the method is explicitly "modelopt"

392 - if quant_method != "modelopt":

393 - return None

394 -

395 - # Look for ModelOpt-specific config structure

396 - if "quantization" in hf_quant_cfg:

397 - quant_config = hf_quant_cfg["quantization"]

398 - if isinstance(quant_config, dict):

399 - quant_algo = str(quant_config.get("quant_algo", ""))

400 - if quant_algo.upper() == "FP8":

401 - return "modelopt"

402 - else:

403 - # Check for compressed-tensors style config with specific quant_algo

404 - quant_algo = str(hf_quant_cfg.get("quant_algo", ""))

405 - if quant_algo.upper() == "FP8":

406 - return "modelopt"

407 -

414 + algo = cls._extract_modelopt_quant_algo(hf_quant_cfg)

415 + if algo is not None and algo == "FP8":

416 + return "modelopt"

417 return None

418

419 @classmethod

@@ -1031,32 +1040,9 @@ class ModelOptNvFp4Config(ModelOptQuantConfigBase):

1040 def override_quantization_method(

1041 cls, hf_quant_cfg, user_quant

1042 ) -> QuantizationMethods | None:

1034 - """Detect if this ModelOpt FP4 config should be used based on

1035 - quantization config."""

1036 - if hf_quant_cfg is None:

1037 - return None

1038 -

1039 - # Use the community standard 'quant_method'

1040 - quant_method = hf_quant_cfg.get("quant_method", "").lower()

1041 -

1042 - # Only proceed if the method is explicitly "modelopt"

1043 - if quant_method != "modelopt":

1044 - return None

1045 -

1046 - # Look for ModelOpt-specific config structure

1047 - if "quantization" in hf_quant_cfg:

1048 - quant_config = hf_quant_cfg["quantization"]

1049 - if isinstance(quant_config, dict):

1050 - quant_algo = quant_config.get("quant_algo", "")

1051 - if "NVFP4" in quant_algo:

1052 - return "modelopt_fp4"

1053 - else:

1054 - # Check for compressed-tensors style config with specific

1055 - # quant_algo field

1056 - quant_algo = hf_quant_cfg.get("quant_algo", "")

1057 - if isinstance(quant_algo, str) and "FP4" in quant_algo.upper():

1058 - return "modelopt_fp4"

1059 -

1043 + algo = cls._extract_modelopt_quant_algo(hf_quant_cfg)

1044 + if algo is not None and ("NVFP4" in algo or "FP4" in algo):

1045 + return "modelopt_fp4"

1046 return None

1047

1048 @classmethod

@@ -1619,31 +1605,9 @@ class ModelOptMxFp8Config(ModelOptQuantConfigBase):

1605 def override_quantization_method(

1606 cls, hf_quant_cfg, user_quant

1607 ) -> QuantizationMethods | None:

1622 - """Detect if this ModelOpt MXFP8 config should be used based on

1623 - quantization config."""

1624 - if hf_quant_cfg is None:

1625 - return None

1626 -

1627 - # Use the community standard 'quant_method'

1628 - quant_method = hf_quant_cfg.get("quant_method", "").lower()

1629 -

1630 - # Only proceed if the method is explicitly "modelopt"

1631 - if quant_method != "modelopt":

1632 - return None

1633 -

1634 - # Look for ModelOpt-specific config structure

1635 - if "quantization" in hf_quant_cfg:

1636 - quant_config = hf_quant_cfg["quantization"]

1637 - if isinstance(quant_config, dict):

1638 - quant_algo = str(quant_config.get("quant_algo", "")).upper()

1639 - if "MXFP8" in quant_algo:

1640 - return "modelopt_mxfp8"

1641 - else:

1642 - # Check for compressed-tensors style config with specific quant_algo

1643 - quant_algo = str(hf_quant_cfg.get("quant_algo", "")).upper()

1644 - if "MXFP8" in quant_algo:

1645 - return "modelopt_mxfp8"

1646 -

1608 + algo = cls._extract_modelopt_quant_algo(hf_quant_cfg)

1609 + if algo is not None and "MXFP8" in algo:

1610 + return "modelopt_mxfp8"

1611 return None

1612

1613 @classmethod

@@ -1841,3 +1805,188 @@ class ModelOptMxFp8LinearMethod(LinearMethodBase):

1805 # Register the method classes for ModelOptMxFp8Config

1806 ModelOptMxFp8Config.LinearMethodCls = ModelOptMxFp8LinearMethod

1807 ModelOptMxFp8Config.KVCacheMethodCls = ModelOptFp8KVCacheMethod

1808 +

1809 +

1810 +class ModelOptMixedPrecisionConfig(ModelOptQuantConfigBase):

1811 + """Config class for ModelOpt MIXED_PRECISION.

1812 +

1813 + Supports checkpoints where different layers use different quantization

1814 + algorithms (e.g., FP8 for dense layers and NVFP4 for MoE experts).

1815 + The per-layer algorithm is specified in the ``quantized_layers`` dict

1816 + inside ``config.json``'s ``quantization_config`` (preferred) or the

1817 + legacy ``hf_quant_config.json``.

1818 + """

1819 +

1820 + def __init__(

1821 + self,

1822 + kv_cache_quant_method: str | None,

1823 + exclude_modules: list[str],

1824 + quantized_layers: dict[str, dict[str, Any]],

1825 + fp8_config: ModelOptFp8Config,

1826 + nvfp4_config: ModelOptNvFp4Config,

1827 + ) -> None:

1828 + super().__init__(exclude_modules)

1829 + self.kv_cache_quant_method = kv_cache_quant_method

1830 + self.quantized_layers = quantized_layers

1831 + self.fp8_config = fp8_config

1832 + self.nvfp4_config = nvfp4_config

1833 +

1834 + def get_name(self) -> QuantizationMethods:

1835 + return "modelopt_mixed"

1836 +

1837 + def get_supported_act_dtypes(self) -> list[torch.dtype]:

1838 + return [torch.bfloat16, torch.half]

1839 +

1840 + @classmethod

1841 + def get_min_capability(cls) -> int:

1842 + return 89

1843 +

1844 + @classmethod

1845 + def override_quantization_method(

1846 + cls, hf_quant_cfg, user_quant

1847 + ) -> QuantizationMethods | None:

1848 + algo = cls._extract_modelopt_quant_algo(hf_quant_cfg)

1849 + if algo is not None and algo == "MIXED_PRECISION":

1850 + return "modelopt_mixed"

1851 + return None

1852 +

1853 + @classmethod

1854 + def _from_config(

1855 + cls,

1856 + *,

1857 + quant_method: str,

1858 + kv_cache_quant_method: str | None,

1859 + exclude_modules: list[str],

1860 + original_config: dict[str, Any],

1861 + group_size: int | None,

1862 + **kwargs: Any,

1863 + ) -> "ModelOptMixedPrecisionConfig":

1864 + if "quantization" in original_config:

1865 + quantized_layers = original_config["quantization"].get(

1866 + "quantized_layers", {}

1867 + )

1868 + else:

1869 + quantized_layers = original_config.get("quantized_layers", {})

1870 +

1871 + if not quantized_layers:

1872 + raise ValueError(

1873 + "MIXED_PRECISION quant_algo requires a non-empty "

1874 + "'quantized_layers' mapping in the quantization config."

1875 + )

1876 +

1877 + # Determine group_size from the first NVFP4 entry if not provided.

1878 + if group_size is None:

1879 + for layer_info in quantized_layers.values():

1880 + if layer_info.get("quant_algo", "").upper() == "NVFP4":

1881 + group_size = layer_info.get("group_size", 16)

1882 + break

1883 + if group_size is None:

1884 + group_size = 16

1885 +

1886 + fp8_config = ModelOptFp8Config(

1887 + quant_method="FP8",

1888 + is_checkpoint_fp8_serialized=True,

1889 + kv_cache_quant_method=kv_cache_quant_method,

1890 + exclude_modules=[],

1891 + )

1892 + nvfp4_config = ModelOptNvFp4Config(

1893 + is_checkpoint_nvfp4_serialized=True,

1894 + kv_cache_quant_algo=kv_cache_quant_method,

1895 + exclude_modules=[],

1896 + group_size=group_size,

1897 + )

1898 +

1899 + return cls(

1900 + kv_cache_quant_method=kv_cache_quant_method,

1901 + exclude_modules=exclude_modules,

1902 + quantized_layers=quantized_layers,

1903 + fp8_config=fp8_config,

1904 + nvfp4_config=nvfp4_config,

1905 + )

1906 +

1907 + def _resolve_quant_algo(self, prefix: str) -> str | None:

1908 + """Look up the quant_algo for a vLLM-side layer prefix.

1909 +

1910 + Tries three strategies in order:

1911 + 1. Direct lookup in ``quantized_layers``.

1912 + 2. Packed/fused-layer lookup (unfuse via ``packed_modules_mapping``).

1913 + 3. Prefix-based lookup for FusedMoE (any child key starts with

1914 + ``prefix + "."``).

1915 +

1916 + Returns the upper-cased quant_algo string, or *None* if the prefix

1917 + is not found.

1918 + """

1919 + # 1. Direct lookup

1920 + if prefix in self.quantized_layers:

1921 + return self.quantized_layers[prefix]["quant_algo"].upper()

1922 +

1923 + # 2. Packed / fused layer lookup

1924 + proj_name = prefix.rsplit(".", 1)[-1]

1925 + if self.packed_modules_mapping and proj_name in self.packed_modules_mapping:

1926 + algos: set[str] = set()

1927 + base = prefix.rsplit(".", 1)[0]

1928 + for shard_name in self.packed_modules_mapping[proj_name]:

1929 + shard_prefix = f"{base}.{shard_name}"

1930 + if shard_prefix in self.quantized_layers:

1931 + algos.add(self.quantized_layers[shard_prefix]["quant_algo"].upper())

1932 + if len(algos) == 1:

1933 + return algos.pop()

1934 + if len(algos) > 1:

1935 + raise ValueError(

1936 + f"Mixed quant_algo within fused layer {prefix}: "

1937 + f"{algos}. All shards must use the same quantization."

1938 + )

1939 +

1940 + # 3. Prefix-based lookup (for FusedMoE / parent modules)

1941 + prefix_dot = prefix + "."

1942 + for key, info in self.quantized_layers.items():

1943 + if key.startswith(prefix_dot):

1944 + return info["quant_algo"].upper()

1945 +

1946 + return None

1947 +

1948 + def get_quant_method(

1949 + self, layer: torch.nn.Module, prefix: str

1950 + ) -> "QuantizeMethodBase | None":

1951 + """Return quantize-method based on layer."""

1952 + # KV-cache quantization

1953 + if isinstance(layer, Attention):

1954 + if self.kv_cache_quant_method:

1955 + return ModelOptFp8KVCacheMethod(self)

1956 + return None

1957 +

1958 + # Excluded layers

1959 + if self.is_layer_excluded(prefix):

1960 + if isinstance(layer, LinearBase):

1961 + return UnquantizedLinearMethod()

1962 + return None

1963 +

1964 + quant_algo = self._resolve_quant_algo(prefix)

1965 +

1966 + if isinstance(layer, LinearBase):

1967 + if quant_algo == "FP8":

1968 + return ModelOptFp8LinearMethod(self.fp8_config)

1969 + if quant_algo == "NVFP4":

1970 + return ModelOptNvFp4LinearMethod(self.nvfp4_config)

1971 + # Layer not in quantized_layers — leave unquantized

1972 + return UnquantizedLinearMethod()

1973 +

1974 + if isinstance(layer, FusedMoE):

1975 + if quant_algo == "FP8":

1976 + return ModelOptFp8MoEMethod(

1977 + quant_config=self.fp8_config,

1978 + moe_config=layer.moe_config,

1979 + )

1980 + if quant_algo == "NVFP4":

1981 + return ModelOptNvFp4FusedMoE(

1982 + quant_config=self.nvfp4_config,

1983 + moe_config=layer.moe_config,

1984 + )

1985 + return None

1986 +

1987 + return None

1988 +

1989 + def apply_vllm_mapper(self, hf_to_vllm_mapper: "WeightsMapper"):

1990 + super().apply_vllm_mapper(hf_to_vllm_mapper)

1991 + if self.quantized_layers:

1992 + self.quantized_layers = hf_to_vllm_mapper.apply_dict(self.quantized_layers)

vllm/model_executor/model_loader/weight_utils.py

+13 -3

@@ -287,7 +287,17 @@ def get_quant_config(

287 )

288

289 if hf_quant_config is not None:

290 - return quant_cls.from_config(hf_quant_config)

290 + # For modelopt_mixed, config.json's quantization_config may or may

291 + # not contain the per-layer quantized_layers map. Newer checkpoints

292 + # embed it directly; older ones keep it only in hf_quant_config.json.

293 + # If it is missing, fall through to the file-based loading path.

294 + if (

295 + model_config.quantization == "modelopt_mixed"

296 + and "quantized_layers" not in hf_quant_config

297 + ):

298 + pass # fall through to file-based loading below

299 + else:

300 + return quant_cls.from_config(hf_quant_config)

301

302 # if hf_quant_config is None, we will try to get config from

303 # hf_overrides

@@ -365,8 +375,8 @@ def get_quant_config(

375

376 if model_config.quantization == "bitsandbytes":

377 config["adapter_name_or_path"] = model_config.model

368 - elif model_config.quantization == "modelopt":

369 - if config["producer"]["name"] == "modelopt":

378 + elif model_config.quantization in ("modelopt", "modelopt_mixed"):

379 + if config.get("producer", {}).get("name") == "modelopt":

380 return quant_cls.from_config(config)

381 else:

382 raise ValueError(