From b9632954894d9fa334db6500c32d8ba73cb7603e Mon Sep 17 00:00:00 2001 From: "minke.yu" Date: Mon, 21 Sep 2026 11:53:26 +0800 Subject: [PATCH] mooncake: support host-memory payloads with intra-node NVLink transport 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 --- .../srt/disaggregation/mooncake/conn.py | 304 +++++++++++++++++- 1 file changed, 298 insertions(+), 6 deletions(-) diff --git a/python/sglang/srt/disaggregation/mooncake/conn.py b/python/sglang/srt/disaggregation/mooncake/conn.py index 25cbf4c80..7dbc59e65 100644 --- a/python/sglang/srt/disaggregation/mooncake/conn.py +++ b/python/sglang/srt/disaggregation/mooncake/conn.py @@ -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":