[Diffusion] Add FLUX.2-dev CPU Offload (Layerwise) (#5256)

Signed-off-by: dpeng123 <diyapeng123@gmail.com>

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"` |

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**        |     ❌     |     ✅      |           ❌           |       ❌        |         ✅         |          ❌          |   ❌    |             ❌             |          ❌           |       ✅        |        ❌         |

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

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,

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),