diff --git a/docs/design/feature/offloader/distributed_layerwise_offload.md b/docs/design/feature/offloader/distributed_layerwise_offload.md index c47b9187b..2b5f929ed 100644 --- a/docs/design/feature/offloader/distributed_layerwise_offload.md +++ b/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; diff --git a/tests/diffusion/model_loader/test_diffusers_loader.py b/tests/diffusion/model_loader/test_diffusers_loader.py index 09acae72d..bd59ee94b 100644 --- a/tests/diffusion/model_loader/test_diffusers_loader.py +++ b/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 diff --git a/tests/diffusion/offloader/test_distributed_layerwise_backend.py b/tests/diffusion/offloader/test_distributed_layerwise_backend.py index ca44f321d..934fe3299 100644 --- a/tests/diffusion/offloader/test_distributed_layerwise_backend.py +++ b/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)) diff --git a/vllm_omni/diffusion/model_loader/diffusers_loader.py b/vllm_omni/diffusion/model_loader/diffusers_loader.py index d582270cc..7069e5012 100644 --- a/vllm_omni/diffusion/model_loader/diffusers_loader.py +++ b/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))