[diffusion] Allow Cache-DiT with DiT layerwise offload

docs/cookbook/diffusion/MiniMax/MiniMax-H3.mdx

@@ -765,9 +765,11 @@ name, for example `sglang generate --quality high`.
 </Note>
 
 For manually tuned Cache-DiT experiments outside that validated path, omit
-the request `quality` field and set the process-wide environment controls
-directly. An explicit `quality: "lossless"` request overrides those controls
-and restores native denoising:
+the request `quality` field and set `--enable-cache-dit` or the
+process-wide `SGLANG_CACHE_DIT_*` defaults. An explicit `quality`
+(including `"lossless"`) takes H3 off the generic Cache-DiT path. The
+24 GB layerwise recipe above can use the same switch; skipped blocks are
+not streamed.
 
 ```bash Command
 SGLANG_CACHE_DIT_ENABLED=true \
@@ -787,12 +789,14 @@ sglang serve \
 ```
 
 <Warning>
-Cache-DiT skips selected block computation and is approximate. It cannot be
-combined with FSDP inference or DiT layerwise offload. Breakable CUDA graph
-execution takes precedence and leaves Cache-DiT disabled. Tune the cache
-thresholds only after comparing both video and audio quality on the target
-task profile. A real B200 request has completed, but the `quality: "high"`
-path above remains fail-closed to the audited 4×H200 workload.
+Cache-DiT skips selected block computation and is approximate. It cannot
+be combined with FSDP inference. DiT layerwise offload is compatible:
+skipped blocks are not streamed, and the first layer after a skip may
+sync-load. Breakable CUDA graph execution takes precedence and leaves
+Cache-DiT disabled. Tune the cache thresholds only after comparing both
+video and audio quality on the target task profile. A real B200 request
+has completed, but the `quality: "high"` path above remains fail-closed
+to the audited 4×H200 workload.
 </Warning>
 
 ## 7. Feature contracts and advanced recipes
@@ -816,7 +820,7 @@ listed hardware and topology; it is not inherited by a similar GPU family.
 | Tensor parallelism | Verified: B200 TP2 + Ulysses4; H100 TP2 + Ulysses2 and TP4 + Ulysses1 | `--tp-size` may be combined with Ulysses when the TP-local head count remains divisible by the Ulysses degree. On 4×H100, TP2 + Ulysses2 is the measured speed default. |
 | FSDP inference | Verified: 4× B200 and 4× H100 + Ulysses4 | Preserves H3's mixed BF16/FP32 parameter policy. B200 completed the exact eager comparison; H100 completed consecutive real requests at about 57 GB peak memory per GPU. |
 | Resident components | Verified: B200, H200, 4×H100 with TP, and 1/2/4/8× MI300X and MI355X | This is the recommended single-request latency path when the complete workload fits. |
-| CPU and layerwise offload | Verified: 2× RTX 5090 TP2; 1× RTX 4090 24 GB | The 5090 lossless recipe keeps 20 DiT blocks plus both VAE encoders resident, streams the remaining DiT blocks, text encoder, and video VAE decoder blocks, and leaves the small audio VAE resident. The 4090 recipe streams DiT and the text encoder with zero resident DiT layers and **omits `vae`** from `--layerwise-offload-components`. |
+| CPU and layerwise offload | Verified: 2× RTX 5090 TP2; 1× RTX 4090 24 GB | The 5090 lossless recipe keeps 20 DiT blocks plus both VAE encoders resident, streams the remaining DiT blocks, text encoder, and video VAE decoder blocks, and leaves the small audio VAE resident. The 4090 recipe streams DiT and the text encoder with zero resident DiT layers and **omits `vae`** from `--layerwise-offload-components`. Compatible with Cache-DiT; skipped blocks are not streamed. |
 | Breakable CUDA graph | Verified: B200 Ref2VA, opt-in | Matching eager output was observed for the captured signature, without a measured speedup. Re-capture for other shapes and reference sets. |
 | `torch.compile` | Measured: H200, opt-in | Steady-state benefit was below measurement noise, while startup increased and numerical output changed. Do not use it for consistency ground truth. |
 
