[diffusion] fix: stabilize ltx-2.3 two-stage cold requests (#35997)

python/sglang/multimodal_gen/runtime/warmup_request_builder.py

@@ -39,6 +39,7 @@ SERVER_WARMUP_IMAGE_MAX_AREA = 768 * 768
 SERVER_WARMUP_DIFFUSERS_IMAGE_MAX_AREA = 512 * 512
 SERVER_WARMUP_VIDEO_MAX_AREA = 832 * 480
 SERVER_WARMUP_MAX_VIDEO_FRAMES = 17
+SERVER_WARMUP_LTX2_TWO_STAGE_MAX_VIDEO_FRAMES = 25
 SERVER_WARMUP_IMAGE_STEPS = 2
 SERVER_WARMUP_VIDEO_STEPS = 2
 
@@ -245,7 +246,15 @@ def _resolve_warmup_num_frames(
     ):
         warmup_num_frames = num_frames
     else:
-        warmup_num_frames = min(num_frames, SERVER_WARMUP_MAX_VIDEO_FRAMES)
+        # Multi-GPU LTX two-stage aligns a one-second request to 25 frames;
+        # cover its latent shape during warmup
+        frame_budget = (
+            SERVER_WARMUP_LTX2_TWO_STAGE_MAX_VIDEO_FRAMES
+            if is_ltx2_two_stage_pipeline_name(server_args.pipeline_class_name)
+            and server_args.num_gpus > 1
+            else SERVER_WARMUP_MAX_VIDEO_FRAMES
+        )
+        warmup_num_frames = min(num_frames, frame_budget)
 
     return server_args.pipeline_config.adjust_num_frames(warmup_num_frames)
 

python/sglang/multimodal_gen/test/server/perf_baselines/h100.json

@@ -2652,7 +2652,7 @@
                 "LTX2AVLatentPreparationStage": 0.13,
                 "LTX2ImageEncodingStage": 27.58,
                 "LTX2AVDenoisingStage": 7506.38,
-                "LTX2UpsampleStage": 2.31,
+                "LTX2UpsampleStage": 172.53,
                 "LTX2RefinementStage": 670.53,
                 "LTX2AVDecodingStage": 250.93,
                 "per_frame_generation": null

python/sglang/multimodal_gen/test/unit/test_cfg_parallel_warmup.py

@@ -502,6 +502,7 @@ class TestWarmupReqCfgParallel(unittest.TestCase):
         server_args = SimpleNamespace(
             pipeline_config=pipeline_config,
             enable_breakable_cuda_graph=False,
+            pipeline_class_name=None,
         )
 
         num_frames = _resolve_warmup_num_frames(
@@ -563,7 +564,8 @@ class TestWarmupReqCfgParallel(unittest.TestCase):
 
         server_args.pipeline_config.task_type = ModelTaskType.T2V
         server_args.pipeline_config.vae_scale_factor = 32
-        server_args.pipeline_config.adjust_num_frames.side_effect = lambda value: value
+        server_args.pipeline_config.adjust_num_frames.return_value = 25
+        server_args.num_gpus = 2
 
         sampling_defaults = SamplingParams(
             width=1920,
@@ -584,6 +586,23 @@ class TestWarmupReqCfgParallel(unittest.TestCase):
         self.assertEqual((reqs[0].width, reqs[0].height), (832, 448))
         self.assertEqual(reqs[0].width % 64, 0)
         self.assertEqual(reqs[0].height % 64, 0)
+        self.assertEqual(reqs[0].num_frames, 25)
+        server_args.pipeline_config.adjust_num_frames.assert_called_once_with(25)
+
+    def test_ltx2_two_stage_single_gpu_keeps_generic_frame_cap(self):
+        server_args = MagicMock()
+        server_args.pipeline_class_name = "LTX2TwoStagePipeline"
+        server_args.num_gpus = 1
+        server_args.pipeline_config.task_type = ModelTaskType.T2V
+        server_args.pipeline_config.adjust_num_frames.side_effect = lambda value: value
+
+        num_frames = _resolve_warmup_num_frames(
+            server_args,
+            SamplingParams(num_frames=121),
+            server_based_warmup=True,
+        )
+
+        self.assertEqual(num_frames, 17)
 
     def test_server_based_warmup_uses_representative_image_fallback(self):
         server_args = MagicMock()