From 1903cee56a09db5fdd52229bc974ed9e51d5ebac Mon Sep 17 00:00:00 2001 From: "Chendi.Xue" Date: Wed, 27 May 2026 21:47:22 -0500 Subject: [PATCH] [MXFP8][XPU] enable mxfp8 using vLLM main repo method (#3782) Signed-off-by: Chendi Xue --- .../scripts/hardware_ci/run-xpu-test.sh | 1 + docs/user_guide/quantization/mxfp8.md | 63 ++++++-- docs/user_guide/quantization/online.md | 4 +- .../quantization/test_mxfp8_config.py | 53 +++++-- .../model_loader/diffusers_loader.py | 25 +++- vllm_omni/quantization/factory.py | 6 + vllm_omni/quantization/inc_config.py | 139 +++++++++++++++++- vllm_omni/quantization/mxfp8_config.py | 72 ++++++++- 8 files changed, 328 insertions(+), 35 deletions(-) diff --git a/.buildkite/scripts/hardware_ci/run-xpu-test.sh b/.buildkite/scripts/hardware_ci/run-xpu-test.sh index ab8c2fd5d..fa4f00e29 100755 --- a/.buildkite/scripts/hardware_ci/run-xpu-test.sh +++ b/.buildkite/scripts/hardware_ci/run-xpu-test.sh @@ -52,6 +52,7 @@ time timeout -k 30 30m docker run \ pip install tblib==3.1.0 cd /workspace/vllm-omni pytest -v -s -m "core_model and xpu and B60" + pytest -v -s tests/diffusion/quantization/test_mxfp8_config.py export VLLM_XPU_USE_SAMPLER_KERNEL=0 # NOTE: Remove this after vLLM v0.21.1 is merged. Fixes Qwen2-5 omni-expansion tests. pytest -v -s -m "advanced_model and xpu and B60" ' diff --git a/docs/user_guide/quantization/mxfp8.md b/docs/user_guide/quantization/mxfp8.md index 418aad369..823e898fd 100644 --- a/docs/user_guide/quantization/mxfp8.md +++ b/docs/user_guide/quantization/mxfp8.md @@ -7,26 +7,28 @@ using the OCP MX format: groups of 32 K-dimension elements share a single `float8_e8m0fnu` exponent scale. This gives better accuracy than channel-wise FP8 while keeping the same 8-bit weight footprint. -This method supports two modes: +This method supports three modes: | Mode | Description | |------|-------------| | **Online** | BF16 weights are quantized to MXFP8 at load time — no pre-processing needed | -| **Offline** | msModelSlim-exported MXFP8 weights converted to diffusers format via `merge_mxfp8_checkpoint.py` — weights and scales are loaded directly from the preprocessed checkpoint | +| **Offline (Native)** | msModelSlim-exported MXFP8 weights converted to diffusers format via `merge_mxfp8_checkpoint.py` — weights and scales are loaded directly from the preprocessed checkpoint | +| **Offline (AutoRound)** | AutoRound MXFP8 checkpoints with `data_type="mx_fp"` — auto-detected from `config.json` | ## Hardware Support -| Device | Support | -|--------|---------| -| NVIDIA Blackwell GPU (SM 100+) | ⭕ | -| NVIDIA Ada/Hopper GPU (SM 89+) | ⭕ | -| NVIDIA Ampere GPU (SM 80+) | ⭕ | -| AMD ROCm | ⭕ | -| Intel XPU | ⭕ | -| Ascend NPU (Atlas 950 A5) | ✅ | +| Device | Online | Offline (Native) | Offline (AutoRound) | +|--------|--------|------------------|---------------------| +| NVIDIA Blackwell GPU (SM 100+) | ⭕ | ⭕ | ⭕ | +| NVIDIA Ada/Hopper GPU (SM 89+) | ⭕ | ⭕ | ⭕ | +| NVIDIA Ampere GPU (SM 80+) | ⭕ | ⭕ | ⭕ | +| AMD ROCm | ⭕ | ⭕ | ⭕ | +| Intel XPU | ✅ | ❌ | ✅ | +| Ascend NPU (Atlas 950 A5) | ✅ | ✅ | ⭕ | -Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this -guide. +Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this guide. + +**Note**: Intel XPU only supports AutoRound MXFP8 for offline mode. Use AutoRound quantized checkpoints or online mode for XPU. ## Model Type Support @@ -82,9 +84,9 @@ python text_to_video.py --model --quantization mxfp8 vllm serve --omni --quantization mxfp8 ``` -### Offline Mode +### Offline Mode (Native) -Offline mode loads a pre-quantized checkpoint from msModelSlim. A preprocessing +Native offline mode loads a pre-quantized checkpoint from msModelSlim. A preprocessing step converts the raw quantized output to the diffusers format expected by vLLM-Omni and injects the quantization config into `transformer/config.json` so that vLLM-Omni auto-detects the offline path without a `--quantization` flag. @@ -159,10 +161,41 @@ omni = Omni(model="/path/to/Wan2.2-TI2V-5B-MXFP8") ``` !!! note - No `--quantization` flag is needed for offline mode. The preprocessing + No `--quantization` flag is needed for native offline mode. The preprocessing script injects `quantization_config` into each `transformer/config.json`, which vLLM-Omni reads automatically to activate the offline MXFP8 method. +### Offline Mode (AutoRound) + +AutoRound MXFP8 checkpoints declare `quant_method="auto-round"` with `data_type="mx_fp"` in their `config.json`. These are automatically detected and use the `IncMxfp8OfflineLinearMethod` backend. + +To use an AutoRound MXFP8 checkpoint: + +```bash +python text_to_video.py --model + +# Online serving +vllm serve --omni +``` + +Python API: + +```python +omni = Omni(model="") +``` + +!!! note + AutoRound MXFP8 checkpoints are auto-detected from `config.json` and do not require a `--quantization` flag. The config must include: + ```json + { + "quantization_config": { + "quant_method": "auto-round", + "data_type": "mx_fp", + ... + } + } + ``` + ## Parameters | Parameter | Type | Default | Description | diff --git a/docs/user_guide/quantization/online.md b/docs/user_guide/quantization/online.md index b9a6861d1..6103be9a0 100644 --- a/docs/user_guide/quantization/online.md +++ b/docs/user_guide/quantization/online.md @@ -21,7 +21,7 @@ by msModelSlim and the merge tools. | NVIDIA Ada/Hopper GPU (SM 89+) | ✅ | ✅ | ⭕ | ⭕ | | NVIDIA Ampere GPU (SM 80+) | ✅ | ✅ | ⭕ | ⭕ | | AMD ROCm | ⭕ | ⭕ | ⭕ | ⭕ | -| Intel XPU | ⭕ | ⭕ | ⭕ | ⭕ | +| Intel XPU | ⭕ | ⭕ | ✅ | ⭕ | | Ascend NPU | ❌ | ✅ | ✅ | ✅ | Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this @@ -36,7 +36,7 @@ MXFP4 are documented for the Ascend NPU path. |--------|-------|----------------|--------| | FP8 W8A8 | [FP8](fp8.md) | Qwen-Image; Wan2.2 is not validated | Validated for Qwen-Image family and other DiT models | | Int8 W8A8 | [Int8](int8.md) | Qwen-Image; Wan2.2 is not validated | Validated for Qwen-Image and Z-Image | -| MXFP8 W8A8 | [MXFP8](mxfp8.md) | Wan2.2-T2V-A14B, Wan2.2-I2V-A14B, Wan2.2-TI2V-5B | Ascend NPU only | +| MXFP8 W8A8 | [MXFP8](mxfp8.md) | Wan2.2-T2V-A14B, Wan2.2-I2V-A14B, Wan2.2-TI2V-5B | Validated on Ascend NPU and Intel XPU | | MXFP4 W4A4 | [MXFP4](mxfp4.md) | Wan2.2-T2V-A14B, Wan2.2-I2V-A14B | Ascend NPU only; TI2V-5B is not supported | ### Multi-Stage Omni/TTS Model (Qwen3-Omni, Qwen3-TTS) diff --git a/tests/diffusion/quantization/test_mxfp8_config.py b/tests/diffusion/quantization/test_mxfp8_config.py index d831b5d47..cc34fb232 100644 --- a/tests/diffusion/quantization/test_mxfp8_config.py +++ b/tests/diffusion/quantization/test_mxfp8_config.py @@ -115,37 +115,64 @@ def test_mxfp8_quant_config_structure(): # --------------------------------------------------------------------------- -def test_get_quant_method_npu_offline(mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch): - """Offline (serialized) path on NPU must return NPUMxfp8LinearMethod.""" - from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config, NPUMxfp8LinearMethod +@pytest.mark.skipif(not current_omni_platform.is_npu(), reason="Native MXFP8 offline only supported on NPU") +def test_get_quant_method_offline_npu(mocker: MockerFixture): + """Offline (serialized) path must return NPUMxfp8LinearMethod on NPU.""" + from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config config = DiffusionMXFP8Config(is_checkpoint_mxfp8_serialized=True) layer = mocker.Mock(spec=LinearBase) - monkeypatch.setattr(current_omni_platform, "is_npu", lambda: True) method = config.get_quant_method(layer, "blocks.0.attn1.to_q") - assert isinstance(method, NPUMxfp8LinearMethod) + assert type(method).__name__ == "NPUMxfp8LinearMethod" -def test_get_quant_method_npu_online(mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch): - """Online (BF16 checkpoint) path on NPU must return NPUMxfp8OnlineLinearMethod.""" - from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config, NPUMxfp8OnlineLinearMethod +@pytest.mark.skipif(not current_omni_platform.is_xpu(), reason="XPU platform not available") +def test_get_quant_method_offline_xpu_raises(mocker: MockerFixture): + """XPU offline mode must raise NotImplementedError (use AutoRound MXFP8 instead).""" + from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config + + config = DiffusionMXFP8Config(is_checkpoint_mxfp8_serialized=True) + layer = mocker.Mock(spec=LinearBase) + + with pytest.raises(NotImplementedError, match="Native MXFP8 offline mode is not supported on XPU"): + config.get_quant_method(layer, "blocks.0.attn1.to_q") + + +@pytest.mark.skipif( + not (current_omni_platform.is_npu() or current_omni_platform.is_xpu()), reason="MXFP8 only supported on NPU and XPU" +) +def test_get_quant_method_online(mocker: MockerFixture): + """Online (BF16 checkpoint) path must return platform-specific method on current platform.""" + from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config + + # Mock the vLLM online method to avoid config dependency for XPU + if current_omni_platform.is_xpu(): + mock_inner = mocker.Mock() + mocker.patch( + "vllm_omni.quantization.mxfp8_config.VllmMxfp8OnlineLinearMethod.__init__", + lambda self: setattr(self, "_inner", mock_inner), + ) config = DiffusionMXFP8Config(is_checkpoint_mxfp8_serialized=False) layer = mocker.Mock(spec=LinearBase) - monkeypatch.setattr(current_omni_platform, "is_npu", lambda: True) method = config.get_quant_method(layer, "blocks.0.attn1.to_q") - assert isinstance(method, NPUMxfp8OnlineLinearMethod) + + if current_omni_platform.is_npu(): + assert type(method).__name__ == "NPUMxfp8OnlineLinearMethod" + elif current_omni_platform.is_xpu(): + assert type(method).__name__ == "VllmMxfp8OnlineLinearMethod" -def test_get_quant_method_non_npu_raises(mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch): - """Non-NPU platform must raise NotImplementedError (CUDA not yet supported).""" +def test_get_quant_method_unsupported_platform_raises(mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch): + """Unsupported platform (CUDA, ROCm) must raise NotImplementedError.""" from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config config = DiffusionMXFP8Config() layer = mocker.Mock(spec=LinearBase) monkeypatch.setattr(current_omni_platform, "is_npu", lambda: False) + monkeypatch.setattr(current_omni_platform, "is_xpu", lambda: False) monkeypatch.setattr(current_omni_platform, "is_cuda", lambda: False) with pytest.raises(NotImplementedError): @@ -159,6 +186,8 @@ def test_get_quant_method_ignored_layer(mocker: MockerFixture, monkeypatch: pyte config = DiffusionMXFP8Config(ignored_layers=["proj_out"]) layer = mocker.Mock(spec=LinearBase) monkeypatch.setattr(current_omni_platform, "is_npu", lambda: True) + monkeypatch.setattr(current_omni_platform, "is_xpu", lambda: False) + monkeypatch.setattr(current_omni_platform, "is_cuda", lambda: False) method = config.get_quant_method(layer, "proj_out") assert isinstance(method, UnquantizedLinearMethod) diff --git a/vllm_omni/diffusion/model_loader/diffusers_loader.py b/vllm_omni/diffusion/model_loader/diffusers_loader.py index f4ebfcb24..9fee2b67e 100644 --- a/vllm_omni/diffusion/model_loader/diffusers_loader.py +++ b/vllm_omni/diffusion/model_loader/diffusers_loader.py @@ -301,10 +301,25 @@ class DiffusersPipelineLoader: if load_format is None: load_format = "default" self.od_config = od_config - # CPU offload + FP8: load weights on device for FP8 quantization + # CPU offload + quantization: for offline-quantized models (e.g., AutoRound MXFP8), + # weights are already quantized in the checkpoint — load directly on CPU. + # For online quantization, load on device so quantization can run on accelerator, + # then move back to CPU afterward. + offload_after_quant = False if load_device == "cpu" and od_config.quantization_config is not None: - load_device = device.type - logger.info(f"Quantization enabled with CPU offload, using {load_device} for weight loading") + quant_cfg = od_config.quantization_config + is_offline = getattr(quant_cfg, "data_type", None) == "mx_fp" or getattr( + quant_cfg, "is_checkpoint_quantized", False + ) + if not is_offline: + load_device = device.type + offload_after_quant = True + logger.info( + "Online quantization with CPU offload, using %s for weight loading (will offload back to CPU)", + load_device, + ) + else: + logger.info("Offline-quantized model with CPU offload, loading weights directly on CPU") target_device = torch.device(load_device) with set_default_torch_dtype(od_config.dtype): @@ -342,6 +357,10 @@ class DiffusersPipelineLoader: # This is needed for vLLM's quantization methods that need to transform weights self._process_weights_after_loading(model, target_device) + if offload_after_quant: + model.to("cpu") + logger.info("Quantization complete, offloaded model back to CPU") + return model.eval() def _process_weights_after_loading(self, model: nn.Module, target_device: torch.device) -> None: diff --git a/vllm_omni/quantization/factory.py b/vllm_omni/quantization/factory.py index 19eafab22..2a0326c2e 100644 --- a/vllm_omni/quantization/factory.py +++ b/vllm_omni/quantization/factory.py @@ -399,6 +399,12 @@ def resolve_quant_config_from_disk( ) return build_quant_config(qc_method, **qc_kwargs) + # AutoRound MXFP8 checkpoints use data_type="mx_fp" instead of + # is_checkpoint_*_serialized; rebuild so the offline path is selected. + if qc_kwargs.get("data_type") == "mx_fp": + logger.info("config.json declares data_type='mx_fp'; rebuilding as offline AutoRound MXFP8.") + return build_quant_config(qc_method, **qc_kwargs) + if ( "ignored_layers" in qc_kwargs and hasattr(quant_config, "ignored_layers") diff --git a/vllm_omni/quantization/inc_config.py b/vllm_omni/quantization/inc_config.py index fb4851cca..d7dad3d71 100644 --- a/vllm_omni/quantization/inc_config.py +++ b/vllm_omni/quantization/inc_config.py @@ -7,8 +7,13 @@ from __future__ import annotations from os.path import commonprefix from typing import TYPE_CHECKING, Any +import torch +from torch.nn import Module +from vllm.model_executor.layers.linear import LinearMethodBase from vllm.model_executor.layers.quantization.inc import INCConfig from vllm.model_executor.models.utils import WeightsMapper +from vllm.model_executor.parameter import ModelWeightParameter +from vllm.model_executor.utils import replace_parameter if TYPE_CHECKING: from vllm.model_executor.layers.quantization.base_config import ( @@ -45,12 +50,39 @@ def _map_with_stage_prefix( class OmniINCConfig(INCConfig): - """INCConfig extended with multi-stage prefix remapping.""" + """INCConfig extended with multi-stage prefix remapping and MXFP8 support. + + AutoRound MXFP8 checkpoints declare quant_method="auto-round" with data_type="mx_fp". + This config detects that case and dispatches to IncMxfp8OfflineLinearMethod instead + of the standard INT quantization path. + + Architecture: + - AutoRound INT4/INT8 → standard INCConfig path (xpu_w4a16_quant_layer, etc.) + - AutoRound MXFP8 (data_type="mx_fp") → IncMxfp8OfflineLinearMethod + - Native MXFP8 (quant_method="mxfp8") → DiffusionMXFP8Config → NPUMxfp8LinearMethod (NPU only) + """ + + # Extend supported data types and formats to include MXFP8 + SUPPORTED_DTYPES = {"int", "mx_fp"} + SUPPORTED_FORMATS = {"auto_round:auto_gptq", "auto_round:auto_awq", "auto_round:llm_compressor"} # ------------------------------------------------------------------ # Core integration: called by vLLM's configure_quant_config() # ------------------------------------------------------------------ + def get_quant_method(self, layer, prefix: str): + """Get quantization method, handling AutoRound MXFP8 as a special case.""" + # Check if this is an AutoRound MXFP8 checkpoint (data_type="mx_fp") + if hasattr(self, "data_type") and self.data_type == "mx_fp" and self.weight_bits == 8: + from vllm.model_executor.layers.linear import LinearBase + + if isinstance(layer, LinearBase): + return IncMxfp8OfflineLinearMethod() + return None + + # Otherwise, use parent INCConfig logic + return super().get_quant_method(layer, prefix) + def apply_vllm_mapper(self, hf_to_vllm_mapper: WeightsMapper) -> None: """Remap HF checkpoint names to vLLM runtime prefixes.""" prefix_map = getattr(hf_to_vllm_mapper, "orig_to_new_prefix", None) or {} @@ -138,3 +170,108 @@ class OmniINCConfig(INCConfig): if isinstance(quant_config, INCConfig): return cls.from_inc_config(quant_config) return quant_config + + +# --------------------------------------------------------------------------- +# IncMxfp8OfflineLinearMethod - used by AutoRound MXFP8 checkpoints +# --------------------------------------------------------------------------- + + +class IncMxfp8OfflineLinearMethod(LinearMethodBase): + """Offline MXFP8 linear method for AutoRound MXFP8 checkpoints. + + For pre-quantized AutoRound MXFP8 checkpoints (data_type="mx_fp") where weights + are stored as BF16 (losslessly representing FP8 E4M3 values) with uint8 MX block scales. + Delegates to the platform-appropriate kernel (XPUMxFp8LinearKernel, etc.). + + Used by OmniINCConfig when data_type="mx_fp". + """ + + def __init__(self) -> None: + from vllm.model_executor.kernels.linear import init_mxfp8_linear_kernel + from vllm.model_executor.layers.quantization.utils.mxfp8_utils import ( + MXFP8_BLOCK_SIZE, + MXFP8_SCALE_DTYPE, + ) + + self.kernel = init_mxfp8_linear_kernel() + self.block_size = MXFP8_BLOCK_SIZE + self.scale_dtype = MXFP8_SCALE_DTYPE + + def create_weights( + self, + layer: torch.nn.Module, + input_size_per_partition: int, + output_partition_sizes: list[int], + input_size: int, + output_size: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ) -> None: + if input_size_per_partition % self.block_size != 0: + raise ValueError( + f"MXFP8 requires input_size_per_partition " + f"({input_size_per_partition}) to be divisible by " + f"{self.block_size}." + ) + + output_size_per_partition = sum(output_partition_sizes) + weight_loader = extra_weight_attrs.get("weight_loader") + + layer.logical_widths = output_partition_sizes + layer.input_size_per_partition = input_size_per_partition + layer.output_size_per_partition = output_size_per_partition + layer.orig_dtype = params_dtype + + # Allocate as params_dtype (BF16) to match checkpoint dtype. + # Will be cast to FP8 in process_weights_after_loading. + layer.register_parameter( + "weight", + ModelWeightParameter( + data=torch.empty(output_size_per_partition, input_size_per_partition, dtype=params_dtype), + input_dim=1, + output_dim=0, + weight_loader=weight_loader, + ), + ) + + num_groups = input_size_per_partition // self.block_size + layer.register_parameter( + "weight_scale", + ModelWeightParameter( + data=torch.empty(output_size_per_partition, num_groups, dtype=self.scale_dtype), + input_dim=1, + output_dim=0, + weight_loader=weight_loader, + ), + ) + + def process_weights_after_loading(self, layer: Module) -> None: + if getattr(layer, "_already_called_process_weights_after_loading", False): + return + + # Cast BF16 weight to FP8. + # Checkpoint stores weights as BF16 (losslessly representing FP8 E4M3 values). + w = layer.weight.data + if w.dtype != torch.float8_e4m3fn: + # Cast to FP8 E4M3. AutoRound stores FP8 values as BF16, so this is lossless. + w = w.to(torch.float8_e4m3fn).contiguous() + replace_parameter(layer, "weight", w) + + # Delegate layout transforms (transpose, scale reinterpret) to the platform kernel. + self.kernel.process_weights_after_loading(layer) + layer._already_called_process_weights_after_loading = True + + def apply( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + ori_shape = x.shape + if x.dim() > 2: + x = x.reshape(-1, ori_shape[-1]) + output = self.kernel.apply_weights(layer, x, bias) + if len(ori_shape) > 2: + output = output.reshape(*ori_shape[:-1], -1) + return output diff --git a/vllm_omni/quantization/mxfp8_config.py b/vllm_omni/quantization/mxfp8_config.py index 544c6a4ca..9d3216000 100644 --- a/vllm_omni/quantization/mxfp8_config.py +++ b/vllm_omni/quantization/mxfp8_config.py @@ -133,10 +133,17 @@ class DiffusionMXFP8Config(QuantizationConfig): if self.is_checkpoint_mxfp8_serialized: return NPUMxfp8LinearMethod(self) return NPUMxfp8OnlineLinearMethod(self) - # Placeholder for future platforms: add `elif is_cuda(): ...` here. + if current_omni_platform.is_xpu(): + if self.is_checkpoint_mxfp8_serialized: + raise NotImplementedError( + "Native MXFP8 offline mode is not supported on XPU. " + "Use AutoRound MXFP8 checkpoints (quant_method='auto-round', data_type='mx_fp') instead, " + "or use online MXFP8 mode without is_checkpoint_mxfp8_serialized." + ) + return VllmMxfp8OnlineLinearMethod() raise NotImplementedError( "DiffusionMXFP8Config (W8A8 MXFP8) is currently only supported " - "on NPU (Ascend) platforms. CUDA support is not yet implemented." + "on NPU (Ascend) and XPU (Intel) platforms." ) return None @@ -450,3 +457,64 @@ class NPUMxfp8OnlineLinearMethod(_LazyWeightMixin, NPUMxfp8LinearMethod): replace_parameter(layer, "weight", weight_fp8) replace_parameter(layer, "weight_scale", weight_scale) layer._already_called_process_weights_after_loading = True + + +# --------------------------------------------------------------------------- +# vLLM kernel-based methods (XPU, CUDA, ROCm — any platform with Mxfp8LinearKernel) +# --------------------------------------------------------------------------- +# Note: VllmMxfp8OfflineLinearMethod removed - XPU only supports AutoRound MXFP8 offline +# (see IncMxfp8OfflineLinearMethod in inc_config.py) + +try: + from vllm.model_executor.layers.quantization.online.mxfp8 import ( + Mxfp8OnlineLinearMethod as _VllmMxfp8OnlineBase, + ) +except ImportError as e: + raise ImportError( + "vLLM MXFP8 support is not available. " + "Online MXFP8 quantization requires vLLM with MXFP8 kernel support. " + "Please upgrade vLLM or use a compatible build." + ) from e + + +class VllmMxfp8OnlineLinearMethod(_VllmMxfp8OnlineBase): + """Online MXFP8 linear method extending vLLM's implementation with 3D tensor support. + + Loads BF16 weights, quantizes to MXFP8 at load time using vLLM's kernel. + Adds reshape logic to handle 3D tensors from diffusion models. + + Inherits from vLLM's Mxfp8OnlineLinearMethod and only overrides __init__ and apply(). + __init__ directly initializes the kernel without calling super().__init__() to avoid + vLLM config dependency in unit tests. All other methods (create_weights, + process_weights_after_loading) are inherited directly from vLLM. + """ + + def __init__(self) -> None: + """Initialize kernel directly without calling parent __init__ to avoid config dependency. + + TODO: skipping super().__init__() couples this to vLLM internals. If the parent + adds required initialization, consider a version-aware assertion or calling + super().__init__() with a fallback config. + """ + from vllm.model_executor.kernels.linear import init_mxfp8_linear_kernel + + self.kernel = init_mxfp8_linear_kernel() + + def apply( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + """Apply with 3D tensor support for diffusion models. + + vLLM's kernel expects 2D tensors (M, K), but diffusion models pass + 3D tensors (batch, seq, hidden). Flatten before kernel, reshape after. + """ + ori_shape = x.shape + if x.dim() > 2: + x = x.reshape(-1, ori_shape[-1]) + output = super().apply(layer, x, bias) + if len(ori_shape) > 2: + output = output.reshape(*ori_shape[:-1], -1) + return output