From d1a7237053e5b7eb5cd0039075d969bb1e62d21e Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Thu, 23 Jul 2026 06:45:38 +0000 Subject: [PATCH 01/59] docs: design no-IMEX DeepEP transport Co-authored-by: Rahul Chalamala --- .../specs/2026-07-23-no-imex-deepep-design.md | 211 ++++++++++++++++++ 1 file changed, 211 insertions(+) create mode 100644 docs/superpowers/specs/2026-07-23-no-imex-deepep-design.md diff --git a/docs/superpowers/specs/2026-07-23-no-imex-deepep-design.md b/docs/superpowers/specs/2026-07-23-no-imex-deepep-design.md new file mode 100644 index 00000000..65aa8432 --- /dev/null +++ b/docs/superpowers/specs/2026-07-23-no-imex-deepep-design.md @@ -0,0 +1,211 @@ +# No-IMEX Intranode DeepEP Design + +## Goal + +Allow Foundry graph SAVE and LOAD for intranode expert parallelism in both +vLLM and SGLang without NVIDIA IMEX fabric. Existing fabric behavior remains +the default. + +## Scope + +This design covers: + +- vLLM and SGLang DeepEP expert-parallel workloads on one host. +- An explicit `fabric` or `nvl_ipc` transport choice in each engine's Foundry + TOML configuration. +- Cross-process CUDA VMM IPC that remains valid when a graph is restored in a + fresh process. +- SAVE/LOAD transport-profile validation. +- Standalone DeepEP and engine-level GPU validation. + +It does not cover multi-host expert parallelism, removing NVSHMEM entirely, +enabling NCCL CUMEM/NVLS paths, or dense tensor parallelism. DeepEP continues +to use NVSHMEM for its RDMA buffer; `nvl_ipc` removes the IMEX-dependent +fabric allocation used for the intranode NVLink buffer. + +## Configuration + +Both `foundry.integration.vllm.config.CUDAGraphExtensionConfig` and +`foundry.integration.sglang.config.CUDAGraphExtensionConfig` gain: + +```python +class DeepEPTransport(str, enum.Enum): + FABRIC = "fabric" + NVL_IPC = "nvl_ipc" + + +deepep_transport: DeepEPTransport = DeepEPTransport.FABRIC +``` + +TOML examples: + +```toml +# Existing behavior; this line may be omitted. +deepep_transport = "fabric" +``` + +```toml +# Intranode EP without IMEX fabric. +deepep_transport = "nvl_ipc" +``` + +Unknown values fail while parsing the TOML. The default is `fabric`, so +existing configurations and archives retain their current behavior. + +## Archive Contract + +The effective transport is part of `warmup_state.json` for both integrations: + +```json +{ + "deepep_transport": "nvl_ipc", + "deepep_num_nvl_bytes": 134217728 +} +``` + +SAVE writes the configured value. LOAD compares the archive value with the +current TOML before DeepEP buffers or graphs are initialized. A mismatch +raises a `RuntimeError` naming the archived and configured values. + +Older archives have no transport field. They are interpreted as `fabric`, +matching Foundry's historical vLLM behavior and preserving backward +compatibility. An older archive cannot be loaded with `nvl_ipc`. + +The engine command line still determines the NVLink buffer size. Foundry does +not add a second buffer-size setting. After DeepEP computes +`num_nvl_bytes`, SAVE updates the archive with that effective value and LOAD +compares it before accepting the communicator. Fabric mode records zero. +This catches SAVE and LOAD commands that select different buffer geometry. + +## Engine Integration + +### vLLM + +`_patch_deepep` continues to wrap +`DeepEPLLAll2AllManager._make_all2all_kwargs`. + +- `fabric`: set `use_fabric=True` and `num_nvl_bytes=0`, preserving current + behavior. +- `nvl_ipc`: set `use_fabric=False` and preserve vLLM's computed + `num_nvl_bytes`. + +The wrapper records the effective transport profile during SAVE and validates +it during LOAD. + +### SGLang + +Foundry wraps `deep_ep.Buffer.__init__`, because SGLang does not expose +vLLM's kwargs-builder seam. + +- `fabric`: set `use_fabric=True` and `num_nvl_bytes=0`. +- `nvl_ipc`: set `use_fabric=False` and preserve SGLang's computed + `num_nvl_bytes`. + +The wrapper requires a DeepEP version whose `Buffer` constructor accepts +`use_fabric`. If that argument is unavailable, Foundry raises an actionable +version error instead of silently selecting another transport. + +The existing NVSHMEM bootstrap ordering remains unchanged in both engines: +load captured modules, construct the DeepEP buffer, initialize queued +NVSHMEM modules, then restore graphs. + +## Cross-Process VMM IPC + +The current hook embeds a raw POSIX file-descriptor number in +`CUipcMemHandle`. File-descriptor numbers are process-local, so that protocol +cannot restore an NVL buffer in a fresh process. + +The hook protocol is upgraded to VMM IPC v2: + +1. `cuIpcGetMemHandle` exports the VMM allocation as a POSIX file descriptor. +2. The exporting process parks the descriptor behind a per-process abstract + Unix-domain socket. +3. The 64-byte CUDA IPC handle carries protocol magic, exporter PID, random + process token, allocation/chunk identity, address, size, and generation; + it never carries a raw descriptor number. +4. `cuIpcOpenMemHandle` authenticates to the exporter socket, receives the + descriptor with `SCM_RIGHTS`, and imports it with the CUDA VMM API. +5. The importer first requests the exporter's address. On collision it maps + in a dedicated import range and returns the relocated address. DeepEP's + buffer synchronization distributes the effective peer addresses before + graph use. +6. `cuIpcCloseMemHandle` reference-counts shared chunk mappings and unmaps + them only after the final close. + +Security and lifecycle requirements: + +- The socket accepts only peers with the same UID via `SO_PEERCRED`. +- A random process token prevents stale handles from matching a reused PID. +- Fork detection closes inherited listeners and parked descriptors, then + creates a new server identity. +- Export entries are invalidated when their allocation is freed. +- LOAD preallocation exports the owning VMM chunk plus an interior offset. + Chunk generation participates in cache identity so a freed and recreated + chunk cannot reuse a stale import. +- Driver calls used to reserve an import address bypass Foundry's interposed + allocator, preventing IPC mappings from moving the deterministic graph + allocation cursor. + +Legacy VMM IPC v1 handles are not accepted as Foundry handles; non-Foundry +CUDA IPC handles continue to pass through to the real CUDA driver. + +## Failure Handling + +- Invalid TOML transport: configuration parse error. +- Archive/config transport mismatch: fail before model or graph restoration. +- Effective NVLink buffer-size mismatch: fail when DeepEP creates its buffer. +- Missing `use_fabric` support in DeepEP: actionable runtime error. +- Missing exporter, authentication failure, expired token, or failed FD + transfer: return the corresponding CUDA IPC error and emit one diagnostic. +- Import or map failure: release the imported handle and reserved address; + do not leave a partial cache entry. + +No error path silently falls back from `nvl_ipc` to fabric, because that +would change allocation geometry and invalidate the captured graph. + +## Documentation and Recipes + +The vLLM and SGLang integration documentation will explain: + +- `deepep_transport` values and the default. +- That `nvl_ipc` is intranode-only and still requires DeepEP/NVSHMEM. +- Matching SAVE and LOAD configuration. +- Required `NCCL_CUMEM_ENABLE=0` and `NCCL_NVLS_ENABLE=0`. +- The supported DeepEP revisions. + +Experimental no-IMEX recipes will use `deepep_transport = "nvl_ipc"` in +their SAVE and LOAD TOMLs rather than an environment-variable switch. + +## Testing + +### CPU-only tests + +- Parse the default and both explicit transport values for vLLM and SGLang. +- Reject an unknown transport. +- Round-trip the transport through each `WarmupState`. +- Treat a missing archived field as `fabric`. +- Reject archive/config transport mismatches. +- Verify both DeepEP wrappers select the expected kwargs and preserve + `num_nvl_bytes` only for `nvl_ipc`. +- Verify recipe TOMLs use the first-class setting. + +### Native and GPU tests + +Run on an H100/H200 host without IMEX: + +1. Build Foundry with the CUDA hook and run existing non-EP graph tests to + guard allocator and graph restoration behavior. +2. Run the two-rank standalone DeepEP matrix: + - no-fabric low-latency/RDMA-only path; + - NVL IPC with a nonzero NVLink buffer; + - NVL IPC while LOAD serves allocations from a preallocated VMM chunk. +3. Run two-pass SAVE then LOAD for the vLLM EP recipe and issue an inference + request. Require coherent output and no CUDA error 999. +4. Run SAVE then LOAD for the SGLang EP recipe and issue an inference request + with the same requirements. +5. As a negative control, change LOAD to `fabric` and verify Foundry rejects + the archive before graph restoration. + +Success requires identical per-rank allocation offsets across SAVE/LOAD, +successful VMM IPC imports, graph replay, and coherent model output in both +engines. From 6d8ae26838291d80cfc6ec1ef9b467877df53263 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Thu, 23 Jul 2026 06:53:18 +0000 Subject: [PATCH 02/59] docs: plan no-IMEX DeepEP support Co-authored-by: Rahul Chalamala --- .../plans/2026-07-23-no-imex-deepep.md | 963 ++++++++++++++++++ 1 file changed, 963 insertions(+) create mode 100644 docs/superpowers/plans/2026-07-23-no-imex-deepep.md diff --git a/docs/superpowers/plans/2026-07-23-no-imex-deepep.md b/docs/superpowers/plans/2026-07-23-no-imex-deepep.md new file mode 100644 index 00000000..a0d5b89d --- /dev/null +++ b/docs/superpowers/plans/2026-07-23-no-imex-deepep.md @@ -0,0 +1,963 @@ +# No-IMEX Intranode DeepEP Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Support fresh-process Foundry SAVE/LOAD for intranode vLLM and SGLang DeepEP workloads without NVIDIA IMEX fabric. + +**Architecture:** Add a shared `fabric`/`nvl_ipc` transport contract consumed by both engine integrations and persisted in each archive. Upgrade the CUDA hook from raw-FD VMM IPC v1 to authenticated SCM_RIGHTS VMM IPC v2 so DeepEP's NVLink buffer can be imported by a fresh LOAD process, including buffers carved from Foundry's preallocated VMM chunk. + +**Tech Stack:** Python 3.12, C++17, CUDA Driver API/VMM, PyTorch 2.11/cu130, DeepEP/NVSHMEM, vLLM, SGLang, pytest, Ruff, CMake/Ninja, Modal H100. + +## Global Constraints + +- Scope is intranode expert parallelism for both vLLM and SGLang; multi-host EP and tensor parallelism are out of scope. +- `deepep_transport` accepts exactly `fabric` and `nvl_ipc`; omitted means `fabric`. +- Existing archives without DeepEP fields are interpreted as `fabric`. +- LOAD must reject a transport or effective `num_nvl_bytes` mismatch; it must never silently fall back. +- `nvl_ipc` still uses DeepEP/NVSHMEM for the RDMA buffer and requires `NCCL_CUMEM_ENABLE=0` and `NCCL_NVLS_ENABLE=0`. +- Imports stay at module scope unless an optional engine dependency requires a guarded import and the reason is documented. +- CUDA tests run on Modal H100 with `--env rahul-dev`, `CC=gcc`, and `CXX=g++`. +- CUDA orchestration parents must not initialize CUDA before spawning children that preload `libcuda_hook.so`. +- Every task is committed and pushed independently. + +--- + +## File Map + +### Create + +- [`python/foundry/integration/deepep_transport.py`](../../../python/foundry/integration/deepep_transport.py) — shared transport enum, kwargs policy, and archive validation. +- [`tests/test_deepep_transport.py`](../../../tests/test_deepep_transport.py) — CPU-only shared contract and SGLang config tests. +- [`tests/test_deepep_transport_vllm.py`](../../../tests/test_deepep_transport_vllm.py) — vLLM config/runtime tests in an environment with the Foundry vLLM fork. +- [`recipe/experimental/README.md`](../../../recipe/experimental/README.md) and no-IMEX recipe assets — intranode EP examples for both engines. + +### Modify + +- [`python/foundry/graph.py`](../../../python/foundry/graph.py) — pass the default graph pool explicitly. +- [`csrc/hook.cpp`](../../../csrc/hook.cpp) — VMM IPC v2, preallocated-chunk export/import, fork safety, and alignment. +- [`python/foundry/integration/vllm/config.py`](../../../python/foundry/integration/vllm/config.py) — TOML field and getter. +- [`python/foundry/integration/vllm/runtime.py`](../../../python/foundry/integration/vllm/runtime.py) — archive fields and early LOAD validation. +- [`python/foundry/integration/vllm/hooks.py`](../../../python/foundry/integration/vllm/hooks.py) — transport-aware DeepEP kwargs and profile record/validation. +- [`python/foundry/integration/sglang/config.py`](../../../python/foundry/integration/sglang/config.py) — TOML field and getter. +- [`python/foundry/integration/sglang/runtime.py`](../../../python/foundry/integration/sglang/runtime.py) — archive fields and early LOAD validation. +- [`python/foundry/integration/sglang/hooks.py`](../../../python/foundry/integration/sglang/hooks.py) — transport-aware `deep_ep.Buffer` wrapper. +- [`tests/test_deepep_fabric.py`](../../../tests/test_deepep_fabric.py) — no-fabric, NVL-buffer, and preallocated LOAD cases. +- [`tests/run_deepep_matrix.sh`](../../../tests/run_deepep_matrix.sh) — portable two-GPU matrix driver. +- [`tests/test_vmm_alloc.py`](../../../tests/test_vmm_alloc.py) — zero-alignment VMM regression. +- [`docs/vllm/moe-and-deepep.md`](../../../docs/vllm/moe-and-deepep.md) — transport modes and archive contract. +- [`docs/sglang/hooks.md`](../../../docs/sglang/hooks.md) — SGLang transport patch and DeepEP version floor. + +--- + +### Task 1: Unblock Default Graph Capture + +**Files:** +- Modify: `python/foundry/graph.py` +- Test: `tests/test_load_graph.py` +- Test: `tests/test_async_load.py` + +**Interfaces:** +- Consumes: `graph(cuda_graph, pool: _POOL_HANDLE | None = None, ...)` +- Produces: `graph.pool: tuple[_POOL_HANDLE | None]`, always containing the positional `pool` argument expected by `CUDAGraph.capture_begin`. + +- [ ] **Step 1: Reproduce the missing-pool failure on H100** + +Run: + +```bash +FOUNDRY_REF="$(git rev-parse HEAD)" \ + modal run --env rahul-dev tests/modal_native.py --suite core +``` + +If the Modal harness has not been ported yet, use the repository build/test commands from `README.md` in a two-stage Modal function and run: + +```bash +pytest tests/test_load_graph.py tests/test_async_load.py -q +``` + +Expected before the fix: capture stops because the binding requires the `pool` positional argument. + +- [ ] **Step 2: Add the regression assertion** + +Add a small CPU-safe unit around `graph.__init__` by constructing the object with `__new__`, stubbing `torch.cuda.Stream`, and asserting: + +```python +def test_graph_default_pool_is_forwarded(monkeypatch): + sentinel_stream = object() + monkeypatch.setattr(torch.cuda, "Stream", lambda: sentinel_stream) + cuda_graph = object() + + context = foundry.graph.__new__(foundry.graph) + foundry.graph.__init__(context, cuda_graph) + + assert context.pool == (None,) +``` + +Expected before the fix: `context.pool == ()`. + +- [ ] **Step 3: Pass the default pool explicitly** + +Change: + +```python +self.pool = () if pool is None else (pool,) +``` + +to: + +```python +self.pool = (pool,) +``` + +This is the focused change from source commit `57e1dba`. + +- [ ] **Step 4: Verify the focused tests** + +Run: + +```bash +pytest tests/test_load_graph.py tests/test_async_load.py -q +``` + +Expected: PASS on Modal H100. + +- [ ] **Step 5: Commit and push** + +```bash +git add python/foundry/graph.py tests/test_load_graph.py +git commit -m "fix: pass default CUDA graph pool" +git push -u origin cursor/no-imex-deepep-f10b +``` + +--- + +### Task 2: Port Hardened Cross-Process VMM IPC + +**Files:** +- Modify: `csrc/hook.cpp` +- Modify: `tests/test_deepep_fabric.py` +- Create: `tests/run_deepep_matrix.sh` +- Modify: `tests/test_vmm_alloc.py` +- Optional test harness create: `tests/modal_native.py` +- Optional test harness create: `tests/test_modal_native.py` + +**Interfaces:** +- Consumes: hooked `cuIpcGetMemHandle`, `cuIpcOpenMemHandle`, and `cuIpcCloseMemHandle`. +- Produces: VMM IPC v2 handles containing `{magic, pid, original_ptr, size, token, chunk_base, chunk_size, generation}` and SCM_RIGHTS descriptor transfer over `\0foundry-vmm-ipc.`. + +- [ ] **Step 1: Capture the failing no-fabric baseline** + +On a two-GPU host without IMEX, run the current standalone test with a nonzero NVLink buffer: + +```bash +export TEST_USE_FABRIC=0 +export TEST_NVL_BYTES_MB=64 +export NCCL_CUMEM_ENABLE=0 +export NCCL_NVLS_ENABLE=0 + +LD_PRELOAD="$NVSHMEM_HOST:$FOUNDRY_HOOK" \ + python tests/test_deepep_fabric.py --save --no-fabric --num-processes=2 +LD_PRELOAD="$NVSHMEM_HOST:$FOUNDRY_HOOK" \ + python tests/test_deepep_fabric.py --load --no-fabric --num-processes=2 +``` + +Expected before the fix: LOAD cannot import the raw process-local FD, usually with CUDA error 999. + +- [ ] **Step 2: Port the tested VMM IPC v2 core** + +Port the `csrc/hook.cpp` portions from source commit `84d2b95`: + +```bash +git show 84d2b95 -- csrc/hook.cpp +``` + +The resulting implementation must define: + +```cpp +static constexpr uint32_t VMM_IPC_MAGIC = 0x564D4D32; // "VMM2" + +struct VmmIpcRequest { + uint64_t ptr; + uint64_t token; +}; + +static bool vmm_ipc_ensure_server(); +static int vmm_ipc_fetch_fd(pid_t exporter_pid, uint64_t token, + CUdeviceptr original_ptr); +static void vmm_ipc_invalidate_export(CUdeviceptr dptr); +static void vmm_ipc_server_loop(int listen_fd); +``` + +The implementation must: + +- create an abstract Unix socket named `\0foundry-vmm-ipc.`; +- authenticate peers with `SO_PEERCRED` and require the same UID; +- generate a per-process random token from `/dev/urandom`; +- send FDs only with `SCM_RIGHTS`; +- store no raw FD number in `CUipcMemHandle`; +- use the real, uninterposed `cuMemAddressReserve` for import mappings; +- retry a colliding exporter address in the dedicated import range beginning at `0x300000000000`; +- pass non-Foundry IPC handles to the real CUDA driver. + +- [ ] **Step 3: Port fork and preallocation hardening** + +Port the `csrc/hook.cpp` changes from `eaf3fea` and `ab22a21` without the unrelated tensor-parallel files: + +```cpp +struct PreallocatedChunk { + size_t size; + CUmemGenericAllocationHandle handle; + uint64_t generation; +}; + +using VmmIpcChunkKey = + std::tuple; + +struct VmmIpcChunkMapping { + CUdeviceptr local_base; + size_t size; + CUmemGenericAllocationHandle handle; + int refcount; +}; +``` + +Required behavior: + +- track every preallocated chunk under a mutex; +- include chunk generation at handle offset 48; +- export the owning chunk plus an interior offset for carved allocations; +- cache imported chunks by PID, token, chunk base, and generation; +- reference-count interior pointers; +- clear inherited listeners and parked descriptors after fork; +- invalidate an export when the owning allocation is freed. + +- [ ] **Step 4: Port zero-alignment normalization** + +Port the focused `csrc/hook.cpp` changes from `f947886` and `0f337a1`. + +Use: + +```cpp +const size_t effective_alignment = + alignment == 0 ? kAllocAlignment : alignment; +``` + +for Foundry-region reservation math while preserving the caller-visible CUDA semantics. + +Add/port: + +```python +def test_zero_alignment_reserve_stays_in_region(): + # Reserve with alignment=0 through the driver API, then assert the + # returned address is 2 MiB aligned and remains inside the Foundry region. +``` + +from source commit `77a28e8`. + +- [ ] **Step 5: Port the portable DeepEP matrix** + +Port the non-engine test changes from `84d2b95`, `059aa42`, and `ab22a21`: + +```bash +bash tests/run_deepep_matrix.sh ll_nofabric both +bash tests/run_deepep_matrix.sh nvl_ipc both +bash tests/run_deepep_matrix.sh nvl_ipc_prealloc both +``` + +Matrix definitions: + +```bash +ll_nofabric: TEST_USE_FABRIC=0 TEST_NVL_BYTES_MB=0 +nvl_ipc: TEST_USE_FABRIC=0 TEST_NVL_BYTES_MB=64 +nvl_ipc_prealloc: TEST_USE_FABRIC=0 TEST_NVL_BYTES_MB=64 TEST_LOAD_PREALLOC_MB=1024 +``` + +Expected after the port: all three cases PASS. `nvl_fabric` remains a negative control and is expected to fail on a host without IMEX. + +- [ ] **Step 6: Build and run focused native regressions** + +Run on Modal: + +```bash +CC=gcc CXX=g++ pip install -e . --no-build-isolation +pytest tests/test_vmm_alloc.py tests/test_load_graph.py tests/test_async_load.py -q +``` + +Then run the three two-GPU matrix cases from Step 5. Expected: PASS, no CUDA error 999, and deterministic dispatch output. + +- [ ] **Step 7: Commit and push** + +```bash +git add csrc/hook.cpp tests/test_deepep_fabric.py \ + tests/run_deepep_matrix.sh tests/test_vmm_alloc.py +git commit -m "fix: support cross-process VMM IPC" +git push -u origin cursor/no-imex-deepep-f10b +``` + +--- + +### Task 3: Add the First-Class Transport Configuration + +**Files:** +- Create: `python/foundry/integration/deepep_transport.py` +- Modify: `python/foundry/integration/vllm/config.py` +- Modify: `python/foundry/integration/sglang/config.py` +- Create: `tests/test_deepep_transport.py` +- Create: `tests/test_deepep_transport_vllm.py` + +**Interfaces:** +- Produces: `DeepEPTransport`, `parse_deepep_transport`, `configure_deepep_buffer`, `validate_transport_match`, and `validate_num_nvl_bytes_match`. +- Produces: `get_deepep_transport() -> DeepEPTransport` in both engine config modules. + +- [ ] **Step 1: Write failing shared-contract tests** + +Create `tests/test_deepep_transport.py` with: + +```python +from __future__ import annotations + +import pytest + +from foundry.integration.deepep_transport import ( + DeepEPTransport, + configure_deepep_buffer, + parse_deepep_transport, + validate_num_nvl_bytes_match, + validate_transport_match, +) + + +@pytest.mark.parametrize( + ("raw", "expected"), + [ + (None, DeepEPTransport.FABRIC), + ("fabric", DeepEPTransport.FABRIC), + ("nvl_ipc", DeepEPTransport.NVL_IPC), + ], +) +def test_parse_deepep_transport(raw, expected): + assert parse_deepep_transport(raw) is expected + + +def test_parse_deepep_transport_rejects_unknown_value(): + with pytest.raises(ValueError, match="deepep_transport"): + parse_deepep_transport("rdma") + + +@pytest.mark.parametrize( + ("transport", "expected_bytes", "expected_fabric"), + [ + (DeepEPTransport.FABRIC, 0, True), + (DeepEPTransport.NVL_IPC, 128, False), + ], +) +def test_configure_deepep_buffer(transport, expected_bytes, expected_fabric): + num_nvl_bytes, kwargs = configure_deepep_buffer( + transport, + num_nvl_bytes=128, + kwargs={"other": "kept"}, + ) + assert num_nvl_bytes == expected_bytes + assert kwargs == {"other": "kept", "use_fabric": expected_fabric} + + +def test_transport_mismatch_names_both_values(): + with pytest.raises(RuntimeError, match="archived=nvl_ipc.*configured=fabric"): + validate_transport_match( + archived="nvl_ipc", + configured=DeepEPTransport.FABRIC, + ) + + +def test_nvl_bytes_mismatch_names_both_values(): + with pytest.raises(RuntimeError, match="archived=64.*effective=128"): + validate_num_nvl_bytes_match( + archived=64, + effective=128, + transport=DeepEPTransport.NVL_IPC, + ) +``` + +Expected: import failure because the shared module does not exist. + +- [ ] **Step 2: Implement the shared transport module** + +Create: + +```python +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the Foundry project +"""Shared DeepEP transport policy for Foundry engine integrations.""" + +from __future__ import annotations + +import enum +from typing import Any + + +class DeepEPTransport(str, enum.Enum): + FABRIC = "fabric" + NVL_IPC = "nvl_ipc" + + +def parse_deepep_transport(raw: str | None) -> DeepEPTransport: + value = DeepEPTransport.FABRIC.value if raw is None else raw + try: + return DeepEPTransport(value) + except ValueError as exc: + choices = ", ".join(item.value for item in DeepEPTransport) + raise ValueError( + f"deepep_transport must be one of: {choices}; got {value!r}" + ) from exc + + +def configure_deepep_buffer( + transport: DeepEPTransport, + *, + num_nvl_bytes: int, + kwargs: dict[str, Any], +) -> tuple[int, dict[str, Any]]: + configured = dict(kwargs) + if transport is DeepEPTransport.FABRIC: + configured["use_fabric"] = True + return 0, configured + configured["use_fabric"] = False + return num_nvl_bytes, configured + + +def validate_transport_match( + *, + archived: str | None, + configured: DeepEPTransport, +) -> None: + archived_transport = parse_deepep_transport(archived) + if archived_transport is not configured: + raise RuntimeError( + "Foundry DeepEP transport mismatch: " + f"archived={archived_transport.value}, configured={configured.value}" + ) + + +def validate_num_nvl_bytes_match( + *, + archived: int, + effective: int, + transport: DeepEPTransport, +) -> None: + if archived != effective: + raise RuntimeError( + "Foundry DeepEP num_nvl_bytes mismatch for " + f"{transport.value}: archived={archived}, effective={effective}" + ) +``` + +- [ ] **Step 3: Write failing engine-config tests** + +For SGLang, parse temporary TOMLs directly. For vLLM, place the equivalent tests in `tests/test_deepep_transport_vllm.py` and use `pytest.importorskip("vllm")` so the Foundry vLLM environment runs them. + +Required assertions: + +```python +assert CUDAGraphExtensionConfig.from_toml(default_path).deepep_transport is DeepEPTransport.FABRIC +assert CUDAGraphExtensionConfig.from_toml(nvl_path).deepep_transport is DeepEPTransport.NVL_IPC +``` + +and `"rdma"` raises `ValueError` naming `deepep_transport`. + +- [ ] **Step 4: Add the field and getter to both configs** + +In both config dataclasses: + +```python +deepep_transport: DeepEPTransport = DeepEPTransport.FABRIC +``` + +In both `from_toml` constructors: + +```python +deepep_transport=parse_deepep_transport(data.get("deepep_transport")), +``` + +In both modules: + +```python +def get_deepep_transport() -> DeepEPTransport: + if _config is None: + return DeepEPTransport.FABRIC + return _config.deepep_transport +``` + +- [ ] **Step 5: Run CPU tests** + +```bash +pytest tests/test_deepep_transport.py -q +pytest tests/test_deepep_transport_vllm.py -q +``` + +Expected: shared and SGLang tests PASS; vLLM tests PASS in the fork environment and skip only when vLLM is not installed. + +- [ ] **Step 6: Commit and push** + +```bash +git add python/foundry/integration/deepep_transport.py \ + python/foundry/integration/vllm/config.py \ + python/foundry/integration/sglang/config.py \ + tests/test_deepep_transport.py tests/test_deepep_transport_vllm.py +git commit -m "feat: configure DeepEP transport in TOML" +git push -u origin cursor/no-imex-deepep-f10b +``` + +--- + +### Task 4: Persist and Validate the Transport Profile + +**Files:** +- Modify: `python/foundry/integration/vllm/runtime.py` +- Modify: `python/foundry/integration/sglang/runtime.py` +- Modify: `tests/test_deepep_transport.py` +- Modify: `tests/test_deepep_transport_vllm.py` + +**Interfaces:** +- Produces identical archive fields in both `WarmupState` dataclasses: + `deepep_transport: str` and `deepep_num_nvl_bytes: int`. +- Produces `update_warmup_state_deepep_profile(...)` and early LOAD validation in each runtime. + +- [ ] **Step 1: Write failing archive tests** + +Add tests that serialize and reload: + +```python +state.deepep_transport = "nvl_ipc" +state.deepep_num_nvl_bytes = 64 * 1024 * 1024 +``` + +Assert the resulting JSON contains both fields. + +Add a legacy JSON fixture without either field and assert: + +```python +assert loaded.deepep_transport == "fabric" +assert loaded.deepep_num_nvl_bytes == 0 +``` + +Add a LOAD mismatch test that stubs CUDA operations and asserts the mismatch is raised before `load_cuda_modules_and_libraries`. + +- [ ] **Step 2: Extend both WarmupState dataclasses** + +Add: + +```python +deepep_transport: str = DeepEPTransport.FABRIC.value +deepep_num_nvl_bytes: int = 0 +``` + +The existing field-filtering loaders then make missing legacy fields use the dataclass defaults. + +- [ ] **Step 3: Add profile update helpers** + +vLLM: + +```python +def update_warmup_state_deepep_profile( + workspace_root: str, + transport: DeepEPTransport, + num_nvl_bytes: int, +) -> None: + state = load_warmup_state(workspace_root) + state.deepep_transport = transport.value + state.deepep_num_nvl_bytes = num_nvl_bytes + save_warmup_state(workspace_root, state) +``` + +SGLang: + +```python +def update_warmup_state_deepep_profile( + transport: DeepEPTransport, + num_nvl_bytes: int, +) -> None: + state = load_warmup_state() + state.deepep_transport = transport.value + state.deepep_num_nvl_bytes = num_nvl_bytes + save_warmup_state(state) +``` + +Only workspace rank 0 writes the shared state. Preserve SGLang's existing rank guard; for vLLM, return without writing when `Path(cfg.workspace_dir).name != "rank_0"`. + +- [ ] **Step 4: Fail early on LOAD transport mismatch** + +In each `setup_graph_extension`, after resolving the rank workspace but before loading modules or reserving the VMM region: + +```python +if cfg.mode == CUDAGraphExtensionMode.LOAD: + warmup_state = load_warmup_state(...) + validate_transport_match( + archived=warmup_state.deepep_transport, + configured=cfg.deepep_transport, + ) +``` + +This guarantees a mismatch fails before DeepEP, NVSHMEM, or graph restoration. + +- [ ] **Step 5: Run archive tests** + +```bash +pytest tests/test_deepep_transport.py tests/test_deepep_transport_vllm.py -q +``` + +Expected: PASS, including legacy archive behavior and fail-fast call ordering. + +- [ ] **Step 6: Commit and push** + +```bash +git add python/foundry/integration/vllm/runtime.py \ + python/foundry/integration/sglang/runtime.py \ + tests/test_deepep_transport.py tests/test_deepep_transport_vllm.py +git commit -m "feat: validate archived DeepEP transport" +git push -u origin cursor/no-imex-deepep-f10b +``` + +--- + +### Task 5: Apply the Transport in vLLM and SGLang + +**Files:** +- Modify: `python/foundry/integration/vllm/hooks.py` +- Modify: `python/foundry/integration/sglang/hooks.py` +- Modify: `tests/test_deepep_transport.py` +- Modify: `tests/test_deepep_transport_vllm.py` + +**Interfaces:** +- vLLM wraps `DeepEPLLAll2AllManager._make_all2all_kwargs`. +- SGLang wraps `deep_ep.Buffer.__init__`. +- Both record the effective `num_nvl_bytes` on SAVE and validate it on LOAD. + +- [ ] **Step 1: Write failing policy tests** + +Test these exact outcomes: + +```python +# fabric +{"num_nvl_bytes": 0, "use_fabric": True} + +# nvl_ipc, given upstream num_nvl_bytes=67108864 +{"num_nvl_bytes": 67108864, "use_fabric": False} +``` + +Test NONE mode leaves upstream values unchanged. + +Test LOAD with an archived 64 MiB buffer and an effective 128 MiB buffer raises before the original `Buffer.__init__` or communicator setup completes. + +- [ ] **Step 2: Replace the vLLM fabric-only policy** + +In `_patch_deepep`: + +```python +transport = get_deepep_transport() +upstream_num_nvl_bytes = int(out.get("num_nvl_bytes", 0)) +effective_num_nvl_bytes, configured = configure_deepep_buffer( + transport, + num_nvl_bytes=upstream_num_nvl_bytes, + kwargs=out, +) +out.clear() +out.update(configured) +out["num_nvl_bytes"] = effective_num_nvl_bytes +``` + +Then: + +```python +if mode == CUDAGraphExtensionMode.SAVE: + rt.update_warmup_state_deepep_profile( + root, + transport, + effective_num_nvl_bytes, + ) +elif mode == CUDAGraphExtensionMode.LOAD: + state = rt.load_warmup_state(root) + validate_num_nvl_bytes_match( + archived=state.deepep_num_nvl_bytes, + effective=effective_num_nvl_bytes, + transport=transport, + ) +``` + +Remove every `FOUNDRY_DEEPEP_NVL_IPC` check. + +- [ ] **Step 3: Port and adapt the SGLang Buffer wrapper** + +Port `_patch_deepep` from source commits `afa9d36`, `b18b186`, `2328691`, and `05df555`, then replace its environment-variable policy with `get_deepep_transport()`. + +At module scope: + +```python +import inspect +``` + +The wrapper must: + +```python +supports_use_fabric = "use_fabric" in inspect.signature(original_init).parameters +``` + +and raise: + +```python +RuntimeError( + "Foundry SGLang DeepEP requires a deep_ep.Buffer constructor " + "with the use_fabric parameter" +) +``` + +before calling an incompatible DeepEP version. + +Call `_patch_deepep()` from `install_hooks()` after configuration is loaded. + +- [ ] **Step 4: Record or validate SGLang's effective profile** + +Inside the wrapper, after applying the configured policy but before calling the original constructor: + +```python +if mode == CUDAGraphExtensionMode.SAVE: + rt.update_warmup_state_deepep_profile( + transport, + effective_num_nvl_bytes, + ) +elif mode == CUDAGraphExtensionMode.LOAD: + state = rt.load_warmup_state() + validate_num_nvl_bytes_match( + archived=state.deepep_num_nvl_bytes, + effective=effective_num_nvl_bytes, + transport=transport, + ) +``` + +Do not alter existing NVSHMEM initialization ordering. + +- [ ] **Step 5: Run policy and import tests** + +```bash +pytest tests/test_deepep_transport.py tests/test_deepep_transport_vllm.py \ + tests/test_imports.py -q +``` + +Expected: PASS in the engine environment. The tests must demonstrate both kwargs profiles, mismatch rejection, old-DeepEP rejection, and NONE-mode no-op behavior. + +- [ ] **Step 6: Commit and push** + +```bash +git add python/foundry/integration/vllm/hooks.py \ + python/foundry/integration/sglang/hooks.py \ + tests/test_deepep_transport.py tests/test_deepep_transport_vllm.py +git commit -m "feat: apply DeepEP transport in engine hooks" +git push -u origin cursor/no-imex-deepep-f10b +``` + +--- + +### Task 6: Add No-IMEX Recipes and Documentation + +**Files:** +- Create/modify: `recipe/experimental/README.md` +- Create/modify: `recipe/experimental/foundry_save.toml` +- Create/modify: `recipe/experimental/foundry_load.toml` +- Create/modify: `recipe/experimental/sglang_foundry_save.toml` +- Create/modify: `recipe/experimental/sglang_foundry_load.toml` +- Create/modify: `recipe/experimental/serve_qwen3-30ba3b_ipc_ep.sh` +- Create/modify: `recipe/experimental/serve_qwen3-30ba3bfp8_sglang_ipc_ep.sh` +- Modify: `docs/vllm/moe-and-deepep.md` +- Modify: `docs/sglang/hooks.md` +- Create/modify: `tests/test_experimental_recipe.py` + +**Interfaces:** +- Produces reproducible SAVE/LOAD commands that select `nvl_ipc` only through TOML. + +- [ ] **Step 1: Write failing recipe contract tests** + +Assert every no-IMEX SAVE/LOAD TOML contains: + +```toml +deepep_transport = "nvl_ipc" +``` + +Assert each engine's SAVE and LOAD pair shares `workspace_root`, `base_addr`, and `region_size`. + +Assert no experimental script contains: + +```bash +FOUNDRY_DEEPEP_NVL_IPC +``` + +- [ ] **Step 2: Port focused experimental recipes** + +Port the EP-only recipe files from `84d2b95`, `403fed2`, and `eb7a0d2`. Exclude all tensor-parallel recipes. + +Replace the source branch's environment-variable selection with the TOML field: + +```toml +deepep_transport = "nvl_ipc" +``` + +Keep: + +```bash +export NCCL_CUMEM_ENABLE=0 +export NCCL_NVLS_ENABLE=0 +``` + +and, for SGLang: + +```bash +export NVSHMEM_CUMEM_HANDLE_TYPE=FILE_DESCRIPTOR +``` + +- [ ] **Step 3: Update integration documentation** + +Document: + +- the two allowed values and `fabric` default; +- intranode-only scope of `nvl_ipc`; +- continued DeepEP/NVSHMEM requirement; +- SAVE/LOAD archive parity and fail-fast errors; +- SGLang's required `Buffer(use_fabric=...)` API; +- no silent fallback; +- the two-GPU validation commands. + +- [ ] **Step 4: Run recipe and formatting tests** + +```bash +pytest tests/test_experimental_recipe.py -q +ruff check python tests +ruff format --check python tests +git diff --check +``` + +Expected: PASS. + +- [ ] **Step 5: Commit and push** + +```bash +git add recipe/experimental docs/vllm/moe-and-deepep.md \ + docs/sglang/hooks.md tests/test_experimental_recipe.py +git commit -m "docs: add no-IMEX DeepEP recipes" +git push -u origin cursor/no-imex-deepep-f10b +``` + +--- + +### Task 7: Verify Native and Engine End-to-End Behavior + +**Files:** +- Test: `tests/test_deepep_fabric.py` +- Test: `tests/run_deepep_matrix.sh` +- Test: vLLM and SGLang experimental recipes + +**Interfaces:** +- Consumes the completed Python, native hook, archive, and recipe work. +- Produces runtime evidence for standalone DeepEP, vLLM inference, SGLang inference, and negative parity validation. + +- [ ] **Step 1: Run the CPU and static suite** + +```bash +pytest tests/test_deepep_transport.py tests/test_deepep_transport_vllm.py \ + tests/test_experimental_recipe.py tests/test_imports.py -q +ruff check python tests +ruff format --check python tests +git diff --check +``` + +Expected: PASS. + +- [ ] **Step 2: Build the exact branch on Modal H100** + +```bash +CC=gcc CXX=g++ pip install -e . --no-build-isolation +``` + +Confirm `libcuda_hook.so` loads and run: + +```bash +pytest tests/test_load_graph.py tests/test_async_load.py \ + tests/test_vmm_alloc.py tests/test_preallocation.py -q +``` + +Expected: PASS. + +- [ ] **Step 3: Run the two-GPU standalone matrix** + +Without initializing CUDA in the orchestration parent: + +```bash +bash tests/run_deepep_matrix.sh ll_nofabric both +bash tests/run_deepep_matrix.sh nvl_ipc both +bash tests/run_deepep_matrix.sh nvl_ipc_prealloc both +``` + +Expected: all PASS, deterministic dispatch output, no CUDA error 999. + +- [ ] **Step 4: Validate vLLM SAVE/LOAD and inference** + +Run two SAVE passes using the same working directory: + +```bash +bash recipe/experimental/serve_qwen3-30ba3b_ipc_ep.sh 2 --save +bash recipe/experimental/serve_qwen3-30ba3b_ipc_ep.sh 2 --save +bash recipe/experimental/serve_qwen3-30ba3b_ipc_ep.sh 2 --load +``` + +Issue one completion request and require coherent output. Check logs for: + +- successful VMM IPC imports or expected relocation messages; +- identical per-rank final allocation offsets; +- no `cuMemImportFromShareableHandle` error 999; +- no graph replay or NVSHMEM initialization error. + +- [ ] **Step 5: Validate SGLang SAVE/LOAD and inference** + +Run: + +```bash +bash recipe/experimental/serve_qwen3-30ba3bfp8_sglang_ipc_ep.sh 2 --save +bash recipe/experimental/serve_qwen3-30ba3bfp8_sglang_ipc_ep.sh 2 --load +``` + +Issue one completion request and apply the same success checks as vLLM. + +- [ ] **Step 6: Run the negative archive-parity control** + +After producing an `nvl_ipc` archive, point LOAD at an otherwise identical TOML containing: + +```toml +deepep_transport = "fabric" +``` + +Expected: startup exits before CUDA modules, DeepEP buffers, or graphs are restored, with: + +```text +Foundry DeepEP transport mismatch: archived=nvl_ipc, configured=fabric +``` + +- [ ] **Step 7: Review the final diff** + +Verify: + +```bash +git status --short +git diff main...HEAD --stat +git diff main...HEAD --check +``` + +Confirm there are no TP-only files, no `FOUNDRY_DEEPEP_NVL_IPC` references, no temporary logs, and no debug-only instrumentation. + +- [ ] **Step 8: Commit any test-driven fixes, push, and update the PR** + +If validation required a fix, return to the task that owns the affected file, +rerun that task's focused test command, and use its explicit `git add` command. +Write a concrete commit subject naming the observed defect, then push: + +```bash +git push -u origin cursor/no-imex-deepep-f10b +``` + +Update the PR description with the passing command matrix and engine inference evidence. From 785aadf1c7ac8fd12aacbc3b0712a2c226fea7ce Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Thu, 23 Jul 2026 07:13:44 +0000 Subject: [PATCH 03/59] fix: pass default CUDA graph pool Co-authored-by: Rahul Chalamala --- python/foundry/graph.py | 2 +- tests/test_load_graph.py | 12 ++++++++++++ 2 files changed, 13 insertions(+), 1 deletion(-) diff --git a/python/foundry/graph.py b/python/foundry/graph.py index 427d1599..cfeffb7d 100644 --- a/python/foundry/graph.py +++ b/python/foundry/graph.py @@ -135,7 +135,7 @@ def __init__( if self.__class__.default_capture_stream is None: self.__class__.default_capture_stream = torch.cuda.Stream() - self.pool = () if pool is None else (pool,) + self.pool = (pool,) self.capture_stream = ( stream if stream is not None else self.__class__.default_capture_stream ) diff --git a/tests/test_load_graph.py b/tests/test_load_graph.py index 0347f5b3..ba799388 100644 --- a/tests/test_load_graph.py +++ b/tests/test_load_graph.py @@ -6,6 +6,7 @@ import sys from pathlib import Path +import foundry import pytest import torch import torch.nn as nn @@ -170,6 +171,17 @@ def _spawn_with_preload(launch_mode): subprocess.check_call(cmd, env=env) +def test_graph_default_pool_is_forwarded(monkeypatch): + sentinel_stream = object() + monkeypatch.setattr(torch.cuda, "Stream", lambda: sentinel_stream) + cuda_graph = object() + + context = foundry.graph.__new__(foundry.graph) + foundry.graph.__init__(context, cuda_graph) + + assert context.pool == (None,) + + @pytest.mark.filterwarnings("ignore:TORCH_CUDA_ARCH_LIST is not set") def test_graph_save_and_load(): """Test CUDA graph save/load with binary dumping and allocator replay""" From 7ed212dda510c7a6c4a668588c17fdcc1b65d989 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Thu, 23 Jul 2026 07:21:23 +0000 Subject: [PATCH 04/59] test: cover cross-process VMM IPC Co-authored-by: Rahul Chalamala --- tests/run_deepep_matrix.sh | 105 ++++++++++++++++++++++++++++++++++++ tests/test_deepep_fabric.py | 57 ++++++++++++++++++-- tests/test_vmm_alloc.py | 55 +++++++++++++++++-- 3 files changed, 209 insertions(+), 8 deletions(-) create mode 100755 tests/run_deepep_matrix.sh diff --git a/tests/run_deepep_matrix.sh b/tests/run_deepep_matrix.sh new file mode 100755 index 00000000..c0db658a --- /dev/null +++ b/tests/run_deepep_matrix.sh @@ -0,0 +1,105 @@ +#!/usr/bin/env bash +# Matrix driver for tests/test_deepep_fabric.py — standalone DeepEP save/load probes +# (no vLLM/sglang init; needs ~4 GB free on each of GPUs 0,1, NOT a fully empty GPU). +# +# ll_nofabric pure low-latency (NVSHMEM heap, FD handles), NCCL fast paths off +# -> the "EP without fabric" baseline; expected PASS +# ll_ncclcumem same, but NCCL_CUMEM_ENABLE=1 -> probe: can NCCL cumem coexist? +# ll_ncclnvls same, but CUMEM=1 + NVLS=1 -> probe: can NCCL NVLS coexist? +# nvl_ipc + 64 MB NVL buffer, use_fabric=0 -> legacy cudaIpc path through the +# hook's VMM-IPC translation (SCM_RIGHTS fd transport; peer mappings +# relocate under shared region bases); expected PASS +# nvl_fabric + 64 MB NVL buffer, use_fabric=1 -> fabric cuMemCreate on a no-IMEX +# machine; documents the fabric dependency; expected FAIL +# +# Usage: run_deepep_matrix.sh +# +# Environment overrides: +# PYTHON_BIN Python with Foundry installed (default: python3) +# VLLM_ROOT vLLM checkout used to build NVSHMEM +# NVSHMEM_SO explicit libnvshmem_host.so path +# DEEPEP_MATRIX_WORK_ROOT per-case work directory root +# DEEPEP_MATRIX_LOG_DIR log output directory +set -euo pipefail + +SCRIPT_DIR=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd) +REPO_ROOT=$(cd -- "$SCRIPT_DIR/.." && pwd) +PYTHON_BIN=${PYTHON_BIN:-python3} +VLLM_ROOT=${VLLM_ROOT:-"$REPO_ROOT/../vllm"} +NVSHMEM_SO=${NVSHMEM_SO:-"$VLLM_ROOT/tools/ep_kernels/ep_kernels_workspace/nvshmem/lib/libnvshmem_host.so"} +DEEPEP_MATRIX_WORK_ROOT=${DEEPEP_MATRIX_WORK_ROOT:-"$REPO_ROOT/tests/deepep_matrix_work"} +DEEPEP_MATRIX_LOG_DIR=${DEEPEP_MATRIX_LOG_DIR:-"$REPO_ROOT/logs/deepep_matrix"} +TEST="$REPO_ROOT/tests/test_deepep_fabric.py" + +CASE=${1:?usage: run_deepep_matrix.sh } +PHASE=${2:?usage: run_deepep_matrix.sh } + +export CUDA_VISIBLE_DEVICES=0,1 +# Pin NVSHMEM heap chunks to POSIX FD handles — explicit "no fabric handles anywhere". +# (Auto-detect picks FD on non-MNNVL machines anyway; pinning removes the variable.) +export NVSHMEM_CUMEM_HANDLE_TYPE=FILE_DESCRIPTOR + +case "$CASE" in + ll_nofabric) export TEST_USE_FABRIC=0 TEST_NVL_BYTES_MB=0 NCCL_CUMEM_ENABLE=0 NCCL_NVLS_ENABLE=0; PORT=29611 ;; + ll_ncclcumem) export TEST_USE_FABRIC=0 TEST_NVL_BYTES_MB=0 NCCL_CUMEM_ENABLE=1 NCCL_NVLS_ENABLE=0; PORT=29621 ;; + ll_ncclcumem_preinit) export TEST_USE_FABRIC=0 TEST_NVL_BYTES_MB=0 NCCL_CUMEM_ENABLE=1 NCCL_NVLS_ENABLE=0 TEST_NCCL_WARMUP_PRE_REGION=1; PORT=29661 ;; + ll_ncclnvls) export TEST_USE_FABRIC=0 TEST_NVL_BYTES_MB=0 NCCL_CUMEM_ENABLE=1 NCCL_NVLS_ENABLE=1; PORT=29631 ;; + ll_ncclnvls_preinit) export TEST_USE_FABRIC=0 TEST_NVL_BYTES_MB=0 NCCL_CUMEM_ENABLE=1 NCCL_NVLS_ENABLE=1 TEST_NCCL_WARMUP_PRE_REGION=1; PORT=29671 ;; + nvl_ipc) export TEST_USE_FABRIC=0 TEST_NVL_BYTES_MB=64 NCCL_CUMEM_ENABLE=0 NCCL_NVLS_ENABLE=0; PORT=29641 ;; + nvl_fabric) export TEST_USE_FABRIC=1 TEST_NVL_BYTES_MB=64 NCCL_CUMEM_ENABLE=0 NCCL_NVLS_ENABLE=0; PORT=29651 ;; + nvl_ipc_prealloc) export TEST_USE_FABRIC=0 TEST_NVL_BYTES_MB=64 NCCL_CUMEM_ENABLE=0 NCCL_NVLS_ENABLE=0 TEST_LOAD_PREALLOC_MB=1024; PORT=29681 ;; + *) echo "unknown case: $CASE" >&2; exit 2 ;; +esac + +case "$PHASE" in + save|load|both) ;; + *) echo "unknown phase: $PHASE" >&2; exit 2 ;; +esac + +if [[ ! -f "$NVSHMEM_SO" ]]; then + echo "NVSHMEM library not found: $NVSHMEM_SO" >&2 + exit 4 +fi + +HOOK_SO=$( + "$PYTHON_BIN" -c \ + "import foundry.ops, pathlib; print(pathlib.Path(foundry.ops.__file__).parent / 'libcuda_hook.so')" +) +if [[ ! -f "$HOOK_SO" ]]; then + echo "Foundry CUDA hook not found: $HOOK_SO" >&2 + exit 5 +fi + +# GPUs are shared with other users — refuse to run without headroom. +for i in 0 1; do + free=$(nvidia-smi --query-gpu=memory.free --format=csv,noheader,nounits -i "$i") + if [ "$free" -lt 6000 ]; then + echo "GPU $i has only ${free} MiB free — aborting (need ~6 GB headroom)"; exit 3 + fi +done + +# main() overwrites TEST_USE_FABRIC from the CLI flag, so pass the flag too. +FABRIC_ARGS=() +[ "$TEST_USE_FABRIC" = "0" ] && FABRIC_ARGS+=("--no-fabric") + +WORK="$DEEPEP_MATRIX_WORK_ROOT/$CASE" +LOGDIR="$DEEPEP_MATRIX_LOG_DIR" +mkdir -p "$WORK" "$LOGDIR" +cd "$WORK" # deepep_fabric_archive/ is CWD-relative + +export LD_PRELOAD="$NVSHMEM_SO:$HOOK_SO${LD_PRELOAD:+:$LD_PRELOAD}" + +run_phase() { + local p=$1 + # distinct rendezvous port per phase to dodge TIME_WAIT + if [ "$p" = save ]; then export MASTER_PORT=$PORT; else export MASTER_PORT=$((PORT + 1)); fi + echo "=== case=$CASE phase=$p $(date '+%F %T') ===" + echo "env: TEST_USE_FABRIC=$TEST_USE_FABRIC TEST_NVL_BYTES_MB=$TEST_NVL_BYTES_MB" \ + "NCCL_CUMEM_ENABLE=$NCCL_CUMEM_ENABLE NCCL_NVLS_ENABLE=$NCCL_NVLS_ENABLE" \ + "NVSHMEM_CUMEM_HANDLE_TYPE=$NVSHMEM_CUMEM_HANDLE_TYPE MASTER_PORT=$MASTER_PORT" + if [ "$p" = save ]; then rm -rf deepep_fabric_archive; fi + "$PYTHON_BIN" "$TEST" "--$p" "${FABRIC_ARGS[@]}" --num-processes=2 \ + 2>&1 | tee "$LOGDIR/${CASE}_${p}.log" +} + +if [ "$PHASE" = both ]; then run_phase save; run_phase load; else run_phase "$PHASE"; fi diff --git a/tests/test_deepep_fabric.py b/tests/test_deepep_fabric.py index be6be723..b1395582 100644 --- a/tests/test_deepep_fabric.py +++ b/tests/test_deepep_fabric.py @@ -77,6 +77,27 @@ def _init_dist(local_rank: int, num_local_ranks: int): ) +def _maybe_warm_nccl_pre_region(group): + """Diagnostic knob (TEST_NCCL_WARMUP_PRE_REGION=1): force NCCL transport + setup BEFORE the foundry region exists, so NCCL's cuMem allocations + (NCCL_CUMEM_ENABLE=1) pass through untracked. Localizes whether the + NCCL-CUMEM incompatibility is "NCCL cuMem while region enabled" only. + The warmup tensor's caching-allocator segment is drained afterwards so + post-region allocations don't reuse an untracked (non-deterministic) + segment. + """ + if os.environ.get("TEST_NCCL_WARMUP_PRE_REGION", "0") != "1": + return + t = torch.ones(1024, 1024, device="cuda") + dist.all_reduce(t) # default group: ring/p2p transport setup + dist.all_reduce(t, group=group) # subgroup used by the DeepEP Buffer + dist.all_gather_object([None] * dist.get_world_size(), 0) # object-coll path + torch.cuda.synchronize() + del t + torch.cuda.empty_cache() # legal here: region not set yet, no captures + print("[TEST] NCCL pre-region warmup done (transports up before region)", flush=True) + + def _create_deterministic_inputs( rank: int, num_tokens: int, hidden: int, num_experts: int, num_topk: int ): @@ -263,6 +284,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) @@ -487,11 +510,22 @@ def _run_load(local_rank: int, num_processes: int): print(f"[Rank {rank}] LOAD: Loading CUDA modules from {rank_archive}", flush=True) fdry.load_cuda_modules_and_libraries(rank_archive) + _maybe_warm_nccl_pre_region(group) + # Step 2: Setup allocation region region_size = fdry.parse_size(REGION_SIZE_STR) print(f"[Rank {rank}] LOAD: Setting up allocation region at {hex(BASE_ADDR)}", flush=True) fdry.set_allocation_region(BASE_ADDR, region_size) + # Optional (TEST_LOAD_PREALLOC_MB=N): reproduce the production LOAD shape - + # one upfront-preallocated chunk that subsequent allocations (including the + # DeepEP NVL buffer) carve from with NO individual cuMem handle. This + # exercises the hook's whole-chunk VMM-IPC export path. + prealloc_mb = int(os.environ.get("TEST_LOAD_PREALLOC_MB", "0")) + if prealloc_mb > 0: + assert fdry.preallocate_region(prealloc_mb * 1024 * 1024), "preallocate_region failed" + print(f"[Rank {rank}] LOAD: preallocated {prealloc_mb} MB chunk", flush=True) + # Step 3: Create DeepEP buffer (initializes NVSHMEM runtime) num_local_experts = num_experts // num_ranks num_rdma_bytes = deep_ep.Buffer.get_low_latency_rdma_size_hint( @@ -499,13 +533,17 @@ def _run_load(local_rank: int, num_processes: int): ) use_fabric = os.environ.get("TEST_USE_FABRIC", "1") == "1" + # Must match SAVE: an NVL buffer allocated on SAVE but not LOAD (or vice versa) + # diverges the allocation trajectory and the peer-pointer setup. + num_nvl_bytes = int(os.environ.get("TEST_NVL_BYTES_MB", "0")) * 1024 * 1024 print( - f"[Rank {rank}] LOAD: Creating DeepEP buffer with use_fabric={use_fabric}, size={num_rdma_bytes / 1e6:.2f} MB", + f"[Rank {rank}] LOAD: Creating DeepEP buffer with use_fabric={use_fabric}, num_nvl_bytes={num_nvl_bytes / 1e6:.2f} MB, num_rdma_bytes={num_rdma_bytes / 1e6:.2f} MB", flush=True, ) buffer = deep_ep.Buffer( group, + num_nvl_bytes=num_nvl_bytes, num_rdma_bytes=num_rdma_bytes, low_latency_mode=True, num_qps_per_rank=num_local_experts, @@ -593,10 +631,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 diff --git a/tests/test_vmm_alloc.py b/tests/test_vmm_alloc.py index 5df909f0..7f928926 100644 --- a/tests/test_vmm_alloc.py +++ b/tests/test_vmm_alloc.py @@ -1,10 +1,12 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the Foundry project +import ctypes import os import subprocess import sys from pathlib import Path +import foundry as fdry import pytest import torch @@ -26,8 +28,6 @@ def _get_hook_so_path(): def _run_core(): - import foundry as fdry - torch.cuda.init() device = torch.device("cuda:0") @@ -104,8 +104,6 @@ def _run_core_without_region(): def _run_core_size_parsing(): - import foundry as fdry - torch.cuda.init() device = torch.device("cuda:0") @@ -146,6 +144,46 @@ def _run_core_size_parsing(): print("[TEST] test_size_parsing PASSED") +def _run_core_zero_alignment_reserve(): + torch.cuda.init() + + base_addr = 0x7F4000000000 + region_size = 1024 * 1024 * 1024 + default_alignment = 2 * 1024 * 1024 + deliberately_unaligned_offset = default_alignment + 12345 + reserve_size = 2 * 1024 * 1024 + + driver = ctypes.CDLL(None) + reserve = driver.cuMemAddressReserve + reserve.argtypes = [ + ctypes.POINTER(ctypes.c_uint64), + ctypes.c_size_t, + ctypes.c_size_t, + ctypes.c_uint64, + ctypes.c_ulonglong, + ] + reserve.restype = ctypes.c_int + address_free = driver.cuMemAddressFree + address_free.argtypes = [ctypes.c_uint64, ctypes.c_size_t] + address_free.restype = ctypes.c_int + + ptr = ctypes.c_uint64() + with fdry.allocation_region(base_addr, region_size): + fdry.set_current_alloc_offset(deliberately_unaligned_offset) + result = reserve(ctypes.byref(ptr), reserve_size, 0, 0, 0) + assert result == 0 + expected_ptr = (base_addr + deliberately_unaligned_offset + default_alignment - 1) & ~( + default_alignment - 1 + ) + assert ptr.value == expected_ptr, ( + f"zero-alignment reserve returned {hex(ptr.value)}, expected " + f"default-aligned {hex(expected_ptr)}" + ) + expected_offset = expected_ptr - base_addr + reserve_size + assert fdry.get_current_alloc_offset() == expected_offset + assert address_free(ptr.value, reserve_size) == 0 + + def _spawn_with_preload(test_mode): so_path = _get_hook_so_path() env = os.environ.copy() @@ -176,6 +214,12 @@ def test_size_parsing(): _spawn_with_preload("size-parsing") +@pytest.mark.filterwarnings("ignore:TORCH_CUDA_ARCH_LIST is not set") +def test_zero_alignment_reserve_stays_in_region(): + """CUDA's alignment=0 convention uses Foundry's default 2 MiB alignment.""" + _spawn_with_preload("zero-alignment-reserve") + + if __name__ == "__main__": if "--core" in sys.argv: _run_core() @@ -183,7 +227,10 @@ def test_size_parsing(): _run_core_without_region() elif "--size-parsing" in sys.argv: _run_core_size_parsing() + elif "--zero-alignment-reserve" in sys.argv: + _run_core_zero_alignment_reserve() else: test_vmm_allocation() test_vmm_allocation_without_region() test_size_parsing() + test_zero_alignment_reserve_stays_in_region() From c7fd76620173d74a510d7a091c7cdc01256bfaeb Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Thu, 23 Jul 2026 07:42:56 +0000 Subject: [PATCH 05/59] fix: support cross-process VMM IPC Co-authored-by: Rahul Chalamala --- csrc/hook.cpp | 649 +++++++++++++++++++++++++++++++++++++++++++++++--- 1 file changed, 612 insertions(+), 37 deletions(-) diff --git a/csrc/hook.cpp b/csrc/hook.cpp index d552d08d..360457b8 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; }); @@ -2702,9 +2732,10 @@ CUresult cuMemAddressReserve(CUdeviceptr* ptr, size_t size, size_t alignment, CU return result; } + const size_t effective_alignment = alignment == 0 ? kAllocAlignment : alignment; constexpr size_t kSmallAllocThreshold = 2ULL << 30; if (size < kSmallAllocThreshold) { - CUdeviceptr aligned_addr = align_to(tls_storage.current_alloc_base_addr, alignment); + CUdeviceptr aligned_addr = align_to(tls_storage.current_alloc_base_addr, effective_alignment); size_t region_base = (size_t)tls_storage.region.base; size_t region_end = region_base + tls_storage.region.size; @@ -2744,7 +2775,7 @@ CUresult cuMemAddressReserve(CUdeviceptr* ptr, size_t size, size_t alignment, CU // constexpr size_t kReserveAlignment = 2ULL << 37; // FIXME(liuxs): This value should be equal to // the first reservation size - CUdeviceptr hint_addr = align_to(tls_storage.current_vmm_reserve_addr, alignment); + CUdeviceptr hint_addr = align_to(tls_storage.current_vmm_reserve_addr, effective_alignment); CUresult result = real_func(ptr, size, alignment, hint_addr, flags); if (result != CUDA_SUCCESS) { fprintf(stderr, "[HOOK] ERROR: cuMemAddressReserve failed with error 0x%x\n", result); @@ -2759,7 +2790,7 @@ CUresult cuMemAddressReserve(CUdeviceptr* ptr, size_t size, size_t alignment, CU (unsigned long long)tls_storage.current_vmm_reserve_addr); } - tls_storage.current_vmm_reserve_addr = align_to(*ptr + size, alignment); + tls_storage.current_vmm_reserve_addr = align_to(*ptr + size, effective_alignment); #ifdef HOOK_DEBUG fprintf(stderr, @@ -2860,8 +2891,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 +3155,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 +3206,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 +3305,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 +3382,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 +3397,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 +3468,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 +3566,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 +3614,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 +3829,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 +3856,12 @@ void free_preallocated_region() { size_t preallocated_size = tls_storage.preallocated_end_addr - tls_storage.preallocated_start_addr; + { + std::lock_guard lock(preallocated_chunks_mutex); + preallocated_chunks.erase(tls_storage.preallocated_start_addr); + } + vmm_ipc_invalidate_export((CUdeviceptr)tls_storage.preallocated_start_addr); + mem_unmap_func(tls_storage.preallocated_start_addr, preallocated_size); mem_release_func(tls_storage.preallocated_handle); @@ -3464,7 +4036,10 @@ void replay_hook_events_from_json(const boost::json::object& events_obj) { size_t alignment = event_obj.at("alignment").to_number(); uint64_t expected_ptr = event_obj.at("ptr").to_number(); - CUdeviceptr aligned_addr = align_to(tls_storage.current_alloc_base_addr, alignment); + // Archives recorded before alignment normalization stored the caller's + // alignment (including CUDA's default-alignment sentinel, zero). + const size_t effective_alignment = alignment == 0 ? kAllocAlignment : alignment; + CUdeviceptr aligned_addr = align_to(tls_storage.current_alloc_base_addr, effective_alignment); if (aligned_addr != expected_ptr) { fprintf(stderr, "[REPLAY] FATAL: Reserve address mismatch!\n"); From ccda6c5ab431986444646475940608738ace9dc7 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Thu, 23 Jul 2026 08:31:25 +0000 Subject: [PATCH 06/59] test: cover DeepEP transport configuration Co-authored-by: Rahul Chalamala --- tests/test_deepep_transport.py | 101 ++++++++++++++++++++++++++++ tests/test_deepep_transport_vllm.py | 86 +++++++++++++++++++++++ 2 files changed, 187 insertions(+) create mode 100644 tests/test_deepep_transport.py create mode 100644 tests/test_deepep_transport_vllm.py diff --git a/tests/test_deepep_transport.py b/tests/test_deepep_transport.py new file mode 100644 index 00000000..70632019 --- /dev/null +++ b/tests/test_deepep_transport.py @@ -0,0 +1,101 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the Foundry project +from __future__ import annotations + +import pytest + +from foundry.integration.deepep_transport import ( + DeepEPTransport, + configure_deepep_buffer, + parse_deepep_transport, + validate_num_nvl_bytes_match, + validate_transport_match, +) +from foundry.integration.sglang import config as sglang_config + + +@pytest.mark.parametrize( + ("raw", "expected"), + [ + (None, DeepEPTransport.FABRIC), + ("fabric", DeepEPTransport.FABRIC), + ("nvl_ipc", DeepEPTransport.NVL_IPC), + ], +) +def test_parse_deepep_transport(raw, expected): + assert parse_deepep_transport(raw) is expected + + +def test_parse_deepep_transport_rejects_unknown_value(): + with pytest.raises(ValueError, match="deepep_transport"): + parse_deepep_transport("rdma") + + +@pytest.mark.parametrize( + ("transport", "expected_bytes", "expected_fabric"), + [ + (DeepEPTransport.FABRIC, 0, True), + (DeepEPTransport.NVL_IPC, 128, False), + ], +) +def test_configure_deepep_buffer(transport, expected_bytes, expected_fabric): + num_nvl_bytes, kwargs = configure_deepep_buffer( + transport, + num_nvl_bytes=128, + kwargs={"other": "kept"}, + ) + assert num_nvl_bytes == expected_bytes + assert kwargs == {"other": "kept", "use_fabric": expected_fabric} + + +def test_transport_mismatch_names_both_values(): + with pytest.raises(RuntimeError, match="archived=nvl_ipc.*configured=fabric"): + validate_transport_match( + archived="nvl_ipc", + configured=DeepEPTransport.FABRIC, + ) + + +def test_nvl_bytes_mismatch_names_both_values(): + with pytest.raises(RuntimeError, match="archived=64.*effective=128"): + validate_num_nvl_bytes_match( + archived=64, + effective=128, + transport=DeepEPTransport.NVL_IPC, + ) + + +@pytest.mark.parametrize( + ("contents", "expected"), + [ + ("", DeepEPTransport.FABRIC), + ('deepep_transport = "nvl_ipc"\n', DeepEPTransport.NVL_IPC), + ], +) +def test_sglang_config_parses_deepep_transport(tmp_path, contents, expected): + config_path = tmp_path / "config.toml" + config_path.write_text(contents) + + config = sglang_config.CUDAGraphExtensionConfig.from_toml(config_path) + + assert config.deepep_transport is expected + + +def test_sglang_config_rejects_unknown_deepep_transport(tmp_path): + config_path = tmp_path / "config.toml" + config_path.write_text('deepep_transport = "rdma"\n') + + with pytest.raises(ValueError, match="deepep_transport"): + sglang_config.CUDAGraphExtensionConfig.from_toml(config_path) + + +def test_sglang_get_deepep_transport_uses_loaded_config(tmp_path, monkeypatch): + config_path = tmp_path / "config.toml" + config_path.write_text('deepep_transport = "nvl_ipc"\n') + monkeypatch.setattr(sglang_config, "_config", None) + + assert sglang_config.get_deepep_transport() is DeepEPTransport.FABRIC + + sglang_config.load_graph_extension_config(str(config_path)) + + assert sglang_config.get_deepep_transport() is DeepEPTransport.NVL_IPC diff --git a/tests/test_deepep_transport_vllm.py b/tests/test_deepep_transport_vllm.py new file mode 100644 index 00000000..03eb6b1d --- /dev/null +++ b/tests/test_deepep_transport_vllm.py @@ -0,0 +1,86 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the Foundry project +from __future__ import annotations + +import importlib.util +import logging +import sys +import types +from pathlib import Path + +import pytest + +from foundry.integration.deepep_transport import DeepEPTransport + +VLLM_CONFIG_PATH = ( + Path(__file__).parents[1] / "python" / "foundry" / "integration" / "vllm" / "config.py" +) + + +def _init_stub_logger(name: str) -> logging.Logger: + return logging.getLogger(name) + + +@pytest.fixture +def vllm_config_module(monkeypatch): + if importlib.util.find_spec("vllm") is None: + vllm_module = types.ModuleType("vllm") + vllm_module.__path__ = [] + logger_module = types.ModuleType("vllm.logger") + logger_module.init_logger = _init_stub_logger + monkeypatch.setitem(sys.modules, "vllm", vllm_module) + monkeypatch.setitem(sys.modules, "vllm.logger", logger_module) + + spec = importlib.util.spec_from_file_location( + "foundry_vllm_config_under_test", + VLLM_CONFIG_PATH, + ) + assert spec is not None + assert spec.loader is not None + module = importlib.util.module_from_spec(spec) + monkeypatch.setitem(sys.modules, spec.name, module) + spec.loader.exec_module(module) + return module + + +@pytest.mark.parametrize( + ("contents", "expected"), + [ + ("", DeepEPTransport.FABRIC), + ('deepep_transport = "nvl_ipc"\n', DeepEPTransport.NVL_IPC), + ], +) +def test_vllm_config_parses_deepep_transport( + tmp_path, + vllm_config_module, + contents, + expected, +): + config_path = tmp_path / "config.toml" + config_path.write_text(contents) + + config = vllm_config_module.CUDAGraphExtensionConfig.from_toml(config_path) + + assert config.deepep_transport is expected + + +def test_vllm_config_rejects_unknown_deepep_transport(tmp_path, vllm_config_module): + config_path = tmp_path / "config.toml" + config_path.write_text('deepep_transport = "rdma"\n') + + with pytest.raises(ValueError, match="deepep_transport"): + vllm_config_module.CUDAGraphExtensionConfig.from_toml(config_path) + + +def test_vllm_get_deepep_transport_uses_loaded_config( + tmp_path, + vllm_config_module, +): + config_path = tmp_path / "config.toml" + config_path.write_text('deepep_transport = "nvl_ipc"\n') + + assert vllm_config_module.get_deepep_transport() is DeepEPTransport.FABRIC + + vllm_config_module.load_graph_extension_config(str(config_path)) + + assert vllm_config_module.get_deepep_transport() is DeepEPTransport.NVL_IPC From 42875919c64597905cf10257250e7856afa6a7bb Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Thu, 23 Jul 2026 08:40:15 +0000 Subject: [PATCH 07/59] feat: configure DeepEP transport in TOML Co-authored-by: Rahul Chalamala --- .../foundry/integration/deepep_transport.py | 62 +++++++++++++++++++ python/foundry/integration/sglang/config.py | 10 +++ python/foundry/integration/vllm/config.py | 10 +++ tests/test_deepep_transport.py | 1 - tests/test_deepep_transport_vllm.py | 1 - 5 files changed, 82 insertions(+), 2 deletions(-) create mode 100644 python/foundry/integration/deepep_transport.py diff --git a/python/foundry/integration/deepep_transport.py b/python/foundry/integration/deepep_transport.py new file mode 100644 index 00000000..435950e7 --- /dev/null +++ b/python/foundry/integration/deepep_transport.py @@ -0,0 +1,62 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the Foundry project +"""Shared DeepEP transport policy for Foundry engine integrations.""" + +from __future__ import annotations + +import enum +from typing import Any + + +class DeepEPTransport(str, enum.Enum): + FABRIC = "fabric" + NVL_IPC = "nvl_ipc" + + +def parse_deepep_transport(raw: str | None) -> DeepEPTransport: + value = DeepEPTransport.FABRIC.value if raw is None else raw + try: + return DeepEPTransport(value) + except ValueError as exc: + choices = ", ".join(item.value for item in DeepEPTransport) + raise ValueError(f"deepep_transport must be one of: {choices}; got {value!r}") from exc + + +def configure_deepep_buffer( + transport: DeepEPTransport, + *, + num_nvl_bytes: int, + kwargs: dict[str, Any], +) -> tuple[int, dict[str, Any]]: + configured = dict(kwargs) + if transport is DeepEPTransport.FABRIC: + configured["use_fabric"] = True + return 0, configured + configured["use_fabric"] = False + return num_nvl_bytes, configured + + +def validate_transport_match( + *, + archived: str | None, + configured: DeepEPTransport, +) -> None: + archived_transport = parse_deepep_transport(archived) + if archived_transport is not configured: + raise RuntimeError( + "Foundry DeepEP transport mismatch: " + f"archived={archived_transport.value}, configured={configured.value}" + ) + + +def validate_num_nvl_bytes_match( + *, + archived: int, + effective: int, + transport: DeepEPTransport, +) -> None: + if archived != effective: + raise RuntimeError( + "Foundry DeepEP num_nvl_bytes mismatch for " + f"{transport.value}: archived={archived}, effective={effective}" + ) diff --git a/python/foundry/integration/sglang/config.py b/python/foundry/integration/sglang/config.py index a0e1d41d..ee9d50a5 100644 --- a/python/foundry/integration/sglang/config.py +++ b/python/foundry/integration/sglang/config.py @@ -11,6 +11,8 @@ import tomllib +from foundry.integration.deepep_transport import DeepEPTransport, parse_deepep_transport + class CUDAGraphExtensionMode(str, enum.Enum): NONE = "none" @@ -23,6 +25,7 @@ class CUDAGraphExtensionConfig: mode: CUDAGraphExtensionMode = CUDAGraphExtensionMode.NONE hook_library_path: str | None = None nvshmem_host_path: str | None = None + deepep_transport: DeepEPTransport = DeepEPTransport.FABRIC base_addr: int = 0x600000000000 region_size: str = "64GB" workspace_root: str = "foundry_archive" @@ -49,6 +52,7 @@ def from_toml(cls, path: str | Path) -> CUDAGraphExtensionConfig: mode=CUDAGraphExtensionMode(data.get("mode", cls.mode.value)), hook_library_path=hook_library_path, nvshmem_host_path=nvshmem_host_path, + deepep_transport=parse_deepep_transport(data.get("deepep_transport")), base_addr=base_addr, region_size=data.get("region_size", cls.region_size), workspace_root=data.get("workspace_root", cls.workspace_root), @@ -131,6 +135,12 @@ def get_nvshmem_host_path() -> str | None: return _config.nvshmem_host_path +def get_deepep_transport() -> DeepEPTransport: + if _config is None: + return DeepEPTransport.FABRIC + return _config.deepep_transport + + def compute_workspace_rank(server_args, tp_rank: int, pp_rank: int, dp_rank: int | None) -> int: if getattr(server_args, "enable_dp_attention", False): return pp_rank * server_args.tp_size + tp_rank diff --git a/python/foundry/integration/vllm/config.py b/python/foundry/integration/vllm/config.py index b03c29e3..552260b6 100644 --- a/python/foundry/integration/vllm/config.py +++ b/python/foundry/integration/vllm/config.py @@ -17,6 +17,8 @@ import tomllib from vllm.logger import init_logger +from foundry.integration.deepep_transport import DeepEPTransport, parse_deepep_transport + if TYPE_CHECKING: from vllm.config import ParallelConfig @@ -35,6 +37,7 @@ class CUDAGraphExtensionConfig: hook_library_path: str | None = None # Path to libnvshmem_host.so for NVSHMEM-based kernels (e.g. DeepEP). nvshmem_host_path: str | None = None + deepep_transport: DeepEPTransport = DeepEPTransport.FABRIC base_addr: int = 0x600000000000 region_size: str = "64GB" workspace_root: str = "foundry_archive" @@ -63,6 +66,7 @@ def from_toml(cls, path: str | Path) -> CUDAGraphExtensionConfig: mode=mode, hook_library_path=hook_library_path, nvshmem_host_path=data.get("nvshmem_host_path"), + deepep_transport=parse_deepep_transport(data.get("deepep_transport")), base_addr=base_addr, region_size=data.get("region_size", cls.region_size), workspace_root=data.get("workspace_root", cls.workspace_root), @@ -127,6 +131,12 @@ def get_nvshmem_host_path() -> str | None: return _config.nvshmem_host_path +def get_deepep_transport() -> DeepEPTransport: + if _config is None: + return DeepEPTransport.FABRIC + return _config.deepep_transport + + def _compute_rank_from_parallel_config( parallel_config: ParallelConfig, ) -> int: diff --git a/tests/test_deepep_transport.py b/tests/test_deepep_transport.py index 70632019..1eb8068c 100644 --- a/tests/test_deepep_transport.py +++ b/tests/test_deepep_transport.py @@ -3,7 +3,6 @@ from __future__ import annotations import pytest - from foundry.integration.deepep_transport import ( DeepEPTransport, configure_deepep_buffer, diff --git a/tests/test_deepep_transport_vllm.py b/tests/test_deepep_transport_vllm.py index 03eb6b1d..1c49d59b 100644 --- a/tests/test_deepep_transport_vllm.py +++ b/tests/test_deepep_transport_vllm.py @@ -9,7 +9,6 @@ from pathlib import Path import pytest - from foundry.integration.deepep_transport import DeepEPTransport VLLM_CONFIG_PATH = ( From d6165b75b5d383bba8e3e711240be095f10a5859 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Thu, 23 Jul 2026 08:54:17 +0000 Subject: [PATCH 08/59] test: cover archived DeepEP transport profiles Signed-off-by: Cursor Agent Co-authored-by: Rahul Chalamala --- tests/test_deepep_transport.py | 116 +++++++++++++++++++++ tests/test_deepep_transport_vllm.py | 153 ++++++++++++++++++++++++++-- 2 files changed, 262 insertions(+), 7 deletions(-) diff --git a/tests/test_deepep_transport.py b/tests/test_deepep_transport.py index 1eb8068c..ed44a334 100644 --- a/tests/test_deepep_transport.py +++ b/tests/test_deepep_transport.py @@ -2,6 +2,9 @@ # SPDX-FileCopyrightText: Copyright contributors to the Foundry project from __future__ import annotations +import json +from types import SimpleNamespace + import pytest from foundry.integration.deepep_transport import ( DeepEPTransport, @@ -11,6 +14,7 @@ validate_transport_match, ) from foundry.integration.sglang import config as sglang_config +from foundry.integration.sglang import runtime as sglang_runtime @pytest.mark.parametrize( @@ -98,3 +102,115 @@ def test_sglang_get_deepep_transport_uses_loaded_config(tmp_path, monkeypatch): sglang_config.load_graph_extension_config(str(config_path)) assert sglang_config.get_deepep_transport() is DeepEPTransport.NVL_IPC + + +def test_sglang_warmup_state_round_trips_deepep_profile(tmp_path, monkeypatch): + config = sglang_config.CUDAGraphExtensionConfig(workspace_root=str(tmp_path)) + monkeypatch.setattr(sglang_config, "_config", config) + monkeypatch.setattr(sglang_runtime, "_state", None) + state = sglang_runtime.WarmupState() + state.deepep_transport = "nvl_ipc" + state.deepep_num_nvl_bytes = 64 * 1024 * 1024 + + sglang_runtime.save_warmup_state(state) + + contents = json.loads((tmp_path / "warmup_state.json").read_text()) + assert contents["deepep_transport"] == "nvl_ipc" + assert contents["deepep_num_nvl_bytes"] == 64 * 1024 * 1024 + loaded = sglang_runtime.load_warmup_state() + assert loaded.deepep_transport == "nvl_ipc" + assert loaded.deepep_num_nvl_bytes == 64 * 1024 * 1024 + + +def test_sglang_legacy_warmup_state_uses_fabric_defaults(tmp_path, monkeypatch): + config = sglang_config.CUDAGraphExtensionConfig(workspace_root=str(tmp_path)) + monkeypatch.setattr(sglang_config, "_config", config) + (tmp_path / "warmup_state.json").write_text('{"sglang_version": "legacy"}') + + loaded = sglang_runtime.load_warmup_state() + + assert loaded.deepep_transport == "fabric" + assert loaded.deepep_num_nvl_bytes == 0 + + +def test_sglang_profile_update_preserves_existing_rank_guard(tmp_path, monkeypatch): + config = sglang_config.CUDAGraphExtensionConfig(workspace_root=str(tmp_path)) + monkeypatch.setattr(sglang_config, "_config", config) + sglang_runtime.save_warmup_state(sglang_runtime.WarmupState()) + monkeypatch.setattr( + sglang_runtime, + "_state", + sglang_runtime.CUDAGraphExtensionState(rank=1), + ) + + sglang_runtime.update_warmup_state_deepep_profile( + DeepEPTransport.NVL_IPC, + 64 * 1024 * 1024, + ) + + rank_one_state = sglang_runtime.load_warmup_state() + assert rank_one_state.deepep_transport == "fabric" + assert rank_one_state.deepep_num_nvl_bytes == 0 + + sglang_runtime._state.rank = 0 + sglang_runtime.update_warmup_state_deepep_profile( + DeepEPTransport.NVL_IPC, + 64 * 1024 * 1024, + ) + rank_zero_state = sglang_runtime.load_warmup_state() + assert rank_zero_state.deepep_transport == "nvl_ipc" + assert rank_zero_state.deepep_num_nvl_bytes == 64 * 1024 * 1024 + + +def test_sglang_load_rejects_transport_mismatch_before_cuda_operations( + tmp_path, + monkeypatch, +): + (tmp_path / "rank_0").mkdir() + (tmp_path / "warmup_state.json").write_text( + json.dumps( + { + "deepep_transport": "nvl_ipc", + "deepep_num_nvl_bytes": 64 * 1024 * 1024, + } + ) + ) + config = sglang_config.CUDAGraphExtensionConfig( + mode=sglang_config.CUDAGraphExtensionMode.LOAD, + deepep_transport=DeepEPTransport.FABRIC, + workspace_root=str(tmp_path), + ) + monkeypatch.setattr(sglang_config, "_config", config) + monkeypatch.setattr(sglang_runtime, "_state", None) + cuda_operations = [] + monkeypatch.setattr(sglang_runtime.cge, "set_skip_fatbin_processing", lambda value: None) + monkeypatch.setattr( + sglang_runtime.cge, + "load_cuda_modules_and_libraries", + lambda path: cuda_operations.append("load_cuda_modules_and_libraries"), + ) + monkeypatch.setattr( + sglang_runtime.cge, + "set_allocation_region", + lambda base_addr, size: cuda_operations.append("set_allocation_region"), + ) + monkeypatch.setattr( + sglang_runtime.torch._C, + "_cuda_getCurrentBlasHandle", + lambda: None, + ) + server_args = SimpleNamespace( + enable_dp_attention=False, + tp_size=1, + pp_size=1, + ) + + with pytest.raises(RuntimeError, match="archived=nvl_ipc.*configured=fabric"): + sglang_runtime.setup_graph_extension( + server_args, + tp_rank=0, + pp_rank=0, + dp_rank=None, + ) + + assert cuda_operations == [] diff --git a/tests/test_deepep_transport_vllm.py b/tests/test_deepep_transport_vllm.py index 1c49d59b..ce75f6a5 100644 --- a/tests/test_deepep_transport_vllm.py +++ b/tests/test_deepep_transport_vllm.py @@ -3,6 +3,7 @@ from __future__ import annotations import importlib.util +import json import logging import sys import types @@ -14,21 +15,30 @@ VLLM_CONFIG_PATH = ( Path(__file__).parents[1] / "python" / "foundry" / "integration" / "vllm" / "config.py" ) +VLLM_RUNTIME_PATH = ( + Path(__file__).parents[1] / "python" / "foundry" / "integration" / "vllm" / "runtime.py" +) def _init_stub_logger(name: str) -> logging.Logger: return logging.getLogger(name) +def _stub_vllm_if_missing(monkeypatch) -> None: + if importlib.util.find_spec("vllm") is not None: + return + vllm_module = types.ModuleType("vllm") + vllm_module.__path__ = [] + vllm_module.__version__ = "test" + logger_module = types.ModuleType("vllm.logger") + logger_module.init_logger = _init_stub_logger + monkeypatch.setitem(sys.modules, "vllm", vllm_module) + monkeypatch.setitem(sys.modules, "vllm.logger", logger_module) + + @pytest.fixture def vllm_config_module(monkeypatch): - if importlib.util.find_spec("vllm") is None: - vllm_module = types.ModuleType("vllm") - vllm_module.__path__ = [] - logger_module = types.ModuleType("vllm.logger") - logger_module.init_logger = _init_stub_logger - monkeypatch.setitem(sys.modules, "vllm", vllm_module) - monkeypatch.setitem(sys.modules, "vllm.logger", logger_module) + _stub_vllm_if_missing(monkeypatch) spec = importlib.util.spec_from_file_location( "foundry_vllm_config_under_test", @@ -42,6 +52,22 @@ def vllm_config_module(monkeypatch): return module +@pytest.fixture +def vllm_runtime_module(monkeypatch): + _stub_vllm_if_missing(monkeypatch) + + spec = importlib.util.spec_from_file_location( + "foundry_vllm_runtime_under_test", + VLLM_RUNTIME_PATH, + ) + assert spec is not None + assert spec.loader is not None + module = importlib.util.module_from_spec(spec) + monkeypatch.setitem(sys.modules, spec.name, module) + spec.loader.exec_module(module) + return module + + @pytest.mark.parametrize( ("contents", "expected"), [ @@ -83,3 +109,116 @@ def test_vllm_get_deepep_transport_uses_loaded_config( vllm_config_module.load_graph_extension_config(str(config_path)) assert vllm_config_module.get_deepep_transport() is DeepEPTransport.NVL_IPC + + +def test_vllm_warmup_state_round_trips_deepep_profile(tmp_path, vllm_runtime_module): + state = vllm_runtime_module.WarmupState() + state.deepep_transport = "nvl_ipc" + state.deepep_num_nvl_bytes = 64 * 1024 * 1024 + + vllm_runtime_module.save_warmup_state(str(tmp_path), state) + + contents = json.loads((tmp_path / "warmup_state.json").read_text()) + assert contents["deepep_transport"] == "nvl_ipc" + assert contents["deepep_num_nvl_bytes"] == 64 * 1024 * 1024 + loaded = vllm_runtime_module.load_warmup_state(str(tmp_path)) + assert loaded.deepep_transport == "nvl_ipc" + assert loaded.deepep_num_nvl_bytes == 64 * 1024 * 1024 + + +def test_vllm_legacy_warmup_state_uses_fabric_defaults(tmp_path, vllm_runtime_module): + (tmp_path / "warmup_state.json").write_text('{"vllm_version": "legacy"}') + + loaded = vllm_runtime_module.load_warmup_state(str(tmp_path)) + + assert loaded.deepep_transport == "fabric" + assert loaded.deepep_num_nvl_bytes == 0 + + +def test_vllm_profile_update_is_rank_zero_only( + tmp_path, + monkeypatch, + vllm_runtime_module, +): + vllm_runtime_module.save_warmup_state( + str(tmp_path), + vllm_runtime_module.WarmupState(), + ) + config = types.SimpleNamespace(workspace_dir=str(tmp_path / "rank_1")) + monkeypatch.setattr(vllm_runtime_module, "get_config", lambda: config) + + vllm_runtime_module.update_warmup_state_deepep_profile( + str(tmp_path), + DeepEPTransport.NVL_IPC, + 64 * 1024 * 1024, + ) + + rank_one_state = vllm_runtime_module.load_warmup_state(str(tmp_path)) + assert rank_one_state.deepep_transport == "fabric" + assert rank_one_state.deepep_num_nvl_bytes == 0 + + config.workspace_dir = str(tmp_path / "rank_0") + vllm_runtime_module.update_warmup_state_deepep_profile( + str(tmp_path), + DeepEPTransport.NVL_IPC, + 64 * 1024 * 1024, + ) + rank_zero_state = vllm_runtime_module.load_warmup_state(str(tmp_path)) + assert rank_zero_state.deepep_transport == "nvl_ipc" + assert rank_zero_state.deepep_num_nvl_bytes == 64 * 1024 * 1024 + + +def test_vllm_load_rejects_transport_mismatch_before_cuda_operations( + tmp_path, + monkeypatch, + vllm_runtime_module, +): + (tmp_path / "rank_0").mkdir() + (tmp_path / "warmup_state.json").write_text( + json.dumps( + { + "deepep_transport": "nvl_ipc", + "deepep_num_nvl_bytes": 64 * 1024 * 1024, + } + ) + ) + config = types.SimpleNamespace( + mode=vllm_runtime_module.CUDAGraphExtensionMode.LOAD, + deepep_transport=DeepEPTransport.FABRIC, + workspace_root=str(tmp_path), + workspace_dir=None, + base_addr=0x600000000000, + region_size="1GB", + ) + monkeypatch.setattr(vllm_runtime_module, "get_config", lambda: config) + monkeypatch.setattr(vllm_runtime_module, "_state", None) + cuda_operations = [] + monkeypatch.setattr( + vllm_runtime_module.fops, + "set_skip_fatbin_processing", + lambda value: None, + ) + monkeypatch.setattr( + vllm_runtime_module.fops, + "load_cuda_modules_and_libraries", + lambda path: cuda_operations.append("load_cuda_modules_and_libraries"), + ) + monkeypatch.setattr( + vllm_runtime_module.fops, + "set_allocation_region", + lambda base_addr, size: cuda_operations.append("set_allocation_region"), + ) + parallel_config = types.SimpleNamespace( + rank=0, + world_size=1, + data_parallel_index=0, + tensor_parallel_size=1, + pipeline_parallel_size=1, + data_parallel_size=1, + nnodes=1, + ) + + with pytest.raises(RuntimeError, match="archived=nvl_ipc.*configured=fabric"): + vllm_runtime_module.setup_graph_extension(parallel_config) + + assert cuda_operations == [] From d74f290a0fff1af3970821da0e60a219e091856e Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Thu, 23 Jul 2026 08:59:55 +0000 Subject: [PATCH 09/59] feat: validate archived DeepEP transport Signed-off-by: Cursor Agent Co-authored-by: Rahul Chalamala --- python/foundry/integration/sglang/runtime.py | 20 ++++++++++++++- python/foundry/integration/vllm/runtime.py | 26 +++++++++++++++++++- 2 files changed, 44 insertions(+), 2 deletions(-) diff --git a/python/foundry/integration/sglang/runtime.py b/python/foundry/integration/sglang/runtime.py index 1f9032c5..146977f5 100644 --- a/python/foundry/integration/sglang/runtime.py +++ b/python/foundry/integration/sglang/runtime.py @@ -17,6 +17,7 @@ from foundry import ops as cge from foundry.allocation_region import parse_size +from foundry.integration.deepep_transport import DeepEPTransport, validate_transport_match from foundry.integration.sglang.config import ( CUDAGraphExtensionMode, compute_workspace_rank, @@ -38,6 +39,8 @@ class WarmupState: gpu_total_memory: int = 0 memory_pool_config: dict = field(default_factory=dict) final_alloc_offset: int = 0 + deepep_transport: str = DeepEPTransport.FABRIC.value + deepep_num_nvl_bytes: int = 0 @dataclass @@ -109,6 +112,16 @@ def load_warmup_state() -> WarmupState: return WarmupState(**{k: v for k, v in data.items() if k in valid}) +def update_warmup_state_deepep_profile( + transport: DeepEPTransport, + num_nvl_bytes: int, +) -> None: + state = load_warmup_state() + state.deepep_transport = transport.value + state.deepep_num_nvl_bytes = num_nvl_bytes + save_warmup_state(state) + + def setup_graph_extension(server_args, tp_rank: int, pp_rank: int, dp_rank: int | None) -> None: """Set up the VMM region before SGLang initializes NCCL/process groups.""" global _state @@ -128,9 +141,14 @@ def setup_graph_extension(server_args, tp_rank: int, pp_rank: int, dp_rank: int shutil.rmtree(workspace_dir) workspace_dir.mkdir(parents=True, exist_ok=True) elif cfg.mode == CUDAGraphExtensionMode.LOAD: - cge.set_skip_fatbin_processing(True) if not workspace_dir.exists(): raise RuntimeError(f"Foundry workspace for rank {rank} does not exist: {workspace_dir}") + warmup_state = load_warmup_state() + validate_transport_match( + archived=warmup_state.deepep_transport, + configured=cfg.deepep_transport, + ) + cge.set_skip_fatbin_processing(True) cge.load_cuda_modules_and_libraries(str(workspace_dir)) region_size = parse_size(cfg.region_size) diff --git a/python/foundry/integration/vllm/runtime.py b/python/foundry/integration/vllm/runtime.py index c81e64cb..de4471fd 100644 --- a/python/foundry/integration/vllm/runtime.py +++ b/python/foundry/integration/vllm/runtime.py @@ -25,6 +25,7 @@ from foundry import ops as fops from foundry.allocation_region import parse_size +from foundry.integration.deepep_transport import DeepEPTransport, validate_transport_match from foundry.integration.vllm.config import ( CUDAGraphExtensionMode, _compute_rank_from_parallel_config, @@ -57,6 +58,8 @@ class WarmupState: num_cpu_blocks: int = 0 gpu_memory_utilization: float = 0.0 final_alloc_offset: int = 0 + deepep_transport: str = DeepEPTransport.FABRIC.value + deepep_num_nvl_bytes: int = 0 _final_alloc_offset: int = 0 @@ -99,6 +102,22 @@ def load_warmup_state(workspace_root: str) -> WarmupState: return WarmupState(**{k: v for k, v in data.items() if k in valid}) +def update_warmup_state_deepep_profile( + workspace_root: str, + transport: DeepEPTransport, + num_nvl_bytes: int, +) -> None: + cfg = get_config() + if cfg is None or cfg.workspace_dir is None: + return + if Path(cfg.workspace_dir).name != "rank_0": + return + state = load_warmup_state(workspace_root) + state.deepep_transport = transport.value + state.deepep_num_nvl_bytes = num_nvl_bytes + save_warmup_state(workspace_root, state) + + # --------------------------------------------------------------------------- # Worker-scope state (per-process, set by setup_graph_extension) # --------------------------------------------------------------------------- @@ -147,7 +166,6 @@ def setup_graph_extension(parallel_config: ParallelConfig) -> None: workspace_dir.mkdir(parents=True, exist_ok=True) elif cfg.mode == CUDAGraphExtensionMode.LOAD: - fops.set_skip_fatbin_processing(True) if not workspace_dir.exists(): # rank_0 fallback is only safe for single-rank setups. # Under TP/PP/DP > 1 each rank has its own device addresses, @@ -173,6 +191,12 @@ def setup_graph_extension(parallel_config: ParallelConfig) -> None: f"exist. For multi-rank LOAD, re-run SAVE with the same " f"parallelism topology." ) + warmup_state = load_warmup_state(cfg.workspace_root) + validate_transport_match( + archived=warmup_state.deepep_transport, + configured=cfg.deepep_transport, + ) + fops.set_skip_fatbin_processing(True) log.info("[foundry] LOAD: loading CUDA modules from %s", cfg.workspace_dir) t = time.perf_counter() fops.load_cuda_modules_and_libraries(cfg.workspace_dir) From b0234e585f1fdd4f185172cba2270b915ecd213e Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Thu, 23 Jul 2026 09:13:18 +0000 Subject: [PATCH 10/59] test: cover DeepEP transport engine hooks Signed-off-by: Cursor Agent Co-authored-by: Rahul Chalamala --- tests/test_deepep_transport.py | 236 ++++++++++++++++++++++++++++ tests/test_deepep_transport_vllm.py | 211 +++++++++++++++++++++++++ 2 files changed, 447 insertions(+) diff --git a/tests/test_deepep_transport.py b/tests/test_deepep_transport.py index ed44a334..696adcb6 100644 --- a/tests/test_deepep_transport.py +++ b/tests/test_deepep_transport.py @@ -3,6 +3,8 @@ from __future__ import annotations import json +import sys +import types from types import SimpleNamespace import pytest @@ -14,9 +16,17 @@ validate_transport_match, ) from foundry.integration.sglang import config as sglang_config +from foundry.integration.sglang import hooks as sglang_hooks from foundry.integration.sglang import runtime as sglang_runtime +def _install_fake_deep_ep(monkeypatch, buffer_cls): + deep_ep_module = types.ModuleType("deep_ep") + deep_ep_module.Buffer = buffer_cls + monkeypatch.setitem(sys.modules, "deep_ep", deep_ep_module) + monkeypatch.setattr(sglang_hooks, "_DEEPEP_PATCHED", False, raising=False) + + @pytest.mark.parametrize( ("raw", "expected"), [ @@ -214,3 +224,229 @@ def test_sglang_load_rejects_transport_mismatch_before_cuda_operations( ) assert cuda_operations == [] + + +@pytest.mark.parametrize( + ("transport", "expected"), + [ + ( + DeepEPTransport.FABRIC, + {"num_nvl_bytes": 0, "use_fabric": True}, + ), + ( + DeepEPTransport.NVL_IPC, + {"num_nvl_bytes": 64 * 1024 * 1024, "use_fabric": False}, + ), + ], +) +def test_sglang_hook_applies_transport_profile_on_save( + monkeypatch, + transport, + expected, +): + constructor_calls = [] + + class Buffer: + def __init__( + self, + group, + num_nvl_bytes=0, + num_rdma_bytes=0, + *, + use_fabric=None, + ): + constructor_calls.append( + { + "num_nvl_bytes": num_nvl_bytes, + "use_fabric": use_fabric, + } + ) + + _install_fake_deep_ep(monkeypatch, Buffer) + updates = [] + monkeypatch.setattr( + sglang_hooks, + "get_graph_extension_mode", + lambda: sglang_hooks.CUDAGraphExtensionMode.SAVE, + ) + monkeypatch.setattr(sglang_hooks, "get_deepep_transport", lambda: transport) + monkeypatch.setattr( + sglang_hooks.rt, + "update_warmup_state_deepep_profile", + lambda selected, num_nvl_bytes: updates.append((selected, num_nvl_bytes)), + ) + + sglang_hooks._patch_deepep() + Buffer("group", 64 * 1024 * 1024, 32 * 1024 * 1024) + + assert constructor_calls == [expected] + assert updates == [(transport, expected["num_nvl_bytes"])] + + +def test_sglang_hook_leaves_none_mode_profile_unchanged(monkeypatch): + constructor_calls = [] + + class Buffer: + def __init__( + self, + group, + num_nvl_bytes=0, + num_rdma_bytes=0, + *, + use_fabric=None, + ): + constructor_calls.append( + { + "num_nvl_bytes": num_nvl_bytes, + "num_rdma_bytes": num_rdma_bytes, + "use_fabric": use_fabric, + } + ) + + _install_fake_deep_ep(monkeypatch, Buffer) + monkeypatch.setattr( + sglang_hooks, + "get_graph_extension_mode", + lambda: sglang_hooks.CUDAGraphExtensionMode.NONE, + ) + monkeypatch.setattr( + sglang_hooks.rt, + "update_warmup_state_deepep_profile", + lambda *args: pytest.fail("NONE mode must not update the archive"), + ) + monkeypatch.setattr( + sglang_hooks.rt, + "load_warmup_state", + lambda: pytest.fail("NONE mode must not load the archive"), + ) + + sglang_hooks._patch_deepep() + Buffer( + "group", + 64 * 1024 * 1024, + 32 * 1024 * 1024, + use_fabric=True, + ) + + assert constructor_calls == [ + { + "num_nvl_bytes": 64 * 1024 * 1024, + "num_rdma_bytes": 32 * 1024 * 1024, + "use_fabric": True, + } + ] + + +def test_sglang_load_rejects_buffer_size_before_original_constructor(monkeypatch): + constructor_calls = [] + + class Buffer: + def __init__( + self, + group, + num_nvl_bytes=0, + num_rdma_bytes=0, + *, + use_fabric=None, + ): + constructor_calls.append((group, num_nvl_bytes, num_rdma_bytes, use_fabric)) + + _install_fake_deep_ep(monkeypatch, Buffer) + monkeypatch.setattr( + sglang_hooks, + "get_graph_extension_mode", + lambda: sglang_hooks.CUDAGraphExtensionMode.LOAD, + ) + monkeypatch.setattr( + sglang_hooks, + "get_deepep_transport", + lambda: DeepEPTransport.NVL_IPC, + ) + monkeypatch.setattr( + sglang_hooks.rt, + "load_warmup_state", + lambda: SimpleNamespace(deepep_num_nvl_bytes=64 * 1024 * 1024), + ) + + sglang_hooks._patch_deepep() + + with pytest.raises(RuntimeError, match="archived=67108864.*effective=134217728"): + Buffer("group", 128 * 1024 * 1024, 32 * 1024 * 1024) + + assert constructor_calls == [] + + +def test_sglang_rejects_old_deepep_before_original_constructor(monkeypatch): + constructor_calls = [] + + class Buffer: + def __init__(self, group, num_nvl_bytes=0, num_rdma_bytes=0): + constructor_calls.append((group, num_nvl_bytes, num_rdma_bytes)) + + _install_fake_deep_ep(monkeypatch, Buffer) + monkeypatch.setattr( + sglang_hooks, + "get_graph_extension_mode", + lambda: sglang_hooks.CUDAGraphExtensionMode.SAVE, + ) + + sglang_hooks._patch_deepep() + + with pytest.raises( + RuntimeError, + match=( + "Foundry SGLang DeepEP requires a deep_ep.Buffer constructor " + "with the use_fabric parameter" + ), + ): + Buffer("group", 64 * 1024 * 1024, 32 * 1024 * 1024) + + assert constructor_calls == [] + + +def test_sglang_save_fails_before_buffer_if_shared_archive_is_missing( + tmp_path, + monkeypatch, +): + constructor_calls = [] + + class Buffer: + def __init__( + self, + group, + num_nvl_bytes=0, + num_rdma_bytes=0, + *, + use_fabric=None, + ): + constructor_calls.append((group, num_nvl_bytes, num_rdma_bytes, use_fabric)) + + _install_fake_deep_ep(monkeypatch, Buffer) + config = sglang_config.CUDAGraphExtensionConfig( + mode=sglang_config.CUDAGraphExtensionMode.SAVE, + deepep_transport=DeepEPTransport.NVL_IPC, + workspace_root=str(tmp_path), + ) + monkeypatch.setattr(sglang_config, "_config", config) + monkeypatch.setattr( + sglang_runtime, + "_state", + sglang_runtime.CUDAGraphExtensionState(rank=0), + ) + monkeypatch.setattr( + sglang_hooks, + "get_graph_extension_mode", + lambda: sglang_hooks.CUDAGraphExtensionMode.SAVE, + ) + monkeypatch.setattr( + sglang_hooks, + "get_deepep_transport", + lambda: DeepEPTransport.NVL_IPC, + ) + + sglang_hooks._patch_deepep() + + with pytest.raises(RuntimeError, match="warmup state file not found"): + Buffer("group", 64 * 1024 * 1024, 32 * 1024 * 1024) + + assert constructor_calls == [] diff --git a/tests/test_deepep_transport_vllm.py b/tests/test_deepep_transport_vllm.py index ce75f6a5..cda50230 100644 --- a/tests/test_deepep_transport_vllm.py +++ b/tests/test_deepep_transport_vllm.py @@ -18,6 +18,9 @@ VLLM_RUNTIME_PATH = ( Path(__file__).parents[1] / "python" / "foundry" / "integration" / "vllm" / "runtime.py" ) +VLLM_HOOKS_PATH = ( + Path(__file__).parents[1] / "python" / "foundry" / "integration" / "vllm" / "hooks.py" +) def _init_stub_logger(name: str) -> logging.Logger: @@ -68,6 +71,50 @@ def vllm_runtime_module(monkeypatch): return module +@pytest.fixture +def vllm_hooks_module(monkeypatch): + _stub_vllm_if_missing(monkeypatch) + + spec = importlib.util.spec_from_file_location( + "foundry_vllm_hooks_under_test", + VLLM_HOOKS_PATH, + ) + assert spec is not None + assert spec.loader is not None + module = importlib.util.module_from_spec(spec) + monkeypatch.setitem(sys.modules, spec.name, module) + spec.loader.exec_module(module) + return module + + +def _install_fake_all2all_module(monkeypatch, manager_cls): + distributed_module = types.ModuleType("vllm.distributed") + distributed_module.__path__ = [] + communicators_module = types.ModuleType("vllm.distributed.device_communicators") + communicators_module.__path__ = [] + all2all_module = types.ModuleType("vllm.distributed.device_communicators.all2all") + all2all_module.DeepEPLLAll2AllManager = manager_cls + communicators_module.all2all = all2all_module + distributed_module.device_communicators = communicators_module + monkeypatch.setattr( + sys.modules["vllm"], + "distributed", + distributed_module, + raising=False, + ) + monkeypatch.setitem(sys.modules, "vllm.distributed", distributed_module) + monkeypatch.setitem( + sys.modules, + "vllm.distributed.device_communicators", + communicators_module, + ) + monkeypatch.setitem( + sys.modules, + "vllm.distributed.device_communicators.all2all", + all2all_module, + ) + + @pytest.mark.parametrize( ("contents", "expected"), [ @@ -222,3 +269,167 @@ def test_vllm_load_rejects_transport_mismatch_before_cuda_operations( vllm_runtime_module.setup_graph_extension(parallel_config) assert cuda_operations == [] + + +@pytest.mark.parametrize( + ("transport", "expected"), + [ + ( + DeepEPTransport.FABRIC, + {"num_nvl_bytes": 0, "use_fabric": True}, + ), + ( + DeepEPTransport.NVL_IPC, + {"num_nvl_bytes": 64 * 1024 * 1024, "use_fabric": False}, + ), + ], +) +def test_vllm_hook_applies_transport_profile_on_save( + tmp_path, + monkeypatch, + vllm_hooks_module, + transport, + expected, +): + class Manager: + def _make_all2all_kwargs(self): + return {"num_nvl_bytes": 64 * 1024 * 1024} + + _install_fake_all2all_module(monkeypatch, Manager) + updates = [] + monkeypatch.setattr( + vllm_hooks_module, + "get_graph_extension_mode", + lambda: vllm_hooks_module.CUDAGraphExtensionMode.SAVE, + ) + monkeypatch.setattr(vllm_hooks_module, "get_deepep_transport", lambda: transport) + monkeypatch.setattr(vllm_hooks_module, "get_workspace_root", lambda: str(tmp_path)) + monkeypatch.setattr( + vllm_hooks_module.rt, + "update_warmup_state_deepep_profile", + lambda root, selected, num_nvl_bytes: updates.append( + (root, selected, num_nvl_bytes) + ), + ) + + vllm_hooks_module._patch_deepep() + result = Manager()._make_all2all_kwargs() + + assert result == expected + assert updates == [(str(tmp_path), transport, expected["num_nvl_bytes"])] + + +def test_vllm_hook_leaves_none_mode_profile_unchanged(monkeypatch, vllm_hooks_module): + upstream = { + "num_nvl_bytes": 64 * 1024 * 1024, + "use_fabric": "upstream", + } + + class Manager: + def _make_all2all_kwargs(self): + return dict(upstream) + + _install_fake_all2all_module(monkeypatch, Manager) + monkeypatch.setattr( + vllm_hooks_module, + "get_graph_extension_mode", + lambda: vllm_hooks_module.CUDAGraphExtensionMode.NONE, + ) + monkeypatch.setattr( + vllm_hooks_module.rt, + "update_warmup_state_deepep_profile", + lambda *args: pytest.fail("NONE mode must not update the archive"), + ) + monkeypatch.setattr( + vllm_hooks_module.rt, + "load_warmup_state", + lambda *args: pytest.fail("NONE mode must not load the archive"), + ) + + vllm_hooks_module._patch_deepep() + + assert Manager()._make_all2all_kwargs() == upstream + + +def test_vllm_load_rejects_buffer_size_before_communicator_setup( + tmp_path, + monkeypatch, + vllm_hooks_module, +): + class Manager: + def _make_all2all_kwargs(self): + return {"num_nvl_bytes": 128 * 1024 * 1024} + + _install_fake_all2all_module(monkeypatch, Manager) + communicator_calls = [] + monkeypatch.setattr( + vllm_hooks_module, + "get_graph_extension_mode", + lambda: vllm_hooks_module.CUDAGraphExtensionMode.LOAD, + ) + monkeypatch.setattr( + vllm_hooks_module, + "get_deepep_transport", + lambda: DeepEPTransport.NVL_IPC, + ) + monkeypatch.setattr(vllm_hooks_module, "get_workspace_root", lambda: str(tmp_path)) + monkeypatch.setattr( + vllm_hooks_module.rt, + "load_warmup_state", + lambda root: types.SimpleNamespace(deepep_num_nvl_bytes=64 * 1024 * 1024), + ) + vllm_hooks_module._patch_deepep() + + def construct_communicator(): + kwargs = Manager()._make_all2all_kwargs() + communicator_calls.append(kwargs) + + with pytest.raises(RuntimeError, match="archived=67108864.*effective=134217728"): + construct_communicator() + + assert communicator_calls == [] + + +def test_vllm_save_bootstraps_shared_archive_before_deepep_patch( + tmp_path, + monkeypatch, + vllm_hooks_module, +): + monkeypatch.setattr(vllm_hooks_module, "_INSTALLED", False) + monkeypatch.delenv("FOUNDRY_SPAWN_T0_NS", raising=False) + monkeypatch.setattr(vllm_hooks_module, "_quarantine_get_device_capability", lambda: None) + monkeypatch.setattr(vllm_hooks_module, "load_graph_extension_config", lambda path: None) + monkeypatch.setattr( + vllm_hooks_module, + "get_graph_extension_mode", + lambda: vllm_hooks_module.CUDAGraphExtensionMode.SAVE, + ) + monkeypatch.setattr(vllm_hooks_module, "get_workspace_root", lambda: str(tmp_path)) + patch_names = ( + "_patch_init_worker_distributed_environment", + "_patch_kernel_warmup", + "_patch_dummy_runs", + "_patch_load_model", + "_patch_prepare_comm_buffer", + "_patch_capture_model", + "_patch_cuda_graph_wrapper_call", + "_patch_initialize_kv_caches", + "_patch_subprocess_spawn_sites", + ) + for patch_name in patch_names: + monkeypatch.setattr(vllm_hooks_module, patch_name, lambda: None) + + def assert_archive_exists(): + assert (tmp_path / "warmup_state.json").is_file() + + monkeypatch.setattr(vllm_hooks_module, "_patch_deepep", assert_archive_exists) + compilation_config = types.SimpleNamespace( + graph_extension_config_path=str(tmp_path / "config.toml"), + cudagraph_num_of_warmups=1, + ) + + vllm_hooks_module.install_hooks(compilation_config) + + contents = json.loads((tmp_path / "warmup_state.json").read_text()) + assert contents["deepep_transport"] == "fabric" + assert contents["deepep_num_nvl_bytes"] == 0 From 60477ff51103d5af679767179855f7237b36bebf Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Thu, 23 Jul 2026 09:24:55 +0000 Subject: [PATCH 11/59] test: preserve vLLM first-save profiling Signed-off-by: Cursor Agent Co-authored-by: Rahul Chalamala --- tests/test_deepep_transport_vllm.py | 104 ++++++++++++++++++++++++++++ 1 file changed, 104 insertions(+) diff --git a/tests/test_deepep_transport_vllm.py b/tests/test_deepep_transport_vllm.py index cda50230..5c28da98 100644 --- a/tests/test_deepep_transport_vllm.py +++ b/tests/test_deepep_transport_vllm.py @@ -115,6 +115,36 @@ def _install_fake_all2all_module(monkeypatch, manager_cls): ) +def _install_fake_kv_cache_modules(monkeypatch, engine_core_cls, scheduler_config): + v1_module = types.ModuleType("vllm.v1") + v1_module.__path__ = [] + engine_module = types.ModuleType("vllm.v1.engine") + engine_module.__path__ = [] + core_module = types.ModuleType("vllm.v1.engine.core") + core_module.EngineCore = engine_core_cls + core_utils_module = types.ModuleType("vllm.v1.core.kv_cache_utils") + core_utils_module.get_kv_cache_configs = ( + lambda vllm_config, specs, available_memory: [ + types.SimpleNamespace( + specs=specs, + available_memory=available_memory, + ) + ] + ) + core_utils_module.generate_scheduler_kv_cache_config = lambda configs: scheduler_config + envs_module = types.ModuleType("vllm.envs") + envs_module.VLLM_ELASTIC_EP_SCALE_UP_LAUNCH = False + engine_module.core = core_module + v1_module.engine = engine_module + monkeypatch.setattr(sys.modules["vllm"], "v1", v1_module, raising=False) + monkeypatch.setattr(sys.modules["vllm"], "envs", envs_module, raising=False) + monkeypatch.setitem(sys.modules, "vllm.v1", v1_module) + monkeypatch.setitem(sys.modules, "vllm.v1.engine", engine_module) + monkeypatch.setitem(sys.modules, "vllm.v1.engine.core", core_module) + monkeypatch.setitem(sys.modules, "vllm.v1.core.kv_cache_utils", core_utils_module) + monkeypatch.setitem(sys.modules, "vllm.envs", envs_module) + + @pytest.mark.parametrize( ("contents", "expected"), [ @@ -433,3 +463,77 @@ def assert_archive_exists(): contents = json.loads((tmp_path / "warmup_state.json").read_text()) assert contents["deepep_transport"] == "fabric" assert contents["deepep_num_nvl_bytes"] == 0 + + +def test_vllm_bootstrap_archive_does_not_skip_first_save_profile( + tmp_path, + monkeypatch, + vllm_hooks_module, +): + profile_calls = [] + initialized_configs = [] + + class ModelExecutor: + def get_kv_cache_specs(self): + return [object()] + + def determine_available_memory(self): + profile_calls.append("determine_available_memory") + return [123456] + + def initialize_from_config(self, configs): + initialized_configs.append(configs) + + class EngineCore: + def __init__(self): + self.model_executor = ModelExecutor() + self.available_gpu_memory_for_kv_cache = 0 + + def _initialize_kv_caches(self, vllm_config): + pytest.fail("SAVE mode must use Foundry's KV-cache initialization") + + scheduler_config = types.SimpleNamespace(num_blocks=7, kv_cache_groups=[]) + _install_fake_kv_cache_modules( + monkeypatch, + EngineCore, + scheduler_config, + ) + bootstrap = vllm_hooks_module.rt.WarmupState( + deepep_transport="nvl_ipc", + deepep_num_nvl_bytes=64 * 1024 * 1024, + ) + vllm_hooks_module.rt.save_warmup_state(str(tmp_path), bootstrap) + monkeypatch.setattr( + vllm_hooks_module, + "get_graph_extension_mode", + lambda: vllm_hooks_module.CUDAGraphExtensionMode.SAVE, + ) + monkeypatch.setattr(vllm_hooks_module, "get_workspace_root", lambda: str(tmp_path)) + monkeypatch.setattr( + vllm_hooks_module.rt, + "create_warmup_state", + vllm_hooks_module.rt.WarmupState, + ) + monkeypatch.setattr(vllm_hooks_module.rt, "get_final_alloc_offset", lambda: 0) + vllm_hooks_module._patch_initialize_kv_caches() + vllm_config = types.SimpleNamespace( + model_config=types.SimpleNamespace(max_model_len=1024), + parallel_config=types.SimpleNamespace(data_parallel_index=0), + cache_config=types.SimpleNamespace( + num_gpu_blocks=0, + block_size=0, + gpu_memory_utilization=0.9, + ), + validate_block_size=lambda: None, + ) + + result = EngineCore()._initialize_kv_caches(vllm_config) + + assert result is scheduler_config + assert profile_calls == ["determine_available_memory"] + assert len(initialized_configs) == 1 + saved = vllm_hooks_module.rt.load_warmup_state(str(tmp_path)) + assert saved.available_gpu_memory == [123456] + assert saved.num_gpu_blocks == 7 + assert saved.deepep_transport == "nvl_ipc" + assert saved.deepep_num_nvl_bytes == 64 * 1024 * 1024 From db3555cae006c00e3798d1e88a5af9a26a1e4260 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Thu, 23 Jul 2026 09:28:11 +0000 Subject: [PATCH 12/59] feat: apply DeepEP transport in engine hooks Signed-off-by: Cursor Agent Co-authored-by: Rahul Chalamala --- python/foundry/integration/sglang/hooks.py | 71 ++++++++++++++ python/foundry/integration/vllm/hooks.py | 107 +++++++++++++++------ 2 files changed, 149 insertions(+), 29 deletions(-) diff --git a/python/foundry/integration/sglang/hooks.py b/python/foundry/integration/sglang/hooks.py index 22cd7cfa..9078acf3 100644 --- a/python/foundry/integration/sglang/hooks.py +++ b/python/foundry/integration/sglang/hooks.py @@ -5,14 +5,20 @@ from __future__ import annotations import functools +import inspect import logging import os import time from dataclasses import asdict +from foundry.integration.deepep_transport import ( + configure_deepep_buffer, + validate_num_nvl_bytes_match, +) from foundry.integration.sglang import runtime as rt from foundry.integration.sglang.config import ( CUDAGraphExtensionMode, + get_deepep_transport, get_graph_extension_mode, get_workspace_root, load_graph_extension_config, @@ -20,6 +26,7 @@ logger = logging.getLogger(__name__) _INSTALLED = False +_DEEPEP_PATCHED = False def _ep_lazy_init_needed() -> bool: @@ -92,11 +99,75 @@ def install_hooks(server_args) -> None: _patch_kernel_warmup() _patch_cuda_graph_capture() _patch_spawn_sites() + _patch_deepep() _INSTALLED = True logger.info("[Foundry] SGLang hooks installed") +def _patch_deepep() -> None: + """Apply Foundry's configured transport to SGLang's DeepEP buffer.""" + global _DEEPEP_PATCHED + if _DEEPEP_PATCHED: + return + + try: + # DeepEP is an optional SGLang dependency, so import it only when + # installing the engine-specific patch. + from deep_ep import Buffer + except Exception: + return + + original_init = Buffer.__init__ + supports_use_fabric = "use_fabric" in inspect.signature(original_init).parameters + + @functools.wraps(original_init) + def patched(self, group, num_nvl_bytes=0, num_rdma_bytes=0, *args, **kwargs): + mode = get_graph_extension_mode() + if mode != CUDAGraphExtensionMode.NONE: + if not supports_use_fabric: + raise RuntimeError( + "Foundry SGLang DeepEP requires a deep_ep.Buffer constructor " + "with the use_fabric parameter" + ) + + transport = get_deepep_transport() + effective_num_nvl_bytes, configured = configure_deepep_buffer( + transport, + num_nvl_bytes=int(num_nvl_bytes), + kwargs=kwargs, + ) + kwargs.clear() + kwargs.update(configured) + num_nvl_bytes = effective_num_nvl_bytes + + if mode == CUDAGraphExtensionMode.SAVE: + rt.update_warmup_state_deepep_profile( + transport, + effective_num_nvl_bytes, + ) + elif mode == CUDAGraphExtensionMode.LOAD: + state = rt.load_warmup_state() + validate_num_nvl_bytes_match( + archived=state.deepep_num_nvl_bytes, + effective=effective_num_nvl_bytes, + transport=transport, + ) + + return original_init( + self, + group, + num_nvl_bytes, + num_rdma_bytes, + *args, + **kwargs, + ) + + Buffer.__init__ = patched + _DEEPEP_PATCHED = True + logger.info("[Foundry] SGLang DeepEP transport patch installed") + + def _patch_init_torch_distributed() -> None: from sglang.srt.model_executor import model_runner as mr diff --git a/python/foundry/integration/vllm/hooks.py b/python/foundry/integration/vllm/hooks.py index f7768630..4747056f 100644 --- a/python/foundry/integration/vllm/hooks.py +++ b/python/foundry/integration/vllm/hooks.py @@ -19,9 +19,14 @@ import torch from vllm.logger import init_logger +from foundry.integration.deepep_transport import ( + configure_deepep_buffer, + validate_num_nvl_bytes_match, +) from foundry.integration.vllm import runtime as rt from foundry.integration.vllm.config import ( CUDAGraphExtensionMode, + get_deepep_transport, get_graph_extension_mode, get_workspace_root, load_graph_extension_config, @@ -112,13 +117,23 @@ def install_hooks(compilation_config: CompilationConfig) -> None: ) load_graph_extension_config(cfg_path) + mode = get_graph_extension_mode() + root = get_workspace_root() + if mode == CUDAGraphExtensionMode.SAVE and root is not None: + warmup_state_path = os.path.join(root, "warmup_state.json") + if not os.path.exists(warmup_state_path): + # DeepEP communicator kwargs are built during worker model load, + # before EngineCore profiles the KV cache. Seed a CPU-only state in + # the orchestration parent so every rank can update it safely. + os.makedirs(root, exist_ok=True) + rt.save_warmup_state(root, rt.WarmupState()) log.info( "[foundry] install_hooks: mode=%s workspace=%s", - get_graph_extension_mode().value, - get_workspace_root(), + mode.value, + root, ) - if get_graph_extension_mode() != CUDAGraphExtensionMode.NONE: + if mode != CUDAGraphExtensionMode.NONE: # NOTE(liuxs): vllm doesn't allow zero cudagraph_num_of_warmups, hardcode here compilation_config.cudagraph_num_of_warmups = 0 log.info( @@ -557,22 +572,31 @@ def patched(self, vllm_config): root = get_workspace_root() ws = None + has_saved_kv_profile = False if root and os.path.exists(os.path.join(root, "warmup_state.json")): ws = rt.load_warmup_state(root) - log.info( - "[foundry] %s mode: warmup_state.json found, skipping KV profiling", - mode.value.upper(), - ) + has_saved_kv_profile = bool(ws.available_gpu_memory) + if has_saved_kv_profile: + log.info( + "[foundry] %s mode: warmup_state.json found, skipping KV profiling", + mode.value.upper(), + ) + else: + log.info( + "[foundry] %s mode: bootstrap warmup_state.json found; " + "running KV profiling", + mode.value.upper(), + ) specs = self.model_executor.get_kv_cache_specs() has_kv = any(s for s in specs) - if ws is not None: + if has_saved_kv_profile: available_mem = ws.available_gpu_memory num_gpu = ws.num_gpu_blocks num_cpu = ws.num_cpu_blocks elif mode == CUDAGraphExtensionMode.LOAD: - raise RuntimeError("foundry LOAD requires warmup_state.json") + raise RuntimeError("foundry LOAD requires a profiled warmup_state.json") elif has_kv: if envs.VLLM_ELASTIC_EP_SCALE_UP_LAUNCH: assert self.available_gpu_memory_for_kv_cache > 0 @@ -597,7 +621,7 @@ def patched(self, vllm_config): args=(vllm_config.model_config.max_model_len,), ) sched = generate_scheduler_kv_cache_config(cfgs) - if ws is None: + if not has_saved_kv_profile: num_gpu = sched.num_blocks elif sched.num_blocks != num_gpu: log.warning( @@ -621,20 +645,19 @@ def patched(self, vllm_config): dp_shard = getattr(vllm_config.parallel_config, "data_parallel_index", 0) or 0 if mode == CUDAGraphExtensionMode.SAVE and root is not None and dp_shard == 0: os.makedirs(root, exist_ok=True) - if ws is not None: - ws.final_alloc_offset = rt.get_final_alloc_offset() - rt.save_warmup_state(root, ws) + state = ws if ws is not None else rt.create_warmup_state() + state.available_gpu_memory = available_mem + state.num_gpu_blocks = num_gpu + state.num_cpu_blocks = num_cpu + state.gpu_memory_utilization = vllm_config.cache_config.gpu_memory_utilization + state.final_alloc_offset = rt.get_final_alloc_offset() + rt.save_warmup_state(root, state) + if has_saved_kv_profile: log.info( - "[foundry] Updated warmup state: final_alloc_offset=%d", ws.final_alloc_offset + "[foundry] Updated warmup state: final_alloc_offset=%d", + state.final_alloc_offset, ) else: - new = rt.create_warmup_state() - new.available_gpu_memory = available_mem - new.num_gpu_blocks = num_gpu - new.num_cpu_blocks = num_cpu - new.gpu_memory_utilization = vllm_config.cache_config.gpu_memory_utilization - new.final_alloc_offset = rt.get_final_alloc_offset() - rt.save_warmup_state(root, new) log.info("[foundry] Saved warmup state: num_gpu_blocks=%d", num_gpu) return sched @@ -675,12 +698,14 @@ def patched_init(self, *args, **kwargs): # --------------------------------------------------------------------------- -# DeepEP buffer kwargs: use_fabric=True, num_nvl_bytes=0 on any foundry mode. +# DeepEP buffer kwargs: apply the configured transport on SAVE/LOAD. # --------------------------------------------------------------------------- def _patch_deepep() -> None: try: + # vLLM is an optional integration dependency, so keep its engine + # import inside the patch installer. from vllm.distributed.device_communicators import all2all except Exception: return @@ -694,13 +719,37 @@ def _patch_deepep() -> None: @functools.wraps(orig) def patched(self, *args, **kwargs): out = orig(self, *args, **kwargs) - if isinstance(out, dict) and get_graph_extension_mode() != CUDAGraphExtensionMode.NONE: - # Force VMM-fabric buffers - # so the foundry cuMem hook tracks them, and disable the - # NVLink buffer (LL-mode uses RDMA) to avoid fabric-cuMemCreate - # collisions with preallocated region. - out["use_fabric"] = True - out["num_nvl_bytes"] = 0 + mode = get_graph_extension_mode() + if not isinstance(out, dict) or mode == CUDAGraphExtensionMode.NONE: + return out + + transport = get_deepep_transport() + upstream_num_nvl_bytes = int(out.get("num_nvl_bytes", 0)) + effective_num_nvl_bytes, configured = configure_deepep_buffer( + transport, + num_nvl_bytes=upstream_num_nvl_bytes, + kwargs=out, + ) + out.clear() + out.update(configured) + out["num_nvl_bytes"] = effective_num_nvl_bytes + + root = get_workspace_root() + if root is None: + raise RuntimeError("Foundry vLLM DeepEP requires workspace_root") + if mode == CUDAGraphExtensionMode.SAVE: + rt.update_warmup_state_deepep_profile( + root, + transport, + effective_num_nvl_bytes, + ) + elif mode == CUDAGraphExtensionMode.LOAD: + state = rt.load_warmup_state(root) + validate_num_nvl_bytes_match( + archived=state.deepep_num_nvl_bytes, + effective=effective_num_nvl_bytes, + transport=transport, + ) return out cls._make_all2all_kwargs = patched From d8164580b9b9b603f5b6e8ae8c33eda68cdba4fa Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Thu, 23 Jul 2026 09:30:44 +0000 Subject: [PATCH 13/59] style: format DeepEP hook changes Signed-off-by: Cursor Agent Co-authored-by: Rahul Chalamala --- python/foundry/integration/vllm/hooks.py | 3 +-- tests/test_deepep_transport_vllm.py | 18 +++++++----------- 2 files changed, 8 insertions(+), 13 deletions(-) diff --git a/python/foundry/integration/vllm/hooks.py b/python/foundry/integration/vllm/hooks.py index 4747056f..66a20e40 100644 --- a/python/foundry/integration/vllm/hooks.py +++ b/python/foundry/integration/vllm/hooks.py @@ -583,8 +583,7 @@ def patched(self, vllm_config): ) else: log.info( - "[foundry] %s mode: bootstrap warmup_state.json found; " - "running KV profiling", + "[foundry] %s mode: bootstrap warmup_state.json found; running KV profiling", mode.value.upper(), ) diff --git a/tests/test_deepep_transport_vllm.py b/tests/test_deepep_transport_vllm.py index 5c28da98..2b086399 100644 --- a/tests/test_deepep_transport_vllm.py +++ b/tests/test_deepep_transport_vllm.py @@ -123,14 +123,12 @@ def _install_fake_kv_cache_modules(monkeypatch, engine_core_cls, scheduler_confi core_module = types.ModuleType("vllm.v1.engine.core") core_module.EngineCore = engine_core_cls core_utils_module = types.ModuleType("vllm.v1.core.kv_cache_utils") - core_utils_module.get_kv_cache_configs = ( - lambda vllm_config, specs, available_memory: [ - types.SimpleNamespace( - specs=specs, - available_memory=available_memory, - ) - ] - ) + core_utils_module.get_kv_cache_configs = lambda vllm_config, specs, available_memory: [ + types.SimpleNamespace( + specs=specs, + available_memory=available_memory, + ) + ] core_utils_module.generate_scheduler_kv_cache_config = lambda configs: scheduler_config envs_module = types.ModuleType("vllm.envs") envs_module.VLLM_ELASTIC_EP_SCALE_UP_LAUNCH = False @@ -337,9 +335,7 @@ def _make_all2all_kwargs(self): monkeypatch.setattr( vllm_hooks_module.rt, "update_warmup_state_deepep_profile", - lambda root, selected, num_nvl_bytes: updates.append( - (root, selected, num_nvl_bytes) - ), + lambda root, selected, num_nvl_bytes: updates.append((root, selected, num_nvl_bytes)), ) vllm_hooks_module._patch_deepep() From efc8728ad83f94b30955de4545b57c461ca9db30 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Thu, 23 Jul 2026 10:03:08 +0000 Subject: [PATCH 14/59] test: preserve vLLM first-save metadata Signed-off-by: Cursor Agent Co-authored-by: Rahul Chalamala --- tests/test_deepep_transport_vllm.py | 56 ++++++++++++++++++++++++----- 1 file changed, 47 insertions(+), 9 deletions(-) diff --git a/tests/test_deepep_transport_vllm.py b/tests/test_deepep_transport_vllm.py index 2b086399..88084d10 100644 --- a/tests/test_deepep_transport_vllm.py +++ b/tests/test_deepep_transport_vllm.py @@ -416,7 +416,7 @@ def construct_communicator(): assert communicator_calls == [] -def test_vllm_save_bootstraps_shared_archive_before_deepep_patch( +def test_vllm_install_does_not_write_shared_archive( tmp_path, monkeypatch, vllm_hooks_module, @@ -445,10 +445,12 @@ def test_vllm_save_bootstraps_shared_archive_before_deepep_patch( for patch_name in patch_names: monkeypatch.setattr(vllm_hooks_module, patch_name, lambda: None) - def assert_archive_exists(): - assert (tmp_path / "warmup_state.json").is_file() - - monkeypatch.setattr(vllm_hooks_module, "_patch_deepep", assert_archive_exists) + deepep_patch_calls = [] + monkeypatch.setattr( + vllm_hooks_module, + "_patch_deepep", + lambda: deepep_patch_calls.append("patched"), + ) compilation_config = types.SimpleNamespace( graph_extension_config_path=str(tmp_path / "config.toml"), cudagraph_num_of_warmups=1, @@ -456,9 +458,8 @@ def assert_archive_exists(): vllm_hooks_module.install_hooks(compilation_config) - contents = json.loads((tmp_path / "warmup_state.json").read_text()) - assert contents["deepep_transport"] == "fabric" - assert contents["deepep_num_nvl_bytes"] == 0 + assert deepep_patch_calls == ["patched"] + assert not (tmp_path / "warmup_state.json").exists() def test_vllm_bootstrap_archive_does_not_skip_first_save_profile( @@ -505,10 +506,17 @@ def _initialize_kv_caches(self, vllm_config): lambda: vllm_hooks_module.CUDAGraphExtensionMode.SAVE, ) monkeypatch.setattr(vllm_hooks_module, "get_workspace_root", lambda: str(tmp_path)) + runtime_metadata = { + "vllm_version": "test-vllm", + "timestamp": "2026-07-23T10:00:00", + "cuda_version": "13.0", + "gpu_name": "H100", + "gpu_total_memory": 80 * 1024**3, + } monkeypatch.setattr( vllm_hooks_module.rt, "create_warmup_state", - vllm_hooks_module.rt.WarmupState, + lambda: vllm_hooks_module.rt.WarmupState(**runtime_metadata), ) monkeypatch.setattr(vllm_hooks_module.rt, "get_final_alloc_offset", lambda: 0) vllm_hooks_module._patch_initialize_kv_caches() @@ -531,5 +539,35 @@ def _initialize_kv_caches(self, vllm_config): saved = vllm_hooks_module.rt.load_warmup_state(str(tmp_path)) assert saved.available_gpu_memory == [123456] assert saved.num_gpu_blocks == 7 + for field, expected in runtime_metadata.items(): + assert getattr(saved, field) == expected + assert saved.deepep_transport == "nvl_ipc" + assert saved.deepep_num_nvl_bytes == 64 * 1024 * 1024 + + +def test_vllm_profile_update_creates_missing_archive_only_on_rank_zero( + tmp_path, + monkeypatch, + vllm_runtime_module, +): + config = types.SimpleNamespace(workspace_dir=str(tmp_path / "rank_1")) + monkeypatch.setattr(vllm_runtime_module, "get_config", lambda: config) + + vllm_runtime_module.update_warmup_state_deepep_profile( + str(tmp_path), + DeepEPTransport.NVL_IPC, + 64 * 1024 * 1024, + ) + + assert not (tmp_path / "warmup_state.json").exists() + + config.workspace_dir = str(tmp_path / "rank_0") + vllm_runtime_module.update_warmup_state_deepep_profile( + str(tmp_path), + DeepEPTransport.NVL_IPC, + 64 * 1024 * 1024, + ) + + saved = vllm_runtime_module.load_warmup_state(str(tmp_path)) assert saved.deepep_transport == "nvl_ipc" assert saved.deepep_num_nvl_bytes == 64 * 1024 * 1024 From 69efcd0e51b54c15ba217ea47a0337709bc200fd Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Thu, 23 Jul 2026 10:05:13 +0000 Subject: [PATCH 15/59] fix: preserve vLLM first-save metadata Signed-off-by: Cursor Agent Co-authored-by: Rahul Chalamala --- python/foundry/integration/vllm/hooks.py | 16 +++++++--------- python/foundry/integration/vllm/runtime.py | 4 +++- 2 files changed, 10 insertions(+), 10 deletions(-) diff --git a/python/foundry/integration/vllm/hooks.py b/python/foundry/integration/vllm/hooks.py index 66a20e40..d292b896 100644 --- a/python/foundry/integration/vllm/hooks.py +++ b/python/foundry/integration/vllm/hooks.py @@ -119,14 +119,6 @@ def install_hooks(compilation_config: CompilationConfig) -> None: load_graph_extension_config(cfg_path) mode = get_graph_extension_mode() root = get_workspace_root() - if mode == CUDAGraphExtensionMode.SAVE and root is not None: - warmup_state_path = os.path.join(root, "warmup_state.json") - if not os.path.exists(warmup_state_path): - # DeepEP communicator kwargs are built during worker model load, - # before EngineCore profiles the KV cache. Seed a CPU-only state in - # the orchestration parent so every rank can update it safely. - os.makedirs(root, exist_ok=True) - rt.save_warmup_state(root, rt.WarmupState()) log.info( "[foundry] install_hooks: mode=%s workspace=%s", mode.value, @@ -644,7 +636,13 @@ def patched(self, vllm_config): dp_shard = getattr(vllm_config.parallel_config, "data_parallel_index", 0) or 0 if mode == CUDAGraphExtensionMode.SAVE and root is not None and dp_shard == 0: os.makedirs(root, exist_ok=True) - state = ws if ws is not None else rt.create_warmup_state() + if has_saved_kv_profile: + state = ws + else: + state = rt.create_warmup_state() + if ws is not None: + state.deepep_transport = ws.deepep_transport + state.deepep_num_nvl_bytes = ws.deepep_num_nvl_bytes state.available_gpu_memory = available_mem state.num_gpu_blocks = num_gpu state.num_cpu_blocks = num_cpu diff --git a/python/foundry/integration/vllm/runtime.py b/python/foundry/integration/vllm/runtime.py index de4471fd..7af22ff2 100644 --- a/python/foundry/integration/vllm/runtime.py +++ b/python/foundry/integration/vllm/runtime.py @@ -112,9 +112,11 @@ def update_warmup_state_deepep_profile( return if Path(cfg.workspace_dir).name != "rank_0": return - state = load_warmup_state(workspace_root) + path = os.path.join(workspace_root, "warmup_state.json") + state = load_warmup_state(workspace_root) if os.path.exists(path) else WarmupState() state.deepep_transport = transport.value state.deepep_num_nvl_bytes = num_nvl_bytes + os.makedirs(workspace_root, exist_ok=True) save_warmup_state(workspace_root, state) From 694f22fb4415822f249882ee78b461473f10d33b Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Thu, 23 Jul 2026 10:15:47 +0000 Subject: [PATCH 16/59] test: cover no-IMEX experimental recipes Signed-off-by: Cursor Agent Co-authored-by: Rahul Chalamala --- tests/test_experimental_recipe.py | 58 +++++++++++++++++++++++++++++++ 1 file changed, 58 insertions(+) create mode 100644 tests/test_experimental_recipe.py diff --git a/tests/test_experimental_recipe.py b/tests/test_experimental_recipe.py new file mode 100644 index 00000000..6be7e5e6 --- /dev/null +++ b/tests/test_experimental_recipe.py @@ -0,0 +1,58 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the Foundry project +"""Static contracts for the experimental no-IMEX DeepEP recipes.""" + +import tomllib +from pathlib import Path + +ROOT = Path(__file__).resolve().parents[1] +EXPERIMENTAL_DIR = ROOT / "recipe" / "experimental" +CONFIG_PAIRS = { + "vLLM": ("foundry_save.toml", "foundry_load.toml"), + "SGLang": ("sglang_foundry_save.toml", "sglang_foundry_load.toml"), +} +SERVE_SCRIPTS = ( + "serve_qwen3-30ba3b_ipc_ep.sh", + "serve_qwen3-30ba3bfp8_sglang_ipc_ep.sh", +) + + +def _load_config(filename: str) -> dict[str, object]: + path = EXPERIMENTAL_DIR / filename + assert path.is_file(), f"missing no-IMEX recipe config: {path}" + return tomllib.loads(path.read_text()) + + +def test_no_imex_configs_select_nvl_ipc_transport() -> None: + for config_pair in CONFIG_PAIRS.values(): + for filename in config_pair: + config = _load_config(filename) + assert config.get("deepep_transport") == "nvl_ipc", filename + + +def test_save_and_load_configs_share_archive_layout() -> None: + for engine, (save_filename, load_filename) in CONFIG_PAIRS.items(): + save_config = _load_config(save_filename) + load_config = _load_config(load_filename) + + for key in ("workspace_root", "base_addr", "region_size"): + assert save_config.get(key) == load_config.get(key), f"{engine}: {key}" + + +def test_experimental_scripts_do_not_use_environment_transport_switch() -> None: + for filename in SERVE_SCRIPTS: + path = EXPERIMENTAL_DIR / filename + assert path.is_file(), f"missing no-IMEX serve script: {path}" + assert "FOUNDRY_DEEPEP_NVL_IPC" not in path.read_text(), filename + + +def test_experimental_configs_have_portable_paths() -> None: + for config_pair in CONFIG_PAIRS.values(): + for filename in config_pair: + path = EXPERIMENTAL_DIR / filename + text = path.read_text() + config = tomllib.loads(text) + + assert "nvshmem_host_path" not in config, filename + assert "/data/" not in text, filename + assert not Path(str(config["workspace_root"])).is_absolute(), filename From 2ec7250390459528511029175284639813e464de Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Thu, 23 Jul 2026 10:16:52 +0000 Subject: [PATCH 17/59] test: cover no-IMEX recipe documentation Signed-off-by: Cursor Agent Co-authored-by: Rahul Chalamala --- tests/test_experimental_recipe.py | 31 +++++++++++++++++++++++++++++++ 1 file changed, 31 insertions(+) diff --git a/tests/test_experimental_recipe.py b/tests/test_experimental_recipe.py index 6be7e5e6..ec9dc998 100644 --- a/tests/test_experimental_recipe.py +++ b/tests/test_experimental_recipe.py @@ -15,6 +15,15 @@ "serve_qwen3-30ba3b_ipc_ep.sh", "serve_qwen3-30ba3bfp8_sglang_ipc_ep.sh", ) +INTEGRATION_DOCS = ( + ROOT / "docs" / "vllm" / "moe-and-deepep.md", + ROOT / "docs" / "sglang" / "hooks.md", +) +TWO_GPU_COMMANDS = ( + "bash tests/run_deepep_matrix.sh ll_nofabric both", + "bash tests/run_deepep_matrix.sh nvl_ipc both", + "bash tests/run_deepep_matrix.sh nvl_ipc_prealloc both", +) def _load_config(filename: str) -> dict[str, object]: @@ -56,3 +65,25 @@ def test_experimental_configs_have_portable_paths() -> None: assert "nvshmem_host_path" not in config, filename assert "/data/" not in text, filename assert not Path(str(config["workspace_root"])).is_absolute(), filename + + +def test_integration_docs_define_the_deepep_transport_contract() -> None: + for path in INTEGRATION_DOCS: + text = path.read_text() + + for term in ("deepep_transport", "fabric", "nvl_ipc", "intranode", "DeepEP", "NVSHMEM"): + assert term in text, f"{path}: missing {term}" + assert "default" in text.lower(), f"{path}: missing default transport" + assert "silent fallback" in text.lower(), f"{path}: missing no-fallback guarantee" + assert ( + "Foundry DeepEP transport mismatch: archived=nvl_ipc, configured=fabric" in text + ), f"{path}: missing archive parity error" + for command in TWO_GPU_COMMANDS: + assert command in text, f"{path}: missing {command}" + + +def test_sglang_docs_define_required_deepep_api() -> None: + text = (ROOT / "docs" / "sglang" / "hooks.md").read_text() + + assert "Buffer(use_fabric=...)" in text + assert "29d31c0" in text From f743855c37dc392589858b6df4696a1faebb96d0 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Thu, 23 Jul 2026 10:21:45 +0000 Subject: [PATCH 18/59] docs: add no-IMEX DeepEP recipes Signed-off-by: Cursor Agent Co-authored-by: Rahul Chalamala --- docs/sglang/hooks.md | 52 +++++++- docs/vllm/moe-and-deepep.md | 71 ++++++++--- recipe/experimental/README.md | 111 ++++++++++++++++++ recipe/experimental/foundry_load.toml | 9 ++ recipe/experimental/foundry_save.toml | 9 ++ .../experimental/serve_qwen3-30ba3b_ipc_ep.sh | 66 +++++++++++ .../serve_qwen3-30ba3bfp8_sglang_ipc_ep.sh | 58 +++++++++ recipe/experimental/sglang_foundry_load.toml | 9 ++ recipe/experimental/sglang_foundry_save.toml | 9 ++ tests/test_experimental_recipe.py | 9 +- 10 files changed, 381 insertions(+), 22 deletions(-) create mode 100644 recipe/experimental/README.md create mode 100644 recipe/experimental/foundry_load.toml create mode 100644 recipe/experimental/foundry_save.toml create mode 100755 recipe/experimental/serve_qwen3-30ba3b_ipc_ep.sh create mode 100755 recipe/experimental/serve_qwen3-30ba3bfp8_sglang_ipc_ep.sh create mode 100644 recipe/experimental/sglang_foundry_load.toml create mode 100644 recipe/experimental/sglang_foundry_save.toml diff --git a/docs/sglang/hooks.md b/docs/sglang/hooks.md index cf5b4c82..28a41ec2 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. @@ -177,6 +178,32 @@ def patched_start(self, *args, **kwargs): Active only when `moe_a2a_backend == deepep`. EP runs DP-attention + DeepEP (NCCL-free); TP attention is unsupported (its NCCL all-reduce is incompatible with the VMM region). +- **DeepEP transport policy** (`_patch_deepep`, hooks). `deepep_transport` + accepts only `fabric` and `nvl_ipc`; omitting it selects the `fabric` default. + The integration wraps `deep_ep.Buffer.__init__` and applies the TOML policy: + `fabric` sets `use_fabric=True` and forces `num_nvl_bytes=0`, while + `nvl_ipc` sets `use_fabric=False` and preserves SGLang's computed nonzero + `num_nvl_bytes`. The latter is intranode-only and requires all ranks to share + the needed NVLink connectivity. +- **Required DeepEP API.** Foundry requires a DeepEP release whose + `Buffer(use_fabric=...)` API accepts the `use_fabric` keyword. DeepEP + `29d31c0` or newer provides that API while retaining SGLang's positional + constructor shape. Foundry raises an actionable error before calling an + older constructor; there is no silent fallback to another transport. +- **DeepEP and NVSHMEM remain required.** `nvl_ipc` changes only the NVLink + buffer. DeepEP's RDMA buffer still uses the NVSHMEM symmetric heap. The + no-IMEX recipe sets `NVSHMEM_CUMEM_HANDLE_TYPE=FILE_DESCRIPTOR`, + `NCCL_CUMEM_ENABLE=0`, and `NCCL_NVLS_ENABLE=0` for matching SAVE and LOAD + runs. +- **Archive parity.** SAVE records the effective transport and NVLink buffer + size. LOAD validates the transport before loading CUDA modules or restoring + graphs, then validates the buffer size when constructing DeepEP. A mismatched + LOAD config fails fast with: + + ```text + Foundry DeepEP transport mismatch: archived=nvl_ipc, configured=fabric + ``` + - **DeepEP buffer pre-capture bootstrap** (`bootstrap_deepep_buffer`, graph_ops). sglang creates the singleton NVSHMEM `Buffer` lazily on the first MoE dispatch — normally during the warmup forwards foundry suppresses, which would push creation *inside* the @@ -203,7 +230,26 @@ TP attention is unsupported (its NCCL all-reduce is incompatible with the VMM re surfaced only on sglang EP. See [`../../recipe/sglang/README.md`](../../recipe/sglang/README.md) for the EP serve -config and required kernel versions (deep_ep `9af0e0d`, sgl-deep-gemm ≥0.1.2, fa3). +config and required kernel versions (DeepEP `29d31c0` or newer, +sgl-deep-gemm ≥0.1.2, fa3). + +### No-IMEX validation + +On a two-GPU host without IMEX, run the standalone RDMA-only, NVLink IPC, and +LOAD-preallocation cases: + +```bash +bash tests/run_deepep_matrix.sh ll_nofabric both +bash tests/run_deepep_matrix.sh nvl_ipc both +bash tests/run_deepep_matrix.sh nvl_ipc_prealloc both +``` + +Then run the focused SGLang EP recipe: + +```bash +bash recipe/experimental/serve_qwen3-30ba3bfp8_sglang_ipc_ep.sh 2 --save +bash recipe/experimental/serve_qwen3-30ba3bfp8_sglang_ipc_ep.sh 2 --load +``` ## Patch idiom @@ -211,4 +257,6 @@ All patches use the `wrap-and-call` idiom — short-circuit on `mode == NONE`, r ## Install order -`install_hooks` calls the patch helpers in the listed order. Order doesn't matter for correctness here — every patch attaches to a different attribute — but spawn sites go last so the install-completion log line appears after every other patch has registered. +`install_hooks` calls the patch helpers in the listed order. Order does not +matter for correctness here because every patch attaches to a different +attribute. diff --git a/docs/vllm/moe-and-deepep.md b/docs/vllm/moe-and-deepep.md index e76b48d8..63fdbf76 100644 --- a/docs/vllm/moe-and-deepep.md +++ b/docs/vllm/moe-and-deepep.md @@ -1,6 +1,9 @@ # MoE / DeepEP Support -vLLM-side MoE / expert-parallelism support is what makes the integration interesting beyond dense decode. The remaining MoE-specific concern is that DeepEP all-to-all communicator buffers must be set up in fabric mode for foundry to interpose them, and NVSHMEM module init must fire after the runtime is bootstrapped. +vLLM-side MoE / expert-parallelism support is what makes the integration +interesting beyond dense decode. Foundry configures the DeepEP all-to-all +communicator for either fabric or intranode NVLink IPC, and NVSHMEM module init +must fire after the runtime is bootstrapped. The dense-only baseline doesn't exercise any of this; turn it on by running an MoE model with EP. @@ -8,32 +11,51 @@ The dense-only baseline doesn't exercise any of this; turn it on by running an M vLLM's recent versions use modular kernels for both quantized and unquantized MoE — the same `model_loader.load_model()` path produces deterministic packing regardless of foundry mode. Foundry's integration takes the unified preallocate → `start_graph_builds` → weight-load → `prepare_communication_buffer_for_model` path for all MoE variants (and for dense models), so no quant-metadata roundtrip through `WarmupState` is needed. The earlier `collect_moe_quant_metadata` / `inject_moe_quant_metadata` / `reset_moe_quant_config` / `reinit_moe_quant_config` / `split_load_model` ceremony has been removed. -## DeepEP fabric mode +## DeepEP transport ### Problem -`DeepEPLLAll2AllManager` builds the all-to-all communicator with one of several backends: +`DeepEPLLAll2AllManager` builds the all-to-all communicator with one of two +Foundry-supported profiles: -- NVL (NVLink direct) — uses `cudaIpcGetMemHandle` for buffer sharing. -- Fabric — uses NVSHMEM + symmetric memory. +- `fabric` uses fabric handles and an NVSHMEM symmetric heap. +- `nvl_ipc` uses `cudaIpcGetMemHandle` for the nonzero intranode NVLink buffer + and Foundry's VMM-IPC translation to share it between ranks. -NVL's IPC handles are process-local and not persistable; the captured graph can't reference them on LOAD. Fabric mode allocates buffers via NVSHMEM, which foundry's hook can interpose (`init_nvshmem_for_loaded_modules` re-binds them on LOAD). +`deepep_transport` accepts only those two values and defaults to `fabric` when +omitted. Select the no-IMEX path in both SAVE and LOAD TOMLs: + +```toml +deepep_transport = "nvl_ipc" +``` + +`nvl_ipc` is intranode-only; all ranks must be on one host with the required +NVLink connectivity. It does not remove the DeepEP or NVSHMEM dependency. +DeepEP's RDMA buffer remains on the NVSHMEM symmetric heap, while only its +NVLink buffer uses IPC. The no-IMEX serve script also requires +`NCCL_CUMEM_ENABLE=0` and `NCCL_NVLS_ENABLE=0`. ### Solution -`_patch_deepep` wraps `DeepEPLLAll2AllManager._make_all2all_kwargs`: +`_patch_deepep` wraps +`DeepEPLLAll2AllManager._make_all2all_kwargs`. In Foundry mode it applies the +TOML policy after vLLM computes the upstream arguments: -```python -@functools.wraps(orig) -def patched(self, *args, **kwargs): - kwargs = orig(self, *args, **kwargs) - if get_graph_extension_mode() != CUDAGraphExtensionMode.NONE: - kwargs["use_fabric"] = True - kwargs["num_nvl_bytes"] = 0 # RDMA-only path - return kwargs +- `fabric`: set `use_fabric=True` and force `num_nvl_bytes=0`. +- `nvl_ipc`: set `use_fabric=False` and preserve the computed nonzero + `num_nvl_bytes`. + +The effective transport and NVLink buffer size are recorded in every SAVE +archive. LOAD validates the transport before loading archived CUDA modules and +validates the buffer size when DeepEP creates the communicator. For example, +loading an `nvl_ipc` archive with a `fabric` TOML fails fast with: + +```text +Foundry DeepEP transport mismatch: archived=nvl_ipc, configured=fabric ``` -Forces every foundry-mode DeepEP setup onto the fabric path. The patch lives behind a try/except in case the import path moves. +There is no silent fallback between transports. Falling back would change the +allocation trajectory and invalidate the captured graphs. ## NVSHMEM initialization order @@ -56,3 +78,20 @@ For a working MoE SAVE→LOAD cycle: - All ranks produce `rank_{N}/warmup_state.json`-equivalent per-rank artifacts. - All ranks' `final_alloc_offset` are reproducible across SAVE pass 1 → pass 2. - On LOAD, the first inference call returns coherent tokens. + +On a two-GPU host without IMEX, validate the RDMA-only baseline, direct NVLink +IPC, and LOAD preallocation: + +```bash +bash tests/run_deepep_matrix.sh ll_nofabric both +bash tests/run_deepep_matrix.sh nvl_ipc both +bash tests/run_deepep_matrix.sh nvl_ipc_prealloc both +``` + +Then run the focused vLLM recipe twice in SAVE mode and once in LOAD mode: + +```bash +bash recipe/experimental/serve_qwen3-30ba3b_ipc_ep.sh 2 --save +bash recipe/experimental/serve_qwen3-30ba3b_ipc_ep.sh 2 --save +bash recipe/experimental/serve_qwen3-30ba3b_ipc_ep.sh 2 --load +``` diff --git a/recipe/experimental/README.md b/recipe/experimental/README.md new file mode 100644 index 00000000..ef78e433 --- /dev/null +++ b/recipe/experimental/README.md @@ -0,0 +1,111 @@ +# No-IMEX intranode DeepEP recipes + +These experimental recipes run vLLM or SGLang expert parallelism on one host +without NVIDIA IMEX. They select Foundry's NVLink IPC transport in matching +SAVE and LOAD TOMLs: + +```toml +deepep_transport = "nvl_ipc" +``` + +`deepep_transport` accepts only `fabric` and `nvl_ipc`; omitting it selects the +`fabric` default. `nvl_ipc` is intranode-only and requires all participating +GPUs to share the needed NVLink connectivity. It does not remove DeepEP or +NVSHMEM: DeepEP's RDMA buffer still uses the NVSHMEM symmetric heap while its +nonzero NVLink buffer uses Foundry's VMM-IPC translation. + +## Files + +```text +recipe/experimental/ +├── README.md +├── foundry_save.toml +├── foundry_load.toml +├── serve_qwen3-30ba3b_ipc_ep.sh +├── sglang_foundry_save.toml +├── sglang_foundry_load.toml +└── serve_qwen3-30ba3bfp8_sglang_ipc_ep.sh +``` + +This directory intentionally contains only expert-parallel recipes. The vLLM +files use `foundry_archive_ipc`; the SGLang files use +`foundry_archive_sglang_ipc`. Both paths are relative to the directory from +which the server is started. + +## Requirements + +- Build and install this Foundry revision so `libcuda_hook.so` includes the + SCM_RIGHTS VMM-IPC transport: + + ```bash + CC=gcc CXX=g++ pip install -e . --no-build-isolation + ``` + +- Install DeepEP and NVSHMEM. Set `nvshmem_host_path` in both TOMLs for an + engine only when Foundry cannot auto-detect a custom NVSHMEM host library. + Use a machine-local absolute path; no checkout-specific path is assumed. +- The SGLang recipe requires DeepEP `29d31c0` or newer because + `deep_ep.Buffer` must accept `use_fabric`, plus sgl-deep-gemm 0.1.2 or newer + and FlashAttention 3. +- The scripts keep `NCCL_CUMEM_ENABLE=0` and `NCCL_NVLS_ENABLE=0` for every + Foundry phase. The SGLang script also sets + `NVSHMEM_CUMEM_HANDLE_TYPE=FILE_DESCRIPTOR`. + +## vLLM SAVE and LOAD + +Run every phase from the same repository path so vLLM sees an identical TOML +path and reuses its compile cache: + +```bash +rm -rf foundry_archive_ipc +bash recipe/experimental/serve_qwen3-30ba3b_ipc_ep.sh 2 --save +bash recipe/experimental/serve_qwen3-30ba3b_ipc_ep.sh 2 --save +bash recipe/experimental/serve_qwen3-30ba3b_ipc_ep.sh 2 --load +``` + +Stop each SAVE server after startup. The first pass profiles and captures; the +second pass checks deterministic re-capture. Leave the LOAD server running for +inference. + +## SGLang SAVE and LOAD + +SGLang needs one SAVE pass: + +```bash +rm -rf foundry_archive_sglang_ipc +bash recipe/experimental/serve_qwen3-30ba3bfp8_sglang_ipc_ep.sh 2 --save +bash recipe/experimental/serve_qwen3-30ba3bfp8_sglang_ipc_ep.sh 2 --load +``` + +Keep the model, EP size, low-latency token cap, chunked-prefill size, and TOML +settings identical between SAVE and LOAD. + +## Archive parity and failure behavior + +The SAVE and LOAD config for an engine must match on `deepep_transport`, +`workspace_root`, `base_addr`, and `region_size`. Foundry stores the effective +transport and NVLink buffer size in the archive. A transport mismatch fails +before CUDA modules, DeepEP buffers, or graphs are restored: + +```text +Foundry DeepEP transport mismatch: archived=nvl_ipc, configured=fabric +``` + +An effective NVLink buffer-size mismatch also fails when DeepEP constructs the +buffer. There is no silent fallback from `nvl_ipc` to `fabric`; changing the +transport changes allocation geometry and would invalidate captured graphs. + +## Two-GPU validation + +On a two-GPU host without IMEX, run all three standalone cases from the +repository root: + +```bash +bash tests/run_deepep_matrix.sh ll_nofabric both +bash tests/run_deepep_matrix.sh nvl_ipc both +bash tests/run_deepep_matrix.sh nvl_ipc_prealloc both +``` + +The cases cover the RDMA-only baseline, a nonzero NVLink IPC buffer, and LOAD +from a preallocated Foundry VMM chunk. They require deterministic dispatch, +successful SAVE and LOAD, and no CUDA error 999. diff --git a/recipe/experimental/foundry_load.toml b/recipe/experimental/foundry_load.toml new file mode 100644 index 00000000..9ed5ec0e --- /dev/null +++ b/recipe/experimental/foundry_load.toml @@ -0,0 +1,9 @@ +mode = "load" +base_addr = 0x600000000000 +region_size = "256GB" +workspace_root = "foundry_archive_ipc" +scratch_space_size = "4096MB" +deepep_transport = "nvl_ipc" +# DeepEP low-latency still allocates its RDMA buffer on the NVSHMEM symmetric +# heap. Foundry auto-detects libnvshmem_host.so from the installed NVSHMEM +# package; set nvshmem_host_path only to use a custom build. diff --git a/recipe/experimental/foundry_save.toml b/recipe/experimental/foundry_save.toml new file mode 100644 index 00000000..3232db54 --- /dev/null +++ b/recipe/experimental/foundry_save.toml @@ -0,0 +1,9 @@ +mode = "save" +base_addr = 0x600000000000 +region_size = "256GB" +workspace_root = "foundry_archive_ipc" +scratch_space_size = "4096MB" +deepep_transport = "nvl_ipc" +# DeepEP low-latency still allocates its RDMA buffer on the NVSHMEM symmetric +# heap. Foundry auto-detects libnvshmem_host.so from the installed NVSHMEM +# package; set nvshmem_host_path only to use a custom build. diff --git a/recipe/experimental/serve_qwen3-30ba3b_ipc_ep.sh b/recipe/experimental/serve_qwen3-30ba3b_ipc_ep.sh new file mode 100755 index 00000000..9d0011c5 --- /dev/null +++ b/recipe/experimental/serve_qwen3-30ba3b_ipc_ep.sh @@ -0,0 +1,66 @@ +#!/bin/bash +# Usage: bash serve_qwen3-30ba3b_ipc_ep.sh [--save|--load] +# +# EXPERIMENTAL: vLLM DeepEP expert parallel using the intranode NVLink IPC +# transport selected by the matching Foundry TOML. This requires the Foundry +# VMM-IPC hook and a host whose GPUs share an NVLink domain. +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" + +EP_SIZE=${1:?Usage: $0 [--save|--load]} +MODEL_NAME="Qwen/Qwen3-30B-A3B" +HOST="0.0.0.0" +PORT=12000 +GPU_MEMORY_UTILIZATION=0.8 + +FOUNDRY_ARGS=() +if [[ "$2" == "--save" ]]; then + FOUNDRY_TOML="${SCRIPT_DIR}/foundry_save.toml" +elif [[ "$2" == "--load" ]]; then + FOUNDRY_TOML="${SCRIPT_DIR}/foundry_load.toml" +elif [[ -n "$2" ]]; then + echo "Usage: $0 [--save|--load]" + exit 1 +fi + +if [[ -n "${FOUNDRY_TOML:-}" ]]; then + # The CUMEM P2P and NVLS multicast paths request mapping capabilities that + # Foundry's VMM allocation region does not carry. + export NCCL_CUMEM_ENABLE=0 + export NCCL_NVLS_ENABLE=0 + export VLLM_USE_V2_MODEL_RUNNER=0 + FOUNDRY_ARGS+=( + --compilation-config.graph_extension_config_path "$FOUNDRY_TOML" + ) + echo "Using Foundry DeepEP NVL/IPC: ${FOUNDRY_TOML}" +else + echo "Running without Foundry (baseline vLLM)" +fi + +# Foundry currently patches the V1 model runner. +export VLLM_USE_V2_MODEL_RUNNER=0 +export VLLM_USE_FLASHINFER_SAMPLER=1 +export VLLM_DISABLE_SHARED_EXPERTS_STREAM=1 + +CUDAGRAPH_CAPTURE_SIZES=($(seq 1 256)) + +ARGS=( + --trust-remote-code + --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..4c55c2e2 --- /dev/null +++ b/recipe/experimental/serve_qwen3-30ba3bfp8_sglang_ipc_ep.sh @@ -0,0 +1,58 @@ +#!/bin/bash +# Usage: bash serve_qwen3-30ba3bfp8_sglang_ipc_ep.sh [--save|--load] +# +# EXPERIMENTAL: SGLang DeepEP expert parallel using the intranode NVLink IPC +# transport selected by the matching Foundry TOML. +# +# Requires DeepEP 29d31c0 or newer, sgl-deep-gemm >= 0.1.2, flash-attn-3, +# NVSHMEM, and the Foundry VMM-IPC hook. +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" + +EP_SIZE=${1:?Usage: $0 [--save|--load]} +MODEL_NAME="${SGL_MODEL:-Qwen/Qwen3-30B-A3B-FP8}" +HOST="0.0.0.0" +PORT=12000 +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 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/sglang_foundry_load.toml b/recipe/experimental/sglang_foundry_load.toml new file mode 100644 index 00000000..5c8c82fa --- /dev/null +++ b/recipe/experimental/sglang_foundry_load.toml @@ -0,0 +1,9 @@ +mode = "load" +base_addr = 0x600000000000 +region_size = "256GB" +workspace_root = "foundry_archive_sglang_ipc" +scratch_space_size = "4096MB" +deepep_transport = "nvl_ipc" +# DeepEP low-latency still allocates its RDMA buffer on the NVSHMEM symmetric +# heap. Foundry auto-detects libnvshmem_host.so from the installed NVSHMEM +# package; set nvshmem_host_path only to use a custom build. diff --git a/recipe/experimental/sglang_foundry_save.toml b/recipe/experimental/sglang_foundry_save.toml new file mode 100644 index 00000000..78669275 --- /dev/null +++ b/recipe/experimental/sglang_foundry_save.toml @@ -0,0 +1,9 @@ +mode = "save" +base_addr = 0x600000000000 +region_size = "256GB" +workspace_root = "foundry_archive_sglang_ipc" +scratch_space_size = "4096MB" +deepep_transport = "nvl_ipc" +# DeepEP low-latency still allocates its RDMA buffer on the NVSHMEM symmetric +# heap. Foundry auto-detects libnvshmem_host.so from the installed NVSHMEM +# package; set nvshmem_host_path only to use a custom build. diff --git a/tests/test_experimental_recipe.py b/tests/test_experimental_recipe.py index ec9dc998..893bd117 100644 --- a/tests/test_experimental_recipe.py +++ b/tests/test_experimental_recipe.py @@ -2,9 +2,10 @@ # SPDX-FileCopyrightText: Copyright contributors to the Foundry project """Static contracts for the experimental no-IMEX DeepEP recipes.""" -import tomllib from pathlib import Path +import tomllib + ROOT = Path(__file__).resolve().parents[1] EXPERIMENTAL_DIR = ROOT / "recipe" / "experimental" CONFIG_PAIRS = { @@ -75,9 +76,9 @@ def test_integration_docs_define_the_deepep_transport_contract() -> None: assert term in text, f"{path}: missing {term}" assert "default" in text.lower(), f"{path}: missing default transport" assert "silent fallback" in text.lower(), f"{path}: missing no-fallback guarantee" - assert ( - "Foundry DeepEP transport mismatch: archived=nvl_ipc, configured=fabric" in text - ), f"{path}: missing archive parity error" + assert "Foundry DeepEP transport mismatch: archived=nvl_ipc, configured=fabric" in text, ( + f"{path}: missing archive parity error" + ) for command in TWO_GPU_COMMANDS: assert command in text, f"{path}: missing {command}" From 0a5d36b80da8a2de1a063bade61fecc7006c20d9 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Thu, 23 Jul 2026 11:03:05 +0000 Subject: [PATCH 19/59] test: cover modern SGLang attention metadata API Signed-off-by: Cursor Agent Co-authored-by: Rahul Chalamala --- tests/test_deepep_transport.py | 65 ++++++++++++++++++++++++++++++++++ 1 file changed, 65 insertions(+) diff --git a/tests/test_deepep_transport.py b/tests/test_deepep_transport.py index 696adcb6..216943cc 100644 --- a/tests/test_deepep_transport.py +++ b/tests/test_deepep_transport.py @@ -27,6 +27,25 @@ def _install_fake_deep_ep(monkeypatch, buffer_cls): monkeypatch.setattr(sglang_hooks, "_DEEPEP_PATCHED", False, raising=False) +def _install_fake_sglang_cuda_graph_runner(monkeypatch, runner_cls): + sglang_module = types.ModuleType("sglang") + srt_module = types.ModuleType("sglang.srt") + model_executor_module = types.ModuleType("sglang.srt.model_executor") + cuda_graph_runner_module = types.ModuleType("sglang.srt.model_executor.cuda_graph_runner") + cuda_graph_runner_module.CudaGraphRunner = runner_cls + model_executor_module.cuda_graph_runner = cuda_graph_runner_module + srt_module.model_executor = model_executor_module + sglang_module.srt = srt_module + monkeypatch.setitem(sys.modules, "sglang", sglang_module) + monkeypatch.setitem(sys.modules, "sglang.srt", srt_module) + monkeypatch.setitem(sys.modules, "sglang.srt.model_executor", model_executor_module) + monkeypatch.setitem( + sys.modules, + "sglang.srt.model_executor.cuda_graph_runner", + cuda_graph_runner_module, + ) + + @pytest.mark.parametrize( ("raw", "expected"), [ @@ -114,6 +133,52 @@ def test_sglang_get_deepep_transport_uses_loaded_config(tmp_path, monkeypatch): assert sglang_config.get_deepep_transport() is DeepEPTransport.NVL_IPC +def test_sglang_save_capture_supports_out_graph_attention_api(monkeypatch): + def original_out_graph(forward_batch, in_capture=False): + return None + + class CudaGraphRunner: + def __init__(self): + self.attn_backend = SimpleNamespace( + init_forward_metadata_out_graph=original_out_graph, + ) + + def capture(self): + return "captured" + + def _create_device_graph(self): + return None + + def _capture_graph(self, graph, pool, stream, run_once_fn): + return run_once_fn() + + def capture_one_batch_size(self, bs, forward, stream_idx=None): + return None, forward() + + _install_fake_sglang_cuda_graph_runner(monkeypatch, CudaGraphRunner) + graph_ops = types.ModuleType("foundry.integration.sglang.graph_ops") + graph_ops.pack_fatbins = lambda: None + graph_ops.save_graph_manifest = lambda: None + monkeypatch.setitem( + sys.modules, + "foundry.integration.sglang.graph_ops", + graph_ops, + ) + monkeypatch.setattr(sglang_hooks, "_ep_lazy_init_needed", lambda: False) + monkeypatch.setattr( + sglang_hooks, + "get_graph_extension_mode", + lambda: sglang_hooks.CUDAGraphExtensionMode.SAVE, + ) + monkeypatch.setattr(sglang_hooks.rt, "capture_final_alloc_offset", lambda: 0) + + sglang_hooks._patch_cuda_graph_capture() + runner = CudaGraphRunner() + + assert runner.capture() == "captured" + assert runner.attn_backend.init_forward_metadata_out_graph is original_out_graph + + def test_sglang_warmup_state_round_trips_deepep_profile(tmp_path, monkeypatch): config = sglang_config.CUDAGraphExtensionConfig(workspace_root=str(tmp_path)) monkeypatch.setattr(sglang_config, "_config", config) From 4c30d5eb2f76a8a3cd4f84190722012d30cc8f18 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Thu, 23 Jul 2026 11:15:39 +0000 Subject: [PATCH 20/59] fix: support modern SGLang attention metadata API Signed-off-by: Cursor Agent Co-authored-by: Rahul Chalamala --- .../foundry/integration/sglang/graph_ops.py | 34 +++++++++++- python/foundry/integration/sglang/hooks.py | 53 +++++++++++++++++-- 2 files changed, 83 insertions(+), 4 deletions(-) diff --git a/python/foundry/integration/sglang/graph_ops.py b/python/foundry/integration/sglang/graph_ops.py index 18ab1f0e..927a95d0 100644 --- a/python/foundry/integration/sglang/graph_ops.py +++ b/python/foundry/integration/sglang/graph_ops.py @@ -8,6 +8,7 @@ import os import re import time +from types import SimpleNamespace from typing import Any import torch @@ -248,6 +249,31 @@ def bootstrap_deepep_buffer(cuda_graph_runner) -> bool: return False +def build_capture_fb_view(cuda_graph_runner, bs: int): + """Build a ForwardBatch-like view for backend capture-side metadata init. + + Mirrors the padded per-bs inputs SGLang's own capture loop feeds the + attention backend as a lightweight namespace. This drives the modern + ``init_forward_metadata_out_graph`` API before a real ``ForwardBatch`` + exists. + """ + buffers = cuda_graph_runner.buffers + num_tokens = bs * cuda_graph_runner.num_tokens_per_bs + encoder_lens = buffers.encoder_lens[:bs] if cuda_graph_runner.is_encoder_decoder else None + seq_lens = buffers.seq_lens[:bs] + return SimpleNamespace( + batch_size=bs, + forward_mode=cuda_graph_runner.capture_forward_mode, + req_pool_indices=buffers.req_pool_indices[:bs], + seq_lens=seq_lens, + seq_lens_cpu=buffers.seq_lens_cpu[:bs], + seq_lens_sum=int(seq_lens.sum().item()), + encoder_lens=encoder_lens, + spec_info=cuda_graph_runner.get_spec_info(num_tokens), + positions=buffers.positions[:num_tokens], + ) + + def initialize_attention_metadata_for_bs(cuda_graph_runner, bs: int) -> None: """Populate ``decode_cuda_graph_metadata[bs]`` for runtime replay. @@ -257,11 +283,17 @@ def initialize_attention_metadata_for_bs(cuda_graph_runner, bs: int) -> None: re-run the same call before runtime replay so the wrappers exist at deterministic VMM addresses. """ + attn_backend = cuda_graph_runner.attn_backend + if hasattr(attn_backend, "init_forward_metadata_out_graph"): + fb_view = build_capture_fb_view(cuda_graph_runner, bs) + attn_backend.init_forward_metadata_out_graph(fb_view, in_capture=True) + return + buffers = cuda_graph_runner.buffers num_tokens = bs * cuda_graph_runner.num_tokens_per_bs encoder_lens = buffers.encoder_lens[:bs] if cuda_graph_runner.is_encoder_decoder else None spec_info = cuda_graph_runner.get_spec_info(num_tokens) - cuda_graph_runner.attn_backend.init_forward_metadata_capture_cuda_graph( + attn_backend.init_forward_metadata_capture_cuda_graph( bs, num_tokens, buffers.req_pool_indices[:bs], diff --git a/python/foundry/integration/sglang/hooks.py b/python/foundry/integration/sglang/hooks.py index 9078acf3..f63cc1ff 100644 --- a/python/foundry/integration/sglang/hooks.py +++ b/python/foundry/integration/sglang/hooks.py @@ -493,7 +493,47 @@ def patched(self, *args, **kwargs): # FlashInfer by its per-bs indices_updater_decode. attn_backend = self.attn_backend use_fi_prepass = hasattr(attn_backend, "indices_updater_decode") - real_init = attn_backend.init_forward_metadata_capture_cuda_graph + # SGLang #26735 replaced the legacy capture-init entry point with + # ``init_forward_metadata_out_graph(forward_batch, in_capture)``. + # Shim whichever API the installed SGLang fork exposes. + use_new_api = hasattr(attn_backend, "init_forward_metadata_out_graph") + real_out_graph = attn_backend.init_forward_metadata_out_graph if use_new_api else None + real_init = ( + None if use_new_api else attn_backend.init_forward_metadata_capture_cuda_graph + ) + + def reuse_out_graph(forward_batch, in_capture=False): + if not in_capture: + return real_out_graph(forward_batch, in_capture=in_capture) + + from sglang.srt.layers.attention.flashinfer_backend import ( + DecodeMetadata, + PrefillMetadata, + ) + + bs = forward_batch.batch_size + forward_mode = forward_batch.forward_mode + if forward_mode.is_decode_or_idle(): + wrappers = attn_backend.decode_cuda_graph_metadata.get(bs) + if wrappers is None: + return real_out_graph(forward_batch, in_capture=True) + attn_backend.forward_metadata = DecodeMetadata(wrappers) + elif ( + forward_mode.is_target_verify() + or forward_mode.is_draft_extend() + or forward_mode.is_dllm_extend() + ): + wrappers = attn_backend.prefill_cuda_graph_metadata.get(bs) + if wrappers is None: + return real_out_graph(forward_batch, in_capture=True) + attn_backend.forward_metadata = PrefillMetadata( + wrappers, forward_mode.is_dllm_extend(), False + ) + else: + return real_out_graph(forward_batch, in_capture=True) + # Re-plan against the preallocated wrappers without allocating + # another capture metadata workspace. + real_out_graph(forward_batch, in_capture=False) def reuse_pre_pass_init( bs, @@ -607,11 +647,18 @@ def reuse_pre_pass_init( # Drop the pre-pass's last forward_metadata ref so popping the dict # entry doesn't keep the wrapper alive at refcount 1. attn_backend.forward_metadata = None - attn_backend.init_forward_metadata_capture_cuda_graph = reuse_pre_pass_init + if use_new_api: + attn_backend.init_forward_metadata_out_graph = reuse_out_graph + else: + attn_backend.init_forward_metadata_capture_cuda_graph = reuse_pre_pass_init try: result = orig_capture(self, *args, **kwargs) finally: - attn_backend.init_forward_metadata_capture_cuda_graph = real_init + if use_fi_prepass: + if use_new_api: + attn_backend.init_forward_metadata_out_graph = real_out_graph + else: + attn_backend.init_forward_metadata_capture_cuda_graph = real_init from foundry.integration.sglang.graph_ops import ( pack_fatbins, From f48ef04401eba7b9e3555a01a38cc91389ac1948 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Thu, 23 Jul 2026 14:45:09 +0000 Subject: [PATCH 21/59] test: cover vLLM NVSHMEM host auto-detection Signed-off-by: Cursor Agent Co-authored-by: Rahul Chalamala --- tests/test_deepep_transport_vllm.py | 92 +++++++++++++++++++++++++++++ 1 file changed, 92 insertions(+) diff --git a/tests/test_deepep_transport_vllm.py b/tests/test_deepep_transport_vllm.py index 88084d10..967e315f 100644 --- a/tests/test_deepep_transport_vllm.py +++ b/tests/test_deepep_transport_vllm.py @@ -5,6 +5,7 @@ import importlib.util import json import logging +import os import sys import types from pathlib import Path @@ -23,6 +24,20 @@ ) +def _installed_nvshmem_host_path() -> Path: + spec = importlib.util.find_spec("nvidia.nvshmem") + assert spec is not None, "nvidia.nvshmem must be installed for this integration test" + roots = [Path(root) for root in spec.submodule_search_locations or ()] + if spec.origin: + roots.append(Path(spec.origin).parent) + for root in roots: + for name in ("libnvshmem_host.so.3", "libnvshmem_host.so"): + candidate = root / "lib" / name + if candidate.exists(): + return candidate.resolve() + pytest.fail("installed nvidia.nvshmem package has no libnvshmem_host shared library") + + def _init_stub_logger(name: str) -> logging.Logger: return logging.getLogger(name) @@ -186,6 +201,83 @@ def test_vllm_get_deepep_transport_uses_loaded_config( assert vllm_config_module.get_deepep_transport() is DeepEPTransport.NVL_IPC +def test_vllm_config_autodetects_installed_nvshmem_host_path( + tmp_path, + vllm_config_module, +): + config_path = tmp_path / "config.toml" + config_path.write_text("") + expected = _installed_nvshmem_host_path() + + config = vllm_config_module.CUDAGraphExtensionConfig.from_toml(config_path) + + assert config.nvshmem_host_path == str(expected) + assert Path(config.nvshmem_host_path).is_absolute() + + +def test_vllm_runtime_preloads_autodetected_nvshmem_host_path( + tmp_path, + monkeypatch, + vllm_runtime_module, +): + config_path = tmp_path / "config.toml" + config_path.write_text("") + expected = _installed_nvshmem_host_path() + config_module = sys.modules[vllm_runtime_module.get_nvshmem_host_path.__module__] + config = config_module.CUDAGraphExtensionConfig.from_toml(config_path) + monkeypatch.setattr(config_module, "_config", config) + monkeypatch.setattr(vllm_runtime_module, "get_hook_library_path", lambda: None) + monkeypatch.delenv("LD_PRELOAD", raising=False) + + vllm_runtime_module.setup_ld_preload_env() + + preload_paths = [Path(path) for path in os.environ["LD_PRELOAD"].split(":")] + assert expected in preload_paths + assert all(path.is_absolute() for path in preload_paths) + + +def test_vllm_config_explicit_nvshmem_host_path_overrides_detection( + tmp_path, + monkeypatch, + vllm_config_module, +): + explicit = tmp_path / "custom" / "libnvshmem_host.so" + config_path = tmp_path / "config.toml" + config_path.write_text(f'nvshmem_host_path = "{explicit}"\n') + monkeypatch.setattr( + vllm_config_module.CUDAGraphExtensionConfig, + "_detect_nvshmem_host_path", + staticmethod(lambda: pytest.fail("explicit path must bypass auto-detection")), + raising=False, + ) + + config = vllm_config_module.CUDAGraphExtensionConfig.from_toml(config_path) + + assert config.nvshmem_host_path == str(explicit) + + +def test_vllm_config_missing_nvshmem_package_is_none_for_non_ep( + tmp_path, + monkeypatch, + vllm_config_module, +): + config_path = tmp_path / "config.toml" + config_path.write_text("") + real_find_spec = vllm_config_module.importlib.util.find_spec + + def find_spec_without_nvshmem(name): + if name == "nvidia.nvshmem": + return None + return real_find_spec(name) + + monkeypatch.setattr(vllm_config_module.importlib.util, "find_spec", find_spec_without_nvshmem) + + config = vllm_config_module.CUDAGraphExtensionConfig.from_toml(config_path) + + assert config.nvshmem_host_path is None + assert config.deepep_transport is DeepEPTransport.FABRIC + + def test_vllm_warmup_state_round_trips_deepep_profile(tmp_path, vllm_runtime_module): state = vllm_runtime_module.WarmupState() state.deepep_transport = "nvl_ipc" From 384017d69c354acb54b02a05190ef47418f9d987 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Thu, 23 Jul 2026 14:49:56 +0000 Subject: [PATCH 22/59] fix: detect installed NVSHMEM host library for vLLM Signed-off-by: Cursor Agent Co-authored-by: Rahul Chalamala --- python/foundry/integration/vllm/config.py | 25 ++++++++++++++++++++++- 1 file changed, 24 insertions(+), 1 deletion(-) diff --git a/python/foundry/integration/vllm/config.py b/python/foundry/integration/vllm/config.py index 552260b6..9011412f 100644 --- a/python/foundry/integration/vllm/config.py +++ b/python/foundry/integration/vllm/config.py @@ -62,10 +62,14 @@ def from_toml(cls, path: str | Path) -> CUDAGraphExtensionConfig: if hook_library_path is None: hook_library_path = cls._detect_hook_so_path() + nvshmem_host_path = data.get("nvshmem_host_path") + if nvshmem_host_path is None: + nvshmem_host_path = cls._detect_nvshmem_host_path() + return cls( mode=mode, hook_library_path=hook_library_path, - nvshmem_host_path=data.get("nvshmem_host_path"), + nvshmem_host_path=nvshmem_host_path, deepep_transport=parse_deepep_transport(data.get("deepep_transport")), base_addr=base_addr, region_size=data.get("region_size", cls.region_size), @@ -83,6 +87,25 @@ def _detect_hook_so_path() -> str | None: return str(hook_so_path) return None + @staticmethod + def _detect_nvshmem_host_path() -> str | None: + """Locate DeepEP's NVSHMEM host lib from the installed wheel.""" + try: + spec = importlib.util.find_spec("nvidia.nvshmem") + except (ImportError, ValueError): + return None + if spec is None: + return None + roots = list(spec.submodule_search_locations or []) + if spec.origin: + roots.append(str(Path(spec.origin).parent)) + for root in roots: + for name in ("libnvshmem_host.so.3", "libnvshmem_host.so"): + candidate = Path(root) / "lib" / name + if candidate.exists(): + return str(candidate) + return None + # Module-global config singleton (process-local). _config: CUDAGraphExtensionConfig | None = None From 8a4a0bda245184cb354072d48686099d29538f0a Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Fri, 24 Jul 2026 00:00:28 +0000 Subject: [PATCH 23/59] test: cover CUDA graph dependency edge metadata Signed-off-by: Cursor Agent Co-authored-by: Rahul Chalamala --- tests/cuda_graph_edge_data_test.cu | 173 +++++++++++++++ tests/test_graph_edge_data.py | 335 +++++++++++++++++++++++++++++ 2 files changed, 508 insertions(+) create mode 100644 tests/cuda_graph_edge_data_test.cu create mode 100644 tests/test_graph_edge_data.py diff --git a/tests/cuda_graph_edge_data_test.cu b/tests/cuda_graph_edge_data_test.cu new file mode 100644 index 00000000..99b3cf4e --- /dev/null +++ b/tests/cuda_graph_edge_data_test.cu @@ -0,0 +1,173 @@ +// Test-only CUDA helpers for graph dependency edge-data integration coverage. + +#include +#include +#include +#include +#include + +#include +#include +#include +#include + +#include "BinaryGraphFormat.h" +#include "CUDAGraph.h" + +namespace { + +void check_cuda(cudaError_t result, const char* operation) { + TORCH_CHECK(result == cudaSuccess, operation, " failed: ", cudaGetErrorString(result)); +} + +void check_driver(CUresult result, const char* operation) { + const char* error = nullptr; + if (result != CUDA_SUCCESS) { + cuGetErrorString(result, &error); + } + TORCH_CHECK(result == CUDA_SUCCESS, operation, " failed: ", error ? error : "unknown"); +} + +__global__ void write_values(float* values, int count) { + int index = blockIdx.x * blockDim.x + threadIdx.x; + if (index < count) { + values[index] = static_cast(index + 1); + } + cudaTriggerProgrammaticLaunchCompletion(); +} + +__global__ void double_values_programmatic(const float* values, float* output, int count) { + cudaGridDependencySynchronize(); + int index = blockIdx.x * blockDim.x + threadIdx.x; + if (index < count) { + output[index] = values[index] * 2.0f; + } +} + +__global__ void double_values_default(const float* values, float* output, int count) { + int index = blockIdx.x * blockDim.x + threadIdx.x; + if (index < count) { + output[index] = values[index] * 2.0f; + } +} + +void validate_tensors(const torch::Tensor& values, const torch::Tensor& output) { + TORCH_CHECK(values.is_cuda() && output.is_cuda(), "expected CUDA tensors"); + TORCH_CHECK(values.scalar_type() == torch::kFloat32, "expected float32 input"); + TORCH_CHECK(output.scalar_type() == torch::kFloat32, "expected float32 output"); + TORCH_CHECK(values.is_contiguous() && output.is_contiguous(), "expected contiguous tensors"); + TORCH_CHECK(values.numel() == output.numel(), "input and output sizes must match"); +} + +} // namespace + +void launch_programmatic(torch::Tensor values, torch::Tensor output) { + validate_tensors(values, output); + const int count = static_cast(values.numel()); + constexpr int threads = 128; + const int blocks = (count + threads - 1) / threads; + cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + + write_values<<>>(values.data_ptr(), count); + check_cuda(cudaGetLastError(), "write_values launch"); + + cudaLaunchAttribute attribute{}; + attribute.id = cudaLaunchAttributeProgrammaticStreamSerialization; + attribute.val.programmaticStreamSerializationAllowed = 1; + + cudaLaunchConfig_t config{}; + config.gridDim = dim3(blocks); + config.blockDim = dim3(threads); + config.stream = stream; + config.attrs = &attribute; + config.numAttrs = 1; + check_cuda(cudaLaunchKernelEx(&config, double_values_programmatic, values.data_ptr(), + output.data_ptr(), count), + "double_values_programmatic launch"); +} + +void launch_default(torch::Tensor values, torch::Tensor output) { + validate_tensors(values, output); + const int count = static_cast(values.numel()); + constexpr int threads = 128; + const int blocks = (count + threads - 1) / threads; + cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + + write_values<<>>(values.data_ptr(), count); + check_cuda(cudaGetLastError(), "write_values launch"); + double_values_default<<>>(values.data_ptr(), + output.data_ptr(), count); + check_cuda(cudaGetLastError(), "double_values_default launch"); +} + +std::vector> graph_edges( + uint64_t graph_address) { + CUgraph graph = reinterpret_cast(graph_address); + + size_t node_count = 0; + check_driver(cuGraphGetNodes(graph, nullptr, &node_count), "cuGraphGetNodes(count)"); + std::vector nodes(node_count); + check_driver(cuGraphGetNodes(graph, nodes.data(), &node_count), "cuGraphGetNodes"); + + std::unordered_map node_indices; + for (size_t index = 0; index < node_count; ++index) { + node_indices[nodes[index]] = static_cast(index); + } + + size_t edge_count = 0; +#if CUDA_VERSION >= 13000 + check_driver(cuGraphGetEdges(graph, nullptr, nullptr, nullptr, &edge_count), + "cuGraphGetEdges(count)"); +#else + check_driver(cuGraphGetEdges_v2(graph, nullptr, nullptr, nullptr, &edge_count), + "cuGraphGetEdges_v2(count)"); +#endif + + std::vector from_nodes(edge_count); + std::vector to_nodes(edge_count); + std::vector edge_data(edge_count); +#if CUDA_VERSION >= 13000 + check_driver(cuGraphGetEdges(graph, from_nodes.data(), to_nodes.data(), edge_data.data(), + &edge_count), + "cuGraphGetEdges"); +#else + check_driver(cuGraphGetEdges_v2(graph, from_nodes.data(), to_nodes.data(), edge_data.data(), + &edge_count), + "cuGraphGetEdges_v2"); +#endif + + std::vector> result; + result.reserve(edge_count); + for (size_t index = 0; index < edge_count; ++index) { + result.emplace_back(node_indices.at(from_nodes[index]), node_indices.at(to_nodes[index]), + edge_data[index].from_port, edge_data[index].to_port, + edge_data[index].type); + } + return result; +} + +uint64_t shared_raw_cuda_graph(uint64_t foundry_graph_address) { + auto* graph = reinterpret_cast(foundry_graph_address); + TORCH_CHECK(graph->on_demand_data_ != nullptr, "graph has no on-demand data"); + TORCH_CHECK(graph->on_demand_data_->shared_exec != nullptr, "graph has no shared executor"); + TORCH_CHECK(graph->on_demand_data_->shared_exec->graph != nullptr, + "shared executor has no CUDA graph"); + return reinterpret_cast(graph->on_demand_data_->shared_exec->graph); +} + +std::string read_binary_graph_json(const std::string& path) { + return boost::json::serialize(foundry::read_and_parse_binary_graph(path)); +} + +bool binary_graph_valid(const std::string& path) { + return foundry::read_binary_graph_file(path).valid(); +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, module) { + module.def("launch_programmatic", &launch_programmatic); + module.def("launch_default", &launch_default); + module.def("graph_edges", &graph_edges); + module.def("shared_raw_cuda_graph", &shared_raw_cuda_graph); + module.def("read_binary_graph_json", &read_binary_graph_json); + module.def("binary_graph_valid", &binary_graph_valid); +} diff --git a/tests/test_graph_edge_data.py b/tests/test_graph_edge_data.py new file mode 100644 index 00000000..e428a61b --- /dev/null +++ b/tests/test_graph_edge_data.py @@ -0,0 +1,335 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the Foundry project +from __future__ import annotations + +import copy +import importlib.util +import json +import os +import shutil +import struct +import subprocess +import sys +from pathlib import Path + +import foundry as fdry +import pytest +import torch +from foundry import ops +from torch.utils.cpp_extension import load + +BASE_ADDR = 0x4F0000000000 +REGION_SIZE = "1GB" +TENSOR_SIZE = 256 +TEST_ROOT_ENV = "FOUNDRY_EDGE_DATA_TEST_ROOT" + +REPO_ROOT = Path(__file__).resolve().parents[1] +CUDA_HELPER = Path(__file__).with_name("cuda_graph_edge_data_test.cu") +BINARY_GRAPH_IO = REPO_ROOT / "csrc" / "BinaryGraphIO.cpp" + +FILE_HEADER_SIZE = 64 +SECTION_ENTRY_SIZE = 24 +DEPENDENCY_SECTION = 2 +LEGACY_DEPENDENCY_SIZE = 8 +EDGE_DEPENDENCY_SIZE = 16 + + +def _test_root() -> Path: + return Path(os.environ.get(TEST_ROOT_ENV, "test_graph_edge_data")) + + +def _archive_dir() -> Path: + return _test_root() / "graphs" + + +def _hook_archive_dir() -> Path: + return _test_root() / "hook_archive" + + +def _load_cuda_helper(): + return load( + name="foundry_cuda_graph_edge_data_test", + sources=[str(CUDA_HELPER), str(BINARY_GRAPH_IO)], + extra_include_paths=[str(REPO_ROOT / "include")], + extra_cflags=["-std=c++17"], + extra_cuda_cflags=["-std=c++17"], + extra_ldflags=["-lboost_json", "-lboost_filesystem"], + verbose=False, + ) + + +def _hook_so_path() -> str: + spec = importlib.util.find_spec("foundry.ops") + if spec is None or spec.origin is None: + raise RuntimeError("foundry.ops is unavailable") + path = Path(spec.origin).resolve().parent / "libcuda_hook.so" + if not path.is_file(): + raise RuntimeError(f"Foundry hook is unavailable: {path}") + return str(path) + + +def _dependency_section(data: bytes | bytearray) -> tuple[int, int, int]: + (num_dependencies,) = struct.unpack_from(" None: + data = bytearray(source.read_bytes()) + num_dependencies, section_offset, section_size = _dependency_section(data) + if num_dependencies == 0: + raise AssertionError("test graph unexpectedly has no dependencies") + entry_size = section_size // num_dependencies + if section_size != entry_size * num_dependencies: + raise AssertionError("source dependency section is malformed") + + if entry_size == LEGACY_DEPENDENCY_SIZE: + destination.write_bytes(data) + return + if entry_size != EDGE_DEPENDENCY_SIZE: + raise AssertionError(f"unexpected dependency entry size: {entry_size}") + + legacy_records = b"".join( + data[ + section_offset + index * EDGE_DEPENDENCY_SIZE : section_offset + + index * EDGE_DEPENDENCY_SIZE + + LEGACY_DEPENDENCY_SIZE + ] + for index in range(num_dependencies) + ) + old_end = section_offset + section_size + new_data = data[:section_offset] + legacy_records + data[old_end:] + removed = section_size - len(legacy_records) + + (num_sections,) = struct.unpack_from(" section_offset: + struct.pack_into(" None: + data = bytearray(source.read_bytes()) + num_dependencies, _section_offset, section_size = _dependency_section(data) + if num_dependencies == 0 or section_size < 2: + raise AssertionError("test graph unexpectedly has no dependency payload") + entry_offset = FILE_HEADER_SIZE + DEPENDENCY_SECTION * SECTION_ENTRY_SIZE + struct.pack_into(" tuple[int, int, int, int, int]: + return ( + value["from"], + value["to"], + value["from_port"], + value["to_port"], + value["type"], + ) + + +def _expected_values() -> torch.Tensor: + return torch.arange(1, TENSOR_SIZE + 1, dtype=torch.float32) * 2 + + +def _allocate_graph_tensors() -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + return tuple(torch.zeros(TENSOR_SIZE, device="cuda", dtype=torch.float32) for _ in range(4)) + + +def _capture_graph(helper, launch, values: torch.Tensor, output: torch.Tensor): + stream = torch.cuda.Stream() + graph = fdry.CUDAGraph(keep_graph=True) + with fdry.graph(graph, stream=stream): + launch(values, output) + torch.cuda.synchronize() + return graph + + +def _run_saving_process() -> None: + helper = _load_cuda_helper() + torch.cuda.init() + torch.cuda.set_device(0) + torch.set_default_device("cuda") + + archive_dir = _archive_dir() + archive_dir.mkdir(parents=True, exist_ok=True) + fdry.set_allocation_region(BASE_ADDR, fdry.parse_size(REGION_SIZE)) + programmatic_values, programmatic_output, default_values, default_output = ( + _allocate_graph_tensors() + ) + + programmatic_graph = _capture_graph( + helper, + helper.launch_programmatic, + programmatic_values, + programmatic_output, + ) + default_graph = _capture_graph( + helper, + helper.launch_default, + default_values, + default_output, + ) + + captured_programmatic_edges = [ + tuple(edge) for edge in helper.graph_edges(programmatic_graph.raw_cuda_graph()) + ] + captured_default_edges = [ + tuple(edge) for edge in helper.graph_edges(default_graph.raw_cuda_graph()) + ] + assert any(edge[2:] == (1, 0, 1) for edge in captured_programmatic_edges) + assert captured_default_edges + assert all(edge[2:] == (0, 0, 0) for edge in captured_default_edges) + + programmatic_path = archive_dir / "graph_0.json" + default_path = archive_dir / "graph_1.json" + programmatic_graph.save(str(programmatic_path), output_tensors=programmatic_output) + default_graph.save(str(default_path), output_tensors=default_output) + + programmatic_json = json.loads(programmatic_path.read_text()) + default_json = json.loads(default_path.read_text()) + serialized_edges = [_edge_tuple(dependency) for dependency in programmatic_json["dependencies"]] + assert sorted(serialized_edges) == sorted(captured_programmatic_edges) + assert [node["type"] for node in programmatic_json["nodes"]] == [ + node["type"] for node in default_json["nodes"] + ] + assert programmatic_json["topology_key"] != default_json["topology_key"] + + fdry.save_graph_manifest(str(archive_dir)) + manifest = json.loads((archive_dir / "graph_manifest.json").read_text()) + assert len(manifest["topology_groups"]) == 2 + + binary_path = programmatic_path.with_suffix(".cugraph") + num_dependencies, _section_offset, section_size = _dependency_section(binary_path.read_bytes()) + assert section_size == num_dependencies * EDGE_DEPENDENCY_SIZE + binary_json = json.loads(helper.read_binary_graph_json(str(binary_path))) + assert sorted(_edge_tuple(value) for value in binary_json["dependencies"]) == sorted( + captured_programmatic_edges + ) + + legacy_json = copy.deepcopy(programmatic_json) + for dependency in legacy_json["dependencies"]: + dependency.pop("from_port") + dependency.pop("to_port") + dependency.pop("type") + legacy_path = archive_dir / "legacy.json" + legacy_path.write_text(json.dumps(legacy_json)) + legacy_binary = legacy_path.with_suffix(".cugraph") + _write_legacy_binary(binary_path, legacy_binary) + legacy_binary_json = json.loads(helper.read_binary_graph_json(str(legacy_binary))) + assert sorted(_edge_tuple(value) for value in legacy_binary_json["dependencies"]) == sorted( + (edge[0], edge[1], 0, 0, 0) for edge in captured_programmatic_edges + ) + + malformed_binary = archive_dir / "malformed.cugraph" + _write_malformed_binary(binary_path, malformed_binary) + assert not helper.binary_graph_valid(str(malformed_binary)) + assert json.loads(helper.read_binary_graph_json(str(malformed_binary)) or "null") is None + + programmatic_graph.replay() + default_graph.replay() + torch.cuda.synchronize() + torch.testing.assert_close(programmatic_output.cpu(), _expected_values()) + torch.testing.assert_close(default_output.cpu(), _expected_values()) + fdry.stop_allocation_region() + + +def _load_graph_path(path: Path, *, parallel: bool): + if parallel: + pending = fdry.CUDAGraph.start_graph_builds([str(path)], num_threads=1) + graph, output = fdry.CUDAGraph.finish_one_graph_load(pending, 0) + return graph, output + return fdry.CUDAGraph.load(str(path)) + + +def _run_loading_process(*, legacy: bool, parallel: bool) -> None: + torch.cuda.init() + torch.cuda.set_device(0) + torch.set_default_device("cuda") + fdry.load_cuda_modules_and_libraries(str(_hook_archive_dir())) + helper = _load_cuda_helper() + + fdry.set_allocation_region(BASE_ADDR, fdry.parse_size(REGION_SIZE)) + _allocate_graph_tensors() + path = _archive_dir() / ("legacy.json" if legacy else "graph_0.json") + graph, output = _load_graph_path(path, parallel=parallel) + + if parallel: + graph_address = ops.get_graph_ptr(graph) + raw_graph = helper.shared_raw_cuda_graph(graph_address) + else: + raw_graph = graph.raw_cuda_graph() + loaded_edges = [tuple(edge) for edge in helper.graph_edges(raw_graph)] + + source_json = json.loads(path.read_text()) + expected_edges = [] + for dependency in source_json["dependencies"]: + expected_edges.append( + ( + dependency["from"], + dependency["to"], + dependency.get("from_port", 0), + dependency.get("to_port", 0), + dependency.get("type", 0), + ) + ) + assert sorted(loaded_edges) == sorted(expected_edges) + if legacy: + assert all(edge[2:] == (0, 0, 0) for edge in loaded_edges) + else: + assert any(edge[2:] == (1, 0, 1) for edge in loaded_edges) + + graph.replay() + torch.cuda.synchronize() + torch.testing.assert_close(output.cpu(), _expected_values()) + fdry.stop_allocation_region() + + +def _spawn(mode: str, root: Path) -> None: + environment = os.environ.copy() + environment[TEST_ROOT_ENV] = str(root) + hook = _hook_so_path() + environment["LD_PRELOAD"] = ( + f"{hook}:{environment['LD_PRELOAD']}" if environment.get("LD_PRELOAD") else hook + ) + subprocess.check_call( + [sys.executable, str(Path(__file__).resolve()), f"--{mode}"], + cwd=root, + env=environment, + ) + + +@pytest.mark.filterwarnings("ignore:TORCH_CUDA_ARCH_LIST is not set") +def test_cuda_graph_edge_data_round_trip(tmp_path): + root = tmp_path / "edge-data-round-trip" + root.mkdir() + try: + _spawn("saving-run", root) + _spawn("serial-load", root) + _spawn("serial-legacy-load", root) + _spawn("parallel-load", root) + _spawn("parallel-legacy-load", root) + finally: + shutil.rmtree(root, ignore_errors=True) + + +if __name__ == "__main__": + if "--saving-run" in sys.argv: + _run_saving_process() + elif "--serial-load" in sys.argv: + _run_loading_process(legacy=False, parallel=False) + elif "--serial-legacy-load" in sys.argv: + _run_loading_process(legacy=True, parallel=False) + elif "--parallel-load" in sys.argv: + _run_loading_process(legacy=False, parallel=True) + elif "--parallel-legacy-load" in sys.argv: + _run_loading_process(legacy=True, parallel=True) + else: + raise SystemExit("expected a test subprocess mode") From 721f7834b9c973687b9feae3638b89092dfba8ab Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Fri, 24 Jul 2026 00:09:48 +0000 Subject: [PATCH 24/59] test: link edge metadata helper to CUDA driver Signed-off-by: Cursor Agent Co-authored-by: Rahul Chalamala --- tests/test_graph_edge_data.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/tests/test_graph_edge_data.py b/tests/test_graph_edge_data.py index e428a61b..03f54696 100644 --- a/tests/test_graph_edge_data.py +++ b/tests/test_graph_edge_data.py @@ -53,7 +53,12 @@ def _load_cuda_helper(): extra_include_paths=[str(REPO_ROOT / "include")], extra_cflags=["-std=c++17"], extra_cuda_cflags=["-std=c++17"], - extra_ldflags=["-lboost_json", "-lboost_filesystem"], + extra_ldflags=[ + "-L/usr/local/cuda/lib64/stubs", + "-lcuda", + "-lboost_json", + "-lboost_filesystem", + ], verbose=False, ) From d81340eeec0efdda8881399b7c13cd4a9fc55da3 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Fri, 24 Jul 2026 00:18:11 +0000 Subject: [PATCH 25/59] fix: preserve CUDA graph dependency edge metadata Signed-off-by: Cursor Agent Co-authored-by: Rahul Chalamala --- csrc/BinaryGraphIO.cpp | 78 ++++++++++++++++++++++++++++++++++--- csrc/CUDAGraph.cpp | 42 ++++++++++++++++++-- csrc/CUDAGraphParallel.cpp | 19 ++++++++- include/BinaryGraphFormat.h | 14 ++++++- include/metadata.h | 1 + 5 files changed, 141 insertions(+), 13 deletions(-) diff --git a/csrc/BinaryGraphIO.cpp b/csrc/BinaryGraphIO.cpp index 79c8ac81..419e056d 100644 --- a/csrc/BinaryGraphIO.cpp +++ b/csrc/BinaryGraphIO.cpp @@ -12,6 +12,46 @@ namespace foundry { +namespace { + +bool validate_binary_layout(const uint8_t* data, size_t file_size, + const binary_format::FileHeader& header) { + namespace bf = binary_format; + + if (header.num_sections != bf::SECTION_COUNT) { + return false; + } + + const uint64_t directory_size = + static_cast(header.num_sections) * sizeof(bf::SectionEntry); + if (directory_size > file_size - sizeof(bf::FileHeader)) { + return false; + } + + const auto* sections = reinterpret_cast(data + sizeof(bf::FileHeader)); + for (uint32_t i = 0; i < header.num_sections; ++i) { + const auto& section = sections[i]; + if (section.section_type != i || section.offset > file_size || + section.size > file_size - section.offset) { + return false; + } + } + + const uint64_t legacy_dependency_size = + static_cast(header.num_dependencies) * sizeof(bf::LegacyBinDependency); + const uint64_t dependency_size = + static_cast(header.num_dependencies) * sizeof(bf::BinDependency); + const uint64_t stored_dependency_size = sections[bf::SECTION_DEPENDENCY_TABLE].size; + const bool has_dependencies = (header.flags & bf::FLAG_HAS_DEPENDENCIES) != 0; + if (has_dependencies != (header.num_dependencies != 0)) { + return false; + } + return stored_dependency_size == legacy_dependency_size || + stored_dependency_size == dependency_size; +} + +} // namespace + void CUDAGraph::save_binary(const std::string& bin_path, const boost::json::object& root) { namespace json = boost::json; namespace bf = binary_format; @@ -270,7 +310,19 @@ void CUDAGraph::save_binary(const std::string& bin_path, const boost::json::obje deps.reserve(deps_array.size()); for (const auto& dep_val : deps_array) { const json::object& d = dep_val.as_object(); - deps.push_back({d.at("from").to_number(), d.at("to").to_number()}); + bf::BinDependency dep = {}; + dep.from_id = d.at("from").to_number(); + dep.to_id = d.at("to").to_number(); + if (d.contains("from_port")) { + dep.from_port = static_cast(d.at("from_port").to_number()); + } + if (d.contains("to_port")) { + dep.to_port = static_cast(d.at("to_port").to_number()); + } + if (d.contains("type")) { + dep.type = static_cast(d.at("type").to_number()); + } + deps.push_back(dep); } // ---- Generators ---- @@ -397,7 +449,7 @@ boost::json::value read_and_parse_binary_graph(const std::string& bin_path) { if (file_size < sizeof(bf::FileHeader)) return json::value{}; const bf::FileHeader& header = *reinterpret_cast(data); - if (!bf::validate_header(header)) + if (!bf::validate_header(header) || !validate_binary_layout(data, file_size, header)) return json::value{}; // Read section directory @@ -653,13 +705,26 @@ boost::json::value read_and_parse_binary_graph(const std::string& bin_path) { // Dependencies if (header.flags & bf::FLAG_HAS_DEPENDENCIES) { auto [dep_data, dep_size] = section(bf::SECTION_DEPENDENCY_TABLE); - const bf::BinDependency* deps = reinterpret_cast(dep_data); + const size_t dependency_entry_size = dep_size / header.num_dependencies; + const bool has_edge_data = dependency_entry_size == sizeof(bf::BinDependency); json::array deps_array; deps_array.reserve(header.num_dependencies); for (uint32_t i = 0; i < header.num_dependencies; i++) { + const uint8_t* record = dep_data + i * dependency_entry_size; + const auto* legacy = reinterpret_cast(record); json::object d; - d["from"] = deps[i].from_id; - d["to"] = deps[i].to_id; + d["from"] = legacy->from_id; + d["to"] = legacy->to_id; + if (has_edge_data) { + const auto* dep = reinterpret_cast(record); + d["from_port"] = static_cast(dep->from_port); + d["to_port"] = static_cast(dep->to_port); + d["type"] = static_cast(dep->type); + } else { + d["from_port"] = 0; + d["to_port"] = 0; + d["type"] = 0; + } deps_array.push_back(d); } root["dependencies"] = std::move(deps_array); @@ -725,7 +790,8 @@ BinaryGraphFile read_binary_graph_file(const std::string& bin_path) { in.close(); memcpy(&result.header, result.data.data(), sizeof(bf::FileHeader)); - if (!bf::validate_header(result.header)) { + if (!bf::validate_header(result.header) || + !validate_binary_layout(result.data.data(), file_size, result.header)) { result.data.clear(); return result; } diff --git a/csrc/CUDAGraph.cpp b/csrc/CUDAGraph.cpp index 25df3524..03f20d4a 100644 --- a/csrc/CUDAGraph.cpp +++ b/csrc/CUDAGraph.cpp @@ -934,6 +934,7 @@ void CUDAGraph::analyze_captured_graph() { GraphDependency dep; dep.from_index = from_it->second; dep.to_index = to_it->second; + dep.edge_data = edges[i]; graph_dependencies.push_back(dep); } } @@ -1335,7 +1336,8 @@ void CUDAGraph::save(const std::string& json_path, const OutputTensors& output_t nodes_array.push_back(node_obj); } - // Compute topology key = node types + cluster dim values per kernel node. + // Compute topology key = node types + cluster dim values per kernel node + + // dependency endpoints and edge data. // cuGraphExecUpdate does NOT propagate kernel node attribute changes // (set via cuGraphKernelNodeSetAttribute) to the CUgraphExec — only // CUDA_KERNEL_NODE_PARAMS fields are updated. This means cluster dim @@ -1381,6 +1383,22 @@ void CUDAGraph::save(const std::string& json_path, const OutputTensors& output_t } } } + + std::vector sorted_dependencies = graph_dependencies; + std::sort(sorted_dependencies.begin(), sorted_dependencies.end(), + [](const GraphDependency& lhs, const GraphDependency& rhs) { + return std::make_tuple(lhs.from_index, lhs.to_index, lhs.edge_data.from_port, + lhs.edge_data.to_port, lhs.edge_data.type) < + std::make_tuple(rhs.from_index, rhs.to_index, rhs.edge_data.from_port, + rhs.edge_data.to_port, rhs.edge_data.type); + }); + topology_key += "|D"; + for (const auto& dep : sorted_dependencies) { + topology_key += ";" + std::to_string(dep.from_index) + ">" + std::to_string(dep.to_index) + + ":" + std::to_string(static_cast(dep.edge_data.from_port)) + ":" + + std::to_string(static_cast(dep.edge_data.to_port)) + ":" + + std::to_string(static_cast(dep.edge_data.type)); + } root["topology_key"] = topology_key; } @@ -1447,6 +1465,9 @@ void CUDAGraph::save(const std::string& json_path, const OutputTensors& output_t json::object dep_obj; dep_obj["from"] = dep.from_index; dep_obj["to"] = dep.to_index; + dep_obj["from_port"] = static_cast(dep.edge_data.from_port); + dep_obj["to_port"] = static_cast(dep.edge_data.to_port); + dep_obj["type"] = static_cast(dep.edge_data.type); deps_array.push_back(dep_obj); } root["dependencies"] = deps_array; @@ -2031,24 +2052,39 @@ GraphLoadResult CUDAGraph::load(const std::string& json_path, MempoolId_t pool) if (!deps_array.empty()) { std::vector from_nodes; std::vector to_nodes; + std::vector edge_data; from_nodes.reserve(deps_array.size()); to_nodes.reserve(deps_array.size()); + edge_data.reserve(deps_array.size()); for (const auto& dep_val : deps_array) { const json::object& dep_obj = dep_val.as_object(); int from_id = dep_obj.at("from").to_number(); int to_id = dep_obj.at("to").to_number(); + CUgraphEdgeData data{}; + if (dep_obj.contains("from_port")) { + data.from_port = static_cast( + dep_obj.at("from_port").to_number()); + } + if (dep_obj.contains("to_port")) { + data.to_port = + static_cast(dep_obj.at("to_port").to_number()); + } + if (dep_obj.contains("type")) { + data.type = static_cast(dep_obj.at("type").to_number()); + } from_nodes.push_back(id_to_node[from_id]); to_nodes.push_back(id_to_node[to_id]); + edge_data.push_back(data); } #if (defined(CUDA_VERSION) && CUDA_VERSION >= 13000) CUresult dep_result = cuGraphAddDependencies(cuGraph, from_nodes.data(), to_nodes.data(), - nullptr, deps_array.size()); + edge_data.data(), deps_array.size()); #else CUresult dep_result = cuGraphAddDependencies_v2(cuGraph, from_nodes.data(), to_nodes.data(), - nullptr, deps_array.size()); + edge_data.data(), deps_array.size()); #endif if (dep_result != CUDA_SUCCESS) { fprintf(stderr, "[foundry LOAD ERROR] cuGraphAddDependencies FAILED with error %d\n", diff --git a/csrc/CUDAGraphParallel.cpp b/csrc/CUDAGraphParallel.cpp index 6ebfc0f6..af5cd774 100644 --- a/csrc/CUDAGraphParallel.cpp +++ b/csrc/CUDAGraphParallel.cpp @@ -776,21 +776,36 @@ GraphLoadResult CUDAGraph::build_graph_from_parsed(ParsedGraphData&& parsed, CUc if (!deps_array.empty()) { std::vector from_nodes; std::vector to_nodes; + std::vector edge_data; from_nodes.reserve(deps_array.size()); to_nodes.reserve(deps_array.size()); + edge_data.reserve(deps_array.size()); for (const auto& dep_val : deps_array) { const json::object& dep_obj = dep_val.as_object(); from_nodes.push_back(id_to_node[dep_obj.at("from").to_number()]); to_nodes.push_back(id_to_node[dep_obj.at("to").to_number()]); + CUgraphEdgeData data{}; + if (dep_obj.contains("from_port")) { + data.from_port = static_cast( + dep_obj.at("from_port").to_number()); + } + if (dep_obj.contains("to_port")) { + data.to_port = + static_cast(dep_obj.at("to_port").to_number()); + } + if (dep_obj.contains("type")) { + data.type = static_cast(dep_obj.at("type").to_number()); + } + edge_data.push_back(data); } #if (defined(CUDA_VERSION) && CUDA_VERSION >= 13000) CUresult dep_result = cuGraphAddDependencies(cuGraph, from_nodes.data(), to_nodes.data(), - nullptr, deps_array.size()); + edge_data.data(), deps_array.size()); #else CUresult dep_result = cuGraphAddDependencies_v2(cuGraph, from_nodes.data(), to_nodes.data(), - nullptr, deps_array.size()); + edge_data.data(), deps_array.size()); #endif if (dep_result != CUDA_SUCCESS) { fprintf(stderr, "[foundry LOAD ERROR] cuGraphAddDependencies FAILED with error %d\n", diff --git a/include/BinaryGraphFormat.h b/include/BinaryGraphFormat.h index ca791640..4a002851 100644 --- a/include/BinaryGraphFormat.h +++ b/include/BinaryGraphFormat.h @@ -21,7 +21,7 @@ class value; // Sections: // STRING_TABLE — deduplicated null-terminated strings // NODE_TABLE — fixed-size node descriptors (160 bytes each) -// DEPENDENCY_TABLE — (from_id, to_id) pairs +// DEPENDENCY_TABLE — fixed-width endpoint and CUgraphEdgeData records // KERNEL_PARAM_INDEX — per-kernel param descriptors // KERNEL_PARAM_DATA — raw binary param bytes // ARG_BUFFER_DATA — raw binary extra arg buffers @@ -220,11 +220,21 @@ static_assert(sizeof(BinGenerator) == 24, "BinGenerator must be 24 bytes"); // ---- Dependency ---- +struct LegacyBinDependency { + uint32_t from_id; + uint32_t to_id; +}; +static_assert(sizeof(LegacyBinDependency) == 8, "LegacyBinDependency must be 8 bytes"); + struct BinDependency { uint32_t from_id; uint32_t to_id; + uint8_t from_port; + uint8_t to_port; + uint8_t type; + uint8_t reserved[5]; }; -static_assert(sizeof(BinDependency) == 8, "BinDependency must be 8 bytes"); +static_assert(sizeof(BinDependency) == 16, "BinDependency must be 16 bytes"); #pragma pack(pop) diff --git a/include/metadata.h b/include/metadata.h index 4d867ab9..bdda637e 100644 --- a/include/metadata.h +++ b/include/metadata.h @@ -202,6 +202,7 @@ struct EmptyNodeMetadata {}; struct GraphDependency { int from_index; int to_index; + CUgraphEdgeData edge_data{}; }; using GraphNodeMetadata = From 52ca812ccd474b349c023a46c90319b853e17d6f Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Fri, 24 Jul 2026 00:33:00 +0000 Subject: [PATCH 26/59] test: use canonical graph archive names Signed-off-by: Cursor Agent Co-authored-by: Rahul Chalamala --- tests/test_graph_edge_data.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/test_graph_edge_data.py b/tests/test_graph_edge_data.py index 03f54696..5aaed5d8 100644 --- a/tests/test_graph_edge_data.py +++ b/tests/test_graph_edge_data.py @@ -193,8 +193,8 @@ def _run_saving_process() -> None: assert captured_default_edges assert all(edge[2:] == (0, 0, 0) for edge in captured_default_edges) - programmatic_path = archive_dir / "graph_0.json" - default_path = archive_dir / "graph_1.json" + programmatic_path = archive_dir / "graph_0_FULL_t1_r1_UX_pcN.json" + default_path = archive_dir / "graph_1_FULL_t1_r1_UX_pcN.json" programmatic_graph.save(str(programmatic_path), output_tensors=programmatic_output) default_graph.save(str(default_path), output_tensors=default_output) @@ -263,7 +263,7 @@ def _run_loading_process(*, legacy: bool, parallel: bool) -> None: fdry.set_allocation_region(BASE_ADDR, fdry.parse_size(REGION_SIZE)) _allocate_graph_tensors() - path = _archive_dir() / ("legacy.json" if legacy else "graph_0.json") + path = _archive_dir() / ("legacy.json" if legacy else "graph_0_FULL_t1_r1_UX_pcN.json") graph, output = _load_graph_path(path, parallel=parallel) if parallel: From d98f9b363614618f292883a9fa65bcec428bc2d8 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Fri, 24 Jul 2026 00:42:07 +0000 Subject: [PATCH 27/59] test: compare edge replay output on CPU Signed-off-by: Cursor Agent Co-authored-by: Rahul Chalamala --- tests/test_graph_edge_data.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_graph_edge_data.py b/tests/test_graph_edge_data.py index 5aaed5d8..f8386d6f 100644 --- a/tests/test_graph_edge_data.py +++ b/tests/test_graph_edge_data.py @@ -141,7 +141,7 @@ def _edge_tuple(value: dict) -> tuple[int, int, int, int, int]: def _expected_values() -> torch.Tensor: - return torch.arange(1, TENSOR_SIZE + 1, dtype=torch.float32) * 2 + return torch.arange(1, TENSOR_SIZE + 1, dtype=torch.float32, device="cpu") * 2 def _allocate_graph_tensors() -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: From aa6a920797b87bbbde3ff5723b097b3649152302 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Fri, 24 Jul 2026 00:55:33 +0000 Subject: [PATCH 28/59] test: isolate legacy parallel graph fixture Signed-off-by: Cursor Agent Co-authored-by: Rahul Chalamala --- tests/test_graph_edge_data.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/tests/test_graph_edge_data.py b/tests/test_graph_edge_data.py index f8386d6f..e4bf8339 100644 --- a/tests/test_graph_edge_data.py +++ b/tests/test_graph_edge_data.py @@ -264,6 +264,12 @@ def _run_loading_process(*, legacy: bool, parallel: bool) -> None: fdry.set_allocation_region(BASE_ADDR, fdry.parse_size(REGION_SIZE)) _allocate_graph_tensors() path = _archive_dir() / ("legacy.json" if legacy else "graph_0_FULL_t1_r1_UX_pcN.json") + if legacy and parallel: + isolated_dir = _archive_dir() / "legacy_parallel" + isolated_dir.mkdir(exist_ok=True) + path = isolated_dir / "graph_0_FULL_t1_r1_UX_pcN.json" + shutil.copyfile(_archive_dir() / "legacy.json", path) + shutil.copyfile(_archive_dir() / "legacy.cugraph", path.with_suffix(".cugraph")) graph, output = _load_graph_path(path, parallel=parallel) if parallel: From 60ac2a6eec13e8f55a3fc89ab9a97cb87f5dca75 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Fri, 24 Jul 2026 01:04:02 +0000 Subject: [PATCH 29/59] style: format edge metadata CUDA helper Signed-off-by: Cursor Agent Co-authored-by: Rahul Chalamala --- tests/cuda_graph_edge_data_test.cu | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/tests/cuda_graph_edge_data_test.cu b/tests/cuda_graph_edge_data_test.cu index 99b3cf4e..2095feec 100644 --- a/tests/cuda_graph_edge_data_test.cu +++ b/tests/cuda_graph_edge_data_test.cu @@ -96,7 +96,7 @@ void launch_default(torch::Tensor values, torch::Tensor output) { write_values<<>>(values.data_ptr(), count); check_cuda(cudaGetLastError(), "write_values launch"); double_values_default<<>>(values.data_ptr(), - output.data_ptr(), count); + output.data_ptr(), count); check_cuda(cudaGetLastError(), "double_values_default launch"); } @@ -127,13 +127,13 @@ std::vector> graph_edges std::vector to_nodes(edge_count); std::vector edge_data(edge_count); #if CUDA_VERSION >= 13000 - check_driver(cuGraphGetEdges(graph, from_nodes.data(), to_nodes.data(), edge_data.data(), - &edge_count), - "cuGraphGetEdges"); + check_driver( + cuGraphGetEdges(graph, from_nodes.data(), to_nodes.data(), edge_data.data(), &edge_count), + "cuGraphGetEdges"); #else - check_driver(cuGraphGetEdges_v2(graph, from_nodes.data(), to_nodes.data(), edge_data.data(), - &edge_count), - "cuGraphGetEdges_v2"); + check_driver( + cuGraphGetEdges_v2(graph, from_nodes.data(), to_nodes.data(), edge_data.data(), &edge_count), + "cuGraphGetEdges_v2"); #endif std::vector> result; From fc173a3581ea501ab7c4f5a4d1697bf5dd09abe8 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Fri, 24 Jul 2026 01:55:35 +0000 Subject: [PATCH 30/59] test: cover mixed CUDA graph edge metadata Signed-off-by: Cursor Agent Co-authored-by: Rahul Chalamala --- tests/cuda_graph_edge_data_test.cu | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/tests/cuda_graph_edge_data_test.cu b/tests/cuda_graph_edge_data_test.cu index 2095feec..5844cdaf 100644 --- a/tests/cuda_graph_edge_data_test.cu +++ b/tests/cuda_graph_edge_data_test.cu @@ -51,6 +51,8 @@ __global__ void double_values_default(const float* values, float* output, int co } } +__global__ void complete_after_values() {} + void validate_tensors(const torch::Tensor& values, const torch::Tensor& output) { TORCH_CHECK(values.is_cuda() && output.is_cuda(), "expected CUDA tensors"); TORCH_CHECK(values.scalar_type() == torch::kFloat32, "expected float32 input"); @@ -84,6 +86,8 @@ void launch_programmatic(torch::Tensor values, torch::Tensor output) { check_cuda(cudaLaunchKernelEx(&config, double_values_programmatic, values.data_ptr(), output.data_ptr(), count), "double_values_programmatic launch"); + complete_after_values<<<1, 1, 0, stream>>>(); + check_cuda(cudaGetLastError(), "complete_after_values launch"); } void launch_default(torch::Tensor values, torch::Tensor output) { @@ -98,6 +102,8 @@ void launch_default(torch::Tensor values, torch::Tensor output) { double_values_default<<>>(values.data_ptr(), output.data_ptr(), count); check_cuda(cudaGetLastError(), "double_values_default launch"); + complete_after_values<<<1, 1, 0, stream>>>(); + check_cuda(cudaGetLastError(), "complete_after_values launch"); } std::vector> graph_edges( From f7f27a45510d430fb7a3c32e4a65d7ee994f855d Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Fri, 24 Jul 2026 02:07:33 +0000 Subject: [PATCH 31/59] fix: batch CUDA dependencies by edge metadata Signed-off-by: Cursor Agent Co-authored-by: Rahul Chalamala --- csrc/CUDAGraph.cpp | 53 ++++++++++++++++++++++++++++++++----- csrc/CUDAGraphParallel.cpp | 8 +----- include/CUDAGraphInternal.h | 5 ++++ 3 files changed, 52 insertions(+), 14 deletions(-) diff --git a/csrc/CUDAGraph.cpp b/csrc/CUDAGraph.cpp index 03f20d4a..871d9bf6 100644 --- a/csrc/CUDAGraph.cpp +++ b/csrc/CUDAGraph.cpp @@ -16,6 +16,8 @@ #include #include #include +#include +#include #include #include #include @@ -127,6 +129,49 @@ MempoolId_t torch_graph_pool_handle(bool is_user_created = true) { } // namespace +CUresult add_graph_dependencies(CUgraph graph, + const std::vector& from_nodes, + const std::vector& to_nodes, + const std::vector& edge_data) { + TORCH_INTERNAL_ASSERT(from_nodes.size() == to_nodes.size()); + TORCH_INTERNAL_ASSERT(from_nodes.size() == edge_data.size()); + + struct DependencyBatch { + std::vector from_nodes; + std::vector to_nodes; + std::vector edge_data; + }; + using EdgeDataKey = std::tuple; + std::map batches; + + // CUDA 13.0 silently normalizes non-default entries when one call mixes edge annotations. + // Group identical annotations so each driver call remains lossless without adding edges singly. + for (size_t index = 0; index < edge_data.size(); ++index) { + const auto& data = edge_data[index]; + auto& batch = batches[{data.from_port, data.to_port, data.type}]; + batch.from_nodes.push_back(from_nodes[index]); + batch.to_nodes.push_back(to_nodes[index]); + batch.edge_data.push_back(data); + } + + for (auto& [key, batch] : batches) { + (void)key; +#if (defined(CUDA_VERSION) && CUDA_VERSION >= 13000) + CUresult result = + cuGraphAddDependencies(graph, batch.from_nodes.data(), batch.to_nodes.data(), + batch.edge_data.data(), batch.edge_data.size()); +#else + CUresult result = + cuGraphAddDependencies_v2(graph, batch.from_nodes.data(), batch.to_nodes.data(), + batch.edge_data.data(), batch.edge_data.size()); +#endif + if (result != CUDA_SUCCESS) { + return result; + } + } + return CUDA_SUCCESS; +} + static bool _cuda_graphs_debug = false; static CUDAGeneratorStateRegistry global_generator_state_registry; @@ -2079,13 +2124,7 @@ GraphLoadResult CUDAGraph::load(const std::string& json_path, MempoolId_t pool) edge_data.push_back(data); } -#if (defined(CUDA_VERSION) && CUDA_VERSION >= 13000) - CUresult dep_result = cuGraphAddDependencies(cuGraph, from_nodes.data(), to_nodes.data(), - edge_data.data(), deps_array.size()); -#else - CUresult dep_result = cuGraphAddDependencies_v2(cuGraph, from_nodes.data(), to_nodes.data(), - edge_data.data(), deps_array.size()); -#endif + CUresult dep_result = add_graph_dependencies(cuGraph, from_nodes, to_nodes, edge_data); if (dep_result != CUDA_SUCCESS) { fprintf(stderr, "[foundry LOAD ERROR] cuGraphAddDependencies FAILED with error %d\n", dep_result); diff --git a/csrc/CUDAGraphParallel.cpp b/csrc/CUDAGraphParallel.cpp index af5cd774..05bfce7e 100644 --- a/csrc/CUDAGraphParallel.cpp +++ b/csrc/CUDAGraphParallel.cpp @@ -800,13 +800,7 @@ GraphLoadResult CUDAGraph::build_graph_from_parsed(ParsedGraphData&& parsed, CUc edge_data.push_back(data); } -#if (defined(CUDA_VERSION) && CUDA_VERSION >= 13000) - CUresult dep_result = cuGraphAddDependencies(cuGraph, from_nodes.data(), to_nodes.data(), - edge_data.data(), deps_array.size()); -#else - CUresult dep_result = cuGraphAddDependencies_v2(cuGraph, from_nodes.data(), to_nodes.data(), - edge_data.data(), deps_array.size()); -#endif + CUresult dep_result = add_graph_dependencies(cuGraph, from_nodes, to_nodes, edge_data); if (dep_result != CUDA_SUCCESS) { fprintf(stderr, "[foundry LOAD ERROR] cuGraphAddDependencies FAILED with error %d\n", dep_result); diff --git a/include/CUDAGraphInternal.h b/include/CUDAGraphInternal.h index 52c53789..94c7470b 100644 --- a/include/CUDAGraphInternal.h +++ b/include/CUDAGraphInternal.h @@ -8,6 +8,11 @@ namespace foundry { +CUresult add_graph_dependencies(CUgraph graph, + const std::vector& from_nodes, + const std::vector& to_nodes, + const std::vector& edge_data); + // Holds deferred metadata for the split start/finish graph loading flow. // Returned by start_graph_builds_impl, consumed by finish_graph_loads_impl. struct PendingGraphLoads { From 622e49b5639bff8bcdfe07932cbfa5534d1347d7 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Fri, 24 Jul 2026 02:17:44 +0000 Subject: [PATCH 32/59] style: format dependency batching helper Signed-off-by: Cursor Agent Co-authored-by: Rahul Chalamala --- csrc/CUDAGraph.cpp | 8 +++----- include/CUDAGraphInternal.h | 3 +-- 2 files changed, 4 insertions(+), 7 deletions(-) diff --git a/csrc/CUDAGraph.cpp b/csrc/CUDAGraph.cpp index 871d9bf6..47ee2d9e 100644 --- a/csrc/CUDAGraph.cpp +++ b/csrc/CUDAGraph.cpp @@ -129,8 +129,7 @@ MempoolId_t torch_graph_pool_handle(bool is_user_created = true) { } // namespace -CUresult add_graph_dependencies(CUgraph graph, - const std::vector& from_nodes, +CUresult add_graph_dependencies(CUgraph graph, const std::vector& from_nodes, const std::vector& to_nodes, const std::vector& edge_data) { TORCH_INTERNAL_ASSERT(from_nodes.size() == to_nodes.size()); @@ -157,9 +156,8 @@ CUresult add_graph_dependencies(CUgraph graph, for (auto& [key, batch] : batches) { (void)key; #if (defined(CUDA_VERSION) && CUDA_VERSION >= 13000) - CUresult result = - cuGraphAddDependencies(graph, batch.from_nodes.data(), batch.to_nodes.data(), - batch.edge_data.data(), batch.edge_data.size()); + CUresult result = cuGraphAddDependencies(graph, batch.from_nodes.data(), batch.to_nodes.data(), + batch.edge_data.data(), batch.edge_data.size()); #else CUresult result = cuGraphAddDependencies_v2(graph, batch.from_nodes.data(), batch.to_nodes.data(), diff --git a/include/CUDAGraphInternal.h b/include/CUDAGraphInternal.h index 94c7470b..6b0d1a0c 100644 --- a/include/CUDAGraphInternal.h +++ b/include/CUDAGraphInternal.h @@ -8,8 +8,7 @@ namespace foundry { -CUresult add_graph_dependencies(CUgraph graph, - const std::vector& from_nodes, +CUresult add_graph_dependencies(CUgraph graph, const std::vector& from_nodes, const std::vector& to_nodes, const std::vector& edge_data); From 550b9d55c2f392ace0c6a31fb163b813ed7f4feb Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Fri, 24 Jul 2026 11:05:33 +0000 Subject: [PATCH 33/59] test: cover SGLang LOAD logits shard pre-copy Signed-off-by: Cursor Agent Co-authored-by: Rahul Chalamala --- tests/test_deepep_transport.py | 162 +++++++++++++++++++++++++++++++++ 1 file changed, 162 insertions(+) diff --git a/tests/test_deepep_transport.py b/tests/test_deepep_transport.py index 216943cc..be149529 100644 --- a/tests/test_deepep_transport.py +++ b/tests/test_deepep_transport.py @@ -8,6 +8,7 @@ from types import SimpleNamespace import pytest +import torch from foundry.integration.deepep_transport import ( DeepEPTransport, configure_deepep_buffer, @@ -46,6 +47,34 @@ def _install_fake_sglang_cuda_graph_runner(monkeypatch, runner_cls): ) +def _install_fake_sglang_attention_tp_gather(monkeypatch, gather, group): + sglang_module = types.ModuleType("sglang") + srt_module = types.ModuleType("sglang.srt") + layers_module = types.ModuleType("sglang.srt.layers") + dp_attention_module = types.ModuleType("sglang.srt.layers.dp_attention") + logits_processor_module = types.ModuleType("sglang.srt.layers.logits_processor") + dp_attention_module.get_attention_tp_group = lambda: group + logits_processor_module.attn_tp_all_gather_into_tensor = gather + layers_module.dp_attention = dp_attention_module + layers_module.logits_processor = logits_processor_module + srt_module.layers = layers_module + sglang_module.srt = srt_module + monkeypatch.setitem(sys.modules, "sglang", sglang_module) + monkeypatch.setitem(sys.modules, "sglang.srt", srt_module) + monkeypatch.setitem(sys.modules, "sglang.srt.layers", layers_module) + monkeypatch.setitem( + sys.modules, + "sglang.srt.layers.dp_attention", + dp_attention_module, + ) + monkeypatch.setitem( + sys.modules, + "sglang.srt.layers.logits_processor", + logits_processor_module, + ) + return logits_processor_module + + @pytest.mark.parametrize( ("raw", "expected"), [ @@ -179,6 +208,139 @@ def capture_one_batch_size(self, bs, forward, stream_idx=None): assert runner.attn_backend.init_forward_metadata_out_graph is original_out_graph +def test_sglang_load_precopies_local_attention_tp_logits_before_gather(monkeypatch): + group = SimpleNamespace(rank_in_group=1, world_size=3) + calls = [] + delegate_result = object() + + def fake_gather(global_logits, logits): + calls.append( + { + "global_logits": global_logits, + "logits": logits, + "local_slice_before_delegate": global_logits[1].clone(), + } + ) + global_logits[0].fill_(10) + global_logits[2].fill_(30) + return delegate_result + + logits_processor = _install_fake_sglang_attention_tp_gather( + monkeypatch, + fake_gather, + group, + ) + monkeypatch.setattr( + sglang_hooks, + "get_graph_extension_mode", + lambda: sglang_hooks.CUDAGraphExtensionMode.LOAD, + ) + logits = torch.tensor([[1.0, 2.0], [3.0, 4.0]]) + global_logits = torch.full((3, 2, 2), -1.0) + + sglang_hooks._patch_attn_tp_logits_gather() + result = logits_processor.attn_tp_all_gather_into_tensor(global_logits, logits) + + assert result is delegate_result + assert len(calls) == 1 + assert calls[0]["global_logits"] is global_logits + assert calls[0]["logits"] is logits + torch.testing.assert_close(calls[0]["local_slice_before_delegate"], logits) + torch.testing.assert_close(global_logits[1], logits) + torch.testing.assert_close(global_logits[0], torch.full((2, 2), 10.0)) + torch.testing.assert_close(global_logits[2], torch.full((2, 2), 30.0)) + + +@pytest.mark.parametrize( + "mode", + [ + sglang_hooks.CUDAGraphExtensionMode.NONE, + sglang_hooks.CUDAGraphExtensionMode.SAVE, + ], +) +def test_sglang_non_load_modes_leave_attention_tp_gather_unchanged(monkeypatch, mode): + group = SimpleNamespace(rank_in_group=1, world_size=3) + delegate_result = object() + + def fake_gather(global_logits, logits): + return delegate_result + + logits_processor = _install_fake_sglang_attention_tp_gather( + monkeypatch, + fake_gather, + group, + ) + monkeypatch.setattr(sglang_hooks, "get_graph_extension_mode", lambda: mode) + + sglang_hooks._patch_attn_tp_logits_gather() + + assert logits_processor.attn_tp_all_gather_into_tensor is fake_gather + + +def test_sglang_load_attention_tp_gather_patch_is_idempotent(monkeypatch): + group = SimpleNamespace(rank_in_group=0, world_size=2) + calls = [] + + def fake_gather(global_logits, logits): + calls.append((global_logits, logits)) + return "delegate-result" + + logits_processor = _install_fake_sglang_attention_tp_gather( + monkeypatch, + fake_gather, + group, + ) + monkeypatch.setattr( + sglang_hooks, + "get_graph_extension_mode", + lambda: sglang_hooks.CUDAGraphExtensionMode.LOAD, + ) + + sglang_hooks._patch_attn_tp_logits_gather() + installed = logits_processor.attn_tp_all_gather_into_tensor + sglang_hooks._patch_attn_tp_logits_gather() + + assert logits_processor.attn_tp_all_gather_into_tensor is installed + assert installed(torch.empty((2, 1, 2)), torch.ones((1, 2))) == "delegate-result" + assert len(calls) == 1 + + +def test_sglang_load_attention_tp_gather_rejects_invalid_rank_and_shape(monkeypatch): + group = SimpleNamespace(rank_in_group=1, world_size=3) + calls = [] + + def fake_gather(global_logits, logits): + calls.append((global_logits, logits)) + + logits_processor = _install_fake_sglang_attention_tp_gather( + monkeypatch, + fake_gather, + group, + ) + monkeypatch.setattr( + sglang_hooks, + "get_graph_extension_mode", + lambda: sglang_hooks.CUDAGraphExtensionMode.LOAD, + ) + sglang_hooks._patch_attn_tp_logits_gather() + gather = logits_processor.attn_tp_all_gather_into_tensor + + with pytest.raises( + RuntimeError, + match=r"Foundry SGLang LOAD attention-TP gather.*output shape", + ): + gather(torch.empty((3, 2, 3)), torch.empty((2, 2))) + + group.rank_in_group = 3 + with pytest.raises( + RuntimeError, + match=r"Foundry SGLang LOAD attention-TP gather.*rank_in_group=3.*world_size=3", + ): + gather(torch.empty((3, 2, 2)), torch.empty((2, 2))) + + assert calls == [] + + def test_sglang_warmup_state_round_trips_deepep_profile(tmp_path, monkeypatch): config = sglang_config.CUDAGraphExtensionConfig(workspace_root=str(tmp_path)) monkeypatch.setattr(sglang_config, "_config", config) From 7c8323ac00ba5eb0bc5e79297803e0063afda42d Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Fri, 24 Jul 2026 11:14:32 +0000 Subject: [PATCH 34/59] fix: pre-copy SGLang LOAD local logits shard Signed-off-by: Cursor Agent Co-authored-by: Rahul Chalamala --- python/foundry/integration/sglang/hooks.py | 59 ++++++++++++++++++++++ 1 file changed, 59 insertions(+) diff --git a/python/foundry/integration/sglang/hooks.py b/python/foundry/integration/sglang/hooks.py index f63cc1ff..180e7b3e 100644 --- a/python/foundry/integration/sglang/hooks.py +++ b/python/foundry/integration/sglang/hooks.py @@ -97,6 +97,7 @@ def install_hooks(server_args) -> None: _patch_init_memory_pool() _patch_load_model() _patch_kernel_warmup() + _patch_attn_tp_logits_gather() _patch_cuda_graph_capture() _patch_spawn_sites() _patch_deepep() @@ -105,6 +106,64 @@ def install_hooks(server_args) -> None: logger.info("[Foundry] SGLang hooks installed") +def _patch_attn_tp_logits_gather() -> None: + """Pre-copy LOAD's local logits shard before the attention-TP gather.""" + if get_graph_extension_mode() != CUDAGraphExtensionMode.LOAD: + return + + # SGLang is optional; keep its imports at the hook-install boundary. + from sglang.srt.layers import dp_attention, logits_processor + + original = logits_processor.attn_tp_all_gather_into_tensor + if getattr(original, "_foundry_load_local_shard_precopy", False): + return + + @functools.wraps(original) + def patched(global_logits, logits): + group = dp_attention.get_attention_tp_group() + try: + rank = int(group.rank_in_group) + world_size = int(group.world_size) + except (AttributeError, TypeError, ValueError) as exc: + raise RuntimeError( + "Foundry SGLang LOAD attention-TP gather requires " + "get_attention_tp_group() to expose integer rank_in_group " + "and world_size" + ) from exc + + if world_size <= 0 or rank < 0 or rank >= world_size: + raise RuntimeError( + "Foundry SGLang LOAD attention-TP gather received invalid " + f"rank bounds: rank_in_group={rank}, world_size={world_size}" + ) + + try: + output_shape = tuple(int(size) for size in global_logits.shape) + input_shape = tuple(int(size) for size in logits.shape) + except (AttributeError, TypeError, ValueError) as exc: + raise RuntimeError( + "Foundry SGLang LOAD attention-TP gather requires tensor " + "inputs with concrete shapes" + ) from exc + + expected_output_shape = (world_size, *input_shape) + if output_shape != expected_output_shape: + raise RuntimeError( + "Foundry SGLang LOAD attention-TP gather cannot derive a safe " + "local destination slice: " + f"output shape={output_shape}, input shape={input_shape}, " + f"rank_in_group={rank}, world_size={world_size}, " + f"expected output shape={expected_output_shape}" + ) + + global_logits[rank].copy_(logits) + return original(global_logits, logits) + + patched._foundry_load_local_shard_precopy = True + logits_processor.attn_tp_all_gather_into_tensor = patched + logger.info("[Foundry] SGLang LOAD attention-TP logits pre-copy installed") + + def _patch_deepep() -> None: """Apply Foundry's configured transport to SGLang's DeepEP buffer.""" global _DEEPEP_PATCHED From 5cf1906c6f96eb880d8716e87dec436d1a63f268 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Fri, 24 Jul 2026 12:15:44 +0000 Subject: [PATCH 35/59] test: cover SGLang regular TP logits shard pre-copy Signed-off-by: Cursor Agent Co-authored-by: Rahul Chalamala --- tests/test_deepep_transport.py | 83 ++++++++++++++++++++++++++++++++++ 1 file changed, 83 insertions(+) diff --git a/tests/test_deepep_transport.py b/tests/test_deepep_transport.py index be149529..cd3f1922 100644 --- a/tests/test_deepep_transport.py +++ b/tests/test_deepep_transport.py @@ -75,6 +75,31 @@ def _install_fake_sglang_attention_tp_gather(monkeypatch, gather, group): return logits_processor_module +def _install_fake_sglang_regular_tp_gather(monkeypatch, gather, group): + logits_processor_module = _install_fake_sglang_attention_tp_gather( + monkeypatch, + lambda output, input_: group._all_gather_into_tensor(output, input_), + group, + ) + distributed_module = types.ModuleType("sglang.srt.distributed") + parallel_state_module = types.ModuleType("sglang.srt.distributed.parallel_state") + parallel_state_module.GroupCoordinator = type(group) + distributed_module.parallel_state = parallel_state_module + sys.modules["sglang.srt"].distributed = distributed_module + logits_processor_module.tensor_model_parallel_all_gather = gather + monkeypatch.setitem( + sys.modules, + "sglang.srt.distributed", + distributed_module, + ) + monkeypatch.setitem( + sys.modules, + "sglang.srt.distributed.parallel_state", + parallel_state_module, + ) + return logits_processor_module + + @pytest.mark.parametrize( ("raw", "expected"), [ @@ -251,6 +276,64 @@ def fake_gather(global_logits, logits): torch.testing.assert_close(global_logits[2], torch.full((2, 2), 30.0)) +def test_sglang_load_precopies_local_regular_tp_logits_before_gather(monkeypatch): + calls = [] + gather_result = object() + lower_result = object() + + class GroupCoordinator: + rank_in_group = 1 + world_size = 3 + + def _all_gather_into_tensor(self, output, input_): + output_shards = output.reshape(self.world_size, *input_.shape) + calls.append( + { + "output": output, + "input": input_, + "local_slice_before_delegate": output_shards[1].clone(), + } + ) + output_shards[0].fill_(10) + output_shards[2].fill_(30) + return lower_result + + group = GroupCoordinator() + + def fake_gather(logits, dim=-1): + assert dim == -1 + output = torch.full( + (group.world_size * logits.shape[0], *logits.shape[1:]), + -1.0, + ) + assert group._all_gather_into_tensor(output, logits) is lower_result + return gather_result + + logits_processor = _install_fake_sglang_regular_tp_gather( + monkeypatch, + fake_gather, + group, + ) + monkeypatch.setattr( + sglang_hooks, + "get_graph_extension_mode", + lambda: sglang_hooks.CUDAGraphExtensionMode.LOAD, + ) + logits = torch.tensor([[1.0, 2.0], [3.0, 4.0]]) + + sglang_hooks._patch_attn_tp_logits_gather() + result = logits_processor.tensor_model_parallel_all_gather(logits) + + assert result is gather_result + assert len(calls) == 1 + assert calls[0]["input"] is logits + torch.testing.assert_close(calls[0]["local_slice_before_delegate"], logits) + output_shards = calls[0]["output"].reshape(group.world_size, *logits.shape) + torch.testing.assert_close(output_shards[1], logits) + torch.testing.assert_close(output_shards[0], torch.full((2, 2), 10.0)) + torch.testing.assert_close(output_shards[2], torch.full((2, 2), 30.0)) + + @pytest.mark.parametrize( "mode", [ From c39b59506a9c025e358ac11e1f33de2fb095825c Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Fri, 24 Jul 2026 12:32:11 +0000 Subject: [PATCH 36/59] test: complete regular TP pre-copy coverage Signed-off-by: Cursor Agent Co-authored-by: Rahul Chalamala --- tests/test_deepep_transport.py | 84 ++++++++++++++++++++++++++++++++++ 1 file changed, 84 insertions(+) diff --git a/tests/test_deepep_transport.py b/tests/test_deepep_transport.py index cd3f1922..741f49f2 100644 --- a/tests/test_deepep_transport.py +++ b/tests/test_deepep_transport.py @@ -322,6 +322,12 @@ def fake_gather(logits, dim=-1): logits = torch.tensor([[1.0, 2.0], [3.0, 4.0]]) sglang_hooks._patch_attn_tp_logits_gather() + installed_gather = logits_processor.tensor_model_parallel_all_gather + installed_lower_gather = GroupCoordinator._all_gather_into_tensor + sglang_hooks._patch_attn_tp_logits_gather() + + assert logits_processor.tensor_model_parallel_all_gather is installed_gather + assert GroupCoordinator._all_gather_into_tensor is installed_lower_gather result = logits_processor.tensor_model_parallel_all_gather(logits) assert result is gather_result @@ -334,6 +340,84 @@ def fake_gather(logits, dim=-1): torch.testing.assert_close(output_shards[2], torch.full((2, 2), 30.0)) +@pytest.mark.parametrize( + "mode", + [ + sglang_hooks.CUDAGraphExtensionMode.NONE, + sglang_hooks.CUDAGraphExtensionMode.SAVE, + ], +) +def test_sglang_non_load_modes_leave_regular_tp_gather_unchanged(monkeypatch, mode): + class GroupCoordinator: + rank_in_group = 0 + world_size = 2 + + def _all_gather_into_tensor(self, output, input_): + return None + + group = GroupCoordinator() + + def fake_gather(logits, dim=-1): + return logits + + logits_processor = _install_fake_sglang_regular_tp_gather( + monkeypatch, + fake_gather, + group, + ) + original_lower_gather = GroupCoordinator._all_gather_into_tensor + monkeypatch.setattr(sglang_hooks, "get_graph_extension_mode", lambda: mode) + + sglang_hooks._patch_attn_tp_logits_gather() + + assert logits_processor.tensor_model_parallel_all_gather is fake_gather + assert GroupCoordinator._all_gather_into_tensor is original_lower_gather + + +def test_sglang_load_regular_tp_gather_rejects_invalid_rank_and_shape(monkeypatch): + calls = [] + + class GroupCoordinator: + rank_in_group = 1 + world_size = 3 + + def _all_gather_into_tensor(self, output, input_): + calls.append((output, input_)) + + group = GroupCoordinator() + + def fake_gather(logits, dim=-1): + return logits + + _install_fake_sglang_regular_tp_gather( + monkeypatch, + fake_gather, + group, + ) + monkeypatch.setattr( + sglang_hooks, + "get_graph_extension_mode", + lambda: sglang_hooks.CUDAGraphExtensionMode.LOAD, + ) + sglang_hooks._patch_attn_tp_logits_gather() + gather = GroupCoordinator._all_gather_into_tensor + + with pytest.raises( + RuntimeError, + match=r"Foundry SGLang LOAD regular-TP gather.*output shape", + ): + gather(group, torch.empty((3, 2, 3)), torch.empty((2, 2))) + + group.rank_in_group = 3 + with pytest.raises( + RuntimeError, + match=r"Foundry SGLang LOAD regular-TP gather.*rank_in_group=3.*world_size=3", + ): + gather(group, torch.empty((6, 2)), torch.empty((2, 2))) + + assert calls == [] + + @pytest.mark.parametrize( "mode", [ From b2d4079e57f087ffa7b2c37c4d7c7b0428b97a2f Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Fri, 24 Jul 2026 12:38:52 +0000 Subject: [PATCH 37/59] fix: pre-copy SGLang regular TP local shard Signed-off-by: Cursor Agent Co-authored-by: Rahul Chalamala --- python/foundry/integration/sglang/hooks.py | 61 ++++++++++++++++++++++ 1 file changed, 61 insertions(+) diff --git a/python/foundry/integration/sglang/hooks.py b/python/foundry/integration/sglang/hooks.py index 180e7b3e..ee0b286e 100644 --- a/python/foundry/integration/sglang/hooks.py +++ b/python/foundry/integration/sglang/hooks.py @@ -112,6 +112,7 @@ def _patch_attn_tp_logits_gather() -> None: return # SGLang is optional; keep its imports at the hook-install boundary. + from sglang.srt.distributed.parallel_state import GroupCoordinator from sglang.srt.layers import dp_attention, logits_processor original = logits_processor.attn_tp_all_gather_into_tensor @@ -161,6 +162,66 @@ def patched(global_logits, logits): patched._foundry_load_local_shard_precopy = True logits_processor.attn_tp_all_gather_into_tensor = patched + + original_group_gather = GroupCoordinator._all_gather_into_tensor + if not getattr(original_group_gather, "_foundry_load_local_shard_precopy", False): + + @functools.wraps(original_group_gather) + def patched_group_gather(self, output, input): + try: + rank = int(self.rank_in_group) + world_size = int(self.world_size) + except (AttributeError, TypeError, ValueError) as exc: + raise RuntimeError( + "Foundry SGLang LOAD regular-TP gather requires " + "GroupCoordinator to expose integer rank_in_group and world_size" + ) from exc + + if world_size <= 0 or rank < 0 or rank >= world_size: + raise RuntimeError( + "Foundry SGLang LOAD regular-TP gather received invalid " + f"rank bounds: rank_in_group={rank}, world_size={world_size}" + ) + + try: + output_shape = tuple(int(size) for size in output.shape) + input_shape = tuple(int(size) for size in input.shape) + except (AttributeError, TypeError, ValueError) as exc: + raise RuntimeError( + "Foundry SGLang LOAD regular-TP gather requires tensor " + "inputs with concrete shapes" + ) from exc + + expected_output_shapes = [(world_size, *input_shape)] + if input_shape: + concatenated_shape = ( + world_size * input_shape[0], + *input_shape[1:], + ) + if concatenated_shape != expected_output_shapes[0]: + expected_output_shapes.append(concatenated_shape) + if output_shape not in expected_output_shapes: + raise RuntimeError( + "Foundry SGLang LOAD regular-TP gather cannot derive a safe " + "local destination slice: " + f"output shape={output_shape}, input shape={input_shape}, " + f"rank_in_group={rank}, world_size={world_size}, " + f"expected output shape in {expected_output_shapes}" + ) + + try: + local_output = output.view(world_size, *input_shape)[rank] + except (AttributeError, RuntimeError, TypeError) as exc: + raise RuntimeError( + "Foundry SGLang LOAD regular-TP gather requires a contiguous " + "rank-major output layout" + ) from exc + local_output.copy_(input) + return original_group_gather(self, output, input) + + patched_group_gather._foundry_load_local_shard_precopy = True + GroupCoordinator._all_gather_into_tensor = patched_group_gather + logger.info("[Foundry] SGLang LOAD attention-TP logits pre-copy installed") From 3e32b730f284dbcf8881bc9ae35aa69f2de74dd3 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Fri, 24 Jul 2026 12:43:55 +0000 Subject: [PATCH 38/59] test: complete fake SGLang gather modules Signed-off-by: Cursor Agent Co-authored-by: Rahul Chalamala --- tests/test_deepep_transport.py | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/tests/test_deepep_transport.py b/tests/test_deepep_transport.py index 741f49f2..b3e8420d 100644 --- a/tests/test_deepep_transport.py +++ b/tests/test_deepep_transport.py @@ -48,19 +48,34 @@ def _install_fake_sglang_cuda_graph_runner(monkeypatch, runner_cls): def _install_fake_sglang_attention_tp_gather(monkeypatch, gather, group): + class GroupCoordinator: + def _all_gather_into_tensor(self, output, input_): + return None + sglang_module = types.ModuleType("sglang") srt_module = types.ModuleType("sglang.srt") + distributed_module = types.ModuleType("sglang.srt.distributed") + parallel_state_module = types.ModuleType("sglang.srt.distributed.parallel_state") layers_module = types.ModuleType("sglang.srt.layers") dp_attention_module = types.ModuleType("sglang.srt.layers.dp_attention") logits_processor_module = types.ModuleType("sglang.srt.layers.logits_processor") + parallel_state_module.GroupCoordinator = GroupCoordinator + distributed_module.parallel_state = parallel_state_module dp_attention_module.get_attention_tp_group = lambda: group logits_processor_module.attn_tp_all_gather_into_tensor = gather layers_module.dp_attention = dp_attention_module layers_module.logits_processor = logits_processor_module + srt_module.distributed = distributed_module srt_module.layers = layers_module sglang_module.srt = srt_module monkeypatch.setitem(sys.modules, "sglang", sglang_module) monkeypatch.setitem(sys.modules, "sglang.srt", srt_module) + monkeypatch.setitem(sys.modules, "sglang.srt.distributed", distributed_module) + monkeypatch.setitem( + sys.modules, + "sglang.srt.distributed.parallel_state", + parallel_state_module, + ) monkeypatch.setitem(sys.modules, "sglang.srt.layers", layers_module) monkeypatch.setitem( sys.modules, From 67e2db30e8e96b794504a513e5dc69f081e538e8 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Fri, 24 Jul 2026 13:45:27 +0000 Subject: [PATCH 39/59] test: cover optional SGLang gather APIs Signed-off-by: Cursor Agent Co-authored-by: Rahul Chalamala --- tests/test_deepep_transport.py | 178 +++++++++++++++++++++++++++++++++ 1 file changed, 178 insertions(+) diff --git a/tests/test_deepep_transport.py b/tests/test_deepep_transport.py index b3e8420d..2d5a730c 100644 --- a/tests/test_deepep_transport.py +++ b/tests/test_deepep_transport.py @@ -355,6 +355,184 @@ def fake_gather(logits, dim=-1): torch.testing.assert_close(output_shards[2], torch.full((2, 2), 30.0)) +def test_sglang_load_patches_regular_tp_when_attention_helper_is_absent(monkeypatch): + calls = [] + delegate_result = object() + + class GroupCoordinator: + rank_in_group = 1 + world_size = 2 + + def _all_gather_into_tensor(self, output, input_): + calls.append((output, input_)) + return delegate_result + + group = GroupCoordinator() + logits_processor = _install_fake_sglang_regular_tp_gather( + monkeypatch, + lambda logits, dim=-1: logits, + group, + ) + monkeypatch.delattr(logits_processor, "attn_tp_all_gather_into_tensor") + original_lower_gather = GroupCoordinator._all_gather_into_tensor + monkeypatch.setattr( + sglang_hooks, + "get_graph_extension_mode", + lambda: sglang_hooks.CUDAGraphExtensionMode.LOAD, + ) + + sglang_hooks._patch_attn_tp_logits_gather() + + assert not hasattr(logits_processor, "attn_tp_all_gather_into_tensor") + assert GroupCoordinator._all_gather_into_tensor is not original_lower_gather + logits = torch.tensor([[1.0, 2.0]]) + output = torch.full((2, 1, 2), -1.0) + assert group._all_gather_into_tensor(output, logits) is delegate_result + torch.testing.assert_close(output[1], logits) + assert len(calls) == 1 + assert calls[0][0] is output + assert calls[0][1] is logits + + +def test_sglang_load_patches_attention_tp_when_group_coordinator_is_absent(monkeypatch): + group = SimpleNamespace(rank_in_group=1, world_size=2) + delegate_result = object() + + def fake_gather(global_logits, logits): + return delegate_result + + logits_processor = _install_fake_sglang_attention_tp_gather( + monkeypatch, + fake_gather, + group, + ) + parallel_state = sys.modules["sglang.srt.distributed.parallel_state"] + monkeypatch.delattr(parallel_state, "GroupCoordinator") + monkeypatch.setattr( + sglang_hooks, + "get_graph_extension_mode", + lambda: sglang_hooks.CUDAGraphExtensionMode.LOAD, + ) + + sglang_hooks._patch_attn_tp_logits_gather() + + installed = logits_processor.attn_tp_all_gather_into_tensor + assert installed is not fake_gather + logits = torch.tensor([[1.0, 2.0]]) + global_logits = torch.full((2, 1, 2), -1.0) + assert installed(global_logits, logits) is delegate_result + torch.testing.assert_close(global_logits[1], logits) + + +def test_sglang_load_skips_attention_tp_when_group_accessor_is_absent(monkeypatch): + calls = [] + + class GroupCoordinator: + rank_in_group = 0 + world_size = 2 + + def _all_gather_into_tensor(self, output, input_): + calls.append((output, input_)) + return None + + group = GroupCoordinator() + logits_processor = _install_fake_sglang_regular_tp_gather( + monkeypatch, + lambda logits, dim=-1: logits, + group, + ) + dp_attention = sys.modules["sglang.srt.layers.dp_attention"] + monkeypatch.delattr(dp_attention, "get_attention_tp_group") + original_attention_gather = logits_processor.attn_tp_all_gather_into_tensor + original_lower_gather = GroupCoordinator._all_gather_into_tensor + monkeypatch.setattr( + sglang_hooks, + "get_graph_extension_mode", + lambda: sglang_hooks.CUDAGraphExtensionMode.LOAD, + ) + + sglang_hooks._patch_attn_tp_logits_gather() + + assert logits_processor.attn_tp_all_gather_into_tensor is original_attention_gather + assert GroupCoordinator._all_gather_into_tensor is not original_lower_gather + logits = torch.tensor([[1.0, 2.0]]) + output = torch.full((2, 1, 2), -1.0) + group._all_gather_into_tensor(output, logits) + torch.testing.assert_close(output[0], logits) + assert len(calls) == 1 + assert calls[0][0] is output + assert calls[0][1] is logits + + +@pytest.mark.parametrize("broken_api", ["accessor", "delegate"]) +def test_sglang_load_propagates_present_broken_attention_tp_api( + monkeypatch, + broken_api, +): + message = f"broken attention {broken_api}" + + def broken(*args, **kwargs): + raise AttributeError(message) + + group = SimpleNamespace(rank_in_group=0, world_size=2) + gather = broken if broken_api == "delegate" else lambda global_logits, logits: None + logits_processor = _install_fake_sglang_attention_tp_gather( + monkeypatch, + gather, + group, + ) + if broken_api == "accessor": + dp_attention = sys.modules["sglang.srt.layers.dp_attention"] + monkeypatch.setattr(dp_attention, "get_attention_tp_group", broken) + parallel_state = sys.modules["sglang.srt.distributed.parallel_state"] + monkeypatch.delattr(parallel_state, "GroupCoordinator") + monkeypatch.setattr( + sglang_hooks, + "get_graph_extension_mode", + lambda: sglang_hooks.CUDAGraphExtensionMode.LOAD, + ) + + sglang_hooks._patch_attn_tp_logits_gather() + + with pytest.raises(AttributeError, match=message): + logits_processor.attn_tp_all_gather_into_tensor( + torch.empty((2, 1, 2)), + torch.ones((1, 2)), + ) + + +def test_sglang_load_propagates_present_broken_regular_tp_api(monkeypatch): + message = "broken regular delegate" + + class GroupCoordinator: + rank_in_group = 0 + world_size = 2 + + def _all_gather_into_tensor(self, output, input_): + raise AttributeError(message) + + group = GroupCoordinator() + logits_processor = _install_fake_sglang_regular_tp_gather( + monkeypatch, + lambda logits, dim=-1: logits, + group, + ) + monkeypatch.delattr(logits_processor, "attn_tp_all_gather_into_tensor") + monkeypatch.setattr( + sglang_hooks, + "get_graph_extension_mode", + lambda: sglang_hooks.CUDAGraphExtensionMode.LOAD, + ) + + sglang_hooks._patch_attn_tp_logits_gather() + + with pytest.raises(AttributeError, match=message): + group._all_gather_into_tensor( + torch.empty((2, 1, 2)), + torch.ones((1, 2)), + ) + + @pytest.mark.parametrize( "mode", [ From b4e63456c2adb038567c4385d3c8f34d707b3f2c Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Fri, 24 Jul 2026 13:48:41 +0000 Subject: [PATCH 40/59] fix: tolerate optional SGLang gather APIs Signed-off-by: Cursor Agent Co-authored-by: Rahul Chalamala --- python/foundry/integration/sglang/hooks.py | 109 ++++++++++++--------- 1 file changed, 62 insertions(+), 47 deletions(-) diff --git a/python/foundry/integration/sglang/hooks.py b/python/foundry/integration/sglang/hooks.py index ee0b286e..a6ed68e6 100644 --- a/python/foundry/integration/sglang/hooks.py +++ b/python/foundry/integration/sglang/hooks.py @@ -112,59 +112,74 @@ def _patch_attn_tp_logits_gather() -> None: return # SGLang is optional; keep its imports at the hook-install boundary. - from sglang.srt.distributed.parallel_state import GroupCoordinator + from sglang.srt.distributed import parallel_state from sglang.srt.layers import dp_attention, logits_processor - original = logits_processor.attn_tp_all_gather_into_tensor - if getattr(original, "_foundry_load_local_shard_precopy", False): - return + missing = object() + original = getattr(logits_processor, "attn_tp_all_gather_into_tensor", missing) + get_attention_tp_group = getattr(dp_attention, "get_attention_tp_group", missing) + group_coordinator = getattr(parallel_state, "GroupCoordinator", missing) + original_group_gather = ( + missing + if group_coordinator is missing + else getattr(group_coordinator, "_all_gather_into_tensor", missing) + ) - @functools.wraps(original) - def patched(global_logits, logits): - group = dp_attention.get_attention_tp_group() - try: - rank = int(group.rank_in_group) - world_size = int(group.world_size) - except (AttributeError, TypeError, ValueError) as exc: - raise RuntimeError( - "Foundry SGLang LOAD attention-TP gather requires " - "get_attention_tp_group() to expose integer rank_in_group " - "and world_size" - ) from exc - - if world_size <= 0 or rank < 0 or rank >= world_size: - raise RuntimeError( - "Foundry SGLang LOAD attention-TP gather received invalid " - f"rank bounds: rank_in_group={rank}, world_size={world_size}" - ) + if ( + original is not missing + and get_attention_tp_group is not missing + and not getattr(original, "_foundry_load_local_shard_precopy", False) + ): - try: - output_shape = tuple(int(size) for size in global_logits.shape) - input_shape = tuple(int(size) for size in logits.shape) - except (AttributeError, TypeError, ValueError) as exc: - raise RuntimeError( - "Foundry SGLang LOAD attention-TP gather requires tensor " - "inputs with concrete shapes" - ) from exc - - expected_output_shape = (world_size, *input_shape) - if output_shape != expected_output_shape: - raise RuntimeError( - "Foundry SGLang LOAD attention-TP gather cannot derive a safe " - "local destination slice: " - f"output shape={output_shape}, input shape={input_shape}, " - f"rank_in_group={rank}, world_size={world_size}, " - f"expected output shape={expected_output_shape}" - ) + @functools.wraps(original) + def patched(global_logits, logits): + group = get_attention_tp_group() + try: + rank = int(group.rank_in_group) + world_size = int(group.world_size) + except (AttributeError, TypeError, ValueError) as exc: + raise RuntimeError( + "Foundry SGLang LOAD attention-TP gather requires " + "get_attention_tp_group() to expose integer rank_in_group " + "and world_size" + ) from exc + + if world_size <= 0 or rank < 0 or rank >= world_size: + raise RuntimeError( + "Foundry SGLang LOAD attention-TP gather received invalid " + f"rank bounds: rank_in_group={rank}, world_size={world_size}" + ) + + try: + output_shape = tuple(int(size) for size in global_logits.shape) + input_shape = tuple(int(size) for size in logits.shape) + except (AttributeError, TypeError, ValueError) as exc: + raise RuntimeError( + "Foundry SGLang LOAD attention-TP gather requires tensor " + "inputs with concrete shapes" + ) from exc + + expected_output_shape = (world_size, *input_shape) + if output_shape != expected_output_shape: + raise RuntimeError( + "Foundry SGLang LOAD attention-TP gather cannot derive a safe " + "local destination slice: " + f"output shape={output_shape}, input shape={input_shape}, " + f"rank_in_group={rank}, world_size={world_size}, " + f"expected output shape={expected_output_shape}" + ) - global_logits[rank].copy_(logits) - return original(global_logits, logits) + global_logits[rank].copy_(logits) + return original(global_logits, logits) - patched._foundry_load_local_shard_precopy = True - logits_processor.attn_tp_all_gather_into_tensor = patched + patched._foundry_load_local_shard_precopy = True + logits_processor.attn_tp_all_gather_into_tensor = patched - original_group_gather = GroupCoordinator._all_gather_into_tensor - if not getattr(original_group_gather, "_foundry_load_local_shard_precopy", False): + if original_group_gather is not missing and not getattr( + original_group_gather, + "_foundry_load_local_shard_precopy", + False, + ): @functools.wraps(original_group_gather) def patched_group_gather(self, output, input): @@ -220,7 +235,7 @@ def patched_group_gather(self, output, input): return original_group_gather(self, output, input) patched_group_gather._foundry_load_local_shard_precopy = True - GroupCoordinator._all_gather_into_tensor = patched_group_gather + group_coordinator._all_gather_into_tensor = patched_group_gather logger.info("[Foundry] SGLang LOAD attention-TP logits pre-copy installed") From 397375649a37f77a6f61822bf228e48e5f723d06 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Fri, 24 Jul 2026 14:23:14 +0000 Subject: [PATCH 41/59] test: cover final no-IMEX review regressions Signed-off-by: Cursor Agent Co-authored-by: Rahul Chalamala --- tests/test_deepep_transport.py | 140 +++++++++++++++++++++++++++--- tests/test_experimental_recipe.py | 23 +++++ tests/test_vmm_alloc.py | 119 ++++++++++++++++++++++++- 3 files changed, 268 insertions(+), 14 deletions(-) diff --git a/tests/test_deepep_transport.py b/tests/test_deepep_transport.py index 2d5a730c..c570dbb9 100644 --- a/tests/test_deepep_transport.py +++ b/tests/test_deepep_transport.py @@ -3,6 +3,7 @@ from __future__ import annotations import json +import logging import sys import types from types import SimpleNamespace @@ -291,7 +292,11 @@ def fake_gather(global_logits, logits): torch.testing.assert_close(global_logits[2], torch.full((2, 2), 30.0)) -def test_sglang_load_precopies_local_regular_tp_logits_before_gather(monkeypatch): +@pytest.mark.parametrize("output_layout", ["rank_major", "concatenated"]) +def test_sglang_load_precopies_local_regular_tp_logits_before_gather( + monkeypatch, + output_layout, +): calls = [] gather_result = object() lower_result = object() @@ -317,10 +322,14 @@ def _all_gather_into_tensor(self, output, input_): def fake_gather(logits, dim=-1): assert dim == -1 - output = torch.full( - (group.world_size * logits.shape[0], *logits.shape[1:]), - -1.0, - ) + if output_layout == "rank_major": + output_shape = (group.world_size, *logits.shape) + else: + output_shape = ( + group.world_size * logits.shape[0], + *logits.shape[1:], + ) + output = torch.full(output_shape, -1.0) assert group._all_gather_into_tensor(output, logits) is lower_result return gather_result @@ -464,6 +473,35 @@ def _all_gather_into_tensor(self, output, input_): assert calls[0][1] is logits +def test_sglang_load_logs_only_gather_seams_installed(monkeypatch, caplog): + class GroupCoordinator: + rank_in_group = 0 + world_size = 2 + + def _all_gather_into_tensor(self, output, input_): + return None + + group = GroupCoordinator() + logits_processor = _install_fake_sglang_regular_tp_gather( + monkeypatch, + lambda logits, dim=-1: logits, + group, + ) + monkeypatch.delattr(logits_processor, "attn_tp_all_gather_into_tensor") + monkeypatch.setattr( + sglang_hooks, + "get_graph_extension_mode", + lambda: sglang_hooks.CUDAGraphExtensionMode.LOAD, + ) + + with caplog.at_level(logging.INFO, logger=sglang_hooks.__name__): + sglang_hooks._patch_attn_tp_logits_gather() + + assert caplog.messages[-1] == ( + "[Foundry] SGLang LOAD logits pre-copy patches installed: regular_tp" + ) + + @pytest.mark.parametrize("broken_api", ["accessor", "delegate"]) def test_sglang_load_propagates_present_broken_attention_tp_api( monkeypatch, @@ -567,8 +605,11 @@ def fake_gather(logits, dim=-1): assert GroupCoordinator._all_gather_into_tensor is original_lower_gather -def test_sglang_load_regular_tp_gather_rejects_invalid_rank_and_shape(monkeypatch): +def test_sglang_load_regular_tp_gather_delegates_unsupported_shape_unchanged( + monkeypatch, +): calls = [] + delegate_result = object() class GroupCoordinator: rank_in_group = 1 @@ -576,6 +617,7 @@ class GroupCoordinator: def _all_gather_into_tensor(self, output, input_): calls.append((output, input_)) + return delegate_result group = GroupCoordinator() @@ -594,14 +636,88 @@ def fake_gather(logits, dim=-1): ) sglang_hooks._patch_attn_tp_logits_gather() gather = GroupCoordinator._all_gather_into_tensor + output = torch.full((3, 2, 3), -1.0) + input_ = torch.ones((2, 2)) - with pytest.raises( - RuntimeError, - match=r"Foundry SGLang LOAD regular-TP gather.*output shape", - ): - gather(group, torch.empty((3, 2, 3)), torch.empty((2, 2))) + result = gather(group, output, input_) + + assert result is delegate_result + assert len(calls) == 1 + assert calls[0][0] is output + assert calls[0][1] is input_ + torch.testing.assert_close(output, torch.full((3, 2, 3), -1.0)) + + +def test_sglang_load_regular_tp_gather_delegates_unviewable_layout_unchanged( + monkeypatch, +): + calls = [] + delegate_result = object() + + class Input: + shape = (2, 2) + + class Output: + shape = (6, 2) + + def view(self, *shape): + raise RuntimeError(f"unsupported layout for {shape}") + + class GroupCoordinator: + rank_in_group = 1 + world_size = 3 + + def _all_gather_into_tensor(self, output, input_): + calls.append((output, input_)) + return delegate_result + + group = GroupCoordinator() + _install_fake_sglang_regular_tp_gather( + monkeypatch, + lambda logits, dim=-1: logits, + group, + ) + monkeypatch.setattr( + sglang_hooks, + "get_graph_extension_mode", + lambda: sglang_hooks.CUDAGraphExtensionMode.LOAD, + ) + sglang_hooks._patch_attn_tp_logits_gather() + output = Output() + input_ = Input() + + result = group._all_gather_into_tensor(output, input_) + + assert result is delegate_result + assert len(calls) == 1 + assert calls[0][0] is output + assert calls[0][1] is input_ + + +def test_sglang_load_regular_tp_gather_rejects_invalid_rank(monkeypatch): + calls = [] + + class GroupCoordinator: + rank_in_group = 3 + world_size = 3 + + def _all_gather_into_tensor(self, output, input_): + calls.append((output, input_)) + + group = GroupCoordinator() + _install_fake_sglang_regular_tp_gather( + monkeypatch, + lambda logits, dim=-1: logits, + group, + ) + monkeypatch.setattr( + sglang_hooks, + "get_graph_extension_mode", + lambda: sglang_hooks.CUDAGraphExtensionMode.LOAD, + ) + sglang_hooks._patch_attn_tp_logits_gather() + gather = GroupCoordinator._all_gather_into_tensor - group.rank_in_group = 3 with pytest.raises( RuntimeError, match=r"Foundry SGLang LOAD regular-TP gather.*rank_in_group=3.*world_size=3", diff --git a/tests/test_experimental_recipe.py b/tests/test_experimental_recipe.py index 893bd117..ae249dac 100644 --- a/tests/test_experimental_recipe.py +++ b/tests/test_experimental_recipe.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the Foundry project """Static contracts for the experimental no-IMEX DeepEP recipes.""" +import subprocess from pathlib import Path import tomllib @@ -56,6 +57,28 @@ def test_experimental_scripts_do_not_use_environment_transport_switch() -> None: assert "FOUNDRY_DEEPEP_NVL_IPC" not in path.read_text(), filename +def test_experimental_scripts_are_strict_and_syntax_valid() -> None: + for filename in SERVE_SCRIPTS: + path = EXPERIMENTAL_DIR / filename + text = path.read_text() + + assert text.splitlines()[1] == "set -euo pipefail", filename + subprocess.run(["bash", "-n", str(path)], check=True) + + vllm_script = (EXPERIMENTAL_DIR / SERVE_SCRIPTS[0]).read_text() + assert vllm_script.count("export VLLM_USE_V2_MODEL_RUNNER=0") == 1 + + +def test_experimental_scripts_pin_validated_deterministic_defaults() -> None: + vllm_script = (EXPERIMENTAL_DIR / SERVE_SCRIPTS[0]).read_text() + sglang_script = (EXPERIMENTAL_DIR / SERVE_SCRIPTS[1]).read_text() + + assert "--seed 42" in vllm_script + assert "MEM_FRACTION_STATIC=0.65" in sglang_script + assert "--random-seed 42" in sglang_script + assert "0.65" in (EXPERIMENTAL_DIR / "README.md").read_text() + + def test_experimental_configs_have_portable_paths() -> None: for config_pair in CONFIG_PAIRS.values(): for filename in config_pair: diff --git a/tests/test_vmm_alloc.py b/tests/test_vmm_alloc.py index 7f928926..44c6b789 100644 --- a/tests/test_vmm_alloc.py +++ b/tests/test_vmm_alloc.py @@ -1,7 +1,9 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the Foundry project import ctypes +import importlib.util import os +import struct import subprocess import sys from pathlib import Path @@ -10,10 +12,16 @@ import pytest import torch +CUDA_SUCCESS = 0 +CUDA_ERROR_INVALID_VALUE = 1 +VMM_IPC_MAGIC = 0x564D4D32 -def _get_hook_so_path(): - import importlib.util +class CUipcMemHandle(ctypes.Structure): + _fields_ = [("reserved", ctypes.c_ubyte * 64)] + + +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") @@ -184,6 +192,104 @@ def _run_core_zero_alignment_reserve(): assert address_free(ptr.value, reserve_size) == 0 +def _run_core_malformed_preallocated_ipc_handle(): + torch.cuda.init() + + base_addr = 0x574000000000 + region_size = 1024 * 1024 * 1024 + prealloc_size = 64 * 1024 * 1024 + allocation_size = 4 * 1024 * 1024 + uint64_max = (1 << 64) - 1 + + driver = ctypes.CDLL(None) + mem_alloc = driver.cuMemAlloc_v2 + mem_alloc.argtypes = [ctypes.POINTER(ctypes.c_uint64), ctypes.c_size_t] + mem_alloc.restype = ctypes.c_int + mem_free = driver.cuMemFree_v2 + mem_free.argtypes = [ctypes.c_uint64] + mem_free.restype = ctypes.c_int + ipc_get = driver.cuIpcGetMemHandle + ipc_get.argtypes = [ctypes.POINTER(CUipcMemHandle), ctypes.c_uint64] + ipc_get.restype = ctypes.c_int + ipc_open = driver.cuIpcOpenMemHandle + ipc_open.argtypes = [ + ctypes.POINTER(ctypes.c_uint64), + CUipcMemHandle, + ctypes.c_uint, + ] + ipc_open.restype = ctypes.c_int + ipc_close = driver.cuIpcCloseMemHandle + ipc_close.argtypes = [ctypes.c_uint64] + ipc_close.restype = ctypes.c_int + + padding = ctypes.c_uint64() + target = ctypes.c_uint64() + results = {} + valid_open_result = None + valid_close_result = None + fdry.set_allocation_region(base_addr, region_size) + try: + assert fdry.preallocate_region(prealloc_size) + assert mem_alloc(ctypes.byref(padding), allocation_size) == CUDA_SUCCESS + assert mem_alloc(ctypes.byref(target), allocation_size) == CUDA_SUCCESS + assert target.value > padding.value + + handle = CUipcMemHandle() + assert ipc_get(ctypes.byref(handle), target.value) == CUDA_SUCCESS + blob = bytearray(bytes(handle)) + ( + magic, + _exporter_pid, + original_ptr, + size, + _token, + chunk_base, + chunk_size, + _generation, + ) = struct.unpack_from("=IIQQQQQQ", blob) + assert magic == VMM_IPC_MAGIC + assert original_ptr == target.value + assert size > 0 + assert chunk_base == base_addr + assert chunk_size == prealloc_size + assert original_ptr > chunk_base + + remaining = chunk_base + chunk_size - original_ptr + mutations = { + "zero_original_ptr": (8, 0), + "zero_size": (16, 0), + "original_ptr_before_chunk": (8, chunk_base - 1), + "allocation_exceeds_chunk": (16, remaining + 1), + "allocation_end_overflows": (16, uint64_max - original_ptr + 2), + "chunk_size_without_base": (32, 0), + } + for name, (offset, value) in mutations.items(): + malformed_blob = bytearray(blob) + struct.pack_into("=Q", malformed_blob, offset, value) + malformed = CUipcMemHandle.from_buffer_copy(malformed_blob) + opened = ctypes.c_uint64() + result = ipc_open(ctypes.byref(opened), malformed, 1) + results[name] = result + if result == CUDA_SUCCESS: + assert ipc_close(opened.value) == CUDA_SUCCESS + + opened = ctypes.c_uint64() + valid_open_result = ipc_open(ctypes.byref(opened), handle, 1) + if valid_open_result == CUDA_SUCCESS: + valid_close_result = ipc_close(opened.value) + finally: + if target.value: + assert mem_free(target.value) == CUDA_SUCCESS + if padding.value: + assert mem_free(padding.value) == CUDA_SUCCESS + fdry.free_preallocated_region() + fdry.stop_allocation_region() + + assert results == {name: CUDA_ERROR_INVALID_VALUE for name in mutations} + assert valid_open_result == CUDA_SUCCESS + assert valid_close_result == CUDA_SUCCESS + + def _spawn_with_preload(test_mode): so_path = _get_hook_so_path() env = os.environ.copy() @@ -220,6 +326,12 @@ def test_zero_alignment_reserve_stays_in_region(): _spawn_with_preload("zero-alignment-reserve") +@pytest.mark.filterwarnings("ignore:TORCH_CUDA_ARCH_LIST is not set") +def test_vmm_ipc_rejects_malformed_preallocated_chunk_handle(): + """Malformed VMM2 chunk metadata is rejected by the real hooked IPC path.""" + _spawn_with_preload("malformed-preallocated-ipc-handle") + + if __name__ == "__main__": if "--core" in sys.argv: _run_core() @@ -229,8 +341,11 @@ def test_zero_alignment_reserve_stays_in_region(): _run_core_size_parsing() elif "--zero-alignment-reserve" in sys.argv: _run_core_zero_alignment_reserve() + elif "--malformed-preallocated-ipc-handle" in sys.argv: + _run_core_malformed_preallocated_ipc_handle() else: test_vmm_allocation() test_vmm_allocation_without_region() test_size_parsing() test_zero_alignment_reserve_stays_in_region() + test_vmm_ipc_rejects_malformed_preallocated_chunk_handle() From 7517509af63f5a11937ec866c7a79784311e21de Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Fri, 24 Jul 2026 14:38:21 +0000 Subject: [PATCH 42/59] fix: delegate unsupported SGLang regular-TP gather layouts Retain the validated local-shard pre-copy for recognized rank-major and concatenated LOAD gather layouts, but pass any unrecognized output layout through to the original collective unchanged instead of raising, preserving the delegate's return value and exceptions. Invalid rank bounds still fail fast. Log which optional gather seams were actually installed. Co-authored-by: Rahul Chalamala Signed-off-by: Cursor Agent --- python/foundry/integration/sglang/hooks.py | 84 ++++++++++++---------- 1 file changed, 48 insertions(+), 36 deletions(-) diff --git a/python/foundry/integration/sglang/hooks.py b/python/foundry/integration/sglang/hooks.py index a6ed68e6..c2654222 100644 --- a/python/foundry/integration/sglang/hooks.py +++ b/python/foundry/integration/sglang/hooks.py @@ -107,7 +107,7 @@ def install_hooks(server_args) -> None: def _patch_attn_tp_logits_gather() -> None: - """Pre-copy LOAD's local logits shard before the attention-TP gather.""" + """Pre-copy LOAD's local logits shard before the recognized TP gather seams.""" if get_graph_extension_mode() != CUDAGraphExtensionMode.LOAD: return @@ -125,6 +125,8 @@ def _patch_attn_tp_logits_gather() -> None: else getattr(group_coordinator, "_all_gather_into_tensor", missing) ) + installed: list[str] = [] + if ( original is not missing and get_attention_tp_group is not missing @@ -174,6 +176,7 @@ def patched(global_logits, logits): patched._foundry_load_local_shard_precopy = True logits_processor.attn_tp_all_gather_into_tensor = patched + installed.append("attention_tp") if original_group_gather is not missing and not getattr( original_group_gather, @@ -198,46 +201,55 @@ def patched_group_gather(self, output, input): f"rank bounds: rank_in_group={rank}, world_size={world_size}" ) - try: - output_shape = tuple(int(size) for size in output.shape) - input_shape = tuple(int(size) for size in input.shape) - except (AttributeError, TypeError, ValueError) as exc: - raise RuntimeError( - "Foundry SGLang LOAD regular-TP gather requires tensor " - "inputs with concrete shapes" - ) from exc - - expected_output_shapes = [(world_size, *input_shape)] - if input_shape: - concatenated_shape = ( - world_size * input_shape[0], - *input_shape[1:], - ) - if concatenated_shape != expected_output_shapes[0]: - expected_output_shapes.append(concatenated_shape) - if output_shape not in expected_output_shapes: - raise RuntimeError( - "Foundry SGLang LOAD regular-TP gather cannot derive a safe " - "local destination slice: " - f"output shape={output_shape}, input shape={input_shape}, " - f"rank_in_group={rank}, world_size={world_size}, " - f"expected output shape in {expected_output_shapes}" - ) - - try: - local_output = output.view(world_size, *input_shape)[rank] - except (AttributeError, RuntimeError, TypeError) as exc: - raise RuntimeError( - "Foundry SGLang LOAD regular-TP gather requires a contiguous " - "rank-major output layout" - ) from exc - local_output.copy_(input) + # Only pre-copy the local shard for recognized rank-major layouts. + # Any other caller (a different collective on this shared + # coordinator method, or an unexpected output layout) is passed + # through to the original collective unchanged so the delegate's + # own result and exceptions are preserved. + local_output = _regular_tp_local_shard(output, input, rank, world_size) + if local_output is not None: + local_output.copy_(input) return original_group_gather(self, output, input) patched_group_gather._foundry_load_local_shard_precopy = True group_coordinator._all_gather_into_tensor = patched_group_gather + installed.append("regular_tp") + + if installed: + logger.info( + "[Foundry] SGLang LOAD logits pre-copy patches installed: %s", + ", ".join(installed), + ) + else: + logger.info("[Foundry] SGLang LOAD logits pre-copy found no supported gather seams") + + +def _regular_tp_local_shard(output, input, rank: int, world_size: int): + """Return the local rank slice of ``output`` for a recognized rank-major + gather layout, or ``None`` when the layout is unsupported. - logger.info("[Foundry] SGLang LOAD attention-TP logits pre-copy installed") + Returning ``None`` lets the caller delegate unrecognized layouts to the + original collective unchanged instead of raising, so unrelated callers of + the shared coordinator method are never disturbed. + """ + try: + output_shape = tuple(int(size) for size in output.shape) + input_shape = tuple(int(size) for size in input.shape) + except (AttributeError, TypeError, ValueError): + return None + + expected_output_shapes = [(world_size, *input_shape)] + if input_shape: + concatenated_shape = (world_size * input_shape[0], *input_shape[1:]) + if concatenated_shape != expected_output_shapes[0]: + expected_output_shapes.append(concatenated_shape) + if output_shape not in expected_output_shapes: + return None + + try: + return output.view(world_size, *input_shape)[rank] + except (AttributeError, RuntimeError, TypeError): + return None def _patch_deepep() -> None: From 6cde90481920d78e3f95c1a7d9adca16b678ae7e Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Fri, 24 Jul 2026 14:38:28 +0000 Subject: [PATCH 43/59] fix: validate VMM IPC chunk-import range before interior pointer Reject corrupt VMM2 IPC chunk descriptors in cuIpcOpenMemHandle before deriving the interior pointer: chunk_base/chunk_size must be nonzero together, and the requested carve must be fully contained within the advertised chunk. All range arithmetic is overflow-safe because the descriptor is untrusted. Invalid metadata returns CUDA_ERROR_INVALID_VALUE before any fd fetch, import, reservation, or mapping, so existing cleanup paths are unaffected. Co-authored-by: Rahul Chalamala Signed-off-by: Cursor Agent --- csrc/hook.cpp | 41 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 41 insertions(+) diff --git a/csrc/hook.cpp b/csrc/hook.cpp index 360457b8..61a52103 100644 --- a/csrc/hook.cpp +++ b/csrc/hook.cpp @@ -26,6 +26,7 @@ #include #include #include +#include #include #include #include @@ -3321,6 +3322,46 @@ CUresult cuIpcOpenMemHandle(CUdeviceptr* pdptr, CUipcMemHandle handle, unsigned memcpy(&chunk_size, handle.reserved + 40, sizeof(uint64_t)); memcpy(&chunk_generation, handle.reserved + 48, sizeof(uint64_t)); pid_t exporter_pid = (pid_t)exporter_pid_u32; + + // The descriptor came from a peer process, so validate the chunk geometry + // before deriving any interior pointer. chunk_base/chunk_size are nonzero + // together iff this is a carve from the preallocated chunk; a mismatch is + // corrupt metadata. All range math below is overflow-safe because the + // values are untrusted. Rejecting here happens before any fd fetch, + // import, reservation, or mapping, so there is nothing to clean up. + const bool has_chunk_base = (chunk_base != 0); + const bool has_chunk_size = (chunk_size != 0); + if (has_chunk_base != has_chunk_size) { + fprintf(stderr, + "[HOOK] ERROR: VMM-IPC handle has inconsistent chunk metadata: " + "chunk_base=0x%llx chunk_size=%llu\n", + (unsigned long long)chunk_base, (unsigned long long)chunk_size); + return CUDA_ERROR_INVALID_VALUE; + } + if (has_chunk_base) { + const uint64_t carve_ptr = (uint64_t)original_ptr; + const uint64_t carve_size = (uint64_t)size; + if (carve_ptr == 0 || carve_size == 0) { + fprintf(stderr, + "[HOOK] ERROR: VMM-IPC chunk handle has zero pointer/size: " + "original_ptr=0x%llx size=%llu\n", + (unsigned long long)carve_ptr, (unsigned long long)carve_size); + return CUDA_ERROR_INVALID_VALUE; + } + // The carve must start at or after the chunk base and end at or before + // the chunk end, with neither endpoint sum wrapping around 2^64. + if (carve_ptr < chunk_base || chunk_size > (UINT64_MAX - chunk_base) || + carve_size > (UINT64_MAX - carve_ptr) || + (carve_ptr + carve_size) > (chunk_base + chunk_size)) { + fprintf(stderr, + "[HOOK] ERROR: VMM-IPC carve [0x%llx,+0x%llx) is not contained in " + "chunk [0x%llx,+0x%llx)\n", + (unsigned long long)carve_ptr, (unsigned long long)carve_size, + (unsigned long long)chunk_base, (unsigned long long)chunk_size); + return CUDA_ERROR_INVALID_VALUE; + } + } + // For chunk carves the fd registry on the exporter is keyed by the chunk // base, and what we import/map is the whole chunk. CUdeviceptr fetch_key = (chunk_base != 0) ? (CUdeviceptr)chunk_base : original_ptr; From 750368e5e625dee217c52841bf8012e5ca0ea8fa Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Fri, 24 Jul 2026 14:38:39 +0000 Subject: [PATCH 44/59] chore: harden experimental recipes and pin deterministic defaults Add set -euo pipefail to both experimental serve scripts (guarding the optional second argument to preserve defaults), remove the duplicate VLLM_USE_V2_MODEL_RUNNER export, pin --seed 42 / --random-seed 42, and set the validated SGLang static memory fraction default to 0.65. Document the 0.65 rationale in the recipe README. Co-authored-by: Rahul Chalamala Signed-off-by: Cursor Agent --- recipe/experimental/README.md | 5 ++++- recipe/experimental/serve_qwen3-30ba3b_ipc_ep.sh | 9 +++++---- .../serve_qwen3-30ba3bfp8_sglang_ipc_ep.sh | 13 +++++++++---- 3 files changed, 18 insertions(+), 9 deletions(-) diff --git a/recipe/experimental/README.md b/recipe/experimental/README.md index ef78e433..1cd0da5f 100644 --- a/recipe/experimental/README.md +++ b/recipe/experimental/README.md @@ -78,7 +78,10 @@ bash recipe/experimental/serve_qwen3-30ba3bfp8_sglang_ipc_ep.sh 2 --load ``` Keep the model, EP size, low-latency token cap, chunked-prefill size, and TOML -settings identical between SAVE and LOAD. +settings identical between SAVE and LOAD. Both phases pin `--random-seed 42` and +`--mem-fraction-static 0.65`: at `0.8`, LOAD failed before graph restoration with +`cuMemCreate failed with error 2` while reserving each rank's saved remainder, +whereas `0.65` restores every graph and reaches the endpoint on both ranks. ## Archive parity and failure behavior diff --git a/recipe/experimental/serve_qwen3-30ba3b_ipc_ep.sh b/recipe/experimental/serve_qwen3-30ba3b_ipc_ep.sh index 9d0011c5..da874654 100755 --- a/recipe/experimental/serve_qwen3-30ba3b_ipc_ep.sh +++ b/recipe/experimental/serve_qwen3-30ba3b_ipc_ep.sh @@ -1,4 +1,5 @@ #!/bin/bash +set -euo pipefail # Usage: bash serve_qwen3-30ba3b_ipc_ep.sh [--save|--load] # # EXPERIMENTAL: vLLM DeepEP expert parallel using the intranode NVLink IPC @@ -13,11 +14,11 @@ PORT=12000 GPU_MEMORY_UTILIZATION=0.8 FOUNDRY_ARGS=() -if [[ "$2" == "--save" ]]; then +if [[ "${2:-}" == "--save" ]]; then FOUNDRY_TOML="${SCRIPT_DIR}/foundry_save.toml" -elif [[ "$2" == "--load" ]]; then +elif [[ "${2:-}" == "--load" ]]; then FOUNDRY_TOML="${SCRIPT_DIR}/foundry_load.toml" -elif [[ -n "$2" ]]; then +elif [[ -n "${2:-}" ]]; then echo "Usage: $0 [--save|--load]" exit 1 fi @@ -27,7 +28,6 @@ if [[ -n "${FOUNDRY_TOML:-}" ]]; then # Foundry's VMM allocation region does not carry. export NCCL_CUMEM_ENABLE=0 export NCCL_NVLS_ENABLE=0 - export VLLM_USE_V2_MODEL_RUNNER=0 FOUNDRY_ARGS+=( --compilation-config.graph_extension_config_path "$FOUNDRY_TOML" ) @@ -45,6 +45,7 @@ CUDAGRAPH_CAPTURE_SIZES=($(seq 1 256)) ARGS=( --trust-remote-code + --seed 42 --host "$HOST" --port "$PORT" --tensor-parallel-size 1 diff --git a/recipe/experimental/serve_qwen3-30ba3bfp8_sglang_ipc_ep.sh b/recipe/experimental/serve_qwen3-30ba3bfp8_sglang_ipc_ep.sh index 4c55c2e2..77fa7853 100755 --- a/recipe/experimental/serve_qwen3-30ba3bfp8_sglang_ipc_ep.sh +++ b/recipe/experimental/serve_qwen3-30ba3bfp8_sglang_ipc_ep.sh @@ -1,4 +1,5 @@ #!/bin/bash +set -euo pipefail # Usage: bash serve_qwen3-30ba3bfp8_sglang_ipc_ep.sh [--save|--load] # # EXPERIMENTAL: SGLang DeepEP expert parallel using the intranode NVLink IPC @@ -12,14 +13,17 @@ 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 +# Validated static memory fraction: LOAD reserves each rank's saved remainder +# during graph restoration, which failed at 0.8 with cuMemCreate error 2 but +# succeeds at 0.65 (both ranks restore every graph and reach the endpoint). +MEM_FRACTION_STATIC=0.65 FOUNDRY_ARGS=() -if [[ "$2" == "--save" ]]; then +if [[ "${2:-}" == "--save" ]]; then FOUNDRY_TOML="${SCRIPT_DIR}/sglang_foundry_save.toml" -elif [[ "$2" == "--load" ]]; then +elif [[ "${2:-}" == "--load" ]]; then FOUNDRY_TOML="${SCRIPT_DIR}/sglang_foundry_load.toml" -elif [[ -n "$2" ]]; then +elif [[ -n "${2:-}" ]]; then echo "Usage: $0 [--save|--load]" exit 1 fi @@ -55,4 +59,5 @@ sglang serve \ --chunked-prefill-size 256 \ --attention-backend fa3 \ --cuda-graph-max-bs 128 \ + --random-seed 42 \ "${FOUNDRY_ARGS[@]}" From 743d2f3500252ee74573b17779a3aeae550af0d7 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Fri, 24 Jul 2026 14:38:39 +0000 Subject: [PATCH 45/59] docs: document binary dependency record widths and compatibility Explain the 8-byte legacy and current 16-byte dependency record layouts, how current readers accept both by stored section size, and why FORMAT_VERSION stays at 1. Note that compatibility is one-directional: legacy readers cannot parse the 16-byte edge-metadata records this writer emits. Co-authored-by: Rahul Chalamala Signed-off-by: Cursor Agent --- include/BinaryGraphFormat.h | 22 ++++++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/include/BinaryGraphFormat.h b/include/BinaryGraphFormat.h index 4a002851..0a1dc248 100644 --- a/include/BinaryGraphFormat.h +++ b/include/BinaryGraphFormat.h @@ -35,6 +35,12 @@ namespace foundry { namespace binary_format { static constexpr uint8_t MAGIC[8] = {'C', 'U', 'G', 'R', 'A', 'P', 'H', '\0'}; +// FORMAT_VERSION stays at 1 across the dependency-record widening from 8 to 16 +// bytes (see the Dependency section below). The section table records each +// section's byte size and the header records num_dependencies, so the two +// record widths are unambiguous and a version bump is unnecessary; bumping it +// would only make current readers needlessly reject the many valid +// endpoint-only (8-byte) archives written before edge metadata existed. static constexpr uint32_t FORMAT_VERSION = 1; enum SectionType : uint32_t { @@ -219,6 +225,22 @@ struct BinGenerator { static_assert(sizeof(BinGenerator) == 24, "BinGenerator must be 24 bytes"); // ---- Dependency ---- +// +// The dependency table has two on-disk record layouts, distinguished purely by +// the stored section size (num_dependencies * sizeof(record)): +// +// - LegacyBinDependency (8 bytes): endpoints only, written by readers that +// predate CUDA graph edge metadata. Current readers accept it and default +// the edge-metadata fields (from_port/to_port/type) to zero. +// - BinDependency (16 bytes): the current record, adding the CUgraphEdgeData +// fields plus five reserved zero bytes. This is what the current writer +// always emits. +// +// Compatibility is one-directional: current readers read both widths, but a +// legacy (8-byte-only) reader cannot understand the 16-byte records this writer +// produces and is not expected to. Because FORMAT_VERSION is unchanged, such a +// reader would misparse a current archive; there is no old-reader forward +// compatibility for edge-metadata archives. struct LegacyBinDependency { uint32_t from_id; From 3bc1d1f36ec511126c3947714aa3303262b172d4 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Fri, 24 Jul 2026 14:53:53 +0000 Subject: [PATCH 46/59] test: scope VMM IPC regression to invalid-range rejection The malformed-chunk regression's in-process positive control opened a self-exported handle, which fails at cuMemSetAccess (CUDA IPC is inter-process; a second same-process mapping of a self-imported handle returns CUDA_ERROR_INVALID_VALUE) independent of the new range validation. Drop that unsupported same-process open and keep the six invalid-metadata rejection assertions, which are the regression's purpose. Valid cross-process chunk import plus interior-pointer derivation stays covered by the two-H100 nvl_ipc_prealloc DeepEP matrix case. Co-authored-by: Rahul Chalamala Signed-off-by: Cursor Agent --- tests/test_vmm_alloc.py | 17 ++++++++--------- 1 file changed, 8 insertions(+), 9 deletions(-) diff --git a/tests/test_vmm_alloc.py b/tests/test_vmm_alloc.py index 44c6b789..be87fa0d 100644 --- a/tests/test_vmm_alloc.py +++ b/tests/test_vmm_alloc.py @@ -225,8 +225,6 @@ def _run_core_malformed_preallocated_ipc_handle(): padding = ctypes.c_uint64() target = ctypes.c_uint64() results = {} - valid_open_result = None - valid_close_result = None fdry.set_allocation_region(base_addr, region_size) try: assert fdry.preallocate_region(prealloc_size) @@ -272,11 +270,6 @@ def _run_core_malformed_preallocated_ipc_handle(): results[name] = result if result == CUDA_SUCCESS: assert ipc_close(opened.value) == CUDA_SUCCESS - - opened = ctypes.c_uint64() - valid_open_result = ipc_open(ctypes.byref(opened), handle, 1) - if valid_open_result == CUDA_SUCCESS: - valid_close_result = ipc_close(opened.value) finally: if target.value: assert mem_free(target.value) == CUDA_SUCCESS @@ -285,9 +278,15 @@ def _run_core_malformed_preallocated_ipc_handle(): fdry.free_preallocated_region() fdry.stop_allocation_region() + # Every malformed chunk descriptor is rejected before the interior pointer + # is derived. The accepting path (a valid descriptor that maps the chunk and + # returns base + (ptr - chunk_base)) cannot be exercised in-process: CUDA + # IPC is inter-process, and importing our own exported handle then calling + # cuMemSetAccess on a second same-process mapping returns + # CUDA_ERROR_INVALID_VALUE regardless of this validation. Valid cross-process + # chunk import plus interior-pointer derivation is covered end-to-end by the + # two-H100 nvl_ipc_prealloc DeepEP matrix case. assert results == {name: CUDA_ERROR_INVALID_VALUE for name in mutations} - assert valid_open_result == CUDA_SUCCESS - assert valid_close_result == CUDA_SUCCESS def _spawn_with_preload(test_mode): From af38a6a71030c513ae91fbb4f4c4c9b27b064403 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Fri, 24 Jul 2026 16:34:01 +0000 Subject: [PATCH 47/59] test: cover cached-chunk, stale-generation, and deepep-mode rereview findings Signed-off-by: Cursor Agent Co-authored-by: Rahul Chalamala --- tests/test_experimental_recipe.py | 22 ++++ tests/test_vmm_alloc.py | 203 ++++++++++++++++++++++++++++++ 2 files changed, 225 insertions(+) diff --git a/tests/test_experimental_recipe.py b/tests/test_experimental_recipe.py index ae249dac..d2a26f94 100644 --- a/tests/test_experimental_recipe.py +++ b/tests/test_experimental_recipe.py @@ -111,3 +111,25 @@ def test_sglang_docs_define_required_deepep_api() -> None: assert "Buffer(use_fabric=...)" in text assert "29d31c0" in text + + +def test_sglang_script_accepts_documented_deepep_mode_override() -> None: + script = (EXPERIMENTAL_DIR / SERVE_SCRIPTS[1]).read_text() + + # The DeepEP dispatch mode is an override with the product default preserved. + assert 'DEEPEP_MODE="${DEEPEP_MODE:-low_latency}"' in script + # The mode is threaded through to sglang rather than hardcoded. + assert '--deepep-mode "$DEEPEP_MODE"' in script + assert "--deepep-mode low_latency" not in script + # The override is documented in the script itself. + assert "DEEPEP_MODE" in script.split("sglang serve")[0] + + +def test_sglang_readme_documents_deepep_mode_override() -> None: + readme = (EXPERIMENTAL_DIR / "README.md").read_text() + + assert "DEEPEP_MODE" in readme + # Both the graph-capturable nonzero-NVL mode and the normal-mode limitation + # are documented so the override is used correctly. + assert "auto" in readme + assert "num_nvl_bytes" in readme diff --git a/tests/test_vmm_alloc.py b/tests/test_vmm_alloc.py index be87fa0d..5aba0877 100644 --- a/tests/test_vmm_alloc.py +++ b/tests/test_vmm_alloc.py @@ -289,6 +289,185 @@ def _run_core_malformed_preallocated_ipc_handle(): assert results == {name: CUDA_ERROR_INVALID_VALUE for name in mutations} +def _load_driver_ipc(): + """Bind the driver VMM + IPC entry points used by the cross-process probes.""" + driver = ctypes.CDLL(None) + mem_alloc = driver.cuMemAlloc_v2 + mem_alloc.argtypes = [ctypes.POINTER(ctypes.c_uint64), ctypes.c_size_t] + mem_alloc.restype = ctypes.c_int + mem_free = driver.cuMemFree_v2 + mem_free.argtypes = [ctypes.c_uint64] + mem_free.restype = ctypes.c_int + ipc_get = driver.cuIpcGetMemHandle + ipc_get.argtypes = [ctypes.POINTER(CUipcMemHandle), ctypes.c_uint64] + ipc_get.restype = ctypes.c_int + ipc_open = driver.cuIpcOpenMemHandle + ipc_open.argtypes = [ctypes.POINTER(ctypes.c_uint64), CUipcMemHandle, ctypes.c_uint] + ipc_open.restype = ctypes.c_int + ipc_close = driver.cuIpcCloseMemHandle + ipc_close.argtypes = [ctypes.c_uint64] + ipc_close.restype = ctypes.c_int + return mem_alloc, mem_free, ipc_get, ipc_open, ipc_close + + +# Offsets of the packed VMM2 handle fields (see cuIpcGetMemHandle in csrc/hook.cpp). +_BLOB_ORIGINAL_PTR = 8 +_BLOB_SIZE = 16 +_BLOB_CHUNK_SIZE = 40 +_BLOB_GENERATION = 48 + + +def _open_handle(ipc_open, blob): + """Open a (possibly mutated) VMM2 handle blob, returning (result, opened_ptr).""" + handle = CUipcMemHandle.from_buffer_copy(bytes(blob)) + opened = ctypes.c_uint64() + result = ipc_open(ctypes.byref(opened), handle, 1) + return result, opened + + +def _export_chunk_handle_and_run_importer(child_mode, base_addr): + """Exporter side: preallocate a chunk, carve an allocation, export its VMM2 + chunk handle, and drive a *separate* importer process against that live + export. Returns after asserting the importer child passed.""" + torch.cuda.init() + + region_size = 1024 * 1024 * 1024 + prealloc_size = 64 * 1024 * 1024 + allocation_size = 4 * 1024 * 1024 + + mem_alloc, mem_free, ipc_get, _ipc_open, _ipc_close = _load_driver_ipc() + + padding = ctypes.c_uint64() + target = ctypes.c_uint64() + fdry.set_allocation_region(base_addr, region_size) + try: + assert fdry.preallocate_region(prealloc_size) + assert mem_alloc(ctypes.byref(padding), allocation_size) == CUDA_SUCCESS + assert mem_alloc(ctypes.byref(target), allocation_size) == CUDA_SUCCESS + assert target.value > padding.value + + handle = CUipcMemHandle() + assert ipc_get(ctypes.byref(handle), target.value) == CUDA_SUCCESS + blob = bytes(handle) + (magic, _pid, original_ptr, size, _token, chunk_base, chunk_size, generation) = ( + struct.unpack_from("=IIQQQQQQ", blob) + ) + assert magic == VMM_IPC_MAGIC + assert chunk_base == base_addr + assert chunk_size == prealloc_size + assert original_ptr == target.value + assert size > 0 + assert generation != 0 + + # A distinct process imports the live export. The child inherits our + # LD_PRELOAD (set by _spawn_with_preload) so it runs the same hook. + cmd = [ + sys.executable, + str(Path(__file__).resolve()), + f"--{child_mode}", + blob.hex(), + str(allocation_size), + ] + completed = subprocess.run(cmd, env=os.environ.copy()) + assert completed.returncode == 0, ( + f"importer child {child_mode!r} failed with rc={completed.returncode}" + ) + finally: + if target.value: + assert mem_free(target.value) == CUDA_SUCCESS + if padding.value: + assert mem_free(padding.value) == CUDA_SUCCESS + fdry.free_preallocated_region() + fdry.stop_allocation_region() + + +def _run_core_cached_chunk_export(): + _export_chunk_handle_and_run_importer("cached-chunk-importer", 0x566000000000) + + +def _run_core_cached_chunk_importer(blob_hex, allocation_size): + """Importer side: first import the live chunk (establishing a cache entry), + then present same-key handles with inconsistent chunk geometry. Each must be + rejected by validating against the *actual cached mapping*, not the incoming + handle's own (attacker-controlled) metadata.""" + torch.cuda.init() + _mem_alloc, _mem_free, _ipc_get, ipc_open, ipc_close = _load_driver_ipc() + + blob = bytes.fromhex(blob_hex) + (_magic, _pid, _original_ptr, _size, _token, chunk_base, chunk_size, _generation) = ( + struct.unpack_from("=IIQQQQQQ", blob) + ) + + # 1) Valid cross-process import establishes the cached whole-chunk mapping + # and returns an interior pointer. + valid_result, valid_ptr = _open_handle(ipc_open, blob) + assert valid_result == CUDA_SUCCESS, f"valid cross-process chunk import failed: {valid_result}" + + results = {} + + # 2) Same key (pid, token, chunk_base, generation) but an inflated chunk + # size while the carve stays inside the real mapping. The incoming + # metadata is self-consistent, so only validating the carve against the + # *cached* mapping size can catch the chunk-size mismatch. + size_mismatch = bytearray(blob) + struct.pack_into("=Q", size_mismatch, _BLOB_CHUNK_SIZE, chunk_size * 2) + struct.pack_into("=Q", size_mismatch, _BLOB_ORIGINAL_PTR, chunk_base) + struct.pack_into("=Q", size_mismatch, _BLOB_SIZE, allocation_size) + result, opened = _open_handle(ipc_open, size_mismatch) + if result == CUDA_SUCCESS: + ipc_close(opened.value) + results["chunk_size_mismatch_within_real"] = result + + # 3) Same key with inflated chunk size and a carve that starts at the real + # chunk end — inside the (bogus) incoming chunk but outside the cached + # mapping. Must be rejected before an interior pointer is handed out. + carve_outside = bytearray(blob) + struct.pack_into("=Q", carve_outside, _BLOB_CHUNK_SIZE, chunk_size * 2) + struct.pack_into("=Q", carve_outside, _BLOB_ORIGINAL_PTR, chunk_base + chunk_size) + struct.pack_into("=Q", carve_outside, _BLOB_SIZE, allocation_size) + result, opened = _open_handle(ipc_open, carve_outside) + if result == CUDA_SUCCESS: + ipc_close(opened.value) + results["carve_outside_cached_mapping"] = result + + assert ipc_close(valid_ptr.value) == CUDA_SUCCESS + + assert results == {name: CUDA_ERROR_INVALID_VALUE for name in results}, results + + +def _run_core_stale_generation_export(): + _export_chunk_handle_and_run_importer("stale-generation-importer", 0x578000000000) + + +def _run_core_stale_generation_importer(blob_hex, _allocation_size): + """Importer side: prove the exporter refuses to transfer a chunk descriptor + for a generation that does not match the live export, then confirm the + matching generation is still served (the chunk is genuinely live).""" + torch.cuda.init() + _mem_alloc, _mem_free, _ipc_get, ipc_open, ipc_close = _load_driver_ipc() + + blob = bytes.fromhex(blob_hex) + (_magic, _pid, _original_ptr, _size, _token, _chunk_base, _chunk_size, generation) = ( + struct.unpack_from("=IIQQQQQQ", blob) + ) + + # Stale generation: same base, different generation. The exporter must + # refuse to hand over the fd, so the import fails before mapping. + stale = bytearray(blob) + struct.pack_into("=Q", stale, _BLOB_GENERATION, generation + 1) + stale_result, stale_ptr = _open_handle(ipc_open, stale) + if stale_result == CUDA_SUCCESS: + ipc_close(stale_ptr.value) + assert stale_result == CUDA_ERROR_INVALID_VALUE, ( + f"stale-generation chunk import was not rejected: {stale_result}" + ) + + # Positive control: the live generation is still importable. + live_result, live_ptr = _open_handle(ipc_open, blob) + assert live_result == CUDA_SUCCESS, f"live-generation chunk import failed: {live_result}" + assert ipc_close(live_ptr.value) == CUDA_SUCCESS + + def _spawn_with_preload(test_mode): so_path = _get_hook_so_path() env = os.environ.copy() @@ -331,6 +510,20 @@ def test_vmm_ipc_rejects_malformed_preallocated_chunk_handle(): _spawn_with_preload("malformed-preallocated-ipc-handle") +@pytest.mark.filterwarnings("ignore:TORCH_CUDA_ARCH_LIST is not set") +def test_vmm_ipc_cached_chunk_rejects_inconsistent_metadata(): + """A cache-hit reopen with same-key but inconsistent chunk metadata is + validated against the actual cached mapping and rejected (cross-process).""" + _spawn_with_preload("cached-chunk-export") + + +@pytest.mark.filterwarnings("ignore:TORCH_CUDA_ARCH_LIST is not set") +def test_vmm_ipc_rejects_stale_generation_import(): + """The exporter refuses to serve a chunk descriptor for a stale generation, + while the live generation still imports (cross-process).""" + _spawn_with_preload("stale-generation-export") + + if __name__ == "__main__": if "--core" in sys.argv: _run_core() @@ -342,9 +535,19 @@ def test_vmm_ipc_rejects_malformed_preallocated_chunk_handle(): _run_core_zero_alignment_reserve() elif "--malformed-preallocated-ipc-handle" in sys.argv: _run_core_malformed_preallocated_ipc_handle() + elif "--cached-chunk-export" in sys.argv: + _run_core_cached_chunk_export() + elif "--cached-chunk-importer" in sys.argv: + _run_core_cached_chunk_importer(sys.argv[2], int(sys.argv[3])) + elif "--stale-generation-export" in sys.argv: + _run_core_stale_generation_export() + elif "--stale-generation-importer" in sys.argv: + _run_core_stale_generation_importer(sys.argv[2], int(sys.argv[3])) else: test_vmm_allocation() test_vmm_allocation_without_region() test_size_parsing() test_zero_alignment_reserve_stays_in_region() test_vmm_ipc_rejects_malformed_preallocated_chunk_handle() + test_vmm_ipc_cached_chunk_rejects_inconsistent_metadata() + test_vmm_ipc_rejects_stale_generation_import() From 9f619c1c82b4e256a3839c5a157c0802e5c6029b Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Fri, 24 Jul 2026 16:38:13 +0000 Subject: [PATCH 48/59] fix: validate cached VMM IPC chunk carve against actual mapping Signed-off-by: Cursor Agent Co-authored-by: Rahul Chalamala --- csrc/hook.cpp | 25 +++++++++++++++++++++++++ 1 file changed, 25 insertions(+) diff --git a/csrc/hook.cpp b/csrc/hook.cpp index 61a52103..952e4b7a 100644 --- a/csrc/hook.cpp +++ b/csrc/hook.cpp @@ -3373,6 +3373,31 @@ CUresult cuIpcOpenMemHandle(CUdeviceptr* pdptr, CUipcMemHandle handle, unsigned VmmIpcChunkKey key = std::make_tuple(exporter_pid, token, chunk_base, chunk_generation); auto it = vmm_ipc_chunk_mappings.find(key); if (it != vmm_ipc_chunk_mappings.end()) { + // Cache hit. The incoming handle's chunk_size/carve were only checked + // against the handle's own (untrusted) metadata above. Re-validate the + // request against the ACTUAL cached mapping before deriving an interior + // pointer: the key does not include chunk_size, so a same-key handle + // could otherwise smuggle a larger chunk_size and a carve that lands + // outside the real mapping. Reject before touching the refcount so all + // bookkeeping is preserved on the error path. + const size_t cached_size = it->second.size; + if ((uint64_t)chunk_size != (uint64_t)cached_size) { + fprintf(stderr, + "[HOOK] ERROR: VMM-IPC cache hit chunk_size mismatch: handle=%llu cached=%zu\n", + (unsigned long long)chunk_size, cached_size); + return CUDA_ERROR_INVALID_VALUE; + } + const uint64_t carve_ptr = (uint64_t)original_ptr; + const uint64_t carve_size = (uint64_t)size; + if (carve_ptr < chunk_base || carve_size > (UINT64_MAX - carve_ptr) || + (carve_ptr + carve_size) > (chunk_base + (uint64_t)cached_size)) { + fprintf(stderr, + "[HOOK] ERROR: VMM-IPC cache hit carve [0x%llx,+0x%llx) outside cached mapping " + "[0x%llx,+0x%zx)\n", + (unsigned long long)carve_ptr, (unsigned long long)carve_size, + (unsigned long long)chunk_base, cached_size); + return CUDA_ERROR_INVALID_VALUE; + } it->second.refcount++; CUdeviceptr mapped = it->second.local_base + (original_ptr - (CUdeviceptr)chunk_base); auto sub_it = vmm_ipc_chunk_subptrs.find(mapped); From 7dd065b85c48c797289cb606282432ee94651175 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Fri, 24 Jul 2026 16:38:32 +0000 Subject: [PATCH 49/59] fix: verify chunk generation against live export before fd transfer Signed-off-by: Cursor Agent Co-authored-by: Rahul Chalamala --- csrc/hook.cpp | 100 +++++++++++++++++++++++++++++--------------------- 1 file changed, 58 insertions(+), 42 deletions(-) diff --git a/csrc/hook.cpp b/csrc/hook.cpp index 952e4b7a..ad3b6ad9 100644 --- a/csrc/hook.cpp +++ b/csrc/hook.cpp @@ -2907,11 +2907,17 @@ static constexpr uint32_t VMM_IPC_MAGIC = 0x564D4D32; // "VMM2" // 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; +// exporter VA -> parked export. The allocation handle is kept so VA reuse +// after free/realloc invalidates the cached fd. The generation lets the server +// reject requests carrying a stale chunk generation (free/recreate at one +// base): it is the preallocated-chunk generation for chunk exports and 0 for +// direct (non-chunk) allocations. +struct VmmIpcExportedFd { + CUmemGenericAllocationHandle handle; + int fd; + uint64_t generation; +}; +static boost::unordered::concurrent_flat_map vmm_ipc_exported_fds; static std::mutex vmm_ipc_server_mutex; static int vmm_ipc_listen_fd = -1; // pid that owns the running server thread. fork() copies this .so's state but @@ -2924,10 +2930,15 @@ static pid_t vmm_ipc_server_pid = -1; // process's allocation). static uint64_t vmm_ipc_token = 0; -// Wire format of a fetch request. +// Wire format of a fetch request. Internal to the per-process VMM-IPC socket +// (both peers run the same hook build within a run), so extending it does not +// touch the exported CUipcMemHandle blob layout. generation carries the chunk +// generation the importer expects (0 for direct allocations); the server +// refuses to transfer a descriptor whose live generation differs. struct VmmIpcRequest { uint64_t ptr; uint64_t token; + uint64_t generation; }; static void vmm_ipc_set_timeouts(int fd) { @@ -2998,10 +3009,13 @@ static void vmm_ipc_server_loop(int listen_fd) { 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); }); + // Only serve when the requested generation matches the live export, so a + // stale handle for a freed/recreated chunk at the same base is rejected. + vmm_ipc_exported_fds.visit((CUdeviceptr)req.ptr, + [&](const std::pair& kv) { + if (kv.second.generation == req.generation) + fd = fcntl(kv.second.fd, F_DUPFD_CLOEXEC, 0); + }); } char status = (fd >= 0) ? 0 : 1; struct iovec iov = {&status, 1}; @@ -3037,11 +3051,10 @@ static bool vmm_ipc_ensure_server() { 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; - }); + vmm_ipc_exported_fds.erase_if([](const std::pair& kv) { + close(kv.second.fd); + return true; + }); } // First call in this process, or first after fork (the parent's accept // thread did not survive). Close any inherited listen fd so the parent's @@ -3086,7 +3099,8 @@ static bool vmm_ipc_ensure_server() { // 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) { +static int vmm_ipc_fetch_fd(pid_t exporter_pid, uint64_t token, uint64_t generation, + CUdeviceptr original_ptr) { int s = socket(AF_UNIX, SOCK_STREAM | SOCK_CLOEXEC, 0); if (s < 0) return -1; @@ -3101,7 +3115,7 @@ static int vmm_ipc_fetch_fd(pid_t exporter_pid, uint64_t token, CUdeviceptr orig close(s); return -1; } - VmmIpcRequest req = {(uint64_t)original_ptr, token}; + VmmIpcRequest req = {(uint64_t)original_ptr, token, generation}; if (send(s, &req, sizeof(req), MSG_NOSIGNAL) != (ssize_t)sizeof(req)) { close(s); return -1; @@ -3132,13 +3146,12 @@ static int vmm_ipc_fetch_fd(pid_t exporter_pid, uint64_t token, CUdeviceptr orig // 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; - }); + vmm_ipc_exported_fds.erase_if([dptr](const std::pair& kv) { + if (kv.first != dptr) + return false; + close(kv.second.fd); + return true; + }); } // Hook for cuIpcGetMemHandle - intercept Driver API to support VMM allocations @@ -3219,12 +3232,11 @@ CUresult cuIpcGetMemHandle(CUipcMemHandle* pHandle, CUdeviceptr dptr) { // 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; - }); + vmm_ipc_exported_fds.visit(export_key, + [&](const std::pair& kv) { + if (kv.second.handle == export_handle) + fd = kv.second.fd; + }); if (fd < 0) { int new_fd = -1; CUresult result = @@ -3240,16 +3252,17 @@ CUresult cuIpcGetMemHandle(CUipcMemHandle* pHandle, CUdeviceptr dptr) { // 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) { + std::make_pair(export_key, VmmIpcExportedFd{export_handle, new_fd, chunk_generation}), + [&](std::pair& kv) { // Entry exists: either a racing thread won (same handle - drop - // ours) or it is stale from a freed allocation (replace). - if (kv.second.first == export_handle) { + // ours) or it is stale from a freed allocation (replace, refreshing + // the generation for the new export). + if (kv.second.handle == export_handle) { close(new_fd); - new_fd = kv.second.second; + new_fd = kv.second.fd; } else { - close(kv.second.second); - kv.second = std::make_pair(export_handle, new_fd); + close(kv.second.fd); + kv.second = VmmIpcExportedFd{export_handle, new_fd, chunk_generation}; } }); fd = new_fd; @@ -3427,12 +3440,15 @@ CUresult cuIpcOpenMemHandle(CUdeviceptr* pdptr, CUipcMemHandle handle, unsigned // 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); }); + // Same-process import: mirror the server's generation check so a stale + // generation is rejected here too. + vmm_ipc_exported_fds.visit(fetch_key, + [&](const std::pair& kv) { + if (kv.second.generation == chunk_generation) + fd = fcntl(kv.second.fd, F_DUPFD_CLOEXEC, 0); + }); } else { - fd = vmm_ipc_fetch_fd(exporter_pid, token, fetch_key); + fd = vmm_ipc_fetch_fd(exporter_pid, token, chunk_generation, fetch_key); } if (fd < 0) { fprintf(stderr, "[HOOK] ERROR: VMM-IPC could not obtain fd for ptr=0x%llx from pid %d\n", From 1b59dbe00553ce3c3bae12f2da69733f628de3e0 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Fri, 24 Jul 2026 16:45:52 +0000 Subject: [PATCH 50/59] test: run cross-process VMM IPC regressions across two GPUs Signed-off-by: Cursor Agent Co-authored-by: Rahul Chalamala --- tests/test_vmm_alloc.py | 283 +++++++++++++++++++++++----------------- 1 file changed, 160 insertions(+), 123 deletions(-) diff --git a/tests/test_vmm_alloc.py b/tests/test_vmm_alloc.py index 5aba0877..fda71ac9 100644 --- a/tests/test_vmm_alloc.py +++ b/tests/test_vmm_alloc.py @@ -325,147 +325,177 @@ def _open_handle(ipc_open, blob): return result, opened -def _export_chunk_handle_and_run_importer(child_mode, base_addr): - """Exporter side: preallocate a chunk, carve an allocation, export its VMM2 - chunk handle, and drive a *separate* importer process against that live - export. Returns after asserting the importer child passed.""" - torch.cuda.init() - - region_size = 1024 * 1024 * 1024 - prealloc_size = 64 * 1024 * 1024 - allocation_size = 4 * 1024 * 1024 - - mem_alloc, mem_free, ipc_get, _ipc_open, _ipc_close = _load_driver_ipc() - - padding = ctypes.c_uint64() - target = ctypes.c_uint64() - fdry.set_allocation_region(base_addr, region_size) - try: - assert fdry.preallocate_region(prealloc_size) - assert mem_alloc(ctypes.byref(padding), allocation_size) == CUDA_SUCCESS - assert mem_alloc(ctypes.byref(target), allocation_size) == CUDA_SUCCESS - assert target.value > padding.value - - handle = CUipcMemHandle() - assert ipc_get(ctypes.byref(handle), target.value) == CUDA_SUCCESS - blob = bytes(handle) - (magic, _pid, original_ptr, size, _token, chunk_base, chunk_size, generation) = ( - struct.unpack_from("=IIQQQQQQ", blob) - ) - assert magic == VMM_IPC_MAGIC - assert chunk_base == base_addr - assert chunk_size == prealloc_size - assert original_ptr == target.value - assert size > 0 - assert generation != 0 - - # A distinct process imports the live export. The child inherits our - # LD_PRELOAD (set by _spawn_with_preload) so it runs the same hook. - cmd = [ - sys.executable, - str(Path(__file__).resolve()), - f"--{child_mode}", - blob.hex(), - str(allocation_size), - ] - completed = subprocess.run(cmd, env=os.environ.copy()) - assert completed.returncode == 0, ( - f"importer child {child_mode!r} failed with rc={completed.returncode}" - ) - finally: - if target.value: - assert mem_free(target.value) == CUDA_SUCCESS - if padding.value: - assert mem_free(padding.value) == CUDA_SUCCESS - fdry.free_preallocated_region() - fdry.stop_allocation_region() - - -def _run_core_cached_chunk_export(): - _export_chunk_handle_and_run_importer("cached-chunk-importer", 0x566000000000) - - -def _run_core_cached_chunk_importer(blob_hex, allocation_size): - """Importer side: first import the live chunk (establishing a cache entry), - then present same-key handles with inconsistent chunk geometry. Each must be - rejected by validating against the *actual cached mapping*, not the incoming - handle's own (attacker-controlled) metadata.""" - torch.cuda.init() - _mem_alloc, _mem_free, _ipc_get, ipc_open, ipc_close = _load_driver_ipc() - - blob = bytes.fromhex(blob_hex) - (_magic, _pid, _original_ptr, _size, _token, chunk_base, chunk_size, _generation) = ( - struct.unpack_from("=IIQQQQQQ", blob) +# CUDA VMM IPC is inter-process; a fully successful import (needed to establish +# a cached chunk mapping, and as the stale-generation positive control) must map +# a peer's allocation with cuMemSetAccess, which returns CUDA_ERROR_INVALID_VALUE +# for a *same-GPU* import of an allocation this process/GPU already owns. So the +# real cross-process regressions run rank 0 (exporter) and rank 1 (importer) on +# two distinct GPUs, exactly like the DeepEP nvl_ipc_prealloc matrix case. +_REGRESSION_BASE_ADDR = 0x566000000000 +_REGRESSION_REGION_SIZE = 1024 * 1024 * 1024 +_REGRESSION_PREALLOC_SIZE = 64 * 1024 * 1024 +_REGRESSION_ALLOCATION_SIZE = 4 * 1024 * 1024 + + +def _importer_check_cached_chunk(blob, ipc_open, ipc_close): + """Establish a cached whole-chunk mapping via a valid import, then present + same-key handles with inconsistent chunk geometry. Each reopen must be + validated against the ACTUAL cached mapping (not the incoming handle's own + attacker-controlled metadata) and rejected.""" + (_magic, _pid, _optr, _size, _token, chunk_base, chunk_size, _gen) = struct.unpack_from( + "=IIQQQQQQ", blob ) - # 1) Valid cross-process import establishes the cached whole-chunk mapping - # and returns an interior pointer. valid_result, valid_ptr = _open_handle(ipc_open, blob) - assert valid_result == CUDA_SUCCESS, f"valid cross-process chunk import failed: {valid_result}" + if valid_result != CUDA_SUCCESS: + return False, f"valid cross-process chunk import failed: {valid_result}" results = {} - # 2) Same key (pid, token, chunk_base, generation) but an inflated chunk - # size while the carve stays inside the real mapping. The incoming - # metadata is self-consistent, so only validating the carve against the - # *cached* mapping size can catch the chunk-size mismatch. + # Inflated chunk size while the carve stays inside the real mapping. The + # incoming metadata is self-consistent, so only checking the carve against + # the cached mapping size can catch the chunk-size mismatch. size_mismatch = bytearray(blob) struct.pack_into("=Q", size_mismatch, _BLOB_CHUNK_SIZE, chunk_size * 2) struct.pack_into("=Q", size_mismatch, _BLOB_ORIGINAL_PTR, chunk_base) - struct.pack_into("=Q", size_mismatch, _BLOB_SIZE, allocation_size) + struct.pack_into("=Q", size_mismatch, _BLOB_SIZE, _REGRESSION_ALLOCATION_SIZE) result, opened = _open_handle(ipc_open, size_mismatch) if result == CUDA_SUCCESS: ipc_close(opened.value) results["chunk_size_mismatch_within_real"] = result - # 3) Same key with inflated chunk size and a carve that starts at the real - # chunk end — inside the (bogus) incoming chunk but outside the cached - # mapping. Must be rejected before an interior pointer is handed out. + # Inflated chunk size with a carve starting at the real chunk end: inside + # the bogus incoming chunk but outside the cached mapping. carve_outside = bytearray(blob) struct.pack_into("=Q", carve_outside, _BLOB_CHUNK_SIZE, chunk_size * 2) struct.pack_into("=Q", carve_outside, _BLOB_ORIGINAL_PTR, chunk_base + chunk_size) - struct.pack_into("=Q", carve_outside, _BLOB_SIZE, allocation_size) + struct.pack_into("=Q", carve_outside, _BLOB_SIZE, _REGRESSION_ALLOCATION_SIZE) result, opened = _open_handle(ipc_open, carve_outside) if result == CUDA_SUCCESS: ipc_close(opened.value) results["carve_outside_cached_mapping"] = result - assert ipc_close(valid_ptr.value) == CUDA_SUCCESS - - assert results == {name: CUDA_ERROR_INVALID_VALUE for name in results}, results - - -def _run_core_stale_generation_export(): - _export_chunk_handle_and_run_importer("stale-generation-importer", 0x578000000000) + ipc_close(valid_ptr.value) + for name, code in results.items(): + if code != CUDA_ERROR_INVALID_VALUE: + return False, f"{name} not rejected on cache hit: {code}" + return True, "cached-chunk inconsistent metadata rejected" -def _run_core_stale_generation_importer(blob_hex, _allocation_size): - """Importer side: prove the exporter refuses to transfer a chunk descriptor - for a generation that does not match the live export, then confirm the - matching generation is still served (the chunk is genuinely live).""" - torch.cuda.init() - _mem_alloc, _mem_free, _ipc_get, ipc_open, ipc_close = _load_driver_ipc() - blob = bytes.fromhex(blob_hex) - (_magic, _pid, _original_ptr, _size, _token, _chunk_base, _chunk_size, generation) = ( - struct.unpack_from("=IIQQQQQQ", blob) +def _importer_check_stale_generation(blob, ipc_open, ipc_close): + """Prove the exporter refuses to transfer a chunk descriptor for a stale + generation, while the live generation still imports (the chunk is live).""" + (_magic, _pid, _optr, _size, _token, _chunk_base, _chunk_size, generation) = struct.unpack_from( + "=IIQQQQQQ", blob ) - # Stale generation: same base, different generation. The exporter must - # refuse to hand over the fd, so the import fails before mapping. stale = bytearray(blob) struct.pack_into("=Q", stale, _BLOB_GENERATION, generation + 1) stale_result, stale_ptr = _open_handle(ipc_open, stale) if stale_result == CUDA_SUCCESS: ipc_close(stale_ptr.value) - assert stale_result == CUDA_ERROR_INVALID_VALUE, ( - f"stale-generation chunk import was not rejected: {stale_result}" - ) + return False, "stale-generation chunk import was not rejected" + if stale_result != CUDA_ERROR_INVALID_VALUE: + return False, f"stale-generation import failed with unexpected code: {stale_result}" - # Positive control: the live generation is still importable. live_result, live_ptr = _open_handle(ipc_open, blob) - assert live_result == CUDA_SUCCESS, f"live-generation chunk import failed: {live_result}" - assert ipc_close(live_ptr.value) == CUDA_SUCCESS + if live_result != CUDA_SUCCESS: + return False, f"live-generation chunk import failed: {live_result}" + ipc_close(live_ptr.value) + return True, "stale generation rejected; live generation served" + + +_IMPORTER_CHECKS = { + "cached-chunk": _importer_check_cached_chunk, + "stale-generation": _importer_check_stale_generation, +} + + +def _run_2gpu_regression(scenario, local_rank): + """Two-GPU cross-process regression. Rank 0 exports a live preallocated-chunk + VMM2 handle and keeps its fd server alive; rank 1 (a different GPU) imports + it through the real hook and runs the scenario's checks.""" + import torch.distributed as dist + + # gloo (CPU) so the blob/status handoff never touches the VMM region. + dist.init_process_group("gloo", rank=local_rank, world_size=2) + mem_alloc, mem_free, ipc_get, ipc_open, ipc_close = _load_driver_ipc() + + padding = ctypes.c_uint64() + target = ctypes.c_uint64() + region_active = False + blob_box = [None] + status_box = [None] + try: + torch.cuda.set_device(local_rank) + torch.cuda.init() + + if local_rank == 0: + fdry.set_allocation_region(_REGRESSION_BASE_ADDR, _REGRESSION_REGION_SIZE) + region_active = True + assert fdry.preallocate_region(_REGRESSION_PREALLOC_SIZE) + assert mem_alloc(ctypes.byref(padding), _REGRESSION_ALLOCATION_SIZE) == CUDA_SUCCESS + assert mem_alloc(ctypes.byref(target), _REGRESSION_ALLOCATION_SIZE) == CUDA_SUCCESS + assert target.value > padding.value + handle = CUipcMemHandle() + assert ipc_get(ctypes.byref(handle), target.value) == CUDA_SUCCESS + blob = bytes(handle) + (magic, _pid, original_ptr, size, _token, chunk_base, chunk_size, generation) = ( + struct.unpack_from("=IIQQQQQQ", blob) + ) + assert magic == VMM_IPC_MAGIC + assert chunk_base == _REGRESSION_BASE_ADDR + assert chunk_size == _REGRESSION_PREALLOC_SIZE + assert original_ptr == target.value + assert size > 0 + assert generation != 0 + blob_box[0] = blob + + # Rank 0 hands the live handle to rank 1, keeps its server thread alive + # (blocked receiving the status), then both barrier before teardown. + dist.broadcast_object_list(blob_box, src=0) + if local_rank == 1: + status_box[0] = _IMPORTER_CHECKS[scenario](blob_box[0], ipc_open, ipc_close) + dist.broadcast_object_list(status_box, src=1) + dist.barrier() + finally: + if local_rank == 0: + if target.value: + mem_free(target.value) + if padding.value: + mem_free(padding.value) + if region_active: + fdry.free_preallocated_region() + fdry.stop_allocation_region() + dist.destroy_process_group() + + ok, message = status_box[0] + assert ok, message + + +def _spawn_2gpu_regression(scenario): + so_path = _get_hook_so_path() + base_env = os.environ.copy() + if base_env.get("LD_PRELOAD"): + base_env["LD_PRELOAD"] = f"{so_path}:{base_env['LD_PRELOAD']}" + else: + base_env["LD_PRELOAD"] = so_path + base_env["MASTER_ADDR"] = "127.0.0.1" + base_env["MASTER_PORT"] = str({"cached-chunk": 29753, "stale-generation": 29757}[scenario]) + + procs = [] + for rank in range(2): + env = base_env.copy() + cmd = [ + sys.executable, + str(Path(__file__).resolve()), + f"--2gpu-{scenario}", + str(rank), + ] + procs.append(subprocess.Popen(cmd, env=env)) + returncodes = [proc.wait() for proc in procs] + assert all(code == 0 for code in returncodes), f"regression ranks failed: {returncodes}" def _spawn_with_preload(test_mode): @@ -510,18 +540,31 @@ def test_vmm_ipc_rejects_malformed_preallocated_chunk_handle(): _spawn_with_preload("malformed-preallocated-ipc-handle") -@pytest.mark.filterwarnings("ignore:TORCH_CUDA_ARCH_LIST is not set") +def _cuda_device_count() -> int: + """GPU count without initializing CUDA in this orchestration parent.""" + result = subprocess.run( + ["nvidia-smi", "-L"], + stdout=subprocess.PIPE, + stderr=subprocess.DEVNULL, + text=True, + ) + if result.returncode != 0: + return 0 + return sum(1 for line in result.stdout.splitlines() if line.startswith("GPU ")) + + +@pytest.mark.skipif(_cuda_device_count() < 2, reason="requires 2 GPUs for cross-process IPC") def test_vmm_ipc_cached_chunk_rejects_inconsistent_metadata(): """A cache-hit reopen with same-key but inconsistent chunk metadata is - validated against the actual cached mapping and rejected (cross-process).""" - _spawn_with_preload("cached-chunk-export") + validated against the actual cached mapping and rejected (2-GPU cross-process).""" + _spawn_2gpu_regression("cached-chunk") -@pytest.mark.filterwarnings("ignore:TORCH_CUDA_ARCH_LIST is not set") +@pytest.mark.skipif(_cuda_device_count() < 2, reason="requires 2 GPUs for cross-process IPC") def test_vmm_ipc_rejects_stale_generation_import(): """The exporter refuses to serve a chunk descriptor for a stale generation, - while the live generation still imports (cross-process).""" - _spawn_with_preload("stale-generation-export") + while the live generation still imports (2-GPU cross-process).""" + _spawn_2gpu_regression("stale-generation") if __name__ == "__main__": @@ -535,19 +578,13 @@ def test_vmm_ipc_rejects_stale_generation_import(): _run_core_zero_alignment_reserve() elif "--malformed-preallocated-ipc-handle" in sys.argv: _run_core_malformed_preallocated_ipc_handle() - elif "--cached-chunk-export" in sys.argv: - _run_core_cached_chunk_export() - elif "--cached-chunk-importer" in sys.argv: - _run_core_cached_chunk_importer(sys.argv[2], int(sys.argv[3])) - elif "--stale-generation-export" in sys.argv: - _run_core_stale_generation_export() - elif "--stale-generation-importer" in sys.argv: - _run_core_stale_generation_importer(sys.argv[2], int(sys.argv[3])) + elif "--2gpu-cached-chunk" in sys.argv: + _run_2gpu_regression("cached-chunk", int(sys.argv[2])) + elif "--2gpu-stale-generation" in sys.argv: + _run_2gpu_regression("stale-generation", int(sys.argv[2])) else: test_vmm_allocation() test_vmm_allocation_without_region() test_size_parsing() test_zero_alignment_reserve_stays_in_region() test_vmm_ipc_rejects_malformed_preallocated_chunk_handle() - test_vmm_ipc_cached_chunk_rejects_inconsistent_metadata() - test_vmm_ipc_rejects_stale_generation_import() From 632df90c32e8cfe169fc86c086e5b77c54e9067a Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Fri, 24 Jul 2026 16:57:34 +0000 Subject: [PATCH 51/59] feat: accept documented DEEPEP_MODE override in SGLang recipe The SGLang experimental serve script now threads --deepep-mode through a DEEPEP_MODE env override (default low_latency) instead of hardcoding low_latency. DEEPEP_MODE=auto is the supported nonzero-num_nvl_bytes mode that keeps low-latency decode graphs Foundry captures (verified against pinned sglang@4cfcc25: DeepEPMode.AUTO enables both normal and low_latency so get_deepep_buffer computes a nonzero NVL buffer, and server_args only sets disable_cuda_graph for deepep_mode=normal, not auto). README documents the override and the normal-mode limitation. Co-authored-by: Rahul Chalamala Signed-off-by: Cursor Agent --- recipe/experimental/README.md | 20 +++++++++++++++++++ .../serve_qwen3-30ba3bfp8_sglang_ipc_ep.sh | 10 +++++++++- 2 files changed, 29 insertions(+), 1 deletion(-) diff --git a/recipe/experimental/README.md b/recipe/experimental/README.md index 1cd0da5f..6a63eeb4 100644 --- a/recipe/experimental/README.md +++ b/recipe/experimental/README.md @@ -83,6 +83,26 @@ settings identical between SAVE and LOAD. Both phases pin `--random-seed 42` and `cuMemCreate failed with error 2` while reserving each rank's saved remainder, whereas `0.65` restores every graph and reaches the endpoint on both ranks. +### DeepEP dispatch mode override (`DEEPEP_MODE`) + +The SGLang script threads `--deepep-mode "$DEEPEP_MODE"` through to SGLang and +defaults to `low_latency`, whose graph-capturable decode dispatch computes +`num_nvl_bytes=0` (no intranode NVLink buffer, so the VMM-IPC import path is not +exercised end to end). Override the mode with the same value on SAVE and LOAD: + +```bash +DEEPEP_MODE=auto bash recipe/experimental/serve_qwen3-30ba3bfp8_sglang_ipc_ep.sh 2 --save +DEEPEP_MODE=auto bash recipe/experimental/serve_qwen3-30ba3bfp8_sglang_ipc_ep.sh 2 --load +``` + +`DEEPEP_MODE=auto` keeps the low-latency decode graphs Foundry captures while +also allocating DeepEP's normal-path intranode NVLink buffer with a nonzero +`num_nvl_bytes` (from `get_nvl_buffer_size_hint`). With `nvl_ipc` transport that +buffer flows through Foundry's VMM-IPC translation, so LOAD observes a real VMM +import. `DEEPEP_MODE=normal` is intentionally unsupported for SAVE/LOAD: SGLang +forces `disable_cuda_graph=True` when `deepep_mode=normal`, leaving Foundry no +captured graphs to persist or restore. + ## Archive parity and failure behavior The SAVE and LOAD config for an engine must match on `deepep_transport`, diff --git a/recipe/experimental/serve_qwen3-30ba3bfp8_sglang_ipc_ep.sh b/recipe/experimental/serve_qwen3-30ba3bfp8_sglang_ipc_ep.sh index 77fa7853..a6b57c44 100755 --- a/recipe/experimental/serve_qwen3-30ba3bfp8_sglang_ipc_ep.sh +++ b/recipe/experimental/serve_qwen3-30ba3bfp8_sglang_ipc_ep.sh @@ -13,6 +13,14 @@ EP_SIZE=${1:?Usage: $0 [--save|--load]} MODEL_NAME="${SGL_MODEL:-Qwen/Qwen3-30B-A3B-FP8}" HOST="0.0.0.0" PORT=12000 +# DeepEP dispatch mode (SGLang --deepep-mode). Default low_latency keeps the +# product's graph-capturable decode dispatch with num_nvl_bytes=0. Override with +# DEEPEP_MODE=auto to also allocate DeepEP's nonzero-num_nvl_bytes intranode +# NVLink buffer (exercising the Foundry VMM-IPC import) while keeping the +# low-latency decode graphs. DEEPEP_MODE=normal is unsupported here: SGLang sets +# disable_cuda_graph=True for deepep_mode=normal, so Foundry has no graphs to +# SAVE/LOAD. +DEEPEP_MODE="${DEEPEP_MODE:-low_latency}" # Validated static memory fraction: LOAD reserves each rank's saved remainder # during graph restoration, which failed at 0.8 with cuMemCreate error 2 but # succeeds at 0.65 (both ranks restore every graph and reach the endpoint). @@ -51,7 +59,7 @@ sglang serve \ --ep-size "$EP_SIZE" \ --enable-dp-attention \ --moe-a2a-backend deepep \ - --deepep-mode low_latency \ + --deepep-mode "$DEEPEP_MODE" \ --moe-runner-backend deep_gemm \ --mem-fraction-static "$MEM_FRACTION_STATIC" \ --disable-radix-cache \ From dea454b176446b26bc9770b2ad889690afad389d Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Fri, 24 Jul 2026 16:58:04 +0000 Subject: [PATCH 52/59] test: assert SGLang hook doc lists _patch_attn_tp_logits_gather Co-authored-by: Rahul Chalamala Signed-off-by: Cursor Agent --- tests/test_experimental_recipe.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/tests/test_experimental_recipe.py b/tests/test_experimental_recipe.py index d2a26f94..b65e49ec 100644 --- a/tests/test_experimental_recipe.py +++ b/tests/test_experimental_recipe.py @@ -113,6 +113,16 @@ def test_sglang_docs_define_required_deepep_api() -> None: assert "29d31c0" in text +def test_sglang_hooks_doc_lists_attn_tp_logits_gather_patch() -> None: + text = (ROOT / "docs" / "sglang" / "hooks.md").read_text() + + # The LOAD logits pre-copy patch must be listed in the install order and + # described, not omitted from the hook-install documentation. + assert "_patch_attn_tp_logits_gather" in text + assert text.count("_patch_attn_tp_logits_gather") >= 2 + assert "pre-cop" in text.lower() + + def test_sglang_script_accepts_documented_deepep_mode_override() -> None: script = (EXPERIMENTAL_DIR / SERVE_SCRIPTS[1]).read_text() From ec2687f9047ce76ece39e76d1ca06a203a505964 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Fri, 24 Jul 2026 16:58:11 +0000 Subject: [PATCH 53/59] docs: document _patch_attn_tp_logits_gather in SGLang hook surface Add the LOAD-only logits pre-copy patch to the install-order list and a dedicated section describing both wrapped seams (attention-TP and regular TP gather), renumbering the CudaGraphRunner and spawn-site sections and fixing the internal cross-references accordingly. Co-authored-by: Rahul Chalamala Signed-off-by: Cursor Agent --- docs/sglang/hooks.md | 47 ++++++++++++++++++++++++++++++++++---------- 1 file changed, 37 insertions(+), 10 deletions(-) diff --git a/docs/sglang/hooks.md b/docs/sglang/hooks.md index 28a41ec2..10360fd0 100644 --- a/docs/sglang/hooks.md +++ b/docs/sglang/hooks.md @@ -10,9 +10,10 @@ In install order: 2. `_patch_init_memory_pool` 3. `_patch_load_model` 4. `_patch_kernel_warmup` -5. `_patch_cuda_graph_capture` -6. `_patch_spawn_sites` -7. `_patch_deepep` +5. `_patch_attn_tp_logits_gather` +6. `_patch_cuda_graph_capture` +7. `_patch_spawn_sites` +8. `_patch_deepep` > Removed: an earlier `_patch_model_runner_init` wrapped `ModelRunner.__init__` to stash `dp_rank` on `self` (upstream didn't expose it as an attribute). Upstream sglang now does `self.dp_rank = dp_rank` in the constructor, so our wrapper is no longer needed; `_resolve_dp_rank` reads `self.dp_rank` directly. @@ -72,19 +73,45 @@ No-op on SAVE/LOAD: The FlashInfer autotune path inside is shut off both by this no-op and by the forced `disable_flashinfer_autotune` flag. -### 5. `CudaGraphRunner` (the largest patch) +### 5. `attn_tp_all_gather_into_tensor` / `GroupCoordinator._all_gather_into_tensor` + +**LOAD only.** `_patch_attn_tp_logits_gather` pre-copies each rank's local logits +shard into the gather output buffer before the recognized TP all-gather seams +run. On LOAD the captured graph replays a gather whose peer inputs are not +re-executed, so without seeding the local shard the rank's own slice of the +gathered logits would be stale. Two seams are wrapped when present: + +- `sglang.srt.layers.logits_processor.attn_tp_all_gather_into_tensor` (DP/EP + attention-TP gather). Requires `get_attention_tp_group()` to expose integer + `rank_in_group` / `world_size`; validates the `(world_size, *input.shape)` + output layout and fails fast on invalid rank bounds. +- `GroupCoordinator._all_gather_into_tensor` (regular TP gather). Pre-copies the + local shard only for recognized rank-major `(world_size, *input.shape)` and + concatenated `(world_size*input.shape[0], *input.shape[1:])` layouts via + `_regular_tp_local_shard`; any other caller or unrecognized/unviewable layout + is delegated to the original collective unchanged, preserving its return value + and exceptions. Invalid rank bounds still raise `RuntimeError`. + +`NONE`/`SAVE` modes are untouched. The wrappers are idempotent +(`_foundry_load_local_shard_precopy` guard), and the installed seams are logged: + +```text +[Foundry] SGLang LOAD logits pre-copy patches installed: attention_tp, regular_tp +``` + +### 6. `CudaGraphRunner` (the largest patch) Four sub-patches: -#### 5a. `_create_device_graph` +#### 6a. `_create_device_graph` On SAVE returns `foundry.CUDAGraph()` instead of `torch.cuda.CUDAGraph()`. -#### 5b. `_capture_graph` +#### 6b. `_capture_graph` On SAVE enters `foundry.graph(graph, pool=pool, stream=stream)` instead of `torch.cuda.graph(...)`, captures the `run_once` callable's allocations into foundry's hook event log. -#### 5c. `capture_one_batch_size` +#### 6c. `capture_one_batch_size` On SAVE wraps the `forward` callable with a 3-call counter: @@ -108,7 +135,7 @@ key = bs if stream_idx is None else f"{stream_idx}_{bs}" (Sglang removed the `_make_graph_key` / `get_capture_lora_variant` helpers in commit `ce2506e1c`, the same commit that deprecated `record_nolora_graph` dual MoE graph capture.) -#### 5d. `capture` +#### 6d. `capture` The outermost replacement. @@ -153,7 +180,7 @@ self.output_buffers = {k: v[1] for k, v in state.loaded_graphs.items()} `load_all_graphs` calls `start_graph_builds(all_paths)` and `finish_graph_loads(pending)` exactly once each. This is required for the manifest's template + on-demand linking to work; per-graph `start_graph_builds([single_path])` calls would leave on-demand graphs without a `shared_exec`, and runtime replay would abort with `Called CUDAGraph::replay without a preceding successful capture or load`. -### 6. Spawn-site patches +### 7. Spawn-site patches Two parent-side wrappers: @@ -218,7 +245,7 @@ TP attention is unsupported (its NCCL all-reduce is incompatible with the VMM re - **`deepep_adapter` mode on LOAD.** LOAD replaces the capture loop, so the adapter's `_captured_deepep_mode` is never set; replay asserts on it. The hook calls `deepep_adapter.capture(is_extend_in_batch=False)` after load. -- **FlashInfer pre-pass gated to FlashInfer.** The §5d pre-pass + `reuse_pre_pass_init` +- **FlashInfer pre-pass gated to FlashInfer.** The §6d pre-pass + `reuse_pre_pass_init` shim handle FlashInfer's per-bs wrappers. fa3 (`FlashAttentionBackend`) uses a single fixed `init_cuda_graph_state` workspace, so the shim is skipped (detected via absence of `indices_updater_decode`); but fa3's per-bs `decode_cuda_graph_metadata[bs]` is still From 287eb14050bd79cc67cef1a7939ad7021bef965b Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Fri, 24 Jul 2026 17:05:39 +0000 Subject: [PATCH 54/59] test: assert DeepEP dispatch routing instead of printing Rewrite _verify_dispatch_result to decode every valid received token (bounded by the per-expert recv_count DeepEP returns, since padding slots are not zeroed) and assert the decoded (src_rank, token_id) is in range and that this rank's global expert id is in the token's deterministic top-k selection. Route the LOAD post-replay check through the same helper. The token payload still encodes rank*1000+token here, which is not exactly representable in bfloat16 for rank>0 (>=1000 rounds to a multiple of 4/8), so the new assertions expose the rounding: the decoded token id no longer matches the expert it landed on. This is the RED baseline for the encoding fix. Co-authored-by: Rahul Chalamala Signed-off-by: Cursor Agent --- tests/test_deepep_fabric.py | 209 ++++++++++++++++++------------------ 1 file changed, 104 insertions(+), 105 deletions(-) diff --git a/tests/test_deepep_fabric.py b/tests/test_deepep_fabric.py index b1395582..1f57c9ee 100644 --- a/tests/test_deepep_fabric.py +++ b/tests/test_deepep_fabric.py @@ -30,6 +30,23 @@ ARCHIVE_DIR = "deepep_fabric_archive" HOOK_ARCHIVE_DIR = "hook_archive" +# Token payload encoding used by the deterministic dispatch check: +# x[token, :] = _TOKEN_OFFSET + rank * _RANK_STRIDE + token +# _verify_dispatch_result decodes each received payload back to (src_rank, +# token_id) and asserts that this rank's global expert id is in the token's +# deterministic top-k selection {(token_id + k) % num_experts}. +_RANK_STRIDE = 1000 +_TOKEN_OFFSET = 0 + + +def _encode_token(rank: int, token: int) -> float: + return float(_TOKEN_OFFSET + rank * _RANK_STRIDE + token) + + +def _decode_token(value: float) -> tuple[int, int]: + raw = int(round(float(value))) - _TOKEN_OFFSET + return raw // _RANK_STRIDE, raw % _RANK_STRIDE + def _get_hook_so_path(): import importlib.util @@ -108,10 +125,10 @@ def _create_deterministic_inputs( """ import deep_ep - # x[token_i, :] = rank * 1000 + token_i (easy to trace which rank/token) + # x[token_i, :] = _encode_token(rank, token_i) (traceable to which rank/token) x = torch.zeros((num_tokens, hidden), dtype=torch.bfloat16, device="cuda") for i in range(num_tokens): - x[i, :] = float(rank * 1000 + i) + x[i, :] = _encode_token(rank, i) # Deterministic topk_idx: token i selects experts [i % num_experts, (i+1) % num_experts, ...] topk_idx = torch.zeros((num_tokens, num_topk), dtype=deep_ep.topk_idx_t, device="cuda") @@ -169,85 +186,75 @@ def _print_buffer_info(buffer, rank: int, prefix: str): def _verify_dispatch_result( recv_x, - recv_topk_idx, - recv_src_idx, - num_recv, + recv_count, rank: int, num_ranks: int, num_tokens: int, - hidden: int, num_experts: int, + num_topk: int, prefix: str, ): - """Verify dispatch results are correct based on our deterministic pattern.""" + """Assert dispatched tokens obey the deterministic routing invariant. + + Inputs are seeded via ``_encode_token`` and token ``t`` selects experts + ``{(t + k) % num_experts : k in range(num_topk)}``. ``recv_x`` is shaped + ``[num_local_experts, num_max_dispatch_tokens_per_rank * num_ranks, hidden]`` + and only the first ``recv_count[e]`` slots of local expert ``e`` are valid + (DeepEP does not zero the padding). For every valid slot, decode the payload + back to ``(src_rank, token_id)`` and assert it is in range and that this + rank's global expert id is in the token's top-k selection. Returns the number + of verified tokens and requires at least one. + """ + assert recv_x is not None, f"{prefix}: recv_x is None" + assert recv_x.dim() == 3, f"{prefix}: unexpected recv_x shape {tuple(recv_x.shape)}" + num_local_experts = num_experts // num_ranks local_expert_start = rank * num_local_experts - local_expert_end = local_expert_start + num_local_experts - - # num_recv might be an EventOverlap object (async), need to sync and get value - actual_num_recv = num_recv - if hasattr(num_recv, "wait"): - # It's an async object, wait for it - num_recv.wait() - if hasattr(num_recv, "value"): - actual_num_recv = num_recv.value - elif hasattr(num_recv, "item"): - actual_num_recv = num_recv.item() - elif not isinstance(num_recv, int): - # Try to get from recv_x shape - actual_num_recv = ( - recv_x.shape[0] * recv_x.shape[1] - if recv_x is not None and len(recv_x.shape) >= 2 - else 0 - ) - print( - f"[Rank {rank}] {prefix}: num_recv is {type(num_recv)}, using recv_x shape to estimate", - flush=True, - ) - print(f"[Rank {rank}] {prefix}: Verification - num_recv = {actual_num_recv}", flush=True) + counts = [int(c) for c in recv_count.detach().to("cpu").reshape(-1).tolist()] + assert len(counts) == num_local_experts, ( + f"{prefix}: recv_count has {len(counts)} entries, expected {num_local_experts}" + ) + print( - f"[Rank {rank}] {prefix}: Verification - local experts range = [{local_expert_start}, {local_expert_end})", + f"[Rank {rank}] {prefix}: recv_x shape = {tuple(recv_x.shape)}, " + f"local experts [{local_expert_start}, {local_expert_start + num_local_experts}), " + f"counts = {counts}", flush=True, ) - # recv_x shape is [num_local_experts, max_tokens_per_expert, hidden] - # Print shape info for debugging - if recv_x is not None: - print(f"[Rank {rank}] {prefix}: recv_x shape = {recv_x.shape}", flush=True) - - # Check some values in recv_x to verify the pattern - if recv_x is not None and recv_x.numel() > 0: - # Check first expert's first few tokens - num_experts_in_recv = recv_x.shape[0] - tokens_per_expert = recv_x.shape[1] - check_experts = min(2, num_experts_in_recv) - check_tokens = min(3, tokens_per_expert) - - print( - f"[Rank {rank}] {prefix}: Checking first {check_experts} experts, {check_tokens} tokens each:", - flush=True, + verified = 0 + for expert_i in range(num_local_experts): + global_expert = local_expert_start + expert_i + valid = counts[expert_i] + assert 0 <= valid <= recv_x.shape[1], ( + f"{prefix}: expert {global_expert} count {valid} out of range [0, {recv_x.shape[1]}]" ) + for slot in range(valid): + value = recv_x[expert_i, slot, 0].item() + src_rank, token_id = _decode_token(value) + assert 0 <= src_rank < num_ranks, ( + f"{prefix}: expert {global_expert} slot {slot} value {value} " + f"decodes to src_rank {src_rank} outside [0, {num_ranks})" + ) + assert 0 <= token_id < num_tokens, ( + f"{prefix}: expert {global_expert} slot {slot} value {value} " + f"decodes to token {token_id} outside [0, {num_tokens})" + ) + selected = {(token_id + k) % num_experts for k in range(num_topk)} + assert global_expert in selected, ( + f"{prefix}: expert {global_expert} received token {token_id} from " + f"rank {src_rank} whose top-k {sorted(selected)} does not include it" + ) + verified += 1 - for expert_i in range(check_experts): - expert_global_idx = local_expert_start + expert_i - for token_i in range(check_tokens): - recv_val = recv_x[expert_i, token_i, 0].item() - - # 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 + assert verified > 0, f"{prefix}: no valid tokens verified (counts={counts})" + print( + f"[Rank {rank}] {prefix}: verified {verified} dispatched tokens " + f"across {num_local_experts} local experts", + flush=True, + ) + return verified def _run_save(local_rank: int, num_processes: int): @@ -327,11 +334,11 @@ def _run_save(local_rank: int, num_processes: int): print(f"[Rank {rank}] SAVE: Input tensors:", flush=True) print(f"[Rank {rank}] SAVE: x address = {hex(x.data_ptr())}, shape = {x.shape}", flush=True) print( - f"[Rank {rank}] SAVE: x[0, :5] = {x[0, :5].tolist()} (expect {rank * 1000 + 0})", + f"[Rank {rank}] SAVE: x[0, :5] = {x[0, :5].tolist()} (expect {_encode_token(rank, 0)})", flush=True, ) print( - f"[Rank {rank}] SAVE: x[1, :5] = {x[1, :5].tolist()} (expect {rank * 1000 + 1})", + f"[Rank {rank}] SAVE: x[1, :5] = {x[1, :5].tolist()} (expect {_encode_token(rank, 1)})", flush=True, ) print(f"[Rank {rank}] SAVE: topk_idx address = {hex(topk_idx.data_ptr())}", flush=True) @@ -353,8 +360,10 @@ def _run_save(local_rank: int, num_processes: int): ) torch.cuda.synchronize() - # Unpack dispatch result - recv_x, recv_topk_idx, recv_src_idx, num_recv, *rest = result + # Unpack dispatch result. low_latency_dispatch returns + # (recv_x, recv_count, handle, event, hook); recv_count is a + # [num_local_experts] tensor of per-expert valid token counts. + recv_x, recv_count, *rest = result print(f"[Rank {rank}] SAVE: Warmup dispatch result:", flush=True) print( f"[Rank {rank}] SAVE: recv_x address = {hex(recv_x.data_ptr()) if recv_x is not None else 'None'}", @@ -364,19 +373,17 @@ def _run_save(local_rank: int, num_processes: int): f"[Rank {rank}] SAVE: recv_x shape = {recv_x.shape if recv_x is not None else 'None'}", flush=True, ) - print(f"[Rank {rank}] SAVE: num_recv = {num_recv}", flush=True) + print(f"[Rank {rank}] SAVE: recv_count = {recv_count.tolist()}", flush=True) # Verify results _verify_dispatch_result( recv_x, - recv_topk_idx, - recv_src_idx, - num_recv, + recv_count, rank, num_ranks, num_tokens, - hidden, num_experts, + num_topk, "SAVE-warmup", ) @@ -568,11 +575,11 @@ def _run_load(local_rank: int, num_processes: int): print(f"[Rank {rank}] LOAD: Input tensors:", flush=True) print(f"[Rank {rank}] LOAD: x address = {hex(x.data_ptr())}, shape = {x.shape}", flush=True) print( - f"[Rank {rank}] LOAD: x[0, :5] = {x[0, :5].tolist()} (expect {rank * 1000 + 0})", + f"[Rank {rank}] LOAD: x[0, :5] = {x[0, :5].tolist()} (expect {_encode_token(rank, 0)})", flush=True, ) print( - f"[Rank {rank}] LOAD: x[1, :5] = {x[1, :5].tolist()} (expect {rank * 1000 + 1})", + f"[Rank {rank}] LOAD: x[1, :5] = {x[1, :5].tolist()} (expect {_encode_token(rank, 1)})", flush=True, ) print(f"[Rank {rank}] LOAD: topk_idx address = {hex(topk_idx.data_ptr())}", flush=True) @@ -594,8 +601,9 @@ def _run_load(local_rank: int, num_processes: int): ) torch.cuda.synchronize() - # Unpack and verify warmup result - recv_x, recv_topk_idx, recv_src_idx, num_recv, *rest = warmup_result + # Unpack and verify warmup result. Same layout as SAVE: + # (recv_x, recv_count, handle, event, hook). + recv_x, recv_count, *rest = warmup_result print(f"[Rank {rank}] LOAD: Warmup dispatch result:", flush=True) print( f"[Rank {rank}] LOAD: recv_x address = {hex(recv_x.data_ptr()) if recv_x is not None else 'None'}", @@ -605,19 +613,17 @@ def _run_load(local_rank: int, num_processes: int): f"[Rank {rank}] LOAD: recv_x shape = {recv_x.shape if recv_x is not None else 'None'}", flush=True, ) - print(f"[Rank {rank}] LOAD: num_recv = {num_recv}", flush=True) + print(f"[Rank {rank}] LOAD: recv_count = {recv_count.tolist()}", flush=True) # Verify warmup results _verify_dispatch_result( recv_x, - recv_topk_idx, - recv_src_idx, - num_recv, + recv_count, rank, num_ranks, num_tokens, - hidden, num_experts, + num_topk, "LOAD-warmup", ) @@ -673,29 +679,22 @@ def _run_load(local_rank: int, num_processes: int): graph.replay() torch.cuda.synchronize() - # Verify after replay - check the LOADED output tensors (not warmup recv_x) - # recv_x shape is [num_local_experts, max_tokens_per_expert, hidden] + # Verify after replay against the LOADED output tensor (not warmup recv_x). + # The dispatch is deterministic, so the replayed graph reproduces the same + # per-expert routing as the warmup pass; reuse the warmup recv_count to bound + # the valid slots and assert the same routing invariant. print(f"[Rank {rank}] LOAD: After graph replay:", flush=True) verify_tensor = loaded_recv_x if loaded_recv_x is not None else recv_x - if verify_tensor is not None and verify_tensor.numel() > 0: - print(f"[Rank {rank}] LOAD: verify_tensor shape = {verify_tensor.shape}", flush=True) - num_local_experts = num_experts // num_ranks - local_expert_start = rank * num_local_experts - - # Check first expert's first few tokens - check_experts = min(2, verify_tensor.shape[0]) - check_tokens = min(3, verify_tensor.shape[1]) - - for expert_i in range(check_experts): - for token_i in range(check_tokens): - val = verify_tensor[expert_i, token_i, 0].item() - decoded_src_rank = int(val) // 1000 - decoded_token_id = int(val) % 1000 - print( - f"[Rank {rank}] LOAD: expert[{expert_i}], token[{token_i}]: " - f"value={val:.1f} (rank={decoded_src_rank}, token={decoded_token_id})", - flush=True, - ) + _verify_dispatch_result( + verify_tensor, + recv_count, + rank, + num_ranks, + num_tokens, + num_experts, + num_topk, + "LOAD-replay", + ) print(f"[Rank {rank}] LOAD: Graph replay successful", flush=True) From bb32d60fda0cda6938be92ef71a4e99a9e9ab284 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Fri, 24 Jul 2026 17:07:17 +0000 Subject: [PATCH 55/59] test: encode DeepEP dispatch payloads exactly in bfloat16 Switch the token payload stride/offset so every rank*token id maps to an integer in [1, 128], which bfloat16 represents exactly (integers <= 256). The previous rank*1000 encoding rounded rank>0 payloads to nearby multiples of 4/8, which the new routing assertions correctly reject. With the exact encoding the decoded (src_rank, token_id) round-trips, and the +1 offset keeps real tokens non-zero so they are distinguishable from DeepEP's non-zeroed padding slots. Makes the standalone dispatch matrix pass with the assertions in place. Co-authored-by: Rahul Chalamala Signed-off-by: Cursor Agent --- tests/test_deepep_fabric.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/tests/test_deepep_fabric.py b/tests/test_deepep_fabric.py index 1f57c9ee..57476a68 100644 --- a/tests/test_deepep_fabric.py +++ b/tests/test_deepep_fabric.py @@ -35,8 +35,15 @@ # _verify_dispatch_result decodes each received payload back to (src_rank, # token_id) and asserts that this rank's global expert id is in the token's # deterministic top-k selection {(token_id + k) % num_experts}. -_RANK_STRIDE = 1000 -_TOKEN_OFFSET = 0 +# +# The payload must be exactly representable in bfloat16 (the dispatch dtype), +# which only holds for integers up to 256. With the 64 tokens/rank and 2 ranks +# this test dispatches, _RANK_STRIDE=64 (== tokens per rank) and a +1 offset keep +# every payload in [1, 128]: exact in bf16 and never zero, so real tokens are +# distinguishable from DeepEP's (non-zeroed) padding slots. The stride must be +# >= tokens per rank so token ids stay recoverable via modulo. +_RANK_STRIDE = 64 +_TOKEN_OFFSET = 1 def _encode_token(rank: int, token: int) -> float: From 184bc7601fcf0a1c49d6ed90d4e7f940ebd25078 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Fri, 24 Jul 2026 17:54:56 +0000 Subject: [PATCH 56/59] test: race two importer threads on a peer chunk's first open Add a 2-GPU cross-process regression where rank 1 launches two importer threads that contend on the SAME peer chunk's first open. Both miss the mapping cache and import a temporary mapping; the winner establishes a narrow self-consistent view while the insertion-race loser claims the full chunk and carves past the winner's mapping. The loser must be rejected against the actual cached mapping (CUDA_ERROR_INVALID_VALUE) rather than returning an out-of-mapping interior pointer, and a subsequent valid open/close must still work. Signed-off-by: Cursor Agent Co-authored-by: Rahul Chalamala --- tests/test_vmm_alloc.py | 105 +++++++++++++++++++++++++++++++++++++++- 1 file changed, 104 insertions(+), 1 deletion(-) diff --git a/tests/test_vmm_alloc.py b/tests/test_vmm_alloc.py index fda71ac9..55e47cb1 100644 --- a/tests/test_vmm_alloc.py +++ b/tests/test_vmm_alloc.py @@ -406,9 +406,99 @@ def _importer_check_stale_generation(blob, ipc_open, ipc_close): return True, "stale generation rejected; live generation served" +def _importer_check_concurrent_race(blob, ipc_open, ipc_close): + """Two importer threads contend on the SAME peer chunk's FIRST open, so both + miss the mapping cache and both import their own temporary mapping before one + wins the insertion. The winner establishes a narrow, self-consistent view of + the chunk; the loser carries inconsistent geometry (it claims the full chunk + and carves past the winner's narrow mapping). The insertion-race-loser path + must re-validate the request against the ACTUAL cached mapping and reject it + before adopting the winner's mapping - otherwise it returns an interior + pointer that lands outside the real mapping. + + The winner maps half the chunk and is released a hair before the loser (via + the entered Event), and it maps strictly less than the loser, so it reliably + wins the insertion while the loser still misses the cache on its first find + and lands on the loser path.""" + import threading + + (_magic, _pid, _optr, _size, _token, chunk_base, chunk_size, _gen) = struct.unpack_from( + "=IIQQQQQQ", blob + ) + narrow_size = chunk_size // 2 + + # Winner: a self-consistent view of the first half of the chunk. Maps + # narrow_size bytes and hands out the interior at chunk_base -> valid. + narrow = bytearray(blob) + struct.pack_into("=Q", narrow, _BLOB_ORIGINAL_PTR, chunk_base) + struct.pack_into("=Q", narrow, _BLOB_SIZE, _REGRESSION_ALLOCATION_SIZE) + struct.pack_into("=Q", narrow, _BLOB_CHUNK_SIZE, narrow_size) + + # Loser: same cache key, but claims the full chunk and carves at 3/4 of it - + # inside its own (inflated) chunk so the incoming-handle checks pass and its + # own import succeeds, but past the winner's narrow mapping. + carve_offset = narrow_size + narrow_size // 2 + wide = bytearray(blob) + struct.pack_into("=Q", wide, _BLOB_ORIGINAL_PTR, chunk_base + carve_offset) + struct.pack_into("=Q", wide, _BLOB_SIZE, _REGRESSION_ALLOCATION_SIZE) + struct.pack_into("=Q", wide, _BLOB_CHUNK_SIZE, chunk_size) + + entered = threading.Event() + results = {} + + def open_narrow(): + entered.set() + results["narrow"] = _open_handle(ipc_open, narrow) + + def open_wide(): + entered.wait() + results["wide"] = _open_handle(ipc_open, wide) + + winner = threading.Thread(target=open_narrow) + loser = threading.Thread(target=open_wide) + winner.start() + loser.start() + winner.join() + loser.join() + + narrow_result, narrow_ptr = results["narrow"] + wide_result, wide_ptr = results["wide"] + + if narrow_result != CUDA_SUCCESS: + return False, f"valid (narrow) concurrent open failed: {narrow_result}" + if narrow_ptr.value == 0: + return False, "valid (narrow) concurrent open returned a null pointer" + + # The inconsistent open must be rejected before any pointer is returned. If + # it was accepted, prove the returned interior is out of the winner's real + # mapping (this is exactly the bug: base + carve_offset >= base + narrow). + if wide_result == CUDA_SUCCESS: + offset = wide_ptr.value - narrow_ptr.value + ipc_close(wide_ptr.value) + ipc_close(narrow_ptr.value) + return ( + False, + "inconsistent concurrent open was accepted with an out-of-mapping " + f"pointer (offset 0x{offset:x} into a 0x{narrow_size:x}-byte mapping)", + ) + if wide_result != CUDA_ERROR_INVALID_VALUE: + ipc_close(narrow_ptr.value) + return False, f"inconsistent concurrent open failed with unexpected code: {wide_result}" + + # A subsequent valid open/close of the genuine full-chunk handle still works + # after the loser mapping/handle was cleaned up. + ipc_close(narrow_ptr.value) + live_result, live_ptr = _open_handle(ipc_open, blob) + if live_result != CUDA_SUCCESS: + return False, f"post-race valid open failed: {live_result}" + ipc_close(live_ptr.value) + return True, "insertion-race loser with inconsistent geometry rejected" + + _IMPORTER_CHECKS = { "cached-chunk": _importer_check_cached_chunk, "stale-generation": _importer_check_stale_generation, + "concurrent-race": _importer_check_concurrent_race, } @@ -482,7 +572,9 @@ def _spawn_2gpu_regression(scenario): else: base_env["LD_PRELOAD"] = so_path base_env["MASTER_ADDR"] = "127.0.0.1" - base_env["MASTER_PORT"] = str({"cached-chunk": 29753, "stale-generation": 29757}[scenario]) + base_env["MASTER_PORT"] = str( + {"cached-chunk": 29753, "stale-generation": 29757, "concurrent-race": 29761}[scenario] + ) procs = [] for rank in range(2): @@ -567,6 +659,15 @@ def test_vmm_ipc_rejects_stale_generation_import(): _spawn_2gpu_regression("stale-generation") +@pytest.mark.skipif(_cuda_device_count() < 2, reason="requires 2 GPUs for cross-process IPC") +def test_vmm_ipc_concurrent_first_open_race_rejects_inconsistent_geometry(): + """Two importer threads contend on a peer chunk's first open; the + insertion-race loser carrying inconsistent geometry is rejected against the + actual cached mapping instead of returning an out-of-mapping interior + pointer (2-GPU cross-process).""" + _spawn_2gpu_regression("concurrent-race") + + if __name__ == "__main__": if "--core" in sys.argv: _run_core() @@ -582,6 +683,8 @@ def test_vmm_ipc_rejects_stale_generation_import(): _run_2gpu_regression("cached-chunk", int(sys.argv[2])) elif "--2gpu-stale-generation" in sys.argv: _run_2gpu_regression("stale-generation", int(sys.argv[2])) + elif "--2gpu-concurrent-race" in sys.argv: + _run_2gpu_regression("concurrent-race", int(sys.argv[2])) else: test_vmm_allocation() test_vmm_allocation_without_region() From 8acb26d87aed937440295a453400ffac8add54c9 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Fri, 24 Jul 2026 18:09:17 +0000 Subject: [PATCH 57/59] test: exercise the VMM IPC insertion-race loser path under contention Rework the concurrent regression: two importer threads contend on the same peer chunk's first open so one lands on the insertion-race-loser path and must adopt the winner's whole-chunk mapping (same interior). Because the driver cannot map a sub-range of an imported handle, an inconsistent handle cannot itself reach the loser path, so inconsistent same-key geometry is presented against the live cached mapping (a cache hit) and must be rejected with CUDA_ERROR_INVALID_VALUE and no out-of-mapping pointer, sharing the validator the loser path also uses as defense-in-depth. Signed-off-by: Cursor Agent Co-authored-by: Rahul Chalamala --- tests/test_vmm_alloc.py | 162 ++++++++++++++++++++-------------------- 1 file changed, 83 insertions(+), 79 deletions(-) diff --git a/tests/test_vmm_alloc.py b/tests/test_vmm_alloc.py index 55e47cb1..c0d14e8e 100644 --- a/tests/test_vmm_alloc.py +++ b/tests/test_vmm_alloc.py @@ -407,92 +407,95 @@ def _importer_check_stale_generation(blob, ipc_open, ipc_close): def _importer_check_concurrent_race(blob, ipc_open, ipc_close): - """Two importer threads contend on the SAME peer chunk's FIRST open, so both - miss the mapping cache and both import their own temporary mapping before one - wins the insertion. The winner establishes a narrow, self-consistent view of - the chunk; the loser carries inconsistent geometry (it claims the full chunk - and carves past the winner's narrow mapping). The insertion-race-loser path - must re-validate the request against the ACTUAL cached mapping and reject it - before adopting the winner's mapping - otherwise it returns an interior - pointer that lands outside the real mapping. - - The winner maps half the chunk and is released a hair before the loser (via - the entered Event), and it maps strictly less than the loser, so it reliably - wins the insertion while the loser still misses the cache on its first find - and lands on the loser path.""" + """Contend two importer threads on the SAME peer chunk's FIRST open so both + miss the mapping cache and both import the whole chunk; one wins the + insertion and the other lands on the insertion-race-loser path, which must + adopt the winner's mapping and hand out an interior derived from the winner's + base (not leak its own temporary). Then, against the live cached mapping, + present same-key handles whose geometry is inconsistent with the ACTUAL + mapping: each must be validated against the real mapping and rejected with + CUDA_ERROR_INVALID_VALUE and no out-of-mapping pointer. A subsequent valid + open/close still works. + + Why the inconsistent request is a cache-hit rather than a third racer: the + CUDA driver cannot map a sub-range of an imported handle (see the whole-chunk + mapping comments in csrc/hook.cpp), so every importer maps the FULL chunk. + Any handle that survives the incoming-metadata check and imports + successfully has therefore mapped exactly the cached size with a carve inside + it, i.e. it is already consistent with the cached mapping by the time it + reaches the loser path; an inconsistent chunk_size instead fails its own + full-chunk import before ever reaching that branch. The loser-path + re-validation shares the exact validator proven reachable-rejectable here on + the cache-hit path, as defense-in-depth for that whole-chunk invariant.""" import threading (_magic, _pid, _optr, _size, _token, chunk_base, chunk_size, _gen) = struct.unpack_from( "=IIQQQQQQ", blob ) - narrow_size = chunk_size // 2 - - # Winner: a self-consistent view of the first half of the chunk. Maps - # narrow_size bytes and hands out the interior at chunk_base -> valid. - narrow = bytearray(blob) - struct.pack_into("=Q", narrow, _BLOB_ORIGINAL_PTR, chunk_base) - struct.pack_into("=Q", narrow, _BLOB_SIZE, _REGRESSION_ALLOCATION_SIZE) - struct.pack_into("=Q", narrow, _BLOB_CHUNK_SIZE, narrow_size) - - # Loser: same cache key, but claims the full chunk and carves at 3/4 of it - - # inside its own (inflated) chunk so the incoming-handle checks pass and its - # own import succeeds, but past the winner's narrow mapping. - carve_offset = narrow_size + narrow_size // 2 - wide = bytearray(blob) - struct.pack_into("=Q", wide, _BLOB_ORIGINAL_PTR, chunk_base + carve_offset) - struct.pack_into("=Q", wide, _BLOB_SIZE, _REGRESSION_ALLOCATION_SIZE) - struct.pack_into("=Q", wide, _BLOB_CHUNK_SIZE, chunk_size) - - entered = threading.Event() - results = {} - def open_narrow(): - entered.set() - results["narrow"] = _open_handle(ipc_open, narrow) - - def open_wide(): - entered.wait() - results["wide"] = _open_handle(ipc_open, wide) - - winner = threading.Thread(target=open_narrow) - loser = threading.Thread(target=open_wide) - winner.start() - loser.start() - winner.join() - loser.join() - - narrow_result, narrow_ptr = results["narrow"] - wide_result, wide_ptr = results["wide"] - - if narrow_result != CUDA_SUCCESS: - return False, f"valid (narrow) concurrent open failed: {narrow_result}" - if narrow_ptr.value == 0: - return False, "valid (narrow) concurrent open returned a null pointer" - - # The inconsistent open must be rejected before any pointer is returned. If - # it was accepted, prove the returned interior is out of the winner's real - # mapping (this is exactly the bug: base + carve_offset >= base + narrow). - if wide_result == CUDA_SUCCESS: - offset = wide_ptr.value - narrow_ptr.value - ipc_close(wide_ptr.value) - ipc_close(narrow_ptr.value) - return ( - False, - "inconsistent concurrent open was accepted with an out-of-mapping " - f"pointer (offset 0x{offset:x} into a 0x{narrow_size:x}-byte mapping)", - ) - if wide_result != CUDA_ERROR_INVALID_VALUE: - ipc_close(narrow_ptr.value) - return False, f"inconsistent concurrent open failed with unexpected code: {wide_result}" + barrier = threading.Barrier(2) + opens = {} + + def racer(idx): + barrier.wait() + opens[idx] = _open_handle(ipc_open, blob) + + threads = [threading.Thread(target=racer, args=(i,)) for i in range(2)] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + + for idx, (result, ptr) in opens.items(): + if result != CUDA_SUCCESS: + return False, f"concurrent valid open {idx} failed: {result}" + if ptr.value == 0: + return False, f"concurrent valid open {idx} returned a null pointer" + + # Both racers open the same carve, so the loser must have adopted the + # winner's whole-chunk mapping (same interior), not a leaked temporary. + ptr0 = opens[0][1].value + ptr1 = opens[1][1].value + if ptr0 != ptr1: + return False, f"concurrent opens of the same carve diverged: 0x{ptr0:x} vs 0x{ptr1:x}" + + # The mapping is live (refcount 2). Present same-key handles whose geometry + # is inconsistent with the ACTUAL mapping; each must be rejected on the + # cache-hit path before any pointer is handed back. + rejects = {} + + size_mismatch = bytearray(blob) + struct.pack_into("=Q", size_mismatch, _BLOB_CHUNK_SIZE, chunk_size * 2) + struct.pack_into("=Q", size_mismatch, _BLOB_ORIGINAL_PTR, chunk_base) + struct.pack_into("=Q", size_mismatch, _BLOB_SIZE, _REGRESSION_ALLOCATION_SIZE) + result, opened = _open_handle(ipc_open, size_mismatch) + if result == CUDA_SUCCESS: + ipc_close(opened.value) + rejects["chunk_size_mismatch"] = result + + carve_outside = bytearray(blob) + struct.pack_into("=Q", carve_outside, _BLOB_CHUNK_SIZE, chunk_size * 2) + struct.pack_into("=Q", carve_outside, _BLOB_ORIGINAL_PTR, chunk_base + chunk_size) + struct.pack_into("=Q", carve_outside, _BLOB_SIZE, _REGRESSION_ALLOCATION_SIZE) + result, opened = _open_handle(ipc_open, carve_outside) + if result == CUDA_SUCCESS: + ipc_close(opened.value) + rejects["carve_outside_mapping"] = result + + # Release both concurrent opens (refcount -> 0, mapping torn down). + ipc_close(ptr0) + ipc_close(ptr1) + + for name, code in rejects.items(): + if code != CUDA_ERROR_INVALID_VALUE: + return False, f"{name} not rejected against the live cached mapping: {code}" - # A subsequent valid open/close of the genuine full-chunk handle still works - # after the loser mapping/handle was cleaned up. - ipc_close(narrow_ptr.value) + # Subsequent valid open/close of the genuine handle still works. live_result, live_ptr = _open_handle(ipc_open, blob) if live_result != CUDA_SUCCESS: return False, f"post-race valid open failed: {live_result}" ipc_close(live_ptr.value) - return True, "insertion-race loser with inconsistent geometry rejected" + return True, "loser path adopted winner mapping; inconsistent geometry rejected" _IMPORTER_CHECKS = { @@ -661,10 +664,11 @@ def test_vmm_ipc_rejects_stale_generation_import(): @pytest.mark.skipif(_cuda_device_count() < 2, reason="requires 2 GPUs for cross-process IPC") def test_vmm_ipc_concurrent_first_open_race_rejects_inconsistent_geometry(): - """Two importer threads contend on a peer chunk's first open; the - insertion-race loser carrying inconsistent geometry is rejected against the - actual cached mapping instead of returning an out-of-mapping interior - pointer (2-GPU cross-process).""" + """Two importer threads contend on a peer chunk's first open so the + insertion-race-loser path is exercised (the loser adopts the winner's + whole-chunk mapping); inconsistent same-key geometry against the live mapping + is validated against the actual cached mapping and rejected with no + out-of-mapping pointer (2-GPU cross-process).""" _spawn_2gpu_regression("concurrent-race") From 0093e38edfc35e0278d5ff2cbad8db997e3de40b Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Fri, 24 Jul 2026 18:13:12 +0000 Subject: [PATCH 58/59] fix: factor VMM IPC chunk-request validator for both cache paths The mapping cache key omits chunk_size, so a same-key handle can carry a different (untrusted) chunk_size and carve than the live mapping. The cache-hit fast path already re-validated the request against the actual cached mapping, but the insertion-race-loser path adopted the winner's mapping without the same check. Factor a single vmm_ipc_chunk_request_matches_mapping() validator (chunk-size equality + overflow-safe carve containment against VmmIpcChunkMapping.size) and call it in BOTH paths before any refcount bump, subpointer registration, or pointer return. The loser path releases its losing temporary mapping/handle first and only then re-validates, so all cleanup and lock/refcount behavior is preserved on the reject path. The driver cannot map a sub-range of an imported handle, so in practice every handle reaching the loser path has mapped the full chunk and is already consistent; the loser-path check is defense-in-depth that keeps the two adopt paths symmetric. Signed-off-by: Cursor Agent Co-authored-by: Rahul Chalamala --- csrc/hook.cpp | 73 ++++++++++++++++++++++++++++++++++++--------------- 1 file changed, 52 insertions(+), 21 deletions(-) diff --git a/csrc/hook.cpp b/csrc/hook.cpp index ad3b6ad9..35909151 100644 --- a/csrc/hook.cpp +++ b/csrc/hook.cpp @@ -2969,6 +2969,44 @@ static std::map vmm_ipc_chunk_mappings; // Interior pointer handed to a caller -> owning chunk key and number of opens. static std::map vmm_ipc_chunk_subptrs; +// Validate an incoming chunk-carve open request against the ACTUAL cached +// whole-chunk mapping. The cache key (exporter pid, token, chunk_base, +// generation) deliberately omits chunk_size, so a same-key handle can carry a +// different (untrusted) chunk_size and carve than the live mapping. EVERY path +// that is about to adopt an existing mapping - the initial cache-hit fast path +// AND the insertion-race loser path, which can both race two first opens of the +// same peer chunk - must call this before mutating a refcount, registering an +// interior subpointer, or returning a pointer. Without it a same-key handle +// could smuggle a larger chunk_size and a carve landing outside the real +// mapping, yielding an out-of-mapping interior pointer. All arithmetic is +// overflow-safe because the handle values are untrusted. Must be called with +// vmm_ipc_chunk_mutex held (it reads the live mapping). Returns true iff the +// request is consistent with `mapping`; logs and returns false otherwise, and +// the caller returns CUDA_ERROR_INVALID_VALUE. +static bool vmm_ipc_chunk_request_matches_mapping(uint64_t chunk_base, uint64_t chunk_size, + CUdeviceptr carve_ptr, size_t carve_size, + const VmmIpcChunkMapping& mapping, + const char* context) { + const uint64_t cached_size = (uint64_t)mapping.size; + if (chunk_size != cached_size) { + fprintf(stderr, + "[HOOK] ERROR: VMM-IPC %s chunk_size mismatch: handle=%llu cached=%llu\n", + context, (unsigned long long)chunk_size, (unsigned long long)cached_size); + return false; + } + const uint64_t p = (uint64_t)carve_ptr; + const uint64_t s = (uint64_t)carve_size; + if (p < chunk_base || s > (UINT64_MAX - p) || (p + s) > (chunk_base + cached_size)) { + fprintf(stderr, + "[HOOK] ERROR: VMM-IPC %s carve [0x%llx,+0x%llx) outside cached mapping " + "[0x%llx,+0x%llx)\n", + context, (unsigned long long)p, (unsigned long long)s, + (unsigned long long)chunk_base, (unsigned long long)cached_size); + return false; + } + return true; +} + // Dedicated VA zone for relocated peer imports, far below the foundry region // (default 0x600000000000) and the large-reserve zone right above region end. // Without this, the driver tends to place a relocated import immediately @@ -3389,26 +3427,10 @@ CUresult cuIpcOpenMemHandle(CUdeviceptr* pdptr, CUipcMemHandle handle, unsigned // Cache hit. The incoming handle's chunk_size/carve were only checked // against the handle's own (untrusted) metadata above. Re-validate the // request against the ACTUAL cached mapping before deriving an interior - // pointer: the key does not include chunk_size, so a same-key handle - // could otherwise smuggle a larger chunk_size and a carve that lands - // outside the real mapping. Reject before touching the refcount so all - // bookkeeping is preserved on the error path. - const size_t cached_size = it->second.size; - if ((uint64_t)chunk_size != (uint64_t)cached_size) { - fprintf(stderr, - "[HOOK] ERROR: VMM-IPC cache hit chunk_size mismatch: handle=%llu cached=%zu\n", - (unsigned long long)chunk_size, cached_size); - return CUDA_ERROR_INVALID_VALUE; - } - const uint64_t carve_ptr = (uint64_t)original_ptr; - const uint64_t carve_size = (uint64_t)size; - if (carve_ptr < chunk_base || carve_size > (UINT64_MAX - carve_ptr) || - (carve_ptr + carve_size) > (chunk_base + (uint64_t)cached_size)) { - fprintf(stderr, - "[HOOK] ERROR: VMM-IPC cache hit carve [0x%llx,+0x%llx) outside cached mapping " - "[0x%llx,+0x%zx)\n", - (unsigned long long)carve_ptr, (unsigned long long)carve_size, - (unsigned long long)chunk_base, cached_size); + // pointer, and reject before touching the refcount so all bookkeeping is + // preserved on the error path. + if (!vmm_ipc_chunk_request_matches_mapping(chunk_base, chunk_size, original_ptr, size, + it->second, "cache hit")) { return CUDA_ERROR_INVALID_VALUE; } it->second.refcount++; @@ -3577,13 +3599,22 @@ CUresult cuIpcOpenMemHandle(CUdeviceptr* pdptr, CUipcMemHandle handle, unsigned 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. + // winner's mapping, drop ours. Release the losing temporary mapping and + // handle first (unconditionally, whether or not we adopt the winner). typedef CUresult (*cuMemUnmap_t)(CUdeviceptr, size_t); auto mem_unmap_func = (cuMemUnmap_t)CUDA_DRIVER_CALL(cuda_driver_entry_table, CUDA_ENTRY_cuMemUnmap); mem_unmap_func(mapped_ptr, map_size); real_address_free_func(mapped_ptr, map_size); mem_release_func(imported_handle); + // Same authoritative re-validation as the cache-hit fast path: the key + // omits chunk_size, so this same-key handle may not match the winner's + // live mapping. Reject before adopting it - our temporary is already + // released and the winner's refcount is untouched. + if (!vmm_ipc_chunk_request_matches_mapping(chunk_base, chunk_size, original_ptr, size, + it->second, "insertion race")) { + return CUDA_ERROR_INVALID_VALUE; + } it->second.refcount++; interior = it->second.local_base + (original_ptr - (CUdeviceptr)chunk_base); } else { From 02af2b2c6408bb48a99e8a1b5e6541b514bfd379 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Fri, 24 Jul 2026 18:13:19 +0000 Subject: [PATCH 59/59] test: assert exact DeepEP dispatch delivery per expert Strengthen _verify_dispatch_result beyond routing-only: the deterministic top-k selection depends only on the token id, so each local expert must receive exactly {(src_rank, token) : rank in ranks, token selects expert}, once each. Assert recv_count[e] equals the expected delivery count and that the decoded (src_rank, token_id) multiset equals the expected set with no duplicate, missing, or spurious deliveries. Signed-off-by: Cursor Agent Co-authored-by: Rahul Chalamala --- tests/test_deepep_fabric.py | 45 ++++++++++++++++++++++++++++++++----- 1 file changed, 40 insertions(+), 5 deletions(-) diff --git a/tests/test_deepep_fabric.py b/tests/test_deepep_fabric.py index 57476a68..913b569b 100644 --- a/tests/test_deepep_fabric.py +++ b/tests/test_deepep_fabric.py @@ -201,16 +201,25 @@ def _verify_dispatch_result( num_topk: int, prefix: str, ): - """Assert dispatched tokens obey the deterministic routing invariant. + """Assert dispatched tokens obey the deterministic routing invariant, with + EXACT-delivery checking (no missing, no duplicate, no spurious tokens). Inputs are seeded via ``_encode_token`` and token ``t`` selects experts - ``{(t + k) % num_experts : k in range(num_topk)}``. ``recv_x`` is shaped + ``{(t + k) % num_experts : k in range(num_topk)}`` identically on every rank + (the selection depends only on the token id). So for a local expert ``e`` the + exact set of deliveries it must receive is + ``{(r, t) : r in range(num_ranks), t in range(num_tokens), e in selected(t)}`` + - each source rank contributes every token that selected ``e``, exactly once. + + ``recv_x`` is shaped ``[num_local_experts, num_max_dispatch_tokens_per_rank * num_ranks, hidden]`` and only the first ``recv_count[e]`` slots of local expert ``e`` are valid (DeepEP does not zero the padding). For every valid slot, decode the payload - back to ``(src_rank, token_id)`` and assert it is in range and that this - rank's global expert id is in the token's top-k selection. Returns the number - of verified tokens and requires at least one. + back to ``(src_rank, token_id)``, assert it is in range and routed correctly, + then require the decoded multiset for ``e`` to equal the expected set exactly: + ``recv_count[e]`` must match the expected delivery count and there must be no + duplicate or missing ``(src_rank, token_id)`` pair. Returns the number of + verified tokens and requires at least one. """ assert recv_x is not None, f"{prefix}: recv_x is None" assert recv_x.dim() == 3, f"{prefix}: unexpected recv_x shape {tuple(recv_x.shape)}" @@ -237,6 +246,20 @@ def _verify_dispatch_result( assert 0 <= valid <= recv_x.shape[1], ( f"{prefix}: expert {global_expert} count {valid} out of range [0, {recv_x.shape[1]}]" ) + + # The exact set of (src_rank, token_id) deliveries this expert must get. + expected = { + (r, t) + for r in range(num_ranks) + for t in range(num_tokens) + if global_expert in {(t + k) % num_experts for k in range(num_topk)} + } + assert valid == len(expected), ( + f"{prefix}: expert {global_expert} recv_count {valid} != expected " + f"delivery count {len(expected)} (missing or duplicate dispatches)" + ) + + received = [] for slot in range(valid): value = recv_x[expert_i, slot, 0].item() src_rank, token_id = _decode_token(value) @@ -253,8 +276,20 @@ def _verify_dispatch_result( f"{prefix}: expert {global_expert} received token {token_id} from " f"rank {src_rank} whose top-k {sorted(selected)} does not include it" ) + received.append((src_rank, token_id)) verified += 1 + received_set = set(received) + assert len(received_set) == len(received), ( + f"{prefix}: expert {global_expert} received duplicate deliveries: " + f"{sorted(p for p in received_set if received.count(p) > 1)}" + ) + assert received_set == expected, ( + f"{prefix}: expert {global_expert} delivery set mismatch; " + f"missing={sorted(expected - received_set)} " + f"unexpected={sorted(received_set - expected)}" + ) + assert verified > 0, f"{prefix}: no valid tokens verified (counts={counts})" print( f"[Rank {rank}] {prefix}: verified {verified} dispatched tokens "