[Model] Add BAGEL step execution and continuous batching (#6359)

docs/design/feature/diffusion_continuous_batching.md

@@ -150,8 +150,17 @@ that do not belong in the shared contract should be stored in `state.extra`.
 Queueing and lifecycle metadata remain in the scheduler's request state.
 
 Current native pipelines that explicitly enable step execution include
-Qwen-Image, HunyuanImage3, and Helios. Step execution alone does not imply
-continuous-batching support: Qwen-Image accepts batched step states.
+Qwen-Image, BAGEL, HunyuanImage3, and Helios. Step execution alone does not
+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.
 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

@@ -94,8 +94,16 @@ validating a model or debugging correctness, then increase it for
 multi-request throughput.
 
 Step execution is capability-based, not a generic switch for every diffusion
-model. Qwen-Image supports step-wise continuous batching. HunyuanImage3 also
-supports it, but only when its resolved self-attention backend is `TORCH_SDPA`;
+model. Qwen-Image supports step-wise continuous batching. BAGEL supports
+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
+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
 `--max-num-seqs >1`. See the

docs/user_guide/diffusion_features.md

@@ -116,7 +116,7 @@ The following tables show which models support each feature:
 
 | Model                    | ⚡TeaCache | ⚡Cache-DiT | 🔀SP (Ulysses & Ring) | 🔀CFG-Parallel | 🔀Tensor-Parallel | 🔀Pipeline-Parallel | 🔀HSDP | 💾CPU Offload (Layerwise) | 💾VAE-Patch-Parallel | 💾Quantization | 🔄Step Execution |
 |--------------------------|:---------:|:----------:|:---------------------:|:--------------:|:-----------------:|:-------------------:|:------:|:-------------------------:|:--------------------:|:--------------:|:----------------:|
-| **Bagel**                |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ✅    |             ✅             |          ❌           |       ❌        |        ❌         |
+| **Bagel**                |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ✅    |             ✅             |          ❌           |       ❌        |        ✅         |
 | **Boogu-Image**          |     ❌     |     ❌      |           ❌           |       ✅        |         ❌         |          ❌          |   ❌    |             ❌             |          ❌           |       ❌        |        ❌         |
 | **FLUX.1-dev**           |     ✅     |     ✅      |           ❌           |       ✅        |         ✅         |          ❌          |   ✅    |             ❌             |          ❌           |       ✅        |        ❌         |
 | **FLUX.1-schnell**       |     ❌     |     ✅      |           ❌           |       ✅        |         ✅         |          ❌          |   ✅    |             ❌             |          ❌           |       ✅        |        ❌         |
@@ -152,6 +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).
 
 ### VideoGen
 

docs/user_guide/examples/online_serving/bagel.md

