From da7a8f801480d8fc8c00c1f67bcd481d5d8aa6e8 Mon Sep 17 00:00:00 2001 From: Wu JIAZHEN <83007646+asukaqaq-s@users.noreply.github.com> Date: Fri, 20 Mar 2026 19:21:50 +0800 Subject: [PATCH 1/2] reafator pipeline stage/step pipeline (#1368) Signed-off-by: asukaqaq-s <1311722138@qq.com> --- docs/.nav.yml | 2 + .../model/adding_diffusion_model.md | 16 + .../feature/diffusion_step_execution.md | 121 ++++ docs/user_guide/diffusion/step_execution.md | 61 ++ .../diffusion/test_diffusion_step_pipeline.py | 507 +++++++++++++++++ vllm_omni/diffusion/data.py | 3 + .../diffusion/distributed/cfg_parallel.py | 47 +- vllm_omni/diffusion/ipc.py | 36 +- vllm_omni/diffusion/models/interface.py | 41 ++ .../models/qwen_image/pipeline_qwen_image.py | 532 +++++++++++++----- vllm_omni/diffusion/sched/interface.py | 54 ++ .../worker/diffusion_model_runner.py | 117 ++++ .../diffusion/worker/diffusion_worker.py | 21 +- vllm_omni/diffusion/worker/utils.py | 124 ++++ 14 files changed, 1540 insertions(+), 142 deletions(-) create mode 100644 docs/design/feature/diffusion_step_execution.md create mode 100644 docs/user_guide/diffusion/step_execution.md create mode 100644 tests/diffusion/test_diffusion_step_pipeline.py create mode 100644 vllm_omni/diffusion/sched/interface.py create mode 100644 vllm_omni/diffusion/worker/utils.py diff --git a/docs/.nav.yml b/docs/.nav.yml index b7d08e77e..bfa9365f6 100644 --- a/docs/.nav.yml +++ b/docs/.nav.yml @@ -60,6 +60,7 @@ nav: - FP8: user_guide/diffusion/quantization/fp8.md - Int8: user_guide/diffusion/quantization/int8.md - GGUF: user_guide/diffusion/quantization/gguf.md + - Step Execution: user_guide/diffusion/step_execution.md - Parallelism Acceleration: user_guide/diffusion/parallelism_acceleration.md - CPU Offloading: user_guide/diffusion/cpu_offload_diffusion.md - LoRA: user_guide/diffusion/lora.md @@ -91,6 +92,7 @@ nav: - design/feature/cache_dit.md - design/feature/teacache.md - design/feature/async_chunk_design.md + - design/feature/diffusion_step_execution.md - Module Design: - design/module/ar_module.md - design/module/dit_module.md diff --git a/docs/contributing/model/adding_diffusion_model.md b/docs/contributing/model/adding_diffusion_model.md index d044e74d7..366903433 100644 --- a/docs/contributing/model/adding_diffusion_model.md +++ b/docs/contributing/model/adding_diffusion_model.md @@ -739,6 +739,22 @@ See detailed guide: [How to add Sequence Parallel support](../../design/feature/ omni = Omni(model="your-model", ulysses_degree=2, ring_degree=2) ``` +### Step Execution + +See detailed design guide: [How to add step execution support](../../design/feature/diffusion_step_execution.md) + +Use this only when your pipeline can be split into stable request-scoped and +step-scoped phases. The reference implementation is +`QwenImagePipeline`, which maps its request-level `forward()` into: + +1. `prepare_encode()` for prompt encoding, latent init, timestep prep, and per-request scheduler setup. +2. `denoise_step()` for one transformer/noise prediction. +3. `step_scheduler()` for one scheduler update and `step_index` advance. +4. `post_decode()` for the final VAE decode. + +Do not enable `step_execution=True` until those four methods are implemented +and validated against the request-level path. + ### Cache Acceleration #### TeaCache diff --git a/docs/design/feature/diffusion_step_execution.md b/docs/design/feature/diffusion_step_execution.md new file mode 100644 index 000000000..b8c81f04f --- /dev/null +++ b/docs/design/feature/diffusion_step_execution.md @@ -0,0 +1,121 @@ +# Diffusion Step Execution + +This guide documents vLLM-Omni's stepwise diffusion contract for model authors +and contributors implementing `step_execution=True` support for a diffusion +pipeline. + +For end-user enablement, supported models, and current limitations, see +[Step Execution](../../user_guide/diffusion/step_execution.md). + +## Current Support Scope + +`step_execution` is **not** a generic diffusion toggle. It only works for +pipelines that implement the segmented stateful contract in +[`vllm_omni/diffusion/models/interface.py`](gh-file:vllm_omni/diffusion/models/interface.py). + +Current in-tree support: + +| Pipeline | Example models | Step execution | +|----------|----------------|----------------| +| `QwenImagePipeline` | `Qwen/Qwen-Image`, `Qwen/Qwen-Image-2512` | Yes | +| All other diffusion pipelines | `QwenImageEditPipeline`, `QwenImageEditPlusPipeline`, `QwenImageLayeredPipeline`, GLM-Image, Wan, Flux, etc. | No | + +Current engine/runtime limitations: + +- `StepScheduler` only schedules `batch_size=1`. +- `cache_backend` is not supported in step mode. +- Request-mode extras such as KV transfer are not wired into step mode yet. +- Unsupported pipelines now fail early during model loading instead of failing on the first request. + +## Execution Contract + +Step mode is driven by four pipeline methods plus the shared mutable request +state object: + +- `prepare_encode(state)`: one-time request preparation. +- `denoise_step(state)`: compute the noise prediction for the current step. +- `step_scheduler(state, noise_pred)`: mutate latents and advance step state. +- `post_decode(state)`: decode the final output after denoising is complete. + +The state lives in +[`vllm_omni/diffusion/worker/utils.py`](gh-file:vllm_omni/diffusion/worker/utils.py) +as `DiffusionRequestState`. Store request-scoped tensors there, or use +`state.extra` for model-specific fields that do not justify extending the core +dataclass. + +The worker-side step loop lives in +[`vllm_omni/diffusion/worker/diffusion_model_runner.py`](gh-file:vllm_omni/diffusion/worker/diffusion_model_runner.py): + +1. `prepare_encode()` runs once for a new request. +2. `denoise_step()` runs every scheduler tick. +3. `step_scheduler()` mutates `state.latents` and advances `state.step_index`. +4. `post_decode()` runs exactly once after `state.denoise_completed` becomes true. + +## Recommended Split + +When converting an existing request-level `forward()` pipeline, keep the split +strict and mechanical: + +| Request-level phase | Stepwise method | What belongs there | +|---------------------|-----------------|--------------------| +| Input validation, prompt encoding, latent init, timestep prep, per-request scheduler creation | `prepare_encode()` | Anything that should happen once per request | +| Transformer forward / noise prediction | `denoise_step()` | Pure denoise computation for the current timestep | +| `scheduler.step(...)` and `step_index += 1` | `step_scheduler()` | Only latent/state mutation for one step | +| VAE decode / postprocess | `post_decode()` | Final decode only | + +Keep the stepwise path reusing the same helpers as the request-level path +whenever possible. Reimplementing the denoise loop from scratch is the easiest +way to introduce behavioral drift. + +## Qwen-Image Reference + +[`pipeline_qwen_image.py`](gh-file:vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py) +is the reference implementation and is split correctly for the current +contract: + +- `prepare_encode()` reuses `_prepare_generation_context()` so prompt encoding, + latent init, timestep creation, CFG setup, and shape bookkeeping stay aligned + with `forward()`. +- `prepare_encode()` deep-copies `self.scheduler` **after** + `prepare_timesteps()` so request-specific scheduler state is isolated. +- `denoise_step()` reuses `_build_denoise_kwargs()` plus + `predict_noise_maybe_with_cfg()`, so sequential CFG, CFG-parallel, and + non-CFG behavior stay identical to the request-level path. +- `step_scheduler()` only calls + `scheduler_step_maybe_with_cfg(..., per_request_scheduler=state.scheduler)` + and increments `state.step_index`. +- `post_decode()` reuses `_decode_latents()`, so the final image decode matches + the normal `forward()` path. + +That decomposition is the target pattern for future models. + +## Rules For New Pipelines + +- Do not keep request-scoped scheduler state on `self.scheduler`. Copy it into + `state.scheduler` during `prepare_encode()`. +- Do not mutate `state.step_index` inside `denoise_step()`. Only + `step_scheduler()` should advance the step. +- Do not decode partial outputs in `denoise_step()` or `step_scheduler()`. +- If the request-level pipeline has condition latents, masks, or edit-specific + tensors, store them in `state` or `state.extra`, not in global pipeline + attributes. +- Preserve CFG behavior by sharing the same helper path used by `forward()`. +- Keep `post_decode()` equivalent to the tail of `forward()`. + +## Validation Checklist + +Before marking a pipeline as `supports_step_execution = True`, verify: + +- Stepwise output matches request-level output for the same seed and sampling params. +- Per-request scheduler state is isolated across concurrent requests. +- Abort during denoise does not leak cached state. +- `step_index` reported by `RunnerOutput` matches the scheduler progress. +- CFG-parallel and non-CFG paths both work if the request-level pipeline supports them. + +## Related Files + +- Contract: [`vllm_omni/diffusion/models/interface.py`](gh-file:vllm_omni/diffusion/models/interface.py) +- State: [`vllm_omni/diffusion/worker/utils.py`](gh-file:vllm_omni/diffusion/worker/utils.py) +- Runner loop: [`vllm_omni/diffusion/worker/diffusion_model_runner.py`](gh-file:vllm_omni/diffusion/worker/diffusion_model_runner.py) +- Scheduler transport: [`vllm_omni/diffusion/sched/interface.py`](gh-file:vllm_omni/diffusion/sched/interface.py) +- Reference pipeline: [`vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py`](gh-file:vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py) diff --git a/docs/user_guide/diffusion/step_execution.md b/docs/user_guide/diffusion/step_execution.md new file mode 100644 index 000000000..99c287850 --- /dev/null +++ b/docs/user_guide/diffusion/step_execution.md @@ -0,0 +1,61 @@ +# Step Execution + +Step execution is an opt-in diffusion execution mode enabled with +`step_execution=True` when constructing `Omni`. + +It is not a generic diffusion toggle for every pipeline. Only pipelines that +implement the stepwise contract support it today. + +## Quick Start + +```python +from vllm_omni import Omni +from vllm_omni.inputs.data import OmniDiffusionSamplingParams + +omni = Omni( + model="Qwen/Qwen-Image", + step_execution=True, +) + +outputs = omni.generate( + "A cat sitting on a windowsill", + OmniDiffusionSamplingParams( + num_inference_steps=50, + ), +) +``` + +## Supported Pipelines + +| Pipeline | Example models | Step execution | +|----------|----------------|----------------| +| `QwenImagePipeline` | `Qwen/Qwen-Image`, `Qwen/Qwen-Image-2512` | Yes | +| All other diffusion pipelines | `QwenImageEditPipeline`, `QwenImageEditPlusPipeline`, `QwenImageLayeredPipeline`, GLM-Image, Wan, Flux, etc. | No | + +## Current Limitations + +- `step_execution` currently supports `batch_size=1` only. +- `cache_backend` is not supported together with step execution. +- Unsupported pipelines fail early during model loading. +- Request-mode extras such as KV transfer are not wired into step mode yet. + +## When To Use It + +Use step execution only when you specifically need the pipeline to run through +its stepwise request state machine. For normal diffusion inference, leave it +disabled unless your workflow depends on this mode. + +If you are looking for general diffusion speedups, see +[Diffusion Acceleration Overview](../diffusion_acceleration.md). + +## Troubleshooting + +If model loading fails with a message mentioning `prepare_encode()`, +`denoise_step()`, `step_scheduler()`, and `post_decode()`, the selected +pipeline does not support step execution. + +## For Model Authors + +If you want to add step execution support to a new diffusion pipeline, see the +implementation guide: +[Diffusion Step Execution Design](../../design/feature/diffusion_step_execution.md). diff --git a/tests/diffusion/test_diffusion_step_pipeline.py b/tests/diffusion/test_diffusion_step_pipeline.py new file mode 100644 index 000000000..121fbc0d5 --- /dev/null +++ b/tests/diffusion/test_diffusion_step_pipeline.py @@ -0,0 +1,507 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for step-level diffusion runner and worker execution.""" + +import os +from contextlib import contextmanager +from types import SimpleNamespace + +import pytest +import torch + +import vllm_omni.diffusion.worker.diffusion_model_runner as model_runner_module +from tests.utils import hardware_test +from vllm_omni.diffusion.data import DiffusionOutput +from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin +from vllm_omni.diffusion.distributed.comm import RingComm, SeqAllToAll4D +from vllm_omni.diffusion.distributed.parallel_state import ( + destroy_distributed_env, + get_sp_group, + init_distributed_environment, + initialize_model_parallel, +) +from vllm_omni.diffusion.ipc import ( + pack_diffusion_output_shm, + unpack_diffusion_output_shm, +) +from vllm_omni.diffusion.sched.interface import ( + DiffusionRequestState as SchedulerRequestState, +) +from vllm_omni.diffusion.sched.interface import ( + DiffusionSchedulerOutput, +) +from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner +from vllm_omni.diffusion.worker.diffusion_worker import DiffusionWorker +from vllm_omni.diffusion.worker.utils import RunnerOutput +from vllm_omni.platforms import current_omni_platform + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion] + +# --------------------------------------------------------------------------- +# Helpers & fixtures +# --------------------------------------------------------------------------- + + +@contextmanager +def _noop_forward_context(*args, **kwargs): + del args, kwargs + yield + + +def _update_environment_variables(envs_dict: dict[str, str]) -> None: + for key, value in envs_dict.items(): + os.environ[key] = value + + +class _StepPipeline: + """Minimal pipeline stub that supports step-wise execution.""" + + supports_step_execution = True + + def __init__(self): + self.prepare_calls = 0 + self.denoise_calls = 0 + self.scheduler_calls = 0 + self.decode_calls = 0 + + def prepare_encode(self, state, **kwargs): + del kwargs + self.prepare_calls += 1 + state.timesteps = [torch.tensor(10), torch.tensor(5)] + state.latents = torch.tensor([0.0]) + return state + + def denoise_step(self, state, **kwargs): + del state, kwargs + self.denoise_calls += 1 + return torch.tensor([1.0]) + + def step_scheduler(self, state, noise_pred, **kwargs): + del noise_pred, kwargs + self.scheduler_calls += 1 + state.step_index += 1 + + def post_decode(self, state, **kwargs): + del kwargs + self.decode_calls += 1 + return DiffusionOutput(output=torch.tensor([state.step_index], dtype=torch.float32)) + + +class _IdentityNoiseTransformer(torch.nn.Module): + def forward(self, x: torch.Tensor, **kwargs): + del kwargs + return (x,) + + +class _AdditiveScheduler: + def step(self, noise_pred: torch.Tensor, t: torch.Tensor, latents: torch.Tensor, return_dict: bool = False): + del t, return_dict + return (latents + noise_pred,) + + +class _DistributedStepPipeline(CFGParallelMixin): + supports_step_execution = True + + def __init__(self, mode: str, device: torch.device): + self.mode = mode + self.device = device + self._interrupt = False + self.scheduler = _AdditiveScheduler() + self.transformer = _IdentityNoiseTransformer() + + @property + def interrupt(self): + return self._interrupt + + def prepare_encode(self, state, **kwargs): + del kwargs + state.timesteps = [torch.tensor(1.0, device=self.device)] + state.latents = torch.ones((1, 1), device=self.device) + state.step_index = 0 + state.scheduler = self.scheduler + state.do_true_cfg = self.mode == "cfg" + return state + + def denoise_step(self, state, **kwargs): + del kwargs + if self.mode == "ulysses": + sp_group = get_sp_group().ulysses_group + seq_world_size = torch.distributed.get_world_size(sp_group) + input_tensor = torch.randn(1, 2, 2 * seq_world_size, 2, device=self.device) + original = input_tensor.clone() + intermediate = SeqAllToAll4D.apply(sp_group, input_tensor, 2, 1, False) + output = SeqAllToAll4D.apply(sp_group, intermediate, 1, 2, False) + torch.testing.assert_close(output, original, rtol=1e-5, atol=1e-5) + return torch.ones_like(state.latents) + + if self.mode == "ring": + ring_group = get_sp_group().ring_group + rank = torch.distributed.get_rank(ring_group) + world_size = torch.distributed.get_world_size(ring_group) + comm = RingComm(ring_group) + input_tensor = torch.full((1, 2, 2), float(rank + 1), device=self.device) + recv_tensor = comm.send_recv(input_tensor) + comm.commit() + comm.wait() + expected = torch.full_like(recv_tensor, float(((rank - 1) % world_size) + 1)) + torch.testing.assert_close(recv_tensor, expected, rtol=1e-5, atol=1e-5) + return torch.ones_like(state.latents) + + positive_kwargs = {"x": state.latents + 1} + negative_kwargs = {"x": state.latents - 1} + return self.predict_noise_maybe_with_cfg( + do_true_cfg=True, + true_cfg_scale=1.0, + positive_kwargs=positive_kwargs, + negative_kwargs=negative_kwargs, + cfg_normalize=False, + ) + + def step_scheduler(self, state, noise_pred, **kwargs): + del kwargs + if self.mode == "cfg": + state.latents = self.scheduler_step_maybe_with_cfg( + noise_pred, + state.current_timestep, + state.latents, + do_true_cfg=True, + per_request_scheduler=state.scheduler, + ) + else: + state.latents = state.latents + noise_pred + state.step_index += 1 + + def post_decode(self, state, **kwargs): + del kwargs + return DiffusionOutput(output=state.latents.detach().cpu()) + + +def _make_step_request(num_inference_steps: int = 2): + return SimpleNamespace( + prompts=["a prompt"], + request_ids=["req-1"], + sampling_params=SimpleNamespace( + generator=None, + seed=None, + generator_device=None, + num_inference_steps=num_inference_steps, + ), + ) + + +def _make_runner(): + runner = object.__new__(DiffusionModelRunner) + runner.vllm_config = object() + runner.od_config = SimpleNamespace( + cache_backend=None, + parallel_config=SimpleNamespace(use_hsdp=False), + ) + runner.device = torch.device("cpu") + runner.pipeline = _StepPipeline() + runner.cache_backend = None + runner.offload_backend = None + runner.state_cache = {} + runner.kv_transfer_manager = SimpleNamespace() + return runner + + +def _make_distributed_runner(mode: str, device: torch.device): + runner = object.__new__(DiffusionModelRunner) + runner.vllm_config = object() + runner.od_config = SimpleNamespace( + cache_backend=None, + parallel_config=SimpleNamespace(use_hsdp=False), + ) + runner.device = device + runner.pipeline = _DistributedStepPipeline(mode=mode, device=device) + runner.cache_backend = None + runner.offload_backend = None + runner.state_cache = {} + runner.kv_transfer_manager = SimpleNamespace() + return runner + + +def _make_scheduler_output(req, sched_req_id="req-1", step_id=0, finished_req_ids=None): + return DiffusionSchedulerOutput( + step_id=step_id, + req_states=[SchedulerRequestState(sched_req_id=sched_req_id, req=req)], + finished_req_ids=set() if finished_req_ids is None else set(finished_req_ids), + num_running_reqs=1, + num_waiting_reqs=0, + ) + + +def _expected_output_for_mode(mode: str) -> torch.Tensor: + if mode == "cfg": + return torch.tensor([[3.0]]) + return torch.tensor([[2.0]]) + + +def _distributed_step_worker(local_rank: int, world_size: int, mode: str, master_port: str): + device = torch.device(f"{current_omni_platform.device_type}:{local_rank}") + current_omni_platform.set_device(device) + _update_environment_variables( + { + "RANK": str(local_rank), + "LOCAL_RANK": str(local_rank), + "WORLD_SIZE": str(world_size), + "MASTER_ADDR": "localhost", + "MASTER_PORT": master_port, + } + ) + model_runner_module.set_forward_context = _noop_forward_context + + try: + init_distributed_environment() + if mode == "ulysses": + initialize_model_parallel(ulysses_degree=world_size) + elif mode == "ring": + initialize_model_parallel(ring_degree=world_size) + elif mode == "cfg": + initialize_model_parallel(cfg_parallel_size=world_size) + else: + raise ValueError(f"Unsupported distributed test mode: {mode}") + + runner = _make_distributed_runner(mode, device) + output = DiffusionModelRunner.execute_stepwise( + runner, + _make_scheduler_output(_make_step_request(num_inference_steps=1), step_id=0), + ) + + assert output.finished is True + assert output.result is not None + torch.testing.assert_close(output.result.output, _expected_output_for_mode(mode), rtol=1e-5, atol=1e-5) + assert "req-1" not in runner.state_cache + finally: + destroy_distributed_env() + + +# --------------------------------------------------------------------------- +# Runner / Worker +# --------------------------------------------------------------------------- + + +@pytest.mark.cpu +class TestRunner: + """DiffusionModelRunner.execute_stepwise""" + + def test_completes_request_and_clears_state(self, monkeypatch): + runner = _make_runner() + req = _make_step_request() + monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context) + + first = DiffusionModelRunner.execute_stepwise(runner, _make_scheduler_output(req, step_id=0)) + assert first.req_id == "req-1" + assert first.step_index == 1 + assert first.finished is False + assert first.result is None + assert "req-1" in runner.state_cache + + second = DiffusionModelRunner.execute_stepwise(runner, _make_scheduler_output(req, step_id=1)) + assert second.req_id == "req-1" + assert second.step_index == 2 + assert second.finished is True + assert second.result is not None + assert second.result.error is None + assert torch.equal(second.result.output, torch.tensor([2.0])) + assert "req-1" not in runner.state_cache + + assert runner.pipeline.prepare_calls == 1 + assert runner.pipeline.denoise_calls == 2 + assert runner.pipeline.scheduler_calls == 2 + assert runner.pipeline.decode_calls == 1 + + def test_load_model_rejects_unsupported_step_execution(self, monkeypatch): + class _RequestOnlyPipeline: + pass + + class _FakeLoader: + def __init__(self, *args, **kwargs): + del args, kwargs + + def load_model(self, **kwargs): + del kwargs + return _RequestOnlyPipeline() + + class _FakeProfiler: + consumed_memory = 0 + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc, tb): + del exc_type, exc, tb + return False + + runner = object.__new__(DiffusionModelRunner) + runner.vllm_config = object() + runner.od_config = SimpleNamespace( + enable_cpu_offload=False, + enable_layerwise_offload=False, + enforce_eager=True, + cache_backend=None, + cache_config=None, + step_execution=True, + model_class_name="RequestOnlyPipeline", + parallel_config=SimpleNamespace(use_hsdp=False), + ) + runner.device = torch.device("cpu") + runner.pipeline = None + runner.cache_backend = None + runner.offload_backend = None + runner.state_cache = {} + runner.kv_transfer_manager = SimpleNamespace() + + monkeypatch.setattr(model_runner_module, "DiffusersPipelineLoader", _FakeLoader) + monkeypatch.setattr(model_runner_module, "DeviceMemoryProfiler", _FakeProfiler) + monkeypatch.setattr(model_runner_module, "get_offload_backend", lambda *args, **kwargs: None) + monkeypatch.setattr(model_runner_module, "get_cache_backend", lambda *args, **kwargs: None) + + with pytest.raises(ValueError, match="RequestOnlyPipeline"): + DiffusionModelRunner.load_model(runner) + + +@pytest.mark.cpu +class TestWorker: + """DiffusionWorker.execute_stepwise""" + + def test_delegates_to_model_runner(self): + worker = object.__new__(DiffusionWorker) + expected = RunnerOutput(req_id="req-1", step_index=1, finished=False, result=None) + scheduler_output = SimpleNamespace( + req_states=[ + SimpleNamespace( + req=SimpleNamespace( + sampling_params=SimpleNamespace(lora_request=None), + ) + ) + ] + ) + worker.lora_manager = None + worker.model_runner = SimpleNamespace( + execute_stepwise=lambda arg: expected if arg is scheduler_output else None + ) + + output = DiffusionWorker.execute_stepwise(worker, scheduler_output) + + assert output is expected + + def test_clears_active_lora_before_stepwise_execution(self): + worker = object.__new__(DiffusionWorker) + scheduler_output = SimpleNamespace( + req_states=[ + SimpleNamespace( + req=SimpleNamespace( + sampling_params=SimpleNamespace(lora_request=None), + ) + ) + ] + ) + calls: list[object | None] = [] + + class _FakeLoRAManager: + def set_active_adapter(self, adapter): + calls.append(adapter) + + worker.lora_manager = _FakeLoRAManager() + worker.model_runner = SimpleNamespace(execute_stepwise=lambda arg: RunnerOutput(req_id="req-1")) + + DiffusionWorker.execute_stepwise(worker, scheduler_output) + + assert calls == [None] + + def test_rejects_lora_requests_in_step_mode(self): + worker = object.__new__(DiffusionWorker) + scheduler_output = SimpleNamespace( + req_states=[ + SimpleNamespace( + req=SimpleNamespace( + sampling_params=SimpleNamespace(lora_request=object()), + ) + ) + ] + ) + worker.lora_manager = None + worker.model_runner = SimpleNamespace(execute_stepwise=lambda arg: RunnerOutput(req_id="req-1")) + + with pytest.raises(ValueError, match="does not support LoRA"): + DiffusionWorker.execute_stepwise(worker, scheduler_output) + + +@pytest.mark.cpu +class TestIPC: + def test_pack_unpack_runner_output_shm(self): + tensor = torch.zeros(300_000, dtype=torch.float32) + output = RunnerOutput(req_id="req-1", finished=True, result=DiffusionOutput(output=tensor)) + + packed = pack_diffusion_output_shm(output) + assert isinstance(packed.result.output, dict) + assert packed.result.output["__tensor_shm__"] is True + + unpacked = unpack_diffusion_output_shm(packed) + assert isinstance(unpacked.result.output, torch.Tensor) + torch.testing.assert_close(unpacked.result.output, tensor) + + +@pytest.mark.cpu +class TestSupportedPipelines: + """Step-execution protocol checks for supported pipelines.""" + + def test_qwen_image_supports_step_execution(self): + from vllm_omni.diffusion.models.interface import SupportsStepExecution, supports_step_execution + from vllm_omni.diffusion.models.qwen_image.pipeline_qwen_image import QwenImagePipeline + + # Avoid loading model weights; protocol membership depends on the class contract. + pipeline = object.__new__(QwenImagePipeline) + + assert pipeline.supports_step_execution is True + assert supports_step_execution(pipeline) is True + assert isinstance(pipeline, SupportsStepExecution) is True + + +@hardware_test( + res={"cuda": "L4"}, + num_cards=2, +) +def test_execute_stepwise_with_ulysses_parallel(): + world_size = 2 + if current_omni_platform.get_device_count() < world_size: + pytest.skip(f"Test requires {world_size} devices") + + torch.multiprocessing.spawn( + _distributed_step_worker, + args=(world_size, "ulysses", "29540"), + nprocs=world_size, + ) + + +@hardware_test( + res={"cuda": "L4"}, + num_cards=2, +) +def test_execute_stepwise_with_ring_parallel(): + world_size = 2 + if current_omni_platform.get_device_count() < world_size: + pytest.skip(f"Test requires {world_size} devices") + + torch.multiprocessing.spawn( + _distributed_step_worker, + args=(world_size, "ring", "29541"), + nprocs=world_size, + ) + + +@hardware_test( + res={"cuda": "L4"}, + num_cards=2, +) +def test_execute_stepwise_with_cfg_parallel(): + world_size = 2 + if current_omni_platform.get_device_count() < world_size: + pytest.skip(f"Test requires {world_size} devices") + + torch.multiprocessing.spawn( + _distributed_step_worker, + args=(world_size, "cfg", "29542"), + nprocs=world_size, + ) diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index c22b163f6..f15aae6a0 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -464,6 +464,9 @@ class OmniDiffusionConfig: # Diffusion pipeline Profiling config enable_diffusion_pipeline_profiler: bool = False + # Step mode settings + step_execution: bool = False + @property def is_moe(self) -> bool: num_experts = self.tf_model_config.get("num_experts", None) diff --git a/vllm_omni/diffusion/distributed/cfg_parallel.py b/vllm_omni/diffusion/distributed/cfg_parallel.py index 9f86bce22..0743a00d4 100644 --- a/vllm_omni/diffusion/distributed/cfg_parallel.py +++ b/vllm_omni/diffusion/distributed/cfg_parallel.py @@ -182,7 +182,13 @@ class CFGParallelMixin(metaclass=ABCMeta): """ raise NotImplementedError("Subclasses must implement diffuse") - def scheduler_step(self, noise_pred: torch.Tensor, t: torch.Tensor, latents: torch.Tensor) -> torch.Tensor: + def scheduler_step( + self, + noise_pred: torch.Tensor, + t: torch.Tensor, + latents: torch.Tensor, + per_request_scheduler: Any | None = None, + ) -> torch.Tensor: """ Step the scheduler. @@ -190,14 +196,30 @@ class CFGParallelMixin(metaclass=ABCMeta): noise_pred: Predicted noise t: Current timestep latents: Current latents + per_request_scheduler: Optional request-scoped scheduler that + overrides ``self.scheduler`` for this call. This is + primarily used by step-wise execution, where each request + may keep scheduler state in its own runner-managed state + object. Request-level execution should usually leave this + as ``None`` and continue using ``self.scheduler``. Returns: Updated latents after scheduler step """ - return self.scheduler.step(noise_pred, t, latents, return_dict=False)[0] + sched = per_request_scheduler if per_request_scheduler is not None else getattr(self, "scheduler", None) + if sched is None: + raise ValueError("No scheduler is available. Set self.scheduler or pass per_request_scheduler.") + if not callable(getattr(sched, "step", None)): + raise TypeError("per_request_scheduler must provide a callable step(...) method.") + return sched.step(noise_pred, t, latents, return_dict=False)[0] def scheduler_step_maybe_with_cfg( - self, noise_pred: torch.Tensor, t: torch.Tensor, latents: torch.Tensor, do_true_cfg: bool + self, + noise_pred: torch.Tensor, + t: torch.Tensor, + latents: torch.Tensor, + do_true_cfg: bool, + per_request_scheduler: Any | None = None, ) -> torch.Tensor: """ Step the scheduler with (maybe) automatic CFG parallel synchronization. @@ -210,6 +232,11 @@ class CFGParallelMixin(metaclass=ABCMeta): t: Current timestep latents: Current latents do_true_cfg: Whether CFG is enabled + per_request_scheduler: Optional request-scoped scheduler that + overrides ``self.scheduler`` for this call. This is mainly + needed by step-wise execution, where scheduler state may be + stored per request. Request-level execution should normally + leave this as ``None``. Returns: Updated latents (synchronized across all CFG ranks) @@ -223,13 +250,23 @@ class CFGParallelMixin(metaclass=ABCMeta): # Only rank 0 computes the scheduler step if cfg_rank == 0: - latents = self.scheduler_step(noise_pred, t, latents) + latents = self.scheduler_step( + noise_pred, + t, + latents, + per_request_scheduler=per_request_scheduler, + ) # Broadcast the updated latents to all ranks latents = latents.contiguous() cfg_group.broadcast(latents, src=0) else: # No CFG parallel: directly compute scheduler step - latents = self.scheduler_step(noise_pred, t, latents) + latents = self.scheduler_step( + noise_pred, + t, + latents, + per_request_scheduler=per_request_scheduler, + ) return latents diff --git a/vllm_omni/diffusion/ipc.py b/vllm_omni/diffusion/ipc.py index cac406ef4..d3d7b3aff 100644 --- a/vllm_omni/diffusion/ipc.py +++ b/vllm_omni/diffusion/ipc.py @@ -65,12 +65,7 @@ def _tensor_from_shm(handle: dict[str, Any]) -> torch.Tensor: return tensor -def pack_diffusion_output_shm(output: DiffusionOutput) -> DiffusionOutput: - """Replace large tensors in *output* with shared-memory handles. - - The DiffusionOutput is modified **in-place** so that the (now lightweight) - object can be serialised cheaply through a MessageQueue. - """ +def _pack_diffusion_fields(output: DiffusionOutput) -> DiffusionOutput: if output.output is not None and isinstance(output.output, torch.Tensor): if output.output.nelement() * output.output.element_size() > _SHM_TENSOR_THRESHOLD: output.output = _tensor_to_shm(output.output) @@ -80,10 +75,35 @@ def pack_diffusion_output_shm(output: DiffusionOutput) -> DiffusionOutput: return output -def unpack_diffusion_output_shm(output: DiffusionOutput) -> DiffusionOutput: - """Reconstruct tensors from shared-memory handles produced by ``pack_diffusion_output_shm``.""" +def pack_diffusion_output_shm(output: object) -> object: + """Replace large tensors in diffusion worker outputs with SHM handles. + + Supports either a bare ``DiffusionOutput`` or a wrapper object carrying one + in ``.result`` (for example ``RunnerOutput``). + """ + if isinstance(output, DiffusionOutput): + return _pack_diffusion_fields(output) + + result = getattr(output, "result", None) + if isinstance(result, DiffusionOutput): + output.result = _pack_diffusion_fields(result) + return output + + +def _unpack_diffusion_fields(output: DiffusionOutput) -> DiffusionOutput: if isinstance(output.output, dict) and output.output.get("__tensor_shm__"): output.output = _tensor_from_shm(output.output) if isinstance(output.trajectory_latents, dict) and output.trajectory_latents.get("__tensor_shm__"): output.trajectory_latents = _tensor_from_shm(output.trajectory_latents) return output + + +def unpack_diffusion_output_shm(output: object) -> object: + """Reconstruct tensors from SHM handles in diffusion worker outputs.""" + if isinstance(output, DiffusionOutput): + return _unpack_diffusion_fields(output) + + result = getattr(output, "result", None) + if isinstance(result, DiffusionOutput): + output.result = _unpack_diffusion_fields(result) + return output diff --git a/vllm_omni/diffusion/models/interface.py b/vllm_omni/diffusion/models/interface.py index ef90e9804..ef906472b 100644 --- a/vllm_omni/diffusion/models/interface.py +++ b/vllm_omni/diffusion/models/interface.py @@ -1,11 +1,21 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from __future__ import annotations + from typing import ( + TYPE_CHECKING, + Any, ClassVar, Protocol, runtime_checkable, ) +if TYPE_CHECKING: + import torch + + from vllm_omni.diffusion.data import DiffusionOutput + from vllm_omni.diffusion.worker.utils import DiffusionRequestState + @runtime_checkable class SupportImageInput(Protocol): @@ -21,3 +31,34 @@ class SupportAudioInput(Protocol): @runtime_checkable class SupportAudioOutput(Protocol): support_audio_output: ClassVar[bool] = True + + +@runtime_checkable +class SupportsStepExecution(Protocol): + """State-driven step-level execution protocol for diffusion pipelines. + + Pipelines should split request-level ``forward()`` into: + ``prepare_encode()`` (one-time request setup), ``denoise_step()`` + (one denoise forward), ``step_scheduler()`` (one scheduler update), + and ``post_decode()`` (final decode). + """ + + supports_step_execution: ClassVar[bool] = True + + def prepare_encode(self, state: DiffusionRequestState, **kwargs: Any) -> DiffusionRequestState: + """Prepare request-level inputs and return initialized state.""" + + def denoise_step(self, state: DiffusionRequestState, **kwargs: Any) -> torch.Tensor | None: + """Run one denoise step.""" + + def step_scheduler(self, state: DiffusionRequestState, noise_pred: torch.Tensor, **kwargs: Any) -> None: + """Run one scheduler step.""" + + def post_decode(self, state: DiffusionRequestState, **kwargs: Any) -> DiffusionOutput: + """Decode output after denoise loop.""" + + +def supports_step_execution(pipeline: object) -> bool: + """Return whether `pipeline` implements :class:`SupportsStepExecution`.""" + + return isinstance(pipeline, SupportsStepExecution) diff --git a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py index 8ec3d6b9f..49e6f780b 100644 --- a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py +++ b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py @@ -1,13 +1,14 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import copy import inspect import json import logging import math import os from collections.abc import Iterable -from typing import Any +from typing import TYPE_CHECKING, Any, ClassVar import numpy as np import torch @@ -35,6 +36,10 @@ from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPi from vllm_omni.diffusion.quantization import get_vllm_quant_config_for_layers from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.diffusion.utils.tf_utils import get_transformer_config_kwargs + +if TYPE_CHECKING: + from vllm_omni.diffusion.worker.utils import DiffusionRequestState + from vllm_omni.model_executor.model_loader.weight_utils import ( download_weights_from_hf_specific, ) @@ -238,6 +243,8 @@ def apply_rotary_emb_qwen( class QwenImagePipeline(nn.Module, QwenImageCFGParallelMixin, DiffusionPipelineProfilerMixin): + supports_step_execution: ClassVar[bool] = True + def __init__( self, *, @@ -539,6 +546,372 @@ class QwenImagePipeline(nn.Module, QwenImageCFGParallelMixin, DiffusionPipelineP def interrupt(self): return self._interrupt + def _extract_prompts(self, prompts): + """Extract prompt and negative_prompt from OmniPromptType list.""" + prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in prompts] or None + if all(isinstance(p, str) or p.get("negative_prompt") is None for p in prompts): + negative_prompt = None + elif prompts: + negative_prompt = ["" if isinstance(p, str) else (p.get("negative_prompt") or "") for p in prompts] + else: + negative_prompt = None + return prompt, negative_prompt + + def _prepare_generation_context( + self, + *, + prompt, + negative_prompt, + height, + width, + num_inference_steps, + sigmas, + guidance_scale, + num_images_per_prompt, + generator, + true_cfg_scale, + max_sequence_length, + prompt_embeds=None, + prompt_embeds_mask=None, + negative_prompt_embeds=None, + negative_prompt_embeds_mask=None, + latents=None, + attention_kwargs=None, + callback_on_step_end_tensor_inputs=None, + ): + """Shared preparation logic for forward() and prepare_encode(). + + Validates inputs, encodes prompts, prepares latents, computes timesteps, + and returns all intermediate values as a dict. + """ + self.check_inputs( + prompt, + height, + width, + negative_prompt, + prompt_embeds, + negative_prompt_embeds, + prompt_embeds_mask, + negative_prompt_embeds_mask, + callback_on_step_end_tensor_inputs, + max_sequence_length, + ) + + self._guidance_scale = guidance_scale + self._attention_kwargs = attention_kwargs or {} + self._current_timestep = None + self._interrupt = False + + if prompt is not None and isinstance(prompt, str): + batch_size = 1 + elif prompt is not None and isinstance(prompt, list): + batch_size = len(prompt) + elif prompt_embeds is not None: + batch_size = prompt_embeds.shape[0] + else: + batch_size = 1 + + has_neg_prompt = negative_prompt is not None or ( + negative_prompt_embeds is not None and negative_prompt_embeds_mask is not None + ) + do_true_cfg = true_cfg_scale > 1 and has_neg_prompt + self.check_cfg_parallel_validity(true_cfg_scale, has_neg_prompt) + + prompt_embeds, prompt_embeds_mask = self.encode_prompt( + prompt=prompt, + prompt_embeds=prompt_embeds, + prompt_embeds_mask=prompt_embeds_mask, + num_images_per_prompt=num_images_per_prompt, + max_sequence_length=max_sequence_length, + ) + if do_true_cfg: + negative_prompt_embeds, negative_prompt_embeds_mask = self.encode_prompt( + prompt=negative_prompt, + prompt_embeds=negative_prompt_embeds, + prompt_embeds_mask=negative_prompt_embeds_mask, + num_images_per_prompt=num_images_per_prompt, + max_sequence_length=max_sequence_length, + ) + else: + negative_prompt_embeds = None + negative_prompt_embeds_mask = None + + num_channels_latents = self.transformer.in_channels // 4 + latents = self.prepare_latents( + batch_size * num_images_per_prompt, + num_channels_latents, + height, + width, + prompt_embeds.dtype, + self.device, + generator, + latents, + ) + + img_shapes = [[(1, height // self.vae_scale_factor // 2, width // self.vae_scale_factor // 2)]] * batch_size + + timesteps, num_inference_steps = self.prepare_timesteps( + num_inference_steps, + sigmas, + latents.shape[1], + ) + self._num_timesteps = len(timesteps) + + if self.transformer.guidance_embeds: + guidance = torch.full([1], guidance_scale, dtype=torch.float32) + guidance = guidance.expand(latents.shape[0]) + else: + guidance = None + + txt_seq_lens = prompt_embeds_mask.sum(dim=1).tolist() if prompt_embeds_mask is not None else None + negative_txt_seq_lens = ( + negative_prompt_embeds_mask.sum(dim=1).tolist() if negative_prompt_embeds_mask is not None else None + ) + + return { + "prompt_embeds": prompt_embeds, + "prompt_embeds_mask": prompt_embeds_mask, + "negative_prompt_embeds": negative_prompt_embeds, + "negative_prompt_embeds_mask": negative_prompt_embeds_mask, + "latents": latents, + "img_shapes": img_shapes, + "timesteps": timesteps, + "do_true_cfg": do_true_cfg, + "guidance": guidance, + "txt_seq_lens": txt_seq_lens, + "negative_txt_seq_lens": negative_txt_seq_lens, + } + + def prepare_encode( + self, + state: "DiffusionRequestState", + **kwargs: Any, + ) -> "DiffusionRequestState": + """Populate *state* with encoded prompts, latents, timesteps, and CFG config.""" + sampling = state.sampling + prompt, negative_prompt = self._extract_prompts(state.prompts or []) + + ctx = self._prepare_generation_context( + prompt=prompt, + negative_prompt=negative_prompt, + height=sampling.height or self.default_sample_size * self.vae_scale_factor, + width=sampling.width or self.default_sample_size * self.vae_scale_factor, + num_inference_steps=sampling.num_inference_steps or 50, + sigmas=sampling.sigmas, + guidance_scale=sampling.guidance_scale if sampling.guidance_scale_provided else 1.0, + num_images_per_prompt=sampling.num_outputs_per_prompt if sampling.num_outputs_per_prompt > 0 else 1, + generator=sampling.generator, + true_cfg_scale=sampling.true_cfg_scale or 4.0, + max_sequence_length=sampling.max_sequence_length or 512, + attention_kwargs=kwargs.get("attention_kwargs"), + ) + + # prepare_timesteps() has already materialized request-specific timestep + # state on self.scheduler, so deepcopy preserves dynamic-shifting state + # without replaying set_timesteps() on the per-request scheduler. + # Per-request scheduler (must not share state with self.scheduler) + req_scheduler = copy.deepcopy(self.scheduler) + req_scheduler.set_begin_index(0) + + # Populate state from generation context + state.prompt_embeds = ctx["prompt_embeds"] + state.prompt_embeds_mask = ctx["prompt_embeds_mask"] + state.negative_prompt_embeds = ctx["negative_prompt_embeds"] + state.negative_prompt_embeds_mask = ctx["negative_prompt_embeds_mask"] + state.latents = ctx["latents"] + state.timesteps = ctx["timesteps"] + state.step_index = 0 + state.scheduler = req_scheduler + state.do_true_cfg = ctx["do_true_cfg"] + state.guidance = ctx["guidance"] + state.img_shapes = ctx["img_shapes"] + state.txt_seq_lens = ctx["txt_seq_lens"] + state.negative_txt_seq_lens = ctx["negative_txt_seq_lens"] + # QwenImage always normalizes CFG output (matching forward()) + state.sampling.cfg_normalize = True + + return state + + def _build_denoise_kwargs( + self, + latents: torch.Tensor, + timestep: torch.Tensor, + guidance: torch.Tensor | None, + prompt_embeds: torch.Tensor, + prompt_embeds_mask: torch.Tensor, + img_shapes: list, + txt_seq_lens: list[int] | None, + do_true_cfg: bool, + negative_prompt_embeds: torch.Tensor | None, + negative_prompt_embeds_mask: torch.Tensor | None, + negative_txt_seq_lens: list[int] | None, + image_latents: torch.Tensor | None = None, + extra_transformer_kwargs: dict[str, Any] | None = None, + ) -> tuple[dict[str, Any], dict[str, Any] | None, int | None]: + """Build positive/negative kwargs and output_slice for one denoise step. + + Returns: + (positive_kwargs, negative_kwargs, output_slice) + """ + extra_transformer_kwargs = extra_transformer_kwargs or {} + + # Broadcast timestep to match batch size + t_for_model = timestep.expand(latents.shape[0]).to( + device=latents.device, + dtype=latents.dtype, + ) + + # Concatenate image latents if available (editing pipelines) + latent_model_input = latents + if image_latents is not None: + latent_model_input = torch.cat([latents, image_latents], dim=1) + + positive_kwargs = { + "hidden_states": latent_model_input, + "timestep": t_for_model / 1000, + "guidance": guidance, + "encoder_hidden_states_mask": prompt_embeds_mask, + "encoder_hidden_states": prompt_embeds, + "img_shapes": img_shapes, + "txt_seq_lens": txt_seq_lens, + **extra_transformer_kwargs, + } + if do_true_cfg: + negative_kwargs = { + "hidden_states": latent_model_input, + "timestep": t_for_model / 1000, + "guidance": guidance, + "encoder_hidden_states_mask": negative_prompt_embeds_mask, + "encoder_hidden_states": negative_prompt_embeds, + "img_shapes": img_shapes, + "txt_seq_lens": negative_txt_seq_lens, + **extra_transformer_kwargs, + } + else: + negative_kwargs = None + + output_slice = latents.size(1) if image_latents is not None else None + return positive_kwargs, negative_kwargs, output_slice + + def _decode_latents( + self, + latents: torch.Tensor, + height: int, + width: int, + output_type: str = "pil", + ) -> DiffusionOutput: + """Unpack, normalize, and VAE-decode latents into a DiffusionOutput.""" + if output_type == "latent": + return DiffusionOutput( + output=latents, + stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None, + ) + + latents = self._unpack_latents(latents, height, width, self.vae_scale_factor) + latents = latents.to(self.vae.dtype) + latents_mean = ( + torch.tensor(self.vae.config.latents_mean) + .view(1, self.vae.config.z_dim, 1, 1, 1) + .to(latents.device, latents.dtype) + ) + latents_std = 1.0 / torch.tensor(self.vae.config.latents_std).view(1, self.vae.config.z_dim, 1, 1, 1).to( + latents.device, latents.dtype + ) + latents = latents / latents_std + latents_mean + image = self.vae.decode(latents, return_dict=False)[0][:, :, 0] + return DiffusionOutput( + output=image, + stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None, + ) + + def denoise_step( + self, + state: "DiffusionRequestState", + **kwargs: Any, + ) -> torch.Tensor | None: + """One denoise step: read from *state*, delegate to CFGParallelMixin. + + Reuses ``predict_noise_maybe_with_cfg`` so that CFG-parallel, + sequential-CFG, and no-CFG paths are handled identically to + ``diffuse()``. + """ + if self.interrupt: + return None + + t = state.current_timestep + self._current_timestep = t + self.transformer.do_true_cfg = state.do_true_cfg + + # Normalize timestep to [batch_size] tensor + if not torch.is_tensor(t): + t = torch.tensor(t, device=state.latents.device, dtype=state.latents.dtype) + + positive_kwargs, negative_kwargs, output_slice = self._build_denoise_kwargs( + latents=state.latents, + timestep=t, + guidance=state.guidance, + prompt_embeds=state.prompt_embeds, + prompt_embeds_mask=state.prompt_embeds_mask, + img_shapes=state.img_shapes, + txt_seq_lens=state.txt_seq_lens, + do_true_cfg=state.do_true_cfg, + negative_prompt_embeds=state.negative_prompt_embeds, + negative_prompt_embeds_mask=state.negative_prompt_embeds_mask, + negative_txt_seq_lens=state.negative_txt_seq_lens, + image_latents=state.sampling.image_latent, + extra_transformer_kwargs={ + "attention_kwargs": self.attention_kwargs, + "return_dict": False, + }, + ) + + true_cfg_scale = state.sampling.true_cfg_scale or 4.0 + cfg_normalize = state.sampling.cfg_normalize + + return self.predict_noise_maybe_with_cfg( + state.do_true_cfg, + true_cfg_scale, + positive_kwargs, + negative_kwargs, + cfg_normalize, + output_slice, + ) + + def step_scheduler( + self, + state: "DiffusionRequestState", + noise_pred: torch.Tensor, + **kwargs: Any, + ) -> None: + """One scheduler step: update ``state.latents`` and advance ``step_index``.""" + if self.interrupt: + return + + t = state.current_timestep + state.latents = self.scheduler_step_maybe_with_cfg( + noise_pred, + t, + state.latents, + state.do_true_cfg, + per_request_scheduler=state.scheduler, + ) + + state.step_index += 1 + + def post_decode( + self, + state: "DiffusionRequestState", + **kwargs: Any, + ) -> DiffusionOutput: + """Decode final latents from *state*.""" + self._current_timestep = None + + height = state.sampling.height or self.default_sample_size * self.vae_scale_factor + width = state.sampling.width or self.default_sample_size * self.vae_scale_factor + output_type = kwargs.get("output_type", "pil") + + return self._decode_latents(state.latents, height, width, output_type) + def forward( self, req: OmniDiffusionRequest, @@ -562,13 +935,8 @@ class QwenImagePipeline(nn.Module, QwenImageCFGParallelMixin, DiffusionPipelineP callback_on_step_end_tensor_inputs: list[str] = ["latents"], max_sequence_length: int = 512, ) -> DiffusionOutput: - # TODO: In online mode, sometimes it receives [{"negative_prompt": None}, {...}], so cannot use .get("...", "") - # TODO: May be some data formatting operations on the API side. Hack for now. - prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in req.prompts] or prompt - if all(isinstance(p, str) or p.get("negative_prompt") is None for p in req.prompts): - negative_prompt = None - elif req.prompts: - negative_prompt = ["" if isinstance(p, str) else (p.get("negative_prompt") or "") for p in req.prompts] + extracted_prompt, negative_prompt = self._extract_prompts(req.prompts) + prompt = extracted_prompt or prompt height = req.sampling_params.height or self.default_sample_size * self.vae_scale_factor width = req.sampling_params.width or self.default_sample_size * self.vae_scale_factor @@ -584,113 +952,40 @@ class QwenImagePipeline(nn.Module, QwenImageCFGParallelMixin, DiffusionPipelineP if req.sampling_params.num_outputs_per_prompt > 0 else num_images_per_prompt ) - # 1. check inputs - # 2. encode prompts - # 3. prepare latents and timesteps - # 4. diffusion process - # 5. decode latents - # 6. post-process outputs - self.check_inputs( - prompt, - height, - width, - negative_prompt, - prompt_embeds, - negative_prompt_embeds, - prompt_embeds_mask, - negative_prompt_embeds_mask, - callback_on_step_end_tensor_inputs, - max_sequence_length, - ) - self._guidance_scale = guidance_scale - self._attention_kwargs = attention_kwargs - self._current_timestep = None - self._interrupt = False - - if prompt is not None and isinstance(prompt, str): - batch_size = 1 - elif prompt is not None and isinstance(prompt, list): - batch_size = len(prompt) - else: - batch_size = prompt_embeds.shape[0] - - has_neg_prompt = negative_prompt is not None or ( - negative_prompt_embeds is not None and negative_prompt_embeds_mask is not None - ) - - do_true_cfg = true_cfg_scale > 1 and has_neg_prompt - self.check_cfg_parallel_validity(true_cfg_scale, has_neg_prompt) - - prompt_embeds, prompt_embeds_mask = self.encode_prompt( + ctx = self._prepare_generation_context( prompt=prompt, + negative_prompt=negative_prompt, + height=height, + width=width, + num_inference_steps=num_inference_steps, + sigmas=sigmas, + guidance_scale=guidance_scale, + num_images_per_prompt=num_images_per_prompt, + generator=generator, + true_cfg_scale=true_cfg_scale, + max_sequence_length=max_sequence_length, prompt_embeds=prompt_embeds, prompt_embeds_mask=prompt_embeds_mask, - num_images_per_prompt=num_images_per_prompt, - max_sequence_length=max_sequence_length, + negative_prompt_embeds=negative_prompt_embeds, + negative_prompt_embeds_mask=negative_prompt_embeds_mask, + latents=latents, + attention_kwargs=attention_kwargs, + callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs, ) - if do_true_cfg: - negative_prompt_embeds, negative_prompt_embeds_mask = self.encode_prompt( - prompt=negative_prompt, - prompt_embeds=negative_prompt_embeds, - prompt_embeds_mask=negative_prompt_embeds_mask, - num_images_per_prompt=num_images_per_prompt, - max_sequence_length=max_sequence_length, - ) - - num_channels_latents = self.transformer.in_channels // 4 - latents = self.prepare_latents( - batch_size * num_images_per_prompt, - num_channels_latents, - height, - width, - prompt_embeds.dtype, - self.device, - generator, - latents, - ) - img_shapes = [ - [ - ( - 1, - height // self.vae_scale_factor // 2, - width // self.vae_scale_factor // 2, - ) - ] - ] * batch_size - - timesteps, num_inference_steps = self.prepare_timesteps(num_inference_steps, sigmas, latents.shape[1]) - # num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0) - self._num_timesteps = len(timesteps) - - # handle guidance - if self.transformer.guidance_embeds: - guidance = torch.full([1], guidance_scale, dtype=torch.float32) - guidance = guidance.expand(latents.shape[0]) - else: - guidance = None - - if self.attention_kwargs is None: - self._attention_kwargs = {} - - txt_seq_lens = prompt_embeds_mask.sum(dim=1).tolist() if prompt_embeds_mask is not None else None - negative_txt_seq_lens = ( - negative_prompt_embeds_mask.sum(dim=1).tolist() if negative_prompt_embeds_mask is not None else None - ) - # print inputp params latents = self.diffuse( - prompt_embeds, - prompt_embeds_mask, - negative_prompt_embeds, - negative_prompt_embeds_mask, - latents, - img_shapes, - txt_seq_lens, - negative_txt_seq_lens, - timesteps, - do_true_cfg, - guidance, + ctx["prompt_embeds"], + ctx["prompt_embeds_mask"], + ctx["negative_prompt_embeds"], + ctx["negative_prompt_embeds_mask"], + ctx["latents"], + ctx["img_shapes"], + ctx["txt_seq_lens"], + ctx["negative_txt_seq_lens"], + ctx["timesteps"], + ctx["do_true_cfg"], + ctx["guidance"], true_cfg_scale, image_latents=None, cfg_normalize=True, @@ -701,26 +996,7 @@ class QwenImagePipeline(nn.Module, QwenImageCFGParallelMixin, DiffusionPipelineP ) self._current_timestep = None - if output_type == "latent": - image = latents - else: - latents = self._unpack_latents(latents, height, width, self.vae_scale_factor) - latents = latents.to(self.vae.dtype) - latents_mean = ( - torch.tensor(self.vae.config.latents_mean) - .view(1, self.vae.config.z_dim, 1, 1, 1) - .to(latents.device, latents.dtype) - ) - latents_std = 1.0 / torch.tensor(self.vae.config.latents_std).view(1, self.vae.config.z_dim, 1, 1, 1).to( - latents.device, latents.dtype - ) - latents = latents / latents_std + latents_mean - image = self.vae.decode(latents, return_dict=False)[0][:, :, 0] - # processed_image = self.image_processor.postprocess(image, output_type=output_type) - - return DiffusionOutput( - output=image, stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None - ) + return self._decode_latents(latents, height, width, output_type) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: loader = AutoWeightsLoader(self) diff --git a/vllm_omni/diffusion/sched/interface.py b/vllm_omni/diffusion/sched/interface.py new file mode 100644 index 000000000..d8b3a39c5 --- /dev/null +++ b/vllm_omni/diffusion/sched/interface.py @@ -0,0 +1,54 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +import enum +from dataclasses import dataclass + +from vllm_omni.diffusion.request import OmniDiffusionRequest + + +class DiffusionRequestStatus(enum.IntEnum): + """Request status tracked by diffusion scheduler.""" + + WAITING = enum.auto() + RUNNING = enum.auto() + PREEMPTED = enum.auto() + + # if any status is after or equal to FINISHED_COMPLETED, it is considered finished + FINISHED_COMPLETED = enum.auto() + FINISHED_ABORTED = enum.auto() + FINISHED_ERROR = enum.auto() + + @staticmethod + def is_finished(status: DiffusionRequestStatus) -> bool: + return status >= DiffusionRequestStatus.FINISHED_COMPLETED + + +@dataclass +class DiffusionRequestState: + """Scheduler-owned state for one queued OmniDiffusionRequest.""" + + sched_req_id: str + req: OmniDiffusionRequest + status: DiffusionRequestStatus = DiffusionRequestStatus.WAITING + error: str | None = None + + def is_finished(self) -> bool: + return DiffusionRequestStatus.is_finished(self.status) + + +@dataclass +class DiffusionSchedulerOutput: + """Output of a single scheduling cycle. + + Kept intentionally small so step-execution components can share a stable + transport shape while scheduler policy continues to evolve. + """ + + step_id: int + req_states: list[DiffusionRequestState] + finished_req_ids: set[str] + num_running_reqs: int + num_waiting_reqs: int diff --git a/vllm_omni/diffusion/worker/diffusion_model_runner.py b/vllm_omni/diffusion/worker/diffusion_model_runner.py index 972c95c29..e3c27f945 100644 --- a/vllm_omni/diffusion/worker/diffusion_model_runner.py +++ b/vllm_omni/diffusion/worker/diffusion_model_runner.py @@ -10,6 +10,7 @@ model-related operations. from __future__ import annotations +import copy import time from collections.abc import Iterable from contextlib import nullcontext @@ -26,9 +27,12 @@ from vllm_omni.diffusion.compile import regionally_compile from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig from vllm_omni.diffusion.forward_context import set_forward_context from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader +from vllm_omni.diffusion.models.interface import supports_step_execution from vllm_omni.diffusion.offloader import get_offload_backend from vllm_omni.diffusion.registry import _NO_CACHE_ACCELERATION from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.diffusion.sched.interface import DiffusionSchedulerOutput +from vllm_omni.diffusion.worker.utils import DiffusionRequestState, RunnerOutput from vllm_omni.distributed.omni_connectors.kv_transfer_manager import OmniKVTransferManager from vllm_omni.platforms import current_omni_platform @@ -65,6 +69,9 @@ class DiffusionModelRunner: self.cache_backend = None self.offload_backend = None + # Cache for per-request stepwise state. + self.state_cache: dict[str, DiffusionRequestState] = {} + # Initialize KV cache manager for connector management self.kv_transfer_manager = OmniKVTransferManager.from_od_config(od_config) @@ -138,6 +145,13 @@ class DiffusionModelRunner: ) logger.info("Model runner: Model loaded successfully.") + if getattr(self.od_config, "step_execution", False) and not self.supports_step_mode(): + raise ValueError( + "step_execution=True requires a pipeline implementing " + "prepare_encode(), denoise_step(), step_scheduler(), and post_decode(); " + f"{self.od_config.model_class_name} does not support that contract." + ) + # Apply CPU offloading self.offload_backend = get_offload_backend(self.od_config, device=self.device) if self.offload_backend is not None: @@ -239,3 +253,106 @@ class DiffusionModelRunner: cache_summary(self.pipeline, details=True) return output + + # ------------------------------------------------------------------ + # Step-wise execution + # ------------------------------------------------------------------ + + def supports_step_mode(self) -> bool: + """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) -> DiffusionRequestState: + """Step-before update: cleanup finished requests and get/create one running state.""" + for req_id in scheduler_output.finished_req_ids: + self.state_cache.pop(req_id, None) + + req_states = scheduler_output.req_states + if len(req_states) != 1: + raise ValueError(f"Step mode currently supports batch_size=1, but got {len(req_states)} req_states.") + + # TODO: remove req state from SchedulerOutput + # Stepwise mode currently trusts runner-owned cached state more than + # re-validating scheduler-provided request content on every step. + sched_req_state = req_states[0] + req_id = sched_req_state.sched_req_id + if req_id in self.state_cache: + return self.state_cache[req_id] + + req = sched_req_state.req + request_ids = req.request_ids or [req_id] + if len(request_ids) != len(req.prompts): + raise ValueError( + f"request_ids length ({len(request_ids)}) does not match prompts length ({len(req.prompts)})" + ) + + state = DiffusionRequestState( + req_id=req_id, + sampling=copy.deepcopy(req.sampling_params), + prompts=req.prompts, + ) + self.state_cache[req_id] = state + return state + + def _update_states_after(self, state: DiffusionRequestState, finished: bool) -> None: + """Step-after update: clear cached state for completed request.""" + if finished: + self.state_cache.pop(state.req_id, None) + + def execute_stepwise(self, scheduler_output: DiffusionSchedulerOutput) -> RunnerOutput: + """Execute one step for one scheduled request and return runner output.""" + assert self.pipeline is not None, "Model not loaded. Call load_model() first." + if not self.supports_step_mode(): + raise ValueError("Current pipeline does not support step execution.") + # Stepwise mode only supports the basic state-driven denoise path for now. + # Request-mode extras such as cache backends, KV transfer, editing inputs, + # and similar features are not supported here yet. + if self.od_config.cache_backend not in (None, "none"): + raise ValueError("Step mode does not support cache_backend yet.") + + use_hsdp = self.od_config.parallel_config.use_hsdp + grad_context = torch.no_grad() if use_hsdp else torch.inference_mode() + with grad_context: + state = self._update_states(scheduler_output) + + if state.new_request: + # TODO: support kv manager recv + # TODO: support cache backend + if state.sampling.generator is None and state.sampling.seed is not None: + if state.sampling.generator_device is not None: + gen_device = state.sampling.generator_device + elif self.device.type == "cpu": + gen_device = "cpu" + else: + gen_device = self.device + state.sampling.generator = torch.Generator(device=gen_device).manual_seed(state.sampling.seed) + + with set_forward_context(vllm_config=self.vllm_config, omni_diffusion_config=self.od_config): + # step0/new request: encode + if state.new_request: + self.pipeline.prepare_encode(state) + + noise_pred = self.pipeline.denoise_step(state) + finished = False + + # In CFG parallel mode, only rank 0 gets the actual noise_pred; non-rank-0 workers receive None. + # A true interrupt (all ranks return None) is detected by checking self.pipeline.interrupt. + if noise_pred is None and getattr(self.pipeline, "interrupt", False): + finished = True + result = DiffusionOutput(error="stepwise denoise interrupted") + else: + self.pipeline.step_scheduler(state, noise_pred) + finished = state.denoise_completed + if finished: + result = self.pipeline.post_decode(state) + else: + result = None + + self._update_states_after(state, finished) + + return RunnerOutput( + req_id=state.req_id, + step_index=state.step_index, + finished=finished, + result=result, + ) diff --git a/vllm_omni/diffusion/worker/diffusion_worker.py b/vllm_omni/diffusion/worker/diffusion_worker.py index 1b9dc5667..3ca12d8e6 100644 --- a/vllm_omni/diffusion/worker/diffusion_worker.py +++ b/vllm_omni/diffusion/worker/diffusion_worker.py @@ -38,7 +38,9 @@ from vllm_omni.diffusion.ipc import pack_diffusion_output_shm from vllm_omni.diffusion.lora.manager import DiffusionLoRAManager from vllm_omni.diffusion.profiler import CurrentProfiler from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.diffusion.sched.interface import DiffusionSchedulerOutput from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner +from vllm_omni.diffusion.worker.utils import RunnerOutput from vllm_omni.lora.request import LoRARequest from vllm_omni.platforms import current_omni_platform from vllm_omni.worker.gpu_memory_utils import get_process_gpu_memory @@ -194,6 +196,19 @@ class DiffusionWorker: logger.warning("LoRA activation skipped: %s", exc) return self.model_runner.execute_model(req) + def execute_stepwise(self, scheduler_output: DiffusionSchedulerOutput) -> RunnerOutput: + """Execute one diffusion step by delegating to the model runner.""" + assert self.model_runner is not None, "Model runner not initialized" + if self.lora_manager is not None: + # Step mode does not support LoRA yet. Clear any previously active + # adapter first so worker-local LoRA state cannot leak in. + self.lora_manager.set_active_adapter(None) + + if any(req_state.req.sampling_params.lora_request is not None for req_state in scheduler_output.req_states): + raise ValueError("Step mode does not support LoRA yet.") + + return self.model_runner.execute_stepwise(scheduler_output) + def load_weights(self, weights) -> set[str]: """Load weights by delegating to the model runner.""" assert self.model_runner is not None, "Model runner not initialized" @@ -376,7 +391,7 @@ class WorkerProc: ) return wrapper - def return_result(self, output: DiffusionOutput): + def return_result(self, output: object): """Reply to client, only on rank 0.""" if self.result_mq is not None: try: @@ -623,6 +638,10 @@ class WorkerWrapperBase: """ return self.worker.execute_model(reqs, od_config) + def execute_stepwise(self, scheduler_output: DiffusionSchedulerOutput) -> RunnerOutput: + """Execute one diffusion step.""" + return self.worker.execute_stepwise(scheduler_output) + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: """ Load model weights. diff --git a/vllm_omni/diffusion/worker/utils.py b/vllm_omni/diffusion/worker/utils.py new file mode 100644 index 000000000..af88913e7 --- /dev/null +++ b/vllm_omni/diffusion/worker/utils.py @@ -0,0 +1,124 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Per-request mutable state for step-wise diffusion execution.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Any + +import torch + +if TYPE_CHECKING: + from vllm_omni.diffusion.data import DiffusionOutput + from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniPromptType + + +@dataclass +class DiffusionRequestState: + """Per-request mutable state across all pipeline stages. + + Owned by Runner and passed through all step-execution stages: + ``prepare_encode()`` initializes/updates fields, ``denoise_step()`` and + ``step_scheduler()`` mutate per-step fields, and ``post_decode()`` + consumes final latents. This state object is also the cache unit for + future continuous batching. + + This dataclass keeps only the minimal cross-model state required by the + step-execution contract. Pipeline-specific state should be stored in + ``extra`` and promoted here only when it becomes shared across models. + + Examples: + - Wan-style pipelines may keep ``condition``, ``first_frame_mask``, or + ``image_embeds`` in ``extra``. + - Bagel-style pipelines may keep ``gen_context``, + ``cfg_text_context``, ``cfg_img_context``, or ``image_shape`` in + ``extra``. + """ + + # ── Identity / request-level inputs ── + req_id: str + sampling: OmniDiffusionSamplingParams + prompts: list[OmniPromptType] | None = None + + # ── Encoded prompts (set once by prepare_encode) ── + prompt_embeds: torch.Tensor | None = None + prompt_embeds_mask: torch.Tensor | None = None + negative_prompt_embeds: torch.Tensor | None = None + negative_prompt_embeds_mask: torch.Tensor | None = None + + # ── Latent state (mutated every step by step_scheduler) ── + latents: torch.Tensor | None = None + + # ── Timestep schedule (set once by prepare_encode) ── + timesteps: torch.Tensor | list[torch.Tensor] | None = None + step_index: int = 0 + + # ── Per-request scheduler instance (set once by prepare_encode) ── + scheduler: Any | None = None + + # ── CFG config (set once by prepare_encode) ── + do_true_cfg: bool = False + guidance: torch.Tensor | None = None + + # ── Spatial / sequence metadata (set once by prepare_encode) ── + img_shapes: list | None = None + txt_seq_lens: list[int] | None = None + negative_txt_seq_lens: list[int] | None = None + + # Pipeline-specific extras. Keep model-private fields here unless they + # become part of the shared step-execution contract. + # For example: Wan condition tensors / masks, or Bagel KV contexts. + extra: dict[str, Any] = field(default_factory=dict) + + # ── Properties ── + + @property + def current_timestep(self) -> torch.Tensor | None: + if self.timesteps is None: + return None + if self.step_index >= self.total_steps: + return None + if isinstance(self.timesteps, torch.Tensor): + if self.timesteps.ndim == 0: + return self.timesteps + return self.timesteps[self.step_index] + return self.timesteps[self.step_index] + + @property + def total_steps(self) -> int: + if self.timesteps is None: + return 0 + if isinstance(self.timesteps, torch.Tensor): + if self.timesteps.ndim == 0: + return 1 + return int(self.timesteps.shape[0]) + return len(self.timesteps) + + @property + def denoise_completed(self) -> bool: + total_steps = self.total_steps + if total_steps == 0: + return False + return self.step_index >= total_steps + + @property + def new_request(self) -> bool: + # TODO: this is only an approximation for current stepwise mode. + # A real "new request" signal should eventually come from scheduler/runner state transitions. + return self.step_index == 0 or self.timesteps is None + + +@dataclass +class RunnerOutput: + """Output of a single denoising step for a request. + + NOTE: `latents` may be None when returned through IPC to avoid + serialization overhead. The actual latents are kept in Worker's + _request_state_cache. + """ + + req_id: str + step_index: int | None = None + finished: bool = False + result: DiffusionOutput | None = None From 6901ba426e28b1bdd90bde626ae064f7a55b2b28 Mon Sep 17 00:00:00 2001 From: Junhong Liu Date: Fri, 20 Mar 2026 19:33:52 +0800 Subject: [PATCH 2/2] [Fixbug][Perf] Qwen3-omni: code predictor with re-prefill + SDPA and eliminate decode hot-path CPU round-trips (#2012) Signed-off-by: Junhong Liu <98734602+LJH-LBJ@users.noreply.github.com> --- .../models/qwen3_omni/qwen3_omni.py | 52 +- .../qwen3_omni_moe_code_predictor_mtp.py | 561 ++++++++---------- .../qwen3_omni/qwen3_omni_moe_talker.py | 113 ++-- vllm_omni/worker/gpu_model_runner.py | 9 +- 4 files changed, 310 insertions(+), 425 deletions(-) diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py index ec928195a..21f0185aa 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py @@ -159,6 +159,12 @@ class Qwen3OmniMoeForConditionalGeneration( # suppress tokens by setting their probability to ~1e-9 (finite very small) self.suppressed_tokens = self._get_talker_suppressed_tokens() self.requires_raw_input_tokens = True + # Keys that should stay on GPU in model_intermediate_buffer to avoid CPU↔GPU round-trips + self.gpu_resident_buffer_keys: set[str] = { + "last_talker_hidden", + "trailing_text_hidden", + "tts_pad_embed_projected", + } elif self.model_stage == "code2wav": self.enable_update_additional_information = True @@ -223,14 +229,16 @@ class Qwen3OmniMoeForConditionalGeneration( # ==================== Forward Pass ==================== def _get_talker_suppressed_tokens(self): - return [ - i - for i in range( - self.config.talker_config.text_config.vocab_size - 1024, - self.config.talker_config.text_config.vocab_size, - ) - if i != self.config.talker_config.codec_eos_token_id - ] + """Return a boolean mask on GPU for suppressed token positions.""" + vocab_size = self.config.talker_config.text_config.vocab_size + mask = torch.zeros(vocab_size, dtype=torch.bool) + start = vocab_size - 1024 + eos_id = self.config.talker_config.codec_eos_token_id + for i in range(start, vocab_size): + if i != eos_id: + mask[i] = True + # Will be moved to the correct device on first use + return mask def get_mrope_input_positions( self, @@ -578,7 +586,7 @@ class Qwen3OmniMoeForConditionalGeneration( Postprocess the talker hidden states. """ update_dict = {} - update_dict["last_talker_hidden"] = hidden_states[-1, :].detach().to("cpu").contiguous() + update_dict["last_talker_hidden"] = hidden_states[-1, :].detach() return update_dict def talker_preprocess(self, input_ids: torch.Tensor, input_embeds: torch.Tensor, **info_dict: dict): @@ -632,9 +640,11 @@ class Qwen3OmniMoeForConditionalGeneration( if inputs_embeds.shape[-1] == 2048: inputs_embeds = self.text_projection(inputs_embeds) code_predictor_codes, summed_embeddings = self.talker.code_predictor_forward( - input_ids, inputs_embeds.clone(), last_talker_hidden=last_talker_hidden + input_ids, inputs_embeds, last_talker_hidden=last_talker_hidden ) - inputs_embeds = summed_embeddings.clone() + # summed_embeddings is [B, seq_len, H] (3D) while text_step is [B, H] (2D). + # Flatten to 2D first to avoid wrong broadcasting: [B,1,H]+[B,H] → [B,B,H] + inputs_embeds = summed_embeddings.reshape(-1, self.talker_config.text_config.hidden_size) inputs_embeds = (inputs_embeds + text_step).reshape(-1, self.talker_config.text_config.hidden_size) return inputs_embeds, code_predictor_codes.squeeze(-1) @@ -758,7 +768,7 @@ class Qwen3OmniMoeForConditionalGeneration( # compatible with old shape [1,S,D] rem_tail = trailing_text_hidden.squeeze(0) if rem_tail.shape[0] > 0: - update_dict["trailing_text_hidden"] = rem_tail.detach().to("cpu").contiguous() + update_dict["trailing_text_hidden"] = rem_tail.detach() # Also persist projected tts_pad for decode fallback if needed if isinstance(tts_pad_thinker, torch.Tensor): pad_in = tts_pad_thinker @@ -767,7 +777,7 @@ class Qwen3OmniMoeForConditionalGeneration( if pad_in.ndim == 1: pad_in = pad_in.view(1, 1, -1) pad_proj = self.talker.text_projection(pad_in.to(self._module_device(self.talker))) - update_dict["tts_pad_embed_projected"] = pad_proj.detach().to("cpu").contiguous() + update_dict["tts_pad_embed_projected"] = pad_proj.detach() except Exception: pass self._talker_cache_thinker_decode_embeds(info_dict, update_dict) @@ -926,7 +936,7 @@ class Qwen3OmniMoeForConditionalGeneration( if isinstance(q_tail, torch.Tensor) and q_tail.numel() > 0: use_vec = q_tail[0:1, :] new_q_tail = ( - q_tail[1:, :].detach().to("cpu").contiguous() + q_tail[1:, :].detach() if q_tail.shape[0] > 1 else self.tts_pad_embed.to(input_embeds.device, dtype=input_embeds.dtype) ) @@ -1122,16 +1132,10 @@ class Qwen3OmniMoeForConditionalGeneration( # implemented by assigning their logits to log(1e-9). if getattr(self, "model_stage", None) == "talker" and isinstance(logits, torch.Tensor): - try: - logits_cpu = logits.cpu() - logits_cpu[:, self.suppressed_tokens] = -1e9 - logits = logits_cpu.to(logits.device) - except Exception as e: - print(f"Error in logits suppression: {e}") - print(f"logits.shape: {logits.shape}") - print(f"self.suppressed_tokens: {self.suppressed_tokens}") - raise e - logits[:, self.suppressed_tokens] = -1e9 + # Move mask to device once (lazy), then reuse every step + if self.suppressed_tokens.device != logits.device: + self.suppressed_tokens = self.suppressed_tokens.to(logits.device) + logits.masked_fill_(self.suppressed_tokens.unsqueeze(0), -1e9) return logits def sample( diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_code_predictor_mtp.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_code_predictor_mtp.py index fc7402890..b1f0b2237 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_code_predictor_mtp.py +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_code_predictor_mtp.py @@ -1,27 +1,16 @@ -"""Qwen3-Omni Code Predictor with MTP (Multi-Token Prediction) support. +"""Qwen3-Omni Code Predictor -- optimized re-prefill, no KV cache. -This module implements the code predictor component for Qwen3-Omni talker models. - -The code predictor generates residual RVQ (Residual Vector Quantization) codes -autoregressively, predicting layers 1 to N based on layer-0 codes from the talker. +* SDPA attention (F.scaled_dot_product_attention) with native GQA support +* Per-call embedding buffer to avoid cross-request aliasing +* Pre-allocated position_ids (read-only, safe to persist) +* torch.compile on inner transformer by default +* Inline sampling (top-k + top-p) -- no custom op overhead """ -from collections import namedtuple -from typing import Any - import torch import torch.nn as nn import torch.nn.functional as F -from transformers import Cache, PretrainedConfig -from transformers.generation.logits_process import ( - LogitsProcessorList, - TopKLogitsWarper, - TopPLogitsWarper, -) -from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS -from vllm.compilation.decorators import support_torch_compile -from vllm.config import CacheConfig, ModelConfig, VllmConfig, get_current_vllm_config -from vllm.forward_context import get_forward_context +from vllm.config import VllmConfig from vllm.logger import init_logger from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import ( @@ -32,47 +21,53 @@ from vllm.model_executor.layers.linear import ( from vllm.model_executor.layers.quantization import QuantizationConfig from vllm.model_executor.layers.rotary_embedding import get_rope from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding -from vllm.utils.torch_utils import direct_register_custom_op + +from vllm_omni.platforms import current_omni_platform logger = init_logger(__name__) -# ============================================================================ -# Code Predictor Attention Layer -# ============================================================================ + +# =================================================================== +# Standalone Attention (SDPA, no KV cache, no HF backend fallback) +# =================================================================== class Qwen3OmniCodePredictorAttention(nn.Module): - """Multi-head self-attention for code predictor with vLLM optimization.""" + """Multi-head self-attention for code predictor. + + Uses ``F.scaled_dot_product_attention`` directly. No KV cache -- the code + predictor always re-prefills the full (short) sequence each AR step. + + Input : [B, seq_len, hidden_size] + Output: [B, seq_len, hidden_size] + """ def __init__( self, config, - layer_idx: int, - vllm_config: VllmConfig = None, quant_config: QuantizationConfig | None = None, prefix: str = "", ): super().__init__() - - self.num_heads = config.code_predictor_config.num_attention_heads - self.num_key_value_heads = config.code_predictor_config.num_key_value_heads + cp_cfg = config.code_predictor_config + self.num_heads = cp_cfg.num_attention_heads + self.num_kv_heads = cp_cfg.num_key_value_heads self.head_dim = getattr( - config.code_predictor_config, + cp_cfg, "head_dim", - config.code_predictor_config.hidden_size // config.code_predictor_config.num_attention_heads, + cp_cfg.hidden_size // cp_cfg.num_attention_heads, ) - self.hidden_size = config.code_predictor_config.hidden_size - - if self.num_heads % self.num_key_value_heads != 0: - raise ValueError("num_attention_heads must be divisible by num_key_value_heads") - - self.num_key_value_groups = self.num_heads // self.num_key_value_heads + self.hidden_size = cp_cfg.hidden_size + self.q_size = self.num_heads * self.head_dim + self.kv_size = self.num_kv_heads * self.head_dim + self.scaling = self.head_dim**-0.5 + self._use_gqa = self.num_kv_heads != self.num_heads self.qkv_proj = QKVParallelLinear( hidden_size=self.hidden_size, head_size=self.head_dim, total_num_heads=self.num_heads, - total_num_kv_heads=self.num_key_value_heads, + total_num_kv_heads=self.num_kv_heads, bias=False, quant_config=quant_config, prefix=f"{prefix}.qkv_proj", @@ -88,120 +83,55 @@ class Qwen3OmniCodePredictorAttention(nn.Module): ) self.rotary_emb = get_rope( self.head_dim, - max_position=config.code_predictor_config.max_position_embeddings, + max_position=cp_cfg.max_position_embeddings, rope_parameters=None, dual_chunk_attention_config=None, ) - - self.q_size = self.num_heads * self.head_dim - self.kv_size = self.num_key_value_heads * self.head_dim - - # Query/Key normalization - self.q_norm = RMSNorm(self.head_dim, eps=config.code_predictor_config.rms_norm_eps) - self.k_norm = RMSNorm(self.head_dim, eps=config.code_predictor_config.rms_norm_eps) - self.is_causal = True - self.config = config - - self.attention_backends = ["flash_attention_2", "xformers", "eager", "sdpa"] - cudagraph_mode = get_current_vllm_config().compilation_config.cudagraph_mode - if "flash_attention_2" in ALL_ATTENTION_FUNCTIONS and cudagraph_mode.has_full_cudagraphs(): - logger.warning( - f"CUDAGraphMode.{cudagraph_mode.name} is currently not supported " - f"with flash attention for Qwen3-Omni talker MTP." - f"removing flash attention from attention_backends" - ) - self.attention_backends.remove("flash_attention_2") + self.q_norm = RMSNorm(self.head_dim, eps=cp_cfg.rms_norm_eps) + self.k_norm = RMSNorm(self.head_dim, eps=cp_cfg.rms_norm_eps) def forward( self, hidden_states: torch.Tensor, - past_key_values: Cache | None = None, - cache_position: torch.LongTensor | None = None, - use_cache: bool = False, - position_ids: torch.LongTensor | None = None, + position_ids: torch.Tensor, ) -> torch.Tensor: bsz, seq_len, _ = hidden_states.shape - qkv, _ = self.qkv_proj(hidden_states) + # Flatten to 2-D so vLLM rotary_emb gets [num_tokens, size] + qkv, _ = self.qkv_proj(hidden_states.reshape(bsz * seq_len, -1)) q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) - # Reshape for attention - q = q.reshape(bsz, seq_len, self.num_heads, self.head_dim) - k = k.reshape(bsz, seq_len, self.num_key_value_heads, self.head_dim) - v = v.reshape(bsz, seq_len, self.num_key_value_heads, self.head_dim) - - # Apply normalization - q = self.q_norm(q).contiguous() - k = self.k_norm(k).contiguous() - q = q.reshape(-1, self.q_size) - k = k.reshape(-1, self.kv_size) - - # Apply RoPE + # QK-norm -> RoPE (both 2-D) + q = self.q_norm(q.view(-1, self.num_heads, self.head_dim)).view(q.shape) + k = self.k_norm(k.view(-1, self.num_kv_heads, self.head_dim)).view(k.shape) q, k = self.rotary_emb(position_ids, q, k) - # Reshape for attention - q = q.reshape(bsz, seq_len, self.num_heads, self.head_dim) - k = k.reshape(bsz, seq_len, self.num_key_value_heads, self.head_dim) + # [B, heads, seq, head_dim] + q = q.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2) + k = k.view(bsz, seq_len, self.num_kv_heads, self.head_dim).transpose(1, 2) + v = v.view(bsz, seq_len, self.num_kv_heads, self.head_dim).transpose(1, 2) - v_heads = v.transpose(1, 2).contiguous() - q_heads = q.transpose(1, 2).contiguous() - k_heads = k.transpose(1, 2).contiguous() + attn_out = F.scaled_dot_product_attention( + q, + k, + v, + scale=self.scaling, + is_causal=True, + enable_gqa=self._use_gqa, + ) - if past_key_values is not None: - sin, cos = self.rotary_emb.get_cos_sin(seq_len) - # sin and cos are specific to RoPE models; cache_position needed for the static cache - cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} - k_heads, v_heads = past_key_values.update(k_heads, v_heads, self.layer_idx, cache_kwargs) - - # Try attention backends in order of preference, with runtime error handling - # This handles cases where the backend is registered but not actually available - attn_output = None - last_error = None - - for backend_name in self.attention_backends: - if backend_name not in ALL_ATTENTION_FUNCTIONS: - continue - - try: - attention_interface = ALL_ATTENTION_FUNCTIONS[backend_name] - attn_output, _ = attention_interface( - self, - q_heads, - k_heads, - v_heads, - None, - dropout=0.0 if not self.training else getattr(self, "attention_dropout", 0.0), - scaling=self.head_dim**-0.5, - sliding_window=None, - use_cache=use_cache, - position_ids=position_ids[:seq_len].unsqueeze(0), - output_hidden_states=True, - output_attentions=False, - ) - break - except (ValueError, ImportError, RuntimeError, AttributeError) as e: - # Store error and try next backend - last_error = e - continue - - if attn_output is None: - raise RuntimeError( - f"All attention backends failed. Last error: {last_error}. " - "Please install flash-attn, or ensure PyTorch's scaled_dot_product_attention is available." - ) - attn_output = attn_output.reshape(*(hidden_states.shape[:-1]), -1).contiguous() - - attn_output, _ = self.o_proj(attn_output) - return attn_output + attn_out = attn_out.transpose(1, 2).reshape(bsz * seq_len, -1) + output, _ = self.o_proj(attn_out) + return output.view(bsz, seq_len, -1) -# ============================================================================ -# Code Predictor MLP Layer -# ============================================================================ +# =================================================================== +# MLP +# =================================================================== class Qwen3OmniCodePredictorMLP(nn.Module): - """Feed-forward network for code predictor with fused gate/up projection.""" + """SiLU-gated MLP for code predictor.""" def __init__( self, @@ -221,7 +151,6 @@ class Qwen3OmniCodePredictorMLP(nn.Module): prefix=f"{prefix}.gate_up_proj", disable_tp=True, ) - self.down_proj = RowParallelLinear( input_size=intermediate_size, output_size=hidden_size, @@ -238,35 +167,23 @@ class Qwen3OmniCodePredictorMLP(nn.Module): return down -# ============================================================================ -# MTP Layer (Multi-Token Prediction Layer) -# ============================================================================ +# =================================================================== +# Decoder Layer +# =================================================================== -class Qwen3OmniCodePredictorMTPLayer(nn.Module): - """MTP layer for speculative decoding - predicts next residual code layer.""" +class Qwen3OmniCodePredictorDecoderLayer(nn.Module): + """Transformer decoder layer (SDPA, no KV cache).""" def __init__( self, - config: PretrainedConfig, - prefix: str, - model_config: ModelConfig, - layer_idx: int, - cache_config: CacheConfig | None = None, + config, quant_config: QuantizationConfig | None = None, + prefix: str = "", ) -> None: super().__init__() - self.layer_idx = layer_idx - self.config = config - self.self_attn = Qwen3OmniCodePredictorAttention( config, - layer_idx, - vllm_config=type( - "VllmConfig", - (), - {"cache_config": cache_config, "quant_config": quant_config, "model_config": model_config}, - )(), quant_config=quant_config, prefix=f"{prefix}.self_attn", ) @@ -275,260 +192,258 @@ class Qwen3OmniCodePredictorMTPLayer(nn.Module): quant_config=quant_config, prefix=f"{prefix}.mlp", ) - self.input_layernorm = RMSNorm( - config.code_predictor_config.hidden_size, eps=config.code_predictor_config.rms_norm_eps - ) - self.post_attention_layernorm = RMSNorm( - config.code_predictor_config.hidden_size, eps=config.code_predictor_config.rms_norm_eps - ) + cp_cfg = config.code_predictor_config + self.input_layernorm = RMSNorm(cp_cfg.hidden_size, eps=cp_cfg.rms_norm_eps) + self.post_attention_layernorm = RMSNorm(cp_cfg.hidden_size, eps=cp_cfg.rms_norm_eps) - def mtp_block( + def forward( self, hidden_states: torch.Tensor, - past_key_values: Cache | None = None, - cache_position: torch.LongTensor | None = None, - use_cache: bool = False, - position_ids: torch.LongTensor | None = None, + position_ids: torch.Tensor, ) -> torch.Tensor: - # Self-attention with residual residual = hidden_states hidden_states = self.input_layernorm(hidden_states) - hidden_states = self.self_attn(hidden_states, past_key_values, cache_position, use_cache, position_ids) + hidden_states = self.self_attn(hidden_states, position_ids) hidden_states = residual + hidden_states - # MLP with residual residual = hidden_states hidden_states = self.post_attention_layernorm(hidden_states) hidden_states = self.mlp(hidden_states) hidden_states = residual + hidden_states - return hidden_states -class Qwen3OmniCodePredictorBaseModel(nn.Module): - """ - Base model for code predictor - matches HF Qwen3OmniMoeTalkerCodePredictorModel structure. +# =================================================================== +# Base Transformer Model (re-prefill, no KV cache) +# =================================================================== - This is a simple transformer that processes inputs_embeds and outputs hidden states. + +class Qwen3OmniCodePredictorBaseModel(nn.Module): + """Inner transformer for code predictor. + + Signature: ``forward(inputs_embeds, position_ids) -> hidden_states`` + -- plain Tensor in, plain Tensor out (no namedtuple). """ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() config = vllm_config.model_config.hf_config.code_predictor_config - self.config = config - self.vocab_size = config.vocab_size - self.num_code_groups = config.num_code_groups - # Codec embeddings (for layers 1-num_code_groups-1) self.codec_embedding = nn.ModuleList( - [ - VocabParallelEmbedding( - config.vocab_size, - config.hidden_size, - ) - for _ in range(config.num_code_groups - 1) - ] + [VocabParallelEmbedding(config.vocab_size, config.hidden_size) for _ in range(config.num_code_groups - 1)] ) - # Decoder layers self.layers = nn.ModuleList( [ - Qwen3OmniCodePredictorMTPLayer( + Qwen3OmniCodePredictorDecoderLayer( vllm_config.model_config.hf_config, - f"{prefix}.layers.{idx}", - model_config=vllm_config.model_config, - layer_idx=idx, - cache_config=vllm_config.cache_config, quant_config=vllm_config.quant_config, + prefix=f"{prefix}.layers.{idx}", ) for idx in range(config.num_hidden_layers) ] ) - self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) def forward( self, inputs_embeds: torch.Tensor, - attention_mask: torch.Tensor | None = None, - position_ids: torch.LongTensor | None = None, - past_key_values: Any | None = None, - use_cache: bool | None = False, - cache_position: torch.LongTensor | None = None, - **kwargs: Any, - ) -> Any: - """ - Forward pass matching HF structure. - - Args: - inputs_embeds: [batch, seq_len, hidden_size] - position_ids: Optional position IDs tensor - past_key_values: Optional cached key-value pairs - use_cache: Whether to use cache - cache_position: Optional cache position tensor - **kwargs: Additional keyword arguments - - Returns: - Named tuple with .last_hidden_state and .past_key_values attributes - """ - batch_size, seq_len, _ = inputs_embeds.shape - # Forward through decoder layers + position_ids: torch.Tensor, + ) -> torch.Tensor: hidden_states = inputs_embeds - for layer in self.layers: - hidden_states = layer.mtp_block(hidden_states, past_key_values, cache_position, use_cache, position_ids) - - # Final norm + hidden_states = layer(hidden_states, position_ids) hidden_states = self.norm(hidden_states) - - # Return in HF-compatible format - Output = namedtuple("Output", ["last_hidden_state", "past_key_values"]) - return Output(last_hidden_state=hidden_states, past_key_values=None) # [batch, num_code_groups-1, hidden_size] - - def get_input_embeddings(self): - """Return codec embeddings for HF compatibility.""" - return self.codec_embedding + return hidden_states -def code_predictor_sample( - logits: torch.Tensor, - layer_name: str, -) -> torch.Tensor: - forward_context = get_forward_context() - self = forward_context.no_compile_layers[layer_name] - logits = self.logits_processors(None, logits[:, -1]) - probs = F.softmax(logits, dim=-1) - code = torch.multinomial(probs.squeeze(1), num_samples=1) # [batch, 1] - return code +# =================================================================== +# Code Predictor Wrapper (optimized re-prefill, persistent buffers) +# =================================================================== -def code_predictor_sample_fake( - logits: torch.Tensor, - layer_name: str, -) -> torch.Tensor: - return torch.empty((logits.shape[0], 1), dtype=torch.int64, device=logits.device) - - -direct_register_custom_op( - op_name="qwen3_omni_code_predictor_sample", - op_func=code_predictor_sample, - fake_impl=code_predictor_sample_fake, -) - - -@support_torch_compile class Qwen3OmniMoeTalkerCodePredictor(nn.Module): - """ - Code predictor wrapper matching HF structure. + """Optimized code predictor -- re-prefill approach, no KV cache. - Structure: - - self.model: Qwen3OmniCodePredictorBaseModel (transformer) - - self.lm_head: ModuleList of output heads + Each AR step forwards the full growing sequence (len 2 -> num_code_groups+1) + through the transformer. The extra O(T^2) FLOPs are negligible for + short sequences, and this avoids all KV-cache management overhead. + + Optimizations: + 1. Per-call embedding buffer -- avoids cross-request aliasing. + 2. Pre-allocated position_ids -- no torch.arange per step. + 3. Cached module references -- bypass ModuleList indexing. + 4. torch.compile on inner transformer. + 5. Inline sampling (top-k + top-p) -- no custom op overhead. """ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() - talker_code_predictor_config = vllm_config.model_config.hf_config + config = vllm_config.model_config.hf_config + self.config = config self.quant_config = vllm_config.quant_config self.prefix = prefix - self.config = talker_code_predictor_config - self.vocab_size = self.config.code_predictor_config.vocab_size - self.num_code_groups = self.config.code_predictor_config.num_code_groups + self.num_code_groups = config.code_predictor_config.num_code_groups + self._hidden_size = config.code_predictor_config.hidden_size - # Base transformer model (matches HF structure) - self.model = Qwen3OmniCodePredictorBaseModel(vllm_config=vllm_config, prefix=prefix) + self.model = Qwen3OmniCodePredictorBaseModel( + vllm_config=vllm_config, + prefix=prefix, + ) - # Output heads for each residual layer (1-num_layers-1) + # One lm_head per residual layer (layers 1 .. G-1) self.lm_head = nn.ModuleList( [ nn.Linear( - self.config.code_predictor_config.hidden_size, - self.config.code_predictor_config.vocab_size, + config.code_predictor_config.hidden_size, + config.code_predictor_config.vocab_size, bias=False, ) for _ in range(self.num_code_groups - 1) ] ) - self.logits_processors = LogitsProcessorList( - [ - TopKLogitsWarper(top_k=50), - TopPLogitsWarper(top_p=0.8), - ] - ) - compilation_config = get_current_vllm_config().compilation_config - if prefix in compilation_config.static_forward_context: - raise ValueError(f"Duplicate layer name: {prefix}") - compilation_config.static_forward_context[prefix] = self - self.layer_name = prefix + self.set_sampling_params() + # Lazily initialised position ids (read-only, safe to persist) + self._pos_ids: torch.Tensor | None = None + + # Cached plain-list refs (set once) + self._lm_heads: list | None = None + self._codec_embeds: list | None = None + + # Model forward (optionally compiled) + self._model_fwd: object | None = None + + def set_sampling_params(self, top_k: int = 50, top_p: float = 0.8): + """Configure sampling parameters to maintain consistency with previous implementation.""" + self._top_k = top_k + self._top_p = top_p + logger.debug(f"Sampling parameters updated: top_k={top_k}, top_p={top_p}s") + + # ------------------------------------------------------------------ + # Lazy-init helpers + # ------------------------------------------------------------------ + + def _ensure_pos_ids(self, device: torch.device) -> None: + if self._pos_ids is not None and self._pos_ids.device == device: + return + max_seq = self.num_code_groups + 1 + self._pos_ids = torch.arange(max_seq, dtype=torch.long, device=device) + + def _ensure_cached_refs(self) -> None: + if self._lm_heads is not None: + return + self._lm_heads = list(self.lm_head) + self._codec_embeds = list(self.model.codec_embedding) + + def _ensure_model_fwd(self) -> None: + if self._model_fwd is not None: + return + if current_omni_platform.supports_torch_inductor(): + # torch.compile fuses the 5-layer transformer's small kernels, + # reducing BF16 intermediate round-trips and improving precision. + # mode="default" avoids Inductor's own CUDA graph capture so it + # doesn't conflict with the outer CUDAGraphWrapper. + self._model_fwd = torch.compile( + self.model.forward, + mode="default", + dynamic=True, + ) + logger.info("code_predictor: torch.compile enabled (mode=default)") + else: + self._model_fwd = self.model.forward + logger.info("code_predictor: using eager mode (no torch.compile)") + + # ------------------------------------------------------------------ + # Forward -- re-prefill + inline sampling + # ------------------------------------------------------------------ + + @torch.inference_mode() def forward( self, layer0_code: torch.Tensor, layer0_embed: torch.Tensor, last_talker_hidden: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]: - """ - Forward pass for code predictor. + """Predict residual codebooks 1..G-1 autoregressively via re-prefill. Args: - layer0_code: - Code index for code-group (layer) 0. - Shape: [batch_size, 1], dtype typically int64. - - last_talker_hidden: - - Shape: [batch_size, hidden_size]. + layer0_code: [bsz, 1] int64 + layer0_embed: [bsz, 1, hidden_size] + last_talker_hidden: [bsz, 1, hidden_size] Returns: - pos_all_layers: - Predicted codes for all code groups, including `layer0_code`. - Shape: [batch_size, num_code_groups, 1]. - - current_input: - The final input embedding sequence after appending embeddings of all - predicted codes (one token per predicted layer). - Shape: [batch_size, num_code_groups + 2, hidden_size]. + all_codes: [bsz, num_code_groups, 1] + proj_buf: [bsz, num_code_groups + 1, hidden_size] + pos 0 = last_talker_hidden (NOT a codec embed) + pos 1 = layer0_embed + pos 2.. = codec_embedding[i](predicted_code_i) """ - pos_codes = [layer0_code] # Start with layer 0: [batch, 1] - try: - current_input = torch.cat([last_talker_hidden, layer0_embed], dim=1) # [batch, 2, hidden_size] - except Exception as e: - print(f"Error in current_input: {e}") - print(f"last_talker_hidden shape: {last_talker_hidden.shape}") - print(f"prev_embed shape: {layer0_embed.shape}") - raise e - batch_size = current_input.shape[0] + bsz = int(layer0_code.shape[0]) + device = layer0_code.device + dtype = last_talker_hidden.dtype + num_groups = self.num_code_groups - # Predict all residual layers (layers 1 to num_code_groups-1) autoregressively - for layer_idx in range(self.num_code_groups - 1): - seq_len = layer_idx + 2 - # Compute position_ids dynamically to avoid torch.compile specializing batch_size - position_ids = torch.arange(seq_len, device=current_input.device, dtype=torch.int64).repeat(batch_size) - # Forward through code_predictor model - outputs = self.model( - inputs_embeds=current_input, - attention_mask=None, - position_ids=position_ids, - past_key_values=None, - use_cache=False, - cache_position=None, - ) - hidden_state = outputs.last_hidden_state # [batch, 2, hidden_size] + # Lazy init (read-only caches only) + self._ensure_pos_ids(device) + self._ensure_model_fwd() + self._ensure_cached_refs() - # Use the corresponding lm_head for this layer - logits = self.lm_head[layer_idx](hidden_state[:, -1:, :]) - code = torch.ops.vllm.qwen3_omni_code_predictor_sample(logits, self.layer_name) - pos_codes.append(code) - # Update prev_embed for next layer (if not last layer) - # layer_idx=0 predicts layer 1, embed with codec_embedding[1] - new_embed = self.model.codec_embedding[layer_idx](code) # [batch, 1, hidden_size] - current_input = torch.cat([current_input, new_embed], dim=1) # [batch, 3~n, hidden_size] - pos_all_layers = torch.stack(pos_codes, dim=1) # [batch, num_code_groups, 1] - return pos_all_layers, current_input + # Allocate proj_buf locally each call to avoid cross-call aliasing + max_seq = num_groups + 1 + proj_buf = torch.zeros(bsz, max_seq, self._hidden_size, dtype=dtype, device=device) + pos_ids = self._pos_ids + model_fwd = self._model_fwd + lm_heads = self._lm_heads + codec_embeds = self._codec_embeds + + # Output codes + all_codes = torch.empty(bsz, num_groups, 1, dtype=torch.int64, device=device) + all_codes[:, 0] = layer0_code + + # Fill buffer positions 0 & 1 + proj_buf[:bsz, 0:1, :] = last_talker_hidden + proj_buf[:bsz, 1:2, :] = layer0_embed + + # Autoregressive loop: predict layers 1..G-1 + for step in range(1, num_groups): + seq_len = step + 1 + projected = proj_buf[:bsz, :seq_len, :] + step_pos_ids = pos_ids[:seq_len] if bsz == 1 else pos_ids[:seq_len].repeat(bsz) + assert step_pos_ids.shape[0] == bsz * seq_len + + hidden_out = model_fwd(projected, step_pos_ids) + + # Inline sampling: top-k -> top-p -> softmax -> multinomial + logits = lm_heads[step - 1](hidden_out[:, -1, :]) # [bsz, vocab] + if self._top_k > 0: + topk_vals, _ = logits.topk(self._top_k, dim=-1) + logits = logits.masked_fill(logits < topk_vals[:, -1:], float("-inf")) + if self._top_p < 1.0: + sorted_logits, sorted_idx = logits.sort(dim=-1, descending=True) + cumulative_probs = F.softmax(sorted_logits, dim=-1).cumsum(dim=-1) + # Remove tokens with cumulative probability above top_p + remove_mask = cumulative_probs - F.softmax(sorted_logits, dim=-1) >= self._top_p + sorted_logits[remove_mask] = float("-inf") + logits = sorted_logits.scatter(1, sorted_idx, sorted_logits) + probs = F.softmax(logits, dim=-1) + code = torch.multinomial(probs, num_samples=1) # [bsz, 1] + + all_codes[:, step] = code + + # Embed predicted code -> next buffer position + new_embed = codec_embeds[step - 1](code) # [batch, 1, hidden_size] + proj_buf[:bsz, step + 1 : step + 2, :] = new_embed + + return all_codes, proj_buf[:bsz] + + # ------------------------------------------------------------------ + # Weight loading + # ------------------------------------------------------------------ def load_weights(self, weights: list[tuple[str, torch.Tensor]]) -> set[str]: """Load weights with mapping for fused QKV and gate_up projections. diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py index 3362ae17f..9f5c20cf7 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py @@ -129,111 +129,70 @@ class Qwen3OmniMoeTalkerForConditionalGeneration( self.code_predictor = Qwen3OmniMoeTalkerCodePredictor( vllm_config=vllm_config, prefix=maybe_prefix(prefix, "code_predictor") ) - max_batch_size = max( - vllm_config.scheduler_config.max_num_seqs, vllm_config.compilation_config.max_cudagraph_capture_size - ) - self.layer0_embed_buffer = torch.zeros( - (max_batch_size, 1, self.config.text_config.hidden_size), - dtype=vllm_config.model_config.dtype, - ) def code_predictor_forward( self, input_ids: torch.Tensor, inputs_embeds: torch.Tensor | None = None, *, - temperature: float = 1.0, - top_k: int = 50, # Match transformers default - top_p: float = 0.8, # Match transformers default - generation_steps: int | None = None, last_talker_hidden: torch.Tensor | None = None, **_: object, ) -> tuple[torch.Tensor, torch.Tensor]: - """ - Generate full RVQ codec codes for the provided sequence. + """Generate full RVQ codes + summed embeddings (single-loop, no KV cache). - The code predictor consumes the layer-0 codec codes produced by the talker - alongside the talker's hidden states, and autoregressively predicts the remaining - residual layers (to num_codec_groups). + The code predictor uses re-prefill: each AR step re-forwards the full + (short) sequence through the transformer. The returned ``proj_buf`` + already contains all codec embeddings at positions 1..G, + so summed_embeddings = proj_buf[:, 1:, :].sum(dim=1) — no second + loop or re-embedding needed. Returns: - tuple containing: - - residual_codes: A tensor of shape [batch, num_code_groups, seq_len] containing - the complete set of codec codes - - summed_embeddings: A tensor of shape [batch, seq_len, hidden_size] - Sum of all layer embeddings at each position (like Transformers) + result_codes: [batch, num_code_groups, seq_len] + summed_embeddings: [batch, seq_len, hidden_size] """ if input_ids is None: - raise ValueError("`input_ids` containing layer-0 codec codes must be provided.") + raise ValueError("`input_ids` (layer-0 codes) must be provided.") if inputs_embeds is None: - raise ValueError("`inputs_embeds` containing talker hidden states must be provided.") + raise ValueError("`inputs_embeds` (talker hidden states) must be provided.") if inputs_embeds.ndim == 2: inputs_embeds = inputs_embeds.unsqueeze(0) if input_ids.ndim == 1: input_ids = input_ids.unsqueeze(0) - # Ensure the tensors are contiguous for the autoregressive sampling loop - inputs_embeds = inputs_embeds.contiguous() - input_ids = input_ids.contiguous() - - # Generate full codec codes using MTP - # This will be the parallel prediction implementation batch_size, seq_len = input_ids.shape + device = input_ids.device + embed_fn = self.language_model.model.codec_embedding + hidden_size = self.config.code_predictor_config.hidden_size - # For now, use sequential generation (TODO: implement parallel) - # Result will be [batch, num_code_groups, seq_len] - # - all_codes_per_position will collect [batch, num_code_groups, 1] for each position - all_codes_per_position = [] - middle_hidden_states = [] # Collect hidden states for each position - - # Generate residual layers for each position - for pos in range(seq_len): - layer0_code = input_ids[:, pos : pos + 1] # [batch, 1] - - # Initial input: [last_talker_hidden, layer0_embed] - layer0_embed = self.embed_input_ids(layer0_code) - self.layer0_embed_buffer[:batch_size].copy_(layer0_embed) - pos_all_layers, current_input = self.code_predictor( - layer0_code, self.layer0_embed_buffer[:batch_size], last_talker_hidden - ) - - # Stack all layers for this position: [batch, num_code_groups, 1] - all_codes_per_position.append(pos_all_layers) - middle_hidden_states.append(current_input[:, 2:-1, :]) - - # Concatenate across positions: [batch, num_code_groups, seq_len] - result_codes = torch.cat(all_codes_per_position, dim=2) - - # Build summed embeddings for each position (like Transformers) - # This combines layer-0 embed, mid layers hidden states, and last layer embed - all_summed_embeddings = [] + result_codes = torch.empty( + batch_size, + self.num_code_groups, + seq_len, + dtype=torch.int64, + device=device, + ) + summed_embeddings = torch.empty( + batch_size, + seq_len, + hidden_size, + dtype=inputs_embeds.dtype, + device=device, + ) for pos in range(seq_len): - # Layer 0 embedding - layer0_code = result_codes[:, 0, pos : pos + 1] # [batch, 1] - layer0_embed = self.embed_input_ids(layer0_code) # [batch, 1, hidden_size] + layer0_code = input_ids[:, pos : pos + 1] + layer0_embed = embed_fn(layer0_code) - # mid layers hidden states (from CodePredictor) - mid_residual_hiddens = middle_hidden_states[pos] # [batch, num_code_groups-2, hidden_size] - mid_list = list(mid_residual_hiddens.split(1, dim=1)) - - # last layer embedding - last_layer_code = result_codes[:, -1, pos : pos + 1] # [batch, 1] - last_residual_hidden = self.code_predictor.model.codec_embedding[-1](last_layer_code) - - # Concatenate all layers: [batch, num_code_groups, hidden_size] - pos_codec_hiddens = torch.cat( - [layer0_embed] + mid_list + [last_residual_hidden], - dim=1, + pos_all_layers, proj_buf = self.code_predictor( + layer0_code, + layer0_embed, + last_talker_hidden, ) - # Sum across layers: [batch, 1, hidden_size] (like Transformers) - pos_summed = pos_codec_hiddens.sum(dim=1, keepdim=True) - all_summed_embeddings.append(pos_summed) - - # Concatenate across positions: [batch, seq_len, hidden_size] - summed_embeddings = torch.cat(all_summed_embeddings, dim=1).squeeze(1) + result_codes[:, :, pos : pos + 1] = pos_all_layers + # proj_buf layout: [0]=talker_hidden, [1..G]=codec embeds + summed_embeddings[:, pos, :] = proj_buf[:, 1:, :].sum(dim=1) return result_codes, summed_embeddings diff --git a/vllm_omni/worker/gpu_model_runner.py b/vllm_omni/worker/gpu_model_runner.py index 9e9674458..ddc6c1190 100644 --- a/vllm_omni/worker/gpu_model_runner.py +++ b/vllm_omni/worker/gpu_model_runner.py @@ -1362,10 +1362,17 @@ class OmniGPUModelRunner(GPUModelRunner): req_state = self.requests.get(req_id) if req_state is None: return + # Check if the model declares keys that should stay on GPU + gpu_keys: set[str] = set() + if hasattr(self, "model") and hasattr(self.model, "gpu_resident_buffer_keys"): + gpu_keys = self.model.gpu_resident_buffer_keys existing = self.model_intermediate_buffer.setdefault(req_id, {}) for k, v in upd.items(): if isinstance(v, torch.Tensor): - existing[k] = v.detach().to("cpu").contiguous() + if k in gpu_keys: + existing[k] = v.detach().clone() + else: + existing[k] = v.detach().to("cpu").contiguous() elif isinstance(v, list): existing[k] = [ (item.detach().to("cpu").contiguous() if isinstance(item, torch.Tensor) else item) for item in v