[Bugfix] Fix HSDP compatibility with the new online FP8 linear method (#5677)

Signed-off-by: roG0d <baonudesifeizhai@gmail.com>

docs/user_guide/diffusion/parallelism/hsdp.md

@@ -64,6 +64,41 @@ omni = Omni(
 )
 ```
 
+### Combined with Online FP8 Quantization
+
+HSDP can be combined with online FP8 quantization for diffusion models that
+support both features. HSDP shards the quantized weights across GPUs to reduce
+the per-GPU weight memory footprint.
+
+```python
+from vllm_omni import Omni
+from vllm_omni.diffusion.data import DiffusionParallelConfig
+
+omni = Omni(
+    model="<your-model>",
+    quantization="fp8",
+    parallel_config=DiffusionParallelConfig(
+        use_hsdp=True,
+        hsdp_shard_size=2,
+    ),
+)
+```
+
+For online serving:
+
+```bash
+vllm serve <your-model> \
+    --omni \
+    --quantization fp8 \
+    --use-hsdp \
+    --hsdp-shard-size 2
+```
+
+The model must support both HSDP and online FP8 quantization. HSDP remains
+incompatible with Tensor Parallelism. See the
+[FP8 Quantization Guide](../../quantization/fp8.md) for supported models and
+hardware.
+
 ---
 
 ## Example Script
@@ -147,3 +182,4 @@ For detailed instructions on adding HSDP support to new models, see the [HSDP Co
 1. ✅ **Enable HSDP** - Set `use_hsdp=True` and `hsdp_shard_size` to reduce per-GPU memory for large models
 2. ✅ **Combine with SP** - Use together with `ulysses_degree` for video models requiring both memory reduction and sequence parallelism
 3. ⚠️ **Incompatible with TP** - `tensor_parallel_size` must be 1 when HSDP is enabled
+4. ✅ **Combine with online FP8** - Use `quantization="fp8"` or `--quantization fp8` to reduce the sharded weight memory footprint

tests/diffusion/quantization/test_hsdp_fp8.py

@@ -36,15 +36,23 @@ class _ToyKernel:
 # --- helper to build a toy module simulating online-FP8 post-load state ---
 
 
-def _make_fp8_toy_module(out_features: int = 16, in_features: int = 32):
+def _make_fp8_toy_module(
+    out_features: int = 16,
+    in_features: int = 32,
+    *,
+    quant_method_cls=None,
+):
     """Return a module whose weight is ``qweight.t()`` (non-contiguous) and whose
     ``quant_method`` is an ``Fp8LinearMethod`` with a ``_ToyKernel``.
     """
     from vllm.model_executor.layers.quantization.fp8 import Fp8LinearMethod
 
+    if quant_method_cls is None:
+        quant_method_cls = Fp8LinearMethod
+
     # Create the quant-method instance without calling __init__ (avoids
     # heavyweight upstream config dependencies).  isinstance() still works.
-    qm = object.__new__(Fp8LinearMethod)
+    qm = object.__new__(quant_method_cls)
     qm.fp8_linear = _ToyKernel()
 
     # randn does not support float8 on CPU; zeros does and the test only
@@ -104,6 +112,28 @@ def test_rewrites_non_contiguous_weight():
     assert module.weight.shape == (16, 32)
 
 
+def test_rewrites_per_tensor_online_fp8_weight():
+    """Regression test for vLLM's new online FP8 frontend (#45463)."""
+    from vllm.model_executor.layers.quantization.online.fp8 import (
+        Fp8PerTensorOnlineLinearMethod,
+    )
+
+    module = _make_fp8_toy_module(
+        16,
+        32,
+        quant_method_cls=Fp8PerTensorOnlineLinearMethod,
+    )
+
+    n_rewritten_layers = prepare_fp8_layers_for_fsdp(module)
+
+    assert n_rewritten_layers == 1
+    assert module.weight.is_contiguous()
+    assert module.weight.shape == (16, 32)
+    w, *_ = module.quant_method.fp8_linear._get_layer_params(module)
+    assert w.shape == (32, 16)
+    assert w.stride() == (1, 32)
+
+
 def test_patched_get_layer_params_returns_transposed_view():
     module = _make_fp8_toy_module(16, 32)
     prepare_fp8_layers_for_fsdp(module)

vllm_omni/diffusion/quantization/hsdp_fp8.py

@@ -33,9 +33,17 @@ import types
 from torch import nn
 from vllm.logger import init_logger
 from vllm.model_executor.layers.quantization.fp8 import Fp8LinearMethod
+from vllm.model_executor.layers.quantization.online.fp8 import (
+    Fp8PerTensorOnlineLinearMethod,
+)
 
 logger = init_logger(__name__)
 
+_FP8_TRANSPOSED_WEIGHT_METHODS = (
+    Fp8LinearMethod,
+    Fp8PerTensorOnlineLinearMethod,
+)
+
 
 def _build_transposed_get_layer_params(original_bound_method):
     """Return a ``_get_layer_params`` replacement that transposes ``w``.
@@ -55,9 +63,8 @@ def _build_transposed_get_layer_params(original_bound_method):
 def prepare_fp8_layers_for_fsdp(model: nn.Module) -> int:
     """Make online-FP8 linear layers in ``model`` FSDP2-compatible.
 
-    For every layer whose ``quant_method`` is an :class:`Fp8LinearMethod`
-    and whose weight is currently a non-contiguous transpose view, this
-    function:
+    For every layer whose quantization method stores FP8 weights as a
+    non-contiguous transpose view, this function:
 
     1. Replaces ``layer.weight`` with the underlying ``(out, in)`` row-major
        contiguous storage so FSDP2 ``fully_shard`` accepts it.
@@ -76,7 +83,7 @@ def prepare_fp8_layers_for_fsdp(model: nn.Module) -> int:
     patched_kernel_ids: set[int] = set()
     for module in model.modules():
         qm = getattr(module, "quant_method", None)
-        if not isinstance(qm, Fp8LinearMethod):
+        if not isinstance(qm, _FP8_TRANSPOSED_WEIGHT_METHODS):
             continue
 
         weight = getattr(module, "weight", None)