diff --git a/docs/design/feature/diffusion_continuous_batching.md b/docs/design/feature/diffusion_continuous_batching.md index 2ec89b9d1..e12455c3d 100644 --- a/docs/design/feature/diffusion_continuous_batching.md +++ b/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 diff --git a/docs/user_guide/diffusion/execution_modes.md b/docs/user_guide/diffusion/execution_modes.md index f9fbc1be0..a4f3f0f50 100644 --- a/docs/user_guide/diffusion/execution_modes.md +++ b/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 diff --git a/docs/user_guide/diffusion_features.md b/docs/user_guide/diffusion_features.md index 9e867dd47..c616a284c 100644 --- a/docs/user_guide/diffusion_features.md +++ b/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 diff --git a/docs/user_guide/examples/online_serving/bagel.md b/docs/user_guide/examples/online_serving/bagel.md index 4b5c6dde6..83d49ceb9 100644 --- a/docs/user_guide/examples/online_serving/bagel.md +++ b/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 diff --git a/recipes/Bagel/BAGEL-7B-MoT.md b/recipes/Bagel/BAGEL-7B-MoT.md index 3b4b6c594..1044ad66b 100644 --- a/recipes/Bagel/BAGEL-7B-MoT.md +++ b/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 diff --git a/tests/config/test_omni_config.py b/tests/config/test_omni_config.py index 126b51c18..f99de9bf3 100644 --- a/tests/config/test_omni_config.py +++ b/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) diff --git a/tests/diffusion/kernels/mot/test_mot_norm_npu.py b/tests/diffusion/kernels/mot/test_mot_norm_npu.py new file mode 100644 index 000000000..928afc9fc --- /dev/null +++ b/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) diff --git a/tests/diffusion/models/bagel/test_step_execution.py b/tests/diffusion/models/bagel/test_step_execution.py new file mode 100644 index 000000000..a17463ef6 --- /dev/null +++ b/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)) diff --git a/tests/diffusion/models/bagel/test_trajectory_recording.py b/tests/diffusion/models/bagel/test_trajectory_recording.py index f3dbc0df2..908b1b732 100644 --- a/tests/diffusion/models/bagel/test_trajectory_recording.py +++ b/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] diff --git a/tests/diffusion/test_diffusion_engine_dummy_run.py b/tests/diffusion/test_diffusion_engine_dummy_run.py index 247fee029..6bf0e83ea 100644 --- a/tests/diffusion/test_diffusion_engine_dummy_run.py +++ b/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 diff --git a/tests/diffusion/test_diffusion_model_runner.py b/tests/diffusion/test_diffusion_model_runner.py index fcb823c92..fee348a3a 100644 --- a/tests/diffusion/test_diffusion_model_runner.py +++ b/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: diff --git a/tests/diffusion/test_diffusion_request.py b/tests/diffusion/test_diffusion_request.py index 7476ec9c1..a3a2e410d 100644 --- a/tests/diffusion/test_diffusion_request.py +++ b/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 diff --git a/tests/diffusion/test_diffusion_scheduler.py b/tests/diffusion/test_diffusion_scheduler.py index ae9b86f80..0a317f1fb 100644 --- a/tests/diffusion/test_diffusion_scheduler.py +++ b/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)) diff --git a/vllm_omni/config/omni_config.py b/vllm_omni/config/omni_config.py index 2ef6bf276..88efc8d60 100644 --- a/vllm_omni/config/omni_config.py +++ b/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, ) diff --git a/vllm_omni/diffusion/diffusion_engine.py b/vllm_omni/diffusion/diffusion_engine.py index 0929e9419..b96b82f07 100644 --- a/vllm_omni/diffusion/diffusion_engine.py +++ b/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)") diff --git a/vllm_omni/diffusion/layers/mot/mot_layernorm.py b/vllm_omni/diffusion/layers/mot/mot_layernorm.py index 0c6bd5b2e..0ff67d7e3 100644 --- a/vllm_omni/diffusion/layers/mot/mot_layernorm.py +++ b/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 # ------------------------------------------------------------------ diff --git a/vllm_omni/diffusion/models/bagel/bagel_transformer.py b/vllm_omni/diffusion/models/bagel/bagel_transformer.py index de21ddccf..84b91b182 100644 --- a/vllm_omni/diffusion/models/bagel/bagel_transformer.py +++ b/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 diff --git a/vllm_omni/diffusion/models/bagel/pipeline_bagel.py b/vllm_omni/diffusion/models/bagel/pipeline_bagel.py index 6ca306c60..8d71bf6c2 100644 --- a/vllm_omni/diffusion/models/bagel/pipeline_bagel.py +++ b/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()) diff --git a/vllm_omni/diffusion/models/lance/lance_transformer.py b/vllm_omni/diffusion/models/lance/lance_transformer.py index 12959eb64..b3a281de0 100644 --- a/vllm_omni/diffusion/models/lance/lance_transformer.py +++ b/vllm_omni/diffusion/models/lance/lance_transformer.py @@ -278,9 +278,9 @@ class LanceBagel(Bagel): ``forward_cache_update_vae`` β€” is reused unchanged. """ - # Upstream Lance 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 diff --git a/vllm_omni/diffusion/registry.py b/vllm_omni/diffusion/registry.py index ed6a816a9..297177052 100644 --- a/vllm_omni/diffusion/registry.py +++ b/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", diff --git a/vllm_omni/diffusion/request.py b/vllm_omni/diffusion/request.py index b2eae3fca..c9dad66f2 100644 --- a/vllm_omni/diffusion/request.py +++ b/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 diff --git a/vllm_omni/diffusion/sched/base_scheduler.py b/vllm_omni/diffusion/sched/base_scheduler.py index 8fe79a11c..ad6fbfadd 100644 --- a/vllm_omni/diffusion/sched/base_scheduler.py +++ b/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}, ) diff --git a/vllm_omni/diffusion/sched/interface.py b/vllm_omni/diffusion/sched/interface.py index 62e5eddef..8cbfeb739 100644 --- a/vllm_omni/diffusion/sched/interface.py +++ b/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 diff --git a/vllm_omni/diffusion/worker/diffusion_model_runner.py b/vllm_omni/diffusion/worker/diffusion_model_runner.py index de415297e..4fea1ed17 100644 --- a/vllm_omni/diffusion/worker/diffusion_model_runner.py +++ b/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: