[Feature][Diffusion] Allow online INT8 quantization with DLO AllGather (#6573)

Signed-off-by: brandneway <gyuan4892@gmail.com>

docs/design/feature/offloader/distributed_layerwise_offload.md

@@ -26,6 +26,8 @@ Legend: ✅ supported, ⚠️ compatibility path or limited validation, ❌ unsu
 | **TP > 1** | ⚠️ Ordinary TP-aware loader only; no direct checkpoint mmap. | ⚠️ Ordinary TP-aware loader only; no direct checkpoint mmap. |
 | **HSDP** | ❌ Rejected to avoid double-sharding parameters. | ⚠️ Limited end-to-end coverage. |
 | **Per-tensor online FP8 linears** | ✅ Ordinary loader finalizes weights and scales before DLO sharding. | ✅ Ordinary loader retains complete rank-local tensors. |
+| **Online INT8 linears** | ✅ Ordinary loader finalizes the int8 weight (contiguous or transposed view) and fp32 scale before DLO sharding. | ✅ Ordinary loader retains complete rank-local tensors. |
+| **Online MXFP8 linears** | ✅ Ordinary loader finalizes the fp8 weight (contiguous) and e8m0 block scale (contiguous or transposed view) before DLO sharding. | ✅ Ordinary loader retains complete rank-local tensors. |
 | **Other online quantization methods** | ❌ Rejected until runtime packing and scale layouts are validated. | ⚠️ Allowed through the ordinary loader; validation is method-specific. |
 | **Model-level or standard layerwise CPU offload** | ❌ Disabled because DLO takes priority. | ❌ Disabled because DLO takes priority. |
 | **Resident leading layers** | ❌ Rejected. | ✅ Requires eligible resident paths in the model's `OffloadPlan`. |
@@ -438,11 +440,15 @@ requirements.
 
 Direct checkpoint mmap can back either transfer path. It is currently limited
 to proven TP1, non-HSDP, non-online-quantized layouts. Other layouts use the
-ordinary loader. Per-tensor online FP8 linears can use DLO AllGather after the
-ordinary loader finalizes their runtime weights and scales; DLO then shards and
-reconstructs those tensors with their recorded layouts. Other online methods
-must use `--dlo-no-use-allgather` or disable online quantization until their
-runtime layouts are validated.
+ordinary loader. Per-tensor online FP8, online INT8, and online MXFP8 linears
+can use DLO AllGather after the ordinary loader finalizes their runtime
+weights and scales; DLO then shards and reconstructs those tensors with their
+recorded layouts (they all keep plain 1-byte dtypes over ordinary strided
+views: online INT8 has an int8 weight plus fp32 scale, online MXFP8 has an
+fp8 weight plus e8m0 block scale, each with a contiguous or transposed-view
+finalized tensor, all covered by the same physical-order packing as online
+FP8). Other online methods must use `--dlo-no-use-allgather` or disable
+online quantization until their runtime layouts are validated.
 
 The Host Weight Runtime representation, publication, and no-AllGather consumer
 contracts are merged; see
@@ -459,8 +465,8 @@ Current source-level validation includes:
 - HSDP + DLO without AllGather acceptance at configuration level;
 - loader preflight fallback for TP, HSDP, online quantization, unknown custom
   loaders, missing keys, and shape/dtype mismatches;
-- ordinary-loader fallback for per-tensor online FP8 linears followed by DLO
-  sharding of finalized weights and scales;
+- ordinary-loader fallback for per-tensor online FP8, online INT8, and online
+  MXFP8 linears followed by DLO sharding of finalized weights and scales;
 - exact loader-to-backend plan transfer and ordinary-loader fallback;
 - ordered publication hashing, overlapped durability, unordered parallel
   checksum fallback, and node-local source-digest reuse/invalidation;

tests/diffusion/model_loader/test_diffusers_loader.py

@@ -273,12 +273,15 @@ def test_required_hwr_miss_fails_before_ordinary_loading_or_publication(
     assert loader.take_host_weight_plan() is None
 
 
-def _make_dlo_online_quant_config() -> OmniDiffusionConfig:
+def _make_dlo_online_quant_config(dp_size: int = 2) -> OmniDiffusionConfig:
     return OmniDiffusionConfig(
         model="",
         dtype=torch.float32,
         quantization_config="fp8",
-        parallel_config=DiffusionParallelConfig(data_parallel_size=2),
+        parallel_config=DiffusionParallelConfig(
+            data_parallel_size=dp_size,
+            sequence_parallel_size=1,
+        ),
         enable_distributed_layerwise_offload=True,
         dlo_use_allgather=True,
     )
@@ -824,6 +827,93 @@ def test_dlo_allgather_online_fp8_uses_ordinary_loader(monkeypatch):
     assert loader.take_host_weight_plan() is None
 
 
+def test_dlo_allgather_online_int8_uses_ordinary_loader(monkeypatch):
+    import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
+    from vllm_omni.quantization.int8_config import NPUInt8OnlineLinearMethod
+
+    od_config = _make_dlo_online_quant_config()
+    loader = DiffusersPipelineLoader(LoadConfig(), od_config)
+    model = nn.Module()
+    model.transformer = nn.Linear(2, 2, bias=False)
+    model.transformer.quant_method = object.__new__(NPUInt8OnlineLinearMethod)
+    model.transformer.quant_method.uses_meta_device = True
+    calls: list[object] = []
+    allowlist_models: list[nn.Module] = []
+
+    original_allowlist_check = loader._unsupported_dlo_allgather_online_quant_methods
+
+    def check_allowlist(candidate: nn.Module) -> tuple[str, ...]:
+        allowlist_models.append(candidate)
+        return original_allowlist_check(candidate)
+
+    loader._init_from_load_format = lambda *_args, **_kwargs: model  # type: ignore[method-assign]
+    monkeypatch.setattr(loader, "_unsupported_dlo_allgather_online_quant_methods", check_allowlist)
+    loader._request_offload_after_quant = lambda _model: 1  # type: ignore[method-assign]
+    loader.load_weights = (  # type: ignore[method-assign]
+        lambda _model, *, stream_online_quant_to_cpu=False: calls.append(("load", stream_online_quant_to_cpu))
+    )
+    loader._process_weights_after_loading = lambda *_args: calls.append("process")  # type: ignore[method-assign]
+    loader._apply_skip_softmax_calibration = lambda _model: None  # type: ignore[method-assign]
+    monkeypatch.setattr(
+        loader_mod,
+        "build_checkpoint_mmap_plan",
+        lambda *_args, **_kwargs: HostWeightPlanResult(
+            None,
+            "online quantization requires the ordinary loader",
+        ),
+    )
+
+    assert loader.load_model(load_device="cpu", device=torch.device("cpu")) is model
+    assert allowlist_models == [model]
+    assert calls == [("load", True), "process"]
+    assert loader.take_host_weight_plan() is None
+
+
+def test_dlo_online_quant_group_size_one_skips_allgather_gate(monkeypatch):
+    """A DLO group of one runs no weight collective, so an otherwise
+    unvalidated online method must not be rejected by the AllGather gate."""
+
+    class UnsupportedOnlineMethod:
+        uses_meta_device = True
+
+    import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
+
+    od_config = _make_dlo_online_quant_config(dp_size=1)
+    loader = DiffusersPipelineLoader(LoadConfig(), od_config)
+    model = nn.Module()
+    model.transformer = nn.Linear(2, 2, bias=False)
+    model.transformer.quant_method = UnsupportedOnlineMethod()
+    calls: list[object] = []
+    allowlist_models: list[nn.Module] = []
+
+    original_allowlist_check = loader._unsupported_dlo_allgather_online_quant_methods
+
+    def check_allowlist(candidate: nn.Module) -> tuple[str, ...]:
+        allowlist_models.append(candidate)
+        return original_allowlist_check(candidate)
+
+    loader._init_from_load_format = lambda *_args, **_kwargs: model  # type: ignore[method-assign]
+    monkeypatch.setattr(loader, "_unsupported_dlo_allgather_online_quant_methods", check_allowlist)
+    loader._request_offload_after_quant = lambda _model: 1  # type: ignore[method-assign]
+    loader.load_weights = (  # type: ignore[method-assign]
+        lambda _model, *, stream_online_quant_to_cpu=False: calls.append(("load", stream_online_quant_to_cpu))
+    )
+    loader._process_weights_after_loading = lambda *_args: calls.append("process")  # type: ignore[method-assign]
+    loader._apply_skip_softmax_calibration = lambda _model: None  # type: ignore[method-assign]
+    monkeypatch.setattr(
+        loader_mod,
+        "build_checkpoint_mmap_plan",
+        lambda *_args, **_kwargs: HostWeightPlanResult(
+            None,
+            "online quantization requires the ordinary loader",
+        ),
+    )
+
+    assert loader.load_model(load_device="cpu", device=torch.device("cpu")) is model
+    assert allowlist_models == []
+    assert calls == [("load", True), "process"]
+
+
 def test_dlo_allgather_rejects_unvalidated_online_quant_method(monkeypatch):
     import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
 
@@ -847,10 +937,54 @@ def test_dlo_allgather_rejects_unvalidated_online_quant_method(monkeypatch):
         ),
     )
 
-    with pytest.raises(ValueError, match="per-tensor FP8 linears"):
+    with pytest.raises(ValueError, match="per-tensor FP8, INT8, and MXFP8 linears"):
         loader.load_model(load_device="cpu")
 
 
+def test_dlo_allgather_online_mxfp8_uses_ordinary_loader(monkeypatch):
+    mxfp8_config = pytest.importorskip("vllm_omni.quantization.mxfp8_config")
+    NPUMxfp8OnlineLinearMethod = mxfp8_config.NPUMxfp8OnlineLinearMethod
+
+    import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
+
+    od_config = _make_dlo_online_quant_config()
+    loader = DiffusersPipelineLoader(LoadConfig(), od_config)
+    model = nn.Module()
+    model.transformer = nn.Linear(2, 2, bias=False)
+    model.transformer.quant_method = object.__new__(NPUMxfp8OnlineLinearMethod)
+    model.transformer.quant_method.uses_meta_device = True
+    calls: list[object] = []
+    allowlist_models: list[nn.Module] = []
+
+    original_allowlist_check = loader._unsupported_dlo_allgather_online_quant_methods
+
+    def check_allowlist(candidate: nn.Module) -> tuple[str, ...]:
+        allowlist_models.append(candidate)
+        return original_allowlist_check(candidate)
+
+    loader._init_from_load_format = lambda *_args, **_kwargs: model  # type: ignore[method-assign]
+    monkeypatch.setattr(loader, "_unsupported_dlo_allgather_online_quant_methods", check_allowlist)
+    loader._request_offload_after_quant = lambda _model: 1  # type: ignore[method-assign]
+    loader.load_weights = (  # type: ignore[method-assign]
+        lambda _model, *, stream_online_quant_to_cpu=False: calls.append(("load", stream_online_quant_to_cpu))
+    )
+    loader._process_weights_after_loading = lambda *_args: calls.append("process")  # type: ignore[method-assign]
+    loader._apply_skip_softmax_calibration = lambda _model: None  # type: ignore[method-assign]
+    monkeypatch.setattr(
+        loader_mod,
+        "build_checkpoint_mmap_plan",
+        lambda *_args, **_kwargs: HostWeightPlanResult(
+            None,
+            "online quantization requires the ordinary loader",
+        ),
+    )
+
+    assert loader.load_model(load_device="cpu", device=torch.device("cpu")) is model
+    assert allowlist_models == [model]
+    assert calls == [("load", True), "process"]
+    assert loader.take_host_weight_plan() is None
+
+
 def test_hsdp_processes_quantized_weights_before_sharding(mocker):
     import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
     from vllm_omni.diffusion.offloader.module_collector import PipelineModules

