[diffusion] Generalize layerwise offload residency mixin to all components

[mickqian]

Summary

• Rename the DiT-specific layerwise offload mixin to LayerwiseOffloadableModuleMixin.
• Resolve layerwise residency by module capability before falling back to existing component CPU-offload flags.
• Add --layerwise-offload-components to select layerwise offload by pipeline component name, with --layerwise-offload-modules accepted as an alias.
• Keep --dit-layerwise-offload legacy behavior: when no component is named, only default DiT components are configured; encoder / VAE / bridge / upsampler / vocoder must be selected explicitly.
• Disable conflicting component CPU/FSDP offload flags when the same component is explicitly selected for layerwise offload, and keep layer buffers resident to avoid releasing shared buffers such as RoPE caches.

Validation

Remote H200 container /sgl-workspace/sglang at fdf022713606dc2e6262975145d94e7f7d504a0d:

• PYTHONPATH=/sgl-workspace/sglang/python python -m pytest python/sglang/multimodal_gen/test/unit/test_layerwise_offload.py python/sglang/multimodal_gen/test/unit/test_server_args.py -> 41 passed, 2 warnings
• Z-Image 512x512, 1 step, seed 0, torch_sdpa backend:
  • --dit-layerwise-offload true --dit-offload-prefetch-size 0 -> PASS, enabled ['transformer']
  • --layerwise-offload-components transformer --dit-offload-prefetch-size 0 -> PASS, enabled ['transformer']
  • --layerwise-offload-components transformer text_encoder --dit-offload-prefetch-size 0 -> PASS, enabled ['text_encoder', 'transformer']
  • --layerwise-offload-components all --dit-offload-prefetch-size 0 -> PASS, enabled ['text_encoder', 'vae', 'transformer']
  • --layerwise-offload-modules transformer --dit-offload-prefetch-size 0 -> PASS, enabled ['transformer']
  • --layerwise-offload-components missing_component --dit-offload-prefetch-size 0 -> PASS with warning and no layerwise component
• All six generated PNG files have identical sha256: 0f0dba3e7d97aa3be19ef7d6d1cd3ea0e727c322153a8a4f9904089b8e9ee4c1
• No local tests were run.

---

CI States

Latest PR Test: Run #25928755979
Latest PR Test (Extra): ⚠️ Not enabled — add run-ci-extra label to opt in.

[gemini-code-assist]

Code Review

This pull request refactors the layerwise offload mechanism by renaming OffloadableDiTMixin to LayerwiseOffloadableModuleMixin and centralizing residency strategy logic in component_manager.py. It introduces helper functions like is_layerwise_offloaded_module and should_cpu_offload_component to simplify offload decisions across various model components. Feedback includes suggestions to remove a redundant bool() call, simplify multi-line tuple unpacking for better readability, and remove an unnecessary trailing comma in a tuple assignment.

[gemini-code-assist]

The bool() call here is redundant. In Python, an empty list is evaluated as False in a boolean context, so you can check for non-emptiness directly. The pythonic way is to use the list itself in the condition.

Suggested change

	and bool(module.layerwise_offload_managers)
	and module.layerwise_offload_managers

[gemini-code-assist]

This multi-line tuple unpacking seems unnecessary as the line is not excessively long. It could be simplified to a single line for better readability and to reduce vertical space.

        (shift_msa, scale_msa, gate_msa), (shift_mlp, scale_mlp, gate_mlp) = temb_mod_params_img

[gemini-code-assist]

The trailing comma in this tuple unpacking is unnecessary. While valid syntax, it's typically used to define a single-element tuple. For multi-element tuples, it's unconventional and can be removed for clarity.

Suggested change

	(x, cap_feats, x_size, x_valid_lens, cap_valid_lens,) = self.patchify_and_embed(
	(x, cap_feats, x_size, x_valid_lens, cap_valid_lens) = self.patchify_and_embed(