257 lines
11 KiB
Python
257 lines
11 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
|
|
"""Build native vLLM configuration views for the Omni diffusion runtime."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass, field
|
|
from typing import Any, cast
|
|
|
|
import torch
|
|
from vllm.config import CompilationConfig, DeviceConfig, KVTransferConfig, VllmConfig
|
|
from vllm.transformers_utils.config import get_hf_text_config
|
|
|
|
from vllm_omni.diffusion.data import OmniDiffusionConfig
|
|
from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
|
|
from vllm_omni.platforms import current_omni_platform
|
|
|
|
|
|
def resolve_diffusion_max_model_len(od_config: OmniDiffusionConfig) -> int:
|
|
"""Resolve the native KV admission bound for a diffusion stage."""
|
|
|
|
configured = getattr(od_config, "max_model_len", None)
|
|
if configured is not None:
|
|
configured = int(configured)
|
|
if configured == -1:
|
|
configured = None
|
|
elif configured <= 0:
|
|
raise ValueError(f"max_model_len must be positive for Scheduler-managed Diffusion KV, got {configured}")
|
|
else:
|
|
return configured
|
|
|
|
# ``max_model_len`` keeps the native vLLM meaning: the per-sequence KV
|
|
# admission ceiling consumed by ``KVCacheManager``. HunyuanImage3 exposes
|
|
# that ceiling directly as ``max_position_embeddings`` because its text,
|
|
# reference-image, and generated-image tokens share one Transformer
|
|
# sequence. Other diffusion models commonly express their serving limits
|
|
# through resolution, duration, or frame-count constraints instead. When
|
|
# paged KV support expands to those models, the model integration must
|
|
# derive ``max_model_len`` from its serving limits, while model-owned
|
|
# preprocessing validates each request and derives its matching
|
|
# ``seq_len``. Text-encoder length fields must not be mistaken for the DiT
|
|
# KV limit.
|
|
hf_config = getattr(od_config, "tf_model_config", None)
|
|
for field_name in ("max_position_embeddings", "max_sequence_length", "seq_length"):
|
|
value = getattr(hf_config, field_name, None)
|
|
if value is None and isinstance(hf_config, dict):
|
|
value = hf_config.get(field_name)
|
|
if value is not None and int(value) > 0:
|
|
return int(value)
|
|
|
|
mode = getattr(getattr(od_config, "diffusion_kv_mode", None), "value", None)
|
|
if mode == "paged_scheduler":
|
|
raise ValueError(
|
|
"paged_scheduler Diffusion KV requires max_model_len or a model config "
|
|
"with max_position_embeddings/max_sequence_length/seq_length"
|
|
)
|
|
# Dense execution does not consume the native KV admission bound, but the
|
|
# shared Worker VllmConfig still requires a positive placeholder.
|
|
return 1
|
|
|
|
|
|
@dataclass
|
|
class _DiffusionVllmModelConfig:
|
|
"""Private adapter for the native ``VllmConfig.model_config`` surface."""
|
|
|
|
model: str | None
|
|
dtype: torch.dtype
|
|
max_model_len: int
|
|
original_max_model_len: int | None
|
|
runner_type: str = "generate"
|
|
quantization: str | None = None
|
|
quantization_config: Any | None = None
|
|
hf_config: Any | None = None
|
|
hf_text_config: Any | None = None
|
|
multimodal_config: Any | None = None
|
|
enforce_eager: bool = False
|
|
disable_cascade_attn: bool = False
|
|
enable_return_routed_experts: bool = False
|
|
enable_prompt_embeds: bool = False
|
|
logits_processors: tuple[str, ...] = ()
|
|
use_mla: bool = False
|
|
is_moe: bool = False
|
|
# Native vLLM attention layers/builders read these fields directly. The
|
|
# diffusion model does not have a full vLLM ModelConfig, so keep the
|
|
# narrow compatibility surface here instead of manufacturing a second
|
|
# model configuration object in the Worker backend.
|
|
is_mm_prefix_lm: bool = False
|
|
rswa_window: int | None = None
|
|
_attention_num_heads: int | None = field(default=None, init=False, repr=False)
|
|
_attention_num_kv_heads: int | None = field(default=None, init=False, repr=False)
|
|
_attention_head_size: int | None = field(default=None, init=False, repr=False)
|
|
|
|
# Needed for models that bundle things like LogitsProcessors, e.g., SenseNova
|
|
head_dtype: torch.dtype | None = None
|
|
|
|
@property
|
|
def is_quantized(self) -> bool:
|
|
return self.quantization is not None
|
|
|
|
def is_nvfp4_quantized(self) -> bool:
|
|
return self.quantization == "modelopt_fp4"
|
|
|
|
def set_attention_geometry(
|
|
self,
|
|
*,
|
|
num_heads: int,
|
|
num_kv_heads: int,
|
|
head_size: int,
|
|
) -> None:
|
|
"""Record rank-local geometry needed by native metadata builders.
|
|
|
|
Diffusion attention modules are loaded outside vLLM's model loader, so
|
|
their dimensions are not available when this lightweight config is
|
|
created. The paged Worker backend fills them in after model loading
|
|
and before calling ``init_attn_backend``.
|
|
"""
|
|
|
|
values = {
|
|
"num_heads": num_heads,
|
|
"num_kv_heads": num_kv_heads,
|
|
"head_size": head_size,
|
|
}
|
|
invalid = {name: value for name, value in values.items() if type(value) is not int or value <= 0}
|
|
if invalid:
|
|
raise ValueError(f"Diffusion attention geometry must contain positive integers: {invalid!r}")
|
|
self._attention_num_heads = num_heads
|
|
self._attention_num_kv_heads = num_kv_heads
|
|
self._attention_head_size = head_size
|
|
|
|
def _require_attention_geometry(self) -> tuple[int, int, int]:
|
|
geometry = (
|
|
self._attention_num_heads,
|
|
self._attention_num_kv_heads,
|
|
self._attention_head_size,
|
|
)
|
|
if any(value is None for value in geometry):
|
|
raise RuntimeError("Diffusion attention geometry is unavailable before paged cache layer registration")
|
|
num_heads, num_kv_heads, head_size = geometry
|
|
assert num_heads is not None and num_kv_heads is not None and head_size is not None
|
|
return num_heads, num_kv_heads, head_size
|
|
|
|
def get_num_attention_heads(self, parallel_config: Any) -> int:
|
|
del parallel_config
|
|
return self._require_attention_geometry()[0]
|
|
|
|
def get_num_kv_heads(self, parallel_config: Any) -> int:
|
|
del parallel_config
|
|
return self._require_attention_geometry()[1]
|
|
|
|
def get_head_size(self) -> int:
|
|
return self._require_attention_geometry()[2]
|
|
|
|
def get_total_num_kv_heads(self) -> int:
|
|
assert self.hf_text_config is not None
|
|
return self.hf_text_config.num_key_value_heads
|
|
|
|
@property
|
|
def is_diffusion(self) -> bool:
|
|
# This config is also installed as the current VllmConfig for the
|
|
# existing Omni diffusion runtime, which does not use vLLM's native
|
|
# diffusion runner. Preserve that established backend-selection flag.
|
|
return False
|
|
|
|
|
|
def _make_diffusion_vllm_model_config(od_config: OmniDiffusionConfig) -> _DiffusionVllmModelConfig:
|
|
quant_config = getattr(od_config, "quantization_config", None)
|
|
quantization = quant_config.get_name() if quant_config is not None and hasattr(quant_config, "get_name") else None
|
|
hf_config = getattr(od_config, "tf_model_config", None)
|
|
hf_text_config = get_hf_text_config(hf_config) if hasattr(hf_config, "get_text_config") else hf_config
|
|
configured_max_model_len = getattr(od_config, "max_model_len", None)
|
|
max_model_len = resolve_diffusion_max_model_len(od_config)
|
|
paged_scheduler = (
|
|
getattr(od_config, "diffusion_kv_mode", DiffusionKVCacheMode.DENSE_LEGACY)
|
|
is DiffusionKVCacheMode.PAGED_SCHEDULER
|
|
)
|
|
return _DiffusionVllmModelConfig(
|
|
model=od_config.model,
|
|
dtype=od_config.dtype,
|
|
max_model_len=max_model_len,
|
|
# Preserve the user's -1 sentinel only for the native paged-cache
|
|
# sizing flow. Dense keeps its established resolved model-config view.
|
|
original_max_model_len=configured_max_model_len if paged_scheduler else max_model_len,
|
|
quantization=quantization,
|
|
quantization_config=quant_config,
|
|
hf_config=hf_config,
|
|
hf_text_config=hf_text_config,
|
|
enforce_eager=getattr(od_config, "enforce_eager", False),
|
|
is_moe=bool(getattr(od_config, "is_moe", False)),
|
|
)
|
|
|
|
|
|
def create_base_diffusion_vllm_config(
|
|
device: torch.device,
|
|
od_config: OmniDiffusionConfig,
|
|
) -> VllmConfig:
|
|
"""Create the native vLLM 0.27 config used by diffusion."""
|
|
|
|
return VllmConfig(
|
|
compilation_config=CompilationConfig(),
|
|
device_config=DeviceConfig(device=device),
|
|
additional_config=od_config.additional_config,
|
|
)
|
|
|
|
|
|
def configure_diffusion_vllm_config(vllm_config: VllmConfig, od_config: OmniDiffusionConfig) -> VllmConfig:
|
|
"""Populate native model, parallel, cache, and Scheduler configuration."""
|
|
|
|
parallel_config = od_config.parallel_config
|
|
assert parallel_config.data_parallel_size is not None
|
|
vllm_config.parallel_config.tensor_parallel_size = parallel_config.tensor_parallel_size
|
|
vllm_config.parallel_config.data_parallel_size = parallel_config.data_parallel_size
|
|
if parallel_config.enable_expert_parallel and od_config.is_moe:
|
|
vllm_config.parallel_config.data_parallel_size = (
|
|
parallel_config.data_parallel_size * parallel_config.cfg_parallel_size
|
|
)
|
|
vllm_config.parallel_config.prefill_context_parallel_size = parallel_config.sequence_parallel_size
|
|
vllm_config.parallel_config.enable_expert_parallel = parallel_config.enable_expert_parallel
|
|
|
|
vllm_config.model_config = _make_diffusion_vllm_model_config(od_config) # type: ignore[assignment]
|
|
vllm_config.quant_config = od_config.quantization_config
|
|
vllm_config.profiler_config = od_config.profiler_config
|
|
if (
|
|
getattr(od_config, "diffusion_kv_mode", DiffusionKVCacheMode.DENSE_LEGACY)
|
|
is DiffusionKVCacheMode.PAGED_SCHEDULER
|
|
):
|
|
vllm_config.parallel_config.pipeline_parallel_size = parallel_config.pipeline_parallel_size
|
|
vllm_config.cache_config.gpu_memory_utilization = float(getattr(od_config, "gpu_memory_utilization", 0.9))
|
|
vllm_config.cache_config.kv_cache_memory_bytes = getattr(od_config, "kv_cache_memory_bytes", None)
|
|
vllm_config.cache_config.enable_prefix_caching = bool(getattr(od_config, "enable_prefix_caching", False))
|
|
max_rows_per_request = getattr(od_config, "diffusion_kv_max_rows_per_request", None)
|
|
if type(max_rows_per_request) is not int or max_rows_per_request <= 0:
|
|
raise ValueError("paged_scheduler requires a positive diffusion_kv_max_rows_per_request")
|
|
vllm_config.scheduler_config.max_num_seqs = int(od_config.max_num_seqs) * max_rows_per_request
|
|
max_num_batched_tokens = getattr(od_config, "max_num_batched_tokens", None)
|
|
if max_num_batched_tokens is not None:
|
|
vllm_config.scheduler_config.max_num_batched_tokens = int(max_num_batched_tokens)
|
|
current_omni_platform.configure_diffusion_vllm_config(vllm_config, od_config)
|
|
|
|
kv_transfer_config = getattr(od_config, "kv_transfer_config", None)
|
|
if kv_transfer_config is not None:
|
|
vllm_config.kv_transfer_config = cast(KVTransferConfig, kv_transfer_config)
|
|
return vllm_config
|
|
|
|
|
|
def create_diffusion_vllm_config(
|
|
device: torch.device,
|
|
od_config: OmniDiffusionConfig,
|
|
) -> VllmConfig:
|
|
"""Build the same fully populated native config in Engine and Workers."""
|
|
|
|
vllm_config = create_base_diffusion_vllm_config(
|
|
device,
|
|
od_config,
|
|
)
|
|
return configure_diffusion_vllm_config(vllm_config, od_config)
|