`core.inference.disaggregation.transfer_backends.base`#

KV transfer backend registry and the buffer geometry shared by backends.

Backends are selected explicitly by the caller’s launcher configuration, never from the environment.

Module Contents#

Classes#

BufferGeometry

Address geometry of one registered paged buffer.

Functions#

`construct_kv_transfer_backend_class`	Return the backend class registered under `name`.
`compute_buffer_geometry`	Locate the blocks axis, derive the slice strides, and validate the canonical KV layout when the full geometry is provided.
`export_geometry_meta`	The wire schema shared by every backend’s export_meta.

Data#

KVTransportBackend

API#

core.inference.disaggregation.transfer_backends.base.KVTransportBackend#: None

core.inference.disaggregation.transfer_backends.base.construct_kv_transfer_backend_class( name: str, ) → core.inference.disaggregation.transfer_backends.base.KVTransportBackend#: Return the backend class registered under name.

class core.inference.disaggregation.transfer_backends.base.BufferGeometry#

Address geometry of one registered paged buffer.

Each (outer, block) pair is one contiguous slice; the outer stride skips over the full block pool for that outer index. layout is the canonical KV shard layout when the buffer is a KV cache; SSM pools carry their typed layout separately.

buf_ptr: int#: None

element_size: int#: None

device_id: int#: None

blocks_axis: int#: None

num_blocks: int#: None

num_outer: int#: None

bytes_per_slice: int#: None

outer_stride_bytes: int#: None

heads_per_partition: Optional[int]#: None

head_dim: Optional[int]#: None

tokens_per_block: Optional[int]#: None

layout: Optional[megatron.core.inference.disaggregation.kv_reshard.KVShardLayout]#: None

core.inference.disaggregation.transfer_backends.base.compute_buffer_geometry( memory_buffer: torch.Tensor, expected_num_blocks: int, *, backend_name: str, 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, ) → core.inference.disaggregation.transfer_backends.base.BufferGeometry#

Locate the blocks axis, derive the slice strides, and validate the canonical KV layout when the full geometry is provided.

Shared by every transfer backend so they agree on addressing and on the exported metadata schema. The inference KV layout is [2, L, B, T, H, d].

core.inference.disaggregation.transfer_backends.base.export_geometry_meta( geometry: core.inference.disaggregation.transfer_backends.base.BufferGeometry, ssm_layout=None, ) → dict#: The wire schema shared by every backend’s export_meta.

core.inference.disaggregation.transfer_backends.base#