from .cuda_ipc import CudaRTLibrary, create_shared_buffer, free_shared_buffer
from .dlpack_utils import pack_strided_memory
from .mapping import Mapping
from .trtllm_ar import AllReduceFusionOp as AllReduceFusionOp
from .trtllm_ar import AllReduceFusionPattern as AllReduceFusionPattern
from .trtllm_ar import AllReduceStrategyConfig as AllReduceStrategyConfig
from .trtllm_ar import AllReduceStrategyType as AllReduceStrategyType
from .trtllm_ar import QuantizationSFLayout as QuantizationSFLayout
from .trtllm_ar import (
    compute_fp4_swizzled_layout_sf_size as compute_fp4_swizzled_layout_sf_size,
)
from .trtllm_ar import gen_trtllm_comm_module as gen_trtllm_comm_module
from .trtllm_ar import trtllm_allreduce_fusion as trtllm_allreduce_fusion
from .trtllm_ar import (
    trtllm_create_ipc_workspace_for_all_reduce as trtllm_create_ipc_workspace_for_all_reduce,
)
from .trtllm_ar import (
    trtllm_create_ipc_workspace_for_all_reduce_fusion as trtllm_create_ipc_workspace_for_all_reduce_fusion,
)
from .trtllm_ar import trtllm_custom_all_reduce as trtllm_custom_all_reduce
from .trtllm_ar import (
    trtllm_destroy_ipc_workspace_for_all_reduce as trtllm_destroy_ipc_workspace_for_all_reduce,
)
from .trtllm_ar import (
    trtllm_destroy_ipc_workspace_for_all_reduce_fusion as trtllm_destroy_ipc_workspace_for_all_reduce_fusion,
)
from .trtllm_ar import trtllm_lamport_initialize as trtllm_lamport_initialize
from .trtllm_ar import trtllm_lamport_initialize_all as trtllm_lamport_initialize_all
from .trtllm_ar import trtllm_moe_allreduce_fusion as trtllm_moe_allreduce_fusion
from .trtllm_ar import (
    trtllm_moe_finalize_allreduce_fusion as trtllm_moe_finalize_allreduce_fusion,
)
from .vllm_ar import all_reduce as vllm_all_reduce
from .vllm_ar import dispose as vllm_dispose
from .vllm_ar import gen_vllm_comm_module as gen_vllm_comm_module
from .vllm_ar import get_graph_buffer_ipc_meta as vllm_get_graph_buffer_ipc_meta
from .vllm_ar import init_custom_ar as vllm_init_custom_ar
from .vllm_ar import meta_size as vllm_meta_size
from .vllm_ar import register_buffer as vllm_register_buffer
from .vllm_ar import register_graph_buffers as vllm_register_graph_buffers
from .ulysses import UlyssesCommunicator as UlyssesCommunicator
from .ulysses import dispose_ulysses_a2a as dispose_ulysses_a2a
from .ulysses import gen_ulysses_a2a_module as gen_ulysses_a2a_module
from .ulysses import get_ulysses_a2a_module as get_ulysses_a2a_module
from .ulysses import init_ulysses_a2a as init_ulysses_a2a
from .ulysses import ulysses_a2a as ulysses_a2a
from .ulysses_topology import ULYSSES_BACKENDS as ULYSSES_BACKENDS
from .ulysses_topology import UlyssesBackendDecision as UlyssesBackendDecision
from .ulysses_topology import UlyssesBackendError as UlyssesBackendError
from .ulysses_topology import UlyssesRankTopology as UlyssesRankTopology
from .ulysses_topology import decide_ulysses_backend as decide_ulysses_backend
from .ulysses_topology import (
    probe_ulysses_rank_topology as probe_ulysses_rank_topology,
)
from .ulysses_topology import resolve_ulysses_backend as resolve_ulysses_backend

# Unified AllReduce Fusion API
from .allreduce import AllReduceFusionWorkspace as AllReduceFusionWorkspace
from .trtllm_mnnvl_ar import (
    MNNVLAllReduceFusionWorkspace as MNNVLAllReduceFusionWorkspace,
)
from .allreduce import TRTLLMAllReduceFusionWorkspace as TRTLLMAllReduceFusionWorkspace
from .allreduce import allreduce_fusion as allreduce_fusion
from .allreduce import (
    create_allreduce_fusion_workspace as create_allreduce_fusion_workspace,
)

# MNNVL A2A (Throughput Backend)
from .trtllm_moe_alltoall import MoeAlltoAll as MoeAlltoAll
from .trtllm_moe_alltoall import moe_a2a_active_rank_mask as moe_a2a_active_rank_mask
from .trtllm_moe_alltoall import moe_a2a_combine as moe_a2a_combine
from .trtllm_moe_alltoall import moe_a2a_dispatch as moe_a2a_dispatch
from .trtllm_moe_alltoall import moe_a2a_initialize as moe_a2a_initialize
from .trtllm_moe_alltoall import (
    moe_a2a_get_workspace_size_per_rank as moe_a2a_get_workspace_size_per_rank,
)
from .trtllm_moe_alltoall import (
    moe_a2a_sanitize_expert_ids as moe_a2a_sanitize_expert_ids,
)
from .trtllm_moe_alltoall import (
    moe_a2a_wrap_payload_tensor_in_workspace as moe_a2a_wrap_payload_tensor_in_workspace,
)

# DCP A2A (Decode Context Parallel Attention Reduction)
from .dcp_alltoall import decode_cp_a2a_alltoall as decode_cp_a2a_alltoall
from .dcp_alltoall import (
    decode_cp_a2a_allocate_mnnvl_workspace as decode_cp_a2a_allocate_mnnvl_workspace,
)
from .dcp_alltoall import decode_cp_a2a_init_workspace as decode_cp_a2a_init_workspace
from .dcp_alltoall import decode_cp_a2a_workspace_size as decode_cp_a2a_workspace_size

# from .mnnvl import MnnvlMemory, MnnvlMoe, MoEAlltoallInfo


def __getattr__(name: str):
    if name == "all_gather_matmul":
        from .all_gather_matmul import all_gather_matmul

        return all_gather_matmul
    if name == "quantized_all_reduce":
        from .quantized_allreduce import quantized_all_reduce

        return quantized_all_reduce
    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
