feat(lingbot): compile causal fast DiT

Configure LingBot DiT through ModelRuntimeConfig, enable compile in the v2 H100 example, and retain worker-local KV cache state across causal chunks.

Add fixed 2/4/5/6-GPU VAE and DiT placement strategies, document compile benchmark results, and cover the runtime configuration and placement mappings with focused tests.

Verification: .venv/bin/python -m pytest tests/unit/pipelines/lingbot_world_fast/test_module_loading.py tests/unit/pipelines/lingbot_world_fast/test_runtime_baseline.py tests/unit/pipelines/lingbot_world_fast/test_parallelism.py tests/unit/pipelines/lingbot_world_fast/test_streaming.py tests/unit/pipelines/lingbot_world_fast/test_stream_example.py tests/unit/pipelines/lingbot_world_v2/test_service.py -q; git diff --check.

Author: lzx1413

examples/lingbot/README.md

@@ -129,7 +129,7 @@ python examples/lingbot/lingbot_world_fast_image_to_video_h100.py \
 
 | Option | Default | Description |
 | --- | --- | --- |
-| `--gpu_num` | `1` | Number of GPUs used for Ulysses sequence parallelism |
+| `--gpu_num` | `1` | Total visible GPUs; selects the LingBot VAE and DiT placement strategy |
 | `--model_root` | `${TF_MODEL_ZOO_PATH}/Wan2.2-I2V-A14B` | Wan2.2 I2V base model directory |
 | `--fast_model_root` | `${TF_MODEL_ZOO_PATH}/lingbot/lingbot-world-fast` | LingBot-World-Fast model directory |
 | `--image_path` | Bundled `image.jpg` | Input image path |
