From cc835028746f19fbbfd3638ff3710a5f0177a945 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 23 Jul 2026 00:00:00 +0000 Subject: [PATCH 01/22] [sglang] add experimental dense tensor-parallel recipe Co-Authored-By: Rahul Chalamala <22563365+rchalamala@users.noreply.github.com> Signed-off-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- recipe/experimental/README.md | 51 +++++++++++++++- .../experimental/serve_qwen3-8b_sglang_tp.sh | 61 +++++++++++++++++++ .../experimental/sglang_foundry_tp_load.toml | 8 +++ .../experimental/sglang_foundry_tp_save.toml | 8 +++ 4 files changed, 125 insertions(+), 3 deletions(-) create mode 100755 recipe/experimental/serve_qwen3-8b_sglang_tp.sh create mode 100644 recipe/experimental/sglang_foundry_tp_load.toml create mode 100644 recipe/experimental/sglang_foundry_tp_save.toml diff --git a/recipe/experimental/README.md b/recipe/experimental/README.md index 23a092d9..ec27f715 100644 --- a/recipe/experimental/README.md +++ b/recipe/experimental/README.md @@ -19,9 +19,12 @@ recipe/experimental/ ├── foundry_save.toml # SAVE config (workspace_root = "foundry_archive_ipc") ├── foundry_load.toml # LOAD config (same workspace_root) ├── serve_qwen3-30ba3b_ipc_ep.sh # Qwen3-30B-A3B (MoE) vLLM EP, NVL/IPC -├── sglang_foundry_save.toml # SGLang SAVE config -├── sglang_foundry_load.toml # SGLang LOAD config -└── serve_qwen3-30ba3bfp8_sglang_ipc_ep.sh # SGLang EP, DeepEP NVL/IPC +├── sglang_foundry_save.toml # SGLang EP SAVE config +├── sglang_foundry_load.toml # SGLang EP LOAD config +├── serve_qwen3-30ba3bfp8_sglang_ipc_ep.sh # SGLang EP, DeepEP NVL/IPC +├── sglang_foundry_tp_save.toml # SGLang dense-TP SAVE config +├── sglang_foundry_tp_load.toml # SGLang dense-TP LOAD config +└── serve_qwen3-8b_sglang_tp.sh # SGLang dense tensor parallel (NCCL all-reduce) ``` It is the standard `recipe/vllm` EP recipe plus one switch — `FOUNDRY_DEEPEP_NVL_IPC=1` @@ -137,6 +140,43 @@ SAVE and LOAD. The hook then: A healthy run logs two `VMM-IPC import relocated` lines per phase (one per peer) and **no** `error 999`. +## SGLang dense tensor parallel + +`serve_qwen3-8b_sglang_tp.sh` runs a **dense** model tensor-parallel across +`` GPUs (`--tp-size N`, no DP-attention, no expert parallel) — the +first recipe here that captures a **cross-rank collective into the graph +itself**. The DP recipe replicates the whole model per rank (no in-graph +collective) and the EP recipe puts attention on DP-attention and the MoE on +DeepEP's all-to-all; neither exercises a per-layer TP all-reduce. + +```bash +rm -rf foundry_archive_sglang_tp +CUDA_VISIBLE_DEVICES=0,1 bash serve_qwen3-8b_sglang_tp.sh 2 --save +CUDA_VISIBLE_DEVICES=0,1 bash serve_qwen3-8b_sglang_tp.sh 2 --load +bash ../../../experimental/query.sh 12000 Qwen/Qwen3-8B +``` + +How the TP all-reduce survives SAVE/LOAD: + +- `--disable-custom-all-reduce` routes the per-layer all-reduce through **NCCL** + (not SGLang's custom one-shot/two-shot kernel). Both are CUDA-IPC based; NCCL + is the path validated against the VMM-IPC bridge. +- `NCCL_CUMEM_ENABLE=0` / `NCCL_NVLS_ENABLE=0` keep NCCL off the CUMEM P2P and + NVLS multicast fast paths (whose driver-capability flags the foundry VMM + region doesn't carry) and on the **legacy CUDA-IPC** intra-node transport. +- NCCL's peer buffers are foundry VMM allocations shared via + `cuIpcGetMemHandle` / `cuIpcOpenMemHandle`. The same VMM-IPC bridge the EP/NVL + path uses interposes them: it transports the backing fd between ranks over + `SCM_RIGHTS` and re-maps each peer buffer, so the peer addresses baked into the + captured all-reduce resolve on LOAD instead of dangling. + +No NVSHMEM is involved (dense, no DeepEP), so the TP TOMLs leave +`nvshmem_host_path` unset. Keep `--tp-size`, `--cuda-graph-max-bs`, and the NCCL +environment identical between SAVE and LOAD so the captured graphs and the NCCL +buffer trajectory match. Custom all-reduce (SGLang's own IPC kernel) uses the +same `cuIpc` primitives and is expected to work over this bridge too, but is +left disabled here until GPU-validated. + ## Validation status SAVE → SAVE → LOAD → query verified on Qwen3-30B-A3B EP=2 (2× H200, no IMEX): @@ -146,6 +186,11 @@ offset, query returns coherent completions. LOAD reaches a serving server in has no steady-state serving cost — peer addresses are resolved once at init via the device pointer table, not per token. +The **dense TP** recipe (`serve_qwen3-8b_sglang_tp.sh`) is **not yet +GPU-validated** — it reuses the same VMM-IPC bridge as the EP/NVL path, but the +in-graph NCCL all-reduce is a new code path here. GPU validation (SAVE → LOAD → +temperature-0 A/B vs baseline on 2× H200) is the follow-up. + ## IPC-specific troubleshooting | Symptom | Likely cause | diff --git a/recipe/experimental/serve_qwen3-8b_sglang_tp.sh b/recipe/experimental/serve_qwen3-8b_sglang_tp.sh new file mode 100755 index 00000000..71eccb68 --- /dev/null +++ b/recipe/experimental/serve_qwen3-8b_sglang_tp.sh @@ -0,0 +1,61 @@ +#!/bin/bash +# Usage: CUDA_VISIBLE_DEVICES=0,1 bash serve_qwen3-8b_sglang_tp.sh [--save|--load] +# +# EXPERIMENTAL: SGLang dense tensor parallel (tp-size > 1, no DP-attention, no +# expert parallel). Unlike the DP recipe — where each rank is an independent +# replica with no cross-rank collective inside the captured graph — pure TP +# captures a per-layer all-reduce across the TP ranks into every CUDA graph. +# +# That all-reduce is served by NCCL (custom all-reduce is disabled): with +# NCCL_CUMEM_ENABLE=0 NCCL's intra-node transport shares its peer buffers over +# legacy CUDA IPC (cuIpcGetMemHandle / cuIpcOpenMemHandle). Those handles are +# process-local and would be dangling on LOAD; the Foundry VMM-IPC bridge in +# libcuda_hook.so (the ep-ipc commit) intercepts them, transports the backing +# fd between ranks with SCM_RIGHTS, and re-maps each peer buffer at its recorded +# address — the same mechanism the DeepEP NVL/IPC path relies on. +# +# Requires the Foundry hook from the ep-ipc commit (cuIpc VMM-IPC bridge). +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" + +TP_SIZE=${1:?Usage: $0 [--save|--load]} +MODEL_NAME="${SGL_MODEL:-Qwen/Qwen3-8B}" +HOST="0.0.0.0" +PORT=12000 +MEM_FRACTION_STATIC=0.8 + +FOUNDRY_ARGS=() +if [[ "$2" == "--save" ]]; then + FOUNDRY_TOML="${SCRIPT_DIR}/sglang_foundry_tp_save.toml" +elif [[ "$2" == "--load" ]]; then + FOUNDRY_TOML="${SCRIPT_DIR}/sglang_foundry_tp_load.toml" +elif [[ -n "$2" ]]; then + echo "Usage: $0 [--save|--load]" + exit 1 +fi + +if [[ -n "${FOUNDRY_TOML:-}" ]]; then + # CUMEM P2P / NVLS multicast cuMemMap with driver-capability flags the + # foundry VMM region doesn't carry — pin NCCL to the IPC/ring transports so + # its peer buffers go through cuIpc (which the VMM-IPC bridge interposes). + export NCCL_CUMEM_ENABLE=0 + export NCCL_NVLS_ENABLE=0 + FOUNDRY_ARGS+=( --foundry-graph-extension-config-path "$FOUNDRY_TOML" ) + echo "Using Foundry TP (NCCL all-reduce over VMM-IPC): ${FOUNDRY_TOML}" +else + echo "Running without Foundry (baseline SGLang)" +fi + +# --disable-custom-all-reduce: route the TP all-reduce through NCCL rather than +# SGLang's custom one-shot/two-shot kernel. Both are CUDA-IPC based, but the +# NCCL path is the one validated against the Foundry VMM-IPC bridge here. +sglang serve \ + --model-path "$MODEL_NAME" \ + --trust-remote-code \ + --host "$HOST" --port "$PORT" \ + --tp-size "$TP_SIZE" \ + --mem-fraction-static "$MEM_FRACTION_STATIC" \ + --disable-radix-cache \ + --disable-custom-all-reduce \ + --attention-backend flashinfer \ + --cuda-graph-max-bs 256 \ + "${FOUNDRY_ARGS[@]}" diff --git a/recipe/experimental/sglang_foundry_tp_load.toml b/recipe/experimental/sglang_foundry_tp_load.toml new file mode 100644 index 00000000..938f2ad2 --- /dev/null +++ b/recipe/experimental/sglang_foundry_tp_load.toml @@ -0,0 +1,8 @@ +mode = "load" +base_addr = 0x600000000000 +region_size = "256GB" +workspace_root = "foundry_archive_sglang_tp" +scratch_space_size = "1024MB" +# Dense tensor parallel needs no NVSHMEM (no DeepEP). The per-layer TP +# all-reduce is served by NCCL; its intra-node peer buffers go over legacy +# CUDA IPC, which the Foundry VMM-IPC bridge in libcuda_hook.so interposes. diff --git a/recipe/experimental/sglang_foundry_tp_save.toml b/recipe/experimental/sglang_foundry_tp_save.toml new file mode 100644 index 00000000..0c8380a9 --- /dev/null +++ b/recipe/experimental/sglang_foundry_tp_save.toml @@ -0,0 +1,8 @@ +mode = "save" +base_addr = 0x600000000000 +region_size = "256GB" +workspace_root = "foundry_archive_sglang_tp" +scratch_space_size = "1024MB" +# Dense tensor parallel needs no NVSHMEM (no DeepEP). The per-layer TP +# all-reduce is served by NCCL; its intra-node peer buffers go over legacy +# CUDA IPC, which the Foundry VMM-IPC bridge in libcuda_hook.so interposes. From ee7d61a32a4d4732e5f14978975b96ab03bcbede Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 23 Jul 2026 01:32:26 +0000 Subject: [PATCH 02/22] [sglang] record dense-TP GPU validation results (Qwen3-8B TP=2, 2x H200) SAVE->SAVE->LOAD validated on 2x H200 (no IMEX): no MMU/Xid/CUDA fault, per-rank final_alloc_offset reproducible across passes, LOAD ~52s vs ~172s cold SAVE, and temperature-0 LOAD-engine tokens match the warmed SAVE engine exactly. scratch_space_size=1024MB sufficient. Document the required pre-#26735 sglang base for the foundry hook. Co-Authored-By: Rahul Chalamala <22563365+rchalamala@users.noreply.github.com> Signed-off-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- recipe/experimental/README.md | 31 +++++++++++++++++++++++++++---- 1 file changed, 27 insertions(+), 4 deletions(-) diff --git a/recipe/experimental/README.md b/recipe/experimental/README.md index ec27f715..1d781bd1 100644 --- a/recipe/experimental/README.md +++ b/recipe/experimental/README.md @@ -186,10 +186,33 @@ offset, query returns coherent completions. LOAD reaches a serving server in has no steady-state serving cost — peer addresses are resolved once at init via the device pointer table, not per token. -The **dense TP** recipe (`serve_qwen3-8b_sglang_tp.sh`) is **not yet -GPU-validated** — it reuses the same VMM-IPC bridge as the EP/NVL path, but the -in-graph NCCL all-reduce is a new code path here. GPU validation (SAVE → LOAD → -temperature-0 A/B vs baseline on 2× H200) is the follow-up. +The **dense TP** recipe (`serve_qwen3-8b_sglang_tp.sh`) is **GPU-validated** on +**Qwen3-8B, `--tp-size 2`, 2× H200 (no IMEX/fabric)**, SAVE → SAVE → LOAD: + +- **No faults** — no MMU fault / Xid / illegal-memory-access / CUDA error in any + phase; NCCL initialises and the in-graph TP all-reduce replays over the + VMM-IPC bridge. +- **Deterministic** — per-rank `final_alloc_offset` identical across both SAVE + passes (`121831948288` ≈ 113.5 GiB on each rank), i.e. the allocation + trajectory LOAD replays is reproducible. +- **Fast cold start** — LOAD reaches a serving server in ~52 s vs ~172 s for a + cold SAVE capture (graph capture skipped, replaced by preallocate + replay). +- **Correct** — at temperature 0 the restored LOAD engine reproduces the warmed + SAVE engine's tokens **exactly** on every probe prompt. (A no-foundry baseline + diverges on some prompts — expected: `foundry_shim` disables FlashInfer + autotune + piecewise cuda-graphs, changing kernel selection vs upstream + defaults, so the correct reference is the foundry SAVE engine, which LOAD + matches bit-for-bit at the token level.) +- **Scratch** — `scratch_space_size = "1024MB"` is sufficient for dense TP here. + +SGLang build note: this was validated against the foundry-org/sglang integration +built on the pre-`init_forward_metadata`-ABC-rename base +(`6c69756fa8`, #26735's parent) with the foundry integration commit +(`76ac2f575b`) applied on top. The fork's `foundry` branch HEAD sits on a newer +sglang whose `FlashInferAttnBackend` renamed +`init_forward_metadata_capture_cuda_graph` → `init_forward_metadata_out_graph`, +which foundry 0.0.2's SGLang hook does not yet call — build the matching base +until the hook is ported. ## IPC-specific troubleshooting From 427978596e9d1a84245ab5f033dbe0fbe52fbe5f Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 23 Jul 2026 02:08:02 +0000 Subject: [PATCH 03/22] [sglang] port graph-capture hook to init_forward_metadata 3-method ABC (#26735) sglang PR #26735 removed init_forward_metadata_capture_cuda_graph in favor of init_forward_metadata_out_graph(forward_batch, in_capture) + _in_graph. Foundry's FlashInfer reuse-prepass shim and the LOAD-side metadata pre-pass called the removed method, so SAVE aborted with AttributeError on any post-rename sglang (all modes, not just TP). Detect the API by hasattr and drive the new out_graph(in_capture=True) path (building a ForwardBatch-like view for the pre-pass); the capture-time reuse shim now intercepts out_graph, sets forward_metadata for the pre-allocated wrappers, and re-plans via out_graph(in_capture=False) with no re-allocation. Falls back to the legacy method on pre-rename sglang, so no version pin is required. Co-Authored-By: Rahul Chalamala <22563365+rchalamala@users.noreply.github.com> Signed-off-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../foundry/integration/sglang/graph_ops.py | 40 ++++++++++++- python/foundry/integration/sglang/hooks.py | 60 ++++++++++++++++++- 2 files changed, 96 insertions(+), 4 deletions(-) diff --git a/python/foundry/integration/sglang/graph_ops.py b/python/foundry/integration/sglang/graph_ops.py index 18ab1f0e..9f2a98b4 100644 --- a/python/foundry/integration/sglang/graph_ops.py +++ b/python/foundry/integration/sglang/graph_ops.py @@ -248,6 +248,34 @@ def bootstrap_deepep_buffer(cuda_graph_runner) -> bool: return False +def build_capture_fb_view(cuda_graph_runner, bs: int): + """Build a ForwardBatch-like view for backend capture-side metadata init. + + Mirrors the padded per-bs inputs sglang's own capture loop feeds the + attention backend (``cuda_graph_runner.py::capture_one_batch_size``), as a + lightweight ``SimpleNamespace`` since the pre-pass runs before a real + ``ForwardBatch`` exists. Used to drive the post-#26735 3-method init ABC + (``init_forward_metadata_out_graph``). + """ + from types import SimpleNamespace + + buffers = cuda_graph_runner.buffers + num_tokens = bs * cuda_graph_runner.num_tokens_per_bs + encoder_lens = buffers.encoder_lens[:bs] if cuda_graph_runner.is_encoder_decoder else None + seq_lens = buffers.seq_lens[:bs] + return SimpleNamespace( + batch_size=bs, + forward_mode=cuda_graph_runner.capture_forward_mode, + req_pool_indices=buffers.req_pool_indices[:bs], + seq_lens=seq_lens, + seq_lens_cpu=buffers.seq_lens_cpu[:bs], + seq_lens_sum=int(seq_lens.sum().item()), + encoder_lens=encoder_lens, + spec_info=cuda_graph_runner.get_spec_info(num_tokens), + positions=buffers.positions[:num_tokens], + ) + + def initialize_attention_metadata_for_bs(cuda_graph_runner, bs: int) -> None: """Populate ``decode_cuda_graph_metadata[bs]`` for runtime replay. @@ -257,11 +285,21 @@ def initialize_attention_metadata_for_bs(cuda_graph_runner, bs: int) -> None: re-run the same call before runtime replay so the wrappers exist at deterministic VMM addresses. """ + attn_backend = cuda_graph_runner.attn_backend + # sglang #26735 replaced the single ``init_forward_metadata_capture_cuda_graph`` + # with a 3-method ABC; the capture-side allocation now lives in + # ``init_forward_metadata_out_graph(fb, in_capture=True)``. Prefer the new + # API, fall back to the legacy method for pre-rename sglang. + if hasattr(attn_backend, "init_forward_metadata_out_graph"): + fb_view = build_capture_fb_view(cuda_graph_runner, bs) + attn_backend.init_forward_metadata_out_graph(fb_view, in_capture=True) + return + buffers = cuda_graph_runner.buffers num_tokens = bs * cuda_graph_runner.num_tokens_per_bs encoder_lens = buffers.encoder_lens[:bs] if cuda_graph_runner.is_encoder_decoder else None spec_info = cuda_graph_runner.get_spec_info(num_tokens) - cuda_graph_runner.attn_backend.init_forward_metadata_capture_cuda_graph( + attn_backend.init_forward_metadata_capture_cuda_graph( bs, num_tokens, buffers.req_pool_indices[:bs], diff --git a/python/foundry/integration/sglang/hooks.py b/python/foundry/integration/sglang/hooks.py index 13f022a9..c15c55cc 100644 --- a/python/foundry/integration/sglang/hooks.py +++ b/python/foundry/integration/sglang/hooks.py @@ -472,7 +472,54 @@ 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 split the single capture-init entry point into the + # 3-method ABC; the per-iter capture path is now + # ``init_forward_metadata_out_graph(fb, in_capture=True)``. Detect + # which API this sglang exposes and shim the matching method. + use_new_api = hasattr(attn_backend, "init_forward_metadata_out_graph") + real_out_graph = 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): + # Only the capture path (in_capture=True) must avoid + # re-allocating the per-bs wrappers the pre-pass already built; + # replay / eager (in_capture=False) pass straight through. + if not in_capture: + return real_out_graph(forward_batch, in_capture=in_capture) + + from sglang.srt.layers.attention.flashinfer_backend import ( + DecodeMetadata, + PrefillMetadata, + ) + + bs = forward_batch.batch_size + forward_mode = forward_batch.forward_mode + if forward_mode.is_decode_or_idle(): + wrappers = 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-run the planner against the pre-allocated wrappers with + # in_capture=False: no ``_prepare_cuda_graph_metadata`` (so no + # second torch.empty for ``_int_workspace_buffer``) and no + # begin_forward re-install (the pre-pass did that). We set + # forward_metadata above since that path skips it. + real_out_graph(forward_batch, in_capture=False) def reuse_pre_pass_init( bs, @@ -586,11 +633,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 c95648aac08135c90e2f1f49d61490ef277bd90a Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 23 Jul 2026 03:14:43 +0000 Subject: [PATCH 04/22] [sglang] Run pre-capture warmup on all SAVE, not just EP Foundry skips sglang kernel_warmup and suppresses the pre-capture warmup forwards in SAVE for VMM-cursor determinism, and previously only ran its own neutered pre-capture warmup (_run_warmup_pass) for the DeepEP/EP path. On the dense path the model's first forward therefore ran *inside* the CUDA-graph capture; on Qwen3.5 that first forward lazily triggers torch.compile/inductor (vocab-parallel get_masked_input_and_mask -> inductor SFDP lazy_init), whose CPU<->CUDA copy is illegal mid-capture and aborts capture with 'Graph contains no nodes'. Run the warmup on all SAVE so lazy inductor/JIT init happens outside capture; LOAD still preallocates absolute offsets and never warms up, so the cursor stays symmetric. Co-Authored-By: Rahul Chalamala <22563365+rchalamala@users.noreply.github.com> Signed-off-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- python/foundry/integration/sglang/hooks.py | 50 +++++++++++----------- 1 file changed, 26 insertions(+), 24 deletions(-) diff --git a/python/foundry/integration/sglang/hooks.py b/python/foundry/integration/sglang/hooks.py index c15c55cc..5ffbf150 100644 --- a/python/foundry/integration/sglang/hooks.py +++ b/python/foundry/integration/sglang/hooks.py @@ -375,31 +375,33 @@ def _run_warmup_pass(self): def patched(self, *args, **kwargs): mode = get_graph_extension_mode() - # DeepEP/EP only (gated so the validated dense / single-GPU / DP paths - # are untouched): - # SAVE: sglang triggers pre-capture lazy init (DeepGEMM per-shape JIT, - # DeepEP buffer, ...) via warmup forwards that foundry suppresses; - # left lazy they fire inside the captured stream and abort - # ("operation not permitted when stream is capturing"). Run a real - # forward per bs up front so all that happens outside capture. - # LOAD: no capture happens — preallocate_for_load_mode reserves the - # whole region up to final_alloc_offset and replay places each alloc - # at its recorded absolute offset, so LOAD need NOT replay the warmup - # cursor trajectory. Running the warmup here is not just unnecessary - # but harmful: its orig_capture re-enters graph_capture() and leaves - # the context in a state that breaks the threaded finish_graph_loads - # ("invalid device context"). So warmup is SAVE-only. - # bootstrap_deepep_buffer runs on both (cheap singleton) to guarantee the - # NVSHMEM runtime is up before replay on LOAD / before capture on SAVE. - if _ep_lazy_init_needed(): - if mode == CUDAGraphExtensionMode.SAVE: - _run_warmup_pass(self) - if mode != CUDAGraphExtensionMode.NONE: - from foundry.integration.sglang.graph_ops import ( - bootstrap_deepep_buffer, - ) + # Pre-capture warmup (SAVE only): run one real forward per bs with graph + # capture neutered, so every lazy init sglang normally does in its + # (foundry-suppressed) warmup forwards + skipped ``kernel_warmup`` fires + # OUTSIDE the captured stream. Two independent paths need it: + # - EP: DeepGEMM per-shape JIT + DeepEP buffer creation, which left + # lazy fire inside the captured stream and abort ("operation not + # permitted when stream is capturing"). + # - any model whose forward lazily triggers torch.compile/inductor + # (e.g. Qwen3.5's vocab-parallel ``get_masked_input_and_mask``): + # inductor's one-time SFDP ``lazy_init`` does a CPU->CUDA copy that + # is illegal mid-capture and yields a "Graph contains no nodes" + # abort. + # LOAD never warms up: preallocate_for_load_mode reserves the whole + # region up to final_alloc_offset and replay places each alloc at its + # recorded absolute offset, so LOAD need NOT reproduce the warmup cursor + # trajectory — and running orig_capture there re-enters graph_capture() + # and breaks the threaded finish_graph_loads ("invalid device context"). + if mode == CUDAGraphExtensionMode.SAVE: + _run_warmup_pass(self) + # bootstrap_deepep_buffer runs on both SAVE and LOAD (cheap singleton) to + # guarantee the NVSHMEM runtime is up before replay / before capture. + if _ep_lazy_init_needed() and mode != CUDAGraphExtensionMode.NONE: + from foundry.integration.sglang.graph_ops import ( + bootstrap_deepep_buffer, + ) - bootstrap_deepep_buffer(self) + bootstrap_deepep_buffer(self) if mode == CUDAGraphExtensionMode.LOAD: from sglang.srt.distributed.device_communicators.pynccl_allocator import ( From 0dfe8ef52f72c55915728e41aa50bae75f7baf3d Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 23 Jul 2026 03:32:42 +0000 Subject: [PATCH 05/22] Revert "[sglang] Run pre-capture warmup on all SAVE, not just EP" This reverts commit c95648aac08135c90e2f1f49d61490ef277bd90a. --- python/foundry/integration/sglang/hooks.py | 50 +++++++++++----------- 1 file changed, 24 insertions(+), 26 deletions(-) diff --git a/python/foundry/integration/sglang/hooks.py b/python/foundry/integration/sglang/hooks.py index 5ffbf150..c15c55cc 100644 --- a/python/foundry/integration/sglang/hooks.py +++ b/python/foundry/integration/sglang/hooks.py @@ -375,33 +375,31 @@ def _run_warmup_pass(self): def patched(self, *args, **kwargs): mode = get_graph_extension_mode() - # Pre-capture warmup (SAVE only): run one real forward per bs with graph - # capture neutered, so every lazy init sglang normally does in its - # (foundry-suppressed) warmup forwards + skipped ``kernel_warmup`` fires - # OUTSIDE the captured stream. Two independent paths need it: - # - EP: DeepGEMM per-shape JIT + DeepEP buffer creation, which left - # lazy fire inside the captured stream and abort ("operation not - # permitted when stream is capturing"). - # - any model whose forward lazily triggers torch.compile/inductor - # (e.g. Qwen3.5's vocab-parallel ``get_masked_input_and_mask``): - # inductor's one-time SFDP ``lazy_init`` does a CPU->CUDA copy that - # is illegal mid-capture and yields a "Graph contains no nodes" - # abort. - # LOAD never warms up: preallocate_for_load_mode reserves the whole - # region up to final_alloc_offset and replay places each alloc at its - # recorded absolute offset, so LOAD need NOT reproduce the warmup cursor - # trajectory — and running orig_capture there re-enters graph_capture() - # and breaks the threaded finish_graph_loads ("invalid device context"). - if mode == CUDAGraphExtensionMode.SAVE: - _run_warmup_pass(self) - # bootstrap_deepep_buffer runs on both SAVE and LOAD (cheap singleton) to - # guarantee the NVSHMEM runtime is up before replay / before capture. - if _ep_lazy_init_needed() and mode != CUDAGraphExtensionMode.NONE: - from foundry.integration.sglang.graph_ops import ( - bootstrap_deepep_buffer, - ) + # DeepEP/EP only (gated so the validated dense / single-GPU / DP paths + # are untouched): + # SAVE: sglang triggers pre-capture lazy init (DeepGEMM per-shape JIT, + # DeepEP buffer, ...) via warmup forwards that foundry suppresses; + # left lazy they fire inside the captured stream and abort + # ("operation not permitted when stream is capturing"). Run a real + # forward per bs up front so all that happens outside capture. + # LOAD: no capture happens — preallocate_for_load_mode reserves the + # whole region up to final_alloc_offset and replay places each alloc + # at its recorded absolute offset, so LOAD need NOT replay the warmup + # cursor trajectory. Running the warmup here is not just unnecessary + # but harmful: its orig_capture re-enters graph_capture() and leaves + # the context in a state that breaks the threaded finish_graph_loads + # ("invalid device context"). So warmup is SAVE-only. + # bootstrap_deepep_buffer runs on both (cheap singleton) to guarantee the + # NVSHMEM runtime is up before replay on LOAD / before capture on SAVE. + if _ep_lazy_init_needed(): + if mode == CUDAGraphExtensionMode.SAVE: + _run_warmup_pass(self) + if mode != CUDAGraphExtensionMode.NONE: + from foundry.integration.sglang.graph_ops import ( + bootstrap_deepep_buffer, + ) - bootstrap_deepep_buffer(self) + bootstrap_deepep_buffer(self) if mode == CUDAGraphExtensionMode.LOAD: from sglang.srt.distributed.device_communicators.pynccl_allocator import ( From d36f647d0c8a64428b2d0975b022bae752127951 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 23 Jul 2026 03:47:03 +0000 Subject: [PATCH 06/22] [sglang] docs: port done + document Qwen3.5-122B hybrid MoE as unsupported Update the dense-TP README: the hook is ported to sglang's new init_forward_metadata 3-method ABC and validated on foundry HEAD (8B/TP=2 reproduces bit-for-bit), so the version-pin note is removed. Add a section documenting that Qwen3.5-122B-A10B (hybrid Mamba/GatedDeltaNet + MoE + multimodal) is not supported by this dense recipe, with the exact blockers (SAVE capture 'Graph contains no nodes' from lazy torch.compile inside capture; FlashInfer AllReduce Fusion bypassing the NCCL/VMM-IPC path; LOAD OOM from Mamba SSM state). Co-Authored-By: Rahul Chalamala <22563365+rchalamala@users.noreply.github.com> Signed-off-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- recipe/experimental/README.md | 43 ++++++++++++++++++++++++++++------- 1 file changed, 35 insertions(+), 8 deletions(-) diff --git a/recipe/experimental/README.md b/recipe/experimental/README.md index 1d781bd1..900f0c13 100644 --- a/recipe/experimental/README.md +++ b/recipe/experimental/README.md @@ -205,14 +205,41 @@ The **dense TP** recipe (`serve_qwen3-8b_sglang_tp.sh`) is **GPU-validated** on matches bit-for-bit at the token level.) - **Scratch** — `scratch_space_size = "1024MB"` is sufficient for dense TP here. -SGLang build note: this was validated against the foundry-org/sglang integration -built on the pre-`init_forward_metadata`-ABC-rename base -(`6c69756fa8`, #26735's parent) with the foundry integration commit -(`76ac2f575b`) applied on top. The fork's `foundry` branch HEAD sits on a newer -sglang whose `FlashInferAttnBackend` renamed -`init_forward_metadata_capture_cuda_graph` → `init_forward_metadata_out_graph`, -which foundry 0.0.2's SGLang hook does not yet call — build the matching base -until the hook is ported. +SGLang build note: this is validated against the foundry-org/sglang `foundry` +branch HEAD, which sits on a newer sglang whose `FlashInferAttnBackend` renamed +`init_forward_metadata_capture_cuda_graph` → the 3-method +`init_forward_metadata_out_graph` / `_in_graph` / `init_cuda_graph_state` ABC +(#26735). The foundry SGLang hook has been ported to that API (it still falls +back to the pre-rename method when present), so no sglang version pin is needed; +the same 8B/TP=2 result above reproduces bit-for-bit on the HEAD build. + +## Not supported: Qwen3.5-122B-A10B (hybrid Mamba/MoE) + +This dense-TP recipe is for **dense transformer** models. It does **not** work on +`Qwen/Qwen3.5-122B-A10B` (H200:4, `--tp-size 4`), which is a hybrid +Mamba/`Qwen3_5GatedDeltaNet` linear-attention + 256-expert MoE + multimodal +model. Vanilla sglang serves it fine (baseline is healthy), but foundry SAVE +does not: + +- **SAVE capture aborts with `Graph contains no nodes`.** On the dense path + foundry skips sglang's `kernel_warmup` and suppresses the pre-capture warmup + forwards (for VMM-cursor determinism), running its own pre-capture warmup only + for the DeepEP/EP path. So this model's *first* forward runs inside the CUDA + graph capture, and that first forward lazily triggers `torch.compile`/inductor + (vocab-parallel `get_masked_input_and_mask` → inductor SFDP `lazy_init`), whose + CPU↔CUDA copy is illegal mid-capture and yields an empty graph. (Naively + enabling the pre-capture warmup on the dense path lets capture succeed but + perturbs kernel/graph selection and **corrupts the SAVE↔LOAD correspondence** + — LOAD then emits garbage — so warmup suppression must stay.) +- **Other blockers seen even past capture:** sglang auto-enables *FlashInfer + AllReduce Fusion* for `Qwen3_5MoeForConditionalGeneration`, a fused all-reduce + that bypasses the plain-NCCL path this recipe routes over the VMM-IPC bridge; + and the ~24.6 GB/rank Mamba SSM state + KV + preallocated graph region makes + LOAD OOM (`rc=-9`) on H200:4. + +Supporting this model needs dedicated foundry work (a capture-safe warmup that +preserves SAVE↔LOAD kernel identity, Mamba-state handling, and routing/handling +the fused all-reduce through VMM-IPC) — out of scope for this dense-TP recipe. ## IPC-specific troubleshooting From 2807b6bcf959ada282dbb87f2399f20148941577 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 23 Jul 2026 05:34:58 +0000 Subject: [PATCH 07/22] [sglang] Cursor-neutral pre-capture warmup for dense SAVE (Qwen3.5 hybrid MoE) Run one real forward per capture_bs before graph capture on the dense/DP SAVE path, with the allocation region STOPPED and the caching allocator emptied afterwards, so lazy torch.compile/inductor + per-shape GEMM JIT happen outside the captured stream while the VMM cursor entering the FlashInfer pre-pass is byte-for-byte what LOAD sees. Fixes 'Graph contains no nodes' on models that defer compilation to first forward without the SAVE-only warmup drift that corrupts LOAD. Gated by FOUNDRY_SGLANG_SAVE_WARMUP (default on); EP path unchanged. Co-Authored-By: Rahul Chalamala <22563365+rchalamala@users.noreply.github.com> Signed-off-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- python/foundry/integration/sglang/hooks.py | 60 ++++++++++++++++++++++ 1 file changed, 60 insertions(+) diff --git a/python/foundry/integration/sglang/hooks.py b/python/foundry/integration/sglang/hooks.py index c15c55cc..fb5a9447 100644 --- a/python/foundry/integration/sglang/hooks.py +++ b/python/foundry/integration/sglang/hooks.py @@ -35,6 +35,16 @@ def _ep_lazy_init_needed() -> bool: return False +def _dense_save_warmup_enabled() -> bool: + """Whether the dense/DP SAVE path runs a cursor-neutral pre-capture warmup. + + Defaults on: models that defer torch.compile / per-shape GEMM JIT to their + first forward (e.g. Qwen3.5 hybrid MoE) would otherwise trigger that lazy + init inside CUDA-graph capture and abort. Set ``FOUNDRY_SGLANG_SAVE_WARMUP=0`` + to restore the old suppress-only behavior.""" + return os.environ.get("FOUNDRY_SGLANG_SAVE_WARMUP", "1") not in ("0", "false", "False") + + def _resolve_dp_rank(model_runner) -> int | None: dp_rank = getattr(model_runner, "dp_rank", None) if dp_rank is not None: @@ -371,6 +381,50 @@ def _run_warmup_pass(self): time.perf_counter() - t0, ) + def _run_dense_save_warmup(self): + """Cursor-neutral pre-capture warmup for the dense / DP FlashInfer path. + + Some models (e.g. Qwen3.5 hybrid MoE) defer torch.compile / inductor and + per-shape GEMM JIT to their first forward. On the foundry SAVE path that + first forward is *inside* CUDA-graph capture (sglang's own warmup forwards + are suppressed for VMM determinism), where a lazy CPU<->CUDA copy or a JIT + allocation aborts capture ("Graph contains no nodes"). Run one real forward + per capture_bs up front to force all that lazy init — but with the + allocation region STOPPED (allocations use the ordinary CUDA allocator and + the VMM cursor does not advance) and the caching allocator emptied after, + so the cursor entering the FlashInfer pre-pass is byte-for-byte what a + no-warmup run (i.e. LOAD) sees. This keeps SAVE↔LOAD wrapper offsets + aligned — a plain SAVE-only warmup drifts them and corrupts LOAD. The + per-bs FlashInfer wrappers the warmup populates are dropped so the real + pre-pass re-allocates them in-region. + """ + import torch + + from foundry import ops as fops + + attn_backend = self.attn_backend + has_decode = hasattr(attn_backend, "decode_cuda_graph_metadata") + has_prefill = hasattr(attn_backend, "prefill_cuda_graph_metadata") + if has_decode: + attn_backend.decode_cuda_graph_metadata = {} + if has_prefill: + attn_backend.prefill_cuda_graph_metadata = {} + rt.log_alloc_offset("save_before_warmup") + fops.stop_allocation_region() + try: + _run_warmup_pass(self) + finally: + fops.resume_allocation_region() + attn_backend.forward_metadata = None + # Drop the out-of-region wrappers the warmup populated so the real + # pre-pass re-allocates them inside the VMM region. + if has_decode: + attn_backend.decode_cuda_graph_metadata = {} + if has_prefill: + attn_backend.prefill_cuda_graph_metadata = {} + torch.cuda.empty_cache() + rt.log_alloc_offset("save_after_warmup") + @functools.wraps(orig_capture) def patched(self, *args, **kwargs): mode = get_graph_extension_mode() @@ -462,6 +516,12 @@ def patched(self, *args, **kwargs): return None if mode == CUDAGraphExtensionMode.SAVE: + # Force model-level lazy init (torch.compile/inductor, per-shape GEMM + # JIT) outside capture for the dense/DP path, cursor-neutrally so the + # FlashInfer pre-pass below still lands at LOAD's offsets. EP already + # warmed up above; skip there. + if not _ep_lazy_init_needed() and _dense_save_warmup_enabled(): + _run_dense_save_warmup(self) # FlashInfer allocates a per-bs metadata wrapper (each with its own # _int_workspace_buffer) on every capture init, so foundry pre-allocates # them up front and installs a reuse shim that makes the inner init From 2b6f25df848d7d697abaa565703f3f37ad3dccc1 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 23 Jul 2026 06:43:04 +0000 Subject: [PATCH 08/22] [sglang] Route FlashInfer prepass through hybrid linear-attn wrapper Qwen3.5-style hybrid models nest FlashInfer inside a HybridLinearAttnBackend (FlashInfer full-attn layers + mamba/linear-attn layers). Foundry detected FlashInfer by indices_updater_decode on the top-level attn_backend, which the hybrid wrapper does not expose, so the per-bs FlashInfer wrapper prepass was skipped: SAVE allocated wrappers interleaved with graph memory during capture while LOAD allocated them after load_all_graphs, diverging the VMM offsets and making the decode graph replay hit an illegal memory access on LOAD. Add _flashinfer_backend() to resolve the FlashInfer backend through the hybrid wrapper's full_attn_backend, and install the reuse shim / run the prepass on it. The mamba backend keeps all its graph buffers in init_cuda_graph_state (pre-capture, already SAVE/LOAD-symmetric), so no extra handling is needed there. Signed-off-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-Authored-By: Rahul Chalamala <22563365+rchalamala@users.noreply.github.com> Signed-off-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- python/foundry/integration/sglang/hooks.py | 118 +++++++++++++++------ 1 file changed, 85 insertions(+), 33 deletions(-) diff --git a/python/foundry/integration/sglang/hooks.py b/python/foundry/integration/sglang/hooks.py index fb5a9447..ab18653f 100644 --- a/python/foundry/integration/sglang/hooks.py +++ b/python/foundry/integration/sglang/hooks.py @@ -35,6 +35,29 @@ def _ep_lazy_init_needed() -> bool: return False +def _flashinfer_backend(attn_backend): + """Return the FlashInfer backend that owns the per-bs capture wrappers, or + None if this attention backend doesn't use FlashInfer's per-bs prepass. + + FlashInfer allocates a fresh per-bs metadata wrapper (each with its own + ``_int_workspace_buffer``) on every capture init, so foundry pre-allocates + them all up front and installs a reuse shim to keep the VMM cursor + deterministic vs LOAD. Hybrid linear-attn / mamba models (e.g. Qwen3.5) + wrap FlashInfer inside a ``HybridLinearAttnBackend`` whose full-attention + layers use FlashInfer while its linear-attn (mamba) layers keep all their + graph buffers in ``init_cuda_graph_state`` (pre-capture, so already + SAVE/LOAD-symmetric). Detect the FlashInfer backend by its per-bs + ``indices_updater_decode``, seeing through the hybrid wrapper's + ``full_attn_backend`` so the wrappers get the deterministic prepass too.""" + if hasattr(attn_backend, "indices_updater_decode"): + return attn_backend + if hasattr(attn_backend, "full_attn_backend"): + inner = attn_backend.full_attn_backend + if hasattr(inner, "indices_updater_decode"): + return inner + return None + + def _dense_save_warmup_enabled() -> bool: """Whether the dense/DP SAVE path runs a cursor-neutral pre-capture warmup. @@ -402,26 +425,32 @@ def _run_dense_save_warmup(self): from foundry import ops as fops - attn_backend = self.attn_backend - has_decode = hasattr(attn_backend, "decode_cuda_graph_metadata") - has_prefill = hasattr(attn_backend, "prefill_cuda_graph_metadata") + # The per-bs FlashInfer wrappers the warmup populates live on the + # FlashInfer backend — which is ``self.attn_backend.full_attn_backend`` + # for hybrid linear-attn / mamba models, not the wrapper itself. Clear + # them there so the real (in-region) pre-pass re-allocates them. + fi_backend = _flashinfer_backend(self.attn_backend) + has_decode = fi_backend is not None and hasattr(fi_backend, "decode_cuda_graph_metadata") + has_prefill = fi_backend is not None and hasattr(fi_backend, "prefill_cuda_graph_metadata") if has_decode: - attn_backend.decode_cuda_graph_metadata = {} + fi_backend.decode_cuda_graph_metadata = {} if has_prefill: - attn_backend.prefill_cuda_graph_metadata = {} + fi_backend.prefill_cuda_graph_metadata = {} rt.log_alloc_offset("save_before_warmup") fops.stop_allocation_region() try: _run_warmup_pass(self) finally: fops.resume_allocation_region() - attn_backend.forward_metadata = None + self.attn_backend.forward_metadata = None + if fi_backend is not None: + fi_backend.forward_metadata = None # Drop the out-of-region wrappers the warmup populated so the real # pre-pass re-allocates them inside the VMM region. if has_decode: - attn_backend.decode_cuda_graph_metadata = {} + fi_backend.decode_cuda_graph_metadata = {} if has_prefill: - attn_backend.prefill_cuda_graph_metadata = {} + fi_backend.prefill_cuda_graph_metadata = {} torch.cuda.empty_cache() rt.log_alloc_offset("save_after_warmup") @@ -486,7 +515,13 @@ def patched(self, *args, **kwargs): # ``start_base_addr_0`` when graph load begins. # FlashInfer-only pre-pass (see SAVE branch). fa3 etc. allocate their # cuda-graph metadata once in init_cuda_graph_state, so skip it. - if hasattr(self.attn_backend, "indices_updater_decode"): + # ``_flashinfer_backend`` sees through the hybrid linear-attn wrapper + # (Qwen3.5) so its FlashInfer full-attn wrappers get pre-allocated + # here — BEFORE load_all_graphs — landing at SAVE's offsets. Its + # mamba layers keep all buffers in init_cuda_graph_state, so no + # divergent alloc there. + fi_backend = _flashinfer_backend(self.attn_backend) + if fi_backend is not None: initialize_all_attention_metadata(self) rt.log_alloc_offset("after_pre_init") # Single ``start_graph_builds(all_paths)`` call so templates @@ -504,7 +539,7 @@ def patched(self, *args, **kwargs): # now. Run AFTER load_all_graphs so the (already-correct) loaded-graph # VMM offsets are unaffected — fa3's metadata are lightweight views # over the fixed init_cuda_graph_state workspace, not graph memory. - if not hasattr(self.attn_backend, "indices_updater_decode"): + if fi_backend is None: initialize_all_attention_metadata(self) # Initialize the DeepEP cuda-graph adapter's captured mode. Upstream # sets this in deepep_adapter.capture() during the capture loop, @@ -528,18 +563,32 @@ def patched(self, *args, **kwargs): # reuse them — keeping the VMM cursor deterministic vs LOAD. Backends # with a single fixed cuda-graph metadata workspace allocated once in # init_cuda_graph_state (e.g. fa3 / FlashAttentionBackend) don't need - # this and the plain capture is already SAVE/LOAD-deterministic. Detect - # FlashInfer by its per-bs indices_updater_decode. + # this and the plain capture is already SAVE/LOAD-deterministic. + # ``_flashinfer_backend`` returns the FlashInfer backend, seeing + # through the hybrid linear-attn wrapper (Qwen3.5) whose FlashInfer + # full-attn layers own the per-bs wrappers while its mamba layers + # keep all buffers in init_cuda_graph_state. The shim is installed on + # that FlashInfer backend directly; the hybrid wrapper's own + # init_forward_metadata_out_graph then dispatches into it during + # capture, so the mamba backend still runs its normal (alloc-free) + # metadata copy. attn_backend = self.attn_backend - use_fi_prepass = hasattr(attn_backend, "indices_updater_decode") + fi_backend = _flashinfer_backend(attn_backend) + use_fi_prepass = fi_backend is not None # sglang #26735 split the single capture-init entry point into the # 3-method ABC; the per-iter capture path is now # ``init_forward_metadata_out_graph(fb, in_capture=True)``. Detect # which API this sglang exposes and shim the matching method. use_new_api = hasattr(attn_backend, "init_forward_metadata_out_graph") - real_out_graph = attn_backend.init_forward_metadata_out_graph if use_new_api else None + real_out_graph = ( + fi_backend.init_forward_metadata_out_graph + if use_new_api and use_fi_prepass + else None + ) real_init = ( - None if use_new_api else attn_backend.init_forward_metadata_capture_cuda_graph + fi_backend.init_forward_metadata_capture_cuda_graph + if not use_new_api and use_fi_prepass + else None ) def reuse_out_graph(forward_batch, in_capture=False): @@ -557,19 +606,19 @@ def reuse_out_graph(forward_batch, in_capture=False): 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) + wrappers = fi_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) + fi_backend.forward_metadata = DecodeMetadata(wrappers) elif ( forward_mode.is_target_verify() or forward_mode.is_draft_extend() or forward_mode.is_dllm_extend() ): - wrappers = attn_backend.prefill_cuda_graph_metadata.get(bs) + wrappers = fi_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( + fi_backend.forward_metadata = PrefillMetadata( wrappers, forward_mode.is_dllm_extend(), False ) else: @@ -607,7 +656,7 @@ def reuse_pre_pass_init( ) if forward_mode.is_decode_or_idle(): - wrappers = attn_backend.decode_cuda_graph_metadata.get(bs) + wrappers = fi_backend.decode_cuda_graph_metadata.get(bs) if wrappers is None: return real_init( bs, @@ -619,7 +668,7 @@ def reuse_pre_pass_init( spec_info, ) seq_lens_sum = seq_lens.sum().item() - attn_backend.indices_updater_decode.update( + fi_backend.indices_updater_decode.update( req_pool_indices, seq_lens, seq_lens.cpu(), @@ -628,16 +677,16 @@ def reuse_pre_pass_init( encoder_lens=encoder_lens, spec_info=spec_info, fixed_split_size=None, - disable_split_kv=attn_backend.disable_cuda_graph_kv_split, + disable_split_kv=fi_backend.disable_cuda_graph_kv_split, ) - attn_backend.forward_metadata = DecodeMetadata(wrappers) + fi_backend.forward_metadata = DecodeMetadata(wrappers) return if ( forward_mode.is_target_verify() or forward_mode.is_draft_extend() or forward_mode.is_dllm_extend() ): - wrappers = attn_backend.prefill_cuda_graph_metadata.get(bs) + wrappers = fi_backend.prefill_cuda_graph_metadata.get(bs) if wrappers is None: return real_init( bs, @@ -651,12 +700,12 @@ def reuse_pre_pass_init( seq_lens_sum = seq_lens.sum().item() use_ragged = forward_mode.is_dllm_extend() prefix_lens = ( - seq_lens - attn_backend.dllm_config.block_size + seq_lens - fi_backend.dllm_config.block_size if forward_mode.is_dllm_extend() else None ) spec_info_arg = None if forward_mode.is_dllm_extend() else spec_info - attn_backend.indices_updater_prefill.update( + fi_backend.indices_updater_prefill.update( req_pool_indices, seq_lens, seq_lens.cpu(), @@ -667,7 +716,7 @@ def reuse_pre_pass_init( encoder_lens=encoder_lens, spec_info=spec_info_arg, ) - attn_backend.forward_metadata = PrefillMetadata(wrappers, use_ragged, False) + fi_backend.forward_metadata = PrefillMetadata(wrappers, use_ragged, False) return # Unknown mode — fall back to real init. return real_init( @@ -691,20 +740,23 @@ def reuse_pre_pass_init( initialize_all_attention_metadata(self) rt.log_alloc_offset("save_after_pre_init") # Drop the pre-pass's last forward_metadata ref so popping the dict - # entry doesn't keep the wrapper alive at refcount 1. - attn_backend.forward_metadata = None + # entry doesn't keep the wrapper alive at refcount 1. Installed on + # the FlashInfer backend (== attn_backend, or its full_attn_backend + # for the hybrid wrapper); the hybrid's own out_graph dispatches + # into the shimmed FlashInfer method during capture. + fi_backend.forward_metadata = None if use_new_api: - attn_backend.init_forward_metadata_out_graph = reuse_out_graph + fi_backend.init_forward_metadata_out_graph = reuse_out_graph else: - attn_backend.init_forward_metadata_capture_cuda_graph = reuse_pre_pass_init + fi_backend.init_forward_metadata_capture_cuda_graph = reuse_pre_pass_init try: result = orig_capture(self, *args, **kwargs) finally: if use_fi_prepass: if use_new_api: - attn_backend.init_forward_metadata_out_graph = real_out_graph + fi_backend.init_forward_metadata_out_graph = real_out_graph else: - attn_backend.init_forward_metadata_capture_cuda_graph = real_init + fi_backend.init_forward_metadata_capture_cuda_graph = real_init from foundry.integration.sglang.graph_ops import ( pack_fatbins, From 7403994600f2c685f8fac3d56ebc5182a937f413 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 23 Jul 2026 07:44:06 +0000 Subject: [PATCH 09/22] [sglang] In-region warmup on both SAVE and LOAD for hybrid model graph replay Models with lazy-init persistent buffers (Qwen3.5 MoE + mamba) need the warmup to run with the VMM allocation region ACTIVE so those buffers land at deterministic in-region addresses. The previous cursor-neutral approach (region stopped during warmup) put them out-of-region, which worked for SAVE (same process) but failed on LOAD: the graph replay referenced stale out-of-region addresses from a different process. Run the same in-region warmup on both SAVE and LOAD so the VMM cursor trajectory matches, then clear FlashInfer metadata and empty the caching allocator cache before the prepass/capture proceeds. This makes all lazy-init model buffers (torch.compile codegen, Triton workspace, etc.) available at the same in-region addresses on LOAD as on SAVE. Co-Authored-By: Rahul Chalamala <22563365+rchalamala@users.noreply.github.com> Signed-off-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- python/foundry/integration/sglang/hooks.py | 81 ++++++++++++---------- 1 file changed, 44 insertions(+), 37 deletions(-) diff --git a/python/foundry/integration/sglang/hooks.py b/python/foundry/integration/sglang/hooks.py index ab18653f..5b7b29f4 100644 --- a/python/foundry/integration/sglang/hooks.py +++ b/python/foundry/integration/sglang/hooks.py @@ -404,31 +404,30 @@ def _run_warmup_pass(self): time.perf_counter() - t0, ) - def _run_dense_save_warmup(self): - """Cursor-neutral pre-capture warmup for the dense / DP FlashInfer path. - - Some models (e.g. Qwen3.5 hybrid MoE) defer torch.compile / inductor and - per-shape GEMM JIT to their first forward. On the foundry SAVE path that - first forward is *inside* CUDA-graph capture (sglang's own warmup forwards - are suppressed for VMM determinism), where a lazy CPU<->CUDA copy or a JIT - allocation aborts capture ("Graph contains no nodes"). Run one real forward - per capture_bs up front to force all that lazy init — but with the - allocation region STOPPED (allocations use the ordinary CUDA allocator and - the VMM cursor does not advance) and the caching allocator emptied after, - so the cursor entering the FlashInfer pre-pass is byte-for-byte what a - no-warmup run (i.e. LOAD) sees. This keeps SAVE↔LOAD wrapper offsets - aligned — a plain SAVE-only warmup drifts them and corrupts LOAD. The - per-bs FlashInfer wrappers the warmup populates are dropped so the real - pre-pass re-allocates them in-region. + def _run_dense_warmup(self): + """In-region pre-capture warmup for the dense / DP path. + + Some models (e.g. Qwen3.5 hybrid MoE) defer torch.compile / inductor, + per-shape GEMM JIT, or Triton kernel compilation to their first forward. + On the foundry SAVE path that first forward is *inside* CUDA-graph capture + (sglang's own warmup forwards are suppressed for VMM determinism), where a + lazy CPU<->CUDA copy aborts capture ("Graph contains no nodes"). More + critically, hybrid models (MoE + mamba) lazily initialize persistent + module-level GPU buffers (e.g. Triton codegen workspace, torch.compile + inductor intermediates) on the first real forward — if those buffers land + outside the VMM region, the CUDA graph captures their out-of-region + addresses, which are stale on LOAD (different process → different addresses). + + Run one real forward per capture_bs with the allocation region ACTIVE so all + lazy-init persistent buffers land IN-REGION at deterministic VMM addresses. + The cursor advances during warmup; LOAD runs the same warmup (same model, + same forward, same capture_bs) to produce the same cursor trajectory. After + warmup, clear FlashInfer per-bs wrappers and empty the caching allocator + cache so the subsequent FlashInfer pre-pass re-allocates wrappers at the + post-warmup cursor — identical on both SAVE and LOAD. """ import torch - from foundry import ops as fops - - # The per-bs FlashInfer wrappers the warmup populates live on the - # FlashInfer backend — which is ``self.attn_backend.full_attn_backend`` - # for hybrid linear-attn / mamba models, not the wrapper itself. Clear - # them there so the real (in-region) pre-pass re-allocates them. fi_backend = _flashinfer_backend(self.attn_backend) has_decode = fi_backend is not None and hasattr(fi_backend, "decode_cuda_graph_metadata") has_prefill = fi_backend is not None and hasattr(fi_backend, "prefill_cuda_graph_metadata") @@ -436,23 +435,21 @@ def _run_dense_save_warmup(self): fi_backend.decode_cuda_graph_metadata = {} if has_prefill: fi_backend.prefill_cuda_graph_metadata = {} - rt.log_alloc_offset("save_before_warmup") - fops.stop_allocation_region() + rt.log_alloc_offset("before_warmup") try: _run_warmup_pass(self) finally: - fops.resume_allocation_region() self.attn_backend.forward_metadata = None if fi_backend is not None: fi_backend.forward_metadata = None - # Drop the out-of-region wrappers the warmup populated so the real - # pre-pass re-allocates them inside the VMM region. + # Drop warmup's FlashInfer wrappers so the real pre-pass + # re-allocates them at the post-warmup cursor. if has_decode: fi_backend.decode_cuda_graph_metadata = {} if has_prefill: fi_backend.prefill_cuda_graph_metadata = {} torch.cuda.empty_cache() - rt.log_alloc_offset("save_after_warmup") + rt.log_alloc_offset("after_warmup") @functools.wraps(orig_capture) def patched(self, *args, **kwargs): @@ -467,11 +464,11 @@ def patched(self, *args, **kwargs): # forward per bs up front so all that happens outside capture. # LOAD: no capture happens — preallocate_for_load_mode reserves the # whole region up to final_alloc_offset and replay places each alloc - # at its recorded absolute offset, so LOAD need NOT replay the warmup - # cursor trajectory. Running the warmup here is not just unnecessary - # but harmful: its orig_capture re-enters graph_capture() and leaves - # the context in a state that breaks the threaded finish_graph_loads - # ("invalid device context"). So warmup is SAVE-only. + # at its recorded absolute offset. EP warmup is SAVE-only (its + # orig_capture re-enters graph_capture() and breaks the threaded + # finish_graph_loads with "invalid device context"). Dense warmup + # runs on BOTH SAVE and LOAD (in the LOAD branch below) so lazy-init + # persistent buffers land at deterministic in-region addresses. # bootstrap_deepep_buffer runs on both (cheap singleton) to guarantee the # NVSHMEM runtime is up before replay on LOAD / before capture on SAVE. if _ep_lazy_init_needed(): @@ -503,6 +500,14 @@ def patched(self, *args, **kwargs): if cgr.get_global_graph_memory_pool() is None: cgr.set_global_graph_memory_pool(self.device_module.graph_pool_handle()) set_graph_pool_id(cgr.get_global_graph_memory_pool()) + # Run the same in-region warmup as SAVE so lazy-init model + # buffers (torch.compile codegen, Triton kernel workspace, + # MoE/mamba persistent state) land at the same VMM addresses. + # Without this, SAVE's graph captures out-of-region pointers to + # buffers that were lazily allocated during SAVE's warmup — stale + # on LOAD (different process). + if not _ep_lazy_init_needed() and _dense_save_warmup_enabled(): + _run_dense_warmup(self) rt.log_alloc_offset("before_preallocate") rt.preallocate_for_load_mode() rt.log_alloc_offset("after_preallocate") @@ -552,11 +557,13 @@ def patched(self, *args, **kwargs): if mode == CUDAGraphExtensionMode.SAVE: # Force model-level lazy init (torch.compile/inductor, per-shape GEMM - # JIT) outside capture for the dense/DP path, cursor-neutrally so the - # FlashInfer pre-pass below still lands at LOAD's offsets. EP already - # warmed up above; skip there. + # JIT) outside capture for the dense/DP path. The warmup runs + # in-region so lazy-init persistent model buffers land at + # deterministic VMM addresses; LOAD runs the same warmup to + # reproduce the cursor trajectory. EP already warmed up above; + # skip there. if not _ep_lazy_init_needed() and _dense_save_warmup_enabled(): - _run_dense_save_warmup(self) + _run_dense_warmup(self) # FlashInfer allocates a per-bs metadata wrapper (each with its own # _int_workspace_buffer) on every capture init, so foundry pre-allocates # them up front and installs a reuse shim that makes the inner init From dfdddad83539abc22c286a477a7f7be2e3727ddf Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 23 Jul 2026 08:26:47 +0000 Subject: [PATCH 10/22] [sglang] Update README: Qwen3.5-122B-A10B GPU-validated on H200:4 TP=4 16/16 exact SAVE==LOAD temperature-0 token match, zero faults, byte-identical SAVE/LOAD VMM offset trajectory, SAVE offsets reproducible across passes on all 4 ranks. Co-Authored-By: Rahul Chalamala <22563365+rchalamala@users.noreply.github.com> Signed-off-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- recipe/experimental/README.md | 69 +++++++++++++++++++++-------------- 1 file changed, 41 insertions(+), 28 deletions(-) diff --git a/recipe/experimental/README.md b/recipe/experimental/README.md index 900f0c13..6730ac64 100644 --- a/recipe/experimental/README.md +++ b/recipe/experimental/README.md @@ -193,7 +193,7 @@ The **dense TP** recipe (`serve_qwen3-8b_sglang_tp.sh`) is **GPU-validated** on phase; NCCL initialises and the in-graph TP all-reduce replays over the VMM-IPC bridge. - **Deterministic** — per-rank `final_alloc_offset` identical across both SAVE - passes (`121831948288` ≈ 113.5 GiB on each rank), i.e. the allocation + passes (e.g. `120517033984` on each rank), i.e. the allocation trajectory LOAD replays is reproducible. - **Fast cold start** — LOAD reaches a serving server in ~52 s vs ~172 s for a cold SAVE capture (graph capture skipped, replaced by preallocate + replay). @@ -213,33 +213,46 @@ branch HEAD, which sits on a newer sglang whose `FlashInferAttnBackend` renamed back to the pre-rename method when present), so no sglang version pin is needed; the same 8B/TP=2 result above reproduces bit-for-bit on the HEAD build. -## Not supported: Qwen3.5-122B-A10B (hybrid Mamba/MoE) - -This dense-TP recipe is for **dense transformer** models. It does **not** work on -`Qwen/Qwen3.5-122B-A10B` (H200:4, `--tp-size 4`), which is a hybrid -Mamba/`Qwen3_5GatedDeltaNet` linear-attention + 256-expert MoE + multimodal -model. Vanilla sglang serves it fine (baseline is healthy), but foundry SAVE -does not: - -- **SAVE capture aborts with `Graph contains no nodes`.** On the dense path - foundry skips sglang's `kernel_warmup` and suppresses the pre-capture warmup - forwards (for VMM-cursor determinism), running its own pre-capture warmup only - for the DeepEP/EP path. So this model's *first* forward runs inside the CUDA - graph capture, and that first forward lazily triggers `torch.compile`/inductor - (vocab-parallel `get_masked_input_and_mask` → inductor SFDP `lazy_init`), whose - CPU↔CUDA copy is illegal mid-capture and yields an empty graph. (Naively - enabling the pre-capture warmup on the dense path lets capture succeed but - perturbs kernel/graph selection and **corrupts the SAVE↔LOAD correspondence** - — LOAD then emits garbage — so warmup suppression must stay.) -- **Other blockers seen even past capture:** sglang auto-enables *FlashInfer - AllReduce Fusion* for `Qwen3_5MoeForConditionalGeneration`, a fused all-reduce - that bypasses the plain-NCCL path this recipe routes over the VMM-IPC bridge; - and the ~24.6 GB/rank Mamba SSM state + KV + preallocated graph region makes - LOAD OOM (`rc=-9`) on H200:4. - -Supporting this model needs dedicated foundry work (a capture-safe warmup that -preserves SAVE↔LOAD kernel identity, Mamba-state handling, and routing/handling -the fused all-reduce through VMM-IPC) — out of scope for this dense-TP recipe. +## Qwen3.5-122B-A10B (hybrid Mamba/MoE) — GPU-validated + +The dense-TP recipe also works on `Qwen/Qwen3.5-122B-A10B`, a hybrid +Mamba/`Qwen3_5GatedDeltaNet` linear-attention + 256-expert MoE model. +GPU-validated on **H200:4, `--tp-size 4`** (no IMEX/fabric): + +- **No faults** — zero MMU fault / Xid / illegal-memory-access / CUDA error in + any phase (baseline, SAVE, SAVE2, LOAD). +- **Deterministic** — per-rank `final_alloc_offset` identical across both SAVE + passes on all 4 ranks (`120376524800`). +- **Fast cold start** — LOAD serves in ~127 s vs ~249 s cold SAVE capture. +- **Correct** — at temperature 0 the restored LOAD engine reproduces the warmed + SAVE engine's tokens **exactly** on all 16 probe prompts. +- **Scratch** — `scratch_space_size = "1024MB"` sufficient. + +Byte-identical SAVE/LOAD VMM offset trajectory (TP0): + +| Checkpoint | SAVE | LOAD | +|---|---|---| +| after\_init\_memory\_pool | 119676076032 | 119676076032 | +| before\_warmup | 120120672256 | 120120672256 | +| after\_warmup | 120271667200 | 120271667200 | +| after\_pre\_init | 120313610240 | 120313610240 | +| final / after\_load\_all\_graphs | 120376524800 | 120376524800 | + +Key model-specific requirements: + +- `--enforce-disable-flashinfer-allreduce-fusion` — Qwen3.5 auto-enables + FlashInfer AllReduce Fusion on SM90, which bypasses the NCCL-over-VMM-IPC + path. The flag keeps all-reduce on plain NCCL. +- `--max-mamba-cache-size N` — pin explicitly (matching + `--max-running-requests`). Foundry LOAD bypasses sglang's memory profiling so + `max_mamba_cache_size` defaults to None without the flag, crashing MambaPool. +- `--disable-radix-cache` — required for hybrid (Mamba) models. +- In-region warmup — the foundry hook runs a pre-capture warmup with the VMM + region active so torch.compile codegen, Triton kernel workspace, and + MoE/mamba persistent buffers land at deterministic in-region addresses (the + same warmup runs on LOAD, producing the same cursor trajectory). Without this, + lazy-init model buffers allocated during warmup land out-of-region and cause + `CUDA error: illegal memory access` on LOAD graph replay. ## IPC-specific troubleshooting From e2dbdf2d686182abf38701fa443e45d2287ca881 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 23 Jul 2026 08:50:15 +0000 Subject: [PATCH 11/22] [sglang] Optional LOAD warmup skip: reserve warmup alloc gap instead of re-running forwards MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The dense warmup on LOAD re-runs model forwards (~28s on Qwen3.5-122B) purely to place lazy-init persistent buffers in-region at the addresses the decode graphs reference. When those buffers are write-before-read scratch within the captured graph, LOAD can skip the forwards and simply reserve the SAVE-side warmup allocation gap (recorded in warmup_alloc_offset.json), leaving it mapped but uninitialized. Gated behind FOUNDRY_SGLANG_SKIP_LOAD_WARMUP=1 (off by default) since correctness depends on the scratch-only assumption — must be validated with exact SAVE==LOAD token equality per model. Co-Authored-By: Rahul Chalamala <22563365+rchalamala@users.noreply.github.com> Signed-off-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- python/foundry/integration/sglang/hooks.py | 36 +++++++++++++++++++- python/foundry/integration/sglang/runtime.py | 24 +++++++++++++ 2 files changed, 59 insertions(+), 1 deletion(-) diff --git a/python/foundry/integration/sglang/hooks.py b/python/foundry/integration/sglang/hooks.py index 5b7b29f4..0d4b84fc 100644 --- a/python/foundry/integration/sglang/hooks.py +++ b/python/foundry/integration/sglang/hooks.py @@ -68,6 +68,19 @@ def _dense_save_warmup_enabled() -> bool: return os.environ.get("FOUNDRY_SGLANG_SAVE_WARMUP", "1") not in ("0", "false", "False") +def _load_warmup_skip_enabled() -> bool: + """Whether LOAD skips re-running the dense warmup forwards and instead just + reserves the SAVE-side warmup allocation gap (``warmup_alloc_offset.json``). + + Off by default. The warmup places ~model-specific persistent buffers + in-region; re-running it on LOAD reproduces them at the recorded addresses. + Skipping it (and leaving that gap mapped-but-uninitialized) is only correct + when those buffers are write-before-read scratch within the captured decode + graph — set ``FOUNDRY_SGLANG_SKIP_LOAD_WARMUP=1`` to try it, but validate + exact SAVE==LOAD token equality before relying on it.""" + return os.environ.get("FOUNDRY_SGLANG_SKIP_LOAD_WARMUP", "0") not in ("0", "false", "False") + + def _resolve_dp_rank(model_runner) -> int | None: dp_rank = getattr(model_runner, "dp_rank", None) if dp_rank is not None: @@ -450,6 +463,9 @@ def _run_dense_warmup(self): fi_backend.prefill_cuda_graph_metadata = {} torch.cuda.empty_cache() rt.log_alloc_offset("after_warmup") + # Record the post-warmup cursor so LOAD can optionally reserve this + # gap instead of re-running the warmup (see _load_warmup_skip_enabled). + rt.save_warmup_alloc_offset() @functools.wraps(orig_capture) def patched(self, *args, **kwargs): @@ -506,11 +522,29 @@ def patched(self, *args, **kwargs): # Without this, SAVE's graph captures out-of-region pointers to # buffers that were lazily allocated during SAVE's warmup — stale # on LOAD (different process). - if not _ep_lazy_init_needed() and _dense_save_warmup_enabled(): + warmup_needed = not _ep_lazy_init_needed() and _dense_save_warmup_enabled() + skip_warmup = warmup_needed and _load_warmup_skip_enabled() + if warmup_needed and not skip_warmup: _run_dense_warmup(self) rt.log_alloc_offset("before_preallocate") rt.preallocate_for_load_mode() rt.log_alloc_offset("after_preallocate") + if skip_warmup: + # Instead of re-running the warmup forwards, jump the cursor + # over the SAVE-side warmup gap. preallocate_for_load_mode + # already mapped the whole region (including this gap) from the + # current cursor to final_alloc_offset, so the decode graphs' + # workspace pointers land on mapped (if uninitialized) memory. + warmup_off = rt.load_warmup_alloc_offset() + if warmup_off <= 0: + raise RuntimeError( + "FOUNDRY_SGLANG_SKIP_LOAD_WARMUP=1 but no " + "warmup_alloc_offset.json in the archive" + ) + from foundry import ops as fops + + fops.set_current_alloc_offset(warmup_off) + rt.log_alloc_offset("after_warmup_skip") # Pre-pass: allocate every per-bs FlashInfer wrapper in # ``reversed(capture_bs)`` order, matching the order SAVE used. # SAVE's patched ``init_forward_metadata_capture_cuda_graph`` diff --git a/python/foundry/integration/sglang/runtime.py b/python/foundry/integration/sglang/runtime.py index 1f9032c5..b781ae7f 100644 --- a/python/foundry/integration/sglang/runtime.py +++ b/python/foundry/integration/sglang/runtime.py @@ -180,6 +180,30 @@ def capture_final_alloc_offset() -> int: return _final_alloc_offset +def save_warmup_alloc_offset() -> None: + """Persist the VMM cursor immediately after the dense warmup so LOAD can + reserve the same gap without re-running the warmup forwards.""" + cfg = get_config() + if cfg is None or cfg.mode != CUDAGraphExtensionMode.SAVE or cfg.workspace_dir is None: + return + offset = cge.get_current_alloc_offset() + path = os.path.join(cfg.workspace_dir, "warmup_alloc_offset.json") + with open(path, "w") as f: + json.dump({"warmup_alloc_offset": offset}, f) + + +def load_warmup_alloc_offset() -> int: + """Return the SAVE-side post-warmup cursor, or 0 if not recorded.""" + cfg = get_config() + if cfg is None or cfg.workspace_dir is None: + return 0 + path = os.path.join(cfg.workspace_dir, "warmup_alloc_offset.json") + if not os.path.exists(path): + return 0 + with open(path) as f: + return json.load(f).get("warmup_alloc_offset", 0) + + def preallocate_for_load_mode() -> None: cfg = get_config() if cfg is None or cfg.mode != CUDAGraphExtensionMode.LOAD: From 76b8a84fa2247d686cc62061e0495da86c5fe03f Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 23 Jul 2026 22:16:41 +0000 Subject: [PATCH 12/22] [sglang] Optional warmup serialize+restore: memcpy dumped gap bytes on LOAD Instead of re-running the dense warmup forwards on LOAD, dump the exact post-warmup VMM gap bytes during SAVE and cudaMemcpy them back into the same addresses on LOAD (opt-in via FOUNDRY_SGLANG_WARMUP_RESTORE=1). The decode graphs replay by address and read these restored persistent buffers; decode is never run eagerly on LOAD, so nothing re-triggers the lazy init. Off by default; the validated re-run-warmup path is unchanged. Co-Authored-By: Rahul Chalamala <22563365+rchalamala@users.noreply.github.com> Signed-off-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- python/foundry/integration/sglang/hooks.py | 40 +++++- python/foundry/integration/sglang/runtime.py | 143 +++++++++++++++++++ 2 files changed, 181 insertions(+), 2 deletions(-) diff --git a/python/foundry/integration/sglang/hooks.py b/python/foundry/integration/sglang/hooks.py index 0d4b84fc..31df7d47 100644 --- a/python/foundry/integration/sglang/hooks.py +++ b/python/foundry/integration/sglang/hooks.py @@ -81,6 +81,20 @@ def _load_warmup_skip_enabled() -> bool: return os.environ.get("FOUNDRY_SGLANG_SKIP_LOAD_WARMUP", "0") not in ("0", "false", "False") +def _warmup_restore_enabled() -> bool: + """Whether the dense path serializes the warmup gap on SAVE and restores its + bytes on LOAD instead of re-running the warmup forwards. + + Off by default. Unlike ``FOUNDRY_SGLANG_SKIP_LOAD_WARMUP`` (which leaves the + gap uninitialized and is incorrect for models whose warmup buffers hold + read-before-write state), this dumps the exact post-warmup gap bytes during + SAVE and memcpys them back into the same VMM addresses on LOAD — equivalent + to re-running the deterministic in-region warmup, but ~warmup-time faster. + Set ``FOUNDRY_SGLANG_WARMUP_RESTORE=1`` (must be set on BOTH the SAVE run, + to produce the dump, and the LOAD run, to consume it).""" + return os.environ.get("FOUNDRY_SGLANG_WARMUP_RESTORE", "0") not in ("0", "false", "False") + + def _resolve_dp_rank(model_runner) -> int | None: dp_rank = getattr(model_runner, "dp_rank", None) if dp_rank is not None: @@ -448,6 +462,9 @@ def _run_dense_warmup(self): fi_backend.decode_cuda_graph_metadata = {} if has_prefill: fi_backend.prefill_cuda_graph_metadata = {} + from foundry import ops as fops + + before_warmup = fops.get_current_alloc_offset() rt.log_alloc_offset("before_warmup") try: _run_warmup_pass(self) @@ -462,10 +479,15 @@ def _run_dense_warmup(self): if has_prefill: fi_backend.prefill_cuda_graph_metadata = {} torch.cuda.empty_cache() + after_warmup = fops.get_current_alloc_offset() rt.log_alloc_offset("after_warmup") # Record the post-warmup cursor so LOAD can optionally reserve this # gap instead of re-running the warmup (see _load_warmup_skip_enabled). rt.save_warmup_alloc_offset() + # Serialize the warmup gap bytes so LOAD can restore them instead of + # re-running the warmup forwards (see _warmup_restore_enabled). + if get_graph_extension_mode() == CUDAGraphExtensionMode.SAVE and _warmup_restore_enabled(): + rt.dump_warmup_region(before_warmup, after_warmup) @functools.wraps(orig_capture) def patched(self, *args, **kwargs): @@ -523,12 +545,26 @@ def patched(self, *args, **kwargs): # buffers that were lazily allocated during SAVE's warmup — stale # on LOAD (different process). warmup_needed = not _ep_lazy_init_needed() and _dense_save_warmup_enabled() - skip_warmup = warmup_needed and _load_warmup_skip_enabled() - if warmup_needed and not skip_warmup: + restore_warmup = warmup_needed and _warmup_restore_enabled() + skip_warmup = warmup_needed and not restore_warmup and _load_warmup_skip_enabled() + if warmup_needed and not skip_warmup and not restore_warmup: _run_dense_warmup(self) rt.log_alloc_offset("before_preallocate") rt.preallocate_for_load_mode() rt.log_alloc_offset("after_preallocate") + if restore_warmup: + # Instead of re-running the warmup forwards, memcpy the exact + # post-warmup gap bytes SAVE dumped back into the (now mapped) + # VMM region and advance the cursor to the post-warmup offset. + # The decode graphs replay by address and read these restored + # buffers; decode is never run eagerly on LOAD (foundry replays + # the captured graph), so nothing re-triggers the lazy init. + if not rt.restore_warmup_region(): + raise RuntimeError( + "FOUNDRY_SGLANG_WARMUP_RESTORE=1 but no warmup_region.bin " + "in the archive (was SAVE run with the flag set?)" + ) + rt.log_alloc_offset("after_warmup_restore") if skip_warmup: # Instead of re-running the warmup forwards, jump the cursor # over the SAVE-side warmup gap. preallocate_for_load_mode diff --git a/python/foundry/integration/sglang/runtime.py b/python/foundry/integration/sglang/runtime.py index b781ae7f..14aba84c 100644 --- a/python/foundry/integration/sglang/runtime.py +++ b/python/foundry/integration/sglang/runtime.py @@ -204,6 +204,149 @@ def load_warmup_alloc_offset() -> int: return json.load(f).get("warmup_alloc_offset", 0) +_WARMUP_REGION_BIN = "warmup_region.bin" +_WARMUP_REGION_JSON = "warmup_region.json" +_CUDA_MEMCPY_D2H = 2 +_CUDA_MEMCPY_H2D = 1 +_libcudart_cache = None + + +def _libcudart(): + """Return a ctypes handle to the libcudart already loaded by torch. + + We copy raw bytes into / out of the VMM region (a plain device pointer at + ``base_addr + offset``) with ``cudaMemcpy``. torch has libcudart mapped in + this process, so resolve its path from ``/proc/self/maps`` rather than + relying on the loader search path (the cu13 wheel dir isn't on it).""" + global _libcudart_cache + if _libcudart_cache is not None: + return _libcudart_cache + import ctypes + import re + + lib_path = None + try: + with open("/proc/self/maps") as f: + for line in f: + m = re.search(r"(\S*/libcudart\.so[^\s]*)", line) + if m: + lib_path = m.group(1) + break + except OSError: + lib_path = None + lib = ctypes.CDLL(lib_path) if lib_path else ctypes.CDLL("libcudart.so") + lib.cudaMemcpy.restype = ctypes.c_int + lib.cudaMemcpy.argtypes = [ + ctypes.c_void_p, + ctypes.c_void_p, + ctypes.c_size_t, + ctypes.c_int, + ] + lib.cudaDeviceSynchronize.restype = ctypes.c_int + lib.cudaGetErrorString.restype = ctypes.c_char_p + lib.cudaGetErrorString.argtypes = [ctypes.c_int] + _libcudart_cache = lib + return lib + + +def _cuda_check(lib, rc: int, what: str) -> None: + if rc != 0: + msg = lib.cudaGetErrorString(rc) + detail = msg.decode() if isinstance(msg, bytes) else str(msg) + raise RuntimeError(f"[Foundry] {what} failed: cuda error {rc} ({detail})") + + +def dump_warmup_region(before_offset: int, after_offset: int) -> None: + """Dump the VMM bytes the dense warmup produced (``[before, after)``) so LOAD + can restore them instead of re-running the warmup forwards. + + The decode graphs replay by absolute address and read persistent buffers the + warmup wrote in this gap (torch.compile constants, Triton/MoE/mamba workspace, + RoPE caches). Because the in-region warmup is deterministic, SAVE's post-warmup + gap bytes equal what LOAD's warmup would produce — restoring them is equivalent + to (but far cheaper than) re-running the warmup.""" + import ctypes + + cfg = get_config() + if cfg is None or cfg.mode != CUDAGraphExtensionMode.SAVE or cfg.workspace_dir is None: + return + size = after_offset - before_offset + if size <= 0: + return + lib = _libcudart() + src = cfg.base_addr + before_offset + buf = ctypes.create_string_buffer(size) + _cuda_check(lib, lib.cudaDeviceSynchronize(), "cudaDeviceSynchronize(pre-dump)") + _cuda_check( + lib, + lib.cudaMemcpy( + ctypes.cast(buf, ctypes.c_void_p), ctypes.c_void_p(src), size, _CUDA_MEMCPY_D2H + ), + "cudaMemcpy(warmup region D2H)", + ) + bin_path = os.path.join(cfg.workspace_dir, _WARMUP_REGION_BIN) + with open(bin_path, "wb") as f: + f.write(buf.raw) + json_path = os.path.join(cfg.workspace_dir, _WARMUP_REGION_JSON) + with open(json_path, "w") as f: + json.dump({"before_offset": before_offset, "after_offset": after_offset}, f) + logger.info( + "[Foundry] SGLang dumped warmup region [%d, %d) (%.1f MB) to %s", + before_offset, + after_offset, + size / (1024 * 1024), + bin_path, + ) + + +def restore_warmup_region() -> bool: + """Restore the SAVE-side warmup gap bytes into the mapped VMM region and set + the cursor to the post-warmup offset. Returns False if no dump is present. + + Must be called AFTER ``preallocate_for_load_mode`` has physically mapped the + region (including this gap), so the H2D copy targets backed memory.""" + import ctypes + + cfg = get_config() + if cfg is None or cfg.mode != CUDAGraphExtensionMode.LOAD or cfg.workspace_dir is None: + return False + json_path = os.path.join(cfg.workspace_dir, _WARMUP_REGION_JSON) + bin_path = os.path.join(cfg.workspace_dir, _WARMUP_REGION_BIN) + if not (os.path.exists(json_path) and os.path.exists(bin_path)): + return False + with open(json_path) as f: + meta = json.load(f) + before_offset = meta["before_offset"] + after_offset = meta["after_offset"] + size = after_offset - before_offset + with open(bin_path, "rb") as f: + data = f.read() + if len(data) != size: + raise RuntimeError( + f"[Foundry] warmup region dump size mismatch: file={len(data)} expected={size}" + ) + lib = _libcudart() + dst = cfg.base_addr + before_offset + buf = (ctypes.c_char * size).from_buffer_copy(data) + _cuda_check( + lib, + lib.cudaMemcpy( + ctypes.c_void_p(dst), ctypes.cast(buf, ctypes.c_void_p), size, _CUDA_MEMCPY_H2D + ), + "cudaMemcpy(warmup region H2D)", + ) + _cuda_check(lib, lib.cudaDeviceSynchronize(), "cudaDeviceSynchronize(post-restore)") + cge.set_current_alloc_offset(after_offset) + logger.info( + "[Foundry] SGLang restored warmup region [%d, %d) (%.1f MB); cursor=%d", + before_offset, + after_offset, + size / (1024 * 1024), + after_offset, + ) + return True + + def preallocate_for_load_mode() -> None: cfg = get_config() if cfg is None or cfg.mode != CUDAGraphExtensionMode.LOAD: From dcbece3a0cff1d970aad638eabee5f7d25f12305 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 23 Jul 2026 22:31:09 +0000 Subject: [PATCH 13/22] [sglang] Use CUDA driver API (cuMemcpyDtoH/HtoD) for warmup region dump/restore The VMM region is driver-mapped (cuMemMap on CUdeviceptr); the CUDA runtime cudaMemcpy returns cudaErrorInvalidValue for it. Copy the warmup gap bytes with the driver API instead, operating on CUdeviceptr integers in the current context. Co-Authored-By: Rahul Chalamala <22563365+rchalamala@users.noreply.github.com> Signed-off-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- python/foundry/integration/sglang/runtime.py | 90 ++++++++++---------- 1 file changed, 44 insertions(+), 46 deletions(-) diff --git a/python/foundry/integration/sglang/runtime.py b/python/foundry/integration/sglang/runtime.py index 14aba84c..84f99bbe 100644 --- a/python/foundry/integration/sglang/runtime.py +++ b/python/foundry/integration/sglang/runtime.py @@ -206,21 +206,22 @@ def load_warmup_alloc_offset() -> int: _WARMUP_REGION_BIN = "warmup_region.bin" _WARMUP_REGION_JSON = "warmup_region.json" -_CUDA_MEMCPY_D2H = 2 -_CUDA_MEMCPY_H2D = 1 -_libcudart_cache = None - - -def _libcudart(): - """Return a ctypes handle to the libcudart already loaded by torch. - - We copy raw bytes into / out of the VMM region (a plain device pointer at - ``base_addr + offset``) with ``cudaMemcpy``. torch has libcudart mapped in - this process, so resolve its path from ``/proc/self/maps`` rather than - relying on the loader search path (the cu13 wheel dir isn't on it).""" - global _libcudart_cache - if _libcudart_cache is not None: - return _libcudart_cache +_libcuda_cache = None + + +def _libcuda(): + """Return a ctypes handle to the CUDA *driver* library (libcuda). + + The VMM region is driver-mapped (``cuMemMap`` on a ``CUdeviceptr``), which + the CUDA *runtime* (``cudaMemcpy``) cannot address — it returns + ``cudaErrorInvalidValue`` for pointers not in its allocation tables. So copy + the gap bytes with the driver API (``cuMemcpyDtoH_v2`` / ``cuMemcpyHtoD_v2``), + which operate directly on ``CUdeviceptr`` integers in the current context (the + same context torch/foundry mapped the region in). Resolve libcuda from + ``/proc/self/maps`` (torch has it loaded) rather than the loader path.""" + global _libcuda_cache + if _libcuda_cache is not None: + return _libcuda_cache import ctypes import re @@ -228,32 +229,33 @@ def _libcudart(): try: with open("/proc/self/maps") as f: for line in f: - m = re.search(r"(\S*/libcudart\.so[^\s]*)", line) + m = re.search(r"(\S*/libcuda\.so[^\s]*)", line) if m: lib_path = m.group(1) break except OSError: lib_path = None - lib = ctypes.CDLL(lib_path) if lib_path else ctypes.CDLL("libcudart.so") - lib.cudaMemcpy.restype = ctypes.c_int - lib.cudaMemcpy.argtypes = [ - ctypes.c_void_p, - ctypes.c_void_p, - ctypes.c_size_t, - ctypes.c_int, - ] - lib.cudaDeviceSynchronize.restype = ctypes.c_int - lib.cudaGetErrorString.restype = ctypes.c_char_p - lib.cudaGetErrorString.argtypes = [ctypes.c_int] - _libcudart_cache = lib + lib = ctypes.CDLL(lib_path) if lib_path else ctypes.CDLL("libcuda.so.1") + # CUdeviceptr is a 64-bit integer; host pointers are void*. + lib.cuMemcpyDtoH_v2.restype = ctypes.c_int + lib.cuMemcpyDtoH_v2.argtypes = [ctypes.c_void_p, ctypes.c_ulonglong, ctypes.c_size_t] + lib.cuMemcpyHtoD_v2.restype = ctypes.c_int + lib.cuMemcpyHtoD_v2.argtypes = [ctypes.c_ulonglong, ctypes.c_void_p, ctypes.c_size_t] + lib.cuCtxSynchronize.restype = ctypes.c_int + lib.cuGetErrorString.restype = ctypes.c_int + lib.cuGetErrorString.argtypes = [ctypes.c_int, ctypes.POINTER(ctypes.c_char_p)] + _libcuda_cache = lib return lib -def _cuda_check(lib, rc: int, what: str) -> None: +def _cu_check(lib, rc: int, what: str) -> None: if rc != 0: - msg = lib.cudaGetErrorString(rc) - detail = msg.decode() if isinstance(msg, bytes) else str(msg) - raise RuntimeError(f"[Foundry] {what} failed: cuda error {rc} ({detail})") + import ctypes + + msg = ctypes.c_char_p() + lib.cuGetErrorString(rc, ctypes.byref(msg)) + detail = msg.value.decode() if msg.value else "unknown" + raise RuntimeError(f"[Foundry] {what} failed: cuda driver error {rc} ({detail})") def dump_warmup_region(before_offset: int, after_offset: int) -> None: @@ -273,16 +275,14 @@ def dump_warmup_region(before_offset: int, after_offset: int) -> None: size = after_offset - before_offset if size <= 0: return - lib = _libcudart() + lib = _libcuda() src = cfg.base_addr + before_offset buf = ctypes.create_string_buffer(size) - _cuda_check(lib, lib.cudaDeviceSynchronize(), "cudaDeviceSynchronize(pre-dump)") - _cuda_check( + _cu_check(lib, lib.cuCtxSynchronize(), "cuCtxSynchronize(pre-dump)") + _cu_check( lib, - lib.cudaMemcpy( - ctypes.cast(buf, ctypes.c_void_p), ctypes.c_void_p(src), size, _CUDA_MEMCPY_D2H - ), - "cudaMemcpy(warmup region D2H)", + lib.cuMemcpyDtoH_v2(ctypes.cast(buf, ctypes.c_void_p), src, size), + "cuMemcpyDtoH(warmup region)", ) bin_path = os.path.join(cfg.workspace_dir, _WARMUP_REGION_BIN) with open(bin_path, "wb") as f: @@ -325,17 +325,15 @@ def restore_warmup_region() -> bool: raise RuntimeError( f"[Foundry] warmup region dump size mismatch: file={len(data)} expected={size}" ) - lib = _libcudart() + lib = _libcuda() dst = cfg.base_addr + before_offset buf = (ctypes.c_char * size).from_buffer_copy(data) - _cuda_check( + _cu_check( lib, - lib.cudaMemcpy( - ctypes.c_void_p(dst), ctypes.cast(buf, ctypes.c_void_p), size, _CUDA_MEMCPY_H2D - ), - "cudaMemcpy(warmup region H2D)", + lib.cuMemcpyHtoD_v2(dst, ctypes.cast(buf, ctypes.c_void_p), size), + "cuMemcpyHtoD(warmup region)", ) - _cuda_check(lib, lib.cudaDeviceSynchronize(), "cudaDeviceSynchronize(post-restore)") + _cu_check(lib, lib.cuCtxSynchronize(), "cuCtxSynchronize(post-restore)") cge.set_current_alloc_offset(after_offset) logger.info( "[Foundry] SGLang restored warmup region [%d, %d) (%.1f MB); cursor=%d", From 8fd119f1393a067079fb229065777782d0fc1577 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 23 Jul 2026 22:49:51 +0000 Subject: [PATCH 14/22] [sglang] Dump only mapped sub-intervals of the warmup gap foundry unmaps freed allocations (cuMemFree -> cuMemUnmap), so after the dense warmup + empty_cache the [before, after) gap is fragmented: persistent buffers stay mapped, freed transients are holes. A single cuMemcpyDtoH over the whole span hit unmapped VA (CUDA_ERROR_INVALID_VALUE). Probe each 2MiB page with cuPointerGetAttribute(CU_POINTER_ATTRIBUTE_MAPPED), dump each mapped run, and record intervals so LOAD restores exactly those bytes into the (contiguously remapped) region. Co-Authored-By: Rahul Chalamala <22563365+rchalamala@users.noreply.github.com> Signed-off-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- python/foundry/integration/sglang/runtime.py | 103 +++++++++++++++---- 1 file changed, 81 insertions(+), 22 deletions(-) diff --git a/python/foundry/integration/sglang/runtime.py b/python/foundry/integration/sglang/runtime.py index 84f99bbe..5bb900f0 100644 --- a/python/foundry/integration/sglang/runtime.py +++ b/python/foundry/integration/sglang/runtime.py @@ -242,12 +242,53 @@ def _libcuda(): lib.cuMemcpyHtoD_v2.restype = ctypes.c_int lib.cuMemcpyHtoD_v2.argtypes = [ctypes.c_ulonglong, ctypes.c_void_p, ctypes.c_size_t] lib.cuCtxSynchronize.restype = ctypes.c_int + lib.cuPointerGetAttribute.restype = ctypes.c_int + lib.cuPointerGetAttribute.argtypes = [ + ctypes.c_void_p, + ctypes.c_int, + ctypes.c_ulonglong, + ] lib.cuGetErrorString.restype = ctypes.c_int lib.cuGetErrorString.argtypes = [ctypes.c_int, ctypes.POINTER(ctypes.c_char_p)] _libcuda_cache = lib return lib +# cuPointerGetAttribute(CU_POINTER_ATTRIBUTE_MAPPED) reports whether a VA inside +# a reserved region is physically backed. Probe at the VMM minimum granularity +# (2 MiB on all current NVIDIA GPUs, which is also what foundry aligns +# allocations to), so mapped/unmapped runs land on probe boundaries. +_CU_POINTER_ATTRIBUTE_MAPPED = 15 +_VMM_GRANULARITY = 2 * 1024 * 1024 + + +def _mapped_intervals(lib, base_ptr: int, size: int) -> list[tuple[int, int]]: + """Return ``[(rel_offset, length), ...]`` of the physically-backed runs in + ``[base_ptr, base_ptr + size)``. foundry unmaps freed allocations, so the + dense warmup gap is fragmented: persistent buffers stay mapped, transient + ones become holes.""" + import ctypes + + intervals: list[tuple[int, int]] = [] + val = ctypes.c_int(0) + run_start = None + off = 0 + while off < size: + rc = lib.cuPointerGetAttribute( + ctypes.byref(val), _CU_POINTER_ATTRIBUTE_MAPPED, base_ptr + off + ) + mapped = rc == 0 and val.value == 1 + if mapped and run_start is None: + run_start = off + elif not mapped and run_start is not None: + intervals.append((run_start, off - run_start)) + run_start = None + off += _VMM_GRANULARITY + if run_start is not None: + intervals.append((run_start, size - run_start)) + return intervals + + def _cu_check(lib, rc: int, what: str) -> None: if rc != 0: import ctypes @@ -276,24 +317,37 @@ def dump_warmup_region(before_offset: int, after_offset: int) -> None: if size <= 0: return lib = _libcuda() - src = cfg.base_addr + before_offset - buf = ctypes.create_string_buffer(size) + base_ptr = cfg.base_addr + before_offset _cu_check(lib, lib.cuCtxSynchronize(), "cuCtxSynchronize(pre-dump)") - _cu_check( - lib, - lib.cuMemcpyDtoH_v2(ctypes.cast(buf, ctypes.c_void_p), src, size), - "cuMemcpyDtoH(warmup region)", - ) + intervals = _mapped_intervals(lib, base_ptr, size) + total = sum(length for _, length in intervals) bin_path = os.path.join(cfg.workspace_dir, _WARMUP_REGION_BIN) with open(bin_path, "wb") as f: - f.write(buf.raw) + for rel, length in intervals: + buf = ctypes.create_string_buffer(length) + _cu_check( + lib, + lib.cuMemcpyDtoH_v2(ctypes.cast(buf, ctypes.c_void_p), base_ptr + rel, length), + "cuMemcpyDtoH(warmup region)", + ) + f.write(buf.raw) json_path = os.path.join(cfg.workspace_dir, _WARMUP_REGION_JSON) with open(json_path, "w") as f: - json.dump({"before_offset": before_offset, "after_offset": after_offset}, f) + json.dump( + { + "before_offset": before_offset, + "after_offset": after_offset, + "intervals": intervals, + }, + f, + ) logger.info( - "[Foundry] SGLang dumped warmup region [%d, %d) (%.1f MB) to %s", + "[Foundry] SGLang dumped warmup region [%d, %d): %d mapped bytes across " + "%d intervals (of %.1f MB span) to %s", before_offset, after_offset, + total, + len(intervals), size / (1024 * 1024), bin_path, ) @@ -318,28 +372,33 @@ def restore_warmup_region() -> bool: meta = json.load(f) before_offset = meta["before_offset"] after_offset = meta["after_offset"] - size = after_offset - before_offset + intervals = [(int(rel), int(length)) for rel, length in meta["intervals"]] + total = sum(length for _, length in intervals) with open(bin_path, "rb") as f: data = f.read() - if len(data) != size: + if len(data) != total: raise RuntimeError( - f"[Foundry] warmup region dump size mismatch: file={len(data)} expected={size}" + f"[Foundry] warmup region dump size mismatch: file={len(data)} expected={total}" ) lib = _libcuda() - dst = cfg.base_addr + before_offset - buf = (ctypes.c_char * size).from_buffer_copy(data) - _cu_check( - lib, - lib.cuMemcpyHtoD_v2(dst, ctypes.cast(buf, ctypes.c_void_p), size), - "cuMemcpyHtoD(warmup region)", - ) + base_ptr = cfg.base_addr + before_offset + pos = 0 + for rel, length in intervals: + buf = (ctypes.c_char * length).from_buffer_copy(data[pos : pos + length]) + _cu_check( + lib, + lib.cuMemcpyHtoD_v2(base_ptr + rel, ctypes.cast(buf, ctypes.c_void_p), length), + "cuMemcpyHtoD(warmup region)", + ) + pos += length _cu_check(lib, lib.cuCtxSynchronize(), "cuCtxSynchronize(post-restore)") cge.set_current_alloc_offset(after_offset) logger.info( - "[Foundry] SGLang restored warmup region [%d, %d) (%.1f MB); cursor=%d", + "[Foundry] SGLang restored warmup region [%d, %d): %d bytes across %d intervals; cursor=%d", before_offset, after_offset, - size / (1024 * 1024), + total, + len(intervals), after_offset, ) return True From 7bfe52ba4e42513eb51be04d8c879e0fdf60f92d Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 23 Jul 2026 23:06:48 +0000 Subject: [PATCH 15/22] [sglang] Dump warmup gap before empty_cache; anchor LOAD cursor post-empty_cache The gap is only contiguously mapped BEFORE empty_cache (the caching allocator still holds every block the warmup touched). After empty_cache the freed blocks are returned to foundry and unmapped, so the probe found 0 mapped bytes. Move the dump into the finally block before empty_cache and anchor the LOAD cursor on the post-empty_cache warmup offset the validated path uses. Co-Authored-By: Rahul Chalamala <22563365+rchalamala@users.noreply.github.com> Signed-off-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- python/foundry/integration/sglang/hooks.py | 21 +++++++++++++------- python/foundry/integration/sglang/runtime.py | 9 +++++++-- 2 files changed, 21 insertions(+), 9 deletions(-) diff --git a/python/foundry/integration/sglang/hooks.py b/python/foundry/integration/sglang/hooks.py index 31df7d47..21b5bbb3 100644 --- a/python/foundry/integration/sglang/hooks.py +++ b/python/foundry/integration/sglang/hooks.py @@ -466,6 +466,9 @@ def _run_dense_warmup(self): before_warmup = fops.get_current_alloc_offset() rt.log_alloc_offset("before_warmup") + dump_restore = ( + get_graph_extension_mode() == CUDAGraphExtensionMode.SAVE and _warmup_restore_enabled() + ) try: _run_warmup_pass(self) finally: @@ -478,16 +481,20 @@ def _run_dense_warmup(self): fi_backend.decode_cuda_graph_metadata = {} if has_prefill: fi_backend.prefill_cuda_graph_metadata = {} + # Serialize the warmup gap bytes BEFORE empty_cache: the caching + # allocator still holds (and keeps mapped) every block the warmup + # touched, so [before_warmup, cursor) is contiguously backed. After + # empty_cache the freed blocks are returned to foundry and unmapped, + # leaving nothing to copy. LOAD restores these bytes instead of + # re-running the warmup (see _warmup_restore_enabled). + if dump_restore: + rt.dump_warmup_region(before_warmup, fops.get_current_alloc_offset()) torch.cuda.empty_cache() - after_warmup = fops.get_current_alloc_offset() rt.log_alloc_offset("after_warmup") - # Record the post-warmup cursor so LOAD can optionally reserve this - # gap instead of re-running the warmup (see _load_warmup_skip_enabled). + # Record the post-warmup cursor so LOAD can reserve this gap + # (_load_warmup_skip_enabled) or anchor the warmup-restore cursor + # (_warmup_restore_enabled) without re-running the warmup forwards. rt.save_warmup_alloc_offset() - # Serialize the warmup gap bytes so LOAD can restore them instead of - # re-running the warmup forwards (see _warmup_restore_enabled). - if get_graph_extension_mode() == CUDAGraphExtensionMode.SAVE and _warmup_restore_enabled(): - rt.dump_warmup_region(before_warmup, after_warmup) @functools.wraps(orig_capture) def patched(self, *args, **kwargs): diff --git a/python/foundry/integration/sglang/runtime.py b/python/foundry/integration/sglang/runtime.py index 5bb900f0..829da37c 100644 --- a/python/foundry/integration/sglang/runtime.py +++ b/python/foundry/integration/sglang/runtime.py @@ -392,14 +392,19 @@ def restore_warmup_region() -> bool: ) pos += length _cu_check(lib, lib.cuCtxSynchronize(), "cuCtxSynchronize(post-restore)") - cge.set_current_alloc_offset(after_offset) + # The dump interval end is the pre-empty_cache cursor; the validated + # trajectory anchors on the post-empty_cache cursor (warmup_alloc_offset.json). + # Set the cursor to that so the subsequent FlashInfer pre-pass / graph load + # allocates from the exact SAVE offset. + cursor = load_warmup_alloc_offset() or after_offset + cge.set_current_alloc_offset(cursor) logger.info( "[Foundry] SGLang restored warmup region [%d, %d): %d bytes across %d intervals; cursor=%d", before_offset, after_offset, total, len(intervals), - after_offset, + cursor, ) return True From 90b5f3f2104c5e25693ffcf720c5269e1097b170 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 23 Jul 2026 23:24:41 +0000 Subject: [PATCH 16/22] [sglang] Dump warmup gap via ground-truth trial copies The pointer-attribute probe used the wrong CU_POINTER_ATTRIBUTE value and found 0 mapped bytes. Replace it with trial cuMemcpyDtoH copies: one contiguous copy first (succeeds when the gap is fully mapped, i.e. before empty_cache), falling back to per-2MiB-page copies that keep only cleanly-copied runs. No dependence on pointer-attribute semantics. Co-Authored-By: Rahul Chalamala <22563365+rchalamala@users.noreply.github.com> Signed-off-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- python/foundry/integration/sglang/runtime.py | 73 ++++++++++---------- 1 file changed, 36 insertions(+), 37 deletions(-) diff --git a/python/foundry/integration/sglang/runtime.py b/python/foundry/integration/sglang/runtime.py index 829da37c..7a246a0f 100644 --- a/python/foundry/integration/sglang/runtime.py +++ b/python/foundry/integration/sglang/runtime.py @@ -242,51 +242,59 @@ def _libcuda(): lib.cuMemcpyHtoD_v2.restype = ctypes.c_int lib.cuMemcpyHtoD_v2.argtypes = [ctypes.c_ulonglong, ctypes.c_void_p, ctypes.c_size_t] lib.cuCtxSynchronize.restype = ctypes.c_int - lib.cuPointerGetAttribute.restype = ctypes.c_int - lib.cuPointerGetAttribute.argtypes = [ - ctypes.c_void_p, - ctypes.c_int, - ctypes.c_ulonglong, - ] lib.cuGetErrorString.restype = ctypes.c_int lib.cuGetErrorString.argtypes = [ctypes.c_int, ctypes.POINTER(ctypes.c_char_p)] _libcuda_cache = lib return lib -# cuPointerGetAttribute(CU_POINTER_ATTRIBUTE_MAPPED) reports whether a VA inside -# a reserved region is physically backed. Probe at the VMM minimum granularity -# (2 MiB on all current NVIDIA GPUs, which is also what foundry aligns -# allocations to), so mapped/unmapped runs land on probe boundaries. -_CU_POINTER_ATTRIBUTE_MAPPED = 15 +# VMM minimum allocation granularity (2 MiB on all current NVIDIA GPUs, which is +# also what foundry aligns allocations to), used as the probe step when the gap +# turns out to be fragmented. _VMM_GRANULARITY = 2 * 1024 * 1024 -def _mapped_intervals(lib, base_ptr: int, size: int) -> list[tuple[int, int]]: - """Return ``[(rel_offset, length), ...]`` of the physically-backed runs in - ``[base_ptr, base_ptr + size)``. foundry unmaps freed allocations, so the - dense warmup gap is fragmented: persistent buffers stay mapped, transient - ones become holes.""" +def _dump_mapped(lib, base_ptr: int, size: int) -> tuple[list[tuple[int, int]], bytes]: + """Copy the physically-backed bytes of ``[base_ptr, base_ptr + size)`` to host + and return ``([(rel_offset, length), ...], concatenated_bytes)``. + + Ground truth via trial copies rather than pointer-attribute queries: try one + contiguous copy first (the gap is contiguously mapped as long as the caching + allocator still holds every block, i.e. before ``empty_cache``); if that + fails (foundry unmaps freed blocks, so a hole may exist mid-warmup), fall + back to per-granularity copies and keep only the pages that copy cleanly. A + failed ``cuMemcpyDtoH`` is a synchronous argument-validation error that + leaves the context healthy, so probing this way is safe.""" import ctypes + buf = ctypes.create_string_buffer(size) + rc = lib.cuMemcpyDtoH_v2(ctypes.cast(buf, ctypes.c_void_p), base_ptr, size) + if rc == 0: + return [(0, size)], buf.raw + intervals: list[tuple[int, int]] = [] - val = ctypes.c_int(0) + parts: list[bytes] = [] run_start = None + run_parts: list[bytes] = [] off = 0 while off < size: - rc = lib.cuPointerGetAttribute( - ctypes.byref(val), _CU_POINTER_ATTRIBUTE_MAPPED, base_ptr + off - ) - mapped = rc == 0 and val.value == 1 - if mapped and run_start is None: - run_start = off - elif not mapped and run_start is not None: + length = min(_VMM_GRANULARITY, size - off) + page = ctypes.create_string_buffer(length) + prc = lib.cuMemcpyDtoH_v2(ctypes.cast(page, ctypes.c_void_p), base_ptr + off, length) + if prc == 0: + if run_start is None: + run_start = off + run_parts = [] + run_parts.append(page.raw) + elif run_start is not None: intervals.append((run_start, off - run_start)) + parts.append(b"".join(run_parts)) run_start = None - off += _VMM_GRANULARITY + off += length if run_start is not None: intervals.append((run_start, size - run_start)) - return intervals + parts.append(b"".join(run_parts)) + return intervals, b"".join(parts) def _cu_check(lib, rc: int, what: str) -> None: @@ -308,8 +316,6 @@ def dump_warmup_region(before_offset: int, after_offset: int) -> None: RoPE caches). Because the in-region warmup is deterministic, SAVE's post-warmup gap bytes equal what LOAD's warmup would produce — restoring them is equivalent to (but far cheaper than) re-running the warmup.""" - import ctypes - cfg = get_config() if cfg is None or cfg.mode != CUDAGraphExtensionMode.SAVE or cfg.workspace_dir is None: return @@ -319,18 +325,11 @@ def dump_warmup_region(before_offset: int, after_offset: int) -> None: lib = _libcuda() base_ptr = cfg.base_addr + before_offset _cu_check(lib, lib.cuCtxSynchronize(), "cuCtxSynchronize(pre-dump)") - intervals = _mapped_intervals(lib, base_ptr, size) + intervals, data = _dump_mapped(lib, base_ptr, size) total = sum(length for _, length in intervals) bin_path = os.path.join(cfg.workspace_dir, _WARMUP_REGION_BIN) with open(bin_path, "wb") as f: - for rel, length in intervals: - buf = ctypes.create_string_buffer(length) - _cu_check( - lib, - lib.cuMemcpyDtoH_v2(ctypes.cast(buf, ctypes.c_void_p), base_ptr + rel, length), - "cuMemcpyDtoH(warmup region)", - ) - f.write(buf.raw) + f.write(data) json_path = os.path.join(cfg.workspace_dir, _WARMUP_REGION_JSON) with open(json_path, "w") as f: json.dump( From 00fb57c0ade8cc05cc289666ab4f9cdad6a735b4 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 23 Jul 2026 23:40:31 +0000 Subject: [PATCH 17/22] [sglang] Write shared warmup_state.json atomically to avoid torn reads All TP ranks read workspace_root/warmup_state.json at end of capture while rank 0 rewrites it in place; a reader could hit the truncated file mid-write and die with JSONDecodeError. Write to a pid-tmp file and os.replace so readers always see a complete file. Co-Authored-By: Rahul Chalamala <22563365+rchalamala@users.noreply.github.com> Signed-off-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- python/foundry/integration/sglang/runtime.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/python/foundry/integration/sglang/runtime.py b/python/foundry/integration/sglang/runtime.py index 7a246a0f..a5e6f9fc 100644 --- a/python/foundry/integration/sglang/runtime.py +++ b/python/foundry/integration/sglang/runtime.py @@ -91,8 +91,13 @@ def save_warmup_state(state: WarmupState) -> None: if ext_state is not None and ext_state.rank != 0 and os.path.exists(path): return os.makedirs(workspace_root, exist_ok=True) - with open(path, "w") as f: + # warmup_state.json lives at the shared workspace_root; every TP rank calls + # load_warmup_state() at end of capture. Write atomically (temp + replace) + # so a concurrent reader never observes the truncated/empty file. + tmp = f"{path}.{os.getpid()}.tmp" + with open(tmp, "w") as f: json.dump(asdict(state), f, indent=2) + os.replace(tmp, path) logger.info("[Foundry] Saved SGLang warmup state to %s", path) From 114d4cd8ff939b771d3cfd87c101e9824136fdb3 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 23 Jul 2026 23:58:25 +0000 Subject: [PATCH 18/22] [sglang] Document validated opt-in FOUNDRY_SGLANG_WARMUP_RESTORE path Co-Authored-By: Rahul Chalamala <22563365+rchalamala@users.noreply.github.com> Signed-off-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- recipe/experimental/README.md | 32 ++++++++++++++++++++++++++++++++ 1 file changed, 32 insertions(+) diff --git a/recipe/experimental/README.md b/recipe/experimental/README.md index 6730ac64..2f2419c0 100644 --- a/recipe/experimental/README.md +++ b/recipe/experimental/README.md @@ -254,6 +254,38 @@ Key model-specific requirements: lazy-init model buffers allocated during warmup land out-of-region and cause `CUDA error: illegal memory access` on LOAD graph replay. +### Optional: skip the LOAD warmup by restoring its bytes (`FOUNDRY_SGLANG_WARMUP_RESTORE=1`) + +The in-region warmup re-runs real model forwards on LOAD (~20 s on 122B) only to +reproduce the persistent buffers it wrote on SAVE. Because that warmup is +deterministic, those bytes are identical on SAVE and LOAD, so LOAD can restore +them instead of recomputing them. + +Set `FOUNDRY_SGLANG_WARMUP_RESTORE=1` on **both** the SAVE run (to dump the gap) +and the LOAD run (to consume it). On SAVE, just before the post-warmup +`empty_cache`, foundry copies the warmup allocation gap +(`[before_warmup, after_warmup)`) out of the VMM region to +`rank_N/warmup_region.bin` (+ `.json` interval metadata). On LOAD it memcpys +those bytes back into the same VMM addresses (after `preallocate_for_load_mode` +maps the region) and sets the alloc cursor to the post-warmup offset — skipping +the warmup forwards entirely. + +GPU-validated on `Qwen/Qwen3.5-122B-A10B`, H200:4, `--tp-size 4`: + +- **Exact** — temperature-0 tokens from the restored LOAD engine match the + warmed SAVE engine on all 16/16 probe prompts (bit-identical, not "close"), + same as the warmup-running path; `final_alloc_offset` `120376524800` + reproducible across both SAVE passes on all 4 ranks. +- **Faster LOAD** — ~106 s vs ~127 s with the warmup re-run (vs ~249 s cold + SAVE). +- **Archive cost** — 144 MiB/rank (one contiguous mapped interval) written + during SAVE. + +Off by default. Distinct from `FOUNDRY_SGLANG_SKIP_LOAD_WARMUP=1`, which merely +reserves the gap without restoring its contents — proven **incorrect** here +(coherent but divergent LOAD tokens), since the gap holds read-before-write +state. + ## IPC-specific troubleshooting | Symptom | Likely cause | From 15b969bfdc29a78c1933360f364a738418cad10b Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 24 Jul 2026 01:12:45 +0000 Subject: [PATCH 19/22] [sglang] Clamp warmup restore to final_alloc_offset (LOAD-mapped region) The dump end is the pre-empty_cache peak, which can exceed final_alloc_offset for models whose warmup transients peak above the final captured offset. preallocate_for_load_mode only maps [.., final_alloc_offset), so restoring those higher bytes would hit unmapped VA. Those bytes are transient (never referenced by the replayed graphs), so clamp each interval to the mapped region. No-op on the validated path (peak == after_warmup < final). Co-Authored-By: Rahul Chalamala <22563365+rchalamala@users.noreply.github.com> Signed-off-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- python/foundry/integration/sglang/runtime.py | 57 ++++++++++++++------ 1 file changed, 42 insertions(+), 15 deletions(-) diff --git a/python/foundry/integration/sglang/runtime.py b/python/foundry/integration/sglang/runtime.py index a5e6f9fc..abd9a77c 100644 --- a/python/foundry/integration/sglang/runtime.py +++ b/python/foundry/integration/sglang/runtime.py @@ -357,6 +357,21 @@ def dump_warmup_region(before_offset: int, after_offset: int) -> None: ) +def _load_final_alloc_offset() -> int: + """Return the SAVE-side final alloc offset (the top of the LOAD-mapped + region), or 0 if not recorded.""" + cfg = get_config() + if cfg is None or cfg.workspace_dir is None: + return 0 + path = os.path.join(cfg.workspace_dir, "final_alloc_offset.json") + if os.path.exists(path): + with open(path) as f: + final = json.load(f).get("final_alloc_offset", 0) + if final > 0: + return final + return load_warmup_state().final_alloc_offset + + def restore_warmup_region() -> bool: """Restore the SAVE-side warmup gap bytes into the mapped VMM region and set the cursor to the post-warmup offset. Returns False if no dump is present. @@ -384,16 +399,32 @@ def restore_warmup_region() -> bool: raise RuntimeError( f"[Foundry] warmup region dump size mismatch: file={len(data)} expected={total}" ) + # preallocate_for_load_mode only maps [.., final_alloc_offset); the dump end + # is the pre-empty_cache peak, which can exceed it for models whose warmup + # transients peak above the final captured offset. Bytes at/above final are + # transient (never referenced by the replayed graphs), so clamp the restore + # to the mapped region — writing past it would hit unmapped VA. + final = _load_final_alloc_offset() lib = _libcuda() base_ptr = cfg.base_addr + before_offset + restored = 0 pos = 0 for rel, length in intervals: - buf = (ctypes.c_char * length).from_buffer_copy(data[pos : pos + length]) - _cu_check( - lib, - lib.cuMemcpyHtoD_v2(base_ptr + rel, ctypes.cast(buf, ctypes.c_void_p), length), - "cuMemcpyHtoD(warmup region)", - ) + copy_len = length + if final > 0: + abs_end = before_offset + rel + length + if before_offset + rel >= final: + copy_len = 0 + elif abs_end > final: + copy_len = final - (before_offset + rel) + if copy_len > 0: + buf = (ctypes.c_char * copy_len).from_buffer_copy(data[pos : pos + copy_len]) + _cu_check( + lib, + lib.cuMemcpyHtoD_v2(base_ptr + rel, ctypes.cast(buf, ctypes.c_void_p), copy_len), + "cuMemcpyHtoD(warmup region)", + ) + restored += copy_len pos += length _cu_check(lib, lib.cuCtxSynchronize(), "cuCtxSynchronize(post-restore)") # The dump interval end is the pre-empty_cache cursor; the validated @@ -403,11 +434,14 @@ def restore_warmup_region() -> bool: cursor = load_warmup_alloc_offset() or after_offset cge.set_current_alloc_offset(cursor) logger.info( - "[Foundry] SGLang restored warmup region [%d, %d): %d bytes across %d intervals; cursor=%d", + "[Foundry] SGLang restored warmup region [%d, %d): %d of %d bytes across " + "%d intervals (clamped to final=%d); cursor=%d", before_offset, after_offset, + restored, total, len(intervals), + final, cursor, ) return True @@ -417,14 +451,7 @@ def preallocate_for_load_mode() -> None: cfg = get_config() if cfg is None or cfg.mode != CUDAGraphExtensionMode.LOAD: return - final = 0 - if cfg.workspace_dir is not None: - path = os.path.join(cfg.workspace_dir, "final_alloc_offset.json") - if os.path.exists(path): - with open(path) as f: - final = json.load(f).get("final_alloc_offset", 0) - if final <= 0: - final = load_warmup_state().final_alloc_offset + final = _load_final_alloc_offset() remaining = final - cge.get_current_alloc_offset() if remaining > 0: cge.preallocate_region(remaining) From 11f1a966c8453041959ea7a25d71f51b2f6aedc5 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 24 Jul 2026 01:23:30 +0000 Subject: [PATCH 20/22] [sglang] Default warmup-restore on, with safe fallback to re-run warmup FOUNDRY_SGLANG_WARMUP_RESTORE now defaults to 1: SAVE dumps the warmup gap, LOAD restores it. LOAD only takes the restore path when the archive actually carries the dump (warmup_region_available); otherwise it falls back to re-running the deterministic warmup forwards, so archives produced before this change (or with restore disabled) still load. Set the flag to 0 to force the re-run path on both save and load. Co-Authored-By: Rahul Chalamala <22563365+rchalamala@users.noreply.github.com> Signed-off-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- python/foundry/integration/sglang/hooks.py | 26 +++++++++++++------- python/foundry/integration/sglang/runtime.py | 14 +++++++++++ 2 files changed, 31 insertions(+), 9 deletions(-) diff --git a/python/foundry/integration/sglang/hooks.py b/python/foundry/integration/sglang/hooks.py index 21b5bbb3..6963650d 100644 --- a/python/foundry/integration/sglang/hooks.py +++ b/python/foundry/integration/sglang/hooks.py @@ -85,14 +85,15 @@ def _warmup_restore_enabled() -> bool: """Whether the dense path serializes the warmup gap on SAVE and restores its bytes on LOAD instead of re-running the warmup forwards. - Off by default. Unlike ``FOUNDRY_SGLANG_SKIP_LOAD_WARMUP`` (which leaves the + Defaults on. Unlike ``FOUNDRY_SGLANG_SKIP_LOAD_WARMUP`` (which leaves the gap uninitialized and is incorrect for models whose warmup buffers hold read-before-write state), this dumps the exact post-warmup gap bytes during SAVE and memcpys them back into the same VMM addresses on LOAD — equivalent to re-running the deterministic in-region warmup, but ~warmup-time faster. - Set ``FOUNDRY_SGLANG_WARMUP_RESTORE=1`` (must be set on BOTH the SAVE run, - to produce the dump, and the LOAD run, to consume it).""" - return os.environ.get("FOUNDRY_SGLANG_WARMUP_RESTORE", "0") not in ("0", "false", "False") + SAVE writes the dump; LOAD consumes it when present and otherwise falls back + to re-running the warmup (so archives produced without the dump still load). + Set ``FOUNDRY_SGLANG_WARMUP_RESTORE=0`` to force the re-run path on both.""" + return os.environ.get("FOUNDRY_SGLANG_WARMUP_RESTORE", "1") not in ("0", "false", "False") def _resolve_dp_rank(model_runner) -> int | None: @@ -552,7 +553,17 @@ def patched(self, *args, **kwargs): # buffers that were lazily allocated during SAVE's warmup — stale # on LOAD (different process). warmup_needed = not _ep_lazy_init_needed() and _dense_save_warmup_enabled() - restore_warmup = warmup_needed and _warmup_restore_enabled() + # Restore only when the archive actually carries the dump; otherwise + # (e.g. an archive captured before warmup-restore, or with it off) + # fall back to re-running the deterministic warmup forwards. + restore_warmup = ( + warmup_needed and _warmup_restore_enabled() and rt.warmup_region_available() + ) + if warmup_needed and _warmup_restore_enabled() and not restore_warmup: + logger.info( + "[Foundry] SGLang warmup-restore requested but no warmup_region " + "dump in the archive; re-running the warmup forwards instead." + ) skip_warmup = warmup_needed and not restore_warmup and _load_warmup_skip_enabled() if warmup_needed and not skip_warmup and not restore_warmup: _run_dense_warmup(self) @@ -567,10 +578,7 @@ def patched(self, *args, **kwargs): # buffers; decode is never run eagerly on LOAD (foundry replays # the captured graph), so nothing re-triggers the lazy init. if not rt.restore_warmup_region(): - raise RuntimeError( - "FOUNDRY_SGLANG_WARMUP_RESTORE=1 but no warmup_region.bin " - "in the archive (was SAVE run with the flag set?)" - ) + raise RuntimeError("warmup_region dump reported available but restore failed") rt.log_alloc_offset("after_warmup_restore") if skip_warmup: # Instead of re-running the warmup forwards, jump the cursor diff --git a/python/foundry/integration/sglang/runtime.py b/python/foundry/integration/sglang/runtime.py index abd9a77c..4d856d60 100644 --- a/python/foundry/integration/sglang/runtime.py +++ b/python/foundry/integration/sglang/runtime.py @@ -357,6 +357,20 @@ def dump_warmup_region(before_offset: int, after_offset: int) -> None: ) +def warmup_region_available() -> bool: + """Whether the archive contains a serialized warmup region to restore. + + Lets LOAD fall back to re-running the deterministic warmup forwards when an + archive was produced before warmup-restore existed (or with it disabled), + instead of failing.""" + cfg = get_config() + if cfg is None or cfg.workspace_dir is None: + return False + return os.path.exists(os.path.join(cfg.workspace_dir, _WARMUP_REGION_JSON)) and os.path.exists( + os.path.join(cfg.workspace_dir, _WARMUP_REGION_BIN) + ) + + def _load_final_alloc_offset() -> int: """Return the SAVE-side final alloc offset (the top of the LOAD-mapped region), or 0 if not recorded.""" From a3fe80fae6cf81eeff7dcda375b03de93d36b2c1 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 24 Jul 2026 01:46:43 +0000 Subject: [PATCH 21/22] [sglang] Clarify EP-gate comment now that dense/DP run their own warmup Co-Authored-By: Rahul Chalamala <22563365+rchalamala@users.noreply.github.com> Signed-off-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- python/foundry/integration/sglang/hooks.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/python/foundry/integration/sglang/hooks.py b/python/foundry/integration/sglang/hooks.py index 6963650d..7758a5f8 100644 --- a/python/foundry/integration/sglang/hooks.py +++ b/python/foundry/integration/sglang/hooks.py @@ -501,8 +501,9 @@ def _run_dense_warmup(self): def patched(self, *args, **kwargs): mode = get_graph_extension_mode() - # DeepEP/EP only (gated so the validated dense / single-GPU / DP paths - # are untouched): + # DeepEP/EP only (gated so dense / single-GPU / DP paths do not enter + # this EP-specific block; they get their own in-region warmup via + # _run_dense_warmup on both SAVE and LOAD, see below): # SAVE: sglang triggers pre-capture lazy init (DeepGEMM per-shape JIT, # DeepEP buffer, ...) via warmup forwards that foundry suppresses; # left lazy they fire inside the captured stream and abort From 4769c0f694bb497a3a68b81f715a051adc951929 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 24 Jul 2026 01:47:40 +0000 Subject: [PATCH 22/22] [sglang] README: warmup-restore now default-on; document fallback, clamp, frag test, upgrade note Co-Authored-By: Rahul Chalamala <22563365+rchalamala@users.noreply.github.com> Signed-off-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- recipe/experimental/README.md | 36 ++++++++++++++++++++++++++--------- 1 file changed, 27 insertions(+), 9 deletions(-) diff --git a/recipe/experimental/README.md b/recipe/experimental/README.md index 2f2419c0..7873d82e 100644 --- a/recipe/experimental/README.md +++ b/recipe/experimental/README.md @@ -254,21 +254,34 @@ Key model-specific requirements: lazy-init model buffers allocated during warmup land out-of-region and cause `CUDA error: illegal memory access` on LOAD graph replay. -### Optional: skip the LOAD warmup by restoring its bytes (`FOUNDRY_SGLANG_WARMUP_RESTORE=1`) +### Skip the LOAD warmup by restoring its bytes (`FOUNDRY_SGLANG_WARMUP_RESTORE`, default on) The in-region warmup re-runs real model forwards on LOAD (~20 s on 122B) only to reproduce the persistent buffers it wrote on SAVE. Because that warmup is deterministic, those bytes are identical on SAVE and LOAD, so LOAD can restore them instead of recomputing them. -Set `FOUNDRY_SGLANG_WARMUP_RESTORE=1` on **both** the SAVE run (to dump the gap) -and the LOAD run (to consume it). On SAVE, just before the post-warmup -`empty_cache`, foundry copies the warmup allocation gap +This is **on by default** (`FOUNDRY_SGLANG_WARMUP_RESTORE=1`). On SAVE, just +before the post-warmup `empty_cache`, foundry copies the warmup allocation gap (`[before_warmup, after_warmup)`) out of the VMM region to `rank_N/warmup_region.bin` (+ `.json` interval metadata). On LOAD it memcpys those bytes back into the same VMM addresses (after `preallocate_for_load_mode` maps the region) and sets the alloc cursor to the post-warmup offset — skipping -the warmup forwards entirely. +the warmup forwards entirely. LOAD only takes the restore path when the archive +actually carries the dump; otherwise (an archive captured with restore disabled) +it falls back to re-running the warmup forwards, so older archives still load. +Set `FOUNDRY_SGLANG_WARMUP_RESTORE=0` to force the re-run path on both SAVE and +LOAD. + +Robustness details: the restore is clamped to the LOAD-mapped region +(`final_alloc_offset`) — any warmup transient that peaks above the final +captured offset is dropped (it is never referenced by the replayed graphs), so a +larger-transient model can't drive an H2D copy into unmapped VA. The gap is +dumped **before** `empty_cache` (afterwards freed blocks are `cuMemUnmap`ped); +if the gap is fragmented, `_dump_mapped` falls back from one contiguous +`cuMemcpyDtoH` to per-2 MiB-page trial copies, keeping only cleanly-copied runs +(GPU-tested against a deliberately holed VMM range — see the `frag` entrypoint +in the validation harness). GPU-validated on `Qwen/Qwen3.5-122B-A10B`, H200:4, `--tp-size 4`: @@ -281,10 +294,15 @@ GPU-validated on `Qwen/Qwen3.5-122B-A10B`, H200:4, `--tp-size 4`: - **Archive cost** — 144 MiB/rank (one contiguous mapped interval) written during SAVE. -Off by default. Distinct from `FOUNDRY_SGLANG_SKIP_LOAD_WARMUP=1`, which merely -reserves the gap without restoring its contents — proven **incorrect** here -(coherent but divergent LOAD tokens), since the gap holds read-before-write -state. +Distinct from `FOUNDRY_SGLANG_SKIP_LOAD_WARMUP=1`, which merely reserves the gap +without restoring its contents — proven **incorrect** here (coherent but +divergent LOAD tokens), since the gap holds read-before-write state; that flag +stays off by default. + +**Upgrade note:** the dense/single-GPU/DP paths now run an in-region warmup +(they previously did not), which shifts their VMM cursor trajectory. An archive +captured before this change is not LOAD-compatible with a build after it — +re-capture (re-run SAVE) after upgrading. ## IPC-specific troubleshooting