[diffusion] feat: allow cache-dit with dit layerwise offload (#35858)
This commit is contained in:
WenhaoZhang
@@ -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
|
||||
|
||||
|
||||
@@ -257,7 +257,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.
|
||||
|
||||
@@ -325,7 +325,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:
|
||||
|
||||
|
||||
@@ -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.
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user