[diffusion] feat: allow cache-dit with dit layerwise offload (#35858)
This commit is contained in:
WenhaoZhang
@@ -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. |
|
||||
|
||||
@@ -463,6 +463,7 @@ class LayerwiseOffloadManager:
|
||||
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 @@ class LayerwiseOffloadManager:
|
||||
"""
|
||||
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 @@ class LayerwiseOffloadManager:
|
||||
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 @@ class LayerwiseOffloadManager:
|
||||
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 @@ class LayerwiseOffloadManager:
|
||||
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
|
||||
|
||||
@@ -2316,9 +2316,10 @@ class ServerArgs(DisaggServerArgsMixin):
|
||||
action=StoreBoolean,
|
||||
default=ServerArgs.dit_layerwise_offload,
|
||||
help="Enable layerwise CPU offload with async H2D prefetch overlap for DiTs. "
|
||||
"It selects only the DiT layerwise group. Cannot be used together with cache-dit "
|
||||
"(SGLANG_CACHE_DIT_ENABLED) or use_fsdp_inference. If legacy DiT offload "
|
||||
"flags are also provided, layerwise offload is the effective DiT mode.",
|
||||
"It selects only the DiT layerwise group. Compatible with cache-dit: "
|
||||
"skipped blocks are not streamed. Cannot be used together with "
|
||||
"use_fsdp_inference. If legacy DiT offload flags are also provided, "
|
||||
"layerwise offload is the effective DiT mode.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--layerwise-offload-components",
|
||||
@@ -3404,17 +3405,6 @@ class ServerArgs(DisaggServerArgsMixin):
|
||||
if self.dit_offload_prefetch_size < 0.0:
|
||||
raise ValueError("dit_offload_prefetch_size must be non-negative")
|
||||
|
||||
is_dit_layerwise_offload_selected = self.is_dit_layerwise_offload_selected
|
||||
|
||||
if envs.SGLANG_CACHE_DIT_ENABLED and is_dit_layerwise_offload_selected:
|
||||
raise ValueError(
|
||||
"DiT layerwise offload cannot be enabled together with cache-dit. "
|
||||
"cache-dit may reuse skipped blocks whose weights have been released by layerwise offload, "
|
||||
"causing shape mismatch errors. "
|
||||
"Please disable --dit-layerwise-offload, remove DiT from --layerwise-offload-components, "
|
||||
"or disable SGLANG_CACHE_DIT_ENABLED."
|
||||
)
|
||||
|
||||
if (
|
||||
self.performance_mode == "memory"
|
||||
or self.is_arg_explicitly_set("layerwise_offload_components")
|
||||
|
||||
@@ -1883,3 +1883,182 @@ def test_park_placeholders_are_shared(monkeypatch):
|
||||
id(p) for n, p in comp.named_parameters() if n not in managed and p.numel() == 1
|
||||
}
|
||||
assert len(stand_ins) <= len(comp._park_placeholders)
|
||||
|
||||
|
||||
def _layer_weight_ok(layer: torch.nn.Module) -> bool:
|
||||
return tuple(layer.weight.shape) != (1,)
|
||||
|
||||
|
||||
def test_skip_middle_layers_loads_destination_weights(monkeypatch):
|
||||
"""Cache-DiT DBCache shape: run Fn, jump to Bn, middle never forwards.
|
||||
|
||||
The destination layer used to see empty(1,) weights when wraparound
|
||||
prefetch and the sequential i+1 window desynced. The jump must sync-load
|
||||
Bn and leave the skipped gap released.
|
||||
"""
|
||||
_patch_fake_device(monkeypatch)
|
||||
model = _RunnableBlockModel(8)
|
||||
manager = _resident_manager(model, num_layers=8, prefetch_size=1)
|
||||
|
||||
hidden = torch.ones(1, 2)
|
||||
fn_end = 1
|
||||
bn_start = 6
|
||||
for layer in model.blocks[:fn_end]:
|
||||
hidden = layer(hidden)
|
||||
for layer in model.blocks[bn_start:]:
|
||||
hidden = layer(hidden)
|
||||
|
||||
assert hidden.shape == (1, 2)
|
||||
for idx in range(fn_end, bn_start):
|
||||
assert idx not in manager._gpu_layers, idx
|
||||
assert not _layer_weight_ok(model.blocks[idx]), idx
|
||||
|
||||
|
||||
def test_skip_only_fn_releases_speculative_prefetch_on_next_step(monkeypatch):
|
||||
"""Fn-only step (Cache-DiT hit with Bn=0) must not leak the i+1 prefetch."""
|
||||
_patch_fake_device(monkeypatch)
|
||||
model = _RunnableBlockModel(8)
|
||||
manager = _resident_manager(model, num_layers=8, prefetch_size=1)
|
||||
|
||||
hidden = torch.ones(1, 2)
|
||||
hidden = model.blocks[0](hidden)
|
||||
# Layer 0's leading burst prefetches layer 1; that layer never runs.
|
||||
assert 1 in manager._gpu_layers
|
||||
|
||||
# Next denoise step's prepare drops leftovers that never posted.
|
||||
manager.prepare_for_next_req(non_blocking=False)
|
||||
assert 1 not in manager._gpu_layers
|
||||
|
||||
|
||||
def test_last_layer_wraps_to_next_step_head(monkeypatch):
|
||||
"""A full-stack step may hide the next step's layer 0 behind the last layer."""
|
||||
_patch_fake_device(monkeypatch)
|
||||
model = _RunnableBlockModel(8)
|
||||
manager = _resident_manager(model, num_layers=8, prefetch_size=1)
|
||||
|
||||
hidden = torch.ones(1, 2)
|
||||
for layer in model.blocks:
|
||||
hidden = layer(hidden)
|
||||
|
||||
assert hidden.shape == (1, 2)
|
||||
assert 0 in manager._gpu_layers
|
||||
assert 7 not in manager._gpu_layers
|
||||
|
||||
|
||||
def _dbcache_layers(num_layers: int, fn: int, bn: int) -> list[int]:
|
||||
"""Layers CachedBlocks would call for one DBCache step."""
|
||||
fn = min(max(fn, 0), num_layers)
|
||||
bn = min(max(bn, 0), num_layers - fn)
|
||||
layers = list(range(fn))
|
||||
if bn:
|
||||
layers.extend(range(num_layers - bn, num_layers))
|
||||
return layers
|
||||
|
||||
|
||||
def _run_layer_set(model, layer_indices: list[int]) -> torch.Tensor:
|
||||
hidden = torch.ones(1, 2)
|
||||
for idx in layer_indices:
|
||||
hidden = model.blocks[idx](hidden)
|
||||
return hidden
|
||||
|
||||
|
||||
@pytest.mark.parametrize("num_layers", [8, 12])
|
||||
@pytest.mark.parametrize(
|
||||
"fn,bn",
|
||||
[
|
||||
(1, 0), # default Cache-DiT hit
|
||||
(1, 1),
|
||||
(1, 2),
|
||||
(2, 0),
|
||||
(2, 2),
|
||||
(4, 2),
|
||||
(3, 5), # Fn+Bn == 8, no gap when num_layers=8
|
||||
(8, 0), # full stack / miss
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("prefetch_size", [1, 2])
|
||||
@pytest.mark.parametrize(
|
||||
"residency_policy",
|
||||
[RESIDENCY_POLICY_LEADING, RESIDENCY_POLICY_STRIDED],
|
||||
)
|
||||
def test_dbcache_layer_patterns_never_see_empty_weights(
|
||||
monkeypatch, num_layers, fn, bn, prefetch_size, residency_policy
|
||||
):
|
||||
"""Hit / miss / hit-again under several Cache-DiT Fn/Bn and prefetch windows."""
|
||||
if fn + bn > num_layers:
|
||||
pytest.skip("Fn+Bn exceeds this stack")
|
||||
_patch_fake_device(monkeypatch)
|
||||
model = _RunnableBlockModel(num_layers)
|
||||
manager = _resident_manager(
|
||||
model,
|
||||
num_layers=num_layers,
|
||||
prefetch_size=prefetch_size,
|
||||
residency_policy=residency_policy,
|
||||
resident_layers=0,
|
||||
)
|
||||
|
||||
hit_layers = _dbcache_layers(num_layers, fn, bn)
|
||||
miss_layers = list(range(num_layers))
|
||||
gap = [idx for idx in miss_layers if idx not in set(hit_layers)]
|
||||
|
||||
def _assert_gpu_layers_have_real_weights() -> None:
|
||||
for idx in range(num_layers):
|
||||
on_gpu = idx in manager._gpu_layers
|
||||
assert _layer_weight_ok(model.blocks[idx]) is on_gpu, idx
|
||||
|
||||
hidden = _run_layer_set(model, hit_layers)
|
||||
assert hidden.shape == (1, 2)
|
||||
_assert_gpu_layers_have_real_weights()
|
||||
|
||||
# Speculative Mn prefetch may still sit on GPU until the next prepare.
|
||||
manager.prepare_for_next_req(non_blocking=False)
|
||||
keep = set(manager._head_of_stream()) | set(manager._retained_set)
|
||||
for idx in gap:
|
||||
if idx not in keep:
|
||||
assert idx not in manager._gpu_layers, idx
|
||||
assert not _layer_weight_ok(model.blocks[idx]), idx
|
||||
|
||||
hidden = _run_layer_set(model, miss_layers)
|
||||
assert hidden.shape == (1, 2)
|
||||
_assert_gpu_layers_have_real_weights()
|
||||
|
||||
manager.prepare_for_next_req(non_blocking=False)
|
||||
hidden = _run_layer_set(model, hit_layers)
|
||||
assert hidden.shape == (1, 2)
|
||||
_assert_gpu_layers_have_real_weights()
|
||||
|
||||
# Two hits in a row (Bn=0 never reaches last layer; still must rematerialize 0).
|
||||
manager.prepare_for_next_req(non_blocking=False)
|
||||
hidden = _run_layer_set(model, hit_layers)
|
||||
assert hidden.shape == (1, 2)
|
||||
_assert_gpu_layers_have_real_weights()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"step_kinds",
|
||||
[
|
||||
# SCM-style: forced compute (full stack) mixed with DBCache hits.
|
||||
("full", "hit10", "hit10", "full", "hit12", "hit10"),
|
||||
# TaylorSeer does not change which blocks CachedBlocks calls.
|
||||
("full", "full", "hit20", "hit20", "hit12", "full"),
|
||||
],
|
||||
)
|
||||
def test_mixed_scm_and_dbcache_step_schedule(monkeypatch, step_kinds):
|
||||
"""A request is a sequence of full-stack and skip-compute steps."""
|
||||
_patch_fake_device(monkeypatch)
|
||||
num_layers = 8
|
||||
model = _RunnableBlockModel(num_layers)
|
||||
manager = _resident_manager(model, num_layers=num_layers, prefetch_size=2)
|
||||
kind_to_layers = {
|
||||
"full": list(range(num_layers)),
|
||||
"hit10": _dbcache_layers(num_layers, 1, 0),
|
||||
"hit12": _dbcache_layers(num_layers, 1, 2),
|
||||
"hit20": _dbcache_layers(num_layers, 2, 0),
|
||||
}
|
||||
for kind in step_kinds:
|
||||
hidden = _run_layer_set(model, kind_to_layers[kind])
|
||||
assert hidden.shape == (1, 2)
|
||||
for idx in range(num_layers):
|
||||
on_gpu = idx in manager._gpu_layers
|
||||
assert _layer_weight_ok(model.blocks[idx]) is on_gpu, (kind, idx)
|
||||
manager.prepare_for_next_req(non_blocking=False)
|
||||
|
||||
@@ -1644,6 +1644,20 @@ class TestOffloadDefaults(unittest.TestCase):
|
||||
self.assertFalse(args.use_fsdp_inference)
|
||||
self.assertTrue(args.enable_cfg_parallel)
|
||||
|
||||
def test_cache_dit_allows_explicit_dit_layerwise_offload(self):
|
||||
with patch.dict(os.environ, {"SGLANG_CACHE_DIT_ENABLED": "true"}):
|
||||
args = self._from_dict_with_pipeline_config(
|
||||
QwenImagePipelineConfig(),
|
||||
kwargs={
|
||||
"model_path": "/data/my-model",
|
||||
"performance_mode": "manual",
|
||||
"dit_layerwise_offload": True,
|
||||
},
|
||||
)
|
||||
|
||||
self.assertTrue(args.is_dit_layerwise_offload_selected)
|
||||
self.assertEqual(args.layerwise_offload_components, ["dit"])
|
||||
|
||||
def test_auto_multi_gpu_sana_wm_realtime_disables_cfg_parallel(self):
|
||||
args = self._from_dict_with_pipeline_config(
|
||||
SanaWMRealtimeConfig(),
|
||||
|
||||
Reference in New Issue
Block a user