From 026f613e5769fabcd70bf94a62d60ad9c617555e Mon Sep 17 00:00:00 2001 From: catcat <108673086+iwzbi@users.noreply.github.com> Date: Mon, 5 Jan 2026 22:49:03 +0800 Subject: [PATCH] [Misc] Merge diffusion forward context (#582) Signed-off-by: iwzbi Signed-off-by: Hongsheng Liu Co-authored-by: Hongsheng Liu --- .../attention/test_sequence_parallel.py | 4 +- vllm_omni/diffusion/attention/layer.py | 4 +- .../diffusion/attention/parallel/factory.py | 4 +- .../diffusion/attention/parallel/ring.py | 6 +-- vllm_omni/diffusion/data.py | 52 ------------------- .../qwen_image/qwen_image_transformer.py | 4 +- vllm_omni/diffusion/worker/gpu_worker.py | 8 +-- vllm_omni/diffusion/worker/npu/npu_worker.py | 9 ++-- 8 files changed, 16 insertions(+), 75 deletions(-) diff --git a/tests/diffusion/attention/test_sequence_parallel.py b/tests/diffusion/attention/test_sequence_parallel.py index 9601db441..4304de77d 100644 --- a/tests/diffusion/attention/test_sequence_parallel.py +++ b/tests/diffusion/attention/test_sequence_parallel.py @@ -12,13 +12,13 @@ from vllm_omni.diffusion.attention.layer import Attention from vllm_omni.diffusion.data import ( DiffusionParallelConfig, OmniDiffusionConfig, - set_current_omni_diffusion_config, ) from vllm_omni.diffusion.distributed.parallel_state import ( destroy_distributed_env, init_distributed_environment, initialize_model_parallel, ) +from vllm_omni.diffusion.forward_context import set_forward_context from vllm_omni.utils.platform_utils import detect_device_type device_type = detect_device_type() @@ -385,7 +385,7 @@ def ulysses_attention_on_test_model( ) # Set the config so Attention can access it - with set_current_omni_diffusion_config(od_config): + with set_forward_context(omni_diffusion_config=od_config): # Create model hidden_size = num_heads * head_size diff --git a/vllm_omni/diffusion/attention/layer.py b/vllm_omni/diffusion/attention/layer.py index e7fa4c66f..0d3bf8b7d 100644 --- a/vllm_omni/diffusion/attention/layer.py +++ b/vllm_omni/diffusion/attention/layer.py @@ -15,8 +15,8 @@ from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata from vllm_omni.diffusion.attention.parallel import build_parallel_attention_strategy from vllm_omni.diffusion.attention.parallel.ring import RingParallelAttention from vllm_omni.diffusion.attention.selector import get_attn_backend -from vllm_omni.diffusion.data import get_current_omni_diffusion_config from vllm_omni.diffusion.distributed.parallel_state import get_sp_group +from vllm_omni.diffusion.forward_context import get_forward_context logger = init_logger(__name__) @@ -58,7 +58,7 @@ class Attention(nn.Module): self.ring_runner = None try: - config = get_current_omni_diffusion_config() + config = get_forward_context().omni_diffusion_config self.backend_pref = config.attention_backend if config.parallel_config.ring_degree > 1: self.use_ring = True diff --git a/vllm_omni/diffusion/attention/parallel/factory.py b/vllm_omni/diffusion/attention/parallel/factory.py index 19e2c9ae1..2d5ef2317 100644 --- a/vllm_omni/diffusion/attention/parallel/factory.py +++ b/vllm_omni/diffusion/attention/parallel/factory.py @@ -6,8 +6,8 @@ from __future__ import annotations from vllm_omni.diffusion.attention.parallel.base import NoParallelAttention, ParallelAttentionStrategy from vllm_omni.diffusion.attention.parallel.ring import RingParallelAttention from vllm_omni.diffusion.attention.parallel.ulysses import UlyssesParallelAttention -from vllm_omni.diffusion.data import get_current_omni_diffusion_config from vllm_omni.diffusion.distributed.parallel_state import get_sequence_parallel_world_size, get_sp_group +from vllm_omni.diffusion.forward_context import get_forward_context def build_parallel_attention_strategy( @@ -24,7 +24,7 @@ def build_parallel_attention_strategy( and initialized process groups. """ try: - cfg = get_current_omni_diffusion_config() + cfg = get_forward_context().omni_diffusion_config p = cfg.parallel_config except Exception: return NoParallelAttention() diff --git a/vllm_omni/diffusion/attention/parallel/ring.py b/vllm_omni/diffusion/attention/parallel/ring.py index df145ac3b..2d63f40c0 100644 --- a/vllm_omni/diffusion/attention/parallel/ring.py +++ b/vllm_omni/diffusion/attention/parallel/ring.py @@ -16,10 +16,10 @@ from vllm_omni.diffusion.attention.parallel.base import ( ParallelAttentionContext, # ParallelAttentionStrategy, # Not used in type hint below currently ) +from vllm_omni.diffusion.distributed.group_coordinator import SequenceParallelGroupCoordinator # from vllm_omni.diffusion.attention.backends.ring_selector import AttnType # Already imported above -from vllm_omni.diffusion.data import get_current_omni_diffusion_config -from vllm_omni.diffusion.distributed.group_coordinator import SequenceParallelGroupCoordinator +from vllm_omni.diffusion.forward_context import get_forward_context if TYPE_CHECKING: from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata @@ -110,7 +110,7 @@ class RingParallelAttention: backend_pref = self.attn_backend_pref if backend_pref is None: try: - config = get_current_omni_diffusion_config() + config = get_forward_context().omni_diffusion_config # config might not have attention_backend attribute if not updated backend_pref = getattr(config, "attention_backend", None) except Exception: diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index 9beaf9212..195f75aea 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -5,9 +5,7 @@ import enum import os import random from collections.abc import Callable -from contextlib import contextmanager from dataclasses import dataclass, field, fields -from functools import lru_cache from typing import Any import torch @@ -458,56 +456,6 @@ class OmniDiffusionConfig: return cls(**kwargs) -_current_omni_diffusion_config: OmniDiffusionConfig | None = None -_current_prefix: str | None = None - - -@contextmanager -def set_current_omni_diffusion_config( - omni_diffusion_config: OmniDiffusionConfig, check_compile=False, prefix: str | None = None -): - """ - Temporarily set the current vLLM-Omni config. - Used during model initialization. - We save the current vLLM-Omni config in a global variable, - so that all modules can access it, e.g. custom ops - can access the vLLM-Omni config to determine how to dispatch. - """ - global _current_omni_diffusion_config, _current_prefix - old_omni_diffusion_config = _current_omni_diffusion_config - old_prefix = _current_prefix - # from vllm.compilation.counter import compilation_counter - - # num_models_seen = compilation_counter.num_models_seen - try: - _current_omni_diffusion_config = omni_diffusion_config - _current_prefix = prefix - yield - except Exception: - raise - else: - if check_compile: - raise RuntimeError("Compilation is not yet supported for OmniDiffusion") - finally: - _current_omni_diffusion_config = old_omni_diffusion_config - _current_prefix = old_prefix - # Clear the compilation config cache when context changes - get_cached_compilation_config.cache_clear() - - -@lru_cache(maxsize=1) -def get_cached_compilation_config(): - """Cache config to avoid repeated calls to get_current_omni_diffusion_config()""" - return get_current_omni_diffusion_config().compilation_config - - -def get_current_omni_diffusion_config() -> OmniDiffusionConfig: - if _current_omni_diffusion_config is None: - logger.warning_once("Current OmniDiffusionConfig is not set.") - return OmniDiffusionConfig() - return _current_omni_diffusion_config - - @dataclass class DiffusionOutput: """ diff --git a/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py b/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py index 1a52769ff..48e460ea8 100644 --- a/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py +++ b/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py @@ -24,7 +24,7 @@ from vllm_omni.diffusion.attention.backends.abstract import ( ) from vllm_omni.diffusion.attention.layer import Attention from vllm_omni.diffusion.cache.base import CachedTransformer -from vllm_omni.diffusion.data import OmniDiffusionConfig, get_current_omni_diffusion_config +from vllm_omni.diffusion.data import OmniDiffusionConfig from vllm_omni.diffusion.distributed.parallel_state import ( get_sequence_parallel_rank, get_sequence_parallel_world_size, @@ -362,7 +362,7 @@ class QwenImageCrossAttention(nn.Module): self.rope = RotaryEmbedding(is_neox_style=False) try: - config = get_current_omni_diffusion_config() + config = get_forward_context().omni_diffusion_config self.parallel_config = config.parallel_config except Exception: self.parallel_config = None diff --git a/vllm_omni/diffusion/worker/gpu_worker.py b/vllm_omni/diffusion/worker/gpu_worker.py index 64cdc1dc5..0f8a9b855 100644 --- a/vllm_omni/diffusion/worker/gpu_worker.py +++ b/vllm_omni/diffusion/worker/gpu_worker.py @@ -8,7 +8,7 @@ from contextlib import AbstractContextManager, nullcontext import torch import zmq -from vllm.config import LoadConfig, VllmConfig, set_current_vllm_config +from vllm.config import LoadConfig, VllmConfig from vllm.distributed.device_communicators.shm_broadcast import MessageQueue from vllm.logger import init_logger from vllm.utils.mem_utils import DeviceMemoryProfiler, GiB_bytes @@ -17,7 +17,6 @@ from vllm_omni.diffusion.cache.selector import get_cache_backend from vllm_omni.diffusion.data import ( DiffusionOutput, OmniDiffusionConfig, - set_current_omni_diffusion_config, ) from vllm_omni.diffusion.distributed.parallel_state import ( destroy_distributed_env, @@ -69,10 +68,7 @@ class GPUWorker: vllm_config.parallel_config.tensor_parallel_size = self.od_config.parallel_config.tensor_parallel_size vllm_config.parallel_config.data_parallel_size = self.od_config.parallel_config.data_parallel_size self.vllm_config = vllm_config - with ( - set_current_omni_diffusion_config(self.od_config), - set_current_vllm_config(vllm_config), - ): + with set_forward_context(vllm_config=vllm_config, omni_diffusion_config=self.od_config): init_distributed_environment(world_size=world_size, rank=rank) logger.info(f"Worker {self.rank}: Initialized device and distributed environment.") parallel_config = self.od_config.parallel_config diff --git a/vllm_omni/diffusion/worker/npu/npu_worker.py b/vllm_omni/diffusion/worker/npu/npu_worker.py index 6b269e66c..bfeb0d914 100644 --- a/vllm_omni/diffusion/worker/npu/npu_worker.py +++ b/vllm_omni/diffusion/worker/npu/npu_worker.py @@ -5,19 +5,19 @@ import os import time import torch -from vllm.config import LoadConfig, VllmConfig, set_current_vllm_config +from vllm.config import LoadConfig, VllmConfig from vllm.logger import init_logger from vllm.utils.mem_utils import DeviceMemoryProfiler, GiB_bytes from vllm_omni.diffusion.cache.selector import get_cache_backend from vllm_omni.diffusion.data import ( OmniDiffusionConfig, - set_current_omni_diffusion_config, ) from vllm_omni.diffusion.distributed.parallel_state import ( init_distributed_environment, initialize_model_parallel, ) +from vllm_omni.diffusion.forward_context import set_forward_context from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader from vllm_omni.diffusion.worker.gpu_worker import GPUWorker, WorkerProc @@ -49,10 +49,7 @@ class NPUWorker(GPUWorker): vllm_config.parallel_config.tensor_parallel_size = self.od_config.parallel_config.tensor_parallel_size vllm_config.parallel_config.data_parallel_size = self.od_config.parallel_config.data_parallel_size self.vllm_config = vllm_config - with ( - set_current_omni_diffusion_config(self.od_config), - set_current_vllm_config(vllm_config), - ): + with set_forward_context(vllm_config=vllm_config, omni_diffusion_config=self.od_config): init_distributed_environment(world_size=world_size, rank=rank) logger.info(f"Worker {self.rank}: Initialized device and distributed environment.") parallel_config = self.od_config.parallel_config