[sglang] Experimental dense tensor-parallel recipe (NCCL all-reduce over VMM-IPC)

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],

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)