diff --git a/docs/D_TO_P_PHASE1_LINK_ZH.md b/docs/D_TO_P_PHASE1_LINK_ZH.md index fce0c76..8f2fc6b 100644 --- a/docs/D_TO_P_PHASE1_LINK_ZH.md +++ b/docs/D_TO_P_PHASE1_LINK_ZH.md @@ -92,13 +92,25 @@ peer.close() ## 4. 当前覆盖范围（清单） -- ✅ Host CPU 内存的 D→P RDMA byte transfer +- ✅ Host CPU 内存的 D→P RDMA byte transfer (`scripts/smoke_snapshot_link.py`) +- ✅ **GPU 内存** cuda:0 → cuda:1 的 D→P RDMA（`scripts/smoke_snapshot_link_gpu.py`，5/5 size 全 SHA 校验通过，256 MB 8.5 ms / 251 Gbps） - ✅ 单 IB device (mlx5_60) - ✅ 同节点 loopback（127.0.0.1） -- ⏳ GPU 内存（设备指针 + `batch_transfer_write_on_cuda`）—— 现有 `push()` 走 `transfer_sync_write`，对 GPU 指针支持取决于 mooncake 的 protocol；下一步验证 - ⏳ 跨节点（远端 IP）—— 设计上一致，未验证 - ⏳ 多 D → 单 P（多 sender → 共享 recv buffer 的 offset 协调）—— 留给 Phase 3 整合时设计 -- ⏳ ZeroCopy 入 SGLang kv_pool slot —— 留给 Phase 2 +- ⏳ ZeroCopy 入 SGLang kv_pool slot —— 留给 Phase 2/3 + +### GPU smoke 性能 + +| Size | Push duration | Throughput | +|---:|---:|---:| +| 16 KB | 8.27 ms (cold) | 0.016 Gbps | +| 1 MB | 0.096 ms | 87.6 Gbps | +| 16 MB | 0.844 ms | 159 Gbps | +| 64 MB | 2.52 ms | 213 Gbps | +| **256 MB** | **8.54 ms** | **251 Gbps** | + +GPU↔GPU 比 host↔host 慢一些（251 vs 316 Gbps for 64MB），但仍接近 mlx5_60 NDR 400Gb 的 60% 线率。对 KVC 单 session ~50K tokens × ~80 KB/token ≈ 4 GB 量级的 transfer，对应 D→P 时间约 130 ms。 ## 5. 下一步（Phase 2 / Phase 3） diff --git a/scripts/smoke_snapshot_link_gpu.py b/scripts/smoke_snapshot_link_gpu.py new file mode 100644 index 0000000..be1958c --- /dev/null +++ b/scripts/smoke_snapshot_link_gpu.py @@ -0,0 +1,236 @@ +#!/usr/bin/env python3 +"""GPU-aware smoke test for snapshot_link RDMA byte transfer. + +Sender on cuda:0, receiver subprocess on cuda:1. Tests whether +mooncake's transfer_sync_write can move bytes between two GPUs via +RDMA (which is what the real D→P flow will need for KV bytes). + +Usage: + bash scripts/setup_env.sh && uv run --no-sync python scripts/smoke_snapshot_link_gpu.py + +The sender uses cuda:0 (--send-gpu); the receiver subprocess uses +cuda:1 (--recv-gpu) by default. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import os +import subprocess +import sys +import tempfile +import time +from pathlib import Path + +_HERE = Path(__file__).resolve().parent +sys.path.insert(0, str(_HERE.parent / "src")) + + +SIZES_BYTES_DEFAULT = [ + 1 << 14, # 16 KB + 1 << 20, # 1 MB + 1 << 24, # 16 MB + 1 << 26, # 64 MB + 1 << 28, # 256 MB +] + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--host", default=os.environ.get("SNAPSHOT_LINK_HOST", "127.0.0.1")) + ap.add_argument("--ib", default=os.environ.get("SNAPSHOT_LINK_IB", "mlx5_60")) + ap.add_argument("--recv-port", type=int, + default=int(os.environ.get("SNAPSHOT_LINK_RECV_PORT", "17787"))) + ap.add_argument("--send-port", type=int, + default=int(os.environ.get("SNAPSHOT_LINK_SEND_PORT", "17788"))) + ap.add_argument("--max-bytes", type=int, default=256 * 1024 * 1024) + ap.add_argument("--sizes", default=",".join(str(s) for s in SIZES_BYTES_DEFAULT)) + ap.add_argument("--send-gpu", type=int, default=0) + ap.add_argument("--recv-gpu", type=int, default=1) + args = ap.parse_args() + + sizes = [int(s) for s in args.sizes.split(",")] + tmpdir = Path(tempfile.mkdtemp(prefix="snapshot_link_gpu_smoke_")) + control_path = tmpdir / "endpoint.json" + recv_stderr_log = tmpdir / "recv.stderr.log" + + recv_cmd = [ + sys.executable, + str(_HERE / "snapshot_link_receiver_gpu.py"), + "--host", args.host, + "--port", str(args.recv_port), + "--ib", args.ib, + "--max-bytes", str(args.max_bytes), + "--control-path", str(control_path), + "--sizes", args.sizes, + "--gpu-id", str(args.recv_gpu), + ] + recv_stderr = open(recv_stderr_log, "w") + print(f"[sender] receiver cmd: {' '.join(recv_cmd)}", flush=True) + recv_proc = subprocess.Popen( + recv_cmd, stdout=subprocess.PIPE, stderr=recv_stderr, bufsize=1, + universal_newlines=True, + ) + + try: + import torch + if not torch.cuda.is_available(): + print("[sender] FAIL: cuda not available") + return 1 + torch.cuda.set_device(args.send_gpu) + + deadline = time.time() + 90.0 + meta = None + while time.time() < deadline: + if control_path.exists(): + try: + meta = json.loads(control_path.read_text()) + if meta.get("ready"): + break + except Exception: + pass + if recv_proc.poll() is not None: + _dump_recv_stderr(recv_stderr_log) + print(f"[sender] FAIL: receiver exited (rc={recv_proc.returncode})") + return 1 + time.sleep(0.1) + if meta is None: + print("[sender] FAIL: receiver endpoint timeout") + return 1 + print(f"[sender] receiver endpoint: gpu={meta['gpu_id']}, " + f"sid={meta['session_id']}, ptr={hex(int(meta['base_ptr']))}, " + f"cap={meta['capacity_bytes']}", flush=True) + + from agentic_pd_hybrid.snapshot_link import SnapshotPeer, SnapshotEndpoint + + endpoint = SnapshotEndpoint( + session_id=meta["session_id"], + base_ptr=int(meta["base_ptr"]), + capacity_bytes=int(meta["capacity_bytes"]), + ) + + peer = SnapshotPeer( + host=args.host, + port=args.send_port, + ib_device=args.ib, + receive_capacity_bytes=0, + ) + + # Allocate a sender buffer on cuda:0 + send_tensor = torch.zeros(args.max_bytes, dtype=torch.uint8, + device=f"cuda:{args.send_gpu}") + send_ptr = send_tensor.data_ptr() + ret = peer.engine.register_memory(send_ptr, args.max_bytes) + if ret != 0: + print(f"[sender] FAIL: register_memory ret={ret}") + return 1 + print(f"[sender] own gpu={args.send_gpu}, sid={peer.session_id}, " + f"buf @ {hex(send_ptr)} ({args.max_bytes} B)", flush=True) + + transfers = [] + for size in sizes: + if size > args.max_bytes: + continue + # Fill with deterministic pattern on GPU + seed = int(time.time() * 1e6) & 0xFFFFFFFF + # Use a simple seeded pattern via torch ops + gen = torch.Generator(device=f"cuda:{args.send_gpu}") + gen.manual_seed(seed) + send_tensor[:size] = torch.randint(0, 256, (size,), dtype=torch.uint8, + device=f"cuda:{args.send_gpu}", + generator=gen) + torch.cuda.synchronize(args.send_gpu) + # Compute expected hash (host-side) + host_view = send_tensor[:size].cpu().numpy().tobytes() + expected_sha = hashlib.sha256(host_view).hexdigest() + # Push via RDMA + t0 = time.perf_counter() + ret = peer.push(endpoint, send_ptr, 0, size, remote_offset=0) + t1 = time.perf_counter() + dt_ms = (t1 - t0) * 1000.0 + gbps = (size * 8.0 / 1e9) / max(t1 - t0, 1e-9) + print(f"[sender] push size={size:>10d} ret={ret} " + f"dur={dt_ms:>9.3f} ms thru={gbps:>6.3f} Gbps", + flush=True) + + # Signal receiver to verify + signal_path = control_path.with_suffix(f".do{size}") + ack_path = control_path.with_suffix(f".ack{size}") + signal_path.write_text(json.dumps({"sha": expected_sha})) + ack_deadline = time.time() + 90.0 + while time.time() < ack_deadline: + if ack_path.exists(): + break + if recv_proc.poll() is not None: + print(f"[sender] FAIL: receiver died after size={size}") + _dump_recv_stderr(recv_stderr_log) + return 1 + time.sleep(0.05) + transfers.append({ + "size": size, "ret": ret, "dur_ms": round(dt_ms, 3), + "thru_Gbps": round(gbps, 3), "ack": ack_path.exists(), + }) + + try: + recv_proc.wait(timeout=10) + except subprocess.TimeoutExpired: + recv_proc.terminate() + recv_proc.wait(timeout=5) + + events = [] + if recv_proc.stdout is not None: + for raw in recv_proc.stdout: + raw = raw.strip() + if not raw: + continue + try: + events.append(json.loads(raw)) + except json.JSONDecodeError: + events.append({"event": "non-json", "raw": raw}) + + print("=" * 78) + print("[receiver] events:") + verify_ok = 0 + verify_fail = 0 + for ev in events: + print(f" {ev}") + if ev.get("event") == "verify": + if ev.get("ok"): + verify_ok += 1 + else: + verify_fail += 1 + + recv_stderr.close() + _dump_recv_stderr(recv_stderr_log, header="--- receiver stderr ---") + + overall = "PASS" if verify_fail == 0 and verify_ok == len(transfers) else "FAIL" + print("=" * 78) + print(f"OVERALL: {overall} verify_ok={verify_ok} verify_fail={verify_fail} " + f"transfers={len(transfers)} send_gpu={args.send_gpu} recv_gpu={args.recv_gpu}") + return 0 if overall == "PASS" else 1 + + finally: + try: + recv_proc.terminate() + recv_proc.wait(timeout=5) + except Exception: + try: + recv_proc.kill() + except Exception: + pass + + +def _dump_recv_stderr(path: Path, header: str = "--- receiver stderr (last 60) ---") -> None: + try: + text = path.read_text() + except FileNotFoundError: + return + print(header, flush=True) + for line in text.splitlines()[-60:]: + print(f" {line}", flush=True) + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/snapshot_link_receiver_gpu.py b/scripts/snapshot_link_receiver_gpu.py new file mode 100644 index 0000000..5d12ccb --- /dev/null +++ b/scripts/snapshot_link_receiver_gpu.py @@ -0,0 +1,124 @@ +#!/usr/bin/env python3 +"""GPU-side receiver child for snapshot_link smoke test (CUDA mem).""" +from __future__ import annotations + +import argparse +import hashlib +import json +import sys +import time +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "src")) + + +def _emit(d: dict) -> None: + print(json.dumps(d), flush=True) + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--host", required=True) + ap.add_argument("--port", type=int, required=True) + ap.add_argument("--ib", required=True) + ap.add_argument("--max-bytes", type=int, required=True) + ap.add_argument("--control-path", required=True) + ap.add_argument("--sizes", required=True) + ap.add_argument("--gpu-id", type=int, default=1, help="receiver GPU id") + args = ap.parse_args() + + sizes = [int(s) for s in args.sizes.split(",")] + + try: + import torch + if not torch.cuda.is_available(): + _emit({"event": "init-failed", "error": "cuda not available"}) + sys.exit(2) + torch.cuda.set_device(args.gpu_id) + # allocate a GPU buffer of max_bytes + recv_tensor = torch.zeros(args.max_bytes, dtype=torch.uint8, device=f"cuda:{args.gpu_id}") + recv_ptr = recv_tensor.data_ptr() + except Exception as e: + import traceback + _emit({"event": "init-failed", "error": repr(e), "tb": traceback.format_exc()}) + sys.exit(2) + + # Spin up SnapshotPeer with NO internal recv buffer, then register our GPU tensor + from agentic_pd_hybrid.snapshot_link import SnapshotPeer, SnapshotEndpoint + try: + peer = SnapshotPeer( + host=args.host, + port=args.port, + ib_device=args.ib, + receive_capacity_bytes=0, + ) + ret = peer.engine.register_memory(recv_ptr, args.max_bytes) + if ret != 0: + _emit({"event": "init-failed", "error": f"register_memory({hex(recv_ptr)}, {args.max_bytes}) ret={ret}"}) + sys.exit(2) + except Exception as e: + import traceback + _emit({"event": "init-failed", "error": repr(e), "tb": traceback.format_exc()}) + sys.exit(2) + + endpoint = SnapshotEndpoint( + session_id=peer.session_id, + base_ptr=recv_ptr, + capacity_bytes=args.max_bytes, + ) + Path(args.control_path).write_text(json.dumps({ + "session_id": endpoint.session_id, + "base_ptr": endpoint.base_ptr, + "capacity_bytes": endpoint.capacity_bytes, + "gpu_id": args.gpu_id, + "ready": True, + })) + _emit({"event": "endpoint-ready", + "session_id": endpoint.session_id, + "base_ptr": endpoint.base_ptr, + "capacity": endpoint.capacity_bytes, + "gpu_id": args.gpu_id}) + + cp = Path(args.control_path) + for size in sizes: + if size > args.max_bytes: + continue + signal_path = cp.with_suffix(f".do{size}") + ack_path = cp.with_suffix(f".ack{size}") + deadline = time.time() + 120.0 + while time.time() < deadline: + if signal_path.exists(): + break + time.sleep(0.05) + else: + _emit({"event": "no-signal-timeout", "size": size}) + continue + try: + payload = json.loads(signal_path.read_text()) + expected_sha = payload["sha"] + except Exception as e: + _emit({"event": "signal-parse-error", "size": size, "err": repr(e)}) + continue + + # Copy from GPU to CPU and hash + torch.cuda.synchronize(args.gpu_id) + host_bytes = bytes(recv_tensor[:size].cpu().numpy().tobytes()) + recv_sha = hashlib.sha256(host_bytes).hexdigest() + ok = recv_sha == expected_sha + _emit({ + "event": "verify", + "size": size, + "ok": ok, + "expected_sha": expected_sha[:16], + "got_sha": recv_sha[:16], + "first8_recv": host_bytes[:8].hex(), + "last8_recv": host_bytes[-8:].hex(), + }) + ack_path.write_text("done") + + peer.close() + _emit({"event": "receiver-done"}) + + +if __name__ == "__main__": + main()