diff --git a/atom/kv_transfer/disaggregation/factory.py b/atom/kv_transfer/disaggregation/factory.py index e3562166e3..3e49b7b3da 100644 --- a/atom/kv_transfer/disaggregation/factory.py +++ b/atom/kv_transfer/disaggregation/factory.py @@ -135,6 +135,14 @@ def create_connector( scheduler_class="MooncakeConnectorScheduler", ) +KVConnectorFactory.register( + "native", + worker_module="atom.kv_transfer.disaggregation.native.native_connector", + worker_class="NativeConnector", + scheduler_module="atom.kv_transfer.disaggregation.native.native_connector", + scheduler_class="NativeConnectorScheduler", +) + # Composite backend: fans out to several sub-connectors listed under # kv_transfer_config["connectors"] (e.g. moriio P/D + lmcache_offload on one # prefill node). Lightweight import — no heavy deps until a sub is built. diff --git a/atom/kv_transfer/disaggregation/native/README.md b/atom/kv_transfer/disaggregation/native/README.md new file mode 100644 index 0000000000..01b0a6c5ea --- /dev/null +++ b/atom/kv_transfer/disaggregation/native/README.md @@ -0,0 +1,79 @@ +# Native KV connector (`kv_connector="native"`) + +A fully in-tree prefill/decode (P/D) KV-cache connector for the **single-node +(scale-up / XGMI)** case. It depends only on the HIP Virtual Memory Management +(VMM) API — **no third-party transport** (no MoRI, no Mooncake). + +## When to use it + +| Scenario | Connector | +|---|---| +| Single node, P/D split across GPUs on the same box (XGMI fabric) | **`native`** | +| Cross-node P/D (RDMA NICs) | `moriio` | + +The `native` connector moves KV directly GPU→GPU over the fabric with +`hipMemcpy` peer copies (via HIP VMM shareable handles). It requires GPUs with +VMM support (queried per device); it raises a clear error otherwise. + +## How to select it + +Pass `"kv_connector": "native"` in `--kv-transfer-config`. There is **no +`protocol` field** (the transport is always single-node XGMI), and GPUs use +**natural placement** (`device == rank`) — no visibility reorder is needed. + +## Launch (4 prefill GPUs + 4 decode GPUs on one node) + +```bash +# 1) proxy +python -m atom.kv_transfer.disaggregation.proxy --port 10001 + +# 2) prefill engine (producer) on GPUs 0-3 +HIP_VISIBLE_DEVICES=0,1,2,3 python -m atom.entrypoints.openai_server \ + --model deepseek-ai/DeepSeek-V4-Pro --kv_cache_dtype bf16 -tp 4 \ + --gpu-memory-utilization 0.85 --max-num-seqs 128 \ + --host 0.0.0.0 --server-port 8003 \ + --kv-transfer-config '{"kv_connector":"native","kv_role":"kv_producer","proxy_ip":"127.0.0.1","proxy_ping_port":36367,"http_port":8003,"handshake_port":6501}' + +# 3) decode engine (consumer) on GPUs 4-7 +HIP_VISIBLE_DEVICES=4,5,6,7 python -m atom.entrypoints.openai_server \ + --model deepseek-ai/DeepSeek-V4-Pro --kv_cache_dtype bf16 -tp 4 \ + --gpu-memory-utilization 0.85 --max-num-seqs 128 \ + --host 0.0.0.0 --server-port 8004 \ + --kv-transfer-config '{"kv_connector":"native","kv_role":"kv_consumer","proxy_ip":"127.0.0.1","proxy_ping_port":36367,"http_port":8004,"handshake_port":6501}' + +# 4) send requests to the proxy +curl -s http://127.0.0.1:10001/v1/completions \ + -H "Content-Type: application/json" \ + -d '{"model":"deepseek-ai/DeepSeek-V4-Pro","prompt":"The capital of France is","max_tokens":16}' +``` + +`kv_transfer_params.do_remote_prefill` in the response should be `true`. + +## `kv-transfer-config` fields + +| key | meaning | +|---|---| +| `kv_connector` | `"native"` | +| `kv_role` | `"kv_producer"` (prefill) or `"kv_consumer"` (decode) | +| `proxy_ip`, `proxy_ping_port` | proxy address for registration | +| `http_port` | this engine's OpenAI server port | +| `handshake_port` | base port for the UNIX side channel (per-rank offset added) | + +## How it works + +- Each worker allocates a VMM **staging** buffer and (consumer) exports its + POSIX fd over a UNIX side channel (`SCM_RIGHTS`). +- The consumer sends its destination block ids + staging fd to the producer; + the producer imports the staging (granting its own device access), gathers + the request's KV blocks straight into it over XGMI, and replies `WRITE_DONE`. +- The consumer scatters from its staging into its local KV pool. +- One fd import per (producer, consumer) pair; subsequent transfers are direct + device-to-device copies — no RDMA, no IPC-handle churn, no host staging. + +## Status + +v1 wires the VMM transport primitive (validated cross-process by +`tests/test_native_vmm_transfer.py`) into the KVConnector interface. Requests +within a scheduler step are transferred sequentially; a concurrent staging pool +and DeepSeek-V4 slot/index-region fast paths are follow-ups. See ROCm/ATOM#1483 +for end-to-end 4P4D validation. diff --git a/atom/kv_transfer/disaggregation/native/__init__.py b/atom/kv_transfer/disaggregation/native/__init__.py new file mode 100644 index 0000000000..e7abed63bd --- /dev/null +++ b/atom/kv_transfer/disaggregation/native/__init__.py @@ -0,0 +1,8 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2025, Advanced Micro Devices, Inc. All rights reserved. + +"""Single-node (XGMI) scale-up KV-transfer primitives.""" + +from atom.kv_transfer.disaggregation.native.vmm import VmmBuffer, supported + +__all__ = ["VmmBuffer", "supported"] diff --git a/atom/kv_transfer/disaggregation/native/_vmm_ext.cpp b/atom/kv_transfer/disaggregation/native/_vmm_ext.cpp new file mode 100644 index 0000000000..76c2e5e2b9 --- /dev/null +++ b/atom/kv_transfer/disaggregation/native/_vmm_ext.cpp @@ -0,0 +1,159 @@ +// SPDX-License-Identifier: MIT +// Copyright (C) 2024-2025, Advanced Micro Devices, Inc. All rights reserved. +// +// Minimal HIP VMM bridge for cross-process single-node (XGMI) KV sharing: +// allocate an exportable VMM buffer, export/import its POSIX fd, grant peer +// access, and copy. Reliable cross-process where legacy hipIpc is not, and does +// not require the source GPU to be in the consumer's visible set. +#include +#include + +#include +#include +#include + +#define HIPCK(expr) \ + do { \ + hipError_t _e = (expr); \ + if (_e != hipSuccess) \ + throw std::runtime_error(std::string(#expr) + ": " + \ + hipGetErrorString(_e)); \ + } while (0) + +namespace { + +struct Region { + hipMemGenericAllocationHandle_t handle; + void *ptr; + size_t size; +}; + +std::unordered_map g_regions; +int64_t g_next_id = 0; + +hipMemAllocationProp make_prop(int device) { + hipMemAllocationProp prop{}; + prop.type = hipMemAllocationTypePinned; + prop.location.type = hipMemLocationTypeDevice; + prop.location.id = device; + prop.requestedHandleType = hipMemHandleTypePosixFileDescriptor; + return prop; +} + +size_t round_up_to_granularity(size_t nbytes, int device) { + hipMemAllocationProp prop = make_prop(device); + size_t gran = 0; + HIPCK(hipMemGetAllocationGranularity(&gran, &prop, + hipMemAllocationGranularityMinimum)); + return ((nbytes + gran - 1) / gran) * gran; +} + +// Grant `device` read/write peer access to a mapped range. +void grant_access(void *ptr, size_t size, int device) { + hipMemAccessDesc desc{}; + desc.location.type = hipMemLocationTypeDevice; + desc.location.id = device; + desc.flags = hipMemAccessFlagsProtReadWrite; + HIPCK(hipMemSetAccess(ptr, size, &desc, 1)); +} + +Region ®ion(int64_t id) { + auto it = g_regions.find(id); + if (it == g_regions.end()) + throw std::runtime_error("unknown VMM region id " + std::to_string(id)); + return it->second; +} + +} // namespace + +bool vmm_supported(int device) { + int value = 0; + hipDeviceGetAttribute( + &value, hipDeviceAttributeVirtualMemoryManagementSupported, device); + return value != 0; +} + +// Allocate an exportable VMM buffer on `device`, map it and grant `device` +// access. Returns an opaque region id. +int64_t vmm_alloc(int64_t nbytes, int device) { + HIPCK(hipSetDevice(device)); + size_t size = round_up_to_granularity(static_cast(nbytes), device); + hipMemAllocationProp prop = make_prop(device); + Region r{}; + r.size = size; + HIPCK(hipMemCreate(&r.handle, size, &prop, 0)); + HIPCK(hipMemAddressReserve(&r.ptr, size, 0, 0, 0)); + HIPCK(hipMemMap(r.ptr, size, 0, r.handle, 0)); + grant_access(r.ptr, size, device); + int64_t id = g_next_id++; + g_regions[id] = r; + return id; +} + +// Export the region's handle as a POSIX file descriptor (to send over a UNIX +// socket via SCM_RIGHTS). +int vmm_export_fd(int64_t id) { + int fd = -1; + HIPCK(hipMemExportToShareableHandle( + &fd, region(id).handle, hipMemHandleTypePosixFileDescriptor, 0)); + return fd; +} + +// Import a peer's fd, map it on `device` and grant `device` peer access. +// `nbytes` must match the producer's requested size (rounded identically). +int64_t vmm_import(int fd, int64_t nbytes, int device) { + HIPCK(hipSetDevice(device)); + size_t size = round_up_to_granularity(static_cast(nbytes), device); + Region r{}; + r.size = size; + HIPCK(hipMemImportFromShareableHandle( + &r.handle, reinterpret_cast(static_cast(fd)), + hipMemHandleTypePosixFileDescriptor)); + HIPCK(hipMemAddressReserve(&r.ptr, size, 0, 0, 0)); + HIPCK(hipMemMap(r.ptr, size, 0, r.handle, 0)); + grant_access(r.ptr, size, device); + int64_t id = g_next_id++; + g_regions[id] = r; + return id; +} + +// Wrap the first `nbytes` of the mapped region as a (non-owning) uint8 tensor. +int64_t vmm_ptr(int64_t id) { + return reinterpret_cast(region(id).ptr); +} + +torch::Tensor vmm_tensor(int64_t id, int64_t nbytes, int device) { + auto opts = torch::TensorOptions().dtype(torch::kUInt8).device( + torch::kCUDA, device); + return torch::from_blob(region(id).ptr, {nbytes}, opts); +} + +// Device-to-device copy between two raw device pointers (peer-mapped ok). Used +// by the connector to gather/scatter KV blocks to/from the VMM staging region. +void vmm_copy(int64_t dst_ptr, int64_t src_ptr, int64_t nbytes) { + HIPCK(hipMemcpy(reinterpret_cast(dst_ptr), + reinterpret_cast(src_ptr), + static_cast(nbytes), hipMemcpyDeviceToDevice)); +} + +void vmm_free(int64_t id) { + auto it = g_regions.find(id); + if (it == g_regions.end()) + return; + Region &r = it->second; + hipMemUnmap(r.ptr, r.size); + hipMemAddressFree(r.ptr, r.size); + hipMemRelease(r.handle); + g_regions.erase(it); +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("vmm_supported", &vmm_supported); + m.def("vmm_alloc", &vmm_alloc); + m.def("vmm_export_fd", &vmm_export_fd); + m.def("vmm_import", &vmm_import); + m.def("vmm_ptr", &vmm_ptr); + m.def("vmm_tensor", &vmm_tensor); + m.def("vmm_copy", &vmm_copy); + m.def("vmm_free", &vmm_free); +} diff --git a/atom/kv_transfer/disaggregation/native/native_connector.py b/atom/kv_transfer/disaggregation/native/native_connector.py new file mode 100644 index 0000000000..0a4fd74fce --- /dev/null +++ b/atom/kv_transfer/disaggregation/native/native_connector.py @@ -0,0 +1,382 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2025, Advanced Micro Devices, Inc. All rights reserved. + +"""Native single-node KV connector (HIP VMM; no third-party transport). + +Single-node (XGMI) prefill/decode connector, selected with +``kv_connector="native"``. Push path: the consumer sends a VMM staging fd over a +UNIX socket; the producer imports it, gathers the request's KV blocks + SWA +slots + compressor state into it (hipMemcpy peer over XGMI); the consumer +scatters into its KV pool. +""" + +from __future__ import annotations + +import logging +import os +import socket +import threading +import time +from typing import Any + +import msgpack +import torch +import zmq + +from aiter.dist.parallel_state import get_dp_group, get_tp_group +from atom.config import Config +from atom.kv_transfer.disaggregation.base import ( + KVConnectorBase, + KVConnectorSchedulerBase, +) +from atom.kv_transfer.disaggregation.native import vmm +from atom.kv_transfer.disaggregation.types import ConnectorMetadata, ReqMeta +from atom.utils.network import get_ip + +logger = logging.getLogger("atom") + +ReqId = str +TransferId = int + +_MSG_WRITE_REQUEST = b"\x01" +_MSG_WRITE_DONE = b"\x02" +_PREFILL_WAIT_S = 30.0 + + +def _port_offset(dp_rank: int, tp_rank: int, tp_size: int = 1) -> int: + return dp_rank * tp_size + tp_rank + + +def _sock_path(port: int) -> str: + return f"/tmp/atom_native_p_{port}.sock" + + +# --------------------------------------------------------------------------- +# Scheduler (transport-agnostic). +# --------------------------------------------------------------------------- +class NativeConnectorScheduler(KVConnectorSchedulerBase): + def __init__(self, config: Config) -> None: + self.config = config + kv_cfg = config.kv_transfer_config or {} + self.is_producer = kv_cfg.get("kv_role", "kv_producer") == "kv_producer" + self._reqs_to_save: dict[ReqId, ReqMeta] = {} + self._reqs_to_recv: dict[ReqId, ReqMeta] = {} + self.request_id_to_transfer_id: dict[ReqId, TransferId] = {} + self.transfer_id_to_request_id: dict[TransferId, ReqId] = {} + + def get_num_new_matched_tokens(self, seq: Any) -> tuple[int, bool]: + params = seq.kv_transfer_params or {} + if params.get("do_remote_prefill") and not self.is_producer: + return 0, True + return 0, False + + def update_state_after_alloc(self, seq: Any) -> None: + params = seq.kv_transfer_params or {} + tid = params.get("transfer_id") + if tid is not None: + self.transfer_id_to_request_id[tid] = seq.id + self.request_id_to_transfer_id[seq.id] = tid + slot_idx = getattr(seq, "per_req_cache_group", getattr(seq, "slot_index", -1)) + params["local_slot_index"] = slot_idx + meta = ReqMeta( + local_block_ids=list(getattr(seq, "block_ids", []) or []), + remote_block_ids=params.get("remote_block_ids") or [], + remote_host=params.get("remote_host", ""), + remote_port=params.get("remote_port", 0), + remote_handshake_port=params.get("remote_handshake_port", 0), + remote_engine_id=params.get("remote_engine_id", ""), + tp_size=params.get("tp_size", 1), + remote_dp_size=params.get("remote_dp_size", 1), + remote_dp_rank=params.get("remote_dp_rank", 0), + transfer_id=params.get("transfer_id", 0), + local_slot_index=slot_idx, + ) + if params.get("do_remote_prefill"): + assert not self.is_producer + params["do_remote_prefill"] = False + self._reqs_to_recv[seq.id] = meta + elif params.get("do_remote_decode"): + assert self.is_producer + # The transfer handle the consumer will request is the producer's + # own request id (see request_finished), not params["transfer_id"]. + meta.transfer_id = seq.id + self._reqs_to_save[seq.id] = meta + + def build_connector_meta(self) -> ConnectorMetadata: + meta = ConnectorMetadata() + meta.request_id_to_transfer_id = dict(self.request_id_to_transfer_id) + meta.reqs_to_save = dict(self._reqs_to_save) + meta.reqs_to_recv = dict(self._reqs_to_recv) + self._reqs_to_save.clear() + self._reqs_to_recv.clear() + return meta + + def request_finished(self, seq: Any) -> None: + if self.is_producer: + seq.kv_transfer_params_output = { + "do_remote_prefill": True, + "do_remote_decode": False, + "transfer_id": seq.id, + } + tid = self.request_id_to_transfer_id.pop(seq.id, None) + if tid is not None: + self.transfer_id_to_request_id.pop(tid, None) + + +# --------------------------------------------------------------------------- +# Worker (VMM transport). +# --------------------------------------------------------------------------- +class NativeConnector(KVConnectorBase): + def __init__(self, config: Config) -> None: + self.config = config + kv_cfg = config.kv_transfer_config or {} + self.is_producer = kv_cfg.get("kv_role", "kv_producer") == "kv_producer" + self.device = torch.cuda.current_device() + if not vmm.supported(self.device): + raise RuntimeError( + "kv_connector='native' requires HIP VMM (single-node scale-up); " + "use 'moriio' for cross-node RDMA." + ) + self.tp_rank = get_tp_group().rank_in_group + self.tp_size = get_tp_group().world_size + self.dp_rank = get_dp_group().rank_in_group + self.dp_size = get_dp_group().world_size + self.local_ip = get_ip() + self.http_port = kv_cfg.get("http_port", 8000) + self.request_address = f"{self.local_ip}:{self.http_port}" + self.proxy_ip = kv_cfg.get("proxy_ip") + self.proxy_ping_port = kv_cfg.get("proxy_ping_port", 36367) + self.base_handshake_port = kv_cfg.get("handshake_port", 6501) + self._port = self.base_handshake_port + _port_offset( + self.dp_rank, self.tp_rank, self.tp_size + ) + + # KV region layout (filled by register_kv_caches). + self._block_regions: list[tuple[int, int]] = [] # (base, bytes/block) + self._slot_regions: list[tuple[int, int]] = [] # (base, bytes/slot) + self._state_base = 0 + self._state_slot_bytes = 0 + self._state_pool_free: list[int] = [] + self._gather_slot = None + self._scatter_slot = None + + self._lock = threading.Lock() + self.done_sending: set[ReqId] = set() + self.done_recving: set[ReqId] = set() + # producer: transfer_id -> (src_block_ids, src_slot) + self._prefills: dict[TransferId, tuple[list[int], int]] = {} + self._prefills_cv = threading.Condition(self._lock) + self._imported: dict[int, vmm.VmmBuffer] = {} + + self._zmq = zmq.Context() + if self.tp_rank == 0 and self.dp_rank == 0 and self.proxy_ip: + threading.Thread(target=self._ping, daemon=True).start() + + # -- proxy service discovery ------------------------------------------- + + def _ping(self) -> None: + path = f"tcp://{self.proxy_ip}:{self.proxy_ping_port}" + role = "P" if self.is_producer else "D" + with self._zmq.socket(zmq.DEALER) as sock: + sock.connect(path) + i = 1 + while True: + try: + sock.send( + msgpack.dumps( + { + "type": "register", + "role": role, + "index": str(i), + "request_address": f"http://{self.request_address}/v1/completions", + "rpc_port": self._port, + "handshake_port": self.base_handshake_port, + "dp_size": self.dp_size, + "tp_size": self.tp_size, + "transfer_mode": "write", + } + ) + ) + i += 1 + except Exception: + pass + time.sleep(5.0) + + # -- KVConnectorBase ---------------------------------------------------- + + def register_kv_caches(self, kv_caches, transfer_tensors=None) -> None: + tt = transfer_tensors + if tt is None: + raise RuntimeError("native connector requires KV transfer tensors") + self._block_regions = [(r.base_addr, r.unit_bytes) for r in tt.block_regions] + self._slot_regions = [(r.base_addr, r.unit_bytes) for r in tt.slot_regions] + self._gather_slot = tt.gather_slot + self._scatter_slot = tt.scatter_slot + if tt.staging_region is not None: + self._state_base = tt.staging_region.base_addr + self._state_slot_bytes = tt.staging_region.unit_bytes + self._state_pool_free = list(range(tt.staging_pool_size)) + logger.info( + "[native] registered %d block + %d slot regions, state_slot=%dB " + "(role=%s dev=%d rank=%d)", + len(self._block_regions), + len(self._slot_regions), + self._state_slot_bytes, + "producer" if self.is_producer else "consumer", + self.device, + self.tp_rank, + ) + if self.is_producer: + threading.Thread(target=self._serve, daemon=True).start() + + def start_load_kv(self, metadata: ConnectorMetadata) -> None: + for _, meta in metadata.reqs_to_save.items(): + with self._prefills_cv: + self._prefills[meta.transfer_id] = ( + meta.local_block_ids, + meta.local_slot_index, + ) + self._prefills_cv.notify_all() + for req_id, meta in metadata.reqs_to_recv.items(): + self._recv_request(req_id, meta) + + def get_finished(self) -> tuple[set, set]: + with self._lock: + ds, dr = set(self.done_sending), set(self.done_recving) + self.done_sending.clear() + self.done_recving.clear() + return ds, dr + + def get_finished_recv_blocks(self) -> list[int]: + return [] + + # -- staging layout ----------------------------------------------------- + + def _req_bytes(self, nblocks: int) -> int: + b = sum(bpb for _, bpb in self._block_regions) * nblocks + b += sum(bps for _, bps in self._slot_regions) + b += self._state_slot_bytes + return b + + def _acquire_state_slot(self) -> int: + with self._lock: + return self._state_pool_free.pop() if self._state_pool_free else -1 + + def _release_state_slot(self, idx: int) -> None: + if idx >= 0: + with self._lock: + self._state_pool_free.append(idx) + + # -- consumer ----------------------------------------------------------- + + def _recv_request(self, req_id: ReqId, meta: ReqMeta) -> None: + nblocks = len(meta.local_block_ids) + staging = vmm.VmmBuffer.alloc(self._req_bytes(nblocks), self.device) + payload = msgpack.dumps( + { + "req_id": req_id, + "transfer_id": meta.transfer_id, + "dst_block_ids": meta.local_block_ids, + "dst_slot": meta.local_slot_index, + } + ) + target = _sock_path( + meta.remote_handshake_port + + _port_offset(meta.remote_dp_rank, self.tp_rank, meta.tp_size) + ) + s = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) + s.connect(target) + socket.send_fds(s, [_MSG_WRITE_REQUEST + payload], [staging.export_fd()]) + resp = s.recv(4096) + s.close() + if resp[:1] != _MSG_WRITE_DONE: + logger.error("[native] no WRITE_DONE for req %s", req_id) + return + self._scatter(staging, meta.local_block_ids, meta.local_slot_index) + with self._lock: + self.done_recving.add(req_id) + + def _scatter(self, staging, dst_block_ids, dst_slot) -> None: + off = 0 + for base, bpb in self._block_regions: + for db in dst_block_ids: + vmm.copy(base + db * bpb, staging.data_ptr + off, bpb) + off += bpb + for base, bps in self._slot_regions: + if dst_slot >= 0: + vmm.copy(base + dst_slot * bps, staging.data_ptr + off, bps) + off += bps + if self._state_slot_bytes and dst_slot >= 0 and self._scatter_slot is not None: + pool_idx = self._acquire_state_slot() + if pool_idx >= 0: + vmm.copy( + self._state_base + pool_idx * self._state_slot_bytes, + staging.data_ptr + off, + self._state_slot_bytes, + ) + self._scatter_slot(dst_slot, pool_idx) + self._release_state_slot(pool_idx) + torch.cuda.synchronize(self.device) + + # -- producer ----------------------------------------------------------- + + def _serve(self) -> None: + path = _sock_path(self._port) + if os.path.exists(path): + os.unlink(path) + srv = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) + srv.bind(path) + srv.listen(64) + while True: + conn, _ = srv.accept() + threading.Thread(target=self._handle, args=(conn,), daemon=True).start() + + def _handle(self, conn: socket.socket) -> None: + try: + msg, fds, _, _ = socket.recv_fds(conn, 1 << 16, 1) + if not msg or msg[:1] != _MSG_WRITE_REQUEST: + return + req = msgpack.loads(msg[1:]) + dst = self._imported.get(fds[0]) + if dst is None: + nblocks = len(req["dst_block_ids"]) + dst = vmm.VmmBuffer.import_fd( + fds[0], self._req_bytes(nblocks), self.device + ) + self._imported[fds[0]] = dst + self._gather(dst, req["transfer_id"]) + conn.sendall(_MSG_WRITE_DONE + msgpack.dumps({"req_id": req["req_id"]})) + with self._lock: + self.done_sending.add(req["req_id"]) + except Exception: + logger.exception("[native] producer handler error") + finally: + conn.close() + + def _gather(self, staging, transfer_id: int) -> None: + with self._prefills_cv: + self._prefills_cv.wait_for( + lambda: transfer_id in self._prefills, timeout=_PREFILL_WAIT_S + ) + src_block_ids, src_slot = self._prefills.get(transfer_id, ([], -1)) + off = 0 + for base, bpb in self._block_regions: + for sb in src_block_ids: + vmm.copy(staging.data_ptr + off, base + sb * bpb, bpb) + off += bpb + for base, bps in self._slot_regions: + if src_slot >= 0: + vmm.copy(staging.data_ptr + off, base + src_slot * bps, bps) + off += bps + if self._state_slot_bytes and src_slot >= 0 and self._gather_slot is not None: + pool_idx = self._acquire_state_slot() + if pool_idx >= 0: + self._gather_slot(src_slot, pool_idx) + torch.cuda.current_stream().synchronize() + vmm.copy( + staging.data_ptr + off, + self._state_base + pool_idx * self._state_slot_bytes, + self._state_slot_bytes, + ) + self._release_state_slot(pool_idx) + torch.cuda.synchronize(self.device) diff --git a/atom/kv_transfer/disaggregation/native/vmm.py b/atom/kv_transfer/disaggregation/native/vmm.py new file mode 100644 index 0000000000..fe86c6d4dc --- /dev/null +++ b/atom/kv_transfer/disaggregation/native/vmm.py @@ -0,0 +1,103 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2025, Advanced Micro Devices, Inc. All rights reserved. + +"""HIP VMM cross-process GPU buffer sharing (scale-up KV transfer). + +A producer allocates an exportable VMM buffer and shares its POSIX fd (over a +UNIX socket); a consumer imports it, grants its device peer access, and copies +directly over the fabric. Unlike legacy hipIpc, VMM shareable handles are +reliable cross-process and don't need the source GPU visible to the consumer. +The C++ helper is JIT-compiled lazily (never at import), so importing is cheap +on CPU-only hosts. +""" + +from __future__ import annotations + +import functools +import os + +import torch + +__all__ = ["supported", "VmmBuffer"] + + +@functools.lru_cache(maxsize=1) +def _ext(): + from torch.utils.cpp_extension import load + + rocm = os.environ.get("ROCM_PATH", "/opt/rocm") + src = os.path.join(os.path.dirname(__file__), "_vmm_ext.cpp") + return load( + name="atom_vmm_ext", + sources=[src], + extra_include_paths=[os.path.join(rocm, "include")], + extra_ldflags=[f"-L{os.path.join(rocm, 'lib')}", "-lamdhip64"], + verbose=False, + ) + + +def supported(device: int = 0) -> bool: + """Whether the device supports HIP Virtual Memory Management.""" + return bool(_ext().vmm_supported(device)) + + +def copy(dst_ptr: int, src_ptr: int, nbytes: int) -> None: + """Device-to-device copy between raw device pointers (peer-mapped ok).""" + _ext().vmm_copy(dst_ptr, src_ptr, nbytes) + + +class VmmBuffer: + """An exportable VMM buffer mapped on one device. + + Create with :meth:`alloc` (producer) or :meth:`import_fd` (consumer). Use + :meth:`tensor` to get a (non-owning) view for reads/writes/copies. + """ + + def __init__(self, region_id: int, nbytes: int, device: int): + self._id = region_id + self.nbytes = nbytes + self.device = device + + @classmethod + def alloc(cls, nbytes: int, device: int) -> "VmmBuffer": + return cls(_ext().vmm_alloc(nbytes, device), nbytes, device) + + @classmethod + def import_fd(cls, fd: int, nbytes: int, device: int) -> "VmmBuffer": + """Import a producer's exported fd and map it on ``device``. + + ``nbytes`` must match the producer's ``alloc`` size. + """ + return cls(_ext().vmm_import(fd, nbytes, device), nbytes, device) + + def export_fd(self) -> int: + """POSIX fd to send to a peer (e.g. via ``socket.send_fds``).""" + return _ext().vmm_export_fd(self._id) + + @property + def data_ptr(self) -> int: + """Raw mapped device pointer (for :func:`copy`).""" + return _ext().vmm_ptr(self._id) + + def tensor(self, dtype: torch.dtype, shape) -> torch.Tensor: + """View of the buffer as ``dtype`` reshaped to ``shape``. + + The returned tensor keeps this :class:`VmmBuffer` (and therefore the + underlying VMM mapping) alive for its own lifetime, so callers may + drop the buffer reference and keep only the tensor. + """ + flat = _ext().vmm_tensor(self._id, self.nbytes, self.device) # uint8 + view = flat.view(dtype).view(*shape) + view._vmm_keepalive = self # tie mapping lifetime to the tensor + return view + + def close(self) -> None: + if self._id is not None: + _ext().vmm_free(self._id) + self._id = None + + def __del__(self): + try: + self.close() + except Exception: + pass diff --git a/tests/test_native_vmm_transfer.py b/tests/test_native_vmm_transfer.py new file mode 100644 index 0000000000..f2b753e9a6 --- /dev/null +++ b/tests/test_native_vmm_transfer.py @@ -0,0 +1,95 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2025, Advanced Micro Devices, Inc. All rights reserved. + +"""Cross-process HIP VMM KV-transfer test: a producer exports a VMM buffer on +GPU 0, a consumer imports it on GPU 1 and peer-copies blocks, then verifies. +Requires >= 2 GPUs with VMM support; skips otherwise. +""" + +from __future__ import annotations + +import socket + +import pytest +import torch +import torch.multiprocessing as mp + +NB, BE = 256, 4096 # 256 blocks x 4096 bf16 = 2 MiB; values < 128 are bf16-exact +NBYTES = NB * BE * 2 +NCOPY = 64 + + +def _producer(path, device, ready): + from atom.kv_transfer.disaggregation.native import VmmBuffer + + torch.cuda.set_device(device) + buf = VmmBuffer.alloc(NBYTES, device) + kv = buf.tensor(torch.bfloat16, (NB, BE)) + for i in range(NB): + kv[i].fill_(float(i % 128)) + torch.cuda.synchronize() + + srv = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) + srv.bind(path) + srv.listen(1) + ready.set() + conn, _ = srv.accept() + socket.send_fds(conn, [b"vmm"], [buf.export_fd()]) + conn.recv(1) # keep the buffer alive until the consumer is done + conn.close() + srv.close() + + +def _consumer(path, device, ready, result): + from atom.kv_transfer.disaggregation.native import VmmBuffer + + torch.cuda.set_device(device) + ready.wait(60) + s = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) + s.connect(path) + _, fds, _, _ = socket.recv_fds(s, 16, 1) + + peer = VmmBuffer.import_fd(fds[0], NBYTES, device).tensor(torch.bfloat16, (NB, BE)) + local = torch.empty(NB, BE, dtype=torch.bfloat16, device=device) + + # concurrent per-"request" peer copies across streams (hipMemcpyPeerAsync) + streams = [torch.cuda.Stream(device=device) for _ in range(NCOPY)] + for i, st in enumerate(streams): + with torch.cuda.stream(st): + local[i % NB].copy_(peer[(i * 7) % NB], non_blocking=True) + for st in streams: + st.synchronize() + + ok = all( + local[i % NB][0].item() == float(((i * 7) % NB) % 128) for i in range(NCOPY) + ) + result.put(bool(ok)) + s.send(b"d") + s.close() + + +@pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.device_count() < 2, + reason="requires >= 2 GPUs", +) +def test_vmm_cross_process_transfer(tmp_path): + from atom.kv_transfer.disaggregation.native import supported + + if not supported(0) or not supported(1): + pytest.skip("HIP VMM not supported on these devices") + + ctx = mp.get_context("spawn") + ready = ctx.Event() + result = ctx.Queue() + path = str(tmp_path / "vmm.sock") + + prod = ctx.Process(target=_producer, args=(path, 0, ready)) + cons = ctx.Process(target=_consumer, args=(path, 1, ready, result)) + prod.start() + cons.start() + cons.join(180) + prod.join(30) + + assert cons.exitcode == 0, "consumer process crashed" + assert prod.exitcode == 0, "producer process crashed" + assert result.get(timeout=5) is True, "peer-copied data mismatch"