[diffusion] feat: add an arg for controlling the number of prefetched layers in layerwise-offload (#17693)

python/sglang/multimodal_gen/test/server/testcase_configs.py

@@ -164,6 +164,7 @@ class DiffusionServerArgs:
     enable_warmup: bool = False
 
     dit_layerwise_offload: bool = False
+    dit_offload_prefetch_size: int | float | None = None
     enable_cache_dit: bool = False
     text_encoder_cpu_offload: bool = False
 
@@ -369,6 +370,7 @@ ONE_GPU_CASES_A: list[DiffusionTestCase] = [
             model_path="black-forest-labs/FLUX.2-dev",
             modality="image",
             dit_layerwise_offload=True,
+            dit_offload_prefetch_size=2,
         ),
         T2I_sampling_params,
     ),