[RFC]: Diffusion Chunked Scheduling RFC (Draft)

Overview

The primary objective is to achieve high Service Level Objective (SLO) attainment for short diffusion requests in mixed workloads. The current scheduling suffers from Head-of-Line (HoL) blocking, which violates SLOs. We propose mitigating this via a revised scheduler that implements a DiT chunked inference pipeline, thereby eliminating HoL blocking and ensuring high SLO adherence.

Background

"Recent advancements in generative modeling have shifted towards Diffusion Transformer (DiT) architectures, exemplified by state-of-the-art models like Wan2.2 and Qwen-Image. However, deploying these models in online inference environments presents unique challenges, particularly when handling mixed workloads of varying computational costs (i.e., a combination of long and short requests).

To scope this RFC feasible for the first iteration, we make the following assumptions:

1. Metric: We define SLO (Service Level Objective) as $3 \times$ the standalone execution time on a single GPU. The success metric is the percentage of requests completing within this time.
2. Homogeneous Dimensions (Resolution & Frames): For the initial implementation, we assume all requests in a batch share identical spatial resolution and temporal length (frame count).

Motivation.

DIT inference time

Profiling confirms that DiT runtime scales strictly linearly with inference steps ($R^2 \approx 1$), with the denoising loop consuming >99% of total GPU time. As shown in Figure 1, a 50-step request (~80s) takes ~4.7x longer than a 10-step request (~17s). This disparity proves that in a FIFO queue, short requests will suffer wait times far exceeding the defined 3× SLO threshold due to Head-of-Line blocking, validating the necessity for step-level preemption.

vllm-omni currently implement

"While vllm-omni efficiently schedules Qwen2.5-Omni via a token-based mechanism, this design is incompatible with the distinct workloads of visual DiTs. Consequently, the system lacks the decoupling required for fine-grained scheduling, treating each request as an atomic operation and forcing inefficient serial execution."

The Problem: Scheduling Bottlenecks in Monolithic Execution

Current monolithic execution creates severe Head-of-Line blocking for mixed workloads. As illustrated in Figure 1:

• Late Start (Figure 1a): Short requests queued behind long tasks suffer excessive wait times, often exceeding their SLO deadline before execution even begins.
• Resource Waste (Figure 1b): Delayed requests that start execution but hit their deadline mid-inference result in "wasted compute"—consuming GPU cycles for partial steps that yield no usable result.

Proposed Change.

Chunk definition

Chunk definition: to minimize the CPU overhead of per-step scheduling, we define a chunk as a group of consecutive steps (e.g., 5-10) and only perform scheduling and intermediate cache management at chunk boundaries.

Basic Idea: Chunked Preemption

To address these bottlenecks, we propose breaking monolithic execution into fine-grained chunks to enable preemption. As illustrated in Figure 2:

• Early Injection (Figure 2a): By splitting the long task into chunks, the scheduler can inject the short request immediately after the current chunk completes. This drastically reduces waiting time, allowing the short task to meet its SLO Deadline.
• Clean Finish (Figure 2b): Consequently, the short request finishes execution safely before the deadline, eliminating the "wasted compute" scenario and ensuring valid output.

Architecture

• Control Plane (EDF Scheduler): Implement a chunk-based scheduler on Rank 0 that dispatches tasks in Chunks (e.g., 5-10 steps) rather than full requests, enabling fine-grained preemption and EDF scheduling.
  • • core/dit_scheduler.py
• Data Plane (Re-entrant Execution): Refactor GPU workers from execute_model() to “execute_chunk()”. This allows the diffusion loop to pause and resume states (latents, step index) at chunk boundaries.
  • • gpu_worker.py; vllm_omni/diffusion/models/**/pipeline_*.py
  • • core/request.py
• Resource Plane (Block Manager): Adapt vLLM's memory manager to treat diffusion latents as paged blocks. This enables dynamic allocation, eviction, and swapping to manage concurrency without OOM.
  • • diffusion/block_manager.py

Schedule policy

The scheduler design leverages two fundamental characteristics of DiT models to enable precise and responsive control:

