[Feature] Distributed layerwise offload with DP multi-concurrency + mmap weight… (#5397)

Signed-off-by: evanchueng <evanchueng@outlook.com> Signed-off-by: david6666666 <530634352@qq.com> Co-authored-by: david6666666 <530634352@qq.com>
2026-07-30 13:35:38 +00:00
1 parent b3f4fbf9e2
commit 744c65b730
@@ -179,7 +179,134 @@ class Flux2Transformer2DModel(nn.Module):
 - Support single GPU only for now


-### Implementation Notes
+## Distributed Layerwise Offloading
+
+### How It Works
+
+Distributed layerwise offloading extends single-GPU layerwise offloading to
+multi-device deployments.  Each DP rank stores only **1/dp_size** of the model
+weights in host memory; full layer weights are reconstructed at runtime via
+**AllGather** on a dedicated communication stream, overlapped with computation
+via a fixed double-buffer scheme.
+
+**Key features:**
+- **Weight sharding + AllGather**: each rank stores 1/dp_size of weights,
+  reconstructed per-layer via `all_gather_into_tensor`
+- **Fixed double-buffer**: exactly 2 transformer blocks on each device at any
+  time, regardless of model size
+- **DP multi-concurrency**: N concurrent requests processed in parallel
+  (AllGather only gathers weight shards, which are request-independent)
+- **mmap weight loading**: weights loaded as mmap views pointing to shared OS
+  page cache, eliminating O(dp_size × model_size) RSS during model creation
+- **Platform-agnostic**: works on NVIDIA GPU (CUDA/NCCL) and Ascend NPU
+  (CANN/HCCL) via vLLM-Omni's platform abstraction
+
+**Execution Flow (per device):**
+
+```
+Compute Stream:  [Layer N]          [Layer N+1]          [Layer N+2]
+H2D Stream:      [H2D Shard N+1]   [H2D Shard N+2]
+AllGather:       [AG N+1]          [AG N+2]
+Slot Usage:      Slot0: Layer N    Slot1: Layer N+1
+```
+
+### Usage
+
+**CLI:**
+```bash
+# 4× GPU/NPU with AllGather (recommended)
+vllm serve /path/to/model --omni \
+  --enable-distributed-layerwise-offload \
+  --data-parallel-size 4
+
+# Without AllGather (each rank loads full weights, no sharding)
+vllm serve /path/to/model --omni \
+  --enable-distributed-layerwise-offload \
+  --data-parallel-size 4 \
+  --dlo-no-use-allgather
+
+# With SP instead of DP (long sequences)
+vllm serve /path/to/model --omni \
+  --enable-distributed-layerwise-offload \
+  --usp 4
+```
+
+**Python API:**
+```python
+from vllm_omni import Omni
+
+m = Omni(
+    model="/path/to/model",
+    enable_distributed_layerwise_offload=True,
+    dlo_use_allgather=True,  # default
+)
+```
+
+### CLI Flags
+
+| Flag | Description | Default |
+|------|-------------|---------|
+| `--enable-distributed-layerwise-offload` | Enable DLO | `false` |
+| `--data-parallel-size N` | Number of DP ranks for weight sharding | `1` |
+| `--dlo-use-allgather` | Shard + AllGather (saves CPU, requires concurrent requests) | `true` |
+| `--dlo-no-use-allgather` | Full weights per rank (no sharding, no AllGather) | `false` |
+
+### How model weights are loaded (mmap path)
+
+When DLO + AllGather is active, the offloader:
+
+1. **Saves non-persistent buffers** (e.g. RoPE `inv_freq`, timestep `freqs`)
+   from the normally-created transformer
+2. **Converts to meta device** via `to_empty(device="meta")`, releasing random
+   initialization weights
+3. **Loads checkpoint weights as mmap views** via `safe_open().get_tensor()`,
+   which return views into the OS page cache (shared across all ranks, 0 RSS)
+4. **Calls `post_load_weights()`** to apply model-specific dtype conversions
+   (e.g. Cosmos3's `time_embedder` → FP32)
+5. **Restores non-persistent buffers** from saved copies
+
+This approach requires **zero model-specific code changes** — no pipeline or
+transformer modifications are needed.
+
+### OffloadPlan (declarative topology metadata)
+
+Models can optionally declare an `OffloadPlan` class variable to provide
+topology metadata (block attribute names, submodules to offload) without
+any offload-specific logic:
+
+```python
+from vllm_omni.diffusion.offloader import OffloadPlan
+
+class MyPipeline(nn.Module):
+    _dit_modules = ["transformer"]
+    _offload_plan = OffloadPlan(
+        block_attrs={"transformer": ("blocks",)},
+        offload_submodules={"context_encoder": "layers"},
+    )
+```
+
+When not declared, the offloader falls back to `_layerwise_offload_blocks_attrs`
+and heuristic attribute search.  This is backward-compatible — existing models
+work without any changes.
+
+### DP Multi-concurrency
+
+When `--data-parallel-size > 1` and AllGather is enabled, the scheduler batches
+up to `dp_size` requests per denoise step.  Each DP rank processes a different
+request while AllGather synchronizes only weight shards (request-independent).
+
+**Requirements for concurrent requests:**
+- `num_inference_steps` must be specified explicitly (None is not allowed)
+- All concurrent requests must have the same `num_inference_steps` value
+- These constraints exist because AllGather is a collective that requires
+  every rank to participate at each denoise step
+
+### Limitations
+- Online quantization (FP8) is incompatible with mmap loading — falls back to
+  regular `load_weights()` automatically
+- Tensor Parallel is not supported (DLO uses DP-based sharding)
+- HSDP + AllGather is rejected (would double-shard weights)
+- `num_inference_steps=None` is not allowed in DP multi-concurrency mode

 **Module Discovery**

@@ -208,7 +335,10 @@ OffloadBackend (base class)
 ├── ModelLevelOffloadBackend → uses SequentialOffloadHook (.to() swap)
 │                              (delegates to a pipeline's enable_omni_model_cpu_offload
 │                               for split models like Cosmos3)
-└── LayerWiseOffloadBackend → uses LayerwiseOffloadHook
+├── LayerWiseOffloadBackend → uses LayerwiseOffloadHook
+│                          (single-GPU, full weights on host)
+└── DistributedLayerwiseOffloadBackend → uses DistributedLayerwiseOffloadHook
+                                         (multi-GPU, 1/dp_size sharded weights + AllGather)
 ```

 Factory function `get_offload_backend()` selects the appropriate backend based on
@@ -221,20 +351,22 @@ reasoner/generator components inside the model forward pass.

 ## Supported Models

-| Architecture | Example Models | DiT Class | Model-Level Offload | Layerwise Offload | Blocks Attrs (Layerwise specific) |
-|--------------|----------------|-----------|---------------------|-------------------|-----------------------------------|
-| Flux2Pipeline | `black-forest-labs/FLUX.2-dev` | `Flux2Transformer2DModel` | ✓ | ✓ | `"transformer_blocks"`, `"single_transformer_blocks"` |
-| LongCatImagePipeline | `meituan-longcat/LongCat-Image` | `LongCatImageTransformer2DModel` | - | ✓ | `"transformer_blocks"`, `"single_transformer_blocks"` |
-| NextStep11Pipeline | `stepfun-ai/NextStep-1.1` | `NextStepModel` | - | ✓ | `"layers"` |
-| OvisImagePipeline | `AIDC-AI/Ovis-Image-7B` | `OvisImageTransformer2DModel` | - | ✓ | `"transformer"` |
-| QwenImagePipeline | `Qwen/Qwen-Image` | `QwenImageTransformer2DModel` | ✓ | ✓ | `"transformer_blocks"` |
-| StableDiffusionXLPipeline | `stabilityai/stable-diffusion-xl-base-1.0` | `SDXLUNet2DConditionModel` | ✓ | ✓ | `"down_blocks"`, `"up_blocks"` |
-| StableDiffusion3Pipeline | `stabilityai/stable-diffusion-3.5-medium` | `SD3Transformer2DModel` | - | ✓ | `"transformer_blocks"` |
-| Wan22I2VPipeline | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | `WanTransformer3DModel` | ✓ | ✓ | `"blocks"` |
-| Wan22Pipeline | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | `WanTransformer3DModel` | ✓ | ✓ | `"blocks"` |
-| SoulXSingerPipeline / SoulXSingerSVCPipeline | `Soul-AILab/SoulX-Singer` | `DiffLlama` (`cfm_decoder.model.diff_estimator`) | ✓ | ✓ | `"layers"` |
-| BagelPipeline | `ByteDance-Seed/BAGEL-7B-MoT` | `Qwen2MoTModel` | - | ✓ | `"layers"`, `"customized modules"` |
+| Architecture | Example Models | DiT Class | Model-Level Offload | Layerwise Offload | Distributed Layerwise Offload | Blocks Attrs (Layerwise specific) |
+|--------------|----------------|-----------|---------------------|-------------------|-------------------------------|-----------------------------------|
+| Flux2Pipeline | `black-forest-labs/FLUX.2-dev` | `Flux2Transformer2DModel` | ✓ | ✓ | - | `"transformer_blocks"`, `"single_transformer_blocks"` |
+| LongCatImagePipeline | `meituan-longcat/LongCat-Image` | `LongCatImageTransformer2DModel` | - | ✓ | - | `"transformer_blocks"`, `"single_transformer_blocks"` |
+| NextStep11Pipeline | `stepfun-ai/NextStep-1.1` | `NextStepModel` | - | ✓ | - | `"layers"` |
+| OvisImagePipeline | `AIDC-AI/Ovis-Image-7B` | `OvisImageTransformer2DModel` | - | ✓ | - | `"transformer"` |
+| QwenImagePipeline | `Qwen/Qwen-Image` | `QwenImageTransformer2DModel` | ✓ | ✓ | - | `"transformer_blocks"` |
+| StableDiffusionXLPipeline | `stabilityai/stable-diffusion-xl-base-1.0` | `SDXLUNet2DConditionModel` | ✓ | ✓ | - | `"down_blocks"`, `"up_blocks"` |
+| StableDiffusion3Pipeline | `stabilityai/stable-diffusion-3.5-medium` | `SD3Transformer2DModel` | - | ✓ | - | `"transformer_blocks"` |
+| Wan22I2VPipeline | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | `WanTransformer3DModel` | ✓ | ✓ | - | `"blocks"` |
+| Wan22Pipeline | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | `WanTransformer3DModel` | ✓ | ✓ | - | `"blocks"` |
+| SoulXSingerPipeline / SoulXSingerSVCPipeline | `Soul-AILab/SoulX-Singer` | `DiffLlama` (`cfm_decoder.model.diff_estimator`) | ✓ | ✓ | - | `"layers"` |
+| BagelPipeline | `ByteDance-Seed/BAGEL-7B-MoT` | `Qwen2MoTModel` | - | ✓ | - | `"layers"`, `"customized modules"` |
+| Cosmos3OmniDiffusersPipeline | `nvidia/Cosmos3-Nano`, `nvidia/Cosmos3-Super` | `Cosmos3VFMTransformer`, `Cosmos3LanguageModel` | ✓ | ✓ | ✓ | `"layers"`, `"gen_layers"` |

 **Notes:**
 - Model-Level Offloading is expected to be supported by all common diffusion models (DiT and encoders) naturally
 - Layerwise Offloading requires DiT class to define `_layerwise_offload_blocks_attrs` pointing to transformer blocks
+- Distributed Layerwise Offloading works with any model that supports Layerwise Offloading — no additional model changes required.  See [Cosmos3 DistOffload recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/cosmos3/Cosmos3-DistOffload.md) for usage examples.
@@ -0,0 +1,191 @@
+# Cosmos3 — Distributed Layerwise Offload
+
+> Sharding weights across DP ranks with H2D + AllGather overlap for multi-NPU/GPU deployment
+
+## When to use this recipe
+
+Use this recipe when deploying Cosmos3-Super (64B / 124 GB) or Cosmos3-Nano (17B / 33 GB) across multiple devices where the model does not fit on a single device's HBM, or when you need higher throughput via DP multi-concurrency.
+
+**Key benefits:**
+- Each rank stores only 1/dp_size of the model weights on host memory
+- Only 2 transformer blocks reside on each device at any time (fixed double-buffer)
+- DP multi-concurrency: N concurrent requests with near-linear throughput scaling
+- mmap weight loading: shared page cache eliminates O(dp_size × model_size) RSS
+
+## Prerequisites
+
+- 2+ NPU or GPU devices (Ascend 910B3, NVIDIA B300, etc.)
+- vLLM-Omni with distributed layerwise offload support
+- Cosmos3-Nano or Cosmos3-Super checkpoint (local path or HuggingFace repo ID)
+
+## Serve commands
+
+### Cosmos3-Nano with DP=4 + AllGather (recommended for max throughput)
+
+```bash
+vllm serve /path/to/Cosmos3-Nano \
+  --omni \
+  --host 0.0.0.0 --port 8000 \
+  --no-guardrails --init-timeout 3600 \
+  --vae-use-tiling --vae-patch-parallel-size 1 \
+  --enable-distributed-layerwise-offload \
+  --data-parallel-size 4
+```
+
+### Cosmos3-Super with DP=2 + AllGather (for models > single-card HBM)
+
+```bash
+vllm serve /path/to/Cosmos3-Super \
+  --omni \
+  --host 0.0.0.0 --port 8000 \
+  --no-guardrails --init-timeout 3600 \
+  --vae-use-tiling --vae-patch-parallel-size 1 \
+  --enable-distributed-layerwise-offload \
+  --data-parallel-size 2
+```
+
+### Cosmos3-Super with DP=4 + AllGather
+
+```bash
+vllm serve /path/to/Cosmos3-Super \
+  --omni \
+  --host 0.0.0.0 --port 8000 \
+  --no-guardrails --init-timeout 3600 \
+  --vae-use-tiling --vae-patch-parallel-size 1 \
+  --enable-distributed-layerwise-offload \
+  --data-parallel-size 4
+```
+
+### Without AllGather (full weights per rank, no sharding)
+
+Use `--dlo-no-use-allgather` when you want each rank to load full weights independently —
+no AllGather synchronization, no concurrent request requirement, but N× host memory.
+
+```bash
+vllm serve /path/to/Cosmos3-Nano \
+  --omni \
+  --host 0.0.0.0 --port 8000 \
+  --no-guardrails --init-timeout 3600 \
+  --vae-use-tiling --vae-patch-parallel-size 1 \
+  --enable-distributed-layerwise-offload \
+  --data-parallel-size 4 \
+  --dlo-no-use-allgather
+```
+
+### With SP (sequence parallel) instead of DP
+
+For long-sequence workloads where SP parallelism is preferred:
+
+```bash
+vllm serve /path/to/Cosmos3-Super \
+  --omni \
+  --host 0.0.0.0 --port 8000 \
+  --no-guardrails --init-timeout 3600 \
+  --vae-use-tiling --vae-patch-parallel-size 1 \
+  --enable-distributed-layerwise-offload \
+  --usp 4
+```
+
+## Request examples
+
+### Text-to-Image (T2I)
+
+```bash
+curl -s -o output.png -X POST "http://localhost:8000/v1/images/generations" \
+  -H "Content-Type: multipart/form-data" \
+  -F "model=/path/to/Cosmos3-Nano" \
+  -F "prompt=A robot arm cleaning a plate, cinematic" \
+  -F "size=1024x1024" \
+  -F "num_inference_steps=50" \
+  -F "guidance_scale=6.0" \
+  -F "seed=42"
+```
+
+### Text-to-Video (T2V)
+
+```bash
+curl -s -o output.mp4 -X POST "http://localhost:8000/v1/videos/sync" \
+  -H "Accept: video/mp4" \
+  -F "model=/path/to/Cosmos3-Nano" \
+  -F "prompt=A robot arm cleaning a plate, cinematic shot" \
+  -F "negative_prompt=blurry" \
+  -F "size=832x480" \
+  -F "num_frames=29" \
+  -F "fps=24" \
+  -F "num_inference_steps=35" \
+  -F "guidance_scale=6.0" \
+  -F "max_sequence_length=4096" \
+  -F "flow_shift=10.0" \
+  -F 'extra_params={"use_resolution_template":false,"use_duration_template":false,"guardrails":false}' \
+  -F "seed=17"
+```
+
+> **Important:** When using DP multi-concurrency (AllGather mode), `num_inference_steps`
+> must be specified explicitly and must be the same for all concurrent requests.
+> `num_inference_steps=None` is not allowed because it may resolve differently
+> per request mode (e.g., T2V=35 vs action_mode=30), which would deadlock AllGather.
+
+### Multiple concurrent requests (DP multi-concurrency)
+
+With `--data-parallel-size 4`, send 4 concurrent requests to maximize throughput:
+
+```bash
+for i in 1 2 3 4; do
+  curl -s -o output_${i}.mp4 -X POST "http://localhost:8000/v1/videos/sync" \
+    -H "Accept: video/mp4" \
+    -F "model=/path/to/Cosmos3-Nano" \
+    -F "prompt=A robot arm cleaning a plate, shot ${i}" \
+    -F "negative_prompt=blurry" \
+    -F "size=832x480" \
+    -F "num_frames=29" \
+    -F "fps=24" \
+    -F "num_inference_steps=35" \
+    -F "guidance_scale=6.0" \
+    -F "max_sequence_length=4096" \
+    -F "flow_shift=10.0" \
+    -F 'extra_params={"use_resolution_template":false,"use_duration_template":false,"guardrails":false}' \
+    -F "seed=17" &
+done
+wait
+```
+
+## CLI flags
+
+| Flag | Description | Default |
+|------|-------------|---------|
+| `--enable-distributed-layerwise-offload` | Enable DLO with H2D + AllGather overlap | `false` |
+| `--data-parallel-size N` | Number of DP ranks (weight sharding + concurrent requests) | `1` |
+| `--dlo-use-allgather` | Use shard + AllGather for weight reconstruction (recommended) | `true` |
+| `--dlo-no-use-allgather` | Each rank loads full weights independently (no sharding, no AllGather) | `false` |
+| `--usp N` | Use SP instead of DP for weight sharding (long sequences) | `1` |
+
+## Mode comparison
+
+| Mode | CLI flags | CPU/rank | HBM/card | Throughput | Use case |
+|------|----------|---------|---------|-----------|----------|
+| DLO + AllGather (DP4) | `--enable-distributed-layerwise-offload --data-parallel-size 4` | 1/4 model | 2 blocks | **3.3× HSDP** | Max throughput, short sequences |
+| DLO + AllGather (DP2) | `--enable-distributed-layerwise-offload --data-parallel-size 2` | 1/2 model | 2 blocks | 2× HSDP | Balanced throughput + memory |
+| DLO + AllGather (SP4) | `--enable-distributed-layerwise-offload --usp 4` | 1/4 model | 2 blocks | 1× (single req) | Long sequences (720p+) |
+| DLO no-AllGather | `--enable-distributed-layerwise-offload --dlo-no-use-allgather` | Full model | 2 blocks | 0.4× HSDP | No AllGather overhead, high CPU |
+| HSDP | `--use-hsdp --hsdp-shard-size 4` | 0 | 1/4 model | 1× (baseline) | Weights fit in HBM |
+
+## Memory expectations
+
+| Model | Config | Host RAM (cgroup) | HBM/card |
+|-------|--------|-------------------|----------|
+| Nano (33 GB) | DP2 + AllGather | ~38 GB | ~10 GB |
+| Nano (33 GB) | DP4 + AllGather | ~47 GB | ~10 GB |
+| Super (124 GB) | DP2 + AllGather | ~157 GB | ~15 GB |
+| Super (124 GB) | DP4 + AllGather | ~172 GB | ~10 GB |
+
+> Host RAM is dominated by page cache (1× model_size, shared across ranks).
+> HBM only holds 2 transformer blocks + framework + activations, independent of model size.
+
+## Limitations
+
+- `--dlo-use-allgather` (default) requires all concurrent requests to have the same
+  `num_inference_steps` (AllGather is a collective that needs all ranks synchronized)
+- Online quantization (FP8) is incompatible with mmap loading — falls back to regular
+  `load_weights()` automatically
+- Tensor Parallel is not supported (DLO uses DP-based sharding, a different dimension)
+- HSDP + AllGather is rejected (would double-shard weights)
@@ -92,6 +92,21 @@ class MockScheduler:
        # assume all new req finished
        return [req.request_id for req in sched_output.scheduled_new_reqs]

+    def get_request_state(self, request_id):
+        return None
+
+    def pop_request_state(self, request_id):
+        return None
+
+    def finish_requests(self, request_id, status):
+        pass
+
+    def close(self):
+        pass
+
+    def initialize(self, od_config):
+        pass
+

 class _BatchCapablePipeline:
    supports_request_batch = True
@@ -644,7 +659,7 @@ async def test_async_add_req_and_stream_response():
    engine.step_execution = False
    engine.execution_mode = DiffusionExecutionMode.REQUEST_BATCH

-    def _finalize(rid, out, err=None):
+    def _finalize(rid, out, err=None, **kwargs):
        # Stream consumers stop on ``finished``; keep result_data for assertions.
        return SimpleNamespace(result_data=out.result.result_data, finished=True)

@@ -102,7 +102,10 @@ def _start_worker(req_q, res_q, count=2):

    def _run():
        for _ in range(count):
-            req = req_q.get(timeout=10)
+            try:
+                req = req_q.get(timeout=10)
+            except queue.Empty:
+                break
            method = req.get("method", "")
            args = req.get("args", ())
            if method == "execute_model_batch" and args and isinstance(args[0], DiffusionSchedulerOutput):
@@ -146,7 +149,7 @@ def _inject_interleave(executor):
        orig_enqueue(item)
        if threading.current_thread().name == "thread_a":
            a_enqueued.set()  # tell B: "A has enqueued"
-            b_complete.wait(5)  # block A until B finishes
+            b_complete.wait(1)  # block A until B finishes

    executor._broadcast_mq.enqueue = _controlled
    return a_enqueued, b_complete
@@ -173,8 +176,8 @@ class TestConcurrentRequestExecution:
            results["B"] = engine.add_req_and_wait_for_response(_mock_request("B"))
            b_complete.set()  # release A

-        ta = threading.Thread(target=_a, name="thread_a")
-        tb = threading.Thread(target=_b, name="thread_b")
+        ta = threading.Thread(target=_a, name="thread_a", daemon=True)
+        tb = threading.Thread(target=_b, name="thread_b", daemon=True)
        ta.start()
        tb.start()
        ta.join(10)
@@ -274,8 +277,8 @@ class TestConcurrentCollectiveRpc:
            )
            b_complete.set()

-        ta = threading.Thread(target=_a, name="thread_a")
-        tb = threading.Thread(target=_b, name="thread_b")
+        ta = threading.Thread(target=_a, name="thread_a", daemon=True)
+        tb = threading.Thread(target=_b, name="thread_b", daemon=True)
        ta.start()
        tb.start()
        ta.join(10)
@@ -311,8 +314,8 @@ class TestConcurrentRequestExecutionAndCollectiveRpc:
            )
            b_complete.set()

-        ta = threading.Thread(target=_a, name="thread_a")
-        tb = threading.Thread(target=_b, name="thread_b")
+        ta = threading.Thread(target=_a, name="thread_a", daemon=True)
+        tb = threading.Thread(target=_b, name="thread_b", daemon=True)
        ta.start()
        tb.start()
        ta.join(10)
@@ -683,7 +686,10 @@ class TestMultiprocExecutorStepStreamingOutput:

        def _worker():
            for runner_output in runner_outputs:
-                req_q.get(timeout=10)
+                try:
+                    req_q.get(timeout=10)
+                except queue.Empty:
+                    break
                res_q.put(runner_output)

        thread = threading.Thread(target=_worker, daemon=True)
@@ -474,6 +474,8 @@ class _DiffusionConfigProjection:
    output_type: str = "pil"
    enable_cpu_offload: bool = False
    enable_layerwise_offload: bool = False
+    enable_distributed_layerwise_offload: bool = False
+    dlo_use_allgather: bool = True
    pin_cpu_memory: bool = True
    diffusion_compile_granularity: Literal["regional", "full"] = "regional"
    diffusion_compile_dynamic: bool = Field(default=True, strict=True)
@@ -672,6 +672,12 @@ class OmniDiffusionConfig:
    enable_cpu_offload: bool = False
    # Layer-wise offloading (block-level offloading) parameters
    enable_layerwise_offload: bool = False
+    # Distributed layer-wise offloading with H2D + AllGather overlap (RFC-1)
+    enable_distributed_layerwise_offload: bool = False
+    # If True: shard weights 1/dp_size + AllGather (saves CPU memory, requires
+    # concurrent requests in DP mode). If False: each rank loads full weights
+    # via H2D only (N× CPU memory, but no AllGather synchronization needed).
+    dlo_use_allgather: bool = True

    pin_cpu_memory: bool = True  # Use pinned memory for faster transfers when offloading

@@ -156,6 +156,10 @@ class DiffusionEngine:
    #: remains a pure explicit user override (never mutated by engines).
    default_diffusion_model_runner_cls: str | None = None

+    # Class-level default so tests using object.__new__ (without __init__)
+    # don't hit AttributeError when _busy_loop accesses self.dp_concurrent.
+    dp_concurrent: bool = False
+
    def __init__(
        self,
        od_config: OmniDiffusionConfig,
@@ -224,6 +228,29 @@ class DiffusionEngine:
        self.scheduler.initialize(od_config)

    def _init_runtime_state(self) -> None:
+        # DP multi-concurrency: allow batching dp_size requests so each
+        # worker processes a different request in parallel.  Only enabled
+        # for distributed layerwise offload (which shards weights and
+        # needs all ranks active simultaneously).  Ordinary DP with a
+        # non-batch pipeline should not schedule multiple requests.
+        dp_size = 1
+        if getattr(self.od_config, "parallel_config", None) is not None:
+            dp_size = getattr(self.od_config.parallel_config, "data_parallel_size", 1)
+        dist_offload = getattr(self.od_config, "enable_distributed_layerwise_offload", False)
+        if dp_size > 1 and dist_offload:
+            self.scheduler.max_num_running_reqs = dp_size
+            self.dp_concurrent = True
+            # Set batch admission wait so concurrent requests accumulate
+            # before scheduling.  The scheduler reads this from od_config.
+            if getattr(self.od_config, "request_batch_max_wait_ms", 0) == 0:
+                self.od_config.request_batch_max_wait_ms = 500.0
+            logger.info(
+                "dp_concurrent: max_num_running_reqs=%d, batch_wait=%sms",
+                dp_size,
+                self.od_config.request_batch_max_wait_ms,
+            )
+        else:
+            self.dp_concurrent = False
        self.main_loop: asyncio.AbstractEventLoop | None = None
        self.stop_event: threading.Event | None = None
        self.worker_thread: threading.Thread | None = None
@@ -402,7 +429,7 @@ class DiffusionEngine:
                    # Only RPC / abort work pending; loop back to drain it.
                    continue

-                if self.execution_mode == DiffusionExecutionMode.REQUEST_BATCH:
+                if self.supports_request_batch or self.dp_concurrent:
                    self._wait_for_request_batch_admission_locked()

                sched_output = self.scheduler.schedule()
@@ -442,7 +469,7 @@ class DiffusionEngine:

        Caller must hold ``self._cv``.
        """
-        if self.step_execution or not self.supports_request_batch:
+        if self.step_execution or (not self.supports_request_batch and not self.dp_concurrent):
            return

        max_wait_s = self.od_config.request_batch_max_wait_ms / 1000.0
@@ -460,9 +487,12 @@ class DiffusionEngine:
        deadline = start + max_wait_s
        last_waiting = -1
        stable_since = start
-        # Require a short idle period with no queue growth so bursty HTTP
-        # ingress can land before the first schedule() of a wave.
-        stable_window_s = min(0.05, max_wait_s / 5.0)
+        # For dp_concurrent, use a longer stable window so all dp_size
+        # HTTP requests have time to land before scheduling.
+        if self.dp_concurrent:
+            stable_window_s = min(0.3, max_wait_s / 2.0)
+        else:
+            stable_window_s = min(0.05, max_wait_s / 5.0)

        while not self.stop_event.is_set():
            waiting = self.scheduler.num_waiting_requests()
@@ -783,6 +813,27 @@ class DiffusionEngine:
                raise RuntimeError(f"Could not {action} profiler: {e}") from e

    def run_startup_warmup(self) -> None:
+        dlo_use_allgather = getattr(self.od_config, "dlo_use_allgather", True)
+        # Skip dummy run when AllGather is used with more than 1 rank,
+        # because the dummy run sends only 1 request but AllGather requires
+        # all ranks to participate simultaneously.  This covers both DP > 1
+        # and SP > 1 (where dp_size is derived from sp_size in OffloadConfig).
+        pc = getattr(self.od_config, "parallel_config", None)
+        dp_size = getattr(pc, "data_parallel_size", 1) if pc else 1
+        sp_size = getattr(pc, "sequence_parallel_size", 1) if pc else 1
+        effective_shard_size = max(dp_size, sp_size)
+        skip_dummy = (
+            getattr(self.od_config, "enable_distributed_layerwise_offload", False)
+            and dlo_use_allgather
+            and effective_shard_size > 1
+        )
+        if skip_dummy:
+            logger.info(
+                "Skipping dummy run (dist_offload with AllGather, dp_size=%d, sp_size=%d)",
+                dp_size,
+                sp_size,
+            )
+            return
        try:
            self._dummy_run()
        except Exception as e:
@@ -1,6 +1,7 @@
 from __future__ import annotations

 import concurrent.futures
+import json
 import multiprocessing as mp
 import multiprocessing.connection
 import queue
@@ -64,12 +65,17 @@ class _ExecutorShutdownCleaner:
 class MultiprocDiffusionExecutor(DiffusionExecutor):
    uses_multiproc: bool = True

+    # Class-level defaults so tests using object.__new__ (without _init_executor)
+    # don't hit AttributeError when collective_rpc accesses these.
+    _rpc_wave_id: int = 0
+
    def _init_executor(self) -> None:
        self._processes: list[mp.Process] = []
        self._closed = False
        self._is_failed = False
        self._failure_callbacks: list[Callable[[], None]] = []
        self._result_mq: MessageQueue | None = None
+        self._rpc_wave_id: int = 0

        num_workers = cast(int, self.od_config.num_gpus)
        self.wake_events = [mp.Event() for _ in range(num_workers)]
@@ -215,6 +221,30 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
        MultiprocDiffusionExecutor._raise_for_rpc_error_dict(response)
        return response

+    _MAX_STALE_DISCARDS = 16
+
+    def _validate_wave_id(self, response: Any, expected_wave_id: int, deadline: float | None, method: str) -> Any:
+        """Discard stale RPC responses from a previous wave."""
+        discards = 0
+        while discards < self._MAX_STALE_DISCARDS:
+            if not isinstance(response, dict):
+                return response
+            resp_wave_id = response.get("wave_id")
+            if resp_wave_id is None or resp_wave_id == expected_wave_id:
+                return response
+            logger.warning(
+                "Discarding stale RPC response (wave_id=%s, expected=%s, method=%s)",
+                resp_wave_id,
+                expected_wave_id,
+                method,
+            )
+            discards += 1
+            response = self._dequeue_one_with_failure_polling(deadline, method)
+        raise TimeoutError(
+            f"Discarded {self._MAX_STALE_DISCARDS} stale RPC responses "
+            f"without finding wave_id={expected_wave_id} for method={method}."
+        )
+
    def _launch_workers(
        self,
        broadcast_handle: Handle,
@@ -355,9 +385,84 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
        from vllm_omni.diffusion.worker.utils import BatchRunnerOutput, RunnerOutput

        self._ensure_open()
+        new_reqs = scheduler_output.scheduled_new_reqs
        runner_outputs: list[RunnerOutput] = []

-        for new_req in scheduler_output.scheduled_new_reqs:
+        # DP multi-concurrency: when DLO+AllGather is active and multiple
+        # requests are scheduled, send ALL requests in one broadcast RPC.
+        # Each rank picks req[rank % len(reqs)] and computes independently.
+        # All ranks reply (unique_reply_rank=None) so we collect dp_size
+        # responses and match by dp_rank.
+        if (
+            len(new_reqs) > 1
+            and getattr(self.od_config, "enable_distributed_layerwise_offload", False)
+            and getattr(self.od_config, "dlo_use_allgather", True)
+        ):
+            # Validate: all concurrent requests must share the same
+            # num_inference_steps and identical extra_args, because
+            # AllGather is a collective that requires every rank to
+            # participate at each step.
+            step_counts = {
+                nr.req.sampling_params.num_inference_steps
+                for nr in new_reqs
+                if nr.req.sampling_params.num_inference_steps is not None
+            }
+            has_none = any(nr.req.sampling_params.num_inference_steps is None for nr in new_reqs)
+            if (len(step_counts) > 1) or has_none:
+                raise ValueError(
+                    "DP multi-concurrency requires all concurrent requests to have "
+                    "the same explicit num_inference_steps (None is not allowed), got "
+                    f"{[nr.req.sampling_params.num_inference_steps for nr in new_reqs]}."
+                )
+            extra_args_signatures: set = set()
+            for nr in new_reqs:
+                ea = getattr(nr.req, "extra_args", None)
+                if ea and isinstance(ea, dict):
+                    extra_args_signatures.add(json.dumps(ea, sort_keys=True))
+                else:
+                    extra_args_signatures.add(None)
+            if len(extra_args_signatures) > 1:
+                raise ValueError(
+                    "DP multi-concurrency requires all concurrent requests to "
+                    "share identical extra_args. Different extra_args can change "
+                    "the forward schedule and cause AllGather deadlock."
+                )
+
+            reqs_list = [nr.req for nr in new_reqs]
+            try:
+                results = self.collective_rpc(
+                    "execute_model",
+                    args=(reqs_list, self.od_config, scheduler_output.kv_prefetch_job),
+                    unique_reply_rank=None,
+                    exec_all_ranks=True,
+                )
+                results = results if isinstance(results, list) else [results]
+                for i, new_req in enumerate(new_reqs):
+                    res = results[i] if i < len(results) else results[0]
+                    if isinstance(res, DiffusionOutput):
+                        runner_outputs.append(
+                            RunnerOutput(
+                                request_id=new_req.request_id,
+                                step_index=None,
+                                finished=True,
+                                result=res,
+                            )
+                        )
+                    else:
+                        raise RuntimeError(f"Unexpected response type [{i}]: {type(res)!r}")
+            except Exception as exc:
+                for new_req in new_reqs:
+                    runner_outputs.append(
+                        RunnerOutput(
+                            request_id=new_req.request_id,
+                            step_index=None,
+                            finished=True,
+                            result=DiffusionOutput(error=str(exc)),
+                        )
+                    )
+            return BatchRunnerOutput.from_list(runner_outputs)
+
+        for new_req in new_reqs:
            req = new_req.req
            try:
                result = self.collective_rpc(
@@ -511,24 +616,61 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
        # ── Path 2: step-execution or other non-execute_model RPCs ──
        # _dequeue_one_with_failure_polling reads from _sync_result_buffer
        # when pump is active (request-mode), or result_mq directly otherwise.
+        self._rpc_wave_id += 1
+        wave_id = self._rpc_wave_id
+        rpc_request["wave_id"] = wave_id
+
        try:
            # Broadcast RPC request to all workers via unified message queue
            self._broadcast_mq.enqueue(rpc_request)  # pyright: ignore[reportOptionalMemberAccess] MQ is not None before shutdown

-            num_responses = 1
+            # Determine number of responses to collect:
+            # - unique_reply_rank=None + exec_all_ranks=True: all DP ranks reply
+            #   (N responses, one per DP worker).
+            # - Otherwise: 1 response (only rank 0 or specified rank)
+            if unique_reply_rank is None and exec_all_ranks:
+                dp_size = getattr(self.od_config.parallel_config, "data_parallel_size", 1)
+                num_responses = max(1, dp_size)
+            else:
+                num_responses = 1

-            responses = []
-            for _ in range(num_responses):
-                response = self._dequeue_one_with_failure_polling(deadline, method)
+            responses: list = []
+            if unique_reply_rank is None and exec_all_ranks and num_responses > 1:
+                # DP multi-concurrency: collect num_responses replies, sort by dp_rank.
+                tagged: list[tuple[int, Any]] = []
+                collected_errors: list[str] = []
+                for _ in range(num_responses):
+                    response = self._dequeue_one_with_failure_polling(deadline, method)
+                    response = self._validate_wave_id(response, wave_id, deadline, method)
+                    try:
+                        unpack_diffusion_output_shm(response)
+                    except Exception as e:
+                        logger.warning("SHM unpack failed (data may already be inline): %s", e)
+                    if isinstance(response, dict) and response.get("status") == "error":
+                        collected_errors.append(str(response.get("error", "unknown")))
+                    else:
+                        response = MultiprocDiffusionExecutor._handle_rpc_response(response)
+                        if isinstance(response, dict) and "dp_rank" in response:
+                            tagged.append((response["dp_rank"], response["output"]))
+                        else:
+                            tagged.append((len(tagged), response))
+                if collected_errors:
+                    raise RuntimeError(f"Worker error: {collected_errors[0]}")
+                tagged.sort(key=lambda x: x[0])
+                responses = [r for _, r in tagged]
+            else:
+                for _ in range(num_responses):
+                    response = self._dequeue_one_with_failure_polling(deadline, method)
+                    response = self._validate_wave_id(response, wave_id, deadline, method)

-                try:
-                    unpack_diffusion_output_shm(response)
-                except Exception as e:
-                    logger.warning("SHM unpack failed (data may already be inline): %s", e)
+                    try:
+                        unpack_diffusion_output_shm(response)
+                    except Exception as e:
+                        logger.warning("SHM unpack failed (data may already be inline): %s", e)

-                response = MultiprocDiffusionExecutor._handle_rpc_response(response)
+                    response = MultiprocDiffusionExecutor._handle_rpc_response(response)

-                responses.append(response)
+                    responses.append(response)

            return responses[0] if unique_reply_rank is not None else responses
        except Exception as e:
@@ -172,8 +172,9 @@ def pack_diffusion_output_shm(

    Supports a bare ``DiffusionOutput``, a wrapper object carrying one in
    ``.result`` (for example ``RunnerOutput``), an RPC result envelope carrying
-    the diffusion output in ``["result"]``, or a batch wrapper carrying
-    ``RunnerOutput`` objects in ``.runner_outputs``.
+    the diffusion output in ``["result"]``, a batch wrapper carrying
+    ``RunnerOutput`` objects in ``.runner_outputs``, or a DP-tagged dict
+    ``{"dp_rank": int, "output": DiffusionOutput}`` used by DP multi-concurrency.

    If *d2h_stream* is provided, D2H copies use that stream (non-blocking on
    the default stream).  The caller must synchronize *d2h_stream* afterward.
@@ -181,6 +182,13 @@ def pack_diffusion_output_shm(
    if isinstance(output, DiffusionOutput):
        return _pack_diffusion_fields(output, d2h_stream=d2h_stream)

+    # DP multi-concurrency: {"dp_rank": int, "output": DiffusionOutput}
+    if isinstance(output, dict) and "dp_rank" in output and "output" in output:
+        inner = output["output"]
+        if isinstance(inner, DiffusionOutput):
+            output["output"] = _pack_diffusion_fields(inner)
+        return output
+
    if _is_rpc_result_envelope(output):
        result = output.get("result")
        if isinstance(result, DiffusionOutput):
@@ -211,6 +219,13 @@ def unpack_diffusion_output_shm(output: object) -> object:
    if isinstance(output, DiffusionOutput):
        return _unpack_diffusion_fields(output)

+    # DP multi-concurrency: {"dp_rank": int, "output": DiffusionOutput}
+    if isinstance(output, dict) and "dp_rank" in output and "output" in output:
+        inner = output["output"]
+        if isinstance(inner, DiffusionOutput):
+            output["output"] = _unpack_diffusion_fields(inner)
+        return output
+
    if _is_rpc_result_envelope(output):
        result = output.get("result")
        if isinstance(result, DiffusionOutput):
@@ -362,17 +362,42 @@ class DiffusersPipelineLoader:
                )
            else:
                model = self._init_from_load_format(load_format, target_device, custom_pipeline_name, is_hsdp=False)
-                logger.debug("Loading weights on %s ...", load_device)
-                if load_format == "diffusers":
-                    # DiffusersAdapterPipeline.load_weights() calls
-                    # DiffusionPipeline.from_pretrained() internally; it does
-                    # not use our native customized pipeline classes.
-                    cast(DiffusersAdapterPipeline, model).load_weights()
+
+                # Skip load_weights only for DLO+AllGather when the model
+                # supports mmap loading and no online quantization is active.
+                # This condition MUST match the gate in
+                # DistributedLayerwiseOffloadBackend.enable() so that the
+                # loader skips ⟺ enable() uses mmap.
+                from vllm_omni.diffusion.offloader.offload_plan import supports_mmap_loading
+
+                _dist_offload = getattr(self.od_config, "enable_distributed_layerwise_offload", False)
+                _use_ag = getattr(self.od_config, "dlo_use_allgather", True)
+                _has_online_quant = self._has_online_quant(model)
+                _supports_mmap = supports_mmap_loading(model)
+
+                _skip_load = _dist_offload and _use_ag and _supports_mmap and not _has_online_quant
+
+                if _skip_load:
+                    logger.info("DLO+AllGather active: skipping load_weights (will load via mmap in enable())")
                else:
-                    self.load_weights(model)
-                # HSDP processes quantized weights before wrapping parameters as
-                # DTensors. The non-HSDP path can process them here as usual.
-                self._process_weights_after_loading(model, target_device)
+                    if _dist_offload and _use_ag and _has_online_quant:
+                        raise ValueError(
+                            "Online quantization is incompatible with DLO+AllGather: "
+                            "the sharding + AllGather mechanism flattens weights by "
+                            "dtype, which breaks quantized weight/scale layouts. "
+                            "Please use --dlo-no-use-allgather or disable online "
+                            "quantization."
+                        )
+                    if _dist_offload and _use_ag and not _supports_mmap:
+                        logger.info(
+                            "DLO+AllGather active but model does not support mmap: loading weights via regular loader."
+                        )
+                    logger.debug("Loading weights on %s ...", load_device)
+                    if load_format == "diffusers":
+                        cast(DiffusersAdapterPipeline, model).load_weights()
+                    else:
+                        self.load_weights(model)
+                    self._process_weights_after_loading(model, target_device)

            if offload_after_quant:
                model.to("cpu")
@@ -8,24 +8,54 @@ from vllm_omni.diffusion.data import OmniDiffusionConfig
 from vllm_omni.platforms import current_omni_platform

 from .base import OffloadBackend, OffloadConfig, OffloadStrategy
+from .block_discovery import get_blocks_attr_names, get_blocks_from_dit, set_blocks_attr_names
+from .distributed_layerwise_backend import (
+    DistributedLayerwiseOffloadBackend,
+    DistributedLayerwiseOffloadHook,
+    apply_distributed_block_hook,
+    remove_distributed_block_hook,
+)
 from .layerwise_backend import LayerWiseOffloadBackend
+from .offload_plan import OffloadPlan, get_offload_plan, supports_mmap_loading
 from .sequential_backend import (
    ModelLevelOffloadBackend,
    apply_sequential_offload,
    remove_sequential_offload,
 )
+from .tensor_utils import (
+    dtype_size,
+    is_dtensor,
+    is_materialized_tensor,
+    make_offload_placeholder,
+    set_tensor_storage,
+)

 logger = init_logger(__name__)

 __all__ = [
    "OffloadBackend",
    "OffloadConfig",
+    "OffloadPlan",
    "OffloadStrategy",
    "LayerWiseOffloadBackend",
+    "DistributedLayerwiseOffloadBackend",
+    "DistributedLayerwiseOffloadHook",
    "ModelLevelOffloadBackend",
    "apply_sequential_offload",
    "remove_sequential_offload",
+    "apply_distributed_block_hook",
+    "remove_distributed_block_hook",
    "get_offload_backend",
+    "get_offload_plan",
+    "supports_mmap_loading",
+    "get_blocks_attr_names",
+    "get_blocks_from_dit",
+    "set_blocks_attr_names",
+    "dtype_size",
+    "is_dtensor",
+    "is_materialized_tensor",
+    "make_offload_placeholder",
+    "set_tensor_storage",
 ]


@@ -75,6 +105,8 @@ def get_offload_backend(
        return ModelLevelOffloadBackend(config, device)
    elif config.strategy == OffloadStrategy.LAYER_WISE:
        return LayerWiseOffloadBackend(config, device)
+    elif config.strategy == OffloadStrategy.DISTRIBUTED_LAYER_WISE:
+        return DistributedLayerwiseOffloadBackend(config, device)
    else:
        logger.error("Unknown offload strategy: %s", config.strategy)
        return None
@@ -18,6 +18,7 @@ class OffloadStrategy(Enum):
    NONE = "none"
    MODEL_LEVEL = "model_level"  # Sequential offloading between DiT and encoders
    LAYER_WISE = "layer_wise"  # Block-level
+    DISTRIBUTED_LAYER_WISE = "distributed_layer_wise"  # Block-level with DP sharding + H2D/AllGather overlap


@dataclass
@@ -25,13 +26,21 @@ class OffloadConfig:
    strategy: OffloadStrategy
    pin_cpu_memory: bool = True
    use_hsdp: bool = False
+    dp_size: int = 1  # derived from parallel_config, not user-configurable
+    dlo_use_allgather: bool = True  # if True: shard + AllGather; if False: full weights H2D only
+    model_path: str | None = None  # checkpoint path for mmap weight loading

    @classmethod
    def from_od_config(cls, od_config: OmniDiffusionConfig) -> "OffloadConfig":
        """Extract and validate offload settings from OmniDiffusionConfig.

-        For now, enforces mutual exclusion between model-level and layer-wise offloading.
-        Layer-wise takes priority if both are enabled.
+        Enforces mutual exclusion among the three offload strategies.
+        Distributed layer-wise takes the highest priority, then layer-wise,
+        then model-level.
+
+        The ``dp_size`` is automatically derived from ``parallel_config`` —
+        it is NOT a user-configurable parameter. The distributed layerwise
+        offload works with whatever DP/SP parallelism is already set up.

        Args:
            od_config: OmniDiffusionConfig with offload settings
@@ -41,13 +50,39 @@ class OffloadConfig:
        """
        enable_cpu_offload = getattr(od_config, "enable_cpu_offload", False)
        enable_layerwise_offload = getattr(od_config, "enable_layerwise_offload", False)
+        enable_distributed_layerwise_offload = getattr(od_config, "enable_distributed_layerwise_offload", False)
        pin_cpu_memory = getattr(od_config, "pin_cpu_memory", True)

        parallel_config = getattr(od_config, "parallel_config", None)
        use_hsdp = getattr(parallel_config, "use_hsdp", False) if parallel_config else False

-        # Determine strategy (mutual exclusion, layer-wise takes priority)
-        if enable_layerwise_offload:
+        # Derive dp_size from parallel_config — not user-configurable.
+        # The offload adapts to whatever DP/SP is already configured.
+        dp_size = 1
+        if parallel_config is not None:
+            dp_size = getattr(parallel_config, "data_parallel_size", 1)
+            # HSDP's fully_shard_degree also contributes to effective DP
+            hsdp_shard_size = getattr(parallel_config, "hsdp_shard_size", -1) if use_hsdp else -1
+            hsdp_replicate_size = getattr(parallel_config, "hsdp_replicate_size", 1) if use_hsdp else 1
+            if use_hsdp and hsdp_shard_size > 0:
+                dp_size = hsdp_shard_size * hsdp_replicate_size
+
+            # When there is no DP but SP > 1, shard weights across SP ranks.
+            # AllGather reconstructs full weights per layer; each rank then
+            # computes on its SP portion of the sequence.  This gives N×
+            # compute parallelism with 1/N H2D transfer, reusing the exact
+            # same AllGather code path — only the process group changes.
+            if dp_size <= 1:
+                sp_size = getattr(parallel_config, "sequence_parallel_size", 1)
+                if sp_size and sp_size > 1:
+                    dp_size = sp_size
+
+        # Determine strategy (mutual exclusion, distributed layer-wise takes priority)
+        if enable_distributed_layerwise_offload:
+            strategy = OffloadStrategy.DISTRIBUTED_LAYER_WISE
+            if enable_layerwise_offload or enable_cpu_offload:
+                logger.info("Distributed layer-wise offloading takes priority, disabling other offloading strategies.")
+        elif enable_layerwise_offload:
            strategy = OffloadStrategy.LAYER_WISE
            if enable_cpu_offload:
                logger.info(
@@ -59,10 +94,37 @@ class OffloadConfig:
        else:
            strategy = OffloadStrategy.NONE

+        # When dlo_use_allgather=False, each rank loads full weights (no sharding,
+        # no AllGather).  This avoids AllGather synchronization requirements
+        # (concurrent requests, dummy run skip) at the cost of N× CPU memory.
+        dlo_use_allgather = getattr(od_config, "dlo_use_allgather", True)
+
+        # If dlo_use_allgather=False, force dp_size=1 (each rank independent)
+        if enable_distributed_layerwise_offload and not dlo_use_allgather:
+            dp_size = 1
+            logger.info(
+                "Distributed layerwise offload: dlo_use_allgather=False, "
+                "each rank loads full weights (no shard, no AllGather)"
+            )
+
+        # HSDP already shards parameters into DTensors.  Running distributed
+        # layerwise offload on top would shard each to_local() again, producing
+        # incorrect reconstruction after AllGather.  Reject this combination.
+        if enable_distributed_layerwise_offload and use_hsdp and dlo_use_allgather:
+            raise ValueError(
+                "Distributed layerwise offload with AllGather is incompatible with "
+                "HSDP: HSDP parameters are already sharded DTensors, and the offloader "
+                "would double-shard them. Use --dlo-no-use-allgather (full weights per "
+                "rank) or disable HSDP."
+            )
+
        return cls(
            strategy=strategy,
            pin_cpu_memory=pin_cpu_memory,
            use_hsdp=use_hsdp,
+            dp_size=dp_size,
+            dlo_use_allgather=dlo_use_allgather,
+            model_path=getattr(od_config, "model", None),
        )


@@ -0,0 +1,85 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Block discovery for layerwise offload.
+
+Shared between LayerWiseOffloadBackend and DistributedLayerwiseOffloadBackend.
+"""
+
+from __future__ import annotations
+
+from torch import nn
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+def get_blocks_attr_names(model: nn.Module) -> list[str]:
+    """Get block attribute names from model class."""
+    attrs: list[str] = getattr(model.__class__, "_layerwise_offload_blocks_attrs", [])
+
+    if not attrs:
+        old_attr = getattr(model.__class__, "_layerwise_offload_blocks_attr", None)
+        if old_attr is not None:
+            logger.warning(
+                "'_layerwise_offload_blocks_attr' is deprecated, "
+                "please use '_layerwise_offload_blocks_attrs' instead. "
+                "Example: _layerwise_offload_blocks_attrs = ['blocks']"
+            )
+            attrs = [old_attr] if isinstance(old_attr, str) else list(old_attr)
+
+    return attrs
+
+
+def set_blocks_attr_names(model: nn.Module, names: list[str]) -> None:
+    if not hasattr(model.__class__, "_layerwise_offload_blocks_attrs"):
+        setattr(model.__class__, "_layerwise_offload_blocks_attrs", names)
+
+
+def get_blocks_from_dit(model: nn.Module) -> tuple[list[str], list[nn.Module]]:
+    """Retrieve blocks and attribute names from provided DiT model."""
+    blocks_attr_names = get_blocks_attr_names(model)
+    if not blocks_attr_names:
+        logger.warning(
+            f"No _layerwise_offload_blocks_attrs defined for {model.__class__.__name__}, "
+            "skipping distributed layerwise offloading"
+        )
+        return [], []
+
+    blocks: list[nn.Module] = []
+    for name in blocks_attr_names:
+        attr = getattr(model, name, None)
+        if attr is None:
+            raise AttributeError(
+                f"Attribute '{name}' declared in _layerwise_offload_blocks_attrs "
+                f"does not exist on model {model.__class__.__name__}"
+            )
+        try:
+            attr_iter = iter(attr)
+        except TypeError:
+            if isinstance(attr, nn.Module):
+                logger.warning(
+                    "Attribute '%s' on %s is not iterable; treating it as one block.",
+                    name,
+                    model.__class__.__name__,
+                )
+                blocks.append(attr)
+                continue
+
+            logger.warning(
+                "Attribute '%s' on %s is not iterable (got %s); skipping it.",
+                name,
+                model.__class__.__name__,
+                type(attr).__name__,
+            )
+        else:
+            blocks.extend(attr_iter)
+
+    if not blocks:
+        logger.warning(
+            "No blocks found in %s for %s, skipping distributed layerwise offloading",
+            blocks_attr_names,
+            model.__class__.__name__,
+        )
+        return [], []
+
+    return blocks_attr_names, blocks
@@ -80,11 +80,13 @@ class ModuleDiscovery:
        "text_encoder_2",
        "text_encoder_3",
        "image_encoder",
+        "vision_encoder",
        "mllm",
    ]
    _FALLBACK_VAE_ATTRS = [
        "vae",
        "audio_vae",
+        "sound_tokenizer",
    ]

    @staticmethod
@@ -0,0 +1,61 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Declarative OffloadPlan for distributed layerwise offload.
+
+Models declare this as a class attribute ``_offload_plan`` on the
+pipeline class. When present, the offloader uses it instead of
+heuristic block discovery.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+from torch import nn
+
+
+@dataclass(frozen=True)
+class OffloadPlan:
+    """Optional declarative metadata for distributed layerwise offload.
+
+    Models declare this as a class attribute ``_offload_plan`` on the
+    pipeline class.  When present, the offloader uses it instead of
+    heuristic block discovery, making it easier to support new models
+    without offload-specific logic in constructor or forward code.
+
+    If not declared, the offloader falls back to:
+    1. ``_layerwise_offload_blocks_attrs`` on each DiT module class.
+    2. Heuristic search for ``layers`` / ``blocks`` / ``h`` attributes.
+
+    Attributes:
+        block_attrs: Maps DiT path → tuple of block-list attribute names.
+            e.g. ``{"transformer": ("gen_layers",),
+                    "transformer.language_model": ("layers",)}``
+        offload_submodules: Maps child name → block-list attribute name,
+            for large non-DiT submodules within a DiT that should be
+            independently offloaded with their own hooks.
+            e.g. ``{"context_encoder": "layers"}``
+    """
+
+    block_attrs: dict[str, tuple[str, ...]] = field(default_factory=dict)
+    offload_submodules: dict[str, str] = field(default_factory=dict)
+
+
+def get_offload_plan(pipeline: nn.Module) -> OffloadPlan | None:
+    """Retrieve the OffloadPlan declared by the pipeline, if any."""
+    return getattr(pipeline, "_offload_plan", None)
+
+
+def supports_mmap_loading(pipeline: nn.Module) -> bool:
+    """Whether the pipeline supports mmap weight loading.
+
+    A pipeline supports mmap if any module defines ``_remap_ckpt_key``,
+    which maps checkpoint keys to model parameter names.  Models without
+    this method must use the regular ``load_weights()`` path.
+
+    This check is shared between ``diffusers_loader.py`` (to decide whether
+    to skip ``load_weights``) and ``DistributedLayerwiseOffloadBackend.enable()``
+    (to decide whether to call ``_load_weights_via_mmap``), ensuring both
+    paths use the same gate condition.
+    """
+    return any(callable(getattr(type(m), "_remap_ckpt_key", None)) for m in pipeline.modules())
@@ -0,0 +1,47 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Shared tensor utilities for distributed layerwise offload.
+
+These helpers are used by both DistributedLayerwiseOffloadHook and
+DistributedLayerwiseOffloadBackend, and can be reused by other
+offload backends.
+"""
+
+from __future__ import annotations
+
+import torch
+from torch.distributed.tensor import DTensor
+
+
+def dtype_size(dtype: torch.dtype) -> int:
+    """Return element size in bytes for a torch.dtype."""
+    return torch.empty(1, dtype=dtype).element_size()
+
+
+def is_dtensor(t: torch.Tensor) -> bool:
+    """Check if tensor is a DTensor."""
+    return isinstance(t, DTensor)
+
+
+def set_tensor_storage(target: torch.Tensor, value: torch.Tensor) -> None:
+    """Replace target's underlying storage with value (zero-copy)."""
+    if is_dtensor(target):
+        target._local_tensor = value
+    else:
+        target.data = value
+
+
+def make_offload_placeholder(tensor: torch.Tensor) -> torch.Tensor:
+    """Create a zero-element placeholder to free GPU memory."""
+    if is_dtensor(tensor):
+        local_shape = tuple(tensor.to_local().shape)
+        return torch.empty(local_shape, device="meta", dtype=tensor.dtype)
+    return torch.empty((0,), device=tensor.device, dtype=tensor.dtype)
+
+
+def is_materialized_tensor(t: torch.Tensor) -> bool:
+    """Check if tensor holds real data (not meta or empty placeholder)."""
+    if is_dtensor(t):
+        local_t = t.to_local()
+        return not local_t.is_meta
+    return not t.is_meta and t.data.numel() > 0
@@ -212,7 +212,11 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
        )

        load_device = (
-            "cpu" if self.od_config.enable_cpu_offload or self.od_config.enable_layerwise_offload else str(self.device)
+            "cpu"
+            if self.od_config.enable_cpu_offload
+            or self.od_config.enable_layerwise_offload
+            or getattr(self.od_config, "enable_distributed_layerwise_offload", False)
+            else str(self.device)
        )

        def get_memory_context() -> AbstractContextManager[Any]:
@@ -478,7 +482,8 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
        # better perf. HSDP2's fully_shard pre-forward hooks need tensor version
        # counters, which inference tensors do not track.
        use_hsdp = od_config.parallel_config.use_hsdp
-        grad_context = torch.no_grad() if use_hsdp else torch.inference_mode()
+        use_distributed_offload = getattr(self.od_config, "enable_distributed_layerwise_offload", False)
+        grad_context = torch.no_grad() if (use_hsdp or use_distributed_offload) else torch.inference_mode()
        with grad_context:
            for req in reqs:
                self._prepare_request_for_forward(
@@ -417,12 +417,34 @@ class DiffusionWorker:

    def execute_model(
        self,
-        req: OmniDiffusionRequest,
+        req: OmniDiffusionRequest | list[OmniDiffusionRequest],
        od_config: OmniDiffusionConfig,
        kv_prefetch_job: KVPrefetchJob | None = None,
    ) -> DiffusionOutput:
-        """Execute a forward pass by delegating to the model runner."""
+        """Execute a forward pass by delegating to the model runner.
+
+        If *req* is a list (DP multi-concurrency), each rank picks one
+        request based on its distributed rank.  AllGather in the layerwise
+        offload only gathers weight shards (request-independent), so all
+        ranks stay synchronised at each AllGather call while computing
+        different activations.
+
+        Each rank returns its OWN DiffusionOutput (no gather).  The executor
+        collects N responses via result_mq (all ranks share one queue).
+        """
        assert self.model_runner is not None, "Model runner not initialized"
+
+        # DP multi-concurrency: pick one request per DP rank.
+        # Use rank_in_group (DP rank) not global rank, so that SP/TP/CFG
+        # ranks within the same DP replica select the same request.
+        is_batch = isinstance(req, list)
+        if is_batch:
+            from vllm_omni.diffusion.distributed.parallel_state import get_data_parallel_rank
+
+            dp_rank = get_data_parallel_rank()
+            idx = dp_rank % len(req)
+            req = req[idx]
+
        if self.lora_manager is not None:
            try:
                self.lora_manager.set_active_adapter(req.sampling_params.lora_request, req.sampling_params.lora_scale)
@@ -436,6 +458,14 @@ class DiffusionWorker:
            output = self.model_runner.execute_model(req, kv_prefetch_job=kv_prefetch_job)
        if profiler:
            profiler.step()
+
+        # Each primary rank returns its own output tagged with its DP rank
+        # so the executor can match results to requests by dp_rank.
+        # Non-primary ranks (SP/TP > 0) do not reply.
+        if is_batch:
+            from vllm_omni.diffusion.distributed.parallel_state import get_data_parallel_rank
+
+            return {"dp_rank": get_data_parallel_rank(), "output": output}
        return output

    def execute_model_batch(
@@ -781,17 +811,12 @@ class WorkerProc:
        self.result_mq = None
        self.result_mq_handle = None

-        # Setup result sender (only for rank 0)
-        if gpu_id == 0:
-            self.result_mq = MessageQueue(n_reader=1, n_local_reader=1, local_reader_ranks=[0])
-            self.result_mq_handle = self.result_mq.export_handle()
-            WorkerProc._shared_result_handle = self.result_mq_handle
-            logger.info(f"Worker {gpu_id} created result MessageQueue")
-        else:
-            handle = getattr(WorkerProc, "_shared_result_handle", None)
-            if handle:
-                self.result_mq = MessageQueue.create_from_handle(handle, gpu_id)
-                logger.info(f"Worker {gpu_id} attached to shared result MessageQueue")
+        # Each worker creates its own result MessageQueue (as writer).
+        # The executor creates one reader per worker to collect responses.
+        # This supports DP multi-concurrency where all ranks reply independently.
+        self.result_mq = MessageQueue(n_reader=1, n_local_reader=1, local_reader_ranks=[0])
+        self.result_mq_handle = self.result_mq.export_handle()
+        logger.info(f"Worker {gpu_id} created result MessageQueue")

        assert od_config.master_port is not None

@@ -920,12 +945,33 @@ class WorkerProc:
        output_rank = rpc_request.get("output_rank")
        exec_all_ranks = rpc_request.get("exec_all_ranks", False)
        collect_rank_status = rpc_request.get("collect_rank_status", False)
+        wave_id = rpc_request.get("wave_id")

        if collect_rank_status and not exec_all_ranks:
            raise ValueError("collect_rank_status requires exec_all_ranks=True so all ranks enter the status gather")

        should_execute = exec_all_ranks or output_rank is None or output_rank == self.gpu_id
-        should_reply = (output_rank is None or output_rank == self.gpu_id) and self.result_mq is not None
+        # For DP multi-concurrency (output_rank=None), only the primary rank
+        # within each DP replica should reply.  This prevents SP/TP/CFG/PP
+        # ranks from enqueuing extra replies that the executor doesn't drain.
+        if output_rank is None and exec_all_ranks:
+            from vllm.distributed.parallel_state import get_tensor_model_parallel_rank
+
+            from vllm_omni.diffusion.distributed.parallel_state import (
+                get_classifier_free_guidance_rank,
+                get_pipeline_parallel_rank,
+                get_sequence_parallel_rank,
+            )
+
+            is_primary_in_replica = (
+                get_sequence_parallel_rank() == 0
+                and get_classifier_free_guidance_rank() == 0
+                and get_tensor_model_parallel_rank() == 0
+                and get_pipeline_parallel_rank() == 0
+            )
+            should_reply = is_primary_in_replica and self.result_mq is not None
+        else:
+            should_reply = (output_rank is None or output_rank == self.gpu_id) and self.result_mq is not None

        if not should_execute:
            return None, False
@@ -968,6 +1014,7 @@ class WorkerProc:
                        "method": method,
                        "result": result,
                        "rank_statuses": rank_statuses,
+                        "wave_id": wave_id,
                    },
                    True,
                )
@@ -975,6 +1022,8 @@ class WorkerProc:

        if rpc_exception is not None:
            raise rpc_exception
+        if isinstance(result, dict) and wave_id is not None:
+            result["wave_id"] = wave_id
        return result, should_reply

    def _worker_busy_loop(self) -> None:
@@ -1016,10 +1065,16 @@ class WorkerProc:
                        self._return_result(result, rpc_id=rpc_id)
                except Exception as e:
                    logger.error(f"Error processing RPC: {e}", exc_info=True)
+                    # Apply the same reply gate as the success path so
+                    # non-output ranks don't enqueue stale error replies
+                    # that compete with the expected responder's message.
+                    output_rank = msg.get("output_rank")
+                    exec_all_ranks = msg.get("exec_all_ranks", False)
+                    wave_id = msg.get("wave_id")
                    if self.result_mq is not None:
                        if rpc_id is not None:
-                            # Async RPC: must complete the executor''s pending
-                            # future so collective_rpc() doesn''t hang.
+                            # Async RPC: must complete the executor's pending
+                            # future so collective_rpc() doesn't hang.
                            self.result_mq.enqueue(
                                AsyncDiffusionOutput(
                                    kind=AsyncOutputKind.RPC_RESULT,
@@ -1027,8 +1082,36 @@ class WorkerProc:
                                    error=str(e),
                                )
                            )
-                        else:
-                            self._return_result({"status": "error", "error": str(e)})
+                        elif output_rank is None and exec_all_ranks:
+                            # DP multi-concurrency: primary ranks reply, tagged
+                            from vllm.distributed.parallel_state import (
+                                get_tensor_model_parallel_rank,
+                            )
+
+                            from vllm_omni.diffusion.distributed.parallel_state import (
+                                get_classifier_free_guidance_rank,
+                                get_data_parallel_rank,
+                                get_pipeline_parallel_rank,
+                                get_sequence_parallel_rank,
+                            )
+
+                            is_primary = (
+                                get_sequence_parallel_rank() == 0
+                                and get_classifier_free_guidance_rank() == 0
+                                and get_tensor_model_parallel_rank() == 0
+                                and get_pipeline_parallel_rank() == 0
+                            )
+                            if is_primary:
+                                try:
+                                    dp_rank = get_data_parallel_rank()
+                                except Exception:
+                                    dp_rank = self.gpu_id
+                                self._return_result(
+                                    {"status": "error", "error": str(e), "dp_rank": dp_rank, "wave_id": wave_id}
+                                )
+                        elif output_rank is None or output_rank == self.gpu_id:
+                            # Normal RPC: only the expected rank replies
+                            self._return_result({"status": "error", "error": str(e), "wave_id": wave_id})

            elif isinstance(msg, dict) and msg.get("type") == "shutdown":
                logger.info("Worker %s: Received shutdown message", self.gpu_id)
@@ -1103,7 +1186,7 @@ class WorkerProc:
            pipe_writer.send(
                {
                    "status": "ready",
-                    "result_handle": worker_proc.result_mq_handle if rank == 0 else None,
+                    "result_handle": worker_proc.result_mq_handle,
                }
            )
            worker_proc._worker_busy_loop()
@@ -480,6 +480,8 @@ class OrchestratorArgs:
    num_weight_load_threads: int = 4
    enable_cpu_offload: bool = False
    enable_layerwise_offload: bool = False
+    enable_distributed_layerwise_offload: bool = False
+    dlo_use_allgather: bool = True
    boundary_ratio: float | None = None
    flow_shift: float | None = None
    diffusion_kv_cache_dtype: str | None = None
@@ -1008,6 +1008,8 @@ class AsyncOmniEngine:
            "enable_cache_dit_summary": kwargs.get("enable_cache_dit_summary", False),
            "enable_cpu_offload": kwargs.get("enable_cpu_offload", False),
            "enable_layerwise_offload": kwargs.get("enable_layerwise_offload", False),
+            "enable_distributed_layerwise_offload": kwargs.get("enable_distributed_layerwise_offload", False),
+            "dlo_use_allgather": kwargs.get("dlo_use_allgather", True),
            "enforce_eager": False if kwargs.get("enforce_eager") is None else kwargs.get("enforce_eager"),
            "diffusion_compile_granularity": (
                "regional"
@@ -641,6 +641,30 @@ class OmniServeCommand(CLISubcommand):
            action="store_true",
            help="Enable layerwise (blockwise) offloading on DiT modules.",
        )
+        omni_config_group.add_argument(
+            "--enable-distributed-layerwise-offload",
+            action="store_true",
+            help="Enable distributed layerwise offloading with H2D + AllGather overlap. "
+            "Shards weights across DP ranks, stores only 1/DP_size on each host, "
+            "and overlaps H2D transfers and AllGather with computation. "
+            "DP size is automatically derived from the parallel configuration.",
+        )
+        omni_config_group.add_argument(
+            "--dlo-use-allgather",
+            dest="dlo_use_allgather",
+            action="store_true",
+            default=True,
+            help="Use shard + AllGather for weight reconstruction (default: True). "
+            "When disabled (--dlo-no-use-allgather), each rank loads full weights "
+            "via H2D only — no sharding, no AllGather, no concurrent request "
+            "requirement, but N× CPU memory.",
+        )
+        omni_config_group.add_argument(
+            "--dlo-no-use-allgather",
+            dest="dlo_use_allgather",
+            action="store_false",
+            help="Disable AllGather: each rank loads full weights independently.",
+        )
        # Video model parameters (e.g., Wan2.2) - engine-level
        omni_config_group.add_argument(
            "--boundary-ratio",