41 lines
1.9 KiB
Python
41 lines
1.9 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""The AR-Diffusion Engine (AR-Diffusion)."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from vllm.logger import init_logger
|
|
|
|
from vllm_omni.diffusion.diffusion_engine import DiffusionEngine
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
#: Import path of the runner the AR-Diffusion engine routes its workers to.
|
|
AR_DIFFUSION_MODEL_RUNNER_CLS = "vllm_omni.experimental.ar_diffusion.runner.ARDiffusionModelRunner"
|
|
|
|
|
|
class ARDiffusionEngine(DiffusionEngine):
|
|
"""AR-Diffusion engine with engine-level KV cache management.
|
|
|
|
AR-Diffusion serves autoregressive / chunked blockwise-causal diffusion models
|
|
(world models, AR-DiT) that materialize persistent attention KV. It reuses
|
|
vLLM's paged KV stack (``KVCacheManager`` / ``BlockPool`` / ``BlockTables``)
|
|
as a library, driven from the engine rather than hand-rolled inside each
|
|
model.
|
|
|
|
It is selected per model via ``OmniDiffusionConfig.engine_backend = "ar_diffusion"``
|
|
(resolved by :meth:`DiffusionEngine.make_engine`), so models that do not opt
|
|
in keep using the base ``DiffusionEngine`` unchanged.
|
|
|
|
Architecture note: in the multiproc setup the KV cache lives in the worker /
|
|
runner process (GPU side), co-located with the model and KV tensors — so the
|
|
actual KV *body* is :class:`~vllm_omni.experimental.ar_diffusion.kv_cache.manager.ARDiffusionKVCache`, owned
|
|
by :class:`~vllm_omni.experimental.ar_diffusion.runner.ARDiffusionModelRunner`. ``ARDiffusionEngine`` itself is the
|
|
thin selection / injection seam; it wires the AR-Diffusion executor → worker → runner
|
|
so capable pipelines run against runner-owned KV session state.
|
|
"""
|
|
|
|
#: Workers of this engine build the AR-Diffusion runner (unless the config
|
|
#: sets an explicit ``diffusion_model_runner_cls`` override). Declared on
|
|
#: the class so routing needs no od_config mutation at engine init.
|
|
default_diffusion_model_runner_cls = AR_DIFFUSION_MODEL_RUNNER_CLS
|