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..7758a5f8 100644 --- a/python/foundry/integration/sglang/hooks.py +++ b/python/foundry/integration/sglang/hooks.py @@ -35,6 +35,67 @@ def _ep_lazy_init_needed() -> bool: return False +def _flashinfer_backend(attn_backend): + """Return the FlashInfer backend that owns the per-bs capture wrappers, or + None if this attention backend doesn't use FlashInfer's per-bs prepass. + + FlashInfer allocates a fresh per-bs metadata wrapper (each with its own + ``_int_workspace_buffer``) on every capture init, so foundry pre-allocates + them all up front and installs a reuse shim to keep the VMM cursor + deterministic vs LOAD. Hybrid linear-attn / mamba models (e.g. Qwen3.5) + wrap FlashInfer inside a ``HybridLinearAttnBackend`` whose full-attention + layers use FlashInfer while its linear-attn (mamba) layers keep all their + graph buffers in ``init_cuda_graph_state`` (pre-capture, so already + SAVE/LOAD-symmetric). Detect the FlashInfer backend by its per-bs + ``indices_updater_decode``, seeing through the hybrid wrapper's + ``full_attn_backend`` so the wrappers get the deterministic prepass too.""" + if hasattr(attn_backend, "indices_updater_decode"): + return attn_backend + if hasattr(attn_backend, "full_attn_backend"): + inner = attn_backend.full_attn_backend + if hasattr(inner, "indices_updater_decode"): + return inner + return None + + +def _dense_save_warmup_enabled() -> bool: + """Whether the dense/DP SAVE path runs a cursor-neutral pre-capture warmup. + + Defaults on: models that defer torch.compile / per-shape GEMM JIT to their + first forward (e.g. Qwen3.5 hybrid MoE) would otherwise trigger that lazy + init inside CUDA-graph capture and abort. Set ``FOUNDRY_SGLANG_SAVE_WARMUP=0`` + to restore the old suppress-only behavior.""" + return os.environ.get("FOUNDRY_SGLANG_SAVE_WARMUP", "1") not in ("0", "false", "False") + + +def _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 _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. + + 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. + 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: dp_rank = getattr(model_runner, "dp_rank", None) if dp_rank is not None: @@ -371,12 +432,78 @@ def _run_warmup_pass(self): time.perf_counter() - t0, ) + def _run_dense_warmup(self): + """In-region pre-capture warmup for the dense / DP path. + + Some models (e.g. Qwen3.5 hybrid MoE) defer torch.compile / inductor, + per-shape GEMM JIT, or Triton kernel compilation to their first forward. + On the foundry SAVE path that first forward is *inside* CUDA-graph capture + (sglang's own warmup forwards are suppressed for VMM determinism), where a + lazy CPU<->CUDA copy aborts capture ("Graph contains no nodes"). More + critically, hybrid models (MoE + mamba) lazily initialize persistent + module-level GPU buffers (e.g. Triton codegen workspace, torch.compile + inductor intermediates) on the first real forward — if those buffers land + outside the VMM region, the CUDA graph captures their out-of-region + addresses, which are stale on LOAD (different process → different addresses). + + Run one real forward per capture_bs with the allocation region ACTIVE so all + lazy-init persistent buffers land IN-REGION at deterministic VMM addresses. + The cursor advances during warmup; LOAD runs the same warmup (same model, + same forward, same capture_bs) to produce the same cursor trajectory. After + warmup, clear FlashInfer per-bs wrappers and empty the caching allocator + cache so the subsequent FlashInfer pre-pass re-allocates wrappers at the + post-warmup cursor — identical on both SAVE and LOAD. + """ + import torch + + fi_backend = _flashinfer_backend(self.attn_backend) + has_decode = fi_backend is not None and hasattr(fi_backend, "decode_cuda_graph_metadata") + has_prefill = fi_backend is not None and hasattr(fi_backend, "prefill_cuda_graph_metadata") + if has_decode: + fi_backend.decode_cuda_graph_metadata = {} + if has_prefill: + fi_backend.prefill_cuda_graph_metadata = {} + from foundry import ops as fops + + 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: + self.attn_backend.forward_metadata = None + if fi_backend is not None: + fi_backend.forward_metadata = None + # Drop warmup's FlashInfer wrappers so the real pre-pass + # re-allocates them at the post-warmup cursor. + if has_decode: + fi_backend.decode_cuda_graph_metadata = {} + if has_prefill: + fi_backend.prefill_cuda_graph_metadata = {} + # 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() + rt.log_alloc_offset("after_warmup") + # 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() + @functools.wraps(orig_capture) def patched(self, *args, **kwargs): mode = get_graph_extension_mode() - # DeepEP/EP only (gated so the validated dense / single-GPU / DP paths - # are untouched): + # 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 @@ -384,11 +511,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(): @@ -420,9 +547,56 @@ 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). + warmup_needed = not _ep_lazy_init_needed() and _dense_save_warmup_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) 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("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 + # 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`` @@ -432,7 +606,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 @@ -450,7 +630,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, @@ -462,17 +642,86 @@ def patched(self, *args, **kwargs): return None if mode == CUDAGraphExtensionMode.SAVE: + # Force model-level lazy init (torch.compile/inductor, per-shape GEMM + # JIT) outside capture for the dense/DP path. The warmup runs + # in-region so lazy-init persistent model buffers land at + # deterministic VMM addresses; LOAD runs the same warmup to + # reproduce the cursor trajectory. EP already warmed up above; + # skip there. + if not _ep_lazy_init_needed() and _dense_save_warmup_enabled(): + _run_dense_warmup(self) # FlashInfer allocates a per-bs metadata wrapper (each with its own # _int_workspace_buffer) on every capture init, so foundry pre-allocates # them up front and installs a reuse shim that makes the inner init # reuse them — keeping the VMM cursor deterministic vs LOAD. Backends # with a single fixed cuda-graph metadata workspace allocated once in # init_cuda_graph_state (e.g. fa3 / FlashAttentionBackend) don't need - # this and the plain capture is already SAVE/LOAD-deterministic. Detect - # FlashInfer by its per-bs indices_updater_decode. + # this and the plain capture is already SAVE/LOAD-deterministic. + # ``_flashinfer_backend`` returns the FlashInfer backend, seeing + # through the hybrid linear-attn wrapper (Qwen3.5) whose FlashInfer + # full-attn layers own the per-bs wrappers while its mamba layers + # keep all buffers in init_cuda_graph_state. The shim is installed on + # that FlashInfer backend directly; the hybrid wrapper's own + # init_forward_metadata_out_graph then dispatches into it during + # capture, so the mamba backend still runs its normal (alloc-free) + # metadata copy. attn_backend = self.attn_backend - use_fi_prepass = hasattr(attn_backend, "indices_updater_decode") - real_init = attn_backend.init_forward_metadata_capture_cuda_graph + fi_backend = _flashinfer_backend(attn_backend) + use_fi_prepass = fi_backend is not None + # sglang #26735 split the single capture-init entry point into the + # 3-method ABC; the per-iter capture path is now + # ``init_forward_metadata_out_graph(fb, in_capture=True)``. Detect + # which API this sglang exposes and shim the matching method. + use_new_api = hasattr(attn_backend, "init_forward_metadata_out_graph") + real_out_graph = ( + fi_backend.init_forward_metadata_out_graph + if use_new_api and use_fi_prepass + else None + ) + real_init = ( + fi_backend.init_forward_metadata_capture_cuda_graph + if not use_new_api and use_fi_prepass + else None + ) + + def reuse_out_graph(forward_batch, in_capture=False): + # Only the capture path (in_capture=True) must avoid + # re-allocating the per-bs wrappers the pre-pass already built; + # replay / eager (in_capture=False) pass straight through. + if not in_capture: + return real_out_graph(forward_batch, in_capture=in_capture) + + from sglang.srt.layers.attention.flashinfer_backend import ( + DecodeMetadata, + PrefillMetadata, + ) + + bs = forward_batch.batch_size + forward_mode = forward_batch.forward_mode + if forward_mode.is_decode_or_idle(): + wrappers = fi_backend.decode_cuda_graph_metadata.get(bs) + if wrappers is None: + return real_out_graph(forward_batch, in_capture=True) + fi_backend.forward_metadata = DecodeMetadata(wrappers) + elif ( + forward_mode.is_target_verify() + or forward_mode.is_draft_extend() + or forward_mode.is_dllm_extend() + ): + wrappers = fi_backend.prefill_cuda_graph_metadata.get(bs) + if wrappers is None: + return real_out_graph(forward_batch, in_capture=True) + fi_backend.forward_metadata = PrefillMetadata( + wrappers, forward_mode.is_dllm_extend(), False + ) + else: + return real_out_graph(forward_batch, in_capture=True) + # Re-run the planner against the pre-allocated wrappers with + # in_capture=False: no ``_prepare_cuda_graph_metadata`` (so no + # second torch.empty for ``_int_workspace_buffer``) and no + # begin_forward re-install (the pre-pass did that). We set + # forward_metadata above since that path skips it. + real_out_graph(forward_batch, in_capture=False) def reuse_pre_pass_init( bs, @@ -500,7 +749,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, @@ -512,7 +761,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(), @@ -521,16 +770,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, @@ -544,12 +793,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(), @@ -560,7 +809,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( @@ -584,13 +833,23 @@ def reuse_pre_pass_init( initialize_all_attention_metadata(self) rt.log_alloc_offset("save_after_pre_init") # Drop the pre-pass's last forward_metadata ref so popping the dict - # entry doesn't keep the wrapper alive at refcount 1. - attn_backend.forward_metadata = None - attn_backend.init_forward_metadata_capture_cuda_graph = reuse_pre_pass_init + # entry doesn't keep the wrapper alive at refcount 1. Installed on + # the FlashInfer backend (== attn_backend, or its full_attn_backend + # for the hybrid wrapper); the hybrid's own out_graph dispatches + # into the shimmed FlashInfer method during capture. + fi_backend.forward_metadata = None + if use_new_api: + fi_backend.init_forward_metadata_out_graph = reuse_out_graph + else: + fi_backend.init_forward_metadata_capture_cuda_graph = reuse_pre_pass_init try: result = orig_capture(self, *args, **kwargs) finally: - attn_backend.init_forward_metadata_capture_cuda_graph = real_init + if use_fi_prepass: + if use_new_api: + fi_backend.init_forward_metadata_out_graph = real_out_graph + else: + fi_backend.init_forward_metadata_capture_cuda_graph = real_init from foundry.integration.sglang.graph_ops import ( pack_fatbins, diff --git a/python/foundry/integration/sglang/runtime.py b/python/foundry/integration/sglang/runtime.py index 1f9032c5..4d856d60 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) @@ -180,18 +185,287 @@ 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) + + +_WARMUP_REGION_BIN = "warmup_region.bin" +_WARMUP_REGION_JSON = "warmup_region.json" +_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 + + lib_path = None + try: + with open("/proc/self/maps") as f: + for line in f: + 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("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 + + +# 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 _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]] = [] + parts: list[bytes] = [] + run_start = None + run_parts: list[bytes] = [] + off = 0 + while off < size: + 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 += length + if run_start is not None: + intervals.append((run_start, size - run_start)) + parts.append(b"".join(run_parts)) + return intervals, b"".join(parts) + + +def _cu_check(lib, rc: int, what: str) -> None: + if rc != 0: + 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: + """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.""" + 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 = _libcuda() + base_ptr = cfg.base_addr + before_offset + _cu_check(lib, lib.cuCtxSynchronize(), "cuCtxSynchronize(pre-dump)") + 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: + f.write(data) + 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, + "intervals": intervals, + }, + f, + ) + logger.info( + "[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, + ) + + +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.""" + 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. + + 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"] + 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) != total: + 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: + 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 + # 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 of %d bytes across " + "%d intervals (clamped to final=%d); cursor=%d", + before_offset, + after_offset, + restored, + total, + len(intervals), + final, + cursor, + ) + return True + + 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) diff --git a/recipe/experimental/README.md b/recipe/experimental/README.md index 23a092d9..7873d82e 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,124 @@ 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 **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 (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). +- **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 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. + +## 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. + +### 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. + +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. 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`: + +- **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. + +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 | 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.