397 lines
15 KiB
Python
397 lines
15 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
|
|
from contextlib import nullcontext
|
|
from functools import cache
|
|
from typing import Any
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
from vllm.logger import init_logger
|
|
from vllm_ascend.platform import NPUPlatform
|
|
|
|
from vllm_omni.diffusion.attention.backends.registry import DiffusionAttentionBackendEnum
|
|
from vllm_omni.platforms.interface import OmniPlatform, OmniPlatformEnum
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
_DIFFUSION_PACKED_MODULES_MAPPING = {
|
|
"HunyuanImage3Pipeline": {
|
|
"experts": ["experts.0.gate_up_proj", "experts.0.down_proj"],
|
|
},
|
|
}
|
|
|
|
|
|
@cache
|
|
def _get_strict_ulysses_paged_backend() -> type:
|
|
"""Return an Ascend backend that bypasses vLLM PCP dispatch."""
|
|
|
|
from vllm_ascend.attention.attention_v1 import (
|
|
AscendAttentionBackend,
|
|
AscendAttentionBackendImpl,
|
|
AscendAttentionMetadataBuilder,
|
|
)
|
|
|
|
class AscendStrictUlyssesPagedBackend(AscendAttentionBackend):
|
|
@staticmethod
|
|
def get_impl_cls() -> type:
|
|
return AscendAttentionBackendImpl
|
|
|
|
@staticmethod
|
|
def get_builder_cls() -> type:
|
|
return AscendAttentionMetadataBuilder
|
|
|
|
return AscendStrictUlyssesPagedBackend
|
|
|
|
|
|
class NPUOmniPlatform(OmniPlatform, NPUPlatform):
|
|
"""NPU/Ascend implementation of OmniPlatform.
|
|
|
|
Inherits all NPU-specific implementations from vllm-ascend's NPUPlatform,
|
|
and adds Omni-specific interfaces from OmniPlatform.
|
|
"""
|
|
|
|
_omni_enum = OmniPlatformEnum.NPU
|
|
dist_backend: str = "hccl"
|
|
|
|
# conv2d convolution operator in the code2wav module of Qwen3-TTS not being able to run on Aclnn
|
|
def __init__(self) -> None:
|
|
from vllm_ascend.utils import adapt_patch
|
|
|
|
from vllm_omni.platforms.npu._310p import apply_patches as apply_310p_patches
|
|
from vllm_omni.platforms.npu.models.minicpmo_4_5_code2wav import (
|
|
apply_minicpmo_4_5_code2wav_patch,
|
|
)
|
|
from vllm_omni.platforms.npu.models.qwen3_tts import (
|
|
apply_qwen3_tts_patches,
|
|
)
|
|
from vllm_omni.platforms.npu.models.qwen3_tts_tokenizer_v2 import (
|
|
apply_qwen3_tts_tokenizer_v2_patch,
|
|
)
|
|
|
|
adapt_patch(is_global_patch=True)
|
|
apply_minicpmo_4_5_code2wav_patch()
|
|
apply_qwen3_tts_patches()
|
|
apply_qwen3_tts_tokenizer_v2_patch()
|
|
apply_310p_patches()
|
|
|
|
@classmethod
|
|
def set_device(cls, device: torch.device) -> None:
|
|
super().set_device(device)
|
|
|
|
# Register vllm_ascend custom ops (torch.ops._C_ascend.*).
|
|
from vllm_ascend.utils import enable_custom_op
|
|
|
|
enable_custom_op()
|
|
|
|
# Ascend quantized weights are converted from ND to FRACTAL_NZ
|
|
# after loading. Enable internal format so the NZ storage layout
|
|
# is preserved for fused NPU kernels.
|
|
torch.npu.config.allow_internal_format = True
|
|
|
|
@classmethod
|
|
def get_omni_ar_worker_cls(cls) -> str:
|
|
return "vllm_omni.platforms.npu.worker.npu_ar_worker.NPUARWorker"
|
|
|
|
@classmethod
|
|
def get_omni_generation_worker_cls(cls) -> str:
|
|
return "vllm_omni.platforms.npu.worker.npu_generation_worker.NPUGenerationWorker"
|
|
|
|
@classmethod
|
|
def init_diffusion_worker_vllm_config(cls, vllm_config: Any) -> None:
|
|
from vllm_ascend.ascend_config import init_ascend_config
|
|
from vllm_ascend.utils import adapt_patch
|
|
|
|
# Omni's custom DiffusionWorker does not pass through vLLM-Ascend's
|
|
# NPUWorker constructor, where worker-local patches are normally
|
|
# installed. In particular, AscendBlockTables needs the patched
|
|
# non-UVA buffer implementation on NPU.
|
|
adapt_patch()
|
|
init_ascend_config(vllm_config)
|
|
|
|
@classmethod
|
|
def configure_diffusion_vllm_config(cls, vllm_config: Any, od_config: Any) -> None:
|
|
"""Use the block geometry required by Ascend's native paged kernel."""
|
|
if getattr(od_config, "diffusion_kv_mode", None) is None:
|
|
return
|
|
from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
|
|
|
|
if od_config.diffusion_kv_mode is not DiffusionKVCacheMode.PAGED_SCHEDULER:
|
|
return
|
|
from vllm_ascend.attention.attention_v1 import AscendAttentionBackend
|
|
|
|
supported_sizes = [
|
|
size for size in AscendAttentionBackend.get_supported_kernel_block_sizes() if type(size) is int and size > 0
|
|
]
|
|
if not supported_sizes:
|
|
raise RuntimeError("Ascend paged attention did not expose an integer kernel block size")
|
|
# vLLM's generic default is 16, while the Ascend FIA backend stores
|
|
# cache pages as 128-token blocks. Set the Manager geometry
|
|
# before KV specs are collected so Scheduler and Worker agree.
|
|
vllm_config.cache_config.block_size = supported_sizes[0]
|
|
|
|
@classmethod
|
|
def requires_diffusion_paged_kv_prewrite(cls) -> bool:
|
|
"""Write the full K/V span once before piecewise FIA segments."""
|
|
|
|
return True
|
|
|
|
@classmethod
|
|
def get_diffusion_paged_kv_attn_backend(cls, attn_backend: type, *, ulysses_degree: int) -> type:
|
|
"""Keep strict Ulysses paged FIA out of vLLM's PCP implementation."""
|
|
|
|
del cls
|
|
if ulysses_degree <= 1:
|
|
return attn_backend
|
|
from vllm_ascend.attention.attention_v1 import AscendAttentionBackend
|
|
|
|
if not isinstance(attn_backend, type) or not issubclass(attn_backend, AscendAttentionBackend):
|
|
return attn_backend
|
|
return _get_strict_ulysses_paged_backend()
|
|
|
|
@classmethod
|
|
def get_diffusion_kv_block_tables_cls(cls) -> type:
|
|
from vllm_ascend.worker.v2.block_table import AscendBlockTables
|
|
|
|
return AscendBlockTables
|
|
|
|
@classmethod
|
|
def build_diffusion_kv_attn_metadata(cls, **kwargs: Any) -> dict[str, Any]:
|
|
"""Build the Ascend metadata required by the native NPU backend."""
|
|
from vllm_ascend.attention.attention_v1 import AscendAttentionState
|
|
from vllm_ascend.worker.v2.attn_utils import build_attn_metadata
|
|
|
|
kwargs = dict(kwargs)
|
|
seq_lens_cpu = kwargs.pop("seq_lens_cpu")
|
|
kwargs["seq_lens_np"] = seq_lens_cpu.detach().cpu().numpy()
|
|
# The diffusion adapter always supplies a paged cache and the current
|
|
# K/V write span. ChunkedPrefill is Ascend's cache-backed FIA state for
|
|
# both multi-token updates and single-token updates in this path.
|
|
kwargs["attn_state"] = AscendAttentionState.ChunkedPrefill
|
|
return build_attn_metadata(**kwargs)
|
|
|
|
@classmethod
|
|
def init_diffusion_model_runner_runtime(cls, vllm_config: Any, od_config: Any, device: torch.device) -> None:
|
|
from vllm_ascend.ascend_forward_context import set_mc2_mask, set_mc2_tokens_capacity
|
|
|
|
from vllm_omni.platforms.npu.models.minimax_h3 import (
|
|
apply_minimax_h3_qwen3vl_patch,
|
|
apply_minimax_h3_qwen3vl_sdpa_patch,
|
|
apply_minimax_h3_qwen3vl_swiglu_patch,
|
|
)
|
|
|
|
# These patches import the MiniMax encoder package, whose __init__ loads
|
|
# pipeline_minimax_h3 → diffusion.data. Doing that during platform
|
|
# construction races vllm_omni/__init__.py (patch before config) and
|
|
# closes a cycle through pipeline_registry → PI0_PIPELINE →
|
|
# DiffusionOutput. Apply them only after the platform exists, before
|
|
# the diffusion pipeline is loaded.
|
|
apply_minimax_h3_qwen3vl_patch()
|
|
apply_minimax_h3_qwen3vl_sdpa_patch()
|
|
apply_minimax_h3_qwen3vl_swiglu_patch()
|
|
set_mc2_tokens_capacity(vllm_config, od_config.max_num_seqs, 1)
|
|
set_mc2_mask(vllm_config, device)
|
|
|
|
@classmethod
|
|
def get_default_stage_config_path(cls) -> str:
|
|
return "vllm_omni/deploy"
|
|
|
|
@classmethod
|
|
def prepare_diffusion_op_runtime(cls, op_name: str, **kwargs: Any) -> None:
|
|
if op_name != "fused_moe":
|
|
return
|
|
|
|
from vllm_omni.platforms.npu.layers.fused_moe import prepare_fused_moe_runtime
|
|
|
|
prepare_fused_moe_runtime()
|
|
|
|
@classmethod
|
|
def register_additional_diffusion_fused_moe_hooks(cls, moe_runner: Any) -> None:
|
|
from vllm_omni.platforms.npu.layers.fused_moe import fused_moe_forward_context_pre_hook
|
|
|
|
moe_runner.register_forward_pre_hook(
|
|
fused_moe_forward_context_pre_hook,
|
|
with_kwargs=True,
|
|
)
|
|
|
|
@classmethod
|
|
def reset_diffusion_fused_moe_forward_context(cls) -> None:
|
|
from vllm_omni.platforms.npu.layers.fused_moe import reset_fused_moe_forward_context
|
|
|
|
reset_fused_moe_forward_context()
|
|
|
|
@classmethod
|
|
def get_diffusion_packed_modules_mapping(
|
|
cls,
|
|
model_class: type[nn.Module],
|
|
) -> dict[str, list[str]] | None:
|
|
return _DIFFUSION_PACKED_MODULES_MAPPING.get(model_class.__name__, None)
|
|
|
|
@classmethod
|
|
def get_diffusion_attn_backend_cls(
|
|
cls,
|
|
selected_backend: str | None,
|
|
head_size: int,
|
|
allow_trtllm_default: bool = True,
|
|
) -> str:
|
|
# NPU has no TRTLLM backend; accepted for signature parity, ignored.
|
|
from importlib.util import find_spec
|
|
|
|
if selected_backend is not None:
|
|
backend_upper = selected_backend.upper()
|
|
cls.validate_diffusion_attn_backend(backend_upper)
|
|
if backend_upper in ("FLASH_ATTN_HUB", "FLASH_ATTN_3_HUB"):
|
|
logger.warning(
|
|
"HuggingFace kernels-backed FlashAttention is "
|
|
"not supported on NPU. Falling back to local "
|
|
"FLASH_ATTN."
|
|
)
|
|
backend_upper = "FLASH_ATTN"
|
|
|
|
if backend_upper in ("FLASH_ATTN", "RAINFUSION_ATTN") and find_spec("mindiesd"):
|
|
# Eager-import mindiesd only for backends that actually reach
|
|
# mindiesd kernels: FLASH_ATTN directly, and RAINFUSION_ATTN
|
|
# via its dense FlashAttention fallback (used before
|
|
# start_step and on any layer without a sparsifiable video
|
|
# segment). Other backends (e.g. TORCH_SDPA) never touch
|
|
# mindiesd, so a broken optional install must not block them.
|
|
# CANN snapshots the custom-op registry at the first
|
|
# custom-op regInfo lookup in the process (e.g. a
|
|
# vllm-ascend custom op during model load/warmup). Import
|
|
# mindiesd here so its env.py prepends the mindiesd vendor
|
|
# dirs (aie_ascendc etc.) to ASCEND_CUSTOM_OPP_PATH before
|
|
# that snapshot; otherwise aclnnLaserAttention /
|
|
# FusedAttentionScore fail with EZ1001 "does not support
|
|
# opType" for the rest of the process.
|
|
import mindiesd # noqa: F401
|
|
|
|
backend = DiffusionAttentionBackendEnum[backend_upper]
|
|
logger.debug("Using diffusion attention backend '%s'", backend_upper)
|
|
return backend.get_path()
|
|
|
|
# Try FLASH_ATTN if mindiesd is available, otherwise fall back to SDPA
|
|
if find_spec("mindiesd"):
|
|
# Configure ASCEND_CUSTOM_OPP_PATH for mindiesd custom ops upon import
|
|
import mindiesd # noqa: F401
|
|
|
|
logger.debug("Defaulting to diffusion attention backend FLASH_ATTN")
|
|
return DiffusionAttentionBackendEnum.FLASH_ATTN.get_path()
|
|
|
|
logger.debug("Falling back to diffusion attention backend SDPA")
|
|
return DiffusionAttentionBackendEnum.TORCH_SDPA.get_path()
|
|
|
|
@classmethod
|
|
def supports_diffusion_dense_flash_attention(cls) -> bool:
|
|
"""Return whether MindIE-SD is installed for dense NPU FlashAttention."""
|
|
|
|
from importlib.util import find_spec
|
|
|
|
return find_spec("mindiesd") is not None
|
|
|
|
@classmethod
|
|
def supports_torch_inductor(cls) -> bool:
|
|
return False
|
|
|
|
@classmethod
|
|
def get_torch_device(cls, local_rank: int | None = None) -> torch.device:
|
|
if local_rank is None:
|
|
return torch.device("npu")
|
|
return torch.device("npu", local_rank)
|
|
|
|
@classmethod
|
|
def get_device_count(cls) -> int:
|
|
return torch.npu.device_count()
|
|
|
|
@classmethod
|
|
def get_device_version(cls) -> str | None:
|
|
return None
|
|
|
|
@classmethod
|
|
def synchronize(cls) -> None:
|
|
torch.npu.synchronize()
|
|
|
|
@classmethod
|
|
def record_device_event(cls) -> torch.Event | None:
|
|
"""Record a NPU event on the default stream to mark tensor readiness.
|
|
|
|
On NPU/Ascend with HCCL, distributed communication may use internal
|
|
streams not visible to the default stream. Synchronize the default
|
|
stream first so that HCCL results are written back before we record
|
|
the event, ensuring d2h_stream.wait_event() captures the complete
|
|
output data.
|
|
"""
|
|
try:
|
|
torch.npu.current_stream().synchronize()
|
|
# The async output worker uses the public ``torch.Stream`` API.
|
|
# With torch_npu 2.10, a native ``torch.npu.Event`` cannot be
|
|
# consumed by that wrapper stream (the reverse direction is
|
|
# supported), while ``torch.Event`` is dispatched to the active
|
|
# NPU backend and remains cross-stream compatible.
|
|
event = torch.Event()
|
|
event.record()
|
|
return event
|
|
except Exception:
|
|
logger.warning("Failed to record NPU event for cross-stream sync")
|
|
return None
|
|
|
|
@classmethod
|
|
def get_free_memory(cls, device: torch.device | None = None) -> int:
|
|
free, _ = torch.npu.mem_get_info(device)
|
|
return free
|
|
|
|
@classmethod
|
|
def get_device_memory(cls, device: torch.device | None = None) -> tuple[int, int]:
|
|
free, total = torch.npu.mem_get_info(device)
|
|
return free, total
|
|
|
|
@classmethod
|
|
def get_device_total_memory(cls, device_id: int = 0) -> int:
|
|
device_props = torch.npu.get_device_properties(device_id)
|
|
return device_props.total_memory
|
|
|
|
@classmethod
|
|
def create_autocast_context(cls, *, device_type, dtype, enabled=True):
|
|
if device_type != "npu":
|
|
return super().create_autocast_context(
|
|
device_type=device_type,
|
|
dtype=dtype,
|
|
enabled=enabled,
|
|
)
|
|
if not enabled:
|
|
return nullcontext()
|
|
|
|
# NPU-specific fallback
|
|
try:
|
|
return torch.npu.amp.autocast(dtype=dtype)
|
|
except (RuntimeError, TypeError, ValueError) as exc:
|
|
logger.warning("autocast unavailable for device_type=%s dtype=%s: %s", device_type, dtype, exc)
|
|
return nullcontext()
|
|
|
|
@classmethod
|
|
def get_profiler_cls(cls) -> str:
|
|
return "vllm_omni.platforms.npu.profiler.NPUTorchProfilerWrapper"
|
|
|
|
@classmethod
|
|
def get_graph_wrapper_cls(cls) -> type:
|
|
from vllm_ascend.compilation.acl_graph import ACLGraphWrapper
|
|
|
|
return ACLGraphWrapper
|
|
|
|
@classmethod
|
|
def set_forward_context(
|
|
cls,
|
|
attn_metadata,
|
|
vllm_config,
|
|
*,
|
|
cudagraph_runtime_mode,
|
|
batch_descriptor,
|
|
):
|
|
from vllm_ascend.ascend_forward_context import set_ascend_forward_context
|
|
|
|
return set_ascend_forward_context(
|
|
attn_metadata,
|
|
vllm_config,
|
|
aclgraph_runtime_mode=cudagraph_runtime_mode,
|
|
batch_descriptor=batch_descriptor,
|
|
)
|