a8c536829c · tduncan/Nvidia-Optimization-Stack-VLLM

@@ -2,6 +2,7 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 from collections.abc import Callable

5 +from fnmatch import fnmatch

6 from typing import TYPE_CHECKING, Any, Optional

7

8 import torch

@@ -13,7 +14,6 @@ import vllm.model_executor.layers.fused_moe.modular_kernel as mk

14 from vllm._custom_ops import cutlass_scaled_fp4_mm, scaled_fp4_quant

15 from vllm.logger import init_logger

16 from vllm.model_executor.layers.fused_moe.config import (

16 - FusedMoEConfig,

17 FusedMoEQuantConfig,

18 RoutingMethodType,

19 fp8_w8a8_moe_quant_config,

@@ -86,45 +86,218 @@ QUANT_ALGOS = ["FP8", "NVFP4"]

86 KV_CACHE_QUANT_ALGOS = ["FP8"]

87

88

89 -class ModelOptFp8Config(QuantizationConfig):

89 +class ModelOptFp8KVCacheMethod(BaseKVCacheMethod):

90 + """

91 + Supports loading kv-cache scaling factors from FP8 checkpoints.

92 + """

93 +

94 + def __init__(self, quant_config: "ModelOptQuantConfigBase"):

95 + super().__init__(quant_config)

96 +

97 +

98 +class ModelOptQuantConfigBase(QuantizationConfig):

99 + LinearMethodCls: type = LinearMethodBase

100 + FusedMoEMethodCls: type = FusedMoEMethodBase

101 + KVCacheMethodCls: type = BaseKVCacheMethod

102 +

103 + def __init__(

104 + self,

105 + exclude_modules: list[str],

106 + ):

107 + super().__init__()

108 + self.exclude_modules: list[str] = exclude_modules

109 +

110 + def is_layer_excluded(self, prefix: str) -> bool:

111 + """

112 + Check if a layer should be excluded from quantization.

113 +

114 + Handles both exact matching (for fused layers) and ModelOpt wildcard matching.

115 +

116 + The ModelOpt exclude_modules list is a list of wildcards.

117 + """

118 + if len(self.exclude_modules) == 0:

119 + return False

120 +

121 + # First check exact matching with fused layer support

122 + if is_layer_skipped(prefix, self.exclude_modules, self.packed_modules_mapping):

123 + return True

124 +

125 + # TODO: This special hard coded logic is not needed for quantized checkpoints

126 + # generated by ModelOpt >= 0.39.0 where they are handled natually by the

127 + # exclude_modules config. But need to keep them for loading quantized

128 + # checkpoints generated by older versions. Then check substring matching

129 + # for patterns not caught by exact match

130 + for exclude_module in self.exclude_modules:

131 + # Skip exact matches already handled above

132 + if exclude_module != prefix and (

133 + exclude_module in prefix

134 + or (

135 + prefix.startswith("language_model.")

136 + and exclude_module in prefix.removeprefix("language_model.")

137 + )

138 + ):

139 + return True

140 +

141 + # modelopt exclude modules are not simple strings, they are wildcards

142 + for wildcard_pattern in self.exclude_modules:

143 + if fnmatch(prefix, wildcard_pattern):

144 + return True

145 +

146 + return False

147 +

148 + def get_quant_method(

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

150 + ) -> Optional["QuantizeMethodBase"]:

151 + from vllm.attention.layer import Attention # Avoid circular import

152 +

153 + # handle kv-cache first so we can focus only on weight quantization thereafter

154 + if isinstance(layer, Attention):

155 + return self.KVCacheMethodCls(self)

156 +

157 + # handle exclusion

158 + if self.is_layer_excluded(prefix):

159 + if isinstance(layer, LinearBase):

160 + return UnquantizedLinearMethod()

161 + return None

162 +

163 + # TODO: This special hard coded logic is not needed for quantized checkpoints

164 + # generated by ModelOpt >= 0.39.0 where they are handled natually by the

165 + # exclude_modules config. But need to keep them for loading quantized

166 + # checkpoints generated by older versions. Then check substring matching

167 + # for patterns not caught by exact match

168 + if "vision_tower" in prefix or "vision_model" in prefix:

169 + return UnquantizedLinearMethod()

170 +

171 + # now, the layer is quantized, handle it here

172 + if isinstance(layer, LinearBase):

173 + return self.LinearMethodCls(self)

174 + elif isinstance(layer, FusedMoE):

175 + return self.FusedMoEMethodCls(quant_config=self, layer=layer)

176 +

177 + return None

178 +

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

180 + if len(self.exclude_modules) > 0:

181 + self.exclude_modules = hf_to_vllm_mapper.apply_list(self.exclude_modules)

182 +

183 + @staticmethod

184 + def get_config_filenames() -> list[str]:

185 + return ["hf_quant_config.json"]

186 +

187 + @classmethod

188 + def _from_config(

189 + cls,

190 + *,

191 + quant_method: str,

192 + kv_cache_quant_method: str | None,

193 + exclude_modules: list[str],

194 + original_config: dict[str, Any],

195 + group_size: int | None,

196 + ) -> "ModelOptQuantConfigBase":

197 + raise NotImplementedError("Please implement this function in sub classes")

198 +

199 + @classmethod

200 + def from_config(cls, config: dict[str, Any]) -> "ModelOptQuantConfigBase":

201 + # Handle both ModelOpt format and compressed-tensors style format

202 + if "quantization" in config:

203 + # Traditional ModelOpt format:

204 + # {"quantization": {"quant_algo": "..."}}

205 + quant_config = cls.get_from_keys(config, ["quantization"])

206 + if not isinstance(quant_config, dict):

207 + raise ValueError("Expected 'quantization' to be a dictionary in config")

208 +

209 + quant_method = quant_config.get("quant_algo")

210 +

211 + # Handle kv_cache_quant_algo with proper type validation

212 + kv_cache_quant_method = quant_config.get("kv_cache_quant_algo")

213 +

214 + # Handle group_size with proper type validation

215 + group_size_raw = quant_config.get("group_size")

216 +

217 + # "exclude_modules" is the key in the legacy hf_quant_config.json

218 + exclude_modules = quant_config.get("exclude_modules", [])

219 + else:

220 + # Compressed-tensors style format:

221 + # {"quant_algo": "...", "quant_method": "modelopt"}

222 + quant_method = config.get("quant_algo")

223 + kv_cache_quant_method = config.get("kv_cache_quant_algo")

224 + # "ignore" is the key in config.json

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

226 + group_size_raw = config.get("group_size")

227 +

228 + if not quant_method:

229 + raise ValueError("Missing 'quant_algo' in quantization config")

230 +

231 + if kv_cache_quant_method is None:

232 + # No KV cache quantization, keep this branch just to have this comment

233 + pass

234 + elif not isinstance(kv_cache_quant_method, str):

235 + raise ValueError(

236 + f"kv_cache_quant_algo must be a string, got "

237 + f"{type(kv_cache_quant_method)}"

238 + )

239 +

240 + if not isinstance(exclude_modules, list):

241 + raise ValueError(

242 + f"exclude_modules must be a list, got {type(exclude_modules)}"

243 + )

244 +

245 + if group_size_raw is None:

246 + group_size = None

247 + elif isinstance(group_size_raw, int):

248 + group_size = group_size_raw

249 + else:

250 + try:

251 + group_size = int(group_size_raw)

252 + except (ValueError, TypeError):

253 + raise ValueError(

254 + f"group_size must be an integer, got {type(group_size_raw)}"

255 + ) from None

256 +

257 + if quant_method not in QUANT_ALGOS:

258 + raise ValueError(

259 + f"ModelOpt currently only supports: {QUANT_ALGOS} "

260 + "quantizations in vLLM. Please check the "

261 + "`hf_quant_config.json` file for your model's "

262 + "quant configuration."

263 + )

264 + return cls._from_config(

265 + quant_method=quant_method,

266 + kv_cache_quant_method=kv_cache_quant_method,

267 + exclude_modules=exclude_modules,

268 + group_size=group_size,

269 + original_config=config,

270 + )

271 +

272 +

273 +class ModelOptFp8Config(ModelOptQuantConfigBase):

274 """Config class for ModelOpt FP8."""

275

276 def __init__(

277 self,

94 - is_checkpoint_fp8_serialized: bool = False,

95 - kv_cache_quant_method: str | None = None,

96 - exclude_modules: list[str] | None = None,

278 + is_checkpoint_fp8_serialized: bool,

279 + kv_cache_quant_method: str | None,

280 + exclude_modules: list[str],

281 ) -> None:

98 - super().__init__()

282 + super().__init__(exclude_modules)

283 self.is_checkpoint_fp8_serialized = is_checkpoint_fp8_serialized

284 self.kv_cache_quant_method = kv_cache_quant_method

101 - self.exclude_modules = exclude_modules or []

285 if is_checkpoint_fp8_serialized:

286 logger.warning(

287 "Detected ModelOpt fp8 checkpoint. Please note that"

288 " the format is experimental and could change."

289 )

290

108 - @classmethod

109 - def get_name(cls) -> QuantizationMethods:

291 + def get_name(self) -> QuantizationMethods:

292 return "modelopt"

293

112 - @classmethod

113 - def get_supported_act_dtypes(cls) -> list[torch.dtype]:

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

295 return [torch.bfloat16, torch.half]

296

297 @classmethod

298 def get_min_capability(cls) -> int:

299 return 89

300

120 - @classmethod

121 - def get_config_filenames(cls) -> list[str]:

122 - return ["hf_quant_config.json"]

123 -

124 - def apply_vllm_mapper(self, hf_to_vllm_mapper: "WeightsMapper"):

125 - if self.exclude_modules is not None:

126 - self.exclude_modules = hf_to_vllm_mapper.apply_list(self.exclude_modules)

127 -

301 @classmethod

302 def override_quantization_method(

303 cls, hf_quant_cfg, user_quant

@@ -158,88 +331,19 @@ class ModelOptFp8Config(QuantizationConfig):

331 return None

332

333 @classmethod

161 - def from_config(cls, config: dict[str, Any]) -> "ModelOptFp8Config":

162 - # Handle both ModelOpt format and compressed-tensors style format

163 - if "quantization" in config:

164 - # ModelOpt format: {"quantization": {"quant_algo": "..."}}

165 - quant_config = cls.get_from_keys(config, ["quantization"])

166 - if not isinstance(quant_config, dict):

167 - raise ValueError("Expected 'quantization' to be a dictionary in config")

168 - quant_method = quant_config.get("quant_algo", "")

169 - if not quant_method:

170 - raise ValueError("Missing 'quant_algo' in quantization config")

171 - kv_cache_quant_method = quant_config.get("kv_cache_quant_algo")

172 - # "exclude_modules" is the key in the legacy hf_quant_config.json

173 - exclude_modules = quant_config.get("exclude_modules")

174 - else:

175 - # Compressed-tensors style format:

176 - # {"quant_algo": "...", "quant_method": "modelopt"}

177 - quant_method = config.get("quant_algo", "")

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

179 - # "ignore" is the key in config.json

180 - exclude_modules = config.get("ignore")

181 -

182 - if quant_method not in QUANT_ALGOS:

183 - raise ValueError(

184 - f"ModelOpt currently only supports: {QUANT_ALGOS} "

185 - "quantizations in vLLM. Please check the "

186 - "`hf_quant_config.json` file for your model's "

187 - "quant configuration."

188 - )

334 + def _from_config(

335 + cls,

336 + *,

337 + quant_method: str,

338 + kv_cache_quant_method: str | None,

339 + exclude_modules: list[str],

340 + original_config: dict[str, Any],

341 + **kwargs: Any,

342 + ) -> "ModelOptFp8Config":

343 is_checkpoint_fp8_serialized = "FP8" in quant_method

344

345 return cls(is_checkpoint_fp8_serialized, kv_cache_quant_method, exclude_modules)

346

193 - def is_layer_excluded(self, prefix: str) -> bool:

194 - """

195 - Check if a layer should be excluded from quantization.

196 - Handles both exact matching (for fused layers) and substring matching.

197 -

198 - This method handles both regular models and multimodal models that use

199 - the language_model prefix. For multimodal models, it checks if the

200 - module name (without the language_model prefix) is in the exclude list.

201 - """

202 - if self.exclude_modules is None:

203 - return False

204 -

205 - # First check exact matching with fused layer support

206 - if is_layer_skipped(prefix, self.exclude_modules, self.packed_modules_mapping):

207 - return True

208 -

209 - # Then check substring matching for patterns not caught by exact match

210 - for module in self.exclude_modules:

211 - # Skip exact matches already handled above

212 - if module != prefix and (

213 - module in prefix

214 - or (

215 - prefix.startswith("language_model.")

216 - and module in prefix.removeprefix("language_model.")

217 - )

218 - ):

219 - return True

220 - return False

221 -

222 - def get_quant_method(

223 - self, layer: torch.nn.Module, prefix: str

224 - ) -> Optional["QuantizeMethodBase"]:

225 - from vllm.attention.layer import ( # Avoid circular import

226 - Attention,

227 - MLAAttention,

228 - )

229 -

230 - if isinstance(layer, LinearBase):

231 - if self.is_layer_excluded(prefix):

232 - return UnquantizedLinearMethod()

233 - # Check if this is a vision model layer that should not be quantized

234 - if "vision_tower" in prefix or "vision_model" in prefix:

235 - return UnquantizedLinearMethod()

236 - return ModelOptFp8LinearMethod(self)

237 - elif isinstance(layer, (Attention, MLAAttention)):

238 - return ModelOptFp8KVCacheMethod(self)

239 - elif isinstance(layer, FusedMoE):

240 - return ModelOptFp8MoEMethod(self, layer)

241 - return None

242 -

347

348 class ModelOptFp8LinearMethod(LinearMethodBase):

349 """Linear method for Model Optimizer static quantization.

@@ -344,7 +448,7 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):

448 def __init__(

449 self,

450 quant_config: ModelOptFp8Config,

347 - layer: torch.nn.Module,

451 + layer: FusedMoE,

452 ) -> None:

453 super().__init__(layer.moe_config)

454 self.layer = layer

@@ -686,7 +790,12 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):

790 )

791

792

689 -class ModelOptNvFp4Config(QuantizationConfig):

793 +ModelOptFp8Config.LinearMethodCls = ModelOptFp8LinearMethod

794 +ModelOptFp8Config.FusedMoEMethodCls = ModelOptFp8MoEMethod

795 +ModelOptFp8Config.KVCacheMethodCls = ModelOptFp8KVCacheMethod

796 +

797 +

798 +class ModelOptNvFp4Config(ModelOptQuantConfigBase):

799 """Config class for ModelOpt FP4."""

800

801 def __init__(

@@ -696,7 +805,7 @@ class ModelOptNvFp4Config(QuantizationConfig):

805 exclude_modules: list[str],

806 group_size: int = 16,

807 ) -> None:

699 - super().__init__()

808 + super().__init__(exclude_modules)

809 self.is_checkpoint_nvfp4_serialized = is_checkpoint_nvfp4_serialized

810 if is_checkpoint_nvfp4_serialized:

811 logger.warning(

@@ -706,28 +815,17 @@ class ModelOptNvFp4Config(QuantizationConfig):

815

816 self.group_size = group_size

817 self.kv_cache_quant_algo = kv_cache_quant_algo

709 - self.exclude_modules = exclude_modules

818

711 - @classmethod

712 - def get_name(cls) -> QuantizationMethods:

819 + def get_name(self) -> QuantizationMethods:

820 return "modelopt_fp4"

821

715 - @classmethod

716 - def get_supported_act_dtypes(cls) -> list[torch.dtype]:

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

823 return [torch.bfloat16, torch.half, torch.float8_e4m3fn]

824

825 @classmethod

826 def get_min_capability(cls) -> int:

827 return 80

828

723 - @classmethod

724 - def get_config_filenames(cls) -> list[str]:

725 - return ["hf_quant_config.json"]

726 -

727 - def apply_vllm_mapper(self, hf_to_vllm_mapper: "WeightsMapper"):

728 - if self.exclude_modules is not None:

729 - self.exclude_modules = hf_to_vllm_mapper.apply_list(self.exclude_modules)

730 -

829 @classmethod

830 def override_quantization_method(

831 cls, hf_quant_cfg, user_quant

@@ -761,105 +859,25 @@ class ModelOptNvFp4Config(QuantizationConfig):

859 return None

860

861 @classmethod

764 - def from_config(cls, config: dict[str, Any]) -> "ModelOptNvFp4Config":

765 - # Handle both traditional ModelOpt format and compressed-tensors

766 - # style format

767 - if "quantization" in config:

768 - # Traditional ModelOpt format:

769 - # {"quantization": {"quant_algo": "..."}}

770 - quant_config = cls.get_from_keys(config, ["quantization"])

771 - if not isinstance(quant_config, dict):

772 - raise ValueError("Expected 'quantization' to be a dictionary in config")

773 -

774 - quant_method = quant_config.get("quant_algo", "")

775 - if not quant_method:

776 - raise ValueError("Missing 'quant_algo' in quantization config")

777 -

778 - # Handle kv_cache_quant_algo with proper type validation

779 - kv_cache_quant_algo_raw = quant_config.get("kv_cache_quant_algo")

780 - if kv_cache_quant_algo_raw is None:

781 - # No KV cache quantization by default

782 - kv_cache_quant_algo = None

783 - elif isinstance(kv_cache_quant_algo_raw, str):

784 - kv_cache_quant_algo = kv_cache_quant_algo_raw

785 - else:

786 - raise ValueError(

787 - f"kv_cache_quant_algo must be a string, got "

788 - f"{type(kv_cache_quant_algo_raw)}"

789 - )

790 -

791 - # Handle group_size with proper type validation

792 - group_size_raw = quant_config.get("group_size")

793 - if group_size_raw is None:

794 - group_size = 16 # Default value

795 - elif isinstance(group_size_raw, int):

796 - group_size = group_size_raw

797 - else:

798 - try:

799 - group_size = int(group_size_raw)

800 - except (ValueError, TypeError):

801 - raise ValueError(

802 - f"group_size must be an integer, got {type(group_size_raw)}"

803 - ) from None

804 -

805 - # "exclude_modules" is the key in the legacy hf_quant_config.json

806 - exclude_modules = quant_config.get("exclude_modules", [])

807 - if not isinstance(exclude_modules, list):

808 - raise ValueError(

809 - f"exclude_modules must be a list, got {type(exclude_modules)}"

810 - )

811 - else:

812 - # Compressed-tensors style format:

813 - # {"quant_algo": "...", "quant_method": "modelopt"}

814 - quant_method = config.get("quant_algo", "")

815 -

816 - # Handle kv_cache_quant_algo with proper type validation

817 - kv_cache_quant_algo_raw = config.get("kv_cache_quant_algo")

818 - if kv_cache_quant_algo_raw is None:

819 - # No KV cache quantization by default

820 - kv_cache_quant_algo = None

821 - elif isinstance(kv_cache_quant_algo_raw, str):

822 - kv_cache_quant_algo = kv_cache_quant_algo_raw

823 - else:

824 - raise ValueError(

825 - f"kv_cache_quant_algo must be a string, got "

826 - f"{type(kv_cache_quant_algo_raw)}"

827 - )

828 -

829 - # Handle group_size with proper type validation

830 - group_size_raw = config.get("group_size")

831 - if group_size_raw is None:

832 - group_size = 16 # Default value

833 - elif isinstance(group_size_raw, int):

834 - group_size = group_size_raw

835 - else:

836 - try:

837 - group_size = int(group_size_raw)

838 - except (ValueError, TypeError):

839 - raise ValueError(

840 - f"group_size must be an integer, got {type(group_size_raw)}"

841 - ) from None

842 -

843 - # "ignore" is the key in config.json

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

845 - if not isinstance(exclude_modules, list):

846 - raise ValueError(

847 - f"exclude_modules must be a list, got {type(exclude_modules)}"

848 - )

849 -

850 - if quant_method not in QUANT_ALGOS:

851 - raise ValueError(

852 - f"ModelOpt currently only supports: {QUANT_ALGOS} "

853 - "quantizations in vLLM. Please check the "

854 - "`hf_quant_config.json` file for your model's "

855 - "quant configuration."

856 - )

862 + def _from_config(

863 + cls,

864 + *,

865 + quant_method: str,

866 + kv_cache_quant_method: str | None,

867 + exclude_modules: list[str],

868 + original_config: dict[str, Any],

869 + group_size: int | None,

870 + **kwargs: Any,

871 + ) -> "ModelOptNvFp4Config":

872 is_checkpoint_nvfp4_serialized = "NVFP4" in quant_method

873

874 + if group_size is None:

875 + group_size = 16 # Default value

876 +

877 # For FP4, these fields are required

860 - if is_checkpoint_nvfp4_serialized and "quantization" in config:

878 + if is_checkpoint_nvfp4_serialized and "quantization" in original_config:

879 # Check if required fields are present in the quantization config

862 - quant_config = config["quantization"]

880 + quant_config = original_config["quantization"]

881 required_fields = ["group_size", "kv_cache_quant_algo", "exclude_modules"]

882 missing_fields = [

883 field for field in required_fields if field not in quant_config

@@ -872,64 +890,11 @@ class ModelOptNvFp4Config(QuantizationConfig):

890

891 return cls(

892 is_checkpoint_nvfp4_serialized,

875 - kv_cache_quant_algo,

893 + kv_cache_quant_method,

894 exclude_modules,

895 group_size,

896 )

897

880 - def is_layer_excluded(self, prefix: str) -> bool:

881 - """

882 - Check if a layer should be excluded from quantization.

883 - Handles both exact matching (for fused layers) and pattern matching.

884 - """

885 - # First check exact matching with fused layer support

886 - if is_layer_skipped(prefix, self.exclude_modules, self.packed_modules_mapping):

887 - return True

888 -

889 - # Check regex pattern matching for patterns not caught by exact match

890 - import regex as re

891 -

892 - for pattern in self.exclude_modules:

893 - # Skip patterns that would be caught by exact matching

894 - if "*" in pattern or "." in pattern:

895 - regex_str = pattern.replace(".", r"\.").replace("*", r".*")

896 - if re.fullmatch(regex_str, prefix):

897 - return True

898 - return False

899 -

900 - def get_quant_method(

901 - self, layer: torch.nn.Module, prefix: str

902 - ) -> Optional["QuantizeMethodBase"]:

903 - from vllm.attention.layer import ( # Avoid circular import

904 - Attention,

905 - MLAAttention,

906 - )

907 -

908 - skip_layer = self.is_layer_excluded(prefix)

909 - if isinstance(layer, LinearBase):

910 - if skip_layer:

911 - return UnquantizedLinearMethod()

912 - # Check if this is a vision model layer that should not be quantized

913 - if "vision_tower" in prefix or "vision_model" in prefix:

914 - return UnquantizedLinearMethod()

915 - return ModelOptNvFp4LinearMethod(self)

916 - elif isinstance(layer, (Attention, MLAAttention)):

917 - return ModelOptFp8KVCacheMethod(self)

918 - elif isinstance(layer, FusedMoE):

919 - if skip_layer:

920 - return None

921 - return ModelOptNvFp4FusedMoE(self, layer.moe_config, layer)

922 - return None

923 -

924 -

925 -class ModelOptFp8KVCacheMethod(BaseKVCacheMethod):

926 - """

927 - Supports loading kv-cache scaling factors from FP8 checkpoints.

928 - """

929 -

930 - def __init__(self, quant_config: ModelOptFp8Config | ModelOptNvFp4Config):

931 - super().__init__(quant_config)

932 -

898

899 class ModelOptNvFp4LinearMethod(LinearMethodBase):

900 """Linear method for Model Optimizer NVFP4.

@@ -1157,14 +1122,13 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):

1122 def __init__(

1123 self,

1124 quant_config: ModelOptNvFp4Config,

1160 - moe: FusedMoEConfig,

1161 - layer: torch.nn.Module,

1125 + layer: FusedMoE,

1126 ) -> None:

1127 from vllm.model_executor.layers.quantization.utils.nvfp4_moe_support import (

1128 detect_nvfp4_moe_support, # noqa: E501

1129 )

1130

1167 - super().__init__(moe)

1131 + super().__init__(layer.moe_config)

1132 self.quant_config = quant_config

1133 self.layer = layer

1134 _nvfp4 = detect_nvfp4_moe_support(self.__class__.__name__)

@@ -1802,3 +1766,8 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):

1766 k=x.shape[1],

1767 e=layer.w13_weight.shape[0],

1768 )

1769 +

1770 +

1771 +ModelOptNvFp4Config.LinearMethodCls = ModelOptNvFp4LinearMethod

1772 +ModelOptNvFp4Config.FusedMoEMethodCls = ModelOptNvFp4FusedMoE

1773 +ModelOptNvFp4Config.KVCacheMethodCls = ModelOptFp8KVCacheMethod

Consolidate Nvidia ModelOpt quant config handling for all quantization methods (#28076)