minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 13 Packages Projects Releases Wiki Activity

[diffusion] refactor: unify component residency controls (#34736)

Browse Source
This commit is contained in:
Mick
2026-08-16 11:24:48 +08:00
committed by GitHub
parent f68517f644
commit e9fe58139f
70 changed files with 2629 additions and 1423 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/cookbook/diffusion/Krea/Krea-2.mdx
+7 -5
View File
@@ -245,16 +245,18 @@ SGLANG_CACHE_DIT_MC=3 \
sglang serve --model-path krea/Krea-2-Raw
```
#### 4.2.2 Memory & CPU Offload
#### 4.2.2 Memory and Component Residency
Krea-2's DiT is ~24 GB in bf16 (the bulk of the model). On memory-constrained GPUs you can keep less of it resident:
- `--dit-layerwise-offload`: stream the DiT's transformer blocks layer-by-layer with async host-to-device prefetch overlap, so only a small working set stays on the GPU. This is the primary way to fit Krea-2 on a single consumer / 32 GB-class card, at a modest latency cost. Tune the memory/latency trade-off with `--dit-offload-prefetch-size` (`0.0` prefetches one layer for the lowest memory; larger values prefetch more layers -- faster but more memory).
- `--dit-cpu-offload`: keep the whole DiT in host memory. Combine it with `--dit-layerwise-offload` for the lowest peak GPU memory (weights stay on host and only the layers needed for the current step are brought on-device).
- `--text-encoder-cpu-offload`: offload the Qwen3-VL text encoder (it is idle during the denoise loop).
- `--vae-cpu-offload`: offload the VAE.
- `--component-residency dit=layerwise-offload`: stream the DiT's transformer blocks layer-by-layer with async host-to-device prefetch overlap, so only a small working set stays on the GPU. This is the primary way to fit Krea-2 on a single consumer / 32 GB-class card, at a modest latency cost. Tune the memory/latency trade-off with `--dit-offload-prefetch-size` (`0.0` prefetches one layer for the lowest memory; larger values prefetch more layers -- faster but more memory).
- `--component-residency dit=component-offload`: keep the complete DiT on CPU between denoising uses. This and layerwise offload are distinct modes; do not combine them for the same component.
- `--component-residency text_encoder=component-offload`: offload the Qwen3-VL text encoder while it is idle during denoising.
- `--component-residency vae=component-offload`: offload the VAE between uses.
- `--pin-cpu-memory`: pin host memory for offload. Add only as a temporary workaround if you hit `CUDA error: invalid argument`.
The legacy `--dit-layerwise-offload`, `--dit-cpu-offload`, `--text-encoder-cpu-offload`, and `--vae-cpu-offload` forms remain accepted. If both legacy DiT offload flags are enabled, layerwise offload is the effective DiT mode.
On large-VRAM GPUs (e.g. H200), keep everything resident (offloads off) for the fastest latency.
## 5. Benchmark
docs/cookbook/diffusion/README.mdx
+8
View File
@@ -32,6 +32,14 @@ Each recipe provides step-by-step instructions to help you quickly implement SGL
3. Adapt configurations to your specific hardware and requirements
4. Join our community to share feedback and improvements
For memory placement, prefer the unified
[`--component-residency`](/docs/sglang-diffusion/api/cli#component-residency)
selector. Each component resolves to exactly one of `resident`,
`component-offload`, or `layerwise-offload`. Existing options such as
`--dit-cpu-offload`, `--text-encoder-cpu-offload`,
`--image-encoder-cpu-offload`, and `--vae-cpu-offload` remain supported by all
recipes that already use them.
The sglang diffusion cookbook directory structure are shown below:
```text Example