@@ -1440,7 +1444,8 @@ timed (seed 42); only the timed pass is reported. GPU peak stayed about
 `kitchen_int8` + FA changes Linear numerics only. The `sol_attn` /
 `sage_attn` / hybrid rows also change the attention algorithm, so speed
 and pixel fidelity rank in opposite orders there. Default remains
-`kitchen_int8` + `fa`.
+`kitchen_int8` + `fa`. Cache-DiT can share this layerwise recipe; omit
+`quality` and see the quality-level section.
 
 ### AMD Instinct task and scaling runs
 

docs/docs/sglang-diffusion/api/cli.mdx

@@ -248,7 +248,7 @@ sglang generate \
 HTTP server-only arguments are ignored by `sglang generate`.
 </Note>
 
-For supported native pipelines, set `SGLANG_CACHE_DIT_ENABLED=true` to enable Cache-DiT. For the diffusers backend, use `--backend diffusers --cache-dit-config ...`. See [Cache-DiT](../cache_dit).
+For supported native pipelines, set `SGLANG_CACHE_DIT_ENABLED=true` to enable Cache-DiT. It can run with DiT layerwise offload; it cannot run with FSDP. For the diffusers backend, use `--backend diffusers --cache-dit-config ...`. See [Cache-DiT](../cache_dit).
 
 For supported image pipelines, breakable CUDA graph can be enabled with `--enable-breakable-cuda-graph`, but you must declare every served resolution in `--warmup-resolutions` so warmup captures matching graph signatures.
 
@@ -316,7 +316,7 @@ sglang generate \
 
 Values passed to the compatibility option `--layerwise-offload-components` must match loaded component keys, such as `transformer`, `text_encoder`, `image_encoder`, `vae`, `condition_image_encoder`, `spatial_upsampler`, or `vocoder`. Its `default` group selects text encoders, image encoders, and VAEs. Use `all` to select every layerwise-offloadable component.
 
-Layerwise tuning options such as `--dit-offload-prefetch-size`, `--dit-layerwise-resident-layers`, and `--dit-layerwise-residency-policy` continue to control the streamed layer working set. Prefer the smallest component set that solves the memory issue because layerwise offload can increase latency.
+Layerwise tuning options such as `--dit-offload-prefetch-size`, `--dit-layerwise-resident-layers`, and `--dit-layerwise-residency-policy` continue to control the streamed layer working set. Prefer the smallest component set that solves the memory issue because layerwise offload can increase latency. DiT layerwise offload can run with Cache-DiT: skipped blocks are not streamed, and the first layer after a skip may sync-load. Cache-DiT remains incompatible with FSDP.
 
 Those three set the default for every streamed component. To give one component its own values, use the `component=value` forms, which also accept JSON:
 

docs/docs/sglang-diffusion/cache_dit.mdx

@@ -605,7 +605,8 @@ SGLang Diffusion x Cache-DiT supports almost all models originally supported in
 ## Limitations
 
 - **SGLang-native pipelines**: Distributed Cache-DiT paths exist for supported pipelines. Hybrid SP+TP configurations add communication and cache coordination overhead, so validate them on the target model and hardware before using them as production defaults.
-- **SCM minimum steps**: SCM requires >= 8 inference steps to be effective
+- **DiT layerwise offload**: Compatible. Skipped blocks are not streamed, and the first layer after a skip may sync-load. Still incompatible with `--use-fsdp-inference`.
+- **SCM minimum steps**: SCM requires >= 8 inference steps to be effective. Some pipelines report `steps - 1` NFEs (for example MiniMax-H3 `num_inference_steps=8` is 7 NFEs), which trips the upstream `steps_mask` assertion; use at least 9 requested steps or custom bins.
 - **Model support**: The model must be registered in Cache-DiT's
   `BlockAdapterRegister` or have an SGLang custom block adapter.
 

docs/docs/sglang-diffusion/caching-acceleration.mdx

@@ -54,7 +54,7 @@ advanced strategies like DBCache and TaylorSeer. It can achieve up to **1.69x sp
 See [Cache-DiT](./cache_dit) for detailed configuration.
 
 <Note>