@@ -60,6 +60,50 @@ vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091 \
 
 See [`bagel_single_stage.yaml`](https://github.com/vllm-project/vllm-omni/tree/main/vllm_omni/deploy/bagel_single_stage.yaml) for configuration. The `pipeline: bagel_single_stage` field selects the single-stage topology from the pipeline registry.
 
+### Step Execution and Continuous Batching
+
+Image generation supports step execution with all three BAGEL deploy configs.
+The global flag is applied only to diffusion stages:
+
+- With `bagel.yaml` or `bagel_think.yaml`, Stage 1 uses step execution and the
+  Stage 0 Thinker remains autoregressive.
+- With `bagel_single_stage.yaml`, image requests use step execution in Stage 0.
+  Explicit `text2text` and `img2text` requests use the same complete-request
+  text path as they do without `--step-execution`.
+
+Start with the deploy file's default diffusion capacity of one:
+
+```bash
+vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091 \
+    --deploy-config vllm_omni/deploy/bagel.yaml \
+    --step-execution
+```
+
+To enable step-wise continuous batching, override `max_num_seqs` on the
+diffusion stage. Use Stage 1 for either two-stage config:
+
+```bash
+vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091 \
+    --deploy-config vllm_omni/deploy/bagel.yaml \
+    --step-execution \
+    --stage-overrides '{"1":{"max_num_seqs":4}}'
+```
+
+Use Stage 0 for the single-stage config:
+
+```bash
+vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091 \
+    --deploy-config vllm_omni/deploy/bagel_single_stage.yaml \
+    --step-execution \
+    --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
+Cache-DiT are not currently supported together with BAGEL step execution.
+
 ### Tensor Parallelism (TP)
 
 For larger models or multi-GPU environments, enable TP via CLI:
@@ -203,6 +247,13 @@ curl http://localhost:8091/v1/chat/completions \
   -d @payload.json
 ```
 
+BAGEL img2img derives its generated image size from the input image. It
+preserves the input aspect ratio, aligns dimensions to the latent stride, and
+applies the checkpoint size limit; explicit `height` and `width` values do not
+override this model-specific resize policy. In step mode, the effective size is
+resolved before scheduler admission so requests with different resulting
+shapes are not combined.
+
 ### Image to Text (img2text)
 
 **Python client:**
@@ -267,8 +318,8 @@ curl http://localhost:8091/v1/chat/completions \
 | `--server` / `-s` | `http://localhost:8091` | Server URL |
 | `--image-url` / `-i` | `None` | Input image URL or local path (img2img/img2text) |
 | `--modality` / `-m` | `text2img` | `text2img`, `img2img`, `img2text`, `text2text` |
-| `--height` | `512` | Image height in pixels |
-| `--width` | `512` | Image width in pixels |
+| `--height` | `512` | Text-to-image output height; BAGEL img2img derives its size from the input image |
+| `--width` | `512` | Text-to-image output width; BAGEL img2img derives its size from the input image |
 | `--steps` | `25` | Number of inference steps |
 | `--seed` | `42` | Random seed |
 | `--negative` | `None` | Negative prompt for CFG |
@@ -293,6 +344,7 @@ python openai_chat_client.py \
 | File | Description |
 | :--- | :---------- |
 | [`bagel.yaml`](https://github.com/vllm-project/vllm-omni/tree/main/vllm_omni/deploy/bagel.yaml) | Two-stage default (Thinker + DiT on GPU 0) |
+| [`bagel_think.yaml`](https://github.com/vllm-project/vllm-omni/tree/main/vllm_omni/deploy/bagel_think.yaml) | Two-stage think mode (Thinker + DiT) |
 | [`bagel_single_stage.yaml`](https://github.com/vllm-project/vllm-omni/tree/main/vllm_omni/deploy/bagel_single_stage.yaml) | Single-stage (DiT only) |
 
 ### Key Deploy YAML Fields

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

@@ -34,6 +34,7 @@ parameters through the pipeline-declared `extra_args` contract.
   [`examples/online_serving/image_to_image/openai_chat_client.py`](../../examples/online_serving/image_to_image/openai_chat_client.py)
 - Default deploy configs:
   [`vllm_omni/deploy/bagel.yaml`](../../vllm_omni/deploy/bagel.yaml),
+  [`vllm_omni/deploy/bagel_think.yaml`](../../vllm_omni/deploy/bagel_think.yaml),
   [`vllm_omni/deploy/bagel_single_stage.yaml`](../../vllm_omni/deploy/bagel_single_stage.yaml)
 
 ## Hardware Support
@@ -119,6 +120,12 @@ python examples/offline_inference/image_to_image/image_edit.py \
 The `--extra-args` JSON forwards BAGEL-specific parameters (e.g. `cfg_text_scale`,
 `cfg_img_scale`, `cfg_interval`, `cfg_renorm_type`) into
 `OmniDiffusionSamplingParams.extra_args` via the model-extras registry.
+For BAGEL img2img, the pipeline derives the generated image size from the input
+image: it preserves the input aspect ratio, aligns dimensions to the latent
+stride, and applies the checkpoint size limit. Explicit `--height` and
+`--width` values do not override that BAGEL-specific resize policy. Step-mode
+admission resolves this effective size before batching so differently shaped
+requests are scheduled separately.
 
 Run text-to-text with the shared understanding example. BAGEL's default
 `bagel.yaml` deploy config is discovered from the checkpoint, so no model-specific
@@ -156,6 +163,99 @@ vllm serve ByteDance-Seed/BAGEL-7B-MoT \
   --deploy-config vllm_omni/deploy/bagel_single_stage.yaml
 ```
 
+#### Step Execution and Continuous Batching
+
+BAGEL supports step execution for image generation with all three in-repo
+deploy configs:
+
+- `bagel.yaml`: the flag applies to the Stage 1 diffusion engine; Stage 0
+  remains on the normal autoregressive path.
+- `bagel_think.yaml`: the flag also applies only to the Stage 1 diffusion
+  engine; Thinker decoding in Stage 0 is unchanged.
+- `bagel_single_stage.yaml`: image requests use step execution in Stage 0.
+  Explicit text-output requests keep the existing complete-request path, so
+  `text2text` and `img2text` behavior is unchanged.
+
+Start with one active request when validating correctness. For example, the
+following command uses the single-stage topology:
+
+```bash
+vllm serve ByteDance-Seed/BAGEL-7B-MoT \
+  --omni \
+  --port 8091 \
+  --deploy-config vllm_omni/deploy/bagel_single_stage.yaml \
+  --step-execution
+```
+
+To let compatible requests share denoising waves, override the diffusion-stage
+capacity. The deploy files set `max_num_seqs: 1` for their diffusion stages, so
+the per-stage value takes precedence over the global `--max-num-seqs` flag.
+Use Stage 0 for `bagel_single_stage.yaml`:
+
+```bash
+vllm serve ByteDance-Seed/BAGEL-7B-MoT \
+  --omni \
+  --port 8091 \
+  --deploy-config vllm_omni/deploy/bagel_single_stage.yaml \
+  --step-execution \
+  --stage-overrides '{"0":{"max_num_seqs":4}}'
+```
+
+For either two-stage config, use Stage 1 instead:
+
+```bash
+vllm serve ByteDance-Seed/BAGEL-7B-MoT \
+  --omni \
+  --port 8091 \
+  --deploy-config vllm_omni/deploy/bagel.yaml \
+  --step-execution \
+  --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
+requests with incompatible image shapes or BAGEL CFG/renormalization settings.
+
+Step execution currently does not support BAGEL sequence parallelism or a
+diffusion cache backend. Do not combine it with BAGEL SP, TeaCache, or
+Cache-DiT. Tensor parallelism and CFG parallelism are separate features. For
+best utilization, match the CFG parallel world size to the number of active
+BAGEL CFG branches. A larger world size is supported, but extra ranks run dummy
+forward passes to participate in collectives and their outputs are ignored.
+BAGEL currently does not support two CFG ranks when three CFG branches are
+active.
+
+For a matched throughput comparison, start the first server with the default
+stage capacity of one and run:
+
+```bash
+python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
+  --base-url http://localhost:8091 \
+  --endpoint /v1/chat/completions \
+  --model ByteDance-Seed/BAGEL-7B-MoT \
+  --task t2i \
+  --dataset random \
+  --num-prompts 20 \
+  --max-concurrency 4 \
+  --width 512 \
+  --height 512 \
+  --num-inference-steps 20 \
+  --seed 42 \
+  --warmup-requests 4 \
+  --warmup-concurrency 4 \
+  --warmup-num-inference-steps 2 \
+  --output-file /tmp/bagel_step_seq1.json \
+  --disable-tqdm
+```
+
+Restart the single-stage server with `--stage-overrides
+'{"0":{"max_num_seqs":4}}'` (or use Stage 1 for either two-stage config),
+rerun the same benchmark with
+`--output-file /tmp/bagel_step_seq4.json`, and compare completed requests,
+duration, throughput, and latency percentiles. Keep the model, device topology,
+request count, client concurrency, resolution, steps, seed, and warmup settings
+fixed between runs.
+
 Send a text-to-image request with BAGEL-specific generation parameters:
 
 ```bash

tests/config/test_omni_config.py

@@ -339,6 +339,37 @@ def test_stage_cli_field_selection_defers_ownership_validation_until_sources_are
     assert omni_config_module._stage_cli_overrides(0, {"enable_lora": True}) == {"enable_lora": True}
 
 
+@pytest.mark.parametrize(
+    ("model_type", "diffusion_stage_ids"),
+    [
+        ("bagel", {1}),
+        ("bagel_think", {1}),
+        ("bagel_single_stage", {0}),
+    ],
+)
+def test_step_execution_cli_is_scoped_to_diffusion_stages(model_type, diffusion_stage_ids):
+    omni_config = _from_pipeline_key(model_type, cli_overrides={"step_execution": True})
+
+    for stage in omni_config.stage_configs:
+        if stage.stage_id in diffusion_stage_ids:
+            assert isinstance(stage, VllmOmniDiffusionStageConfig)
+            assert stage.diffusion_config.step_execution is True
+        else:
+            assert not hasattr(stage, "diffusion_config")
+
+
+def test_stage_scoped_step_execution_ignores_bagel_ar_stage():
+    omni_config = _from_pipeline_key(
+        "bagel",
+        cli_overrides={
+            "stage_0_step_execution": True,
+            "stage_1_step_execution": True,
+        },
+    )
+
+    assert omni_config.stage_by_id(1).diffusion_config.step_execution is True
+
+
 def test_runtime_num_gpus_is_derived_from_parallel_world_size():
     omni_config = _from_pipeline_key("hunyuan_image3_dit")
     stage = omni_config.stage_by_id(0)

tests/diffusion/kernels/mot/test_mot_norm_npu.py

@@ -0,0 +1,36 @@
+import pytest
+import torch
+
+import vllm_omni.diffusion.layers.custom_op as custom_op
+from vllm_omni.diffusion.layers.mot.mot_layernorm import MoTRMSNorm
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+class _NpuPlatform:
+    is_rocm = staticmethod(lambda: False)
+    is_cuda = staticmethod(lambda: False)
+    is_npu = staticmethod(lambda: True)
+    is_xpu = staticmethod(lambda: False)
+    is_musa = staticmethod(lambda: False)
+
+
+@pytest.mark.parametrize("use_mot_routing", [False, True])
+def test_mot_rms_norm_npu_falls_back_to_native(monkeypatch, use_mot_routing: bool):
+    monkeypatch.setattr(
+        custom_op,
+        "current_omni_platform",
+        _NpuPlatform(),
+    )
+
+    layer = MoTRMSNorm(hidden_size=8)
+    assert layer._forward_method.__func__ is MoTRMSNorm.forward_npu
+
+    x = torch.randn(6, 8)
+    text_indices = torch.tensor([0, 2, 4]) if use_mot_routing else None
+    vae_indices = torch.tensor([1, 3, 5]) if use_mot_routing else None
+
+    expected = layer.forward_native(x, text_indices, vae_indices)
+    actual = layer(x, text_indices, vae_indices)
+
+    torch.testing.assert_close(actual, expected)

tests/diffusion/models/bagel/test_step_execution.py

@@ -0,0 +1,753 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Focused CPU tests for BAGEL step execution and packed batching."""
+
+from __future__ import annotations
+
+import types
+from dataclasses import dataclass, field
+from unittest.mock import MagicMock, patch
+
+import pytest
+import torch
+from PIL import Image
+from safetensors.torch import save_file
+from torch import nn
+
+from vllm_omni.diffusion.models.bagel.bagel_transformer import Bagel, NaiveCache
+from vllm_omni.diffusion.models.bagel.pipeline_bagel import (
+    BagelGenParams,
+    BagelPipeline,
+    get_bagel_pre_process_func,
+)
+from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.diffusion.sched import StepScheduler
+from vllm_omni.diffusion.worker.input_batch import InputBatch
+from vllm_omni.diffusion.worker.utils import StepRequestState
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+@dataclass
+class _Sampling:
+    return_trajectory_latents: bool
+    return_trajectory_decoded: bool = False
+    true_cfg_scale: float = 4.0
+    cfg_normalize: bool = False
+    image_latent: torch.Tensor | None = None
+    extra_args: dict = field(default_factory=dict)
+
+
+@dataclass
+class _ScheduleConfig:
+    _denoise_schedule_extra_step: bool = True
+
+
+@dataclass
+class _ParallelConfig:
+    cfg_parallel_size: int
+
+
+@dataclass
+class _LanguageModelOutput:
+    packed_query_sequence: torch.Tensor
+
+
+def _sampling(*, trajectory: bool = False) -> _Sampling:
+    return _Sampling(
+        return_trajectory_latents=trajectory,
+    )
+
+
+@pytest.mark.parametrize(
+    ("prompt", "use_step_execution"),
+    [
+        ({"prompt": "describe this image", "modalities": ["text"]}, False),
+        ({"prompt": "draw a cat", "modalities": ["image"]}, True),
+        ("describe this image", True),
+    ],
+)
+def test_bagel_step_preprocessor_only_falls_back_for_explicit_text(prompt, use_step_execution):
+    pre_process = get_bagel_pre_process_func(types.SimpleNamespace(step_execution=True))
+    request = OmniDiffusionRequest(
+        prompt=prompt,
+        sampling_params=OmniDiffusionSamplingParams(num_inference_steps=2),
+        request_id="req",
+    )
+
+    assert pre_process(request) is request
+    assert request.use_step_execution is use_step_execution
+
+
+def test_bagel_step_preprocessor_defaults_to_non_step_when_omitted():
+    pre_process = get_bagel_pre_process_func(types.SimpleNamespace())
+    request = OmniDiffusionRequest(
+        prompt={"prompt": "describe this image", "modalities": ["text"]},
+        sampling_params=OmniDiffusionSamplingParams(num_inference_steps=2),
+        request_id="req",
+    )
+
+    assert pre_process(request) is request
+    assert request.use_step_execution is True
+
+
+@pytest.mark.parametrize(
+    ("first_extra_args", "second_extra_args"),
+    [
+        ({"cfg_text_scale": 1.0}, {"cfg_text_scale": 4.0}),
+        ({"cfg_img_scale": 1.0}, {"cfg_img_scale": 1.5}),
+        ({"cfg_interval": (0.0, 1.0)}, {"cfg_interval": (0.4, 1.0)}),
+        ({"cfg_renorm_type": "global"}, {"cfg_renorm_type": "channel"}),
+        ({"cfg_renorm_min": 0.0}, {"cfg_renorm_min": 0.5}),
+    ],
+)
+def test_bagel_step_preprocessor_separates_incompatible_cfg_settings(first_extra_args, second_extra_args):
+    pre_process = get_bagel_pre_process_func(types.SimpleNamespace(step_execution=True))
+    requests = [
+        OmniDiffusionRequest(
+            prompt={"prompt": "draw a cat", "modalities": ["image"]},
+            sampling_params=OmniDiffusionSamplingParams(num_inference_steps=2, extra_args=extra_args),
+            request_id=request_id,
+        )
+        for request_id, extra_args in (("a", first_extra_args), ("b", second_extra_args))
+    ]
+
+    first, second = (pre_process(request) for request in requests)
+
+    assert first.batch_compatibility_key != second.batch_compatibility_key
+
+
+def test_bagel_step_preprocessor_buckets_effective_img2img_sizes(tmp_path):
+    (tmp_path / "config.json").write_text(
+        '{"vae_config":{"downsample":8},"latent_patch_size":2,"max_latent_size":32}',
+        encoding="utf-8",
+    )
+    save_file(
+        {"latent_pos_embed.pos_embed": torch.zeros(64 * 64, 1)},
+        tmp_path / "ema.safetensors",
+    )
+    (tmp_path / "model.safetensors.index.json").write_text(
+        '{"weight_map":{"latent_pos_embed.pos_embed":"ema.safetensors"}}',
+        encoding="utf-8",
+    )
+    pre_process = get_bagel_pre_process_func(
+        types.SimpleNamespace(model=str(tmp_path), revision=None, step_execution=True)
+    )
+    requests = [
+        OmniDiffusionRequest(
+            prompt={
+                "prompt": "edit this image",
+                "modalities": ["img2img"],
+                "multi_modal_data": {"img2img": Image.new("RGB", size)},
+            },
+            sampling_params=OmniDiffusionSamplingParams(num_inference_steps=2),
+            request_id=request_id,
+        )
+        for request_id, size in (("a", (800, 400)), ("b", (1024, 512)))
+    ]
+    first, second = (pre_process(request) for request in requests)
+
+    assert (first.sampling_params.height, first.sampling_params.width) == (400, 800)
+    assert (second.sampling_params.height, second.sampling_params.width) == (512, 1024)
+
+    scheduler = StepScheduler()
+    scheduler.initialize(types.SimpleNamespace(max_num_seqs=2))
+    scheduler.add_request(first)
+    scheduler.add_request(second)
+
+    schedule = scheduler.schedule()
+
+    assert schedule.scheduled_request_ids == ["a"]
+    assert schedule.num_running_reqs == 1
+    assert schedule.num_waiting_reqs == 1
+
+
+def _cache(length: int, value: float, num_layers: int = 1) -> NaiveCache:
+    cache = NaiveCache(num_layers)
+    if length:
+        for layer in range(num_layers):
+            cache.key_cache[layer] = torch.full((length, 1, 2), value)
+            cache.value_cache[layer] = torch.full((length, 1, 2), value)
+    return cache
+
+
+def _generation_input(num_latents: int = 4) -> dict[str, torch.Tensor]:
+    seq_len = num_latents + 2
+    return {
+        "packed_text_ids": torch.tensor([10, 11]),
+        "packed_text_indexes": torch.tensor([0, seq_len - 1]),
+        "packed_init_noises": torch.arange(num_latents * 3, dtype=torch.float32).reshape(num_latents, 3),
+        "packed_vae_position_ids": torch.arange(num_latents),
+        "packed_vae_token_indexes": torch.arange(1, seq_len - 1),
+        "packed_seqlens": torch.tensor([seq_len], dtype=torch.int32),
+        "packed_position_ids": torch.arange(seq_len),
+    }
+
+
+def _prepared_context(
+    *,
+    kv_len: int = 2,
+    cfg_text_scale: float = 1.0,
+    trajectory: bool = False,
+    think_text: str | None = None,
+) -> dict:
+    del trajectory
+    generation_input = _generation_input()
+    context = {
+        "kv_lens": [kv_len],
+        "ropes": [kv_len],
+        "past_key_values": _cache(kv_len, float(kv_len)),
+    }
+    return {
+        "generation_input": generation_input,
+        "cfg_text_packed_position_ids": generation_input["packed_position_ids"] + 20,
+        "cfg_img_packed_position_ids": generation_input["packed_position_ids"] + 40,
+        "gen_context": context,
+        "cfg_text_context": {
+            "kv_lens": [kv_len + 1],
+            "ropes": [kv_len + 1],
+            "past_key_values": _cache(kv_len + 1, float(kv_len + 1)),
+        },
+        "cfg_img_context": context,
+        "gen_params": BagelGenParams(num_timesteps=5, timestep_shift=2.0, cfg_text_scale=cfg_text_scale),
+        "image_shape": (64, 64),
+        "think_text": think_text,
+    }
+
+
+def _pipeline_for_prepare(prepared: dict) -> BagelPipeline:
+    pipeline = MagicMock()
+    pipeline.scheduler = None
+    pipeline.scheduler_kwargs = {}
+    pipeline.bagel._sp_size = 1
+    pipeline._forward_single.return_value = prepared
+    pipeline.bagel.prepare_denoise_schedule.side_effect = lambda latents, steps, shift: Bagel.prepare_denoise_schedule(
+        _ScheduleConfig(), latents, steps, shift
+    )
+    pipeline.prepare_encode = types.MethodType(BagelPipeline.prepare_encode, pipeline)
+    return pipeline
+
+
+def _prepare_state(request_id: str, *, kv_len: int, cfg_text_scale: float = 1.0) -> StepRequestState:
+    state = StepRequestState(request_id=request_id, sampling=_sampling(trajectory=True), prompt="draw a cat")
+    pipeline = _pipeline_for_prepare(_prepared_context(kv_len=kv_len, cfg_text_scale=cfg_text_scale))
+    return pipeline.prepare_encode(state)
+
+
+def test_prepare_encode_populates_request_local_state_and_schedule():
+    prepared = _prepared_context(kv_len=3)
+    pipeline = _pipeline_for_prepare(prepared)
+    state = StepRequestState(request_id="req", sampling=_sampling(trajectory=True), prompt="draw a cat")
+
+    result = pipeline.prepare_encode(state)
+
+    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)
+    expected_points = 2.0 * expected_points / (1 + expected_points)
+    assert torch.equal(state.timesteps, expected_points[:-1])
+    assert state.timesteps.dtype == torch.float32
+    assert state.timesteps.device == state.latents.device
+    assert torch.equal(state.extra["bagel_dts"], expected_points[:-1] - expected_points[1:])
+    assert state.extra["bagel_dts"].dtype == torch.float32
+    assert state.extra["bagel_generation_input"]["packed_text_ids"].dtype == torch.long
+    assert state.extra["bagel_generation_input"]["packed_text_indexes"].dtype == torch.long
+    assert state.extra["bagel_generation_input"]["packed_vae_position_ids"].dtype == torch.long
+    assert state.extra["bagel_generation_input"]["packed_vae_token_indexes"].dtype == torch.long
+    assert state.extra["bagel_generation_input"]["packed_position_ids"].dtype == torch.long
+    assert state.extra["bagel_generation_input"]["packed_seqlens"].dtype == torch.int32
+    assert state.extra["bagel_gen_context"] is prepared["gen_context"]
+    assert len(state.extra["bagel_trajectory_latents"]) == 1
+    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():
+    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)
+
+    assert state.total_steps == 1
+    assert len(state.timesteps) == 1
+    assert len(state.extra["bagel_dts"]) == 1
+
+
+def test_prepare_encode_rejects_sequence_parallel_before_forward():
+    pipeline = _pipeline_for_prepare(_prepared_context())
+    pipeline.bagel._sp_size = 2
+    state = StepRequestState(request_id="req", sampling=_sampling(), prompt="draw a cat")
+
+    with pytest.raises(NotImplementedError, match="does not currently support sequence parallelism"):
+        pipeline.prepare_encode(state)
+
+    pipeline._forward_single.assert_not_called()
+
+
+def test_step_protocol_is_enabled_for_bagel():
+    from vllm_omni.diffusion.models.interface import supports_step_execution
+
+    bagel = object.__new__(BagelPipeline)
+
+    assert supports_step_execution(bagel)
+
+
+def test_prepare_encode_reuses_image_and_think_generation_context():
+    pipeline = _pipeline_for_prepare(_prepared_context(think_text="reasoning"))
+    sampling = _sampling()
+    sampling.extra_args["think"] = True
+    state = StepRequestState(
+        request_id="req",
+        sampling=sampling,
+        prompt={"prompt": "edit", "multi_modal_data": {"image": "input.png"}},
+    )
+
+    result = pipeline.prepare_encode(state)
+
+    assert result is state
+    pipeline._forward_single.assert_called_once_with(state.prompt, sampling, prepare_only=True)
+    assert state.extra["bagel_think_text"] == "reasoning"
+
+
+def test_two_request_packed_indexes_are_rebased():
+    state_a = _prepare_state("a", kv_len=2)
+    state_b = _prepare_state("b", kv_len=3)
+
+    packed = BagelPipeline._pack_step_generation_inputs([state_a, state_b])
+
+    assert packed["packed_seqlens"].tolist() == [6, 6]
+    assert packed["packed_text_indexes"].tolist() == [0, 5, 6, 11]
+    assert packed["packed_vae_token_indexes"].tolist() == [1, 2, 3, 4, 7, 8, 9, 10]
+
+
+def test_different_packed_sequence_lengths_are_rejected():
+    states = []
+    for index, num_latents in enumerate((2, 4)):
+        state = StepRequestState(request_id=str(index), sampling=_sampling(), prompt="draw a cat")
+        state.extra["bagel_generation_input"] = _generation_input(num_latents)
+        states.append(state)
+
+    with pytest.raises(ValueError, match="matching packed sequence lengths"):
+        BagelPipeline._pack_step_generation_inputs(states)
+
+
+def test_denoise_step_uses_one_forward_and_preserves_request_kv_lengths():
+    states = [
+        _prepare_state("a", kv_len=2, cfg_text_scale=4.0),
+        _prepare_state("b", kv_len=3, cfg_text_scale=4.0),
+    ]
+    input_batch = InputBatch.make_batch(states)
+    expected = torch.ones_like(input_batch.latents)
+
+    pipeline = MagicMock()
+    pipeline.device = torch.device("cpu")
+    pipeline.od_config.dtype = torch.float32
+    pipeline.bagel._sp_size = 1
+    pipeline.bagel.parallel_config = _ParallelConfig(cfg_parallel_size=1)
+    pipeline.bagel.forward.return_value = expected
+    pipeline._pack_step_generation_inputs = BagelPipeline._pack_step_generation_inputs
+    pipeline._build_denoise_kwargs = types.MethodType(BagelPipeline._build_denoise_kwargs, pipeline)
+    pipeline.denoise_step = types.MethodType(BagelPipeline.denoise_step, pipeline)
+
+    actual = pipeline.denoise_step(input_batch, states=states)
+
+    assert actual is expected
+    pipeline.bagel.forward.assert_called_once()
+    call = pipeline.bagel.forward.call_args.kwargs
+    assert call["packed_text_indexes"].tolist() == [0, 5, 6, 11]
+    assert call["past_key_values"].key_values_lens == [2, 3]
+    assert call["cfg_branch_caches"][1].key_values_lens == [3, 4]
+    assert call["cfg_vae_lengths"] == [4, 4]
+    assert len(call["cfg_branch_pids"]) == 3
+
+
+def test_denoise_step_rejects_mixed_cfg_settings_defensively():
+    states = [
+        _prepare_state("a", kv_len=2, cfg_text_scale=2.0),
+        _prepare_state("b", kv_len=3, cfg_text_scale=4.0),
+    ]
+    input_batch = InputBatch.make_batch(states)
+    pipeline = MagicMock()
+    pipeline.bagel.parallel_config = _ParallelConfig(cfg_parallel_size=1)
+    pipeline._pack_step_generation_inputs = BagelPipeline._pack_step_generation_inputs
+    pipeline._build_denoise_kwargs = types.MethodType(BagelPipeline._build_denoise_kwargs, pipeline)
+
+    with pytest.raises(ValueError, match="Mixed BAGEL CFG settings"):
+        pipeline._build_denoise_kwargs(input_batch, states)
+
+
+def test_denoise_step_dispatches_three_cfg_parallel_branches():
+    states = [
+        _prepare_state("a", kv_len=2, cfg_text_scale=4.0),
+        _prepare_state("b", kv_len=3, cfg_text_scale=4.0),
+    ]
+    input_batch = InputBatch.make_batch(states)
+    expected = torch.ones_like(input_batch.latents)
+
+    pipeline = MagicMock()
+    pipeline.device = torch.device("cpu")
+    pipeline.od_config.dtype = torch.float32
+    pipeline.bagel._sp_size = 1
+    pipeline.bagel.parallel_config = _ParallelConfig(cfg_parallel_size=3)
+    pipeline.bagel.predict_noise_with_multi_branch_cfg.return_value = expected
+    pipeline._pack_step_generation_inputs = BagelPipeline._pack_step_generation_inputs
+    pipeline._build_denoise_kwargs = types.MethodType(BagelPipeline._build_denoise_kwargs, pipeline)
+    pipeline._denoise_step_cfg_parallel = types.MethodType(BagelPipeline._denoise_step_cfg_parallel, pipeline)
+    pipeline.denoise_step = types.MethodType(BagelPipeline.denoise_step, pipeline)
+    cfg_group = MagicMock()
+
+    with (
+        patch(
+            "vllm_omni.diffusion.models.bagel.pipeline_bagel.get_classifier_free_guidance_world_size",
+            return_value=3,
+        ),
+        patch("vllm_omni.diffusion.models.bagel.pipeline_bagel.get_cfg_group", return_value=cfg_group),
+    ):
+        actual = pipeline.denoise_step(input_batch, states=states)
+
+    assert actual is expected
+    assert cfg_group.broadcast.call_args.args[0] is input_batch.latents
+    call = pipeline.bagel.predict_noise_with_multi_branch_cfg.call_args.kwargs
+    assert call["do_true_cfg"]
+    assert len(call["branches_kwargs"]) == 3
+    assert call["true_cfg_scale"]["cfg_vae_lengths"] == [4, 4]
+
+
+def test_denoise_step_allows_idle_cfg_rank_for_two_branches():
+    states = [
+        _prepare_state("a", kv_len=2, cfg_text_scale=4.0),
+        _prepare_state("b", kv_len=3, cfg_text_scale=4.0),
+    ]
+    for state in states:
+        state.extra["bagel_gen_params"].cfg_img_scale = 1.0
+    input_batch = InputBatch.make_batch(states)
+    expected = torch.ones_like(input_batch.latents)
+
+    pipeline = MagicMock()
+    pipeline.device = torch.device("cpu")
+    pipeline.od_config.dtype = torch.float32
+    pipeline.bagel.parallel_config = _ParallelConfig(cfg_parallel_size=3)
+    pipeline.bagel.predict_noise_with_multi_branch_cfg.return_value = expected
+    pipeline._pack_step_generation_inputs = BagelPipeline._pack_step_generation_inputs
+    pipeline._build_denoise_kwargs = types.MethodType(BagelPipeline._build_denoise_kwargs, pipeline)
+    pipeline._denoise_step_cfg_parallel = types.MethodType(BagelPipeline._denoise_step_cfg_parallel, pipeline)
+    pipeline.denoise_step = types.MethodType(BagelPipeline.denoise_step, pipeline)
+    cfg_group = MagicMock()
+
+    with (
+        patch(
+            "vllm_omni.diffusion.models.bagel.pipeline_bagel.get_classifier_free_guidance_world_size",
+            return_value=3,
+        ),
+        patch("vllm_omni.diffusion.models.bagel.pipeline_bagel.get_cfg_group", return_value=cfg_group),
+    ):
+        actual = pipeline.denoise_step(input_batch, states=states)
+
+    assert actual is expected
+    call = pipeline.bagel.predict_noise_with_multi_branch_cfg.call_args.kwargs
+    assert len(call["branches_kwargs"]) == 2
+
+
+def test_denoise_step_rejects_three_cfg_branches_on_two_ranks():
+    states = [
+        _prepare_state("a", kv_len=2, cfg_text_scale=4.0),
+        _prepare_state("b", kv_len=3, cfg_text_scale=4.0),
+    ]
+    input_batch = InputBatch.make_batch(states)
+
+    pipeline = MagicMock()
+    pipeline.device = torch.device("cpu")
+    pipeline.od_config.dtype = torch.float32
+    pipeline.bagel.parallel_config = _ParallelConfig(cfg_parallel_size=2)
+    pipeline._pack_step_generation_inputs = BagelPipeline._pack_step_generation_inputs
+    pipeline._build_denoise_kwargs = types.MethodType(BagelPipeline._build_denoise_kwargs, pipeline)
+    pipeline._denoise_step_cfg_parallel = types.MethodType(BagelPipeline._denoise_step_cfg_parallel, pipeline)
+    pipeline.denoise_step = types.MethodType(BagelPipeline.denoise_step, pipeline)
+
+    with (
+        patch(
+            "vllm_omni.diffusion.models.bagel.pipeline_bagel.get_classifier_free_guidance_world_size",
+            return_value=2,
+        ),
+        patch("vllm_omni.diffusion.models.bagel.pipeline_bagel.get_cfg_group", return_value=MagicMock()),
+        pytest.raises(ValueError, match="requires cfg_parallel_size=3"),
+    ):
+        pipeline.denoise_step(input_batch, states=states)
+
+
+@pytest.mark.parametrize("cfg_text_scale", [1.0, 4.0], ids=["no_cfg", "text_cfg"])
+def test_step_scheduler_matches_bagel_euler_update(cfg_text_scale: float):
+    state = _prepare_state("req", kv_len=2, cfg_text_scale=cfg_text_scale)
+    initial = state.latents.clone()
+    velocity = torch.full_like(initial, 0.25)
+    expected_trajectory = [initial.clone()]
+
+    pipeline = MagicMock()
+    pipeline.step_scheduler = types.MethodType(BagelPipeline.step_scheduler, pipeline)
+    for step, dt in enumerate(state.extra["bagel_dts"]):
+        pipeline.step_scheduler(state, velocity)
+        expected = initial - velocity * state.extra["bagel_dts"][: step + 1].sum()
+        expected_trajectory.append(expected)
+        assert torch.allclose(state.latents, expected)
+        assert state.step_index == step + 1
+
+    assert state.denoise_completed
+    assert len(state.extra["bagel_trajectory_latents"]) == state.total_steps + 1
+    assert len(state.extra["bagel_trajectory_timesteps"]) == state.total_steps
+    for actual, expected in zip(state.extra["bagel_trajectory_latents"], expected_trajectory, strict=True):
+        assert torch.allclose(actual, expected)
+
+
+@pytest.mark.parametrize("cfg_text_scale", [1.0, 4.0], ids=["no_cfg", "text_cfg"])
+@pytest.mark.parametrize("velocity_dtype", [torch.float32, torch.bfloat16], ids=["fp32_velocity", "bf16_velocity"])
+def test_complete_step_updates_match_full_bagel_loop(cfg_text_scale: float, velocity_dtype: torch.dtype):
+    state = _prepare_state("req", kv_len=2, cfg_text_scale=cfg_text_scale)
+    initial = state.latents.clone()
+
+    bagel = MagicMock()
+    bagel._sp_size = 1
+    bagel._denoise_schedule_extra_step = True
+    bagel.prepare_denoise_schedule = types.MethodType(Bagel.prepare_denoise_schedule, bagel)
+    bagel.forward.side_effect = lambda x_t, **_kwargs: torch.full(
+        x_t.shape,
+        0.25,
+        dtype=velocity_dtype,
+        device=x_t.device,
+    )
+    bagel.generate_image = types.MethodType(Bagel.generate_image, bagel)
+    generation_input = dict(state.extra["bagel_generation_input"])
+    generation_input["packed_init_noises"] = initial.clone()
+
+    with patch(
+        "vllm_omni.diffusion.models.bagel.bagel_transformer.get_classifier_free_guidance_world_size",
+        return_value=1,
+    ):
+        full_latents, full_trajectory, full_timesteps, _ = bagel.generate_image(
+            **generation_input,
+            past_key_values=state.extra["bagel_gen_context"]["past_key_values"],
+            num_timesteps=state.extra["bagel_gen_params"].num_timesteps,
+            timestep_shift=state.extra["bagel_gen_params"].timestep_shift,
+            cfg_text_scale=cfg_text_scale,
+            cfg_img_scale=1.0,
+            cfg_text_packed_position_ids=state.extra["bagel_cfg_text_packed_position_ids"],
+            cfg_text_past_key_values=state.extra["bagel_cfg_text_context"]["past_key_values"],
+            return_trajectory_latents=True,
+        )
+
+    pipeline = MagicMock()
+    pipeline.step_scheduler = types.MethodType(BagelPipeline.step_scheduler, pipeline)
+    velocity = torch.full(initial.shape, 0.25, dtype=velocity_dtype, device=initial.device)
+    while not state.denoise_completed:
+        pipeline.step_scheduler(state, velocity)
+
+    assert torch.allclose(state.latents, full_latents[0])
+    assert state.latents.dtype == full_latents[0].dtype == torch.float32
+    assert len(state.extra["bagel_trajectory_latents"]) == len(full_trajectory)
+    assert len(state.extra["bagel_trajectory_timesteps"]) == len(full_timesteps)
+    for actual, expected in zip(state.extra["bagel_trajectory_latents"], full_trajectory, strict=True):
+        assert torch.allclose(actual, expected)
+
+
+@dataclass
+class _SchedulerOutput:
+    prev_sample: torch.Tensor
+    log_prob: torch.Tensor
+
+
+class _Scheduler:
+    def step(self, velocity, _timestep, latents, dt, **_kwargs):
+        return _SchedulerOutput(latents - velocity * dt, torch.tensor(-0.5))
+
+
+def test_step_scheduler_records_log_probability():
+    state = _prepare_state("req", kv_len=2)
+    state.scheduler = _Scheduler()
+    pipeline = MagicMock()
+    pipeline.step_scheduler = types.MethodType(BagelPipeline.step_scheduler, pipeline)
+
+    pipeline.step_scheduler(state, torch.full_like(state.latents, 0.25))
+
+    assert len(state.extra["bagel_trajectory_log_probs"]) == 1
+    assert state.extra["bagel_trajectory_log_probs"][0].item() == -0.5
+
+
+def test_post_decode_reuses_full_path_output_formatter():
+    state = _prepare_state("req", kv_len=2)
+    expected = object()
+    pipeline = MagicMock()
+    pipeline._build_image_output.return_value = expected
+    pipeline.post_decode = types.MethodType(BagelPipeline.post_decode, pipeline)
+
+    result = pipeline.post_decode(state)
+
+    assert result is expected
+    call = pipeline._build_image_output.call_args
+    assert call.args == (state.latents, (64, 64))
+    assert call.kwargs["trajectory_latents"] is state.extra["bagel_trajectory_latents"]
+    assert call.kwargs["trajectory_timesteps"] is state.extra["bagel_trajectory_timesteps"]
+    assert call.kwargs["think_text"] is None
+
+
+def test_nested_cache_merge_keeps_per_request_branch_lengths():
+    gen = NaiveCache.merge([_cache(2, 2.0), _cache(3, 3.0)])
+    cfg = NaiveCache.merge([_cache(0, 0.0), _cache(4, 4.0)])
+
+    merged = NaiveCache.merge([gen, cfg])
+
+    assert merged.key_values_lens == [2, 3, 0, 4]
+    assert merged.key_cache[0].shape[0] == 9
+    assert torch.equal(merged.key_cache[0][:2], gen.key_cache[0][:2])
+    assert torch.equal(merged.key_cache[0][-4:], cfg.key_cache[0])
+
+
+class _PositionAddingLanguageModel(nn.Module):
+    def __init__(self, hidden_size: int):
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.denoise_calls = 0
+
+    def forward(
+        self,
+        *,
+        return_embeddings_only: bool = False,
+        packed_text_ids: torch.Tensor | None = None,
+        packed_query_sequence: torch.Tensor | None = None,
+        packed_query_position_ids: torch.Tensor | None = None,
+        **_kwargs: object,
+    ) -> _LanguageModelOutput:
+        if return_embeddings_only:
+            assert packed_text_ids is not None
+            count = packed_text_ids.numel()
+            return _LanguageModelOutput(packed_query_sequence=torch.zeros(count, self.hidden_size))
+        self.denoise_calls += 1
+        assert packed_query_sequence is not None
+        assert packed_query_position_ids is not None
+        positions = packed_query_position_ids.to(packed_query_sequence).unsqueeze(-1)
+        return _LanguageModelOutput(packed_query_sequence=packed_query_sequence + positions)
+
+
+class _ZeroEmbedding(nn.Module):
+    def __init__(self, hidden_size: int):
+        super().__init__()
+        self.hidden_size = hidden_size
+
+    def forward(self, values: torch.Tensor) -> torch.Tensor:
+        return torch.zeros(values.shape[0], self.hidden_size, device=values.device)
+
+
+@pytest.mark.parametrize(
+    "cfg_kwargs",
+    [
+        {"cfg_text_scale": 3.0},
+        {
+            "cfg_text_scale": 1.0,
+            "cfg_vae_lengths": [2],
+            "cfg_text_scales": [3.0],
+            "cfg_img_scales": [1.0],
+        },
+    ],
+    ids=["scalar-scale", "per-request-scales"],
+)
+def test_bagel_forward_skips_cfg_without_branch_contexts(cfg_kwargs):
+    bagel = Bagel.__new__(Bagel)
+    nn.Module.__init__(bagel)
+    bagel.hidden_size = 2
+    bagel.use_moe = False
+    bagel.language_model = _PositionAddingLanguageModel(hidden_size=2)
+    bagel.latent_pos_embed = _ZeroEmbedding(hidden_size=2)
+    bagel.time_embedder = _ZeroEmbedding(hidden_size=2)
+    bagel.vae2llm = nn.Identity()
+    bagel.llm2vae = nn.Identity()
+
+    result = bagel.forward(
+        x_t=torch.zeros(2, 2),
+        timestep=torch.zeros(2),
+        packed_vae_token_indexes=torch.tensor([1, 2]),
+        packed_vae_position_ids=torch.zeros(2, dtype=torch.long),
+        packed_text_ids=torch.zeros(2, dtype=torch.long),
+        packed_text_indexes=torch.tensor([0, 3]),
+        packed_position_ids=torch.ones(4, dtype=torch.long),
+        packed_seqlens=torch.tensor([4]),
+        past_key_values=NaiveCache(1),
+        cfg_renorm_min=1.0,
+        **cfg_kwargs,
+    )
+
+    assert bagel.language_model.denoise_calls == 1
+    assert torch.equal(result, torch.ones(2, 2))
+
+
+def test_bagel_forward_combines_same_cfg_per_request_in_one_model_call():
+    bagel = Bagel.__new__(Bagel)
+    nn.Module.__init__(bagel)
+    bagel.hidden_size = 2
+    bagel.use_moe = False
+    bagel.language_model = _PositionAddingLanguageModel(hidden_size=2)
+    bagel.latent_pos_embed = _ZeroEmbedding(hidden_size=2)
+    bagel.time_embedder = _ZeroEmbedding(hidden_size=2)
+    bagel.vae2llm = nn.Identity()
+    bagel.llm2vae = nn.Identity()
+
+    result = bagel.forward(
+        x_t=torch.zeros(4, 2),
+        timestep=torch.zeros(4),
+        packed_vae_token_indexes=torch.tensor([1, 2, 5, 6]),
+        packed_vae_position_ids=torch.zeros(4, dtype=torch.long),
+        packed_text_ids=torch.zeros(4, dtype=torch.long),
+        packed_text_indexes=torch.tensor([0, 3, 4, 7]),
+        packed_position_ids=torch.tensor([1, 1, 1, 1, 2, 2, 2, 2]),
+        packed_seqlens=torch.tensor([4, 4]),
+        past_key_values=NaiveCache(1),
+        cfg_renorm_min=1.0,
+        cfg_text_scale=3.0,
+        cfg_img_scale=1.0,
+        cfg_branch_pids=[
+            torch.tensor([1, 1, 1, 1, 2, 2, 2, 2]),
+            torch.zeros(8, dtype=torch.long),
+        ],
+        cfg_branch_caches=[NaiveCache(1), NaiveCache(1)],
+        cfg_vae_lengths=[2, 2],
+        cfg_text_scales=[3.0, 3.0],
+        cfg_img_scales=[1.0, 1.0],
+    )
+
+    assert bagel.language_model.denoise_calls == 1
+    assert torch.equal(result[:2], torch.full((2, 2), 3.0))
+    assert torch.equal(result[2:], torch.full((2, 2), 6.0))
+
+
+def test_multi_branch_cfg_combines_same_cfg_for_each_request_independently():
+    bagel = Bagel.__new__(Bagel)
+    nn.Module.__init__(bagel)
+    positive = torch.tensor([[1.0], [1.0], [2.0], [2.0]])
+    negative = torch.zeros_like(positive)
+
+    result = bagel.combine_multi_branch_cfg_noise(
+        [positive, negative],
+        {
+            "cfg_text_scale": 3.0,
+            "cfg_img_scale": 1.0,
+            "cfg_renorm_type": "global",
+            "cfg_renorm_min": 1.0,
+            "cfg_vae_lengths": [2, 2],
+            "cfg_text_scales": [3.0, 3.0],
+            "cfg_img_scales": [1.0, 1.0],
+        },
+    )
+
+    assert torch.equal(result[:2], torch.full((2, 1), 3.0))
+    assert torch.equal(result[2:], torch.full((2, 1), 6.0))

tests/diffusion/models/bagel/test_trajectory_recording.py

@@ -19,18 +19,14 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 NUM_TOKENS = 8
 HIDDEN_DIM = 16
 NUM_TIMESTEPS = 5
-# Official BAGEL samples num_timesteps points then drops the terminal t=0,
-# yielding num_timesteps - 1 denoise steps.
-EXPECTED_STEPS = NUM_TIMESTEPS - 1
+EXPECTED_STEPS = NUM_TIMESTEPS
 
 
 def _make_mock_bagel(mocker: MockerFixture):
     """Create a mock Bagel with forward returning constant velocity."""
     mock = mocker.MagicMock(spec=Bagel)
     mock._sp_size = 1
-    # 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
+    mock._denoise_schedule_extra_step = True
 
     # forward returns a small constant velocity so x_t changes each step
     def fake_forward(self, x_t, **kwargs):
@@ -182,7 +178,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"])
+        ts = torch.linspace(1, 0, args["num_timesteps"] + 1)
         shift = args.get("timestep_shift", 1.0)
         ts = shift * ts / (1 + (shift - 1) * ts)
         expected_timesteps = ts[:-1]  # last element is dropped
@@ -200,7 +196,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"])
+        ts = torch.linspace(1, 0, args["num_timesteps"] + 1)
         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

@@ -16,6 +16,43 @@ 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)])
+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(
+        "Config",
+        (),
+        {"model_class_name": "mock_model", "diffusion_load_format": "default"},
+    )()
+    engine.step_execution = step_execution
+    engine.pre_process_func = None
+    captured_requests = []
+
+    monkeypatch.setattr(
+        "vllm_omni.diffusion.diffusion_engine.supports_multimodal_input",
+        lambda od_config: (False, False),
+    )
+    monkeypatch.setattr(
+        "vllm_omni.diffusion.diffusion_engine.get_dummy_run_num_frames",
+        lambda model_class_name, supports_audio_input: 1,
+    )
+
+    def run_request(request):
+        captured_requests.append(request)
+        return type("Output", (), {"error": None})()
+
+    engine.add_req_and_wait_for_response = run_request
+
+    engine._dummy_run()
+
+    assert len(captured_requests) == 1
+    assert captured_requests[0].sampling_params.num_inference_steps == expected_steps
+
+
 def test_dummy_run_num_frames_uses_explicit_model_setting(monkeypatch: pytest.MonkeyPatch) -> None:
     class JointAudioVideoModel:
         dummy_run_num_frames = 2

tests/diffusion/test_diffusion_model_runner.py

@@ -11,6 +11,7 @@ import torch
 import vllm_omni.diffusion.worker.diffusion_model_runner as model_runner_module
 from tests.helpers.mark import hardware_test
 from vllm_omni.diffusion.data import DiffusionOutput
+from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
 from vllm_omni.diffusion.diffusion_kv.model_runner_backend import DiffusionKVModelRunnerBackend
 from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner
 from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch, split_diffusion_output_by_request
@@ -297,6 +298,106 @@ def test_refresh_cache_without_request_or_pipeline_default_warns(caplog):
     assert "requires num_inference_steps to be passed explicitly" in caplog.text
 
 
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_execute_stepwise_falls_back_to_full_forward_without_creating_step_state(monkeypatch):
+    runner = _make_runner(cache_backend=None, cache_backend_name="none")
+    runner.pipeline = _SingleRequestDiffusionOutputPipeline()
+    request = _make_request()
+    request.request_id = "text-request"
+    request.use_step_execution = False
+    scheduler_output = SimpleNamespace(
+        finished_req_ids=set(),
+        scheduled_new_reqs=[
+            SimpleNamespace(
+                request_id=request.request_id,
+                req=request,
+                diffusion_kv_metadata=None,
+            )
+        ],
+        scheduled_cached_reqs=SimpleNamespace(request_ids=[]),
+    )
+    monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
+    monkeypatch.setattr(model_runner_module, "current_omni_platform", _fake_platform_for_peak_memory())
+
+    result = DiffusionModelRunner.execute_stepwise(runner, scheduler_output)
+    output = result.get_request_output(request.request_id)
+
+    assert output.finished is True
+    assert output.step_index == 0
+    assert output.result is not None
+    assert output.result.output == request.prompt
+    assert isinstance(runner.pipeline.last_req, DiffusionRequestBatch)
+    assert runner.pipeline.last_req.num_reqs == 1
+    assert runner.state_cache == {}
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_non_step_fallback_cleans_finished_step_wave_state_and_paged_kv(monkeypatch):
+    """A finished step request must not survive into a text fallback wave."""
+    runner = _make_runner(cache_backend=None, cache_backend_name="none")
+    runner.pipeline = _SingleRequestDiffusionOutputPipeline()
+    runner.state_cache = {"aborted-step-request": object()}
+    old_step_batch = object()
+    runner.input_batch = old_step_batch
+    runner.od_config.diffusion_kv_mode = DiffusionKVCacheMode.PAGED_SCHEDULER
+    runner.remove_diffusion_kv_requests = Mock()
+
+    text_request = _make_request()
+    text_request.request_id = "text-fallback-request"
+    text_request.use_step_execution = False
+    scheduler_output = SimpleNamespace(
+        finished_req_ids={"aborted-step-request"},
+        scheduled_new_reqs=[
+            SimpleNamespace(
+                request_id=text_request.request_id,
+                req=text_request,
+                diffusion_kv_metadata=None,
+            )
+        ],
+        scheduled_cached_reqs=SimpleNamespace(request_ids=[]),
+    )
+    monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
+    monkeypatch.setattr(model_runner_module, "current_omni_platform", _fake_platform_for_peak_memory())
+
+    result = DiffusionModelRunner.execute_stepwise(runner, scheduler_output)
+
+    assert result.get_request_output(text_request.request_id).finished is True
+    assert runner.state_cache == {}
+    runner.remove_diffusion_kv_requests.assert_called_once_with(["aborted-step-request"])
+    assert runner.input_batch is None
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_execute_stepwise_rejects_mixed_step_and_full_forward_batch():
+    runner = _make_runner(cache_backend=None, cache_backend_name="none")
+    non_step_request = _make_request()
+    non_step_request.use_step_execution = False
+    step_request = _make_request()
+    step_request.request_id = "step-request"
+    scheduler_output = SimpleNamespace(
+        finished_req_ids=set(),
+        scheduled_new_reqs=[
+            SimpleNamespace(
+                request_id=non_step_request.request_id,
+                req=non_step_request,
+                diffusion_kv_metadata=None,
+            ),
+            SimpleNamespace(
+                request_id=step_request.request_id,
+                req=step_request,
+                diffusion_kv_metadata=None,
+            ),
+        ],
+        scheduled_cached_reqs=SimpleNamespace(request_ids=[]),
+    )
+
+    with pytest.raises(ValueError, match="Cannot mix stepwise and non-step fallback"):
+        DiffusionModelRunner.execute_stepwise(runner, scheduler_output)
+
+
 @pytest.mark.core_model
 @pytest.mark.cpu
 def test_update_states_carries_prepared_layout() -> None:

tests/diffusion/test_diffusion_request.py

@@ -2,10 +2,12 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 import random
+from types import SimpleNamespace
 
 import pytest
 
 from vllm_omni.diffusion.request import DUMMY_DIFFUSION_REQUEST_ID, OmniDiffusionRequest
+from vllm_omni.diffusion.sched.step_scheduler import StepScheduler
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 
 pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
@@ -63,6 +65,23 @@ def test_non_dummy_request_is_not_identified_as_dummy_run():
     assert not OmniDiffusionRequest.is_dummy_run_request_id(None)
 
 
+def test_step_scheduler_separates_full_forward_fallback_from_step_requests():
+    scheduler = StepScheduler()
+    scheduler.initialize(SimpleNamespace())
+    scheduler.max_num_running_reqs = 2
+    text_request = _make_request()
+    text_request.use_step_execution = False
+    image_request = _make_request()
+    image_request.request_id = "image-request"
+
+    scheduler.add_request(text_request)
+    scheduler.add_request(image_request)
+    output = scheduler.schedule()
+
+    assert output.scheduled_request_ids == [text_request.request_id]
+    assert output.num_waiting_reqs == 1
+
+
 def test_tp_seed_same_across_ranks_and_varies_across_requests():
     random.seed(0)
     n_requests = 5

tests/diffusion/test_diffusion_scheduler.py

@@ -255,6 +255,15 @@ class TestGetStepBatchSamplingParamsKey:
         b = scheduler._build_sampling_params_key(self._make(lora_int_id=1, lora_scale=0.5))
         assert a == b
 
+    def test_distinguishes_pipeline_condition_structure(self) -> None:
+        scheduler = _ConcreteScheduler()
+        a = self._make()
+        b = self._make()
+        a.batch_compatibility_key = ("bagel_cfg", 1.0)
+        b.batch_compatibility_key = ("bagel_cfg", 4.0)
+
+        assert scheduler._build_sampling_params_key(a) != scheduler._build_sampling_params_key(b)
+
 
 class TestGetRequestBatchSamplingParamsKey:
     """Tests for the request-batch compatibility key builder on RequestScheduler."""
@@ -1559,6 +1568,22 @@ class TestStepScheduler:
         assert sched_output.num_running_reqs == 2
         assert sched_output.num_waiting_reqs == 0
 
+    def test_batches_incompatible_pipeline_conditions_separately(self) -> None:
+        scheduler = StepScheduler()
+        scheduler.initialize(SimpleNamespace(max_num_seqs=2))
+        request_a = _make_step_request("a")
+        request_b = _make_step_request("b")
+        request_a.batch_compatibility_key = ("bagel_cfg", 1.0)
+        request_b.batch_compatibility_key = ("bagel_cfg", 4.0)
+
+        req_a = scheduler.add_request(request_a)
+        scheduler.add_request(request_b)
+        sched_output = scheduler.schedule()
+
+        assert _new_ids(sched_output) == [req_a]
+        assert sched_output.num_running_reqs == 1
+        assert sched_output.num_waiting_reqs == 1
+
     def test_step_batch_allows_different_num_inference_steps(self) -> None:
         scheduler = StepScheduler()
         scheduler.initialize(SimpleNamespace(max_num_seqs=2))

vllm_omni/config/omni_config.py

@@ -320,13 +320,23 @@ def _resolve_scheduler_path(execution_type: StageExecutionType, async_scheduling
     return _scheduler_path(_resolve_scheduler(execution_type, async_scheduling))
 
 
-def _stage_cli_overrides(stage_id: int, cli_overrides: Mapping[str, Any]) -> dict[str, Any]:
+def _stage_cli_overrides(
+    stage_id: int,
+    cli_overrides: Mapping[str, Any],
+    *,
+    execution_type: StageExecutionType | None = None,
+) -> dict[str, Any]:
     runtime_overrides = build_stage_runtime_overrides(stage_id, dict(cli_overrides))
     global_stage_fields = _global_stage_cli_fields()
     result: dict[str, Any] = {}
     for key, value in runtime_overrides.items():
         if key in global_stage_fields or f"stage_{stage_id}_{key}" in cli_overrides:
             result[key] = _copy_value(value)
+
+    # step_execution is a diffusion execution protocol, not an LLM engine
+    # argument. Keep global and stage-scoped CLI values off AR/generation stages.
+    if execution_type is not None and execution_type is not StageExecutionType.DIFFUSION:
+        result.pop("step_execution", None)
     return result
 
 
@@ -1878,7 +1888,11 @@ class VllmOmniConfig:
                 _stage_engine_values(
                     deploy_by_id.get(topology.stage_id),
                     topology,
-                    _stage_cli_overrides(topology.stage_id, cli_overrides),
+                    _stage_cli_overrides(
+                        topology.stage_id,
+                        cli_overrides,
+                        execution_type=topology.execution_type,
+                    ),
                 ),
                 model=model,
             )

vllm_omni/diffusion/diffusion_engine.py

@@ -1047,9 +1047,9 @@ class DiffusionEngine:
         width: int,
         guidance_scale: float,
         num_image_inputs: int = 1,
+        num_inference_steps: int = 1,
     ) -> OmniDiffusionRequest | None:
-        """Build a one-step model request for startup profiling or warmup."""
-
+        """Build a minimal model request for startup profiling or warmup."""
         prompt: OmniTextPrompt = {"prompt": "dummy run"}
         supports_image_input, supports_audio_input = supports_multimodal_input(self.od_config)
         if supports_image_input:
@@ -1070,7 +1070,7 @@ class DiffusionEngine:
             sampling_params=OmniDiffusionSamplingParams(
                 height=height,
                 width=width,
-                num_inference_steps=1,
+                num_inference_steps=num_inference_steps,
                 num_frames=num_frames,
                 guidance_scale=guidance_scale,
                 num_outputs_per_prompt=1,
@@ -1141,6 +1141,7 @@ class DiffusionEngine:
             height=512,
             width=512,
             guidance_scale=0.0,
+            num_inference_steps=1,
         )
         if req is None:
             logger.info("Skipping dummy warmup run (num_frames=0)")

vllm_omni/diffusion/layers/mot/mot_layernorm.py

@@ -82,6 +82,17 @@ class MoTRMSNorm(CustomOp):
             eps=self.variance_epsilon,
         )
 
+    # ------------------------------------------------------------------
+    # NPU fast-path (not implemented yet; reserved for future optimization)
+    # ------------------------------------------------------------------
+    def forward_npu(
+        self,
+        x: torch.Tensor,
+        text_indices: torch.Tensor | None = None,
+        vae_indices: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        return self.forward_native(x, text_indices, vae_indices)
+
     # ------------------------------------------------------------------
     # Helpers
     # ------------------------------------------------------------------

vllm_omni/diffusion/models/bagel/bagel_transformer.py

@@ -9,6 +9,7 @@
 
 from collections.abc import Iterable, Sequence
 from dataclasses import dataclass
+from typing import Any
 
 import numpy as np
 import torch
@@ -364,15 +365,19 @@ class NaiveCache:
 
     @staticmethod
     def merge(caches: Sequence["NaiveCache"]) -> "NaiveCache":
-        """Merge per-branch NaiveCaches into one for batched attention;
-        this lets us do the forward passes for CFG in one batched pass,
-        although it's worth noting that this is currently only used
-        for single request. We need this so that gen mode knows the
-        respective kv lengths, and can split things back out as needed.
+        """Merge NaiveCaches into one for packed request/CFG attention.
+
+        Already-merged inputs keep their per-sequence lengths so nested
+        request-by-CFG packing can split cache rows back out correctly.
         """
         num_layers = caches[0].num_layers
         merged = NaiveCache(num_layers)
-        lens = [c.seq_lens for c in caches]
+        lens = []
+        for cache in caches:
+            if cache.key_values_lens is None:
+                lens.append(cache.seq_lens)
+            else:
+                lens.extend(cache.key_values_lens)
         merged.key_values_lens = lens
 
         nonempty = [c for c in caches if c.key_cache[0] is not None]
@@ -1251,12 +1256,22 @@ class Bagel(CFGParallelMixin, nn.Module):
     config_class = BagelConfig
     base_model_prefix = "bagel"
 
-    # 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
+    # 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
+
+    def prepare_denoise_schedule(
+        self,
+        x_t: torch.Tensor,
+        num_timesteps: int,
+        timestep_shift: float,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Build BAGEL's flow-matching timesteps and per-step deltas."""
+        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)
+        return schedule[:-1], schedule[:-1] - schedule[1:]
 
     def __init__(
         self,
@@ -1819,14 +1834,8 @@ 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 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.
-        num_sample_points = num_timesteps + 1 if self._denoise_schedule_extra_step else num_timesteps
-        timesteps = torch.linspace(1, 0, num_sample_points, device=x_t.device)
-        timesteps = timestep_shift * timesteps / (1 + (timestep_shift - 1) * timesteps)
-        dts = timesteps[:-1] - timesteps[1:]
-        timesteps = timesteps[:-1]
+        # Build ``num_timesteps`` flow-matching Euler updates over [1, 0].
+        timesteps, dts = Bagel.prepare_denoise_schedule(self, x_t, num_timesteps, timestep_shift)
 
         # Optional trajectory recording for RL rollout data collection
         trajectory_latents: list[torch.Tensor] | None = [] if return_trajectory_latents else None
@@ -2228,6 +2237,52 @@ class Bagel(CFGParallelMixin, nn.Module):
 
         return v_t
 
+    def _combine_cfg_per_request(
+        self,
+        v_t: torch.Tensor,
+        cfg_text_v_t: torch.Tensor,
+        cfg_img_v_t: torch.Tensor | None,
+        cfg_vae_lengths: Sequence[int],
+        cfg_text_scales: Sequence[float],
+        cfg_img_scales: Sequence[float],
+        cfg_renorm_type: str,
+        cfg_renorm_min: float,
+    ) -> torch.Tensor:
+        """Apply BAGEL CFG independently to each packed request."""
+        if len(cfg_text_scales) != len(cfg_vae_lengths):
+            raise ValueError("cfg_text_scales must match cfg_vae_lengths.")
+        if len(cfg_img_scales) != len(cfg_vae_lengths):
+            raise ValueError("cfg_img_scales must match cfg_vae_lengths.")
+
+        gen_parts = v_t.split(list(cfg_vae_lengths))
+        text_parts = cfg_text_v_t.split(list(cfg_vae_lengths))
+        img_parts = (
+            cfg_img_v_t.split(list(cfg_vae_lengths)) if cfg_img_v_t is not None else [None] * len(cfg_vae_lengths)
+        )
+        return torch.cat(
+            [
+                gen_part
+                if text_scale <= 1.0
+                else self._combine_cfg(
+                    gen_part,
+                    text_part,
+                    img_part,
+                    text_scale,
+                    img_scale,
+                    cfg_renorm_type,
+                    cfg_renorm_min,
+                )
+                for gen_part, text_part, img_part, text_scale, img_scale in zip(
+                    gen_parts,
+                    text_parts,
+                    img_parts,
+                    cfg_text_scales,
+                    cfg_img_scales,
+                    strict=True,
+                )
+            ]
+        )
+
     # ── CFGParallelMixin hooks ──
     # Bagel mounts CFGParallelMixin (see class declaration) to reuse the shared
     # N-branch CFG dispatch/all_gather logic in predict_noise_with_multi_branch_cfg.
@@ -2247,7 +2302,7 @@ class Bagel(CFGParallelMixin, nn.Module):
     def combine_multi_branch_cfg_noise(
         self,
         predictions: list[torch.Tensor],
-        true_cfg_scale: dict[str, float],
+        true_cfg_scale: dict[str, Any],
         cfg_normalize: bool = False,
     ) -> torch.Tensor:
         """Combine gen/text/img branch velocities via Bagel's renorm CFG.
@@ -2257,6 +2312,18 @@ class Bagel(CFGParallelMixin, nn.Module):
         because renormalization is folded into ``_combine_cfg`` itself.
         """
         cfg_img_v_t = predictions[2] if len(predictions) > 2 else None
+        cfg_vae_lengths = true_cfg_scale.get("cfg_vae_lengths")
+        if cfg_vae_lengths is not None:
+            return self._combine_cfg_per_request(
+                predictions[0],
+                predictions[1],
+                cfg_img_v_t,
+                cfg_vae_lengths,
+                true_cfg_scale["cfg_text_scales"],
+                true_cfg_scale["cfg_img_scales"],
+                true_cfg_scale["cfg_renorm_type"],
+                true_cfg_scale["cfg_renorm_min"],
+            )
         return self._combine_cfg(
             predictions[0],
             predictions[1],
@@ -2270,7 +2337,7 @@ class Bagel(CFGParallelMixin, nn.Module):
     def forward_single_branch(
         self,
         x_t: torch.Tensor,
-        timestep: torch.LongTensor,
+        timestep: torch.Tensor,
         packed_vae_token_indexes: torch.LongTensor,
         packed_vae_position_ids: torch.LongTensor,
         packed_text_ids: torch.LongTensor,
@@ -2379,7 +2446,7 @@ class Bagel(CFGParallelMixin, nn.Module):
     def forward(
         self,
         x_t: torch.Tensor,
-        timestep: torch.LongTensor,
+        timestep: torch.Tensor,
         packed_vae_token_indexes: torch.LongTensor,
         packed_vae_position_ids: torch.LongTensor,
         packed_text_ids: torch.LongTensor,
@@ -2393,6 +2460,9 @@ class Bagel(CFGParallelMixin, nn.Module):
         cfg_img_scale: float = 1.0,
         cfg_branch_pids: list[torch.Tensor] | None = None,
         cfg_branch_caches: list[NaiveCache] | None = None,
+        cfg_vae_lengths: list[int] | None = None,
+        cfg_text_scales: list[float] | None = None,
+        cfg_img_scales: list[float] | None = None,
     ):
         # Build query sequence (identical for all CFG branches)
         packed_text_embedding = self.language_model.forward(
@@ -2416,11 +2486,16 @@ class Bagel(CFGParallelMixin, nn.Module):
             extra_inputs["packed_vae_token_indexes"] = packed_vae_token_indexes
             extra_inputs["packed_text_indexes"] = packed_text_indexes
 
-        use_cfg = cfg_text_scale > 1.0
+        has_cfg_branches = cfg_branch_pids is not None and cfg_branch_caches is not None
+        use_cfg = has_cfg_branches and (
+            cfg_text_scale > 1.0 or (cfg_text_scales is not None and any(scale > 1.0 for scale in cfg_text_scales))
+        )
         cfg_text_v_t = None
         cfg_img_v_t = None
 
-        if use_cfg and cfg_branch_pids is not None and cfg_branch_caches is not None:
+        if use_cfg:
+            assert cfg_branch_pids is not None
+            assert cfg_branch_caches is not None
             num_branches = len(cfg_branch_pids)
             seq_len = int(packed_seqlens.sum())
 
@@ -2466,14 +2541,30 @@ class Bagel(CFGParallelMixin, nn.Module):
 
         # ── CFG combination ──
         if use_cfg:
-            v_t = self._combine_cfg(
-                v_t,
-                cfg_text_v_t,
-                cfg_img_v_t,
-                cfg_text_scale,
-                cfg_img_scale,
-                cfg_renorm_type,
-                cfg_renorm_min,
-            )
+            if cfg_vae_lengths is None:
+                v_t = self._combine_cfg(
+                    v_t,
+                    cfg_text_v_t,
+                    cfg_img_v_t,
+                    cfg_text_scale,
+                    cfg_img_scale,
+                    cfg_renorm_type,
+                    cfg_renorm_min,
+                )
+            else:
+                if cfg_text_scales is None:
+                    raise ValueError("cfg_text_scales must be provided with cfg_vae_lengths.")
+                if cfg_img_scales is None:
+                    cfg_img_scales = [cfg_img_scale] * len(cfg_vae_lengths)
+                v_t = self._combine_cfg_per_request(
+                    v_t,
+                    cfg_text_v_t,
+                    cfg_img_v_t,
+                    cfg_vae_lengths,
+                    cfg_text_scales,
+                    cfg_img_scales,
+                    cfg_renorm_type,
+                    cfg_renorm_min,
+                )
 
         return v_t

vllm_omni/diffusion/models/bagel/pipeline_bagel.py

@@ -8,11 +8,11 @@ from __future__ import annotations
 
 import json
 import os
-from collections.abc import Iterable
-from copy import deepcopy
+from collections.abc import Iterable, Sequence
+from copy import copy, deepcopy
 from dataclasses import dataclass
 from math import isqrt
-from typing import ClassVar
+from typing import TYPE_CHECKING, Any, ClassVar
 
 import numpy as np
 import torch
@@ -25,10 +25,15 @@ from vllm.model_executor.models.utils import AutoWeightsLoader
 from vllm.transformers_utils.configs.bagel import BagelConfig
 
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
+from vllm_omni.diffusion.distributed.parallel_state import (
+    get_cfg_group,
+    get_classifier_free_guidance_world_size,
+)
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
 from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery
 from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
+from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
 from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific
 
@@ -37,6 +42,11 @@ from .bagel_transformer import Bagel, NaiveCache, Qwen2MoTConfig, Qwen2MoTForCau
 
 logger = init_logger(__name__)
 
+if TYPE_CHECKING:
+    from vllm_omni.diffusion.worker.input_batch import InputBatch
+    from vllm_omni.diffusion.worker.utils import StepRequestState
+    from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniPromptType
+
 
 @dataclass
 class BagelGenParams:
@@ -95,6 +105,144 @@ def get_bagel_post_process_func(od_config: OmniDiffusionConfig):
     return post_process_func
 
 
+def _resolve_bagel_image_geometry(od_config: OmniDiffusionConfig) -> tuple[int, int]:
+    """Return the effective ``(latent_downsample, max_latent_size)``.
+
+    Some BAGEL checkpoints advertise a stale ``max_latent_size`` in
+    ``config.json``.  Model loading already corrects it from the positional
+    embedding weight, so admission preprocessing must do the same or it can
+    compute a different img2img shape from the Worker.
+    """
+
+    model = getattr(od_config, "model", None)
+    if not model:
+        raise ValueError("BAGEL img2img preprocessing requires od_config.model.")
+    if os.path.exists(model):
+        model_path = model
+    else:
+        model_path = download_weights_from_hf_specific(
+            model,
+            None,
+            ["*"],
+            revision=getattr(od_config, "revision", None),
+        )
+
+    config_path = os.path.join(model_path, "config.json")
+    with open(config_path, encoding="utf-8") as f:
+        bagel_cfg = json.load(f)
+
+    vae_config = bagel_cfg.get("vae_config") or {}
+    latent_downsample = int(vae_config.get("downsample", 8)) * int(bagel_cfg.get("latent_patch_size", 2))
+    max_latent_size = int(bagel_cfg.get("max_latent_size", 32))
+
+    index_path = os.path.join(model_path, "model.safetensors.index.json")
+    if os.path.exists(index_path):
+        with open(index_path, encoding="utf-8") as f:
+            weight_map = json.load(f).get("weight_map") or {}
+        position_key = next(
+            (key for key in ("latent_pos_embed.pos_embed", "bagel.latent_pos_embed.pos_embed") if key in weight_map),
+            None,
+        )
+        if position_key is not None:
+            from safetensors import safe_open
+
+            shard_path = os.path.join(model_path, weight_map[position_key])
+            with safe_open(shard_path, framework="pt") as f:
+                num_positions = int(f.get_slice(position_key).get_shape()[0])
+            inferred_size = isqrt(num_positions)
+            if inferred_size * inferred_size != num_positions:
+                raise ValueError(f"BAGEL latent position embedding length must be a square, got {num_positions}.")
+            max_latent_size = inferred_size
+
+    return latent_downsample, max_latent_size
+
+
+def _bagel_effective_image_size(
+    image_size: tuple[int, int],
+    *,
+    latent_downsample: int,
+    max_latent_size: int,
+) -> tuple[int, int]:
+    """Match BAGEL's img2img resize and return ``(height, width)``."""
+
+    width, height = image_size
+    if width <= 0 or height <= 0:
+        raise ValueError(f"BAGEL img2img input must have positive dimensions, got {width}x{height}.")
+
+    max_image_size = int(max_latent_size * latent_downsample)
+    scale = min(max_image_size / max(width, height), 1.0)
+    min_image_size = min(256, max_image_size)
+    scale = max(scale, min_image_size / min(width, height))
+    resized_width = max(
+        latent_downsample,
+        int(round(width * scale / latent_downsample) * latent_downsample),
+    )
+    resized_height = max(
+        latent_downsample,
+        int(round(height * scale / latent_downsample) * latent_downsample),
+    )
+    return min(resized_height, max_image_size), min(resized_width, max_image_size)
+
+
+def get_bagel_pre_process_func(od_config: OmniDiffusionConfig):
+    """Resolve BAGEL execution mode and step-batch compatibility."""
+
+    step_execution = bool(getattr(od_config, "step_execution", False))
+    image_geometry: tuple[int, int] | None = None
+
+    def pre_process_func(request: OmniDiffusionRequest):
+        nonlocal image_geometry
+
+        sampling = request.sampling_params
+        kv_metadata = getattr(sampling, "kv_metadata", None) or {}
+        image_shape = kv_metadata.get("image_shape")
+        if image_shape is not None:
+            sampling.height, sampling.width = (int(value) for value in image_shape)
+        elif isinstance(request.prompt, dict):
+            modalities = request.prompt.get("modalities") or []
+            multi_modal_data = request.prompt.get("multi_modal_data") or {}
+            image_input = multi_modal_data.get("img2img")
+            if image_input is None and "text" not in modalities:
+                image_input = multi_modal_data.get("image")
+            if isinstance(image_input, list):
+                image_input = image_input[0] if image_input else None
+            if image_input is not None:
+                if isinstance(image_input, str):
+                    with Image.open(image_input) as image:
+                        input_size = image.size
+                else:
+                    input_size = image_input.size
+                if image_geometry is None:
+                    image_geometry = _resolve_bagel_image_geometry(od_config)
+                latent_downsample, max_latent_size = image_geometry
+                sampling.height, sampling.width = _bagel_effective_image_size(
+                    input_size,
+                    latent_downsample=latent_downsample,
+                    max_latent_size=max_latent_size,
+                )
+
+        extra_args = request.sampling_params.extra_args or {}
+        cfg_interval = tuple(extra_args.get("cfg_interval", (0.4, 1.0)))
+        request.batch_compatibility_key = (
+            "bagel_cfg",
+            float(extra_args.get("cfg_text_scale", 4.0)),
+            float(extra_args.get("cfg_img_scale", 1.5)),
+            cfg_interval,
+            extra_args.get("cfg_renorm_type", "global"),
+            float(extra_args.get("cfg_renorm_min", 0.0)),
+        )
+
+        # Preserve legacy plain-string behavior: only an explicit modality
+        # request selects Bagel.generate_text() instead of stepwise denoising.
+        if step_execution and isinstance(request.prompt, dict):
+            modalities = request.prompt.get("modalities") or []
+            if "text" in modalities:
+                request.use_step_execution = False
+        return request
+
+    return pre_process_func
+
+
 @dataclass
 class _VaeCfg:
     z_channels: int = 16
@@ -168,6 +316,7 @@ class BagelPipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipelineProf
         "bagel.connector",
         "bagel.vit_pos_embed",
     ]
+    supports_step_execution: ClassVar[bool] = True
 
     def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""):
         super().__init__()
@@ -347,17 +496,26 @@ class BagelPipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipelineProf
                 """This model only supports a single prompt, not a batched request.""",
                 """Taking only the first image for now.""",
             )
+        return self._forward_single(req.prompts[0], req.sampling_params)
+
+    @torch.inference_mode()
+    def _forward_single(
+        self,
+        first_prompt: OmniPromptType,
+        sampling: OmniDiffusionSamplingParams,
+        *,
+        prepare_only: bool = False,
+    ) -> DiffusionOutput | dict[str, Any]:
         # TODO: In online mode, sometimes it receives [{"prompts": None}, {...}], so cannot use .get("...", "")
         # TODO: May be some data formatting operations on the API side. Hack for now.
-        first_prompt = req.prompts[0]
-        prompt = first_prompt if isinstance(req.prompts[0], str) else (req.prompts[0].get("prompt") or "")
+        prompt = first_prompt if isinstance(first_prompt, str) else (first_prompt.get("prompt") or "")
 
         max_hw = int(self.bagel.max_latent_size * self.bagel.latent_downsample)
-        if req.sampling_params.height is None and req.sampling_params.width is None:
+        if sampling.height is None and sampling.width is None:
             height = width = max_hw
         else:
-            height = int(req.sampling_params.height) if req.sampling_params.height is not None else max_hw
-            width = int(req.sampling_params.width) if req.sampling_params.width is not None else max_hw
+            height = int(sampling.height) if sampling.height is not None else max_hw
+            width = int(sampling.width) if sampling.width is not None else max_hw
         if height > max_hw or width > max_hw:
             raise ValueError(
                 f"Requested resolution {height}x{width} exceeds Bagel checkpoint limit "
@@ -366,7 +524,7 @@ class BagelPipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipelineProf
             )
         image_shape = (height, width)
 
-        extra_args = getattr(req.sampling_params, "extra_args", {}) or {}
+        extra_args = getattr(sampling, "extra_args", {}) or {}
         cfg_text_scale = extra_args.get("cfg_text_scale", 4.0)
         cfg_img_scale = extra_args.get("cfg_img_scale", 1.5)
 
@@ -375,7 +533,7 @@ class BagelPipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipelineProf
         cfg_renorm_min = extra_args.get("cfg_renorm_min", 0.0)
 
         gen_params = BagelGenParams(
-            num_timesteps=int(req.sampling_params.num_inference_steps or 50),
+            num_timesteps=int(sampling.num_inference_steps or 50),
             timestep_shift=float(extra_args.get("timestep_shift", 3.0)),
             cfg_text_scale=cfg_text_scale,
             cfg_img_scale=cfg_img_scale,
@@ -392,34 +550,30 @@ class BagelPipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipelineProf
         cfg_text_context = deepcopy(gen_context)
         cfg_img_context = deepcopy(gen_context)
 
-        injected_kv = req.sampling_params.past_key_values
+        injected_kv = sampling.past_key_values
         if injected_kv is not None:
             logger.info("Using injected KV Cache (direct)")
             injected_kv = NaiveCache.from_object(injected_kv)
             gen_context["past_key_values"] = injected_kv
             seq_len = injected_kv.key_cache[0].shape[0]
             gen_context["kv_lens"] = [seq_len]
-            if req.sampling_params.kv_metadata and "ropes" in req.sampling_params.kv_metadata:
-                gen_context["ropes"] = req.sampling_params.kv_metadata["ropes"]
+            if sampling.kv_metadata and "ropes" in sampling.kv_metadata:
+                gen_context["ropes"] = sampling.kv_metadata["ropes"]
             else:
                 gen_context["ropes"] = [seq_len]
 
-            if req.sampling_params.kv_metadata and "image_shape" in req.sampling_params.kv_metadata:
-                image_shape = tuple(req.sampling_params.kv_metadata["image_shape"])
+            if sampling.kv_metadata and "image_shape" in sampling.kv_metadata:
+                image_shape = tuple(sampling.kv_metadata["image_shape"])
 
-            branch_kvs = getattr(req.sampling_params, "cfg_branch_past_key_values", None) or {}
-            branch_metadata = getattr(req.sampling_params, "cfg_branch_kv_metadata", None) or {}
-            active_branch = getattr(req.sampling_params, "cfg_active_branch", None)
-            branch_roles = getattr(req.sampling_params, "cfg_branch_roles", None) or list(branch_kvs.keys())
+            branch_kvs = getattr(sampling, "cfg_branch_past_key_values", None) or {}
+            branch_metadata = getattr(sampling, "cfg_branch_kv_metadata", None) or {}
+            active_branch = getattr(sampling, "cfg_active_branch", None)
+            branch_roles = getattr(sampling, "cfg_branch_roles", None) or list(branch_kvs.keys())
 
-            cfg_text_kv = getattr(req.sampling_params, "cfg_text_past_key_values", None) or branch_kvs.get("cfg_text")
-            cfg_text_metadata = getattr(req.sampling_params, "cfg_text_kv_metadata", None) or branch_metadata.get(
-                "cfg_text"
-            )
-            cfg_img_kv = getattr(req.sampling_params, "cfg_img_past_key_values", None) or branch_kvs.get("cfg_img")
-            cfg_img_metadata = getattr(req.sampling_params, "cfg_img_kv_metadata", None) or branch_metadata.get(
-                "cfg_img"
-            )
+            cfg_text_kv = getattr(sampling, "cfg_text_past_key_values", None) or branch_kvs.get("cfg_text")
+            cfg_text_metadata = getattr(sampling, "cfg_text_kv_metadata", None) or branch_metadata.get("cfg_text")
+            cfg_img_kv = getattr(sampling, "cfg_img_past_key_values", None) or branch_kvs.get("cfg_img")
+            cfg_img_metadata = getattr(sampling, "cfg_img_kv_metadata", None) or branch_metadata.get("cfg_img")
 
             cfg_parallel_contract = (
                 active_branch is not None or bool(branch_roles) or cfg_text_kv is not None or cfg_img_kv is not None
@@ -455,8 +609,8 @@ class BagelPipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipelineProf
                 cfg_img_seq_len = injected_kv.key_cache[0].shape[0]
                 cfg_img_context["past_key_values"] = injected_kv
                 cfg_img_context["kv_lens"] = [cfg_img_seq_len]
-                if req.sampling_params.kv_metadata and "ropes" in req.sampling_params.kv_metadata:
-                    cfg_img_context["ropes"] = req.sampling_params.kv_metadata["ropes"]
+                if sampling.kv_metadata and "ropes" in sampling.kv_metadata:
+                    cfg_img_context["ropes"] = sampling.kv_metadata["ropes"]
                 else:
                     cfg_img_context["ropes"] = [cfg_img_seq_len]
             else:
@@ -748,10 +902,10 @@ class BagelPipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipelineProf
             )
 
         # ---- Image generation (text2img / img2img) ----
-        if req.sampling_params.seed is not None:
-            torch.manual_seed(req.sampling_params.seed)
+        if sampling.seed is not None:
+            torch.manual_seed(sampling.seed)
             if self.device.type == "cuda":
-                torch.cuda.manual_seed(req.sampling_params.seed)
+                torch.cuda.manual_seed(sampling.seed)
 
         generation_input = self.bagel.prepare_vae_latent(
             curr_kvlens=gen_context["kv_lens"],
@@ -787,7 +941,7 @@ class BagelPipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipelineProf
 
         # NOTE: For now we disable device specific noise regeneration so that e2e tests can run
         # on both CUDA and ROCm. Context: https://github.com/vllm-project/vllm-omni/pull/4081
-        # self._regen_init_noise_on_device(generation_input, req.sampling_params.seed)
+        # self._regen_init_noise_on_device(generation_input, sampling.seed)
 
         # text cfg
         generation_input_cfg_text = self.bagel.prepare_vae_latent_cfg(
@@ -808,6 +962,19 @@ class BagelPipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipelineProf
             if torch.is_tensor(v):
                 generation_input_cfg_img[k] = v.to(self.device)
 
+        if prepare_only:
+            return {
+                "generation_input": generation_input,
+                "cfg_text_packed_position_ids": generation_input_cfg_text["cfg_packed_position_ids"],
+                "cfg_img_packed_position_ids": generation_input_cfg_img["cfg_packed_position_ids"],
+                "gen_context": gen_context,
+                "cfg_text_context": cfg_text_context,
+                "cfg_img_context": cfg_img_context,
+                "gen_params": gen_params,
+                "image_shape": image_shape,
+                "think_text": think_text,
+            }
+
         with torch.autocast(
             device_type=self.device.type,
             enabled=self.device.type != "cpu",
@@ -827,12 +994,33 @@ class BagelPipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipelineProf
                 **generation_input,
                 cfg_text_packed_position_ids=generation_input_cfg_text["cfg_packed_position_ids"],
                 cfg_img_packed_position_ids=generation_input_cfg_img["cfg_packed_position_ids"],
-                return_trajectory_latents=req.sampling_params.return_trajectory_latents,
+                return_trajectory_latents=sampling.return_trajectory_latents,
                 scheduler=self.scheduler,
                 scheduler_kwargs=self.scheduler_kwargs,
             )
 
-        img = self._decode_image_from_latent(self.bagel, self.vae, latents[0], image_shape)
+        return self._build_image_output(
+            latents[0],
+            image_shape,
+            trajectory_latents=trajectory_latents,
+            trajectory_timesteps=trajectory_timesteps,
+            trajectory_log_probs=trajectory_log_probs,
+            return_trajectory_decoded=sampling.return_trajectory_decoded,
+            think_text=think_text,
+        )
+
+    def _build_image_output(
+        self,
+        latent: torch.Tensor,
+        image_shape: tuple[int, int],
+        *,
+        trajectory_latents: list[torch.Tensor] | None = None,
+        trajectory_timesteps: list[torch.Tensor] | None = None,
+        trajectory_log_probs: list[torch.Tensor] | None = None,
+        return_trajectory_decoded: bool = False,
+        think_text: str | None = None,
+    ) -> DiffusionOutput:
+        img = self._decode_image_from_latent(self.bagel, self.vae, latent, image_shape)
 
         # Build trajectory output when requested
         trajectory_latents_stacked: torch.Tensor | None = None
@@ -841,7 +1029,7 @@ class BagelPipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipelineProf
         if trajectory_latents:
             trajectory_latents_stacked = torch.stack(trajectory_latents)
             trajectory_timesteps_stacked = torch.stack(trajectory_timesteps)
-            if req.sampling_params.return_trajectory_decoded:
+            if return_trajectory_decoded:
                 trajectory_decoded = [
                     self._decode_image_from_latent(self.bagel, self.vae, lat, image_shape) for lat in trajectory_latents
                 ]
@@ -875,6 +1063,317 @@ class BagelPipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipelineProf
             stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None,
         )
 
+    def prepare_encode(
+        self,
+        state: StepRequestState,
+        **kwargs: object,
+    ) -> StepRequestState:
+        """Populate *state* with BAGEL inputs, latents, timesteps, and CFG config."""
+        del kwargs
+        if self.bagel._sp_size > 1:
+            raise NotImplementedError("BAGEL step execution does not currently support sequence parallelism.")
+
+        sampling = state.sampling
+        prompt = state.prompt if state.prompt is not None else ""
+        if isinstance(prompt, dict):
+            if "text" in (prompt.get("modalities") or []):
+                raise NotImplementedError("BAGEL text output is not supported by step execution.")
+
+        ctx = self._forward_single(prompt, sampling, prepare_only=True)
+        if not isinstance(ctx, dict):
+            raise RuntimeError("BAGEL step preparation did not produce an image-generation context.")
+
+        generation_input = ctx["generation_input"]
+        latents = generation_input["packed_init_noises"]
+        gen_params = ctx["gen_params"]
+        timesteps, dts = self.bagel.prepare_denoise_schedule(
+            latents,
+            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.
+        req_scheduler = copy(self.scheduler)
+
+        # Populate state from generation context.
+        state.latents = latents
+        state.timesteps = timesteps
+        state.step_index = 0
+        state.scheduler = req_scheduler
+        state.do_true_cfg = gen_params.cfg_text_scale > 1.0
+        state.img_shapes = [ctx["image_shape"]]
+        state.extra.update(
+            {
+                "bagel_generation_input": {
+                    key: value for key, value in generation_input.items() if key != "packed_init_noises"
+                },
+                "bagel_cfg_text_packed_position_ids": ctx["cfg_text_packed_position_ids"],
+                "bagel_cfg_img_packed_position_ids": ctx["cfg_img_packed_position_ids"],
+                "bagel_gen_context": ctx["gen_context"],
+                "bagel_cfg_text_context": ctx["cfg_text_context"],
+                "bagel_cfg_img_context": ctx["cfg_img_context"],
+                "bagel_gen_params": gen_params,
+                "bagel_dts": dts,
+                "bagel_image_shape": ctx["image_shape"],
+                "bagel_scheduler_kwargs": dict(self.scheduler_kwargs),
+                "bagel_think_text": ctx["think_text"],
+                "bagel_trajectory_latents": (
+                    # Trajectory output needs immutable snapshots because each
+                    # later scheduler update replaces or mutates the latents.
+                    [latents.clone()] if sampling.return_trajectory_latents and len(timesteps) > 0 else []
+                ),
+                "bagel_trajectory_timesteps": [],
+                "bagel_trajectory_log_probs": [],
+            }
+        )
+        return state
+
+    @staticmethod
+    def _pack_step_generation_inputs(
+        states: Sequence[StepRequestState],
+    ) -> dict[str, torch.Tensor]:
+        generation_inputs = [state.extra["bagel_generation_input"] for state in states]
+        packed_seqlens = [generation_input["packed_seqlens"] for generation_input in generation_inputs]
+        if any(seqlens.numel() != 1 for seqlens in packed_seqlens):
+            raise ValueError("BAGEL step execution expects one packed sequence per request.")
+        seq_lengths = [int(seqlens.item()) for seqlens in packed_seqlens]
+        if len(set(seq_lengths)) != 1:
+            raise ValueError("BAGEL step batching requires matching packed sequence lengths.")
+
+        query_offset = 0
+        text_indexes = []
+        vae_indexes = []
+        for generation_input, seq_len in zip(generation_inputs, seq_lengths, strict=True):
+            text_indexes.append(generation_input["packed_text_indexes"] + query_offset)
+            vae_indexes.append(generation_input["packed_vae_token_indexes"] + query_offset)
+            query_offset += seq_len
+
+        return {
+            "packed_text_ids": torch.cat([item["packed_text_ids"] for item in generation_inputs]),
+            "packed_text_indexes": torch.cat(text_indexes),
+            "packed_vae_position_ids": torch.cat([item["packed_vae_position_ids"] for item in generation_inputs]),
+            "packed_vae_token_indexes": torch.cat(vae_indexes),
+            "packed_seqlens": torch.cat(packed_seqlens),
+            "packed_position_ids": torch.cat([item["packed_position_ids"] for item in generation_inputs]),
+        }
+
+    def _build_denoise_kwargs(
+        self,
+        input_batch: InputBatch,
+        states: Sequence[StepRequestState],
+    ) -> dict[str, Any]:
+        parallel_config = getattr(self.bagel, "parallel_config", None)
+        cfg_parallel_size = getattr(parallel_config, "cfg_parallel_size", 1)
+
+        packed = self._pack_step_generation_inputs(states)
+        vae_lengths = [int(state.latents.shape[0]) for state in states]
+        if sum(vae_lengths) != int(input_batch.latents.shape[0]):
+            raise ValueError("BAGEL packed latent rows do not match InputBatch.latents.")
+
+        cfg_settings = {
+            (
+                state.extra["bagel_gen_params"].cfg_text_scale,
+                state.extra["bagel_gen_params"].cfg_img_scale,
+                tuple(state.extra["bagel_gen_params"].cfg_interval),
+                state.extra["bagel_gen_params"].cfg_renorm_type,
+                state.extra["bagel_gen_params"].cfg_renorm_min,
+            )
+            for state in states
+        }
+        if len(cfg_settings) != 1:
+            raise ValueError("Mixed BAGEL CFG settings cannot share one step batch.")
+        cfg_text_scale, cfg_img_scale, _cfg_interval, cfg_renorm_type, cfg_renorm_min = next(iter(cfg_settings))
+
+        cfg_text_scales = []
+        cfg_img_scales = []
+        for state in states:
+            gen_params = state.extra["bagel_gen_params"]
+            timestep = state.current_timestep
+            if timestep is None:
+                raise ValueError(f"BAGEL request {state.request_id} has no current timestep.")
+            t_value = float(timestep.item())
+            in_cfg_window = t_value > gen_params.cfg_interval[0] and t_value <= gen_params.cfg_interval[1]
+            text_scale = gen_params.cfg_text_scale if in_cfg_window else 1.0
+            cfg_text_scales.append(text_scale)
+            cfg_img_scales.append(gen_params.cfg_img_scale if in_cfg_window and text_scale > 1.0 else 1.0)
+
+        use_cfg_text = any(scale > 1.0 for scale in cfg_text_scales)
+        use_cfg_img = use_cfg_text and any(scale > 1.0 for scale in cfg_img_scales)
+        configured_cfg_text = any(state.extra["bagel_gen_params"].cfg_text_scale > 1.0 for state in states)
+        configured_cfg_img = any(
+            state.extra["bagel_gen_params"].cfg_text_scale > 1.0 and state.extra["bagel_gen_params"].cfg_img_scale > 1.0
+            for state in states
+        )
+        build_cfg_text = use_cfg_text or (cfg_parallel_size > 1 and configured_cfg_text)
+        build_cfg_img = use_cfg_img or (cfg_parallel_size > 1 and configured_cfg_img)
+        gen_cache = NaiveCache.merge([state.extra["bagel_gen_context"]["past_key_values"] for state in states])
+        cfg_branch_pids = None
+        cfg_branch_caches = None
+        if build_cfg_text:
+            cfg_branch_pids = [
+                packed["packed_position_ids"],
+                torch.cat([state.extra["bagel_cfg_text_packed_position_ids"] for state in states]),
+            ]
+            cfg_branch_caches = [
+                gen_cache,
+                NaiveCache.merge([state.extra["bagel_cfg_text_context"]["past_key_values"] for state in states]),
+            ]
+            if build_cfg_img:
+                cfg_branch_pids.append(
+                    torch.cat([state.extra["bagel_cfg_img_packed_position_ids"] for state in states])
+                )
+                cfg_branch_caches.append(
+                    NaiveCache.merge([state.extra["bagel_cfg_img_context"]["past_key_values"] for state in states])
+                )
+
+        return {
+            "x_t": input_batch.latents,
+            "timestep": input_batch.timesteps,
+            "past_key_values": gen_cache,
+            "cfg_renorm_min": cfg_renorm_min,
+            "cfg_renorm_type": cfg_renorm_type,
+            "cfg_text_scale": cfg_text_scale,
+            "cfg_img_scale": cfg_img_scale,
+            "cfg_branch_pids": cfg_branch_pids,
+            "cfg_branch_caches": cfg_branch_caches,
+            "cfg_vae_lengths": vae_lengths if use_cfg_text else None,
+            "cfg_text_scales": cfg_text_scales if use_cfg_text else None,
+            "cfg_img_scales": cfg_img_scales if use_cfg_text else None,
+            **packed,
+        }
+
+    def _denoise_step_cfg_parallel(self, denoise_kwargs: dict[str, Any]) -> torch.Tensor:
+        cfg_branch_pids = denoise_kwargs["cfg_branch_pids"]
+        cfg_branch_caches = denoise_kwargs["cfg_branch_caches"]
+        if cfg_branch_pids is None or cfg_branch_caches is None:
+            return self.bagel.forward(**denoise_kwargs)
+
+        cfg_world_size = get_classifier_free_guidance_world_size()
+        num_branches = len(cfg_branch_pids)
+        if cfg_world_size == 2 and num_branches == 3:
+            raise ValueError(
+                f"Image CFG (cfg_img_scale={denoise_kwargs['cfg_img_scale']}) requires cfg_parallel_size=3, "
+                "but got cfg_parallel_size=2. Use cfg_parallel_size=3 to enable image CFG in parallel mode."
+            )
+
+        common_keys = (
+            "x_t",
+            "timestep",
+            "packed_vae_token_indexes",
+            "packed_vae_position_ids",
+            "packed_text_ids",
+            "packed_text_indexes",
+            "packed_seqlens",
+        )
+        common = {key: denoise_kwargs[key] for key in common_keys}
+        branches_kwargs = [
+            {**common, "packed_position_ids": position_ids, "past_key_values": cache}
+            for position_ids, cache in zip(cfg_branch_pids, cfg_branch_caches, strict=True)
+        ]
+        result = self.bagel.predict_noise_with_multi_branch_cfg(
+            do_true_cfg=denoise_kwargs["cfg_text_scales"] is not None,
+            true_cfg_scale={
+                "cfg_text_scale": denoise_kwargs["cfg_text_scale"],
+                "cfg_img_scale": denoise_kwargs["cfg_img_scale"],
+                "cfg_renorm_type": denoise_kwargs["cfg_renorm_type"],
+                "cfg_renorm_min": denoise_kwargs["cfg_renorm_min"],
+                "cfg_vae_lengths": denoise_kwargs["cfg_vae_lengths"],
+                "cfg_text_scales": denoise_kwargs["cfg_text_scales"],
+                "cfg_img_scales": denoise_kwargs["cfg_img_scales"],
+            },
+            branches_kwargs=branches_kwargs,
+        )
+        if not isinstance(result, torch.Tensor):
+            raise TypeError("BAGEL CFG parallel step must return one velocity tensor.")
+        return result
+
+    def denoise_step(
+        self,
+        input_batch: InputBatch,
+        *,
+        states: Sequence[StepRequestState] | None = None,
+        **kwargs: object,
+    ) -> torch.Tensor:
+        """One denoise step: pack request-local BAGEL state and run the model."""
+        del kwargs
+        states = tuple(states or input_batch.states)
+        if not states:
+            raise ValueError("BAGEL denoise_step requires at least one request state.")
+        denoise_kwargs = self._build_denoise_kwargs(input_batch, states)
+        parallel_config = getattr(self.bagel, "parallel_config", None)
+        use_cfg_parallel = (
+            getattr(parallel_config, "cfg_parallel_size", 1) > 1 and denoise_kwargs["cfg_branch_pids"] is not None
+        )
+
+        with torch.autocast(
+            device_type=self.device.type,
+            enabled=self.device.type != "cpu",
+            dtype=self.od_config.dtype,
+        ):
+            if use_cfg_parallel:
+                if any(state.step_index == 0 for state in states):
+                    get_cfg_group().broadcast(input_batch.latents, src=0)
+                return self._denoise_step_cfg_parallel(denoise_kwargs)
+            return self.bagel.forward(**denoise_kwargs)
+
+    def step_scheduler(
+        self,
+        state: StepRequestState,
+        noise_pred: torch.Tensor,
+        **kwargs: object,
+    ) -> None:
+        """One scheduler step: update ``state.latents`` and advance ``step_index``."""
+        del kwargs
+        t = state.current_timestep
+        if t is None or state.latents is None:
+            raise ValueError(f"BAGEL request {state.request_id} is not ready for a scheduler step.")
+        dt = state.extra["bagel_dts"][state.step_index]
+        log_prob = None
+        if state.scheduler is not None:
+            output = state.scheduler.step(
+                noise_pred.to(state.latents.device),
+                t,
+                state.latents,
+                dt,
+                **state.extra["bagel_scheduler_kwargs"],
+            )
+            state.latents = output.prev_sample
+            log_prob = getattr(output, "log_prob", None)
+        else:
+            state.latents = state.latents - noise_pred.to(state.latents.device) * dt
+
+        if state.sampling.return_trajectory_latents:
+            # Preserve this opt-in step snapshot before the next update.
+            state.extra["bagel_trajectory_latents"].append(state.latents.clone())
+            state.extra["bagel_trajectory_timesteps"].append(t)
+            if log_prob is not None:
+                state.extra["bagel_trajectory_log_probs"].append(log_prob)
+        state.step_index += 1
+
+    def post_decode(
+        self,
+        state: StepRequestState,
+        **kwargs: object,
+    ) -> DiffusionOutput:
+        """Decode final latents from *state*."""
+        del kwargs
+        if state.latents is None:
+            raise ValueError(f"BAGEL request {state.request_id} has no latents to decode.")
+        image_shape = state.extra["bagel_image_shape"]
+        return_trajectory_decoded = state.sampling.return_trajectory_decoded
+        return self._build_image_output(
+            state.latents,
+            image_shape,
+            trajectory_latents=state.extra["bagel_trajectory_latents"],
+            trajectory_timesteps=state.extra["bagel_trajectory_timesteps"],
+            trajectory_log_probs=state.extra["bagel_trajectory_log_probs"],
+            return_trajectory_decoded=return_trajectory_decoded,
+            think_text=state.extra["bagel_think_text"],
+        )
+
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         state = self.state_dict()
         allowed = set(state.keys())

vllm_omni/diffusion/models/lance/lance_transformer.py

@@ -278,9 +278,9 @@ class LanceBagel(Bagel):
     ``forward_cache_update_vae`` — is reused unchanged.
     """
 
-    # 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.
+    # Upstream Lance uses the same ``num_timesteps + 1`` schedule-point
+    # convention as BAGEL, so the denoise loop runs exactly ``num_timesteps``
+    # Euler steps.
     _denoise_schedule_extra_step: bool = True
 
     @staticmethod

vllm_omni/diffusion/registry.py

@@ -611,6 +611,7 @@ _DIFFUSION_PRE_PROCESS_FUNCS = {
     # arch: pre_process_func
     # `pre_process_func` function must be placed in {mod_folder}/{mod_relname}.py,
     # where mod_folder and mod_relname are  defined and mapped using `_DIFFUSION_MODELS` via the `arch` key
+    "BagelPipeline": "get_bagel_pre_process_func",
     "GlmImagePipeline": "get_glm_image_pre_process_func",
     "BooguImagePipeline": "get_boogu_image_pre_process_func",
     "QwenImageEditPipeline": "get_qwen_image_edit_pre_process_func",

vllm_omni/diffusion/request.py

@@ -67,6 +67,9 @@ class OmniDiffusionRequest:
     # This is populated by a pipeline preprocessor before the request reaches
     # the scheduler; ``None`` keeps the default behavior for other pipelines.
     batch_compatibility_key: tuple[Any, ...] | None = None
+    # A model preprocessor may keep selected requests on the legacy full-forward
+    # path even when the engine is globally configured for step execution.
+    use_step_execution: bool = True
     # KV-recv wall-clock (ms), set by the runner's _prepare_request_for_forward
     # and carried to DiffusionOutput for the vllm_omni:diffusion_kv_load_s metric.
     kv_recv_ms: float = 0.0

vllm_omni/diffusion/sched/base_scheduler.py

@@ -34,10 +34,12 @@ logger = init_logger(__name__)
 
 BatchSamplingParamsKey = StepBatchSamplingParamsKey | RequestBatchSamplingParamsKey
 
-# LoRA identity is derived from `sampling.lora_request`, not a same-named field
-# on sampling params, so it must be resolved separately from the bulk lookup.
+# LoRA identity and execution mode are request-owned rather than same-named
+# sampling-param fields, so they must be resolved separately from the bulk lookup.
 _STEP_BATCH_SAMPLING_PARAMS_KEY_FIELD_NAMES = frozenset(field.name for field in fields(StepBatchSamplingParamsKey)) - {
-    "lora_int_id"
+    "condition_key",
+    "lora_int_id",
+    "use_step_execution",
 }
 
 
@@ -426,7 +428,9 @@ class BaseScheduler(ABC):
         # LoRA identity is optional on sampling params (and on test stubs).
         lora_request = getattr(sampling, "lora_request", None)
         return StepBatchSamplingParamsKey(
+            condition_key=getattr(request, "batch_compatibility_key", None),
             lora_int_id=lora_request.lora_int_id if lora_request is not None else None,
+            use_step_execution=getattr(request, "use_step_execution", True),
             **{name: getattr(sampling, name) for name in _STEP_BATCH_SAMPLING_PARAMS_KEY_FIELD_NAMES},
         )
 

vllm_omni/diffusion/sched/interface.py

@@ -73,6 +73,15 @@ class StepBatchSamplingParamsKey:
     # because model acceleration hooks are shared by the whole worker batch.
     quality: str | None = None
 
+    # Step-mode engines can admit a model-specific full-forward fallback. Do
+    # not mix it with state-driven denoising in one scheduler wave.
+    use_step_execution: bool = True
+
+    # Pipeline-specific structure populated during preprocessing. This keeps
+    # model-owned settings that must be homogeneous (for example BAGEL CFG
+    # scales and renormalization) out of the generic sampling-params schema.
+    condition_key: tuple[Any, ...] | None = None
+
     # Output count. Requests with different num_outputs_per_prompt produce
     # differently shaped outputs and cannot share a batch.
     num_outputs_per_prompt: int = 1

vllm_omni/diffusion/worker/diffusion_model_runner.py

@@ -933,11 +933,22 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
         """Return whether current pipeline supports step execution."""
         return self.pipeline is not None and supports_step_execution(self.pipeline)
 
-    def _update_states(self, scheduler_output: DiffusionSchedulerOutput) -> tuple[list[StepRequestState], list[str]]:
-        """Step-before update: cleanup finished requests and get/create one running state."""
-        for request_id in scheduler_output.finished_req_ids:
+    def _cleanup_finished_step_requests(self, scheduler_output: DiffusionSchedulerOutput) -> None:
+        """Retire state and paged-KV rows released by the scheduler wave."""
+        finished_req_ids = scheduler_output.finished_req_ids
+        for request_id in finished_req_ids:
             self.state_cache.pop(request_id, None)
 
+        if (
+            getattr(self.od_config, "diffusion_kv_mode", DiffusionKVCacheMode.DENSE_LEGACY)
+            is DiffusionKVCacheMode.PAGED_SCHEDULER
+            and finished_req_ids
+        ):
+            self.remove_diffusion_kv_requests(list(finished_req_ids))
+
+    def _update_states(self, scheduler_output: DiffusionSchedulerOutput) -> tuple[list[StepRequestState], list[str]]:
+        """Resolve cached state and create state for newly admitted requests."""
+
         resolved: list[StepRequestState] = []
         new_request_ids: list[str] = []
         try:
@@ -1074,6 +1085,38 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
             record_output_peak_memory=True,
         )
 
+    def _execute_non_step_requests(self, scheduler_output: DiffusionSchedulerOutput) -> BatchRunnerOutput:
+        """Run a complete legacy forward for requests excluded from step mode."""
+        if scheduler_output.scheduled_cached_reqs.request_ids:
+            raise ValueError("A non-step fallback batch cannot contain cached stepwise requests.")
+
+        runner_outputs: list[RunnerOutput] = []
+        for index, new_req in enumerate(scheduler_output.scheduled_new_reqs):
+            try:
+                result = self.execute_model(
+                    new_req.req,
+                    kv_prefetch_job=getattr(scheduler_output, "kv_prefetch_job", None) if index == 0 else None,
+                    diffusion_kv_metadata=new_req.diffusion_kv_metadata,
+                )
+            except Exception as exc:
+                logger.error(
+                    "Non-step fallback execution failed for %s",
+                    new_req.request_id,
+                    exc_info=True,
+                )
+                result = DiffusionOutput(error=str(exc))
+
+            step_index = getattr(new_req.req.sampling_params, "step_index", None)
+            runner_outputs.append(
+                RunnerOutput(
+                    request_id=new_req.request_id,
+                    step_index=0 if step_index is None else step_index,
+                    finished=True,
+                    result=result,
+                )
+            )
+        return BatchRunnerOutput.from_list(runner_outputs)
+
     def _execute_stepwise(
         self,
         scheduler_output: DiffusionSchedulerOutput,
@@ -1085,8 +1128,28 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
         """Execute one step with explicit validation and profiling policy."""
 
         assert self.pipeline is not None, "Model not loaded. Call load_model() first."
+        # A scheduler wave can release a previous stepwise request while it
+        # admits a full-forward fallback request. Do this before dispatch so
+        # the fallback does not bypass normal state/KV retirement.
+        self._cleanup_finished_step_requests(scheduler_output)
         for new_req in scheduler_output.scheduled_new_reqs:
             validate_new_request_data_identity(new_req)
+        non_step_requests = [
+            new_req
+            for new_req in scheduler_output.scheduled_new_reqs
+            if not getattr(new_req.req, "use_step_execution", True)
+        ]
+        if non_step_requests:
+            scheduled_request_count = len(scheduler_output.scheduled_new_reqs) + len(
+                scheduler_output.scheduled_cached_reqs.request_ids
+            )
+            if len(non_step_requests) != scheduled_request_count:
+                raise ValueError("Cannot mix stepwise and non-step fallback requests in one scheduler batch.")
+            # This wave has no stepwise requests, so the previous batch cannot
+            # be reused and must not keep its step tensors alive.
+            self.input_batch = None
+            return self._execute_non_step_requests(scheduler_output)
+        for new_req in scheduler_output.scheduled_new_reqs:
             if validate_kv_metadata:
                 self._validate_diffusion_kv_metadata(
                     request_id=new_req.req.request_id,
@@ -1102,15 +1165,8 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
 
         # Scheduler metadata is installed only for newly admitted requests;
         # cached requests continue to use the row installed on their first
-        # step. Paged rows are retired here for scheduler-finished requests so
-        # a subsequent wave can reuse the native table slots. Dense execution
-        # must not issue Worker KV cleanup side effects.
-        if (
-            getattr(self.od_config, "diffusion_kv_mode", DiffusionKVCacheMode.DENSE_LEGACY)
-            is DiffusionKVCacheMode.PAGED_SCHEDULER
-            and scheduler_output.finished_req_ids
-        ):
-            self.remove_diffusion_kv_requests(list(scheduler_output.finished_req_ids))
+        # step. Finished-state retirement is shared with full-forward
+        # fallback dispatch above.
         installed_request_ids: list[str] = []
         try:
             for new_req in scheduler_output.scheduled_new_reqs: