File size: 867 Bytes
5c93746 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 | """
Communication module for distributed inference pipeline.
This module provides abstractions for distributed communication operations,
model data transfer, and buffer management in the StreamDiffusionV2 pipeline.
"""
from .distributed_communicator import DistributedCommunicator
from .model_data_transfer import ModelDataTransfer
from .buffer_manager import BufferManager
from .data_containers import LatentData, KVCacheData, CommunicationConfig
from .kv_cache_manager import KVCacheManager
from .utils import CommunicationTags, init_distributed, setup_logging, compute_balanced_split
__all__ = [
'DistributedCommunicator',
'ModelDataTransfer',
'BufferManager',
'LatentData',
'KVCacheData',
'CommunicationConfig',
'KVCacheManager',
'CommunicationTags',
'init_distributed',
'setup_logging',
'compute_balanced_split'
]
|