nemo_automodel.components.moe.uccl_ep.buffer

UCCLBuffer: a DeepEP-compatible Buffer backed by UCCL-EP.

This module re-exports the canonical Buffer implementation under the UCCLBuffer alias expected by nemo_automodel, with automatic intranode detection.

Module Contents

Classes

Name	Description
`EventHandle`	-
`UCCLBuffer`	Buffer subclass that auto-detects intranode mode.

Data

__all__

API

class nemo_automodel.components.moe.uccl_ep.buffer.EventHandle()

class nemo_automodel.components.moe.uccl_ep.buffer.UCCLBuffer(
    group,
    num_nvl_bytes: int = 0,
    num_rdma_bytes: int = 0,
    low_latency_mode: bool = False,
    num_qps_per_rank: int = 24,
    allow_nvlink_for_low_latency_mode: bool = True,
    allow_mnnvl: bool = False,
    explicitly_destroy: bool = False,
    is_intranode: bool = False
)

Bases: Buffer

Buffer subclass that auto-detects intranode mode.

When all EP ranks fit on a single node (group_size <= LOCAL_WORLD_SIZE), RDMA is disabled and only NVLink is used, avoiding RDMA MR registration failures on single-node setups.

nemo_automodel.components.moe.uccl_ep.buffer.__all__ = ['UCCLBuffer', 'Buffer', 'EventOverlap', 'EventHandle']