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..f2d44d3c 100644 --- a/README.md +++ b/README.md @@ -82,11 +82,30 @@ Foundry ships engine integrations under `foundry/python/foundry/integration/`. P | Engine | Single GPU | DP | TP | EP | |---|:---:|:---:|:---:|:---:| | vLLM | โœ… | โœ… | ๐Ÿšง | โœ… | -| SGLang | โœ… | โœ… | ๐Ÿšง | โœ… | +| SGLang | โœ… | โœ… | โœ… | โœ… | | TensorRT-LLM | ๐Ÿšง | ๐Ÿšง | ๐Ÿšง | ๐Ÿšง | โœ… validated end-to-end (SAVE โ†’ LOAD โ†’ query)  ยท  ๐Ÿšง not yet +SGLang-main validation is pinned to +`9b853e6832e71a3058212df02a025232a453e146`; the current evidence uses immutable +Foundry `eb1ec00b7004339c5150ce6c3ba057ceda047833`. Plain TP uses one explicit +padded symmetric-memory decode shape; the Full prefill backend may additionally +persist multiple token buckets. Qwen3.5-122B-A10B-GPTQ-Int4 TP=4 restored +decode-8 and prefill-16/64 graphs on every H100 rank in Modal run +`eb1ec00b7004339c5150ce6c-1784907539317556437`. Combined online-FP8 TP=2/EP=2, +NEXTN, DeepEP low-latency, and Full prefill-16/64 was validated with +Qwen3.5-35B-A3B in run +`eb1ec00b7004339c5150ce6c-1784906247386993159`. + +NEXTN is the GPU-validated speculative algorithm. DeepEP on the pinned Hopper +stack requires FA3 and the documented NVSHMEM transport settings; a separate +legacy NVL CUDA-IPC path is available for no-fabric hosts. Multi-shape decode, +other speculative algorithms, other fabrics, breakable/`tc` prefill +persistence, PP, LoRA, PDMux, diffusion, and recapture remain outside the +validated scope. Broader TP collective support is tracked in +[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`. ### Performance diff --git a/RELEASE.md b/RELEASE.md index 6586a781..12cdedf9 100644 --- a/RELEASE.md +++ b/RELEASE.md @@ -1,10 +1,48 @@ -# Foundry 0.0.2 +# Release Notes + +## Unreleased + +- **SGLang-main tensor parallel and Full prefill.** One explicit padded + torch-symmetric-memory decode shape is validated end-to-end with Qwen3-8B + TP=2 and Qwen3.5-122B-A10B-GPTQ-Int4 TP=4 on H100. The TP4 regression also + persisted Full prefill buckets 16/64: fresh LOAD restored three graphs per + rank at offset `59301167104`, with exact outputs, collectives, catalogs, and + SAVE/SAVE2 fingerprints in run + `eb1ec00b7004339c5150ce6c-1784907539317556437`. Multi-shape decode and + alternate-collective support remain tracked in + [upstream issue #6](https://github.com/foundry-org/foundry/issues/6). +- **SGLang graph-mode composition.** With SGLang + `9b853e6832e71a3058212df02a025232a453e146`, Qwen3.5-35B-A3B online FP8 + combined TP=2/EP=2, NEXTN (3 steps, top-k 1, 4 draft tokens), DeepEP + low-latency, and Full prefill 16/64 in one fresh LOAD. Run + `eb1ec00b7004339c5150ce6c-1784906247386993159` at immutable Foundry + `eb1ec00b7004339c5150ce6c3ba057ceda047833` restored four graphs and three + sessions per rank, reproduced rank-local offsets and archives, and retained + exact token, speculative-acceptance, DeepEP-node, bootstrap, NVSHMEM-module, + and no-recapture evidence. NEXTN is the only GPU-validated speculative + algorithm; Full is the only validated prefill-persistence backend. +- **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. +- **A100/FA2 load parity.** 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). + +--- + +## 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,12 +61,13 @@ 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 +- **SGLang** โ€” integration for SGLang v0.5.13. Working configurations in this + historical release: single GPU, data parallel (DP), expert parallel (EP, DeepEP low-latency + DP-attention with fa3). Self-contained recipes under `recipe/sglang/` (shared TOML pair + per-config serve scripts, mirroring `recipe/vllm/`) for Qwen3-1.7B (single), - Qwen3-1.7B (DP), and Qwen3-30B-A3B-FP8 (EP). TP attention stays unsupported + Qwen3-1.7B (DP), and Qwen3-30B-A3B-FP8 (EP). TP attention was not yet supported (NCCL all-reduce vs the VMM region) โ€” EP uses DP-attention instead. EP requires SGLang's pinned kernels โ€” DeepEP `9af0e0d` (not vLLM's `29d31c0`), `sgl-deep-gemm >= 0.1.2`, and `flash-attn-3` (fa3 sidesteps a FlashInfer @@ -45,7 +84,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 +102,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 +110,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 +117,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 +134,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 +146,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..6efaaa1e 100644 --- a/ROADMAP.md +++ b/ROADMAP.md @@ -30,8 +30,17 @@ - [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 + - [x] Single-shape symmetric-memory TP on SGLang main + - [x] Qwen3-8B TP=2 on 2ร—H100 + - [x] Qwen3.5-122B-A10B TP=4 on 4ร—H100 + - [x] Full-backend prefill persistence with 16/64-token buckets + - [x] NEXTN speculative graph persistence on TP=2 + - [x] DeepEP low-latency graph persistence on TP=2/EP=2 + - [x] Combined TP=2/EP=2 + NEXTN + DeepEP + Full prefill + - [ ] Multi-shape and alternate-collective TP ([upstream issue](https://github.com/foundry-org/foundry/issues/6)) + - [ ] Other speculative algorithms and breakable/`tc` prefill persistence + - [ ] PP, LoRA, PDMux, diffusion, and recapture workflows ## 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..6c18d2f9 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); } @@ -2030,11 +2040,12 @@ std::shared_ptr start_graph_builds_impl( // SAVE mode. // ============================================================================ -// Per-entry finish. Idempotent on build_complete_ wait (shared_future). +// Per-entry finish. shared_future::get() is repeatable and propagates any +// background build exception to every caller. GraphLoadResult finish_one_graph_load_impl(std::shared_ptr pending, size_t index, ReconstructTensorFn reconstruct_fn) { if (pending->build_complete_.valid()) { - pending->build_complete_.wait(); + pending->build_complete_.get(); } if (index >= pending->entries.size()) { @@ -2045,9 +2056,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..b8ce2edc 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; @@ -2657,6 +2686,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 +2890,254 @@ void* dlsym(void* handle, const char* symbol) { // VMM allocations by translating to cuMemExportToShareableHandle API // ============================================================================= -// Magic marker to identify VMM IPC handles in CUipcMemHandle -static constexpr uint32_t VMM_IPC_MAGIC = 0x564D4D49; // "VMMI" +// Magic marker to identify VMM IPC handles in CUipcMemHandle. +// v2 ("VMM2"): the blob carries {magic, exporter pid, original ptr, aligned size}. +// The POSIX fd itself is transferred via SCM_RIGHTS over a per-process abstract +// unix socket (a raw fd integer is meaningless in another process - the v1 bug +// behind "cuMemImportFromShareableHandle failed with error 999"). +static constexpr uint32_t VMM_IPC_MAGIC = 0x564D4D32; // "VMM2" + +// --------------------------------------------------------------------------- +// VMM-IPC fd transport: each exporting process runs one server thread on an +// abstract unix socket "\0foundry-vmm-ipc.". Importers connect, send the +// 8-byte exporter VA they want, and receive the exported fd via SCM_RIGHTS. +// The server does no CUDA calls: fds are exported on the cuIpcGetMemHandle +// caller's thread and parked in vmm_ipc_exported_fds. +// --------------------------------------------------------------------------- + +// exporter VA -> (allocation handle it was exported from, fd). The handle is +// kept so VA reuse after free/realloc invalidates the cached fd. +static boost::unordered::concurrent_flat_map> + vmm_ipc_exported_fds; +static std::mutex vmm_ipc_server_mutex; +static int vmm_ipc_listen_fd = -1; +// pid that owns the running server thread. fork() copies this .so's state but +// not threads, so a forked child must rebind its own socket (vLLM's default +// worker start method is fork). +static pid_t vmm_ipc_server_pid = -1; +// Per-process random token carried in the handle blob and verified by the +// server: defends against pid reuse (a recycled pid + the deterministic +// shared-base VAs would otherwise let an importer silently fetch a different +// process's allocation). +static uint64_t vmm_ipc_token = 0; + +// Wire format of a fetch request. +struct VmmIpcRequest { + uint64_t ptr; + uint64_t token; +}; + +static void vmm_ipc_set_timeouts(int fd) { + struct timeval tv = {30, 0}; + setsockopt(fd, SOL_SOCKET, SO_RCVTIMEO, &tv, sizeof(tv)); + setsockopt(fd, SOL_SOCKET, SO_SNDTIMEO, &tv, sizeof(tv)); +} + +// Imported whole-chunk mappings (LOAD-mode exports: a buffer carved from the +// peer's preallocated chunk is shared by exporting the chunk handle plus an +// offset; cuMemMap cannot map a sub-range of an imported handle, so we map +// the entire peer chunk once and hand out interior pointers). The per-process +// token is part of the key so PID reuse cannot return an old exporter's mapping. +// Refcounted by open/close calls. +struct VmmIpcChunkMapping { + CUdeviceptr local_base; + size_t size; + CUmemGenericAllocationHandle handle; + int refcount; +}; +using VmmIpcChunkKey = std::tuple; +struct VmmIpcChunkSubptr { + VmmIpcChunkKey key; + int open_count; +}; +static std::mutex vmm_ipc_chunk_mutex; +static std::map vmm_ipc_chunk_mappings; +// Interior pointer handed to a caller -> owning chunk key and number of opens. +static std::map vmm_ipc_chunk_subptrs; + +// Dedicated VA zone for relocated peer imports, far below the foundry region +// (default 0x600000000000) and the large-reserve zone right above region end. +// Without this, the driver tends to place a relocated import immediately +// after the region reservation - exactly where the hooked large-reserve hint +// sends the NVSHMEM symmetric heap next, which silently breaks the heap's +// SAVE/LOAD address determinism (observed: heap displaced by a peer-chunk +// import landing at region_end). +static std::atomic vmm_ipc_import_hint{0x300000000000ULL}; + +static socklen_t vmm_ipc_socket_addr(pid_t pid, struct sockaddr_un* addr) { + memset(addr, 0, sizeof(*addr)); + addr->sun_family = AF_UNIX; + // Abstract namespace (sun_path[0] == '\0'): no filesystem entry, vanishes + // with the process. + int n = snprintf(addr->sun_path + 1, sizeof(addr->sun_path) - 2, "foundry-vmm-ipc.%d", (int)pid); + return (socklen_t)(offsetof(struct sockaddr_un, sun_path) + 1 + n); +} + +static void vmm_ipc_server_loop(int listen_fd) { + while (true) { + int conn = accept(listen_fd, nullptr, nullptr); + if (conn < 0) { + if (errno == EBADF || errno == EINVAL) + break; // socket gone + continue; // EINTR/ECONNABORTED/EMFILE/... are transient + } + vmm_ipc_set_timeouts(conn); + // Abstract sockets have no filesystem permissions: only serve same-uid peers. + struct ucred cred; + socklen_t cred_len = sizeof(cred); + if (getsockopt(conn, SOL_SOCKET, SO_PEERCRED, &cred, &cred_len) != 0 || cred.uid != getuid()) { + close(conn); + continue; + } + VmmIpcRequest req = {}; + int fd = -1; + if (recv(conn, &req, sizeof(req), MSG_WAITALL) == (ssize_t)sizeof(req) && + req.token == vmm_ipc_token) { + // Dup under the map lock: the export path may concurrently close and + // replace the parked fd (stale-entry refresh); the dup we send is ours. + vmm_ipc_exported_fds.visit( + (CUdeviceptr)req.ptr, + [&](const std::pair>& + kv) { fd = fcntl(kv.second.second, F_DUPFD_CLOEXEC, 0); }); + } + char status = (fd >= 0) ? 0 : 1; + struct iovec iov = {&status, 1}; + char cbuf[CMSG_SPACE(sizeof(int))] = {}; + struct msghdr msg = {}; + msg.msg_iov = &iov; + msg.msg_iovlen = 1; + if (fd >= 0) { + msg.msg_control = cbuf; + msg.msg_controllen = CMSG_SPACE(sizeof(int)); + struct cmsghdr* c = CMSG_FIRSTHDR(&msg); + c->cmsg_level = SOL_SOCKET; + c->cmsg_type = SCM_RIGHTS; + c->cmsg_len = CMSG_LEN(sizeof(int)); + memcpy(CMSG_DATA(c), &fd, sizeof(int)); + } + // MSG_NOSIGNAL: a peer dying mid-handshake must not SIGPIPE-kill us. + if (sendmsg(conn, &msg, MSG_NOSIGNAL) != 1) { + fprintf(stderr, "[HOOK] WARNING: VMM-IPC fd server sendmsg failed (errno=%d)\n", errno); + } + if (fd >= 0) + close(fd); + close(conn); + } +} + +static bool vmm_ipc_ensure_server() { + std::lock_guard lock(vmm_ipc_server_mutex); + pid_t pid = getpid(); + if (vmm_ipc_server_pid == pid) { + return vmm_ipc_listen_fd >= 0; + } + if (vmm_ipc_server_pid != -1) { + // fork() inherits parked descriptors even with FD_CLOEXEC. The child has + // no server thread and must not pin or serve the parent's allocations. + vmm_ipc_exported_fds.erase_if( + [](const std::pair>& kv) { + close(kv.second.second); + return true; + }); + } + // First call in this process, or first after fork (the parent's accept + // thread did not survive). Close any inherited listen fd so the parent's + // abstract name is not pinned alive by us after the parent exits. + if (vmm_ipc_listen_fd >= 0) { + close(vmm_ipc_listen_fd); + vmm_ipc_listen_fd = -1; + } + int s = socket(AF_UNIX, SOCK_STREAM | SOCK_CLOEXEC, 0); + struct sockaddr_un addr; + socklen_t len = vmm_ipc_socket_addr(pid, &addr); + if (s < 0 || bind(s, (struct sockaddr*)&addr, len) != 0 || listen(s, 64) != 0) { + fprintf(stderr, "[HOOK] ERROR: VMM-IPC fd server setup failed (errno=%d)\n", errno); + if (s >= 0) + close(s); + vmm_ipc_listen_fd = -1; + vmm_ipc_server_pid = pid; // don't retry-spam; exports in this process fail loudly + return false; + } + // Per-process token (re-derived after fork). /dev/urandom, with a clock^pid + // fallback - this is anti-accident (pid reuse), not a security boundary; + // same-uid access control is SO_PEERCRED above. + uint64_t tok = 0; + int ur = open("/dev/urandom", O_RDONLY | O_CLOEXEC); + if (ur >= 0) { + if (read(ur, &tok, sizeof(tok)) != (ssize_t)sizeof(tok)) + tok = 0; + close(ur); + } + if (tok == 0) { + struct timeval tv; + gettimeofday(&tv, nullptr); + tok = ((uint64_t)tv.tv_sec << 32) ^ (uint64_t)tv.tv_usec ^ ((uint64_t)pid << 16) ^ + 0x9e3779b97f4a7c15ULL; + } + vmm_ipc_token = tok; + vmm_ipc_listen_fd = s; + vmm_ipc_server_pid = pid; + std::thread(vmm_ipc_server_loop, s).detach(); + return true; +} + +// Importer side: fetch the fd for `original_ptr` from `exporter_pid`'s server. +// Returns a live fd in THIS process, or -1. +static int vmm_ipc_fetch_fd(pid_t exporter_pid, uint64_t token, CUdeviceptr original_ptr) { + int s = socket(AF_UNIX, SOCK_STREAM | SOCK_CLOEXEC, 0); + if (s < 0) + return -1; + vmm_ipc_set_timeouts(s); + struct sockaddr_un addr; + socklen_t len = vmm_ipc_socket_addr(exporter_pid, &addr); + if (connect(s, (struct sockaddr*)&addr, len) != 0) { + fprintf(stderr, + "[HOOK] ERROR: VMM-IPC connect to exporter pid %d failed (errno=%d) - exporter dead " + "or hook version mismatch\n", + (int)exporter_pid, errno); + close(s); + return -1; + } + VmmIpcRequest req = {(uint64_t)original_ptr, token}; + if (send(s, &req, sizeof(req), MSG_NOSIGNAL) != (ssize_t)sizeof(req)) { + close(s); + return -1; + } + char status = 1; + struct iovec iov = {&status, 1}; + char cbuf[CMSG_SPACE(sizeof(int))] = {}; + struct msghdr msg = {}; + msg.msg_iov = &iov; + msg.msg_iovlen = 1; + msg.msg_control = cbuf; + msg.msg_controllen = sizeof(cbuf); + ssize_t r = recvmsg(s, &msg, MSG_CMSG_CLOEXEC); + close(s); + if (r != 1 || status != 0) + return -1; + for (struct cmsghdr* c = CMSG_FIRSTHDR(&msg); c != nullptr; c = CMSG_NXTHDR(&msg, c)) { + if (c->cmsg_level == SOL_SOCKET && c->cmsg_type == SCM_RIGHTS) { + int fd = -1; + memcpy(&fd, CMSG_DATA(c), sizeof(int)); + return fd; + } + } + return -1; +} + +// Close and forget the parked exported fd for a VA (called from the free +// path). Without this, the fd keeps the freed allocation's physical pages +// alive and the server would hand importers a stale mapping. +static void vmm_ipc_invalidate_export(CUdeviceptr dptr) { + vmm_ipc_exported_fds.erase_if( + [dptr](const std::pair>& kv) { + if (kv.first != dptr) + return false; + close(kv.second.second); + return true; + }); +} // Hook for cuIpcGetMemHandle - intercept Driver API to support VMM allocations CUresult cuIpcGetMemHandle(CUipcMemHandle* pHandle, CUdeviceptr dptr) { @@ -2878,7 +3154,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 +3205,81 @@ CUresult cuIpcGetMemHandle(CUipcMemHandle* pHandle, CUdeviceptr dptr) { return CUDA_ERROR_NOT_SUPPORTED; } - int fd = -1; - CUresult result = - export_func(&fd, metadata.handle, CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR, 0); + // The server also owns the per-process token packed into the blob, so it + // must exist before we hand out any handle. + if (!vmm_ipc_ensure_server()) { + return CUDA_ERROR_NOT_SUPPORTED; + } - if (result != CUDA_SUCCESS) { - fprintf(stderr, - "[HOOK] ERROR: cuMemExportToShareableHandle failed with error %d for ptr=0x%llx\n", - result, (unsigned long long)dptr); - return result; + // Export once per allocation and park the fd for the server thread. + // The cached entry is keyed by VA and validated against the allocation + // handle, so VA reuse after free/realloc re-exports instead of serving a + // stale fd (the free path also eagerly invalidates via + // vmm_ipc_invalidate_export). + int fd = -1; + vmm_ipc_exported_fds.visit( + export_key, + [&](const std::pair>& kv) { + if (kv.second.first == export_handle) + fd = kv.second.second; + }); + if (fd < 0) { + int new_fd = -1; + CUresult result = + export_func(&new_fd, export_handle, CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR, 0); + if (result != CUDA_SUCCESS) { + fprintf(stderr, + "[HOOK] ERROR: cuMemExportToShareableHandle failed with error %d for ptr=0x%llx\n", + result, (unsigned long long)dptr); + return result; + } + // Keep parked fds out of exec'd children. Plain fork still inherits + // them; vmm_ipc_ensure_server closes the inherited registry before the + // child starts its own server. SCM_RIGHTS transfer is unaffected. + fcntl(new_fd, F_SETFD, FD_CLOEXEC); + vmm_ipc_exported_fds.insert_or_visit( + std::make_pair(export_key, std::make_pair(export_handle, new_fd)), + [&](std::pair>& kv) { + // Entry exists: either a racing thread won (same handle - drop + // ours) or it is stale from a freed allocation (replace). + if (kv.second.first == export_handle) { + close(new_fd); + new_fd = kv.second.second; + } else { + close(kv.second.second); + kv.second = std::make_pair(export_handle, new_fd); + } + }); + fd = new_fd; } - // Pack our custom data into the handle structure - // CUipcMemHandle has 64 reserved bytes - we use them to store our info: - // - Magic marker (4 bytes) to identify VMM handles - // - File descriptor (4 bytes) - // - Original pointer address (8 bytes) - critical for CUDA graph replay! - // - Size (8 bytes) + // Pack our custom data into the handle structure. + // CUipcMemHandle has 64 reserved bytes: + // - Magic marker (4 bytes, "VMM2") + // - Exporter pid (4 bytes) - importer fetches the fd from this process's + // VMM-IPC socket via SCM_RIGHTS (a raw fd integer is not portable) + // - Original pointer address (8 bytes) - fd lookup key + placement hint + // - Aligned size (8 bytes) + // - Per-process token (8 bytes) - server rejects mismatches (pid reuse) + // - Chunk base + chunk size (8+8 bytes) - nonzero iff the pointer is a + // carve from the preallocated chunk; the importer then maps the whole + // chunk and returns base + (ptr - chunk_base) + // - Chunk generation (8 bytes) - distinguishes free/recreate at one base memset(pHandle, 0, sizeof(CUipcMemHandle)); + uint32_t pid_u32 = (uint32_t)getpid(); memcpy(pHandle->reserved, &VMM_IPC_MAGIC, sizeof(uint32_t)); - memcpy(pHandle->reserved + 4, &fd, sizeof(int)); + memcpy(pHandle->reserved + 4, &pid_u32, sizeof(uint32_t)); memcpy(pHandle->reserved + 8, &dptr, sizeof(CUdeviceptr)); memcpy(pHandle->reserved + 16, &metadata.size, sizeof(size_t)); + memcpy(pHandle->reserved + 24, &vmm_ipc_token, sizeof(uint64_t)); + memcpy(pHandle->reserved + 32, &chunk_base, sizeof(uint64_t)); + memcpy(pHandle->reserved + 40, &chunk_size, sizeof(uint64_t)); + memcpy(pHandle->reserved + 48, &chunk_generation, sizeof(uint64_t)); #ifdef HOOK_DEBUG - fprintf(stderr, "[HOOK] cuIpcGetMemHandle: VMM ptr=0x%llx, fd=%d, size=%zu\n", + fprintf(stderr, + "[HOOK] cuIpcGetMemHandle: VMM ptr=0x%llx, fd=%d (served via socket), size=%zu\n", (unsigned long long)dptr, fd, metadata.size); #endif return CUDA_SUCCESS; @@ -2940,19 +3304,75 @@ CUresult cuIpcOpenMemHandle(CUdeviceptr* pdptr, CUipcMemHandle handle, unsigned if (magic == VMM_IPC_MAGIC) { // This is a VMM IPC handle - extract the packed data - int fd; + uint32_t exporter_pid_u32; CUdeviceptr original_ptr; size_t size; + uint64_t token; + uint64_t chunk_base; + uint64_t chunk_size; + uint64_t chunk_generation; - memcpy(&fd, handle.reserved + 4, sizeof(int)); + memcpy(&exporter_pid_u32, handle.reserved + 4, sizeof(uint32_t)); memcpy(&original_ptr, handle.reserved + 8, sizeof(CUdeviceptr)); memcpy(&size, handle.reserved + 16, sizeof(size_t)); + memcpy(&token, handle.reserved + 24, sizeof(uint64_t)); + memcpy(&chunk_base, handle.reserved + 32, sizeof(uint64_t)); + memcpy(&chunk_size, handle.reserved + 40, sizeof(uint64_t)); + memcpy(&chunk_generation, handle.reserved + 48, sizeof(uint64_t)); + pid_t exporter_pid = (pid_t)exporter_pid_u32; + // For chunk carves the fd registry on the exporter is keyed by the chunk + // base, and what we import/map is the whole chunk. + CUdeviceptr fetch_key = (chunk_base != 0) ? (CUdeviceptr)chunk_base : original_ptr; + size_t map_size = (chunk_base != 0) ? (size_t)chunk_size : size; + + if (chunk_base != 0) { + // Fast path: peer chunk already mapped -> hand out an interior pointer. + std::lock_guard lock(vmm_ipc_chunk_mutex); + VmmIpcChunkKey key = std::make_tuple(exporter_pid, token, chunk_base, chunk_generation); + auto it = vmm_ipc_chunk_mappings.find(key); + if (it != vmm_ipc_chunk_mappings.end()) { + it->second.refcount++; + CUdeviceptr mapped = it->second.local_base + (original_ptr - (CUdeviceptr)chunk_base); + auto sub_it = vmm_ipc_chunk_subptrs.find(mapped); + if (sub_it == vmm_ipc_chunk_subptrs.end()) { + vmm_ipc_chunk_subptrs[mapped] = VmmIpcChunkSubptr{key, 1}; + } else if (sub_it->second.key == key) { + sub_it->second.open_count++; + } else { + it->second.refcount--; + fprintf(stderr, + "[HOOK] ERROR: VMM-IPC interior pointer collision at 0x%llx across chunks\n", + (unsigned long long)mapped); + return CUDA_ERROR_INVALID_VALUE; + } + *pdptr = mapped; + return CUDA_SUCCESS; + } + } #ifdef HOOK_DEBUG - fprintf(stderr, "[HOOK] cuIpcOpenMemHandle: VMM fd=%d, original_ptr=0x%llx, size=%zu\n", fd, - (unsigned long long)original_ptr, size); + fprintf(stderr, + "[HOOK] cuIpcOpenMemHandle: VMM exporter_pid=%d, original_ptr=0x%llx, size=%zu\n", + (int)exporter_pid, (unsigned long long)original_ptr, size); #endif + // Obtain a live fd in THIS process: dup from our own registry for + // same-process opens, SCM_RIGHTS fetch from the exporter otherwise. + int fd = -1; + if (exporter_pid == getpid() && token == vmm_ipc_token) { + vmm_ipc_exported_fds.visit( + fetch_key, + [&](const std::pair>& + kv) { fd = fcntl(kv.second.second, F_DUPFD_CLOEXEC, 0); }); + } else { + fd = vmm_ipc_fetch_fd(exporter_pid, token, fetch_key); + } + if (fd < 0) { + fprintf(stderr, "[HOOK] ERROR: VMM-IPC could not obtain fd for ptr=0x%llx from pid %d\n", + (unsigned long long)fetch_key, (int)exporter_pid); + return CUDA_ERROR_INVALID_VALUE; + } + // Import the allocation handle from file descriptor typedef CUresult (*cuMemImportFromShareableHandle_t)(CUmemGenericAllocationHandle*, void*, CUmemAllocationHandleType); @@ -2961,12 +3381,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 +3396,58 @@ CUresult cuIpcOpenMemHandle(CUdeviceptr* pdptr, CUipcMemHandle handle, unsigned return result; } - // Map to the SAME address as original (critical for CUDA graph replay!) + typedef CUresult (*cuMemAddressReserve_t)(CUdeviceptr*, size_t, size_t, CUdeviceptr, + unsigned long long); + auto real_reserve_func = (cuMemAddressReserve_t)CUDA_DRIVER_CALL( + cuda_driver_entry_table, CUDA_ENTRY_cuMemAddressReserve); + typedef CUresult (*cuMemRelease_t)(CUmemGenericAllocationHandle); + auto mem_release_func = + (cuMemRelease_t)CUDA_DRIVER_CALL(cuda_driver_entry_table, CUDA_ENTRY_cuMemRelease); + typedef CUresult (*cuMemAddressFree_t)(CUdeviceptr, size_t); + auto real_address_free_func = + (cuMemAddressFree_t)CUDA_DRIVER_CALL(cuda_driver_entry_table, CUDA_ENTRY_cuMemAddressFree); + + // Reserve our own VA range for the peer mapping (real driver call, NOT the + // hooked wrapper - peer imports must never advance the deterministic + // cursor). The exporter's VA is passed as a hint: with per-rank disjoint + // regions it is honored and the mapping is address-stable; with today's + // shared-base symmetric layouts that range is occupied by our own region + // reservation, the driver picks elsewhere, and the caller gets a valid but + // relocated peer pointer. That is correct for table-indirect consumers + // (DeepEP buffer_ptrs_gpu, custom-allreduce RankData contents): peer VAs + // are never baked into captured graph nodes on those paths. + // First preference: the exporter's own VA (address-stable whenever it is + // locally free, e.g. future per-rank disjoint-base configs). + CUdeviceptr mapped_ptr = 0; + result = real_reserve_func(&mapped_ptr, map_size, 0, fetch_key, 0); + if (result == CUDA_SUCCESS && mapped_ptr != fetch_key) { + // Hint not honored. Do NOT keep the driver's fallback choice - it tends + // to sit immediately after existing reservations, i.e. on the NVSHMEM + // heap hint. Re-reserve inside the dedicated import zone. + real_address_free_func(mapped_ptr, map_size); + uint64_t zone_hint = + vmm_ipc_import_hint.fetch_add((map_size + ((1ULL << 30) - 1)) & ~((1ULL << 30) - 1)); + mapped_ptr = 0; + result = real_reserve_func(&mapped_ptr, map_size, 0, (CUdeviceptr)zone_hint, 0); + } + if (result != CUDA_SUCCESS) { + fprintf(stderr, "[HOOK] ERROR: cuMemAddressReserve for IPC import failed with error %d\n", + result); + mem_release_func(imported_handle); + return result; + } + typedef CUresult (*cuMemMap_t)(CUdeviceptr, size_t, size_t, CUmemGenericAllocationHandle, unsigned long long); auto map_func = (cuMemMap_t)CUDA_DRIVER_CALL(cuda_driver_entry_table, CUDA_ENTRY_cuMemMap); - result = map_func(original_ptr, size, 0, imported_handle, 0); + result = map_func(mapped_ptr, map_size, 0, imported_handle, 0); if (result != CUDA_SUCCESS) { fprintf(stderr, "[HOOK] ERROR: cuMemMap for IPC failed with error %d at addr=0x%llx size=%zu\n", - result, (unsigned long long)original_ptr, size); + result, (unsigned long long)mapped_ptr, map_size); + real_address_free_func(mapped_ptr, map_size); + mem_release_func(imported_handle); return result; } @@ -3003,27 +3467,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 +3565,36 @@ CUresult cuIpcCloseMemHandle(CUdeviceptr dptr) { auto real_func = (cuIpcCloseMemHandle_t)CUDA_DRIVER_CALL(cuda_driver_entry_table, CUDA_ENTRY_cuIpcCloseMemHandle); + // Interior pointer into an imported peer chunk? Refcounted: the chunk + // mapping is torn down when its last interior pointer closes. + { + std::lock_guard lock(vmm_ipc_chunk_mutex); + auto sub_it = vmm_ipc_chunk_subptrs.find(dptr); + if (sub_it != vmm_ipc_chunk_subptrs.end()) { + VmmIpcChunkKey key = sub_it->second.key; + if (--sub_it->second.open_count == 0) { + vmm_ipc_chunk_subptrs.erase(sub_it); + } + auto it = vmm_ipc_chunk_mappings.find(key); + if (it != vmm_ipc_chunk_mappings.end() && --it->second.refcount == 0) { + typedef CUresult (*cuMemUnmap_t)(CUdeviceptr, size_t); + auto mem_unmap_func = + (cuMemUnmap_t)CUDA_DRIVER_CALL(cuda_driver_entry_table, CUDA_ENTRY_cuMemUnmap); + typedef CUresult (*cuMemRelease_t)(CUmemGenericAllocationHandle); + auto mem_release_func = + (cuMemRelease_t)CUDA_DRIVER_CALL(cuda_driver_entry_table, CUDA_ENTRY_cuMemRelease); + typedef CUresult (*cuMemAddressFree_t)(CUdeviceptr, size_t); + auto addr_free_func = (cuMemAddressFree_t)CUDA_DRIVER_CALL(cuda_driver_entry_table, + CUDA_ENTRY_cuMemAddressFree); + mem_unmap_func(it->second.local_base, it->second.size); + mem_release_func(it->second.handle); + addr_free_func(it->second.local_base, it->second.size); + vmm_ipc_chunk_mappings.erase(it); + } + return CUDA_SUCCESS; + } + } + AllocMetadata metadata; bool found = false; @@ -3059,8 +3613,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 +3828,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 +3855,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..e6a32e56 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,48 @@ 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. +## Tensor parallel + +SGLang-main TP uses one explicit padded **decode** graph shape and torch +symmetric memory. +`_patch_torch_symm_mem` creates the buffer outside Foundry's VMM, and +`_patch_multimem_all_gather` reuses that stable buffer for PyTorch's +`multimem_all_gather_out`. The plugin rejects TP without symmetric memory and +multi-shape decode archives. Full-backend prefill is a separate phase and may +persist multiple token buckets. + +Validated behavior: + +- Qwen3-8B TP=2 on 2ร—H100; +- Qwen3.5-122B-A10B-GPTQ-Int4 TP=4 on 4ร—H100; +- one decode shape plus Full prefill 16/64 in the TP4 regression; +- three graphs and identical offset `59301167104` on every TP4 rank; +- sequential and fixed-batch LOAD outputs match baseline. + +The repeatable validation entry point is `tests/modal_sglang_tp.py`. +It inspects each rank's graph archive to require symmetric all-reduce and +all-gather kernels with no captured NCCL kernels. + +### Full prefill and NEXTN ordering + +Only the Full prefill backend is GPU-validated for persistence. Foundry defers +its capture until target decode state exists, pads hybrid Mamba replay metadata, +and preallocates shape-specific base metadata before any prefill graph. For each +plain GDN shape, it then primes only the FLA-derived tensors required by that +shape immediately before SAVE capture or LOAD preparation. Eagle/NEXTN retains +SGLang's native allocation order. This preserves separate 16/64-token identities +without depending on the capacity of FLA's identity cache. + +NEXTN is the only GPU-validated speculative algorithm. In the combined +Qwen3.5-35B-A3B run, the catalog contained target verify/decode, target Full +prefill 16/64, and draft decode sessions on both ranks. No claim is made for +other speculative algorithms or breakable/`tc` prefill persistence. + ## 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`. Plain EP can run DP-attention plus +DeepEP without TP attention collectives. Combined TP=2/EP=2 installs these +additions alongside the symmetric TP path described above. - **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 +222,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 @@ -202,8 +248,21 @@ TP attention is unsupported (its NCCL all-reduce is incompatible with the VMM re (mirrors the bg thread). Dense graphs never hit it (no event nodes), which is why it surfaced only on sglang EP. +The combined Hopper proof is pinned to SGLang +`9b853e6832e71a3058212df02a025232a453e146` and Modal run +`eb1ec00b7004339c5150ce6c-1784906247386993159` at immutable Foundry +`eb1ec00b7004339c5150ce6c3ba057ceda047833`. It uses FA3, +`SGLANG_MEM_FRACTION_STATIC=0.7`, `NVSHMEM_REMOTE_TRANSPORT=none`, +`NVSHMEM_IB_ENABLE=0`, and `NVSHMEM_IB_ENABLE_IBGDA=0`. DeepEP bootstrap +precedes graph builds, and LOAD initializes one archived NVSHMEM module per +rank. Legacy NCCL `P2P/IPC` channels were observed; other fabrics were not +validated. For no-fabric DeepEP NVL, `FOUNDRY_DEEPEP_NVL_IPC=1` selects the +separately validated VMM-backed CUDA-IPC path while NVSHMEM still owns the RDMA +symmetric heap. + 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..bf4a911e 100644 --- a/docs/sglang/overview.md +++ b/docs/sglang/overview.md @@ -2,7 +2,9 @@ Foundry persists SGLang's `CudaGraphRunner` graphs to disk on SAVE and restores them on LOAD, skipping graph capture, kernel warmup, and the per-batch-size attention metadata setup costs. -Tested on single-GPU Qwen3-1.7B / 4B / 14B and **data-parallel (DP=2)** Qwen3-1.7B with the FlashInfer attention backend. +Tested on single-GPU Qwen3-1.7B / 4B / 14B, data-parallel Qwen3-1.7B, +tensor-parallel Qwen3-8B TP=2, Qwen3.5-122B-A10B TP=4, and combined +Qwen3.5-35B-A3B online-FP8 TP=2/EP=2. ## Parallelism @@ -10,16 +12,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. | -| Expert parallel (DeepEP) | โœ… | Validated EP=2 on Qwen3-30B-A3B-FP8 (SAVE/SAVE2/LOAD/query); restored decode graphs match baseline throughput. See **Expert parallel** below. | - -**Expert parallel (DeepEP).** EP runs DP-attention (each rank its own attention โ€” no -NCCL all-reduce) + DeepEP for the MoE all-to-all (NVSHMEM, foundry-compatible). The +| Tensor parallel (TP) | โœ… | One explicit padded torch-symmetric-memory decode shape; validated with Qwen3-8B TP=2 and Qwen3.5-122B-A10B TP=4 on H100. Full prefill may add multiple token buckets. | +| Expert parallel (DeepEP) | โœ… | Validated EP=2 on Qwen3-30B-A3B-FP8 and in the Qwen3.5-35B-A3B combined TP=2/EP=2 gate. See **Expert parallel** below. | + +**Tensor parallel.** The TP recipe is +`recipe/experimental/serve_qwen3-8b_sglang_tp.sh +[--save|--load]`. SGLang main defaults to torch symmetric memory and one +explicit padded decode shape. Foundry reuses the stable symmetric buffer for +all-reduce and all-gather, disables other custom/fused collective paths, and +rejects multi-shape decode archives. Full-backend prefill persistence is +separate and may carry multiple token buckets. The checked-in proof verifies +sequential and fixed-batch outputs, every rank's collective kernels, archive +fingerprints, allocation offsets, replay, and absence of recapture: +`modal run --env tests/modal_sglang_tp.py`. + +Validation is pinned to SGLang +`9b853e6832e71a3058212df02a025232a453e146`, immutable Foundry +`eb1ec00b7004339c5150ce6c3ba057ceda047833`, CUDA 13, PyTorch 2.11, and H100. The +Qwen3.5-122B-A10B-GPTQ-Int4 TP4 gate used `moe_wna16`, FA3, Triton linear +attention/MoE, an 8-request decode graph, and Full prefill buckets 16/64. Run +`eb1ec00b7004339c5150ce6c-1784907539317556437` restored three graphs and both +decode/prefill catalog sessions on each rank at the symmetric offset +`59301167104`. +Broader decode-shape and collective support remains tracked in +[issue #6](https://github.com/foundry-org/foundry/issues/6). + +## Persisted graph modes + +| Mode | GPU-validated scope | +|---|---| +| Prefill | Full backend with 16/64-token buckets. Persisted hybrid-linear prefill is validated only for GDN. Non-GDN hybrid-linear backends and configurations with forward hooks fall back to eager prefill. Breakable and `tc` persistence are not validated. | +| Speculative decoding | NEXTN only. The validated setting is 3 steps, top-k 1, and 4 draft tokens. | +| DeepEP | `low_latency`; validated standalone and with TP, NEXTN, and Full prefill. | +| Combined | Qwen3.5-35B-A3B online FP8, TP=2/EP=2, FA3, one decode-8 shape, NEXTN, DeepEP low-latency, and Full prefill 16/64. | + +When Foundry is active and Full prefill is requested, GDN is the only +GPU-validated persisted hybrid-linear backend. Mamba2, Lightning, ShortConv, +and other non-GDN hybrid-linear backends use SGLang eager prefill instead. +Configurations with nonempty `forward_hooks` also use eager prefill so hooks +are not traced into a deferred persisted graph. These fallbacks leave the +configured prefill backend value intact after runner setup and do not disable +Foundry decode persistence. + +The combined proof is Modal run +`eb1ec00b7004339c5150ce6c-1784906247386993159`. Every rank restored target +verify/decode, target prefill, and draft decode sessions (four graphs total); +SAVE/SAVE2 semantic fingerprints and rank-local offsets `75768004608` and +`56774098944` matched. The fresh LOAD reproduced exact batched token IDs, +average NEXTN acceptance `2.775`, DeepEP dispatch/combine nodes, +bootstrap-before-build ordering, one loaded NVSHMEM module per rank, and no +SAVE/recapture/runtime errors. + +**Expert parallel (DeepEP).** Plain EP can run DP-attention (each rank its own +attention) plus DeepEP for the MoE all-to-all. The combined TP=2/EP=2 gate +instead composes torch-symmetric-memory TP collectives with DeepEP. 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 +81,25 @@ 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. +On the pinned Hopper combined path, use FA3 and keep +`SGLANG_MEM_FRACTION_STATIC=0.7`, `NVSHMEM_REMOTE_TRANSPORT=none`, +`NVSHMEM_IB_ENABLE=0`, and `NVSHMEM_IB_ENABLE_IBGDA=0` identical across phases. +The observed intranode NCCL channels were legacy `P2P/IPC`. These settings do +not validate InfiniBand or another NVSHMEM fabric. + +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. + +## Retained non-goals + +The current GPU evidence does not cover speculative algorithms other than +NEXTN, multi-shape decode TP, arbitrary collective or fabric backends, +breakable/`tc` prefill persistence, pipeline parallelism, LoRA, PDMux, +diffusion, or restoring a graph after runtime recapture. + **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 @@ -99,8 +169,8 @@ 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. -5. `self.graphs` / `self.output_buffers` are populated from `state.loaded_graphs`; the rest of SGLang's serving path runs unchanged. +4. `CudaGraphRunner.capture` is replaced with: preallocate the deterministic range up to `final_alloc_offset` once, then call `start_graph_builds(session_paths)` for each catalog session. +5. As SGLang requests that session's shapes in archive order, Foundry calls `finish_one_graph_load(pending, index)` for each graph and populates the backend's graph/output maps. After the final session is validated, replay performs no catalog I/O; the rest of SGLang's serving path runs unchanged. ## Doc set diff --git a/docs/superpowers/plans/2026-07-23-sglang-main-graph-modes.md b/docs/superpowers/plans/2026-07-23-sglang-main-graph-modes.md new file mode 100644 index 00000000..ac4f9bd0 --- /dev/null +++ b/docs/superpowers/plans/2026-07-23-sglang-main-graph-modes.md @@ -0,0 +1,1138 @@ +# SGLang Main Graph Modes 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:** Persist SGLang-main full CUDA graph sessions for DeepEP, Qwen3.5 +NEXTN speculative decoding, and prefill, while retaining verified TP4 behavior. + +**Architecture:** Add a versioned per-rank SGLang graph catalog and a +deterministic tensor-tree output codec, then generalize the existing +`BaseCudaGraphBackend` replacement to consume both. Route full decode and full +prefill to that backend, preserve upstream non-full backends, and bootstrap +DeepEP before graph capture/linking. SGLang continues to own static inputs, +attention metadata, speculative runner behavior, and the eager prefill tail. + +**Tech Stack:** Python 3.12, PyTorch 2.11/cu130, Foundry C++/CUDA bindings, +SGLang `9b853e6832e71a3058212df02a025232a453e146`, pytest, Ruff, Modal H100. + +## Global Constraints + +- Keep the SGLang integration out-of-tree and activated through + `sglang.srt.plugins`. +- Persist only SGLang's `full` graph backend. `breakable`, `tc_piecewise`, and + `disabled` remain upstream pass-through modes. +- Preserve the current guards for pipeline parallelism, LoRA, PD multiplexing, + explicit return-hidden-states mode, diffusion models, and graph recapture. +- DeepEP graph persistence supports `auto`/`low_latency`; SGLang's eager + fallback for `normal` remains authoritative. +- Archives are rank-topology-specific and use complete graph dependencies + (`strip_dependencies=False`). +- Every production behavior starts with a focused failing test and a verified + RED result. +- Before each test phase on the cloud branch, commit and push the pre-test + revision and update the draft pull request. Commit the implementation + separately after GREEN. +- Do not claim a model/configuration in user docs until a fresh Modal run + produces exact output and restored-replay evidence. +- Run GPU work only in Modal's `rahul-dev` environment; the local VM has no + CUDA device or toolkit. + +## File Structure + +### New files + +- `python/foundry/integration/sglang/output_codec.py`: flatten and reconstruct + allowlisted SGLang output trees. +- `python/foundry/integration/sglang/archive.py`: versioned graph-session + catalog and atomic persistence. +- `python/foundry/integration/sglang/main_deepep.py`: SGLang-main DeepEP + buffer bootstrap and graph-mode selection. +- `tests/test_sglang_output_codec.py`: CPU output codec contracts. +- `tests/test_sglang_graph_catalog.py`: CPU catalog contracts. +- `tests/test_sglang_main_deepep.py`: CPU DeepEP bootstrap contracts with + SGLang stubs. +- `tests/modal_deepep_fabric.py`: two-H100 Modal wrapper for the standalone + semantic DeepEP SAVE/LOAD oracle. + +### Modified files + +- `python/foundry/integration/sglang/runtime.py`: process-local next-session + cursor. +- `python/foundry/integration/sglang/main_backend.py`: generalized + phase/session capture and load. +- `python/foundry/integration/sglang/hooks_main.py`: decode and prefill + resolver routing. +- `python/foundry/integration/sglang/plugin.py`: remove the three feature + rejections and stop forcing prefill off. +- `tests/test_sglang_main_compat.py`: plugin and resolver compatibility + contracts. +- `recipe/experimental/serve_qwen3-8b_sglang_tp.sh`: explicit prefill, + speculative, and DeepEP options. +- `recipe/sglang/serve_qwen3-mini.sh`: opt-in full-prefill smoke options. +- `tests/test_sglang_tp_recipe.py`: recipe argument contracts. +- `tests/modal_sglang_main_smoke.py`: full-prefill session/replay checks. +- `tests/modal_sglang_tp.py`: catalog, speculative, and DeepEP observability. +- `README.md`, `RELEASE.md`, `ROADMAP.md`, `docs/sglang/overview.md`, + `docs/sglang/hooks.md`, `recipe/sglang/README.md`, and + `recipe/experimental/README.md`: evidence-backed support matrix. + +--- + +### Task 1: Tensor-tree output codec + +**Files:** + +- Create: `python/foundry/integration/sglang/output_codec.py` +- Create: `tests/test_sglang_output_codec.py` + +**Interfaces:** + +- Produces: + `pack_output(output: Any, *, allowed_types: tuple[type, ...]) -> PackedOutput` +- Produces: + `unpack_output(tensors: Any, schema: dict[str, Any], *, allowed_types: tuple[type, ...]) -> Any` +- Produces: + `PackedOutput(tensors: list[torch.Tensor], schema: dict[str, Any])` +- Consumes no SGLang imports; callers supply the allowed output classes. + +- [ ] **Step 1: Write codec round-trip and error tests** + +Create the test file with SPDX headers and these concrete contracts: + +```python +from __future__ import annotations + +import torch +import pytest + +from foundry.integration.sglang.output_codec import pack_output, unpack_output + + +class FakeLogitsOutput: + def __init__(self, next_token_logits): + self.next_token_logits = next_token_logits + self.hidden_states = None + + +def test_output_codec_round_trips_dynamic_nested_tensor_tree() -> None: + logits = torch.arange(12).reshape(3, 4) + hidden = torch.arange(6).reshape(3, 2) + output = FakeLogitsOutput(logits) + output.hidden_states = hidden + output.topk_p = logits + value = [output, (hidden, {"shared": logits, "constant": 7})] + + packed = pack_output(value, allowed_types=(FakeLogitsOutput,)) + restored = unpack_output( + tuple(packed.tensors), + packed.schema, + allowed_types=(FakeLogitsOutput,), + ) + + assert len(packed.tensors) == 2 + assert isinstance(restored[0], FakeLogitsOutput) + assert restored[0].next_token_logits is restored[0].topk_p + assert restored[0].next_token_logits is restored[1][1]["shared"] + assert torch.equal(restored[0].hidden_states, hidden) + assert restored[1][1]["constant"] == 7 + + +def test_output_codec_rejects_an_unsupported_leaf_with_path() -> None: + with pytest.raises(TypeError, match=r"\$\.bad"): + pack_output({"bad": object()}, allowed_types=()) + + +def test_output_codec_rejects_an_untrusted_object_schema() -> None: + with pytest.raises(ValueError, match="not allowlisted"): + unpack_output( + [], + {"kind": "object", "type": "evil.Payload", "attributes": {}}, + allowed_types=(), + ) + + +def test_output_codec_rejects_tensor_count_mismatch() -> None: + tensor = torch.ones(1) + packed = pack_output(tensor, allowed_types=()) + with pytest.raises(ValueError, match="tensor count"): + unpack_output([], packed.schema, allowed_types=()) +``` + +- [ ] **Step 2: Commit and push the RED test** + +```bash +git add tests/test_sglang_output_codec.py +git commit -m "test: define SGLang graph output codec" +git push -u origin cursor/tp-consolidation-2e2c +``` + +- [ ] **Step 3: Run the focused test and verify RED** + +Run: + +```bash +PYTHONPATH=python pytest tests/test_sglang_output_codec.py -q +``` + +Expected: collection fails because +`foundry.integration.sglang.output_codec` does not exist. + +- [ ] **Step 4: Implement the codec** + +Implement a `PackedOutput` frozen dataclass and recursive encoder/decoder with +these exact schema kinds: + +```python +@dataclass(frozen=True) +class PackedOutput: + tensors: list[torch.Tensor] + schema: dict[str, Any] + + +def pack_output( + output: Any, + *, + allowed_types: tuple[type, ...], +) -> PackedOutput: + tensors: list[torch.Tensor] = [] + tensor_indices: dict[int, int] = {} + active_containers: set[int] = set() + allowed = {_type_name(cls): cls for cls in allowed_types} + + def encode(value: Any, path: str) -> dict[str, Any]: + if isinstance(value, torch.Tensor): + identity = id(value) + index = tensor_indices.get(identity) + if index is None: + index = len(tensors) + tensor_indices[identity] = index + tensors.append(value) + return {"kind": "tensor", "index": index} + if value is None or isinstance(value, (bool, int, float, str)): + return {"kind": "constant", "value": value} + if isinstance(value, (list, tuple)): + return _encode_sequence(value, path, encode, active_containers) + if isinstance(value, dict): + return _encode_mapping(value, path, encode, active_containers) + + type_name = _type_name(type(value)) + if type_name not in allowed: + raise TypeError( + f"Unsupported SGLang graph output at {path}: {type(value)!r}" + ) + attributes = { + name: encode(attribute, f"{path}.{name}") + for name, attribute in sorted(vars(value).items()) + } + return { + "kind": "object", + "type": type_name, + "attributes": attributes, + } + + return PackedOutput(tensors=tensors, schema=encode(output, "$")) +``` + +Implement `_encode_sequence` and `_encode_mapping` so they detect recursive +containers, preserve list versus tuple, require string dictionary keys, and +sort dictionary keys. Implement `unpack_output` with the inverse traversal, +`cls.__new__(cls)` plus `setattr` for allowlisted objects, tensor input +normalization, and an exact tensor-count check based on referenced indices. + +- [ ] **Step 5: Commit and push the implementation** + +```bash +git add python/foundry/integration/sglang/output_codec.py +git commit -m "feat: serialize SGLang graph output trees" +git push -u origin cursor/tp-consolidation-2e2c +``` + +- [ ] **Step 6: Run GREEN and lint** + +Run: + +```bash +PYTHONPATH=python pytest tests/test_sglang_output_codec.py -q +ruff check python/foundry/integration/sglang/output_codec.py tests/test_sglang_output_codec.py +ruff format --check python/foundry/integration/sglang/output_codec.py tests/test_sglang_output_codec.py +``` + +Expected: four tests pass and both Ruff commands exit zero. + +--- + +### Task 2: Versioned graph-session catalog + +**Files:** + +- Create: `python/foundry/integration/sglang/archive.py` +- Create: `tests/test_sglang_graph_catalog.py` + +**Interfaces:** + +- Produces: + `shape_key_record(shape_key: Any) -> dict[str, int | str | None]` +- Produces: + `runner_descriptor(runner: Any, phase: str) -> dict[str, Any]` +- Produces: `GraphCatalog.open_save_session(index, descriptor) -> dict` +- Produces: `GraphCatalog.claim_load_session(index, descriptor) -> dict` +- Produces: `GraphCatalog.append_graph(index, graph_record) -> None` +- Produces: `GraphCatalog.session_count -> int` +- Consumed by `FoundryMainCudaGraphBackend`. + +- [ ] **Step 1: Write catalog behavior tests** + +Cover all of the following in `tests/test_sglang_graph_catalog.py`: + +```python +def test_catalog_distinguishes_equal_shapes_in_separate_sessions(tmp_path): + catalog = GraphCatalog(tmp_path) + target = {"class": "Target", "phase": "decode", "role": "target", + "forward_mode": "TARGET_VERIFY", "step": None} + draft = {"class": "Draft", "phase": "decode", "role": "draft", + "forward_mode": "DECODE", "step": None} + + catalog.open_save_session(0, target) + catalog.append_graph( + 0, + { + "capture_index": 0, + "filename": "graph_0_decode_target_s0_k8.json", + "shape_key": {"size": 8, "stream_idx": None, "variant_label": None}, + "output_schema": {"kind": "tensor", "index": 0}, + }, + ) + catalog.open_save_session(1, draft) + catalog.append_graph( + 1, + { + "capture_index": 1, + "filename": "graph_1_decode_draft_s1_k8.json", + "shape_key": {"size": 8, "stream_idx": None, "variant_label": None}, + "output_schema": {"kind": "tuple", "items": []}, + }, + ) + + reloaded = GraphCatalog(tmp_path) + assert reloaded.claim_load_session(0, target)["graphs"][0]["capture_index"] == 0 + assert reloaded.claim_load_session(1, draft)["graphs"][0]["capture_index"] == 1 + + +def test_catalog_rejects_runner_order_mismatch(tmp_path): + catalog = GraphCatalog(tmp_path) + descriptor = {"class": "Target", "phase": "decode", "role": "target", + "forward_mode": "DECODE", "step": None} + catalog.open_save_session(0, descriptor) + + with pytest.raises(RuntimeError, match="session 0 descriptor mismatch"): + GraphCatalog(tmp_path).claim_load_session( + 0, {**descriptor, "forward_mode": "TARGET_VERIFY"} + ) + + +def test_catalog_writes_valid_json_after_every_append(tmp_path): + catalog = GraphCatalog(tmp_path) + descriptor = {"class": "Prefill", "phase": "prefill", "role": "target", + "forward_mode": "EXTEND", "step": None} + catalog.open_save_session(0, descriptor) + json.loads((tmp_path / CATALOG_FILENAME).read_text()) + catalog.append_graph(0, _graph_record(0, 16)) + assert json.loads((tmp_path / CATALOG_FILENAME).read_text())["version"] == 1 +``` + +Also test unknown catalog versions, non-contiguous save session indices, and +`shape_key_record` with `size`, `stream_idx`, and `variant_label`. + +- [ ] **Step 2: Commit/push and verify RED** + +Commit the tests, push, then run: + +```bash +PYTHONPATH=python pytest tests/test_sglang_graph_catalog.py -q +``` + +Expected: import failure because `archive.py` does not exist. + +- [ ] **Step 3: Implement `GraphCatalog`** + +Use `CATALOG_FILENAME = "sglang_graph_catalog.json"` and +`CATALOG_VERSION = 1`. `GraphCatalog.__init__` loads an existing file or starts +with `{"version": 1, "sessions": []}`. `_write` must serialize to +`.tmp.`, flush, and replace the destination with `os.replace`. + +`open_save_session` must require `index == len(sessions)`, append +`{"session_index": index, "runner": descriptor, "graphs": []}`, write, and +return that dictionary. `claim_load_session` must validate the index exists +and compare the complete descriptor. `append_graph` must require the graph's +`capture_index` to be greater than every existing graph index and append to +the named session. + +`runner_descriptor` uses an explicit `phase` argument and emits: + +```python +{ + "class": f"{type(runner).__module__}.{type(runner).__qualname__}", + "phase": phase, + "role": ( + "draft" + if getattr(runner.model_runner, "is_draft_worker", False) + else "target" + ), + "forward_mode": getattr( + getattr(runner, "capture_forward_mode", None), "name", None + ), + "step": getattr(runner, "step", None), +} +``` + +- [ ] **Step 4: Commit/push and verify GREEN** + +Run: + +```bash +PYTHONPATH=python pytest tests/test_sglang_graph_catalog.py -q +ruff check python/foundry/integration/sglang/archive.py tests/test_sglang_graph_catalog.py +ruff format --check python/foundry/integration/sglang/archive.py tests/test_sglang_graph_catalog.py +``` + +Expected: all catalog tests and lint pass. + +Commit message: `feat: catalog SGLang graph sessions`. + +--- + +### Task 3: Generalize the Foundry full backend + +**Files:** + +- Modify: `python/foundry/integration/sglang/runtime.py:44-49` +- Modify: `python/foundry/integration/sglang/main_backend.py` +- Test: `tests/test_sglang_graph_catalog.py` +- Test: `tests/test_sglang_output_codec.py` + +**Interfaces:** + +- Consumes `GraphCatalog`, `runner_descriptor`, `shape_key_record`, + `pack_output`, and `unpack_output`. +- Changes constructor to + `FoundryMainCudaGraphBackend(cuda_graph_runner, *, phase: str)`. +- Produces per-session graph filenames and catalog records. +- Preserves SGLang's `BaseCudaGraphBackend` methods unchanged. + +- [ ] **Step 1: Add failing state/session contracts** + +Extend catalog tests to create two fake runners with the same shape and assert +their descriptors differ by phase/role/step. Extend codec tests with a fake +logits object containing `next_token_logits`, `hidden_states`, `topk_p`, and +`topk_index`, plus a list of two such objects. + +Add to `CUDAGraphExtensionState` test coverage through a direct assertion that +a new state starts with `session_index == 0`. + +- [ ] **Step 2: Commit/push and verify RED** + +Run: + +```bash +PYTHONPATH=python pytest \ + tests/test_sglang_graph_catalog.py \ + tests/test_sglang_output_codec.py -q +``` + +Expected: failures for the absent `session_index` field and missing extended +tree behavior. + +- [ ] **Step 3: Add the runtime session cursor** + +Change the state dataclass to: + +```python +@dataclass +class CUDAGraphExtensionState: + capture_index: int = 0 + session_index: int = 0 + rank: int = 0 + loaded_graphs: dict = field(default_factory=dict) +``` + +- [ ] **Step 4: Replace size-only backend persistence** + +In `main_backend.py`: + +1. Import `LogitsProcessorOutput` as the only initially allowlisted object + class. +2. Store `phase`, `descriptor`, `catalog`, and `session` on construction. +3. In `capture_session`, create/claim the current session using + `state.session_index`, increment the cursor once, and on LOAD start builds + only for that session's graph files. +4. Replace `_validate_shape_key` with `shape_key_record`; continue rejecting + stream and LoRA variants through the existing plugin guards, not through + archive identity. +5. On SAVE, call: + +```python +packed = pack_output(output, allowed_types=(LogitsProcessorOutput,)) +filename = ( + f"graph_{state.capture_index}_{self._phase}_{self._descriptor['role']}" + f"_s{self._session['session_index']}_k{shape['size']}.json" +) +graph.save(os.path.join(cfg.workspace_dir, filename), packed.tensors) +self._catalog.append_graph( + self._session["session_index"], + { + "capture_index": state.capture_index, + "filename": filename, + "shape_key": shape, + "output_schema": packed.schema, + }, +) +``` + +1. On LOAD, compare the requested shape dictionary with the next graph record, + finish that session-local graph index, and call `unpack_output` using the + record schema. +2. Keep graph and reconstructed output strongly referenced in `_graphs` and + `_outputs`. +3. In LOAD `replay_session`, require + `state.session_index == catalog.session_count` before the first replay. All + target/draft runners have been constructed by then, so this rejects a + trailing archived session before a response is served. +4. Finalize manifest/fatbins/final offset after every completed SAVE session. +5. Include phase, role, session, and shape in save/load log lines. + +- [ ] **Step 5: Commit/push and run CPU GREEN** + +Run: + +```bash +PYTHONPATH=python pytest \ + tests/test_sglang_graph_catalog.py \ + tests/test_sglang_output_codec.py \ + tests/test_sglang_main_compat.py -q +ruff check \ + python/foundry/integration/sglang/archive.py \ + python/foundry/integration/sglang/output_codec.py \ + python/foundry/integration/sglang/runtime.py \ + python/foundry/integration/sglang/main_backend.py +``` + +Expected: all focused CPU tests pass. Commit message: +`feat: persist multiple SGLang graph sessions`. + +--- + +### Task 4: Route full decode and full prefill without mode rejection + +**Files:** + +- Modify: `python/foundry/integration/sglang/hooks_main.py:12-24,201-219` +- Modify: `python/foundry/integration/sglang/main_backend.py:45` +- Modify: `python/foundry/integration/sglang/plugin.py:18-68` +- Modify: `tests/test_sglang_main_compat.py` + +**Interfaces:** + +- Decode resolver creates `FoundryMainCudaGraphBackend(runner, phase="decode")`. +- Prefill resolver creates `FoundryMainCudaGraphBackend(runner, phase="prefill")`. +- Non-full resolvers call their captured original function. +- EAGLE/NEXTN full-prefill construction retains SGLang's target `FULL` and + draft `LAST` hidden-state capture modes. + +- [ ] **Step 1: Change plugin tests first** + +Update the fake server args with `cuda_graph_backend_prefill=None`. Assert the +BEFORE hook: + +```python +assert server_args.cuda_graph_backend_decode == "full" +assert server_args.disable_prefill_cuda_graph is False +``` + +Set resolved config to full decode/full prefill, then set: + +```python +server_args.speculative_algorithm = "NEXTN" +server_args.moe_a2a_backend = "deepep" +after_post_init(None, server_args) +``` + +The call must not raise. Retain the TP symmetric-memory and one-decode-shape +negative tests. + +Add a resolver-routing test with stub `backend_utils`, `decode_runner`, and +`prefill_runner` modules. Verify full decode/full prefill construct Foundry +with the expected phase, while `breakable` returns the original sentinel. + +Add a pinned-compatibility test for the speculative prefill constructor +wrapper. Its fake original function must observe `prefill.backend == +"breakable"` and return a runner whose `prefill_backend_name` initially +matches that value. Verify the wrapper restores the server config and returned +runner to `"full"`. Run the same test for target and draft model runners, and +verify non-EAGLE or non-full calls are unchanged. + +- [ ] **Step 2: Commit/push and verify RED** + +Run: + +```bash +PYTHONPATH=python pytest tests/test_sglang_main_compat.py -q +``` + +Expected: the old hook still disables prefill and rejects NEXTN/DeepEP. + +- [ ] **Step 3: Implement plugin and resolver behavior** + +In `_configure_server_args`, set decode full only when the convenience field is +unset, leave `disable_prefill_cuda_graph` untouched, and retain autotune/profile +settings. Remove the prefill-disabled, speculative, and DeepEP rejection +branches. + +Condition TP-specific requirements on at least one resolved full phase: + +```python +uses_foundry_full = ( + server_args.cuda_graph_config.decode.backend == "full" + or server_args.cuda_graph_config.prefill.backend == "full" +) +``` + +Keep the one-explicit-decode-shape requirement only when decode is full. + +Patch both resolver symbols in their utility module and runner modules: + +```python +def foundry_backend_or_original(runner, *, phase, original): + if get_graph_extension_mode() == CUDAGraphExtensionMode.NONE: + return original(runner) + phase_config = getattr(runner.model_runner.server_args.cuda_graph_config, phase) + if phase_config.backend != "full": + logger.info("[Foundry] SGLang %s backend=%s is not persisted", phase, phase_config.backend) + return original(runner) + return FoundryMainCudaGraphBackend(runner, phase=phase) +``` + +Make `FoundryMainCudaGraphBackend` inherit SGLang's +`FullCudaGraphBackend` while continuing to implement its own constructor and +all backend methods. This gives `PrefillCudaGraphRunner` the correct +`isinstance(..., FullCudaGraphBackend)` behavior. + +Patch `cuda_graph_setup.capture_prefill_graph` with a context-variable guard. +When Foundry is active, the configured prefill backend is full, and +`model_runner.spec_algorithm.is_eagle()` is true: + +1. set the guard; +2. temporarily set `cuda_graph_config.prefill.backend` to `Backend.BREAKABLE`; +3. call the original function so SGLang bypasses its target skip and selects + target FULL or draft LAST hidden capture; +4. restore the configured backend in `finally`; and +5. set the returned runner's `prefill_backend_name` to `Backend.FULL`. + +The patched prefill resolver checks the guard before its normal backend-name +branch and returns Foundry's full-style backend while the temporary value is +breakable. This is the only backend-name compatibility override. + +- [ ] **Step 4: Commit/push and verify GREEN** + +Run: + +```bash +PYTHONPATH=python pytest tests/test_sglang_main_compat.py -q +ruff check \ + python/foundry/integration/sglang/hooks_main.py \ + python/foundry/integration/sglang/plugin.py \ + tests/test_sglang_main_compat.py +``` + +Expected: compatibility tests pass. Commit message: +`feat: route SGLang full prefill graphs through Foundry`. + +--- + +### Task 5: Bootstrap SGLang-main DeepEP symmetrically + +**Files:** + +- Create: `python/foundry/integration/sglang/main_deepep.py` +- Create: `tests/test_sglang_main_deepep.py` +- Modify: `python/foundry/integration/sglang/main_backend.py` + +**Interfaces:** + +- Produces: `bootstrap_deepep_buffer(cuda_graph_runner: Any) -> bool` +- Produces: `set_deepep_graph_mode() -> bool` +- Consumed before SAVE/LOAD session setup and inside `replay_session`. + +- [ ] **Step 1: Write DeepEP state/bootstrap tests** + +Build fake top-level SGLang modules in `sys.modules` before loading +`main_deepep.py`. The fake API must model: + +```python +class FakeDeepEPBuffer: + state = SimpleNamespace(buffer=None, dispatch_mode=None) + + @classmethod + def _state(cls): + return cls.state + + @classmethod + def set_dispatch_mode(cls, mode): + cls.state.dispatch_mode = mode +``` + +Use a fake dispatcher wrapper whose `_inners` contains a fake +`DeepEPDispatcher` with `_low_latency_dispatcher._get_buffer`. Assert: + +- inactive MoE A2A returns `False`; +- active bootstrap calls `_get_buffer` exactly once and is idempotent; +- mode resolution receives `is_extend_in_batch=False`; +- no dispatcher raises a descriptive `RuntimeError` rather than entering + stream capture uninitialized. + +- [ ] **Step 2: Commit/push and verify RED** + +Run: + +```bash +PYTHONPATH=python pytest tests/test_sglang_main_deepep.py -q +``` + +Expected: import failure because `main_deepep.py` does not exist. + +- [ ] **Step 3: Implement main DeepEP helpers** + +Use top-level imports from pinned SGLang main: + +```python +from sglang.srt.layers.moe.token_dispatcher.deepep import ( + DeepEPBuffer, + DeepEPDispatcher, +) +from sglang.srt.layers.moe.utils import get_deepep_mode, get_moe_a2a_backend +``` + +`set_deepep_graph_mode` resolves and applies: + +```python +mode = get_deepep_mode().resolve(is_extend_in_batch=False) +DeepEPBuffer.set_dispatch_mode(mode) +``` + +`bootstrap_deepep_buffer` checks `DeepEPBuffer._state().buffer`, scans +`runner.model_runner.model.modules()`, unwraps `_inners`, chooses the +low-latency implementation, calls `_get_buffer`, verifies the state now owns a +buffer, and applies graph mode. + +- [ ] **Step 4: Integrate with the backend** + +Call bootstrap before LOAD preallocation/linking and before entering SAVE +capture. On LOAD, call `cge.init_nvshmem_for_loaded_modules()` after bootstrap +and before `start_graph_builds`. `replay_session` calls +`set_deepep_graph_mode()` before yielding. + +- [ ] **Step 5: Commit/push and verify GREEN** + +Run: + +```bash +PYTHONPATH=python pytest tests/test_sglang_main_deepep.py -q +ruff check \ + python/foundry/integration/sglang/main_deepep.py \ + python/foundry/integration/sglang/main_backend.py \ + tests/test_sglang_main_deepep.py +``` + +Expected: all DeepEP CPU contracts pass. Commit message: +`feat: bootstrap DeepEP for SGLang graph restore`. + +--- + +### Task 6: Expose reproducible graph-mode recipe options + +**Files:** + +- Modify: `recipe/experimental/serve_qwen3-8b_sglang_tp.sh` +- Modify: `recipe/sglang/serve_qwen3-mini.sh` +- Modify: `tests/test_sglang_tp_recipe.py` + +**Interfaces:** + +- Consumes environment variables: + `SGLANG_CUDA_GRAPH_BACKEND_PREFILL`, + `SGLANG_CUDA_GRAPH_BS_PREFILL`, + `SGLANG_SPECULATIVE_ALGORITHM`, + `SGLANG_SPECULATIVE_NUM_STEPS`, + `SGLANG_SPECULATIVE_EAGLE_TOPK`, + `SGLANG_SPECULATIVE_NUM_DRAFT_TOKENS`, + `SGLANG_MOE_A2A_BACKEND`, and `SGLANG_DEEPEP_MODE`. + +- [ ] **Step 1: Write recipe argument tests** + +Allow `_run_recipe` to accept environment overrides. Add a test that supplies: + +```python +{ + "SGLANG_CUDA_GRAPH_BACKEND_PREFILL": "full", + "SGLANG_CUDA_GRAPH_BS_PREFILL": "16 64", + "SGLANG_SPECULATIVE_ALGORITHM": "NEXTN", + "SGLANG_SPECULATIVE_NUM_STEPS": "3", + "SGLANG_SPECULATIVE_EAGLE_TOPK": "1", + "SGLANG_SPECULATIVE_NUM_DRAFT_TOKENS": "4", + "SGLANG_MOE_A2A_BACKEND": "deepep", + "SGLANG_DEEPEP_MODE": "low_latency", +} +``` + +Assert every corresponding CLI flag/value is present and +`--disable-piecewise-cuda-graph` is absent. Keep a default-mode test that +explicitly selects prefill `disabled`. + +- [ ] **Step 2: Commit/push and verify RED** + +Run: + +```bash +PYTHONPATH=python pytest tests/test_sglang_tp_recipe.py -q +``` + +Expected: the script ignores new variables and always disables prefill. + +- [ ] **Step 3: Implement shell argument arrays** + +Set prefill backend default to `disabled`. Split +`SGLANG_CUDA_GRAPH_BS_PREFILL` on spaces and append each integer after +`--cuda-graph-bs-prefill`. Append speculative and DeepEP flags only when their +environment variables are non-empty. Replace the unconditional +`--disable-piecewise-cuda-graph` flag with explicit +`--cuda-graph-backend-prefill "$PREFILL_BACKEND"`. + +Apply the same explicit prefill backend/buckets to the mini recipe, defaulting +to disabled unless its environment selects full. + +- [ ] **Step 4: Commit/push and verify GREEN** + +Run: + +```bash +PYTHONPATH=python pytest tests/test_sglang_tp_recipe.py -q +bash -n recipe/experimental/serve_qwen3-8b_sglang_tp.sh +bash -n recipe/sglang/serve_qwen3-mini.sh +``` + +Expected: recipe tests pass and both scripts parse. Commit message: +`feat: configure SGLang persistent graph modes`. + +--- + +### Task 7: Full-prefill H100 smoke + +**Files:** + +- Modify: `tests/modal_sglang_main_smoke.py` +- Modify: `python/foundry/integration/sglang/log_scan.py` only if a real + non-error is misclassified. + +**Interfaces:** + +- Reads each rank's `sglang_graph_catalog.json`. +- Reports session phase/role/mode and graph count. + +- [ ] **Step 1: Add failing acceptance checks before backend fixes** + +Configure the mini recipe through the launch environment: + +```python +env["SGLANG_CUDA_GRAPH_BACKEND_PREFILL"] = "full" +env["SGLANG_CUDA_GRAPH_BS_PREFILL"] = "16 64" +``` + +Generate one short and one long prompt. Return exact token IDs using output +logprobs. Parse the catalog and require at least one `prefill` session and one +`decode` session. Require LOAD logs for both phases and no SAVE log on LOAD. + +- [ ] **Step 2: Commit/push, update the draft PR, and run RED on Modal** + +Run: + +```bash +FOUNDRY_REF=$(git rev-parse HEAD) \ +SGLANG_REF=9b853e6832e71a3058212df02a025232a453e146 \ +modal run --env rahul-dev tests/modal_sglang_main_smoke.py +``` + +Expected before completion: a focused prefill session/catalog/load failure, +not an environment setup failure. + +- [ ] **Step 3: Fix only evidence-backed prefill defects** + +Preserve SGLang's eager tail and static-buffer handling. Changes in Foundry are +limited to catalog/session sequencing, tensor output reconstruction, shared +pool use, and allocation-order parity demonstrated by the failing run. + +- [ ] **Step 4: Commit/push and rerun to GREEN** + +Use the same immutable commit and Modal command. Require: + +- exact baseline/SAVE/LOAD token IDs for both prompts; +- at least two prefill graph buckets in the archive; +- prefill and decode sessions restored; +- no native capture during LOAD; +- no CUDA, allocator, or log-scan errors. + +Commit message: `test: validate SGLang full prefill restore`. + +--- + +### Task 8: Qwen3.5 NEXTN speculative graph acceptance + +**Files:** + +- Modify: `tests/modal_sglang_tp.py` +- Modify: `python/foundry/integration/sglang/main_backend.py` only for a + reproduced speculative output/session defect. + +**Interfaces:** + +- Adds catalog session summaries and `/server_info` speculative acceptance + metrics to the TP report. + +- [ ] **Step 1: Add speculative observability and assertions** + +Extend `SERVE_ENV_NAMES` with the recipe variables from Task 6. Include +`sglang_graph_catalog.json` in archive fingerprints. Parse all catalog sessions +per rank and report tuples of `(phase, role, forward_mode, step, graph_count)`. + +When `SGLANG_SPECULATIVE_ALGORITHM` is set, query `/server_info` and require: + +- a target `TARGET_VERIFY` session; +- at least one draft session; +- at least one restored output schema containing `hidden_states`; +- `avg_spec_accept_length > 1.0`; +- exact seeded baseline/LOAD batched token IDs. + +- [ ] **Step 2: Commit/push and run Qwen3.5 NEXTN RED** + +Run: + +```bash +FOUNDRY_REF=$(git rev-parse HEAD) \ +SGLANG_REF=9b853e6832e71a3058212df02a025232a453e146 \ +SGLANG_MODEL=Qwen/Qwen3.5-9B \ +TP_SIZE=2 MODAL_GPU=H100:2 \ +SGLANG_SPECULATIVE_ALGORITHM=NEXTN \ +SGLANG_SPECULATIVE_NUM_STEPS=3 \ +SGLANG_SPECULATIVE_EAGLE_TOPK=1 \ +SGLANG_SPECULATIVE_NUM_DRAFT_TOKENS=4 \ +SGLANG_CUDA_GRAPH_BACKEND_PREFILL=full \ +SGLANG_CUDA_GRAPH_BS_PREFILL="16 64" \ +modal run --env rahul-dev tests/modal_sglang_tp.py +``` + +Classify the first failure as target-verify output, draft output, session +ordering, RNG, allocation order, or upstream SGLang behavior. + +- [ ] **Step 3: Apply focused fixes and rerun** + +Do not add algorithm-name branches to Foundry's backend. Extend only the +generic codec/catalog/backend invariant demonstrated by the failing artifact. + +- [ ] **Step 4: Commit/push and obtain GREEN** + +The run passes only with exact token IDs, nontrivial acceptance, all expected +target/draft/prefill sessions loaded on both ranks, reproducible SAVE archives, +and no graph recapture on LOAD. + +Commit message: `test: validate Qwen3.5 NEXTN graph restore`. + +--- + +### Task 9: SGLang-main DeepEP graph acceptance + +**Files:** + +- Modify: `tests/modal_sglang_tp.py` +- Create: `tests/modal_deepep_fabric.py` +- Reuse: `tests/test_deepep_fabric.py` + +**Interfaces:** + +- Modal image installs the DeepEP revision selected by pinned SGLang's + `scripts/ci/cuda/ci_install_deepep.sh`. +- Report requires DeepEP dispatch/combine kernels in saved graphs and restored + replay on every rank. + +- [ ] **Step 1: Add DeepEP image and runtime evidence** + +Conditionally install DeepEP when `SGLANG_MOE_A2A_BACKEND=deepep` using the +pinned SGLang CI script. Add log/catalog checks for: + +- Foundry DeepEP transport patch installation; +- buffer bootstrap before graph builds; +- NVSHMEM loaded-module initialization; +- per-rank graph sessions containing DeepEP kernels; +- no low-latency dispatch/combine errors. + +Keep the standalone semantic fabric test as the payload oracle; do not replace +it with log-only checks. + +- [ ] **Step 2: Add the standalone Modal wrapper** + +Create a Modal app using the same CUDA 13.0/PyTorch 2.11/Foundry/SGLang image +as `modal_sglang_tp.py`, install DeepEP through the pinned SGLang CI script, +request `gpu="H100:2"`, and execute: + +```python +subprocess.run( + [ + sys.executable, + "/foundry/tests/test_deepep_fabric.py", + "--run", + ], + check=True, + env={ + **os.environ, + "TEST_USE_FABRIC": "1", + "TEST_LOAD_API": "parallel", + }, +) +``` + +Mount a dedicated archive volume at `/data`, set the subprocess working +directory to `/data`, and delete its test archive after a passing run. + +- [ ] **Step 3: Commit/push and run standalone semantic verification** + +Run on two H100s: + +```bash +modal run --env rahul-dev tests/modal_deepep_fabric.py +``` + +Require the injective BF16 per-expert payload assertions to pass. + +- [ ] **Step 4: Run SGLang DeepEP RED** + +Use an upstream-supported FP8 MoE model: + +```bash +FOUNDRY_REF=$(git rev-parse HEAD) \ +SGLANG_REF=9b853e6832e71a3058212df02a025232a453e146 \ +SGLANG_MODEL=Qwen/Qwen3-30B-A3B-FP8 \ +TP_SIZE=2 MODAL_GPU=H100:2 \ +SGLANG_MOE_A2A_BACKEND=deepep \ +SGLANG_DEEPEP_MODE=low_latency \ +SGLANG_CUDA_GRAPH_BACKEND_PREFILL=full \ +SGLANG_CUDA_GRAPH_BS_PREFILL="16 64" \ +modal run --env rahul-dev tests/modal_sglang_tp.py +``` + +- [ ] **Step 5: Fix only reproduced DeepEP lifecycle defects** + +Use SAVE/LOAD buffer addresses, VMM offsets, NVSHMEM module counts, and +dispatch payload evidence. Do not weaken catalog mismatch or CUDA error checks. + +- [ ] **Step 6: Commit/push and obtain GREEN** + +Require semantic DeepEP payload correctness, exact SGLang token IDs, +reproducible per-rank archives/offsets, and restored prefill/decode sessions. + +Commit message: `test: validate SGLang DeepEP graph restore`. + +--- + +### Task 10: Combined mode and target TP4 regression + +**Files:** + +- Modify: `tests/modal_sglang_tp.py` only for generic combined assertions. +- Modify documentation files listed in File Structure after evidence exists. + +**Interfaces:** + +- No new runtime API. This task proves composition and updates the support + matrix. + +- [ ] **Step 1: Run combined Qwen3.5 MoE mode** + +Run Qwen3.5-35B-A3B with TP2/EP2, NEXTN, DeepEP low latency, and full prefill: + +```bash +FOUNDRY_REF=$(git rev-parse HEAD) \ +SGLANG_REF=9b853e6832e71a3058212df02a025232a453e146 \ +SGLANG_MODEL=Qwen/Qwen3.5-35B-A3B \ +SGLANG_QUANTIZATION=fp8 \ +TP_SIZE=2 MODAL_GPU=H100:2 \ +SGLANG_SPECULATIVE_ALGORITHM=NEXTN \ +SGLANG_SPECULATIVE_NUM_STEPS=3 \ +SGLANG_SPECULATIVE_EAGLE_TOPK=1 \ +SGLANG_SPECULATIVE_NUM_DRAFT_TOKENS=4 \ +SGLANG_MOE_A2A_BACKEND=deepep \ +SGLANG_DEEPEP_MODE=low_latency \ +SGLANG_CUDA_GRAPH_BACKEND_PREFILL=full \ +SGLANG_CUDA_GRAPH_BS_PREFILL="16 64" \ +modal run --env rahul-dev tests/modal_sglang_tp.py +``` + +Require the union of all Task 7-9 checks in one fresh LOAD process. + +- [ ] **Step 2: Re-run Qwen3.5-122B TP4 with full prefill** + +Use the existing GPTQ-Int4 environment plus: + +```bash +SGLANG_CUDA_GRAPH_BACKEND_PREFILL=full +SGLANG_CUDA_GRAPH_BS_PREFILL="16 64" +``` + +Require exact batched tokens, symmetric collectives on every rank, full-prefill +and decode session restore, reproducible archives/offsets, and no regression +from the previously validated TP4 path. + +- [ ] **Step 3: Run the complete CPU verification** + +Run: + +```bash +PYTHONPATH=python pytest \ + tests/test_sglang_output_codec.py \ + tests/test_sglang_graph_catalog.py \ + tests/test_sglang_main_deepep.py \ + tests/test_sglang_main_compat.py \ + tests/test_sglang_tp_recipe.py -q +ruff check python/foundry/integration/sglang tests/test_sglang_*.py +ruff format --check python/foundry/integration/sglang tests/test_sglang_*.py +pre-commit run --all-files +``` + +Expected: zero test, lint, format, or pre-commit failures. + +- [ ] **Step 4: Update evidence-backed documentation** + +Update support tables and launch examples with: + +- exact SGLang commit; +- exact model, topology, and graph modes; +- DeepEP mode/transport constraints; +- full-prefill-only persistence; +- which speculative algorithm received GPU validation; +- retained non-goals. + +Remove the stale statements that prefill, speculative decoding, and +SGLang-main DeepEP are categorically unsupported. + +- [ ] **Step 5: Commit/push and update the draft PR** + +Commit documentation and any final harness-only changes with: + +```bash +git add README.md RELEASE.md ROADMAP.md docs/sglang \ + recipe/sglang/README.md recipe/experimental/README.md tests +git commit -m "docs: record SGLang graph mode validation" +git push -u origin cursor/tp-consolidation-2e2c +``` + +Update the draft PR with the immutable Modal run identifiers and exact check +results. Do not mark it ready or merge it without an explicit user request. diff --git a/docs/superpowers/specs/2026-07-23-sglang-main-graph-modes-design.md b/docs/superpowers/specs/2026-07-23-sglang-main-graph-modes-design.md new file mode 100644 index 00000000..52707abc --- /dev/null +++ b/docs/superpowers/specs/2026-07-23-sglang-main-graph-modes-design.md @@ -0,0 +1,394 @@ +# SGLang Main DeepEP, Speculative Decode, and Prefill Graph Persistence + +## Status + +Approved on 2026-07-23. + +This design targets SGLang commit +`9b853e6832e71a3058212df02a025232a453e146` and extends the existing Foundry +SGLang-main plugin. The implementation remains out-of-tree: no SGLang source +patch is required at installation time. + +## Context + +The current SGLang-main integration persists one plain target-decode graph +whose output is exactly `LogitsProcessorOutput.next_token_logits`. The plugin +also: + +- disables prefill CUDA graphs; +- rejects every speculative algorithm; +- rejects DeepEP; and +- names and discovers graphs using only a decode batch size. + +Those assumptions do not hold for the requested workloads: + +- SGLang's target-verify, EAGLE/NEXTN draft, draft-extend, multi-layer MTP, and + frozen-KV MTP runners all resolve a decode `BaseCudaGraphBackend`, but return + different tensor trees. +- Full prefill capture resolves the same backend interface but keys graphs by + token bucket and returns transformer hidden states; SGLang runs the logits + tail eagerly after graph replay. +- DeepEP owns process-wide communication buffers and a dispatch mode that must + exist before capture or graph loading and must match at replay. +- Multiple runners can use the same shape size in one process. Batch size is + therefore not a unique archive identity. + +## Goals + +1. Persist every SGLang full-backend graph session used by the supported + process, including target decode/verify, speculative draft variants, and + full prefill. +2. Reconstruct each runner's output with the same Python structure and stable + tensor addresses that SGLang received after native capture. +3. Initialize DeepEP symmetrically on SAVE and LOAD and preserve its graph + dispatch mode. +4. Keep non-full SGLang graph backends operational without claiming that + Foundry persisted them. +5. Prove correctness with exact token outputs, archive/session evidence, and + restored graph replay on H100 GPUs. + +## Non-goals + +- Persisting breakable or tc-piecewise prefill graphs. +- Adding graph support where SGLang itself disables it, including DeepEP + `normal` mode. +- Pipeline parallelism, LoRA graph variants, PD multiplexing, diffusion + models, or runtime graph recapture. +- TensorRT-LLM validation. +- DP topology portability. Archives remain tied to the rank topology used + during SAVE. + +## Chosen Architecture + +Use one generalized Foundry implementation of SGLang's +`BaseCudaGraphBackend`. Do not create separate persistence backends for +DeepEP, speculative decoding, and prefill. + +The common backend is split into three responsibilities: + +1. An archive catalog identifies graph sessions and shape keys. +2. An output codec flattens supported Python tensor trees for Foundry and + reconstructs them after LOAD. +3. The backend owns capture/load sequencing, graph pools, DeepEP preparation, + and strong references to restored outputs. + +This boundary follows SGLang's runner/backend refactor: target, draft, and +prefill runners retain responsibility for static input buffers, attention +metadata, eager tails, and output slicing. Foundry only replaces full graph +capture and replay. + +## Archive Catalog + +Each rank directory gains `sglang_graph_catalog.json`. Graph JSON files remain +direct children of the rank directory so Foundry's graph manifest and fatbin +packaging continue to work. + +The catalog has a versioned schema: + +```json +{ + "version": 1, + "sessions": [ + { + "session_index": 0, + "runner": { + "class": "sglang.srt.model_executor.runner.decode_cuda_graph_runner.DecodeCudaGraphRunner", + "phase": "decode", + "role": "target", + "forward_mode": "TARGET_VERIFY", + "step": null + }, + "graphs": [ + { + "capture_index": 0, + "filename": "graph_0_decode_target_s0_k8.json", + "shape_key": { + "size": 8, + "stream_idx": null, + "variant_label": null + }, + "output_schema": { + "kind": "object", + "type": "LogitsProcessorOutput", + "attributes": {} + } + } + ] + } + ] +} +``` + +The exact output schema is generated by the codec and can contain nested +nodes. The example abbreviates it. + +Session identity is based on backend construction/capture order plus a runner +descriptor. The descriptor includes: + +- fully qualified runner class; +- phase (`decode` or `prefill`); +- target or draft role; +- capture forward mode; +- speculative step when the runner exposes one. + +On LOAD, each backend claims the next catalog session and verifies the complete +descriptor before loading any graph. Within a session, it verifies each full +`ShapeKey` and capture order. This makes equal batch sizes in target, draft, +draft-extend, and prefill sessions unambiguous. + +Catalog writes use a temporary file and `os.replace`, so every completed graph +append leaves a valid catalog. The graph filename remains human-readable, but +the catalog is authoritative. + +## Output Codec + +Foundry graph serialization accepts tensors or a flat list/tuple of tensors. +SGLang full backends can return: + +- a tensor; +- `LogitsProcessorOutput`, including hidden states and dynamically attached + speculative fields such as `topk_p` and `topk_index`; +- tuples returned by draft runners; +- lists of per-step outputs from a single multi-step graph; and +- nested string-keyed dictionaries used by supported SGLang output wrappers. + +The codec performs a deterministic depth-first traversal and emits: + +1. a flat tensor list passed to `CUDAGraph.save`; and +2. a JSON schema that describes containers, object attributes, constants, and + tensor indices. + +Repeated references to the same tensor use one tensor index, preserving +aliasing when the tree is reconstructed. Supported leaf constants are `None`, +booleans, integers, floats, and strings. Supported containers are lists, +tuples, and string-keyed dictionaries. + +Object reconstruction is allowlisted to SGLang output types needed by the +backend contract. It does not import arbitrary classes named by archive data. +For `LogitsProcessorOutput`, all instance attributes are encoded, not only +dataclass fields, so speculative runners retain fields attached after model +forward. Restored tensors and their enclosing output objects are held strongly +by the backend for the server lifetime. + +An unsupported output node is a capture-time integration error with its full +tree path in the message. Silently dropping a field or substituting eager +output is unsafe because SGLang may consume that field after replay. + +## Backend Resolution and Plugin Behavior + +The Foundry resolver patch covers both: + +- `resolve_decode_backend`; and +- `resolve_prefill_backend`. + +When Foundry is active and a phase resolves to `full`, the resolver returns the +generalized Foundry backend. For `breakable`, `tc_piecewise`, or `disabled`, it +returns SGLang's original resolution result. Foundry logs that such a session +is not persisted; it does not reject the server configuration. + +The plugin no longer disables prefill graphs and no longer rejects speculative +algorithms or DeepEP. It also does not silently change an explicitly selected +prefill backend. Foundry recipes and validation harnesses select +`--cuda-graph-backend-prefill full` when persistent prefill graphs are +required. + +The existing guards for pipeline parallelism, LoRA, PD multiplexing, +return-hidden-states recapture, and diffusion models remain until those +separate contracts are implemented. + +## Capture and Load Lifecycle + +For each full-backend session: + +### SAVE + +1. Build or reuse SGLang's process-global CUDA graph memory pool. +2. Prepare DeepEP if active. +3. Let the SGLang runner prepare static inputs and attention metadata. +4. Capture each shape through Foundry in SGLang's requested order. +5. Flatten and save the output, append its catalog record, and retain the live + graph/output. +6. Refresh Foundry's graph manifest, fatbin bundle, and final allocation + offset after the session. Repeating finalization is intentional because + later speculative or prefill runners may add sessions. + +### LOAD + +1. Build or reuse the same process-global graph memory pool. +2. Prepare DeepEP at the same lifecycle point as SAVE. +3. Claim and validate the next archive session. +4. Preallocate to the saved final offset and initialize loaded NVSHMEM modules + before linking graphs that contain DeepEP nodes. +5. Start graph builds for this session with the shared pool. +6. Let SGLang repeat its normal per-shape input and metadata preparation. +7. Finish each graph load in catalog order, reconstruct its output tree, and + retain strong references. +8. Reject missing, extra, reordered, or mismatched sessions before serving. + +The archive keeps complete dependencies (`strip_dependencies=False`). +Distributed graph replay is not switched to Foundry's shared-executor +optimization in this change. + +## DeepEP + +Supported DeepEP graph operation follows SGLang's own low-latency CUDA-graph +path. `deepep_mode=auto` resolves graph capture to low latency; explicit +low-latency mode is also supported. SGLang disables CUDA graphs for normal +mode, so Foundry leaves that configuration to SGLang rather than failing it. + +The existing bootstrap helper is updated for SGLang main's resource-backed +`DeepEPBuffer._state().buffer` API. It locates a real DeepEP dispatcher and +creates the process-wide buffer outside stream capture on every rank. + +Before each Foundry session and replay session, the backend applies the mode +that SGLang resolves with `is_extend_in_batch=False`, matching SGLang's graph +runner capture state. This fills the prefill-runner gap while remaining +consistent with `DeepEPCudaGraphRunnerAdapter` on decode and speculative +runners. + +The existing Foundry transport hook remains responsible for fabric versus NVL +IPC selection. SAVE and LOAD must use the same transport setting, EP group, +rank topology, and DeepEP buffer sizing. + +## Speculative Decoding + +All current CUDA speculative runners on the pinned SGLang main route their +full graph capture through `resolve_decode_backend`. The generalized resolver +therefore covers: + +- target `TARGET_VERIFY`; +- EAGLE/NEXTN draft decode; +- draft extend; +- multi-layer per-step draft extend; +- the single-graph multi-step variant; and +- frozen-KV MTP. + +The catalog separates these sessions by runner descriptor and step. The output +codec handles logits, hidden states, top-k tensors, tuples, and multi-step +lists without adding runner-specific persistence code. + +The first GPU correctness target is Qwen3.5 NEXTN because it exercises target +verify, draft outputs, hidden states, and draft-extend graphs. Other +speculative algorithms are not rejected, but are not labeled GPU-validated +until their own acceptance case passes. + +On the pinned SGLang revision, `capture_prefill_graph` intentionally skips an +EAGLE/NEXTN target when the prefill backend is not breakable, and +`PrefillCudaGraphRunner` selects the required FULL/LAST hidden-state capture +modes only through its breakable-EAGLE branch. Foundry must preserve those +hidden-state semantics while supplying its monolithic full backend. The plugin +therefore wraps prefill-runner construction for EAGLE/NEXTN: construction sees +the breakable-EAGLE mode long enough to select target FULL or draft LAST, +backend resolution still returns Foundry's full-style backend, and the runner +is restored to full-prefill runtime semantics before construction returns. +This compatibility seam is pinned and covered by CPU contracts; it does not +modify SGLang source files. + +Foundry's graph-level RNG/generator restoration remains the source of truth +for stochastic graph nodes. Validation uses explicit request seeds and also +checks that repeated replay advances RNG rather than freezing one sampled +result. + +## Prefill Graphs + +Foundry persists only SGLang's full prefill backend. This backend captures the +transformer body per token bucket and returns hidden states. SGLang continues +to: + +- refresh attention metadata outside the graph; +- populate static prefill buffers; +- bucket and pad token counts; +- run the LM head, logits processor, and logprob tail eagerly; and +- trim outputs to the raw token/request counts. + +Foundry does not duplicate those responsibilities. It saves and restores the +body's hidden-state tensor and returns it to SGLang's existing eager-tail +path. + +Hybrid Qwen3.5 memory configuration remains deterministic through the saved +memory-pool configuration and `max_mamba_cache_size`. Full-prefill acceptance +must include a hybrid Qwen3.5 model, prefix-cache reuse, and more than one +prefill token bucket. + +SGLang currently labels full prefill experimental. The Foundry implementation +is pinned to the tested SGLang commit and records upstream regressions as +compatibility failures rather than hiding them. + +## Error Handling + +Foundry fails LOAD before serving when continuing would be unsafe: + +- catalog version is unknown; +- runner/session descriptors differ; +- graph shape keys or order differ; +- graph or output-schema files are missing; +- output tensor count differs from the schema; +- DeepEP rank topology or initialization is inconsistent; or +- not every graph in the claimed session was consumed. + +These are archive compatibility errors, not unsupported-mode policy. SGLang +configurations whose backend is not persisted continue through the upstream +backend. + +## Tests + +### CPU-only contracts + +- Catalog round-trip, atomic update, duplicate-size sessions, and mismatch + diagnostics. +- Output-codec round-trip for tensors, aliases, nested tuples/lists/dicts, + `LogitsProcessorOutput` hidden/dynamic fields, and a multi-step output list. +- Rejection of unsupported leaves and malformed/untrusted object schemas. +- Decode and prefill resolver routing: full uses Foundry; other backends use + SGLang. +- Plugin validation accepts speculative decoding, DeepEP, and full prefill. +- DeepEP bootstrap uses SGLang main's resource-backed state and is idempotent. + +Every production behavior is introduced with a failing test first. + +### H100 acceptance + +1. **Full prefill smoke:** a small SGLang-main model, at least two prefill + buckets, exact baseline/LOAD token IDs, and restored prefill replay evidence. +2. **Qwen3.5 NEXTN:** Qwen3.5 NEXTN on a configuration supported upstream, + exact seeded baseline/LOAD token IDs, target and draft session evidence, + hidden-state output restoration, and nontrivial speculative acceptance. +3. **DeepEP:** a two-GPU MoE EP configuration in low-latency/auto mode, + semantic dispatch/combine payload checks, exact baseline/LOAD tokens, and + per-rank restored DeepEP graph evidence. +4. **Combined mode:** Qwen3.5 MoE with NEXTN, DeepEP, and full prefill enabled, + proving all catalog sessions restore in one fresh process. +5. **Target deployment:** Qwen3.5-122B-A10B-GPTQ-Int4 TP4 with persistent + decode and full-prefill graphs, exact batched token IDs, reproducible + per-rank archives/allocation offsets, and restored replay on every rank. + +Each acceptance run compares a native SGLang baseline, Foundry SAVE, a second +SAVE for reproducibility, and Foundry LOAD in a fresh process. Passing requires +no unexpected eager fallback for the exercised buckets, no CUDA/NCCL/DeepEP +errors, and the expected session/graph counts on every rank. + +## Upstream Basis + +The design deliberately follows active SGLang interfaces and known changes: + +- PR #23906: CUDA graph runner/backend refactor. +- PR #27988: experimental full prefill CUDA graphs. +- PR #28973: one graph memory pool shared by prefill and decode. +- PR #28870: speculative hidden-state plumbing in prefill. +- PR #11666: explicit DeepEP mode switching for CUDA graphs. +- PR #30853: draft-extend CUDA graph support expansion. + +Open upstream work, including multi-request prefill legality, multimodal +prefill input handling, and speculative replay WAR-event fixes, is treated as +compatibility risk. Foundry tests pin the SGLang commit so a moving upstream +change cannot silently alter archive semantics. + +## Delivery Order + +1. Add catalog and output-codec CPU contracts. +2. Generalize the full backend and archive sequencing. +3. Route full prefill and remove the three plugin rejections. +4. Update and integrate DeepEP bootstrap/mode handling. +5. Add focused Modal smoke tests, then the combined and TP4 acceptance runs. +6. Update user documentation only with configurations proven by fresh GPU + evidence. 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/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..3c97a943 100644 --- a/python/foundry/__init__.py +++ b/python/foundry/__init__.py @@ -1,5 +1,10 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the Foundry project +# ruff: noqa: I001 +# The ops wildcard must come FIRST: .graph's CUDAGraph (and the +# allocation_region helpers) intentionally override the raw pybind names. +from .ops import * + from .allocation_region import ( allocation_region, free_preallocated_region, @@ -14,7 +19,6 @@ graph, save_graph_manifest, ) -from .ops import * # Re-exports. Listed here so ruff's --fix doesn't strip them as unused. # (We also configure per-file-ignores for F401 in pyproject.toml as a backstop.) diff --git a/python/foundry/graph.py b/python/foundry/graph.py index 427d1599..64675381 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 ) @@ -167,7 +170,11 @@ def __exit__(self, *args) -> None: self.stream_ctx.__exit__(*args) -def save_graph_manifest(archive_dir: str) -> None: +def save_graph_manifest( + archive_dir: str, + *, + strip_dependencies: bool = True, +) -> None: """Write graph_manifest.json with topology groups and template assignments. Reads topology_key from each saved graph JSON, groups graphs by topology, @@ -213,22 +220,21 @@ def save_graph_manifest(archive_dir: str) -> None: with open(manifest_path, "w") as f: json.dump(manifest, f, indent=2) - # Strip dependencies from on-demand graph files. - # On-demand graphs only need nodes + common_kernel_node_attrs for - # prepare_on_demand_graph; dependencies are only used by - # build_graph_from_parsed (template builds). num_stripped = 0 - for filename in graph_files: - if filename in templates: - continue - json_path = os.path.join(archive_dir, filename) - with open(json_path) as f: - data = json.load(f) - if "dependencies" in data: - del data["dependencies"] - with open(json_path, "w") as f: - json.dump(data, f) - num_stripped += 1 + if strip_dependencies: + # Shared-executor members only need node parameters. Independent graph + # loading, however, needs the original edges to rebuild each graph. + for filename in graph_files: + if filename in templates: + continue + json_path = os.path.join(archive_dir, filename) + with open(json_path) as f: + data = json.load(f) + if "dependencies" in data: + del data["dependencies"] + with open(json_path, "w") as f: + json.dump(data, f) + num_stripped += 1 num_templates = len(groups) num_on_demand = sum(len(g["members"]) - 1 for g in groups) diff --git a/python/foundry/integration/sglang/archive.py b/python/foundry/integration/sglang/archive.py new file mode 100644 index 00000000..3620a59f --- /dev/null +++ b/python/foundry/integration/sglang/archive.py @@ -0,0 +1,104 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the Foundry project +"""Versioned per-rank catalog for SGLang CUDA graph sessions.""" + +from __future__ import annotations + +import json +import os +from pathlib import Path +from typing import Any + +CATALOG_FILENAME = "sglang_graph_catalog.json" +CATALOG_VERSION = 1 + + +def shape_key_record(shape_key: Any) -> dict[str, int | str | None]: + return { + "size": int(shape_key.size), + "stream_idx": getattr(shape_key, "stream_idx", None), + "variant_label": getattr(shape_key, "variant_label", None), + } + + +def runner_descriptor(runner: Any, phase: str) -> dict[str, Any]: + return { + "class": f"{type(runner).__module__}.{type(runner).__qualname__}", + "phase": phase, + "role": ("draft" if getattr(runner.model_runner, "is_draft_worker", False) else "target"), + "forward_mode": getattr( + getattr(runner, "capture_forward_mode", None), + "name", + None, + ), + "step": getattr(runner, "step", None), + } + + +class GraphCatalog: + def __init__(self, directory: Path | str) -> None: + self._path = Path(directory) / CATALOG_FILENAME + self._data: dict[str, Any] = {} + self._reload() + + def _reload(self) -> None: + if self._path.exists(): + payload = json.loads(self._path.read_text()) + version = payload.get("version") + if version != CATALOG_VERSION: + raise RuntimeError(f"unknown catalog version: {version!r}") + self._data = payload + else: + self._data = {"version": CATALOG_VERSION, "sessions": []} + + @property + def session_count(self) -> int: + self._reload() + return len(self._data["sessions"]) + + def _write(self) -> None: + tmp_path = Path(f"{self._path}.tmp.{os.getpid()}") + with tmp_path.open("w") as handle: + json.dump(self._data, handle, indent=2) + handle.flush() + os.fsync(handle.fileno()) + os.replace(tmp_path, self._path) + + def open_save_session(self, index: int, descriptor: dict[str, Any]) -> dict[str, Any]: + self._reload() + expected_index = len(self._data["sessions"]) + if index != expected_index: + raise RuntimeError(f"expected session index {expected_index}, got {index}") + session = { + "session_index": index, + "runner": descriptor, + "graphs": [], + } + self._data["sessions"].append(session) + self._write() + return session + + def claim_load_session(self, index: int, descriptor: dict[str, Any]) -> dict[str, Any]: + self._reload() + sessions = self._data["sessions"] + if index < 0 or index >= len(sessions): + raise RuntimeError(f"session {index} not found in catalog") + session = sessions[index] + if session["runner"] != descriptor: + raise RuntimeError(f"session {index} descriptor mismatch") + return session + + def append_graph(self, index: int, graph_record: dict[str, Any]) -> None: + self._reload() + sessions = self._data["sessions"] + if index < 0 or index >= len(sessions): + raise RuntimeError(f"session {index} not found in catalog") + capture_index = graph_record["capture_index"] + for session in sessions: + for graph in session["graphs"]: + if capture_index <= graph["capture_index"]: + raise RuntimeError( + f"capture_index {capture_index} must exceed all existing graph indices" + ) + sessions[index]["graphs"].append(graph_record) + self._write() diff --git a/python/foundry/integration/sglang/graph_ops.py b/python/foundry/integration/sglang/graph_ops.py index 18ab1f0e..0320c231 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(): diff --git a/python/foundry/integration/sglang/hooks.py b/python/foundry/integration/sglang/hooks.py index 22cd7cfa..1ff85555 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: @@ -64,6 +82,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 +114,76 @@ 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. + + When supported, 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``. Older DeepEP revisions have no + ``use_fabric`` switch and retain their native NVL/IPC transport. + """ + 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 + transport_reported = False + + @functools.wraps(original_init) + def patched(self, group, num_nvl_bytes=0, num_rdma_bytes=0, *args, **kwargs): + nonlocal transport_reported + if supports_use_fabric and get_graph_extension_mode() != CUDAGraphExtensionMode.NONE: + if os.environ.get("FOUNDRY_DEEPEP_NVL_IPC", "0") == "1": + kwargs["use_fabric"] = False + else: + num_nvl_bytes = 0 + kwargs["use_fabric"] = True + result = original_init( + self, + group, + num_nvl_bytes, + num_rdma_bytes, + *args, + **kwargs, + ) + if not transport_reported: + if not supports_use_fabric: + logger.info( + "[Foundry] SGLang DeepEP Buffer lacks use_fabric; preserving its " + "NVL/IPC transport and allocation sizes" + ) + logger.info("[Foundry] SGLang DeepEP transport patch installed") + transport_reported = True + return result + + Buffer.__init__ = patched + _DEEPEP_PATCHED = True + + def _patch_init_torch_distributed() -> None: from sglang.srt.model_executor import model_runner as mr diff --git a/python/foundry/integration/sglang/hooks_main.py b/python/foundry/integration/sglang/hooks_main.py new file mode 100644 index 00000000..5ef417da --- /dev/null +++ b/python/foundry/integration/sglang/hooks_main.py @@ -0,0 +1,609 @@ +# 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 copy +import functools +import logging +from contextlib import contextmanager +from contextvars import ContextVar +from dataclasses import asdict +from typing import Any + +import torch +from sglang.kernels.ops.attention.fla.chunk import CHUNK_SIZE as GDN_CHUNK_SIZE +from sglang.kernels.ops.attention.fla.index import ( + prepare_chunk_indices, + prepare_chunk_offsets, + prepare_lens, +) +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.layers.attention.hybrid_linear_attn_backend import ( + HybridLinearAttnBackend, + MambaAttnBackendBase, +) +from sglang.srt.layers.attention.linear.gdn_backend import GDNAttnBackend +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.cuda_graph_config import Backend +from sglang.srt.model_executor.model_runner_components import cuda_graph_setup +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 import prefill_cuda_graph_runner as prefill_runner +from sglang.srt.model_executor.runner_backend import utils as backend_utils +from sglang.srt.speculative import ( + eagle_draft_cuda_graph_runner, + eagle_draft_extend_cuda_graph_runner, + frozen_kv_mtp_cuda_graph_runner, + multi_layer_eagle_draft_extend_cuda_graph_runner, +) +from sglang.srt.utils.common import next_power_of_2 + +from foundry import ops as cge +from foundry.integration.sglang import runtime as rt +from foundry.integration.sglang.config import ( + CUDAGraphExtensionMode, + get_graph_extension_mode, +) +from foundry.integration.sglang.main_backend import FoundryMainCudaGraphBackend + +logger = logging.getLogger(__name__) + +_torch_symm_mem_communicators: dict[tuple[int, ...], Any] = {} +_full_prefill_eagle_construction: ContextVar[bool] = ContextVar( + "foundry_full_prefill_eagle_construction", + default=False, +) +_eager_prefill_fallback_warnings: set[str] = set() + + +def install_hooks_main() -> None: + _patch_init_torch_distributed() + _patch_memory_pool() + _patch_torch_symm_mem() + _patch_multimem_all_gather() + _patch_backend_resolver() + _patch_prefill_graph_capture() + _patch_cuda_graph_capture_order() + _patch_hybrid_linear_prefill_metadata() + _patch_spawn_sites() + + +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") + if state.max_mamba_cache_size is not None: + self.server_args.override( + "foundry.warmup_state", + max_mamba_cache_size=state.max_mamba_cache_size, + ) + 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)) + state.max_mamba_cache_size = self.server_args.max_mamba_cache_size + 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) + + group = kwargs.get("group", args[0] if args else None) + if group is None: + raise RuntimeError("Foundry could not identify the torch symmetric-memory group") + world_size = torch.distributed.get_world_size(group) + signal_pad_bytes = triton_symm_mem_ag._MAX_BLOCKS * world_size * 4 + symm_mem = torch_symm_mem_module.torch_symm_mem + symm_mem.set_signal_pad_size(max(symm_mem.get_signal_pad_size(), signal_pad_bytes)) + + # 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 not None: + group_ranks = tuple(torch.distributed.get_process_group_ranks(self.group)) + _torch_symm_mem_communicators[group_ranks] = self + logger.info( + "[Foundry] Torch symmetric-memory buffer address=0x%x", + self.buffer.data_ptr(), + ) + return result + + torch_symm_mem_module.TorchSymmMemCommunicator.__init__ = patched + + +def _patch_multimem_all_gather() -> None: + original = triton_symm_mem_ag.MultimemAllGatherer.__call__ + + @functools.wraps(original) + def patched(self, x): + if get_graph_extension_mode() == CUDAGraphExtensionMode.NONE: + return original(self, x) + + if len(_torch_symm_mem_communicators) != 1: + raise RuntimeError( + "Foundry requires exactly one torch symmetric-memory TP communicator, " + f"found {len(_torch_symm_mem_communicators)}" + ) + communicator = next(iter(_torch_symm_mem_communicators.values())) + if x.ndim != 2 or x.dtype != torch.bfloat16 or not x.is_contiguous(): + raise RuntimeError( + "Foundry multimem all-gather requires contiguous rank-2 bfloat16 input" + ) + world_size = communicator.world_size + token_count, local_hidden_size = x.shape + required_numel = world_size * x.numel() + if required_numel > communicator.buffer.numel(): + raise RuntimeError( + "Foundry torch symmetric-memory buffer is too small for " + f"multimem all-gather: need {required_numel} bf16 elements, " + f"have {communicator.buffer.numel()}" + ) + comm_buff = communicator.buffer[:required_numel].view( + world_size * token_count, + local_hidden_size, + ) + torch.ops.symm_mem.multimem_all_gather_out( + x, + communicator.group.group_name, + comm_buff, + ) + if not getattr(self, "_foundry_observed", False): + logger.info( + "[Foundry] Multimem all-gather buffer address=0x%x", + comm_buff.data_ptr(), + ) + self._foundry_observed = True + return ( + comm_buff.view(world_size, token_count, local_hidden_size) + .movedim(0, 1) + .reshape(token_count, world_size * local_hidden_size) + ) + + triton_symm_mem_ag.MultimemAllGatherer.__call__ = patched + + +def _patch_backend_resolver() -> None: + original_decode = backend_utils.resolve_decode_backend + original_prefill = backend_utils.resolve_prefill_backend + + @functools.wraps(original_decode) + def patched_decode(cuda_graph_runner): + return _foundry_backend_or_original( + cuda_graph_runner, + phase="decode", + original=original_decode, + ) + + @functools.wraps(original_prefill) + def patched_prefill(cuda_graph_runner): + return _foundry_backend_or_original( + cuda_graph_runner, + phase="prefill", + original=original_prefill, + ) + + backend_utils.resolve_decode_backend = patched_decode + for runner_module in ( + decode_runner, + eagle_draft_cuda_graph_runner, + eagle_draft_extend_cuda_graph_runner, + frozen_kv_mtp_cuda_graph_runner, + multi_layer_eagle_draft_extend_cuda_graph_runner, + ): + runner_module.resolve_decode_backend = patched_decode + backend_utils.resolve_prefill_backend = patched_prefill + prefill_runner.resolve_prefill_backend = patched_prefill + + +def _foundry_backend_or_original(cuda_graph_runner, *, phase: str, original): + if get_graph_extension_mode() == CUDAGraphExtensionMode.NONE: + return original(cuda_graph_runner) + if phase == "prefill" and _full_prefill_eagle_construction.get(): + return FoundryMainCudaGraphBackend(cuda_graph_runner, phase=phase) + + phase_config = getattr( + cuda_graph_runner.model_runner.server_args.cuda_graph_config, + phase, + ) + if phase_config.backend != Backend.FULL: + logger.info( + "[Foundry] SGLang %s backend=%s is not persisted", + phase, + phase_config.backend, + ) + return original(cuda_graph_runner) + return FoundryMainCudaGraphBackend(cuda_graph_runner, phase=phase) + + +def _full_prefill_eager_fallback_reason(model_runner) -> str | None: + server_args = model_runner.server_args + if ( + get_graph_extension_mode() == CUDAGraphExtensionMode.NONE + or server_args.cuda_graph_config.prefill.backend != Backend.FULL + ): + return None + if getattr(server_args, "forward_hooks", None): + return "server_args.forward_hooks is nonempty" + + attn_backend = getattr(model_runner, "attn_backend", None) + if not isinstance(attn_backend, HybridLinearAttnBackend): + return None + linear_backend = attn_backend.linear_attn_backend + if isinstance(linear_backend, GDNAttnBackend): + return None + backend_name = f"{type(linear_backend).__module__}.{type(linear_backend).__qualname__}" + return f"hybrid linear attention backend {backend_name} is not GDNAttnBackend" + + +def _warn_eager_prefill_fallback(reason: str) -> None: + if reason in _eager_prefill_fallback_warnings: + return + _eager_prefill_fallback_warnings.add(reason) + logger.warning( + "[Foundry] Full prefill persistence is unavailable because %s; " + "using SGLang eager prefill for this phase while preserving decode persistence", + reason, + ) + + +@contextmanager +def _temporary_prefill_backend(model_runner, backend): + prefill_config = model_runner.server_args.cuda_graph_config.prefill + original_backend = prefill_config.backend + prefill_config.backend = backend + try: + yield + finally: + prefill_config.backend = original_backend + + +def _patch_prefill_graph_capture() -> None: + original = cuda_graph_setup.capture_prefill_graph + + @functools.wraps(original) + def patched(*, model_runner, eager_runner, force_for_draft_worker=False): + prefill_config = model_runner.server_args.cuda_graph_config.prefill + fallback_reason = _full_prefill_eager_fallback_reason(model_runner) + if fallback_reason is not None: + _warn_eager_prefill_fallback(fallback_reason) + with _temporary_prefill_backend(model_runner, Backend.DISABLED): + return original( + model_runner=model_runner, + eager_runner=eager_runner, + force_for_draft_worker=force_for_draft_worker, + ) + + use_compatibility_override = ( + get_graph_extension_mode() != CUDAGraphExtensionMode.NONE + and prefill_config.backend == Backend.FULL + and model_runner.spec_algorithm.is_eagle() + ) + if not use_compatibility_override: + return original( + model_runner=model_runner, + eager_runner=eager_runner, + force_for_draft_worker=force_for_draft_worker, + ) + + token = _full_prefill_eagle_construction.set(True) + try: + with _temporary_prefill_backend(model_runner, Backend.BREAKABLE): + runner = original( + model_runner=model_runner, + eager_runner=eager_runner, + force_for_draft_worker=force_for_draft_worker, + ) + finally: + _full_prefill_eagle_construction.reset(token) + + if runner is not None and hasattr(runner, "prefill_backend_name"): + runner.prefill_backend_name = Backend.FULL + return runner + + cuda_graph_setup.capture_prefill_graph = patched + mr.capture_prefill_graph = patched + + +def _patch_cuda_graph_capture_order() -> None: + original = cuda_graph_setup.capture_cuda_graphs + + @functools.wraps(original) + def patched(*, model_runner, capture_decode_cuda_graph=True): + prefill_config = model_runner.server_args.cuda_graph_config.prefill + fallback_reason = _full_prefill_eager_fallback_reason(model_runner) + if fallback_reason is not None: + _warn_eager_prefill_fallback(fallback_reason) + with _temporary_prefill_backend(model_runner, Backend.DISABLED): + return original( + model_runner=model_runner, + capture_decode_cuda_graph=capture_decode_cuda_graph, + ) + + defer_prefill = ( + capture_decode_cuda_graph + and get_graph_extension_mode() != CUDAGraphExtensionMode.NONE + and prefill_config.backend == Backend.FULL + ) + if not defer_prefill: + return original( + model_runner=model_runner, + capture_decode_cuda_graph=capture_decode_cuda_graph, + ) + + with _temporary_prefill_backend(model_runner, Backend.DISABLED): + capture = original( + model_runner=model_runner, + capture_decode_cuda_graph=capture_decode_cuda_graph, + ) + + _preallocate_mamba_prefill_metadata_buffers(model_runner) + prefill_runner = cuda_graph_setup.capture_prefill_graph( + model_runner=model_runner, + eager_runner=capture.eager_runner, + ) + return type(capture)( + eager_runner=capture.eager_runner, + prefill_runner=prefill_runner, + decode=capture.decode, + ) + + cuda_graph_setup.capture_cuda_graphs = patched + mr.capture_cuda_graphs = patched + + +def _preallocate_mamba_prefill_metadata_buffers(model_runner) -> None: + if model_runner.spec_algorithm.is_eagle(): + return + attn_backend = getattr(model_runner, "attn_backend", None) + if not isinstance(attn_backend, HybridLinearAttnBackend): + return + linear_backend = attn_backend.linear_attn_backend + if not isinstance(linear_backend, GDNAttnBackend): + return + + prefill_config = model_runner.server_args.cuda_graph_config.prefill + capture_tokens = prefill_config.bs + if not capture_tokens: + return + max_req = prefill_config.full_prefill_max_req + if max_req is None: + max_req = max(model_runner.server_args.chunked_prefill_size // 512, 1) + batch_size = min(max_req, model_runner.req_to_token_pool.size) + + buffers = getattr(linear_backend, "_foundry_prefill_metadata_buffers", None) + if buffers is None: + buffers = {} + linear_backend._foundry_prefill_metadata_buffers = buffers + for num_tokens in sorted(capture_tokens, reverse=True): + key = (batch_size, num_tokens) + if key in buffers: + continue + buffers[key] = ( + torch.empty( + batch_size + 1, + dtype=torch.int32, + device=getattr(linear_backend, "device", None), + ), + torch.empty( + batch_size, + dtype=torch.int32, + device=getattr(linear_backend, "device", None), + ), + None, + ) + linear_backend._foundry_prefill_metadata_preallocated = True + + +def _prime_gdn_prefill_metadata(query_start_loc, num_tokens: int) -> None: + output_chunk_size = min( + GDN_CHUNK_SIZE, + max(16, next_power_of_2(num_tokens)), + ) + prepare_lens(query_start_loc) + prepare_chunk_indices(query_start_loc, GDN_CHUNK_SIZE) + if output_chunk_size != GDN_CHUNK_SIZE: + prepare_chunk_indices(query_start_loc, output_chunk_size) + prepare_chunk_offsets(query_start_loc, GDN_CHUNK_SIZE) + + +def _patch_hybrid_linear_prefill_metadata() -> None: + original = MambaAttnBackendBase.init_forward_metadata_out_graph + + @functools.wraps(original) + def patched(self, forward_batch, in_capture=False): + forward_mode = forward_batch.forward_mode + is_plain_extend = ( + forward_mode.is_extend(include_draft_extend_v2=True) + and not forward_mode.is_target_verify() + and not forward_mode.is_draft_extend_v2() + ) + if ( + get_graph_extension_mode() == CUDAGraphExtensionMode.NONE + or not is_plain_extend + or not isinstance(self, GDNAttnBackend) + ): + return original(self, forward_batch, in_capture=in_capture) + + metadata_batch = forward_batch + batch_size = forward_batch.batch_size + extend_start_loc = getattr(forward_batch, "extend_start_loc", None) + if extend_start_loc is not None and extend_start_loc.shape[0] < batch_size: + metadata_batch = copy.copy(forward_batch) + padded_extend_start_loc = torch.empty_like(forward_batch.extend_seq_lens[:batch_size]) + padded_extend_start_loc[0] = 0 + torch.cumsum( + forward_batch.extend_seq_lens[: batch_size - 1], + dim=0, + out=padded_extend_start_loc[1:], + ) + metadata_batch.extend_start_loc = padded_extend_start_loc + + metadata = self._forward_metadata(metadata_batch) + dynamic_tracking_fields = ( + "track_conv_indices", + "track_ssm_h_src", + "track_ssm_h_dst", + "track_ssm_final_src", + "track_ssm_final_dst", + ) + if any(getattr(metadata, name) is not None for name in dynamic_tracking_fields): + raise RuntimeError( + "Foundry full prefill does not support dynamic Mamba tracking metadata" + ) + + buffers = getattr(self, "_foundry_prefill_metadata_buffers", None) + if buffers is None: + buffers = {} + self._foundry_prefill_metadata_buffers = buffers + num_tokens = int(forward_batch.input_ids.shape[0]) + use_preallocated_gdn_metadata = ( + in_capture + and isinstance(self, GDNAttnBackend) + and getattr(self, "_foundry_prefill_metadata_preallocated", False) + ) + if in_capture: + key = (batch_size, num_tokens) + shape_was_preallocated = key in buffers + if key not in buffers: + buffers[key] = ( + torch.empty_like(metadata.query_start_loc), + torch.empty_like(metadata.mamba_cache_indices), + ( + torch.empty_like(metadata.mamba_track_indices) + if metadata.mamba_track_indices is not None + else None + ), + ) + query_buffer, state_buffer, track_buffer = buffers[key] + logger.info( + "[Foundry] Allocated Mamba prefill metadata batch_size=%d " + "num_tokens=%d query=0x%x state=0x%x track=%s", + batch_size, + num_tokens, + query_buffer.data_ptr(), + state_buffer.data_ptr(), + f"0x{track_buffer.data_ptr():x}" if track_buffer is not None else "none", + ) + else: + candidates = [key for key in buffers if key[0] == batch_size and key[1] >= num_tokens] + if not candidates: + raise RuntimeError( + "Foundry has no captured Mamba prefill metadata buffer for " + f"batch_size={batch_size}, num_tokens={num_tokens}" + ) + key = min(candidates, key=lambda candidate: candidate[1]) + + static_query_start_loc, static_state_indices, static_track_indices = buffers[key] + static_query_start_loc.copy_(metadata.query_start_loc) + static_state_indices.copy_(metadata.mamba_cache_indices) + static_state_indices.masked_fill_( + forward_batch.extend_seq_lens[:batch_size] == 0, + self.pad_slot_id, + ) + metadata.query_start_loc = static_query_start_loc + metadata.mamba_cache_indices = static_state_indices + if metadata.mamba_track_indices is not None: + if static_track_indices is None: + raise RuntimeError("Foundry Mamba prefill tracking metadata changed") + static_track_indices.copy_(metadata.mamba_track_indices) + metadata.mamba_track_indices = static_track_indices + self.forward_metadata = metadata + if use_preallocated_gdn_metadata and shape_was_preallocated: + _prime_gdn_prefill_metadata(static_query_start_loc, num_tokens) + return None + + MambaAttnBackendBase.init_forward_metadata_out_graph = 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/log_scan.py b/python/foundry/integration/sglang/log_scan.py new file mode 100644 index 00000000..db7dd93d --- /dev/null +++ b/python/foundry/integration/sglang/log_scan.py @@ -0,0 +1,25 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the Foundry project +"""Runtime error scanning shared by SGLang validation harnesses.""" + +from __future__ import annotations + +import re + +RUNTIME_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, +) +IGNORED_ERROR_BLOCK_PATTERN = re.compile( + r"\[start of libtorchcodec loading traceback\].*?" + r"\[end of libtorchcodec loading traceback\]\.?", + re.IGNORECASE | re.DOTALL, +) + + +def find_runtime_errors(text: str) -> list[str]: + """Return distinct error markers after removing known caught import traces.""" + error_text = IGNORED_ERROR_BLOCK_PATTERN.sub("", text) + return sorted({match.group(0) for match in RUNTIME_ERROR_PATTERN.finditer(error_text)}) diff --git a/python/foundry/integration/sglang/main_backend.py b/python/foundry/integration/sglang/main_backend.py new file mode 100644 index 00000000..25e85d2a --- /dev/null +++ b/python/foundry/integration/sglang/main_backend.py @@ -0,0 +1,356 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the Foundry project +"""Persisted full-graph backend for the current SGLang runner API.""" + +from __future__ import annotations + +import logging +import os +import time +from collections.abc import Callable, Iterator +from contextlib import contextmanager +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.full_cuda_graph_backend import ( + FullCudaGraphBackend, +) +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.archive import ( + GraphCatalog, + runner_descriptor, + shape_key_record, +) +from foundry.integration.sglang.config import ( + CUDAGraphExtensionMode, + get_config, + get_graph_extension_mode, +) +from foundry.integration.sglang.main_deepep import ( + bootstrap_deepep_buffer, + set_deepep_graph_mode, +) +from foundry.integration.sglang.output_codec import pack_output, unpack_output + +logger = logging.getLogger(__name__) + + +class FoundryMainCudaGraphBackend(FullCudaGraphBackend): + """Materialize SGLang main's full prefill and decode graphs through Foundry.""" + + def __init__(self, cuda_graph_runner, *, phase: str = "decode") -> None: + cfg = get_config() + if cfg is None or cfg.workspace_dir is None: + raise RuntimeError("Foundry SGLang graph extension is not initialized") + + self._runner = cuda_graph_runner + self._phase = phase + self._descriptor = runner_descriptor(cuda_graph_runner, phase) + self._catalog = GraphCatalog(cfg.workspace_dir) + self._session: dict[str, Any] | None = None + 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 + if cuda_graph_runner.model_runner.server_args.enable_torch_symm_mem: + logger.info("[Foundry] SGLang-main communication backend=torch_symmetric_memory") + + @contextmanager + def capture_session(self, stream) -> Iterator[None]: + if self._closed: + raise RuntimeError("Foundry SGLang-main does not support CUDA graph recapture") + bootstrap_deepep_buffer(self._runner) + 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() + state = rt.get_state() + if state is None: + raise RuntimeError("Foundry SGLang graph extension is not initialized") + + session_index = state.session_index + if mode == CUDAGraphExtensionMode.SAVE: + self._session = self._catalog.open_save_session( + session_index, + self._descriptor, + ) + elif mode == CUDAGraphExtensionMode.LOAD: + self._session = self._catalog.claim_load_session( + session_index, + self._descriptor, + ) + else: + raise RuntimeError(f"Foundry backend used in unsupported mode: {mode.value}") + state.session_index += 1 + + completed = False + if mode == CUDAGraphExtensionMode.LOAD: + self._prepare_load() + + try: + yield + completed = True + finally: + self._stream = None + 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 phase=%s role=%s session=%d", + self._load_index, + self._phase, + self._descriptor["role"], + self._session["session_index"], + ) + + 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") + + rt.log_alloc_offset("before_preallocate") + rt.preallocate_for_load_mode() + rt.log_alloc_offset("after_preallocate") + + if self._session is None: + raise RuntimeError("Foundry SGLang graph session is not initialized") + self._graph_files = self._session["graphs"] + if not self._graph_files: + raise RuntimeError( + "No Foundry SGLang graph files found for " + f"phase={self._phase} role={self._descriptor['role']} " + f"session={self._session['session_index']}" + ) + paths = [ + os.path.join(cfg.workspace_dir, graph_record["filename"]) + for graph_record in self._graph_files + ] + if self._pool is None: + raise RuntimeError("Foundry SGLang graph pool is not initialized") + nvshmem_module_count = cge.init_nvshmem_for_loaded_modules() + logger.info( + "[Foundry] Initialized NVSHMEM for %d loaded modules before graph builds", + nvshmem_module_count, + ) + self._pending = FoundryCUDAGraph.start_graph_builds( + paths, + pool=self._pool, + num_threads=4, + ) + self._load_index = 0 + logger.info( + "[Foundry] Started %d SGLang-main graph builds phase=%s role=%s session=%d", + len(paths), + self._phase, + self._descriptor["role"], + self._session["session_index"], + ) + + 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 + shape = shape_key_record(shape_key) + mode = get_graph_extension_mode() + + if mode == CUDAGraphExtensionMode.SAVE: + for _ in range(2): + self._runner.device_module.synchronize() + self._runner.model_runner.tp_group.barrier() + forward_fn() + if post_warmup_hook is not None: + post_warmup_hook() + self._capture_one(shape_key, shape, forward_fn) + return + if mode == CUDAGraphExtensionMode.LOAD: + self._load_one(shape_key, shape) + return + raise RuntimeError(f"Foundry backend used in unsupported mode: {mode.value}") + + def _capture_one( + self, + shape_key, + shape: dict[str, int | str | None], + forward_fn: Callable[[], Any], + ) -> None: + if self._stream is None: + raise RuntimeError("Foundry capture session has no CUDA stream") + + graph = FoundryCUDAGraph() + with foundry_graph( + graph, + pool=self._pool, + stream=self._stream, + ): + output = forward_fn() + + self._save_graph(graph, output, shape) + self._graphs[shape_key] = graph + self._outputs[shape_key] = output + + def _save_graph( + self, + graph: FoundryCUDAGraph, + output: Any, + shape: dict[str, int | str | None], + ) -> None: + cfg = get_config() + state = rt.get_state() + if cfg is None or state is None or cfg.workspace_dir is None or self._session is None: + raise RuntimeError("Foundry SGLang graph extension is not initialized") + + packed = pack_output(output, allowed_types=(LogitsProcessorOutput,)) + filename = ( + f"graph_{state.capture_index}_{self._phase}_{self._descriptor['role']}" + f"_s{self._session['session_index']}_k{shape['size']}.json" + ) + graph.save(os.path.join(cfg.workspace_dir, filename), packed.tensors) + self._catalog.append_graph( + self._session["session_index"], + { + "capture_index": state.capture_index, + "filename": filename, + "shape_key": shape, + "output_schema": packed.schema, + }, + ) + state.capture_index += 1 + logger.info( + "[Foundry] Saved SGLang-main CUDA graph %s phase=%s role=%s session=%d shape=%s", + filename, + self._phase, + self._descriptor["role"], + self._session["session_index"], + shape, + ) + + def _load_one( + self, + shape_key, + shape: dict[str, int | str | None], + ) -> None: + if self._pending is None: + 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") + + graph_record = self._graph_files[self._load_index] + filename = graph_record["filename"] + archived_shape = graph_record["shape_key"] + if archived_shape != shape: + raise RuntimeError( + f"SGLang graph order mismatch: requested {shape}, archive has " + f"{archived_shape} ({filename})" + ) + + started_at = time.perf_counter() + graph, tensors = FoundryCUDAGraph.finish_one_graph_load( + self._pending, + self._load_index, + ) + self._load_index += 1 + self._graphs[shape_key] = graph + self._outputs[shape_key] = unpack_output( + tensors, + graph_record["output_schema"], + allowed_types=(LogitsProcessorOutput,), + ) + logger.info( + "[Foundry] Loaded SGLang-main graph %s in %.3fs phase=%s role=%s session=%d shape=%s", + filename, + time.perf_counter() - started_at, + self._phase, + self._descriptor["role"], + self._session["session_index"], + shape, + ) + + @staticmethod + def _finish_save() -> None: + cfg = get_config() + if cfg is None or cfg.workspace_dir is None: + raise RuntimeError("Foundry SGLang graph extension is not initialized") + manifest_path = os.path.join(cfg.workspace_dir, "graph_manifest.json") + if os.path.exists(manifest_path): + os.remove(manifest_path) + # The shared-executor optimizer currently corrupts distributed replay. + # Keep complete per-graph edges so LOAD can reconstruct each graph in + # SAVE order until that generic optimization is safe for collectives. + # Rebuild the aggregate manifest after each prefill/decode session. + foundry_pkg.save_graph_manifest( + cfg.workspace_dir, + strip_dependencies=False, + ) + 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]: + set_deepep_graph_mode() + if get_graph_extension_mode() == CUDAGraphExtensionMode.LOAD: + state = rt.get_state() + if state is None: + raise RuntimeError("Foundry SGLang graph extension is not initialized") + if not state.load_sessions_validated: + session_count = self._catalog.session_count + if state.session_index != session_count: + raise RuntimeError( + "Foundry loaded " + f"{state.session_index}/{session_count} " + "SGLang graph sessions" + ) + state.load_sessions_validated = True + yield + + def replay(self, shape_key, static_forward_batch, **kwargs) -> Any: + del static_forward_batch, kwargs + self._graphs[shape_key].replay() + return self._outputs[shape_key] + + def cleanup(self) -> None: + if self._graphs: + self._closed = True + self._graphs.clear() + self._outputs.clear() + self._pool = None + self._pending = None + self._graph_files = [] + self._load_index = 0 + self._session = None diff --git a/python/foundry/integration/sglang/main_deepep.py b/python/foundry/integration/sglang/main_deepep.py new file mode 100644 index 00000000..77d9b11a --- /dev/null +++ b/python/foundry/integration/sglang/main_deepep.py @@ -0,0 +1,64 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the Foundry project +"""DeepEP setup for the current SGLang CUDA graph runner.""" + +from __future__ import annotations + +import logging +from typing import Any + +from sglang.srt.layers.moe.token_dispatcher.deepep import ( + DeepEPBuffer, + DeepEPDispatcher, +) +from sglang.srt.layers.moe.utils import get_deepep_mode, get_moe_a2a_backend + +logger = logging.getLogger(__name__) + + +def set_deepep_graph_mode() -> bool: + """Select the low-latency DeepEP mode used by decode graphs.""" + if not get_moe_a2a_backend().is_deepep(): + return False + + mode = get_deepep_mode().resolve(is_extend_in_batch=False) + DeepEPBuffer.set_dispatch_mode(mode) + return True + + +def bootstrap_deepep_buffer(cuda_graph_runner: Any) -> bool: + """Create SGLang's resource-backed DeepEP buffer before graph work.""" + if not get_moe_a2a_backend().is_deepep(): + return False + + if DeepEPBuffer._state().buffer is not None: + set_deepep_graph_mode() + logger.info("[Foundry] SGLang DeepEP buffer bootstrap ready") + return True + + for module in cuda_graph_runner.model_runner.model.modules(): + dispatcher = getattr(module, "dispatcher", None) + if dispatcher is None: + continue + + candidates = (dispatcher, *getattr(dispatcher, "_inners", ())) + for candidate in candidates: + if not isinstance(candidate, DeepEPDispatcher): + continue + implementation = getattr(candidate, "_low_latency_dispatcher", None) + if implementation is None: + continue + + implementation._get_buffer() + if DeepEPBuffer._state().buffer is None: + raise RuntimeError( + "SGLang DeepEP dispatcher bootstrap did not create its state buffer" + ) + set_deepep_graph_mode() + logger.info("[Foundry] SGLang DeepEP buffer bootstrap ready") + return True + + raise RuntimeError( + "SGLang DeepEP backend is active, but no low-latency DeepEP dispatcher " + "was found on the CUDA graph runner model" + ) diff --git a/python/foundry/integration/sglang/output_codec.py b/python/foundry/integration/sglang/output_codec.py new file mode 100644 index 00000000..3117cc72 --- /dev/null +++ b/python/foundry/integration/sglang/output_codec.py @@ -0,0 +1,177 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the Foundry project +"""Flatten and reconstruct nested SGLang CUDA graph output trees.""" + +from __future__ import annotations + +from collections.abc import Callable +from dataclasses import dataclass +from typing import Any + +import torch + + +@dataclass(frozen=True) +class PackedOutput: + tensors: list[torch.Tensor] + schema: dict[str, Any] + + +def _type_name(cls: type) -> str: + module = cls.__module__ + qualname = cls.__qualname__ + if module and module != "__main__": + return f"{module}.{qualname}" + return qualname + + +def _encode_sequence( + value: list[Any] | tuple[Any, ...], + path: str, + encode: Callable[[Any, str], dict[str, Any]], + active_containers: set[int], +) -> dict[str, Any]: + identity = id(value) + if identity in active_containers: + raise TypeError(f"Recursive SGLang graph output at {path}") + active_containers.add(identity) + try: + items = [encode(item, f"{path}[{index}]") for index, item in enumerate(value)] + kind = "tuple" if isinstance(value, tuple) else "list" + return {"kind": kind, "items": items} + finally: + active_containers.discard(identity) + + +def _encode_mapping( + value: dict[str, Any], + path: str, + encode: Callable[[Any, str], dict[str, Any]], + active_containers: set[int], +) -> dict[str, Any]: + identity = id(value) + if identity in active_containers: + raise TypeError(f"Recursive SGLang graph output at {path}") + for key in value: + if not isinstance(key, str): + raise TypeError( + f"Unsupported SGLang graph output at {path}: non-string dict key {key!r}" + ) + active_containers.add(identity) + try: + entries = {key: encode(value[key], f"{path}.{key}") for key in sorted(value)} + return {"kind": "dict", "entries": entries} + finally: + active_containers.discard(identity) + + +def pack_output( + output: Any, + *, + allowed_types: tuple[type, ...], +) -> PackedOutput: + tensors: list[torch.Tensor] = [] + tensor_indices: dict[int, int] = {} + active_containers: set[int] = set() + allowed = {_type_name(cls): cls for cls in allowed_types} + + def encode(value: Any, path: str) -> dict[str, Any]: + if isinstance(value, torch.Tensor): + identity = id(value) + index = tensor_indices.get(identity) + if index is None: + index = len(tensors) + tensor_indices[identity] = index + tensors.append(value) + return {"kind": "tensor", "index": index} + if value is None or isinstance(value, (bool, int, float, str)): + return {"kind": "constant", "value": value} + if isinstance(value, (list, tuple)): + return _encode_sequence(value, path, encode, active_containers) + if isinstance(value, dict): + return _encode_mapping(value, path, encode, active_containers) + + type_name = _type_name(type(value)) + if type_name not in allowed: + raise TypeError(f"Unsupported SGLang graph output at {path}: {type(value)!r}") + attributes = { + name: encode(attribute, f"{path}.{name}") + for name, attribute in sorted(vars(value).items()) + } + return { + "kind": "object", + "type": type_name, + "attributes": attributes, + } + + return PackedOutput(tensors=tensors, schema=encode(output, "$")) + + +def _collect_tensor_indices(schema: dict[str, Any]) -> set[int]: + kind = schema["kind"] + if kind == "tensor": + return {schema["index"]} + if kind in ("list", "tuple"): + indices: set[int] = set() + for item in schema["items"]: + indices.update(_collect_tensor_indices(item)) + return indices + if kind == "dict": + indices = set() + for entry in schema["entries"].values(): + indices.update(_collect_tensor_indices(entry)) + return indices + if kind == "object": + indices = set() + for attribute in schema["attributes"].values(): + indices.update(_collect_tensor_indices(attribute)) + return indices + return set() + + +def unpack_output( + tensors: Any, + schema: dict[str, Any], + *, + allowed_types: tuple[type, ...], +) -> Any: + allowed = {_type_name(cls): cls for cls in allowed_types} + + if isinstance(tensors, tuple): + tensor_list = list(tensors) + elif isinstance(tensors, list): + tensor_list = tensors + else: + tensor_list = list(tensors) + + referenced_indices = _collect_tensor_indices(schema) + expected_count = max(referenced_indices) + 1 if referenced_indices else 0 + if len(tensor_list) != expected_count: + raise ValueError( + f"tensor count mismatch: expected {expected_count}, got {len(tensor_list)}" + ) + + def decode(node: dict[str, Any]) -> Any: + kind = node["kind"] + if kind == "tensor": + return tensor_list[node["index"]] + if kind == "constant": + return node["value"] + if kind == "list": + return [decode(item) for item in node["items"]] + if kind == "tuple": + return tuple(decode(item) for item in node["items"]) + if kind == "dict": + return {key: decode(value) for key, value in node["entries"].items()} + if kind == "object": + type_name = node["type"] + if type_name not in allowed: + raise ValueError(f"Object type {type_name!r} is not allowlisted") + cls = allowed[type_name] + obj = cls.__new__(cls) + for name, attribute_schema in node["attributes"].items(): + setattr(obj, name, decode(attribute_schema)) + return obj + raise ValueError(f"Unsupported output schema node kind: {kind!r}") + + return decode(schema) diff --git a/python/foundry/integration/sglang/plugin.py b/python/foundry/integration/sglang/plugin.py new file mode 100644 index 00000000..ddbe96f8 --- /dev/null +++ b/python/foundry/integration/sglang/plugin.py @@ -0,0 +1,94 @@ +# 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") + + if server_args.cuda_graph_backend_decode is None: + server_args.cuda_graph_backend_decode = "full" + 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.pp_size != 1: + raise RuntimeError("Foundry SGLang-main does not yet support pipeline parallelism") + uses_foundry_full = ( + server_args.cuda_graph_config.decode.backend == "full" + or server_args.cuda_graph_config.prefill.backend == "full" + ) + if server_args.tp_size > 1 and uses_foundry_full and not server_args.enable_torch_symm_mem: + raise RuntimeError( + "Foundry SGLang-main tensor parallelism requires --enable-torch-symm-mem" + ) + if ( + server_args.tp_size > 1 + and server_args.cuda_graph_config.decode.backend == "full" + and ( + server_args.cuda_graph_config.decode.bs is None + or len(server_args.cuda_graph_config.decode.bs) != 1 + ) + ): + raise RuntimeError( + "Foundry SGLang-main tensor parallelism requires one explicit " + "--cuda-graph-bs-decode shape" + ) + 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") + + +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/runtime.py b/python/foundry/integration/sglang/runtime.py index 1f9032c5..30287e70 100644 --- a/python/foundry/integration/sglang/runtime.py +++ b/python/foundry/integration/sglang/runtime.py @@ -37,14 +37,18 @@ class WarmupState: gpu_name: str = "" gpu_total_memory: int = 0 memory_pool_config: dict = field(default_factory=dict) + max_mamba_cache_size: int | None = None final_alloc_offset: int = 0 @dataclass class CUDAGraphExtensionState: capture_index: int = 0 + session_index: int = 0 rank: int = 0 loaded_graphs: dict = field(default_factory=dict) + load_preallocated: bool = False + load_sessions_validated: bool = False _state: CUDAGraphExtensionState | None = None @@ -71,7 +75,7 @@ def create_warmup_state(memory_pool_config: dict | None = None) -> WarmupState: except Exception: sglang_version = "unknown" - props = torch.cuda.get_device_properties(0) + props = torch.cuda.get_device_properties(torch.cuda.current_device()) return WarmupState( sglang_version=sglang_version, timestamp=datetime.now().isoformat(), @@ -87,12 +91,14 @@ def save_warmup_state(state: WarmupState) -> None: if workspace_root is None: return ext_state = get_state() - path = os.path.join(workspace_root, "warmup_state.json") - if ext_state is not None and ext_state.rank != 0 and os.path.exists(path): + if ext_state is not None and ext_state.rank != 0: return + path = os.path.join(workspace_root, "warmup_state.json") os.makedirs(workspace_root, exist_ok=True) - with open(path, "w") as f: + temp_path = f"{path}.tmp.{os.getpid()}" + with open(temp_path, "w") as f: json.dump(asdict(state), f, indent=2) + os.replace(temp_path, path) logger.info("[Foundry] Saved SGLang warmup state to %s", path) @@ -120,6 +126,22 @@ def setup_graph_extension(server_args, tp_rank: int, pp_rank: int, dp_rank: int rank = compute_workspace_rank(server_args, tp_rank, pp_rank, dp_rank) Path(cfg.workspace_root).mkdir(parents=True, exist_ok=True) workspace_dir = Path(cfg.workspace_root) / f"rank_{rank}" + if _state is not None: + if _state.rank != rank: + raise RuntimeError( + "Foundry SGLang graph extension is already initialized for " + f"rank {_state.rank}, cannot reinitialize for rank {rank}" + ) + if cfg.workspace_dir != str(workspace_dir): + raise RuntimeError( + "Foundry SGLang graph extension workspace changed after initialization: " + f"{cfg.workspace_dir!r} != {str(workspace_dir)!r}" + ) + logger.info( + "[Foundry] Reusing process-local SGLang graph extension for rank=%d", + rank, + ) + return cfg.workspace_dir = str(workspace_dir) logger.info("[Foundry] SGLang rank=%d workspace_dir=%s", rank, workspace_dir) @@ -184,6 +206,11 @@ def preallocate_for_load_mode() -> None: cfg = get_config() if cfg is None or cfg.mode != CUDAGraphExtensionMode.LOAD: return + state = get_state() + if state is None: + raise RuntimeError("Foundry SGLang graph extension is not initialized") + if state.load_preallocated: + return final = 0 if cfg.workspace_dir is not None: path = os.path.join(cfg.workspace_dir, "final_alloc_offset.json") @@ -195,6 +222,7 @@ def preallocate_for_load_mode() -> None: remaining = final - cge.get_current_alloc_offset() if remaining > 0: cge.preallocate_region(remaining) + state.load_preallocated = True def log_alloc_offset(label: str) -> None: 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..dbcddc27 --- /dev/null +++ b/recipe/experimental/README.md @@ -0,0 +1,241 @@ +# 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-main tensor-parallel recipe, +which captures one padded torch symmetric-memory 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 TP SAVE config +โ”œโ”€โ”€ sglang_foundry_tp_load.toml # SGLang TP LOAD config +โ””โ”€โ”€ serve_qwen3-8b_sglang_tp.sh # SGLang-main symmetric-memory TP +``` + +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. + +## 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 tensor parallel + +`serve_qwen3-8b_sglang_tp.sh` runs tensor parallelism across `` GPUs +(`--tp-size N`, no DP-attention, no expert parallel). The default Qwen3-8B is +dense; the same TP-only path also validates Qwen3.5-122B-A10B without enabling +expert parallelism. This is the first recipe here that captures a cross-rank +collective into the graph itself. + +```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 current-main TP survives SAVE/LOAD: + +- The recipe enables torch symmetric memory by default on SGLang main. Foundry + reuses its stable buffer for per-layer all-reduce and logits all-gather. +- One explicit padded decode graph is captured. `SGLANG_CUDA_GRAPH_MAX_BS` + selects the default shape; `SGLANG_CUDA_GRAPH_BS` can override it. + Multi-shape **decode** archives are rejected because later shapes retain + collective capture-order state that is not portable to a fresh process. +- Full-backend prefill persistence can additionally capture multiple token + buckets through `SGLANG_CUDA_GRAPH_BS_PREFILL`; 16/64 is GPU-validated. + Breakable and `tc` prefill persistence are not validated. +- SGLang custom all-reduce and FlashInfer all-reduce fusion are disabled so + they cannot silently bypass the validated symmetric path. +- `NCCL_CUMEM_ENABLE=0` / `NCCL_NVLS_ENABLE=0` keep NCCL off the CUMEM P2P and + NVLS multicast fast paths when SGLang initializes its fallback communicator. + +No NVSHMEM is involved in the plain dense TP recipe, so its TOMLs leave +`nvshmem_host_path` unset. Keep TP size, model settings, graph shape, and +communication settings identical between SAVE and LOAD. The script pins +`--random-seed 42` by default (`SGL_RANDOM_SEED` overrides it). + +## 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 recipe is validated on SGLang main +`9b853e6832e71a3058212df02a025232a453e146`: + +- Qwen3-8B TP=2 on 2ร—H100: baseline, two SAVE passes, and fresh LOAD passed + every check. One graph per rank restored at offset `60442017792`; sequential + and fixed-batch outputs matched baseline. +- Qwen3.5-122B-A10B-GPTQ-Int4 TP=4 on 4ร—H100 with `moe_wna16`, FA3, Triton + linear attention/MoE, decode batch 8, and Full prefill 16/64: three graphs + per rank restored at offset `59301167104`; exact batched tokens, symmetric + collectives, catalogs, archive fingerprints, and offsets reproduced across + baseline, SAVE, SAVE2, and LOAD. Modal run: + `eb1ec00b7004339c5150ce6c-1784907539317556437`. +- Qwen3.5-35B-A3B online FP8 on 2ร—H100 combined TP=2/EP=2, NEXTN (3 steps, + top-k 1, 4 draft tokens), DeepEP low-latency, decode batch 8, and Full + prefill 16/64. Run `eb1ec00b7004339c5150ce6c-1784906247386993159` at + immutable Foundry `eb1ec00b7004339c5150ce6c3ba057ceda047833` restored four + graphs and target-decode, target-prefill, and draft-decode sessions on each + rank. It reproduced exact token IDs, average acceptance `2.775`, semantic + DeepEP nodes, rank-local offsets, and clean no-recapture LOAD. + +The checked-in proof runs the complete baseline โ†’ SAVE โ†’ SAVE โ†’ LOAD matrix: + +```bash +modal run --env tests/modal_sglang_tp.py +``` + +The combined Hopper run requires FA3, `SGLANG_MEM_FRACTION_STATIC=0.7`, +`NVSHMEM_REMOTE_TRANSPORT=none`, `NVSHMEM_IB_ENABLE=0`, and +`NVSHMEM_IB_ENABLE_IBGDA=0`; it observed legacy NCCL `P2P/IPC` channels. The +no-fabric SGLang DeepEP variant above separately validates legacy NVL +CUDA-IPC through `FOUNDRY_DEEPEP_NVL_IPC=1`, while NVSHMEM remains required +for the RDMA symmetric heap. + +This is pinned single-shape-decode symmetric-memory TP support, not arbitrary +collective or topology support. NEXTN is the only GPU-validated speculative +algorithm. Other speculative algorithms, other fabrics, breakable/`tc` +prefill persistence, PP, LoRA, PDMux, diffusion, and recapture remain +non-goals. Track broader support in +[issue #6](https://github.com/foundry-org/foundry/issues/6). + +## 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..18c453b3 --- /dev/null +++ b/recipe/experimental/serve_qwen3-8b_sglang_tp.sh @@ -0,0 +1,124 @@ +#!/bin/bash +# Usage: CUDA_VISIBLE_DEVICES=0,1 bash serve_qwen3-8b_sglang_tp.sh [--save|--load] +# +# SGLang tensor parallel (tp-size > 1, no DP-attention or expert parallel). +# Current main captures one padded graph with torch symmetric-memory all-reduce +# and all-gather. Legacy builds retain the NCCL fallback. +# +# Multi-shape TP is intentionally rejected because later captures retain +# collective state that is not portable to a fresh process. +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}" +CUDA_GRAPH_BS="${SGLANG_CUDA_GRAPH_BS:-$CUDA_GRAPH_MAX_BS}" +ATTENTION_BACKEND="${SGLANG_ATTENTION_BACKEND:-flashinfer}" +RANDOM_SEED="${SGL_RANDOM_SEED:-42}" +PREFILL_BACKEND="${SGLANG_CUDA_GRAPH_BACKEND_PREFILL:-disabled}" +GRAPH_ARGS=( --cuda-graph-backend-prefill "$PREFILL_BACKEND" ) +if [[ -n "${SGLANG_CUDA_GRAPH_BS_PREFILL:-}" ]]; then + read -ra PREFILL_BS_BUCKETS <<< "${SGLANG_CUDA_GRAPH_BS_PREFILL}" + GRAPH_ARGS+=( --cuda-graph-bs-prefill "${PREFILL_BS_BUCKETS[@]}" ) +fi +SPEC_ARGS=() +if [[ -n "${SGLANG_SPECULATIVE_ALGORITHM:-}" ]]; then + SPEC_ARGS+=( --speculative-algorithm "$SGLANG_SPECULATIVE_ALGORITHM" ) +fi +if [[ -n "${SGLANG_SPECULATIVE_NUM_STEPS:-}" ]]; then + SPEC_ARGS+=( --speculative-num-steps "$SGLANG_SPECULATIVE_NUM_STEPS" ) +fi +if [[ -n "${SGLANG_SPECULATIVE_EAGLE_TOPK:-}" ]]; then + SPEC_ARGS+=( --speculative-eagle-topk "$SGLANG_SPECULATIVE_EAGLE_TOPK" ) +fi +if [[ -n "${SGLANG_SPECULATIVE_NUM_DRAFT_TOKENS:-}" ]]; then + SPEC_ARGS+=( --speculative-num-draft-tokens "$SGLANG_SPECULATIVE_NUM_DRAFT_TOKENS" ) +fi +DEEPEP_ARGS=() +if [[ -n "${SGLANG_MOE_A2A_BACKEND:-}" ]]; then + DEEPEP_ARGS+=( --moe-a2a-backend "$SGLANG_MOE_A2A_BACKEND" ) +fi +if [[ -n "${SGLANG_DEEPEP_MODE:-}" ]]; then + DEEPEP_ARGS+=( --deepep-mode "$SGLANG_DEEPEP_MODE" ) +fi +USES_MAIN_RUNNER=0 +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 + USES_MAIN_RUNNER=1 +fi +ENABLE_TORCH_SYMM_MEM="${SGLANG_ENABLE_TORCH_SYMM_MEM:-$USES_MAIN_RUNNER}" + +# Keep NCCL fallback initialization identical across baseline, SAVE, and LOAD. +# Current-main Foundry TP captures torch symmetric-memory collectives instead. +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 [[ -n "${SGLANG_MAX_RUNNING_REQUESTS:-}" ]]; then + MODEL_ARGS+=( --max-running-requests "$SGLANG_MAX_RUNNING_REQUESTS" ) +fi +if [[ -n "${SGLANG_MAX_TOTAL_TOKENS:-}" ]]; then + MODEL_ARGS+=( --max-total-tokens "$SGLANG_MAX_TOTAL_TOKENS" ) +fi +if [[ "$ENABLE_TORCH_SYMM_MEM" == "1" ]]; then + MODEL_ARGS+=( --enable-torch-symm-mem ) +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 [[ "$USES_MAIN_RUNNER" != "1" ]]; then + FOUNDRY_ARGS+=( --foundry-graph-extension-config-path "$FOUNDRY_TOML" ) + fi + echo "Using Foundry TP: ${FOUNDRY_TOML}" +else + echo "Running without Foundry (baseline SGLang)" +fi + +# Disable SGLang's unrelated custom/fused paths so current main consistently +# uses torch symmetric memory and legacy builds consistently use NCCL. +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 \ + "${GRAPH_ARGS[@]}" \ + --disable-custom-all-reduce \ + --enforce-disable-flashinfer-allreduce-fusion \ + --attention-backend "$ATTENTION_BACKEND" \ + --cuda-graph-max-bs "$CUDA_GRAPH_MAX_BS" \ + --cuda-graph-bs "$CUDA_GRAPH_BS" \ + --decode-log-interval 1 \ + "${SPEC_ARGS[@]}" \ + "${DEEPEP_ARGS[@]}" \ + "${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..d6fbc037 100644 --- a/recipe/sglang/README.md +++ b/recipe/sglang/README.md @@ -39,7 +39,9 @@ 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` | -| Expert parallel | `serve_qwen3-30ba3bfp8_ep.sh` | Qwen3-30B-A3B-FP8 | DP-attention + DeepEP; fa3 backend | +| Tensor parallel | `../experimental/serve_qwen3-8b_sglang_tp.sh` | Qwen3-8B / Qwen3.5-122B-A10B | one symmetric-memory decode shape; TP=2/TP=4 | +| Expert parallel | `serve_qwen3-30ba3bfp8_ep.sh` | Qwen3-30B-A3B-FP8 | DP-attention + DeepEP low-latency; FA3 | +| Combined validation | Modal harness | Qwen3.5-35B-A3B online FP8 | TP=2/EP=2 + NEXTN + DeepEP + Full prefill 16/64 | ## Installation @@ -92,6 +94,43 @@ 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 (tensor parallel) + +Current-main TP captures one padded decode shape using torch symmetric-memory +all-reduce and all-gather. Keep the topology, graph shape, and model settings +identical across baseline, SAVE, and LOAD. The plugin rejects multi-shape +decode archives and TP without symmetric memory. Full-backend prefill +persistence may add multiple token buckets; breakable and `tc` prefill are not +validated. + +```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 baseline โ†’ SAVE โ†’ SAVE โ†’ LOAD proof, run +`modal run --env tests/modal_sglang_tp.py` from the +repository root. The harness checks sequential and fixed-batch outputs, +per-rank archive fingerprints and allocation offsets, graph replay without +recapture, and the actual collective kernels in every rank archive. Validated +configurations are Qwen3-8B TP=2 and Qwen3.5-122B-A10B-GPTQ-Int4 TP=4 on H100. +The final TP4 Full-prefill gate used decode batch 8 and prefill buckets 16/64; +run `eb1ec00b7004339c5150ce6c-1784907539317556437` restored three graphs per +rank at the symmetric offset `59301167104`. + +NEXTN is the only GPU-validated speculative algorithm. The composed proof used +Qwen3.5-35B-A3B online FP8, TP=2/EP=2, NEXTN 3 steps/top-k 1/4 draft tokens, +DeepEP low-latency, one decode-8 shape, and Full prefill 16/64. Modal run +`eb1ec00b7004339c5150ce6c-1784906247386993159` at immutable Foundry +`eb1ec00b7004339c5150ce6c3ba057ceda047833` passed +baseline/SAVE/SAVE2/fresh-LOAD token, catalog, graph-node, archive, offset, and +no-recapture checks. + ## 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 .. ``` @@ -131,6 +170,19 @@ DeepEP low-latency caps dispatch at that per-rank token count (and asserts `(n+1)*2 <= NVSHMEM_QP_DEPTH`); keep it and `--chunked-prefill-size` identical between SAVE and LOAD so the captured graphs match. +For the pinned Hopper combined configuration, FA3 is required and the +evidence-backed settings are `SGLANG_MEM_FRACTION_STATIC=0.7`, +`NVSHMEM_REMOTE_TRANSPORT=none`, `NVSHMEM_IB_ENABLE=0`, and +`NVSHMEM_IB_ENABLE_IBGDA=0`. The run observed legacy NCCL `P2P/IPC` channels; +it does not validate another NVSHMEM fabric. On no-fabric hosts, +`FOUNDRY_DEEPEP_NVL_IPC=1` selects the separately validated legacy CUDA-IPC +path for DeepEP's NVL buffer; NVSHMEM remains required for the RDMA symmetric +heap. + +The current validation scope excludes other speculative algorithms, +multi-shape decode TP, arbitrary collective/fabric backends, breakable/`tc` +prefill persistence, PP, LoRA, PDMux, diffusion, and recapture. + ## Archive layout ``` diff --git a/recipe/sglang/serve_qwen3-mini.sh b/recipe/sglang/serve_qwen3-mini.sh index 77449482..1a68ab05 100755 --- a/recipe/sglang/serve_qwen3-mini.sh +++ b/recipe/sglang/serve_qwen3-mini.sh @@ -4,10 +4,36 @@ 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 +PREFILL_BACKEND="${SGLANG_CUDA_GRAPH_BACKEND_PREFILL:-disabled}" +GRAPH_ARGS=( --cuda-graph-backend-prefill "$PREFILL_BACKEND" ) +if [[ -n "${SGLANG_CUDA_GRAPH_BS_PREFILL:-}" ]]; then + read -ra PREFILL_BS_BUCKETS <<< "${SGLANG_CUDA_GRAPH_BS_PREFILL}" + GRAPH_ARGS+=( --cuda-graph-bs-prefill "${PREFILL_BS_BUCKETS[@]}" ) +fi +SPEC_ARGS=() +if [[ -n "${SGLANG_SPECULATIVE_ALGORITHM:-}" ]]; then + SPEC_ARGS+=( --speculative-algorithm "$SGLANG_SPECULATIVE_ALGORITHM" ) +fi +if [[ -n "${SGLANG_SPECULATIVE_NUM_STEPS:-}" ]]; then + SPEC_ARGS+=( --speculative-num-steps "$SGLANG_SPECULATIVE_NUM_STEPS" ) +fi +if [[ -n "${SGLANG_SPECULATIVE_EAGLE_TOPK:-}" ]]; then + SPEC_ARGS+=( --speculative-eagle-topk "$SGLANG_SPECULATIVE_EAGLE_TOPK" ) +fi +if [[ -n "${SGLANG_SPECULATIVE_NUM_DRAFT_TOKENS:-}" ]]; then + SPEC_ARGS+=( --speculative-num-draft-tokens "$SGLANG_SPECULATIVE_NUM_DRAFT_TOKENS" ) +fi +DEEPEP_ARGS=() +if [[ -n "${SGLANG_MOE_A2A_BACKEND:-}" ]]; then + DEEPEP_ARGS+=( --moe-a2a-backend "$SGLANG_MOE_A2A_BACKEND" ) +fi +if [[ -n "${SGLANG_DEEPEP_MODE:-}" ]]; then + DEEPEP_ARGS+=( --deepep-mode "$SGLANG_DEEPEP_MODE" ) +fi if [[ "$1" == "--save" ]]; then FOUNDRY_TOML="${SCRIPT_DIR}/foundry_save.toml" @@ -24,7 +50,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 +67,13 @@ sglang serve \ --trust-remote-code \ --host "$HOST" --port "$PORT" \ --tp-size 1 \ + --random-seed 42 \ --mem-fraction-static "$MEM_FRACTION_STATIC" \ --disable-radix-cache \ + "${GRAPH_ARGS[@]}" \ --attention-backend flashinfer \ --cuda-graph-max-bs 512 \ + --decode-log-interval 1 \ + "${SPEC_ARGS[@]}" \ + "${DEEPEP_ARGS[@]}" \ "${FOUNDRY_ARGS[@]}" diff --git a/tests/modal_deepep_fabric.py b/tests/modal_deepep_fabric.py new file mode 100644 index 00000000..a1e5149d --- /dev/null +++ b/tests/modal_deepep_fabric.py @@ -0,0 +1,90 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the Foundry project +"""Run the standalone DeepEP fabric graph oracle on two Modal H100s.""" + +from __future__ import annotations + +import importlib.util +import os +import re +import shutil +import subprocess +import sys +import time +from pathlib import Path + +import modal + +_TP_HARNESS_PATH = Path(__file__).with_name("modal_sglang_tp.py") +if not _TP_HARNESS_PATH.exists(): + _TP_HARNESS_PATH = Path("/foundry/tests/modal_sglang_tp.py") +_TP_HARNESS_SPEC = importlib.util.spec_from_file_location( + "foundry_modal_sglang_tp", + _TP_HARNESS_PATH, +) +if _TP_HARNESS_SPEC is None or _TP_HARNESS_SPEC.loader is None: + raise RuntimeError(f"Cannot load shared SGLang Modal image from {_TP_HARNESS_PATH}") +_TP_HARNESS = importlib.util.module_from_spec(_TP_HARNESS_SPEC) +_TP_HARNESS_SPEC.loader.exec_module(_TP_HARNESS) +FOUNDRY_REF = _TP_HARNESS.FOUNDRY_REF +build_image = _TP_HARNESS.build_image + +DATA_ROOT = Path("/data") +TEST_ARCHIVE = DATA_ROOT / "deepep_fabric_archive" +FAILED_RUNS_ROOT = DATA_ROOT / "archive" / "runs" + +image = build_image(install_deepep=True) +app = modal.App("foundry-deepep-fabric-validation") +archive_volume = modal.Volume.from_name( + "foundry-deepep-fabric-validation-archive", + create_if_missing=True, +) + + +@app.function( + image=image, + gpu="H100:2", + timeout=3600, + volumes={"/data": archive_volume}, +) +def run_oracle(run_id: str) -> None: + try: + subprocess.run( + [ + sys.executable, + "/foundry/tests/test_deepep_fabric.py", + "--run", + ], + check=True, + env={ + **os.environ, + "TEST_USE_FABRIC": "1", + "TEST_LOAD_API": "parallel", + # Legacy DeepEP performs NCCL object collectives during + # Buffer construction. Initialize those allocations before + # Foundry reserves its deterministic address region. + "TEST_NCCL_WARMUP_PRE_REGION": "1", + }, + cwd="/data", + ) + except BaseException: + if TEST_ARCHIVE.exists(): + retained_path = FAILED_RUNS_ROOT / run_id + retained_path.parent.mkdir(parents=True, exist_ok=True) + shutil.rmtree(retained_path, ignore_errors=True) + shutil.move(str(TEST_ARCHIVE), retained_path) + print(f"[DeepEP fabric] retained failed archive at {retained_path}", flush=True) + archive_volume.commit() + raise + else: + shutil.rmtree(TEST_ARCHIVE, ignore_errors=True) + archive_volume.commit() + print("[DeepEP fabric] semantic oracle passed; test archive deleted", flush=True) + + +@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()}" + print(f"[DeepEP fabric] run={run_id} failed_artifacts={FAILED_RUNS_ROOT / run_id}") + run_oracle.remote(run_id) diff --git a/tests/modal_sglang_main_smoke.py b/tests/modal_sglang_main_smoke.py new file mode 100644 index 00000000..2cbca724 --- /dev/null +++ b/tests/modal_sglang_main_smoke.py @@ -0,0 +1,366 @@ +# 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 importlib.util +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 + +_LOG_SCAN_PATH = ( + Path(__file__).parents[1] / "python" / "foundry" / "integration" / "sglang" / "log_scan.py" +) +if not _LOG_SCAN_PATH.exists(): + _LOG_SCAN_PATH = Path("/foundry/python/foundry/integration/sglang/log_scan.py") +_LOG_SCAN_SPEC = importlib.util.spec_from_file_location( + "foundry_sglang_log_scan", + _LOG_SCAN_PATH, +) +if _LOG_SCAN_SPEC is None or _LOG_SCAN_SPEC.loader is None: + raise RuntimeError(f"Cannot load SGLang log scanner from {_LOG_SCAN_PATH}") +_LOG_SCAN_MODULE = importlib.util.module_from_spec(_LOG_SCAN_SPEC) +_LOG_SCAN_SPEC.loader.exec_module(_LOG_SCAN_MODULE) +find_runtime_errors = _LOG_SCAN_MODULE.find_runtime_errors + + +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", + "9b853e6832e71a3058212df02a025232a453e146", +) +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" +RECIPE = "/foundry/recipe/sglang/serve_qwen3-mini.sh" + +PROMPTS = [ + "Name a primary color.", + ( + "In one concise sentence, explain why replaying a previously captured CUDA graph " + "can reduce inference overhead while preserving the model's deterministic token " + "sequence for an identical prompt and sampling configuration." + ), +] +MAX_NEW_TOKENS = 16 + +SAVED_GRAPH_LOG_PATTERN = re.compile( + r"\[Foundry\] Saved SGLang-main CUDA graph .* phase=(?P\w+)" +) +LOADED_GRAPH_LOG_PATTERN = re.compile( + r"\[Foundry\] Loaded SGLang-main graph .* phase=(?P\w+)" +) + +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() -> list[list[int]]: + request = urllib.request.Request( + f"http://127.0.0.1:{PORT}/generate", + data=json.dumps( + { + "text": PROMPTS, + "return_logprob": True, + "sampling_params": { + "temperature": 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()) + if not isinstance(result, list) or len(result) != len(PROMPTS): + raise RuntimeError(f"SGLang returned {type(result).__name__} for {len(PROMPTS)} prompts") + token_ids = [] + for item in result: + output_logprobs = item.get("meta_info", {}).get("output_token_logprobs") + if not output_logprobs: + raise RuntimeError("SGLang response omitted output token logprobs") + token_ids.append([int(entry[1]) for entry in output_logprobs]) + return token_ids + + +def _launch( + phase: str, + log_path: Path, + run_root: 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 + env["SGLANG_CUDA_GRAPH_BACKEND_PREFILL"] = "full" + env["SGLANG_CUDA_GRAPH_BS_PREFILL"] = "16 64" + 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() + + +def _catalog_summaries(workspace: Path) -> dict[str, list[dict]]: + summaries = {} + for rank_dir in sorted(workspace.glob("rank_*")): + catalog_path = rank_dir / "sglang_graph_catalog.json" + if not catalog_path.exists(): + continue + catalog = json.loads(catalog_path.read_text()) + summaries[rank_dir.name] = [ + { + "session_index": session["session_index"], + "phase": session["runner"].get("phase"), + "role": session["runner"].get("role"), + "mode": session["runner"].get("forward_mode"), + "graph_count": len(session.get("graphs", [])), + } + for session in catalog.get("sessions", []) + ] + return summaries + + +def _phase_log_counts(pattern: re.Pattern, text: str) -> dict[str, int]: + counts: dict[str, int] = {} + for match in pattern.finditer(text): + phase = match.group("phase") + counts[phase] = counts.get(phase, 0) + 1 + return counts + + +@app.function( + image=image, + gpu=GPU, + timeout=3600, + volumes={ARCHIVE_ROOT: archive_volume, HF_ROOT: hf_volume}, +) +def run_phase(phase: str, run_id: str) -> dict: + run_root = Path(ARCHIVE_ROOT) / run_id + workspace = run_root / "foundry_archive" + 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, run_root) + token_ids: list[list[int]] = [] + phase_error = "" + text = "" + try: + _wait_healthy(process, log_path) + token_ids = _generate() + if process.poll() is not None: + raise RuntimeError(f"SGLang exited unexpectedly\n{log_path}") + except Exception as error: # noqa: BLE001 - retain phase artifacts before propagating + phase_error = f"{type(error).__name__}: {error}" + finally: + log_file.flush() + text = log_path.read_text(errors="replace") + _stop(process, log_file) + + result = { + "phase": phase, + "artifact_dir": str(run_root), + "log_path": str(log_path), + "token_ids": token_ids, + "errors": find_runtime_errors(text), + "phase_error": phase_error, + "saved_graphs": _phase_log_counts(SAVED_GRAPH_LOG_PATTERN, text), + "loaded_graphs": _phase_log_counts(LOADED_GRAPH_LOG_PATTERN, text), + "native_capture_progress": text.count("Capturing batches"), + "replay_observed": "cuda graph: True" in text, + "plugin_observed": "SGLang graph extension setup completed" in text, + "rank_catalogs": _catalog_summaries(workspace), + } + archive_volume.commit() + print(json.dumps(result, sort_keys=True)) + if phase_error: + raise RuntimeError( + f"SGLang {phase} phase failed; artifacts retained at {run_root}: {phase_error}" + ) + return result + + +@app.function( + image=image, + volumes={ARCHIVE_ROOT: archive_volume}, +) +def cleanup(run_id: str) -> None: + shutil.rmtree(Path(ARCHIVE_ROOT) / run_id, ignore_errors=True) + archive_volume.commit() + + +@app.local_entrypoint() +def main() -> None: + 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}" + results = {phase: run_phase.remote(phase, run_id) for phase in ("baseline", "save", "load")} + save_catalogs = results["save"]["rank_catalogs"] + load_catalogs = results["load"]["rank_catalogs"] + save_sessions = [session for sessions in save_catalogs.values() for session in sessions] + prefill_graph_count = sum( + session["graph_count"] for session in save_sessions if session["phase"] == "prefill" + ) + checks = { + "exact_token_ids_match": len(results["baseline"]["token_ids"]) == len(PROMPTS) + and all(results["baseline"]["token_ids"]) + and results["baseline"]["token_ids"] + == results["save"]["token_ids"] + == results["load"]["token_ids"], + "token_ids_discriminate_prompts": len( + {tuple(token_ids) for token_ids in results["baseline"]["token_ids"]} + ) + == len(PROMPTS), + "rank_catalogs_present": set(save_catalogs) == {"rank_0"} + and save_catalogs == load_catalogs, + "prefill_and_decode_sessions": {session["phase"] for session in save_sessions} + >= {"prefill", "decode"}, + "at_least_two_prefill_graphs": prefill_graph_count >= 2, + "save_logged_prefill_and_decode": all( + results["save"]["saved_graphs"].get(phase, 0) > 0 for phase in ("prefill", "decode") + ), + "load_logged_prefill_and_decode": all( + results["load"]["loaded_graphs"].get(phase, 0) > 0 for phase in ("prefill", "decode") + ), + "load_did_not_capture": not results["load"]["saved_graphs"], + "load_replayed_graph": results["load"]["replay_observed"], + "plugin_loaded": results["save"]["plugin_observed"] and results["load"]["plugin_observed"], + "no_runtime_errors": not any(result["errors"] for result in results.values()), + } + report = { + "run_id": run_id, + "artifact_dir": run_root, + "artifacts_retained_on_failure": True, + "foundry_ref": FOUNDRY_REF, + "sglang_ref": SGLANG_REF, + "gpu": GPU, + "prompts": PROMPTS, + "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(run_id) diff --git a/tests/modal_sglang_tp.py b/tests/modal_sglang_tp.py new file mode 100644 index 00000000..fb6d2385 --- /dev/null +++ b/tests/modal_sglang_tp.py @@ -0,0 +1,1171 @@ +# 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 normal SGLang TP baseline and deterministic queries. +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 token-identical baseline/LOAD responses, reproducible graph +archives and per-rank allocation offsets, restored graph replay at multiple +batch sizes, NCCL P2P/IPC on every rank, and no CUDA/NCCL errors. + +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 importlib.util +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 + +_LOG_SCAN_PATH = ( + Path(__file__).parents[1] / "python" / "foundry" / "integration" / "sglang" / "log_scan.py" +) +if not _LOG_SCAN_PATH.exists(): + _LOG_SCAN_PATH = Path("/foundry/python/foundry/integration/sglang/log_scan.py") +_LOG_SCAN_SPEC = importlib.util.spec_from_file_location( + "foundry_sglang_log_scan", + _LOG_SCAN_PATH, +) +if _LOG_SCAN_SPEC is None or _LOG_SCAN_SPEC.loader is None: + raise RuntimeError(f"Cannot load SGLang log scanner from {_LOG_SCAN_PATH}") +_LOG_SCAN_MODULE = importlib.util.module_from_spec(_LOG_SCAN_SPEC) +_LOG_SCAN_SPEC.loader.exec_module(_LOG_SCAN_MODULE) +find_runtime_errors = _LOG_SCAN_MODULE.find_runtime_errors + + +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", + "9b853e6832e71a3058212df02a025232a453e146", +) +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 = ( + "NVSHMEM_IB_ENABLE", + "NVSHMEM_IB_ENABLE_IBGDA", + "NVSHMEM_REMOTE_TRANSPORT", + "SGLANG_QUANTIZATION", + "SGLANG_CONTEXT_LENGTH", + "SGLANG_REASONING_PARSER", + "SGLANG_LINEAR_ATTN_BACKEND", + "SGLANG_MOE_RUNNER_BACKEND", + "SGLANG_ATTENTION_BACKEND", + "SGLANG_MEM_FRACTION_STATIC", + "SGLANG_MAX_RUNNING_REQUESTS", + "SGLANG_MAX_TOTAL_TOKENS", + "SGLANG_CUDA_GRAPH_MAX_BS", + "SGLANG_CUDA_GRAPH_BS", + "SGLANG_ENABLE_TORCH_SYMM_MEM", + "SGLANG_CUDA_GRAPH_BACKEND_PREFILL", + "SGLANG_CUDA_GRAPH_BS_PREFILL", + "SGLANG_SPECULATIVE_ALGORITHM", + "SGLANG_SPECULATIVE_NUM_STEPS", + "SGLANG_SPECULATIVE_EAGLE_TOPK", + "SGLANG_SPECULATIVE_NUM_DRAFT_TOKENS", + "SGLANG_MOE_A2A_BACKEND", + "SGLANG_DEEPEP_MODE", + "TP_KEEP_ARTIFACTS", +) +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" +SPECULATIVE_ALGORITHM = SERVE_ENV.get("SGLANG_SPECULATIVE_ALGORITHM") +DEEPEP_ENABLED = SERVE_ENV.get("SGLANG_MOE_A2A_BACKEND") == "deepep" +FULL_PREFILL_ENABLED = SERVE_ENV.get("SGLANG_CUDA_GRAPH_BACKEND_PREFILL") == "full" +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 +BATCH_MAX_NEW_TOKENS = 8 +EXPECTED_GRAPH_COUNT = int(os.environ.get("EXPECTED_GRAPH_COUNT", "1")) +EXPECTED_NCCL_CHANNELS = 24 +KEEP_ARTIFACTS = os.environ.get("TP_KEEP_ARTIFACTS", "0") == "1" + +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 batch, #running-req: (?P\d+).*cuda graph: True" +) +TP_LOG_RANK_PATTERN = re.compile(r"\bTP(?P\d+)\b") +NVSHMEM_MODULE_COUNT_PATTERN = re.compile( + r"Initialized NVSHMEM for (?P\d+) loaded modules before graph builds" +) +DEEPEP_ERROR_PATTERN = re.compile( + r"(?:DeepEP|deep_ep|low[_ -]?latency|NVSHMEM|dispatch|combine).*" + r"(?:error|failed|failure|exception|assert|timeout|invalid|illegal)|" + r"(?:error|failed|failure|exception|assert|timeout|invalid|illegal).*" + r"(?:DeepEP|deep_ep|low[_ -]?latency|NVSHMEM|dispatch|combine)", + re.IGNORECASE, +) +DEEPEP_INSTALL_COMMAND = ( + "cd /sglang && FORCE_REBUILD_DEEPEP=1 bash scripts/ci/cuda/ci_install_deepep.sh" +) +DEEPEP_DRIVER_CLEANUP_COMMAND = ( + "apt-get purge -y nvidia-dkms-580 nvidia-kernel-common-580 " + "nvidia-kernel-source-580 nvidia-firmware-580 nvidia-modprobe && " + "rm -rf /var/lib/dkms/nvidia*" +) + + +def build_image(*, install_deepep: bool) -> modal.Image: + result = ( + 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", + ) + ) + if install_deepep: + result = result.run_commands( + DEEPEP_INSTALL_COMMAND, + DEEPEP_DRIVER_CLEANUP_COMMAND, + ) + return result.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), + "EXPECTED_GRAPH_COUNT": str(EXPECTED_GRAPH_COUNT), + **SERVE_ENV, + } + ) + + +image = build_image(install_deepep=DEEPEP_ENABLED) + +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(prompt: str) -> 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 _generate_batch() -> list[list[int]]: + request = urllib.request.Request( + f"http://127.0.0.1:{PORT}/generate", + data=json.dumps( + { + "text": CONCURRENT_PROMPTS, + "return_logprob": True, + "sampling_params": { + "temperature": 0.0, + "max_new_tokens": BATCH_MAX_NEW_TOKENS, + }, + } + ).encode(), + headers={"Content-Type": "application/json"}, + ) + with urllib.request.urlopen(request, timeout=120) as response: + result = json.loads(response.read()) + if not isinstance(result, list) or len(result) != len(CONCURRENT_PROMPTS): + raise RuntimeError(f"SGLang returned {type(result).__name__} for a batched request") + token_ids = [] + for item in result: + output_logprobs = item.get("meta_info", {}).get("output_token_logprobs") + if not output_logprobs: + raise RuntimeError("SGLang batched response omitted output token logprobs") + token_ids.append([int(entry[1]) for entry in output_logprobs]) + return token_ids + + +def _spec_accept_lengths_from_info(server_info: dict) -> list[float]: + return [ + float(state["avg_spec_accept_length"]) + for state in server_info.get("internal_states", []) + if state.get("avg_spec_accept_length") is not None + ] + + +def _spec_accept_lengths() -> list[float]: + with urllib.request.urlopen(f"http://127.0.0.1:{PORT}/server_info", timeout=30) as response: + server_info = json.loads(response.read()) + return _spec_accept_lengths_from_info(server_info) + + +def _triton_cache_dir(run_root: Path, phase: str) -> Path: + cache_phase = "save2" if phase == "load" else phase + return run_root / f"triton_cache_{cache_phase}" + + +def _serve_env_for_phase(phase: str) -> dict[str, str]: + serve_env = dict(SERVE_ENV) + if ( + phase == "baseline" + and serve_env.get("SGLANG_CUDA_GRAPH_BACKEND_PREFILL") == "full" + and not SPECULATIVE_ALGORITHM + ): + serve_env["SGLANG_CUDA_GRAPH_BACKEND_PREFILL"] = "disabled" + serve_env.pop("SGLANG_CUDA_GRAPH_BS_PREFILL", None) + return serve_env + + +def _launch( + phase: str, + mode: str | None, + log_path: str, + run_root: Path, +) -> tuple[subprocess.Popen, IO[str]]: + env = dict(os.environ) + env["SGLANG_MODEL"] = MODEL + for name in SERVE_ENV_NAMES: + env.pop(name, None) + env.update(_serve_env_for_phase(phase)) + env["NCCL_DEBUG"] = "INFO" + env["TRITON_CACHE_DIR"] = str(_triton_cache_dir(run_root, phase)) + + # 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 _archive_collective_counts(workspace: Path) -> dict[str, dict[str, int]]: + counts = {} + for rank in range(TP_SIZE): + rank_counts = { + "symmetric_all_reduce": 0, + "symmetric_all_gather": 0, + "nccl": 0, + } + rank_dir = workspace / f"rank_{rank}" + for path in rank_dir.glob("graph_[0-9]*.json"): + graph = json.loads(path.read_text()) + for node in graph.get("nodes", []): + function_name = node.get("params", {}).get("function_name", "") + if "ncclDevKernel" in function_name: + rank_counts["nccl"] += 1 + elif "CUDASymmetricMemoryOps" in function_name and "all_reduce" in function_name: + rank_counts["symmetric_all_reduce"] += 1 + elif "CUDASymmetricMemoryOps" in function_name and "all_gather" in function_name: + rank_counts["symmetric_all_gather"] += 1 + counts[str(rank)] = rank_counts + return counts + + +def _rank_offsets_valid(offsets: dict[str, int], *, deepep: bool) -> bool: + return deepep or len(set(offsets.values())) == 1 + + +def _collective_counts_valid( + counts_by_rank: dict[str, dict[str, int]], + *, + deepep: bool, +) -> bool: + if deepep: + return all(counts["symmetric_all_reduce"] > 0 for counts in counts_by_rank.values()) + return all( + counts["symmetric_all_reduce"] > 0 + and counts["symmetric_all_gather"] > 0 + and counts["nccl"] == 0 + for counts in counts_by_rank.values() + ) + + +def _deepep_kernel_kind(function_name: str) -> str | None: + lowered = function_name.lower() + if "internode_ll" not in lowered: + return None + if "dispatch" in lowered: + return "dispatch" + if "combine" in lowered: + return "combine" + return None + + +def _deepep_graph_observations(workspace: Path) -> dict[str, dict]: + observations = {} + for rank in range(TP_SIZE): + rank_dir = workspace / f"rank_{rank}" + catalog_path = rank_dir / "sglang_graph_catalog.json" + if not catalog_path.exists(): + continue + catalog = json.loads(catalog_path.read_text()) + sessions = [] + total_dispatch = 0 + total_combine = 0 + for session in catalog.get("sessions", []): + dispatch_count = 0 + combine_count = 0 + for graph_record in session.get("graphs", []): + graph_path = rank_dir / graph_record["filename"] + if not graph_path.exists(): + continue + graph = json.loads(graph_path.read_text()) + for node in graph.get("nodes", []): + if node.get("type") != "KernelNode": + continue + kind = _deepep_kernel_kind(node.get("params", {}).get("function_name", "")) + dispatch_count += kind == "dispatch" + combine_count += kind == "combine" + sessions.append( + ( + session["runner"].get("phase"), + session.get("session_index"), + dispatch_count, + combine_count, + ) + ) + total_dispatch += dispatch_count + total_combine += combine_count + observations[str(rank)] = { + "dispatch_kernel_count": total_dispatch, + "combine_kernel_count": total_combine, + "sessions": sessions, + } + return observations + + +def _schema_contains_key(value: object, key: str) -> bool: + if isinstance(value, dict): + return key in value or any(_schema_contains_key(item, key) for item in value.values()) + if isinstance(value, list): + return any(_schema_contains_key(item, key) for item in value) + return False + + +def _catalog_observations(workspace: Path) -> dict[str, dict]: + observations = {} + for rank in range(TP_SIZE): + catalog_path = workspace / f"rank_{rank}" / "sglang_graph_catalog.json" + if not catalog_path.exists(): + continue + catalog = json.loads(catalog_path.read_text()) + sessions = [] + hidden_states_schema_count = 0 + graph_count = 0 + for session in catalog.get("sessions", []): + runner = session["runner"] + graphs = session.get("graphs", []) + session_graph_count = len(graphs) + sessions.append( + ( + runner.get("phase"), + runner.get("role"), + runner.get("forward_mode"), + runner.get("step"), + session_graph_count, + ) + ) + graph_count += session_graph_count + hidden_states_schema_count += sum( + _schema_contains_key(graph.get("output_schema"), "hidden_states") + for graph in graphs + ) + observations[str(rank)] = { + "sessions": sessions, + "graph_count": graph_count, + "hidden_states_schema_count": hidden_states_schema_count, + } + return observations + + +def _expected_graph_counts( + catalogs: dict[str, dict], + *, + expected_ranks: set[str], + speculative: bool, + full_prefill: bool, + deepep: bool = False, +) -> dict[str, int]: + if speculative or deepep or full_prefill: + return { + rank: int(catalogs[rank]["graph_count"]) + for rank in sorted(expected_ranks) + if rank in catalogs + } + return {rank: EXPECTED_GRAPH_COUNT for rank in sorted(expected_ranks)} + + +def _deterministic_outputs_match( + baseline_outputs: list[str], + save_outputs: list[str], + save2_outputs: list[str], + loaded_outputs: list[str], + *, + speculative: bool, +) -> bool: + if speculative: + return save_outputs == save2_outputs == loaded_outputs + return baseline_outputs == save_outputs == save2_outputs == loaded_outputs + + +def _is_optional_local_ibgda_probe(line: str) -> bool: + lowered = line.lower() + return ( + SERVE_ENV.get("NVSHMEM_REMOTE_TRANSPORT", "").lower() == "none" + and "nvshmem" in lowered + and ( + "init failed for transport: IBGDA" in line + or ("ibgda.cpp" in lowered and "get_device_list failed" in line) + ) + ) + + +def _speculative_catalog_checks( + catalogs: dict[str, dict], + expected_ranks: set[str], +) -> dict[str, bool]: + def every_rank_has(predicate) -> bool: + return set(catalogs) == expected_ranks and all( + any(predicate(session) for session in catalogs[rank]["sessions"]) + for rank in expected_ranks + ) + + return { + "target_verify_session_each_rank": every_rank_has( + lambda session: session[1] == "target" and session[2] == "TARGET_VERIFY" + ), + "draft_session_each_rank": every_rank_has(lambda session: session[1] == "draft"), + "prefill_session_each_rank": every_rank_has(lambda session: session[0] == "prefill"), + "hidden_states_schema_each_rank": set(catalogs) == expected_ranks + and all(catalogs[rank]["hidden_states_schema_count"] > 0 for rank in expected_ranks), + } + + +def _full_prefill_catalog_checks( + catalogs: dict[str, dict], + expected_ranks: set[str], +) -> dict[str, bool]: + def every_rank_has_phase(phase: str) -> bool: + return set(catalogs) == expected_ranks and all( + any(session[0] == phase and session[4] > 0 for session in catalogs[rank]["sessions"]) + for rank in expected_ranks + ) + + return { + "prefill_session_each_rank": every_rank_has_phase("prefill"), + "decode_session_each_rank": every_rank_has_phase("decode"), + } + + +def _deepep_catalog_checks( + observations: dict[str, dict], + expected_ranks: set[str], +) -> dict[str, bool]: + def every_rank_has_phase(phase: str) -> bool: + return set(observations) == expected_ranks and all( + any( + session_phase == phase and dispatch_count > 0 and combine_count > 0 + for session_phase, _session_index, dispatch_count, combine_count in ( + observations[rank]["sessions"] + ) + ) + for rank in expected_ranks + ) + + return { + "deepep_dispatch_combine_kernels_each_rank": set(observations) == expected_ranks + and all( + observations[rank]["dispatch_kernel_count"] > 0 + and observations[rank]["combine_kernel_count"] > 0 + for rank in expected_ranks + ), + "deepep_decode_session_kernels_each_rank": every_rank_has_phase("decode"), + "deepep_prefill_session_kernels_each_rank": every_rank_has_phase("prefill"), + } + + +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_graph_payload(path: Path) -> dict: + 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 + return graph + + +def _semantic_graph_sha256(path: Path) -> str: + graph = _semantic_graph_payload(path) + canonical = json.dumps(graph, sort_keys=True, separators=(",", ":")).encode() + return hashlib.sha256(canonical).hexdigest() + + +def _json_difference_paths( + left, + right, + *, + path: str = "$", + limit: int = 64, +) -> list[str]: + differences: list[str] = [] + + def compare(left_value, right_value, current_path: str) -> None: + if len(differences) >= limit: + return + if type(left_value) is not type(right_value): + differences.append(current_path) + return + if isinstance(left_value, dict): + for key in sorted(set(left_value) | set(right_value)): + child_path = f"{current_path}.{key}" + if key not in left_value or key not in right_value: + differences.append(child_path) + else: + compare(left_value[key], right_value[key], child_path) + if len(differences) >= limit: + return + return + if isinstance(left_value, list): + if len(left_value) != len(right_value): + differences.append(f"{current_path}.length") + for index, (left_item, right_item) in enumerate(zip(left_value, right_value)): + compare(left_item, right_item, f"{current_path}[{index}]") + if len(differences) >= limit: + return + return + if left_value != right_value: + differences.append(current_path) + + compare(left, right, path) + return differences + + +def _semantic_graph_difference_paths(left: Path, right: Path) -> list[str]: + return _json_difference_paths( + _semantic_graph_payload(left), + _semantic_graph_payload(right), + ) + + +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_catalog.json", + ) + if path.exists() + ) + fingerprints[str(rank)] = { + path.relative_to(rank_dir).as_posix(): ( + _semantic_graph_sha256(path) + if path.name.startswith("graph_") and path.name != "graph_manifest.json" + else _sha256(path) + ) + for path in sorted(files) + } + return fingerprints + + +def _snapshot_graph_jsons(workspace: Path, destination: Path) -> None: + shutil.rmtree(destination, ignore_errors=True) + for rank in range(TP_SIZE): + rank_dir = workspace / f"rank_{rank}" + graph_paths = sorted(rank_dir.glob("graph_*.json")) + if not graph_paths: + continue + destination_rank_dir = destination / f"rank_{rank}" + destination_rank_dir.mkdir(parents=True, exist_ok=True) + for graph_path in graph_paths: + shutil.copy2(graph_path, destination_rank_dir / graph_path.name) + + +def _archive_graph_difference_paths( + left_workspace: Path, + right_workspace: Path, +) -> dict[str, dict[str, list[str]]]: + differences = {} + for rank in range(TP_SIZE): + left_rank_dir = left_workspace / f"rank_{rank}" + right_rank_dir = right_workspace / f"rank_{rank}" + rank_differences = {} + for left_path in sorted(left_rank_dir.glob("graph_*.json")): + right_path = right_rank_dir / left_path.name + if not right_path.exists(): + rank_differences[left_path.name] = ["$.__file_missing__"] + elif _semantic_graph_sha256(left_path) != _semantic_graph_sha256(right_path): + rank_differences[left_path.name] = _semantic_graph_difference_paths( + left_path, + right_path, + ) + if rank_differences: + differences[str(rank)] = rank_differences + return differences + + +def _scan_log(log_path: str) -> dict: + text = Path(log_path).read_text(errors="replace") + errors = find_runtime_errors(text) + deepep_transport_patch_ranks = set() + deepep_bootstrap_positions: dict[str, list[int]] = {} + deepep_graph_work_positions: dict[str, list[int]] = {} + deepep_loaded_nvshmem_module_counts: dict[str, list[int]] = {} + deepep_errors = [] + deepep_configuration_lines = [] + deepep_optional_transport_probes = [] + for line_index, line in enumerate(text.splitlines()): + rank_match = TP_LOG_RANK_PATTERN.search(line) + rank = rank_match.group("rank") if rank_match else None + if rank is not None: + if "SGLang DeepEP transport patch installed" in line: + deepep_transport_patch_ranks.add(rank) + if "SGLang DeepEP buffer bootstrap ready" in line: + deepep_bootstrap_positions.setdefault(rank, []).append(line_index) + if ( + "Started " in line and "SGLang-main graph builds" in line + ) or "Saved SGLang-main CUDA graph" in line: + deepep_graph_work_positions.setdefault(rank, []).append(line_index) + nvshmem_match = NVSHMEM_MODULE_COUNT_PATTERN.search(line) + if nvshmem_match: + deepep_loaded_nvshmem_module_counts.setdefault(rank, []).append( + int(nvshmem_match.group("count")) + ) + if DEEPEP_ERROR_PATTERN.search(line): + if "server_args=ServerArgs(" in line: + deepep_configuration_lines.append(line) + elif _is_optional_local_ibgda_probe(line): + deepep_optional_transport_probes.append(line) + else: + deepep_errors.append(line) + deepep_bootstrap_before_graph_builds = { + rank: bool(deepep_bootstrap_positions.get(rank)) + and min(deepep_bootstrap_positions[rank]) < min(graph_positions) + for rank, graph_positions in deepep_graph_work_positions.items() + } + load_offsets = { + match.group("rank"): int(match.group("offset")) + for match in LOAD_OFFSET_PATTERN.finditer(text) + } + loaded_graphs: dict[str, int] = {} + for match in LOADED_GRAPHS_PATTERN.finditer(text): + rank = match.group("rank") + loaded_graphs[rank] = loaded_graphs.get(rank, 0) + int(match.group("count")) + 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 = [ + int(match.group("batch_size")) for match in GRAPH_REPLAY_PATTERN.finditer(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, + "native_capture_progress_count": text.count("Capturing batches"), + "torch_symm_mem_observed": ("communication backend=torch_symmetric_memory" in text), + "multimem_all_gather_observed": ("Multimem all-gather buffer address=" in text), + "saved_graph_log_count": text.count("[Foundry] Saved SGLang CUDA graph") + + text.count("[Foundry] Saved SGLang-main CUDA graph"), + "deepep_transport_patch_ranks": sorted(deepep_transport_patch_ranks), + "deepep_bootstrap_before_graph_builds": deepep_bootstrap_before_graph_builds, + "deepep_loaded_nvshmem_module_counts": deepep_loaded_nvshmem_module_counts, + "deepep_configuration_lines": deepep_configuration_lines, + "deepep_optional_transport_probes": deepep_optional_transport_probes, + "deepep_errors": deepep_errors, + } + + +@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) + + triton_cache_dir = _triton_cache_dir(run_root, phase) + if phase != "load": + shutil.rmtree(triton_cache_dir, ignore_errors=True) + triton_cache_dir.mkdir(parents=True, exist_ok=True) + + started_at = time.time() + process, log_file = _launch(phase, 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 {"baseline", "load"}: + result["batched_outputs"] = _generate_batch() + if SPECULATIVE_ALGORITHM: + result["avg_spec_accept_lengths"] = _spec_accept_lengths() + 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) + archive_volume.commit() + + 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_collective_counts"] = _archive_collective_counts(workspace) + result["archive_fingerprints"] = _archive_fingerprints(workspace) + if KEEP_ARTIFACTS: + diagnostics_root = run_root / "graph_diagnostics" + phase_snapshot = diagnostics_root / phase + _snapshot_graph_jsons(workspace, phase_snapshot) + if phase == "save2": + result["archive_graph_difference_paths"] = _archive_graph_difference_paths( + diagnostics_root / "save", + phase_snapshot, + ) + archive_volume.commit() + if DEEPEP_ENABLED: + result["deepep_graph_observations"] = _deepep_graph_observations(workspace) + if phase in {"save", "save2", "load"}: + result["rank_catalogs"] = _catalog_observations(workspace) + + print(json.dumps(result, sort_keys=True)) + return result + + +@app.function( + image=image, + volumes={ARCHIVE_ROOT: archive_volume}, +) +def cleanup_run(run_id: str) -> None: + shutil.rmtree(Path(RUNS_ROOT) / run_id, ignore_errors=True) + archive_volume.commit() + + +@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", "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 = { + "graph_manifest.json", + "fatbin_image_packed.img", + "final_alloc_offset.json", + "sglang_graph_catalog.json", + } + 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_batched_outputs = results["baseline"].get("batched_outputs", []) + loaded_batched_outputs = results["load"].get("batched_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", {}) + save_collective_counts = results["save"].get("archive_collective_counts", {}) + save2_collective_counts = results["save2"].get("archive_collective_counts", {}) + loaded_graphs = results["load"].get("loaded_graphs", {}) + replay_batch_sizes = results["load"].get("graph_replay_batch_sizes", []) + save_catalogs = results["save"].get("rank_catalogs", {}) + save2_catalogs = results["save2"].get("rank_catalogs", {}) + load_catalogs = results["load"].get("rank_catalogs", {}) + expected_graph_counts = _expected_graph_counts( + save_catalogs, + expected_ranks=expected_ranks, + speculative=bool(SPECULATIVE_ALGORITHM), + full_prefill=FULL_PREFILL_ENABLED, + deepep=DEEPEP_ENABLED, + ) + + checks = { + "deterministic_outputs": len(baseline_outputs) == len(PROMPTS) + and all(output.strip() for output in baseline_outputs) + and _deterministic_outputs_match( + baseline_outputs, + save_outputs, + save2_outputs, + loaded_outputs, + speculative=bool(SPECULATIVE_ALGORITHM), + ), + "batched_outputs_nonempty": len(baseline_batched_outputs) == len(CONCURRENT_PROMPTS) + and all(output for output in baseline_batched_outputs) + and len(loaded_batched_outputs) == len(CONCURRENT_PROMPTS) + and all(output for output in loaded_batched_outputs), + "batched_outputs_discriminating": len( + {tuple(output) for output in baseline_batched_outputs} + ) + > 1 + and len({tuple(output) for output in loaded_batched_outputs}) > 1, + "batched_outputs_match": baseline_batched_outputs == loaded_batched_outputs, + "save_offsets_present": set(save_offsets) == expected_ranks and all(save_offsets.values()), + "save_rank_offsets_symmetric": _rank_offsets_valid( + save_offsets, + deepep=DEEPEP_ENABLED, + ), + "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(rank_fingerprints) + and sum( + re.fullmatch(r"graph_\d+.*\.json", name) is not None for name in rank_fingerprints + ) + == expected_graph_counts.get(rank) + for rank, rank_fingerprints in save_fingerprints.items() + ), + "archive_fingerprints_reproducible": save_fingerprints == save2_fingerprints, + "archives_complete": set(graph_counts) == expected_ranks + and save_graph_counts == graph_counts + and graph_counts == expected_graph_counts + and all(count > 0 for count in graph_counts.values()), + "catalogs_present_and_reproducible": set(save_catalogs) == expected_ranks + and save_catalogs == save2_catalogs == load_catalogs, + "symmetric_collectives_each_rank": save_collective_counts == save2_collective_counts + and set(save_collective_counts) == expected_ranks + and _collective_counts_valid( + save_collective_counts, + deepep=DEEPEP_ENABLED, + ), + "graphs_restored_each_rank": loaded_graphs == graph_counts, + "restored_graph_replay_observed": 1 in replay_batch_sizes + and any(batch_size > 1 for batch_size in replay_batch_sizes), + "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"] + and results["save"]["multimem_all_gather_observed"] + and results["load"]["multimem_all_gather_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), + } + if FULL_PREFILL_ENABLED: + checks.update(_full_prefill_catalog_checks(save_catalogs, expected_ranks)) + if SPECULATIVE_ALGORITHM: + checks.update(_speculative_catalog_checks(save_catalogs, expected_ranks)) + checks["exact_seeded_batched_token_ids"] = ( + len(baseline_batched_outputs) == len(CONCURRENT_PROMPTS) + and all(baseline_batched_outputs) + and baseline_batched_outputs == loaded_batched_outputs + ) + checks["nontrivial_speculative_acceptance"] = all( + results[phase].get("avg_spec_accept_lengths") + and all( + accept_length > 1.0 for accept_length in results[phase]["avg_spec_accept_lengths"] + ) + for phase in phases + ) + if DEEPEP_ENABLED: + deepep_observations = results["save"].get("deepep_graph_observations", {}) + checks.update(_deepep_catalog_checks(deepep_observations, expected_ranks)) + checks["deepep_graph_nodes_reproducible"] = deepep_observations == results["save2"].get( + "deepep_graph_observations", {} + ) + checks["deepep_transport_patch_each_rank"] = all( + set(results[phase]["deepep_transport_patch_ranks"]) == expected_ranks + for phase in ("save", "save2", "load") + ) + checks["deepep_bootstrap_precedes_graph_work_each_rank"] = all( + results[phase]["deepep_bootstrap_before_graph_builds"] + == {rank: True for rank in expected_ranks} + for phase in ("save", "save2", "load") + ) + load_nvshmem_counts = results["load"]["deepep_loaded_nvshmem_module_counts"] + checks["deepep_loaded_nvshmem_modules_each_rank"] = set( + load_nvshmem_counts + ) == expected_ranks and all( + any(count > 0 for count in load_nvshmem_counts[rank]) for rank in expected_ranks + ) + checks["no_deepep_runtime_errors"] = not any( + results[phase]["deepep_errors"] for phase in phases + ) + checks["deepep_load_has_no_save_records"] = results["load"]["saved_graph_log_count"] == 0 + + report = { + "run_id": run_id, + "artifacts_retained": KEEP_ARTIFACTS, + "foundry_ref": FOUNDRY_REF, + "sglang_ref": SGLANG_REF, + "speculative_algorithm": SPECULATIVE_ALGORITHM, + "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)}") + if not KEEP_ARTIFACTS: + cleanup_run.remote(run_id) diff --git a/tests/run_deepep_matrix.sh b/tests/run_deepep_matrix.sh new file mode 100755 index 00000000..d7524ccb --- /dev/null +++ b/tests/run_deepep_matrix.sh @@ -0,0 +1,76 @@ +#!/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 + +SCRIPT_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd) +ROOT=$(cd "$SCRIPT_DIR/.." && pwd) +PY=${PYTHON:-python3} +TEST=$ROOT/tests/test_deepep_fabric.py +NVSHMEM_SO=${NVSHMEM_HOST_PATH:?set NVSHMEM_HOST_PATH to libnvshmem_host.so} +HOOK_SO=$($PY -c "import foundry.ops, pathlib; print(pathlib.Path(foundry.ops.__file__).parent / 'libcuda_hook.so')") + +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:-$ROOT/tests/deepep_matrix_work}/$CASE +LOGDIR=${DEEPEP_MATRIX_LOG_ROOT:-$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..7df0c037 100644 --- a/tests/test_deepep_fabric.py +++ b/tests/test_deepep_fabric.py @@ -13,6 +13,8 @@ - Multiple GPUs for distributed testing """ +import importlib.util +import inspect import os import shutil import subprocess @@ -23,17 +25,36 @@ import torch import torch.distributed as dist -# Use lower base address (matching working deepep_args example) -# Note: 0x60000000000 (6TB) NOT 0x600000000000 (96TB) - the extra zero matters! -BASE_ADDR = 0x600000000000 # 6TB +# 96 TiB, matching the deterministic Foundry integration region. +BASE_ADDR = 0x600000000000 REGION_SIZE_STR = "512GB" ARCHIVE_DIR = "deepep_fabric_archive" HOOK_ARCHIVE_DIR = "hook_archive" -def _get_hook_so_path(): - import importlib.util +def _deepep_buffer_kwargs( + buffer_type, + kwargs: dict, + *, + use_fabric: bool, +) -> dict: + result = dict(kwargs) + if "use_fabric" in inspect.signature(buffer_type.__init__).parameters: + result["use_fabric"] = use_fabric + elif use_fabric: + print( + "[TEST] DeepEP Buffer lacks use_fabric; falling back to pinned " + "DeepEP NVL/IPC transport while retaining semantic payload checks", + flush=True, + ) + return result + + +def _deepep_topk_dtype(deep_ep_module): + return getattr(deep_ep_module, "topk_idx_t", torch.int64) + +def _get_hook_so_path(): spec = importlib.util.find_spec("foundry.ops") if not spec or not spec.origin: raise RuntimeError("foundry.ops not found; ensure setup.py develop/pip install completed") @@ -47,6 +68,38 @@ def _get_hook_so_path(): return str(hook_so_path) +def _get_nvshmem_host_path() -> str: + spec = importlib.util.find_spec("nvidia.nvshmem") + if spec is None: + raise RuntimeError("nvidia.nvshmem is not installed") + roots = list(spec.submodule_search_locations or []) + if spec.origin: + roots.append(str(Path(spec.origin).parent)) + for root in roots: + for name in ("libnvshmem_host.so.3", "libnvshmem_host.so"): + candidate = Path(root) / "lib" / name + if candidate.exists(): + return str(candidate) + raise RuntimeError("libnvshmem_host.so not found in the installed nvidia.nvshmem wheel") + + +def test_nvshmem_host_path_resolves_installed_wheel( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + library = tmp_path / "lib" / "libnvshmem_host.so.3" + library.parent.mkdir() + library.touch() + + class Spec: + submodule_search_locations = [str(tmp_path)] + origin = None + + monkeypatch.setattr(importlib.util, "find_spec", lambda _name: Spec()) + + assert _get_nvshmem_host_path() == str(library) + + def _init_dist(local_rank: int, num_local_ranks: int): """Initialize distributed environment for multi-GPU testing.""" ip = os.getenv("MASTER_ADDR", "127.0.0.1") @@ -77,23 +130,49 @@ 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 ): """Create deterministic input tensors with traceable patterns. - Pattern for x: each token i has value (rank * 1000 + i) in all hidden dims + Pattern for x: each token i has value (rank * num_tokens + i) in all hidden + dims. This stays injective and exactly representable in BF16. Pattern for topk_idx: deterministic expert selection based on token id """ import deep_ep - # x[token_i, :] = rank * 1000 + token_i (easy to trace which rank/token) + # x[token_i, :] uniquely encodes source rank and token. x = torch.zeros((num_tokens, hidden), dtype=torch.bfloat16, device="cuda") for i in range(num_tokens): - x[i, :] = float(rank * 1000 + i) + x[i, :] = float(rank * num_tokens + i) # Deterministic topk_idx: token i selects experts [i % num_experts, (i+1) % num_experts, ...] - topk_idx = torch.zeros((num_tokens, num_topk), dtype=deep_ep.topk_idx_t, device="cuda") + topk_idx = torch.zeros( + (num_tokens, num_topk), + dtype=_deepep_topk_dtype(deep_ep), + device="cuda", + ) for i in range(num_tokens): for k in range(num_topk): topk_idx[i, k] = (i + k) % num_experts @@ -148,9 +227,7 @@ def _print_buffer_info(buffer, rank: int, prefix: str): def _verify_dispatch_result( recv_x, - recv_topk_idx, - recv_src_idx, - num_recv, + recv_count, rank: int, num_ranks: int, num_tokens: int, @@ -163,70 +240,44 @@ 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, - ) - - print(f"[Rank {rank}] {prefix}: Verification - num_recv = {actual_num_recv}", flush=True) print( f"[Rank {rank}] {prefix}: Verification - local experts range = [{local_expert_start}, {local_expert_end})", flush=True, ) - # recv_x shape is [num_local_experts, max_tokens_per_expert, hidden] - # Print shape info for debugging - if recv_x is not None: - print(f"[Rank {rank}] {prefix}: recv_x shape = {recv_x.shape}", flush=True) - - # Check some values in recv_x to verify the pattern - if recv_x is not None and recv_x.numel() > 0: - # Check first expert's first few tokens - num_experts_in_recv = recv_x.shape[0] - tokens_per_expert = recv_x.shape[1] - check_experts = min(2, num_experts_in_recv) - check_tokens = min(3, tokens_per_expert) - - print( - f"[Rank {rank}] {prefix}: Checking first {check_experts} experts, {check_tokens} tokens each:", - flush=True, + assert recv_x is not None, f"{prefix}: recv_x is None" + assert recv_x.ndim == 3, f"{prefix}: expected rank-3 recv_x, got {recv_x.shape}" + assert recv_x.shape[0] == num_local_experts, ( + f"{prefix}: local expert count differs: {recv_x.shape[0]} != {num_local_experts}" + ) + assert recv_x.shape[2] == hidden, ( + f"{prefix}: hidden size differs: {recv_x.shape[2]} != {hidden}" + ) + assert recv_count.shape == (num_local_experts,), ( + f"{prefix}: recv_count shape differs: {recv_count.shape}" + ) + print(f"[Rank {rank}] {prefix}: recv_x shape = {recv_x.shape}", flush=True) + + for local_expert, global_expert in enumerate(range(local_expert_start, local_expert_end)): + expected_values = [ + source_rank * num_tokens + token + for source_rank in range(num_ranks) + for token in range(num_tokens) + for topk_slot in range(2) + if (token + topk_slot) % num_experts == global_expert + ] + expected = torch.tensor(expected_values, dtype=recv_x.dtype, device="cpu").sort().values + actual_count = int(recv_count[local_expert].item()) + assert actual_count == len(expected_values), ( + f"{prefix}: expert {global_expert} received {actual_count}, " + f"expected {len(expected_values)}" ) - - 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, - ) - else: - print(f"[Rank {rank}] {prefix}: recv_x is empty or None", flush=True) - - return True + actual_rows = recv_x[local_expert, :actual_count] + assert torch.equal(actual_rows, actual_rows[:, :1].expand_as(actual_rows)), ( + f"{prefix}: expert {global_expert} has nonuniform hidden rows" + ) + actual = actual_rows[:, 0].cpu().sort().values + torch.testing.assert_close(actual, expected, rtol=0, atol=0) def _run_save(local_rank: int, num_processes: int): @@ -263,6 +314,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) @@ -285,13 +338,18 @@ def _run_save(local_rank: int, num_processes: int): 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, - allow_nvlink_for_low_latency_mode=True, - explicitly_destroy=True, - use_fabric=use_fabric, + **_deepep_buffer_kwargs( + deep_ep.Buffer, + { + "num_nvl_bytes": num_nvl_bytes, + "num_rdma_bytes": num_rdma_bytes, + "low_latency_mode": True, + "num_qps_per_rank": num_local_experts, + "allow_nvlink_for_low_latency_mode": True, + "explicitly_destroy": True, + }, + use_fabric=use_fabric, + ), ) # Print buffer info including addresses @@ -304,11 +362,11 @@ def _run_save(local_rank: int, num_processes: int): print(f"[Rank {rank}] SAVE: Input tensors:", flush=True) print(f"[Rank {rank}] SAVE: x address = {hex(x.data_ptr())}, shape = {x.shape}", flush=True) print( - f"[Rank {rank}] SAVE: x[0, :5] = {x[0, :5].tolist()} (expect {rank * 1000 + 0})", + f"[Rank {rank}] SAVE: x[0, :5] = {x[0, :5].tolist()} (expect {rank * num_tokens})", flush=True, ) print( - f"[Rank {rank}] SAVE: x[1, :5] = {x[1, :5].tolist()} (expect {rank * 1000 + 1})", + f"[Rank {rank}] SAVE: x[1, :5] = {x[1, :5].tolist()} (expect {rank * num_tokens + 1})", flush=True, ) print(f"[Rank {rank}] SAVE: topk_idx address = {hex(topk_idx.data_ptr())}", flush=True) @@ -331,7 +389,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, recv_count, _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 +399,12 @@ def _run_save(local_rank: int, num_processes: int): f"[Rank {rank}] SAVE: recv_x shape = {recv_x.shape if recv_x is not None else 'None'}", flush=True, ) - print(f"[Rank {rank}] SAVE: num_recv = {num_recv}", flush=True) + print(f"[Rank {rank}] SAVE: recv_count = {recv_count.tolist()}", flush=True) # Verify results _verify_dispatch_result( recv_x, - recv_topk_idx, - recv_src_idx, - num_recv, + recv_count, rank, num_ranks, num_tokens, @@ -413,8 +469,16 @@ 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_dispatch_result( + graph_result[0], + graph_result[1], + rank, + num_ranks, + num_tokens, + hidden, + num_experts, + "SAVE-replay", + ) print(f"[Rank {rank}] SAVE: Graph replay completed", flush=True) # Save graph and fatbins to per-rank archive @@ -487,11 +551,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,19 +574,28 @@ 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_rdma_bytes=num_rdma_bytes, - low_latency_mode=True, - num_qps_per_rank=num_local_experts, - allow_nvlink_for_low_latency_mode=True, - explicitly_destroy=True, - use_fabric=use_fabric, + **_deepep_buffer_kwargs( + deep_ep.Buffer, + { + "num_nvl_bytes": num_nvl_bytes, + "num_rdma_bytes": num_rdma_bytes, + "low_latency_mode": True, + "num_qps_per_rank": num_local_experts, + "allow_nvlink_for_low_latency_mode": True, + "explicitly_destroy": True, + }, + use_fabric=use_fabric, + ), ) # Print buffer info including addresses @@ -530,11 +614,11 @@ def _run_load(local_rank: int, num_processes: int): print(f"[Rank {rank}] LOAD: Input tensors:", flush=True) print(f"[Rank {rank}] LOAD: x address = {hex(x.data_ptr())}, shape = {x.shape}", flush=True) print( - f"[Rank {rank}] LOAD: x[0, :5] = {x[0, :5].tolist()} (expect {rank * 1000 + 0})", + f"[Rank {rank}] LOAD: x[0, :5] = {x[0, :5].tolist()} (expect {rank * num_tokens})", flush=True, ) print( - f"[Rank {rank}] LOAD: x[1, :5] = {x[1, :5].tolist()} (expect {rank * 1000 + 1})", + f"[Rank {rank}] LOAD: x[1, :5] = {x[1, :5].tolist()} (expect {rank * num_tokens + 1})", flush=True, ) print(f"[Rank {rank}] LOAD: topk_idx address = {hex(topk_idx.data_ptr())}", flush=True) @@ -557,7 +641,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, recv_count, _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 +651,12 @@ def _run_load(local_rank: int, num_processes: int): f"[Rank {rank}] LOAD: recv_x shape = {recv_x.shape if recv_x is not None else 'None'}", flush=True, ) - print(f"[Rank {rank}] LOAD: num_recv = {num_recv}", flush=True) + print(f"[Rank {rank}] LOAD: recv_count = {recv_count.tolist()}", flush=True) # Verify warmup results _verify_dispatch_result( recv_x, - recv_topk_idx, - recv_src_idx, - num_recv, + recv_count, rank, num_ranks, num_tokens, @@ -593,10 +675,21 @@ def _run_load(local_rank: int, num_processes: int): flush=True, ) - # Load graph from per-rank archive + # Load graph from per-rank archive. + # Default: the production path (start_graph_builds + finish_graph_loads), + # same machinery the vLLM/sglang integrations use. TEST_LOAD_API=single + # selects the legacy CUDAGraph.load path โ€” KNOWN BROKEN for DeepEP graphs: + # it does not merge root-level common_kernel_node_attrs, so factored-out + # cooperative/clusterDim attrs are dropped and the dispatch kernel faults + # with "unspecified launch failure" at replay. graph_json = os.path.join(rank_archive, "deepep_dispatch_graph.json") - print(f"[Rank {rank}] LOAD: Loading graph from {graph_json}", flush=True) - graph, output_tensors = fdry.CUDAGraph.load(graph_json) + load_api = os.environ.get("TEST_LOAD_API", "parallel") + print(f"[Rank {rank}] LOAD: Loading graph from {graph_json} (api={load_api})", flush=True) + if load_api == "single": + graph, output_tensors = fdry.CUDAGraph.load(graph_json) + else: + pending = fdry.CUDAGraph.start_graph_builds([graph_json]) + ((graph, output_tensors),) = fdry.CUDAGraph.finish_graph_loads(pending) print(f"[Rank {rank}] LOAD: Graph loaded successfully", flush=True) # Print loaded output tensor info @@ -624,7 +717,20 @@ def _run_load(local_rank: int, num_processes: int): graph.replay() torch.cuda.synchronize() - # Verify after replay - check the LOADED output tensors (not warmup recv_x) + if not isinstance(output_tensors, (list, tuple)) or len(output_tensors) < 2: + raise AssertionError("LOAD replay did not restore recv_x and recv_count") + _verify_dispatch_result( + output_tensors[0], + output_tensors[1], + rank, + num_ranks, + num_tokens, + hidden, + num_experts, + "LOAD-replay", + ) + + # Print a small sample after the exact comparison for diagnostics. # recv_x shape is [num_local_experts, max_tokens_per_expert, hidden] 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 @@ -640,8 +746,8 @@ def _run_load(local_rank: int, num_processes: int): 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 + decoded_src_rank = int(val) // num_tokens + decoded_token_id = int(val) % num_tokens print( f"[Rank {rank}] LOAD: expert[{expert_i}], token[{token_i}]: " f"value={val:.1f} (rank={decoded_src_rank}, token={decoded_token_id})", @@ -674,11 +780,13 @@ def _cleanup_archive(): def _spawn_with_preload(mode: str, num_processes: int): """Spawn subprocess with LD_PRELOAD for CGE hook.""" so_path = _get_hook_so_path() + nvshmem_host_path = _get_nvshmem_host_path() env = os.environ.copy() + preloads = [nvshmem_host_path, so_path] if env.get("LD_PRELOAD"): - env["LD_PRELOAD"] = f"{so_path}:{env['LD_PRELOAD']}" - else: - env["LD_PRELOAD"] = so_path + preloads.append(env["LD_PRELOAD"]) + env["LD_PRELOAD"] = ":".join(preloads) + print(f"[TEST] Preloading NVSHMEM host library: {nvshmem_host_path}", flush=True) # # Set NVSHMEM environment variables to disable IBGDA and remote transports # env['NVSHMEM_REMOTE_TRANSPORT'] = 'none' @@ -709,6 +817,47 @@ def _check_multi_gpu(): return torch.cuda.device_count() >= 2 +def test_buffer_kwargs_include_supported_fabric_transport() -> None: + class Buffer: + def __init__(self, group, *, use_fabric=False): + del group, use_fabric + + original = {"num_nvl_bytes": 123, "num_rdma_bytes": 456} + + assert _deepep_buffer_kwargs(Buffer, original, use_fabric=True) == { + "num_nvl_bytes": 123, + "num_rdma_bytes": 456, + "use_fabric": True, + } + + +def test_buffer_kwargs_preserve_legacy_transport( + capsys: pytest.CaptureFixture[str], +) -> None: + class Buffer: + def __init__(self, group, *, num_nvl_bytes=0, num_rdma_bytes=0): + del group, num_nvl_bytes, num_rdma_bytes + + original = {"num_nvl_bytes": 123, "num_rdma_bytes": 456} + + assert _deepep_buffer_kwargs(Buffer, original, use_fabric=True) == original + assert "falling back to pinned DeepEP NVL/IPC transport" in capsys.readouterr().out + + +def test_topk_dtype_uses_installed_deepep_contract_when_exported() -> None: + class DeepEP: + topk_idx_t = torch.int32 + + assert _deepep_topk_dtype(DeepEP) is torch.int32 + + +def test_topk_dtype_matches_pinned_deepep_int64_contract() -> None: + class DeepEP: + pass + + assert _deepep_topk_dtype(DeepEP) is torch.int64 + + @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(): diff --git a/tests/test_graph_manifest.py b/tests/test_graph_manifest.py new file mode 100644 index 00000000..4304ac7e --- /dev/null +++ b/tests/test_graph_manifest.py @@ -0,0 +1,55 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the Foundry project +"""CPU-only tests for CUDA graph manifest generation.""" + +from __future__ import annotations + +import importlib.util +import json +import sys +from pathlib import Path +from types import ModuleType + + +def _load_graph_module(monkeypatch): + torch = ModuleType("torch") + torch.Tensor = object + monkeypatch.setitem(sys.modules, "torch", torch) + + foundry = ModuleType("foundry") + foundry.__path__ = [] + monkeypatch.setitem(sys.modules, "foundry", foundry) + + ops = ModuleType("foundry.ops") + ops.CUDAGraph = object + monkeypatch.setitem(sys.modules, "foundry.ops", ops) + + graph_path = Path(__file__).parents[1] / "python" / "foundry" / "graph.py" + spec = importlib.util.spec_from_file_location("foundry.graph", graph_path) + assert spec is not None and spec.loader is not None + graph_module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(graph_module) + return graph_module + + +def test_manifest_can_retain_dependencies_for_independent_loading( + monkeypatch, + tmp_path: Path, +) -> None: + graph_module = _load_graph_module(monkeypatch) + dependencies = [{"from": 0, "to": 1}] + for index in range(2): + (tmp_path / f"graph_{index}_FULL_t{2 - index}_r1_UX_pcN.json").write_text( + json.dumps( + { + "topology_key": "shared-topology", + "nodes": [], + "dependencies": dependencies, + } + ) + ) + + graph_module.save_graph_manifest(tmp_path, strip_dependencies=False) + + for graph_path in sorted(tmp_path.glob("graph_[0-9]*.json")): + assert json.loads(graph_path.read_text())["dependencies"] == dependencies diff --git a/tests/test_rng_graph.py b/tests/test_rng_graph.py new file mode 100644 index 00000000..314f96ef --- /dev/null +++ b/tests/test_rng_graph.py @@ -0,0 +1,118 @@ +# 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" +EXPECTED_PATH = f"{ARCHIVE_DIR}/expected.pt" + + +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) + expected = [] + for _ in range(2): + graph.replay() + torch.cuda.synchronize() + expected.append(output.cpu()) + torch.save(expected, EXPECTED_PATH) + fdry.stop_allocation_region() + + +def _load(load_api: str) -> 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)) + if load_api == "single": + graph, output = fdry.CUDAGraph.load(GRAPH_PATH, pool=(0, 0)) + elif load_api == "parallel": + pending = fdry.CUDAGraph.start_graph_builds([GRAPH_PATH], pool=(0, 0)) + ((graph, output),) = fdry.CUDAGraph.finish_graph_loads(pending) + else: + raise ValueError(f"Unknown load API: {load_api}") + + graph.replay() + torch.cuda.synchronize() + first = output.cpu() + graph.replay() + torch.cuda.synchronize() + second = output.cpu() + + expected_first, expected_second = torch.load(EXPECTED_PATH, weights_only=True) + torch.testing.assert_close(first, expected_first, rtol=0, atol=0) + torch.testing.assert_close(second, expected_second, rtol=0, atol=0) + 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-single") + _spawn("load-parallel") + finally: + _cleanup() + + +if __name__ == "__main__": + if "--save" in sys.argv: + _save() + elif "--load-single" in sys.argv: + _load("single") + elif "--load-parallel" in sys.argv: + _load("parallel") + else: + test_rng_graph_replays_in_fresh_process() diff --git a/tests/test_sglang_graph_catalog.py b/tests/test_sglang_graph_catalog.py new file mode 100644 index 00000000..cbd1e99a --- /dev/null +++ b/tests/test_sglang_graph_catalog.py @@ -0,0 +1,438 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the Foundry project +"""CPU-only contracts for the SGLang graph session catalog.""" + +from __future__ import annotations + +import importlib.util +import json +import sys +from enum import Enum +from pathlib import Path +from types import ModuleType, SimpleNamespace + +import pytest + +_ARCHIVE_PATH = ( + Path(__file__).parents[1] / "python" / "foundry" / "integration" / "sglang" / "archive.py" +) + + +def _load_archive_module(): + spec = importlib.util.spec_from_file_location( + "foundry.integration.sglang.archive", + _ARCHIVE_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 + + +_ARCHIVE = _load_archive_module() +CATALOG_FILENAME = _ARCHIVE.CATALOG_FILENAME +GraphCatalog = _ARCHIVE.GraphCatalog +runner_descriptor = _ARCHIVE.runner_descriptor +shape_key_record = _ARCHIVE.shape_key_record + + +def _graph_record(capture_index: int, size: int) -> dict: + return { + "capture_index": capture_index, + "filename": f"graph_{capture_index}_prefill_target_s0_k{size}.json", + "shape_key": {"size": size, "stream_idx": None, "variant_label": None}, + "output_schema": {"kind": "tensor", "index": 0}, + } + + +def test_catalog_distinguishes_equal_shapes_in_separate_sessions(tmp_path) -> None: + catalog = GraphCatalog(tmp_path) + target = { + "class": "Target", + "phase": "decode", + "role": "target", + "forward_mode": "TARGET_VERIFY", + "step": None, + } + draft = { + "class": "Draft", + "phase": "decode", + "role": "draft", + "forward_mode": "DECODE", + "step": None, + } + + catalog.open_save_session(0, target) + catalog.append_graph( + 0, + { + "capture_index": 0, + "filename": "graph_0_decode_target_s0_k8.json", + "shape_key": {"size": 8, "stream_idx": None, "variant_label": None}, + "output_schema": {"kind": "tensor", "index": 0}, + }, + ) + catalog.open_save_session(1, draft) + catalog.append_graph( + 1, + { + "capture_index": 1, + "filename": "graph_1_decode_draft_s1_k8.json", + "shape_key": {"size": 8, "stream_idx": None, "variant_label": None}, + "output_schema": {"kind": "tuple", "items": []}, + }, + ) + + reloaded = GraphCatalog(tmp_path) + assert reloaded.claim_load_session(0, target)["graphs"][0]["capture_index"] == 0 + assert reloaded.claim_load_session(1, draft)["graphs"][0]["capture_index"] == 1 + + +def test_catalog_rejects_claim_load_session_for_missing_session(tmp_path) -> None: + catalog = GraphCatalog(tmp_path) + descriptor = { + "class": "Target", + "phase": "decode", + "role": "target", + "forward_mode": "DECODE", + "step": None, + } + with pytest.raises(RuntimeError, match="session 0 not found in catalog"): + catalog.claim_load_session(0, descriptor) + + +def test_catalog_rejects_append_graph_for_missing_session(tmp_path) -> None: + catalog = GraphCatalog(tmp_path) + with pytest.raises(RuntimeError, match="session 0 not found in catalog"): + catalog.append_graph(0, _graph_record(0, 8)) + + +def test_catalog_rejects_non_increasing_global_capture_index(tmp_path) -> None: + catalog = GraphCatalog(tmp_path) + descriptor = { + "class": "Target", + "phase": "decode", + "role": "target", + "forward_mode": "DECODE", + "step": None, + } + catalog.open_save_session(0, descriptor) + catalog.append_graph(0, _graph_record(0, 8)) + catalog.open_save_session(1, descriptor) + with pytest.raises( + RuntimeError, + match="capture_index 0 must exceed all existing graph indices", + ): + catalog.append_graph(1, _graph_record(0, 16)) + + +def test_catalog_rejects_runner_order_mismatch(tmp_path) -> None: + catalog = GraphCatalog(tmp_path) + descriptor = { + "class": "Target", + "phase": "decode", + "role": "target", + "forward_mode": "DECODE", + "step": None, + } + catalog.open_save_session(0, descriptor) + + with pytest.raises(RuntimeError, match="session 0 descriptor mismatch"): + GraphCatalog(tmp_path).claim_load_session( + 0, + {**descriptor, "forward_mode": "TARGET_VERIFY"}, + ) + + +def test_catalog_writes_valid_json_after_every_append(tmp_path) -> None: + catalog = GraphCatalog(tmp_path) + descriptor = { + "class": "Prefill", + "phase": "prefill", + "role": "target", + "forward_mode": "EXTEND", + "step": None, + } + catalog.open_save_session(0, descriptor) + json.loads((tmp_path / CATALOG_FILENAME).read_text()) + catalog.append_graph(0, _graph_record(0, 16)) + assert json.loads((tmp_path / CATALOG_FILENAME).read_text())["version"] == 1 + + +def test_catalog_rejects_unknown_version(tmp_path) -> None: + (tmp_path / CATALOG_FILENAME).write_text(json.dumps({"version": 99, "sessions": []})) + with pytest.raises(RuntimeError, match="unknown catalog version"): + GraphCatalog(tmp_path) + + +def test_catalog_rejects_non_contiguous_save_session_indices(tmp_path) -> None: + catalog = GraphCatalog(tmp_path) + descriptor = { + "class": "Target", + "phase": "decode", + "role": "target", + "forward_mode": "DECODE", + "step": None, + } + with pytest.raises(RuntimeError, match="expected session index 0"): + catalog.open_save_session(1, descriptor) + + +def test_shape_key_record() -> None: + class FakeShapeKey: + def __init__( + self, + size: int, + stream_idx: int | None = None, + variant_label: str | None = None, + ) -> None: + self.size = size + self.stream_idx = stream_idx + self.variant_label = variant_label + + assert shape_key_record(FakeShapeKey(8)) == { + "size": 8, + "stream_idx": None, + "variant_label": None, + } + assert shape_key_record(FakeShapeKey(16, stream_idx=2)) == { + "size": 16, + "stream_idx": 2, + "variant_label": None, + } + assert shape_key_record(FakeShapeKey(32, variant_label="lora")) == { + "size": 32, + "stream_idx": None, + "variant_label": "lora", + } + + +def test_runner_descriptor_distinguishes_target_and_draft() -> None: + class ForwardMode(Enum): + DECODE = "DECODE" + + class ModelRunner: + def __init__(self, *, is_draft_worker: bool) -> None: + self.is_draft_worker = is_draft_worker + + class Runner: + def __init__( + self, + *, + is_draft_worker: bool, + shape_key, + step: int | None = None, + ) -> None: + self.model_runner = ModelRunner(is_draft_worker=is_draft_worker) + self.capture_forward_mode = ForwardMode.DECODE + self.shape_key = shape_key + self.step = step + + class ShapeKey: + size = 8 + stream_idx = None + variant_label = None + + target_runner = Runner(is_draft_worker=False, shape_key=ShapeKey()) + draft_runner = Runner(is_draft_worker=True, shape_key=ShapeKey(), step=1) + target = runner_descriptor(target_runner, "prefill") + draft = runner_descriptor(draft_runner, "decode") + + assert shape_key_record(target_runner.shape_key) == shape_key_record(draft_runner.shape_key) + assert (target["phase"], target["role"], target["step"]) != ( + draft["phase"], + draft["role"], + draft["step"], + ) + assert target["role"] == "target" + assert draft["role"] == "draft" + assert target["phase"] == "prefill" + assert draft["phase"] == "decode" + assert target["step"] is None + assert draft["step"] == 1 + assert target["forward_mode"] == "DECODE" + assert target["class"].endswith("Runner") + + +def test_catalog_session_count(tmp_path) -> None: + catalog = GraphCatalog(tmp_path) + assert catalog.session_count == 0 + descriptor = { + "class": "Target", + "phase": "decode", + "role": "target", + "forward_mode": "DECODE", + "step": None, + } + catalog.open_save_session(0, descriptor) + assert catalog.session_count == 1 + catalog.open_save_session(1, descriptor) + assert catalog.session_count == 2 + + +def test_catalog_instances_reload_before_open_and_claim(tmp_path) -> None: + first = GraphCatalog(tmp_path) + second = GraphCatalog(tmp_path) + reader = GraphCatalog(tmp_path) + target = { + "class": "Target", + "phase": "decode", + "role": "target", + "forward_mode": "DECODE", + "step": None, + } + draft = { + "class": "Draft", + "phase": "decode", + "role": "draft", + "forward_mode": "DECODE", + "step": 1, + } + + first.open_save_session(0, target) + second.open_save_session(1, draft) + + assert reader.claim_load_session(1, draft)["session_index"] == 1 + + +def test_cuda_graph_extension_state_starts_at_session_zero(monkeypatch) -> None: + foundry = ModuleType("foundry") + foundry_ops = ModuleType("foundry.ops") + foundry.ops = foundry_ops + allocation_region = ModuleType("foundry.allocation_region") + allocation_region.parse_size = lambda value: value + config = ModuleType("foundry.integration.sglang.config") + config.CUDAGraphExtensionMode = Enum( + "CUDAGraphExtensionMode", + ["NONE", "SAVE", "LOAD"], + ) + config.compute_workspace_rank = lambda *args: 0 + config.get_config = lambda: None + config.get_graph_extension_mode = lambda: config.CUDAGraphExtensionMode.NONE + config.get_hook_library_path = lambda: None + config.get_nvshmem_host_path = lambda: None + + monkeypatch.setitem(sys.modules, "foundry", foundry) + monkeypatch.setitem(sys.modules, "foundry.ops", foundry_ops) + monkeypatch.setitem(sys.modules, "foundry.allocation_region", allocation_region) + monkeypatch.setitem(sys.modules, "foundry.integration.sglang.config", config) + + runtime_path = ( + Path(__file__).parents[1] / "python" / "foundry" / "integration" / "sglang" / "runtime.py" + ) + spec = importlib.util.spec_from_file_location( + "foundry.integration.sglang.runtime", + runtime_path, + ) + assert spec is not None and spec.loader is not None + runtime = importlib.util.module_from_spec(spec) + monkeypatch.setitem(sys.modules, spec.name, runtime) + spec.loader.exec_module(runtime) + + state = runtime.CUDAGraphExtensionState() + assert state.session_index == 0 + assert state.load_sessions_validated is False + + +def test_graph_extension_setup_is_process_rank_idempotent(monkeypatch, tmp_path) -> None: + foundry = ModuleType("foundry") + foundry_ops = ModuleType("foundry.ops") + foundry.ops = foundry_ops + region_calls = [] + foundry_ops.set_allocation_region = lambda base, size: region_calls.append((base, size)) + foundry_ops.set_skip_fatbin_processing = lambda enabled: None + foundry_ops.load_cuda_modules_and_libraries = lambda directory: None + + allocation_region = ModuleType("foundry.allocation_region") + allocation_region.parse_size = lambda value: int(value) + config = ModuleType("foundry.integration.sglang.config") + mode = Enum("CUDAGraphExtensionMode", ["NONE", "SAVE", "LOAD"]) + cfg = SimpleNamespace( + mode=mode.SAVE, + workspace_root=str(tmp_path), + workspace_dir=None, + base_addr=1234, + region_size="4096", + ) + config.CUDAGraphExtensionMode = mode + config.compute_workspace_rank = lambda *args: 0 + config.get_config = lambda: cfg + config.get_graph_extension_mode = lambda: mode.SAVE + config.get_hook_library_path = lambda: None + config.get_nvshmem_host_path = lambda: None + + monkeypatch.setitem(sys.modules, "foundry", foundry) + monkeypatch.setitem(sys.modules, "foundry.ops", foundry_ops) + monkeypatch.setitem(sys.modules, "foundry.allocation_region", allocation_region) + monkeypatch.setitem(sys.modules, "foundry.integration.sglang.config", config) + + runtime_path = ( + Path(__file__).parents[1] / "python" / "foundry" / "integration" / "sglang" / "runtime.py" + ) + spec = importlib.util.spec_from_file_location( + "foundry.integration.sglang.runtime", + runtime_path, + ) + assert spec is not None and spec.loader is not None + runtime = importlib.util.module_from_spec(spec) + monkeypatch.setitem(sys.modules, spec.name, runtime) + spec.loader.exec_module(runtime) + monkeypatch.setattr(runtime.torch._C, "_cuda_getCurrentBlasHandle", lambda: None) + + server_args = object() + runtime.setup_graph_extension(server_args, tp_rank=0, pp_rank=0, dp_rank=None) + first_state = runtime.get_state() + sentinel = tmp_path / "rank_0" / "first-model-runner" + sentinel.write_text("preserve") + + runtime.setup_graph_extension(server_args, tp_rank=0, pp_rank=0, dp_rank=None) + + assert runtime.get_state() is first_state + assert sentinel.read_text() == "preserve" + assert region_calls == [(1234, 4096)] + + +def test_load_preallocates_only_once_across_graph_sessions(monkeypatch, tmp_path) -> None: + foundry = ModuleType("foundry") + foundry_ops = ModuleType("foundry.ops") + foundry.ops = foundry_ops + preallocations = [] + foundry_ops.get_current_alloc_offset = lambda: 100 + foundry_ops.preallocate_region = preallocations.append + + allocation_region = ModuleType("foundry.allocation_region") + allocation_region.parse_size = lambda value: value + config = ModuleType("foundry.integration.sglang.config") + mode = Enum("CUDAGraphExtensionMode", ["NONE", "SAVE", "LOAD"]) + cfg = SimpleNamespace(mode=mode.LOAD, workspace_dir=str(tmp_path)) + config.CUDAGraphExtensionMode = mode + config.compute_workspace_rank = lambda *args: 0 + config.get_config = lambda: cfg + config.get_graph_extension_mode = lambda: mode.LOAD + config.get_hook_library_path = lambda: None + config.get_nvshmem_host_path = lambda: None + (tmp_path / "final_alloc_offset.json").write_text(json.dumps({"final_alloc_offset": 1000})) + + monkeypatch.setitem(sys.modules, "foundry", foundry) + monkeypatch.setitem(sys.modules, "foundry.ops", foundry_ops) + monkeypatch.setitem(sys.modules, "foundry.allocation_region", allocation_region) + monkeypatch.setitem(sys.modules, "foundry.integration.sglang.config", config) + + runtime_path = ( + Path(__file__).parents[1] / "python" / "foundry" / "integration" / "sglang" / "runtime.py" + ) + spec = importlib.util.spec_from_file_location( + "foundry.integration.sglang.runtime", + runtime_path, + ) + assert spec is not None and spec.loader is not None + runtime = importlib.util.module_from_spec(spec) + monkeypatch.setitem(sys.modules, spec.name, runtime) + spec.loader.exec_module(runtime) + runtime._state = runtime.CUDAGraphExtensionState() + + runtime.preallocate_for_load_mode() + runtime.preallocate_for_load_mode() + + assert preallocations == [900] diff --git a/tests/test_sglang_main_backend.py b/tests/test_sglang_main_backend.py new file mode 100644 index 00000000..84b41b8e --- /dev/null +++ b/tests/test_sglang_main_backend.py @@ -0,0 +1,655 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the Foundry project +"""CPU-only contracts for the generalized SGLang full graph backend.""" + +from __future__ import annotations + +import importlib.util +import logging +import os +import sys +from contextlib import contextmanager +from enum import Enum +from pathlib import Path +from types import ModuleType, SimpleNamespace + +import pytest +import torch + +_ROOT = Path(__file__).parents[1] +_SGLANG_INTEGRATION = _ROOT / "python" / "foundry" / "integration" / "sglang" + + +def _load_module(name: str, path: Path): + spec = importlib.util.spec_from_file_location(name, path) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + sys.modules[name] = module + spec.loader.exec_module(module) + return module + + +_ARCHIVE = _load_module("test_sglang_main_backend_archive", _SGLANG_INTEGRATION / "archive.py") +_CODEC = _load_module( + "test_sglang_main_backend_output_codec", + _SGLANG_INTEGRATION / "output_codec.py", +) + + +class FakeLogitsProcessorOutput: + def __init__( + self, + next_token_logits, + hidden_states=None, + topk_p=None, + topk_index=None, + ) -> None: + self.next_token_logits = next_token_logits + self.hidden_states = hidden_states + self.topk_p = topk_p + self.topk_index = topk_index + + +class FakeShapeKey: + def __init__( + self, + size: int, + stream_idx: int | None = None, + variant_label: str | None = None, + ) -> None: + self.size = size + self.stream_idx = stream_idx + self.variant_label = variant_label + + +class FakeRunner: + def __init__( + self, + calls, + *, + is_draft_worker: bool = False, + step: int | None = None, + ) -> None: + self.device_module = SimpleNamespace( + synchronize=lambda: calls.synchronize.append(True), + ) + self.model_runner = SimpleNamespace( + is_draft_worker=is_draft_worker, + server_args=SimpleNamespace(enable_torch_symm_mem=False), + tp_group=SimpleNamespace( + barrier=lambda: calls.barrier.append(True), + ), + ) + self.capture_forward_mode = SimpleNamespace(name="DECODE") + self.step = step + + +@pytest.fixture +def backend_env(monkeypatch, tmp_path): + calls = SimpleNamespace( + barrier=[], + catalog_reload=[], + catalog_session_count=[], + final_offset=[], + finish=[], + graph_context=[], + lifecycle=[], + manifest=[], + pack_fatbins=[], + pack_on_exit=[], + pool_ids=[], + preallocate=[], + saved=[], + start=[], + synchronize=[], + ) + + class Mode(Enum): + NONE = "none" + SAVE = "save" + LOAD = "load" + + mode = SimpleNamespace(value=Mode.SAVE) + cfg = SimpleNamespace(workspace_dir=str(tmp_path)) + state = SimpleNamespace( + capture_index=0, + session_index=0, + load_sessions_validated=False, + ) + load_outputs = {} + + class FakeLoadedGraph: + def __init__(self, path: str) -> None: + self.path = path + self.replay_count = 0 + + def replay(self) -> None: + self.replay_count += 1 + + class FakeFoundryCUDAGraph: + def save(self, path: str, tensors) -> None: + calls.saved.append((path, tensors)) + + @staticmethod + def start_graph_builds(paths, *, pool, num_threads): + calls.start.append((paths, pool, num_threads)) + calls.lifecycle.append("start_graph_builds") + return tuple(paths) + + @staticmethod + def finish_one_graph_load(pending, index): + path = pending[index] + graph = FakeLoadedGraph(path) + calls.finish.append((pending, index, graph)) + return graph, load_outputs[path] + + @contextmanager + def fake_foundry_graph(graph, *, pool, stream): + calls.graph_context.append((graph, pool, stream)) + yield + + foundry = ModuleType("foundry") + foundry_ops = ModuleType("foundry.ops") + foundry_graph = ModuleType("foundry.graph") + foundry_integration = ModuleType("foundry.integration") + foundry_sglang = ModuleType("foundry.integration.sglang") + runtime = ModuleType("foundry.integration.sglang.runtime") + config = ModuleType("foundry.integration.sglang.config") + main_deepep = ModuleType("foundry.integration.sglang.main_deepep") + + class RecordingGraphCatalog(_ARCHIVE.GraphCatalog): + def _reload(self): + calls.catalog_reload.append(True) + return super()._reload() + + @property + def session_count(self): + calls.catalog_session_count.append(True) + return super().session_count + + def open_save_session(self, index, descriptor): + calls.lifecycle.append("open_save_session") + return super().open_save_session(index, descriptor) + + def claim_load_session(self, index, descriptor): + calls.lifecycle.append("claim_load_session") + return super().claim_load_session(index, descriptor) + + def save_graph_manifest(directory, **kwargs) -> None: + manifest_path = Path(directory) / "graph_manifest.json" + if manifest_path.exists(): + raise ValueError("graph_manifest.json was treated as a numbered graph") + calls.manifest.append((directory, kwargs)) + manifest_path.write_text("{}") + + def pack_fatbins_to_folder(directory) -> None: + calls.pack_fatbins.append(directory) + + def set_pack_fatbins_on_exit(enabled) -> None: + calls.pack_on_exit.append(enabled) + + def init_nvshmem_for_loaded_modules() -> int: + calls.lifecycle.append("init_nvshmem") + return 3 + + def set_graph_pool_id(pool) -> None: + calls.pool_ids.append(pool) + + def get_graph_pool(device_module): + return (7, 11) + + def get_state(): + return state + + def log_alloc_offset(label) -> None: + return None + + def preallocate_for_load_mode() -> None: + calls.preallocate.append(True) + calls.lifecycle.append("preallocate") + + def capture_final_alloc_offset() -> None: + calls.final_offset.append(True) + + def get_config(): + return cfg + + def get_graph_extension_mode(): + return mode.value + + def bootstrap_deepep_buffer(runner) -> bool: + calls.lifecycle.append("bootstrap") + return False + + def set_deepep_graph_mode() -> bool: + calls.lifecycle.append("set_graph_mode") + return False + + archive = ModuleType("foundry.integration.sglang.archive") + archive.GraphCatalog = RecordingGraphCatalog + archive.runner_descriptor = _ARCHIVE.runner_descriptor + archive.shape_key_record = _ARCHIVE.shape_key_record + foundry.ops = foundry_ops + foundry.save_graph_manifest = save_graph_manifest + foundry_ops.init_nvshmem_for_loaded_modules = init_nvshmem_for_loaded_modules + foundry_ops.pack_fatbins_to_folder = pack_fatbins_to_folder + foundry_ops.set_pack_fatbins_on_exit = set_pack_fatbins_on_exit + foundry_graph.CUDAGraph = FakeFoundryCUDAGraph + foundry_graph.graph = fake_foundry_graph + foundry_sglang.runtime = runtime + runtime.get_state = get_state + runtime.log_alloc_offset = log_alloc_offset + runtime.preallocate_for_load_mode = preallocate_for_load_mode + runtime.capture_final_alloc_offset = capture_final_alloc_offset + config.CUDAGraphExtensionMode = Mode + config.get_config = get_config + config.get_graph_extension_mode = get_graph_extension_mode + main_deepep.bootstrap_deepep_buffer = bootstrap_deepep_buffer + main_deepep.set_deepep_graph_mode = set_deepep_graph_mode + + pynccl = ModuleType( + "sglang.srt.distributed.device_communicators.pynccl_allocator", + ) + logits = ModuleType("sglang.srt.layers.logits_processor") + base_backend = ModuleType( + "sglang.srt.model_executor.runner_backend.base_cuda_graph_backend", + ) + full_backend = ModuleType( + "sglang.srt.model_executor.runner_backend.full_cuda_graph_backend", + ) + runner_pool = ModuleType("sglang.srt.model_executor.runner_utils.pool") + pynccl.set_graph_pool_id = set_graph_pool_id + logits.LogitsProcessorOutput = FakeLogitsProcessorOutput + base_backend.BaseCudaGraphBackend = object + full_backend.FullCudaGraphBackend = object + runner_pool.get_or_create_global_graph_memory_pool = get_graph_pool + + modules = { + "foundry": foundry, + "foundry.ops": foundry_ops, + "foundry.graph": foundry_graph, + "foundry.integration": foundry_integration, + "foundry.integration.sglang": foundry_sglang, + "foundry.integration.sglang.runtime": runtime, + "foundry.integration.sglang.archive": archive, + "foundry.integration.sglang.config": config, + "foundry.integration.sglang.main_deepep": main_deepep, + "foundry.integration.sglang.output_codec": _CODEC, + "sglang.srt.distributed.device_communicators.pynccl_allocator": pynccl, + "sglang.srt.layers.logits_processor": logits, + "sglang.srt.model_executor.runner_backend.base_cuda_graph_backend": base_backend, + "sglang.srt.model_executor.runner_backend.full_cuda_graph_backend": full_backend, + "sglang.srt.model_executor.runner_utils.pool": runner_pool, + } + for name, module in modules.items(): + monkeypatch.setitem(sys.modules, name, module) + + backend_spec = importlib.util.spec_from_file_location( + "foundry.integration.sglang.main_backend", + _SGLANG_INTEGRATION / "main_backend.py", + ) + assert backend_spec is not None and backend_spec.loader is not None + backend_module = importlib.util.module_from_spec(backend_spec) + monkeypatch.setitem(sys.modules, backend_spec.name, backend_module) + backend_spec.loader.exec_module(backend_module) + return SimpleNamespace( + Backend=backend_module.FoundryMainCudaGraphBackend, + GraphCatalog=_ARCHIVE.GraphCatalog, + Mode=Mode, + calls=calls, + cfg=cfg, + load_outputs=load_outputs, + mode=mode, + state=state, + ) + + +def _shape_record(shape_key: FakeShapeKey) -> dict: + return _ARCHIVE.shape_key_record(shape_key) + + +def _runner_descriptor(runner: FakeRunner) -> dict: + return _ARCHIVE.runner_descriptor(runner, "decode") + + +def _graph_record( + capture_index: int, + filename: str, + shape_key: FakeShapeKey, + schema: dict, +) -> dict: + return { + "capture_index": capture_index, + "filename": filename, + "shape_key": _shape_record(shape_key), + "output_schema": schema, + } + + +def _save_catalog_session( + env, + index: int, + runner: FakeRunner, + graph_record: dict | None = None, +) -> None: + catalog = env.GraphCatalog(env.cfg.workspace_dir) + catalog.open_save_session(index, _runner_descriptor(runner)) + if graph_record is not None: + catalog.append_graph(index, graph_record) + + +def test_backend_opens_save_sessions_in_runner_order_with_stale_constructors( + backend_env, +) -> None: + target_runner = FakeRunner(backend_env.calls) + draft_runner = FakeRunner(backend_env.calls, is_draft_worker=True, step=1) + target = backend_env.Backend(target_runner) + draft = backend_env.Backend(draft_runner) + + with target.capture_session(object()): + pass + with draft.capture_session(object()): + pass + + assert target._session["session_index"] == 0 + assert draft._session["session_index"] == 1 + assert backend_env.state.session_index == 2 + + +def test_backend_bootstraps_before_opening_save_session(backend_env) -> None: + backend = backend_env.Backend(FakeRunner(backend_env.calls)) + + with backend.capture_session(object()): + pass + + assert backend_env.calls.lifecycle[:2] == ["bootstrap", "open_save_session"] + + +def test_backend_claims_load_sessions_in_runner_order(backend_env) -> None: + target_runner = FakeRunner(backend_env.calls) + draft_runner = FakeRunner(backend_env.calls, is_draft_worker=True, step=1) + target_shape = FakeShapeKey(8) + draft_shape = FakeShapeKey(16) + target_filename = "graph_0_decode_target_s0_k8.json" + draft_filename = "graph_1_decode_draft_s1_k16.json" + schema = {"kind": "constant", "value": 7} + _save_catalog_session( + backend_env, + 0, + target_runner, + _graph_record(0, target_filename, target_shape, schema), + ) + _save_catalog_session( + backend_env, + 1, + draft_runner, + _graph_record(1, draft_filename, draft_shape, schema), + ) + backend_env.load_outputs[os.path.join(backend_env.cfg.workspace_dir, target_filename)] = [] + backend_env.load_outputs[os.path.join(backend_env.cfg.workspace_dir, draft_filename)] = [] + backend_env.mode.value = backend_env.Mode.LOAD + target = backend_env.Backend(target_runner) + draft = backend_env.Backend(draft_runner) + + with target.capture_session(object()): + target.capture_one(target_shape, lambda: None) + with draft.capture_session(object()): + draft.capture_one(draft_shape, lambda: None) + + assert target._session["session_index"] == 0 + assert draft._session["session_index"] == 1 + assert backend_env.state.session_index == 2 + + +def test_backend_load_orders_bootstrap_nvshmem_and_graph_linking( + backend_env, + caplog: pytest.LogCaptureFixture, +) -> None: + caplog.set_level(logging.INFO) + runner = FakeRunner(backend_env.calls) + shape_key = FakeShapeKey(8) + filename = "graph_0_decode_target_s0_k8.json" + _save_catalog_session( + backend_env, + 0, + runner, + _graph_record( + 0, + filename, + shape_key, + {"kind": "constant", "value": None}, + ), + ) + path = os.path.join(backend_env.cfg.workspace_dir, filename) + backend_env.load_outputs[path] = [] + backend_env.mode.value = backend_env.Mode.LOAD + backend = backend_env.Backend(runner) + + with backend.capture_session(object()): + backend.capture_one(shape_key, lambda: None) + + assert backend_env.calls.lifecycle[:5] == [ + "bootstrap", + "claim_load_session", + "preallocate", + "init_nvshmem", + "start_graph_builds", + ] + assert ( + "[Foundry] Initialized NVSHMEM for 3 loaded modules before graph builds" in caplog.messages + ) + + +def test_backend_save_writes_catalog_record_with_output_schema(backend_env) -> None: + runner = FakeRunner(backend_env.calls) + backend = backend_env.Backend(runner, phase="decode") + shape_key = FakeShapeKey(8) + logits = torch.arange(8) + output = FakeLogitsProcessorOutput(logits, topk_p=logits) + + with backend.capture_session(object()): + backend.capture_one(shape_key, lambda: output) + + session = backend_env.GraphCatalog( + backend_env.cfg.workspace_dir, + ).claim_load_session(0, _runner_descriptor(runner)) + record = session["graphs"][0] + assert record["capture_index"] == 0 + assert record["filename"] == "graph_0_decode_target_s0_k8.json" + assert record["shape_key"] == { + "size": 8, + "stream_idx": None, + "variant_label": None, + } + assert record["output_schema"]["kind"] == "object" + assert backend_env.calls.saved[0][1] == [logits] + + +def test_backend_runs_full_backend_warmups_before_capture(backend_env) -> None: + runner = FakeRunner(backend_env.calls) + backend = backend_env.Backend(runner, phase="decode") + shape_key = FakeShapeKey(8) + events = [] + + def forward(): + events.append("forward") + return FakeLogitsProcessorOutput(torch.arange(2)) + + def post_warmup(): + events.append("post_warmup") + + with backend.capture_session(object()): + backend.capture_one( + shape_key, + forward, + post_warmup_hook=post_warmup, + ) + + assert events == [ + "forward", + "post_warmup", + "forward", + "post_warmup", + "forward", + ] + assert backend_env.calls.synchronize == [True, True] + assert backend_env.calls.barrier == [True, True] + + +def test_backend_load_rejects_full_shape_key_mismatch(backend_env) -> None: + runner = FakeRunner(backend_env.calls) + archived_shape = FakeShapeKey(8, stream_idx=2, variant_label="lora-a") + filename = "graph_0_decode_target_s0_k8.json" + _save_catalog_session( + backend_env, + 0, + runner, + _graph_record( + 0, + filename, + archived_shape, + {"kind": "constant", "value": None}, + ), + ) + path = os.path.join(backend_env.cfg.workspace_dir, filename) + backend_env.load_outputs[path] = [] + backend_env.mode.value = backend_env.Mode.LOAD + backend = backend_env.Backend(runner) + requested_shape = FakeShapeKey(8, stream_idx=2, variant_label="lora-b") + + with ( + pytest.raises(RuntimeError, match="lora-b.*lora-a"), + backend.capture_session(object()), + ): + backend.capture_one(requested_shape, lambda: None) + + assert backend_env.calls.finish == [] + + +def test_backend_load_reconstructs_schema_and_retains_graph_and_output( + backend_env, +) -> None: + runner = FakeRunner(backend_env.calls) + shape_key = FakeShapeKey(8) + filename = "graph_0_decode_target_s0_k8.json" + original_logits = torch.arange(4) + packed = _CODEC.pack_output( + FakeLogitsProcessorOutput(original_logits, topk_p=original_logits), + allowed_types=(FakeLogitsProcessorOutput,), + ) + _save_catalog_session( + backend_env, + 0, + runner, + _graph_record(0, filename, shape_key, packed.schema), + ) + restored_logits = torch.arange(4) + 10 + path = os.path.join(backend_env.cfg.workspace_dir, filename) + backend_env.load_outputs[path] = [restored_logits] + backend_env.mode.value = backend_env.Mode.LOAD + backend = backend_env.Backend(runner) + + with backend.capture_session(object()): + backend.capture_one(shape_key, lambda: None) + + loaded_graph = backend_env.calls.finish[0][2] + restored = backend._outputs[shape_key] + assert backend._graphs[shape_key] is loaded_graph + assert isinstance(restored, FakeLogitsProcessorOutput) + assert restored.next_token_logits is restored_logits + assert restored.topk_p is restored_logits + + +def test_backend_replay_validates_final_catalog_only_once_per_process(backend_env) -> None: + target_runner = FakeRunner(backend_env.calls) + draft_runner = FakeRunner(backend_env.calls, is_draft_worker=True, step=1) + _save_catalog_session(backend_env, 0, target_runner) + _save_catalog_session(backend_env, 1, draft_runner) + backend_env.mode.value = backend_env.Mode.LOAD + backend_env.state.session_index = 2 + target = backend_env.Backend(target_runner) + draft = backend_env.Backend(draft_runner) + backend_env.calls.catalog_reload.clear() + backend_env.calls.catalog_session_count.clear() + + with target.replay_session(): + pass + with target.replay_session(): + pass + with draft.replay_session(): + pass + + assert backend_env.state.load_sessions_validated is True + assert backend_env.calls.catalog_session_count == [True] + assert backend_env.calls.catalog_reload == [True] + + +def test_backend_replay_catalog_mismatch_remains_fail_closed(backend_env) -> None: + target_runner = FakeRunner(backend_env.calls) + draft_runner = FakeRunner(backend_env.calls, is_draft_worker=True, step=1) + _save_catalog_session(backend_env, 0, target_runner) + backend_env.mode.value = backend_env.Mode.LOAD + backend_env.state.session_index = 1 + target = backend_env.Backend(target_runner) + draft = backend_env.Backend(draft_runner) + _save_catalog_session(backend_env, 1, draft_runner) + backend_env.calls.catalog_reload.clear() + backend_env.calls.catalog_session_count.clear() + + for backend in (target, draft): + with ( + pytest.raises(RuntimeError, match="loaded 1/2 SGLang graph sessions"), + backend.replay_session(), + ): + pass + assert backend_env.state.load_sessions_validated is False + + assert backend_env.calls.catalog_session_count == [True, True] + assert backend_env.calls.catalog_reload == [True, True] + + +def test_backend_replay_reapplies_deepep_graph_mode_before_yield(backend_env) -> None: + backend = backend_env.Backend(FakeRunner(backend_env.calls)) + + with backend.replay_session(): + backend_env.calls.lifecycle.append("yield") + + assert backend_env.calls.lifecycle == ["set_graph_mode", "yield"] + + +def test_backend_save_session_finalizes_archive(backend_env) -> None: + backend = backend_env.Backend(FakeRunner(backend_env.calls)) + + with backend.capture_session(object()): + pass + + assert backend_env.calls.manifest == [ + ( + backend_env.cfg.workspace_dir, + {"strip_dependencies": False}, + ), + ] + assert backend_env.calls.pack_fatbins == [backend_env.cfg.workspace_dir] + assert backend_env.calls.pack_on_exit == [False] + assert backend_env.calls.final_offset == [True] + + +def test_backend_rebuilds_manifest_for_each_graph_session(backend_env) -> None: + prefill = backend_env.Backend(FakeRunner(backend_env.calls), phase="prefill") + decode = backend_env.Backend(FakeRunner(backend_env.calls), phase="decode") + + with prefill.capture_session(object()): + pass + with decode.capture_session(object()): + pass + + assert backend_env.calls.manifest == [ + ( + backend_env.cfg.workspace_dir, + {"strip_dependencies": False}, + ), + ( + backend_env.cfg.workspace_dir, + {"strip_dependencies": False}, + ), + ] diff --git a/tests/test_sglang_main_compat.py b/tests/test_sglang_main_compat.py new file mode 100644 index 00000000..fbac8065 --- /dev/null +++ b/tests/test_sglang_main_compat.py @@ -0,0 +1,1331 @@ +# 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 ast +import importlib.util +import logging +import sys +from pathlib import Path +from types import ModuleType, SimpleNamespace + +import pytest +import tomllib +import torch + +_ROOT = Path(__file__).parents[1] +_SGLANG_INTEGRATION = _ROOT / "python" / "foundry" / "integration" / "sglang" + + +def _module(monkeypatch, name: str) -> ModuleType: + module = ModuleType(name) + monkeypatch.setitem(sys.modules, name, module) + return module + + +def _load_hooks_main(monkeypatch): + modules = { + name: _module(monkeypatch, name) + for name in ( + "sglang", + "sglang.srt", + "sglang.srt.distributed", + "sglang.srt.distributed.device_communicators", + "sglang.srt.entrypoints", + "sglang.srt.entrypoints.engine", + "sglang.srt.managers", + "sglang.srt.managers.data_parallel_controller", + "sglang.srt.layers", + "sglang.srt.layers.attention", + "sglang.srt.layers.attention.hybrid_linear_attn_backend", + "sglang.srt.layers.attention.linear", + "sglang.srt.layers.attention.linear.gdn_backend", + "sglang.kernels", + "sglang.kernels.ops", + "sglang.kernels.ops.attention", + "sglang.kernels.ops.attention.fla", + "sglang.kernels.ops.attention.fla.chunk", + "sglang.kernels.ops.attention.fla.index", + "sglang.srt.utils", + "sglang.srt.utils.common", + "sglang.srt.mem_cache", + "sglang.srt.mem_cache.kv_cache_configurator", + "sglang.srt.model_executor", + "sglang.srt.model_executor.cuda_graph_config", + "sglang.srt.model_executor.model_runner", + "sglang.srt.model_executor.model_runner_components", + "sglang.srt.model_executor.model_runner_components.cuda_graph_setup", + "sglang.srt.model_executor.pool_configurator", + "sglang.srt.model_executor.runner", + "sglang.srt.model_executor.runner.decode_cuda_graph_runner", + "sglang.srt.model_executor.runner.prefill_cuda_graph_runner", + "sglang.srt.model_executor.runner_backend", + "sglang.srt.model_executor.runner_backend.utils", + "sglang.srt.speculative", + "sglang.srt.speculative.eagle_draft_cuda_graph_runner", + "sglang.srt.speculative.eagle_draft_extend_cuda_graph_runner", + "sglang.srt.speculative.frozen_kv_mtp_cuda_graph_runner", + "sglang.srt.speculative.multi_layer_eagle_draft_extend_cuda_graph_runner", + "foundry", + "foundry.ops", + "foundry.integration", + "foundry.integration.sglang", + "foundry.integration.sglang.config", + "foundry.integration.sglang.main_backend", + "foundry.integration.sglang.runtime", + ) + } + + class Backend: + FULL = "full" + BREAKABLE = "breakable" + DISABLED = "disabled" + + class Mode: + NONE = "none" + SAVE = "save" + LOAD = "load" + + mode = SimpleNamespace(value=Mode.SAVE) + foundry_backends = [] + + class MambaAttnBackendBase: + def init_forward_metadata_out_graph(self, forward_batch, in_capture=False): + self.original_calls.append((forward_batch, in_capture)) + + class HybridLinearAttnBackend: + def init_forward_metadata_out_graph(self, forward_batch, in_capture=False): + return None + + class GDNAttnBackend(MambaAttnBackendBase): + pass + + class FoundryMainCudaGraphBackend: + def __init__(self, runner, *, phase): + self.runner = runner + self.phase = phase + foundry_backends.append(self) + + device_communicators = modules["sglang.srt.distributed.device_communicators"] + device_communicators.torch_symm_mem = _module( + monkeypatch, + "sglang.srt.distributed.device_communicators.torch_symm_mem", + ) + device_communicators.triton_symm_mem_ag = _module( + monkeypatch, + "sglang.srt.distributed.device_communicators.triton_symm_mem_ag", + ) + modules["sglang.srt.entrypoints"].engine = modules["sglang.srt.entrypoints.engine"] + modules["sglang.srt.managers"].data_parallel_controller = modules[ + "sglang.srt.managers.data_parallel_controller" + ] + modules["sglang.srt.layers"].attention = modules["sglang.srt.layers.attention"] + modules["sglang.srt.layers.attention"].hybrid_linear_attn_backend = modules[ + "sglang.srt.layers.attention.hybrid_linear_attn_backend" + ] + modules[ + "sglang.srt.layers.attention.hybrid_linear_attn_backend" + ].MambaAttnBackendBase = MambaAttnBackendBase + modules[ + "sglang.srt.layers.attention.hybrid_linear_attn_backend" + ].HybridLinearAttnBackend = HybridLinearAttnBackend + modules["sglang.srt.layers.attention.linear.gdn_backend"].GDNAttnBackend = GDNAttnBackend + modules["sglang.kernels.ops.attention.fla.chunk"].CHUNK_SIZE = 64 + fla_index_calls = [] + modules["sglang.kernels.ops.attention.fla.index"].prepare_lens = lambda query_start_loc: ( + fla_index_calls.append(("lens", query_start_loc)) + ) + modules["sglang.kernels.ops.attention.fla.index"].prepare_chunk_indices = ( + lambda query_start_loc, chunk_size: fla_index_calls.append( + ("indices", query_start_loc, chunk_size) + ) + ) + modules["sglang.kernels.ops.attention.fla.index"].prepare_chunk_offsets = ( + lambda query_start_loc, chunk_size: fla_index_calls.append( + ("offsets", query_start_loc, chunk_size) + ) + ) + modules["sglang.srt.utils.common"].next_power_of_2 = lambda value: ( + 1 << (value - 1).bit_length() if value > 0 else 1 + ) + modules["sglang.srt.model_executor"].model_runner = modules[ + "sglang.srt.model_executor.model_runner" + ] + modules["sglang.srt.model_executor.runner"].decode_cuda_graph_runner = modules[ + "sglang.srt.model_executor.runner.decode_cuda_graph_runner" + ] + modules["sglang.srt.model_executor.runner"].prefill_cuda_graph_runner = modules[ + "sglang.srt.model_executor.runner.prefill_cuda_graph_runner" + ] + modules["sglang.srt.model_executor.runner_backend"].utils = modules[ + "sglang.srt.model_executor.runner_backend.utils" + ] + modules["sglang.srt.model_executor.model_runner_components"].cuda_graph_setup = modules[ + "sglang.srt.model_executor.model_runner_components.cuda_graph_setup" + ] + modules["sglang.srt.model_executor.cuda_graph_config"].Backend = Backend + modules["sglang.srt.mem_cache.kv_cache_configurator"].KVCacheConfigurator = type( + "KVCacheConfigurator", + (), + {}, + ) + modules["sglang.srt.model_executor.pool_configurator"].MemoryPoolConfig = type( + "MemoryPoolConfig", + (), + {}, + ) + modules["foundry"].ops = modules["foundry.ops"] + modules["foundry.integration.sglang"].runtime = modules["foundry.integration.sglang.runtime"] + modules["foundry.integration.sglang.config"].CUDAGraphExtensionMode = Mode + modules["foundry.integration.sglang.config"].get_graph_extension_mode = lambda: mode.value + modules[ + "foundry.integration.sglang.main_backend" + ].FoundryMainCudaGraphBackend = FoundryMainCudaGraphBackend + + hooks_spec = importlib.util.spec_from_file_location( + "test_foundry_sglang_hooks_main", + _SGLANG_INTEGRATION / "hooks_main.py", + ) + assert hooks_spec is not None and hooks_spec.loader is not None + hooks = importlib.util.module_from_spec(hooks_spec) + monkeypatch.setitem(sys.modules, hooks_spec.name, hooks) + hooks_spec.loader.exec_module(hooks) + return SimpleNamespace( + Backend=Backend, + Mode=Mode, + backend_utils=modules["sglang.srt.model_executor.runner_backend.utils"], + cuda_graph_setup=modules[ + "sglang.srt.model_executor.model_runner_components.cuda_graph_setup" + ], + decode_runner=modules["sglang.srt.model_executor.runner.decode_cuda_graph_runner"], + decode_runner_modules=[ + modules["sglang.srt.model_executor.runner.decode_cuda_graph_runner"], + modules["sglang.srt.speculative.eagle_draft_cuda_graph_runner"], + modules["sglang.srt.speculative.eagle_draft_extend_cuda_graph_runner"], + modules["sglang.srt.speculative.frozen_kv_mtp_cuda_graph_runner"], + modules["sglang.srt.speculative.multi_layer_eagle_draft_extend_cuda_graph_runner"], + ], + foundry_backends=foundry_backends, + fla_index_calls=fla_index_calls, + gdn=modules["sglang.srt.layers.attention.linear.gdn_backend"], + hooks=hooks, + hybrid_linear=modules["sglang.srt.layers.attention.hybrid_linear_attn_backend"], + mode=mode, + model_runner=modules["sglang.srt.model_executor.model_runner"], + prefill_runner=modules["sglang.srt.model_executor.runner.prefill_cuda_graph_runner"], + ) + + +def _load_legacy_hooks(monkeypatch, buffer_cls): + for name in ( + "foundry", + "foundry.integration", + "foundry.integration.sglang", + ): + _module(monkeypatch, name) + runtime = _module(monkeypatch, "foundry.integration.sglang.runtime") + config = _module(monkeypatch, "foundry.integration.sglang.config") + deep_ep = _module(monkeypatch, "deep_ep") + + class Mode: + NONE = "none" + SAVE = "save" + + config.CUDAGraphExtensionMode = Mode + config.get_graph_extension_mode = lambda: Mode.SAVE + config.get_workspace_root = lambda: Path("archive") + config.load_graph_extension_config = lambda _path: None + deep_ep.Buffer = buffer_cls + sys.modules["foundry.integration.sglang"].runtime = runtime + + hooks_spec = importlib.util.spec_from_file_location( + "test_foundry_sglang_legacy_hooks", + _SGLANG_INTEGRATION / "hooks.py", + ) + assert hooks_spec is not None and hooks_spec.loader is not None + hooks = importlib.util.module_from_spec(hooks_spec) + monkeypatch.setitem(sys.modules, hooks_spec.name, hooks) + hooks_spec.loader.exec_module(hooks) + return hooks + + +@pytest.mark.parametrize( + ("ipc_enabled", "expected_nvl_bytes", "expected_use_fabric"), + [(False, 0, True), (True, 123, False)], +) +def test_deepep_transport_patch_selects_supported_fabric_transport( + monkeypatch, + ipc_enabled: bool, + expected_nvl_bytes: int, + expected_use_fabric: bool, +) -> None: + calls = [] + + class Buffer: + def __init__( + self, + group, + num_nvl_bytes=0, + num_rdma_bytes=0, + *, + low_latency_mode=False, + use_fabric=False, + ): + calls.append( + ( + group, + num_nvl_bytes, + num_rdma_bytes, + low_latency_mode, + use_fabric, + ) + ) + + hooks = _load_legacy_hooks(monkeypatch, Buffer) + if ipc_enabled: + monkeypatch.setenv("FOUNDRY_DEEPEP_NVL_IPC", "1") + + hooks._patch_deepep() + Buffer("group", 123, 456, low_latency_mode=True) + + assert calls == [ + ( + "group", + expected_nvl_bytes, + 456, + True, + expected_use_fabric, + ) + ] + + +@pytest.mark.parametrize("ipc_enabled", [False, True]) +def test_deepep_transport_patch_preserves_legacy_buffer_arguments( + monkeypatch, + caplog: pytest.LogCaptureFixture, + ipc_enabled: bool, +) -> None: + calls = [] + + class Buffer: + def __init__( + self, + group, + num_nvl_bytes=0, + num_rdma_bytes=0, + *, + low_latency_mode=False, + ): + calls.append((group, num_nvl_bytes, num_rdma_bytes, low_latency_mode)) + + hooks = _load_legacy_hooks(monkeypatch, Buffer) + if ipc_enabled: + monkeypatch.setenv("FOUNDRY_DEEPEP_NVL_IPC", "1") + + with caplog.at_level(logging.INFO): + hooks._patch_deepep() + caplog.clear() + Buffer("group", 123, 456, low_latency_mode=True) + hooks._patch_deepep() + + assert calls == [("group", 123, 456, True)] + fallback_records = [ + record for record in caplog.records if "preserving its NVL/IPC transport" in record.message + ] + assert len(fallback_records) == 1 + assert ( + sum( + "SGLang DeepEP transport patch installed" in record.message for record in caplog.records + ) + == 1 + ) + + +def test_foundry_registers_an_sglang_main_plugin() -> None: + config = tomllib.loads((_ROOT / "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, + cuda_graph_backend_prefill=None, + disable_prefill_cuda_graph=False, + disable_flashinfer_autotune=False, + enable_profile_cuda_graph=True, + pp_size=1, + tp_size=1, + enable_torch_symm_mem=False, + speculative_algorithm=None, + enable_lora=False, + enable_pdmux=False, + enable_return_hidden_states=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 False + assert server_args.disable_flashinfer_autotune is True + assert server_args.enable_profile_cuda_graph is False + + server_args.cuda_graph_backend_decode = "breakable" + before_post_init(server_args) + assert server_args.cuda_graph_backend_decode == "breakable" + + server_args.cuda_graph_config = SimpleNamespace( + decode=SimpleNamespace(backend="full", bs=[8]), + prefill=SimpleNamespace(backend="full"), + ) + server_args.speculative_algorithm = "NEXTN" + server_args.moe_a2a_backend = "deepep" + after_post_init(None, server_args) + + server_args.tp_size = 2 + with pytest.raises( + RuntimeError, + match="requires --enable-torch-symm-mem", + ): + after_post_init(None, server_args) + + server_args.enable_torch_symm_mem = True + server_args.cuda_graph_config.decode.bs = [4, 8] + with pytest.raises( + RuntimeError, + match="requires one explicit", + ): + after_post_init(None, server_args) + + server_args.cuda_graph_config.decode.backend = "breakable" + after_post_init(None, server_args) + + server_args.cuda_graph_config.prefill.backend = "breakable" + server_args.enable_torch_symm_mem = False + after_post_init(None, server_args) + + +def test_main_hooks_do_not_install_served_prefill_pointer_dump() -> None: + source = (_SGLANG_INTEGRATION / "hooks_main.py").read_text() + tree = ast.parse(source) + patch = next( + node + for node in tree.body + if isinstance(node, ast.FunctionDef) + and node.name == "_patch_hybrid_linear_prefill_metadata" + ) + + assert not any( + isinstance(node, ast.Call) and isinstance(node.func, ast.Name) and node.func.id == "vars" + for node in ast.walk(patch) + ) + assert not any( + isinstance(node, (ast.Assign, ast.AnnAssign)) + and any( + isinstance(target, ast.Attribute) + and isinstance(target.value, ast.Name) + and target.value.id == "HybridLinearAttnBackend" + and target.attr == "init_forward_metadata_out_graph" + for target in (node.targets if isinstance(node, ast.Assign) else [node.target]) + ) + for node in ast.walk(patch) + ) + + +def test_main_resolvers_route_only_full_phases_through_foundry(monkeypatch) -> None: + env = _load_hooks_main(monkeypatch) + decode_sentinel = object() + prefill_sentinel = object() + env.backend_utils.resolve_decode_backend = lambda runner: decode_sentinel + env.backend_utils.resolve_prefill_backend = lambda runner: prefill_sentinel + for runner_module in env.decode_runner_modules: + runner_module.resolve_decode_backend = env.backend_utils.resolve_decode_backend + env.prefill_runner.resolve_prefill_backend = env.backend_utils.resolve_prefill_backend + + env.hooks._patch_backend_resolver() + + config = SimpleNamespace( + decode=SimpleNamespace(backend=env.Backend.FULL), + prefill=SimpleNamespace(backend=env.Backend.FULL), + ) + runner = SimpleNamespace( + model_runner=SimpleNamespace( + server_args=SimpleNamespace(cuda_graph_config=config), + ) + ) + decode_backend = env.backend_utils.resolve_decode_backend(runner) + prefill_backend = env.backend_utils.resolve_prefill_backend(runner) + + assert decode_backend.phase == "decode" + assert prefill_backend.phase == "prefill" + assert all( + runner_module.resolve_decode_backend is env.backend_utils.resolve_decode_backend + for runner_module in env.decode_runner_modules + ) + assert env.prefill_runner.resolve_prefill_backend is env.backend_utils.resolve_prefill_backend + + config.decode.backend = env.Backend.BREAKABLE + config.prefill.backend = env.Backend.BREAKABLE + assert env.backend_utils.resolve_decode_backend(runner) is decode_sentinel + assert env.backend_utils.resolve_prefill_backend(runner) is prefill_sentinel + + +@pytest.mark.parametrize("is_draft_worker", [False, True], ids=["target", "draft"]) +def test_eagle_full_prefill_construction_preserves_hidden_mode_compatibility( + monkeypatch, + is_draft_worker: bool, +) -> None: + env = _load_hooks_main(monkeypatch) + observed_backends = [] + original_backend = object() + env.backend_utils.resolve_decode_backend = lambda runner: object() + env.backend_utils.resolve_prefill_backend = lambda runner: original_backend + env.decode_runner.resolve_decode_backend = env.backend_utils.resolve_decode_backend + env.prefill_runner.resolve_prefill_backend = env.backend_utils.resolve_prefill_backend + env.hooks._patch_backend_resolver() + + config = SimpleNamespace( + decode=SimpleNamespace(backend=env.Backend.FULL), + prefill=SimpleNamespace(backend=env.Backend.FULL), + ) + hybrid_backend = env.hybrid_linear.HybridLinearAttnBackend() + hybrid_backend.linear_attn_backend = env.gdn.GDNAttnBackend() + model_runner = SimpleNamespace( + attn_backend=hybrid_backend, + is_draft_worker=is_draft_worker, + server_args=SimpleNamespace(cuda_graph_config=config, forward_hooks=[]), + spec_algorithm=SimpleNamespace(is_eagle=lambda: True), + ) + + def capture_prefill_graph(*, model_runner, eager_runner, force_for_draft_worker=False): + del eager_runner, force_for_draft_worker + observed_backends.append(model_runner.server_args.cuda_graph_config.prefill.backend) + cuda_graph_runner = SimpleNamespace(model_runner=model_runner) + backend = env.prefill_runner.resolve_prefill_backend(cuda_graph_runner) + assert backend is not original_backend + return SimpleNamespace( + backend=backend, + prefill_backend_name=model_runner.server_args.cuda_graph_config.prefill.backend, + ) + + env.cuda_graph_setup.capture_prefill_graph = capture_prefill_graph + env.model_runner.capture_prefill_graph = capture_prefill_graph + env.hooks._patch_prefill_graph_capture() + + result = env.cuda_graph_setup.capture_prefill_graph( + model_runner=model_runner, + eager_runner=object(), + force_for_draft_worker=is_draft_worker, + ) + + assert observed_backends == [env.Backend.BREAKABLE] + assert config.prefill.backend == env.Backend.FULL + assert result.prefill_backend_name == env.Backend.FULL + assert result.backend.phase == "prefill" + assert env.model_runner.capture_prefill_graph is env.cuda_graph_setup.capture_prefill_graph + + +@pytest.mark.parametrize( + ("backend", "is_eagle", "mode"), + [ + ("breakable", True, "save"), + ("full", False, "save"), + ("full", True, "none"), + ], +) +def test_prefill_construction_override_is_narrowly_scoped( + monkeypatch, + backend: str, + is_eagle: bool, + mode: str, +) -> None: + env = _load_hooks_main(monkeypatch) + observed_backends = [] + config = SimpleNamespace(prefill=SimpleNamespace(backend=backend)) + model_runner = SimpleNamespace( + server_args=SimpleNamespace(cuda_graph_config=config), + spec_algorithm=SimpleNamespace(is_eagle=lambda: is_eagle), + ) + env.mode.value = mode + + def capture_prefill_graph(*, model_runner, eager_runner, force_for_draft_worker=False): + del eager_runner, force_for_draft_worker + observed_backends.append(model_runner.server_args.cuda_graph_config.prefill.backend) + return SimpleNamespace( + prefill_backend_name=model_runner.server_args.cuda_graph_config.prefill.backend, + ) + + env.cuda_graph_setup.capture_prefill_graph = capture_prefill_graph + env.model_runner.capture_prefill_graph = capture_prefill_graph + env.hooks._patch_prefill_graph_capture() + + result = env.cuda_graph_setup.capture_prefill_graph( + model_runner=model_runner, + eager_runner=object(), + ) + + assert observed_backends == [backend] + assert config.prefill.backend == backend + assert result.prefill_backend_name == backend + + +def test_eagle_full_prefill_restores_config_when_construction_fails(monkeypatch) -> None: + env = _load_hooks_main(monkeypatch) + config = SimpleNamespace(prefill=SimpleNamespace(backend=env.Backend.FULL)) + model_runner = SimpleNamespace( + server_args=SimpleNamespace(cuda_graph_config=config), + spec_algorithm=SimpleNamespace(is_eagle=lambda: True), + ) + + def capture_prefill_graph(*, model_runner, eager_runner, force_for_draft_worker=False): + del eager_runner, force_for_draft_worker + assert model_runner.server_args.cuda_graph_config.prefill.backend == env.Backend.BREAKABLE + raise RuntimeError("constructor failed") + + env.cuda_graph_setup.capture_prefill_graph = capture_prefill_graph + env.model_runner.capture_prefill_graph = capture_prefill_graph + env.hooks._patch_prefill_graph_capture() + + with pytest.raises(RuntimeError, match="constructor failed"): + env.cuda_graph_setup.capture_prefill_graph( + model_runner=model_runner, + eager_runner=object(), + ) + + assert config.prefill.backend == env.Backend.FULL + assert env.hooks._full_prefill_eagle_construction.get() is False + + +@pytest.mark.parametrize("force_for_draft_worker", [False, True], ids=["target", "forced-draft"]) +@pytest.mark.parametrize("fallback", ["unsupported-hybrid", "forward-hooks"]) +def test_full_prefill_direct_construction_falls_back_to_eager( + monkeypatch, + caplog: pytest.LogCaptureFixture, + force_for_draft_worker: bool, + fallback: str, +) -> None: + env = _load_hooks_main(monkeypatch) + observed_backends = [] + eager_runner = object() + config = SimpleNamespace(prefill=SimpleNamespace(backend=env.Backend.FULL)) + hybrid_backend = env.hybrid_linear.HybridLinearAttnBackend() + hybrid_backend.linear_attn_backend = env.hybrid_linear.MambaAttnBackendBase() + model_runner = SimpleNamespace( + attn_backend=hybrid_backend if fallback == "unsupported-hybrid" else object(), + server_args=SimpleNamespace( + cuda_graph_config=config, + forward_hooks=[object()] if fallback == "forward-hooks" else [], + ), + spec_algorithm=SimpleNamespace(is_eagle=lambda: True), + ) + + def capture_prefill_graph(*, model_runner, eager_runner, force_for_draft_worker=False): + del force_for_draft_worker + observed_backends.append(model_runner.server_args.cuda_graph_config.prefill.backend) + if model_runner.server_args.cuda_graph_config.prefill.backend == env.Backend.DISABLED: + return eager_runner + return object() + + env.cuda_graph_setup.capture_prefill_graph = capture_prefill_graph + env.model_runner.capture_prefill_graph = capture_prefill_graph + env.hooks._patch_prefill_graph_capture() + caplog.set_level(logging.WARNING) + + result = env.cuda_graph_setup.capture_prefill_graph( + model_runner=model_runner, + eager_runner=eager_runner, + force_for_draft_worker=force_for_draft_worker, + ) + repeated = env.cuda_graph_setup.capture_prefill_graph( + model_runner=model_runner, + eager_runner=eager_runner, + force_for_draft_worker=force_for_draft_worker, + ) + + assert result is eager_runner + assert repeated is eager_runner + assert observed_backends == [env.Backend.DISABLED, env.Backend.DISABLED] + assert config.prefill.backend == env.Backend.FULL + assert env.foundry_backends == [] + assert sum("eager prefill" in message for message in caplog.messages) == 1 + + +def test_unsupported_hybrid_full_prefill_skips_deferred_capture_and_preallocation( + monkeypatch, +) -> None: + env = _load_hooks_main(monkeypatch) + events = [] + config = SimpleNamespace(prefill=SimpleNamespace(backend=env.Backend.FULL)) + hybrid_backend = env.hybrid_linear.HybridLinearAttnBackend() + hybrid_backend.linear_attn_backend = env.hybrid_linear.MambaAttnBackendBase() + model_runner = SimpleNamespace( + attn_backend=hybrid_backend, + server_args=SimpleNamespace(cuda_graph_config=config, forward_hooks=[]), + spec_algorithm=SimpleNamespace(is_eagle=lambda: False), + ) + eager_runner = object() + decode_capture = object() + + def capture_prefill_graph(*, model_runner, eager_runner, force_for_draft_worker=False): + del force_for_draft_worker + backend = model_runner.server_args.cuda_graph_config.prefill.backend + events.append(("prefill-construction", backend)) + return eager_runner if backend == env.Backend.DISABLED else object() + + def capture_cuda_graphs(*, model_runner, capture_decode_cuda_graph=True): + prefill = env.cuda_graph_setup.capture_prefill_graph( + model_runner=model_runner, + eager_runner=eager_runner, + ) + if capture_decode_cuda_graph: + events.append(("decode", env.Backend.FULL)) + return SimpleNamespace( + eager_runner=eager_runner, + prefill_runner=prefill, + decode=decode_capture, + ) + + env.cuda_graph_setup.capture_prefill_graph = capture_prefill_graph + env.model_runner.capture_prefill_graph = capture_prefill_graph + env.cuda_graph_setup.capture_cuda_graphs = capture_cuda_graphs + env.model_runner.capture_cuda_graphs = capture_cuda_graphs + monkeypatch.setattr( + env.hooks, + "_preallocate_mamba_prefill_metadata_buffers", + lambda model_runner: events.append(("preallocate", model_runner)), + ) + env.hooks._patch_prefill_graph_capture() + env.hooks._patch_cuda_graph_capture_order() + + result = env.cuda_graph_setup.capture_cuda_graphs(model_runner=model_runner) + + assert events == [ + ("prefill-construction", env.Backend.DISABLED), + ("decode", env.Backend.FULL), + ] + assert result.prefill_runner is eager_runner + assert result.decode is decode_capture + assert config.prefill.backend == env.Backend.FULL + assert env.foundry_backends == [] + + +def test_forward_hooks_full_prefill_uses_eager_runner_without_tracing_hooks( + monkeypatch, +) -> None: + env = _load_hooks_main(monkeypatch) + events = [] + config = SimpleNamespace(prefill=SimpleNamespace(backend=env.Backend.FULL)) + + def forward_hook(): + events.append("hook-traced") + + model_runner = SimpleNamespace( + attn_backend=object(), + server_args=SimpleNamespace( + cuda_graph_config=config, + forward_hooks=[forward_hook], + ), + spec_algorithm=SimpleNamespace(is_eagle=lambda: False), + ) + eager_runner = object() + + def capture_prefill_graph(*, model_runner, eager_runner, force_for_draft_worker=False): + del force_for_draft_worker + backend = model_runner.server_args.cuda_graph_config.prefill.backend + events.append(("prefill-construction", backend)) + if backend == env.Backend.FULL: + for hook in model_runner.server_args.forward_hooks: + hook() + return object() + return eager_runner + + def capture_cuda_graphs(*, model_runner, capture_decode_cuda_graph=True): + prefill = env.cuda_graph_setup.capture_prefill_graph( + model_runner=model_runner, + eager_runner=eager_runner, + ) + if capture_decode_cuda_graph: + events.append("decode") + return SimpleNamespace( + eager_runner=eager_runner, + prefill_runner=prefill, + decode=object(), + ) + + env.cuda_graph_setup.capture_prefill_graph = capture_prefill_graph + env.model_runner.capture_prefill_graph = capture_prefill_graph + env.cuda_graph_setup.capture_cuda_graphs = capture_cuda_graphs + env.model_runner.capture_cuda_graphs = capture_cuda_graphs + env.hooks._patch_prefill_graph_capture() + env.hooks._patch_cuda_graph_capture_order() + + result = env.cuda_graph_setup.capture_cuda_graphs(model_runner=model_runner) + + assert events == [ + ("prefill-construction", env.Backend.DISABLED), + "decode", + ] + assert result.prefill_runner is eager_runner + assert config.prefill.backend == env.Backend.FULL + assert env.foundry_backends == [] + + +@pytest.mark.parametrize("is_eagle", [True, False]) +def test_full_prefill_capture_runs_after_target_decode(monkeypatch, is_eagle: bool) -> None: + env = _load_hooks_main(monkeypatch) + events = [] + config = SimpleNamespace(prefill=SimpleNamespace(backend=env.Backend.FULL)) + model_runner = SimpleNamespace( + server_args=SimpleNamespace(cuda_graph_config=config, forward_hooks=[]), + spec_algorithm=SimpleNamespace(is_eagle=lambda: is_eagle), + ) + eager_runner = object() + prefill_runner = object() + decode_capture = object() + + def capture_prefill_graph(*, model_runner, eager_runner, force_for_draft_worker=False): + del force_for_draft_worker + if model_runner.server_args.cuda_graph_config.prefill.backend == env.Backend.DISABLED: + return eager_runner + events.append("prefill") + return prefill_runner + + def capture_cuda_graphs(*, model_runner, capture_decode_cuda_graph=True): + eager = object() + prefill = env.cuda_graph_setup.capture_prefill_graph( + model_runner=model_runner, + eager_runner=eager, + ) + if capture_decode_cuda_graph: + events.append("decode") + return SimpleNamespace( + eager_runner=eager, + prefill_runner=prefill, + decode=decode_capture, + ) + + env.cuda_graph_setup.capture_prefill_graph = capture_prefill_graph + env.model_runner.capture_prefill_graph = capture_prefill_graph + env.cuda_graph_setup.capture_cuda_graphs = capture_cuda_graphs + env.model_runner.capture_cuda_graphs = capture_cuda_graphs + env.hooks._patch_prefill_graph_capture() + env.hooks._patch_cuda_graph_capture_order() + + result = env.cuda_graph_setup.capture_cuda_graphs(model_runner=model_runner) + + assert events == ["decode", "prefill"] + assert result.prefill_runner is prefill_runner + assert result.decode is decode_capture + assert config.prefill.backend == env.Backend.FULL + + +def test_foundry_stabilizes_plain_extend_mamba_metadata(monkeypatch) -> None: + env = _load_hooks_main(monkeypatch) + + class ForwardMode: + def is_extend(self, include_draft_extend_v2=False): + return include_draft_extend_v2 + + def is_target_verify(self): + return False + + def is_draft_extend_v2(self): + return False + + query = torch.tensor([0, 10, 10, 10], dtype=torch.int32) + state_indices = torch.tensor([5, 6, 7], dtype=torch.int32) + metadata = SimpleNamespace( + query_start_loc=query, + mamba_cache_indices=state_indices, + mamba_track_indices=None, + track_conv_indices=None, + track_ssm_h_src=None, + track_ssm_h_dst=None, + track_ssm_final_src=None, + track_ssm_final_dst=None, + ) + backend = env.gdn.GDNAttnBackend() + backend.original_calls = [] + backend.pad_slot_id = -1 + backend.query_start_loc_list = [ + torch.empty(0), + torch.empty(0), + torch.zeros(4, dtype=torch.int32), + ] + backend.state_indices_list = [ + torch.empty(0), + torch.empty(0), + torch.zeros(3, dtype=torch.int32), + ] + backend.mamba_track_indices_buf = torch.empty(3, dtype=torch.int64) + backend._forward_metadata = lambda forward_batch: metadata + forward_batch = SimpleNamespace( + batch_size=3, + extend_seq_lens=torch.tensor([10, 0, 0]), + forward_mode=ForwardMode(), + input_ids=torch.empty(10, dtype=torch.int64), + ) + + env.hooks._patch_hybrid_linear_prefill_metadata() + backend.init_forward_metadata_out_graph(forward_batch, in_capture=True) + + assert backend.original_calls == [] + assert torch.equal(metadata.query_start_loc, query) + assert torch.equal( + metadata.mamba_cache_indices, + torch.tensor([5, -1, -1], dtype=torch.int32), + ) + assert backend.forward_metadata is metadata + + +def test_foundry_leaves_unsupported_hybrid_metadata_on_eager_path(monkeypatch) -> None: + env = _load_hooks_main(monkeypatch) + + class ForwardMode: + def is_extend(self, include_draft_extend_v2=False): + return include_draft_extend_v2 + + def is_target_verify(self): + return False + + def is_draft_extend_v2(self): + return False + + backend = env.hybrid_linear.MambaAttnBackendBase() + backend.original_calls = [] + forward_batch = SimpleNamespace(forward_mode=ForwardMode()) + + env.hooks._patch_hybrid_linear_prefill_metadata() + backend.init_forward_metadata_out_graph(forward_batch, in_capture=True) + + assert backend.original_calls == [(forward_batch, True)] + + +def test_foundry_preallocates_shape_specific_mamba_prefill_metadata(monkeypatch) -> None: + env = _load_hooks_main(monkeypatch) + linear_backend = env.gdn.GDNAttnBackend() + hybrid_backend = env.hybrid_linear.HybridLinearAttnBackend() + hybrid_backend.linear_attn_backend = linear_backend + model_runner = SimpleNamespace( + attn_backend=hybrid_backend, + req_to_token_pool=SimpleNamespace(size=32), + spec_algorithm=SimpleNamespace(is_eagle=lambda: False), + server_args=SimpleNamespace( + chunked_prefill_size=8192, + cuda_graph_config=SimpleNamespace( + prefill=SimpleNamespace( + bs=[16, 64], + full_prefill_max_req=None, + ), + ), + ), + ) + + env.hooks._preallocate_mamba_prefill_metadata_buffers(model_runner) + + buffers = linear_backend._foundry_prefill_metadata_buffers + assert list(buffers) == [(16, 64), (16, 16)] + query_64, state_64, track_64 = buffers[(16, 64)] + query_16, state_16, track_16 = buffers[(16, 16)] + assert query_64.shape == query_16.shape == (17,) + assert state_64.shape == state_16.shape == (16,) + assert query_64 is not query_16 + assert state_64 is not state_16 + assert track_64 is track_16 is None + + env.hooks._preallocate_mamba_prefill_metadata_buffers(model_runner) + assert buffers[(16, 64)][0] is query_64 + assert buffers[(16, 16)][0] is query_16 + + +def test_foundry_preallocation_does_not_prime_capacity_limited_fla_caches( + monkeypatch, +) -> None: + env = _load_hooks_main(monkeypatch) + linear_backend = env.gdn.GDNAttnBackend() + hybrid_backend = env.hybrid_linear.HybridLinearAttnBackend() + hybrid_backend.linear_attn_backend = linear_backend + model_runner = SimpleNamespace( + attn_backend=hybrid_backend, + req_to_token_pool=SimpleNamespace(size=32), + spec_algorithm=SimpleNamespace(is_eagle=lambda: False), + server_args=SimpleNamespace( + chunked_prefill_size=8192, + cuda_graph_config=SimpleNamespace( + prefill=SimpleNamespace( + bs=[16, 64], + full_prefill_max_req=None, + ), + ), + ), + ) + + env.hooks._preallocate_mamba_prefill_metadata_buffers(model_runner) + + assert env.fla_index_calls == [] + + +@pytest.mark.parametrize("mode", ["save", "load"]) +@pytest.mark.parametrize( + ("num_tokens", "expected_chunk_sizes"), + [ + (16, [64, 16]), + (17, [64, 32]), + (64, [64]), + ], +) +def test_foundry_primes_only_current_plain_gdn_shape_during_runner_preparation( + monkeypatch, + mode: str, + num_tokens: int, + expected_chunk_sizes: list[int], +) -> None: + env = _load_hooks_main(monkeypatch) + env.mode.value = mode + linear_backend = env.gdn.GDNAttnBackend() + hybrid_backend = env.hybrid_linear.HybridLinearAttnBackend() + hybrid_backend.linear_attn_backend = linear_backend + model_runner = SimpleNamespace( + attn_backend=hybrid_backend, + req_to_token_pool=SimpleNamespace(size=3), + spec_algorithm=SimpleNamespace(is_eagle=lambda: False), + server_args=SimpleNamespace( + chunked_prefill_size=8192, + cuda_graph_config=SimpleNamespace( + prefill=SimpleNamespace( + bs=[num_tokens], + full_prefill_max_req=3, + ), + ), + ), + ) + env.hooks._preallocate_mamba_prefill_metadata_buffers(model_runner) + assert env.fla_index_calls == [] + + class ForwardMode: + def is_extend(self, include_draft_extend_v2=False): + return include_draft_extend_v2 + + def is_target_verify(self): + return False + + def is_draft_extend_v2(self): + return False + + metadata = SimpleNamespace( + query_start_loc=torch.tensor( + [0, num_tokens, num_tokens, num_tokens], + dtype=torch.int32, + ), + mamba_cache_indices=torch.tensor([5, 6, 7], dtype=torch.int32), + mamba_track_indices=None, + track_conv_indices=None, + track_ssm_h_src=None, + track_ssm_h_dst=None, + track_ssm_final_src=None, + track_ssm_final_dst=None, + ) + linear_backend.original_calls = [] + linear_backend.pad_slot_id = -1 + linear_backend._forward_metadata = lambda forward_batch: metadata + forward_batch = SimpleNamespace( + batch_size=3, + extend_seq_lens=torch.tensor([num_tokens, 0, 0], dtype=torch.int32), + forward_mode=ForwardMode(), + input_ids=torch.empty(num_tokens, dtype=torch.int64), + ) + + env.hooks._patch_hybrid_linear_prefill_metadata() + linear_backend.init_forward_metadata_out_graph(forward_batch, in_capture=True) + + query = linear_backend._foundry_prefill_metadata_buffers[(3, num_tokens)][0] + expected_calls = [("lens", query)] + expected_calls.extend(("indices", query, size) for size in expected_chunk_sizes) + expected_calls.append(("offsets", query, 64)) + assert env.fla_index_calls == expected_calls + assert linear_backend.forward_metadata.query_start_loc is query + + +def test_foundry_does_not_prime_fla_metadata_for_eagle_capture(monkeypatch) -> None: + env = _load_hooks_main(monkeypatch) + linear_backend = env.gdn.GDNAttnBackend() + hybrid_backend = env.hybrid_linear.HybridLinearAttnBackend() + hybrid_backend.linear_attn_backend = linear_backend + model_runner = SimpleNamespace( + attn_backend=hybrid_backend, + req_to_token_pool=SimpleNamespace(size=3), + spec_algorithm=SimpleNamespace(is_eagle=lambda: True), + server_args=SimpleNamespace( + chunked_prefill_size=8192, + cuda_graph_config=SimpleNamespace( + prefill=SimpleNamespace( + bs=[16], + full_prefill_max_req=3, + ), + ), + ), + ) + env.hooks._preallocate_mamba_prefill_metadata_buffers(model_runner) + + assert not hasattr(linear_backend, "_foundry_prefill_metadata_buffers") + assert env.fla_index_calls == [] + + class ForwardMode: + def is_extend(self, include_draft_extend_v2=False): + return include_draft_extend_v2 + + def is_target_verify(self): + return False + + def is_draft_extend_v2(self): + return False + + metadata = SimpleNamespace( + query_start_loc=torch.tensor([0, 16, 16, 16], dtype=torch.int32), + mamba_cache_indices=torch.tensor([5, 6, 7], dtype=torch.int32), + mamba_track_indices=None, + track_conv_indices=None, + track_ssm_h_src=None, + track_ssm_h_dst=None, + track_ssm_final_src=None, + track_ssm_final_dst=None, + ) + linear_backend.original_calls = [] + linear_backend.pad_slot_id = -1 + linear_backend._forward_metadata = lambda forward_batch: metadata + forward_batch = SimpleNamespace( + batch_size=3, + extend_seq_lens=torch.tensor([16, 0, 0], dtype=torch.int32), + forward_mode=ForwardMode(), + input_ids=torch.empty(16, dtype=torch.int64), + ) + + env.hooks._patch_hybrid_linear_prefill_metadata() + linear_backend.init_forward_metadata_out_graph(forward_batch, in_capture=True) + + assert env.fla_index_calls == [] + + +def test_foundry_leaves_eagle_prefill_metadata_allocation_in_capture_order( + monkeypatch, +) -> None: + env = _load_hooks_main(monkeypatch) + linear_backend = env.gdn.GDNAttnBackend() + hybrid_backend = env.hybrid_linear.HybridLinearAttnBackend() + hybrid_backend.linear_attn_backend = linear_backend + model_runner = SimpleNamespace( + attn_backend=hybrid_backend, + req_to_token_pool=SimpleNamespace(size=32), + spec_algorithm=SimpleNamespace(is_eagle=lambda: True), + server_args=SimpleNamespace( + chunked_prefill_size=8192, + cuda_graph_config=SimpleNamespace( + prefill=SimpleNamespace( + bs=[16, 64], + full_prefill_max_req=None, + ), + ), + ), + ) + + env.hooks._preallocate_mamba_prefill_metadata_buffers(model_runner) + + assert not hasattr(linear_backend, "_foundry_prefill_metadata_buffers") + assert env.fla_index_calls == [] + + +def test_foundry_uses_shape_specific_mamba_prefill_metadata(monkeypatch) -> None: + env = _load_hooks_main(monkeypatch) + log_messages = [] + + class StrictLogger: + def info(self, message, *args): + log_messages.append(message % args) + + monkeypatch.setattr(env.hooks, "logger", StrictLogger()) + + class ForwardMode: + def is_extend(self, include_draft_extend_v2=False): + return include_draft_extend_v2 + + def is_target_verify(self): + return False + + def is_draft_extend_v2(self): + return False + + def make_batch(num_tokens: int): + return SimpleNamespace( + batch_size=3, + extend_seq_lens=torch.tensor([num_tokens, 0, 0]), + forward_mode=ForwardMode(), + input_ids=torch.empty(num_tokens, dtype=torch.int64), + ) + + def make_metadata(forward_batch): + num_tokens = len(forward_batch.input_ids) + return SimpleNamespace( + query_start_loc=torch.tensor( + [0, num_tokens, num_tokens, num_tokens], + dtype=torch.int32, + ), + mamba_cache_indices=torch.tensor([5, 6, 7], dtype=torch.int32), + mamba_track_indices=None, + track_conv_indices=None, + track_ssm_h_src=None, + track_ssm_h_dst=None, + track_ssm_final_src=None, + track_ssm_final_dst=None, + ) + + backend = env.gdn.GDNAttnBackend() + backend.original_calls = [] + backend.pad_slot_id = -1 + backend.query_start_loc_list = [ + torch.empty(0), + torch.empty(0), + torch.zeros(4, dtype=torch.int32), + ] + backend.state_indices_list = [ + torch.empty(0), + torch.empty(0), + torch.zeros(3, dtype=torch.int32), + ] + backend.mamba_track_indices_buf = torch.empty(3, dtype=torch.int64) + backend._forward_metadata = make_metadata + + env.hooks._patch_hybrid_linear_prefill_metadata() + backend.init_forward_metadata_out_graph(make_batch(64), in_capture=True) + query_64 = backend.forward_metadata.query_start_loc + state_64 = backend.forward_metadata.mamba_cache_indices + backend.init_forward_metadata_out_graph(make_batch(16), in_capture=True) + query_16 = backend.forward_metadata.query_start_loc + state_16 = backend.forward_metadata.mamba_cache_indices + + assert query_64 is not query_16 + assert state_64 is not state_16 + + backend.init_forward_metadata_out_graph(make_batch(9), in_capture=False) + assert backend.forward_metadata.query_start_loc is query_16 + assert "batch_size=3 num_tokens=64" in log_messages[0] + + +def test_foundry_pads_mamba_prefill_start_locations_for_replay(monkeypatch) -> None: + env = _load_hooks_main(monkeypatch) + + class ForwardMode: + def is_extend(self, include_draft_extend_v2=False): + return include_draft_extend_v2 + + def is_target_verify(self): + return False + + def is_draft_extend_v2(self): + return False + + def make_batch( + extend_seq_lens: list[int], + extend_start_loc: list[int], + num_tokens: int, + ): + return SimpleNamespace( + batch_size=4, + extend_seq_lens=torch.tensor(extend_seq_lens, dtype=torch.int32), + extend_start_loc=torch.tensor(extend_start_loc, dtype=torch.int32), + forward_mode=ForwardMode(), + input_ids=torch.empty(num_tokens, dtype=torch.int64), + ) + + def make_metadata(forward_batch): + batch_size = forward_batch.batch_size + query_start_loc = torch.empty(batch_size + 1, dtype=torch.int32) + query_start_loc[:batch_size] = forward_batch.extend_start_loc + query_start_loc[batch_size] = ( + forward_batch.extend_start_loc[-1] + forward_batch.extend_seq_lens[-1] + ) + return SimpleNamespace( + query_start_loc=query_start_loc, + mamba_cache_indices=torch.arange(batch_size, dtype=torch.int32), + mamba_track_indices=None, + track_conv_indices=None, + track_ssm_h_src=None, + track_ssm_h_dst=None, + track_ssm_final_src=None, + track_ssm_final_dst=None, + ) + + backend = env.gdn.GDNAttnBackend() + backend.original_calls = [] + backend.pad_slot_id = -1 + backend._forward_metadata = make_metadata + + env.hooks._patch_hybrid_linear_prefill_metadata() + capture_batch = make_batch([16, 0, 0, 0], [0, 16, 16, 16], 16) + backend.init_forward_metadata_out_graph(capture_batch, in_capture=True) + + replay_batch = make_batch([4, 3, 0, 0], [0, 4], 7) + backend.init_forward_metadata_out_graph(replay_batch) + + assert torch.equal( + backend.forward_metadata.query_start_loc, + torch.tensor([0, 4, 7, 7, 7], dtype=torch.int32), + ) diff --git a/tests/test_sglang_main_deepep.py b/tests/test_sglang_main_deepep.py new file mode 100644 index 00000000..4d2b97e1 --- /dev/null +++ b/tests/test_sglang_main_deepep.py @@ -0,0 +1,141 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the Foundry project +"""CPU-only contracts for SGLang-main DeepEP graph setup.""" + +from __future__ import annotations + +import importlib.util +import logging +import sys +from pathlib import Path +from types import ModuleType, SimpleNamespace + +import pytest + +_ROOT = Path(__file__).parents[1] +_MAIN_DEEPEP = _ROOT / "python" / "foundry" / "integration" / "sglang" / "main_deepep.py" + + +@pytest.fixture +def deepep_env(monkeypatch): + calls = SimpleNamespace(buffer=0, normal_buffer=0, resolve=[]) + active = SimpleNamespace(value=True) + creates_buffer = SimpleNamespace(value=True) + resolved_mode = object() + + class FakeDeepEPBuffer: + state = SimpleNamespace(buffer=None, dispatch_mode=None) + + @classmethod + def _state(cls): + return cls.state + + @classmethod + def set_dispatch_mode(cls, mode): + cls.state.dispatch_mode = mode + + class FakeBufferImpl: + def _get_buffer(self): + calls.buffer += 1 + if creates_buffer.value: + FakeDeepEPBuffer.state.buffer = object() + + class FakeNormalBufferImpl: + def _get_buffer(self): + calls.normal_buffer += 1 + + class FakeDeepEPDispatcher: + def __init__(self) -> None: + self._low_latency_dispatcher = FakeBufferImpl() + self._normal_dispatcher = FakeNormalBufferImpl() + + class FakeMode: + def resolve(self, *, is_extend_in_batch): + calls.resolve.append(is_extend_in_batch) + return resolved_mode + + class FakeBackend: + def is_deepep(self): + return active.value + + deepep = ModuleType("sglang.srt.layers.moe.token_dispatcher.deepep") + deepep.DeepEPBuffer = FakeDeepEPBuffer + deepep.DeepEPDispatcher = FakeDeepEPDispatcher + utils = ModuleType("sglang.srt.layers.moe.utils") + utils.get_deepep_mode = lambda: FakeMode() + utils.get_moe_a2a_backend = lambda: FakeBackend() + monkeypatch.setitem( + sys.modules, + "sglang.srt.layers.moe.token_dispatcher.deepep", + deepep, + ) + monkeypatch.setitem(sys.modules, "sglang.srt.layers.moe.utils", utils) + + spec = importlib.util.spec_from_file_location( + "test_sglang_main_deepep_module", + _MAIN_DEEPEP, + ) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + monkeypatch.setitem(sys.modules, spec.name, module) + spec.loader.exec_module(module) + + def runner_with_modules(*modules): + model = SimpleNamespace(modules=lambda: modules) + return SimpleNamespace(model_runner=SimpleNamespace(model=model)) + + return SimpleNamespace( + DeepEPBuffer=FakeDeepEPBuffer, + DeepEPDispatcher=FakeDeepEPDispatcher, + active=active, + calls=calls, + creates_buffer=creates_buffer, + module=module, + resolved_mode=resolved_mode, + runner_with_modules=runner_with_modules, + ) + + +def test_inactive_deepep_is_a_noop(deepep_env) -> None: + deepep_env.active.value = False + + assert deepep_env.module.bootstrap_deepep_buffer(object()) is False + assert deepep_env.module.set_deepep_graph_mode() is False + assert deepep_env.calls.buffer == 0 + assert deepep_env.calls.resolve == [] + + +def test_bootstrap_uses_low_latency_dispatcher_once_and_is_idempotent( + deepep_env, + caplog: pytest.LogCaptureFixture, +) -> None: + caplog.set_level(logging.INFO) + dispatcher = deepep_env.DeepEPDispatcher() + wrapper = SimpleNamespace(_inners=[dispatcher]) + runner = deepep_env.runner_with_modules(SimpleNamespace(dispatcher=wrapper)) + + assert deepep_env.module.bootstrap_deepep_buffer(runner) is True + assert deepep_env.module.bootstrap_deepep_buffer(runner) is True + + assert deepep_env.calls.buffer == 1 + assert deepep_env.calls.normal_buffer == 0 + assert deepep_env.calls.resolve == [False, False] + assert deepep_env.DeepEPBuffer.state.dispatch_mode is deepep_env.resolved_mode + assert caplog.messages.count("[Foundry] SGLang DeepEP buffer bootstrap ready") == 2 + + +def test_active_deepep_without_dispatcher_fails_descriptively(deepep_env) -> None: + runner = deepep_env.runner_with_modules(SimpleNamespace()) + + with pytest.raises(RuntimeError, match="DeepEP backend is active.*dispatcher"): + deepep_env.module.bootstrap_deepep_buffer(runner) + + +def test_bootstrap_requires_dispatcher_to_create_state_buffer(deepep_env) -> None: + deepep_env.creates_buffer.value = False + dispatcher = deepep_env.DeepEPDispatcher() + wrapper = SimpleNamespace(_inners=[dispatcher]) + runner = deepep_env.runner_with_modules(SimpleNamespace(dispatcher=wrapper)) + + with pytest.raises(RuntimeError, match="DeepEP dispatcher.*did not create.*buffer"): + deepep_env.module.bootstrap_deepep_buffer(runner) diff --git a/tests/test_sglang_output_codec.py b/tests/test_sglang_output_codec.py new file mode 100644 index 00000000..fa8973ab --- /dev/null +++ b/tests/test_sglang_output_codec.py @@ -0,0 +1,141 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the Foundry project +"""CPU-only contracts for the SGLang graph output codec.""" + +from __future__ import annotations + +import importlib.util +import sys +from pathlib import Path + +import pytest +import torch + +_CODEC_PATH = ( + Path(__file__).parents[1] / "python" / "foundry" / "integration" / "sglang" / "output_codec.py" +) + + +def _load_output_codec_module(): + spec = importlib.util.spec_from_file_location( + "foundry.integration.sglang.output_codec", + _CODEC_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 + + +_CODEC = _load_output_codec_module() +pack_output = _CODEC.pack_output +unpack_output = _CODEC.unpack_output + + +class FakeLogitsOutput: + def __init__( + self, + next_token_logits, + hidden_states, + topk_p, + topk_index, + ): + self.next_token_logits = next_token_logits + self.hidden_states = hidden_states + self.topk_p = topk_p + self.topk_index = topk_index + + +def test_output_codec_round_trips_dynamic_nested_tensor_tree() -> None: + logits = torch.arange(12).reshape(3, 4) + hidden = torch.arange(6).reshape(3, 2) + topk_index = torch.arange(3) + first = FakeLogitsOutput(logits, hidden, logits, topk_index) + second = FakeLogitsOutput(hidden, None, logits, topk_index) + value = [ + first, + second, + (hidden, {"shared": logits, "constant": 7}), + ] + + packed = pack_output(value, allowed_types=(FakeLogitsOutput,)) + restored = unpack_output( + tuple(packed.tensors), + packed.schema, + allowed_types=(FakeLogitsOutput,), + ) + + assert len(packed.tensors) == 3 + assert isinstance(restored[0], FakeLogitsOutput) + assert restored[0].next_token_logits is restored[0].topk_p + assert torch.equal(restored[0].hidden_states, hidden) + assert restored[0].topk_index is restored[1].topk_index + assert restored[1].next_token_logits is restored[0].hidden_states + assert restored[1].hidden_states is None + assert restored[1].topk_p is restored[0].next_token_logits + assert restored[2][0] is restored[0].hidden_states + assert restored[2][1]["shared"] is restored[0].next_token_logits + assert restored[2][1]["constant"] == 7 + + +def test_output_codec_rejects_an_unsupported_leaf_with_path() -> None: + with pytest.raises(TypeError, match=r"\$\.bad"): + pack_output({"bad": object()}, allowed_types=()) + + +def test_output_codec_rejects_an_untrusted_object_schema() -> None: + with pytest.raises(ValueError, match="not allowlisted"): + unpack_output( + [], + {"kind": "object", "type": "evil.Payload", "attributes": {}}, + allowed_types=(), + ) + + +def test_output_codec_rejects_tensor_count_mismatch() -> None: + tensor = torch.ones(1) + packed = pack_output(tensor, allowed_types=()) + with pytest.raises(ValueError, match="tensor count"): + unpack_output([], packed.schema, allowed_types=()) + + +def test_output_codec_rejects_recursive_list_container() -> None: + value: list = [] + value.append(value) + with pytest.raises(TypeError, match=r"Recursive SGLang graph output at \$"): + pack_output(value, allowed_types=()) + + +def test_output_codec_rejects_recursive_tuple_container() -> None: + value = ([],) + value[0].append(value) + with pytest.raises(TypeError, match=r"Recursive SGLang graph output at \$"): + pack_output(value, allowed_types=()) + + +def test_output_codec_rejects_recursive_dict_container() -> None: + value: dict = {} + value["self"] = value + with pytest.raises(TypeError, match=r"Recursive SGLang graph output at \$"): + pack_output(value, allowed_types=()) + + +def test_output_codec_rejects_non_string_dictionary_keys() -> None: + with pytest.raises(TypeError, match=r"non-string dict key"): + pack_output({1: "value"}, allowed_types=()) + + +def test_output_codec_rejects_unknown_schema_node_kind() -> None: + with pytest.raises(ValueError, match="Unsupported output schema node kind"): + unpack_output([], {"kind": "evil"}, allowed_types=()) + + +def test_output_codec_rejects_out_of_range_tensor_index() -> None: + tensor = torch.ones(1) + with pytest.raises(ValueError, match="tensor count"): + unpack_output( + [tensor], + {"kind": "tensor", "index": 1}, + allowed_types=(), + ) diff --git a/tests/test_sglang_tp_recipe.py b/tests/test_sglang_tp_recipe.py new file mode 100644 index 00000000..c857e013 --- /dev/null +++ b/tests/test_sglang_tp_recipe.py @@ -0,0 +1,834 @@ +# 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 ast +import importlib.util +import json +import os +import subprocess +import sys +from pathlib import Path + +import pytest +import tomllib + +REPO_ROOT = Path(__file__).parents[1] +RECIPE_DIR = REPO_ROOT / "recipe" / "experimental" +SERVE_SCRIPT = RECIPE_DIR / "serve_qwen3-8b_sglang_tp.sh" +TP_HARNESS_PATH = REPO_ROOT / "tests" / "modal_sglang_tp.py" +DEEPEP_HARNESS_PATH = REPO_ROOT / "tests" / "modal_deepep_fabric.py" + +GRAPH_MODE_ENV_KEYS = ( + "SGLANG_CUDA_GRAPH_BACKEND_PREFILL", + "SGLANG_CUDA_GRAPH_BS_PREFILL", + "SGLANG_SPECULATIVE_ALGORITHM", + "SGLANG_SPECULATIVE_NUM_STEPS", + "SGLANG_SPECULATIVE_EAGLE_TOPK", + "SGLANG_SPECULATIVE_NUM_DRAFT_TOKENS", + "SGLANG_MOE_A2A_BACKEND", + "SGLANG_DEEPEP_MODE", + "SGLANG_MAX_RUNNING_REQUESTS", + "SGLANG_MAX_TOTAL_TOKENS", +) + +OPTIONAL_GRAPH_MODE_FLAGS = ( + "--cuda-graph-bs-prefill", + "--speculative-algorithm", + "--speculative-num-steps", + "--speculative-eagle-topk", + "--speculative-num-draft-tokens", + "--moe-a2a-backend", + "--deepep-mode", + "--max-running-requests", + "--max-total-tokens", +) + +GRAPH_MODE_ENV = { + "SGLANG_CUDA_GRAPH_BACKEND_PREFILL": "full", + "SGLANG_CUDA_GRAPH_BS_PREFILL": "16 64", + "SGLANG_SPECULATIVE_ALGORITHM": "NEXTN", + "SGLANG_SPECULATIVE_NUM_STEPS": "3", + "SGLANG_SPECULATIVE_EAGLE_TOPK": "1", + "SGLANG_SPECULATIVE_NUM_DRAFT_TOKENS": "4", + "SGLANG_MOE_A2A_BACKEND": "deepep", + "SGLANG_DEEPEP_MODE": "low_latency", + "SGLANG_MAX_RUNNING_REQUESTS": "8", + "SGLANG_MAX_TOTAL_TOKENS": "1048576", +} + + +def _load_tp_harness(): + spec = importlib.util.spec_from_file_location("modal_sglang_tp_contract", TP_HARNESS_PATH) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def test_tp_cleanup_uses_image_containing_module_imports() -> None: + tree = ast.parse(TP_HARNESS_PATH.read_text()) + cleanup = next( + node + for node in tree.body + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) and node.name == "cleanup_run" + ) + decorator = next( + decorator + for decorator in cleanup.decorator_list + if isinstance(decorator, ast.Call) + and isinstance(decorator.func, ast.Attribute) + and decorator.func.attr == "function" + ) + image_keyword = next(keyword for keyword in decorator.keywords if keyword.arg == "image") + + assert isinstance(image_keyword.value, ast.Name) + assert image_keyword.value.id == "image" + + +def test_tp_phase_commits_failed_run_artifacts() -> None: + tree = ast.parse(TP_HARNESS_PATH.read_text()) + run_phase = next( + node + for node in tree.body + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) and node.name == "run_phase" + ) + finally_block = next( + node.finalbody for node in ast.walk(run_phase) if isinstance(node, ast.Try) + ) + + assert any( + isinstance(node, ast.Call) + and isinstance(node.func, ast.Attribute) + and isinstance(node.func.value, ast.Name) + and node.func.value.id == "archive_volume" + and node.func.attr == "commit" + for statement in finally_block + for node in ast.walk(statement) + ) + + +def _flag_value_tail(args: list[str], flag: str) -> list[str]: + index = args.index(flag) + tail: list[str] = [] + for arg in args[index + 1 :]: + if arg.startswith("--"): + break + tail.append(arg) + return tail + + +def _run_recipe( + tmp_path: Path, + mode: str | None, + env_overrides: dict[str, str] | None = 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) + fake_python = bin_dir / "python" + fake_python.write_text("#!/usr/bin/env bash\nexit 0\n") + fake_python.chmod(0o755) + + env_path = tmp_path / "env" + args_path = tmp_path / "args" + env = os.environ.copy() + for key in GRAPH_MODE_ENV_KEYS: + env.pop(key, None) + env.update( + { + "CAPTURE_ENV": str(env_path), + "CAPTURE_ARGS": str(args_path), + "PATH": f"{bin_dir}:{env['PATH']}", + } + ) + if env_overrides is not None: + env.update(env_overrides) + + 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" not in args + assert args[args.index("--cuda-graph-backend-prefill") + 1] == "disabled" + for flag in OPTIONAL_GRAPH_MODE_FLAGS: + assert flag not in args + assert "--enforce-disable-flashinfer-allreduce-fusion" in args + assert "--enable-torch-symm-mem" in args + assert args[args.index("--random-seed") + 1] == "42" + assert args[args.index("--cuda-graph-bs") + 1] == "256" + + +@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) + + assert Path(env["FOUNDRY_SGLANG_CONFIG_PATH"]) == RECIPE_DIR / config_name + assert "--foundry-graph-extension-config-path" not in args + assert args[args.index("--tp-size") + 1] == "2" + assert "--disable-custom-all-reduce" in args + + +@pytest.mark.parametrize("mode", [None, "--save", "--load"]) +def test_tp_recipe_exposes_graph_mode_options_from_environment( + tmp_path: Path, + mode: str | None, +) -> None: + env, args = _run_recipe(tmp_path, mode, GRAPH_MODE_ENV) + + assert env["SGLANG_CUDA_GRAPH_BACKEND_PREFILL"] == "full" + assert args[args.index("--cuda-graph-backend-prefill") + 1] == "full" + assert "--disable-piecewise-cuda-graph" not in args + assert args.count("--cuda-graph-bs-prefill") == 1 + assert _flag_value_tail(args, "--cuda-graph-bs-prefill") == ["16", "64"] + + assert args[args.index("--speculative-algorithm") + 1] == "NEXTN" + assert args[args.index("--speculative-num-steps") + 1] == "3" + assert args[args.index("--speculative-eagle-topk") + 1] == "1" + assert args[args.index("--speculative-num-draft-tokens") + 1] == "4" + assert args[args.index("--moe-a2a-backend") + 1] == "deepep" + assert args[args.index("--deepep-mode") + 1] == "low_latency" + assert args[args.index("--max-running-requests") + 1] == "8" + assert args[args.index("--max-total-tokens") + 1] == "1048576" + + +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_tp_log_scan_ignores_optional_torchcodec_import_tracebacks( + tmp_path: Path, +) -> None: + log_path = tmp_path / "server.log" + log_path.write_text( + "Ignore import error when loading a multimodal processor\n" + "[start of libtorchcodec loading traceback]\n" + "Traceback (most recent call last):\n" + "OSError: libavutil.so.60: cannot open shared object file\n" + "[end of libtorchcodec loading traceback].\n" + ) + module_path = REPO_ROOT / "python" / "foundry" / "integration" / "sglang" / "log_scan.py" + spec = importlib.util.spec_from_file_location("foundry_sglang_log_scan", module_path) + assert spec is not None and spec.loader is not None + log_scan = importlib.util.module_from_spec(spec) + spec.loader.exec_module(log_scan) + + assert log_scan.find_runtime_errors(log_path.read_text()) == [] + + +def test_tp_harness_forwards_every_recipe_graph_mode_variable() -> None: + harness = _load_tp_harness() + + assert set(GRAPH_MODE_ENV_KEYS) <= set(harness.SERVE_ENV_NAMES) + + +def test_tp_harness_forwards_artifact_retention_to_remote_phases() -> None: + harness = _load_tp_harness() + + assert "TP_KEEP_ARTIFACTS" in harness.SERVE_ENV_NAMES + + +def test_tp_harness_reports_all_catalog_sessions_and_hidden_state_schemas( + tmp_path: Path, +) -> None: + harness = _load_tp_harness() + rank_dir = tmp_path / "rank_0" + rank_dir.mkdir() + (rank_dir / "sglang_graph_catalog.json").write_text( + json.dumps( + { + "version": 1, + "sessions": [ + { + "session_index": 0, + "runner": { + "phase": "prefill", + "role": "target", + "forward_mode": "EXTEND", + "step": None, + }, + "graphs": [ + { + "output_schema": { + "kind": "object", + "attributes": {"hidden_states": {"kind": "tensor", "index": 0}}, + } + }, + {"output_schema": {"kind": "tensor", "index": 0}}, + ], + }, + { + "session_index": 1, + "runner": { + "phase": "decode", + "role": "target", + "forward_mode": "TARGET_VERIFY", + "step": None, + }, + "graphs": [{"output_schema": {"kind": "tensor", "index": 0}}], + }, + { + "session_index": 2, + "runner": { + "phase": "decode", + "role": "draft", + "forward_mode": "DECODE", + "step": 1, + }, + "graphs": [{"output_schema": {"kind": "tensor", "index": 0}}], + }, + ], + } + ) + ) + + assert harness._catalog_observations(tmp_path) == { + "0": { + "sessions": [ + ("prefill", "target", "EXTEND", None, 2), + ("decode", "target", "TARGET_VERIFY", None, 1), + ("decode", "draft", "DECODE", 1, 1), + ], + "graph_count": 4, + "hidden_states_schema_count": 1, + } + } + + +def test_tp_harness_fingerprints_catalog_and_sums_loaded_sessions( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + harness = _load_tp_harness() + monkeypatch.setattr(harness, "TP_SIZE", 1) + rank_dir = tmp_path / "rank_0" + rank_dir.mkdir() + for name, content in { + "graph_0.json": "{}", + "graph_manifest.json": "{}", + "fatbin_image_packed.img": "fatbin", + "final_alloc_offset.json": '{"final_alloc_offset": 1}', + "sglang_graph_catalog.json": '{"version": 1, "sessions": []}', + }.items(): + (rank_dir / name).write_text(content) + + fingerprints = harness._archive_fingerprints(tmp_path) + assert "sglang_graph_catalog.json" in fingerprints["0"] + + log_path = tmp_path / "load.log" + log_path.write_text( + "[Foundry TP0] Loaded 2 SGLang graphs phase=prefill role=target session=0\n" + "[Foundry TP0] Loaded 3 SGLang graphs phase=decode role=target session=1\n" + "[Foundry TP0] Loaded 5 SGLang graphs phase=decode role=draft session=2\n" + ) + assert harness._scan_log(str(log_path))["loaded_graphs"] == {"0": 10} + + +def test_tp_harness_reports_normalized_graph_difference_paths(tmp_path: Path) -> None: + harness = _load_tp_harness() + save_graph = tmp_path / "save.json" + save2_graph = tmp_path / "save2.json" + save_graph.write_text( + json.dumps( + { + "generators": [{"id": 1, "seed": 2}], + "nodes": [ + { + "id": 5, + "type": "KernelNode", + "params": { + "kernelParams": [{"index": 0, "value_hex": "0000000000600000"}], + "extra_argBuffer_hex": "", + "extra": [1], + "tag": "save", + }, + } + ], + } + ) + ) + save2_graph.write_text( + json.dumps( + { + "generators": [{"id": 9, "seed": 10}], + "nodes": [ + { + "id": 6, + "type": "KernelNode", + "params": { + "kernelParams": [{"index": 0, "value_hex": "0000000000600000"}], + "extra_argBuffer_hex": "", + "extra": [2], + "tag": "save2", + }, + } + ], + } + ) + ) + + assert harness._semantic_graph_difference_paths(save_graph, save2_graph) == [ + "$.nodes[0].id", + "$.nodes[0].params.tag", + "$.vmm_kernel_pointers[0][0]", + ] + + +def test_tp_harness_retains_and_compares_graph_diagnostics( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + harness = _load_tp_harness() + monkeypatch.setattr(harness, "TP_SIZE", 1) + workspace = tmp_path / "workspace" + rank_dir = workspace / "rank_0" + rank_dir.mkdir(parents=True) + graph_path = rank_dir / "graph_0.json" + graph_path.write_text('{"nodes": [{"id": 1, "type": "EmptyNode"}]}') + (rank_dir / "fatbin_image_packed.img").write_text("not copied") + snapshot = tmp_path / "snapshot" + + harness._snapshot_graph_jsons(workspace, snapshot) + graph_path.write_text('{"nodes": [{"id": 2, "type": "EmptyNode"}]}') + + assert not (snapshot / "rank_0" / "fatbin_image_packed.img").exists() + assert harness._archive_graph_difference_paths(snapshot, workspace) == { + "0": {"graph_0.json": ["$.nodes[0].id"]} + } + + +def test_tp_harness_derives_composed_graph_counts_without_weakening_plain_tp() -> None: + harness = _load_tp_harness() + catalogs = { + "0": {"graph_count": 11}, + "1": {"graph_count": 13}, + } + + assert harness._expected_graph_counts( + catalogs, + expected_ranks={"0", "1"}, + speculative=True, + full_prefill=False, + ) == {"0": 11, "1": 13} + assert harness._expected_graph_counts( + catalogs, + expected_ranks={"0", "1"}, + speculative=False, + full_prefill=True, + ) == {"0": 11, "1": 13} + assert harness._expected_graph_counts( + catalogs, + expected_ranks={"0", "1"}, + speculative=False, + full_prefill=False, + ) == {"0": harness.EXPECTED_GRAPH_COUNT, "1": harness.EXPECTED_GRAPH_COUNT} + + +def test_tp_harness_derives_deepep_counts_without_weakening_plain_tp() -> None: + harness = _load_tp_harness() + catalogs = { + "0": {"graph_count": 3}, + "1": {"graph_count": 3}, + } + + assert harness._expected_graph_counts( + catalogs, + expected_ranks={"0", "1"}, + speculative=False, + deepep=True, + full_prefill=False, + ) == {"0": 3, "1": 3} + assert harness._expected_graph_counts( + catalogs, + expected_ranks={"0", "1"}, + speculative=False, + deepep=False, + full_prefill=False, + ) == {"0": harness.EXPECTED_GRAPH_COUNT, "1": harness.EXPECTED_GRAPH_COUNT} + + +def test_tp_harness_requires_full_prefill_and_decode_sessions_on_every_rank() -> None: + harness = _load_tp_harness() + complete = { + "sessions": [ + ("prefill", "target", "EXTEND", None, 2), + ("decode", "target", "DECODE", None, 1), + ], + "graph_count": 3, + "hidden_states_schema_count": 0, + } + + assert harness._full_prefill_catalog_checks( + {"0": complete, "1": complete}, + {"0", "1"}, + ) == { + "prefill_session_each_rank": True, + "decode_session_each_rank": True, + } + + decode_only = {**complete, "sessions": complete["sessions"][1:]} + checks = harness._full_prefill_catalog_checks( + {"0": complete, "1": decode_only}, + {"0", "1"}, + ) + assert checks["prefill_session_each_rank"] is False + assert checks["decode_session_each_rank"] is True + + +def test_plain_full_prefill_keeps_baseline_as_native_eager_oracle( + monkeypatch: pytest.MonkeyPatch, +) -> None: + harness = _load_tp_harness() + full_prefill_env = { + "SGLANG_CUDA_GRAPH_BACKEND_PREFILL": "full", + "SGLANG_CUDA_GRAPH_BS_PREFILL": "16 64", + } + monkeypatch.setattr(harness, "SERVE_ENV", full_prefill_env) + monkeypatch.setattr(harness, "SPECULATIVE_ALGORITHM", None) + + assert harness._serve_env_for_phase("baseline") == { + "SGLANG_CUDA_GRAPH_BACKEND_PREFILL": "disabled" + } + assert harness._serve_env_for_phase("save") == full_prefill_env + assert harness._serve_env_for_phase("save2") == full_prefill_env + assert harness._serve_env_for_phase("load") == full_prefill_env + + monkeypatch.setattr(harness, "SPECULATIVE_ALGORITHM", "NEXTN") + assert harness._serve_env_for_phase("baseline") == full_prefill_env + + +def test_tp_harness_requires_actual_deepep_dispatch_and_combine_graph_nodes( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + harness = _load_tp_harness() + monkeypatch.setattr(harness, "TP_SIZE", 1) + rank_dir = tmp_path / "rank_0" + rank_dir.mkdir() + (rank_dir / "graph_0_decode.json").write_text( + json.dumps( + { + "nodes": [ + { + "type": "KernelNode", + "params": { + "function_name": "deep_ep::internode_ll::dispatch<__nv_bfloat16>" + }, + }, + { + "type": "KernelNode", + "params": { + "function_name": "deep_ep::internode_ll::combine<__nv_bfloat16>" + }, + }, + ] + } + ) + ) + (rank_dir / "graph_1_prefill.json").write_text( + json.dumps( + { + "nodes": [ + { + "type": "KernelNode", + "params": { + "function_name": "deep_ep::internode_ll::dispatch<__nv_fp8_e4m3>" + }, + }, + { + "type": "KernelNode", + "params": { + "function_name": "deep_ep::internode_ll::combine<__nv_bfloat16>" + }, + }, + ] + } + ) + ) + (rank_dir / "sglang_graph_catalog.json").write_text( + json.dumps( + { + "sessions": [ + { + "session_index": 0, + "runner": {"phase": "decode"}, + "graphs": [{"filename": "graph_0_decode.json"}], + }, + { + "session_index": 1, + "runner": {"phase": "prefill"}, + "graphs": [{"filename": "graph_1_prefill.json"}], + }, + ] + } + ) + ) + + assert harness._deepep_graph_observations(tmp_path) == { + "0": { + "dispatch_kernel_count": 2, + "combine_kernel_count": 2, + "sessions": [ + ("decode", 0, 1, 1), + ("prefill", 1, 1, 1), + ], + } + } + + +def test_tp_harness_reports_per_rank_deepep_lifecycle_evidence( + tmp_path: Path, +) -> None: + harness = _load_tp_harness() + log_path = tmp_path / "load.log" + log_path.write_text( + "\n".join( + [ + "[2026-07-24 00:00:00 TP0 EP0] [Foundry] SGLang DeepEP transport patch installed", + "[2026-07-24 00:00:01 TP0 EP0] [Foundry] SGLang DeepEP buffer bootstrap ready", + "[2026-07-24 00:00:02 TP0 EP0] [Foundry] Initialized NVSHMEM for 3 loaded modules before graph builds", + "[2026-07-24 00:00:03 TP0 EP0] [Foundry] Started 3 SGLang-main graph builds phase=decode", + "[2026-07-24 00:00:00 TP1 EP1] [Foundry] SGLang DeepEP transport patch installed", + "[2026-07-24 00:00:01 TP1 EP1] [Foundry] SGLang DeepEP buffer bootstrap ready", + "[2026-07-24 00:00:02 TP1 EP1] [Foundry] Initialized NVSHMEM for 2 loaded modules before graph builds", + "[2026-07-24 00:00:03 TP1 EP1] [Foundry] Started 3 SGLang-main graph builds phase=decode", + ] + ) + ) + + scan = harness._scan_log(str(log_path)) + assert scan["deepep_transport_patch_ranks"] == ["0", "1"] + assert scan["deepep_bootstrap_before_graph_builds"] == {"0": True, "1": True} + assert scan["deepep_loaded_nvshmem_module_counts"] == {"0": [3], "1": [2]} + assert scan["deepep_errors"] == [] + + +def test_deepep_harness_separates_local_ibgda_probe_from_kernel_failures( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + harness = _load_tp_harness() + monkeypatch.setattr( + harness, + "SERVE_ENV", + {"NVSHMEM_REMOTE_TRANSPORT": "none"}, + ) + device_probe = ( + "/nvshmem/src/modules/transport/ibgda/ibgda.cpp:3714: NULL value get_device_list failed" + ) + transport_probe = ( + "/nvshmem/src/host/transport/transport.cpp:nvshmemi_transport_init:282: " + "init failed for transport: IBGDA" + ) + assertion = ( + "Assertion failed: /root/.cache/deepep/csrc/kernels/internode_ll.cu:181, " + "condition: ibgda_get_state()->num_rc_per_pe >= num_local_experts" + ) + configuration = ( + "[2026-07-24 TP0 EP0] server_args=ServerArgs(" + "deepep_mode='low_latency', crash_on_processing_error=True)" + ) + log_path = tmp_path / "sglang.log" + log_path.write_text(f"{configuration}\n{device_probe}\n{transport_probe}\n{assertion}\n") + + scan = harness._scan_log(str(log_path)) + + assert scan["deepep_configuration_lines"] == [configuration] + assert scan["deepep_optional_transport_probes"] == [ + device_probe, + transport_probe, + ] + assert scan["deepep_errors"] == [assertion] + + +def test_deepep_keeps_plain_tp_offset_and_collective_checks_conditional() -> None: + harness = _load_tp_harness() + asymmetric_offsets = {"0": 73966551040, "1": 62671290368} + deepep_collectives = { + "0": {"symmetric_all_reduce": 1, "symmetric_all_gather": 0, "nccl": 96}, + "1": {"symmetric_all_reduce": 1, "symmetric_all_gather": 0, "nccl": 96}, + } + plain_collectives = { + "0": {"symmetric_all_reduce": 1, "symmetric_all_gather": 1, "nccl": 0}, + "1": {"symmetric_all_reduce": 1, "symmetric_all_gather": 1, "nccl": 0}, + } + + assert not harness._rank_offsets_valid(asymmetric_offsets, deepep=False) + assert harness._rank_offsets_valid(asymmetric_offsets, deepep=True) + assert not harness._collective_counts_valid(deepep_collectives, deepep=False) + assert harness._collective_counts_valid(deepep_collectives, deepep=True) + assert harness._collective_counts_valid(plain_collectives, deepep=False) + + +def test_deepep_harness_forwards_nvshmem_transport_controls() -> None: + harness = _load_tp_harness() + + assert { + "NVSHMEM_REMOTE_TRANSPORT", + "NVSHMEM_IB_ENABLE_IBGDA", + "NVSHMEM_IB_ENABLE", + } <= set(harness.SERVE_ENV_NAMES) + + +def test_standalone_deepep_modal_wrapper_uses_pinned_installer_and_exact_oracle( + monkeypatch: pytest.MonkeyPatch, +) -> None: + source = DEEPEP_HARNESS_PATH.read_text() + tree = ast.parse(source) + monkeypatch.delitem(sys.modules, "modal_sglang_tp", raising=False) + monkeypatch.setattr( + sys, + "path", + [entry for entry in sys.path if Path(entry or ".").resolve() != DEEPEP_HARNESS_PATH.parent], + ) + spec = importlib.util.spec_from_file_location( + "modal_deepep_fabric_contract", + DEEPEP_HARNESS_PATH, + ) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + + assert 'gpu="H100:2"' in source + assert '"/data": archive_volume' in source + assert "scripts/ci/cuda/ci_install_deepep.sh" in TP_HARNESS_PATH.read_text() + assert '"/foundry/tests/test_deepep_fabric.py"' in source + assert '"--run"' in source + assert '"TEST_USE_FABRIC": "1"' in source + assert '"TEST_LOAD_API": "parallel"' in source + assert '"TEST_NCCL_WARMUP_PRE_REGION": "1"' in source + assert 'cwd="/data"' in source + assert any( + isinstance(node, ast.Call) + and isinstance(node.func, ast.Attribute) + and node.func.attr == "rmtree" + for node in ast.walk(tree) + ) + + +def test_deepep_image_removes_container_kernel_driver_packages_after_install() -> None: + source = TP_HARNESS_PATH.read_text() + installer_position = source.rindex("DEEPEP_INSTALL_COMMAND") + cleanup_call_position = source.rindex("DEEPEP_DRIVER_CLEANUP_COMMAND") + cleanup_position = source.index("apt-get purge -y nvidia-dkms-580") + + assert installer_position < cleanup_call_position + assert "nvidia-kernel-common-580" in source + assert "nvidia-kernel-source-580" in source + assert "nvidia-firmware-580" in source + assert "nvidia-modprobe" in source + assert "gdrdrv-dkms" not in source[cleanup_position:] + + +def test_tp_harness_requires_speculative_sessions_and_hidden_states_on_every_rank() -> None: + harness = _load_tp_harness() + complete = { + "sessions": [ + ("prefill", "target", "EXTEND", None, 2), + ("decode", "target", "TARGET_VERIFY", None, 4), + ("decode", "draft", "DECODE", None, 4), + ], + "graph_count": 10, + "hidden_states_schema_count": 1, + } + catalogs = {"0": complete, "1": complete} + + assert harness._speculative_catalog_checks(catalogs, {"0", "1"}) == { + "target_verify_session_each_rank": True, + "draft_session_each_rank": True, + "prefill_session_each_rank": True, + "hidden_states_schema_each_rank": True, + } + + missing_draft = {**complete, "sessions": complete["sessions"][:2]} + checks = harness._speculative_catalog_checks( + {"0": complete, "1": missing_draft}, + {"0", "1"}, + ) + assert checks["draft_session_each_rank"] is False + + +def test_tp_harness_extracts_reported_speculative_accept_lengths() -> None: + harness = _load_tp_harness() + + assert harness._spec_accept_lengths_from_info( + { + "internal_states": [ + {"avg_spec_accept_length": 2.75}, + {"other": "state"}, + ] + } + ) == [2.75] + + +def test_tp_harness_uses_cold_save_caches_and_reuses_save2_cache_for_load( + tmp_path: Path, +) -> None: + harness = _load_tp_harness() + + assert harness._triton_cache_dir(tmp_path, "save") == tmp_path / "triton_cache_save" + assert harness._triton_cache_dir(tmp_path, "save2") == tmp_path / "triton_cache_save2" + assert harness._triton_cache_dir(tmp_path, "load") == tmp_path / "triton_cache_save2" + + +def test_speculative_text_check_keeps_plain_tp_strict_and_requires_restore_stability() -> None: + harness = _load_tp_harness() + baseline = ["native"] + archived = ["restored"] + + assert not harness._deterministic_outputs_match( + baseline, + archived, + archived, + archived, + speculative=False, + ) + assert harness._deterministic_outputs_match( + baseline, + archived, + archived, + archived, + speculative=True, + ) + assert not harness._deterministic_outputs_match( + baseline, + archived, + archived, + ["different load"], + speculative=True, + )