minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

mooncake: support host-memory payloads with intra-node NVLink transport

Browse Source 
IntraNodeNvlinkTransport can only register/reach device memory. sglang's
register list includes host regions (aux buffers, some state components);
one host region failed the whole registerLocalMemoryBatch and the engine
rolled back every region, leaving the segment descriptor empty so all KV
transfers failed with 'Requested address ... not found'.

When the intra-node NVLink transport is active (MC_INTRANODE_NVLINK etc.):
- register only device-memory regions (probe via cudaPointerGetAttributes)
- route transfer blocks whose local source is host memory over the ordered
  per-endpoint zmq channel (new STATE_DATA message) instead of the engine;
  the receiver validates and writes the block into its buffer, relying on
  the same ordering guarantee as the aux TCP path
- force aux over the TCP socket path automatically
This commit is contained in:
minke.yu
2026-09-21 11:53:26 +08:00
parent fa826e08b1
commit b963295489
1 changed files with 298 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/disaggregation/mooncake/conn.py
+298 -6
View File
@@ -82,6 +82,138 @@ FAILED_SESSION_RECOVERIES = Counter(
)
# ---------------------------------------------------------------------------
# Intra-node NVLink transport helpers.
#
# Mooncake's IntraNodeNvlinkTransport can only register and reach *device*
# memory (it IPC-opens the remote cudaMalloc segments). Host-resident regions
# (aux buffers, some state components) cannot be registered: one host region
# makes the whole registerLocalMemoryBatch fail, and the engine then rolls
# back *every* region, leaving the segment descriptor empty and all KV
# transfers failing with "Requested address ... not found". When the
# intra-node NVLink transport is active we therefore
# 1. register only device-memory regions, and
# 2. route blocks whose source is host memory over the ordered zmq channel
# (same ordering guarantee the aux TCP path relies on) instead of the
# transfer engine.
# ---------------------------------------------------------------------------
import ctypes as _ctypes
class _CudaPointerAttributes(_ctypes.Structure):
_fields_ = [
("type", _ctypes.c_int),
("device", _ctypes.c_int),
("devicePointer", _ctypes.c_void_p),
("hostPointer", _ctypes.c_void_p),
]
_CUDA_MEMORY_TYPE_DEVICE = 2
_CUDART = None
def _get_cudart():
global _CUDART
if _CUDART is None:
for name in ("libcudart.so", "libcudart.so.13", "libcudart.so.12"):
try:
_CUDART = _ctypes.CDLL(name)
break
except OSError:
continue
else:
_CUDART = False
return _CUDART or None
def _is_device_pointer(ptr: int) -> bool:
"""Probe a *local* pointer with cudaPointerGetAttributes.
Only valid for pointers owned by this process (never probe remote
segment addresses). Returns False on any error so the caller falls back
to the safe host path.
"""
cudart = _get_cudart()
if cudart is None:
# Cannot tell; assume device so behavior stays unchanged.
return True
attr = _CudaPointerAttributes()
ret = cudart.cudaPointerGetAttributes(
_ctypes.byref(attr), _ctypes.c_void_p(int(ptr))
)
if ret != 0:
# Clear the error so subsequent CUDA calls are not poisoned.
cudart.cudaGetLastError()
return False
return attr.type == _CUDA_MEMORY_TYPE_DEVICE
def _read_bytes_from_address(addr: int, length: int) -> Optional[bytes]:
if length <= 0:
return b""
if _is_device_pointer(addr):
cudart = _get_cudart()
buf = (_ctypes.c_char * length)()
# cudaMemcpyDeviceToHost = 2; synchronous default-stream copy.
ret = cudart.cudaMemcpy(
buf, _ctypes.c_void_p(int(addr)), _ctypes.c_size_t(length), 2
)
if ret != 0:
logger.error(
f"cudaMemcpy D2H failed (ret={ret}) for addr {hex(addr)} len {length}"
)
return None
return bytes(buf)
return _ctypes.string_at(int(addr), length)
def _write_bytes_to_address(addr: int, data: bytes) -> bool:
if not data:
return True
if _is_device_pointer(addr):
cudart = _get_cudart()
buf = _ctypes.create_string_buffer(data, len(data))
# cudaMemcpyHostToDevice = 1; synchronous default-stream copy.
ret = cudart.cudaMemcpy(
_ctypes.c_void_p(int(addr)), buf, _ctypes.c_size_t(len(data)), 1
)
if ret != 0:
logger.error(
f"cudaMemcpy H2D failed (ret={ret}) for addr {hex(addr)} "
f"len {len(data)}"
)
return False
return True
_ctypes.memmove(int(addr), data, len(data))
return True
_NVLINK_INTRA_ACTIVE = None
def _nvlink_intra_transport_active() -> bool:
"""Whether mooncake installed the intra-node NVLink transport.
Mirrors the env probing in mooncake's transfer_engine_impl.cpp: the
transport is installed iff MC_INTRANODE_NVLINK is set (any value), or an
equivalent protocol selection was made.
"""
global _NVLINK_INTRA_ACTIVE
if _NVLINK_INTRA_ACTIVE is None:
active = bool(
os.environ.get("MC_INTRANODE_NVLINK")
or os.environ.get("MC_INTRA_NVLINK")
)
if not active:
proto = (os.environ.get("MOONCAKE_PROTOCOL") or "").strip().lower()
active = proto in ("nvlink_intra", "nvlink-intra", "intra_nvlink")
_NVLINK_INTRA_ACTIVE = active
return _NVLINK_INTRA_ACTIVE
# decode
@dataclasses.dataclass
class TransferInfo:
@@ -214,6 +346,7 @@ class KVArgsRegisterInfo:
class MooncakeKVManager(StagingManagerMixin, CommonKVManager):
AUX_DATA_HEADER = b"AUX_DATA"
STATE_DATA_HEADER = b"STATE_DATA"
# Implements teardown() below, so runtime PD role switching is supported.
supports_role_switch = True
@@ -227,6 +360,10 @@ class MooncakeKVManager(StagingManagerMixin, CommonKVManager):
super().__init__(args, disaggregation_mode, server_args, is_mla_backend)
self.init_engine()
self.register_buffer_to_engine()
# session_id -> (endpoint, dst_port, room), used to route host-memory
# transfer blocks over zmq when the intra-node NVLink transport is
# active (it cannot reach host memory). Populated on bootstrap.
self._session_endpoint_map = {}
self.enable_staging = envs.SGLANG_DISAGG_STAGING_BUFFER.get()
self.max_transfer_batch_indices = (
envs.SGLANG_MOONCAKE_MAX_TRANSFER_BATCH_INDICES.get()
@@ -322,6 +459,13 @@ class MooncakeKVManager(StagingManagerMixin, CommonKVManager):
Deduped because the unified memory pool reports one raw buffer as both
its KV and its mamba state component, and double registration fails in
the engine.
When the intra-node NVLink transport is active, host-memory regions
(aux buffers, some state components) are skipped: the transport only
accepts device memory, and a single host region fails the whole batch
and triggers a full engine-side rollback that would unregister the KV
pools too. Host-resident payloads are instead exchanged over the
ordered zmq channel (see _transfer_data / send_aux).
"""
regions: List[Tuple[int, int]] = []
seen: Set[Tuple[int, int]] = set()
@@ -338,6 +482,24 @@ class MooncakeKVManager(StagingManagerMixin, CommonKVManager):
self.kv_args.state_data_ptrs, self.kv_args.state_data_lens
):
add(ptrs, lens)
if _nvlink_intra_transport_active():
device_regions = []
skipped = []
for ptr, length in regions:
if _is_device_pointer(ptr):
device_regions.append((ptr, length))
else:
skipped.append((ptr, length))
if skipped:
logger.info(
"Intra-node NVLink transport: skipping %d host-memory "
"regions from engine registration (they will be exchanged "
"over the zmq channel instead): %s",
len(skipped),
[(hex(p), l) for p, l in skipped[:8]],
)
regions = device_regions
return regions
def register_buffer_to_engine(self):
@@ -748,10 +910,63 @@ class MooncakeKVManager(StagingManagerMixin, CommonKVManager):
if not transfer_blocks:
return 0
src_addrs, dst_addrs, lengths = zip(*transfer_blocks)
return self.engine.batch_transfer_sync(
mooncake_session_id, list(src_addrs), list(dst_addrs), list(lengths)
)
if not _nvlink_intra_transport_active():
src_addrs, dst_addrs, lengths = zip(*transfer_blocks)
return self.engine.batch_transfer_sync(
mooncake_session_id, list(src_addrs), list(dst_addrs), list(lengths)
)
# Intra-node NVLink transport can only move device memory. Partition
# blocks by the *local source* pointer (probing a local pointer is
# safe; the remote dst is never probed): device-sourced blocks go
# through the engine as usual, host-sourced blocks are shipped over
# the ordered zmq channel and written into the peer's buffer by the
# receiver (see _handle_state_data). This mirrors the aux TCP path.
device_blocks = []
host_blocks = []
for src, dst, length in transfer_blocks:
if _is_device_pointer(src):
device_blocks.append((src, dst, length))
else:
host_blocks.append((src, dst, length))
rc = 0
if device_blocks:
src_addrs, dst_addrs, lengths = zip(*device_blocks)
rc = self.engine.batch_transfer_sync(
mooncake_session_id, list(src_addrs), list(dst_addrs), list(lengths)
)
if rc == 0 and host_blocks:
rc = self._send_host_blocks_tcp(mooncake_session_id, host_blocks)
return rc
def _send_host_blocks_tcp(self, mooncake_session_id, host_blocks):
target = self._session_endpoint_map.get(mooncake_session_id)
if target is None:
logger.error(
f"No zmq endpoint known for mooncake session "
f"{mooncake_session_id}; cannot deliver {len(host_blocks)} "
"host-memory transfer blocks"
)
return -1
endpoint, dst_port, room = target
na = NetworkAddress(endpoint, dst_port)
for src, dst, length in host_blocks:
data = _read_bytes_from_address(src, length)
if data is None:
return -1
self._send_multipart_locked(
na.to_tcp(),
[
MooncakeKVManager.STATE_DATA_HEADER,
str(room).encode("ascii"),
str(int(dst)).encode("ascii"),
struct.pack(">I", len(data)),
data,
],
is_ipv6=na.is_ipv6,
)
return 0
def _send_kvcache_generic(
self,
@@ -1444,8 +1659,10 @@ class MooncakeKVManager(StagingManagerMixin, CommonKVManager):
):
# TODO(shangming): Fix me when nvlink_transport of Mooncake is bug-free
if (
self.enable_custom_mem_pool and self.custom_mem_pool_type == "NVLINK"
) or envs.SGLANG_MOONCAKE_SEND_AUX_TCP.get():
(self.enable_custom_mem_pool and self.custom_mem_pool_type == "NVLINK")
or envs.SGLANG_MOONCAKE_SEND_AUX_TCP.get()
or _nvlink_intra_transport_active()
):
return self.send_aux_tcp(req, prefill_aux_index, dst_aux_ptrs)
transfer_blocks = []
@@ -1528,6 +1745,71 @@ class MooncakeKVManager(StagingManagerMixin, CommonKVManager):
f"Received AUX_DATA for bootstrap_room {room} with length:{len(data)}"
)
def _host_transfer_regions(self):
"""Address ranges this process published as transfer targets.
Used to validate STATE_DATA writes. Built lazily because kv_args is
fully populated only after registration.
"""
regions = getattr(self, "_host_transfer_regions_cache", None)
if regions is None:
regions = []
for ptr, length in zip(
self.kv_args.kv_data_ptrs or [], self.kv_args.kv_data_lens or []
):
regions.append((int(ptr), int(ptr) + int(length)))
for ptr, length in zip(
self.kv_args.aux_data_ptrs or [], self.kv_args.aux_data_lens or []
):
regions.append((int(ptr), int(ptr) + int(length)))
for ptrs, lens in zip(
self.kv_args.state_data_ptrs or [], self.kv_args.state_data_lens or []
):
for ptr, length in zip(ptrs or [], lens or []):
regions.append((int(ptr), int(ptr) + int(length)))
self._host_transfer_regions_cache = regions
return regions
def _handle_state_data(self, msg: List[bytes]):
"""Handle STATE_DATA messages received by the decode thread.
Carries one host-memory transfer block that could not go through the
intra-node NVLink transport. Written directly into the local buffer at
the destination address; ordering against the final status message is
guaranteed by the shared per-endpoint zmq socket.
"""
room = int(msg[1].decode("ascii"))
dst_addr = int(msg[2].decode("ascii"))
data_length = struct.unpack(">I", msg[3])[0]
data = msg[4]
if len(data) != data_length:
logger.error(f"STATE_DATA length mismatch for bootstrap_room {room}")
return
in_region = any(
start <= dst_addr and dst_addr + len(data) <= end
for start, end in self._host_transfer_regions()
)
if not in_region:
logger.error(
f"STATE_DATA for bootstrap_room {room} targets unknown region "
f"{hex(dst_addr)}..{hex(dst_addr + len(data))}; dropping"
)
return
if not _write_bytes_to_address(dst_addr, data):
logger.error(
f"STATE_DATA write failed for bootstrap_room {room} at "
f"{hex(dst_addr)} len {len(data)}"
)
return
logger.debug(
f"Received STATE_DATA for bootstrap_room {room} at {hex(dst_addr)} "
f"with length:{len(data)}"
)
def _get_dsa_cache_transfer_skip_flags(
self, info: Optional[KVArgsRegisterInfo]
) -> Tuple[bool, bool]:
@@ -2369,6 +2651,8 @@ class MooncakeKVManager(StagingManagerMixin, CommonKVManager):
):
self._staging_outstanding.pop(kv_chunk.room, None)
if kv_chunk.room in self.transfer_infos:
for sid in self.transfer_infos[kv_chunk.room]:
self._session_endpoint_map.pop(sid, None)
self.transfer_infos.pop(kv_chunk.room)
self.req_to_decode_prefix_len.pop(kv_chunk.room, None)
if self.enable_staging:
@@ -2512,6 +2796,11 @@ class MooncakeKVManager(StagingManagerMixin, CommonKVManager):
self.transfer_infos[room][mooncake_session_id] = (
TransferInfo.from_zmq(waiting_req_bytes)
)
self._session_endpoint_map[mooncake_session_id] = (
self.transfer_infos[room][mooncake_session_id].endpoint,
self.transfer_infos[room][mooncake_session_id].dst_port,
room,
)
# NOTE: after bootstrapping we can mark the req as waiting for input
if len(self.transfer_infos[room]) == required_dst_info_num:
self.resolve_kv_replica_factor(self.transfer_infos[room])
@@ -2540,6 +2829,9 @@ class MooncakeKVManager(StagingManagerMixin, CommonKVManager):
if msg[0] == MooncakeKVManager.AUX_DATA_HEADER:
self._handle_aux_data(msg)
continue
if msg[0] == MooncakeKVManager.STATE_DATA_HEADER:
self._handle_state_data(msg)
continue
# Staging: prefill notifies a chunk written to staging buffer
if msg[0] == b"CHUNK_READY":