Didan DengandJiangyun Zhu
ee2b78e826
Signed-off-by: Didan Deng <33117903+wtomin@users.noreply.github.com> Signed-off-by: zjy0516 <riverclouds.zhu@qq.com> Co-authored-by: Jiangyun Zhu <riverclouds.zhu@qq.com>
91 lines
2.8 KiB
Python
91 lines
2.8 KiB
Python
from contextlib import contextmanager
|
|
from dataclasses import dataclass
|
|
|
|
from vllm.config import VllmConfig
|
|
|
|
from vllm_omni.diffusion.attention.backends.abstract import (
|
|
AttentionMetadata,
|
|
)
|
|
from vllm_omni.diffusion.data import OmniDiffusionConfig
|
|
|
|
|
|
@dataclass
|
|
class ForwardContext:
|
|
"""
|
|
set forward context for diffusion models
|
|
"""
|
|
|
|
vllm_config: VllmConfig | None = None
|
|
omni_diffusion_config: OmniDiffusionConfig | None = None
|
|
attn_metadata: dict[str, AttentionMetadata] | list[dict[str, AttentionMetadata]] | None = None
|
|
split_text_embed_in_sp: bool = False
|
|
# whether to split the text embed in sequence parallel, if True, the text embed will be split in sequence parallel
|
|
|
|
def __post_init__(self):
|
|
pass
|
|
|
|
|
|
_forward_context: ForwardContext | None = None
|
|
|
|
|
|
def get_forward_context() -> ForwardContext:
|
|
"""Get the current forward context."""
|
|
assert _forward_context is not None, (
|
|
"Forward context is not set. Please use `set_forward_context` to set the forward context."
|
|
)
|
|
return _forward_context
|
|
|
|
|
|
def is_forward_context_available() -> bool:
|
|
return _forward_context is not None
|
|
|
|
|
|
def create_forward_context(
|
|
vllm_config: VllmConfig | None = None,
|
|
omni_diffusion_config: OmniDiffusionConfig | None = None,
|
|
attn_metadata: dict[str, AttentionMetadata] | list[dict[str, AttentionMetadata]] | None = None,
|
|
split_text_embed_in_sp: bool = False,
|
|
):
|
|
return ForwardContext(
|
|
vllm_config=vllm_config,
|
|
omni_diffusion_config=omni_diffusion_config,
|
|
attn_metadata=attn_metadata,
|
|
split_text_embed_in_sp=split_text_embed_in_sp,
|
|
)
|
|
|
|
|
|
@contextmanager
|
|
def override_forward_context(forward_context: ForwardContext | None):
|
|
"""A context manager that overrides the current forward context.
|
|
This is used to override the forward context for a specific
|
|
forward pass.
|
|
"""
|
|
global _forward_context
|
|
prev_context = _forward_context
|
|
_forward_context = forward_context
|
|
try:
|
|
yield
|
|
finally:
|
|
_forward_context = prev_context
|
|
|
|
|
|
@contextmanager
|
|
def set_forward_context(
|
|
vllm_config: VllmConfig | None = None,
|
|
omni_diffusion_config: OmniDiffusionConfig | None = None,
|
|
attn_metadata: dict[str, AttentionMetadata] | list[dict[str, AttentionMetadata]] | None = None,
|
|
split_text_embed_in_sp: bool = False,
|
|
):
|
|
"""A context manager that stores the current forward context,
|
|
can be attention metadata, split_text_embed_in_sp, etc.
|
|
Here we can inject common logic for every model forward pass.
|
|
"""
|
|
forward_context = create_forward_context(
|
|
vllm_config=vllm_config,
|
|
omni_diffusion_config=omni_diffusion_config,
|
|
attn_metadata=attn_metadata,
|
|
split_text_embed_in_sp=split_text_embed_in_sp,
|
|
)
|
|
with override_forward_context(forward_context):
|
|
yield
|