[diffusion] refactor: introduce component residency manager

[mickqian]

Motivation

Component management is crucial for the latency of diffusion pipeline inference, especially with modern diffusion models requiring larger sub-components (mistral, qwen as encoders) and more components (dual-dit for Wan and LTX) which have parameter size sum up to be larger than most modern GPUs.

Currently the model management code is scattered in each pipeline's pre and post hooks, making it hard to apply advance coordination.

Modifications

1. introduce ComponentResidencyStrategy to abstract and cover all pre-existing module management techniques (including layerwise-offload, snapshot and resident mode for LTX pre-merged LoRA)
2. introduce the ComponentResidencyManager to serve as a global manager to coordinate the components' placement, to maximize the latency while making full use of VRAM. The manager calls the preparation, pre-fetch and release hooks defined by the aforementioned strategies to coordinate

Accuracy Tests

Speed Tests and Profiling

Checklist

TODO

1. use warmup request to build the component relationship as a graph, for manager to coordinate the module placement in a smaller granularity
2. provide argument for user to switch between different memory mode (balanced, high-vram, low-vram)

• 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.

[gemini-code-assist]

Code Review

This pull request introduces a centralized PipelineResidencyManager to coordinate the loading and offloading of model components (such as VAEs, DiTs, and text encoders) across different pipeline stages. This system replaces legacy stage-local behavior with configurable residency strategies (static, dynamic, or disabled) and integrates hooks into the pipeline executors to manage component lifecycles. Additionally, the PR adds a compatibility wrapper for Flash Attention v3 kernels and updates various stages to declare their component usage. Feedback focuses on a potential type error in text encoding attention masks, the lack of prefetching in the vanilla D2H strategy, and the use of hardcoded memory buffer values.

[gemini-code-assist]

Prefetching is explicitly disabled for VanillaD2HStrategy. This strategy is used for components like text encoders, which can be large and would benefit from asynchronous H2D transfers to hide latency. Unless there is a specific reason to avoid overlapping these transfers (e.g., memory pressure concerns that aren't already handled by the manager), prefetching should be enabled for this strategy as well.

[gemini-code-assist]

The memory buffer value 2.0 (presumably GB) is hardcoded. It would be better to define this as a named constant or make it configurable via ServerArgs to improve maintainability and allow tuning for different hardware or workloads.

[mickqian]

/tag-and-rerun-ci