1370 lines
54 KiB
Python
1370 lines
54 KiB
Python
from __future__ import annotations
|
|
|
|
import dataclasses
|
|
import json
|
|
import logging
|
|
import struct
|
|
import threading
|
|
import time
|
|
import uuid
|
|
from collections import defaultdict
|
|
from typing import Dict, List, Optional, Set
|
|
|
|
import numpy as np
|
|
import numpy.typing as npt
|
|
|
|
from sglang.srt.disaggregation.base.conn import KVArgs, KVPoll
|
|
from sglang.srt.disaggregation.common.conn import (
|
|
CommonKVBootstrapServer,
|
|
CommonKVManager,
|
|
CommonKVReceiver,
|
|
CommonKVSender,
|
|
)
|
|
from sglang.srt.disaggregation.common.utils import group_concurrent_contiguous
|
|
from sglang.srt.disaggregation.utils import (
|
|
DisaggregationMode,
|
|
filter_kv_indices_for_cp_rank,
|
|
)
|
|
from sglang.srt.environ import envs
|
|
from sglang.srt.server_args import ServerArgs
|
|
|
|
try:
|
|
from nixl._bindings import (
|
|
nixlBackendError,
|
|
nixlCancelledError,
|
|
nixlRemoteDisconnectError,
|
|
)
|
|
|
|
_NIXL_TRANSPORT_ERRORS = (
|
|
nixlRemoteDisconnectError,
|
|
nixlBackendError,
|
|
nixlCancelledError,
|
|
)
|
|
except ImportError:
|
|
_NIXL_TRANSPORT_ERRORS = (RuntimeError,)
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
GUARD = "NixlMsgGuard".encode("ascii")
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class TransferInfo:
|
|
"""Contains indices for a transfer, sent by KVReceiver. Received by prefill bootstrap thread."""
|
|
|
|
room: int
|
|
endpoint: str
|
|
dst_port: int
|
|
agent_name: str
|
|
dst_kv_indices: npt.NDArray[np.int32]
|
|
dst_aux_index: int
|
|
required_dst_info_num: int
|
|
dst_state_indices: List[int]
|
|
decode_prefix_len: Optional[int] = None # for decode radix cache
|
|
|
|
def is_dummy(self):
|
|
# A transfer is "dummy" only for CP non-authoritative ranks.
|
|
# When dst_kv_indices is empty due to a decode-side radix cache
|
|
# full hit (decode_prefix_len > 0), the transfer is NOT dummy --
|
|
# aux/state data still needs to be sent.
|
|
if self.dst_kv_indices.size == 0 and self.decode_prefix_len:
|
|
return False
|
|
return self.dst_kv_indices.size == 0
|
|
|
|
@classmethod
|
|
def from_zmq(cls, msg: List[bytes]):
|
|
# Parse state_indices from msg[7] if present
|
|
if len(msg) > 7 and msg[7] != b"":
|
|
dst_state_indices = list(np.frombuffer(msg[7], dtype=np.int32))
|
|
else:
|
|
dst_state_indices = []
|
|
|
|
return cls(
|
|
room=int(msg[0].decode("ascii")),
|
|
endpoint=msg[1].decode("ascii"),
|
|
dst_port=int(msg[2].decode("ascii")),
|
|
agent_name=msg[3].decode("ascii"),
|
|
dst_kv_indices=np.frombuffer(msg[4], dtype=np.int32),
|
|
dst_aux_index=int(msg[5].decode("ascii")),
|
|
required_dst_info_num=int(msg[6].decode("ascii")),
|
|
dst_state_indices=dst_state_indices,
|
|
decode_prefix_len=(
|
|
int(msg[8].decode("ascii")) if len(msg) > 8 and msg[8] != b"" else None
|
|
), # hacky just add it into the message that will be sent
|
|
)
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class KVArgsRegisterInfo:
|
|
"""Contains base pointers and other info which only needs to be sent once by KVReceiver. Received by prefill bootstrap thread."""
|
|
|
|
room: str
|
|
endpoint: str
|
|
dst_port: int
|
|
agent_name: str
|
|
agent_metadata: bytes
|
|
dst_kv_ptrs: list[int]
|
|
dst_aux_ptrs: list[int]
|
|
dst_state_data_ptrs: list[int]
|
|
gpu_id: int
|
|
decode_tp_size: int
|
|
decode_tp_rank: int
|
|
dst_kv_item_len: int
|
|
dst_state_item_lens: list[int] = dataclasses.field(default_factory=list)
|
|
dst_state_dim_per_tensor: list[int] = dataclasses.field(default_factory=list)
|
|
|
|
@classmethod
|
|
def from_zmq(cls, msg: List[bytes]):
|
|
# Parse state_data_ptrs from msg[7] if present
|
|
if len(msg) > 7 and msg[7] != b"":
|
|
dst_state_data_ptrs = list(struct.unpack(f"{len(msg[7]) // 8}Q", msg[7]))
|
|
else:
|
|
dst_state_data_ptrs = []
|
|
|
|
dst_state_item_lens = []
|
|
dst_state_dim_per_tensor = []
|
|
if len(msg) > 12 and len(msg[12]) > 0:
|
|
dst_state_item_lens = list(struct.unpack(f"{len(msg[12]) // 4}I", msg[12]))
|
|
if len(msg) > 13 and len(msg[13]) > 0:
|
|
dst_state_dim_per_tensor = list(
|
|
struct.unpack(f"{len(msg[13]) // 4}I", msg[13])
|
|
)
|
|
|
|
return cls(
|
|
room=str(msg[0].decode("ascii")),
|
|
endpoint=msg[1].decode("ascii"),
|
|
dst_port=int(msg[2].decode("ascii")),
|
|
agent_name=msg[3].decode("ascii"),
|
|
agent_metadata=msg[4],
|
|
dst_kv_ptrs=list(struct.unpack(f"{len(msg[5]) // 8}Q", msg[5])),
|
|
dst_aux_ptrs=list(struct.unpack(f"{len(msg[6]) // 8}Q", msg[6])),
|
|
dst_state_data_ptrs=dst_state_data_ptrs,
|
|
gpu_id=int(msg[8].decode("ascii")),
|
|
decode_tp_size=int(msg[9].decode("ascii")),
|
|
decode_tp_rank=int(msg[10].decode("ascii")),
|
|
dst_kv_item_len=int(msg[11].decode("ascii")),
|
|
dst_state_item_lens=dst_state_item_lens,
|
|
dst_state_dim_per_tensor=dst_state_dim_per_tensor,
|
|
)
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class TransferStatus:
|
|
"""Used by KV Receiver to know when a transfer is done."""
|
|
|
|
# KV chunks received per pp_rank: {pp_rank: set of chunk_ids}
|
|
received_kvs_per_pp: Dict[int, Set[int]] = dataclasses.field(
|
|
default_factory=lambda: defaultdict(set)
|
|
)
|
|
# Expected chunk count per pp_rank (set when is_last=True): {pp_rank: expected_count}
|
|
expected_kvs_per_pp: Dict[int, int] = dataclasses.field(default_factory=dict)
|
|
# Number of PP ranks expected to send data.
|
|
num_pp_ranks_expected: Optional[int] = None
|
|
# Whether aux data has been received.
|
|
received_aux: bool = False
|
|
# PP ranks that have sent state data (state is layer-specific, each PP rank sends its portion).
|
|
received_state_per_pp: Set[int] = dataclasses.field(default_factory=set)
|
|
# Whether state data is expected (set based on state_type).
|
|
expects_state: bool = False
|
|
# Mark as failed
|
|
is_failure: bool = False
|
|
|
|
def is_done(self):
|
|
if self.is_failure:
|
|
return True
|
|
if self.num_pp_ranks_expected is None or not self.received_aux:
|
|
return False
|
|
# If state data is expected, check all PP ranks have sent it
|
|
if (
|
|
self.expects_state
|
|
and len(self.received_state_per_pp) < self.num_pp_ranks_expected
|
|
):
|
|
return False
|
|
# All PP ranks must have reported their expected count
|
|
if len(self.expected_kvs_per_pp) < self.num_pp_ranks_expected:
|
|
return False
|
|
# Each PP rank must have received all expected chunks
|
|
for pp_rank, expected in self.expected_kvs_per_pp.items():
|
|
if len(self.received_kvs_per_pp[pp_rank]) != expected:
|
|
return False
|
|
return True
|
|
|
|
def is_failed(self):
|
|
return self.is_failure
|
|
|
|
|
|
class NixlKVManager(CommonKVManager):
|
|
def __init__(
|
|
self,
|
|
args: KVArgs,
|
|
disaggregation_mode: DisaggregationMode,
|
|
server_args: ServerArgs,
|
|
is_mla_backend: Optional[bool] = False,
|
|
):
|
|
super().__init__(args, disaggregation_mode, server_args, is_mla_backend)
|
|
try:
|
|
from nixl._api import nixl_agent, nixl_agent_config
|
|
except ImportError as e:
|
|
raise ImportError(
|
|
"Please install NIXL by following the instructions at "
|
|
"https://github.com/ai-dynamo/nixl/blob/main/README.md "
|
|
"to run SGLang with NixlTransferEngine."
|
|
) from e
|
|
|
|
backend = envs.SGLANG_DISAGGREGATION_NIXL_BACKEND.get()
|
|
num_threads = 8 if disaggregation_mode == DisaggregationMode.PREFILL else 0
|
|
backend_params = json.loads(
|
|
envs.SGLANG_DISAGGREGATION_NIXL_BACKEND_PARAMS.get()
|
|
)
|
|
if not isinstance(backend_params, dict) or not all(
|
|
isinstance(key, str) and isinstance(value, str)
|
|
for key, value in backend_params.items()
|
|
):
|
|
raise ValueError(
|
|
"SGLANG_DISAGGREGATION_NIXL_BACKEND_PARAMS must be a JSON object "
|
|
"with string keys and string values"
|
|
)
|
|
agent_config = nixl_agent_config(backends=[], num_threads=num_threads)
|
|
self.agent = nixl_agent(str(uuid.uuid4()), agent_config)
|
|
if num_threads > 0:
|
|
# TODO: Remove this once NIXL passes thread parameters from
|
|
# nixl_agent_config to explicitly-created backends.
|
|
if backend == "UCX" or backend == "OBJ":
|
|
backend_params.setdefault("num_threads", str(num_threads))
|
|
elif backend == "GDS_MT":
|
|
backend_params.setdefault("thread_count", str(num_threads))
|
|
elif backend == "UCCL":
|
|
backend_params.setdefault("num_cpus", str(num_threads))
|
|
self.agent.create_backend(backend, backend_params)
|
|
|
|
available_plugins = self.agent.get_plugin_list()
|
|
if backend not in available_plugins:
|
|
raise ValueError(
|
|
f"NIXL backend '{backend}' not found. Available: {available_plugins}. "
|
|
f"Please install the required NIXL plugin or choose from: {available_plugins}"
|
|
)
|
|
logger.info(f"NIXL KVManager initialized with backend: {backend}")
|
|
|
|
self.register_buffer_to_engine()
|
|
|
|
if self.disaggregation_mode == DisaggregationMode.PREFILL:
|
|
self._start_bootstrap_thread()
|
|
elif self.disaggregation_mode == DisaggregationMode.DECODE:
|
|
self.transfer_statuses: Dict[int, TransferStatus] = defaultdict(
|
|
TransferStatus
|
|
)
|
|
self._start_heartbeat_checker_thread()
|
|
else:
|
|
raise ValueError(
|
|
f"Unsupported DisaggregationMode: {self.disaggregation_mode}"
|
|
)
|
|
|
|
def _start_heartbeat_checker_thread(self):
|
|
"""
|
|
Start the heartbeat checker thread for Decode worker.
|
|
TODO (smor): unite nixl heartbeat checker with mooncake's.
|
|
"""
|
|
|
|
def heartbeat_checker():
|
|
while True:
|
|
time.sleep(self.heartbeat_interval)
|
|
with self.connection_lock:
|
|
addresses = list(self.prefill_info_table.keys())
|
|
|
|
for bootstrap_addr in addresses:
|
|
session = None
|
|
try:
|
|
with self.session_pool_lock:
|
|
session = self.session_pool[bootstrap_addr]
|
|
response = session.get(
|
|
f"http://{bootstrap_addr}/health",
|
|
timeout=(2, 3),
|
|
headers={"Connection": "keep-alive"},
|
|
)
|
|
if response.status_code == 200:
|
|
self.heartbeat_failures[bootstrap_addr] = 0
|
|
|
|
else:
|
|
logger.info(
|
|
f"Attempting to reconnect to {bootstrap_addr}..."
|
|
)
|
|
self.heartbeat_failures[bootstrap_addr] = (
|
|
self.heartbeat_failures.get(bootstrap_addr, 0) + 1
|
|
)
|
|
with self.session_pool_lock:
|
|
if bootstrap_addr in self.session_pool:
|
|
del self.session_pool[bootstrap_addr]
|
|
except Exception:
|
|
logger.info(f"Attempting to reconnect to {bootstrap_addr}...")
|
|
self.heartbeat_failures[bootstrap_addr] = (
|
|
self.heartbeat_failures.get(bootstrap_addr, 0) + 1
|
|
)
|
|
|
|
if (
|
|
self.heartbeat_failures.get(bootstrap_addr, 0)
|
|
>= self.max_failures
|
|
):
|
|
self._handle_node_failure(bootstrap_addr)
|
|
with self.session_pool_lock:
|
|
if bootstrap_addr in self.session_pool:
|
|
del self.session_pool[bootstrap_addr]
|
|
|
|
threading.Thread(target=heartbeat_checker, daemon=True).start()
|
|
|
|
def _handle_node_failure(self, failed_bootstrap_addr):
|
|
"""Handle failure of a prefill node."""
|
|
with self.connection_lock:
|
|
keys_to_remove = [
|
|
k for k in self.connection_pool if k.startswith(failed_bootstrap_addr)
|
|
]
|
|
for k in keys_to_remove:
|
|
del self.connection_pool[k]
|
|
self.prefill_info_table.pop(failed_bootstrap_addr, None)
|
|
|
|
possible_affected_rooms = self.addr_to_rooms_tracker.get(
|
|
failed_bootstrap_addr, []
|
|
)
|
|
self.addr_to_rooms_tracker.pop(failed_bootstrap_addr, None)
|
|
|
|
# Mark all pending transfers associated with the failed node as failed
|
|
affected_rooms = []
|
|
for room in possible_affected_rooms:
|
|
if (
|
|
room in self.transfer_statuses
|
|
and not self.transfer_statuses[room].is_done()
|
|
):
|
|
# Mark the transfer as failed
|
|
self.transfer_statuses[room].is_failure = True
|
|
affected_rooms.append(room)
|
|
|
|
logger.error(
|
|
f"Lost connection with prefill instance (bootstrap_addr: {failed_bootstrap_addr}), "
|
|
f"{len(affected_rooms)} transfers affected"
|
|
)
|
|
for room in possible_affected_rooms:
|
|
logger.error(f"Let room {room} be failed due to prefill down")
|
|
self.update_status(room, KVPoll.Failed)
|
|
|
|
def register_buffer_to_engine(self):
|
|
kv_addrs = []
|
|
for kv_data_ptr, kv_data_len in zip(
|
|
self.kv_args.kv_data_ptrs, self.kv_args.kv_data_lens
|
|
):
|
|
kv_addrs.append((kv_data_ptr, kv_data_len, self.kv_args.gpu_id, ""))
|
|
self.kv_descs = self.agent.register_memory(kv_addrs, "VRAM")
|
|
logger.debug(f"Register kv tensors, len(kv_addr)= {len(kv_addrs)}")
|
|
if not self.kv_descs:
|
|
raise Exception("NIXL memory registration failed for kv tensors")
|
|
aux_addrs = []
|
|
for aux_data_ptr, aux_data_len in zip(
|
|
self.kv_args.aux_data_ptrs, self.kv_args.aux_data_lens
|
|
):
|
|
aux_addrs.append((aux_data_ptr, aux_data_len, 0, ""))
|
|
self.aux_descs = self.agent.register_memory(aux_addrs, "DRAM")
|
|
logger.debug(f"Register aux tensors, len(aux_addrs)= {len(aux_addrs)}")
|
|
if not self.aux_descs:
|
|
raise Exception("NIXL memory registration failed for aux tensors")
|
|
|
|
# Register state/extra pool data buffers if present
|
|
if self.kv_args.state_data_ptrs and self.kv_args.state_data_lens:
|
|
state_addrs = []
|
|
for state_data_ptr, state_data_len in zip(
|
|
self.kv_args.state_data_ptrs, self.kv_args.state_data_lens
|
|
):
|
|
state_addrs.append(
|
|
(state_data_ptr, state_data_len, self.kv_args.gpu_id, "")
|
|
)
|
|
self.state_descs = self.agent.register_memory(state_addrs, "VRAM")
|
|
logger.debug(
|
|
f"Register state tensors, len(state_addrs)= {len(state_addrs)}"
|
|
)
|
|
if not self.state_descs:
|
|
raise Exception("NIXL memory registration failed for state tensors")
|
|
|
|
def _add_remote_peer(self, decode_kv_args: KVArgsRegisterInfo):
|
|
agent_name = decode_kv_args.agent_name
|
|
if agent_name in self.decode_kv_args_table:
|
|
logger.info(f"Peer {agent_name} was already registered, ignoring.")
|
|
return
|
|
self.decode_kv_args_table[agent_name] = decode_kv_args
|
|
self.agent.add_remote_agent(decode_kv_args.agent_metadata)
|
|
|
|
def _send_kvcache_generic(
|
|
self,
|
|
peer_name: str,
|
|
src_data_ptrs: list[int],
|
|
dst_data_ptrs: list[int],
|
|
item_lens: list[int],
|
|
prefill_data_indices: npt.NDArray[np.int32],
|
|
dst_data_indices: npt.NDArray[np.int32],
|
|
dst_gpu_id: int,
|
|
notif: str,
|
|
):
|
|
"""Generic KV cache transfer supporting both MHA and MLA architectures.
|
|
Used by both send_kvcache and maybe_send_extra."""
|
|
# Convert pointer lists to np.uint64 arrays up front.
|
|
# torch.int exceeds np.int64 range on Intel XPU (addresses have bit 63 set, e.g.
|
|
# 0xffff81ab54e01000). Casting here prevents overflow when these values
|
|
# are later used in numpy arithmetic.
|
|
src_data_ptrs = np.array(src_data_ptrs, dtype=np.uint64)
|
|
dst_data_ptrs = np.array(dst_data_ptrs, dtype=np.uint64)
|
|
item_lens = np.array(item_lens, dtype=np.uint64)
|
|
|
|
# group by indices
|
|
prefill_kv_blocks, dst_kv_blocks = group_concurrent_contiguous(
|
|
prefill_data_indices, dst_data_indices
|
|
)
|
|
|
|
logger.debug(f"sending kvcache to {peer_name} with notif {notif}")
|
|
# Make descs
|
|
if self.is_mla_backend:
|
|
src_kv_ptrs, dst_kv_ptrs, layers_current_pp_stage = (
|
|
self.get_mla_kv_ptrs_with_pp(src_data_ptrs, dst_data_ptrs)
|
|
)
|
|
layers_params = [
|
|
(
|
|
src_kv_ptrs[layer_id],
|
|
dst_kv_ptrs[layer_id],
|
|
item_lens[layer_id],
|
|
)
|
|
for layer_id in range(layers_current_pp_stage)
|
|
]
|
|
else:
|
|
src_k_ptrs, src_v_ptrs, dst_k_ptrs, dst_v_ptrs, layers_current_pp_stage = (
|
|
self.get_mha_kv_ptrs_with_pp(src_data_ptrs, dst_data_ptrs)
|
|
)
|
|
|
|
layers_params = [
|
|
(
|
|
src_k_ptrs[layer_id],
|
|
dst_k_ptrs[layer_id],
|
|
item_lens[layer_id],
|
|
)
|
|
for layer_id in range(layers_current_pp_stage)
|
|
] + [
|
|
(
|
|
src_v_ptrs[layer_id],
|
|
dst_v_ptrs[layer_id],
|
|
item_lens[layer_id],
|
|
)
|
|
for layer_id in range(layers_current_pp_stage)
|
|
]
|
|
|
|
src_addrs = []
|
|
src_lens = []
|
|
dst_addrs = []
|
|
dst_lens = []
|
|
|
|
# Precompute block starts/lengths to reduce Python-level loops.
|
|
prefill_starts = np.fromiter(
|
|
(block[0] for block in prefill_kv_blocks), dtype=np.uint64
|
|
)
|
|
dst_starts = np.fromiter((block[0] for block in dst_kv_blocks), dtype=np.uint64)
|
|
block_lens = np.fromiter(
|
|
(len(block) for block in prefill_kv_blocks), dtype=np.uint64
|
|
)
|
|
|
|
for src_ptr, dst_ptr, item_len in layers_params:
|
|
lengths = item_len * block_lens
|
|
src_addrs.append(src_ptr + prefill_starts * item_len)
|
|
src_lens.append(lengths)
|
|
dst_addrs.append(dst_ptr + dst_starts * item_len)
|
|
dst_lens.append(lengths)
|
|
|
|
def make_req_array(addr_chunks, len_chunks, gpu):
|
|
if not addr_chunks:
|
|
return np.empty((0, 3), dtype=np.uint64)
|
|
flat_addrs = np.concatenate(addr_chunks).astype(np.uint64, copy=False)
|
|
flat_lens = np.concatenate(len_chunks).astype(np.uint64, copy=False)
|
|
return np.column_stack(
|
|
(
|
|
flat_addrs,
|
|
flat_lens,
|
|
np.full_like(flat_addrs, gpu, dtype=np.uint64),
|
|
)
|
|
)
|
|
|
|
src_reqs = make_req_array(src_addrs, src_lens, self.kv_args.gpu_id)
|
|
dst_reqs = make_req_array(dst_addrs, dst_lens, dst_gpu_id)
|
|
|
|
logger.debug(
|
|
f"len(src_addrs): before group: {len(prefill_data_indices)}, after group: {len(src_addrs)}"
|
|
)
|
|
src_descs = self.agent.get_xfer_descs(src_reqs, "VRAM")
|
|
dst_descs = self.agent.get_xfer_descs(dst_reqs, "VRAM")
|
|
# Transfer data
|
|
xfer_handle = self.agent.initialize_xfer(
|
|
"WRITE",
|
|
src_descs,
|
|
dst_descs,
|
|
peer_name,
|
|
notif.encode("ascii"), # type: ignore
|
|
)
|
|
if not xfer_handle:
|
|
raise Exception("KVSender failed to create transfer")
|
|
state = self.agent.transfer(xfer_handle)
|
|
if state == "ERR":
|
|
raise Exception("KVSender failed to post transfer")
|
|
return xfer_handle
|
|
|
|
def send_kvcache(
|
|
self,
|
|
peer_name: str,
|
|
prefill_kv_indices: npt.NDArray[np.int32],
|
|
dst_kv_ptrs: list[int],
|
|
dst_kv_indices: npt.NDArray[np.int32],
|
|
dst_gpu_id: int,
|
|
notif: str,
|
|
):
|
|
return self._send_kvcache_generic(
|
|
peer_name=peer_name,
|
|
src_data_ptrs=self.kv_args.kv_data_ptrs,
|
|
dst_data_ptrs=dst_kv_ptrs,
|
|
item_lens=self.kv_args.kv_item_lens,
|
|
prefill_data_indices=prefill_kv_indices,
|
|
dst_data_indices=dst_kv_indices,
|
|
dst_gpu_id=dst_gpu_id,
|
|
notif=notif,
|
|
)
|
|
|
|
def send_kvcache_slice(
|
|
self,
|
|
peer_name: str,
|
|
prefill_kv_indices: npt.NDArray[np.int32],
|
|
dst_kv_ptrs: list[int],
|
|
dst_kv_indices: npt.NDArray[np.int32],
|
|
dst_gpu_id: int,
|
|
notif: str,
|
|
prefill_tp_size: int,
|
|
decode_tp_size: int,
|
|
decode_tp_rank: int,
|
|
dst_kv_item_len: int,
|
|
):
|
|
# Get configuration from kv_args
|
|
local_tp_rank_in_group = self.kv_args.engine_rank % prefill_tp_size
|
|
dst_tp_rank_in_group = decode_tp_rank % decode_tp_size
|
|
|
|
src_kv_item_len = self.kv_args.kv_item_lens[0]
|
|
page_size = self.kv_args.page_size
|
|
|
|
# Use total KV head count (not per-rank) for correct head distribution.
|
|
# Per-rank kv_head_num is max(1, total//tp) which loses info when total < tp.
|
|
total_kv_heads = getattr(self.kv_args, "total_kv_head_num", 0)
|
|
if total_kv_heads <= 0:
|
|
total_kv_heads = self.kv_args.kv_head_num * prefill_tp_size
|
|
|
|
src_heads_per_rank = max(1, total_kv_heads // prefill_tp_size)
|
|
dst_heads_per_rank = max(1, total_kv_heads // decode_tp_size)
|
|
|
|
bytes_per_head_slice_to_send = (
|
|
dst_kv_item_len // page_size // dst_heads_per_rank
|
|
)
|
|
|
|
# GQA replication: how many prefill ranks share the same KV head
|
|
src_replication = max(1, prefill_tp_size // total_kv_heads)
|
|
|
|
# Determine which heads to send
|
|
if prefill_tp_size > decode_tp_size:
|
|
# Multiple prefill ranks to one decode rank
|
|
src_head_start_offset = 0
|
|
num_heads_to_send = src_heads_per_rank
|
|
unique_head_idx = local_tp_rank_in_group // src_replication
|
|
dst_head_start_offset = (
|
|
unique_head_idx * src_heads_per_rank
|
|
) % dst_heads_per_rank
|
|
else:
|
|
# Send KVCache from 1 prefill instance to multiple decode instances
|
|
src_head_start_offset = (
|
|
dst_tp_rank_in_group * dst_heads_per_rank
|
|
) % src_heads_per_rank
|
|
num_heads_to_send = dst_heads_per_rank
|
|
dst_head_start_offset = 0
|
|
|
|
src_k_ptrs, src_v_ptrs, dst_k_ptrs, dst_v_ptrs, layers_current_pp_stage = (
|
|
self.get_mha_kv_ptrs_with_pp(self.kv_args.kv_data_ptrs, dst_kv_ptrs)
|
|
)
|
|
# Calculate precise byte offset and length for the sub-slice within the token
|
|
src_head_slice_offset = src_head_start_offset * bytes_per_head_slice_to_send
|
|
dst_head_slice_offset = dst_head_start_offset * bytes_per_head_slice_to_send
|
|
heads_bytes_per_token_to_send = num_heads_to_send * bytes_per_head_slice_to_send
|
|
|
|
src_dst_ptr_pairs = [
|
|
(
|
|
src_k_ptrs[layer_id],
|
|
dst_k_ptrs[layer_id],
|
|
)
|
|
for layer_id in range(layers_current_pp_stage)
|
|
] + [
|
|
(
|
|
src_v_ptrs[layer_id],
|
|
dst_v_ptrs[layer_id],
|
|
)
|
|
for layer_id in range(layers_current_pp_stage)
|
|
]
|
|
|
|
prefill_indices = np.asarray(prefill_kv_indices, dtype=np.int64)
|
|
dst_indices = np.asarray(dst_kv_indices, dtype=np.int64)
|
|
bytes_per_token_prefill = src_kv_item_len // page_size
|
|
bytes_per_token_decode = dst_kv_item_len // page_size
|
|
token_offsets = np.arange(page_size, dtype=np.int64)
|
|
|
|
src_addrs = []
|
|
dst_addrs = []
|
|
|
|
for src_ptr, dst_ptr in src_dst_ptr_pairs:
|
|
src_page_bases = src_ptr + prefill_indices * src_kv_item_len
|
|
dst_page_bases = dst_ptr + dst_indices * dst_kv_item_len
|
|
|
|
src_all = (
|
|
src_page_bases[:, None]
|
|
+ token_offsets[None, :] * bytes_per_token_prefill
|
|
+ src_head_slice_offset
|
|
).ravel()
|
|
dst_all = (
|
|
dst_page_bases[:, None]
|
|
+ token_offsets[None, :] * bytes_per_token_decode
|
|
+ dst_head_slice_offset
|
|
).ravel()
|
|
|
|
src_addrs.append(src_all)
|
|
dst_addrs.append(dst_all)
|
|
|
|
def make_req_array(addr_chunks, size, gpu):
|
|
if not addr_chunks:
|
|
return np.empty((0, 3), dtype=np.uint64)
|
|
flat_addrs = np.concatenate(addr_chunks).astype(np.uint64, copy=False)
|
|
return np.column_stack(
|
|
(
|
|
flat_addrs,
|
|
np.full_like(flat_addrs, size, dtype=np.uint64),
|
|
np.full_like(flat_addrs, gpu, dtype=np.uint64),
|
|
)
|
|
)
|
|
|
|
src_reqs = make_req_array(
|
|
src_addrs, heads_bytes_per_token_to_send, self.kv_args.gpu_id
|
|
)
|
|
dst_reqs = make_req_array(dst_addrs, heads_bytes_per_token_to_send, dst_gpu_id)
|
|
|
|
# Use NIXL agent for transfer
|
|
src_descs = self.agent.get_xfer_descs(src_reqs, "VRAM")
|
|
dst_descs = self.agent.get_xfer_descs(dst_reqs, "VRAM")
|
|
|
|
xfer_handle = self.agent.initialize_xfer(
|
|
"WRITE", src_descs, dst_descs, peer_name, notif.encode("ascii")
|
|
)
|
|
if not xfer_handle:
|
|
raise Exception("Failed to create sliced KV transfer")
|
|
|
|
state = self.agent.transfer(xfer_handle)
|
|
if state == "ERR":
|
|
raise Exception("Failed to post sliced KV transfer")
|
|
|
|
return xfer_handle
|
|
|
|
def send_aux(
|
|
self,
|
|
peer_name: str,
|
|
prefill_aux_index: int,
|
|
dst_aux_ptrs: list[int],
|
|
dst_aux_index: int,
|
|
notif: str,
|
|
):
|
|
src_addrs = []
|
|
dst_addrs = []
|
|
|
|
prefill_aux_ptrs = self.kv_args.aux_data_ptrs
|
|
prefill_aux_item_lens = self.kv_args.aux_item_lens
|
|
|
|
for i, _ in enumerate(dst_aux_ptrs):
|
|
length = prefill_aux_item_lens[i]
|
|
src_addr = prefill_aux_ptrs[i] + length * prefill_aux_index
|
|
dst_addr = dst_aux_ptrs[i] + length * dst_aux_index
|
|
src_addrs.append((src_addr, length, 0))
|
|
dst_addrs.append((dst_addr, length, 0))
|
|
|
|
src_descs = self.agent.get_xfer_descs(src_addrs, "DRAM")
|
|
dst_descs = self.agent.get_xfer_descs(dst_addrs, "DRAM")
|
|
# Transfer data
|
|
xfer_handle = self.agent.initialize_xfer(
|
|
"WRITE",
|
|
src_descs,
|
|
dst_descs,
|
|
peer_name,
|
|
notif.encode("ascii"), # type: ignore
|
|
)
|
|
if not xfer_handle:
|
|
raise Exception("KVSender failed to create transfer")
|
|
state = self.agent.transfer(xfer_handle)
|
|
if state == "ERR":
|
|
raise Exception("KVSender failed to post transfer")
|
|
return xfer_handle
|
|
|
|
def _send_mamba_state(
|
|
self,
|
|
peer_name: str,
|
|
prefill_state_indices: List[int],
|
|
dst_state_data_ptrs: list[int],
|
|
dst_state_indices: List[int],
|
|
dst_gpu_id: int,
|
|
notif: str,
|
|
):
|
|
"""Transfer Mamba states via RDMA."""
|
|
assert len(prefill_state_indices) == 1, "Mamba should have single state index"
|
|
assert len(dst_state_indices) == len(
|
|
prefill_state_indices
|
|
), "State indices count mismatch between Prefill and Decode"
|
|
|
|
src_addrs = []
|
|
dst_addrs = []
|
|
|
|
prefill_state_data_ptrs = self.kv_args.state_data_ptrs
|
|
prefill_state_item_lens = self.kv_args.state_item_lens
|
|
|
|
for i, dst_state_ptr in enumerate(dst_state_data_ptrs):
|
|
length = prefill_state_item_lens[i]
|
|
src_addr = prefill_state_data_ptrs[i] + length * int(
|
|
prefill_state_indices[0]
|
|
)
|
|
dst_addr = dst_state_ptr + length * int(dst_state_indices[0])
|
|
src_addrs.append((src_addr, length, self.kv_args.gpu_id))
|
|
dst_addrs.append((dst_addr, length, dst_gpu_id))
|
|
|
|
src_descs = self.agent.get_xfer_descs(src_addrs, "VRAM")
|
|
dst_descs = self.agent.get_xfer_descs(dst_addrs, "VRAM")
|
|
|
|
xfer_handle = self.agent.initialize_xfer(
|
|
"WRITE",
|
|
src_descs,
|
|
dst_descs,
|
|
peer_name,
|
|
notif.encode("ascii"),
|
|
)
|
|
if not xfer_handle:
|
|
raise Exception("Failed to create Mamba state transfer")
|
|
state = self.agent.transfer(xfer_handle)
|
|
if state == "ERR":
|
|
raise Exception("Failed to post Mamba state transfer")
|
|
return xfer_handle
|
|
|
|
def _send_mamba_state_slice(
|
|
self,
|
|
peer_name: str,
|
|
prefill_state_indices: List[int],
|
|
dst_state_data_ptrs: list[int],
|
|
dst_state_indices: List[int],
|
|
dst_gpu_id: int,
|
|
notif: str,
|
|
dst_state_item_lens: list[int],
|
|
dst_state_dim_per_tensor: list[int],
|
|
decode_tp_size: int,
|
|
decode_tp_rank: int,
|
|
):
|
|
"""Transfer Mamba states with TP slice support via RDMA.
|
|
|
|
When prefill and decode have different attn_tp_size, we slice the
|
|
TP-sharded dimension (3rd dim) of conv_state and temporal_state
|
|
accordingly, mirroring Mooncake's _send_mamba_state_slice.
|
|
"""
|
|
logger.warning_once(
|
|
"Using Mamba state slice transfer for different TP sizes. "
|
|
f"Prefill attn_tp_size={self.attn_tp_size}, "
|
|
f"Decode attn_tp_size={decode_tp_size}."
|
|
)
|
|
assert len(prefill_state_indices) == 1, "Mamba should have single state index"
|
|
|
|
prefill_state_data_ptrs = self.kv_args.state_data_ptrs
|
|
prefill_state_item_lens = self.kv_args.state_item_lens
|
|
src_state_dim_per_tensor = getattr(self.kv_args, "state_dim_per_tensor", [])
|
|
|
|
if not src_state_dim_per_tensor or not dst_state_dim_per_tensor:
|
|
return self._send_mamba_state(
|
|
peer_name,
|
|
prefill_state_indices,
|
|
dst_state_data_ptrs,
|
|
dst_state_indices,
|
|
dst_gpu_id,
|
|
notif,
|
|
)
|
|
|
|
local_tp_rank_in_group = self.kv_args.engine_rank % self.attn_tp_size
|
|
dst_tp_rank_in_group = decode_tp_rank % decode_tp_size
|
|
|
|
src_addrs = []
|
|
dst_addrs = []
|
|
|
|
for i, dst_state_ptr in enumerate(dst_state_data_ptrs):
|
|
src_item_len = prefill_state_item_lens[i]
|
|
dst_item_len = dst_state_item_lens[i]
|
|
src_dim = src_state_dim_per_tensor[i]
|
|
dst_dim = dst_state_dim_per_tensor[i]
|
|
|
|
src_bytes_per_dim = src_item_len // src_dim
|
|
dst_bytes_per_dim = dst_item_len // dst_dim
|
|
|
|
if self.attn_tp_size > decode_tp_size:
|
|
src_dim_start = 0
|
|
num_dims_to_send = src_dim
|
|
writers_per_decode = self.attn_tp_size // decode_tp_size
|
|
local_writer_idx = local_tp_rank_in_group % writers_per_decode
|
|
dst_dim_start = local_writer_idx * src_dim
|
|
else:
|
|
src_dim_start = (dst_tp_rank_in_group * dst_dim) % src_dim
|
|
num_dims_to_send = dst_dim
|
|
dst_dim_start = 0
|
|
|
|
src_dim_offset = src_dim_start * src_bytes_per_dim
|
|
dst_dim_offset = dst_dim_start * dst_bytes_per_dim
|
|
bytes_to_send = num_dims_to_send * src_bytes_per_dim
|
|
|
|
src_addr = (
|
|
prefill_state_data_ptrs[i]
|
|
+ src_item_len * int(prefill_state_indices[0])
|
|
+ src_dim_offset
|
|
)
|
|
dst_addr = (
|
|
dst_state_ptr
|
|
+ dst_item_len * int(dst_state_indices[0])
|
|
+ dst_dim_offset
|
|
)
|
|
src_addrs.append((src_addr, bytes_to_send, self.kv_args.gpu_id))
|
|
dst_addrs.append((dst_addr, bytes_to_send, dst_gpu_id))
|
|
|
|
src_descs = self.agent.get_xfer_descs(src_addrs, "VRAM")
|
|
dst_descs = self.agent.get_xfer_descs(dst_addrs, "VRAM")
|
|
|
|
xfer_handle = self.agent.initialize_xfer(
|
|
"WRITE",
|
|
src_descs,
|
|
dst_descs,
|
|
peer_name,
|
|
notif.encode("ascii"),
|
|
)
|
|
if not xfer_handle:
|
|
raise Exception("Failed to create Mamba state slice transfer")
|
|
state = self.agent.transfer(xfer_handle)
|
|
if state == "ERR":
|
|
raise Exception("Failed to post Mamba state slice transfer")
|
|
return xfer_handle
|
|
|
|
def maybe_send_extra(
|
|
self,
|
|
peer_name: str,
|
|
prefill_state_indices: List[int],
|
|
dst_state_data_ptrs: list[int],
|
|
dst_state_indices: List[int],
|
|
dst_gpu_id: int,
|
|
notif: str,
|
|
decode_tp_size: int,
|
|
decode_tp_rank: int = 0,
|
|
dst_state_item_lens: list[int] | None = None,
|
|
dst_state_dim_per_tensor: list[int] | None = None,
|
|
):
|
|
"""Send state or extra pool data with type-specific handling."""
|
|
state_type = getattr(self.kv_args, "state_type", "none")
|
|
|
|
if state_type == "mamba":
|
|
if self.attn_tp_size != decode_tp_size:
|
|
return self._send_mamba_state_slice(
|
|
peer_name,
|
|
prefill_state_indices,
|
|
dst_state_data_ptrs,
|
|
dst_state_indices,
|
|
dst_gpu_id,
|
|
notif,
|
|
dst_state_item_lens or [],
|
|
dst_state_dim_per_tensor or [],
|
|
decode_tp_size,
|
|
decode_tp_rank,
|
|
)
|
|
return self._send_mamba_state(
|
|
peer_name,
|
|
prefill_state_indices,
|
|
dst_state_data_ptrs,
|
|
dst_state_indices,
|
|
dst_gpu_id,
|
|
notif,
|
|
)
|
|
elif state_type in ["swa", "nsa"]:
|
|
if not self.is_mla_backend and self.attn_tp_size != decode_tp_size:
|
|
raise RuntimeError(
|
|
f"PD Disaggregation does NOT support PD different TP sizes for non-MLA {state_type.upper()} hybrid models yet."
|
|
)
|
|
if len(prefill_state_indices) != len(dst_state_indices):
|
|
raise RuntimeError(
|
|
f"State index length mismatch: prefill={len(prefill_state_indices)}, "
|
|
f"dst={len(dst_state_indices)}"
|
|
)
|
|
return self._send_kvcache_generic(
|
|
peer_name=peer_name,
|
|
src_data_ptrs=self.kv_args.state_data_ptrs,
|
|
dst_data_ptrs=dst_state_data_ptrs,
|
|
item_lens=self.kv_args.state_item_lens,
|
|
prefill_data_indices=np.array(prefill_state_indices, dtype=np.int32),
|
|
dst_data_indices=np.array(dst_state_indices, dtype=np.int32),
|
|
dst_gpu_id=dst_gpu_id,
|
|
notif=notif,
|
|
)
|
|
else:
|
|
if state_type != "none":
|
|
raise RuntimeError(
|
|
f"PD Disaggregation via NIXL does NOT support {state_type} hybrid models yet."
|
|
)
|
|
return None
|
|
|
|
def add_transfer_request(
|
|
self,
|
|
bootstrap_room: int,
|
|
kv_indices: npt.NDArray[np.int32],
|
|
index_slice: slice,
|
|
is_last: bool,
|
|
chunk_id: int,
|
|
aux_index: Optional[int] = None,
|
|
state_indices: Optional[List[int]] = None,
|
|
):
|
|
assert self.disaggregation_mode == DisaggregationMode.PREFILL
|
|
assert not is_last or (is_last and aux_index is not None)
|
|
|
|
reqs_to_be_processed = self.transfer_infos[bootstrap_room].values()
|
|
handles = []
|
|
for req in reqs_to_be_processed:
|
|
assert bootstrap_room == req.room
|
|
if req.is_dummy():
|
|
continue
|
|
|
|
chunked_dst_kv_indice = req.dst_kv_indices[index_slice]
|
|
assert len(chunked_dst_kv_indice) == len(kv_indices)
|
|
assert req.agent_name in self.decode_kv_args_table
|
|
|
|
decode_tp_size = self.decode_kv_args_table[req.agent_name].decode_tp_size
|
|
|
|
# Skip KV RDMA transfer when there are no pages to send
|
|
# (e.g., decode-side radix cache matched the entire prefix).
|
|
# Aux data is still sent below when is_last=True.
|
|
if len(kv_indices) > 0:
|
|
notif = (
|
|
f"{req.room}_kv_{chunk_id}_{int(is_last)}_{self.kv_args.pp_rank}"
|
|
)
|
|
|
|
if self.is_mla_backend or (decode_tp_size == self.attn_tp_size):
|
|
kv_xfer_handle = self.send_kvcache(
|
|
req.agent_name,
|
|
kv_indices,
|
|
self.decode_kv_args_table[req.agent_name].dst_kv_ptrs,
|
|
chunked_dst_kv_indice,
|
|
self.decode_kv_args_table[req.agent_name].gpu_id,
|
|
notif,
|
|
)
|
|
else:
|
|
kv_xfer_handle = self.send_kvcache_slice(
|
|
req.agent_name,
|
|
kv_indices,
|
|
self.decode_kv_args_table[req.agent_name].dst_kv_ptrs,
|
|
chunked_dst_kv_indice,
|
|
self.decode_kv_args_table[req.agent_name].gpu_id,
|
|
notif,
|
|
prefill_tp_size=self.attn_tp_size,
|
|
decode_tp_size=decode_tp_size,
|
|
decode_tp_rank=self.decode_kv_args_table[
|
|
req.agent_name
|
|
].decode_tp_rank,
|
|
dst_kv_item_len=self.decode_kv_args_table[
|
|
req.agent_name
|
|
].dst_kv_item_len,
|
|
)
|
|
|
|
handles.append(kv_xfer_handle)
|
|
# Only the last chunk we need to send the aux data.
|
|
if is_last:
|
|
if state_indices is not None:
|
|
dst_info = self.decode_kv_args_table[req.agent_name]
|
|
state_xfer_handle = self.maybe_send_extra(
|
|
req.agent_name,
|
|
state_indices,
|
|
dst_info.dst_state_data_ptrs,
|
|
req.dst_state_indices,
|
|
dst_info.gpu_id,
|
|
f"{req.room}_state_{self.kv_args.engine_rank}",
|
|
decode_tp_size,
|
|
decode_tp_rank=dst_info.decode_tp_rank,
|
|
dst_state_item_lens=dst_info.dst_state_item_lens,
|
|
dst_state_dim_per_tensor=dst_info.dst_state_dim_per_tensor,
|
|
)
|
|
if state_xfer_handle is not None:
|
|
handles.append(state_xfer_handle)
|
|
|
|
assert aux_index is not None
|
|
# When no KV pages were sent (decode-side cache hit),
|
|
# encode pp_rank in aux notif so receiver can mark
|
|
# expected_kvs_per_pp[pp_rank] = 0.
|
|
if len(kv_indices) == 0:
|
|
aux_notif = f"{req.room}_aux_nokv_{self.kv_args.pp_rank}"
|
|
else:
|
|
aux_notif = f"{req.room}_aux"
|
|
aux_xfer_handle = self.send_aux(
|
|
req.agent_name,
|
|
aux_index,
|
|
self.decode_kv_args_table[req.agent_name].dst_aux_ptrs,
|
|
req.dst_aux_index,
|
|
aux_notif,
|
|
)
|
|
handles.append(aux_xfer_handle)
|
|
if is_last:
|
|
del self.transfer_infos[bootstrap_room]
|
|
self.req_to_decode_prefix_len.pop(bootstrap_room, None)
|
|
return handles
|
|
|
|
def update_transfer_status(self):
|
|
# Process notifications from received transfers.
|
|
notif_map = self.agent.get_new_notifs()
|
|
for peer_name, messages in notif_map.items():
|
|
# We could also check that self.bootstrap_info['agent_name'] matches
|
|
# the message sender. But the bootstrap room alone should be
|
|
# sufficient to map the status.
|
|
for msg in messages:
|
|
components = msg.decode("ascii").split("_", 4)
|
|
room = int(components[0])
|
|
if components[1] == "kv":
|
|
chunk_id = int(components[2])
|
|
is_last = bool(int(components[3]))
|
|
pp_rank = int(components[4]) if len(components) > 4 else 0
|
|
# Track received chunks per pp_rank
|
|
self.transfer_statuses[room].received_kvs_per_pp[pp_rank].add(
|
|
chunk_id
|
|
)
|
|
if is_last:
|
|
# Record expected chunk count for this pp_rank
|
|
self.transfer_statuses[room].expected_kvs_per_pp[pp_rank] = (
|
|
chunk_id + 1
|
|
)
|
|
# Set num_pp_ranks_expected from table (or default to 1)
|
|
if self.transfer_statuses[room].num_pp_ranks_expected is None:
|
|
self.transfer_statuses[room].num_pp_ranks_expected = (
|
|
self.required_prefill_response_num_table.get(room, 1)
|
|
)
|
|
elif components[1] == "aux":
|
|
self.transfer_statuses[room].received_aux = True
|
|
# Handle "nokv" marker: no KV pages were sent for
|
|
# this pp_rank (decode-side radix cache hit).
|
|
if len(components) > 3 and components[2] == "nokv":
|
|
pp_rank = int(components[3])
|
|
self.transfer_statuses[room].expected_kvs_per_pp[pp_rank] = 0
|
|
if self.transfer_statuses[room].num_pp_ranks_expected is None:
|
|
self.transfer_statuses[room].num_pp_ranks_expected = (
|
|
self.required_prefill_response_num_table.get(room, 1)
|
|
)
|
|
elif components[1] == "state":
|
|
pp_rank = int(components[2]) if len(components) > 2 else 0
|
|
self.transfer_statuses[room].received_state_per_pp.add(pp_rank)
|
|
|
|
def check_transfer_done(self, room: int):
|
|
if room not in self.transfer_statuses:
|
|
return False
|
|
return self.transfer_statuses[room].is_done()
|
|
|
|
def _start_bootstrap_thread(self):
|
|
def bootstrap_thread():
|
|
"""This thread recvs transfer info from the decode engine"""
|
|
while True:
|
|
waiting_req_bytes = self.server_socket.recv_multipart()
|
|
logger.debug(
|
|
f"Received multipart with total byte size {sum(len(x) for x in waiting_req_bytes)}"
|
|
)
|
|
assert (
|
|
waiting_req_bytes[0] == GUARD
|
|
), f"First message should be {GUARD}. Foreign traffic?"
|
|
waiting_req_bytes = waiting_req_bytes[1:]
|
|
room = waiting_req_bytes[0].decode("ascii")
|
|
agent_name = waiting_req_bytes[3].decode("ascii")
|
|
if room == "None":
|
|
# Register new peer and save KV base pointers.
|
|
self._add_remote_peer(
|
|
KVArgsRegisterInfo.from_zmq(waiting_req_bytes)
|
|
)
|
|
logger.debug(f"Register KVArgs from {agent_name} successfully")
|
|
continue
|
|
room = int(room)
|
|
if room not in self.transfer_infos:
|
|
self.transfer_infos[room] = {}
|
|
self.transfer_infos[room][agent_name] = TransferInfo.from_zmq(
|
|
waiting_req_bytes
|
|
)
|
|
required_dst_info_num = self.transfer_infos[room][
|
|
agent_name
|
|
].required_dst_info_num
|
|
logger.debug(f"got info {room=} {agent_name=} {required_dst_info_num=}")
|
|
if len(self.transfer_infos[room]) == required_dst_info_num:
|
|
self.req_to_decode_prefix_len[room] = next(
|
|
(
|
|
info.decode_prefix_len
|
|
for info in self.transfer_infos[room].values()
|
|
if info.decode_prefix_len is not None
|
|
),
|
|
0,
|
|
)
|
|
logger.debug(f"{room=} is bootstrapped")
|
|
self.update_status(room, KVPoll.WaitingForInput)
|
|
|
|
threading.Thread(target=bootstrap_thread).start()
|
|
|
|
|
|
class NixlKVSender(CommonKVSender):
|
|
def __init__(
|
|
self,
|
|
mgr: NixlKVManager,
|
|
bootstrap_addr: str,
|
|
bootstrap_room: int,
|
|
dest_tp_ranks: List[int],
|
|
pp_rank: int,
|
|
):
|
|
super().__init__(mgr, bootstrap_addr, bootstrap_room, dest_tp_ranks, pp_rank)
|
|
self.xfer_handles = []
|
|
self.has_sent = False
|
|
self.chunk_id = 0
|
|
self._send_failed = False
|
|
self._send_error: Optional[Exception] = None
|
|
|
|
def pop_decode_prefix_len(self) -> int:
|
|
return self.kv_mgr.req_to_decode_prefix_len.pop(self.bootstrap_room, 0)
|
|
|
|
def should_send_kv_chunk(self, num_pages: int, last_chunk: bool) -> bool:
|
|
return num_pages > 0 or last_chunk
|
|
|
|
def send(
|
|
self,
|
|
kv_indices: npt.NDArray[np.int32],
|
|
state_indices: Optional[List[int]] = None,
|
|
):
|
|
if self._send_failed:
|
|
return
|
|
|
|
index_slice = slice(self.curr_idx, self.curr_idx + len(kv_indices))
|
|
self.curr_idx += len(kv_indices)
|
|
is_last = self.curr_idx == self.num_kv_indices
|
|
|
|
# Special handling for cp
|
|
if self.kv_mgr.enable_all_cp_ranks_for_transfer:
|
|
kv_indices, index_slice = filter_kv_indices_for_cp_rank(
|
|
self.kv_mgr,
|
|
kv_indices,
|
|
index_slice,
|
|
)
|
|
elif self.kv_mgr.is_dummy_cp_rank:
|
|
if not is_last:
|
|
return
|
|
else:
|
|
self.kv_mgr.update_status(self.bootstrap_room, KVPoll.Success)
|
|
return
|
|
|
|
try:
|
|
new_xfer_handles = self.kv_mgr.add_transfer_request(
|
|
self.bootstrap_room,
|
|
kv_indices,
|
|
index_slice,
|
|
is_last,
|
|
self.chunk_id,
|
|
self.aux_index,
|
|
state_indices,
|
|
)
|
|
except _NIXL_TRANSPORT_ERRORS as e:
|
|
logger.warning(
|
|
f"KVSender transfer request failed for room {self.bootstrap_room}: {e}"
|
|
)
|
|
self._send_failed = True
|
|
self._send_error = e
|
|
return
|
|
|
|
self.xfer_handles.extend(new_xfer_handles)
|
|
self.chunk_id += 1
|
|
if is_last:
|
|
self.has_sent = True
|
|
del self.kv_mgr.request_status[self.bootstrap_room]
|
|
|
|
def poll(self) -> KVPoll:
|
|
if self._send_failed:
|
|
return KVPoll.Failed # type: ignore
|
|
if not self.has_sent:
|
|
return self.kv_mgr.check_status(self.bootstrap_room)
|
|
try:
|
|
states = [self.kv_mgr.agent.check_xfer_state(x) for x in self.xfer_handles]
|
|
except _NIXL_TRANSPORT_ERRORS as e:
|
|
logger.warning(
|
|
f"KVSender check_xfer_state failed for room {self.bootstrap_room}: {e}"
|
|
)
|
|
self._send_failed = True
|
|
self._send_error = e
|
|
return KVPoll.Failed # type: ignore
|
|
if all(x == "DONE" for x in states):
|
|
return KVPoll.Success # type: ignore
|
|
if any(x == "ERR" for x in states):
|
|
self._send_failed = True
|
|
self._send_error = RuntimeError(
|
|
f"NIXL transfer error for room {self.bootstrap_room}"
|
|
)
|
|
return KVPoll.Failed # type: ignore
|
|
return KVPoll.WaitingForInput # type: ignore
|
|
|
|
def failure_exception(self):
|
|
if self._send_error is not None:
|
|
raise self._send_error
|
|
raise RuntimeError("NIXL KVSender Exception")
|
|
|
|
|
|
class NixlKVReceiver(CommonKVReceiver):
|
|
def __init__(
|
|
self,
|
|
mgr: NixlKVManager,
|
|
bootstrap_addr: str,
|
|
bootstrap_room: Optional[int] = None,
|
|
):
|
|
self.started_transfer = False
|
|
super().__init__(mgr, bootstrap_addr, bootstrap_room)
|
|
self.init_time = None
|
|
|
|
def init(
|
|
self,
|
|
prefill_dp_rank: int,
|
|
):
|
|
super().init(prefill_dp_rank)
|
|
|
|
def send_metadata(
|
|
self,
|
|
kv_indices: npt.NDArray[np.int32],
|
|
aux_index: Optional[int] = None,
|
|
state_indices: Optional[List[int]] = None,
|
|
decode_prefix_len: Optional[int] = None,
|
|
):
|
|
if self.bootstrap_infos is None:
|
|
logger.error(
|
|
f"Could not fetch prefill parallel info from bootstrap_addr: {self.bootstrap_addr}",
|
|
)
|
|
self.kv_mgr.update_status(self.bootstrap_room, KVPoll.Failed)
|
|
return
|
|
|
|
for bootstrap_info in self.bootstrap_infos:
|
|
logger.debug(
|
|
f"Fetched bootstrap info: {bootstrap_info} for engine rank: {self.kv_mgr.kv_args.engine_rank}"
|
|
)
|
|
sock, lock = self._connect_to_bootstrap_server(bootstrap_info)
|
|
is_dummy = bootstrap_info["is_dummy"]
|
|
logger.debug(
|
|
f"Sending to prefill server with bootstrap room {self.bootstrap_room} {is_dummy=}"
|
|
)
|
|
with lock:
|
|
sock.send_multipart(
|
|
[
|
|
GUARD,
|
|
str(self.bootstrap_room).encode("ascii"),
|
|
self.kv_mgr.local_ip.encode("ascii"),
|
|
str(self.kv_mgr.rank_port).encode("ascii"),
|
|
self.kv_mgr.agent.name.encode("ascii"),
|
|
kv_indices.tobytes() if not is_dummy else b"",
|
|
str(aux_index).encode("ascii"),
|
|
str(self.required_dst_info_num).encode("ascii"),
|
|
(
|
|
np.array(state_indices, dtype=np.int32).tobytes()
|
|
if not is_dummy and state_indices is not None
|
|
else b""
|
|
),
|
|
str(decode_prefix_len or 0).encode("ascii"),
|
|
]
|
|
)
|
|
|
|
# Mark that we expect state data if state_indices was provided
|
|
if state_indices is not None:
|
|
self.kv_mgr.transfer_statuses[self.bootstrap_room].expects_state = True
|
|
|
|
self.started_transfer = True
|
|
self.init_time = time.time()
|
|
|
|
def poll(self) -> KVPoll:
|
|
if self.conclude_state is not None:
|
|
return self.conclude_state
|
|
status = self.kv_mgr.check_status(self.bootstrap_room)
|
|
if status in (KVPoll.Success, KVPoll.Failed):
|
|
self.conclude_state = status
|
|
return status
|
|
if not self.started_transfer:
|
|
return status
|
|
|
|
now = time.time()
|
|
elapsed = now - self.init_time
|
|
|
|
if elapsed >= self.kv_mgr.waiting_timeout:
|
|
logger.error(f"Request {self.bootstrap_room} waiting_timeout")
|
|
self.kv_mgr.record_failure(
|
|
self.bootstrap_room,
|
|
f"Request {self.bootstrap_room} timed out after {elapsed:.1f}s in KVPoll.WaitingForInput",
|
|
)
|
|
self.conclude_state = KVPoll.Failed
|
|
return KVPoll.Failed
|
|
|
|
self.kv_mgr.update_transfer_status()
|
|
if self.kv_mgr.check_transfer_done(self.bootstrap_room): # type: ignore
|
|
self.kv_mgr.addr_to_rooms_tracker[self.bootstrap_addr].discard(
|
|
self.bootstrap_room
|
|
)
|
|
# Check if the transfer failed
|
|
if self.kv_mgr.transfer_statuses[self.bootstrap_room].is_failed():
|
|
self.conclude_state = KVPoll.Failed
|
|
logger.error(
|
|
f"Transfer for room {self.bootstrap_room} failed due to node failure"
|
|
)
|
|
else:
|
|
self.conclude_state = KVPoll.Success
|
|
del self.kv_mgr.transfer_statuses[self.bootstrap_room]
|
|
return self.conclude_state # type: ignore
|
|
return KVPoll.WaitingForInput # type: ignore
|
|
|
|
def _register_kv_args(self):
|
|
for bootstrap_info in self.bootstrap_infos:
|
|
sock, lock = self._connect_to_bootstrap_server(bootstrap_info)
|
|
packed_kv_data_ptrs = b"".join(
|
|
struct.pack("Q", ptr) for ptr in self.kv_mgr.kv_args.kv_data_ptrs
|
|
)
|
|
packed_aux_data_ptrs = b"".join(
|
|
struct.pack("Q", ptr) for ptr in self.kv_mgr.kv_args.aux_data_ptrs
|
|
)
|
|
packed_state_data_ptrs = b"".join(
|
|
struct.pack("Q", ptr) for ptr in self.kv_mgr.kv_args.state_data_ptrs
|
|
)
|
|
|
|
packed_state_item_lens = b"".join(
|
|
struct.pack("I", item_len)
|
|
for item_len in self.kv_mgr.kv_args.state_item_lens
|
|
)
|
|
state_dim_per_tensor = getattr(
|
|
self.kv_mgr.kv_args, "state_dim_per_tensor", []
|
|
)
|
|
packed_state_dim_per_tensor = b"".join(
|
|
struct.pack("I", dim) for dim in state_dim_per_tensor
|
|
)
|
|
|
|
with lock:
|
|
sock.send_multipart(
|
|
[
|
|
GUARD,
|
|
"None".encode("ascii"),
|
|
self.kv_mgr.local_ip.encode("ascii"),
|
|
str(self.kv_mgr.rank_port).encode("ascii"),
|
|
self.kv_mgr.agent.name.encode("ascii"),
|
|
self.kv_mgr.agent.get_agent_metadata(),
|
|
packed_kv_data_ptrs,
|
|
packed_aux_data_ptrs,
|
|
packed_state_data_ptrs,
|
|
str(self.kv_mgr.kv_args.gpu_id).encode("ascii"),
|
|
str(self.kv_mgr.attn_tp_size).encode("ascii"),
|
|
str(self.kv_mgr.kv_args.engine_rank).encode("ascii"),
|
|
str(self.kv_mgr.kv_args.kv_item_lens[0]).encode("ascii"),
|
|
packed_state_item_lens,
|
|
packed_state_dim_per_tensor,
|
|
]
|
|
)
|
|
|
|
def failure_exception(self):
|
|
raise RuntimeError("NIXL KVReceiver Exception")
|
|
|
|
|
|
class NixlKVBootstrapServer(CommonKVBootstrapServer):
|
|
pass
|