diff --git a/.gitignore b/.gitignore index d958933a..31d1f0a4 100644 --- a/.gitignore +++ b/.gitignore @@ -362,4 +362,7 @@ marimo/_lsp/ __marimo__/ # Streamlit -.streamlit/secrets.toml \ No newline at end of file +.streamlit/secrets.toml +tests/deepep_matrix_work/ +deepep_fabric_archive/ +hook_archive/ diff --git a/README.md b/README.md index bec64ee4..e15e7109 100644 --- a/README.md +++ b/README.md @@ -85,7 +85,35 @@ Foundry ships engine integrations under `foundry/python/foundry/integration/`. P | SGLang | โœ… | โœ… | ๐Ÿšง | โœ… | | TensorRT-LLM | ๐Ÿšง | ๐Ÿšง | ๐Ÿšง | ๐Ÿšง | -โœ… validated end-to-end (SAVE โ†’ LOAD โ†’ query)  ยท  ๐Ÿšง not yet +โœ… officially validated end-to-end  ยท  ๐Ÿšง not official/general + +An experimental SGLang dense-TP recipe is validated with Qwen3-8B at TP=2 on +2ร—H100. Both ranks restore 36 graphs at their saved allocation offsets, and +deterministic LOAD responses match a non-Foundry baseline using the same pinned +settings. This is not general TP support: upstream +[Discussion #5](https://github.com/orgs/foundry-org/discussions/5) reports a +successful internal torch symmetric-memory prototype but no official +implementation; support remains tracked in +[issue #6](https://github.com/foundry-org/foundry/issues/6). + +SGLang-main also has an opt-in torch symmetric-memory TP=2 path. Foundry records +the process-local communication buffer and pointer-table operands, restores the +per-shape warmups, and relocates those operands to the fresh LOAD communicator. +The checked-in 2ร—H100 proof validates one self-contained batch-1 CUDA graph per +rank at `final_alloc_offset=68555898880`: baseline, two SAVE passes, and fresh +LOAD return byte-identical deterministic responses with no recapture or runtime +errors. Larger batches fall back to eager execution. Multi-shape symmetric +capture remains rejected because later graphs depend on mutable FlashInfer state +owned by earlier shapes. + +An experimental vLLM dense-TP recipe is validated with Qwen3-8B at TP=2 on +2ร—H100. Both ranks restore 64 graphs at the same per-rank offset produced by +repeated deterministic SAVE passes, without recapture. Sequential and +concurrent LOAD responses match a non-Foundry baseline using identical pinned +NCCL settings. As with SGLang, this is a pinned result rather than general TP +support; upstream reports an unpublished torch symmetric-memory prototype +([Discussion #5](https://github.com/orgs/foundry-org/discussions/5), +[issue #6](https://github.com/foundry-org/foundry/issues/6)). The adapted vLLM / SGLang / TensorRT-LLM forks will be released alongside this repo at `foundry-org/vllm`, `foundry-org/sglang`, `foundry-org/TensorRT-LLM`. diff --git a/RELEASE.md b/RELEASE.md index 6586a781..7a6149c5 100644 --- a/RELEASE.md +++ b/RELEASE.md @@ -1,10 +1,55 @@ -# Foundry 0.0.2 +# Release Notes + +## Unreleased + +- **Experimental vLLM dense tensor parallel.** Qwen3-8B TP=2 is validated + end-to-end on 2ร—H100 with NCCL P2P/IPC. Deterministic SAVE passes produce + identical per-rank offsets and graph inventories; fresh LOAD restores 64 + graphs per rank without capture and returns byte-identical sequential and + concurrent factual completions. General TP remains open in + [upstream issue #6](https://github.com/foundry-org/foundry/issues/6); + [Discussion #5](https://github.com/orgs/foundry-org/discussions/5) reports an + internal torch symmetric-memory prototype but no official implementation. +- **Experimental SGLang dense tensor parallel.** Qwen3-8B TP=2 is validated + end-to-end on 2ร—H100 with NCCL P2P/IPC. Fresh LOAD restores 36 graphs per + rank at the recorded allocation offset and returns byte-identical + temperature-0 outputs versus a non-Foundry baseline using identical pinned + settings. General TP remains open in + [upstream issue #6](https://github.com/foundry-org/foundry/issues/6); + [Discussion #5](https://github.com/orgs/foundry-org/discussions/5) reports an + unpublished torch symmetric-memory prototype. SGLang-main now has a validated + TP=2 symmetric-memory option for one batch-1 CUDA graph: SAVE records the + external communication operands and full graph JSON, while LOAD rebuilds + required warmup state and relocates those operands to the live communicator. + A 2ร—H100 baseline/SAVE/SAVE/LOAD matrix returned byte-identical outputs at + `final_alloc_offset=68555898880`. Larger batches use eager execution; + multi-shape symmetric graph capture remains fail-closed. +- **No-fabric DeepEP IPC.** The VMM-IPC bridge transports shareable file + descriptors with `SCM_RIGHTS`, allowing vLLM and SGLang DeepEP NVL buffers to + work without fabric/IMEX. This extends + [upstream PR #3](https://github.com/foundry-org/foundry/pull/3). The verified + regression is BF16 on 2ร—H100; GB300/aarch64 FP8 and RDC relinking reported in + [upstream issue #1](https://github.com/foundry-org/foundry/issues/1) remain + outside that validation scope. +- **FA2 workspace-parity fix.** CUDA generator registration is skipped on LOAD + when every graph has zero RNG increment, preventing unused seed/offset tensors + from drifting the tracked workspace cursor. This incorporates + [upstream PR #4](https://github.com/foundry-org/foundry/pull/4) and the root + cause from [issue #2](https://github.com/foundry-org/foundry/issues/2); + this branch did not independently rerun the A100 hardware case. +- **CUDA default-alignment VMM reservations.** `cuMemAddressReserve` accepts + `alignment=0` to request the driver default. Foundry now preserves the current + VMM cursor for that value instead of bit-masking it to address zero. + +--- + +## Foundry 0.0.2 SGLang graduates to a fully validated engine. This release brings the SGLang integration to parity with vLLM across single GPU, data parallel, and expert parallel โ€” with a self-contained recipe and no vLLM build dependency for EP. -## Highlights +### Highlights - **SGLang single GPU / DP / EP all validated end-to-end.** SAVE โ†’ LOAD โ†’ query verified on single-GPU Qwen3-1.7B / 4B / 14B, @@ -23,7 +68,7 @@ parallel โ€” with a self-contained recipe and no vLLM build dependency for EP. `data_parallel_controller.py`). All save/load logic lives in the integration layer; the edits are inert unless `--foundry-graph-extension-config-path` is set. -## Engine integrations +### Engine integrations - **SGLang** โ€” integration for SGLang v0.5.13. Working configurations: single GPU, data parallel (DP), expert parallel (EP, DeepEP low-latency + DP-attention with fa3). Self-contained recipes under `recipe/sglang/` (shared TOML pair + @@ -45,7 +90,7 @@ parallel โ€” with a self-contained recipe and no vLLM build dependency for EP. on LOAD, with `init_nvshmem_for_loaded_modules` run once on LOAD before any NVSHMEM-kernel graph replays. -## Fixes +### Fixes - **Per-rank VMM device binding (DP/TP/EP).** `set_allocation_region` binds to the current CUDA device, so the integration now calls `set_device(gpu_id)` @@ -63,7 +108,7 @@ parallel โ€” with a self-contained recipe and no vLLM build dependency for EP. capture loop never sets), and the FlashInfer per-bs pre-pass gated off for fa3 while still populating `decode_cuda_graph_metadata` post-load for replay. -## Docs +### Docs - New `docs/sglang/` set (overview, direct-edits, hooks, memory-lifecycle, save-load-workflow, memory-consistency) and a self-contained `recipe/sglang/` @@ -71,10 +116,6 @@ parallel โ€” with a self-contained recipe and no vLLM build dependency for EP. - Top-level README parallelism status table updated to mark SGLang single GPU, DP, and EP as validated. ---- - -## Previous Releases - ## Foundry 0.0.1 First public release of Foundry โ€” a CUDA-graph persistence library that @@ -82,7 +123,7 @@ captures an entire model's CUDA graphs (plus their device context: modules, workspaces, VMM layout) once and replays them at startup, eliminating compile, warmup, and capture from cold-start time. -## Highlights +### Highlights - **Deterministic memory layout โ€” zero patching on graph load.** Foundry indirects memory allocation to the same reserved memory region @@ -99,7 +140,7 @@ warmup, and capture from cold-start time. On load, the template is rebuilt once per group and instances are reconstructed on demand, keeping load fast and asynchronous. -## Engine integrations +### Engine integrations - **vLLM** โ€” compatible with vLLM v0.21. Working configurations: single GPU, data parallel (DP), expert parallel (EP, DeepEP low-latency). End-to-end @@ -111,31 +152,31 @@ warmup, and capture from cold-start time. integration layer (`foundry/integration//`); engine forks contain only minimal hook calls. -## Verified kernel & comm support +### Verified kernel & comm support - **cuBLAS NVJET** kernels (Hopper+). - **torch.compile** modules. - **NVSHMEM / DeepEP** validated. - **DeepGEMM FP8 MoE** validated. -## Dependency +### Dependency - **PyTorch 2.11.0** (compatible with 2.9 โ€“ 2.11). - **CUDA 12+**, CMake 4.0+, Boost. -## Documentation +### Documentation - Integration design notes and per-engine recipes under `docs/` and `recipe/`. - vLLM recipe README covers save/load workflow, archive layout, and required env settings. -## Repository hygiene +### Repository hygiene - Open-source pre-commit hooks (ruff, ruff-format, clang-format, markdownlint, actionlint, DCO sign-off). - Smoke tests covering re-export imports and archive round-trip. -## Roadmap +### Roadmap - Adapted vLLM and SGLang forks published alongside the release. - Tensor parallel support. diff --git a/ROADMAP.md b/ROADMAP.md index aa139d65..ec7264f3 100644 --- a/ROADMAP.md +++ b/ROADMAP.md @@ -26,12 +26,18 @@ - [x] Sync with latest vLLM release - [x] EP on vLLM with DeepEP - - [ ] TP on vLLM + - [ ] Official TP on vLLM + - [x] Experimental NCCL TP=2 recipe on 2ร—H100 + - [ ] Evaluate/publish the torch symmetric-memory prototype ([upstream discussion](https://github.com/orgs/foundry-org/discussions/5)) - [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 + - [ ] Official TP on SGLang + - [x] Experimental NCCL TP=2 recipe on 2ร—H100 + - [x] Reproduce torch symmetric-memory fresh-LOAD corruption on 2ร—H100 + - [x] Relocate process-local symmetric-memory operands for a TP=2 batch-1 graph + - [ ] Generalize symmetric-memory replay to multiple graph shapes ([upstream discussion](https://github.com/orgs/foundry-org/discussions/5)) ## Stage 5: Disaggregated and Large-Scale Serving diff --git a/csrc/CUDAGraph.cpp b/csrc/CUDAGraph.cpp index 25df3524..c1135668 100644 --- a/csrc/CUDAGraph.cpp +++ b/csrc/CUDAGraph.cpp @@ -1528,15 +1528,30 @@ GraphLoadResult CUDAGraph::load(const std::string& json_path, MempoolId_t pool) } const json::array& generators_array = root.at("generators").as_array(); + bool has_rng_generator = false; for (const auto& gen_val : generators_array) { const json::object& gen_obj = gen_val.as_object(); - uint64_t state_id = gen_obj.at("id").to_number(); - uint64_t seed = gen_obj.at("seed").to_number(); - uint64_t wholegraph_increment = gen_obj.at("wholegraph_increment").to_number(); + auto increment_it = gen_obj.find("wholegraph_increment"); + // Missing metadata is treated conservatively: registering may expose an + // allocation mismatch, but skipping could silently leave RNG kernel + // arguments pointing at the SAVE process's seed/offset tensors. + if (increment_it == gen_obj.end() || increment_it->value().to_number() != 0) { + has_rng_generator = true; + break; + } + } - auto state = global_generator_state_registry.get_state_from_id(state_id, seed); - state->register_graph(reinterpret_cast(graph.get())); - graph->captured_generator_states_[state] = wholegraph_increment; + if (has_rng_generator) { + for (const auto& gen_val : generators_array) { + const json::object& gen_obj = gen_val.as_object(); + uint64_t state_id = gen_obj.at("id").to_number(); + uint64_t seed = gen_obj.at("seed").to_number(); + uint64_t wholegraph_increment = gen_obj.at("wholegraph_increment").to_number(); + + auto state = global_generator_state_registry.get_state_from_id(state_id, seed); + state->register_graph(reinterpret_cast(graph.get())); + graph->captured_generator_states_[state] = wholegraph_increment; + } } const json::object& allocator_events = root.at("allocator_events").as_object(); diff --git a/csrc/CUDAGraphParallel.cpp b/csrc/CUDAGraphParallel.cpp index 6ebfc0f6..71fb18c6 100644 --- a/csrc/CUDAGraphParallel.cpp +++ b/csrc/CUDAGraphParallel.cpp @@ -1658,6 +1658,7 @@ std::shared_ptr start_graph_builds_impl( go["seed"] = gens[g].seed; go["wholegraph_increment"] = gens[g].wholegraph_increment; gens_array.push_back(go); + pending->has_rng_generators |= gens[g].wholegraph_increment != 0; } pending->entries[i].generators_meta = boost::json::value(std::move(gens_array)); } @@ -1690,6 +1691,15 @@ std::shared_ptr start_graph_builds_impl( auto gen_it = pre_root.find("generators"); if (gen_it != pre_root.end()) { + const boost::json::array& generators = gen_it->value().as_array(); + for (const auto& gen_val : generators) { + const boost::json::object& generator = gen_val.as_object(); + auto increment_it = generator.find("wholegraph_increment"); + if (increment_it == generator.end() || increment_it->value().to_number() != 0) { + pending->has_rng_generators = true; + break; + } + } pending->entries[i].generators_meta = std::move(gen_it->value()); pre_root.erase(gen_it); } @@ -2045,9 +2055,14 @@ GraphLoadResult finish_one_graph_load_impl(std::shared_ptr pe auto& entry = pending->entries[index]; - // Register deferred generators (before allocator replay to match - // SAVE mode timing where generators are created before graph capture). - if (pending->registry && !entry.generators_meta.is_null()) { + // register_graph lazily creates seed/offset tensors before allocator replay. + // They must stay in the deterministic region when a captured kernel consumes + // Philox state, because their pointers are embedded in that kernel's args. + // When the entire pending set has zero RNG increment, no captured node + // references those tensors; skip registration entirely. This avoids the + // caching-allocator segment growth that caused a 2 MiB A100/FA2 cursor drift + // without replacing deterministic pointers with ordinary CUDA allocations. + if (pending->has_rng_generators && pending->registry && !entry.generators_meta.is_null()) { const boost::json::array& gen_array = entry.generators_meta.as_array(); for (const auto& gen_val : gen_array) { const boost::json::object& gen_obj = gen_val.as_object(); diff --git a/csrc/hook.cpp b/csrc/hook.cpp index d552d08d..2d399b0d 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); } @@ -2543,6 +2574,7 @@ CUresult cuMemAllocPitch_v2(CUdeviceptr* dptr, size_t* pPitch, size_t WidthInByt prop.type = CU_MEM_ALLOCATION_TYPE_PINNED; prop.location.type = CU_MEM_LOCATION_TYPE_DEVICE; prop.location.id = device; + prop.requestedHandleTypes = CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR; result = mem_create_func(&allocHandle, aligned_size, &prop, 0); if (result != CUDA_SUCCESS) { @@ -2657,6 +2689,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 +2893,276 @@ 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 this process's VMM-IPC state. Once IPC starts, a fork child +// cannot safely use inherited C++ mutexes/concurrent-map locks or CUDA handles; +// reject it before touching that state. Fork-before-first-IPC remains safe. +static std::atomic vmm_ipc_owner_pid{-1}; +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; + +static bool vmm_ipc_process_matches_owner(const char* operation) { + pid_t owner = vmm_ipc_owner_pid.load(std::memory_order_acquire); + pid_t pid = getpid(); + if (owner == -1 || owner == pid) + return true; + fprintf(stderr, + "[HOOK] ERROR: %s is unsupported in fork child pid %d after VMM-IPC started in pid %d; " + "use spawn/exec\n", + operation, (int)pid, (int)owner); + return false; +} + +static bool vmm_ipc_claim_process(const char* operation) { + pid_t pid = getpid(); + pid_t expected = -1; + if (vmm_ipc_owner_pid.compare_exchange_strong(expected, pid, std::memory_order_acq_rel, + std::memory_order_acquire) || + expected == pid) { + return true; + } + fprintf(stderr, + "[HOOK] ERROR: %s is unsupported in fork child pid %d after VMM-IPC started in pid %d; " + "use spawn/exec\n", + operation, (int)pid, (int)expected); + return false; +} + +// 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() { + if (!vmm_ipc_claim_process("VMM-IPC server initialization")) + return false; + std::lock_guard lock(vmm_ipc_server_mutex); + pid_t pid = getpid(); + if (vmm_ipc_server_pid == pid) { + return vmm_ipc_listen_fd >= 0; + } + // First call in this process. VMM-IPC after fork is rejected before this + // mutex because inherited mutex/map state and CUDA handles are not safe. + 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) { + if (!vmm_ipc_process_matches_owner("VMM-IPC export invalidation")) + return; + 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) { @@ -2869,6 +3170,9 @@ CUresult cuIpcGetMemHandle(CUipcMemHandle* pHandle, CUdeviceptr dptr) { auto real_func = (cuIpcGetMemHandle_t)CUDA_DRIVER_CALL(cuda_driver_entry_table, CUDA_ENTRY_cuIpcGetMemHandle); + if (!vmm_ipc_process_matches_owner("cuIpcGetMemHandle")) + return CUDA_ERROR_NOT_SUPPORTED; + // Try to find this pointer in our VMM allocations AllocMetadata metadata; bool found = false; @@ -2878,7 +3182,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 +3233,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; @@ -2939,20 +3331,79 @@ CUresult cuIpcOpenMemHandle(CUdeviceptr* pdptr, CUipcMemHandle handle, unsigned memcpy(&magic, handle.reserved, sizeof(uint32_t)); if (magic == VMM_IPC_MAGIC) { + if (!vmm_ipc_claim_process("cuIpcOpenMemHandle")) + return CUDA_ERROR_NOT_SUPPORTED; + // 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 +3412,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 +3427,67 @@ 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 && mapped_ptr != (CUdeviceptr)zone_hint) { + fprintf(stderr, + "[HOOK] ERROR: VMM-IPC import zone reservation returned 0x%llx instead of " + "required 0x%llx\n", + (unsigned long long)mapped_ptr, (unsigned long long)zone_hint); + real_address_free_func(mapped_ptr, map_size); + mem_release_func(imported_handle); + return CUDA_ERROR_OUT_OF_MEMORY; + } + } + 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 +3507,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 +3605,39 @@ CUresult cuIpcCloseMemHandle(CUdeviceptr dptr) { auto real_func = (cuIpcCloseMemHandle_t)CUDA_DRIVER_CALL(cuda_driver_entry_table, CUDA_ENTRY_cuIpcCloseMemHandle); + if (!vmm_ipc_process_matches_owner("cuIpcCloseMemHandle")) + return CUDA_ERROR_NOT_SUPPORTED; + + // 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 +3656,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 +3871,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 +3898,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/sglang/hooks.md b/docs/sglang/hooks.md index cf5b4c82..7ec6bcbf 100644 --- a/docs/sglang/hooks.md +++ b/docs/sglang/hooks.md @@ -12,6 +12,7 @@ In install order: 4. `_patch_kernel_warmup` 5. `_patch_cuda_graph_capture` 6. `_patch_spawn_sites` +7. `_patch_deepep` > Removed: an earlier `_patch_model_runner_init` wrapped `ModelRunner.__init__` to stash `dp_rank` on `self` (upstream didn't expose it as an attribute). Upstream sglang now does `self.dp_rank = dp_rank` in the constructor, so our wrapper is no longer needed; `_resolve_dp_rank` reads `self.dp_rank` directly. @@ -172,10 +173,50 @@ def patched_start(self, *args, **kwargs): `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. +## Experimental tensor parallel + +Dense TP needs no TP-specific monkey-patch beyond the common multi-rank hooks +above. SGLang initializes the NCCL communicator while Foundry is still in the +scratch portion of each rank's VMM region, then both SAVE and LOAD skip to the +same 1 GiB boundary before model allocations. The recipe forces NCCL P2P/IPC +(`NCCL_CUMEM_ENABLE=0`, `NCCL_NVLS_ENABLE=0`) and disables SGLang custom +all-reduce so the collective topology is identical in every phase. + +Validated TP=2 behavior on 2ร—H100: + +- both SAVE passes wrote 36 graphs per rank at `final_alloc_offset=69736595456`; +- LOAD restored 36 graphs per rank and reached that same offset; +- decoded requests reported `cuda graph: True`; +- four temperature-0 responses matched the non-Foundry TP baseline using the + same pinned settings byte for byte. + +The repeatable validation entry point is `tests/modal_sglang_tp.py`. +This is a configuration-specific NCCL result, not the project's general TP +backend. Upstream [Discussion #5](https://github.com/orgs/foundry-org/discussions/5) +documents NCCL initialization nondeterminism and a successful unpublished torch +symmetric-memory prototype, but no official implementation. + +For SGLang main, `_patch_torch_symm_mem` stops Foundry allocation tracking while +`TorchSymmMemCommunicator` creates its cross-rank buffer, resumes the VMM region +before model allocations, and retains the live rendezvous handle. SAVE +preserves full graph JSON plus a versioned inventory of the process-local +buffer, device pointer tables, buffer capacity, dtype, rank, and world size. +LOAD rejects backend or ABI drift, rebuilds SGLang's warmup state, and rewrites +the two-shot all-reduce kernel operands and paired memcpy endpoints to the live +communicator before graph construction. + +This support is fail-closed outside the validated scope: TP must be 2 and the +archive must contain exactly one batch-1 graph. Multi-shape capture is rejected +because later shapes depend on mutable FlashInfer planner state owned by earlier +captures. The 2ร—H100 proof run restored the graph at +`final_alloc_offset=68555898880` with byte-identical baseline/LOAD responses and +no runtime errors. + ## 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). +Active only when `moe_a2a_backend == deepep`. EP runs DP-attention + DeepEP +(NCCL-free); dense TP uses the common path described above and does not install +these additions. - **DeepEP buffer pre-capture bootstrap** (`bootstrap_deepep_buffer`, graph_ops). sglang creates the singleton NVSHMEM `Buffer` lazily on the first MoE dispatch โ€” normally @@ -183,6 +224,13 @@ TP attention is unsupported (its NCCL all-reduce is incompatible with the VMM re 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. +- **DeepEP transport policy** (`_patch_deepep`, hooks). On Foundry SAVE/LOAD, + the integration wraps `deep_ep.Buffer.__init__` and defaults to + `use_fabric=True, num_nvl_bytes=0`. Setting `FOUNDRY_DEEPEP_NVL_IPC=1` + instead sets `use_fabric=False` and preserves SGLang's computed nonzero + `num_nvl_bytes`, selecting the VMM-backed CUDA-IPC NVL path. DeepEP + `29d31c0` or newer is required because it adds this keyword while retaining + the older positional constructor shape. - **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 @@ -203,7 +251,8 @@ TP attention is unsupported (its NCCL all-reduce is incompatible with the VMM re surfaced only on sglang EP. See [`../../recipe/sglang/README.md`](../../recipe/sglang/README.md) for the EP serve -config and required kernel versions (deep_ep `9af0e0d`, sgl-deep-gemm โ‰ฅ0.1.2, fa3). +config and required kernel versions (deep_ep `29d31c0` or newer, +sgl-deep-gemm โ‰ฅ0.1.2, fa3). ## Patch idiom diff --git a/docs/sglang/overview.md b/docs/sglang/overview.md index f8ea7cce..a52c88ba 100644 --- a/docs/sglang/overview.md +++ b/docs/sglang/overview.md @@ -1,8 +1,13 @@ # 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. +Foundry persists SGLang's `CudaGraphRunner` graphs to disk on SAVE and restores +them on LOAD. The standard path skips capture and graph warmup. The opt-in +torch symmetric-memory path replays SGLang's required eager warmups before +loading one or more configured graph shapes. -Tested on single-GPU Qwen3-1.7B / 4B / 14B and **data-parallel (DP=2)** Qwen3-1.7B with the FlashInfer attention backend. +Tested on single-GPU Qwen3-1.7B / 4B / 14B, **data-parallel (DP=2)** +Qwen3-1.7B, and **tensor-parallel (TP=2)** Qwen3-8B with the FlashInfer +attention backend. ## Parallelism @@ -10,16 +15,65 @@ Tested on single-GPU Qwen3-1.7B / 4B / 14B and **data-parallel (DP=2)** Qwen3-1. |---|:---:|---| | 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. | +| Tensor parallel (TP) | ๐Ÿšง | Experimental dense Qwen3-8B recipe validated at TP=2 on 2ร—H100. General support is still under development. | | 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. | +**Experimental tensor parallel.** The dense TP recipe is +`recipe/experimental/serve_qwen3-8b_sglang_tp.sh +[--save|--load]`. It keeps the baseline, SAVE, and LOAD topology identical with +`NCCL_CUMEM_ENABLE=0`, `NCCL_NVLS_ENABLE=0`, +`--disable-custom-all-reduce`, `--disable-piecewise-cuda-graph`, and a fixed +server seed (42 by default). NCCL reports P2P/IPC for every channel. On 2ร—H100, +two SAVE runs produced 36 graphs per rank at the same `final_alloc_offset` +(`69736595456`); a fresh LOAD restored all 36 graphs on both ranks at that +offset, performed no recapture, and returned byte-identical temperature-0 +responses for four sequential prompts. Re-run the checked-in proof with +`modal run --env tests/modal_sglang_tp.py`. + +This validation is deliberately scoped to the pinned model, SGLang revision, +NCCL version, topology, and 2ร—H100 environment. Upstream +[Discussion #5](https://github.com/orgs/foundry-org/discussions/5) records that +NCCL initialization is not generally deterministic and reports a successful +internal torch symmetric-memory prototype. No official implementation has been +published; TP support remains open in +[issue #6](https://github.com/foundry-org/foundry/issues/6). + +Set `SGLANG_ENABLE_TORCH_SYMM_MEM=1` to select the SGLang-main +torch symmetric-memory adaptation. Foundry creates the communication buffer +outside its VMM region, records the buffer and pointer-table graph operands, +preserves full graph JSON before manifest compaction, and relocates those +operands to the live communicator during LOAD. It also mirrors SGLang's two +warmup forwards before SAVE capture and LOAD materialization. + +This path requires TP=2 and defaults to exactly one batch-1 CUDA graph +(`SGLANG_CUDA_GRAPH_MAX_BS=1`); larger batches execute eagerly. Multi-shape +capture remains opt-in: set `FOUNDRY_SGLANG_SYMM_MULTI_SHAPE=1` and provide an +ascending `FOUNDRY_SGLANG_SYMM_GRAPH_BATCHES` list. The pinned 2ร—H100 milestones +are: + +| Captured graph keys | Retained run ID | Final offset per rank | Added VMM bytes vs batch 1 | +|---|---|---:|---:| +| `1` | `8383c4db4d9e6459a6238960-1784839292861701133` | `68555898880` | `0` | +| `1,8` | `sglang-symm-1-8-742970b` | `68597841920` | `41943040` (40 MiB) | +| `1,8,32` | `sglang-symm-1-8-32-742970b` | `68664950784` | `109051904` (104 MiB) | + +Both multi-shape runs passed every harness check. SAVE and SAVE2 produced +identical semantic fingerprints and rank-symmetric offsets; fresh LOAD restored +two and three graphs per rank, respectively, without capture. LOAD replayed the +exact configured graph keys, then proved the uncaptured boundary with explicit +eager decode batches 9 and 33 (`cuda graph: False`) before deterministic graph +outputs matched again. No CUDA/NCCL runtime errors occurred. Concurrent request +text is recorded but not required to be byte-identical because scheduling +changes batch formation. These are pinned configuration results; multi-shape +mode remains disabled by default. + **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 +pinned commit (`29d31c0` or newer; it adds the `use_fabric` API), `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 @@ -30,6 +84,12 @@ SAVE-only warmup pass (triggers DeepGEMM JIT + buffer creation outside the captu 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. +For no-fabric H200 systems, use +`recipe/experimental/serve_qwen3-30ba3bfp8_sglang_ipc_ep.sh`. With +`FOUNDRY_DEEPEP_NVL_IPC=1`, Foundry preserves DeepEP's nonzero NVL buffer and +sets `use_fabric=False`, while the RDMA buffer continues to use NVSHMEM's +symmetric heap. Use identical SAVE and LOAD settings. + **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 @@ -90,8 +150,13 @@ SAVE: 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. +5. `CudaGraphRunner.capture` runs a pre-pass that pre-allocates every per-bs + FlashInfer wrapper. The NCCL path suppresses the normal inner warmups. The + symmetric-memory path runs both warmups and defaults to the batch-1 shape; + opt-in multi-shape mode retains explicit per-shape FlashInfer replay state. +6. Each captured graph is written to disk; a manifest groups topologically + equivalent graphs. Symmetric SAVE additionally preserves the full JSON and a + versioned inventory of its external communication operands. 7. The final VMM cursor is recorded as `final_alloc_offset`. LOAD: @@ -99,7 +164,12 @@ LOAD: 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. +4. On NCCL, `CudaGraphRunner.capture` preallocates through + `final_alloc_offset`, runs the same pre-pass, and loads all graphs through + the manifest's template/on-demand path. On symmetric memory, it validates + the saved backend/ABI and per-shape state, runs the required warmups, + relocates the saved buffer and device-table operands to the live + communicator, and loads each preserved full graph JSON. 5. `self.graphs` / `self.output_buffers` are populated from `state.loaded_graphs`; the rest of SGLang's serving path runs unchanged. ## Doc set diff --git a/docs/superpowers/plans/2026-07-23-sglang-multishape-symmetric-state.md b/docs/superpowers/plans/2026-07-23-sglang-multishape-symmetric-state.md new file mode 100644 index 00000000..e5f8dde7 --- /dev/null +++ b/docs/superpowers/plans/2026-07-23-sglang-multishape-symmetric-state.md @@ -0,0 +1,2982 @@ +# SGLang Multi-Shape Symmetric Replay Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Restore multiple SGLang TP=2 torch symmetric-memory CUDA graph shapes by retaining and refreshing explicit per-shape FlashInfer state while preserving the shared graph pool. + +**Architecture:** Add a `ShapeReplayState` capsule for each graph key, a versioned per-shape manifest, and a pinned SGLang adapter that owns only explicit FlashInfer wrapper fields. SAVE and LOAD create capsules in the same shape order; replay activates and refreshes the selected capsule before launching its relocated graph. The existing batch-1 path remains the default until the complete H100 matrix passes. + +**Tech Stack:** Python 3.12, PyTorch 2.11/cu130, CUDA 13.0 driver API, SGLang commit `1b63155efead7494843f6db8681851ba94611f73`, FlashInfer, Boost JSON/C++ Foundry graph loader, pytest, Modal 2ร—H100. + +## Global Constraints + +- Keep TP fixed at 2 and the two-shot BF16/alignment-16 symmetric-memory kernel. +- Keep `FOUNDRY_SGLANG_SYMM_MULTI_SHAPE` disabled by default until the full acceptance matrix passes. +- Preserve the shared graph pool; do not duplicate model weights, KV-cache storage, or one complete graph pool per shape. +- Do not reduce the Qwen3-8B KV capacity or change `mem_fraction_static=0.8`. +- Inventory only explicit, version-pinned FlashInfer fields; never scan `__dict__`, closures, GC objects, or arbitrary pointer ranges for owners. +- Reject unknown external graph operands before graph construction or replay. +- Keep all imports at module scope. +- Add no third-party dependencies. +- Run native/CUDA validation on Modal H100s with `--env rahul-dev`. +- Keep the current NCCL SGLang path, vLLM TP path, and batch-1 symmetric path behavior unchanged while the feature flag is off. +- End every task with focused tests, review, one logical commit, and `git push -u origin cursor/vllm-tp-consolidation-5d04`. + +## File Map + +### Create + +- `python/foundry/integration/sglang/shape_replay_state.py` โ€” capsule lifecycle and tensor-owner snapshots. +- `python/foundry/integration/sglang/state_manifest.py` โ€” versioned per-shape persistence schema and semantic fingerprint. +- `python/foundry/integration/sglang/sglang_shape_adapter.py` โ€” pinned SGLang/FlashInfer create, prepare, activate, and close operations. +- `python/foundry/integration/sglang/flashinfer_graph_abi.py` โ€” exact pinned FlashInfer kernel ABI and typed operand extraction. +- `tests/test_sglang_shape_replay_state.py` โ€” lifecycle, owner, activation, and eager-isolation contracts. +- `tests/test_sglang_state_manifest.py` โ€” manifest ordering, round-trip, version, and fingerprint contracts. +- `tests/test_sglang_flashinfer_graph_abi.py` โ€” PagedParams/merge ABI and owner-range contracts. + +### Modify + +- `python/foundry/integration/sglang/config.py` โ€” opt-in flag and exact capture-shape parsing. +- `python/foundry/integration/sglang/hooks_main.py` โ€” optional exact capture schedule patch. +- `python/foundry/integration/sglang/main_backend.py` โ€” capsule registry and SAVE/LOAD/replay orchestration. +- `python/foundry/integration/sglang/symm_mem_graph.py` โ€” manifest-driven multi-graph relocation while retaining the single-graph default. +- `recipe/experimental/serve_qwen3-8b_sglang_tp.sh` โ€” guarded multi-shape opt-in. +- `tests/test_sglang_symm_mem_graph.py` โ€” positive and negative multi-shape relocation contracts. +- `tests/test_sglang_tp_recipe.py` โ€” default guard and opt-in schedule contracts. +- `tests/test_sglang_main_compat.py` โ€” hook/config compatibility. +- `tests/modal_sglang_tp.py` โ€” exact expected replay-shape set and milestone matrices. +- `docs/sglang/overview.md`, `docs/sglang/hooks.md`, `recipe/experimental/README.md`, `ROADMAP.md`, `RELEASE.md` โ€” rollout evidence and final scope. + +--- + +### Task 1: Add the shape-capsule lifecycle + +**Files:** + +- Create: `python/foundry/integration/sglang/shape_replay_state.py` +- Create: `tests/test_sglang_shape_replay_state.py` + +**Interfaces:** + +- Produces: `ShapeStateLifecycle`, `TensorOwnerSlot`, and `ShapeReplayState`. +- `ShapeReplayState.attach_graph(graph: Any, outputs: Any) -> None` +- `ShapeReplayState.validate_owners() -> None` +- `ShapeReplayState.mark_planned(plan_fingerprint: tuple[int, ...]) -> None` +- `ShapeReplayState.close() -> None` +- Consumed by: Tasks 2, 3, and 5. + +- [ ] **Step 1: Write failing lifecycle and owner-retention tests** + +```python +# tests/test_sglang_shape_replay_state.py +from __future__ import annotations + +import gc +import weakref + +import pytest + +from foundry.integration.sglang.shape_replay_state import ( + ShapeReplayState, + ShapeStateLifecycle, + TensorOwnerSlot, +) + + +class FakeTensor: + def __init__(self, ptr: int, numel: int = 16) -> None: + self._ptr = ptr + self._numel = numel + self.dtype = "torch.int32" + self.device = "cuda:0" + + def data_ptr(self) -> int: + return self._ptr + + def numel(self) -> int: + return self._numel + + def element_size(self) -> int: + return 4 + + +def make_state(tensor: FakeTensor) -> ShapeReplayState: + slot = TensorOwnerSlot.capture("int_workspace", tensor, ownership="shape") + return ShapeReplayState( + shape_key=1, + batch_size=1, + capture_index=0, + attention_backend=object(), + wrappers=(object(),), + metadata=object(), + owners=(slot,), + shared_owners=(), + shared_objects=(), + communicator=object(), + symm_handle=object(), + ) + + +def test_owner_survives_until_state_close() -> None: + tensor = FakeTensor(0x600000001000) + reference = weakref.ref(tensor) + state = make_state(tensor) + del tensor + gc.collect() + assert reference() is not None + + state.mark_planned(tuple(range(15))) + state.attach_graph(object(), object()) + state.close() + gc.collect() + + assert state.lifecycle is ShapeStateLifecycle.CLOSED + assert reference() is None + + +def test_owner_pointer_replacement_is_rejected() -> None: + tensor = FakeTensor(0x600000001000) + state = make_state(tensor) + tensor._ptr = 0x600000002000 + + with pytest.raises(RuntimeError, match="int_workspace"): + state.validate_owners() + + +def test_invalid_lifecycle_transition_is_rejected() -> None: + state = make_state(FakeTensor(0x600000001000)) + + with pytest.raises(RuntimeError, match="planned"): + state.attach_graph(object(), object()) +``` + +- [ ] **Step 2: Run the tests and verify RED** + +Run: + +```bash +python3 -m pytest tests/test_sglang_shape_replay_state.py -q +``` + +Expected: collection fails because `shape_replay_state` does not exist. + +- [ ] **Step 3: Implement the capsule and owner snapshots** + +```python +# python/foundry/integration/sglang/shape_replay_state.py +from __future__ import annotations + +from dataclasses import dataclass, field +from enum import StrEnum +from typing import Any, Literal + + +class ShapeStateLifecycle(StrEnum): + CREATED = "created" + PLANNED = "planned" + MATERIALIZED = "materialized" + CLOSED = "closed" + + +@dataclass(frozen=True) +class TensorOwnerSlot: + name: str + tensor: Any + ownership: Literal["shape", "shared"] + data_ptr: int + numel: int + element_size: int + dtype: str + device: str + + @classmethod + def capture( + cls, + name: str, + tensor: Any, + *, + ownership: Literal["shape", "shared"], + ) -> TensorOwnerSlot: + return cls( + name=name, + tensor=tensor, + ownership=ownership, + data_ptr=int(tensor.data_ptr()), + numel=int(tensor.numel()), + element_size=int(tensor.element_size()), + dtype=str(tensor.dtype), + device=str(tensor.device), + ) + + def validate_live(self) -> None: + actual = ( + int(self.tensor.data_ptr()), + int(self.tensor.numel()), + int(self.tensor.element_size()), + str(self.tensor.dtype), + str(self.tensor.device), + ) + expected = ( + self.data_ptr, + self.numel, + self.element_size, + self.dtype, + self.device, + ) + if actual != expected: + raise RuntimeError( + f"SGLang shape-state owner changed for {self.name}: " + f"{actual!r} != {expected!r}" + ) + + +@dataclass +class ShapeReplayState: + shape_key: Any + batch_size: int + capture_index: int + attention_backend: Any + wrappers: tuple[Any, ...] + metadata: Any + owners: tuple[TensorOwnerSlot, ...] + shared_owners: tuple[TensorOwnerSlot, ...] + shared_objects: tuple[Any, ...] + communicator: Any + symm_handle: Any + plan_fingerprint: tuple[int, ...] = () + graph: Any = None + outputs: Any = None + lifecycle: ShapeStateLifecycle = ShapeStateLifecycle.CREATED + active: bool = False + operand_inventory: dict[str, int] = field(default_factory=dict) + + def validate_owners(self) -> None: + for slot in (*self.owners, *self.shared_owners): + slot.validate_live() + + def mark_planned(self, plan_fingerprint: tuple[int, ...]) -> None: + if self.lifecycle not in { + ShapeStateLifecycle.CREATED, + ShapeStateLifecycle.PLANNED, + }: + raise RuntimeError( + f"Cannot plan SGLang shape state from {self.lifecycle.value}" + ) + self.plan_fingerprint = plan_fingerprint + self.lifecycle = ShapeStateLifecycle.PLANNED + + def attach_graph(self, graph: Any, outputs: Any) -> None: + if self.lifecycle is not ShapeStateLifecycle.PLANNED: + raise RuntimeError("SGLang shape state must be planned before graph attach") + self.graph = graph + self.outputs = outputs + self.lifecycle = ShapeStateLifecycle.MATERIALIZED + + def close(self) -> None: + if self.active: + raise RuntimeError("Cannot close an active SGLang shape state") + self.graph = None + self.outputs = None + self.metadata = None + self.wrappers = () + self.owners = () + self.shared_owners = () + self.shared_objects = () + self.attention_backend = None + self.symm_handle = None + self.communicator = None + self.lifecycle = ShapeStateLifecycle.CLOSED +``` + +- [ ] **Step 4: Run focused tests and verify GREEN** + +Run: + +```bash +python3 -m pytest tests/test_sglang_shape_replay_state.py -q +``` + +Expected: `3 passed`. + +- [ ] **Step 5: Run formatting and commit** + +```bash +pre-commit run --files \ + python/foundry/integration/sglang/shape_replay_state.py \ + tests/test_sglang_shape_replay_state.py +git add \ + python/foundry/integration/sglang/shape_replay_state.py \ + tests/test_sglang_shape_replay_state.py +git commit -s -m "feat: add SGLang shape replay state" +git push -u origin cursor/vllm-tp-consolidation-5d04 +``` + +--- + +### Task 2: Add the versioned per-shape state manifest + +**Files:** + +- Create: `python/foundry/integration/sglang/state_manifest.py` +- Create: `tests/test_sglang_state_manifest.py` + +**Interfaces:** + +- Consumes: `TensorOwnerSlot`, `ShapeReplayState` from Task 1. +- Produces: `OperandSlotRecord`, `TensorSlotRecord`, `ShapeStateRecord`, `ShapeStateManifest`. +- `write_shape_state_manifest(workspace: str | Path, manifest: ShapeStateManifest) -> None` +- `read_shape_state_manifest(workspace: str | Path) -> ShapeStateManifest` +- `build_shape_record(state, graph_filename, plan_schema, operand_slots) -> ShapeStateRecord` +- `validate_shape_record(state, record, plan_schema) -> None` +- Consumed by: Tasks 4 and 5. + +- [ ] **Step 1: Write failing manifest tests** + +```python +# tests/test_sglang_state_manifest.py +from __future__ import annotations + +import json +from dataclasses import replace +from types import SimpleNamespace + +import pytest + +from foundry.integration.sglang.state_manifest import ( + MANIFEST_FILENAME, + OperandSlotRecord, + ShapeStateManifest, + ShapeStateRecord, + TensorSlotRecord, + read_shape_state_manifest, + write_shape_state_manifest, + build_shape_record, + validate_shape_record, +) + + +def make_manifest() -> ShapeStateManifest: + tensor_slot = TensorSlotRecord( + name="int_workspace", + ownership="shape", + dtype="torch.uint8", + device="cuda:0", + numel=8 * 1024 * 1024, + element_size=1, + ) + operand = OperandSlotRecord( + name="symm.buffer_ptrs_dev", + kernel_abi="torch.symm_mem.two_shot.bf16.align16.tp2", + owner_slot="communicator", + node_id=1, + source="kernelParams", + parameter_index=0, + cuda_parameter_offset=0, + value_byte_offset=0, + ctype="BFloat16**", + owner_relative_offset=0, + span_bytes=16, + saved_value=0x600006400400, + ) + shape = ShapeStateRecord( + graph_filename="graph_0_FULL_t1_r1_UX_pcN.json", + batch_size=1, + capture_index=0, + plan_schema="flashinfer.prefill.v15", + plan_fingerprint=tuple(range(15)), + tensor_slots=(tensor_slot,), + operand_slots=(operand,), + ) + return ShapeStateManifest( + backend="torch_symmetric_memory", + rank=0, + world_size=2, + capture_order=(1,), + shapes=(shape,), + ) + + +def test_manifest_round_trip(tmp_path) -> None: + manifest = make_manifest() + write_shape_state_manifest(tmp_path, manifest) + + assert read_shape_state_manifest(tmp_path) == manifest + + +def test_manifest_rejects_duplicate_shape(tmp_path) -> None: + manifest = make_manifest() + duplicate = ShapeStateManifest( + backend=manifest.backend, + rank=manifest.rank, + world_size=manifest.world_size, + capture_order=(1, 1), + shapes=(manifest.shapes[0], manifest.shapes[0]), + ) + + with pytest.raises(RuntimeError, match="duplicate"): + write_shape_state_manifest(tmp_path, duplicate) + + +def test_manifest_rejects_unknown_version(tmp_path) -> None: + path = tmp_path / MANIFEST_FILENAME + path.write_text(json.dumps({"version": 999})) + + with pytest.raises(RuntimeError, match="version"): + read_shape_state_manifest(tmp_path) + + +def test_live_shape_must_match_manifest_record() -> None: + owner = SimpleNamespace( + name="int_workspace", + ownership="shape", + dtype="torch.uint8", + device="cuda:0", + numel=1024, + element_size=1, + ) + state = SimpleNamespace( + batch_size=1, + capture_index=0, + plan_fingerprint=tuple(range(15)), + owners=(owner,), + shared_owners=(), + ) + record = build_shape_record( + state, + graph_filename="graph_0_FULL_t1_r1_UX_pcN.json", + plan_schema="flashinfer.prefill.v15", + operand_slots=(), + ) + validate_shape_record( + state, + record, + plan_schema="flashinfer.prefill.v15", + ) + + with pytest.raises(RuntimeError, match="live shape state"): + validate_shape_record( + state, + replace(record, plan_fingerprint=(99,)), + plan_schema="flashinfer.prefill.v15", + ) +``` + +- [ ] **Step 2: Run tests and verify RED** + +Run: + +```bash +python3 -m pytest tests/test_sglang_state_manifest.py -q +``` + +Expected: collection fails because `state_manifest` does not exist. + +- [ ] **Step 3: Implement the manifest schema and canonical fingerprint** + +```python +# python/foundry/integration/sglang/state_manifest.py +from __future__ import annotations + +import hashlib +import json +from dataclasses import asdict, dataclass +from pathlib import Path +from typing import Literal + +from foundry.integration.sglang.shape_replay_state import ShapeReplayState + +MANIFEST_FILENAME = "sglang_shape_state.json" +MANIFEST_VERSION = 1 + + +@dataclass(frozen=True) +class TensorSlotRecord: + name: str + ownership: Literal["shape", "shared"] + dtype: str + device: str + numel: int + element_size: int + + +@dataclass(frozen=True) +class OperandSlotRecord: + name: str + kernel_abi: str + owner_slot: str + node_id: int + source: str + parameter_index: int + cuda_parameter_offset: int + value_byte_offset: int + ctype: str + owner_relative_offset: int + span_bytes: int + saved_value: int + + +@dataclass(frozen=True) +class ShapeStateRecord: + graph_filename: str + batch_size: int + capture_index: int + plan_schema: str + plan_fingerprint: tuple[int, ...] + tensor_slots: tuple[TensorSlotRecord, ...] + operand_slots: tuple[OperandSlotRecord, ...] + + +@dataclass(frozen=True) +class ShapeStateManifest: + backend: str + rank: int + world_size: int + capture_order: tuple[int, ...] + shapes: tuple[ShapeStateRecord, ...] + + +def build_shape_record( + state: ShapeReplayState, + *, + graph_filename: str, + plan_schema: str, + operand_slots: tuple[OperandSlotRecord, ...], +) -> ShapeStateRecord: + tensor_slots = tuple( + TensorSlotRecord( + name=slot.name, + ownership=slot.ownership, + dtype=slot.dtype, + device=slot.device, + numel=slot.numel, + element_size=slot.element_size, + ) + for slot in (*state.owners, *state.shared_owners) + ) + return ShapeStateRecord( + graph_filename=graph_filename, + batch_size=state.batch_size, + capture_index=state.capture_index, + plan_schema=plan_schema, + plan_fingerprint=state.plan_fingerprint, + tensor_slots=tensor_slots, + operand_slots=operand_slots, + ) + + +def validate_shape_record( + state: ShapeReplayState, + record: ShapeStateRecord, + *, + plan_schema: str, +) -> None: + expected = build_shape_record( + state, + graph_filename=record.graph_filename, + plan_schema=plan_schema, + operand_slots=record.operand_slots, + ) + if expected != record: + raise RuntimeError( + f"SGLang live shape state does not match batch {record.batch_size}" + ) + + +def _shape_to_dict(shape: ShapeStateRecord) -> dict: + return { + "graph_filename": shape.graph_filename, + "batch_size": shape.batch_size, + "capture_index": shape.capture_index, + "plan_schema": shape.plan_schema, + "plan_fingerprint": list(shape.plan_fingerprint), + "tensor_slots": [asdict(slot) for slot in shape.tensor_slots], + "operand_slots": [asdict(slot) for slot in shape.operand_slots], + } + + +def _semantic_fingerprint(manifest: ShapeStateManifest) -> str: + semantic = { + "backend": manifest.backend, + "rank": manifest.rank, + "world_size": manifest.world_size, + "capture_order": list(manifest.capture_order), + "shapes": [ + { + **_shape_to_dict(shape), + "operand_slots": [ + { + key: value + for key, value in asdict(slot).items() + if key != "saved_value" + } + for slot in shape.operand_slots + ], + } + for shape in manifest.shapes + ], + } + encoded = json.dumps( + semantic, + sort_keys=True, + separators=(",", ":"), + ).encode() + return hashlib.sha256(encoded).hexdigest() + + +def _validate_manifest(manifest: ShapeStateManifest) -> None: + if manifest.backend != "torch_symmetric_memory": + raise RuntimeError(f"Unsupported SGLang shape-state backend: {manifest.backend}") + if manifest.world_size != 2 or manifest.rank not in (0, 1): + raise RuntimeError("SGLang multi-shape symmetric replay requires TP=2") + batches = tuple(shape.batch_size for shape in manifest.shapes) + if len(batches) != len(set(batches)): + raise RuntimeError("SGLang shape-state manifest contains duplicate shapes") + if batches != manifest.capture_order: + raise RuntimeError( + f"SGLang shape-state order mismatch: {batches} != {manifest.capture_order}" + ) + indices = tuple(shape.capture_index for shape in manifest.shapes) + if indices != tuple(range(len(indices))): + raise RuntimeError(f"SGLang capture indices are not contiguous: {indices}") + for shape in manifest.shapes: + names = tuple(slot.name for slot in shape.tensor_slots) + if len(names) != len(set(names)): + raise RuntimeError( + f"Duplicate tensor slot in batch {shape.batch_size}: {names}" + ) + locations = tuple( + ( + slot.node_id, + slot.source, + slot.parameter_index, + slot.value_byte_offset, + ) + for slot in shape.operand_slots + ) + if len(locations) != len(set(locations)): + raise RuntimeError( + f"Duplicate operand slot in batch {shape.batch_size}" + ) + + +def write_shape_state_manifest( + workspace: str | Path, + manifest: ShapeStateManifest, +) -> None: + _validate_manifest(manifest) + payload = { + "version": MANIFEST_VERSION, + "schema_fingerprint": _semantic_fingerprint(manifest), + "manifest": { + "backend": manifest.backend, + "rank": manifest.rank, + "world_size": manifest.world_size, + "capture_order": list(manifest.capture_order), + "shapes": [_shape_to_dict(shape) for shape in manifest.shapes], + }, + } + path = Path(workspace) / MANIFEST_FILENAME + path.write_text(json.dumps(payload, sort_keys=True)) + + +def read_shape_state_manifest(workspace: str | Path) -> ShapeStateManifest: + path = Path(workspace) / MANIFEST_FILENAME + try: + payload = json.loads(path.read_text()) + except (OSError, json.JSONDecodeError) as error: + raise RuntimeError("Invalid SGLang shape-state manifest") from error + if payload.get("version") != MANIFEST_VERSION: + raise RuntimeError( + f"Unsupported SGLang shape-state version: {payload.get('version')!r}" + ) + raw = payload["manifest"] + shapes = tuple( + ShapeStateRecord( + graph_filename=shape["graph_filename"], + batch_size=int(shape["batch_size"]), + capture_index=int(shape["capture_index"]), + plan_schema=shape["plan_schema"], + plan_fingerprint=tuple(int(value) for value in shape["plan_fingerprint"]), + tensor_slots=tuple( + TensorSlotRecord(**slot) for slot in shape["tensor_slots"] + ), + operand_slots=tuple( + OperandSlotRecord(**slot) for slot in shape["operand_slots"] + ), + ) + for shape in raw["shapes"] + ) + manifest = ShapeStateManifest( + backend=raw["backend"], + rank=int(raw["rank"]), + world_size=int(raw["world_size"]), + capture_order=tuple(int(value) for value in raw["capture_order"]), + shapes=shapes, + ) + _validate_manifest(manifest) + if payload.get("schema_fingerprint") != _semantic_fingerprint(manifest): + raise RuntimeError("SGLang shape-state schema fingerprint mismatch") + return manifest +``` + +- [ ] **Step 4: Run manifest tests and verify GREEN** + +```bash +python3 -m pytest tests/test_sglang_state_manifest.py -q +``` + +Expected: `4 passed`. + +- [ ] **Step 5: Format and commit** + +```bash +pre-commit run --files \ + python/foundry/integration/sglang/state_manifest.py \ + tests/test_sglang_state_manifest.py +git add \ + python/foundry/integration/sglang/state_manifest.py \ + tests/test_sglang_state_manifest.py +git commit -s -m "feat: persist SGLang shape state manifest" +git push -u origin cursor/vllm-tp-consolidation-5d04 +``` + +--- + +### Task 3: Implement the pinned FlashInfer shape adapter + +**Files:** + +- Create: `python/foundry/integration/sglang/sglang_shape_adapter.py` +- Modify: `python/foundry/integration/sglang/shape_replay_state.py` +- Modify: `tests/test_sglang_shape_replay_state.py` + +**Interfaces:** + +- Consumes: `ShapeReplayState`, `TensorOwnerSlot`. +- Produces: + - `create_shape_state(cuda_graph_runner, shape_key, capture_index, communicator, symm_handle) -> ShapeReplayState` + - `prepare_shape_state(state, forward_batch, *, for_capture: bool) -> None` + - `activate_shape_state(state) -> ContextManager[None]` + - `close_shape_state(state) -> None` +- Consumed by: Task 5. + +- [ ] **Step 1: Add failing adapter tests with explicit fake wrapper fields** + +```python +# append to tests/test_sglang_shape_replay_state.py +from types import SimpleNamespace + +from foundry.integration.sglang.sglang_shape_adapter import ( + PLAN_SCHEMA, + activate_shape_state, + create_shape_state, + prepare_shape_state, +) + + +class FakeWrapper: + def __init__( + self, + base: int, + shared_base: int = 0x600000100000, + ) -> None: + self._int_workspace_buffer = FakeTensor(base) + self._pin_memory_int_workspace_buffer = FakeTensor(base + 0x1000) + self._qo_indptr_buf = FakeTensor(base + 0x2000) + self._paged_kv_indptr_buf = FakeTensor(shared_base + 0x1000) + self._paged_kv_indices_buf = FakeTensor(shared_base + 0x2000) + self._paged_kv_last_page_len_buf = FakeTensor(shared_base + 0x3000) + self._float_workspace_buffer = FakeTensor(shared_base + 0x4000) + self._fixed_batch_size = 8 + self._workspace_size = 8 * 1024 * 1024 + self._cached_module = object() + self._plan_info = list(range(15)) + + +class FakeAttentionBackend: + def __init__(self, wrapper: FakeWrapper) -> None: + self.decode_cuda_graph_metadata = {8: [wrapper]} + self.forward_metadata = object() + self.prepared = [] + + def init_forward_metadata_out_graph(self, forward_batch, in_capture=False): + self.prepared.append((forward_batch, in_capture)) + + +def make_runner(wrapper: FakeWrapper): + backend = FakeAttentionBackend(wrapper) + model_runner = SimpleNamespace( + attn_backend=backend, + graph_shared_output=object(), + ) + return SimpleNamespace( + captured_req_width=1, + attn_backend=backend, + model_runner=model_runner, + buffers=object(), + ) + + +def test_adapter_retains_explicit_owner_allowlist() -> None: + wrapper = FakeWrapper(0x600001000000) + runner = make_runner(wrapper) + state = create_shape_state( + runner, + shape_key=8, + capture_index=0, + communicator=object(), + symm_handle=object(), + ) + + assert tuple(slot.name for slot in state.owners) == ( + "int_workspace", + "pin_memory_int_workspace", + "qo_indptr", + ) + assert tuple(slot.name for slot in state.shared_owners) == ( + "paged_kv_indptr", + "paged_kv_indices", + "paged_kv_last_page_len", + "float_workspace", + ) + assert state.plan_fingerprint == tuple(range(15)) + + +def test_activation_restores_previous_backend_bindings() -> None: + wrapper = FakeWrapper(0x600001000000) + runner = make_runner(wrapper) + state = create_shape_state( + runner, + shape_key=8, + capture_index=0, + communicator=object(), + symm_handle=object(), + ) + previous_metadata = runner.attn_backend.forward_metadata + + with activate_shape_state(state): + assert runner.attn_backend.decode_cuda_graph_metadata[8][0] is wrapper + assert runner.attn_backend.forward_metadata is state.metadata + + assert runner.attn_backend.forward_metadata is previous_metadata + + +def test_replay_prepare_reuses_wrapper_identity() -> None: + wrapper = FakeWrapper(0x600001000000) + runner = make_runner(wrapper) + state = create_shape_state( + runner, + shape_key=8, + capture_index=0, + communicator=object(), + symm_handle=object(), + ) + batch = SimpleNamespace(batch_size=8) + + with activate_shape_state(state): + prepare_shape_state(state, batch, for_capture=False) + + assert runner.attn_backend.prepared == [(batch, False)] + assert state.wrappers == (wrapper,) +``` + +- [ ] **Step 2: Run adapter tests and verify RED** + +```bash +python3 -m pytest tests/test_sglang_shape_replay_state.py -q +``` + +Expected: import fails because `sglang_shape_adapter` does not exist. + +- [ ] **Step 3: Implement exact-field extraction and activation** + +```python +# python/foundry/integration/sglang/sglang_shape_adapter.py +from __future__ import annotations + +from contextlib import contextmanager +from typing import Any, Iterator + +from foundry.integration.sglang.shape_replay_state import ( + ShapeReplayState, + ShapeStateLifecycle, + TensorOwnerSlot, +) + +PLAN_SCHEMA = "flashinfer.prefill.v15" +_SHAPE_FIELDS = ( + ("int_workspace", "_int_workspace_buffer"), + ("pin_memory_int_workspace", "_pin_memory_int_workspace_buffer"), + ("qo_indptr", "_qo_indptr_buf"), +) +_SHARED_FIELDS = ( + ("paged_kv_indptr", "_paged_kv_indptr_buf"), + ("paged_kv_indices", "_paged_kv_indices_buf"), + ("paged_kv_last_page_len", "_paged_kv_last_page_len_buf"), + ("float_workspace", "_float_workspace_buffer"), +) + + +def _require_field(owner: Any, field_name: str) -> Any: + if not hasattr(owner, field_name): + raise RuntimeError( + "Pinned FlashInfer wrapper ABI changed: " + f"missing required field {field_name}" + ) + return getattr(owner, field_name) + + +def _plan_fingerprint(wrapper: Any) -> tuple[int, ...]: + plan_info = tuple(int(value) for value in _require_field(wrapper, "_plan_info")) + if len(plan_info) != 15: + raise RuntimeError( + "Pinned FlashInfer plan schema changed: " + f"expected 15 fields, got {len(plan_info)}" + ) + if _require_field(wrapper, "_cached_module") is None: + raise RuntimeError("FlashInfer wrapper has no cached CUDA graph module") + return plan_info + + +def create_shape_state( + cuda_graph_runner: Any, + shape_key: Any, + capture_index: int, + communicator: Any, + symm_handle: Any, +) -> ShapeReplayState: + if int(cuda_graph_runner.captured_req_width) != 1: + raise RuntimeError("SGLang multi-shape replay requires captured_req_width=1") + batch_size = int(shape_key.size if hasattr(shape_key, "size") else shape_key) + backend = cuda_graph_runner.attn_backend + wrappers = tuple(backend.decode_cuda_graph_metadata.get(batch_size, ())) + if len(wrappers) != 1: + raise RuntimeError( + f"Expected one FlashInfer wrapper for batch {batch_size}, " + f"got {len(wrappers)}" + ) + wrapper = wrappers[0] + owners = tuple( + TensorOwnerSlot.capture( + name, + _require_field(wrapper, field_name), + ownership="shape", + ) + for name, field_name in _SHAPE_FIELDS + ) + shared_owners = tuple( + TensorOwnerSlot.capture( + name, + _require_field(wrapper, field_name), + ownership="shared", + ) + for name, field_name in _SHARED_FIELDS + ) + state = ShapeReplayState( + shape_key=shape_key, + batch_size=batch_size, + capture_index=capture_index, + attention_backend=backend, + wrappers=wrappers, + metadata=backend.forward_metadata, + owners=owners, + shared_owners=shared_owners, + shared_objects=( + cuda_graph_runner.buffers, + cuda_graph_runner.model_runner.graph_shared_output, + ), + communicator=communicator, + symm_handle=symm_handle, + ) + state.mark_planned(_plan_fingerprint(wrapper)) + return state + + +@contextmanager +def activate_shape_state(state: ShapeReplayState) -> Iterator[None]: + if state.active: + raise RuntimeError(f"SGLang shape {state.batch_size} is already active") + if state.lifecycle is ShapeStateLifecycle.CLOSED: + raise RuntimeError(f"SGLang shape {state.batch_size} is closed") + backend = state.attention_backend + previous_wrappers = backend.decode_cuda_graph_metadata.get(state.batch_size) + previous_metadata = backend.forward_metadata + state.active = True + backend.decode_cuda_graph_metadata[state.batch_size] = list(state.wrappers) + backend.forward_metadata = state.metadata + try: + yield + state.metadata = backend.forward_metadata + finally: + if previous_wrappers is None: + backend.decode_cuda_graph_metadata.pop(state.batch_size, None) + else: + backend.decode_cuda_graph_metadata[state.batch_size] = previous_wrappers + backend.forward_metadata = previous_metadata + state.active = False + + +def prepare_shape_state( + state: ShapeReplayState, + forward_batch: Any, + *, + for_capture: bool, +) -> None: + if not state.active: + raise RuntimeError( + f"SGLang shape {state.batch_size} must be active before preparation" + ) + if not for_capture: + state.attention_backend.init_forward_metadata_out_graph( + forward_batch, + in_capture=False, + ) + state.validate_owners() + actual = _plan_fingerprint(state.wrappers[0]) + if actual != state.plan_fingerprint: + raise RuntimeError( + f"FlashInfer plan topology changed for batch {state.batch_size}" + ) + + +def close_shape_state(state: ShapeReplayState) -> None: + state.close() +``` + +- [ ] **Step 4: Run adapter and capsule tests** + +```bash +python3 -m pytest tests/test_sglang_shape_replay_state.py -q +``` + +Expected: all tests pass. + +- [ ] **Step 5: Add a shared-owner consistency test** + +```python +# append to tests/test_sglang_shape_replay_state.py +def test_shapes_share_only_approved_backend_buffers() -> None: + first = create_shape_state( + make_runner(FakeWrapper(0x600001000000)), + shape_key=1, + capture_index=0, + communicator=object(), + symm_handle=object(), + ) + second_wrapper = FakeWrapper(0x600002000000) + second = create_shape_state( + make_runner(second_wrapper), + shape_key=8, + capture_index=1, + communicator=object(), + symm_handle=object(), + ) + + assert {slot.data_ptr for slot in first.owners}.isdisjoint( + slot.data_ptr for slot in second.owners + ) + assert tuple(slot.data_ptr for slot in first.shared_owners) == tuple( + slot.data_ptr for slot in second.shared_owners + ) +``` + +Run: + +```bash +python3 -m pytest tests/test_sglang_shape_replay_state.py -q +``` + +Expected: all tests pass. + +- [ ] **Step 6: Format and commit** + +```bash +pre-commit run --files \ + python/foundry/integration/sglang/shape_replay_state.py \ + python/foundry/integration/sglang/sglang_shape_adapter.py \ + tests/test_sglang_shape_replay_state.py +git add \ + python/foundry/integration/sglang/shape_replay_state.py \ + python/foundry/integration/sglang/sglang_shape_adapter.py \ + tests/test_sglang_shape_replay_state.py +git commit -s -m "feat: adapt SGLang FlashInfer shape state" +git push -u origin cursor/vllm-tp-consolidation-5d04 +``` + +--- + +### Task 4: Make symmetric graph relocation shape-aware + +**Files:** + +- Create: `python/foundry/integration/sglang/flashinfer_graph_abi.py` +- Create: `tests/test_sglang_flashinfer_graph_abi.py` +- Modify: `python/foundry/integration/sglang/symm_mem_graph.py:68-244` +- Modify: `python/foundry/integration/sglang/state_manifest.py` +- Modify: `tests/test_sglang_symm_mem_graph.py` + +**Interfaces:** + +- Consumes: `ShapeReplayState`, `ShapeStateManifest`, `ShapeStateRecord`, `OperandSlotRecord`. +- Produces: `GraphOperandInventory`. +- Produces: `extract_flashinfer_operand_records(graph_data, state) -> tuple[OperandSlotRecord, ...]`. +- Changes: + - `preserve_symmetric_graphs(..., allow_multi_shape: bool = False) -> tuple[GraphOperandInventory, ...]` + - `relocate_symmetric_graphs(..., manifest: ShapeStateManifest | None = None, allow_multi_shape: bool = False) -> list[str]` +- Produces: one validated operand inventory per graph shape. +- Consumed by: Task 5. + +- [ ] **Step 1: Replace the existing negative multi-shape test with opt-in RED tests** + +```python +# tests/test_sglang_symm_mem_graph.py +def test_preserve_allows_ordered_multi_shape_with_flag(tmp_path) -> None: + module = _load_module() + state = _state(module) + filenames = [ + "graph_0_FULL_t8_r8_UX_pcN.json", + "graph_1_FULL_t1_r1_UX_pcN.json", + ] + for filename in filenames: + (tmp_path / filename).write_text(json.dumps(_graph(state))) + + records = module.preserve_symmetric_graphs( + tmp_path, + filenames, + state, + allow_multi_shape=True, + ) + + assert tuple(record.batch_size for record in records) == (8, 1) + + +def test_preserve_still_rejects_multi_shape_by_default(tmp_path) -> None: + module = _load_module() + state = _state(module) + filenames = [ + "graph_0_FULL_t8_r8_UX_pcN.json", + "graph_1_FULL_t1_r1_UX_pcN.json", + ] + for filename in filenames: + (tmp_path / filename).write_text(json.dumps(_graph(state))) + + with pytest.raises(RuntimeError, match="exactly one"): + module.preserve_symmetric_graphs(tmp_path, filenames, state) +``` + +- [ ] **Step 2: Run focused tests and verify RED** + +```bash +python3 -m pytest tests/test_sglang_symm_mem_graph.py -q \ + -k "multi_shape" +``` + +Expected: the opt-in test fails because `allow_multi_shape` is not accepted. + +- [ ] **Step 3: Add strict filename/order parsing** + +```python +# python/foundry/integration/sglang/symm_mem_graph.py +import re +from dataclasses import dataclass + +from foundry.integration.sglang.state_manifest import ( + OperandSlotRecord, +) + +_GRAPH_NAME_RE = re.compile( + r"^graph_(?P\d+)_FULL_t(?P\d+)_r(?P=batch)_UX_pcN\.json$" +) + + +@dataclass(frozen=True) +class GraphOperandInventory: + graph_filename: str + batch_size: int + capture_index: int + operand_slots: tuple[OperandSlotRecord, ...] + + +def _validate_graph_files( + graph_filenames: list[str], + *, + allow_multi_shape: bool, +) -> tuple[tuple[int, int, str], ...]: + parsed = [] + for filename in graph_filenames: + match = _GRAPH_NAME_RE.fullmatch(filename) + if match is None: + raise RuntimeError(f"Invalid SGLang graph filename: {filename}") + parsed.append( + (int(match.group("index")), int(match.group("batch")), filename) + ) + parsed.sort() + indices = tuple(index for index, _batch, _filename in parsed) + batches = tuple(batch for _index, batch, _filename in parsed) + if indices != tuple(range(len(indices))): + raise RuntimeError(f"SGLang graph indices are not contiguous: {indices}") + if len(batches) != len(set(batches)): + raise RuntimeError(f"SGLang graph batches are duplicated: {batches}") + if not allow_multi_shape and batches != (1,): + raise RuntimeError( + "Foundry SGLang torch symmetric memory currently supports exactly " + "one batch-1 CUDA graph" + ) + if allow_multi_shape and batches != tuple(sorted(batches, reverse=True)): + raise RuntimeError( + f"SGLang graph capture order is not largest-to-smallest: {batches}" + ) + return tuple(parsed) +``` + +- [ ] **Step 4: Write failing typed FlashInfer ABI tests** + +```python +# tests/test_sglang_flashinfer_graph_abi.py +from __future__ import annotations + +from types import SimpleNamespace + +import pytest + +from foundry.integration.sglang.flashinfer_graph_abi import ( + MERGE_KERNEL_NAME, + PAGED_KERNEL_NAME, + extract_flashinfer_operand_records, +) + + +class Slot: + def __init__(self, name: str, ptr: int, numel: int) -> None: + self.name = name + self.data_ptr = ptr + self.numel = numel + self.element_size = 1 + + +def write_u64(raw: bytearray, offset: int, value: int) -> None: + raw[offset : offset + 8] = value.to_bytes(8, "little") + + +def make_state(): + plan = (66, 8, 864, 16, 0, 272, 544, 880, 816, 852, 0, 8650752, 928, 1, 1) + owners = ( + Slot("int_workspace", 0x600100000000, 8 * 1024 * 1024), + Slot("qo_indptr", 0x600110000000, 4096), + ) + shared = ( + Slot("paged_kv_indptr", 0x600120000000, 4096), + Slot("paged_kv_indices", 0x600130000000, 4096), + Slot("paged_kv_last_page_len", 0x600140000000, 4096), + Slot("float_workspace", 0x600150000000, 16 * 1024 * 1024), + ) + return SimpleNamespace( + batch_size=8, + plan_fingerprint=plan, + owners=owners, + shared_owners=shared, + ) + + +def make_graph(state) -> dict: + slots = { + slot.name: slot for slot in (*state.owners, *state.shared_owners) + } + plan = state.plan_fingerprint + paged = bytearray(376) + values = { + 0: 0x600180000000, + 56: 0x600190000000, + 64: 0x6001A0000000, + 72: slots["paged_kv_indices"].data_ptr, + 80: slots["paged_kv_indptr"].data_ptr, + 88: slots["paged_kv_last_page_len"].data_ptr, + 104: slots["qo_indptr"].data_ptr, + 112: slots["float_workspace"].data_ptr + plan[10], + 120: slots["float_workspace"].data_ptr + plan[11], + 296: slots["int_workspace"].data_ptr + plan[4], + 304: slots["int_workspace"].data_ptr + plan[5], + 312: slots["int_workspace"].data_ptr + plan[6], + 320: slots["int_workspace"].data_ptr + plan[7], + 328: slots["int_workspace"].data_ptr + plan[8], + 336: slots["int_workspace"].data_ptr + plan[12], + 344: slots["int_workspace"].data_ptr + plan[9], + 360: slots["int_workspace"].data_ptr + plan[2], + } + for offset, value in values.items(): + write_u64(paged, offset, value) + paged[352:356] = state.batch_size.to_bytes(4, "little") + paged[256:260] = (16).to_bytes(4, "little") + paged[372] = 1 + + merge_values = ( + values[112], + values[120], + values[320], + 0x600160000000, + 0, + state.batch_size, + values[360], + 16, + ) + merge_layout = ((0, 0, 8), (1, 8, 8), (2, 16, 8), (3, 24, 8), + (4, 32, 8), (5, 40, 4), (6, 48, 8), (7, 56, 4)) + + nodes = [] + for layer in range(36): + nodes.append( + { + "id": layer * 2, + "type": "KernelNode", + "params": { + "function_name": PAGED_KERNEL_NAME, + "kernelParams": [ + { + "index": 0, + "offset": 0, + "size": 376, + "value_hex": paged.hex(), + } + ], + "extra_argBuffer_hex": "", + }, + } + ) + nodes.append( + { + "id": layer * 2 + 1, + "type": "KernelNode", + "params": { + "function_name": MERGE_KERNEL_NAME, + "kernelParams": [ + { + "index": index, + "offset": offset, + "size": size, + "value_hex": int(value).to_bytes(size, "little").hex(), + } + for (index, offset, size), value in zip( + merge_layout, + merge_values, + strict=True, + ) + ], + "extra_argBuffer_hex": "", + }, + } + ) + return {"nodes": nodes} + + +def test_extracts_all_pinned_flashinfer_operands() -> None: + state = make_state() + records = extract_flashinfer_operand_records(make_graph(state), state) + + assert len(records) == 36 * (3 + 1 + 14 + 4 + 1) + assert {record.owner_slot for record in records} == { + "foundry_vmm", + "normalized_null", + "int_workspace", + "qo_indptr", + "paged_kv_indptr", + "paged_kv_indices", + "paged_kv_last_page_len", + "float_workspace", + } + + +def test_rejects_changed_paged_pointer() -> None: + state = make_state() + graph = make_graph(state) + raw = bytearray.fromhex( + graph["nodes"][0]["params"]["kernelParams"][0]["value_hex"] + ) + write_u64(raw, 296, 0x600170000000) + graph["nodes"][0]["params"]["kernelParams"][0]["value_hex"] = raw.hex() + + with pytest.raises(RuntimeError, match="request_indices"): + extract_flashinfer_operand_records(graph, state) +``` + +- [ ] **Step 5: Implement the exact pinned FlashInfer ABI** + +Create `flashinfer_graph_abi.py` with the exact two function names from the +pinned H100 graph, `PLAN_SCHEMA = "flashinfer.prefill.v15"`, and +`KERNEL_ABI_SCHEMA = +"flashinfer-0.6.15.post1.fa2.qwen3-8b-tp2-bf16-paged376"`. + +Define these typed descriptors: + +```python +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +from foundry.integration.sglang.sglang_shape_adapter import PLAN_SCHEMA +from foundry.integration.sglang.state_manifest import OperandSlotRecord + +KERNEL_ABI_SCHEMA = ( + "flashinfer-0.6.15.post1.fa2.qwen3-8b-tp2-bf16-paged376" +) +PAGED_KERNEL_NAME = ( + "_ZN10flashinfer34BatchPrefillWithPagedKVCacheKernelINS_12KernelTraits" + "ILNS_8MaskModeE0ELj16ELj1ELj2ELj8ELj8ELj1ELj4ELNS_15PosEncodingModeE0E" + "13__nv_bfloat16S4_S4_fiNS_16DefaultAttentionILb0ELb0ELb0ELb0EEEEE" + "11PagedParamsEEvT0_" +) +MERGE_KERNEL_NAME = ( + "_ZN10flashinfer41PersistentVariableLengthMergeStatesKernel" + "ILj8ELj16ELj8ELj4E13__nv_bfloat16S1_iEEvPT3_PfPT5_PT4_S4_jPjj" +) +PAGED_PARAM_LAYOUT = ((0, 0, 376),) +MERGE_PARAM_LAYOUT = ( + (0, 0, 8), (1, 8, 8), (2, 16, 8), (3, 24, 8), + (4, 32, 8), (5, 40, 4), (6, 48, 8), (7, 56, 4), +) +PREFILL_PLAN_INDEX = { + "padded_batch_size": 0, + "total_num_rows": 1, + "total_num_rows_offset": 2, + "cta_tile_q": 3, + "request_indices_offset": 4, + "qo_tile_indices_offset": 5, + "kv_tile_indices_offset": 6, + "merge_indptr_offset": 7, + "o_indptr_offset": 8, + "kv_chunk_size_ptr_offset": 9, + "v_offset": 10, + "s_offset": 11, + "block_valid_mask_offset": 12, + "enable_cuda_graph": 13, + "split_kv": 14, +} + + +@dataclass(frozen=True) +class PointerSpec: + name: str + parameter_index: int + value_byte_offset: int + ctype: str + owner_slot: str + plan_index: int | None + span_bytes: int + + +PAGED_POINTER_SPECS = ( + PointerSpec("paged_kv.indices", 0, 72, "int32_t*", "paged_kv_indices", None, 4), + PointerSpec("paged_kv.indptr", 0, 80, "int32_t*", "paged_kv_indptr", None, 4), + PointerSpec("paged_kv.last_page_len", 0, 88, "int32_t*", "paged_kv_last_page_len", None, 4), + PointerSpec("q_indptr", 0, 104, "int32_t*", "qo_indptr", None, 4), + PointerSpec("tmp_v", 0, 112, "nv_bfloat16*", "float_workspace", 10, 131072), + PointerSpec("tmp_s", 0, 120, "float*", "float_workspace", 11, 1024), + PointerSpec("request_indices", 0, 296, "int32_t*", "int_workspace", 4, 4), + PointerSpec("qo_tile_indices", 0, 304, "int32_t*", "int_workspace", 5, 4), + PointerSpec("kv_tile_indices", 0, 312, "int32_t*", "int_workspace", 6, 4), + PointerSpec("merge_indptr", 0, 320, "int32_t*", "int_workspace", 7, 4), + PointerSpec("o_indptr", 0, 328, "int32_t*", "int_workspace", 8, 4), + PointerSpec("block_valid_mask", 0, 336, "bool*", "int_workspace", 12, 1), + PointerSpec("kv_chunk_size", 0, 344, "int32_t*", "int_workspace", 9, 4), + PointerSpec("total_num_rows", 0, 360, "uint32_t*", "int_workspace", 2, 4), +) +MERGE_POINTER_SPECS = ( + PointerSpec("tmp_v", 0, 0, "nv_bfloat16*", "float_workspace", 10, 131072), + PointerSpec("tmp_s", 1, 0, "float*", "float_workspace", 11, 1024), + PointerSpec("merge_indptr", 2, 0, "int32_t*", "int_workspace", 7, 4), + PointerSpec("total_num_rows", 6, 0, "uint32_t*", "int_workspace", 2, 4), +) +_NULL_PAGED_POINTER_OFFSETS = (96, 152, 160, 176, 184, 192, 200, 208) +_STATIC_PAGED_POINTER_OFFSETS = ( + ("query", 0), + ("kv_cache_k", 56), + ("kv_cache_v", 64), +) +_FOUNDRY_BASE = 0x600000000000 +_FOUNDRY_END = _FOUNDRY_BASE + 256 * 1024**3 + + +def _read_u64(raw: bytes | bytearray, offset: int) -> int: + return int.from_bytes(raw[offset : offset + 8], "little") + + +def _write_u64(raw: bytearray, offset: int, value: int) -> None: + raw[offset : offset + 8] = value.to_bytes(8, "little") + + +def _owner_map(state: Any) -> dict[str, Any]: + return { + slot.name: slot for slot in (*state.owners, *state.shared_owners) + } + + +def _validate_layout(params: list[dict], expected: tuple[tuple[int, int, int], ...]) -> None: + actual = tuple( + (int(param["index"]), int(param["offset"]), int(param["size"])) + for param in params + ) + if actual != expected: + raise RuntimeError(f"FlashInfer kernel ABI changed: {actual} != {expected}") + + +def _span_bytes(spec: PointerSpec, state: Any) -> int: + padded = state.plan_fingerprint[0] + batch = state.batch_size + if spec.name in {"tmp_v", "tmp_s"}: + return spec.span_bytes * padded + if spec.name in { + "request_indices", + "qo_tile_indices", + "kv_tile_indices", + }: + return 4 * padded + if spec.name in {"merge_indptr", "o_indptr"}: + return 4 * (batch + 1) + if spec.name == "block_valid_mask": + return padded + return spec.span_bytes + + +def _record_specs( + node: dict, + specs: tuple[PointerSpec, ...], + state: Any, +) -> tuple[OperandSlotRecord, ...]: + owners = _owner_map(state) + params = node["params"]["kernelParams"] + records = [] + for spec in specs: + param = params[spec.parameter_index] + raw = bytes.fromhex(param["value_hex"]) + owner = owners[spec.owner_slot] + relative = ( + 0 + if spec.plan_index is None + else int(state.plan_fingerprint[spec.plan_index]) + ) + span = _span_bytes(spec, state) + capacity = owner.numel * owner.element_size + if relative < 0 or relative + span > capacity: + raise RuntimeError( + f"FlashInfer owner capacity is too small for {spec.name}" + ) + expected = owner.data_ptr + relative + actual = _read_u64(raw, spec.value_byte_offset) + if actual != expected: + raise RuntimeError( + f"FlashInfer pointer changed for {spec.name}: " + f"0x{actual:x} != 0x{expected:x}" + ) + records.append( + OperandSlotRecord( + name=spec.name, + kernel_abi=KERNEL_ABI_SCHEMA, + owner_slot=spec.owner_slot, + node_id=int(node["id"]), + source="kernelParams", + parameter_index=spec.parameter_index, + cuda_parameter_offset=int(param["offset"]), + value_byte_offset=spec.value_byte_offset, + ctype=spec.ctype, + owner_relative_offset=relative, + span_bytes=span, + saved_value=actual, + ) + ) + return tuple(records) + + +def extract_flashinfer_operand_records( + graph_data: dict, + state: Any, +) -> tuple[OperandSlotRecord, ...]: + plan = state.plan_fingerprint + if len(plan) != 15 or plan[13:] != (1, 1): + raise RuntimeError(f"Unsupported FlashInfer prefill plan: {plan}") + paged_nodes = [ + node + for node in graph_data["nodes"] + if node["type"] == "KernelNode" + and node["params"]["function_name"] == PAGED_KERNEL_NAME + ] + merge_nodes = [ + node + for node in graph_data["nodes"] + if node["type"] == "KernelNode" + and node["params"]["function_name"] == MERGE_KERNEL_NAME + ] + if len(paged_nodes) != 36 or len(merge_nodes) != 36: + raise RuntimeError( + "Expected 36 FlashInfer paged and merge nodes, got " + f"{len(paged_nodes)} and {len(merge_nodes)}" + ) + + records = [] + for paged_node, merge_node in zip(paged_nodes, merge_nodes, strict=True): + paged_params = paged_node["params"]["kernelParams"] + merge_params = merge_node["params"]["kernelParams"] + _validate_layout(paged_params, PAGED_PARAM_LAYOUT) + _validate_layout(merge_params, MERGE_PARAM_LAYOUT) + if paged_node["params"].get("extra_argBuffer_hex") or merge_node[ + "params" + ].get("extra_argBuffer_hex"): + raise RuntimeError("FlashInfer kernels use an unsupported arg buffer") + + paged_raw = bytes.fromhex(paged_params[0]["value_hex"]) + for offset in _NULL_PAGED_POINTER_OFFSETS: + if _read_u64(paged_raw, offset) != 0: + raise RuntimeError( + f"FlashInfer pointer at PagedParams+{offset} must be null" + ) + for name, offset in _STATIC_PAGED_POINTER_OFFSETS: + value = _read_u64(paged_raw, offset) + if not _FOUNDRY_BASE <= value < _FOUNDRY_END: + raise RuntimeError(f"FlashInfer {name} is outside Foundry VMM") + records.append( + OperandSlotRecord( + name=name, + kernel_abi=KERNEL_ABI_SCHEMA, + owner_slot="foundry_vmm", + node_id=int(paged_node["id"]), + source="kernelParams", + parameter_index=0, + cuda_parameter_offset=0, + value_byte_offset=offset, + ctype="void*", + owner_relative_offset=0, + span_bytes=1, + saved_value=value, + ) + ) + alibi = _read_u64(paged_raw, 168) + records.append( + OperandSlotRecord( + name="unused_alibi", + kernel_abi=KERNEL_ABI_SCHEMA, + owner_slot="normalized_null", + node_id=int(paged_node["id"]), + source="kernelParams", + parameter_index=0, + cuda_parameter_offset=0, + value_byte_offset=168, + ctype="float*", + owner_relative_offset=0, + span_bytes=0, + saved_value=alibi, + ) + ) + records.extend(_record_specs(paged_node, PAGED_POINTER_SPECS, state)) + records.extend(_record_specs(merge_node, MERGE_POINTER_SPECS, state)) + + merge_values = [ + int.from_bytes(bytes.fromhex(param["value_hex"]), "little") + for param in merge_params + ] + if _read_u64(paged_raw, 112) != merge_values[0]: + raise RuntimeError("FlashInfer tmp_v cross-node pointer mismatch") + if _read_u64(paged_raw, 120) != merge_values[1]: + raise RuntimeError("FlashInfer tmp_s cross-node pointer mismatch") + if _read_u64(paged_raw, 320) != merge_values[2]: + raise RuntimeError("FlashInfer merge_indptr pointer mismatch") + if _read_u64(paged_raw, 360) != merge_values[6]: + raise RuntimeError("FlashInfer total_num_rows pointer mismatch") + if int.from_bytes(paged_raw[352:356], "little") != state.batch_size: + raise RuntimeError("FlashInfer PagedParams batch size changed") + if merge_values[5] != state.batch_size: + raise RuntimeError("FlashInfer merge max_seq_len changed") + if int.from_bytes(paged_raw[256:260], "little") != 16: + raise RuntimeError("FlashInfer local query-head count changed") + if merge_values[7] != 16 or merge_values[4] != 0: + raise RuntimeError("FlashInfer merge scalar ABI changed") + if not _FOUNDRY_BASE <= merge_values[3] < _FOUNDRY_END: + raise RuntimeError("FlashInfer merge output is outside Foundry VMM") + records.append( + OperandSlotRecord( + name="merge_output", + kernel_abi=KERNEL_ABI_SCHEMA, + owner_slot="foundry_vmm", + node_id=int(merge_node["id"]), + source="kernelParams", + parameter_index=3, + cuda_parameter_offset=int(merge_params[3]["offset"]), + value_byte_offset=0, + ctype="nv_bfloat16*", + owner_relative_offset=0, + span_bytes=1, + saved_value=merge_values[3], + ) + ) + if paged_raw[372] != 1: + raise RuntimeError("FlashInfer partition_kv flag changed") + return tuple(records) + + +def relocate_flashinfer_operands( + graph_data: dict, + state: Any, + records: tuple[OperandSlotRecord, ...], +) -> None: + nodes = {int(node["id"]): node for node in graph_data["nodes"]} + owners = _owner_map(state) + for record in records: + if record.kernel_abi != KERNEL_ABI_SCHEMA: + raise RuntimeError( + f"Unsupported FlashInfer kernel ABI: {record.kernel_abi}" + ) + node = nodes[record.node_id] + param = node["params"]["kernelParams"][record.parameter_index] + raw = bytearray.fromhex(param["value_hex"]) + actual = _read_u64(raw, record.value_byte_offset) + if actual != record.saved_value: + raise RuntimeError(f"Saved FlashInfer operand changed: {record.name}") + if record.owner_slot == "normalized_null": + replacement = 0 + elif record.owner_slot == "foundry_vmm": + replacement = record.saved_value + else: + owner = owners[record.owner_slot] + replacement = owner.data_ptr + record.owner_relative_offset + capacity = owner.numel * owner.element_size + if record.owner_relative_offset + record.span_bytes > capacity: + raise RuntimeError( + f"Live FlashInfer owner is too small for {record.name}" + ) + _write_u64(raw, record.value_byte_offset, replacement) + param["value_hex"] = raw.hex() +``` + +- [ ] **Step 6: Run the FlashInfer ABI tests** + +```bash +python3 -m pytest tests/test_sglang_flashinfer_graph_abi.py -q +``` + +Expected: both tests pass. + +- [ ] **Step 7: Return typed communication operand records from graph validation** + +Update `_relocate_graph` so each allowed communication location is recorded: + +```python +def _communication_operand_records( + graph_filename: str, + graph_data: dict, +) -> tuple[OperandSlotRecord, ...]: + records = [] + for node in graph_data["nodes"]: + if node["type"] != "KernelNode": + continue + params = node["params"] + if _TWO_SHOT_SPECIALIZATION not in params["function_name"]: + continue + records.extend( + ( + OperandSlotRecord( + name="symm.buffer_ptrs_dev", + kernel_abi="torch.symm_mem.two_shot.bf16.align16.tp2", + owner_slot="communicator", + node_id=int(node["id"]), + source="kernelParams", + parameter_index=0, + cuda_parameter_offset=int( + params["kernelParams"][0]["offset"] + ), + value_byte_offset=0, + ctype="BFloat16**", + owner_relative_offset=0, + span_bytes=16, + saved_value=_decode_param(params["kernelParams"][0]), + ), + OperandSlotRecord( + name="symm.signal_pad_ptrs_dev", + kernel_abi="torch.symm_mem.two_shot.bf16.align16.tp2", + owner_slot="communicator", + node_id=int(node["id"]), + source="kernelParams", + parameter_index=3, + cuda_parameter_offset=int( + params["kernelParams"][3]["offset"] + ), + value_byte_offset=0, + ctype="uint32_t**", + owner_relative_offset=0, + span_bytes=16, + saved_value=_decode_param(params["kernelParams"][3]), + ), + ) + ) + if not records: + raise RuntimeError( + f"SGLang symmetric-memory graph has no two-shot kernels: " + f"{graph_filename}" + ) + return tuple(records) +``` + +Change `preserve_symmetric_graphs` to parse every full graph, retain existing +ABI validation, copy it, and return one `GraphOperandInventory` per filename. +Task 5 combines each inventory with the matching capsule: + +```python +records.append( + GraphOperandInventory( + graph_filename=filename, + batch_size=batch_size, + capture_index=index, + operand_slots=_communication_operand_records(filename, graph_data), + ) +) +``` + +- [ ] **Step 5: Make relocation validate the manifest order** + +```python +parsed = _validate_graph_files( + graph_filenames, + allow_multi_shape=allow_multi_shape, +) +if manifest is not None: + manifest_names = tuple(shape.graph_filename for shape in manifest.shapes) + parsed_names = tuple(filename for _index, _batch, filename in parsed) + if manifest_names != parsed_names: + raise RuntimeError( + f"SGLang shape manifest order mismatch: " + f"{manifest_names} != {parsed_names}" + ) +``` + +Keep all existing kernel-specialization, memcpy-layout, dependency, unique-ID, +buffer-capacity, access-policy, memset, and interior-pointer validation. + +- [ ] **Step 6: Run relocation tests** + +```bash +python3 -m pytest tests/test_sglang_symm_mem_graph.py -q +``` + +Expected: all tests pass, including default batch-1 rejection and opt-in +multi-shape acceptance. + +- [ ] **Step 7: Format and commit** + +```bash +pre-commit run --files \ + python/foundry/integration/sglang/symm_mem_graph.py \ + python/foundry/integration/sglang/state_manifest.py \ + tests/test_sglang_symm_mem_graph.py +git add \ + python/foundry/integration/sglang/symm_mem_graph.py \ + python/foundry/integration/sglang/state_manifest.py \ + tests/test_sglang_symm_mem_graph.py +git commit -s -m "feat: relocate multiple SGLang graph shapes" +git push -u origin cursor/vllm-tp-consolidation-5d04 +``` + +--- + +### Task 5: Wire shape capsules into the SGLang main backend + +**Files:** + +- Modify: `python/foundry/integration/sglang/main_backend.py:56-364` +- Modify: `python/foundry/integration/sglang/config.py` +- Create: `tests/test_sglang_main_backend.py` + +**Interfaces:** + +- Consumes: Tasks 1โ€“4. +- Produces: + - `_shape_states: dict[Any, ShapeReplayState]` + - `_multi_shape_enabled: bool` + - ordered SAVE/LOAD capsule creation + - replay-time capsule activation and planner refresh. +- Consumed by: Task 6 and H100 validation. + +- [ ] **Step 1: Add the feature-flag helper with failing tests** + +```python +# append to tests/test_sglang_main_backend.py +from foundry.integration.sglang.config import ( + is_symmetric_multi_shape_enabled, +) + + +def test_multi_shape_flag_defaults_off(monkeypatch) -> None: + monkeypatch.delenv("FOUNDRY_SGLANG_SYMM_MULTI_SHAPE", raising=False) + assert not is_symmetric_multi_shape_enabled() + + +def test_multi_shape_flag_accepts_only_one(monkeypatch) -> None: + monkeypatch.setenv("FOUNDRY_SGLANG_SYMM_MULTI_SHAPE", "1") + assert is_symmetric_multi_shape_enabled() + monkeypatch.setenv("FOUNDRY_SGLANG_SYMM_MULTI_SHAPE", "true") + with pytest.raises(RuntimeError, match="must be 0 or 1"): + is_symmetric_multi_shape_enabled() +``` + +Run: + +```bash +python3 -m pytest tests/test_sglang_main_backend.py -q +``` + +Expected: import fails because the helper is absent. + +- [ ] **Step 2: Implement the strict feature flag** + +```python +# python/foundry/integration/sglang/config.py +MULTI_SHAPE_ENV = "FOUNDRY_SGLANG_SYMM_MULTI_SHAPE" + + +def is_symmetric_multi_shape_enabled() -> bool: + value = os.environ.get(MULTI_SHAPE_ENV, "0") + if value not in {"0", "1"}: + raise RuntimeError(f"{MULTI_SHAPE_ENV} must be 0 or 1, got {value!r}") + return value == "1" +``` + +- [ ] **Step 3: Add a backend cleanup test with explicit SGLang stubs** + +```python +# tests/test_sglang_main_backend.py +from __future__ import annotations + +import importlib.util +import sys +from pathlib import Path +from types import ModuleType + + +def _stub(monkeypatch, name: str) -> ModuleType: + module = ModuleType(name) + monkeypatch.setitem(sys.modules, name, module) + return module + + +def _load_backend_module(monkeypatch): + for name in ( + "sglang", + "sglang.srt", + "sglang.srt.distributed", + "sglang.srt.distributed.device_communicators", + "sglang.srt.layers", + "sglang.srt.model_executor", + "sglang.srt.model_executor.runner_backend", + "sglang.srt.model_executor.runner_utils", + ): + _stub(monkeypatch, name) + + pynccl = _stub( + monkeypatch, + "sglang.srt.distributed.device_communicators.pynccl_allocator", + ) + pynccl.set_graph_pool_id = lambda pool: None + logits = _stub(monkeypatch, "sglang.srt.layers.logits_processor") + logits.LogitsProcessorOutput = type("LogitsProcessorOutput", (), {}) + base = _stub( + monkeypatch, + "sglang.srt.model_executor.runner_backend.base_cuda_graph_backend", + ) + base.BaseCudaGraphBackend = object + pool = _stub( + monkeypatch, + "sglang.srt.model_executor.runner_utils.pool", + ) + pool.get_or_create_global_graph_memory_pool = lambda device: (0, 0) + + path = ( + Path(__file__).parents[1] + / "python" + / "foundry" + / "integration" + / "sglang" + / "main_backend.py" + ) + spec = importlib.util.spec_from_file_location("tested_main_backend", 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 + + +class FakeState: + def __init__(self, capture_index: int) -> None: + self.capture_index = capture_index + + +def test_backend_cleanup_closes_every_shape_in_capture_order(monkeypatch) -> None: + module = _load_backend_module(monkeypatch) + events = [] + monkeypatch.setattr( + module, + "close_shape_state", + lambda state: events.append(state.capture_index), + ) + backend = module.FoundryMainCudaGraphBackend.__new__( + module.FoundryMainCudaGraphBackend + ) + backend._relocated_graph_dir = None + backend._graphs = {} + backend._outputs = {} + backend._pool = None + backend._pending = None + backend._graph_files = [] + backend._graph_load_paths = [] + backend._load_index = 0 + backend._shape_states = { + 8: FakeState(0), + 1: FakeState(1), + } + + backend.cleanup() + + assert events == [0, 1] + assert backend._shape_states == {} +``` + +Expected before implementation: FAIL because `_shape_states` is ignored. + +- [ ] **Step 4: Add capsule fields in `FoundryMainCudaGraphBackend.__init__`** + +```python +import json +from pathlib import Path + +from foundry.integration.sglang.config import is_symmetric_multi_shape_enabled +from foundry.integration.sglang.flashinfer_graph_abi import ( + KERNEL_ABI_SCHEMA, + extract_flashinfer_operand_records, + relocate_flashinfer_operands, +) +from foundry.integration.sglang.sglang_shape_adapter import ( + activate_shape_state, + close_shape_state, + create_shape_state, + prepare_shape_state, +) +from foundry.integration.sglang.shape_replay_state import ShapeReplayState +from foundry.integration.sglang.state_manifest import ( + ShapeStateManifest, + build_shape_record, + read_shape_state_manifest, + validate_shape_record, + write_shape_state_manifest, +) +from foundry.integration.sglang.symm_mem_graph import GraphOperandInventory +from foundry.integration.sglang.symm_mem_graph import FULL_GRAPH_DIR + +# in __init__ +self._multi_shape_enabled = ( + self._symmetric_memory_enabled and is_symmetric_multi_shape_enabled() +) +self._shape_states: dict[Any, ShapeReplayState] = {} +self._saved_shape_manifest: ShapeStateManifest | None = None +self._saved_records_by_batch = {} +self._operand_inventories: dict[str, GraphOperandInventory] = {} +``` + +- [ ] **Step 5: Create and retain a capsule inside `capture_one`** + +Refactor `_capture_one` and `_load_one` to return `(graph, outputs)` instead of +only mutating backend dictionaries: + +```python +def _capture_one( + self, + shape_key, + size: int, + forward_fn: Callable[[], Any], +) -> tuple[FoundryCUDAGraph, Any]: + if self._stream is None: + raise RuntimeError("Foundry capture session has no CUDA stream") + self._runner.device_module.synchronize() + self._runner.model_runner.tp_group.barrier() + graph = FoundryCUDAGraph() + with foundry_graph( + graph, + pool=self._pool, + stream=self._stream, + ): + output = forward_fn() + self._save_graph(graph, output, size) + return graph, output + + +def _load_one(self, shape_key, size: int) -> tuple[FoundryCUDAGraph, Any]: + if self._load_index >= len(self._graph_files): + raise RuntimeError("SGLang requested more graph shapes than were saved") + _index, filename, meta = self._graph_files[self._load_index] + if int(meta["key"]) != size: + raise RuntimeError( + f"SGLang graph order mismatch: requested {size}, " + f"archive has {meta['key']} ({filename})" + ) + result = FoundryCUDAGraph.load( + self._graph_load_paths[self._load_index], + pool=self._pool, + ) + if not isinstance(result, tuple): + raise RuntimeError(f"Loaded graph {filename} has no output tensors") + graph, tensors = result + self._load_index += 1 + return graph, _unpack_output(tensors) +``` + +Then use: + +```python +state = create_shape_state( + self._runner, + shape_key=shape_key, + capture_index=self._load_index + if mode == CUDAGraphExtensionMode.LOAD + else len(self._shape_states), + communicator=self._runner.model_runner.tp_group.torch_symm_mem_comm, + symm_handle=( + self._runner.model_runner.tp_group.torch_symm_mem_comm + ._foundry_symm_mem_handle + ), +) +if mode == CUDAGraphExtensionMode.LOAD: + validate_shape_record( + state, + self._saved_records_by_batch[state.batch_size], + plan_schema=PLAN_SCHEMA, + ) +with activate_shape_state(state): + run_graph_warmups( + synchronize=self._runner.device_module.synchronize, + barrier=self._runner.model_runner.tp_group.barrier, + forward=forward_fn, + post_warmup=post_warmup_hook, + ) + prepare_shape_state(state, None, for_capture=True) + if mode == CUDAGraphExtensionMode.SAVE: + graph, outputs = self._capture_one(shape_key, size, forward_fn) + elif mode == CUDAGraphExtensionMode.LOAD: + record = self._saved_records_by_batch[state.batch_size] + graph_path = Path(self._graph_load_paths[self._load_index]) + graph_data = json.loads(graph_path.read_text()) + flashinfer_records = tuple( + operand + for operand in record.operand_slots + if operand.kernel_abi == KERNEL_ABI_SCHEMA + ) + relocate_flashinfer_operands( + graph_data, + state, + flashinfer_records, + ) + graph_path.write_text(json.dumps(graph_data)) + graph, outputs = self._load_one(shape_key, size) + else: + raise RuntimeError(f"Unsupported Foundry mode: {mode.value}") +state.attach_graph(graph, outputs) +self._shape_states[shape_key] = state +self._graphs[shape_key] = graph +self._outputs[shape_key] = outputs +``` + +Use this block only when `_multi_shape_enabled`; retain the current batch-1 +branch byte-for-byte when the flag is off. + +- [ ] **Step 6: Write/read the per-shape manifest** + +In `_finish_save`, merge relocation records from Task 4 with capsule owner +records: + +```python +inventories = preserve_symmetric_graphs( + cfg.workspace_dir, + graph_filenames, + self._symmetric_memory_state(), + allow_multi_shape=True, +) +states_by_batch = { + state.batch_size: state for state in self._shape_states.values() +} +if set(states_by_batch) != { + inventory.batch_size for inventory in inventories +}: + raise RuntimeError("SGLang capsule and graph shape inventories differ") +shape_records_list = [] +for inventory in inventories: + state = states_by_batch[inventory.batch_size] + graph_path = ( + Path(cfg.workspace_dir) + / FULL_GRAPH_DIR + / inventory.graph_filename + ) + graph_data = json.loads(graph_path.read_text()) + flashinfer_operands = extract_flashinfer_operand_records(graph_data, state) + graph_path.write_text(json.dumps(graph_data)) + shape_records_list.append( + build_shape_record( + state, + graph_filename=inventory.graph_filename, + plan_schema=PLAN_SCHEMA, + operand_slots=( + inventory.operand_slots + flashinfer_operands + ), + ) + ) +shape_records = tuple(shape_records_list) +manifest = ShapeStateManifest( + backend="torch_symmetric_memory", + rank=self._symmetric_memory_state().rank, + world_size=self._symmetric_memory_state().world_size, + capture_order=tuple(record.batch_size for record in shape_records), + shapes=shape_records, +) +write_shape_state_manifest(cfg.workspace_dir, manifest) +``` + +In `_prepare_load`, read it before relocating: + +```python +self._saved_shape_manifest = read_shape_state_manifest(cfg.workspace_dir) +self._saved_records_by_batch = { + record.batch_size: record for record in self._saved_shape_manifest.shapes +} +self._graph_load_paths = relocate_symmetric_graphs( + cfg.workspace_dir, + filenames, + self._symmetric_memory_state(), + output_dir=self._relocated_graph_dir, + manifest=self._saved_shape_manifest, + allow_multi_shape=True, +) +``` + +- [ ] **Step 7: Refresh the selected capsule in `replay`** + +```python +def replay(self, shape_key, static_forward_batch, **kwargs) -> Any: + del kwargs + if not self._multi_shape_enabled: + self._graphs[shape_key].replay() + return self._outputs[shape_key] + + state = self._shape_states[shape_key] + with activate_shape_state(state): + prepare_shape_state( + state, + static_forward_batch, + for_capture=False, + ) + state.graph.replay() + return state.outputs +``` + +- [ ] **Step 8: Close capsules before clearing backend dictionaries** + +```python +def cleanup(self) -> None: + self._cleanup_relocated_graphs() + if self._graphs: + self._closed = True + for state in sorted( + self._shape_states.values(), + key=lambda item: item.capture_index, + ): + close_shape_state(state) + self._shape_states.clear() + self._graphs.clear() + self._outputs.clear() + self._pool = None + self._pending = None + self._graph_files = [] + self._graph_load_paths = [] + self._saved_shape_manifest = None + self._saved_records_by_batch = {} + self._operand_inventories = {} + self._load_index = 0 +``` + +- [ ] **Step 9: Run backend and existing compatibility tests** + +```bash +python3 -m pytest \ + tests/test_sglang_main_backend.py \ + tests/test_sglang_main_compat.py \ + tests/test_sglang_shape_replay_state.py \ + tests/test_sglang_state_manifest.py \ + tests/test_sglang_symm_mem_graph.py -q +``` + +Expected: all tests pass. + +- [ ] **Step 10: Confirm the default batch-1 path is unchanged** + +```bash +python3 -m pytest \ + tests/test_sglang_tp_recipe.py \ + tests/test_sglang_main_compat.py \ + tests/test_sglang_symm_mem_graph.py -q +``` + +Expected: existing batch-1 guard and contracts remain green with the flag +unset. + +- [ ] **Step 11: Format and commit** + +```bash +pre-commit run --files \ + python/foundry/integration/sglang/config.py \ + python/foundry/integration/sglang/main_backend.py \ + python/foundry/integration/sglang/shape_replay_state.py \ + python/foundry/integration/sglang/sglang_shape_adapter.py \ + python/foundry/integration/sglang/state_manifest.py \ + tests/test_sglang_main_backend.py +git add \ + python/foundry/integration/sglang/config.py \ + python/foundry/integration/sglang/main_backend.py \ + python/foundry/integration/sglang/shape_replay_state.py \ + python/foundry/integration/sglang/sglang_shape_adapter.py \ + python/foundry/integration/sglang/state_manifest.py \ + tests/test_sglang_main_backend.py +git commit -s -m "feat: materialize SGLang shape replay capsules" +git push -u origin cursor/vllm-tp-consolidation-5d04 +``` + +--- + +### Task 6: Wire the opt-in capture schedule and recipe guard + +**Files:** + +- Modify: `python/foundry/integration/sglang/config.py` +- Modify: `python/foundry/integration/sglang/hooks_main.py:35-41,145-180` +- Modify: `recipe/experimental/serve_qwen3-8b_sglang_tp.sh:20-60` +- Modify: `tests/test_sglang_tp_recipe.py` +- Modify: `tests/test_sglang_main_compat.py` + +**Interfaces:** + +- Produces: + - `symmetric_graph_batch_sizes() -> tuple[int, ...] | None` + - `_patch_capture_batch_sizes() -> None` +- Consumed by: SGLang's `DecodeCudaGraphRunner.__init__` through the patched + `get_batch_sizes_to_capture`. + +- [ ] **Step 1: Add failing exact-schedule parser tests** + +```python +def test_symmetric_graph_batch_sizes_are_strict(monkeypatch) -> None: + monkeypatch.setenv("FOUNDRY_SGLANG_SYMM_GRAPH_BATCHES", "1,8,32") + assert symmetric_graph_batch_sizes() == (1, 8, 32) + + monkeypatch.setenv("FOUNDRY_SGLANG_SYMM_GRAPH_BATCHES", "8,1") + with pytest.raises(RuntimeError, match="ascending"): + symmetric_graph_batch_sizes() + + monkeypatch.setenv("FOUNDRY_SGLANG_SYMM_GRAPH_BATCHES", "1,1") + with pytest.raises(RuntimeError, match="duplicate"): + symmetric_graph_batch_sizes() +``` + +- [ ] **Step 2: Implement the parser** + +```python +# config.py +MULTI_SHAPE_BATCHES_ENV = "FOUNDRY_SGLANG_SYMM_GRAPH_BATCHES" + + +def symmetric_graph_batch_sizes() -> tuple[int, ...] | None: + raw = os.environ.get(MULTI_SHAPE_BATCHES_ENV) + if raw is None: + return None + try: + batches = tuple(int(value) for value in raw.split(",")) + except ValueError as error: + raise RuntimeError( + f"{MULTI_SHAPE_BATCHES_ENV} must be comma-separated integers" + ) from error + if not batches or any(batch <= 0 for batch in batches): + raise RuntimeError(f"{MULTI_SHAPE_BATCHES_ENV} requires positive values") + if len(batches) != len(set(batches)): + raise RuntimeError(f"{MULTI_SHAPE_BATCHES_ENV} contains a duplicate") + if batches != tuple(sorted(batches)): + raise RuntimeError(f"{MULTI_SHAPE_BATCHES_ENV} must be ascending") + if batches[0] != 1: + raise RuntimeError(f"{MULTI_SHAPE_BATCHES_ENV} must include batch 1") + return batches +``` + +- [ ] **Step 3: Add the capture-size monkeypatch** + +```python +# hooks_main.py +from foundry.integration.sglang.config import ( + is_symmetric_multi_shape_enabled, + symmetric_graph_batch_sizes, +) + + +def install_hooks_main() -> None: + _patch_capture_batch_sizes() + _patch_init_torch_distributed() + _patch_memory_pool() + _patch_torch_symm_mem() + _patch_multimem_all_gather() + _patch_backend_resolver() + _patch_spawn_sites() + + +def _patch_capture_batch_sizes() -> None: + original = decode_runner.get_batch_sizes_to_capture + + @functools.wraps(original) + def patched(model_runner, captured_req_width): + capture_bs, compile_bs = original(model_runner, captured_req_width) + if not ( + model_runner.server_args.enable_torch_symm_mem + and is_symmetric_multi_shape_enabled() + ): + return capture_bs, compile_bs + requested = symmetric_graph_batch_sizes() + if requested is None: + return capture_bs, compile_bs + unknown = set(requested) - set(capture_bs) + if unknown: + raise RuntimeError( + f"Requested SGLang graph batches are unavailable: {sorted(unknown)}" + ) + filtered_compile = [batch for batch in compile_bs if batch in requested] + return list(requested), filtered_compile + + decode_runner.get_batch_sizes_to_capture = patched +``` + +- [ ] **Step 4: Add hook tests** + +Extend the fake `decode_runner` module in `tests/test_sglang_main_compat.py` with +`get_batch_sizes_to_capture`, run `install_hooks_main`, and assert: + +```python +monkeypatch.setenv("FOUNDRY_SGLANG_SYMM_MULTI_SHAPE", "1") +monkeypatch.setenv("FOUNDRY_SGLANG_SYMM_GRAPH_BATCHES", "1,8") +capture_bs, compile_bs = decode_runner.get_batch_sizes_to_capture( + SimpleNamespace( + server_args=SimpleNamespace(enable_torch_symm_mem=True) + ), + 1, +) +assert capture_bs == [1, 8] +assert compile_bs == [1, 8] +``` + +- [ ] **Step 5: Make the recipe opt-in guard conditional** + +```bash +if [[ "${SGLANG_ENABLE_TORCH_SYMM_MEM:-0}" == "1" ]]; then + if [[ "${FOUNDRY_SGLANG_SYMM_MULTI_SHAPE:-0}" == "1" ]]; then + if [[ -z "${FOUNDRY_SGLANG_SYMM_GRAPH_BATCHES:-}" ]]; then + echo "FOUNDRY_SGLANG_SYMM_GRAPH_BATCHES is required for multi-shape replay" >&2 + exit 1 + fi + CUDA_GRAPH_MAX_BS="${FOUNDRY_SGLANG_SYMM_GRAPH_BATCHES##*,}" + else + if [[ -n "${SGLANG_CUDA_GRAPH_MAX_BS:-}" && "${SGLANG_CUDA_GRAPH_MAX_BS}" != "1" ]]; then + echo "Foundry SGLang torch symmetric memory currently requires SGLANG_CUDA_GRAPH_MAX_BS=1" >&2 + exit 1 + fi + CUDA_GRAPH_MAX_BS=1 + fi + MODEL_ARGS+=( --enable-torch-symm-mem ) + COMMUNICATION_BACKEND="torch symmetric memory" +else + COMMUNICATION_BACKEND="NCCL P2P/IPC" +fi +``` + +- [ ] **Step 6: Add recipe RED/GREEN contracts** + +```python +def test_tp_recipe_allows_opt_in_multi_shape( + monkeypatch, + tmp_path, +) -> None: + monkeypatch.setenv("SGLANG_ENABLE_TORCH_SYMM_MEM", "1") + monkeypatch.setenv("FOUNDRY_SGLANG_SYMM_MULTI_SHAPE", "1") + monkeypatch.setenv("FOUNDRY_SGLANG_SYMM_GRAPH_BATCHES", "1,8") + + _env, args = _run_recipe(tmp_path, "--save") + + assert args[args.index("--cuda-graph-max-bs") + 1] == "8" + + +def test_tp_recipe_keeps_batch_one_default(monkeypatch, tmp_path) -> None: + monkeypatch.setenv("SGLANG_ENABLE_TORCH_SYMM_MEM", "1") + + _env, args = _run_recipe(tmp_path, "--save") + + assert args[args.index("--cuda-graph-max-bs") + 1] == "1" +``` + +- [ ] **Step 7: Run CPU contracts** + +```bash +python3 -m pytest \ + tests/test_sglang_tp_recipe.py \ + tests/test_sglang_main_compat.py \ + tests/test_sglang_main_backend.py -q +``` + +Expected: all tests pass. + +- [ ] **Step 8: Format and commit** + +```bash +pre-commit run --files \ + python/foundry/integration/sglang/config.py \ + python/foundry/integration/sglang/hooks_main.py \ + recipe/experimental/serve_qwen3-8b_sglang_tp.sh \ + tests/test_sglang_tp_recipe.py \ + tests/test_sglang_main_compat.py +git add \ + python/foundry/integration/sglang/config.py \ + python/foundry/integration/sglang/hooks_main.py \ + recipe/experimental/serve_qwen3-8b_sglang_tp.sh \ + tests/test_sglang_tp_recipe.py \ + tests/test_sglang_main_compat.py +git commit -s -m "feat: gate SGLang multi-shape replay" +git push -u origin cursor/vllm-tp-consolidation-5d04 +``` + +--- + +### Task 7: Extend the Modal harness to verify exact replay shapes + +**Files:** + +- Modify: `tests/modal_sglang_tp.py:83-123,398-432,499-608` +- Test: `tests/test_sglang_tp_recipe.py` + +**Interfaces:** + +- Consumes: `FOUNDRY_SGLANG_SYMM_GRAPH_BATCHES`. +- Produces: + - `EXPECTED_GRAPH_BATCHES: tuple[int, ...]` + - exact `replayed_graph_batches` acceptance check. + +- [ ] **Step 1: Add strict expected-batch parsing** + +```python +SERVE_ENV_NAMES = ( + "SGLANG_QUANTIZATION", + "SGLANG_CONTEXT_LENGTH", + "SGLANG_REASONING_PARSER", + "SGLANG_LINEAR_ATTN_BACKEND", + "SGLANG_MOE_RUNNER_BACKEND", + "SGLANG_ATTENTION_BACKEND", + "SGLANG_MEM_FRACTION_STATIC", + "SGLANG_CUDA_GRAPH_MAX_BS", + "SGLANG_ENABLE_TORCH_SYMM_MEM", + "FOUNDRY_SGLANG_SYMM_MULTI_SHAPE", + "FOUNDRY_SGLANG_SYMM_GRAPH_BATCHES", +) + + +def _expected_graph_batches() -> tuple[int, ...]: + raw = os.environ.get("FOUNDRY_SGLANG_SYMM_GRAPH_BATCHES") + if raw: + return tuple(int(value) for value in raw.split(",")) + if USES_TORCH_SYMM_MEM: + return (1,) + return ( + 1, 2, 4, 8, 12, 16, 24, 32, 40, 48, 56, 64, 72, 80, 88, 96, + 104, 112, 120, 128, 136, 144, 152, 160, 168, 176, 184, 192, + 200, 208, 216, 224, 232, 240, 248, 256, + ) + + +EXPECTED_GRAPH_BATCHES = _expected_graph_batches() +EXPECTED_GRAPH_COUNT = len(EXPECTED_GRAPH_BATCHES) +RUN_ID_OVERRIDE = os.environ.get("TP_RUN_ID") +USES_MULTI_SHAPE = ( + SERVE_ENV.get("FOUNDRY_SGLANG_SYMM_MULTI_SHAPE") == "1" +) +``` + +- [ ] **Step 2: Persist one deterministic request per expected graph batch** + +Add an internal request helper that submits exactly `batch_size` copies of a +short factual prompt concurrently and record the replay log after each batch: + +```python +def _generate( + prompt: str, + *, + max_new_tokens: int = MAX_NEW_TOKENS, +) -> str: + request = urllib.request.Request( + f"http://127.0.0.1:{PORT}/generate", + data=json.dumps( + { + "text": prompt, + "sampling_params": { + "temperature": 0.0, + "max_new_tokens": max_new_tokens, + }, + } + ).encode(), + headers={"Content-Type": "application/json"}, + ) + with urllib.request.urlopen(request, timeout=120) as response: + result = json.loads(response.read()) + return result["text"] if isinstance(result, dict) else result[0]["text"] + + +def _exercise_graph_batch(batch_size: int) -> None: + prompts = ["Answer with the number 7."] * batch_size + barrier = threading.Barrier(batch_size) + + def generate_after_barrier(prompt: str) -> str: + barrier.wait() + return _generate(prompt, max_new_tokens=1) + + with ThreadPoolExecutor(max_workers=batch_size) as executor: + outputs = list(executor.map(generate_after_barrier, prompts)) + if len(outputs) != batch_size or not all(output.strip() for output in outputs): + raise RuntimeError(f"Graph batch {batch_size} returned empty output") +``` + +Call `_exercise_graph_batch` for every expected batch in SAVE, SAVE2, and LOAD +after the standard deterministic prompts. + +- [ ] **Step 3: Replace the loose replay assertion** + +Include the new manifests in `_archive_fingerprints` and the required archive +set: + +```python +required_archive_files = { + "graph_manifest.json", + "fatbin_image_packed.img", + "final_alloc_offset.json", +} +if USES_TORCH_SYMM_MEM: + required_archive_files.update( + {"sglang_graph_backend.json", "symmetric_memory_state.json"} + ) +if USES_MULTI_SHAPE: + required_archive_files.add("sglang_shape_state.json") + +files.extend( + path + for path in ( + rank_dir / "graph_manifest.json", + rank_dir / "fatbin_image_packed.img", + rank_dir / "final_alloc_offset.json", + rank_dir / "sglang_graph_backend.json", + rank_dir / "symmetric_memory_state.json", + rank_dir / "sglang_shape_state.json", + ) + if path.exists() +) +``` + +Then replace the loose batch-size condition: + +```python +replayed_graph_batches = set(results["load"].get("graph_replay_batch_sizes", [])) + +"restored_graph_replay_observed": ( + set(EXPECTED_GRAPH_BATCHES) <= replayed_graph_batches +), +``` + +Also include `EXPECTED_GRAPH_BATCHES` in the JSON report. + +Persist the aggregate report to the retained run directory: + +```python +@app.function( + image=modal.Image.debian_slim(), + volumes={ARCHIVE_ROOT: archive_volume}, +) +def persist_validation_report(run_id: str, report: dict) -> None: + report_path = Path(RUNS_ROOT) / run_id / "validation_report.json" + report_path.write_text(json.dumps(report, indent=2, sort_keys=True)) + archive_volume.commit() + + +# Replace the current run-id assignment inside `main`. +run_id = RUN_ID_OVERRIDE or f"{safe_ref}-{time.time_ns()}" + +# Insert immediately after the complete `report` dictionary is built. +persist_validation_report.remote(run_id, report) +``` + +- [ ] **Step 4: Keep eager isolation checks** + +Replace the fixed eight-request fallback with one request above the largest +captured shape, using one generated token to bound test cost: + +```python +EAGER_FALLBACK_BATCH = max(EXPECTED_GRAPH_BATCHES) + 1 +EAGER_FALLBACK_PROMPTS = ["Answer with the number 11."] * EAGER_FALLBACK_BATCH + + +def _generate_eager_fallback() -> list[str]: + barrier = threading.Barrier(EAGER_FALLBACK_BATCH) + + def generate_after_barrier(prompt: str) -> str: + barrier.wait() + return _generate(prompt, max_new_tokens=1) + + with ThreadPoolExecutor(max_workers=EAGER_FALLBACK_BATCH) as executor: + return list(executor.map(generate_after_barrier, EAGER_FALLBACK_PROMPTS)) +``` + +Retain `post_concurrent_graph_outputs_match` after this eager request and keep +byte-identical concurrent output text as a non-gating observation. + +- [ ] **Step 5: Run syntax and local tests** + +```bash +python3 -m compileall -q tests/modal_sglang_tp.py +python3 -m pytest tests/test_sglang_tp_recipe.py -q +pre-commit run --files tests/modal_sglang_tp.py tests/test_sglang_tp_recipe.py +``` + +Expected: all commands exit 0. + +- [ ] **Step 6: Commit** + +```bash +git add tests/modal_sglang_tp.py tests/test_sglang_tp_recipe.py +git commit -s -m "test: verify every SGLang replay shape" +git push -u origin cursor/vllm-tp-consolidation-5d04 +``` + +--- + +### Task 8: Validate the two- and three-shape H100 milestones + +**Files:** + +- Modify after successful runs: `docs/sglang/overview.md` +- Modify after successful runs: `recipe/experimental/README.md` + +**Interfaces:** + +- Consumes: Tasks 1โ€“7. +- Produces: retained Modal run IDs and measured state-memory overhead. + +- [ ] **Step 1: Run the 1+8 milestone** + +```bash +FOUNDRY_SGLANG_SYMM_MULTI_SHAPE=1 \ +FOUNDRY_SGLANG_SYMM_GRAPH_BATCHES=1,8 \ +SGLANG_ENABLE_TORCH_SYMM_MEM=1 \ +TP_KEEP_ARTIFACTS=1 \ +TP_RUN_ID="sglang-symm-1-8-$(git rev-parse --short HEAD)" \ +FOUNDRY_REF=$(git rev-parse HEAD) \ +modal run --env rahul-dev tests/modal_sglang_tp.py +``` + +Expected: exit 0; every check is true; graph batches 1 and 8 both appear in +LOAD replay logs; SAVE/SAVE2 fingerprints and offsets match. + +- [ ] **Step 2: Record the retained run and memory delta** + +Retrieve the retained `validation_report.json`, then read the exact values that +must be copied into the milestone evidence: + +```bash +RUN_ID="sglang-symm-1-8-$(git rev-parse --short HEAD)" +modal volume get --env rahul-dev --force \ + foundry-sglang-tp-validation-archive \ + "runs/${RUN_ID}/validation_report.json" \ + /tmp/sglang-1-8-validation.json +python3 - /tmp/sglang-1-8-validation.json <<'PY' +import json +import sys + +report = json.load(open(sys.argv[1])) +offset = report["results"]["save2"]["rank_offsets"]["0"] +batch_one_offset = 68555898880 +print("run_id:", report["run_id"]) +print("final_offset:", offset) +print("added_vmm_bytes_vs_batch1:", offset - batch_one_offset) +print("checks:", report["checks"]) +PY +``` + +Add those exact run ID and numeric values to the milestone table in +`docs/sglang/overview.md`. Do not commit milestone evidence unless every check +printed above is `true`. + +- [ ] **Step 3: Run the 1+8+32 milestone** + +```bash +FOUNDRY_SGLANG_SYMM_MULTI_SHAPE=1 \ +FOUNDRY_SGLANG_SYMM_GRAPH_BATCHES=1,8,32 \ +SGLANG_ENABLE_TORCH_SYMM_MEM=1 \ +TP_KEEP_ARTIFACTS=1 \ +TP_RUN_ID="sglang-symm-1-8-32-$(git rev-parse --short HEAD)" \ +FOUNDRY_REF=$(git rev-parse HEAD) \ +modal run --env rahul-dev tests/modal_sglang_tp.py +``` + +Expected: exit 0 with graph batches 1, 8, and 32 observed. + +- [ ] **Step 4: Run the focused CPU and quality gates after H100 success** + +```bash +python3 -m pytest \ + tests/test_sglang_shape_replay_state.py \ + tests/test_sglang_state_manifest.py \ + tests/test_sglang_main_backend.py \ + tests/test_sglang_symm_mem_graph.py \ + tests/test_sglang_tp_recipe.py \ + tests/test_sglang_main_compat.py -q +pre-commit run --all-files +git diff --check +``` + +Expected: all commands exit 0. + +- [ ] **Step 5: Request code review and commit milestone evidence** + +Request a fresh code-review subagent using the current branch diff. Resolve +every Critical or Important finding, rerun the focused gates, then: + +```bash +git add docs/sglang/overview.md recipe/experimental/README.md +git commit -s -m "docs: record SGLang multi-shape milestones" +git push -u origin cursor/vllm-tp-consolidation-5d04 +``` + +Do not proceed to Task 9 unless both milestone matrices pass. + +--- + +### Task 9: Run the full inventory and promote multi-shape support + +**Files:** + +- Modify: `recipe/experimental/serve_qwen3-8b_sglang_tp.sh` +- Modify: `tests/modal_sglang_tp.py` +- Modify: `README.md` +- Modify: `ROADMAP.md` +- Modify: `RELEASE.md` +- Modify: `docs/sglang/overview.md` +- Modify: `docs/sglang/hooks.md` +- Modify: `recipe/experimental/README.md` +- Modify: `recipe/sglang/README.md` + +**Interfaces:** + +- Consumes: successful Tasks 1โ€“8. +- Produces: full shape support under the opt-in flag, then removal of the + batch-1 default only after every acceptance gate passes. + +- [ ] **Step 1: Run the full pinned shape inventory** + +```bash +FOUNDRY_SGLANG_SYMM_MULTI_SHAPE=1 \ +FOUNDRY_SGLANG_SYMM_GRAPH_BATCHES=1,2,4,8,12,16,24,32,40,48,56,64,72,80,88,96,104,112,120,128,136,144,152,160,168,176,184,192,200,208,216,224,232,240,248,256 \ +SGLANG_ENABLE_TORCH_SYMM_MEM=1 \ +TP_KEEP_ARTIFACTS=1 \ +FOUNDRY_REF=$(git rev-parse HEAD) \ +modal run --env rahul-dev tests/modal_sglang_tp.py +``` + +Expected: exit 0; all 36 graph batches restored and replayed on both ranks; +all checks true; no runtime errors; `mem_fraction_static=0.8` unchanged. + +- [ ] **Step 2: Run the unchanged batch-1 symmetric matrix** + +```bash +SGLANG_ENABLE_TORCH_SYMM_MEM=1 \ +TP_KEEP_ARTIFACTS=1 \ +FOUNDRY_REF=$(git rev-parse HEAD) \ +modal run --env rahul-dev tests/modal_sglang_tp.py +``` + +Expected: exit 0 and the current one-graph proof remains green. + +- [ ] **Step 3: Run the default NCCL SGLang matrix** + +```bash +TP_KEEP_ARTIFACTS=1 \ +FOUNDRY_REF=$(git rev-parse HEAD) \ +modal run --env rahul-dev tests/modal_sglang_tp.py +``` + +Expected: exit 0 with the full NCCL graph inventory and P2P/IPC channels. + +- [ ] **Step 4: Run the vLLM TP regression** + +```bash +TP_KEEP_ARTIFACTS=1 \ +FOUNDRY_REF=$(git rev-parse HEAD) \ +modal run --env rahul-dev tests/modal_vllm_tp.py +``` + +Expected: exit 0; 64 graphs restored per rank; deterministic outputs and +offsets remain reproducible. + +- [ ] **Step 5: Promote the recipe only after all four H100 matrices pass** + +Change the symmetric recipe branch so multi-shape is selected without the +temporary feature flag, while preserving explicit TP=2 and pinned-version +guards. Remove `FOUNDRY_SGLANG_SYMM_MULTI_SHAPE` from the recipe, but keep the +environment variable accepted for one release as a no-op compatibility input. + +Update the roadmap item to: + +```markdown +- [x] Validate torch symmetric-memory TP=2 across the pinned graph inventory +``` + +Document: + +- the retained full-inventory run ID; +- graph count and final offsets; +- semantic fingerprint result; +- exact model/SGLang/PyTorch/CUDA/H100 scope; +- measured planner-memory overhead; +- the continued TP=2 limitation. + +- [ ] **Step 6: Run final local verification** + +```bash +python3 -m pytest \ + tests/test_sglang_shape_replay_state.py \ + tests/test_sglang_state_manifest.py \ + tests/test_sglang_main_backend.py \ + tests/test_sglang_symm_mem_graph.py \ + tests/test_sglang_tp_recipe.py \ + tests/test_sglang_main_compat.py \ + tests/test_vllm_tp_recipe.py -q +pre-commit run --all-files +git diff --check +git status --short --branch +``` + +Expected: pytest has zero failures, every pre-commit hook passes, diff check +passes, and only the intended final documentation/recipe changes are present. + +- [ ] **Step 7: Request final review, commit, push, and update evidence** + +Run a fresh code review. Resolve every Critical or Important issue and repeat +Step 6. Then: + +```bash +git add \ + README.md ROADMAP.md RELEASE.md \ + docs/sglang/overview.md docs/sglang/hooks.md \ + recipe/experimental/README.md recipe/sglang/README.md \ + recipe/experimental/serve_qwen3-8b_sglang_tp.sh \ + tests/modal_sglang_tp.py +git commit -s -m "feat: validate multi-shape SGLang symmetric replay" +git push -u origin cursor/vllm-tp-consolidation-5d04 +``` + +Update the existing pull-request description with the final retained run IDs, +exact verification commands, supported scope, and measured memory overhead. diff --git a/docs/superpowers/specs/2026-07-23-sglang-multishape-symmetric-state-design.md b/docs/superpowers/specs/2026-07-23-sglang-multishape-symmetric-state-design.md new file mode 100644 index 00000000..e98ff960 --- /dev/null +++ b/docs/superpowers/specs/2026-07-23-sglang-multishape-symmetric-state-design.md @@ -0,0 +1,263 @@ +# SGLang Multi-Shape Symmetric Replay Design + +**Date:** 2026-07-23 +**Status:** Approved + +## Context + +Foundry currently supports SGLang torch symmetric memory for TP=2 with one +batch-1 CUDA graph. SAVE preserves the full graph and records process-local +communication operands. LOAD rebuilds the SGLang warmup state and relocates +those operands to the live communicator. + +The one-shape path passes the complete 2ร—H100 baseline/SAVE/SAVE2/LOAD matrix. +The same batch-1 graph becomes unsafe when captured after larger shapes. Those +later captures refer to mutable FlashInfer planner and static metadata state +created by earlier captures in SGLang's process-wide graph pool. Recreating the +same allocation cursor is insufficient: it reproduces addresses, but not +ownership, lifetime, or the active contents of every mutable external buffer. + +The current batch-1 restriction remains the safe production behavior until +this design passes the multi-shape acceptance tests. + +## Goals + +1. Restore all SGLang decode graph shapes used by the pinned Qwen3-8B TP=2 + recipe without recapture or silent corruption. +2. Preserve SGLang's shared graph execution pool so graph memory does not scale + linearly with the number of shapes. +3. Give each shape explicit ownership of every mutable attention/planner object + referenced by its graph. +4. Rebuild and rebind shape state through typed integration code, not generic + pointer guessing. +5. Fail before replay when an upstream ABI, backend, graph topology, or state + inventory is incompatible. +6. Keep the validated NCCL, vLLM TP, and batch-1 symmetric paths unchanged. + +## Non-Goals + +- TP sizes other than 2. +- Multicast symmetric-memory kernels used by TP=4/6/8. +- Cross-node tensor parallelism. +- Generic serialization of arbitrary CUDA library internals. +- Replacing FlashInfer or changing SGLang's attention algorithm. +- Removing the current batch-1 safety gate before full H100 validation passes. + +## Considered Approaches + +### Dedicated graph pool per shape + +This provides strong isolation, but duplicates graph execution memory for every +shape. SGLang deliberately captures large-to-small in one pool so smaller +graphs can reuse the largest allocation. Replacing that with 36 pools risks +exhausting the memory left after model weights and the KV cache. + +### One padded maximum-batch graph + +This removes cross-shape state, but a batch-1 request would execute work sized +for the maximum batch. The predictable correctness does not justify the +small-batch latency and throughput loss. + +### Shared graph pool with per-shape state capsules + +This is the selected design. Compute and graph execution memory remain shared. +Only mutable planner metadata, owner objects, and explicit relocation metadata +are isolated by shape. This follows SGLang's existing per-batch wrapper model +while making ownership and replay preparation explicit. + +## Architecture + +### `ShapeReplayState` + +Add an integration-owned state capsule for each SGLang graph key. It contains: + +- the graph shape key and capture order; +- the SGLang/FlashInfer wrapper objects used by that shape; +- all tensors that own planner arrays or static metadata referenced by the + graph; +- the live symmetric-memory communicator handle; +- the saved external-operand inventory and its live relocation map; +- the loaded graph and output tensor owners; +- a lifecycle state: `created`, `planned`, `captured`/`loaded`, or `closed`. + +The capsule retains owners for at least as long as its loaded CUDA graph. +Closing the backend releases graphs before releasing capsule-owned tensors. + +### Shared versus isolated memory + +The following remain shared: + +- model weights; +- KV-cache pools; +- SGLang's graph execution pool; +- immutable compiled kernels and packed fatbins; +- immutable model configuration. + +The following become shape-owned: + +- FlashInfer integer planner workspaces and arrays; +- mutable planner metadata consumed by captured attention kernels; +- shape-specific static input views whose addresses appear in the graph; +- any mutable library workspace proven to be read by only one shape. + +The design must not duplicate model weights, KV-cache storage, or the full graph +pool. + +### SGLang adapter + +Add a narrow adapter around the pinned SGLang attention backend: + +```python +create_shape_state(shape_key, forward_batch) -> ShapeReplayState +prepare_shape_state(state, forward_batch, *, for_capture: bool) -> None +activate_shape_state(state) -> context manager +close_shape_state(state) -> None +``` + +The adapter uses SGLang's existing `decode_cuda_graph_metadata[bs]` and planner +APIs. If the pinned SGLang revision cannot expose an owner through a stable +attribute, the adapter fails with a version-specific error. It does not scan +Python objects heuristically for tensors. + +### State manifest + +Each rank writes a versioned state manifest alongside the graph manifest. For +each shape it records: + +- graph filename and shape key; +- capture-order index; +- backend, dtype, rank, and world size; +- named external operand slots and their saved values; +- expected kernel specialization and parameter layout; +- expected copy/dependency topology; +- buffer capacities; +- a semantic fingerprint of the state schema. + +Raw mutable buffer contents are not treated as durable serialization. LOAD +rebuilds them through the SGLang planner, then verifies the live owners and +relocates graph operands. + +### Operand relocation + +Relocation remains allowlist-based: + +1. Identify an operand by graph node, parameter index or struct offset, and a + typed state-manifest slot. +2. Validate the saved function specialization, parameter layout, dependency + edges, copy geometry, pointer range, rank, world size, dtype, and capacity. +3. Replace it with the corresponding live capsule operand. +4. Reject any saved or live external pointer found outside the complete + allowlist. + +Foundry never infers pointer semantics from a numeric address alone. + +## SAVE Flow + +1. Initialize the common model, KV cache, communicator, and shared graph pool. +2. Walk graph shapes in SGLang's normal largest-to-smallest order. +3. Create the shape capsule and its dedicated mutable planner state. +4. Run the two SGLang warmups and the post-warmup hook. +5. Synchronize the device and TP ranks. +6. Capture the graph using the shared graph pool while the shape capsule is + active. +7. Save full graph JSON and build the typed external-operand inventory. +8. Retain the capsule so later captures cannot invalidate its owners. +9. After all shapes, write the state manifest, graph manifest, packed fatbins, + and final allocation offset. + +SAVE fails before publishing the manifest if a graph references mutable state +that is not owned by either the shared immutable state or its shape capsule. + +## LOAD Flow + +1. Validate the saved backend and state-manifest versions before GPU + preallocation. +2. Recreate common model, KV-cache, communicator, and shared graph-pool state. +3. Walk shapes in the saved capture order. +4. Create the live shape capsule. +5. Run the same two warmups and planner preparation for that shape. +6. Synchronize the device and TP ranks. +7. Validate and relocate the saved graph to the live capsule and communicator. +8. Load the graph and retain graph outputs and capsule owners together. +9. Verify the final allocation offset and the number of loaded shapes. + +LOAD does not jump directly to the final allocation offset on this path. It +replays shape initialization and graph allocation events in SAVE order. + +## Replay Flow + +1. Select the capsule by the padded SGLang graph key. +2. Enter `activate_shape_state(state)`. +3. Refresh planner metadata from the live request using + `prepare_shape_state(..., for_capture=False)`. +4. Copy request data into the normal static graph inputs. +5. Replay the loaded graph on SGLang's current forward stream. +6. Keep the capsule active until output consumers have observed the graph + completion. + +Eager fallback uses separate transient planner state and must not mutate a graph +capsule. A graph replay immediately after eager fallback is an acceptance test. + +## Error Handling + +The integration raises before graph construction or replay for: + +- missing, duplicate, or reordered shape records; +- backend, rank, world-size, dtype, or capacity mismatch; +- unsupported kernel specialization; +- changed parameter or copy layout; +- unknown external operands; +- a live owner whose address or size does not match its manifest slot; +- planner initialization that does not reach a synchronized ready state; +- final allocation-offset or graph-count mismatch. + +Errors identify the rank, shape key, graph file, and named state slot without +logging raw GPU addresses. + +## Validation Strategy + +### CPU contracts + +- state-capsule lifecycle and owner retention; +- manifest round-trip and version rejection; +- exact shape ordering; +- relocation allowlist completeness; +- malformed kernel, copy, dependency, and pointer-range rejection; +- eager-state isolation; +- cleanup ordering. + +### Focused H100 progression + +1. Shapes 1 and 8. +2. Shapes 1, 8, and 32. +3. The full pinned SGLang shape inventory. + +Each step runs baseline, SAVE, SAVE2, and fresh LOAD on 2ร—H100. + +### Full acceptance criteria + +- TP=2 Qwen3-8B on the pinned SGLang/PyTorch/CUDA stack. +- All configured graph shapes saved and restored on both ranks. +- SAVE and SAVE2 have identical semantic fingerprints and per-rank offsets. +- LOAD performs no graph recapture. +- Deterministic sequential outputs match baseline byte for byte. +- Concurrent outputs are nonempty; a deterministic graph request after + concurrent eager work still matches baseline byte for byte. +- Every configured graph batch size appears in replay logs. +- No CUDA error, Xid, illegal access, NCCL warning/error, or scheduler + exception. +- The existing model/KV-cache configuration still fits at + `mem_fraction_static=0.8`; the implementation does not reduce KV capacity to + buy planner isolation. +- Default NCCL SGLang TP, current batch-1 symmetric TP, vLLM TP, and local + quality gates remain green. + +## Rollout + +1. Implement behind `FOUNDRY_SGLANG_SYMM_MULTI_SHAPE=1`. +2. Keep the current TP=2/batch-1 restriction as the default. +3. Validate the 1+8 and 1+8+32 milestones. +4. Run the full retained-artifact H100 matrix and code review. +5. Remove the batch-1 recipe guard only after every acceptance criterion + passes. +6. Keep explicit TP=2 and pinned-version guards until separately generalized. diff --git a/docs/vllm/direct-edits.md b/docs/vllm/direct-edits.md index a1e63f1f..a18cf1e7 100644 --- a/docs/vllm/direct-edits.md +++ b/docs/vllm/direct-edits.md @@ -55,6 +55,8 @@ class CompilationConfig: This is the earliest activation point in the parent process. As soon as a `CompilationConfig` is fully constructed, foundry's runtime patches are in place. +> Note: vLLM folds `graph_extension_config_path` into `CompilationConfig.compute_hash()` (the torch.compile cache key) even though it never affects codegen. This is harmless **as long as every SAVE/LOAD phase passes the same path string** โ€” run all phases from one canonical CWD (or always an absolute path). If pass 1 and pass 2 see the path via different mount aliases (`/home/...` vs `/data/...`), pass 2 misses pass 1's warm cache and inductor recompiles inside the cuda-graph capture window โ†’ `cudaErrorStreamCaptureInvalidated`. We keep vLLM unpatched and rely on consistent invocation instead; the EP recipe scripts are launched from the workspace root. (An optional `ignored_factors` exclusion would harden this at the source but is not applied.) + ## 3. `vllm/v1/engine/core.py` (~6 lines) ```python diff --git a/docs/vllm/overview.md b/docs/vllm/overview.md index a96e9e87..67d9922b 100644 --- a/docs/vllm/overview.md +++ b/docs/vllm/overview.md @@ -8,6 +8,9 @@ Validated on: - DP > 1 (multi-GPU dense) - Mixture-of-experts with DeepEP all-to-all (expert parallelism) +Experimentally validated on Qwen3-8B TP=2 with NCCL P2P/IPC. This is a pinned +configuration, not official/general TP support. + ## Critical invariants (read this first) Hard-won lessons. Each of these failure modes cost real debugging time. Future changes touching the integration must respect them. @@ -38,6 +41,29 @@ In reality, NVSHMEM's symmetric heap goes through the **large-alloc branch** of All MoE variants (quantized, unquantized) and dense models now share the same path: `preallocate_for_load_mode` โ†’ `do_original_load()` (weight load + `prepare_communication_buffer_for_model` + NVSHMEM init post-hook) โ†’ `start_graph_builds`. The background template builds were previously kicked off *before* `do_original_load` to overlap with weight IO; that was net-negative due to driver contention, so they now overlap with the cheaper post-load init work instead. +### 5. The experimental TP recipe must use one NCCL path in every phase + +vLLM has several higher-priority collective implementations besides PyNCCL: +custom all-reduce, PyTorch symmetric memory, NCCL symmetric memory, standalone +FlashInfer all-reduce, and an O2 compiler pass that fuses all-reduce with +RMSNorm. Changing implementations between baseline, SAVE, and LOAD changes both +the captured kernels and their memory/IPC setup. + +The validated TP recipe disables those alternatives and keeps +`NCCL_CUMEM_ENABLE=0` plus `NCCL_NVLS_ENABLE=0` in every phase. It also sets +`pass_config.fuse_allreduce_rms=false`; otherwise a cached fused graph can lazily +create its FlashInfer workspace from inside the second SAVE's capture and abort. +Pass all nested settings in one `--compilation-config` JSON objectโ€”mixing the +`-cc.*` and `--compilation-config.*` forms can replace earlier values and silently +drop `graph_extension_config_path`. + +This makes one TP=2 configuration reproducible, but it does not remove NCCL's +general initialization nondeterminism. Upstream +[Discussion #5](https://github.com/orgs/foundry-org/discussions/5) reports a +successful internal torch symmetric-memory prototype but no official +implementation; support is tracked in +[issue #6](https://github.com/foundry-org/foundry/issues/6). + ## How to use ### 1. Install foundry + vLLM @@ -95,8 +121,11 @@ vLLM's lifecycle splits across the engine-core process and worker processes. Fou 1. Same VMM setup as pass 1. 2. `_initialize_kv_caches` skips the profile forward and uses saved block counts. -3. Same deterministic allocation order โ‡’ same `final_alloc_offset`. -4. Captured graphs are byte-identical to pass 1's (already on disk โ€” pass 2 overwrites them). +3. The deterministic allocation order produces the archive used by LOAD. +4. Captured graphs overwrite pass 1's seed archive. Repeating this deterministic + pass is optional validation: graph inventories and `final_alloc_offset` must + match, although debugging JSON contains process-specific values and need not + be byte-identical. **LOAD**: @@ -126,6 +155,7 @@ vLLM's lifecycle splits across the engine-core process and worker processes. Fou |---|---| | Dense decode (Qwen3, Llama, โ€ฆ) | primary test target | | DP > 1 | multi-GPU dense | +| Experimental TP=2 | dense NCCL P2P/IPC; not official/general TP support | | MoE with DeepEP | unified path; no split-load | | `torch.compile` (full-graph) | runs on SAVE; `do_not_compile=True` on LOAD (faster startup; graphs replay via `CUDAGraphWrapper`) | diff --git a/include/CUDAGraphInternal.h b/include/CUDAGraphInternal.h index 52c53789..d3e53942 100644 --- a/include/CUDAGraphInternal.h +++ b/include/CUDAGraphInternal.h @@ -21,6 +21,10 @@ struct PendingGraphLoads { MempoolId_t pool; c10::DeviceIndex dev; CUDAGeneratorStateRegistry* registry = nullptr; // for deferred generator registration + // True if any graph in this pending set consumed Philox state. Generator + // state is shared across entries, so this must be decided for the full set + // before the first lazy seed/offset tensor allocation. + bool has_rng_generators = false; // Signaled when background graph building (Phase 2) completes. // finish_graph_loads_impl waits on this before allocator replay. diff --git a/pyproject.toml b/pyproject.toml index b70e5572..718d5008 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -23,6 +23,9 @@ dev = [ "pytest", ] +[project.entry-points."sglang.srt.plugins"] +foundry = "foundry.integration.sglang.plugin:register" + [tool.setuptools] packages = [ "foundry", diff --git a/python/foundry/__init__.py b/python/foundry/__init__.py index bb88e00e..226db5e9 100644 --- a/python/foundry/__init__.py +++ b/python/foundry/__init__.py @@ -1,5 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the Foundry project +# The ops wildcard must come FIRST: .graph's CUDAGraph (and the +# allocation_region helpers) intentionally override the raw pybind names. from .allocation_region import ( allocation_region, free_preallocated_region, diff --git a/python/foundry/graph.py b/python/foundry/graph.py index 427d1599..6a3b4c1c 100644 --- a/python/foundry/graph.py +++ b/python/foundry/graph.py @@ -135,7 +135,10 @@ def __init__( if self.__class__.default_capture_stream is None: self.__class__.default_capture_stream = torch.cuda.Stream() - self.pool = () if pool is None else (pool,) + # The C++ binding accepts a pool tuple or explicit None, but the + # argument itself is required. Keep one positional element so the + # default path passes None instead of omitting `pool` entirely. + self.pool = (pool,) self.capture_stream = ( stream if stream is not None else self.__class__.default_capture_stream ) diff --git a/python/foundry/integration/sglang/config.py b/python/foundry/integration/sglang/config.py index a0e1d41d..adca4902 100644 --- a/python/foundry/integration/sglang/config.py +++ b/python/foundry/integration/sglang/config.py @@ -6,11 +6,15 @@ import enum import importlib.util +import os from dataclasses import dataclass from pathlib import Path import tomllib +MULTI_SHAPE_ENV = "FOUNDRY_SGLANG_SYMM_MULTI_SHAPE" +MULTI_SHAPE_BATCHES_ENV = "FOUNDRY_SGLANG_SYMM_GRAPH_BATCHES" + class CUDAGraphExtensionMode(str, enum.Enum): NONE = "none" @@ -131,6 +135,32 @@ def get_nvshmem_host_path() -> str | None: return _config.nvshmem_host_path +def is_symmetric_multi_shape_enabled() -> bool: + value = os.environ.get(MULTI_SHAPE_ENV, "0") + if value not in {"0", "1"}: + raise RuntimeError(f"{MULTI_SHAPE_ENV} must be 0 or 1, got {value!r}") + return value == "1" + + +def symmetric_graph_batch_sizes() -> tuple[int, ...] | None: + raw = os.environ.get(MULTI_SHAPE_BATCHES_ENV) + if raw is None: + return None + try: + batches = tuple(int(value) for value in raw.split(",")) + except ValueError as error: + raise RuntimeError(f"{MULTI_SHAPE_BATCHES_ENV} must be comma-separated integers") from error + if not batches or any(batch <= 0 for batch in batches): + raise RuntimeError(f"{MULTI_SHAPE_BATCHES_ENV} requires positive values") + if len(batches) != len(set(batches)): + raise RuntimeError(f"{MULTI_SHAPE_BATCHES_ENV} contains a duplicate") + if batches != tuple(sorted(batches)): + raise RuntimeError(f"{MULTI_SHAPE_BATCHES_ENV} must be ascending") + if batches[0] != 1: + raise RuntimeError(f"{MULTI_SHAPE_BATCHES_ENV} must include batch 1") + return batches + + def compute_workspace_rank(server_args, tp_rank: int, pp_rank: int, dp_rank: int | None) -> int: if getattr(server_args, "enable_dp_attention", False): return pp_rank * server_args.tp_size + tp_rank diff --git a/python/foundry/integration/sglang/flashinfer_graph_abi.py b/python/foundry/integration/sglang/flashinfer_graph_abi.py new file mode 100644 index 00000000..5991a9c3 --- /dev/null +++ b/python/foundry/integration/sglang/flashinfer_graph_abi.py @@ -0,0 +1,505 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the Foundry project +"""Pinned FlashInfer CUDA graph ABI for SGLang shape replay.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +from foundry.integration.sglang.sglang_shape_adapter import PLAN_SCHEMA +from foundry.integration.sglang.state_manifest import OperandSlotRecord + +KERNEL_ABI_SCHEMA = "flashinfer-0.6.15.post1.fa2.qwen3-8b-tp2-bf16-paged376" +PAGED_KERNEL_NAME = ( + "_ZN10flashinfer34BatchPrefillWithPagedKVCacheKernelINS_12KernelTraits" + "ILNS_8MaskModeE0ELj16ELj1ELj2ELj8ELj8ELj1ELj4ELNS_15PosEncodingModeE0E" + "13__nv_bfloat16S4_S4_fiNS_16DefaultAttentionILb0ELb0ELb0ELb0EEEEE" + "11PagedParamsEEvT0_" +) +MERGE_KERNEL_NAME = ( + "_ZN10flashinfer41PersistentVariableLengthMergeStatesKernel" + "ILj8ELj16ELj8ELj4E13__nv_bfloat16S1_iEEvPT3_PfPT5_PT4_S4_jPjj" +) +PAGED_PARAM_LAYOUT = ((0, 0, 376),) +MERGE_PARAM_LAYOUT = ( + (0, 0, 8), + (1, 8, 8), + (2, 16, 8), + (3, 24, 8), + (4, 32, 8), + (5, 40, 4), + (6, 48, 8), + (7, 56, 4), +) +PREFILL_PLAN_INDEX = { + "padded_batch_size": 0, + "total_num_rows": 1, + "total_num_rows_offset": 2, + "cta_tile_q": 3, + "request_indices_offset": 4, + "qo_tile_indices_offset": 5, + "kv_tile_indices_offset": 6, + "merge_indptr_offset": 7, + "o_indptr_offset": 8, + "kv_chunk_size_ptr_offset": 9, + "v_offset": 10, + "s_offset": 11, + "block_valid_mask_offset": 12, + "enable_cuda_graph": 13, + "split_kv": 14, +} +_EXPECTED_PLAN_SCHEMA = "flashinfer.prefill.v15" +_NODE_PAIR_COUNT = 36 +_NULL_PAGED_POINTER_OFFSETS = (96, 152, 160, 176, 184, 192, 200, 208) +_STATIC_PAGED_POINTER_OFFSETS = ( + ("query", 0), + ("kv_cache_k", 56), + ("kv_cache_v", 64), +) +_FOUNDRY_BASE = 0x600000000000 +_FOUNDRY_END = _FOUNDRY_BASE + 256 * 1024**3 + + +@dataclass(frozen=True) +class PointerSpec: + name: str + parameter_index: int + value_byte_offset: int + ctype: str + owner_slot: str + plan_index: int | None + span_bytes: int + + +PAGED_POINTER_SPECS = ( + PointerSpec("paged_kv.indices", 0, 72, "int32_t*", "paged_kv_indices", None, 4), + PointerSpec("paged_kv.indptr", 0, 80, "int32_t*", "paged_kv_indptr", None, 4), + PointerSpec( + "paged_kv.last_page_len", + 0, + 88, + "int32_t*", + "paged_kv_last_page_len", + None, + 4, + ), + PointerSpec("q_indptr", 0, 104, "int32_t*", "qo_indptr", None, 4), + PointerSpec("tmp_v", 0, 112, "nv_bfloat16*", "float_workspace", 10, 131072), + PointerSpec("tmp_s", 0, 120, "float*", "float_workspace", 11, 1024), + PointerSpec("request_indices", 0, 296, "int32_t*", "int_workspace", 4, 4), + PointerSpec("qo_tile_indices", 0, 304, "int32_t*", "int_workspace", 5, 4), + PointerSpec("kv_tile_indices", 0, 312, "int32_t*", "int_workspace", 6, 4), + PointerSpec("merge_indptr", 0, 320, "int32_t*", "int_workspace", 7, 4), + PointerSpec("o_indptr", 0, 328, "int32_t*", "int_workspace", 8, 4), + PointerSpec("block_valid_mask", 0, 336, "bool*", "int_workspace", 12, 1), + PointerSpec("kv_chunk_size", 0, 344, "int32_t*", "int_workspace", 9, 4), + PointerSpec("total_num_rows", 0, 360, "uint32_t*", "int_workspace", 2, 4), +) +MERGE_POINTER_SPECS = ( + PointerSpec("tmp_v", 0, 0, "nv_bfloat16*", "float_workspace", 10, 131072), + PointerSpec("tmp_s", 1, 0, "float*", "float_workspace", 11, 1024), + PointerSpec("merge_indptr", 2, 0, "int32_t*", "int_workspace", 7, 4), + PointerSpec("total_num_rows", 6, 0, "uint32_t*", "int_workspace", 2, 4), +) + + +def _validate_plan(state: Any) -> tuple[int, ...]: + if PLAN_SCHEMA != _EXPECTED_PLAN_SCHEMA: + raise RuntimeError(f"Unsupported FlashInfer plan schema: {PLAN_SCHEMA!r}") + plan = tuple(int(value) for value in state.plan_fingerprint) + if ( + len(plan) != len(PREFILL_PLAN_INDEX) + or plan[PREFILL_PLAN_INDEX["enable_cuda_graph"]] != 1 + or plan[PREFILL_PLAN_INDEX["split_kv"]] != 1 + ): + raise RuntimeError(f"Unsupported FlashInfer prefill plan: {plan}") + return plan + + +def _read_u64(raw: bytes | bytearray, offset: int) -> int: + return int.from_bytes(raw[offset : offset + 8], "little") + + +def _write_u64(raw: bytearray, offset: int, value: int) -> None: + raw[offset : offset + 8] = value.to_bytes(8, "little") + + +def _owner_map(state: Any) -> dict[str, Any]: + return {slot.name: slot for slot in (*state.owners, *state.shared_owners)} + + +def _read_record_value(raw: bytes | bytearray, offset: int, label: str) -> int: + if offset < 0 or offset + 8 > len(raw): + raise RuntimeError(f"FlashInfer {label} offset is out of bounds") + return _read_u64(raw, offset) + + +def _build_record( + *, + node_id: int, + param: dict, + name: str, + owner_slot: str, + value_byte_offset: int, + ctype: str, + owner_relative_offset: int, + span_bytes: int, + saved_value: int, +) -> OperandSlotRecord: + return OperandSlotRecord( + name=name, + kernel_abi=KERNEL_ABI_SCHEMA, + owner_slot=owner_slot, + node_id=node_id, + source="kernelParams", + parameter_index=int(param["index"]), + cuda_parameter_offset=int(param["offset"]), + value_byte_offset=value_byte_offset, + ctype=ctype, + owner_relative_offset=owner_relative_offset, + span_bytes=span_bytes, + saved_value=saved_value, + ) + + +def _validate_kernel_node( + node: dict, + *, + function_name: str, + layout: tuple[tuple[int, int, int], ...], +) -> list[dict]: + if node["type"] != "KernelNode": + raise RuntimeError("FlashInfer node is not a KernelNode") + params = node["params"] + if params["function_name"] != function_name: + raise RuntimeError( + f"FlashInfer kernel symbol changed: {params['function_name']} != {function_name}" + ) + kernel_params = params["kernelParams"] + _validate_layout(kernel_params, layout) + if params.get("extra_argBuffer_hex"): + raise RuntimeError("FlashInfer kernels use an unsupported arg buffer") + return kernel_params + + +def _collect_flashinfer_nodes(graph_data: dict) -> tuple[list[dict], list[dict]]: + paged_nodes = [ + node + for node in graph_data["nodes"] + if node["type"] == "KernelNode" and node["params"]["function_name"] == PAGED_KERNEL_NAME + ] + merge_nodes = [ + node + for node in graph_data["nodes"] + if node["type"] == "KernelNode" and node["params"]["function_name"] == MERGE_KERNEL_NAME + ] + if len(paged_nodes) != _NODE_PAIR_COUNT or len(merge_nodes) != _NODE_PAIR_COUNT: + raise RuntimeError( + "Expected 36 FlashInfer paged and merge nodes, got " + f"{len(paged_nodes)} and {len(merge_nodes)}" + ) + return paged_nodes, merge_nodes + + +def _validate_layout(params: list[dict], expected: tuple[tuple[int, int, int], ...]) -> None: + actual = tuple( + (int(param["index"]), int(param["offset"]), int(param["size"])) for param in params + ) + if actual != expected: + raise RuntimeError(f"FlashInfer kernel ABI changed: {actual} != {expected}") + + +def _span_bytes(spec: PointerSpec, state: Any) -> int: + padded_batch_size = int(state.plan_fingerprint[PREFILL_PLAN_INDEX["padded_batch_size"]]) + batch_size = int(state.batch_size) + if spec.name in {"tmp_v", "tmp_s"}: + return spec.span_bytes * padded_batch_size + if spec.name in {"request_indices", "qo_tile_indices", "kv_tile_indices"}: + return 4 * padded_batch_size + if spec.name in {"merge_indptr", "o_indptr"}: + return 4 * (batch_size + 1) + if spec.name == "block_valid_mask": + return padded_batch_size + return spec.span_bytes + + +def _record_specs( + node: dict, + specs: tuple[PointerSpec, ...], + state: Any, + *, + validate_saved_pointers: bool, +) -> tuple[OperandSlotRecord, ...]: + owners = _owner_map(state) + params = node["params"]["kernelParams"] + records = [] + for spec in specs: + param = params[spec.parameter_index] + raw = bytes.fromhex(param["value_hex"]) + owner = owners[spec.owner_slot] + owner_relative_offset = ( + 0 if spec.plan_index is None else int(state.plan_fingerprint[spec.plan_index]) + ) + span_bytes = _span_bytes(spec, state) + capacity = int(owner.numel) * int(owner.element_size) + if owner_relative_offset < 0 or owner_relative_offset + span_bytes > capacity: + raise RuntimeError(f"FlashInfer owner capacity is too small for {spec.name}") + actual = _read_record_value(raw, spec.value_byte_offset, spec.name) + expected = int(owner.data_ptr) + owner_relative_offset + if validate_saved_pointers and actual != expected: + raise RuntimeError( + f"FlashInfer pointer changed for {spec.name}: 0x{actual:x} != 0x{expected:x}" + ) + records.append( + _build_record( + node_id=int(node["id"]), + param=param, + name=spec.name, + owner_slot=spec.owner_slot, + value_byte_offset=spec.value_byte_offset, + ctype=spec.ctype, + owner_relative_offset=owner_relative_offset, + span_bytes=span_bytes, + saved_value=actual, + ) + ) + return tuple(records) + + +def _build_flashinfer_records( + graph_data: dict, + state: Any, + *, + validate_saved_pointers: bool, +) -> tuple[OperandSlotRecord, ...]: + _validate_plan(state) + paged_nodes, merge_nodes = _collect_flashinfer_nodes(graph_data) + + records = [] + for paged_node, merge_node in zip(paged_nodes, merge_nodes, strict=True): + paged_params = _validate_kernel_node( + paged_node, + function_name=PAGED_KERNEL_NAME, + layout=PAGED_PARAM_LAYOUT, + ) + merge_params = _validate_kernel_node( + merge_node, + function_name=MERGE_KERNEL_NAME, + layout=MERGE_PARAM_LAYOUT, + ) + + paged_raw = bytes.fromhex(paged_params[0]["value_hex"]) + for offset in _NULL_PAGED_POINTER_OFFSETS: + if _read_record_value(paged_raw, offset, f"PagedParams+{offset}") != 0: + raise RuntimeError(f"FlashInfer pointer at PagedParams+{offset} must be null") + for name, offset in _STATIC_PAGED_POINTER_OFFSETS: + value = _read_record_value(paged_raw, offset, name) + if not _FOUNDRY_BASE <= value < _FOUNDRY_END: + raise RuntimeError(f"FlashInfer {name} is outside Foundry VMM") + records.append( + _build_record( + node_id=int(paged_node["id"]), + param=paged_params[0], + name=name, + owner_slot="foundry_vmm", + value_byte_offset=offset, + ctype="void*", + owner_relative_offset=0, + span_bytes=1, + saved_value=value, + ) + ) + records.append( + _build_record( + node_id=int(paged_node["id"]), + param=paged_params[0], + name="unused_alibi", + owner_slot="normalized_null", + value_byte_offset=168, + ctype="float*", + owner_relative_offset=0, + span_bytes=0, + saved_value=_read_record_value(paged_raw, 168, "unused_alibi"), + ) + ) + records.extend( + _record_specs( + paged_node, + PAGED_POINTER_SPECS, + state, + validate_saved_pointers=validate_saved_pointers, + ) + ) + records.extend( + _record_specs( + merge_node, + MERGE_POINTER_SPECS, + state, + validate_saved_pointers=validate_saved_pointers, + ) + ) + + merge_values = [ + int.from_bytes(bytes.fromhex(param["value_hex"]), "little") for param in merge_params + ] + if _read_record_value(paged_raw, 112, "tmp_v") != merge_values[0]: + raise RuntimeError("FlashInfer tmp_v cross-node pointer mismatch") + if _read_record_value(paged_raw, 120, "tmp_s") != merge_values[1]: + raise RuntimeError("FlashInfer tmp_s cross-node pointer mismatch") + if _read_record_value(paged_raw, 320, "merge_indptr") != merge_values[2]: + raise RuntimeError("FlashInfer merge_indptr pointer mismatch") + if _read_record_value(paged_raw, 360, "total_num_rows") != merge_values[6]: + raise RuntimeError("FlashInfer total_num_rows pointer mismatch") + if int.from_bytes(paged_raw[352:356], "little") != int(state.batch_size): + raise RuntimeError("FlashInfer PagedParams batch size changed") + if merge_values[5] != int(state.batch_size): + raise RuntimeError("FlashInfer merge max_seq_len changed") + if int.from_bytes(paged_raw[256:260], "little") != 16: + raise RuntimeError("FlashInfer local query-head count changed") + if merge_values[7] != 16 or merge_values[4] != 0: + raise RuntimeError("FlashInfer merge scalar ABI changed") + if not _FOUNDRY_BASE <= merge_values[3] < _FOUNDRY_END: + raise RuntimeError("FlashInfer merge output is outside Foundry VMM") + records.append( + _build_record( + node_id=int(merge_node["id"]), + param=merge_params[3], + name="merge_output", + owner_slot="foundry_vmm", + value_byte_offset=0, + ctype="nv_bfloat16*", + owner_relative_offset=0, + span_bytes=1, + saved_value=merge_values[3], + ) + ) + if paged_raw[372] != 1: + raise RuntimeError("FlashInfer partition_kv flag changed") + return tuple(records) + + +def extract_flashinfer_operand_records( + graph_data: dict, + state: Any, +) -> tuple[OperandSlotRecord, ...]: + return _build_flashinfer_records( + graph_data, + state, + validate_saved_pointers=True, + ) + + +def _record_location(record: OperandSlotRecord) -> tuple[int, str, int, int]: + return ( + record.node_id, + record.source, + record.parameter_index, + record.value_byte_offset, + ) + + +def _sorted_record_locations( + records: tuple[OperandSlotRecord, ...], +) -> tuple[tuple[int, str, int, int], ...]: + return tuple(sorted(_record_location(record) for record in records)) + + +def relocate_flashinfer_operands( + graph_data: dict, + state: Any, + records: tuple[OperandSlotRecord, ...], +) -> None: + expected_records = _build_flashinfer_records( + graph_data, + state, + validate_saved_pointers=False, + ) + expected_locations = _sorted_record_locations(expected_records) + nodes = {int(node["id"]): node for node in graph_data["nodes"]} + owners = _owner_map(state) + for record in records: + if record.kernel_abi != KERNEL_ABI_SCHEMA: + raise RuntimeError(f"Unsupported FlashInfer kernel ABI: {record.kernel_abi}") + if record.source != "kernelParams": + raise RuntimeError(f"FlashInfer record source changed for {record.name}") + node = nodes.get(record.node_id) + if node is None: + raise RuntimeError(f"FlashInfer record node is missing for {record.name}") + if node["type"] != "KernelNode": + raise RuntimeError(f"FlashInfer record node is not a KernelNode for {record.name}") + params = node["params"]["kernelParams"] + if record.parameter_index < 0 or record.parameter_index >= len(params): + raise RuntimeError(f"FlashInfer record parameter index changed for {record.name}") + raw = bytearray.fromhex(params[record.parameter_index]["value_hex"]) + if record.value_byte_offset < 0 or record.value_byte_offset + 8 > len(raw): + raise RuntimeError(f"FlashInfer record value offset is out of bounds for {record.name}") + if record.owner_relative_offset < 0: + raise RuntimeError( + f"FlashInfer record has negative owner_relative_offset for {record.name}" + ) + if record.span_bytes < 0: + raise RuntimeError(f"FlashInfer record has negative span for {record.name}") + if record.owner_slot not in {"normalized_null", "foundry_vmm"}: + owner = owners.get(record.owner_slot) + if owner is None: + raise RuntimeError(f"FlashInfer record owner slot changed for {record.name}") + capacity = int(owner.numel) * int(owner.element_size) + if ( + record.owner_relative_offset > capacity + or record.owner_relative_offset + record.span_bytes > capacity + ): + raise RuntimeError(f"FlashInfer owner range is invalid for {record.name}") + + actual_locations = _sorted_record_locations(records) + if actual_locations != expected_locations: + raise RuntimeError("FlashInfer operand record set changed") + + expected_by_location = {_record_location(record): record for record in expected_records} + pending_writes: dict[tuple[int, int], bytearray] = {} + for record in records: + expected = expected_by_location[_record_location(record)] + node = nodes[record.node_id] + param = node["params"]["kernelParams"][record.parameter_index] + key = (record.node_id, record.parameter_index) + raw = pending_writes.get(key) + if raw is None: + raw = bytearray.fromhex(param["value_hex"]) + pending_writes[key] = raw + if record.name != expected.name: + raise RuntimeError(f"FlashInfer record name changed for {record.name}") + if record.cuda_parameter_offset != int(param["offset"]): + raise RuntimeError(f"FlashInfer record cuda parameter offset changed for {record.name}") + if record.cuda_parameter_offset != expected.cuda_parameter_offset: + raise RuntimeError(f"FlashInfer record cuda parameter offset changed for {record.name}") + if record.ctype != expected.ctype: + raise RuntimeError(f"FlashInfer record ctype changed for {record.name}") + if record.owner_slot != expected.owner_slot: + raise RuntimeError(f"FlashInfer record owner slot changed for {record.name}") + if record.owner_relative_offset != expected.owner_relative_offset: + raise RuntimeError(f"FlashInfer record owner relative offset changed for {record.name}") + if record.span_bytes != expected.span_bytes: + raise RuntimeError(f"FlashInfer record span changed for {record.name}") + + actual = _read_record_value(raw, record.value_byte_offset, record.name) + if actual != record.saved_value: + raise RuntimeError(f"Saved FlashInfer operand changed: {record.name}") + if record.owner_slot == "normalized_null": + replacement = 0 + elif record.owner_slot == "foundry_vmm": + replacement = record.saved_value + else: + owner = owners.get(record.owner_slot) + if owner is None: + raise RuntimeError(f"FlashInfer record owner slot changed for {record.name}") + capacity = int(owner.numel) * int(owner.element_size) + if ( + record.owner_relative_offset > capacity + or record.owner_relative_offset + record.span_bytes > capacity + ): + raise RuntimeError(f"FlashInfer owner range is invalid for {record.name}") + replacement = int(owner.data_ptr) + record.owner_relative_offset + _write_u64(raw, record.value_byte_offset, replacement) + for (node_id, parameter_index), raw in pending_writes.items(): + param = nodes[node_id]["params"]["kernelParams"][parameter_index] + param["value_hex"] = raw.hex() diff --git a/python/foundry/integration/sglang/graph_ops.py b/python/foundry/integration/sglang/graph_ops.py index 18ab1f0e..d9aa9998 100644 --- a/python/foundry/integration/sglang/graph_ops.py +++ b/python/foundry/integration/sglang/graph_ops.py @@ -32,6 +32,9 @@ def _batch_size_from_key(key: Any) -> int: if isinstance(key, int): return key + shape_size = getattr(key, "size", None) + if isinstance(shape_size, int): + return shape_size key_str = str(key) for part in reversed(key_str.split("_")): if part.isdigit(): @@ -248,6 +251,34 @@ def bootstrap_deepep_buffer(cuda_graph_runner) -> bool: return False +def build_capture_fb_view(cuda_graph_runner, bs: int): + """Build a ForwardBatch-like view for backend capture-side metadata init. + + Mirrors the padded per-bs inputs sglang's own capture loop feeds the + attention backend (``cuda_graph_runner.py::capture_one_batch_size``), as a + lightweight ``SimpleNamespace`` since the pre-pass runs before a real + ``ForwardBatch`` exists. Used to drive the post-#26735 3-method init ABC + (``init_forward_metadata_out_graph``). + """ + from types import SimpleNamespace + + buffers = cuda_graph_runner.buffers + num_tokens = bs * cuda_graph_runner.num_tokens_per_bs + encoder_lens = buffers.encoder_lens[:bs] if cuda_graph_runner.is_encoder_decoder else None + seq_lens = buffers.seq_lens[:bs] + return SimpleNamespace( + batch_size=bs, + forward_mode=cuda_graph_runner.capture_forward_mode, + req_pool_indices=buffers.req_pool_indices[:bs], + seq_lens=seq_lens, + seq_lens_cpu=buffers.seq_lens_cpu[:bs], + seq_lens_sum=int(seq_lens.sum().item()), + encoder_lens=encoder_lens, + spec_info=cuda_graph_runner.get_spec_info(num_tokens), + positions=buffers.positions[:num_tokens], + ) + + def initialize_attention_metadata_for_bs(cuda_graph_runner, bs: int) -> None: """Populate ``decode_cuda_graph_metadata[bs]`` for runtime replay. @@ -257,11 +288,21 @@ def initialize_attention_metadata_for_bs(cuda_graph_runner, bs: int) -> None: re-run the same call before runtime replay so the wrappers exist at deterministic VMM addresses. """ + attn_backend = cuda_graph_runner.attn_backend + # sglang #26735 replaced the single ``init_forward_metadata_capture_cuda_graph`` + # with a 3-method ABC; the capture-side allocation now lives in + # ``init_forward_metadata_out_graph(fb, in_capture=True)``. Prefer the new + # API, fall back to the legacy method for pre-rename sglang. + if hasattr(attn_backend, "init_forward_metadata_out_graph"): + fb_view = build_capture_fb_view(cuda_graph_runner, bs) + attn_backend.init_forward_metadata_out_graph(fb_view, in_capture=True) + return + buffers = cuda_graph_runner.buffers num_tokens = bs * cuda_graph_runner.num_tokens_per_bs encoder_lens = buffers.encoder_lens[:bs] if cuda_graph_runner.is_encoder_decoder else None spec_info = cuda_graph_runner.get_spec_info(num_tokens) - cuda_graph_runner.attn_backend.init_forward_metadata_capture_cuda_graph( + attn_backend.init_forward_metadata_capture_cuda_graph( bs, num_tokens, buffers.req_pool_indices[:bs], diff --git a/python/foundry/integration/sglang/hooks.py b/python/foundry/integration/sglang/hooks.py index 22cd7cfa..5a8dc61d 100644 --- a/python/foundry/integration/sglang/hooks.py +++ b/python/foundry/integration/sglang/hooks.py @@ -5,10 +5,14 @@ from __future__ import annotations import functools +import importlib +import importlib.util +import inspect import logging import os import time from dataclasses import asdict +from types import SimpleNamespace from foundry.integration.sglang import runtime as rt from foundry.integration.sglang.config import ( @@ -20,6 +24,20 @@ logger = logging.getLogger(__name__) _INSTALLED = False +_DEEPEP_PATCHED = False + + +def install_hooks_from_path(config_path: str) -> None: + """Install hooks without requiring a Foundry-specific ServerArgs field. + + SGLang main discovers Foundry as a plugin in every process and carries the + config path through the environment. Legacy Foundry forks continue to call + :func:`install_hooks` with their declared ``ServerArgs`` field. + """ + + install_hooks( + SimpleNamespace(foundry_graph_extension_config_path=config_path), + ) def _ep_lazy_init_needed() -> bool: @@ -33,6 +51,39 @@ def _ep_lazy_init_needed() -> bool: return False +def _flashinfer_backend(attn_backend): + """Return the FlashInfer backend that owns the per-bs capture wrappers, or + None if this attention backend doesn't use FlashInfer's per-bs prepass. + + FlashInfer allocates a fresh per-bs metadata wrapper (each with its own + ``_int_workspace_buffer``) on every capture init, so foundry pre-allocates + them all up front and installs a reuse shim to keep the VMM cursor + deterministic vs LOAD. Hybrid linear-attn / mamba models (e.g. Qwen3.5) + wrap FlashInfer inside a ``HybridLinearAttnBackend`` whose full-attention + layers use FlashInfer while its linear-attn (mamba) layers keep all their + graph buffers in ``init_cuda_graph_state`` (pre-capture, so already + SAVE/LOAD-symmetric). Detect the FlashInfer backend by its per-bs + ``indices_updater_decode``, seeing through the hybrid wrapper's + ``full_attn_backend`` so the wrappers get the deterministic prepass too.""" + if hasattr(attn_backend, "indices_updater_decode"): + return attn_backend + if hasattr(attn_backend, "full_attn_backend"): + inner = attn_backend.full_attn_backend + if hasattr(inner, "indices_updater_decode"): + return inner + return None + + +def _dense_save_warmup_enabled() -> bool: + """Whether the dense/DP SAVE path runs a cursor-neutral pre-capture warmup. + + Defaults on: models that defer torch.compile / per-shape GEMM JIT to their + first forward (e.g. Qwen3.5 hybrid MoE) would otherwise trigger that lazy + init inside CUDA-graph capture and abort. Set ``FOUNDRY_SGLANG_SAVE_WARMUP=0`` + to restore the old suppress-only behavior.""" + return os.environ.get("FOUNDRY_SGLANG_SAVE_WARMUP", "1") not in ("0", "false", "False") + + def _resolve_dp_rank(model_runner) -> int | None: dp_rank = getattr(model_runner, "dp_rank", None) if dp_rank is not None: @@ -64,6 +115,16 @@ def _resolve_dp_rank(model_runner) -> int | None: return None +def _uses_main_runner_api() -> bool: + try: + return ( + importlib.util.find_spec("sglang.srt.model_executor.runner.decode_cuda_graph_runner") + is not None + ) + except (ImportError, ModuleNotFoundError, ValueError): + return False + + def install_hooks(server_args) -> None: global _INSTALLED cfg_path = getattr(server_args, "foundry_graph_extension_config_path", None) @@ -86,17 +147,71 @@ def install_hooks(server_args) -> None: get_workspace_root(), ) - _patch_init_torch_distributed() - _patch_init_memory_pool() - _patch_load_model() - _patch_kernel_warmup() - _patch_cuda_graph_capture() - _patch_spawn_sites() + if _uses_main_runner_api(): + # Imported only after API detection so legacy SGLang installations do + # not resolve modules that do not exist in their runner layout. + hooks_main = importlib.import_module("foundry.integration.sglang.hooks_main") + hooks_main.install_hooks_main() + else: + _patch_init_torch_distributed() + _patch_init_memory_pool() + _patch_load_model() + _patch_kernel_warmup() + _patch_cuda_graph_capture() + _patch_spawn_sites() + _patch_deepep() _INSTALLED = True logger.info("[Foundry] SGLang hooks installed") +def _patch_deepep() -> None: + """Select Foundry's DeepEP memory transport for graph SAVE/LOAD. + + DeepEP 29d31c0 added ``use_fabric`` while preserving the older Buffer + positional call shape. Keep the default Foundry mode on the fabric + RDMA-only path; the IPC experiment retains DeepEP's NVL allocation and + uses the VMM CUDA-IPC bridge in ``libcuda_hook.so``. + """ + global _DEEPEP_PATCHED + if _DEEPEP_PATCHED: + return + + try: + from deep_ep import Buffer + except Exception: + return + + original_init = Buffer.__init__ + supports_use_fabric = "use_fabric" in inspect.signature(original_init).parameters + + @functools.wraps(original_init) + def patched(self, group, num_nvl_bytes=0, num_rdma_bytes=0, *args, **kwargs): + if get_graph_extension_mode() != CUDAGraphExtensionMode.NONE: + if not supports_use_fabric: + raise RuntimeError( + "[Foundry] SGLang DeepEP requires Buffer(use_fabric=...). " + "Install DeepEP 29d31c0 or newer." + ) + if os.environ.get("FOUNDRY_DEEPEP_NVL_IPC", "0") == "1": + kwargs["use_fabric"] = False + else: + num_nvl_bytes = 0 + kwargs["use_fabric"] = True + return original_init( + self, + group, + num_nvl_bytes, + num_rdma_bytes, + *args, + **kwargs, + ) + + Buffer.__init__ = patched + _DEEPEP_PATCHED = True + logger.info("[Foundry] SGLang DeepEP transport patch installed") + + def _patch_init_torch_distributed() -> None: from sglang.srt.model_executor import model_runner as mr @@ -321,6 +436,53 @@ def _run_warmup_pass(self): time.perf_counter() - t0, ) + def _run_dense_warmup(self): + """In-region pre-capture warmup for the dense / DP path. + + Some models (e.g. Qwen3.5 hybrid MoE) defer torch.compile / inductor, + per-shape GEMM JIT, or Triton kernel compilation to their first forward. + On the foundry SAVE path that first forward is *inside* CUDA-graph capture + (sglang's own warmup forwards are suppressed for VMM determinism), where a + lazy CPU<->CUDA copy aborts capture ("Graph contains no nodes"). More + critically, hybrid models (MoE + mamba) lazily initialize persistent + module-level GPU buffers (e.g. Triton codegen workspace, torch.compile + inductor intermediates) on the first real forward โ€” if those buffers land + outside the VMM region, the CUDA graph captures their out-of-region + addresses, which are stale on LOAD (different process โ†’ different addresses). + + Run one real forward per capture_bs with the allocation region ACTIVE so all + lazy-init persistent buffers land IN-REGION at deterministic VMM addresses. + The cursor advances during warmup; LOAD runs the same warmup (same model, + same forward, same capture_bs) to produce the same cursor trajectory. After + warmup, clear FlashInfer per-bs wrappers and empty the caching allocator + cache so the subsequent FlashInfer pre-pass re-allocates wrappers at the + post-warmup cursor โ€” identical on both SAVE and LOAD. + """ + import torch + + fi_backend = _flashinfer_backend(self.attn_backend) + has_decode = fi_backend is not None and hasattr(fi_backend, "decode_cuda_graph_metadata") + has_prefill = fi_backend is not None and hasattr(fi_backend, "prefill_cuda_graph_metadata") + if has_decode: + fi_backend.decode_cuda_graph_metadata = {} + if has_prefill: + fi_backend.prefill_cuda_graph_metadata = {} + rt.log_alloc_offset("before_warmup") + try: + _run_warmup_pass(self) + finally: + self.attn_backend.forward_metadata = None + if fi_backend is not None: + fi_backend.forward_metadata = None + # Drop warmup's FlashInfer wrappers so the real pre-pass + # re-allocates them at the post-warmup cursor. + if has_decode: + fi_backend.decode_cuda_graph_metadata = {} + if has_prefill: + fi_backend.prefill_cuda_graph_metadata = {} + torch.cuda.empty_cache() + rt.log_alloc_offset("after_warmup") + @functools.wraps(orig_capture) def patched(self, *args, **kwargs): mode = get_graph_extension_mode() @@ -334,11 +496,11 @@ def patched(self, *args, **kwargs): # 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. + # at its recorded absolute offset. EP warmup is SAVE-only (its + # orig_capture re-enters graph_capture() and breaks the threaded + # finish_graph_loads with "invalid device context"). Dense warmup + # runs on BOTH SAVE and LOAD (in the LOAD branch below) so lazy-init + # persistent buffers land at deterministic in-region addresses. # 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(): @@ -370,6 +532,14 @@ def patched(self, *args, **kwargs): 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()) + # Run the same in-region warmup as SAVE so lazy-init model + # buffers (torch.compile codegen, Triton kernel workspace, + # MoE/mamba persistent state) land at the same VMM addresses. + # Without this, SAVE's graph captures out-of-region pointers to + # buffers that were lazily allocated during SAVE's warmup โ€” stale + # on LOAD (different process). + if not _ep_lazy_init_needed() and _dense_save_warmup_enabled(): + _run_dense_warmup(self) rt.log_alloc_offset("before_preallocate") rt.preallocate_for_load_mode() rt.log_alloc_offset("after_preallocate") @@ -382,7 +552,13 @@ def patched(self, *args, **kwargs): # ``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"): + # ``_flashinfer_backend`` sees through the hybrid linear-attn wrapper + # (Qwen3.5) so its FlashInfer full-attn wrappers get pre-allocated + # here โ€” BEFORE load_all_graphs โ€” landing at SAVE's offsets. Its + # mamba layers keep all buffers in init_cuda_graph_state, so no + # divergent alloc there. + fi_backend = _flashinfer_backend(self.attn_backend) + if fi_backend is not None: initialize_all_attention_metadata(self) rt.log_alloc_offset("after_pre_init") # Single ``start_graph_builds(all_paths)`` call so templates @@ -400,7 +576,7 @@ def patched(self, *args, **kwargs): # 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"): + if fi_backend is None: 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, @@ -412,17 +588,86 @@ def patched(self, *args, **kwargs): return None if mode == CUDAGraphExtensionMode.SAVE: + # Force model-level lazy init (torch.compile/inductor, per-shape GEMM + # JIT) outside capture for the dense/DP path. The warmup runs + # in-region so lazy-init persistent model buffers land at + # deterministic VMM addresses; LOAD runs the same warmup to + # reproduce the cursor trajectory. EP already warmed up above; + # skip there. + if not _ep_lazy_init_needed() and _dense_save_warmup_enabled(): + _run_dense_warmup(self) # 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. + # this and the plain capture is already SAVE/LOAD-deterministic. + # ``_flashinfer_backend`` returns the FlashInfer backend, seeing + # through the hybrid linear-attn wrapper (Qwen3.5) whose FlashInfer + # full-attn layers own the per-bs wrappers while its mamba layers + # keep all buffers in init_cuda_graph_state. The shim is installed on + # that FlashInfer backend directly; the hybrid wrapper's own + # init_forward_metadata_out_graph then dispatches into it during + # capture, so the mamba backend still runs its normal (alloc-free) + # metadata copy. attn_backend = self.attn_backend - use_fi_prepass = hasattr(attn_backend, "indices_updater_decode") - real_init = attn_backend.init_forward_metadata_capture_cuda_graph + fi_backend = _flashinfer_backend(attn_backend) + use_fi_prepass = fi_backend is not None + # sglang #26735 split the single capture-init entry point into the + # 3-method ABC; the per-iter capture path is now + # ``init_forward_metadata_out_graph(fb, in_capture=True)``. Detect + # which API this sglang exposes and shim the matching method. + use_new_api = hasattr(attn_backend, "init_forward_metadata_out_graph") + real_out_graph = ( + fi_backend.init_forward_metadata_out_graph + if use_new_api and use_fi_prepass + else None + ) + real_init = ( + fi_backend.init_forward_metadata_capture_cuda_graph + if not use_new_api and use_fi_prepass + else None + ) + + def reuse_out_graph(forward_batch, in_capture=False): + # Only the capture path (in_capture=True) must avoid + # re-allocating the per-bs wrappers the pre-pass already built; + # replay / eager (in_capture=False) pass straight through. + if not in_capture: + return real_out_graph(forward_batch, in_capture=in_capture) + + from sglang.srt.layers.attention.flashinfer_backend import ( + DecodeMetadata, + PrefillMetadata, + ) + + bs = forward_batch.batch_size + forward_mode = forward_batch.forward_mode + if forward_mode.is_decode_or_idle(): + wrappers = fi_backend.decode_cuda_graph_metadata.get(bs) + if wrappers is None: + return real_out_graph(forward_batch, in_capture=True) + fi_backend.forward_metadata = DecodeMetadata(wrappers) + elif ( + forward_mode.is_target_verify() + or forward_mode.is_draft_extend() + or forward_mode.is_dllm_extend() + ): + wrappers = fi_backend.prefill_cuda_graph_metadata.get(bs) + if wrappers is None: + return real_out_graph(forward_batch, in_capture=True) + fi_backend.forward_metadata = PrefillMetadata( + wrappers, forward_mode.is_dllm_extend(), False + ) + else: + return real_out_graph(forward_batch, in_capture=True) + # Re-run the planner against the pre-allocated wrappers with + # in_capture=False: no ``_prepare_cuda_graph_metadata`` (so no + # second torch.empty for ``_int_workspace_buffer``) and no + # begin_forward re-install (the pre-pass did that). We set + # forward_metadata above since that path skips it. + real_out_graph(forward_batch, in_capture=False) def reuse_pre_pass_init( bs, @@ -450,7 +695,7 @@ def reuse_pre_pass_init( ) if forward_mode.is_decode_or_idle(): - wrappers = attn_backend.decode_cuda_graph_metadata.get(bs) + wrappers = fi_backend.decode_cuda_graph_metadata.get(bs) if wrappers is None: return real_init( bs, @@ -462,7 +707,7 @@ def reuse_pre_pass_init( spec_info, ) seq_lens_sum = seq_lens.sum().item() - attn_backend.indices_updater_decode.update( + fi_backend.indices_updater_decode.update( req_pool_indices, seq_lens, seq_lens.cpu(), @@ -471,16 +716,16 @@ def reuse_pre_pass_init( encoder_lens=encoder_lens, spec_info=spec_info, fixed_split_size=None, - disable_split_kv=attn_backend.disable_cuda_graph_kv_split, + disable_split_kv=fi_backend.disable_cuda_graph_kv_split, ) - attn_backend.forward_metadata = DecodeMetadata(wrappers) + fi_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) + wrappers = fi_backend.prefill_cuda_graph_metadata.get(bs) if wrappers is None: return real_init( bs, @@ -494,12 +739,12 @@ def reuse_pre_pass_init( seq_lens_sum = seq_lens.sum().item() use_ragged = forward_mode.is_dllm_extend() prefix_lens = ( - seq_lens - attn_backend.dllm_config.block_size + seq_lens - fi_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( + fi_backend.indices_updater_prefill.update( req_pool_indices, seq_lens, seq_lens.cpu(), @@ -510,7 +755,7 @@ def reuse_pre_pass_init( encoder_lens=encoder_lens, spec_info=spec_info_arg, ) - attn_backend.forward_metadata = PrefillMetadata(wrappers, use_ragged, False) + fi_backend.forward_metadata = PrefillMetadata(wrappers, use_ragged, False) return # Unknown mode โ€” fall back to real init. return real_init( @@ -534,13 +779,23 @@ def reuse_pre_pass_init( 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 + # entry doesn't keep the wrapper alive at refcount 1. Installed on + # the FlashInfer backend (== attn_backend, or its full_attn_backend + # for the hybrid wrapper); the hybrid's own out_graph dispatches + # into the shimmed FlashInfer method during capture. + fi_backend.forward_metadata = None + if use_new_api: + fi_backend.init_forward_metadata_out_graph = reuse_out_graph + else: + fi_backend.init_forward_metadata_capture_cuda_graph = reuse_pre_pass_init try: result = orig_capture(self, *args, **kwargs) finally: - attn_backend.init_forward_metadata_capture_cuda_graph = real_init + if use_fi_prepass: + if use_new_api: + fi_backend.init_forward_metadata_out_graph = real_out_graph + else: + fi_backend.init_forward_metadata_capture_cuda_graph = real_init from foundry.integration.sglang.graph_ops import ( pack_fatbins, diff --git a/python/foundry/integration/sglang/hooks_main.py b/python/foundry/integration/sglang/hooks_main.py new file mode 100644 index 00000000..4fdee403 --- /dev/null +++ b/python/foundry/integration/sglang/hooks_main.py @@ -0,0 +1,232 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the Foundry project +"""Monkey-patches for the current SGLang runner architecture.""" + +from __future__ import annotations + +import functools +import logging +from dataclasses import asdict + +import torch +from sglang.srt.distributed.device_communicators import ( + torch_symm_mem as torch_symm_mem_module, +) +from sglang.srt.distributed.device_communicators import triton_symm_mem_ag +from sglang.srt.entrypoints import engine as engine_mod +from sglang.srt.managers import data_parallel_controller as dpc +from sglang.srt.mem_cache.kv_cache_configurator import KVCacheConfigurator +from sglang.srt.model_executor import model_runner as mr +from sglang.srt.model_executor.pool_configurator import MemoryPoolConfig +from sglang.srt.model_executor.runner import decode_cuda_graph_runner as decode_runner +from sglang.srt.model_executor.runner_backend import utils as backend_utils + +from foundry import ops as cge +from foundry.integration.sglang import runtime as rt +from foundry.integration.sglang.config import ( + MULTI_SHAPE_BATCHES_ENV, + CUDAGraphExtensionMode, + get_graph_extension_mode, + is_symmetric_multi_shape_enabled, + symmetric_graph_batch_sizes, +) +from foundry.integration.sglang.main_backend import FoundryMainCudaGraphBackend + +logger = logging.getLogger(__name__) + + +def install_hooks_main() -> None: + _patch_capture_batch_sizes() + _patch_init_torch_distributed() + _patch_memory_pool() + _patch_torch_symm_mem() + _patch_multimem_all_gather() + _patch_backend_resolver() + _patch_spawn_sites() + + +def _patch_capture_batch_sizes() -> None: + original = decode_runner.get_batch_sizes_to_capture + + @functools.wraps(original) + def patched(model_runner, captured_req_width): + capture_bs, compile_bs = original(model_runner, captured_req_width) + if not ( + model_runner.server_args.enable_torch_symm_mem and is_symmetric_multi_shape_enabled() + ): + return capture_bs, compile_bs + requested = symmetric_graph_batch_sizes() + if requested is None: + raise RuntimeError(f"{MULTI_SHAPE_BATCHES_ENV} is required for multi-shape replay") + unknown = set(requested) - set(capture_bs) + if unknown: + raise RuntimeError(f"Requested SGLang graph batches are unavailable: {sorted(unknown)}") + filtered_compile = [batch for batch in compile_bs if batch in requested] + return list(requested), filtered_compile + + decode_runner.get_batch_sizes_to_capture = patched + + +def _patch_init_torch_distributed() -> None: + cls = mr.ModelRunner + original = cls.init_torch_distributed + + @functools.wraps(original) + def patched(self, *args, **kwargs): + mode = get_graph_extension_mode() + if mode == CUDAGraphExtensionMode.NONE: + return original(self, *args, **kwargs) + + if self.device == "cuda": + torch.get_device_module(self.device).set_device(self.gpu_id) + + ps = self.ps + rt.setup_graph_extension( + self.server_args, + tp_rank=ps.tp_rank, + pp_rank=ps.pp_rank, + dp_rank=ps.dp_rank, + ) + rt.log_alloc_offset("after_setup_graph_ext") + result = original(self, *args, **kwargs) + rt.log_alloc_offset("after_init_torch_dist") + rt.skip_to_scratch_boundary() + rt.log_alloc_offset("after_scratch_skip") + return result + + cls.init_torch_distributed = patched + + +def _patch_memory_pool() -> None: + original_resolve = KVCacheConfigurator._resolve_memory_pool_config + original_configure = KVCacheConfigurator.configure + + @functools.wraps(original_resolve) + def patched_resolve(self, pre_model_load_memory): + if get_graph_extension_mode() != CUDAGraphExtensionMode.LOAD: + return original_resolve(self, pre_model_load_memory) + + state = rt.load_warmup_state() + if not state.memory_pool_config: + raise RuntimeError("Foundry LOAD requires memory_pool_config") + torch.cuda.empty_cache() + logger.info("[Foundry] SGLang-main reused saved memory pool config") + return MemoryPoolConfig(**state.memory_pool_config) + + @functools.wraps(original_configure) + def patched_configure(self, *args, **kwargs): + mode = get_graph_extension_mode() + rt.log_alloc_offset("before_configure_memory_pool") + result = original_configure(self, *args, **kwargs) + rt.log_alloc_offset("after_configure_memory_pool") + if mode == CUDAGraphExtensionMode.SAVE: + state = rt.create_warmup_state(asdict(result.memory_pool_config)) + rt.save_warmup_state(state) + return result + + KVCacheConfigurator._resolve_memory_pool_config = patched_resolve + KVCacheConfigurator.configure = patched_configure + + +def _patch_torch_symm_mem() -> None: + original = torch_symm_mem_module.TorchSymmMemCommunicator.__init__ + + @functools.wraps(original) + def patched(self, *args, **kwargs): + if get_graph_extension_mode() == CUDAGraphExtensionMode.NONE: + return original(self, *args, **kwargs) + + # PyTorch symmetric memory owns its own cross-rank virtual-address + # contract. Allocate it outside Foundry's monotonic region, then resume + # Foundry before model weights and KV pools are created. + cge.stop_allocation_region() + try: + result = original(self, *args, **kwargs) + finally: + cge.resume_allocation_region() + if self.buffer is None: + raise RuntimeError( + "Foundry requested SGLang torch symmetric memory, but the " + "communicator is unavailable on this topology" + ) + handle = torch_symm_mem_module.torch_symm_mem.rendezvous( + self.buffer, + self.group.group_name, + ) + if handle.world_size != 2: + raise RuntimeError( + "Foundry SGLang torch symmetric-memory replay currently requires TP=2" + ) + self._foundry_symm_mem_handle = handle + logger.info( + "[Foundry] Torch symmetric-memory communicator ready: rank=%d world_size=%d", + handle.rank, + handle.world_size, + ) + return result + + torch_symm_mem_module.TorchSymmMemCommunicator.__init__ = patched + + +def _patch_multimem_all_gather() -> None: + original = triton_symm_mem_ag.MultimemAllGatherer.__init__ + + @functools.wraps(original) + def patched(self, max_tokens, *, enabled=True, skip_entry_sync=False): + if get_graph_extension_mode() != CUDAGraphExtensionMode.NONE: + enabled = False + return original( + self, + max_tokens, + enabled=enabled, + skip_entry_sync=skip_entry_sync, + ) + + triton_symm_mem_ag.MultimemAllGatherer.__init__ = patched + + +def _patch_backend_resolver() -> None: + original = backend_utils.resolve_decode_backend + + @functools.wraps(original) + def patched(cuda_graph_runner): + mode = get_graph_extension_mode() + if mode == CUDAGraphExtensionMode.NONE: + return original(cuda_graph_runner) + + backend_name = cuda_graph_runner.model_runner.server_args.cuda_graph_config.decode.backend + if backend_name != "full": + raise RuntimeError( + "Foundry SGLang-main requires the full decode CUDA-graph backend, " + f"got {backend_name!r}" + ) + return FoundryMainCudaGraphBackend(cuda_graph_runner) + + backend_utils.resolve_decode_backend = patched + decode_runner.resolve_decode_backend = patched + + +def _patch_spawn_sites() -> None: + launch_descriptor = engine_mod.Engine.__dict__["_launch_scheduler_processes"] + is_classmethod = isinstance(launch_descriptor, classmethod) + original_launch = launch_descriptor.__func__ if is_classmethod else launch_descriptor + + @functools.wraps(original_launch) + def patched_launch(owner, *args, **kwargs): + if get_graph_extension_mode() != CUDAGraphExtensionMode.NONE: + rt.setup_ld_preload_env() + return original_launch(owner, *args, **kwargs) + + engine_mod.Engine._launch_scheduler_processes = ( + classmethod(patched_launch) if is_classmethod else patched_launch + ) + + original_start = dpc.DataParallelController.launch_tensor_parallel_group + + @functools.wraps(original_start) + def patched_start(self, *args, **kwargs): + if get_graph_extension_mode() != CUDAGraphExtensionMode.NONE: + rt.setup_ld_preload_env() + return original_start(self, *args, **kwargs) + + dpc.DataParallelController.launch_tensor_parallel_group = patched_start diff --git a/python/foundry/integration/sglang/main_backend.py b/python/foundry/integration/sglang/main_backend.py new file mode 100644 index 00000000..e85ba019 --- /dev/null +++ b/python/foundry/integration/sglang/main_backend.py @@ -0,0 +1,513 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the Foundry project +"""Full decode-graph backend for the current SGLang runner API.""" + +from __future__ import annotations + +import json +import logging +import os +import shutil +import tempfile +import time +from collections.abc import Callable, Iterator +from contextlib import contextmanager +from dataclasses import fields +from pathlib import Path +from typing import Any + +from sglang.srt.distributed.device_communicators.pynccl_allocator import ( + set_graph_pool_id, +) +from sglang.srt.layers.logits_processor import LogitsProcessorOutput +from sglang.srt.model_executor.runner_backend.base_cuda_graph_backend import ( + BaseCudaGraphBackend, +) +from sglang.srt.model_executor.runner_utils.pool import ( + get_or_create_global_graph_memory_pool, +) + +import foundry as foundry_pkg +from foundry import ops as cge +from foundry.graph import CUDAGraph as FoundryCUDAGraph +from foundry.graph import graph as foundry_graph +from foundry.integration.sglang import runtime as rt +from foundry.integration.sglang.config import ( + CUDAGraphExtensionMode, + get_config, + get_graph_extension_mode, + is_symmetric_multi_shape_enabled, +) +from foundry.integration.sglang.flashinfer_graph_abi import ( + KERNEL_ABI_SCHEMA, + extract_flashinfer_operand_records, + relocate_flashinfer_operands, +) +from foundry.integration.sglang.graph_ops import ( + _pack_output, + _scan_graph_files, + _unpack_output, +) +from foundry.integration.sglang.sglang_shape_adapter import ( + PLAN_SCHEMA, + activate_shape_state, + adopt_prepared_shape_state, + close_shape_state, + create_shape_state, + prepare_shape_state, +) +from foundry.integration.sglang.shape_replay_state import ShapeReplayState +from foundry.integration.sglang.state_manifest import ( + ShapeStateManifest, + build_shape_record, + clear_shape_state_manifest, + read_shape_state_manifest, + validate_shape_record, + write_shape_state_manifest, +) +from foundry.integration.sglang.symm_mem_graph import ( + FULL_GRAPH_DIR, + GraphOperandInventory, + SymmetricMemoryGraphState, + clear_symmetric_graph_state, + preserve_symmetric_graphs, + relocate_symmetric_graphs, + run_graph_warmups, + validate_graph_backend, + write_graph_backend_metadata, +) + +logger = logging.getLogger(__name__) + + +class FoundryMainCudaGraphBackend(BaseCudaGraphBackend): + """Materialize SGLang main's full decode graphs through Foundry.""" + + def __init__(self, cuda_graph_runner) -> None: + self._runner = cuda_graph_runner + self._graphs: dict[Any, FoundryCUDAGraph] = {} + self._outputs: dict[Any, Any] = {} + self._pool = None + self._stream = None + self._pending = None + self._graph_files = [] + self._load_index = 0 + self._closed = False + self._symmetric_memory_enabled = ( + cuda_graph_runner.model_runner.server_args.enable_torch_symm_mem + ) + self._multi_shape_enabled = ( + self._symmetric_memory_enabled and is_symmetric_multi_shape_enabled() + ) + self._graph_load_paths: list[str] = [] + self._relocated_graph_dir: str | None = None + self._shape_states: dict[Any, ShapeReplayState] = {} + self._saved_shape_manifest: ShapeStateManifest | None = None + self._saved_records_by_batch = {} + self._operand_inventories: dict[str, GraphOperandInventory] = {} + if self._symmetric_memory_enabled: + logger.info("[Foundry] SGLang-main communication backend=torch_symmetric_memory") + + @staticmethod + def _validate_shape_key(shape_key) -> int: + if shape_key.stream_idx is not None or shape_key.variant_label is not None: + raise RuntimeError( + "Foundry SGLang-main currently supports one decode stream " + "without LoRA graph variants" + ) + return int(shape_key.size) + + def _symmetric_memory_state(self) -> SymmetricMemoryGraphState: + communicator = self._runner.model_runner.tp_group.torch_symm_mem_comm + handle = communicator._foundry_symm_mem_handle + return SymmetricMemoryGraphState( + buffer=communicator.buffer.data_ptr(), + buffer_ptrs_dev=handle.buffer_ptrs_dev, + signal_pad_ptrs_dev=handle.signal_pad_ptrs_dev, + multicast_ptr=handle.multicast_ptr, + buffer_numel=communicator.buffer.numel(), + element_size=communicator.buffer.element_size(), + dtype=str(communicator.buffer.dtype), + rank=handle.rank, + world_size=handle.world_size, + ) + + def _cleanup_relocated_graphs(self) -> None: + if self._relocated_graph_dir is not None: + shutil.rmtree(self._relocated_graph_dir, ignore_errors=True) + self._relocated_graph_dir = None + + @contextmanager + def capture_session(self, stream) -> Iterator[None]: + if self._closed: + raise RuntimeError("Foundry SGLang-main does not support CUDA graph recapture") + if self._pool is None: + self._pool = get_or_create_global_graph_memory_pool( + self._runner.device_module, + ) + set_graph_pool_id(self._pool) + self._stream = stream + + mode = get_graph_extension_mode() + completed = False + if mode == CUDAGraphExtensionMode.LOAD: + self._prepare_load() + + try: + yield + completed = True + finally: + self._stream = None + try: + if completed and mode == CUDAGraphExtensionMode.SAVE: + self._finish_save() + elif ( + completed + and mode == CUDAGraphExtensionMode.LOAD + and self._load_index != len(self._graph_files) + ): + raise RuntimeError( + f"Foundry loaded {self._load_index}/{len(self._graph_files)} SGLang graphs" + ) + elif completed and mode == CUDAGraphExtensionMode.LOAD: + rt.log_alloc_offset("after_load_all_graphs") + logger.info( + "[Foundry] Loaded %d SGLang graphs", + self._load_index, + ) + finally: + if mode == CUDAGraphExtensionMode.LOAD: + self._cleanup_relocated_graphs() + + def _prepare_load(self) -> None: + cfg = get_config() + if cfg is None or cfg.workspace_dir is None: + raise RuntimeError("Foundry SGLang graph extension is not initialized") + + validate_graph_backend( + cfg.workspace_dir, + self._symmetric_memory_enabled, + ) + + self._graph_files = _scan_graph_files(cfg.workspace_dir) + if not self._graph_files: + raise RuntimeError(f"No Foundry SGLang graph files found in {cfg.workspace_dir}") + paths = [ + os.path.join(cfg.workspace_dir, filename) + for _index, filename, _meta in self._graph_files + ] + if self._symmetric_memory_enabled: + rt.log_alloc_offset("before_symmetric_graph_load") + filenames = [filename for _index, filename, _meta in self._graph_files] + self._relocated_graph_dir = tempfile.mkdtemp(prefix="foundry-sglang-symmetric-") + try: + if self._multi_shape_enabled: + self._saved_shape_manifest = read_shape_state_manifest(cfg.workspace_dir) + self._saved_records_by_batch = { + record.batch_size: record for record in self._saved_shape_manifest.shapes + } + self._graph_load_paths = relocate_symmetric_graphs( + cfg.workspace_dir, + filenames, + self._symmetric_memory_state(), + output_dir=self._relocated_graph_dir, + manifest=self._saved_shape_manifest if self._multi_shape_enabled else None, + allow_multi_shape=self._multi_shape_enabled, + ) + except BaseException: + self._cleanup_relocated_graphs() + raise + self._pending = None + logger.info( + "[Foundry] Relocated %d SGLang symmetric-memory graphs", + len(self._graph_load_paths), + ) + else: + rt.log_alloc_offset("before_preallocate") + rt.preallocate_for_load_mode() + rt.log_alloc_offset("after_preallocate") + self._graph_load_paths = paths + self._pending = FoundryCUDAGraph.start_graph_builds(paths, num_threads=4) + self._load_index = 0 + logger.info( + "[Foundry] Started %d SGLang-main graph builds", + len(paths), + ) + + def capture_one( + self, + shape_key, + forward_fn: Callable[[], Any], + capture_inputs: Any = None, + post_warmup_hook: Callable[[], None] | None = None, + ) -> None: + del capture_inputs + size = self._validate_shape_key(shape_key) + mode = get_graph_extension_mode() + + if not self._multi_shape_enabled: + if self._symmetric_memory_enabled: + run_graph_warmups( + synchronize=self._runner.device_module.synchronize, + barrier=self._runner.model_runner.tp_group.barrier, + forward=forward_fn, + post_warmup=post_warmup_hook, + ) + + if mode == CUDAGraphExtensionMode.SAVE: + graph, output = self._capture_one(shape_key, size, forward_fn) + self._graphs[shape_key] = graph + self._outputs[shape_key] = output + return + if mode == CUDAGraphExtensionMode.LOAD: + graph, output = self._load_one(shape_key, size) + self._graphs[shape_key] = graph + self._outputs[shape_key] = output + return + raise RuntimeError(f"Foundry backend used in unsupported mode: {mode.value}") + + communicator = self._runner.model_runner.tp_group.torch_symm_mem_comm + state = create_shape_state( + self._runner, + shape_key=shape_key, + capture_index=self._load_index + if mode == CUDAGraphExtensionMode.LOAD + else len(self._shape_states), + communicator=communicator, + symm_handle=communicator._foundry_symm_mem_handle, + ) + if mode == CUDAGraphExtensionMode.LOAD: + validate_shape_record( + state, + self._saved_records_by_batch[state.batch_size], + plan_schema=PLAN_SCHEMA, + ) + with activate_shape_state(state): + run_graph_warmups( + synchronize=self._runner.device_module.synchronize, + barrier=self._runner.model_runner.tp_group.barrier, + forward=forward_fn, + post_warmup=post_warmup_hook, + ) + prepare_shape_state(state, None, for_capture=True) + if mode == CUDAGraphExtensionMode.SAVE: + graph, outputs = self._capture_one(shape_key, size, forward_fn) + elif mode == CUDAGraphExtensionMode.LOAD: + record = self._saved_records_by_batch[state.batch_size] + graph_path = Path(self._graph_load_paths[self._load_index]) + graph_data = json.loads(graph_path.read_text()) + flashinfer_records = tuple( + operand + for operand in record.operand_slots + if operand.kernel_abi == KERNEL_ABI_SCHEMA + ) + relocate_flashinfer_operands( + graph_data, + state, + flashinfer_records, + ) + graph_path.write_text(json.dumps(graph_data)) + graph, outputs = self._load_one(shape_key, size) + else: + raise RuntimeError(f"Unsupported Foundry mode: {mode.value}") + state.attach_graph(graph, outputs) + self._shape_states[shape_key] = state + self._graphs[shape_key] = graph + self._outputs[shape_key] = outputs + + def _capture_one( + self, + shape_key, + size: int, + forward_fn: Callable[[], Any], + ) -> tuple[FoundryCUDAGraph, Any]: + if self._stream is None: + raise RuntimeError("Foundry capture session has no CUDA stream") + + self._runner.device_module.synchronize() + self._runner.model_runner.tp_group.barrier() + + graph = FoundryCUDAGraph() + with foundry_graph( + graph, + pool=self._pool, + stream=self._stream, + ): + output = forward_fn() + + self._save_graph(graph, output, size) + return graph, output + + @staticmethod + def _save_graph(graph: FoundryCUDAGraph, output: Any, size: int) -> None: + cfg = get_config() + 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") + + if not isinstance(output, LogitsProcessorOutput): + raise RuntimeError( + "Foundry SGLang-main currently supports decode logits output only, " + f"got {type(output)!r}" + ) + active = [ + field.name + for field in fields(output) + if field.name != "next_token_logits" and getattr(output, field.name) is not None + ] + if active: + raise RuntimeError( + "Foundry SGLang-main cannot serialize output fields: " + ", ".join(active) + ) + + packed_output = _pack_output(output) + filename = f"graph_{state.capture_index}_FULL_t{size}_r{size}_UX_pcN.json" + graph_path = os.path.join(cfg.workspace_dir, filename) + graph.save(graph_path, packed_output) + state.capture_index += 1 + logger.info( + "[Foundry] Saved SGLang-main CUDA graph %s key=%s", + filename, + size, + ) + + def _load_one(self, shape_key, size: int) -> tuple[FoundryCUDAGraph, Any]: + if self._pending is None and not self._symmetric_memory_enabled: + raise RuntimeError("Foundry graph builds were not started") + if self._load_index >= len(self._graph_files): + raise RuntimeError("SGLang requested more graph shapes than were saved") + + _index, filename, meta = self._graph_files[self._load_index] + if int(meta["key"]) != size: + raise RuntimeError( + f"SGLang graph order mismatch: requested {size}, archive has " + f"{meta['key']} ({filename})" + ) + + started_at = time.perf_counter() + if self._symmetric_memory_enabled: + result = FoundryCUDAGraph.load( + self._graph_load_paths[self._load_index], + pool=self._pool, + ) + if not isinstance(result, tuple): + raise RuntimeError(f"Loaded graph {filename} has no output tensors") + graph, tensors = result + else: + graph, tensors = FoundryCUDAGraph.finish_one_graph_load( + self._pending, + self._load_index, + ) + self._load_index += 1 + output = _unpack_output(tensors) + logger.info( + "[Foundry] Loaded SGLang-main graph %s in %.3fs", + filename, + time.perf_counter() - started_at, + ) + return graph, output + + def _finish_save(self) -> None: + cfg = get_config() + if cfg is None or cfg.workspace_dir is None: + raise RuntimeError("Foundry SGLang graph extension is not initialized") + clear_shape_state_manifest(cfg.workspace_dir) + if self._symmetric_memory_enabled: + graph_filenames = [ + filename for _index, filename, _meta in _scan_graph_files(cfg.workspace_dir) + ] + if self._multi_shape_enabled: + inventories = preserve_symmetric_graphs( + cfg.workspace_dir, + graph_filenames, + self._symmetric_memory_state(), + allow_multi_shape=True, + ) + self._operand_inventories = { + inventory.graph_filename: inventory for inventory in inventories + } + states_by_batch = {state.batch_size: state for state in self._shape_states.values()} + if set(states_by_batch) != {inventory.batch_size for inventory in inventories}: + raise RuntimeError("SGLang capsule and graph shape inventories differ") + shape_records_list = [] + for inventory in inventories: + state = states_by_batch[inventory.batch_size] + graph_path = Path(cfg.workspace_dir) / FULL_GRAPH_DIR / inventory.graph_filename + graph_data = json.loads(graph_path.read_text()) + flashinfer_operands = extract_flashinfer_operand_records(graph_data, state) + graph_path.write_text(json.dumps(graph_data)) + shape_records_list.append( + build_shape_record( + state, + graph_filename=inventory.graph_filename, + plan_schema=PLAN_SCHEMA, + operand_slots=inventory.operand_slots + flashinfer_operands, + ) + ) + shape_records = tuple(shape_records_list) + symmetric_state = self._symmetric_memory_state() + manifest = ShapeStateManifest( + backend="torch_symmetric_memory", + rank=symmetric_state.rank, + world_size=symmetric_state.world_size, + capture_order=tuple(record.batch_size for record in shape_records), + shapes=shape_records, + ) + write_shape_state_manifest(cfg.workspace_dir, manifest) + else: + preserve_symmetric_graphs( + cfg.workspace_dir, + graph_filenames, + self._symmetric_memory_state(), + ) + else: + clear_symmetric_graph_state(cfg.workspace_dir) + write_graph_backend_metadata( + cfg.workspace_dir, + self._symmetric_memory_enabled, + ) + foundry_pkg.save_graph_manifest(cfg.workspace_dir) + cge.pack_fatbins_to_folder(cfg.workspace_dir) + cge.set_pack_fatbins_on_exit(False) + rt.capture_final_alloc_offset() + + def can_run(self, forward_batch, shape_key) -> bool: + del forward_batch + return shape_key in self._graphs + + @contextmanager + def replay_session(self) -> Iterator[None]: + yield + + def replay(self, shape_key, static_forward_batch, **kwargs) -> Any: + del static_forward_batch, kwargs + if not self._multi_shape_enabled: + self._graphs[shape_key].replay() + return self._outputs[shape_key] + + state = self._shape_states[shape_key] + adopt_prepared_shape_state(state) + with activate_shape_state(state): + state.graph.replay() + return state.outputs + + def cleanup(self) -> None: + self._cleanup_relocated_graphs() + if self._graphs: + self._closed = True + self._graphs.clear() + self._outputs.clear() + for state in sorted( + self._shape_states.values(), + key=lambda item: item.capture_index, + ): + close_shape_state(state) + self._shape_states.clear() + self._pool = None + self._pending = None + self._graph_files = [] + self._graph_load_paths = [] + self._saved_shape_manifest = None + self._saved_records_by_batch = {} + self._operand_inventories = {} + self._load_index = 0 diff --git a/python/foundry/integration/sglang/plugin.py b/python/foundry/integration/sglang/plugin.py new file mode 100644 index 00000000..760184da --- /dev/null +++ b/python/foundry/integration/sglang/plugin.py @@ -0,0 +1,82 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the Foundry project +"""SGLang-main plugin entry point for Foundry activation.""" + +from __future__ import annotations + +import os + +from sglang.srt.plugins.hook_registry import HookRegistry, HookType + +from foundry.integration.sglang.hooks import install_hooks_from_path + +CONFIG_PATH_ENV = "FOUNDRY_SGLANG_CONFIG_PATH" +_installed = False +_installation_error: BaseException | None = None + + +def _configure_server_args(server_args, *args, **kwargs) -> None: + """Select the main APIs Foundry currently materializes. + + This hook runs before ``ServerArgs.__post_init__`` resolves + ``cuda_graph_config``, so the convenience fields below feed the normal + SGLang parser rather than mutating a finalized configuration. + """ + + if _installation_error is not None: + raise RuntimeError("Foundry SGLang-main hook installation failed") from _installation_error + if not _installed: + raise RuntimeError("Foundry SGLang-main plugin was not installed") + + server_args.cuda_graph_backend_decode = "full" + server_args.disable_prefill_cuda_graph = True + server_args.disable_flashinfer_autotune = True + server_args.enable_profile_cuda_graph = False + + +def _validate_server_args(result, server_args, *args, **kwargs) -> None: + del result + if server_args.cuda_graph_config.decode.backend != "full": + raise RuntimeError("Foundry requires SGLang's full decode graph backend") + if server_args.cuda_graph_config.prefill.backend != "disabled": + raise RuntimeError("Foundry does not support SGLang prefill graph capture") + if server_args.pp_size != 1: + raise RuntimeError("Foundry SGLang-main does not yet support pipeline parallelism") + if server_args.speculative_algorithm is not None: + raise RuntimeError("Foundry SGLang-main does not yet support speculative decoding") + if server_args.enable_lora: + raise RuntimeError("Foundry SGLang-main does not yet support LoRA graph variants") + if server_args.enable_pdmux: + raise RuntimeError("Foundry SGLang-main does not yet support PD multiplexing") + if server_args.enable_return_hidden_states: + raise RuntimeError("Foundry SGLang-main does not serialize hidden-state output") + if server_args.dllm_algorithm is not None: + raise RuntimeError("Foundry SGLang-main does not yet support diffusion models") + if "deepep" in str(server_args.moe_a2a_backend).lower(): + raise RuntimeError("Foundry SGLang-main DeepEP port is not yet complete") + + +def register() -> None: + """Register Foundry when the config-path environment variable is set.""" + + global _installed, _installation_error + config_path = os.environ.get(CONFIG_PATH_ENV) + if not config_path: + return + + HookRegistry.register( + "sglang.srt.server_args.ServerArgs.__post_init__", + _configure_server_args, + HookType.BEFORE, + ) + HookRegistry.register( + "sglang.srt.server_args.ServerArgs.__post_init__", + _validate_server_args, + HookType.AFTER, + ) + try: + install_hooks_from_path(config_path) + except BaseException as error: + _installation_error = error + raise + _installed = True diff --git a/python/foundry/integration/sglang/sglang_shape_adapter.py b/python/foundry/integration/sglang/sglang_shape_adapter.py new file mode 100644 index 00000000..2bc1948e --- /dev/null +++ b/python/foundry/integration/sglang/sglang_shape_adapter.py @@ -0,0 +1,186 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the Foundry project +from __future__ import annotations + +from collections.abc import Iterator +from contextlib import contextmanager +from typing import Any + +from foundry.integration.sglang.shape_replay_state import ( + ShapeReplayState, + ShapeStateLifecycle, + TensorOwnerSlot, +) + +PLAN_SCHEMA = "flashinfer.prefill.v15" + +_SHAPE_FIELDS = ( + ("int_workspace", "_int_workspace_buffer"), + ("pin_memory_int_workspace", "_pin_memory_int_workspace_buffer"), + ("qo_indptr", "_qo_indptr_buf"), +) + +_SHARED_FIELDS = ( + ("paged_kv_indptr", "_paged_kv_indptr_buf"), + ("paged_kv_indices", "_paged_kv_indices_buf"), + ("paged_kv_last_page_len", "_paged_kv_last_page_len_buf"), + ("float_workspace", "_float_workspace_buffer"), +) + + +def _require_field(owner: Any, field_name: str) -> Any: + if not hasattr(owner, field_name): + raise RuntimeError( + f"Pinned FlashInfer wrapper ABI changed: missing required field {field_name}" + ) + return getattr(owner, field_name) + + +def _plan_fingerprint(wrapper: Any) -> tuple[int, ...]: + plan_info = tuple(int(value) for value in _require_field(wrapper, "_plan_info")) + if len(plan_info) != 15: + raise RuntimeError( + f"Pinned FlashInfer plan schema changed: expected 15 fields, got {len(plan_info)}" + ) + if _require_field(wrapper, "_cached_module") is None: + raise RuntimeError("FlashInfer wrapper has no cached CUDA graph module") + return plan_info + + +def _restore_wrapper_binding(backend: Any, batch_size: int, previous_wrappers: Any) -> None: + if previous_wrappers is None: + backend.decode_cuda_graph_metadata.pop(batch_size, None) + else: + backend.decode_cuda_graph_metadata[batch_size] = previous_wrappers + + +def _require_live_wrapper_identity(state: ShapeReplayState) -> None: + live_wrappers = tuple( + state.attention_backend.decode_cuda_graph_metadata.get(state.batch_size, ()) + ) + if len(live_wrappers) != len(state.wrappers) or any( + live_wrapper is not expected_wrapper + for live_wrapper, expected_wrapper in zip(live_wrappers, state.wrappers, strict=False) + ): + raise RuntimeError( + "SGLang backend replaced " + f"decode_cuda_graph_metadata[{state.batch_size}] during replay preparation; " + "pinned FlashInfer replay must preserve wrapper identity" + ) + + +def create_shape_state( + cuda_graph_runner: Any, + shape_key: Any, + capture_index: int, + communicator: Any, + symm_handle: Any, +) -> ShapeReplayState: + if int(cuda_graph_runner.captured_req_width) != 1: + raise RuntimeError("SGLang multi-shape replay requires captured_req_width=1") + + batch_size = int(shape_key.size if hasattr(shape_key, "size") else shape_key) + backend = cuda_graph_runner.attn_backend + wrappers = tuple(backend.decode_cuda_graph_metadata.get(batch_size, ())) + if len(wrappers) != 1: + raise RuntimeError( + f"Expected one FlashInfer wrapper for batch {batch_size}, got {len(wrappers)}" + ) + + wrapper = wrappers[0] + owners = tuple( + TensorOwnerSlot.capture( + name, + _require_field(wrapper, field_name), + ownership="shape", + ) + for name, field_name in _SHAPE_FIELDS + ) + shared_owners = tuple( + TensorOwnerSlot.capture( + name, + _require_field(wrapper, field_name), + ownership="shared", + ) + for name, field_name in _SHARED_FIELDS + ) + state = ShapeReplayState( + shape_key=shape_key, + batch_size=batch_size, + capture_index=capture_index, + attention_backend=backend, + wrappers=wrappers, + metadata=backend.forward_metadata, + owners=owners, + shared_owners=shared_owners, + shared_objects=( + cuda_graph_runner.buffers, + cuda_graph_runner.model_runner.graph_shared_output, + ), + communicator=communicator, + symm_handle=symm_handle, + ) + state.mark_planned(_plan_fingerprint(wrapper)) + return state + + +@contextmanager +def activate_shape_state(state: ShapeReplayState) -> Iterator[None]: + if state.active: + raise RuntimeError(f"SGLang shape {state.batch_size} is already active") + if state.lifecycle is ShapeStateLifecycle.CLOSED: + raise RuntimeError(f"SGLang shape {state.batch_size} is closed") + + backend = state.attention_backend + previous_wrappers = backend.decode_cuda_graph_metadata.get(state.batch_size) + previous_metadata = backend.forward_metadata + state.active = True + try: + backend.decode_cuda_graph_metadata[state.batch_size] = list(state.wrappers) + backend.forward_metadata = state.metadata + yield + state.metadata = backend.forward_metadata + finally: + try: + _restore_wrapper_binding(backend, state.batch_size, previous_wrappers) + finally: + try: + backend.forward_metadata = previous_metadata + finally: + state.active = False + + +def prepare_shape_state( + state: ShapeReplayState, + forward_batch: Any, + *, + for_capture: bool, +) -> None: + if not state.active: + raise RuntimeError(f"SGLang shape {state.batch_size} must be active before preparation") + if not for_capture: + state.attention_backend.init_forward_metadata_out_graph( + forward_batch, + in_capture=False, + ) + _require_live_wrapper_identity(state) + state.validate_owners() + actual = _plan_fingerprint(state.wrappers[0]) + if actual != state.plan_fingerprint: + raise RuntimeError(f"FlashInfer plan topology changed for batch {state.batch_size}") + + +def adopt_prepared_shape_state(state: ShapeReplayState) -> None: + _require_live_wrapper_identity(state) + state.validate_owners() + actual = _plan_fingerprint(state.wrappers[0]) + if actual != state.plan_fingerprint: + raise RuntimeError(f"FlashInfer plan topology changed for batch {state.batch_size}") + live_metadata = state.attention_backend.forward_metadata + if live_metadata is None: + raise RuntimeError(f"SGLang shape {state.batch_size} has no live forward_metadata") + state.metadata = live_metadata + + +def close_shape_state(state: ShapeReplayState) -> None: + state.close() diff --git a/python/foundry/integration/sglang/shape_replay_state.py b/python/foundry/integration/sglang/shape_replay_state.py new file mode 100644 index 00000000..7acdea27 --- /dev/null +++ b/python/foundry/integration/sglang/shape_replay_state.py @@ -0,0 +1,121 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the Foundry project +from __future__ import annotations + +from dataclasses import dataclass, field +from enum import StrEnum +from typing import Any, Literal + + +class ShapeStateLifecycle(StrEnum): + CREATED = "created" + PLANNED = "planned" + MATERIALIZED = "materialized" + CLOSED = "closed" + + +@dataclass(frozen=True) +class TensorOwnerSlot: + name: str + tensor: Any + ownership: Literal["shape", "shared"] + data_ptr: int + numel: int + element_size: int + dtype: str + device: str + + @classmethod + def capture( + cls, + name: str, + tensor: Any, + *, + ownership: Literal["shape", "shared"], + ) -> TensorOwnerSlot: + return cls( + name=name, + tensor=tensor, + ownership=ownership, + data_ptr=int(tensor.data_ptr()), + numel=int(tensor.numel()), + element_size=int(tensor.element_size()), + dtype=str(tensor.dtype), + device=str(tensor.device), + ) + + def validate_live(self) -> None: + actual = ( + int(self.tensor.data_ptr()), + int(self.tensor.numel()), + int(self.tensor.element_size()), + str(self.tensor.dtype), + str(self.tensor.device), + ) + expected = ( + self.data_ptr, + self.numel, + self.element_size, + self.dtype, + self.device, + ) + if actual != expected: + raise RuntimeError( + f"SGLang shape-state owner changed for {self.name}: {actual!r} != {expected!r}" + ) + + +@dataclass +class ShapeReplayState: + shape_key: Any + batch_size: int + capture_index: int + attention_backend: Any + wrappers: tuple[Any, ...] + metadata: Any + owners: tuple[TensorOwnerSlot, ...] + shared_owners: tuple[TensorOwnerSlot, ...] + shared_objects: tuple[Any, ...] + communicator: Any + symm_handle: Any + plan_fingerprint: tuple[int, ...] = () + graph: Any = None + outputs: Any = None + lifecycle: ShapeStateLifecycle = ShapeStateLifecycle.CREATED + active: bool = False + operand_inventory: dict[str, int] = field(default_factory=dict) + + def validate_owners(self) -> None: + for slot in (*self.owners, *self.shared_owners): + slot.validate_live() + + def mark_planned(self, plan_fingerprint: tuple[int, ...]) -> None: + if self.lifecycle not in { + ShapeStateLifecycle.CREATED, + ShapeStateLifecycle.PLANNED, + }: + raise RuntimeError(f"Cannot plan SGLang shape state from {self.lifecycle.value}") + self.plan_fingerprint = plan_fingerprint + self.lifecycle = ShapeStateLifecycle.PLANNED + + def attach_graph(self, graph: Any, outputs: Any) -> None: + if self.lifecycle is not ShapeStateLifecycle.PLANNED: + raise RuntimeError("SGLang shape state must be planned before graph attach") + self.graph = graph + self.outputs = outputs + self.lifecycle = ShapeStateLifecycle.MATERIALIZED + + def close(self) -> None: + if self.active: + raise RuntimeError("Cannot close an active SGLang shape state") + self.graph = None + self.outputs = None + self.metadata = None + self.wrappers = () + self.owners = () + self.shared_owners = () + self.shared_objects = () + self.attention_backend = None + self.symm_handle = None + self.communicator = None + self.lifecycle = ShapeStateLifecycle.CLOSED diff --git a/python/foundry/integration/sglang/state_manifest.py b/python/foundry/integration/sglang/state_manifest.py new file mode 100644 index 00000000..c93a89fd --- /dev/null +++ b/python/foundry/integration/sglang/state_manifest.py @@ -0,0 +1,244 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the Foundry project +from __future__ import annotations + +import hashlib +import json +from dataclasses import asdict, dataclass +from pathlib import Path +from typing import TYPE_CHECKING, Literal +from uuid import uuid4 + +if TYPE_CHECKING: + from foundry.integration.sglang.shape_replay_state import ShapeReplayState + +MANIFEST_FILENAME = "sglang_shape_state.json" +MANIFEST_VERSION = 1 + + +@dataclass(frozen=True) +class TensorSlotRecord: + name: str + ownership: Literal["shape", "shared"] + dtype: str + device: str + numel: int + element_size: int + + +@dataclass(frozen=True) +class OperandSlotRecord: + name: str + kernel_abi: str + owner_slot: str + node_id: int + source: str + parameter_index: int + cuda_parameter_offset: int + value_byte_offset: int + ctype: str + owner_relative_offset: int + span_bytes: int + saved_value: int + + +@dataclass(frozen=True) +class ShapeStateRecord: + graph_filename: str + batch_size: int + capture_index: int + plan_schema: str + plan_fingerprint: tuple[int, ...] + tensor_slots: tuple[TensorSlotRecord, ...] + operand_slots: tuple[OperandSlotRecord, ...] + + +@dataclass(frozen=True) +class ShapeStateManifest: + backend: str + rank: int + world_size: int + capture_order: tuple[int, ...] + shapes: tuple[ShapeStateRecord, ...] + + +def build_shape_record( + state: ShapeReplayState, + *, + graph_filename: str, + plan_schema: str, + operand_slots: tuple[OperandSlotRecord, ...], +) -> ShapeStateRecord: + tensor_slots = tuple( + TensorSlotRecord( + name=slot.name, + ownership=slot.ownership, + dtype=slot.dtype, + device=slot.device, + numel=slot.numel, + element_size=slot.element_size, + ) + for slot in (*state.owners, *state.shared_owners) + ) + return ShapeStateRecord( + graph_filename=graph_filename, + batch_size=state.batch_size, + capture_index=state.capture_index, + plan_schema=plan_schema, + plan_fingerprint=state.plan_fingerprint, + tensor_slots=tensor_slots, + operand_slots=operand_slots, + ) + + +def validate_shape_record( + state: ShapeReplayState, + record: ShapeStateRecord, + *, + plan_schema: str, +) -> None: + expected = build_shape_record( + state, + graph_filename=record.graph_filename, + plan_schema=plan_schema, + operand_slots=record.operand_slots, + ) + if expected != record: + raise RuntimeError(f"SGLang live shape state does not match batch {record.batch_size}") + + +def _shape_to_dict(shape: ShapeStateRecord) -> dict: + return { + "graph_filename": shape.graph_filename, + "batch_size": shape.batch_size, + "capture_index": shape.capture_index, + "plan_schema": shape.plan_schema, + "plan_fingerprint": list(shape.plan_fingerprint), + "tensor_slots": [asdict(slot) for slot in shape.tensor_slots], + "operand_slots": [asdict(slot) for slot in shape.operand_slots], + } + + +def _semantic_fingerprint(manifest: ShapeStateManifest) -> str: + semantic = { + "backend": manifest.backend, + "rank": manifest.rank, + "world_size": manifest.world_size, + "capture_order": list(manifest.capture_order), + "shapes": [ + { + **_shape_to_dict(shape), + "operand_slots": [ + {key: value for key, value in asdict(slot).items() if key != "saved_value"} + for slot in shape.operand_slots + ], + } + for shape in manifest.shapes + ], + } + encoded = json.dumps( + semantic, + sort_keys=True, + separators=(",", ":"), + ).encode() + return hashlib.sha256(encoded).hexdigest() + + +def _validate_manifest(manifest: ShapeStateManifest) -> None: + if manifest.backend != "torch_symmetric_memory": + raise RuntimeError(f"Unsupported SGLang shape-state backend: {manifest.backend}") + if manifest.world_size != 2 or manifest.rank not in (0, 1): + raise RuntimeError("SGLang multi-shape symmetric replay requires TP=2") + batches = tuple(shape.batch_size for shape in manifest.shapes) + if len(batches) != len(set(batches)): + raise RuntimeError("SGLang shape-state manifest contains duplicate shapes") + if batches != manifest.capture_order: + raise RuntimeError( + f"SGLang shape-state order mismatch: {batches} != {manifest.capture_order}" + ) + indices = tuple(shape.capture_index for shape in manifest.shapes) + if indices != tuple(range(len(indices))): + raise RuntimeError(f"SGLang capture indices are not contiguous: {indices}") + for shape in manifest.shapes: + names = tuple(slot.name for slot in shape.tensor_slots) + if len(names) != len(set(names)): + raise RuntimeError(f"Duplicate tensor slot in batch {shape.batch_size}: {names}") + locations = tuple( + ( + slot.node_id, + slot.source, + slot.parameter_index, + slot.value_byte_offset, + ) + for slot in shape.operand_slots + ) + if len(locations) != len(set(locations)): + raise RuntimeError(f"Duplicate operand slot in batch {shape.batch_size}") + + +def validate_shape_state_manifest(manifest: ShapeStateManifest) -> None: + _validate_manifest(manifest) + + +def clear_shape_state_manifest(workspace: str | Path) -> None: + path = Path(workspace) / MANIFEST_FILENAME + path.unlink(missing_ok=True) + + +def write_shape_state_manifest( + workspace: str | Path, + manifest: ShapeStateManifest, +) -> None: + validate_shape_state_manifest(manifest) + payload = { + "version": MANIFEST_VERSION, + "schema_fingerprint": _semantic_fingerprint(manifest), + "manifest": { + "backend": manifest.backend, + "rank": manifest.rank, + "world_size": manifest.world_size, + "capture_order": list(manifest.capture_order), + "shapes": [_shape_to_dict(shape) for shape in manifest.shapes], + }, + } + path = Path(workspace) / MANIFEST_FILENAME + temp_path = path.with_name(f".{path.name}.{uuid4().hex}.tmp") + try: + temp_path.write_text(json.dumps(payload, sort_keys=True)) + temp_path.replace(path) + finally: + temp_path.unlink(missing_ok=True) + + +def read_shape_state_manifest(workspace: str | Path) -> ShapeStateManifest: + path = Path(workspace) / MANIFEST_FILENAME + try: + payload = json.loads(path.read_text()) + except (OSError, json.JSONDecodeError) as error: + raise RuntimeError("Invalid SGLang shape-state manifest") from error + if payload.get("version") != MANIFEST_VERSION: + raise RuntimeError(f"Unsupported SGLang shape-state version: {payload.get('version')!r}") + raw = payload["manifest"] + shapes = tuple( + ShapeStateRecord( + graph_filename=shape["graph_filename"], + batch_size=int(shape["batch_size"]), + capture_index=int(shape["capture_index"]), + plan_schema=shape["plan_schema"], + plan_fingerprint=tuple(int(value) for value in shape["plan_fingerprint"]), + tensor_slots=tuple(TensorSlotRecord(**slot) for slot in shape["tensor_slots"]), + operand_slots=tuple(OperandSlotRecord(**slot) for slot in shape["operand_slots"]), + ) + for shape in raw["shapes"] + ) + manifest = ShapeStateManifest( + backend=raw["backend"], + rank=int(raw["rank"]), + world_size=int(raw["world_size"]), + capture_order=tuple(int(value) for value in raw["capture_order"]), + shapes=shapes, + ) + validate_shape_state_manifest(manifest) + if payload.get("schema_fingerprint") != _semantic_fingerprint(manifest): + raise RuntimeError("SGLang shape-state schema fingerprint mismatch") + return manifest diff --git a/python/foundry/integration/sglang/symm_mem_graph.py b/python/foundry/integration/sglang/symm_mem_graph.py new file mode 100644 index 00000000..14901500 --- /dev/null +++ b/python/foundry/integration/sglang/symm_mem_graph.py @@ -0,0 +1,676 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the Foundry project +"""Relocate process-local PyTorch symmetric-memory graph operands.""" + +from __future__ import annotations + +import json +import re +import shutil +from collections.abc import Callable +from dataclasses import dataclass +from pathlib import Path +from typing import NamedTuple + +from foundry.integration.sglang.state_manifest import ( + OperandSlotRecord, + ShapeStateManifest, + validate_shape_state_manifest, +) + +FULL_GRAPH_DIR = "symmetric_full_graphs" +RELOCATED_GRAPH_DIR = "symmetric_relocated_graphs" +STATE_FILENAME = "symmetric_memory_state.json" +BACKEND_FILENAME = "sglang_graph_backend.json" +STATE_VERSION = 1 +_TWO_SHOT_KERNEL = "two_shot_all_reduce_kernel_inplace" +_TWO_SHOT_SPECIALIZATION = ( + "two_shot_all_reduce_kernel_inplaceIN3c108BFloat16ELi16ELi2EEEvPPT_mmPPjmm" +) +_GRAPH_NAME_RE = re.compile(r"^graph_(?P\d+)_FULL_t(?P\d+)_r(?P=batch)_UX_pcN\.json$") +_PARAM_OFFSETS = (0, 8, 16, 24, 32, 40) +_COPY_ZERO_FIELDS = ( + "srcXInBytes", + "srcY", + "srcZ", + "srcLOD", + "srcPitch", + "srcHeight", + "dstXInBytes", + "dstY", + "dstZ", + "dstLOD", + "dstPitch", + "dstHeight", +) + + +class SymmetricMemoryGraphState(NamedTuple): + buffer: int + buffer_ptrs_dev: int + signal_pad_ptrs_dev: int + multicast_ptr: int + buffer_numel: int + element_size: int + dtype: str + rank: int + world_size: int + + +@dataclass(frozen=True) +class GraphOperandInventory: + graph_filename: str + batch_size: int + capture_index: int + operand_slots: tuple[OperandSlotRecord, ...] + + +def run_graph_warmups( + *, + synchronize: Callable[[], None], + barrier: Callable[[], None], + forward: Callable[[], object], + post_warmup: Callable[[], None] | None, +) -> None: + for _ in range(2): + synchronize() + barrier() + forward() + if post_warmup is not None: + post_warmup() + synchronize() + barrier() + + +def preserve_symmetric_graphs( + workspace: str | Path, + graph_filenames: list[str], + state: SymmetricMemoryGraphState, + *, + allow_multi_shape: bool = False, +) -> tuple[GraphOperandInventory, ...]: + workspace = Path(workspace) + _validate_state(state) + parsed = _validate_graph_files( + graph_filenames, + allow_multi_shape=allow_multi_shape, + ) + full_graph_dir = workspace / FULL_GRAPH_DIR + shutil.rmtree(full_graph_dir, ignore_errors=True) + full_graph_dir.mkdir(exist_ok=True) + inventories = [] + for index, batch_size, filename in parsed: + graph_data = json.loads((workspace / filename).read_text()) + kernel_count = _relocate_graph(graph_data, state, state) + if kernel_count == 0: + raise RuntimeError(f"SGLang symmetric-memory graph has no two-shot kernels: {filename}") + inventories.append( + GraphOperandInventory( + graph_filename=filename, + batch_size=batch_size, + capture_index=index, + operand_slots=_communication_operand_records(filename, graph_data), + ) + ) + shutil.copy2(workspace / filename, full_graph_dir / filename) + metadata = { + "version": STATE_VERSION, + "backend": "torch_symmetric_memory", + "state": state._asdict(), + } + (workspace / STATE_FILENAME).write_text(json.dumps(metadata, sort_keys=True)) + return tuple(inventories) + + +def clear_symmetric_graph_state(workspace: str | Path) -> None: + workspace = Path(workspace) + (workspace / STATE_FILENAME).unlink(missing_ok=True) + shutil.rmtree(workspace / FULL_GRAPH_DIR, ignore_errors=True) + shutil.rmtree(workspace / RELOCATED_GRAPH_DIR, ignore_errors=True) + + +def write_graph_backend_metadata( + workspace: str | Path, + symmetric_memory_enabled: bool, +) -> None: + workspace = Path(workspace) + backend = "torch_symmetric_memory" if symmetric_memory_enabled else "nccl" + if symmetric_memory_enabled: + _load_symmetric_state(workspace) + metadata = { + "version": STATE_VERSION, + "backend": backend, + } + (workspace / BACKEND_FILENAME).write_text(json.dumps(metadata, sort_keys=True)) + + +def validate_graph_backend( + workspace: str | Path, + symmetric_memory_enabled: bool, +) -> None: + workspace = Path(workspace) + try: + metadata = json.loads((workspace / BACKEND_FILENAME).read_text()) + except (OSError, json.JSONDecodeError) as error: + raise RuntimeError("SGLang graph archive has no valid backend metadata") from error + if metadata.get("version") != STATE_VERSION: + raise RuntimeError(f"Unsupported SGLang graph backend version: {metadata.get('version')!r}") + archive_backend = metadata.get("backend") + if archive_backend not in {"nccl", "torch_symmetric_memory"}: + raise RuntimeError(f"Invalid SGLang graph backend: {archive_backend!r}") + expected_backend = "torch_symmetric_memory" if symmetric_memory_enabled else "nccl" + if archive_backend != expected_backend: + raise RuntimeError( + "SGLang communication backend changed between SAVE and LOAD: " + f"archive={archive_backend}, load={expected_backend}" + ) + if archive_backend == "torch_symmetric_memory": + _load_symmetric_state(workspace) + elif (workspace / STATE_FILENAME).exists(): + raise RuntimeError("NCCL SGLang graph archive contains symmetric-memory state") + + +def _load_symmetric_state(workspace: str | Path) -> SymmetricMemoryGraphState: + workspace = Path(workspace) + try: + metadata = json.loads((workspace / STATE_FILENAME).read_text()) + except (OSError, json.JSONDecodeError) as error: + raise RuntimeError("SGLang graph archive has no valid symmetric-memory state") from error + if metadata.get("version") != STATE_VERSION: + raise RuntimeError( + f"Unsupported SGLang symmetric-memory graph state version: {metadata.get('version')!r}" + ) + if metadata.get("backend") != "torch_symmetric_memory": + raise RuntimeError( + f"Invalid SGLang symmetric-memory graph backend: {metadata.get('backend')!r}" + ) + try: + return SymmetricMemoryGraphState(**metadata["state"]) + except (KeyError, TypeError) as error: + raise RuntimeError("Invalid SGLang symmetric-memory graph state") from error + + +def _validate_state_compatibility( + saved_state: SymmetricMemoryGraphState, + live_state: SymmetricMemoryGraphState, +) -> None: + _validate_state(saved_state) + _validate_state(live_state) + if saved_state.rank != live_state.rank: + raise RuntimeError( + "SGLang symmetric-memory rank changed between SAVE and LOAD: " + f"{saved_state.rank} != {live_state.rank}" + ) + if saved_state.world_size != live_state.world_size: + raise RuntimeError( + "SGLang symmetric-memory world size changed between SAVE and LOAD: " + f"{saved_state.world_size} != {live_state.world_size}" + ) + if saved_state.dtype != live_state.dtype: + raise RuntimeError( + "SGLang symmetric-memory dtype changed between SAVE and LOAD: " + f"{saved_state.dtype} != {live_state.dtype}" + ) + if saved_state.element_size != live_state.element_size: + raise RuntimeError( + "SGLang symmetric-memory element size changed between SAVE and LOAD: " + f"{saved_state.element_size} != {live_state.element_size}" + ) + if live_state.buffer_numel < saved_state.buffer_numel: + raise RuntimeError( + "SGLang symmetric-memory LOAD buffer is smaller than SAVE: " + f"{live_state.buffer_numel} < {saved_state.buffer_numel}" + ) + + +def _validate_state(state: SymmetricMemoryGraphState) -> None: + if state.world_size != 2 or state.rank not in (0, 1): + raise RuntimeError("Foundry SGLang symmetric-memory replay currently requires TP=2") + if state.dtype != "torch.bfloat16" or state.element_size != 2: + raise RuntimeError("Foundry SGLang symmetric-memory replay requires bfloat16 buffers") + if state.buffer_numel <= 0: + raise RuntimeError("SGLang symmetric-memory buffer is empty") + if not all( + ( + state.buffer, + state.buffer_ptrs_dev, + state.signal_pad_ptrs_dev, + state.multicast_ptr, + ) + ): + raise RuntimeError("SGLang symmetric-memory state contains a null pointer") + + +def relocate_symmetric_graphs( + workspace: str | Path, + graph_filenames: list[str], + live_state: SymmetricMemoryGraphState, + *, + manifest: ShapeStateManifest | None = None, + allow_multi_shape: bool = False, + output_dir: str | Path | None = None, +) -> list[str]: + workspace = Path(workspace) + parsed = _validate_graph_files( + graph_filenames, + allow_multi_shape=allow_multi_shape, + ) + if manifest is not None: + validate_shape_state_manifest(manifest) + manifest_names = tuple(shape.graph_filename for shape in manifest.shapes) + parsed_names = tuple(filename for _index, _batch, filename in parsed) + if manifest_names != parsed_names: + raise RuntimeError( + f"SGLang shape manifest order mismatch: {manifest_names} != {parsed_names}" + ) + saved_state = _load_symmetric_state(workspace) + _validate_state_compatibility(saved_state, live_state) + + relocated_dir = Path(output_dir) if output_dir is not None else workspace / RELOCATED_GRAPH_DIR + relocated_dir.mkdir(parents=True, exist_ok=True) + relocated_paths = [] + for _index, _batch, filename in parsed: + source_path = workspace / FULL_GRAPH_DIR / filename + graph_data = json.loads(source_path.read_text()) + kernel_count = _relocate_graph(graph_data, saved_state, live_state) + if kernel_count == 0: + raise RuntimeError(f"SGLang symmetric-memory graph has no two-shot kernels: {filename}") + relocated_path = relocated_dir / filename + relocated_path.write_text(json.dumps(graph_data)) + relocated_paths.append(str(relocated_path)) + return relocated_paths + + +def _validate_graph_files( + graph_filenames: list[str], + *, + allow_multi_shape: bool, +) -> tuple[tuple[int, int, str], ...]: + parsed = [] + for filename in graph_filenames: + match = _GRAPH_NAME_RE.fullmatch(filename) + if match is None: + raise RuntimeError(f"Invalid SGLang graph filename: {filename}") + parsed.append( + ( + int(match.group("index")), + int(match.group("batch")), + filename, + ) + ) + parsed.sort() + indices = tuple(index for index, _batch, _filename in parsed) + batches = tuple(batch for _index, batch, _filename in parsed) + if indices != tuple(range(len(indices))): + raise RuntimeError(f"SGLang graph indices are not contiguous: {indices}") + if len(batches) != len(set(batches)): + raise RuntimeError(f"SGLang graph batches are duplicated: {batches}") + if not allow_multi_shape and batches != (1,): + raise RuntimeError( + "Foundry SGLang torch symmetric memory currently supports exactly one " + "batch-1 CUDA graph" + ) + if allow_multi_shape and batches != tuple(sorted(batches, reverse=True)): + raise RuntimeError(f"SGLang graph capture order is not largest-to-smallest: {batches}") + return tuple(parsed) + + +def _communication_operand_records( + graph_filename: str, + graph_data: dict, +) -> tuple[OperandSlotRecord, ...]: + records = [] + for node in graph_data["nodes"]: + if node["type"] != "KernelNode": + continue + params = node["params"] + if _TWO_SHOT_SPECIALIZATION not in params["function_name"]: + continue + records.extend( + ( + OperandSlotRecord( + name="symm.buffer_ptrs_dev", + kernel_abi="torch.symm_mem.two_shot.bf16.align16.tp2", + owner_slot="communicator", + node_id=int(node["id"]), + source="kernelParams", + parameter_index=0, + cuda_parameter_offset=int(params["kernelParams"][0]["offset"]), + value_byte_offset=0, + ctype="BFloat16**", + owner_relative_offset=0, + span_bytes=16, + saved_value=_decode_param(params["kernelParams"][0]), + ), + OperandSlotRecord( + name="symm.signal_pad_ptrs_dev", + kernel_abi="torch.symm_mem.two_shot.bf16.align16.tp2", + owner_slot="communicator", + node_id=int(node["id"]), + source="kernelParams", + parameter_index=3, + cuda_parameter_offset=int(params["kernelParams"][3]["offset"]), + value_byte_offset=0, + ctype="uint32_t**", + owner_relative_offset=0, + span_bytes=16, + saved_value=_decode_param(params["kernelParams"][3]), + ), + ) + ) + if not records: + raise RuntimeError( + f"SGLang symmetric-memory graph has no two-shot kernels: {graph_filename}" + ) + return tuple(records) + + +def _relocate_graph( + graph_data: dict, + saved_state: SymmetricMemoryGraphState, + live_state: SymmetricMemoryGraphState, +) -> int: + nodes = graph_data["nodes"] + node_ids = [int(node["id"]) for node in nodes] + if len(node_ids) != len(set(node_ids)): + raise RuntimeError("SGLang symmetric-memory graph has duplicate node IDs") + dependencies = { + (int(dependency["from"]), int(dependency["to"])) + for dependency in graph_data["dependencies"] + } + operand_values = { + saved_state.buffer, + saved_state.buffer_ptrs_dev, + saved_state.signal_pad_ptrs_dev, + saved_state.multicast_ptr, + live_state.buffer, + live_state.buffer_ptrs_dev, + live_state.signal_pad_ptrs_dev, + live_state.multicast_ptr, + } + operand_ranges = tuple( + operand_range + for state in (saved_state, live_state) + for operand_range in ( + ( + state.buffer, + state.buffer + state.buffer_numel * state.element_size, + ), + ( + state.buffer_ptrs_dev, + state.buffer_ptrs_dev + state.world_size * 8, + ), + ( + state.signal_pad_ptrs_dev, + state.signal_pad_ptrs_dev + state.world_size * 8, + ), + ( + state.multicast_ptr, + state.multicast_ptr + state.buffer_numel * state.element_size, + ), + ) + ) + allowed_operands: set[tuple[int, str, int, int]] = set() + used_copy_ids: set[int] = set() + kernel_count = 0 + for position, node in enumerate(nodes): + params = node["params"] + if node["type"] != "KernelNode" or _TWO_SHOT_KERNEL not in params["function_name"]: + continue + if _TWO_SHOT_SPECIALIZATION not in params["function_name"]: + raise RuntimeError("SGLang symmetric-memory two-shot kernel specialization changed") + + kernel_count += 1 + kernel_id = int(node["id"]) + encoded_params = params["kernelParams"] + _validate_kernel_abi(encoded_params) + allowed_operands.update( + { + (kernel_id, "kernelParams", 0, 0), + (kernel_id, "kernelParams", 3, 0), + } + ) + numel = _decode_param(encoded_params[2]) + _relocate_pointer_param( + encoded_params, + index=0, + expected=saved_state.buffer_ptrs_dev, + replacement=live_state.buffer_ptrs_dev, + ) + _relocate_pointer_param( + encoded_params, + index=3, + expected=saved_state.signal_pad_ptrs_dev, + replacement=live_state.signal_pad_ptrs_dev, + ) + _validate_scalar_param(encoded_params, 4, saved_state.rank, "rank") + _validate_scalar_param( + encoded_params, + 5, + saved_state.world_size, + "world size", + ) + if _decode_param(encoded_params[1]) != 0: + raise RuntimeError("SGLang symmetric-memory input offset is not zero") + if numel <= 0: + raise RuntimeError("SGLang symmetric-memory input is empty") + if numel > saved_state.buffer_numel or numel > live_state.buffer_numel: + raise RuntimeError( + "SGLang symmetric-memory collective exceeds the communication buffer" + ) + if position == 0 or position + 1 >= len(nodes): + raise RuntimeError("SGLang symmetric-memory kernel has no copy pair") + ingress = nodes[position - 1] + egress = nodes[position + 1] + _relocate_copy( + ingress, + field="dstDevice", + other_field="srcDevice", + expected=saved_state.buffer, + replacement=live_state.buffer, + width=numel * saved_state.element_size, + ) + _relocate_copy( + egress, + field="srcDevice", + other_field="dstDevice", + expected=saved_state.buffer, + replacement=live_state.buffer, + width=numel * saved_state.element_size, + ) + ingress_id = int(ingress["id"]) + egress_id = int(egress["id"]) + allowed_operands.update( + { + (ingress_id, "dstDevice", -1, 0), + (egress_id, "srcDevice", -1, 0), + } + ) + if ingress_id in used_copy_ids or egress_id in used_copy_ids: + raise RuntimeError("SGLang symmetric-memory copy node is reused") + used_copy_ids.update((ingress_id, egress_id)) + if (ingress_id, kernel_id) not in dependencies or ( + kernel_id, + egress_id, + ) not in dependencies: + raise RuntimeError("SGLang symmetric-memory copy/kernel dependencies changed") + + _validate_operand_inventory( + nodes, + operand_values, + operand_ranges, + allowed_operands, + ) + if len(used_copy_ids) != kernel_count * 2: + raise RuntimeError("SGLang symmetric-memory graph copy/kernel count mismatch") + return kernel_count + + +def _validate_operand_inventory( + nodes: list[dict], + operand_values: set[int], + operand_ranges: tuple[tuple[int, int], ...], + allowed_operands: set[tuple[int, str, int, int]], +) -> None: + for node in nodes: + node_id = int(node["id"]) + params = node["params"] + if node["type"] == "MemcpyNode": + for field in ("srcDevice", "dstDevice"): + value = int(params[field]) + location = (node_id, field, -1, 0) + if ( + _is_external_operand(value, operand_values, operand_ranges) + and location not in allowed_operands + ): + raise RuntimeError( + "SGLang graph contains an unrecognized symmetric-memory " + f"operand in memcpy node {node_id}" + ) + continue + if node["type"] == "MemsetNode": + if _is_external_operand( + int(params["dst"]), + operand_values, + operand_ranges, + ): + raise RuntimeError( + "SGLang graph contains an unrecognized symmetric-memory " + f"operand in memset node {node_id}" + ) + continue + if node["type"] == "KernelNode": + for param in params["kernelParams"]: + encoded = param.get("value_hex") + if not encoded: + continue + raw = bytes.fromhex(encoded) + for offset, value in _iter_qwords(raw): + location = ( + node_id, + "kernelParams", + int(param["index"]), + offset, + ) + if ( + _is_external_operand( + value, + operand_values, + operand_ranges, + ) + and location not in allowed_operands + ): + raise RuntimeError( + "SGLang graph contains an unrecognized symmetric-memory " + f"operand in kernel node {node_id}" + ) + extra = bytes.fromhex(params.get("extra_argBuffer_hex") or "") + for offset, value in _iter_qwords(extra): + if _is_external_operand(value, operand_values, operand_ranges): + raise RuntimeError( + "SGLang graph contains an unrecognized symmetric-memory " + f"argument-buffer operand in kernel node {node_id}" + ) + access_policy_ptr = int( + params.get("kernel_node_attrs", {}).get( + "accessPolicyWindowBasePtr", + 0, + ) + ) + if _is_external_operand( + access_policy_ptr, + operand_values, + operand_ranges, + ): + raise RuntimeError( + "SGLang graph contains an unrecognized symmetric-memory " + f"access-policy operand in kernel node {node_id}" + ) + + +def _is_external_operand( + value: int, + operand_values: set[int], + operand_ranges: tuple[tuple[int, int], ...], +) -> bool: + return value in operand_values or any(start <= value < end for start, end in operand_ranges) + + +def _iter_qwords(raw: bytes): + for offset in range(0, len(raw) - 7, 8): + yield offset, int.from_bytes(raw[offset : offset + 8], "little") + + +def _decode_param(param: dict) -> int: + raw = bytes.fromhex(param["value_hex"]) + if len(raw) != 8: + raise RuntimeError(f"Expected an 8-byte CUDA graph parameter, got {len(raw)} bytes") + return int.from_bytes(raw, "little") + + +def _relocate_pointer_param( + params: list[dict], + *, + index: int, + expected: int, + replacement: int, +) -> None: + actual = _decode_param(params[index]) + if actual != expected: + raise RuntimeError( + f"Symmetric-memory kernel parameter {index} changed: 0x{actual:x} != 0x{expected:x}" + ) + params[index]["value_hex"] = replacement.to_bytes(8, "little").hex() + + +def _validate_scalar_param( + params: list[dict], + index: int, + expected: int, + label: str, +) -> None: + actual = _decode_param(params[index]) + if actual != expected: + raise RuntimeError(f"Symmetric-memory graph {label} changed: {actual} != {expected}") + + +def _validate_kernel_abi(params: list[dict]) -> None: + if len(params) != len(_PARAM_OFFSETS): + raise RuntimeError( + f"SGLang symmetric-memory kernel ABI changed: expected 6 parameters, got {len(params)}" + ) + for position, (param, expected_offset) in enumerate(zip(params, _PARAM_OFFSETS, strict=True)): + if int(param["index"]) != position: + raise RuntimeError( + "SGLang symmetric-memory kernel parameters are reordered or duplicated" + ) + if int(param["offset"]) != expected_offset or int(param["size"]) != 8: + raise RuntimeError( + f"SGLang symmetric-memory kernel parameter {position} layout changed" + ) + + +def _relocate_copy( + node: dict, + *, + field: str, + other_field: str, + expected: int, + replacement: int, + width: int, +) -> None: + if node["type"] != "MemcpyNode": + raise RuntimeError("SGLang symmetric-memory kernel copy pair changed") + params = node["params"] + if int(params[field]) != expected or int(params[other_field]) == expected: + raise RuntimeError("SGLang symmetric-memory copy direction changed") + if ( + int(params["srcMemoryType"]) != 2 + or int(params["dstMemoryType"]) != 2 + or int(params["WidthInBytes"]) != width + or int(params["Height"]) != 1 + or int(params["Depth"]) != 1 + or any(int(params[name]) != 0 for name in _COPY_ZERO_FIELDS) + ): + raise RuntimeError("SGLang symmetric-memory copy layout changed") + params[field] = replacement diff --git a/python/foundry/integration/vllm/hooks.py b/python/foundry/integration/vllm/hooks.py index f7768630..b330663d 100644 --- a/python/foundry/integration/vllm/hooks.py +++ b/python/foundry/integration/vllm/hooks.py @@ -695,12 +695,22 @@ def _patch_deepep() -> None: def patched(self, *args, **kwargs): out = orig(self, *args, **kwargs) if isinstance(out, dict) and get_graph_extension_mode() != CUDAGraphExtensionMode.NONE: - # Force VMM-fabric buffers - # so the foundry cuMem hook tracks them, and disable the - # NVLink buffer (LL-mode uses RDMA) to avoid fabric-cuMemCreate - # collisions with preallocated region. - out["use_fabric"] = True - out["num_nvl_bytes"] = 0 + if os.environ.get("FOUNDRY_DEEPEP_NVL_IPC", "0") == "1": + # Keep upstream's NVLink buffer (legacy cudaIpc sharing). The + # hook's VMM-IPC layer transports the fds via SCM_RIGHTS and + # maps peers at relocated VAs, which DeepEP absorbs through + # its buffer_ptrs_gpu device table - peer VAs are never baked + # into captured graphs. use_fabric stays off: fabric + # cuMemCreate needs IMEX, and exercising the IPC path is the + # point of this mode. + out["use_fabric"] = False + else: + # Default: force VMM-fabric buffers + # so the foundry cuMem hook tracks them, and disable the + # NVLink buffer (LL-mode uses RDMA) to avoid fabric-cuMemCreate + # collisions with preallocated region. + out["use_fabric"] = True + out["num_nvl_bytes"] = 0 return out cls._make_all2all_kwargs = patched diff --git a/recipe/experimental/README.md b/recipe/experimental/README.md new file mode 100644 index 00000000..2ea6a305 --- /dev/null +++ b/recipe/experimental/README.md @@ -0,0 +1,269 @@ +# Foundry recipe โ€” experimental multi-GPU IPC paths + +End-to-end SAVE / LOAD recipe for **Qwen3-30B-A3B** expert-parallel where DeepEP's +intranode NVLink buffer stays on the **legacy CUDA-IPC** path +(`cudaIpcGetMemHandle` / `cudaIpcOpenMemHandle`) under foundry, instead of the +default fabric/NVSHMEM-only path used by `recipe/vllm/serve_qwen3-30ba3b_ep.sh`. +This directory also contains the validated SGLang **dense tensor-parallel** +recipe, which captures NCCL P2P/IPC all-reduce in every decode graph. + +This exercises foundry's **VMM-IPC translation layer**: DeepEP's NVL buffer is a +foundry VMM (`cuMemCreate`) allocation that legacy IPC can't share, so the hook +exports it as a POSIX fd and transports that fd to the peer rank via +`SCM_RIGHTS` over a per-process unix socket, then maps it into the importer's +own VA range. This is the path that lets foundry SAVE/LOAD work on machines +**without fabric / IMEX**, where DeepEP would otherwise fail at `Buffer.sync` +with `cuMemImportFromShareableHandle ... error 999`. + +``` +recipe/experimental/ +โ”œโ”€โ”€ README.md # this file +โ”œโ”€โ”€ foundry_save.toml # SAVE config (workspace_root = "foundry_archive_ipc") +โ”œโ”€โ”€ foundry_load.toml # LOAD config (same workspace_root) +โ”œโ”€โ”€ serve_qwen3-30ba3b_ipc_ep.sh # Qwen3-30B-A3B (MoE) vLLM EP, NVL/IPC +โ”œโ”€โ”€ sglang_foundry_save.toml # SGLang EP SAVE config +โ”œโ”€โ”€ sglang_foundry_load.toml # SGLang EP LOAD config +โ”œโ”€โ”€ serve_qwen3-30ba3bfp8_sglang_ipc_ep.sh # SGLang EP, DeepEP NVL/IPC +โ”œโ”€โ”€ sglang_foundry_tp_save.toml # SGLang dense-TP SAVE config +โ”œโ”€โ”€ sglang_foundry_tp_load.toml # SGLang dense-TP LOAD config +โ””โ”€โ”€ serve_qwen3-8b_sglang_tp.sh # SGLang dense tensor parallel (NCCL all-reduce) +``` + +It is the standard `recipe/vllm` EP recipe plus one switch โ€” `FOUNDRY_DEEPEP_NVL_IPC=1` +โ€” so read [`../vllm/README.md`](../vllm/README.md) first for installation, the +two-pass SAVE workflow, the archive layout, and the shared EP flags. Only the +IPC-specific deltas are documented here. + +## Required code + +The IPC path needs only **Foundry** changes beyond the standard install โ€” no +vLLM edit: + +- **Foundry** โ€” `foundry/csrc/hook.cpp`: the SCM_RIGHTS VMM-IPC fd transport + + whole-chunk peer mapping (handles LOAD-mode buffers carved from the + preallocated chunk). **C++ change โ†’ rebuild required:** + `uv pip install -e . --no-build-isolation` in `foundry/`. +- **Foundry** โ€” `foundry/python/foundry/integration/vllm/hooks.py`: the + `FOUNDRY_DEEPEP_NVL_IPC` env knob in `_patch_deepep` (Python, no rebuild). +- **Foundry** โ€” `foundry/python/foundry/integration/sglang/hooks.py`: the + matching SGLang `deep_ep.Buffer` policy patch (Python, no rebuild). + +### Run all phases from one consistent path (no vLLM edit needed) + +vLLM folds the foundry TOML path (`graph_extension_config_path`) into its +torch.compile cache key even though the path never affects codegen. If SAVE +pass 1 and pass 2 see that path with *different spellings* โ€” e.g. two mount +aliases of the same dir (`/data/...` vs `/home/...`), or you move the TOML +between passes โ€” pass 2 misses pass 1's warm compile cache and inductor +recompiles **inside** the cuda-graph capture window, where its combo-kernel +benchmark does an illegal `torch.randn` โ†’ `cudaErrorStreamCaptureInvalidated`. + +The fix is operational, not code: **invoke the script the same way for every +phase** (same shell / same `cd`, or always an absolute path), so SAVE pass 1, +SAVE pass 2, and LOAD all pass the identical path string โ†’ identical cache hash +โ†’ pass 2 reuses pass 1's compiled kernels. Run from the workspace root, e.g.: + +```bash +cd # one canonical path for all three phases +bash foundry/recipe/experimental/serve_qwen3-30ba3b_ipc_ep.sh 2 --save +bash foundry/recipe/experimental/serve_qwen3-30ba3b_ipc_ep.sh 2 --save +bash foundry/recipe/experimental/serve_qwen3-30ba3b_ipc_ep.sh 2 --load +``` + +## Workflow + +Same two-pass SAVE โ†’ LOAD as the base recipe; `` is the first arg: + +```bash +# 0. Fresh start (distinct workspace from the base recipe) +rm -rf foundry_archive_ipc + +# 1. SAVE pass 1 โ€” memory profile + capture +bash serve_qwen3-30ba3b_ipc_ep.sh 2 --save +# wait for "Application startup complete", then Ctrl-C + +# 2. SAVE pass 2 โ€” deterministic re-capture +bash serve_qwen3-30ba3b_ipc_ep.sh 2 --save +# wait for "Application startup complete", then Ctrl-C + +# 3. LOAD โ€” preallocate, re-import IPC buffers, replay graphs +bash serve_qwen3-30ba3b_ipc_ep.sh 2 --load +# leave running + +# 4. Query (separate shell) +curl -s http://0.0.0.0:12000/v1/completions \ + -H 'Content-Type: application/json' \ + -d '{"model":"Qwen/Qwen3-30B-A3B","prompt":"The capital of France is","max_tokens":12,"temperature":0}' +``` + +Uncomment `nvshmem_host_path` in both TOMLs first (EP still needs NVSHMEM for the +DeepEP RDMA buffer; the IPC path only changes the NVL buffer). + +## SGLang variant + +The SGLang variant is +`serve_qwen3-30ba3bfp8_sglang_ipc_ep.sh`. It uses the same Foundry VMM-IPC core +but the SGLang integration wraps `deep_ep.Buffer.__init__` rather than vLLM's +all-to-all kwargs builder. Use the SGLang-specific TOMLs and run the same +`--save`/`--load` settings for every phase: + +```bash +rm -rf foundry_archive_sglang_ipc +bash serve_qwen3-30ba3bfp8_sglang_ipc_ep.sh 2 --save +bash serve_qwen3-30ba3bfp8_sglang_ipc_ep.sh 2 --load +``` + +The SGLang variant requires DeepEP `29d31c0` or newer. This is the vLLM +installer's validated revision and is API-compatible with SGLang's older +`Buffer(group, num_nvl_bytes, num_rdma_bytes, ...)` call shape while adding +`use_fabric`. NVSHMEM remains required for DeepEP's RDMA symmetric heap; only +the NVL buffer uses CUDA IPC. Keep the model, EP size, low-latency token cap, +chunked-prefill size, and IPC environment settings identical between SAVE and +LOAD. + +Validation is currently BF16 on 2ร—H100. Upstream +[PR #3](https://github.com/foundry-org/foundry/pull/3) still calls FP8 +experimental, and the GB300/aarch64 FP8 plus RDC relinking paths reported in +[issue #1](https://github.com/foundry-org/foundry/issues/1) are outside this +recipe's verified scope. + +## What `FOUNDRY_DEEPEP_NVL_IPC=1` does + +`_patch_deepep` (`foundry/python/foundry/integration/vllm/hooks.py`) normally +forces `use_fabric=True, num_nvl_bytes=0` so the only cross-GPU buffer is the +NVSHMEM symmetric heap. With `FOUNDRY_DEEPEP_NVL_IPC=1` it instead keeps +upstream's nonzero `num_nvl_bytes` with `use_fabric=False`, so the DeepEP Buffer +allocates the NVLink buffer via `cudaMalloc` + `cudaIpcGetMemHandle` on **both** +SAVE and LOAD. The hook then: + +- **SAVE**: exports each VMM-backed NVL buffer as a POSIX fd, served to peers + over `\0foundry-vmm-ipc.` via `SCM_RIGHTS` (same-uid `SO_PEERCRED` check, + per-process token vs PID reuse). +- **LOAD**: buffers are carved from the preallocated chunk (no individual + handle), so the hook exports the **whole chunk** fd + offset; the peer maps the + entire chunk once and returns an interior pointer. +- Peer mappings land at a **relocated** VA (logged `[HOOK] INFO: VMM-IPC import + relocated`) โ€” correct, because DeepEP resolves peers through its device-side + `buffer_ptrs_gpu` table (refreshed by `Buffer.sync`), never through addresses + baked into captured graphs. + +A healthy run logs two `VMM-IPC import relocated` lines per phase (one per peer) +and **no** `error 999`. + +## SGLang dense tensor parallel + +`serve_qwen3-8b_sglang_tp.sh` runs a **dense** model tensor-parallel across +`` GPUs (`--tp-size N`, no DP-attention, no expert parallel) โ€” the +first recipe here that captures a **cross-rank collective into the graph +itself**. The DP recipe replicates the whole model per rank (no in-graph +collective) and the EP recipe puts attention on DP-attention and the MoE on +DeepEP's all-to-all; neither exercises a per-layer TP all-reduce. + +```bash +rm -rf foundry_archive_sglang_tp +CUDA_VISIBLE_DEVICES=0,1 bash serve_qwen3-8b_sglang_tp.sh 2 --save +CUDA_VISIBLE_DEVICES=0,1 bash serve_qwen3-8b_sglang_tp.sh 2 --load +curl -s http://0.0.0.0:12000/v1/completions \ + -H 'Content-Type: application/json' \ + -d '{"model":"Qwen/Qwen3-8B","prompt":"The capital of France is","max_tokens":12,"temperature":0}' +``` + +How the TP all-reduce survives SAVE/LOAD: + +- `--disable-custom-all-reduce` routes the per-layer all-reduce through **NCCL** + (not SGLang's custom one-shot/two-shot kernel). NCCL reports P2P/IPC on all 24 + channels in the validated run. +- `NCCL_CUMEM_ENABLE=0` / `NCCL_NVLS_ENABLE=0` keep NCCL off the CUMEM P2P and + NVLS multicast fast paths (whose driver-capability flags the foundry VMM + region doesn't carry) and on the **legacy CUDA-IPC** intra-node transport. +- NCCL communicator initialization runs before the 1 GiB scratch boundary on + both SAVE and LOAD. LOAD therefore constructs fresh P2P connection state + before restoring the captured graphs, while Foundry reproduces the same + rank-local allocation layout. + +No NVSHMEM is involved (dense, no DeepEP), so the TP TOMLs leave +`nvshmem_host_path` unset. Keep `--tp-size`, `--cuda-graph-max-bs`, and the NCCL +environment identical between SAVE and LOAD so the captured graphs and the NCCL +buffer trajectory match. The script also pins `--random-seed 42` by default +(`SGL_RANDOM_SEED` overrides it), making SAVE-pass metadata directly +comparable. Custom all-reduce (SGLang's own IPC kernel) uses the same `cuIpc` +primitives and is expected to work over this bridge too, but is left disabled +here until GPU-validated. + +## Validation status + +SAVE โ†’ SAVE โ†’ LOAD โ†’ query verified on Qwen3-30B-A3B EP=2 (2ร— H200, no IMEX): +per-rank `final_alloc_offset` identical across passes, LOAD replays at the saved +offset, query returns coherent completions. LOAD reaches a serving server in +~27 s; the IPC import itself is sub-second (inside the ~7.5 s `load_model`) and +has no steady-state serving cost โ€” peer addresses are resolved once at init via +the device pointer table, not per token. + +The **dense TP** recipe (`serve_qwen3-8b_sglang_tp.sh`) is validated with +Qwen3-8B at TP=2 on 2ร—H100. Two SAVE passes each wrote 36 graphs per rank at +`final_alloc_offset=69736595456`; a fresh LOAD restored all 36 graphs on both +ranks at that offset without recapture or CUDA/NCCL errors. Four +temperature-0 responses matched the non-Foundry TP baseline byte for byte. + +The checked-in proof runs the complete baseline โ†’ SAVE โ†’ SAVE โ†’ LOAD matrix: + +```bash +modal run --env tests/modal_sglang_tp.py +``` + +Set `SGLANG_ENABLE_TORCH_SYMM_MEM=1` to run the validated SGLang-main +torch symmetric-memory mode. It is restricted to TP=2 and defaults to one +batch-1 graph; the script sets `SGLANG_CUDA_GRAPH_MAX_BS=1`, and larger batches +fall back to eager execution. Foundry preserves full graphs before manifest +compaction, validates the two-shot ABI, records process-local communication +operands, retains explicit per-shape FlashInfer replay state, and relocates +those operands during fresh LOAD. + +Multi-shape capture is guarded by +`FOUNDRY_SGLANG_SYMM_MULTI_SHAPE=1`; configure the ascending graph-key schedule +with `FOUNDRY_SGLANG_SYMM_GRAPH_BATCHES`. The retained 2ร—H100 milestones are: + +| Graph keys | Retained run ID | Final offset per rank | Delta vs batch-1 offset `68555898880` | +|---|---|---:|---:| +| `1,8` | `sglang-symm-1-8-742970b` | `68597841920` | `41943040` bytes (40 MiB) | +| `1,8,32` | `sglang-symm-1-8-32-742970b` | `68664950784` | `109051904` bytes (104 MiB) | + +Both matrices passed every validation check. SAVE/SAVE2 fingerprints and +offsets matched, LOAD restored every graph on both ranks without recapture, +configured replay keys were observed exactly, and deterministic outputs matched +the baseline. The isolation gate also observed explicit uncaptured eager decode +batches 9 and 33 with `cuda graph: False`; subsequent graph-backed outputs still +matched. No runtime errors were reported. This remains a pinned, opt-in result, +not the default recipe configuration. + +This is a pinned configuration result, not general NCCL TP support. Upstream +[Discussion #5](https://github.com/orgs/foundry-org/discussions/5) calls out +NCCL initialization nondeterminism and motivated the symmetric-memory direction. +General TP and multi-shape support remain tracked in +[issue #6](https://github.com/foundry-org/foundry/issues/6). + +## Not supported: Qwen3.5-122B-A10B + +The dense-TP recipe is validated for dense transformer models. The hybrid +Mamba/MoE `Qwen/Qwen3.5-122B-A10B` configuration remains unsupported: + +- On the legacy SGLang fallback, `FOUNDRY_SGLANG_SAVE_WARMUP=1` runs a + cursor-neutral eager warmup and avoids the model's lazy + `torch.compile`/inductor initialization inside CUDA capture. The SGLang-main + backend does not yet have an equivalent warmup hook. +- SGLang auto-enables FlashInfer AllReduce Fusion for this architecture, which + bypasses the recipe's plain-NCCL path. +- The per-rank Mamba state, KV cache, and preallocated graph region exceeded + H200 memory in the tested TP=4 configuration. + +Supporting this model requires main-runner warmup parity, fused-collective +handling, and dedicated Mamba-state memory work. + +## IPC-specific troubleshooting + +| Symptom | Likely cause | +|---|---| +| `[HOOK] ERROR: cuMemImportFromShareableHandle failed with error 999` at `deep_ep.cpp` `Buffer::sync` | Foundry hook not rebuilt with the SCM_RIGHTS transport โ€” it's still packing a raw fd. Rebuild `foundry/` (`uv pip install -e . --no-build-isolation`). | +| `operation failed due to a previous error during capture` on SAVE **pass 2** | vLLM compile-cache over-keying not applied (`graph_extension_config_path` still hashed) โ†’ recompile-in-capture. Apply the `compilation.py` `ignored_factors` edit, or set `FOUNDRY_DISABLE_COMBO_KERNELS=1` (opt-in belt-and-suspenders, wired in the experimental serve script under `experimental/expert-parallel/`). | +| LOAD `illegal memory access` at replay with an NVL buffer present | Relocated peer import collided with the NVSHMEM heap hint โ€” ensure the hook build includes the dedicated import-VA zone (`0x300000000000`). | +| `error 999` only with `--deepep-mode auto`/`normal` | That's expected for non-LL modes here; this recipe pins `deepep_low_latency`. | diff --git a/recipe/experimental/foundry_load.toml b/recipe/experimental/foundry_load.toml new file mode 100644 index 00000000..09677bb2 --- /dev/null +++ b/recipe/experimental/foundry_load.toml @@ -0,0 +1,9 @@ +mode = "load" +base_addr = 0x600000000000 +region_size = "256GB" +workspace_root = "foundry_archive_ipc" +scratch_space_size = "4096MB" +# REQUIRED for the EP/IPC recipe โ€” DeepEP low-latency still allocates its RDMA +# buffer on the NVSHMEM symmetric heap (the NVL buffer is the cudaIpc path). +# Point this at the libnvshmem_host.so built by vllm/tools/ep_kernels. +nvshmem_host_path = "/data/liuxs/workarea/foundry-org/vllm/tools/ep_kernels/ep_kernels_workspace/nvshmem/lib/libnvshmem_host.so" diff --git a/recipe/experimental/foundry_load_fp8.toml b/recipe/experimental/foundry_load_fp8.toml new file mode 100644 index 00000000..a035b6ae --- /dev/null +++ b/recipe/experimental/foundry_load_fp8.toml @@ -0,0 +1,6 @@ +mode = "load" +base_addr = 0x600000000000 +region_size = "256GB" +workspace_root = "foundry_archive_ipc_fp8" +scratch_space_size = "4096MB" +nvshmem_host_path = "/data/liuxs/workarea/foundry-org/vllm/tools/ep_kernels/ep_kernels_workspace/nvshmem/lib/libnvshmem_host.so" diff --git a/recipe/experimental/foundry_save.toml b/recipe/experimental/foundry_save.toml new file mode 100644 index 00000000..da07d1a8 --- /dev/null +++ b/recipe/experimental/foundry_save.toml @@ -0,0 +1,9 @@ +mode = "save" +base_addr = 0x600000000000 +region_size = "256GB" +workspace_root = "foundry_archive_ipc" +scratch_space_size = "4096MB" +# REQUIRED for the EP/IPC recipe โ€” DeepEP low-latency still allocates its RDMA +# buffer on the NVSHMEM symmetric heap (the NVL buffer is the cudaIpc path). +# Point this at the libnvshmem_host.so built by vllm/tools/ep_kernels. +nvshmem_host_path = "/data/liuxs/workarea/foundry-org/vllm/tools/ep_kernels/ep_kernels_workspace/nvshmem/lib/libnvshmem_host.so" diff --git a/recipe/experimental/foundry_save_fp8.toml b/recipe/experimental/foundry_save_fp8.toml new file mode 100644 index 00000000..74a8adc2 --- /dev/null +++ b/recipe/experimental/foundry_save_fp8.toml @@ -0,0 +1,6 @@ +mode = "save" +base_addr = 0x600000000000 +region_size = "256GB" +workspace_root = "foundry_archive_ipc_fp8" +scratch_space_size = "4096MB" +nvshmem_host_path = "/data/liuxs/workarea/foundry-org/vllm/tools/ep_kernels/ep_kernels_workspace/nvshmem/lib/libnvshmem_host.so" diff --git a/recipe/experimental/serve_qwen3-30ba3b_ipc_ep.sh b/recipe/experimental/serve_qwen3-30ba3b_ipc_ep.sh new file mode 100755 index 00000000..949f18f6 --- /dev/null +++ b/recipe/experimental/serve_qwen3-30ba3b_ipc_ep.sh @@ -0,0 +1,92 @@ +#!/bin/bash +# Usage: bash serve_qwen3-30ba3b_ipc_ep.sh [--save|--load] +# +# EXPERIMENTAL: DeepEP expert-parallel with the legacy CUDA-IPC NVLink buffer +# (cudaIpcGet/OpenMemHandle) kept ON under foundry, instead of the default +# fabric/NVSHMEM-only path. This exercises foundry's VMM-IPC translation layer +# (SCM_RIGHTS fd transport + whole-chunk peer mapping) โ€” the path that lets +# foundry SAVE/LOAD work on machines without fabric/IMEX, where DeepEP shares +# its intranode NVL buffers via legacy IPC. +# +# Differs from recipe/vllm/serve_qwen3-30ba3b_ep.sh by exactly one env var: +# FOUNDRY_DEEPEP_NVL_IPC=1 -> _patch_deepep keeps upstream's nonzero +# num_nvl_bytes with use_fabric=False (legacy cudaIpc NVL buffer on SAVE and +# LOAD), instead of forcing use_fabric=True / num_nvl_bytes=0. +# +# Requires the foundry hook rebuilt with the SCM_RIGHTS VMM-IPC transport and +# the vLLM compile-cache fix (see this directory's README "Required code"). +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" + +EP_SIZE=${1:?Usage: $0 [--save|--load]} +MODEL_NAME="Qwen/Qwen3-30B-A3B" +HOST="0.0.0.0" +PORT=12000 +GPU_MEMORY_UTILIZATION=0.8 + +FOUNDRY_ARGS=() +if [[ "$2" == "--save" ]]; then + FOUNDRY_ARGS+=( --compilation-config.graph_extension_config_path "${SCRIPT_DIR}/foundry_save.toml" ) + # Foundry pins NCCL to the IPC/ring transports โ€” the CUMEM P2P and NVLS + # multicast fast paths cuMemMap with driver-capability flags the foundry + # VMM region doesn't carry. + export NCCL_CUMEM_ENABLE=0 + export NCCL_NVLS_ENABLE=0 + # Foundry only patches the V1 model runner; pin V1 explicitly so vLLM + # doesn't quietly route Qwen3ForCausalLM-class architectures to V2. + export VLLM_USE_V2_MODEL_RUNNER=0 + # Keep DeepEP's NVLink buffer on the legacy cudaIpc path under foundry. + # Must be identical on SAVE and LOAD (it changes the comm-buffer alloc + # trajectory and the captured graphs). + export FOUNDRY_DEEPEP_NVL_IPC=1 + echo "Using foundry SAVE (DeepEP NVL/IPC): ${SCRIPT_DIR}/foundry_save.toml" +elif [[ "$2" == "--load" ]]; then + FOUNDRY_ARGS+=( --compilation-config.graph_extension_config_path "${SCRIPT_DIR}/foundry_load.toml" ) + export NCCL_CUMEM_ENABLE=0 + export NCCL_NVLS_ENABLE=0 + export VLLM_USE_V2_MODEL_RUNNER=0 + export FOUNDRY_DEEPEP_NVL_IPC=1 + echo "Using foundry LOAD (DeepEP NVL/IPC): ${SCRIPT_DIR}/foundry_load.toml" +elif [[ -n "$2" ]]; then + echo "Usage: $0 [--save|--load]" + exit 1 +else + echo "Running without foundry (baseline vLLM)" +fi + +# LD_PRELOAD of libcuda_hook.so / libnvshmem_host.so is set by foundry's +# setup_ld_preload_env at worker spawn time (uses the paths in the TOML +# config). Baseline runs don't need either preloaded by the shell. + +# Foundry only patches the V1 model runner. vLLM defaults certain +# architectures (e.g. Qwen3ForCausalLM) to the V2 runner, which our +# patches don't touch โ€” pin V1 here. +export VLLM_USE_V2_MODEL_RUNNER=0 + +export VLLM_USE_FLASHINFER_SAMPLER=1 +export VLLM_DISABLE_SHARED_EXPERTS_STREAM=1 + +CUDAGRAPH_CAPTURE_SIZES=($(seq 1 256)) + +ARGS=( + --trust-remote-code + --host "$HOST" + --port "$PORT" + --tensor-parallel-size 1 + --data-parallel-size "$EP_SIZE" + --gpu-memory-utilization "$GPU_MEMORY_UTILIZATION" + --distributed-executor-backend uni + --enable-expert-parallel + --all2all-backend deepep_low_latency + --no-enable-prefix-caching + # max-num-batched-tokens == max-num-seqs satisfies vLLM's batched>=seqs check + # and keeps DeepEP's (tokens+1)*2 <= NVSHMEM_QP_DEPTH(=1024) โ€” no QP override. + --max-num-batched-tokens 256 + --max-num-seqs 256 + --attention-config.backend FLASH_ATTN + --compilation-config.cudagraph_mode FULL_DECODE_ONLY + --compilation-config.cudagraph_num_of_warmups 0 + --chat-template-content-format string + --cudagraph-capture-sizes ${CUDAGRAPH_CAPTURE_SIZES[@]} +) + +vllm serve "$MODEL_NAME" "${ARGS[@]}" "${FOUNDRY_ARGS[@]}" diff --git a/recipe/experimental/serve_qwen3-30ba3bfp8_ipc_ep.sh b/recipe/experimental/serve_qwen3-30ba3bfp8_ipc_ep.sh new file mode 100755 index 00000000..c5da14d8 --- /dev/null +++ b/recipe/experimental/serve_qwen3-30ba3bfp8_ipc_ep.sh @@ -0,0 +1,69 @@ +#!/bin/bash +# Usage: bash serve_qwen3-30ba3bfp8_ipc_ep.sh [--save|--load] +# +# EXPERIMENTAL: FP8 Qwen3-30B-A3B with DeepGEMM MoE + DeepEP expert-parallel, +# keeping the legacy CUDA-IPC NVLink buffer ON under foundry (FOUNDRY_DEEPEP_NVL_IPC=1). +# Same as serve_qwen3-30ba3b_ipc_ep.sh but: FP8 model + VLLM_USE_DEEP_GEMM=1 +# (blockscale FP8 MoE), and its own archive (foundry_archive_ipc_fp8). +# +# Run every phase from one consistent path (or absolute path) so SAVE pass 1/2 +# and LOAD pass the identical graph_extension_config_path โ€” see README. +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" + +EP_SIZE=${1:?Usage: $0 [--save|--load]} +MODEL_NAME="Qwen/Qwen3-30B-A3B-FP8" +HOST="0.0.0.0" +PORT=12000 +GPU_MEMORY_UTILIZATION=0.8 + +FOUNDRY_ARGS=() +if [[ "$2" == "--save" ]]; then + FOUNDRY_ARGS+=( --compilation-config.graph_extension_config_path "${SCRIPT_DIR}/foundry_save_fp8.toml" ) + export NCCL_CUMEM_ENABLE=0 + export NCCL_NVLS_ENABLE=0 + export VLLM_USE_V2_MODEL_RUNNER=0 + export FOUNDRY_DEEPEP_NVL_IPC=1 + echo "Using foundry SAVE (FP8/DeepGEMM, DeepEP NVL/IPC): ${SCRIPT_DIR}/foundry_save_fp8.toml" +elif [[ "$2" == "--load" ]]; then + FOUNDRY_ARGS+=( --compilation-config.graph_extension_config_path "${SCRIPT_DIR}/foundry_load_fp8.toml" ) + export NCCL_CUMEM_ENABLE=0 + export NCCL_NVLS_ENABLE=0 + export VLLM_USE_V2_MODEL_RUNNER=0 + export FOUNDRY_DEEPEP_NVL_IPC=1 + echo "Using foundry LOAD (FP8/DeepGEMM, DeepEP NVL/IPC): ${SCRIPT_DIR}/foundry_load_fp8.toml" +elif [[ -n "$2" ]]; then + echo "Usage: $0 [--save|--load]" + exit 1 +else + echo "Running without foundry (baseline vLLM)" +fi + +export VLLM_USE_V2_MODEL_RUNNER=0 +export VLLM_USE_FLASHINFER_SAMPLER=1 +export VLLM_DISABLE_SHARED_EXPERTS_STREAM=1 +# FP8 blockscale MoE via DeepGEMM. +export VLLM_USE_DEEP_GEMM=1 + +CUDAGRAPH_CAPTURE_SIZES=($(seq 1 256)) + +ARGS=( + --trust-remote-code + --host "$HOST" + --port "$PORT" + --tensor-parallel-size 1 + --data-parallel-size "$EP_SIZE" + --gpu-memory-utilization "$GPU_MEMORY_UTILIZATION" + --distributed-executor-backend uni + --enable-expert-parallel + --all2all-backend deepep_low_latency + --no-enable-prefix-caching + --max-num-batched-tokens 256 + --max-num-seqs 256 + --attention-config.backend FLASH_ATTN + --compilation-config.cudagraph_mode FULL_DECODE_ONLY + --compilation-config.cudagraph_num_of_warmups 0 + --chat-template-content-format string + --cudagraph-capture-sizes ${CUDAGRAPH_CAPTURE_SIZES[@]} +) + +vllm serve "$MODEL_NAME" "${ARGS[@]}" "${FOUNDRY_ARGS[@]}" diff --git a/recipe/experimental/serve_qwen3-30ba3bfp8_sglang_ipc_ep.sh b/recipe/experimental/serve_qwen3-30ba3bfp8_sglang_ipc_ep.sh new file mode 100755 index 00000000..6b9bb242 --- /dev/null +++ b/recipe/experimental/serve_qwen3-30ba3bfp8_sglang_ipc_ep.sh @@ -0,0 +1,61 @@ +#!/bin/bash +# Usage: bash serve_qwen3-30ba3bfp8_sglang_ipc_ep.sh [--save|--load] +# +# EXPERIMENTAL: SGLang DeepEP expert parallel with a legacy CUDA-IPC NVLink +# buffer. FOUNDRY_DEEPEP_NVL_IPC=1 preserves DeepEP's nonzero NVL allocation +# and selects use_fabric=False; the Foundry VMM-IPC hook transports handles +# between ranks with SCM_RIGHTS and supports LOAD preallocation chunks. +# +# Requires DeepEP 29d31c0 or newer, sgl-deep-gemm >= 0.1.2, flash-attn-3, +# and the Foundry hook from the ep-ipc commit. +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" + +EP_SIZE=${1:?Usage: $0 [--save|--load]} +MODEL_NAME="${SGL_MODEL:-Qwen/Qwen3-30B-A3B-FP8}" +HOST="0.0.0.0" +PORT=12000 +MEM_FRACTION_STATIC=0.8 + +FOUNDRY_ARGS=() +if [[ "$2" == "--save" ]]; then + FOUNDRY_TOML="${SCRIPT_DIR}/sglang_foundry_save.toml" +elif [[ "$2" == "--load" ]]; then + FOUNDRY_TOML="${SCRIPT_DIR}/sglang_foundry_load.toml" +elif [[ -n "$2" ]]; then + echo "Usage: $0 [--save|--load]" + exit 1 +fi + +if [[ -n "${FOUNDRY_TOML:-}" ]]; then + export FOUNDRY_DEEPEP_NVL_IPC=1 + export NVSHMEM_CUMEM_HANDLE_TYPE=FILE_DESCRIPTOR + export NCCL_CUMEM_ENABLE=0 + export NCCL_NVLS_ENABLE=0 + FOUNDRY_ARGS+=( --foundry-graph-extension-config-path "$FOUNDRY_TOML" ) + echo "Using Foundry DeepEP NVL/IPC: ${FOUNDRY_TOML}" +else + echo "Running without Foundry (baseline SGLang)" +fi + +# Keep this identical between SAVE and LOAD: it affects DeepEP's allocation +# trajectory and the captured low-latency graph shapes. +export SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK=256 + +sglang serve \ + --model-path "$MODEL_NAME" \ + --trust-remote-code \ + --host "$HOST" --port "$PORT" \ + --tp-size "$EP_SIZE" \ + --dp-size "$EP_SIZE" \ + --ep-size "$EP_SIZE" \ + --enable-dp-attention \ + --moe-a2a-backend deepep \ + --deepep-mode low_latency \ + --moe-runner-backend deep_gemm \ + --mem-fraction-static "$MEM_FRACTION_STATIC" \ + --disable-radix-cache \ + --disable-custom-all-reduce \ + --chunked-prefill-size 256 \ + --attention-backend fa3 \ + --cuda-graph-max-bs 128 \ + "${FOUNDRY_ARGS[@]}" diff --git a/recipe/experimental/serve_qwen3-8b_sglang_tp.sh b/recipe/experimental/serve_qwen3-8b_sglang_tp.sh new file mode 100755 index 00000000..9b950c15 --- /dev/null +++ b/recipe/experimental/serve_qwen3-8b_sglang_tp.sh @@ -0,0 +1,106 @@ +#!/bin/bash +# Usage: CUDA_VISIBLE_DEVICES=0,1 bash serve_qwen3-8b_sglang_tp.sh [--save|--load] +# +# EXPERIMENTAL: SGLang dense tensor parallel (tp-size > 1, no DP-attention, no +# expert parallel). Unlike the DP recipe โ€” where each rank is an independent +# replica with no cross-rank collective inside the captured graph โ€” pure TP +# captures a per-layer all-reduce across the TP ranks into every CUDA graph. +# +# By default that all-reduce is served by NCCL (custom all-reduce is disabled). +# With NCCL_CUMEM_ENABLE=0, NCCL shares peer buffers over legacy CUDA IPC, which +# Foundry's VMM-IPC bridge translates. Set SGLANG_ENABLE_TORCH_SYMM_MEM=1 to use +# SGLang's torch symmetric-memory all-reduce; Foundry relocates its process-local +# buffer and pointer-table graph operands during fresh-process LOAD. +# +# Requires the Foundry hook from the ep-ipc commit (cuIpc VMM-IPC bridge). +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" + +TP_SIZE=${1:?Usage: $0 [--save|--load]} +MODEL_NAME="${SGLANG_MODEL:-Qwen/Qwen3-8B}" +HOST="0.0.0.0" +PORT=12000 +MEM_FRACTION_STATIC="${SGLANG_MEM_FRACTION_STATIC:-0.8}" +CUDA_GRAPH_MAX_BS="${SGLANG_CUDA_GRAPH_MAX_BS:-256}" +ATTENTION_BACKEND="${SGLANG_ATTENTION_BACKEND:-flashinfer}" +RANDOM_SEED="${SGL_RANDOM_SEED:-42}" + +# Keep the NCCL transport identical across baseline, SAVE, and LOAD. CUMEM P2P +# and NVLS multicast use mapping capabilities the Foundry VMM region does not +# carry, so this recipe consistently exercises the IPC/ring transport. +export NCCL_CUMEM_ENABLE=0 +export NCCL_NVLS_ENABLE=0 + +FOUNDRY_ARGS=() +MODEL_ARGS=() +if [[ -n "${SGLANG_QUANTIZATION:-}" ]]; then + MODEL_ARGS+=( --quantization "$SGLANG_QUANTIZATION" ) +fi +if [[ -n "${SGLANG_CONTEXT_LENGTH:-}" ]]; then + MODEL_ARGS+=( --context-length "$SGLANG_CONTEXT_LENGTH" ) +fi +if [[ -n "${SGLANG_REASONING_PARSER:-}" ]]; then + MODEL_ARGS+=( --reasoning-parser "$SGLANG_REASONING_PARSER" ) +fi +if [[ -n "${SGLANG_LINEAR_ATTN_BACKEND:-}" ]]; then + MODEL_ARGS+=( --linear-attn-backend "$SGLANG_LINEAR_ATTN_BACKEND" ) +fi +if [[ -n "${SGLANG_MOE_RUNNER_BACKEND:-}" ]]; then + MODEL_ARGS+=( --moe-runner-backend "$SGLANG_MOE_RUNNER_BACKEND" ) +fi +if [[ "${SGLANG_ENABLE_TORCH_SYMM_MEM:-0}" == "1" ]]; then + if [[ "${FOUNDRY_SGLANG_SYMM_MULTI_SHAPE:-0}" == "1" ]]; then + if [[ -z "${FOUNDRY_SGLANG_SYMM_GRAPH_BATCHES:-}" ]]; then + echo "FOUNDRY_SGLANG_SYMM_GRAPH_BATCHES is required for multi-shape replay" >&2 + exit 1 + fi + CUDA_GRAPH_MAX_BS="${FOUNDRY_SGLANG_SYMM_GRAPH_BATCHES##*,}" + else + if [[ -n "${SGLANG_CUDA_GRAPH_MAX_BS:-}" && "${SGLANG_CUDA_GRAPH_MAX_BS}" != "1" ]]; then + echo "Foundry SGLang torch symmetric memory currently requires SGLANG_CUDA_GRAPH_MAX_BS=1" >&2 + exit 1 + fi + CUDA_GRAPH_MAX_BS=1 + fi + MODEL_ARGS+=( --enable-torch-symm-mem ) + COMMUNICATION_BACKEND="torch symmetric memory" +else + COMMUNICATION_BACKEND="NCCL P2P/IPC" +fi + +if [[ "$2" == "--save" ]]; then + FOUNDRY_TOML="${SCRIPT_DIR}/sglang_foundry_tp_save.toml" +elif [[ "$2" == "--load" ]]; then + FOUNDRY_TOML="${SCRIPT_DIR}/sglang_foundry_tp_load.toml" +elif [[ -n "$2" ]]; then + echo "Usage: $0 [--save|--load]" + exit 1 +fi + +if [[ -n "${FOUNDRY_TOML:-}" ]]; then + export FOUNDRY_SGLANG_CONFIG_PATH="$FOUNDRY_TOML" + if ! python -c 'import importlib.util; raise SystemExit(0 if importlib.util.find_spec("sglang.srt.model_executor.runner.decode_cuda_graph_runner") else 1)'; then + FOUNDRY_ARGS+=( --foundry-graph-extension-config-path "$FOUNDRY_TOML" ) + fi + echo "Using Foundry TP (${COMMUNICATION_BACKEND}): ${FOUNDRY_TOML}" +else + echo "Running without Foundry (baseline SGLang, ${COMMUNICATION_BACKEND})" +fi + +# --disable-custom-all-reduce: route the TP all-reduce through NCCL rather than +# SGLang's custom one-shot/two-shot kernel. Both are CUDA-IPC based, but the +# NCCL path is validated end-to-end by tests/modal_sglang_tp.py. +sglang serve \ + --model-path "$MODEL_NAME" \ + --trust-remote-code \ + --host "$HOST" --port "$PORT" \ + --tp-size "$TP_SIZE" \ + --random-seed "$RANDOM_SEED" \ + --mem-fraction-static "$MEM_FRACTION_STATIC" \ + --disable-radix-cache \ + --disable-piecewise-cuda-graph \ + --disable-custom-all-reduce \ + --attention-backend "$ATTENTION_BACKEND" \ + --cuda-graph-max-bs "$CUDA_GRAPH_MAX_BS" \ + --decode-log-interval 1 \ + "${MODEL_ARGS[@]}" \ + "${FOUNDRY_ARGS[@]}" diff --git a/recipe/experimental/sglang_foundry_load.toml b/recipe/experimental/sglang_foundry_load.toml new file mode 100644 index 00000000..8484af73 --- /dev/null +++ b/recipe/experimental/sglang_foundry_load.toml @@ -0,0 +1,8 @@ +mode = "load" +base_addr = 0x600000000000 +region_size = "256GB" +workspace_root = "foundry_archive_sglang_ipc" +scratch_space_size = "4096MB" +# DeepEP low-latency still uses NVSHMEM for its RDMA symmetric heap. +# Foundry auto-detects libnvshmem_host.so from the nvidia-nvshmem wheel. +# Set nvshmem_host_path here only when using a custom NVSHMEM build. diff --git a/recipe/experimental/sglang_foundry_save.toml b/recipe/experimental/sglang_foundry_save.toml new file mode 100644 index 00000000..2c9d18eb --- /dev/null +++ b/recipe/experimental/sglang_foundry_save.toml @@ -0,0 +1,8 @@ +mode = "save" +base_addr = 0x600000000000 +region_size = "256GB" +workspace_root = "foundry_archive_sglang_ipc" +scratch_space_size = "4096MB" +# DeepEP low-latency still uses NVSHMEM for its RDMA symmetric heap. +# Foundry auto-detects libnvshmem_host.so from the nvidia-nvshmem wheel. +# Set nvshmem_host_path here only when using a custom NVSHMEM build. diff --git a/recipe/experimental/sglang_foundry_tp_load.toml b/recipe/experimental/sglang_foundry_tp_load.toml new file mode 100644 index 00000000..8ed9f23a --- /dev/null +++ b/recipe/experimental/sglang_foundry_tp_load.toml @@ -0,0 +1,8 @@ +mode = "load" +base_addr = 0x600000000000 +region_size = "256GB" +workspace_root = "foundry_archive_sglang_tp" +scratch_space_size = "1024MB" +# Dense tensor parallel needs no NVSHMEM (no DeepEP). The per-layer TP +# all-reduce is served by NCCL P2P/IPC. NCCL initialization consumes about +# 302 MiB on the validated TP=2 H100 path, leaving headroom before this boundary. diff --git a/recipe/experimental/sglang_foundry_tp_save.toml b/recipe/experimental/sglang_foundry_tp_save.toml new file mode 100644 index 00000000..428ade3d --- /dev/null +++ b/recipe/experimental/sglang_foundry_tp_save.toml @@ -0,0 +1,8 @@ +mode = "save" +base_addr = 0x600000000000 +region_size = "256GB" +workspace_root = "foundry_archive_sglang_tp" +scratch_space_size = "1024MB" +# Dense tensor parallel needs no NVSHMEM (no DeepEP). The per-layer TP +# all-reduce is served by NCCL P2P/IPC. NCCL initialization consumes about +# 302 MiB on the validated TP=2 H100 path, leaving headroom before this boundary. diff --git a/recipe/sglang/README.md b/recipe/sglang/README.md index af4ce528..83ca1662 100644 --- a/recipe/sglang/README.md +++ b/recipe/sglang/README.md @@ -39,6 +39,7 @@ model or topology before SAVE. |---|---|---|---| | 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` | +| Tensor parallel | `../experimental/serve_qwen3-8b_sglang_tp.sh` | Qwen3-8B | experimental TP=2 validation on 2ร—H100; NCCL P2P/IPC | | Expert parallel | `serve_qwen3-30ba3bfp8_ep.sh` | Qwen3-30B-A3B-FP8 | DP-attention + DeepEP; fa3 backend | ## Installation @@ -92,6 +93,44 @@ CUDA_VISIBLE_DEVICES=0,1 bash serve_qwen3-1.7b_dp.sh 2 --save CUDA_VISIBLE_DEVICES=0,1 bash serve_qwen3-1.7b_dp.sh 2 --load ``` +## Run (experimental tensor parallel) + +Dense TP captures NCCL all-reduce in every decode graph. Keep NCCL on P2P/IPC +and use the same topology for baseline, SAVE, and LOAD; the script sets the +required environment, graph flags, and server seed consistently. + +```bash +cd ../experimental +rm -rf foundry_archive_sglang_tp +CUDA_VISIBLE_DEVICES=0,1 bash serve_qwen3-8b_sglang_tp.sh 2 --save +CUDA_VISIBLE_DEVICES=0,1 bash serve_qwen3-8b_sglang_tp.sh 2 --load +curl -s http://0.0.0.0:12000/v1/completions \ + -H 'Content-Type: application/json' \ + -d '{"model":"Qwen/Qwen3-8B","prompt":"The capital of France is","max_tokens":12,"temperature":0}' +``` + +For a repeatable 2ร—H100 baseline โ†’ SAVE โ†’ SAVE โ†’ LOAD proof, run +`modal run --env tests/modal_sglang_tp.py` from the +repository root. The harness checks byte-identical temperature-0 outputs, +per-rank allocation offsets, graph counts, and CUDA/NCCL errors. +Prefix the command with `SGLANG_ENABLE_TORCH_SYMM_MEM=1` to run the +symmetric-memory matrix. + +For the validated SGLang-main torch symmetric-memory path, set +`SGLANG_ENABLE_TORCH_SYMM_MEM=1` in baseline, SAVE, and LOAD. This mode requires +TP=2 and forces `SGLANG_CUDA_GRAPH_MAX_BS=1`; batch sizes above one fall back to +eager execution. Foundry preserves full graph JSON, records the process-local +communication operands, rebuilds the required SGLang warmup state, and +relocates those operands on fresh LOAD. The 2ร—H100 proof restored one graph per +rank at `final_alloc_offset=68555898880` and returned byte-identical +temperature-0 outputs without runtime errors. Multi-shape symmetric capture is +rejected because its graphs share mutable FlashInfer planner state. + +Do not extrapolate this result to arbitrary NCCL versions, GPU counts, or +topologies. Upstream [Discussion #5](https://github.com/orgs/foundry-org/discussions/5) +documents the remaining NCCL determinism concern and a successful unpublished +torch symmetric-memory prototype, but no official implementation. + ## Run (expert parallel / DeepEP) EP needs three kernel packages โ€” all SGLang-native, no vLLM involved: @@ -100,14 +139,14 @@ EP needs three kernel packages โ€” all SGLang-native, no vLLM involved: 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 +- **DeepEP** @ `29d31c0` or newer โ€” required for Foundry's fabric/IPC policy + switch. Build via SGLang's own installer, overriding its old `9af0e0d` pin, + and build against the NVSHMEM wheel above. For a single node you can skip the script's gdrcopy/RDMA apt steps and just build the wheel: ```bash git clone https://github.com/deepseek-ai/DeepEP.git && cd DeepEP - git checkout 9af0e0d0e74f3577af1979c9b9e1ac2cad0104ee + git checkout 29d31c095796f3c8ece47ee9cdcc167051bbeed9 TORCH_CUDA_ARCH_LIST="9.0" python setup.py install # Hopper; "9.0;10.0" for Blackwell cd .. ``` diff --git a/recipe/sglang/serve_qwen3-30ba3bfp8_ep.sh b/recipe/sglang/serve_qwen3-30ba3bfp8_ep.sh index a185fda3..d3035f2d 100755 --- a/recipe/sglang/serve_qwen3-30ba3bfp8_ep.sh +++ b/recipe/sglang/serve_qwen3-30ba3bfp8_ep.sh @@ -2,7 +2,7 @@ # 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, +# Requires (see README ยงEP): deep_ep @ 29d31c0+, sgl-deep-gemm >= 0.1.2, flash-attn-3, # and the nvshmem_host_path uncommented in foundry_{save,load}.toml. SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" diff --git a/recipe/sglang/serve_qwen3-mini.sh b/recipe/sglang/serve_qwen3-mini.sh index 77449482..a584d53a 100755 --- a/recipe/sglang/serve_qwen3-mini.sh +++ b/recipe/sglang/serve_qwen3-mini.sh @@ -4,7 +4,7 @@ SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" -MODEL_NAME="Qwen/Qwen3-1.7B" +MODEL_NAME="${SGL_MODEL:-Qwen/Qwen3-1.7B}" HOST="0.0.0.0" PORT=12000 MEM_FRACTION_STATIC=0.6 @@ -24,7 +24,11 @@ fi FOUNDRY_ARGS=() if [[ -n "${FOUNDRY_TOML:-}" ]]; then - FOUNDRY_ARGS+=( --foundry-graph-extension-config-path "${FOUNDRY_TOML}" ) + export FOUNDRY_SGLANG_CONFIG_PATH="${FOUNDRY_TOML}" + if ! python -c 'import importlib.util; raise SystemExit(0 if importlib.util.find_spec("sglang.srt.model_executor.runner.decode_cuda_graph_runner") else 1)'; then + # Legacy Foundry fork: activation predates SGLang's plugin framework. + FOUNDRY_ARGS+=( --foundry-graph-extension-config-path "${FOUNDRY_TOML}" ) + fi fi # LD_PRELOAD of libcuda_hook.so is set by foundry's setup_ld_preload_env at @@ -37,8 +41,11 @@ sglang serve \ --trust-remote-code \ --host "$HOST" --port "$PORT" \ --tp-size 1 \ + --random-seed 42 \ --mem-fraction-static "$MEM_FRACTION_STATIC" \ --disable-radix-cache \ + --disable-piecewise-cuda-graph \ --attention-backend flashinfer \ --cuda-graph-max-bs 512 \ + --decode-log-interval 1 \ "${FOUNDRY_ARGS[@]}" diff --git a/recipe/vllm/README.md b/recipe/vllm/README.md index ea55d2f6..4a24069e 100644 --- a/recipe/vllm/README.md +++ b/recipe/vllm/README.md @@ -11,6 +11,7 @@ recipe/vllm/ โ”œโ”€โ”€ foundry_load.toml # shared LOAD config (same workspace_root) โ”œโ”€โ”€ serve_qwen3-mini.sh # Qwen3-1.7B single GPU โ”œโ”€โ”€ serve_qwen3-14b_dp.sh # Qwen3-14B data parallel +โ”œโ”€โ”€ serve_qwen3-8b_tp.sh # Qwen3-8B experimental tensor parallel โ”œโ”€โ”€ serve_qwen3-30ba3b_ep.sh # Qwen3-30B-A3B (MoE) expert parallel (DeepEP) โ””โ”€โ”€ serve_qwen3-30ba3bfp8_ep.sh # Qwen3-30B-A3B FP8 expert parallel (DeepGEMM) ``` @@ -20,6 +21,7 @@ Every script accepts the same trailing `--save` / `--load` flag. Scripts that sc ```bash bash serve_qwen3-mini.sh [--save|--load] bash serve_qwen3-14b_dp.sh [--save|--load] +bash serve_qwen3-8b_tp.sh [--save|--load] bash serve_qwen3-30ba3b_ep.sh [--save|--load] bash serve_qwen3-30ba3bfp8_ep.sh [--save|--load] ``` @@ -133,16 +135,29 @@ bash serve_qwen3-mini.sh --load bash ../../../experimental/query.sh 12000 Qwen/Qwen3-1.7B ``` -For DP / EP, prepend the parallel size: +For DP / TP / EP, prepend the parallel size: ```bash rm -rf foundry_archive +bash serve_qwen3-8b_tp.sh 2 --save # SAVE pass 1, TP=2 +bash serve_qwen3-8b_tp.sh 2 --save # SAVE pass 2 +bash serve_qwen3-8b_tp.sh 2 --load + +# Or expert parallel: +rm -rf foundry_archive bash serve_qwen3-30ba3b_ep.sh 2 --save # SAVE pass 1, EP=2 bash serve_qwen3-30ba3b_ep.sh 2 --save # SAVE pass 2 bash serve_qwen3-30ba3b_ep.sh 2 --load bash ../../../experimental/query.sh 12000 Qwen/Qwen3-30B-A3B ``` +For the automated 2ร—H100 TP proof, pin an immutable Foundry revision: + +```bash +FOUNDRY_REF=$(git rev-parse HEAD) \ + modal run --env tests/modal_vllm_tp.py +``` + ## Archive layout Every SAVE writes into `foundry_archive/` next to the script's working directory: @@ -166,6 +181,31 @@ All scripts, when invoked with `--save` or `--load`, export one env-var override - `VLLM_USE_V2_MODEL_RUNNER=0` โ€” pin the V1 model runner. Foundry's monkey-patches are on `vllm.v1.worker.gpu_model_runner.GPUModelRunner`; vLLM's `VllmConfig.use_v2_model_runner` property silently routes architectures in `DEFAULT_V2_MODEL_RUNNER_ARCHITECTURES` (currently `Qwen3ForCausalLM`) to `vllm.v1.worker.gpu.model_runner` (V2), which our patches don't touch. Without this pin, `capture_model` runs unhooked, no archive is written, and SAVE silently produces an empty workspace. +The tensor-parallel script is a pinned experimental NCCL configuration, not +official/general TP support. Upstream +[Discussion #5](https://github.com/orgs/foundry-org/discussions/5) notes that +NCCL initialization is not generally deterministic and reports a successful +internal torch symmetric-memory prototype, but no official implementation. +Follow official support in +[issue #6](https://github.com/foundry-org/foundry/issues/6). + +For this validated configuration, the script keeps one collective +implementation across baseline, SAVE, and LOAD: + +- `--disable-custom-all-reduce`, `VLLM_ALLREDUCE_USE_SYMM_MEM=0`, + `VLLM_ALLREDUCE_USE_FLASHINFER=0`, and `VLLM_USE_NCCL_SYMM_MEM=0` route + collectives through PyNCCL. +- `NCCL_CUMEM_ENABLE=0` and `NCCL_NVLS_ENABLE=0` select NCCL P2P/IPC, whose + cross-process VMM mappings Foundry restores. +- `pass_config.fuse_allreduce_rms=false` disables vLLM's optional FlashInfer + all-reduce/RMSNorm compiler rewrite. Otherwise the second SAVE can try to + create its workspace inside CUDA graph capture. +- One atomic `--compilation-config` JSON value carries both that pass override + and the Foundry TOML path. Mixing `-cc.*` and `--compilation-config.*` + overrides can silently replace the config and leave an empty archive. +- `VLLM_WORKER_MULTIPROC_METHOD=spawn` selects the multiprocessing mode accepted + by the pinned Foundry vLLM fork. + The two EP scripts (`*_ep.sh`) on top of the basic flags also set: - `--enable-expert-parallel --all2all-backend deepep_low_latency`. @@ -201,5 +241,7 @@ nvshmem_qp_depth >= (num_max_dispatch_tokens_per_rank + 1) * 2 | `[HOOK] WARNING: Neither nvshmemx_cumodule_init nor nvshmemx_culibrary_init found` | `libnvshmem_host.so` is not in `LD_PRELOAD`. EP scripts need both `libcuda_hook.so` *and* `libnvshmem_host.so` preloaded. | | EP graph replay fails with `illegal memory access` | Almost always `libnvshmem_host.so` isn't preloaded. | | `[foundry] LOAD requires warmup_state.json` raised on LOAD | SAVE pass 1 didn't write `warmup_state.json`. Make sure SAVE finishes "Application startup complete" before you Ctrl-C, and re-run pass 1, then pass 2. | +| TP SAVE is healthy but `rank_*/` archives are empty | The nested compilation flags replaced `graph_extension_config_path`; use the atomic config from `serve_qwen3-8b_tp.sh`. | +| TP capture raises `Flashinfer allreduce workspace must be initialized` | The all-reduce/RMSNorm fusion is active; keep `pass_config.fuse_allreduce_rms=false` in every phase. | | `AssertionError ... nvshmem_qp_depth >= (num_max_dispatch_tokens_per_rank + 1) * 2` in `deep_ep/buffer.py:602` | `max_num_batched_tokens > 511` with default `NVSHMEM_QP_DEPTH=1024`. Lower `--max-num-batched-tokens` or `export NVSHMEM_QP_DEPTH` higher โ€” see the constraint above. | | `[TIMING] NCCL init: 122.992 s` on Foundry SAVE | Foundry set `LD_PRELOAD` to apply cuda driver hook which extends module load time, but it won't affect LOAD because all modules are preloaded. | diff --git a/recipe/vllm/serve_qwen3-8b_tp.sh b/recipe/vllm/serve_qwen3-8b_tp.sh new file mode 100644 index 00000000..c36756db --- /dev/null +++ b/recipe/vllm/serve_qwen3-8b_tp.sh @@ -0,0 +1,77 @@ +#!/bin/bash +# Usage: CUDA_VISIBLE_DEVICES=0,1 bash serve_qwen3-8b_tp.sh [--save|--load] +# +# EXPERIMENTAL: this is a validated TP=2 NCCL configuration, not official TP. +# Upstream has reported an unpublished torch symmetric-memory prototype. +# +# Dense tensor parallelism captures a per-layer all-reduce in every CUDA graph. +# Keep that collective on NCCL's legacy P2P/IPC path in baseline, SAVE, and LOAD +# so all phases exercise the same transport and Foundry can restore its VMM +# mappings at their recorded addresses. +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" + +TP_SIZE="${1:-}" +MODE="${2:-}" +if [[ ! "$TP_SIZE" =~ ^[0-9]+$ ]] || (( TP_SIZE < 2 )); then + echo "Usage: $0 =2+ [--save|--load]" >&2 + exit 1 +fi + +MODEL_NAME="${VLLM_MODEL:-Qwen/Qwen3-8B}" +HOST="${VLLM_HOST:-0.0.0.0}" +PORT="${VLLM_PORT:-12000}" +GPU_MEMORY_UTILIZATION="${VLLM_GPU_MEMORY_UTILIZATION:-0.8}" + +# Route TP collectives through one reproducible NCCL implementation. vLLM's +# symmetric-memory and FlashInfer paths take precedence over PyNCCL unless they +# are disabled explicitly, even with --disable-custom-all-reduce. +export NCCL_CUMEM_ENABLE=0 +export NCCL_NVLS_ENABLE=0 +export VLLM_ALLREDUCE_USE_SYMM_MEM=0 +export VLLM_ALLREDUCE_USE_FLASHINFER=0 +export VLLM_USE_NCCL_SYMM_MEM=0 + +# Foundry patches the V1 model runner. The pinned vLLM fork accepts spawn (not +# forkserver) for worker processes, and Foundry injects LD_PRELOAD at each spawn. +export VLLM_USE_V2_MODEL_RUNNER=0 +export VLLM_WORKER_MULTIPROC_METHOD=spawn + +if [[ "$MODE" == "--save" ]]; then + FOUNDRY_TOML="${SCRIPT_DIR}/foundry_save.toml" +elif [[ "$MODE" == "--load" ]]; then + FOUNDRY_TOML="${SCRIPT_DIR}/foundry_load.toml" +elif [[ -n "$MODE" ]]; then + echo "Usage: $0 =2+ [--save|--load]" >&2 + exit 1 +fi + +if [[ -n "${FOUNDRY_TOML:-}" ]]; then + COMPILATION_CONFIG="{\"cudagraph_mode\":\"FULL_DECODE_ONLY\",\"cudagraph_num_of_warmups\":0,\"pass_config\":{\"fuse_allreduce_rms\":false},\"graph_extension_config_path\":\"${FOUNDRY_TOML}\"}" + echo "Using Foundry TP (NCCL P2P/IPC): ${FOUNDRY_TOML}" +else + COMPILATION_CONFIG='{"cudagraph_mode":"FULL_DECODE_ONLY","cudagraph_num_of_warmups":0,"pass_config":{"fuse_allreduce_rms":false}}' + echo "Running without Foundry (baseline vLLM TP)" +fi + +CUDAGRAPH_CAPTURE_SIZES=($(seq 1 64)) + +ARGS=( + --trust-remote-code + --host "$HOST" + --port "$PORT" + --tensor-parallel-size "$TP_SIZE" + --distributed-executor-backend mp + --disable-custom-all-reduce + --gpu-memory-utilization "$GPU_MEMORY_UTILIZATION" + --no-enable-prefix-caching + --max-num-batched-tokens 64 + --max-num-seqs 64 + --attention-config.backend FLASH_ATTN + --compilation-config "$COMPILATION_CONFIG" + --chat-template-content-format string + --cudagraph-capture-sizes "${CUDAGRAPH_CAPTURE_SIZES[@]}" +) + +vllm serve "$MODEL_NAME" "${ARGS[@]}" diff --git a/tests/modal_sglang_main_smoke.py b/tests/modal_sglang_main_smoke.py new file mode 100644 index 00000000..0f55f6a5 --- /dev/null +++ b/tests/modal_sglang_main_smoke.py @@ -0,0 +1,255 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the Foundry project +"""Single-GPU SAVE/LOAD smoke test against current SGLang main.""" + +from __future__ import annotations + +import contextlib +import json +import os +import re +import shutil +import signal +import subprocess +import time +import urllib.request +from pathlib import Path +from typing import IO + +import modal + + +def _local_revision() -> str: + try: + return subprocess.check_output( + ["git", "rev-parse", "HEAD"], + cwd=Path(__file__).parents[1], + stderr=subprocess.DEVNULL, + text=True, + ).strip() + except (OSError, subprocess.CalledProcessError): + return "main" + + +FOUNDRY_REF = os.environ.get("FOUNDRY_REF") or _local_revision() +SGLANG_REF = os.environ.get( + "SGLANG_REF", + "1b63155efead7494843f6db8681851ba94611f73", +) +MODEL = os.environ.get("SGL_MODEL", "Qwen/Qwen3-1.7B") +GPU = os.environ.get("MODAL_GPU", "H100") +PORT = 12000 + +ARCHIVE_ROOT = "/data/archive" +HF_ROOT = "/data/hf" +RUN_ID = os.environ.get("SMOKE_RUN_ID") or ( + f"{FOUNDRY_REF[:12]}-{SGLANG_REF[:12]}-{time.time_ns()}" +) +RUN_ROOT = f"{ARCHIVE_ROOT}/{RUN_ID}" +WORKSPACE = f"{RUN_ROOT}/foundry_archive" +RECIPE = "/foundry/recipe/sglang/serve_qwen3-mini.sh" + +ERROR_PATTERN = re.compile( + r"\[(?:HOOK|REPLAY)\] ERROR|Traceback|CUDA error|illegal memory access|" + r"Scheduler hit an exception|segmentation fault", + re.IGNORECASE, +) + +image = ( + modal.Image.from_registry( + "nvidia/cuda:13.0.1-cudnn-devel-ubuntu24.04", + add_python="3.12", + ) + .env({"CC": "gcc", "CXX": "g++"}) + .apt_install( + "git", + "build-essential", + "cmake", + "ninja-build", + "libboost-all-dev", + "libnuma-dev", + ) + .pip_install( + "cmake>=4.0.0", + "wheel", + "packaging", + "ninja", + "setuptools>=80", + "setuptools-scm", + ) + .pip_install( + "torch==2.11.0", + "torchvision==0.26.0", + "torchaudio==2.11.0", + index_url="https://download.pytorch.org/whl/cu130", + ) + .run_commands( + "git clone https://github.com/sgl-project/sglang.git /sglang", + f"cd /sglang && git checkout {SGLANG_REF}", + "cd /sglang && pip install -e 'python[all]' --no-build-isolation", + ) + .run_commands( + "git clone https://github.com/modal-projects/foundry.git /foundry", + f"cd /foundry && git checkout {FOUNDRY_REF}", + "cd /foundry && pip install -e . --no-build-isolation", + ) + .env( + { + "HF_HOME": HF_ROOT, + "HUGGINGFACE_HUB_CACHE": f"{HF_ROOT}/hub", + "PYTHONUNBUFFERED": "1", + "FOUNDRY_REF": FOUNDRY_REF, + "SGLANG_REF": SGLANG_REF, + } + ) +) + +app = modal.App("foundry-sglang-main-smoke") +archive_volume = modal.Volume.from_name( + "foundry-sglang-main-smoke-archive", + create_if_missing=True, +) +hf_volume = modal.Volume.from_name( + "foundry-sglang-tp-hf-cache", + create_if_missing=True, +) + + +def _wait_healthy(process: subprocess.Popen, log_path: Path) -> None: + deadline = time.time() + 1200 + url = f"http://127.0.0.1:{PORT}/health_generate" + while time.time() < deadline: + if process.poll() is not None: + raise RuntimeError( + f"SGLang exited with {process.returncode}\n" + + "\n".join(log_path.read_text(errors="replace").splitlines()[-100:]) + ) + try: + with urllib.request.urlopen(url, timeout=5) as response: + if response.status == 200: + return + except Exception: + pass + time.sleep(2) + raise RuntimeError("SGLang main did not become healthy") + + +def _generate() -> str: + request = urllib.request.Request( + f"http://127.0.0.1:{PORT}/generate", + data=json.dumps( + { + "text": "The capital of France is", + "sampling_params": { + "temperature": 0, + "max_new_tokens": 32, + }, + } + ).encode(), + headers={"Content-Type": "application/json"}, + ) + with urllib.request.urlopen(request, timeout=120) as response: + return json.loads(response.read())["text"] + + +def _launch(phase: str, log_path: Path) -> tuple[subprocess.Popen, IO[str]]: + args = ["bash", RECIPE] + if phase in {"save", "load"}: + args.append(f"--{phase}") + env = os.environ.copy() + env["SGL_MODEL"] = MODEL + log_file = log_path.open("w") + process = subprocess.Popen( + args, + cwd=RUN_ROOT, + env=env, + stdout=log_file, + stderr=subprocess.STDOUT, + preexec_fn=os.setsid, + ) + return process, log_file + + +def _stop(process: subprocess.Popen, log_file: IO[str]) -> None: + try: + os.killpg(os.getpgid(process.pid), signal.SIGINT) + process.wait(timeout=60) + except Exception: + with contextlib.suppress(Exception): + os.killpg(os.getpgid(process.pid), signal.SIGKILL) + finally: + log_file.close() + + +@app.function( + image=image, + gpu=GPU, + timeout=3600, + volumes={ARCHIVE_ROOT: archive_volume, HF_ROOT: hf_volume}, +) +def run_phase(phase: str) -> dict: + run_root = Path(RUN_ROOT) + run_root.mkdir(parents=True, exist_ok=True) + if phase == "save": + shutil.rmtree(WORKSPACE, ignore_errors=True) + + log_path = run_root / f"{phase}.log" + process, log_file = _launch(phase, log_path) + try: + _wait_healthy(process, log_path) + output = _generate() + log_file.flush() + text = log_path.read_text(errors="replace") + finally: + _stop(process, log_file) + + result = { + "phase": phase, + "output": output, + "errors": sorted(set(ERROR_PATTERN.findall(text))), + "saved_graphs": text.count("Saved SGLang-main CUDA graph"), + "loaded_graphs": text.count("Loaded SGLang-main graph"), + "native_capture_progress": text.count("Capturing batches"), + "replay_observed": "cuda graph: True" in text, + "plugin_observed": "SGLang graph extension setup completed" in text, + } + archive_volume.commit() + print(json.dumps(result, sort_keys=True)) + return result + + +@app.function( + image=modal.Image.debian_slim(), + volumes={ARCHIVE_ROOT: archive_volume}, +) +def cleanup() -> None: + shutil.rmtree(RUN_ROOT, ignore_errors=True) + archive_volume.commit() + + +@app.local_entrypoint() +def main() -> None: + results = {phase: run_phase.remote(phase) for phase in ("baseline", "save", "load")} + checks = { + "outputs_match": results["baseline"]["output"] + == results["save"]["output"] + == results["load"]["output"], + "save_wrote_graphs": results["save"]["saved_graphs"] > 0, + "load_restored_graphs": results["load"]["loaded_graphs"] > 0, + "load_did_not_capture": results["load"]["saved_graphs"] == 0 + and results["load"]["loaded_graphs"] > 0, + "load_replayed_graph": results["load"]["replay_observed"], + "plugin_loaded": results["save"]["plugin_observed"] and results["load"]["plugin_observed"], + "no_errors": not any(result["errors"] for result in results.values()), + } + report = { + "foundry_ref": FOUNDRY_REF, + "sglang_ref": SGLANG_REF, + "checks": checks, + "results": results, + } + print(json.dumps(report, indent=2, sort_keys=True)) + failures = [name for name, passed in checks.items() if not passed] + if failures: + raise RuntimeError(f"SGLang-main smoke failed: {', '.join(failures)}") + cleanup.remote() diff --git a/tests/modal_sglang_tp.py b/tests/modal_sglang_tp.py new file mode 100644 index 00000000..9877be39 --- /dev/null +++ b/tests/modal_sglang_tp.py @@ -0,0 +1,1061 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the Foundry project +"""End-to-end SGLang tensor-parallel validation on Modal. + +The harness runs four phases on two GPUs: + +1. A non-Foundry SGLang TP baseline using the same pinned collective settings. +2. Foundry SAVE into a fresh archive. +3. A second SAVE to check per-rank allocation reproducibility. +4. Foundry LOAD in a fresh server process and the same deterministic queries. + +Success requires byte-identical decoded baseline/LOAD responses, reproducible +semantic graph metadata/inventories and per-rank allocation offsets, replay of +every configured graph shape, observation of the selected collective backend, +and no CUDA/NCCL errors. Torch symmetric memory uses the validated TP=2, +batch-1-only graph scope; NCCL exercises the full shape inventory. + +Run the current checkout: + + modal run --env tests/modal_sglang_tp.py + +Run a specific commit: + + FOUNDRY_REF= \ + modal run --env tests/modal_sglang_tp.py + +Use an immutable commit rather than a moving branch name so Modal's image cache +cannot reuse a build from an older branch head. +""" + +from __future__ import annotations + +import contextlib +import hashlib +import json +import os +import re +import shutil +import signal +import subprocess +import threading +import time +import urllib.request +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path +from typing import IO + +import modal + + +def _local_foundry_revision() -> str: + try: + return subprocess.check_output( + ["git", "rev-parse", "HEAD"], + cwd=Path(__file__).parents[1], + text=True, + ).strip() + except (OSError, subprocess.CalledProcessError): + return "main" + + +CUDA_TAG = "13.0.1-cudnn-devel-ubuntu24.04" + +FOUNDRY_REPO = "https://github.com/modal-projects/foundry.git" +FOUNDRY_REF = os.environ.get("FOUNDRY_REF") or _local_foundry_revision() + +SGLANG_REPO = os.environ.get( + "SGLANG_REPO", + "https://github.com/sgl-project/sglang.git", +) +SGLANG_REF = os.environ.get( + "SGLANG_REF", + "1b63155efead7494843f6db8681851ba94611f73", +) +SGLANG_FOUNDRY_COMMIT = os.environ.get("SGLANG_FOUNDRY_COMMIT", "") +SGLANG_PATCH_COMMAND = ( + f"cd /sglang && git fetch origin foundry && git cherry-pick --no-commit {SGLANG_FOUNDRY_COMMIT}" + if SGLANG_FOUNDRY_COMMIT + else "true" +) +USES_MAIN_PLUGIN = not SGLANG_FOUNDRY_COMMIT + +MODEL = os.environ.get("SGLANG_MODEL", "Qwen/Qwen3-8B") +SERVE_ENV_NAMES = ( + "SGLANG_QUANTIZATION", + "SGLANG_CONTEXT_LENGTH", + "SGLANG_REASONING_PARSER", + "SGLANG_LINEAR_ATTN_BACKEND", + "SGLANG_MOE_RUNNER_BACKEND", + "SGLANG_ATTENTION_BACKEND", + "SGLANG_MEM_FRACTION_STATIC", + "SGLANG_CUDA_GRAPH_MAX_BS", + "SGLANG_ENABLE_TORCH_SYMM_MEM", + "FOUNDRY_SGLANG_SYMM_MULTI_SHAPE", + "FOUNDRY_SGLANG_SYMM_GRAPH_BATCHES", +) +SERVE_ENV = {name: os.environ[name] for name in SERVE_ENV_NAMES if os.environ.get(name)} +USES_TORCH_SYMM_MEM = SERVE_ENV.get("SGLANG_ENABLE_TORCH_SYMM_MEM") == "1" +USES_MULTI_SHAPE = SERVE_ENV.get("FOUNDRY_SGLANG_SYMM_MULTI_SHAPE") == "1" +TP_SIZE = int(os.environ.get("TP_SIZE", "2")) +if TP_SIZE < 2: + raise ValueError("Tensor-parallel validation requires TP_SIZE >= 2") +MODAL_GPU = os.environ.get("MODAL_GPU", f"H100:{TP_SIZE}") +PORT = 12000 + +ARCHIVE_ROOT = "/data/archive" +RUNS_ROOT = f"{ARCHIVE_ROOT}/runs" +HF_CACHE = "/data/hf" +RECIPE_SCRIPT = "/foundry/recipe/experimental/serve_qwen3-8b_sglang_tp.sh" +FOUNDRY_BASE_ADDR = 0x600000000000 +FOUNDRY_REGION_END = FOUNDRY_BASE_ADDR + 256 * 1024**3 + +PROMPTS = [ + "Explain what tensor parallelism is in one paragraph.", + "Write a haiku about GPUs.", + "List three prime numbers greater than 50.", + "Summarize why CUDA graphs speed up inference.", +] +CONCURRENT_PROMPTS = PROMPTS * 2 +MAX_NEW_TOKENS = 64 +GRAPH_BATCH_EXERCISE_PROMPT = "Answer with the number 7." +GRAPH_BATCH_EXERCISE_MAX_NEW_TOKENS = 2 +NCCL_DEFAULT_GRAPH_BATCHES = ( + 1, + 2, + 4, + 8, + 12, + 16, + 24, + 32, + 40, + 48, + 56, + 64, + 72, + 80, + 88, + 96, + 104, + 112, + 120, + 128, + 136, + 144, + 152, + 160, + 168, + 176, + 184, + 192, + 200, + 208, + 216, + 224, + 232, + 240, + 248, + 256, +) +RUN_ID_OVERRIDE = os.environ.get("TP_RUN_ID") + + +def expected_graph_batches_from_env( + *, + graph_batches: str | None, + uses_torch_symm_mem: bool, +) -> tuple[int, ...]: + if graph_batches: + return tuple(int(value) for value in graph_batches.split(",")) + if uses_torch_symm_mem: + return (1,) + return NCCL_DEFAULT_GRAPH_BATCHES + + +def _expected_graph_batches() -> tuple[int, ...]: + return expected_graph_batches_from_env( + graph_batches=os.environ.get("FOUNDRY_SGLANG_SYMM_GRAPH_BATCHES"), + uses_torch_symm_mem=USES_TORCH_SYMM_MEM, + ) + + +def required_archive_files( + *, + uses_torch_symm_mem: bool, + uses_multi_shape: bool, +) -> set[str]: + files = { + "graph_manifest.json", + "fatbin_image_packed.img", + "final_alloc_offset.json", + } + if uses_torch_symm_mem: + files.update({"sglang_graph_backend.json", "symmetric_memory_state.json"}) + if uses_multi_shape: + files.add("sglang_shape_state.json") + return files + + +def restored_graph_replay_observed( + expected_batches: tuple[int, ...], + replay_batch_sizes: list[int] | set[int], +) -> bool: + return set(expected_batches) <= set(replay_batch_sizes) + + +def graph_batch_exercise_order(expected_batches: tuple[int, ...]) -> tuple[int, ...]: + if len(set(expected_batches)) != len(expected_batches): + raise ValueError("expected_batches must contain each graph key exactly once") + return tuple(sorted(expected_batches, reverse=True)) + + +def parse_replay_batch_sizes(log_text: str) -> set[int]: + return {int(match.group("batch_size")) for match in GRAPH_REPLAY_PATTERN.finditer(log_text)} + + +def parse_eager_decode_batch_sizes(log_text: str) -> set[int]: + return {int(match.group("batch_size")) for match in EAGER_DECODE_PATTERN.finditer(log_text)} + + +def eager_fallback_observed( + eager_batch_sizes: list[int] | set[int], + eager_fallback_batch: int, +) -> bool: + return any(batch_size >= eager_fallback_batch for batch_size in eager_batch_sizes) + + +def build_sampling_params( + *, + max_new_tokens: int, + min_new_tokens: int | None = None, + ignore_eos: bool | None = None, +) -> dict[str, float | int | bool]: + params: dict[str, float | int | bool] = { + "temperature": 0.0, + "max_new_tokens": max_new_tokens, + } + if min_new_tokens is not None: + params["min_new_tokens"] = min_new_tokens + if ignore_eos is not None: + params["ignore_eos"] = ignore_eos + return params + + +def eager_fallback_sampling_params() -> dict[str, float | int | bool]: + return build_sampling_params( + max_new_tokens=EAGER_FALLBACK_MAX_NEW_TOKENS, + min_new_tokens=EAGER_FALLBACK_MAX_NEW_TOKENS, + ignore_eos=True, + ) + + +def graph_batch_exercise_sampling_params() -> dict[str, float | int | bool]: + return build_sampling_params( + max_new_tokens=GRAPH_BATCH_EXERCISE_MAX_NEW_TOKENS, + min_new_tokens=GRAPH_BATCH_EXERCISE_MAX_NEW_TOKENS, + ignore_eos=True, + ) + + +def build_generate_request_payload( + prompts: list[str], + *, + max_new_tokens: int, + min_new_tokens: int | None = None, + ignore_eos: bool | None = None, +) -> dict[str, object]: + return { + "text": prompts, + "sampling_params": build_sampling_params( + max_new_tokens=max_new_tokens, + min_new_tokens=min_new_tokens, + ignore_eos=ignore_eos, + ), + } + + +def parse_generate_list_response(result: object, *, expected_count: int) -> list[str]: + if isinstance(result, list): + outputs = [item["text"] if isinstance(item, dict) else str(item) for item in result] + elif isinstance(result, dict): + text = result.get("text") + if isinstance(text, list): + outputs = [str(item) for item in text] + elif text is not None: + if expected_count != 1: + raise RuntimeError( + f"Batched generate returned one output, expected {expected_count}" + ) + outputs = [str(text)] + else: + raise RuntimeError("Batched generate response missing text field") + else: + raise RuntimeError(f"Unexpected generate response type: {type(result)!r}") + + if len(outputs) != expected_count: + raise RuntimeError( + f"Batched generate returned {len(outputs)} outputs, expected {expected_count}" + ) + if not all(output.strip() for output in outputs): + raise RuntimeError("Batched generate returned empty output") + return outputs + + +def canonical_symmetric_memory_state_payload(payload: dict) -> dict: + state = dict(payload["state"]) + for field in SYMMETRIC_MEMORY_POINTER_FIELDS: + state.pop(field, None) + return { + "version": payload["version"], + "backend": payload["backend"], + "state": state, + } + + +def semantic_symmetric_memory_state_fingerprint(payload: dict) -> str: + canonical = json.dumps( + canonical_symmetric_memory_state_payload(payload), + sort_keys=True, + separators=(",", ":"), + ) + return hashlib.sha256(canonical.encode()).hexdigest() + + +def canonical_shape_state_payload(payload: dict) -> dict: + manifest = payload.get("manifest", {}) + shapes = [] + for shape in manifest.get("shapes", []): + operand_slots = [ + {key: value for key, value in slot.items() if key != "saved_value"} + for slot in shape.get("operand_slots", []) + ] + shapes.append( + { + key: value + for key, value in shape.items() + if key not in {"operand_slots", "saved_value"} + } + | {"operand_slots": operand_slots} + ) + return { + "version": payload["version"], + "schema_fingerprint": payload["schema_fingerprint"], + "backend": manifest.get("backend"), + "rank": manifest.get("rank"), + "world_size": manifest.get("world_size"), + "capture_order": manifest.get("capture_order"), + "shapes": shapes, + } + + +def semantic_shape_state_fingerprint(payload: dict) -> str: + canonical = json.dumps( + canonical_shape_state_payload(payload), + sort_keys=True, + separators=(",", ":"), + ) + return hashlib.sha256(canonical.encode()).hexdigest() + + +def should_skip_run_cleanup(*, keep_artifacts: bool, run_id_override: str | None) -> bool: + return keep_artifacts or run_id_override is not None + + +def cleanup_run_artifacts(run_root: Path) -> None: + report_path = run_root / VALIDATION_REPORT_FILENAME + preserved_report = report_path.read_text() if report_path.exists() else None + shutil.rmtree(run_root, ignore_errors=True) + if preserved_report is not None: + run_root.mkdir(parents=True, exist_ok=True) + report_path.write_text(preserved_report) + + +EXPECTED_GRAPH_BATCHES = _expected_graph_batches() +EXPECTED_GRAPH_COUNT = len(EXPECTED_GRAPH_BATCHES) +EAGER_FALLBACK_BATCH = max(EXPECTED_GRAPH_BATCHES) + 1 +EAGER_FALLBACK_MAX_NEW_TOKENS = 4 +EAGER_FALLBACK_PROMPTS = ["Answer with the number 11."] * EAGER_FALLBACK_BATCH +EXPECTED_NCCL_CHANNELS = 24 +KEEP_ARTIFACTS = os.environ.get("TP_KEEP_ARTIFACTS", "0") == "1" + +ERROR_PATTERN = re.compile( + r"\[(?:HOOK|REPLAY)\] ERROR|Traceback|MMU fault|Xid|CUDA error|" + r"illegal memory access|an illegal memory|NCCL WARN|NCCL error|" + r"Scheduler hit an exception|segmentation fault", + re.IGNORECASE, +) +LOAD_OFFSET_PATTERN = re.compile( + r"TP(?P\d+).*alloc_offset\[after_load_all_graphs\]=" + r"(?P\d+)" +) +LOADED_GRAPHS_PATTERN = re.compile(r"TP(?P\d+).*Loaded (?P\d+) SGLang graphs") +NCCL_TRANSPORT_PATTERN = re.compile( + r"\[(?P\d+)\] NCCL INFO Channel (?P\d+)/\d+ : " + r".* via (?P\S+)" +) +GRAPH_REPLAY_PATTERN = re.compile(r"Decode graph replay: .*? key_size=(?P\d+) \(bs\)") +EAGER_DECODE_PATTERN = re.compile( + r"Decode batch, #running-req: (?P\d+).*cuda graph: False" +) +SYMMETRIC_MEMORY_POINTER_FIELDS = ( + "buffer", + "buffer_ptrs_dev", + "signal_pad_ptrs_dev", + "multicast_ptr", +) +VALIDATION_REPORT_FILENAME = "validation_report.json" + +image = ( + modal.Image.from_registry(f"nvidia/cuda:{CUDA_TAG}", add_python="3.12") + .env({"CC": "gcc", "CXX": "g++"}) + .apt_install( + "git", + "build-essential", + "cmake", + "ninja-build", + "libboost-all-dev", + "libnuma-dev", + ) + .pip_install( + "cmake>=4.0.0", + "wheel", + "packaging", + "ninja", + "setuptools>=80", + "setuptools-scm", + ) + .pip_install( + "torch==2.11.0", + "torchvision==0.26.0", + "torchaudio==2.11.0", + index_url="https://download.pytorch.org/whl/cu130", + ) + .run_commands( + f"git clone {SGLANG_REPO} /sglang", + f"cd /sglang && git checkout {SGLANG_REF}", + SGLANG_PATCH_COMMAND, + "cd /sglang && pip install -e 'python[all]' --no-build-isolation", + ) + .run_commands( + f"git clone {FOUNDRY_REPO} /foundry", + f"cd /foundry && git checkout {FOUNDRY_REF}", + "cd /foundry && pip install -e . --no-build-isolation", + ) + .env( + { + "HF_HOME": HF_CACHE, + "HUGGINGFACE_HUB_CACHE": f"{HF_CACHE}/hub", + "PYTHONUNBUFFERED": "1", + "FOUNDRY_REF": FOUNDRY_REF, + "SGLANG_REF": SGLANG_REF, + "SGLANG_MODEL": MODEL, + "TP_SIZE": str(TP_SIZE), + **SERVE_ENV, + } + ) +) + +app = modal.App("foundry-sglang-tp-validation") +archive_volume = modal.Volume.from_name( + os.environ.get( + "TP_ARCHIVE_VOLUME", + "foundry-sglang-tp-validation-archive", + ), + create_if_missing=True, +) +hf_volume = modal.Volume.from_name( + os.environ.get("TP_HF_VOLUME", "foundry-sglang-tp-hf-cache"), + create_if_missing=True, +) + + +def _recipe_command(mode: str | None) -> list[str]: + command = ["bash", RECIPE_SCRIPT, str(TP_SIZE)] + if mode is not None: + command.append(f"--{mode}") + return command + + +def _log_tail(log_path: str, lines: int = 80) -> str: + path = Path(log_path) + if not path.exists(): + return "" + return "\n".join(path.read_text(errors="replace").splitlines()[-lines:]) + + +def _wait_until_healthy( + process: subprocess.Popen, + log_path: str, + timeout: float, +) -> None: + url = f"http://127.0.0.1:{PORT}/health_generate" + deadline = time.time() + timeout + last_error = "" + + while time.time() < deadline: + if process.poll() is not None: + raise RuntimeError( + f"SGLang exited with {process.returncode} before becoming healthy\n" + f"{_log_tail(log_path)}" + ) + try: + with urllib.request.urlopen(url, timeout=5) as response: + if response.status == 200: + return + except Exception as error: # noqa: BLE001 - preserve the last health error + last_error = str(error) + time.sleep(3) + + raise RuntimeError( + f"SGLang was not healthy after {timeout}s: {last_error}\n{_log_tail(log_path)}" + ) + + +def _generate_batched( + prompts: list[str], + *, + max_new_tokens: int, + min_new_tokens: int | None = None, + ignore_eos: bool | None = None, +) -> list[str]: + payload = build_generate_request_payload( + prompts, + max_new_tokens=max_new_tokens, + min_new_tokens=min_new_tokens, + ignore_eos=ignore_eos, + ) + request = urllib.request.Request( + f"http://127.0.0.1:{PORT}/generate", + data=json.dumps(payload).encode(), + headers={"Content-Type": "application/json"}, + ) + with urllib.request.urlopen(request, timeout=120) as response: + result = json.loads(response.read()) + return parse_generate_list_response(result, expected_count=len(prompts)) + + +def _generate( + prompt: str, + *, + max_new_tokens: int = MAX_NEW_TOKENS, + min_new_tokens: int | None = None, + ignore_eos: bool | None = None, +) -> str: + request = urllib.request.Request( + f"http://127.0.0.1:{PORT}/generate", + data=json.dumps( + { + "text": prompt, + "sampling_params": build_sampling_params( + max_new_tokens=max_new_tokens, + min_new_tokens=min_new_tokens, + ignore_eos=ignore_eos, + ), + } + ).encode(), + headers={"Content-Type": "application/json"}, + ) + with urllib.request.urlopen(request, timeout=120) as response: + result = json.loads(response.read()) + return result["text"] if isinstance(result, dict) else result[0]["text"] + + +def _exercise_graph_batch(batch_size: int) -> None: + prompts = [GRAPH_BATCH_EXERCISE_PROMPT] * batch_size + _generate_batched( + prompts, + max_new_tokens=GRAPH_BATCH_EXERCISE_MAX_NEW_TOKENS, + min_new_tokens=GRAPH_BATCH_EXERCISE_MAX_NEW_TOKENS, + ignore_eos=True, + ) + + +def _generate_eager_fallback() -> list[str]: + return _generate_batched( + EAGER_FALLBACK_PROMPTS, + max_new_tokens=EAGER_FALLBACK_MAX_NEW_TOKENS, + min_new_tokens=EAGER_FALLBACK_MAX_NEW_TOKENS, + ignore_eos=True, + ) + + +def _generate_concurrently() -> list[str]: + barrier = threading.Barrier(len(CONCURRENT_PROMPTS)) + + def generate_after_barrier(prompt: str) -> str: + barrier.wait() + return _generate(prompt) + + with ThreadPoolExecutor(max_workers=len(CONCURRENT_PROMPTS)) as executor: + return list(executor.map(generate_after_barrier, CONCURRENT_PROMPTS)) + + +def _launch( + mode: str | None, + log_path: str, + run_root: Path, +) -> tuple[subprocess.Popen, IO[str]]: + env = dict(os.environ) + env["SGLANG_MODEL"] = MODEL + env.update(SERVE_ENV) + env["NCCL_DEBUG"] = "INFO" + env["SGLANG_LOG_DECODE_GRAPH_KEY"] = "1" + + # The handle intentionally spans the child lifetime and is closed by _stop. + log_file = open(log_path, "w") # noqa: SIM115 + process = subprocess.Popen( + _recipe_command(mode), + stdout=log_file, + stderr=subprocess.STDOUT, + env=env, + cwd=run_root, + preexec_fn=os.setsid, + ) + return process, log_file + + +def _stop(process: subprocess.Popen, log_file: IO[str]) -> None: + try: + os.killpg(os.getpgid(process.pid), signal.SIGINT) + process.wait(timeout=60) + except Exception: # noqa: BLE001 - best-effort teardown after validation + with contextlib.suppress(Exception): + os.killpg(os.getpgid(process.pid), signal.SIGKILL) + finally: + log_file.close() + + +def _read_rank_offsets(workspace: Path) -> dict[str, int]: + offsets = {} + for rank in range(TP_SIZE): + path = workspace / f"rank_{rank}" / "final_alloc_offset.json" + if not path.exists(): + continue + offsets[str(rank)] = int(json.loads(path.read_text())["final_alloc_offset"]) + return offsets + + +def _archive_graph_counts(workspace: Path) -> dict[str, int]: + counts = {} + for rank in range(TP_SIZE): + rank_dir = workspace / f"rank_{rank}" + counts[str(rank)] = len(list(rank_dir.glob("graph_*.cugraph"))) + return counts + + +def _sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as file: + for chunk in iter(lambda: file.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def _semantic_symmetric_memory_state_sha256(path: Path) -> str: + payload = json.loads(path.read_text()) + return semantic_symmetric_memory_state_fingerprint(payload) + + +def _semantic_shape_state_sha256(path: Path) -> str: + payload = json.loads(path.read_text()) + return semantic_shape_state_fingerprint(payload) + + +def _archive_file_fingerprint(path: Path) -> str: + if path.name.startswith("graph_") and path.name != "graph_manifest.json": + return _semantic_graph_sha256(path) + if path.name == "symmetric_memory_state.json": + return _semantic_symmetric_memory_state_sha256(path) + if path.name == "sglang_shape_state.json": + return _semantic_shape_state_sha256(path) + return _sha256(path) + + +def _semantic_graph_sha256(path: Path) -> str: + graph = json.loads(path.read_text()) + # SGLang chooses a fresh server seed on each launch. Generator identity and + # seed may therefore differ even when the captured topology, kernel params, + # tensor addresses, and allocator events are identical. + for generator in graph.get("generators", []): + generator.pop("id", None) + generator.pop("seed", None) + vmm_kernel_pointers = [] + for node in graph.get("nodes", []): + if node.get("type") != "KernelNode": + continue + params = node["params"] + encoded_params = params.pop("kernelParams", []) + encoded_arg_buffer = params.pop("extra_argBuffer_hex", "") + # `extra` contains launch-API host pointers. Its shape is stable, but + # the pointer values are process-local and never consumed by the GPU. + params.pop("extra", None) + encoded_values = [ + (f"param:{param['index']}", param.get("value_hex", "")) for param in encoded_params + ] + encoded_values.append(("arg_buffer", encoded_arg_buffer)) + for source, encoded in encoded_values: + raw = bytes.fromhex(encoded) + for offset in range(0, max(0, len(raw) - 7), 8): + value = int.from_bytes(raw[offset : offset + 8], "little") + if FOUNDRY_BASE_ADDR <= value < FOUNDRY_REGION_END: + vmm_kernel_pointers.append([node["id"], source, offset, value]) + graph["vmm_kernel_pointers"] = vmm_kernel_pointers + canonical = json.dumps(graph, sort_keys=True, separators=(",", ":")).encode() + return hashlib.sha256(canonical).hexdigest() + + +def _archive_fingerprints(workspace: Path) -> dict[str, dict[str, str]]: + fingerprints = {} + for rank in range(TP_SIZE): + rank_dir = workspace / f"rank_{rank}" + files = [*rank_dir.glob("graph_*.json")] + files.extend( + path + for path in ( + rank_dir / "graph_manifest.json", + rank_dir / "fatbin_image_packed.img", + rank_dir / "final_alloc_offset.json", + rank_dir / "sglang_graph_backend.json", + rank_dir / "symmetric_memory_state.json", + rank_dir / "sglang_shape_state.json", + ) + if path.exists() + ) + fingerprints[str(rank)] = { + path.relative_to(rank_dir).as_posix(): _archive_file_fingerprint(path) + for path in sorted(files) + } + return fingerprints + + +def _wait_for_replay_batches( + process: subprocess.Popen, + log_path: str, + expected_batches: tuple[int, ...], + timeout: float, +) -> None: + deadline = time.time() + timeout + expected = set(expected_batches) + while time.time() < deadline: + if process.poll() is not None: + raise RuntimeError( + f"SGLang exited with {process.returncode} before replay batches were observed\n" + f"{_log_tail(log_path)}" + ) + path = Path(log_path) + if path.exists(): + observed = parse_replay_batch_sizes(path.read_text(errors="replace")) + if expected <= observed: + return + time.sleep(0.5) + + path = Path(log_path) + observed = ( + parse_replay_batch_sizes(path.read_text(errors="replace")) if path.exists() else set() + ) + missing = expected - observed + raise RuntimeError( + "Timed out waiting for graph replay batches. " + f"expected={sorted(expected)} observed={sorted(observed)} missing={sorted(missing)}\n" + f"{_log_tail(log_path)}" + ) + + +def _wait_for_eager_fallback( + process: subprocess.Popen, + log_path: str, + eager_fallback_batch: int, + timeout: float, +) -> None: + deadline = time.time() + timeout + while time.time() < deadline: + if process.poll() is not None: + path = Path(log_path) + observed = ( + parse_eager_decode_batch_sizes(path.read_text(errors="replace")) + if path.exists() + else set() + ) + satisfying = sorted(batch for batch in observed if batch >= eager_fallback_batch) + raise RuntimeError( + "SGLang exited before eager fallback was observed. " + f"required_batch>={eager_fallback_batch} observed={sorted(observed)} " + f"satisfying={satisfying}\n" + f"{_log_tail(log_path)}" + ) + path = Path(log_path) + if path.exists(): + observed = parse_eager_decode_batch_sizes(path.read_text(errors="replace")) + if eager_fallback_observed(observed, eager_fallback_batch): + return + time.sleep(0.5) + + path = Path(log_path) + observed = ( + parse_eager_decode_batch_sizes(path.read_text(errors="replace")) if path.exists() else set() + ) + satisfying = sorted(batch for batch in observed if batch >= eager_fallback_batch) + raise RuntimeError( + "Timed out waiting for eager fallback decode batch. " + f"required_batch>={eager_fallback_batch} observed={sorted(observed)} " + f"satisfying={satisfying}\n" + f"{_log_tail(log_path)}" + ) + + +def _scan_log(log_path: str) -> dict: + text = Path(log_path).read_text(errors="replace") + errors = sorted({match.group(0) for match in ERROR_PATTERN.finditer(text)}) + load_offsets = { + match.group("rank"): int(match.group("offset")) + for match in LOAD_OFFSET_PATTERN.finditer(text) + } + loaded_graphs = { + match.group("rank"): int(match.group("count")) + for match in LOADED_GRAPHS_PATTERN.finditer(text) + } + transport_channels: dict[str, dict[str, set[int]]] = {} + for match in NCCL_TRANSPORT_PATTERN.finditer(text): + rank_transports = transport_channels.setdefault(match.group("rank"), {}) + rank_transports.setdefault(match.group("transport"), set()).add(int(match.group("channel"))) + graph_replay_batch_sizes = sorted(parse_replay_batch_sizes(text)) + eager_decode_batch_sizes = sorted(parse_eager_decode_batch_sizes(text)) + return { + "errors": errors, + "load_offsets": load_offsets, + "loaded_graphs": loaded_graphs, + "nccl_transport_channels": { + rank: { + transport: sorted(channels) + for transport, channels in sorted(rank_transports.items()) + } + for rank, rank_transports in sorted(transport_channels.items()) + }, + "graph_replay_batch_sizes": graph_replay_batch_sizes, + "eager_decode_batch_sizes": eager_decode_batch_sizes, + "native_capture_progress_count": text.count("Capturing batches"), + "torch_symm_mem_observed": ("communication backend=torch_symmetric_memory" in text), + "saved_graph_log_count": text.count("[Foundry] Saved SGLang CUDA graph") + + text.count("[Foundry] Saved SGLang-main CUDA graph"), + } + + +@app.function( + image=image, + gpu=MODAL_GPU, + timeout=3600, + volumes={ARCHIVE_ROOT: archive_volume, HF_CACHE: hf_volume}, +) +def run_phase(phase: str, run_id: str) -> dict: + mode = { + "baseline": None, + "save": "save", + "save2": "save", + "load": "load", + }[phase] + run_root = Path(RUNS_ROOT) / run_id + workspace = run_root / "foundry_archive_sglang_tp" + run_root.mkdir(parents=True, exist_ok=True) + log_path = str(run_root / f"{phase}.log") + print(f"[TP validation] run={run_id} phase={phase} mode={mode} gpu={MODAL_GPU}") + + if phase == "save": + shutil.rmtree(workspace, ignore_errors=True) + + started_at = time.time() + process, log_file = _launch(mode, log_path, run_root) + result = {"phase": phase, "artifact_dir": str(run_root)} + runtime_scan = None + try: + _wait_until_healthy(process, log_path, timeout=1800) + result["time_to_healthy_s"] = round(time.time() - started_at, 1) + result["outputs"] = [_generate(prompt) for prompt in PROMPTS] + if phase in {"save", "save2", "load"}: + for batch_size in graph_batch_exercise_order(EXPECTED_GRAPH_BATCHES): + _exercise_graph_batch(batch_size) + if phase == "load": + log_file.flush() + _wait_for_replay_batches( + process, + log_path, + EXPECTED_GRAPH_BATCHES, + timeout=300, + ) + if phase in {"baseline", "load"}: + result["concurrent_outputs"] = _generate_concurrently() + result["eager_fallback_outputs"] = _generate_eager_fallback() + if phase == "load": + log_file.flush() + _wait_for_eager_fallback( + process, + log_path, + EAGER_FALLBACK_BATCH, + timeout=300, + ) + result["post_concurrent_outputs"] = [_generate(prompt) for prompt in PROMPTS] + if process.poll() is not None: + raise RuntimeError( + f"SGLang exited unexpectedly after validation requests\n{_log_tail(log_path)}" + ) + log_file.flush() + runtime_scan = _scan_log(log_path) + finally: + _stop(process, log_file) + + if runtime_scan is None: + raise RuntimeError(f"Runtime log scan was not captured\n{_log_tail(log_path)}") + result.update(runtime_scan) + if phase in {"save", "save2"}: + result["rank_offsets"] = _read_rank_offsets(workspace) + result["archive_graph_counts"] = _archive_graph_counts(workspace) + result["archive_fingerprints"] = _archive_fingerprints(workspace) + + archive_volume.commit() + print(json.dumps(result, sort_keys=True)) + return result + + +@app.function( + image=modal.Image.debian_slim(), + volumes={ARCHIVE_ROOT: archive_volume}, +) +def cleanup_run(run_id: str) -> None: + cleanup_run_artifacts(Path(RUNS_ROOT) / run_id) + archive_volume.commit() + + +@app.function( + image=modal.Image.debian_slim(), + volumes={ARCHIVE_ROOT: archive_volume}, +) +def persist_validation_report(run_id: str, report: dict) -> None: + report_path = Path(RUNS_ROOT) / run_id / VALIDATION_REPORT_FILENAME + report_path.parent.mkdir(parents=True, exist_ok=True) + report_path.write_text(json.dumps(report, indent=2, sort_keys=True)) + archive_volume.commit() + + +@app.local_entrypoint() +def main() -> None: + safe_ref = re.sub(r"[^A-Za-z0-9_.-]", "-", FOUNDRY_REF)[:24] + run_id = RUN_ID_OVERRIDE or f"{safe_ref}-{time.time_ns()}" + phases = ("baseline", "save", "save2", "load") + results = {phase: run_phase.remote(phase, run_id) for phase in phases} + + expected_ranks = {str(rank) for rank in range(TP_SIZE)} + expected_transport_channels = { + str(rank): {"P2P/IPC": list(range(EXPECTED_NCCL_CHANNELS))} for rank in range(TP_SIZE) + } + required_archive_files_set = required_archive_files( + uses_torch_symm_mem=USES_TORCH_SYMM_MEM, + uses_multi_shape=USES_MULTI_SHAPE, + ) + baseline_outputs = results["baseline"].get("outputs", []) + save_outputs = results["save"].get("outputs", []) + save2_outputs = results["save2"].get("outputs", []) + loaded_outputs = results["load"].get("outputs", []) + baseline_concurrent_outputs = results["baseline"].get("concurrent_outputs", []) + loaded_concurrent_outputs = results["load"].get("concurrent_outputs", []) + baseline_post_concurrent_outputs = results["baseline"].get( + "post_concurrent_outputs", + [], + ) + loaded_post_concurrent_outputs = results["load"].get( + "post_concurrent_outputs", + [], + ) + save_offsets = results["save"].get("rank_offsets", {}) + save2_offsets = results["save2"].get("rank_offsets", {}) + load_offsets = results["load"].get("load_offsets", {}) + save_fingerprints = results["save"].get("archive_fingerprints", {}) + save2_fingerprints = results["save2"].get("archive_fingerprints", {}) + save_graph_counts = results["save"].get("archive_graph_counts", {}) + graph_counts = results["save2"].get("archive_graph_counts", {}) + loaded_graphs = results["load"].get("loaded_graphs", {}) + replay_batch_sizes = results["load"].get("graph_replay_batch_sizes", []) + replayed_graph_batches = set(replay_batch_sizes) + eager_decode_batch_sizes = results["load"].get("eager_decode_batch_sizes", []) + + checks = { + "deterministic_outputs": len(baseline_outputs) == len(PROMPTS) + and all(output.strip() for output in baseline_outputs) + and baseline_outputs == save_outputs == save2_outputs == loaded_outputs, + "concurrent_outputs_nonempty": len(baseline_concurrent_outputs) == len(CONCURRENT_PROMPTS) + and all(output.strip() for output in baseline_concurrent_outputs) + and len(loaded_concurrent_outputs) == len(CONCURRENT_PROMPTS) + and all(output.strip() for output in loaded_concurrent_outputs), + "post_concurrent_graph_outputs_match": ( + baseline_post_concurrent_outputs + == loaded_post_concurrent_outputs + == baseline_outputs + == loaded_outputs + ), + "save_offsets_present": set(save_offsets) == expected_ranks and all(save_offsets.values()), + "save_rank_offsets_symmetric": len(set(save_offsets.values())) == 1, + "save_offsets_reproducible": save_offsets == save2_offsets, + "load_offsets_match_save": load_offsets == save2_offsets, + "archive_fingerprints_present": set(save_fingerprints) == expected_ranks + and all( + required_archive_files_set <= set(rank_fingerprints) + and sum( + name.startswith("graph_") and name.endswith(".json") for name in rank_fingerprints + ) + == EXPECTED_GRAPH_COUNT + 1 + for rank_fingerprints in save_fingerprints.values() + ), + "archive_fingerprints_reproducible": save_fingerprints == save2_fingerprints, + "archives_complete": set(graph_counts) == expected_ranks + and save_graph_counts == graph_counts + and all(count == EXPECTED_GRAPH_COUNT for count in graph_counts.values()), + "graphs_restored_each_rank": loaded_graphs == graph_counts, + "restored_graph_replay_observed": restored_graph_replay_observed( + EXPECTED_GRAPH_BATCHES, + replayed_graph_batches, + ), + "eager_fallback_observed": eager_fallback_observed( + eager_decode_batch_sizes, + EAGER_FALLBACK_BATCH, + ), + "load_did_not_capture": results["load"]["saved_graph_log_count"] == 0 + and (USES_MAIN_PLUGIN or results["load"]["native_capture_progress_count"] == 0), + "communication_backend_observed": ( + results["save"]["torch_symm_mem_observed"] + and results["load"]["torch_symm_mem_observed"] + if USES_TORCH_SYMM_MEM + else all( + results[phase]["nccl_transport_channels"] == expected_transport_channels + for phase in phases + ) + ), + "no_runtime_errors": not any(results[phase]["errors"] for phase in phases), + } + + report = { + "run_id": run_id, + "artifacts_retained": KEEP_ARTIFACTS, + "foundry_ref": FOUNDRY_REF, + "sglang_ref": SGLANG_REF, + "gpu": MODAL_GPU, + "expected_graph_batches": list(EXPECTED_GRAPH_BATCHES), + "eager_fallback_batch": EAGER_FALLBACK_BATCH, + "observed_eager_decode_batch_sizes": eager_decode_batch_sizes, + "checks": checks, + "observations": { + "concurrent_outputs_byte_identical": ( + baseline_concurrent_outputs == loaded_concurrent_outputs + ), + }, + "results": results, + } + persist_validation_report.remote(run_id, report) + print(json.dumps(report, indent=2, sort_keys=True)) + + failed = [name for name, passed in checks.items() if not passed] + if failed: + raise RuntimeError(f"TP validation failed: {', '.join(failed)}") + if not should_skip_run_cleanup( + keep_artifacts=KEEP_ARTIFACTS, + run_id_override=RUN_ID_OVERRIDE, + ): + cleanup_run.remote(run_id) diff --git a/tests/modal_vllm_tp.py b/tests/modal_vllm_tp.py new file mode 100644 index 00000000..e5a5b550 --- /dev/null +++ b/tests/modal_vllm_tp.py @@ -0,0 +1,430 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the Foundry project +"""End-to-end vLLM tensor-parallel validation on Modal. + +The harness runs five phases on two GPUs: + +1. A non-Foundry vLLM TP baseline using the same pinned collective settings. +2. A seed SAVE that records vLLM's memory-profile warmup state. +3. A deterministic SAVE using that warmup state. +4. A second deterministic SAVE to prove offset and inventory reproducibility. +5. Foundry LOAD in a fresh server process and the same deterministic queries. + +Success requires byte-identical decoded baseline/LOAD responses, reproducible +graph inventories and per-rank allocation offsets, restored graphs on every +rank, NCCL P2P/IPC in every phase, no graph recapture on LOAD, and no CUDA/NCCL +errors. + +Run the current checkout: + + modal run --env rahul-dev tests/modal_vllm_tp.py + +Run a specific commit: + + FOUNDRY_REF= \ + modal run --env rahul-dev tests/modal_vllm_tp.py + +Both Foundry and vLLM resolve to immutable commits before the remote image is +built so Modal cannot reuse a stale image for a moving branch. +""" + +from __future__ import annotations + +import contextlib +import json +import os +import re +import shutil +import signal +import subprocess +import threading +import time +import urllib.request +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path +from typing import IO + +import modal + + +def _local_foundry_revision() -> str: + try: + return subprocess.check_output( + ["git", "rev-parse", "HEAD"], + cwd=Path(__file__).parents[1], + stderr=subprocess.DEVNULL, + text=True, + ).strip() + except (OSError, subprocess.CalledProcessError): + return "main" + + +CUDA_TAG = "13.0.1-cudnn-devel-ubuntu24.04" + +FOUNDRY_REPO = "https://github.com/modal-projects/foundry.git" +FOUNDRY_REF = os.environ.get("FOUNDRY_REF") or _local_foundry_revision() + +VLLM_REPO = "https://github.com/foundry-org/vllm.git" +VLLM_REF = os.environ.get("VLLM_REF", "4309c257d3f639e5490d3811293c890c61c76f29") +# The Foundry fork commit is one Python-only commit on top of this vLLM base. +# Pin its matching CUDA wheel: falling back to the moving nightly wheel can +# silently omit extensions that this older source tree imports as ``vllm._C``. +VLLM_WHEEL_COMMIT = "6cbe448eed751824d608faf9078ef84724d621c1" + +MODEL = os.environ.get("VLLM_MODEL", "Qwen/Qwen3-8B") +TP_SIZE = int(os.environ.get("TP_SIZE", "2")) +if TP_SIZE < 2: + raise ValueError("Tensor-parallel validation requires TP_SIZE >= 2") +MODAL_GPU = os.environ.get("MODAL_GPU", f"H100:{TP_SIZE}") +PORT = 12000 +EXPECTED_GRAPH_COUNT = 64 + +ARCHIVE_ROOT = "/data/archive" +RUNS_ROOT = f"{ARCHIVE_ROOT}/runs" +HF_CACHE = "/data/hf" +RECIPE_SCRIPT = "/foundry/recipe/vllm/serve_qwen3-8b_tp.sh" + +PROMPTS = [ + "The capital of France is", + "Reply with only the result: 2 + 2 =", + "The chemical formula for water is", + "A GPU is a graphics processing", +] +MAX_TOKENS = 4 + +ERROR_PATTERN = re.compile( + r"\[HOOK\] ERROR|MMU fault|Xid|CUDA error|illegal memory access|" + r"an illegal memory|NCCL WARN|NCCL error|Engine core proc died|" + r"segmentation fault", + re.IGNORECASE, +) +EARLY_BUILDS_PATTERN = re.compile( + r"\[foundry\] Starting early graph builds \((?P\d+) graphs" +) +NCCL_TRANSPORT_PATTERN = re.compile( + r"\[(?P\d+)\] NCCL INFO Channel \d+/\d+ : .* via (?P\S+)" +) + +image = ( + modal.Image.from_registry(f"nvidia/cuda:{CUDA_TAG}", add_python="3.12") + .env({"CC": "gcc", "CXX": "g++"}) + .apt_install( + "git", + "build-essential", + "cmake", + "ninja-build", + "libboost-all-dev", + "libnuma-dev", + ) + .pip_install( + "cmake>=4.0.0", + "wheel", + "packaging", + "ninja", + "setuptools>=80", + "setuptools-scm", + "uv", + ) + .pip_install( + "torch==2.11.0", + "torchvision==0.26.0", + "torchaudio==2.11.0", + index_url="https://download.pytorch.org/whl/cu130", + ) + .run_commands( + f"git clone {VLLM_REPO} /vllm", + f"cd /vllm && git checkout {VLLM_REF}", + "cd /vllm && VLLM_USE_PRECOMPILED=1 " + f"VLLM_PRECOMPILED_WHEEL_COMMIT={VLLM_WHEEL_COMMIT} " + "VLLM_PRECOMPILED_WHEEL_VARIANT=cu130 " + "uv pip install --system --editable . " + "--extra-index-url https://wheels.vllm.ai/nightly/cu130", + "test -f /vllm/vllm/_C.abi3.so", + ) + .run_commands( + f"git clone {FOUNDRY_REPO} /foundry", + f"cd /foundry && git checkout {FOUNDRY_REF}", + "cd /foundry && pip install -e . --no-build-isolation", + ) + .env( + { + "HF_HOME": HF_CACHE, + "HUGGINGFACE_HUB_CACHE": f"{HF_CACHE}/hub", + "PYTHONUNBUFFERED": "1", + } + ) +) + +app = modal.App("foundry-vllm-tp-validation") +archive_volume = modal.Volume.from_name( + os.environ.get("VLLM_TP_ARCHIVE_VOLUME", "foundry-vllm-tp-validation-archive"), + create_if_missing=True, +) +hf_volume = modal.Volume.from_name( + os.environ.get("VLLM_TP_HF_VOLUME", "foundry-vllm-tp-hf-cache"), + create_if_missing=True, +) + + +def _recipe_command(mode: str | None) -> list[str]: + command = ["bash", RECIPE_SCRIPT, str(TP_SIZE)] + if mode is not None: + command.append(f"--{mode}") + return command + + +def _log_tail(log_path: str, lines: int = 100) -> str: + path = Path(log_path) + if not path.exists(): + return "" + return "\n".join(path.read_text(errors="replace").splitlines()[-lines:]) + + +def _wait_until_healthy( + process: subprocess.Popen, + log_path: str, + timeout: float, +) -> None: + url = f"http://127.0.0.1:{PORT}/health" + deadline = time.time() + timeout + last_error = "" + + while time.time() < deadline: + if process.poll() is not None: + raise RuntimeError( + f"vLLM exited with {process.returncode} before becoming healthy\n" + f"{_log_tail(log_path)}" + ) + try: + with urllib.request.urlopen(url, timeout=5) as response: + if response.status == 200: + return + except Exception as error: # noqa: BLE001 - preserve the last health error + last_error = str(error) + time.sleep(3) + + raise RuntimeError( + f"vLLM was not healthy after {timeout}s: {last_error}\n{_log_tail(log_path)}" + ) + + +def _generate(prompt: str) -> str: + request = urllib.request.Request( + f"http://127.0.0.1:{PORT}/v1/completions", + data=json.dumps( + { + "model": MODEL, + "prompt": prompt, + "max_tokens": MAX_TOKENS, + "temperature": 0.0, + "seed": 0, + "stop": ["\n"], + } + ).encode(), + headers={"Content-Type": "application/json"}, + ) + with urllib.request.urlopen(request, timeout=180) as response: + result = json.loads(response.read()) + return result["choices"][0]["text"] + + +def _generate_concurrently() -> list[str]: + barrier = threading.Barrier(len(PROMPTS)) + + def generate_after_barrier(prompt: str) -> str: + barrier.wait() + return _generate(prompt) + + with ThreadPoolExecutor(max_workers=len(PROMPTS)) as executor: + return list(executor.map(generate_after_barrier, PROMPTS)) + + +def _launch( + mode: str | None, + log_path: str, + run_root: Path, +) -> tuple[subprocess.Popen, IO[str]]: + env = dict(os.environ) + env["CUDA_VISIBLE_DEVICES"] = ",".join(str(rank) for rank in range(TP_SIZE)) + env["VLLM_MODEL"] = MODEL + env["VLLM_HOST"] = "127.0.0.1" + env["VLLM_PORT"] = str(PORT) + env["NCCL_DEBUG"] = "INFO" + + # The handle intentionally spans the child lifetime and is closed by _stop. + log_file = open(log_path, "w") # noqa: SIM115 + process = subprocess.Popen( + _recipe_command(mode), + stdout=log_file, + stderr=subprocess.STDOUT, + env=env, + cwd=run_root, + start_new_session=True, + ) + return process, log_file + + +def _stop(process: subprocess.Popen, log_file: IO[str]) -> None: + try: + os.killpg(os.getpgid(process.pid), signal.SIGINT) + process.wait(timeout=120) + except Exception: # noqa: BLE001 - best-effort teardown after validation + with contextlib.suppress(Exception): + os.killpg(os.getpgid(process.pid), signal.SIGKILL) + finally: + log_file.close() + + +def _read_rank_offsets(workspace: Path) -> dict[str, int]: + offsets = {} + for rank in range(TP_SIZE): + path = workspace / f"rank_{rank}" / "final_alloc_offset.json" + if not path.exists(): + continue + offsets[str(rank)] = int(json.loads(path.read_text())["final_alloc_offset"]) + return offsets + + +def _archive_graph_counts(workspace: Path) -> dict[str, int]: + counts = {} + for rank in range(TP_SIZE): + rank_dir = workspace / f"rank_{rank}" + counts[str(rank)] = len(list(rank_dir.glob("graph_*.cugraph"))) + return counts + + +def _archive_inventory(workspace: Path) -> dict[str, list[str]]: + inventories = {} + for rank in range(TP_SIZE): + rank_dir = workspace / f"rank_{rank}" + files = sorted(rank_dir.glob("graph_*.json")) + inventories[str(rank)] = [path.relative_to(rank_dir).as_posix() for path in sorted(files)] + return inventories + + +def _scan_log(log_path: str) -> dict: + text = Path(log_path).read_text(errors="replace") + errors = sorted({match.group(0) for match in ERROR_PATTERN.finditer(text)}) + transports: dict[str, set[str]] = {} + for match in NCCL_TRANSPORT_PATTERN.finditer(text): + transports.setdefault(match.group("rank"), set()).add(match.group("transport")) + return { + "errors": errors, + "early_build_graph_counts": [ + int(match.group("count")) for match in EARLY_BUILDS_PATTERN.finditer(text) + ], + "nccl_transports": { + rank: sorted(rank_transports) for rank, rank_transports in transports.items() + }, + "captured_offset_count": text.count("[foundry] Captured final_alloc_offset:"), + "saved_manifest_count": text.count("[foundry] Saved graph_manifest.json:"), + } + + +@app.function( + image=image, + gpu=MODAL_GPU, + timeout=3600, + volumes={ARCHIVE_ROOT: archive_volume, HF_CACHE: hf_volume}, +) +def run_phase(phase: str, run_id: str) -> dict: + mode = { + "baseline": None, + "seed": "save", + "save": "save", + "save2": "save", + "load": "load", + }[phase] + run_root = Path(RUNS_ROOT) / run_id + workspace = run_root / "foundry_archive" + run_root.mkdir(parents=True, exist_ok=True) + log_path = str(run_root / f"{phase}.log") + print(f"[TP validation] run={run_id} phase={phase} mode={mode} gpu={MODAL_GPU}") + + if phase == "seed": + shutil.rmtree(workspace, ignore_errors=True) + + started_at = time.time() + process, log_file = _launch(mode, log_path, run_root) + result = {"phase": phase, "artifact_dir": str(run_root)} + try: + _wait_until_healthy(process, log_path, timeout=1800) + result["time_to_healthy_s"] = round(time.time() - started_at, 1) + if phase in {"baseline", "load"}: + result["outputs"] = [_generate(prompt) for prompt in PROMPTS] + result["concurrent_outputs"] = _generate_concurrently() + if process.poll() is not None: + raise RuntimeError( + f"vLLM exited unexpectedly after validation requests\n{_log_tail(log_path)}" + ) + finally: + _stop(process, log_file) + + result.update(_scan_log(log_path)) + if phase in {"seed", "save", "save2"}: + result["rank_offsets"] = _read_rank_offsets(workspace) + result["archive_graph_counts"] = _archive_graph_counts(workspace) + result["archive_inventory"] = _archive_inventory(workspace) + + archive_volume.commit() + hf_volume.commit() + print(json.dumps(result, sort_keys=True)) + return result + + +@app.local_entrypoint() +def main() -> None: + safe_ref = re.sub(r"[^A-Za-z0-9_.-]", "-", FOUNDRY_REF)[:24] + run_id = f"{safe_ref}-{time.time_ns()}" + phases = ("baseline", "seed", "save", "save2", "load") + results = {phase: run_phase.remote(phase, run_id) for phase in phases} + + expected_ranks = {str(rank) for rank in range(TP_SIZE)} + expected_transports = {str(rank): ["P2P/IPC"] for rank in range(TP_SIZE)} + baseline_outputs = results["baseline"].get("outputs", []) + loaded_outputs = results["load"].get("outputs", []) + baseline_concurrent_outputs = results["baseline"].get("concurrent_outputs", []) + loaded_concurrent_outputs = results["load"].get("concurrent_outputs", []) + save_offsets = results["save"].get("rank_offsets", {}) + save2_offsets = results["save2"].get("rank_offsets", {}) + save_inventory = results["save"].get("archive_inventory", {}) + save2_inventory = results["save2"].get("archive_inventory", {}) + save_graph_counts = results["save"].get("archive_graph_counts", {}) + save2_graph_counts = results["save2"].get("archive_graph_counts", {}) + loaded_graph_counts = sorted(results["load"].get("early_build_graph_counts", [])) + + checks = { + "deterministic_outputs": bool(baseline_outputs) and baseline_outputs == loaded_outputs, + "deterministic_concurrent_outputs": bool(baseline_concurrent_outputs) + and baseline_concurrent_outputs == loaded_concurrent_outputs, + "save_offsets_present": set(save_offsets) == expected_ranks and all(save_offsets.values()), + "save_rank_offsets_symmetric": len(set(save_offsets.values())) == 1, + "save_offsets_reproducible": save_offsets == save2_offsets, + "archive_inventory_present": set(save_inventory) == expected_ranks + and all(save_inventory.values()), + "archive_inventory_reproducible": save_inventory == save2_inventory, + "archives_complete": set(save2_graph_counts) == expected_ranks + and save_graph_counts == save2_graph_counts + and set(save2_graph_counts.values()) == {EXPECTED_GRAPH_COUNT}, + "graphs_restored_each_rank": loaded_graph_counts == sorted(save2_graph_counts.values()), + "load_did_not_capture": results["load"]["captured_offset_count"] == 0 + and results["load"]["saved_manifest_count"] == 0, + "nccl_p2p_ipc_every_phase": all( + results[phase]["nccl_transports"] == expected_transports for phase in phases + ), + "no_runtime_errors": not any(results[phase]["errors"] for phase in phases), + } + + report = { + "run_id": run_id, + "foundry_ref": FOUNDRY_REF, + "vllm_ref": VLLM_REF, + "gpu": MODAL_GPU, + "checks": checks, + "results": results, + } + print(json.dumps(report, indent=2, sort_keys=True)) + + failed = [name for name, passed in checks.items() if not passed] + if failed: + raise RuntimeError(f"TP validation failed: {', '.join(failed)}") diff --git a/tests/run_deepep_matrix.sh b/tests/run_deepep_matrix.sh new file mode 100755 index 00000000..d2a1cfe3 --- /dev/null +++ b/tests/run_deepep_matrix.sh @@ -0,0 +1,82 @@ +#!/usr/bin/env bash +# Matrix driver for tests/test_deepep_fabric.py โ€” standalone DeepEP save/load probes +# (no vLLM/sglang init; needs ~4 GB free on each of GPUs 0,1, NOT a fully empty GPU). +# +# ll_nofabric pure low-latency (NVSHMEM heap, FD handles), NCCL fast paths off +# -> the "EP without fabric" baseline; expected PASS +# ll_ncclcumem same, but NCCL_CUMEM_ENABLE=1 -> probe: can NCCL cumem coexist? +# ll_ncclnvls same, but CUMEM=1 + NVLS=1 -> probe: can NCCL NVLS coexist? +# nvl_ipc + 64 MB NVL buffer, use_fabric=0 -> legacy cudaIpc path through the +# hook's VMM-IPC translation (SCM_RIGHTS fd transport; peer mappings +# relocate under shared region bases); expected PASS +# nvl_fabric + 64 MB NVL buffer, use_fabric=1 -> fabric cuMemCreate on a no-IMEX +# machine; documents the fabric dependency; expected FAIL +# +# Usage: run_deepep_matrix.sh +set -euo pipefail + +REPO_ROOT=$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd) +PY=${PYTHON:-python3} +TEST=$REPO_ROOT/tests/test_deepep_fabric.py +HOOK_SO=${HOOK_SO:-$("$PY" -c "import foundry.ops, pathlib; print(pathlib.Path(foundry.ops.__file__).parent / 'libcuda_hook.so')")} +NVSHMEM_SO=${NVSHMEM_SO:-$("$PY" -c "from foundry.integration.sglang.config import CUDAGraphExtensionConfig; print(CUDAGraphExtensionConfig._detect_nvshmem_host_path() or '')")} + +if [ ! -f "$HOOK_SO" ]; then + echo "Foundry hook not found: $HOOK_SO"; exit 4 +fi +if [ ! -f "$NVSHMEM_SO" ]; then + echo "NVSHMEM host library not found; set NVSHMEM_SO=/path/to/libnvshmem_host.so"; exit 4 +fi + +CASE=${1:?usage: run_deepep_matrix.sh } +PHASE=${2:?usage: run_deepep_matrix.sh } + +export CUDA_VISIBLE_DEVICES=0,1 +# Pin NVSHMEM heap chunks to POSIX FD handles โ€” explicit "no fabric handles anywhere". +# (Auto-detect picks FD on non-MNNVL machines anyway; pinning removes the variable.) +export NVSHMEM_CUMEM_HANDLE_TYPE=FILE_DESCRIPTOR + +case "$CASE" in + ll_nofabric) export TEST_USE_FABRIC=0 TEST_NVL_BYTES_MB=0 NCCL_CUMEM_ENABLE=0 NCCL_NVLS_ENABLE=0; PORT=29611 ;; + ll_ncclcumem) export TEST_USE_FABRIC=0 TEST_NVL_BYTES_MB=0 NCCL_CUMEM_ENABLE=1 NCCL_NVLS_ENABLE=0; PORT=29621 ;; + ll_ncclcumem_preinit) export TEST_USE_FABRIC=0 TEST_NVL_BYTES_MB=0 NCCL_CUMEM_ENABLE=1 NCCL_NVLS_ENABLE=0 TEST_NCCL_WARMUP_PRE_REGION=1; PORT=29661 ;; + ll_ncclnvls) export TEST_USE_FABRIC=0 TEST_NVL_BYTES_MB=0 NCCL_CUMEM_ENABLE=1 NCCL_NVLS_ENABLE=1; PORT=29631 ;; + ll_ncclnvls_preinit) export TEST_USE_FABRIC=0 TEST_NVL_BYTES_MB=0 NCCL_CUMEM_ENABLE=1 NCCL_NVLS_ENABLE=1 TEST_NCCL_WARMUP_PRE_REGION=1; PORT=29671 ;; + nvl_ipc) export TEST_USE_FABRIC=0 TEST_NVL_BYTES_MB=64 NCCL_CUMEM_ENABLE=0 NCCL_NVLS_ENABLE=0; PORT=29641 ;; + nvl_fabric) export TEST_USE_FABRIC=1 TEST_NVL_BYTES_MB=64 NCCL_CUMEM_ENABLE=0 NCCL_NVLS_ENABLE=0; PORT=29651 ;; + nvl_ipc_prealloc) export TEST_USE_FABRIC=0 TEST_NVL_BYTES_MB=64 NCCL_CUMEM_ENABLE=0 NCCL_NVLS_ENABLE=0 TEST_LOAD_PREALLOC_MB=1024; PORT=29681 ;; + *) echo "unknown case: $CASE"; exit 2 ;; +esac + +# GPUs are shared with other users โ€” refuse to run without headroom. +for i in 0 1; do + free=$(nvidia-smi --query-gpu=memory.free --format=csv,noheader,nounits -i $i) + if [ "$free" -lt 6000 ]; then + echo "GPU $i has only ${free} MiB free โ€” aborting (need ~6 GB headroom)"; exit 3 + fi +done + +# main() overwrites TEST_USE_FABRIC from the CLI flag, so pass the flag too. +FABRIC_ARG="" +[ "$TEST_USE_FABRIC" = "0" ] && FABRIC_ARG="--no-fabric" + +WORK=${DEEPEP_MATRIX_WORK_ROOT:-$REPO_ROOT/tests/deepep_matrix_work}/$CASE +LOGDIR=${DEEPEP_MATRIX_LOG_DIR:-$REPO_ROOT/tests/deepep_matrix_logs} +mkdir -p "$WORK" "$LOGDIR" +cd "$WORK" # deepep_fabric_archive/ is CWD-relative + +export LD_PRELOAD=$NVSHMEM_SO:$HOOK_SO${LD_PRELOAD:+:$LD_PRELOAD} + +run_phase() { + local p=$1 + # distinct rendezvous port per phase to dodge TIME_WAIT + if [ "$p" = save ]; then export MASTER_PORT=$PORT; else export MASTER_PORT=$((PORT + 1)); fi + echo "=== case=$CASE phase=$p $(date '+%F %T') ===" + echo "env: TEST_USE_FABRIC=$TEST_USE_FABRIC TEST_NVL_BYTES_MB=$TEST_NVL_BYTES_MB" \ + "NCCL_CUMEM_ENABLE=$NCCL_CUMEM_ENABLE NCCL_NVLS_ENABLE=$NCCL_NVLS_ENABLE" \ + "NVSHMEM_CUMEM_HANDLE_TYPE=$NVSHMEM_CUMEM_HANDLE_TYPE MASTER_PORT=$MASTER_PORT" + if [ "$p" = save ]; then rm -rf deepep_fabric_archive; fi + "$PY" "$TEST" --$p $FABRIC_ARG --num-processes=2 2>&1 | tee "$LOGDIR/${CASE}_${p}.log" +} + +if [ "$PHASE" = both ]; then run_phase save; run_phase load; else run_phase "$PHASE"; fi diff --git a/tests/test_deepep_fabric.py b/tests/test_deepep_fabric.py index be6be723..9db689aa 100644 --- a/tests/test_deepep_fabric.py +++ b/tests/test_deepep_fabric.py @@ -22,6 +22,7 @@ import pytest import torch import torch.distributed as dist +from foundry.integration.sglang.config import CUDAGraphExtensionConfig # Use lower base address (matching working deepep_args example) # Note: 0x60000000000 (6TB) NOT 0x600000000000 (96TB) - the extra zero matters! @@ -77,6 +78,27 @@ def _init_dist(local_rank: int, num_local_ranks: int): ) +def _maybe_warm_nccl_pre_region(group): + """Diagnostic knob (TEST_NCCL_WARMUP_PRE_REGION=1): force NCCL transport + setup BEFORE the foundry region exists, so NCCL's cuMem allocations + (NCCL_CUMEM_ENABLE=1) pass through untracked. Localizes whether the + NCCL-CUMEM incompatibility is "NCCL cuMem while region enabled" only. + The warmup tensor's caching-allocator segment is drained afterwards so + post-region allocations don't reuse an untracked (non-deterministic) + segment. + """ + if os.environ.get("TEST_NCCL_WARMUP_PRE_REGION", "0") != "1": + return + t = torch.ones(1024, 1024, device="cuda") + dist.all_reduce(t) # default group: ring/p2p transport setup + dist.all_reduce(t, group=group) # subgroup used by the DeepEP Buffer + dist.all_gather_object([None] * dist.get_world_size(), 0) # object-coll path + torch.cuda.synchronize() + del t + torch.cuda.empty_cache() # legal here: region not set yet, no captures + print("[TEST] NCCL pre-region warmup done (transports up before region)", flush=True) + + def _create_deterministic_inputs( rank: int, num_tokens: int, hidden: int, num_experts: int, num_topk: int ): @@ -148,9 +170,9 @@ def _print_buffer_info(buffer, rank: int, prefix: str): def _verify_dispatch_result( recv_x, - recv_topk_idx, - recv_src_idx, - num_recv, + expert_num_tokens, + dispatch_handle, + dispatch_event, rank: int, num_ranks: int, num_tokens: int, @@ -163,70 +185,56 @@ def _verify_dispatch_result( local_expert_start = rank * num_local_experts local_expert_end = local_expert_start + num_local_experts - # num_recv might be an EventOverlap object (async), need to sync and get value - actual_num_recv = num_recv - if hasattr(num_recv, "wait"): - # It's an async object, wait for it - num_recv.wait() - if hasattr(num_recv, "value"): - actual_num_recv = num_recv.value - elif hasattr(num_recv, "item"): - actual_num_recv = num_recv.item() - elif not isinstance(num_recv, int): - # Try to get from recv_x shape - actual_num_recv = ( - recv_x.shape[0] * recv_x.shape[1] - if recv_x is not None and len(recv_x.shape) >= 2 - else 0 - ) - print( - f"[Rank {rank}] {prefix}: num_recv is {type(num_recv)}, using recv_x shape to estimate", - flush=True, - ) + if hasattr(dispatch_event, "wait"): + dispatch_event.wait() + + assert dispatch_handle is not None, f"[Rank {rank}] {prefix}: missing dispatch handle" + assert isinstance(recv_x, torch.Tensor), f"[Rank {rank}] {prefix}: recv_x is not a tensor" + assert recv_x.ndim == 3 + assert recv_x.shape[0] == num_local_experts + assert recv_x.shape[2] == hidden + assert isinstance(expert_num_tokens, torch.Tensor) + assert expert_num_tokens.numel() == num_local_experts - print(f"[Rank {rank}] {prefix}: Verification - num_recv = {actual_num_recv}", flush=True) + counts = [int(value) for value in expert_num_tokens.tolist()] + assert all(0 <= count <= recv_x.shape[1] for count in counts) + assert sum(counts) > 0 + + print(f"[Rank {rank}] {prefix}: Verification - expert counts = {counts}", flush=True) print( f"[Rank {rank}] {prefix}: Verification - local experts range = [{local_expert_start}, {local_expert_end})", flush=True, ) + print(f"[Rank {rank}] {prefix}: recv_x shape = {recv_x.shape}", flush=True) + + for expert_i, count in enumerate(counts): + expert_global_idx = local_expert_start + expert_i + for token_i in range(min(3, count)): + token = recv_x[expert_i, token_i] + recv_val = token[0].item() + assert torch.all(token == token[0]) + + decoded_src_rank = int(recv_val) // 1000 + decoded_token_id = int(recv_val) % 1000 + assert 0 <= decoded_src_rank < num_ranks + assert 0 <= decoded_token_id < num_tokens + selected_experts = { + decoded_token_id % num_experts, + (decoded_token_id + 1) % num_experts, + } + assert expert_global_idx in selected_experts - # recv_x shape is [num_local_experts, max_tokens_per_expert, hidden] - # Print shape info for debugging - if recv_x is not None: - print(f"[Rank {rank}] {prefix}: recv_x shape = {recv_x.shape}", flush=True) - - # Check some values in recv_x to verify the pattern - if recv_x is not None and recv_x.numel() > 0: - # Check first expert's first few tokens - num_experts_in_recv = recv_x.shape[0] - tokens_per_expert = recv_x.shape[1] - check_experts = min(2, num_experts_in_recv) - check_tokens = min(3, tokens_per_expert) - - print( - f"[Rank {rank}] {prefix}: Checking first {check_experts} experts, {check_tokens} tokens each:", - flush=True, - ) - - for expert_i in range(check_experts): - expert_global_idx = local_expert_start + expert_i - for token_i in range(check_tokens): - recv_val = recv_x[expert_i, token_i, 0].item() - - # Decode: src_rank = recv_val // 1000, token_id = recv_val % 1000 - decoded_src_rank = int(recv_val) // 1000 - decoded_token_id = int(recv_val) % 1000 + print( + f"[Rank {rank}] {prefix}: expert[{expert_i}] (global {expert_global_idx}), " + f"token[{token_i}]: value={recv_val:.1f} " + f"(decoded: src_rank={decoded_src_rank}, token_id={decoded_token_id})", + flush=True, + ) - print( - f"[Rank {rank}] {prefix}: expert[{expert_i}] (global {expert_global_idx}), " - f"token[{token_i}]: value={recv_val:.1f} " - f"(decoded: src_rank={decoded_src_rank}, token_id={decoded_token_id})", - flush=True, - ) - else: - print(f"[Rank {rank}] {prefix}: recv_x is empty or None", flush=True) - return True +def _assert_valid_dispatch_equal(actual, expected, expert_counts): + for expert_i, count in enumerate(expert_counts): + assert torch.equal(actual[expert_i, :count], expected[expert_i, :count]) def _run_save(local_rank: int, num_processes: int): @@ -263,6 +271,8 @@ def _run_save(local_rank: int, num_processes: int): num_experts = num_ranks * 4 # 4 experts per rank num_topk = 2 + _maybe_warm_nccl_pre_region(group) + # Setup allocation region region_size = fdry.parse_size(REGION_SIZE_STR) print(f"[Rank {rank}] SAVE: Setting up allocation region at {hex(BASE_ADDR)}", flush=True) @@ -331,7 +341,7 @@ def _run_save(local_rank: int, num_processes: int): torch.cuda.synchronize() # Unpack dispatch result - recv_x, recv_topk_idx, recv_src_idx, num_recv, *rest = result + recv_x, expert_num_tokens, dispatch_handle, dispatch_event, *rest = result print(f"[Rank {rank}] SAVE: Warmup dispatch result:", flush=True) print( f"[Rank {rank}] SAVE: recv_x address = {hex(recv_x.data_ptr()) if recv_x is not None else 'None'}", @@ -341,14 +351,14 @@ def _run_save(local_rank: int, num_processes: int): f"[Rank {rank}] SAVE: recv_x shape = {recv_x.shape if recv_x is not None else 'None'}", flush=True, ) - print(f"[Rank {rank}] SAVE: num_recv = {num_recv}", flush=True) + print(f"[Rank {rank}] SAVE: expert_num_tokens = {expert_num_tokens}", flush=True) # Verify results _verify_dispatch_result( recv_x, - recv_topk_idx, - recv_src_idx, - num_recv, + expert_num_tokens, + dispatch_handle, + dispatch_event, rank, num_ranks, num_tokens, @@ -356,6 +366,8 @@ def _run_save(local_rank: int, num_processes: int): num_experts, "SAVE-warmup", ) + expected_recv_x = recv_x.clone() + expected_expert_num_tokens = expert_num_tokens.clone() dist.barrier(group) print(f"[Rank {rank}] SAVE: Warmup completed and verified", flush=True) @@ -413,8 +425,26 @@ def _run_save(local_rank: int, num_processes: int): graph.replay() torch.cuda.synchronize() - # Verify graph replay result - graph_num_recv = cumulative_stats.sum().item() # Approximate + # Verify the captured graph produced the same deterministic dispatch. + graph_recv_x, graph_expert_num_tokens, graph_handle, graph_event, *rest = graph_result + _verify_dispatch_result( + graph_recv_x, + graph_expert_num_tokens, + graph_handle, + graph_event, + rank, + num_ranks, + num_tokens, + hidden, + num_experts, + "SAVE-replay", + ) + assert torch.equal(graph_expert_num_tokens, expected_expert_num_tokens) + _assert_valid_dispatch_equal( + graph_recv_x, + expected_recv_x, + [int(value) for value in graph_expert_num_tokens.tolist()], + ) print(f"[Rank {rank}] SAVE: Graph replay completed", flush=True) # Save graph and fatbins to per-rank archive @@ -487,11 +517,22 @@ def _run_load(local_rank: int, num_processes: int): print(f"[Rank {rank}] LOAD: Loading CUDA modules from {rank_archive}", flush=True) fdry.load_cuda_modules_and_libraries(rank_archive) + _maybe_warm_nccl_pre_region(group) + # Step 2: Setup allocation region region_size = fdry.parse_size(REGION_SIZE_STR) print(f"[Rank {rank}] LOAD: Setting up allocation region at {hex(BASE_ADDR)}", flush=True) fdry.set_allocation_region(BASE_ADDR, region_size) + # Optional (TEST_LOAD_PREALLOC_MB=N): reproduce the production LOAD shape - + # one upfront-preallocated chunk that subsequent allocations (including the + # DeepEP NVL buffer) carve from with NO individual cuMem handle. This + # exercises the hook's whole-chunk VMM-IPC export path. + prealloc_mb = int(os.environ.get("TEST_LOAD_PREALLOC_MB", "0")) + if prealloc_mb > 0: + assert fdry.preallocate_region(prealloc_mb * 1024 * 1024), "preallocate_region failed" + print(f"[Rank {rank}] LOAD: preallocated {prealloc_mb} MB chunk", flush=True) + # Step 3: Create DeepEP buffer (initializes NVSHMEM runtime) num_local_experts = num_experts // num_ranks num_rdma_bytes = deep_ep.Buffer.get_low_latency_rdma_size_hint( @@ -499,13 +540,17 @@ def _run_load(local_rank: int, num_processes: int): ) use_fabric = os.environ.get("TEST_USE_FABRIC", "1") == "1" + # Must match SAVE: an NVL buffer allocated on SAVE but not LOAD (or vice versa) + # diverges the allocation trajectory and the peer-pointer setup. + num_nvl_bytes = int(os.environ.get("TEST_NVL_BYTES_MB", "0")) * 1024 * 1024 print( - f"[Rank {rank}] LOAD: Creating DeepEP buffer with use_fabric={use_fabric}, size={num_rdma_bytes / 1e6:.2f} MB", + f"[Rank {rank}] LOAD: Creating DeepEP buffer with use_fabric={use_fabric}, num_nvl_bytes={num_nvl_bytes / 1e6:.2f} MB, num_rdma_bytes={num_rdma_bytes / 1e6:.2f} MB", flush=True, ) buffer = deep_ep.Buffer( group, + num_nvl_bytes=num_nvl_bytes, num_rdma_bytes=num_rdma_bytes, low_latency_mode=True, num_qps_per_rank=num_local_experts, @@ -557,7 +602,7 @@ def _run_load(local_rank: int, num_processes: int): torch.cuda.synchronize() # Unpack and verify warmup result - recv_x, recv_topk_idx, recv_src_idx, num_recv, *rest = warmup_result + recv_x, expert_num_tokens, dispatch_handle, dispatch_event, *rest = warmup_result print(f"[Rank {rank}] LOAD: Warmup dispatch result:", flush=True) print( f"[Rank {rank}] LOAD: recv_x address = {hex(recv_x.data_ptr()) if recv_x is not None else 'None'}", @@ -567,14 +612,14 @@ def _run_load(local_rank: int, num_processes: int): f"[Rank {rank}] LOAD: recv_x shape = {recv_x.shape if recv_x is not None else 'None'}", flush=True, ) - print(f"[Rank {rank}] LOAD: num_recv = {num_recv}", flush=True) + print(f"[Rank {rank}] LOAD: expert_num_tokens = {expert_num_tokens}", flush=True) # Verify warmup results _verify_dispatch_result( recv_x, - recv_topk_idx, - recv_src_idx, - num_recv, + expert_num_tokens, + dispatch_handle, + dispatch_event, rank, num_ranks, num_tokens, @@ -582,6 +627,8 @@ def _run_load(local_rank: int, num_processes: int): num_experts, "LOAD-warmup", ) + expected_recv_x = recv_x.clone() + expected_counts = [int(value) for value in expert_num_tokens.tolist()] dist.barrier(group) print(f"[Rank {rank}] LOAD: Warmup completed", flush=True) @@ -593,10 +640,21 @@ def _run_load(local_rank: int, num_processes: int): flush=True, ) - # Load graph from per-rank archive + # Load graph from per-rank archive. + # Default: the production path (start_graph_builds + finish_graph_loads), + # same machinery the vLLM/sglang integrations use. TEST_LOAD_API=single + # selects the legacy CUDAGraph.load path โ€” KNOWN BROKEN for DeepEP graphs: + # it does not merge root-level common_kernel_node_attrs, so factored-out + # cooperative/clusterDim attrs are dropped and the dispatch kernel faults + # with "unspecified launch failure" at replay. graph_json = os.path.join(rank_archive, "deepep_dispatch_graph.json") - print(f"[Rank {rank}] LOAD: Loading graph from {graph_json}", flush=True) - graph, output_tensors = fdry.CUDAGraph.load(graph_json) + load_api = os.environ.get("TEST_LOAD_API", "parallel") + print(f"[Rank {rank}] LOAD: Loading graph from {graph_json} (api={load_api})", flush=True) + if load_api == "single": + graph, output_tensors = fdry.CUDAGraph.load(graph_json) + else: + pending = fdry.CUDAGraph.start_graph_builds([graph_json]) + ((graph, output_tensors),) = fdry.CUDAGraph.finish_graph_loads(pending) print(f"[Rank {rank}] LOAD: Graph loaded successfully", flush=True) # Print loaded output tensor info @@ -619,34 +677,21 @@ def _run_load(local_rank: int, num_processes: int): ) loaded_recv_x = output_tensors + assert loaded_recv_x is not None, "Loaded graph did not reconstruct recv_x" + assert loaded_recv_x.shape == expected_recv_x.shape + for expert_i, count in enumerate(expected_counts): + loaded_recv_x[expert_i, :count].fill_(float("nan")) + torch.cuda.synchronize() + # Replay loaded graph print(f"[Rank {rank}] LOAD: Replaying loaded graph", flush=True) graph.replay() torch.cuda.synchronize() - # Verify after replay - check the LOADED output tensors (not warmup recv_x) - # recv_x shape is [num_local_experts, max_tokens_per_expert, hidden] + # Verify after replay against the deterministic warmup output. The sentinel + # fill above ensures a no-op replay cannot pass by reusing warmup contents. print(f"[Rank {rank}] LOAD: After graph replay:", flush=True) - verify_tensor = loaded_recv_x if loaded_recv_x is not None else recv_x - if verify_tensor is not None and verify_tensor.numel() > 0: - print(f"[Rank {rank}] LOAD: verify_tensor shape = {verify_tensor.shape}", flush=True) - num_local_experts = num_experts // num_ranks - local_expert_start = rank * num_local_experts - - # Check first expert's first few tokens - check_experts = min(2, verify_tensor.shape[0]) - check_tokens = min(3, verify_tensor.shape[1]) - - for expert_i in range(check_experts): - for token_i in range(check_tokens): - val = verify_tensor[expert_i, token_i, 0].item() - decoded_src_rank = int(val) // 1000 - decoded_token_id = int(val) % 1000 - print( - f"[Rank {rank}] LOAD: expert[{expert_i}], token[{token_i}]: " - f"value={val:.1f} (rank={decoded_src_rank}, token={decoded_token_id})", - flush=True, - ) + _assert_valid_dispatch_equal(loaded_recv_x, expected_recv_x, expected_counts) print(f"[Rank {rank}] LOAD: Graph replay successful", flush=True) @@ -671,14 +716,17 @@ def _cleanup_archive(): print(f"[TEST] Cleaned up {HOOK_ARCHIVE_DIR}") -def _spawn_with_preload(mode: str, num_processes: int): +def _spawn_with_preload(mode: str, num_processes: int, *, use_fabric: bool = True): """Spawn subprocess with LD_PRELOAD for CGE hook.""" so_path = _get_hook_so_path() env = os.environ.copy() + preload_paths = [so_path] + nvshmem_path = CUDAGraphExtensionConfig._detect_nvshmem_host_path() + if nvshmem_path: + preload_paths.insert(0, nvshmem_path) if env.get("LD_PRELOAD"): - env["LD_PRELOAD"] = f"{so_path}:{env['LD_PRELOAD']}" - else: - env["LD_PRELOAD"] = so_path + preload_paths.append(env["LD_PRELOAD"]) + env["LD_PRELOAD"] = ":".join(preload_paths) # # Set NVSHMEM environment variables to disable IBGDA and remote transports # env['NVSHMEM_REMOTE_TRANSPORT'] = 'none' @@ -691,6 +739,8 @@ def _spawn_with_preload(mode: str, num_processes: int): f"--{mode}", f"--num-processes={num_processes}", ] + if not use_fabric: + cmd.append("--no-fabric") subprocess.check_call(cmd, env=env) @@ -711,24 +761,40 @@ def _check_multi_gpu(): @pytest.mark.skipif(not _check_deepep_available(), reason="DeepEP not installed") @pytest.mark.skipif(not _check_multi_gpu(), reason="Requires at least 2 GPUs") -def test_deepep_fabric_graph_save_load(): - """Test DeepEP with use_fabric=True graph save/load.""" +@pytest.mark.parametrize( + ("use_fabric", "nvl_bytes_mb"), + [ + (True, 0), + (False, 64), + ], + ids=["fabric-rdma", "nvl-scm-rights"], +) +def test_deepep_graph_save_load(monkeypatch, use_fabric: bool, nvl_bytes_mb: int): + """Test both fabric and legacy NVL/SCM_RIGHTS graph save/load paths.""" pytest.importorskip("deep_ep") num_processes = min(torch.cuda.device_count(), 2) # Use 2 GPUs for testing + monkeypatch.setenv("TEST_NVL_BYTES_MB", str(nvl_bytes_mb)) + monkeypatch.setenv("NCCL_CUMEM_ENABLE", "0") + monkeypatch.setenv("NCCL_NVLS_ENABLE", "0") + if not use_fabric: + monkeypatch.setenv("NVSHMEM_CUMEM_HANDLE_TYPE", "FILE_DESCRIPTOR") - print(f"\n[TEST] Starting DeepEP fabric graph save/load test with {num_processes} processes") + print( + f"\n[TEST] Starting DeepEP graph save/load test with {num_processes} processes " + f"(use_fabric={use_fabric}, nvl_bytes_mb={nvl_bytes_mb})" + ) _cleanup_archive() print("[TEST] Running SAVE mode (create buffer, capture graph, save)") - _spawn_with_preload("save", num_processes) + _spawn_with_preload("save", num_processes, use_fabric=use_fabric) print("[TEST] Running LOAD mode (load graph, replay)") - _spawn_with_preload("load", num_processes) + _spawn_with_preload("load", num_processes, use_fabric=use_fabric) _cleanup_archive() - print("[TEST] test_deepep_fabric_graph_save_load PASSED") + print("[TEST] test_deepep_graph_save_load PASSED") def main(): @@ -795,10 +861,10 @@ def main(): _cleanup_archive() print("[TEST] Running SAVE mode") - _spawn_with_preload("save", num_processes) + _spawn_with_preload("save", num_processes, use_fabric=not args.no_fabric) print("[TEST] Running LOAD mode") - _spawn_with_preload("load", num_processes) + _spawn_with_preload("load", num_processes, use_fabric=not args.no_fabric) _cleanup_archive() print("[TEST] PASSED") diff --git a/tests/test_load_cublas_ws.py b/tests/test_load_cublas_ws.py index 49f79406..a7f9ece4 100644 --- a/tests/test_load_cublas_ws.py +++ b/tests/test_load_cublas_ws.py @@ -119,7 +119,7 @@ def _run_loading_run(): print("[LOADING] cuBLAS initialized") graph_json = os.path.join(ARCHIVE_DIR, "graph.json") - graph = fdry.CUDAGraph.load(graph_json) + graph, _loaded_output = fdry.CUDAGraph.load(graph_json) print(f"[LOADING] Graph loaded from {graph_json}") graph.replay() diff --git a/tests/test_nvjet_graph.py b/tests/test_nvjet_graph.py index 08fc81c3..bdea48db 100644 --- a/tests/test_nvjet_graph.py +++ b/tests/test_nvjet_graph.py @@ -121,7 +121,7 @@ def _run_loading_run(): print("[LOADING] cuBLAS initialized") graph_json = os.path.join(ARCHIVE_DIR, "graph.json") - graph = fdry.CUDAGraph.load(graph_json) + graph, _loaded_output = fdry.CUDAGraph.load(graph_json) print(f"[LOADING] Graph loaded from {graph_json}") graph.replay() diff --git a/tests/test_rng_graph.py b/tests/test_rng_graph.py new file mode 100644 index 00000000..1ad60ec8 --- /dev/null +++ b/tests/test_rng_graph.py @@ -0,0 +1,104 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the Foundry project +"""Fresh-process replay coverage for CUDA graphs that consume Philox state.""" + +from __future__ import annotations + +import importlib.util +import os +import shutil +import subprocess +import sys +from pathlib import Path + +import foundry as fdry +import pytest +import torch + +BASE_ADDR = 0x650000000000 +REGION_SIZE = "2GB" +ARCHIVE_DIR = "rng_graph_archive" +HOOK_ARCHIVE_DIR = "hook_archive" +GRAPH_PATH = f"{ARCHIVE_DIR}/graph.json" + + +def _hook_path() -> str: + spec = importlib.util.find_spec("foundry.ops") + if not spec or not spec.origin: + raise RuntimeError("foundry.ops is not installed") + path = Path(spec.origin).resolve().parent / "libcuda_hook.so" + if not path.exists(): + raise RuntimeError(f"CUDA hook not found at {path}") + return str(path) + + +def _save() -> None: + torch.cuda.init() + torch.cuda.set_device(0) + torch.cuda.manual_seed_all(1234) + Path(ARCHIVE_DIR).mkdir(exist_ok=True) + + fdry.set_allocation_region(BASE_ADDR, fdry.parse_size(REGION_SIZE)) + graph = fdry.CUDAGraph() + with fdry.graph(graph, pool=(0, 0)): + output = torch.rand(4096, device="cuda") + graph.save(GRAPH_PATH, output_tensors=output) + fdry.stop_allocation_region() + + +def _load() -> None: + torch.cuda.init() + torch.cuda.set_device(0) + torch.cuda.manual_seed_all(1234) + + fdry.load_cuda_modules_and_libraries(HOOK_ARCHIVE_DIR) + fdry.set_allocation_region(BASE_ADDR, fdry.parse_size(REGION_SIZE)) + graph, output = fdry.CUDAGraph.load(GRAPH_PATH, pool=(0, 0)) + + graph.replay() + torch.cuda.synchronize() + first = output.clone() + graph.replay() + torch.cuda.synchronize() + second = output.clone() + + assert torch.isfinite(first).all() + assert torch.isfinite(second).all() + assert ((first >= 0) & (first < 1)).all() + assert ((second >= 0) & (second < 1)).all() + assert not torch.equal(first, second) + fdry.stop_allocation_region() + + +def _spawn(mode: str) -> None: + env = os.environ.copy() + hook = _hook_path() + env["LD_PRELOAD"] = f"{hook}:{env['LD_PRELOAD']}" if env.get("LD_PRELOAD") else hook + subprocess.check_call( + [sys.executable, str(Path(__file__).resolve()), f"--{mode}"], + env=env, + ) + + +def _cleanup() -> None: + shutil.rmtree(ARCHIVE_DIR, ignore_errors=True) + shutil.rmtree(HOOK_ARCHIVE_DIR, ignore_errors=True) + + +@pytest.mark.filterwarnings("ignore:TORCH_CUDA_ARCH_LIST is not set") +def test_rng_graph_replays_in_fresh_process() -> None: + _cleanup() + try: + _spawn("save") + _spawn("load") + finally: + _cleanup() + + +if __name__ == "__main__": + if "--save" in sys.argv: + _save() + elif "--load" in sys.argv: + _load() + else: + test_rng_graph_replays_in_fresh_process() diff --git a/tests/test_sglang_flashinfer_graph_abi.py b/tests/test_sglang_flashinfer_graph_abi.py new file mode 100644 index 00000000..7fe08255 --- /dev/null +++ b/tests/test_sglang_flashinfer_graph_abi.py @@ -0,0 +1,441 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the Foundry project +"""CPU-only contract coverage for pinned FlashInfer graph ABI extraction.""" + +from __future__ import annotations + +import copy +import importlib.util +import sys +from dataclasses import replace +from pathlib import Path +from types import ModuleType, SimpleNamespace + +import pytest + +SGLANG_DIR = Path(__file__).parents[1] / "python" / "foundry" / "integration" / "sglang" + + +def _ensure_package(name: str) -> ModuleType: + package = sys.modules.get(name) + if package is None: + package = ModuleType(name) + package.__path__ = [] # type: ignore[attr-defined] + sys.modules[name] = package + return package + + +def _load_module(module_name: str, filename: str): + module_path = SGLANG_DIR / filename + spec = importlib.util.spec_from_file_location(module_name, module_path) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + sys.modules[module_name] = module + spec.loader.exec_module(module) + return module + + +foundry_pkg = _ensure_package("foundry") +integration_pkg = _ensure_package("foundry.integration") +sglang_pkg = _ensure_package("foundry.integration.sglang") +foundry_pkg.integration = integration_pkg +integration_pkg.sglang = sglang_pkg + +_shape_replay_state = _load_module( + "foundry.integration.sglang.shape_replay_state", + "shape_replay_state.py", +) +sglang_pkg.shape_replay_state = _shape_replay_state +_state_manifest = _load_module( + "foundry.integration.sglang.state_manifest", + "state_manifest.py", +) +sglang_pkg.state_manifest = _state_manifest +_shape_adapter = _load_module( + "foundry.integration.sglang.sglang_shape_adapter", + "sglang_shape_adapter.py", +) +sglang_pkg.sglang_shape_adapter = _shape_adapter +_flashinfer_graph_abi = _load_module( + "foundry.integration.sglang.flashinfer_graph_abi", + "flashinfer_graph_abi.py", +) +sglang_pkg.flashinfer_graph_abi = _flashinfer_graph_abi + +MERGE_KERNEL_NAME = _flashinfer_graph_abi.MERGE_KERNEL_NAME +PAGED_KERNEL_NAME = _flashinfer_graph_abi.PAGED_KERNEL_NAME +extract_flashinfer_operand_records = _flashinfer_graph_abi.extract_flashinfer_operand_records +relocate_flashinfer_operands = _flashinfer_graph_abi.relocate_flashinfer_operands + + +class Slot: + def __init__(self, name: str, ptr: int, numel: int) -> None: + self.name = name + self.data_ptr = ptr + self.numel = numel + self.element_size = 1 + + +def write_u64(raw: bytearray, offset: int, value: int) -> None: + raw[offset : offset + 8] = value.to_bytes(8, "little") + + +def make_state(): + plan = (66, 8, 864, 16, 0, 272, 544, 880, 816, 852, 0, 8650752, 928, 1, 1) + owners = ( + Slot("int_workspace", 0x600100000000, 8 * 1024 * 1024), + Slot("qo_indptr", 0x600110000000, 4096), + ) + shared = ( + Slot("paged_kv_indptr", 0x600120000000, 4096), + Slot("paged_kv_indices", 0x600130000000, 4096), + Slot("paged_kv_last_page_len", 0x600140000000, 4096), + Slot("float_workspace", 0x600150000000, 16 * 1024 * 1024), + ) + return SimpleNamespace( + batch_size=8, + plan_fingerprint=plan, + owners=owners, + shared_owners=shared, + ) + + +def make_live_state(): + plan = (66, 8, 864, 16, 0, 272, 544, 880, 816, 852, 0, 8650752, 928, 1, 1) + owners = ( + Slot("int_workspace", 0x600200000000, 8 * 1024 * 1024), + Slot("qo_indptr", 0x600210000000, 4096), + ) + shared = ( + Slot("paged_kv_indptr", 0x600220000000, 4096), + Slot("paged_kv_indices", 0x600230000000, 4096), + Slot("paged_kv_last_page_len", 0x600240000000, 4096), + Slot("float_workspace", 0x600250000000, 16 * 1024 * 1024), + ) + return SimpleNamespace( + batch_size=8, + plan_fingerprint=plan, + owners=owners, + shared_owners=shared, + ) + + +def make_graph(state) -> dict: + slots = {slot.name: slot for slot in (*state.owners, *state.shared_owners)} + plan = state.plan_fingerprint + paged = bytearray(376) + values = { + 0: 0x600180000000, + 56: 0x600190000000, + 64: 0x6001A0000000, + 72: slots["paged_kv_indices"].data_ptr, + 80: slots["paged_kv_indptr"].data_ptr, + 88: slots["paged_kv_last_page_len"].data_ptr, + 104: slots["qo_indptr"].data_ptr, + 112: slots["float_workspace"].data_ptr + plan[10], + 120: slots["float_workspace"].data_ptr + plan[11], + 296: slots["int_workspace"].data_ptr + plan[4], + 304: slots["int_workspace"].data_ptr + plan[5], + 312: slots["int_workspace"].data_ptr + plan[6], + 320: slots["int_workspace"].data_ptr + plan[7], + 328: slots["int_workspace"].data_ptr + plan[8], + 336: slots["int_workspace"].data_ptr + plan[12], + 344: slots["int_workspace"].data_ptr + plan[9], + 360: slots["int_workspace"].data_ptr + plan[2], + } + for offset, value in values.items(): + write_u64(paged, offset, value) + paged[352:356] = state.batch_size.to_bytes(4, "little") + paged[256:260] = (16).to_bytes(4, "little") + paged[372] = 1 + + merge_values = ( + values[112], + values[120], + values[320], + 0x600160000000, + 0, + state.batch_size, + values[360], + 16, + ) + merge_layout = ( + (0, 0, 8), + (1, 8, 8), + (2, 16, 8), + (3, 24, 8), + (4, 32, 8), + (5, 40, 4), + (6, 48, 8), + (7, 56, 4), + ) + + nodes = [] + for layer in range(36): + nodes.append( + { + "id": layer * 2, + "type": "KernelNode", + "params": { + "function_name": PAGED_KERNEL_NAME, + "kernelParams": [ + { + "index": 0, + "offset": 0, + "size": 376, + "value_hex": paged.hex(), + } + ], + "extra_argBuffer_hex": "", + }, + } + ) + nodes.append( + { + "id": layer * 2 + 1, + "type": "KernelNode", + "params": { + "function_name": MERGE_KERNEL_NAME, + "kernelParams": [ + { + "index": index, + "offset": offset, + "size": size, + "value_hex": int(value).to_bytes(size, "little").hex(), + } + for (index, offset, size), value in zip( + merge_layout, + merge_values, + strict=True, + ) + ], + "extra_argBuffer_hex": "", + }, + } + ) + return {"nodes": nodes} + + +def first_record(records, name: str): + return next(record for record in records if record.name == name) + + +def test_extracts_all_pinned_flashinfer_operands() -> None: + state = make_state() + records = extract_flashinfer_operand_records(make_graph(state), state) + + assert len(records) == 36 * (3 + 1 + 14 + 4 + 1) + assert {record.owner_slot for record in records} == { + "foundry_vmm", + "normalized_null", + "int_workspace", + "qo_indptr", + "paged_kv_indptr", + "paged_kv_indices", + "paged_kv_last_page_len", + "float_workspace", + } + + +def test_rejects_changed_paged_pointer() -> None: + state = make_state() + graph = make_graph(state) + raw = bytearray.fromhex(graph["nodes"][0]["params"]["kernelParams"][0]["value_hex"]) + write_u64(raw, 296, 0x600170000000) + graph["nodes"][0]["params"]["kernelParams"][0]["value_hex"] = raw.hex() + + with pytest.raises(RuntimeError, match="request_indices"): + extract_flashinfer_operand_records(graph, state) + + +def test_relocates_extracted_flashinfer_operands() -> None: + saved_state = make_state() + live_state = make_live_state() + graph = make_graph(saved_state) + records = extract_flashinfer_operand_records(graph, saved_state) + + relocate_flashinfer_operands(graph, live_state, records) + + paged_raw = bytes.fromhex(graph["nodes"][0]["params"]["kernelParams"][0]["value_hex"]) + merge_tmp_v = int.from_bytes( + bytes.fromhex(graph["nodes"][1]["params"]["kernelParams"][0]["value_hex"]), + "little", + ) + assert int.from_bytes(paged_raw[72:80], "little") == 0x600230000000 + assert int.from_bytes(paged_raw[80:88], "little") == 0x600220000000 + assert int.from_bytes(paged_raw[296:304], "little") == 0x600200000000 + assert merge_tmp_v == 0x600250000000 + + +def test_relocation_rejects_changed_function_symbol() -> None: + saved_state = make_state() + graph = make_graph(saved_state) + records = extract_flashinfer_operand_records(graph, saved_state) + graph["nodes"][0]["params"]["function_name"] = "renamed_kernel" + + with pytest.raises(RuntimeError, match="Expected 36 FlashInfer paged and merge nodes"): + relocate_flashinfer_operands(graph, make_live_state(), records) + + +@pytest.mark.parametrize( + ("mutate", "message"), + [ + ( + lambda graph: graph["nodes"][1]["params"]["kernelParams"][0].__setitem__( + "index", + 9, + ), + "ABI changed", + ), + ( + lambda graph: graph["nodes"][1]["params"]["kernelParams"][0].__setitem__( + "offset", + 16, + ), + "ABI changed", + ), + ( + lambda graph: graph["nodes"][1]["params"]["kernelParams"][0].__setitem__( + "size", + 4, + ), + "ABI changed", + ), + ], +) +def test_relocation_rejects_changed_live_parameter_layout(mutate, message: str) -> None: + saved_state = make_state() + graph = make_graph(saved_state) + records = extract_flashinfer_operand_records(graph, saved_state) + mutate(graph) + + with pytest.raises(RuntimeError, match=message): + relocate_flashinfer_operands(graph, make_live_state(), records) + + +def test_relocation_rejects_nonempty_arg_buffer() -> None: + saved_state = make_state() + graph = make_graph(saved_state) + records = extract_flashinfer_operand_records(graph, saved_state) + graph["nodes"][1]["params"]["extra_argBuffer_hex"] = "00" + + with pytest.raises(RuntimeError, match="unsupported arg buffer"): + relocate_flashinfer_operands(graph, make_live_state(), records) + + +@pytest.mark.parametrize( + ("build_records", "message"), + [ + (lambda records: records[:-1], "record set"), + (lambda records: records + (records[0],), "record set"), + ], +) +def test_relocation_rejects_missing_or_extra_records(build_records, message: str) -> None: + saved_state = make_state() + graph = make_graph(saved_state) + records = extract_flashinfer_operand_records(graph, saved_state) + + with pytest.raises(RuntimeError, match=message): + relocate_flashinfer_operands(graph, make_live_state(), build_records(records)) + + +@pytest.mark.parametrize( + ("mutate", "message"), + [ + (lambda record: replace(record, source="extra_argBuffer_hex"), "source"), + ( + lambda record: replace( + record, + cuda_parameter_offset=record.cuda_parameter_offset + 8, + ), + "cuda parameter offset", + ), + (lambda record: replace(record, ctype="float*"), "ctype"), + (lambda record: replace(record, owner_slot="communicator"), "owner slot"), + ( + lambda record: replace( + record, + owner_relative_offset=record.owner_relative_offset + 4, + ), + "owner relative offset", + ), + (lambda record: replace(record, span_bytes=record.span_bytes + 1), "span"), + ], +) +def test_relocation_rejects_record_metadata_mismatch(mutate, message: str) -> None: + saved_state = make_state() + graph = make_graph(saved_state) + records = extract_flashinfer_operand_records(graph, saved_state) + target = first_record(records, "request_indices") + mutated = tuple(mutate(record) if record is target else record for record in records) + + with pytest.raises(RuntimeError, match=message): + relocate_flashinfer_operands(graph, make_live_state(), mutated) + + +def test_relocation_rejects_negative_owner_relative_offset() -> None: + saved_state = make_state() + graph = make_graph(saved_state) + records = extract_flashinfer_operand_records(graph, saved_state) + target = first_record(records, "request_indices") + mutated = tuple( + replace(record, owner_relative_offset=-1) if record is target else record + for record in records + ) + + with pytest.raises(RuntimeError, match="negative owner_relative_offset"): + relocate_flashinfer_operands(graph, make_live_state(), mutated) + + +def test_relocation_rejects_negative_span() -> None: + saved_state = make_state() + graph = make_graph(saved_state) + records = extract_flashinfer_operand_records(graph, saved_state) + target = first_record(records, "request_indices") + mutated = tuple( + replace(record, span_bytes=-1) if record is target else record for record in records + ) + + with pytest.raises(RuntimeError, match="negative span"): + relocate_flashinfer_operands(graph, make_live_state(), mutated) + + +def test_relocation_rejects_out_of_bounds_value_offset() -> None: + saved_state = make_state() + graph = make_graph(saved_state) + records = extract_flashinfer_operand_records(graph, saved_state) + target = first_record(records, "request_indices") + mutated = tuple( + replace(record, value_byte_offset=400) if record is target else record for record in records + ) + + with pytest.raises(RuntimeError, match="out of bounds"): + relocate_flashinfer_operands(graph, make_live_state(), mutated) + + +def test_relocation_rejects_owner_range_violation() -> None: + saved_state = make_state() + graph = make_graph(saved_state) + records = extract_flashinfer_operand_records(graph, saved_state) + target = first_record(records, "request_indices") + mutated = tuple( + replace(record, owner_relative_offset=8 * 1024 * 1024) if record is target else record + for record in records + ) + + with pytest.raises(RuntimeError, match="owner range"): + relocate_flashinfer_operands(graph, make_live_state(), mutated) + + +def test_relocation_leaves_graph_unchanged_on_late_failure() -> None: + saved_state = make_state() + graph = make_graph(saved_state) + original = copy.deepcopy(graph) + records = extract_flashinfer_operand_records(graph, saved_state) + last = records[-1] + mutated = (*records[:-1], replace(last, saved_value=last.saved_value + 8)) + + with pytest.raises(RuntimeError, match="Saved FlashInfer operand changed"): + relocate_flashinfer_operands(graph, make_live_state(), mutated) + + assert graph == original diff --git a/tests/test_sglang_main_backend.py b/tests/test_sglang_main_backend.py new file mode 100644 index 00000000..5a0105c3 --- /dev/null +++ b/tests/test_sglang_main_backend.py @@ -0,0 +1,539 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the Foundry project +"""CPU-only contracts for the SGLang main backend capsule wiring.""" + +from __future__ import annotations + +import importlib.util +import json +import sys +from contextlib import contextmanager +from pathlib import Path +from types import ModuleType, SimpleNamespace + +import pytest + +SGLANG_DIR = Path(__file__).parents[1] / "python" / "foundry" / "integration" / "sglang" + + +def _ensure_package(name: str) -> ModuleType: + package = sys.modules.get(name) + if package is None: + package = ModuleType(name) + package.__path__ = [] # type: ignore[attr-defined] + sys.modules[name] = package + return package + + +def _load_sglang_module(module_name: str, filename: str) -> ModuleType: + spec = importlib.util.spec_from_file_location(module_name, SGLANG_DIR / filename) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + sys.modules[module_name] = module + spec.loader.exec_module(module) + return module + + +foundry_pkg = _ensure_package("foundry") +integration_pkg = _ensure_package("foundry.integration") +sglang_pkg = _ensure_package("foundry.integration.sglang") +foundry_pkg.integration = integration_pkg +integration_pkg.sglang = sglang_pkg + +_config = _load_sglang_module("foundry.integration.sglang.config", "config.py") +sglang_pkg.config = _config +is_symmetric_multi_shape_enabled = _config.is_symmetric_multi_shape_enabled + + +def test_multi_shape_flag_defaults_off(monkeypatch) -> None: + monkeypatch.delenv("FOUNDRY_SGLANG_SYMM_MULTI_SHAPE", raising=False) + assert not is_symmetric_multi_shape_enabled() + + +def test_multi_shape_flag_accepts_only_one(monkeypatch) -> None: + monkeypatch.setenv("FOUNDRY_SGLANG_SYMM_MULTI_SHAPE", "1") + assert is_symmetric_multi_shape_enabled() + monkeypatch.setenv("FOUNDRY_SGLANG_SYMM_MULTI_SHAPE", "true") + with pytest.raises(RuntimeError, match="must be 0 or 1"): + is_symmetric_multi_shape_enabled() + + +def _stub(monkeypatch, name: str) -> ModuleType: + module = ModuleType(name) + monkeypatch.setitem(sys.modules, name, module) + return module + + +def _load_backend_module(monkeypatch): + foundry_pkg = _ensure_package("foundry") + integration_pkg = _ensure_package("foundry.integration") + sglang_pkg = _ensure_package("foundry.integration.sglang") + foundry_pkg.integration = integration_pkg + integration_pkg.sglang = sglang_pkg + sglang_pkg.config = _config + + for module_name, filename in ( + ("foundry.integration.sglang.shape_replay_state", "shape_replay_state.py"), + ("foundry.integration.sglang.state_manifest", "state_manifest.py"), + ("foundry.integration.sglang.sglang_shape_adapter", "sglang_shape_adapter.py"), + ("foundry.integration.sglang.flashinfer_graph_abi", "flashinfer_graph_abi.py"), + ("foundry.integration.sglang.symm_mem_graph", "symm_mem_graph.py"), + ): + module = _load_sglang_module(module_name, filename) + setattr(sglang_pkg, module_name.rsplit(".", maxsplit=1)[-1], module) + + for name in ( + "sglang", + "sglang.srt", + "sglang.srt.distributed", + "sglang.srt.distributed.device_communicators", + "sglang.srt.layers", + "sglang.srt.model_executor", + "sglang.srt.model_executor.runner_backend", + "sglang.srt.model_executor.runner_utils", + ): + package = _stub(monkeypatch, name) + package.__path__ = [] # type: ignore[attr-defined] + + pynccl = _stub( + monkeypatch, + "sglang.srt.distributed.device_communicators.pynccl_allocator", + ) + pynccl.set_graph_pool_id = lambda pool: None + + logits = _stub(monkeypatch, "sglang.srt.layers.logits_processor") + logits.LogitsProcessorOutput = type("LogitsProcessorOutput", (), {}) + + base = _stub( + monkeypatch, + "sglang.srt.model_executor.runner_backend.base_cuda_graph_backend", + ) + base.BaseCudaGraphBackend = object + + pool = _stub(monkeypatch, "sglang.srt.model_executor.runner_utils.pool") + pool.get_or_create_global_graph_memory_pool = lambda device: (0, 0) + + foundry_ops = _stub(monkeypatch, "foundry.ops") + foundry_ops.pack_fatbins_to_folder = lambda workspace_dir: None + foundry_ops.set_pack_fatbins_on_exit = lambda enabled: None + foundry_pkg.ops = foundry_ops + foundry_pkg.save_graph_manifest = lambda workspace_dir: None + + foundry_graph = _stub(monkeypatch, "foundry.graph") + + class _StubCUDAGraph: + @staticmethod + def start_graph_builds(paths, num_threads=4): + del paths, num_threads + raise AssertionError("graph builds should not start in this CPU cleanup test") + + @contextmanager + def _graph_context(graph, *, pool=None, stream=None): + del graph, pool, stream + yield + + foundry_graph.CUDAGraph = _StubCUDAGraph + foundry_graph.graph = _graph_context + foundry_pkg.graph = foundry_graph + + runtime = _stub(monkeypatch, "foundry.integration.sglang.runtime") + runtime.log_alloc_offset = lambda label: None + runtime.preallocate_for_load_mode = lambda: None + runtime.capture_final_alloc_offset = lambda: None + runtime.get_state = lambda: None + sglang_pkg.runtime = runtime + + graph_ops = _stub(monkeypatch, "foundry.integration.sglang.graph_ops") + graph_ops._pack_output = lambda output: output + graph_ops._scan_graph_files = lambda workspace_dir: [] + graph_ops._unpack_output = lambda tensors: tensors + sglang_pkg.graph_ops = graph_ops + + spec = importlib.util.spec_from_file_location( + "tested_main_backend", + SGLANG_DIR / "main_backend.py", + ) + 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 + + +class FakeState: + def __init__(self, capture_index: int) -> None: + self.capture_index = capture_index + + +class FakeShapeKey: + def __init__(self, size: int) -> None: + self.size = size + self.stream_idx = None + self.variant_label = None + + +class RecordingDict(dict): + def __init__(self, events: list[tuple[str, object, object]], label: str, *args, **kwargs): + self._events = events + self._label = label + super().__init__(*args, **kwargs) + + def __setitem__(self, key, value) -> None: + self._events.append((self._label, key, value)) + super().__setitem__(key, value) + + +class TrackingState: + def __init__(self, batch_size: int, capture_index: int, events: list[object]) -> None: + self.batch_size = batch_size + self.capture_index = capture_index + self.events = events + self.graph = None + self.outputs = None + self.binding = "restored" + + def attach_graph(self, graph, outputs) -> None: + self.events.append(("attach", graph, outputs)) + self.graph = graph + self.outputs = outputs + + +def _make_backend(module, *, multi_shape_enabled: bool, symmetric_memory_enabled: bool = True): + communicator = SimpleNamespace(_foundry_symm_mem_handle=object()) + backend = module.FoundryMainCudaGraphBackend.__new__(module.FoundryMainCudaGraphBackend) + backend._runner = SimpleNamespace( + device_module=SimpleNamespace(synchronize=lambda: None), + model_runner=SimpleNamespace( + tp_group=SimpleNamespace( + barrier=lambda: None, + torch_symm_mem_comm=communicator, + ) + ), + ) + backend._graphs = {} + backend._outputs = {} + backend._pool = object() + backend._stream = object() + backend._pending = None + backend._graph_files = [] + backend._graph_load_paths = [] + backend._load_index = 0 + backend._closed = False + backend._symmetric_memory_enabled = symmetric_memory_enabled + backend._multi_shape_enabled = multi_shape_enabled + backend._relocated_graph_dir = None + backend._shape_states = {} + backend._saved_shape_manifest = None + backend._saved_records_by_batch = {} + backend._operand_inventories = {} + return backend, communicator + + +def test_backend_cleanup_closes_every_shape_in_capture_order(monkeypatch) -> None: + module = _load_backend_module(monkeypatch) + events = [] + monkeypatch.setattr( + module, + "close_shape_state", + lambda state: events.append( + ( + state.capture_index, + dict(backend._graphs), + dict(backend._outputs), + ) + ), + raising=False, + ) + backend = module.FoundryMainCudaGraphBackend.__new__(module.FoundryMainCudaGraphBackend) + backend._relocated_graph_dir = None + backend._graphs = {8: "graph-eight", 1: "graph-one"} + backend._outputs = {8: "output-eight", 1: "output-one"} + backend._pool = None + backend._pending = None + backend._graph_files = [] + backend._graph_load_paths = [] + backend._load_index = 0 + backend._shape_states = { + 8: FakeState(0), + 1: FakeState(1), + } + + backend.cleanup() + + assert events == [ + (0, {}, {}), + (1, {}, {}), + ] + assert backend._shape_states == {} + + +def test_multi_shape_save_capture_one_orchestrates_capsule_lifecycle(monkeypatch) -> None: + module = _load_backend_module(monkeypatch) + backend, communicator = _make_backend(module, multi_shape_enabled=True) + events = [] + shape_key = FakeShapeKey(8) + graph = object() + outputs = object() + state = TrackingState(batch_size=8, capture_index=0, events=events) + backend._shape_states = RecordingDict(events, "shape_states") + backend._graphs = RecordingDict(events, "graphs") + backend._outputs = RecordingDict(events, "outputs") + monkeypatch.setattr( + module, + "get_graph_extension_mode", + lambda: module.CUDAGraphExtensionMode.SAVE, + ) + + monkeypatch.setattr( + module, + "create_shape_state", + lambda runner, **kwargs: events.append( + ("create", runner, kwargs["shape_key"], kwargs["capture_index"], kwargs["communicator"]) + ) + or state, + ) + + @contextmanager + def _activate(captured_state): + assert captured_state is state + state.binding = "active" + events.append("activate") + try: + yield + finally: + state.binding = "restored" + + monkeypatch.setattr(module, "activate_shape_state", _activate) + monkeypatch.setattr(module, "run_graph_warmups", lambda **kwargs: events.append("warmups")) + monkeypatch.setattr( + module, + "prepare_shape_state", + lambda captured_state, batch, *, for_capture: events.append( + ("prepare", captured_state, batch, for_capture) + ), + ) + backend._capture_one = lambda captured_shape_key, size, forward_fn: events.append( + ("capture", captured_shape_key, size) + ) or (graph, outputs) + + backend.capture_one(shape_key, lambda: "forward-output") + + assert events == [ + ("create", backend._runner, shape_key, 0, communicator), + "activate", + "warmups", + ("prepare", state, None, True), + ("capture", shape_key, 8), + ("attach", graph, outputs), + ("shape_states", shape_key, state), + ("graphs", shape_key, graph), + ("outputs", shape_key, outputs), + ] + assert state.binding == "restored" + assert state.graph is graph + assert state.outputs is outputs + assert backend._shape_states[shape_key] is state + assert backend._graphs[shape_key] is graph + assert backend._outputs[shape_key] is outputs + + +def test_multi_shape_load_capture_one_relocates_before_load_and_registers( + monkeypatch, tmp_path +) -> None: + module = _load_backend_module(monkeypatch) + backend, communicator = _make_backend(module, multi_shape_enabled=True) + events = [] + shape_key = FakeShapeKey(8) + graph = object() + outputs = object() + state = TrackingState(batch_size=8, capture_index=0, events=events) + backend._shape_states = RecordingDict(events, "shape_states") + backend._graphs = RecordingDict(events, "graphs") + backend._outputs = RecordingDict(events, "outputs") + flashinfer_record = SimpleNamespace(kernel_abi=module.KERNEL_ABI_SCHEMA, name="flash") + communication_record = SimpleNamespace(kernel_abi="torch.symm", name="comm") + backend._saved_records_by_batch = { + 8: SimpleNamespace( + batch_size=8, + operand_slots=(communication_record, flashinfer_record), + ) + } + graph_path = tmp_path / "graph_0_FULL_t8_r8_UX_pcN.json" + graph_path.write_text(json.dumps({"stage": "before"})) + backend._graph_load_paths = [str(graph_path)] + + monkeypatch.setattr( + module, + "get_graph_extension_mode", + lambda: module.CUDAGraphExtensionMode.LOAD, + ) + monkeypatch.setattr( + module, + "create_shape_state", + lambda runner, **kwargs: events.append( + ("create", runner, kwargs["shape_key"], kwargs["capture_index"], kwargs["communicator"]) + ) + or state, + ) + monkeypatch.setattr( + module, + "validate_shape_record", + lambda captured_state, record, *, plan_schema: events.append( + ("validate", captured_state, record, plan_schema) + ), + ) + + @contextmanager + def _activate(captured_state): + assert captured_state is state + state.binding = "active" + events.append("activate") + try: + yield + finally: + state.binding = "restored" + + monkeypatch.setattr(module, "activate_shape_state", _activate) + monkeypatch.setattr(module, "run_graph_warmups", lambda **kwargs: events.append("warmups")) + monkeypatch.setattr( + module, + "prepare_shape_state", + lambda captured_state, batch, *, for_capture: events.append( + ("prepare", captured_state, batch, for_capture) + ), + ) + monkeypatch.setattr( + module, + "relocate_flashinfer_operands", + lambda graph_data, captured_state, records: events.append( + ("relocate", graph_data["stage"], captured_state, records) + ) + or graph_data.__setitem__("stage", "relocated"), + ) + + def _load_one(captured_shape_key, size): + events.append(("load", captured_shape_key, size, backend._load_index)) + assert json.loads(graph_path.read_text())["stage"] == "relocated" + backend._load_index += 1 + return graph, outputs + + backend._load_one = _load_one + + backend.capture_one(shape_key, lambda: "forward-output") + + assert events == [ + ("create", backend._runner, shape_key, 0, communicator), + ( + "validate", + state, + backend._saved_records_by_batch[8], + module.PLAN_SCHEMA, + ), + "activate", + "warmups", + ("prepare", state, None, True), + ("relocate", "before", state, (flashinfer_record,)), + ("load", shape_key, 8, 0), + ("attach", graph, outputs), + ("shape_states", shape_key, state), + ("graphs", shape_key, graph), + ("outputs", shape_key, outputs), + ] + assert state.binding == "restored" + assert state.graph is graph + assert state.outputs is outputs + assert backend._load_index == 1 + assert backend._shape_states[shape_key] is state + assert backend._graphs[shape_key] is graph + assert backend._outputs[shape_key] is outputs + + +def test_multi_shape_replay_adopts_live_padded_metadata_without_replanning(monkeypatch) -> None: + module = _load_backend_module(monkeypatch) + backend, _communicator = _make_backend(module, multi_shape_enabled=True) + events = [] + shape_key = FakeShapeKey(8) + state = TrackingState(batch_size=8, capture_index=0, events=events) + padded_metadata = object() + state.metadata = "stale-metadata" + state.attention_backend = SimpleNamespace( + forward_metadata=padded_metadata, + decode_cuda_graph_metadata={8: ["wrapper"]}, + ) + + class _ReplayGraph: + def replay(self) -> None: + events.append(("replay", state.binding, state.metadata)) + + state.graph = _ReplayGraph() + state.outputs = object() + backend._shape_states = {shape_key: state} + + @contextmanager + def _activate(captured_state): + assert captured_state is state + state.binding = "active" + events.append("activate") + try: + yield + finally: + state.binding = "restored" + + monkeypatch.setattr(module, "activate_shape_state", _activate) + monkeypatch.setattr( + module, + "adopt_prepared_shape_state", + lambda captured_state: events.append( + ( + "adopt", + captured_state, + captured_state.attention_backend.forward_metadata, + ) + ) + or setattr(captured_state, "metadata", captured_state.attention_backend.forward_metadata), + ) + monkeypatch.setattr( + module, + "prepare_shape_state", + lambda *args, **kwargs: (_ for _ in ()).throw( + AssertionError("replay must not call prepare_shape_state") + ), + ) + batch = SimpleNamespace(batch_size=7) + + assert backend.replay(shape_key, batch) is state.outputs + assert events == [ + ("adopt", state, padded_metadata), + "activate", + ("replay", "active", padded_metadata), + ] + assert state.binding == "restored" + assert state.metadata is padded_metadata + + +def test_finish_save_failure_removes_old_shape_manifest(monkeypatch, tmp_path) -> None: + module = _load_backend_module(monkeypatch) + backend, _communicator = _make_backend( + module, + multi_shape_enabled=False, + symmetric_memory_enabled=False, + ) + manifest_path = tmp_path / "sglang_shape_state.json" + manifest_path.write_text('{"stale": true}') + monkeypatch.setattr( + module, + "get_config", + lambda: SimpleNamespace(workspace_dir=str(tmp_path)), + ) + monkeypatch.setattr(module, "clear_symmetric_graph_state", lambda workspace: None) + monkeypatch.setattr( + module, + "write_graph_backend_metadata", + lambda workspace, symmetric_memory_enabled: (_ for _ in ()).throw( + RuntimeError("injected publish failure") + ), + ) + + with pytest.raises(RuntimeError, match="injected publish failure"): + backend._finish_save() + + assert not manifest_path.exists() diff --git a/tests/test_sglang_main_compat.py b/tests/test_sglang_main_compat.py new file mode 100644 index 00000000..8b4885b5 --- /dev/null +++ b/tests/test_sglang_main_compat.py @@ -0,0 +1,346 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the Foundry project +"""CPU-only contracts for the SGLang-main integration surface.""" + +from __future__ import annotations + +import importlib.util +import sys +from pathlib import Path +from types import ModuleType, SimpleNamespace + +import pytest +import tomllib + +SGLANG_DIR = Path(__file__).parents[1] / "python" / "foundry" / "integration" / "sglang" + + +def _ensure_package(name: str) -> ModuleType: + package = sys.modules.get(name) + if package is None: + package = ModuleType(name) + package.__path__ = [] # type: ignore[attr-defined] + sys.modules[name] = package + return package + + +def _load_sglang_module(module_name: str, filename: str) -> ModuleType: + spec = importlib.util.spec_from_file_location(module_name, SGLANG_DIR / filename) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + sys.modules[module_name] = module + spec.loader.exec_module(module) + return module + + +def _stub(monkeypatch: pytest.MonkeyPatch, name: str) -> ModuleType: + module = ModuleType(name) + monkeypatch.setitem(sys.modules, name, module) + return module + + +def _load_config_module() -> ModuleType: + foundry_pkg = _ensure_package("foundry") + integration_pkg = _ensure_package("foundry.integration") + sglang_pkg = _ensure_package("foundry.integration.sglang") + foundry_pkg.integration = integration_pkg + integration_pkg.sglang = sglang_pkg + config = _load_sglang_module("foundry.integration.sglang.config", "config.py") + sglang_pkg.config = config + return config + + +def _load_hooks_main(monkeypatch: pytest.MonkeyPatch) -> tuple[ModuleType, ModuleType]: + foundry_pkg = _ensure_package("foundry") + integration_pkg = _ensure_package("foundry.integration") + sglang_pkg = _ensure_package("foundry.integration.sglang") + foundry_pkg.integration = integration_pkg + integration_pkg.sglang = sglang_pkg + config = _load_sglang_module("foundry.integration.sglang.config", "config.py") + sglang_pkg.config = config + + for name in ( + "sglang", + "sglang.srt", + "sglang.srt.distributed", + "sglang.srt.distributed.device_communicators", + "sglang.srt.entrypoints", + "sglang.srt.managers", + "sglang.srt.mem_cache", + "sglang.srt.model_executor", + "sglang.srt.model_executor.runner", + "sglang.srt.model_executor.runner_backend", + ): + package = _stub(monkeypatch, name) + package.__path__ = [] # type: ignore[attr-defined] + + decode_runner = _stub( + monkeypatch, + "sglang.srt.model_executor.runner.decode_cuda_graph_runner", + ) + + def get_batch_sizes_to_capture(model_runner, captured_req_width): + del model_runner, captured_req_width + return [1, 2, 4, 8], [1, 2, 4, 8] + + decode_runner.get_batch_sizes_to_capture = get_batch_sizes_to_capture + decode_runner.resolve_decode_backend = lambda cuda_graph_runner: cuda_graph_runner + + torch_symm_mem = _stub( + monkeypatch, + "sglang.srt.distributed.device_communicators.torch_symm_mem", + ) + torch_symm_mem.TorchSymmMemCommunicator = type( + "TorchSymmMemCommunicator", + (), + {"__init__": lambda self, *args, **kwargs: None}, + ) + + triton_symm_mem_ag = _stub( + monkeypatch, + "sglang.srt.distributed.device_communicators.triton_symm_mem_ag", + ) + triton_symm_mem_ag.MultimemAllGatherer = type( + "MultimemAllGatherer", + (), + {"__init__": lambda self, *args, **kwargs: None}, + ) + + engine_mod = _stub(monkeypatch, "sglang.srt.entrypoints.engine") + engine_mod.Engine = type( + "Engine", + (), + {"_launch_scheduler_processes": staticmethod(lambda owner, *args, **kwargs: None)}, + ) + + dpc = _stub(monkeypatch, "sglang.srt.managers.data_parallel_controller") + dpc.DataParallelController = type( + "DataParallelController", + (), + {"launch_tensor_parallel_group": lambda self, *args, **kwargs: None}, + ) + + kv_cache_configurator = _stub(monkeypatch, "sglang.srt.mem_cache.kv_cache_configurator") + kv_cache_configurator.KVCacheConfigurator = type( + "KVCacheConfigurator", + (), + { + "_resolve_memory_pool_config": lambda self, pre_model_load_memory: None, + "configure": lambda self, *args, **kwargs: None, + }, + ) + + model_runner = _stub(monkeypatch, "sglang.srt.model_executor.model_runner") + model_runner.ModelRunner = type( + "ModelRunner", + (), + {"init_torch_distributed": lambda self, *args, **kwargs: None}, + ) + + pool_configurator = _stub(monkeypatch, "sglang.srt.model_executor.pool_configurator") + pool_configurator.MemoryPoolConfig = dict + + backend_utils = _stub(monkeypatch, "sglang.srt.model_executor.runner_backend.utils") + backend_utils.resolve_decode_backend = lambda cuda_graph_runner: cuda_graph_runner + + foundry_ops = _stub(monkeypatch, "foundry.ops") + foundry_ops.stop_allocation_region = lambda: None + foundry_ops.resume_allocation_region = lambda: None + foundry_pkg.ops = foundry_ops + + runtime = _stub(monkeypatch, "foundry.integration.sglang.runtime") + runtime.setup_graph_extension = lambda *args, **kwargs: None + runtime.log_alloc_offset = lambda label: None + runtime.skip_to_scratch_boundary = lambda: None + runtime.load_warmup_state = lambda: None + runtime.create_warmup_state = lambda config: None + runtime.save_warmup_state = lambda state: None + runtime.setup_ld_preload_env = lambda: None + sglang_pkg.runtime = runtime + + main_backend = _stub(monkeypatch, "foundry.integration.sglang.main_backend") + main_backend.FoundryMainCudaGraphBackend = object + sglang_pkg.main_backend = main_backend + + spec = importlib.util.spec_from_file_location( + "foundry.integration.sglang.hooks_main", + SGLANG_DIR / "hooks_main.py", + ) + assert spec is not None and spec.loader is not None + hooks_main = importlib.util.module_from_spec(spec) + sys.modules["foundry.integration.sglang.hooks_main"] = hooks_main + spec.loader.exec_module(hooks_main) + return hooks_main, decode_runner + + +_config = _load_config_module() +symmetric_graph_batch_sizes = _config.symmetric_graph_batch_sizes + + +def test_symmetric_graph_batch_sizes_are_strict(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("FOUNDRY_SGLANG_SYMM_GRAPH_BATCHES", "1,8,32") + assert symmetric_graph_batch_sizes() == (1, 8, 32) + + monkeypatch.setenv("FOUNDRY_SGLANG_SYMM_GRAPH_BATCHES", "8,1") + with pytest.raises(RuntimeError, match="ascending"): + symmetric_graph_batch_sizes() + + monkeypatch.setenv("FOUNDRY_SGLANG_SYMM_GRAPH_BATCHES", "1,1") + with pytest.raises(RuntimeError, match="duplicate"): + symmetric_graph_batch_sizes() + + monkeypatch.setenv("FOUNDRY_SGLANG_SYMM_GRAPH_BATCHES", "1,,8") + with pytest.raises(RuntimeError, match="comma-separated integers"): + symmetric_graph_batch_sizes() + + monkeypatch.setenv("FOUNDRY_SGLANG_SYMM_GRAPH_BATCHES", "1,abc") + with pytest.raises(RuntimeError, match="comma-separated integers"): + symmetric_graph_batch_sizes() + + monkeypatch.setenv("FOUNDRY_SGLANG_SYMM_GRAPH_BATCHES", "0,8") + with pytest.raises(RuntimeError, match="positive"): + symmetric_graph_batch_sizes() + + monkeypatch.setenv("FOUNDRY_SGLANG_SYMM_GRAPH_BATCHES", "2,8") + with pytest.raises(RuntimeError, match="batch 1"): + symmetric_graph_batch_sizes() + + +def _symm_mem_model_runner() -> SimpleNamespace: + return SimpleNamespace( + server_args=SimpleNamespace(enable_torch_symm_mem=True), + ) + + +def test_capture_batch_sizes_hook_filters_requested_shapes( + monkeypatch: pytest.MonkeyPatch, +) -> None: + hooks_main, decode_runner = _load_hooks_main(monkeypatch) + hooks_main.install_hooks_main() + + monkeypatch.setenv("FOUNDRY_SGLANG_SYMM_MULTI_SHAPE", "1") + monkeypatch.setenv("FOUNDRY_SGLANG_SYMM_GRAPH_BATCHES", "1,8") + capture_bs, compile_bs = decode_runner.get_batch_sizes_to_capture( + _symm_mem_model_runner(), + 1, + ) + assert capture_bs == [1, 8] + assert compile_bs == [1, 8] + + +def test_capture_batch_sizes_hook_requires_explicit_schedule( + monkeypatch: pytest.MonkeyPatch, +) -> None: + hooks_main, decode_runner = _load_hooks_main(monkeypatch) + hooks_main.install_hooks_main() + + monkeypatch.setenv("FOUNDRY_SGLANG_SYMM_MULTI_SHAPE", "1") + monkeypatch.delenv("FOUNDRY_SGLANG_SYMM_GRAPH_BATCHES", raising=False) + with pytest.raises(RuntimeError, match="FOUNDRY_SGLANG_SYMM_GRAPH_BATCHES"): + decode_runner.get_batch_sizes_to_capture(_symm_mem_model_runner(), 1) + + +def test_capture_batch_sizes_hook_rejects_unknown_batches( + monkeypatch: pytest.MonkeyPatch, +) -> None: + hooks_main, decode_runner = _load_hooks_main(monkeypatch) + hooks_main.install_hooks_main() + + monkeypatch.setenv("FOUNDRY_SGLANG_SYMM_MULTI_SHAPE", "1") + monkeypatch.setenv("FOUNDRY_SGLANG_SYMM_GRAPH_BATCHES", "1,16") + with pytest.raises(RuntimeError, match="unavailable"): + decode_runner.get_batch_sizes_to_capture(_symm_mem_model_runner(), 1) + + +def test_foundry_registers_an_sglang_main_plugin() -> None: + config = tomllib.loads((Path(__file__).parents[1] / "pyproject.toml").read_text()) + + plugins = config["project"]["entry-points"]["sglang.srt.plugins"] + assert plugins["foundry"] == "foundry.integration.sglang.plugin:register" + + +def test_main_plugin_installs_hooks_and_configures_full_decode( + monkeypatch, + tmp_path: Path, +) -> None: + registrations = [] + installs = [] + + class _HookType: + BEFORE = "before" + AFTER = "after" + + class _HookRegistry: + @classmethod + def register(cls, target, hook, hook_type): + registrations.append((target, hook, hook_type)) + + hook_registry = ModuleType("sglang.srt.plugins.hook_registry") + hook_registry.HookRegistry = _HookRegistry + hook_registry.HookType = _HookType + monkeypatch.setitem( + sys.modules, + "sglang.srt.plugins.hook_registry", + hook_registry, + ) + + config_path = tmp_path / "foundry_save.toml" + config_path.write_text( + 'mode = "save"\nworkspace_root = "archive"\n', + ) + monkeypatch.setenv("FOUNDRY_SGLANG_CONFIG_PATH", str(config_path)) + + hooks = ModuleType("foundry.integration.sglang.hooks") + hooks.install_hooks_from_path = lambda path: installs.append(path) + monkeypatch.setitem(sys.modules, "foundry.integration.sglang.hooks", hooks) + + plugin_path = ( + Path(__file__).parents[1] / "python" / "foundry" / "integration" / "sglang" / "plugin.py" + ) + spec = importlib.util.spec_from_file_location( + "foundry.integration.sglang.plugin", + plugin_path, + ) + assert spec is not None and spec.loader is not None + plugin = importlib.util.module_from_spec(spec) + spec.loader.exec_module(plugin) + + plugin.register() + + assert installs == [str(config_path)] + assert len(registrations) == 2 + target, before_post_init, before_type = registrations[0] + after_target, after_post_init, after_type = registrations[1] + assert target == after_target == "sglang.srt.server_args.ServerArgs.__post_init__" + assert before_type == _HookType.BEFORE + assert after_type == _HookType.AFTER + + server_args = SimpleNamespace( + cuda_graph_backend_decode=None, + disable_prefill_cuda_graph=False, + disable_flashinfer_autotune=False, + enable_profile_cuda_graph=True, + pp_size=1, + speculative_algorithm=None, + enable_lora=False, + enable_pdmux=False, + enable_return_hidden_states=False, + enable_torch_symm_mem=False, + dllm_algorithm=None, + moe_a2a_backend="none", + ) + before_post_init(server_args) + + assert server_args.cuda_graph_backend_decode == "full" + assert server_args.disable_prefill_cuda_graph is True + assert server_args.disable_flashinfer_autotune is True + assert server_args.enable_profile_cuda_graph is False + + server_args.cuda_graph_config = SimpleNamespace( + decode=SimpleNamespace(backend="full"), + prefill=SimpleNamespace(backend="disabled"), + ) + after_post_init(None, server_args) + + server_args.enable_torch_symm_mem = True + after_post_init(None, server_args) diff --git a/tests/test_sglang_shape_replay_state.py b/tests/test_sglang_shape_replay_state.py new file mode 100644 index 00000000..34a49cb3 --- /dev/null +++ b/tests/test_sglang_shape_replay_state.py @@ -0,0 +1,572 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the Foundry project +"""CPU-only contract coverage for SGLang shape replay state lifecycle.""" + +from __future__ import annotations + +import gc +import importlib.util +import sys +import weakref +from pathlib import Path +from types import ModuleType, SimpleNamespace + +import pytest + +SGLANG_DIR = Path(__file__).parents[1] / "python" / "foundry" / "integration" / "sglang" + + +def _ensure_package(name: str) -> ModuleType: + package = sys.modules.get(name) + if package is None: + package = ModuleType(name) + package.__path__ = [] # type: ignore[attr-defined] + sys.modules[name] = package + return package + + +def _load_module(module_name: str, filename: str): + module_path = SGLANG_DIR / filename + spec = importlib.util.spec_from_file_location(module_name, module_path) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + sys.modules[module_name] = module + spec.loader.exec_module(module) + return module + + +foundry_pkg = _ensure_package("foundry") +integration_pkg = _ensure_package("foundry.integration") +sglang_pkg = _ensure_package("foundry.integration.sglang") +foundry_pkg.integration = integration_pkg +integration_pkg.sglang = sglang_pkg + +_mod = _load_module( + "foundry.integration.sglang.shape_replay_state", + "shape_replay_state.py", +) +sglang_pkg.shape_replay_state = _mod +ShapeReplayState = _mod.ShapeReplayState +ShapeStateLifecycle = _mod.ShapeStateLifecycle +TensorOwnerSlot = _mod.TensorOwnerSlot + +try: + _adapter_mod = _load_module( + "foundry.integration.sglang.sglang_shape_adapter", + "sglang_shape_adapter.py", + ) +except FileNotFoundError as error: + _adapter_mod = None + _adapter_import_error = error +else: + sglang_pkg.sglang_shape_adapter = _adapter_mod + _adapter_import_error = None + + +def _require_adapter(): + if _adapter_import_error is not None: + pytest.fail(f"Unable to load sglang_shape_adapter.py: {_adapter_import_error}") + assert _adapter_mod is not None + return _adapter_mod + + +class FakeTensor: + def __init__(self, ptr: int, numel: int = 16) -> None: + self._ptr = ptr + self._numel = numel + self.dtype = "torch.int32" + self.device = "cuda:0" + + def data_ptr(self) -> int: + return self._ptr + + def numel(self) -> int: + return self._numel + + def element_size(self) -> int: + return 4 + + +class FakeWrapper: + def __init__( + self, + base: int, + shared_base: int = 0x600000100000, + fixed_batch_size: int = 8, + ) -> None: + self._int_workspace_buffer = FakeTensor(base) + self._pin_memory_int_workspace_buffer = FakeTensor(base + 0x1000) + self._qo_indptr_buf = FakeTensor(base + 0x2000) + self._paged_kv_indptr_buf = FakeTensor(shared_base + 0x1000) + self._paged_kv_indices_buf = FakeTensor(shared_base + 0x2000) + self._paged_kv_last_page_len_buf = FakeTensor(shared_base + 0x3000) + self._float_workspace_buffer = FakeTensor(shared_base + 0x4000) + self._fixed_batch_size = fixed_batch_size + self._workspace_size = 8 * 1024 * 1024 + self._cached_module = object() + self._plan_info = list(range(15)) + + +class FailOnceMapping(dict): + def __init__( + self, + *args, + fail_key: int, + error_message: str = "injected wrapper install failure", + **kwargs, + ) -> None: + self._armed = False + super().__init__(*args, **kwargs) + self._fail_key = fail_key + self._error_message = error_message + self._failed = False + self._armed = True + + def __setitem__(self, key, value) -> None: + super().__setitem__(key, value) + if self._armed and key == self._fail_key and not self._failed: + self._failed = True + raise RuntimeError(self._error_message) + + +class FakeAttentionBackend: + def __init__( + self, + wrappers_by_batch: dict[int, list[FakeWrapper]], + *, + decode_cuda_graph_metadata=None, + initial_metadata: object | None = None, + on_prepare=None, + fail_forward_metadata_set_once: bool = False, + forward_metadata_error_message: str = "injected metadata install failure", + ) -> None: + self.decode_cuda_graph_metadata = ( + decode_cuda_graph_metadata + if decode_cuda_graph_metadata is not None + else {batch_size: list(wrappers) for batch_size, wrappers in wrappers_by_batch.items()} + ) + self._forward_metadata = object() if initial_metadata is None else initial_metadata + self._fail_forward_metadata_set_once = fail_forward_metadata_set_once + self._forward_metadata_error_message = forward_metadata_error_message + self._forward_metadata_set_failed = False + self.prepared = [] + self.on_prepare = on_prepare + + @property + def forward_metadata(self): + return self._forward_metadata + + @forward_metadata.setter + def forward_metadata(self, value) -> None: + self._forward_metadata = value + if self._fail_forward_metadata_set_once and not self._forward_metadata_set_failed: + self._forward_metadata_set_failed = True + raise RuntimeError(self._forward_metadata_error_message) + + def init_forward_metadata_out_graph(self, forward_batch, in_capture=False): + self.prepared.append((forward_batch, in_capture)) + if self.on_prepare is not None: + self.on_prepare(self, forward_batch, in_capture) + + +def make_state(tensor: FakeTensor) -> ShapeReplayState: + slot = TensorOwnerSlot.capture("int_workspace", tensor, ownership="shape") + return ShapeReplayState( + shape_key=1, + batch_size=1, + capture_index=0, + attention_backend=object(), + wrappers=(object(),), + metadata=object(), + owners=(slot,), + shared_owners=(), + shared_objects=(), + communicator=object(), + symm_handle=object(), + ) + + +def make_runner( + wrapper: FakeWrapper | None = None, + *, + wrappers_by_batch: dict[int, FakeWrapper] | None = None, + backend: FakeAttentionBackend | None = None, +): + if backend is None: + if wrappers_by_batch is None: + assert wrapper is not None + wrappers_by_batch = {8: wrapper} + backend = FakeAttentionBackend( + {batch_size: [shape_wrapper] for batch_size, shape_wrapper in wrappers_by_batch.items()} + ) + model_runner = SimpleNamespace( + attn_backend=backend, + graph_shared_output=object(), + ) + return SimpleNamespace( + captured_req_width=1, + attn_backend=backend, + model_runner=model_runner, + buffers=object(), + ) + + +def test_owner_survives_until_state_close() -> None: + tensor = FakeTensor(0x600000001000) + reference = weakref.ref(tensor) + state = make_state(tensor) + del tensor + gc.collect() + assert reference() is not None + + state.mark_planned(tuple(range(15))) + state.attach_graph(object(), object()) + state.close() + gc.collect() + + assert state.lifecycle is ShapeStateLifecycle.CLOSED + assert reference() is None + + +def test_owner_pointer_replacement_is_rejected() -> None: + tensor = FakeTensor(0x600000001000) + state = make_state(tensor) + tensor._ptr = 0x600000002000 + + with pytest.raises(RuntimeError, match="int_workspace"): + state.validate_owners() + + +def test_invalid_lifecycle_transition_is_rejected() -> None: + state = make_state(FakeTensor(0x600000001000)) + + with pytest.raises(RuntimeError, match="planned"): + state.attach_graph(object(), object()) + + +def test_adapter_retains_explicit_owner_allowlist() -> None: + adapter = _require_adapter() + wrapper = FakeWrapper(0x600001000000) + runner = make_runner(wrapper) + state = adapter.create_shape_state( + runner, + shape_key=8, + capture_index=0, + communicator=object(), + symm_handle=object(), + ) + + assert adapter.PLAN_SCHEMA == "flashinfer.prefill.v15" + assert tuple(slot.name for slot in state.owners) == ( + "int_workspace", + "pin_memory_int_workspace", + "qo_indptr", + ) + assert tuple(slot.name for slot in state.shared_owners) == ( + "paged_kv_indptr", + "paged_kv_indices", + "paged_kv_last_page_len", + "float_workspace", + ) + assert state.lifecycle is ShapeStateLifecycle.PLANNED + assert state.plan_fingerprint == tuple(range(15)) + + +def test_activation_restores_previous_backend_bindings() -> None: + adapter = _require_adapter() + wrapper = FakeWrapper(0x600001000000) + runner = make_runner(wrapper) + state = adapter.create_shape_state( + runner, + shape_key=8, + capture_index=0, + communicator=object(), + symm_handle=object(), + ) + previous_wrappers = runner.attn_backend.decode_cuda_graph_metadata[8] + previous_metadata = runner.attn_backend.forward_metadata + refreshed_metadata = object() + + with adapter.activate_shape_state(state): + assert state.active is True + assert runner.attn_backend.decode_cuda_graph_metadata[8][0] is wrapper + assert runner.attn_backend.forward_metadata is state.metadata + runner.attn_backend.forward_metadata = refreshed_metadata + + assert state.active is False + assert state.metadata is refreshed_metadata + assert runner.attn_backend.decode_cuda_graph_metadata[8] is previous_wrappers + assert runner.attn_backend.forward_metadata is previous_metadata + + +def test_activation_rejects_nested_or_closed_state() -> None: + adapter = _require_adapter() + state = adapter.create_shape_state( + make_runner(FakeWrapper(0x600001000000)), + shape_key=8, + capture_index=0, + communicator=object(), + symm_handle=object(), + ) + + with ( + adapter.activate_shape_state(state), + pytest.raises(RuntimeError, match="already active"), + adapter.activate_shape_state(state), + ): + pass + + adapter.close_shape_state(state) + + with pytest.raises(RuntimeError, match="closed"), adapter.activate_shape_state(state): + pass + + +def test_replay_prepare_reuses_wrapper_identity() -> None: + adapter = _require_adapter() + wrapper = FakeWrapper(0x600001000000) + runner = make_runner(wrapper) + state = adapter.create_shape_state( + runner, + shape_key=8, + capture_index=0, + communicator=object(), + symm_handle=object(), + ) + batch = SimpleNamespace(batch_size=8) + + with adapter.activate_shape_state(state): + adapter.prepare_shape_state(state, batch, for_capture=False) + + assert runner.attn_backend.prepared == [(batch, False)] + assert state.wrappers == (wrapper,) + + +def test_adopt_prepared_shape_state_uses_live_padded_metadata_without_replanning() -> None: + adapter = _require_adapter() + wrapper = FakeWrapper(0x600001000000) + runner = make_runner(wrapper) + state = adapter.create_shape_state( + runner, + shape_key=8, + capture_index=0, + communicator=object(), + symm_handle=object(), + ) + padded_metadata = object() + runner.attn_backend.forward_metadata = padded_metadata + + adapter.adopt_prepared_shape_state(state) + + assert state.metadata is padded_metadata + assert runner.attn_backend.prepared == [] + + +def test_adopt_prepared_shape_state_rejects_live_wrapper_replacement() -> None: + adapter = _require_adapter() + wrapper = FakeWrapper(0x600001000000) + replacement = FakeWrapper(0x600003000000) + runner = make_runner(wrapper) + state = adapter.create_shape_state( + runner, + shape_key=8, + capture_index=0, + communicator=object(), + symm_handle=object(), + ) + runner.attn_backend.decode_cuda_graph_metadata[8] = [replacement] + + with pytest.raises(RuntimeError, match="preserve wrapper identity"): + adapter.adopt_prepared_shape_state(state) + + +def test_adopt_prepared_shape_state_rejects_missing_live_metadata() -> None: + adapter = _require_adapter() + wrapper = FakeWrapper(0x600001000000) + runner = make_runner(wrapper) + state = adapter.create_shape_state( + runner, + shape_key=8, + capture_index=0, + communicator=object(), + symm_handle=object(), + ) + runner.attn_backend.forward_metadata = None + + with pytest.raises(RuntimeError, match="no live forward_metadata"): + adapter.adopt_prepared_shape_state(state) + + +def test_capture_prepare_rejects_backend_wrapper_replacement() -> None: + adapter = _require_adapter() + wrapper = FakeWrapper(0x600001000000) + replacement = FakeWrapper(0x600003000000) + runner = make_runner(wrapper) + state = adapter.create_shape_state( + runner, + shape_key=8, + capture_index=0, + communicator=object(), + symm_handle=object(), + ) + + with adapter.activate_shape_state(state): + runner.attn_backend.decode_cuda_graph_metadata[8] = [replacement] + with pytest.raises(RuntimeError, match="preserve wrapper identity"): + adapter.prepare_shape_state(state, None, for_capture=True) + + +def test_prepare_rejects_backend_wrapper_replacement_and_preserves_metadata() -> None: + adapter = _require_adapter() + wrapper = FakeWrapper(0x600001000000) + replacement_wrapper = FakeWrapper(0x600003000000) + replacement_metadata = object() + + def replace_wrapper(backend, _forward_batch, _in_capture) -> None: + backend.decode_cuda_graph_metadata[8] = [replacement_wrapper] + backend.forward_metadata = replacement_metadata + + backend = FakeAttentionBackend( + {8: [wrapper]}, + initial_metadata=object(), + on_prepare=replace_wrapper, + ) + runner = make_runner(backend=backend) + state = adapter.create_shape_state( + runner, + shape_key=8, + capture_index=0, + communicator=object(), + symm_handle=object(), + ) + batch = SimpleNamespace(batch_size=8) + previous_wrappers = backend.decode_cuda_graph_metadata[8] + previous_metadata = backend.forward_metadata + + with ( + pytest.raises(RuntimeError, match="preserve wrapper identity"), + adapter.activate_shape_state(state), + ): + adapter.prepare_shape_state(state, batch, for_capture=False) + + assert backend.decode_cuda_graph_metadata[8] is previous_wrappers + assert backend.forward_metadata is previous_metadata + assert state.metadata is previous_metadata + + +def test_activation_recovers_from_side_effecting_wrapper_install_failure() -> None: + adapter = _require_adapter() + wrapper = FakeWrapper(0x600001000000) + previous_metadata = object() + previous_wrappers = [wrapper] + decode_cuda_graph_metadata = FailOnceMapping( + {8: previous_wrappers}, + fail_key=8, + error_message="injected wrapper install failure", + ) + backend = FakeAttentionBackend( + {8: [wrapper]}, + decode_cuda_graph_metadata=decode_cuda_graph_metadata, + initial_metadata=previous_metadata, + ) + runner = make_runner(backend=backend) + state = adapter.create_shape_state( + runner, + shape_key=8, + capture_index=0, + communicator=object(), + symm_handle=object(), + ) + + with ( + pytest.raises(RuntimeError, match="injected wrapper install failure"), + adapter.activate_shape_state(state), + ): + pass + + assert state.active is False + assert backend.decode_cuda_graph_metadata[8] is previous_wrappers + assert backend.forward_metadata is previous_metadata + assert state.metadata is previous_metadata + + +def test_activation_recovers_from_side_effecting_metadata_install_failure() -> None: + adapter = _require_adapter() + wrapper = FakeWrapper(0x600001000000) + previous_metadata = object() + backend = FakeAttentionBackend( + {8: [wrapper]}, + initial_metadata=previous_metadata, + fail_forward_metadata_set_once=True, + forward_metadata_error_message="injected metadata install failure", + ) + runner = make_runner(backend=backend) + state = adapter.create_shape_state( + runner, + shape_key=8, + capture_index=0, + communicator=object(), + symm_handle=object(), + ) + previous_wrappers = backend.decode_cuda_graph_metadata[8] + + with ( + pytest.raises(RuntimeError, match="injected metadata install failure"), + adapter.activate_shape_state(state), + ): + pass + + assert state.active is False + assert backend.decode_cuda_graph_metadata[8] is previous_wrappers + assert backend.forward_metadata is previous_metadata + assert state.metadata is previous_metadata + + +def test_shapes_share_only_approved_backend_buffers() -> None: + adapter = _require_adapter() + runner = make_runner( + wrappers_by_batch={ + 1: FakeWrapper(0x600001000000, fixed_batch_size=1), + 8: FakeWrapper(0x600002000000, fixed_batch_size=8), + } + ) + first = adapter.create_shape_state( + runner, + shape_key=1, + capture_index=0, + communicator=object(), + symm_handle=object(), + ) + second = adapter.create_shape_state( + runner, + shape_key=8, + capture_index=1, + communicator=object(), + symm_handle=object(), + ) + + assert first.attention_backend is second.attention_backend + assert {slot.data_ptr for slot in first.owners}.isdisjoint( + slot.data_ptr for slot in second.owners + ) + assert tuple(slot.data_ptr for slot in first.shared_owners) == tuple( + slot.data_ptr for slot in second.shared_owners + ) + + +def test_close_shape_state_defers_active_guard_until_context_exit() -> None: + adapter = _require_adapter() + state = adapter.create_shape_state( + make_runner(FakeWrapper(0x600001000000)), + shape_key=8, + capture_index=0, + communicator=object(), + symm_handle=object(), + ) + + with adapter.activate_shape_state(state), pytest.raises(RuntimeError, match="active"): + adapter.close_shape_state(state) + + adapter.close_shape_state(state) + + assert state.lifecycle is ShapeStateLifecycle.CLOSED diff --git a/tests/test_sglang_state_manifest.py b/tests/test_sglang_state_manifest.py new file mode 100644 index 00000000..9e2a4280 --- /dev/null +++ b/tests/test_sglang_state_manifest.py @@ -0,0 +1,202 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the Foundry project +"""CPU-only contract coverage for SGLang shape state manifest.""" + +from __future__ import annotations + +import importlib.util +import json +import sys +from dataclasses import replace +from pathlib import Path +from types import ModuleType, SimpleNamespace + +import pytest + +SGLANG_DIR = Path(__file__).parents[1] / "python" / "foundry" / "integration" / "sglang" + + +def _ensure_package(name: str) -> ModuleType: + package = sys.modules.get(name) + if package is None: + package = ModuleType(name) + package.__path__ = [] # type: ignore[attr-defined] + sys.modules[name] = package + return package + + +def _load_module(): + module_path = SGLANG_DIR / "state_manifest.py" + spec = importlib.util.spec_from_file_location( + "foundry.integration.sglang.state_manifest", + module_path, + ) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) + return module + + +foundry_pkg = _ensure_package("foundry") +integration_pkg = _ensure_package("foundry.integration") +sglang_pkg = _ensure_package("foundry.integration.sglang") +foundry_pkg.integration = integration_pkg +integration_pkg.sglang = sglang_pkg + + +_mod = _load_module() +sglang_pkg.state_manifest = _mod +MANIFEST_FILENAME = _mod.MANIFEST_FILENAME +OperandSlotRecord = _mod.OperandSlotRecord +ShapeStateManifest = _mod.ShapeStateManifest +ShapeStateRecord = _mod.ShapeStateRecord +TensorSlotRecord = _mod.TensorSlotRecord +build_shape_record = _mod.build_shape_record +read_shape_state_manifest = _mod.read_shape_state_manifest +validate_shape_record = _mod.validate_shape_record +validate_shape_state_manifest = _mod.validate_shape_state_manifest +write_shape_state_manifest = _mod.write_shape_state_manifest + + +def make_manifest() -> ShapeStateManifest: + tensor_slot = TensorSlotRecord( + name="int_workspace", + ownership="shape", + dtype="torch.uint8", + device="cuda:0", + numel=8 * 1024 * 1024, + element_size=1, + ) + operand = OperandSlotRecord( + name="symm.buffer_ptrs_dev", + kernel_abi="torch.symm_mem.two_shot.bf16.align16.tp2", + owner_slot="communicator", + node_id=1, + source="kernelParams", + parameter_index=0, + cuda_parameter_offset=0, + value_byte_offset=0, + ctype="BFloat16**", + owner_relative_offset=0, + span_bytes=16, + saved_value=0x600006400400, + ) + shape = ShapeStateRecord( + graph_filename="graph_0_FULL_t1_r1_UX_pcN.json", + batch_size=1, + capture_index=0, + plan_schema="flashinfer.prefill.v15", + plan_fingerprint=tuple(range(15)), + tensor_slots=(tensor_slot,), + operand_slots=(operand,), + ) + return ShapeStateManifest( + backend="torch_symmetric_memory", + rank=0, + world_size=2, + capture_order=(1,), + shapes=(shape,), + ) + + +def test_manifest_round_trip(tmp_path) -> None: + manifest = make_manifest() + write_shape_state_manifest(tmp_path, manifest) + + assert read_shape_state_manifest(tmp_path) == manifest + + +def test_manifest_write_replaces_target_atomically(monkeypatch, tmp_path) -> None: + manifest = make_manifest() + path = tmp_path / MANIFEST_FILENAME + replace_calls = [] + original_replace = Path.replace + + def _record_replace(self, target): + replace_calls.append((self, Path(target))) + return original_replace(self, target) + + monkeypatch.setattr(Path, "replace", _record_replace) + + write_shape_state_manifest(tmp_path, manifest) + + assert read_shape_state_manifest(tmp_path) == manifest + assert len(replace_calls) == 1 + temp_path, final_path = replace_calls[0] + assert final_path == path + assert temp_path.parent == final_path.parent + assert temp_path != final_path + assert not temp_path.exists() + + +def test_manifest_rejects_duplicate_shape(tmp_path) -> None: + manifest = make_manifest() + duplicate = ShapeStateManifest( + backend=manifest.backend, + rank=manifest.rank, + world_size=manifest.world_size, + capture_order=(1, 1), + shapes=(manifest.shapes[0], manifest.shapes[0]), + ) + + with pytest.raises(RuntimeError, match="duplicate"): + write_shape_state_manifest(tmp_path, duplicate) + + +def test_public_manifest_validator_rejects_invalid_manifest() -> None: + manifest = make_manifest() + invalid = ShapeStateManifest( + backend=manifest.backend, + rank=manifest.rank, + world_size=3, + capture_order=manifest.capture_order, + shapes=manifest.shapes, + ) + + with pytest.raises(RuntimeError, match="requires TP=2"): + validate_shape_state_manifest(invalid) + + +def test_manifest_rejects_unknown_version(tmp_path) -> None: + path = tmp_path / MANIFEST_FILENAME + path.write_text(json.dumps({"version": 999})) + + with pytest.raises(RuntimeError, match="version"): + read_shape_state_manifest(tmp_path) + + +def test_live_shape_must_match_manifest_record() -> None: + owner = SimpleNamespace( + name="int_workspace", + ownership="shape", + dtype="torch.uint8", + device="cuda:0", + numel=1024, + element_size=1, + ) + state = SimpleNamespace( + batch_size=1, + capture_index=0, + plan_fingerprint=tuple(range(15)), + owners=(owner,), + shared_owners=(), + ) + record = build_shape_record( + state, + graph_filename="graph_0_FULL_t1_r1_UX_pcN.json", + plan_schema="flashinfer.prefill.v15", + operand_slots=(), + ) + validate_shape_record( + state, + record, + plan_schema="flashinfer.prefill.v15", + ) + + with pytest.raises(RuntimeError, match="live shape state"): + validate_shape_record( + state, + replace(record, plan_fingerprint=(99,)), + plan_schema="flashinfer.prefill.v15", + ) diff --git a/tests/test_sglang_symm_mem_graph.py b/tests/test_sglang_symm_mem_graph.py new file mode 100644 index 00000000..3d19313f --- /dev/null +++ b/tests/test_sglang_symm_mem_graph.py @@ -0,0 +1,606 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the Foundry project +"""CPU contracts for SGLang symmetric-memory graph relocation.""" + +from __future__ import annotations + +import importlib.util +import json +import sys +from pathlib import Path +from types import ModuleType + +import pytest + +SGLANG_DIR = Path(__file__).parents[1] / "python" / "foundry" / "integration" / "sglang" + + +def _ensure_package(name: str) -> ModuleType: + package = sys.modules.get(name) + if package is None: + package = ModuleType(name) + package.__path__ = [] # type: ignore[attr-defined] + sys.modules[name] = package + return package + + +def _load_path_module(module_name: str, filename: str): + module_path = SGLANG_DIR / filename + spec = importlib.util.spec_from_file_location(module_name, module_path) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + sys.modules[module_name] = module + spec.loader.exec_module(module) + return module + + +foundry_pkg = _ensure_package("foundry") +integration_pkg = _ensure_package("foundry.integration") +sglang_pkg = _ensure_package("foundry.integration.sglang") +foundry_pkg.integration = integration_pkg +integration_pkg.sglang = sglang_pkg + + +def _load_state_manifest(): + module = _load_path_module( + "foundry.integration.sglang.state_manifest", + "state_manifest.py", + ) + sglang_pkg.state_manifest = module + return module + + +def _load_module(): + shape_module = _load_path_module( + "foundry.integration.sglang.shape_replay_state", + "shape_replay_state.py", + ) + sglang_pkg.shape_replay_state = shape_module + _load_state_manifest() + adapter_module = _load_path_module( + "foundry.integration.sglang.sglang_shape_adapter", + "sglang_shape_adapter.py", + ) + sglang_pkg.sglang_shape_adapter = adapter_module + flashinfer_path = SGLANG_DIR / "flashinfer_graph_abi.py" + if flashinfer_path.exists(): + flashinfer_module = _load_path_module( + "foundry.integration.sglang.flashinfer_graph_abi", + "flashinfer_graph_abi.py", + ) + sglang_pkg.flashinfer_graph_abi = flashinfer_module + module = _load_path_module( + "foundry.integration.sglang.symm_mem_graph", + "symm_mem_graph.py", + ) + sglang_pkg.symm_mem_graph = module + return module + + +def _pointer_param(index: int, value: int) -> dict: + return { + "index": index, + "offset": index * 8, + "size": 8, + "value_hex": value.to_bytes(8, "little").hex(), + } + + +def _copy_params(src: int, dst: int) -> dict: + return { + "srcDevice": src, + "dstDevice": dst, + "srcMemoryType": 2, + "dstMemoryType": 2, + "WidthInBytes": 8192, + "Height": 1, + "Depth": 1, + "srcXInBytes": 0, + "srcY": 0, + "srcZ": 0, + "srcLOD": 0, + "srcPitch": 0, + "srcHeight": 0, + "dstXInBytes": 0, + "dstY": 0, + "dstZ": 0, + "dstLOD": 0, + "dstPitch": 0, + "dstHeight": 0, + } + + +def _graph(state) -> dict: + return { + "nodes": [ + { + "id": 0, + "type": "MemcpyNode", + "params": _copy_params(0x600000000400, state.buffer), + }, + { + "id": 1, + "type": "KernelNode", + "params": { + "function_name": ( + "_ZN64_GLOBAL__N__6test" + "two_shot_all_reduce_kernel_inplace" + "IN3c108BFloat16ELi16ELi2EEEvPPT_mmPPjmm" + ), + "kernelParams": [ + _pointer_param(0, state.buffer_ptrs_dev), + _pointer_param(1, 0), + _pointer_param(2, 4096), + _pointer_param(3, state.signal_pad_ptrs_dev), + _pointer_param(4, state.rank), + _pointer_param(5, state.world_size), + ], + }, + }, + { + "id": 2, + "type": "MemcpyNode", + "params": _copy_params(state.buffer, 0x600000002400), + }, + ], + "dependencies": [ + {"from": 0, "to": 1}, + {"from": 1, "to": 2}, + ], + } + + +def _state( + module, + *, + buffer: int = 0x2BC0000000, + buffer_ptrs_dev: int = 0x600006400400, + signal_pad_ptrs_dev: int = 0x600006400600, + multicast_ptr: int = 0x2C04000000, + rank: int = 0, + world_size: int = 2, + buffer_numel: int = 4096, +): + return module.SymmetricMemoryGraphState( + buffer=buffer, + buffer_ptrs_dev=buffer_ptrs_dev, + signal_pad_ptrs_dev=signal_pad_ptrs_dev, + multicast_ptr=multicast_ptr, + buffer_numel=buffer_numel, + element_size=2, + dtype="torch.bfloat16", + rank=rank, + world_size=world_size, + ) + + +def test_preserves_and_relocates_symmetric_graphs(tmp_path: Path) -> None: + module = _load_module() + saved = _state( + module, + buffer=0x2BC0000000, + buffer_ptrs_dev=0x600006400400, + signal_pad_ptrs_dev=0x600006400600, + multicast_ptr=0x2C04000000, + rank=1, + ) + live = _state( + module, + buffer=0x4A20000000, + buffer_ptrs_dev=0x600006500400, + signal_pad_ptrs_dev=0x600006500600, + multicast_ptr=0x4A64000000, + rank=1, + ) + filename = "graph_0_FULL_t1_r1_UX_pcN.json" + (tmp_path / filename).write_text(json.dumps(_graph(saved))) + + inventories = module.preserve_symmetric_graphs(tmp_path, [filename], saved) + module.write_graph_backend_metadata(tmp_path, True) + relocated_paths = module.relocate_symmetric_graphs(tmp_path, [filename], live) + + assert len(inventories) == 1 + assert inventories[0].graph_filename == filename + assert inventories[0].batch_size == 1 + relocated = json.loads(Path(relocated_paths[0]).read_text()) + assert relocated["nodes"][0]["params"]["dstDevice"] == live.buffer + assert relocated["nodes"][2]["params"]["srcDevice"] == live.buffer + params = relocated["nodes"][1]["params"]["kernelParams"] + assert int.from_bytes(bytes.fromhex(params[0]["value_hex"]), "little") == (live.buffer_ptrs_dev) + assert int.from_bytes(bytes.fromhex(params[3]["value_hex"]), "little") == ( + live.signal_pad_ptrs_dev + ) + assert relocated["nodes"][0]["params"]["srcDevice"] == 0x600000000400 + assert relocated["nodes"][2]["params"]["dstDevice"] == 0x600000002400 + + +def test_relocation_rejects_incompatible_live_group(tmp_path: Path) -> None: + module = _load_module() + saved = _state(module) + live = _state( + module, + buffer=5, + buffer_ptrs_dev=6, + signal_pad_ptrs_dev=7, + multicast_ptr=8, + rank=1, + ) + filename = "graph_0_FULL_t1_r1_UX_pcN.json" + (tmp_path / filename).write_text(json.dumps(_graph(saved))) + module.preserve_symmetric_graphs(tmp_path, [filename], saved) + + with pytest.raises(RuntimeError, match="rank"): + module.relocate_symmetric_graphs(tmp_path, [filename], live) + + +def test_rejects_save_load_backend_mismatch(tmp_path: Path) -> None: + module = _load_module() + with pytest.raises(RuntimeError, match="no valid backend metadata"): + module.validate_graph_backend(tmp_path, False) + + (tmp_path / module.BACKEND_FILENAME).write_text("{") + with pytest.raises(RuntimeError, match="no valid backend metadata"): + module.validate_graph_backend(tmp_path, False) + + module.write_graph_backend_metadata(tmp_path, False) + module.validate_graph_backend(tmp_path, False) + with pytest.raises(RuntimeError, match="archive=nccl"): + module.validate_graph_backend(tmp_path, True) + + state = _state(module) + filename = "graph_0_FULL_t1_r1_UX_pcN.json" + (tmp_path / filename).write_text(json.dumps(_graph(state))) + module.preserve_symmetric_graphs(tmp_path, [filename], state) + module.write_graph_backend_metadata(tmp_path, True) + module.validate_graph_backend(tmp_path, True) + with pytest.raises(RuntimeError, match="archive=torch_symmetric_memory"): + module.validate_graph_backend(tmp_path, False) + + +@pytest.mark.parametrize( + ("mutation", "message"), + [ + ( + lambda graph: graph["nodes"][1]["params"]["kernelParams"].reverse(), + "reordered", + ), + ( + lambda graph: graph["nodes"][1]["params"].__setitem__( + "function_name", + graph["nodes"][1]["params"]["function_name"].replace( + "ELi2EEE", + "ELi4EEE", + ), + ), + "specialization", + ), + ( + lambda graph: graph["nodes"][0]["params"].__setitem__( + "WidthInBytes", + 4096, + ), + "copy layout", + ), + ( + lambda graph: graph["nodes"][0]["params"].__setitem__( + "dstXInBytes", + 16, + ), + "copy layout", + ), + ( + lambda graph: graph["dependencies"].clear(), + "dependencies", + ), + ( + lambda graph: graph["nodes"][2].__setitem__("id", 0), + "duplicate node IDs", + ), + ( + lambda graph: graph["nodes"].append( + { + "id": 3, + "type": "KernelNode", + "params": { + "function_name": "unknown_collective", + "kernelParams": [_pointer_param(0, 0x600006400400)], + "extra_argBuffer_hex": "", + }, + } + ), + "unrecognized symmetric-memory operand", + ), + ( + lambda graph: graph["nodes"].append( + { + "id": 3, + "type": "KernelNode", + "params": { + "function_name": "unknown_collective", + "kernelParams": [_pointer_param(0, 0x600006400400 + 8)], + "extra_argBuffer_hex": "", + }, + } + ), + "unrecognized symmetric-memory operand", + ), + ( + lambda graph: graph["nodes"].append( + { + "id": 3, + "type": "KernelNode", + "params": { + "function_name": "unknown_collective", + "kernelParams": [_pointer_param(0, 0x600006400600 + 8)], + "extra_argBuffer_hex": "", + }, + } + ), + "unrecognized symmetric-memory operand", + ), + ( + lambda graph: graph["nodes"].append( + { + "id": 3, + "type": "KernelNode", + "params": { + "function_name": "unknown_collective", + "kernelParams": [_pointer_param(0, 0x2C04000000 + 8)], + "extra_argBuffer_hex": "", + }, + } + ), + "unrecognized symmetric-memory operand", + ), + ( + lambda graph: graph["nodes"].append( + { + "id": 3, + "type": "KernelNode", + "params": { + "function_name": "unknown_collective", + "kernelParams": [_pointer_param(0, 0x2BC0000000 + 16)], + "extra_argBuffer_hex": "", + }, + } + ), + "unrecognized symmetric-memory operand", + ), + ( + lambda graph: graph["nodes"].append( + { + "id": 3, + "type": "MemsetNode", + "params": {"dst": 0x2BC0000000}, + } + ), + "memset node", + ), + ( + lambda graph: graph["nodes"][1]["params"].__setitem__( + "kernel_node_attrs", + {"accessPolicyWindowBasePtr": 0x2BC0000000}, + ), + "access-policy operand", + ), + ( + lambda graph: graph["nodes"][1]["params"].__setitem__( + "function_name", + "multimem_all_reduce_kernel", + ), + "unrecognized symmetric-memory operand", + ), + ], +) +def test_relocation_rejects_changed_collective_abi( + tmp_path: Path, + mutation, + message: str, +) -> None: + module = _load_module() + state = _state(module) + filename = "graph_0_FULL_t1_r1_UX_pcN.json" + graph = _graph(state) + mutation(graph) + (tmp_path / filename).write_text(json.dumps(graph)) + + with pytest.raises(RuntimeError, match=message): + module.preserve_symmetric_graphs(tmp_path, [filename], state) + + +def test_preserve_rejects_collective_larger_than_buffer(tmp_path: Path) -> None: + module = _load_module() + state = _state(module, buffer_numel=2048) + filename = "graph_0_FULL_t1_r1_UX_pcN.json" + (tmp_path / filename).write_text(json.dumps(_graph(state))) + + with pytest.raises(RuntimeError, match="exceeds the communication buffer"): + module.preserve_symmetric_graphs(tmp_path, [filename], state) + + +def test_preserve_allows_unencoded_ordinary_kernel_metadata(tmp_path: Path) -> None: + module = _load_module() + state = _state(module) + filename = "graph_0_FULL_t1_r1_UX_pcN.json" + graph = _graph(state) + graph["nodes"].append( + { + "id": 3, + "type": "KernelNode", + "params": { + "function_name": "ordinary_kernel", + "kernelParams": [{"index": 0, "offset": 0, "size": 8}], + "extra_argBuffer_hex": "", + }, + } + ) + (tmp_path / filename).write_text(json.dumps(graph)) + + module.preserve_symmetric_graphs(tmp_path, [filename], state) + + +def test_preserve_allows_ordered_multi_shape_with_flag(tmp_path: Path) -> None: + module = _load_module() + state = _state(module) + filenames = [ + "graph_0_FULL_t8_r8_UX_pcN.json", + "graph_1_FULL_t1_r1_UX_pcN.json", + ] + for filename in filenames: + (tmp_path / filename).write_text(json.dumps(_graph(state))) + + records = module.preserve_symmetric_graphs( + tmp_path, + filenames, + state, + allow_multi_shape=True, + ) + + assert tuple(record.batch_size for record in records) == (8, 1) + + +def test_preserve_still_rejects_multi_shape_by_default(tmp_path: Path) -> None: + module = _load_module() + state = _state(module) + filenames = [ + "graph_0_FULL_t8_r8_UX_pcN.json", + "graph_1_FULL_t1_r1_UX_pcN.json", + ] + for filename in filenames: + (tmp_path / filename).write_text(json.dumps(_graph(state))) + + with pytest.raises(RuntimeError, match="exactly one"): + module.preserve_symmetric_graphs(tmp_path, filenames, state) + + +def test_relocation_rejects_manifest_order_mismatch(tmp_path: Path) -> None: + module = _load_module() + manifest_module = _load_state_manifest() + state = _state(module) + filenames = [ + "graph_0_FULL_t8_r8_UX_pcN.json", + "graph_1_FULL_t1_r1_UX_pcN.json", + ] + for filename in filenames: + (tmp_path / filename).write_text(json.dumps(_graph(state))) + + module.preserve_symmetric_graphs( + tmp_path, + filenames, + state, + allow_multi_shape=True, + ) + manifest = manifest_module.ShapeStateManifest( + backend="torch_symmetric_memory", + rank=state.rank, + world_size=state.world_size, + capture_order=(1, 8), + shapes=( + manifest_module.ShapeStateRecord( + graph_filename="graph_1_FULL_t1_r1_UX_pcN.json", + batch_size=1, + capture_index=0, + plan_schema="flashinfer.prefill.v15", + plan_fingerprint=tuple(range(15)), + tensor_slots=(), + operand_slots=(), + ), + manifest_module.ShapeStateRecord( + graph_filename="graph_0_FULL_t8_r8_UX_pcN.json", + batch_size=8, + capture_index=1, + plan_schema="flashinfer.prefill.v15", + plan_fingerprint=tuple(range(15)), + tensor_slots=(), + operand_slots=(), + ), + ), + ) + + with pytest.raises(RuntimeError, match="manifest order mismatch"): + module.relocate_symmetric_graphs( + tmp_path, + filenames, + state, + manifest=manifest, + allow_multi_shape=True, + ) + + +def test_relocation_rejects_invalid_in_memory_manifest(tmp_path: Path) -> None: + module = _load_module() + manifest_module = _load_state_manifest() + state = _state(module) + filenames = [ + "graph_0_FULL_t8_r8_UX_pcN.json", + "graph_1_FULL_t1_r1_UX_pcN.json", + ] + for filename in filenames: + (tmp_path / filename).write_text(json.dumps(_graph(state))) + + module.preserve_symmetric_graphs( + tmp_path, + filenames, + state, + allow_multi_shape=True, + ) + manifest = manifest_module.ShapeStateManifest( + backend="torch_symmetric_memory", + rank=state.rank, + world_size=3, + capture_order=(8, 1), + shapes=( + manifest_module.ShapeStateRecord( + graph_filename="graph_0_FULL_t8_r8_UX_pcN.json", + batch_size=8, + capture_index=0, + plan_schema="flashinfer.prefill.v15", + plan_fingerprint=tuple(range(15)), + tensor_slots=(), + operand_slots=(), + ), + manifest_module.ShapeStateRecord( + graph_filename="graph_1_FULL_t1_r1_UX_pcN.json", + batch_size=1, + capture_index=1, + plan_schema="flashinfer.prefill.v15", + plan_fingerprint=tuple(range(15)), + tensor_slots=(), + operand_slots=(), + ), + ), + ) + + with pytest.raises(RuntimeError, match="requires TP=2"): + module.relocate_symmetric_graphs( + tmp_path, + filenames, + state, + manifest=manifest, + allow_multi_shape=True, + ) + + +def test_graph_warmups_run_identically_twice() -> None: + module = _load_module() + calls = [] + + module.run_graph_warmups( + synchronize=lambda: calls.append("synchronize"), + barrier=lambda: calls.append("barrier"), + forward=lambda: calls.append("forward"), + post_warmup=lambda: calls.append("post"), + ) + + assert calls == [ + "synchronize", + "barrier", + "forward", + "post", + "synchronize", + "barrier", + "forward", + "post", + "synchronize", + "barrier", + ] diff --git a/tests/test_sglang_tp_recipe.py b/tests/test_sglang_tp_recipe.py new file mode 100644 index 00000000..f0c902fb --- /dev/null +++ b/tests/test_sglang_tp_recipe.py @@ -0,0 +1,513 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the Foundry project +"""CPU-only contract tests for the SGLang tensor-parallel recipe.""" + +from __future__ import annotations + +import os +import subprocess +from pathlib import Path + +import pytest +import tomllib + +from tests.modal_sglang_tp import ( + MAX_NEW_TOKENS, + NCCL_DEFAULT_GRAPH_BATCHES, + VALIDATION_REPORT_FILENAME, + build_generate_request_payload, + build_sampling_params, + canonical_shape_state_payload, + canonical_symmetric_memory_state_payload, + cleanup_run_artifacts, + eager_fallback_observed, + eager_fallback_sampling_params, + expected_graph_batches_from_env, + graph_batch_exercise_order, + graph_batch_exercise_sampling_params, + parse_eager_decode_batch_sizes, + parse_generate_list_response, + parse_replay_batch_sizes, + required_archive_files, + restored_graph_replay_observed, + semantic_shape_state_fingerprint, + semantic_symmetric_memory_state_fingerprint, + should_skip_run_cleanup, +) + +REPO_ROOT = Path(__file__).parents[1] +RECIPE_DIR = REPO_ROOT / "recipe" / "experimental" +SERVE_SCRIPT = RECIPE_DIR / "serve_qwen3-8b_sglang_tp.sh" + + +def _run_recipe(tmp_path: Path, mode: str | None) -> tuple[dict[str, str], list[str]]: + bin_dir = tmp_path / "bin" + bin_dir.mkdir() + fake_sglang = bin_dir / "sglang" + fake_sglang.write_text( + '#!/usr/bin/env bash\nenv > "$CAPTURE_ENV"\nprintf "%s\\n" "$@" > "$CAPTURE_ARGS"\n' + ) + fake_sglang.chmod(0o755) + + env_path = tmp_path / "env" + args_path = tmp_path / "args" + env = os.environ.copy() + env.update( + { + "CAPTURE_ENV": str(env_path), + "CAPTURE_ARGS": str(args_path), + "PATH": f"{bin_dir}:{env['PATH']}", + } + ) + + command = ["bash", str(SERVE_SCRIPT), "2"] + if mode is not None: + command.append(mode) + subprocess.run(command, check=True, env=env, capture_output=True, text=True) + + captured_env = dict(line.split("=", 1) for line in env_path.read_text().splitlines()) + return captured_env, args_path.read_text().splitlines() + + +@pytest.mark.parametrize("mode", [None, "--save", "--load"]) +def test_tp_recipe_uses_same_nccl_transport_for_every_phase( + tmp_path: Path, + mode: str | None, +) -> None: + env, args = _run_recipe(tmp_path, mode) + + assert env["NCCL_CUMEM_ENABLE"] == "0" + assert env["NCCL_NVLS_ENABLE"] == "0" + assert "--disable-piecewise-cuda-graph" in args + assert args[args.index("--random-seed") + 1] == "42" + + +@pytest.mark.parametrize("mode", [None, "--save", "--load"]) +def test_tp_recipe_supports_torch_symmetric_memory( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + mode: str | None, +) -> None: + monkeypatch.setenv("SGLANG_ENABLE_TORCH_SYMM_MEM", "1") + + _env, args = _run_recipe(tmp_path, mode) + + assert "--enable-torch-symm-mem" in args + assert args[args.index("--cuda-graph-max-bs") + 1] == "1" + + +def test_tp_recipe_rejects_multi_shape_symmetric_memory( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + monkeypatch.setenv("SGLANG_ENABLE_TORCH_SYMM_MEM", "1") + monkeypatch.setenv("SGLANG_CUDA_GRAPH_MAX_BS", "8") + + with pytest.raises(subprocess.CalledProcessError): + _run_recipe(tmp_path, "--save") + + +def test_tp_recipe_allows_opt_in_multi_shape( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + monkeypatch.setenv("SGLANG_ENABLE_TORCH_SYMM_MEM", "1") + monkeypatch.setenv("FOUNDRY_SGLANG_SYMM_MULTI_SHAPE", "1") + monkeypatch.setenv("FOUNDRY_SGLANG_SYMM_GRAPH_BATCHES", "1,8") + + _env, args = _run_recipe(tmp_path, "--save") + + assert args[args.index("--cuda-graph-max-bs") + 1] == "8" + + +def test_tp_recipe_requires_graph_batches_for_multi_shape( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + monkeypatch.setenv("SGLANG_ENABLE_TORCH_SYMM_MEM", "1") + monkeypatch.setenv("FOUNDRY_SGLANG_SYMM_MULTI_SHAPE", "1") + monkeypatch.delenv("FOUNDRY_SGLANG_SYMM_GRAPH_BATCHES", raising=False) + + with pytest.raises(subprocess.CalledProcessError) as exc_info: + _run_recipe(tmp_path, "--save") + + assert "FOUNDRY_SGLANG_SYMM_GRAPH_BATCHES" in exc_info.value.stderr + + +def test_tp_recipe_keeps_batch_one_default( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + monkeypatch.setenv("SGLANG_ENABLE_TORCH_SYMM_MEM", "1") + + _env, args = _run_recipe(tmp_path, "--save") + + assert args[args.index("--cuda-graph-max-bs") + 1] == "1" + + +@pytest.mark.parametrize( + ("mode", "config_name"), + [ + ("--save", "sglang_foundry_tp_save.toml"), + ("--load", "sglang_foundry_tp_load.toml"), + ], +) +def test_tp_recipe_selects_matching_foundry_config( + tmp_path: Path, + mode: str, + config_name: str, +) -> None: + _env, args = _run_recipe(tmp_path, mode) + + config_index = args.index("--foundry-graph-extension-config-path") + 1 + assert Path(args[config_index]) == RECIPE_DIR / config_name + assert args[args.index("--tp-size") + 1] == "2" + assert "--disable-custom-all-reduce" in args + + +def test_tp_save_and_load_configs_are_symmetric() -> None: + save = tomllib.loads((RECIPE_DIR / "sglang_foundry_tp_save.toml").read_text()) + load = tomllib.loads((RECIPE_DIR / "sglang_foundry_tp_load.toml").read_text()) + + assert save.pop("mode") == "save" + assert load.pop("mode") == "load" + assert save == load + + +def test_expected_graph_batches_defaults_to_one_for_symm_mem() -> None: + assert expected_graph_batches_from_env( + graph_batches=None, + uses_torch_symm_mem=True, + ) == (1,) + + +def test_expected_graph_batches_defaults_to_nccl_inventory() -> None: + assert ( + expected_graph_batches_from_env( + graph_batches=None, + uses_torch_symm_mem=False, + ) + == NCCL_DEFAULT_GRAPH_BATCHES + ) + assert len(NCCL_DEFAULT_GRAPH_BATCHES) == 36 + + +@pytest.mark.parametrize( + "raw", + ["1,8", "1,8,32", "4,16,64"], +) +def test_expected_graph_batches_parses_explicit_list(raw: str) -> None: + assert expected_graph_batches_from_env( + graph_batches=raw, + uses_torch_symm_mem=True, + ) == tuple(int(value) for value in raw.split(",")) + + +def test_restored_graph_replay_observed_requires_every_expected_batch() -> None: + expected = (1, 8, 32) + assert restored_graph_replay_observed(expected, [1, 8, 32, 64]) + assert not restored_graph_replay_observed(expected, [1, 8]) + assert not restored_graph_replay_observed((1,), [8]) + + +def test_required_archive_files_for_symm_mem_and_multi_shape() -> None: + base = required_archive_files( + uses_torch_symm_mem=False, + uses_multi_shape=False, + ) + assert base == { + "graph_manifest.json", + "fatbin_image_packed.img", + "final_alloc_offset.json", + } + + symm = required_archive_files( + uses_torch_symm_mem=True, + uses_multi_shape=False, + ) + assert symm == base | { + "sglang_graph_backend.json", + "symmetric_memory_state.json", + } + + multi = required_archive_files( + uses_torch_symm_mem=True, + uses_multi_shape=True, + ) + assert multi == symm | {"sglang_shape_state.json"} + + +def _sample_symmetric_memory_state( + *, + buffer: int = 0x600006400000, + buffer_ptrs_dev: int = 0x600006400400, + signal_pad_ptrs_dev: int = 0x600006400600, + multicast_ptr: int = 0x2C04000000, + buffer_numel: int = 1_048_576, +) -> dict: + return { + "version": 1, + "backend": "torch_symmetric_memory", + "state": { + "buffer": buffer, + "buffer_ptrs_dev": buffer_ptrs_dev, + "signal_pad_ptrs_dev": signal_pad_ptrs_dev, + "multicast_ptr": multicast_ptr, + "buffer_numel": buffer_numel, + "element_size": 2, + "dtype": "bfloat16", + "rank": 0, + "world_size": 2, + }, + } + + +def _sample_shape_state( + *, + schema_fingerprint: str = "abc123", + saved_value: int = 0x600006400400, +) -> dict: + return { + "version": 1, + "schema_fingerprint": schema_fingerprint, + "manifest": { + "backend": "torch_symmetric_memory", + "rank": 0, + "world_size": 2, + "capture_order": [8, 1], + "shapes": [ + { + "graph_filename": "graph_0_FULL_t8_r8_UX_pcN.json", + "batch_size": 8, + "capture_index": 0, + "plan_schema": "flashinfer-v1", + "plan_fingerprint": [1, 2, 3], + "tensor_slots": [], + "operand_slots": [ + { + "name": "symm.buffer_ptrs_dev", + "kernel_abi": "two_shot", + "owner_slot": "symm", + "node_id": 1, + "source": "param:0", + "parameter_index": 0, + "cuda_parameter_offset": 0, + "value_byte_offset": 0, + "ctype": "void*", + "owner_relative_offset": 0, + "span_bytes": 8, + "saved_value": saved_value, + } + ], + } + ], + }, + } + + +def test_symmetric_memory_state_fingerprint_ignores_process_local_pointers() -> None: + first = _sample_symmetric_memory_state() + second = _sample_symmetric_memory_state( + buffer=0x700000000000, + buffer_ptrs_dev=0x700000000400, + signal_pad_ptrs_dev=0x700000000600, + multicast_ptr=0x3D05000000, + ) + + assert semantic_symmetric_memory_state_fingerprint( + first + ) == semantic_symmetric_memory_state_fingerprint(second) + + +def test_symmetric_memory_state_fingerprint_detects_capacity_change() -> None: + baseline = _sample_symmetric_memory_state() + changed = _sample_symmetric_memory_state(buffer_numel=2_097_152) + + assert semantic_symmetric_memory_state_fingerprint( + baseline + ) != semantic_symmetric_memory_state_fingerprint(changed) + + +def test_shape_state_fingerprint_ignores_saved_value_bytes() -> None: + first = _sample_shape_state(saved_value=0x600006400400) + second = _sample_shape_state(saved_value=0x700000000400) + + assert semantic_shape_state_fingerprint(first) == semantic_shape_state_fingerprint(second) + + +def test_shape_state_fingerprint_detects_schema_change() -> None: + baseline = _sample_shape_state(schema_fingerprint="abc123") + changed = _sample_shape_state(schema_fingerprint="def456") + + assert semantic_shape_state_fingerprint(baseline) != semantic_shape_state_fingerprint(changed) + + +def test_canonical_payloads_strip_pointer_and_saved_value_fields() -> None: + symm = canonical_symmetric_memory_state_payload(_sample_symmetric_memory_state()) + assert "buffer" not in symm["state"] + assert symm["state"]["buffer_numel"] == 1_048_576 + + shape = canonical_shape_state_payload(_sample_shape_state()) + assert shape["schema_fingerprint"] == "abc123" + assert "saved_value" not in shape["shapes"][0]["operand_slots"][0] + + +def test_parse_replay_batch_sizes_uses_graph_key_not_raw_running_req() -> None: + log_text = ( + "Decode batch, #running-req: 7, cuda graph: True\n" + "Decode graph replay: worker=target key_size=8 (bs) mode=DECODE raw_bs=7\n" + ) + + assert parse_replay_batch_sizes(log_text) == {8} + + +def test_parse_replay_batch_sizes_collects_multiple_bs_key_sizes() -> None: + log_text = ( + "Decode graph replay: worker=target key_size=1 (bs) mode=DECODE raw_bs=1\n" + "Decode graph replay: worker=target key_size=8 (bs) mode=DECODE raw_bs=7\n" + "Decode graph replay: worker=target key_size=32 (bs) mode=DECODE raw_bs=29\n" + ) + + assert parse_replay_batch_sizes(log_text) == {1, 8, 32} + + +def test_parse_replay_batch_sizes_ignores_num_tokens_graph_keys() -> None: + log_text = ( + "Decode graph replay: worker=target key_size=64 (num_tokens) " + "mode=TARGET_VERIFY raw_bs=7 slots=[0, 1, 2]\n" + "Decode graph replay: worker=target key_size=8 (bs) mode=DECODE raw_bs=7\n" + ) + + assert parse_replay_batch_sizes(log_text) == {8} + + +def test_parse_replay_batch_sizes_ignores_running_req_lines_without_graph_key() -> None: + log_text = ( + "Decode batch, #running-req: 1, cuda graph: True\n" + "Decode batch, #running-req: 8, cuda graph: True\n" + "Decode batch, #running-req: 4, cuda graph: False\n" + ) + + assert parse_replay_batch_sizes(log_text) == set() + + +def test_parse_eager_decode_batch_sizes_ignores_graph_true_lines() -> None: + log_text = ( + "Decode batch, #running-req: 9, cuda graph: True\n" + "Decode graph replay: worker=target key_size=8 (bs) mode=DECODE raw_bs=7\n" + "Decode batch, #running-req: 9, cuda graph: False\n" + ) + + assert parse_eager_decode_batch_sizes(log_text) == {9} + + +def test_eager_fallback_observed_requires_batch_at_least_fallback_size() -> None: + fallback_batch = 9 + assert not eager_fallback_observed({8}, fallback_batch) + assert eager_fallback_observed({9}, fallback_batch) + assert eager_fallback_observed({10, 8}, fallback_batch) + + +def test_build_sampling_params_omits_extra_keys_for_normal_generation() -> None: + assert build_sampling_params(max_new_tokens=MAX_NEW_TOKENS) == { + "temperature": 0.0, + "max_new_tokens": MAX_NEW_TOKENS, + } + + +def test_eager_fallback_sampling_params_force_decode_tokens() -> None: + assert eager_fallback_sampling_params() == { + "temperature": 0.0, + "max_new_tokens": 4, + "min_new_tokens": 4, + "ignore_eos": True, + } + + +def test_graph_batch_exercise_sampling_params_force_decode_tokens() -> None: + assert graph_batch_exercise_sampling_params() == { + "temperature": 0.0, + "max_new_tokens": 2, + "min_new_tokens": 2, + "ignore_eos": True, + } + + +def test_graph_batch_exercise_order_descends_for_small_set() -> None: + assert graph_batch_exercise_order((1, 8, 32)) == (32, 8, 1) + + +def test_graph_batch_exercise_order_rejects_duplicate_keys() -> None: + with pytest.raises(ValueError, match="exactly once"): + graph_batch_exercise_order((8, 8, 1)) + + +def test_graph_batch_exercise_order_preserves_full_inventory_keys() -> None: + ordered = graph_batch_exercise_order(NCCL_DEFAULT_GRAPH_BATCHES) + + assert ordered == tuple(sorted(NCCL_DEFAULT_GRAPH_BATCHES, reverse=True)) + assert set(ordered) == set(NCCL_DEFAULT_GRAPH_BATCHES) + assert len(ordered) == len(NCCL_DEFAULT_GRAPH_BATCHES) + + +def test_build_generate_request_payload_uses_list_text_input() -> None: + prompts = ["one", "two", "three"] + payload = build_generate_request_payload( + prompts, + max_new_tokens=2, + min_new_tokens=2, + ignore_eos=True, + ) + + assert payload == { + "text": prompts, + "sampling_params": { + "temperature": 0.0, + "max_new_tokens": 2, + "min_new_tokens": 2, + "ignore_eos": True, + }, + } + + +def test_parse_generate_list_response_preserves_cardinality_and_order() -> None: + result = [ + {"text": "first"}, + {"text": "second"}, + ] + + assert parse_generate_list_response(result, expected_count=2) == ["first", "second"] + + +def test_parse_generate_list_response_accepts_dict_with_text_list() -> None: + result = {"text": ["alpha", "beta", "gamma"]} + + assert parse_generate_list_response(result, expected_count=3) == ["alpha", "beta", "gamma"] + + +def test_parse_generate_list_response_rejects_wrong_cardinality() -> None: + with pytest.raises(RuntimeError, match="expected 3"): + parse_generate_list_response({"text": ["only-one"]}, expected_count=3) + + +def test_parse_generate_list_response_rejects_empty_outputs() -> None: + with pytest.raises(RuntimeError, match="empty output"): + parse_generate_list_response({"text": ["ok", " "]}, expected_count=2) + + +def test_should_skip_run_cleanup_when_artifacts_or_run_id_override() -> None: + assert should_skip_run_cleanup(keep_artifacts=True, run_id_override=None) + assert should_skip_run_cleanup(keep_artifacts=False, run_id_override="manual-run") + assert not should_skip_run_cleanup(keep_artifacts=False, run_id_override=None) + + +def test_cleanup_run_artifacts_preserves_validation_report(tmp_path: Path) -> None: + run_root = tmp_path / "run-123" + run_root.mkdir() + (run_root / VALIDATION_REPORT_FILENAME).write_text('{"checks": {}}') + (run_root / "load.log").write_text("log data") + + cleanup_run_artifacts(run_root) + + assert not (run_root / "load.log").exists() + assert (run_root / VALIDATION_REPORT_FILENAME).read_text() == '{"checks": {}}' diff --git a/tests/test_symm_mem_graph.py b/tests/test_symm_mem_graph.py new file mode 100644 index 00000000..63055f17 --- /dev/null +++ b/tests/test_symm_mem_graph.py @@ -0,0 +1,208 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the Foundry project +"""Fresh-process Foundry replay coverage for PyTorch symmetric-memory graphs.""" + +from __future__ import annotations + +import importlib.util +import os +import shutil +import subprocess +import sys +from pathlib import Path + +import foundry as fdry +import torch +import torch.distributed as dist +import torch.distributed._symmetric_memory as symm_mem +from foundry.integration.sglang.symm_mem_graph import ( + SymmetricMemoryGraphState, + preserve_symmetric_graphs, + relocate_symmetric_graphs, +) + +BASE_ADDR = 0x600000000000 +REGION_SIZE = "2GB" +SCRATCH_OFFSET = 256 * 1024 * 1024 +ARCHIVE_DIR = Path("symm_mem_graph_archive") +NUMEL = 4096 +NUM_COLLECTIVES = 73 + + +def _hook_path() -> str: + spec = importlib.util.find_spec("foundry.ops") + if not spec or not spec.origin: + raise RuntimeError("foundry.ops is not installed") + path = Path(spec.origin).resolve().parent / "libcuda_hook.so" + if not path.exists(): + raise RuntimeError(f"CUDA hook not found at {path}") + return str(path) + + +def _init_dist(local_rank: int, world_size: int) -> dist.ProcessGroup: + torch.cuda.set_device(local_rank) + dist.init_process_group( + backend="nccl", + init_method=f"tcp://127.0.0.1:{os.environ['MASTER_PORT']}", + rank=local_rank, + world_size=world_size, + device_id=torch.device(f"cuda:{local_rank}"), + ) + return dist.group.WORLD + + +def _create_symm_state(group: dist.ProcessGroup, rank: int): + # Seed the caching allocator with a Foundry-backed segment. PyTorch's small + # device pointer tables then reuse this deterministic segment while the + # actual P2P allocation is created with Foundry tracking paused. + seed = torch.empty(1024, dtype=torch.uint8, device="cuda") + del seed + fdry.stop_allocation_region() + try: + buffer = symm_mem.empty(NUMEL, dtype=torch.bfloat16, device="cuda") + handle = symm_mem.rendezvous(buffer, group.group_name) + finally: + fdry.resume_allocation_region() + + print( + f"[rank {rank}] buffer={hex(buffer.data_ptr())} " + f"buffer_ptrs_dev={hex(handle.buffer_ptrs_dev)} " + f"signal_pad_ptrs_dev={hex(handle.signal_pad_ptrs_dev)} " + f"buffer_ptrs={[hex(ptr) for ptr in handle.buffer_ptrs]} " + f"signal_pad_ptrs={[hex(ptr) for ptr in handle.signal_pad_ptrs]}", + flush=True, + ) + return buffer, handle + + +def _eager_all_reduce(buffer: torch.Tensor, group: dist.ProcessGroup, rank: int) -> None: + buffer.fill_(float(rank + 1)) + torch.ops.symm_mem.two_shot_all_reduce_(buffer, "sum", group.group_name) + torch.cuda.synchronize() + torch.testing.assert_close( + buffer, + torch.full_like(buffer, 3.0), + rtol=0, + atol=0, + ) + + +def _worker(local_rank: int, world_size: int, mode: str) -> None: + group = _init_dist(local_rank, world_size) + rank_archive = ARCHIVE_DIR / f"rank_{local_rank}" + + if mode == "load": + fdry.load_cuda_modules_and_libraries(str(rank_archive)) + + fdry.set_allocation_region(BASE_ADDR, fdry.parse_size(REGION_SIZE)) + buffer, _handle = _create_symm_state(group, local_rank) + fdry.set_current_alloc_offset(SCRATCH_OFFSET) + + input_tensor = torch.full( + (NUMEL,), + float(local_rank + 1), + dtype=torch.bfloat16, + device="cuda", + ) + output_tensor = torch.empty_like(input_tensor) + _eager_all_reduce(buffer, group, local_rank) + dist.barrier() + + graph_path = rank_archive / "graph_0_FULL_t1_r1_UX_pcN.json" + state = SymmetricMemoryGraphState( + buffer=buffer.data_ptr(), + buffer_ptrs_dev=_handle.buffer_ptrs_dev, + signal_pad_ptrs_dev=_handle.signal_pad_ptrs_dev, + multicast_ptr=_handle.multicast_ptr, + buffer_numel=buffer.numel(), + element_size=buffer.element_size(), + dtype=str(buffer.dtype), + rank=_handle.rank, + world_size=_handle.world_size, + ) + if mode == "save": + graph = fdry.CUDAGraph() + with fdry.graph(graph, pool=(0, 0)): + for _ in range(NUM_COLLECTIVES): + buffer.copy_(input_tensor) + torch.ops.symm_mem.two_shot_all_reduce_( + buffer, + "sum", + group.group_name, + ) + output_tensor.copy_(buffer) + + dist.barrier() + graph.replay() + torch.cuda.synchronize() + torch.testing.assert_close( + output_tensor, + torch.full_like(output_tensor, 3.0), + rtol=0, + atol=0, + ) + rank_archive.mkdir(parents=True, exist_ok=True) + graph.save(str(graph_path), output_tensors=output_tensor) + preserve_symmetric_graphs(rank_archive, [graph_path.name], state) + fdry.pack_fatbins_to_folder(str(rank_archive)) + fdry.set_pack_fatbins_on_exit(False) + else: + relocated_paths = relocate_symmetric_graphs( + rank_archive, + [graph_path.name], + state, + ) + graph, loaded_output = fdry.CUDAGraph.load(relocated_paths[0], pool=(0, 0)) + dist.barrier() + graph.replay() + torch.cuda.synchronize() + torch.testing.assert_close( + loaded_output, + torch.full_like(loaded_output, 3.0), + rtol=0, + atol=0, + ) + + dist.barrier() + fdry.stop_allocation_region() + dist.destroy_process_group() + + +def _run_phase(mode: str) -> None: + torch.multiprocessing.spawn( + _worker, + args=(2, mode), + nprocs=2, + join=True, + ) + + +def _spawn_phase(mode: str, port: int) -> None: + env = os.environ.copy() + hook = _hook_path() + env["LD_PRELOAD"] = f"{hook}:{env['LD_PRELOAD']}" if env.get("LD_PRELOAD") else hook + env["MASTER_PORT"] = str(port) + subprocess.check_call( + [sys.executable, str(Path(__file__).resolve()), f"--{mode}"], + env=env, + ) + + +def test_symmetric_memory_graph_replays_in_fresh_process() -> None: + if torch.cuda.device_count() < 2: + return + shutil.rmtree(ARCHIVE_DIR, ignore_errors=True) + try: + _spawn_phase("save", 29731) + _spawn_phase("load", 29732) + finally: + shutil.rmtree(ARCHIVE_DIR, ignore_errors=True) + + +if __name__ == "__main__": + if "--save" in sys.argv: + _run_phase("save") + elif "--load" in sys.argv: + _run_phase("load") + else: + test_symmetric_memory_graph_replays_in_fresh_process() diff --git a/tests/test_vllm_tp_recipe.py b/tests/test_vllm_tp_recipe.py new file mode 100644 index 00000000..78d72461 --- /dev/null +++ b/tests/test_vllm_tp_recipe.py @@ -0,0 +1,121 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the Foundry project +"""CPU-only contract tests for the vLLM tensor-parallel recipe.""" + +from __future__ import annotations + +import json +import os +import subprocess +from pathlib import Path + +import pytest +import tomllib + +REPO_ROOT = Path(__file__).parents[1] +RECIPE_DIR = REPO_ROOT / "recipe" / "vllm" +SERVE_SCRIPT = RECIPE_DIR / "serve_qwen3-8b_tp.sh" + + +def _run_recipe(tmp_path: Path, mode: str | None) -> tuple[dict[str, str], list[str]]: + bin_dir = tmp_path / "bin" + bin_dir.mkdir() + fake_vllm = bin_dir / "vllm" + fake_vllm.write_text( + '#!/usr/bin/env bash\nenv > "$CAPTURE_ENV"\nprintf "%s\\n" "$@" > "$CAPTURE_ARGS"\n' + ) + fake_vllm.chmod(0o755) + + env_path = tmp_path / "env" + args_path = tmp_path / "args" + env = os.environ.copy() + env.update( + { + "CAPTURE_ENV": str(env_path), + "CAPTURE_ARGS": str(args_path), + "PATH": f"{bin_dir}:{env['PATH']}", + } + ) + + command = ["bash", str(SERVE_SCRIPT), "2"] + if mode is not None: + command.append(mode) + subprocess.run(command, check=True, env=env, capture_output=True, text=True) + + captured_env = dict(line.split("=", 1) for line in env_path.read_text().splitlines()) + return captured_env, args_path.read_text().splitlines() + + +def _compilation_config(args: list[str]) -> dict: + assert args.count("--compilation-config") == 1 + config_index = args.index("--compilation-config") + 1 + return json.loads(args[config_index]) + + +@pytest.mark.parametrize("mode", [None, "--save", "--load"]) +def test_tp_recipe_uses_same_collective_transport_for_every_phase( + tmp_path: Path, + mode: str | None, +) -> None: + env, args = _run_recipe(tmp_path, mode) + + assert env["NCCL_CUMEM_ENABLE"] == "0" + assert env["NCCL_NVLS_ENABLE"] == "0" + assert env["VLLM_ALLREDUCE_USE_SYMM_MEM"] == "0" + assert env["VLLM_ALLREDUCE_USE_FLASHINFER"] == "0" + assert env["VLLM_USE_NCCL_SYMM_MEM"] == "0" + assert env["VLLM_USE_V2_MODEL_RUNNER"] == "0" + assert env["VLLM_WORKER_MULTIPROC_METHOD"] == "spawn" + assert "--disable-custom-all-reduce" in args + config = _compilation_config(args) + assert config["cudagraph_mode"] == "FULL_DECODE_ONLY" + assert config["cudagraph_num_of_warmups"] == 0 + assert config["pass_config"]["fuse_allreduce_rms"] is False + + +@pytest.mark.parametrize( + ("mode", "config_name"), + [ + ("--save", "foundry_save.toml"), + ("--load", "foundry_load.toml"), + ], +) +def test_tp_recipe_selects_matching_foundry_config( + tmp_path: Path, + mode: str, + config_name: str, +) -> None: + _env, args = _run_recipe(tmp_path, mode) + + config = _compilation_config(args) + assert Path(config["graph_extension_config_path"]) == RECIPE_DIR / config_name + assert args[0] == "serve" + assert args[args.index("--tensor-parallel-size") + 1] == "2" + assert args[args.index("--distributed-executor-backend") + 1] == "mp" + + +def test_tp_baseline_does_not_enable_foundry(tmp_path: Path) -> None: + _env, args = _run_recipe(tmp_path, None) + + assert "graph_extension_config_path" not in _compilation_config(args) + + +def test_tp_save_and_load_configs_are_symmetric() -> None: + save = tomllib.loads((RECIPE_DIR / "foundry_save.toml").read_text()) + load = tomllib.loads((RECIPE_DIR / "foundry_load.toml").read_text()) + + assert save.pop("mode") == "save" + assert load.pop("mode") == "load" + assert save == load + assert save["scratch_space_size"] == "4096MB" + + +@pytest.mark.parametrize("arguments", [["1"], ["2", "--invalid"]]) +def test_tp_recipe_rejects_unsupported_invocations(arguments: list[str]) -> None: + result = subprocess.run( + ["bash", str(SERVE_SCRIPT), *arguments], + capture_output=True, + text=True, + ) + + assert result.returncode != 0 diff --git a/tests/test_vmm_alloc.py b/tests/test_vmm_alloc.py index 5df909f0..be5a2cb3 100644 --- a/tests/test_vmm_alloc.py +++ b/tests/test_vmm_alloc.py @@ -1,14 +1,21 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the Foundry project +import ctypes import os +import signal import subprocess import sys from pathlib import Path +import foundry as fdry import pytest import torch +class CUipcMemHandle(ctypes.Structure): + _fields_ = [("reserved", ctypes.c_byte * 64)] + + def _get_hook_so_path(): import importlib.util @@ -26,8 +33,6 @@ def _get_hook_so_path(): def _run_core(): - import foundry as fdry - torch.cuda.init() device = torch.device("cuda:0") @@ -104,8 +109,6 @@ def _run_core_without_region(): def _run_core_size_parsing(): - import foundry as fdry - torch.cuda.init() device = torch.device("cuda:0") @@ -146,6 +149,102 @@ def _run_core_size_parsing(): print("[TEST] test_size_parsing PASSED") +def _run_core_zero_alignment_reserve(): + torch.cuda.init() + + base_addr = 0x7F4000000000 + region_size = 1024 * 1024 * 1024 + reserve_size = 2 * 1024 * 1024 + + driver = ctypes.CDLL(None) + reserve = driver.cuMemAddressReserve + reserve.argtypes = [ + ctypes.POINTER(ctypes.c_uint64), + ctypes.c_size_t, + ctypes.c_size_t, + ctypes.c_uint64, + ctypes.c_ulonglong, + ] + reserve.restype = ctypes.c_int + address_free = driver.cuMemAddressFree + address_free.argtypes = [ctypes.c_uint64, ctypes.c_size_t] + address_free.restype = ctypes.c_int + + ptr = ctypes.c_uint64() + with fdry.allocation_region(base_addr, region_size): + result = reserve(ctypes.byref(ptr), reserve_size, 0, 0, 0) + assert result == 0 + assert base_addr <= ptr.value < base_addr + region_size, ( + f"zero-alignment reserve returned {hex(ptr.value)}, outside " + f"[{hex(base_addr)}, {hex(base_addr + region_size)})" + ) + assert fdry.get_current_alloc_offset() == reserve_size + assert address_free(ptr.value, reserve_size) == 0 + + +def _run_core_pitched_ipc_export(): + torch.cuda.init() + torch.cuda.set_device(0) + + driver = ctypes.CDLL(None) + alloc_pitch = driver.cuMemAllocPitch_v2 + alloc_pitch.argtypes = [ + ctypes.POINTER(ctypes.c_uint64), + ctypes.POINTER(ctypes.c_size_t), + ctypes.c_size_t, + ctypes.c_size_t, + ctypes.c_uint, + ] + alloc_pitch.restype = ctypes.c_int + ipc_get_handle = driver.cuIpcGetMemHandle + ipc_get_handle.argtypes = [ctypes.POINTER(CUipcMemHandle), ctypes.c_uint64] + ipc_get_handle.restype = ctypes.c_int + mem_free = driver.cuMemFree_v2 + mem_free.argtypes = [ctypes.c_uint64] + mem_free.restype = ctypes.c_int + + ptr = ctypes.c_uint64() + pitch = ctypes.c_size_t() + handle = CUipcMemHandle() + with fdry.allocation_region(0x600000000000, "1GB"): + assert alloc_pitch(ctypes.byref(ptr), ctypes.byref(pitch), 1024, 1024, 4) == 0 + assert ipc_get_handle(ctypes.byref(handle), ptr.value) == 0 + assert mem_free(ptr.value) == 0 + + +def _run_core_rejects_vmm_ipc_after_fork(): + torch.cuda.init() + torch.cuda.set_device(0) + + driver = ctypes.CDLL(None) + mem_alloc = driver.cuMemAlloc_v2 + mem_alloc.argtypes = [ctypes.POINTER(ctypes.c_uint64), ctypes.c_size_t] + mem_alloc.restype = ctypes.c_int + ipc_get_handle = driver.cuIpcGetMemHandle + ipc_get_handle.argtypes = [ctypes.POINTER(CUipcMemHandle), ctypes.c_uint64] + ipc_get_handle.restype = ctypes.c_int + mem_free = driver.cuMemFree_v2 + mem_free.argtypes = [ctypes.c_uint64] + mem_free.restype = ctypes.c_int + + ptr = ctypes.c_uint64() + parent_handle = CUipcMemHandle() + with fdry.allocation_region(0x600000000000, "1GB"): + assert mem_alloc(ctypes.byref(ptr), 2 * 1024 * 1024) == 0 + assert ipc_get_handle(ctypes.byref(parent_handle), ptr.value) == 0 + + child_pid = os.fork() + if child_pid == 0: + signal.alarm(5) + child_handle = CUipcMemHandle() + result = ipc_get_handle(ctypes.byref(child_handle), ptr.value) + os._exit(0 if result == 801 else 1) # CUDA_ERROR_NOT_SUPPORTED + + _, status = os.waitpid(child_pid, 0) + assert os.waitstatus_to_exitcode(status) == 0 + assert mem_free(ptr.value) == 0 + + def _spawn_with_preload(test_mode): so_path = _get_hook_so_path() env = os.environ.copy() @@ -176,6 +275,24 @@ def test_size_parsing(): _spawn_with_preload("size-parsing") +@pytest.mark.filterwarnings("ignore:TORCH_CUDA_ARCH_LIST is not set") +def test_zero_alignment_reserve_stays_in_region(): + """CUDA's alignment=0 convention must not reset the VMM cursor to zero.""" + _spawn_with_preload("zero-alignment-reserve") + + +@pytest.mark.filterwarnings("ignore:TORCH_CUDA_ARCH_LIST is not set") +def test_pitched_vmm_allocation_can_be_exported(): + """Tracked pitched allocations must request a shareable POSIX handle.""" + _spawn_with_preload("pitched-ipc-export") + + +@pytest.mark.filterwarnings("ignore:TORCH_CUDA_ARCH_LIST is not set") +def test_vmm_ipc_rejects_use_after_fork(): + """A fork child must fail promptly instead of touching inherited IPC locks.""" + _spawn_with_preload("reject-vmm-ipc-after-fork") + + if __name__ == "__main__": if "--core" in sys.argv: _run_core() @@ -183,7 +300,16 @@ def test_size_parsing(): _run_core_without_region() elif "--size-parsing" in sys.argv: _run_core_size_parsing() + elif "--zero-alignment-reserve" in sys.argv: + _run_core_zero_alignment_reserve() + elif "--pitched-ipc-export" in sys.argv: + _run_core_pitched_ipc_export() + elif "--reject-vmm-ipc-after-fork" in sys.argv: + _run_core_rejects_vmm_ipc_after_fork() else: test_vmm_allocation() test_vmm_allocation_without_region() test_size_parsing() + test_zero_alignment_reserve_stays_in_region() + test_pitched_vmm_allocation_can_be_exported() + test_vmm_ipc_rejects_use_after_fork()