[KVConnector] native scale-up KV connector (HIP VMM, kv_connector=native)

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.

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.

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

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 <torch/extension.h>
+#include <hip/hip_runtime.h>
+
+#include <stdexcept>
+#include <string>
+#include <unordered_map>
+
+#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<int64_t, Region> 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 &region(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<size_t>(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<size_t>(nbytes), device);
+  Region r{};
+  r.size = size;
+  HIPCK(hipMemImportFromShareableHandle(
+      &r.handle, reinterpret_cast<void *>(static_cast<intptr_t>(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<int64_t>(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<void *>(dst_ptr),
+                  reinterpret_cast<void *>(src_ptr),
+                  static_cast<size_t>(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);
+}