tests/diffusion/offloader/test_distributed_layerwise_backend.py

@@ -245,6 +245,160 @@ class TestDistributedLayerwiseOffloadHook:
         )
         torch.testing.assert_close(rank0_block.weight_scale, expected_scale)
 
+    def test_allgather_reconstructs_online_int8_weight_and_scale(
+        self,
+        patched_offload_runtime,
+        monkeypatch,
+    ):
+        """DLO must reconstruct finalized online-INT8 tensors byte-for-byte.
+
+        One block carries both finalized online-INT8 layouts: the NPU method
+        leaves a contiguous pre-transposed (K, N) int8 weight with a 1-D fp32
+        per-channel scale, while the CUDA method leaves a transposed-view
+        weight (stride (1, K)) with a scalar fp32 scale.  The int8 and fp32
+        dtype groups must survive the flat-shard AllGather independently.
+        """
+
+        class OnlineInt8Block(nn.Module):
+            def __init__(self):
+                super().__init__()
+                logical = torch.arange(1.0, 13.0).reshape(3, 4)
+                # NPU online int8: qweight.t().contiguous() -> (K, N), 1-D scale
+                self.npu_weight = nn.Parameter(logical.t().contiguous().to(torch.int8), requires_grad=False)
+                self.npu_weight_scale = nn.Parameter(torch.tensor([2.0, 3.0, 4.0]), requires_grad=False)
+                # CUDA online int8: Parameter(qweight.t().data) -> transposed
+                # view (stride (1, K)) over the (N, K) storage, scalar scale
+                self.cuda_weight = nn.Parameter((logical + 100).to(torch.int8).t(), requires_grad=False)
+                self.cuda_weight_scale = nn.Parameter(torch.tensor([0.5]), requires_grad=False)
+
+        current_block = OnlineInt8Block()
+        rank0_block = OnlineInt8Block()
+        rank1_block = OnlineInt8Block()
+        expected = {name: tensor.detach().clone() for name, tensor in rank0_block.named_parameters()}
+        expected_strides = {name: tensor.stride() for name, tensor in rank0_block.named_parameters()}
+
+        def make_hook(block: OnlineInt8Block, rank: int) -> DistributedLayerwiseOffloadHook:
+            return DistributedLayerwiseOffloadHook(
+                next_block=block,
+                device=torch.device("cpu"),
+                dp_group=object(),
+                dp_size=2,
+                rank=rank,
+                copy_stream=DummyStream(),
+                comm_stream=DummyStream(),
+                pin_memory=False,
+            )
+
+        rank0_hook = make_hook(rank0_block, 0)
+        rank1_hook = make_hook(rank1_block, 1)
+        rank0_hook.initialize_hook(current_block)
+        rank1_hook.initialize_hook(OnlineInt8Block())
+        rank0_hook.gpu_shard_buffers = [
+            {dtype: torch.empty_like(shard) for dtype, shard in rank0_hook.cpu_shards.items()} for _ in range(2)
+        ]
+
+        def fake_allgather(output, local_shard, *, group):
+            del group
+            remote_shard = rank1_hook.cpu_shards[local_shard.dtype]
+            shard_size = local_shard.numel()
+            output[:shard_size].copy_(local_shard)
+            output[shard_size : 2 * shard_size].copy_(remote_shard)
+
+        monkeypatch.setattr(torch.distributed, "all_gather_into_tensor", fake_allgather)
+        rank0_hook.prefetch_layer(slot=0, non_blocking=False)
+
+        for name, tensor in rank0_block.named_parameters():
+            assert tensor.stride() == expected_strides[name]
+            torch.testing.assert_close(
+                tensor.float(),
+                expected[name].float(),
+                rtol=0,
+                atol=0,
+            )
+
+    def test_allgather_reconstructs_online_mxfp8_weight_and_scale(
+        self,
+        patched_offload_runtime,
+        monkeypatch,
+    ):
+        """DLO must reconstruct finalized online-MXFP8 tensors byte-for-byte.
+
+        One block carries both finalized online-MXFP8 layouts: the NPU method
+        leaves a contiguous pre-transposed (K, N) fp8 weight with a contiguous
+        (K_groups/2, N, 2) e8m0 block scale, while the vLLM-kernel method
+        (XPU) leaves an (N, K) fp8 weight whose e8m0 scale is a transposed
+        view over a contiguous (K/32, N) buffer.  The fp8 and e8m0 dtype
+        groups must survive the flat-shard AllGather independently.
+        """
+
+        class OnlineMxfp8Block(nn.Module):
+            def __init__(self):
+                super().__init__()
+                logical = torch.arange(1.0, 13.0).reshape(3, 4)
+                # Build scales from raw e8m0 bit patterns (1 byte/elem).
+                npu_scale_bits = torch.arange(1, 13, dtype=torch.uint8).reshape(2, 3, 2)
+                kernel_scale_bits = torch.arange(1, 7, dtype=torch.uint8).reshape(2, 3)
+                # NPU online mxfp8: (N, K).t().contiguous() -> (K, N) fp8
+                # weight, contiguous (K_groups/2, N, 2) e8m0 scale
+                self.npu_weight = nn.Parameter(logical.t().contiguous().to(torch.float8_e4m3fn), requires_grad=False)
+                self.npu_weight_scale = nn.Parameter(npu_scale_bits.view(torch.float8_e8m0fnu), requires_grad=False)
+                # vLLM-kernel online mxfp8 (XPU): (N, K) fp8 weight, e8m0
+                # scale kept as a transposed view (stride (1, K/32)) over the
+                # contiguous (K/32, N) storage
+                self.kernel_weight = nn.Parameter((logical + 100).to(torch.float8_e4m3fn), requires_grad=False)
+                self.kernel_weight_scale = nn.Parameter(
+                    kernel_scale_bits.view(torch.float8_e8m0fnu).t(), requires_grad=False
+                )
+
+        current_block = OnlineMxfp8Block()
+        rank0_block = OnlineMxfp8Block()
+        rank1_block = OnlineMxfp8Block()
+        expected = {name: tensor.detach().clone() for name, tensor in rank0_block.named_parameters()}
+        expected_strides = {name: tensor.stride() for name, tensor in rank0_block.named_parameters()}
+
+        def make_hook(block: OnlineMxfp8Block, rank: int) -> DistributedLayerwiseOffloadHook:
+            return DistributedLayerwiseOffloadHook(
+                next_block=block,
+                device=torch.device("cpu"),
+                dp_group=object(),
+                dp_size=2,
+                rank=rank,
+                copy_stream=DummyStream(),
+                comm_stream=DummyStream(),
+                pin_memory=False,
+            )
+
+        rank0_hook = make_hook(rank0_block, 0)
+        rank1_hook = make_hook(rank1_block, 1)
+        rank0_hook.initialize_hook(current_block)
+        rank1_hook.initialize_hook(OnlineMxfp8Block())
+        rank0_hook.gpu_shard_buffers = [
+            {dtype: torch.empty_like(shard) for dtype, shard in rank0_hook.cpu_shards.items()} for _ in range(2)
+        ]
+
+        def fake_allgather(output, local_shard, *, group):
+            del group
+            remote_shard = rank1_hook.cpu_shards[local_shard.dtype]
+            shard_size = local_shard.numel()
+            output[:shard_size].copy_(local_shard)
+            output[shard_size : 2 * shard_size].copy_(remote_shard)
+
+        monkeypatch.setattr(torch.distributed, "all_gather_into_tensor", fake_allgather)
+        rank0_hook.prefetch_layer(slot=0, non_blocking=False)
+
+        for name, tensor in rank0_block.named_parameters():
+            assert tensor.stride() == expected_strides[name]
+            # Compare raw bytes in logical order: float8/e8m0 CPU upcasts are
+            # version-sensitive, while contiguous().view(uint8) is an exact
+            # byte-level comparison for both contiguous and transposed-view
+            # parameters.
+            torch.testing.assert_close(
+                tensor.detach().contiguous().view(torch.uint8),
+                expected[name].contiguous().view(torch.uint8),
+                rtol=0,
+                atol=0,
+            )
+
     def test_dtensor_wrapper_preserved_across_prefetch_and_offload(self, dist_group, patched_offload_runtime):
         """DTensor wrapper should be preserved through prefetch/offload cycle."""
         current_block = TinyBlock(_make_values(1.0))

