vLLM-Omni Distributed Layerwise Offload

vLLM-Omni's Distributed Layerwise Offload (DLO) allows video generation models that exceed single-device HBM capacity—such as the 64B Cosmos3-Super—to run across multiple NPUs or GPUs with minimal host memory overhead. This system solves the memory bottlenecks associated with large Diffusion Transformer (DiT) models by combining meta-device initialization, weight sharding, and a double-buffered prefetch pipeline.

Solving the HBM and Host Memory Bottleneck

Large DiT models often fail to fit in device HBM, and traditional offloading or parallelism strategies introduce significant trade-offs. Pure layerwise offload in data-parallel (DP) configurations typically requires each rank to store a full copy of the model in host memory, leading to host RAM requirements that scale linearly with the number of devices (O(dp_size × model_size)).

Distributed Layerwise Offload addresses these limitations through four primary technical mechanisms:

1. Meta-Device Initialization and mmap Weight Loading

To prevent the massive Resident Set Size (RSS) spikes during model creation, vLLM-Omni uses meta-device initialization and memory-mapped (mmap) loading.

  • Mechanism: The system converts DiT modules to a meta device using to_empty(device="meta"), releasing parameter storage while keeping metadata. It then replaces these parameters with mmap views that point directly to the shared OS page cache.
  • Result: Because all ranks mmap the same safetensors files, the OS maintains only one copy of the file page in the cache. For Cosmos3-Nano DP4, this reduced the cold-start cgroup-visible peak memory from 178 GB to 47 GB (a 73% reduction).

2. Weight Sharding and AllGather Reconstruction

To eliminate the need for each rank to hold a full model copy in pinned CPU memory, vLLM-Omni implements weight sharding.

  • Mechanism: Each rank stores only 1/dp_size of the model weights. At runtime, the full weights for the current layer are reconstructed on the device using all_gather_into_tensor on a dedicated communication stream.
  • Result: Total pinned host memory is reduced from dp_size × model_size to simply model_size across all ranks. For Cosmos3-Super DP4, this reduced per-rank pinned memory from 124 GB to 31 GB.

3. Double-Buffered Prefetch Pipeline

To prevent the GPU from sitting idle during data movement and to keep HBM usage constant regardless of model depth, vLLM-Omni uses a double-buffer scheme.

  • Mechanism: The system maintains exactly two device buffers sized to the largest block in the model. While the compute stream processes layer N in slot 0, background streams handle the H2D (Host-to-Device) transfer and AllGather for layer N+1 into slot 1.
  • Result: HBM usage for weights is bounded by 2 × max_block_size. In tests with Cosmos3-Nano and Cosmos3-Super, peak HBM grew only 22% (23.1 GB to 28.1 GB) despite the model size increasing 3.8×.

4. DP Multi-Concurrency for Throughput

Since AllGather is request-independent (it gathers weights, not activations), vLLM-Omni allows different DP ranks to process different requests in parallel while remaining synchronized for weight reconstruction.

  • Mechanism: The scheduler batches up to dp_size requests. Each DP rank picks one request from the list and executes it through the pipeline's single-request forward path.
  • Result: This amortizes the AllGather overhead. In measurements, 4 concurrent requests achieved 3.3× the throughput of the HSDP single-request baseline, reaching approximately 83% of ideal linear scaling.

Performance and Validation Results

Platform Agnosticism

Distributed Layerwise Offload is platform-agnostic, functioning on both NVIDIA GPUs (CUDA/NCCL) and Ascend NPUs (CANN/HCCL).

On NVIDIA B300 GPUs, DLO+AG DP4 with 4 concurrent requests achieved 1.39× the throughput of HSDP+USP4 while using only 30% of the HBM (12.6 GiB vs 42.0 GiB) for 1024×1024 T2I tasks. For 720p 10s video generation, DLO+AG+USP4 performed within 2.13% of HSDP's latency while using only 47% of the HBM.

Topology-Aware Policy (MiniMax-H3 Study)

Research on 8× NVIDIA B300 GPUs using the MiniMax-H3 model indicates that the optimal DLO mode depends on the hardware topology:

  • DP1×SP8: AllGather is preferred for the lowest latency (P50 34.55s).
  • DP4×SP2: AllGather provides a balanced point for throughput (151.89 videos/h).
  • DP8×SP1: Rank-local DLO is preferred for the highest throughput (183.78 videos/h) and lowest energy consumption (43.97 Wh/video).

Memory Accounting on Ascend NPU

On Ascend hardware, pin_memory() allocates via /dev/davinci_manager, placing shards in CPU kernel DMA memory. This memory is invisible to the cgroup memory controller. Consequently, cgroup-visible memory scales as O(model_size + dp_size × constant), though total physical RAM still includes these pinned DMA shards.

Summary of Memory Scaling (Extrapolated)

Based on the measured memory model, vLLM-Omni can scale to very large models within a 2 TB system RAM limit:

Model Size dp_size Est. cgroup Peak Est. Total RAM Fits 2 TB?
33 GB 4 47 GB ~80 GB Yes
124 GB 4 172 GB ~296 GB Yes
185 GB 4 ~220 GB ~405 GB Yes
400 GB 4 ~423 GB ~823 GB Yes
400 GB 8 ~443 GB ~843 GB Yes

Sources

Related

  • Dispatch
    vLLM-Omni Accelerates Inference with AutoRound Quantization

    vLLM-Omni now integrates Intel's AutoRound post-training quantization, enabling W4A16 quantization that cuts model size up to 62% while preserving accuracy and unlocking performance gains on Intel XPU and NVIDIA GPUs.

  • Dispatch
    vLLM-Omni Optimizations for Qwen3-Omni-30B-A3B-Instruct Serving

    vLLM-Omni serves Qwen3-Omni-30B-A3B-Instruct via a three-stage pipeline (Thinker, Talker, Code2Wav) and improves throughput and latency using stage decomposition, CUDA Graphs, async chunk handoffs, async output, stage replicas, and hot‑path cleanup.

  • Dispatch
    VeRL-Omni v0.2.0 release notes / what's new

    VeRL-Omni v0.2.0 introduces request-level batching for faster diffusion RL and a reusable omni training stack for stable multimodal autoregressive training.

  • Project
    kacper-daftcode/vLLM-Moet

    A patched vLLM runtime for NVIDIA Blackwell GPUs that uses 2-bit quantization and tiered NVMe offloading to serve massive MoE models like GLM-5.2 and DeepSeek-V4-Flash on consumer hardware.

  • Dispatch
    vLLM DeepSeek V4 Support: Efficient Long-context Attention

    vLLM now supports DeepSeek V4-Pro and V4-Flash, implementing a new attention mechanism that enables context lengths up to one million tokens with significant KV cache memory savings.