`core.inference.disaggregation.transfer_backends.nccl`#

Two-sided (NCCL) KV transfer backend for disaggregated prefill/decode.

Unlike the one-sided NIXL backend, both peers participate: the decode posts receives when the hand-off request arrives (begin_pull_blocks) and the prefill posts the matching sends when the coordinator’s SEND_KV names the decode instance (begin_push_blocks). Both sides enumerate the same reshard plan in the same deterministic order, so the point-to-point operations match by post order per peer pair. Data moves straight out of the prefill’s pinned blocks; there is no staging copy on the send side.

Module Contents#

Classes#

`NcclTransferHandle`	Pollable handle for one batched NCCL transfer.
`NcclTransferBackend`	Per-buffer NCCL transport over the default process group.

Functions#

`_make_copy`
`_kv_layout_from_meta`	Rebuild a peer’s KVShardLayout from its exported metadata.

Data#

logger

API#

core.inference.disaggregation.transfer_backends.nccl.logger#: ‘getLogger(…)’

class core.inference.disaggregation.transfer_backends.nccl.NcclTransferHandle( works: List[Any], keepalive: List[torch.Tensor], scatters: List[Any], )#

Pollable handle for one batched NCCL transfer.

Receives land in temporary contiguous buffers; on completion the handle runs its scatter closures once to place the data into the paged buffers. Send handles keep the gathered source slices alive until reaped.

Initialization

poll() → bool#: Return True if the transfer has settled, scattering received data into the paged buffers on first completion.

wait() → None#: Block until the transfer completes, then scatter.

_finish() → None#

core.inference.disaggregation.transfer_backends.nccl._make_copy(view: torch.Tensor, buf: torch.Tensor)#

core.inference.disaggregation.transfer_backends.nccl._kv_layout_from_meta( meta: Dict[str, Any], ) → megatron.core.inference.disaggregation.kv_reshard.KVShardLayout#: Rebuild a peer’s KVShardLayout from its exported metadata.

class core.inference.disaggregation.transfer_backends.nccl.NcclTransferBackend( agent_name: str, memory_buffer: torch.Tensor, expected_num_blocks: int, tp_size: Optional[int] = None, tp_rank: Optional[int] = None, num_kv_heads_global: Optional[int] = None, heads_per_partition: Optional[int] = None, head_dim: Optional[int] = None, tokens_per_block: Optional[int] = None, global_rank: Optional[int] = None, pp_size: Optional[int] = None, pp_rank: Optional[int] = None, num_layers_global: Optional[int] = None, layer_start: Optional[int] = None, layer_end: Optional[int] = None, ssm_layout: Optional[megatron.core.inference.disaggregation.ssm_reshard.SSMShardLayout] = None, ssm_state_kind: Optional[str] = None, )#

Per-buffer NCCL transport over the default process group.

Mirrors the NIXL backend’s construction and metadata schema so the hand-off layer treats the two interchangeably; only the transfer calls differ (two-sided matched send/recv instead of one-sided reads).