vllm_omni/diffusion/model_loader/diffusers_loader.py

@@ -263,7 +263,7 @@ class DiffusersPipelineLoader(HWRLoaderMixin):
             allow_patterns = allow_patterns_overrides
 
         if not is_local and indexed_weight_files is not None:
-            hf_folder = download_weights_from_hf_specific(
+            hf_folder: Path | str = download_weights_from_hf_specific(
                 model_name_or_path=str(model_name_or_path),
                 cache_dir=self.load_config.download_dir,
                 allow_patterns=[self._repo_relative_path(subfolder, filename) for filename in indexed_weight_files],
@@ -357,8 +357,9 @@ class DiffusersPipelineLoader(HWRLoaderMixin):
 
     def _get_source_quant_config(self, source: "ComponentSource") -> object | None:
         quant_config = self.quant_config
-        if hasattr(quant_config, "resolve"):
-            return quant_config.resolve(source.prefix.rstrip("."))
+        resolve = getattr(quant_config, "resolve", None)
+        if resolve is not None:
+            return resolve(source.prefix.rstrip("."))
         return quant_config
 
     def _get_checkpoint_adapter(
@@ -562,19 +563,19 @@ class DiffusersPipelineLoader(HWRLoaderMixin):
                     )
                     self.host_weight_plan = plan_result.plan
 
-                _skip_load = self.host_weight_plan is not None
+                host_weight_plan = self.host_weight_plan
 
-                if _skip_load:
+                if host_weight_plan is not None:
                     logger.info(
                         "DLO host-weight plan active (%s, %s): skipping ordinary materialization for %s",
                         "AllGather" if _use_ag and _dlo_group_size > 1 else "rank-local",
-                        self.host_weight_plan.backing_kind,
-                        sorted(self.host_weight_plan.planned_source_prefixes) or "legacy DiT sources",
+                        host_weight_plan.backing_kind,
+                        sorted(host_weight_plan.planned_source_prefixes) or "legacy DiT sources",
                     )
                     ordinary_sources = tuple(
                         source
                         for source in weight_sources
-                        if source.prefix not in self.host_weight_plan.planned_source_prefixes
+                        if source.prefix not in host_weight_plan.planned_source_prefixes
                     )
                     if ordinary_sources:
                         logger.info(
@@ -584,21 +585,25 @@ class DiffusersPipelineLoader(HWRLoaderMixin):
                         self.load_weights(
                             model,
                             sources=ordinary_sources,
-                            planned_weights=self.host_weight_plan.bindings,
+                            planned_weights=host_weight_plan.bindings,
                         )
                 else:
-                    if _dist_offload and _use_ag and _has_online_quant:
+                    if _dist_offload and _use_ag and _dlo_group_size > 1 and _has_online_quant:
+                        # An effective DLO group size of one performs no weight
+                        # collective, so the AllGather layout allowlist does not
+                        # apply there even with dlo_use_allgather=True.
                         unsupported_methods = self._unsupported_dlo_allgather_online_quant_methods(model)
                         if unsupported_methods:
                             raise ValueError(
                                 "DLO+AllGather supports online quantization only for "
-                                "per-tensor FP8 linears; unsupported online methods: "
-                                f"{', '.join(unsupported_methods)}. Please use "
+                                "per-tensor FP8, INT8, and MXFP8 linears; unsupported online "
+                                f"methods: {', '.join(unsupported_methods)}. Please use "
                                 "--dlo-no-use-allgather or disable online quantization."
                             )
                         logger.info(
-                            "Online per-tensor FP8 with DLO+AllGather: using the "
-                            "ordinary loader before sharding finalized weights and scales"
+                            "Validated online methods (per-tensor FP8, INT8, MXFP8) with "
+                            "DLO+AllGather: using the ordinary loader before sharding "
+                            "finalized weights and scales"
                         )
                     if _dist_offload and plan_result is not None:
                         logger.info(
@@ -678,9 +683,10 @@ class DiffusersPipelineLoader(HWRLoaderMixin):
         marked = 0
         for module in model.modules():
             quant_method = getattr(module, "quant_method", None)
-            if getattr(quant_method, "supports_offload_after_quant", False):
-                quant_method.enable_offload_after_quant()
-                marked += 1
+            if quant_method is None or not getattr(quant_method, "supports_offload_after_quant", False):
+                continue
+            quant_method.enable_offload_after_quant()
+            marked += 1
         return marked
 
     @staticmethod
@@ -694,22 +700,62 @@ class DiffusersPipelineLoader(HWRLoaderMixin):
     def _unsupported_dlo_allgather_online_quant_methods(model: nn.Module) -> tuple[str, ...]:
         """Return unsupported online-quant methods for DLO AllGather.
 
-        Per-tensor online FP8 is safe after the ordinary loader has finalized
-        its weight and scale parameters. DLO shards those runtime tensors by
-        dtype and reconstructs their recorded shapes and strides before the
-        kernel consumes them. Other online methods may create different scale,
-        packing, or aliasing layouts and remain fail-closed until validated.
+        Per-tensor online FP8, online INT8, and online MXFP8 are safe after
+        the ordinary loader has finalized their weight and scale parameters.
+        DLO shards those runtime tensors by dtype and reconstructs their
+        recorded shapes and strides before the kernel consumes them. They all
+        keep plain transportable 1-byte dtypes over ordinary strided views:
+
+        - online INT8: int8 weight plus fp32 scale, either contiguous (NPU,
+          pre-transposed (K, N)) or a transposed view (CUDA, stride (1, K));
+        - online MXFP8: fp8 weight plus e8m0 block scale, either contiguous
+          (NPU: (K, N) weight with (K_groups/2, N, 2) scale) or with the
+          scale stored as a transposed view (vLLM kernel, .t() over a
+          contiguous (K/32, N) buffer).
+
+        Both shape families are already covered by the physical-order packing
+        that online FP8 requires. Other online methods may create different
+        scale, packing, or aliasing layouts (e.g. dual-scale fp4 pairs,
+        swizzled or NZ hardware formats) and remain fail-closed until
+        validated.
         """
         from vllm.model_executor.layers.quantization.online.fp8 import (
             Fp8PerTensorOnlineLinearMethod,
         )
 
+        from vllm_omni.quantization.int8_config import (
+            Int8OnlineLinearMethod,
+            NPUInt8OnlineLinearMethod,
+        )
+
+        try:
+            from vllm_omni.quantization.mxfp8_config import (
+                NPUMxfp8OnlineLinearMethod,
+                VllmMxfp8OnlineLinearMethod,
+            )
+
+            mxfp8_online_methods: tuple[type, ...] = (
+                NPUMxfp8OnlineLinearMethod,
+                VllmMxfp8OnlineLinearMethod,
+            )
+        except ImportError:
+            # MXFP8 requires a vLLM build with MXFP8 kernel support; treat it
+            # as absent when the module cannot be imported.
+            mxfp8_online_methods = ()
+
+        allowed_online_methods: tuple[type, ...] = (
+            Fp8PerTensorOnlineLinearMethod,
+            Int8OnlineLinearMethod,
+            NPUInt8OnlineLinearMethod,
+            *mxfp8_online_methods,
+        )
+
         unsupported: set[str] = set()
         for module in model.modules():
             quant_method = getattr(module, "quant_method", None)
             if not getattr(quant_method, "uses_meta_device", False):
                 continue
-            if not isinstance(quant_method, Fp8PerTensorOnlineLinearMethod):
+            if not isinstance(quant_method, allowed_online_methods):
                 unsupported.add(type(quant_method).__name__)
         return tuple(sorted(unsupported))