[Bug][CI] Pin the LTX FP8 quality gate to its established sigma schedule (#5302)

Signed-off-by: mglyn <1203789601@qq.com>

docs/user_guide/examples/online_serving/text_to_video.md

@@ -296,14 +296,14 @@ done
 
 ```bash
 vllm serve Lightricks/LTX-2 --omni --port 8098 \
-    --enforce-eager --flow-shift 1.0 --boundary-ratio 1.0
+    --enforce-eager --boundary-ratio 1.0
 ```
 
 For multi-GPU memory reduction, you can enable HSDP:
 
 ```bash
 vllm serve Lightricks/LTX-2 --omni --port 8098 \
-    --enforce-eager --flow-shift 1.0 --boundary-ratio 1.0 \
+    --enforce-eager --boundary-ratio 1.0 \
     --use-hsdp --hsdp-shard-size 2
 ```
 

examples/online_serving/text_to_video/README.md

@@ -254,7 +254,7 @@ done
 
 ```bash
 vllm serve Lightricks/LTX-2 --omni --port 8098 \
-    --enforce-eager --flow-shift 1.0 --boundary-ratio 1.0
+    --enforce-eager --boundary-ratio 1.0
 ```
 
 #### Start with Optimization Presets

tests/diffusion/models/ltx2/test_ltx2_pipeline.py

@@ -17,7 +17,11 @@ from vllm_omni.diffusion.models.ltx2.ltx2_components import (
     LTX23_COMPONENT_PROFILE,
 )
 from vllm_omni.diffusion.models.ltx2.ltx2_conditioning import LTXI2VConditioningMixin
-from vllm_omni.diffusion.models.ltx2.ltx2_denoise import LTXDenoiseExecutor, LTXPhaseResult
+from vllm_omni.diffusion.models.ltx2.ltx2_denoise import (
+    LTXDenoiseExecutor,
+    LTXPhaseResult,
+    prepare_scheduler_stage,
+)
 from vllm_omni.diffusion.models.ltx2.ltx2_guidance import (
     LTX_LEGACY_VELOCITY_GUIDANCE,
     LTX_OFFICIAL_X0_GUIDANCE,
@@ -324,6 +328,38 @@ def test_ltx2_two_stage_reuses_prompt_context_between_phases():
     torch.testing.assert_close(output.output[1], torch.tensor([4.0]))
 
 
+def test_ltx_custom_sigmas_bypass_scheduler_shifting():
+    from diffusers import FlowMatchEulerDiscreteScheduler
+
+    scheduler = FlowMatchEulerDiscreteScheduler(
+        use_dynamic_shifting=True,
+        base_shift=0.95,
+        max_shift=2.05,
+        shift_terminal=0.1,
+    )
+    pipeline = SimpleNamespace(
+        scheduler=scheduler,
+        _make_video_audio_scheduler=lambda *args, **kwargs: object(),
+    )
+    sigmas = [1.0, 0.75, 0.25, 0.0]
+
+    audio_scheduler, _, timesteps = prepare_scheduler_stage(
+        pipeline,
+        SimpleNamespace(num_inference_steps=3),
+        device=torch.device("cpu"),
+        sigmas=sigmas,
+        timesteps=None,
+        latent_num_frames=1,
+        latent_height=1,
+        latent_width=1,
+    )
+
+    expected = torch.tensor(sigmas, dtype=torch.float32)
+    torch.testing.assert_close(pipeline.scheduler.sigmas, expected)
+    torch.testing.assert_close(audio_scheduler.sigmas, expected)
+    torch.testing.assert_close(timesteps, expected[:-1] * scheduler.config.num_train_timesteps)
+
+
 class TestLTXRequestParsing:
     @pytest.mark.parametrize(
         ("t2v_cls", "i2v_cls"),
@@ -513,6 +549,7 @@ class TestLTX23ForwardStages:
         from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 
         pipe = _make_ltx_request_pipe(LTX23Pipeline)
+        request_sigmas = [1.0, 0.4]
         req = DiffusionRequestBatch(
             [
                 OmniDiffusionRequest(
@@ -525,12 +562,13 @@ class TestLTX23ForwardStages:
                         num_inference_steps=3,
                         guidance_scale=4.5,
                         output_type="latent",
+                        sigmas=request_sigmas,
                     ),
                     request_id="ltx23-t2v-forward-stage-delegation",
                 )
             ]
         )
-        sigmas = [1.0, 0.5]
+        fallback_sigmas = [1.0, 0.5]
         timesteps = [1000, 500]
         attention_kwargs = {"scale": 1.0}
         seen = {}
@@ -545,7 +583,7 @@ class TestLTX23ForwardStages:
 
         output = pipe.forward(
             req,
-            sigmas=sigmas,
+            sigmas=fallback_sigmas,
             timesteps=timesteps,
             noise_scale=0.25,
             attention_kwargs=attention_kwargs,
@@ -561,7 +599,7 @@ class TestLTX23ForwardStages:
         assert seen["request_inputs"].output_type == "latent"
         assert seen["kwargs"] == {
             "noise_scale": 0.25,
-            "sigmas": sigmas,
+            "sigmas": request_sigmas,
             "timesteps": timesteps,
             "attention_kwargs": attention_kwargs,
         }

tests/diffusion/quantization/test_quantization_quality.py

@@ -82,6 +82,7 @@ class QualityTestConfig:
     gpu: str = "H100"  # minimum GPU requirement
     negative_prompt: str = ""
     guidance_scale: float | None = None
+    sigmas: list[float] | None = None
     enable_cpu_offload: bool = False
 
     def baseline_ref(self) -> str:
@@ -215,6 +216,19 @@ QUALITY_CONFIGS = [
         width=256,
         num_frames=25,
         num_inference_steps=8,
+        # Preserve the final scheduler trajectory used when this FP8 quality
+        # gate was established. Explicit LTX sigmas bypass dynamic shifting.
+        sigmas=[
+            1.0,
+            0.92185378074646,
+            0.8327768445014954,
+            0.7303033471107483,
+            0.6111654043197632,
+            0.4709382653236389,
+            0.30347853899002075,
+            0.10000002384185791,
+            0.0,
+        ],
     ),
 ]
 
@@ -322,6 +336,7 @@ def _generate_video(omni, config: QualityTestConfig):
             num_inference_steps=config.num_inference_steps,
             num_frames=config.num_frames,
             guidance_scale=config.guidance_scale,
+            sigmas=config.sigmas,
         ),
     )
 