1. Iterative Nature: The sequential denoising loop allows for safe preemption at any step boundary, enabling the scheduler to yield execution without state corruption.
2. Linear Predictability: Unlike variable-length LLM decoding, diffusion steps incur constant computational cost. Profiling confirms strict linear scaling ($R^2 \approx 1$), enabling precise chunk sizing and deterministic scheduling predictions. Building on these properties, we implement an SLO-aware Earliest Deadline First (EDF) policy that operates at Chunk boundaries. By continuously evaluating the deadline slack for each active request: $$\text{Slack} = T_{\text{deadline}} - T_{\text{now}} - T_{\text{remaining}}$$ The system dynamically prioritizes tasks with minimal slack. This mechanism ensures that urgent generation requests preempt those with ample slack, strictly guaranteeing latency SLOs while maximizing aggregate system throughput through fine-grained preemption.

Memory manage

1. Unified Resource Object We define the per-request GPU footprint as a ${ \text{latents}, \text{embeds}, \text{cachestate} }$ ; Here, latents represents the mandatory, recoverable state, while cache_state (e.g., TeaCache/Cache-DiT tensors) represents disposable, transient acceleration artifacts.
2. HBM Management & Accounting The HBM Manager operates at block granularity. It enforces strict lifecycle management (allocate/evict/resume) for latents. In contrast, cache_state is tracked solely for budget accounting; its memory footprint influences admission and eviction decisions but is not required to be serialized or preserved during swaps.
3. Preemption Semantics A PAUSE action suspends scheduling and preserves latents (optionally swapping blocks to host memory). To prioritize correctness and implementation simplicity, cache_state is discarded (flushed via refresh()) upon context switching/preemption.

Feedback Period.

No response

CC List.

None

---

Any Other Things.

Other side benefits (with careful engineering)

At $$256^2$$ resolution, the GPU Compute Units (SMs) are significantly underutilized with a single request.

• Data Evidence: Increasing the batch size from 1 to 4 increases the total workload by 400%, yet the execution time only rises by ~107%.This latency increase (~2.07x total) fits comfortably within our SLO requirements (which allow for a latency budget of up to 3x execution time for batched tasks).
• Implication: Aggregated batching is effective at low resolutions, enabling improved throughput while meeting SLO (Service Level Objective) requirements.

The Problem: GPU Underutilization (Rigid Batching) Static batching forces lock-step execution, meaning completed slots cannot be refilled until the entire batch finishes. As shown in Figure 3(a), mixing requests with different step counts (e.g., 10 vs. 50) creates massive "compute bubbles"—idle periods where GPU cores sit underutilized for the remainder of the long task, directly degrading cluster throughput. The Solution: Continuous Batching To maximize utilization, we propose Continuous Batching. As illustrated in Figure 3(b), this mechanism immediately injects new requests into freed slots during the inference loop, ensuring GPU compute remains saturated regardless of individual request variance.

Other potential discussions

Heterogeneous Resolution & Frames

• Mixed steps Batching Validity: Mixing requests at different generation stages (e.g., Step 50 vs. Step 10) is mathematically valid. DiT models are stateless functions F(x,t,c) where the timestep t is injected as a batch vector, not a global scalar. Since DiTs use per-sample normalization (LayerNorm/RMSNorm) rather than Batch Normalization, there is no statistical contamination between high-noise and low-noise samples within the same batch.
• Mixed resolution Batching: Support batching a 512x512 image request with a 1024x1024 image request. This requires advanced memory management (like PagedAttention/Block Tables) to handle non-contiguous memory or dynamic padding.
• Frame: For video, handling different frame counts in the same batch efficiently.

Implementation Constraints

• Attention Isolation: We must ensure strict attention masking (e.g., correct cu_seqlens in FlashAttention) so that tokens from separate requests do not attend to each other.
• CFG Alignment: Mixing requests that require Classifier-Free Guidance (2x batch expansion) with those that do not (e.g., distilled models) introduces complex indexing overhead. Initial implementations will enforce uniform CFG settings per batch.

Before submitting a new issue...

• Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the documentation page, which can answer lots of frequently asked questions.