diff --git a/docs/design/feature/diffusion_continuous_batching.md b/docs/design/feature/diffusion_continuous_batching.md index e12455c3d..9bb41fdd9 100644 --- a/docs/design/feature/diffusion_continuous_batching.md +++ b/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 diff --git a/docs/user_guide/diffusion/execution_modes.md b/docs/user_guide/diffusion/execution_modes.md index a4f3f0f50..3aa625295 100644 --- a/docs/user_guide/diffusion/execution_modes.md +++ b/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 diff --git a/docs/user_guide/diffusion_features.md b/docs/user_guide/diffusion_features.md index c616a284c..a06f53d3e 100644 --- a/docs/user_guide/diffusion_features.md +++ b/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 diff --git a/docs/user_guide/examples/online_serving/bagel.md b/docs/user_guide/examples/online_serving/bagel.md index 83d49ceb9..6c10c4349 100644 --- a/docs/user_guide/examples/online_serving/bagel.md +++ b/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) diff --git a/recipes/Bagel/BAGEL-7B-MoT.md b/recipes/Bagel/BAGEL-7B-MoT.md index 1044ad66b..19ace3d33 100644 --- a/recipes/Bagel/BAGEL-7B-MoT.md +++ b/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 diff --git a/tests/diffusion/models/bagel/test_step_execution.py b/tests/diffusion/models/bagel/test_step_execution.py index a17463ef6..40d4154e3 100644 --- a/tests/diffusion/models/bagel/test_step_execution.py +++ b/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, diff --git a/tests/diffusion/models/bagel/test_trajectory_recording.py b/tests/diffusion/models/bagel/test_trajectory_recording.py index 908b1b732..38169d750 100644 --- a/tests/diffusion/models/bagel/test_trajectory_recording.py +++ b/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] diff --git a/tests/diffusion/test_diffusion_engine_dummy_run.py b/tests/diffusion/test_diffusion_engine_dummy_run.py index 6bf0e83ea..ea37b2fcc 100644 --- a/tests/diffusion/test_diffusion_engine_dummy_run.py +++ b/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: diff --git a/vllm_omni/diffusion/diffusion_engine.py b/vllm_omni/diffusion/diffusion_engine.py index b96b82f07..f0dc4b0e4 100644 --- a/vllm_omni/diffusion/diffusion_engine.py +++ b/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)") diff --git a/vllm_omni/diffusion/models/bagel/bagel_transformer.py b/vllm_omni/diffusion/models/bagel/bagel_transformer.py index 84b91b182..ddaa5693d 100644 --- a/vllm_omni/diffusion/models/bagel/bagel_transformer.py +++ b/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 diff --git a/vllm_omni/diffusion/models/bagel/pipeline_bagel.py b/vllm_omni/diffusion/models/bagel/pipeline_bagel.py index 8d71bf6c2..f27601148 100644 --- a/vllm_omni/diffusion/models/bagel/pipeline_bagel.py +++ b/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. diff --git a/vllm_omni/diffusion/models/lance/lance_transformer.py b/vllm_omni/diffusion/models/lance/lance_transformer.py index b3a281de0..12959eb64 100644 --- a/vllm_omni/diffusion/models/lance/lance_transformer.py +++ b/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