From b471a5cf236cff2419007549c1867a14774e09a6 Mon Sep 17 00:00:00 2001 From: Diya Peng <50388704+dpeng123@users.noreply.github.com> Date: Wed, 22 Jul 2026 23:53:04 -0700 Subject: [PATCH] [Diffusion] Add FLUX.2-dev CPU Offload (Layerwise) (#5256) Signed-off-by: dpeng123 --- .../diffusion/cpu_offload_diffusion.md | 1 + docs/user_guide/diffusion_features.md | 2 +- .../models/flux2/test_flux2_transformer_tp.py | 25 +++++++++++++++++++ .../test_flux_2_dev_expansion.py | 13 +++++++++- .../models/flux2/flux2_transformer.py | 1 + 5 files changed, 40 insertions(+), 2 deletions(-) diff --git a/docs/user_guide/diffusion/cpu_offload_diffusion.md b/docs/user_guide/diffusion/cpu_offload_diffusion.md index d6fcd8944..31a2670ba 100644 --- a/docs/user_guide/diffusion/cpu_offload_diffusion.md +++ b/docs/user_guide/diffusion/cpu_offload_diffusion.md @@ -223,6 +223,7 @@ reasoner/generator components inside the model forward pass. | Architecture | Example Models | DiT Class | Model-Level Offload | Layerwise Offload | Blocks Attrs (Layerwise specific) | |--------------|----------------|-----------|---------------------|-------------------|-----------------------------------| +| Flux2Pipeline | `black-forest-labs/FLUX.2-dev` | `Flux2Transformer2DModel` | ✓ | ✓ | `"transformer_blocks"`, `"single_transformer_blocks"` | | LongCatImagePipeline | `meituan-longcat/LongCat-Image` | `LongCatImageTransformer2DModel` | - | ✓ | `"transformer_blocks"`, `"single_transformer_blocks"` | | NextStep11Pipeline | `stepfun-ai/NextStep-1.1` | `NextStepModel` | - | ✓ | `"layers"` | | OvisImagePipeline | `AIDC-AI/Ovis-Image-7B` | `OvisImageTransformer2DModel` | - | ✓ | `"transformer"` | diff --git a/docs/user_guide/diffusion_features.md b/docs/user_guide/diffusion_features.md index e5da3c294..bf714b3a0 100644 --- a/docs/user_guide/diffusion_features.md +++ b/docs/user_guide/diffusion_features.md @@ -117,7 +117,7 @@ The following tables show which models support each feature: | **FLUX.1-schnell** | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | | **FLUX.2-klein** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | | **FLUX.1-Kontext-dev** | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | -| **FLUX.2-dev** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | +| **FLUX.2-dev** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | | **GLM-Image** | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | | **Hidream-I1-Full** | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | | **HunyuanImage3** | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | diff --git a/tests/diffusion/models/flux2/test_flux2_transformer_tp.py b/tests/diffusion/models/flux2/test_flux2_transformer_tp.py index c613bb0b4..c935871a0 100644 --- a/tests/diffusion/models/flux2/test_flux2_transformer_tp.py +++ b/tests/diffusion/models/flux2/test_flux2_transformer_tp.py @@ -253,3 +253,28 @@ class TestFlux2PackedModuleMapping: assert len(loaded_params) == 2, "Should load two weights" assert "x_embedder.weight" in loaded_params assert any("to_qkv" in p for p in loaded_params), "to_q should be mapped to to_qkv" + + +@pytest.mark.core_model +@pytest.mark.cpu +class TestFlux2TransformerLayerwiseOffload: + def test_layerwise_offload_blocks_attrs(self): + assert Flux2Transformer2DModel._layerwise_offload_blocks_attrs == [ + "transformer_blocks", + "single_transformer_blocks", + ] + + def test_get_blocks_from_dit(self): + from vllm_omni.diffusion.offloader.layerwise_backend import LayerWiseOffloadBackend + + # Block discovery only needs the ModuleList structure, not production dims. + model = Flux2Transformer2DModel( + num_layers=2, + num_single_layers=2, + num_attention_heads=2, + attention_head_dim=4, + joint_attention_dim=16, + ) + attr_names, blocks = LayerWiseOffloadBackend.get_blocks_from_dit(model) + assert attr_names == ["transformer_blocks", "single_transformer_blocks"] + assert len(blocks) == 4 diff --git a/tests/e2e/online_serving/test_flux_2_dev_expansion.py b/tests/e2e/online_serving/test_flux_2_dev_expansion.py index 77a7a9271..868a1e3bd 100644 --- a/tests/e2e/online_serving/test_flux_2_dev_expansion.py +++ b/tests/e2e/online_serving/test_flux_2_dev_expansion.py @@ -2,7 +2,8 @@ End-to-end diffusion coverage for FLUX.2-dev in online serving mode. Coverage: -- CPU offload +- CPU offload (model-level) +- Layerwise CPU offload - Ulysses sequence parallelism - Ring sequence parallelism @@ -43,6 +44,16 @@ def _get_flux_2_dev_feature_cases(model: str): id="cpu_offload", marks=SINGLE_CARD_FEATURE_MARKS, ), + pytest.param( + OmniServerParams( + model=model, + server_args=[ + "--enable-layerwise-offload", + ], + ), + id="layerwise_offload", + marks=SINGLE_CARD_FEATURE_MARKS, + ), pytest.param( OmniServerParams( model=model, diff --git a/vllm_omni/diffusion/models/flux2/flux2_transformer.py b/vllm_omni/diffusion/models/flux2/flux2_transformer.py index 52e0a4430..2beb1fadb 100644 --- a/vllm_omni/diffusion/models/flux2/flux2_transformer.py +++ b/vllm_omni/diffusion/models/flux2/flux2_transformer.py @@ -765,6 +765,7 @@ class Flux2Transformer2DModel(nn.Module): ) _repeated_blocks = ["Flux2TransformerBlock", "Flux2SingleTransformerBlock"] + _layerwise_offload_blocks_attrs = ["transformer_blocks", "single_transformer_blocks"] _sp_plan = { "": { "hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3, auto_pad=True),