@@ -446,6 +461,33 @@ def test_benchmark_generate_image_unwraps_nested_omni_request_output(monkeypatch
     assert peak_mem == 0.0
 
 
+def test_generate_video_forwards_sigmas(monkeypatch):
+    from vllm_omni.platforms import current_omni_platform
+
+    monkeypatch.setattr(current_omni_platform, "device_type", "cpu", raising=False)
+    monkeypatch.setattr(torch.accelerator, "reset_peak_memory_stats", lambda: None, raising=False)
+    monkeypatch.setattr(torch.accelerator, "max_memory_allocated", lambda: 0, raising=False)
+    captured = SimpleNamespace(sampling=None)
+
+    class DummyOmni:
+        def generate(self, _prompt, sampling):
+            captured.sampling = sampling
+            return [SimpleNamespace(images=[np.zeros((1, 2, 2, 3), dtype=np.float32)])]
+
+    config = QualityTestConfig(
+        id="ltx-sigmas",
+        model="unused",
+        quantization="fp8",
+        task="t2v",
+        prompt="test",
+        max_lpips=0.1,
+        sigmas=[1.0, 0.5],
+    )
+    _generate_video(DummyOmni(), config)
+
+    assert captured.sampling.sigmas == [1.0, 0.5]
+
+
 _marks = hardware_marks(res={"cuda": "H100"})
 _OUTPUT_DIR = Path(os.environ["VLLM_OMNI_QUALITY_OUTPUT_DIR"]) if "VLLM_OMNI_QUALITY_OUTPUT_DIR" in os.environ else None
 

vllm_omni/diffusion/models/ltx2/ltx2_denoise.py

@@ -177,6 +177,17 @@ class I2VVideoAudioScheduler:
         return ((video_out, audio_out),)
 
 
+def _set_scheduler_sigmas(scheduler: Any, sigmas: torch.Tensor) -> torch.Tensor:
+    sigmas = sigmas.to(torch.float32)
+    timesteps = sigmas[:-1] * scheduler.config.get("num_train_timesteps", 1000)
+    scheduler.sigmas = sigmas
+    scheduler.timesteps = timesteps
+    scheduler.num_inference_steps = len(timesteps)
+    scheduler._step_index = None
+    scheduler._begin_index = None
+    return timesteps
+
+
 def prepare_scheduler_stage(
     pipeline: Any,
     request_inputs: LTXRequestInputs,
@@ -188,18 +199,9 @@ def prepare_scheduler_stage(
     latent_height: int,
     latent_width: int,
 ) -> tuple[Any, Any, torch.Tensor]:
-    sigmas = (
-        np.linspace(1.0, 1 / request_inputs.num_inference_steps, request_inputs.num_inference_steps)
-        if sigmas is None
-        else sigmas
-    )
-    mu = calculate_shift(
-        pipeline.scheduler.config.get("max_image_seq_len", 4096),
-        pipeline.scheduler.config.get("base_image_seq_len", 1024),
-        pipeline.scheduler.config.get("max_image_seq_len", 4096),
-        pipeline.scheduler.config.get("base_shift", 0.95),
-        pipeline.scheduler.config.get("max_shift", 2.05),
-    )
+    if sigmas is not None and timesteps is not None:
+        raise ValueError("Only one of `sigmas` or `timesteps` may be provided.")
+
     audio_scheduler = copy.deepcopy(pipeline.scheduler)
     video_audio_scheduler = pipeline._make_video_audio_scheduler(
         audio_scheduler,
@@ -207,6 +209,27 @@ def prepare_scheduler_stage(
         latent_height,
         latent_width,
     )
+    if sigmas is not None:
+        scheduler_sigmas = torch.as_tensor(sigmas, dtype=torch.float32, device=device)
+        if scheduler_sigmas.ndim != 1 or scheduler_sigmas.numel() < 2:
+            raise ValueError("An LTX custom sigma schedule must contain at least two boundary values.")
+        if scheduler_sigmas[-1] != 0:
+            scheduler_sigmas = torch.cat([scheduler_sigmas, scheduler_sigmas.new_zeros(1)])
+        timesteps_tensor = _set_scheduler_sigmas(pipeline.scheduler, scheduler_sigmas)
+        _set_scheduler_sigmas(audio_scheduler, scheduler_sigmas.clone())
+        return audio_scheduler, video_audio_scheduler, timesteps_tensor
+
+    sigmas = np.linspace(1.0, 1 / request_inputs.num_inference_steps, request_inputs.num_inference_steps)
+    # Official LTX2 one-stage omits the latent, so mu stays at max_shift across resolutions.
+    # Mu varies with latent token count only when the scheduler receives a latent, as in official HQ two-stage;
+    # explicit sigma schedules bypass mu entirely.
+    mu = calculate_shift(
+        pipeline.scheduler.config.get("max_image_seq_len", 4096),
+        pipeline.scheduler.config.get("base_image_seq_len", 1024),
+        pipeline.scheduler.config.get("max_image_seq_len", 4096),
+        pipeline.scheduler.config.get("base_shift", 0.95),
+        pipeline.scheduler.config.get("max_shift", 2.05),
+    )
     retrieve_timesteps(
         audio_scheduler,
         request_inputs.num_inference_steps,

vllm_omni/diffusion/models/ltx2/ltx2_request.py

@@ -71,6 +71,17 @@ class LTXRequestMixin:
     supports_request_batch = False
     supports_guidance_rescale = False
 
+    @staticmethod
+    def _resolve_request_sigmas(
+        req: DiffusionRequestBatch,
+        fallback: list[float] | None,
+    ) -> list[float] | None:
+        request_sigmas = [sampling.sigmas for sampling in req.sampling_params_list]
+        first = request_sigmas[0] if request_sigmas else None
+        if any(item != first for item in request_sigmas[1:]):
+            raise ValueError("Batched LTX requests must use identical custom sigmas.")
+        return first if first is not None else fallback
+
     def check_inputs(
         self,
         prompt: str | list[str] | None,

vllm_omni/diffusion/models/ltx2/pipeline_ltx2.py

@@ -96,6 +96,7 @@ class LTXOneStagePipeline(LTXPipelineRuntime):
         image: Any | None = None,
     ) -> DiffusionOutput | list[DiffusionOutput]:
         del return_dict
+        sigmas = self._resolve_request_sigmas(req, sigmas)
         request_inputs = self._resolve_request_inputs(
             req,
             prompt=prompt,