class NixlBaseConnectorWorker:
"""Base implementation of Worker side methods shared by pull and push."""
# Transfer mode included in the NIXL compatibility hash so that a push
# (WRITE) connector and a pull (READ) connector never handshake together.
# Overridden by NixlPushConnectorWorker.
_TRANSFER_MODE: str = "pull"
# Layer-name routing is supported only by NixlPushConnector for HMA or
# packed KV caches under PP.
_supports_pp_hma = False
def _compute_desc_ids(
self,
block_ids: BlockIds,
dst_num_blocks: int,
block_size_ratio: float | None,
physical_blocks_per_logical: int,
region_num_blocks: list[int] | None = None,
region_group_ids: list[int] | None = None,
uses_region_group_mapping: bool | None = None,
) -> np.ndarray:
"""Compute NIXL descriptor IDs for given block IDs."""
num_ssm_regions = 0
if self._has_mamba:
assert self._conv_decomp is not None
# NIXL regions per SSM layer = conv sub-projections + 1 SSM temporal
# (Mamba2/GDN: 3+1=4; Mamba1: 1+1=2).
ssm_regions_per_layer = len(self._conv_decomp.local_conv_offsets) + 1
# Count only the regions that actually hold SSM state; must match
# the descriptors emitted by _build_mamba_local.
num_ssm_regions = (
len(self._ssm_region_indices or self.block_len_per_layer)
* ssm_regions_per_layer
)
num_blocks = dst_num_blocks
if block_size_ratio is not None:
num_blocks = int(num_blocks * block_size_ratio)
use_layer_name_routing = bool(self._transfer_layer_group_ids)
num_desc_regions = (
len(self._transfer_layer_group_ids)
if use_layer_name_routing
else self.num_regions or len(self.block_len_per_layer)
)
if region_num_blocks is None:
region_num_blocks = [num_blocks] * num_desc_regions
elif block_size_ratio is not None:
region_num_blocks = [
int(count * block_size_ratio) for count in region_num_blocks
]
if region_num_blocks and not use_layer_name_routing:
num_desc_regions = len(region_num_blocks)
assert len(region_num_blocks) == num_desc_regions
region_offsets = np.cumsum([0, *region_num_blocks[:-1]])
num_fa_descs = sum(region_num_blocks)
if use_layer_name_routing:
# Descriptors in layer order, using each layer's KV-cache group.
return np.concatenate(
[
np.asarray(block_ids[g], dtype=np.int32) + region_offsets[k]
for k, g in enumerate(self._transfer_layer_group_ids)
]
)
# All-attention fast path.
if num_ssm_regions == 0:
if region_group_ids is None:
region_group_ids = self.region_group_ids
if uses_region_group_mapping is None:
uses_region_group_mapping = self._uses_region_group_mapping
if not region_group_ids and num_desc_regions == 1:
region_group_ids = [0]
assert len(region_group_ids) == num_desc_regions
if uses_region_group_mapping is None:
uses_region_group_mapping = len(set(region_group_ids)) > 1
region_group_ids_array = np.asarray(region_group_ids, dtype=np.int32)
# NOTE (NickLucche) With HMA, every kv group has the same number of layers
# and layers from different groups share the same kv tensor.
# eg block_ids=[[1, 2], [3]]->blocks [1, 2] need to be
# read across all regions, same for [3], but group0-group1 blocks will
# always differ (different areas). Therefore we can just flatten the
# block_ids and compute the descs ids for all groups at once.
if not uses_region_group_mapping:
block_arr = np.concatenate(
[np.asarray(group, dtype=np.int32) for group in block_ids]
)[None, :]
return (region_offsets[:, None] + block_arr).ravel()
desc_ids = []
for group_id, group in enumerate(block_ids):
if not group:
continue
region_ids = np.flatnonzero(
(region_group_ids_array == group_id)
| (region_group_ids_array == _SHARED_REGION_GROUP_ID)
)[:, None]
assert region_ids.size > 0
group_blocks = np.asarray(group, dtype=np.int32)[None, :]
desc_ids.append((region_offsets[region_ids] + group_blocks).ravel())
if not desc_ids:
return np.array([], dtype=np.int64)
return np.concatenate(desc_ids)
# Compute desc ids per group using the right stride: FA descs have
# num_blocks entries per region (kernel granularity, expanded by
# block_size_ratio for heterogeneous block sizes), SSM descs have
# logical_blocks entries per region (no kernel splitting, and never
# ratio-expanded since state blocks are indivisible).
logical_blocks = dst_num_blocks // physical_blocks_per_logical
all_descs: list[np.ndarray] = []
for i, group in enumerate(block_ids):
group_arr = np.asarray(group, dtype=np.int32)
spec_type = self._group_spec_types[i]
if _is_attention_spec(spec_type):
# A scratch cache lives only in its own regions; every other
# attention group spans all of them.
fa_region_ids = (
np.asarray(self._scratch_region_indices, dtype=np.int32)
if spec_type is CircularBufferSpec
else np.arange(self.num_regions, dtype=np.int32)
)[:, None]
all_descs.append(
(fa_region_ids * num_blocks + group_arr[None, :]).ravel()
)
elif _is_ssm_spec(spec_type):
# NOTE (NickLucche) SSM and Attention block regions can
# be exchanged arbitrarily by manager. Therefore, descs
# are laid out as:
# [descs_fa (all regions) | descs_ssm (all regions)].
# num_fa_descs offset must be computed per-engine since
# P and D can have different num_blocks (and thus
# different FA desc counts).
ssm_region_ids = (
np.arange(num_ssm_regions, num_ssm_regions + 1, dtype=np.int32)
if i == self._ple_group_index
else np.arange(num_ssm_regions, dtype=np.int32)
)[:, None]
all_descs.append(
(
ssm_region_ids * logical_blocks
+ group_arr[None, :]
+ num_fa_descs
).ravel()
)
else:
raise ValueError(
f"Unknown spec type {self._group_spec_types[i]} at index {i}"
)
return np.concatenate(all_descs)
def _build_local_splits_from_plan(
self,
plan: TPMapping,
src_blocks_data: np.ndarray,
num_fa_descs: int,
block_size_ratio: int = 1,
) -> Iterator[np.ndarray]:
"""Build split handle data for P_TP > D_TP scenario.
num_fa_descs is the boundary between FA and SSM descriptors.
Split counts are derived from source_ranks_per_group lengths.
FA uses rank_to_attention_slot for the slot offset;
SSM uses the rank's positional index.
Each Nx5 strided desc covers one region, so slicing it slices every block.
"""
fa_idx = next(
i for i, t in enumerate(self._group_spec_types) if _is_attention_spec(t)
)
fa_num_splits = len(plan.source_ranks_per_group[fa_idx])
starts = self._strided_desc_starts(src_blocks_data)
num_descs = int(src_blocks_data[:, 4].sum())
has_ssm_descs = num_fa_descs < num_descs
ssm_idx = next(
(i for i, t in enumerate(self._group_spec_types) if _is_ssm_spec(t)),
None,
)
ssm_num_splits = (
len(plan.source_ranks_per_group[ssm_idx])
if has_ssm_descs and ssm_idx is not None
else 0
)
sharded_desc_end = num_descs - (
self._logical_num_blocks if self._ple_region_index is not None else 0
)
assert num_fa_descs <= sharded_desc_end
# REPLICATE (MLA) descs are read whole on every rank; SPLIT (full-attn)
# descs take this rank's head slice. SSM descs before the PLE pages are
# sharded by the rank's positional index; PLE pages are copied whole.
is_fa = starts < num_fa_descs
fa_replicated = self._fa_desc_replicated(
starts[is_fa], num_fa_descs, block_size_ratio
)
assert block_size_ratio == 1 or fa_num_splits == 1 or fa_replicated.all(), (
"Head-sharded attention reads with P_TP > D_TP and heterogeneous "
"block sizes are not supported"
)
fa_split = is_fa.copy()
fa_split[is_fa] = ~fa_replicated
ssm_split = ~is_fa & (starts < sharded_desc_end)
for p_idx, p_rank in enumerate(plan.all_source_ranks):
fa_slot = plan.rank_to_attention_slot.get(p_rank, 0)
handle = src_blocks_data.copy()
handle[fa_split, 1] //= fa_num_splits
handle[fa_split, 0] += fa_slot * handle[fa_split, 1]
if ssm_num_splits:
handle[ssm_split, 1] //= ssm_num_splits
handle[ssm_split, 0] += p_idx * handle[ssm_split, 1]
yield handle
def _needs_split_local_xfer_handles(self, tp_ratio: int, plan: TPMapping) -> bool:
"""Whether reads need per-source slices of the local KV region.
Pure MLA attention is replicated across TP ranks and writes the whole
local region. Multiple physical remote workers may still participate
because DCP assigns them disjoint blocks, but that does not require
splitting the local region. Hybrid MLA+SSM is different: its mapping
contains multiple source ranks for the sharded SSM state.
"""
return tp_ratio < 0 and (not self.use_mla or len(plan.all_source_ranks) > 1)
def _fa_desc_replicated(
self,
desc_starts: np.ndarray,
num_fa_descs: int,
block_size_ratio: int = 1,
) -> np.ndarray:
"""Replicate flag per FA strided desc, given each one's flat block index.
_build_fa_local emits regions in order, with
``region_num_blocks[i] * block_size_ratio`` blocks per region. A strided
desc never spans regions, so the region owning its first block decides.
"""
assert self.transfer_topo is not None
n_regions = len(self.block_len_per_layer)
if n_regions == 0 or self.num_regions == 0:
return np.zeros(len(desc_starts), dtype=bool)
region_indices = self._transfer_layer_region_indices or range(n_regions)
region_ends = np.cumsum(
[self.region_num_blocks[i] * block_size_ratio for i in region_indices]
)
assert region_ends[-1] == num_fa_descs, (
f"FA region blocks {region_ends[-1]} != num_fa_descs {num_fa_descs}"
)
replicated = np.array([self._is_region_replicated(i) for i in region_indices])
return replicated[np.searchsorted(region_ends, desc_starts, side="right")]
def _is_region_replicated(self, region_idx: int) -> bool:
"""Whether region ``region_idx`` is transferred REPLICATE vs SPLIT.
REPLICATE (MLA): identical on every rank, whole block read from one
rank at offset 0, key-only. SPLIT (full-attn): head-sharded across TP.
Defaults to SPLIT when the per-region map is unset (e.g. tests that set
block_len_per_layer without register_kv_caches).
"""
return region_idx < len(self._region_is_mla) and self._region_is_mla[region_idx]
def _set_region_layers(self, region_layers: list[list[str]]) -> None:
layer_names = [layer_name for region in region_layers for layer_name in region]
assert len(layer_names) == len(set(layer_names)), (
"A KV cache layer spans multiple NIXL regions"
)
self.region_members: list[list[str]] = region_layers
if not self._requires_layer_name_routing():
return
group_by_layer = self.kv_cache_config.transfer_group_index_by_layer
ungrouped = [name for name in layer_names if name not in group_by_layer]
assert not ungrouped, f"KV cache layers outside any local group: {ungrouped}"
self._transfer_layer_names = tuple(layer_names)
self._transfer_layer_region_indices = tuple(
region_idx
for region_idx, region in enumerate(region_layers)
for _ in region
)
self._transfer_layer_group_ids = tuple(
group_by_layer[name] for name in layer_names
)
def _tracks_region_layers(self) -> bool:
"""Whether regions advertise their member layer names.
Push workers do so for HMA and packed layouts, so a PP producer can
match layers by name instead of by region index.
"""
return (
self._supports_pp_hma
and (self._is_hma_required or self._has_packed_cache)
and not self._has_mamba
)
def _requires_layer_name_routing(self) -> bool:
"""Whether PP push must match HMA/packed layers by name, not region index."""
return self.pp_size > 1 and self._tracks_region_layers()
def _align_remote_regions_by_layer(
self, nixl_agent_meta: NixlAgentMetadata
) -> None:
"""Select remote regions for this PP stage in local layer order."""
remote_region_layers = nixl_agent_meta.region_members
# Region-index routing here would silently transfer stale KV.
assert remote_region_layers, "Remote advertised no region_members"
assert (
len(nixl_agent_meta.kv_caches_base_addr)
== len(nixl_agent_meta.block_lens)
== len(nixl_agent_meta.block_strides)
== len(remote_region_layers)
), "Remote region metadata lengths disagree"
assert all(
values is None or len(values) == len(remote_region_layers)
for values in (
nixl_agent_meta.region_num_blocks,
nixl_agent_meta.region_group_ids,
nixl_agent_meta.region_names,
nixl_agent_meta.region_mem_types,
)
), "Remote region metadata lengths disagree"
remote_region_by_layer: dict[str, int] = {}
for region_idx, layer_names in enumerate(remote_region_layers):
for layer_name in layer_names:
assert layer_name not in remote_region_by_layer, (
f"Remote advertised layer {layer_name!r} in multiple regions"
)
remote_region_by_layer[layer_name] = region_idx
missing = [
name
for name in self._transfer_layer_names
if name not in remote_region_by_layer
]
assert not missing, f"Remote is missing locally owned layers: {missing}"
packed_regions = {
remote_region_by_layer[name]
for name in nixl_agent_meta.packed_member_layouts
}
remote_regions = [
remote_region_by_layer[name] for name in self._transfer_layer_names
]
nixl_agent_meta.kv_caches_base_addr = [
nixl_agent_meta.kv_caches_base_addr[i] for i in remote_regions
]
nixl_agent_meta.block_lens = [
nixl_agent_meta.block_lens[i] for i in remote_regions
]
nixl_agent_meta.block_strides = [
nixl_agent_meta.block_strides[i] for i in remote_regions
]
if nixl_agent_meta.region_num_blocks is not None:
nixl_agent_meta.region_num_blocks = [
nixl_agent_meta.region_num_blocks[i] for i in remote_regions
]
if nixl_agent_meta.region_group_ids is not None:
nixl_agent_meta.region_group_ids = [
nixl_agent_meta.region_group_ids[i] for i in remote_regions
]
if nixl_agent_meta.region_mem_types is not None:
nixl_agent_meta.region_mem_types = [
nixl_agent_meta.region_mem_types[i] for i in remote_regions
]
if nixl_agent_meta.region_names is not None:
nixl_agent_meta.region_names = list(self._transfer_layer_names)
if nixl_agent_meta.packed_member_layouts:
for i, layer_name in enumerate(self._transfer_layer_names):
if remote_regions[i] not in packed_regions:
continue
assert layer_name in nixl_agent_meta.packed_member_layouts, (
f"Remote packed layer {layer_name!r} has no layout"
)
offset, page_size = nixl_agent_meta.packed_member_layouts[layer_name]
assert (
0 <= offset < offset + page_size <= nixl_agent_meta.block_lens[i]
), f"Remote packed layer {layer_name!r} escapes its block"
local_region = self._transfer_layer_region_indices[i]
assert page_size == self.block_len_per_layer[local_region], (
f"Packed MLA page sizes must match for layer {layer_name!r}"
)
nixl_agent_meta.kv_caches_base_addr[i] += offset
nixl_agent_meta.block_lens[i] = page_size
nixl_agent_meta.packed_member_layouts = {}
# One layer per region keeps a second alignment pass a no-op.
nixl_agent_meta.region_members = [[name] for name in self._transfer_layer_names]
def _split_remote_packed_rows(self, nixl_agent_meta: NixlAgentMetadata) -> None:
"""Map the remote's whole-row region onto our per-layer regions.
E.g. a BLHNC prefiller registers one region whose block is a full row
holding every layer's page, while a HiSparse decoder registers one
region per layer. Using the advertised `packed_member_layouts`, the
remote row is rewritten in place as one virtual region per local layer:
remote row b: [ L0 page | L1 page | ... ] base=B, stride=S
^off0 ^off1
remote metadata, before: [(B, S, S)] # (base, len, stride)
remote metadata, after: [(B + off0, page0, S), # L0
(B + off1, page1, S)] # L1
Remote regions outside the row (e.g. a separate indexer) pass through.
"""
layouts = nixl_agent_meta.packed_member_layouts
remote_names = nixl_agent_meta.region_names
assert remote_names is not None, "Remote advertised no region_names"
(row,) = [i for i, name in enumerate(remote_names) if name in layouts]
if nixl_agent_meta.region_group_ids is not None and (
nixl_agent_meta.region_group_ids[row] == _SHARED_REGION_GROUP_ID
):
raise NotImplementedError(
"Reading multi-group packed KV rows into per-layer regions is not "
"supported; use the same KV cache layout on P and D."
)
remote_index = {name: i for i, name in enumerate(remote_names)}
regions = [
(row, *layouts[name])
if name in layouts
else (remote_index[name], 0, nixl_agent_meta.block_lens[remote_index[name]])
for name in self.region_names
]
nixl_agent_meta.kv_caches_base_addr = [
nixl_agent_meta.kv_caches_base_addr[i] + offset for i, offset, _ in regions
]
nixl_agent_meta.block_lens = [page_size for _, _, page_size in regions]
nixl_agent_meta.block_strides = [
nixl_agent_meta.block_strides[i] for i, _, _ in regions
]
for attr in ("region_num_blocks", "region_group_ids", "region_mem_types"):
values = getattr(nixl_agent_meta, attr)
if values is not None:
setattr(nixl_agent_meta, attr, [values[i] for i, _, _ in regions])
nixl_agent_meta.region_names = list(self.region_names)
nixl_agent_meta.packed_member_layouts = {}
def __init__(
self,
vllm_config: "VllmConfig",
engine_id: str,
kv_cache_config: "KVCacheConfig",
):
nixl_wrapper_cls = NixlWrapper
if nixl_wrapper_cls is None:
logger.error("NIXL is not available")
raise RuntimeError("NIXL is not available")
# Descriptor lists are prepared in NIXL's strided (Nx5) form.
check_nixl_version("1.5.0")
logger.info("Initializing NIXL wrapper")
logger.info("Initializing NIXL worker %s", engine_id)
# Config.
self.vllm_config = vllm_config
if vllm_config.kv_transfer_config is None:
raise ValueError("kv_transfer_config must be set for NixlConnector")
self.kv_transfer_config = vllm_config.kv_transfer_config
self.nixl_backends = vllm_config.kv_transfer_config.get_from_extra_config(
"backends", ["UCX"]
)
kv_lease_duration: int = vllm_config.kv_transfer_config.get_from_extra_config(
"kv_lease_duration", 30
)
# NOTE (NickLucche): For now we use a hardcoded value for a simpler interface.
self._lease_extension = kv_lease_duration * 2 // 3
self._bidirectional_kv_xfer_enabled: bool = (
vllm_config.kv_transfer_config.get_from_extra_config(
"bidirectional_kv_xfer", False
)
)
self.kv_cache_config = kv_cache_config
transfer_block_sizes = [
group.kv_cache_spec.block_size
for group in kv_cache_config.transfer_groups
if get_representative_spec_type(group.kv_cache_spec)
not in (MambaSpec, CircularBufferSpec)
]
self.block_size = (
math.lcm(*transfer_block_sizes)
if transfer_block_sizes
else cast(int, vllm_config.cache_config.block_size)
)
# Per-layer specs, unwrapping UniformTypeKVCacheSpecs group wrappers.
self._layer_specs: dict[str, KVCacheSpec] = {}
for group in kv_cache_config.transfer_groups:
group_spec = group.kv_cache_spec
if isinstance(group_spec, UniformTypeKVCacheSpecs):
self._layer_specs.update(group_spec.kv_cache_specs)
else:
self._layer_specs.update(dict.fromkeys(group.layer_names, group_spec))
self._is_hma_required = (
not vllm_config.scheduler_config.disable_hybrid_kv_cache_manager
and any(
not isinstance(spec, FullAttentionSpec)
for spec in self._layer_specs.values()
)
)
# ---- Model state (derived from model config) ----
mamba_ssm_size = (0, 0)
# Conv state sub-projection decomposition (None when no Mamba).
# The transfer requires DS (dim, state_len) conv layout so that
# conv sub-projections are contiguous in memory.
self._conv_decomp: MambaConvSplitInfo | None = None
self._has_mamba = any(
isinstance(spec, MambaSpec) for spec in self._layer_specs.values()
)
self._is_csa_linear = any(
get_representative_spec_type(group.kv_cache_spec) is CircularBufferSpec
for group in kv_cache_config.kv_cache_groups
)
self._ple_group_index: int | None = None
if self._is_csa_linear:
ple_groups = [
(index, group)
for index, group in enumerate(kv_cache_config.kv_cache_groups)
if group.layer_names
and all(
isinstance(spec, MambaSpec) and spec.tp_replicated
for spec in iter_layer_specs(group.kv_cache_spec)
)
]
if len(ple_groups) > 1 or any(
len(group.layer_names) != 1 for _, group in ple_groups
):
raise ValueError(
"CSA-linear NIXL requires at most one PLE cache owner."
)
self._ple_group_index = ple_groups[0][0] if ple_groups else None
if self._has_mamba:
assert self._is_hma_required
from vllm.model_executor.layers.mamba.mamba_utils import (
is_conv_state_dim_first,
)
assert is_conv_state_dim_first(), (
"3-read Mamba conv transfer requires DS conv state layout. "
"Set VLLM_SSM_CONV_STATE_LAYOUT=DS"
)
mamba_spec = next(
spec
for spec in self._layer_specs.values()
if isinstance(spec, MambaSpec)
and (not self._is_csa_linear or not spec.tp_replicated)
)
self._conv_decomp = derive_mamba_conv_split(
mamba_spec,
vllm_config.parallel_config.tensor_parallel_size,
)
mamba_ssm_size = self._conv_decomp.ssm_sizes
self._mamba_ssm_size = mamba_ssm_size
# Agent.
non_ucx_backends = [b for b in self.nixl_backends if b != "UCX"]
# Configure NIXL num_threads to avoid UAR exhaustion on Mellanox NICs.
# Each UCX thread allocates UARs (doorbell pages) via DevX, and
# excessive NIXL UAR usage can exhaust NIC UAR space. This can cause
# components like NVSHMEM (used by DeepEP kernels) to fail during RDMA
# initialization with "mlx5dv_devx_alloc_uar" errors.
# Ref: https://network.nvidia.com/files/doc-2020/ethernet-adapters-programming-manual.pdf#page=63
num_threads = vllm_config.kv_transfer_config.get_from_extra_config(
"num_threads", 4
)
if nixl_agent_config is None:
config = None
else:
# Enable telemetry by default for NIXL 0.7.1 and above.
config = (
nixl_agent_config(backends=self.nixl_backends, capture_telemetry=True)
if len(non_ucx_backends) > 0
else nixl_agent_config(num_threads=num_threads, capture_telemetry=True)
)
self.nixl_wrapper = nixl_wrapper_cls(str(uuid.uuid4()), config)
# Map of engine_id -> {(pp_rank, tp_rank): agent_name, ...}.
# non-PP remote uses pp_rank 0, i.e. (0, tp_rank).
self._remote_agents: dict[EngineId, dict[tuple[int, int], str]] = defaultdict(
dict
)
# Failed reads need a local metadata probe; invalid peers wait for drain.
self._failed_remote_engines: set[EngineId] = set()
self._invalid_remote_engines: set[EngineId] = set()
# Invalid engine -> (backoff, earliest release time). Kept across release
# so that repeat invalidations back off.
self._invalid_engine_backoff: dict[EngineId, tuple[float, float]] = {}
# Map of engine_id -> clock offset.
self._engine_clock_offset: dict[EngineId, float] = {}
# Metadata.
self.engine_id: EngineId = engine_id
self.tp_rank = get_tensor_model_parallel_rank()
self.world_size = get_tensor_model_parallel_world_size()
self.pcp_size = vllm_config.parallel_config.prefill_context_parallel_size
self.dcp_size = vllm_config.parallel_config.decode_context_parallel_size
self.pcp_rank = get_pcp_group().rank_in_group if self.pcp_size > 1 else 0
# TP1 PCP+DCP owns distinct KV shards; replicated PCP uses rank zero.
self.pcp_dcp_sharded = self.pcp_size > 1 and self.dcp_size > 1
self.transfer_tp_size = (
self.pcp_size if self.pcp_dcp_sharded else self.world_size
)
self.transfer_tp_rank = self.pcp_rank if self.pcp_dcp_sharded else self.tp_rank
# MLA is either fully replicated or fully sharded across transfer ranks.
self.dcp_rank = self.transfer_tp_rank % self.dcp_size
self.num_blocks = kv_cache_config.num_blocks
self.enable_permute_local_kv = False
self.enable_heterogeneous_attn_post_process = False
# KV Caches and nixl tracking data.
self.device_type = current_platform.device_type
self.kv_buffer_device: str = vllm_config.kv_transfer_config.kv_buffer_device
if self.device_type not in _NIXL_SUPPORTED_DEVICE:
raise RuntimeError(f"{self.device_type} is not supported.")
elif self.kv_buffer_device not in _NIXL_SUPPORTED_DEVICE[self.device_type]:
raise RuntimeError(
f"{self.device_type} with {self.kv_buffer_device} kv_buffer "
"is not supported."
)
self.device_kv_caches: dict[str, torch.Tensor] = {}
# cpu kv buffer for xfer
# used when device memory can not be registered under nixl
self.host_xfer_buffers: dict[str, torch.Tensor] = {}
if self.device_type == "cpu":
self.use_host_buffer = False
else:
self.use_host_buffer = self.kv_buffer_device == "cpu"
# reserve different cores for start_load_kv() from model_forward()
if self.device_type == "cpu":
numa_core_list = current_platform.discover_numa_topology()
# setup one last core in each numa for kv transfer.
rsv_cores_for_kv = [
max(each_numa_core_list) for each_numa_core_list in numa_core_list
]
if rsv_cores_for_kv:
if not hasattr(os, "sched_setaffinity"):
raise NotImplementedError(
"os.sched_setaffinity is not available on this platform"
)
os.sched_setaffinity(0, rsv_cores_for_kv)
# support for oot platform which can't register nixl memory
# type based on kv_buffer_device
nixl_memory_type = current_platform.get_nixl_memory_type()
if nixl_memory_type is None:
if self.kv_buffer_device in ["cuda", "xpu"]:
nixl_memory_type = "VRAM"
elif self.kv_buffer_device == "cpu":
nixl_memory_type = "DRAM"
if nixl_memory_type is None:
raise RuntimeError(
f"{self.device_type} with {self.kv_buffer_device} kv_buffer "
"is not supported."
)
self.nixl_memory_type = nixl_memory_type
# Note: host xfer buffer ops when use_host_buffer is True
self.copy_blocks: CopyBlocksOp | None = None
# Map of engine_id -> kv_caches_base_addr. For TP case, each local
self.device_id: int = 0
# Current rank may pull from multiple remote TP workers.
# EngineId, dict[int, list[int]] -> engine_id, tp_rank, base_addr_for_layer
self.kv_caches_base_addr = defaultdict[EngineId, dict[int, list[int]]](dict)
# Number of NIXL regions. Currently one region per cache
# (so 1 per layer for MLA, otherwise 2 per layer)
self.num_regions = 0
self.region_mem_types: list[str] = []
self.region_group_ids: list[int] = []
self._uses_region_group_mapping = False
self.region_names: list[str] = []
self.region_num_blocks: list[int] = []
self._mixed_mem_types = False
# True on TP ranks that share a HiSparse host pool with TP rank 0 and
# leave its NIXL registration and DRAM reads to rank 0.
self._skip_dram_xfer = False
self._desc_is_dram_by_block_size: dict[int, np.ndarray] = {}
self._desc_pos_by_block_size: dict[int, np.ndarray] = {}
self._dram_src_handles_by_block_size: dict[int, int | None] = {}
# PP>1 (push mode): this worker holds a contiguous layer slice and
# transfers into the matching sub-range of a PP=1 remote's regions.
self.pp_size = vllm_config.parallel_config.pipeline_parallel_size
self._remote_region_offset = 0
if self.pp_size > 1:
if self._is_hma_required and not self._supports_pp_hma:
raise NotImplementedError(
"NixlConnector (pull) does not support pipeline_parallel_size "
"> 1 with hybrid KV cache layouts (HMA); use NixlPushConnector "
"for PP + HMA."
)
# PP push routes HMA (hybrid) attention layouts by layer name;
# Mamba/SSM hybrids are not yet supported under PP.
if self._has_mamba:
raise NotImplementedError(
"NixlPushConnector does not support pipeline_parallel_size > 1 "
"with Mamba/SSM hybrid KV cache layouts yet."
)
# Decode-side PP is unsupported (completions counted per consumer rank).
if vllm_config.kv_transfer_config.kv_role == "kv_consumer":
raise NotImplementedError(
"NixlPushConnector consumer (decode) does not support "
"pipeline_parallel_size > 1."
)
# Keep heartbeat handshakes to a PP-sharded producer notif-only.
self._hb_handshake_notif_only = False
# nixl_prepped_dlist_handle.
self.src_xfer_handles_by_block_size: dict[int, int] = {}
# Local descriptor arrays per remote block size (block_size_ratio>1),
# kept for building per-tp-ratio splits at the same granularity.
self.src_blocks_data_by_block_size: dict[int, np.ndarray] = {}
# Populated dynamically during handshake based on remote configuration.
# Per-source split handles, keyed by (tp_ratio, remote_block_size).
self.src_xfer_handles_by_tp_ratio: dict[tuple[int, int], list[int]] = {}
self._dram_src_handles_by_tp_ratio: dict[tuple[int, int], list[int | None]] = {}
# Map of engine_id -> {tp_rank: nixl_prepped_dlist_handle (int)}.
self.dst_xfer_side_handles = defaultdict[EngineId, dict[int, int]](dict)
# Map of engine_id -> num_blocks. All ranks in the same deployment will
# have the same number of blocks.
self.dst_num_blocks: dict[EngineId, int] = {}
self.dst_region_num_blocks: dict[EngineId, list[int]] = {}
self.dst_region_group_ids: dict[EngineId, list[int]] = {}
self.dst_uses_region_group_mapping: dict[EngineId, bool] = {}
self.dst_region_mem_types: dict[EngineId, list[str]] = {}
self._registered_descs: list[Any] = []
# In progress transfers.
# [req_id -> list[handle]]
self._recving_metadata: dict[ReqId, ReqMeta] = {}
self._recving_transfers = defaultdict[ReqId, list[TransferHandle]](list)
# Track the expiration time of requests that are waiting to be sent.
self._reqs_to_send: dict[ReqId, float] = {}
# Replicated-KV PCP ranks > 0 never transfer; requests they are asked
# to send are reported finished immediately (see get_finished).
self._replicated_pcp_done_sending: set[ReqId] = set()
# Set of requests that have been part of a batch, regardless of status.
self._reqs_to_process: set[ReqId] = set()
# Invalid blocks from failed NIXL operations (thread-safe queue of block ids)
self._invalid_block_ids: queue.Queue[set[int]] = queue.Queue()
# requests that skipped transfer (handshake or transfer failures)
# Uses Queue for thread-safe cross-thread coordination with the
# background handshake thread, matching the _ready_requests pattern.
self._failed_recv_reqs: queue.Queue[ReqId] = queue.Queue()
self._recv_failures: set[ReqId] = set()
self._pending_recv_notifs: dict[ReqId, list[tuple[str, bytes]]] = {}
# Handshake metadata of this worker for NIXL transfers.
self.xfer_handshake_metadata: NixlHandshakePayload | None = None
# Background thread for initializing new NIXL handshakes.
self._handshake_initiation_executor = ThreadPoolExecutor(
# NIXL is not guaranteed to be thread-safe, limit 1 worker.
max_workers=1,
thread_name_prefix="vllm-nixl-handshake-initiator",
)
self._ready_requests = queue.Queue[tuple[ReqId, ReqMeta]]()
self._handshake_futures: dict[
EngineId, Future[tuple[dict[tuple[int, int], str], float]]
] = {}
# Protects _handshake_futures and _remote_agents.
self._handshake_lock = threading.RLock()
# TTL-based eviction of stale remote engine state.
self._engine_last_active: dict[EngineId, float] = {}
self._engine_ttl: float = vllm_config.kv_transfer_config.get_from_extra_config(
"engine_ttl", 3600.0
)
# Map of pull peer (host, port) -> engine_id currently serving it.
self._engine_by_address: dict[tuple[str, int], EngineId] = {}
self.model_config = vllm_config.model_config
self.use_mla = self.model_config.use_mla
# Get the attention backend from the first layer
# NOTE (NickLucche) models with multiple backends are not supported yet
self.attn_backends = get_current_attn_backends(vllm_config)
self.backend_name = self.attn_backends[0].get_name()
self.kv_cache_layout = (
vllm_config.cache_config.get_resolved_kv_cache_layout().name
)
self.host_buffer_kv_cache_layout = self.kv_cache_layout
logger.info(
"Detected attention backend(s) %s",
[backend.get_name() for backend in self.attn_backends],
)
logger.info("Detected kv cache layout %s", self.kv_cache_layout)
# lazy initialized in register_kv_caches
self.compat_hash: str | None = None
self.transfer_topo: TransferTopology | None = None
# With heterogeneous TP (or DCP), P must wait for all assigned D
# workers to finish reading before safely freeing the blocks.
self.consumer_notification_counts_by_req = defaultdict[ReqId, int](int)
self.expected_consumer_notifications_by_req: dict[ReqId, int] = {}
self.xfer_stats = NixlKVConnectorStats()
self._physical_blocks_per_logical_kv_block = 1
self._sync_block_size_with_kernel()
# Unwrap UniformTypeKVCacheSpecs to get the representative spec type
self._group_spec_types = tuple(
get_representative_spec_type(g.kv_cache_spec)
for g in self.kv_cache_config.transfer_groups
)
# Per-region MLA flag, 1:1 with block_len_per_layer. True -> REPLICATE
# (MLA), False -> SPLIT (head-sharded full-attn). Mixed only for models
# combining both (e.g. GQA main + MLA Eagle-3 draft).
self._region_is_mla = list[bool]()
self._ssm_region_indices = list[int]()
# Regions holding a scratch cache (the CSA compressor circular buffer).
# Tracked explicitly because a scratch page shares its address with the
# page it overlays, so its regions cannot be recovered from spec type.
self._scratch_region_indices = list[int]()
self._ple_region_index: int | None = None
# Page length of the PLE short-conv state; its region may be shared with
# pages of a different (smaller) size, so it cannot use the region's.
self._ple_block_len: int | None = None
# Enable different block lengths for different layers *only* when MLA is used.
# This is not used for SSM layers, which use the counterpart `mamba_ssm_size`.
self.block_len_per_layer = list[int]()
# Per-region block stride in bytes. Taken from the registered tensor's
# stride(0) so it stays correct under layouts that interleave layers
# within a block (BLHNC/BHLNC), where stride > block_len.
self.block_stride_per_layer = list[int]()
# Block-outermost layouts (BLHNC/BHLNC) pack every layer's page into one
# block row of a shared allocation.
self._has_packed_cache = KVCacheLayout[self.kv_cache_layout].is_block_outermost
# Region layer membership is populated for HMA and packed layouts.
self.region_members = []
self._packed_member_layouts: dict[str, tuple[int, int]] = {}
# Local layer order, set by ``_set_region_layers`` only when
# routing by layer name. Empty means region-index routing.
self._transfer_layer_names = tuple[str, ...]()
self._transfer_layer_region_indices = tuple[int, ...]()
self._transfer_layer_group_ids = tuple[int, ...]()
# Per-engine TP mappings. Generated during handshake.
self.tp_mappings: dict[EngineId, TPMapping] = {}
self.enforce_compat_hash = self.kv_transfer_config.get_from_extra_config(
"enforce_handshake_compat", True
)
def _validate_remote_parallel_config(
self, agent_metadata: NixlAgentMetadata
) -> None:
local_pcp_size = self.pcp_size
local_dcp_size = self.dcp_size
remote_pcp_size = agent_metadata.pcp_size
remote_dcp_size = agent_metadata.dcp_size
if remote_pcp_size > 1 and remote_dcp_size not in (1, remote_pcp_size):
raise NotImplementedError(
"Remote NixlConnector PCP+DCP does not span the full PCP group. "
f"Remote PCP/DCP={remote_pcp_size}/{remote_dcp_size}."
)
if (local_pcp_size > 1 and local_dcp_size == 1 and remote_dcp_size > 1) or (
remote_pcp_size > 1 and remote_dcp_size == 1 and local_dcp_size > 1
):
raise NotImplementedError(
"Replicated PCP cannot be paired with a DCP-sharded NIXL peer. "
f"Local PCP/DCP={local_pcp_size}/{local_dcp_size}; "
f"remote PCP/DCP={remote_pcp_size}/{remote_dcp_size}."
)
def _sync_block_size_with_kernel(self) -> None:
backends, specs = get_current_attn_backends_and_specs(
self.vllm_config, self.kv_cache_config, self.attn_backends
)
kernel_block_size = select_common_block_size(self.block_size, backends, specs)
# Number of blocks not accounting for kernel block mismatches
self._logical_num_blocks = self.num_blocks
if self.block_size != kernel_block_size:
logger.info_once(
"User-specified logical block size (%s) does not match"
" physical kernel block size (%s). Using the latter.",
self.block_size,
kernel_block_size,
)
assert self.block_size > kernel_block_size
self._physical_blocks_per_logical_kv_block = (
self.block_size // kernel_block_size
)
self.block_size = kernel_block_size
self.num_blocks *= self._physical_blocks_per_logical_kv_block
def _validate_csa_linear_tp_layout(self, remote_tp_size: int) -> None:
"""Reject P/D pairs whose main-KV pages have different head layouts.
A rank at TP < total KV heads holds several heads per page, while a
rank at TP >= total KV heads holds exactly one (replicated beyond
that). CSA-linear pages are shared tensors that transfer as whole
block units, so endpoints on opposite sides of that boundary lay
them out differently and cannot exchange them.
"""
assert self.transfer_topo is not None
total_kv_heads = self.transfer_topo.total_num_kv_heads
if (self.world_size < total_kv_heads) != (remote_tp_size < total_kv_heads):
raise ValueError(
"CSA-linear NIXL requires both endpoints on the same side of "
f"the KV-head sharding boundary ({total_kv_heads} KV heads): "
f"got local TP {self.world_size}, remote TP {remote_tp_size}."
)
def _nixl_handshake(
self,
host: str,
port: int,
remote_tp_size: int,
expected_engine_id: str,
remote_dcp_size: int = 1,
remote_pp_size: int = 1,
notif_agents_only: bool = False,
) -> tuple[dict[tuple[int, int], str], float]:
"""Do a NIXL handshake with a remote instance."""
if self._is_csa_linear:
self._validate_csa_linear_tp_layout(remote_tp_size)
# the first time we connect to a remote agent.
# be careful, the handshake happens in a background thread.
# it does not have an active cuda context until any cuda runtime
# call is made. when UCX fails to find a valid cuda context, it will
# disable any cuda ipc communication, essentially disabling any NVLink
# communication.
# when we are using device buffers, we need to set the device
# explicitly to make sure the handshake background thread has a valid
# cuda context.
if not self.use_host_buffer:
current_platform.set_device(self.device_id)
# When target instance TP > local TP, we need to perform multiple
# handshakes. Do it in a single background job for simplicity.
# Regardless, only handshake with the remote TP rank(s) that current
# local rank will read from. Note that With homogeneous TP,
# this happens to be the same single rank_i.
assert self.transfer_topo is not None
p_remote_ranks = self.transfer_topo.handshake_target_ranks(
remote_tp_size, remote_dcp_size
)
remote_rank_to_agent_name: dict[tuple[int, int], str] = {}
path = make_zmq_path("tcp", host, port)
# Clock offset to the peer, estimated from the handshake round-trip.
# Keep the lowest-RTT sample: hop cost is ~uniform across ranks, so a
# higher RTT is just noise that skews the midpoint estimate.
best_rtt = float("inf")
best_offset: float | None = None
with zmq_ctx(zmq.REQ, path) as sock:
for remote_pp_rank, remote_rank in itertools.product(
range(remote_pp_size), p_remote_ranks
):
logger.debug(
"Querying metadata on path: %s at remote pp rank %s, tp rank %s",
path,
remote_pp_rank,
remote_rank,
)
# Send query for the request.
msg = msgspec.msgpack.encode(
(GET_META_MSG, remote_pp_rank, remote_rank)
)
# Set receive timeout to 5 seconds to avoid hanging on dead server
sock.setsockopt(zmq.RCVTIMEO, 5000) # milliseconds
start_time = time.perf_counter()
sock.send(msg)
reply_parts = sock.recv_multipart()
recv_time = time.perf_counter()
assert len(reply_parts) == 2
handshake_bytes = reply_parts[0]
remote_perf = msgspec.msgpack.decode(reply_parts[1])
rtt = recv_time - start_time
if rtt < best_rtt:
best_rtt = rtt
best_offset = remote_perf - (start_time + recv_time) / 2
# Decode handshake payload to get compatibility hash
handshake_decoder = msgspec.msgpack.Decoder(NixlHandshakePayload)
try:
handshake_payload = handshake_decoder.decode(handshake_bytes)
except (msgspec.DecodeError, msgspec.ValidationError) as e:
raise RuntimeError(
f"Failed to decode NixlHandshakePayload. This likely indicates "
f"an incompatibility between connector version. Error: {e}"
) from e
got_metadata_time = time.perf_counter()
logger.debug(
"NIXL handshake: get metadata took: %s",
got_metadata_time - start_time,
)
# Check compatibility hash BEFORE decoding agent metadata
assert self.compat_hash is not None
if (
self.enforce_compat_hash
and handshake_payload.compatibility_hash != self.compat_hash
):
raise RuntimeError(
f"NIXL compatibility hash mismatch. "
f"Local: {self.compat_hash}, "
f"Remote: {handshake_payload.compatibility_hash}. "
f"Prefill and decode instances have incompatible "
f"configurations. This may be due to: different vLLM versions,"
f" models, dtypes, KV cache layouts, attention backends, etc. "
f"Both instances must use identical configurations."
f"Disable this check using "
f'--kv-transfer-config \'{{"kv_connector_extra_config": '
f'{{"enforce_handshake_compat": false}}}}\''
)
logger.info(
"NIXL compatibility check passed (hash: %s)",
handshake_payload.compatibility_hash,
)
# Decode agent metadata
metadata_decoder = msgspec.msgpack.Decoder(NixlAgentMetadata)
try:
metadata = metadata_decoder.decode(
handshake_payload.agent_metadata_bytes
)
except (msgspec.DecodeError, msgspec.ValidationError) as e:
# This should not happen if hash matched
raise RuntimeError(
f"Failed to decode NixlAgentMetadata. Error: {e}"
) from e
self._validate_remote_parallel_config(metadata)
# Ensure engine id matches.
if metadata.engine_id != expected_engine_id:
raise RuntimeError(
f"Remote NIXL agent engine ID mismatch. "
f"Expected {expected_engine_id},"
f"received {metadata.engine_id}."
)
# Register Remote agent.
if notif_agents_only:
remote_agent_name = self._add_notif_only_remote_agent(
metadata, remote_tp_size, metadata.dcp_size
)
else:
remote_agent_name = self.add_remote_agent(
metadata, remote_rank, remote_tp_size, metadata.dcp_size
)
setup_agent_time = time.perf_counter()
logger.debug(
"NIXL handshake: add agent took: %s (notif_agents_only=%s)",
setup_agent_time - got_metadata_time,
notif_agents_only,
)
remote_ranks = (remote_pp_rank, remote_rank)
remote_rank_to_agent_name[remote_ranks] = remote_agent_name
assert best_offset is not None
return remote_rank_to_agent_name, best_offset
def _add_notif_only_remote_agent(
self, metadata: NixlAgentMetadata, remote_tp_size: int, remote_dcp_size: int = 1
) -> str:
"""Load a remote agent for notifs only on the push-mode decode side.
Skips descriptor setup but records engine info for block accounting.
"""
assert self.transfer_topo is not None
self.transfer_topo.register_remote_engine(
metadata.engine_id,
EngineTransferInfo(
remote_tp_size=remote_tp_size,
remote_block_size=metadata.block_size,
remote_block_len=metadata.block_lens[0],
remote_physical_blocks_per_logical=(
metadata.physical_blocks_per_logical_kv_block
),
remote_dcp_size=remote_dcp_size,
),
)
return self.nixl_wrapper.add_remote_agent(metadata.agent_metadata)
def initialize_host_xfer_buffer(self, kv_caches: dict[str, torch.Tensor]) -> None:
"""Initialize transfer buffer in CPU mem for accelerators
NOT directly supported by NIXL (e.g., tpu)
"""
xfer_buffers: dict[str, torch.Tensor] = {}
try:
for layer_name, kv_cache in kv_caches.items():
kv_shape = kv_cache.shape
kv_dtype = kv_cache.dtype
permute_shape = False
inv_order = (0, 2, 1, 3)
if not self.use_mla:
assert kv_cache.ndim == 4
if self.kv_cache_layout == "LBNHC":
if self.kv_transfer_config.enable_permute_local_kv:
logger.info_once(
"'enable_permute_local_kv' flag is enabled while "
"device KV Layout is LBNHC. Init host buffer with"
" LBHNC to better support Decode/Prefill TP_ratio > 1."
)
# Since LBNHC will not support Decode/Prefill TP_ratio > 1,
# we can leverage host_buffer for permute.
self.host_buffer_kv_cache_layout = "LBHNC"
else:
# Packed KV layout is logical (B, H, N, 2*D). Allocate
# (B, N, H, 2*D) and view it as logical (B, H, N, 2*D)
# so raw NIXL transfers see NHD physical strides.
kv_shape = tuple(kv_shape[i] for i in inv_order)
permute_shape = True
xfer_buffers[layer_name] = torch.empty(
kv_shape, dtype=kv_dtype, device="cpu"
)
if permute_shape:
xfer_buffers[layer_name] = xfer_buffers[layer_name].permute(
inv_order
)
except MemoryError as e:
logger.error("NIXLConnectorWorker gets %s.", e)
raise
self.host_xfer_buffers = xfer_buffers
def set_host_xfer_buffer_ops(self, copy_operation: CopyBlocksOp):
"""Assign copy (d2h, h2d) operations when host buffer is used."""
# Set a no-op if the host buffer is not cpu.
if self.kv_buffer_device != "cpu":
return
# Set a no-op if self.device_type is 'cpu'.
if self.device_type == "cpu":
return
assert self.use_host_buffer
self.copy_blocks = copy_operation
def _log_failure(
self,
failure_type: str,
req_id: str | None,
msg: str = "",
error: BaseException | None = None,
meta: ReqMeta | None = None,
**extra_context,
):
"""Log transfer failure with structured context for easier debugging."""
context: dict[str, Any] = {
"failure_type": failure_type,
"request_id": req_id,
"engine_id": self.engine_id,
}
if meta is None and req_id is not None:
# Try to get metadata from in progress transfers when not provided
meta = self._recving_metadata.get(req_id)
if meta and meta.remote:
context.update(
{
"remote_engine_id": meta.remote.engine_id,
"remote_request_id": meta.remote.request_id,
"remote_host": meta.remote.host,
"remote_port": meta.remote.port,
"num_local_blocks": sum(
len(group) for group in meta.local_block_ids
),
"num_remote_blocks": sum(
len(group) for group in meta.remote.block_ids
),
"local_block_ids_sample": meta.local_block_ids[0][:10]
if meta.local_block_ids
else [],
}
)
context.update(extra_context)
if msg:
failure_type = f"{failure_type}. {msg}"
logger.error(
"NIXL transfer failure: %s | Context: %s",
failure_type,
context,
exc_info=error is not None,
stacklevel=2,
)
def _ensure_handshake(
self,
engine_id: EngineId,
host: str,
port: int,
tp_size: int,
dcp_size: int = 1,
pp_size: int = 1,
notif_agents_only: bool = False,
) -> Future[tuple[dict[tuple[int, int], str], float]] | None:
"""Ensure a handshake is in-flight (or already done) for *engine_id*.
Returns the ``Future`` if a handshake is pending (or was just
started), or ``None`` if the handshake already completed
successfully. Callers can attach per-request callbacks to the
returned future.
Failures to handshake are logged and the request is marked as failed.
"""
self._evict_stale_engines()
with self._handshake_lock:
if engine_id in self._remote_agents:
return None
fut = self._handshake_futures.get(engine_id)
if fut is not None:
return fut
fut = self._handshake_initiation_executor.submit(
self._nixl_handshake,
host,
port,
tp_size,
engine_id,
dcp_size,
pp_size,
notif_agents_only,
)
self._handshake_futures[engine_id] = fut
def done_callback(
f: Future[tuple[dict[tuple[int, int], str], float]],
eid=engine_id,
):
with self._handshake_lock:
del self._handshake_futures[eid]
try:
remote_agents, clock_offset = f.result()
self._remote_agents[eid] = remote_agents
self._engine_clock_offset[eid] = clock_offset
self._engine_last_active[eid] = time.perf_counter()
self._track_remote_engine_replacement(eid, host, port)
except Exception as e:
self._log_failure(
failure_type="handshake_setup_failed",
req_id=None,
error=e,
remote_engine_id=eid,
)
# Count once per handshake, regardless of waiting requests.
self.xfer_stats.record_failed_handshake()
fut.add_done_callback(done_callback)
return fut
def _background_nixl_handshake(
self, req_id: str, remote_engine_id: EngineId, meta: ReqMeta
):
# Do NIXL handshake in background and add to _ready_requests when done.
assert meta.remote is not None
fut = self._ensure_handshake(
remote_engine_id,
meta.remote.host,
meta.remote.port,
meta.tp_size,
meta.dcp_size,
meta.pp_size,
)
if fut is None:
# Already handshaked — only happens if caller does not pre-check.
self._ready_requests.put((req_id, meta))
return
# Check handshake success before proceeding with request.
def request_ready(f: Future[Any], entry=(req_id, meta)):
try:
f.result()
self._ready_requests.put(entry)
except Exception as e:
self._log_failure(
failure_type="handshake_failed",
req_id=req_id,
error=e,
meta=meta,
)
self._failed_recv_reqs.put(req_id)
fut.add_done_callback(request_ready)
def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
"""Register the KV Cache data in nixl."""
use_layer_name_routing = self._requires_layer_name_routing()
route_packed_layers = self._has_packed_cache and use_layer_name_routing
self.transfer_topo = TransferTopology(
tp_rank=self.transfer_tp_rank,
tp_size=self.transfer_tp_size,
block_size=self.block_size,
engine_id=self.engine_id,
is_mla=self.use_mla,
total_num_kv_heads=1
if self.use_mla
else self.model_config.get_total_num_kv_heads(),
attn_backends=self.attn_backends,
dcp_size=self.dcp_size,
# SSM States come in tuples (ssm, conv)
tensor_shape=next(iter(kv_caches.values())).shape
if not self._has_mamba
else None,
is_mamba=self._has_mamba,
)
backend_name = self.backend_name
if self._supports_pp_hma and self._has_packed_cache:
# PP can reorder the same backends across stages of a packed model.
backend_name = ",".join(
sorted({backend.get_name() for backend in self.attn_backends})
)
self.compat_hash = compute_nixl_compatibility_hash(
self.vllm_config,
backend_name,
transfer_mode=self._TRANSFER_MODE,
)
if self._is_csa_linear and self.use_host_buffer:
raise NotImplementedError(
"NIXL host staging does not preserve CSA-linear shared tensors."
)
if self.use_host_buffer:
self.initialize_host_xfer_buffer(kv_caches=kv_caches)
assert len(self.host_xfer_buffers) == len(kv_caches), (
f"host_buffer: {len(self.host_xfer_buffers)}, "
f"kv_caches: {len(kv_caches)}"
)
xfer_buffers = self.host_xfer_buffers
else:
xfer_buffers = kv_caches
assert not self.host_xfer_buffers, (
"host_xfer_buffer should not be initialized when "
f"kv_buffer_device is {self.kv_buffer_device}"
)
logger.info(
"Registering KV_Caches. use_mla: %s, kv_buffer_device: %s, "
"use_host_buffer: %s",
self.use_mla,
self.kv_buffer_device,
self.use_host_buffer,
)
registration_ranges: dict[tuple[int, str], tuple[int, int, int]] = {}
region_mem_types: list[str] = []
seen_base_addresses: list[int] = []
self._ssm_region_indices = []
self._scratch_region_indices = []
self._ple_block_len = None
self._ple_region_index = None
# A shared packed row may transfer whole only if every page in it is
# replicated across TP ranks: MLA without SSM state. Head-sharded
# attention or SSM state (e.g. Qwen3.8-Flash-Next) must stay split.
packed_storage = (
_share_storage_and_block_stride(list(xfer_buffers.values()))
and self.use_mla
and not self._has_mamba
)
if route_packed_layers and any(
not isinstance(spec, (MLAAttentionSpec, SlidingWindowMLASpec))
for spec in self._layer_specs.values()
):
raise NotImplementedError("PP push with packed KV caches requires MLA")
track_region_layers = self._tracks_region_layers()
region_layers: list[list[str]] = []
packed_member_layouts: dict[str, tuple[int, int]] = {}
# K and V are packed into the content dim, so each attention layer is a
# single NIXL region whose block transfers as one unit. Mamba layers instead
# register separate conv/ssm sub-regions (see `_build_mamba_local`).
# P and D may allocate equivalent transferable layers in different
# cache-group orders. Keep their region lists aligned without putting
# layers.10 before layers.2, which would break PP region slicing.
layer_names = (
xfer_buffers
if self._is_csa_linear
else sorted(
xfer_buffers,
key=_region_sort_key,
)
)
for layer_name in layer_names:
cache = xfer_buffers[layer_name]
# NOTE (NickLucche) Hybrid SSM mamba/FA physical page_size may differ when
# kernel requires a specific block size. This leads to SSM and FA layers
# having different num_blocks.
# `_physical_blocks_per_logical_kv_block` ratio is used to adjust for this.
layer_spec = self._layer_specs.get(layer_name)
if layer_spec is None:
logger.debug(
"Skipping layer %s as no KVCache spec is present. "
"This is likely because the layer is sharing its KV cache",
layer_name,
)
continue
group_index = self.kv_cache_config.transfer_group_index_by_layer[layer_name]
# `layer_spec.page_size_bytes` only accounts for logical page_size, that is
# the page_size assuming constant `self._logical_num_blocks`.
physical_page_size = (
layer_spec.page_size_bytes
if isinstance(layer_spec, MambaSpec)
else layer_spec.page_size_bytes
// self._physical_blocks_per_logical_kv_block
)
group = self.kv_cache_config.transfer_groups[group_index]
if group.host_resident:
logical_num_blocks = self.kv_cache_config.hisparse_host_num_blocks
assert logical_num_blocks is not None
else:
logical_num_blocks = self.kv_cache_config.num_blocks
group_id = group_index
num_blocks = (
logical_num_blocks
if isinstance(layer_spec, MambaSpec)
else logical_num_blocks * self._physical_blocks_per_logical_kv_block
)
is_mla_region = isinstance(
layer_spec, (MLAAttentionSpec, SlidingWindowMLASpec)
)
logger.debug(
"Registering layer %s with cache shape: %s", layer_name, cache.shape
)
storage = cache.untyped_storage()
storage_addr = storage.data_ptr()
is_host_resident = group.host_resident
if cache.device.type == "cpu":
mem_type = "DRAM"
region_device_id = 0
else:
mem_type = self.nixl_memory_type
region_device_id = max(cache.get_device(), 0)
self.device_id = region_device_id
if is_host_resident:
registration_base = cache.data_ptr()
registration_len = cache.nbytes
else:
registration_base = storage_addr
registration_len = storage.nbytes()
storage_key = (registration_base, mem_type)
if storage_key not in registration_ranges:
registration_ranges[storage_key] = (
registration_base,
registration_base + registration_len,
region_device_id,
)
if isinstance(layer_spec, MambaSpec):
physical_ratio = self._physical_blocks_per_logical_kv_block
block_len = physical_page_size // physical_ratio
block_stride = physical_page_size
if self._is_csa_linear and cache.ndim != 1:
# CSA-linear packs multiple cache owners in each block, so
# state views advance by the packed block stride.
block_stride = cache.stride(0) * cache.element_size()
assert block_stride % physical_ratio == 0
region_specs = [
(cache.data_ptr(), block_len, block_stride // physical_ratio)
]
else:
if cache.ndim == 1:
# Flat byte view: HMA tensors shared between layer types carry
# no block dimension, so stride(0) is 1 byte. Blocks abut, so
# the stride is the page size.
block_stride = physical_page_size
else:
block_stride = cache.stride(0) * cache.element_size()
# A connector that registers the backing as one memory region
# may need its length rounded up to a page boundary. Such a tail
# is shorter than one block, so it cannot make the storage
# segmented.
storage_is_block_major = (
0 <= registration_len - num_blocks * block_stride < block_stride
)
# A view with a contiguous page can own a strided region even
# when sibling layers share the same backing allocation.
page_contiguous = (
cache.ndim > 1
and cache[0].is_contiguous()
and cache[0].nbytes == physical_page_size
)
# Block-outermost MLA: one block row holds every layer's page.
packed_mla = (
self._has_packed_cache
and not self._has_mamba
and packed_storage
and storage_is_block_major
and is_mla_region
)
if packed_mla and use_layer_name_routing:
# PP>1 prefill stage: register only this layer's page of each row.
region_specs = [
(cache.data_ptr(), physical_page_size, block_stride)
]
elif packed_mla and (self.pp_size == 1 or not page_contiguous):
# Register whole rows so a block is a single descriptor. Also,
# advertise where this layer's page sits so a remote with
# LB* can still address it (HiSparse).
storage_block_len = block_stride
region_specs = [
(registration_base, storage_block_len, storage_block_len)
]
offset = cache.data_ptr() - storage_addr
assert offset >= 0 and offset + physical_page_size <= block_stride
packed_member_layouts[layer_name] = (offset, physical_page_size)
elif storage_is_block_major and (
not page_contiguous
and ((packed_storage and is_mla_region) or not self._is_csa_linear)
):
# TODO(Lucas): handle TP slicing for packed_storage; for now
# restrict to MLA (DSv4) where kv is replicated.
# The complete storage row is one block stride. Dividing the
# registered length instead would fold the padding tail in.
storage_block_len = block_stride
region_specs = [
(registration_base, storage_block_len, storage_block_len)
]
elif storage_is_block_major:
region_specs = [
(cache.data_ptr(), physical_page_size, block_stride)
]
else:
segment_bytes = num_blocks * block_stride
if cache.nbytes % segment_bytes != 0:
raise AssertionError(
"KV cache view cannot be partitioned into NIXL regions: "
f"layer={layer_name}, cache_nbytes={cache.nbytes}, "
f"num_blocks={num_blocks}, block_stride={block_stride}, "
f"physical_page_size={physical_page_size}, "
f"cache_shape={tuple(cache.shape)}, "
f"cache_stride={tuple(cache.stride())}"
)
num_segments = cache.nbytes // segment_bytes
region_block_len = (
block_stride if num_segments > 1 else physical_page_size
)
region_specs = [
(
cache.data_ptr() + segment_idx * segment_bytes,
region_block_len,
block_stride,
)
for segment_idx in range(num_segments)
]
for base_addr, block_len, block_stride in region_specs:
# A packed PP producer keeps one region per layer, so aliases at
# one address (a layer and its SWA view) stay distinct.
if base_addr in seen_base_addresses and not route_packed_layers:
region_index = seen_base_addresses.index(base_addr)
assert region_mem_types[region_index] == mem_type
self._region_is_mla[region_index] |= is_mla_region
if is_mla_region:
self.block_len_per_layer[region_index] = block_len
self.block_stride_per_layer[region_index] = block_stride
self.region_num_blocks[region_index] = num_blocks
if self.region_group_ids[region_index] != group_id:
self.region_group_ids[region_index] = _SHARED_REGION_GROUP_ID
else:
region_index = len(seen_base_addresses)
seen_base_addresses.append(base_addr)
self.block_len_per_layer.append(block_len)
self.block_stride_per_layer.append(block_stride)
self.region_group_ids.append(group_id)
self.region_names.append(layer_name)
self.region_num_blocks.append(num_blocks)
self._region_is_mla.append(is_mla_region)
region_mem_types.append(mem_type)
if track_region_layers:
if region_index == len(region_layers):
region_layers.append([layer_name])
else:
region_layers[region_index].append(layer_name)
if isinstance(layer_spec, MambaSpec):
if layer_spec.tp_replicated:
assert self._is_csa_linear
assert self._ple_region_index in (None, region_index)
self._ple_region_index = region_index
self._ple_block_len = block_len
elif region_index not in self._ssm_region_indices:
self._ssm_region_indices.append(region_index)
elif (
isinstance(layer_spec, CircularBufferSpec)
and region_index not in self._scratch_region_indices
):
self._scratch_region_indices.append(region_index)
# When there's a mismatch between kbs<>bs, we rely on HMA to ensure
# caches are either [NB, PS] or [NB*r, PS/r] where r is bs/kbs.
if (
self._physical_blocks_per_logical_kv_block == 1
and cache.shape[0] < num_blocks
):
raise AssertionError(
"KV cache tensor has too few "
f"blocks; layer={layer_name}, "
f"expected_num_blocks={num_blocks}, "
f"cache_shape={tuple(cache.shape)}, "
f"cache_stride={tuple(cache.stride())}, "
f"layer_spec={type(layer_spec).__name__}, "
f"backend={self.backend_name}, "
"all_backends="
f"{[backend.get_name() for backend in self.attn_backends]}, "
f"kv_cache_layout={self.kv_cache_layout}"
)
logger.debug(
"Different block lengths collected: %s", set(self.block_len_per_layer)
)
assert (
len(self.block_len_per_layer)
== len(seen_base_addresses)
== len(self._region_is_mla)
== len(self.block_stride_per_layer)
== len(self.region_group_ids)
== len(self.region_names)
== len(self.region_num_blocks)
)
# Descriptor ids must be region-ordered, matching the remote side.
self._scratch_region_indices.sort()
self.kv_caches_base_addr[self.engine_id][self.tp_rank] = seen_base_addresses
self.num_regions = len(seen_base_addresses)
self.region_mem_types = region_mem_types
self._uses_region_group_mapping = len(set(self.region_group_ids)) > 1
self._mixed_mem_types = len(set(region_mem_types)) > 1
if self._has_mamba:
self.region_num_blocks = [self.num_blocks] * self.num_regions
self._set_region_layers(region_layers)
self._packed_member_layouts = packed_member_layouts
if self.pp_size > 1 and not use_layer_name_routing:
start_layer, end_layer = self.model_config.get_layers_start_end_indices(
self.vllm_config.parallel_config
)
num_local_layers = end_layer - start_layer
assert num_local_layers > 0 and self.num_regions % num_local_layers == 0
regions_per_layer = self.num_regions // num_local_layers
self._remote_region_offset = regions_per_layer * start_layer
# Total local FA descriptors (boundary between FA and mamba descs).
xfer_region_num_blocks = (
[self.region_num_blocks[i] for i in self._transfer_layer_region_indices]
if self._transfer_layer_region_indices
else self.region_num_blocks
)
self.num_descs = sum(xfer_region_num_blocks)
self._mixed_mem_types = len(set(region_mem_types)) > 1
self._skip_dram_xfer = self._leaves_dram_xfer_to_rank0()
if self._mixed_mem_types:
assert self.use_mla and not self._has_mamba, (
"Mixed-device KV registration is only supported for MLA "
"models without Mamba layers."
)
for mem_type in sorted(set(region_mem_types)):
if mem_type == "DRAM" and self._skip_dram_xfer:
continue
ranges_for_mem_type = [
(start, end - start, device_id, "")
for (_, cache_mem_type), (
start,
end,
device_id,
) in registration_ranges.items()
if cache_mem_type == mem_type
]
descs = self.nixl_wrapper.get_reg_descs(ranges_for_mem_type, mem_type)
self.nixl_wrapper.register_memory(descs, backends=self.nixl_backends)
self._registered_descs.append(descs)
self.device_kv_caches = kv_caches
self.dst_num_blocks[self.engine_id] = self.num_blocks
self.dst_region_num_blocks[self.engine_id] = xfer_region_num_blocks
self.dst_region_group_ids[self.engine_id] = self.region_group_ids
self.dst_uses_region_group_mapping[self.engine_id] = (
self._uses_region_group_mapping
)
self.dst_region_mem_types[self.engine_id] = self.region_mem_types
if self._has_mamba:
logger.info(
"Hybrid SSM registration: num_blocks=%s, "
"logical_num_blocks=%s, ratio=%s, num_regions=%s, "
"num_descs=%s, mamba_ssm_size=%s, block_len_per_layer=%s",
self.num_blocks,
self._logical_num_blocks,
self._physical_blocks_per_logical_kv_block,
self.num_regions,
self.num_descs,
self._mamba_ssm_size,
set(self.block_len_per_layer),
)
# Register local/src descr for NIXL xfer.
self.src_xfer_handles_by_block_size[self.block_size], self.src_blocks_data = (
self.register_local_xfer_handler(self.block_size)
)
# After KV Caches registered, listen for new connections.
agent_metadata = NixlAgentMetadata(
engine_id=self.engine_id,
agent_metadata=self.nixl_wrapper.get_agent_metadata(),
device_id=self.device_id,
kv_caches_base_addr=self.kv_caches_base_addr[self.engine_id][self.tp_rank],
num_blocks=self.num_blocks,
block_lens=self.block_len_per_layer,
block_strides=self.block_stride_per_layer,
kv_cache_layout=self.kv_cache_layout
if not self.use_host_buffer
else self.host_buffer_kv_cache_layout,
block_size=self.block_size,
ssm_sizes=self._mamba_ssm_size,
ple_block_len=self._ple_block_len,
attn_backend_name=self.backend_name,
physical_blocks_per_logical_kv_block=(
self._physical_blocks_per_logical_kv_block
),
region_num_blocks=self.region_num_blocks,
region_group_ids=self.region_group_ids,
region_names=self.region_names,
region_mem_types=self.region_mem_types,
dcp_size=self.dcp_size,
pcp_size=self.pcp_size,
region_members=self.region_members,
packed_member_layouts=self._packed_member_layouts,
)
# Wrap metadata in payload with hash for defensive decoding
assert self.compat_hash is not None
encoder = msgspec.msgpack.Encoder()
self.xfer_handshake_metadata = NixlHandshakePayload(
compatibility_hash=self.compat_hash,
agent_metadata_bytes=encoder.encode(agent_metadata),
)
def _ple_page_len(self) -> int:
"""Transfer length of one PLE state page.
The PLE page shares its region with the first page of other cache
groups (block-outer layout), whose block_len may be smaller; the
descriptors must cover the whole PLE page or part of the conv state
never reaches the decode side."""
assert self._ple_region_index is not None
assert self._ple_block_len is not None, (
"PLE region registered without its page length"
)
return self._ple_block_len
def _build_mamba_local(self, base_addresses: list[int]) -> np.ndarray:
"""Build desc regions (conv sub-projections + ssm) per layer for
local mamba blocks with DS conv layout, as Nx5 uint64 strided descs.
A Mamba block interleaves conv and SSM state, which crucially differ in
size, so the two are indexed as separate sub-regions. Attention blocks
instead pack K and V into the content dim and transfer as a single unit.
Reference diagram:
KVCacheTensor (Shared)
/ \\
/ \\
/ \\
Attention (FlashInfer) View Mamba View
| |
| |
+-------------------+ +-------------------+
| KVCacheTensor | | KVCacheTensor |
| | | |
|<----- page ------>| |<----- page ------->|
| size | | size |
| Key 0 | Val 0 | |Conv 0 | SSM 0 |
| Key 1 | Val 1 | |Conv 1 | SSM 1 |
| ... | ... | | ... | ... |
| Key N-2 | Val N-2 | |Conv N-2| SSM N-2 |
| Key N-1 | Val N-1 | |Conv N-1| SSM N-1 |
+-------------------+ +--------------------+
|1st_split-2nd_split| |1st_split-2nd_split |
Mamba state blocks are indivisible (not token-extent data), so the
descriptors always use the local page geometry regardless of any
attention block-size ratio; their desc ids are likewise never
ratio-expanded (see _compute_desc_ids).
"""
assert base_addresses, "Local KV cache base addresses must not be empty."
assert self._conv_decomp is not None
conv_offsets = self._conv_decomp.local_conv_offsets
conv_size, ssm_size = self._mamba_ssm_size
num_blocks = self._logical_num_blocks
physical_per_logical = self._physical_blocks_per_logical_kv_block
device_id = self.device_id
strided_descs: list[tuple[int, int, int, int, int]] = []
region_indices = self._ssm_region_indices or range(len(base_addresses))
for i in region_indices:
base_addr = base_addresses[i]
block_stride = self.block_stride_per_layer[i] * physical_per_logical
for off, sz in conv_offsets:
strided_descs.append(
(base_addr + off, sz, device_id, block_stride, num_blocks)
)
# SSM temporal state follows the conv state.
strided_descs.append(
(base_addr + conv_size, ssm_size, device_id, block_stride, num_blocks)
)
if (region_index := self._ple_region_index) is not None:
block_len = self._ple_page_len() * physical_per_logical
block_stride = (
self.block_stride_per_layer[region_index] * physical_per_logical
)
base_addr = base_addresses[region_index]
strided_descs.append(
(base_addr, block_len, device_id, block_stride, num_blocks)
)
return self._pack_strided_descs(strided_descs)
def _build_mamba_remote(
self,
nixl_agent_meta: NixlAgentMetadata,
tp_ratio: int,
transfer_info: EngineTransferInfo,
) -> np.ndarray:
"""Build remote desc regions (conv sub-projections + ssm) per layer.
For hetero-TP, each D rank reads only its sub-projection slice from
the P rank. Returns Nx5 uint64 strided descs."""
assert nixl_agent_meta.kv_caches_base_addr, (
"Remote KV cache base addresses must not be empty."
)
assert self._conv_decomp is not None
effective_ratio = max(tp_ratio, 1)
# Mamba conv state is always TP-sharded, even when attention KV
# is replicated (num_kv_heads < tp_size).
local_offset = self.tp_rank % effective_ratio
conv_size_remote = nixl_agent_meta.ssm_sizes[0]
conv_offsets = self._conv_decomp.remote_conv_offsets(local_offset, tp_ratio)
if tp_ratio >= 1:
ssm_read_size = self._mamba_ssm_size[1]
else:
ssm_read_size = nixl_agent_meta.ssm_sizes[1]
remote_physical_per_logical = transfer_info.remote_physical_blocks_per_logical
num_blocks = nixl_agent_meta.num_blocks // remote_physical_per_logical
device_id = nixl_agent_meta.device_id
strided_descs: list[tuple[int, int, int, int, int]] = []
# NOTE (ZhanqiuHu): use per-layer block_lens[i], not [0], in case
# block lengths vary across layers (e.g. MLA).
region_indices = self._ssm_region_indices or range(
len(nixl_agent_meta.kv_caches_base_addr)
)
for i in region_indices:
base_addr = nixl_agent_meta.kv_caches_base_addr[i]
block_stride = (
nixl_agent_meta.block_strides[i] * remote_physical_per_logical
)
for off, sz in conv_offsets:
strided_descs.append(
(base_addr + off, sz, device_id, block_stride, num_blocks)
)
# SSM temporal state is also TP-sharded on the heads dimension.
ssm_addr = base_addr + conv_size_remote + local_offset * ssm_read_size
strided_descs.append(
(ssm_addr, ssm_read_size, device_id, block_stride, num_blocks)
)
if (region_index := self._ple_region_index) is not None:
local_block_len = (
self._ple_page_len() * self._physical_blocks_per_logical_kv_block
)
# Same NIXL_CONNECTOR_VERSION on both sides, so the field is set.
assert nixl_agent_meta.ple_block_len is not None
remote_block_len = (
nixl_agent_meta.ple_block_len * remote_physical_per_logical
)
if local_block_len != remote_block_len:
raise ValueError(
"PLE pages require identical P/D geometry: "
f"local={local_block_len}, remote={remote_block_len}."
)
remote_block_stride = (
nixl_agent_meta.block_strides[region_index]
* remote_physical_per_logical
)
addr = nixl_agent_meta.kv_caches_base_addr[region_index]
strided_descs.append(
(addr, remote_block_len, device_id, remote_block_stride, num_blocks)
)
return self._pack_strided_descs(strided_descs)
@staticmethod
def _pack_strided_descs(
strided_descs: list[tuple[int, int, int, int, int]],
) -> np.ndarray:
"""Pack (addr, len, dev_id, stride, count) rows into an Nx5 uint64 array.
Each row describes `count` blocks of `len` bytes whose starts are
`stride` bytes apart; NIXL >= 1.5 prepares this layout directly.
"""
return np.array(strided_descs, dtype=np.uint64).reshape(-1, 5)
@staticmethod
def _block_strided_descs(
base_addr: int,
block_len: int,
block_stride: int,
num_blocks: int,
block_size_ratio: int,
device_id: int,
) -> np.ndarray:
"""Nx5 strided descs for `num_blocks` blocks, each split into
`block_size_ratio` sub-blocks of `block_len` bytes.
Evenly spaced sub-blocks collapse into one strided desc. In a padded
block the sub-blocks abut only within the block, so each block gets
its own.
"""
if block_size_ratio == 1 or block_stride == block_size_ratio * block_len:
stride = block_stride // block_size_ratio
count = num_blocks * block_size_ratio
return np.array(
[[base_addr, block_len, device_id, stride, count]], dtype=np.uint64
)
strided_descs = np.empty((num_blocks, 5), dtype=np.uint64)
strided_descs[:, 0] = (
base_addr + np.arange(num_blocks, dtype=np.uint64) * block_stride
)
strided_descs[:, 1] = block_len
strided_descs[:, 2] = device_id
strided_descs[:, 3] = block_len
strided_descs[:, 4] = block_size_ratio
return strided_descs
@staticmethod
def _strided_desc_starts(strided_descs: np.ndarray) -> np.ndarray:
"""Flat block index of each strided desc's first block."""
counts = strided_descs[:, 4]
return np.cumsum(counts) - counts
def _build_fa_local(
self,
base_addresses: list[int],
block_size_ratio: int,
) -> np.ndarray:
"""Build local FA descriptors as Nx5 uint64 strided descs."""
assert self.transfer_topo is not None
assert base_addresses, "Local KV cache base addresses must not be empty."
device_id = self.device_id
region_indices = self._transfer_layer_region_indices or range(
len(base_addresses)
)
parts = [
self._block_strided_descs(
base_addresses[i],
self.block_len_per_layer[i] // block_size_ratio,
self.block_stride_per_layer[i],
self.region_num_blocks[i],
block_size_ratio,
device_id,
)
for i in region_indices
]
return np.concatenate(parts)
def _build_fa_remote(
self,
plan: TPMapping,
nixl_agent_meta: NixlAgentMetadata,
block_size_ratio: int,
) -> np.ndarray:
"""Build remote FA descriptors as Nx5 uint64 strided descs."""
assert self.transfer_topo is not None
assert nixl_agent_meta.kv_caches_base_addr, (
"Remote KV cache base addresses must not be empty."
)
fa_group_idx = next(
i for i, t in enumerate(self._group_spec_types) if _is_attention_spec(t)
)
# SPLIT regions read their head slice from this many remote ranks at a
# per-rank offset; REPLICATE regions read the whole block once.
split_reads = len(plan.source_ranks_per_group[fa_group_idx])
region_num_blocks = nixl_agent_meta.region_num_blocks or [
nixl_agent_meta.num_blocks
] * len(nixl_agent_meta.kv_caches_base_addr)
device_id = nixl_agent_meta.device_id
block_strides = nixl_agent_meta.block_strides
strided_descs: list[tuple[int, int, int, int, int]] = []
use_layer_name_routing = bool(self._transfer_layer_region_indices)
for i, base_addr in enumerate(nixl_agent_meta.kv_caches_base_addr):
local_region = (
self._transfer_layer_region_indices[i] if use_layer_name_routing else i
)
replicated = self._is_region_replicated(local_region)
# Read our whole local region size from remote..
local_block_len = self.block_len_per_layer[local_region]
remote_kv_block_len = local_block_len // block_size_ratio
if block_size_ratio > 1:
# ..using remote kv_block_len as transfer unit
local_block_len = remote_kv_block_len
# REPLICATE reads the whole block once at offset 0; SPLIT gathers
# its head slice from `split_reads` remote ranks at a per-rank offset.
num_reads = 1 if replicated else split_reads
rank_offset = (
0 if replicated else plan.rank_offset_factor * remote_kv_block_len
)
local_block_len = local_block_len // num_reads
addr = base_addr + rank_offset
count = region_num_blocks[i]
strided_descs.append(
(addr, local_block_len, device_id, block_strides[i], count)
)
return self._pack_strided_descs(strided_descs)
def register_local_xfer_handler(
self,
block_size: int,
) -> tuple[int, np.ndarray]:
"""Function used for register local xfer handler with local block_size or
Remote block_size.
When local block_size is same as remote block_size, we use local block_size
to register local_xfer_handler during init.
When remote block size is less than local block size, we need to use
register another local_xfer_handler using remote block len to ensure
data copy correctness.
"""
assert self.transfer_topo is not None
block_size_ratio = self.block_size // block_size
local_base_addresses = self.kv_caches_base_addr[self.engine_id][self.tp_rank]
blocks_data = self._build_fa_local(local_base_addresses, block_size_ratio)
logger.debug(
"Created %s blocks for src engine %s and rank %s on device id %s",
blocks_data[:, 4].sum(),
self.engine_id,
self.tp_rank,
self.device_id,
)
if self._has_mamba:
assert self.num_descs * block_size_ratio == blocks_data[:, 4].sum()
# TODO (ZhanqiuHu): For homogeneous TP (tp_ratio == 1), the 3-descs split
# is unnecessary — a single conv desc per block suffices. Consider
# adding a fast path that falls back to the standard 2-region
# registration (_build_fa_local mamba=True) when no hetero-TP
# remote has been seen. Currently we always register 4 regions
# because local descs are created before knowing the remote TP.
logger.debug("Registering local Mamba descriptors (4 regions/layer)")
mamba = self._build_mamba_local(local_base_addresses)
blocks_data = np.concatenate([blocks_data, mamba])
if self._mixed_mem_types:
desc_is_dram = np.concatenate(
[
np.full(
count * block_size_ratio,
mem_type == "DRAM",
dtype=bool,
)
for mem_type, count in zip(
self.region_mem_types,
self.region_num_blocks,
strict=True,
)
]
)
assert len(desc_is_dram) == blocks_data[:, 4].sum()
desc_pos = np.empty(len(desc_is_dram), dtype=np.int64)
dram_idx = np.where(desc_is_dram)[0]
vram_idx = np.where(~desc_is_dram)[0]
desc_pos[dram_idx] = np.arange(len(dram_idx), dtype=np.int64)
desc_pos[vram_idx] = np.arange(len(vram_idx), dtype=np.int64)
self._desc_is_dram_by_block_size[block_size] = desc_is_dram
self._desc_pos_by_block_size[block_size] = desc_pos
# DRAM descriptors are registered under CPU device 0.
# A strided desc never spans regions, so its first block decides.
strided_is_dram = desc_is_dram[self._strided_desc_starts(blocks_data)]
blocks_data[strided_is_dram, 2] = 0
if self._skip_dram_xfer:
self._dram_src_handles_by_block_size[block_size] = None
else:
self._dram_src_handles_by_block_size[block_size] = (
self.nixl_wrapper.prep_xfer_dlist(
"NIXL_INIT_AGENT",
blocks_data[strided_is_dram],
mem_type="DRAM",
)
)
device_handle = self.nixl_wrapper.prep_xfer_dlist(
"NIXL_INIT_AGENT",
blocks_data[~strided_is_dram],
mem_type=self.nixl_memory_type,
)
return device_handle, blocks_data
# NIXL_INIT_AGENT to be used for preparations of local descs.
handle = self.nixl_wrapper.prep_xfer_dlist(
"NIXL_INIT_AGENT", blocks_data, mem_type=self.region_mem_types[0]
)
return handle, blocks_data
def add_remote_agent(
self,
nixl_agent_meta: NixlAgentMetadata,
remote_tp_rank: int = 0,
remote_tp_size: int = 1,
remote_dcp_size: int = 1,
) -> str:
"""Add the remote NIXL agent and prepare the descriptors for reading cache
blocks from remote.
In particular, handle both homogeneous and heterogeneous TP. The former
requires local rank_i to read from remote rank_i.
The latter, in the case of D.world_size < P.world_size, requires that a
local (D) TP worker reads from multiple remote (P) TP workers.
Conversely, assuming D.world_size > P.world_size, two or more local TP
workers will read from a single remote TP worker.
Here's an example for the last case described above (non-MLA):
rank_offset p_remote_tp_rank
(kv split no)
--------------------------------
0 0 Worker0 ---- 1st half of KV ----> Worker0 [ KV Cache ]
/
1 0 Worker1 ---- 2nd half of KV -----/
0 1 Worker2 ---- 1st half of KV ----> Worker1 [ KV Cache ]
/
1 1 Worker3 ---- 2nd half of KV -----/
Decoder TP workers Prefix TP workers
(world_size=4) (world_size=2)
tp_ratio = 4 // 2 = 2
Considering the KV Caches, if P-Worker_i has cache size [2, num_blocksP, kv_heads, block_size, head_dim]
then D-Worker_j has [2, num_blocksD, kv_heads//tp_ratio, block_size, head_dim]. Mind the "LBHNC" layout format.
Assuming num_blocksD >= num_blocksP, D-Worker0 reads from P-Worker0 by preparing the kv_heads//tp_ratio
first heads from all the slots of all the blocks. D-Worker1 will do the same, but reading the second split
along the kv_heads dimension, and so forth until "tp_ratio" D TP workers have pulled from P-Worker0.
Note that the above will also hold true for the homogeneous TP case, where tp_ratio evaluates to 1.
Regarding MLA case, the cache is replicated across TP workers so the rank_offset will just always be 0
so that the whole cache is shared by "tp_ratio" D TP workers.
For Mamba hetero-TP, both tp_ratio > 0 (D_TP > P_TP) and
tp_ratio < 0 (P_TP > D_TP) are supported by the 3-read transfer.
""" # noqa: E501
engine_id = nixl_agent_meta.engine_id
# TODO re-evaluate refreshing for scaling/recovery
if (0, remote_tp_rank) in self._remote_agents.get(engine_id, {}):
logger.debug(
"Remote agent with engine_id %s and rank"
"%s already exchanged metadata, skip handshake.",
engine_id,
remote_tp_rank,
)
return self._remote_agents[engine_id][(0, remote_tp_rank)]
assert self.transfer_topo is not None
transfer_topo = self.transfer_topo
# Number of physical regions registered locally (one per layer/tensor).
num_local_regions = len(self.block_len_per_layer)
if self._transfer_layer_region_indices:
if self.block_size != nixl_agent_meta.block_size:
raise NotImplementedError(
"Attention-HMA push requires identical P/D block sizes."
)
if remote_tp_size > transfer_topo.tp_size and not self.use_mla:
raise NotImplementedError(
"Attention-HMA push does not support decode TP greater than "
"prefill TP yet"
)
self._align_remote_regions_by_layer(nixl_agent_meta)
elif nixl_agent_meta.packed_member_layouts and not self._packed_member_layouts:
# With HiSparse, P may be on BL* while D is LB*, as only indexer is on HBM.
# TODO (NickLucche) we may consider not allowing this case at all (P BL*)
# and just clear this function.
self._split_remote_packed_rows(nixl_agent_meta)
elif (
self.pp_size > 1
and len(nixl_agent_meta.kv_caches_base_addr) > num_local_regions
):
# This worker holds a PP layer-slice; the PP=1 remote registered
# the full model. Slice its regions to our layer window so the
# logic below sees congruent local/remote lists.
start = self._remote_region_offset
end = start + num_local_regions
assert len(nixl_agent_meta.kv_caches_base_addr) >= end
nixl_agent_meta.kv_caches_base_addr = nixl_agent_meta.kv_caches_base_addr[
start:end
]
nixl_agent_meta.block_lens = nixl_agent_meta.block_lens[start:end]
nixl_agent_meta.block_strides = nixl_agent_meta.block_strides[start:end]
if nixl_agent_meta.region_num_blocks is not None:
nixl_agent_meta.region_num_blocks = nixl_agent_meta.region_num_blocks[
start:end
]
if nixl_agent_meta.region_group_ids is not None:
nixl_agent_meta.region_group_ids = nixl_agent_meta.region_group_ids[
start:end
]
if nixl_agent_meta.region_names is not None:
nixl_agent_meta.region_names = nixl_agent_meta.region_names[start:end]
if nixl_agent_meta.region_mem_types is not None:
nixl_agent_meta.region_mem_types = nixl_agent_meta.region_mem_types[
start:end
]
if (
self._packed_member_layouts
and not self._transfer_layer_region_indices
and self._packed_member_layouts != nixl_agent_meta.packed_member_layouts
):
# Whole rows transfer as one unit, so both sides must pack them alike.
raise RuntimeError(
"Packed KV block rows differ between P and D. Use the same "
"VLLM_KV_CACHE_LAYOUT and block size on both sides."
)
### Register remote engine in TransferTopology (idempotent).
physical_blocks_per_logical = (
nixl_agent_meta.physical_blocks_per_logical_kv_block
)
transfer_info = EngineTransferInfo(
remote_tp_size=remote_tp_size,
remote_block_size=nixl_agent_meta.block_size,
remote_block_len=nixl_agent_meta.block_lens[0],
remote_physical_blocks_per_logical=physical_blocks_per_logical,
remote_dcp_size=remote_dcp_size,
)
transfer_topo.register_remote_engine(engine_id, transfer_info)
logger.info("Transfer plan: %s", transfer_topo.describe(engine_id))
self.tp_mappings[engine_id] = compute_tp_mapping(
transfer_topology=transfer_topo,
remote_tp_size=remote_tp_size,
group_spec_types=self._group_spec_types,
remote_dcp_size=remote_dcp_size,
)
remote_agent_name = self.nixl_wrapper.add_remote_agent(
nixl_agent_meta.agent_metadata
)
# Create dst descs and xfer side handles. TP workers have same #blocks
# so we only register once per engine_id.
# Example:
# block_size_ratio > 1:
# remote: | 0| 1| 2| 3| 4| 5| 6| 7| 8| 9|10|11|12|
# local origin:| 0| 1| 8| 12|
# local mapped:| 0| 1| 2| 3| 4| 5| 6| 7| 8| 9|10|11|12|13|14|15|
block_size_ratio = transfer_topo.block_size_ratio(nixl_agent_meta.block_size)
if engine_id not in self.dst_num_blocks:
num_remote_regions = len(nixl_agent_meta.kv_caches_base_addr)
self.dst_num_blocks[engine_id] = nixl_agent_meta.num_blocks
self.dst_region_num_blocks[engine_id] = (
nixl_agent_meta.region_num_blocks
or [nixl_agent_meta.num_blocks] * num_remote_regions
)
remote_region_group_ids = nixl_agent_meta.region_group_ids or (
self.region_group_ids
if len(self.region_group_ids) == num_remote_regions
else [0] * num_remote_regions
)
self.dst_region_group_ids[engine_id] = remote_region_group_ids
self.dst_uses_region_group_mapping[engine_id] = (
len(set(self.dst_region_group_ids[engine_id])) > 1
)
self.dst_region_mem_types[engine_id] = (
nixl_agent_meta.region_mem_types
or [self.nixl_memory_type] * num_remote_regions
)
# Keep track of remote agent kv caches base addresses.
self.kv_caches_base_addr[engine_id][remote_tp_rank] = (
nixl_agent_meta.kv_caches_base_addr
)
self._validate_remote_agent_handshake(
nixl_agent_meta, remote_tp_size, remote_dcp_size
)
# This is 1 when P and D `--tensor-parallel-size` match. Otherwise,
# this is the ratio between the two sizes.
tp_ratio = transfer_topo.tp_ratio(remote_tp_size)
logger.debug(
"Registering remote agent (%s, rank %s) memory regions with tp_ratio %s",
engine_id,
remote_tp_rank,
tp_ratio,
)
plan = self.tp_mappings[engine_id]
### (Optional) Register a local handler at the remote engine's block
### granularity (remote/prefill blocks smaller than local).
remote_block_size = nixl_agent_meta.block_size
src_blocks_data = self.src_blocks_data
if block_size_ratio > 1:
if remote_block_size not in self.src_xfer_handles_by_block_size:
handle, blocks_data = self.register_local_xfer_handler(
remote_block_size
)
self.src_xfer_handles_by_block_size[remote_block_size] = handle
self.src_blocks_data_by_block_size[remote_block_size] = blocks_data
src_blocks_data = self.src_blocks_data_by_block_size[remote_block_size]
### (Optional) Register local agent memory regions. MLA is not split.
split_key = (tp_ratio, remote_block_size)
if self._needs_split_local_xfer_handles(tp_ratio, plan) and (
split_key not in self.src_xfer_handles_by_tp_ratio
):
# Remote tp_size > local tp_size: read from multiple remote ranks.
# Logically "split" own regions into per-source chunks. Hybrid
# MLA+SSM also needs this path: MLA is replicated and read once,
# while the SSM state is sharded across every remote TP rank.
# We only do this once per remote (tp_size, block_size).
self.src_xfer_handles_by_tp_ratio[split_key] = []
if self._mixed_mem_types:
self._dram_src_handles_by_tp_ratio[split_key] = []
for handle_data in self._build_local_splits_from_plan(
plan,
src_blocks_data,
self.num_descs * block_size_ratio,
block_size_ratio,
):
if self._mixed_mem_types:
desc_is_dram = self._desc_is_dram_by_block_size[remote_block_size]
strided_is_dram = desc_is_dram[
self._strided_desc_starts(handle_data)
]
if self._skip_dram_xfer:
dram_handle = None
else:
dram_handle = self.nixl_wrapper.prep_xfer_dlist(
"NIXL_INIT_AGENT",
handle_data[strided_is_dram],
mem_type="DRAM",
)
self._dram_src_handles_by_tp_ratio[split_key].append(dram_handle)
handle_data = handle_data[~strided_is_dram]
handle = self.nixl_wrapper.prep_xfer_dlist(
"NIXL_INIT_AGENT", handle_data, mem_type=self.nixl_memory_type
)
self.src_xfer_handles_by_tp_ratio[split_key].append(handle)
### Register remote agent memory regions
# With homogeneous TP, D pulls the whole kv cache from corresponding rank. With
# heterogeneous TP, prepare the descriptors by splitting the P KV cache along
# kv_head dim, of D worker's kv_head size (D>P).
# Eg. PTP1 DTP2 => P0 KV:[block0-KV_0 | block0-KV_1..].
# Register all remote blocks, but only the corresponding kv heads.
blocks_data = self._build_fa_remote(
plan,
nixl_agent_meta,
block_size_ratio,
)
logger.debug(
"Created %s blocks for dst engine %s with remote rank %s and local rank %s",
blocks_data[:, 4].sum(),
engine_id,
remote_tp_rank,
self.tp_rank,
)
if self._has_mamba:
logger.debug(
"Registering remote Mamba blocks for engine %s rank %s",
engine_id,
remote_tp_rank,
)
mamba = self._build_mamba_remote(nixl_agent_meta, tp_ratio, transfer_info)
blocks_data = np.concatenate([blocks_data, mamba])
remote_mem_types = set(self.dst_region_mem_types[engine_id])
if len(remote_mem_types) != 1:
raise NotImplementedError(
"NIXL pull does not yet support a mixed-memory producer"
)
self.dst_xfer_side_handles[engine_id][remote_tp_rank] = (
self.nixl_wrapper.prep_xfer_dlist(
remote_agent_name, blocks_data, mem_type=remote_mem_types.pop()
)
)
return remote_agent_name
def _validate_remote_agent_handshake(
self,
nixl_agent_meta: NixlAgentMetadata,
remote_tp_size: int,
remote_dcp_size: int = 1,
):
"""Validate the remote agent handshake metadata ensuring the
invariants hold true.
"""
remote_engine_id = nixl_agent_meta.engine_id
assert self.transfer_topo is not None
remote_info = self.transfer_topo.get_engine_info(remote_engine_id)
assert remote_info.remote_tp_size == remote_tp_size
assert remote_info.remote_dcp_size == remote_dcp_size
# DCP sizes must divide one another; this is what keeps the
# read-slicing math in pull_worker a closed form.
assert (
self.dcp_size % remote_dcp_size == 0 or remote_dcp_size % self.dcp_size == 0
), (
f"DCP sizes must divide one another: local={self.dcp_size}, "
f"remote={remote_dcp_size} (engine {remote_engine_id})."
)
if self._has_mamba and self.dcp_size != remote_dcp_size:
raise RuntimeError(
"Hybrid MLA+Mamba NIXL transfers require matching DCP sizes, "
f"got local={self.dcp_size}, remote={remote_dcp_size}."
)
tp_ratio = self.transfer_topo.tp_ratio(remote_tp_size)
block_size_ratio = self.transfer_topo.block_size_ratio(
nixl_agent_meta.block_size
)
# num_kv_heads > tp_size with P_TP > D_TP not supported for non-mamba.
# Mamba models can have replicated FA KV with tp_ratio < 0.
# MLA models do not need to handle kv replication.
if not self.use_mla and not self._has_mamba:
assert not (
tp_ratio < 0 and self.transfer_topo.is_kv_replicated(remote_engine_id)
)
remote_physical_per_logical = (
nixl_agent_meta.physical_blocks_per_logical_kv_block
)
if (
(self.dcp_size > 1 or remote_dcp_size > 1)
and self.region_group_ids
and self.region_group_ids != self.dst_region_group_ids[remote_engine_id]
and (
self._physical_blocks_per_logical_kv_block != 1
or remote_physical_per_logical != 1
)
):
raise NotImplementedError(
"DCP region pulls require matching logical and physical block sizes"
)
if (
self._has_mamba
and remote_physical_per_logical
!= self._physical_blocks_per_logical_kv_block
and self.vllm_config.cache_config.enable_prefix_caching
):
raise RuntimeError(
"Prefix caching with heterogeneous physical_blocks_per_logical "
"is not supported for Mamba hybrid models. "
f"Local: {self._physical_blocks_per_logical_kv_block}, "
f"Remote: {remote_physical_per_logical}. "
"Disable prefix caching with --disable-prefix-caching."
)
if block_size_ratio != 1:
# Heterogeneous block sizes transfer at remote-block granularity;
# the untransferred tail of the last local attention block is
# zeroed in the receive post-process, and mamba state pages
# transfer 1:1 (never sub-split).
assert not self.use_host_buffer, (
"Heterogeneous block sizes are not supported with host buffer"
)
kv_cache_layout = (
self.kv_cache_layout
if not self.use_host_buffer
else self.host_buffer_kv_cache_layout
)
if not self.use_mla and nixl_agent_meta.kv_cache_layout != kv_cache_layout:
if (
self.kv_transfer_config.enable_permute_local_kv
and nixl_agent_meta.kv_cache_layout == "LBHNC"
and kv_cache_layout == "LBNHC"
):
logger.info(
"Remote is LBHNC and local is LBNHC, enabled additional permute "
"on local device KV."
)
assert not self._is_hma_required, (
"HMA does not support block size post processing"
)
self.enable_permute_local_kv = True
else:
raise RuntimeError(
"Heterogeneous TP expects same kv_cache_layout (local "
f"{kv_cache_layout}, remote {nixl_agent_meta.kv_cache_layout}). "
"Or enable experimental feature to receive LBHNC into a local "
"LBNHC cache by setting 'enable_permute_local_kv'=True "
"in --kv-transfer-config."
)
# if remote_agent used attn is not same as local,
# hint heterogenuous attn post process
if (
nixl_agent_meta.attn_backend_name != self.backend_name
and self.backend_name in ["CPU_ATTN"]
):
if self._is_hma_required:
raise RuntimeError(
"heterogeneous attn post process is not supported with HMA"
)
logger.info(
"[Experimental] CPU_ATTN backend is used, "
"hint heterogeneous attn post process"
)
self.enable_heterogeneous_attn_post_process = True
# Heterogeneous TP requires head-splitting, which only works with
# block-contiguous layouts (e.g. LBHNC). MLA and replicated-KV cases
# don't split on heads. Mamba doesn't support heterogeneous TP.
if (
abs(tp_ratio) != 1
and not self.use_mla
and not self.transfer_topo.is_kv_replicated(remote_engine_id)
and not KVCacheLayout[kv_cache_layout].is_block_contiguous
and not self.enable_permute_local_kv
):
raise RuntimeError(
"Heterogeneous TP head-dimension splitting requires contiguous heads. "
"Use a block-contiguous layout (e.g. LBHNC) on the prefill side."
)
# NVFP4 pages store all heads' data before all heads' scales, so a
# head's bytes are not a contiguous slice of the page.
if (
not self.use_mla
and self.transfer_topo.local_physical_heads
!= max(1, self.transfer_topo.total_num_kv_heads // remote_tp_size)
and any(
isinstance(spec, AttentionSpec) and spec.kv_quant_mode.is_nvfp4
for spec in self._layer_specs.values()
)
):
raise RuntimeError(
"NIXL KV-head re-sharding is not supported for NVFP4 KV cache "
f"(local TP {self.transfer_topo.tp_size}, remote TP "
f"{remote_tp_size}). Use matching TP sizes on prefill and decode."
)
# Per-region block_len validation enforcing the P/D invariant.
# REPLICATE regions (MLA, or a whole-model MLA / replicated-KV transfer)
# only allow the number of blocks to differ; SPLIT regions scale with
# the per-rank KV head ratio rather than the raw tp_ratio, because GQA
# replication caps per-rank heads at 1 when tp > total_kv_heads
# (issue #45330). Mamba uses the ssm_sizes counterpart, so skip here.
local_regions = self._transfer_layer_region_indices or range(
len(self.block_len_per_layer)
)
if self._has_mamba and self.use_mla:
# Hybrid MLA+SSM (e.g. KimiLinear's KDA+MLA): regions are
# kernel-granularity views of the mamba-unified page. The MLA
# per-token page is TP-independent, so the block lengths must
# match up to the kernel block size ratio even under
# heterogeneous TP (remote kernel blocks may be smaller).
# SSM geometry is validated via ssm_sizes/conv offsets instead.
assert self.block_len_per_layer == [
block_len * block_size_ratio for block_len in nixl_agent_meta.block_lens
], (
"Hybrid MLA kernel-granularity block lengths must match "
f"between P and D (block_size_ratio={block_size_ratio}): "
f"local={self.block_len_per_layer}, "
f"remote={nixl_agent_meta.block_lens}."
)
elif not self._has_mamba:
assert len(local_regions) == len(nixl_agent_meta.block_lens), (
"Number of KV layers must match between prefill and decode"
)
model_replicated = self.use_mla or self.transfer_topo.is_kv_replicated(
remote_engine_id
)
total_kv_heads = self.transfer_topo.total_num_kv_heads
local_heads = self.transfer_topo.local_physical_heads
remote_heads = max(1, total_kv_heads // remote_tp_size)
for i, region_idx in enumerate(local_regions):
local_len = self.block_len_per_layer[region_idx]
replicated = model_replicated or self._is_region_replicated(region_idx)
remote_len = nixl_agent_meta.block_lens[i]
if replicated:
assert local_len // block_size_ratio == remote_len, (
"KV cache sizes must match between P and D when "
f"replicated (region {i}: local={local_len}, "
f"remote={remote_len}, bsr={block_size_ratio})."
)
elif tp_ratio > 0:
assert (
remote_len
== (local_len * remote_heads // local_heads) // block_size_ratio
), (
f"SPLIT region {i}: remote P KV block_len {remote_len} "
f"must equal local {local_len} * remote_heads "
f"{remote_heads} // local_heads {local_heads} "
f"// block_size_ratio {block_size_ratio}."
)
else:
assert block_size_ratio == 1, (
"Different local/remote block sizes are not supported "
"when P TP > D TP."
)
assert remote_len == local_len * remote_heads // local_heads, (
f"SPLIT region {i}: remote P KV block_len {remote_len} "
f"must equal local {local_len} * remote_heads "
f"{remote_heads} // local_heads {local_heads}."
)
# TP workers that handhshake with same remote have same #blocks.
assert self.dst_num_blocks[remote_engine_id] == nixl_agent_meta.num_blocks
assert len(nixl_agent_meta.kv_caches_base_addr) == len(local_regions)
num_remote_regions = len(nixl_agent_meta.kv_caches_base_addr)
if nixl_agent_meta.region_num_blocks is not None:
assert len(nixl_agent_meta.region_num_blocks) == num_remote_regions
assert len(nixl_agent_meta.block_strides) == num_remote_regions
assert all(
stride >= block_len
for stride, block_len in zip(
nixl_agent_meta.block_strides,
nixl_agent_meta.block_lens,
strict=True,
)
)
if nixl_agent_meta.region_group_ids is not None:
assert len(nixl_agent_meta.region_group_ids) == num_remote_regions
if nixl_agent_meta.region_names is not None:
expected_names = (
list(self._transfer_layer_names)
if self._transfer_layer_names
else self.region_names
)
assert nixl_agent_meta.region_names == expected_names, (
"NIXL transfer regions must name the same model caches in the "
"same order"
)
if nixl_agent_meta.region_mem_types is not None:
assert len(nixl_agent_meta.region_mem_types) == num_remote_regions
def sync_recved_kv_to_device(self, req_id: str, meta: ReqMeta):
"""Copy recved kv from host buffer to device."""
assert self.use_host_buffer
assert self.copy_blocks is not None
local_block_ids = meta.local_physical_block_ids
# Every KV cache group allocates from the same BlockPool, so block ids
# are unique across groups and the per-group copies can be issued as one.
block_ids = [block_id for group in local_block_ids for block_id in group]
# The h2d block copy below is intentionally synchronous.
with gpu_sync_allowed():
self.copy_blocks(
self.host_xfer_buffers,
self.device_kv_caches,
block_ids,
block_ids,
"h2d",
)
if logger.isEnabledFor(logging.DEBUG):
logger.debug(
"synced recved kv of request[%s] to device kv buffer,"
"local_block_ids: %s. ",
req_id,
",".join(map(str, local_block_ids)),
)
def save_kv_to_host(self, metadata: NixlConnectorMetadata):
"""Copy kv from device to host buffer."""
assert self.use_host_buffer
assert self.copy_blocks is not None
# The d2h block copies below are intentionally synchronous.
with gpu_sync_allowed():
for req_id, meta in metadata.reqs_to_save.items():
meta.local_physical_block_ids = self._logical_to_kernel_block_ids(
meta.local_block_ids, self._physical_blocks_per_logical_kv_block
)
if logger.isEnabledFor(logging.DEBUG):
logger.debug(
"save_load_kv for request[%s] to host xfer buffer."
"local_block_ids: %s. ",
req_id,
",".join(map(str, meta.local_physical_block_ids)),
)
# blocking
block_ids = [
block_id
for group in meta.local_physical_block_ids
for block_id in group
]
self.copy_blocks(
self.device_kv_caches,
self.host_xfer_buffers,
block_ids,
block_ids,
"d2h",
)
def post_process_device_kv_on_receive(
self,
block_size_ratio: int,
block_ids_list: list[tuple[int, list[int], int]],
convert: bool = True,
):
"""Post process device kv cache after receiving from remote.
3 types of conversion supported (``convert``):
* kv_cache_postprocess_layout => convert from HND to NHD
* kv_cache_postprocess_blksize => convert from small block size
to large block size
* kv_cache_postprocess_blksize_and_layout => convert from small
block size to large block size and convert from HND to NHD
The transfer only covers ``covered_sub_blocks`` remote-sized
sub-blocks of each request's local attention blocks; the rest was
clipped, either by remote-block pairing (block-size ratio) or by the
hetero-ppl front trim in ``_apply_prefix_caching``. Those blocks were
excluded from the scheduler's alloc-time KV zeroing (which would race
the RDMA write), so everything past the covered range is zeroed here.
Stale bytes would otherwise surface as garbage or NaNs once decode
grows into the untransferred tail.
"""
if len(self.device_kv_caches) == 0:
return
assert block_size_ratio >= 1, "Only nP < nD supported currently."
assert self.transfer_topo is not None
if not convert:
logger.debug(
"Post-processing device kv cache on receive by zeroing "
"untransferred blocks."
)
elif self.enable_permute_local_kv and block_size_ratio > 1:
logger.debug(
"Post-processing device kv cache on receive by converting "
"block_size with %sx bigger and permuting layout from HND"
" to NHD.",
block_size_ratio,
)
elif self.enable_permute_local_kv:
logger.debug(
"Post-processing device kv cache on receive by permuting layout"
"from HND to NHD."
)
else:
logger.debug(
"Post-processing device kv cache on receive by converting "
"block_size with %sx bigger.",
block_size_ratio,
)
for group, block_ids, covered_sub_blocks in block_ids_list:
attn_caches = [
self.device_kv_caches[name]
for name in self.kv_cache_config.transfer_groups[group].layer_names
]
# Blocks the transfer didn't write: the token tail of the last
# partially covered block, then everything beyond it.
covered_blocks, sub_blocks_in_last = divmod(
covered_sub_blocks, block_size_ratio
)
first_stale = covered_blocks + (1 if sub_blocks_in_last else 0)
has_stale = first_stale < len(block_ids)
indices = None
if convert or has_stale:
indices = async_tensor_h2d(block_ids, attn_caches[0].device, torch.long)
if convert:
for cache in attn_caches:
if self.enable_permute_local_kv:
kv_postprocess_blksize_and_layout_on_receive(
cache, indices, block_size_ratio
)
elif KVCacheLayout[self.kv_cache_layout].is_block_contiguous:
# Token-major blocks get the sub-blocks in token order.
kv_postprocess_blksize_on_receive(
cache, indices, block_size_ratio
)
if sub_blocks_in_last:
last_block_id = block_ids[covered_blocks]
for cache in attn_caches:
# Attention caches are [B, H, N, C]: tokens on dim 2.
sub_block_tokens = cache.shape[2] // block_size_ratio
zero_from = sub_blocks_in_last * sub_block_tokens
cache[last_block_id, :, zero_from:].zero_()
if has_stale:
assert indices is not None
stale_ids = indices[first_stale:]
for cache in attn_caches:
cache.index_fill_(0, stale_ids, 0)
def post_process_device_kv_on_receive_heterogeneous_attn(
self, block_ids: list[int]
):
"""Post process device kv cache after receiving from remote
for heterogeneous attention.
"""
assert self.enable_heterogeneous_attn_post_process
indices = torch.tensor(block_ids, device=self.device_type, dtype=torch.long)
for _, cache_or_caches in self.device_kv_caches.items():
current_platform.pack_kv_cache(
kv_cache=cache_or_caches,
indices=indices,
)
def _leaves_dram_xfer_to_rank0(self) -> bool:
"""Whether this rank leaves the DRAM regions' NIXL work to TP rank 0.
A HiSparse host pool shared by the local TP ranks (one mmap per DP
rank) is registered with NIXL and filled by TP rank 0 alone: every
rank would otherwise register the same pages on every NIC and READ the
same host KV into them. Private pools are unaffected.
"""
if not (
self._mixed_mem_types and self.kv_cache_config.hisparse_shared_host_pool
):
return False
skip = self.tp_rank != 0
logger.info(
"HiSparse host pool: NIXL DRAM registration and reads on TP rank 0 "
"only (tp_rank=%d, skip=%s)",
self.tp_rank,
skip,
)
return skip
def _zero_region_blocks(self, block_ids: BlockIds) -> None:
"""Clear clipped physical pages in their owning region only."""
if not any(block_ids):
return
bases = self.kv_caches_base_addr[self.engine_id][self.tp_rank]
for region, blocks in enumerate(block_ids):
if not blocks:
continue
if self._skip_dram_xfer and self.region_mem_types[region] == "DRAM":
# The shared pool is rank 0's to write.
continue
cache = self.device_kv_caches[self.region_names[region]]
storage = cache.untyped_storage()
pages = torch.empty(0, dtype=torch.uint8, device=cache.device).set_(
storage,
bases[region] - storage.data_ptr(),
(self.region_num_blocks[region], self.block_len_per_layer[region]),
(self.block_stride_per_layer[region], 1),
)
for block in blocks:
pages[block].zero_()
def get_transfer_results(self) -> KVConnectorTransferResults:
"""Get transfers that completed on this specific worker.
The scheduler process (via the MultiprocExecutor) will use this output
to track which workers are done.
"""
assert self.transfer_topo is not None
done_sending = self._get_new_notifs()
done_recving, newly_failed = self._pop_done_transfers(self._recving_transfers)
if newly_failed:
self._recv_failures.update(newly_failed)
done_sending.update(self._replicated_pcp_done_sending)
self._replicated_pcp_done_sending.clear()
# Process receive failures reported by background threads. Each
# producer already recorded a specific metric (handshake or
# notification), so only the failure bookkeeping runs here.
while not self._failed_recv_reqs.empty():
try:
req_id = self._failed_recv_reqs.get_nowait()
except queue.Empty:
break
self._handle_failed_transfer(
req_id, None, self._recv_failures, record_failed_transfer=False
)
failed_recv_reqs: set[ReqId] = set()
if self._recv_failures:
# Keep failure bookkeeping independent of healthy request count.
for req_id in tuple(self._recv_failures):
if req_id not in self._recving_metadata:
self._recv_failures.remove(req_id)
elif req_id not in self._recving_transfers:
self._recv_failures.remove(req_id)
failed_recv_reqs.add(req_id)
done_recving.update(failed_recv_reqs)
if len(done_sending) > 0 or len(done_recving) > 0:
logger.debug(
"Rank %s, get_finished: %s requests done sending "
"and %s requests done recving (%s failed)",
self.tp_rank,
len(done_sending),
len(done_recving),
len(failed_recv_reqs),
)
block_ids_for_blocksize_post_process = defaultdict(list)
block_ids_for_heterogeneous_attn_post_process = list[list[int]]()
direct_device_recving = set[str]()
for req_id in done_recving:
# clean up metadata for completed requests
meta = self._recving_metadata.pop(req_id, None)
assert meta is not None, f"{req_id} not found in recving_metadata list"
# Skip KV sync and post-processing for failed requests
if req_id in failed_recv_reqs:
self._pending_recv_notifs.pop(req_id, None)
# TODO (NickLucche) handle failed transfer for HMA.
if not self._is_hma_required:
self._invalid_block_ids.put(set(meta.local_block_ids[0]))
logger.warning(
"Skipping KV post-processing for failed request %s",
req_id,
)
continue
self._send_pending_recv_notifs(req_id)
assert meta.remote is not None
if self.use_host_buffer:
self.sync_recved_kv_to_device(req_id, meta)
direct_device_recving.add(req_id)
if meta.region_blocks_to_zero is not None:
# P/D group positions differ. Use the actual region read plan,
# including any local allocation padding the read did not cover.
self._zero_region_blocks(meta.region_blocks_to_zero)
continue
# Post processing for heteroblocksize/layout, and for blocks the
# transfer clipped. The latter happens either at remote-block
# granularity (block_size_ratio > 1) or at kernel-block
# granularity, when equal kernel pages meet differing logical
# block sizes and _apply_prefix_caching front-trims to the
# minimum count (hybrid heterogeneous TP).
remote_info = self.transfer_topo.get_engine_info(meta.remote.engine_id)
block_size_ratio = self.transfer_topo.block_size_ratio(
remote_info.remote_block_size
)
hetero_ppl = (
remote_info.remote_physical_blocks_per_logical
!= self._physical_blocks_per_logical_kv_block
)
if block_size_ratio > 1 or self.enable_permute_local_kv or hetero_ppl:
for g, local_group in enumerate(meta.local_physical_block_ids):
if not local_group or _is_ssm_spec(self._group_spec_types[g]):
continue
# Number of remote-sized sub-blocks the transfer covered;
# everything past this was clipped and must be zeroed.
covered_sub_blocks = min(
len(local_group) * block_size_ratio,
len(meta.remote.block_ids[g]),
)
block_ids_for_blocksize_post_process[block_size_ratio].append(
(g, local_group, covered_sub_blocks)
)
# post processing for heterogeneous attention
if self.enable_heterogeneous_attn_post_process:
block_ids_for_heterogeneous_attn_post_process.append(
meta.local_physical_block_ids[0]
)
for (
block_size_ratio,
block_ids_list,
) in block_ids_for_blocksize_post_process.items():
# MLA never needs the block-size/layout conversion, but its
# clipped blocks still need zeroing.
convert = not self.use_mla and (
block_size_ratio > 1 or self.enable_permute_local_kv
)
self.post_process_device_kv_on_receive(
block_size_ratio, block_ids_list, convert
)
for block_ids in block_ids_for_heterogeneous_attn_post_process:
self.post_process_device_kv_on_receive_heterogeneous_attn(block_ids)
self._sync_device_after_direct_recv(direct_device_recving)
# Handle timeout to avoid stranding blocks on remote.
self._reap_expired_send_leases(done_sending)
self._cleanup_remote_engines()
return KVConnectorTransferResults(
finished_sending=done_sending,
finished_recving=done_recving,
failed_recving=failed_recv_reqs,
)
def get_finished(self) -> tuple[set[str], set[str]]:
"""Compatibility wrapper for the legacy completion API."""
results = self.get_transfer_results()
return results.finished_sending, results.finished_recving
def _sync_device_after_direct_recv(self, done_recving: set[str]) -> None:
"""Make direct NIXL writes visible before model execution."""
requires_sync = current_platform.is_rocm() and self._has_mamba
if self.use_host_buffer or not done_recving or not requires_sync:
return
torch.accelerator.synchronize()
def _get_new_notifs(self) -> set[str]:
"""Get req_ids which got a remote xfer notification.
Subclasses must implement this to handle mode-specific notifications.
"""
raise NotImplementedError
def _reap_expired_send_leases(self, done_sending: set[str]) -> None:
"""Reclaim expired send-side KV leases into ``done_sending``.
``_reqs_to_send`` is not ordered by expiry: heartbeats update the
deadline in place, and mixed TTLs share the map, so a live head
entry can sit in front of already-expired ones. Scan every entry
rather than stopping at the first still-live request.
"""
if not self._reqs_to_send:
return
now = time.perf_counter()
expired = [
req_id for req_id, expires in self._reqs_to_send.items() if now >= expires
]
for req_id in expired:
count = self.consumer_notification_counts_by_req.pop(req_id, 0)
self.expected_consumer_notifications_by_req.pop(req_id, None)
self.xfer_stats.record_kv_expired_req()
logger.warning(
"Releasing expired KV blocks for request %s which were "
"retrieved by %d remote worker(s) before lease expired.",
req_id,
count,
)
self._reqs_to_process.remove(req_id)
del self._reqs_to_send[req_id]
done_sending.add(req_id)
def _handle_heartbeat(self, payload: str) -> None:
"""Extend leases for requests referenced in a heartbeat.
Args:
payload: comma-separated P-side request IDs, e.g.
"req_abc,req_def".
"""
new_expiry = time.perf_counter() + self._lease_extension
for req_id in payload.split(","):
if req_id in self._reqs_to_send:
old = self._reqs_to_send[req_id]
self._reqs_to_send[req_id] = max(old, new_expiry)
logger.debug(
"Heartbeat extended lease for request %s "
"by %ds (old_expiry=%.1f, new_expiry=%.1f)",
req_id,
self._lease_extension,
old,
new_expiry,
)
def _pop_done_transfers(
self, transfers: dict[str, list[int]]
) -> tuple[set[str], set[str]]:
"""Poll transfers, retaining handles until they can be released.
Args:
transfers: Outstanding handles by request, updated in place.
Returns:
Requests with no outstanding handles and requests with observed
failures. Failed requests may still have outstanding handles.
"""
done_req_ids: set[str] = set()
failed_req_ids: set[str] = set()
for req_id, handles in list(transfers.items()):
in_progress = []
for handle in handles:
try:
xfer_state = self.nixl_wrapper.check_xfer_state(handle)
if xfer_state == "DONE":
res = self.nixl_wrapper.get_xfer_telemetry(handle)
self.xfer_stats.record_transfer(res)
self.nixl_wrapper.release_xfer_handle(handle)
elif xfer_state == "PROC":
in_progress.append(handle)
else:
self._log_failure(
failure_type="transfer_failed",
req_id=req_id,
xfer_state=xfer_state,
)
if not self._handle_failed_transfer(
req_id, handle, failed_req_ids
):
in_progress.append(handle)
except Exception as e:
self._log_failure(
failure_type="transfer_exception",
req_id=req_id,
error=e,
)
if not self._handle_failed_transfer(req_id, handle, failed_req_ids):
in_progress.append(handle)
if not in_progress:
done_req_ids.add(req_id)
del transfers[req_id]
else:
transfers[req_id] = in_progress
return done_req_ids, failed_req_ids
def _try_release_xfer_handle(self, req_id: str, handle: int) -> bool:
"""Release a handle, returning False if the caller must retain it."""
try:
self.nixl_wrapper.release_xfer_handle(handle)
except Exception as e:
# A status error does not guarantee that the backend stopped DMA.
self._log_failure(
failure_type="transfer_release_failed",
msg="Retaining handle and blocks until release succeeds",
req_id=req_id,
error=e,
)
return False
return True
def _handle_failed_transfer(
self,
req_id: str,
handle: int | None,
failed_req_ids: set[str] | None = None,
record_failed_transfer: bool = True,
) -> bool:
"""Record a failure and release its handle, returning False to retain it.
Args:
req_id: The request ID.
handle: The transfer handle.
failed_req_ids: Requests observed as failed; each stays listed
until it has no outstanding handles.
record_failed_transfer: Whether to count the failure toward the
transport-failure metric. Callers that already recorded a more
specific metric (eg KV expiry, reported separately from
transport failures) pass False.
"""
if record_failed_transfer:
self.xfer_stats.record_failed_transfer()
if (
self._TRANSFER_MODE == "pull"
and (meta := self._recving_metadata.get(req_id)) is not None
and meta.remote is not None
):
self._failed_remote_engines.add(meta.remote.engine_id)
if failed_req_ids is not None:
failed_req_ids.add(req_id)
return handle is None or self._try_release_xfer_handle(req_id, handle)
def _send_pending_recv_notifs(self, req_id: str) -> None:
"""Send notifications deferred by split DRAM/VRAM reads."""
for agent_name, notif_id in self._pending_recv_notifs.pop(req_id, []):
try:
self.nixl_wrapper.send_notif(agent_name, notif_msg=notif_id)
except Exception as e:
self._log_failure(
failure_type="notification_failed",
msg="P worker blocks will be freed after timeout. "
"This may indicate network issues.",
req_id=req_id,
error=e,
remote_agent_name=agent_name,
)
self.xfer_stats.record_failed_notification()
def _send_heartbeats(self, metadata: NixlConnectorMetadata) -> None:
"""Send heartbeats to remote engines, extending the lease on KV blocks."""
for engine_id, hb_info in metadata.heartbeat_by_engine.items():
# Proactive handshake (this request may still be in waiting queue) so
# the **next** heartbeat for this remote can go through.
if (
self._ensure_handshake(
engine_id,
hb_info.host,
hb_info.port,
hb_info.tp_size,
hb_info.dcp_size,
hb_info.pp_size,
self._hb_handshake_notif_only and hb_info.pp_size > 1,
)
is not None
):
continue # handshake is still pending
# Refresh the TTL: a heartbeat means the router still paired this P-D,
# so requests are still waiting on this engine
self._engine_last_active[engine_id] = time.perf_counter()
# Build the heartbeat message: "HB:req1,req2,..."
hb_msg = ("HB:" + ",".join(hb_info.req_ids)).encode()
for agent_name in self._remote_agents[engine_id].values():
try:
self.nixl_wrapper.send_notif(agent_name, notif_msg=hb_msg)
except Exception:
logger.debug(
"Failed to send heartbeat to engine %s",
engine_id,
exc_info=True,
)
def get_mapped_blocks(
self, block_ids: np.ndarray, block_size_ratio: int
) -> np.ndarray:
"""Calculates the new set of block IDs by mapping every element
in the (potentially sparse) input array.
Example: block_ids=[0, 2], block_size_ratio=2
get_mapped_blocks 0 1 [2 3] 4 5
# remote is |h0-b0|h1-b0||h0-b1|h1-b1||h0-b1|h1-b1||
# local is |h0-b0......||h1-b0......||h2-b0........
local_block_ids 0 [1] 2
"""
if block_ids.size == 0:
return np.array([], dtype=np.int64)
start_ids = block_ids * block_size_ratio
offsets = np.arange(block_size_ratio)
mapped_2d = start_ids[:, None] + offsets[None, :]
return mapped_2d.flatten().astype(np.int64)
def _map_block_ids_for_block_size_ratio(
self,
local_block_ids: BlockIds,
remote_block_ids: BlockIds,
block_size_ratio: int,
) -> tuple[BlockIds, BlockIds]:
"""Map attention-group block ids to remote-block granularity.
Each local attention block is split into ``block_size_ratio``
sub-blocks paired 1:1 with remote blocks. Sub-blocks beyond the
remote list — the untransferred tail of the last local block — are
clipped here and zeroed in the receive post-process. Mamba state
blocks are indivisible and transfer 1:1, unexpanded.
ex: remote (prefill) block ids with block_size 4:
[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
Local (decode) block ids with block_size 16: [1, 2, 3] expand to
[4, 5, ..., 15], then clip to the first 10 to pair 1:1 with remote.
"""
mapped_local: list[list[int]] = []
mapped_remote: list[list[int]] = []
for i, remote_group in enumerate(remote_block_ids):
local_group = local_block_ids[i] if local_block_ids else []
if _is_ssm_spec(self._group_spec_types[i]):
mapped_local.append(list(local_group))
mapped_remote.append(list(remote_group))
continue
mapped = self.get_mapped_blocks(
np.asarray(local_group), block_size_ratio
).tolist()
if len(mapped) > len(remote_group):
mapped = mapped[: len(remote_group)]
mapped_local.append(mapped)
mapped_remote.append(list(remote_group))
if not any(mapped_local):
# Full prefix cache hit is indicated with an empty list.
return [], mapped_remote
return mapped_local, mapped_remote
def _logical_to_kernel_block_ids(self, block_ids: BlockIds, ratio: int) -> BlockIds:
"""Convert block ids to kernel physical block ids.
This is required when the logical block size (the one set by the user)
does not match the one required by the attn backend.
`ratio` is the number of physical blocks per logical block.
We always receive logical blocks from the engine, so we expand them here eg:
logical block ids: [(SW-clipped) [1], (FA) [2, 3]], ratio=2
physical block ids: [(SW-clipped) [2, 3], (FA) [4, 5, 6, 7]]
"""
if ratio == 1:
# Noop when physical and logical block sizes are the same
return block_ids
block_arange = np.arange(0, ratio).reshape(1, -1)
# Mamba blocks have no logical<>physical discrepancy (block-size=1)
physical_block_ids = []
for i, group in enumerate(block_ids):
spec = self.kv_cache_config.transfer_groups[i].kv_cache_spec
if _is_ssm_spec(get_representative_spec_type(spec)):
physical_block_ids.append(group)
else:
physical_block_ids.append(
(np.array(group)[:, None] * ratio + block_arange)
.reshape(-1)
.tolist()
)
return physical_block_ids
def _apply_dcp_prefix_caching(
self,
local_ids: list[int],
remote_ids: list[int],
remote_rank: int,
local_dcp_size: int,
local_dcp_rank: int,
remote_dcp_size: int,
local_num_computed_blocks: int,
num_remote_blocks: int | None = None,
) -> tuple[list[int], list[int]]:
"""Match local and remote block IDs by global DCP position.
``local_ids`` excludes blocks already satisfied by the local prefix
cache, while ``remote_ids`` contains the full transferable remote list.
``num_remote_blocks``, when provided, is the global block count before
DCP partitioning and excludes allocation padding.
Example:
local DCP size 2, rank 0 owns *global positions*=[0, 2, 4, 6]
cached_blocks=[0,2,4]
to transfer: [6] (the rest is cached)
remote DCP size 4
remote rank 0 owns [0, 4] -> no more match, skip the read
remote rank 2 owns [2, 6] -> match position 6 and read
An empty result (rank 0 above) takes the notification-only path, allowing the
remote to release its blocks without performing a transfer.
"""
local_size, remote_size = local_dcp_size, remote_dcp_size
if num_remote_blocks is not None:
num_local = cdiv(num_remote_blocks - local_dcp_rank, local_size)
local_ids = local_ids[: max(0, num_local - local_num_computed_blocks)]
if local_size == remote_size:
local_slice = local_ids
remote_slice = remote_ids[local_num_computed_blocks:]
elif local_size < remote_size:
k = remote_size // local_size
p = (remote_rank - local_dcp_rank) // local_size
start_local = (p - local_num_computed_blocks) % k
start_remote = (local_num_computed_blocks + start_local - p) // k
local_slice = local_ids[start_local::k]
remote_slice = remote_ids[start_remote:]
else:
k = local_size // remote_size
remote_dcp_rank = remote_rank % remote_size
c = (local_dcp_rank - remote_dcp_rank) // remote_size
start_remote = c + local_num_computed_blocks * k
local_slice = local_ids
remote_slice = remote_ids[start_remote::k]
matched_blocks = min(len(local_slice), len(remote_slice))
if num_remote_blocks is not None and matched_blocks < len(local_slice):
raise ValueError("Remote KV pages do not cover the requested range")
return local_slice[:matched_blocks], remote_slice[:matched_blocks]
@staticmethod
def _block_ids_by_region(
block_ids: BlockIds, region_group_ids: list[int]
) -> BlockIds:
"""Expand group block IDs into the corresponding region order."""
shared_block_ids = list(itertools.chain.from_iterable(block_ids))
block_ids_by_region = []
for group_id in region_group_ids:
if group_id == _SHARED_REGION_GROUP_ID:
block_ids_by_region.append(shared_block_ids.copy())
continue
block_ids_by_region.append(list(block_ids[group_id]))
return block_ids_by_region
def _apply_prefix_caching_by_region(
self,
decode_block_ids: BlockIds,
prefill_block_ids: BlockIds,
*,
num_computed_blocks: list[int] | None = None,
num_remote_blocks: int | None = None,
remote_rank: int = 0,
remote_dcp_size: int = 1,
) -> tuple[BlockIds, BlockIds]:
"""Pair an uncached decode suffix with the same prefill regions."""
assert len(decode_block_ids) == len(prefill_block_ids)
if not any(decode_block_ids):
empty_regions: list[list[int]] = [[] for _ in decode_block_ids]
return empty_regions, empty_regions.copy()
if num_computed_blocks is not None:
assert num_remote_blocks is not None
matched_decode, matched_prefill = [], []
for decode_region, prefill_region, start in zip(
decode_block_ids, prefill_block_ids, num_computed_blocks, strict=True
):
local, remote = self._apply_dcp_prefix_caching(
decode_region,
prefill_region,
remote_rank,
self.dcp_size,
self.dcp_rank,
remote_dcp_size,
start,
num_remote_blocks=num_remote_blocks,
)
matched_decode.append(local)
matched_prefill.append(remote)
return matched_decode, matched_prefill
trimmed_prefill: list[list[int]] = []
for decode_region, prefill_region in zip(
decode_block_ids, prefill_block_ids, strict=True
):
if len(decode_region) > len(prefill_region):
raise ValueError(
"Decode allocated more KV pages than the prefill worker supplied"
)
trimmed_prefill.append(
list(prefill_region[-len(decode_region) :]) if decode_region else []
)
return decode_block_ids, trimmed_prefill
def _apply_prefix_caching(
self,
decode_block_ids: BlockIds,
prefill_block_ids: BlockIds,
decode_physical_per_logical: int,
prefill_physical_per_logical: int,
) -> tuple[BlockIds, BlockIds]:
"""Trim block ID lists so only the uncomputed suffix is transferred.
Inputs are *kernel* (physical) block IDs, already expanded from logical IDs
with each side's physical-per-logical ratio. Both pull and push call this after
that expansion so the trim happens at kernel granularity.
The prefix-cache hit is always on the decode (D) side, so ``decode``
holds only its uncomputed blocks while prefill (P) holds the full
sequence. This is mode-independent: pull passes its own (D) blocks as
``decode``, push passes the remote D registration as ``decode``.
For non-Mamba models: end-trim ``prefill`` to match ``decode`` count, so
already-cached prefix blocks are skipped in the transfer.
For Mamba hybrid: SSM groups pair state slots by position and FA
groups end-trim to the uncomputed suffix when physical-per-logical
matches.
"""
# Partial prefix cache hit: just transfer uncomputed blocks.
# Skip mamba groups — their blocks represent full state (conv+ssm),
# not per-token data, so trimming would corrupt the transfer.
prefill_block_ids = list(prefill_block_ids)
if not self._has_mamba:
for i, prefill_group in enumerate(prefill_block_ids):
num_decode_blocks = len(decode_block_ids[i])
assert num_decode_blocks <= len(prefill_group)
if num_decode_blocks < len(prefill_group):
prefill_block_ids[i] = prefill_group[-num_decode_blocks:]
else:
# (NOTE: ZhanqiuHu) HeteroTP can cause different kernel block counts
# due to logical block rounding.
# Example: 640 prompt tokens, kernel_block_size=64
# prefill physical_per_logical=10, decode physical_per_logical=6
# prefill logical ids from kv_transfer_params = [0]
# decode logical ids allocated = [0, 1]
# prefill kernel blocks: [0..9] (1*10=10)
# decode kernel blocks: [0..11] (2*6=12)
# actual data blocks = ceil(640/64) = 10, trim both to 10
# Vice versa (prefill physical_per_logical=6, decode=10):
# prefill logical ids = [0, 1], decode logical ids = [0]
# prefill kernel blocks: [0..11] (2*6=12)
# decode kernel blocks: [0..9] (1*10=10)
# actual data blocks = ceil(640/64) = 10, trim both to 10
decode_block_ids = list(decode_block_ids)
for i, prefill_group in enumerate(prefill_block_ids):
num_decode_blocks = len(decode_block_ids[i])
num_prefill_blocks = len(prefill_group)
if _is_ssm_spec(self._group_spec_types[i]):
if num_decode_blocks == num_prefill_blocks:
continue
# Only state-bearing slots reach here, normally just one
# (see get_exchange_clipped_blocks). A longer prefill list
# carries earlier positions the decode side already has
# (prefix hit) -> take its tail; a longer decode list holds
# the position D recomputes itself, which gets no prefill
# state.
assert num_decode_blocks - num_prefill_blocks <= 1, (
f"Group {i}: unpairable SSM state slots, "
f"decode={num_decode_blocks} prefill={num_prefill_blocks}"
)
num_blocks = min(num_decode_blocks, num_prefill_blocks)
if num_decode_blocks < num_prefill_blocks:
prefill_block_ids[i] = prefill_group[-num_blocks:]
else:
decode_block_ids[i] = decode_block_ids[i][:num_blocks]
elif (
decode_physical_per_logical == prefill_physical_per_logical
and num_decode_blocks < num_prefill_blocks
):
# Partial prefix cache hit for FA group.
prefill_block_ids[i] = prefill_group[-num_decode_blocks:]
else:
# TODO Handle prefix caching with different block_sizes
# Allocation rounding legitimately leaves up to
# ppl - 1 trailing dead kernel blocks per side (plus one
# extra decode block for the recomputed final token), so
# the counts may differ by up to the sum of the two
# ratios; anything larger indicates mismatched lists.
max_padding = (
decode_physical_per_logical + prefill_physical_per_logical
)
assert abs(num_decode_blocks - num_prefill_blocks) <= max_padding, (
f"Group {i}: |{num_decode_blocks} - "
f"{num_prefill_blocks}| > {max_padding}"
)
num_blocks = min(num_decode_blocks, num_prefill_blocks)
decode_block_ids[i] = decode_block_ids[i][:num_blocks]
prefill_block_ids[i] = prefill_group[:num_blocks]
return decode_block_ids, prefill_block_ids
def get_kv_connector_stats(self) -> KVConnectorStats | None:
"""Get the KV transfer stats for the connector."""
# Clear stats for next iteration
if not self.xfer_stats.is_empty():
return self.xfer_stats.clone_and_reset()
return None
def get_block_ids_with_load_errors(self) -> set[int]:
"""Return and clear the set of block IDs that failed to load.
This is called by the scheduler to identify blocks that need
to be retried after a NIXL transfer failure.
"""
# Drain the queue (thread-safe, no lock needed).
result: set[int] = set()
while not self._invalid_block_ids.empty():
try:
result.update(self._invalid_block_ids.get_nowait())
except queue.Empty:
break
return result
def _evict_stale_engines(self) -> None:
"""Scan for and evict remote engines that have exceeded their TTL.
Called from the main thread in when a new remote engine appears.
We can only go OOM as we discover and register a new remote, therefore we make
sure we clean up stale engine data structures before then. This invariant
prevents us from using background threads, though memory usage is not guaranteed
to be "optimal" until a new handshake is performed.
Engines with active transfers, heartbeats or pending handshakes cannot
be stale:
- Reads stamp _engine_last_active when they are issued, and engines a
transfer is still reading from are held back explicitly, since the
stamp is not refreshed while the read runs.
- Heartbeats stamp it too, so an engine stays registered while this
instance holds requests waiting on it.
- Pending handshakes don't have an _engine_last_active entry yet
"""
# NOTE (NickLucche): This does NOT currently prevent OOMing if a huge number
# of remote engines is registered all at once (adding a background cleanup
# thread wouldnt help either).
# If that scenario is plausible, we can follow up with an LRU eviction policy.
if self._engine_ttl <= 0:
return
now = time.perf_counter()
busy = self._engines_with_inflight_transfers()
for eid, last_active in list(self._engine_last_active.items()):
if now - last_active > self._engine_ttl and eid not in busy:
self._cleanup_remote_engine(eid)
def _engines_with_inflight_transfers(self) -> set[EngineId]:
"""Remote engines a transfer is still reading from.
The timestamp is stamped when a read is issued and not refreshed while
it runs, so a transfer that outlives the TTL leaves its engine looking
idle. A peer that has lost its NIC holds one indefinitely.
"""
return {
meta.remote.engine_id
for req_id in self._recving_transfers
if (meta := self._recving_metadata.get(req_id)) is not None
and meta.remote is not None
}
def _track_remote_engine_replacement(
self, engine_id: EngineId, host: str, port: int
) -> None:
"""Record the engine currently serving a pull peer address."""
# TODO: Also handle push mode, which handshakes in both directions.
if self._TRANSFER_MODE != "pull":
return
previous = self._engine_by_address.get((host, port))
self._engine_by_address[(host, port)] = engine_id
if previous is None or previous == engine_id:
return
logger.info(
"Remote engine at %s:%d changed from %s to %s.",
host,
port,
previous,
engine_id,
)
def _cleanup_remote_engines(self) -> None:
"""Release replaced or locally invalid pull peers after their reads drain."""
# TODO: Also handle push mode, which handshakes in both directions.
if self._TRANSFER_MODE != "pull":
return
# A background handshake may be loading native metadata with the GIL
# released. Defer both probing and cleanup until it has finished.
with self._handshake_lock:
if self._handshake_futures:
return
replaced = self._remote_agents.keys() - self._engine_by_address.values()
now = time.perf_counter()
for engine_id in self._failed_remote_engines - self._invalid_remote_engines:
agents = self._remote_agents.get(engine_id, {}).values()
try:
# No descriptors: query local existence, not peer health.
if not all(map(self.nixl_wrapper.check_remote_metadata, agents)):
self._mark_remote_engine_invalid(engine_id, now)
except Exception:
logger.warning(
"Could not check local NIXL metadata for engine %s.",
engine_id,
exc_info=True,
)
self._failed_remote_engines.clear()
# Invalid peers keep rejecting new reads until their backoff has elapsed.
invalid = {
engine_id
for engine_id in self._invalid_remote_engines
if self._invalid_engine_backoff.get(engine_id, (0.0, 0.0))[1] <= now
}
if not replaced and not invalid:
return
# _recving_metadata is only accessed from this (main) thread.
busy = {
meta.remote.engine_id
for meta in self._recving_metadata.values()
if meta.remote is not None
}
# Replaced peers can still serve queued requests. Invalid peers reject
# new reads, but must retain descriptors until outstanding handles end.
invalid -= self._engines_with_inflight_transfers()
for engine_id in (replaced - busy) | invalid:
self._cleanup_remote_engine(engine_id, log_eviction=False)
logger.info(
"Released NIXL state for replaced or locally invalid remote engine %s.",
engine_id,
)
def _mark_remote_engine_invalid(self, engine_id: EngineId, now: float) -> None:
"""Schedule release of an engine whose local NIXL metadata is gone.
The first invalidation is released as soon as reads drain. Repeats
before the peer has stayed healthy for the max backoff double the
delay, during which reads fail fast instead of re-handshaking.
"""
# Forget engines that have stayed healthy long enough.
for eid, (_, release_at) in list(self._invalid_engine_backoff.items()):
if now - release_at > _INVALID_PEER_MAX_BACKOFF_S:
del self._invalid_engine_backoff[eid]
prev = self._invalid_engine_backoff.get(engine_id)
backoff = 0.0
if prev is not None:
backoff = min(
max(2 * prev[0], _INVALID_PEER_MIN_BACKOFF_S),
_INVALID_PEER_MAX_BACKOFF_S,
)
self._invalid_engine_backoff[engine_id] = (backoff, now + backoff)
self._invalid_remote_engines.add(engine_id)
logger.warning(
"Local NIXL metadata for remote engine %s was invalidated; "
"re-handshaking after reads drain%s.",
engine_id,
f" and a {backoff:.0f}s backoff" if backoff else "",
)
def _cleanup_remote_engine(
self, engine_id: EngineId, *, log_eviction: bool = True
) -> None:
"""Remove all state for a single remote engine.
Releases NIXL resources (dlist handles, remote agents) and clears
all per-engine data structures. Used by TTL eviction, replaced/invalid
peer cleanup and shutdown.
"""
assert engine_id in self._remote_agents
# Notif-only engines (push-mode D side) have no descriptor state.
for handle in self.dst_xfer_side_handles.pop(engine_id, {}).values():
self.nixl_wrapper.release_dlist_handle(handle)
# Pop under the handshake lock; NIXL teardown stays outside it.
with self._handshake_lock:
agents = self._remote_agents.pop(engine_id)
for address, eid in list(self._engine_by_address.items()):
if eid == engine_id:
del self._engine_by_address[address]
for agent_name in agents.values():
self.nixl_wrapper.remove_remote_agent(agent_name)
self._failed_remote_engines.discard(engine_id)
self._invalid_remote_engines.discard(engine_id)
self.kv_caches_base_addr.pop(engine_id, None)
self.dst_num_blocks.pop(engine_id, None)
self.dst_region_num_blocks.pop(engine_id, None)
self.dst_region_group_ids.pop(engine_id, None)
self.dst_uses_region_group_mapping.pop(engine_id, None)
self.dst_region_mem_types.pop(engine_id, None)
self.tp_mappings.pop(engine_id, None)
if self.transfer_topo is not None:
self.transfer_topo.unregister_remote_engine(engine_id)
# Drop the cached clock offset; it is re-measured on the next handshake.
self._engine_clock_offset.pop(engine_id, None)
# A just-completed handshake may not have recorded activity yet, so
# tolerate a missing entry.
last_active = self._engine_last_active.pop(engine_id, None)
if log_eviction and last_active is not None:
logger.info(
"Evicted stale remote engine %s (inactive for %.1fs).",
engine_id,
time.perf_counter() - last_active,
)
def __del__(self):
with contextlib.suppress(Exception):
self.shutdown()
def _finish_shutdown(self) -> None:
self._recving_transfers.clear()
try:
for handle in self.src_xfer_handles_by_block_size.values():
self.nixl_wrapper.release_dlist_handle(handle)
for handles in self.src_xfer_handles_by_tp_ratio.values():
for handle in handles:
self.nixl_wrapper.release_dlist_handle(handle)
for dram_handles in self._dram_src_handles_by_tp_ratio.values():
for dram_handle in dram_handles:
if dram_handle is not None:
self.nixl_wrapper.release_dlist_handle(dram_handle)
for dram_handle in self._dram_src_handles_by_block_size.values():
if dram_handle is not None:
self.nixl_wrapper.release_dlist_handle(dram_handle)
except Exception:
logger.exception("NIXL dlist-handle release failed at shutdown.")
self.src_xfer_handles_by_block_size.clear()
self.src_xfer_handles_by_tp_ratio.clear()
self._dram_src_handles_by_tp_ratio.clear()
self._dram_src_handles_by_block_size.clear()
try:
for engine_id in list(self._remote_agents):
self._cleanup_remote_engine(engine_id, log_eviction=False)
except Exception:
logger.exception("NIXL remote-engine cleanup failed at shutdown.")
try:
for desc in self._registered_descs:
self.nixl_wrapper.deregister_memory(desc)
finally:
self._registered_descs.clear()
# Drop cache references before their owners release registered
# host memory; handshake futures may outlive model-runner shutdown.
self.device_kv_caches = {}
self.host_xfer_buffers = {}
def shutdown(self) -> None:
"""Shutdown the connector worker."""
if not hasattr(self, "_handshake_initiation_executor"):
# error happens during init, no need to shutdown
return
self._handshake_initiation_executor.shutdown(wait=False)
for handles in self._recving_transfers.values():
for handle in handles:
self.nixl_wrapper.release_xfer_handle(handle)
self._recving_transfers.clear()
self._finish_shutdown()