diff --git a/megatron/core/inference/disaggregation/mamba_reshard.py b/megatron/core/inference/disaggregation/mamba_reshard.py deleted file mode 100644 index 8a23735154a..00000000000 --- a/megatron/core/inference/disaggregation/mamba_reshard.py +++ /dev/null @@ -1,222 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. - -"""Heterogeneous TP/PP reshard of Mamba conv/ssm state between prefill and -decode shard layouts (the Mamba analog of the attention KV reshard).""" - -from __future__ import annotations - -from dataclasses import dataclass -from typing import List, Tuple - -from megatron.core.inference.disaggregation.utils import intersect - -# Channel bands of a Mamba layer's state, in the order the conv state -# concatenates them on its channel axis (x, B, C); ssm is the head axis. -# (name, lives_in_conv). conv bands share one tensor; ssm is its own tensor. -_CONV_BANDS = ("x", "B", "C") - - -@dataclass(frozen=True) -class MambaStateDims: - """The model's (global, unsharded) Mamba structural dims. - - These belong to the MambaMixer / model config -- carried as one unit (rather - than loose constants spread across the layout) so there's a single source - and they can't drift apart. The producer should read them straight from the - model config (e.g. ``ngroups = config.mamba_num_groups``) rather than - reverse-deriving from tensor shapes. TP shards ``nheads``/``ngroups``; the - rest are unsharded. - """ - - nheads: int - headdim: int - d_state: int - ngroups: int - d_conv: int - - -@dataclass(frozen=True) -class MambaShardLayout: - """One rank's Mamba-state ownership: which global layers + TP rank, plus the - model's structural dims (:class:`MambaStateDims`). Per-rank locals follow by - dividing by ``tp_size``.""" - - global_rank: int - tp_size: int - tp_rank: int - layer_start: int # global Mamba-layer index of this rank's first layer - num_layers: int # Mamba layers held locally (this PP stage) - dims: MambaStateDims - - def __post_init__(self) -> None: - # Wire reconstruction (MambaShardLayout(**dict)) hands ``dims`` as a - # plain dict; coerce it back to MambaStateDims. - if isinstance(self.dims, dict): - object.__setattr__(self, "dims", MambaStateDims(**self.dims)) - # TP shards heads and groups; both must divide evenly or the local - # conv/ssm band sizes truncate to the wrong (or zero) width silently. - if self.dims.nheads % self.tp_size != 0: - raise ValueError(f"nheads={self.dims.nheads} not divisible by tp_size={self.tp_size}") - if self.dims.ngroups % self.tp_size != 0: - raise ValueError(f"ngroups={self.dims.ngroups} not divisible by tp_size={self.tp_size}") - - # Convenience proxies onto the dims so callers read ``layout.headdim`` etc. - @property - def nheads(self) -> int: - """Global (unsharded) number of Mamba heads.""" - return self.dims.nheads - - @property - def headdim(self) -> int: - """Dimension of each Mamba head.""" - return self.dims.headdim - - @property - def d_state(self) -> int: - """SSM state size per head.""" - return self.dims.d_state - - @property - def ngroups(self) -> int: - """Global (unsharded) number of B/C groups.""" - return self.dims.ngroups - - @property - def d_conv(self) -> int: - """Convolution kernel width.""" - return self.dims.d_conv - - def mamba_shard_key(self) -> Tuple[int, int]: - """The Mamba shard this rank holds: ``(tp_rank, layer_start)``. Ranks - sharing a key hold identical state (e.g. EP/DP replicas of it).""" - return (self.tp_rank, self.layer_start) - - @property - def d_inner(self) -> int: - """Global inner dimension (nheads * headdim).""" - return self.dims.nheads * self.dims.headdim - - @property - def nheads_local(self) -> int: - """Number of Mamba heads held by this TP rank.""" - return self.dims.nheads // self.tp_size - - @property - def d_inner_local(self) -> int: - """Local inner dimension for this TP rank.""" - return self.d_inner // self.tp_size - - @property - def ngroups_local(self) -> int: - """Number of B/C groups held by this TP rank.""" - return self.dims.ngroups // self.tp_size - - @property - def conv_dim_local(self) -> int: - """Total local conv channel width (x + B + C bands).""" - return self.d_inner_local + 2 * self.ngroups_local * self.dims.d_state - - def layer_range(self) -> Tuple[int, int]: - """Global Mamba-layer range ``[lo, hi)`` owned by this rank.""" - return (self.layer_start, self.layer_start + self.num_layers) - - def _band(self, name: str) -> Tuple[int, int, int]: - """``(global_total, local_size, conv_local_offset)`` for a band. - - ``conv_local_offset`` is the band's start on the local conv channel - axis; for the ``ssm`` (head) band it is the start on the local head - axis (always 0, heads are the whole tensor).""" - if name == "x": - g = self.d_inner - return g, self.d_inner_local, 0 - if name == "B": - g = self.dims.ngroups * self.dims.d_state - return g, self.ngroups_local * self.dims.d_state, self.d_inner_local - if name == "C": - g = self.dims.ngroups * self.dims.d_state - return ( - g, - self.ngroups_local * self.dims.d_state, - self.d_inner_local + self.ngroups_local * self.dims.d_state, - ) - if name == "ssm": - return self.dims.nheads, self.nheads_local, 0 - raise KeyError(name) - - -@dataclass(frozen=True) -class MambaReshardTransfer: - """One sub-block move for the reshard. - - ``band`` is ``"x"``/``"B"``/``"C"`` (conv channel axis) or ``"ssm"`` (head - axis). ``src_layer``/``dst_layer`` are local layer indices on each side; - ``*_lo``/``*_hi`` are the local channel/head slice bounds. - """ - - src_rank: int - dst_rank: int - band: str - global_layer: int - src_layer: int - dst_layer: int - src_lo: int - src_hi: int - dst_lo: int - dst_hi: int - - @property - def is_conv(self) -> bool: - """True if this transfer targets the conv state; False for ssm.""" - return self.band in _CONV_BANDS - - -def plan_mamba_reshard( - src_layouts: List[MambaShardLayout], dst_layouts: List[MambaShardLayout] -) -> List[MambaReshardTransfer]: - """Plan the conv/ssm sub-block moves from the prefill (src) layouts to the - decode (dst) layouts. One transfer per (src rank, dst rank, global layer, - band) where both the layer ranges and the channel ranges overlap.""" - # Dedupe replica sources: ranks sharing (tp_rank, layer_start) hold identical - # Mamba state (e.g. EP/DP replicas), so source each shard from exactly one of - # them -- the smallest global_rank -- to avoid duplicate sends. - rep_rank: dict = {} - for s in src_layouts: - key = s.mamba_shard_key() - if key not in rep_rank or s.global_rank < rep_rank[key]: - rep_rank[key] = s.global_rank - source_ranks = set(rep_rank.values()) - - out: List[MambaReshardTransfer] = [] - for s in src_layouts: - if s.global_rank not in source_ranks: - continue - s_lr = s.layer_range() - for d in dst_layouts: - layer_ov = intersect(s_lr, d.layer_range()) - if layer_ov is None: - continue - for band in (*_CONV_BANDS, "ssm"): - _, s_size, s_off = s._band(band) - _, d_size, d_off = d._band(band) - s_glo = (s.tp_rank * s_size, s.tp_rank * s_size + s_size) - d_glo = (d.tp_rank * d_size, d.tp_rank * d_size + d_size) - chan_ov = intersect(s_glo, d_glo) - if chan_ov is None: - continue - lo, hi = chan_ov - for g in range(layer_ov[0], layer_ov[1]): - out.append( - MambaReshardTransfer( - src_rank=s.global_rank, - dst_rank=d.global_rank, - band=band, - global_layer=g, - src_layer=g - s.layer_start, - dst_layer=g - d.layer_start, - src_lo=s_off + (lo - s_glo[0]), - src_hi=s_off + (hi - s_glo[0]), - dst_lo=d_off + (lo - d_glo[0]), - dst_hi=d_off + (hi - d_glo[0]), - ) - ) - return out diff --git a/megatron/core/inference/disaggregation/ssm_reshard.py b/megatron/core/inference/disaggregation/ssm_reshard.py new file mode 100644 index 00000000000..9c06b9f70f9 --- /dev/null +++ b/megatron/core/inference/disaggregation/ssm_reshard.py @@ -0,0 +1,205 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""Heterogeneous TP/PP reshard of SSM boundary-snapshot state between +prefill and decode shard layouts (the SSM analog of attention KV resharding). + +A snapshot's conv state packs three channel bands, [x | B | C], on one axis: +x is head-sharded (d_inner) and B/C are group-sharded (ngroups * d_state). +The recurrent state is head-sharded. ``plan_ssm_reshard`` emits one transfer per +(src rank, dst rank, global layer, band) whose layer and channel ranges +overlap; both sides compute the same plan from the same layout lists, so the +send and receive orders match. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import List, Tuple + +from megatron.core.inference.disaggregation.utils import intersect + +# Channel bands of an SSM layer's conv state, in the order the conv state +# concatenates them on its channel axis. The recurrent band is the head axis +# of its own tensor. +_CONV_BANDS = ("x", "B", "C") + + +@dataclass(frozen=True) +class SSMStateDims: + """The model's global (unsharded) SSM structural dimensions. + + Carried as one unit so there is a single source and the dims cannot drift + apart. The producer should read them from the model config rather than + deriving them from tensor shapes. TP shards nheads/ngroups; the rest are + unsharded. + """ + + nheads: int + headdim: int + d_state: int + ngroups: int + d_conv: int + + +@dataclass(frozen=True) +class SSMShardLayout: + """One rank's SSM-state ownership: which global layers and TP rank, + plus the model's structural dims. Per-rank local sizes follow by dividing + by tp_size.""" + + global_rank: int + tp_size: int + tp_rank: int + layer_start: int # global SSM-layer index of this rank's first layer + num_layers: int # SSM layers held locally (this PP stage) + dims: SSMStateDims + + def __post_init__(self) -> None: + # Wire reconstruction (SSMShardLayout(**dict)) hands dims as a plain + # dict; coerce it back to SSMStateDims. + if isinstance(self.dims, dict): + object.__setattr__(self, "dims", SSMStateDims(**self.dims)) + # TP shards heads and groups; both must divide evenly or the local + # band widths are wrong. + if self.dims.nheads % self.tp_size != 0: + raise ValueError(f"nheads={self.dims.nheads} not divisible by tp_size={self.tp_size}") + if self.dims.ngroups % self.tp_size != 0: + raise ValueError(f"ngroups={self.dims.ngroups} not divisible by tp_size={self.tp_size}") + + @property + def d_inner(self) -> int: + """Global inner dimension (nheads * headdim).""" + return self.dims.nheads * self.dims.headdim + + @property + def nheads_local(self) -> int: + """SSM heads held by this TP rank.""" + return self.dims.nheads // self.tp_size + + @property + def d_inner_local(self) -> int: + """Local inner dimension for this TP rank.""" + return self.d_inner // self.tp_size + + @property + def ngroups_local(self) -> int: + """B/C groups held by this TP rank.""" + return self.dims.ngroups // self.tp_size + + @property + def conv_dim_local(self) -> int: + """Total local conv channel width (x + B + C bands).""" + return self.d_inner_local + 2 * self.ngroups_local * self.dims.d_state + + def shard_key(self) -> Tuple[int, int]: + """The SSM shard this rank holds: (tp_rank, layer_start). Ranks + sharing a key hold identical state (e.g. EP/DP replicas).""" + return (self.tp_rank, self.layer_start) + + def layer_range(self) -> Tuple[int, int]: + """Global SSM-layer range [lo, hi) owned by this rank.""" + return (self.layer_start, self.layer_start + self.num_layers) + + def band(self, name: str) -> Tuple[int, int, int]: + """Return (global_total, local_size, local_offset) for a band. + + local_offset is the band's start on the local conv channel axis; for + the "recurrent" (head) band it is the start on the local head axis + (always 0, heads are the whole tensor). + """ + if name == "x": + return self.d_inner, self.d_inner_local, 0 + if name == "B": + g = self.dims.ngroups * self.dims.d_state + return g, self.ngroups_local * self.dims.d_state, self.d_inner_local + if name == "C": + g = self.dims.ngroups * self.dims.d_state + return ( + g, + self.ngroups_local * self.dims.d_state, + self.d_inner_local + self.ngroups_local * self.dims.d_state, + ) + if name == "recurrent": + return self.dims.nheads, self.nheads_local, 0 + raise KeyError(name) + + +@dataclass(frozen=True) +class SSMReshardTransfer: + """One sub-block move of the snapshot reshard. + + band is "x"/"B"/"C" (conv channel axis) or "recurrent" (head axis). + src_layer/dst_layer are local layer indices on each side; *_lo/*_hi are + the local channel or head slice bounds. + """ + + src_rank: int + dst_rank: int + band: str + global_layer: int + src_layer: int + dst_layer: int + src_lo: int + src_hi: int + dst_lo: int + dst_hi: int + + @property + def is_conv(self) -> bool: + """True if this transfer targets the conv state; False for recurrent.""" + return self.band in _CONV_BANDS + + +def plan_ssm_reshard( + src_layouts: List[SSMShardLayout], dst_layouts: List[SSMShardLayout] +) -> List[SSMReshardTransfer]: + """Plan the conv/recurrent sub-block moves from the prefill (src) layouts to + the decode (dst) layouts: one transfer per (src rank, dst rank, global + layer, band) where both the layer ranges and the channel ranges overlap. + + Ranks sharing (tp_rank, layer_start) hold identical SSM state (e.g. + EP/DP replicas), so each shard is sourced from exactly one of them, the + smallest global_rank. Deterministic given the layout lists, so both sides + enumerate the same transfers in the same order. + """ + rep_rank: dict = {} + for s in src_layouts: + key = s.shard_key() + if key not in rep_rank or s.global_rank < rep_rank[key]: + rep_rank[key] = s.global_rank + source_ranks = set(rep_rank.values()) + + out: List[SSMReshardTransfer] = [] + for s in src_layouts: + if s.global_rank not in source_ranks: + continue + s_lr = s.layer_range() + for d in dst_layouts: + layer_ov = intersect(s_lr, d.layer_range()) + if layer_ov is None: + continue + for band in (*_CONV_BANDS, "recurrent"): + _, s_size, s_off = s.band(band) + _, d_size, d_off = d.band(band) + s_glo = (s.tp_rank * s_size, s.tp_rank * s_size + s_size) + d_glo = (d.tp_rank * d_size, d.tp_rank * d_size + d_size) + chan_ov = intersect(s_glo, d_glo) + if chan_ov is None: + continue + lo, hi = chan_ov + for g in range(layer_ov[0], layer_ov[1]): + out.append( + SSMReshardTransfer( + src_rank=s.global_rank, + dst_rank=d.global_rank, + band=band, + global_layer=g, + src_layer=g - s.layer_start, + dst_layer=g - d.layer_start, + src_lo=s_off + (lo - s_glo[0]), + src_hi=s_off + (hi - s_glo[0]), + dst_lo=d_off + (lo - d_glo[0]), + dst_hi=d_off + (hi - d_glo[0]), + ) + ) + return out diff --git a/megatron/core/inference/disaggregation/transfer_backends/__init__.py b/megatron/core/inference/disaggregation/transfer_backends/__init__.py new file mode 100644 index 00000000000..262675e55c0 --- /dev/null +++ b/megatron/core/inference/disaggregation/transfer_backends/__init__.py @@ -0,0 +1,7 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""KV transfer backends for disaggregated inference.""" + +from .base import KVTransportBackend, construct_kv_transfer_backend_class + +__all__ = ["KVTransportBackend", "construct_kv_transfer_backend_class"] diff --git a/megatron/core/inference/disaggregation/transfer_backends/base.py b/megatron/core/inference/disaggregation/transfer_backends/base.py new file mode 100644 index 00000000000..fa8df492e44 --- /dev/null +++ b/megatron/core/inference/disaggregation/transfer_backends/base.py @@ -0,0 +1,212 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""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. +""" + +from __future__ import annotations + +from dataclasses import asdict, dataclass +from typing import Any, Optional + +import torch + +from megatron.core.inference.disaggregation.kv_reshard import KVShardLayout +from megatron.core.inference.disaggregation.ssm_reshard import SSMShardLayout + +KVTransportBackend = Any + + +def construct_kv_transfer_backend_class(name: str) -> KVTransportBackend: + """Return the backend class registered under ``name``.""" + + normalized = name.lower().replace("_", "-") + if normalized == "nixl": + from .nixl import NixlTransferBackend + + return NixlTransferBackend + if normalized == "nccl": + from .nccl import NcclTransferBackend + + return NcclTransferBackend + raise ValueError("Unsupported KV transfer backend %r; expected 'nixl' or 'nccl'." % name) + + +@dataclass +class 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 + element_size: int + device_id: int + blocks_axis: int + num_blocks: int + num_outer: int + bytes_per_slice: int + outer_stride_bytes: int + heads_per_partition: Optional[int] + head_dim: Optional[int] + tokens_per_block: Optional[int] + layout: Optional[KVShardLayout] + + +def 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[SSMShardLayout] = None, + ssm_state_kind: Optional[str] = None, +) -> 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]. + """ + if (ssm_layout is None) != (ssm_state_kind is None): + raise ValueError("ssm_layout and ssm_state_kind must be provided together") + if ssm_state_kind not in (None, "conv", "recurrent"): + raise ValueError("ssm_state_kind must be 'conv' or 'recurrent'") + + layout_capable = ( + None + not in ( + global_rank, + tp_size, + tp_rank, + pp_size, + pp_rank, + num_layers_global, + num_kv_heads_global, + heads_per_partition, + head_dim, + tokens_per_block, + layer_start, + layer_end, + ) + and heads_per_partition * tp_size == num_kv_heads_global + ) + + shape = list(memory_buffer.shape) + candidates = [i for i, dim in enumerate(shape) if dim == expected_num_blocks] + if not candidates: + raise RuntimeError( + f"{backend_name}: no axis in memory_buffer shape {shape} matches " + f"expected_num_blocks={expected_num_blocks}. Layout is unrecognized; " + "bug in caller or new Megatron tensor shape." + ) + if len(candidates) > 1: + raise RuntimeError( + f"{backend_name}: ambiguous blocks axis in shape {shape} " + f"(expected_num_blocks={expected_num_blocks} matches multiple axes " + f"{candidates}). Caller must pass a more distinctive value." + ) + blocks_axis = candidates[0] + + elements_per_slice = 1 + for dim in shape[blocks_axis + 1 :]: + elements_per_slice *= dim + element_size = memory_buffer.element_size() + bytes_per_slice = element_size * elements_per_slice + num_outer = 1 + for dim in shape[:blocks_axis]: + num_outer *= dim + + layout = None + if layout_capable: + layout = KVShardLayout( + num_layers=int(num_layers_global), + num_heads=int(num_kv_heads_global), + tp_size=int(tp_size), + tp_rank=int(tp_rank), + pp_size=int(pp_size), + pp_rank=int(pp_rank), + global_rank=int(global_rank), + layer_start=int(layer_start), + num_local_layers=int(layer_end) - int(layer_start), + ) + if blocks_axis != 2: + raise ValueError("inference KV transfers require the [2, L, B, T, H, d] layout") + if layout.local_num_heads() != heads_per_partition: + raise ValueError( + "heads_per_partition does not match the canonical KV layout: " + f"{heads_per_partition} vs {layout.local_num_heads()}" + ) + if num_outer % layout.local_num_layers() != 0: + raise ValueError( + f"num_outer={num_outer} is not divisible by local layers=" + f"{layout.local_num_layers()}" + ) + if layout is not None and ssm_layout is not None: + raise ValueError("a transfer backend cannot have both KV and SSM layouts") + + return BufferGeometry( + buf_ptr=memory_buffer.data_ptr(), + element_size=element_size, + device_id=memory_buffer.device.index if memory_buffer.is_cuda else 0, + blocks_axis=blocks_axis, + num_blocks=expected_num_blocks, + num_outer=num_outer, + bytes_per_slice=bytes_per_slice, + outer_stride_bytes=expected_num_blocks * bytes_per_slice, + heads_per_partition=heads_per_partition, + head_dim=head_dim, + tokens_per_block=tokens_per_block, + layout=layout, + ) + + +def export_geometry_meta(geometry: BufferGeometry, ssm_layout=None) -> dict: + """The wire schema shared by every backend's export_meta.""" + meta = { + "base_addr": geometry.buf_ptr, + "outer_stride_bytes": geometry.outer_stride_bytes, + "device_id": geometry.device_id, + "num_outer": geometry.num_outer, + "bytes_per_slice": geometry.bytes_per_slice, + "blocks_axis": geometry.blocks_axis, + "num_blocks": geometry.num_blocks, + "heads_per_partition": geometry.heads_per_partition, + "head_dim": geometry.head_dim, + "tokens_per_block": geometry.tokens_per_block, + "element_size": geometry.element_size, + } + if geometry.layout is not None: + layer_start, layer_end = geometry.layout.layer_range() + meta.update( + { + "global_rank": geometry.layout.global_rank, + "tp_size": geometry.layout.tp_size, + "tp_rank": geometry.layout.tp_rank, + "pp_size": geometry.layout.pp_size, + "pp_rank": geometry.layout.pp_rank, + "num_layers_global": geometry.layout.num_layers, + "num_kv_heads_global": geometry.layout.num_heads, + "layer_start": layer_start, + "layer_end": layer_end, + } + ) + if ssm_layout is not None: + meta["ssm_layout"] = asdict(ssm_layout) + return meta diff --git a/megatron/core/inference/disaggregation/transfer_backends/nccl.py b/megatron/core/inference/disaggregation/transfer_backends/nccl.py new file mode 100644 index 00000000000..96350fb54bf --- /dev/null +++ b/megatron/core/inference/disaggregation/transfer_backends/nccl.py @@ -0,0 +1,293 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""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. +""" + +from __future__ import annotations + +import logging +from typing import Any, Dict, List, Optional + +import torch +import torch.distributed as dist + +from megatron.core.inference.disaggregation.kv_reshard import KVShardLayout, plan_kv_reshard +from megatron.core.inference.disaggregation.ssm_reshard import SSMShardLayout, plan_ssm_reshard +from megatron.core.inference.disaggregation.transfer_backends.base import ( + compute_buffer_geometry, + export_geometry_meta, +) +from megatron.core.inference.disaggregation.utils import transfer_peer_records + +logger = logging.getLogger(__name__) + + +class NcclTransferHandle: + """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. + """ + + def __init__(self, works: List[Any], keepalive: List[torch.Tensor], scatters: List[Any]): + self._works = works + self._keepalive = keepalive + self._scatters = scatters + self._done = not works and not scatters + + def poll(self) -> bool: + """Return True if the transfer has settled, scattering received data + into the paged buffers on first completion.""" + if self._done: + return True + if not all(w.is_completed() for w in self._works): + return False + self._finish() + return True + + def wait(self) -> None: + """Block until the transfer completes, then scatter.""" + for w in self._works: + w.wait() + self._finish() + + def _finish(self) -> None: + if self._done: + return + with torch.inference_mode(): + for scatter in self._scatters: + scatter() + self._keepalive.clear() + self._works = [] + self._done = True + + +def _make_copy(view: torch.Tensor, buf: torch.Tensor): + def _copy(): + view.copy_(buf.view(view.shape)) + + return _copy + + +def _kv_layout_from_meta(meta: Dict[str, Any]) -> KVShardLayout: + """Rebuild a peer's KVShardLayout from its exported metadata.""" + return KVShardLayout( + num_layers=int(meta["num_layers_global"]), + num_heads=int(meta["num_kv_heads_global"]), + tp_size=int(meta["tp_size"]), + tp_rank=int(meta["tp_rank"]), + pp_size=int(meta["pp_size"]), + pp_rank=int(meta["pp_rank"]), + global_rank=int(meta["global_rank"]), + layer_start=int(meta["layer_start"]), + num_local_layers=int(meta["layer_end"]) - int(meta["layer_start"]), + ) + + +class NcclTransferBackend: + """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). + """ + + name = "nccl" + is_push = True + + def __init__( + self, + 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[SSMShardLayout] = None, + ssm_state_kind: Optional[str] = None, + ): + if not (dist.is_available() and dist.is_initialized()): + raise RuntimeError( + "NcclTransferBackend requires torch.distributed to be initialized; " + "the prefill and decode workers must share a process group." + ) + self.agent_name = agent_name + self._memory_buffer = memory_buffer + geometry = compute_buffer_geometry( + memory_buffer, + expected_num_blocks, + backend_name="NcclTransferBackend", + tp_size=tp_size, + tp_rank=tp_rank, + num_kv_heads_global=num_kv_heads_global, + heads_per_partition=heads_per_partition, + head_dim=head_dim, + tokens_per_block=tokens_per_block, + global_rank=global_rank, + pp_size=pp_size, + pp_rank=pp_rank, + num_layers_global=num_layers_global, + layer_start=layer_start, + layer_end=layer_end, + ssm_layout=ssm_layout, + ssm_state_kind=ssm_state_kind, + ) + self._geometry = geometry + self._layout = geometry.layout + self._ssm_layout = ssm_layout + self._ssm_state_kind = ssm_state_kind + logger.info( + "NcclTransferBackend[%s] over %d-block buffer (rank=%d, shape=%s)", + agent_name, + geometry.num_blocks, + dist.get_rank(), + list(memory_buffer.shape), + ) + + def export_meta(self) -> Dict[str, Any]: + """The shared geometry schema plus this rank's NCCL address.""" + meta = export_geometry_meta(self._geometry, self._ssm_layout) + meta["transport"] = "nccl" + meta["nccl_rank"] = dist.get_rank() + return meta + + # --- shared enumeration ------------------------------------------------- + def _kv_transfers(self, peer_records, mine_is_src: bool): + """Yield (peer_meta, layers, heads) for this rank's part of the KV + reshard plan, in deterministic plan order.""" + sources: list = [] + peers_by_rank: dict = {} + for meta, blocks in peer_records: + layout = _kv_layout_from_meta(meta) + if layout.global_rank in peers_by_rank: + raise ValueError(f"duplicate peer global_rank={layout.global_rank} in KV metadata") + peers_by_rank[layout.global_rank] = meta + sources.append(layout) + if mine_is_src: + plan = plan_kv_reshard([self._layout], sources) + else: + plan = plan_kv_reshard(sources, [self._layout]) + for transfer in plan: + peer_rank = transfer.dst_rank if mine_is_src else transfer.src_rank + meta = peers_by_rank[peer_rank] + if mine_is_src: + layers = transfer.src_layer_slice(self._layout) + heads = transfer.src_head_slice(self._layout) + else: + layers = transfer.dst_layer_slice(self._layout) + heads = transfer.dst_head_slice(self._layout) + yield meta, layers, heads + + def _ssm_transfers(self, peer_records, mine_is_src: bool): + """Yield (peer_meta, lo, hi) band slices of this rank's SSM state, + in deterministic plan order.""" + for meta, _ in peer_records: + raw_layout = meta.get("ssm_layout") + if not isinstance(raw_layout, dict): + raise ValueError("peer metadata is missing ssm_layout") + peer_layout = SSMShardLayout(**raw_layout) + if mine_is_src: + plan = plan_ssm_reshard([self._ssm_layout], [peer_layout]) + else: + plan = plan_ssm_reshard([peer_layout], [self._ssm_layout]) + for t in plan: + if t.is_conv != (self._ssm_state_kind == "conv"): + continue + if mine_is_src: + yield meta, t.src_layer, t.src_lo, t.src_hi + else: + yield meta, t.dst_layer, t.dst_lo, t.dst_hi + + def _kv_block_view(self, block_id: int, layers: slice, heads: slice) -> torch.Tensor: + """One block's (kv, layer, token, head, dim) fragment in the + [2, L, B, T, H, d] paged buffer.""" + return self._memory_buffer[:, layers, block_id, :, heads, :] + + # --- decode side --------------------------------------------------------- + def begin_pull_blocks( + self, peer_meta: Any, src_block_ids: List[int], dst_block_ids: List[int] + ) -> NcclTransferHandle: + """Post the receives matching the prefill's sends; the handle scatters + into the destination blocks (or SSM slots) on completion.""" + if not dst_block_ids: + return NcclTransferHandle([], [], []) + records = transfer_peer_records(peer_meta, src_block_ids) + + ops: List[Any] = [] + buffers: List[torch.Tensor] = [] + scatters: List[Any] = [] + device = self._memory_buffer.device + dtype = self._memory_buffer.dtype + + if self._ssm_layout is not None: + for meta, layer, lo, hi in self._ssm_transfers(records, mine_is_src=False): + for slot in dst_block_ids: + view = self._memory_buffer[layer, int(slot), lo:hi] + buf = torch.empty(view.shape, dtype=dtype, device=device) + buffers.append(buf) + ops.append(dist.P2POp(dist.irecv, buf, int(meta["nccl_rank"]))) + scatters.append(_make_copy(view, buf)) + else: + geo = self._geometry + for meta, layers, heads in self._kv_transfers(records, mine_is_src=False): + n_layers = layers.stop - layers.start + n_heads = heads.stop - heads.start + for block in dst_block_ids: + buf = torch.empty( + (2, n_layers, geo.tokens_per_block, n_heads, geo.head_dim), + dtype=dtype, + device=device, + ) + buffers.append(buf) + ops.append(dist.P2POp(dist.irecv, buf, int(meta["nccl_rank"]))) + scatters.append(_make_copy(self._kv_block_view(int(block), layers, heads), buf)) + + works = dist.batch_isend_irecv(ops) if ops else [] + return NcclTransferHandle(works, buffers, scatters) + + # --- prefill side ---------------------------------------------------------- + def begin_push_blocks(self, peer_meta: Any, src_block_ids: List[int]) -> NcclTransferHandle: + """Post the sends matching the decode's receives, straight out of the + pinned source blocks (or SSM slots). `peer_meta` is the decode + instance's per-rank metadata in the same nested shape as a hand-off's + kv_meta.""" + if not src_block_ids: + return NcclTransferHandle([], [], []) + records = transfer_peer_records(peer_meta, []) + + ops: List[Any] = [] + keep: List[torch.Tensor] = [] + + if self._ssm_layout is not None: + for meta, layer, lo, hi in self._ssm_transfers(records, mine_is_src=True): + for slot in src_block_ids: + sub = self._memory_buffer[layer, int(slot), lo:hi].contiguous() + keep.append(sub) + ops.append(dist.P2POp(dist.isend, sub, int(meta["nccl_rank"]))) + else: + for meta, layers, heads in self._kv_transfers(records, mine_is_src=True): + for block in src_block_ids: + sub = self._kv_block_view(int(block), layers, heads).contiguous() + keep.append(sub) + ops.append(dist.P2POp(dist.isend, sub, int(meta["nccl_rank"]))) + + works = dist.batch_isend_irecv(ops) if ops else [] + return NcclTransferHandle(works, keep, []) diff --git a/megatron/core/inference/disaggregation/transfer_backends/nixl.py b/megatron/core/inference/disaggregation/transfer_backends/nixl.py new file mode 100644 index 00000000000..af355e25045 --- /dev/null +++ b/megatron/core/inference/disaggregation/transfer_backends/nixl.py @@ -0,0 +1,625 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""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. +""" + +from __future__ import annotations + +import base64 +import logging +import os +import time +from dataclasses import asdict, dataclass +from typing import Any, Dict, List, Optional + +import torch + +from megatron.core.inference.disaggregation.kv_reshard import KVShardLayout, plan_kv_reshard +from megatron.core.inference.disaggregation.ssm_reshard import SSMShardLayout, plan_ssm_reshard +from megatron.core.inference.disaggregation.transfer_backends.base import compute_buffer_geometry +from megatron.core.inference.disaggregation.utils import transfer_peer_records + +logger = logging.getLogger(__name__) + +try: + from nixl._api import nixl_agent # type: ignore[import-not-found] + + _HAVE_NIXL = True +except ImportError: + nixl_agent = None # type: ignore[assignment] + _HAVE_NIXL = False + + +# NIXL exposes polling, not a blocking wait. A long stall usually means peer or +# fabric failure, so cap the wait. +_POLL_INTERVAL_S = 0.0005 # 0.5 ms +_POLL_TIMEOUT_S = 30.0 + + +@dataclass +class NixlPullHandle: + """Pollable handle for one logical pull made of one or more NIXL transfers.""" + + agent: Any + xfers: List[Any] + contexts: List[str] + submitted_at: float + timeout_s: float = _POLL_TIMEOUT_S + done: bool = False + error: Optional[str] = None + + def poll(self) -> bool: + """Return True if every transfer has settled, without blocking.""" + if self.done: + if self.error is not None: + raise RuntimeError(self.error) + return True + if not self.xfers: + self.done = True + return True + + errors: List[str] = [] + pending: List[str] = [] + for xfer, ctx in zip(self.xfers, self.contexts): + state = self.agent.check_xfer_state(xfer) + if state == "DONE": + continue + if state == "ERR": + errors.append(ctx) + continue + pending.append(f"{ctx}: {state}") + + if not pending: + self.done = True + if errors: + self.error = f"NIXL transfer failed ({', '.join(errors)})" + raise RuntimeError(self.error) + return True + if time.perf_counter() - self.submitted_at > self.timeout_s: + raise TimeoutError( + f"NIXL transfer timed out after {self.timeout_s}s; pending={pending}" + ) + return False + + def wait(self) -> None: + """Block until the transfer completes; NIXL has no blocking wait, so + poll with a short sleep to avoid monopolizing a CPU core.""" + while not self.poll(): + time.sleep(_POLL_INTERVAL_S) + + +class NixlTransferBackend: + """Per-rank NIXL agent owning a registration over the paged KV buffer. + + Per-block transfers are descriptor ranges over that registration. Peer + metadata is exchanged by the control plane and registered lazily on first + pull. + """ + + name = "nixl" + + def __init__( + self, + 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[SSMShardLayout] = None, + ssm_state_kind: Optional[str] = None, + ): + if not _HAVE_NIXL: + raise RuntimeError( + "NixlTransferBackend requires the nixl Python package. Install the " + "NIXL runtime and `pip install nixl` before launching " + "disaggregated workers." + ) + self.agent_name = agent_name + self._memory_buffer = memory_buffer + + # Addressing geometry shared with the other backends. + geometry = compute_buffer_geometry( + memory_buffer, + expected_num_blocks, + backend_name="NixlTransferBackend", + tp_size=tp_size, + tp_rank=tp_rank, + num_kv_heads_global=num_kv_heads_global, + heads_per_partition=heads_per_partition, + head_dim=head_dim, + tokens_per_block=tokens_per_block, + global_rank=global_rank, + pp_size=pp_size, + pp_rank=pp_rank, + num_layers_global=num_layers_global, + layer_start=layer_start, + layer_end=layer_end, + ssm_layout=ssm_layout, + ssm_state_kind=ssm_state_kind, + ) + self._geometry = geometry + shape = list(memory_buffer.shape) + self._buf_ptr = geometry.buf_ptr + self._element_size = geometry.element_size + self._device_id = geometry.device_id + self._outer_stride_bytes = geometry.outer_stride_bytes + self._num_outer = geometry.num_outer + self._bytes_per_slice = geometry.bytes_per_slice + self._blocks_axis = geometry.blocks_axis + self._num_blocks = geometry.num_blocks + self._heads_per_partition = geometry.heads_per_partition + self._head_dim = geometry.head_dim + self._tokens_per_block = geometry.tokens_per_block + self._layout = geometry.layout + self._ssm_layout = ssm_layout + self._ssm_state_kind = ssm_state_kind + + # Configure UCX before agent construction. Avoid TCP for VRAM addresses; + # operators may override this by setting UCX_TLS before launch. + os.environ.setdefault("UCX_TLS", "cuda_ipc,cuda_copy,cma,shm,self") + # Explicit registration makes the UCX memtype cache unnecessary and + # avoids stale VRAM/host classifications. + os.environ.setdefault("UCX_MEMTYPE_CACHE", "n") + + self._agent = nixl_agent(agent_name) + self._reg_handle = self._agent.register_memory(memory_buffer) + + # Base64 keeps NIXL metadata safe for msgpack/json control messages. + self._agent_metadata = self._agent.get_agent_metadata() + + # Peer agent_name -> id returned by add_remote_agent. + self._known_peers: Dict[str, Any] = {} + + logger.info( + "NixlTransferBackend[%s] registered %d-block buffer " + "(blocks_axis=%d, %d outer-slices/block × %d bytes/slice = " + "%d bytes/block, device=%d, shape=%s)", + agent_name, + self._num_blocks, + self._blocks_axis, + self._num_outer, + self._bytes_per_slice, + self._num_outer * self._bytes_per_slice, + self._device_id, + shape, + ) + + def export_meta(self) -> 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. + """ + meta = { + "agent_name": self.agent_name, + "agent_metadata_b64": base64.b64encode(self._agent_metadata).decode("ascii"), + "base_addr": self._buf_ptr, + "outer_stride_bytes": self._outer_stride_bytes, + "device_id": self._device_id, + "num_outer": self._num_outer, + "bytes_per_slice": self._bytes_per_slice, + "blocks_axis": self._blocks_axis, + "num_blocks": self._num_blocks, + "heads_per_partition": self._heads_per_partition, + "head_dim": self._head_dim, + "tokens_per_block": self._tokens_per_block, + "element_size": self._element_size, + } + if self._layout is not None: + layer_start, layer_end = self._layout.layer_range() + meta.update( + { + "global_rank": self._layout.global_rank, + "tp_size": self._layout.tp_size, + "tp_rank": self._layout.tp_rank, + "pp_size": self._layout.pp_size, + "pp_rank": self._layout.pp_rank, + "num_layers_global": self._layout.num_layers, + "num_kv_heads_global": self._layout.num_heads, + "layer_start": layer_start, + "layer_end": layer_end, + } + ) + if self._ssm_layout is not None: + meta["ssm_layout"] = asdict(self._ssm_layout) + return meta + + def _ensure_peer_registered(self, peer_meta: Dict[str, Any]) -> str: + """Register the peer with NIXL on first use; return its agent id.""" + peer_name = peer_meta["agent_name"] + existing = self._known_peers.get(peer_name) + if existing is not None: + return existing + metadata_b64 = peer_meta.get("agent_metadata_b64") + if not metadata_b64: + raise ValueError(f"peer_meta for {peer_name!r} is missing agent_metadata_b64") + peer_id = self._agent.add_remote_agent(base64.b64decode(metadata_b64)) + resolved = peer_id if peer_id else peer_name + self._known_peers[peer_name] = resolved + logger.info("NixlTransferBackend[%s] registered peer %s", self.agent_name, peer_name) + return resolved + + def _validate_peer( + self, + 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.""" + + if len(src_block_ids) != len(dst_block_ids): + raise ValueError( + f"source/destination block_id length mismatch for peer " + f"{meta.get('agent_name')!r}: {len(src_block_ids)} vs {len(dst_block_ids)}" + ) + for block in src_block_ids: + if not 0 <= block < int(meta["num_blocks"]): + raise ValueError(f"source block {block} is outside pool [0, {meta['num_blocks']})") + for block in dst_block_ids: + if not 0 <= block < self._num_blocks: + raise ValueError( + f"destination block {block} is outside pool [0, {self._num_blocks})" + ) + + local = { + "head_dim": self._head_dim, + "tokens_per_block": self._tokens_per_block, + "element_size": self._element_size, + "num_outer": self._num_outer, + "bytes_per_slice": self._bytes_per_slice, + "blocks_axis": self._blocks_axis, + "heads_per_partition": self._heads_per_partition, + } + fields = ["head_dim", "tokens_per_block", "element_size"] + if matched_layout: + fields.extend(["num_outer", "bytes_per_slice", "blocks_axis", "heads_per_partition"]) + mismatches = [ + f"{field}: peer={meta.get(field)} local={local[field]}" + for field in fields + if meta.get(field) is not None + and local[field] is not None + and meta.get(field) != local[field] + ] + if mismatches: + kind = "matched-layout" if matched_layout else "transfer" + raise ValueError(f"{kind} geometry mismatch: {', '.join(mismatches)}") + + @staticmethod + def _kv_layout_from_meta(meta: Dict[str, Any]) -> KVShardLayout: + """Reconstruct a main-planner KV layout from peer wire metadata.""" + + keys = ( + "global_rank", + "tp_size", + "tp_rank", + "pp_size", + "pp_rank", + "num_layers_global", + "num_kv_heads_global", + "layer_start", + "layer_end", + ) + missing = [key for key in keys if meta.get(key) is None] + if missing: + raise ValueError(f"peer metadata missing KV layout fields: {missing}") + return KVShardLayout( + num_layers=int(meta["num_layers_global"]), + num_heads=int(meta["num_kv_heads_global"]), + tp_size=int(meta["tp_size"]), + tp_rank=int(meta["tp_rank"]), + pp_size=int(meta["pp_size"]), + pp_rank=int(meta["pp_rank"]), + global_rank=int(meta["global_rank"]), + layer_start=int(meta["layer_start"]), + num_local_layers=int(meta["layer_end"]) - int(meta["layer_start"]), + ) + + def begin_pull_blocks( + self, peer_meta: Any, src_block_ids: List[int], dst_block_ids: List[int] + ) -> NixlPullHandle: + """Submit a pull and return a handle that can be polled later.""" + if not isinstance(peer_meta, dict) or "pp_metas" not in peer_meta: + if not src_block_ids and not dst_block_ids: + return NixlPullHandle( + agent=self._agent, + xfers=[], + contexts=[], + submitted_at=time.perf_counter(), + done=True, + ) + + xfers: List[Any] = [] + contexts: List[str] = [] + submitted_at = time.perf_counter() + try: + if self._ssm_layout is not None: + state_kind = self._ssm_state_kind + assert state_kind is not None + width = ( + self._ssm_layout.conv_dim_local + if state_kind == "conv" + else self._ssm_layout.nheads_local + ) + if ( + self._heads_per_partition != width + or self._num_outer != self._ssm_layout.num_layers + or self._blocks_axis != 1 + ): + raise ValueError(f"local {state_kind} geometry does not match its SSM layout") + + sources = [] + peers_by_rank = {} + for meta, blocks in transfer_peer_records(peer_meta, src_block_ids): + raw_layout = meta.get("ssm_layout") + if not isinstance(raw_layout, dict): + raise ValueError("peer metadata is missing ssm_layout") + layout = SSMShardLayout(**raw_layout) + self._validate_peer(meta, blocks, dst_block_ids) + peer_width = ( + layout.conv_dim_local if state_kind == "conv" else layout.nheads_local + ) + if ( + meta.get("heads_per_partition") != peer_width + or int(meta["num_outer"]) != layout.num_layers + or int(meta["blocks_axis"]) != 1 + ): + raise ValueError( + f"peer {state_kind} geometry does not match its SSM layout" + ) + if layout.global_rank in peers_by_rank: + raise ValueError( + f"duplicate source global_rank={layout.global_rank} " "in SSM metadata" + ) + sources.append(layout) + peers_by_rank[layout.global_rank] = (meta, blocks) + if not sources: + raise ValueError("SSM handoff contains no source peer metadata") + + transfers = [ + transfer + for transfer in plan_ssm_reshard(sources, [self._ssm_layout]) + if transfer.is_conv == (state_kind == "conv") + ] + for layer in range(self._ssm_layout.num_layers): + intervals = sorted( + (transfer.dst_lo, transfer.dst_hi) + for transfer in transfers + if transfer.dst_layer == layer + ) + if not intervals or intervals[0][0] != 0 or intervals[-1][1] != width: + raise ValueError(f"incomplete SSM {state_kind} coverage for layer {layer}") + if any(a[1] != b[0] for a, b in zip(intervals, intervals[1:])): + raise ValueError( + f"non-contiguous SSM {state_kind} coverage for layer {layer}" + ) + + for transfer in transfers: + meta, blocks = peers_by_rank[transfer.src_rank] + xfer, ctx = self._begin_transfer( + meta, + blocks, + dst_block_ids, + transfer.src_layer, + transfer.dst_layer, + 1, + transfer.src_lo, + transfer.dst_lo, + transfer.src_hi - transfer.src_lo, + ) + xfers.append(xfer) + contexts.append(ctx) + elif self._layout is not None: + sources = [] + peers_by_rank = {} + for meta, blocks in transfer_peer_records(peer_meta, src_block_ids): + layout = self._kv_layout_from_meta(meta) + self._validate_peer(meta, blocks, dst_block_ids) + if meta.get("heads_per_partition") != layout.local_num_heads(): + raise ValueError("peer heads_per_partition does not match its KV layout") + if int(meta["num_outer"]) % layout.local_num_layers(): + raise ValueError("peer num_outer is not divisible by its local layer count") + if layout.global_rank in peers_by_rank: + raise ValueError( + f"duplicate source global_rank={layout.global_rank} in KV metadata" + ) + sources.append(layout) + peers_by_rank[layout.global_rank] = (meta, blocks, layout) + if not sources: + raise ValueError("KV handoff contains no source peer metadata") + + local_planes = self._num_outer // self._layout.local_num_layers() + for transfer in plan_kv_reshard(sources, [self._layout]): + meta, blocks, source_layout = peers_by_rank[transfer.src_rank] + source_planes = int(meta["num_outer"]) // source_layout.local_num_layers() + if source_planes != local_planes: + raise ValueError( + f"outer-plane mismatch peer={source_planes} local={local_planes}" + ) + + src_layers = transfer.src_layer_slice(source_layout) + dst_layers = transfer.dst_layer_slice(self._layout) + src_heads = transfer.src_head_slice(source_layout) + dst_heads = transfer.dst_head_slice(self._layout) + layer_count = src_layers.stop - src_layers.start + head_count = src_heads.stop - src_heads.start + full_heads = ( + src_heads.start == 0 + and src_heads.stop == source_layout.local_num_heads() + and dst_heads.start == 0 + and dst_heads.stop == self._layout.local_num_heads() + and int(meta["bytes_per_slice"]) == self._bytes_per_slice + ) + full_layers = ( + src_layers.start == 0 + and src_layers.stop == source_layout.local_num_layers() + and dst_layers.start == 0 + and dst_layers.stop == self._layout.local_num_layers() + ) + if full_heads and full_layers: + xfer, ctx = self._begin_transfer( + meta, blocks, dst_block_ids, 0, 0, self._num_outer + ) + xfers.append(xfer) + contexts.append(ctx) + continue + if not full_heads and (int(meta["blocks_axis"]) != 2 or self._blocks_axis != 2): + raise NotImplementedError( + "KV head resharding requires the [2, L, B, T, H, d] layout" + ) + + for plane in range(local_planes): + xfer, ctx = self._begin_transfer( + meta, + blocks, + dst_block_ids, + plane * source_layout.local_num_layers() + src_layers.start, + plane * self._layout.local_num_layers() + dst_layers.start, + layer_count, + 0 if full_heads else src_heads.start, + 0 if full_heads else dst_heads.start, + 0 if full_heads else head_count, + ) + xfers.append(xfer) + contexts.append(ctx) + else: + records = transfer_peer_records(peer_meta, src_block_ids) + if len(records) != 1: + raise ValueError("matched-layout transfer requires exactly one source peer") + meta, blocks = records[0] + self._validate_peer(meta, blocks, dst_block_ids, matched_layout=True) + xfer, ctx = self._begin_transfer(meta, blocks, dst_block_ids, 0, 0, self._num_outer) + xfers.append(xfer) + contexts.append(ctx) + except Exception as exc: + if xfers: + cleanup = NixlPullHandle( + agent=self._agent, xfers=xfers, contexts=contexts, submitted_at=submitted_at + ) + try: + cleanup.wait() + except TimeoutError: + # Tell the owner not to recycle the destination storage while + # an already-submitted transfer may still write to it. + setattr(exc, "transfer_destinations_safe", False) + except Exception: + # Transfer errors are reported only after every submitted + # transfer has reached a terminal state. + pass + raise + return NixlPullHandle( + agent=self._agent, xfers=xfers, contexts=contexts, submitted_at=submitted_at + ) + + def _begin_transfer( + self, + 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.""" + pm = peer_meta + peer_base = pm["base_addr"] + peer_device_id = pm.get("device_id", 0) + peer_bps = pm["bytes_per_slice"] + peer_os = pm["outer_stride_bytes"] + peer_id = self._ensure_peer_registered(pm) + + bps = self._bytes_per_slice + local_os = self._outer_stride_bytes + + src_tuples: List[Any] = [] + dst_tuples: List[Any] = [] + + if n_heads == 0: + # One descriptor per block and outer slice. + for src_b, dst_b in zip(src_block_ids, dst_block_ids): + for i in range(n_outer): + src_o = src_o_start + i + dst_o = dst_o_start + i + src_tuples.append( + (peer_base + src_o * peer_os + src_b * peer_bps, peer_bps, peer_device_id) + ) + dst_tuples.append( + (self._buf_ptr + dst_o * local_os + dst_b * bps, bps, self._device_id) + ) + ctx = ( + f"matched peer={peer_id} outer[{src_o_start}:+{n_outer}] " + f"blocks={len(src_block_ids)}" + ) + else: + # Head sub-range copy: one descriptor per token. + assert self._head_dim is not None + assert self._heads_per_partition is not None + assert self._tokens_per_block is not None + d_bytes = self._head_dim * self._element_size + local_token_stride = self._heads_per_partition * d_bytes + peer_token_stride = pm["heads_per_partition"] * d_bytes + T = self._tokens_per_block + frag_bytes = n_heads * d_bytes + src_h_off = src_h0 * d_bytes + dst_h_off = dst_h0 * d_bytes + + for src_b, dst_b in zip(src_block_ids, dst_block_ids): + for i in range(n_outer): + src_o = src_o_start + i + dst_o = dst_o_start + i + src_slice = peer_base + src_o * peer_os + src_b * peer_bps + src_h_off + dst_slice = self._buf_ptr + dst_o * local_os + dst_b * bps + dst_h_off + for t in range(T): + src_tuples.append( + (src_slice + t * peer_token_stride, frag_bytes, peer_device_id) + ) + dst_tuples.append( + (dst_slice + t * local_token_stride, frag_bytes, self._device_id) + ) + ctx = ( + f"reshard peer={peer_id} outer[{src_o_start}:+{n_outer}] " + f"heads[{src_h0}:+{n_heads}] blocks={len(src_block_ids)}" + ) + + src_descs = self._agent.get_xfer_descs(src_tuples, mem_type="VRAM") + dst_descs = self._agent.get_xfer_descs(dst_tuples, mem_type="VRAM") + # READ pulls remote -> local. Signature is (op, local, remote, peer). + xfer = self._agent.initialize_xfer("READ", dst_descs, src_descs, peer_id) + try: + self._agent.transfer(xfer) + except Exception as exc: + # The transport may have accepted the operation before surfacing an + # error, so its destination cannot be proven safe for immediate reuse. + setattr(exc, "transfer_destinations_safe", False) + raise + return xfer, ctx + + def close(self) -> None: + """Release the registration and agent.""" + if self._agent is None: + return + try: + self._agent.deregister_memory(self._reg_handle) + except Exception: # noqa: BLE001 - shutdown path + logger.exception("NixlTransferBackend: deregister_memory failed") + self._agent = None diff --git a/megatron/core/inference/disaggregation/utils.py b/megatron/core/inference/disaggregation/utils.py index 9b5e153b443..779207e40f6 100644 --- a/megatron/core/inference/disaggregation/utils.py +++ b/megatron/core/inference/disaggregation/utils.py @@ -4,7 +4,7 @@ from __future__ import annotations -from typing import Optional, Tuple +from typing import Any, List, Optional, Tuple def intersect(a: Tuple[int, int], b: Tuple[int, int]) -> Optional[Tuple[int, int]]: @@ -14,7 +14,7 @@ def intersect(a: Tuple[int, int], b: Tuple[int, int]) -> Optional[Tuple[int, int def transfers_for_src(plan, src_rank): - """Transfers in ``plan`` originating from ``src_rank`` (any KV/Mamba + """Transfers in ``plan`` originating from ``src_rank`` (any KV/SSM reshard transfer -- both expose a ``src_rank`` field).""" return [t for t in plan if t.src_rank == src_rank] @@ -22,3 +22,36 @@ def transfers_for_src(plan, src_rank): def transfers_for_dst(plan, dst_rank): """Transfers in ``plan`` destined for ``dst_rank``.""" return [t for t in plan if t.dst_rank == dst_rank] + + +def transfer_peer_records(peer_meta: Any, src_block_ids: List[int]) -> List[Tuple[dict, List[int]]]: + """Normalize flat/TP/PP transfer metadata into peer/block records.""" + + def append_metas(raw_metas: Any, default_blocks: List[int]) -> None: + metas = raw_metas if isinstance(raw_metas, list) else [raw_metas] + for meta in metas: + if not isinstance(meta, dict): + raise ValueError("transfer peer metadata entries must be dictionaries") + blocks = meta.get("block_ids", default_blocks) + records.append((meta, [int(block) for block in blocks])) + + records: List[Tuple[dict, List[int]]] = [] + if isinstance(peer_meta, dict) and "pp_metas" in peer_meta: + for entry in peer_meta["pp_metas"]: + raw_metas = entry.get("tp_metas", entry) + blocks = [int(block) for block in entry.get("block_ids", [])] + append_metas(raw_metas, blocks) + return records + + if isinstance(peer_meta, dict) and "tp_metas" in peer_meta: + peer_meta = peer_meta["tp_metas"] + blocks = [int(block) for block in src_block_ids] + append_metas(peer_meta, blocks) + return records + + +def transfer_block_count(peer_meta: Any, src_block_ids: List[int]) -> int: + """Return the sequence-block count represented by transfer metadata.""" + + records = transfer_peer_records(peer_meta, src_block_ids) + return len(records[0][1]) if records else 0 diff --git a/tests/unit_tests/inference/test_kv_transfer_backends.py b/tests/unit_tests/inference/test_kv_transfer_backends.py new file mode 100644 index 00000000000..0f32bc2e5c9 --- /dev/null +++ b/tests/unit_tests/inference/test_kv_transfer_backends.py @@ -0,0 +1,197 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +import pytest +import torch + +from megatron.core.inference.disaggregation.ssm_reshard import SSMShardLayout, SSMStateDims +from megatron.core.inference.disaggregation.transfer_backends import base + + +def test_backend_registry_selects_by_explicit_name(): + assert base.construct_kv_transfer_backend_class("nixl").name == "nixl" + + try: + base.construct_kv_transfer_backend_class("unsupported") + except ValueError as exc: + assert "expected 'nixl'" in str(exc) + else: + raise AssertionError("unsupported backend should raise") + + +def test_ssm_geometry_uses_conv_and_recurrent_state_names(): + layout = SSMShardLayout( + global_rank=0, + tp_size=1, + tp_rank=0, + layer_start=0, + num_layers=1, + dims=SSMStateDims(nheads=2, headdim=4, d_state=5, ngroups=1, d_conv=3), + ) + memory_buffer = torch.zeros(1, 3, 2, 4, 5) + geometry = base.compute_buffer_geometry( + memory_buffer, + expected_num_blocks=3, + backend_name="test", + heads_per_partition=2, + ssm_layout=layout, + ssm_state_kind="recurrent", + ) + + metadata = base.export_geometry_meta(geometry, layout) + assert metadata["ssm_layout"]["dims"]["nheads"] == 2 + assert "mamba_layout" not in metadata + + with pytest.raises(ValueError, match="'conv' or 'recurrent'"): + base.compute_buffer_geometry( + memory_buffer, + expected_num_blocks=3, + backend_name="test", + heads_per_partition=2, + ssm_layout=layout, + ssm_state_kind="ssm", + ) + + +def test_nixl_direct_backend_exports_metadata_with_fake_agent(monkeypatch): + from megatron.core.inference.disaggregation.transfer_backends import nixl as nixl_mod + + class FakeAgent: + def __init__(self, name): + self.name = name + + def get_agent_metadata(self): + return b"agent-meta" + + def register_memory(self, tensor): + return ("reg", tuple(tensor.shape)) + + monkeypatch.setattr(nixl_mod, "_HAVE_NIXL", True) + monkeypatch.setattr(nixl_mod, "nixl_agent", FakeAgent) + + backend = nixl_mod.NixlTransferBackend( + "prefill", torch.zeros(2, 3, 5, dtype=torch.float32), expected_num_blocks=3 + ) + metadata = backend.export_meta() + + assert metadata["agent_name"] == "prefill" + assert metadata["bytes_per_slice"] == 20 + assert metadata["num_outer"] == 2 + assert metadata["num_blocks"] == 3 + assert metadata["blocks_axis"] == 1 + + +def test_nixl_begin_pull_blocks_uses_remote_metadata_with_fake_agent(monkeypatch): + from megatron.core.inference.disaggregation.transfer_backends import nixl as nixl_mod + + class FakeAgent: + def __init__(self, name): + self.name = name + self.transferred = False + + def get_agent_metadata(self): + return b"local" + + def register_memory(self, tensor): + return ("reg", tuple(tensor.shape)) + + def add_remote_agent(self, metadata): + assert metadata == b"remote" + return "peer" + + def get_xfer_descs(self, tuples, mem_type): + assert mem_type == "VRAM" + return tuples + + def initialize_xfer(self, op, local_desc, remote_desc, peer_id): + assert op == "READ" + assert peer_id == "peer" + return (local_desc, remote_desc) + + def transfer(self, xfer): + self.transferred = True + + def check_xfer_state(self, xfer): + assert self.transferred + return "DONE" + + monkeypatch.setattr(nixl_mod, "_HAVE_NIXL", True) + monkeypatch.setattr(nixl_mod, "nixl_agent", FakeAgent) + + backend = nixl_mod.NixlTransferBackend( + "decode", torch.zeros(2, 3, 5, dtype=torch.float32), expected_num_blocks=3 + ) + peer_meta = { + "agent_name": "prefill", + "agent_metadata_b64": "cmVtb3Rl", + "base_addr": 1234, + "bytes_per_slice": 20, + "num_outer": 2, + "outer_stride_bytes": 60, + "num_blocks": 3, + "device_id": 0, + "blocks_axis": 1, + } + backend.begin_pull_blocks(peer_meta, [1], [2]).wait() + + assert backend._agent.transferred is True + + +def test_nixl_begin_pull_blocks_returns_pollable_handle(monkeypatch): + from megatron.core.inference.disaggregation.transfer_backends import nixl as nixl_mod + + class FakeAgent: + def __init__(self, name): + self.name = name + self.transfers = 0 + self.polls = 0 + + def get_agent_metadata(self): + return b"local" + + def register_memory(self, tensor): + return ("reg", tuple(tensor.shape)) + + def add_remote_agent(self, metadata): + assert metadata == b"remote" + return "peer" + + def get_xfer_descs(self, tuples, mem_type): + assert mem_type == "VRAM" + return tuples + + def initialize_xfer(self, op, local_desc, remote_desc, peer_id): + assert op == "READ" + assert peer_id == "peer" + return {"local": local_desc, "remote": remote_desc} + + def transfer(self, xfer): + self.transfers += 1 + + def check_xfer_state(self, xfer): + self.polls += 1 + return "DONE" if self.polls >= 2 else "PENDING" + + monkeypatch.setattr(nixl_mod, "_HAVE_NIXL", True) + monkeypatch.setattr(nixl_mod, "nixl_agent", FakeAgent) + + backend = nixl_mod.NixlTransferBackend( + "decode", torch.zeros(2, 3, 5, dtype=torch.float32), expected_num_blocks=3 + ) + peer_meta = { + "agent_name": "prefill", + "agent_metadata_b64": "cmVtb3Rl", + "base_addr": 1234, + "bytes_per_slice": 20, + "num_outer": 2, + "outer_stride_bytes": 60, + "num_blocks": 3, + "device_id": 0, + "blocks_axis": 1, + } + + handle = backend.begin_pull_blocks(peer_meta, [1], [2]) + + assert backend._agent.transfers == 1 + assert backend._agent.polls == 0 + assert handle.poll() is False + assert handle.poll() is True diff --git a/tests/unit_tests/inference/test_nccl_transfer_backend.py b/tests/unit_tests/inference/test_nccl_transfer_backend.py new file mode 100644 index 00000000000..ecb5fb44b27 --- /dev/null +++ b/tests/unit_tests/inference/test_nccl_transfer_backend.py @@ -0,0 +1,138 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""Distributed unit test of the two-sided NCCL transfer backend. + +Prefill TP2 {0,1} -> decode TP1 {2} on real GPUs: the decode posts +begin_pull_blocks, the prefills post the matching begin_push_blocks, and the +decode's paged buffer must end up byte-identical to a direct shard of a known +global KV. Exercises the hetero head-merge through the same reshard plan the +NIXL backend uses. The test uses the process group provided by the unit-test +runner instead of spawning a nested distributed job. +""" + +import os + +import pytest +import torch + +L, H, HD, T, NB = 4, 8, 16, 8, 6 # layers, kv heads, head dim, tokens/block, pool blocks +BLOCKS = [1, 3] # the request's blocks (same ids both sides for simplicity) + + +def _global_blocks(): + """Global KV for the request's blocks: (block, kv, layer, token, head, dim) + with a distinct value per (block, kv, layer, head).""" + g = torch.zeros(len(BLOCKS), 2, L, T, H, HD) + for b in range(len(BLOCKS)): + for kv in range(2): + for l in range(L): + for h in range(H): + g[b, kv, l, :, h, :] = ((b * 2 + kv) * L + l) * 100 + h + return g + + +def _backend(rank, tp_size, tp_rank, device): + from megatron.core.inference.disaggregation.transfer_backends.nccl import NcclTransferBackend + + heads_local = H // tp_size + buf = torch.zeros(2, L, NB, T, heads_local, HD, device=device) + backend = NcclTransferBackend( + agent_name=f"test-rank{rank}", + memory_buffer=buf, + expected_num_blocks=NB, + tp_size=tp_size, + tp_rank=tp_rank, + num_kv_heads_global=H, + heads_per_partition=heads_local, + head_dim=HD, + tokens_per_block=T, + global_rank=rank, + pp_size=1, + pp_rank=0, + num_layers_global=L, + layer_start=0, + layer_end=L, + ) + return backend, buf + + +def _meta_stub(rank, tp_size, tp_rank): + """A rank's export_meta, built without its backend; the address fields are + unused by NCCL and the geometry is deterministic.""" + heads_local = H // tp_size + return { + "transport": "nccl", + "nccl_rank": rank, + "num_blocks": NB, + "blocks_axis": 2, + "num_outer": 2 * L, + "heads_per_partition": heads_local, + "head_dim": HD, + "tokens_per_block": T, + "element_size": 4, + "bytes_per_slice": T * heads_local * HD * 4, + "outer_stride_bytes": NB * T * heads_local * HD * 4, + "base_addr": 0, + "device_id": rank, + "global_rank": rank, + "tp_size": tp_size, + "tp_rank": tp_rank, + "pp_size": 1, + "pp_rank": 0, + "num_layers_global": L, + "num_kv_heads_global": H, + "layer_start": 0, + "layer_end": L, + } + + +@pytest.mark.skipif( + not ( + torch.cuda.is_available() + and torch.cuda.device_count() >= 3 + and int(os.environ.get("WORLD_SIZE", "1")) >= 3 + ), + reason="requires torchrun with >=3 CUDA ranks (prefill TP2 {0,1} + decode TP1 {2})", +) +def test_nccl_push_pull_tp2_to_tp1(): + import torch.distributed as dist + + local_rank = int(os.environ["LOCAL_RANK"]) + torch.cuda.set_device(local_rank) + device = f"cuda:{local_rank}" + if not dist.is_initialized(): + dist.init_process_group("nccl") + + rank = dist.get_rank() + control_group = dist.new_group(backend="gloo") + + # Initialize the default NCCL communicator collectively before ranks 0–2 + # use it for point-to-point transfers. Extra CI ranks synchronize through + # the Gloo control group and do not issue conflicting NCCL collectives. + dist.barrier() + + transfer_ok = True + g = _global_blocks().to(device) + if rank in (0, 1): # prefill TP2 + backend, buf = _backend(rank, 2, rank, device) + heads = slice(rank * (H // 2), (rank + 1) * (H // 2)) + for i, block in enumerate(BLOCKS): + # buffer layout [2, L, B, T, h, d] + buf[:, :, block] = g[i, :, :, :, heads, :] + # In production the decode's metas arrive in SEND_KV. + handle = backend.begin_push_blocks({"tp_metas": [_meta_stub(2, 1, 0)]}, BLOCKS) + handle.wait() + elif rank == 2: # decode TP1 + backend, buf = _backend(rank, 1, 0, device) + # In production the prefills' metas arrive in the hand-off kv_meta. + metas = [_meta_stub(0, 2, 0), _meta_stub(1, 2, 1)] + handle = backend.begin_pull_blocks({"tp_metas": metas}, BLOCKS, BLOCKS) + handle.wait() + expected = torch.zeros_like(buf) + for i, block in enumerate(BLOCKS): + expected[:, :, block] = g[i] + transfer_ok = torch.equal(buf, expected) + + result = torch.tensor(int(transfer_ok)) + dist.all_reduce(result, op=dist.ReduceOp.MIN, group=control_group) + assert result.item() == 1 diff --git a/tests/unit_tests/inference/test_mamba_reshard.py b/tests/unit_tests/inference/test_ssm_reshard.py similarity index 66% rename from tests/unit_tests/inference/test_mamba_reshard.py rename to tests/unit_tests/inference/test_ssm_reshard.py index 4a197813ab9..584f3a0213e 100644 --- a/tests/unit_tests/inference/test_mamba_reshard.py +++ b/tests/unit_tests/inference/test_ssm_reshard.py @@ -1,10 +1,10 @@ # Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. -"""Hetero TP/PP reshard of Mamba conv/ssm state (pure, CPU). +"""Hetero TP/PP reshard of SSM conv/recurrent state (pure, CPU). -Builds a known global Mamba state, shards it to a source (tp,pp) the exact way -mamba_mixer does ([x|B|C] conv bands + head-sharded ssm, layers split by PP), -runs plan_mamba_reshard to a different destination (tp,pp), and asserts every +Builds a known global SSM state, shards it to a source (tp,pp) the exact way +MambaMixer does ([x|B|C] conv bands + head-sharded recurrent state, layers +split by PP), runs ``plan_ssm_reshard`` to a different destination (tp,pp), and asserts every destination rank ends up byte-identical to a direct shard of the global state. This validates the band/layer index math against the real sharding model without a hybrid checkpoint (the residual gap is a real-model functional run). @@ -13,10 +13,10 @@ import pytest import torch -from megatron.core.inference.disaggregation.mamba_reshard import ( - MambaShardLayout, - MambaStateDims, - plan_mamba_reshard, +from megatron.core.inference.disaggregation.ssm_reshard import ( + SSMShardLayout, + SSMStateDims, + plan_ssm_reshard, ) @@ -26,17 +26,17 @@ def apply_conv_transfer(t, src_conv, dst_conv): dst_conv[t.dst_layer, t.dst_lo : t.dst_hi, :] = src_conv[t.src_layer, t.src_lo : t.src_hi, :] -def apply_ssm_transfer(t, src_ssm, dst_ssm): - """Copy an ssm sub-block in-memory; ssm is +def apply_recurrent_transfer(t, src_recurrent, dst_recurrent): + """Copy a recurrent-state sub-block in-memory; recurrent state is ``(num_layers, nheads_local, headdim, d_state)`` -- the band slices heads.""" - dst_ssm[t.dst_layer, t.dst_lo : t.dst_hi, :, :] = src_ssm[ + dst_recurrent[t.dst_layer, t.dst_lo : t.dst_hi, :, :] = src_recurrent[ t.src_layer, t.src_lo : t.src_hi, :, : ] # Global model dims (chosen divisible by the tp values under test). NHEADS, HEADDIM, DSTATE, NGROUPS, DCONV = 8, 4, 2, 2, 3 -M = 4 # global Mamba layers +M = 4 # global SSM layers D_INNER = NHEADS * HEADDIM # 32 G = NGROUPS * DSTATE # 4 (B and C band global size) CONV_DIM = D_INNER + 2 * G # 40 @@ -45,38 +45,38 @@ def apply_ssm_transfer(t, src_ssm, dst_ssm): def _global_state(): """Distinct value per (layer, channel, ...) so any mis-slice is caught.""" conv = torch.arange(M * CONV_DIM * DCONV, dtype=torch.float32).reshape(M, CONV_DIM, DCONV) - ssm = ( + recurrent = ( torch.arange(M * NHEADS * HEADDIM * DSTATE, dtype=torch.float32).reshape( M, NHEADS, HEADDIM, DSTATE ) + 10_000.0 ) - return conv, ssm + return conv, recurrent def _layouts(tp, pp): - """One MambaShardLayout per rank for a (tp, pp) instance; rank = p*tp + r. + """One SSMShardLayout per rank for a (tp, pp) instance; rank = p*tp + r. PP splits the M layers evenly (contiguous per stage).""" per = M // pp out = {} for p in range(pp): for r in range(tp): rank = p * tp + r - out[rank] = MambaShardLayout( + out[rank] = SSMShardLayout( global_rank=rank, tp_size=tp, tp_rank=r, layer_start=p * per, num_layers=per, - dims=MambaStateDims( + dims=SSMStateDims( nheads=NHEADS, headdim=HEADDIM, d_state=DSTATE, ngroups=NGROUPS, d_conv=DCONV ), ) return out -def _shard(conv_g, ssm_g, lay: MambaShardLayout): - """Shard the global state to one rank exactly as mamba_mixer does.""" +def _shard(conv_g, recurrent_g, lay: SSMShardLayout): + """Shard the global state to one rank exactly as MambaMixer does.""" s, e = lay.layer_range() r, tp = lay.tp_rank, lay.tp_size di_l = D_INNER // tp @@ -86,8 +86,8 @@ def _shard(conv_g, ssm_g, lay: MambaShardLayout): c = conv_g[s:e, D_INNER + G : D_INNER + 2 * G][:, r * g_l : (r + 1) * g_l] conv_l = torch.cat([x, b, c], dim=1).contiguous() nh_l = NHEADS // tp - ssm_l = ssm_g[s:e, r * nh_l : (r + 1) * nh_l, :, :].contiguous() - return conv_l, ssm_l + recurrent_l = recurrent_g[s:e, r * nh_l : (r + 1) * nh_l, :, :].contiguous() + return conv_l, recurrent_l @pytest.mark.parametrize( @@ -101,12 +101,12 @@ def _shard(conv_g, ssm_g, lay: MambaShardLayout): ((2, 1), (2, 1)), # identity ], ) -def test_mamba_reshard_reconstructs_destination(src, dst): - conv_g, ssm_g = _global_state() +def test_ssm_reshard_reconstructs_destination(src, dst): + conv_g, recurrent_g = _global_state() src_lay, dst_lay = _layouts(*src), _layouts(*dst) # Source per-rank tensors (as a prefill instance would hold them). - src_t = {rk: _shard(conv_g, ssm_g, lay) for rk, lay in src_lay.items()} + src_t = {rk: _shard(conv_g, recurrent_g, lay) for rk, lay in src_lay.items()} # Destination buffers, zero-filled at each rank's local shape. dst_t = {} for rk, lay in dst_lay.items(): @@ -115,71 +115,73 @@ def test_mamba_reshard_reconstructs_destination(src, dst): torch.zeros(lay.num_layers, lay.nheads_local, HEADDIM, DSTATE), ) - plan = plan_mamba_reshard(list(src_lay.values()), list(dst_lay.values())) + plan = plan_ssm_reshard(list(src_lay.values()), list(dst_lay.values())) for t in plan: if t.is_conv: apply_conv_transfer(t, src_t[t.src_rank][0], dst_t[t.dst_rank][0]) else: - apply_ssm_transfer(t, src_t[t.src_rank][1], dst_t[t.dst_rank][1]) + apply_recurrent_transfer(t, src_t[t.src_rank][1], dst_t[t.dst_rank][1]) # Every destination rank must match a direct shard of the global state. for rk, lay in dst_lay.items(): - want_conv, want_ssm = _shard(conv_g, ssm_g, lay) + want_conv, want_recurrent = _shard(conv_g, recurrent_g, lay) assert torch.equal(dst_t[rk][0], want_conv), f"conv mismatch at rank {rk} ({src}->{dst})" - assert torch.equal(dst_t[rk][1], want_ssm), f"ssm mismatch at rank {rk} ({src}->{dst})" + assert torch.equal( + dst_t[rk][1], want_recurrent + ), f"recurrent mismatch at rank {rk} ({src}->{dst})" -def test_mamba_rejects_indivisible_groups(): +def test_ssm_rejects_indivisible_groups(): """ngroups < tp_size would truncate the B/C bands to zero width; reject it up front instead of silently dropping state.""" with pytest.raises(ValueError): - MambaShardLayout( + SSMShardLayout( global_rank=0, tp_size=4, tp_rank=0, layer_start=0, num_layers=1, - dims=MambaStateDims(nheads=8, headdim=HEADDIM, d_state=DSTATE, ngroups=2, d_conv=DCONV), + dims=SSMStateDims(nheads=8, headdim=HEADDIM, d_state=DSTATE, ngroups=2, d_conv=DCONV), ) -def test_mamba_dedupes_replica_sources(): - """Two source ranks holding the same Mamba shard (same tp_rank+layer_start, +def test_ssm_dedupes_replica_sources(): + """Two source ranks holding the same SSM shard (same tp_rank+layer_start, e.g. EP/DP replicas) are deduped: the shard is sourced from exactly one of them (smallest global_rank), so no duplicate sends.""" def _lay(gr): - return MambaShardLayout( + return SSMShardLayout( global_rank=gr, tp_size=1, tp_rank=0, layer_start=0, num_layers=M, - dims=MambaStateDims( + dims=SSMStateDims( nheads=NHEADS, headdim=HEADDIM, d_state=DSTATE, ngroups=NGROUPS, d_conv=DCONV ), ) - plan = plan_mamba_reshard([_lay(0), _lay(1)], [_lay(2)]) + plan = plan_ssm_reshard([_lay(0), _lay(1)], [_lay(2)]) assert {t.src_rank for t in plan} == {0} # only the smallest-rank replica sources def test_layout_wire_roundtrip(): - """Layouts cross the coordinator as plain dicts (asdict) and are rebuilt via - MambaShardLayout(**dict); the nested dims dict must coerce back to - MambaStateDims so proxies (.headdim/.d_conv/...) keep working.""" + """Layouts cross the coordinator as plain dicts (asdict) and are rebuilt + via SSMShardLayout(**dict); the nested dims dict must coerce back to + SSMStateDims.""" import dataclasses - lay = MambaShardLayout( + lay = SSMShardLayout( global_rank=1, tp_size=2, tp_rank=1, layer_start=0, num_layers=M, - dims=MambaStateDims( + dims=SSMStateDims( nheads=NHEADS, headdim=HEADDIM, d_state=DSTATE, ngroups=NGROUPS, d_conv=DCONV ), ) - rebuilt = MambaShardLayout(**dataclasses.asdict(lay)) + rebuilt = SSMShardLayout(**dataclasses.asdict(lay)) assert rebuilt == lay - assert rebuilt.headdim == HEADDIM and rebuilt.d_conv == DCONV + assert rebuilt.conv_dim_local == lay.conv_dim_local