[Bugfix][BAGEL] Restore original denoise step count (#7049)

Signed-off-by: Trigger <Meng.Bo.Wang@outlook.com>
Signed-off-by: Gao Han <hgaoaf@connect.ust.hk>
Co-authored-by: Gao Han <hgaoaf@connect.ust.hk>

docs/design/feature/diffusion_continuous_batching.md

@@ -155,12 +155,12 @@ imply continuous-batching support: Qwen-Image accepts batched step states.
 BAGEL accepts batched step states for image generation with its default,
 think-mode, and single-stage deploy configs. Only the diffusion stage uses the
 step protocol in two-stage deployments; explicit single-stage text output
-retains complete-request execution. BAGEL schedules exactly
-`num_inference_steps` denoising updates, including a one-step request. When a
-finished stepwise request is followed by an explicit text fallback wave, the
-runner retires its step state and paged diffusion-KV rows before the full
-forward dispatch. BAGEL's step path does not currently support sequence
-parallelism or a diffusion cache backend.
+retains complete-request execution. BAGEL requires at least two inference steps
+because it schedules `num_inference_steps - 1` denoising updates, and its step
+path does not currently support sequence parallelism or a diffusion cache
+backend. When a finished stepwise request is followed by an explicit text
+fallback wave, the runner retires its step state and paged diffusion-KV rows
+before the full forward dispatch.
 HunyuanImage3 accepts batched step states only when its resolved self-attention
 backend is `TORCH_SDPA`; otherwise it rejects groups larger than one request.
 Configure `DIFFUSION_ATTENTION_BACKEND=TORCH_SDPA` or

docs/user_guide/diffusion/execution_modes.md

@@ -99,10 +99,10 @@ step-wise continuous batching for image generation with `bagel.yaml`,
 `bagel_think.yaml`, and `bagel_single_stage.yaml`. In the two-stage topologies,
 only the diffusion stage uses step execution; the autoregressive Thinker is
 unchanged. Explicit text-output requests in the single-stage topology retain
-the complete-request path. BAGEL schedules exactly `num_inference_steps`
-denoising updates, so one-step image requests are supported. BAGEL step
-execution does not currently support sequence parallelism or a diffusion cache
-backend. HunyuanImage3 also supports step execution, but only
+the complete-request path. Use at least two inference steps for BAGEL image
+requests because BAGEL schedules `num_inference_steps - 1` denoising updates.
+BAGEL step execution does not currently support sequence parallelism or a
+diffusion cache backend. HunyuanImage3 also supports step execution, but only
 when its resolved self-attention backend is `TORCH_SDPA`;
 set `DIFFUSION_ATTENTION_BACKEND=TORCH_SDPA` or configure
 `diffusion_attention_config.default.backend=TORCH_SDPA` before using

docs/user_guide/diffusion_features.md

@@ -152,7 +152,7 @@ The following tables show which models support each feature:
 > 3. Cosmos3 T2I uses `Cosmos3OmniDiffusersPipeline` with `modalities=["image"]`. Model-level CPU offload swaps the nested UND reasoner and GEN generator pathways; layerwise offload remains available for blockwise GEN/UND offload.
 > 4. Krea 2 currently supports single-GPU inference plus LoRA, Cache-DiT, HSDP, CPU/layerwise offload, and VAE-patch-parallel (decode). TP/SP/CFG-Parallel are not yet wired. The few-step distilled (Turbo) checkpoint uses `is_distilled=true` (fixed timestep shift `mu=1.15`); generate at 2048x2048 by default with `num_inference_steps≈8` and `guidance_scale=0`. The Raw checkpoint uses 1024x1024, `num_inference_steps=28`, and `guidance_scale=4.5`.
 > 5. HunyuanImage3 supports step execution. Multi-request step execution requires `TORCH_SDPA`; see [Diffusion Execution Modes](diffusion/execution_modes.md#step-execution).
-> 6. BAGEL step execution supports image generation with `bagel.yaml`, `bagel_think.yaml`, and `bagel_single_stage.yaml`; two-stage Thinker execution and explicit single-stage text output remain on their existing complete-request paths. BAGEL executes exactly `num_inference_steps` denoising updates, including one-step requests. BAGEL step execution cannot currently be combined with sequence parallelism or a diffusion cache backend; see [Diffusion Execution Modes](diffusion/execution_modes.md#step-execution).
+> 6. BAGEL step execution supports image generation with `bagel.yaml`, `bagel_think.yaml`, and `bagel_single_stage.yaml`; two-stage Thinker execution and explicit single-stage text output remain on their existing complete-request paths. Image requests require `num_inference_steps >= 2`. BAGEL step execution cannot currently be combined with sequence parallelism or a diffusion cache backend; see [Diffusion Execution Modes](diffusion/execution_modes.md#step-execution).
 
 ### VideoGen
 

docs/user_guide/examples/online_serving/bagel.md

@@ -98,10 +98,10 @@ vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091 \
     --stage-overrides '{"0":{"max_num_seqs":4}}'
 ```
 
-BAGEL image requests use exactly `num_inference_steps` denoising updates,
-including one-step requests. The scheduler batches only requests with
-compatible effective image sizes and BAGEL CFG/renormalization settings.
-Sequence parallelism and diffusion cache backends such as TeaCache or
+BAGEL image requests require `num_inference_steps >= 2`: the original BAGEL
+schedule uses `num_inference_steps - 1` denoising updates. The scheduler batches
+only requests with compatible effective image sizes and BAGEL CFG/renormalization
+settings. Sequence parallelism and diffusion cache backends such as TeaCache or
 Cache-DiT are not currently supported together with BAGEL step execution.
 
 ### Tensor Parallelism (TP)

recipes/Bagel/BAGEL-7B-MoT.md

@@ -212,8 +212,9 @@ vllm serve ByteDance-Seed/BAGEL-7B-MoT \
   --stage-overrides '{"1":{"max_num_seqs":4}}'
 ```
 
-For image requests, BAGEL builds exactly `num_inference_steps` scheduler
-updates, including one-step requests. The scheduler automatically separates
+Use at least two inference steps for image requests. BAGEL builds
+`num_inference_steps - 1` scheduler updates, so a one-step image request has no
+denoising iteration to schedule. The scheduler automatically separates
 requests with incompatible image shapes or BAGEL CFG/renormalization settings.
 
 Step execution currently does not support BAGEL sequence parallelism or a

tests/diffusion/models/bagel/test_step_execution.py

@@ -41,7 +41,7 @@ class _Sampling:
 
 @dataclass
 class _ScheduleConfig:
-    _denoise_schedule_extra_step: bool = True
+    _denoise_schedule_extra_step: bool = False
 
 
 @dataclass
@@ -245,8 +245,8 @@ def test_prepare_encode_populates_request_local_state_and_schedule():
     assert result is state
     assert state.latents is prepared["generation_input"]["packed_init_noises"]
     assert state.latents.dtype == torch.float32
-    assert state.total_steps == 5
-    expected_points = torch.linspace(1, 0, 6)
+    assert state.total_steps == 4
+    expected_points = torch.linspace(1, 0, 5)
     expected_points = 2.0 * expected_points / (1 + expected_points)
     assert torch.equal(state.timesteps, expected_points[:-1])
     assert state.timesteps.dtype == torch.float32
@@ -264,27 +264,28 @@ def test_prepare_encode_populates_request_local_state_and_schedule():
     assert "packed_init_noises" not in state.extra["bagel_generation_input"]
 
 
-def test_prepare_encode_rejects_non_positive_step_schedule():
-    prepared = _prepared_context()
-    prepared["gen_params"].num_timesteps = 0
-    pipeline = _pipeline_for_prepare(prepared)
-    state = StepRequestState(request_id="req", sampling=_sampling(), prompt="draw a cat")
-
-    with pytest.raises(ValueError, match="num_inference_steps >= 1"):
-        pipeline.prepare_encode(state)
-
-
-def test_prepare_encode_supports_one_denoising_update():
+def test_prepare_encode_rejects_one_step_schedule():
     prepared = _prepared_context()
     prepared["gen_params"].num_timesteps = 1
     pipeline = _pipeline_for_prepare(prepared)
     state = StepRequestState(request_id="req", sampling=_sampling(), prompt="draw a cat")
 
-    pipeline.prepare_encode(state)
+    with pytest.raises(ValueError, match="num_inference_steps >= 2"):
+        pipeline.prepare_encode(state)
 
-    assert state.total_steps == 1
-    assert len(state.timesteps) == 1
-    assert len(state.extra["bagel_dts"]) == 1
+
+def test_lance_style_extra_schedule_point_supports_one_denoising_update():
+    latents = torch.zeros(1, 3)
+
+    timesteps, dts = Bagel.prepare_denoise_schedule(
+        _ScheduleConfig(_denoise_schedule_extra_step=True),
+        latents,
+        num_timesteps=1,
+        timestep_shift=1.0,
+    )
+
+    assert len(timesteps) == 1
+    assert len(dts) == 1
 
 
 def test_prepare_encode_rejects_sequence_parallel_before_forward():
@@ -520,7 +521,7 @@ def test_complete_step_updates_match_full_bagel_loop(cfg_text_scale: float, velo
 
     bagel = MagicMock()
     bagel._sp_size = 1
-    bagel._denoise_schedule_extra_step = True
+    bagel._denoise_schedule_extra_step = False
     bagel.prepare_denoise_schedule = types.MethodType(Bagel.prepare_denoise_schedule, bagel)
     bagel.forward.side_effect = lambda x_t, **_kwargs: torch.full(
         x_t.shape,

tests/diffusion/models/bagel/test_trajectory_recording.py

@@ -19,14 +19,18 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 NUM_TOKENS = 8
 HIDDEN_DIM = 16
 NUM_TIMESTEPS = 5
-EXPECTED_STEPS = NUM_TIMESTEPS
+# Official BAGEL samples num_timesteps points then drops the terminal t=0,
+# yielding num_timesteps - 1 denoise steps.
+EXPECTED_STEPS = NUM_TIMESTEPS - 1
 
 
 def _make_mock_bagel(mocker: MockerFixture):
     """Create a mock Bagel with forward returning constant velocity."""
     mock = mocker.MagicMock(spec=Bagel)
     mock._sp_size = 1
-    mock._denoise_schedule_extra_step = True
+    # MagicMock would otherwise auto-return a truthy stub for this flag; pin it
+    # to the official BAGEL schedule convention (num_timesteps points).
+    mock._denoise_schedule_extra_step = False
 
     # forward returns a small constant velocity so x_t changes each step
     def fake_forward(self, x_t, **kwargs):
@@ -93,6 +97,14 @@ def bagel_and_args(
 class TestTrajectoryRecording:
     """Tests for trajectory latent/timestep recording in generate_image."""
 
+    def test_complete_request_rejects_one_step_schedule(self, bagel_and_args):
+        """The non-step image path must not decode untouched initial noise."""
+        bagel, args = bagel_and_args
+        args["num_timesteps"] = 1
+
+        with pytest.raises(ValueError, match="num_inference_steps >= 2"):
+            bagel.generate_image(**args)
+
     def test_trajectory_disabled_returns_none(self, bagel_and_args):
         bagel, args = bagel_and_args
 
@@ -178,7 +190,7 @@ class TestTrajectoryRecording:
         bagel, args = bagel_and_args
 
         # Recompute the expected timestep schedule (mirrors generate_image logic)
-        ts = torch.linspace(1, 0, args["num_timesteps"] + 1)
+        ts = torch.linspace(1, 0, args["num_timesteps"])
         shift = args.get("timestep_shift", 1.0)
         ts = shift * ts / (1 + (shift - 1) * ts)
         expected_timesteps = ts[:-1]  # last element is dropped
@@ -196,7 +208,7 @@ class TestTrajectoryRecording:
         bagel, args = bagel_and_args
 
         # Compute the number of denoising steps the same way generate_image does
-        ts = torch.linspace(1, 0, args["num_timesteps"] + 1)
+        ts = torch.linspace(1, 0, args["num_timesteps"])
         shift = args.get("timestep_shift", 1.0)
         ts = shift * ts / (1 + (shift - 1) * ts)
         num_steps = len(ts) - 1  # timesteps = ts[:-1]

tests/diffusion/test_diffusion_engine_dummy_run.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 from types import SimpleNamespace
 from unittest.mock import Mock
@@ -16,11 +16,10 @@ from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
 
 
-@pytest.mark.parametrize(("step_execution", "expected_steps"), [(False, 1), (True, 1)])
+@pytest.mark.parametrize("step_execution", [False, True])
 def test_dummy_run_uses_enough_steps_for_execution_mode(
     monkeypatch: pytest.MonkeyPatch,
     step_execution: bool,
-    expected_steps: int,
 ) -> None:
     engine = DiffusionEngine.__new__(DiffusionEngine)
     engine.od_config = type(
@@ -50,7 +49,7 @@ def test_dummy_run_uses_enough_steps_for_execution_mode(
     engine._dummy_run()
 
     assert len(captured_requests) == 1
-    assert captured_requests[0].sampling_params.num_inference_steps == expected_steps
+    assert captured_requests[0].sampling_params.num_inference_steps == 2
 
 
 def test_dummy_run_num_frames_uses_explicit_model_setting(monkeypatch: pytest.MonkeyPatch) -> None:

vllm_omni/diffusion/diffusion_engine.py

@@ -1141,7 +1141,11 @@ class DiffusionEngine:
             height=512,
             width=512,
             guidance_scale=0.0,
-            num_inference_steps=1,
+            # Dummy warmup must exercise at least one denoising iteration in
+            # every execution mode. Some pipelines (for example BAGEL) perform
+            # ``num_inference_steps - 1`` scheduler updates and reject an empty
+            # one-step schedule.
+            num_inference_steps=2,
         )
         if req is None:
             logger.info("Skipping dummy warmup run (num_frames=0)")

vllm_omni/diffusion/models/bagel/bagel_transformer.py

@@ -1256,10 +1256,12 @@ class Bagel(CFGParallelMixin, nn.Module):
     config_class = BagelConfig
     base_model_prefix = "bagel"
 
-    # Include the terminal t=0 sample so ``num_timesteps`` means exactly
-    # ``num_timesteps`` Euler updates. Subclasses may override this for a
-    # model-specific schedule convention.
-    _denoise_schedule_extra_step: bool = True
+    # Flow-matching denoise schedule convention. Official BAGEL samples
+    # ``num_timesteps`` points over [1, 0] and drops the terminal t=0, yielding
+    # ``num_timesteps - 1`` Euler steps. Lance samples one extra point
+    # (``num_timesteps + 1``) for ``num_timesteps`` steps; ``LanceBagel`` flips
+    # this on. See https://github.com/vllm-project/vllm-omni/issues/4470.
+    _denoise_schedule_extra_step: bool = False
 
     def prepare_denoise_schedule(
         self,
@@ -1268,6 +1270,8 @@ class Bagel(CFGParallelMixin, nn.Module):
         timestep_shift: float,
     ) -> tuple[torch.Tensor, torch.Tensor]:
         """Build BAGEL's flow-matching timesteps and per-step deltas."""
+        if num_timesteps < 2 and not self._denoise_schedule_extra_step:
+            raise ValueError("BAGEL image generation requires num_inference_steps >= 2.")
         num_sample_points = num_timesteps + 1 if self._denoise_schedule_extra_step else num_timesteps
         schedule = torch.linspace(1, 0, num_sample_points, device=x_t.device)
         schedule = timestep_shift * schedule / (1 + (timestep_shift - 1) * schedule)
@@ -1834,7 +1838,9 @@ class Bagel(CFGParallelMixin, nn.Module):
             frame_condition_token_indexes = frame_condition_token_indexes.to(x_t.device).long()
             pinned_x_t = x_t[frame_condition_token_indexes].clone()
 
-        # Build ``num_timesteps`` flow-matching Euler updates over [1, 0].
+        # Build the flow-matching schedule. BAGEL drops the terminal t=0 for
+        # ``num_timesteps - 1`` Euler steps; Lance keeps it for ``num_timesteps``.
+        # ``_denoise_schedule_extra_step`` (overridden by ``LanceBagel``) selects which.
         timesteps, dts = Bagel.prepare_denoise_schedule(self, x_t, num_timesteps, timestep_shift)
 
         # Optional trajectory recording for RL rollout data collection

vllm_omni/diffusion/models/bagel/pipeline_bagel.py

@@ -1091,8 +1091,6 @@ class BagelPipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipelineProf
             gen_params.num_timesteps,
             gen_params.timestep_shift,
         )
-        if len(timesteps) == 0:
-            raise ValueError("BAGEL step execution requires num_inference_steps >= 1.")
 
         # Keep request-local scheduler progress while sharing immutable config
         # and tensor buffers with the pipeline scheduler.

vllm_omni/diffusion/models/lance/lance_transformer.py

@@ -278,9 +278,9 @@ class LanceBagel(Bagel):
     ``forward_cache_update_vae`` — is reused unchanged.
     """
 
-    # Upstream Lance uses the same ``num_timesteps + 1`` schedule-point
-    # convention as BAGEL, so the denoise loop runs exactly ``num_timesteps``
-    # Euler steps.
+    # Upstream Lance samples ``num_timesteps + 1`` schedule points (one extra
+    # over official BAGEL) so the denoise loop runs exactly ``num_timesteps``
+    # Euler steps. See Bagel.generate_image and issue #4470.
     _denoise_schedule_extra_step: bool = True
 
     @staticmethod