Modifications

One-stage

• Disable DiT CPU offload by default on high-memory Hopper GPUs
• Disable VAE CPU offload by default on high-memory Hopper GPUs

Two-stage

• Prebuild and premerge the stage-2 distilled transformer (transformer_2, for snapshot and resident mode)
• Pin the stage-1 transformer on GPU at startup when beneficial
• Replace request-time module.to(cpu/cuda) switching with CPU snapshot based release
• Fix an initialization regression where CPU snapshots redundantly cloned tensors that were already on CPU

Pre-merged LoRA (only applicable for two-stage)

when running two-stage pipeline, LoRA will be applied between two denoising stages, this would include:

1. cuda.synchronize to wait for ongoing tasks
2. layerwise_offload.disable_offload(): move the dit to GPU, h2d(required when dit-layerwise-offload is enabled), ~67.1s
3. apply lora (~18.5s)
4. layerwise_offload.offload(): optional, to resume the offload behavior

while with new approach snapshot mode:

1. h2d (for dit_2), but for dit_1, points param.data directly to cpu_snapshot instead of d2h

for even more aggressive mode resident, all two dits are always kept resident.

Motivation

Statistics

Stage Breakdown

Peak VRAM

NOTES

1. LTX2LoRASwitchStage no longer exists in non-legacy mode
2. low_vram snapshot mode is default enabled on h100, while may be sub-optimal with sp > 1

Modifications

Accuracy Tests

Speed Tests and Profiling

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

Review and Merge Process

1. Ping Merge Oncalls to start the process. See the PR Merge Process.
2. Get approvals from CODEOWNERS and other reviewers.
3. Trigger CI tests with comments or contact authorized users to do so.
   • Common commands include /tag-and-rerun-ci, /tag-run-ci-label, /rerun-failed-ci
4. After green CI and required approvals, ask Merge Oncalls or people with Write permission to merge the PR.