[diffusion] feat: generalize layerwise offloader to flux1 (#15633)

python/sglang/multimodal_gen/runtime/loader/component_loader.py

@@ -712,16 +712,17 @@ class TransformerLoader(ComponentLoader):
 
         model = model.eval()
 
-        if server_args.dit_layerwise_offload and hasattr(model, "blocks"):
+        if server_args.dit_layerwise_offload and hasattr(model, "dit_module_names"):
+            # TODO(will): support multiple module names
+            module_name = getattr(model, "dit_module_names", ["transformer_blocks"])[0]
             try:
-                num_layers = len(getattr(model, "blocks"))
+                num_layers = len(getattr(model, module_name))
             except Exception:
                 num_layers = None
-
             if isinstance(num_layers, int) and num_layers > 0:
                 mgr = LayerwiseOffloadManager(
                     model,
-                    module_list_attr="blocks",
+                    module_list_attr=module_name,
                     num_layers=num_layers,
                     enabled=True,
                     pin_cpu_memory=server_args.pin_cpu_memory,

python/sglang/multimodal_gen/runtime/models/dits/flux.py

@@ -391,6 +391,10 @@ class FluxTransformer2DModel(CachableDiT):
         self.inner_dim = (
             self.config.num_attention_heads * self.config.attention_head_dim
         )
+        self.dit_module_names = [
+            "transformer_blocks",
+            "single_transformer_blocks",
+        ]
 
         self.rotary_emb = FluxPosEmbed(theta=10000, axes_dim=self.config.axes_dims_rope)
 
@@ -496,23 +500,46 @@ class FluxTransformer2DModel(CachableDiT):
             ip_hidden_states = self.encoder_hid_proj(ip_adapter_image_embeds)
             joint_attention_kwargs.update({"ip_hidden_states": ip_hidden_states})
 
-        for index_block, block in enumerate(self.transformer_blocks):
-            encoder_hidden_states, hidden_states = block(
-                hidden_states=hidden_states,
-                encoder_hidden_states=encoder_hidden_states,
-                temb=temb,
-                freqs_cis=freqs_cis,
-                joint_attention_kwargs=joint_attention_kwargs,
-            )
-
-        for index_block, block in enumerate(self.single_transformer_blocks):
-            encoder_hidden_states, hidden_states = block(
-                hidden_states=hidden_states,
-                encoder_hidden_states=encoder_hidden_states,
-                temb=temb,
-                freqs_cis=freqs_cis,
-                joint_attention_kwargs=joint_attention_kwargs,
-            )
+        offload_mgr = getattr(self, "_layerwise_offload_manager", None)
+        if offload_mgr is not None and getattr(offload_mgr, "enabled", False):
+            for i, block in enumerate(self.transformer_blocks):
+                with offload_mgr.layer_scope(
+                    prefetch_layer_idx=i + 1,
+                    release_layer_idx=i,
+                    non_blocking=True,
+                ):
+                    encoder_hidden_states, hidden_states = block(
+                        hidden_states=hidden_states,
+                        encoder_hidden_states=encoder_hidden_states,
+                        temb=temb,
+                        freqs_cis=freqs_cis,
+                        joint_attention_kwargs=joint_attention_kwargs,
+                    )
+            for block in self.single_transformer_blocks:
+                encoder_hidden_states, hidden_states = block(
+                    hidden_states=hidden_states,
+                    encoder_hidden_states=encoder_hidden_states,
+                    temb=temb,
+                    freqs_cis=freqs_cis,
+                    joint_attention_kwargs=joint_attention_kwargs,
+                )
+        else:
+            for block in self.transformer_blocks:
+                encoder_hidden_states, hidden_states = block(
+                    hidden_states=hidden_states,
+                    encoder_hidden_states=encoder_hidden_states,
+                    temb=temb,
+                    freqs_cis=freqs_cis,
+                    joint_attention_kwargs=joint_attention_kwargs,
+                )
+            for block in self.single_transformer_blocks:
+                encoder_hidden_states, hidden_states = block(
+                    hidden_states=hidden_states,
+                    encoder_hidden_states=encoder_hidden_states,
+                    temb=temb,
+                    freqs_cis=freqs_cis,
+                    joint_attention_kwargs=joint_attention_kwargs,
+                )
 
         hidden_states = self.norm_out(hidden_states, temb)
 

python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py

@@ -610,6 +610,7 @@ class WanTransformer3DModel(CachableDiT):
         self.num_channels_latents = config.num_channels_latents
         self.patch_size = config.patch_size
         self.text_len = config.text_len
+        self.dit_module_names = ["blocks"]
 
         # 1. Patch & position embedding
         self.patch_embedding = PatchEmbed(