diff --git a/README.md b/README.md index bec64ee4..32f7abd7 100644 --- a/README.md +++ b/README.md @@ -82,12 +82,32 @@ Foundry ships engine integrations under `foundry/python/foundry/integration/`. P | Engine | Single GPU | DP | TP | EP | |---|:---:|:---:|:---:|:---:| | vLLM | โœ… | โœ… | ๐Ÿšง | โœ… | -| SGLang | โœ… | โœ… | ๐Ÿšง | โœ… | +| SGLang `0.5.15.post1` | โœ… | โœ… | โœ…[^sglang-tp] | โœ… | | TensorRT-LLM | ๐Ÿšง | ๐Ÿšง | ๐Ÿšง | ๐Ÿšง | โœ… validated end-to-end (SAVE โ†’ LOAD โ†’ query)  ยท  ๐Ÿšง not yet -The adapted vLLM / SGLang / TensorRT-LLM forks will be released alongside this repo at `foundry-org/vllm`, `foundry-org/sglang`, `foundry-org/TensorRT-LLM`. +[^sglang-tp]: SGLang conventional tensor parallelism is validated only for + dense TP=2 ร— DP=2 with PP=1 and EP=1. All other conventional TP topologies + remain unsupported. + +The SGLang integration is latest-only for upstream `v0.5.15.post1` +(`0b3bb0c`), CUDA 13.0, and PyTorch 2.11/cu130 on H100. Full decode with +prefill graphs disabled is validated for single-GPU FlashInfer and regular +DP=2 FlashInfer on Qwen3-1.7B (52 graphs/rank), exact dense TP=2 ร— DP=2 +FlashInfer on Qwen3-8B (36 graphs/rank), and DP-attention EP=2 and EP=4 with +DeepEP low-latency, DeepGEMM, and FA3 on Qwen3-30B-A3B-FP8 (20 graphs/rank). +All passed baseline/SAVE/fresh LOAD/query with text/token parity and no +recapture. EP=4 uses DP-attention/EP rather than conventional TP attention, +even though SGLang's process topology sets `tp_size=4`. Existing SGLang +archives are incompatible and require a fresh SAVE. See +[`recipe/sglang/README.md`](recipe/sglang/README.md). + +For the current Foundry-only SGLang workflow, check out upstream +`sgl-project/sglang` at the tagged release and apply the documented four-file +patch. `foundry-org/sglang` should carry that same patch when its release branch +is published; the current instructions do not assume that branch already +exists. Adapted engine forks will be released alongside this repository. ### Performance diff --git a/ROADMAP.md b/ROADMAP.md index aa139d65..a68941f7 100644 --- a/ROADMAP.md +++ b/ROADMAP.md @@ -30,8 +30,8 @@ - [x] Quantized MoE with DeepGemm - [x] Drop-in integration layer for CUDA graph persistence in vLLM - [x] Sync with latest SGLang release - - [ ] EP on SGLang - - [ ] TP on SGLang + - [x] EP on SGLang (DP-attention + DeepEP, validated at EP=2 and EP=4) + - [x] Exact dense TP=2 ร— DP=2 on SGLang (all other TP topologies unsupported) ## Stage 5: Disaggregated and Large-Scale Serving diff --git a/csrc/hook.cpp b/csrc/hook.cpp index d552d08d..a78a35ef 100644 --- a/csrc/hook.cpp +++ b/csrc/hook.cpp @@ -14,12 +14,20 @@ #include #include #include +#include #include #include #include #include #include #include +#include +#include +#include +#include +#include +#include +#include #include #include #include @@ -253,6 +261,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; @@ -267,6 +296,8 @@ static std::mutex hook_events_mutex; static CUdeviceptr hook_recording_start_base_addr{0}; static inline size_t align_to(size_t addr, size_t alignment) { + if (alignment == 0) + return addr; return (addr + alignment - 1) & ~(alignment - 1); } @@ -2657,6 +2688,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; }); @@ -2860,8 +2892,254 @@ 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 -> (allocation handle it was exported from, fd). The handle is +// kept so VA reuse after free/realloc invalidates the cached fd. +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. +struct VmmIpcRequest { + uint64_t ptr; + uint64_t token; +}; + +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; + +// 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. + vmm_ipc_exported_fds.visit( + (CUdeviceptr)req.ptr, + [&](const std::pair>& + kv) { fd = fcntl(kv.second.second, 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.second); + 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, 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}; + 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.second); + return true; + }); +} // Hook for cuIpcGetMemHandle - intercept Driver API to support VMM allocations CUresult cuIpcGetMemHandle(CUipcMemHandle* pHandle, CUdeviceptr dptr) { @@ -2878,7 +3156,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 +3207,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.first == export_handle) + fd = kv.second.second; + }); + 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, std::make_pair(export_handle, new_fd)), + [&](std::pair>& kv) { + // Entry exists: either a racing thread won (same handle - drop + // ours) or it is stale from a freed allocation (replace). + if (kv.second.first == export_handle) { + close(new_fd); + new_fd = kv.second.second; + } else { + close(kv.second.second); + kv.second = std::make_pair(export_handle, new_fd); + } + }); + 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 +3306,75 @@ 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; + // 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()) { + 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) { + vmm_ipc_exported_fds.visit( + fetch_key, + [&](const std::pair>& + kv) { fd = fcntl(kv.second.second, F_DUPFD_CLOEXEC, 0); }); + } else { + fd = vmm_ipc_fetch_fd(exporter_pid, token, 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 +3383,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 +3398,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 +3469,87 @@ 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. + 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); + 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 +3567,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 +3615,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 +3830,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 +3857,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); diff --git a/docs/overview.md b/docs/overview.md index 46a92063..86253dff 100644 --- a/docs/overview.md +++ b/docs/overview.md @@ -1,72 +1,76 @@ -# Foundry Inference-Engine Integrations +# Foundry inference-engine integrations -Foundry ships two inference-engine integrations that wrap an unmodified serving stack with cold-start graph persistence: +Foundry supplies runtime integration packages for vLLM and SGLang. A small +engine-side shim activates each package; deterministic memory, graph +persistence, packed module restoration, and lifecycle hooks remain in Foundry. | Engine | Integration code | Documentation | |---|---|---| -| vLLM | `foundry/python/foundry/integration/vllm/` | [`vllm/overview.md`](vllm/overview.md) | -| SGLang | `foundry/python/foundry/integration/sglang/` | [`sglang/overview.md`](sglang/overview.md) | +| vLLM | `python/foundry/integration/vllm/` | [`vllm/overview.md`](vllm/overview.md) | +| SGLang `0.5.15.post1` | `python/foundry/integration/sglang/` | [`sglang/overview.md`](sglang/overview.md) | -Both follow the same shape: a tiny shim in the host engine's source tree re-exports `install_hooks(...)` and `get_graph_extension_mode()`, and the shim is called once from the host engine's config object's `__post_init__`. `install_hooks` then installs a small set of runtime monkey-patches that wrap the engine's lifecycle methods. Everything substantive โ€” VMM region setup, graph save/load, warmup skipping โ€” lives in the foundry package. +## Shared model -## Mental model +Both integrations: -```mermaid -flowchart BT - subgraph engine["Host engine (unmodified except shim + minimal guards)"] - VLLM["vLLM core
~97 lines across 5 files"] - SGL["SGLang core
~47 lines across 4 files"] - end +1. reserve a fixed-address VMM region before model/runtime allocations; +2. capture full or piecewise CUDA graphs during SAVE; +3. persist graph structure, packed device code, topology groups, and allocation + state; +4. recreate model/runtime allocations at the same offsets in a fresh LOAD + process; +5. batch-build graph templates and link on-demand members; +6. place `libcuda_hook.so` in child-process `LD_PRELOAD` before CUDA starts. - subgraph shim["Engine-side shims (re-export + activation)"] - VLLM_SHIM["vllm/compilation/foundry_shim.py
re-export + ImportError fallback"] - SGL_SHIM["sglang/srt/foundry_shim.py
re-export + force-disable flags"] - end +Each engine owns different capture and warmup seams, so the engine-specific +hooks and archive metadata are intentionally separate. - subgraph foundry["foundry/python/foundry/integration/"] - CFG["config.py โ€” CUDAGraphExtensionConfig"] - RT["runtime.py โ€” WarmupState, VMM region,
final_alloc_offset"] - GOPS["graph_ops.py โ€” save/load, manifest, fatbins"] - HOOKS["hooks.py โ€” install_hooks() + monkey-patches"] - end +## Integration comparison - VLLM -- "CompilationConfig.__post_init__" --> VLLM_SHIM - SGL -- "ServerArgs.__post_init__" --> SGL_SHIM - VLLM_SHIM -- "install_hooks(compilation_config)" --> HOOKS - SGL_SHIM -- "install_hooks(server_args)" --> HOOKS - - HOOKS --> CFG - HOOKS --> RT - HOOKS --> GOPS -``` - -## What both integrations share - -- **SAVEโ†’LOAD cycle.** SAVE runs the full cold-start path once, captures CUDA graphs + kernel fatbins + a warmup-state JSON, and writes them to a per-rank workspace. LOAD restores from the workspace, skipping graph capture and kernel warmup. -- **Deterministic VMM region.** Both integrations call `foundry::set_allocation_region(base_addr, region_size)` before NCCL / distributed init. Every allocation inside the region (model weights, KV pool, attention workspaces, captured-graph alloc events) lands at a byte-deterministic offset, so SAVE and LOAD address the same VMM offsets. Only memory *referenced by captured CUDA graphs* needs the region; the vLLM integration now calls `foundry::stop_allocation_region()` at the end of `capture_model` so subsequent runtime allocations (sampler warmup, scheduler buffers, inference workspaces) fall back to the standard CUDA caching allocator instead of paying VMM overhead. The SGLang side will follow. -- **`final_alloc_offset` watermark.** SAVE records the final cursor position; LOAD preallocates the same range via a single `cuMemCreate+cuMemMap`, then bump-allocates within it. Because `stop_allocation_region` fires before sampler warmup, the watermark reflects only the VMM-tracked working set, not transient post-capture allocations. -- **Manifest-driven graph load.** Captured graphs are grouped into topology equivalence classes; one graph per group is built as a *template* and the rest become *on-demand* graphs that share the template's executor via per-graph node updates. -- **Spawn-time `LD_PRELOAD`.** `libcuda_hook.so` must be in `LD_PRELOAD` before any CUDA call so its `cuMemAlloc_v2` / `cuModuleLoad*` interposers run in every process. The integrations set the env var before child spawn; the serve scripts also export it from the shell as defense-in-depth. - -## What differs between the two - -| Aspect | vLLM | SGLang | +| Aspect | vLLM | SGLang `0.5.15.post1` | |---|---|---| -| Graph capture seam | `CUDAGraphWrapper.__call__` (per piecewise wrapper) | `CudaGraphRunner.capture` (centralized full-graph) | -| Pre-graph warmup | `kernel_warmup` (no-op), compile-warmup `_dummy_run` loop (empty under our config), sampler-warmup full forward (skipped via signature-targeted `_dummy_run` patch โ€” see vllm/hooks.md ยง3) | `kernel_warmup` + 2 in-place warmup forwards inside `capture_one_batch_size` | -| Profile forward | yes (`memory_profiling` runs full forward) โ†’ two-pass SAVE required | no (only `_profile_available_bytes` samples free memory) โ†’ single-pass SAVE | -| Compile interaction | `torch.compile` integrated; disabled on LOAD via `do_not_compile=True` | `torch.compile` is the piecewise path; force-disabled via `disable_piecewise_cuda_graph=True` | -| Per-bs metadata | allocated **inside** captured graph (recorded as alloc events) | allocated **outside** captured graph (FlashInfer wrappers) โ†’ pre-pass init + `reuse_pre_pass_init` shim needed | -| MoE / EP | DeepEP all2all + NVSHMEM module init handled in `hooks.py` (`_patch_deepep` forces `use_fabric=True`, `_patch_prepare_comm_buffer` post-hook flushes `init_nvshmem_for_loaded_modules`) | coming soon | -| DP empty-shard phantom | `_patch_dummy_runs` lets `Worker.execute_dummy_batch`'s `_dummy_run(uniform_decode=True)` pass through to the real forward so the empty shard joins `coordinate_batch_across_dp` and the EP all2all (see [`vllm/hooks.md`](vllm/hooks.md) ยง3 and `claude-doc/invariants.md` ยง5) | n/a | -| Async graph builds | kicked off **after** weight load (driver contention with `load_weights` was net-negative); overlaps with KV-cache + scheduler bring-up instead | timing-independent: launched inside the capture context | -| Process model | uniproc / MP / Ray executors | `spawn`-based `mp.Process`; needs per-child `install_hooks` | -| Direct edits in host repo | 5 files, ~97 lines | 4 files, ~47 lines | - -The per-engine `overview.md` pages explain those differences in detail and how to use each integration. +| Graph scope | Configured vLLM piecewise/full integration modes | Validated: full decode only; prefill disabled | +| Capture seam | `CUDAGraphWrapper.__call__` | `DecodeCudaGraphRunner.capture` orchestration and `FullCudaGraphBackend.capture_one` storage | +| Run-once warmup seam | vLLM compile/kernel/sampler lifecycle hooks | `BaseRunner.warmup` | +| Graph identity | vLLM piecewise graph metadata | `ShapeKey(size, stream_idx, variant_label)` | +| Backend storage | vLLM wrapper state | `FullCudaGraphBackend._graphs`, `_outputs`, `_pool` | +| Attention preparation | vLLM-specific compile/capture path | Separate out-of-graph planning and in-graph GPU metadata work | +| Profile behavior | Full profile forward; two-pass SAVE recipe | Pool-memory measurement only; one SAVE process | +| Process model | Uniprocess, multiprocessing, or Ray executors | Multiprocessing `spawn`; explicit hook install in scheduler and DP-controller children | +| Parallel scope | Single GPU, DP, and documented EP configuration | Validated on H100: single GPU FlashInfer, regular DP=2 FlashInfer, exact dense TP=2 ร— DP=2 FlashInfer, and DP-attention EP=2/EP=4 + DeepEP low-latency + FA3 | + +SGLang validation used upstream commit `0b3bb0c`, CUDA 13.0, and PyTorch +2.11/cu130. Qwen3-1.7B restored 52 graphs for single GPU and 52 per rank for +DP=2. Dense Qwen3-8B TP=2 ร— DP=2 restored 36 graphs per rank across rank +workspaces 0โ€“3 at the identical `69822578688` final offset, with parity on +explicitly routed DP0 and DP1 queries. Qwen3-30B-A3B-FP8 restored 20 graphs per +rank for both EP=2 and EP=4; EP=4 restored 416 event nodes and showed +four-rank NVSHMEM/bootstrap/adapter restoration. Every mode passed +baseline/SAVE/fresh LOAD/query with text/token parity and no recapture. + +For SGLang, every conventional TP topology except exact dense TP=2 ร— DP=2 +with PP=1 and EP=1 is unsupported. Custom all-reduce, NCCL CUMEM/NVLS, +pipeline parallelism, speculative decoding, LoRA/PDMux, dLLM, hidden-state +capture, memory-saver graphs, and prefill graphs are not supported. EP=4 sets +`tp_size=4` as part of SGLang's DP-attention process topology; it is expert +parallelism, not conventional TP attention. + +## Archive authority + +The SGLang integration adds a versioned `sglang_graph_index.json` to each rank +workspace. It records capture order, phase, output kind, filename, and complete +shape identity. LOAD validates that index plus `warmup_state.json` schema, +SGLang version, and resolved decode/prefill backends before restoring graphs. +Earlier SGLang archives require a fresh SAVE. + +vLLM uses its own integration-specific graph records. Both engines use +Foundry's `graph_manifest.json` for native topology/template relationships. ## Reading order -- New to the integrations: read `vllm/overview.md` or `sglang/overview.md` end-to-end. -- Debugging a SAVEโ†”LOAD bug: read `vllm/memory-consistency.md` or `sglang/memory-consistency.md`. -- Porting to a new engine: this file + `/direct-edits.md` + `/hooks.md` show the minimal contract. +- Start with [`vllm/overview.md`](vllm/overview.md) or + [`sglang/overview.md`](sglang/overview.md). +- Use the corresponding `direct-edits.md` for the external engine shim. +- Use `memory-lifecycle.md` and `memory-consistency.md` when debugging + allocation parity. +- Use `save-load-workflow.md` for commands and archive validation. diff --git a/docs/sglang/direct-edits.md b/docs/sglang/direct-edits.md index 49c983c9..35e19891 100644 --- a/docs/sglang/direct-edits.md +++ b/docs/sglang/direct-edits.md @@ -1,93 +1,160 @@ -# SGLang Direct Edits +# SGLang v0.5.15.post1 direct edits -Total footprint in the `sglang/` repo: **47 lines across 4 files**. Every other line of integration logic lives in `foundry.integration.sglang`. +Foundry's runtime integration lives in this repository. For the current +Foundry-only workflow, clone upstream `sgl-project/sglang` at commit +`0b3bb0cbe31873994c9f989fddfe2f87ca839fdd` (`v0.5.15.post1`) and apply the +following four-file shim. -## 1. `sglang/srt/foundry_shim.py` (new, ~27 lines) +`foundry-org/sglang` should carry the same patch once its release branch is +published. This guide does not imply that the current fork already contains +the patch. -The activation entry point. Re-exports the foundry public API and forces flags that conflict with SAVE/LOAD determinism. +## 1. Add `python/sglang/srt/foundry_shim.py` ```python +from foundry.integration.sglang import install_hooks +from sglang.srt.model_executor.cuda_graph_config import Backend + + def apply_server_args(server_args) -> None: - cfg_path = getattr(server_args, "foundry_graph_extension_config_path", None) - if not cfg_path: + if not server_args.foundry_graph_extension_config_path: return - server_args.disable_piecewise_cuda_graph = True + server_args.cuda_graph_config.decode.backend = Backend.FULL + server_args.cuda_graph_config.prefill.backend = Backend.DISABLED server_args.enable_profile_cuda_graph = False server_args.disable_flashinfer_autotune = True - - from foundry.integration.sglang.hooks import install_hooks install_hooks(server_args) ``` -`install_hooks` is idempotent (guarded by a module-level `_INSTALLED` flag) so calling it multiple times across processes is safe. +The shim receives already resolved `cuda_graph_config` objects. It forces the +only accepted graph shape: full decode and disabled prefill. It also disables +graph profiling and FlashInfer autotuning before calling Foundry's public +`foundry.integration.sglang.install_hooks` entry point. + +## 2. Edit `python/sglang/srt/server_args.py` + +Add the field in the `Custom hooks, probe, and plugins` section: + +```python +foundry_graph_extension_config_path: A[ + Optional[str], "Path to the Foundry CUDA graph extension TOML config." +] = None +``` + +SGLang v0.5.15.post1 derives `--foundry-graph-extension-config-path` +automatically from `A[...]` dataclass metadata through +`add_cli_args_from_dataclass`. Do not add a manual `parser.add_argument`; +manual entries in this release are reserved for deprecated or dynamic flags. + +Keep `__post_init__` as an ordered dispatcher. Add a narrow helper: + +```python +def _handle_foundry_graph_extension(self): + if self.foundry_graph_extension_config_path: + # Foundry is optional; import its SGLang shim only when configured. + from sglang.srt.foundry_shim import apply_server_args -## 2. `sglang/srt/server_args.py` (~12 lines) + apply_server_args(self) +``` -Adds the config-path field, the CLI argument, and the `__post_init__` activation: +Call it after graph resolution and every platform-default pass, but before GPU +memory sizing: ```python -class ServerArgs: - ... - foundry_graph_extension_config_path: Optional[str] = None - ... - - def __post_init__(self): - ... - if self.foundry_graph_extension_config_path: - from sglang.srt.foundry_shim import apply_server_args - apply_server_args(self) - ... - - @staticmethod - def add_cli_args(parser): - ... - parser.add_argument( - "--foundry-graph-extension-config-path", - type=str, - default=ServerArgs.foundry_graph_extension_config_path, - help="Path to Foundry CUDA graph extension TOML config.", - ) +self._handle_cuda_graph_config() + +self._handle_hpu_backends() +self._handle_cpu_backends() +self._handle_npu_backends() +self._handle_mps_backends() +self._handle_xpu_backends() +current_platform.apply_server_args_defaults(self) + +self._handle_foundry_graph_extension() + +gpu_mem = get_device_memory_capacity(self.device) +self._handle_gpu_memory_settings(gpu_mem) ``` -The field must live on `ServerArgs` (not be monkey-attached) because: +This order is required. Foundry must mutate resolved +`cuda_graph_config.decode.backend` and +`cuda_graph_config.prefill.backend`; `_handle_gpu_memory_settings(gpu_mem)` +must then size decode graph batches and memory from those final values. -- `argparse` consumes it from the CLI. -- `ServerArgs` is pickled across the `spawn` boundary; only declared fields survive serialization. +The declared field is needed because SGLang's dataclass-generated CLI and +`ServerArgs.from_cli_args` handle declared fields. It also makes the value part +of the normal `ServerArgs` object passed to spawned children. The reason is not +that Python pickle cannot serialize arbitrary attributes. -## 3. `sglang/srt/managers/scheduler.py` (~4 lines) +## 3. Edit `python/sglang/srt/managers/scheduler.py` + +Install hooks explicitly in every scheduler child immediately after plugins +load: ```python -def run_scheduler_process(server_args, ...): +def run_scheduler_process( + server_args: ServerArgs, + # ... existing parameters ... +): load_plugins() if server_args.foundry_graph_extension_config_path: from sglang.srt.foundry_shim import apply_server_args + apply_server_args(server_args) - ... + + dp_rank = configure_scheduler_process( + # ... existing arguments ... + ) ``` -This is the first thing every spawned scheduler child runs. Without it, foundry's runtime monkey-patches aren't installed in the scheduler process and the model-loading lifecycle takes the upstream path. +## 4. Edit `python/sglang/srt/managers/data_parallel_controller.py` -## 4. `sglang/srt/managers/data_parallel_controller.py` (~4 lines) +Install hooks explicitly in the DP-controller child after logging is +configured and before it constructs child schedulers: ```python -def run_data_parallel_controller_process(server_args, ...): +def run_data_parallel_controller_process( + server_args: ServerArgs, + port_args: PortArgs, + pipe_writer, + run_scheduler_process_func: Callable = run_scheduler_process, +): + setproctitle.setproctitle("sglang::data_parallel_controller") + faulthandler.enable() + kill_itself_when_parent_died() parent_process = psutil.Process().parent() + configure_logger(server_args) if server_args.foundry_graph_extension_config_path: from sglang.srt.foundry_shim import apply_server_args + apply_server_args(server_args) - ... + + if server_args.enable_trace: + # ... existing code ... ``` -Same purpose as the scheduler call, but for the DP-controller child. This child later spawns its own scheduler subprocesses, so it must have foundry's spawn-site patches installed before doing so โ€” otherwise its nested `proc.start()` calls won't set `LD_PRELOAD` in their environments. +Both child-process calls are mandatory. SGLang uses the multiprocessing +`spawn` model, which transfers data but does not inherit parent-process Python +monkey-patches. The DP controller later creates scheduler children, so both +levels must install Foundry's hooks before their launch paths execute. -## Why these can't be moved into the integration layer +## Verify the external checkout + +After applying the four edits, install the checkout rather than exposing it as +source only: + +```bash +SETUPTOOLS_SCM_PRETEND_VERSION=0.5.15.post1 \ + pip install -e sglang/python \ + --extra-index-url https://download.pytorch.org/whl/cu130 +python -c 'from importlib.metadata import version; assert version("sglang") == "0.5.15.post1"' +sglang serve --help | rg -- '--foundry-graph-extension-config-path' +``` -| Edit | Why direct | -|---|---| -| `foundry_shim.py` | Has to be importable as `sglang.srt.foundry_shim`. It's the only file foundry's hooks can rely on being present in the sglang import path. | -| `ServerArgs.foundry_graph_extension_config_path` | Needs to be a real dataclass field (argparse + pickle). | -| `ServerArgs.__post_init__` call | Earliest activation point in the parent process. | -| `run_scheduler_process` install | `spawn` does not inherit Python state; the child needs an explicit re-install. | -| `run_data_parallel_controller_process` install | Same reason; this child also spawns further children. | +The tracked shim makes the tagged checkout dirty. Because SGLang uses +`setuptools_scm`, an unqualified editable install would derive a dirty/local +version. The pretend-version variable deterministically writes exact +`0.5.15.post1` distribution metadata at install time. SAVE and LOAD continue +to reject missing or non-exact metadata. diff --git a/docs/sglang/hooks.md b/docs/sglang/hooks.md index cf5b4c82..f04178f7 100644 --- a/docs/sglang/hooks.md +++ b/docs/sglang/hooks.md @@ -1,214 +1,192 @@ -# SGLang Hook Surface - -Every monkey-patch `install_hooks(server_args)` installs. +# SGLang hook surface + +`foundry.integration.sglang.install_hooks(server_args)` is the public, +idempotent entry point for SGLang `0.5.15.post1`. It loads the Foundry TOML, +validates the required graph configuration, and installs the hooks below. + +## Installation order + +```text +ModelRunner.init_torch_distributed +ModelRunnerKVCacheMixin.init_memory_pool +ModelRunner.load_model +BaseRunner.warmup +DecodeCudaGraphRunner.capture +FullCudaGraphBackend.capture_one +process launch sites +``` -## Patches installed by `install_hooks` +The latest runner classes are imported only while hooks install, so Foundry +itself remains usable without the optional SGLang dependency. -In install order: +## Configuration validation -1. `_patch_init_torch_distributed` -2. `_patch_init_memory_pool` -3. `_patch_load_model` -4. `_patch_kernel_warmup` -5. `_patch_cuda_graph_capture` -6. `_patch_spawn_sites` +Before patching, Foundry requires: -> 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. +- resolved decode backend `full`; +- resolved prefill backend `disabled`; +- FlashInfer for single-GPU and regular-DP decode; +- for conventional tensor parallelism, exactly + `tp_size=2, dp_size=2, pp_size=1, ep_size=1` with FlashInfer, + `--disable-custom-all-reduce`, `NCCL_CUMEM_ENABLE=0`, and + `NCCL_NVLS_ENABLE=0`; +- FA3 plus DeepEP low-latency for DP-attention. -### 1. `ModelRunner.init_torch_distributed` +It rejects every other conventional TP topology, PP, speculative, LoRA/PDMux, +dLLM, hidden-state-output, memory-saver, and debug-graph configurations. It +also rejects custom/symmetric-memory collectives and NCCL NVLS for the exact +TP=2 ร— DP=2 mode. -**Bind the device first (multi-GPU correctness).** Upstream `init_torch_distributed` -calls `torch.get_device_module(self.device).set_device(self.gpu_id)` *inside* -itself. But `setup_graph_extension` reserves the VMM region with -`set_allocation_region`, which binds to **whatever CUDA device is current at call -time** โ€” and we call it *before* upstream. For DP/TP/EP rank > 0 the current -device is still `cuda:0` at that point, so without intervention the region lands -on the wrong GPU and the rank's later allocations fault with an **async illegal -memory access** that surfaces at an unrelated op (e.g. the first -`torch.cuda.Stream()` in `ModelRunner.__init__`). The patch therefore mirrors -upstream's `set_device(self.gpu_id)` up front (guarded on `self.device == "cuda"`) -before reserving the region. Single-GPU and rank 0 are unaffected (`gpu_id == 0`), -which is why this only appeared once DP was exercised. +The validated DeepEP scope is DP-attention at EP=2 and EP=4. EP=4 carries +`tp_size=4` in SGLang's process topology, but this branch follows the +DP-attention path and is not conventional TP attention. -Before upstream: `set_device(self.gpu_id)`, then `setup_graph_extension(...)` -reserves the VMM region, loads cached fatbins (LOAD only via -`load_cuda_modules_and_libraries`), and eagerly initializes the cuBLAS handle into -scratch space. +## Distributed initialization -After upstream: `skip_to_scratch_boundary()` forces the cursor to `cfg.scratch_space_size`. +The `ModelRunner.init_torch_distributed` wrapper sets the runner's CUDA device +before Foundry reserves its VMM range. This is required for ranks above zero: +the allocation region binds to the device current at reservation time. -``` -[Foundry] SGLang alloc_offset[after_setup_graph_ext]=โ€ฆ -[Foundry] SGLang alloc_offset[after_init_torch_dist]=โ€ฆ -[Foundry] SGLang alloc_offset[after_scratch_skip]=โ€ฆ -``` +Before upstream initialization, the hook: -Validated on DP=2 (Qwen3-1.7B, one full replica per rank): both ranks reach an -identical `final_alloc_offset` across both SAVE passes and LOAD replays each -rank's graphs at that offset. Multi-rank runs also export -`NCCL_CUMEM_ENABLE=0` / `NCCL_NVLS_ENABLE=0` from the serve script โ€” the CUMEM -P2P and NVLS multicast fast paths `cuMemMap` with driver-capability flags the -foundry VMM region doesn't carry. +1. computes the rank workspace; +2. prepares packed CUDA modules on LOAD; +3. reserves the configured VMM region; +4. eagerly creates the cuBLAS handle. -### 2. `ModelRunnerKVCacheMixin.init_memory_pool` +After upstream initialization and communicator warmup, it advances the VMM +cursor to the configured scratch boundary. Nondeterministic startup +allocations remain below that boundary. -- **SAVE**: call upstream `init_memory_pool` unchanged; after it returns, serialize the resolved `MemoryPoolConfig` (via `dataclasses.asdict`) into `warmup_state.json`. -- **LOAD**: skip upstream profiling. Load `MemoryPoolConfig` from `warmup_state.json`. Call `torch.cuda.empty_cache()` to mirror SAVE's `_resolve_memory_pool_config โ†’ get_available_gpu_memory(empty_cache=True)` side effect. Then call `_apply_memory_pool_config(config)` directly. +### Dense TP VMM-IPC bridge -The `empty_cache()` mirror is load-bearing โ€” see [`memory-consistency.md`](memory-consistency.md) ยง5. +The exact TP=2 ร— DP=2 topology shares graph-referenced VMM allocations across +scheduler processes through CUDA IPC. Foundry's driver hook exports the whole +physical VMM chunk, records the requested tensor's interior offset, and sends +the live POSIX file descriptor over a same-UID abstract Unix socket with +`SCM_RIGHTS`. Exporter PID/token and chunk-generation checks reject stale +handles; an importer can reserve a relocated VA when the exporter's address +collides locally. -### 3. `ModelRunner.load_model` +Four-H100 validation mapped DP0/TP0, DP0/TP1, DP1/TP0, and DP1/TP1 to +workspaces `rank_0` through `rank_3`. All ranks restored 36 graphs at final +offset `69822578688` without VMM-IPC, CUDA, or preallocation errors. -Currently a passthrough wrapper. Kept as a future hook point for a LOAD-time `start_graph_builds(all_paths)` overlap with weight IO; the current LOAD path is fast enough (~80 ms total) that overlap isn't needed. +## Memory-pool initialization -### 4. `ModelRunner.kernel_warmup` +On SAVE, upstream `ModelRunnerKVCacheMixin.init_memory_pool` runs unchanged. +Foundry then serializes `dataclasses.asdict(self.memory_pool_config)` and +compatibility metadata to `warmup_state.json`. -No-op on SAVE/LOAD: - -``` -[Foundry] SGLang kernel_warmup skipped in mode -``` +On LOAD, Foundry first validates the warmup state against the installed +distribution and current resolved graph configuration. It reconstructs +`MemoryPoolConfig`, calls `torch.cuda.empty_cache()` to mirror SAVE's profiling +side effect, and invokes `_apply_memory_pool_config` without profiling again. -The FlashInfer autotune path inside is shut off both by this no-op and by the forced `disable_flashinfer_autotune` flag. +## Model loading -### 5. `CudaGraphRunner` (the largest patch) +`ModelRunner.load_model` currently delegates directly to upstream. It is kept +as a stable integration point; graph builds are not overlapped with weight +loading. -Four sub-patches: +## `BaseRunner.warmup` -#### 5a. `_create_device_graph` +SAVE and LOAD return without executing the upstream warmup/autotune body: -On SAVE returns `foundry.CUDAGraph()` instead of `torch.cuda.CUDAGraph()`. +```text +[Foundry] SGLang runner warmup skipped in mode +``` -#### 5b. `_capture_graph` +The external shim also sets `disable_flashinfer_autotune = True`. -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. +## `DecodeCudaGraphRunner.capture` -#### 5c. `capture_one_batch_size` +Mode `none` delegates to upstream. -On SAVE wraps the `forward` callable with a 3-call counter: +### SAVE -```python -counter = [0]; real_forward = forward -def warmup_skipping_forward(*args, **kwargs): - counter[0] += 1 - if counter[0] <= 2: - return None - return real_forward(*args, **kwargs) -forward = warmup_skipping_forward -``` +For FlashInfer, Foundry first walks `reversed(runner.capture_bs)`. For each +size, `capture_prepare(size)` constructs the canonical `ForwardBatch`, then +`init_forward_metadata_out_graph(forward_batch, in_capture=True)` creates the +allocation-bearing metadata. -This suppresses the two pre-capture warmup forwards SGLang does in `for _ in range(2): run_once()`. Only the third invocation โ€” inside `_capture_graph`'s graph capture context โ€” runs the real forward. JIT and autotune allocations that those warmups would normally trigger now happen inside the captured graph and become foundry alloc events that replay verbatim on LOAD. (See `memory-consistency.md` ยง2.) +During the subsequent upstream orchestration, a temporary wrapper: -After upstream returns, calls `save_graph(graph, output, key)`. The key matches the inline shape upstream uses for `self.graphs[key]` in `_capture_one_stream`: +1. selects the prepared metadata by batch size; +2. sets it as the active `forward_metadata`; +3. calls the current out-of-graph planner with `in_capture=False`. -```python -key = bs if stream_idx is None else f"{stream_idx}_{bs}" -``` +The normal `run_once` still invokes +`init_forward_metadata_in_graph(forward_batch)` inside graph capture, so +graph-recordable attention operations remain part of the graph. FA3 does not +need the FlashInfer allocation-reuse wrapper. -(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.) +After capture, Foundry writes the graph index, topology manifest, packed +fatbins, and final allocation watermark. -#### 5d. `capture` +### LOAD -The outermost replacement. +LOAD rejects a second capture request after hydration. On the first call it: -**SAVE**: +1. requires `FullCudaGraphBackend`; +2. creates or reuses the graph pool and registers it with both SGLang pool + registries; +3. bootstraps DeepEP when active; +4. pre-maps memory through SAVE's final watermark; +5. recreates per-shape attention metadata in SAVE order; +6. batch-loads every graph named by `sglang_graph_index.json`; +7. reconstructs each real `ShapeKey`; +8. fills backend `_graphs`, `_outputs`, and `_pool`; +9. initializes the DeepEP adapter for decode replay. -```python -initialize_all_attention_metadata(self) # pre-pass -attn_backend.forward_metadata = None -real_init = attn_backend.init_forward_metadata_capture_cuda_graph -attn_backend.init_forward_metadata_capture_cuda_graph = reuse_pre_pass_init -try: - result = orig_capture(self, *args, **kwargs) # upstream capture loop -finally: - attn_backend.init_forward_metadata_capture_cuda_graph = real_init -save_graph_manifest() -pack_fatbins() -capture_final_alloc_offset() -``` +## `FullCudaGraphBackend.capture_one` -`initialize_all_attention_metadata` walks `reversed(self.capture_bs)` and pre-allocates every per-bs FlashInfer wrapper. The wrappers are stored in `attn_backend.decode_cuda_graph_metadata[bs]`. +On SAVE, the replacement bypasses the backend's two eager warmup forwards. It +creates a Foundry graph, executes `forward_fn` once inside Foundry's graph +context, and stores the resulting graph/output in backend `_graphs` and +`_outputs`. -`reuse_pre_pass_init` is a drop-in replacement for the upstream inner init that runs inside `capture_one_batch_size`. For decode mode it: +Every archive record preserves: -1. Looks up the pre-pass wrapper from `decode_cuda_graph_metadata[bs]`. -2. Re-runs `indices_updater_decode.update(...)` with the same buffer slices (idempotent โ€” writes plan info to the same `_int_workspace_buffer`). -3. Sets `attn_backend.forward_metadata = DecodeMetadata(wrappers)` so the captured forward sees the right metadata for this iter. +```text +phase = decode +capture_index +filename +ShapeKey(size, stream_idx, variant_label) +output_kind = next_token_logits +``` -No allocation. The captured graph references the pre-pass wrapper's address; LOAD's pre-pass produces a wrapper at the same address. +Only `LogitsProcessorOutput.next_token_logits` is supported. Other populated +output fields fail explicitly. -**LOAD**: +LOAD never enters this capture method; attempting to do so is treated as an +incompatible archive/configuration request. -```python -if cgr.get_global_graph_memory_pool() is None: - cgr.set_global_graph_memory_pool(self.device_module.graph_pool_handle()) -set_graph_pool_id(cgr.get_global_graph_memory_pool()) -preallocate_for_load_mode() # cuMemCreate+cuMemMap up to final_alloc_offset -initialize_all_attention_metadata(self) # pre-pass (same as SAVE) -load_all_graphs(self) # ONE start_graph_builds + finish_graph_loads -self.graphs = {k: v[0] for k, v in state.loaded_graphs.items()} -self.output_buffers = {k: v[1] for k, v in state.loaded_graphs.items()} -``` +## DeepEP lifecycle -`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`. +For DP-attention + DeepEP low-latency + FA3: -### 6. Spawn-site patches +- SAVE runs one eager orchestration pass with backend capture neutralized to + initialize DeepGEMM and other lazy kernel state outside stream capture. +- SAVE and LOAD bootstrap the singleton `DeepEPBuffer` before persisted graph + capture/load. +- Graph loading initializes queued NVSHMEM modules once before native graph + builds. +- LOAD invokes `deepep_adapter.capture(is_extend_in_batch=False)` after + backend hydration. -Two parent-side wrappers: +## Process launch hooks -```python -# Engine._launch_scheduler_processes -def patched_launch(self, *args, **kwargs): - if get_graph_extension_mode() != CUDAGraphExtensionMode.NONE: - rt.setup_ld_preload_env() - return orig_launch(self, *args, **kwargs) +Before the parent launches schedulers, +`Engine._launch_scheduler_processes` calls +`setup_ld_preload_env()`. Before a DP controller launches its scheduler group, +`DataParallelController.launch_tensor_parallel_group` does the same. -# DataParallelController.launch_tensor_parallel_group -def patched_start(self, *args, **kwargs): - if get_graph_extension_mode() != CUDAGraphExtensionMode.NONE: - rt.setup_ld_preload_env() - return orig_start(self, *args, **kwargs) -``` +The helper prepends `libcuda_hook.so` and, when configured, the NVSHMEM host +library to `LD_PRELOAD`, and exports Foundry mode/timing metadata. Child +processes inherit this environment before any CUDA call. -`setup_ld_preload_env()` prepends `libcuda_hook.so` (and optionally `libnvshmem_host.so`) to `os.environ["LD_PRELOAD"]`, sets `FOUNDRY_MODE`, and records a wall-clock marker (`FOUNDRY_SPAWN_T0_NS`). All children spawned from these methods inherit the env. - -## Expert parallel (DeepEP) additions - -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 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 - captured stream (`deep_ep_cpp.Buffer(...)` โ†’ "operation not permitted when stream is - capturing"). The hook forces it before the capture loop, unwrapping the - `MaybeTboDeepEPDispatcher._inners` to reach a `DeepEPDispatcher`. Runs on SAVE and LOAD. -- **SAVE-only warmup pass** (`_run_warmup_pass`, `capture()` patch). Reuses the upstream - capture loop with graph capture neutered (run forwards only) to trigger every - pre-capture lazy init โ€” DeepGEMM per-shape JIT (`stream.synchronize()` is illegal in - capture), buffer creation, etc. โ€” outside the captured stream. LOAD doesn't need it - (preallocate + replay place allocations at recorded offsets; bootstrap covers NVSHMEM). -- **`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` - 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 - populated post-load for the replay lookup. -- **C++: bind context on graph-build pool workers** (`CUDAGraphParallel.cpp`). EP graphs - carry `NODE_EVENT_RECORD/WAIT` nodes โ†’ `cuEventCreate` during on-demand prep, which runs - on `SimpleThreadPool` workers that never `cuCtxSetCurrent(main_ctx)`. Added that call - (mirrors the bg thread). Dense graphs never hit it (no event nodes), which is why it - 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). - -## Patch idiom - -All patches use the `wrap-and-call` idiom โ€” short-circuit on `mode == NONE`, run foundry pre-work, call `orig`, run foundry post-work. The single exception is `capture()` on LOAD, which replaces the upstream method entirely (does not call `orig`). - -## 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. +Environment inheritance is not Python-state inheritance. The external +four-file shim therefore also calls `apply_server_args` inside +`run_scheduler_process` and `run_data_parallel_controller_process`; see +[`direct-edits.md`](direct-edits.md). diff --git a/docs/sglang/memory-consistency.md b/docs/sglang/memory-consistency.md index d1e21d02..d51d161c 100644 --- a/docs/sglang/memory-consistency.md +++ b/docs/sglang/memory-consistency.md @@ -1,123 +1,154 @@ -# SGLang Memory-Consistency Post-Mortem +# SGLang memory-consistency invariants -The hardest correctness problem in the SGLang integration was making LOAD produce the same output as SAVE. Cursors at every checkpoint matched, foundry's replay reported success, the graph launched cleanly โ€” and the model emitted garbage tokens. +Foundry can successfully launch a restored graph while still producing wrong +tokens if a graph-referenced buffer occupies a different VMM address on LOAD. +The integration therefore treats allocation parity as a correctness property, +not only a performance property. -This doc walks every divergence we hit and the fix that landed. The common cause: anything that allocates on one path but not the other shifts the VMM cursor or changes the address an `_int_workspace_buffer` lands at, and the captured graph kernels read addresses fixed at SAVE time. +## What Foundry validates -## The general invariant +LOAD validates: -**SAVE and LOAD must walk an identical VMM cursor trajectory.** Foundry validates `start_base_addr_X` for each captured graph; the validation only catches asymmetries up to the first divergence. Captured-kernel argument addresses (e.g. wrapper `_int_workspace_buffer.data_ptr()`) are **not** validated โ€” they just point at whatever VMM offset SAVE recorded. If LOAD's allocator lands a different buffer at the same offset, kernels read it; if no allocation lands at the offset, kernels read the preallocated zero-initialized memory. +- warmup-state schema 1; +- archived and installed SGLang version `0.5.15.post1`; +- archived and current decode backend `full`; +- archived and current prefill backend `disabled`; +- graph-index schema 1 and exact entry fields; +- contiguous capture order beginning at zero; +- every indexed graph file; +- the number of native load results; +- the final graph replay allocation trajectory through saved graph metadata. -## Bug 1 โ€” `CUDAGraph::capture_begin/end` toggling the VMM region +These checks reject known incompatible inputs, but deterministic preparation is +still required to keep graph argument addresses correct. -### Symptom +## 1. Keep the VMM region active -`replay_hook_events_from_json` errored on iter 2 with `Memory offset mismatch during replay`: current cursor 20 MB ahead of saved `start_base_addr_2`. +The deterministic allocation region remains active through metadata +preparation and the complete capture/load sequence. Per-graph recording +windows may start and stop, but they must not toggle the allocation region. +Otherwise allocations between shapes bypass the VMM on one path and shift +every later graph address. -### Root cause +## 2. Remove SAVE-only backend warmup allocations -`foundry/csrc/CUDAGraph.cpp` used to bracket every capture with `resume_allocation_region()` / `stop_allocation_region()`. So on SAVE: +Upstream `FullCudaGraphBackend.capture_one` executes `forward_fn` twice before +capturing it. Those forwards allocate and release activations, changing the +CUDA caching allocator's free-segment layout. LOAD does not run those forwards. -- iter 0 enters with region ON. bs=24 init allocates inside VMM; cursor advances. -- iter 0 `capture_end` โ†’ region OFF. -- iter 1 init runs with region OFF. Allocations bypass VMM, no cursor advance. -- iter 1 `capture_begin` โ†’ region ON; `start_base_addr_1` recorded. -- โ€ฆ +Foundry replaces this method on SAVE and executes the real forward only once, +inside the Foundry capture context. This keeps SAVE's allocator history aligned +with LOAD while still persisting every allocation made by the captured +forward. -LOAD never calls `_capture_graph`, so nothing flipped the region off. LOAD's per-bs inits all advanced the cursor โ€” past the saved `start_base_addr` values from iter 1 onward. +`BaseRunner.warmup` is also suppressed in SAVE and LOAD, and the external shim +disables FlashInfer autotuning. -### Fix +## 3. Recreate attention metadata in one order -Removed both calls in `CUDAGraph::capture_begin/end`. Only the hook recording window (`start_hook_record` / `end_hook_record`) brackets each capture. Region toggling is now each integration's responsibility, and both SGLang and vLLM keep the region ON across the whole capture loop. +SGLang v0.5.15.post1 separates attention metadata preparation: -(vLLM initially tried a Python-side `resume/stop_allocation_region` bracket around each capture to suppress the `final_alloc_offset` inflation seen on Qwen3-30B-A3B EP2 โ€” ~194 GB observed vs ~80 GB actual. That was wrong: any tensor a between-capture allocation produced that was later referenced by a captured kernel would have ended up outside the deterministic region, so LOAD's preallocation couldn't recreate it at the right address. The actual root cause was the `empty_cache()` inside `foundry/graph.py:CUDAGraph.__enter__`, which dropped torch caching-allocator segments and forced subsequent inter-capture allocations to take a fresh `cuMemAlloc_v2` path. Removing it brought `final_alloc_offset` back down to ~80 GB. The vLLM integration now also wraps the whole capture loop in `torch.cuda.use_mem_pool(graph_pool)` so SAVE-2 โ†” LOAD caching state stays symmetric for the per-graph `finish_one_graph_load` replay path.) - -## Bug 2 โ€” Pre-capture warmup forwards on SAVE only - -### Symptom - -After bug 1: iter 2 still drifted 20 MB. - -### Root cause - -`CudaGraphRunner.capture_one_batch_size` runs: - -```python -for _ in range(2): - self.device_module.synchronize() - self.model_runner.tp_group.barrier() - run_once() +```text +out-of-graph: host planning and allocation-bearing metadata +in-graph: static GPU work recorded in the graph ``` -before `_capture_graph`. These forwards allocate activation tensors and immediately free them. Torch's caching allocator keeps the freed segments as cached free blocks. Later per-bs init allocations may hit or miss cache depending on segment fragmentation history โ€” and the cache-hit pattern is not reproducible across SAVE and LOAD. - -LOAD never runs these warmups โ†’ different cache state โ†’ different cache hit/miss decisions โ†’ cursor drift. - -### Fix - -In `patched_capture_one_batch_size` wrap `forward` with a 3-call counter; first two invocations return `None`, the third (inside `_capture_graph`) runs the real forward. The `synchronize` and `barrier` calls in the loop are harmless. JIT and autotune that the warmups normally trigger now run inside the captured graph and are recorded as alloc events โ€” foundry replays them verbatim on LOAD. - -## Bug 3 โ€” Per-bs init outside graph capture - -### Symptom - -After bugs 1 and 2: cursors matched everywhere, no foundry-side error, but the model produced garbage tokens. +For FlashInfer, Foundry calls `capture_prepare(size)` and the out-of-graph +method for every size in `reversed(capture_bs)` before SAVE capture and before +LOAD graph hydration. This gives both paths the same allocation order. -### Root cause +During SAVE's normal inner orchestration, the wrapper selects the pre-created +metadata and runs the latest planner in non-capture mode. It does not create a +second wrapper. The in-graph method remains at the start of `forward_fn`, so +its GPU operations are recorded and replayed normally. -`init_forward_metadata_capture_cuda_graph(bs, ...)` is called by `capture_one_batch_size` **before** `_capture_graph`, so its allocations are not recorded in the captured graph's allocator events. Each call creates a new `BatchDecodeWithPagedKVCacheWrapper`; the wrapper's constructor allocates `_int_workspace_buffer` via `torch.empty(...)`. The captured graph kernels reference `_int_workspace_buffer.data_ptr()` directly. +This preserves both allocation addresses and the release's latest planning +semantics, including the prepared metadata's output-location binding. -For LOAD's wrapper to land at the same VMM address as SAVE's, the cursor at each iter's init time must match. With per-iter init alone, that depends on caching-allocator behavior, which is non-deterministic even after bug 2. +## 4. Load the complete indexed graph set together -### Fix (two parts) +Foundry's `graph_manifest.json` groups equal graph topologies. One member is a +template with a built executable; other members share that executable and +carry per-graph node updates. -**Both sides**: pre-pass init for every bs in `reversed(capture_bs)` order before the capture / load loop. `initialize_all_attention_metadata` walks the bs list and calls `initialize_attention_metadata_for_bs` for each. Same call sequence on both sides โ†’ same cursor trajectory โ†’ same wrapper addresses. +LOAD must submit every indexed graph path in one +`start_graph_builds(paths)` call and finish them with one +`finish_graph_loads(pending)` call. Splitting the set prevents on-demand +members from seeing their template. -**SAVE only**: the upstream `capture_one_batch_size` still calls the inner `init_forward_metadata_capture_cuda_graph(bs, ...)` per iter. Replaced it with `reuse_pre_pass_init` which: +`sglang_graph_index.json`, sorted by `capture_index`, supplies that path list. +The index also supplies the complete `ShapeKey(size, stream_idx, +variant_label)`. Graph filenames intentionally carry no shape identity. -- Looks up the existing wrapper from `decode_cuda_graph_metadata[bs]`. -- Re-runs `indices_updater_decode.update(...)` with the same buffer slices โ€” idempotent write to the same `_int_workspace_buffer`. -- Sets `attn_backend.forward_metadata = DecodeMetadata(wrappers)` for the current iter. +## 5. Mirror the pool profiler's cache drain -No allocation. Captured graph references the pre-pass wrapper's address; LOAD's identical pre-pass produces a wrapper at the same address. +SAVE's normal memory-pool resolution drains the CUDA caching allocator while +measuring available memory. LOAD skips the profiler and directly applies the +saved pool config, so it must mirror that side effect: -(Two earlier approaches didn't work: relying on torch's caching allocator to reuse a freed segment after popping the dict entry, and forcing GC. The caching allocator doesn't deterministically reuse segments when request sizes don't exactly match cached blocks. Explicit reuse via dict lookup is the only reliable path.) - -The `attn_backend.forward_metadata = None` between the pre-pass and the patch install drops the last bs's `DecodeMetadata` reference โ€” defensive, not strictly required by `reuse_pre_pass_init`. - -## Bug 4 โ€” Per-graph `start_graph_builds` broke template/on-demand linking - -### Symptom - -After bugs 1-3: cursors matched perfectly, no foundry replay error โ€” but the runtime crashed at first inference with `Called CUDAGraph::replay without a preceding successful capture or load`. +```python +torch.cuda.empty_cache() +self._apply_memory_pool_config(self.memory_pool_config) +``` -### Root cause +Without the drain, subsequent attention initialization can hit cached segments +on LOAD but request new VMM memory on SAVE, or the reverse. -`graph_manifest.json` groups captured graphs by topology. One graph per group is built as a **template** (full cuGraph + cuGraphExec instantiated); the rest are **on-demand** graphs that link to the template's `shared_exec` and patch in per-graph node updates at link time. +## 6. Transfer dense-TP VMM allocations by live file descriptor -An earlier (interleaved) LOAD called `start_graph_builds([single_path])` once per graph. Every call had 1 input โ†’ 1 topology group โ†’ graph was its own template. But the manifest marks graphs 2-N as on-demand sharing graph 0's or graph 1's template. Those templates were loaded in earlier `start_graph_builds` calls and no longer in scope โ€” so the on-demand graphs ended up with `shared_exec=null`. Runtime `CUDAGraph::replay` takes the on-demand path when `on_demand_data_` is set; with `shared_exec=null` it falls through to a `capture_ended_ || has_graph_exec_` check that errors. +The exact dense TP=2 ร— DP=2 path uses CUDA IPC between scheduler processes. +A raw exporter file-descriptor integer is meaningless in another process. +Foundry therefore exports the whole preallocated VMM chunk and transfers its +live file descriptor through a same-UID abstract Unix socket with +`SCM_RIGHTS`. The IPC handle carries exporter PID/token, chunk generation, and +the allocation's interior offset; imports use a relocated reservation if the +original VA collides. -### Fix +Qwen3-8B validation mapped DP0/TP0, DP0/TP1, DP1/TP0, and DP1/TP1 to rank +workspaces 0โ€“3. All four ranks saved and restored 36 graphs at final offset +`69822578688`; explicitly routed DP0/DP1 queries had text/token parity and LOAD +had no IPC, CUDA, or preallocation errors. -`load_all_graphs(self)` calls `start_graph_builds(all_paths)` exactly once, then `finish_graph_loads(pending)` once. All N graphs go through the manifest's template/on-demand linking in a single pass. `finish_graph_loads` then replays each graph's allocator events sequentially, walking the cursor through every `start_base_addr_X` in order. +## 7. Bootstrap DeepEP outside graph capture -## Bug 5 โ€” `_resolve_memory_pool_config`'s hidden `empty_cache` +DeepEP constructs its singleton NVSHMEM buffer lazily, and DeepGEMM may compile +shape-specific kernels lazily. Creating either inside stream capture is +invalid. -### Symptom +For the H100-validated EP=2 and EP=4 DP-attention + DeepEP low-latency + +DeepGEMM + FA3 configurations: -After bugs 1-4: cursors at `after_setup_graph_ext`, `after_init_torch_dist`, `after_scratch_skip`, `before_init_memory_pool`, and `after_init_memory_pool` all matched between SAVE and LOAD. But `save_before_pre_init` (29698 MB on SAVE) was 20 MB ahead of LOAD's `before_preallocate` (29678 MB). +- SAVE runs a dedicated eager orchestration pass to initialize lazy state; +- SAVE and LOAD bootstrap the DeepEP buffer before capture/load; +- LOAD initializes queued NVSHMEM modules before native graph builds; +- LOAD initializes the DeepEP adapter after backend hydration. -### Root cause +Qwen3-30B-A3B-FP8 validation restored 20 graphs per rank with text/token +parity, no recapture, and DeepEP/NVSHMEM/adapter/event restoration evidence. +EP=4 restored 416 event nodes and confirmed bootstrap/adapter behavior on all +four ranks. -SAVE's `init_memory_pool` calls `_resolve_memory_pool_config` โ†’ `_profile_available_bytes` โ†’ `get_available_gpu_memory(...)` (in `sglang/srt/utils/common.py`), which calls `torch.cuda.empty_cache()`. `empty_cache` releases torch caching-allocator segments back to the driver via `cuMemFree_v2`, which the foundry hook unmaps. +EP=4 uses DP-attention and expert parallelism. Its SGLang process topology has +`tp_size=4`, but it is not conventional TP attention. -LOAD's `_patch_init_memory_pool` skipped upstream profiling entirely and went directly to `_apply_memory_pool_config` โ€” no equivalent drain. Caching allocator on LOAD retained segments that SAVE released, so the subsequent attention-backend init (between `init_memory_pool` and `capture()`) took a different `cuMemAlloc` path on each side. 20 MB drift by the time `capture()` was reached. +The DeepEP token limit, chunked-prefill size, model, topology, and environment +must remain identical between SAVE and LOAD. -### Fix +## Diagnosing a mismatch -One line in `_patch_init_memory_pool`'s LOAD branch: +Compare SAVE and LOAD allocation-offset logs at: -```python -torch.cuda.empty_cache() -self._apply_memory_pool_config(self.memory_pool_config) +```text +after_setup_graph_ext +after_init_torch_dist +after_scratch_skip +before_init_memory_pool +after_init_memory_pool +before graph preallocation/preparation +after graph load ``` + +The post-load offset must equal SAVE's `final_alloc_offset`. If compatibility +validation or any offset differs, remove the archive and perform a fresh SAVE +with the exact package, flags, model, topology, and GPU environment. Do not +reuse an archive from an earlier integration. diff --git a/docs/sglang/memory-lifecycle.md b/docs/sglang/memory-lifecycle.md index 45b034e2..0652794c 100644 --- a/docs/sglang/memory-lifecycle.md +++ b/docs/sglang/memory-lifecycle.md @@ -1,119 +1,156 @@ -# SGLang Memory & Pool Lifecycle +# SGLang memory and pool lifecycle -## The invariant +## Core invariant -**SAVE and LOAD must walk an identical VMM cursor trajectory** from `setup_graph_extension` through end-of-capture / end-of-load. Anything that allocates on one path but not the other shifts the cursor by some amount. The captured graph kernels reference VMM addresses fixed at SAVE time; if LOAD's allocations don't land at those exact offsets, the buffers the kernels read are unmapped or contain stale data โ€” **silently**, with no foundry-side error. +SAVE and LOAD must allocate graph-referenced state at identical VMM offsets. +Captured kernels retain device addresses from SAVE; a different LOAD +allocation order can silently bind those addresses to the wrong buffers. -## TOML config schema +The integration controls both the allocation trajectory and the compatibility +contract for SGLang `0.5.15.post1`. + +## TOML schema ```toml -mode = "save" | "load" # required -base_addr = 0x600000000000 # required; start of VMM region -region_size = "256GB" # required; total VMM reservation -workspace_root = "foundry_archive_โ€ฆ" # required; on-disk workspace -scratch_space_size = "1024MB" # required; pre-deterministic scratch budget - -hook_library_path = "โ€ฆ/libcuda_hook.so" # optional; auto-discovered -nvshmem_host_path = "โ€ฆ/libnvshmem_host.so" # optional; for EP parity +mode = "save" # or "load" +base_addr = 0x600000000000 +region_size = "256GB" +workspace_root = "foundry_archive" +scratch_space_size = "1024MB" + +# Optional overrides; normally auto-detected. +# hook_library_path = ".../libcuda_hook.so" +# nvshmem_host_path = ".../libnvshmem_host.so.3" ``` -- `base_addr` must be high enough not to collide with model weights or PyTorch's default heap. `0x600000000000` is a safe pick on x86_64. -- `region_size` must exceed `final_alloc_offset` for the model + KV pool + graphs. 256 GiB is generous for a 14 B model. -- `scratch_space_size` should be larger than the sum of NCCL workspace, cuBLAS workspace, and any other one-time allocations that aren't deterministic across runs. -- `workspace_root` is shared across ranks; per-rank artifacts go in `workspace_root/rank_{N}/`. Delete `workspace_root` before any SAVE. +- `base_addr` and `region_size` define the deterministic VMM reservation. +- `scratch_space_size` separates nondeterministic startup allocations from + deterministic runtime allocations. +- `workspace_root` contains shared warmup state and one graph directory per + rank. +- Delete `workspace_root` before every SAVE. ## Lifecycle ```mermaid flowchart TD - INIT["ModelRunner.__init__
(upstream stashes self.dp_rank)"] - TORCH["ModelRunner.init_torch_distributed"] - TORCH_PRE["setup_graph_extension
(reserve VMM region, eager cuBLAS,
LOAD: load cached fatbins)"] - TORCH_ORIG["<upstream init_torch_distributed>
NCCL warmup + comm bring-up"] - TORCH_POST["skip_to_scratch_boundary
cursor โ†’ scratch_space_size"] - LOAD_MODEL["ModelRunner.load_model
(passthrough today)"] - POOL["ModelRunnerKVCacheMixin.init_memory_pool"] - POOL_SAVE["SAVE: orig + save resolved
MemoryPoolConfig to warmup_state.json"] - POOL_LOAD["LOAD: load MemoryPoolConfig
+ empty_cache()
+ _apply_memory_pool_config"] - KWARM["ModelRunner.kernel_warmup
(no-op on SAVE/LOAD)"] - GRAPHS["ModelRunner.init_device_graphs
โ†’ CudaGraphRunner(self).capture()"] - SAVE_CAP["SAVE: pre-pass init
+ reuse_pre_pass_init shim
+ capture loop
+ save_graph_manifest
+ pack_fatbins
+ capture_final_alloc_offset"] - LOAD_CAP["LOAD: preallocate_for_load_mode
+ pre-pass init
+ load_all_graphs"] - - INIT --> TORCH - TORCH --> TORCH_PRE --> TORCH_ORIG --> TORCH_POST - TORCH_POST --> LOAD_MODEL --> POOL - POOL --> POOL_SAVE - POOL --> POOL_LOAD - POOL_SAVE --> KWARM - POOL_LOAD --> KWARM - KWARM --> GRAPHS - GRAPHS --> SAVE_CAP - GRAPHS --> LOAD_CAP + SHIM["resolve full decode / disabled prefill"] + DIST["init_torch_distributed wrapper"] + VMM["bind rank device
reserve VMM
load packed modules on LOAD"] + SCRATCH["upstream distributed init
cursor to scratch boundary"] + MODEL["load model"] + POOL["init memory pool"] + POOLS["SAVE: resolve and record pool"] + POOLL["LOAD: validate archive
empty cache and apply saved pool"] + WARM["BaseRunner.warmup suppressed"] + DECODE["DecodeCudaGraphRunner.capture"] + SAVEC["SAVE: metadata prep
capture and index graphs"] + LOADC["LOAD: pre-map span
metadata prep and hydrate backend"] + + SHIM --> DIST --> VMM --> SCRATCH --> MODEL --> POOL + POOL --> POOLS --> WARM + POOL --> POOLL --> WARM + WARM --> DECODE + DECODE --> SAVEC + DECODE --> LOADC ``` -## Allocation buckets +## Allocation regions -### Bucket A โ€” pre-deterministic scratch +### Startup scratch -CUDA context, cuBLAS handle, NCCL warmup, distributed bring-up. These run **inside** the VMM region but are released by `skip_to_scratch_boundary` โ€” the cursor is forced to `cfg.scratch_space_size` regardless of what landed below it. Any non-determinism in this region is invisible to the rest of the lifecycle. +CUDA context creation, cuBLAS, NCCL/communicator setup, and related one-time +work happen before the scratch boundary. After upstream distributed +initialization, Foundry sets the cursor to `scratch_space_size` regardless of +the exact amount consumed below it. -### Bucket B โ€” deterministic runtime state +For multi-GPU operation, the hook sets `self.gpu_id` as current before VMM +reservation. Otherwise rank-local allocations can bind to GPU 0. -Model weights, KV pools, attention workspace buffers, FlashInfer wrapper `_int_workspace_buffer`s, captured-graph alloc events. These must allocate in identical order on SAVE and LOAD. Their final cursor position is the `final_alloc_offset` persisted to `warmup_state.json`. +### Deterministic runtime state -### Bucket C โ€” forbidden divergence +After the scratch boundary, model weights, KV pools, attention workspaces, +captured graph allocations, and output buffers must follow the same order. +SAVE records the terminal cursor as `final_alloc_offset`. -Anything that runs on one path but not the other. The fixes we landed identify and align the three known cases: +LOAD calls `preallocate_region(final_alloc_offset - current_offset)` before +native graph loading. This maps physical memory but does not advance the +logical cursor; replayed allocation events advance it in SAVE order. -1. The two pre-capture warmup forwards in `capture_one_batch_size` โ€” skipped on SAVE so they don't pollute the caching allocator with freed activations LOAD can't reproduce. (See doc 06.) -2. The per-iter inner `init_forward_metadata_capture_cuda_graph(bs)` call โ€” replaced on SAVE with `reuse_pre_pass_init` so it doesn't re-allocate the wrappers the pre-pass already built. (See doc 03 / doc 05.) -3. `_resolve_memory_pool_config` calls `get_available_gpu_memory(empty_cache=True)` on SAVE; LOAD's `_patch_init_memory_pool` mirrors it with an explicit `torch.cuda.empty_cache()` before `_apply_memory_pool_config`. (See below.) +### State that must stay eager -## The `_resolve_memory_pool_config` mirror +DeepEP/NVSHMEM buffer construction and DeepGEMM lazy initialization cannot +start inside stream capture. SAVE's dedicated eager bootstrap performs that +work before persisted graph capture. LOAD initializes its adapter and queued +NVSHMEM module state before replay. -`init_memory_pool` on SAVE calls `_resolve_memory_pool_config` โ†’ `_profile_available_bytes` โ†’ `get_available_gpu_memory(...)` which calls `torch.cuda.empty_cache()` (`sglang/srt/utils/common.py`). `empty_cache` releases torch caching-allocator segments back to the driver via `cuMemFree_v2`, which the foundry hook unmaps. +## Memory-pool parity -LOAD skips upstream `init_memory_pool` and goes straight to `_apply_memory_pool_config`, so it had no equivalent drain. Caching-allocator segments retained on LOAD changed which torch.empty calls hit cache vs. miss during the subsequent attention-backend init โ€” pushing the cursor 20 MB behind SAVE by the time `capture()` was reached. +SAVE resolves SGLang's `MemoryPoolConfig` through normal profiling and records +the complete dataclass dictionary in `warmup_state.json`. -The fix is one line: +LOAD validates archive compatibility before changing pool state, then: ```python -torch.cuda.empty_cache() # mirror SAVE's _resolve_memory_pool_config side effect +self.memory_pool_config = MemoryPoolConfig(**state.memory_pool_config) +torch.cuda.empty_cache() self._apply_memory_pool_config(self.memory_pool_config) ``` -## Why saving raw memory numbers isn't enough +The cache drain mirrors a side effect of SAVE's memory profiling. Omitting it +changes whether later allocations reuse cached segments or advance the VMM +cursor. + +Persisting a raw free-memory count would be insufficient because SGLang also +applies model constraints, page alignment, request limits, and topology when +resolving the final pool. + +## Attention metadata parity -The pool config is computed from a profile and a configurator: +The release separates attention preparation into: -1. `_profile_available_bytes(pre_model_load_memory)` -2. `MemoryPoolConfigurator.calculate_pool_sizes(...)` -3. `_apply_token_constraints(...)` (user caps, page alignment, PP sync) -4. `_resolve_max_num_reqs(...)` -5. `_apply_memory_pool_config(...)` +- `init_forward_metadata_out_graph`, for host planning and allocation-bearing + per-iteration state; +- `init_forward_metadata_in_graph`, for static GPU operations recorded into + the CUDA graph. -Persisting the resolved `MemoryPoolConfig` (via `dataclasses.asdict`) and re-applying it on LOAD is necessary and sufficient. Persisting raw free-memory numbers would re-run the configurator on LOAD with stale inputs. +For FlashInfer, Foundry runs the out-of-graph allocation pass for every decode +size in reverse capture order on both SAVE and LOAD. During SAVE capture, the +inner planner reuses the prepared object rather than allocating a second one. +The in-graph method remains inside `forward_fn`. -`WarmupState.memory_pool_config` is exactly this dict. +FA3 uses its initialized graph state without the FlashInfer reuse wrapper. -## VMM region setup +## Graph-pool ownership -`setup_graph_extension` (in `runtime.py`): +`FullCudaGraphBackend` owns: -- Computes the per-rank workspace path (`{workspace_root}/rank_{compute_workspace_rank(...)}`) -- On SAVE: removes the rank workspace if it exists, then creates it fresh -- On LOAD: - - `cge.set_skip_fatbin_processing(True)` - - `cge.load_cuda_modules_and_libraries(workspace_dir)` โ€” restores the fatbins SAVE wrote, into device code memory -- `cge.set_allocation_region(cfg.base_addr, parse_size(cfg.region_size))` โ€” reserves a VMM address range; the cursor starts at `base_addr` -- `torch._C._cuda_getCurrentBlasHandle()` โ€” eager cuBLAS handle so the workspace it wants lands in scratch instead of post-scratch territory -- creates the `CUDAGraphExtensionState` singleton +- `_graphs`: graph objects keyed by complete `ShapeKey`; +- `_outputs`: reconstructed `LogitsProcessorOutput` objects; +- `_pool`: the shared graph memory-pool handle. -After upstream `init_torch_distributed` returns, `skip_to_scratch_boundary` forces the cursor to `cfg.scratch_space_size` (default 1 GiB). Allocations below that line are scratch and don't need to be deterministic. +LOAD registers `_pool` with SGLang's process-wide graph-pool and PyNCCL pool +registries before restoring graphs, then hydrates the backend maps from the +authoritative graph index. -## Final watermark +## Persisted compatibility state -`capture_final_alloc_offset` runs after the SAVE-side capture loop completes (after `save_graph_manifest` and `pack_fatbins`). It writes `final_alloc_offset` to both `rank_{N}/final_alloc_offset.json` and the shared `warmup_state.json`. +Schema-1 `warmup_state.json` records: + +```text +schema_version +sglang_version +decode_cuda_graph_backend +prefill_cuda_graph_backend +timestamp +cuda_version +gpu_name +gpu_total_memory +memory_pool_config +``` -`preallocate_for_load_mode` uses this to call `cge.preallocate_region(final - current)` so the entire deterministic range is mapped to physical memory in one shot. The cursor is **not** advanced โ€” preallocate just pre-maps; the cursor advances naturally as cuMemAllocs land within the preallocated range (fast-path: pointer bump, no driver calls). +LOAD requires archive schema 1, SGLang `0.5.15.post1`, decode `full`, prefill +`disabled`, matching installed package metadata, and matching current resolved +backends. A mismatch requires a fresh SAVE; no archive migration is attempted. +The VMM watermark is rank-specific and is persisted separately in +`rank_/final_alloc_offset.json`. diff --git a/docs/sglang/overview.md b/docs/sglang/overview.md index f8ea7cce..316bd000 100644 --- a/docs/sglang/overview.md +++ b/docs/sglang/overview.md @@ -1,112 +1,137 @@ -# SGLang Integration Overview - -Foundry persists SGLang's `CudaGraphRunner` graphs to disk on SAVE and restores them on LOAD, skipping graph capture, kernel warmup, and the per-batch-size attention metadata setup costs. - -Tested on single-GPU Qwen3-1.7B / 4B / 14B and **data-parallel (DP=2)** Qwen3-1.7B with the FlashInfer attention backend. - -## Parallelism - -| Mode | Status | Notes | -|---|:---:|---| -| Single GPU | โœ… | Qwen3-1.7B / 4B / 14B | -| Data parallel (DP) | โœ… | One full replica per rank; validated DP=2. Requires the per-rank device binding (below) and `NCCL_CUMEM_ENABLE=0` / `NCCL_NVLS_ENABLE=0`. | -| Tensor parallel (TP) | ๐Ÿšง | Deterministic NCCL memory layout is under construction. | -| Expert parallel (DeepEP) | โœ… | Validated EP=2 on Qwen3-30B-A3B-FP8 (SAVE/SAVE2/LOAD/query); restored decode graphs match baseline throughput. See **Expert parallel** below. | - -**Expert parallel (DeepEP).** EP runs DP-attention (each rank its own attention โ€” no -NCCL all-reduce) + DeepEP for the MoE all-to-all (NVSHMEM, foundry-compatible). The -serve script is `recipe/sglang/serve_qwen3-30ba3bfp8_ep.sh - [--save|--load]` with: `--enable-dp-attention --moe-a2a-backend deepep ---deepep-mode low_latency --moe-runner-backend deep_gemm --attention-backend fa3 ---disable-custom-all-reduce`. Required kernel builds in the env: `deep_ep` at sglang's -pinned commit (`9af0e0d`, not vLLM's), `sgl-deep-gemm>=0.1.2` (0.1.0 lacks -`m_grouped_bf16_gemm_nt_masked`), `flash-attn-3`. `fa3` is required because the -flashinfer ragged-prefill path has an off-by-one (`q.shape != qo_indptr`) under this -config. DeepEP low-latency caps dispatch at -`SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK` (default 128); raise it (+ chunk -prefill) for larger batches and keep it identical across SAVE/LOAD. Foundry-specific -EP handling is in [`hooks.md`](hooks.md): a DeepEP buffer pre-capture bootstrap, a -SAVE-only warmup pass (triggers DeepGEMM JIT + buffer creation outside the capture -stream), `deepep_adapter` mode init on LOAD, the FlashInfer pre-pass gated off for -fa3, and a C++ fix binding the CUDA context on the graph-build pool workers. - -**Per-rank device binding (DP/TP/EP).** Foundry's `set_allocation_region` binds the -VMM region to the CUDA device current at call time. Upstream sets the device -*inside* `init_torch_distributed`, which the integration wraps and front-runs, so -the hook explicitly calls `set_device(self.gpu_id)` before reserving the region โ€” -otherwise rank > 0 reserves on `cuda:0` and faults. See [`hooks.md`](hooks.md) ยง1. -The DP serve script lives at `recipe/sglang/serve_qwen3-1.7b_dp.sh` -(` [--save|--load]`); pick GPUs with `CUDA_VISIBLE_DEVICES`. - -## How to use - -### 1. Install foundry +# SGLang integration overview + +Foundry supports exactly SGLang `0.5.15.post1`, upstream commit `0b3bb0c`. +Full decode CUDA graph SAVE and fresh-process LOAD are validated on H100 with +CUDA 13.0 and PyTorch 2.11/cu130. Prefill CUDA graphs are disabled. + +## Validated configurations + +| Configuration | Attention/MoE backend | Status | +|---|---|:---:| +| Single GPU | FlashInfer | Validated on H100 | +| Regular data parallelism | One full replica per rank, FlashInfer | Validated at DP=2 on H100 | +| Dense tensor ร— data parallelism | FlashInfer | Validated only at TP=2 ร— DP=2, PP=1, EP=1 on H100 | +| DP-attention expert parallelism | DeepEP low-latency, DeepGEMM, and FA3 | Validated at EP=2 and EP=4 on H100 | + +Every conventional TP topology except exact dense TP=2 ร— DP=2 with PP=1 and +EP=1 is rejected at startup. The validated dense topology also requires custom +all-reduce off and `NCCL_CUMEM_ENABLE=0`/`NCCL_NVLS_ENABLE=0`. Pipeline +parallelism, speculative decoding, LoRA/PDMux, dLLM, hidden-state capture, and +memory-saver graphs remain unsupported. Foundry does not claim prefill-graph +support. + +Validation results: + +- Single-GPU Qwen3-1.7B passed baseline, SAVE, fresh LOAD, and query with + text/token parity, no recapture, and 52 restored graphs. +- Regular DP=2 Qwen3-1.7B passed on both FlashInfer ranks with text/token + parity, no recapture, and 52 restored graphs per rank. +- Dense TP=2 ร— DP=2 Qwen3-8B passed on four H100s with explicitly routed DP0 + and DP1 parity, no recapture, rank workspaces 0โ€“3, 36 restored graphs per + rank, and identical final offsets of `69822578688`. No VMM-IPC, CUDA, or + preallocation errors occurred. +- EP=2 and EP=4 Qwen3-30B-A3B-FP8 passed with DP-attention, DeepEP + low-latency, DeepGEMM, and FA3, with text/token parity, no recapture, and 20 + restored graphs per rank. EP=4 restored 416 event nodes and showed + four-rank DeepEP/NVSHMEM bootstrap and adapter evidence. + +EP=4 is not conventional TP attention. SGLang sets `tp_size=dp_size=ep_size=4` +for this DP-attention/EP process topology, while each rank owns its attention +work and DeepEP handles expert all-to-all. + +## Upstream execution model + +This SGLang release separates orchestration from graph storage: + +- `BaseRunner.warmup` owns run-once kernel warmup and autotuning. +- `DecodeCudaGraphRunner.capture` prepares decode shapes and orchestrates + capture. +- `FullCudaGraphBackend.capture_one` owns one full-model graph capture. +- `ShapeKey(size, stream_idx, variant_label)` is the complete graph identity. +- `FullCudaGraphBackend._graphs`, `_outputs`, and `_pool` own restored graph + state. + +Attention metadata is also split. Per-iteration host work and dynamic planning +run in `init_forward_metadata_out_graph`; graph-recordable static GPU work runs +in `init_forward_metadata_in_graph`. + +Foundry patches these latest seams rather than rebuilding SGLang's runner. + +## SAVE lifecycle + +1. The shim resolves graph settings to decode `full` and prefill `disabled`, + then installs Foundry hooks. +2. Before distributed initialization, Foundry binds the rank's CUDA device, + reserves the deterministic VMM region, restores its scratch boundary, and + initializes the rank workspace. +3. SGLang allocates model weights, KV pools, and attention state inside that + region. Foundry records the resolved `MemoryPoolConfig` and exact package + and graph compatibility fields in `warmup_state.json`. +4. Foundry suppresses `BaseRunner.warmup`. +5. For FlashInfer, Foundry prepares every decode shape in reverse capture + order. The allocation-bearing out-of-graph metadata pass runs once per + shape; the normal inner path reuses that object while retaining current + planning behavior. +6. The replacement for `FullCudaGraphBackend.capture_one` skips its two eager + forwards, captures one real forward with Foundry, and records the complete + `ShapeKey`. +7. SAVE writes schema-1 `sglang_graph_index.json`, the Foundry topology + manifest, packed device code, and the final VMM allocation watermark. + +For DeepEP, SAVE first runs one eager capture-orchestration pass with native +graph capture neutralized. This initializes DeepGEMM and other lazy state +outside stream capture. Foundry also bootstraps the singleton DeepEP/NVSHMEM +buffer before the persisted capture. + +## LOAD lifecycle + +1. A fresh process installs hooks in the parent, DP controller when present, + and every scheduler child. +2. Before applying saved pool sizing, LOAD validates warmup schema 1, archived + and installed SGLang version `0.5.15.post1`, archived decode/prefill + backends, and the current resolved backends. +3. Foundry restores packed device code, recreates model/KV state at + deterministic VMM offsets, applies the saved `MemoryPoolConfig`, and + pre-maps memory through the final SAVE watermark. +4. The same reverse-order attention metadata preparation recreates + allocation-bearing objects. +5. LOAD reads schema-1 `sglang_graph_index.json` in `capture_index` order and + submits all indexed graph paths through one + `start_graph_builds`/`finish_graph_loads` pair. This preserves the + topology manifest's template and on-demand relationships. +6. Each index record becomes a real + `ShapeKey(size, stream_idx, variant_label)`. Foundry fills backend + `_graphs` and `_outputs`, restores `_pool`, and initializes the DeepEP + adapter before serving. + +The graph index, not filenames, is authoritative for phase, capture order, +shape identity, and output kind. + +## Archive compatibility + +This integration is latest-only. Archives from earlier SGLang versions or +earlier Foundry formats are not migrated. Delete the workspace and perform a +fresh SAVE: ```bash -pushd foundry && pip install -e . --no-build-isolation && popd -``` - -An editable install (`pip install -e .`) is the supported path โ€” the recipe serve scripts set no `PYTHONPATH` and rely on `foundry` being importable. To run straight from a source checkout, export `PYTHONPATH=โ€ฆ/foundry/python:โ€ฆ/sglang/python` yourself. - -### 2. Write a TOML config - -```toml -# recipe/sglang/foundry_save.toml -mode = "save" -base_addr = 0x600000000000 -region_size = "256GB" -workspace_root = "foundry_archive" -scratch_space_size = "1024MB" -``` - -The matching `foundry_load.toml` just changes `mode = "load"`. See [`memory-lifecycle.md`](memory-lifecycle.md) for what each field controls. - -### 3. Run SAVE, then LOAD - -```bash -# SAVE rm -rf foundry_archive -bash recipe/sglang/serve_qwen3-mini.sh --save -# Wait for "Application startup complete", then SIGTERM the server. - -# LOAD -bash recipe/sglang/serve_qwen3-mini.sh --load - -# Query -curl -s http://0.0.0.0:12000/v1/completions \ - -H 'Content-Type: application/json' \ - -d '{"model":"Qwen/Qwen3-1.7B","prompt":"The capital of France is", - "max_tokens":12,"temperature":0}' -# โ†’ "Paris. The capital of the United States is Washington, D" ``` -A single SAVE pass is enough โ€” SGLang doesn't run a profile-forward at startup, so there is no non-determinism that requires a second pass. - -## What the integration does - -SAVE: - -1. `setup_graph_extension(...)` reserves a VMM region at `base_addr` and creates the per-rank workspace. -2. Distributed init / NCCL warmup runs in scratch space; the cursor is then forced to `scratch_space_size`. -3. Model weights, KV pool, and FlashInfer workspace buffers allocate inside the VMM region at byte-deterministic offsets. -4. `kernel_warmup` is a no-op. -5. `CudaGraphRunner.capture` runs a pre-pass that pre-allocates every per-bs FlashInfer wrapper, then enters the upstream capture loop with an idempotent inner-init shim (`reuse_pre_pass_init`) and a wrapper on `forward` that suppresses the two pre-capture warmup forwards. -6. Each captured graph is written to disk; a manifest groups topologically equivalent graphs. -7. The final VMM cursor is recorded as `final_alloc_offset`. - -LOAD: +SGLang must be installed as a distribution whose metadata reports exactly +`0.5.15.post1`. Apply the four-file tracked shim, then install with +`SETUPTOOLS_SCM_PRETEND_VERSION=0.5.15.post1 pip install -e sglang/python`. +The override is required because `setuptools_scm` otherwise derives a +dirty/local version from the patched tag. A checkout made importable only +through a source path is not supported. -1. `setup_graph_extension(...)` restores the VMM region and replays captured fatbins into device code memory. -2. Distributed init runs as usual; the cursor advances to the same `scratch_space_size`. -3. Model weights and KV pool re-allocate at the same deterministic offsets. `init_memory_pool` reuses the saved `MemoryPoolConfig` (and calls `torch.cuda.empty_cache()` to mirror SAVE's `_resolve_memory_pool_config` side effect). -4. `CudaGraphRunner.capture` is replaced with: preallocate the entire deterministic range up to `final_alloc_offset`; run the same pre-pass init; call `start_graph_builds(all_paths) + finish_graph_loads(pending)` exactly once. All N graphs are loaded in one shot so the manifest's template/on-demand linking works. -5. `self.graphs` / `self.output_buffers` are populated from `state.loaded_graphs`; the rest of SGLang's serving path runs unchanged. +## Run it -## Doc set +The recipe includes baseline, SAVE, fresh-process LOAD, and query commands for +all validated configurations: -- [`overview.md`](overview.md) โ€” this file -- [`direct-edits.md`](direct-edits.md) โ€” every line that changed in `sglang/` -- [`hooks.md`](hooks.md) โ€” every monkey-patch foundry installs and what it does on SAVE / LOAD -- [`memory-lifecycle.md`](memory-lifecycle.md) โ€” VMM region setup, the SAVEโ†”LOAD allocation parity contract, and the `final_alloc_offset` watermark -- [`save-load-workflow.md`](save-load-workflow.md) โ€” serve scripts, TOML schema, expected logs, validation checks -- [`memory-consistency.md`](memory-consistency.md) โ€” the five known divergences that caused silent LOAD failures and how each was fixed +- [`../../recipe/sglang/README.md`](../../recipe/sglang/README.md) +- [`direct-edits.md`](direct-edits.md) for the external four-file shim +- [`hooks.md`](hooks.md) for every runtime patch +- [`memory-lifecycle.md`](memory-lifecycle.md) for deterministic allocation +- [`memory-consistency.md`](memory-consistency.md) for parity invariants +- [`save-load-workflow.md`](save-load-workflow.md) for archive and log checks diff --git a/docs/sglang/save-load-workflow.md b/docs/sglang/save-load-workflow.md index 78be4b49..24a71220 100644 --- a/docs/sglang/save-load-workflow.md +++ b/docs/sglang/save-load-workflow.md @@ -1,35 +1,37 @@ -# SGLang Save / Load Workflow +# SGLang SAVE/LOAD workflow -## Serve scripts +This workflow targets SGLang `0.5.15.post1` only. Install the SGLang checkout +as a package, apply the external four-file shim, and delete every older archive +before SAVE. -Key knobs (1.7B example): +## Recipe flags -```bash -MODEL_NAME="Qwen/Qwen3-1.7B" -HOST="0.0.0.0" -PORT=12000 -MEM_FRACTION_STATIC=0.6 - -# No LD_PRELOAD / PYTHONPATH here: foundry + the sglang fork are pip-installed, and -# foundry's setup_ld_preload_env auto-detects libcuda_hook.so and LD_PRELOADs it -# into every worker at spawn time. (Running from a source checkout instead? Export -# PYTHONPATH=.../foundry/python:.../sglang/python yourself.) +The single-GPU recipe expands to: +```bash sglang serve \ - --model-path "$MODEL_NAME" \ - --trust-remote-code \ - --host "$HOST" --port "$PORT" \ - --tp-size 1 \ - --mem-fraction-static "$MEM_FRACTION_STATIC" \ - --disable-radix-cache \ - --attention-backend flashinfer \ - --cuda-graph-max-bs 512 \ - --foundry-graph-extension-config-path "$FOUNDRY_TOML" + --model-path Qwen/Qwen3-1.7B \ + --trust-remote-code \ + --host 0.0.0.0 \ + --port 12000 \ + --tp-size 1 \ + --mem-fraction-static 0.6 \ + --disable-radix-cache \ + --attention-backend flashinfer \ + --cuda-graph-max-bs-decode 512 \ + --disable-prefill-cuda-graph \ + --foundry-graph-extension-config-path "$FOUNDRY_TOML" ``` -`--cuda-graph-max-bs 512` is the closest analogue to vLLM's `--max-num-seqs 512` โ€” it drives `capture_bs` to span a similar range of decode batch sizes (52 batch sizes from 1 โ†’ 512). +`--cuda-graph-max-bs-decode` is the release's phase-specific CLI field. +`--disable-prefill-cuda-graph` ensures that only decode graphs are created. +The recipes use both flags in baseline, SAVE, and LOAD because they define the +selected full-decode configuration validated on H100. + +The Foundry option is present only for SAVE/LOAD. It is generated from the +external shim's annotated `ServerArgs` field. -## TOML configs +## TOML files `recipe/sglang/foundry_save.toml`: @@ -41,106 +43,173 @@ workspace_root = "foundry_archive" scratch_space_size = "1024MB" ``` -`foundry_load.toml` is identical except `mode = "load"`. See [`memory-lifecycle.md`](memory-lifecycle.md) for the field semantics. +`foundry_load.toml` uses the same values with `mode = "load"`. -## Run sequence +The exact dense TP=2 ร— DP=2 recipe uses dedicated +`foundry_tp2_dp2_save.toml`/`foundry_tp2_dp2_load.toml` files and workspace +`foundry_archive_sglang_qwen3_8b_tp2_dp2`. Its launcher fixes PP=1, EP=1, +FlashInfer, full decode, disabled prefill, disabled custom all-reduce, and +`NCCL_CUMEM_ENABLE=0`/`NCCL_NVLS_ENABLE=0` for baseline, SAVE, and LOAD. + +## Verify package metadata ```bash -cd /path/to/foundry/recipe/sglang +SETUPTOOLS_SCM_PRETEND_VERSION=0.5.15.post1 \ + pip install -e sglang/python \ + --extra-index-url https://download.pytorch.org/whl/cu130 +python -c 'from importlib.metadata import version; assert version("sglang") == "0.5.15.post1"' +``` -# 1. Clean any prior workspace -rm -rf foundry_archive +The required tracked-file shim makes the tagged checkout dirty. SGLang uses +`setuptools_scm`, so the pretend-version variable is required to make an +editable install's distribution metadata deterministically equal +`0.5.15.post1`. A source-only checkout does not satisfy the gate. SAVE refuses +to create an archive when package metadata is missing or differs from that +exact version. -# 2. SAVE +## Baseline to fresh-process LOAD + +Run each server as a separate process: + +```bash +cd recipe/sglang + +# Baseline +bash serve_qwen3-mini.sh +# Wait for startup, query if desired, then stop it. + +# SAVE +rm -rf foundry_archive bash serve_qwen3-mini.sh --save -# wait for "Application startup complete", then Ctrl-C / SIGTERM +# Wait for "Application startup complete", then stop it. -# 3. LOAD +# LOAD in a fresh process bash serve_qwen3-mini.sh --load -# wait for "Application startup complete" -# 4. Query -curl -s http://0.0.0.0:12000/v1/completions \ +# Query from another shell +curl -s http://127.0.0.1:12000/v1/completions \ -H 'Content-Type: application/json' \ - -d '{"model":"Qwen/Qwen3-1.7B","prompt":"The capital of France is", - "max_tokens":12,"temperature":0}' + -d '{"model":"Qwen/Qwen3-1.7B","prompt":"The capital of France is","max_tokens":12,"temperature":0}' ``` -For Qwen3-1.7B, expected output: +Use the DP and DeepEP commands in +[`../../recipe/sglang/README.md`](../../recipe/sglang/README.md) for the other +validated configurations. -``` -" Paris. The capital of the United States is Washington, D" -``` +Validation used upstream SGLang `v0.5.15.post1` (`0b3bb0c`), CUDA 13.0, +PyTorch 2.11/cu130, and H100. Single-GPU and DP=2 Qwen3-1.7B FlashInfer +restored 52 graphs per applicable rank. Dense TP=2 ร— DP=2 Qwen3-8B restored +36 graphs per rank across rank workspaces 0โ€“3 at identical final offsets of +`69822578688`; explicitly routed DP0 and DP1 queries matched and no IPC, CUDA, +or preallocation errors occurred. EP=2 and EP=4 Qwen3-30B-A3B-FP8 with +DP-attention, DeepEP low-latency, DeepGEMM, and FA3 restored 20 graphs per +rank; EP=4 restored 416 event nodes with four-rank NVSHMEM/bootstrap/adapter +evidence. Baseline/SAVE/fresh LOAD/query passed with text/token parity and no +recapture in every configuration. -A single SAVE pass is sufficient. SGLang does not run a profile-forward at startup (only `_profile_available_bytes` samples free GPU memory), so there is no allocation non-determinism that a second SAVE pass would need to mask. +Although SGLang sets `tp_size=4` for EP=4, that mode uses DP-attention and +expert parallelism rather than conventional TP attention. Conventional TP is +validated only for exact dense TP=2 ร— DP=2 with PP=1 and EP=1. -## Per-rank workspace layout +## Archive files -``` +```text foundry_archive/ - warmup_state.json # shared; MemoryPoolConfig + final_alloc_offset - rank_0/ - graph_{0..N-1}_FULL_t{bs}_r{bs}_UX_pcN.json # one per captured graph - graph_{0..N-1}_FULL_t{bs}_r{bs}_UX_pcN.cugraph # binary cuGraph blob - graph_manifest.json # topology groups for template + on-demand linking - final_alloc_offset.json # per-rank VMM watermark - fatbin_image_packed.img # packed kernel fatbins - fatbin_entrypoint_packed.txt # fatbin entry-point index +โ”œโ”€โ”€ warmup_state.json +โ””โ”€โ”€ rank_0/ + โ”œโ”€โ”€ sglang_graph_index.json + โ”œโ”€โ”€ graph_000000_sglang.json + โ”œโ”€โ”€ graph_000000_sglang.cugraph + โ”œโ”€โ”€ ... + โ”œโ”€โ”€ graph_manifest.json + โ”œโ”€โ”€ final_alloc_offset.json + โ”œโ”€โ”€ fatbin_image_packed.img + โ””โ”€โ”€ fatbin_entrypoint_packed.txt ``` -Filename scheme (defined in `graph_ops._graph_filename`): - +Each additional DP/EP rank has its own `rank_/`. + +`warmup_state.json` includes compatibility fields: + +```json +{ + "schema_version": 1, + "sglang_version": "0.5.15.post1", + "decode_cuda_graph_backend": "full", + "prefill_cuda_graph_backend": "disabled", + "timestamp": "2026-07-23T00:00:00", + "cuda_version": "13.0", + "gpu_name": "NVIDIA H100", + "gpu_total_memory": 0, + "memory_pool_config": {} +} ``` -graph_{state.capture_index}_FULL_t{bs}_r{bs}_UX_pcN.json + +The timestamp, GPU values, and pool dictionary shown above are illustrative; +SAVE records the actual values. Each rank's VMM watermark lives in its +`rank_/final_alloc_offset.json`. + +`sglang_graph_index.json` is authoritative: + +```json +{ + "schema_version": 1, + "graphs": [ + { + "capture_index": 0, + "filename": "graph_000000_sglang.json", + "output_kind": "next_token_logits", + "phase": "decode", + "shape_key": { + "size": 512, + "stream_idx": null, + "variant_label": null + } + } + ] +} ``` -`state.capture_index` increments per `save_graph` call so files sort in SAVE-time order. `_GRAPH_FILENAME_RE` in `graph_ops.py` parses them on LOAD. +LOAD validates exact fields and uses `capture_index`, not lexical filename +order. `graph_manifest.json` separately records native topology groups, +templates, and on-demand members. -## Expected logs +## Expected checkpoints -SAVE (success): +SAVE logs include: -``` -[Foundry] SGLang hooks installing: mode=save workspace=foundry_archive_qwen_1.7b -[Foundry] SGLang hooks installed -Foundry SGLang integration activated from .../save_qwen_1.7b.toml -[Foundry] SGLang rank=0 workspace_dir=foundry_archive_qwen_1.7b/rank_0 -[Foundry] SGLang graph extension setup completed in 0.x s -[Foundry] SGLang skipped allocator to scratch boundary 1073741824 -[Foundry] SGLang reused saved memory pool config โ† only on LOAD -[Foundry] SGLang kernel_warmup skipped in save mode -โ€ฆ -[Foundry] Saved SGLang CUDA graph graph_0_FULL_t512_r512_UX_pcN.json key=512 -โ€ฆ -[Foundry] Saved SGLang CUDA graph graph_51_FULL_t1_r1_UX_pcN.json key=1 -[foundry] Saved graph_manifest.json: 9 topology groups (9 templates, 43 on-demand, 43 stripped) -[Foundry] Saved SGLang warmup state to foundry_archive_qwen_1.7b/warmup_state.json -[Foundry] SGLang final_alloc_offset=22785556480 -โ€ฆ -INFO: Application startup complete. +```text +[Foundry] SGLang hooks installing: mode=save ... +[Foundry] SGLang skipped allocator to scratch boundary ... +[Foundry] SGLang runner warmup skipped in save mode +[Foundry] Saved SGLang CUDA graph graph_000000_sglang.json shape_key=ShapeKey(...) +[Foundry] Saved SGLang warmup state to .../warmup_state.json +[Foundry] SGLang final_alloc_offset=... +INFO: Application startup complete. ``` -LOAD (success): +LOAD logs include: -``` -[Foundry] SGLang hooks installing: mode=load workspace=foundry_archive_qwen_1.7b -[Foundry] SGLang hooks installed -[Foundry] SGLang graph extension setup completed in 0.x s +```text +[Foundry] SGLang hooks installing: mode=load ... [Foundry] SGLang reused saved memory pool config -[Foundry] SGLang kernel_warmup skipped in load mode -[Foundry] SGLang alloc_offset[before_preallocate]=โ€ฆ (โ€ฆ MB) -[Foundry] SGLang alloc_offset[after_preallocate]=โ€ฆ (โ€ฆ MB) -[Foundry] SGLang alloc_offset[after_pre_init]=โ€ฆ (โ€ฆ MB) -[CGE] Using graph_manifest.json (9 topology groups) -[CGE] Phase 1: 52 graphs parsed in 0.x ms, 9 topologies, 4 threads, 52 binary + 0 json -[CGE BUILD] Template 0 (...): N nodes, done in X.X ms -โ€ฆ -[CGE] Phase 2: 9 templates + 43 on-demand = 52 graphs built in xx.x ms -[CGE] finish_graph_loads: 52 graphs, xx.x ms -[Foundry] Loaded 52 SGLang graphs in 0.0x s -[Foundry] SGLang alloc_offset[after_load_all_graphs]=22785556480 (โ€ฆ MB) -โ€ฆ -INFO: Application startup complete. +[Foundry] SGLang runner warmup skipped in load mode +[CGE] Using graph_manifest.json ... +[CGE] finish_graph_loads: ... +[Foundry] Loaded ... SGLang graphs in ...s +INFO: Application startup complete. ``` -The `after_load_all_graphs` value **must** equal SAVE's `final_alloc_offset`. If it doesn't, see [`memory-consistency.md`](memory-consistency.md). +Compatibility validation occurs before the saved pool config is applied. +Missing index files, legacy schemas, version mismatches, backend mismatches, or +unsupported output/shape records stop LOAD and request a fresh SAVE. + +## Migration rule + +There is no archive migration path. When changing SGLang, Foundry, graph +flags, model, topology, or relevant kernel packages: + +```bash +rm -rf foundry_archive +bash serve_qwen3-mini.sh --save +``` diff --git a/python/foundry/integration/sglang/__init__.py b/python/foundry/integration/sglang/__init__.py index 7a771998..ab4b28b0 100644 --- a/python/foundry/integration/sglang/__init__.py +++ b/python/foundry/integration/sglang/__init__.py @@ -1,3 +1,15 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the Foundry project """Foundry integration for SGLang.""" + +from foundry.integration.sglang.config import ( + CUDAGraphExtensionMode, + get_graph_extension_mode, +) +from foundry.integration.sglang.hooks import install_hooks + +__all__ = [ + "CUDAGraphExtensionMode", + "get_graph_extension_mode", + "install_hooks", +] diff --git a/python/foundry/integration/sglang/cuda_graph.py b/python/foundry/integration/sglang/cuda_graph.py new file mode 100644 index 00000000..c8abc003 --- /dev/null +++ b/python/foundry/integration/sglang/cuda_graph.py @@ -0,0 +1,570 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the Foundry project +"""Foundry hooks for SGLang's v0.5.15.post1 split CUDA-graph runner. + +This module is the optional-dependency boundary for SGLang. It imports +SGLang's private runner APIs at module scope, so callers must load it only +while installing hooks in a process where the supported SGLang version is +available. The public ``foundry.integration.sglang`` package does not import +this module eagerly. +""" + +from __future__ import annotations + +import functools +import importlib +import logging +import os +import sys +import time +from collections.abc import Callable, Iterable, Iterator +from contextlib import contextmanager +from dataclasses import fields +from typing import Any + +from sglang.srt.model_executor.runner.base_runner import BaseRunner +from sglang.srt.model_executor.runner.decode_cuda_graph_runner import ( + DecodeCudaGraphRunner, +) +from sglang.srt.model_executor.runner.shape_key import ShapeKey +from sglang.srt.model_executor.runner_backend.full_cuda_graph_backend import ( + FullCudaGraphBackend, +) + +from foundry.integration.sglang import graph_ops +from foundry.integration.sglang.config import ( + CUDAGraphExtensionMode, +) + +logger = logging.getLogger(__name__) + +_ORIGINAL_ATTR = "__foundry_sglang_original__" +_EAGER_DEEPEP_WARMUP = False + + +def _module(name: str): + """Resolve a boundary module, including test/runtime replacements.""" + + return sys.modules.get(name) or importlib.import_module(name) + + +def _mode() -> CUDAGraphExtensionMode: + config = _module("foundry.integration.sglang.config") + return config.get_graph_extension_mode() + + +def _value(value: Any) -> Any: + """Return the value of an enum-like setting without requiring its type.""" + + return getattr(value, "value", value) + + +def _decode_attention_backend(server_args: Any) -> str | None: + backend = getattr(server_args, "decode_attention_backend", None) + if backend is None: + backend = getattr(server_args, "attention_backend", None) + return _value(backend) + + +def validate_configuration(server_args: Any) -> None: + """Reject SGLang configurations outside Foundry's validated graph scope.""" + + cuda_graph_config = getattr(server_args, "cuda_graph_config", None) + decode = getattr(getattr(cuda_graph_config, "decode", None), "backend", None) + prefill = getattr(getattr(cuda_graph_config, "prefill", None), "backend", None) + if _value(decode) != "full": + raise ValueError( + f"Foundry requires cuda_graph_config decode backend 'full'; got {_value(decode)!r}." + ) + if _value(prefill) != "disabled": + raise ValueError( + "Foundry requires cuda_graph_config prefill backend 'disabled'; " + f"got {_value(prefill)!r}." + ) + + tp_size = getattr(server_args, "tp_size", 1) + dp_size = getattr(server_args, "dp_size", 1) + pp_size = getattr(server_args, "pp_size", 1) + ep_size = getattr(server_args, "ep_size", 1) + dp_attention = bool(getattr(server_args, "enable_dp_attention", False)) + regular_tp = tp_size > 1 and not dp_attention + if regular_tp: + topology = (tp_size, dp_size, pp_size, ep_size) + if topology != (2, 2, 1, 1): + raise ValueError( + "Foundry tensor-parallel decode supports only " + "tp_size=2, dp_size=2, pp_size=1, ep_size=1; " + f"got tp_size={tp_size}, dp_size={dp_size}, " + f"pp_size={pp_size}, ep_size={ep_size}." + ) + if not bool(getattr(server_args, "disable_custom_all_reduce", False)): + raise ValueError("Foundry TP=2 x DP=2 requires --disable-custom-all-reduce.") + if bool(getattr(server_args, "enable_nccl_nvls", False)): + raise ValueError( + "Foundry TP=2 x DP=2 does not support NCCL NVLS; leave --enable-nccl-nvls disabled." + ) + if bool(getattr(server_args, "enable_symm_mem", False)) or bool( + getattr(server_args, "enable_torch_symm_mem", False) + ): + raise ValueError( + "Foundry TP=2 x DP=2 does not support symmetric memory collectives; " + "leave --enable-symm-mem and --enable-torch-symm-mem disabled." + ) + for variable in ("NCCL_CUMEM_ENABLE", "NCCL_NVLS_ENABLE"): + value = os.environ.get(variable) + if value != "0": + raise ValueError(f"Foundry TP=2 x DP=2 requires {variable}='0'; got {value!r}.") + if pp_size > 1: + raise ValueError("Foundry does not support pipeline parallel CUDA graphs (pp_size > 1).") + if getattr(server_args, "speculative_algorithm", None) is not None: + raise ValueError("Foundry does not support speculative decoding.") + if bool(getattr(server_args, "enable_lora", False)) or bool( + getattr(server_args, "lora_paths", None) + ): + raise ValueError("Foundry does not support LoRA CUDA graphs.") + if bool(getattr(server_args, "enable_pdmux", False)): + raise ValueError("Foundry does not support PDMux CUDA graphs.") + if getattr(server_args, "dllm_algorithm", None) is not None: + raise ValueError("Foundry does not support dLLM or diffusion CUDA graphs.") + if bool(getattr(server_args, "enable_return_hidden_states", False)): + raise ValueError("Foundry does not support hidden-state capture.") + if bool(getattr(server_args, "enable_memory_saver", False)): + raise ValueError("Foundry does not support memory-saver CUDA graphs.") + if bool(getattr(server_args, "debug_cuda_graph", False)): + raise ValueError("Foundry requires the full backend and cannot use debug CUDA graphs.") + + attention_backend = _decode_attention_backend(server_args) + moe_backend = _value(getattr(server_args, "moe_a2a_backend", "none")) + if dp_attention: + if attention_backend != "fa3" or moe_backend != "deepep": + raise ValueError( + "Foundry DP-attention requires the FA3 attention backend and DeepEP MoE all-to-all." + ) + deepep_mode = _value(getattr(server_args, "deepep_mode", "auto")) + if deepep_mode not in ("auto", "low_latency"): + raise ValueError( + "Foundry DP-attention requires DeepEP low-latency decode mode " + "('--deepep-mode auto' or 'low_latency')." + ) + elif attention_backend != "flashinfer": + raise ValueError( + "Foundry non-DP-attention decode requires the FlashInfer " + f"attention backend; got {attention_backend!r}." + ) + elif moe_backend != "none": + raise ValueError( + "Foundry only supports a MoE all-to-all backend with the " + "DP-attention + DeepEP + FA3 configuration." + ) + + +def register_graph_pool(pool: Any) -> None: + """Register one pool with both latest split-runner pool registries.""" + + pool_module = _module("sglang.srt.model_executor.runner_utils.pool") + pynccl_allocator = _module("sglang.srt.distributed.device_communicators.pynccl_allocator") + pool_module.set_global_graph_memory_pool(pool) + pynccl_allocator.set_graph_pool_id(pool) + + +def hydrate_backend( + backend: FullCudaGraphBackend, + loaded: Iterable[tuple[ShapeKey, Any, Any]], + *, + pool: Any, +) -> None: + """Install loaded graphs and outputs in a full split-runner backend.""" + + if not isinstance(backend, FullCudaGraphBackend): + raise RuntimeError( + f"Foundry LOAD requires SGLang's FullCudaGraphBackend; got {type(backend).__name__}." + ) + backend._graphs.clear() + backend._outputs.clear() + for shape_key, graph, output in loaded: + if not isinstance(shape_key, ShapeKey): + raise TypeError( + "Foundry LOAD requires complete SGLang ShapeKey objects, " + f"got {type(shape_key).__name__}." + ) + backend._graphs[shape_key] = graph + backend._outputs[shape_key] = output + backend._pool = pool + + +def _mark_wrapper(wrapper: Callable[..., Any], original: Callable[..., Any]): + setattr(wrapper, _ORIGINAL_ATTR, original) + return wrapper + + +def _patch_once(cls: type, name: str, factory: Callable[[Callable[..., Any]], Callable[..., Any]]): + current = getattr(cls, name) + if hasattr(current, _ORIGINAL_ATTR): + return + setattr(cls, name, _mark_wrapper(factory(current), current)) + + +def _warmup_wrapper(original: Callable[..., Any]): + @functools.wraps(original) + def patched(self): + mode = _mode() + if mode == CUDAGraphExtensionMode.NONE: + return original(self) + logger.info("[Foundry] SGLang runner warmup skipped in %s mode", mode.value) + return None + + return patched + + +def _ensure_full_backend(runner: DecodeCudaGraphRunner) -> FullCudaGraphBackend: + backend = runner.backend + if not isinstance(backend, FullCudaGraphBackend): + raise RuntimeError( + "Foundry requires SGLang's FullCudaGraphBackend and will not " + f"recapture with {type(backend).__name__}." + ) + return backend + + +def _uses_deepep(runner: DecodeCudaGraphRunner) -> bool: + return _value(getattr(runner.model_runner.server_args, "moe_a2a_backend", "none")) == "deepep" + + +def _pack_output(output: Any) -> tuple[Any, str]: + logits_processor = _module("sglang.srt.layers.logits_processor") + logits_output_type = logits_processor.LogitsProcessorOutput + if not isinstance(output, logits_output_type): + raise TypeError( + "Foundry supports only SGLang LogitsProcessorOutput CUDA graph outputs; " + f"got {type(output).__name__}." + ) + if output.next_token_logits is None: + raise TypeError("SGLang decode CUDA graph output has no next_token_logits.") + + populated_unsupported_fields = [ + field.name + for field in fields(output) + if field.name != "next_token_logits" and getattr(output, field.name) is not None + ] + if populated_unsupported_fields: + raise TypeError( + "Foundry supports only next_token_logits in SGLang decode CUDA graph output; " + f"unsupported populated fields: {populated_unsupported_fields}." + ) + return output.next_token_logits, "next_token_logits" + + +def _unpack_output(tensors: Any, output_kind: str) -> Any: + if output_kind != "next_token_logits": + raise RuntimeError( + f"Unsupported SGLang CUDA graph output kind during LOAD: {output_kind!r}." + ) + if isinstance(tensors, (tuple, list)): + if len(tensors) != 1: + raise RuntimeError(f"Expected one SGLang CUDA graph output tensor, got {len(tensors)}.") + tensors = tensors[0] + logits_processor = _module("sglang.srt.layers.logits_processor") + return logits_processor.LogitsProcessorOutput(next_token_logits=tensors) + + +def bootstrap_deepep_buffer(cuda_graph_runner: DecodeCudaGraphRunner) -> bool: + """Initialize DeepEP's singleton buffer before CUDA graph capture.""" + + moe_utils = _module("sglang.srt.layers.moe.utils") + if not moe_utils.get_moe_a2a_backend().is_deepep(): + return False + deepep_module = _module("sglang.srt.layers.moe.token_dispatcher.deepep") + deep_ep_buffer = deepep_module.DeepEPBuffer + deep_ep_dispatcher = deepep_module.DeepEPDispatcher + + if deep_ep_buffer._buffer is not None: + return True + + model = cuda_graph_runner.model_runner.model + for module in model.modules(): + dispatcher = getattr(module, "dispatcher", None) + if dispatcher is None: + continue + candidates = [dispatcher, *getattr(dispatcher, "_inners", [])] + deepep = next( + (item for item in candidates if isinstance(item, deep_ep_dispatcher)), + None, + ) + if deepep is None: + continue + implementation = getattr(deepep, "_low_latency_dispatcher", None) or getattr( + deepep, + "_normal_dispatcher", + None, + ) + if implementation is None: + continue + start = time.perf_counter() + implementation._get_buffer() + logger.info( + "[Foundry] Bootstrapped DeepEP buffer pre-capture in %.3fs", + time.perf_counter() - start, + ) + return True + + logger.warning( + "[Foundry] DeepEP backend active but no DeepEPDispatcher found on the " + "model; buffer not bootstrapped (capture may fail inside stream capture)." + ) + return False + + +@contextmanager +def _deepep_eager_warmup() -> Iterator[None]: + global _EAGER_DEEPEP_WARMUP + previous = _EAGER_DEEPEP_WARMUP + compile_utils = _module("sglang.srt.layers.deep_gemm_wrapper.compile_utils") + original_compile = compile_utils._compile_deep_gemm_one_type_all + + @functools.wraps(original_compile) + def compile_outside_allocation_region(*args, **kwargs): + # DeepGEMM compiles all shapes only on the first local rank. Its warmup + # executor is explicitly deleted and empty_cache()'d by upstream, so + # only that executor is transient JIT state. The surrounding eager + # forward still owns persistent lazy allocations used by saved graphs. + graph_ops.cge.stop_allocation_region() + try: + return original_compile(*args, **kwargs) + finally: + graph_ops.cge.resume_allocation_region() + + compile_utils._compile_deep_gemm_one_type_all = compile_outside_allocation_region + _EAGER_DEEPEP_WARMUP = True + try: + yield + finally: + _EAGER_DEEPEP_WARMUP = previous + compile_utils._compile_deep_gemm_one_type_all = original_compile + + +def _is_flashinfer_backend(attn_backend: Any) -> bool: + return hasattr(attn_backend, "indices_updater_decode") and hasattr( + attn_backend, "decode_cuda_graph_metadata" + ) + + +def _prepare_attention_metadata( + runner: DecodeCudaGraphRunner, +) -> list[tuple[Any, int, Any]]: + """Create per-shape metadata in upstream's reverse capture order.""" + + prepared = [] + for size in reversed(runner.capture_bs): + forward_batch, attn_backend, _pp_proxy_tensors = runner.capture_prepare(size) + attn_backend.init_forward_metadata_out_graph( + forward_batch, + in_capture=True, + ) + prepared.append((attn_backend, size, attn_backend.forward_metadata)) + return prepared + + +@contextmanager +def _reuse_flashinfer_metadata( + prepared: list[tuple[Any, int, Any]], +) -> Iterator[None]: + """Reuse pre-created wrappers while retaining the latest planner path. + + The pre-pass performs the allocation-bearing ``in_capture=True`` call. The + inner upstream capture then selects the saved metadata object and invokes + the same planner with ``in_capture=False``. This avoids a second wrapper + allocation while retaining FlashInfer planning and the SWA output-location + binding installed by the pre-pass. + """ + + by_backend: dict[int, tuple[Any, dict[int, Any]]] = {} + for attn_backend, size, metadata in prepared: + entry = by_backend.setdefault(id(attn_backend), (attn_backend, {})) + entry[1][size] = metadata + + restorations = [] + for attn_backend, metadata_by_size in by_backend.values(): + original = attn_backend.init_forward_metadata_out_graph + had_instance_override = "init_forward_metadata_out_graph" in vars(attn_backend) + previous_override = vars(attn_backend).get("init_forward_metadata_out_graph") + + @functools.wraps(original) + def reuse( + forward_batch, + in_capture=False, + *, + _attn_backend=attn_backend, + _metadata_by_size=metadata_by_size, + _original=original, + ): + if in_capture: + _attn_backend.forward_metadata = _metadata_by_size[forward_batch.batch_size] + return _original(forward_batch, in_capture=False) + return _original(forward_batch, in_capture=in_capture) + + attn_backend.init_forward_metadata_out_graph = reuse + restorations.append((attn_backend, had_instance_override, previous_override)) + + try: + yield + finally: + for attn_backend, had_instance_override, previous_override in restorations: + if had_instance_override: + attn_backend.init_forward_metadata_out_graph = previous_override + else: + del attn_backend.init_forward_metadata_out_graph + + +def _save_capture(runner: DecodeCudaGraphRunner, original: Callable[..., Any]): + if _uses_deepep(runner): + start = time.perf_counter() + # The DeepEP/NVSHMEM buffer is a live graph dependency. Allocate it + # symmetrically inside Foundry's region before excluding transient JIT + # warmup allocations from the recorded cursor. + bootstrap_deepep_buffer(runner) + with _deepep_eager_warmup(): + original(runner) + logger.info( + "[Foundry] SGLang DeepEP eager lazy-init warmup completed in %.3fs", + time.perf_counter() - start, + ) + + if _is_flashinfer_backend(runner.attn_backend): + prepared = _prepare_attention_metadata(runner) + with _reuse_flashinfer_metadata(prepared): + result = original(runner) + else: + result = original(runner) + + graph_ops.save_graph_manifest() + graph_ops.pack_fatbins() + graph_ops.capture_final_alloc_offset() + return result + + +def _load_capture( + runner: DecodeCudaGraphRunner, + original: Callable[..., Any], +) -> None: + if getattr(runner, "_foundry_graphs_loaded", False): + raise RuntimeError( + "Foundry LOAD cannot satisfy an SGLang CUDA-graph recapture request; " + "restart with a fresh compatible archive." + ) + + backend = _ensure_full_backend(runner) + pool = backend._pool + if pool is None: + pool = runner.device_module.graph_pool_handle() + register_graph_pool(pool) + backend._pool = pool + + if _uses_deepep(runner): + bootstrap_deepep_buffer(runner) + with _deepep_eager_warmup(): + original(runner) + + graph_ops.preallocate_for_load_mode() + _prepare_attention_metadata(runner) + loaded = [] + for entry, graph, tensors in graph_ops.load_all_graphs(runner): + archive_key = entry.shape_key + shape_key = ShapeKey( + size=archive_key.size, + stream_idx=archive_key.stream_idx, + variant_label=archive_key.variant_label, + ) + loaded.append( + ( + shape_key, + graph, + _unpack_output(tensors, entry.output_kind), + ) + ) + hydrate_backend(backend, loaded, pool=pool) + runner.deepep_adapter.capture(is_extend_in_batch=False) + runner._foundry_graphs_loaded = True + + +def _capture_wrapper(original: Callable[..., Any]): + @functools.wraps(original) + def patched(self): + mode = _mode() + if mode == CUDAGraphExtensionMode.NONE: + return original(self) + + _ensure_full_backend(self) + if mode == CUDAGraphExtensionMode.SAVE: + return _save_capture(self, original) + if mode == CUDAGraphExtensionMode.LOAD: + return _load_capture(self, original) + raise RuntimeError(f"Unsupported Foundry SGLang mode: {mode!r}") + + return patched + + +def _capture_one_wrapper(original: Callable[..., Any]): + @functools.wraps(original) + def patched( + self, + shape_key, + forward_fn, + dummies=None, + post_warmup_hook=None, + ): + mode = _mode() + if mode == CUDAGraphExtensionMode.NONE: + return original( + self, + shape_key, + forward_fn, + dummies=dummies, + post_warmup_hook=post_warmup_hook, + ) + if not isinstance(self, FullCudaGraphBackend): + raise RuntimeError("Foundry graph capture requires SGLang's FullCudaGraphBackend.") + if _EAGER_DEEPEP_WARMUP: + output = forward_fn() + if post_warmup_hook is not None: + post_warmup_hook() + return output + if mode == CUDAGraphExtensionMode.LOAD: + raise RuntimeError( + "Foundry LOAD attempted an upstream graph capture; the archive " + "is incompatible with this SGLang configuration." + ) + if not isinstance(shape_key, ShapeKey): + raise TypeError("SGLang split-runner capture must supply a complete ShapeKey.") + + graph = graph_ops.create_device_graph() + output = graph_ops.capture_graph( + graph, + self._pool, + self._capture_stream, + forward_fn, + ) + packed_output, output_kind = _pack_output(output) + archive_key = graph_ops.ArchiveShapeKey( + size=shape_key.size, + stream_idx=shape_key.stream_idx, + variant_label=shape_key.variant_label, + ) + graph_ops.save_graph( + graph, + packed_output, + archive_key, + output_kind, + ) + self._graphs[shape_key] = graph + self._outputs[shape_key] = output + return None + + return patched + + +def install_cuda_graph_hooks() -> None: + """Install idempotent hooks on SGLang's latest split-runner classes.""" + + _patch_once(BaseRunner, "warmup", _warmup_wrapper) + _patch_once(DecodeCudaGraphRunner, "capture", _capture_wrapper) + _patch_once(FullCudaGraphBackend, "capture_one", _capture_one_wrapper) diff --git a/python/foundry/integration/sglang/graph_ops.py b/python/foundry/integration/sglang/graph_ops.py index 18ab1f0e..f9f80e45 100644 --- a/python/foundry/integration/sglang/graph_ops.py +++ b/python/foundry/integration/sglang/graph_ops.py @@ -2,12 +2,14 @@ # SPDX-FileCopyrightText: Copyright contributors to the Foundry project """Foundry CUDA graph save/load helpers for SGLang.""" -from __future__ import annotations - +import importlib +import json import logging import os -import re +import tempfile import time +from dataclasses import asdict, dataclass +from pathlib import Path from typing import Any import torch @@ -21,66 +23,223 @@ get_config, get_graph_extension_mode, ) -from foundry.integration.sglang.runtime import get_state + +rt = importlib.import_module("foundry.integration.sglang.runtime") logger = logging.getLogger(__name__) -_pending_graph_builds: tuple[Any, list[tuple[int, str, dict[str, Any]]]] | None = None -_GRAPH_FILENAME_RE = re.compile(r"^graph_(?P\d+)_FULL_t(?P\d+)_r\d+_UX_pcN\.json$") +GRAPH_INDEX_SCHEMA_VERSION = 1 +GRAPH_INDEX_FILENAME = "sglang_graph_index.json" +_SUPPORTED_PHASE = "decode" +_SUPPORTED_OUTPUT_KIND = "next_token_logits" -def _batch_size_from_key(key: Any) -> int: - if isinstance(key, int): - return key - key_str = str(key) - for part in reversed(key_str.split("_")): - if part.isdigit(): - return int(part) - raise ValueError(f"Cannot derive batch size from SGLang CUDA graph key: {key!r}") +@dataclass(frozen=True) +class ArchiveShapeKey: + size: int + stream_idx: int | None = None + variant_label: str | None = None -def _graph_filename(index: int, key: Any) -> str: - batch_size = _batch_size_from_key(key) - return f"graph_{index}_FULL_t{batch_size}_r{batch_size}_UX_pcN.json" +@dataclass(frozen=True) +class GraphIndexEntry: + phase: str + capture_index: int + filename: str + shape_key: ArchiveShapeKey + output_kind: str -def _pack_output(output: Any) -> torch.Tensor: - from sglang.srt.layers.logits_processor import LogitsProcessorOutput +def _fresh_save_error(message: str) -> ValueError: + return ValueError(f"{message} Perform a fresh Foundry SAVE with SGLang 0.5.15.post1.") - if isinstance(output, LogitsProcessorOutput): - if output.next_token_logits is None: - raise TypeError("SGLang decode CUDA graph output has no next_token_logits") - return output.next_token_logits - if isinstance(output, torch.Tensor): - return output +def _require_int(value: Any, field_name: str, *, minimum: int = 0) -> int: + if isinstance(value, bool) or not isinstance(value, int) or value < minimum: + raise _fresh_save_error( + f"SGLang graph index field {field_name!r} must be an integer >= {minimum}; " + f"got {value!r}." + ) + return value + + +def _normalize_shape_key(shape_key: Any) -> ArchiveShapeKey: + if isinstance(shape_key, dict): + expected = {"size", "stream_idx", "variant_label"} + if set(shape_key) != expected: + raise _fresh_save_error( + "SGLang graph index shape_key must contain exactly size, stream_idx, " + "and variant_label." + ) + size = shape_key["size"] + stream_idx = shape_key["stream_idx"] + variant_label = shape_key["variant_label"] + elif all( + hasattr(shape_key, field_name) for field_name in ("size", "stream_idx", "variant_label") + ): + size = shape_key.size + stream_idx = shape_key.stream_idx + variant_label = shape_key.variant_label + else: + raise _fresh_save_error( + "SGLang graph index shape_key must be a complete record with size, " + "stream_idx, and variant_label." + ) - raise TypeError(f"Unsupported SGLang CUDA graph output type: {type(output)!r}") + size = _require_int(size, "shape_key.size", minimum=1) + if stream_idx is not None: + stream_idx = _require_int(stream_idx, "shape_key.stream_idx") + if variant_label is not None and not isinstance(variant_label, str): + raise _fresh_save_error( + "SGLang graph index field 'shape_key.variant_label' must be a string or null; " + f"got {variant_label!r}." + ) + return ArchiveShapeKey( + size=size, + stream_idx=stream_idx, + variant_label=variant_label, + ) -def _unpack_output(tensors: Any) -> Any: - from sglang.srt.layers.logits_processor import LogitsProcessorOutput +def _normalize_entry(entry: Any) -> GraphIndexEntry: + if isinstance(entry, dict): + expected = { + "phase", + "capture_index", + "filename", + "shape_key", + "output_kind", + } + if set(entry) != expected: + raise _fresh_save_error( + "SGLang graph index entries must contain exactly phase, capture_index, " + "filename, shape_key, and output_kind." + ) + values = entry + elif isinstance(entry, GraphIndexEntry): + values = { + "phase": entry.phase, + "capture_index": entry.capture_index, + "filename": entry.filename, + "shape_key": entry.shape_key, + "output_kind": entry.output_kind, + } + else: + raise _fresh_save_error( + f"SGLang graph index entry must be an object; got {type(entry).__name__}." + ) - if isinstance(tensors, (tuple, list)): - if len(tensors) != 1: - raise RuntimeError(f"Expected one SGLang CUDA graph output tensor, got {len(tensors)}") - tensors = tensors[0] - return LogitsProcessorOutput(next_token_logits=tensors) + phase = values["phase"] + if phase != _SUPPORTED_PHASE: + raise _fresh_save_error( + f"Unsupported SGLang graph phase {phase!r}; only {_SUPPORTED_PHASE!r} is supported." + ) + capture_index = _require_int(values["capture_index"], "capture_index") + filename = values["filename"] + if ( + not isinstance(filename, str) + or not filename + or filename in {".", ".."} + or "/" in filename + or "\\" in filename + or Path(filename).name != filename + ): + raise _fresh_save_error( + f"SGLang graph index filename must be a non-empty basename; got {filename!r}." + ) + output_kind = values["output_kind"] + if output_kind != _SUPPORTED_OUTPUT_KIND: + raise _fresh_save_error( + f"Unsupported SGLang graph output kind {output_kind!r}; " + f"only {_SUPPORTED_OUTPUT_KIND!r} is supported." + ) + return GraphIndexEntry( + phase=phase, + capture_index=capture_index, + filename=filename, + shape_key=_normalize_shape_key(values["shape_key"]), + output_kind=output_kind, + ) -def _scan_graph_files(workspace_dir: str) -> list[tuple[int, str, dict[str, Any]]]: - graph_files = [] - for filename in os.listdir(workspace_dir): - match = _GRAPH_FILENAME_RE.match(filename) - if not match: - continue - meta = { - "index": int(match.group("index")), - "key": int(match.group("bs")), - } - graph_files.append((int(meta["index"]), filename, meta)) - graph_files.sort(key=lambda x: x[0]) - return graph_files +def _normalize_entries(entries: Any) -> list[GraphIndexEntry]: + if not isinstance(entries, (tuple, list)): + raise _fresh_save_error("SGLang graph index 'graphs' must be a list.") + normalized = sorted( + (_normalize_entry(entry) for entry in entries), + key=lambda entry: entry.capture_index, + ) + if normalized and normalized[0].capture_index != 0: + raise _fresh_save_error("SGLang graph capture_index values must start at zero.") + filenames = [entry.filename for entry in normalized] + if len(filenames) != len(set(filenames)): + raise _fresh_save_error("SGLang graph index filenames must be unique.") + for previous, current in zip(normalized, normalized[1:]): + if current.capture_index == previous.capture_index: + raise _fresh_save_error( + f"Duplicate SGLang graph capture_index {current.capture_index}." + ) + if current.capture_index != previous.capture_index + 1: + raise _fresh_save_error( + "SGLang graph capture_index values must be contiguous in capture order." + ) + return normalized + + +def save_graph_index( + path: str | os.PathLike[str], + entries: list[GraphIndexEntry], +) -> None: + index_path = Path(path) + normalized = _normalize_entries(entries) + payload = { + "schema_version": GRAPH_INDEX_SCHEMA_VERSION, + "graphs": [asdict(entry) for entry in normalized], + } + temporary_path = None + try: + with tempfile.NamedTemporaryFile( + "w", + encoding="utf-8", + dir=index_path.parent, + prefix=f".{index_path.name}.", + suffix=".tmp", + delete=False, + ) as temporary: + temporary_path = temporary.name + json.dump(payload, temporary, indent=2, sort_keys=True) + temporary.write("\n") + temporary.flush() + os.fsync(temporary.fileno()) + os.replace(temporary_path, index_path) + finally: + if temporary_path is not None and os.path.exists(temporary_path): + os.unlink(temporary_path) + + +def load_graph_index(path: str | os.PathLike[str]) -> list[GraphIndexEntry]: + index_path = Path(path) + try: + with index_path.open(encoding="utf-8") as index_file: + payload = json.load(index_file) + except (OSError, json.JSONDecodeError) as exc: + raise _fresh_save_error( + f"Cannot read the latest SGLang graph index at {index_path}: {exc}" + ) from exc + + if not isinstance(payload, dict): + raise _fresh_save_error("SGLang graph index root must be an object.") + schema_version = payload.get("schema_version") + if isinstance(schema_version, bool) or schema_version != GRAPH_INDEX_SCHEMA_VERSION: + raise _fresh_save_error( + "Unsupported or legacy SGLang graph index schema " + f"{schema_version!r}; expected {GRAPH_INDEX_SCHEMA_VERSION}." + ) + if set(payload) != {"schema_version", "graphs"}: + raise _fresh_save_error( + "SGLang graph index must contain exactly schema_version and graphs." + ) + return _normalize_entries(payload["graphs"]) def create_device_graph(): @@ -98,25 +257,50 @@ def capture_graph(graph, pool, stream, run_once_fn): return None -def save_graph(graph, output: Any, key: Any) -> None: +def save_graph( + graph: Any, + packed_output: Any, + shape_key: ArchiveShapeKey, + output_kind: str, +) -> None: cfg = get_config() - state = get_state() + state = rt.get_state() if cfg is None or state is None or cfg.workspace_dir is None: raise RuntimeError("Foundry SGLang graph extension is not initialized") - packed_output = _pack_output(output) - filename = _graph_filename(state.capture_index, key) + entry = _normalize_entry( + GraphIndexEntry( + phase=_SUPPORTED_PHASE, + capture_index=state.capture_index, + filename=f"graph_{state.capture_index:06d}_sglang.json", + shape_key=shape_key, + output_kind=output_kind, + ) + ) + filename = entry.filename graph_path = os.path.join(cfg.workspace_dir, filename) graph.save(graph_path, packed_output) + state.graph_index_entries.append(entry) state.capture_index += 1 - logger.info("[Foundry] Saved SGLang CUDA graph %s key=%s", filename, key) + logger.info( + "[Foundry] Saved SGLang CUDA graph %s shape_key=%s", + filename, + shape_key, + ) def save_graph_manifest() -> None: cfg = get_config() if cfg is None or cfg.workspace_dir is None: return + state = rt.get_state() + if state is None or not state.graph_index_entries: + raise RuntimeError("Foundry SAVE captured no SGLang CUDA graphs") + save_graph_index( + os.path.join(cfg.workspace_dir, GRAPH_INDEX_FILENAME), + state.graph_index_entries, + ) foundry_pkg.save_graph_manifest(cfg.workspace_dir) @@ -128,168 +312,13 @@ def pack_fatbins() -> None: cge.set_pack_fatbins_on_exit(False) -def start_graph_builds() -> None: - global _pending_graph_builds - cfg = get_config() - if cfg is None or cfg.workspace_dir is None or cfg.mode != CUDAGraphExtensionMode.LOAD: - return - - graph_files = _scan_graph_files(cfg.workspace_dir) - if not graph_files: - raise RuntimeError(f"No Foundry SGLang graph files found in {cfg.workspace_dir}") - - paths = [os.path.join(cfg.workspace_dir, filename) for _, filename, _ in graph_files] - t0 = time.perf_counter() - pending = FoundryCUDAGraph.start_graph_builds(paths, num_threads=4) - _pending_graph_builds = (pending, graph_files) - logger.info( - "[Foundry] Started SGLang graph builds for %d graphs in %.3fs", - len(paths), - time.perf_counter() - t0, - ) - - -def preload_all_graphs() -> None: - global _pending_graph_builds - cfg = get_config() - state = get_state() - if cfg is None or state is None or cfg.workspace_dir is None: - raise RuntimeError("Foundry SGLang graph extension is not initialized") - - if _pending_graph_builds is None: - start_graph_builds() - assert _pending_graph_builds is not None - - cge.init_nvshmem_for_loaded_modules() - - pending, graph_files = _pending_graph_builds - _pending_graph_builds = None - - t0 = time.perf_counter() - results = FoundryCUDAGraph.finish_graph_loads(pending) - logger.info( - "[Foundry] Finished SGLang graph loads for %d graphs in %.3fs", - len(results), - time.perf_counter() - t0, - ) - - for i, (_index, _filename, meta) in enumerate(graph_files): - graph, tensors = results[i] - state.loaded_graphs[meta["key"]] = (graph, _unpack_output(tensors)) - - -def bootstrap_deepep_buffer(cuda_graph_runner) -> bool: - """Force the singleton DeepEP ``Buffer`` (NVSHMEM runtime + symmetric heap) - to be created BEFORE the cuda-graph capture loop. - - sglang creates the DeepEP buffer lazily on the first MoE dispatch โ€” normally - during the two pre-capture warmup forwards. Foundry suppresses those warmups - (for allocation determinism), which would push buffer creation into the - captured forward, where ``deep_ep_cpp.Buffer(...)`` aborts with - ``operation not permitted when stream is capturing``. - - Triggering it here (outside any stream capture) creates only the NVSHMEM - runtime + symmetric heap โ€” no model activations โ€” so it stays symmetric - across SAVE and LOAD and lands at the same VMM offset on both. The buffer is - a process-wide singleton (``DeepEPBuffer._buffer``), so one creation per rank - is enough. It is a collective over the EP group, so every rank must reach - this point together โ€” which they do, since ``capture`` runs on all ranks. - - Returns True if a buffer was (or already is) created, False if DeepEP is off. - """ - try: - from sglang.srt.layers.moe.utils import get_moe_a2a_backend - - if not get_moe_a2a_backend().is_deepep(): - return False - except Exception: - return False - - from sglang.srt.layers.moe.token_dispatcher.deepep import ( - DeepEPBuffer, - DeepEPDispatcher, - ) - - if DeepEPBuffer._buffer is not None: - return True - - model = cuda_graph_runner.model_runner.model - for module in model.modules(): - dispatcher = getattr(module, "dispatcher", None) - if dispatcher is None: - continue - # ``module.dispatcher`` is normally a MaybeTboDeepEPDispatcher wrapper - # whose ``_inners`` hold the real DeepEPDispatcher(s); unwrap it. (Also - # handle a bare DeepEPDispatcher for safety.) - candidates = [dispatcher, *getattr(dispatcher, "_inners", [])] - deepep = next((d for d in candidates if isinstance(d, DeepEPDispatcher)), None) - if deepep is None: - continue - # Prefer the low-latency impl (the mode foundry captures); the buffer - # is sized for whichever impls exist, so either bootstraps the shared - # singleton. - impl = getattr(deepep, "_low_latency_dispatcher", None) or getattr( - deepep, "_normal_dispatcher", None - ) - if impl is None: - continue - t0 = time.perf_counter() - impl._get_buffer() - logger.info( - "[Foundry] Bootstrapped DeepEP buffer pre-capture in %.3fs", - time.perf_counter() - t0, - ) - return True - - logger.warning( - "[Foundry] DeepEP backend active but no DeepEPDispatcher found on the " - "model; buffer not bootstrapped (capture may fail inside stream capture)." - ) - return False - - -def initialize_attention_metadata_for_bs(cuda_graph_runner, bs: int) -> None: - """Populate ``decode_cuda_graph_metadata[bs]`` for runtime replay. - - The FlashInfer wrappers and their internal ``_int_workspace_buffer`` - are constructed here, outside the captured graph. The graph's - runtime kernels reference these buffer addresses, so LOAD must - re-run the same call before runtime replay so the wrappers exist - at deterministic VMM addresses. - """ - 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( - bs, - num_tokens, - buffers.req_pool_indices[:bs], - buffers.seq_lens[:bs], - encoder_lens, - cuda_graph_runner.capture_forward_mode, - spec_info, - ) - - -def initialize_all_attention_metadata(cuda_graph_runner) -> None: - """Pre-pass: populate ``decode_cuda_graph_metadata`` for all bs at once. - - Called on both SAVE and LOAD before the capture/load loop. Walking - ``reversed(self.capture_bs)`` (largest first) matches SAVE's natural - capture order; same order on both sides keeps the VMM cursor - trajectory identical. - """ - for bs in reversed(cuda_graph_runner.capture_bs): - initialize_attention_metadata_for_bs(cuda_graph_runner, bs) - - -def load_all_graphs(cuda_graph_runner) -> None: +def load_all_graphs( + _cuda_graph_runner: Any, +) -> list[tuple[GraphIndexEntry, Any, Any]]: """LOAD-time replacement for the upstream capture loop. - All FlashInfer wrappers are pre-allocated by - ``initialize_all_attention_metadata`` (called by the capture hook - before this function), so the VMM cursor sits where SAVE recorded + The capture hook recreates per-shape attention metadata before this + function, so the VMM cursor sits where SAVE recorded ``start_base_addr_0``. Load every graph in one ``start_graph_builds`` call โ€” this is what enables template + on-demand linking in the manifest. ``finish_graph_loads`` then @@ -297,29 +326,57 @@ def load_all_graphs(cuda_graph_runner) -> None: cursor exactly the way SAVE did inside its capture loop. """ cfg = get_config() - state = get_state() + state = rt.get_state() if cfg is None or state is None or cfg.workspace_dir is None: raise RuntimeError("Foundry SGLang graph extension is not initialized") - graph_files = _scan_graph_files(cfg.workspace_dir) - if not graph_files: - raise RuntimeError(f"No Foundry SGLang graph files found in {cfg.workspace_dir}") + entries = load_graph_index(os.path.join(cfg.workspace_dir, GRAPH_INDEX_FILENAME)) + if not entries: + raise _fresh_save_error("The SGLang graph index contains no decode graphs.") + missing_files = [ + entry.filename + for entry in entries + if not os.path.isfile(os.path.join(cfg.workspace_dir, entry.filename)) + ] + if missing_files: + raise _fresh_save_error( + f"The SGLang graph index references missing graph files: {missing_files}." + ) # NVSHMEM init runs once before any graph loads โ€” graphs may reference # NVSHMEM symbols. Single-GPU dense models have 0 NVSHMEM modules, so # this is a no-op there but kept for EP parity. cge.init_nvshmem_for_loaded_modules() - paths = [os.path.join(cfg.workspace_dir, filename) for _, filename, _ in graph_files] + paths = [os.path.join(cfg.workspace_dir, entry.filename) for entry in entries] t0 = time.perf_counter() pending = FoundryCUDAGraph.start_graph_builds(paths, num_threads=4) results = FoundryCUDAGraph.finish_graph_loads(pending) + if len(results) != len(entries): + raise RuntimeError( + "Foundry loaded an unexpected number of SGLang graphs: " + f"expected {len(entries)}, got {len(results)}" + ) logger.info( "[Foundry] Loaded %d SGLang graphs in %.3fs", len(results), time.perf_counter() - t0, ) - for i, (_index, _filename, meta) in enumerate(graph_files): - graph, tensors = results[i] - state.loaded_graphs[meta["key"]] = (graph, _unpack_output(tensors)) + loaded = [] + for entry, result in zip(entries, results): + graph, tensors = result + loaded.append((entry, graph, tensors)) + return loaded + + +def capture_final_alloc_offset() -> int: + """Persist the final deterministic VMM cursor through runtime state.""" + + return rt.capture_final_alloc_offset() + + +def preallocate_for_load_mode() -> None: + """Preallocate the saved deterministic VMM span before graph loading.""" + + rt.preallocate_for_load_mode() diff --git a/python/foundry/integration/sglang/hooks.py b/python/foundry/integration/sglang/hooks.py index 22cd7cfa..183fa2e0 100644 --- a/python/foundry/integration/sglang/hooks.py +++ b/python/foundry/integration/sglang/hooks.py @@ -5,11 +5,14 @@ from __future__ import annotations import functools +import importlib import logging import os import time from dataclasses import asdict +import torch + from foundry.integration.sglang import runtime as rt from foundry.integration.sglang.config import ( CUDAGraphExtensionMode, @@ -22,17 +25,6 @@ _INSTALLED = False -def _ep_lazy_init_needed() -> bool: - """True when the DeepEP all-to-all backend is active, so pre-capture lazy - init (NVSHMEM buffer, DeepGEMM JIT) must be warmed up outside stream capture.""" - try: - from sglang.srt.layers.moe.utils import get_moe_a2a_backend - - return get_moe_a2a_backend().is_deepep() - except Exception: - return False - - def _resolve_dp_rank(model_runner) -> int | None: dp_rank = getattr(model_runner, "dp_rank", None) if dp_rank is not None: @@ -40,12 +32,11 @@ def _resolve_dp_rank(model_runner) -> int | None: server_args = model_runner.server_args if getattr(server_args, "enable_dp_attention", False): - from sglang.srt.layers.dp_attention import compute_dp_attention_world_info - + dp_attention = importlib.import_module("sglang.srt.layers.dp_attention") # Returns (attn_tp_rank, attn_tp_size, attn_dp_rank, attn_dp_size) โ€” a # 4-tuple. Take attn_dp_rank (3rd) and discard the rest. (Unpacking into # 3 targets would raise ValueError.) - _, _, dp_rank, _ = compute_dp_attention_world_info( + _, _, dp_rank, _ = dp_attention.compute_dp_attention_world_info( server_args.enable_dp_attention, model_runner.tp_rank, server_args.tp_size, @@ -80,6 +71,10 @@ def install_hooks(server_args) -> None: ) load_graph_extension_config(cfg_path) + # Optional-dependency boundary: cuda_graph imports SGLang's private runner + # APIs at module scope and is therefore loaded only during hook installation. + cuda_graph = importlib.import_module("foundry.integration.sglang.cuda_graph") + cuda_graph.validate_configuration(server_args) logger.info( "[Foundry] SGLang hooks installing: mode=%s workspace=%s", get_graph_extension_mode().value, @@ -88,9 +83,7 @@ def install_hooks(server_args) -> None: _patch_init_torch_distributed() _patch_init_memory_pool() - _patch_load_model() - _patch_kernel_warmup() - _patch_cuda_graph_capture() + cuda_graph.install_cuda_graph_hooks() _patch_spawn_sites() _INSTALLED = True @@ -98,8 +91,7 @@ def install_hooks(server_args) -> None: def _patch_init_torch_distributed() -> None: - from sglang.srt.model_executor import model_runner as mr - + mr = importlib.import_module("sglang.srt.model_executor.model_runner") cls = mr.ModelRunner orig = cls.init_torch_distributed @@ -119,8 +111,6 @@ def patched(self, *args, **kwargs): # (surfacing at the first Stream()/kernel). Mirrors model_runner's own # set_device(self.gpu_id). Single-GPU is unaffected (gpu_id == 0). if self.device == "cuda": - import torch - torch.get_device_module(self.device).set_device(self.gpu_id) rt.setup_graph_extension( @@ -140,9 +130,8 @@ def patched(self, *args, **kwargs): def _patch_init_memory_pool() -> None: - from sglang.srt.model_executor import model_runner_kv_cache_mixin as kv_mixin - from sglang.srt.model_executor.pool_configurator import MemoryPoolConfig - + kv_mixin = importlib.import_module("sglang.srt.model_executor.model_runner_kv_cache_mixin") + pool_configurator = importlib.import_module("sglang.srt.model_executor.pool_configurator") cls = kv_mixin.ModelRunnerKVCacheMixin orig = cls.init_memory_pool @@ -153,13 +142,12 @@ def patched(self, pre_model_load_memory): return orig(self, pre_model_load_memory) if mode == CUDAGraphExtensionMode.LOAD: - import torch - rt.log_alloc_offset("before_init_memory_pool") state = rt.load_warmup_state() + rt.validate_warmup_state(state, self.server_args) if not state.memory_pool_config: raise RuntimeError("Foundry LOAD requires memory_pool_config") - self.memory_pool_config = MemoryPoolConfig(**state.memory_pool_config) + self.memory_pool_config = pool_configurator.MemoryPoolConfig(**state.memory_pool_config) # Mirror SAVE's ``_resolve_memory_pool_config`` -> # ``get_available_gpu_memory(empty_cache=True)`` side # effect. Without this, torch's caching allocator retains @@ -176,409 +164,36 @@ def patched(self, pre_model_load_memory): rt.log_alloc_offset("before_init_memory_pool") result = orig(self, pre_model_load_memory) rt.log_alloc_offset("after_init_memory_pool") - state = rt.create_warmup_state(asdict(self.memory_pool_config)) + state = rt.create_warmup_state( + asdict(self.memory_pool_config), + self.server_args, + ) rt.save_warmup_state(state) return result cls.init_memory_pool = patched -def _patch_load_model() -> None: - from sglang.srt.model_executor import model_runner as mr - - cls = mr.ModelRunner - orig = cls.load_model - - @functools.wraps(orig) - def patched(self, *args, **kwargs): - return orig(self, *args, **kwargs) - - cls.load_model = patched - - -def _patch_kernel_warmup() -> None: - from sglang.srt.model_executor import model_runner as mr - - cls = mr.ModelRunner - orig = cls.kernel_warmup - - @functools.wraps(orig) - def patched(self, *args, **kwargs): - mode = get_graph_extension_mode() - if mode == CUDAGraphExtensionMode.NONE: - return orig(self, *args, **kwargs) - # Phase 1 keeps pre-graph model-forward warmups out of SAVE/LOAD. - logger.info("[Foundry] SGLang kernel_warmup skipped in %s mode", mode.value) - return None - - cls.kernel_warmup = patched - - -def _patch_cuda_graph_capture() -> None: - from sglang.srt.model_executor import cuda_graph_runner as cgr - - cls = cgr.CudaGraphRunner - orig_capture = cls.capture - orig_create_device_graph = cls._create_device_graph - orig_capture_graph = cls._capture_graph - orig_capture_one_batch_size = cls.capture_one_batch_size - - # When set, the capture machinery is being reused as a foundry-driven WARMUP - # pass: run a real forward per bs (no graph capture, no save) to trigger all - # of sglang's pre-capture lazy init โ€” DeepEP buffer, DeepGEMM per-shape JIT, - # etc. โ€” that would otherwise fire inside the captured stream and abort with - # "operation not permitted when stream is capturing". See `_run_warmup_pass`. - warmup_active = [False] - - @functools.wraps(orig_create_device_graph) - def patched_create_device_graph(self, *args, **kwargs): - mode = get_graph_extension_mode() - if warmup_active[0]: - # Throwaway graph object; the warmup never captures into it. - return orig_create_device_graph(self, *args, **kwargs) - if mode == CUDAGraphExtensionMode.SAVE: - from foundry.integration.sglang.graph_ops import create_device_graph - - return create_device_graph() - return orig_create_device_graph(self, *args, **kwargs) - - @functools.wraps(orig_capture_graph) - def patched_capture_graph(self, graph, pool, stream, run_once_fn): - mode = get_graph_extension_mode() - if warmup_active[0]: - # Warmup: run the forward eagerly (NO torch.cuda.graph), so DeepGEMM - # JIT compile / NVSHMEM buffer creation happen outside stream capture. - return run_once_fn() - if mode == CUDAGraphExtensionMode.SAVE: - from foundry.integration.sglang.graph_ops import capture_graph - - return capture_graph(graph, pool, stream, run_once_fn) - return orig_capture_graph(self, graph, pool, stream, run_once_fn) - - @functools.wraps(orig_capture_one_batch_size) - def patched_capture_one_batch_size(self, bs, forward, stream_idx=None): - mode = get_graph_extension_mode() - if warmup_active[0]: - # Warmup pass: run upstream unchanged (real warmup forwards + a - # neutered "capture" that is just another real forward). No - # suppression, no save_graph. - return orig_capture_one_batch_size(self, bs, forward, stream_idx) - if mode == CUDAGraphExtensionMode.SAVE: - # Suppress the two pre-capture warmup forwards - # (cuda_graph_runner.py: ``for _ in range(2): run_once()``). - # Their non-deterministic activation allocations would pollute - # the torch caching allocator with freed segments that LOAD - # cannot reproduce โ€” causing the per-bs init's cache-miss vs - # cache-hit asymmetry that drifts the VMM cursor away from - # each saved ``start_base_addr``. JIT / autotune still happens - # inside the graph capture (3rd run_once invocation) and is - # recorded as alloc events, mirroring vLLM doc 04 ยง2. - counter = [0] - real_forward = forward - - def warmup_skipping_forward(*args, **kwargs): - counter[0] += 1 - if counter[0] <= 2: - return None - return real_forward(*args, **kwargs) - - forward = warmup_skipping_forward - graph, output = orig_capture_one_batch_size(self, bs, forward, stream_idx) - if mode == CUDAGraphExtensionMode.SAVE: - from foundry.integration.sglang.graph_ops import save_graph - - # Mirror the inline key shape upstream uses for self.graphs[key] - # in `_capture_one_stream`. ``_make_graph_key`` and - # ``get_capture_lora_variant`` were removed in sglang - # commit ce2506e1c (record_nolora_graph deprecation). - key = bs if stream_idx is None else f"{stream_idx}_{bs}" - save_graph(graph, output, key) - return graph, output - - def _run_warmup_pass(self): - """Foundry-driven pre-capture warmup for the DeepEP/EP path. - - Reuses the upstream capture loop with graph capture neutered (see - ``warmup_active``) to run one real forward per ``capture_bs`` BEFORE the - real capture/replay work โ€” triggering every pre-capture lazy init sglang - normally does in its (foundry-suppressed) warmup forwards: DeepEP buffer - creation, DeepGEMM per-shape JIT compile, etc. Run on BOTH SAVE and LOAD - so the VMM cursor trajectory stays symmetric. ``self.graphs`` / - ``self.output_buffers`` are saved/cleared/restored so the throwaway - warmup entries don't leak into the real pass. - """ - warmup_active[0] = True - saved_graphs, saved_buffers = self.graphs, self.output_buffers - self.graphs, self.output_buffers = {}, {} - t0 = time.perf_counter() - try: - orig_capture(self) - finally: - warmup_active[0] = False - self.graphs, self.output_buffers = saved_graphs, saved_buffers - logger.info( - "[Foundry] SGLang EP warmup pass (lazy-init) completed in %.3fs", - time.perf_counter() - t0, - ) - - @functools.wraps(orig_capture) - def patched(self, *args, **kwargs): - mode = get_graph_extension_mode() - - # DeepEP/EP only (gated so the validated dense / single-GPU / DP paths - # are untouched): - # SAVE: sglang triggers pre-capture lazy init (DeepGEMM per-shape JIT, - # DeepEP buffer, ...) via warmup forwards that foundry suppresses; - # left lazy they fire inside the captured stream and abort - # ("operation not permitted when stream is capturing"). Run a real - # forward per bs up front so all that happens outside capture. - # LOAD: no capture happens โ€” preallocate_for_load_mode reserves the - # whole region up to final_alloc_offset and replay places each alloc - # at its recorded absolute offset, so LOAD need NOT replay the warmup - # cursor trajectory. Running the warmup here is not just unnecessary - # but harmful: its orig_capture re-enters graph_capture() and leaves - # the context in a state that breaks the threaded finish_graph_loads - # ("invalid device context"). So warmup is SAVE-only. - # bootstrap_deepep_buffer runs on both (cheap singleton) to guarantee the - # NVSHMEM runtime is up before replay on LOAD / before capture on SAVE. - if _ep_lazy_init_needed(): - if mode == CUDAGraphExtensionMode.SAVE: - _run_warmup_pass(self) - if mode != CUDAGraphExtensionMode.NONE: - from foundry.integration.sglang.graph_ops import ( - bootstrap_deepep_buffer, - ) - - bootstrap_deepep_buffer(self) - - if mode == CUDAGraphExtensionMode.LOAD: - from sglang.srt.distributed.device_communicators.pynccl_allocator import ( - set_graph_pool_id, - ) - - from foundry.integration.sglang.graph_ops import ( - initialize_all_attention_metadata, - load_all_graphs, - ) - - state = rt.get_state() - if state is None: - raise RuntimeError("Foundry SGLang state is not initialized") - # Set up the graph memory pool once โ€” sglang shares one pool - # across all captured graphs, and runtime replay also requires - # it to be set so pynccl knows which pool the graph belongs to. - if cgr.get_global_graph_memory_pool() is None: - cgr.set_global_graph_memory_pool(self.device_module.graph_pool_handle()) - set_graph_pool_id(cgr.get_global_graph_memory_pool()) - rt.log_alloc_offset("before_preallocate") - rt.preallocate_for_load_mode() - rt.log_alloc_offset("after_preallocate") - # Pre-pass: allocate every per-bs FlashInfer wrapper in - # ``reversed(capture_bs)`` order, matching the order SAVE used. - # SAVE's patched ``init_forward_metadata_capture_cuda_graph`` - # is idempotent so the inner per-iter call inside the SAVE - # capture loop does not re-allocate. Same upfront allocation - # sequence on both sides โ†’ cursor sits at SAVE's - # ``start_base_addr_0`` when graph load begins. - # FlashInfer-only pre-pass (see SAVE branch). fa3 etc. allocate their - # cuda-graph metadata once in init_cuda_graph_state, so skip it. - if hasattr(self.attn_backend, "indices_updater_decode"): - initialize_all_attention_metadata(self) - rt.log_alloc_offset("after_pre_init") - # Single ``start_graph_builds(all_paths)`` call so templates - # and on-demand graphs link via ``shared_exec`` in the - # manifest. ``finish_graph_loads`` replays alloc events - # graph-by-graph in the same order SAVE captured them. - load_all_graphs(self) - rt.log_alloc_offset("after_load_all_graphs") - self.graphs = {k: v[0] for k, v in state.loaded_graphs.items()} - self.output_buffers = {k: v[1] for k, v in state.loaded_graphs.items()} - # Non-FlashInfer backends (e.g. fa3) populate per-bs decode metadata - # โ€” looked up at replay as attn_backend.decode_cuda_graph_metadata[bs] - # โ€” inside the capture loop, which LOAD replaces. The FlashInfer - # pre-pass above (gated) handled flashinfer; for the rest, populate it - # now. Run AFTER load_all_graphs so the (already-correct) loaded-graph - # VMM offsets are unaffected โ€” fa3's metadata are lightweight views - # over the fixed init_cuda_graph_state workspace, not graph memory. - if not hasattr(self.attn_backend, "indices_updater_decode"): - initialize_all_attention_metadata(self) - # Initialize the DeepEP cuda-graph adapter's captured mode. Upstream - # sets this in deepep_adapter.capture() during the capture loop, - # which LOAD replaces โ€” so without this, runtime replay() asserts - # `_captured_deepep_mode is not None` on the first decode. capture() - # self-gates on the DeepEP backend (no-op otherwise) and sets the - # decode dispatch mode the captured graphs expect. - self.deepep_adapter.capture(is_extend_in_batch=False) - return None - - if mode == CUDAGraphExtensionMode.SAVE: - # FlashInfer allocates a per-bs metadata wrapper (each with its own - # _int_workspace_buffer) on every capture init, so foundry pre-allocates - # them up front and installs a reuse shim that makes the inner init - # reuse them โ€” keeping the VMM cursor deterministic vs LOAD. Backends - # with a single fixed cuda-graph metadata workspace allocated once in - # init_cuda_graph_state (e.g. fa3 / FlashAttentionBackend) don't need - # this and the plain capture is already SAVE/LOAD-deterministic. Detect - # 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 - - def reuse_pre_pass_init( - bs, - num_tokens, - req_pool_indices, - seq_lens, - encoder_lens, - forward_mode, - spec_info, - ): - # The pre-pass already allocated a wrapper for this - # bs and stored it in - # ``decode_cuda_graph_metadata`` / - # ``prefill_cuda_graph_metadata``. Reuse it directly - # โ€” no second torch.empty for ``_int_workspace_buffer``. - # Re-run the planner with the same buffer slices the - # capture forward uses, then point - # ``forward_metadata`` at the same wrappers. Same - # plan call on LOAD via the symmetric pre-pass, so - # the captured graph kernels reference VMM addresses - # that LOAD's wrappers actually occupy. - from sglang.srt.layers.attention.flashinfer_backend import ( - DecodeMetadata, - PrefillMetadata, - ) - - if forward_mode.is_decode_or_idle(): - wrappers = attn_backend.decode_cuda_graph_metadata.get(bs) - if wrappers is None: - return real_init( - bs, - num_tokens, - req_pool_indices, - seq_lens, - encoder_lens, - forward_mode, - spec_info, - ) - seq_lens_sum = seq_lens.sum().item() - attn_backend.indices_updater_decode.update( - req_pool_indices, - seq_lens, - seq_lens.cpu(), - seq_lens_sum, - decode_wrappers=wrappers, - encoder_lens=encoder_lens, - spec_info=spec_info, - fixed_split_size=None, - disable_split_kv=attn_backend.disable_cuda_graph_kv_split, - ) - attn_backend.forward_metadata = DecodeMetadata(wrappers) - return - if ( - 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_init( - bs, - num_tokens, - req_pool_indices, - seq_lens, - encoder_lens, - forward_mode, - spec_info, - ) - seq_lens_sum = seq_lens.sum().item() - use_ragged = forward_mode.is_dllm_extend() - prefix_lens = ( - seq_lens - attn_backend.dllm_config.block_size - if forward_mode.is_dllm_extend() - else None - ) - spec_info_arg = None if forward_mode.is_dllm_extend() else spec_info - attn_backend.indices_updater_prefill.update( - req_pool_indices, - seq_lens, - seq_lens.cpu(), - seq_lens_sum, - prefix_lens=prefix_lens, - prefill_wrappers=wrappers, - use_ragged=use_ragged, - encoder_lens=encoder_lens, - spec_info=spec_info_arg, - ) - attn_backend.forward_metadata = PrefillMetadata(wrappers, use_ragged, False) - return - # Unknown mode โ€” fall back to real init. - return real_init( - bs, - num_tokens, - req_pool_indices, - seq_lens, - encoder_lens, - forward_mode, - spec_info, - ) - - if use_fi_prepass: - from foundry.integration.sglang.graph_ops import ( - initialize_all_attention_metadata, - ) - - rt.log_alloc_offset("save_before_pre_init") - # Pre-pass: allocate every per-bs FlashInfer wrapper up front, in - # the same ``reversed(capture_bs)`` order LOAD uses. - initialize_all_attention_metadata(self) - rt.log_alloc_offset("save_after_pre_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 - try: - result = orig_capture(self, *args, **kwargs) - finally: - attn_backend.init_forward_metadata_capture_cuda_graph = real_init - - from foundry.integration.sglang.graph_ops import ( - pack_fatbins, - save_graph_manifest, - ) - - save_graph_manifest() - pack_fatbins() - rt.capture_final_alloc_offset() - return result - - return orig_capture(self, *args, **kwargs) - - cls._create_device_graph = patched_create_device_graph - cls._capture_graph = patched_capture_graph - cls.capture_one_batch_size = patched_capture_one_batch_size - cls.capture = patched - - def _patch_spawn_sites() -> None: try: - from sglang.srt.entrypoints import engine as engine_mod + engine_mod = importlib.import_module("sglang.srt.entrypoints.engine") except Exception: engine_mod = None if engine_mod is not None: - orig_launch = engine_mod.Engine._launch_scheduler_processes + engine_cls = engine_mod.Engine + orig_launch = engine_cls.__dict__["_launch_scheduler_processes"].__func__ @functools.wraps(orig_launch) - def patched_launch(self, *args, **kwargs): + def patched_launch(cls, *args, **kwargs): if get_graph_extension_mode() != CUDAGraphExtensionMode.NONE: rt.setup_ld_preload_env() - return orig_launch(self, *args, **kwargs) + return orig_launch(cls, *args, **kwargs) - engine_mod.Engine._launch_scheduler_processes = patched_launch + engine_cls._launch_scheduler_processes = classmethod(patched_launch) try: - from sglang.srt.managers import data_parallel_controller as dpc + dpc = importlib.import_module("sglang.srt.managers.data_parallel_controller") except Exception: dpc = None diff --git a/python/foundry/integration/sglang/runtime.py b/python/foundry/integration/sglang/runtime.py index 1f9032c5..b4e54a3c 100644 --- a/python/foundry/integration/sglang/runtime.py +++ b/python/foundry/integration/sglang/runtime.py @@ -8,9 +8,12 @@ import logging import os import shutil +import tempfile import time from dataclasses import asdict, dataclass, field from datetime import datetime +from importlib.metadata import PackageNotFoundError +from importlib.metadata import version as package_version from pathlib import Path import torch @@ -28,23 +31,30 @@ logger = logging.getLogger(__name__) +WARMUP_STATE_SCHEMA_VERSION = 1 +SUPPORTED_SGLANG_VERSION = "0.5.15.post1" +_SUPPORTED_DECODE_BACKEND = "full" +_SUPPORTED_PREFILL_BACKEND = "disabled" + @dataclass class WarmupState: + schema_version: int = 0 sglang_version: str = "" + decode_cuda_graph_backend: str = "" + prefill_cuda_graph_backend: str = "" timestamp: str = "" cuda_version: str = "" gpu_name: str = "" gpu_total_memory: int = 0 memory_pool_config: dict = field(default_factory=dict) - final_alloc_offset: int = 0 @dataclass class CUDAGraphExtensionState: capture_index: int = 0 rank: int = 0 - loaded_graphs: dict = field(default_factory=dict) + graph_index_entries: list[object] = field(default_factory=list) _state: CUDAGraphExtensionState | None = None @@ -65,15 +75,43 @@ def _workspace_root() -> str | None: return None if cfg is None else cfg.workspace_root -def create_warmup_state(memory_pool_config: dict | None = None) -> WarmupState: - try: - from sglang.version import __version__ as sglang_version - except Exception: - sglang_version = "unknown" +def _setting_value(value): + return getattr(value, "value", value) + + +def _cuda_graph_backends(server_args) -> tuple[str | None, str | None]: + cuda_graph_config = getattr(server_args, "cuda_graph_config", None) + decode = getattr(getattr(cuda_graph_config, "decode", None), "backend", None) + prefill = getattr(getattr(cuda_graph_config, "prefill", None), "backend", None) + return _setting_value(decode), _setting_value(prefill) + +def _installed_sglang_version() -> str: + try: + return package_version("sglang") + except PackageNotFoundError: + return "unknown" + + +def create_warmup_state( + memory_pool_config: dict | None, + server_args, +) -> WarmupState: + sglang_version = _installed_sglang_version() + if sglang_version != SUPPORTED_SGLANG_VERSION: + raise RuntimeError( + "Foundry SAVE requires SGLang package version " + f"{SUPPORTED_SGLANG_VERSION!r}, but installed package metadata reports " + f"{sglang_version!r}. Install SGLang {SUPPORTED_SGLANG_VERSION} as a package " + "before SAVE so its exact version can be recorded." + ) + decode_backend, prefill_backend = _cuda_graph_backends(server_args) props = torch.cuda.get_device_properties(0) return WarmupState( + schema_version=WARMUP_STATE_SCHEMA_VERSION, sglang_version=sglang_version, + decode_cuda_graph_backend=decode_backend or "", + prefill_cuda_graph_backend=prefill_backend or "", timestamp=datetime.now().isoformat(), cuda_version=torch.version.cuda or "unknown", gpu_name=props.name, @@ -82,6 +120,58 @@ def create_warmup_state(memory_pool_config: dict | None = None) -> WarmupState: ) +def _compatibility_error(message: str) -> RuntimeError: + return RuntimeError(f"{message} Perform a fresh Foundry SAVE before using LOAD.") + + +def validate_warmup_state(state: WarmupState, server_args) -> None: + if ( + isinstance(state.schema_version, bool) + or state.schema_version != WARMUP_STATE_SCHEMA_VERSION + ): + raise _compatibility_error( + "Unsupported or legacy SGLang archive schema " + f"{state.schema_version!r}; expected {WARMUP_STATE_SCHEMA_VERSION}." + ) + + if state.sglang_version != SUPPORTED_SGLANG_VERSION: + raise _compatibility_error( + "SGLang archive version mismatch: " + f"recorded {state.sglang_version!r}, expected {SUPPORTED_SGLANG_VERSION!r}." + ) + installed_version = _installed_sglang_version() + if installed_version != SUPPORTED_SGLANG_VERSION: + raise _compatibility_error( + "Installed SGLang version mismatch: " + f"found {installed_version!r}, expected {SUPPORTED_SGLANG_VERSION!r}." + ) + + if state.decode_cuda_graph_backend != _SUPPORTED_DECODE_BACKEND: + raise _compatibility_error( + "SGLang archive decode CUDA graph backend mismatch: " + f"recorded {state.decode_cuda_graph_backend!r}, " + f"expected {_SUPPORTED_DECODE_BACKEND!r}." + ) + if state.prefill_cuda_graph_backend != _SUPPORTED_PREFILL_BACKEND: + raise _compatibility_error( + "SGLang archive prefill CUDA graph backend mismatch: " + f"recorded {state.prefill_cuda_graph_backend!r}, " + f"expected {_SUPPORTED_PREFILL_BACKEND!r}." + ) + + decode_backend, prefill_backend = _cuda_graph_backends(server_args) + if decode_backend != state.decode_cuda_graph_backend: + raise _compatibility_error( + "Current SGLang decode CUDA graph backend does not match the archive: " + f"current {decode_backend!r}, archive {state.decode_cuda_graph_backend!r}." + ) + if prefill_backend != state.prefill_cuda_graph_backend: + raise _compatibility_error( + "Current SGLang prefill CUDA graph backend does not match the archive: " + f"current {prefill_backend!r}, archive {state.prefill_cuda_graph_backend!r}." + ) + + def save_warmup_state(state: WarmupState) -> None: workspace_root = _workspace_root() if workspace_root is None: @@ -91,8 +181,25 @@ def save_warmup_state(state: WarmupState) -> None: if ext_state is not None and ext_state.rank != 0 and os.path.exists(path): return os.makedirs(workspace_root, exist_ok=True) - with open(path, "w") as f: - json.dump(asdict(state), f, indent=2) + temporary_path = None + try: + with tempfile.NamedTemporaryFile( + "w", + encoding="utf-8", + dir=workspace_root, + prefix=".warmup_state.", + suffix=".tmp", + delete=False, + ) as temporary: + temporary_path = temporary.name + json.dump(asdict(state), temporary, indent=2, sort_keys=True) + temporary.write("\n") + temporary.flush() + os.fsync(temporary.fileno()) + os.replace(temporary_path, path) + finally: + if temporary_path is not None and os.path.exists(temporary_path): + os.unlink(temporary_path) logger.info("[Foundry] Saved SGLang warmup state to %s", path) @@ -102,9 +209,14 @@ def load_warmup_state() -> WarmupState: raise RuntimeError("Foundry workspace_root is not initialized") path = os.path.join(workspace_root, "warmup_state.json") if not os.path.exists(path): - raise RuntimeError(f"Foundry warmup state file not found: {path}") - with open(path) as f: - data = json.load(f) + raise _compatibility_error(f"Foundry warmup state file not found: {path}.") + try: + with open(path) as f: + data = json.load(f) + except (OSError, json.JSONDecodeError) as exc: + raise _compatibility_error(f"Cannot read the SGLang warmup state at {path}: {exc}") from exc + if not isinstance(data, dict): + raise _compatibility_error("SGLang warmup state root must be an object.") valid = set(WarmupState.__dataclass_fields__.keys()) return WarmupState(**{k: v for k, v in data.items() if k in valid}) @@ -170,12 +282,6 @@ def capture_final_alloc_offset() -> int: path = os.path.join(cfg.workspace_dir, "final_alloc_offset.json") with open(path, "w") as f: json.dump({"final_alloc_offset": _final_alloc_offset}, f) - if cfg.workspace_root is not None: - warmup_state_path = os.path.join(cfg.workspace_root, "warmup_state.json") - if os.path.exists(warmup_state_path): - state = load_warmup_state() - state.final_alloc_offset = _final_alloc_offset - save_warmup_state(state) logger.info("[Foundry] SGLang final_alloc_offset=%d", _final_alloc_offset) return _final_alloc_offset @@ -184,17 +290,35 @@ def preallocate_for_load_mode() -> None: cfg = get_config() if cfg is None or cfg.mode != CUDAGraphExtensionMode.LOAD: return - final = 0 - if cfg.workspace_dir is not None: - path = os.path.join(cfg.workspace_dir, "final_alloc_offset.json") - if os.path.exists(path): - with open(path) as f: - final = json.load(f).get("final_alloc_offset", 0) - if final <= 0: - final = load_warmup_state().final_alloc_offset - remaining = final - cge.get_current_alloc_offset() + if cfg.workspace_dir is None: + raise _compatibility_error( + "Foundry LOAD has no rank workspace for its per-rank final allocation offset." + ) + path = os.path.join(cfg.workspace_dir, "final_alloc_offset.json") + try: + with open(path) as f: + data = json.load(f) + except (OSError, json.JSONDecodeError) as exc: + raise _compatibility_error( + f"Cannot read the per-rank final allocation offset at {path}: {exc}" + ) from exc + final = data.get("final_alloc_offset") if isinstance(data, dict) else None + if isinstance(final, bool) or not isinstance(final, int) or final <= 0: + raise _compatibility_error( + f"Invalid per-rank final allocation offset at {path}: {final!r}; " + "expected a positive integer." + ) + current = cge.get_current_alloc_offset() + remaining = final - current if remaining > 0: - cge.preallocate_region(remaining) + success = cge.preallocate_region(remaining) + if not success: + rank = None if _state is None else _state.rank + raise RuntimeError( + "Foundry LOAD failed to preallocate the recorded CUDA allocation region: " + f"rank={rank}, requested={remaining}, current={current}, final={final}, " + f"source='per_rank', path={path!r}." + ) def log_alloc_offset(label: str) -> None: diff --git a/recipe/sglang/README.md b/recipe/sglang/README.md index af4ce528..d638d9af 100644 --- a/recipe/sglang/README.md +++ b/recipe/sglang/README.md @@ -1,154 +1,267 @@ -# Foundry recipe โ€” SGLang +# Foundry recipe โ€” SGLang v0.5.15.post1 -End-to-end serve scripts for SAVE / LOAD of CUDA graphs through the foundry SGLang -integration. All scripts in this directory share the same pair of foundry TOML files -(`foundry_save.toml` / `foundry_load.toml`) โ€” pick a script for your model + parallelism, -run `--save`, then `--load`, then query. The integration code is in -[`../../python/foundry/integration/sglang/`](../../python/foundry/integration/sglang/); -design notes are under [`../../docs/sglang/`](../../docs/sglang/). +These recipes persist and restore full decode CUDA graphs for exactly SGLang +`0.5.15.post1` (upstream commit `0b3bb0c`). Validation used CUDA 13.0, PyTorch +2.11/cu130, and H100. Prefill graphs are always disabled. -## Files in this directory +| Script | Validated configuration | +|---|---| +| `serve_qwen3-mini.sh` | Single GPU, Qwen3-1.7B, FlashInfer | +| `serve_qwen3-1.7b_dp.sh` | Regular data parallelism, one full replica per GPU, FlashInfer | +| `serve_qwen3-8b_tp2_dp2.sh` | Exact dense TP=2 ร— DP=2, PP=1, EP=1, Qwen3-8B, FlashInfer | +| `serve_qwen3-30ba3bfp8_ep.sh` | DP-attention EP=2/EP=4, DeepEP low-latency, DeepGEMM, FA3, Qwen3-30B-A3B-FP8 | -``` -recipe/sglang/ -โ”œโ”€โ”€ README.md # this file -โ”œโ”€โ”€ foundry_save.toml # shared SAVE config (workspace_root = "foundry_archive") -โ”œโ”€โ”€ foundry_load.toml # shared LOAD config (same workspace_root) -โ”œโ”€โ”€ serve_qwen3-mini.sh # Qwen3-1.7B single GPU -โ”œโ”€โ”€ serve_qwen3-1.7b_dp.sh # Qwen3-1.7B data parallel -โ””โ”€โ”€ serve_qwen3-30ba3bfp8_ep.sh # Qwen3-30B-A3B FP8 expert parallel (DeepEP) -``` +Single GPU and DP=2 passed baseline/SAVE/fresh LOAD/query with text/token +parity, no recapture, and 52 graphs per applicable rank. Dense TP=2 ร— DP=2 +passed on four H100s with explicitly routed DP0/DP1 parity, rank workspaces +0โ€“3, 36 graphs per rank, identical final offsets of `69822578688`, and no +IPC/CUDA/preallocation errors. EP=2 and EP=4 passed with 20 graphs per rank +plus DeepEP/NVSHMEM/bootstrap/adapter/event restoration evidence; EP=4 +restored 416 event nodes. + +The dense TP recipe has dedicated `foundry_tp2_dp2_save.toml` and +`foundry_tp2_dp2_load.toml` files. Every other conventional TP topology, +custom all-reduce, NCCL CUMEM/NVLS, pipeline parallelism, speculative +decoding, LoRA/PDMux, dLLM, hidden-state capture, and memory-saver graphs are +not supported. EP=4 is DP-attention/expert parallelism, not conventional TP +attention, even though SGLang sets `tp_size=4` in that process topology. + +## Requirements -Every script accepts the same trailing `--save` / `--load` flag. Scripts that scale -across GPUs take the parallel-size as the first positional argument: +Use a dedicated environment with: + +- Python 3.12; +- CUDA 13; +- PyTorch 2.11.0 with cu130; +- `flashinfer_python[cu13]==0.6.12` and `flashinfer_cubin==0.6.12`; +- `sglang-kernel==0.4.4`; +- `sgl-deep-gemm==0.1.4`; +- the SGLang FA3 package, `kernels-community/sgl-flash-attn3` revision 1, + declared by SGLang through the `kernels` package. + +The expert-parallel recipe additionally requires DeepEP commit `9af0e0d` and +the NVSHMEM library supplied with the cu13 PyTorch environment. + +## Check out and install SGLang + +Keep Foundry and SGLang as sibling checkouts. For this Foundry-only workflow, +clone the authoritative upstream tag: ```bash -bash serve_qwen3-mini.sh [--save|--load] -bash serve_qwen3-1.7b_dp.sh [--save|--load] -bash serve_qwen3-30ba3bfp8_ep.sh [--save|--load] +git clone --branch v0.5.15.post1 --depth 1 \ + https://github.com/sgl-project/sglang.git sglang +test "$(git -C sglang rev-parse HEAD)" = \ + "0b3bb0cbe31873994c9f989fddfe2f87ca839fdd" ``` -A single SAVE pass is enough โ€” SGLang has no startup profile-forward, so there is no -non-determinism that requires a two-pass save (unlike the vLLM recipe). +Apply the four-file shim in +[`../../docs/sglang/direct-edits.md`](../../docs/sglang/direct-edits.md), then +install both projects as packages from the directory containing the sibling +checkouts: -Because the two TOMLs are shared (single `workspace_root = "foundry_archive"`), one -archive is written per host; run a fresh `rm -rf foundry_archive` whenever you change -model or topology before SAVE. +```bash +conda create -p venv python=3.12 +conda activate venv +conda install -c conda-forge boost-cpp boost -| Mode | Script | Model | Notes | -|---|---|---|---| -| Single GPU | `serve_qwen3-mini.sh` | Qwen3-1.7B | FlashInfer backend | -| Data parallel | `serve_qwen3-1.7b_dp.sh` | Qwen3-1.7B | one full replica/rank; `NCCL_CUMEM_ENABLE=0`/`NCCL_NVLS_ENABLE=0` | -| Expert parallel | `serve_qwen3-30ba3bfp8_ep.sh` | Qwen3-30B-A3B-FP8 | DP-attention + DeepEP; fa3 backend | +pip install torch==2.11.0 torchvision==0.26.0 torchaudio==2.11.0 \ + --index-url https://download.pytorch.org/whl/cu130 +SETUPTOOLS_SCM_PRETEND_VERSION=0.5.15.post1 \ + pip install -e sglang/python \ + --extra-index-url https://download.pytorch.org/whl/cu130 +pushd foundry +CC=gcc CXX=g++ pip install -e . --no-build-isolation +popd -## Installation +python -c 'from importlib.metadata import version; assert version("sglang") == "0.5.15.post1"' +``` + +SGLang uses `setuptools_scm`. The required tracked-file shim makes the tagged +checkout dirty, so an unqualified editable install derives a dirty/local +version instead of exactly `0.5.15.post1`. +`SETUPTOOLS_SCM_PRETEND_VERSION=0.5.15.post1` deterministically embeds the +required distribution version during installation. The exact runtime metadata +gate remains in force. + +An editable SGLang install with that override is supported because it creates +exact distribution metadata. A source checkout exposed only through +`PYTHONPATH` is not supported: SAVE and LOAD both require installed package +metadata to report exactly `0.5.15.post1`. -The recipes assume `foundry` and the SGLang fork are **pip-installed** (editable is -fine) so both import without any `PYTHONPATH`, and foundry's spawn-site patch -auto-detects `libcuda_hook.so` from its install โ€” the scripts set no `LD_PRELOAD` -themselves. The standard workspace layout has `foundry/` (this repo) and `sglang/` -(the foundry-org SGLang fork) as siblings: +`foundry-org/sglang` should carry the same four-file patch once its +`v0.5.15.post1` release branch is published. These instructions do not assume +that branch exists or already contains the shim. +The tagged release pins the kernel stack above in `python/pyproject.toml`. +Confirm the resolved environment before capture: + +```bash +python -m pip show sglang torch flashinfer-python flashinfer-cubin \ + sglang-kernel sgl-deep-gemm kernels ``` -/ -โ”œโ”€โ”€ foundry/ # this repo -โ”‚ โ”œโ”€โ”€ python/foundry/ # `pip install -e .` builds libcuda_hook.so here -โ”‚ โ”œโ”€โ”€ recipe/sglang/ # <-- you are here -โ”‚ โ””โ”€โ”€ ... -โ””โ”€โ”€ sglang/ # foundry-org/sglang fork (with direct edits applied) + +For DeepEP on H100: + +```bash +git clone https://github.com/deepseek-ai/DeepEP.git +cd DeepEP +git checkout 9af0e0d0e74f3577af1979c9b9e1ac2cad0104ee +TORCH_CUDA_ARCH_LIST="9.0" python setup.py install +cd .. ``` -Use a dedicated env, kept separate from the vLLM env so kernel pins don't clash: +Foundry auto-detects `libcuda_hook.so` and +`site-packages/nvidia/nvshmem/lib/libnvshmem_host.so.3` before worker spawn. +Use `hook_library_path` or `nvshmem_host_path` in both TOMLs only when +auto-detection does not fit the environment. + +## Archive compatibility + +Archives created by earlier SGLang integrations are incompatible. Delete the +entire workspace and perform a new SAVE after upgrading: ```bash -conda create -p venv/ python=3.12 -conda activate venv/ -conda install -c conda-forge boost-cpp boost # foundry C++ deps - -# in-tree sglang fork, editable -pip install -e sglang/python --extra-index-url https://download.pytorch.org/whl/cu130 -# foundry -pushd foundry && pip install -e . --no-build-isolation && popd +rm -rf foundry_archive ``` -`libcuda_hook.so` finds boost via a baked rpath; if it can't, add the conda lib dir to -`LD_LIBRARY_PATH` (`export LD_LIBRARY_PATH=$CONDA_PREFIX/lib:$LD_LIBRARY_PATH`). +LOAD validates the archive schema, archived SGLang version, installed SGLang +version, and resolved graph configuration before applying the saved memory-pool +configuration. Both phases must resolve to decode `full` and prefill +`disabled`. + +## Baseline, SAVE, LOAD, query -## Run (single GPU / DP) +Run from this directory. Each script accepts an optional trailing `--save` or +`--load`; omitting it starts a baseline server. ```bash -# single GPU +cd foundry/recipe/sglang + +# 1. Baseline. Wait for startup, optionally query it, then stop the process. +bash serve_qwen3-mini.sh + +# 2. Fresh SAVE. Remove every older archive first. rm -rf foundry_archive -bash serve_qwen3-mini.sh --save # wait for "Application startup complete", then SIGTERM -bash serve_qwen3-mini.sh --load # leave running +bash serve_qwen3-mini.sh --save +# Wait for "Application startup complete", then stop the process. -# query (separate shell) -curl -s http://0.0.0.0:12000/v1/completions -H 'Content-Type: application/json' \ +# 3. Fresh-process LOAD. Do not reuse the SAVE process. +bash serve_qwen3-mini.sh --load + +# 4. Query LOAD from another shell. +curl -s http://127.0.0.1:12000/v1/completions \ + -H 'Content-Type: application/json' \ -d '{"model":"Qwen/Qwen3-1.7B","prompt":"The capital of France is","max_tokens":12,"temperature":0}' +``` -# data parallel (pick GPUs with CUDA_VISIBLE_DEVICES) +Regular DP: + +```bash rm -rf foundry_archive CUDA_VISIBLE_DEVICES=0,1 bash serve_qwen3-1.7b_dp.sh 2 --save +# Stop SAVE after startup. CUDA_VISIBLE_DEVICES=0,1 bash serve_qwen3-1.7b_dp.sh 2 --load ``` -## Run (expert parallel / DeepEP) +Exact dense TP=2 ร— DP=2: -EP needs three kernel packages โ€” all SGLang-native, no vLLM involved: +```bash +# 1. Baseline. After startup, run the routed queries below, then stop it. +CUDA_VISIBLE_DEVICES=0,1,2,3 bash serve_qwen3-8b_tp2_dp2.sh + +# 2. Fresh SAVE using the dedicated workspace/config. +rm -rf foundry_archive_sglang_qwen3_8b_tp2_dp2 +CUDA_VISIBLE_DEVICES=0,1,2,3 bash serve_qwen3-8b_tp2_dp2.sh --save +# After startup, run the routed queries below, then stop SAVE. + +# 3. Fresh-process LOAD. +CUDA_VISIBLE_DEVICES=0,1,2,3 bash serve_qwen3-8b_tp2_dp2.sh --load + +# 4. Run after each ready phase; compare DP0 and DP1 across all three. +for dp_rank in 0 1; do + curl -s http://127.0.0.1:12000/generate \ + -H 'Content-Type: application/json' \ + -d "{\"text\":\"The capital of France is\",\"sampling_params\":{\"temperature\":0,\"max_new_tokens\":12},\"routed_dp_rank\":${dp_rank}}" +done +``` -- **NVSHMEM** โ€” already in the env. cu13 `torch` pulls the `nvidia-nvshmem-cuXX` - wheel as a dependency (`libnvshmem_host.so.3` under `site-packages/nvidia/nvshmem/lib/`). - Foundry auto-detects it from the wheel (just like `libcuda_hook.so`) and the - spawn-site patches preload it into each worker โ€” no manual path, no TOML field. -- **DeepEP** @ `9af0e0d` โ€” SGLang's pin. Build via SGLang's own installer - `sglang/scripts/ci/cuda/ci_install_deepep.sh` (it `git checkout`s exactly `9af0e0d` - and builds against the NVSHMEM wheel above). For a single node you can skip the - script's gdrcopy/RDMA apt steps and just build the wheel: +The TP launcher fixes `--tp-size 2 --dp-size 2 --pp-size 1 --ep-size 1`, +`--attention-backend flashinfer`, full decode, disabled prefill, +`--disable-custom-all-reduce`, `NCCL_CUMEM_ENABLE=0`, and +`NCCL_NVLS_ENABLE=0` in baseline, SAVE, and LOAD. Do not remove or generalize +these safety constraints. - ```bash - git clone https://github.com/deepseek-ai/DeepEP.git && cd DeepEP - git checkout 9af0e0d0e74f3577af1979c9b9e1ac2cad0104ee - TORCH_CUDA_ARCH_LIST="9.0" python setup.py install # Hopper; "9.0;10.0" for Blackwell - cd .. - ``` +SGLang uses CUDA IPC to share tensors between the TP scheduler processes. +Foundry's VMM-IPC bridge exports the whole physical VMM chunk, preserves each +allocation's interior offset, and transfers the live file descriptor over a +same-UID abstract Unix socket with `SCM_RIGHTS`; it never treats a raw +file-descriptor integer as portable across processes. The validated mapping is +DP0/TP0 โ†’ `rank_0`, DP0/TP1 โ†’ `rank_1`, DP1/TP0 โ†’ `rank_2`, and DP1/TP1 โ†’ +`rank_3`. Each rank saved and restored 36 graphs at final offset +`69822578688`. -- **`sgl-deep-gemm >= 0.1.2`** (0.1.0 lacks `m_grouped_bf16_gemm_nt_masked`) and - **`flash-attn-3`** (the fa3 attention backend โ€” flashinfer's ragged-prefill path - has an off-by-one under this config). +DP-attention with DeepEP (EP=2 or EP=4): ```bash +EP_SIZE=4 +CUDA_VISIBLE_DEVICES=0,1,2,3 bash serve_qwen3-30ba3bfp8_ep.sh "$EP_SIZE" +# Stop the baseline after query. rm -rf foundry_archive -CUDA_VISIBLE_DEVICES=0,1 bash serve_qwen3-30ba3bfp8_ep.sh 2 --save -CUDA_VISIBLE_DEVICES=0,1 bash serve_qwen3-30ba3bfp8_ep.sh 2 --load -curl -s http://0.0.0.0:12000/v1/completions -H 'Content-Type: application/json' \ +CUDA_VISIBLE_DEVICES=0,1,2,3 bash serve_qwen3-30ba3bfp8_ep.sh "$EP_SIZE" --save +# Stop SAVE after startup. +CUDA_VISIBLE_DEVICES=0,1,2,3 bash serve_qwen3-30ba3bfp8_ep.sh "$EP_SIZE" --load +curl -s http://127.0.0.1:12000/v1/completions \ + -H 'Content-Type: application/json' \ -d '{"model":"Qwen/Qwen3-30B-A3B-FP8","prompt":"The capital of France is","max_tokens":12,"temperature":0}' ``` -The EP script sets `--enable-dp-attention --moe-a2a-backend deepep --deepep-mode -low_latency --moe-runner-backend deep_gemm --attention-backend fa3 ---disable-custom-all-reduce` and `SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK=256`. -DeepEP low-latency caps dispatch at that per-rank token count (and asserts -`(n+1)*2 <= NVSHMEM_QP_DEPTH`); keep it and `--chunked-prefill-size` identical between -SAVE and LOAD so the captured graphs match. +Set `EP_SIZE=2` and select two GPUs for the validated EP=2 variant. EP=4's +`tp_size=4` is part of SGLang's DP-attention/EP process topology and does not +mean conventional TP attention. + +The regular-DP and EP scripts set `NCCL_CUMEM_ENABLE=0` and +`NCCL_NVLS_ENABLE=0` for Foundry runs; the dense TP launcher sets both in all +three phases. The DeepEP script also keeps +`SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK=256` and +`--chunked-prefill-size 256` identical on SAVE and LOAD. ## Archive layout -``` +```text foundry_archive/ -โ”œโ”€โ”€ warmup_state.json # KV-block sizing + MemoryPoolConfig (rank 0) +โ”œโ”€โ”€ warmup_state.json โ””โ”€โ”€ rank_/ - โ”œโ”€โ”€ graph_*.json + .cugraph # one pair per captured graph - โ”œโ”€โ”€ graph_manifest.json # topology groups + template assignments - โ”œโ”€โ”€ fatbin_image_packed.img # packed CUDA modules - โ””โ”€โ”€ final_alloc_offset.json # per-rank VMM watermark + โ”œโ”€โ”€ sglang_graph_index.json + โ”œโ”€โ”€ graph_000000_sglang.json + โ”œโ”€โ”€ graph_000000_sglang.cugraph + โ”œโ”€โ”€ graph_manifest.json + โ”œโ”€โ”€ fatbin_image_packed.img + โ”œโ”€โ”€ fatbin_entrypoint_packed.txt + โ””โ”€โ”€ final_alloc_offset.json ``` -For DP / EP each rank gets its own `rank_/`. +`warmup_state.json` schema 1 records `schema_version`, `sglang_version`, +`decode_cuda_graph_backend`, `prefill_cuda_graph_backend`, `timestamp`, +`cuda_version`, `gpu_name`, `gpu_total_memory`, and `memory_pool_config`. +Each rank's `final_alloc_offset.json` stores its own VMM watermark. + +`sglang_graph_index.json` schema 1 is the authoritative graph list. Each entry +records capture order, phase, filename, output kind, and the complete +`ShapeKey(size, stream_idx, variant_label)`. Filenames are opaque; LOAD does +not derive graph identity from them. `graph_manifest.json` remains the Foundry +topology/template manifest used by the batched native loader. + +For regular DP, exact TP=2 ร— DP=2, and EP, each rank owns a separate +`rank_/` directory. The dense TP recipe writes the same layout under its +dedicated `foundry_archive_sglang_qwen3_8b_tp2_dp2` root. ## Troubleshooting -| Symptom | Likely cause | +| Symptom | Action | |---|---| -| `Reserved address โ€ฆ != requested base 0x600000000000` | VMM base collided with another allocation. Re-run; non-deterministic, the next run usually succeeds. | -| EP replay `illegal memory access` / `nvshmemx_cumodule_init not found` | `libnvshmem_host.so.3` not preloaded โ€” foundry couldn't auto-detect the `nvidia-nvshmem` wheel. Confirm it's installed (`pip show nvidia-nvshmem-cu13`), or set `nvshmem_host_path` in both TOMLs. | -| `nvshmem_qp_depth >= (num_max_dispatch_tokens_per_rank + 1) * 2` | `SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK` too high for `NVSHMEM_QP_DEPTH`; lower it or raise the QP depth. | +| Package version is `unknown` or not `0.5.15.post1` | Install `sglang/python` with pip; do not rely on a source-path-only checkout. | +| Archive schema, version, or backend mismatch | Delete `foundry_archive` and perform a fresh SAVE with the exact environment and flags. | +| Dense TP topology/collective validation fails | Use exactly TP=2, DP=2, PP=1, EP=1 and retain disabled custom all-reduce plus both NCCL environment guards. Other TP topologies are unsupported. | +| Reserved VMM address differs from `0x600000000000` | Another allocation owns the requested range; restart the process. | +| DeepEP reports missing NVSHMEM symbols | Verify the cu13 NVSHMEM wheel or set `nvshmem_host_path` in both TOMLs. | +| DeepEP token/QP-depth assertion fails | Lower `SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK` or raise `NVSHMEM_QP_DEPTH`, identically for SAVE and LOAD. | diff --git a/recipe/sglang/foundry_tp2_dp2_load.toml b/recipe/sglang/foundry_tp2_dp2_load.toml new file mode 100644 index 00000000..fa256abf --- /dev/null +++ b/recipe/sglang/foundry_tp2_dp2_load.toml @@ -0,0 +1,5 @@ +mode = "load" +base_addr = 0x600000000000 +region_size = "256GB" +workspace_root = "foundry_archive_sglang_qwen3_8b_tp2_dp2" +scratch_space_size = "1024MB" diff --git a/recipe/sglang/foundry_tp2_dp2_save.toml b/recipe/sglang/foundry_tp2_dp2_save.toml new file mode 100644 index 00000000..7410cef3 --- /dev/null +++ b/recipe/sglang/foundry_tp2_dp2_save.toml @@ -0,0 +1,5 @@ +mode = "save" +base_addr = 0x600000000000 +region_size = "256GB" +workspace_root = "foundry_archive_sglang_qwen3_8b_tp2_dp2" +scratch_space_size = "1024MB" diff --git a/recipe/sglang/serve_qwen3-1.7b_dp.sh b/recipe/sglang/serve_qwen3-1.7b_dp.sh index c562a5ac..85e938fe 100755 --- a/recipe/sglang/serve_qwen3-1.7b_dp.sh +++ b/recipe/sglang/serve_qwen3-1.7b_dp.sh @@ -36,8 +36,8 @@ fi # LD_PRELOAD of libcuda_hook.so is set by foundry's setup_ld_preload_env at # worker spawn time (path auto-detected; propagated to the DP controller and -# every rank's scheduler child). Assumes foundry + the sglang fork are -# pip-installed (see README) so both import without PYTHONPATH. +# every rank's scheduler child). Assumes foundry + the patched upstream SGLang +# checkout are pip-installed (see README). sglang serve \ --model-path "$MODEL_NAME" \ @@ -48,5 +48,6 @@ sglang serve \ --mem-fraction-static "$MEM_FRACTION_STATIC" \ --disable-radix-cache \ --attention-backend flashinfer \ - --cuda-graph-max-bs 512 \ + --cuda-graph-max-bs-decode 512 \ + --disable-prefill-cuda-graph \ "${FOUNDRY_ARGS[@]}" diff --git a/recipe/sglang/serve_qwen3-30ba3bfp8_ep.sh b/recipe/sglang/serve_qwen3-30ba3bfp8_ep.sh index a185fda3..c0ccca21 100755 --- a/recipe/sglang/serve_qwen3-30ba3bfp8_ep.sh +++ b/recipe/sglang/serve_qwen3-30ba3bfp8_ep.sh @@ -2,8 +2,8 @@ # Qwen3-30B-A3B-FP8 (MoE), expert parallel via DeepEP low-latency + DP-attention. # Usage: CUDA_VISIBLE_DEVICES=0,1 bash serve_qwen3-30ba3bfp8_ep.sh [--save|--load] # -# Requires (see README ยงEP): deep_ep @ 9af0e0d, sgl-deep-gemm >= 0.1.2, flash-attn-3, -# and the nvshmem_host_path uncommented in foundry_{save,load}.toml. +# Requires (see README ยงEP): deep_ep @ 9af0e0d, sgl-deep-gemm 0.1.4, the +# SGLang FA3 package, and NVSHMEM (auto-detected or configured in the TOMLs). SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" @@ -42,7 +42,7 @@ fi # (the hook from the foundry install, NVSHMEM from the nvidia-nvshmem wheel via # config._detect_nvshmem_host_path), so nothing is preloaded by the shell. Set # nvshmem_host_path in the TOML only to override the NVSHMEM auto-detection. -# Assumes foundry + the sglang fork are pip-installed (see README). +# Assumes foundry + the patched upstream SGLang checkout are pip-installed. # DeepEP low-latency caps tokens/rank at SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK # (default 128; (n+1)*2 must be <= NVSHMEM_QP_DEPTH). Raise to 256 ((256+1)*2=514 @@ -66,5 +66,6 @@ sglang serve \ --disable-custom-all-reduce \ --chunked-prefill-size 256 \ --attention-backend fa3 \ - --cuda-graph-max-bs 128 \ + --cuda-graph-max-bs-decode 128 \ + --disable-prefill-cuda-graph \ "${FOUNDRY_ARGS[@]}" diff --git a/recipe/sglang/serve_qwen3-8b_tp2_dp2.sh b/recipe/sglang/serve_qwen3-8b_tp2_dp2.sh new file mode 100755 index 00000000..806fab71 --- /dev/null +++ b/recipe/sglang/serve_qwen3-8b_tp2_dp2.sh @@ -0,0 +1,45 @@ +#!/bin/bash +# Qwen3-8B, two TP ranks in each of two regular-DP replicas. +# Usage: CUDA_VISIBLE_DEVICES=0,1,2,3 bash serve_qwen3-8b_tp2_dp2.sh [--save|--load] + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" + +# Keep Foundry's validated NCCL transport constraints identical in baseline, +# SAVE, and LOAD so phase comparisons do not change collective behavior. +export NCCL_CUMEM_ENABLE=0 +export NCCL_NVLS_ENABLE=0 + +FOUNDRY_ARGS=() +if [[ "$1" == "--save" ]]; then + FOUNDRY_TOML="${SCRIPT_DIR}/foundry_tp2_dp2_save.toml" + FOUNDRY_ARGS+=( --foundry-graph-extension-config-path "${FOUNDRY_TOML}" ) + echo "Using foundry SAVE: ${FOUNDRY_TOML}" +elif [[ "$1" == "--load" ]]; then + FOUNDRY_TOML="${SCRIPT_DIR}/foundry_tp2_dp2_load.toml" + FOUNDRY_ARGS+=( --foundry-graph-extension-config-path "${FOUNDRY_TOML}" ) + echo "Using foundry LOAD: ${FOUNDRY_TOML}" +elif [[ -n "$1" ]]; then + echo "Usage: $0 [--save|--load]" + exit 1 +else + echo "Running without foundry (baseline SGLang)" +fi + +sglang serve \ + --model-path Qwen/Qwen3-8B \ + --trust-remote-code \ + --host 0.0.0.0 \ + --port 12000 \ + --tp-size 2 \ + --dp-size 2 \ + --pp-size 1 \ + --ep-size 1 \ + --load-balance-method round_robin \ + --mem-fraction-static 0.8 \ + --disable-radix-cache \ + --disable-custom-all-reduce \ + --attention-backend flashinfer \ + --cuda-graph-backend-decode full \ + --cuda-graph-max-bs-decode 256 \ + --disable-prefill-cuda-graph \ + "${FOUNDRY_ARGS[@]}" diff --git a/recipe/sglang/serve_qwen3-mini.sh b/recipe/sglang/serve_qwen3-mini.sh index 77449482..5262bf83 100755 --- a/recipe/sglang/serve_qwen3-mini.sh +++ b/recipe/sglang/serve_qwen3-mini.sh @@ -29,8 +29,8 @@ fi # LD_PRELOAD of libcuda_hook.so is set by foundry's setup_ld_preload_env at # worker spawn time (path auto-detected from the foundry install). Baseline runs -# don't need it preloaded by the shell. Assumes foundry + the sglang fork are -# pip-installed (see README) so both import without PYTHONPATH. +# don't need it preloaded by the shell. Assumes foundry + the patched upstream +# SGLang checkout are pip-installed (see README). sglang serve \ --model-path "$MODEL_NAME" \ @@ -40,5 +40,6 @@ sglang serve \ --mem-fraction-static "$MEM_FRACTION_STATIC" \ --disable-radix-cache \ --attention-backend flashinfer \ - --cuda-graph-max-bs 512 \ + --cuda-graph-max-bs-decode 512 \ + --disable-prefill-cuda-graph \ "${FOUNDRY_ARGS[@]}" diff --git a/tests/test_imports.py b/tests/test_imports.py index 2de75a6e..c83c3c7c 100644 --- a/tests/test_imports.py +++ b/tests/test_imports.py @@ -13,6 +13,10 @@ """ import importlib +import importlib.util +from importlib.metadata import PackageNotFoundError, version + +import pytest PUBLIC_MODULES = [ # Top-level @@ -76,9 +80,22 @@ def test_vllm_integration_public_api(): def test_sglang_integration_submodules_importable(): - # SGLang side currently exposes its API via submodules rather than - # re-exports on `foundry.integration.sglang`. Just verify each submodule - # imports cleanly โ€” a ruff --fix that stripped a critical import would - # break this. for sub in ("config", "hooks", "runtime", "graph_ops"): importlib.import_module(f"foundry.integration.sglang.{sub}") + + +def test_sglang_cuda_graph_importable_when_sglang_is_installed(): + if importlib.util.find_spec("sglang") is None: + pytest.skip("SGLang is an optional dependency") + try: + version("sglang") + except PackageNotFoundError: + pytest.skip("SGLang distribution metadata is unavailable") + + importlib.import_module("foundry.integration.sglang.cuda_graph") + + +def test_sglang_integration_public_api(): + mod = importlib.import_module("foundry.integration.sglang") + for name in ("install_hooks", "CUDAGraphExtensionMode", "get_graph_extension_mode"): + assert hasattr(mod, name), f"foundry.integration.sglang.{name} not exported" diff --git a/tests/test_sglang_api_contract.py b/tests/test_sglang_api_contract.py new file mode 100644 index 00000000..b63849c4 --- /dev/null +++ b/tests/test_sglang_api_contract.py @@ -0,0 +1,97 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the Foundry project +"""Import/signature contract for upstream SGLang v0.5.15.post1.""" + +from __future__ import annotations + +import importlib +import importlib.util +import inspect +from importlib.metadata import PackageNotFoundError, version + +import pytest + + +def _require_sglang(): + if importlib.util.find_spec("sglang") is None: + pytest.skip("SGLang is an optional dependency") + try: + installed = version("sglang") + except PackageNotFoundError: + pytest.skip("SGLang package metadata is unavailable") + assert installed == "0.5.15.post1" + + +def _parameters(callable_object): + return inspect.signature(callable_object).parameters + + +def test_latest_runner_and_backend_symbols_and_signatures(): + _require_sglang() + base_runner = importlib.import_module("sglang.srt.model_executor.runner.base_runner").BaseRunner + decode_runner = importlib.import_module( + "sglang.srt.model_executor.runner.decode_cuda_graph_runner" + ).DecodeCudaGraphRunner + full_backend = importlib.import_module( + "sglang.srt.model_executor.runner_backend.full_cuda_graph_backend" + ).FullCudaGraphBackend + + assert list(_parameters(base_runner.warmup)) == ["self"] + assert list(_parameters(decode_runner.capture)) == ["self"] + assert list(_parameters(decode_runner.capture_prepare)) == [ + "self", + "size", + "stream_idx", + ] + assert list(_parameters(full_backend.capture_one)) == [ + "self", + "shape_key", + "forward_fn", + "dummies", + "post_warmup_hook", + ] + + +def test_shape_key_has_complete_latest_identity(): + _require_sglang() + shape_key = importlib.import_module("sglang.srt.model_executor.runner.shape_key").ShapeKey + + parameters = _parameters(shape_key) + assert list(parameters) == ["size", "stream_idx", "variant_label"] + assert parameters["stream_idx"].default is None + assert parameters["variant_label"].default is None + assert shape_key(8, 1, "nolora") != shape_key(8, 0, "nolora") + + +def test_attention_backend_uses_split_graph_metadata_contract(): + _require_sglang() + attention_backend = importlib.import_module( + "sglang.srt.layers.attention.base_attn_backend" + ).AttentionBackend + + assert list(_parameters(attention_backend.init_forward_metadata_out_graph)) == [ + "self", + "forward_batch", + "in_capture", + ] + assert ( + _parameters(attention_backend.init_forward_metadata_out_graph)["in_capture"].default + is False + ) + assert list(_parameters(attention_backend.init_forward_metadata_in_graph)) == [ + "self", + "forward_batch", + ] + + +def test_deep_gemm_compile_hook_is_available(): + _require_sglang() + compile_utils = importlib.import_module("sglang.srt.layers.deep_gemm_wrapper.compile_utils") + + assert callable(compile_utils._compile_deep_gemm_one_type_all) + + +def test_removed_monolithic_cuda_graph_runner_is_not_part_of_contract(): + _require_sglang() + + assert importlib.util.find_spec("sglang.srt.model_executor.cuda_graph_runner") is None diff --git a/tests/test_sglang_config.py b/tests/test_sglang_config.py new file mode 100644 index 00000000..0f99afae --- /dev/null +++ b/tests/test_sglang_config.py @@ -0,0 +1,100 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the Foundry project +"""Rank-mapping contracts for hybrid SGLang TP and DP workers.""" + +import importlib.util +import sys +from pathlib import Path +from types import ModuleType, SimpleNamespace + +import pytest + +SGLANG_INTEGRATION_DIR = Path(__file__).parents[1] / "python" / "foundry" / "integration" / "sglang" + + +def _load_module(name: str, path: Path): + spec = importlib.util.spec_from_file_location(name, path) + assert spec is not None + assert spec.loader is not None + module = importlib.util.module_from_spec(spec) + sys.modules[name] = module + spec.loader.exec_module(module) + return module + + +def _module(name: str, **attributes): + module = ModuleType(name) + for attribute, value in attributes.items(): + setattr(module, attribute, value) + return module + + +def _load_hooks(monkeypatch: pytest.MonkeyPatch): + foundry = _module("foundry") + foundry.__path__ = [] + integration = _module("foundry.integration") + integration.__path__ = [] + sglang = _module("foundry.integration.sglang") + sglang.__path__ = [] + runtime = _module("foundry.integration.sglang.runtime") + config = _module( + "foundry.integration.sglang.config", + CUDAGraphExtensionMode=SimpleNamespace(NONE="none"), + get_graph_extension_mode=lambda: "none", + get_workspace_root=lambda: None, + load_graph_extension_config=lambda _path: None, + ) + for name, module in { + "torch": _module("torch"), + "foundry": foundry, + "foundry.integration": integration, + "foundry.integration.sglang": sglang, + "foundry.integration.sglang.runtime": runtime, + "foundry.integration.sglang.config": config, + }.items(): + monkeypatch.setitem(sys.modules, name, module) + return _load_module( + "foundry_sglang_hooks_rank_test", + SGLANG_INTEGRATION_DIR / "hooks.py", + ) + + +@pytest.mark.parametrize( + ("dp_rank", "tp_rank", "expected"), + [ + (0, 0, 0), + (0, 1, 1), + (1, 0, 2), + (1, 1, 3), + ], +) +def test_tp2_dp2_workers_have_distinct_workspace_ranks(monkeypatch, dp_rank, tp_rank, expected): + config = _load_module( + "foundry_sglang_config_rank_test", + SGLANG_INTEGRATION_DIR / "config.py", + ) + hooks = _load_hooks(monkeypatch) + server_args = SimpleNamespace( + tp_size=2, + pp_size=1, + dp_size=2, + enable_dp_attention=False, + ) + model_runner = SimpleNamespace( + server_args=server_args, + tp_rank=tp_rank, + dp_rank=dp_rank, + ) + + resolved_dp_rank = hooks._resolve_dp_rank(model_runner) + + assert resolved_dp_rank == dp_rank + assert ( + config.compute_workspace_rank( + server_args, + tp_rank=tp_rank, + pp_rank=0, + dp_rank=resolved_dp_rank, + ) + == expected + ) diff --git a/tests/test_sglang_cuda_graph.py b/tests/test_sglang_cuda_graph.py new file mode 100644 index 00000000..633fdbb7 --- /dev/null +++ b/tests/test_sglang_cuda_graph.py @@ -0,0 +1,560 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the Foundry project +"""Pure-Python contracts for the SGLang v0.5.15.post1 graph hooks.""" + +from __future__ import annotations + +import importlib +import sys +from contextlib import contextmanager +from dataclasses import dataclass +from enum import Enum +from pathlib import Path +from types import ModuleType, SimpleNamespace + +import pytest + + +class Mode(str, Enum): + NONE = "none" + SAVE = "save" + LOAD = "load" + + +@dataclass(frozen=True) +class ShapeKey: + size: int + stream_idx: int | None = None + variant_label: str | None = None + + +class BaseRunner: + def __init__(self): + self.warmup_calls = 0 + + def warmup(self): + self.warmup_calls += 1 + return "warmed" + + +class DecodeCudaGraphRunner(BaseRunner): + """Small latest-API runner; deliberately has no legacy CudaGraphRunner API.""" + + def __init__(self): + super().__init__() + self.capture_calls = 0 + self.prepared = [] + + def capture_prepare(self, size, stream_idx=None): + key = ShapeKey(size=size, stream_idx=stream_idx) + self.prepared.append(key) + return key + + def capture(self): + self.capture_calls += 1 + return "captured" + + +class FullCudaGraphBackend: + def __init__(self): + self._graphs = {} + self._outputs = {} + self._pool = None + self.capture_one_calls = 0 + + @contextmanager + def capture_session(self, stream): + del stream + yield + + def capture_one( + self, + shape_key, + forward_fn, + dummies=None, + post_warmup_hook=None, + ): + del dummies, post_warmup_hook + self.capture_one_calls += 1 + self._graphs[shape_key] = object() + self._outputs[shape_key] = forward_fn() + + +_ORIGINAL_WARMUP = BaseRunner.warmup +_ORIGINAL_CAPTURE = DecodeCudaGraphRunner.capture +_ORIGINAL_CAPTURE_PREPARE = DecodeCudaGraphRunner.capture_prepare +_ORIGINAL_CAPTURE_ONE = FullCudaGraphBackend.capture_one + + +def _module(name: str, **attributes): + module = ModuleType(name) + for attr, value in attributes.items(): + setattr(module, attr, value) + return module + + +def _load_cuda_graph(monkeypatch: pytest.MonkeyPatch): + BaseRunner.warmup = _ORIGINAL_WARMUP + DecodeCudaGraphRunner.capture = _ORIGINAL_CAPTURE + DecodeCudaGraphRunner.capture_prepare = _ORIGINAL_CAPTURE_PREPARE + FullCudaGraphBackend.capture_one = _ORIGINAL_CAPTURE_ONE + + source_root = Path(__file__).parents[1] / "python" / "foundry" + mode = SimpleNamespace(value=Mode.NONE) + registered = SimpleNamespace(global_pool=None, pynccl_pool=None) + + foundry = _module("foundry") + foundry.__path__ = [str(source_root)] + integration = _module("foundry.integration") + integration.__path__ = [str(source_root / "integration")] + sglang_integration = _module("foundry.integration.sglang") + sglang_integration.__path__ = [str(source_root / "integration" / "sglang")] + + config = _module( + "foundry.integration.sglang.config", + CUDAGraphExtensionMode=Mode, + get_graph_extension_mode=lambda: mode.value, + ) + graph_ops = _module("foundry.integration.sglang.graph_ops") + + runner_package = _module("sglang.srt.model_executor.runner") + runner_package.__path__ = [] + base_runner = _module( + "sglang.srt.model_executor.runner.base_runner", + BaseRunner=BaseRunner, + ) + decode_runner = _module( + "sglang.srt.model_executor.runner.decode_cuda_graph_runner", + DecodeCudaGraphRunner=DecodeCudaGraphRunner, + ) + shape_key = _module( + "sglang.srt.model_executor.runner.shape_key", + ShapeKey=ShapeKey, + ) + full_backend = _module( + "sglang.srt.model_executor.runner_backend.full_cuda_graph_backend", + FullCudaGraphBackend=FullCudaGraphBackend, + ) + pool = _module( + "sglang.srt.model_executor.runner_utils.pool", + set_global_graph_memory_pool=lambda value: setattr(registered, "global_pool", value), + ) + pynccl = _module( + "sglang.srt.distributed.device_communicators.pynccl_allocator", + set_graph_pool_id=lambda value: setattr(registered, "pynccl_pool", value), + ) + + modules = { + "foundry": foundry, + "foundry.integration": integration, + "foundry.integration.sglang": sglang_integration, + "foundry.integration.sglang.config": config, + "foundry.integration.sglang.graph_ops": graph_ops, + "sglang": _module("sglang"), + "sglang.srt": _module("sglang.srt"), + "sglang.srt.model_executor": _module("sglang.srt.model_executor"), + "sglang.srt.model_executor.runner": runner_package, + "sglang.srt.model_executor.runner.base_runner": base_runner, + "sglang.srt.model_executor.runner.decode_cuda_graph_runner": decode_runner, + "sglang.srt.model_executor.runner.shape_key": shape_key, + "sglang.srt.model_executor.runner_backend": _module( + "sglang.srt.model_executor.runner_backend" + ), + "sglang.srt.model_executor.runner_backend.full_cuda_graph_backend": full_backend, + "sglang.srt.model_executor.runner_utils": _module("sglang.srt.model_executor.runner_utils"), + "sglang.srt.model_executor.runner_utils.pool": pool, + "sglang.srt.distributed": _module("sglang.srt.distributed"), + "sglang.srt.distributed.device_communicators": _module( + "sglang.srt.distributed.device_communicators" + ), + "sglang.srt.distributed.device_communicators.pynccl_allocator": pynccl, + } + for name, module in modules.items(): + monkeypatch.setitem(sys.modules, name, module) + + cuda_graph = importlib.import_module("foundry.integration.sglang.cuda_graph") + return cuda_graph, mode, registered + + +def _server_args(**overrides): + values = { + "cuda_graph_config": SimpleNamespace( + decode=SimpleNamespace(backend="full"), + prefill=SimpleNamespace(backend="disabled"), + ), + "tp_size": 1, + "pp_size": 1, + "dp_size": 1, + "ep_size": 1, + "enable_dp_attention": False, + "attention_backend": "flashinfer", + "moe_a2a_backend": "none", + "disable_custom_all_reduce": False, + "enable_nccl_nvls": False, + "enable_symm_mem": False, + "enable_torch_symm_mem": False, + "speculative_algorithm": None, + "enable_lora": False, + "enable_pdmux": False, + "dllm_algorithm": None, + "enable_return_hidden_states": False, + "enable_memory_saver": False, + } + values.update(overrides) + return SimpleNamespace(**values) + + +@pytest.mark.parametrize( + "server_args", + [ + _server_args(), + _server_args(dp_size=2), + _server_args( + tp_size=2, + dp_size=2, + ep_size=2, + enable_dp_attention=True, + attention_backend="fa3", + moe_a2a_backend="deepep", + ), + _server_args( + tp_size=4, + dp_size=4, + ep_size=4, + enable_dp_attention=True, + attention_backend="fa3", + moe_a2a_backend="deepep", + ), + ], +) +def test_configuration_accepts_supported_full_decode_modes(monkeypatch, server_args): + cuda_graph, _mode, _registered = _load_cuda_graph(monkeypatch) + + cuda_graph.validate_configuration(server_args) + + +def test_configuration_accepts_exact_tp2_dp2_full_decode(monkeypatch): + monkeypatch.setenv("NCCL_CUMEM_ENABLE", "0") + monkeypatch.setenv("NCCL_NVLS_ENABLE", "0") + cuda_graph, _mode, _registered = _load_cuda_graph(monkeypatch) + + cuda_graph.validate_configuration( + _server_args( + tp_size=2, + dp_size=2, + ep_size=1, + disable_custom_all_reduce=True, + ) + ) + + +@pytest.mark.parametrize( + "overrides", + [ + {"tp_size": 2, "dp_size": 1}, + {"tp_size": 2, "dp_size": 3}, + {"tp_size": 2, "dp_size": 2, "ep_size": 2}, + {"tp_size": 4, "dp_size": 2}, + ], +) +def test_configuration_rejects_nearby_tp_topologies(monkeypatch, overrides): + monkeypatch.setenv("NCCL_CUMEM_ENABLE", "0") + monkeypatch.setenv("NCCL_NVLS_ENABLE", "0") + cuda_graph, _mode, _registered = _load_cuda_graph(monkeypatch) + + with pytest.raises( + ValueError, + match=r"tp_size=2, dp_size=2, pp_size=1, ep_size=1", + ): + cuda_graph.validate_configuration(_server_args(**overrides)) + + +@pytest.mark.parametrize( + ("overrides", "diagnostic"), + [ + ({"disable_custom_all_reduce": False}, "disable.custom.all.reduce"), + ({"enable_nccl_nvls": True}, "NVLS"), + ({"enable_symm_mem": True}, "symmetric memory"), + ({"enable_torch_symm_mem": True}, "symmetric memory"), + ], +) +def test_configuration_rejects_unsafe_tp2_dp2_collectives(monkeypatch, overrides, diagnostic): + monkeypatch.setenv("NCCL_CUMEM_ENABLE", "0") + monkeypatch.setenv("NCCL_NVLS_ENABLE", "0") + cuda_graph, _mode, _registered = _load_cuda_graph(monkeypatch) + values = { + "tp_size": 2, + "dp_size": 2, + "disable_custom_all_reduce": True, + } + values.update(overrides) + server_args = _server_args(**values) + + with pytest.raises(ValueError, match=diagnostic): + cuda_graph.validate_configuration(server_args) + + +@pytest.mark.parametrize( + ("variable", "value"), + [ + ("NCCL_CUMEM_ENABLE", None), + ("NCCL_CUMEM_ENABLE", "1"), + ("NCCL_NVLS_ENABLE", None), + ("NCCL_NVLS_ENABLE", "1"), + ], +) +def test_configuration_requires_disabled_nccl_vmm_fast_paths(monkeypatch, variable, value): + monkeypatch.setenv("NCCL_CUMEM_ENABLE", "0") + monkeypatch.setenv("NCCL_NVLS_ENABLE", "0") + if value is None: + monkeypatch.delenv(variable) + else: + monkeypatch.setenv(variable, value) + cuda_graph, _mode, _registered = _load_cuda_graph(monkeypatch) + + with pytest.raises(ValueError, match=variable): + cuda_graph.validate_configuration( + _server_args( + tp_size=2, + dp_size=2, + disable_custom_all_reduce=True, + ) + ) + + +@pytest.mark.parametrize( + ("overrides", "diagnostic"), + [ + ( + { + "cuda_graph_config": SimpleNamespace( + decode=SimpleNamespace(backend="breakable"), + prefill=SimpleNamespace(backend="disabled"), + ) + }, + "decode.*full", + ), + ( + { + "cuda_graph_config": SimpleNamespace( + decode=SimpleNamespace(backend="tc_piecewise"), + prefill=SimpleNamespace(backend="disabled"), + ) + }, + "decode.*full", + ), + ( + { + "cuda_graph_config": SimpleNamespace( + decode=SimpleNamespace(backend="disabled"), + prefill=SimpleNamespace(backend="disabled"), + ) + }, + "decode.*full", + ), + ( + { + "cuda_graph_config": SimpleNamespace( + decode=SimpleNamespace(backend="full"), + prefill=SimpleNamespace(backend="full"), + ) + }, + "prefill.*disabled", + ), + ({"tp_size": 2}, "tensor parallel|tp_size"), + ({"pp_size": 2}, "pipeline parallel|pp_size"), + ({"speculative_algorithm": "EAGLE"}, "speculative"), + ({"enable_lora": True}, "LoRA"), + ({"enable_pdmux": True}, "PDMux"), + ({"dllm_algorithm": "dream"}, "dLLM|diffusion"), + ({"enable_return_hidden_states": True}, "hidden"), + ({"enable_memory_saver": True}, "memory.saver"), + ({"attention_backend": "triton"}, "FlashInfer|FA3|attention"), + ], +) +def test_configuration_rejects_unsupported_modes(monkeypatch, overrides, diagnostic): + cuda_graph, _mode, _registered = _load_cuda_graph(monkeypatch) + + with pytest.raises((RuntimeError, ValueError), match=diagnostic): + cuda_graph.validate_configuration(_server_args(**overrides)) + + +def test_base_runner_warmup_is_suppressed_only_for_save_and_load(monkeypatch): + cuda_graph, mode, _registered = _load_cuda_graph(monkeypatch) + cuda_graph.install_cuda_graph_hooks() + runner = BaseRunner() + + mode.value = Mode.NONE + assert runner.warmup() == "warmed" + mode.value = Mode.SAVE + assert runner.warmup() is None + mode.value = Mode.LOAD + assert runner.warmup() is None + assert runner.warmup_calls == 1 + + +def test_hook_installer_targets_latest_split_runner_classes(monkeypatch): + original_capture = DecodeCudaGraphRunner.capture + original_capture_one = FullCudaGraphBackend.capture_one + cuda_graph, mode, _registered = _load_cuda_graph(monkeypatch) + + cuda_graph.install_cuda_graph_hooks() + + assert DecodeCudaGraphRunner.capture is not original_capture + assert FullCudaGraphBackend.capture_one is not original_capture_one + mode.value = Mode.NONE + runner = DecodeCudaGraphRunner() + assert runner.capture() == "captured" + assert "sglang.srt.model_executor.cuda_graph_runner" not in sys.modules + + +def test_load_hydrates_full_backend_with_complete_shape_keys(monkeypatch): + cuda_graph, _mode, _registered = _load_cuda_graph(monkeypatch) + backend = FullCudaGraphBackend() + pool = object() + graph_a, graph_b = object(), object() + output_a, output_b = object(), object() + key_a = ShapeKey(size=8, stream_idx=0, variant_label="lora") + key_b = ShapeKey(size=8, stream_idx=1, variant_label="nolora") + + cuda_graph.hydrate_backend( + backend, + [ + (key_a, graph_a, output_a), + (key_b, graph_b, output_b), + ], + pool=pool, + ) + + assert backend._graphs == {key_a: graph_a, key_b: graph_b} + assert backend._outputs == {key_a: output_a, key_b: output_b} + assert backend._pool is pool + + +def test_pool_registration_uses_latest_split_runner_api(monkeypatch): + cuda_graph, _mode, registered = _load_cuda_graph(monkeypatch) + pool = object() + + cuda_graph.register_graph_pool(pool) + + assert registered.global_pool is pool + assert registered.pynccl_pool is pool + assert "sglang.srt.model_executor.cuda_graph_runner" not in sys.modules + + +def test_deepep_jit_warmup_does_not_advance_recorded_allocator(monkeypatch): + cuda_graph, _mode, _registered = _load_cuda_graph(monkeypatch) + events = [] + allocator = SimpleNamespace(enabled=True) + compile_utils = _module("sglang.srt.layers.deep_gemm_wrapper.compile_utils") + + def stop_allocation_region(): + allocator.enabled = False + events.append("stop") + + def resume_allocation_region(): + allocator.enabled = True + events.append("resume") + + def compile_all(): + events.append(("compile", allocator.enabled)) + + compile_utils._compile_deep_gemm_one_type_all = compile_all + monkeypatch.setitem( + sys.modules, + "sglang.srt.layers.deep_gemm_wrapper.compile_utils", + compile_utils, + ) + monkeypatch.setattr( + cuda_graph.graph_ops, + "cge", + SimpleNamespace( + stop_allocation_region=stop_allocation_region, + resume_allocation_region=resume_allocation_region, + ), + raising=False, + ) + monkeypatch.setattr(cuda_graph, "_uses_deepep", lambda _runner: True) + monkeypatch.setattr( + cuda_graph, + "bootstrap_deepep_buffer", + lambda _runner: events.append("bootstrap"), + ) + for name in ("save_graph_manifest", "pack_fatbins", "capture_final_alloc_offset"): + monkeypatch.setattr(cuda_graph.graph_ops, name, lambda: None, raising=False) + + runner = SimpleNamespace(attn_backend=object()) + + def capture(_runner): + eager = cuda_graph._EAGER_DEEPEP_WARMUP + events.append(("capture_start", eager, allocator.enabled)) + if eager: + compile_utils._compile_deep_gemm_one_type_all() + events.append(("capture_end", eager, allocator.enabled)) + return "captured" + + assert cuda_graph._save_capture(runner, capture) == "captured" + assert events == [ + "bootstrap", + ("capture_start", True, True), + "stop", + ("compile", False), + "resume", + ("capture_end", True, True), + ("capture_start", False, True), + ("capture_end", False, True), + ] + + +def test_deepep_load_replays_lazy_initialization_before_preallocation(monkeypatch): + cuda_graph, _mode, _registered = _load_cuda_graph(monkeypatch) + events = [] + compile_utils = _module("sglang.srt.layers.deep_gemm_wrapper.compile_utils") + compile_utils._compile_deep_gemm_one_type_all = lambda: None + monkeypatch.setitem( + sys.modules, + "sglang.srt.layers.deep_gemm_wrapper.compile_utils", + compile_utils, + ) + backend = FullCudaGraphBackend() + backend._pool = object() + runner = SimpleNamespace( + _foundry_graphs_loaded=False, + attn_backend=object(), + backend=backend, + capture_bs=[], + deepep_adapter=SimpleNamespace(capture=lambda **_kwargs: events.append("deepep_mode")), + device_module=SimpleNamespace(graph_pool_handle=lambda: object()), + ) + + monkeypatch.setattr(cuda_graph, "_uses_deepep", lambda _runner: True) + monkeypatch.setattr( + cuda_graph, + "bootstrap_deepep_buffer", + lambda _runner: events.append("bootstrap"), + ) + monkeypatch.setattr( + cuda_graph.graph_ops, + "preallocate_for_load_mode", + lambda: events.append("preallocate"), + raising=False, + ) + monkeypatch.setattr( + cuda_graph.graph_ops, + "load_all_graphs", + lambda _runner: [], + raising=False, + ) + + def eager_capture(_runner): + events.append(("eager_capture", cuda_graph._EAGER_DEEPEP_WARMUP)) + + cuda_graph._load_capture(runner, eager_capture) + + assert runner._foundry_graphs_loaded is True + assert events == [ + "bootstrap", + ("eager_capture", True), + "preallocate", + "deepep_mode", + ] diff --git a/tests/test_sglang_graph_ops.py b/tests/test_sglang_graph_ops.py new file mode 100644 index 00000000..3494fe2c --- /dev/null +++ b/tests/test_sglang_graph_ops.py @@ -0,0 +1,367 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the Foundry project +"""Contracts for the versioned SGLang CUDA-graph archive index.""" + +from __future__ import annotations + +import importlib.util +import json +import sys +from dataclasses import dataclass +from pathlib import Path +from types import ModuleType, SimpleNamespace + +import pytest + + +@dataclass(frozen=True) +class ShapeKey: + """The public key shape in SGLang v0.5.15.post1.""" + + size: int + stream_idx: int | None = None + variant_label: str | None = None + + +def _load_graph_ops(monkeypatch: pytest.MonkeyPatch): + """Load graph_ops without requiring the native Foundry extension.""" + + foundry = ModuleType("foundry") + foundry.__path__ = [] + foundry.save_graph_manifest = lambda _path: None + + ops = ModuleType("foundry.ops") + graph = ModuleType("foundry.graph") + graph.CUDAGraph = object + graph.graph = lambda *_args, **_kwargs: None + + config = ModuleType("foundry.integration.sglang.config") + config.CUDAGraphExtensionMode = object + config.get_config = lambda: None + config.get_graph_extension_mode = lambda: None + + runtime = ModuleType("foundry.integration.sglang.runtime") + runtime.get_state = lambda: None + + torch = ModuleType("torch") + torch.Tensor = object + torch.cuda = type("FakeCuda", (), {"CUDAGraph": object})() + + for name, module in ( + ("torch", torch), + ("foundry", foundry), + ("foundry.ops", ops), + ("foundry.graph", graph), + ("foundry.integration.sglang.config", config), + ("foundry.integration.sglang.runtime", runtime), + ): + monkeypatch.setitem(sys.modules, name, module) + + path = ( + Path(__file__).parents[1] / "python" / "foundry" / "integration" / "sglang" / "graph_ops.py" + ) + spec = importlib.util.spec_from_file_location("_foundry_sglang_graph_ops_contract", path) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def _load_real_manifest_scanner(monkeypatch: pytest.MonkeyPatch): + ops = sys.modules["foundry.ops"] + ops.CUDAGraph = object + path = Path(__file__).parents[1] / "python" / "foundry" / "graph.py" + spec = importlib.util.spec_from_file_location("foundry._graph_manifest_contract", path) + assert spec is not None and 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.save_graph_manifest + + +def _entry(graph_ops, **overrides): + values = { + "phase": "decode", + "capture_index": 0, + "filename": "opaque-name.json", + "shape_key": ShapeKey(size=17, stream_idx=2, variant_label="nolora"), + "output_kind": "next_token_logits", + } + values.update(overrides) + return graph_ops.GraphIndexEntry(**values) + + +def _assert_entry( + actual, + *, + phase, + capture_index, + filename, + size, + stream_idx, + variant_label, + output_kind="next_token_logits", +): + assert actual.phase == phase + assert actual.capture_index == capture_index + assert actual.filename == filename + assert actual.output_kind == output_kind + assert actual.shape_key.size == size + assert actual.shape_key.stream_idx == stream_idx + assert actual.shape_key.variant_label == variant_label + + +def test_versioned_graph_index_round_trips_complete_metadata(tmp_path, monkeypatch): + graph_ops = _load_graph_ops(monkeypatch) + index_path = tmp_path / "sglang_graph_index.json" + entry = _entry(graph_ops) + + graph_ops.save_graph_index(index_path, [entry]) + + payload = json.loads(index_path.read_text()) + assert payload["schema_version"] == graph_ops.GRAPH_INDEX_SCHEMA_VERSION + loaded = graph_ops.load_graph_index(index_path) + assert len(loaded) == 1 + _assert_entry( + loaded[0], + phase="decode", + capture_index=0, + filename="opaque-name.json", + size=17, + stream_idx=2, + variant_label="nolora", + ) + assert payload["graphs"][0] == { + "phase": "decode", + "capture_index": 0, + "filename": "opaque-name.json", + "shape_key": { + "size": 17, + "stream_idx": 2, + "variant_label": "nolora", + }, + "output_kind": "next_token_logits", + } + + +def test_saved_filename_matches_real_manifest_scanner_without_key_identity( + tmp_path, + monkeypatch, +): + graph_ops = _load_graph_ops(monkeypatch) + save_graph_manifest = _load_real_manifest_scanner(monkeypatch) + state = SimpleNamespace(capture_index=0, graph_index_entries=[]) + monkeypatch.setattr( + graph_ops, + "get_config", + lambda: SimpleNamespace(workspace_dir=str(tmp_path)), + ) + monkeypatch.setattr(graph_ops.rt, "get_state", lambda: state) + + class Graph: + def save(self, path, _output): + Path(path).write_text( + json.dumps( + { + "topology_key": "opaque-topology", + "dependencies": [], + } + ) + ) + + shape_key = graph_ops.ArchiveShapeKey( + size=17, + stream_idx=2, + variant_label="nolora", + ) + graph_ops.save_graph( + Graph(), + object(), + shape_key, + "next_token_logits", + ) + filename = state.graph_index_entries[0].filename + + save_graph_manifest(str(tmp_path)) + + assert filename == "graph_000000_sglang.json" + assert int(filename.split("_")[1]) == 0 + assert all(identity not in filename for identity in ("17", "2", "nolora")) + manifest = json.loads((tmp_path / "graph_manifest.json").read_text()) + assert manifest["topology_groups"][0]["members"] == [filename] + + +@pytest.mark.parametrize("schema_version", [None, 999]) +def test_graph_index_rejects_legacy_or_unsupported_schema(tmp_path, monkeypatch, schema_version): + graph_ops = _load_graph_ops(monkeypatch) + index_path = tmp_path / "sglang_graph_index.json" + payload = {"graphs": []} + if schema_version is not None: + payload["schema_version"] = schema_version + index_path.write_text(json.dumps(payload)) + + with pytest.raises((RuntimeError, ValueError), match=r"(?i)fresh.*SAVE|SAVE.*fresh"): + graph_ops.load_graph_index(index_path) + + +def test_shape_identity_does_not_collapse_stream_or_variant(tmp_path, monkeypatch): + graph_ops = _load_graph_ops(monkeypatch) + index_path = tmp_path / "sglang_graph_index.json" + entries = [ + _entry( + graph_ops, + capture_index=0, + filename="first.json", + shape_key=ShapeKey(size=8, stream_idx=0, variant_label="lora"), + ), + _entry( + graph_ops, + capture_index=1, + filename="second.json", + shape_key=ShapeKey(size=8, stream_idx=1, variant_label="nolora"), + ), + ] + + graph_ops.save_graph_index(index_path, entries) + loaded = graph_ops.load_graph_index(index_path) + + assert len(loaded) == 2 + _assert_entry( + loaded[0], + phase="decode", + capture_index=0, + filename="first.json", + size=8, + stream_idx=0, + variant_label="lora", + ) + _assert_entry( + loaded[1], + phase="decode", + capture_index=1, + filename="second.json", + size=8, + stream_idx=1, + variant_label="nolora", + ) + assert ( + loaded[0].shape_key.size, + loaded[0].shape_key.stream_idx, + loaded[0].shape_key.variant_label, + ) != ( + loaded[1].shape_key.size, + loaded[1].shape_key.stream_idx, + loaded[1].shape_key.variant_label, + ) + + +def test_graph_index_load_order_is_capture_index_not_filename(tmp_path, monkeypatch): + graph_ops = _load_graph_ops(monkeypatch) + index_path = tmp_path / "sglang_graph_index.json" + later = _entry( + graph_ops, + capture_index=1, + filename="aaa.json", + shape_key=ShapeKey(size=4), + ) + earlier = _entry( + graph_ops, + capture_index=0, + filename="zzz.json", + shape_key=ShapeKey(size=8), + ) + + graph_ops.save_graph_index(index_path, [later, earlier]) + + loaded = graph_ops.load_graph_index(index_path) + assert [entry.capture_index for entry in loaded] == [0, 1] + assert [entry.filename for entry in loaded] == ["zzz.json", "aaa.json"] + assert [entry.shape_key.size for entry in loaded] == [8, 4] + + +def test_graph_index_requires_capture_indices_to_start_at_zero(tmp_path, monkeypatch): + graph_ops = _load_graph_ops(monkeypatch) + index_path = tmp_path / "sglang_graph_index.json" + entry = _entry(graph_ops, capture_index=3) + + with pytest.raises((RuntimeError, ValueError), match=r"(?i)capture.*zero|fresh.*SAVE"): + graph_ops.save_graph_index(index_path, [entry]) + + +@pytest.mark.parametrize("legacy_key", [8, "1_8"]) +def test_graph_index_rejects_integer_and_string_shape_key_heuristics( + tmp_path, monkeypatch, legacy_key +): + graph_ops = _load_graph_ops(monkeypatch) + index_path = tmp_path / "sglang_graph_index.json" + index_path.write_text( + json.dumps( + { + "schema_version": graph_ops.GRAPH_INDEX_SCHEMA_VERSION, + "graphs": [ + { + "phase": "decode", + "capture_index": 0, + "filename": "graph_0_FULL_t8_r8_UX_pcN.json", + "shape_key": legacy_key, + "output_kind": "next_token_logits", + } + ], + } + ) + ) + + with pytest.raises((RuntimeError, TypeError, ValueError), match=r"(?i)fresh.*SAVE|shape"): + graph_ops.load_graph_index(index_path) + + +def test_index_metadata_not_filename_regex_is_source_of_truth(tmp_path, monkeypatch): + graph_ops = _load_graph_ops(monkeypatch) + index_path = tmp_path / "sglang_graph_index.json" + entry = _entry( + graph_ops, + filename="graph_42_FULL_t999_r999_UX_pcN.json", + shape_key=ShapeKey(size=7, stream_idx=None, variant_label=None), + ) + + graph_ops.save_graph_index(index_path, [entry]) + loaded = graph_ops.load_graph_index(index_path) + + assert len(loaded) == 1 + _assert_entry( + loaded[0], + phase="decode", + capture_index=0, + filename="graph_42_FULL_t999_r999_UX_pcN.json", + size=7, + stream_idx=None, + variant_label=None, + ) + + +def test_unsupported_output_kind_is_rejected(tmp_path, monkeypatch): + graph_ops = _load_graph_ops(monkeypatch) + index_path = tmp_path / "sglang_graph_index.json" + index_path.write_text( + json.dumps( + { + "schema_version": graph_ops.GRAPH_INDEX_SCHEMA_VERSION, + "graphs": [ + { + "phase": "decode", + "capture_index": 0, + "filename": "opaque.json", + "shape_key": { + "size": 4, + "stream_idx": None, + "variant_label": None, + }, + "output_kind": "hidden_states", + } + ], + } + ) + ) + + with pytest.raises((RuntimeError, TypeError, ValueError), match=r"(?i)output.*kind|next_token"): + graph_ops.load_graph_index(index_path) diff --git a/tests/test_sglang_recipes.py b/tests/test_sglang_recipes.py new file mode 100644 index 00000000..f6868598 --- /dev/null +++ b/tests/test_sglang_recipes.py @@ -0,0 +1,67 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the Foundry project +"""Contracts for the supported SGLang v0.5.15.post1 serve recipes.""" + +import re +from pathlib import Path + +import pytest +import tomllib + +RECIPE_DIR = Path(__file__).parents[1] / "recipe" / "sglang" +SERVE_SCRIPTS = ( + "serve_qwen3-mini.sh", + "serve_qwen3-1.7b_dp.sh", + "serve_qwen3-30ba3bfp8_ep.sh", +) +TP2_DP2_SCRIPT = "serve_qwen3-8b_tp2_dp2.sh" +TP2_DP2_SAVE_CONFIG = "foundry_tp2_dp2_save.toml" +TP2_DP2_LOAD_CONFIG = "foundry_tp2_dp2_load.toml" + + +@pytest.mark.parametrize("script_name", SERVE_SCRIPTS) +def test_recipe_uses_phase_specific_decode_graph_flags(script_name): + script = (RECIPE_DIR / script_name).read_text() + + assert "--cuda-graph-max-bs-decode" in script + assert "--disable-prefill-cuda-graph" in script + assert re.search(r"--cuda-graph-max-bs(?:\s|$)", script) is None + + +def test_tp2_dp2_recipe_has_exact_validated_topology_and_collective_guards(): + script = (RECIPE_DIR / TP2_DP2_SCRIPT).read_text() + + for argument in ( + "--model-path Qwen/Qwen3-8B", + "--tp-size 2", + "--dp-size 2", + "--pp-size 1", + "--ep-size 1", + "--load-balance-method round_robin", + "--mem-fraction-static 0.8", + "--disable-radix-cache", + "--disable-custom-all-reduce", + "--attention-backend flashinfer", + "--cuda-graph-backend-decode full", + "--cuda-graph-max-bs-decode 256", + "--disable-prefill-cuda-graph", + ): + assert argument in script + assert "export NCCL_CUMEM_ENABLE=0" in script + assert "export NCCL_NVLS_ENABLE=0" in script + assert script.index("export NCCL_CUMEM_ENABLE=0") < script.index('if [[ "$1" == "--save" ]]') + assert script.index("export NCCL_NVLS_ENABLE=0") < script.index('if [[ "$1" == "--save" ]]') + assert TP2_DP2_SAVE_CONFIG in script + assert TP2_DP2_LOAD_CONFIG in script + + +def test_tp2_dp2_save_and_load_configs_differ_only_by_mode(): + with open(RECIPE_DIR / TP2_DP2_SAVE_CONFIG, "rb") as file: + save_config = tomllib.load(file) + with open(RECIPE_DIR / TP2_DP2_LOAD_CONFIG, "rb") as file: + load_config = tomllib.load(file) + + assert save_config.pop("mode") == "save" + assert load_config.pop("mode") == "load" + assert save_config == load_config + assert save_config["workspace_root"] != "foundry_archive" diff --git a/tests/test_sglang_runtime.py b/tests/test_sglang_runtime.py new file mode 100644 index 00000000..614598ab --- /dev/null +++ b/tests/test_sglang_runtime.py @@ -0,0 +1,304 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the Foundry project +"""Pure-Python contracts for SGLang archive compatibility metadata.""" + +from __future__ import annotations + +import importlib.util +import json +import sys +from pathlib import Path +from types import ModuleType, SimpleNamespace + +import pytest + + +def _module(name: str, **attributes): + module = ModuleType(name) + for attr, value in attributes.items(): + setattr(module, attr, value) + return module + + +def _load_runtime(monkeypatch: pytest.MonkeyPatch, workspace_root: Path | None = None): + foundry = _module("foundry") + foundry.__path__ = [] + ops = _module("foundry.ops") + allocation_region = _module( + "foundry.allocation_region", + parse_size=lambda value: value, + ) + config = _module( + "foundry.integration.sglang.config", + CUDAGraphExtensionMode=SimpleNamespace(NONE="none", SAVE="save", LOAD="load"), + compute_workspace_rank=lambda *_args, **_kwargs: 0, + get_config=lambda: SimpleNamespace( + workspace_root=str(workspace_root) if workspace_root is not None else None, + workspace_dir=None, + ), + get_graph_extension_mode=lambda: "none", + get_hook_library_path=lambda: None, + get_nvshmem_host_path=lambda: None, + ) + torch = _module( + "torch", + cuda=SimpleNamespace( + get_device_properties=lambda _index: SimpleNamespace( + name="stub GPU", + total_memory=1024, + ) + ), + version=SimpleNamespace(cuda="13.0"), + ) + + for name, module in ( + ("torch", torch), + ("foundry", foundry), + ("foundry.ops", ops), + ("foundry.allocation_region", allocation_region), + ("foundry.integration.sglang.config", config), + ): + monkeypatch.setitem(sys.modules, name, module) + + path = ( + Path(__file__).parents[1] / "python" / "foundry" / "integration" / "sglang" / "runtime.py" + ) + spec = importlib.util.spec_from_file_location("_foundry_sglang_runtime_contract", path) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + monkeypatch.setitem(sys.modules, spec.name, module) + spec.loader.exec_module(module) + monkeypatch.setattr( + module, + "package_version", + lambda package: "0.5.15.post1" if package == "sglang" else "unexpected", + raising=False, + ) + return module + + +def _server_args(decode: str = "full", prefill: str = "disabled"): + return SimpleNamespace( + cuda_graph_config=SimpleNamespace( + decode=SimpleNamespace(backend=decode), + prefill=SimpleNamespace(backend=prefill), + ) + ) + + +def _matching_state(runtime): + return runtime.WarmupState( + schema_version=runtime.WARMUP_STATE_SCHEMA_VERSION, + sglang_version=runtime.SUPPORTED_SGLANG_VERSION, + decode_cuda_graph_backend="full", + prefill_cuda_graph_backend="disabled", + ) + + +def test_create_warmup_state_records_latest_compatibility_metadata(monkeypatch): + runtime = _load_runtime(monkeypatch) + + state = runtime.create_warmup_state({"page_size": 16}, _server_args()) + + assert state.schema_version == runtime.WARMUP_STATE_SCHEMA_VERSION + assert runtime.WARMUP_STATE_SCHEMA_VERSION == 1 + assert state.sglang_version == "0.5.15.post1" + assert runtime.SUPPORTED_SGLANG_VERSION == "0.5.15.post1" + assert state.decode_cuda_graph_backend == "full" + assert state.prefill_cuda_graph_backend == "disabled" + assert state.memory_pool_config == {"page_size": 16} + + +def test_create_warmup_state_rejects_missing_sglang_package_metadata(monkeypatch): + runtime = _load_runtime(monkeypatch) + + def missing_package_version(_package): + raise runtime.PackageNotFoundError("sglang") + + monkeypatch.setattr(runtime, "package_version", missing_package_version) + + with pytest.raises( + RuntimeError, + match=r"(?is)SAVE.*install.*0\.5\.15\.post1|install.*0\.5\.15\.post1.*SAVE", + ): + runtime.create_warmup_state({}, _server_args()) + + +@pytest.mark.parametrize( + ("changes", "diagnostic"), + [ + ({"schema_version": 0}, "schema"), + ({"sglang_version": ""}, "metadata|version"), + ({"sglang_version": "0.5.14"}, "version"), + ({"decode_cuda_graph_backend": "breakable"}, "decode"), + ({"prefill_cuda_graph_backend": "full"}, "prefill"), + ], +) +def test_validate_warmup_state_rejects_incompatible_archive( + monkeypatch, + changes, + diagnostic, +): + runtime = _load_runtime(monkeypatch) + state = _matching_state(runtime) + for name, value in changes.items(): + setattr(state, name, value) + + with pytest.raises( + (RuntimeError, ValueError), + match=rf"(?is){diagnostic}.*fresh.*SAVE|fresh.*SAVE.*{diagnostic}", + ): + runtime.validate_warmup_state(state, _server_args()) + + +@pytest.mark.parametrize( + ("server_args", "diagnostic"), + [ + (_server_args(decode="breakable"), "decode"), + (_server_args(prefill="full"), "prefill"), + ], +) +def test_validate_warmup_state_rejects_current_config_mismatch( + monkeypatch, + server_args, + diagnostic, +): + runtime = _load_runtime(monkeypatch) + + with pytest.raises( + (RuntimeError, ValueError), + match=rf"(?is){diagnostic}.*fresh.*SAVE|fresh.*SAVE.*{diagnostic}", + ): + runtime.validate_warmup_state(_matching_state(runtime), server_args) + + +def test_validate_warmup_state_accepts_matching_latest_archive(monkeypatch): + runtime = _load_runtime(monkeypatch) + + runtime.validate_warmup_state(_matching_state(runtime), _server_args()) + + +def test_preallocate_for_load_mode_raises_when_region_cannot_be_reserved( + tmp_path, + monkeypatch, +): + runtime = _load_runtime(monkeypatch) + (tmp_path / "final_alloc_offset.json").write_text(json.dumps({"final_alloc_offset": 4096})) + monkeypatch.setattr( + runtime, + "get_config", + lambda: SimpleNamespace( + mode=runtime.CUDAGraphExtensionMode.LOAD, + workspace_dir=str(tmp_path), + ), + ) + monkeypatch.setattr( + runtime, + "cge", + SimpleNamespace( + get_current_alloc_offset=lambda: 1024, + preallocate_region=lambda _size: False, + ), + ) + + with pytest.raises( + RuntimeError, + match=r"(?i)preallocat.*3072|3072.*preallocat", + ): + runtime.preallocate_for_load_mode() + + +@pytest.mark.parametrize( + ("contents", "case"), + [ + (None, "missing"), + ("{", "malformed"), + (json.dumps({"final_alloc_offset": "4096"}), "wrong_type"), + (json.dumps({"final_alloc_offset": 0}), "zero"), + (json.dumps({"final_alloc_offset": -1}), "negative"), + ], +) +def test_preallocate_for_load_mode_requires_positive_per_rank_final_offset( + tmp_path, + monkeypatch, + contents, + case, +): + runtime = _load_runtime(monkeypatch) + if contents is not None: + (tmp_path / "final_alloc_offset.json").write_text(contents) + monkeypatch.setattr( + runtime, + "get_config", + lambda: SimpleNamespace( + mode=runtime.CUDAGraphExtensionMode.LOAD, + workspace_dir=str(tmp_path), + ), + ) + preallocate_calls = [] + monkeypatch.setattr( + runtime, + "cge", + SimpleNamespace( + get_current_alloc_offset=lambda: 1024, + preallocate_region=lambda size: preallocate_calls.append(size), + ), + ) + monkeypatch.setattr( + runtime, + "load_warmup_state", + lambda: SimpleNamespace(final_alloc_offset=8192), + ) + + with pytest.raises( + RuntimeError, + match=r"(?is)per-rank.*final.*offset.*fresh.*SAVE|fresh.*SAVE.*per-rank.*final.*offset", + ): + runtime.preallocate_for_load_mode() + + assert not preallocate_calls, case + + +def test_save_warmup_state_replaces_existing_file_atomically(tmp_path, monkeypatch): + runtime = _load_runtime(monkeypatch, tmp_path) + state_path = tmp_path / "warmup_state.json" + original = '{"sentinel": "existing"}\n' + state_path.write_text(original) + replace_calls = [] + + def fail_replace(source, destination): + replace_calls.append((Path(source), Path(destination))) + raise OSError("simulated interrupted replacement") + + monkeypatch.setattr(runtime.os, "replace", fail_replace) + + with pytest.raises(OSError, match="interrupted replacement"): + runtime.save_warmup_state(_matching_state(runtime)) + + assert replace_calls + assert replace_calls[0][0].parent == tmp_path + assert replace_calls[0][1] == state_path + assert state_path.read_text() == original + assert list(tmp_path.iterdir()) == [state_path] + + +def test_load_warmup_state_ignores_unknown_fields_but_preserves_missing_metadata( + tmp_path, + monkeypatch, +): + runtime = _load_runtime(monkeypatch, tmp_path) + (tmp_path / "warmup_state.json").write_text( + json.dumps( + { + "timestamp": "legacy", + "unknown_future_field": {"safe": True}, + } + ) + ) + + state = runtime.load_warmup_state() + + assert state.timestamp == "legacy" + assert state.schema_version == 0 + with pytest.raises((RuntimeError, ValueError), match=r"(?i)fresh.*SAVE|SAVE.*fresh"): + runtime.validate_warmup_state(state, _server_args()) diff --git a/tests/test_vmm_alloc.py b/tests/test_vmm_alloc.py index 5df909f0..4a2a46f6 100644 --- a/tests/test_vmm_alloc.py +++ b/tests/test_vmm_alloc.py @@ -1,5 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the Foundry project +import ctypes import os import subprocess import sys @@ -146,6 +147,36 @@ def _run_core_size_parsing(): print("[TEST] test_size_parsing PASSED") +def _run_core_zero_alignment_reserve(): + import foundry as fdry + + torch.cuda.init() + + base_addr = 0x7F0000000000 + region_size = 1024 * 1024 * 1024 + reserved_size = 64 * 1024 + reserve = ctypes.CDLL(None).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 + + with fdry.allocation_region(base_addr, region_size): + ptr = ctypes.c_uint64() + offset_before = fdry.get_current_alloc_offset() + + result = reserve(ctypes.byref(ptr), reserved_size, 0, 0, 0) + + offset_after = fdry.get_current_alloc_offset() + assert result == 0 + assert base_addr <= ptr.value < base_addr + region_size + assert offset_before <= offset_after <= region_size + + def _spawn_with_preload(test_mode): so_path = _get_hook_so_path() env = os.environ.copy() @@ -176,6 +207,12 @@ 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_allocation_region(): + """CUDA's default alignment must not reset the allocation cursor.""" + _spawn_with_preload("zero-alignment-reserve") + + if __name__ == "__main__": if "--core" in sys.argv: _run_core() @@ -183,7 +220,10 @@ 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() else: test_vmm_allocation() test_vmm_allocation_without_region() test_size_parsing() + test_zero_alignment_reserve_stays_in_allocation_region()