@@ -166,13 +166,33 @@ through `PPL_CONFIG`; `stream-serve --gpu-num` is passed to `get_service(gpu_num
 
 ### Scheduler and Stage Placement
 
-LingBot offline and stream-server execution share the actor-based streaming scheduler.
-The bundled examples place VAE encode and decode on GPU 0, while direct
-`LingBotWorldFastPipelineConfig` users may set `vae_encode_config` and
-`vae_decode_config` independently. The scheduler does not infer a resource
-group from overlapping device IDs, so VAE encode, DiT, and VAE decode may
-overlap on one GPU. If a topology runs out of memory, move stages to different
-devices.
+LingBot offline and stream-server execution share the actor-based streaming scheduler. The bundled examples use
+fixed placement for the following total GPU counts:
+
+| Total GPUs | DiT GPUs | VAE encode GPU | VAE decode GPU |
+| --- | --- | --- | --- |
+| 2 | `0-1` | `0` | `1` |
+| 4 | `0-3` | `0` | `1` |
+| 5 | `0-3` | `4` | `4` |
+| 6 | `0-4` | `5` | `5` |
+
+For other counts, the examples retain the PPL-configured VAE devices and assign all visible GPUs to DiT. Direct
+`LingBotWorldFastPipelineConfig` users may set `vae_encode_config`, `vae_decode_config`, and `dit_config` independently.
+
+### H100 Compile Benchmark
+
+The v2 example was measured at 480p (832x464 internal size), 77 output frames, five latent chunks of four frames,
+BF16 DiT, FP32 VAE, SageAttention SM90, `torch.compile` enabled, and FSDP disabled. Each value is the mean from a
+second session after a complete warmup session. Pure DiT measures synchronous `denoise_and_update_cache`; chunk
+period is the mean interval between decoded chunk outputs while encode, DiT, and decode overlap.
+
+| Total H100 GPUs | Pure DiT seconds/chunk | Overlapped chunk period seconds/chunk |
+| --- | --- | --- |
+| 2 | 1.587 | 2.096 |
+| 4 | 0.911 | 1.615 |
+
+The scheduler does not infer a resource group from overlapping device IDs, so VAE encode, DiT, and VAE decode may
+overlap on a shared GPU.
 
 See the [streaming scheduler guide](../../docs/en/stream_scheduler.md) for
 architecture, metric definitions, and lifecycle guarantees.

examples/lingbot/lingbot_world_fast_image_to_video_h100.py

@@ -63,7 +63,7 @@
     target_fps=16,
     max_duration_seconds=5.0,
     attn_impl=AttnImplType.SAGE_ATTN_2_8_8_SM90,
-    enable_fsdp=True,
+    enable_fsdp=False,
     local_attn_size=-1,
     sink_size=0,
     timestep_indices=(0, 179, 358, 679),
@@ -76,14 +76,30 @@
 )
 
 
+def _resolve_stage_devices(total_gpu_count: int) -> tuple[list[int], int, int]:
+    """Return DiT, VAE encode, and VAE decode devices for available GPUs."""
+    if total_gpu_count < 1:
+        raise ValueError(f"parallelism must be positive, got {total_gpu_count}")
+    if total_gpu_count in {2, 4}:
+        return list(range(total_gpu_count)), 0, 1
+    if total_gpu_count == 5:
+        return [0, 1, 2, 3], 4, 4
+    if total_gpu_count == 6:
+        return [0, 1, 2, 3, 4], 5, 5
+    return (
+        list(range(total_gpu_count)),
+        int(PPL_CONFIG["vae_encode_device_id"]),
+        int(PPL_CONFIG["vae_decode_device_id"]),
+    )
+
+
 def get_pipeline(
     parallelism: int = PPL_CONFIG["parallelism"],
     model_root: str | None = None,
     fast_model_root: str | None = None,
 ) -> LingBotWorldFastPipeline:
     """Load LingBot-World-Fast for offline chunked generation."""
-    if parallelism < 1:
-        raise ValueError(f"parallelism must be positive, got {parallelism}")
+    dit_device_ids, vae_encode_device, vae_decode_device = _resolve_stage_devices(parallelism)
     model_root_path = Path(model_root).expanduser() if model_root else None
     fast_model_root_path = Path(fast_model_root).expanduser() if fast_model_root else None
     vae_path = str(model_root_path / "Wan2.1_VAE.pth") if model_root_path else PPL_CONFIG["vae_path"]
@@ -125,29 +141,33 @@ def get_pipeline(
         LingBotWorldFastPipelineConfig(
             vae_encode_config=ModelRuntimeConfig(
                 device_type="cuda",
-                device_id=int(PPL_CONFIG["vae_encode_device_id"]),
+                device_id=vae_encode_device,
                 torch_dtype=PPL_CONFIG["vae_torch_dtype"],
-                parallel_config=ParallelConfig(device_ids=[int(PPL_CONFIG["vae_encode_device_id"])]),
+                parallel_config=ParallelConfig(device_ids=[vae_encode_device]),
             ),
             vae_decode_config=ModelRuntimeConfig(
                 device_type="cuda",
-                device_id=int(PPL_CONFIG["vae_decode_device_id"]),
+                device_id=vae_decode_device,
                 torch_dtype=PPL_CONFIG["vae_torch_dtype"],
-                parallel_config=ParallelConfig(device_ids=[int(PPL_CONFIG["vae_decode_device_id"])]),
+                parallel_config=ParallelConfig(device_ids=[vae_decode_device]),
+            ),
+            text_encoding_config=ModelRuntimeConfig(device_type="cuda", device_id=dit_device_ids[0], torch_dtype=dtype),
+            dit_config=ModelRuntimeConfig(
+                device_type="cuda",
+                device_id=dit_device_ids[0],
+                torch_dtype=dtype,
+                attention_config=AttentionConfig.dense_attention(PPL_CONFIG["attn_impl"]),
+                parallel_config=ParallelConfig(
+                    device_ids=dit_device_ids if len(dit_device_ids) > 1 else None,
+                    sp_ulysses_degree=len(dit_device_ids),
+                    enable_fsdp=PPL_CONFIG["enable_fsdp"] and len(dit_device_ids) > 1,
+                ),
             ),
-            text_encoding_config=ModelRuntimeConfig(device_type="cuda", device_id=0, torch_dtype=dtype),
-            dit_torch_dtype=dtype,
             control_type=PPL_CONFIG["control_mode"],
             max_area=RESOLUTION_AREAS[PPL_CONFIG["resolution"]],
             local_attn_size=PPL_CONFIG["local_attn_size"],
             sink_size=PPL_CONFIG["sink_size"],
             timestep_indices=PPL_CONFIG["timestep_indices"],
-            attention_config=AttentionConfig.dense_attention(PPL_CONFIG["attn_impl"]),
-            parallel_config=ParallelConfig(
-                device_ids=list(range(parallelism)) if parallelism > 1 else None,
-                sp_ulysses_degree=parallelism,
-                enable_fsdp=PPL_CONFIG["enable_fsdp"] and parallelism > 1,
-            ),
         ),
     )
     return pipeline
@@ -216,7 +236,7 @@ def run(
     "--gpu_num",
     default=PPL_CONFIG["parallelism"],
     type=int,
-    help="Number of GPUs used for Ulysses sequence parallelism",
+    help="Total GPUs; selects the LingBot VAE and DiT placement strategy",
 )
 @click.option("--image_path", default=DEFAULT_IMAGE_PATH, type=click.Path(exists=True))
 @click.option("--action_path", default=DEFAULT_ACTION_PATH, type=click.Path(exists=True, file_okay=False))

examples/lingbot/lingbot_world_v2_image_to_video_h100.py

@@ -22,7 +22,7 @@
 import torch
 from PIL import Image
 
-from telefuser.core.config import AttentionConfig, AttnImplType, ModelRuntimeConfig, ParallelConfig
+from telefuser.core.config import AttentionConfig, AttnImplType, CompileConfig, ModelRuntimeConfig, ParallelConfig
 from telefuser.core.module_manager import ModuleManager
 from telefuser.models.lingbot_world_fast_dit import LingBotWorldFastDiT
 from telefuser.models.wan_video_text_encoder import WanTextEncoder
@@ -74,8 +74,9 @@
     seed=42,
     target_fps=16,
     max_duration_seconds=120.0,
-    attn_impl=AttnImplType.TORCH_CUDNN,
-    enable_fsdp=True,
+    attn_impl=AttnImplType.SAGE_ATTN_2_8_8_SM90,
+    compile_config=CompileConfig(enabled=True),
+    enable_fsdp=False,
     local_attn_size=18,
     sink_size=6,
     timestep_indices=(0, 250, 500, 750),
@@ -90,14 +91,30 @@
 )
 
 
+def _resolve_stage_devices(total_gpu_count: int) -> tuple[list[int], int, int]:
+    """Return DiT, VAE encode, and VAE decode devices for available GPUs."""
+    if total_gpu_count < 1:
+        raise ValueError(f"parallelism must be positive, got {total_gpu_count}")
+    if total_gpu_count in {2, 4}:
+        return list(range(total_gpu_count)), 0, 1
+    if total_gpu_count == 5:
+        return [0, 1, 2, 3], 4, 4
+    if total_gpu_count == 6:
+        return [0, 1, 2, 3, 4], 5, 5
+    return (
+        list(range(total_gpu_count)),
+        int(PPL_CONFIG["vae_encode_device_id"]),
+        int(PPL_CONFIG["vae_decode_device_id"]),
+    )
+
+
 def get_pipeline(
     parallelism: int = PPL_CONFIG["parallelism"],
     model_root: str | None = None,
     v2_model_root: str | None = None,
 ) -> LingBotWorldV2Pipeline:
     """Load LingBot-World v2 for offline chunked generation."""
-    if parallelism < 1:
-        raise ValueError(f"parallelism must be positive, got {parallelism}")
+    dit_device_ids, vae_encode_device, vae_decode_device = _resolve_stage_devices(parallelism)
 
     model_root_path = Path(model_root).expanduser() if model_root else None
     v2_model_root_path = Path(v2_model_root).expanduser() if v2_model_root else None
@@ -111,9 +128,6 @@ def get_pipeline(
         else PPL_CONFIG["dit_path_list"]
     )
 
-    vae_encode_device = int(PPL_CONFIG["vae_encode_device_id"])
-    vae_decode_device = int(PPL_CONFIG["vae_decode_device_id"])
-    dit_device_ids = list(range(parallelism))
     dtype = PPL_CONFIG["torch_dtype"]
     module_manager = ModuleManager(device="cpu")
     module_manager.load_model(
@@ -155,17 +169,17 @@ def get_pipeline(
                 parallel_config=ParallelConfig(device_ids=[vae_decode_device]),
             ),
             text_encoding_config=ModelRuntimeConfig(device_type="cuda", device_id=dit_device_ids[0], torch_dtype=dtype),
-            dit_torch_dtype=dtype,
-            control_type=PPL_CONFIG["control_mode"],
-            max_area=RESOLUTION_AREAS[PPL_CONFIG["resolution"]],
-            local_attn_size=PPL_CONFIG["local_attn_size"],
-            sink_size=PPL_CONFIG["sink_size"],
-            timestep_indices=PPL_CONFIG["timestep_indices"],
-            attention_config=AttentionConfig.dense_attention(PPL_CONFIG["attn_impl"]),
-            parallel_config=ParallelConfig(
-                device_ids=dit_device_ids if len(dit_device_ids) > 1 else None,
-                sp_ulysses_degree=len(dit_device_ids),
-                enable_fsdp=PPL_CONFIG["enable_fsdp"] and len(dit_device_ids) > 1,
+            dit_config=ModelRuntimeConfig(
+                device_type="cuda",
+                device_id=dit_device_ids[0],
+                torch_dtype=dtype,
+                attention_config=AttentionConfig.dense_attention(PPL_CONFIG["attn_impl"]),
+                parallel_config=ParallelConfig(
+                    device_ids=dit_device_ids if len(dit_device_ids) > 1 else None,
+                    sp_ulysses_degree=len(dit_device_ids),
+                    enable_fsdp=PPL_CONFIG["enable_fsdp"] and len(dit_device_ids) > 1,
+                ),
+                compile_config=PPL_CONFIG["compile_config"],
             ),
         ),
     )
@@ -241,7 +255,7 @@ def run(
     "--gpu_num",
     default=PPL_CONFIG["parallelism"],
     type=int,
-    help="Number of GPUs used for Ulysses sequence parallelism",
+    help="Total GPUs; selects the LingBot VAE and DiT placement strategy",
 )
 @click.option("--image_path", default=DEFAULT_IMAGE_PATH, type=click.Path(exists=True))
 @click.option("--action_path", default=DEFAULT_ACTION_PATH, type=click.Path(exists=True, file_okay=False))

telefuser/pipelines/lingbot_world_fast/denoising.py

@@ -62,6 +62,9 @@ def __init__(
         self.dit.set_attention_config(model_runtime_config.attention_config)
         self.model_names = ["dit"]
         self._cache_registry: dict[int, _DenoisingCacheState] = {}
+        if model_runtime_config.parallel_config.world_size == 1 and model_runtime_config.compile_config.enabled:
+            logger.info(f"Enabling torch.compile for {self.name}")
+            self.dit = torch.compile(self.dit, **model_runtime_config.compile_config.get_compile_kwargs())
 
     def parallel_models(self) -> None:
         """Configure Ulysses SP and optional FSDP inside a ParallelWorker."""
@@ -81,6 +84,9 @@ def parallel_models(self) -> None:
                 buffer_dtype=self.torch_dtype,
             )
             self.onload_models_flag = True
+        if self.model_runtime_config.compile_config.enabled:
+            logger.info(f"Enabling torch.compile for {self.name}")
+            self.dit = torch.compile(self.dit, **self.model_runtime_config.compile_config.get_compile_kwargs())
 
     def _init_self_kv_cache(
         self,

telefuser/pipelines/lingbot_world_fast/pipeline.py

@@ -14,7 +14,7 @@
 from PIL import Image
 
 from telefuser.core.base_pipeline import BasePipeline
-from telefuser.core.config import AttentionConfig, ModelRuntimeConfig, ParallelConfig
+from telefuser.core.config import ModelRuntimeConfig, ParallelConfig
 from telefuser.core.module_manager import ModuleManager
 from telefuser.models.lingbot_world_fast_dit import LingBotWorldFastDiT
 from telefuser.models.t5_tokenizer import HuggingfaceTokenizer
@@ -53,16 +53,14 @@ class LingBotWorldFastPipelineConfig:
     vae_encode_config: ModelRuntimeConfig = field(default_factory=_default_vae_stage_runtime_config)
     vae_decode_config: ModelRuntimeConfig = field(default_factory=_default_vae_stage_runtime_config)
     text_encoding_config: ModelRuntimeConfig = field(default_factory=ModelRuntimeConfig)
-    dit_torch_dtype: torch.dtype = torch.bfloat16
+    dit_config: ModelRuntimeConfig = field(default_factory=ModelRuntimeConfig)
     control_type: str = "cam"
     orig_height: int = 480
     orig_width: int = 832
     max_area: int = 480 * 832
     local_attn_size: int = -1
     sink_size: int = 0
     timestep_indices: tuple[int, ...] = (0, 179, 358, 679)
-    parallel_config: ParallelConfig = field(default_factory=ParallelConfig)
-    attention_config: AttentionConfig = field(default_factory=AttentionConfig)
 
 
 class LingBotWorldFastPipeline(BasePipeline):
@@ -137,27 +135,22 @@ def init(self, module_manager: ModuleManager, config: LingBotWorldFastPipelineCo
         self.vae_encode_worker = ParallelWorker(vae_encode_stage)
         self.vae_decode_worker = ParallelWorker(vae_decode_stage)
 
-        dit_device = "cpu" if config.parallel_config.world_size > 1 else self.device
+        dit_runtime_config = config.dit_config
+        dit_device = "cpu" if dit_runtime_config.parallel_config.world_size > 1 else self.device
         self.dit = module_manager.fetch_module("lingbot_world_fast_dit")
         if self.dit is None:
             raise ValueError("LingBot requires a loaded lingbot_world_fast_dit module")
-        self.dit = self.dit.to(dit_device, dtype=config.dit_torch_dtype).eval().requires_grad_(False)
+        self.dit = self.dit.to(dit_device, dtype=dit_runtime_config.torch_dtype).eval().requires_grad_(False)
         if self.dit.control_type != config.control_type:
             raise ValueError(
                 f"DiT checkpoint control type is {self.dit.control_type!r}, not requested {config.control_type!r}"
             )
         self.dit.set_causal_attention_window(config.local_attn_size, config.sink_size)
 
-        pipeline_device = torch.device(self.device)
-        dit_runtime_config = ModelRuntimeConfig(
-            device_type=pipeline_device.type,
-            device_id=pipeline_device.index or 0,
-            torch_dtype=config.dit_torch_dtype,
-            attention_config=config.attention_config,
-            parallel_config=config.parallel_config,
-        )
         denoise_stage = LingBotWorldFastDenoisingStage("lingbot_world_fast_denoise", module_manager, dit_runtime_config)
-        self.denoise_stage = ParallelWorker(denoise_stage) if config.parallel_config.world_size > 1 else denoise_stage
+        self.denoise_stage = (
+            ParallelWorker(denoise_stage) if dit_runtime_config.parallel_config.world_size > 1 else denoise_stage
+        )
 
     @staticmethod
     def _validate_vae_stage_runtime_config(runtime_config: ModelRuntimeConfig) -> None:

tests/unit/pipelines/lingbot_world_fast/test_module_loading.py

@@ -2,6 +2,7 @@
 
 import torch
 
+from telefuser.core.config import ModelRuntimeConfig
 from telefuser.pipelines.lingbot_world_fast.pipeline import LingBotWorldFastPipeline, LingBotWorldFastPipelineConfig
 
 
@@ -38,9 +39,7 @@ def fetch_module(name: str, require_model_path: bool = False):
         pipeline = LingBotWorldFastPipeline(device="cpu", torch_dtype=torch.float32)
         pipeline.init(
             module_manager,
-            LingBotWorldFastPipelineConfig(
-                dit_torch_dtype=torch.float32,
-            ),
+            LingBotWorldFastPipelineConfig(dit_config=ModelRuntimeConfig(torch_dtype=torch.float32)),
         )
 
     assert module_manager.load_model.call_count == 0

tests/unit/pipelines/lingbot_world_fast/test_parallelism.py

@@ -4,7 +4,7 @@
 import pytest
 import torch
 
-from telefuser.core.config import AttentionConfig, ModelRuntimeConfig, ParallelConfig
+from telefuser.core.config import AttentionConfig, CompileConfig, ModelRuntimeConfig, ParallelConfig
 from telefuser.pipelines.lingbot_world_fast.denoising import LingBotWorldFastDenoisingStage
 
 
@@ -96,3 +96,36 @@ def test_denoising_stage_rejects_uneven_ulysses_head_partition() -> None:
         pytest.raises(ValueError, match="divisible"),
     ):
         stage._init_self_kv_cache(batch_size=1, kv_size=1)
+
+
+def test_denoising_stage_parallel_models_compiles_after_fsdp() -> None:
+    dit = MagicMock()
+    dit.get_fsdp_module_names.return_value = ["blocks"]
+    parallel_config = ParallelConfig(
+        device_ids=[0, 1, 2, 3],
+        sp_ulysses_degree=4,
+        enable_fsdp=True,
+    )
+    compile_config = CompileConfig(enabled=True, backend="eager")
+    runtime_config = ModelRuntimeConfig(
+        device_type="cuda",
+        device_id=0,
+        attention_config=AttentionConfig(),
+        parallel_config=parallel_config,
+        compile_config=compile_config,
+    )
+    module_manager = MagicMock()
+    module_manager.fetch_module.return_value = dit
+    stage = LingBotWorldFastDenoisingStage("denoise", module_manager, runtime_config)
+    fsdp_model = MagicMock()
+    compiled_model = MagicMock()
+
+    with (
+        patch("telefuser.pipelines.lingbot_world_fast.denoising.create_device_mesh_from_config"),
+        patch("telefuser.pipelines.lingbot_world_fast.denoising.shard_model", return_value=fsdp_model),
+        patch("telefuser.pipelines.lingbot_world_fast.denoising.torch.compile", return_value=compiled_model) as compile,
+    ):
+        stage.parallel_models()
+
+    compile.assert_called_once_with(fsdp_model, **compile_config.get_compile_kwargs())
+    assert stage.dit is compiled_model