`core.inference.disaggregation.transfer_backends.nixl`#

Direct NIXL backend for disaggregated prefill/decode KV transfer.

Each rank registers its paged KV buffer once, exports NIXL peer metadata, and the decode side pulls source block ranges directly into its local KV blocks.

Backend selection belongs in transfer_backends.base and is supplied explicitly by the launcher.

Module Contents#

Classes#

`NixlPullHandle`	Pollable handle for one logical pull made of one or more NIXL transfers.
`_NixlAgentContext`	NIXL resources shared by the state buffers on one rank.
`NixlTransferBackend`	Per-buffer registration on a rank’s NIXL agent.

Functions#

`_detect_nixl_variant`	Return the installed distribution that provides the active NIXL API module.
`_validate_ucx_transport_config`	Configure safe UCX memory detection and reject host-only transports.
`_validate_nixl_cuda_support`	Reject a NIXL/UCX runtime that cannot safely register CUDA memory.

Data#

`logger`
`_POLL_INTERVAL_S`
`_POLL_TIMEOUT_S`

API#

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

core.inference.disaggregation.transfer_backends.nixl._detect_nixl_variant(api_module_name: str) → Optional[str]#: Return the installed distribution that provides the active NIXL API module.

core.inference.disaggregation.transfer_backends.nixl._POLL_INTERVAL_S#: 0.0005

core.inference.disaggregation.transfer_backends.nixl._POLL_TIMEOUT_S#: 30.0

core.inference.disaggregation.transfer_backends.nixl._validate_ucx_transport_config(memory_buffer: torch.Tensor) → None#: Configure safe UCX memory detection and reject host-only transports.

core.inference.disaggregation.transfer_backends.nixl._validate_nixl_cuda_support( agent: Any, memory_buffer: torch.Tensor, ) → None#: Reject a NIXL/UCX runtime that cannot safely register CUDA memory.

class core.inference.disaggregation.transfer_backends.nixl.NixlPullHandle#

Pollable handle for one logical pull made of one or more NIXL transfers.

agent: Any#: None

xfers: List[Any]#: None

contexts: List[str]#: None

submitted_at: float#: None

timeout_s: float#: None

done: bool#: False

error: Optional[str]#: None

poll() → bool#: Return True if every transfer has settled, without blocking.

wait() → None#: Block until the transfer completes; NIXL has no blocking wait, so poll with a short sleep to avoid monopolizing a CPU core.

class core.inference.disaggregation.transfer_backends.nixl._NixlAgentContext(agent_name: str)#

NIXL resources shared by the state buffers on one rank.

Initialization

acquire() → Any#: Retain the shared agent for one buffer registration.

release() → None#: Release one registration and drop the agent after the last user closes.

class core.inference.disaggregation.transfer_backends.nixl.NixlTransferBackend( 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, _shared_context: Optional[core.inference.disaggregation.transfer_backends.nixl._NixlAgentContext] = None, )#

Per-buffer registration on a rank’s NIXL agent.

Per-block transfers are descriptor ranges over that registration. Peer metadata is exchanged by the control plane and registered lazily on first pull.

Initialization

name#: ‘nixl’

new_registered_buffer(

**kwargs,

) → core.inference.disaggregation.transfer_backends.nixl.NixlTransferBackend#: Register another state buffer on this backend’s NIXL agent.

export_meta() → Dict[str, Any]#

Return JSON/msgpack-safe metadata for shipping to a decode peer.

Layout fields describe the scatter-gather address ranges needed to pull source blocks into decode-owned blocks.

_ensure_peer_registered( peer_meta: Dict[str, Any], ) → str#: Register the peer with NIXL on first use; return its agent id.

_validate_peer( meta: Dict[str, Any], src_block_ids: List[int], dst_block_ids: List[int], *, matched_layout: bool = False, ) → None#: Validate block mappings and physical transfer compatibility.

static _kv_layout_from_meta( meta: Dict[str, Any], ) → megatron.core.inference.disaggregation.kv_reshard.KVShardLayout#: Reconstruct a main-planner KV layout from peer wire metadata.

begin_pull_blocks( peer_meta: Any, src_block_ids: List[int], dst_block_ids: List[int], ) → core.inference.disaggregation.transfer_backends.nixl.NixlPullHandle#: Submit a pull and return a handle that can be polled later.

_begin_transfer( peer_meta: Dict[str, Any], src_block_ids: List[int], dst_block_ids: List[int], src_o_start: int, dst_o_start: int, n_outer: int, src_h0: int = 0, dst_h0: int = 0, n_heads: int = 0, ) → tuple[Any, str]#: Submit one full-slice or head-fragment NIXL transfer.

close() → None#: Release this buffer registration and its reference to the shared agent.

core.inference.disaggregation.transfer_backends.nixl#