[MXFP8][XPU] enable mxfp8 using vLLM main repo method (#3782)

Signed-off-by: Chendi Xue <chendi.xue@intel.com>

.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"
 '

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 <your-model> --quantization mxfp8
 vllm serve <your-model> --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 <autoround-mxfp8-model>
+
+# Online serving
+vllm serve <autoround-mxfp8-model> --omni
+```
+
+Python API:
+
+```python
+omni = Omni(model="<autoround-mxfp8-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 |

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)

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)

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:

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")

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

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