[Bugfix] Fix HSDP + FP8 online quantization compatibility (#4494)

Signed-off-by: Kwokho Tsui <kh_tsui@163.com>
Co-authored-by: Kwokho Tsui <kh_tsui@163.com>

tests/diffusion/quantization/test_hsdp_fp8.py

@@ -0,0 +1,147 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for HSDP/FSDP2 compatibility with online FP8 quantization."""
+
+import pytest
+import torch
+from torch import nn
+
+from vllm_omni.diffusion.quantization.hsdp_fp8 import (
+    _build_transposed_get_layer_params,
+    prepare_fp8_layers_for_fsdp,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+# --- shared test doubles ---
+
+
+class _ToyKernel:
+    """Minimal kernel that mirrors ``FP8ScaledMMLinearKernel``."""
+
+    def __init__(self):
+        self.layer_param_names = ("weight", "weight_scale", "input_scale", "input_scale_ub")
+
+    def _get_layer_params(self, layer):
+        w, w_s, x_s, x_s_ub = self.layer_param_names
+        return (
+            getattr(layer, w),
+            getattr(layer, w_s),
+            getattr(layer, x_s, None),
+            getattr(layer, x_s_ub, None),
+        )
+
+
+# --- 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):
+    """Return a module whose weight is ``qweight.t()`` (non-contiguous) and whose
+    ``quant_method`` is an ``Fp8OnlineLinearMethod`` with a ``_ToyKernel``.
+    """
+    from vllm.model_executor.layers.quantization.fp8 import Fp8OnlineLinearMethod
+
+    # Create the quant-method instance without calling __init__ (avoids
+    # heavyweight upstream config dependencies).  isinstance() still works.
+    qm = object.__new__(Fp8OnlineLinearMethod)
+    qm.fp8_linear = _ToyKernel()
+
+    # randn does not support float8 on CPU; zeros does and the test only
+    # cares about shape/stride, not the actual weight values.
+    qweight = torch.zeros(out_features, in_features, dtype=torch.float8_e4m3fn)
+    weight = nn.Parameter(qweight.t())  # (in, out) non-contiguous column-major view
+
+    module = nn.Module()
+    module.quant_method = qm
+    module.weight = weight
+    module.weight_scale = nn.Parameter(torch.ones(1, dtype=torch.float32))
+    # Set input_scale / input_scale_ub to verify _get_layer_params
+    # passes them through unchanged.
+    module.input_scale = nn.Parameter(torch.ones(1, dtype=torch.float32))
+    module.input_scale_ub = nn.Parameter(torch.ones(1, dtype=torch.float32))
+
+    return module
+
+
+# --- tests ---
+
+
+def test_transposed_get_layer_params():
+    import types
+
+    kernel = _ToyKernel()
+
+    # original bound method + patch + re-bind, matching prepare_fp8_layers_for_fsdp
+    original_bound = kernel._get_layer_params
+    patched_func = _build_transposed_get_layer_params(original_bound)
+    kernel._get_layer_params = types.MethodType(patched_func, kernel)
+
+    module = _make_fp8_toy_module(16, 32)
+    # Reproduce the storage rewrite from prepare_fp8_layers_for_fsdp:
+    # weight was qweight.t()  →  .t() recovers qweight (row-major contiguous).
+    module.weight = nn.Parameter(module.weight.data.t(), requires_grad=False)
+    assert module.weight.is_contiguous(), "qweight should be row-major contiguous"
+
+    w, w_s, x_s, x_s_ub = kernel._get_layer_params(module)
+
+    # patched method transposes row-major → column-major for Cutlass
+    assert w.shape == (32, 16)
+    assert not w.is_contiguous()
+    # scales are passed through unchanged
+    assert w_s is module.weight_scale
+    assert x_s is module.input_scale
+    assert x_s_ub is module.input_scale_ub
+
+
+def test_rewrites_non_contiguous_weight():
+    module = _make_fp8_toy_module(16, 32)
+    assert not module.weight.is_contiguous()
+
+    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)
+
+
+def test_patched_get_layer_params_returns_transposed_view():
+    module = _make_fp8_toy_module(16, 32)
+    prepare_fp8_layers_for_fsdp(module)
+
+    kernel = module.quant_method.fp8_linear
+    w, w_s, x_s, x_s_ub = kernel._get_layer_params(module)
+
+    # Cutlass expects column-major B: (in, out), stride (1, in)
+    assert w.shape == (32, 16)
+    assert not w.is_contiguous()
+    assert w.stride() == (1, 32)
+
+    # scales pass through unchanged
+    assert w_s is module.weight_scale
+    assert x_s is module.input_scale
+    assert x_s_ub is module.input_scale_ub
+
+
+def test_contiguous_weights_skipped():
+    module = _make_fp8_toy_module(16, 32)
+    # make weight contiguous before calling
+    module.weight = nn.Parameter(module.weight.data.t().contiguous())
+    assert module.weight.is_contiguous()
+
+    n_rewritten_layers = prepare_fp8_layers_for_fsdp(module)
+    assert n_rewritten_layers == 0
+
+
+def test_kernel_patched_only_once_per_instance():
+    module = _make_fp8_toy_module(16, 32)
+    kernel = module.quant_method.fp8_linear
+    original = kernel._get_layer_params
+
+    prepare_fp8_layers_for_fsdp(module)
+    first_patched = kernel._get_layer_params
+
+    # Second call with same kernel id must not stack patches
+    n_rewritten_layers = prepare_fp8_layers_for_fsdp(module)
+    assert n_rewritten_layers == 0
+    assert kernel._get_layer_params is first_patched
+    assert kernel._get_layer_params is not original