diff --git a/docs/cookbook/diffusion/MiniMax/MiniMax-H3.mdx b/docs/cookbook/diffusion/MiniMax/MiniMax-H3.mdx index 90cb3f0ff..ef92b2efc 100644 --- a/docs/cookbook/diffusion/MiniMax/MiniMax-H3.mdx +++ b/docs/cookbook/diffusion/MiniMax/MiniMax-H3.mdx @@ -765,9 +765,11 @@ name, for example `sglang generate --quality high`. 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 \ ``` -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. ## 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 diff --git a/docs/docs/sglang-diffusion/api/cli.mdx b/docs/docs/sglang-diffusion/api/cli.mdx index 4ce6a6d21..2fe237d94 100644 --- a/docs/docs/sglang-diffusion/api/cli.mdx +++ b/docs/docs/sglang-diffusion/api/cli.mdx @@ -257,7 +257,7 @@ sglang generate \ HTTP server-only arguments are ignored by `sglang generate`. -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: diff --git a/docs/docs/sglang-diffusion/cache_dit.mdx b/docs/docs/sglang-diffusion/cache_dit.mdx index 14d62e0ca..1cffcd6f9 100644 --- a/docs/docs/sglang-diffusion/cache_dit.mdx +++ b/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. diff --git a/docs/docs/sglang-diffusion/caching-acceleration.mdx b/docs/docs/sglang-diffusion/caching-acceleration.mdx index bdcdd041f..041898339 100644 --- a/docs/docs/sglang-diffusion/caching-acceleration.mdx +++ b/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. -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. ### Quick Start diff --git a/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-performance/SKILL.md b/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-performance/SKILL.md index c55b006b7..d922fc792 100644 --- a/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-performance/SKILL.md +++ b/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 ` (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 ` | 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 ` (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 ` | 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. | diff --git a/python/sglang/multimodal_gen/runtime/managers/memory_managers/layerwise_offload.py b/python/sglang/multimodal_gen/runtime/managers/memory_managers/layerwise_offload.py index fb1e37831..48c6619cc 100644 --- a/python/sglang/multimodal_gen/runtime/managers/memory_managers/layerwise_offload.py +++ b/python/sglang/multimodal_gen/runtime/managers/memory_managers/layerwise_offload.py @@ -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 diff --git a/python/sglang/multimodal_gen/runtime/server_args/server_args.py b/python/sglang/multimodal_gen/runtime/server_args/server_args.py index ae457e7c0..9034c562c 100644 --- a/python/sglang/multimodal_gen/runtime/server_args/server_args.py +++ b/python/sglang/multimodal_gen/runtime/server_args/server_args.py @@ -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") diff --git a/python/sglang/multimodal_gen/test/unit/test_layerwise_offload.py b/python/sglang/multimodal_gen/test/unit/test_layerwise_offload.py index 1fbc65cbc..83fce212b 100644 --- a/python/sglang/multimodal_gen/test/unit/test_layerwise_offload.py +++ b/python/sglang/multimodal_gen/test/unit/test_layerwise_offload.py @@ -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) diff --git a/python/sglang/multimodal_gen/test/unit/test_server_args.py b/python/sglang/multimodal_gen/test/unit/test_server_args.py index 646c75b90..38f2e757f 100644 --- a/python/sglang/multimodal_gen/test/unit/test_server_args.py +++ b/python/sglang/multimodal_gen/test/unit/test_server_args.py @@ -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(),