Add normalized online FP8 cache for DLO

Signed-off-by: hsliuustc0106 <hsliuustc0106@users.noreply.github.com>

docs/design/feature/offloader/distributed_layerwise_offload.md

@@ -105,6 +105,31 @@ flags such as `_supports_mmap_loading` or parameter attributes for mmap-only
 transforms. Model-specific direct-layout knowledge, when required, lives in a
 checkpoint adapter beside the ordinary loader.
 
+### Phase-I normalized online-FP8 cache
+
+The direct-checkpoint plan above cannot quantize a BF16 source online. Phase I
+adds an explicit, opt-in normalized cache for the narrower case of one TP1
+transformer with per-tensor online FP8 and DLO AllGather:
+
+```text
+cache miss:  BF16 source -> ordinary online FP8 loader -> atomic FP8 cache
+cache hit:   FP8 cache   -> direct mmap plan -> DLO host shard -> AllGather
+```
+
+The cache stores the serialized runtime representation, including finalized FP8
+weights, generated scales, and persistent transformer buffers. Its manifest
+records the schema, component, quantization contract, tensor count, source
+files, and a source fingerprint. A missing, stale, incomplete, incompatible, or
+invalid artifact is a cache miss; cache publication is best-effort and falls
+back to the ordinary loader on failure.
+
+Phase I keeps the cache binding deliberately simple: one source per transformer
+component, with existing per-tensor transforms applied before publication. The
+cache is not selected for DLO-disabled, no-AllGather, TP>1, or HSDP deployments,
+and it normalizes the transformer only. Cache sharing and compatibility for
+other components, transformed TP coordinates, and broader DP/SP layouts remain
+in [RFC #6231](https://github.com/vllm-project/vllm-omni/issues/6231).
+
 ### AllGather path
 
 With the default `dlo_use_allgather=True`, each rank stores approximately
@@ -211,10 +236,10 @@ 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.
 
-A normalized runtime mmap cache, built through the ordinary loader, is the
+The Phase-I normalized online-FP8 cache is the supported TP1 transformer case.
-proposed general mechanism for sharing transformed TP or quantized layouts.
+The general cache-compatibility contract for additional components, transformed
-That cache and its publication/lifecycle protocol are intentionally outside
+TP coordinates, and cross-DP/SP reuse remains in
-this phase; see [RFC #6195](https://github.com/vllm-project/vllm-omni/issues/6195).
+[RFC #6231](https://github.com/vllm-project/vllm-omni/issues/6231).
 
 ## Validation coverage
 

docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md

@@ -42,6 +42,12 @@ vllm serve /path/to/model --omni \
 vllm serve /path/to/model --omni \
   --enable-distributed-layerwise-offload \
   --usp 4
+
+# Opt-in normalized online-FP8 cache (TP1 + DLO AllGather)
+vllm serve /path/to/model --omni \
+  --quantization fp8 \
+  --enable-distributed-layerwise-offload \
+  --dlo-fp8-cache-dir /path/to/node-local/fp8-cache
 ```
 
 ```python
@@ -63,6 +69,7 @@ omni = Omni(
 | `--dlo-use-allgather` | Shard host weights and reconstruct with AllGather | `true` |
 | `--dlo-no-use-allgather` | Stream complete rank-local blocks without a DLO weight collective | `false` |
 | `--dlo-resident-layers N` | Keep N leading main-DiT blocks on device; requires no-AllGather and model-declared resident paths | `0` |
+| `--dlo-fp8-cache-dir PATH` | Opt into the Phase-I normalized online-FP8 cache for the transformer | unset |
 
 ## Host-weight loading
 
@@ -104,6 +111,24 @@ process.
 When the effective DLO group size is one, `dlo_use_allgather=True` does not
 perform a collective and uses the same rank-local transfer behavior.
 
+## Normalized online-FP8 cache
+
+`--dlo-fp8-cache-dir` is an opt-in Phase-I cache for per-tensor online FP8 with
+TP1 and DLO AllGather. On a cache miss, the ordinary loader quantizes the
+transformer once and publishes finalized FP8 weights and scales as an atomic,
+serialized safetensors artifact. On a valid cache hit, the loader selects that
+serialized source before DLO planning, so the transformer can be loaded through
+the direct checkpoint-mmap path without re-materializing the full BF16 source.
+
+The cache is keyed to the source files through the manifest's schema, component,
+quantization contract, and source fingerprint. Incomplete, incompatible, or
+stale artifacts are treated as cache misses. Cache creation is best-effort: a
+publication failure logs a warning and keeps the ordinary online-FP8 path.
+
+Phase I normalizes the transformer only. Other pipeline components keep their
+ordinary loading and quantization behavior. The cache is not selected when DLO
+is disabled, AllGather is disabled, TP is greater than one, or HSDP is enabled.
+
 ## Declarative topology
 
 Models may declare an `OffloadPlan` instead of embedding offload logic:
@@ -147,10 +172,9 @@ must enter each collective.
   `OffloadPlan` that declares eligible `resident_dit_paths`.
 - DP concurrency requires an explicit, identical inference-step count.
 
-Sharing transformed TP or quantized runtime layouts through a normalized mmap
+Sharing transformed TP or quantized runtime layouts beyond this TP1 transformer
-cache is a follow-up design in
+cache remains a follow-up design in
-[RFC #6195](https://github.com/vllm-project/vllm-omni/issues/6195), not part of
+[RFC #6231](https://github.com/vllm-project/vllm-omni/issues/6231).
-the direct-checkpoint path.
 
 See the [Cosmos3 DistOffload recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/cosmos3/Cosmos3-DistOffload.md)
 for an end-to-end example.

tests/diffusion/model_loader/test_fp8_cache.py

@@ -0,0 +1,153 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Unit tests for the Phase-I normalized online-FP8 cache contract."""
+
+import json
+from types import SimpleNamespace
+
+import pytest
+import torch
+from safetensors.torch import load_file, save_file
+from torch import nn
+
+from vllm_omni.diffusion.model_loader.fp8_cache import (
+    CACHE_COMPLETE_NAME,
+    CACHE_COMPONENT_SUBFOLDER,
+    CACHE_MANIFEST_NAME,
+    _source_fingerprint,
+    _transformer_quant_config,
+    build_normalized_fp8_cache,
+    is_online_fp8_config,
+    load_fp8_cache_manifest,
+)
+
+pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
+
+
+def _write_valid_cache(tmp_path):
+    source = tmp_path / "source.safetensors"
+    save_file({"weight": torch.ones(2, 2)}, str(source))
+
+    cache = tmp_path / "cache"
+    component = cache / CACHE_COMPONENT_SUBFOLDER
+    component.mkdir(parents=True)
+    shard = component / "part-00000.safetensors"
+    save_file({"weight": torch.ones(2, 2, dtype=torch.float8_e4m3fn)}, str(shard))
+    (component / "model.safetensors.index.json").write_text(
+        json.dumps(
+            {
+                "metadata": {"total_size": 4},
+                "weight_map": {"weight": shard.name},
+            }
+        )
+    )
+
+    manifest = {
+        "schema_version": 1,
+        "cache_kind": "normalized_fp8_runtime",
+        "component": "transformer",
+        "component_subfolder": CACHE_COMPONENT_SUBFOLDER,
+        "source_prefix": "transformer.",
+        "quantization": {"method": "fp8", "serialized": True, "activation_scheme": "dynamic"},
+        "tensor_count": 1,
+        "source_files": [str(source)],
+        "source_fingerprint": _source_fingerprint({str(source)}),
+    }
+    cache.mkdir(exist_ok=True)
+    (cache / CACHE_MANIFEST_NAME).write_text(json.dumps(manifest))
+    (cache / CACHE_COMPLETE_NAME).write_text("phase1\n")
+    return cache, source, manifest
+
+
+def test_online_and_serialized_fp8_configs_are_classified():
+    from vllm.model_executor.layers.quantization.fp8 import Fp8Config
+
+    assert is_online_fp8_config(Fp8Config())
+    assert not is_online_fp8_config(Fp8Config(is_checkpoint_fp8_serialized=True))
+
+    cache_config = _transformer_quant_config()
+    assert not is_online_fp8_config(cache_config)
+    assert cache_config.resolve("transformer").is_checkpoint_fp8_serialized
+    assert cache_config.resolve("text_encoder") is None
+
+
+def test_load_fp8_cache_manifest_validates_artifact_and_sources(tmp_path):
+    cache, source, expected = _write_valid_cache(tmp_path)
+
+    assert load_fp8_cache_manifest(cache) == expected
+
+    source.write_bytes(source.read_bytes() + b"changed")
+    assert load_fp8_cache_manifest(cache) is None
+
+
+def test_load_fp8_cache_manifest_rejects_wrong_serialization_contract(tmp_path):
+    cache, _, manifest = _write_valid_cache(tmp_path)
+
+    manifest["quantization"]["serialized"] = False
+    (cache / CACHE_MANIFEST_NAME).write_text(json.dumps(manifest))
+
+    assert load_fp8_cache_manifest(cache) is None
+
+
+def test_load_fp8_cache_manifest_rejects_path_traversal(tmp_path):
+    cache, _, manifest = _write_valid_cache(tmp_path)
+    index_path = cache / CACHE_COMPONENT_SUBFOLDER / "model.safetensors.index.json"
+    index_path.write_text(json.dumps({"weight_map": {"weight": "../outside.safetensors"}}))
+
+    assert load_fp8_cache_manifest(cache) is None
+
+
+def test_load_fp8_cache_manifest_rejects_non_object_json(tmp_path):
+    cache, _, _ = _write_valid_cache(tmp_path)
+    (cache / CACHE_MANIFEST_NAME).write_text("[]")
+
+    assert load_fp8_cache_manifest(cache) is None
+
+
+def test_build_normalized_fp8_cache_serializes_weight_and_scale(tmp_path):
+    from vllm.model_executor.layers.quantization.online.fp8 import (
+        Fp8PerTensorOnlineLinearMethod,
+    )
+
+    source_dir = tmp_path / "source"
+    source_dir.mkdir()
+    source_file = source_dir / "model.safetensors"
+    save_file({"block.weight": torch.ones((3, 2), dtype=torch.bfloat16)}, str(source_file))
+
+    model = nn.Module()
+    model._dit_modules = ("transformer",)
+    model.transformer = nn.Module()
+    model.transformer.block = nn.Module()
+    model.transformer.block.weight = nn.Parameter(torch.ones((2, 3), dtype=torch.float8_e4m3fn), requires_grad=False)
+    model.transformer.block.weight_scale = nn.Parameter(torch.tensor([0.5]), requires_grad=False)
+    model.transformer.block.quant_method = object.__new__(Fp8PerTensorOnlineLinearMethod)
+
+    source = SimpleNamespace(
+        model_or_path=str(source_dir),
+        subfolder=None,
+        revision=None,
+        prefix="transformer.",
+    )
+    cache_dir = tmp_path / "cache"
+
+    manifest = build_normalized_fp8_cache(
+        model,
+        sources=(source,),
+        cache_dir=cache_dir,
+    )
+
+    assert manifest["cache_kind"] == "normalized_fp8_runtime"
+    tensors = load_file(str(cache_dir / CACHE_COMPONENT_SUBFOLDER / "part-00000.safetensors"))
+    scale_tensors = load_file(str(cache_dir / CACHE_COMPONENT_SUBFOLDER / "part-00001.safetensors"))
+    assert tensors["block.weight"].dtype == torch.float8_e4m3fn
+    assert tuple(tensors["block.weight"].shape) == (3, 2)
+    assert torch.equal(scale_tensors["block.weight_scale"], torch.tensor([0.5]))
+
+
+@pytest.mark.parametrize("missing", [CACHE_COMPLETE_NAME, CACHE_MANIFEST_NAME])
+def test_load_fp8_cache_manifest_requires_completion_marker(tmp_path, missing):
+    cache, _, _ = _write_valid_cache(tmp_path)
+    (cache / missing).unlink()
+
+    assert load_fp8_cache_manifest(cache) is None

tests/entrypoints/test_async_omni_diffusion_config.py

@@ -347,6 +347,8 @@ def test_serve_cli_forwards_distributed_offload_residency():
             "--dlo-no-use-allgather",
             "--dlo-resident-layers",
             "20",
+            "--dlo-fp8-cache-dir",
+            "/tmp/dlo-fp8-cache",
         ]
     )
 
@@ -357,9 +359,11 @@ def test_serve_cli_forwards_distributed_offload_residency():
     assert args.enable_distributed_layerwise_offload is True
     assert args.dlo_use_allgather is False
     assert args.dlo_resident_layers == 20
+    assert args.dlo_fp8_cache_dir == "/tmp/dlo-fp8-cache"
     assert engine_args["enable_distributed_layerwise_offload"] is True
     assert engine_args["dlo_use_allgather"] is False
     assert engine_args["dlo_resident_layers"] == 20
+    assert engine_args["dlo_fp8_cache_dir"] == "/tmp/dlo-fp8-cache"
 
 
 def test_serve_cli_accepts_diffusion_compile_controls():

vllm_omni/config/omni_config.py

@@ -693,6 +693,7 @@ class _DiffusionConfigProjection:
     enable_distributed_layerwise_offload: bool = False
     dlo_use_allgather: bool = True
     dlo_resident_layers: int = Field(default=0, ge=0)
+    dlo_fp8_cache_dir: str | None = None
     pin_cpu_memory: bool = True
     diffusion_compile_granularity: Literal["regional", "full"] = "regional"
     diffusion_compile_dynamic: bool = Field(default=True, strict=True)

vllm_omni/config/stage_config.py

@@ -383,6 +383,7 @@ class StageDeployConfig:
     enable_distributed_layerwise_offload: bool | None = None
     dlo_use_allgather: bool | None = None
     dlo_resident_layers: int | None = None
+    dlo_fp8_cache_dir: str | None = None
     # Diffusion-specific debug and observability knobs.
     enable_diffusion_pipeline_profiler: bool | None = None
 

vllm_omni/diffusion/data.py

@@ -752,6 +752,10 @@ class OmniDiffusionConfig:
     dlo_use_allgather: bool = True
     # Leading main-DiT blocks kept resident by distributed layerwise offload.
     dlo_resident_layers: int = 0
+    # Optional Phase-I normalized online-FP8 cache. A cache miss builds the
+    # artifact after the ordinary online loader; a valid cache hit uses the
+    # serialized FP8 checkpoint-mmap path.
+    dlo_fp8_cache_dir: str | None = None
 
     pin_cpu_memory: bool = True  # Use pinned memory for faster transfers when offloading
 

vllm_omni/diffusion/model_loader/checkpoint_adapters/direct_mmap.py

@@ -41,6 +41,8 @@ class _MiniMaxH3DirectMmapAdapter:
         target: torch.Tensor,
     ) -> DirectMmapTensorPolicy:
         del target
+        if getattr(self.pipeline, "_dlo_fp8_cache_normalized", False):
+            return DirectMmapTensorPolicy(allow_custom_loader=True)
         transform = None
         suffix = ".qkv_proj.weight"
         if runtime_name.endswith(suffix):

vllm_omni/diffusion/model_loader/diffusers_loader.py

@@ -8,7 +8,7 @@ import re
 import time
 from collections.abc import Generator, Iterable, Sequence
 from pathlib import Path
-from typing import cast
+from typing import Any, cast
 
 import torch
 from torch import nn
@@ -34,6 +34,12 @@ from vllm_omni.diffusion.distributed.hsdp import HSDPInferenceConfig, apply_hsdp
 from vllm_omni.diffusion.model_loader.checkpoint_adapters import (
     get_checkpoint_adapter,
 )
+from vllm_omni.diffusion.model_loader.fp8_cache import (
+    _transformer_quant_config,
+    build_normalized_fp8_cache,
+    is_online_fp8_config,
+    load_fp8_cache_manifest,
+)
 from vllm_omni.diffusion.model_loader.host_weight_plan import (
     HostWeightPlan,
     build_checkpoint_mmap_plan,
@@ -141,6 +147,10 @@ class DiffusersPipelineLoader:
         self.quant_config = od_config.quantization_config
         self.parallel_config = od_config.parallel_config
         self.host_weight_plan: HostWeightPlan | None = None
+        self._fp8_cache_dir: Path | None = None
+        self._fp8_cache_requested = False
+        self._fp8_cache_hit = False
+        self._fp8_cache_manifest: dict[str, Any] | None = None
 
     def take_host_weight_plan(self) -> HostWeightPlan | None:
         """Transfer the loader-produced plan to the offload backend."""
@@ -367,12 +377,61 @@ class DiffusersPipelineLoader:
         )
 
     def _get_weight_sources(self, model: nn.Module) -> tuple["ComponentSource", ...]:
-        return tuple(
+        sources = tuple(
             cast(
                 Iterable[DiffusersPipelineLoader.ComponentSource],
                 getattr(model, "weights_sources", ()),
             )
         )
+        if not self._fp8_cache_hit or self._fp8_cache_dir is None or self._fp8_cache_manifest is None:
+            return sources
+
+        cache_prefix = str(self._fp8_cache_manifest["source_prefix"])
+        cache_subfolder = str(self._fp8_cache_manifest.get("component_subfolder", "transformer"))
+        return tuple(
+            dataclasses.replace(
+                source,
+                model_or_path=str(self._fp8_cache_dir),
+                subfolder=cache_subfolder,
+                revision=None,
+            )
+            if source.prefix == cache_prefix
+            else source
+            for source in sources
+        )
+
+    def _prepare_fp8_cache(self) -> None:
+        """Select the Phase-I cache path before constructing the pipeline."""
+        cache_dir = getattr(self.od_config, "dlo_fp8_cache_dir", None)
+        if not cache_dir or not is_online_fp8_config(self.quant_config):
+            return
+        if not getattr(self.od_config, "enable_distributed_layerwise_offload", False):
+            logger.warning("Ignoring --dlo-fp8-cache-dir because distributed layerwise offload is disabled")
+            return
+        if not getattr(self.od_config, "dlo_use_allgather", True):
+            logger.warning("Ignoring --dlo-fp8-cache-dir because DLO AllGather is disabled")
+            return
+        if int(getattr(self.parallel_config, "tensor_parallel_size", 1)) != 1:
+            logger.warning("Ignoring --dlo-fp8-cache-dir because Phase-I cache requires TP=1")
+            return
+        if bool(getattr(self.parallel_config, "use_hsdp", False)):
+            logger.warning("Ignoring --dlo-fp8-cache-dir because Phase-I cache does not support HSDP")
+            return
+
+        self._fp8_cache_dir = Path(cache_dir)
+        self._fp8_cache_requested = True
+        self._fp8_cache_manifest = load_fp8_cache_manifest(self._fp8_cache_dir)
+        if self._fp8_cache_manifest is None:
+            logger.info("DLO normalized FP8 cache miss: %s", self._fp8_cache_dir)
+            return
+
+        self._fp8_cache_hit = True
+        self.quant_config = _transformer_quant_config()
+        self.od_config.quantization_config = self.quant_config
+        logger.info(
+            "DLO normalized FP8 cache hit: %s; selecting serialized FP8 checkpoint path",
+            self._fp8_cache_dir,
+        )
 
     def _get_expected_parameter_names(self, model: nn.Module) -> set[str]:
         """Return parameter names that should be covered by strict load checks."""
@@ -403,6 +462,7 @@ class DiffusersPipelineLoader:
         self.host_weight_plan = None
         if load_format is None:
             load_format = "default"
+        self._prepare_fp8_cache()
         # 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,
@@ -413,7 +473,9 @@ class DiffusersPipelineLoader:
             is_offline = getattr(quant_cfg, "data_type", None) == "mx_fp" or getattr(
                 quant_cfg, "is_checkpoint_quantized", False
             )
-            if not is_offline:
+            if self._fp8_cache_hit:
+                logger.info("Normalized FP8 cache is already serialized; keeping model initialization on CPU")
+            elif not is_offline:
                 load_device = device.type
                 offload_after_quant = True
                 logger.info(
@@ -431,6 +493,11 @@ class DiffusersPipelineLoader:
                 )
             else:
                 model = self._init_from_load_format(load_format, target_device, custom_pipeline_name, is_hsdp=False)
+                if self._fp8_cache_hit:
+                    # The direct-mmap adapter must see this before plan
+                    # construction so normalized QKV tensors bypass the raw
+                    # checkpoint grouped-QKV transform.
+                    setattr(model, "_dlo_fp8_cache_normalized", True)
 
                 _dist_offload = getattr(self.od_config, "enable_distributed_layerwise_offload", False)
                 _use_ag = getattr(self.od_config, "dlo_use_allgather", True)
@@ -520,6 +587,20 @@ class DiffusersPipelineLoader:
                 model.to("cpu")
                 logger.info("Quantization complete, offloaded model back to CPU")
 
+        if self._fp8_cache_requested and self._fp8_cache_dir is not None and not self._fp8_cache_hit:
+            try:
+                build_normalized_fp8_cache(
+                    model,
+                    sources=self._get_weight_sources(model),
+                    cache_dir=self._fp8_cache_dir,
+                )
+            except Exception as exc:
+                logger.warning(
+                    "DLO normalized FP8 cache build failed; retaining ordinary-loader path: %s",
+                    exc,
+                    exc_info=True,
+                )
+
         self._apply_skip_softmax_calibration(model)
         return model.eval()
 

vllm_omni/diffusion/model_loader/fp8_cache.py

@@ -0,0 +1,349 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Phase-I normalized FP8 cache support for DLO checkpoint mmap."""
+
+from __future__ import annotations
+
+import hashlib
+import json
+import os
+import shutil
+import uuid
+from pathlib import Path
+from typing import Any
+
+import torch
+from safetensors import safe_open
+from safetensors.torch import save_file
+from torch import nn
+from vllm.logger import init_logger
+
+from vllm_omni.diffusion.model_loader.host_weight_plan import (
+    _build_source_map,
+    _collect_required_dit_tensors,
+)
+from vllm_omni.quantization.component_config import ComponentQuantizationConfig
+
+logger = init_logger(__name__)
+
+CACHE_SCHEMA_VERSION = 1
+CACHE_KIND = "normalized_fp8_runtime"
+CACHE_MANIFEST_NAME = "manifest.json"
+CACHE_COMPLETE_NAME = "COMPLETE"
+CACHE_COMPONENT_SUBFOLDER = "transformer"
+
+
+def _transformer_quant_config() -> ComponentQuantizationConfig:
+    """Build the serialized FP8 config used by a normalized cache."""
+    from vllm.model_executor.layers.quantization.fp8 import Fp8Config
+
+    return ComponentQuantizationConfig(
+        {
+            "transformer": Fp8Config(
+                is_checkpoint_fp8_serialized=True,
+                activation_scheme="dynamic",
+            )
+        }
+    )
+
+
+def is_online_fp8_config(quant_config: Any) -> bool:
+    """Return whether the transformer component requests online FP8."""
+    if quant_config is None:
+        return False
+    if hasattr(quant_config, "resolve"):
+        quant_config = quant_config.resolve("transformer")
+    get_name = getattr(quant_config, "get_name", None)
+    if not callable(get_name) or str(get_name()).lower() != "fp8":
+        return False
+    return not bool(getattr(quant_config, "is_checkpoint_fp8_serialized", False))
+
+
+def _source_fingerprint(files: set[str]) -> str:
+    entries: list[dict[str, int | str]] = []
+    for filename in sorted(files):
+        stat = os.stat(filename)
+        entries.append(
+            {
+                "path": os.path.abspath(filename),
+                "size": stat.st_size,
+                "mtime_ns": stat.st_mtime_ns,
+            }
+        )
+    payload = json.dumps(entries, sort_keys=True, separators=(",", ":")).encode()
+    return hashlib.sha256(payload).hexdigest()
+
+
+def load_fp8_cache_manifest(cache_dir: str | os.PathLike[str] | None) -> dict[str, Any] | None:
+    """Return a validated Phase-I manifest, or ``None`` for a cache miss."""
+    if cache_dir is None:
+        return None
+    root = Path(cache_dir)
+    manifest_path = root / CACHE_MANIFEST_NAME
+    if not (root / CACHE_COMPLETE_NAME).is_file() or not manifest_path.is_file():
+        return None
+    try:
+        manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
+    except (OSError, ValueError):
+        return None
+    if not isinstance(manifest, dict):
+        return None
+    if manifest.get("schema_version") != CACHE_SCHEMA_VERSION:
+        return None
+    if manifest.get("cache_kind") != CACHE_KIND:
+        return None
+    if manifest.get("component") != CACHE_COMPONENT_SUBFOLDER:
+        return None
+    quantization = manifest.get("quantization")
+    if not isinstance(quantization, dict) or quantization != {
+        "method": "fp8",
+        "serialized": True,
+        "activation_scheme": "dynamic",
+    }:
+        return None
+    if manifest.get("component_subfolder") != CACHE_COMPONENT_SUBFOLDER:
+        return None
+    if manifest.get("source_prefix") != f"{CACHE_COMPONENT_SUBFOLDER}.":
+        return None
+
+    index_path = root / CACHE_COMPONENT_SUBFOLDER / "model.safetensors.index.json"
+    if not index_path.is_file():
+        return None
+    try:
+        index = json.loads(index_path.read_text(encoding="utf-8"))
+    except (OSError, ValueError):
+        return None
+    if not isinstance(index, dict):
+        return None
+    weight_map = index.get("weight_map")
+    if not isinstance(weight_map, dict) or not weight_map or not all(isinstance(key, str) for key in weight_map):
+        return None
+    if manifest.get("tensor_count") != len(weight_map):
+        return None
+    for filename in weight_map.values():
+        if not isinstance(filename, str):
+            return None
+        relative_path = Path(filename)
+        if (
+            relative_path.is_absolute()
+            or relative_path.name != filename
+            or relative_path.suffix != ".safetensors"
+            or not (root / CACHE_COMPONENT_SUBFOLDER / relative_path).is_file()
+        ):
+            return None
+
+    source_files = manifest.get("source_files")
+    source_fingerprint = manifest.get("source_fingerprint")
+    if not isinstance(source_files, list) or not source_files or not isinstance(source_fingerprint, str):
+        return None
+    try:
+        if _source_fingerprint({str(filename) for filename in source_files}) != source_fingerprint:
+            return None
+    except (OSError, TypeError):
+        return None
+    return manifest
+
+
+def _cache_key_for_runtime_name(runtime_name: str, prefix: str) -> str:
+    if not runtime_name.startswith(prefix):
+        raise ValueError(f"runtime tensor {runtime_name!r} is outside cache source prefix {prefix!r}")
+    return runtime_name[len(prefix) :]
+
+
+def build_normalized_fp8_cache(
+    model: nn.Module,
+    *,
+    sources: tuple[object, ...],
+    cache_dir: str | os.PathLike[str],
+) -> dict[str, Any]:
+    """Write a normalized FP8 cache from a completed online-FP8 model.
+
+    The builder writes tensors one at a time into individual safetensors shards,
+    so it does not create a second full-model state dict. The first cache build
+    still follows the existing online loader; warm launches consume the cache
+    through the direct DLO mmap path.
+    """
+    cache_root = Path(cache_dir)
+    existing_manifest = load_fp8_cache_manifest(cache_root)
+    if existing_manifest is not None:
+        return existing_manifest
+    if cache_root.exists() and any(cache_root.iterdir()):
+        raise FileExistsError(f"FP8 cache directory is not empty and is not valid: {cache_root}")
+    cache_root.parent.mkdir(parents=True, exist_ok=True)
+
+    dit_names = tuple(getattr(model, "_dit_modules", ()))
+    dit_names = tuple(name for name in dit_names if hasattr(model, name))
+    if len(dit_names) != 1:
+        raise ValueError(f"Phase-I FP8 cache requires exactly one DiT module, got {dit_names}")
+    dit_name = dit_names[0]
+    if dit_name != CACHE_COMPONENT_SUBFOLDER:
+        raise ValueError(
+            f"Phase-I FP8 cache requires a top-level {CACHE_COMPONENT_SUBFOLDER!r} module, got {dit_name!r}"
+        )
+    source_prefix = f"{dit_name}."
+    dit_module = model.get_submodule(dit_name)
+    planned_sources = tuple(source for source in sources if getattr(source, "prefix", "") == source_prefix)
+    if len(planned_sources) != 1:
+        raise ValueError(f"Phase-I FP8 cache requires one source for {source_prefix!r}, got {len(planned_sources)}")
+
+    source_map = _build_source_map(
+        sources=planned_sources,
+        model_path=None,
+        remap_fn=getattr(type(model), "_remap_ckpt_key", None),
+    )
+    required = _collect_required_dit_tensors(((dit_name, dit_module),))
+    quantized_modules: dict[str, nn.Module] = {}
+    try:
+        from vllm.model_executor.layers.quantization.online.fp8 import Fp8PerTensorOnlineLinearMethod
+    except ImportError as exc:  # pragma: no cover - vLLM 0.27 provides this class
+        raise RuntimeError("Phase-I FP8 cache requires vLLM's online FP8 implementation") from exc
+
+    for module_name, module in model.named_modules():
+        if not module_name.startswith(source_prefix):
+            continue
+        if isinstance(getattr(module, "quant_method", None), Fp8PerTensorOnlineLinearMethod):
+            if not hasattr(module, "weight_scale"):
+                raise ValueError(f"online FP8 module {module_name!r} has no generated weight_scale")
+            quantized_modules[module_name] = module
+
+    tmp_root = cache_root.parent / f".{cache_root.name}.tmp-{uuid.uuid4().hex}"
+    component_root = tmp_root / CACHE_COMPONENT_SUBFOLDER
+    component_root.mkdir(parents=True, exist_ok=False)
+    weight_map: dict[str, str] = {}
+    source_files: set[str] = set()
+    total_size = 0
+    part_index = 0
+
+    file_cache: dict[str, Any] = {}
+
+    def write_tensor(key: str, tensor: torch.Tensor) -> None:
+        nonlocal part_index, total_size
+        if key in weight_map:
+            raise ValueError(f"normalized FP8 cache has duplicate tensor key {key!r}")
+        tensor = tensor.detach().to("cpu").contiguous()
+        filename = f"part-{part_index:05d}.safetensors"
+        part_index += 1
+        save_file({key: tensor}, str(component_root / filename))
+        weight_map[key] = filename
+        total_size += tensor.numel() * tensor.element_size()
+
+    try:
+        for runtime_name, target in required.items():
+            binding_info = source_map.get(runtime_name)
+            module_name = runtime_name.removesuffix(".weight")
+            quant_module = quantized_modules.get(module_name) if runtime_name.endswith(".weight") else None
+            if runtime_name.endswith(".weight_scale"):
+                owner_name = runtime_name.removesuffix(".weight_scale")
+                if owner_name in quantized_modules:
+                    # Generated online scales are published in the dedicated
+                    # pass below; they have no raw checkpoint binding.
+                    continue
+            if quant_module is not None:
+                if binding_info is None:
+                    raise ValueError(f"no source binding for quantized weight {runtime_name!r}")
+                qweight = quant_module.weight.detach()
+                if qweight.ndim != 2:
+                    raise ValueError(f"Phase-I FP8 cache only supports 2D weights, got {runtime_name!r}")
+                # Online FP8 stores the final runtime weight as qweight.T. The
+                # serialized FP8 method expects the pre-process qweight shape.
+                qweight = qweight.t().contiguous()
+                # Normalized caches contain the post-checkpoint-adapter QKV
+                # order. The cache-hit adapter bypasses the raw-checkpoint
+                # grouped-QKV transform, then the serialized FP8 finalizer
+                # performs only the standard weight transpose.
+                write_tensor(binding_info[0], qweight)
+                source_files.add(binding_info[1])
+                continue
+
+            if binding_info is not None:
+                checkpoint_key, file_path = binding_info
+                source_files.add(file_path)
+                handle = file_cache.get(file_path)
+                if handle is None:
+                    handle = safe_open(file_path, framework="pt", device="cpu")
+                    file_cache[file_path] = handle
+                write_tensor(checkpoint_key, handle.get_tensor(checkpoint_key))
+            else:
+                # Constructor-derived persistent buffers have no checkpoint
+                # entry; preserve them as canonical cache tensors.
+                write_tensor(_cache_key_for_runtime_name(runtime_name, source_prefix), target)
+
+        for module_name, module in quantized_modules.items():
+            scale = module.weight_scale.detach()
+            logical_widths = getattr(module, "logical_widths", None)
+            if scale.numel() == 1 and logical_widths is not None and len(logical_widths) > 1:
+                # Serialized Fp8LinearMethod represents a fused per-tensor
+                # weight with one scale slot per logical shard. Online FP8
+                # produces one scalar for the fused tensor, so replicate that
+                # scalar to the offline runtime shape without changing its
+                # value.
+                scale = scale.reshape(1).expand(len(logical_widths)).contiguous()
+            write_tensor(
+                _cache_key_for_runtime_name(f"{module_name}.weight_scale", source_prefix),
+                scale,
+            )
+
+        index = {
+            "metadata": {"total_size": total_size},
+            "weight_map": dict(sorted(weight_map.items())),
+        }
+        (component_root / "model.safetensors.index.json").write_text(
+            json.dumps(index, indent=2, sort_keys=True) + "\n",
+            encoding="utf-8",
+        )
+        manifest = {
+            "schema_version": CACHE_SCHEMA_VERSION,
+            "cache_kind": CACHE_KIND,
+            "component": "transformer",
+            "component_subfolder": CACHE_COMPONENT_SUBFOLDER,
+            "source_prefix": source_prefix,
+            "quantization": {
+                "method": "fp8",
+                "serialized": True,
+                "activation_scheme": "dynamic",
+            },
+            "tensor_count": len(weight_map),
+            "source_fingerprint": _source_fingerprint(source_files),
+            "source_files": sorted(source_files),
+        }
+        (tmp_root / CACHE_MANIFEST_NAME).write_text(
+            json.dumps(manifest, indent=2, sort_keys=True) + "\n",
+            encoding="utf-8",
+        )
+        (tmp_root / CACHE_COMPLETE_NAME).write_text("phase1\n", encoding="utf-8")
+        for handle in file_cache.values():
+            close = getattr(handle, "__exit__", None)
+            if callable(close):
+                close(None, None, None)
+        try:
+            os.replace(tmp_root, cache_root)
+        except FileExistsError:
+            # Multiple DLO workers may build the same cache concurrently. If
+            # another worker published a complete artifact first, treat this
+            # publisher as a successful cache race rather than a build error.
+            published = load_fp8_cache_manifest(cache_root)
+            if published is None:
+                raise
+            shutil.rmtree(tmp_root, ignore_errors=True)
+            logger.info("Normalized FP8 cache was published concurrently at %s", cache_root)
+            return published
+        logger.info("Published normalized FP8 cache at %s (%d tensors)", cache_root, len(weight_map))
+        return manifest
+    except BaseException:
+        for handle in file_cache.values():
+            close = getattr(handle, "__exit__", None)
+            if callable(close):
+                close(None, None, None)
+        shutil.rmtree(tmp_root, ignore_errors=True)
+        raise
+
+
+__all__ = [
+    "CACHE_SCHEMA_VERSION",
+    "CACHE_KIND",
+    "CACHE_COMPONENT_SUBFOLDER",
+    "build_normalized_fp8_cache",
+    "is_online_fp8_config",
+    "load_fp8_cache_manifest",
+    "_transformer_quant_config",
+]

vllm_omni/diffusion/offloader/distributed_layerwise_backend.py

@@ -1034,6 +1034,9 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
 
         logger.info("Realized %d loader-planned tensors as mmap views", len(loaded_names))
 
+        if getattr(pipeline, "_dlo_fp8_cache_normalized", False):
+            self._process_normalized_fp8_cache(pipeline)
+
         # Keep file handles open — _shard_and_pin will read from the mmap views.
         # They will be released after _shard_and_pin completes (when params are
         # replaced with offload placeholders).
@@ -1103,6 +1106,52 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
                 f"(first 5: {remaining_meta[:5]})."
             )
 
+    @staticmethod
+    def _process_normalized_fp8_cache(pipeline: nn.Module) -> None:
+        """Finalize serialized FP8 parameters after mmap replacement.
+
+        Phase-I cache files use the serialized FP8 checkpoint layout: the
+        weight is stored before the quantization method's final transpose, and
+        ``weight_scale`` is stored as an explicit cache tensor.  The ordinary
+        loader normally runs this method before DLO owns the plan; mmap loading
+        must run the same finalizer after replacing meta parameters.
+        """
+        from vllm.model_executor.layers.quantization.fp8 import Fp8LinearMethod
+        from vllm.model_executor.utils import replace_parameter
+
+        for module in pipeline.modules():
+            quant_method = getattr(module, "quant_method", None)
+            if not isinstance(quant_method, Fp8LinearMethod):
+                continue
+            if getattr(module, "_already_called_process_weights_after_loading", False):
+                continue
+
+            # ``Fp8LinearMethod.process_weights_after_loading`` also handles
+            # non-serialized checkpoints.  Its common path calls
+            # ``process_fp8_weight_tensor_strategy`` to requantize fused
+            # tensors, which is a CUDA-only operation.  A normalized cache
+            # already contains the serialized FP8 values and the generated
+            # scales, so only perform the layout transition and the kernel's
+            # own preprocessing here.  This keeps the warm path CPU/mmap
+            # compatible and preserves the cache bytes exactly.
+            if quant_method.block_quant:
+                raise ValueError("Phase-I normalized FP8 cache does not support block-quantized weights")
+            replace_parameter(module, "weight", module.weight.t().data)
+            # The online quantizer records one scalar.  A serialized FP8
+            # checkpoint may expose one repeated slot per logical shard in a
+            # fused module, but those slots must agree before they can be
+            # collapsed to the canonical per-tensor ``[1, 1]`` form.
+            weight_scale = module.weight_scale.reshape(-1)
+            if weight_scale.numel() > 1 and not torch.all(weight_scale == weight_scale[0]):
+                raise ValueError("Phase-I normalized FP8 cache contains non-uniform fused weight scales")
+            weight_scale = weight_scale[:1].reshape(1, 1)
+            replace_parameter(module, "weight_scale", weight_scale.data)
+            module.input_scale = None
+            if quant_method.use_marlin and hasattr(quant_method.fp8_linear, "marlin_input_dtype"):
+                quant_method.fp8_linear.marlin_input_dtype = quant_method.marlin_input_dtype
+            quant_method.fp8_linear.process_weights_after_loading(module)
+            module._already_called_process_weights_after_loading = True
+
     def _init_dp_group(self) -> None:
         """Reuse the process group initialized by parallel_state.
 

vllm_omni/engine/arg_utils.py

@@ -546,6 +546,7 @@ class OrchestratorArgs:
     enable_distributed_layerwise_offload: bool = False
     dlo_use_allgather: bool = True
     dlo_resident_layers: int = 0
+    dlo_fp8_cache_dir: str | None = None
     boundary_ratio: float | None = None
     flow_shift: float | None = None
     diffusion_kv_cache_dtype: str | None = None

vllm_omni/engine/async_omni_engine.py

@@ -1046,6 +1046,7 @@ class AsyncOmniEngine:
             "enable_distributed_layerwise_offload": kwargs.get("enable_distributed_layerwise_offload", False),
             "dlo_use_allgather": kwargs.get("dlo_use_allgather", True),
             "dlo_resident_layers": kwargs.get("dlo_resident_layers", 0),
+            "dlo_fp8_cache_dir": kwargs.get("dlo_fp8_cache_dir", None),
             "enforce_eager": False if kwargs.get("enforce_eager") is None else kwargs.get("enforce_eager"),
             "diffusion_compile_granularity": (
                 "regional"

vllm_omni/entrypoints/cli/serve.py

@@ -729,6 +729,14 @@ class OmniServeCommand(CLISubcommand):
             help="Keep this many leading main-DiT blocks resident on the device "
             "while distributed layerwise offload streams the remaining blocks.",
         )
+        omni_config_group.add_argument(
+            "--dlo-fp8-cache-dir",
+            type=str,
+            default=None,
+            help="Optional Phase-I normalized online-FP8 cache directory. "
+            "A cache miss builds the cache after the ordinary loader; a valid "
+            "cache hit uses serialized FP8 checkpoint mmap with DLO.",
+        )
         # Video model parameters (e.g., Wan2.2) - engine-level
         omni_config_group.add_argument(
             "--boundary-ratio",