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 8, 16:30 UTC
Add support for ModelOpt MXFP8 dense models (#33786)
Signed-off-by: Daniel Serebrenik <daserebrenik@nvidia.com>
danisereb committed
Feb 8, 2026 at 21:16 UTC
084aa19f02b00198b36bd8d742d4169d6f5a32ce
6 files changed
+375
-14
docs/features/quantization/modelopt.md
+1
@@ -17,6 +17,7 @@ following `quantization.quant_algo` values:
17
- `FP8_PER_CHANNEL_PER_TOKEN`: per-channel weight scale and dynamic per-token activation quantization.
18
- `FP8_PB_WO` (ModelOpt may emit `fp8_pb_wo`): block-scaled FP8 weight-only (typically 128×128 blocks).
19
- `NVFP4`: ModelOpt NVFP4 checkpoints (use `quantization="modelopt_fp4"`).
20
+- `MXFP8`: ModelOpt MXFP8 checkpoints (use `quantization="modelopt_mxfp8"`).
21
22
## Quantizing HuggingFace Models with PTQ
23
vllm/config/model.py
+1
@@ -878,6 +878,7 @@ class ModelConfig:
878
"moe_wna16",
879
"modelopt",
880
"modelopt_fp4",
881
+ "modelopt_mxfp8",
882
"petit_nvfp4",
883
# Ensure heavy backends are probed last to avoid unnecessary
884
# imports during override detection (e.g., MXFP4 imports Triton)
vllm/model_executor/layers/fused_moe/config.py
+2
@@ -494,6 +494,7 @@ class FusedMoEQuantConfig:
494
"mxfp4",
495
"mxfp6_e3m2",
496
"mxfp6_e2m3",
497
+ "mxfp8",
498
}
499
assert not isinstance(weight_dtype, str) or weight_dtype in {
500
"nvfp4",
@@ -501,6 +502,7 @@ class FusedMoEQuantConfig:
502
"mxfp6_e3m2",
503
"mxfp6_e2m3",
504
"int4",
505
+ "mxfp8",
506
}
507
508
if weight_dtype is None:
vllm/model_executor/layers/quantization/__init__.py
+3
-1
@@ -17,6 +17,7 @@ QuantizationMethods = Literal[
17
"fp_quant",
18
"modelopt",
19
"modelopt_fp4",
20
+ "modelopt_mxfp8",
21
"gguf",
22
"gptq_marlin",
23
"awq_marlin",
@@ -119,7 +120,7 @@ def get_quantization_config(quantization: str) -> type[QuantizationConfig]:
120
from .gptq import GPTQConfig
121
from .gptq_marlin import GPTQMarlinConfig
122
from .inc import INCConfig
122
- from .modelopt import ModelOptFp8Config, ModelOptNvFp4Config
123
+ from .modelopt import ModelOptFp8Config, ModelOptMxFp8Config, ModelOptNvFp4Config
124
from .moe_wna16 import MoeWNA16Config
125
from .mxfp4 import Mxfp4Config
126
from .petit import PetitNvFp4Config
@@ -133,6 +134,7 @@ def get_quantization_config(quantization: str) -> type[QuantizationConfig]:
134
"fp_quant": FPQuantConfig,
135
"modelopt": ModelOptFp8Config,
136
"modelopt_fp4": ModelOptNvFp4Config,
137
+ "modelopt_mxfp8": ModelOptMxFp8Config,
138
"gguf": GGUFConfig,
139
"gptq_marlin": GPTQMarlinConfig,
140
"awq_marlin": AWQMarlinConfig,
vllm/model_executor/layers/quantization/modelopt.py
+247
-2
@@ -63,6 +63,13 @@ from vllm.model_executor.layers.quantization.utils.fp8_utils import (
63
from vllm.model_executor.layers.quantization.utils.marlin_utils import (
64
get_marlin_input_dtype,
65
)
66
+from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
67
+ MXFP8_BLOCK_SIZE,
68
+ MXFP8_SCALE_DTYPE,
69
+ MXFP8_VALUE_DTYPE,
70
+ Mxfp8LinearBackend,
71
+ Mxfp8LinearOp,
72
+)
73
from vllm.model_executor.layers.quantization.utils.nvfp4_utils import (
74
apply_nvfp4_linear,
75
convert_to_nvfp4_linear_kernel_format,
@@ -103,6 +110,8 @@ QUANT_ALGOS = [
110
"FP8_PB_WO",
111
# FP4
112
"NVFP4",
113
+ # MXFP8
114
+ "MXFP8",
115
]
116
KV_CACHE_QUANT_ALGOS = ["FP8"]
117
@@ -386,12 +395,12 @@ class ModelOptFp8Config(ModelOptQuantConfigBase):
395
quant_config = hf_quant_cfg["quantization"]
396
if isinstance(quant_config, dict):
397
quant_algo = str(quant_config.get("quant_algo", ""))
389
- if "FP8" in quant_algo.upper():
398
+ if quant_algo.upper() == "FP8":
399
return "modelopt"
400
else:
401
# Check for compressed-tensors style config with specific quant_algo
402
quant_algo = str(hf_quant_cfg.get("quant_algo", ""))
394
- if "FP8" in quant_algo.upper():
403
+ if quant_algo.upper() == "FP8":
404
return "modelopt"
405
406
return None
@@ -1547,3 +1556,239 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
1556
ModelOptNvFp4Config.LinearMethodCls = ModelOptNvFp4LinearMethod
1557
ModelOptNvFp4Config.FusedMoEMethodCls = ModelOptNvFp4FusedMoE
1558
ModelOptNvFp4Config.KVCacheMethodCls = ModelOptFp8KVCacheMethod
1559
+
1560
+
1561
+class ModelOptMxFp8Config(ModelOptQuantConfigBase):
1562
+ """Config class for ModelOpt MXFP8."""
1563
+
1564
+ def __init__(
1565
+ self,
1566
+ is_checkpoint_mxfp8_serialized: bool,
1567
+ kv_cache_quant_algo: str | None,
1568
+ exclude_modules: list[str],
1569
+ ) -> None:
1570
+ super().__init__(exclude_modules)
1571
+ self.is_checkpoint_mxfp8_serialized = is_checkpoint_mxfp8_serialized
1572
+
1573
+ if not is_checkpoint_mxfp8_serialized:
1574
+ raise ValueError(
1575
+ "MXFP8 quantization requires a serialized checkpoint. "
1576
+ "Dynamic quantization is not supported."
1577
+ )
1578
+
1579
+ logger.warning(
1580
+ "Detected ModelOpt MXFP8 checkpoint. Please note that "
1581
+ "the format is experimental and could change in future."
1582
+ )
1583
+
1584
+ self.kv_cache_quant_algo = kv_cache_quant_algo
1585
+
1586
+ def get_name(self) -> QuantizationMethods:
1587
+ return "modelopt_mxfp8"
1588
+
1589
+ def get_supported_act_dtypes(self) -> list[torch.dtype]:
1590
+ return [torch.bfloat16]
1591
+
1592
+ @classmethod
1593
+ def get_min_capability(cls) -> int:
1594
+ # MXFP8 hardware acceleration requires Blackwell (SM100) or newer
1595
+ return 100
1596
+
1597
+ def get_quant_method(
1598
+ self, layer: torch.nn.Module, prefix: str
1599
+ ) -> "QuantizeMethodBase | None":
1600
+ # MXFP8 does not yet support MoE models
1601
+ if isinstance(layer, FusedMoE):
1602
+ raise NotImplementedError(
1603
+ "MXFP8 quantization does not yet support MoE models. "
1604
+ "Please use FP8 or NVFP4 quantization for MoE models."
1605
+ )
1606
+ return super().get_quant_method(layer, prefix)
1607
+
1608
+ @classmethod
1609
+ def override_quantization_method(
1610
+ cls, hf_quant_cfg, user_quant
1611
+ ) -> QuantizationMethods | None:
1612
+ """Detect if this ModelOpt MXFP8 config should be used based on
1613
+ quantization config."""
1614
+ if hf_quant_cfg is None:
1615
+ return None
1616
+
1617
+ # Use the community standard 'quant_method'
1618
+ quant_method = hf_quant_cfg.get("quant_method", "").lower()
1619
+
1620
+ # Only proceed if the method is explicitly "modelopt"
1621
+ if quant_method != "modelopt":
1622
+ return None
1623
+
1624
+ # Look for ModelOpt-specific config structure
1625
+ if "quantization" in hf_quant_cfg:
1626
+ quant_config = hf_quant_cfg["quantization"]
1627
+ if isinstance(quant_config, dict):
1628
+ quant_algo = str(quant_config.get("quant_algo", "")).upper()
1629
+ if "MXFP8" in quant_algo:
1630
+ return "modelopt_mxfp8"
1631
+ else:
1632
+ # Check for compressed-tensors style config with specific quant_algo
1633
+ quant_algo = str(hf_quant_cfg.get("quant_algo", "")).upper()
1634
+ if "MXFP8" in quant_algo:
1635
+ return "modelopt_mxfp8"
1636
+
1637
+ return None
1638
+
1639
+ @classmethod
1640
+ def _from_config(
1641
+ cls,
1642
+ *,
1643
+ quant_method: str,
1644
+ kv_cache_quant_method: str | None,
1645
+ exclude_modules: list[str],
1646
+ original_config: dict[str, Any],
1647
+ **kwargs: Any,
1648
+ ) -> "ModelOptMxFp8Config":
1649
+ is_checkpoint_mxfp8_serialized = "MXFP8" in quant_method.upper()
1650
+
1651
+ # For MXFP8, validate required fields in the config
1652
+ if is_checkpoint_mxfp8_serialized and "quantization" in original_config:
1653
+ quant_config = original_config["quantization"]
1654
+ required_fields = ["kv_cache_quant_algo", "exclude_modules"]
1655
+ missing_fields = [
1656
+ field for field in required_fields if field not in quant_config
1657
+ ]
1658
+ if missing_fields:
1659
+ raise ValueError(
1660
+ f"MXFP8 quantization requires the following fields in "
1661
+ f"hf_quant_config.json: {missing_fields}"
1662
+ )
1663
+
1664
+ return cls(
1665
+ is_checkpoint_mxfp8_serialized,
1666
+ kv_cache_quant_method,
1667
+ exclude_modules,
1668
+ )
1669
+
1670
+
1671
+class ModelOptMxFp8LinearMethod(LinearMethodBase):
1672
+ """Linear method for ModelOpt MXFP8 quantization."""
1673
+
1674
+ def __init__(self, quant_config: ModelOptMxFp8Config) -> None:
1675
+ self.quant_config = quant_config
1676
+
1677
+ if not self.quant_config.is_checkpoint_mxfp8_serialized:
1678
+ raise ValueError(
1679
+ "MXFP8 currently only supports serialized checkpoints. "
1680
+ "Dynamic quantization is not supported."
1681
+ )
1682
+
1683
+ backend: Mxfp8LinearBackend = Mxfp8LinearBackend.EMULATION
1684
+ self.mxfp8_linear_op = Mxfp8LinearOp(backend=backend)
1685
+ logger.info_once("Using %s backend for MXFP8 GEMM", backend.value)
1686
+
1687
+ def create_weights(
1688
+ self,
1689
+ layer: torch.nn.Module,
1690
+ input_size_per_partition: int,
1691
+ output_partition_sizes: list[int],
1692
+ input_size: int,
1693
+ output_size: int,
1694
+ params_dtype: torch.dtype,
1695
+ **extra_weight_attrs,
1696
+ ):
1697
+ del input_size, output_size
1698
+
1699
+ if not self.quant_config.is_checkpoint_mxfp8_serialized:
1700
+ raise ValueError(
1701
+ "MXFP8 quantization was selected, but checkpoint is not "
1702
+ "MXFP8 serialized. Dynamic quantization is not supported."
1703
+ )
1704
+
1705
+ output_size_per_partition = sum(output_partition_sizes)
1706
+ weight_loader = extra_weight_attrs.get("weight_loader")
1707
+ layer.logical_widths = output_partition_sizes
1708
+ layer.input_size_per_partition = input_size_per_partition
1709
+ layer.output_size_per_partition = output_size_per_partition
1710
+
1711
+ if input_size_per_partition % MXFP8_BLOCK_SIZE != 0:
1712
+ raise ValueError(
1713
+ f"MXFP8 requires input dimension to be divisible by "
1714
+ f"{MXFP8_BLOCK_SIZE}, got {input_size_per_partition}"
1715
+ )
1716
+
1717
+ # Weight tensor: FP8 E4M3 format
1718
+ weight = ModelWeightParameter(
1719
+ data=torch.empty(
1720
+ output_size_per_partition,
1721
+ input_size_per_partition,
1722
+ dtype=MXFP8_VALUE_DTYPE,
1723
+ ),
1724
+ input_dim=1,
1725
+ output_dim=0,
1726
+ weight_loader=weight_loader,
1727
+ )
1728
+ layer.register_parameter("weight", weight)
1729
+
1730
+ # Weight scale tensor (E8M0 encoded as uint8), one scale per block of 32 along K
1731
+ weight_scale = ModelWeightParameter(
1732
+ data=torch.empty(
1733
+ output_size_per_partition,
1734
+ input_size_per_partition // MXFP8_BLOCK_SIZE,
1735
+ dtype=MXFP8_SCALE_DTYPE,
1736
+ ),
1737
+ input_dim=1,
1738
+ output_dim=0,
1739
+ weight_loader=weight_loader,
1740
+ )
1741
+ layer.register_parameter("weight_scale", weight_scale)
1742
+
1743
+ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
1744
+ if layer.weight.ndim != 2:
1745
+ raise ValueError(
1746
+ f"MXFP8 weight must be 2D tensor [N, K], got {layer.weight.ndim}D "
1747
+ f"with shape {tuple(layer.weight.shape)}"
1748
+ )
1749
+
1750
+ if layer.weight.dtype != MXFP8_VALUE_DTYPE:
1751
+ raise ValueError(
1752
+ f"MXFP8 weight must be {MXFP8_VALUE_DTYPE} (FP8 E4M3), "
1753
+ f"got {layer.weight.dtype}. The checkpoint may not be properly "
1754
+ f"quantized with MXFP8."
1755
+ )
1756
+
1757
+ weight = layer.weight.data # [N, K]
1758
+ N, K = weight.shape
1759
+ scale_k = K // MXFP8_BLOCK_SIZE
1760
+
1761
+ # Slice weight_scale to match weight dimensions (handles padding)
1762
+ weight_scale = layer.weight_scale.data[:N, :scale_k].contiguous()
1763
+
1764
+ layer.weight = Parameter(weight.contiguous(), requires_grad=False)
1765
+ layer.weight_scale = Parameter(weight_scale, requires_grad=False)
1766
+
1767
+ def apply(
1768
+ self,
1769
+ layer: torch.nn.Module,
1770
+ x: torch.Tensor,
1771
+ bias: torch.Tensor | None = None,
1772
+ ) -> torch.Tensor:
1773
+ if layer.weight.dtype != MXFP8_VALUE_DTYPE:
1774
+ raise ValueError(
1775
+ f"Weight dtype {layer.weight.dtype} != expected {MXFP8_VALUE_DTYPE}"
1776
+ )
1777
+ if layer.weight_scale.dtype != MXFP8_SCALE_DTYPE:
1778
+ raise ValueError(
1779
+ f"Weight scale dtype {layer.weight_scale.dtype} != "
1780
+ f"expected {MXFP8_SCALE_DTYPE}"
1781
+ )
1782
+
1783
+ return self.mxfp8_linear_op.apply(
1784
+ input=x,
1785
+ weight=layer.weight,
1786
+ weight_scale=layer.weight_scale,
1787
+ out_dtype=x.dtype,
1788
+ bias=bias,
1789
+ )
1790
+
1791
+
1792
+# Register the method classes for ModelOptMxFp8Config
1793
+ModelOptMxFp8Config.LinearMethodCls = ModelOptMxFp8LinearMethod
1794
+ModelOptMxFp8Config.KVCacheMethodCls = ModelOptFp8KVCacheMethod
vllm/model_executor/layers/quantization/utils/mxfp8_utils.py
+121
-11
@@ -1,24 +1,134 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+from enum import Enum
5
+
6
import torch
7
8
from vllm.logger import init_logger
9
+from vllm.utils.torch_utils import direct_register_custom_op
10
11
logger = init_logger(__name__)
12
13
11
-def mxfp8_e4m3_quantize(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
12
- try:
13
- from flashinfer import mxfp8_quantize as mxfp8_e4m3_quantize
14
- except ImportError as err:
15
- raise ImportError(
16
- "The package `flashinfer` is required to do "
17
- "MX-FP8 quantization. Please install it with"
18
- "`pip install flashinfer`"
19
- ) from err
14
+class Mxfp8LinearBackend(Enum):
15
+ EMULATION = "emulation"
16
+
17
+
18
+# MXFP8 constants
19
+MXFP8_VALUE_DTYPE = torch.float8_e4m3fn
20
+MXFP8_SCALE_DTYPE = torch.uint8
21
+MXFP8_BLOCK_SIZE = 32
22
+
23
21
- x_q, x_scales = mxfp8_e4m3_quantize(x, is_sf_swizzled_layout=False)
22
- if x_scales.ndim == 1:
24
+def _mxfp8_e4m3_quantize_impl(
25
+ x: torch.Tensor, is_sf_swizzled_layout: bool = False
26
+) -> tuple[torch.Tensor, torch.Tensor]:
27
+ from flashinfer import mxfp8_quantize as flashinfer_mxfp8_quantize
28
+
29
+ x_q, x_scales = flashinfer_mxfp8_quantize(
30
+ x, is_sf_swizzled_layout=is_sf_swizzled_layout
31
+ )
32
+ if x_scales.ndim == 1 and x.ndim == 2 and not is_sf_swizzled_layout:
33
x_scales = x_scales.view(x.size(0), -1)
34
return x_q, x_scales
35
+
36
+
37
+def mxfp8_e4m3_quantize(
38
+ x: torch.Tensor, is_sf_swizzled_layout: bool = False
39
+) -> tuple[torch.Tensor, torch.Tensor]:
40
+ return torch.ops.vllm.mxfp8_quantize(x, is_sf_swizzled_layout)
41
+
42
+
43
+def dequant_mxfp8_to_bf16(x: torch.Tensor, scales: torch.Tensor) -> torch.Tensor:
44
+ """Dequantize MXFP8 tensor to BF16."""
45
+ x_float = x.to(torch.float32)
46
+
47
+ num_blocks = x.shape[-1] // MXFP8_BLOCK_SIZE
48
+ x_blocked = x_float.view(*x.shape[:-1], num_blocks, MXFP8_BLOCK_SIZE)
49
+
50
+ descale = torch.exp2(scales.to(torch.float32) - 127.0)
51
+
52
+ dequantized = x_blocked * descale.unsqueeze(-1)
53
+
54
+ dequantized = dequantized.view(*x.shape)
55
+
56
+ return dequantized.to(torch.bfloat16)
57
+
58
+
59
+def mxfp8_e4m3_quantize_fake(
60
+ x: torch.Tensor, is_sf_swizzled_layout: bool = False
61
+) -> tuple[torch.Tensor, torch.Tensor]:
62
+ """Fake implementation for torch.compile tracing."""
63
+ fp_data = torch.empty_like(x, dtype=MXFP8_VALUE_DTYPE)
64
+
65
+ block_size = MXFP8_BLOCK_SIZE
66
+
67
+ if x.ndim == 2:
68
+ M, N = x.shape
69
+ K = (N + block_size - 1) // block_size
70
+ if is_sf_swizzled_layout:
71
+ M_padded = ((M + 127) // 128) * 128
72
+ K_padded = ((K + 3) // 4) * 4
73
+ scales = torch.empty(
74
+ M_padded * K_padded, dtype=MXFP8_SCALE_DTYPE, device=x.device
75
+ )
76
+ else:
77
+ scales = torch.empty((M, K), dtype=MXFP8_SCALE_DTYPE, device=x.device)
78
+ elif x.ndim == 3:
79
+ B, M, N = x.shape
80
+ K = (N + block_size - 1) // block_size
81
+ if is_sf_swizzled_layout:
82
+ M_padded = ((M + 127) // 128) * 128
83
+ K_padded = ((K + 3) // 4) * 4
84
+ scales = torch.empty(
85
+ B * M_padded * K_padded, dtype=MXFP8_SCALE_DTYPE, device=x.device
86
+ )
87
+ else:
88
+ scales = torch.empty((B, M, K), dtype=MXFP8_SCALE_DTYPE, device=x.device)
89
+ else:
90
+ scale_shape = list(x.shape)
91
+ scale_shape[-1] = (x.shape[-1] + block_size - 1) // block_size
92
+ scales = torch.empty(scale_shape, dtype=MXFP8_SCALE_DTYPE, device=x.device)
93
+
94
+ return fp_data, scales
95
+
96
+
97
+direct_register_custom_op(
98
+ op_name="mxfp8_quantize",
99
+ op_func=_mxfp8_e4m3_quantize_impl,
100
+ fake_impl=mxfp8_e4m3_quantize_fake,
101
+)
102
+
103
+
104
+class Mxfp8LinearOp:
105
+ def __init__(self, backend: Mxfp8LinearBackend):
106
+ if backend not in Mxfp8LinearBackend:
107
+ raise ValueError(f"Unsupported backend: {backend}")
108
+
109
+ self.backend = backend
110
+
111
+ def apply(
112
+ self,
113
+ input: torch.Tensor,
114
+ weight: torch.Tensor,
115
+ weight_scale: torch.Tensor,
116
+ out_dtype: torch.dtype,
117
+ bias: torch.Tensor | None = None,
118
+ ) -> torch.Tensor:
119
+ # Validate weight_scale dtype and shape (must be 2D for TORCH backend)
120
+ if weight_scale.dtype != MXFP8_SCALE_DTYPE:
121
+ raise ValueError(
122
+ f"TORCH backend requires {MXFP8_SCALE_DTYPE} weight_scale dtype, "
123
+ f"got {weight_scale.dtype}."
124
+ )
125
+ if weight_scale.ndim != 2:
126
+ raise ValueError(
127
+ f"TORCH backend requires 2D weight_scale, got {weight_scale.ndim}D. "
128
+ f"Ensure process_weights_after_loading was called."
129
+ )
130
+
131
+ weight_bf16 = dequant_mxfp8_to_bf16(weight, weight_scale)
132
+
133
+ output = torch.nn.functional.linear(input, weight_bf16, bias)
134
+ return output.to(out_dtype)