diff --git a/csrc/BinaryGraphIO.cpp b/csrc/BinaryGraphIO.cpp index 79c8ac81..419e056d 100644 --- a/csrc/BinaryGraphIO.cpp +++ b/csrc/BinaryGraphIO.cpp @@ -12,6 +12,46 @@ namespace foundry { +namespace { + +bool validate_binary_layout(const uint8_t* data, size_t file_size, + const binary_format::FileHeader& header) { + namespace bf = binary_format; + + if (header.num_sections != bf::SECTION_COUNT) { + return false; + } + + const uint64_t directory_size = + static_cast(header.num_sections) * sizeof(bf::SectionEntry); + if (directory_size > file_size - sizeof(bf::FileHeader)) { + return false; + } + + const auto* sections = reinterpret_cast(data + sizeof(bf::FileHeader)); + for (uint32_t i = 0; i < header.num_sections; ++i) { + const auto& section = sections[i]; + if (section.section_type != i || section.offset > file_size || + section.size > file_size - section.offset) { + return false; + } + } + + const uint64_t legacy_dependency_size = + static_cast(header.num_dependencies) * sizeof(bf::LegacyBinDependency); + const uint64_t dependency_size = + static_cast(header.num_dependencies) * sizeof(bf::BinDependency); + const uint64_t stored_dependency_size = sections[bf::SECTION_DEPENDENCY_TABLE].size; + const bool has_dependencies = (header.flags & bf::FLAG_HAS_DEPENDENCIES) != 0; + if (has_dependencies != (header.num_dependencies != 0)) { + return false; + } + return stored_dependency_size == legacy_dependency_size || + stored_dependency_size == dependency_size; +} + +} // namespace + void CUDAGraph::save_binary(const std::string& bin_path, const boost::json::object& root) { namespace json = boost::json; namespace bf = binary_format; @@ -270,7 +310,19 @@ void CUDAGraph::save_binary(const std::string& bin_path, const boost::json::obje deps.reserve(deps_array.size()); for (const auto& dep_val : deps_array) { const json::object& d = dep_val.as_object(); - deps.push_back({d.at("from").to_number(), d.at("to").to_number()}); + bf::BinDependency dep = {}; + dep.from_id = d.at("from").to_number(); + dep.to_id = d.at("to").to_number(); + if (d.contains("from_port")) { + dep.from_port = static_cast(d.at("from_port").to_number()); + } + if (d.contains("to_port")) { + dep.to_port = static_cast(d.at("to_port").to_number()); + } + if (d.contains("type")) { + dep.type = static_cast(d.at("type").to_number()); + } + deps.push_back(dep); } // ---- Generators ---- @@ -397,7 +449,7 @@ boost::json::value read_and_parse_binary_graph(const std::string& bin_path) { if (file_size < sizeof(bf::FileHeader)) return json::value{}; const bf::FileHeader& header = *reinterpret_cast(data); - if (!bf::validate_header(header)) + if (!bf::validate_header(header) || !validate_binary_layout(data, file_size, header)) return json::value{}; // Read section directory @@ -653,13 +705,26 @@ boost::json::value read_and_parse_binary_graph(const std::string& bin_path) { // Dependencies if (header.flags & bf::FLAG_HAS_DEPENDENCIES) { auto [dep_data, dep_size] = section(bf::SECTION_DEPENDENCY_TABLE); - const bf::BinDependency* deps = reinterpret_cast(dep_data); + const size_t dependency_entry_size = dep_size / header.num_dependencies; + const bool has_edge_data = dependency_entry_size == sizeof(bf::BinDependency); json::array deps_array; deps_array.reserve(header.num_dependencies); for (uint32_t i = 0; i < header.num_dependencies; i++) { + const uint8_t* record = dep_data + i * dependency_entry_size; + const auto* legacy = reinterpret_cast(record); json::object d; - d["from"] = deps[i].from_id; - d["to"] = deps[i].to_id; + d["from"] = legacy->from_id; + d["to"] = legacy->to_id; + if (has_edge_data) { + const auto* dep = reinterpret_cast(record); + d["from_port"] = static_cast(dep->from_port); + d["to_port"] = static_cast(dep->to_port); + d["type"] = static_cast(dep->type); + } else { + d["from_port"] = 0; + d["to_port"] = 0; + d["type"] = 0; + } deps_array.push_back(d); } root["dependencies"] = std::move(deps_array); @@ -725,7 +790,8 @@ BinaryGraphFile read_binary_graph_file(const std::string& bin_path) { in.close(); memcpy(&result.header, result.data.data(), sizeof(bf::FileHeader)); - if (!bf::validate_header(result.header)) { + if (!bf::validate_header(result.header) || + !validate_binary_layout(result.data.data(), file_size, result.header)) { result.data.clear(); return result; } diff --git a/csrc/CUDAGraph.cpp b/csrc/CUDAGraph.cpp index 25df3524..47ee2d9e 100644 --- a/csrc/CUDAGraph.cpp +++ b/csrc/CUDAGraph.cpp @@ -16,6 +16,8 @@ #include #include #include +#include +#include #include #include #include @@ -127,6 +129,47 @@ MempoolId_t torch_graph_pool_handle(bool is_user_created = true) { } // namespace +CUresult add_graph_dependencies(CUgraph graph, const std::vector& from_nodes, + const std::vector& to_nodes, + const std::vector& edge_data) { + TORCH_INTERNAL_ASSERT(from_nodes.size() == to_nodes.size()); + TORCH_INTERNAL_ASSERT(from_nodes.size() == edge_data.size()); + + struct DependencyBatch { + std::vector from_nodes; + std::vector to_nodes; + std::vector edge_data; + }; + using EdgeDataKey = std::tuple; + std::map batches; + + // CUDA 13.0 silently normalizes non-default entries when one call mixes edge annotations. + // Group identical annotations so each driver call remains lossless without adding edges singly. + for (size_t index = 0; index < edge_data.size(); ++index) { + const auto& data = edge_data[index]; + auto& batch = batches[{data.from_port, data.to_port, data.type}]; + batch.from_nodes.push_back(from_nodes[index]); + batch.to_nodes.push_back(to_nodes[index]); + batch.edge_data.push_back(data); + } + + for (auto& [key, batch] : batches) { + (void)key; +#if (defined(CUDA_VERSION) && CUDA_VERSION >= 13000) + CUresult result = cuGraphAddDependencies(graph, batch.from_nodes.data(), batch.to_nodes.data(), + batch.edge_data.data(), batch.edge_data.size()); +#else + CUresult result = + cuGraphAddDependencies_v2(graph, batch.from_nodes.data(), batch.to_nodes.data(), + batch.edge_data.data(), batch.edge_data.size()); +#endif + if (result != CUDA_SUCCESS) { + return result; + } + } + return CUDA_SUCCESS; +} + static bool _cuda_graphs_debug = false; static CUDAGeneratorStateRegistry global_generator_state_registry; @@ -934,6 +977,7 @@ void CUDAGraph::analyze_captured_graph() { GraphDependency dep; dep.from_index = from_it->second; dep.to_index = to_it->second; + dep.edge_data = edges[i]; graph_dependencies.push_back(dep); } } @@ -1335,7 +1379,8 @@ void CUDAGraph::save(const std::string& json_path, const OutputTensors& output_t nodes_array.push_back(node_obj); } - // Compute topology key = node types + cluster dim values per kernel node. + // Compute topology key = node types + cluster dim values per kernel node + + // dependency endpoints and edge data. // cuGraphExecUpdate does NOT propagate kernel node attribute changes // (set via cuGraphKernelNodeSetAttribute) to the CUgraphExec — only // CUDA_KERNEL_NODE_PARAMS fields are updated. This means cluster dim @@ -1381,6 +1426,22 @@ void CUDAGraph::save(const std::string& json_path, const OutputTensors& output_t } } } + + std::vector sorted_dependencies = graph_dependencies; + std::sort(sorted_dependencies.begin(), sorted_dependencies.end(), + [](const GraphDependency& lhs, const GraphDependency& rhs) { + return std::make_tuple(lhs.from_index, lhs.to_index, lhs.edge_data.from_port, + lhs.edge_data.to_port, lhs.edge_data.type) < + std::make_tuple(rhs.from_index, rhs.to_index, rhs.edge_data.from_port, + rhs.edge_data.to_port, rhs.edge_data.type); + }); + topology_key += "|D"; + for (const auto& dep : sorted_dependencies) { + topology_key += ";" + std::to_string(dep.from_index) + ">" + std::to_string(dep.to_index) + + ":" + std::to_string(static_cast(dep.edge_data.from_port)) + ":" + + std::to_string(static_cast(dep.edge_data.to_port)) + ":" + + std::to_string(static_cast(dep.edge_data.type)); + } root["topology_key"] = topology_key; } @@ -1447,6 +1508,9 @@ void CUDAGraph::save(const std::string& json_path, const OutputTensors& output_t json::object dep_obj; dep_obj["from"] = dep.from_index; dep_obj["to"] = dep.to_index; + dep_obj["from_port"] = static_cast(dep.edge_data.from_port); + dep_obj["to_port"] = static_cast(dep.edge_data.to_port); + dep_obj["type"] = static_cast(dep.edge_data.type); deps_array.push_back(dep_obj); } root["dependencies"] = deps_array; @@ -2031,25 +2095,34 @@ GraphLoadResult CUDAGraph::load(const std::string& json_path, MempoolId_t pool) if (!deps_array.empty()) { std::vector from_nodes; std::vector to_nodes; + std::vector edge_data; from_nodes.reserve(deps_array.size()); to_nodes.reserve(deps_array.size()); + edge_data.reserve(deps_array.size()); for (const auto& dep_val : deps_array) { const json::object& dep_obj = dep_val.as_object(); int from_id = dep_obj.at("from").to_number(); int to_id = dep_obj.at("to").to_number(); + CUgraphEdgeData data{}; + if (dep_obj.contains("from_port")) { + data.from_port = static_cast( + dep_obj.at("from_port").to_number()); + } + if (dep_obj.contains("to_port")) { + data.to_port = + static_cast(dep_obj.at("to_port").to_number()); + } + if (dep_obj.contains("type")) { + data.type = static_cast(dep_obj.at("type").to_number()); + } from_nodes.push_back(id_to_node[from_id]); to_nodes.push_back(id_to_node[to_id]); + edge_data.push_back(data); } -#if (defined(CUDA_VERSION) && CUDA_VERSION >= 13000) - CUresult dep_result = cuGraphAddDependencies(cuGraph, from_nodes.data(), to_nodes.data(), - nullptr, deps_array.size()); -#else - CUresult dep_result = cuGraphAddDependencies_v2(cuGraph, from_nodes.data(), to_nodes.data(), - nullptr, deps_array.size()); -#endif + CUresult dep_result = add_graph_dependencies(cuGraph, from_nodes, to_nodes, edge_data); if (dep_result != CUDA_SUCCESS) { fprintf(stderr, "[foundry LOAD ERROR] cuGraphAddDependencies FAILED with error %d\n", dep_result); diff --git a/csrc/CUDAGraphParallel.cpp b/csrc/CUDAGraphParallel.cpp index 6ebfc0f6..05bfce7e 100644 --- a/csrc/CUDAGraphParallel.cpp +++ b/csrc/CUDAGraphParallel.cpp @@ -776,22 +776,31 @@ GraphLoadResult CUDAGraph::build_graph_from_parsed(ParsedGraphData&& parsed, CUc if (!deps_array.empty()) { std::vector from_nodes; std::vector to_nodes; + std::vector edge_data; from_nodes.reserve(deps_array.size()); to_nodes.reserve(deps_array.size()); + edge_data.reserve(deps_array.size()); for (const auto& dep_val : deps_array) { const json::object& dep_obj = dep_val.as_object(); from_nodes.push_back(id_to_node[dep_obj.at("from").to_number()]); to_nodes.push_back(id_to_node[dep_obj.at("to").to_number()]); + CUgraphEdgeData data{}; + if (dep_obj.contains("from_port")) { + data.from_port = static_cast( + dep_obj.at("from_port").to_number()); + } + if (dep_obj.contains("to_port")) { + data.to_port = + static_cast(dep_obj.at("to_port").to_number()); + } + if (dep_obj.contains("type")) { + data.type = static_cast(dep_obj.at("type").to_number()); + } + edge_data.push_back(data); } -#if (defined(CUDA_VERSION) && CUDA_VERSION >= 13000) - CUresult dep_result = cuGraphAddDependencies(cuGraph, from_nodes.data(), to_nodes.data(), - nullptr, deps_array.size()); -#else - CUresult dep_result = cuGraphAddDependencies_v2(cuGraph, from_nodes.data(), to_nodes.data(), - nullptr, deps_array.size()); -#endif + CUresult dep_result = add_graph_dependencies(cuGraph, from_nodes, to_nodes, edge_data); if (dep_result != CUDA_SUCCESS) { fprintf(stderr, "[foundry LOAD ERROR] cuGraphAddDependencies FAILED with error %d\n", dep_result); diff --git a/csrc/hook.cpp b/csrc/hook.cpp index d552d08d..35909151 100644 --- a/csrc/hook.cpp +++ b/csrc/hook.cpp @@ -14,12 +14,21 @@ #include #include #include +#include #include #include #include #include #include #include +#include +#include +#include +#include +#include +#include +#include +#include #include #include #include @@ -253,6 +262,27 @@ static std::once_flag default_allocation_region_flag; static boost::unordered::concurrent_flat_map global_alloc_metadata; static boost::unordered::concurrent_flat_map global_carved_reserve_metadata; +// Defined in the VMM-IPC section below (inside the extern "C" block, hence +// the matching linkage here); closes and forgets the exported fd for a VA +// when its allocation is freed (an exported fd keeps the physical pages +// alive and would otherwise serve stale memory to importers). +extern "C" { +static void vmm_ipc_invalidate_export(CUdeviceptr dptr); +} + +// Process-global mirrors of live LOAD-mode preallocated chunks. The +// authoritative state is thread-local, but cuIpcGetMemHandle may run on a +// different thread. Keep one coherent record per range so concurrent devices +// or regions cannot overwrite each other's handle/base/size tuple. +struct PreallocatedChunk { + size_t size; + CUmemGenericAllocationHandle handle; + uint64_t generation; +}; +static std::mutex preallocated_chunks_mutex; +static std::map preallocated_chunks; +static std::atomic next_preallocated_chunk_generation{1}; + struct HookAllocationEvent { enum class Type { Alloc, Free, Reserve }; Type type; @@ -2657,6 +2687,7 @@ CUresult cuMemFree_v2(CUdeviceptr dptr) { mem_addr_free_func(metadata.ptr, metadata.size); } + vmm_ipc_invalidate_export(dptr); global_alloc_metadata.erase_if( [dptr](const std::pair& kv) { return kv.first == dptr; }); @@ -2702,9 +2733,10 @@ CUresult cuMemAddressReserve(CUdeviceptr* ptr, size_t size, size_t alignment, CU return result; } + const size_t effective_alignment = alignment == 0 ? kAllocAlignment : alignment; constexpr size_t kSmallAllocThreshold = 2ULL << 30; if (size < kSmallAllocThreshold) { - CUdeviceptr aligned_addr = align_to(tls_storage.current_alloc_base_addr, alignment); + CUdeviceptr aligned_addr = align_to(tls_storage.current_alloc_base_addr, effective_alignment); size_t region_base = (size_t)tls_storage.region.base; size_t region_end = region_base + tls_storage.region.size; @@ -2744,7 +2776,7 @@ CUresult cuMemAddressReserve(CUdeviceptr* ptr, size_t size, size_t alignment, CU // constexpr size_t kReserveAlignment = 2ULL << 37; // FIXME(liuxs): This value should be equal to // the first reservation size - CUdeviceptr hint_addr = align_to(tls_storage.current_vmm_reserve_addr, alignment); + CUdeviceptr hint_addr = align_to(tls_storage.current_vmm_reserve_addr, effective_alignment); CUresult result = real_func(ptr, size, alignment, hint_addr, flags); if (result != CUDA_SUCCESS) { fprintf(stderr, "[HOOK] ERROR: cuMemAddressReserve failed with error 0x%x\n", result); @@ -2759,7 +2791,7 @@ CUresult cuMemAddressReserve(CUdeviceptr* ptr, size_t size, size_t alignment, CU (unsigned long long)tls_storage.current_vmm_reserve_addr); } - tls_storage.current_vmm_reserve_addr = align_to(*ptr + size, alignment); + tls_storage.current_vmm_reserve_addr = align_to(*ptr + size, effective_alignment); #ifdef HOOK_DEBUG fprintf(stderr, @@ -2860,8 +2892,305 @@ void* dlsym(void* handle, const char* symbol) { // VMM allocations by translating to cuMemExportToShareableHandle API // ============================================================================= -// Magic marker to identify VMM IPC handles in CUipcMemHandle -static constexpr uint32_t VMM_IPC_MAGIC = 0x564D4D49; // "VMMI" +// Magic marker to identify VMM IPC handles in CUipcMemHandle. +// v2 ("VMM2"): the blob carries {magic, exporter pid, original ptr, aligned size}. +// The POSIX fd itself is transferred via SCM_RIGHTS over a per-process abstract +// unix socket (a raw fd integer is meaningless in another process - the v1 bug +// behind "cuMemImportFromShareableHandle failed with error 999"). +static constexpr uint32_t VMM_IPC_MAGIC = 0x564D4D32; // "VMM2" + +// --------------------------------------------------------------------------- +// VMM-IPC fd transport: each exporting process runs one server thread on an +// abstract unix socket "\0foundry-vmm-ipc.". Importers connect, send the +// 8-byte exporter VA they want, and receive the exported fd via SCM_RIGHTS. +// The server does no CUDA calls: fds are exported on the cuIpcGetMemHandle +// caller's thread and parked in vmm_ipc_exported_fds. +// --------------------------------------------------------------------------- + +// exporter VA -> parked export. The allocation handle is kept so VA reuse +// after free/realloc invalidates the cached fd. The generation lets the server +// reject requests carrying a stale chunk generation (free/recreate at one +// base): it is the preallocated-chunk generation for chunk exports and 0 for +// direct (non-chunk) allocations. +struct VmmIpcExportedFd { + CUmemGenericAllocationHandle handle; + int fd; + uint64_t generation; +}; +static boost::unordered::concurrent_flat_map vmm_ipc_exported_fds; +static std::mutex vmm_ipc_server_mutex; +static int vmm_ipc_listen_fd = -1; +// pid that owns the running server thread. fork() copies this .so's state but +// not threads, so a forked child must rebind its own socket (vLLM's default +// worker start method is fork). +static pid_t vmm_ipc_server_pid = -1; +// Per-process random token carried in the handle blob and verified by the +// server: defends against pid reuse (a recycled pid + the deterministic +// shared-base VAs would otherwise let an importer silently fetch a different +// process's allocation). +static uint64_t vmm_ipc_token = 0; + +// Wire format of a fetch request. Internal to the per-process VMM-IPC socket +// (both peers run the same hook build within a run), so extending it does not +// touch the exported CUipcMemHandle blob layout. generation carries the chunk +// generation the importer expects (0 for direct allocations); the server +// refuses to transfer a descriptor whose live generation differs. +struct VmmIpcRequest { + uint64_t ptr; + uint64_t token; + uint64_t generation; +}; + +static void vmm_ipc_set_timeouts(int fd) { + struct timeval tv = {30, 0}; + setsockopt(fd, SOL_SOCKET, SO_RCVTIMEO, &tv, sizeof(tv)); + setsockopt(fd, SOL_SOCKET, SO_SNDTIMEO, &tv, sizeof(tv)); +} + +// Imported whole-chunk mappings (LOAD-mode exports: a buffer carved from the +// peer's preallocated chunk is shared by exporting the chunk handle plus an +// offset; cuMemMap cannot map a sub-range of an imported handle, so we map +// the entire peer chunk once and hand out interior pointers). The per-process +// token is part of the key so PID reuse cannot return an old exporter's mapping. +// Refcounted by open/close calls. +struct VmmIpcChunkMapping { + CUdeviceptr local_base; + size_t size; + CUmemGenericAllocationHandle handle; + int refcount; +}; +using VmmIpcChunkKey = std::tuple; +struct VmmIpcChunkSubptr { + VmmIpcChunkKey key; + int open_count; +}; +static std::mutex vmm_ipc_chunk_mutex; +static std::map vmm_ipc_chunk_mappings; +// Interior pointer handed to a caller -> owning chunk key and number of opens. +static std::map vmm_ipc_chunk_subptrs; + +// Validate an incoming chunk-carve open request against the ACTUAL cached +// whole-chunk mapping. The cache key (exporter pid, token, chunk_base, +// generation) deliberately omits chunk_size, so a same-key handle can carry a +// different (untrusted) chunk_size and carve than the live mapping. EVERY path +// that is about to adopt an existing mapping - the initial cache-hit fast path +// AND the insertion-race loser path, which can both race two first opens of the +// same peer chunk - must call this before mutating a refcount, registering an +// interior subpointer, or returning a pointer. Without it a same-key handle +// could smuggle a larger chunk_size and a carve landing outside the real +// mapping, yielding an out-of-mapping interior pointer. All arithmetic is +// overflow-safe because the handle values are untrusted. Must be called with +// vmm_ipc_chunk_mutex held (it reads the live mapping). Returns true iff the +// request is consistent with `mapping`; logs and returns false otherwise, and +// the caller returns CUDA_ERROR_INVALID_VALUE. +static bool vmm_ipc_chunk_request_matches_mapping(uint64_t chunk_base, uint64_t chunk_size, + CUdeviceptr carve_ptr, size_t carve_size, + const VmmIpcChunkMapping& mapping, + const char* context) { + const uint64_t cached_size = (uint64_t)mapping.size; + if (chunk_size != cached_size) { + fprintf(stderr, + "[HOOK] ERROR: VMM-IPC %s chunk_size mismatch: handle=%llu cached=%llu\n", + context, (unsigned long long)chunk_size, (unsigned long long)cached_size); + return false; + } + const uint64_t p = (uint64_t)carve_ptr; + const uint64_t s = (uint64_t)carve_size; + if (p < chunk_base || s > (UINT64_MAX - p) || (p + s) > (chunk_base + cached_size)) { + fprintf(stderr, + "[HOOK] ERROR: VMM-IPC %s carve [0x%llx,+0x%llx) outside cached mapping " + "[0x%llx,+0x%llx)\n", + context, (unsigned long long)p, (unsigned long long)s, + (unsigned long long)chunk_base, (unsigned long long)cached_size); + return false; + } + return true; +} + +// Dedicated VA zone for relocated peer imports, far below the foundry region +// (default 0x600000000000) and the large-reserve zone right above region end. +// Without this, the driver tends to place a relocated import immediately +// after the region reservation - exactly where the hooked large-reserve hint +// sends the NVSHMEM symmetric heap next, which silently breaks the heap's +// SAVE/LOAD address determinism (observed: heap displaced by a peer-chunk +// import landing at region_end). +static std::atomic vmm_ipc_import_hint{0x300000000000ULL}; + +static socklen_t vmm_ipc_socket_addr(pid_t pid, struct sockaddr_un* addr) { + memset(addr, 0, sizeof(*addr)); + addr->sun_family = AF_UNIX; + // Abstract namespace (sun_path[0] == '\0'): no filesystem entry, vanishes + // with the process. + int n = snprintf(addr->sun_path + 1, sizeof(addr->sun_path) - 2, "foundry-vmm-ipc.%d", (int)pid); + return (socklen_t)(offsetof(struct sockaddr_un, sun_path) + 1 + n); +} + +static void vmm_ipc_server_loop(int listen_fd) { + while (true) { + int conn = accept(listen_fd, nullptr, nullptr); + if (conn < 0) { + if (errno == EBADF || errno == EINVAL) + break; // socket gone + continue; // EINTR/ECONNABORTED/EMFILE/... are transient + } + vmm_ipc_set_timeouts(conn); + // Abstract sockets have no filesystem permissions: only serve same-uid peers. + struct ucred cred; + socklen_t cred_len = sizeof(cred); + if (getsockopt(conn, SOL_SOCKET, SO_PEERCRED, &cred, &cred_len) != 0 || cred.uid != getuid()) { + close(conn); + continue; + } + VmmIpcRequest req = {}; + int fd = -1; + if (recv(conn, &req, sizeof(req), MSG_WAITALL) == (ssize_t)sizeof(req) && + req.token == vmm_ipc_token) { + // Dup under the map lock: the export path may concurrently close and + // replace the parked fd (stale-entry refresh); the dup we send is ours. + // Only serve when the requested generation matches the live export, so a + // stale handle for a freed/recreated chunk at the same base is rejected. + vmm_ipc_exported_fds.visit((CUdeviceptr)req.ptr, + [&](const std::pair& kv) { + if (kv.second.generation == req.generation) + fd = fcntl(kv.second.fd, F_DUPFD_CLOEXEC, 0); + }); + } + char status = (fd >= 0) ? 0 : 1; + struct iovec iov = {&status, 1}; + char cbuf[CMSG_SPACE(sizeof(int))] = {}; + struct msghdr msg = {}; + msg.msg_iov = &iov; + msg.msg_iovlen = 1; + if (fd >= 0) { + msg.msg_control = cbuf; + msg.msg_controllen = CMSG_SPACE(sizeof(int)); + struct cmsghdr* c = CMSG_FIRSTHDR(&msg); + c->cmsg_level = SOL_SOCKET; + c->cmsg_type = SCM_RIGHTS; + c->cmsg_len = CMSG_LEN(sizeof(int)); + memcpy(CMSG_DATA(c), &fd, sizeof(int)); + } + // MSG_NOSIGNAL: a peer dying mid-handshake must not SIGPIPE-kill us. + if (sendmsg(conn, &msg, MSG_NOSIGNAL) != 1) { + fprintf(stderr, "[HOOK] WARNING: VMM-IPC fd server sendmsg failed (errno=%d)\n", errno); + } + if (fd >= 0) + close(fd); + close(conn); + } +} + +static bool vmm_ipc_ensure_server() { + std::lock_guard lock(vmm_ipc_server_mutex); + pid_t pid = getpid(); + if (vmm_ipc_server_pid == pid) { + return vmm_ipc_listen_fd >= 0; + } + if (vmm_ipc_server_pid != -1) { + // fork() inherits parked descriptors even with FD_CLOEXEC. The child has + // no server thread and must not pin or serve the parent's allocations. + vmm_ipc_exported_fds.erase_if([](const std::pair& kv) { + close(kv.second.fd); + return true; + }); + } + // First call in this process, or first after fork (the parent's accept + // thread did not survive). Close any inherited listen fd so the parent's + // abstract name is not pinned alive by us after the parent exits. + if (vmm_ipc_listen_fd >= 0) { + close(vmm_ipc_listen_fd); + vmm_ipc_listen_fd = -1; + } + int s = socket(AF_UNIX, SOCK_STREAM | SOCK_CLOEXEC, 0); + struct sockaddr_un addr; + socklen_t len = vmm_ipc_socket_addr(pid, &addr); + if (s < 0 || bind(s, (struct sockaddr*)&addr, len) != 0 || listen(s, 64) != 0) { + fprintf(stderr, "[HOOK] ERROR: VMM-IPC fd server setup failed (errno=%d)\n", errno); + if (s >= 0) + close(s); + vmm_ipc_listen_fd = -1; + vmm_ipc_server_pid = pid; // don't retry-spam; exports in this process fail loudly + return false; + } + // Per-process token (re-derived after fork). /dev/urandom, with a clock^pid + // fallback - this is anti-accident (pid reuse), not a security boundary; + // same-uid access control is SO_PEERCRED above. + uint64_t tok = 0; + int ur = open("/dev/urandom", O_RDONLY | O_CLOEXEC); + if (ur >= 0) { + if (read(ur, &tok, sizeof(tok)) != (ssize_t)sizeof(tok)) + tok = 0; + close(ur); + } + if (tok == 0) { + struct timeval tv; + gettimeofday(&tv, nullptr); + tok = ((uint64_t)tv.tv_sec << 32) ^ (uint64_t)tv.tv_usec ^ ((uint64_t)pid << 16) ^ + 0x9e3779b97f4a7c15ULL; + } + vmm_ipc_token = tok; + vmm_ipc_listen_fd = s; + vmm_ipc_server_pid = pid; + std::thread(vmm_ipc_server_loop, s).detach(); + return true; +} + +// Importer side: fetch the fd for `original_ptr` from `exporter_pid`'s server. +// Returns a live fd in THIS process, or -1. +static int vmm_ipc_fetch_fd(pid_t exporter_pid, uint64_t token, uint64_t generation, + CUdeviceptr original_ptr) { + int s = socket(AF_UNIX, SOCK_STREAM | SOCK_CLOEXEC, 0); + if (s < 0) + return -1; + vmm_ipc_set_timeouts(s); + struct sockaddr_un addr; + socklen_t len = vmm_ipc_socket_addr(exporter_pid, &addr); + if (connect(s, (struct sockaddr*)&addr, len) != 0) { + fprintf(stderr, + "[HOOK] ERROR: VMM-IPC connect to exporter pid %d failed (errno=%d) - exporter dead " + "or hook version mismatch\n", + (int)exporter_pid, errno); + close(s); + return -1; + } + VmmIpcRequest req = {(uint64_t)original_ptr, token, generation}; + if (send(s, &req, sizeof(req), MSG_NOSIGNAL) != (ssize_t)sizeof(req)) { + close(s); + return -1; + } + char status = 1; + struct iovec iov = {&status, 1}; + char cbuf[CMSG_SPACE(sizeof(int))] = {}; + struct msghdr msg = {}; + msg.msg_iov = &iov; + msg.msg_iovlen = 1; + msg.msg_control = cbuf; + msg.msg_controllen = sizeof(cbuf); + ssize_t r = recvmsg(s, &msg, MSG_CMSG_CLOEXEC); + close(s); + if (r != 1 || status != 0) + return -1; + for (struct cmsghdr* c = CMSG_FIRSTHDR(&msg); c != nullptr; c = CMSG_NXTHDR(&msg, c)) { + if (c->cmsg_level == SOL_SOCKET && c->cmsg_type == SCM_RIGHTS) { + int fd = -1; + memcpy(&fd, CMSG_DATA(c), sizeof(int)); + return fd; + } + } + return -1; +} + +// Close and forget the parked exported fd for a VA (called from the free +// path). Without this, the fd keeps the freed allocation's physical pages +// alive and the server would hand importers a stale mapping. +static void vmm_ipc_invalidate_export(CUdeviceptr dptr) { + vmm_ipc_exported_fds.erase_if([dptr](const std::pair& kv) { + if (kv.first != dptr) + return false; + close(kv.second.fd); + return true; + }); +} // Hook for cuIpcGetMemHandle - intercept Driver API to support VMM allocations CUresult cuIpcGetMemHandle(CUipcMemHandle* pHandle, CUdeviceptr dptr) { @@ -2878,7 +3207,46 @@ CUresult cuIpcGetMemHandle(CUipcMemHandle* pHandle, CUdeviceptr dptr) { metadata = kv.second; }); - if (found && metadata.handle != 0) { + // Decide what to export: the allocation's own handle (SAVE-mode slow-path + // allocs), or the whole preallocated chunk plus an offset (LOAD-mode fast + // path carves, which have no individual handle). cuMemMap cannot map a + // sub-range of an imported handle, so chunk carves are shared by exporting + // the chunk handle; the importer maps the entire chunk and returns an + // interior pointer. + CUmemGenericAllocationHandle export_handle = 0; + CUdeviceptr export_key = dptr; // fd-registry key == blob lookup key + uint64_t chunk_base = 0; + uint64_t chunk_size = 0; + uint64_t chunk_generation = 0; + std::unique_lock preallocated_chunk_lock; + if (found) { + export_handle = metadata.handle; + if (export_handle == 0 && metadata.from_preallocation) { + // Keep the chunk alive and its record stable until the exported fd is + // parked below. free_preallocated_region takes the same mutex. + preallocated_chunk_lock = std::unique_lock(preallocated_chunks_mutex); + auto chunk_it = preallocated_chunks.upper_bound(dptr); + if (chunk_it != preallocated_chunks.begin()) { + --chunk_it; + CUdeviceptr candidate_base = chunk_it->first; + if (dptr >= candidate_base && (uint64_t)(dptr - candidate_base) < chunk_it->second.size) { + chunk_base = (uint64_t)candidate_base; + chunk_size = chunk_it->second.size; + chunk_generation = chunk_it->second.generation; + export_handle = chunk_it->second.handle; + export_key = candidate_base; + } + } + if (export_handle == 0) { + fprintf(stderr, + "[HOOK] ERROR: cuIpcGetMemHandle: carved ptr=0x%llx is outside all live " + "preallocated chunks\n", + (unsigned long long)dptr); + } + } + } + + if (export_handle != 0) { // This is a VMM allocation - export via shareable handle typedef CUresult (*cuMemExportToShareableHandle_t)( void*, CUmemGenericAllocationHandle, CUmemAllocationHandleType, unsigned long long); @@ -2890,32 +3258,81 @@ CUresult cuIpcGetMemHandle(CUipcMemHandle* pHandle, CUdeviceptr dptr) { return CUDA_ERROR_NOT_SUPPORTED; } - int fd = -1; - CUresult result = - export_func(&fd, metadata.handle, CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR, 0); + // The server also owns the per-process token packed into the blob, so it + // must exist before we hand out any handle. + if (!vmm_ipc_ensure_server()) { + return CUDA_ERROR_NOT_SUPPORTED; + } - if (result != CUDA_SUCCESS) { - fprintf(stderr, - "[HOOK] ERROR: cuMemExportToShareableHandle failed with error %d for ptr=0x%llx\n", - result, (unsigned long long)dptr); - return result; + // Export once per allocation and park the fd for the server thread. + // The cached entry is keyed by VA and validated against the allocation + // handle, so VA reuse after free/realloc re-exports instead of serving a + // stale fd (the free path also eagerly invalidates via + // vmm_ipc_invalidate_export). + int fd = -1; + vmm_ipc_exported_fds.visit(export_key, + [&](const std::pair& kv) { + if (kv.second.handle == export_handle) + fd = kv.second.fd; + }); + if (fd < 0) { + int new_fd = -1; + CUresult result = + export_func(&new_fd, export_handle, CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR, 0); + if (result != CUDA_SUCCESS) { + fprintf(stderr, + "[HOOK] ERROR: cuMemExportToShareableHandle failed with error %d for ptr=0x%llx\n", + result, (unsigned long long)dptr); + return result; + } + // Keep parked fds out of exec'd children. Plain fork still inherits + // them; vmm_ipc_ensure_server closes the inherited registry before the + // child starts its own server. SCM_RIGHTS transfer is unaffected. + fcntl(new_fd, F_SETFD, FD_CLOEXEC); + vmm_ipc_exported_fds.insert_or_visit( + std::make_pair(export_key, VmmIpcExportedFd{export_handle, new_fd, chunk_generation}), + [&](std::pair& kv) { + // Entry exists: either a racing thread won (same handle - drop + // ours) or it is stale from a freed allocation (replace, refreshing + // the generation for the new export). + if (kv.second.handle == export_handle) { + close(new_fd); + new_fd = kv.second.fd; + } else { + close(kv.second.fd); + kv.second = VmmIpcExportedFd{export_handle, new_fd, chunk_generation}; + } + }); + fd = new_fd; } - // Pack our custom data into the handle structure - // CUipcMemHandle has 64 reserved bytes - we use them to store our info: - // - Magic marker (4 bytes) to identify VMM handles - // - File descriptor (4 bytes) - // - Original pointer address (8 bytes) - critical for CUDA graph replay! - // - Size (8 bytes) + // Pack our custom data into the handle structure. + // CUipcMemHandle has 64 reserved bytes: + // - Magic marker (4 bytes, "VMM2") + // - Exporter pid (4 bytes) - importer fetches the fd from this process's + // VMM-IPC socket via SCM_RIGHTS (a raw fd integer is not portable) + // - Original pointer address (8 bytes) - fd lookup key + placement hint + // - Aligned size (8 bytes) + // - Per-process token (8 bytes) - server rejects mismatches (pid reuse) + // - Chunk base + chunk size (8+8 bytes) - nonzero iff the pointer is a + // carve from the preallocated chunk; the importer then maps the whole + // chunk and returns base + (ptr - chunk_base) + // - Chunk generation (8 bytes) - distinguishes free/recreate at one base memset(pHandle, 0, sizeof(CUipcMemHandle)); + uint32_t pid_u32 = (uint32_t)getpid(); memcpy(pHandle->reserved, &VMM_IPC_MAGIC, sizeof(uint32_t)); - memcpy(pHandle->reserved + 4, &fd, sizeof(int)); + memcpy(pHandle->reserved + 4, &pid_u32, sizeof(uint32_t)); memcpy(pHandle->reserved + 8, &dptr, sizeof(CUdeviceptr)); memcpy(pHandle->reserved + 16, &metadata.size, sizeof(size_t)); + memcpy(pHandle->reserved + 24, &vmm_ipc_token, sizeof(uint64_t)); + memcpy(pHandle->reserved + 32, &chunk_base, sizeof(uint64_t)); + memcpy(pHandle->reserved + 40, &chunk_size, sizeof(uint64_t)); + memcpy(pHandle->reserved + 48, &chunk_generation, sizeof(uint64_t)); #ifdef HOOK_DEBUG - fprintf(stderr, "[HOOK] cuIpcGetMemHandle: VMM ptr=0x%llx, fd=%d, size=%zu\n", + fprintf(stderr, + "[HOOK] cuIpcGetMemHandle: VMM ptr=0x%llx, fd=%d (served via socket), size=%zu\n", (unsigned long long)dptr, fd, metadata.size); #endif return CUDA_SUCCESS; @@ -2940,19 +3357,127 @@ CUresult cuIpcOpenMemHandle(CUdeviceptr* pdptr, CUipcMemHandle handle, unsigned if (magic == VMM_IPC_MAGIC) { // This is a VMM IPC handle - extract the packed data - int fd; + uint32_t exporter_pid_u32; CUdeviceptr original_ptr; size_t size; + uint64_t token; + uint64_t chunk_base; + uint64_t chunk_size; + uint64_t chunk_generation; - memcpy(&fd, handle.reserved + 4, sizeof(int)); + memcpy(&exporter_pid_u32, handle.reserved + 4, sizeof(uint32_t)); memcpy(&original_ptr, handle.reserved + 8, sizeof(CUdeviceptr)); memcpy(&size, handle.reserved + 16, sizeof(size_t)); + memcpy(&token, handle.reserved + 24, sizeof(uint64_t)); + memcpy(&chunk_base, handle.reserved + 32, sizeof(uint64_t)); + memcpy(&chunk_size, handle.reserved + 40, sizeof(uint64_t)); + memcpy(&chunk_generation, handle.reserved + 48, sizeof(uint64_t)); + pid_t exporter_pid = (pid_t)exporter_pid_u32; + + // The descriptor came from a peer process, so validate the chunk geometry + // before deriving any interior pointer. chunk_base/chunk_size are nonzero + // together iff this is a carve from the preallocated chunk; a mismatch is + // corrupt metadata. All range math below is overflow-safe because the + // values are untrusted. Rejecting here happens before any fd fetch, + // import, reservation, or mapping, so there is nothing to clean up. + const bool has_chunk_base = (chunk_base != 0); + const bool has_chunk_size = (chunk_size != 0); + if (has_chunk_base != has_chunk_size) { + fprintf(stderr, + "[HOOK] ERROR: VMM-IPC handle has inconsistent chunk metadata: " + "chunk_base=0x%llx chunk_size=%llu\n", + (unsigned long long)chunk_base, (unsigned long long)chunk_size); + return CUDA_ERROR_INVALID_VALUE; + } + if (has_chunk_base) { + const uint64_t carve_ptr = (uint64_t)original_ptr; + const uint64_t carve_size = (uint64_t)size; + if (carve_ptr == 0 || carve_size == 0) { + fprintf(stderr, + "[HOOK] ERROR: VMM-IPC chunk handle has zero pointer/size: " + "original_ptr=0x%llx size=%llu\n", + (unsigned long long)carve_ptr, (unsigned long long)carve_size); + return CUDA_ERROR_INVALID_VALUE; + } + // The carve must start at or after the chunk base and end at or before + // the chunk end, with neither endpoint sum wrapping around 2^64. + if (carve_ptr < chunk_base || chunk_size > (UINT64_MAX - chunk_base) || + carve_size > (UINT64_MAX - carve_ptr) || + (carve_ptr + carve_size) > (chunk_base + chunk_size)) { + fprintf(stderr, + "[HOOK] ERROR: VMM-IPC carve [0x%llx,+0x%llx) is not contained in " + "chunk [0x%llx,+0x%llx)\n", + (unsigned long long)carve_ptr, (unsigned long long)carve_size, + (unsigned long long)chunk_base, (unsigned long long)chunk_size); + return CUDA_ERROR_INVALID_VALUE; + } + } + + // For chunk carves the fd registry on the exporter is keyed by the chunk + // base, and what we import/map is the whole chunk. + CUdeviceptr fetch_key = (chunk_base != 0) ? (CUdeviceptr)chunk_base : original_ptr; + size_t map_size = (chunk_base != 0) ? (size_t)chunk_size : size; + + if (chunk_base != 0) { + // Fast path: peer chunk already mapped -> hand out an interior pointer. + std::lock_guard lock(vmm_ipc_chunk_mutex); + VmmIpcChunkKey key = std::make_tuple(exporter_pid, token, chunk_base, chunk_generation); + auto it = vmm_ipc_chunk_mappings.find(key); + if (it != vmm_ipc_chunk_mappings.end()) { + // Cache hit. The incoming handle's chunk_size/carve were only checked + // against the handle's own (untrusted) metadata above. Re-validate the + // request against the ACTUAL cached mapping before deriving an interior + // pointer, and reject before touching the refcount so all bookkeeping is + // preserved on the error path. + if (!vmm_ipc_chunk_request_matches_mapping(chunk_base, chunk_size, original_ptr, size, + it->second, "cache hit")) { + return CUDA_ERROR_INVALID_VALUE; + } + it->second.refcount++; + CUdeviceptr mapped = it->second.local_base + (original_ptr - (CUdeviceptr)chunk_base); + auto sub_it = vmm_ipc_chunk_subptrs.find(mapped); + if (sub_it == vmm_ipc_chunk_subptrs.end()) { + vmm_ipc_chunk_subptrs[mapped] = VmmIpcChunkSubptr{key, 1}; + } else if (sub_it->second.key == key) { + sub_it->second.open_count++; + } else { + it->second.refcount--; + fprintf(stderr, + "[HOOK] ERROR: VMM-IPC interior pointer collision at 0x%llx across chunks\n", + (unsigned long long)mapped); + return CUDA_ERROR_INVALID_VALUE; + } + *pdptr = mapped; + return CUDA_SUCCESS; + } + } #ifdef HOOK_DEBUG - fprintf(stderr, "[HOOK] cuIpcOpenMemHandle: VMM fd=%d, original_ptr=0x%llx, size=%zu\n", fd, - (unsigned long long)original_ptr, size); + fprintf(stderr, + "[HOOK] cuIpcOpenMemHandle: VMM exporter_pid=%d, original_ptr=0x%llx, size=%zu\n", + (int)exporter_pid, (unsigned long long)original_ptr, size); #endif + // Obtain a live fd in THIS process: dup from our own registry for + // same-process opens, SCM_RIGHTS fetch from the exporter otherwise. + int fd = -1; + if (exporter_pid == getpid() && token == vmm_ipc_token) { + // Same-process import: mirror the server's generation check so a stale + // generation is rejected here too. + vmm_ipc_exported_fds.visit(fetch_key, + [&](const std::pair& kv) { + if (kv.second.generation == chunk_generation) + fd = fcntl(kv.second.fd, F_DUPFD_CLOEXEC, 0); + }); + } else { + fd = vmm_ipc_fetch_fd(exporter_pid, token, chunk_generation, fetch_key); + } + if (fd < 0) { + fprintf(stderr, "[HOOK] ERROR: VMM-IPC could not obtain fd for ptr=0x%llx from pid %d\n", + (unsigned long long)fetch_key, (int)exporter_pid); + return CUDA_ERROR_INVALID_VALUE; + } + // Import the allocation handle from file descriptor typedef CUresult (*cuMemImportFromShareableHandle_t)(CUmemGenericAllocationHandle*, void*, CUmemAllocationHandleType); @@ -2961,12 +3486,14 @@ CUresult cuIpcOpenMemHandle(CUdeviceptr* pdptr, CUipcMemHandle handle, unsigned if (import_func == nullptr) { fprintf(stderr, "[HOOK] ERROR: cuMemImportFromShareableHandle not found\n"); + close(fd); return CUDA_ERROR_NOT_SUPPORTED; } CUmemGenericAllocationHandle imported_handle; CUresult result = import_func(&imported_handle, (void*)(intptr_t)fd, CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR); + close(fd); // the driver does not take ownership of our fd copy if (result != CUDA_SUCCESS) { fprintf(stderr, "[HOOK] ERROR: cuMemImportFromShareableHandle failed with error %d\n", @@ -2974,16 +3501,58 @@ CUresult cuIpcOpenMemHandle(CUdeviceptr* pdptr, CUipcMemHandle handle, unsigned return result; } - // Map to the SAME address as original (critical for CUDA graph replay!) + typedef CUresult (*cuMemAddressReserve_t)(CUdeviceptr*, size_t, size_t, CUdeviceptr, + unsigned long long); + auto real_reserve_func = (cuMemAddressReserve_t)CUDA_DRIVER_CALL( + cuda_driver_entry_table, CUDA_ENTRY_cuMemAddressReserve); + typedef CUresult (*cuMemRelease_t)(CUmemGenericAllocationHandle); + auto mem_release_func = + (cuMemRelease_t)CUDA_DRIVER_CALL(cuda_driver_entry_table, CUDA_ENTRY_cuMemRelease); + typedef CUresult (*cuMemAddressFree_t)(CUdeviceptr, size_t); + auto real_address_free_func = + (cuMemAddressFree_t)CUDA_DRIVER_CALL(cuda_driver_entry_table, CUDA_ENTRY_cuMemAddressFree); + + // Reserve our own VA range for the peer mapping (real driver call, NOT the + // hooked wrapper - peer imports must never advance the deterministic + // cursor). The exporter's VA is passed as a hint: with per-rank disjoint + // regions it is honored and the mapping is address-stable; with today's + // shared-base symmetric layouts that range is occupied by our own region + // reservation, the driver picks elsewhere, and the caller gets a valid but + // relocated peer pointer. That is correct for table-indirect consumers + // (DeepEP buffer_ptrs_gpu, custom-allreduce RankData contents): peer VAs + // are never baked into captured graph nodes on those paths. + // First preference: the exporter's own VA (address-stable whenever it is + // locally free, e.g. future per-rank disjoint-base configs). + CUdeviceptr mapped_ptr = 0; + result = real_reserve_func(&mapped_ptr, map_size, 0, fetch_key, 0); + if (result == CUDA_SUCCESS && mapped_ptr != fetch_key) { + // Hint not honored. Do NOT keep the driver's fallback choice - it tends + // to sit immediately after existing reservations, i.e. on the NVSHMEM + // heap hint. Re-reserve inside the dedicated import zone. + real_address_free_func(mapped_ptr, map_size); + uint64_t zone_hint = + vmm_ipc_import_hint.fetch_add((map_size + ((1ULL << 30) - 1)) & ~((1ULL << 30) - 1)); + mapped_ptr = 0; + result = real_reserve_func(&mapped_ptr, map_size, 0, (CUdeviceptr)zone_hint, 0); + } + if (result != CUDA_SUCCESS) { + fprintf(stderr, "[HOOK] ERROR: cuMemAddressReserve for IPC import failed with error %d\n", + result); + mem_release_func(imported_handle); + return result; + } + typedef CUresult (*cuMemMap_t)(CUdeviceptr, size_t, size_t, CUmemGenericAllocationHandle, unsigned long long); auto map_func = (cuMemMap_t)CUDA_DRIVER_CALL(cuda_driver_entry_table, CUDA_ENTRY_cuMemMap); - result = map_func(original_ptr, size, 0, imported_handle, 0); + result = map_func(mapped_ptr, map_size, 0, imported_handle, 0); if (result != CUDA_SUCCESS) { fprintf(stderr, "[HOOK] ERROR: cuMemMap for IPC failed with error %d at addr=0x%llx size=%zu\n", - result, (unsigned long long)original_ptr, size); + result, (unsigned long long)mapped_ptr, map_size); + real_address_free_func(mapped_ptr, map_size); + mem_release_func(imported_handle); return result; } @@ -3003,27 +3572,96 @@ CUresult cuIpcOpenMemHandle(CUdeviceptr* pdptr, CUipcMemHandle handle, unsigned auto set_access_func = (cuMemSetAccess_t)CUDA_DRIVER_CALL(cuda_driver_entry_table, CUDA_ENTRY_cuMemSetAccess); - result = set_access_func(original_ptr, size, &accessDesc, 1); + result = set_access_func(mapped_ptr, map_size, &accessDesc, 1); if (result != CUDA_SUCCESS) { fprintf(stderr, "[HOOK] ERROR: cuMemSetAccess for IPC failed with error %d\n", result); + typedef CUresult (*cuMemUnmap_t)(CUdeviceptr, size_t); + auto mem_unmap_func = + (cuMemUnmap_t)CUDA_DRIVER_CALL(cuda_driver_entry_table, CUDA_ENTRY_cuMemUnmap); + mem_unmap_func(mapped_ptr, map_size); + real_address_free_func(mapped_ptr, map_size); + mem_release_func(imported_handle); return result; } - *pdptr = original_ptr; + if (mapped_ptr != fetch_key) { + fprintf(stderr, + "[HOOK] INFO: VMM-IPC import relocated: exporter VA 0x%llx -> local VA 0x%llx " + "(expected with shared region bases; peer tables must be refreshed by the caller)\n", + (unsigned long long)fetch_key, (unsigned long long)mapped_ptr); + } + + if (chunk_base != 0) { + // Register the whole-chunk mapping and hand out the interior pointer. + CUdeviceptr interior = mapped_ptr + (original_ptr - (CUdeviceptr)chunk_base); + std::lock_guard lock(vmm_ipc_chunk_mutex); + VmmIpcChunkKey key = std::make_tuple(exporter_pid, token, chunk_base, chunk_generation); + auto it = vmm_ipc_chunk_mappings.find(key); + if (it != vmm_ipc_chunk_mappings.end()) { + // Lost a race to a concurrent open of the same peer chunk: keep the + // winner's mapping, drop ours. Release the losing temporary mapping and + // handle first (unconditionally, whether or not we adopt the winner). + typedef CUresult (*cuMemUnmap_t)(CUdeviceptr, size_t); + auto mem_unmap_func = + (cuMemUnmap_t)CUDA_DRIVER_CALL(cuda_driver_entry_table, CUDA_ENTRY_cuMemUnmap); + mem_unmap_func(mapped_ptr, map_size); + real_address_free_func(mapped_ptr, map_size); + mem_release_func(imported_handle); + // Same authoritative re-validation as the cache-hit fast path: the key + // omits chunk_size, so this same-key handle may not match the winner's + // live mapping. Reject before adopting it - our temporary is already + // released and the winner's refcount is untouched. + if (!vmm_ipc_chunk_request_matches_mapping(chunk_base, chunk_size, original_ptr, size, + it->second, "insertion race")) { + return CUDA_ERROR_INVALID_VALUE; + } + it->second.refcount++; + interior = it->second.local_base + (original_ptr - (CUdeviceptr)chunk_base); + } else { + vmm_ipc_chunk_mappings[key] = VmmIpcChunkMapping{mapped_ptr, map_size, imported_handle, 1}; + } + auto sub_it = vmm_ipc_chunk_subptrs.find(interior); + if (sub_it == vmm_ipc_chunk_subptrs.end()) { + vmm_ipc_chunk_subptrs[interior] = VmmIpcChunkSubptr{key, 1}; + } else if (sub_it->second.key == key) { + sub_it->second.open_count++; + } else { + auto mapping_it = vmm_ipc_chunk_mappings.find(key); + if (mapping_it != vmm_ipc_chunk_mappings.end() && --mapping_it->second.refcount == 0) { + typedef CUresult (*cuMemUnmap_t)(CUdeviceptr, size_t); + auto mem_unmap_func = + (cuMemUnmap_t)CUDA_DRIVER_CALL(cuda_driver_entry_table, CUDA_ENTRY_cuMemUnmap); + mem_unmap_func(mapping_it->second.local_base, mapping_it->second.size); + real_address_free_func(mapping_it->second.local_base, mapping_it->second.size); + mem_release_func(mapping_it->second.handle); + vmm_ipc_chunk_mappings.erase(mapping_it); + } + fprintf(stderr, + "[HOOK] ERROR: VMM-IPC interior pointer collision at 0x%llx across chunks\n", + (unsigned long long)interior); + return CUDA_ERROR_INVALID_VALUE; + } + *pdptr = interior; + return CUDA_SUCCESS; + } + + *pdptr = mapped_ptr; - // Track this imported allocation + // Track this imported allocation (close path unmaps, releases, and frees + // the reservation we created above) AllocMetadata alloc_metadata; - alloc_metadata.ptr = original_ptr; + alloc_metadata.ptr = mapped_ptr; alloc_metadata.size = size; alloc_metadata.handle = imported_handle; alloc_metadata.region_base = 0; // region_base == 0 indicates IPC-imported alloc_metadata.from_preallocation = false; - global_alloc_metadata.emplace(original_ptr, alloc_metadata); + global_alloc_metadata.emplace(mapped_ptr, alloc_metadata); #ifdef HOOK_DEBUG fprintf(stderr, - "[HOOK] cuIpcOpenMemHandle: Successfully mapped VMM fd=%d -> ptr=0x%llx, size=%zu\n", - fd, (unsigned long long)*pdptr, size); + "[HOOK] cuIpcOpenMemHandle: Successfully mapped exporter ptr=0x%llx -> ptr=0x%llx, " + "size=%zu\n", + (unsigned long long)original_ptr, (unsigned long long)*pdptr, size); #endif return CUDA_SUCCESS; } @@ -3041,6 +3679,36 @@ CUresult cuIpcCloseMemHandle(CUdeviceptr dptr) { auto real_func = (cuIpcCloseMemHandle_t)CUDA_DRIVER_CALL(cuda_driver_entry_table, CUDA_ENTRY_cuIpcCloseMemHandle); + // Interior pointer into an imported peer chunk? Refcounted: the chunk + // mapping is torn down when its last interior pointer closes. + { + std::lock_guard lock(vmm_ipc_chunk_mutex); + auto sub_it = vmm_ipc_chunk_subptrs.find(dptr); + if (sub_it != vmm_ipc_chunk_subptrs.end()) { + VmmIpcChunkKey key = sub_it->second.key; + if (--sub_it->second.open_count == 0) { + vmm_ipc_chunk_subptrs.erase(sub_it); + } + auto it = vmm_ipc_chunk_mappings.find(key); + if (it != vmm_ipc_chunk_mappings.end() && --it->second.refcount == 0) { + typedef CUresult (*cuMemUnmap_t)(CUdeviceptr, size_t); + auto mem_unmap_func = + (cuMemUnmap_t)CUDA_DRIVER_CALL(cuda_driver_entry_table, CUDA_ENTRY_cuMemUnmap); + typedef CUresult (*cuMemRelease_t)(CUmemGenericAllocationHandle); + auto mem_release_func = + (cuMemRelease_t)CUDA_DRIVER_CALL(cuda_driver_entry_table, CUDA_ENTRY_cuMemRelease); + typedef CUresult (*cuMemAddressFree_t)(CUdeviceptr, size_t); + auto addr_free_func = (cuMemAddressFree_t)CUDA_DRIVER_CALL(cuda_driver_entry_table, + CUDA_ENTRY_cuMemAddressFree); + mem_unmap_func(it->second.local_base, it->second.size); + mem_release_func(it->second.handle); + addr_free_func(it->second.local_base, it->second.size); + vmm_ipc_chunk_mappings.erase(it); + } + return CUDA_SUCCESS; + } + } + AllocMetadata metadata; bool found = false; @@ -3059,8 +3727,14 @@ CUresult cuIpcCloseMemHandle(CUdeviceptr dptr) { auto mem_release_func = (cuMemRelease_t)CUDA_DRIVER_CALL(cuda_driver_entry_table, CUDA_ENTRY_cuMemRelease); + typedef CUresult (*cuMemAddressFree_t)(CUdeviceptr, size_t); + auto real_address_free_func = + (cuMemAddressFree_t)CUDA_DRIVER_CALL(cuda_driver_entry_table, CUDA_ENTRY_cuMemAddressFree); + mem_unmap_func(dptr, metadata.size); mem_release_func(metadata.handle); + // The import path owns a dedicated VA reservation for this mapping. + real_address_free_func(dptr, metadata.size); global_alloc_metadata.erase_if( [dptr](const std::pair& kv) { return kv.first == dptr; }); @@ -3268,6 +3942,11 @@ bool preallocate_region(size_t size) { tls_storage.preallocated_start_addr = target_addr; tls_storage.preallocated_end_addr = target_addr + aligned_size; tls_storage.has_preallocation = true; + { + std::lock_guard lock(preallocated_chunks_mutex); + uint64_t generation = next_preallocated_chunk_generation.fetch_add(1); + preallocated_chunks[target_addr] = PreallocatedChunk{aligned_size, allocHandle, generation}; + } // Note: we do NOT advance current_alloc_base_addr here. // The alloc calls will advance it as they consume the preallocated memory. @@ -3290,6 +3969,12 @@ void free_preallocated_region() { size_t preallocated_size = tls_storage.preallocated_end_addr - tls_storage.preallocated_start_addr; + { + std::lock_guard lock(preallocated_chunks_mutex); + preallocated_chunks.erase(tls_storage.preallocated_start_addr); + } + vmm_ipc_invalidate_export((CUdeviceptr)tls_storage.preallocated_start_addr); + mem_unmap_func(tls_storage.preallocated_start_addr, preallocated_size); mem_release_func(tls_storage.preallocated_handle); @@ -3464,7 +4149,10 @@ void replay_hook_events_from_json(const boost::json::object& events_obj) { size_t alignment = event_obj.at("alignment").to_number(); uint64_t expected_ptr = event_obj.at("ptr").to_number(); - CUdeviceptr aligned_addr = align_to(tls_storage.current_alloc_base_addr, alignment); + // Archives recorded before alignment normalization stored the caller's + // alignment (including CUDA's default-alignment sentinel, zero). + const size_t effective_alignment = alignment == 0 ? kAllocAlignment : alignment; + CUdeviceptr aligned_addr = align_to(tls_storage.current_alloc_base_addr, effective_alignment); if (aligned_addr != expected_ptr) { fprintf(stderr, "[REPLAY] FATAL: Reserve address mismatch!\n"); diff --git a/docs/sglang/hooks.md b/docs/sglang/hooks.md index cf5b4c82..10360fd0 100644 --- a/docs/sglang/hooks.md +++ b/docs/sglang/hooks.md @@ -10,8 +10,10 @@ In install order: 2. `_patch_init_memory_pool` 3. `_patch_load_model` 4. `_patch_kernel_warmup` -5. `_patch_cuda_graph_capture` -6. `_patch_spawn_sites` +5. `_patch_attn_tp_logits_gather` +6. `_patch_cuda_graph_capture` +7. `_patch_spawn_sites` +8. `_patch_deepep` > Removed: an earlier `_patch_model_runner_init` wrapped `ModelRunner.__init__` to stash `dp_rank` on `self` (upstream didn't expose it as an attribute). Upstream sglang now does `self.dp_rank = dp_rank` in the constructor, so our wrapper is no longer needed; `_resolve_dp_rank` reads `self.dp_rank` directly. @@ -71,19 +73,45 @@ No-op on SAVE/LOAD: The FlashInfer autotune path inside is shut off both by this no-op and by the forced `disable_flashinfer_autotune` flag. -### 5. `CudaGraphRunner` (the largest patch) +### 5. `attn_tp_all_gather_into_tensor` / `GroupCoordinator._all_gather_into_tensor` + +**LOAD only.** `_patch_attn_tp_logits_gather` pre-copies each rank's local logits +shard into the gather output buffer before the recognized TP all-gather seams +run. On LOAD the captured graph replays a gather whose peer inputs are not +re-executed, so without seeding the local shard the rank's own slice of the +gathered logits would be stale. Two seams are wrapped when present: + +- `sglang.srt.layers.logits_processor.attn_tp_all_gather_into_tensor` (DP/EP + attention-TP gather). Requires `get_attention_tp_group()` to expose integer + `rank_in_group` / `world_size`; validates the `(world_size, *input.shape)` + output layout and fails fast on invalid rank bounds. +- `GroupCoordinator._all_gather_into_tensor` (regular TP gather). Pre-copies the + local shard only for recognized rank-major `(world_size, *input.shape)` and + concatenated `(world_size*input.shape[0], *input.shape[1:])` layouts via + `_regular_tp_local_shard`; any other caller or unrecognized/unviewable layout + is delegated to the original collective unchanged, preserving its return value + and exceptions. Invalid rank bounds still raise `RuntimeError`. + +`NONE`/`SAVE` modes are untouched. The wrappers are idempotent +(`_foundry_load_local_shard_precopy` guard), and the installed seams are logged: + +```text +[Foundry] SGLang LOAD logits pre-copy patches installed: attention_tp, regular_tp +``` + +### 6. `CudaGraphRunner` (the largest patch) Four sub-patches: -#### 5a. `_create_device_graph` +#### 6a. `_create_device_graph` On SAVE returns `foundry.CUDAGraph()` instead of `torch.cuda.CUDAGraph()`. -#### 5b. `_capture_graph` +#### 6b. `_capture_graph` On SAVE enters `foundry.graph(graph, pool=pool, stream=stream)` instead of `torch.cuda.graph(...)`, captures the `run_once` callable's allocations into foundry's hook event log. -#### 5c. `capture_one_batch_size` +#### 6c. `capture_one_batch_size` On SAVE wraps the `forward` callable with a 3-call counter: @@ -107,7 +135,7 @@ key = bs if stream_idx is None else f"{stream_idx}_{bs}" (Sglang removed the `_make_graph_key` / `get_capture_lora_variant` helpers in commit `ce2506e1c`, the same commit that deprecated `record_nolora_graph` dual MoE graph capture.) -#### 5d. `capture` +#### 6d. `capture` The outermost replacement. @@ -152,7 +180,7 @@ self.output_buffers = {k: v[1] for k, v in state.loaded_graphs.items()} `load_all_graphs` calls `start_graph_builds(all_paths)` and `finish_graph_loads(pending)` exactly once each. This is required for the manifest's template + on-demand linking to work; per-graph `start_graph_builds([single_path])` calls would leave on-demand graphs without a `shared_exec`, and runtime replay would abort with `Called CUDAGraph::replay without a preceding successful capture or load`. -### 6. Spawn-site patches +### 7. Spawn-site patches Two parent-side wrappers: @@ -177,6 +205,32 @@ def patched_start(self, *args, **kwargs): Active only when `moe_a2a_backend == deepep`. EP runs DP-attention + DeepEP (NCCL-free); TP attention is unsupported (its NCCL all-reduce is incompatible with the VMM region). +- **DeepEP transport policy** (`_patch_deepep`, hooks). `deepep_transport` + accepts only `fabric` and `nvl_ipc`; omitting it selects the `fabric` default. + The integration wraps `deep_ep.Buffer.__init__` and applies the TOML policy: + `fabric` sets `use_fabric=True` and forces `num_nvl_bytes=0`, while + `nvl_ipc` sets `use_fabric=False` and preserves SGLang's computed nonzero + `num_nvl_bytes`. The latter is intranode-only and requires all ranks to share + the needed NVLink connectivity. +- **Required DeepEP API.** Foundry requires a DeepEP release whose + `Buffer(use_fabric=...)` API accepts the `use_fabric` keyword. DeepEP + `29d31c0` or newer provides that API while retaining SGLang's positional + constructor shape. Foundry raises an actionable error before calling an + older constructor; there is no silent fallback to another transport. +- **DeepEP and NVSHMEM remain required.** `nvl_ipc` changes only the NVLink + buffer. DeepEP's RDMA buffer still uses the NVSHMEM symmetric heap. The + no-IMEX recipe sets `NVSHMEM_CUMEM_HANDLE_TYPE=FILE_DESCRIPTOR`, + `NCCL_CUMEM_ENABLE=0`, and `NCCL_NVLS_ENABLE=0` for matching SAVE and LOAD + runs. +- **Archive parity.** SAVE records the effective transport and NVLink buffer + size. LOAD validates the transport before loading CUDA modules or restoring + graphs, then validates the buffer size when constructing DeepEP. A mismatched + LOAD config fails fast with: + + ```text + Foundry DeepEP transport mismatch: archived=nvl_ipc, configured=fabric + ``` + - **DeepEP buffer pre-capture bootstrap** (`bootstrap_deepep_buffer`, graph_ops). sglang creates the singleton NVSHMEM `Buffer` lazily on the first MoE dispatch — normally during the warmup forwards foundry suppresses, which would push creation *inside* the @@ -191,7 +245,7 @@ TP attention is unsupported (its NCCL all-reduce is incompatible with the VMM re - **`deepep_adapter` mode on LOAD.** LOAD replaces the capture loop, so the adapter's `_captured_deepep_mode` is never set; replay asserts on it. The hook calls `deepep_adapter.capture(is_extend_in_batch=False)` after load. -- **FlashInfer pre-pass gated to FlashInfer.** The §5d pre-pass + `reuse_pre_pass_init` +- **FlashInfer pre-pass gated to FlashInfer.** The §6d pre-pass + `reuse_pre_pass_init` shim handle FlashInfer's per-bs wrappers. fa3 (`FlashAttentionBackend`) uses a single fixed `init_cuda_graph_state` workspace, so the shim is skipped (detected via absence of `indices_updater_decode`); but fa3's per-bs `decode_cuda_graph_metadata[bs]` is still @@ -203,7 +257,26 @@ TP attention is unsupported (its NCCL all-reduce is incompatible with the VMM re surfaced only on sglang EP. See [`../../recipe/sglang/README.md`](../../recipe/sglang/README.md) for the EP serve -config and required kernel versions (deep_ep `9af0e0d`, sgl-deep-gemm ≥0.1.2, fa3). +config and required kernel versions (DeepEP `29d31c0` or newer, +sgl-deep-gemm ≥0.1.2, fa3). + +### No-IMEX validation + +On a two-GPU host without IMEX, run the standalone RDMA-only, NVLink IPC, and +LOAD-preallocation cases: + +```bash +bash tests/run_deepep_matrix.sh ll_nofabric both +bash tests/run_deepep_matrix.sh nvl_ipc both +bash tests/run_deepep_matrix.sh nvl_ipc_prealloc both +``` + +Then run the focused SGLang EP recipe: + +```bash +bash recipe/experimental/serve_qwen3-30ba3bfp8_sglang_ipc_ep.sh 2 --save +bash recipe/experimental/serve_qwen3-30ba3bfp8_sglang_ipc_ep.sh 2 --load +``` ## Patch idiom @@ -211,4 +284,6 @@ All patches use the `wrap-and-call` idiom — short-circuit on `mode == NONE`, r ## Install order -`install_hooks` calls the patch helpers in the listed order. Order doesn't matter for correctness here — every patch attaches to a different attribute — but spawn sites go last so the install-completion log line appears after every other patch has registered. +`install_hooks` calls the patch helpers in the listed order. Order does not +matter for correctness here because every patch attaches to a different +attribute. diff --git a/docs/superpowers/plans/2026-07-23-no-imex-deepep.md b/docs/superpowers/plans/2026-07-23-no-imex-deepep.md new file mode 100644 index 00000000..a0d5b89d --- /dev/null +++ b/docs/superpowers/plans/2026-07-23-no-imex-deepep.md @@ -0,0 +1,963 @@ +# No-IMEX Intranode DeepEP Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Support fresh-process Foundry SAVE/LOAD for intranode vLLM and SGLang DeepEP workloads without NVIDIA IMEX fabric. + +**Architecture:** Add a shared `fabric`/`nvl_ipc` transport contract consumed by both engine integrations and persisted in each archive. Upgrade the CUDA hook from raw-FD VMM IPC v1 to authenticated SCM_RIGHTS VMM IPC v2 so DeepEP's NVLink buffer can be imported by a fresh LOAD process, including buffers carved from Foundry's preallocated VMM chunk. + +**Tech Stack:** Python 3.12, C++17, CUDA Driver API/VMM, PyTorch 2.11/cu130, DeepEP/NVSHMEM, vLLM, SGLang, pytest, Ruff, CMake/Ninja, Modal H100. + +## Global Constraints + +- Scope is intranode expert parallelism for both vLLM and SGLang; multi-host EP and tensor parallelism are out of scope. +- `deepep_transport` accepts exactly `fabric` and `nvl_ipc`; omitted means `fabric`. +- Existing archives without DeepEP fields are interpreted as `fabric`. +- LOAD must reject a transport or effective `num_nvl_bytes` mismatch; it must never silently fall back. +- `nvl_ipc` still uses DeepEP/NVSHMEM for the RDMA buffer and requires `NCCL_CUMEM_ENABLE=0` and `NCCL_NVLS_ENABLE=0`. +- Imports stay at module scope unless an optional engine dependency requires a guarded import and the reason is documented. +- CUDA tests run on Modal H100 with `--env rahul-dev`, `CC=gcc`, and `CXX=g++`. +- CUDA orchestration parents must not initialize CUDA before spawning children that preload `libcuda_hook.so`. +- Every task is committed and pushed independently. + +--- + +## File Map + +### Create + +- [`python/foundry/integration/deepep_transport.py`](../../../python/foundry/integration/deepep_transport.py) — shared transport enum, kwargs policy, and archive validation. +- [`tests/test_deepep_transport.py`](../../../tests/test_deepep_transport.py) — CPU-only shared contract and SGLang config tests. +- [`tests/test_deepep_transport_vllm.py`](../../../tests/test_deepep_transport_vllm.py) — vLLM config/runtime tests in an environment with the Foundry vLLM fork. +- [`recipe/experimental/README.md`](../../../recipe/experimental/README.md) and no-IMEX recipe assets — intranode EP examples for both engines. + +### Modify + +- [`python/foundry/graph.py`](../../../python/foundry/graph.py) — pass the default graph pool explicitly. +- [`csrc/hook.cpp`](../../../csrc/hook.cpp) — VMM IPC v2, preallocated-chunk export/import, fork safety, and alignment. +- [`python/foundry/integration/vllm/config.py`](../../../python/foundry/integration/vllm/config.py) — TOML field and getter. +- [`python/foundry/integration/vllm/runtime.py`](../../../python/foundry/integration/vllm/runtime.py) — archive fields and early LOAD validation. +- [`python/foundry/integration/vllm/hooks.py`](../../../python/foundry/integration/vllm/hooks.py) — transport-aware DeepEP kwargs and profile record/validation. +- [`python/foundry/integration/sglang/config.py`](../../../python/foundry/integration/sglang/config.py) — TOML field and getter. +- [`python/foundry/integration/sglang/runtime.py`](../../../python/foundry/integration/sglang/runtime.py) — archive fields and early LOAD validation. +- [`python/foundry/integration/sglang/hooks.py`](../../../python/foundry/integration/sglang/hooks.py) — transport-aware `deep_ep.Buffer` wrapper. +- [`tests/test_deepep_fabric.py`](../../../tests/test_deepep_fabric.py) — no-fabric, NVL-buffer, and preallocated LOAD cases. +- [`tests/run_deepep_matrix.sh`](../../../tests/run_deepep_matrix.sh) — portable two-GPU matrix driver. +- [`tests/test_vmm_alloc.py`](../../../tests/test_vmm_alloc.py) — zero-alignment VMM regression. +- [`docs/vllm/moe-and-deepep.md`](../../../docs/vllm/moe-and-deepep.md) — transport modes and archive contract. +- [`docs/sglang/hooks.md`](../../../docs/sglang/hooks.md) — SGLang transport patch and DeepEP version floor. + +--- + +### Task 1: Unblock Default Graph Capture + +**Files:** +- Modify: `python/foundry/graph.py` +- Test: `tests/test_load_graph.py` +- Test: `tests/test_async_load.py` + +**Interfaces:** +- Consumes: `graph(cuda_graph, pool: _POOL_HANDLE | None = None, ...)` +- Produces: `graph.pool: tuple[_POOL_HANDLE | None]`, always containing the positional `pool` argument expected by `CUDAGraph.capture_begin`. + +- [ ] **Step 1: Reproduce the missing-pool failure on H100** + +Run: + +```bash +FOUNDRY_REF="$(git rev-parse HEAD)" \ + modal run --env rahul-dev tests/modal_native.py --suite core +``` + +If the Modal harness has not been ported yet, use the repository build/test commands from `README.md` in a two-stage Modal function and run: + +```bash +pytest tests/test_load_graph.py tests/test_async_load.py -q +``` + +Expected before the fix: capture stops because the binding requires the `pool` positional argument. + +- [ ] **Step 2: Add the regression assertion** + +Add a small CPU-safe unit around `graph.__init__` by constructing the object with `__new__`, stubbing `torch.cuda.Stream`, and asserting: + +```python +def test_graph_default_pool_is_forwarded(monkeypatch): + sentinel_stream = object() + monkeypatch.setattr(torch.cuda, "Stream", lambda: sentinel_stream) + cuda_graph = object() + + context = foundry.graph.__new__(foundry.graph) + foundry.graph.__init__(context, cuda_graph) + + assert context.pool == (None,) +``` + +Expected before the fix: `context.pool == ()`. + +- [ ] **Step 3: Pass the default pool explicitly** + +Change: + +```python +self.pool = () if pool is None else (pool,) +``` + +to: + +```python +self.pool = (pool,) +``` + +This is the focused change from source commit `57e1dba`. + +- [ ] **Step 4: Verify the focused tests** + +Run: + +```bash +pytest tests/test_load_graph.py tests/test_async_load.py -q +``` + +Expected: PASS on Modal H100. + +- [ ] **Step 5: Commit and push** + +```bash +git add python/foundry/graph.py tests/test_load_graph.py +git commit -m "fix: pass default CUDA graph pool" +git push -u origin cursor/no-imex-deepep-f10b +``` + +--- + +### Task 2: Port Hardened Cross-Process VMM IPC + +**Files:** +- Modify: `csrc/hook.cpp` +- Modify: `tests/test_deepep_fabric.py` +- Create: `tests/run_deepep_matrix.sh` +- Modify: `tests/test_vmm_alloc.py` +- Optional test harness create: `tests/modal_native.py` +- Optional test harness create: `tests/test_modal_native.py` + +**Interfaces:** +- Consumes: hooked `cuIpcGetMemHandle`, `cuIpcOpenMemHandle`, and `cuIpcCloseMemHandle`. +- Produces: VMM IPC v2 handles containing `{magic, pid, original_ptr, size, token, chunk_base, chunk_size, generation}` and SCM_RIGHTS descriptor transfer over `\0foundry-vmm-ipc.`. + +- [ ] **Step 1: Capture the failing no-fabric baseline** + +On a two-GPU host without IMEX, run the current standalone test with a nonzero NVLink buffer: + +```bash +export TEST_USE_FABRIC=0 +export TEST_NVL_BYTES_MB=64 +export NCCL_CUMEM_ENABLE=0 +export NCCL_NVLS_ENABLE=0 + +LD_PRELOAD="$NVSHMEM_HOST:$FOUNDRY_HOOK" \ + python tests/test_deepep_fabric.py --save --no-fabric --num-processes=2 +LD_PRELOAD="$NVSHMEM_HOST:$FOUNDRY_HOOK" \ + python tests/test_deepep_fabric.py --load --no-fabric --num-processes=2 +``` + +Expected before the fix: LOAD cannot import the raw process-local FD, usually with CUDA error 999. + +- [ ] **Step 2: Port the tested VMM IPC v2 core** + +Port the `csrc/hook.cpp` portions from source commit `84d2b95`: + +```bash +git show 84d2b95 -- csrc/hook.cpp +``` + +The resulting implementation must define: + +```cpp +static constexpr uint32_t VMM_IPC_MAGIC = 0x564D4D32; // "VMM2" + +struct VmmIpcRequest { + uint64_t ptr; + uint64_t token; +}; + +static bool vmm_ipc_ensure_server(); +static int vmm_ipc_fetch_fd(pid_t exporter_pid, uint64_t token, + CUdeviceptr original_ptr); +static void vmm_ipc_invalidate_export(CUdeviceptr dptr); +static void vmm_ipc_server_loop(int listen_fd); +``` + +The implementation must: + +- create an abstract Unix socket named `\0foundry-vmm-ipc.`; +- authenticate peers with `SO_PEERCRED` and require the same UID; +- generate a per-process random token from `/dev/urandom`; +- send FDs only with `SCM_RIGHTS`; +- store no raw FD number in `CUipcMemHandle`; +- use the real, uninterposed `cuMemAddressReserve` for import mappings; +- retry a colliding exporter address in the dedicated import range beginning at `0x300000000000`; +- pass non-Foundry IPC handles to the real CUDA driver. + +- [ ] **Step 3: Port fork and preallocation hardening** + +Port the `csrc/hook.cpp` changes from `eaf3fea` and `ab22a21` without the unrelated tensor-parallel files: + +```cpp +struct PreallocatedChunk { + size_t size; + CUmemGenericAllocationHandle handle; + uint64_t generation; +}; + +using VmmIpcChunkKey = + std::tuple; + +struct VmmIpcChunkMapping { + CUdeviceptr local_base; + size_t size; + CUmemGenericAllocationHandle handle; + int refcount; +}; +``` + +Required behavior: + +- track every preallocated chunk under a mutex; +- include chunk generation at handle offset 48; +- export the owning chunk plus an interior offset for carved allocations; +- cache imported chunks by PID, token, chunk base, and generation; +- reference-count interior pointers; +- clear inherited listeners and parked descriptors after fork; +- invalidate an export when the owning allocation is freed. + +- [ ] **Step 4: Port zero-alignment normalization** + +Port the focused `csrc/hook.cpp` changes from `f947886` and `0f337a1`. + +Use: + +```cpp +const size_t effective_alignment = + alignment == 0 ? kAllocAlignment : alignment; +``` + +for Foundry-region reservation math while preserving the caller-visible CUDA semantics. + +Add/port: + +```python +def test_zero_alignment_reserve_stays_in_region(): + # Reserve with alignment=0 through the driver API, then assert the + # returned address is 2 MiB aligned and remains inside the Foundry region. +``` + +from source commit `77a28e8`. + +- [ ] **Step 5: Port the portable DeepEP matrix** + +Port the non-engine test changes from `84d2b95`, `059aa42`, and `ab22a21`: + +```bash +bash tests/run_deepep_matrix.sh ll_nofabric both +bash tests/run_deepep_matrix.sh nvl_ipc both +bash tests/run_deepep_matrix.sh nvl_ipc_prealloc both +``` + +Matrix definitions: + +```bash +ll_nofabric: TEST_USE_FABRIC=0 TEST_NVL_BYTES_MB=0 +nvl_ipc: TEST_USE_FABRIC=0 TEST_NVL_BYTES_MB=64 +nvl_ipc_prealloc: TEST_USE_FABRIC=0 TEST_NVL_BYTES_MB=64 TEST_LOAD_PREALLOC_MB=1024 +``` + +Expected after the port: all three cases PASS. `nvl_fabric` remains a negative control and is expected to fail on a host without IMEX. + +- [ ] **Step 6: Build and run focused native regressions** + +Run on Modal: + +```bash +CC=gcc CXX=g++ pip install -e . --no-build-isolation +pytest tests/test_vmm_alloc.py tests/test_load_graph.py tests/test_async_load.py -q +``` + +Then run the three two-GPU matrix cases from Step 5. Expected: PASS, no CUDA error 999, and deterministic dispatch output. + +- [ ] **Step 7: Commit and push** + +```bash +git add csrc/hook.cpp tests/test_deepep_fabric.py \ + tests/run_deepep_matrix.sh tests/test_vmm_alloc.py +git commit -m "fix: support cross-process VMM IPC" +git push -u origin cursor/no-imex-deepep-f10b +``` + +--- + +### Task 3: Add the First-Class Transport Configuration + +**Files:** +- Create: `python/foundry/integration/deepep_transport.py` +- Modify: `python/foundry/integration/vllm/config.py` +- Modify: `python/foundry/integration/sglang/config.py` +- Create: `tests/test_deepep_transport.py` +- Create: `tests/test_deepep_transport_vllm.py` + +**Interfaces:** +- Produces: `DeepEPTransport`, `parse_deepep_transport`, `configure_deepep_buffer`, `validate_transport_match`, and `validate_num_nvl_bytes_match`. +- Produces: `get_deepep_transport() -> DeepEPTransport` in both engine config modules. + +- [ ] **Step 1: Write failing shared-contract tests** + +Create `tests/test_deepep_transport.py` with: + +```python +from __future__ import annotations + +import pytest + +from foundry.integration.deepep_transport import ( + DeepEPTransport, + configure_deepep_buffer, + parse_deepep_transport, + validate_num_nvl_bytes_match, + validate_transport_match, +) + + +@pytest.mark.parametrize( + ("raw", "expected"), + [ + (None, DeepEPTransport.FABRIC), + ("fabric", DeepEPTransport.FABRIC), + ("nvl_ipc", DeepEPTransport.NVL_IPC), + ], +) +def test_parse_deepep_transport(raw, expected): + assert parse_deepep_transport(raw) is expected + + +def test_parse_deepep_transport_rejects_unknown_value(): + with pytest.raises(ValueError, match="deepep_transport"): + parse_deepep_transport("rdma") + + +@pytest.mark.parametrize( + ("transport", "expected_bytes", "expected_fabric"), + [ + (DeepEPTransport.FABRIC, 0, True), + (DeepEPTransport.NVL_IPC, 128, False), + ], +) +def test_configure_deepep_buffer(transport, expected_bytes, expected_fabric): + num_nvl_bytes, kwargs = configure_deepep_buffer( + transport, + num_nvl_bytes=128, + kwargs={"other": "kept"}, + ) + assert num_nvl_bytes == expected_bytes + assert kwargs == {"other": "kept", "use_fabric": expected_fabric} + + +def test_transport_mismatch_names_both_values(): + with pytest.raises(RuntimeError, match="archived=nvl_ipc.*configured=fabric"): + validate_transport_match( + archived="nvl_ipc", + configured=DeepEPTransport.FABRIC, + ) + + +def test_nvl_bytes_mismatch_names_both_values(): + with pytest.raises(RuntimeError, match="archived=64.*effective=128"): + validate_num_nvl_bytes_match( + archived=64, + effective=128, + transport=DeepEPTransport.NVL_IPC, + ) +``` + +Expected: import failure because the shared module does not exist. + +- [ ] **Step 2: Implement the shared transport module** + +Create: + +```python +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the Foundry project +"""Shared DeepEP transport policy for Foundry engine integrations.""" + +from __future__ import annotations + +import enum +from typing import Any + + +class DeepEPTransport(str, enum.Enum): + FABRIC = "fabric" + NVL_IPC = "nvl_ipc" + + +def parse_deepep_transport(raw: str | None) -> DeepEPTransport: + value = DeepEPTransport.FABRIC.value if raw is None else raw + try: + return DeepEPTransport(value) + except ValueError as exc: + choices = ", ".join(item.value for item in DeepEPTransport) + raise ValueError( + f"deepep_transport must be one of: {choices}; got {value!r}" + ) from exc + + +def configure_deepep_buffer( + transport: DeepEPTransport, + *, + num_nvl_bytes: int, + kwargs: dict[str, Any], +) -> tuple[int, dict[str, Any]]: + configured = dict(kwargs) + if transport is DeepEPTransport.FABRIC: + configured["use_fabric"] = True + return 0, configured + configured["use_fabric"] = False + return num_nvl_bytes, configured + + +def validate_transport_match( + *, + archived: str | None, + configured: DeepEPTransport, +) -> None: + archived_transport = parse_deepep_transport(archived) + if archived_transport is not configured: + raise RuntimeError( + "Foundry DeepEP transport mismatch: " + f"archived={archived_transport.value}, configured={configured.value}" + ) + + +def validate_num_nvl_bytes_match( + *, + archived: int, + effective: int, + transport: DeepEPTransport, +) -> None: + if archived != effective: + raise RuntimeError( + "Foundry DeepEP num_nvl_bytes mismatch for " + f"{transport.value}: archived={archived}, effective={effective}" + ) +``` + +- [ ] **Step 3: Write failing engine-config tests** + +For SGLang, parse temporary TOMLs directly. For vLLM, place the equivalent tests in `tests/test_deepep_transport_vllm.py` and use `pytest.importorskip("vllm")` so the Foundry vLLM environment runs them. + +Required assertions: + +```python +assert CUDAGraphExtensionConfig.from_toml(default_path).deepep_transport is DeepEPTransport.FABRIC +assert CUDAGraphExtensionConfig.from_toml(nvl_path).deepep_transport is DeepEPTransport.NVL_IPC +``` + +and `"rdma"` raises `ValueError` naming `deepep_transport`. + +- [ ] **Step 4: Add the field and getter to both configs** + +In both config dataclasses: + +```python +deepep_transport: DeepEPTransport = DeepEPTransport.FABRIC +``` + +In both `from_toml` constructors: + +```python +deepep_transport=parse_deepep_transport(data.get("deepep_transport")), +``` + +In both modules: + +```python +def get_deepep_transport() -> DeepEPTransport: + if _config is None: + return DeepEPTransport.FABRIC + return _config.deepep_transport +``` + +- [ ] **Step 5: Run CPU tests** + +```bash +pytest tests/test_deepep_transport.py -q +pytest tests/test_deepep_transport_vllm.py -q +``` + +Expected: shared and SGLang tests PASS; vLLM tests PASS in the fork environment and skip only when vLLM is not installed. + +- [ ] **Step 6: Commit and push** + +```bash +git add python/foundry/integration/deepep_transport.py \ + python/foundry/integration/vllm/config.py \ + python/foundry/integration/sglang/config.py \ + tests/test_deepep_transport.py tests/test_deepep_transport_vllm.py +git commit -m "feat: configure DeepEP transport in TOML" +git push -u origin cursor/no-imex-deepep-f10b +``` + +--- + +### Task 4: Persist and Validate the Transport Profile + +**Files:** +- Modify: `python/foundry/integration/vllm/runtime.py` +- Modify: `python/foundry/integration/sglang/runtime.py` +- Modify: `tests/test_deepep_transport.py` +- Modify: `tests/test_deepep_transport_vllm.py` + +**Interfaces:** +- Produces identical archive fields in both `WarmupState` dataclasses: + `deepep_transport: str` and `deepep_num_nvl_bytes: int`. +- Produces `update_warmup_state_deepep_profile(...)` and early LOAD validation in each runtime. + +- [ ] **Step 1: Write failing archive tests** + +Add tests that serialize and reload: + +```python +state.deepep_transport = "nvl_ipc" +state.deepep_num_nvl_bytes = 64 * 1024 * 1024 +``` + +Assert the resulting JSON contains both fields. + +Add a legacy JSON fixture without either field and assert: + +```python +assert loaded.deepep_transport == "fabric" +assert loaded.deepep_num_nvl_bytes == 0 +``` + +Add a LOAD mismatch test that stubs CUDA operations and asserts the mismatch is raised before `load_cuda_modules_and_libraries`. + +- [ ] **Step 2: Extend both WarmupState dataclasses** + +Add: + +```python +deepep_transport: str = DeepEPTransport.FABRIC.value +deepep_num_nvl_bytes: int = 0 +``` + +The existing field-filtering loaders then make missing legacy fields use the dataclass defaults. + +- [ ] **Step 3: Add profile update helpers** + +vLLM: + +```python +def update_warmup_state_deepep_profile( + workspace_root: str, + transport: DeepEPTransport, + num_nvl_bytes: int, +) -> None: + state = load_warmup_state(workspace_root) + state.deepep_transport = transport.value + state.deepep_num_nvl_bytes = num_nvl_bytes + save_warmup_state(workspace_root, state) +``` + +SGLang: + +```python +def update_warmup_state_deepep_profile( + transport: DeepEPTransport, + num_nvl_bytes: int, +) -> None: + state = load_warmup_state() + state.deepep_transport = transport.value + state.deepep_num_nvl_bytes = num_nvl_bytes + save_warmup_state(state) +``` + +Only workspace rank 0 writes the shared state. Preserve SGLang's existing rank guard; for vLLM, return without writing when `Path(cfg.workspace_dir).name != "rank_0"`. + +- [ ] **Step 4: Fail early on LOAD transport mismatch** + +In each `setup_graph_extension`, after resolving the rank workspace but before loading modules or reserving the VMM region: + +```python +if cfg.mode == CUDAGraphExtensionMode.LOAD: + warmup_state = load_warmup_state(...) + validate_transport_match( + archived=warmup_state.deepep_transport, + configured=cfg.deepep_transport, + ) +``` + +This guarantees a mismatch fails before DeepEP, NVSHMEM, or graph restoration. + +- [ ] **Step 5: Run archive tests** + +```bash +pytest tests/test_deepep_transport.py tests/test_deepep_transport_vllm.py -q +``` + +Expected: PASS, including legacy archive behavior and fail-fast call ordering. + +- [ ] **Step 6: Commit and push** + +```bash +git add python/foundry/integration/vllm/runtime.py \ + python/foundry/integration/sglang/runtime.py \ + tests/test_deepep_transport.py tests/test_deepep_transport_vllm.py +git commit -m "feat: validate archived DeepEP transport" +git push -u origin cursor/no-imex-deepep-f10b +``` + +--- + +### Task 5: Apply the Transport in vLLM and SGLang + +**Files:** +- Modify: `python/foundry/integration/vllm/hooks.py` +- Modify: `python/foundry/integration/sglang/hooks.py` +- Modify: `tests/test_deepep_transport.py` +- Modify: `tests/test_deepep_transport_vllm.py` + +**Interfaces:** +- vLLM wraps `DeepEPLLAll2AllManager._make_all2all_kwargs`. +- SGLang wraps `deep_ep.Buffer.__init__`. +- Both record the effective `num_nvl_bytes` on SAVE and validate it on LOAD. + +- [ ] **Step 1: Write failing policy tests** + +Test these exact outcomes: + +```python +# fabric +{"num_nvl_bytes": 0, "use_fabric": True} + +# nvl_ipc, given upstream num_nvl_bytes=67108864 +{"num_nvl_bytes": 67108864, "use_fabric": False} +``` + +Test NONE mode leaves upstream values unchanged. + +Test LOAD with an archived 64 MiB buffer and an effective 128 MiB buffer raises before the original `Buffer.__init__` or communicator setup completes. + +- [ ] **Step 2: Replace the vLLM fabric-only policy** + +In `_patch_deepep`: + +```python +transport = get_deepep_transport() +upstream_num_nvl_bytes = int(out.get("num_nvl_bytes", 0)) +effective_num_nvl_bytes, configured = configure_deepep_buffer( + transport, + num_nvl_bytes=upstream_num_nvl_bytes, + kwargs=out, +) +out.clear() +out.update(configured) +out["num_nvl_bytes"] = effective_num_nvl_bytes +``` + +Then: + +```python +if mode == CUDAGraphExtensionMode.SAVE: + rt.update_warmup_state_deepep_profile( + root, + transport, + effective_num_nvl_bytes, + ) +elif mode == CUDAGraphExtensionMode.LOAD: + state = rt.load_warmup_state(root) + validate_num_nvl_bytes_match( + archived=state.deepep_num_nvl_bytes, + effective=effective_num_nvl_bytes, + transport=transport, + ) +``` + +Remove every `FOUNDRY_DEEPEP_NVL_IPC` check. + +- [ ] **Step 3: Port and adapt the SGLang Buffer wrapper** + +Port `_patch_deepep` from source commits `afa9d36`, `b18b186`, `2328691`, and `05df555`, then replace its environment-variable policy with `get_deepep_transport()`. + +At module scope: + +```python +import inspect +``` + +The wrapper must: + +```python +supports_use_fabric = "use_fabric" in inspect.signature(original_init).parameters +``` + +and raise: + +```python +RuntimeError( + "Foundry SGLang DeepEP requires a deep_ep.Buffer constructor " + "with the use_fabric parameter" +) +``` + +before calling an incompatible DeepEP version. + +Call `_patch_deepep()` from `install_hooks()` after configuration is loaded. + +- [ ] **Step 4: Record or validate SGLang's effective profile** + +Inside the wrapper, after applying the configured policy but before calling the original constructor: + +```python +if mode == CUDAGraphExtensionMode.SAVE: + rt.update_warmup_state_deepep_profile( + transport, + effective_num_nvl_bytes, + ) +elif mode == CUDAGraphExtensionMode.LOAD: + state = rt.load_warmup_state() + validate_num_nvl_bytes_match( + archived=state.deepep_num_nvl_bytes, + effective=effective_num_nvl_bytes, + transport=transport, + ) +``` + +Do not alter existing NVSHMEM initialization ordering. + +- [ ] **Step 5: Run policy and import tests** + +```bash +pytest tests/test_deepep_transport.py tests/test_deepep_transport_vllm.py \ + tests/test_imports.py -q +``` + +Expected: PASS in the engine environment. The tests must demonstrate both kwargs profiles, mismatch rejection, old-DeepEP rejection, and NONE-mode no-op behavior. + +- [ ] **Step 6: Commit and push** + +```bash +git add python/foundry/integration/vllm/hooks.py \ + python/foundry/integration/sglang/hooks.py \ + tests/test_deepep_transport.py tests/test_deepep_transport_vllm.py +git commit -m "feat: apply DeepEP transport in engine hooks" +git push -u origin cursor/no-imex-deepep-f10b +``` + +--- + +### Task 6: Add No-IMEX Recipes and Documentation + +**Files:** +- Create/modify: `recipe/experimental/README.md` +- Create/modify: `recipe/experimental/foundry_save.toml` +- Create/modify: `recipe/experimental/foundry_load.toml` +- Create/modify: `recipe/experimental/sglang_foundry_save.toml` +- Create/modify: `recipe/experimental/sglang_foundry_load.toml` +- Create/modify: `recipe/experimental/serve_qwen3-30ba3b_ipc_ep.sh` +- Create/modify: `recipe/experimental/serve_qwen3-30ba3bfp8_sglang_ipc_ep.sh` +- Modify: `docs/vllm/moe-and-deepep.md` +- Modify: `docs/sglang/hooks.md` +- Create/modify: `tests/test_experimental_recipe.py` + +**Interfaces:** +- Produces reproducible SAVE/LOAD commands that select `nvl_ipc` only through TOML. + +- [ ] **Step 1: Write failing recipe contract tests** + +Assert every no-IMEX SAVE/LOAD TOML contains: + +```toml +deepep_transport = "nvl_ipc" +``` + +Assert each engine's SAVE and LOAD pair shares `workspace_root`, `base_addr`, and `region_size`. + +Assert no experimental script contains: + +```bash +FOUNDRY_DEEPEP_NVL_IPC +``` + +- [ ] **Step 2: Port focused experimental recipes** + +Port the EP-only recipe files from `84d2b95`, `403fed2`, and `eb7a0d2`. Exclude all tensor-parallel recipes. + +Replace the source branch's environment-variable selection with the TOML field: + +```toml +deepep_transport = "nvl_ipc" +``` + +Keep: + +```bash +export NCCL_CUMEM_ENABLE=0 +export NCCL_NVLS_ENABLE=0 +``` + +and, for SGLang: + +```bash +export NVSHMEM_CUMEM_HANDLE_TYPE=FILE_DESCRIPTOR +``` + +- [ ] **Step 3: Update integration documentation** + +Document: + +- the two allowed values and `fabric` default; +- intranode-only scope of `nvl_ipc`; +- continued DeepEP/NVSHMEM requirement; +- SAVE/LOAD archive parity and fail-fast errors; +- SGLang's required `Buffer(use_fabric=...)` API; +- no silent fallback; +- the two-GPU validation commands. + +- [ ] **Step 4: Run recipe and formatting tests** + +```bash +pytest tests/test_experimental_recipe.py -q +ruff check python tests +ruff format --check python tests +git diff --check +``` + +Expected: PASS. + +- [ ] **Step 5: Commit and push** + +```bash +git add recipe/experimental docs/vllm/moe-and-deepep.md \ + docs/sglang/hooks.md tests/test_experimental_recipe.py +git commit -m "docs: add no-IMEX DeepEP recipes" +git push -u origin cursor/no-imex-deepep-f10b +``` + +--- + +### Task 7: Verify Native and Engine End-to-End Behavior + +**Files:** +- Test: `tests/test_deepep_fabric.py` +- Test: `tests/run_deepep_matrix.sh` +- Test: vLLM and SGLang experimental recipes + +**Interfaces:** +- Consumes the completed Python, native hook, archive, and recipe work. +- Produces runtime evidence for standalone DeepEP, vLLM inference, SGLang inference, and negative parity validation. + +- [ ] **Step 1: Run the CPU and static suite** + +```bash +pytest tests/test_deepep_transport.py tests/test_deepep_transport_vllm.py \ + tests/test_experimental_recipe.py tests/test_imports.py -q +ruff check python tests +ruff format --check python tests +git diff --check +``` + +Expected: PASS. + +- [ ] **Step 2: Build the exact branch on Modal H100** + +```bash +CC=gcc CXX=g++ pip install -e . --no-build-isolation +``` + +Confirm `libcuda_hook.so` loads and run: + +```bash +pytest tests/test_load_graph.py tests/test_async_load.py \ + tests/test_vmm_alloc.py tests/test_preallocation.py -q +``` + +Expected: PASS. + +- [ ] **Step 3: Run the two-GPU standalone matrix** + +Without initializing CUDA in the orchestration parent: + +```bash +bash tests/run_deepep_matrix.sh ll_nofabric both +bash tests/run_deepep_matrix.sh nvl_ipc both +bash tests/run_deepep_matrix.sh nvl_ipc_prealloc both +``` + +Expected: all PASS, deterministic dispatch output, no CUDA error 999. + +- [ ] **Step 4: Validate vLLM SAVE/LOAD and inference** + +Run two SAVE passes using the same working directory: + +```bash +bash recipe/experimental/serve_qwen3-30ba3b_ipc_ep.sh 2 --save +bash recipe/experimental/serve_qwen3-30ba3b_ipc_ep.sh 2 --save +bash recipe/experimental/serve_qwen3-30ba3b_ipc_ep.sh 2 --load +``` + +Issue one completion request and require coherent output. Check logs for: + +- successful VMM IPC imports or expected relocation messages; +- identical per-rank final allocation offsets; +- no `cuMemImportFromShareableHandle` error 999; +- no graph replay or NVSHMEM initialization error. + +- [ ] **Step 5: Validate SGLang SAVE/LOAD and inference** + +Run: + +```bash +bash recipe/experimental/serve_qwen3-30ba3bfp8_sglang_ipc_ep.sh 2 --save +bash recipe/experimental/serve_qwen3-30ba3bfp8_sglang_ipc_ep.sh 2 --load +``` + +Issue one completion request and apply the same success checks as vLLM. + +- [ ] **Step 6: Run the negative archive-parity control** + +After producing an `nvl_ipc` archive, point LOAD at an otherwise identical TOML containing: + +```toml +deepep_transport = "fabric" +``` + +Expected: startup exits before CUDA modules, DeepEP buffers, or graphs are restored, with: + +```text +Foundry DeepEP transport mismatch: archived=nvl_ipc, configured=fabric +``` + +- [ ] **Step 7: Review the final diff** + +Verify: + +```bash +git status --short +git diff main...HEAD --stat +git diff main...HEAD --check +``` + +Confirm there are no TP-only files, no `FOUNDRY_DEEPEP_NVL_IPC` references, no temporary logs, and no debug-only instrumentation. + +- [ ] **Step 8: Commit any test-driven fixes, push, and update the PR** + +If validation required a fix, return to the task that owns the affected file, +rerun that task's focused test command, and use its explicit `git add` command. +Write a concrete commit subject naming the observed defect, then push: + +```bash +git push -u origin cursor/no-imex-deepep-f10b +``` + +Update the PR description with the passing command matrix and engine inference evidence. diff --git a/docs/superpowers/specs/2026-07-23-no-imex-deepep-design.md b/docs/superpowers/specs/2026-07-23-no-imex-deepep-design.md new file mode 100644 index 00000000..65aa8432 --- /dev/null +++ b/docs/superpowers/specs/2026-07-23-no-imex-deepep-design.md @@ -0,0 +1,211 @@ +# No-IMEX Intranode DeepEP Design + +## Goal + +Allow Foundry graph SAVE and LOAD for intranode expert parallelism in both +vLLM and SGLang without NVIDIA IMEX fabric. Existing fabric behavior remains +the default. + +## Scope + +This design covers: + +- vLLM and SGLang DeepEP expert-parallel workloads on one host. +- An explicit `fabric` or `nvl_ipc` transport choice in each engine's Foundry + TOML configuration. +- Cross-process CUDA VMM IPC that remains valid when a graph is restored in a + fresh process. +- SAVE/LOAD transport-profile validation. +- Standalone DeepEP and engine-level GPU validation. + +It does not cover multi-host expert parallelism, removing NVSHMEM entirely, +enabling NCCL CUMEM/NVLS paths, or dense tensor parallelism. DeepEP continues +to use NVSHMEM for its RDMA buffer; `nvl_ipc` removes the IMEX-dependent +fabric allocation used for the intranode NVLink buffer. + +## Configuration + +Both `foundry.integration.vllm.config.CUDAGraphExtensionConfig` and +`foundry.integration.sglang.config.CUDAGraphExtensionConfig` gain: + +```python +class DeepEPTransport(str, enum.Enum): + FABRIC = "fabric" + NVL_IPC = "nvl_ipc" + + +deepep_transport: DeepEPTransport = DeepEPTransport.FABRIC +``` + +TOML examples: + +```toml +# Existing behavior; this line may be omitted. +deepep_transport = "fabric" +``` + +```toml +# Intranode EP without IMEX fabric. +deepep_transport = "nvl_ipc" +``` + +Unknown values fail while parsing the TOML. The default is `fabric`, so +existing configurations and archives retain their current behavior. + +## Archive Contract + +The effective transport is part of `warmup_state.json` for both integrations: + +```json +{ + "deepep_transport": "nvl_ipc", + "deepep_num_nvl_bytes": 134217728 +} +``` + +SAVE writes the configured value. LOAD compares the archive value with the +current TOML before DeepEP buffers or graphs are initialized. A mismatch +raises a `RuntimeError` naming the archived and configured values. + +Older archives have no transport field. They are interpreted as `fabric`, +matching Foundry's historical vLLM behavior and preserving backward +compatibility. An older archive cannot be loaded with `nvl_ipc`. + +The engine command line still determines the NVLink buffer size. Foundry does +not add a second buffer-size setting. After DeepEP computes +`num_nvl_bytes`, SAVE updates the archive with that effective value and LOAD +compares it before accepting the communicator. Fabric mode records zero. +This catches SAVE and LOAD commands that select different buffer geometry. + +## Engine Integration + +### vLLM + +`_patch_deepep` continues to wrap +`DeepEPLLAll2AllManager._make_all2all_kwargs`. + +- `fabric`: set `use_fabric=True` and `num_nvl_bytes=0`, preserving current + behavior. +- `nvl_ipc`: set `use_fabric=False` and preserve vLLM's computed + `num_nvl_bytes`. + +The wrapper records the effective transport profile during SAVE and validates +it during LOAD. + +### SGLang + +Foundry wraps `deep_ep.Buffer.__init__`, because SGLang does not expose +vLLM's kwargs-builder seam. + +- `fabric`: set `use_fabric=True` and `num_nvl_bytes=0`. +- `nvl_ipc`: set `use_fabric=False` and preserve SGLang's computed + `num_nvl_bytes`. + +The wrapper requires a DeepEP version whose `Buffer` constructor accepts +`use_fabric`. If that argument is unavailable, Foundry raises an actionable +version error instead of silently selecting another transport. + +The existing NVSHMEM bootstrap ordering remains unchanged in both engines: +load captured modules, construct the DeepEP buffer, initialize queued +NVSHMEM modules, then restore graphs. + +## Cross-Process VMM IPC + +The current hook embeds a raw POSIX file-descriptor number in +`CUipcMemHandle`. File-descriptor numbers are process-local, so that protocol +cannot restore an NVL buffer in a fresh process. + +The hook protocol is upgraded to VMM IPC v2: + +1. `cuIpcGetMemHandle` exports the VMM allocation as a POSIX file descriptor. +2. The exporting process parks the descriptor behind a per-process abstract + Unix-domain socket. +3. The 64-byte CUDA IPC handle carries protocol magic, exporter PID, random + process token, allocation/chunk identity, address, size, and generation; + it never carries a raw descriptor number. +4. `cuIpcOpenMemHandle` authenticates to the exporter socket, receives the + descriptor with `SCM_RIGHTS`, and imports it with the CUDA VMM API. +5. The importer first requests the exporter's address. On collision it maps + in a dedicated import range and returns the relocated address. DeepEP's + buffer synchronization distributes the effective peer addresses before + graph use. +6. `cuIpcCloseMemHandle` reference-counts shared chunk mappings and unmaps + them only after the final close. + +Security and lifecycle requirements: + +- The socket accepts only peers with the same UID via `SO_PEERCRED`. +- A random process token prevents stale handles from matching a reused PID. +- Fork detection closes inherited listeners and parked descriptors, then + creates a new server identity. +- Export entries are invalidated when their allocation is freed. +- LOAD preallocation exports the owning VMM chunk plus an interior offset. + Chunk generation participates in cache identity so a freed and recreated + chunk cannot reuse a stale import. +- Driver calls used to reserve an import address bypass Foundry's interposed + allocator, preventing IPC mappings from moving the deterministic graph + allocation cursor. + +Legacy VMM IPC v1 handles are not accepted as Foundry handles; non-Foundry +CUDA IPC handles continue to pass through to the real CUDA driver. + +## Failure Handling + +- Invalid TOML transport: configuration parse error. +- Archive/config transport mismatch: fail before model or graph restoration. +- Effective NVLink buffer-size mismatch: fail when DeepEP creates its buffer. +- Missing `use_fabric` support in DeepEP: actionable runtime error. +- Missing exporter, authentication failure, expired token, or failed FD + transfer: return the corresponding CUDA IPC error and emit one diagnostic. +- Import or map failure: release the imported handle and reserved address; + do not leave a partial cache entry. + +No error path silently falls back from `nvl_ipc` to fabric, because that +would change allocation geometry and invalidate the captured graph. + +## Documentation and Recipes + +The vLLM and SGLang integration documentation will explain: + +- `deepep_transport` values and the default. +- That `nvl_ipc` is intranode-only and still requires DeepEP/NVSHMEM. +- Matching SAVE and LOAD configuration. +- Required `NCCL_CUMEM_ENABLE=0` and `NCCL_NVLS_ENABLE=0`. +- The supported DeepEP revisions. + +Experimental no-IMEX recipes will use `deepep_transport = "nvl_ipc"` in +their SAVE and LOAD TOMLs rather than an environment-variable switch. + +## Testing + +### CPU-only tests + +- Parse the default and both explicit transport values for vLLM and SGLang. +- Reject an unknown transport. +- Round-trip the transport through each `WarmupState`. +- Treat a missing archived field as `fabric`. +- Reject archive/config transport mismatches. +- Verify both DeepEP wrappers select the expected kwargs and preserve + `num_nvl_bytes` only for `nvl_ipc`. +- Verify recipe TOMLs use the first-class setting. + +### Native and GPU tests + +Run on an H100/H200 host without IMEX: + +1. Build Foundry with the CUDA hook and run existing non-EP graph tests to + guard allocator and graph restoration behavior. +2. Run the two-rank standalone DeepEP matrix: + - no-fabric low-latency/RDMA-only path; + - NVL IPC with a nonzero NVLink buffer; + - NVL IPC while LOAD serves allocations from a preallocated VMM chunk. +3. Run two-pass SAVE then LOAD for the vLLM EP recipe and issue an inference + request. Require coherent output and no CUDA error 999. +4. Run SAVE then LOAD for the SGLang EP recipe and issue an inference request + with the same requirements. +5. As a negative control, change LOAD to `fabric` and verify Foundry rejects + the archive before graph restoration. + +Success requires identical per-rank allocation offsets across SAVE/LOAD, +successful VMM IPC imports, graph replay, and coherent model output in both +engines. diff --git a/docs/vllm/moe-and-deepep.md b/docs/vllm/moe-and-deepep.md index e76b48d8..63fdbf76 100644 --- a/docs/vllm/moe-and-deepep.md +++ b/docs/vllm/moe-and-deepep.md @@ -1,6 +1,9 @@ # MoE / DeepEP Support -vLLM-side MoE / expert-parallelism support is what makes the integration interesting beyond dense decode. The remaining MoE-specific concern is that DeepEP all-to-all communicator buffers must be set up in fabric mode for foundry to interpose them, and NVSHMEM module init must fire after the runtime is bootstrapped. +vLLM-side MoE / expert-parallelism support is what makes the integration +interesting beyond dense decode. Foundry configures the DeepEP all-to-all +communicator for either fabric or intranode NVLink IPC, and NVSHMEM module init +must fire after the runtime is bootstrapped. The dense-only baseline doesn't exercise any of this; turn it on by running an MoE model with EP. @@ -8,32 +11,51 @@ The dense-only baseline doesn't exercise any of this; turn it on by running an M vLLM's recent versions use modular kernels for both quantized and unquantized MoE — the same `model_loader.load_model()` path produces deterministic packing regardless of foundry mode. Foundry's integration takes the unified preallocate → `start_graph_builds` → weight-load → `prepare_communication_buffer_for_model` path for all MoE variants (and for dense models), so no quant-metadata roundtrip through `WarmupState` is needed. The earlier `collect_moe_quant_metadata` / `inject_moe_quant_metadata` / `reset_moe_quant_config` / `reinit_moe_quant_config` / `split_load_model` ceremony has been removed. -## DeepEP fabric mode +## DeepEP transport ### Problem -`DeepEPLLAll2AllManager` builds the all-to-all communicator with one of several backends: +`DeepEPLLAll2AllManager` builds the all-to-all communicator with one of two +Foundry-supported profiles: -- NVL (NVLink direct) — uses `cudaIpcGetMemHandle` for buffer sharing. -- Fabric — uses NVSHMEM + symmetric memory. +- `fabric` uses fabric handles and an NVSHMEM symmetric heap. +- `nvl_ipc` uses `cudaIpcGetMemHandle` for the nonzero intranode NVLink buffer + and Foundry's VMM-IPC translation to share it between ranks. -NVL's IPC handles are process-local and not persistable; the captured graph can't reference them on LOAD. Fabric mode allocates buffers via NVSHMEM, which foundry's hook can interpose (`init_nvshmem_for_loaded_modules` re-binds them on LOAD). +`deepep_transport` accepts only those two values and defaults to `fabric` when +omitted. Select the no-IMEX path in both SAVE and LOAD TOMLs: + +```toml +deepep_transport = "nvl_ipc" +``` + +`nvl_ipc` is intranode-only; all ranks must be on one host with the required +NVLink connectivity. It does not remove the DeepEP or NVSHMEM dependency. +DeepEP's RDMA buffer remains on the NVSHMEM symmetric heap, while only its +NVLink buffer uses IPC. The no-IMEX serve script also requires +`NCCL_CUMEM_ENABLE=0` and `NCCL_NVLS_ENABLE=0`. ### Solution -`_patch_deepep` wraps `DeepEPLLAll2AllManager._make_all2all_kwargs`: +`_patch_deepep` wraps +`DeepEPLLAll2AllManager._make_all2all_kwargs`. In Foundry mode it applies the +TOML policy after vLLM computes the upstream arguments: -```python -@functools.wraps(orig) -def patched(self, *args, **kwargs): - kwargs = orig(self, *args, **kwargs) - if get_graph_extension_mode() != CUDAGraphExtensionMode.NONE: - kwargs["use_fabric"] = True - kwargs["num_nvl_bytes"] = 0 # RDMA-only path - return kwargs +- `fabric`: set `use_fabric=True` and force `num_nvl_bytes=0`. +- `nvl_ipc`: set `use_fabric=False` and preserve the computed nonzero + `num_nvl_bytes`. + +The effective transport and NVLink buffer size are recorded in every SAVE +archive. LOAD validates the transport before loading archived CUDA modules and +validates the buffer size when DeepEP creates the communicator. For example, +loading an `nvl_ipc` archive with a `fabric` TOML fails fast with: + +```text +Foundry DeepEP transport mismatch: archived=nvl_ipc, configured=fabric ``` -Forces every foundry-mode DeepEP setup onto the fabric path. The patch lives behind a try/except in case the import path moves. +There is no silent fallback between transports. Falling back would change the +allocation trajectory and invalidate the captured graphs. ## NVSHMEM initialization order @@ -56,3 +78,20 @@ For a working MoE SAVE→LOAD cycle: - All ranks produce `rank_{N}/warmup_state.json`-equivalent per-rank artifacts. - All ranks' `final_alloc_offset` are reproducible across SAVE pass 1 → pass 2. - On LOAD, the first inference call returns coherent tokens. + +On a two-GPU host without IMEX, validate the RDMA-only baseline, direct NVLink +IPC, and LOAD preallocation: + +```bash +bash tests/run_deepep_matrix.sh ll_nofabric both +bash tests/run_deepep_matrix.sh nvl_ipc both +bash tests/run_deepep_matrix.sh nvl_ipc_prealloc both +``` + +Then run the focused vLLM recipe twice in SAVE mode and once in LOAD mode: + +```bash +bash recipe/experimental/serve_qwen3-30ba3b_ipc_ep.sh 2 --save +bash recipe/experimental/serve_qwen3-30ba3b_ipc_ep.sh 2 --save +bash recipe/experimental/serve_qwen3-30ba3b_ipc_ep.sh 2 --load +``` diff --git a/include/BinaryGraphFormat.h b/include/BinaryGraphFormat.h index ca791640..0a1dc248 100644 --- a/include/BinaryGraphFormat.h +++ b/include/BinaryGraphFormat.h @@ -21,7 +21,7 @@ class value; // Sections: // STRING_TABLE — deduplicated null-terminated strings // NODE_TABLE — fixed-size node descriptors (160 bytes each) -// DEPENDENCY_TABLE — (from_id, to_id) pairs +// DEPENDENCY_TABLE — fixed-width endpoint and CUgraphEdgeData records // KERNEL_PARAM_INDEX — per-kernel param descriptors // KERNEL_PARAM_DATA — raw binary param bytes // ARG_BUFFER_DATA — raw binary extra arg buffers @@ -35,6 +35,12 @@ namespace foundry { namespace binary_format { static constexpr uint8_t MAGIC[8] = {'C', 'U', 'G', 'R', 'A', 'P', 'H', '\0'}; +// FORMAT_VERSION stays at 1 across the dependency-record widening from 8 to 16 +// bytes (see the Dependency section below). The section table records each +// section's byte size and the header records num_dependencies, so the two +// record widths are unambiguous and a version bump is unnecessary; bumping it +// would only make current readers needlessly reject the many valid +// endpoint-only (8-byte) archives written before edge metadata existed. static constexpr uint32_t FORMAT_VERSION = 1; enum SectionType : uint32_t { @@ -219,12 +225,38 @@ struct BinGenerator { static_assert(sizeof(BinGenerator) == 24, "BinGenerator must be 24 bytes"); // ---- Dependency ---- +// +// The dependency table has two on-disk record layouts, distinguished purely by +// the stored section size (num_dependencies * sizeof(record)): +// +// - LegacyBinDependency (8 bytes): endpoints only, written by readers that +// predate CUDA graph edge metadata. Current readers accept it and default +// the edge-metadata fields (from_port/to_port/type) to zero. +// - BinDependency (16 bytes): the current record, adding the CUgraphEdgeData +// fields plus five reserved zero bytes. This is what the current writer +// always emits. +// +// Compatibility is one-directional: current readers read both widths, but a +// legacy (8-byte-only) reader cannot understand the 16-byte records this writer +// produces and is not expected to. Because FORMAT_VERSION is unchanged, such a +// reader would misparse a current archive; there is no old-reader forward +// compatibility for edge-metadata archives. + +struct LegacyBinDependency { + uint32_t from_id; + uint32_t to_id; +}; +static_assert(sizeof(LegacyBinDependency) == 8, "LegacyBinDependency must be 8 bytes"); struct BinDependency { uint32_t from_id; uint32_t to_id; + uint8_t from_port; + uint8_t to_port; + uint8_t type; + uint8_t reserved[5]; }; -static_assert(sizeof(BinDependency) == 8, "BinDependency must be 8 bytes"); +static_assert(sizeof(BinDependency) == 16, "BinDependency must be 16 bytes"); #pragma pack(pop) diff --git a/include/CUDAGraphInternal.h b/include/CUDAGraphInternal.h index 52c53789..6b0d1a0c 100644 --- a/include/CUDAGraphInternal.h +++ b/include/CUDAGraphInternal.h @@ -8,6 +8,10 @@ namespace foundry { +CUresult add_graph_dependencies(CUgraph graph, const std::vector& from_nodes, + const std::vector& to_nodes, + const std::vector& edge_data); + // Holds deferred metadata for the split start/finish graph loading flow. // Returned by start_graph_builds_impl, consumed by finish_graph_loads_impl. struct PendingGraphLoads { diff --git a/include/metadata.h b/include/metadata.h index 4d867ab9..bdda637e 100644 --- a/include/metadata.h +++ b/include/metadata.h @@ -202,6 +202,7 @@ struct EmptyNodeMetadata {}; struct GraphDependency { int from_index; int to_index; + CUgraphEdgeData edge_data{}; }; using GraphNodeMetadata = diff --git a/python/foundry/graph.py b/python/foundry/graph.py index 427d1599..cfeffb7d 100644 --- a/python/foundry/graph.py +++ b/python/foundry/graph.py @@ -135,7 +135,7 @@ def __init__( if self.__class__.default_capture_stream is None: self.__class__.default_capture_stream = torch.cuda.Stream() - self.pool = () if pool is None else (pool,) + self.pool = (pool,) self.capture_stream = ( stream if stream is not None else self.__class__.default_capture_stream ) diff --git a/python/foundry/integration/deepep_transport.py b/python/foundry/integration/deepep_transport.py new file mode 100644 index 00000000..435950e7 --- /dev/null +++ b/python/foundry/integration/deepep_transport.py @@ -0,0 +1,62 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the Foundry project +"""Shared DeepEP transport policy for Foundry engine integrations.""" + +from __future__ import annotations + +import enum +from typing import Any + + +class DeepEPTransport(str, enum.Enum): + FABRIC = "fabric" + NVL_IPC = "nvl_ipc" + + +def parse_deepep_transport(raw: str | None) -> DeepEPTransport: + value = DeepEPTransport.FABRIC.value if raw is None else raw + try: + return DeepEPTransport(value) + except ValueError as exc: + choices = ", ".join(item.value for item in DeepEPTransport) + raise ValueError(f"deepep_transport must be one of: {choices}; got {value!r}") from exc + + +def configure_deepep_buffer( + transport: DeepEPTransport, + *, + num_nvl_bytes: int, + kwargs: dict[str, Any], +) -> tuple[int, dict[str, Any]]: + configured = dict(kwargs) + if transport is DeepEPTransport.FABRIC: + configured["use_fabric"] = True + return 0, configured + configured["use_fabric"] = False + return num_nvl_bytes, configured + + +def validate_transport_match( + *, + archived: str | None, + configured: DeepEPTransport, +) -> None: + archived_transport = parse_deepep_transport(archived) + if archived_transport is not configured: + raise RuntimeError( + "Foundry DeepEP transport mismatch: " + f"archived={archived_transport.value}, configured={configured.value}" + ) + + +def validate_num_nvl_bytes_match( + *, + archived: int, + effective: int, + transport: DeepEPTransport, +) -> None: + if archived != effective: + raise RuntimeError( + "Foundry DeepEP num_nvl_bytes mismatch for " + f"{transport.value}: archived={archived}, effective={effective}" + ) diff --git a/python/foundry/integration/sglang/config.py b/python/foundry/integration/sglang/config.py index a0e1d41d..ee9d50a5 100644 --- a/python/foundry/integration/sglang/config.py +++ b/python/foundry/integration/sglang/config.py @@ -11,6 +11,8 @@ import tomllib +from foundry.integration.deepep_transport import DeepEPTransport, parse_deepep_transport + class CUDAGraphExtensionMode(str, enum.Enum): NONE = "none" @@ -23,6 +25,7 @@ class CUDAGraphExtensionConfig: mode: CUDAGraphExtensionMode = CUDAGraphExtensionMode.NONE hook_library_path: str | None = None nvshmem_host_path: str | None = None + deepep_transport: DeepEPTransport = DeepEPTransport.FABRIC base_addr: int = 0x600000000000 region_size: str = "64GB" workspace_root: str = "foundry_archive" @@ -49,6 +52,7 @@ def from_toml(cls, path: str | Path) -> CUDAGraphExtensionConfig: mode=CUDAGraphExtensionMode(data.get("mode", cls.mode.value)), hook_library_path=hook_library_path, nvshmem_host_path=nvshmem_host_path, + deepep_transport=parse_deepep_transport(data.get("deepep_transport")), base_addr=base_addr, region_size=data.get("region_size", cls.region_size), workspace_root=data.get("workspace_root", cls.workspace_root), @@ -131,6 +135,12 @@ def get_nvshmem_host_path() -> str | None: return _config.nvshmem_host_path +def get_deepep_transport() -> DeepEPTransport: + if _config is None: + return DeepEPTransport.FABRIC + return _config.deepep_transport + + def compute_workspace_rank(server_args, tp_rank: int, pp_rank: int, dp_rank: int | None) -> int: if getattr(server_args, "enable_dp_attention", False): return pp_rank * server_args.tp_size + tp_rank diff --git a/python/foundry/integration/sglang/graph_ops.py b/python/foundry/integration/sglang/graph_ops.py index 18ab1f0e..927a95d0 100644 --- a/python/foundry/integration/sglang/graph_ops.py +++ b/python/foundry/integration/sglang/graph_ops.py @@ -8,6 +8,7 @@ import os import re import time +from types import SimpleNamespace from typing import Any import torch @@ -248,6 +249,31 @@ def bootstrap_deepep_buffer(cuda_graph_runner) -> bool: return False +def build_capture_fb_view(cuda_graph_runner, bs: int): + """Build a ForwardBatch-like view for backend capture-side metadata init. + + Mirrors the padded per-bs inputs SGLang's own capture loop feeds the + attention backend as a lightweight namespace. This drives the modern + ``init_forward_metadata_out_graph`` API before a real ``ForwardBatch`` + exists. + """ + buffers = cuda_graph_runner.buffers + num_tokens = bs * cuda_graph_runner.num_tokens_per_bs + encoder_lens = buffers.encoder_lens[:bs] if cuda_graph_runner.is_encoder_decoder else None + seq_lens = buffers.seq_lens[:bs] + return SimpleNamespace( + batch_size=bs, + forward_mode=cuda_graph_runner.capture_forward_mode, + req_pool_indices=buffers.req_pool_indices[:bs], + seq_lens=seq_lens, + seq_lens_cpu=buffers.seq_lens_cpu[:bs], + seq_lens_sum=int(seq_lens.sum().item()), + encoder_lens=encoder_lens, + spec_info=cuda_graph_runner.get_spec_info(num_tokens), + positions=buffers.positions[:num_tokens], + ) + + def initialize_attention_metadata_for_bs(cuda_graph_runner, bs: int) -> None: """Populate ``decode_cuda_graph_metadata[bs]`` for runtime replay. @@ -257,11 +283,17 @@ def initialize_attention_metadata_for_bs(cuda_graph_runner, bs: int) -> None: re-run the same call before runtime replay so the wrappers exist at deterministic VMM addresses. """ + attn_backend = cuda_graph_runner.attn_backend + if hasattr(attn_backend, "init_forward_metadata_out_graph"): + fb_view = build_capture_fb_view(cuda_graph_runner, bs) + attn_backend.init_forward_metadata_out_graph(fb_view, in_capture=True) + return + buffers = cuda_graph_runner.buffers num_tokens = bs * cuda_graph_runner.num_tokens_per_bs encoder_lens = buffers.encoder_lens[:bs] if cuda_graph_runner.is_encoder_decoder else None spec_info = cuda_graph_runner.get_spec_info(num_tokens) - cuda_graph_runner.attn_backend.init_forward_metadata_capture_cuda_graph( + attn_backend.init_forward_metadata_capture_cuda_graph( bs, num_tokens, buffers.req_pool_indices[:bs], diff --git a/python/foundry/integration/sglang/hooks.py b/python/foundry/integration/sglang/hooks.py index 22cd7cfa..c2654222 100644 --- a/python/foundry/integration/sglang/hooks.py +++ b/python/foundry/integration/sglang/hooks.py @@ -5,14 +5,20 @@ from __future__ import annotations import functools +import inspect import logging import os import time from dataclasses import asdict +from foundry.integration.deepep_transport import ( + configure_deepep_buffer, + validate_num_nvl_bytes_match, +) from foundry.integration.sglang import runtime as rt from foundry.integration.sglang.config import ( CUDAGraphExtensionMode, + get_deepep_transport, get_graph_extension_mode, get_workspace_root, load_graph_extension_config, @@ -20,6 +26,7 @@ logger = logging.getLogger(__name__) _INSTALLED = False +_DEEPEP_PATCHED = False def _ep_lazy_init_needed() -> bool: @@ -90,13 +97,224 @@ def install_hooks(server_args) -> None: _patch_init_memory_pool() _patch_load_model() _patch_kernel_warmup() + _patch_attn_tp_logits_gather() _patch_cuda_graph_capture() _patch_spawn_sites() + _patch_deepep() _INSTALLED = True logger.info("[Foundry] SGLang hooks installed") +def _patch_attn_tp_logits_gather() -> None: + """Pre-copy LOAD's local logits shard before the recognized TP gather seams.""" + if get_graph_extension_mode() != CUDAGraphExtensionMode.LOAD: + return + + # SGLang is optional; keep its imports at the hook-install boundary. + from sglang.srt.distributed import parallel_state + from sglang.srt.layers import dp_attention, logits_processor + + missing = object() + original = getattr(logits_processor, "attn_tp_all_gather_into_tensor", missing) + get_attention_tp_group = getattr(dp_attention, "get_attention_tp_group", missing) + group_coordinator = getattr(parallel_state, "GroupCoordinator", missing) + original_group_gather = ( + missing + if group_coordinator is missing + else getattr(group_coordinator, "_all_gather_into_tensor", missing) + ) + + installed: list[str] = [] + + if ( + original is not missing + and get_attention_tp_group is not missing + and not getattr(original, "_foundry_load_local_shard_precopy", False) + ): + + @functools.wraps(original) + def patched(global_logits, logits): + group = get_attention_tp_group() + try: + rank = int(group.rank_in_group) + world_size = int(group.world_size) + except (AttributeError, TypeError, ValueError) as exc: + raise RuntimeError( + "Foundry SGLang LOAD attention-TP gather requires " + "get_attention_tp_group() to expose integer rank_in_group " + "and world_size" + ) from exc + + if world_size <= 0 or rank < 0 or rank >= world_size: + raise RuntimeError( + "Foundry SGLang LOAD attention-TP gather received invalid " + f"rank bounds: rank_in_group={rank}, world_size={world_size}" + ) + + try: + output_shape = tuple(int(size) for size in global_logits.shape) + input_shape = tuple(int(size) for size in logits.shape) + except (AttributeError, TypeError, ValueError) as exc: + raise RuntimeError( + "Foundry SGLang LOAD attention-TP gather requires tensor " + "inputs with concrete shapes" + ) from exc + + expected_output_shape = (world_size, *input_shape) + if output_shape != expected_output_shape: + raise RuntimeError( + "Foundry SGLang LOAD attention-TP gather cannot derive a safe " + "local destination slice: " + f"output shape={output_shape}, input shape={input_shape}, " + f"rank_in_group={rank}, world_size={world_size}, " + f"expected output shape={expected_output_shape}" + ) + + global_logits[rank].copy_(logits) + return original(global_logits, logits) + + patched._foundry_load_local_shard_precopy = True + logits_processor.attn_tp_all_gather_into_tensor = patched + installed.append("attention_tp") + + if original_group_gather is not missing and not getattr( + original_group_gather, + "_foundry_load_local_shard_precopy", + False, + ): + + @functools.wraps(original_group_gather) + def patched_group_gather(self, output, input): + try: + rank = int(self.rank_in_group) + world_size = int(self.world_size) + except (AttributeError, TypeError, ValueError) as exc: + raise RuntimeError( + "Foundry SGLang LOAD regular-TP gather requires " + "GroupCoordinator to expose integer rank_in_group and world_size" + ) from exc + + if world_size <= 0 or rank < 0 or rank >= world_size: + raise RuntimeError( + "Foundry SGLang LOAD regular-TP gather received invalid " + f"rank bounds: rank_in_group={rank}, world_size={world_size}" + ) + + # Only pre-copy the local shard for recognized rank-major layouts. + # Any other caller (a different collective on this shared + # coordinator method, or an unexpected output layout) is passed + # through to the original collective unchanged so the delegate's + # own result and exceptions are preserved. + local_output = _regular_tp_local_shard(output, input, rank, world_size) + if local_output is not None: + local_output.copy_(input) + return original_group_gather(self, output, input) + + patched_group_gather._foundry_load_local_shard_precopy = True + group_coordinator._all_gather_into_tensor = patched_group_gather + installed.append("regular_tp") + + if installed: + logger.info( + "[Foundry] SGLang LOAD logits pre-copy patches installed: %s", + ", ".join(installed), + ) + else: + logger.info("[Foundry] SGLang LOAD logits pre-copy found no supported gather seams") + + +def _regular_tp_local_shard(output, input, rank: int, world_size: int): + """Return the local rank slice of ``output`` for a recognized rank-major + gather layout, or ``None`` when the layout is unsupported. + + Returning ``None`` lets the caller delegate unrecognized layouts to the + original collective unchanged instead of raising, so unrelated callers of + the shared coordinator method are never disturbed. + """ + try: + output_shape = tuple(int(size) for size in output.shape) + input_shape = tuple(int(size) for size in input.shape) + except (AttributeError, TypeError, ValueError): + return None + + expected_output_shapes = [(world_size, *input_shape)] + if input_shape: + concatenated_shape = (world_size * input_shape[0], *input_shape[1:]) + if concatenated_shape != expected_output_shapes[0]: + expected_output_shapes.append(concatenated_shape) + if output_shape not in expected_output_shapes: + return None + + try: + return output.view(world_size, *input_shape)[rank] + except (AttributeError, RuntimeError, TypeError): + return None + + +def _patch_deepep() -> None: + """Apply Foundry's configured transport to SGLang's DeepEP buffer.""" + global _DEEPEP_PATCHED + if _DEEPEP_PATCHED: + return + + try: + # DeepEP is an optional SGLang dependency, so import it only when + # installing the engine-specific patch. + from deep_ep import Buffer + except Exception: + return + + original_init = Buffer.__init__ + supports_use_fabric = "use_fabric" in inspect.signature(original_init).parameters + + @functools.wraps(original_init) + def patched(self, group, num_nvl_bytes=0, num_rdma_bytes=0, *args, **kwargs): + mode = get_graph_extension_mode() + if mode != CUDAGraphExtensionMode.NONE: + if not supports_use_fabric: + raise RuntimeError( + "Foundry SGLang DeepEP requires a deep_ep.Buffer constructor " + "with the use_fabric parameter" + ) + + transport = get_deepep_transport() + effective_num_nvl_bytes, configured = configure_deepep_buffer( + transport, + num_nvl_bytes=int(num_nvl_bytes), + kwargs=kwargs, + ) + kwargs.clear() + kwargs.update(configured) + num_nvl_bytes = effective_num_nvl_bytes + + if mode == CUDAGraphExtensionMode.SAVE: + rt.update_warmup_state_deepep_profile( + transport, + effective_num_nvl_bytes, + ) + elif mode == CUDAGraphExtensionMode.LOAD: + state = rt.load_warmup_state() + validate_num_nvl_bytes_match( + archived=state.deepep_num_nvl_bytes, + effective=effective_num_nvl_bytes, + transport=transport, + ) + + return original_init( + self, + group, + num_nvl_bytes, + num_rdma_bytes, + *args, + **kwargs, + ) + + Buffer.__init__ = patched + _DEEPEP_PATCHED = True + logger.info("[Foundry] SGLang DeepEP transport patch installed") + + def _patch_init_torch_distributed() -> None: from sglang.srt.model_executor import model_runner as mr @@ -422,7 +640,47 @@ def patched(self, *args, **kwargs): # FlashInfer by its per-bs indices_updater_decode. attn_backend = self.attn_backend use_fi_prepass = hasattr(attn_backend, "indices_updater_decode") - real_init = attn_backend.init_forward_metadata_capture_cuda_graph + # SGLang #26735 replaced the legacy capture-init entry point with + # ``init_forward_metadata_out_graph(forward_batch, in_capture)``. + # Shim whichever API the installed SGLang fork exposes. + use_new_api = hasattr(attn_backend, "init_forward_metadata_out_graph") + real_out_graph = attn_backend.init_forward_metadata_out_graph if use_new_api else None + real_init = ( + None if use_new_api else attn_backend.init_forward_metadata_capture_cuda_graph + ) + + def reuse_out_graph(forward_batch, in_capture=False): + if not in_capture: + return real_out_graph(forward_batch, in_capture=in_capture) + + from sglang.srt.layers.attention.flashinfer_backend import ( + DecodeMetadata, + PrefillMetadata, + ) + + bs = forward_batch.batch_size + forward_mode = forward_batch.forward_mode + if forward_mode.is_decode_or_idle(): + wrappers = attn_backend.decode_cuda_graph_metadata.get(bs) + if wrappers is None: + return real_out_graph(forward_batch, in_capture=True) + attn_backend.forward_metadata = DecodeMetadata(wrappers) + elif ( + forward_mode.is_target_verify() + or forward_mode.is_draft_extend() + or forward_mode.is_dllm_extend() + ): + wrappers = attn_backend.prefill_cuda_graph_metadata.get(bs) + if wrappers is None: + return real_out_graph(forward_batch, in_capture=True) + attn_backend.forward_metadata = PrefillMetadata( + wrappers, forward_mode.is_dllm_extend(), False + ) + else: + return real_out_graph(forward_batch, in_capture=True) + # Re-plan against the preallocated wrappers without allocating + # another capture metadata workspace. + real_out_graph(forward_batch, in_capture=False) def reuse_pre_pass_init( bs, @@ -536,11 +794,18 @@ def reuse_pre_pass_init( # Drop the pre-pass's last forward_metadata ref so popping the dict # entry doesn't keep the wrapper alive at refcount 1. attn_backend.forward_metadata = None - attn_backend.init_forward_metadata_capture_cuda_graph = reuse_pre_pass_init + if use_new_api: + attn_backend.init_forward_metadata_out_graph = reuse_out_graph + else: + attn_backend.init_forward_metadata_capture_cuda_graph = reuse_pre_pass_init try: result = orig_capture(self, *args, **kwargs) finally: - attn_backend.init_forward_metadata_capture_cuda_graph = real_init + if use_fi_prepass: + if use_new_api: + attn_backend.init_forward_metadata_out_graph = real_out_graph + else: + attn_backend.init_forward_metadata_capture_cuda_graph = real_init from foundry.integration.sglang.graph_ops import ( pack_fatbins, diff --git a/python/foundry/integration/sglang/runtime.py b/python/foundry/integration/sglang/runtime.py index 1f9032c5..146977f5 100644 --- a/python/foundry/integration/sglang/runtime.py +++ b/python/foundry/integration/sglang/runtime.py @@ -17,6 +17,7 @@ from foundry import ops as cge from foundry.allocation_region import parse_size +from foundry.integration.deepep_transport import DeepEPTransport, validate_transport_match from foundry.integration.sglang.config import ( CUDAGraphExtensionMode, compute_workspace_rank, @@ -38,6 +39,8 @@ class WarmupState: gpu_total_memory: int = 0 memory_pool_config: dict = field(default_factory=dict) final_alloc_offset: int = 0 + deepep_transport: str = DeepEPTransport.FABRIC.value + deepep_num_nvl_bytes: int = 0 @dataclass @@ -109,6 +112,16 @@ def load_warmup_state() -> WarmupState: return WarmupState(**{k: v for k, v in data.items() if k in valid}) +def update_warmup_state_deepep_profile( + transport: DeepEPTransport, + num_nvl_bytes: int, +) -> None: + state = load_warmup_state() + state.deepep_transport = transport.value + state.deepep_num_nvl_bytes = num_nvl_bytes + save_warmup_state(state) + + def setup_graph_extension(server_args, tp_rank: int, pp_rank: int, dp_rank: int | None) -> None: """Set up the VMM region before SGLang initializes NCCL/process groups.""" global _state @@ -128,9 +141,14 @@ def setup_graph_extension(server_args, tp_rank: int, pp_rank: int, dp_rank: int shutil.rmtree(workspace_dir) workspace_dir.mkdir(parents=True, exist_ok=True) elif cfg.mode == CUDAGraphExtensionMode.LOAD: - cge.set_skip_fatbin_processing(True) if not workspace_dir.exists(): raise RuntimeError(f"Foundry workspace for rank {rank} does not exist: {workspace_dir}") + warmup_state = load_warmup_state() + validate_transport_match( + archived=warmup_state.deepep_transport, + configured=cfg.deepep_transport, + ) + cge.set_skip_fatbin_processing(True) cge.load_cuda_modules_and_libraries(str(workspace_dir)) region_size = parse_size(cfg.region_size) diff --git a/python/foundry/integration/vllm/config.py b/python/foundry/integration/vllm/config.py index b03c29e3..9011412f 100644 --- a/python/foundry/integration/vllm/config.py +++ b/python/foundry/integration/vllm/config.py @@ -17,6 +17,8 @@ import tomllib from vllm.logger import init_logger +from foundry.integration.deepep_transport import DeepEPTransport, parse_deepep_transport + if TYPE_CHECKING: from vllm.config import ParallelConfig @@ -35,6 +37,7 @@ class CUDAGraphExtensionConfig: hook_library_path: str | None = None # Path to libnvshmem_host.so for NVSHMEM-based kernels (e.g. DeepEP). nvshmem_host_path: str | None = None + deepep_transport: DeepEPTransport = DeepEPTransport.FABRIC base_addr: int = 0x600000000000 region_size: str = "64GB" workspace_root: str = "foundry_archive" @@ -59,10 +62,15 @@ def from_toml(cls, path: str | Path) -> CUDAGraphExtensionConfig: if hook_library_path is None: hook_library_path = cls._detect_hook_so_path() + nvshmem_host_path = data.get("nvshmem_host_path") + if nvshmem_host_path is None: + nvshmem_host_path = cls._detect_nvshmem_host_path() + return cls( mode=mode, hook_library_path=hook_library_path, - nvshmem_host_path=data.get("nvshmem_host_path"), + nvshmem_host_path=nvshmem_host_path, + deepep_transport=parse_deepep_transport(data.get("deepep_transport")), base_addr=base_addr, region_size=data.get("region_size", cls.region_size), workspace_root=data.get("workspace_root", cls.workspace_root), @@ -79,6 +87,25 @@ def _detect_hook_so_path() -> str | None: return str(hook_so_path) return None + @staticmethod + def _detect_nvshmem_host_path() -> str | None: + """Locate DeepEP's NVSHMEM host lib from the installed wheel.""" + try: + spec = importlib.util.find_spec("nvidia.nvshmem") + except (ImportError, ValueError): + return None + if spec is None: + return None + roots = list(spec.submodule_search_locations or []) + if spec.origin: + roots.append(str(Path(spec.origin).parent)) + for root in roots: + for name in ("libnvshmem_host.so.3", "libnvshmem_host.so"): + candidate = Path(root) / "lib" / name + if candidate.exists(): + return str(candidate) + return None + # Module-global config singleton (process-local). _config: CUDAGraphExtensionConfig | None = None @@ -127,6 +154,12 @@ def get_nvshmem_host_path() -> str | None: return _config.nvshmem_host_path +def get_deepep_transport() -> DeepEPTransport: + if _config is None: + return DeepEPTransport.FABRIC + return _config.deepep_transport + + def _compute_rank_from_parallel_config( parallel_config: ParallelConfig, ) -> int: diff --git a/python/foundry/integration/vllm/hooks.py b/python/foundry/integration/vllm/hooks.py index f7768630..d292b896 100644 --- a/python/foundry/integration/vllm/hooks.py +++ b/python/foundry/integration/vllm/hooks.py @@ -19,9 +19,14 @@ import torch from vllm.logger import init_logger +from foundry.integration.deepep_transport import ( + configure_deepep_buffer, + validate_num_nvl_bytes_match, +) from foundry.integration.vllm import runtime as rt from foundry.integration.vllm.config import ( CUDAGraphExtensionMode, + get_deepep_transport, get_graph_extension_mode, get_workspace_root, load_graph_extension_config, @@ -112,13 +117,15 @@ def install_hooks(compilation_config: CompilationConfig) -> None: ) load_graph_extension_config(cfg_path) + mode = get_graph_extension_mode() + root = get_workspace_root() log.info( "[foundry] install_hooks: mode=%s workspace=%s", - get_graph_extension_mode().value, - get_workspace_root(), + mode.value, + root, ) - if get_graph_extension_mode() != CUDAGraphExtensionMode.NONE: + if mode != CUDAGraphExtensionMode.NONE: # NOTE(liuxs): vllm doesn't allow zero cudagraph_num_of_warmups, hardcode here compilation_config.cudagraph_num_of_warmups = 0 log.info( @@ -557,22 +564,30 @@ def patched(self, vllm_config): root = get_workspace_root() ws = None + has_saved_kv_profile = False if root and os.path.exists(os.path.join(root, "warmup_state.json")): ws = rt.load_warmup_state(root) - log.info( - "[foundry] %s mode: warmup_state.json found, skipping KV profiling", - mode.value.upper(), - ) + has_saved_kv_profile = bool(ws.available_gpu_memory) + if has_saved_kv_profile: + log.info( + "[foundry] %s mode: warmup_state.json found, skipping KV profiling", + mode.value.upper(), + ) + else: + log.info( + "[foundry] %s mode: bootstrap warmup_state.json found; running KV profiling", + mode.value.upper(), + ) specs = self.model_executor.get_kv_cache_specs() has_kv = any(s for s in specs) - if ws is not None: + if has_saved_kv_profile: available_mem = ws.available_gpu_memory num_gpu = ws.num_gpu_blocks num_cpu = ws.num_cpu_blocks elif mode == CUDAGraphExtensionMode.LOAD: - raise RuntimeError("foundry LOAD requires warmup_state.json") + raise RuntimeError("foundry LOAD requires a profiled warmup_state.json") elif has_kv: if envs.VLLM_ELASTIC_EP_SCALE_UP_LAUNCH: assert self.available_gpu_memory_for_kv_cache > 0 @@ -597,7 +612,7 @@ def patched(self, vllm_config): args=(vllm_config.model_config.max_model_len,), ) sched = generate_scheduler_kv_cache_config(cfgs) - if ws is None: + if not has_saved_kv_profile: num_gpu = sched.num_blocks elif sched.num_blocks != num_gpu: log.warning( @@ -621,20 +636,25 @@ def patched(self, vllm_config): dp_shard = getattr(vllm_config.parallel_config, "data_parallel_index", 0) or 0 if mode == CUDAGraphExtensionMode.SAVE and root is not None and dp_shard == 0: os.makedirs(root, exist_ok=True) - if ws is not None: - ws.final_alloc_offset = rt.get_final_alloc_offset() - rt.save_warmup_state(root, ws) + if has_saved_kv_profile: + state = ws + else: + state = rt.create_warmup_state() + if ws is not None: + state.deepep_transport = ws.deepep_transport + state.deepep_num_nvl_bytes = ws.deepep_num_nvl_bytes + state.available_gpu_memory = available_mem + state.num_gpu_blocks = num_gpu + state.num_cpu_blocks = num_cpu + state.gpu_memory_utilization = vllm_config.cache_config.gpu_memory_utilization + state.final_alloc_offset = rt.get_final_alloc_offset() + rt.save_warmup_state(root, state) + if has_saved_kv_profile: log.info( - "[foundry] Updated warmup state: final_alloc_offset=%d", ws.final_alloc_offset + "[foundry] Updated warmup state: final_alloc_offset=%d", + state.final_alloc_offset, ) else: - new = rt.create_warmup_state() - new.available_gpu_memory = available_mem - new.num_gpu_blocks = num_gpu - new.num_cpu_blocks = num_cpu - new.gpu_memory_utilization = vllm_config.cache_config.gpu_memory_utilization - new.final_alloc_offset = rt.get_final_alloc_offset() - rt.save_warmup_state(root, new) log.info("[foundry] Saved warmup state: num_gpu_blocks=%d", num_gpu) return sched @@ -675,12 +695,14 @@ def patched_init(self, *args, **kwargs): # --------------------------------------------------------------------------- -# DeepEP buffer kwargs: use_fabric=True, num_nvl_bytes=0 on any foundry mode. +# DeepEP buffer kwargs: apply the configured transport on SAVE/LOAD. # --------------------------------------------------------------------------- def _patch_deepep() -> None: try: + # vLLM is an optional integration dependency, so keep its engine + # import inside the patch installer. from vllm.distributed.device_communicators import all2all except Exception: return @@ -694,13 +716,37 @@ def _patch_deepep() -> None: @functools.wraps(orig) def patched(self, *args, **kwargs): out = orig(self, *args, **kwargs) - if isinstance(out, dict) and get_graph_extension_mode() != CUDAGraphExtensionMode.NONE: - # Force VMM-fabric buffers - # so the foundry cuMem hook tracks them, and disable the - # NVLink buffer (LL-mode uses RDMA) to avoid fabric-cuMemCreate - # collisions with preallocated region. - out["use_fabric"] = True - out["num_nvl_bytes"] = 0 + mode = get_graph_extension_mode() + if not isinstance(out, dict) or mode == CUDAGraphExtensionMode.NONE: + return out + + transport = get_deepep_transport() + upstream_num_nvl_bytes = int(out.get("num_nvl_bytes", 0)) + effective_num_nvl_bytes, configured = configure_deepep_buffer( + transport, + num_nvl_bytes=upstream_num_nvl_bytes, + kwargs=out, + ) + out.clear() + out.update(configured) + out["num_nvl_bytes"] = effective_num_nvl_bytes + + root = get_workspace_root() + if root is None: + raise RuntimeError("Foundry vLLM DeepEP requires workspace_root") + if mode == CUDAGraphExtensionMode.SAVE: + rt.update_warmup_state_deepep_profile( + root, + transport, + effective_num_nvl_bytes, + ) + elif mode == CUDAGraphExtensionMode.LOAD: + state = rt.load_warmup_state(root) + validate_num_nvl_bytes_match( + archived=state.deepep_num_nvl_bytes, + effective=effective_num_nvl_bytes, + transport=transport, + ) return out cls._make_all2all_kwargs = patched diff --git a/python/foundry/integration/vllm/runtime.py b/python/foundry/integration/vllm/runtime.py index c81e64cb..7af22ff2 100644 --- a/python/foundry/integration/vllm/runtime.py +++ b/python/foundry/integration/vllm/runtime.py @@ -25,6 +25,7 @@ from foundry import ops as fops from foundry.allocation_region import parse_size +from foundry.integration.deepep_transport import DeepEPTransport, validate_transport_match from foundry.integration.vllm.config import ( CUDAGraphExtensionMode, _compute_rank_from_parallel_config, @@ -57,6 +58,8 @@ class WarmupState: num_cpu_blocks: int = 0 gpu_memory_utilization: float = 0.0 final_alloc_offset: int = 0 + deepep_transport: str = DeepEPTransport.FABRIC.value + deepep_num_nvl_bytes: int = 0 _final_alloc_offset: int = 0 @@ -99,6 +102,24 @@ def load_warmup_state(workspace_root: str) -> WarmupState: return WarmupState(**{k: v for k, v in data.items() if k in valid}) +def update_warmup_state_deepep_profile( + workspace_root: str, + transport: DeepEPTransport, + num_nvl_bytes: int, +) -> None: + cfg = get_config() + if cfg is None or cfg.workspace_dir is None: + return + if Path(cfg.workspace_dir).name != "rank_0": + return + path = os.path.join(workspace_root, "warmup_state.json") + state = load_warmup_state(workspace_root) if os.path.exists(path) else WarmupState() + state.deepep_transport = transport.value + state.deepep_num_nvl_bytes = num_nvl_bytes + os.makedirs(workspace_root, exist_ok=True) + save_warmup_state(workspace_root, state) + + # --------------------------------------------------------------------------- # Worker-scope state (per-process, set by setup_graph_extension) # --------------------------------------------------------------------------- @@ -147,7 +168,6 @@ def setup_graph_extension(parallel_config: ParallelConfig) -> None: workspace_dir.mkdir(parents=True, exist_ok=True) elif cfg.mode == CUDAGraphExtensionMode.LOAD: - fops.set_skip_fatbin_processing(True) if not workspace_dir.exists(): # rank_0 fallback is only safe for single-rank setups. # Under TP/PP/DP > 1 each rank has its own device addresses, @@ -173,6 +193,12 @@ def setup_graph_extension(parallel_config: ParallelConfig) -> None: f"exist. For multi-rank LOAD, re-run SAVE with the same " f"parallelism topology." ) + warmup_state = load_warmup_state(cfg.workspace_root) + validate_transport_match( + archived=warmup_state.deepep_transport, + configured=cfg.deepep_transport, + ) + fops.set_skip_fatbin_processing(True) log.info("[foundry] LOAD: loading CUDA modules from %s", cfg.workspace_dir) t = time.perf_counter() fops.load_cuda_modules_and_libraries(cfg.workspace_dir) diff --git a/recipe/experimental/README.md b/recipe/experimental/README.md new file mode 100644 index 00000000..6a63eeb4 --- /dev/null +++ b/recipe/experimental/README.md @@ -0,0 +1,134 @@ +# No-IMEX intranode DeepEP recipes + +These experimental recipes run vLLM or SGLang expert parallelism on one host +without NVIDIA IMEX. They select Foundry's NVLink IPC transport in matching +SAVE and LOAD TOMLs: + +```toml +deepep_transport = "nvl_ipc" +``` + +`deepep_transport` accepts only `fabric` and `nvl_ipc`; omitting it selects the +`fabric` default. `nvl_ipc` is intranode-only and requires all participating +GPUs to share the needed NVLink connectivity. It does not remove DeepEP or +NVSHMEM: DeepEP's RDMA buffer still uses the NVSHMEM symmetric heap while its +nonzero NVLink buffer uses Foundry's VMM-IPC translation. + +## Files + +```text +recipe/experimental/ +├── README.md +├── foundry_save.toml +├── foundry_load.toml +├── serve_qwen3-30ba3b_ipc_ep.sh +├── sglang_foundry_save.toml +├── sglang_foundry_load.toml +└── serve_qwen3-30ba3bfp8_sglang_ipc_ep.sh +``` + +This directory intentionally contains only expert-parallel recipes. The vLLM +files use `foundry_archive_ipc`; the SGLang files use +`foundry_archive_sglang_ipc`. Both paths are relative to the directory from +which the server is started. + +## Requirements + +- Build and install this Foundry revision so `libcuda_hook.so` includes the + SCM_RIGHTS VMM-IPC transport: + + ```bash + CC=gcc CXX=g++ pip install -e . --no-build-isolation + ``` + +- Install DeepEP and NVSHMEM. Set `nvshmem_host_path` in both TOMLs for an + engine only when Foundry cannot auto-detect a custom NVSHMEM host library. + Use a machine-local absolute path; no checkout-specific path is assumed. +- The SGLang recipe requires DeepEP `29d31c0` or newer because + `deep_ep.Buffer` must accept `use_fabric`, plus sgl-deep-gemm 0.1.2 or newer + and FlashAttention 3. +- The scripts keep `NCCL_CUMEM_ENABLE=0` and `NCCL_NVLS_ENABLE=0` for every + Foundry phase. The SGLang script also sets + `NVSHMEM_CUMEM_HANDLE_TYPE=FILE_DESCRIPTOR`. + +## vLLM SAVE and LOAD + +Run every phase from the same repository path so vLLM sees an identical TOML +path and reuses its compile cache: + +```bash +rm -rf foundry_archive_ipc +bash recipe/experimental/serve_qwen3-30ba3b_ipc_ep.sh 2 --save +bash recipe/experimental/serve_qwen3-30ba3b_ipc_ep.sh 2 --save +bash recipe/experimental/serve_qwen3-30ba3b_ipc_ep.sh 2 --load +``` + +Stop each SAVE server after startup. The first pass profiles and captures; the +second pass checks deterministic re-capture. Leave the LOAD server running for +inference. + +## SGLang SAVE and LOAD + +SGLang needs one SAVE pass: + +```bash +rm -rf foundry_archive_sglang_ipc +bash recipe/experimental/serve_qwen3-30ba3bfp8_sglang_ipc_ep.sh 2 --save +bash recipe/experimental/serve_qwen3-30ba3bfp8_sglang_ipc_ep.sh 2 --load +``` + +Keep the model, EP size, low-latency token cap, chunked-prefill size, and TOML +settings identical between SAVE and LOAD. Both phases pin `--random-seed 42` and +`--mem-fraction-static 0.65`: at `0.8`, LOAD failed before graph restoration with +`cuMemCreate failed with error 2` while reserving each rank's saved remainder, +whereas `0.65` restores every graph and reaches the endpoint on both ranks. + +### DeepEP dispatch mode override (`DEEPEP_MODE`) + +The SGLang script threads `--deepep-mode "$DEEPEP_MODE"` through to SGLang and +defaults to `low_latency`, whose graph-capturable decode dispatch computes +`num_nvl_bytes=0` (no intranode NVLink buffer, so the VMM-IPC import path is not +exercised end to end). Override the mode with the same value on SAVE and LOAD: + +```bash +DEEPEP_MODE=auto bash recipe/experimental/serve_qwen3-30ba3bfp8_sglang_ipc_ep.sh 2 --save +DEEPEP_MODE=auto bash recipe/experimental/serve_qwen3-30ba3bfp8_sglang_ipc_ep.sh 2 --load +``` + +`DEEPEP_MODE=auto` keeps the low-latency decode graphs Foundry captures while +also allocating DeepEP's normal-path intranode NVLink buffer with a nonzero +`num_nvl_bytes` (from `get_nvl_buffer_size_hint`). With `nvl_ipc` transport that +buffer flows through Foundry's VMM-IPC translation, so LOAD observes a real VMM +import. `DEEPEP_MODE=normal` is intentionally unsupported for SAVE/LOAD: SGLang +forces `disable_cuda_graph=True` when `deepep_mode=normal`, leaving Foundry no +captured graphs to persist or restore. + +## Archive parity and failure behavior + +The SAVE and LOAD config for an engine must match on `deepep_transport`, +`workspace_root`, `base_addr`, and `region_size`. Foundry stores the effective +transport and NVLink buffer size in the archive. A transport mismatch fails +before CUDA modules, DeepEP buffers, or graphs are restored: + +```text +Foundry DeepEP transport mismatch: archived=nvl_ipc, configured=fabric +``` + +An effective NVLink buffer-size mismatch also fails when DeepEP constructs the +buffer. There is no silent fallback from `nvl_ipc` to `fabric`; changing the +transport changes allocation geometry and would invalidate captured graphs. + +## Two-GPU validation + +On a two-GPU host without IMEX, run all three standalone cases from the +repository root: + +```bash +bash tests/run_deepep_matrix.sh ll_nofabric both +bash tests/run_deepep_matrix.sh nvl_ipc both +bash tests/run_deepep_matrix.sh nvl_ipc_prealloc both +``` + +The cases cover the RDMA-only baseline, a nonzero NVLink IPC buffer, and LOAD +from a preallocated Foundry VMM chunk. They require deterministic dispatch, +successful SAVE and LOAD, and no CUDA error 999. diff --git a/recipe/experimental/foundry_load.toml b/recipe/experimental/foundry_load.toml new file mode 100644 index 00000000..9ed5ec0e --- /dev/null +++ b/recipe/experimental/foundry_load.toml @@ -0,0 +1,9 @@ +mode = "load" +base_addr = 0x600000000000 +region_size = "256GB" +workspace_root = "foundry_archive_ipc" +scratch_space_size = "4096MB" +deepep_transport = "nvl_ipc" +# DeepEP low-latency still allocates its RDMA buffer on the NVSHMEM symmetric +# heap. Foundry auto-detects libnvshmem_host.so from the installed NVSHMEM +# package; set nvshmem_host_path only to use a custom build. diff --git a/recipe/experimental/foundry_save.toml b/recipe/experimental/foundry_save.toml new file mode 100644 index 00000000..3232db54 --- /dev/null +++ b/recipe/experimental/foundry_save.toml @@ -0,0 +1,9 @@ +mode = "save" +base_addr = 0x600000000000 +region_size = "256GB" +workspace_root = "foundry_archive_ipc" +scratch_space_size = "4096MB" +deepep_transport = "nvl_ipc" +# DeepEP low-latency still allocates its RDMA buffer on the NVSHMEM symmetric +# heap. Foundry auto-detects libnvshmem_host.so from the installed NVSHMEM +# package; set nvshmem_host_path only to use a custom build. diff --git a/recipe/experimental/serve_qwen3-30ba3b_ipc_ep.sh b/recipe/experimental/serve_qwen3-30ba3b_ipc_ep.sh new file mode 100755 index 00000000..da874654 --- /dev/null +++ b/recipe/experimental/serve_qwen3-30ba3b_ipc_ep.sh @@ -0,0 +1,67 @@ +#!/bin/bash +set -euo pipefail +# Usage: bash serve_qwen3-30ba3b_ipc_ep.sh [--save|--load] +# +# EXPERIMENTAL: vLLM DeepEP expert parallel using the intranode NVLink IPC +# transport selected by the matching Foundry TOML. This requires the Foundry +# VMM-IPC hook and a host whose GPUs share an NVLink domain. +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" + +EP_SIZE=${1:?Usage: $0 [--save|--load]} +MODEL_NAME="Qwen/Qwen3-30B-A3B" +HOST="0.0.0.0" +PORT=12000 +GPU_MEMORY_UTILIZATION=0.8 + +FOUNDRY_ARGS=() +if [[ "${2:-}" == "--save" ]]; then + FOUNDRY_TOML="${SCRIPT_DIR}/foundry_save.toml" +elif [[ "${2:-}" == "--load" ]]; then + FOUNDRY_TOML="${SCRIPT_DIR}/foundry_load.toml" +elif [[ -n "${2:-}" ]]; then + echo "Usage: $0 [--save|--load]" + exit 1 +fi + +if [[ -n "${FOUNDRY_TOML:-}" ]]; then + # The CUMEM P2P and NVLS multicast paths request mapping capabilities that + # Foundry's VMM allocation region does not carry. + export NCCL_CUMEM_ENABLE=0 + export NCCL_NVLS_ENABLE=0 + FOUNDRY_ARGS+=( + --compilation-config.graph_extension_config_path "$FOUNDRY_TOML" + ) + echo "Using Foundry DeepEP NVL/IPC: ${FOUNDRY_TOML}" +else + echo "Running without Foundry (baseline vLLM)" +fi + +# Foundry currently patches the V1 model runner. +export VLLM_USE_V2_MODEL_RUNNER=0 +export VLLM_USE_FLASHINFER_SAMPLER=1 +export VLLM_DISABLE_SHARED_EXPERTS_STREAM=1 + +CUDAGRAPH_CAPTURE_SIZES=($(seq 1 256)) + +ARGS=( + --trust-remote-code + --seed 42 + --host "$HOST" + --port "$PORT" + --tensor-parallel-size 1 + --data-parallel-size "$EP_SIZE" + --gpu-memory-utilization "$GPU_MEMORY_UTILIZATION" + --distributed-executor-backend uni + --enable-expert-parallel + --all2all-backend deepep_low_latency + --no-enable-prefix-caching + --max-num-batched-tokens 256 + --max-num-seqs 256 + --attention-config.backend FLASH_ATTN + --compilation-config.cudagraph_mode FULL_DECODE_ONLY + --compilation-config.cudagraph_num_of_warmups 0 + --chat-template-content-format string + --cudagraph-capture-sizes "${CUDAGRAPH_CAPTURE_SIZES[@]}" +) + +vllm serve "$MODEL_NAME" "${ARGS[@]}" "${FOUNDRY_ARGS[@]}" diff --git a/recipe/experimental/serve_qwen3-30ba3bfp8_sglang_ipc_ep.sh b/recipe/experimental/serve_qwen3-30ba3bfp8_sglang_ipc_ep.sh new file mode 100755 index 00000000..a6b57c44 --- /dev/null +++ b/recipe/experimental/serve_qwen3-30ba3bfp8_sglang_ipc_ep.sh @@ -0,0 +1,71 @@ +#!/bin/bash +set -euo pipefail +# Usage: bash serve_qwen3-30ba3bfp8_sglang_ipc_ep.sh [--save|--load] +# +# EXPERIMENTAL: SGLang DeepEP expert parallel using the intranode NVLink IPC +# transport selected by the matching Foundry TOML. +# +# Requires DeepEP 29d31c0 or newer, sgl-deep-gemm >= 0.1.2, flash-attn-3, +# NVSHMEM, and the Foundry VMM-IPC hook. +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" + +EP_SIZE=${1:?Usage: $0 [--save|--load]} +MODEL_NAME="${SGL_MODEL:-Qwen/Qwen3-30B-A3B-FP8}" +HOST="0.0.0.0" +PORT=12000 +# DeepEP dispatch mode (SGLang --deepep-mode). Default low_latency keeps the +# product's graph-capturable decode dispatch with num_nvl_bytes=0. Override with +# DEEPEP_MODE=auto to also allocate DeepEP's nonzero-num_nvl_bytes intranode +# NVLink buffer (exercising the Foundry VMM-IPC import) while keeping the +# low-latency decode graphs. DEEPEP_MODE=normal is unsupported here: SGLang sets +# disable_cuda_graph=True for deepep_mode=normal, so Foundry has no graphs to +# SAVE/LOAD. +DEEPEP_MODE="${DEEPEP_MODE:-low_latency}" +# Validated static memory fraction: LOAD reserves each rank's saved remainder +# during graph restoration, which failed at 0.8 with cuMemCreate error 2 but +# succeeds at 0.65 (both ranks restore every graph and reach the endpoint). +MEM_FRACTION_STATIC=0.65 + +FOUNDRY_ARGS=() +if [[ "${2:-}" == "--save" ]]; then + FOUNDRY_TOML="${SCRIPT_DIR}/sglang_foundry_save.toml" +elif [[ "${2:-}" == "--load" ]]; then + FOUNDRY_TOML="${SCRIPT_DIR}/sglang_foundry_load.toml" +elif [[ -n "${2:-}" ]]; then + echo "Usage: $0 [--save|--load]" + exit 1 +fi + +if [[ -n "${FOUNDRY_TOML:-}" ]]; then + export NVSHMEM_CUMEM_HANDLE_TYPE=FILE_DESCRIPTOR + export NCCL_CUMEM_ENABLE=0 + export NCCL_NVLS_ENABLE=0 + FOUNDRY_ARGS+=(--foundry-graph-extension-config-path "$FOUNDRY_TOML") + echo "Using Foundry DeepEP NVL/IPC: ${FOUNDRY_TOML}" +else + echo "Running without Foundry (baseline SGLang)" +fi + +# Keep this identical between SAVE and LOAD: it affects DeepEP's allocation +# trajectory and the captured low-latency graph shapes. +export SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK=256 + +sglang serve \ + --model-path "$MODEL_NAME" \ + --trust-remote-code \ + --host "$HOST" --port "$PORT" \ + --tp-size "$EP_SIZE" \ + --dp-size "$EP_SIZE" \ + --ep-size "$EP_SIZE" \ + --enable-dp-attention \ + --moe-a2a-backend deepep \ + --deepep-mode "$DEEPEP_MODE" \ + --moe-runner-backend deep_gemm \ + --mem-fraction-static "$MEM_FRACTION_STATIC" \ + --disable-radix-cache \ + --disable-custom-all-reduce \ + --chunked-prefill-size 256 \ + --attention-backend fa3 \ + --cuda-graph-max-bs 128 \ + --random-seed 42 \ + "${FOUNDRY_ARGS[@]}" diff --git a/recipe/experimental/sglang_foundry_load.toml b/recipe/experimental/sglang_foundry_load.toml new file mode 100644 index 00000000..5c8c82fa --- /dev/null +++ b/recipe/experimental/sglang_foundry_load.toml @@ -0,0 +1,9 @@ +mode = "load" +base_addr = 0x600000000000 +region_size = "256GB" +workspace_root = "foundry_archive_sglang_ipc" +scratch_space_size = "4096MB" +deepep_transport = "nvl_ipc" +# DeepEP low-latency still allocates its RDMA buffer on the NVSHMEM symmetric +# heap. Foundry auto-detects libnvshmem_host.so from the installed NVSHMEM +# package; set nvshmem_host_path only to use a custom build. diff --git a/recipe/experimental/sglang_foundry_save.toml b/recipe/experimental/sglang_foundry_save.toml new file mode 100644 index 00000000..78669275 --- /dev/null +++ b/recipe/experimental/sglang_foundry_save.toml @@ -0,0 +1,9 @@ +mode = "save" +base_addr = 0x600000000000 +region_size = "256GB" +workspace_root = "foundry_archive_sglang_ipc" +scratch_space_size = "4096MB" +deepep_transport = "nvl_ipc" +# DeepEP low-latency still allocates its RDMA buffer on the NVSHMEM symmetric +# heap. Foundry auto-detects libnvshmem_host.so from the installed NVSHMEM +# package; set nvshmem_host_path only to use a custom build. diff --git a/tests/cuda_graph_edge_data_test.cu b/tests/cuda_graph_edge_data_test.cu new file mode 100644 index 00000000..5844cdaf --- /dev/null +++ b/tests/cuda_graph_edge_data_test.cu @@ -0,0 +1,179 @@ +// Test-only CUDA helpers for graph dependency edge-data integration coverage. + +#include +#include +#include +#include +#include + +#include +#include +#include +#include + +#include "BinaryGraphFormat.h" +#include "CUDAGraph.h" + +namespace { + +void check_cuda(cudaError_t result, const char* operation) { + TORCH_CHECK(result == cudaSuccess, operation, " failed: ", cudaGetErrorString(result)); +} + +void check_driver(CUresult result, const char* operation) { + const char* error = nullptr; + if (result != CUDA_SUCCESS) { + cuGetErrorString(result, &error); + } + TORCH_CHECK(result == CUDA_SUCCESS, operation, " failed: ", error ? error : "unknown"); +} + +__global__ void write_values(float* values, int count) { + int index = blockIdx.x * blockDim.x + threadIdx.x; + if (index < count) { + values[index] = static_cast(index + 1); + } + cudaTriggerProgrammaticLaunchCompletion(); +} + +__global__ void double_values_programmatic(const float* values, float* output, int count) { + cudaGridDependencySynchronize(); + int index = blockIdx.x * blockDim.x + threadIdx.x; + if (index < count) { + output[index] = values[index] * 2.0f; + } +} + +__global__ void double_values_default(const float* values, float* output, int count) { + int index = blockIdx.x * blockDim.x + threadIdx.x; + if (index < count) { + output[index] = values[index] * 2.0f; + } +} + +__global__ void complete_after_values() {} + +void validate_tensors(const torch::Tensor& values, const torch::Tensor& output) { + TORCH_CHECK(values.is_cuda() && output.is_cuda(), "expected CUDA tensors"); + TORCH_CHECK(values.scalar_type() == torch::kFloat32, "expected float32 input"); + TORCH_CHECK(output.scalar_type() == torch::kFloat32, "expected float32 output"); + TORCH_CHECK(values.is_contiguous() && output.is_contiguous(), "expected contiguous tensors"); + TORCH_CHECK(values.numel() == output.numel(), "input and output sizes must match"); +} + +} // namespace + +void launch_programmatic(torch::Tensor values, torch::Tensor output) { + validate_tensors(values, output); + const int count = static_cast(values.numel()); + constexpr int threads = 128; + const int blocks = (count + threads - 1) / threads; + cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + + write_values<<>>(values.data_ptr(), count); + check_cuda(cudaGetLastError(), "write_values launch"); + + cudaLaunchAttribute attribute{}; + attribute.id = cudaLaunchAttributeProgrammaticStreamSerialization; + attribute.val.programmaticStreamSerializationAllowed = 1; + + cudaLaunchConfig_t config{}; + config.gridDim = dim3(blocks); + config.blockDim = dim3(threads); + config.stream = stream; + config.attrs = &attribute; + config.numAttrs = 1; + check_cuda(cudaLaunchKernelEx(&config, double_values_programmatic, values.data_ptr(), + output.data_ptr(), count), + "double_values_programmatic launch"); + complete_after_values<<<1, 1, 0, stream>>>(); + check_cuda(cudaGetLastError(), "complete_after_values launch"); +} + +void launch_default(torch::Tensor values, torch::Tensor output) { + validate_tensors(values, output); + const int count = static_cast(values.numel()); + constexpr int threads = 128; + const int blocks = (count + threads - 1) / threads; + cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + + write_values<<>>(values.data_ptr(), count); + check_cuda(cudaGetLastError(), "write_values launch"); + double_values_default<<>>(values.data_ptr(), + output.data_ptr(), count); + check_cuda(cudaGetLastError(), "double_values_default launch"); + complete_after_values<<<1, 1, 0, stream>>>(); + check_cuda(cudaGetLastError(), "complete_after_values launch"); +} + +std::vector> graph_edges( + uint64_t graph_address) { + CUgraph graph = reinterpret_cast(graph_address); + + size_t node_count = 0; + check_driver(cuGraphGetNodes(graph, nullptr, &node_count), "cuGraphGetNodes(count)"); + std::vector nodes(node_count); + check_driver(cuGraphGetNodes(graph, nodes.data(), &node_count), "cuGraphGetNodes"); + + std::unordered_map node_indices; + for (size_t index = 0; index < node_count; ++index) { + node_indices[nodes[index]] = static_cast(index); + } + + size_t edge_count = 0; +#if CUDA_VERSION >= 13000 + check_driver(cuGraphGetEdges(graph, nullptr, nullptr, nullptr, &edge_count), + "cuGraphGetEdges(count)"); +#else + check_driver(cuGraphGetEdges_v2(graph, nullptr, nullptr, nullptr, &edge_count), + "cuGraphGetEdges_v2(count)"); +#endif + + std::vector from_nodes(edge_count); + std::vector to_nodes(edge_count); + std::vector edge_data(edge_count); +#if CUDA_VERSION >= 13000 + check_driver( + cuGraphGetEdges(graph, from_nodes.data(), to_nodes.data(), edge_data.data(), &edge_count), + "cuGraphGetEdges"); +#else + check_driver( + cuGraphGetEdges_v2(graph, from_nodes.data(), to_nodes.data(), edge_data.data(), &edge_count), + "cuGraphGetEdges_v2"); +#endif + + std::vector> result; + result.reserve(edge_count); + for (size_t index = 0; index < edge_count; ++index) { + result.emplace_back(node_indices.at(from_nodes[index]), node_indices.at(to_nodes[index]), + edge_data[index].from_port, edge_data[index].to_port, + edge_data[index].type); + } + return result; +} + +uint64_t shared_raw_cuda_graph(uint64_t foundry_graph_address) { + auto* graph = reinterpret_cast(foundry_graph_address); + TORCH_CHECK(graph->on_demand_data_ != nullptr, "graph has no on-demand data"); + TORCH_CHECK(graph->on_demand_data_->shared_exec != nullptr, "graph has no shared executor"); + TORCH_CHECK(graph->on_demand_data_->shared_exec->graph != nullptr, + "shared executor has no CUDA graph"); + return reinterpret_cast(graph->on_demand_data_->shared_exec->graph); +} + +std::string read_binary_graph_json(const std::string& path) { + return boost::json::serialize(foundry::read_and_parse_binary_graph(path)); +} + +bool binary_graph_valid(const std::string& path) { + return foundry::read_binary_graph_file(path).valid(); +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, module) { + module.def("launch_programmatic", &launch_programmatic); + module.def("launch_default", &launch_default); + module.def("graph_edges", &graph_edges); + module.def("shared_raw_cuda_graph", &shared_raw_cuda_graph); + module.def("read_binary_graph_json", &read_binary_graph_json); + module.def("binary_graph_valid", &binary_graph_valid); +} diff --git a/tests/run_deepep_matrix.sh b/tests/run_deepep_matrix.sh new file mode 100755 index 00000000..c0db658a --- /dev/null +++ b/tests/run_deepep_matrix.sh @@ -0,0 +1,105 @@ +#!/usr/bin/env bash +# Matrix driver for tests/test_deepep_fabric.py — standalone DeepEP save/load probes +# (no vLLM/sglang init; needs ~4 GB free on each of GPUs 0,1, NOT a fully empty GPU). +# +# ll_nofabric pure low-latency (NVSHMEM heap, FD handles), NCCL fast paths off +# -> the "EP without fabric" baseline; expected PASS +# ll_ncclcumem same, but NCCL_CUMEM_ENABLE=1 -> probe: can NCCL cumem coexist? +# ll_ncclnvls same, but CUMEM=1 + NVLS=1 -> probe: can NCCL NVLS coexist? +# nvl_ipc + 64 MB NVL buffer, use_fabric=0 -> legacy cudaIpc path through the +# hook's VMM-IPC translation (SCM_RIGHTS fd transport; peer mappings +# relocate under shared region bases); expected PASS +# nvl_fabric + 64 MB NVL buffer, use_fabric=1 -> fabric cuMemCreate on a no-IMEX +# machine; documents the fabric dependency; expected FAIL +# +# Usage: run_deepep_matrix.sh +# +# Environment overrides: +# PYTHON_BIN Python with Foundry installed (default: python3) +# VLLM_ROOT vLLM checkout used to build NVSHMEM +# NVSHMEM_SO explicit libnvshmem_host.so path +# DEEPEP_MATRIX_WORK_ROOT per-case work directory root +# DEEPEP_MATRIX_LOG_DIR log output directory +set -euo pipefail + +SCRIPT_DIR=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd) +REPO_ROOT=$(cd -- "$SCRIPT_DIR/.." && pwd) +PYTHON_BIN=${PYTHON_BIN:-python3} +VLLM_ROOT=${VLLM_ROOT:-"$REPO_ROOT/../vllm"} +NVSHMEM_SO=${NVSHMEM_SO:-"$VLLM_ROOT/tools/ep_kernels/ep_kernels_workspace/nvshmem/lib/libnvshmem_host.so"} +DEEPEP_MATRIX_WORK_ROOT=${DEEPEP_MATRIX_WORK_ROOT:-"$REPO_ROOT/tests/deepep_matrix_work"} +DEEPEP_MATRIX_LOG_DIR=${DEEPEP_MATRIX_LOG_DIR:-"$REPO_ROOT/logs/deepep_matrix"} +TEST="$REPO_ROOT/tests/test_deepep_fabric.py" + +CASE=${1:?usage: run_deepep_matrix.sh } +PHASE=${2:?usage: run_deepep_matrix.sh } + +export CUDA_VISIBLE_DEVICES=0,1 +# Pin NVSHMEM heap chunks to POSIX FD handles — explicit "no fabric handles anywhere". +# (Auto-detect picks FD on non-MNNVL machines anyway; pinning removes the variable.) +export NVSHMEM_CUMEM_HANDLE_TYPE=FILE_DESCRIPTOR + +case "$CASE" in + ll_nofabric) export TEST_USE_FABRIC=0 TEST_NVL_BYTES_MB=0 NCCL_CUMEM_ENABLE=0 NCCL_NVLS_ENABLE=0; PORT=29611 ;; + ll_ncclcumem) export TEST_USE_FABRIC=0 TEST_NVL_BYTES_MB=0 NCCL_CUMEM_ENABLE=1 NCCL_NVLS_ENABLE=0; PORT=29621 ;; + ll_ncclcumem_preinit) export TEST_USE_FABRIC=0 TEST_NVL_BYTES_MB=0 NCCL_CUMEM_ENABLE=1 NCCL_NVLS_ENABLE=0 TEST_NCCL_WARMUP_PRE_REGION=1; PORT=29661 ;; + ll_ncclnvls) export TEST_USE_FABRIC=0 TEST_NVL_BYTES_MB=0 NCCL_CUMEM_ENABLE=1 NCCL_NVLS_ENABLE=1; PORT=29631 ;; + ll_ncclnvls_preinit) export TEST_USE_FABRIC=0 TEST_NVL_BYTES_MB=0 NCCL_CUMEM_ENABLE=1 NCCL_NVLS_ENABLE=1 TEST_NCCL_WARMUP_PRE_REGION=1; PORT=29671 ;; + nvl_ipc) export TEST_USE_FABRIC=0 TEST_NVL_BYTES_MB=64 NCCL_CUMEM_ENABLE=0 NCCL_NVLS_ENABLE=0; PORT=29641 ;; + nvl_fabric) export TEST_USE_FABRIC=1 TEST_NVL_BYTES_MB=64 NCCL_CUMEM_ENABLE=0 NCCL_NVLS_ENABLE=0; PORT=29651 ;; + nvl_ipc_prealloc) export TEST_USE_FABRIC=0 TEST_NVL_BYTES_MB=64 NCCL_CUMEM_ENABLE=0 NCCL_NVLS_ENABLE=0 TEST_LOAD_PREALLOC_MB=1024; PORT=29681 ;; + *) echo "unknown case: $CASE" >&2; exit 2 ;; +esac + +case "$PHASE" in + save|load|both) ;; + *) echo "unknown phase: $PHASE" >&2; exit 2 ;; +esac + +if [[ ! -f "$NVSHMEM_SO" ]]; then + echo "NVSHMEM library not found: $NVSHMEM_SO" >&2 + exit 4 +fi + +HOOK_SO=$( + "$PYTHON_BIN" -c \ + "import foundry.ops, pathlib; print(pathlib.Path(foundry.ops.__file__).parent / 'libcuda_hook.so')" +) +if [[ ! -f "$HOOK_SO" ]]; then + echo "Foundry CUDA hook not found: $HOOK_SO" >&2 + exit 5 +fi + +# GPUs are shared with other users — refuse to run without headroom. +for i in 0 1; do + free=$(nvidia-smi --query-gpu=memory.free --format=csv,noheader,nounits -i "$i") + if [ "$free" -lt 6000 ]; then + echo "GPU $i has only ${free} MiB free — aborting (need ~6 GB headroom)"; exit 3 + fi +done + +# main() overwrites TEST_USE_FABRIC from the CLI flag, so pass the flag too. +FABRIC_ARGS=() +[ "$TEST_USE_FABRIC" = "0" ] && FABRIC_ARGS+=("--no-fabric") + +WORK="$DEEPEP_MATRIX_WORK_ROOT/$CASE" +LOGDIR="$DEEPEP_MATRIX_LOG_DIR" +mkdir -p "$WORK" "$LOGDIR" +cd "$WORK" # deepep_fabric_archive/ is CWD-relative + +export LD_PRELOAD="$NVSHMEM_SO:$HOOK_SO${LD_PRELOAD:+:$LD_PRELOAD}" + +run_phase() { + local p=$1 + # distinct rendezvous port per phase to dodge TIME_WAIT + if [ "$p" = save ]; then export MASTER_PORT=$PORT; else export MASTER_PORT=$((PORT + 1)); fi + echo "=== case=$CASE phase=$p $(date '+%F %T') ===" + echo "env: TEST_USE_FABRIC=$TEST_USE_FABRIC TEST_NVL_BYTES_MB=$TEST_NVL_BYTES_MB" \ + "NCCL_CUMEM_ENABLE=$NCCL_CUMEM_ENABLE NCCL_NVLS_ENABLE=$NCCL_NVLS_ENABLE" \ + "NVSHMEM_CUMEM_HANDLE_TYPE=$NVSHMEM_CUMEM_HANDLE_TYPE MASTER_PORT=$MASTER_PORT" + if [ "$p" = save ]; then rm -rf deepep_fabric_archive; fi + "$PYTHON_BIN" "$TEST" "--$p" "${FABRIC_ARGS[@]}" --num-processes=2 \ + 2>&1 | tee "$LOGDIR/${CASE}_${p}.log" +} + +if [ "$PHASE" = both ]; then run_phase save; run_phase load; else run_phase "$PHASE"; fi diff --git a/tests/test_deepep_fabric.py b/tests/test_deepep_fabric.py index be6be723..913b569b 100644 --- a/tests/test_deepep_fabric.py +++ b/tests/test_deepep_fabric.py @@ -30,6 +30,30 @@ ARCHIVE_DIR = "deepep_fabric_archive" HOOK_ARCHIVE_DIR = "hook_archive" +# Token payload encoding used by the deterministic dispatch check: +# x[token, :] = _TOKEN_OFFSET + rank * _RANK_STRIDE + token +# _verify_dispatch_result decodes each received payload back to (src_rank, +# token_id) and asserts that this rank's global expert id is in the token's +# deterministic top-k selection {(token_id + k) % num_experts}. +# +# The payload must be exactly representable in bfloat16 (the dispatch dtype), +# which only holds for integers up to 256. With the 64 tokens/rank and 2 ranks +# this test dispatches, _RANK_STRIDE=64 (== tokens per rank) and a +1 offset keep +# every payload in [1, 128]: exact in bf16 and never zero, so real tokens are +# distinguishable from DeepEP's (non-zeroed) padding slots. The stride must be +# >= tokens per rank so token ids stay recoverable via modulo. +_RANK_STRIDE = 64 +_TOKEN_OFFSET = 1 + + +def _encode_token(rank: int, token: int) -> float: + return float(_TOKEN_OFFSET + rank * _RANK_STRIDE + token) + + +def _decode_token(value: float) -> tuple[int, int]: + raw = int(round(float(value))) - _TOKEN_OFFSET + return raw // _RANK_STRIDE, raw % _RANK_STRIDE + def _get_hook_so_path(): import importlib.util @@ -77,6 +101,27 @@ def _init_dist(local_rank: int, num_local_ranks: int): ) +def _maybe_warm_nccl_pre_region(group): + """Diagnostic knob (TEST_NCCL_WARMUP_PRE_REGION=1): force NCCL transport + setup BEFORE the foundry region exists, so NCCL's cuMem allocations + (NCCL_CUMEM_ENABLE=1) pass through untracked. Localizes whether the + NCCL-CUMEM incompatibility is "NCCL cuMem while region enabled" only. + The warmup tensor's caching-allocator segment is drained afterwards so + post-region allocations don't reuse an untracked (non-deterministic) + segment. + """ + if os.environ.get("TEST_NCCL_WARMUP_PRE_REGION", "0") != "1": + return + t = torch.ones(1024, 1024, device="cuda") + dist.all_reduce(t) # default group: ring/p2p transport setup + dist.all_reduce(t, group=group) # subgroup used by the DeepEP Buffer + dist.all_gather_object([None] * dist.get_world_size(), 0) # object-coll path + torch.cuda.synchronize() + del t + torch.cuda.empty_cache() # legal here: region not set yet, no captures + print("[TEST] NCCL pre-region warmup done (transports up before region)", flush=True) + + def _create_deterministic_inputs( rank: int, num_tokens: int, hidden: int, num_experts: int, num_topk: int ): @@ -87,10 +132,10 @@ def _create_deterministic_inputs( """ import deep_ep - # x[token_i, :] = rank * 1000 + token_i (easy to trace which rank/token) + # x[token_i, :] = _encode_token(rank, token_i) (traceable to which rank/token) x = torch.zeros((num_tokens, hidden), dtype=torch.bfloat16, device="cuda") for i in range(num_tokens): - x[i, :] = float(rank * 1000 + i) + x[i, :] = _encode_token(rank, i) # Deterministic topk_idx: token i selects experts [i % num_experts, (i+1) % num_experts, ...] topk_idx = torch.zeros((num_tokens, num_topk), dtype=deep_ep.topk_idx_t, device="cuda") @@ -148,85 +193,110 @@ def _print_buffer_info(buffer, rank: int, prefix: str): def _verify_dispatch_result( recv_x, - recv_topk_idx, - recv_src_idx, - num_recv, + recv_count, rank: int, num_ranks: int, num_tokens: int, - hidden: int, num_experts: int, + num_topk: int, prefix: str, ): - """Verify dispatch results are correct based on our deterministic pattern.""" + """Assert dispatched tokens obey the deterministic routing invariant, with + EXACT-delivery checking (no missing, no duplicate, no spurious tokens). + + Inputs are seeded via ``_encode_token`` and token ``t`` selects experts + ``{(t + k) % num_experts : k in range(num_topk)}`` identically on every rank + (the selection depends only on the token id). So for a local expert ``e`` the + exact set of deliveries it must receive is + ``{(r, t) : r in range(num_ranks), t in range(num_tokens), e in selected(t)}`` + - each source rank contributes every token that selected ``e``, exactly once. + + ``recv_x`` is shaped + ``[num_local_experts, num_max_dispatch_tokens_per_rank * num_ranks, hidden]`` + and only the first ``recv_count[e]`` slots of local expert ``e`` are valid + (DeepEP does not zero the padding). For every valid slot, decode the payload + back to ``(src_rank, token_id)``, assert it is in range and routed correctly, + then require the decoded multiset for ``e`` to equal the expected set exactly: + ``recv_count[e]`` must match the expected delivery count and there must be no + duplicate or missing ``(src_rank, token_id)`` pair. Returns the number of + verified tokens and requires at least one. + """ + assert recv_x is not None, f"{prefix}: recv_x is None" + assert recv_x.dim() == 3, f"{prefix}: unexpected recv_x shape {tuple(recv_x.shape)}" + num_local_experts = num_experts // num_ranks local_expert_start = rank * num_local_experts - local_expert_end = local_expert_start + num_local_experts - - # num_recv might be an EventOverlap object (async), need to sync and get value - actual_num_recv = num_recv - if hasattr(num_recv, "wait"): - # It's an async object, wait for it - num_recv.wait() - if hasattr(num_recv, "value"): - actual_num_recv = num_recv.value - elif hasattr(num_recv, "item"): - actual_num_recv = num_recv.item() - elif not isinstance(num_recv, int): - # Try to get from recv_x shape - actual_num_recv = ( - recv_x.shape[0] * recv_x.shape[1] - if recv_x is not None and len(recv_x.shape) >= 2 - else 0 - ) - print( - f"[Rank {rank}] {prefix}: num_recv is {type(num_recv)}, using recv_x shape to estimate", - flush=True, - ) - print(f"[Rank {rank}] {prefix}: Verification - num_recv = {actual_num_recv}", flush=True) + counts = [int(c) for c in recv_count.detach().to("cpu").reshape(-1).tolist()] + assert len(counts) == num_local_experts, ( + f"{prefix}: recv_count has {len(counts)} entries, expected {num_local_experts}" + ) + print( - f"[Rank {rank}] {prefix}: Verification - local experts range = [{local_expert_start}, {local_expert_end})", + f"[Rank {rank}] {prefix}: recv_x shape = {tuple(recv_x.shape)}, " + f"local experts [{local_expert_start}, {local_expert_start + num_local_experts}), " + f"counts = {counts}", flush=True, ) - # recv_x shape is [num_local_experts, max_tokens_per_expert, hidden] - # Print shape info for debugging - if recv_x is not None: - print(f"[Rank {rank}] {prefix}: recv_x shape = {recv_x.shape}", flush=True) - - # Check some values in recv_x to verify the pattern - if recv_x is not None and recv_x.numel() > 0: - # Check first expert's first few tokens - num_experts_in_recv = recv_x.shape[0] - tokens_per_expert = recv_x.shape[1] - check_experts = min(2, num_experts_in_recv) - check_tokens = min(3, tokens_per_expert) - - print( - f"[Rank {rank}] {prefix}: Checking first {check_experts} experts, {check_tokens} tokens each:", - flush=True, + verified = 0 + for expert_i in range(num_local_experts): + global_expert = local_expert_start + expert_i + valid = counts[expert_i] + assert 0 <= valid <= recv_x.shape[1], ( + f"{prefix}: expert {global_expert} count {valid} out of range [0, {recv_x.shape[1]}]" ) - for expert_i in range(check_experts): - expert_global_idx = local_expert_start + expert_i - for token_i in range(check_tokens): - recv_val = recv_x[expert_i, token_i, 0].item() + # The exact set of (src_rank, token_id) deliveries this expert must get. + expected = { + (r, t) + for r in range(num_ranks) + for t in range(num_tokens) + if global_expert in {(t + k) % num_experts for k in range(num_topk)} + } + assert valid == len(expected), ( + f"{prefix}: expert {global_expert} recv_count {valid} != expected " + f"delivery count {len(expected)} (missing or duplicate dispatches)" + ) - # Decode: src_rank = recv_val // 1000, token_id = recv_val % 1000 - decoded_src_rank = int(recv_val) // 1000 - decoded_token_id = int(recv_val) % 1000 + received = [] + for slot in range(valid): + value = recv_x[expert_i, slot, 0].item() + src_rank, token_id = _decode_token(value) + assert 0 <= src_rank < num_ranks, ( + f"{prefix}: expert {global_expert} slot {slot} value {value} " + f"decodes to src_rank {src_rank} outside [0, {num_ranks})" + ) + assert 0 <= token_id < num_tokens, ( + f"{prefix}: expert {global_expert} slot {slot} value {value} " + f"decodes to token {token_id} outside [0, {num_tokens})" + ) + selected = {(token_id + k) % num_experts for k in range(num_topk)} + assert global_expert in selected, ( + f"{prefix}: expert {global_expert} received token {token_id} from " + f"rank {src_rank} whose top-k {sorted(selected)} does not include it" + ) + received.append((src_rank, token_id)) + verified += 1 - print( - f"[Rank {rank}] {prefix}: expert[{expert_i}] (global {expert_global_idx}), " - f"token[{token_i}]: value={recv_val:.1f} " - f"(decoded: src_rank={decoded_src_rank}, token_id={decoded_token_id})", - flush=True, - ) - else: - print(f"[Rank {rank}] {prefix}: recv_x is empty or None", flush=True) + received_set = set(received) + assert len(received_set) == len(received), ( + f"{prefix}: expert {global_expert} received duplicate deliveries: " + f"{sorted(p for p in received_set if received.count(p) > 1)}" + ) + assert received_set == expected, ( + f"{prefix}: expert {global_expert} delivery set mismatch; " + f"missing={sorted(expected - received_set)} " + f"unexpected={sorted(received_set - expected)}" + ) - return True + assert verified > 0, f"{prefix}: no valid tokens verified (counts={counts})" + print( + f"[Rank {rank}] {prefix}: verified {verified} dispatched tokens " + f"across {num_local_experts} local experts", + flush=True, + ) + return verified def _run_save(local_rank: int, num_processes: int): @@ -263,6 +333,8 @@ def _run_save(local_rank: int, num_processes: int): num_experts = num_ranks * 4 # 4 experts per rank num_topk = 2 + _maybe_warm_nccl_pre_region(group) + # Setup allocation region region_size = fdry.parse_size(REGION_SIZE_STR) print(f"[Rank {rank}] SAVE: Setting up allocation region at {hex(BASE_ADDR)}", flush=True) @@ -304,11 +376,11 @@ def _run_save(local_rank: int, num_processes: int): print(f"[Rank {rank}] SAVE: Input tensors:", flush=True) print(f"[Rank {rank}] SAVE: x address = {hex(x.data_ptr())}, shape = {x.shape}", flush=True) print( - f"[Rank {rank}] SAVE: x[0, :5] = {x[0, :5].tolist()} (expect {rank * 1000 + 0})", + f"[Rank {rank}] SAVE: x[0, :5] = {x[0, :5].tolist()} (expect {_encode_token(rank, 0)})", flush=True, ) print( - f"[Rank {rank}] SAVE: x[1, :5] = {x[1, :5].tolist()} (expect {rank * 1000 + 1})", + f"[Rank {rank}] SAVE: x[1, :5] = {x[1, :5].tolist()} (expect {_encode_token(rank, 1)})", flush=True, ) print(f"[Rank {rank}] SAVE: topk_idx address = {hex(topk_idx.data_ptr())}", flush=True) @@ -330,8 +402,10 @@ def _run_save(local_rank: int, num_processes: int): ) torch.cuda.synchronize() - # Unpack dispatch result - recv_x, recv_topk_idx, recv_src_idx, num_recv, *rest = result + # Unpack dispatch result. low_latency_dispatch returns + # (recv_x, recv_count, handle, event, hook); recv_count is a + # [num_local_experts] tensor of per-expert valid token counts. + recv_x, recv_count, *rest = result print(f"[Rank {rank}] SAVE: Warmup dispatch result:", flush=True) print( f"[Rank {rank}] SAVE: recv_x address = {hex(recv_x.data_ptr()) if recv_x is not None else 'None'}", @@ -341,19 +415,17 @@ def _run_save(local_rank: int, num_processes: int): f"[Rank {rank}] SAVE: recv_x shape = {recv_x.shape if recv_x is not None else 'None'}", flush=True, ) - print(f"[Rank {rank}] SAVE: num_recv = {num_recv}", flush=True) + print(f"[Rank {rank}] SAVE: recv_count = {recv_count.tolist()}", flush=True) # Verify results _verify_dispatch_result( recv_x, - recv_topk_idx, - recv_src_idx, - num_recv, + recv_count, rank, num_ranks, num_tokens, - hidden, num_experts, + num_topk, "SAVE-warmup", ) @@ -487,11 +559,22 @@ def _run_load(local_rank: int, num_processes: int): print(f"[Rank {rank}] LOAD: Loading CUDA modules from {rank_archive}", flush=True) fdry.load_cuda_modules_and_libraries(rank_archive) + _maybe_warm_nccl_pre_region(group) + # Step 2: Setup allocation region region_size = fdry.parse_size(REGION_SIZE_STR) print(f"[Rank {rank}] LOAD: Setting up allocation region at {hex(BASE_ADDR)}", flush=True) fdry.set_allocation_region(BASE_ADDR, region_size) + # Optional (TEST_LOAD_PREALLOC_MB=N): reproduce the production LOAD shape - + # one upfront-preallocated chunk that subsequent allocations (including the + # DeepEP NVL buffer) carve from with NO individual cuMem handle. This + # exercises the hook's whole-chunk VMM-IPC export path. + prealloc_mb = int(os.environ.get("TEST_LOAD_PREALLOC_MB", "0")) + if prealloc_mb > 0: + assert fdry.preallocate_region(prealloc_mb * 1024 * 1024), "preallocate_region failed" + print(f"[Rank {rank}] LOAD: preallocated {prealloc_mb} MB chunk", flush=True) + # Step 3: Create DeepEP buffer (initializes NVSHMEM runtime) num_local_experts = num_experts // num_ranks num_rdma_bytes = deep_ep.Buffer.get_low_latency_rdma_size_hint( @@ -499,13 +582,17 @@ def _run_load(local_rank: int, num_processes: int): ) use_fabric = os.environ.get("TEST_USE_FABRIC", "1") == "1" + # Must match SAVE: an NVL buffer allocated on SAVE but not LOAD (or vice versa) + # diverges the allocation trajectory and the peer-pointer setup. + num_nvl_bytes = int(os.environ.get("TEST_NVL_BYTES_MB", "0")) * 1024 * 1024 print( - f"[Rank {rank}] LOAD: Creating DeepEP buffer with use_fabric={use_fabric}, size={num_rdma_bytes / 1e6:.2f} MB", + f"[Rank {rank}] LOAD: Creating DeepEP buffer with use_fabric={use_fabric}, num_nvl_bytes={num_nvl_bytes / 1e6:.2f} MB, num_rdma_bytes={num_rdma_bytes / 1e6:.2f} MB", flush=True, ) buffer = deep_ep.Buffer( group, + num_nvl_bytes=num_nvl_bytes, num_rdma_bytes=num_rdma_bytes, low_latency_mode=True, num_qps_per_rank=num_local_experts, @@ -530,11 +617,11 @@ def _run_load(local_rank: int, num_processes: int): print(f"[Rank {rank}] LOAD: Input tensors:", flush=True) print(f"[Rank {rank}] LOAD: x address = {hex(x.data_ptr())}, shape = {x.shape}", flush=True) print( - f"[Rank {rank}] LOAD: x[0, :5] = {x[0, :5].tolist()} (expect {rank * 1000 + 0})", + f"[Rank {rank}] LOAD: x[0, :5] = {x[0, :5].tolist()} (expect {_encode_token(rank, 0)})", flush=True, ) print( - f"[Rank {rank}] LOAD: x[1, :5] = {x[1, :5].tolist()} (expect {rank * 1000 + 1})", + f"[Rank {rank}] LOAD: x[1, :5] = {x[1, :5].tolist()} (expect {_encode_token(rank, 1)})", flush=True, ) print(f"[Rank {rank}] LOAD: topk_idx address = {hex(topk_idx.data_ptr())}", flush=True) @@ -556,8 +643,9 @@ def _run_load(local_rank: int, num_processes: int): ) torch.cuda.synchronize() - # Unpack and verify warmup result - recv_x, recv_topk_idx, recv_src_idx, num_recv, *rest = warmup_result + # Unpack and verify warmup result. Same layout as SAVE: + # (recv_x, recv_count, handle, event, hook). + recv_x, recv_count, *rest = warmup_result print(f"[Rank {rank}] LOAD: Warmup dispatch result:", flush=True) print( f"[Rank {rank}] LOAD: recv_x address = {hex(recv_x.data_ptr()) if recv_x is not None else 'None'}", @@ -567,19 +655,17 @@ def _run_load(local_rank: int, num_processes: int): f"[Rank {rank}] LOAD: recv_x shape = {recv_x.shape if recv_x is not None else 'None'}", flush=True, ) - print(f"[Rank {rank}] LOAD: num_recv = {num_recv}", flush=True) + print(f"[Rank {rank}] LOAD: recv_count = {recv_count.tolist()}", flush=True) # Verify warmup results _verify_dispatch_result( recv_x, - recv_topk_idx, - recv_src_idx, - num_recv, + recv_count, rank, num_ranks, num_tokens, - hidden, num_experts, + num_topk, "LOAD-warmup", ) @@ -593,10 +679,21 @@ def _run_load(local_rank: int, num_processes: int): flush=True, ) - # Load graph from per-rank archive + # Load graph from per-rank archive. + # Default: the production path (start_graph_builds + finish_graph_loads), + # same machinery the vLLM/sglang integrations use. TEST_LOAD_API=single + # selects the legacy CUDAGraph.load path — KNOWN BROKEN for DeepEP graphs: + # it does not merge root-level common_kernel_node_attrs, so factored-out + # cooperative/clusterDim attrs are dropped and the dispatch kernel faults + # with "unspecified launch failure" at replay. graph_json = os.path.join(rank_archive, "deepep_dispatch_graph.json") - print(f"[Rank {rank}] LOAD: Loading graph from {graph_json}", flush=True) - graph, output_tensors = fdry.CUDAGraph.load(graph_json) + load_api = os.environ.get("TEST_LOAD_API", "parallel") + print(f"[Rank {rank}] LOAD: Loading graph from {graph_json} (api={load_api})", flush=True) + if load_api == "single": + graph, output_tensors = fdry.CUDAGraph.load(graph_json) + else: + pending = fdry.CUDAGraph.start_graph_builds([graph_json]) + ((graph, output_tensors),) = fdry.CUDAGraph.finish_graph_loads(pending) print(f"[Rank {rank}] LOAD: Graph loaded successfully", flush=True) # Print loaded output tensor info @@ -624,29 +721,22 @@ def _run_load(local_rank: int, num_processes: int): graph.replay() torch.cuda.synchronize() - # Verify after replay - check the LOADED output tensors (not warmup recv_x) - # recv_x shape is [num_local_experts, max_tokens_per_expert, hidden] + # Verify after replay against the LOADED output tensor (not warmup recv_x). + # The dispatch is deterministic, so the replayed graph reproduces the same + # per-expert routing as the warmup pass; reuse the warmup recv_count to bound + # the valid slots and assert the same routing invariant. print(f"[Rank {rank}] LOAD: After graph replay:", flush=True) verify_tensor = loaded_recv_x if loaded_recv_x is not None else recv_x - if verify_tensor is not None and verify_tensor.numel() > 0: - print(f"[Rank {rank}] LOAD: verify_tensor shape = {verify_tensor.shape}", flush=True) - num_local_experts = num_experts // num_ranks - local_expert_start = rank * num_local_experts - - # Check first expert's first few tokens - check_experts = min(2, verify_tensor.shape[0]) - check_tokens = min(3, verify_tensor.shape[1]) - - for expert_i in range(check_experts): - for token_i in range(check_tokens): - val = verify_tensor[expert_i, token_i, 0].item() - decoded_src_rank = int(val) // 1000 - decoded_token_id = int(val) % 1000 - print( - f"[Rank {rank}] LOAD: expert[{expert_i}], token[{token_i}]: " - f"value={val:.1f} (rank={decoded_src_rank}, token={decoded_token_id})", - flush=True, - ) + _verify_dispatch_result( + verify_tensor, + recv_count, + rank, + num_ranks, + num_tokens, + num_experts, + num_topk, + "LOAD-replay", + ) print(f"[Rank {rank}] LOAD: Graph replay successful", flush=True) diff --git a/tests/test_deepep_transport.py b/tests/test_deepep_transport.py new file mode 100644 index 00000000..c570dbb9 --- /dev/null +++ b/tests/test_deepep_transport.py @@ -0,0 +1,1155 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the Foundry project +from __future__ import annotations + +import json +import logging +import sys +import types +from types import SimpleNamespace + +import pytest +import torch +from foundry.integration.deepep_transport import ( + DeepEPTransport, + configure_deepep_buffer, + parse_deepep_transport, + validate_num_nvl_bytes_match, + validate_transport_match, +) +from foundry.integration.sglang import config as sglang_config +from foundry.integration.sglang import hooks as sglang_hooks +from foundry.integration.sglang import runtime as sglang_runtime + + +def _install_fake_deep_ep(monkeypatch, buffer_cls): + deep_ep_module = types.ModuleType("deep_ep") + deep_ep_module.Buffer = buffer_cls + monkeypatch.setitem(sys.modules, "deep_ep", deep_ep_module) + monkeypatch.setattr(sglang_hooks, "_DEEPEP_PATCHED", False, raising=False) + + +def _install_fake_sglang_cuda_graph_runner(monkeypatch, runner_cls): + sglang_module = types.ModuleType("sglang") + srt_module = types.ModuleType("sglang.srt") + model_executor_module = types.ModuleType("sglang.srt.model_executor") + cuda_graph_runner_module = types.ModuleType("sglang.srt.model_executor.cuda_graph_runner") + cuda_graph_runner_module.CudaGraphRunner = runner_cls + model_executor_module.cuda_graph_runner = cuda_graph_runner_module + srt_module.model_executor = model_executor_module + sglang_module.srt = srt_module + monkeypatch.setitem(sys.modules, "sglang", sglang_module) + monkeypatch.setitem(sys.modules, "sglang.srt", srt_module) + monkeypatch.setitem(sys.modules, "sglang.srt.model_executor", model_executor_module) + monkeypatch.setitem( + sys.modules, + "sglang.srt.model_executor.cuda_graph_runner", + cuda_graph_runner_module, + ) + + +def _install_fake_sglang_attention_tp_gather(monkeypatch, gather, group): + class GroupCoordinator: + def _all_gather_into_tensor(self, output, input_): + return None + + sglang_module = types.ModuleType("sglang") + srt_module = types.ModuleType("sglang.srt") + distributed_module = types.ModuleType("sglang.srt.distributed") + parallel_state_module = types.ModuleType("sglang.srt.distributed.parallel_state") + layers_module = types.ModuleType("sglang.srt.layers") + dp_attention_module = types.ModuleType("sglang.srt.layers.dp_attention") + logits_processor_module = types.ModuleType("sglang.srt.layers.logits_processor") + parallel_state_module.GroupCoordinator = GroupCoordinator + distributed_module.parallel_state = parallel_state_module + dp_attention_module.get_attention_tp_group = lambda: group + logits_processor_module.attn_tp_all_gather_into_tensor = gather + layers_module.dp_attention = dp_attention_module + layers_module.logits_processor = logits_processor_module + srt_module.distributed = distributed_module + srt_module.layers = layers_module + sglang_module.srt = srt_module + monkeypatch.setitem(sys.modules, "sglang", sglang_module) + monkeypatch.setitem(sys.modules, "sglang.srt", srt_module) + monkeypatch.setitem(sys.modules, "sglang.srt.distributed", distributed_module) + monkeypatch.setitem( + sys.modules, + "sglang.srt.distributed.parallel_state", + parallel_state_module, + ) + monkeypatch.setitem(sys.modules, "sglang.srt.layers", layers_module) + monkeypatch.setitem( + sys.modules, + "sglang.srt.layers.dp_attention", + dp_attention_module, + ) + monkeypatch.setitem( + sys.modules, + "sglang.srt.layers.logits_processor", + logits_processor_module, + ) + return logits_processor_module + + +def _install_fake_sglang_regular_tp_gather(monkeypatch, gather, group): + logits_processor_module = _install_fake_sglang_attention_tp_gather( + monkeypatch, + lambda output, input_: group._all_gather_into_tensor(output, input_), + group, + ) + distributed_module = types.ModuleType("sglang.srt.distributed") + parallel_state_module = types.ModuleType("sglang.srt.distributed.parallel_state") + parallel_state_module.GroupCoordinator = type(group) + distributed_module.parallel_state = parallel_state_module + sys.modules["sglang.srt"].distributed = distributed_module + logits_processor_module.tensor_model_parallel_all_gather = gather + monkeypatch.setitem( + sys.modules, + "sglang.srt.distributed", + distributed_module, + ) + monkeypatch.setitem( + sys.modules, + "sglang.srt.distributed.parallel_state", + parallel_state_module, + ) + return logits_processor_module + + +@pytest.mark.parametrize( + ("raw", "expected"), + [ + (None, DeepEPTransport.FABRIC), + ("fabric", DeepEPTransport.FABRIC), + ("nvl_ipc", DeepEPTransport.NVL_IPC), + ], +) +def test_parse_deepep_transport(raw, expected): + assert parse_deepep_transport(raw) is expected + + +def test_parse_deepep_transport_rejects_unknown_value(): + with pytest.raises(ValueError, match="deepep_transport"): + parse_deepep_transport("rdma") + + +@pytest.mark.parametrize( + ("transport", "expected_bytes", "expected_fabric"), + [ + (DeepEPTransport.FABRIC, 0, True), + (DeepEPTransport.NVL_IPC, 128, False), + ], +) +def test_configure_deepep_buffer(transport, expected_bytes, expected_fabric): + num_nvl_bytes, kwargs = configure_deepep_buffer( + transport, + num_nvl_bytes=128, + kwargs={"other": "kept"}, + ) + assert num_nvl_bytes == expected_bytes + assert kwargs == {"other": "kept", "use_fabric": expected_fabric} + + +def test_transport_mismatch_names_both_values(): + with pytest.raises(RuntimeError, match="archived=nvl_ipc.*configured=fabric"): + validate_transport_match( + archived="nvl_ipc", + configured=DeepEPTransport.FABRIC, + ) + + +def test_nvl_bytes_mismatch_names_both_values(): + with pytest.raises(RuntimeError, match="archived=64.*effective=128"): + validate_num_nvl_bytes_match( + archived=64, + effective=128, + transport=DeepEPTransport.NVL_IPC, + ) + + +@pytest.mark.parametrize( + ("contents", "expected"), + [ + ("", DeepEPTransport.FABRIC), + ('deepep_transport = "nvl_ipc"\n', DeepEPTransport.NVL_IPC), + ], +) +def test_sglang_config_parses_deepep_transport(tmp_path, contents, expected): + config_path = tmp_path / "config.toml" + config_path.write_text(contents) + + config = sglang_config.CUDAGraphExtensionConfig.from_toml(config_path) + + assert config.deepep_transport is expected + + +def test_sglang_config_rejects_unknown_deepep_transport(tmp_path): + config_path = tmp_path / "config.toml" + config_path.write_text('deepep_transport = "rdma"\n') + + with pytest.raises(ValueError, match="deepep_transport"): + sglang_config.CUDAGraphExtensionConfig.from_toml(config_path) + + +def test_sglang_get_deepep_transport_uses_loaded_config(tmp_path, monkeypatch): + config_path = tmp_path / "config.toml" + config_path.write_text('deepep_transport = "nvl_ipc"\n') + monkeypatch.setattr(sglang_config, "_config", None) + + assert sglang_config.get_deepep_transport() is DeepEPTransport.FABRIC + + sglang_config.load_graph_extension_config(str(config_path)) + + assert sglang_config.get_deepep_transport() is DeepEPTransport.NVL_IPC + + +def test_sglang_save_capture_supports_out_graph_attention_api(monkeypatch): + def original_out_graph(forward_batch, in_capture=False): + return None + + class CudaGraphRunner: + def __init__(self): + self.attn_backend = SimpleNamespace( + init_forward_metadata_out_graph=original_out_graph, + ) + + def capture(self): + return "captured" + + def _create_device_graph(self): + return None + + def _capture_graph(self, graph, pool, stream, run_once_fn): + return run_once_fn() + + def capture_one_batch_size(self, bs, forward, stream_idx=None): + return None, forward() + + _install_fake_sglang_cuda_graph_runner(monkeypatch, CudaGraphRunner) + graph_ops = types.ModuleType("foundry.integration.sglang.graph_ops") + graph_ops.pack_fatbins = lambda: None + graph_ops.save_graph_manifest = lambda: None + monkeypatch.setitem( + sys.modules, + "foundry.integration.sglang.graph_ops", + graph_ops, + ) + monkeypatch.setattr(sglang_hooks, "_ep_lazy_init_needed", lambda: False) + monkeypatch.setattr( + sglang_hooks, + "get_graph_extension_mode", + lambda: sglang_hooks.CUDAGraphExtensionMode.SAVE, + ) + monkeypatch.setattr(sglang_hooks.rt, "capture_final_alloc_offset", lambda: 0) + + sglang_hooks._patch_cuda_graph_capture() + runner = CudaGraphRunner() + + assert runner.capture() == "captured" + assert runner.attn_backend.init_forward_metadata_out_graph is original_out_graph + + +def test_sglang_load_precopies_local_attention_tp_logits_before_gather(monkeypatch): + group = SimpleNamespace(rank_in_group=1, world_size=3) + calls = [] + delegate_result = object() + + def fake_gather(global_logits, logits): + calls.append( + { + "global_logits": global_logits, + "logits": logits, + "local_slice_before_delegate": global_logits[1].clone(), + } + ) + global_logits[0].fill_(10) + global_logits[2].fill_(30) + return delegate_result + + logits_processor = _install_fake_sglang_attention_tp_gather( + monkeypatch, + fake_gather, + group, + ) + monkeypatch.setattr( + sglang_hooks, + "get_graph_extension_mode", + lambda: sglang_hooks.CUDAGraphExtensionMode.LOAD, + ) + logits = torch.tensor([[1.0, 2.0], [3.0, 4.0]]) + global_logits = torch.full((3, 2, 2), -1.0) + + sglang_hooks._patch_attn_tp_logits_gather() + result = logits_processor.attn_tp_all_gather_into_tensor(global_logits, logits) + + assert result is delegate_result + assert len(calls) == 1 + assert calls[0]["global_logits"] is global_logits + assert calls[0]["logits"] is logits + torch.testing.assert_close(calls[0]["local_slice_before_delegate"], logits) + torch.testing.assert_close(global_logits[1], logits) + torch.testing.assert_close(global_logits[0], torch.full((2, 2), 10.0)) + torch.testing.assert_close(global_logits[2], torch.full((2, 2), 30.0)) + + +@pytest.mark.parametrize("output_layout", ["rank_major", "concatenated"]) +def test_sglang_load_precopies_local_regular_tp_logits_before_gather( + monkeypatch, + output_layout, +): + calls = [] + gather_result = object() + lower_result = object() + + class GroupCoordinator: + rank_in_group = 1 + world_size = 3 + + def _all_gather_into_tensor(self, output, input_): + output_shards = output.reshape(self.world_size, *input_.shape) + calls.append( + { + "output": output, + "input": input_, + "local_slice_before_delegate": output_shards[1].clone(), + } + ) + output_shards[0].fill_(10) + output_shards[2].fill_(30) + return lower_result + + group = GroupCoordinator() + + def fake_gather(logits, dim=-1): + assert dim == -1 + if output_layout == "rank_major": + output_shape = (group.world_size, *logits.shape) + else: + output_shape = ( + group.world_size * logits.shape[0], + *logits.shape[1:], + ) + output = torch.full(output_shape, -1.0) + assert group._all_gather_into_tensor(output, logits) is lower_result + return gather_result + + logits_processor = _install_fake_sglang_regular_tp_gather( + monkeypatch, + fake_gather, + group, + ) + monkeypatch.setattr( + sglang_hooks, + "get_graph_extension_mode", + lambda: sglang_hooks.CUDAGraphExtensionMode.LOAD, + ) + logits = torch.tensor([[1.0, 2.0], [3.0, 4.0]]) + + sglang_hooks._patch_attn_tp_logits_gather() + installed_gather = logits_processor.tensor_model_parallel_all_gather + installed_lower_gather = GroupCoordinator._all_gather_into_tensor + sglang_hooks._patch_attn_tp_logits_gather() + + assert logits_processor.tensor_model_parallel_all_gather is installed_gather + assert GroupCoordinator._all_gather_into_tensor is installed_lower_gather + result = logits_processor.tensor_model_parallel_all_gather(logits) + + assert result is gather_result + assert len(calls) == 1 + assert calls[0]["input"] is logits + torch.testing.assert_close(calls[0]["local_slice_before_delegate"], logits) + output_shards = calls[0]["output"].reshape(group.world_size, *logits.shape) + torch.testing.assert_close(output_shards[1], logits) + torch.testing.assert_close(output_shards[0], torch.full((2, 2), 10.0)) + torch.testing.assert_close(output_shards[2], torch.full((2, 2), 30.0)) + + +def test_sglang_load_patches_regular_tp_when_attention_helper_is_absent(monkeypatch): + calls = [] + delegate_result = object() + + class GroupCoordinator: + rank_in_group = 1 + world_size = 2 + + def _all_gather_into_tensor(self, output, input_): + calls.append((output, input_)) + return delegate_result + + group = GroupCoordinator() + logits_processor = _install_fake_sglang_regular_tp_gather( + monkeypatch, + lambda logits, dim=-1: logits, + group, + ) + monkeypatch.delattr(logits_processor, "attn_tp_all_gather_into_tensor") + original_lower_gather = GroupCoordinator._all_gather_into_tensor + monkeypatch.setattr( + sglang_hooks, + "get_graph_extension_mode", + lambda: sglang_hooks.CUDAGraphExtensionMode.LOAD, + ) + + sglang_hooks._patch_attn_tp_logits_gather() + + assert not hasattr(logits_processor, "attn_tp_all_gather_into_tensor") + assert GroupCoordinator._all_gather_into_tensor is not original_lower_gather + logits = torch.tensor([[1.0, 2.0]]) + output = torch.full((2, 1, 2), -1.0) + assert group._all_gather_into_tensor(output, logits) is delegate_result + torch.testing.assert_close(output[1], logits) + assert len(calls) == 1 + assert calls[0][0] is output + assert calls[0][1] is logits + + +def test_sglang_load_patches_attention_tp_when_group_coordinator_is_absent(monkeypatch): + group = SimpleNamespace(rank_in_group=1, world_size=2) + delegate_result = object() + + def fake_gather(global_logits, logits): + return delegate_result + + logits_processor = _install_fake_sglang_attention_tp_gather( + monkeypatch, + fake_gather, + group, + ) + parallel_state = sys.modules["sglang.srt.distributed.parallel_state"] + monkeypatch.delattr(parallel_state, "GroupCoordinator") + monkeypatch.setattr( + sglang_hooks, + "get_graph_extension_mode", + lambda: sglang_hooks.CUDAGraphExtensionMode.LOAD, + ) + + sglang_hooks._patch_attn_tp_logits_gather() + + installed = logits_processor.attn_tp_all_gather_into_tensor + assert installed is not fake_gather + logits = torch.tensor([[1.0, 2.0]]) + global_logits = torch.full((2, 1, 2), -1.0) + assert installed(global_logits, logits) is delegate_result + torch.testing.assert_close(global_logits[1], logits) + + +def test_sglang_load_skips_attention_tp_when_group_accessor_is_absent(monkeypatch): + calls = [] + + class GroupCoordinator: + rank_in_group = 0 + world_size = 2 + + def _all_gather_into_tensor(self, output, input_): + calls.append((output, input_)) + return None + + group = GroupCoordinator() + logits_processor = _install_fake_sglang_regular_tp_gather( + monkeypatch, + lambda logits, dim=-1: logits, + group, + ) + dp_attention = sys.modules["sglang.srt.layers.dp_attention"] + monkeypatch.delattr(dp_attention, "get_attention_tp_group") + original_attention_gather = logits_processor.attn_tp_all_gather_into_tensor + original_lower_gather = GroupCoordinator._all_gather_into_tensor + monkeypatch.setattr( + sglang_hooks, + "get_graph_extension_mode", + lambda: sglang_hooks.CUDAGraphExtensionMode.LOAD, + ) + + sglang_hooks._patch_attn_tp_logits_gather() + + assert logits_processor.attn_tp_all_gather_into_tensor is original_attention_gather + assert GroupCoordinator._all_gather_into_tensor is not original_lower_gather + logits = torch.tensor([[1.0, 2.0]]) + output = torch.full((2, 1, 2), -1.0) + group._all_gather_into_tensor(output, logits) + torch.testing.assert_close(output[0], logits) + assert len(calls) == 1 + assert calls[0][0] is output + assert calls[0][1] is logits + + +def test_sglang_load_logs_only_gather_seams_installed(monkeypatch, caplog): + class GroupCoordinator: + rank_in_group = 0 + world_size = 2 + + def _all_gather_into_tensor(self, output, input_): + return None + + group = GroupCoordinator() + logits_processor = _install_fake_sglang_regular_tp_gather( + monkeypatch, + lambda logits, dim=-1: logits, + group, + ) + monkeypatch.delattr(logits_processor, "attn_tp_all_gather_into_tensor") + monkeypatch.setattr( + sglang_hooks, + "get_graph_extension_mode", + lambda: sglang_hooks.CUDAGraphExtensionMode.LOAD, + ) + + with caplog.at_level(logging.INFO, logger=sglang_hooks.__name__): + sglang_hooks._patch_attn_tp_logits_gather() + + assert caplog.messages[-1] == ( + "[Foundry] SGLang LOAD logits pre-copy patches installed: regular_tp" + ) + + +@pytest.mark.parametrize("broken_api", ["accessor", "delegate"]) +def test_sglang_load_propagates_present_broken_attention_tp_api( + monkeypatch, + broken_api, +): + message = f"broken attention {broken_api}" + + def broken(*args, **kwargs): + raise AttributeError(message) + + group = SimpleNamespace(rank_in_group=0, world_size=2) + gather = broken if broken_api == "delegate" else lambda global_logits, logits: None + logits_processor = _install_fake_sglang_attention_tp_gather( + monkeypatch, + gather, + group, + ) + if broken_api == "accessor": + dp_attention = sys.modules["sglang.srt.layers.dp_attention"] + monkeypatch.setattr(dp_attention, "get_attention_tp_group", broken) + parallel_state = sys.modules["sglang.srt.distributed.parallel_state"] + monkeypatch.delattr(parallel_state, "GroupCoordinator") + monkeypatch.setattr( + sglang_hooks, + "get_graph_extension_mode", + lambda: sglang_hooks.CUDAGraphExtensionMode.LOAD, + ) + + sglang_hooks._patch_attn_tp_logits_gather() + + with pytest.raises(AttributeError, match=message): + logits_processor.attn_tp_all_gather_into_tensor( + torch.empty((2, 1, 2)), + torch.ones((1, 2)), + ) + + +def test_sglang_load_propagates_present_broken_regular_tp_api(monkeypatch): + message = "broken regular delegate" + + class GroupCoordinator: + rank_in_group = 0 + world_size = 2 + + def _all_gather_into_tensor(self, output, input_): + raise AttributeError(message) + + group = GroupCoordinator() + logits_processor = _install_fake_sglang_regular_tp_gather( + monkeypatch, + lambda logits, dim=-1: logits, + group, + ) + monkeypatch.delattr(logits_processor, "attn_tp_all_gather_into_tensor") + monkeypatch.setattr( + sglang_hooks, + "get_graph_extension_mode", + lambda: sglang_hooks.CUDAGraphExtensionMode.LOAD, + ) + + sglang_hooks._patch_attn_tp_logits_gather() + + with pytest.raises(AttributeError, match=message): + group._all_gather_into_tensor( + torch.empty((2, 1, 2)), + torch.ones((1, 2)), + ) + + +@pytest.mark.parametrize( + "mode", + [ + sglang_hooks.CUDAGraphExtensionMode.NONE, + sglang_hooks.CUDAGraphExtensionMode.SAVE, + ], +) +def test_sglang_non_load_modes_leave_regular_tp_gather_unchanged(monkeypatch, mode): + class GroupCoordinator: + rank_in_group = 0 + world_size = 2 + + def _all_gather_into_tensor(self, output, input_): + return None + + group = GroupCoordinator() + + def fake_gather(logits, dim=-1): + return logits + + logits_processor = _install_fake_sglang_regular_tp_gather( + monkeypatch, + fake_gather, + group, + ) + original_lower_gather = GroupCoordinator._all_gather_into_tensor + monkeypatch.setattr(sglang_hooks, "get_graph_extension_mode", lambda: mode) + + sglang_hooks._patch_attn_tp_logits_gather() + + assert logits_processor.tensor_model_parallel_all_gather is fake_gather + assert GroupCoordinator._all_gather_into_tensor is original_lower_gather + + +def test_sglang_load_regular_tp_gather_delegates_unsupported_shape_unchanged( + monkeypatch, +): + calls = [] + delegate_result = object() + + class GroupCoordinator: + rank_in_group = 1 + world_size = 3 + + def _all_gather_into_tensor(self, output, input_): + calls.append((output, input_)) + return delegate_result + + group = GroupCoordinator() + + def fake_gather(logits, dim=-1): + return logits + + _install_fake_sglang_regular_tp_gather( + monkeypatch, + fake_gather, + group, + ) + monkeypatch.setattr( + sglang_hooks, + "get_graph_extension_mode", + lambda: sglang_hooks.CUDAGraphExtensionMode.LOAD, + ) + sglang_hooks._patch_attn_tp_logits_gather() + gather = GroupCoordinator._all_gather_into_tensor + output = torch.full((3, 2, 3), -1.0) + input_ = torch.ones((2, 2)) + + result = gather(group, output, input_) + + assert result is delegate_result + assert len(calls) == 1 + assert calls[0][0] is output + assert calls[0][1] is input_ + torch.testing.assert_close(output, torch.full((3, 2, 3), -1.0)) + + +def test_sglang_load_regular_tp_gather_delegates_unviewable_layout_unchanged( + monkeypatch, +): + calls = [] + delegate_result = object() + + class Input: + shape = (2, 2) + + class Output: + shape = (6, 2) + + def view(self, *shape): + raise RuntimeError(f"unsupported layout for {shape}") + + class GroupCoordinator: + rank_in_group = 1 + world_size = 3 + + def _all_gather_into_tensor(self, output, input_): + calls.append((output, input_)) + return delegate_result + + group = GroupCoordinator() + _install_fake_sglang_regular_tp_gather( + monkeypatch, + lambda logits, dim=-1: logits, + group, + ) + monkeypatch.setattr( + sglang_hooks, + "get_graph_extension_mode", + lambda: sglang_hooks.CUDAGraphExtensionMode.LOAD, + ) + sglang_hooks._patch_attn_tp_logits_gather() + output = Output() + input_ = Input() + + result = group._all_gather_into_tensor(output, input_) + + assert result is delegate_result + assert len(calls) == 1 + assert calls[0][0] is output + assert calls[0][1] is input_ + + +def test_sglang_load_regular_tp_gather_rejects_invalid_rank(monkeypatch): + calls = [] + + class GroupCoordinator: + rank_in_group = 3 + world_size = 3 + + def _all_gather_into_tensor(self, output, input_): + calls.append((output, input_)) + + group = GroupCoordinator() + _install_fake_sglang_regular_tp_gather( + monkeypatch, + lambda logits, dim=-1: logits, + group, + ) + monkeypatch.setattr( + sglang_hooks, + "get_graph_extension_mode", + lambda: sglang_hooks.CUDAGraphExtensionMode.LOAD, + ) + sglang_hooks._patch_attn_tp_logits_gather() + gather = GroupCoordinator._all_gather_into_tensor + + with pytest.raises( + RuntimeError, + match=r"Foundry SGLang LOAD regular-TP gather.*rank_in_group=3.*world_size=3", + ): + gather(group, torch.empty((6, 2)), torch.empty((2, 2))) + + assert calls == [] + + +@pytest.mark.parametrize( + "mode", + [ + sglang_hooks.CUDAGraphExtensionMode.NONE, + sglang_hooks.CUDAGraphExtensionMode.SAVE, + ], +) +def test_sglang_non_load_modes_leave_attention_tp_gather_unchanged(monkeypatch, mode): + group = SimpleNamespace(rank_in_group=1, world_size=3) + delegate_result = object() + + def fake_gather(global_logits, logits): + return delegate_result + + logits_processor = _install_fake_sglang_attention_tp_gather( + monkeypatch, + fake_gather, + group, + ) + monkeypatch.setattr(sglang_hooks, "get_graph_extension_mode", lambda: mode) + + sglang_hooks._patch_attn_tp_logits_gather() + + assert logits_processor.attn_tp_all_gather_into_tensor is fake_gather + + +def test_sglang_load_attention_tp_gather_patch_is_idempotent(monkeypatch): + group = SimpleNamespace(rank_in_group=0, world_size=2) + calls = [] + + def fake_gather(global_logits, logits): + calls.append((global_logits, logits)) + return "delegate-result" + + logits_processor = _install_fake_sglang_attention_tp_gather( + monkeypatch, + fake_gather, + group, + ) + monkeypatch.setattr( + sglang_hooks, + "get_graph_extension_mode", + lambda: sglang_hooks.CUDAGraphExtensionMode.LOAD, + ) + + sglang_hooks._patch_attn_tp_logits_gather() + installed = logits_processor.attn_tp_all_gather_into_tensor + sglang_hooks._patch_attn_tp_logits_gather() + + assert logits_processor.attn_tp_all_gather_into_tensor is installed + assert installed(torch.empty((2, 1, 2)), torch.ones((1, 2))) == "delegate-result" + assert len(calls) == 1 + + +def test_sglang_load_attention_tp_gather_rejects_invalid_rank_and_shape(monkeypatch): + group = SimpleNamespace(rank_in_group=1, world_size=3) + calls = [] + + def fake_gather(global_logits, logits): + calls.append((global_logits, logits)) + + logits_processor = _install_fake_sglang_attention_tp_gather( + monkeypatch, + fake_gather, + group, + ) + monkeypatch.setattr( + sglang_hooks, + "get_graph_extension_mode", + lambda: sglang_hooks.CUDAGraphExtensionMode.LOAD, + ) + sglang_hooks._patch_attn_tp_logits_gather() + gather = logits_processor.attn_tp_all_gather_into_tensor + + with pytest.raises( + RuntimeError, + match=r"Foundry SGLang LOAD attention-TP gather.*output shape", + ): + gather(torch.empty((3, 2, 3)), torch.empty((2, 2))) + + group.rank_in_group = 3 + with pytest.raises( + RuntimeError, + match=r"Foundry SGLang LOAD attention-TP gather.*rank_in_group=3.*world_size=3", + ): + gather(torch.empty((3, 2, 2)), torch.empty((2, 2))) + + assert calls == [] + + +def test_sglang_warmup_state_round_trips_deepep_profile(tmp_path, monkeypatch): + config = sglang_config.CUDAGraphExtensionConfig(workspace_root=str(tmp_path)) + monkeypatch.setattr(sglang_config, "_config", config) + monkeypatch.setattr(sglang_runtime, "_state", None) + state = sglang_runtime.WarmupState() + state.deepep_transport = "nvl_ipc" + state.deepep_num_nvl_bytes = 64 * 1024 * 1024 + + sglang_runtime.save_warmup_state(state) + + contents = json.loads((tmp_path / "warmup_state.json").read_text()) + assert contents["deepep_transport"] == "nvl_ipc" + assert contents["deepep_num_nvl_bytes"] == 64 * 1024 * 1024 + loaded = sglang_runtime.load_warmup_state() + assert loaded.deepep_transport == "nvl_ipc" + assert loaded.deepep_num_nvl_bytes == 64 * 1024 * 1024 + + +def test_sglang_legacy_warmup_state_uses_fabric_defaults(tmp_path, monkeypatch): + config = sglang_config.CUDAGraphExtensionConfig(workspace_root=str(tmp_path)) + monkeypatch.setattr(sglang_config, "_config", config) + (tmp_path / "warmup_state.json").write_text('{"sglang_version": "legacy"}') + + loaded = sglang_runtime.load_warmup_state() + + assert loaded.deepep_transport == "fabric" + assert loaded.deepep_num_nvl_bytes == 0 + + +def test_sglang_profile_update_preserves_existing_rank_guard(tmp_path, monkeypatch): + config = sglang_config.CUDAGraphExtensionConfig(workspace_root=str(tmp_path)) + monkeypatch.setattr(sglang_config, "_config", config) + sglang_runtime.save_warmup_state(sglang_runtime.WarmupState()) + monkeypatch.setattr( + sglang_runtime, + "_state", + sglang_runtime.CUDAGraphExtensionState(rank=1), + ) + + sglang_runtime.update_warmup_state_deepep_profile( + DeepEPTransport.NVL_IPC, + 64 * 1024 * 1024, + ) + + rank_one_state = sglang_runtime.load_warmup_state() + assert rank_one_state.deepep_transport == "fabric" + assert rank_one_state.deepep_num_nvl_bytes == 0 + + sglang_runtime._state.rank = 0 + sglang_runtime.update_warmup_state_deepep_profile( + DeepEPTransport.NVL_IPC, + 64 * 1024 * 1024, + ) + rank_zero_state = sglang_runtime.load_warmup_state() + assert rank_zero_state.deepep_transport == "nvl_ipc" + assert rank_zero_state.deepep_num_nvl_bytes == 64 * 1024 * 1024 + + +def test_sglang_load_rejects_transport_mismatch_before_cuda_operations( + tmp_path, + monkeypatch, +): + (tmp_path / "rank_0").mkdir() + (tmp_path / "warmup_state.json").write_text( + json.dumps( + { + "deepep_transport": "nvl_ipc", + "deepep_num_nvl_bytes": 64 * 1024 * 1024, + } + ) + ) + config = sglang_config.CUDAGraphExtensionConfig( + mode=sglang_config.CUDAGraphExtensionMode.LOAD, + deepep_transport=DeepEPTransport.FABRIC, + workspace_root=str(tmp_path), + ) + monkeypatch.setattr(sglang_config, "_config", config) + monkeypatch.setattr(sglang_runtime, "_state", None) + cuda_operations = [] + monkeypatch.setattr(sglang_runtime.cge, "set_skip_fatbin_processing", lambda value: None) + monkeypatch.setattr( + sglang_runtime.cge, + "load_cuda_modules_and_libraries", + lambda path: cuda_operations.append("load_cuda_modules_and_libraries"), + ) + monkeypatch.setattr( + sglang_runtime.cge, + "set_allocation_region", + lambda base_addr, size: cuda_operations.append("set_allocation_region"), + ) + monkeypatch.setattr( + sglang_runtime.torch._C, + "_cuda_getCurrentBlasHandle", + lambda: None, + ) + server_args = SimpleNamespace( + enable_dp_attention=False, + tp_size=1, + pp_size=1, + ) + + with pytest.raises(RuntimeError, match="archived=nvl_ipc.*configured=fabric"): + sglang_runtime.setup_graph_extension( + server_args, + tp_rank=0, + pp_rank=0, + dp_rank=None, + ) + + assert cuda_operations == [] + + +@pytest.mark.parametrize( + ("transport", "expected"), + [ + ( + DeepEPTransport.FABRIC, + {"num_nvl_bytes": 0, "use_fabric": True}, + ), + ( + DeepEPTransport.NVL_IPC, + {"num_nvl_bytes": 64 * 1024 * 1024, "use_fabric": False}, + ), + ], +) +def test_sglang_hook_applies_transport_profile_on_save( + monkeypatch, + transport, + expected, +): + constructor_calls = [] + + class Buffer: + def __init__( + self, + group, + num_nvl_bytes=0, + num_rdma_bytes=0, + *, + use_fabric=None, + ): + constructor_calls.append( + { + "num_nvl_bytes": num_nvl_bytes, + "use_fabric": use_fabric, + } + ) + + _install_fake_deep_ep(monkeypatch, Buffer) + updates = [] + monkeypatch.setattr( + sglang_hooks, + "get_graph_extension_mode", + lambda: sglang_hooks.CUDAGraphExtensionMode.SAVE, + ) + monkeypatch.setattr(sglang_hooks, "get_deepep_transport", lambda: transport) + monkeypatch.setattr( + sglang_hooks.rt, + "update_warmup_state_deepep_profile", + lambda selected, num_nvl_bytes: updates.append((selected, num_nvl_bytes)), + ) + + sglang_hooks._patch_deepep() + Buffer("group", 64 * 1024 * 1024, 32 * 1024 * 1024) + + assert constructor_calls == [expected] + assert updates == [(transport, expected["num_nvl_bytes"])] + + +def test_sglang_hook_leaves_none_mode_profile_unchanged(monkeypatch): + constructor_calls = [] + + class Buffer: + def __init__( + self, + group, + num_nvl_bytes=0, + num_rdma_bytes=0, + *, + use_fabric=None, + ): + constructor_calls.append( + { + "num_nvl_bytes": num_nvl_bytes, + "num_rdma_bytes": num_rdma_bytes, + "use_fabric": use_fabric, + } + ) + + _install_fake_deep_ep(monkeypatch, Buffer) + monkeypatch.setattr( + sglang_hooks, + "get_graph_extension_mode", + lambda: sglang_hooks.CUDAGraphExtensionMode.NONE, + ) + monkeypatch.setattr( + sglang_hooks.rt, + "update_warmup_state_deepep_profile", + lambda *args: pytest.fail("NONE mode must not update the archive"), + ) + monkeypatch.setattr( + sglang_hooks.rt, + "load_warmup_state", + lambda: pytest.fail("NONE mode must not load the archive"), + ) + + sglang_hooks._patch_deepep() + Buffer( + "group", + 64 * 1024 * 1024, + 32 * 1024 * 1024, + use_fabric=True, + ) + + assert constructor_calls == [ + { + "num_nvl_bytes": 64 * 1024 * 1024, + "num_rdma_bytes": 32 * 1024 * 1024, + "use_fabric": True, + } + ] + + +def test_sglang_load_rejects_buffer_size_before_original_constructor(monkeypatch): + constructor_calls = [] + + class Buffer: + def __init__( + self, + group, + num_nvl_bytes=0, + num_rdma_bytes=0, + *, + use_fabric=None, + ): + constructor_calls.append((group, num_nvl_bytes, num_rdma_bytes, use_fabric)) + + _install_fake_deep_ep(monkeypatch, Buffer) + monkeypatch.setattr( + sglang_hooks, + "get_graph_extension_mode", + lambda: sglang_hooks.CUDAGraphExtensionMode.LOAD, + ) + monkeypatch.setattr( + sglang_hooks, + "get_deepep_transport", + lambda: DeepEPTransport.NVL_IPC, + ) + monkeypatch.setattr( + sglang_hooks.rt, + "load_warmup_state", + lambda: SimpleNamespace(deepep_num_nvl_bytes=64 * 1024 * 1024), + ) + + sglang_hooks._patch_deepep() + + with pytest.raises(RuntimeError, match="archived=67108864.*effective=134217728"): + Buffer("group", 128 * 1024 * 1024, 32 * 1024 * 1024) + + assert constructor_calls == [] + + +def test_sglang_rejects_old_deepep_before_original_constructor(monkeypatch): + constructor_calls = [] + + class Buffer: + def __init__(self, group, num_nvl_bytes=0, num_rdma_bytes=0): + constructor_calls.append((group, num_nvl_bytes, num_rdma_bytes)) + + _install_fake_deep_ep(monkeypatch, Buffer) + monkeypatch.setattr( + sglang_hooks, + "get_graph_extension_mode", + lambda: sglang_hooks.CUDAGraphExtensionMode.SAVE, + ) + + sglang_hooks._patch_deepep() + + with pytest.raises( + RuntimeError, + match=( + "Foundry SGLang DeepEP requires a deep_ep.Buffer constructor " + "with the use_fabric parameter" + ), + ): + Buffer("group", 64 * 1024 * 1024, 32 * 1024 * 1024) + + assert constructor_calls == [] + + +def test_sglang_save_fails_before_buffer_if_shared_archive_is_missing( + tmp_path, + monkeypatch, +): + constructor_calls = [] + + class Buffer: + def __init__( + self, + group, + num_nvl_bytes=0, + num_rdma_bytes=0, + *, + use_fabric=None, + ): + constructor_calls.append((group, num_nvl_bytes, num_rdma_bytes, use_fabric)) + + _install_fake_deep_ep(monkeypatch, Buffer) + config = sglang_config.CUDAGraphExtensionConfig( + mode=sglang_config.CUDAGraphExtensionMode.SAVE, + deepep_transport=DeepEPTransport.NVL_IPC, + workspace_root=str(tmp_path), + ) + monkeypatch.setattr(sglang_config, "_config", config) + monkeypatch.setattr( + sglang_runtime, + "_state", + sglang_runtime.CUDAGraphExtensionState(rank=0), + ) + monkeypatch.setattr( + sglang_hooks, + "get_graph_extension_mode", + lambda: sglang_hooks.CUDAGraphExtensionMode.SAVE, + ) + monkeypatch.setattr( + sglang_hooks, + "get_deepep_transport", + lambda: DeepEPTransport.NVL_IPC, + ) + + sglang_hooks._patch_deepep() + + with pytest.raises(RuntimeError, match="warmup state file not found"): + Buffer("group", 64 * 1024 * 1024, 32 * 1024 * 1024) + + assert constructor_calls == [] diff --git a/tests/test_deepep_transport_vllm.py b/tests/test_deepep_transport_vllm.py new file mode 100644 index 00000000..967e315f --- /dev/null +++ b/tests/test_deepep_transport_vllm.py @@ -0,0 +1,665 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the Foundry project +from __future__ import annotations + +import importlib.util +import json +import logging +import os +import sys +import types +from pathlib import Path + +import pytest +from foundry.integration.deepep_transport import DeepEPTransport + +VLLM_CONFIG_PATH = ( + Path(__file__).parents[1] / "python" / "foundry" / "integration" / "vllm" / "config.py" +) +VLLM_RUNTIME_PATH = ( + Path(__file__).parents[1] / "python" / "foundry" / "integration" / "vllm" / "runtime.py" +) +VLLM_HOOKS_PATH = ( + Path(__file__).parents[1] / "python" / "foundry" / "integration" / "vllm" / "hooks.py" +) + + +def _installed_nvshmem_host_path() -> Path: + spec = importlib.util.find_spec("nvidia.nvshmem") + assert spec is not None, "nvidia.nvshmem must be installed for this integration test" + roots = [Path(root) for root in spec.submodule_search_locations or ()] + if spec.origin: + roots.append(Path(spec.origin).parent) + for root in roots: + for name in ("libnvshmem_host.so.3", "libnvshmem_host.so"): + candidate = root / "lib" / name + if candidate.exists(): + return candidate.resolve() + pytest.fail("installed nvidia.nvshmem package has no libnvshmem_host shared library") + + +def _init_stub_logger(name: str) -> logging.Logger: + return logging.getLogger(name) + + +def _stub_vllm_if_missing(monkeypatch) -> None: + if importlib.util.find_spec("vllm") is not None: + return + vllm_module = types.ModuleType("vllm") + vllm_module.__path__ = [] + vllm_module.__version__ = "test" + logger_module = types.ModuleType("vllm.logger") + logger_module.init_logger = _init_stub_logger + monkeypatch.setitem(sys.modules, "vllm", vllm_module) + monkeypatch.setitem(sys.modules, "vllm.logger", logger_module) + + +@pytest.fixture +def vllm_config_module(monkeypatch): + _stub_vllm_if_missing(monkeypatch) + + spec = importlib.util.spec_from_file_location( + "foundry_vllm_config_under_test", + VLLM_CONFIG_PATH, + ) + assert spec is not None + assert spec.loader is not None + module = importlib.util.module_from_spec(spec) + monkeypatch.setitem(sys.modules, spec.name, module) + spec.loader.exec_module(module) + return module + + +@pytest.fixture +def vllm_runtime_module(monkeypatch): + _stub_vllm_if_missing(monkeypatch) + + spec = importlib.util.spec_from_file_location( + "foundry_vllm_runtime_under_test", + VLLM_RUNTIME_PATH, + ) + assert spec is not None + assert spec.loader is not None + module = importlib.util.module_from_spec(spec) + monkeypatch.setitem(sys.modules, spec.name, module) + spec.loader.exec_module(module) + return module + + +@pytest.fixture +def vllm_hooks_module(monkeypatch): + _stub_vllm_if_missing(monkeypatch) + + spec = importlib.util.spec_from_file_location( + "foundry_vllm_hooks_under_test", + VLLM_HOOKS_PATH, + ) + assert spec is not None + assert spec.loader is not None + module = importlib.util.module_from_spec(spec) + monkeypatch.setitem(sys.modules, spec.name, module) + spec.loader.exec_module(module) + return module + + +def _install_fake_all2all_module(monkeypatch, manager_cls): + distributed_module = types.ModuleType("vllm.distributed") + distributed_module.__path__ = [] + communicators_module = types.ModuleType("vllm.distributed.device_communicators") + communicators_module.__path__ = [] + all2all_module = types.ModuleType("vllm.distributed.device_communicators.all2all") + all2all_module.DeepEPLLAll2AllManager = manager_cls + communicators_module.all2all = all2all_module + distributed_module.device_communicators = communicators_module + monkeypatch.setattr( + sys.modules["vllm"], + "distributed", + distributed_module, + raising=False, + ) + monkeypatch.setitem(sys.modules, "vllm.distributed", distributed_module) + monkeypatch.setitem( + sys.modules, + "vllm.distributed.device_communicators", + communicators_module, + ) + monkeypatch.setitem( + sys.modules, + "vllm.distributed.device_communicators.all2all", + all2all_module, + ) + + +def _install_fake_kv_cache_modules(monkeypatch, engine_core_cls, scheduler_config): + v1_module = types.ModuleType("vllm.v1") + v1_module.__path__ = [] + engine_module = types.ModuleType("vllm.v1.engine") + engine_module.__path__ = [] + core_module = types.ModuleType("vllm.v1.engine.core") + core_module.EngineCore = engine_core_cls + core_utils_module = types.ModuleType("vllm.v1.core.kv_cache_utils") + core_utils_module.get_kv_cache_configs = lambda vllm_config, specs, available_memory: [ + types.SimpleNamespace( + specs=specs, + available_memory=available_memory, + ) + ] + core_utils_module.generate_scheduler_kv_cache_config = lambda configs: scheduler_config + envs_module = types.ModuleType("vllm.envs") + envs_module.VLLM_ELASTIC_EP_SCALE_UP_LAUNCH = False + engine_module.core = core_module + v1_module.engine = engine_module + monkeypatch.setattr(sys.modules["vllm"], "v1", v1_module, raising=False) + monkeypatch.setattr(sys.modules["vllm"], "envs", envs_module, raising=False) + monkeypatch.setitem(sys.modules, "vllm.v1", v1_module) + monkeypatch.setitem(sys.modules, "vllm.v1.engine", engine_module) + monkeypatch.setitem(sys.modules, "vllm.v1.engine.core", core_module) + monkeypatch.setitem(sys.modules, "vllm.v1.core.kv_cache_utils", core_utils_module) + monkeypatch.setitem(sys.modules, "vllm.envs", envs_module) + + +@pytest.mark.parametrize( + ("contents", "expected"), + [ + ("", DeepEPTransport.FABRIC), + ('deepep_transport = "nvl_ipc"\n', DeepEPTransport.NVL_IPC), + ], +) +def test_vllm_config_parses_deepep_transport( + tmp_path, + vllm_config_module, + contents, + expected, +): + config_path = tmp_path / "config.toml" + config_path.write_text(contents) + + config = vllm_config_module.CUDAGraphExtensionConfig.from_toml(config_path) + + assert config.deepep_transport is expected + + +def test_vllm_config_rejects_unknown_deepep_transport(tmp_path, vllm_config_module): + config_path = tmp_path / "config.toml" + config_path.write_text('deepep_transport = "rdma"\n') + + with pytest.raises(ValueError, match="deepep_transport"): + vllm_config_module.CUDAGraphExtensionConfig.from_toml(config_path) + + +def test_vllm_get_deepep_transport_uses_loaded_config( + tmp_path, + vllm_config_module, +): + config_path = tmp_path / "config.toml" + config_path.write_text('deepep_transport = "nvl_ipc"\n') + + assert vllm_config_module.get_deepep_transport() is DeepEPTransport.FABRIC + + vllm_config_module.load_graph_extension_config(str(config_path)) + + assert vllm_config_module.get_deepep_transport() is DeepEPTransport.NVL_IPC + + +def test_vllm_config_autodetects_installed_nvshmem_host_path( + tmp_path, + vllm_config_module, +): + config_path = tmp_path / "config.toml" + config_path.write_text("") + expected = _installed_nvshmem_host_path() + + config = vllm_config_module.CUDAGraphExtensionConfig.from_toml(config_path) + + assert config.nvshmem_host_path == str(expected) + assert Path(config.nvshmem_host_path).is_absolute() + + +def test_vllm_runtime_preloads_autodetected_nvshmem_host_path( + tmp_path, + monkeypatch, + vllm_runtime_module, +): + config_path = tmp_path / "config.toml" + config_path.write_text("") + expected = _installed_nvshmem_host_path() + config_module = sys.modules[vllm_runtime_module.get_nvshmem_host_path.__module__] + config = config_module.CUDAGraphExtensionConfig.from_toml(config_path) + monkeypatch.setattr(config_module, "_config", config) + monkeypatch.setattr(vllm_runtime_module, "get_hook_library_path", lambda: None) + monkeypatch.delenv("LD_PRELOAD", raising=False) + + vllm_runtime_module.setup_ld_preload_env() + + preload_paths = [Path(path) for path in os.environ["LD_PRELOAD"].split(":")] + assert expected in preload_paths + assert all(path.is_absolute() for path in preload_paths) + + +def test_vllm_config_explicit_nvshmem_host_path_overrides_detection( + tmp_path, + monkeypatch, + vllm_config_module, +): + explicit = tmp_path / "custom" / "libnvshmem_host.so" + config_path = tmp_path / "config.toml" + config_path.write_text(f'nvshmem_host_path = "{explicit}"\n') + monkeypatch.setattr( + vllm_config_module.CUDAGraphExtensionConfig, + "_detect_nvshmem_host_path", + staticmethod(lambda: pytest.fail("explicit path must bypass auto-detection")), + raising=False, + ) + + config = vllm_config_module.CUDAGraphExtensionConfig.from_toml(config_path) + + assert config.nvshmem_host_path == str(explicit) + + +def test_vllm_config_missing_nvshmem_package_is_none_for_non_ep( + tmp_path, + monkeypatch, + vllm_config_module, +): + config_path = tmp_path / "config.toml" + config_path.write_text("") + real_find_spec = vllm_config_module.importlib.util.find_spec + + def find_spec_without_nvshmem(name): + if name == "nvidia.nvshmem": + return None + return real_find_spec(name) + + monkeypatch.setattr(vllm_config_module.importlib.util, "find_spec", find_spec_without_nvshmem) + + config = vllm_config_module.CUDAGraphExtensionConfig.from_toml(config_path) + + assert config.nvshmem_host_path is None + assert config.deepep_transport is DeepEPTransport.FABRIC + + +def test_vllm_warmup_state_round_trips_deepep_profile(tmp_path, vllm_runtime_module): + state = vllm_runtime_module.WarmupState() + state.deepep_transport = "nvl_ipc" + state.deepep_num_nvl_bytes = 64 * 1024 * 1024 + + vllm_runtime_module.save_warmup_state(str(tmp_path), state) + + contents = json.loads((tmp_path / "warmup_state.json").read_text()) + assert contents["deepep_transport"] == "nvl_ipc" + assert contents["deepep_num_nvl_bytes"] == 64 * 1024 * 1024 + loaded = vllm_runtime_module.load_warmup_state(str(tmp_path)) + assert loaded.deepep_transport == "nvl_ipc" + assert loaded.deepep_num_nvl_bytes == 64 * 1024 * 1024 + + +def test_vllm_legacy_warmup_state_uses_fabric_defaults(tmp_path, vllm_runtime_module): + (tmp_path / "warmup_state.json").write_text('{"vllm_version": "legacy"}') + + loaded = vllm_runtime_module.load_warmup_state(str(tmp_path)) + + assert loaded.deepep_transport == "fabric" + assert loaded.deepep_num_nvl_bytes == 0 + + +def test_vllm_profile_update_is_rank_zero_only( + tmp_path, + monkeypatch, + vllm_runtime_module, +): + vllm_runtime_module.save_warmup_state( + str(tmp_path), + vllm_runtime_module.WarmupState(), + ) + config = types.SimpleNamespace(workspace_dir=str(tmp_path / "rank_1")) + monkeypatch.setattr(vllm_runtime_module, "get_config", lambda: config) + + vllm_runtime_module.update_warmup_state_deepep_profile( + str(tmp_path), + DeepEPTransport.NVL_IPC, + 64 * 1024 * 1024, + ) + + rank_one_state = vllm_runtime_module.load_warmup_state(str(tmp_path)) + assert rank_one_state.deepep_transport == "fabric" + assert rank_one_state.deepep_num_nvl_bytes == 0 + + config.workspace_dir = str(tmp_path / "rank_0") + vllm_runtime_module.update_warmup_state_deepep_profile( + str(tmp_path), + DeepEPTransport.NVL_IPC, + 64 * 1024 * 1024, + ) + rank_zero_state = vllm_runtime_module.load_warmup_state(str(tmp_path)) + assert rank_zero_state.deepep_transport == "nvl_ipc" + assert rank_zero_state.deepep_num_nvl_bytes == 64 * 1024 * 1024 + + +def test_vllm_load_rejects_transport_mismatch_before_cuda_operations( + tmp_path, + monkeypatch, + vllm_runtime_module, +): + (tmp_path / "rank_0").mkdir() + (tmp_path / "warmup_state.json").write_text( + json.dumps( + { + "deepep_transport": "nvl_ipc", + "deepep_num_nvl_bytes": 64 * 1024 * 1024, + } + ) + ) + config = types.SimpleNamespace( + mode=vllm_runtime_module.CUDAGraphExtensionMode.LOAD, + deepep_transport=DeepEPTransport.FABRIC, + workspace_root=str(tmp_path), + workspace_dir=None, + base_addr=0x600000000000, + region_size="1GB", + ) + monkeypatch.setattr(vllm_runtime_module, "get_config", lambda: config) + monkeypatch.setattr(vllm_runtime_module, "_state", None) + cuda_operations = [] + monkeypatch.setattr( + vllm_runtime_module.fops, + "set_skip_fatbin_processing", + lambda value: None, + ) + monkeypatch.setattr( + vllm_runtime_module.fops, + "load_cuda_modules_and_libraries", + lambda path: cuda_operations.append("load_cuda_modules_and_libraries"), + ) + monkeypatch.setattr( + vllm_runtime_module.fops, + "set_allocation_region", + lambda base_addr, size: cuda_operations.append("set_allocation_region"), + ) + parallel_config = types.SimpleNamespace( + rank=0, + world_size=1, + data_parallel_index=0, + tensor_parallel_size=1, + pipeline_parallel_size=1, + data_parallel_size=1, + nnodes=1, + ) + + with pytest.raises(RuntimeError, match="archived=nvl_ipc.*configured=fabric"): + vllm_runtime_module.setup_graph_extension(parallel_config) + + assert cuda_operations == [] + + +@pytest.mark.parametrize( + ("transport", "expected"), + [ + ( + DeepEPTransport.FABRIC, + {"num_nvl_bytes": 0, "use_fabric": True}, + ), + ( + DeepEPTransport.NVL_IPC, + {"num_nvl_bytes": 64 * 1024 * 1024, "use_fabric": False}, + ), + ], +) +def test_vllm_hook_applies_transport_profile_on_save( + tmp_path, + monkeypatch, + vllm_hooks_module, + transport, + expected, +): + class Manager: + def _make_all2all_kwargs(self): + return {"num_nvl_bytes": 64 * 1024 * 1024} + + _install_fake_all2all_module(monkeypatch, Manager) + updates = [] + monkeypatch.setattr( + vllm_hooks_module, + "get_graph_extension_mode", + lambda: vllm_hooks_module.CUDAGraphExtensionMode.SAVE, + ) + monkeypatch.setattr(vllm_hooks_module, "get_deepep_transport", lambda: transport) + monkeypatch.setattr(vllm_hooks_module, "get_workspace_root", lambda: str(tmp_path)) + monkeypatch.setattr( + vllm_hooks_module.rt, + "update_warmup_state_deepep_profile", + lambda root, selected, num_nvl_bytes: updates.append((root, selected, num_nvl_bytes)), + ) + + vllm_hooks_module._patch_deepep() + result = Manager()._make_all2all_kwargs() + + assert result == expected + assert updates == [(str(tmp_path), transport, expected["num_nvl_bytes"])] + + +def test_vllm_hook_leaves_none_mode_profile_unchanged(monkeypatch, vllm_hooks_module): + upstream = { + "num_nvl_bytes": 64 * 1024 * 1024, + "use_fabric": "upstream", + } + + class Manager: + def _make_all2all_kwargs(self): + return dict(upstream) + + _install_fake_all2all_module(monkeypatch, Manager) + monkeypatch.setattr( + vllm_hooks_module, + "get_graph_extension_mode", + lambda: vllm_hooks_module.CUDAGraphExtensionMode.NONE, + ) + monkeypatch.setattr( + vllm_hooks_module.rt, + "update_warmup_state_deepep_profile", + lambda *args: pytest.fail("NONE mode must not update the archive"), + ) + monkeypatch.setattr( + vllm_hooks_module.rt, + "load_warmup_state", + lambda *args: pytest.fail("NONE mode must not load the archive"), + ) + + vllm_hooks_module._patch_deepep() + + assert Manager()._make_all2all_kwargs() == upstream + + +def test_vllm_load_rejects_buffer_size_before_communicator_setup( + tmp_path, + monkeypatch, + vllm_hooks_module, +): + class Manager: + def _make_all2all_kwargs(self): + return {"num_nvl_bytes": 128 * 1024 * 1024} + + _install_fake_all2all_module(monkeypatch, Manager) + communicator_calls = [] + monkeypatch.setattr( + vllm_hooks_module, + "get_graph_extension_mode", + lambda: vllm_hooks_module.CUDAGraphExtensionMode.LOAD, + ) + monkeypatch.setattr( + vllm_hooks_module, + "get_deepep_transport", + lambda: DeepEPTransport.NVL_IPC, + ) + monkeypatch.setattr(vllm_hooks_module, "get_workspace_root", lambda: str(tmp_path)) + monkeypatch.setattr( + vllm_hooks_module.rt, + "load_warmup_state", + lambda root: types.SimpleNamespace(deepep_num_nvl_bytes=64 * 1024 * 1024), + ) + vllm_hooks_module._patch_deepep() + + def construct_communicator(): + kwargs = Manager()._make_all2all_kwargs() + communicator_calls.append(kwargs) + + with pytest.raises(RuntimeError, match="archived=67108864.*effective=134217728"): + construct_communicator() + + assert communicator_calls == [] + + +def test_vllm_install_does_not_write_shared_archive( + tmp_path, + monkeypatch, + vllm_hooks_module, +): + monkeypatch.setattr(vllm_hooks_module, "_INSTALLED", False) + monkeypatch.delenv("FOUNDRY_SPAWN_T0_NS", raising=False) + monkeypatch.setattr(vllm_hooks_module, "_quarantine_get_device_capability", lambda: None) + monkeypatch.setattr(vllm_hooks_module, "load_graph_extension_config", lambda path: None) + monkeypatch.setattr( + vllm_hooks_module, + "get_graph_extension_mode", + lambda: vllm_hooks_module.CUDAGraphExtensionMode.SAVE, + ) + monkeypatch.setattr(vllm_hooks_module, "get_workspace_root", lambda: str(tmp_path)) + patch_names = ( + "_patch_init_worker_distributed_environment", + "_patch_kernel_warmup", + "_patch_dummy_runs", + "_patch_load_model", + "_patch_prepare_comm_buffer", + "_patch_capture_model", + "_patch_cuda_graph_wrapper_call", + "_patch_initialize_kv_caches", + "_patch_subprocess_spawn_sites", + ) + for patch_name in patch_names: + monkeypatch.setattr(vllm_hooks_module, patch_name, lambda: None) + + deepep_patch_calls = [] + monkeypatch.setattr( + vllm_hooks_module, + "_patch_deepep", + lambda: deepep_patch_calls.append("patched"), + ) + compilation_config = types.SimpleNamespace( + graph_extension_config_path=str(tmp_path / "config.toml"), + cudagraph_num_of_warmups=1, + ) + + vllm_hooks_module.install_hooks(compilation_config) + + assert deepep_patch_calls == ["patched"] + assert not (tmp_path / "warmup_state.json").exists() + + +def test_vllm_bootstrap_archive_does_not_skip_first_save_profile( + tmp_path, + monkeypatch, + vllm_hooks_module, +): + profile_calls = [] + initialized_configs = [] + + class ModelExecutor: + def get_kv_cache_specs(self): + return [object()] + + def determine_available_memory(self): + profile_calls.append("determine_available_memory") + return [123456] + + def initialize_from_config(self, configs): + initialized_configs.append(configs) + + class EngineCore: + def __init__(self): + self.model_executor = ModelExecutor() + self.available_gpu_memory_for_kv_cache = 0 + + def _initialize_kv_caches(self, vllm_config): + pytest.fail("SAVE mode must use Foundry's KV-cache initialization") + + scheduler_config = types.SimpleNamespace(num_blocks=7, kv_cache_groups=[]) + _install_fake_kv_cache_modules( + monkeypatch, + EngineCore, + scheduler_config, + ) + bootstrap = vllm_hooks_module.rt.WarmupState( + deepep_transport="nvl_ipc", + deepep_num_nvl_bytes=64 * 1024 * 1024, + ) + vllm_hooks_module.rt.save_warmup_state(str(tmp_path), bootstrap) + monkeypatch.setattr( + vllm_hooks_module, + "get_graph_extension_mode", + lambda: vllm_hooks_module.CUDAGraphExtensionMode.SAVE, + ) + monkeypatch.setattr(vllm_hooks_module, "get_workspace_root", lambda: str(tmp_path)) + runtime_metadata = { + "vllm_version": "test-vllm", + "timestamp": "2026-07-23T10:00:00", + "cuda_version": "13.0", + "gpu_name": "H100", + "gpu_total_memory": 80 * 1024**3, + } + monkeypatch.setattr( + vllm_hooks_module.rt, + "create_warmup_state", + lambda: vllm_hooks_module.rt.WarmupState(**runtime_metadata), + ) + monkeypatch.setattr(vllm_hooks_module.rt, "get_final_alloc_offset", lambda: 0) + vllm_hooks_module._patch_initialize_kv_caches() + vllm_config = types.SimpleNamespace( + model_config=types.SimpleNamespace(max_model_len=1024), + parallel_config=types.SimpleNamespace(data_parallel_index=0), + cache_config=types.SimpleNamespace( + num_gpu_blocks=0, + block_size=0, + gpu_memory_utilization=0.9, + ), + validate_block_size=lambda: None, + ) + + result = EngineCore()._initialize_kv_caches(vllm_config) + + assert result is scheduler_config + assert profile_calls == ["determine_available_memory"] + assert len(initialized_configs) == 1 + saved = vllm_hooks_module.rt.load_warmup_state(str(tmp_path)) + assert saved.available_gpu_memory == [123456] + assert saved.num_gpu_blocks == 7 + for field, expected in runtime_metadata.items(): + assert getattr(saved, field) == expected + assert saved.deepep_transport == "nvl_ipc" + assert saved.deepep_num_nvl_bytes == 64 * 1024 * 1024 + + +def test_vllm_profile_update_creates_missing_archive_only_on_rank_zero( + tmp_path, + monkeypatch, + vllm_runtime_module, +): + config = types.SimpleNamespace(workspace_dir=str(tmp_path / "rank_1")) + monkeypatch.setattr(vllm_runtime_module, "get_config", lambda: config) + + vllm_runtime_module.update_warmup_state_deepep_profile( + str(tmp_path), + DeepEPTransport.NVL_IPC, + 64 * 1024 * 1024, + ) + + assert not (tmp_path / "warmup_state.json").exists() + + config.workspace_dir = str(tmp_path / "rank_0") + vllm_runtime_module.update_warmup_state_deepep_profile( + str(tmp_path), + DeepEPTransport.NVL_IPC, + 64 * 1024 * 1024, + ) + + saved = vllm_runtime_module.load_warmup_state(str(tmp_path)) + assert saved.deepep_transport == "nvl_ipc" + assert saved.deepep_num_nvl_bytes == 64 * 1024 * 1024 diff --git a/tests/test_experimental_recipe.py b/tests/test_experimental_recipe.py new file mode 100644 index 00000000..b65e49ec --- /dev/null +++ b/tests/test_experimental_recipe.py @@ -0,0 +1,145 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the Foundry project +"""Static contracts for the experimental no-IMEX DeepEP recipes.""" + +import subprocess +from pathlib import Path + +import tomllib + +ROOT = Path(__file__).resolve().parents[1] +EXPERIMENTAL_DIR = ROOT / "recipe" / "experimental" +CONFIG_PAIRS = { + "vLLM": ("foundry_save.toml", "foundry_load.toml"), + "SGLang": ("sglang_foundry_save.toml", "sglang_foundry_load.toml"), +} +SERVE_SCRIPTS = ( + "serve_qwen3-30ba3b_ipc_ep.sh", + "serve_qwen3-30ba3bfp8_sglang_ipc_ep.sh", +) +INTEGRATION_DOCS = ( + ROOT / "docs" / "vllm" / "moe-and-deepep.md", + ROOT / "docs" / "sglang" / "hooks.md", +) +TWO_GPU_COMMANDS = ( + "bash tests/run_deepep_matrix.sh ll_nofabric both", + "bash tests/run_deepep_matrix.sh nvl_ipc both", + "bash tests/run_deepep_matrix.sh nvl_ipc_prealloc both", +) + + +def _load_config(filename: str) -> dict[str, object]: + path = EXPERIMENTAL_DIR / filename + assert path.is_file(), f"missing no-IMEX recipe config: {path}" + return tomllib.loads(path.read_text()) + + +def test_no_imex_configs_select_nvl_ipc_transport() -> None: + for config_pair in CONFIG_PAIRS.values(): + for filename in config_pair: + config = _load_config(filename) + assert config.get("deepep_transport") == "nvl_ipc", filename + + +def test_save_and_load_configs_share_archive_layout() -> None: + for engine, (save_filename, load_filename) in CONFIG_PAIRS.items(): + save_config = _load_config(save_filename) + load_config = _load_config(load_filename) + + for key in ("workspace_root", "base_addr", "region_size"): + assert save_config.get(key) == load_config.get(key), f"{engine}: {key}" + + +def test_experimental_scripts_do_not_use_environment_transport_switch() -> None: + for filename in SERVE_SCRIPTS: + path = EXPERIMENTAL_DIR / filename + assert path.is_file(), f"missing no-IMEX serve script: {path}" + assert "FOUNDRY_DEEPEP_NVL_IPC" not in path.read_text(), filename + + +def test_experimental_scripts_are_strict_and_syntax_valid() -> None: + for filename in SERVE_SCRIPTS: + path = EXPERIMENTAL_DIR / filename + text = path.read_text() + + assert text.splitlines()[1] == "set -euo pipefail", filename + subprocess.run(["bash", "-n", str(path)], check=True) + + vllm_script = (EXPERIMENTAL_DIR / SERVE_SCRIPTS[0]).read_text() + assert vllm_script.count("export VLLM_USE_V2_MODEL_RUNNER=0") == 1 + + +def test_experimental_scripts_pin_validated_deterministic_defaults() -> None: + vllm_script = (EXPERIMENTAL_DIR / SERVE_SCRIPTS[0]).read_text() + sglang_script = (EXPERIMENTAL_DIR / SERVE_SCRIPTS[1]).read_text() + + assert "--seed 42" in vllm_script + assert "MEM_FRACTION_STATIC=0.65" in sglang_script + assert "--random-seed 42" in sglang_script + assert "0.65" in (EXPERIMENTAL_DIR / "README.md").read_text() + + +def test_experimental_configs_have_portable_paths() -> None: + for config_pair in CONFIG_PAIRS.values(): + for filename in config_pair: + path = EXPERIMENTAL_DIR / filename + text = path.read_text() + config = tomllib.loads(text) + + assert "nvshmem_host_path" not in config, filename + assert "/data/" not in text, filename + assert not Path(str(config["workspace_root"])).is_absolute(), filename + + +def test_integration_docs_define_the_deepep_transport_contract() -> None: + for path in INTEGRATION_DOCS: + text = path.read_text() + + for term in ("deepep_transport", "fabric", "nvl_ipc", "intranode", "DeepEP", "NVSHMEM"): + assert term in text, f"{path}: missing {term}" + assert "default" in text.lower(), f"{path}: missing default transport" + assert "silent fallback" in text.lower(), f"{path}: missing no-fallback guarantee" + assert "Foundry DeepEP transport mismatch: archived=nvl_ipc, configured=fabric" in text, ( + f"{path}: missing archive parity error" + ) + for command in TWO_GPU_COMMANDS: + assert command in text, f"{path}: missing {command}" + + +def test_sglang_docs_define_required_deepep_api() -> None: + text = (ROOT / "docs" / "sglang" / "hooks.md").read_text() + + assert "Buffer(use_fabric=...)" in text + assert "29d31c0" in text + + +def test_sglang_hooks_doc_lists_attn_tp_logits_gather_patch() -> None: + text = (ROOT / "docs" / "sglang" / "hooks.md").read_text() + + # The LOAD logits pre-copy patch must be listed in the install order and + # described, not omitted from the hook-install documentation. + assert "_patch_attn_tp_logits_gather" in text + assert text.count("_patch_attn_tp_logits_gather") >= 2 + assert "pre-cop" in text.lower() + + +def test_sglang_script_accepts_documented_deepep_mode_override() -> None: + script = (EXPERIMENTAL_DIR / SERVE_SCRIPTS[1]).read_text() + + # The DeepEP dispatch mode is an override with the product default preserved. + assert 'DEEPEP_MODE="${DEEPEP_MODE:-low_latency}"' in script + # The mode is threaded through to sglang rather than hardcoded. + assert '--deepep-mode "$DEEPEP_MODE"' in script + assert "--deepep-mode low_latency" not in script + # The override is documented in the script itself. + assert "DEEPEP_MODE" in script.split("sglang serve")[0] + + +def test_sglang_readme_documents_deepep_mode_override() -> None: + readme = (EXPERIMENTAL_DIR / "README.md").read_text() + + assert "DEEPEP_MODE" in readme + # Both the graph-capturable nonzero-NVL mode and the normal-mode limitation + # are documented so the override is used correctly. + assert "auto" in readme + assert "num_nvl_bytes" in readme diff --git a/tests/test_graph_edge_data.py b/tests/test_graph_edge_data.py new file mode 100644 index 00000000..e4bf8339 --- /dev/null +++ b/tests/test_graph_edge_data.py @@ -0,0 +1,346 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the Foundry project +from __future__ import annotations + +import copy +import importlib.util +import json +import os +import shutil +import struct +import subprocess +import sys +from pathlib import Path + +import foundry as fdry +import pytest +import torch +from foundry import ops +from torch.utils.cpp_extension import load + +BASE_ADDR = 0x4F0000000000 +REGION_SIZE = "1GB" +TENSOR_SIZE = 256 +TEST_ROOT_ENV = "FOUNDRY_EDGE_DATA_TEST_ROOT" + +REPO_ROOT = Path(__file__).resolve().parents[1] +CUDA_HELPER = Path(__file__).with_name("cuda_graph_edge_data_test.cu") +BINARY_GRAPH_IO = REPO_ROOT / "csrc" / "BinaryGraphIO.cpp" + +FILE_HEADER_SIZE = 64 +SECTION_ENTRY_SIZE = 24 +DEPENDENCY_SECTION = 2 +LEGACY_DEPENDENCY_SIZE = 8 +EDGE_DEPENDENCY_SIZE = 16 + + +def _test_root() -> Path: + return Path(os.environ.get(TEST_ROOT_ENV, "test_graph_edge_data")) + + +def _archive_dir() -> Path: + return _test_root() / "graphs" + + +def _hook_archive_dir() -> Path: + return _test_root() / "hook_archive" + + +def _load_cuda_helper(): + return load( + name="foundry_cuda_graph_edge_data_test", + sources=[str(CUDA_HELPER), str(BINARY_GRAPH_IO)], + extra_include_paths=[str(REPO_ROOT / "include")], + extra_cflags=["-std=c++17"], + extra_cuda_cflags=["-std=c++17"], + extra_ldflags=[ + "-L/usr/local/cuda/lib64/stubs", + "-lcuda", + "-lboost_json", + "-lboost_filesystem", + ], + verbose=False, + ) + + +def _hook_so_path() -> str: + spec = importlib.util.find_spec("foundry.ops") + if spec is None or spec.origin is None: + raise RuntimeError("foundry.ops is unavailable") + path = Path(spec.origin).resolve().parent / "libcuda_hook.so" + if not path.is_file(): + raise RuntimeError(f"Foundry hook is unavailable: {path}") + return str(path) + + +def _dependency_section(data: bytes | bytearray) -> tuple[int, int, int]: + (num_dependencies,) = struct.unpack_from(" None: + data = bytearray(source.read_bytes()) + num_dependencies, section_offset, section_size = _dependency_section(data) + if num_dependencies == 0: + raise AssertionError("test graph unexpectedly has no dependencies") + entry_size = section_size // num_dependencies + if section_size != entry_size * num_dependencies: + raise AssertionError("source dependency section is malformed") + + if entry_size == LEGACY_DEPENDENCY_SIZE: + destination.write_bytes(data) + return + if entry_size != EDGE_DEPENDENCY_SIZE: + raise AssertionError(f"unexpected dependency entry size: {entry_size}") + + legacy_records = b"".join( + data[ + section_offset + index * EDGE_DEPENDENCY_SIZE : section_offset + + index * EDGE_DEPENDENCY_SIZE + + LEGACY_DEPENDENCY_SIZE + ] + for index in range(num_dependencies) + ) + old_end = section_offset + section_size + new_data = data[:section_offset] + legacy_records + data[old_end:] + removed = section_size - len(legacy_records) + + (num_sections,) = struct.unpack_from(" section_offset: + struct.pack_into(" None: + data = bytearray(source.read_bytes()) + num_dependencies, _section_offset, section_size = _dependency_section(data) + if num_dependencies == 0 or section_size < 2: + raise AssertionError("test graph unexpectedly has no dependency payload") + entry_offset = FILE_HEADER_SIZE + DEPENDENCY_SECTION * SECTION_ENTRY_SIZE + struct.pack_into(" tuple[int, int, int, int, int]: + return ( + value["from"], + value["to"], + value["from_port"], + value["to_port"], + value["type"], + ) + + +def _expected_values() -> torch.Tensor: + return torch.arange(1, TENSOR_SIZE + 1, dtype=torch.float32, device="cpu") * 2 + + +def _allocate_graph_tensors() -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + return tuple(torch.zeros(TENSOR_SIZE, device="cuda", dtype=torch.float32) for _ in range(4)) + + +def _capture_graph(helper, launch, values: torch.Tensor, output: torch.Tensor): + stream = torch.cuda.Stream() + graph = fdry.CUDAGraph(keep_graph=True) + with fdry.graph(graph, stream=stream): + launch(values, output) + torch.cuda.synchronize() + return graph + + +def _run_saving_process() -> None: + helper = _load_cuda_helper() + torch.cuda.init() + torch.cuda.set_device(0) + torch.set_default_device("cuda") + + archive_dir = _archive_dir() + archive_dir.mkdir(parents=True, exist_ok=True) + fdry.set_allocation_region(BASE_ADDR, fdry.parse_size(REGION_SIZE)) + programmatic_values, programmatic_output, default_values, default_output = ( + _allocate_graph_tensors() + ) + + programmatic_graph = _capture_graph( + helper, + helper.launch_programmatic, + programmatic_values, + programmatic_output, + ) + default_graph = _capture_graph( + helper, + helper.launch_default, + default_values, + default_output, + ) + + captured_programmatic_edges = [ + tuple(edge) for edge in helper.graph_edges(programmatic_graph.raw_cuda_graph()) + ] + captured_default_edges = [ + tuple(edge) for edge in helper.graph_edges(default_graph.raw_cuda_graph()) + ] + assert any(edge[2:] == (1, 0, 1) for edge in captured_programmatic_edges) + assert captured_default_edges + assert all(edge[2:] == (0, 0, 0) for edge in captured_default_edges) + + programmatic_path = archive_dir / "graph_0_FULL_t1_r1_UX_pcN.json" + default_path = archive_dir / "graph_1_FULL_t1_r1_UX_pcN.json" + programmatic_graph.save(str(programmatic_path), output_tensors=programmatic_output) + default_graph.save(str(default_path), output_tensors=default_output) + + programmatic_json = json.loads(programmatic_path.read_text()) + default_json = json.loads(default_path.read_text()) + serialized_edges = [_edge_tuple(dependency) for dependency in programmatic_json["dependencies"]] + assert sorted(serialized_edges) == sorted(captured_programmatic_edges) + assert [node["type"] for node in programmatic_json["nodes"]] == [ + node["type"] for node in default_json["nodes"] + ] + assert programmatic_json["topology_key"] != default_json["topology_key"] + + fdry.save_graph_manifest(str(archive_dir)) + manifest = json.loads((archive_dir / "graph_manifest.json").read_text()) + assert len(manifest["topology_groups"]) == 2 + + binary_path = programmatic_path.with_suffix(".cugraph") + num_dependencies, _section_offset, section_size = _dependency_section(binary_path.read_bytes()) + assert section_size == num_dependencies * EDGE_DEPENDENCY_SIZE + binary_json = json.loads(helper.read_binary_graph_json(str(binary_path))) + assert sorted(_edge_tuple(value) for value in binary_json["dependencies"]) == sorted( + captured_programmatic_edges + ) + + legacy_json = copy.deepcopy(programmatic_json) + for dependency in legacy_json["dependencies"]: + dependency.pop("from_port") + dependency.pop("to_port") + dependency.pop("type") + legacy_path = archive_dir / "legacy.json" + legacy_path.write_text(json.dumps(legacy_json)) + legacy_binary = legacy_path.with_suffix(".cugraph") + _write_legacy_binary(binary_path, legacy_binary) + legacy_binary_json = json.loads(helper.read_binary_graph_json(str(legacy_binary))) + assert sorted(_edge_tuple(value) for value in legacy_binary_json["dependencies"]) == sorted( + (edge[0], edge[1], 0, 0, 0) for edge in captured_programmatic_edges + ) + + malformed_binary = archive_dir / "malformed.cugraph" + _write_malformed_binary(binary_path, malformed_binary) + assert not helper.binary_graph_valid(str(malformed_binary)) + assert json.loads(helper.read_binary_graph_json(str(malformed_binary)) or "null") is None + + programmatic_graph.replay() + default_graph.replay() + torch.cuda.synchronize() + torch.testing.assert_close(programmatic_output.cpu(), _expected_values()) + torch.testing.assert_close(default_output.cpu(), _expected_values()) + fdry.stop_allocation_region() + + +def _load_graph_path(path: Path, *, parallel: bool): + if parallel: + pending = fdry.CUDAGraph.start_graph_builds([str(path)], num_threads=1) + graph, output = fdry.CUDAGraph.finish_one_graph_load(pending, 0) + return graph, output + return fdry.CUDAGraph.load(str(path)) + + +def _run_loading_process(*, legacy: bool, parallel: bool) -> None: + torch.cuda.init() + torch.cuda.set_device(0) + torch.set_default_device("cuda") + fdry.load_cuda_modules_and_libraries(str(_hook_archive_dir())) + helper = _load_cuda_helper() + + fdry.set_allocation_region(BASE_ADDR, fdry.parse_size(REGION_SIZE)) + _allocate_graph_tensors() + path = _archive_dir() / ("legacy.json" if legacy else "graph_0_FULL_t1_r1_UX_pcN.json") + if legacy and parallel: + isolated_dir = _archive_dir() / "legacy_parallel" + isolated_dir.mkdir(exist_ok=True) + path = isolated_dir / "graph_0_FULL_t1_r1_UX_pcN.json" + shutil.copyfile(_archive_dir() / "legacy.json", path) + shutil.copyfile(_archive_dir() / "legacy.cugraph", path.with_suffix(".cugraph")) + graph, output = _load_graph_path(path, parallel=parallel) + + if parallel: + graph_address = ops.get_graph_ptr(graph) + raw_graph = helper.shared_raw_cuda_graph(graph_address) + else: + raw_graph = graph.raw_cuda_graph() + loaded_edges = [tuple(edge) for edge in helper.graph_edges(raw_graph)] + + source_json = json.loads(path.read_text()) + expected_edges = [] + for dependency in source_json["dependencies"]: + expected_edges.append( + ( + dependency["from"], + dependency["to"], + dependency.get("from_port", 0), + dependency.get("to_port", 0), + dependency.get("type", 0), + ) + ) + assert sorted(loaded_edges) == sorted(expected_edges) + if legacy: + assert all(edge[2:] == (0, 0, 0) for edge in loaded_edges) + else: + assert any(edge[2:] == (1, 0, 1) for edge in loaded_edges) + + graph.replay() + torch.cuda.synchronize() + torch.testing.assert_close(output.cpu(), _expected_values()) + fdry.stop_allocation_region() + + +def _spawn(mode: str, root: Path) -> None: + environment = os.environ.copy() + environment[TEST_ROOT_ENV] = str(root) + hook = _hook_so_path() + environment["LD_PRELOAD"] = ( + f"{hook}:{environment['LD_PRELOAD']}" if environment.get("LD_PRELOAD") else hook + ) + subprocess.check_call( + [sys.executable, str(Path(__file__).resolve()), f"--{mode}"], + cwd=root, + env=environment, + ) + + +@pytest.mark.filterwarnings("ignore:TORCH_CUDA_ARCH_LIST is not set") +def test_cuda_graph_edge_data_round_trip(tmp_path): + root = tmp_path / "edge-data-round-trip" + root.mkdir() + try: + _spawn("saving-run", root) + _spawn("serial-load", root) + _spawn("serial-legacy-load", root) + _spawn("parallel-load", root) + _spawn("parallel-legacy-load", root) + finally: + shutil.rmtree(root, ignore_errors=True) + + +if __name__ == "__main__": + if "--saving-run" in sys.argv: + _run_saving_process() + elif "--serial-load" in sys.argv: + _run_loading_process(legacy=False, parallel=False) + elif "--serial-legacy-load" in sys.argv: + _run_loading_process(legacy=True, parallel=False) + elif "--parallel-load" in sys.argv: + _run_loading_process(legacy=False, parallel=True) + elif "--parallel-legacy-load" in sys.argv: + _run_loading_process(legacy=True, parallel=True) + else: + raise SystemExit("expected a test subprocess mode") diff --git a/tests/test_load_graph.py b/tests/test_load_graph.py index 0347f5b3..ba799388 100644 --- a/tests/test_load_graph.py +++ b/tests/test_load_graph.py @@ -6,6 +6,7 @@ import sys from pathlib import Path +import foundry import pytest import torch import torch.nn as nn @@ -170,6 +171,17 @@ def _spawn_with_preload(launch_mode): subprocess.check_call(cmd, env=env) +def test_graph_default_pool_is_forwarded(monkeypatch): + sentinel_stream = object() + monkeypatch.setattr(torch.cuda, "Stream", lambda: sentinel_stream) + cuda_graph = object() + + context = foundry.graph.__new__(foundry.graph) + foundry.graph.__init__(context, cuda_graph) + + assert context.pool == (None,) + + @pytest.mark.filterwarnings("ignore:TORCH_CUDA_ARCH_LIST is not set") def test_graph_save_and_load(): """Test CUDA graph save/load with binary dumping and allocator replay""" diff --git a/tests/test_vmm_alloc.py b/tests/test_vmm_alloc.py index 5df909f0..c0d14e8e 100644 --- a/tests/test_vmm_alloc.py +++ b/tests/test_vmm_alloc.py @@ -1,17 +1,27 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the Foundry project +import ctypes +import importlib.util import os +import struct import subprocess import sys from pathlib import Path +import foundry as fdry import pytest import torch +CUDA_SUCCESS = 0 +CUDA_ERROR_INVALID_VALUE = 1 +VMM_IPC_MAGIC = 0x564D4D32 + + +class CUipcMemHandle(ctypes.Structure): + _fields_ = [("reserved", ctypes.c_ubyte * 64)] -def _get_hook_so_path(): - import importlib.util +def _get_hook_so_path(): spec = importlib.util.find_spec("foundry.ops") if not spec or not spec.origin: raise RuntimeError("foundry.ops not found; ensure setup.py develop/pip install completed") @@ -26,8 +36,6 @@ def _get_hook_so_path(): def _run_core(): - import foundry as fdry - torch.cuda.init() device = torch.device("cuda:0") @@ -104,8 +112,6 @@ def _run_core_without_region(): def _run_core_size_parsing(): - import foundry as fdry - torch.cuda.init() device = torch.device("cuda:0") @@ -146,6 +152,447 @@ def _run_core_size_parsing(): print("[TEST] test_size_parsing PASSED") +def _run_core_zero_alignment_reserve(): + torch.cuda.init() + + base_addr = 0x7F4000000000 + region_size = 1024 * 1024 * 1024 + default_alignment = 2 * 1024 * 1024 + deliberately_unaligned_offset = default_alignment + 12345 + reserve_size = 2 * 1024 * 1024 + + driver = ctypes.CDLL(None) + reserve = driver.cuMemAddressReserve + reserve.argtypes = [ + ctypes.POINTER(ctypes.c_uint64), + ctypes.c_size_t, + ctypes.c_size_t, + ctypes.c_uint64, + ctypes.c_ulonglong, + ] + reserve.restype = ctypes.c_int + address_free = driver.cuMemAddressFree + address_free.argtypes = [ctypes.c_uint64, ctypes.c_size_t] + address_free.restype = ctypes.c_int + + ptr = ctypes.c_uint64() + with fdry.allocation_region(base_addr, region_size): + fdry.set_current_alloc_offset(deliberately_unaligned_offset) + result = reserve(ctypes.byref(ptr), reserve_size, 0, 0, 0) + assert result == 0 + expected_ptr = (base_addr + deliberately_unaligned_offset + default_alignment - 1) & ~( + default_alignment - 1 + ) + assert ptr.value == expected_ptr, ( + f"zero-alignment reserve returned {hex(ptr.value)}, expected " + f"default-aligned {hex(expected_ptr)}" + ) + expected_offset = expected_ptr - base_addr + reserve_size + assert fdry.get_current_alloc_offset() == expected_offset + assert address_free(ptr.value, reserve_size) == 0 + + +def _run_core_malformed_preallocated_ipc_handle(): + torch.cuda.init() + + base_addr = 0x574000000000 + region_size = 1024 * 1024 * 1024 + prealloc_size = 64 * 1024 * 1024 + allocation_size = 4 * 1024 * 1024 + uint64_max = (1 << 64) - 1 + + driver = ctypes.CDLL(None) + mem_alloc = driver.cuMemAlloc_v2 + mem_alloc.argtypes = [ctypes.POINTER(ctypes.c_uint64), ctypes.c_size_t] + mem_alloc.restype = ctypes.c_int + mem_free = driver.cuMemFree_v2 + mem_free.argtypes = [ctypes.c_uint64] + mem_free.restype = ctypes.c_int + ipc_get = driver.cuIpcGetMemHandle + ipc_get.argtypes = [ctypes.POINTER(CUipcMemHandle), ctypes.c_uint64] + ipc_get.restype = ctypes.c_int + ipc_open = driver.cuIpcOpenMemHandle + ipc_open.argtypes = [ + ctypes.POINTER(ctypes.c_uint64), + CUipcMemHandle, + ctypes.c_uint, + ] + ipc_open.restype = ctypes.c_int + ipc_close = driver.cuIpcCloseMemHandle + ipc_close.argtypes = [ctypes.c_uint64] + ipc_close.restype = ctypes.c_int + + padding = ctypes.c_uint64() + target = ctypes.c_uint64() + results = {} + fdry.set_allocation_region(base_addr, region_size) + try: + assert fdry.preallocate_region(prealloc_size) + assert mem_alloc(ctypes.byref(padding), allocation_size) == CUDA_SUCCESS + assert mem_alloc(ctypes.byref(target), allocation_size) == CUDA_SUCCESS + assert target.value > padding.value + + handle = CUipcMemHandle() + assert ipc_get(ctypes.byref(handle), target.value) == CUDA_SUCCESS + blob = bytearray(bytes(handle)) + ( + magic, + _exporter_pid, + original_ptr, + size, + _token, + chunk_base, + chunk_size, + _generation, + ) = struct.unpack_from("=IIQQQQQQ", blob) + assert magic == VMM_IPC_MAGIC + assert original_ptr == target.value + assert size > 0 + assert chunk_base == base_addr + assert chunk_size == prealloc_size + assert original_ptr > chunk_base + + remaining = chunk_base + chunk_size - original_ptr + mutations = { + "zero_original_ptr": (8, 0), + "zero_size": (16, 0), + "original_ptr_before_chunk": (8, chunk_base - 1), + "allocation_exceeds_chunk": (16, remaining + 1), + "allocation_end_overflows": (16, uint64_max - original_ptr + 2), + "chunk_size_without_base": (32, 0), + } + for name, (offset, value) in mutations.items(): + malformed_blob = bytearray(blob) + struct.pack_into("=Q", malformed_blob, offset, value) + malformed = CUipcMemHandle.from_buffer_copy(malformed_blob) + opened = ctypes.c_uint64() + result = ipc_open(ctypes.byref(opened), malformed, 1) + results[name] = result + if result == CUDA_SUCCESS: + assert ipc_close(opened.value) == CUDA_SUCCESS + finally: + if target.value: + assert mem_free(target.value) == CUDA_SUCCESS + if padding.value: + assert mem_free(padding.value) == CUDA_SUCCESS + fdry.free_preallocated_region() + fdry.stop_allocation_region() + + # Every malformed chunk descriptor is rejected before the interior pointer + # is derived. The accepting path (a valid descriptor that maps the chunk and + # returns base + (ptr - chunk_base)) cannot be exercised in-process: CUDA + # IPC is inter-process, and importing our own exported handle then calling + # cuMemSetAccess on a second same-process mapping returns + # CUDA_ERROR_INVALID_VALUE regardless of this validation. Valid cross-process + # chunk import plus interior-pointer derivation is covered end-to-end by the + # two-H100 nvl_ipc_prealloc DeepEP matrix case. + assert results == {name: CUDA_ERROR_INVALID_VALUE for name in mutations} + + +def _load_driver_ipc(): + """Bind the driver VMM + IPC entry points used by the cross-process probes.""" + driver = ctypes.CDLL(None) + mem_alloc = driver.cuMemAlloc_v2 + mem_alloc.argtypes = [ctypes.POINTER(ctypes.c_uint64), ctypes.c_size_t] + mem_alloc.restype = ctypes.c_int + mem_free = driver.cuMemFree_v2 + mem_free.argtypes = [ctypes.c_uint64] + mem_free.restype = ctypes.c_int + ipc_get = driver.cuIpcGetMemHandle + ipc_get.argtypes = [ctypes.POINTER(CUipcMemHandle), ctypes.c_uint64] + ipc_get.restype = ctypes.c_int + ipc_open = driver.cuIpcOpenMemHandle + ipc_open.argtypes = [ctypes.POINTER(ctypes.c_uint64), CUipcMemHandle, ctypes.c_uint] + ipc_open.restype = ctypes.c_int + ipc_close = driver.cuIpcCloseMemHandle + ipc_close.argtypes = [ctypes.c_uint64] + ipc_close.restype = ctypes.c_int + return mem_alloc, mem_free, ipc_get, ipc_open, ipc_close + + +# Offsets of the packed VMM2 handle fields (see cuIpcGetMemHandle in csrc/hook.cpp). +_BLOB_ORIGINAL_PTR = 8 +_BLOB_SIZE = 16 +_BLOB_CHUNK_SIZE = 40 +_BLOB_GENERATION = 48 + + +def _open_handle(ipc_open, blob): + """Open a (possibly mutated) VMM2 handle blob, returning (result, opened_ptr).""" + handle = CUipcMemHandle.from_buffer_copy(bytes(blob)) + opened = ctypes.c_uint64() + result = ipc_open(ctypes.byref(opened), handle, 1) + return result, opened + + +# CUDA VMM IPC is inter-process; a fully successful import (needed to establish +# a cached chunk mapping, and as the stale-generation positive control) must map +# a peer's allocation with cuMemSetAccess, which returns CUDA_ERROR_INVALID_VALUE +# for a *same-GPU* import of an allocation this process/GPU already owns. So the +# real cross-process regressions run rank 0 (exporter) and rank 1 (importer) on +# two distinct GPUs, exactly like the DeepEP nvl_ipc_prealloc matrix case. +_REGRESSION_BASE_ADDR = 0x566000000000 +_REGRESSION_REGION_SIZE = 1024 * 1024 * 1024 +_REGRESSION_PREALLOC_SIZE = 64 * 1024 * 1024 +_REGRESSION_ALLOCATION_SIZE = 4 * 1024 * 1024 + + +def _importer_check_cached_chunk(blob, ipc_open, ipc_close): + """Establish a cached whole-chunk mapping via a valid import, then present + same-key handles with inconsistent chunk geometry. Each reopen must be + validated against the ACTUAL cached mapping (not the incoming handle's own + attacker-controlled metadata) and rejected.""" + (_magic, _pid, _optr, _size, _token, chunk_base, chunk_size, _gen) = struct.unpack_from( + "=IIQQQQQQ", blob + ) + + valid_result, valid_ptr = _open_handle(ipc_open, blob) + if valid_result != CUDA_SUCCESS: + return False, f"valid cross-process chunk import failed: {valid_result}" + + results = {} + + # Inflated chunk size while the carve stays inside the real mapping. The + # incoming metadata is self-consistent, so only checking the carve against + # the cached mapping size can catch the chunk-size mismatch. + size_mismatch = bytearray(blob) + struct.pack_into("=Q", size_mismatch, _BLOB_CHUNK_SIZE, chunk_size * 2) + struct.pack_into("=Q", size_mismatch, _BLOB_ORIGINAL_PTR, chunk_base) + struct.pack_into("=Q", size_mismatch, _BLOB_SIZE, _REGRESSION_ALLOCATION_SIZE) + result, opened = _open_handle(ipc_open, size_mismatch) + if result == CUDA_SUCCESS: + ipc_close(opened.value) + results["chunk_size_mismatch_within_real"] = result + + # Inflated chunk size with a carve starting at the real chunk end: inside + # the bogus incoming chunk but outside the cached mapping. + carve_outside = bytearray(blob) + struct.pack_into("=Q", carve_outside, _BLOB_CHUNK_SIZE, chunk_size * 2) + struct.pack_into("=Q", carve_outside, _BLOB_ORIGINAL_PTR, chunk_base + chunk_size) + struct.pack_into("=Q", carve_outside, _BLOB_SIZE, _REGRESSION_ALLOCATION_SIZE) + result, opened = _open_handle(ipc_open, carve_outside) + if result == CUDA_SUCCESS: + ipc_close(opened.value) + results["carve_outside_cached_mapping"] = result + + ipc_close(valid_ptr.value) + + for name, code in results.items(): + if code != CUDA_ERROR_INVALID_VALUE: + return False, f"{name} not rejected on cache hit: {code}" + return True, "cached-chunk inconsistent metadata rejected" + + +def _importer_check_stale_generation(blob, ipc_open, ipc_close): + """Prove the exporter refuses to transfer a chunk descriptor for a stale + generation, while the live generation still imports (the chunk is live).""" + (_magic, _pid, _optr, _size, _token, _chunk_base, _chunk_size, generation) = struct.unpack_from( + "=IIQQQQQQ", blob + ) + + stale = bytearray(blob) + struct.pack_into("=Q", stale, _BLOB_GENERATION, generation + 1) + stale_result, stale_ptr = _open_handle(ipc_open, stale) + if stale_result == CUDA_SUCCESS: + ipc_close(stale_ptr.value) + return False, "stale-generation chunk import was not rejected" + if stale_result != CUDA_ERROR_INVALID_VALUE: + return False, f"stale-generation import failed with unexpected code: {stale_result}" + + live_result, live_ptr = _open_handle(ipc_open, blob) + if live_result != CUDA_SUCCESS: + return False, f"live-generation chunk import failed: {live_result}" + ipc_close(live_ptr.value) + return True, "stale generation rejected; live generation served" + + +def _importer_check_concurrent_race(blob, ipc_open, ipc_close): + """Contend two importer threads on the SAME peer chunk's FIRST open so both + miss the mapping cache and both import the whole chunk; one wins the + insertion and the other lands on the insertion-race-loser path, which must + adopt the winner's mapping and hand out an interior derived from the winner's + base (not leak its own temporary). Then, against the live cached mapping, + present same-key handles whose geometry is inconsistent with the ACTUAL + mapping: each must be validated against the real mapping and rejected with + CUDA_ERROR_INVALID_VALUE and no out-of-mapping pointer. A subsequent valid + open/close still works. + + Why the inconsistent request is a cache-hit rather than a third racer: the + CUDA driver cannot map a sub-range of an imported handle (see the whole-chunk + mapping comments in csrc/hook.cpp), so every importer maps the FULL chunk. + Any handle that survives the incoming-metadata check and imports + successfully has therefore mapped exactly the cached size with a carve inside + it, i.e. it is already consistent with the cached mapping by the time it + reaches the loser path; an inconsistent chunk_size instead fails its own + full-chunk import before ever reaching that branch. The loser-path + re-validation shares the exact validator proven reachable-rejectable here on + the cache-hit path, as defense-in-depth for that whole-chunk invariant.""" + import threading + + (_magic, _pid, _optr, _size, _token, chunk_base, chunk_size, _gen) = struct.unpack_from( + "=IIQQQQQQ", blob + ) + + barrier = threading.Barrier(2) + opens = {} + + def racer(idx): + barrier.wait() + opens[idx] = _open_handle(ipc_open, blob) + + threads = [threading.Thread(target=racer, args=(i,)) for i in range(2)] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + + for idx, (result, ptr) in opens.items(): + if result != CUDA_SUCCESS: + return False, f"concurrent valid open {idx} failed: {result}" + if ptr.value == 0: + return False, f"concurrent valid open {idx} returned a null pointer" + + # Both racers open the same carve, so the loser must have adopted the + # winner's whole-chunk mapping (same interior), not a leaked temporary. + ptr0 = opens[0][1].value + ptr1 = opens[1][1].value + if ptr0 != ptr1: + return False, f"concurrent opens of the same carve diverged: 0x{ptr0:x} vs 0x{ptr1:x}" + + # The mapping is live (refcount 2). Present same-key handles whose geometry + # is inconsistent with the ACTUAL mapping; each must be rejected on the + # cache-hit path before any pointer is handed back. + rejects = {} + + size_mismatch = bytearray(blob) + struct.pack_into("=Q", size_mismatch, _BLOB_CHUNK_SIZE, chunk_size * 2) + struct.pack_into("=Q", size_mismatch, _BLOB_ORIGINAL_PTR, chunk_base) + struct.pack_into("=Q", size_mismatch, _BLOB_SIZE, _REGRESSION_ALLOCATION_SIZE) + result, opened = _open_handle(ipc_open, size_mismatch) + if result == CUDA_SUCCESS: + ipc_close(opened.value) + rejects["chunk_size_mismatch"] = result + + carve_outside = bytearray(blob) + struct.pack_into("=Q", carve_outside, _BLOB_CHUNK_SIZE, chunk_size * 2) + struct.pack_into("=Q", carve_outside, _BLOB_ORIGINAL_PTR, chunk_base + chunk_size) + struct.pack_into("=Q", carve_outside, _BLOB_SIZE, _REGRESSION_ALLOCATION_SIZE) + result, opened = _open_handle(ipc_open, carve_outside) + if result == CUDA_SUCCESS: + ipc_close(opened.value) + rejects["carve_outside_mapping"] = result + + # Release both concurrent opens (refcount -> 0, mapping torn down). + ipc_close(ptr0) + ipc_close(ptr1) + + for name, code in rejects.items(): + if code != CUDA_ERROR_INVALID_VALUE: + return False, f"{name} not rejected against the live cached mapping: {code}" + + # Subsequent valid open/close of the genuine handle still works. + live_result, live_ptr = _open_handle(ipc_open, blob) + if live_result != CUDA_SUCCESS: + return False, f"post-race valid open failed: {live_result}" + ipc_close(live_ptr.value) + return True, "loser path adopted winner mapping; inconsistent geometry rejected" + + +_IMPORTER_CHECKS = { + "cached-chunk": _importer_check_cached_chunk, + "stale-generation": _importer_check_stale_generation, + "concurrent-race": _importer_check_concurrent_race, +} + + +def _run_2gpu_regression(scenario, local_rank): + """Two-GPU cross-process regression. Rank 0 exports a live preallocated-chunk + VMM2 handle and keeps its fd server alive; rank 1 (a different GPU) imports + it through the real hook and runs the scenario's checks.""" + import torch.distributed as dist + + # gloo (CPU) so the blob/status handoff never touches the VMM region. + dist.init_process_group("gloo", rank=local_rank, world_size=2) + mem_alloc, mem_free, ipc_get, ipc_open, ipc_close = _load_driver_ipc() + + padding = ctypes.c_uint64() + target = ctypes.c_uint64() + region_active = False + blob_box = [None] + status_box = [None] + try: + torch.cuda.set_device(local_rank) + torch.cuda.init() + + if local_rank == 0: + fdry.set_allocation_region(_REGRESSION_BASE_ADDR, _REGRESSION_REGION_SIZE) + region_active = True + assert fdry.preallocate_region(_REGRESSION_PREALLOC_SIZE) + assert mem_alloc(ctypes.byref(padding), _REGRESSION_ALLOCATION_SIZE) == CUDA_SUCCESS + assert mem_alloc(ctypes.byref(target), _REGRESSION_ALLOCATION_SIZE) == CUDA_SUCCESS + assert target.value > padding.value + handle = CUipcMemHandle() + assert ipc_get(ctypes.byref(handle), target.value) == CUDA_SUCCESS + blob = bytes(handle) + (magic, _pid, original_ptr, size, _token, chunk_base, chunk_size, generation) = ( + struct.unpack_from("=IIQQQQQQ", blob) + ) + assert magic == VMM_IPC_MAGIC + assert chunk_base == _REGRESSION_BASE_ADDR + assert chunk_size == _REGRESSION_PREALLOC_SIZE + assert original_ptr == target.value + assert size > 0 + assert generation != 0 + blob_box[0] = blob + + # Rank 0 hands the live handle to rank 1, keeps its server thread alive + # (blocked receiving the status), then both barrier before teardown. + dist.broadcast_object_list(blob_box, src=0) + if local_rank == 1: + status_box[0] = _IMPORTER_CHECKS[scenario](blob_box[0], ipc_open, ipc_close) + dist.broadcast_object_list(status_box, src=1) + dist.barrier() + finally: + if local_rank == 0: + if target.value: + mem_free(target.value) + if padding.value: + mem_free(padding.value) + if region_active: + fdry.free_preallocated_region() + fdry.stop_allocation_region() + dist.destroy_process_group() + + ok, message = status_box[0] + assert ok, message + + +def _spawn_2gpu_regression(scenario): + so_path = _get_hook_so_path() + base_env = os.environ.copy() + if base_env.get("LD_PRELOAD"): + base_env["LD_PRELOAD"] = f"{so_path}:{base_env['LD_PRELOAD']}" + else: + base_env["LD_PRELOAD"] = so_path + base_env["MASTER_ADDR"] = "127.0.0.1" + base_env["MASTER_PORT"] = str( + {"cached-chunk": 29753, "stale-generation": 29757, "concurrent-race": 29761}[scenario] + ) + + procs = [] + for rank in range(2): + env = base_env.copy() + cmd = [ + sys.executable, + str(Path(__file__).resolve()), + f"--2gpu-{scenario}", + str(rank), + ] + procs.append(subprocess.Popen(cmd, env=env)) + returncodes = [proc.wait() for proc in procs] + assert all(code == 0 for code in returncodes), f"regression ranks failed: {returncodes}" + + def _spawn_with_preload(test_mode): so_path = _get_hook_so_path() env = os.environ.copy() @@ -176,6 +623,55 @@ def test_size_parsing(): _spawn_with_preload("size-parsing") +@pytest.mark.filterwarnings("ignore:TORCH_CUDA_ARCH_LIST is not set") +def test_zero_alignment_reserve_stays_in_region(): + """CUDA's alignment=0 convention uses Foundry's default 2 MiB alignment.""" + _spawn_with_preload("zero-alignment-reserve") + + +@pytest.mark.filterwarnings("ignore:TORCH_CUDA_ARCH_LIST is not set") +def test_vmm_ipc_rejects_malformed_preallocated_chunk_handle(): + """Malformed VMM2 chunk metadata is rejected by the real hooked IPC path.""" + _spawn_with_preload("malformed-preallocated-ipc-handle") + + +def _cuda_device_count() -> int: + """GPU count without initializing CUDA in this orchestration parent.""" + result = subprocess.run( + ["nvidia-smi", "-L"], + stdout=subprocess.PIPE, + stderr=subprocess.DEVNULL, + text=True, + ) + if result.returncode != 0: + return 0 + return sum(1 for line in result.stdout.splitlines() if line.startswith("GPU ")) + + +@pytest.mark.skipif(_cuda_device_count() < 2, reason="requires 2 GPUs for cross-process IPC") +def test_vmm_ipc_cached_chunk_rejects_inconsistent_metadata(): + """A cache-hit reopen with same-key but inconsistent chunk metadata is + validated against the actual cached mapping and rejected (2-GPU cross-process).""" + _spawn_2gpu_regression("cached-chunk") + + +@pytest.mark.skipif(_cuda_device_count() < 2, reason="requires 2 GPUs for cross-process IPC") +def test_vmm_ipc_rejects_stale_generation_import(): + """The exporter refuses to serve a chunk descriptor for a stale generation, + while the live generation still imports (2-GPU cross-process).""" + _spawn_2gpu_regression("stale-generation") + + +@pytest.mark.skipif(_cuda_device_count() < 2, reason="requires 2 GPUs for cross-process IPC") +def test_vmm_ipc_concurrent_first_open_race_rejects_inconsistent_geometry(): + """Two importer threads contend on a peer chunk's first open so the + insertion-race-loser path is exercised (the loser adopts the winner's + whole-chunk mapping); inconsistent same-key geometry against the live mapping + is validated against the actual cached mapping and rejected with no + out-of-mapping pointer (2-GPU cross-process).""" + _spawn_2gpu_regression("concurrent-race") + + if __name__ == "__main__": if "--core" in sys.argv: _run_core() @@ -183,7 +679,19 @@ def test_size_parsing(): _run_core_without_region() elif "--size-parsing" in sys.argv: _run_core_size_parsing() + elif "--zero-alignment-reserve" in sys.argv: + _run_core_zero_alignment_reserve() + elif "--malformed-preallocated-ipc-handle" in sys.argv: + _run_core_malformed_preallocated_ipc_handle() + elif "--2gpu-cached-chunk" in sys.argv: + _run_2gpu_regression("cached-chunk", int(sys.argv[2])) + elif "--2gpu-stale-generation" in sys.argv: + _run_2gpu_regression("stale-generation", int(sys.argv[2])) + elif "--2gpu-concurrent-race" in sys.argv: + _run_2gpu_regression("concurrent-race", int(sys.argv[2])) else: test_vmm_allocation() test_vmm_allocation_without_region() test_size_parsing() + test_zero_alignment_reserve_stays_in_region() + test_vmm_ipc_rejects_malformed_preallocated_chunk_handle()