-Cache-DiT currently cannot be combined with `--use-fsdp-inference`. Keep FSDP disabled when enabling Cache-DiT, or use other residency/offload controls instead.
+Cache-DiT currently cannot be combined with `--use-fsdp-inference`. Keep FSDP disabled when enabling Cache-DiT. DiT layerwise offload is compatible: skipped blocks are not streamed, and the first layer after a skip may sync-load.
 </Note>
 
 ### Quick Start

python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-performance/SKILL.md

@@ -46,7 +46,7 @@ These options are intended to preserve output quality. In practice, some paths (
 | **Multi-GPU (SP)** | `--num-gpus N --ulysses-degree N` | Sequence parallelism across GPUs. Shards sequence tokens (not frames) to minimize padding. | Near-linear scaling with N GPUs | Requires NCCL; inter-GPU bandwidth matters. `ulysses_degree * ring_degree = sp_degree`. For Wan2.2 video, start by benchmarking pure Ulysses before assuming a mixed Ulysses/Ring layout is fastest. |
 | **Cross-node SP** | `--nnodes`, `--node-rank`, `--dist-init-addr` with total `--num-gpus`; combine node-local Ulysses with cross-node Ring | Extends sequence parallel groups across multiple nodes. | Capacity and long-sequence scaling beyond one host | Prefer Ulysses within a node and Ring across nodes; all-to-all is usually the least cross-node-friendly. Use `--encoder-parallel replicate` today and verify the model's Ring admission and determinism. MiniMax-H3 is the current end-to-end validated recipe. |
 | **CFG Parallel** | `--enable-cfg-parallel` | Runs conditional and unconditional CFG branches in parallel across GPUs. For CFG models on multi-GPU, benchmark this against pure Ulysses on your topology instead of assuming one always wins. | Often faster than pure SP for CFG models | Requires `num_gpus >= 2`. Halves the Ulysses group size (e.g. 8 GPU → two 4-GPU groups). Only for models that use CFG. Nightly coverage configs may intentionally use smaller Ulysses groups to keep ring behavior exercised; that does not automatically make them the lowest-latency choice. |
-| **Layerwise Offload** | `--dit-layerwise-offload` | Async layer-by-layer H2D prefetch with compute overlap. Only ~2 DiT layers reside on GPU at a time, dramatically reducing VRAM. For some video models the copy stream can be almost fully hidden behind compute. | Saves VRAM (40 GB → ~11 GB for Wan A14B); can be near-zero speed cost on the right workload | Enabled by default for Wan/MOVA video models. Incompatible with Cache-DiT. For **image models** or highly parallelized setups (many GPUs, small per-GPU compute), the copy stream may not be fully hidden and can cause slowdown. |
+| **Layerwise Offload** | `--dit-layerwise-offload` | Async layer-by-layer H2D prefetch with compute overlap. Only ~2 DiT layers reside on GPU at a time, dramatically reducing VRAM. For some video models the copy stream can be almost fully hidden behind compute. | Saves VRAM (40 GB → ~11 GB for Wan A14B); can be near-zero speed cost on the right workload | Enabled by default for Wan/MOVA video models. Compatible with Cache-DiT (skipped blocks are not streamed). For **image models** or highly parallelized setups (many GPUs, small per-GPU compute), the copy stream may not be fully hidden and can cause slowdown. |
 | **Offload Prefetch Size** | `--dit-offload-prefetch-size F` | Fine-grained control over layerwise offload: how many layers to prefetch ahead. `0.0` = 1 layer (min VRAM), `0.1` = 10% of layers, `≥1` = absolute layer count. | Tune for cases where default offload has copy stream interference (e.g. image models). 0.05–0.1 is a good starting point. | Values ≥ 0.5 approach no-offload VRAM with worse performance. Use lower values when copy overlap is weak; disable offload when memory allows and latency dominates. |
 | **FSDP Inference** | `--use-fsdp-inference` | Uses PyTorch FSDP to shard model weights across GPUs with prefetch. Low latency, low VRAM. | Reduces per-GPU VRAM | Mutually exclusive with `--dit-layerwise-offload`. More overhead than SP on high-bandwidth interconnects. |
 | **CPU Offload (components)** | `--text-encoder-cpu-offload`, `--image-encoder-cpu-offload`, `--vae-cpu-offload`, `--dit-cpu-offload` | Offloads specific pipeline components to CPU when not in use. | Reduces peak VRAM | Adds H2D transfer latency when the component is needed. Auto-enabled for low-VRAM GPUs (<30 GB). **Tip:** after the first request completes, the console prints a peak VRAM analysis with suggestions on which offload flags can be safely disabled — look for the `"Components that could stay resident"` log line. |
@@ -64,7 +64,7 @@ These options **trade output quality** for speed or VRAM savings. Results will d
 |---|---|---|---|---|
 | **Request Quality Fast Paths** | `--quality high` (`lossless` is default) | Mounts model-owned accelerated DiT/VAE paths that are validated for high quality but are not bit-exact to the reference path. | Model- and shape-specific | Support is per model and may be a no-op. Keep `--quality lossless` as the A/B ground truth. Report aggregate and worst-frame SSIM/PSNR; defaults are 0.95/28 dB for images and 0.92/24 dB for video unless checked-in model metadata overrides them. Do not confuse this with `--output-quality`, which controls file compression. |
 | **Approximate Attention** | Server-wide: `--attention-backend sage_attn` / `sage_attn_3` / `sliding_tile_attn` / `video_sparse_attn` / `sparse_video_gen_2_attn` / `vmoba_attn` / `sla_attn` / `sage_sla_attn`. Per-request (dense drop-ins only): `--attention-backend-override sage_attn` sampling param / API `extra_body` — valid values `fa`, `torch_sdpa`, `sage_attn`, `sage_attn_3`; rejected (with a log) under BCG, torch.compile, sparse server backends, or a non-ring-capable target with ring parallelism. | Replaces exact attention with approximate or sparse variants. `sage_attn`: INT8/FP8 quantized Q·K; `sliding_tile_attn`: spatial-temporal tile skipping; others: model-specific sparse patterns. | ~1.5–2x on attention (varies by backend) | Quality degradation varies by backend and model. `sage_attn` is the most general; sparse backends (`sliding_tile_attn`, `video_sparse_attn`, etc.) are video-model-specific, may require config files (e.g. `--mask-strategy-file-path` for STA), and are server-level only. Requires corresponding packages installed. |
-| **Cache-DiT** | Native: per-request `--enable-cache-dit true\|false` + `--cache-dit-params <json>` (sampling params; also via API `extra_body`). `SGLANG_CACHE_DIT_ENABLED` / `SGLANG_CACHE_DIT_*` env vars are the server-wide defaults for requests that leave them unset. Diffusers backend: `--backend diffusers --cache-dit-config <yaml-or-json>` | Caches intermediate residuals across denoising steps and skips redundant computations via DBCache, TaylorSeer, and optional SCM. | ~1.5-2x on supported models | Quality depends on cache policy. Incompatible with `--dit-layerwise-offload`. Do not pass `--cache-dit-config` for native SGLang tuning unless you are intentionally using the diffusers backend flow. |
+| **Cache-DiT** | Native: per-request `--enable-cache-dit true\|false` + `--cache-dit-params <json>` (sampling params; also via API `extra_body`). `SGLANG_CACHE_DIT_ENABLED` / `SGLANG_CACHE_DIT_*` env vars are the server-wide defaults for requests that leave them unset. Diffusers backend: `--backend diffusers --cache-dit-config <yaml-or-json>` | Caches intermediate residuals across denoising steps and skips redundant computations via DBCache, TaylorSeer, and optional SCM. | ~1.5-2x on supported models | Quality depends on cache policy. Compatible with `--dit-layerwise-offload`: skipped blocks are not streamed, and the first layer after a skip may sync-load. Models that touch every layer before the block loop (for example a full-stack AdaLN prepass) must keep that prepass off while caching. Do not pass `--cache-dit-config` for native SGLang tuning unless you are intentionally using the diffusers backend flow. |
 | **CFG Gating** | Per-request `--cfg-gate-step 0.5` (sampling param; also via API `extra_body`). `SGLANG_DIFFUSION_CFG_GATE_STEP` is the server-wide default (1.0 = off). | After the given fraction of denoising steps, reuses the cached cond-uncond residual instead of running the unconditional branch each step. | Up to ~2x on the gated tail of CFG models (skips one of two branches) | Lossy; no-op without classifier-free guidance or with `--enable-cfg-parallel`. Lower fractions gate earlier and drift more. |
 | **TeaCache** | `--enable-teacache` (uses model sampling presets) | Reuses residuals when adjacent denoising steps are sufficiently similar. | Model- and threshold-dependent | Approximate and model-specific. Mutually exclusive with Spectrum. Fix prompt/seed/shape/steps and validate temporal consistency, not only single frames. |
 | **Spectrum** | `--enable-spectrum` plus optional `--spectrum-*` controls | Forecasts DiT features and skips selected denoising steps. | Defaults target an accuracy/speed tradeoff; aggressive windows can be much faster | Native `sglang generate` only for FLUX.1, Wan, HunyuanVideo, and SD3; not FLUX.2 or server requests. Mutually exclusive with TeaCache. `--debug` adds shadow validation and is not representative latency. |

python/sglang/multimodal_gen/runtime/managers/memory_managers/layerwise_offload.py

@@ -463,6 +463,7 @@ def __init__(
         self._mapped_regions = MappedRegions()
         # Store forward hooks for removal
         self._forward_hooks: List[Any] = []
+        self._last_forwarded_layer: int | None = None  # skip-compute can jump
 
         if initialize:
             self._initialize()
@@ -872,6 +873,9 @@ def prepare_for_next_req(self, non_blocking=True):
         """
         Prepare for the next round of denoising loop with prefetching the necessary layers
         """
+        self._last_forwarded_layer = None
+        self._release_unneeded_streamed_layers(keep=set(self._head_of_stream()))
+
         # The resident set first: it has to be there for the whole step, and the
         # caller decides whether to block on it.
         for layer_idx in sorted(self._retained_set):
@@ -923,6 +927,22 @@ def _next_streamed(self, *, after: int, count: int) -> List[int]:
             for offset in range(min(count, total))
         ]
 
+    def _release_unneeded_streamed_layers(self, *, keep: Set[int]) -> None:
+        """Free streamed layers that are on GPU but not in ``keep`` or resident."""
+        retain = set(self._retained_set) | keep
+        for layer_idx in list(self._gpu_layers):
+            if layer_idx not in retain:
+                self.release_layer(layer_idx)
+
+    def _release_skip_gap(self, *, last_ran: int, next_ran: int) -> None:
+        """Free speculative prefetches in ``(last_ran, next_ran)`` after a jump."""
+        if next_ran <= last_ran + 1:
+            return
+        retain = set(self._retained_set)
+        for layer_idx in range(last_ran + 1, next_ran):
+            if layer_idx not in retain:
+                self.release_layer(layer_idx)
+
     @torch.compiler.disable
     def _activate_residency(self) -> None:
         """Arm the resident set on the first denoise forward. The pinning itself is
@@ -1382,6 +1402,13 @@ def hook(module, input):
                 if i == 0:
                     self._activate_residency()
                     self.prepare_for_next_req(non_blocking=False)
+                elif (
+                    self._last_forwarded_layer is not None
+                    and i > self._last_forwarded_layer + 1
+                ):
+                    self._release_skip_gap(
+                        last_ran=self._last_forwarded_layer, next_ran=i
+                    )
                 if i not in self._gpu_layers:
                     # LTX audio VAE traverses decoder.up in reverse order
                     self.prefetch_layer(i, non_blocking=False)
@@ -1416,6 +1443,7 @@ def make_post_hook(i):
             def hook(module, input, output):
                 # previous, we wait here, until the copy stream for next layer is finished,
                 # now with any prefetch_size, only wait for the copy stream, when the copy stream is for the next layer
+                self._last_forwarded_layer = i
                 self.release_layer(i)
 
             return hook