26 lines
801 B
Python
26 lines
801 B
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""AR-Diffusion Engine (AR-Diffusion).
|
|
|
|
The AR-Diffusion engine: a ``DiffusionEngine`` subclass that adds engine-level KV
|
|
cache management for autoregressive / chunked diffusion models. Selected via
|
|
``OmniDiffusionConfig.engine_backend = "ar_diffusion"``.
|
|
"""
|
|
|
|
from vllm_omni.experimental.ar_diffusion.capability import (
|
|
ARDiffusionCrossAttentionKVSpec,
|
|
ARDiffusionKVBranchSpec,
|
|
ARDiffusionKVCacheSpec,
|
|
SupportsARDiffusionPipeline,
|
|
SupportsARDiffusionWarmup,
|
|
)
|
|
from vllm_omni.experimental.ar_diffusion.engine import ARDiffusionEngine
|
|
|
|
__all__ = [
|
|
"ARDiffusionKVBranchSpec",
|
|
"ARDiffusionCrossAttentionKVSpec",
|
|
"ARDiffusionEngine",
|
|
"ARDiffusionKVCacheSpec",
|
|
"SupportsARDiffusionPipeline",
|
|
"SupportsARDiffusionWarmup",
|
|
]
|