@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 7a2b596982
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 4, 00:00 UTC

[Quantization] Add ModelOpt NVFP4 W4A16 (4-bit weights, fp16/bf16 activations) support (#41769)

Signed-off-by: Juhi Mittal <juhim@nvidia.com>

Juhi Mittal committed May 9, 2026 at 14:15 UTC 7a2b59698275f30d6168dbac1f2b3c5c3cd5fe94
3 files changed +224 -6
tests/quantization/test_modelopt.py
+46
@@ -239,3 +239,49 @@ def test_modelopt_fp8_pb_wo_checkpoint_setup(default_vllm_config, vllm_runner):
239 output = llm.generate_greedy(["Hello my name is"], max_tokens=4)
240 assert output
241 print(f"ModelOpt FP8_PB_WO output: {output}")
242 +
243 +
244 +def test_modelopt_nvfp4_config_dispatches_w4a4_method():
245 + """``quant_method="NVFP4"`` (W4A4 default) routes to the existing
246 + ``ModelOptNvFp4LinearMethod``."""
247 + from vllm.model_executor.layers.quantization.modelopt import (
248 + ModelOptNvFp4Config,
249 + ModelOptNvFp4LinearMethod,
250 + )
251 +
252 + config = ModelOptNvFp4Config(
253 + quant_method="NVFP4",
254 + is_checkpoint_nvfp4_serialized=True,
255 + kv_cache_quant_algo=None,
256 + exclude_modules=[],
257 + )
258 + assert config.LinearMethodCls is ModelOptNvFp4LinearMethod
259 + assert config.quant_method == "NVFP4"
260 +
261 +
262 +def test_modelopt_nvfp4_config_dispatches_w4a16_method():
263 + """``quant_method="W4A16_NVFP4"`` routes to the new
264 + ``ModelOptNvFp4W4A16LinearMethod`` instead of the W4A4 sibling.
265 +
266 + Mirrors the FP8 dispatch precedent (``ModelOptFp8Config`` selects
267 + one of three FP8 LinearMethods on ``quant_method``); a regression
268 + here would mean a W4A16 NVFP4 checkpoint silently loaded under the
269 + W4A4 method, which would try to register an ``input_scale`` runtime
270 + parameter and (more importantly) call the cutlass W4A4 NVFP4 GEMM
271 + instead of FP4 Marlin.
272 + """
273 + from vllm.model_executor.layers.quantization.modelopt import (
274 + ModelOptNvFp4Config,
275 + ModelOptNvFp4LinearMethod,
276 + ModelOptNvFp4W4A16LinearMethod,
277 + )
278 +
279 + config = ModelOptNvFp4Config(
280 + quant_method="W4A16_NVFP4",
281 + is_checkpoint_nvfp4_serialized=True,
282 + kv_cache_quant_algo=None,
283 + exclude_modules=[],
284 + )
285 + assert config.LinearMethodCls is ModelOptNvFp4W4A16LinearMethod
286 + assert config.LinearMethodCls is not ModelOptNvFp4LinearMethod
287 + assert config.quant_method == "W4A16_NVFP4"
vllm/model_executor/layers/linear.py
+1
@@ -60,6 +60,7 @@ WEIGHT_LOADER_V2_SUPPORTED = [
60 "ModelOptFp8PbWoLinearMethod",
61 "QuarkLinearMethod",
62 "ModelOptNvFp4LinearMethod",
63 + "ModelOptNvFp4W4A16LinearMethod",
64 "HummingLinearMethod",
65 ]
66
vllm/model_executor/layers/quantization/modelopt.py
+177 -6
@@ -11,6 +11,8 @@ import vllm.model_executor.layers.fused_moe.modular_kernel as mk
11 from vllm.config import get_current_vllm_config
12 from vllm.logger import init_logger
13 from vllm.model_executor.kernels.linear import (
14 + MarlinNvFp4LinearKernel,
15 + NvFp4LinearLayerConfig,
16 init_fp8_linear_kernel,
17 init_mxfp8_linear_kernel,
18 init_nvfp4_linear_kernel,
@@ -89,6 +91,7 @@ from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
91 from vllm.model_executor.parameter import (
92 BlockQuantScaleParameter,
93 ChannelQuantScaleParameter,
94 + GroupQuantScaleParameter,
95 ModelWeightParameter,
96 PerTensorScaleParameter,
97 )
@@ -107,8 +110,10 @@ QUANT_ALGOS = [
110 "FP8_PER_CHANNEL_PER_TOKEN",
111 # FP8 per-block weight-only (ModelOpt may emit this as lowercase).
112 "FP8_PB_WO",
110 - # FP4
113 + # NVFP4 W4A4 (4-bit float weights AND 4-bit float activations).
114 "NVFP4",
115 + # W4A16 NVFP4 (4-bit float weights, fp16/bf16 activations).
116 + "W4A16_NVFP4",
117 # MXFP8
118 "MXFP8",
119 # MIXED_PRECISION,
@@ -1003,22 +1008,41 @@ class ModelOptNvFp4Config(ModelOptQuantConfigBase):
1008
1009 def __init__(
1010 self,
1006 - is_checkpoint_nvfp4_serialized: bool,
1007 - kv_cache_quant_algo: str | None,
1008 - exclude_modules: list[str],
1011 + quant_method: str = "NVFP4",
1012 + is_checkpoint_nvfp4_serialized: bool = False,
1013 + kv_cache_quant_algo: str | None = None,
1014 + exclude_modules: list[str] | None = None,
1015 group_size: int = 16,
1016 ) -> None:
1017 + if exclude_modules is None:
1018 + exclude_modules = []
1019 super().__init__(exclude_modules)
1020 + self.quant_method = quant_method
1021 self.is_checkpoint_nvfp4_serialized = is_checkpoint_nvfp4_serialized
1022 if is_checkpoint_nvfp4_serialized:
1023 logger.warning(
1015 - "Detected ModelOpt NVFP4 checkpoint. Please note that"
1016 - " the format is experimental and could change in future."
1024 + "Detected ModelOpt NVFP4 checkpoint (quant_algo=%s). Please "
1025 + "note that the format is experimental and could change in "
1026 + "future.",
1027 + quant_method,
1028 )
1029
1030 self.group_size = group_size
1031 self.kv_cache_quant_algo = kv_cache_quant_algo
1032
1033 + # Select LinearMethod implementation based on quant_algo (FP8 pattern).
1034 + # NVFP4 -> W4A4: cutlass NVFP4 GEMM with input quantization
1035 + # W4A16_NVFP4 -> W4A16: FP4 Marlin GEMM with bf16/fp16 activations
1036 + if quant_method == "NVFP4":
1037 + self.LinearMethodCls = ModelOptNvFp4LinearMethod
1038 + elif quant_method == "W4A16_NVFP4":
1039 + self.LinearMethodCls = ModelOptNvFp4W4A16LinearMethod
1040 + else:
1041 + raise ValueError(
1042 + f"Unsupported ModelOpt NVFP4 quant_algo: {quant_method}. "
1043 + "Supported: NVFP4 / W4A16_NVFP4."
1044 + )
1045 +
1046 def get_name(self) -> QuantizationMethods:
1047 return "modelopt_fp4"
1048
@@ -1069,6 +1093,7 @@ class ModelOptNvFp4Config(ModelOptQuantConfigBase):
1093 )
1094
1095 return cls(
1096 + quant_method,
1097 is_checkpoint_nvfp4_serialized,
1098 kv_cache_quant_method,
1099 exclude_modules,
@@ -1208,6 +1233,152 @@ class ModelOptNvFp4LinearMethod(LinearMethodBase):
1233 return self.kernel.apply_weights(layer=layer, x=x, bias=bias)
1234
1235
1236 +class ModelOptNvFp4W4A16LinearMethod(LinearMethodBase):
1237 + """Linear method for ModelOpt NVFP4 W4A16.
1238 +
1239 + 4-bit NVFP4 weights, fp16/bf16 activations. Loads ModelOpt-style names
1240 + directly (no on-disk conversion) and dispatches to the FP4 Marlin GEMM:
1241 +
1242 + weight uint8 packed NVFP4 (2 nibbles/byte along input dim)
1243 + weight_scale fp8-e4m3 per 16-elem group along input dim
1244 + weight_scale_2 fp32 per-tensor global scale = amax / (6.0 * 448.0)
1245 +
1246 + No activation quantization. Marlin expects the global scale in the same
1247 + form ModelOpt stores (amax/2688), so we rename weight_scale_2 ->
1248 + weight_global_scale **without reciprocation** -- the CT W4A16 path
1249 + reciprocates only because CT stores the inverse on disk.
1250 +
1251 + We also register a placeholder input_scale parameter so that W4A4-shaped
1252 + checkpoints (which contain *_proj.input_scale tensors) can be loaded
1253 + under this method without the per-shard loader hitting a KeyError on
1254 + the merged-name lookup. The placeholder is discarded in
1255 + process_weights_after_loading -- its value is never used.
1256 + """
1257 +
1258 + def __init__(self, quant_config: ModelOptNvFp4Config) -> None:
1259 + self.quant_config = quant_config
1260 + # Vestigial slot mirrored from ModelOptNvFp4LinearMethod: the parent
1261 + # config's get_quant_method only fills marlin_input_dtype when
1262 + # backend == "marlin"; we don't set that since we pin the kernel
1263 + # below, but we keep the attribute for shape parity.
1264 + self.marlin_input_dtype = None
1265 + # Direct-instantiate the Marlin NVFP4 adapter rather than going through
1266 + # init_nvfp4_linear_kernel(): the latter's priority list returns a
1267 + # cutlass W4A4 kernel as first-pick on this hardware, which would
1268 + # silently try to quantize activations (we have no input_scale). For
1269 + # W4A16 there is exactly one valid kernel, so we pin it.
1270 + self.kernel = MarlinNvFp4LinearKernel(NvFp4LinearLayerConfig())
1271 +
1272 + def create_weights(
1273 + self,
1274 + layer: torch.nn.Module,
1275 + input_size_per_partition: int,
1276 + output_partition_sizes: list[int],
1277 + input_size: int,
1278 + output_size: int,
1279 + params_dtype: torch.dtype,
1280 + **extra_weight_attrs,
1281 + ):
1282 + del input_size, output_size
1283 + if not self.quant_config.is_checkpoint_nvfp4_serialized:
1284 + raise ValueError(
1285 + "W4A16_NVFP4 quantization was selected; "
1286 + "dynamic quantization is not supported."
1287 + )
1288 + output_size_per_partition = sum(output_partition_sizes)
1289 + weight_loader = extra_weight_attrs.get("weight_loader")
1290 + layer.logical_widths = output_partition_sizes
1291 + layer.input_size_per_partition = input_size_per_partition
1292 + layer.output_size_per_partition = output_size_per_partition
1293 +
1294 + if input_size_per_partition % 16 != 0:
1295 + raise ValueError(
1296 + "Unsupported model: input feature size is not a multiple of 16."
1297 + )
1298 +
1299 + # Packed NVFP4 weights: uint8, 2 nibbles per byte along the input dim.
1300 + weight = ModelWeightParameter(
1301 + data=torch.empty(
1302 + output_size_per_partition,
1303 + input_size_per_partition // 2,
1304 + dtype=torch.uint8,
1305 + ),
1306 + input_dim=1,
1307 + output_dim=0,
1308 + weight_loader=weight_loader,
1309 + )
1310 + layer.register_parameter("weight", weight)
1311 +
1312 + # Per-tensor global weight scale (fp32). ModelOpt stores
1313 + # amax / (NVFP4_max * fp8_e4m3_max) = amax / 2688. PerTensorScaleParameter
1314 + # holds one entry per fused output partition (e.g. q/k/v in a fused QKV).
1315 + weight_scale_2 = PerTensorScaleParameter(
1316 + data=torch.empty(len(output_partition_sizes), dtype=torch.float32),
1317 + weight_loader=weight_loader,
1318 + )
1319 + layer.register_parameter("weight_scale_2", weight_scale_2)
1320 +
1321 + # Per-group fp8 weight scale.
1322 + weight_scale = GroupQuantScaleParameter(
1323 + data=torch.empty(
1324 + output_size_per_partition,
1325 + input_size_per_partition // self.quant_config.group_size,
1326 + dtype=torch.float8_e4m3fn,
1327 + ),
1328 + input_dim=1,
1329 + output_dim=0,
1330 + weight_loader=weight_loader,
1331 + )
1332 + layer.register_parameter("weight_scale", weight_scale)
1333 +
1334 + # Placeholder input_scale param so W4A4-shaped checkpoints can be
1335 + # loaded under this method without KeyError on the merged-name
1336 + # lookup (qwen2-style stacked-loader path renames *_proj.input_scale
1337 + # to e.g. qkv_proj.input_scale and looks it up unconditionally).
1338 + # Discarded in process_weights_after_loading; never read by the kernel.
1339 + # For native W4A16 checkpoints (no input_scale on disk) the param
1340 + # stays uninitialized and is simply deleted.
1341 + input_scale = PerTensorScaleParameter(
1342 + data=torch.empty(len(output_partition_sizes), dtype=torch.float32),
1343 + weight_loader=weight_loader,
1344 + )
1345 + layer.register_parameter("input_scale", input_scale)
1346 +
1347 + def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
1348 + # Discard the input_scale placeholder. Whether it carries values
1349 + # (W4A4 ckpt loaded as W4A16) or is uninitialized (native W4A16
1350 + # ckpt), W4A16 mode does not quantize activations, so this is unused.
1351 + if hasattr(layer, "input_scale"):
1352 + del layer.input_scale
1353 +
1354 + if torch.unique(layer.weight_scale_2).numel() != 1:
1355 + logger.warning_once(
1356 + "In W4A16_NVFP4 linear, the global weight scale "
1357 + "(weight_scale_2) differs across fused parallel layers "
1358 + "(e.g. q/k/v_proj). This will likely reduce accuracy. "
1359 + "Consider a checkpoint with a shared global scale."
1360 + )
1361 +
1362 + # Rename weight_scale_2 -> weight_global_scale. NO reciprocation:
1363 + # ModelOpt already stores amax/2688, which is exactly what Marlin
1364 + # consumes via nvfp4_marlin_process_global_scale (called inside the
1365 + # Marlin adapter's process_weights_after_loading).
1366 + layer.weight_global_scale = Parameter(
1367 + layer.weight_scale_2.max().to(torch.float32), requires_grad=False
1368 + )
1369 + del layer.weight_scale_2
1370 +
1371 + self.kernel.process_weights_after_loading(layer)
1372 +
1373 + def apply(
1374 + self,
1375 + layer: torch.nn.Module,
1376 + x: torch.Tensor,
1377 + bias: torch.Tensor | None = None,
1378 + ) -> torch.Tensor:
1379 + return self.kernel.apply_weights(layer=layer, x=x, bias=bias)
1380 +
1381 +
1382 class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
1383 """
1384 MoE Method for FP4 Quantization.