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":