[Misc] Merge diffusion forward context (#582)

Signed-off-by: iwzbi <wzbi@zju.edu.cn>
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

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
 

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

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()

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:

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:
     """

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

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

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