From a2f8e2e09728a8cf6faee5537f649c2d5c439c40 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Thu, 23 Jul 2026 06:15:24 +0000 Subject: [PATCH 01/21] test(sglang): define v0.5.15 graph contracts Co-authored-by: Rahul Chalamala --- tests/test_imports.py | 13 +- tests/test_sglang_api_contract.py | 90 ++++++++ tests/test_sglang_cuda_graph.py | 339 ++++++++++++++++++++++++++++++ tests/test_sglang_graph_ops.py | 231 ++++++++++++++++++++ 4 files changed, 668 insertions(+), 5 deletions(-) create mode 100644 tests/test_sglang_api_contract.py create mode 100644 tests/test_sglang_cuda_graph.py create mode 100644 tests/test_sglang_graph_ops.py diff --git a/tests/test_imports.py b/tests/test_imports.py index 2de75a6e..c916069b 100644 --- a/tests/test_imports.py +++ b/tests/test_imports.py @@ -32,6 +32,7 @@ "foundry.integration.sglang.hooks", "foundry.integration.sglang.runtime", "foundry.integration.sglang.graph_ops", + "foundry.integration.sglang.cuda_graph", ] # Names that have historically been ruff --fix victims because they live in @@ -76,9 +77,11 @@ def test_vllm_integration_public_api(): def test_sglang_integration_submodules_importable(): - # SGLang side currently exposes its API via submodules rather than - # re-exports on `foundry.integration.sglang`. Just verify each submodule - # imports cleanly — a ruff --fix that stripped a critical import would - # break this. - for sub in ("config", "hooks", "runtime", "graph_ops"): + for sub in ("config", "hooks", "runtime", "graph_ops", "cuda_graph"): importlib.import_module(f"foundry.integration.sglang.{sub}") + + +def test_sglang_integration_public_api(): + mod = importlib.import_module("foundry.integration.sglang") + for name in ("install_hooks", "CUDAGraphExtensionMode", "get_graph_extension_mode"): + assert hasattr(mod, name), f"foundry.integration.sglang.{name} not exported" diff --git a/tests/test_sglang_api_contract.py b/tests/test_sglang_api_contract.py new file mode 100644 index 00000000..7095c4a3 --- /dev/null +++ b/tests/test_sglang_api_contract.py @@ -0,0 +1,90 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the Foundry project +"""Import/signature contract for upstream SGLang v0.5.15.post1.""" + +from __future__ import annotations + +import importlib +import importlib.util +import inspect +from importlib.metadata import PackageNotFoundError, version + +import pytest + + +def _require_sglang(): + if importlib.util.find_spec("sglang") is None: + pytest.skip("SGLang is an optional dependency") + try: + installed = version("sglang") + except PackageNotFoundError: + pytest.skip("SGLang package metadata is unavailable") + assert installed == "0.5.15.post1" + + +def _parameters(callable_object): + return inspect.signature(callable_object).parameters + + +def test_latest_runner_and_backend_symbols_and_signatures(): + _require_sglang() + base_runner = importlib.import_module("sglang.srt.model_executor.runner.base_runner").BaseRunner + decode_runner = importlib.import_module( + "sglang.srt.model_executor.runner.decode_cuda_graph_runner" + ).DecodeCudaGraphRunner + full_backend = importlib.import_module( + "sglang.srt.model_executor.runner_backend.full_cuda_graph_backend" + ).FullCudaGraphBackend + + assert list(_parameters(base_runner.warmup)) == ["self"] + assert list(_parameters(decode_runner.capture)) == ["self"] + assert list(_parameters(decode_runner.capture_prepare)) == [ + "self", + "size", + "stream_idx", + ] + assert list(_parameters(full_backend.capture_one)) == [ + "self", + "shape_key", + "forward_fn", + "dummies", + "post_warmup_hook", + ] + + +def test_shape_key_has_complete_latest_identity(): + _require_sglang() + shape_key = importlib.import_module("sglang.srt.model_executor.runner.shape_key").ShapeKey + + parameters = _parameters(shape_key) + assert list(parameters) == ["size", "stream_idx", "variant_label"] + assert parameters["stream_idx"].default is None + assert parameters["variant_label"].default is None + assert shape_key(8, 1, "nolora") != shape_key(8, 0, "nolora") + + +def test_attention_backend_uses_split_graph_metadata_contract(): + _require_sglang() + attention_backend = importlib.import_module( + "sglang.srt.layers.attention.base_attn_backend" + ).AttentionBackend + + assert list(_parameters(attention_backend.init_forward_metadata_out_graph)) == [ + "self", + "forward_batch", + "in_capture", + ] + assert ( + _parameters(attention_backend.init_forward_metadata_out_graph)["in_capture"].default + is False + ) + assert list(_parameters(attention_backend.init_forward_metadata_in_graph)) == [ + "self", + "forward_batch", + ] + + +def test_removed_monolithic_cuda_graph_runner_is_not_part_of_contract(): + _require_sglang() + + assert importlib.util.find_spec("sglang.srt.model_executor.cuda_graph_runner") is None diff --git a/tests/test_sglang_cuda_graph.py b/tests/test_sglang_cuda_graph.py new file mode 100644 index 00000000..8a5666f3 --- /dev/null +++ b/tests/test_sglang_cuda_graph.py @@ -0,0 +1,339 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the Foundry project +"""Pure-Python contracts for the SGLang v0.5.15.post1 graph hooks.""" + +from __future__ import annotations + +import importlib +import sys +from contextlib import contextmanager +from dataclasses import dataclass +from enum import Enum +from pathlib import Path +from types import ModuleType, SimpleNamespace + +import pytest + + +class Mode(str, Enum): + NONE = "none" + SAVE = "save" + LOAD = "load" + + +@dataclass(frozen=True) +class ShapeKey: + size: int + stream_idx: int | None = None + variant_label: str | None = None + + +class BaseRunner: + def __init__(self): + self.warmup_calls = 0 + + def warmup(self): + self.warmup_calls += 1 + return "warmed" + + +class DecodeCudaGraphRunner(BaseRunner): + """Small latest-API runner; deliberately has no legacy CudaGraphRunner API.""" + + def __init__(self): + super().__init__() + self.capture_calls = 0 + self.prepared = [] + + def capture_prepare(self, size, stream_idx=None): + key = ShapeKey(size=size, stream_idx=stream_idx) + self.prepared.append(key) + return key + + def capture(self): + self.capture_calls += 1 + return "captured" + + +class FullCudaGraphBackend: + def __init__(self): + self._graphs = {} + self._outputs = {} + self._pool = None + self.capture_one_calls = 0 + + @contextmanager + def capture_session(self, stream): + del stream + yield + + def capture_one( + self, + shape_key, + forward_fn, + dummies=None, + post_warmup_hook=None, + ): + del dummies, post_warmup_hook + self.capture_one_calls += 1 + self._graphs[shape_key] = object() + self._outputs[shape_key] = forward_fn() + + +_ORIGINAL_WARMUP = BaseRunner.warmup +_ORIGINAL_CAPTURE = DecodeCudaGraphRunner.capture +_ORIGINAL_CAPTURE_PREPARE = DecodeCudaGraphRunner.capture_prepare +_ORIGINAL_CAPTURE_ONE = FullCudaGraphBackend.capture_one + + +def _module(name: str, **attributes): + module = ModuleType(name) + for attr, value in attributes.items(): + setattr(module, attr, value) + return module + + +def _load_cuda_graph(monkeypatch: pytest.MonkeyPatch): + BaseRunner.warmup = _ORIGINAL_WARMUP + DecodeCudaGraphRunner.capture = _ORIGINAL_CAPTURE + DecodeCudaGraphRunner.capture_prepare = _ORIGINAL_CAPTURE_PREPARE + FullCudaGraphBackend.capture_one = _ORIGINAL_CAPTURE_ONE + + source_root = Path(__file__).parents[1] / "python" / "foundry" + mode = SimpleNamespace(value=Mode.NONE) + registered = SimpleNamespace(global_pool=None, pynccl_pool=None) + + foundry = _module("foundry") + foundry.__path__ = [str(source_root)] + integration = _module("foundry.integration") + integration.__path__ = [str(source_root / "integration")] + sglang_integration = _module("foundry.integration.sglang") + sglang_integration.__path__ = [str(source_root / "integration" / "sglang")] + + config = _module( + "foundry.integration.sglang.config", + CUDAGraphExtensionMode=Mode, + get_graph_extension_mode=lambda: mode.value, + ) + graph_ops = _module("foundry.integration.sglang.graph_ops") + + runner_package = _module("sglang.srt.model_executor.runner") + runner_package.__path__ = [] + base_runner = _module( + "sglang.srt.model_executor.runner.base_runner", + BaseRunner=BaseRunner, + ) + decode_runner = _module( + "sglang.srt.model_executor.runner.decode_cuda_graph_runner", + DecodeCudaGraphRunner=DecodeCudaGraphRunner, + ) + shape_key = _module( + "sglang.srt.model_executor.runner.shape_key", + ShapeKey=ShapeKey, + ) + full_backend = _module( + "sglang.srt.model_executor.runner_backend.full_cuda_graph_backend", + FullCudaGraphBackend=FullCudaGraphBackend, + ) + pool = _module( + "sglang.srt.model_executor.runner_utils.pool", + set_global_graph_memory_pool=lambda value: setattr(registered, "global_pool", value), + ) + pynccl = _module( + "sglang.srt.distributed.device_communicators.pynccl_allocator", + set_graph_pool_id=lambda value: setattr(registered, "pynccl_pool", value), + ) + + modules = { + "foundry": foundry, + "foundry.integration": integration, + "foundry.integration.sglang": sglang_integration, + "foundry.integration.sglang.config": config, + "foundry.integration.sglang.graph_ops": graph_ops, + "sglang": _module("sglang"), + "sglang.srt": _module("sglang.srt"), + "sglang.srt.model_executor": _module("sglang.srt.model_executor"), + "sglang.srt.model_executor.runner": runner_package, + "sglang.srt.model_executor.runner.base_runner": base_runner, + "sglang.srt.model_executor.runner.decode_cuda_graph_runner": decode_runner, + "sglang.srt.model_executor.runner.shape_key": shape_key, + "sglang.srt.model_executor.runner_backend": _module( + "sglang.srt.model_executor.runner_backend" + ), + "sglang.srt.model_executor.runner_backend.full_cuda_graph_backend": full_backend, + "sglang.srt.model_executor.runner_utils": _module("sglang.srt.model_executor.runner_utils"), + "sglang.srt.model_executor.runner_utils.pool": pool, + "sglang.srt.distributed": _module("sglang.srt.distributed"), + "sglang.srt.distributed.device_communicators": _module( + "sglang.srt.distributed.device_communicators" + ), + "sglang.srt.distributed.device_communicators.pynccl_allocator": pynccl, + } + for name, module in modules.items(): + monkeypatch.setitem(sys.modules, name, module) + + cuda_graph = importlib.import_module("foundry.integration.sglang.cuda_graph") + return cuda_graph, mode, registered + + +def _server_args(**overrides): + values = { + "cuda_graph_config": SimpleNamespace( + decode=SimpleNamespace(backend="full"), + prefill=SimpleNamespace(backend="disabled"), + ), + "tp_size": 1, + "pp_size": 1, + "dp_size": 1, + "enable_dp_attention": False, + "attention_backend": "flashinfer", + "moe_a2a_backend": "none", + "speculative_algorithm": None, + "enable_lora": False, + "enable_pdmux": False, + "dllm_algorithm": None, + "enable_return_hidden_states": False, + "enable_memory_saver": False, + } + values.update(overrides) + return SimpleNamespace(**values) + + +@pytest.mark.parametrize( + "server_args", + [ + _server_args(), + _server_args(dp_size=2), + _server_args( + tp_size=2, + enable_dp_attention=True, + attention_backend="fa3", + moe_a2a_backend="deepep", + ), + ], +) +def test_configuration_accepts_supported_full_decode_modes(monkeypatch, server_args): + cuda_graph, _mode, _registered = _load_cuda_graph(monkeypatch) + + cuda_graph.validate_configuration(server_args) + + +@pytest.mark.parametrize( + ("overrides", "diagnostic"), + [ + ( + { + "cuda_graph_config": SimpleNamespace( + decode=SimpleNamespace(backend="breakable"), + prefill=SimpleNamespace(backend="disabled"), + ) + }, + "decode.*full", + ), + ( + { + "cuda_graph_config": SimpleNamespace( + decode=SimpleNamespace(backend="tc_piecewise"), + prefill=SimpleNamespace(backend="disabled"), + ) + }, + "decode.*full", + ), + ( + { + "cuda_graph_config": SimpleNamespace( + decode=SimpleNamespace(backend="disabled"), + prefill=SimpleNamespace(backend="disabled"), + ) + }, + "decode.*full", + ), + ( + { + "cuda_graph_config": SimpleNamespace( + decode=SimpleNamespace(backend="full"), + prefill=SimpleNamespace(backend="full"), + ) + }, + "prefill.*disabled", + ), + ({"tp_size": 2}, "tensor parallel|tp_size"), + ({"pp_size": 2}, "pipeline parallel|pp_size"), + ({"speculative_algorithm": "EAGLE"}, "speculative"), + ({"enable_lora": True}, "LoRA"), + ({"enable_pdmux": True}, "PDMux"), + ({"dllm_algorithm": "dream"}, "dLLM|diffusion"), + ({"enable_return_hidden_states": True}, "hidden"), + ({"enable_memory_saver": True}, "memory.saver"), + ({"attention_backend": "triton"}, "FlashInfer|FA3|attention"), + ], +) +def test_configuration_rejects_unsupported_modes(monkeypatch, overrides, diagnostic): + cuda_graph, _mode, _registered = _load_cuda_graph(monkeypatch) + + with pytest.raises((RuntimeError, ValueError), match=diagnostic): + cuda_graph.validate_configuration(_server_args(**overrides)) + + +def test_base_runner_warmup_is_suppressed_only_for_save_and_load(monkeypatch): + cuda_graph, mode, _registered = _load_cuda_graph(monkeypatch) + cuda_graph.install_cuda_graph_hooks() + runner = BaseRunner() + + mode.value = Mode.NONE + assert runner.warmup() == "warmed" + mode.value = Mode.SAVE + assert runner.warmup() is None + mode.value = Mode.LOAD + assert runner.warmup() is None + assert runner.warmup_calls == 1 + + +def test_hook_installer_targets_latest_split_runner_classes(monkeypatch): + original_capture = DecodeCudaGraphRunner.capture + original_capture_one = FullCudaGraphBackend.capture_one + cuda_graph, mode, _registered = _load_cuda_graph(monkeypatch) + + cuda_graph.install_cuda_graph_hooks() + + assert DecodeCudaGraphRunner.capture is not original_capture + assert FullCudaGraphBackend.capture_one is not original_capture_one + mode.value = Mode.NONE + runner = DecodeCudaGraphRunner() + assert runner.capture() == "captured" + assert runner.capture_prepare(6, stream_idx=3) == ShapeKey(size=6, stream_idx=3) + assert "sglang.srt.model_executor.cuda_graph_runner" not in sys.modules + + +def test_load_hydrates_full_backend_with_complete_shape_keys(monkeypatch): + cuda_graph, _mode, _registered = _load_cuda_graph(monkeypatch) + backend = FullCudaGraphBackend() + pool = object() + graph_a, graph_b = object(), object() + output_a, output_b = object(), object() + key_a = ShapeKey(size=8, stream_idx=0, variant_label="lora") + key_b = ShapeKey(size=8, stream_idx=1, variant_label="nolora") + + cuda_graph.hydrate_backend( + backend, + [ + (key_a, graph_a, output_a), + (key_b, graph_b, output_b), + ], + pool=pool, + ) + + assert backend._graphs == {key_a: graph_a, key_b: graph_b} + assert backend._outputs == {key_a: output_a, key_b: output_b} + assert backend._pool is pool + + +def test_pool_registration_uses_latest_split_runner_api(monkeypatch): + cuda_graph, _mode, registered = _load_cuda_graph(monkeypatch) + pool = object() + + cuda_graph.register_graph_pool(pool) + + assert registered.global_pool is pool + assert registered.pynccl_pool is pool + assert "sglang.srt.model_executor.cuda_graph_runner" not in sys.modules diff --git a/tests/test_sglang_graph_ops.py b/tests/test_sglang_graph_ops.py new file mode 100644 index 00000000..a6513e85 --- /dev/null +++ b/tests/test_sglang_graph_ops.py @@ -0,0 +1,231 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the Foundry project +"""Contracts for the versioned SGLang CUDA-graph archive index.""" + +from __future__ import annotations + +import importlib.util +import json +import sys +from dataclasses import dataclass +from pathlib import Path +from types import ModuleType + +import pytest + + +@dataclass(frozen=True) +class ShapeKey: + """The public key shape in SGLang v0.5.15.post1.""" + + size: int + stream_idx: int | None = None + variant_label: str | None = None + + +def _load_graph_ops(monkeypatch: pytest.MonkeyPatch): + """Load graph_ops without requiring the native Foundry extension.""" + + foundry = ModuleType("foundry") + foundry.__path__ = [] + foundry.save_graph_manifest = lambda _path: None + + ops = ModuleType("foundry.ops") + graph = ModuleType("foundry.graph") + graph.CUDAGraph = object + graph.graph = lambda *_args, **_kwargs: None + + config = ModuleType("foundry.integration.sglang.config") + config.CUDAGraphExtensionMode = object + config.get_config = lambda: None + config.get_graph_extension_mode = lambda: None + + runtime = ModuleType("foundry.integration.sglang.runtime") + runtime.get_state = lambda: None + + torch = ModuleType("torch") + torch.Tensor = object + torch.cuda = type("FakeCuda", (), {"CUDAGraph": object})() + + for name, module in ( + ("torch", torch), + ("foundry", foundry), + ("foundry.ops", ops), + ("foundry.graph", graph), + ("foundry.integration.sglang.config", config), + ("foundry.integration.sglang.runtime", runtime), + ): + monkeypatch.setitem(sys.modules, name, module) + + path = ( + Path(__file__).parents[1] / "python" / "foundry" / "integration" / "sglang" / "graph_ops.py" + ) + spec = importlib.util.spec_from_file_location("_foundry_sglang_graph_ops_contract", path) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def _entry(graph_ops, **overrides): + values = { + "phase": "decode", + "capture_index": 3, + "filename": "opaque-name.json", + "shape_key": ShapeKey(size=17, stream_idx=2, variant_label="nolora"), + "output_kind": "next_token_logits", + } + values.update(overrides) + return graph_ops.GraphIndexEntry(**values) + + +def test_versioned_graph_index_round_trips_complete_metadata(tmp_path, monkeypatch): + graph_ops = _load_graph_ops(monkeypatch) + index_path = tmp_path / "sglang_graph_index.json" + entry = _entry(graph_ops) + + graph_ops.save_graph_index(index_path, [entry]) + + payload = json.loads(index_path.read_text()) + assert payload["schema_version"] == graph_ops.GRAPH_INDEX_SCHEMA_VERSION + assert graph_ops.load_graph_index(index_path) == [entry] + assert payload["graphs"][0] == { + "phase": "decode", + "capture_index": 3, + "filename": "opaque-name.json", + "shape_key": { + "size": 17, + "stream_idx": 2, + "variant_label": "nolora", + }, + "output_kind": "next_token_logits", + } + + +@pytest.mark.parametrize("schema_version", [None, 999]) +def test_graph_index_rejects_legacy_or_unsupported_schema(tmp_path, monkeypatch, schema_version): + graph_ops = _load_graph_ops(monkeypatch) + index_path = tmp_path / "sglang_graph_index.json" + payload = {"graphs": []} + if schema_version is not None: + payload["schema_version"] = schema_version + index_path.write_text(json.dumps(payload)) + + with pytest.raises((RuntimeError, ValueError), match=r"(?i)fresh.*SAVE|SAVE.*fresh"): + graph_ops.load_graph_index(index_path) + + +def test_shape_identity_does_not_collapse_stream_or_variant(tmp_path, monkeypatch): + graph_ops = _load_graph_ops(monkeypatch) + index_path = tmp_path / "sglang_graph_index.json" + entries = [ + _entry( + graph_ops, + capture_index=0, + filename="first.json", + shape_key=ShapeKey(size=8, stream_idx=0, variant_label="lora"), + ), + _entry( + graph_ops, + capture_index=1, + filename="second.json", + shape_key=ShapeKey(size=8, stream_idx=1, variant_label="nolora"), + ), + ] + + graph_ops.save_graph_index(index_path, entries) + loaded = graph_ops.load_graph_index(index_path) + + assert loaded == entries + assert loaded[0].shape_key != loaded[1].shape_key + + +def test_graph_index_load_order_is_capture_index_not_filename(tmp_path, monkeypatch): + graph_ops = _load_graph_ops(monkeypatch) + index_path = tmp_path / "sglang_graph_index.json" + later = _entry( + graph_ops, + capture_index=1, + filename="aaa.json", + shape_key=ShapeKey(size=4), + ) + earlier = _entry( + graph_ops, + capture_index=0, + filename="zzz.json", + shape_key=ShapeKey(size=8), + ) + + graph_ops.save_graph_index(index_path, [later, earlier]) + + assert graph_ops.load_graph_index(index_path) == [earlier, later] + + +@pytest.mark.parametrize("legacy_key", [8, "1_8"]) +def test_graph_index_rejects_integer_and_string_shape_key_heuristics( + tmp_path, monkeypatch, legacy_key +): + graph_ops = _load_graph_ops(monkeypatch) + index_path = tmp_path / "sglang_graph_index.json" + index_path.write_text( + json.dumps( + { + "schema_version": graph_ops.GRAPH_INDEX_SCHEMA_VERSION, + "graphs": [ + { + "phase": "decode", + "capture_index": 0, + "filename": "graph_0_FULL_t8_r8_UX_pcN.json", + "shape_key": legacy_key, + "output_kind": "next_token_logits", + } + ], + } + ) + ) + + with pytest.raises((RuntimeError, TypeError, ValueError), match=r"(?i)fresh.*SAVE|shape"): + graph_ops.load_graph_index(index_path) + + +def test_index_metadata_not_filename_regex_is_source_of_truth(tmp_path, monkeypatch): + graph_ops = _load_graph_ops(monkeypatch) + index_path = tmp_path / "sglang_graph_index.json" + entry = _entry( + graph_ops, + filename="graph_42_FULL_t999_r999_UX_pcN.json", + shape_key=ShapeKey(size=7, stream_idx=None, variant_label=None), + ) + + graph_ops.save_graph_index(index_path, [entry]) + loaded = graph_ops.load_graph_index(index_path) + + assert loaded[0].shape_key == ShapeKey(size=7, stream_idx=None, variant_label=None) + + +def test_unsupported_output_kind_is_rejected(tmp_path, monkeypatch): + graph_ops = _load_graph_ops(monkeypatch) + index_path = tmp_path / "sglang_graph_index.json" + index_path.write_text( + json.dumps( + { + "schema_version": graph_ops.GRAPH_INDEX_SCHEMA_VERSION, + "graphs": [ + { + "phase": "decode", + "capture_index": 0, + "filename": "opaque.json", + "shape_key": { + "size": 4, + "stream_idx": None, + "variant_label": None, + }, + "output_kind": "hidden_states", + } + ], + } + ) + ) + + with pytest.raises((RuntimeError, TypeError, ValueError), match=r"(?i)output.*kind|next_token"): + graph_ops.load_graph_index(index_path) From f4137ca32f6883ca0425f8c8ee5fccad7a30a476 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Thu, 23 Jul 2026 06:27:31 +0000 Subject: [PATCH 02/21] test(sglang): compare archive records structurally Co-authored-by: Rahul Chalamala --- tests/test_sglang_cuda_graph.py | 1 - tests/test_sglang_graph_ops.py | 78 ++++++++++++++++++++++++++++++--- 2 files changed, 73 insertions(+), 6 deletions(-) diff --git a/tests/test_sglang_cuda_graph.py b/tests/test_sglang_cuda_graph.py index 8a5666f3..5233771a 100644 --- a/tests/test_sglang_cuda_graph.py +++ b/tests/test_sglang_cuda_graph.py @@ -301,7 +301,6 @@ def test_hook_installer_targets_latest_split_runner_classes(monkeypatch): mode.value = Mode.NONE runner = DecodeCudaGraphRunner() assert runner.capture() == "captured" - assert runner.capture_prepare(6, stream_idx=3) == ShapeKey(size=6, stream_idx=3) assert "sglang.srt.model_executor.cuda_graph_runner" not in sys.modules diff --git a/tests/test_sglang_graph_ops.py b/tests/test_sglang_graph_ops.py index a6513e85..42a92956 100644 --- a/tests/test_sglang_graph_ops.py +++ b/tests/test_sglang_graph_ops.py @@ -79,6 +79,26 @@ def _entry(graph_ops, **overrides): return graph_ops.GraphIndexEntry(**values) +def _assert_entry( + actual, + *, + phase, + capture_index, + filename, + size, + stream_idx, + variant_label, + output_kind="next_token_logits", +): + assert actual.phase == phase + assert actual.capture_index == capture_index + assert actual.filename == filename + assert actual.output_kind == output_kind + assert actual.shape_key.size == size + assert actual.shape_key.stream_idx == stream_idx + assert actual.shape_key.variant_label == variant_label + + def test_versioned_graph_index_round_trips_complete_metadata(tmp_path, monkeypatch): graph_ops = _load_graph_ops(monkeypatch) index_path = tmp_path / "sglang_graph_index.json" @@ -88,7 +108,17 @@ def test_versioned_graph_index_round_trips_complete_metadata(tmp_path, monkeypat payload = json.loads(index_path.read_text()) assert payload["schema_version"] == graph_ops.GRAPH_INDEX_SCHEMA_VERSION - assert graph_ops.load_graph_index(index_path) == [entry] + loaded = graph_ops.load_graph_index(index_path) + assert len(loaded) == 1 + _assert_entry( + loaded[0], + phase="decode", + capture_index=3, + filename="opaque-name.json", + size=17, + stream_idx=2, + variant_label="nolora", + ) assert payload["graphs"][0] == { "phase": "decode", "capture_index": 3, @@ -136,8 +166,34 @@ def test_shape_identity_does_not_collapse_stream_or_variant(tmp_path, monkeypatc graph_ops.save_graph_index(index_path, entries) loaded = graph_ops.load_graph_index(index_path) - assert loaded == entries - assert loaded[0].shape_key != loaded[1].shape_key + assert len(loaded) == 2 + _assert_entry( + loaded[0], + phase="decode", + capture_index=0, + filename="first.json", + size=8, + stream_idx=0, + variant_label="lora", + ) + _assert_entry( + loaded[1], + phase="decode", + capture_index=1, + filename="second.json", + size=8, + stream_idx=1, + variant_label="nolora", + ) + assert ( + loaded[0].shape_key.size, + loaded[0].shape_key.stream_idx, + loaded[0].shape_key.variant_label, + ) != ( + loaded[1].shape_key.size, + loaded[1].shape_key.stream_idx, + loaded[1].shape_key.variant_label, + ) def test_graph_index_load_order_is_capture_index_not_filename(tmp_path, monkeypatch): @@ -158,7 +214,10 @@ def test_graph_index_load_order_is_capture_index_not_filename(tmp_path, monkeypa graph_ops.save_graph_index(index_path, [later, earlier]) - assert graph_ops.load_graph_index(index_path) == [earlier, later] + loaded = graph_ops.load_graph_index(index_path) + assert [entry.capture_index for entry in loaded] == [0, 1] + assert [entry.filename for entry in loaded] == ["zzz.json", "aaa.json"] + assert [entry.shape_key.size for entry in loaded] == [8, 4] @pytest.mark.parametrize("legacy_key", [8, "1_8"]) @@ -200,7 +259,16 @@ def test_index_metadata_not_filename_regex_is_source_of_truth(tmp_path, monkeypa graph_ops.save_graph_index(index_path, [entry]) loaded = graph_ops.load_graph_index(index_path) - assert loaded[0].shape_key == ShapeKey(size=7, stream_idx=None, variant_label=None) + assert len(loaded) == 1 + _assert_entry( + loaded[0], + phase="decode", + capture_index=3, + filename="graph_42_FULL_t999_r999_UX_pcN.json", + size=7, + stream_idx=None, + variant_label=None, + ) def test_unsupported_output_kind_is_rejected(tmp_path, monkeypatch): From f271e0f3d94ca802bb9222e4efb84105dcb1c769 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Thu, 23 Jul 2026 06:43:05 +0000 Subject: [PATCH 03/21] feat(sglang): port hooks to split CUDA graph runner Co-authored-by: Rahul Chalamala --- python/foundry/integration/sglang/__init__.py | 12 + .../foundry/integration/sglang/cuda_graph.py | 412 +++++++++++++++++ .../foundry/integration/sglang/graph_ops.py | 132 +++--- python/foundry/integration/sglang/hooks.py | 420 +----------------- 4 files changed, 502 insertions(+), 474 deletions(-) create mode 100644 python/foundry/integration/sglang/cuda_graph.py diff --git a/python/foundry/integration/sglang/__init__.py b/python/foundry/integration/sglang/__init__.py index 7a771998..ab4b28b0 100644 --- a/python/foundry/integration/sglang/__init__.py +++ b/python/foundry/integration/sglang/__init__.py @@ -1,3 +1,15 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the Foundry project """Foundry integration for SGLang.""" + +from foundry.integration.sglang.config import ( + CUDAGraphExtensionMode, + get_graph_extension_mode, +) +from foundry.integration.sglang.hooks import install_hooks + +__all__ = [ + "CUDAGraphExtensionMode", + "get_graph_extension_mode", + "install_hooks", +] diff --git a/python/foundry/integration/sglang/cuda_graph.py b/python/foundry/integration/sglang/cuda_graph.py new file mode 100644 index 00000000..7267d5b6 --- /dev/null +++ b/python/foundry/integration/sglang/cuda_graph.py @@ -0,0 +1,412 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the Foundry project +"""Foundry hooks for SGLang's v0.5.15.post1 split CUDA-graph runner. + +This module is the optional-dependency boundary for SGLang. It imports +SGLang's private runner APIs at module scope, so callers must load it only +while installing hooks in a process where the supported SGLang version is +available. The public ``foundry.integration.sglang`` package does not import +this module eagerly. +""" + +from __future__ import annotations + +import functools +import importlib +import logging +import sys +import time +from collections.abc import Callable, Iterable, Iterator +from contextlib import contextmanager +from typing import Any + +from sglang.srt.model_executor.runner.base_runner import BaseRunner +from sglang.srt.model_executor.runner.decode_cuda_graph_runner import ( + DecodeCudaGraphRunner, +) +from sglang.srt.model_executor.runner.shape_key import ShapeKey +from sglang.srt.model_executor.runner_backend.full_cuda_graph_backend import ( + FullCudaGraphBackend, +) + +from foundry.integration.sglang import graph_ops +from foundry.integration.sglang.config import ( + CUDAGraphExtensionMode, +) + +logger = logging.getLogger(__name__) + +_ORIGINAL_ATTR = "__foundry_sglang_original__" +_EAGER_DEEPEP_WARMUP = False + + +def _module(name: str): + """Resolve a boundary module, including test/runtime replacements.""" + + return sys.modules.get(name) or importlib.import_module(name) + + +def _mode() -> CUDAGraphExtensionMode: + config = _module("foundry.integration.sglang.config") + return config.get_graph_extension_mode() + + +def _value(value: Any) -> Any: + """Return the value of an enum-like setting without requiring its type.""" + + return getattr(value, "value", value) + + +def _decode_attention_backend(server_args: Any) -> str | None: + backend = getattr(server_args, "decode_attention_backend", None) + if backend is None: + backend = getattr(server_args, "attention_backend", None) + return _value(backend) + + +def validate_configuration(server_args: Any) -> None: + """Reject SGLang configurations outside Foundry's validated graph scope.""" + + cuda_graph_config = getattr(server_args, "cuda_graph_config", None) + decode = getattr(getattr(cuda_graph_config, "decode", None), "backend", None) + prefill = getattr(getattr(cuda_graph_config, "prefill", None), "backend", None) + if _value(decode) != "full": + raise ValueError( + f"Foundry requires cuda_graph_config decode backend 'full'; got {_value(decode)!r}." + ) + if _value(prefill) != "disabled": + raise ValueError( + "Foundry requires cuda_graph_config prefill backend 'disabled'; " + f"got {_value(prefill)!r}." + ) + + tp_size = getattr(server_args, "tp_size", 1) + dp_attention = bool(getattr(server_args, "enable_dp_attention", False)) + if tp_size > 1 and not dp_attention: + raise ValueError( + "Foundry does not support tensor parallel decode without DP attention " + f"(tp_size={tp_size})." + ) + if getattr(server_args, "pp_size", 1) > 1: + raise ValueError("Foundry does not support pipeline parallel CUDA graphs (pp_size > 1).") + if getattr(server_args, "speculative_algorithm", None) is not None: + raise ValueError("Foundry does not support speculative decoding.") + if bool(getattr(server_args, "enable_lora", False)) or bool( + getattr(server_args, "lora_paths", None) + ): + raise ValueError("Foundry does not support LoRA CUDA graphs.") + if bool(getattr(server_args, "enable_pdmux", False)): + raise ValueError("Foundry does not support PDMux CUDA graphs.") + if getattr(server_args, "dllm_algorithm", None) is not None: + raise ValueError("Foundry does not support dLLM or diffusion CUDA graphs.") + if bool(getattr(server_args, "enable_return_hidden_states", False)): + raise ValueError("Foundry does not support hidden-state capture.") + if bool(getattr(server_args, "enable_memory_saver", False)): + raise ValueError("Foundry does not support memory-saver CUDA graphs.") + if bool(getattr(server_args, "debug_cuda_graph", False)): + raise ValueError("Foundry requires the full backend and cannot use debug CUDA graphs.") + + attention_backend = _decode_attention_backend(server_args) + moe_backend = _value(getattr(server_args, "moe_a2a_backend", "none")) + if dp_attention: + if attention_backend != "fa3" or moe_backend != "deepep": + raise ValueError( + "Foundry DP-attention requires the FA3 attention backend and DeepEP MoE all-to-all." + ) + deepep_mode = _value(getattr(server_args, "deepep_mode", "auto")) + if deepep_mode not in ("auto", "low_latency"): + raise ValueError( + "Foundry DP-attention requires DeepEP low-latency decode mode " + "('--deepep-mode auto' or 'low_latency')." + ) + elif attention_backend != "flashinfer": + raise ValueError( + "Foundry single-GPU and regular-DP decode require the FlashInfer " + f"attention backend; got {attention_backend!r}." + ) + elif moe_backend != "none": + raise ValueError( + "Foundry only supports a MoE all-to-all backend with the " + "DP-attention + DeepEP + FA3 configuration." + ) + + +def register_graph_pool(pool: Any) -> None: + """Register one pool with both latest split-runner pool registries.""" + + pool_module = _module("sglang.srt.model_executor.runner_utils.pool") + pynccl_allocator = _module("sglang.srt.distributed.device_communicators.pynccl_allocator") + pool_module.set_global_graph_memory_pool(pool) + pynccl_allocator.set_graph_pool_id(pool) + + +def hydrate_backend( + backend: FullCudaGraphBackend, + loaded: Iterable[tuple[ShapeKey, Any, Any]], + *, + pool: Any, +) -> None: + """Install loaded graphs and outputs in a full split-runner backend.""" + + if not isinstance(backend, FullCudaGraphBackend): + raise RuntimeError( + f"Foundry LOAD requires SGLang's FullCudaGraphBackend; got {type(backend).__name__}." + ) + backend._graphs.clear() + backend._outputs.clear() + for shape_key, graph, output in loaded: + if not isinstance(shape_key, ShapeKey): + raise TypeError( + "Foundry LOAD requires complete SGLang ShapeKey objects, " + f"got {type(shape_key).__name__}." + ) + backend._graphs[shape_key] = graph + backend._outputs[shape_key] = output + backend._pool = pool + + +def _mark_wrapper(wrapper: Callable[..., Any], original: Callable[..., Any]): + setattr(wrapper, _ORIGINAL_ATTR, original) + return wrapper + + +def _patch_once(cls: type, name: str, factory: Callable[[Callable[..., Any]], Callable[..., Any]]): + current = getattr(cls, name) + if hasattr(current, _ORIGINAL_ATTR): + return + setattr(cls, name, _mark_wrapper(factory(current), current)) + + +def _warmup_wrapper(original: Callable[..., Any]): + @functools.wraps(original) + def patched(self): + mode = _mode() + if mode == CUDAGraphExtensionMode.NONE: + return original(self) + logger.info("[Foundry] SGLang runner warmup skipped in %s mode", mode.value) + return None + + return patched + + +def _ensure_full_backend(runner: DecodeCudaGraphRunner) -> FullCudaGraphBackend: + backend = runner.backend + if not isinstance(backend, FullCudaGraphBackend): + raise RuntimeError( + "Foundry requires SGLang's FullCudaGraphBackend and will not " + f"recapture with {type(backend).__name__}." + ) + return backend + + +def _uses_deepep(runner: DecodeCudaGraphRunner) -> bool: + return _value(getattr(runner.model_runner.server_args, "moe_a2a_backend", "none")) == "deepep" + + +@contextmanager +def _deepep_eager_warmup() -> Iterator[None]: + global _EAGER_DEEPEP_WARMUP + previous = _EAGER_DEEPEP_WARMUP + _EAGER_DEEPEP_WARMUP = True + try: + yield + finally: + _EAGER_DEEPEP_WARMUP = previous + + +def _is_flashinfer_backend(attn_backend: Any) -> bool: + return hasattr(attn_backend, "indices_updater_decode") and hasattr( + attn_backend, "decode_cuda_graph_metadata" + ) + + +def _prepare_attention_metadata( + runner: DecodeCudaGraphRunner, +) -> list[tuple[Any, int, Any]]: + """Create per-shape metadata in upstream's reverse capture order.""" + + prepared = [] + for size in reversed(runner.capture_bs): + forward_batch, attn_backend, _pp_proxy_tensors = runner.capture_prepare(size) + attn_backend.init_forward_metadata_out_graph( + forward_batch, + in_capture=True, + ) + prepared.append((attn_backend, size, attn_backend.forward_metadata)) + return prepared + + +@contextmanager +def _reuse_flashinfer_metadata( + prepared: list[tuple[Any, int, Any]], +) -> Iterator[None]: + """Reuse pre-created wrappers while retaining the latest planner path. + + The pre-pass performs the allocation-bearing ``in_capture=True`` call. The + inner upstream capture then selects the saved metadata object and invokes + the same planner with ``in_capture=False``. This avoids a second wrapper + allocation while retaining FlashInfer planning and the SWA output-location + binding installed by the pre-pass. + """ + + by_backend: dict[int, tuple[Any, dict[int, Any]]] = {} + for attn_backend, size, metadata in prepared: + entry = by_backend.setdefault(id(attn_backend), (attn_backend, {})) + entry[1][size] = metadata + + restorations = [] + for attn_backend, metadata_by_size in by_backend.values(): + original = attn_backend.init_forward_metadata_out_graph + had_instance_override = "init_forward_metadata_out_graph" in vars(attn_backend) + previous_override = vars(attn_backend).get("init_forward_metadata_out_graph") + + @functools.wraps(original) + def reuse( + forward_batch, + in_capture=False, + *, + _attn_backend=attn_backend, + _metadata_by_size=metadata_by_size, + _original=original, + ): + if in_capture: + _attn_backend.forward_metadata = _metadata_by_size[forward_batch.batch_size] + return _original(forward_batch, in_capture=False) + return _original(forward_batch, in_capture=in_capture) + + attn_backend.init_forward_metadata_out_graph = reuse + restorations.append((attn_backend, had_instance_override, previous_override)) + + try: + yield + finally: + for attn_backend, had_instance_override, previous_override in restorations: + if had_instance_override: + attn_backend.init_forward_metadata_out_graph = previous_override + else: + del attn_backend.init_forward_metadata_out_graph + + +def _save_capture(runner: DecodeCudaGraphRunner, original: Callable[..., Any]): + if _uses_deepep(runner): + start = time.perf_counter() + with _deepep_eager_warmup(): + original(runner) + logger.info( + "[Foundry] SGLang DeepEP eager lazy-init warmup completed in %.3fs", + time.perf_counter() - start, + ) + graph_ops.bootstrap_deepep_buffer(runner) + + if _is_flashinfer_backend(runner.attn_backend): + prepared = _prepare_attention_metadata(runner) + with _reuse_flashinfer_metadata(prepared): + result = original(runner) + else: + result = original(runner) + + graph_ops.save_graph_manifest() + graph_ops.pack_fatbins() + graph_ops.capture_final_alloc_offset() + return result + + +def _load_capture(runner: DecodeCudaGraphRunner) -> None: + if getattr(runner, "_foundry_graphs_loaded", False): + raise RuntimeError( + "Foundry LOAD cannot satisfy an SGLang CUDA-graph recapture request; " + "restart with a fresh compatible archive." + ) + + backend = _ensure_full_backend(runner) + pool = backend._pool + if pool is None: + pool = runner.device_module.graph_pool_handle() + register_graph_pool(pool) + backend._pool = pool + + if _uses_deepep(runner): + graph_ops.bootstrap_deepep_buffer(runner) + + graph_ops.preallocate_for_load_mode() + _prepare_attention_metadata(runner) + loaded = [ + (ShapeKey(size=size, stream_idx=None, variant_label=None), graph, output) + for size, graph, output in graph_ops.load_all_graphs(runner) + ] + hydrate_backend(backend, loaded, pool=pool) + runner.deepep_adapter.capture(is_extend_in_batch=False) + runner._foundry_graphs_loaded = True + + +def _capture_wrapper(original: Callable[..., Any]): + @functools.wraps(original) + def patched(self): + mode = _mode() + if mode == CUDAGraphExtensionMode.NONE: + return original(self) + + _ensure_full_backend(self) + if mode == CUDAGraphExtensionMode.SAVE: + return _save_capture(self, original) + if mode == CUDAGraphExtensionMode.LOAD: + return _load_capture(self) + raise RuntimeError(f"Unsupported Foundry SGLang mode: {mode!r}") + + return patched + + +def _capture_one_wrapper(original: Callable[..., Any]): + @functools.wraps(original) + def patched( + self, + shape_key, + forward_fn, + dummies=None, + post_warmup_hook=None, + ): + mode = _mode() + if mode == CUDAGraphExtensionMode.NONE: + return original( + self, + shape_key, + forward_fn, + dummies=dummies, + post_warmup_hook=post_warmup_hook, + ) + if not isinstance(self, FullCudaGraphBackend): + raise RuntimeError("Foundry graph capture requires SGLang's FullCudaGraphBackend.") + if _EAGER_DEEPEP_WARMUP: + output = forward_fn() + if post_warmup_hook is not None: + post_warmup_hook() + return output + if mode == CUDAGraphExtensionMode.LOAD: + raise RuntimeError( + "Foundry LOAD attempted an upstream graph capture; the archive " + "is incompatible with this SGLang configuration." + ) + if not isinstance(shape_key, ShapeKey): + raise TypeError("SGLang split-runner capture must supply a complete ShapeKey.") + + graph = graph_ops.create_device_graph() + output = graph_ops.capture_graph( + graph, + self._pool, + self._capture_stream, + forward_fn, + ) + graph_ops.save_graph(graph, output, shape_key) + self._graphs[shape_key] = graph + self._outputs[shape_key] = output + return None + + return patched + + +def install_cuda_graph_hooks() -> None: + """Install idempotent hooks on SGLang's latest split-runner classes.""" + + _patch_once(BaseRunner, "warmup", _warmup_wrapper) + _patch_once(DecodeCudaGraphRunner, "capture", _capture_wrapper) + _patch_once(FullCudaGraphBackend, "capture_one", _capture_one_wrapper) diff --git a/python/foundry/integration/sglang/graph_ops.py b/python/foundry/integration/sglang/graph_ops.py index 18ab1f0e..b260337a 100644 --- a/python/foundry/integration/sglang/graph_ops.py +++ b/python/foundry/integration/sglang/graph_ops.py @@ -4,6 +4,7 @@ from __future__ import annotations +import importlib import logging import os import re @@ -21,7 +22,8 @@ get_config, get_graph_extension_mode, ) -from foundry.integration.sglang.runtime import get_state + +rt = importlib.import_module("foundry.integration.sglang.runtime") logger = logging.getLogger(__name__) @@ -29,24 +31,25 @@ _GRAPH_FILENAME_RE = re.compile(r"^graph_(?P\d+)_FULL_t(?P\d+)_r\d+_UX_pcN\.json$") -def _batch_size_from_key(key: Any) -> int: - if isinstance(key, int): - return key - key_str = str(key) - for part in reversed(key_str.split("_")): - if part.isdigit(): - return int(part) - raise ValueError(f"Cannot derive batch size from SGLang CUDA graph key: {key!r}") +def _shape_size(shape_key: Any) -> int: + if not all(hasattr(shape_key, field) for field in ("size", "stream_idx", "variant_label")): + raise TypeError( + "SGLang graph capture requires a complete ShapeKey with size, " + "stream_idx, and variant_label." + ) + if not isinstance(shape_key.size, int): + raise TypeError(f"SGLang ShapeKey.size must be an int, got {shape_key.size!r}") + return shape_key.size -def _graph_filename(index: int, key: Any) -> str: - batch_size = _batch_size_from_key(key) +def _graph_filename(index: int, shape_key: Any) -> str: + batch_size = _shape_size(shape_key) return f"graph_{index}_FULL_t{batch_size}_r{batch_size}_UX_pcN.json" def _pack_output(output: Any) -> torch.Tensor: - from sglang.srt.layers.logits_processor import LogitsProcessorOutput - + logits_processor = importlib.import_module("sglang.srt.layers.logits_processor") + LogitsProcessorOutput = logits_processor.LogitsProcessorOutput if isinstance(output, LogitsProcessorOutput): if output.next_token_logits is None: raise TypeError("SGLang decode CUDA graph output has no next_token_logits") @@ -59,8 +62,8 @@ def _pack_output(output: Any) -> torch.Tensor: def _unpack_output(tensors: Any) -> Any: - from sglang.srt.layers.logits_processor import LogitsProcessorOutput - + logits_processor = importlib.import_module("sglang.srt.layers.logits_processor") + LogitsProcessorOutput = logits_processor.LogitsProcessorOutput if isinstance(tensors, (tuple, list)): if len(tensors) != 1: raise RuntimeError(f"Expected one SGLang CUDA graph output tensor, got {len(tensors)}") @@ -76,7 +79,7 @@ def _scan_graph_files(workspace_dir: str) -> list[tuple[int, str, dict[str, Any] continue meta = { "index": int(match.group("index")), - "key": int(match.group("bs")), + "size": int(match.group("bs")), } graph_files.append((int(meta["index"]), filename, meta)) graph_files.sort(key=lambda x: x[0]) @@ -98,19 +101,23 @@ def capture_graph(graph, pool, stream, run_once_fn): return None -def save_graph(graph, output: Any, key: Any) -> None: +def save_graph(graph, output: Any, shape_key: Any) -> None: cfg = get_config() - state = get_state() + state = rt.get_state() if cfg is None or state is None or cfg.workspace_dir is None: raise RuntimeError("Foundry SGLang graph extension is not initialized") packed_output = _pack_output(output) - filename = _graph_filename(state.capture_index, key) + filename = _graph_filename(state.capture_index, shape_key) graph_path = os.path.join(cfg.workspace_dir, filename) graph.save(graph_path, packed_output) state.capture_index += 1 - logger.info("[Foundry] Saved SGLang CUDA graph %s key=%s", filename, key) + logger.info( + "[Foundry] Saved SGLang CUDA graph %s shape_key=%s", + filename, + shape_key, + ) def save_graph_manifest() -> None: @@ -149,10 +156,10 @@ def start_graph_builds() -> None: ) -def preload_all_graphs() -> None: +def preload_all_graphs() -> list[tuple[int, Any, Any]]: global _pending_graph_builds cfg = get_config() - state = get_state() + state = rt.get_state() if cfg is None or state is None or cfg.workspace_dir is None: raise RuntimeError("Foundry SGLang graph extension is not initialized") @@ -173,9 +180,11 @@ def preload_all_graphs() -> None: time.perf_counter() - t0, ) + loaded = [] for i, (_index, _filename, meta) in enumerate(graph_files): graph, tensors = results[i] - state.loaded_graphs[meta["key"]] = (graph, _unpack_output(tensors)) + loaded.append((meta["size"], graph, _unpack_output(tensors))) + return loaded def bootstrap_deepep_buffer(cuda_graph_runner) -> bool: @@ -197,18 +206,12 @@ def bootstrap_deepep_buffer(cuda_graph_runner) -> bool: Returns True if a buffer was (or already is) created, False if DeepEP is off. """ - try: - from sglang.srt.layers.moe.utils import get_moe_a2a_backend - - if not get_moe_a2a_backend().is_deepep(): - return False - except Exception: + moe_utils = importlib.import_module("sglang.srt.layers.moe.utils") + if not moe_utils.get_moe_a2a_backend().is_deepep(): return False - - from sglang.srt.layers.moe.token_dispatcher.deepep import ( - DeepEPBuffer, - DeepEPDispatcher, - ) + deepep = importlib.import_module("sglang.srt.layers.moe.token_dispatcher.deepep") + DeepEPBuffer = deepep.DeepEPBuffer + DeepEPDispatcher = deepep.DeepEPDispatcher if DeepEPBuffer._buffer is not None: return True @@ -248,48 +251,11 @@ def bootstrap_deepep_buffer(cuda_graph_runner) -> bool: return False -def initialize_attention_metadata_for_bs(cuda_graph_runner, bs: int) -> None: - """Populate ``decode_cuda_graph_metadata[bs]`` for runtime replay. - - The FlashInfer wrappers and their internal ``_int_workspace_buffer`` - are constructed here, outside the captured graph. The graph's - runtime kernels reference these buffer addresses, so LOAD must - re-run the same call before runtime replay so the wrappers exist - at deterministic VMM addresses. - """ - 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( - bs, - num_tokens, - buffers.req_pool_indices[:bs], - buffers.seq_lens[:bs], - encoder_lens, - cuda_graph_runner.capture_forward_mode, - spec_info, - ) - - -def initialize_all_attention_metadata(cuda_graph_runner) -> None: - """Pre-pass: populate ``decode_cuda_graph_metadata`` for all bs at once. - - Called on both SAVE and LOAD before the capture/load loop. Walking - ``reversed(self.capture_bs)`` (largest first) matches SAVE's natural - capture order; same order on both sides keeps the VMM cursor - trajectory identical. - """ - for bs in reversed(cuda_graph_runner.capture_bs): - initialize_attention_metadata_for_bs(cuda_graph_runner, bs) - - -def load_all_graphs(cuda_graph_runner) -> None: +def load_all_graphs(_cuda_graph_runner) -> list[tuple[int, Any, Any]]: """LOAD-time replacement for the upstream capture loop. - All FlashInfer wrappers are pre-allocated by - ``initialize_all_attention_metadata`` (called by the capture hook - before this function), so the VMM cursor sits where SAVE recorded + The capture hook recreates per-shape attention metadata before this + function, so the VMM cursor sits where SAVE recorded ``start_base_addr_0``. Load every graph in one ``start_graph_builds`` call — this is what enables template + on-demand linking in the manifest. ``finish_graph_loads`` then @@ -297,7 +263,7 @@ def load_all_graphs(cuda_graph_runner) -> None: cursor exactly the way SAVE did inside its capture loop. """ cfg = get_config() - state = get_state() + state = rt.get_state() if cfg is None or state is None or cfg.workspace_dir is None: raise RuntimeError("Foundry SGLang graph extension is not initialized") @@ -320,6 +286,20 @@ def load_all_graphs(cuda_graph_runner) -> None: time.perf_counter() - t0, ) + loaded = [] for i, (_index, _filename, meta) in enumerate(graph_files): graph, tensors = results[i] - state.loaded_graphs[meta["key"]] = (graph, _unpack_output(tensors)) + loaded.append((meta["size"], graph, _unpack_output(tensors))) + return loaded + + +def capture_final_alloc_offset() -> int: + """Persist the final deterministic VMM cursor through runtime state.""" + + return rt.capture_final_alloc_offset() + + +def preallocate_for_load_mode() -> None: + """Preallocate the saved deterministic VMM span before graph loading.""" + + rt.preallocate_for_load_mode() diff --git a/python/foundry/integration/sglang/hooks.py b/python/foundry/integration/sglang/hooks.py index 22cd7cfa..3217b416 100644 --- a/python/foundry/integration/sglang/hooks.py +++ b/python/foundry/integration/sglang/hooks.py @@ -5,11 +5,14 @@ from __future__ import annotations import functools +import importlib import logging import os import time from dataclasses import asdict +import torch + from foundry.integration.sglang import runtime as rt from foundry.integration.sglang.config import ( CUDAGraphExtensionMode, @@ -22,17 +25,6 @@ _INSTALLED = False -def _ep_lazy_init_needed() -> bool: - """True when the DeepEP all-to-all backend is active, so pre-capture lazy - init (NVSHMEM buffer, DeepGEMM JIT) must be warmed up outside stream capture.""" - try: - from sglang.srt.layers.moe.utils import get_moe_a2a_backend - - return get_moe_a2a_backend().is_deepep() - except Exception: - return False - - def _resolve_dp_rank(model_runner) -> int | None: dp_rank = getattr(model_runner, "dp_rank", None) if dp_rank is not None: @@ -40,12 +32,11 @@ def _resolve_dp_rank(model_runner) -> int | None: server_args = model_runner.server_args if getattr(server_args, "enable_dp_attention", False): - from sglang.srt.layers.dp_attention import compute_dp_attention_world_info - + dp_attention = importlib.import_module("sglang.srt.layers.dp_attention") # Returns (attn_tp_rank, attn_tp_size, attn_dp_rank, attn_dp_size) — a # 4-tuple. Take attn_dp_rank (3rd) and discard the rest. (Unpacking into # 3 targets would raise ValueError.) - _, _, dp_rank, _ = compute_dp_attention_world_info( + _, _, dp_rank, _ = dp_attention.compute_dp_attention_world_info( server_args.enable_dp_attention, model_runner.tp_rank, server_args.tp_size, @@ -80,6 +71,10 @@ def install_hooks(server_args) -> None: ) load_graph_extension_config(cfg_path) + # Optional-dependency boundary: cuda_graph imports SGLang's private runner + # APIs at module scope and is therefore loaded only during hook installation. + cuda_graph = importlib.import_module("foundry.integration.sglang.cuda_graph") + cuda_graph.validate_configuration(server_args) logger.info( "[Foundry] SGLang hooks installing: mode=%s workspace=%s", get_graph_extension_mode().value, @@ -89,8 +84,7 @@ def install_hooks(server_args) -> None: _patch_init_torch_distributed() _patch_init_memory_pool() _patch_load_model() - _patch_kernel_warmup() - _patch_cuda_graph_capture() + cuda_graph.install_cuda_graph_hooks() _patch_spawn_sites() _INSTALLED = True @@ -98,8 +92,7 @@ def install_hooks(server_args) -> None: def _patch_init_torch_distributed() -> None: - from sglang.srt.model_executor import model_runner as mr - + mr = importlib.import_module("sglang.srt.model_executor.model_runner") cls = mr.ModelRunner orig = cls.init_torch_distributed @@ -119,8 +112,6 @@ def patched(self, *args, **kwargs): # (surfacing at the first Stream()/kernel). Mirrors model_runner's own # set_device(self.gpu_id). Single-GPU is unaffected (gpu_id == 0). if self.device == "cuda": - import torch - torch.get_device_module(self.device).set_device(self.gpu_id) rt.setup_graph_extension( @@ -140,9 +131,8 @@ def patched(self, *args, **kwargs): def _patch_init_memory_pool() -> None: - from sglang.srt.model_executor import model_runner_kv_cache_mixin as kv_mixin - from sglang.srt.model_executor.pool_configurator import MemoryPoolConfig - + kv_mixin = importlib.import_module("sglang.srt.model_executor.model_runner_kv_cache_mixin") + pool_configurator = importlib.import_module("sglang.srt.model_executor.pool_configurator") cls = kv_mixin.ModelRunnerKVCacheMixin orig = cls.init_memory_pool @@ -153,13 +143,11 @@ def patched(self, pre_model_load_memory): return orig(self, pre_model_load_memory) if mode == CUDAGraphExtensionMode.LOAD: - import torch - rt.log_alloc_offset("before_init_memory_pool") state = rt.load_warmup_state() if not state.memory_pool_config: raise RuntimeError("Foundry LOAD requires memory_pool_config") - self.memory_pool_config = MemoryPoolConfig(**state.memory_pool_config) + self.memory_pool_config = pool_configurator.MemoryPoolConfig(**state.memory_pool_config) # Mirror SAVE's ``_resolve_memory_pool_config`` -> # ``get_available_gpu_memory(empty_cache=True)`` side # effect. Without this, torch's caching allocator retains @@ -184,8 +172,7 @@ def patched(self, pre_model_load_memory): def _patch_load_model() -> None: - from sglang.srt.model_executor import model_runner as mr - + mr = importlib.import_module("sglang.srt.model_executor.model_runner") cls = mr.ModelRunner orig = cls.load_model @@ -196,389 +183,26 @@ def patched(self, *args, **kwargs): cls.load_model = patched -def _patch_kernel_warmup() -> None: - from sglang.srt.model_executor import model_runner as mr - - cls = mr.ModelRunner - orig = cls.kernel_warmup - - @functools.wraps(orig) - def patched(self, *args, **kwargs): - mode = get_graph_extension_mode() - if mode == CUDAGraphExtensionMode.NONE: - return orig(self, *args, **kwargs) - # Phase 1 keeps pre-graph model-forward warmups out of SAVE/LOAD. - logger.info("[Foundry] SGLang kernel_warmup skipped in %s mode", mode.value) - return None - - cls.kernel_warmup = patched - - -def _patch_cuda_graph_capture() -> None: - from sglang.srt.model_executor import cuda_graph_runner as cgr - - cls = cgr.CudaGraphRunner - orig_capture = cls.capture - orig_create_device_graph = cls._create_device_graph - orig_capture_graph = cls._capture_graph - orig_capture_one_batch_size = cls.capture_one_batch_size - - # When set, the capture machinery is being reused as a foundry-driven WARMUP - # pass: run a real forward per bs (no graph capture, no save) to trigger all - # of sglang's pre-capture lazy init — DeepEP buffer, DeepGEMM per-shape JIT, - # etc. — that would otherwise fire inside the captured stream and abort with - # "operation not permitted when stream is capturing". See `_run_warmup_pass`. - warmup_active = [False] - - @functools.wraps(orig_create_device_graph) - def patched_create_device_graph(self, *args, **kwargs): - mode = get_graph_extension_mode() - if warmup_active[0]: - # Throwaway graph object; the warmup never captures into it. - return orig_create_device_graph(self, *args, **kwargs) - if mode == CUDAGraphExtensionMode.SAVE: - from foundry.integration.sglang.graph_ops import create_device_graph - - return create_device_graph() - return orig_create_device_graph(self, *args, **kwargs) - - @functools.wraps(orig_capture_graph) - def patched_capture_graph(self, graph, pool, stream, run_once_fn): - mode = get_graph_extension_mode() - if warmup_active[0]: - # Warmup: run the forward eagerly (NO torch.cuda.graph), so DeepGEMM - # JIT compile / NVSHMEM buffer creation happen outside stream capture. - return run_once_fn() - if mode == CUDAGraphExtensionMode.SAVE: - from foundry.integration.sglang.graph_ops import capture_graph - - return capture_graph(graph, pool, stream, run_once_fn) - return orig_capture_graph(self, graph, pool, stream, run_once_fn) - - @functools.wraps(orig_capture_one_batch_size) - def patched_capture_one_batch_size(self, bs, forward, stream_idx=None): - mode = get_graph_extension_mode() - if warmup_active[0]: - # Warmup pass: run upstream unchanged (real warmup forwards + a - # neutered "capture" that is just another real forward). No - # suppression, no save_graph. - return orig_capture_one_batch_size(self, bs, forward, stream_idx) - if mode == CUDAGraphExtensionMode.SAVE: - # Suppress the two pre-capture warmup forwards - # (cuda_graph_runner.py: ``for _ in range(2): run_once()``). - # Their non-deterministic activation allocations would pollute - # the torch caching allocator with freed segments that LOAD - # cannot reproduce — causing the per-bs init's cache-miss vs - # cache-hit asymmetry that drifts the VMM cursor away from - # each saved ``start_base_addr``. JIT / autotune still happens - # inside the graph capture (3rd run_once invocation) and is - # recorded as alloc events, mirroring vLLM doc 04 §2. - counter = [0] - real_forward = forward - - def warmup_skipping_forward(*args, **kwargs): - counter[0] += 1 - if counter[0] <= 2: - return None - return real_forward(*args, **kwargs) - - forward = warmup_skipping_forward - graph, output = orig_capture_one_batch_size(self, bs, forward, stream_idx) - if mode == CUDAGraphExtensionMode.SAVE: - from foundry.integration.sglang.graph_ops import save_graph - - # Mirror the inline key shape upstream uses for self.graphs[key] - # in `_capture_one_stream`. ``_make_graph_key`` and - # ``get_capture_lora_variant`` were removed in sglang - # commit ce2506e1c (record_nolora_graph deprecation). - key = bs if stream_idx is None else f"{stream_idx}_{bs}" - save_graph(graph, output, key) - return graph, output - - def _run_warmup_pass(self): - """Foundry-driven pre-capture warmup for the DeepEP/EP path. - - Reuses the upstream capture loop with graph capture neutered (see - ``warmup_active``) to run one real forward per ``capture_bs`` BEFORE the - real capture/replay work — triggering every pre-capture lazy init sglang - normally does in its (foundry-suppressed) warmup forwards: DeepEP buffer - creation, DeepGEMM per-shape JIT compile, etc. Run on BOTH SAVE and LOAD - so the VMM cursor trajectory stays symmetric. ``self.graphs`` / - ``self.output_buffers`` are saved/cleared/restored so the throwaway - warmup entries don't leak into the real pass. - """ - warmup_active[0] = True - saved_graphs, saved_buffers = self.graphs, self.output_buffers - self.graphs, self.output_buffers = {}, {} - t0 = time.perf_counter() - try: - orig_capture(self) - finally: - warmup_active[0] = False - self.graphs, self.output_buffers = saved_graphs, saved_buffers - logger.info( - "[Foundry] SGLang EP warmup pass (lazy-init) completed in %.3fs", - time.perf_counter() - t0, - ) - - @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): - # 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) - - if mode == CUDAGraphExtensionMode.LOAD: - from sglang.srt.distributed.device_communicators.pynccl_allocator import ( - set_graph_pool_id, - ) - - from foundry.integration.sglang.graph_ops import ( - initialize_all_attention_metadata, - load_all_graphs, - ) - - state = rt.get_state() - if state is None: - raise RuntimeError("Foundry SGLang state is not initialized") - # Set up the graph memory pool once — sglang shares one pool - # across all captured graphs, and runtime replay also requires - # it to be set so pynccl knows which pool the graph belongs to. - 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()) - rt.log_alloc_offset("before_preallocate") - rt.preallocate_for_load_mode() - rt.log_alloc_offset("after_preallocate") - # 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`` - # is idempotent so the inner per-iter call inside the SAVE - # capture loop does not re-allocate. Same upfront allocation - # sequence on both sides → cursor sits at SAVE's - # ``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"): - initialize_all_attention_metadata(self) - rt.log_alloc_offset("after_pre_init") - # Single ``start_graph_builds(all_paths)`` call so templates - # and on-demand graphs link via ``shared_exec`` in the - # manifest. ``finish_graph_loads`` replays alloc events - # graph-by-graph in the same order SAVE captured them. - load_all_graphs(self) - rt.log_alloc_offset("after_load_all_graphs") - self.graphs = {k: v[0] for k, v in state.loaded_graphs.items()} - self.output_buffers = {k: v[1] for k, v in state.loaded_graphs.items()} - # Non-FlashInfer backends (e.g. fa3) populate per-bs decode metadata - # — looked up at replay as attn_backend.decode_cuda_graph_metadata[bs] - # — inside the capture loop, which LOAD replaces. The FlashInfer - # pre-pass above (gated) handled flashinfer; for the rest, populate it - # 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"): - 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, - # which LOAD replaces — so without this, runtime replay() asserts - # `_captured_deepep_mode is not None` on the first decode. capture() - # self-gates on the DeepEP backend (no-op otherwise) and sets the - # decode dispatch mode the captured graphs expect. - self.deepep_adapter.capture(is_extend_in_batch=False) - return None - - if mode == CUDAGraphExtensionMode.SAVE: - # 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. - attn_backend = self.attn_backend - use_fi_prepass = hasattr(attn_backend, "indices_updater_decode") - real_init = attn_backend.init_forward_metadata_capture_cuda_graph - - def reuse_pre_pass_init( - bs, - num_tokens, - req_pool_indices, - seq_lens, - encoder_lens, - forward_mode, - spec_info, - ): - # The pre-pass already allocated a wrapper for this - # bs and stored it in - # ``decode_cuda_graph_metadata`` / - # ``prefill_cuda_graph_metadata``. Reuse it directly - # — no second torch.empty for ``_int_workspace_buffer``. - # Re-run the planner with the same buffer slices the - # capture forward uses, then point - # ``forward_metadata`` at the same wrappers. Same - # plan call on LOAD via the symmetric pre-pass, so - # the captured graph kernels reference VMM addresses - # that LOAD's wrappers actually occupy. - from sglang.srt.layers.attention.flashinfer_backend import ( - DecodeMetadata, - PrefillMetadata, - ) - - if forward_mode.is_decode_or_idle(): - wrappers = attn_backend.decode_cuda_graph_metadata.get(bs) - if wrappers is None: - return real_init( - bs, - num_tokens, - req_pool_indices, - seq_lens, - encoder_lens, - forward_mode, - spec_info, - ) - seq_lens_sum = seq_lens.sum().item() - attn_backend.indices_updater_decode.update( - req_pool_indices, - seq_lens, - seq_lens.cpu(), - seq_lens_sum, - decode_wrappers=wrappers, - encoder_lens=encoder_lens, - spec_info=spec_info, - fixed_split_size=None, - disable_split_kv=attn_backend.disable_cuda_graph_kv_split, - ) - attn_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) - if wrappers is None: - return real_init( - bs, - num_tokens, - req_pool_indices, - seq_lens, - encoder_lens, - forward_mode, - spec_info, - ) - seq_lens_sum = seq_lens.sum().item() - use_ragged = forward_mode.is_dllm_extend() - prefix_lens = ( - seq_lens - attn_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( - req_pool_indices, - seq_lens, - seq_lens.cpu(), - seq_lens_sum, - prefix_lens=prefix_lens, - prefill_wrappers=wrappers, - use_ragged=use_ragged, - encoder_lens=encoder_lens, - spec_info=spec_info_arg, - ) - attn_backend.forward_metadata = PrefillMetadata(wrappers, use_ragged, False) - return - # Unknown mode — fall back to real init. - return real_init( - bs, - num_tokens, - req_pool_indices, - seq_lens, - encoder_lens, - forward_mode, - spec_info, - ) - - if use_fi_prepass: - from foundry.integration.sglang.graph_ops import ( - initialize_all_attention_metadata, - ) - - rt.log_alloc_offset("save_before_pre_init") - # Pre-pass: allocate every per-bs FlashInfer wrapper up front, in - # the same ``reversed(capture_bs)`` order LOAD uses. - 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 - try: - result = orig_capture(self, *args, **kwargs) - finally: - attn_backend.init_forward_metadata_capture_cuda_graph = real_init - - from foundry.integration.sglang.graph_ops import ( - pack_fatbins, - save_graph_manifest, - ) - - save_graph_manifest() - pack_fatbins() - rt.capture_final_alloc_offset() - return result - - return orig_capture(self, *args, **kwargs) - - cls._create_device_graph = patched_create_device_graph - cls._capture_graph = patched_capture_graph - cls.capture_one_batch_size = patched_capture_one_batch_size - cls.capture = patched - - def _patch_spawn_sites() -> None: try: - from sglang.srt.entrypoints import engine as engine_mod + engine_mod = importlib.import_module("sglang.srt.entrypoints.engine") except Exception: engine_mod = None if engine_mod is not None: - orig_launch = engine_mod.Engine._launch_scheduler_processes + engine_cls = engine_mod.Engine + orig_launch = engine_cls.__dict__["_launch_scheduler_processes"].__func__ @functools.wraps(orig_launch) - def patched_launch(self, *args, **kwargs): + def patched_launch(cls, *args, **kwargs): if get_graph_extension_mode() != CUDAGraphExtensionMode.NONE: rt.setup_ld_preload_env() - return orig_launch(self, *args, **kwargs) + return orig_launch(cls, *args, **kwargs) - engine_mod.Engine._launch_scheduler_processes = patched_launch + engine_cls._launch_scheduler_processes = classmethod(patched_launch) try: - from sglang.srt.managers import data_parallel_controller as dpc + dpc = importlib.import_module("sglang.srt.managers.data_parallel_controller") except Exception: dpc = None From 78c399d73f5e6e8952afb57be6defe56ac212d8e Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Thu, 23 Jul 2026 07:07:02 +0000 Subject: [PATCH 04/21] feat(sglang): version graph archives Co-authored-by: Rahul Chalamala --- .../foundry/integration/sglang/cuda_graph.py | 121 ++++- .../foundry/integration/sglang/graph_ops.py | 425 ++++++++++-------- python/foundry/integration/sglang/hooks.py | 6 +- python/foundry/integration/sglang/runtime.py | 105 ++++- tests/test_sglang_runtime.py | 185 ++++++++ 5 files changed, 650 insertions(+), 192 deletions(-) create mode 100644 tests/test_sglang_runtime.py diff --git a/python/foundry/integration/sglang/cuda_graph.py b/python/foundry/integration/sglang/cuda_graph.py index 7267d5b6..bb92659b 100644 --- a/python/foundry/integration/sglang/cuda_graph.py +++ b/python/foundry/integration/sglang/cuda_graph.py @@ -18,6 +18,7 @@ import time from collections.abc import Callable, Iterable, Iterator from contextlib import contextmanager +from dataclasses import fields from typing import Any from sglang.srt.model_executor.runner.base_runner import BaseRunner @@ -203,6 +204,90 @@ def _uses_deepep(runner: DecodeCudaGraphRunner) -> bool: return _value(getattr(runner.model_runner.server_args, "moe_a2a_backend", "none")) == "deepep" +def _pack_output(output: Any) -> tuple[Any, str]: + logits_processor = _module("sglang.srt.layers.logits_processor") + logits_output_type = logits_processor.LogitsProcessorOutput + if not isinstance(output, logits_output_type): + raise TypeError( + "Foundry supports only SGLang LogitsProcessorOutput CUDA graph outputs; " + f"got {type(output).__name__}." + ) + if output.next_token_logits is None: + raise TypeError("SGLang decode CUDA graph output has no next_token_logits.") + + populated_unsupported_fields = [ + field.name + for field in fields(output) + if field.name != "next_token_logits" and getattr(output, field.name) is not None + ] + if populated_unsupported_fields: + raise TypeError( + "Foundry supports only next_token_logits in SGLang decode CUDA graph output; " + f"unsupported populated fields: {populated_unsupported_fields}." + ) + return output.next_token_logits, "next_token_logits" + + +def _unpack_output(tensors: Any, output_kind: str) -> Any: + if output_kind != "next_token_logits": + raise RuntimeError( + f"Unsupported SGLang CUDA graph output kind during LOAD: {output_kind!r}." + ) + if isinstance(tensors, (tuple, list)): + if len(tensors) != 1: + raise RuntimeError(f"Expected one SGLang CUDA graph output tensor, got {len(tensors)}.") + tensors = tensors[0] + logits_processor = _module("sglang.srt.layers.logits_processor") + return logits_processor.LogitsProcessorOutput(next_token_logits=tensors) + + +def bootstrap_deepep_buffer(cuda_graph_runner: DecodeCudaGraphRunner) -> bool: + """Initialize DeepEP's singleton buffer before CUDA graph capture.""" + + moe_utils = _module("sglang.srt.layers.moe.utils") + if not moe_utils.get_moe_a2a_backend().is_deepep(): + return False + deepep_module = _module("sglang.srt.layers.moe.token_dispatcher.deepep") + deep_ep_buffer = deepep_module.DeepEPBuffer + deep_ep_dispatcher = deepep_module.DeepEPDispatcher + + if deep_ep_buffer._buffer is not None: + return True + + model = cuda_graph_runner.model_runner.model + for module in model.modules(): + dispatcher = getattr(module, "dispatcher", None) + if dispatcher is None: + continue + candidates = [dispatcher, *getattr(dispatcher, "_inners", [])] + deepep = next( + (item for item in candidates if isinstance(item, deep_ep_dispatcher)), + None, + ) + if deepep is None: + continue + implementation = getattr(deepep, "_low_latency_dispatcher", None) or getattr( + deepep, + "_normal_dispatcher", + None, + ) + if implementation is None: + continue + start = time.perf_counter() + implementation._get_buffer() + logger.info( + "[Foundry] Bootstrapped DeepEP buffer pre-capture in %.3fs", + time.perf_counter() - start, + ) + return True + + logger.warning( + "[Foundry] DeepEP backend active but no DeepEPDispatcher found on the " + "model; buffer not bootstrapped (capture may fail inside stream capture)." + ) + return False + + @contextmanager def _deepep_eager_warmup() -> Iterator[None]: global _EAGER_DEEPEP_WARMUP @@ -296,7 +381,7 @@ def _save_capture(runner: DecodeCudaGraphRunner, original: Callable[..., Any]): "[Foundry] SGLang DeepEP eager lazy-init warmup completed in %.3fs", time.perf_counter() - start, ) - graph_ops.bootstrap_deepep_buffer(runner) + bootstrap_deepep_buffer(runner) if _is_flashinfer_backend(runner.attn_backend): prepared = _prepare_attention_metadata(runner) @@ -326,14 +411,25 @@ def _load_capture(runner: DecodeCudaGraphRunner) -> None: backend._pool = pool if _uses_deepep(runner): - graph_ops.bootstrap_deepep_buffer(runner) + bootstrap_deepep_buffer(runner) graph_ops.preallocate_for_load_mode() _prepare_attention_metadata(runner) - loaded = [ - (ShapeKey(size=size, stream_idx=None, variant_label=None), graph, output) - for size, graph, output in graph_ops.load_all_graphs(runner) - ] + loaded = [] + for entry, graph, tensors in graph_ops.load_all_graphs(runner): + archive_key = entry.shape_key + shape_key = ShapeKey( + size=archive_key.size, + stream_idx=archive_key.stream_idx, + variant_label=archive_key.variant_label, + ) + loaded.append( + ( + shape_key, + graph, + _unpack_output(tensors, entry.output_kind), + ) + ) hydrate_backend(backend, loaded, pool=pool) runner.deepep_adapter.capture(is_extend_in_batch=False) runner._foundry_graphs_loaded = True @@ -396,7 +492,18 @@ def patched( self._capture_stream, forward_fn, ) - graph_ops.save_graph(graph, output, shape_key) + packed_output, output_kind = _pack_output(output) + archive_key = graph_ops.ArchiveShapeKey( + size=shape_key.size, + stream_idx=shape_key.stream_idx, + variant_label=shape_key.variant_label, + ) + graph_ops.save_graph( + graph, + packed_output, + archive_key, + output_kind, + ) self._graphs[shape_key] = graph self._outputs[shape_key] = output return None diff --git a/python/foundry/integration/sglang/graph_ops.py b/python/foundry/integration/sglang/graph_ops.py index b260337a..f7a0d4cd 100644 --- a/python/foundry/integration/sglang/graph_ops.py +++ b/python/foundry/integration/sglang/graph_ops.py @@ -2,13 +2,14 @@ # SPDX-FileCopyrightText: Copyright contributors to the Foundry project """Foundry CUDA graph save/load helpers for SGLang.""" -from __future__ import annotations - import importlib +import json import logging import os -import re +import tempfile import time +from dataclasses import asdict, dataclass +from pathlib import Path from typing import Any import torch @@ -27,63 +28,216 @@ logger = logging.getLogger(__name__) -_pending_graph_builds: tuple[Any, list[tuple[int, str, dict[str, Any]]]] | None = None -_GRAPH_FILENAME_RE = re.compile(r"^graph_(?P\d+)_FULL_t(?P\d+)_r\d+_UX_pcN\.json$") +GRAPH_INDEX_SCHEMA_VERSION = 1 +GRAPH_INDEX_FILENAME = "sglang_graph_index.json" +_SUPPORTED_PHASE = "decode" +_SUPPORTED_OUTPUT_KIND = "next_token_logits" -def _shape_size(shape_key: Any) -> int: - if not all(hasattr(shape_key, field) for field in ("size", "stream_idx", "variant_label")): - raise TypeError( - "SGLang graph capture requires a complete ShapeKey with size, " - "stream_idx, and variant_label." - ) - if not isinstance(shape_key.size, int): - raise TypeError(f"SGLang ShapeKey.size must be an int, got {shape_key.size!r}") - return shape_key.size +@dataclass(frozen=True) +class ArchiveShapeKey: + size: int + stream_idx: int | None = None + variant_label: str | None = None -def _graph_filename(index: int, shape_key: Any) -> str: - batch_size = _shape_size(shape_key) - return f"graph_{index}_FULL_t{batch_size}_r{batch_size}_UX_pcN.json" +@dataclass(frozen=True) +class GraphIndexEntry: + phase: str + capture_index: int + filename: str + shape_key: ArchiveShapeKey + output_kind: str -def _pack_output(output: Any) -> torch.Tensor: - logits_processor = importlib.import_module("sglang.srt.layers.logits_processor") - LogitsProcessorOutput = logits_processor.LogitsProcessorOutput - if isinstance(output, LogitsProcessorOutput): - if output.next_token_logits is None: - raise TypeError("SGLang decode CUDA graph output has no next_token_logits") - return output.next_token_logits +def _fresh_save_error(message: str) -> ValueError: + return ValueError(f"{message} Perform a fresh Foundry SAVE with SGLang 0.5.15.post1.") - if isinstance(output, torch.Tensor): - return output - - raise TypeError(f"Unsupported SGLang CUDA graph output type: {type(output)!r}") +def _require_int(value: Any, field_name: str, *, minimum: int = 0) -> int: + if isinstance(value, bool) or not isinstance(value, int) or value < minimum: + raise _fresh_save_error( + f"SGLang graph index field {field_name!r} must be an integer >= {minimum}; " + f"got {value!r}." + ) + return value + + +def _normalize_shape_key(shape_key: Any) -> ArchiveShapeKey: + if isinstance(shape_key, dict): + expected = {"size", "stream_idx", "variant_label"} + if set(shape_key) != expected: + raise _fresh_save_error( + "SGLang graph index shape_key must contain exactly size, stream_idx, " + "and variant_label." + ) + size = shape_key["size"] + stream_idx = shape_key["stream_idx"] + variant_label = shape_key["variant_label"] + elif all( + hasattr(shape_key, field_name) for field_name in ("size", "stream_idx", "variant_label") + ): + size = shape_key.size + stream_idx = shape_key.stream_idx + variant_label = shape_key.variant_label + else: + raise _fresh_save_error( + "SGLang graph index shape_key must be a complete record with size, " + "stream_idx, and variant_label." + ) -def _unpack_output(tensors: Any) -> Any: - logits_processor = importlib.import_module("sglang.srt.layers.logits_processor") - LogitsProcessorOutput = logits_processor.LogitsProcessorOutput - if isinstance(tensors, (tuple, list)): - if len(tensors) != 1: - raise RuntimeError(f"Expected one SGLang CUDA graph output tensor, got {len(tensors)}") - tensors = tensors[0] - return LogitsProcessorOutput(next_token_logits=tensors) + size = _require_int(size, "shape_key.size", minimum=1) + if stream_idx is not None: + stream_idx = _require_int(stream_idx, "shape_key.stream_idx") + if variant_label is not None and not isinstance(variant_label, str): + raise _fresh_save_error( + "SGLang graph index field 'shape_key.variant_label' must be a string or null; " + f"got {variant_label!r}." + ) + return ArchiveShapeKey( + size=size, + stream_idx=stream_idx, + variant_label=variant_label, + ) -def _scan_graph_files(workspace_dir: str) -> list[tuple[int, str, dict[str, Any]]]: - graph_files = [] - for filename in os.listdir(workspace_dir): - match = _GRAPH_FILENAME_RE.match(filename) - if not match: - continue - meta = { - "index": int(match.group("index")), - "size": int(match.group("bs")), +def _normalize_entry(entry: Any) -> GraphIndexEntry: + if isinstance(entry, dict): + expected = { + "phase", + "capture_index", + "filename", + "shape_key", + "output_kind", } - graph_files.append((int(meta["index"]), filename, meta)) - graph_files.sort(key=lambda x: x[0]) - return graph_files + if set(entry) != expected: + raise _fresh_save_error( + "SGLang graph index entries must contain exactly phase, capture_index, " + "filename, shape_key, and output_kind." + ) + values = entry + elif isinstance(entry, GraphIndexEntry): + values = { + "phase": entry.phase, + "capture_index": entry.capture_index, + "filename": entry.filename, + "shape_key": entry.shape_key, + "output_kind": entry.output_kind, + } + else: + raise _fresh_save_error( + f"SGLang graph index entry must be an object; got {type(entry).__name__}." + ) + + phase = values["phase"] + if phase != _SUPPORTED_PHASE: + raise _fresh_save_error( + f"Unsupported SGLang graph phase {phase!r}; only {_SUPPORTED_PHASE!r} is supported." + ) + capture_index = _require_int(values["capture_index"], "capture_index") + filename = values["filename"] + if ( + not isinstance(filename, str) + or not filename + or filename in {".", ".."} + or "/" in filename + or "\\" in filename + or Path(filename).name != filename + ): + raise _fresh_save_error( + f"SGLang graph index filename must be a non-empty basename; got {filename!r}." + ) + output_kind = values["output_kind"] + if output_kind != _SUPPORTED_OUTPUT_KIND: + raise _fresh_save_error( + f"Unsupported SGLang graph output kind {output_kind!r}; " + f"only {_SUPPORTED_OUTPUT_KIND!r} is supported." + ) + return GraphIndexEntry( + phase=phase, + capture_index=capture_index, + filename=filename, + shape_key=_normalize_shape_key(values["shape_key"]), + output_kind=output_kind, + ) + + +def _normalize_entries(entries: Any) -> list[GraphIndexEntry]: + if not isinstance(entries, (tuple, list)): + raise _fresh_save_error("SGLang graph index 'graphs' must be a list.") + normalized = sorted( + (_normalize_entry(entry) for entry in entries), + key=lambda entry: entry.capture_index, + ) + filenames = [entry.filename for entry in normalized] + if len(filenames) != len(set(filenames)): + raise _fresh_save_error("SGLang graph index filenames must be unique.") + for previous, current in zip(normalized, normalized[1:]): + if current.capture_index == previous.capture_index: + raise _fresh_save_error( + f"Duplicate SGLang graph capture_index {current.capture_index}." + ) + if current.capture_index != previous.capture_index + 1: + raise _fresh_save_error( + "SGLang graph capture_index values must be contiguous in capture order." + ) + return normalized + + +def save_graph_index( + path: str | os.PathLike[str], + entries: list[GraphIndexEntry], +) -> None: + index_path = Path(path) + normalized = _normalize_entries(entries) + payload = { + "schema_version": GRAPH_INDEX_SCHEMA_VERSION, + "graphs": [asdict(entry) for entry in normalized], + } + temporary_path = None + try: + with tempfile.NamedTemporaryFile( + "w", + encoding="utf-8", + dir=index_path.parent, + prefix=f".{index_path.name}.", + suffix=".tmp", + delete=False, + ) as temporary: + temporary_path = temporary.name + json.dump(payload, temporary, indent=2, sort_keys=True) + temporary.write("\n") + temporary.flush() + os.fsync(temporary.fileno()) + os.replace(temporary_path, index_path) + finally: + if temporary_path is not None and os.path.exists(temporary_path): + os.unlink(temporary_path) + + +def load_graph_index(path: str | os.PathLike[str]) -> list[GraphIndexEntry]: + index_path = Path(path) + try: + with index_path.open(encoding="utf-8") as index_file: + payload = json.load(index_file) + except (OSError, json.JSONDecodeError) as exc: + raise _fresh_save_error( + f"Cannot read the latest SGLang graph index at {index_path}: {exc}" + ) from exc + + if not isinstance(payload, dict): + raise _fresh_save_error("SGLang graph index root must be an object.") + schema_version = payload.get("schema_version") + if isinstance(schema_version, bool) or schema_version != GRAPH_INDEX_SCHEMA_VERSION: + raise _fresh_save_error( + "Unsupported or legacy SGLang graph index schema " + f"{schema_version!r}; expected {GRAPH_INDEX_SCHEMA_VERSION}." + ) + if set(payload) != {"schema_version", "graphs"}: + raise _fresh_save_error( + "SGLang graph index must contain exactly schema_version and graphs." + ) + return _normalize_entries(payload["graphs"]) def create_device_graph(): @@ -101,17 +255,31 @@ def capture_graph(graph, pool, stream, run_once_fn): return None -def save_graph(graph, output: Any, shape_key: Any) -> None: +def save_graph( + graph: Any, + packed_output: Any, + shape_key: ArchiveShapeKey, + output_kind: str, +) -> None: cfg = get_config() state = rt.get_state() if cfg is None or state is None or cfg.workspace_dir is None: raise RuntimeError("Foundry SGLang graph extension is not initialized") - packed_output = _pack_output(output) - filename = _graph_filename(state.capture_index, shape_key) + entry = _normalize_entry( + GraphIndexEntry( + phase=_SUPPORTED_PHASE, + capture_index=state.capture_index, + filename=f"sglang_graph_{state.capture_index:06d}.json", + shape_key=shape_key, + output_kind=output_kind, + ) + ) + filename = entry.filename graph_path = os.path.join(cfg.workspace_dir, filename) graph.save(graph_path, packed_output) + state.graph_index_entries.append(entry) state.capture_index += 1 logger.info( "[Foundry] Saved SGLang CUDA graph %s shape_key=%s", @@ -124,6 +292,13 @@ def save_graph_manifest() -> None: cfg = get_config() if cfg is None or cfg.workspace_dir is None: return + state = rt.get_state() + if state is None or not state.graph_index_entries: + raise RuntimeError("Foundry SAVE captured no SGLang CUDA graphs") + save_graph_index( + os.path.join(cfg.workspace_dir, GRAPH_INDEX_FILENAME), + state.graph_index_entries, + ) foundry_pkg.save_graph_manifest(cfg.workspace_dir) @@ -135,123 +310,9 @@ def pack_fatbins() -> None: cge.set_pack_fatbins_on_exit(False) -def start_graph_builds() -> None: - global _pending_graph_builds - cfg = get_config() - if cfg is None or cfg.workspace_dir is None or cfg.mode != CUDAGraphExtensionMode.LOAD: - return - - graph_files = _scan_graph_files(cfg.workspace_dir) - if not graph_files: - raise RuntimeError(f"No Foundry SGLang graph files found in {cfg.workspace_dir}") - - paths = [os.path.join(cfg.workspace_dir, filename) for _, filename, _ in graph_files] - t0 = time.perf_counter() - pending = FoundryCUDAGraph.start_graph_builds(paths, num_threads=4) - _pending_graph_builds = (pending, graph_files) - logger.info( - "[Foundry] Started SGLang graph builds for %d graphs in %.3fs", - len(paths), - time.perf_counter() - t0, - ) - - -def preload_all_graphs() -> list[tuple[int, Any, Any]]: - global _pending_graph_builds - cfg = get_config() - state = rt.get_state() - if cfg is None or state is None or cfg.workspace_dir is None: - raise RuntimeError("Foundry SGLang graph extension is not initialized") - - if _pending_graph_builds is None: - start_graph_builds() - assert _pending_graph_builds is not None - - cge.init_nvshmem_for_loaded_modules() - - pending, graph_files = _pending_graph_builds - _pending_graph_builds = None - - t0 = time.perf_counter() - results = FoundryCUDAGraph.finish_graph_loads(pending) - logger.info( - "[Foundry] Finished SGLang graph loads for %d graphs in %.3fs", - len(results), - time.perf_counter() - t0, - ) - - loaded = [] - for i, (_index, _filename, meta) in enumerate(graph_files): - graph, tensors = results[i] - loaded.append((meta["size"], graph, _unpack_output(tensors))) - return loaded - - -def bootstrap_deepep_buffer(cuda_graph_runner) -> bool: - """Force the singleton DeepEP ``Buffer`` (NVSHMEM runtime + symmetric heap) - to be created BEFORE the cuda-graph capture loop. - - sglang creates the DeepEP buffer lazily on the first MoE dispatch — normally - during the two pre-capture warmup forwards. Foundry suppresses those warmups - (for allocation determinism), which would push buffer creation into the - captured forward, where ``deep_ep_cpp.Buffer(...)`` aborts with - ``operation not permitted when stream is capturing``. - - Triggering it here (outside any stream capture) creates only the NVSHMEM - runtime + symmetric heap — no model activations — so it stays symmetric - across SAVE and LOAD and lands at the same VMM offset on both. The buffer is - a process-wide singleton (``DeepEPBuffer._buffer``), so one creation per rank - is enough. It is a collective over the EP group, so every rank must reach - this point together — which they do, since ``capture`` runs on all ranks. - - Returns True if a buffer was (or already is) created, False if DeepEP is off. - """ - moe_utils = importlib.import_module("sglang.srt.layers.moe.utils") - if not moe_utils.get_moe_a2a_backend().is_deepep(): - return False - deepep = importlib.import_module("sglang.srt.layers.moe.token_dispatcher.deepep") - DeepEPBuffer = deepep.DeepEPBuffer - DeepEPDispatcher = deepep.DeepEPDispatcher - - if DeepEPBuffer._buffer is not None: - return True - - model = cuda_graph_runner.model_runner.model - for module in model.modules(): - dispatcher = getattr(module, "dispatcher", None) - if dispatcher is None: - continue - # ``module.dispatcher`` is normally a MaybeTboDeepEPDispatcher wrapper - # whose ``_inners`` hold the real DeepEPDispatcher(s); unwrap it. (Also - # handle a bare DeepEPDispatcher for safety.) - candidates = [dispatcher, *getattr(dispatcher, "_inners", [])] - deepep = next((d for d in candidates if isinstance(d, DeepEPDispatcher)), None) - if deepep is None: - continue - # Prefer the low-latency impl (the mode foundry captures); the buffer - # is sized for whichever impls exist, so either bootstraps the shared - # singleton. - impl = getattr(deepep, "_low_latency_dispatcher", None) or getattr( - deepep, "_normal_dispatcher", None - ) - if impl is None: - continue - t0 = time.perf_counter() - impl._get_buffer() - logger.info( - "[Foundry] Bootstrapped DeepEP buffer pre-capture in %.3fs", - time.perf_counter() - t0, - ) - return True - - logger.warning( - "[Foundry] DeepEP backend active but no DeepEPDispatcher found on the " - "model; buffer not bootstrapped (capture may fail inside stream capture)." - ) - return False - - -def load_all_graphs(_cuda_graph_runner) -> list[tuple[int, Any, Any]]: +def load_all_graphs( + _cuda_graph_runner: Any, +) -> list[tuple[GraphIndexEntry, Any, Any]]: """LOAD-time replacement for the upstream capture loop. The capture hook recreates per-shape attention metadata before this @@ -267,19 +328,33 @@ def load_all_graphs(_cuda_graph_runner) -> list[tuple[int, Any, Any]]: if cfg is None or state is None or cfg.workspace_dir is None: raise RuntimeError("Foundry SGLang graph extension is not initialized") - graph_files = _scan_graph_files(cfg.workspace_dir) - if not graph_files: - raise RuntimeError(f"No Foundry SGLang graph files found in {cfg.workspace_dir}") + entries = load_graph_index(os.path.join(cfg.workspace_dir, GRAPH_INDEX_FILENAME)) + if not entries: + raise _fresh_save_error("The SGLang graph index contains no decode graphs.") + missing_files = [ + entry.filename + for entry in entries + if not os.path.isfile(os.path.join(cfg.workspace_dir, entry.filename)) + ] + if missing_files: + raise _fresh_save_error( + f"The SGLang graph index references missing graph files: {missing_files}." + ) # NVSHMEM init runs once before any graph loads — graphs may reference # NVSHMEM symbols. Single-GPU dense models have 0 NVSHMEM modules, so # this is a no-op there but kept for EP parity. cge.init_nvshmem_for_loaded_modules() - paths = [os.path.join(cfg.workspace_dir, filename) for _, filename, _ in graph_files] + paths = [os.path.join(cfg.workspace_dir, entry.filename) for entry in entries] t0 = time.perf_counter() pending = FoundryCUDAGraph.start_graph_builds(paths, num_threads=4) results = FoundryCUDAGraph.finish_graph_loads(pending) + if len(results) != len(entries): + raise RuntimeError( + "Foundry loaded an unexpected number of SGLang graphs: " + f"expected {len(entries)}, got {len(results)}" + ) logger.info( "[Foundry] Loaded %d SGLang graphs in %.3fs", len(results), @@ -287,9 +362,9 @@ def load_all_graphs(_cuda_graph_runner) -> list[tuple[int, Any, Any]]: ) loaded = [] - for i, (_index, _filename, meta) in enumerate(graph_files): - graph, tensors = results[i] - loaded.append((meta["size"], graph, _unpack_output(tensors))) + for entry, result in zip(entries, results): + graph, tensors = result + loaded.append((entry, graph, tensors)) return loaded diff --git a/python/foundry/integration/sglang/hooks.py b/python/foundry/integration/sglang/hooks.py index 3217b416..bb3daed8 100644 --- a/python/foundry/integration/sglang/hooks.py +++ b/python/foundry/integration/sglang/hooks.py @@ -145,6 +145,7 @@ def patched(self, pre_model_load_memory): if mode == CUDAGraphExtensionMode.LOAD: rt.log_alloc_offset("before_init_memory_pool") state = rt.load_warmup_state() + rt.validate_warmup_state(state, self.server_args) if not state.memory_pool_config: raise RuntimeError("Foundry LOAD requires memory_pool_config") self.memory_pool_config = pool_configurator.MemoryPoolConfig(**state.memory_pool_config) @@ -164,7 +165,10 @@ def patched(self, pre_model_load_memory): rt.log_alloc_offset("before_init_memory_pool") result = orig(self, pre_model_load_memory) rt.log_alloc_offset("after_init_memory_pool") - state = rt.create_warmup_state(asdict(self.memory_pool_config)) + state = rt.create_warmup_state( + asdict(self.memory_pool_config), + self.server_args, + ) rt.save_warmup_state(state) return result diff --git a/python/foundry/integration/sglang/runtime.py b/python/foundry/integration/sglang/runtime.py index 1f9032c5..388b624c 100644 --- a/python/foundry/integration/sglang/runtime.py +++ b/python/foundry/integration/sglang/runtime.py @@ -11,6 +11,8 @@ import time from dataclasses import asdict, dataclass, field from datetime import datetime +from importlib.metadata import PackageNotFoundError +from importlib.metadata import version as package_version from pathlib import Path import torch @@ -28,10 +30,18 @@ logger = logging.getLogger(__name__) +WARMUP_STATE_SCHEMA_VERSION = 1 +SUPPORTED_SGLANG_VERSION = "0.5.15.post1" +_SUPPORTED_DECODE_BACKEND = "full" +_SUPPORTED_PREFILL_BACKEND = "disabled" + @dataclass class WarmupState: + schema_version: int = 0 sglang_version: str = "" + decode_cuda_graph_backend: str = "" + prefill_cuda_graph_backend: str = "" timestamp: str = "" cuda_version: str = "" gpu_name: str = "" @@ -44,7 +54,7 @@ class WarmupState: class CUDAGraphExtensionState: capture_index: int = 0 rank: int = 0 - loaded_graphs: dict = field(default_factory=dict) + graph_index_entries: list[object] = field(default_factory=list) _state: CUDAGraphExtensionState | None = None @@ -65,15 +75,35 @@ def _workspace_root() -> str | None: return None if cfg is None else cfg.workspace_root -def create_warmup_state(memory_pool_config: dict | None = None) -> WarmupState: +def _setting_value(value): + return getattr(value, "value", value) + + +def _cuda_graph_backends(server_args) -> tuple[str | None, str | None]: + cuda_graph_config = getattr(server_args, "cuda_graph_config", None) + decode = getattr(getattr(cuda_graph_config, "decode", None), "backend", None) + prefill = getattr(getattr(cuda_graph_config, "prefill", None), "backend", None) + return _setting_value(decode), _setting_value(prefill) + + +def _installed_sglang_version() -> str: try: - from sglang.version import __version__ as sglang_version - except Exception: - sglang_version = "unknown" + return package_version("sglang") + except PackageNotFoundError: + return "unknown" + +def create_warmup_state( + memory_pool_config: dict | None, + server_args, +) -> WarmupState: + decode_backend, prefill_backend = _cuda_graph_backends(server_args) props = torch.cuda.get_device_properties(0) return WarmupState( - sglang_version=sglang_version, + schema_version=WARMUP_STATE_SCHEMA_VERSION, + sglang_version=_installed_sglang_version(), + decode_cuda_graph_backend=decode_backend or "", + prefill_cuda_graph_backend=prefill_backend or "", timestamp=datetime.now().isoformat(), cuda_version=torch.version.cuda or "unknown", gpu_name=props.name, @@ -82,6 +112,58 @@ def create_warmup_state(memory_pool_config: dict | None = None) -> WarmupState: ) +def _compatibility_error(message: str) -> RuntimeError: + return RuntimeError(f"{message} Perform a fresh Foundry SAVE before using LOAD.") + + +def validate_warmup_state(state: WarmupState, server_args) -> None: + if ( + isinstance(state.schema_version, bool) + or state.schema_version != WARMUP_STATE_SCHEMA_VERSION + ): + raise _compatibility_error( + "Unsupported or legacy SGLang archive schema " + f"{state.schema_version!r}; expected {WARMUP_STATE_SCHEMA_VERSION}." + ) + + if state.sglang_version != SUPPORTED_SGLANG_VERSION: + raise _compatibility_error( + "SGLang archive version mismatch: " + f"recorded {state.sglang_version!r}, expected {SUPPORTED_SGLANG_VERSION!r}." + ) + installed_version = _installed_sglang_version() + if installed_version != SUPPORTED_SGLANG_VERSION: + raise _compatibility_error( + "Installed SGLang version mismatch: " + f"found {installed_version!r}, expected {SUPPORTED_SGLANG_VERSION!r}." + ) + + if state.decode_cuda_graph_backend != _SUPPORTED_DECODE_BACKEND: + raise _compatibility_error( + "SGLang archive decode CUDA graph backend mismatch: " + f"recorded {state.decode_cuda_graph_backend!r}, " + f"expected {_SUPPORTED_DECODE_BACKEND!r}." + ) + if state.prefill_cuda_graph_backend != _SUPPORTED_PREFILL_BACKEND: + raise _compatibility_error( + "SGLang archive prefill CUDA graph backend mismatch: " + f"recorded {state.prefill_cuda_graph_backend!r}, " + f"expected {_SUPPORTED_PREFILL_BACKEND!r}." + ) + + decode_backend, prefill_backend = _cuda_graph_backends(server_args) + if decode_backend != state.decode_cuda_graph_backend: + raise _compatibility_error( + "Current SGLang decode CUDA graph backend does not match the archive: " + f"current {decode_backend!r}, archive {state.decode_cuda_graph_backend!r}." + ) + if prefill_backend != state.prefill_cuda_graph_backend: + raise _compatibility_error( + "Current SGLang prefill CUDA graph backend does not match the archive: " + f"current {prefill_backend!r}, archive {state.prefill_cuda_graph_backend!r}." + ) + + def save_warmup_state(state: WarmupState) -> None: workspace_root = _workspace_root() if workspace_root is None: @@ -102,9 +184,14 @@ def load_warmup_state() -> WarmupState: raise RuntimeError("Foundry workspace_root is not initialized") path = os.path.join(workspace_root, "warmup_state.json") if not os.path.exists(path): - raise RuntimeError(f"Foundry warmup state file not found: {path}") - with open(path) as f: - data = json.load(f) + raise _compatibility_error(f"Foundry warmup state file not found: {path}.") + try: + with open(path) as f: + data = json.load(f) + except (OSError, json.JSONDecodeError) as exc: + raise _compatibility_error(f"Cannot read the SGLang warmup state at {path}: {exc}") from exc + if not isinstance(data, dict): + raise _compatibility_error("SGLang warmup state root must be an object.") valid = set(WarmupState.__dataclass_fields__.keys()) return WarmupState(**{k: v for k, v in data.items() if k in valid}) diff --git a/tests/test_sglang_runtime.py b/tests/test_sglang_runtime.py new file mode 100644 index 00000000..4bd168a2 --- /dev/null +++ b/tests/test_sglang_runtime.py @@ -0,0 +1,185 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the Foundry project +"""Pure-Python contracts for SGLang archive compatibility metadata.""" + +from __future__ import annotations + +import importlib.util +import json +import sys +from pathlib import Path +from types import ModuleType, SimpleNamespace + +import pytest + + +def _module(name: str, **attributes): + module = ModuleType(name) + for attr, value in attributes.items(): + setattr(module, attr, value) + return module + + +def _load_runtime(monkeypatch: pytest.MonkeyPatch, workspace_root: Path | None = None): + foundry = _module("foundry") + foundry.__path__ = [] + ops = _module("foundry.ops") + allocation_region = _module( + "foundry.allocation_region", + parse_size=lambda value: value, + ) + config = _module( + "foundry.integration.sglang.config", + CUDAGraphExtensionMode=SimpleNamespace(NONE="none", SAVE="save", LOAD="load"), + compute_workspace_rank=lambda *_args, **_kwargs: 0, + get_config=lambda: SimpleNamespace( + workspace_root=str(workspace_root) if workspace_root is not None else None, + workspace_dir=None, + ), + get_graph_extension_mode=lambda: "none", + get_hook_library_path=lambda: None, + get_nvshmem_host_path=lambda: None, + ) + torch = _module( + "torch", + cuda=SimpleNamespace( + get_device_properties=lambda _index: SimpleNamespace( + name="stub GPU", + total_memory=1024, + ) + ), + version=SimpleNamespace(cuda="13.0"), + ) + + for name, module in ( + ("torch", torch), + ("foundry", foundry), + ("foundry.ops", ops), + ("foundry.allocation_region", allocation_region), + ("foundry.integration.sglang.config", config), + ): + monkeypatch.setitem(sys.modules, name, module) + + path = ( + Path(__file__).parents[1] / "python" / "foundry" / "integration" / "sglang" / "runtime.py" + ) + spec = importlib.util.spec_from_file_location("_foundry_sglang_runtime_contract", path) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + monkeypatch.setitem(sys.modules, spec.name, module) + spec.loader.exec_module(module) + monkeypatch.setattr( + module, + "package_version", + lambda package: "0.5.15.post1" if package == "sglang" else "unexpected", + raising=False, + ) + return module + + +def _server_args(decode: str = "full", prefill: str = "disabled"): + return SimpleNamespace( + cuda_graph_config=SimpleNamespace( + decode=SimpleNamespace(backend=decode), + prefill=SimpleNamespace(backend=prefill), + ) + ) + + +def _matching_state(runtime): + return runtime.WarmupState( + schema_version=runtime.WARMUP_STATE_SCHEMA_VERSION, + sglang_version=runtime.SUPPORTED_SGLANG_VERSION, + decode_cuda_graph_backend="full", + prefill_cuda_graph_backend="disabled", + ) + + +def test_create_warmup_state_records_latest_compatibility_metadata(monkeypatch): + runtime = _load_runtime(monkeypatch) + + state = runtime.create_warmup_state({"page_size": 16}, _server_args()) + + assert state.schema_version == runtime.WARMUP_STATE_SCHEMA_VERSION + assert runtime.WARMUP_STATE_SCHEMA_VERSION == 1 + assert state.sglang_version == "0.5.15.post1" + assert runtime.SUPPORTED_SGLANG_VERSION == "0.5.15.post1" + assert state.decode_cuda_graph_backend == "full" + assert state.prefill_cuda_graph_backend == "disabled" + assert state.memory_pool_config == {"page_size": 16} + + +@pytest.mark.parametrize( + ("changes", "diagnostic"), + [ + ({"schema_version": 0}, "schema"), + ({"sglang_version": ""}, "metadata|version"), + ({"sglang_version": "0.5.14"}, "version"), + ({"decode_cuda_graph_backend": "breakable"}, "decode"), + ({"prefill_cuda_graph_backend": "full"}, "prefill"), + ], +) +def test_validate_warmup_state_rejects_incompatible_archive( + monkeypatch, + changes, + diagnostic, +): + runtime = _load_runtime(monkeypatch) + state = _matching_state(runtime) + for name, value in changes.items(): + setattr(state, name, value) + + with pytest.raises( + (RuntimeError, ValueError), + match=rf"(?is){diagnostic}.*fresh.*SAVE|fresh.*SAVE.*{diagnostic}", + ): + runtime.validate_warmup_state(state, _server_args()) + + +@pytest.mark.parametrize( + ("server_args", "diagnostic"), + [ + (_server_args(decode="breakable"), "decode"), + (_server_args(prefill="full"), "prefill"), + ], +) +def test_validate_warmup_state_rejects_current_config_mismatch( + monkeypatch, + server_args, + diagnostic, +): + runtime = _load_runtime(monkeypatch) + + with pytest.raises( + (RuntimeError, ValueError), + match=rf"(?is){diagnostic}.*fresh.*SAVE|fresh.*SAVE.*{diagnostic}", + ): + runtime.validate_warmup_state(_matching_state(runtime), server_args) + + +def test_validate_warmup_state_accepts_matching_latest_archive(monkeypatch): + runtime = _load_runtime(monkeypatch) + + runtime.validate_warmup_state(_matching_state(runtime), _server_args()) + + +def test_load_warmup_state_ignores_unknown_fields_but_preserves_missing_metadata( + tmp_path, + monkeypatch, +): + runtime = _load_runtime(monkeypatch, tmp_path) + (tmp_path / "warmup_state.json").write_text( + json.dumps( + { + "timestamp": "legacy", + "unknown_future_field": {"safe": True}, + } + ) + ) + + state = runtime.load_warmup_state() + + assert state.timestamp == "legacy" + assert state.schema_version == 0 + with pytest.raises((RuntimeError, ValueError), match=r"(?i)fresh.*SAVE|SAVE.*fresh"): + runtime.validate_warmup_state(state, _server_args()) From 21b06e15b83085f57a426cb90d00a9294f23434e Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Thu, 23 Jul 2026 07:18:47 +0000 Subject: [PATCH 05/21] fix(sglang): validate usable graph archives Co-authored-by: Rahul Chalamala --- .../foundry/integration/sglang/graph_ops.py | 4 +- python/foundry/integration/sglang/runtime.py | 32 +++++++- tests/test_sglang_graph_ops.py | 78 +++++++++++++++++-- tests/test_sglang_runtime.py | 38 +++++++++ 4 files changed, 143 insertions(+), 9 deletions(-) diff --git a/python/foundry/integration/sglang/graph_ops.py b/python/foundry/integration/sglang/graph_ops.py index f7a0d4cd..f9f80e45 100644 --- a/python/foundry/integration/sglang/graph_ops.py +++ b/python/foundry/integration/sglang/graph_ops.py @@ -169,6 +169,8 @@ def _normalize_entries(entries: Any) -> list[GraphIndexEntry]: (_normalize_entry(entry) for entry in entries), key=lambda entry: entry.capture_index, ) + if normalized and normalized[0].capture_index != 0: + raise _fresh_save_error("SGLang graph capture_index values must start at zero.") filenames = [entry.filename for entry in normalized] if len(filenames) != len(set(filenames)): raise _fresh_save_error("SGLang graph index filenames must be unique.") @@ -270,7 +272,7 @@ def save_graph( GraphIndexEntry( phase=_SUPPORTED_PHASE, capture_index=state.capture_index, - filename=f"sglang_graph_{state.capture_index:06d}.json", + filename=f"graph_{state.capture_index:06d}_sglang.json", shape_key=shape_key, output_kind=output_kind, ) diff --git a/python/foundry/integration/sglang/runtime.py b/python/foundry/integration/sglang/runtime.py index 388b624c..b12a1205 100644 --- a/python/foundry/integration/sglang/runtime.py +++ b/python/foundry/integration/sglang/runtime.py @@ -8,6 +8,7 @@ import logging import os import shutil +import tempfile import time from dataclasses import asdict, dataclass, field from datetime import datetime @@ -97,11 +98,19 @@ def create_warmup_state( memory_pool_config: dict | None, server_args, ) -> WarmupState: + sglang_version = _installed_sglang_version() + if sglang_version != SUPPORTED_SGLANG_VERSION: + raise RuntimeError( + "Foundry SAVE requires SGLang package version " + f"{SUPPORTED_SGLANG_VERSION!r}, but installed package metadata reports " + f"{sglang_version!r}. Install SGLang {SUPPORTED_SGLANG_VERSION} as a package " + "before SAVE so its exact version can be recorded." + ) decode_backend, prefill_backend = _cuda_graph_backends(server_args) props = torch.cuda.get_device_properties(0) return WarmupState( schema_version=WARMUP_STATE_SCHEMA_VERSION, - sglang_version=_installed_sglang_version(), + sglang_version=sglang_version, decode_cuda_graph_backend=decode_backend or "", prefill_cuda_graph_backend=prefill_backend or "", timestamp=datetime.now().isoformat(), @@ -173,8 +182,25 @@ 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: - json.dump(asdict(state), f, indent=2) + temporary_path = None + try: + with tempfile.NamedTemporaryFile( + "w", + encoding="utf-8", + dir=workspace_root, + prefix=".warmup_state.", + suffix=".tmp", + delete=False, + ) as temporary: + temporary_path = temporary.name + json.dump(asdict(state), temporary, indent=2, sort_keys=True) + temporary.write("\n") + temporary.flush() + os.fsync(temporary.fileno()) + os.replace(temporary_path, path) + finally: + if temporary_path is not None and os.path.exists(temporary_path): + os.unlink(temporary_path) logger.info("[Foundry] Saved SGLang warmup state to %s", path) diff --git a/tests/test_sglang_graph_ops.py b/tests/test_sglang_graph_ops.py index 42a92956..3494fe2c 100644 --- a/tests/test_sglang_graph_ops.py +++ b/tests/test_sglang_graph_ops.py @@ -9,7 +9,7 @@ import sys from dataclasses import dataclass from pathlib import Path -from types import ModuleType +from types import ModuleType, SimpleNamespace import pytest @@ -67,10 +67,22 @@ def _load_graph_ops(monkeypatch: pytest.MonkeyPatch): return module +def _load_real_manifest_scanner(monkeypatch: pytest.MonkeyPatch): + ops = sys.modules["foundry.ops"] + ops.CUDAGraph = object + path = Path(__file__).parents[1] / "python" / "foundry" / "graph.py" + spec = importlib.util.spec_from_file_location("foundry._graph_manifest_contract", path) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + monkeypatch.setitem(sys.modules, spec.name, module) + spec.loader.exec_module(module) + return module.save_graph_manifest + + def _entry(graph_ops, **overrides): values = { "phase": "decode", - "capture_index": 3, + "capture_index": 0, "filename": "opaque-name.json", "shape_key": ShapeKey(size=17, stream_idx=2, variant_label="nolora"), "output_kind": "next_token_logits", @@ -113,7 +125,7 @@ def test_versioned_graph_index_round_trips_complete_metadata(tmp_path, monkeypat _assert_entry( loaded[0], phase="decode", - capture_index=3, + capture_index=0, filename="opaque-name.json", size=17, stream_idx=2, @@ -121,7 +133,7 @@ def test_versioned_graph_index_round_trips_complete_metadata(tmp_path, monkeypat ) assert payload["graphs"][0] == { "phase": "decode", - "capture_index": 3, + "capture_index": 0, "filename": "opaque-name.json", "shape_key": { "size": 17, @@ -132,6 +144,53 @@ def test_versioned_graph_index_round_trips_complete_metadata(tmp_path, monkeypat } +def test_saved_filename_matches_real_manifest_scanner_without_key_identity( + tmp_path, + monkeypatch, +): + graph_ops = _load_graph_ops(monkeypatch) + save_graph_manifest = _load_real_manifest_scanner(monkeypatch) + state = SimpleNamespace(capture_index=0, graph_index_entries=[]) + monkeypatch.setattr( + graph_ops, + "get_config", + lambda: SimpleNamespace(workspace_dir=str(tmp_path)), + ) + monkeypatch.setattr(graph_ops.rt, "get_state", lambda: state) + + class Graph: + def save(self, path, _output): + Path(path).write_text( + json.dumps( + { + "topology_key": "opaque-topology", + "dependencies": [], + } + ) + ) + + shape_key = graph_ops.ArchiveShapeKey( + size=17, + stream_idx=2, + variant_label="nolora", + ) + graph_ops.save_graph( + Graph(), + object(), + shape_key, + "next_token_logits", + ) + filename = state.graph_index_entries[0].filename + + save_graph_manifest(str(tmp_path)) + + assert filename == "graph_000000_sglang.json" + assert int(filename.split("_")[1]) == 0 + assert all(identity not in filename for identity in ("17", "2", "nolora")) + manifest = json.loads((tmp_path / "graph_manifest.json").read_text()) + assert manifest["topology_groups"][0]["members"] == [filename] + + @pytest.mark.parametrize("schema_version", [None, 999]) def test_graph_index_rejects_legacy_or_unsupported_schema(tmp_path, monkeypatch, schema_version): graph_ops = _load_graph_ops(monkeypatch) @@ -220,6 +279,15 @@ def test_graph_index_load_order_is_capture_index_not_filename(tmp_path, monkeypa assert [entry.shape_key.size for entry in loaded] == [8, 4] +def test_graph_index_requires_capture_indices_to_start_at_zero(tmp_path, monkeypatch): + graph_ops = _load_graph_ops(monkeypatch) + index_path = tmp_path / "sglang_graph_index.json" + entry = _entry(graph_ops, capture_index=3) + + with pytest.raises((RuntimeError, ValueError), match=r"(?i)capture.*zero|fresh.*SAVE"): + graph_ops.save_graph_index(index_path, [entry]) + + @pytest.mark.parametrize("legacy_key", [8, "1_8"]) def test_graph_index_rejects_integer_and_string_shape_key_heuristics( tmp_path, monkeypatch, legacy_key @@ -263,7 +331,7 @@ def test_index_metadata_not_filename_regex_is_source_of_truth(tmp_path, monkeypa _assert_entry( loaded[0], phase="decode", - capture_index=3, + capture_index=0, filename="graph_42_FULL_t999_r999_UX_pcN.json", size=7, stream_idx=None, diff --git a/tests/test_sglang_runtime.py b/tests/test_sglang_runtime.py index 4bd168a2..01f334c8 100644 --- a/tests/test_sglang_runtime.py +++ b/tests/test_sglang_runtime.py @@ -109,6 +109,21 @@ def test_create_warmup_state_records_latest_compatibility_metadata(monkeypatch): assert state.memory_pool_config == {"page_size": 16} +def test_create_warmup_state_rejects_missing_sglang_package_metadata(monkeypatch): + runtime = _load_runtime(monkeypatch) + + def missing_package_version(_package): + raise runtime.PackageNotFoundError("sglang") + + monkeypatch.setattr(runtime, "package_version", missing_package_version) + + with pytest.raises( + RuntimeError, + match=r"(?is)SAVE.*install.*0\.5\.15\.post1|install.*0\.5\.15\.post1.*SAVE", + ): + runtime.create_warmup_state({}, _server_args()) + + @pytest.mark.parametrize( ("changes", "diagnostic"), [ @@ -163,6 +178,29 @@ def test_validate_warmup_state_accepts_matching_latest_archive(monkeypatch): runtime.validate_warmup_state(_matching_state(runtime), _server_args()) +def test_save_warmup_state_replaces_existing_file_atomically(tmp_path, monkeypatch): + runtime = _load_runtime(monkeypatch, tmp_path) + state_path = tmp_path / "warmup_state.json" + original = '{"sentinel": "existing"}\n' + state_path.write_text(original) + replace_calls = [] + + def fail_replace(source, destination): + replace_calls.append((Path(source), Path(destination))) + raise OSError("simulated interrupted replacement") + + monkeypatch.setattr(runtime.os, "replace", fail_replace) + + with pytest.raises(OSError, match="interrupted replacement"): + runtime.save_warmup_state(_matching_state(runtime)) + + assert replace_calls + assert replace_calls[0][0].parent == tmp_path + assert replace_calls[0][1] == state_path + assert state_path.read_text() == original + assert list(tmp_path.iterdir()) == [state_path] + + def test_load_warmup_state_ignores_unknown_fields_but_preserves_missing_metadata( tmp_path, monkeypatch, From 5a57c8816f1514641f32ed6438e2c76ae126a7cf Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Thu, 23 Jul 2026 07:31:50 +0000 Subject: [PATCH 06/21] docs(sglang): update v0.5.15 recipes Co-authored-by: Rahul Chalamala --- README.md | 8 +- docs/overview.md | 108 ++++---- docs/sglang/direct-edits.md | 164 ++++++++---- docs/sglang/hooks.md | 292 +++++++++------------- docs/sglang/memory-consistency.md | 183 +++++++------- docs/sglang/memory-lifecycle.md | 191 ++++++++------ docs/sglang/overview.md | 207 ++++++++------- docs/sglang/save-load-workflow.md | 231 ++++++++++------- recipe/sglang/README.md | 246 ++++++++++-------- recipe/sglang/serve_qwen3-1.7b_dp.sh | 5 +- recipe/sglang/serve_qwen3-30ba3bfp8_ep.sh | 7 +- recipe/sglang/serve_qwen3-mini.sh | 5 +- tests/test_sglang_recipes.py | 24 ++ 13 files changed, 912 insertions(+), 759 deletions(-) create mode 100644 tests/test_sglang_recipes.py diff --git a/README.md b/README.md index bec64ee4..b50c9d75 100644 --- a/README.md +++ b/README.md @@ -82,11 +82,17 @@ Foundry ships engine integrations under `foundry/python/foundry/integration/`. P | Engine | Single GPU | DP | TP | EP | |---|:---:|:---:|:---:|:---:| | vLLM | ✅ | ✅ | 🚧 | ✅ | -| SGLang | ✅ | ✅ | 🚧 | ✅ | +| SGLang `0.5.15.post1` | ✅ | ✅ | 🚧 | ✅ | | TensorRT-LLM | 🚧 | 🚧 | 🚧 | 🚧 | ✅ validated end-to-end (SAVE → LOAD → query)  ·  🚧 not yet +The SGLang integration is latest-only for upstream `v0.5.15.post1` +(`0b3bb0c`). It supports full decode graphs with prefill graphs disabled for +single-GPU FlashInfer, regular DP FlashInfer, and DP-attention with DeepEP +low-latency + FA3. Existing SGLang archives are incompatible and require a +fresh SAVE. See [`recipe/sglang/README.md`](recipe/sglang/README.md). + The adapted vLLM / SGLang / TensorRT-LLM forks will be released alongside this repo at `foundry-org/vllm`, `foundry-org/sglang`, `foundry-org/TensorRT-LLM`. ### Performance diff --git a/docs/overview.md b/docs/overview.md index 46a92063..9e65c7bf 100644 --- a/docs/overview.md +++ b/docs/overview.md @@ -1,72 +1,64 @@ -# Foundry Inference-Engine Integrations +# Foundry inference-engine integrations -Foundry ships two inference-engine integrations that wrap an unmodified serving stack with cold-start graph persistence: +Foundry supplies runtime integration packages for vLLM and SGLang. A small +engine-side shim activates each package; deterministic memory, graph +persistence, packed module restoration, and lifecycle hooks remain in Foundry. | Engine | Integration code | Documentation | |---|---|---| -| vLLM | `foundry/python/foundry/integration/vllm/` | [`vllm/overview.md`](vllm/overview.md) | -| SGLang | `foundry/python/foundry/integration/sglang/` | [`sglang/overview.md`](sglang/overview.md) | +| vLLM | `python/foundry/integration/vllm/` | [`vllm/overview.md`](vllm/overview.md) | +| SGLang `0.5.15.post1` | `python/foundry/integration/sglang/` | [`sglang/overview.md`](sglang/overview.md) | -Both follow the same shape: a tiny shim in the host engine's source tree re-exports `install_hooks(...)` and `get_graph_extension_mode()`, and the shim is called once from the host engine's config object's `__post_init__`. `install_hooks` then installs a small set of runtime monkey-patches that wrap the engine's lifecycle methods. Everything substantive — VMM region setup, graph save/load, warmup skipping — lives in the foundry package. +## Shared model -## Mental model +Both integrations: -```mermaid -flowchart BT - subgraph engine["Host engine (unmodified except shim + minimal guards)"] - VLLM["vLLM core
~97 lines across 5 files"] - SGL["SGLang core
~47 lines across 4 files"] - end +1. reserve a fixed-address VMM region before model/runtime allocations; +2. capture full or piecewise CUDA graphs during SAVE; +3. persist graph structure, packed device code, topology groups, and allocation + state; +4. recreate model/runtime allocations at the same offsets in a fresh LOAD + process; +5. batch-build graph templates and link on-demand members; +6. place `libcuda_hook.so` in child-process `LD_PRELOAD` before CUDA starts. - subgraph shim["Engine-side shims (re-export + activation)"] - VLLM_SHIM["vllm/compilation/foundry_shim.py
re-export + ImportError fallback"] - SGL_SHIM["sglang/srt/foundry_shim.py
re-export + force-disable flags"] - end +Each engine owns different capture and warmup seams, so the engine-specific +hooks and archive metadata are intentionally separate. - subgraph foundry["foundry/python/foundry/integration/"] - CFG["config.py — CUDAGraphExtensionConfig"] - RT["runtime.py — WarmupState, VMM region,
final_alloc_offset"] - GOPS["graph_ops.py — save/load, manifest, fatbins"] - HOOKS["hooks.py — install_hooks() + monkey-patches"] - end +## Integration comparison - VLLM -- "CompilationConfig.__post_init__" --> VLLM_SHIM - SGL -- "ServerArgs.__post_init__" --> SGL_SHIM - VLLM_SHIM -- "install_hooks(compilation_config)" --> HOOKS - SGL_SHIM -- "install_hooks(server_args)" --> HOOKS - - HOOKS --> CFG - HOOKS --> RT - HOOKS --> GOPS -``` - -## What both integrations share - -- **SAVE→LOAD cycle.** SAVE runs the full cold-start path once, captures CUDA graphs + kernel fatbins + a warmup-state JSON, and writes them to a per-rank workspace. LOAD restores from the workspace, skipping graph capture and kernel warmup. -- **Deterministic VMM region.** Both integrations call `foundry::set_allocation_region(base_addr, region_size)` before NCCL / distributed init. Every allocation inside the region (model weights, KV pool, attention workspaces, captured-graph alloc events) lands at a byte-deterministic offset, so SAVE and LOAD address the same VMM offsets. Only memory *referenced by captured CUDA graphs* needs the region; the vLLM integration now calls `foundry::stop_allocation_region()` at the end of `capture_model` so subsequent runtime allocations (sampler warmup, scheduler buffers, inference workspaces) fall back to the standard CUDA caching allocator instead of paying VMM overhead. The SGLang side will follow. -- **`final_alloc_offset` watermark.** SAVE records the final cursor position; LOAD preallocates the same range via a single `cuMemCreate+cuMemMap`, then bump-allocates within it. Because `stop_allocation_region` fires before sampler warmup, the watermark reflects only the VMM-tracked working set, not transient post-capture allocations. -- **Manifest-driven graph load.** Captured graphs are grouped into topology equivalence classes; one graph per group is built as a *template* and the rest become *on-demand* graphs that share the template's executor via per-graph node updates. -- **Spawn-time `LD_PRELOAD`.** `libcuda_hook.so` must be in `LD_PRELOAD` before any CUDA call so its `cuMemAlloc_v2` / `cuModuleLoad*` interposers run in every process. The integrations set the env var before child spawn; the serve scripts also export it from the shell as defense-in-depth. - -## What differs between the two - -| Aspect | vLLM | SGLang | +| Aspect | vLLM | SGLang `0.5.15.post1` | |---|---|---| -| Graph capture seam | `CUDAGraphWrapper.__call__` (per piecewise wrapper) | `CudaGraphRunner.capture` (centralized full-graph) | -| Pre-graph warmup | `kernel_warmup` (no-op), compile-warmup `_dummy_run` loop (empty under our config), sampler-warmup full forward (skipped via signature-targeted `_dummy_run` patch — see vllm/hooks.md §3) | `kernel_warmup` + 2 in-place warmup forwards inside `capture_one_batch_size` | -| Profile forward | yes (`memory_profiling` runs full forward) → two-pass SAVE required | no (only `_profile_available_bytes` samples free memory) → single-pass SAVE | -| Compile interaction | `torch.compile` integrated; disabled on LOAD via `do_not_compile=True` | `torch.compile` is the piecewise path; force-disabled via `disable_piecewise_cuda_graph=True` | -| Per-bs metadata | allocated **inside** captured graph (recorded as alloc events) | allocated **outside** captured graph (FlashInfer wrappers) → pre-pass init + `reuse_pre_pass_init` shim needed | -| MoE / EP | DeepEP all2all + NVSHMEM module init handled in `hooks.py` (`_patch_deepep` forces `use_fabric=True`, `_patch_prepare_comm_buffer` post-hook flushes `init_nvshmem_for_loaded_modules`) | coming soon | -| DP empty-shard phantom | `_patch_dummy_runs` lets `Worker.execute_dummy_batch`'s `_dummy_run(uniform_decode=True)` pass through to the real forward so the empty shard joins `coordinate_batch_across_dp` and the EP all2all (see [`vllm/hooks.md`](vllm/hooks.md) §3 and `claude-doc/invariants.md` §5) | n/a | -| Async graph builds | kicked off **after** weight load (driver contention with `load_weights` was net-negative); overlaps with KV-cache + scheduler bring-up instead | timing-independent: launched inside the capture context | -| Process model | uniproc / MP / Ray executors | `spawn`-based `mp.Process`; needs per-child `install_hooks` | -| Direct edits in host repo | 5 files, ~97 lines | 4 files, ~47 lines | - -The per-engine `overview.md` pages explain those differences in detail and how to use each integration. +| Supported graph scope | Configured vLLM piecewise/full integration modes | Full decode only; prefill disabled | +| Capture seam | `CUDAGraphWrapper.__call__` | `DecodeCudaGraphRunner.capture` orchestration and `FullCudaGraphBackend.capture_one` storage | +| Run-once warmup seam | vLLM compile/kernel/sampler lifecycle hooks | `BaseRunner.warmup` | +| Graph identity | vLLM piecewise graph metadata | `ShapeKey(size, stream_idx, variant_label)` | +| Backend storage | vLLM wrapper state | `FullCudaGraphBackend._graphs`, `_outputs`, `_pool` | +| Attention preparation | vLLM-specific compile/capture path | Separate out-of-graph planning and in-graph GPU metadata work | +| Profile behavior | Full profile forward; two-pass SAVE recipe | Pool-memory measurement only; one SAVE process | +| Process model | Uniprocess, multiprocessing, or Ray executors | Multiprocessing `spawn`; explicit hook install in scheduler and DP-controller children | +| Validated parallel modes | Single GPU, DP, and documented EP configuration | Single GPU FlashInfer, regular DP FlashInfer, and DP-attention + DeepEP low-latency + FA3 | + +For SGLang, TP without DP-attention, PP, speculative decoding, LoRA/PDMux, +dLLM, hidden-state capture, memory-saver graphs, and prefill graphs are not +supported. + +## Archive authority + +The SGLang integration adds a versioned `sglang_graph_index.json` to each rank +workspace. It records capture order, phase, output kind, filename, and complete +shape identity. LOAD validates that index plus `warmup_state.json` schema, +SGLang version, and resolved decode/prefill backends before restoring graphs. +Earlier SGLang archives require a fresh SAVE. + +vLLM uses its own integration-specific graph records. Both engines use +Foundry's `graph_manifest.json` for native topology/template relationships. ## Reading order -- New to the integrations: read `vllm/overview.md` or `sglang/overview.md` end-to-end. -- Debugging a SAVE↔LOAD bug: read `vllm/memory-consistency.md` or `sglang/memory-consistency.md`. -- Porting to a new engine: this file + `/direct-edits.md` + `/hooks.md` show the minimal contract. +- Start with [`vllm/overview.md`](vllm/overview.md) or + [`sglang/overview.md`](sglang/overview.md). +- Use the corresponding `direct-edits.md` for the external engine shim. +- Use `memory-lifecycle.md` and `memory-consistency.md` when debugging + allocation parity. +- Use `save-load-workflow.md` for commands and archive validation. diff --git a/docs/sglang/direct-edits.md b/docs/sglang/direct-edits.md index 49c983c9..86e8d1cf 100644 --- a/docs/sglang/direct-edits.md +++ b/docs/sglang/direct-edits.md @@ -1,93 +1,151 @@ -# SGLang Direct Edits +# SGLang v0.5.15.post1 direct edits -Total footprint in the `sglang/` repo: **47 lines across 4 files**. Every other line of integration logic lives in `foundry.integration.sglang`. +Foundry's runtime integration lives in this repository. The external +`foundry-org/sglang` checkout needs only the following four-file shim, applied +to upstream commit `0b3bb0cbe31873994c9f989fddfe2f87ca839fdd` +(`v0.5.15.post1`). -## 1. `sglang/srt/foundry_shim.py` (new, ~27 lines) - -The activation entry point. Re-exports the foundry public API and forces flags that conflict with SAVE/LOAD determinism. +## 1. Add `python/sglang/srt/foundry_shim.py` ```python +from foundry.integration.sglang import install_hooks +from sglang.srt.model_executor.cuda_graph_config import Backend + + def apply_server_args(server_args) -> None: - cfg_path = getattr(server_args, "foundry_graph_extension_config_path", None) - if not cfg_path: + if not server_args.foundry_graph_extension_config_path: return - server_args.disable_piecewise_cuda_graph = True + server_args.cuda_graph_config.decode.backend = Backend.FULL + server_args.cuda_graph_config.prefill.backend = Backend.DISABLED server_args.enable_profile_cuda_graph = False server_args.disable_flashinfer_autotune = True - - from foundry.integration.sglang.hooks import install_hooks install_hooks(server_args) ``` -`install_hooks` is idempotent (guarded by a module-level `_INSTALLED` flag) so calling it multiple times across processes is safe. +The shim receives already resolved `cuda_graph_config` objects. It forces the +only supported graph shape: full decode and disabled prefill. It also disables +graph profiling and FlashInfer autotuning before calling Foundry's public +`foundry.integration.sglang.install_hooks` entry point. -## 2. `sglang/srt/server_args.py` (~12 lines) +## 2. Edit `python/sglang/srt/server_args.py` -Adds the config-path field, the CLI argument, and the `__post_init__` activation: +Add the field in the `Custom hooks, probe, and plugins` section: ```python -class ServerArgs: - ... - foundry_graph_extension_config_path: Optional[str] = None - ... - - def __post_init__(self): - ... - if self.foundry_graph_extension_config_path: - from sglang.srt.foundry_shim import apply_server_args - apply_server_args(self) - ... - - @staticmethod - def add_cli_args(parser): - ... - parser.add_argument( - "--foundry-graph-extension-config-path", - type=str, - default=ServerArgs.foundry_graph_extension_config_path, - help="Path to Foundry CUDA graph extension TOML config.", - ) +foundry_graph_extension_config_path: A[ + Optional[str], "Path to the Foundry CUDA graph extension TOML config." +] = None ``` -The field must live on `ServerArgs` (not be monkey-attached) because: +SGLang v0.5.15.post1 derives `--foundry-graph-extension-config-path` +automatically from `A[...]` dataclass metadata through +`add_cli_args_from_dataclass`. Do not add a manual `parser.add_argument`; +manual entries in this release are reserved for deprecated or dynamic flags. + +Keep `__post_init__` as an ordered dispatcher. Add a narrow helper: + +```python +def _handle_foundry_graph_extension(self): + if self.foundry_graph_extension_config_path: + # Foundry is optional; import its SGLang shim only when configured. + from sglang.srt.foundry_shim import apply_server_args -- `argparse` consumes it from the CLI. -- `ServerArgs` is pickled across the `spawn` boundary; only declared fields survive serialization. + apply_server_args(self) +``` -## 3. `sglang/srt/managers/scheduler.py` (~4 lines) +Call it after graph resolution and every platform-default pass, but before GPU +memory sizing: ```python -def run_scheduler_process(server_args, ...): +self._handle_cuda_graph_config() + +self._handle_hpu_backends() +self._handle_cpu_backends() +self._handle_npu_backends() +self._handle_mps_backends() +self._handle_xpu_backends() +current_platform.apply_server_args_defaults(self) + +self._handle_foundry_graph_extension() + +gpu_mem = get_device_memory_capacity(self.device) +self._handle_gpu_memory_settings(gpu_mem) +``` + +This order is required. Foundry must mutate resolved +`cuda_graph_config.decode.backend` and +`cuda_graph_config.prefill.backend`; `_handle_gpu_memory_settings(gpu_mem)` +must then size decode graph batches and memory from those final values. + +The declared field is needed because SGLang's dataclass-generated CLI and +`ServerArgs.from_cli_args` handle declared fields. It also makes the value part +of the normal `ServerArgs` object passed to spawned children. The reason is not +that Python pickle cannot serialize arbitrary attributes. + +## 3. Edit `python/sglang/srt/managers/scheduler.py` + +Install hooks explicitly in every scheduler child immediately after plugins +load: + +```python +def run_scheduler_process( + server_args: ServerArgs, + # ... existing parameters ... +): load_plugins() if server_args.foundry_graph_extension_config_path: from sglang.srt.foundry_shim import apply_server_args + apply_server_args(server_args) - ... + + dp_rank = configure_scheduler_process( + # ... existing arguments ... + ) ``` -This is the first thing every spawned scheduler child runs. Without it, foundry's runtime monkey-patches aren't installed in the scheduler process and the model-loading lifecycle takes the upstream path. +## 4. Edit `python/sglang/srt/managers/data_parallel_controller.py` -## 4. `sglang/srt/managers/data_parallel_controller.py` (~4 lines) +Install hooks explicitly in the DP-controller child after logging is +configured and before it constructs child schedulers: ```python -def run_data_parallel_controller_process(server_args, ...): +def run_data_parallel_controller_process( + server_args: ServerArgs, + port_args: PortArgs, + pipe_writer, + run_scheduler_process_func: Callable = run_scheduler_process, +): + setproctitle.setproctitle("sglang::data_parallel_controller") + faulthandler.enable() + kill_itself_when_parent_died() parent_process = psutil.Process().parent() + configure_logger(server_args) if server_args.foundry_graph_extension_config_path: from sglang.srt.foundry_shim import apply_server_args + apply_server_args(server_args) - ... + + if server_args.enable_trace: + # ... existing code ... ``` -Same purpose as the scheduler call, but for the DP-controller child. This child later spawns its own scheduler subprocesses, so it must have foundry's spawn-site patches installed before doing so — otherwise its nested `proc.start()` calls won't set `LD_PRELOAD` in their environments. +Both child-process calls are mandatory. SGLang uses the multiprocessing +`spawn` model, which transfers data but does not inherit parent-process Python +monkey-patches. The DP controller later creates scheduler children, so both +levels must install Foundry's hooks before their launch paths execute. -## Why these can't be moved into the integration layer +## Verify the external checkout + +After applying the four edits, install the checkout rather than exposing it as +source only: + +```bash +pip install -e sglang/python \ + --extra-index-url https://download.pytorch.org/whl/cu130 +python -c 'from importlib.metadata import version; assert version("sglang") == "0.5.15.post1"' +sglang serve --help | rg -- '--foundry-graph-extension-config-path' +``` -| Edit | Why direct | -|---|---| -| `foundry_shim.py` | Has to be importable as `sglang.srt.foundry_shim`. It's the only file foundry's hooks can rely on being present in the sglang import path. | -| `ServerArgs.foundry_graph_extension_config_path` | Needs to be a real dataclass field (argparse + pickle). | -| `ServerArgs.__post_init__` call | Earliest activation point in the parent process. | -| `run_scheduler_process` install | `spawn` does not inherit Python state; the child needs an explicit re-install. | -| `run_data_parallel_controller_process` install | Same reason; this child also spawns further children. | +SAVE and LOAD reject missing or non-exact SGLang distribution metadata. diff --git a/docs/sglang/hooks.md b/docs/sglang/hooks.md index cf5b4c82..a4843105 100644 --- a/docs/sglang/hooks.md +++ b/docs/sglang/hooks.md @@ -1,214 +1,168 @@ -# SGLang Hook Surface - -Every monkey-patch `install_hooks(server_args)` installs. +# SGLang hook surface + +`foundry.integration.sglang.install_hooks(server_args)` is the public, +idempotent entry point for SGLang `0.5.15.post1`. It loads the Foundry TOML, +validates the supported graph configuration, and installs the hooks below. + +## Installation order + +```text +ModelRunner.init_torch_distributed +ModelRunnerKVCacheMixin.init_memory_pool +ModelRunner.load_model +BaseRunner.warmup +DecodeCudaGraphRunner.capture +FullCudaGraphBackend.capture_one +process launch sites +``` -## Patches installed by `install_hooks` +The latest runner classes are imported only while hooks install, so Foundry +itself remains usable without the optional SGLang dependency. -In install order: +## Configuration validation -1. `_patch_init_torch_distributed` -2. `_patch_init_memory_pool` -3. `_patch_load_model` -4. `_patch_kernel_warmup` -5. `_patch_cuda_graph_capture` -6. `_patch_spawn_sites` +Before patching, Foundry requires: -> Removed: an earlier `_patch_model_runner_init` wrapped `ModelRunner.__init__` to stash `dp_rank` on `self` (upstream didn't expose it as an attribute). Upstream sglang now does `self.dp_rank = dp_rank` in the constructor, so our wrapper is no longer needed; `_resolve_dp_rank` reads `self.dp_rank` directly. +- resolved decode backend `full`; +- resolved prefill backend `disabled`; +- FlashInfer for single-GPU and regular-DP decode; +- FA3 plus DeepEP low-latency for DP-attention. -### 1. `ModelRunner.init_torch_distributed` +It rejects unsupported TP, PP, speculative, LoRA/PDMux, dLLM, +hidden-state-output, memory-saver, and debug-graph configurations. -**Bind the device first (multi-GPU correctness).** Upstream `init_torch_distributed` -calls `torch.get_device_module(self.device).set_device(self.gpu_id)` *inside* -itself. But `setup_graph_extension` reserves the VMM region with -`set_allocation_region`, which binds to **whatever CUDA device is current at call -time** — and we call it *before* upstream. For DP/TP/EP rank > 0 the current -device is still `cuda:0` at that point, so without intervention the region lands -on the wrong GPU and the rank's later allocations fault with an **async illegal -memory access** that surfaces at an unrelated op (e.g. the first -`torch.cuda.Stream()` in `ModelRunner.__init__`). The patch therefore mirrors -upstream's `set_device(self.gpu_id)` up front (guarded on `self.device == "cuda"`) -before reserving the region. Single-GPU and rank 0 are unaffected (`gpu_id == 0`), -which is why this only appeared once DP was exercised. +## Distributed initialization -Before upstream: `set_device(self.gpu_id)`, then `setup_graph_extension(...)` -reserves the VMM region, loads cached fatbins (LOAD only via -`load_cuda_modules_and_libraries`), and eagerly initializes the cuBLAS handle into -scratch space. +The `ModelRunner.init_torch_distributed` wrapper sets the runner's CUDA device +before Foundry reserves its VMM range. This is required for ranks above zero: +the allocation region binds to the device current at reservation time. -After upstream: `skip_to_scratch_boundary()` forces the cursor to `cfg.scratch_space_size`. +Before upstream initialization, the hook: -``` -[Foundry] SGLang alloc_offset[after_setup_graph_ext]=… -[Foundry] SGLang alloc_offset[after_init_torch_dist]=… -[Foundry] SGLang alloc_offset[after_scratch_skip]=… -``` +1. computes the rank workspace; +2. prepares packed CUDA modules on LOAD; +3. reserves the configured VMM region; +4. eagerly creates the cuBLAS handle. -Validated on DP=2 (Qwen3-1.7B, one full replica per rank): both ranks reach an -identical `final_alloc_offset` across both SAVE passes and LOAD replays each -rank's graphs at that offset. Multi-rank runs also export -`NCCL_CUMEM_ENABLE=0` / `NCCL_NVLS_ENABLE=0` from the serve script — the CUMEM -P2P and NVLS multicast fast paths `cuMemMap` with driver-capability flags the -foundry VMM region doesn't carry. +After upstream initialization and communicator warmup, it advances the VMM +cursor to the configured scratch boundary. Nondeterministic startup +allocations remain below that boundary. -### 2. `ModelRunnerKVCacheMixin.init_memory_pool` +## Memory-pool initialization -- **SAVE**: call upstream `init_memory_pool` unchanged; after it returns, serialize the resolved `MemoryPoolConfig` (via `dataclasses.asdict`) into `warmup_state.json`. -- **LOAD**: skip upstream profiling. Load `MemoryPoolConfig` from `warmup_state.json`. Call `torch.cuda.empty_cache()` to mirror SAVE's `_resolve_memory_pool_config → get_available_gpu_memory(empty_cache=True)` side effect. Then call `_apply_memory_pool_config(config)` directly. +On SAVE, upstream `ModelRunnerKVCacheMixin.init_memory_pool` runs unchanged. +Foundry then serializes `dataclasses.asdict(self.memory_pool_config)` and +compatibility metadata to `warmup_state.json`. -The `empty_cache()` mirror is load-bearing — see [`memory-consistency.md`](memory-consistency.md) §5. +On LOAD, Foundry first validates the warmup state against the installed +distribution and current resolved graph configuration. It reconstructs +`MemoryPoolConfig`, calls `torch.cuda.empty_cache()` to mirror SAVE's profiling +side effect, and invokes `_apply_memory_pool_config` without profiling again. -### 3. `ModelRunner.load_model` +## Model loading -Currently a passthrough wrapper. Kept as a future hook point for a LOAD-time `start_graph_builds(all_paths)` overlap with weight IO; the current LOAD path is fast enough (~80 ms total) that overlap isn't needed. +`ModelRunner.load_model` currently delegates directly to upstream. It is kept +as a stable integration point; graph builds are not overlapped with weight +loading. -### 4. `ModelRunner.kernel_warmup` +## `BaseRunner.warmup` -No-op on SAVE/LOAD: +SAVE and LOAD return without executing the upstream warmup/autotune body: +```text +[Foundry] SGLang runner warmup skipped in mode ``` -[Foundry] SGLang kernel_warmup skipped in mode -``` - -The FlashInfer autotune path inside is shut off both by this no-op and by the forced `disable_flashinfer_autotune` flag. - -### 5. `CudaGraphRunner` (the largest patch) -Four sub-patches: +The external shim also sets `disable_flashinfer_autotune = True`. -#### 5a. `_create_device_graph` +## `DecodeCudaGraphRunner.capture` -On SAVE returns `foundry.CUDAGraph()` instead of `torch.cuda.CUDAGraph()`. +Mode `none` delegates to upstream. -#### 5b. `_capture_graph` +### SAVE -On SAVE enters `foundry.graph(graph, pool=pool, stream=stream)` instead of `torch.cuda.graph(...)`, captures the `run_once` callable's allocations into foundry's hook event log. +For FlashInfer, Foundry first walks `reversed(runner.capture_bs)`. For each +size, `capture_prepare(size)` constructs the canonical `ForwardBatch`, then +`init_forward_metadata_out_graph(forward_batch, in_capture=True)` creates the +allocation-bearing metadata. -#### 5c. `capture_one_batch_size` +During the subsequent upstream orchestration, a temporary wrapper: -On SAVE wraps the `forward` callable with a 3-call counter: +1. selects the prepared metadata by batch size; +2. sets it as the active `forward_metadata`; +3. calls the current out-of-graph planner with `in_capture=False`. -```python -counter = [0]; real_forward = forward -def warmup_skipping_forward(*args, **kwargs): - counter[0] += 1 - if counter[0] <= 2: - return None - return real_forward(*args, **kwargs) -forward = warmup_skipping_forward -``` +The normal `run_once` still invokes +`init_forward_metadata_in_graph(forward_batch)` inside graph capture, so +graph-recordable attention operations remain part of the graph. FA3 does not +need the FlashInfer allocation-reuse wrapper. -This suppresses the two pre-capture warmup forwards SGLang does in `for _ in range(2): run_once()`. Only the third invocation — inside `_capture_graph`'s graph capture context — runs the real forward. JIT and autotune allocations that those warmups would normally trigger now happen inside the captured graph and become foundry alloc events that replay verbatim on LOAD. (See `memory-consistency.md` §2.) +After capture, Foundry writes the graph index, topology manifest, packed +fatbins, and final allocation watermark. -After upstream returns, calls `save_graph(graph, output, key)`. The key matches the inline shape upstream uses for `self.graphs[key]` in `_capture_one_stream`: +### LOAD -```python -key = bs if stream_idx is None else f"{stream_idx}_{bs}" -``` +LOAD rejects a second capture request after hydration. On the first call it: -(Sglang removed the `_make_graph_key` / `get_capture_lora_variant` helpers in commit `ce2506e1c`, the same commit that deprecated `record_nolora_graph` dual MoE graph capture.) +1. requires `FullCudaGraphBackend`; +2. creates or reuses the graph pool and registers it with both SGLang pool + registries; +3. bootstraps DeepEP when active; +4. pre-maps memory through SAVE's final watermark; +5. recreates per-shape attention metadata in SAVE order; +6. batch-loads every graph named by `sglang_graph_index.json`; +7. reconstructs each real `ShapeKey`; +8. fills backend `_graphs`, `_outputs`, and `_pool`; +9. initializes the DeepEP adapter for decode replay. -#### 5d. `capture` +## `FullCudaGraphBackend.capture_one` -The outermost replacement. +On SAVE, the replacement bypasses the backend's two eager warmup forwards. It +creates a Foundry graph, executes `forward_fn` once inside Foundry's graph +context, and stores the resulting graph/output in backend `_graphs` and +`_outputs`. -**SAVE**: +Every archive record preserves: -```python -initialize_all_attention_metadata(self) # pre-pass -attn_backend.forward_metadata = None -real_init = attn_backend.init_forward_metadata_capture_cuda_graph -attn_backend.init_forward_metadata_capture_cuda_graph = reuse_pre_pass_init -try: - result = orig_capture(self, *args, **kwargs) # upstream capture loop -finally: - attn_backend.init_forward_metadata_capture_cuda_graph = real_init -save_graph_manifest() -pack_fatbins() -capture_final_alloc_offset() +```text +phase = decode +capture_index +filename +ShapeKey(size, stream_idx, variant_label) +output_kind = next_token_logits ``` -`initialize_all_attention_metadata` walks `reversed(self.capture_bs)` and pre-allocates every per-bs FlashInfer wrapper. The wrappers are stored in `attn_backend.decode_cuda_graph_metadata[bs]`. +Only `LogitsProcessorOutput.next_token_logits` is supported. Other populated +output fields fail explicitly. -`reuse_pre_pass_init` is a drop-in replacement for the upstream inner init that runs inside `capture_one_batch_size`. For decode mode it: +LOAD never enters this capture method; attempting to do so is treated as an +incompatible archive/configuration request. -1. Looks up the pre-pass wrapper from `decode_cuda_graph_metadata[bs]`. -2. Re-runs `indices_updater_decode.update(...)` with the same buffer slices (idempotent — writes plan info to the same `_int_workspace_buffer`). -3. Sets `attn_backend.forward_metadata = DecodeMetadata(wrappers)` so the captured forward sees the right metadata for this iter. +## DeepEP lifecycle -No allocation. The captured graph references the pre-pass wrapper's address; LOAD's pre-pass produces a wrapper at the same address. +For DP-attention + DeepEP low-latency + FA3: -**LOAD**: - -```python -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()) -preallocate_for_load_mode() # cuMemCreate+cuMemMap up to final_alloc_offset -initialize_all_attention_metadata(self) # pre-pass (same as SAVE) -load_all_graphs(self) # ONE start_graph_builds + finish_graph_loads -self.graphs = {k: v[0] for k, v in state.loaded_graphs.items()} -self.output_buffers = {k: v[1] for k, v in state.loaded_graphs.items()} -``` +- SAVE runs one eager orchestration pass with backend capture neutralized to + initialize DeepGEMM and other lazy kernel state outside stream capture. +- SAVE and LOAD bootstrap the singleton `DeepEPBuffer` before persisted graph + capture/load. +- Graph loading initializes queued NVSHMEM modules once before native graph + builds. +- LOAD invokes `deepep_adapter.capture(is_extend_in_batch=False)` after + backend hydration. -`load_all_graphs` calls `start_graph_builds(all_paths)` and `finish_graph_loads(pending)` exactly once each. This is required for the manifest's template + on-demand linking to work; per-graph `start_graph_builds([single_path])` calls would leave on-demand graphs without a `shared_exec`, and runtime replay would abort with `Called CUDAGraph::replay without a preceding successful capture or load`. +## Process launch hooks -### 6. Spawn-site patches +Before the parent launches schedulers, +`Engine._launch_scheduler_processes` calls +`setup_ld_preload_env()`. Before a DP controller launches its scheduler group, +`DataParallelController.launch_tensor_parallel_group` does the same. -Two parent-side wrappers: - -```python -# Engine._launch_scheduler_processes -def patched_launch(self, *args, **kwargs): - if get_graph_extension_mode() != CUDAGraphExtensionMode.NONE: - rt.setup_ld_preload_env() - return orig_launch(self, *args, **kwargs) - -# DataParallelController.launch_tensor_parallel_group -def patched_start(self, *args, **kwargs): - if get_graph_extension_mode() != CUDAGraphExtensionMode.NONE: - rt.setup_ld_preload_env() - return orig_start(self, *args, **kwargs) -``` +The helper prepends `libcuda_hook.so` and, when configured, the NVSHMEM host +library to `LD_PRELOAD`, and exports Foundry mode/timing metadata. Child +processes inherit this environment before any CUDA call. -`setup_ld_preload_env()` prepends `libcuda_hook.so` (and optionally `libnvshmem_host.so`) to `os.environ["LD_PRELOAD"]`, sets `FOUNDRY_MODE`, and records a wall-clock marker (`FOUNDRY_SPAWN_T0_NS`). All children spawned from these methods inherit the env. - -## Expert parallel (DeepEP) additions - -Active only when `moe_a2a_backend == deepep`. EP runs DP-attention + DeepEP (NCCL-free); -TP attention is unsupported (its NCCL all-reduce is incompatible with the VMM region). - -- **DeepEP buffer pre-capture bootstrap** (`bootstrap_deepep_buffer`, graph_ops). sglang - creates the singleton NVSHMEM `Buffer` lazily on the first MoE dispatch — normally - during the warmup forwards foundry suppresses, which would push creation *inside* the - captured stream (`deep_ep_cpp.Buffer(...)` → "operation not permitted when stream is - capturing"). The hook forces it before the capture loop, unwrapping the - `MaybeTboDeepEPDispatcher._inners` to reach a `DeepEPDispatcher`. Runs on SAVE and LOAD. -- **SAVE-only warmup pass** (`_run_warmup_pass`, `capture()` patch). Reuses the upstream - capture loop with graph capture neutered (run forwards only) to trigger every - pre-capture lazy init — DeepGEMM per-shape JIT (`stream.synchronize()` is illegal in - capture), buffer creation, etc. — outside the captured stream. LOAD doesn't need it - (preallocate + replay place allocations at recorded offsets; bootstrap covers NVSHMEM). -- **`deepep_adapter` mode on LOAD.** LOAD replaces the capture loop, so the adapter's - `_captured_deepep_mode` is never set; replay asserts on it. The hook calls - `deepep_adapter.capture(is_extend_in_batch=False)` after load. -- **FlashInfer pre-pass gated to FlashInfer.** The §5d pre-pass + `reuse_pre_pass_init` - shim handle FlashInfer's per-bs wrappers. fa3 (`FlashAttentionBackend`) uses a single - fixed `init_cuda_graph_state` workspace, so the shim is skipped (detected via absence of - `indices_updater_decode`); but fa3's per-bs `decode_cuda_graph_metadata[bs]` is still - populated post-load for the replay lookup. -- **C++: bind context on graph-build pool workers** (`CUDAGraphParallel.cpp`). EP graphs - carry `NODE_EVENT_RECORD/WAIT` nodes → `cuEventCreate` during on-demand prep, which runs - on `SimpleThreadPool` workers that never `cuCtxSetCurrent(main_ctx)`. Added that call - (mirrors the bg thread). Dense graphs never hit it (no event nodes), which is why it - surfaced only on sglang EP. - -See [`../../recipe/sglang/README.md`](../../recipe/sglang/README.md) for the EP serve -config and required kernel versions (deep_ep `9af0e0d`, sgl-deep-gemm ≥0.1.2, fa3). - -## Patch idiom - -All patches use the `wrap-and-call` idiom — short-circuit on `mode == NONE`, run foundry pre-work, call `orig`, run foundry post-work. The single exception is `capture()` on LOAD, which replaces the upstream method entirely (does not call `orig`). - -## Install order - -`install_hooks` calls the patch helpers in the listed order. Order doesn't matter for correctness here — every patch attaches to a different attribute — but spawn sites go last so the install-completion log line appears after every other patch has registered. +Environment inheritance is not Python-state inheritance. The external +four-file shim therefore also calls `apply_server_args` inside +`run_scheduler_process` and `run_data_parallel_controller_process`; see +[`direct-edits.md`](direct-edits.md). diff --git a/docs/sglang/memory-consistency.md b/docs/sglang/memory-consistency.md index d1e21d02..017783c3 100644 --- a/docs/sglang/memory-consistency.md +++ b/docs/sglang/memory-consistency.md @@ -1,123 +1,130 @@ -# SGLang Memory-Consistency Post-Mortem +# SGLang memory-consistency invariants -The hardest correctness problem in the SGLang integration was making LOAD produce the same output as SAVE. Cursors at every checkpoint matched, foundry's replay reported success, the graph launched cleanly — and the model emitted garbage tokens. +Foundry can successfully launch a restored graph while still producing wrong +tokens if a graph-referenced buffer occupies a different VMM address on LOAD. +The integration therefore treats allocation parity as a correctness property, +not only a performance property. -This doc walks every divergence we hit and the fix that landed. The common cause: anything that allocates on one path but not the other shifts the VMM cursor or changes the address an `_int_workspace_buffer` lands at, and the captured graph kernels read addresses fixed at SAVE time. +## What Foundry validates -## The general invariant +LOAD validates: -**SAVE and LOAD must walk an identical VMM cursor trajectory.** Foundry validates `start_base_addr_X` for each captured graph; the validation only catches asymmetries up to the first divergence. Captured-kernel argument addresses (e.g. wrapper `_int_workspace_buffer.data_ptr()`) are **not** validated — they just point at whatever VMM offset SAVE recorded. If LOAD's allocator lands a different buffer at the same offset, kernels read it; if no allocation lands at the offset, kernels read the preallocated zero-initialized memory. +- warmup-state schema 1; +- archived and installed SGLang version `0.5.15.post1`; +- archived and current decode backend `full`; +- archived and current prefill backend `disabled`; +- graph-index schema 1 and exact entry fields; +- contiguous capture order beginning at zero; +- every indexed graph file; +- the number of native load results; +- the final graph replay allocation trajectory through saved graph metadata. -## Bug 1 — `CUDAGraph::capture_begin/end` toggling the VMM region +These checks reject known incompatible inputs, but deterministic preparation is +still required to keep graph argument addresses correct. -### Symptom +## 1. Keep the VMM region active -`replay_hook_events_from_json` errored on iter 2 with `Memory offset mismatch during replay`: current cursor 20 MB ahead of saved `start_base_addr_2`. +The deterministic allocation region remains active through metadata +preparation and the complete capture/load sequence. Per-graph recording +windows may start and stop, but they must not toggle the allocation region. +Otherwise allocations between shapes bypass the VMM on one path and shift +every later graph address. -### Root cause +## 2. Remove SAVE-only backend warmup allocations -`foundry/csrc/CUDAGraph.cpp` used to bracket every capture with `resume_allocation_region()` / `stop_allocation_region()`. So on SAVE: +Upstream `FullCudaGraphBackend.capture_one` executes `forward_fn` twice before +capturing it. Those forwards allocate and release activations, changing the +CUDA caching allocator's free-segment layout. LOAD does not run those forwards. -- iter 0 enters with region ON. bs=24 init allocates inside VMM; cursor advances. -- iter 0 `capture_end` → region OFF. -- iter 1 init runs with region OFF. Allocations bypass VMM, no cursor advance. -- iter 1 `capture_begin` → region ON; `start_base_addr_1` recorded. -- … +Foundry replaces this method on SAVE and executes the real forward only once, +inside the Foundry capture context. This keeps SAVE's allocator history aligned +with LOAD while still persisting every allocation made by the captured +forward. -LOAD never calls `_capture_graph`, so nothing flipped the region off. LOAD's per-bs inits all advanced the cursor — past the saved `start_base_addr` values from iter 1 onward. +`BaseRunner.warmup` is also suppressed in SAVE and LOAD, and the external shim +disables FlashInfer autotuning. -### Fix +## 3. Recreate attention metadata in one order -Removed both calls in `CUDAGraph::capture_begin/end`. Only the hook recording window (`start_hook_record` / `end_hook_record`) brackets each capture. Region toggling is now each integration's responsibility, and both SGLang and vLLM keep the region ON across the whole capture loop. +SGLang v0.5.15.post1 separates attention metadata preparation: -(vLLM initially tried a Python-side `resume/stop_allocation_region` bracket around each capture to suppress the `final_alloc_offset` inflation seen on Qwen3-30B-A3B EP2 — ~194 GB observed vs ~80 GB actual. That was wrong: any tensor a between-capture allocation produced that was later referenced by a captured kernel would have ended up outside the deterministic region, so LOAD's preallocation couldn't recreate it at the right address. The actual root cause was the `empty_cache()` inside `foundry/graph.py:CUDAGraph.__enter__`, which dropped torch caching-allocator segments and forced subsequent inter-capture allocations to take a fresh `cuMemAlloc_v2` path. Removing it brought `final_alloc_offset` back down to ~80 GB. The vLLM integration now also wraps the whole capture loop in `torch.cuda.use_mem_pool(graph_pool)` so SAVE-2 ↔ LOAD caching state stays symmetric for the per-graph `finish_one_graph_load` replay path.) - -## Bug 2 — Pre-capture warmup forwards on SAVE only - -### Symptom - -After bug 1: iter 2 still drifted 20 MB. - -### Root cause - -`CudaGraphRunner.capture_one_batch_size` runs: - -```python -for _ in range(2): - self.device_module.synchronize() - self.model_runner.tp_group.barrier() - run_once() +```text +out-of-graph: host planning and allocation-bearing metadata +in-graph: static GPU work recorded in the graph ``` -before `_capture_graph`. These forwards allocate activation tensors and immediately free them. Torch's caching allocator keeps the freed segments as cached free blocks. Later per-bs init allocations may hit or miss cache depending on segment fragmentation history — and the cache-hit pattern is not reproducible across SAVE and LOAD. - -LOAD never runs these warmups → different cache state → different cache hit/miss decisions → cursor drift. - -### Fix - -In `patched_capture_one_batch_size` wrap `forward` with a 3-call counter; first two invocations return `None`, the third (inside `_capture_graph`) runs the real forward. The `synchronize` and `barrier` calls in the loop are harmless. JIT and autotune that the warmups normally trigger now run inside the captured graph and are recorded as alloc events — foundry replays them verbatim on LOAD. - -## Bug 3 — Per-bs init outside graph capture - -### Symptom - -After bugs 1 and 2: cursors matched everywhere, no foundry-side error, but the model produced garbage tokens. - -### Root cause - -`init_forward_metadata_capture_cuda_graph(bs, ...)` is called by `capture_one_batch_size` **before** `_capture_graph`, so its allocations are not recorded in the captured graph's allocator events. Each call creates a new `BatchDecodeWithPagedKVCacheWrapper`; the wrapper's constructor allocates `_int_workspace_buffer` via `torch.empty(...)`. The captured graph kernels reference `_int_workspace_buffer.data_ptr()` directly. +For FlashInfer, Foundry calls `capture_prepare(size)` and the out-of-graph +method for every size in `reversed(capture_bs)` before SAVE capture and before +LOAD graph hydration. This gives both paths the same allocation order. -For LOAD's wrapper to land at the same VMM address as SAVE's, the cursor at each iter's init time must match. With per-iter init alone, that depends on caching-allocator behavior, which is non-deterministic even after bug 2. +During SAVE's normal inner orchestration, the wrapper selects the pre-created +metadata and runs the latest planner in non-capture mode. It does not create a +second wrapper. The in-graph method remains at the start of `forward_fn`, so +its GPU operations are recorded and replayed normally. -### Fix (two parts) +This preserves both allocation addresses and the release's latest planning +semantics, including the prepared metadata's output-location binding. -**Both sides**: pre-pass init for every bs in `reversed(capture_bs)` order before the capture / load loop. `initialize_all_attention_metadata` walks the bs list and calls `initialize_attention_metadata_for_bs` for each. Same call sequence on both sides → same cursor trajectory → same wrapper addresses. +## 4. Load the complete indexed graph set together -**SAVE only**: the upstream `capture_one_batch_size` still calls the inner `init_forward_metadata_capture_cuda_graph(bs, ...)` per iter. Replaced it with `reuse_pre_pass_init` which: +Foundry's `graph_manifest.json` groups equal graph topologies. One member is a +template with a built executable; other members share that executable and +carry per-graph node updates. -- Looks up the existing wrapper from `decode_cuda_graph_metadata[bs]`. -- Re-runs `indices_updater_decode.update(...)` with the same buffer slices — idempotent write to the same `_int_workspace_buffer`. -- Sets `attn_backend.forward_metadata = DecodeMetadata(wrappers)` for the current iter. +LOAD must submit every indexed graph path in one +`start_graph_builds(paths)` call and finish them with one +`finish_graph_loads(pending)` call. Splitting the set prevents on-demand +members from seeing their template. -No allocation. Captured graph references the pre-pass wrapper's address; LOAD's identical pre-pass produces a wrapper at the same address. +`sglang_graph_index.json`, sorted by `capture_index`, supplies that path list. +The index also supplies the complete `ShapeKey(size, stream_idx, +variant_label)`. Graph filenames intentionally carry no shape identity. -(Two earlier approaches didn't work: relying on torch's caching allocator to reuse a freed segment after popping the dict entry, and forcing GC. The caching allocator doesn't deterministically reuse segments when request sizes don't exactly match cached blocks. Explicit reuse via dict lookup is the only reliable path.) +## 5. Mirror the pool profiler's cache drain -The `attn_backend.forward_metadata = None` between the pre-pass and the patch install drops the last bs's `DecodeMetadata` reference — defensive, not strictly required by `reuse_pre_pass_init`. +SAVE's normal memory-pool resolution drains the CUDA caching allocator while +measuring available memory. LOAD skips the profiler and directly applies the +saved pool config, so it must mirror that side effect: -## Bug 4 — Per-graph `start_graph_builds` broke template/on-demand linking - -### Symptom - -After bugs 1-3: cursors matched perfectly, no foundry replay error — but the runtime crashed at first inference with `Called CUDAGraph::replay without a preceding successful capture or load`. - -### Root cause - -`graph_manifest.json` groups captured graphs by topology. One graph per group is built as a **template** (full cuGraph + cuGraphExec instantiated); the rest are **on-demand** graphs that link to the template's `shared_exec` and patch in per-graph node updates at link time. - -An earlier (interleaved) LOAD called `start_graph_builds([single_path])` once per graph. Every call had 1 input → 1 topology group → graph was its own template. But the manifest marks graphs 2-N as on-demand sharing graph 0's or graph 1's template. Those templates were loaded in earlier `start_graph_builds` calls and no longer in scope — so the on-demand graphs ended up with `shared_exec=null`. Runtime `CUDAGraph::replay` takes the on-demand path when `on_demand_data_` is set; with `shared_exec=null` it falls through to a `capture_ended_ || has_graph_exec_` check that errors. - -### Fix - -`load_all_graphs(self)` calls `start_graph_builds(all_paths)` exactly once, then `finish_graph_loads(pending)` once. All N graphs go through the manifest's template/on-demand linking in a single pass. `finish_graph_loads` then replays each graph's allocator events sequentially, walking the cursor through every `start_base_addr_X` in order. +```python +torch.cuda.empty_cache() +self._apply_memory_pool_config(self.memory_pool_config) +``` -## Bug 5 — `_resolve_memory_pool_config`'s hidden `empty_cache` +Without the drain, subsequent attention initialization can hit cached segments +on LOAD but request new VMM memory on SAVE, or the reverse. -### Symptom +## 6. Bootstrap DeepEP outside graph capture -After bugs 1-4: cursors at `after_setup_graph_ext`, `after_init_torch_dist`, `after_scratch_skip`, `before_init_memory_pool`, and `after_init_memory_pool` all matched between SAVE and LOAD. But `save_before_pre_init` (29698 MB on SAVE) was 20 MB ahead of LOAD's `before_preallocate` (29678 MB). +DeepEP constructs its singleton NVSHMEM buffer lazily, and DeepGEMM may compile +shape-specific kernels lazily. Creating either inside stream capture is +invalid. -### Root cause +For the supported DP-attention + DeepEP low-latency + FA3 configuration: -SAVE's `init_memory_pool` calls `_resolve_memory_pool_config` → `_profile_available_bytes` → `get_available_gpu_memory(...)` (in `sglang/srt/utils/common.py`), which calls `torch.cuda.empty_cache()`. `empty_cache` releases torch caching-allocator segments back to the driver via `cuMemFree_v2`, which the foundry hook unmaps. +- SAVE runs a dedicated eager orchestration pass to initialize lazy state; +- SAVE and LOAD bootstrap the DeepEP buffer before capture/load; +- LOAD initializes queued NVSHMEM modules before native graph builds; +- LOAD initializes the DeepEP adapter after backend hydration. -LOAD's `_patch_init_memory_pool` skipped upstream profiling entirely and went directly to `_apply_memory_pool_config` — no equivalent drain. Caching allocator on LOAD retained segments that SAVE released, so the subsequent attention-backend init (between `init_memory_pool` and `capture()`) took a different `cuMemAlloc` path on each side. 20 MB drift by the time `capture()` was reached. +The DeepEP token limit, chunked-prefill size, model, topology, and environment +must remain identical between SAVE and LOAD. -### Fix +## Diagnosing a mismatch -One line in `_patch_init_memory_pool`'s LOAD branch: +Compare SAVE and LOAD allocation-offset logs at: -```python -torch.cuda.empty_cache() -self._apply_memory_pool_config(self.memory_pool_config) +```text +after_setup_graph_ext +after_init_torch_dist +after_scratch_skip +before_init_memory_pool +after_init_memory_pool +before graph preallocation/preparation +after graph load ``` + +The post-load offset must equal SAVE's `final_alloc_offset`. If compatibility +validation or any offset differs, remove the archive and perform a fresh SAVE +with the exact package, flags, model, topology, and GPU environment. Do not +reuse an archive from an earlier integration. diff --git a/docs/sglang/memory-lifecycle.md b/docs/sglang/memory-lifecycle.md index 45b034e2..494b6a38 100644 --- a/docs/sglang/memory-lifecycle.md +++ b/docs/sglang/memory-lifecycle.md @@ -1,119 +1,154 @@ -# SGLang Memory & Pool Lifecycle +# SGLang memory and pool lifecycle -## The invariant +## Core invariant -**SAVE and LOAD must walk an identical VMM cursor trajectory** from `setup_graph_extension` through end-of-capture / end-of-load. Anything that allocates on one path but not the other shifts the cursor by some amount. The captured graph kernels reference VMM addresses fixed at SAVE time; if LOAD's allocations don't land at those exact offsets, the buffers the kernels read are unmapped or contain stale data — **silently**, with no foundry-side error. +SAVE and LOAD must allocate graph-referenced state at identical VMM offsets. +Captured kernels retain device addresses from SAVE; a different LOAD +allocation order can silently bind those addresses to the wrong buffers. -## TOML config schema +The integration controls both the allocation trajectory and the compatibility +contract for SGLang `0.5.15.post1`. + +## TOML schema ```toml -mode = "save" | "load" # required -base_addr = 0x600000000000 # required; start of VMM region -region_size = "256GB" # required; total VMM reservation -workspace_root = "foundry_archive_…" # required; on-disk workspace -scratch_space_size = "1024MB" # required; pre-deterministic scratch budget - -hook_library_path = "…/libcuda_hook.so" # optional; auto-discovered -nvshmem_host_path = "…/libnvshmem_host.so" # optional; for EP parity +mode = "save" # or "load" +base_addr = 0x600000000000 +region_size = "256GB" +workspace_root = "foundry_archive" +scratch_space_size = "1024MB" + +# Optional overrides; normally auto-detected. +# hook_library_path = ".../libcuda_hook.so" +# nvshmem_host_path = ".../libnvshmem_host.so.3" ``` -- `base_addr` must be high enough not to collide with model weights or PyTorch's default heap. `0x600000000000` is a safe pick on x86_64. -- `region_size` must exceed `final_alloc_offset` for the model + KV pool + graphs. 256 GiB is generous for a 14 B model. -- `scratch_space_size` should be larger than the sum of NCCL workspace, cuBLAS workspace, and any other one-time allocations that aren't deterministic across runs. -- `workspace_root` is shared across ranks; per-rank artifacts go in `workspace_root/rank_{N}/`. Delete `workspace_root` before any SAVE. +- `base_addr` and `region_size` define the deterministic VMM reservation. +- `scratch_space_size` separates nondeterministic startup allocations from + deterministic runtime allocations. +- `workspace_root` contains shared warmup state and one graph directory per + rank. +- Delete `workspace_root` before every SAVE. ## Lifecycle ```mermaid flowchart TD - INIT["ModelRunner.__init__
(upstream stashes self.dp_rank)"] - TORCH["ModelRunner.init_torch_distributed"] - TORCH_PRE["setup_graph_extension
(reserve VMM region, eager cuBLAS,
LOAD: load cached fatbins)"] - TORCH_ORIG["<upstream init_torch_distributed>
NCCL warmup + comm bring-up"] - TORCH_POST["skip_to_scratch_boundary
cursor → scratch_space_size"] - LOAD_MODEL["ModelRunner.load_model
(passthrough today)"] - POOL["ModelRunnerKVCacheMixin.init_memory_pool"] - POOL_SAVE["SAVE: orig + save resolved
MemoryPoolConfig to warmup_state.json"] - POOL_LOAD["LOAD: load MemoryPoolConfig
+ empty_cache()
+ _apply_memory_pool_config"] - KWARM["ModelRunner.kernel_warmup
(no-op on SAVE/LOAD)"] - GRAPHS["ModelRunner.init_device_graphs
→ CudaGraphRunner(self).capture()"] - SAVE_CAP["SAVE: pre-pass init
+ reuse_pre_pass_init shim
+ capture loop
+ save_graph_manifest
+ pack_fatbins
+ capture_final_alloc_offset"] - LOAD_CAP["LOAD: preallocate_for_load_mode
+ pre-pass init
+ load_all_graphs"] - - INIT --> TORCH - TORCH --> TORCH_PRE --> TORCH_ORIG --> TORCH_POST - TORCH_POST --> LOAD_MODEL --> POOL - POOL --> POOL_SAVE - POOL --> POOL_LOAD - POOL_SAVE --> KWARM - POOL_LOAD --> KWARM - KWARM --> GRAPHS - GRAPHS --> SAVE_CAP - GRAPHS --> LOAD_CAP + SHIM["resolve full decode / disabled prefill"] + DIST["init_torch_distributed wrapper"] + VMM["bind rank device
reserve VMM
load packed modules on LOAD"] + SCRATCH["upstream distributed init
cursor to scratch boundary"] + MODEL["load model"] + POOL["init memory pool"] + POOLS["SAVE: resolve and record pool"] + POOLL["LOAD: validate archive
empty cache and apply saved pool"] + WARM["BaseRunner.warmup suppressed"] + DECODE["DecodeCudaGraphRunner.capture"] + SAVEC["SAVE: metadata prep
capture and index graphs"] + LOADC["LOAD: pre-map span
metadata prep and hydrate backend"] + + SHIM --> DIST --> VMM --> SCRATCH --> MODEL --> POOL + POOL --> POOLS --> WARM + POOL --> POOLL --> WARM + WARM --> DECODE + DECODE --> SAVEC + DECODE --> LOADC ``` -## Allocation buckets +## Allocation regions -### Bucket A — pre-deterministic scratch +### Startup scratch -CUDA context, cuBLAS handle, NCCL warmup, distributed bring-up. These run **inside** the VMM region but are released by `skip_to_scratch_boundary` — the cursor is forced to `cfg.scratch_space_size` regardless of what landed below it. Any non-determinism in this region is invisible to the rest of the lifecycle. +CUDA context creation, cuBLAS, NCCL/communicator setup, and related one-time +work happen before the scratch boundary. After upstream distributed +initialization, Foundry sets the cursor to `scratch_space_size` regardless of +the exact amount consumed below it. -### Bucket B — deterministic runtime state +For multi-GPU operation, the hook sets `self.gpu_id` as current before VMM +reservation. Otherwise rank-local allocations can bind to GPU 0. -Model weights, KV pools, attention workspace buffers, FlashInfer wrapper `_int_workspace_buffer`s, captured-graph alloc events. These must allocate in identical order on SAVE and LOAD. Their final cursor position is the `final_alloc_offset` persisted to `warmup_state.json`. +### Deterministic runtime state -### Bucket C — forbidden divergence +After the scratch boundary, model weights, KV pools, attention workspaces, +captured graph allocations, and output buffers must follow the same order. +SAVE records the terminal cursor as `final_alloc_offset`. -Anything that runs on one path but not the other. The fixes we landed identify and align the three known cases: +LOAD calls `preallocate_region(final_alloc_offset - current_offset)` before +native graph loading. This maps physical memory but does not advance the +logical cursor; replayed allocation events advance it in SAVE order. -1. The two pre-capture warmup forwards in `capture_one_batch_size` — skipped on SAVE so they don't pollute the caching allocator with freed activations LOAD can't reproduce. (See doc 06.) -2. The per-iter inner `init_forward_metadata_capture_cuda_graph(bs)` call — replaced on SAVE with `reuse_pre_pass_init` so it doesn't re-allocate the wrappers the pre-pass already built. (See doc 03 / doc 05.) -3. `_resolve_memory_pool_config` calls `get_available_gpu_memory(empty_cache=True)` on SAVE; LOAD's `_patch_init_memory_pool` mirrors it with an explicit `torch.cuda.empty_cache()` before `_apply_memory_pool_config`. (See below.) +### State that must stay eager -## The `_resolve_memory_pool_config` mirror +DeepEP/NVSHMEM buffer construction and DeepGEMM lazy initialization cannot +start inside stream capture. SAVE's dedicated eager bootstrap performs that +work before persisted graph capture. LOAD initializes its adapter and queued +NVSHMEM module state before replay. -`init_memory_pool` on SAVE calls `_resolve_memory_pool_config` → `_profile_available_bytes` → `get_available_gpu_memory(...)` which calls `torch.cuda.empty_cache()` (`sglang/srt/utils/common.py`). `empty_cache` releases torch caching-allocator segments back to the driver via `cuMemFree_v2`, which the foundry hook unmaps. +## Memory-pool parity -LOAD skips upstream `init_memory_pool` and goes straight to `_apply_memory_pool_config`, so it had no equivalent drain. Caching-allocator segments retained on LOAD changed which torch.empty calls hit cache vs. miss during the subsequent attention-backend init — pushing the cursor 20 MB behind SAVE by the time `capture()` was reached. +SAVE resolves SGLang's `MemoryPoolConfig` through normal profiling and records +the complete dataclass dictionary in `warmup_state.json`. -The fix is one line: +LOAD validates archive compatibility before changing pool state, then: ```python -torch.cuda.empty_cache() # mirror SAVE's _resolve_memory_pool_config side effect +self.memory_pool_config = MemoryPoolConfig(**state.memory_pool_config) +torch.cuda.empty_cache() self._apply_memory_pool_config(self.memory_pool_config) ``` -## Why saving raw memory numbers isn't enough +The cache drain mirrors a side effect of SAVE's memory profiling. Omitting it +changes whether later allocations reuse cached segments or advance the VMM +cursor. + +Persisting a raw free-memory count would be insufficient because SGLang also +applies model constraints, page alignment, request limits, and topology when +resolving the final pool. + +## Attention metadata parity -The pool config is computed from a profile and a configurator: +The release separates attention preparation into: -1. `_profile_available_bytes(pre_model_load_memory)` -2. `MemoryPoolConfigurator.calculate_pool_sizes(...)` -3. `_apply_token_constraints(...)` (user caps, page alignment, PP sync) -4. `_resolve_max_num_reqs(...)` -5. `_apply_memory_pool_config(...)` +- `init_forward_metadata_out_graph`, for host planning and allocation-bearing + per-iteration state; +- `init_forward_metadata_in_graph`, for static GPU operations recorded into + the CUDA graph. -Persisting the resolved `MemoryPoolConfig` (via `dataclasses.asdict`) and re-applying it on LOAD is necessary and sufficient. Persisting raw free-memory numbers would re-run the configurator on LOAD with stale inputs. +For FlashInfer, Foundry runs the out-of-graph allocation pass for every decode +size in reverse capture order on both SAVE and LOAD. During SAVE capture, the +inner planner reuses the prepared object rather than allocating a second one. +The in-graph method remains inside `forward_fn`. -`WarmupState.memory_pool_config` is exactly this dict. +FA3 uses its initialized graph state without the FlashInfer reuse wrapper. -## VMM region setup +## Graph-pool ownership -`setup_graph_extension` (in `runtime.py`): +`FullCudaGraphBackend` owns: -- Computes the per-rank workspace path (`{workspace_root}/rank_{compute_workspace_rank(...)}`) -- On SAVE: removes the rank workspace if it exists, then creates it fresh -- On LOAD: - - `cge.set_skip_fatbin_processing(True)` - - `cge.load_cuda_modules_and_libraries(workspace_dir)` — restores the fatbins SAVE wrote, into device code memory -- `cge.set_allocation_region(cfg.base_addr, parse_size(cfg.region_size))` — reserves a VMM address range; the cursor starts at `base_addr` -- `torch._C._cuda_getCurrentBlasHandle()` — eager cuBLAS handle so the workspace it wants lands in scratch instead of post-scratch territory -- creates the `CUDAGraphExtensionState` singleton +- `_graphs`: graph objects keyed by complete `ShapeKey`; +- `_outputs`: reconstructed `LogitsProcessorOutput` objects; +- `_pool`: the shared graph memory-pool handle. -After upstream `init_torch_distributed` returns, `skip_to_scratch_boundary` forces the cursor to `cfg.scratch_space_size` (default 1 GiB). Allocations below that line are scratch and don't need to be deterministic. +LOAD registers `_pool` with SGLang's process-wide graph-pool and PyNCCL pool +registries before restoring graphs, then hydrates the backend maps from the +authoritative graph index. -## Final watermark +## Persisted compatibility state -`capture_final_alloc_offset` runs after the SAVE-side capture loop completes (after `save_graph_manifest` and `pack_fatbins`). It writes `final_alloc_offset` to both `rank_{N}/final_alloc_offset.json` and the shared `warmup_state.json`. +Schema-1 `warmup_state.json` records: + +```text +schema_version +sglang_version +decode_cuda_graph_backend +prefill_cuda_graph_backend +cuda_version +gpu_name +gpu_total_memory +memory_pool_config +final_alloc_offset +``` -`preallocate_for_load_mode` uses this to call `cge.preallocate_region(final - current)` so the entire deterministic range is mapped to physical memory in one shot. The cursor is **not** advanced — preallocate just pre-maps; the cursor advances naturally as cuMemAllocs land within the preallocated range (fast-path: pointer bump, no driver calls). +LOAD requires archive schema 1, SGLang `0.5.15.post1`, decode `full`, prefill +`disabled`, matching installed package metadata, and matching current resolved +backends. A mismatch requires a fresh SAVE; no archive migration is attempted. diff --git a/docs/sglang/overview.md b/docs/sglang/overview.md index f8ea7cce..829c735f 100644 --- a/docs/sglang/overview.md +++ b/docs/sglang/overview.md @@ -1,112 +1,111 @@ -# SGLang Integration Overview - -Foundry persists SGLang's `CudaGraphRunner` graphs to disk on SAVE and restores them on LOAD, skipping graph capture, kernel warmup, and the per-batch-size attention metadata setup costs. - -Tested on single-GPU Qwen3-1.7B / 4B / 14B and **data-parallel (DP=2)** Qwen3-1.7B with the FlashInfer attention backend. - -## Parallelism - -| Mode | Status | Notes | -|---|:---:|---| -| Single GPU | ✅ | Qwen3-1.7B / 4B / 14B | -| Data parallel (DP) | ✅ | One full replica per rank; validated DP=2. Requires the per-rank device binding (below) and `NCCL_CUMEM_ENABLE=0` / `NCCL_NVLS_ENABLE=0`. | -| Tensor parallel (TP) | 🚧 | Deterministic NCCL memory layout is under construction. | -| Expert parallel (DeepEP) | ✅ | Validated EP=2 on Qwen3-30B-A3B-FP8 (SAVE/SAVE2/LOAD/query); restored decode graphs match baseline throughput. See **Expert parallel** below. | - -**Expert parallel (DeepEP).** EP runs DP-attention (each rank its own attention — no -NCCL all-reduce) + DeepEP for the MoE all-to-all (NVSHMEM, foundry-compatible). The -serve script is `recipe/sglang/serve_qwen3-30ba3bfp8_ep.sh - [--save|--load]` with: `--enable-dp-attention --moe-a2a-backend deepep ---deepep-mode low_latency --moe-runner-backend deep_gemm --attention-backend fa3 ---disable-custom-all-reduce`. Required kernel builds in the env: `deep_ep` at sglang's -pinned commit (`9af0e0d`, not vLLM's), `sgl-deep-gemm>=0.1.2` (0.1.0 lacks -`m_grouped_bf16_gemm_nt_masked`), `flash-attn-3`. `fa3` is required because the -flashinfer ragged-prefill path has an off-by-one (`q.shape != qo_indptr`) under this -config. DeepEP low-latency caps dispatch at -`SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK` (default 128); raise it (+ chunk -prefill) for larger batches and keep it identical across SAVE/LOAD. Foundry-specific -EP handling is in [`hooks.md`](hooks.md): a DeepEP buffer pre-capture bootstrap, a -SAVE-only warmup pass (triggers DeepGEMM JIT + buffer creation outside the capture -stream), `deepep_adapter` mode init on LOAD, the FlashInfer pre-pass gated off for -fa3, and a C++ fix binding the CUDA context on the graph-build pool workers. - -**Per-rank device binding (DP/TP/EP).** Foundry's `set_allocation_region` binds the -VMM region to the CUDA device current at call time. Upstream sets the device -*inside* `init_torch_distributed`, which the integration wraps and front-runs, so -the hook explicitly calls `set_device(self.gpu_id)` before reserving the region — -otherwise rank > 0 reserves on `cuda:0` and faults. See [`hooks.md`](hooks.md) §1. -The DP serve script lives at `recipe/sglang/serve_qwen3-1.7b_dp.sh` -(` [--save|--load]`); pick GPUs with `CUDA_VISIBLE_DEVICES`. - -## How to use - -### 1. Install foundry +# SGLang integration overview + +Foundry supports exactly SGLang `0.5.15.post1`, upstream commit `0b3bb0c`. +It persists the release's full decode CUDA graphs and restores them in a fresh +process. Prefill CUDA graphs are disabled. + +## Supported configurations + +| Configuration | Attention/MoE backend | Status | +|---|---|:---:| +| Single GPU | FlashInfer | Supported | +| Regular data parallelism | One full replica per rank, FlashInfer | Supported | +| DP-attention expert parallelism | DeepEP low-latency and FA3 | Supported | + +Tensor parallelism without DP-attention, pipeline parallelism, speculative +decoding, LoRA/PDMux, dLLM, hidden-state capture, and memory-saver graphs are +rejected at startup. Foundry does not claim prefill-graph support. + +## Upstream execution model + +The supported SGLang release separates orchestration from graph storage: + +- `BaseRunner.warmup` owns run-once kernel warmup and autotuning. +- `DecodeCudaGraphRunner.capture` prepares decode shapes and orchestrates + capture. +- `FullCudaGraphBackend.capture_one` owns one full-model graph capture. +- `ShapeKey(size, stream_idx, variant_label)` is the complete graph identity. +- `FullCudaGraphBackend._graphs`, `_outputs`, and `_pool` own restored graph + state. + +Attention metadata is also split. Per-iteration host work and dynamic planning +run in `init_forward_metadata_out_graph`; graph-recordable static GPU work runs +in `init_forward_metadata_in_graph`. + +Foundry patches these latest seams rather than rebuilding SGLang's runner. + +## SAVE lifecycle + +1. The shim resolves graph settings to decode `full` and prefill `disabled`, + then installs Foundry hooks. +2. Before distributed initialization, Foundry binds the rank's CUDA device, + reserves the deterministic VMM region, restores its scratch boundary, and + initializes the rank workspace. +3. SGLang allocates model weights, KV pools, and attention state inside that + region. Foundry records the resolved `MemoryPoolConfig` and exact package + and graph compatibility fields in `warmup_state.json`. +4. Foundry suppresses `BaseRunner.warmup`. +5. For FlashInfer, Foundry prepares every decode shape in reverse capture + order. The allocation-bearing out-of-graph metadata pass runs once per + shape; the normal inner path reuses that object while retaining current + planning behavior. +6. The replacement for `FullCudaGraphBackend.capture_one` skips its two eager + forwards, captures one real forward with Foundry, and records the complete + `ShapeKey`. +7. SAVE writes schema-1 `sglang_graph_index.json`, the Foundry topology + manifest, packed device code, and the final VMM allocation watermark. + +For DeepEP, SAVE first runs one eager capture-orchestration pass with native +graph capture neutralized. This initializes DeepGEMM and other lazy state +outside stream capture. Foundry also bootstraps the singleton DeepEP/NVSHMEM +buffer before the persisted capture. + +## LOAD lifecycle + +1. A fresh process installs hooks in the parent, DP controller when present, + and every scheduler child. +2. Before applying saved pool sizing, LOAD validates warmup schema 1, archived + and installed SGLang version `0.5.15.post1`, archived decode/prefill + backends, and the current resolved backends. +3. Foundry restores packed device code, recreates model/KV state at + deterministic VMM offsets, applies the saved `MemoryPoolConfig`, and + pre-maps memory through the final SAVE watermark. +4. The same reverse-order attention metadata preparation recreates + allocation-bearing objects. +5. LOAD reads schema-1 `sglang_graph_index.json` in `capture_index` order and + submits all indexed graph paths through one + `start_graph_builds`/`finish_graph_loads` pair. This preserves the + topology manifest's template and on-demand relationships. +6. Each index record becomes a real + `ShapeKey(size, stream_idx, variant_label)`. Foundry fills backend + `_graphs` and `_outputs`, restores `_pool`, and initializes the DeepEP + adapter before serving. + +The graph index, not filenames, is authoritative for phase, capture order, +shape identity, and output kind. + +## Archive compatibility + +This integration is latest-only. Archives from earlier SGLang versions or +earlier Foundry formats are not migrated. Delete the workspace and perform a +fresh SAVE: ```bash -pushd foundry && pip install -e . --no-build-isolation && popd -``` - -An editable install (`pip install -e .`) is the supported path — the recipe serve scripts set no `PYTHONPATH` and rely on `foundry` being importable. To run straight from a source checkout, export `PYTHONPATH=…/foundry/python:…/sglang/python` yourself. - -### 2. Write a TOML config - -```toml -# recipe/sglang/foundry_save.toml -mode = "save" -base_addr = 0x600000000000 -region_size = "256GB" -workspace_root = "foundry_archive" -scratch_space_size = "1024MB" -``` - -The matching `foundry_load.toml` just changes `mode = "load"`. See [`memory-lifecycle.md`](memory-lifecycle.md) for what each field controls. - -### 3. Run SAVE, then LOAD - -```bash -# SAVE rm -rf foundry_archive -bash recipe/sglang/serve_qwen3-mini.sh --save -# Wait for "Application startup complete", then SIGTERM the server. - -# LOAD -bash recipe/sglang/serve_qwen3-mini.sh --load - -# Query -curl -s http://0.0.0.0:12000/v1/completions \ - -H 'Content-Type: application/json' \ - -d '{"model":"Qwen/Qwen3-1.7B","prompt":"The capital of France is", - "max_tokens":12,"temperature":0}' -# → "Paris. The capital of the United States is Washington, D" ``` -A single SAVE pass is enough — SGLang doesn't run a profile-forward at startup, so there is no non-determinism that requires a second pass. - -## What the integration does - -SAVE: - -1. `setup_graph_extension(...)` reserves a VMM region at `base_addr` and creates the per-rank workspace. -2. Distributed init / NCCL warmup runs in scratch space; the cursor is then forced to `scratch_space_size`. -3. Model weights, KV pool, and FlashInfer workspace buffers allocate inside the VMM region at byte-deterministic offsets. -4. `kernel_warmup` is a no-op. -5. `CudaGraphRunner.capture` runs a pre-pass that pre-allocates every per-bs FlashInfer wrapper, then enters the upstream capture loop with an idempotent inner-init shim (`reuse_pre_pass_init`) and a wrapper on `forward` that suppresses the two pre-capture warmup forwards. -6. Each captured graph is written to disk; a manifest groups topologically equivalent graphs. -7. The final VMM cursor is recorded as `final_alloc_offset`. - -LOAD: +SGLang must be installed as a distribution whose metadata reports exactly +`0.5.15.post1`. `pip install -e sglang/python` is supported; a checkout made +importable only through a source path is not. -1. `setup_graph_extension(...)` restores the VMM region and replays captured fatbins into device code memory. -2. Distributed init runs as usual; the cursor advances to the same `scratch_space_size`. -3. Model weights and KV pool re-allocate at the same deterministic offsets. `init_memory_pool` reuses the saved `MemoryPoolConfig` (and calls `torch.cuda.empty_cache()` to mirror SAVE's `_resolve_memory_pool_config` side effect). -4. `CudaGraphRunner.capture` is replaced with: preallocate the entire deterministic range up to `final_alloc_offset`; run the same pre-pass init; call `start_graph_builds(all_paths) + finish_graph_loads(pending)` exactly once. All N graphs are loaded in one shot so the manifest's template/on-demand linking works. -5. `self.graphs` / `self.output_buffers` are populated from `state.loaded_graphs`; the rest of SGLang's serving path runs unchanged. +## Run it -## Doc set +The recipe includes baseline, SAVE, fresh-process LOAD, and query commands for +all supported configurations: -- [`overview.md`](overview.md) — this file -- [`direct-edits.md`](direct-edits.md) — every line that changed in `sglang/` -- [`hooks.md`](hooks.md) — every monkey-patch foundry installs and what it does on SAVE / LOAD -- [`memory-lifecycle.md`](memory-lifecycle.md) — VMM region setup, the SAVE↔LOAD allocation parity contract, and the `final_alloc_offset` watermark -- [`save-load-workflow.md`](save-load-workflow.md) — serve scripts, TOML schema, expected logs, validation checks -- [`memory-consistency.md`](memory-consistency.md) — the five known divergences that caused silent LOAD failures and how each was fixed +- [`../../recipe/sglang/README.md`](../../recipe/sglang/README.md) +- [`direct-edits.md`](direct-edits.md) for the external four-file shim +- [`hooks.md`](hooks.md) for every runtime patch +- [`memory-lifecycle.md`](memory-lifecycle.md) for deterministic allocation +- [`memory-consistency.md`](memory-consistency.md) for parity invariants +- [`save-load-workflow.md`](save-load-workflow.md) for archive and log checks diff --git a/docs/sglang/save-load-workflow.md b/docs/sglang/save-load-workflow.md index 78be4b49..c09439f0 100644 --- a/docs/sglang/save-load-workflow.md +++ b/docs/sglang/save-load-workflow.md @@ -1,35 +1,37 @@ -# SGLang Save / Load Workflow +# SGLang SAVE/LOAD workflow -## Serve scripts +This workflow targets SGLang `0.5.15.post1` only. Install the SGLang checkout +as a package, apply the external four-file shim, and delete every older archive +before SAVE. -Key knobs (1.7B example): +## Recipe flags -```bash -MODEL_NAME="Qwen/Qwen3-1.7B" -HOST="0.0.0.0" -PORT=12000 -MEM_FRACTION_STATIC=0.6 - -# No LD_PRELOAD / PYTHONPATH here: foundry + the sglang fork are pip-installed, and -# foundry's setup_ld_preload_env auto-detects libcuda_hook.so and LD_PRELOADs it -# into every worker at spawn time. (Running from a source checkout instead? Export -# PYTHONPATH=.../foundry/python:.../sglang/python yourself.) +The single-GPU recipe expands to: +```bash sglang serve \ - --model-path "$MODEL_NAME" \ - --trust-remote-code \ - --host "$HOST" --port "$PORT" \ - --tp-size 1 \ - --mem-fraction-static "$MEM_FRACTION_STATIC" \ - --disable-radix-cache \ - --attention-backend flashinfer \ - --cuda-graph-max-bs 512 \ - --foundry-graph-extension-config-path "$FOUNDRY_TOML" + --model-path Qwen/Qwen3-1.7B \ + --trust-remote-code \ + --host 0.0.0.0 \ + --port 12000 \ + --tp-size 1 \ + --mem-fraction-static 0.6 \ + --disable-radix-cache \ + --attention-backend flashinfer \ + --cuda-graph-max-bs-decode 512 \ + --disable-prefill-cuda-graph \ + --foundry-graph-extension-config-path "$FOUNDRY_TOML" ``` -`--cuda-graph-max-bs 512` is the closest analogue to vLLM's `--max-num-seqs 512` — it drives `capture_bs` to span a similar range of decode batch sizes (52 batch sizes from 1 → 512). +`--cuda-graph-max-bs-decode` is the release's phase-specific CLI field. +`--disable-prefill-cuda-graph` ensures that only decode graphs are created. +The recipes use both flags in baseline, SAVE, and LOAD because they are part of +the supported configuration. + +The Foundry option is present only for SAVE/LOAD. It is generated from the +external shim's annotated `ServerArgs` field. -## TOML configs +## TOML files `recipe/sglang/foundry_save.toml`: @@ -41,106 +43,141 @@ workspace_root = "foundry_archive" scratch_space_size = "1024MB" ``` -`foundry_load.toml` is identical except `mode = "load"`. See [`memory-lifecycle.md`](memory-lifecycle.md) for the field semantics. +`foundry_load.toml` uses the same values with `mode = "load"`. -## Run sequence +## Verify package metadata ```bash -cd /path/to/foundry/recipe/sglang +python -c 'from importlib.metadata import version; assert version("sglang") == "0.5.15.post1"' +``` -# 1. Clean any prior workspace -rm -rf foundry_archive +An editable `pip install -e sglang/python` satisfies this requirement. A +source-only checkout does not. SAVE refuses to create an archive when package +metadata is missing or differs from `0.5.15.post1`. + +## Baseline to fresh-process LOAD -# 2. SAVE +Run each server as a separate process: + +```bash +cd recipe/sglang + +# Baseline +bash serve_qwen3-mini.sh +# Wait for startup, query if desired, then stop it. + +# SAVE +rm -rf foundry_archive bash serve_qwen3-mini.sh --save -# wait for "Application startup complete", then Ctrl-C / SIGTERM +# Wait for "Application startup complete", then stop it. -# 3. LOAD +# LOAD in a fresh process bash serve_qwen3-mini.sh --load -# wait for "Application startup complete" -# 4. Query -curl -s http://0.0.0.0:12000/v1/completions \ +# Query from another shell +curl -s http://127.0.0.1:12000/v1/completions \ -H 'Content-Type: application/json' \ - -d '{"model":"Qwen/Qwen3-1.7B","prompt":"The capital of France is", - "max_tokens":12,"temperature":0}' -``` - -For Qwen3-1.7B, expected output: - -``` -" Paris. The capital of the United States is Washington, D" + -d '{"model":"Qwen/Qwen3-1.7B","prompt":"The capital of France is","max_tokens":12,"temperature":0}' ``` -A single SAVE pass is sufficient. SGLang does not run a profile-forward at startup (only `_profile_available_bytes` samples free GPU memory), so there is no allocation non-determinism that a second SAVE pass would need to mask. +Use the DP and DeepEP commands in +[`../../recipe/sglang/README.md`](../../recipe/sglang/README.md) for the other +supported configurations. -## Per-rank workspace layout +## Archive files -``` +```text foundry_archive/ - warmup_state.json # shared; MemoryPoolConfig + final_alloc_offset - rank_0/ - graph_{0..N-1}_FULL_t{bs}_r{bs}_UX_pcN.json # one per captured graph - graph_{0..N-1}_FULL_t{bs}_r{bs}_UX_pcN.cugraph # binary cuGraph blob - graph_manifest.json # topology groups for template + on-demand linking - final_alloc_offset.json # per-rank VMM watermark - fatbin_image_packed.img # packed kernel fatbins - fatbin_entrypoint_packed.txt # fatbin entry-point index +├── warmup_state.json +└── rank_0/ + ├── sglang_graph_index.json + ├── graph_000000_sglang.json + ├── graph_000000_sglang.cugraph + ├── ... + ├── graph_manifest.json + ├── final_alloc_offset.json + ├── fatbin_image_packed.img + └── fatbin_entrypoint_packed.txt ``` -Filename scheme (defined in `graph_ops._graph_filename`): +Each additional DP/EP rank has its own `rank_/`. + +`warmup_state.json` includes compatibility fields: +```json +{ + "schema_version": 1, + "sglang_version": "0.5.15.post1", + "decode_cuda_graph_backend": "full", + "prefill_cuda_graph_backend": "disabled", + "memory_pool_config": {}, + "final_alloc_offset": 0 +} ``` -graph_{state.capture_index}_FULL_t{bs}_r{bs}_UX_pcN.json + +The actual pool dictionary and final offset are populated by SAVE. + +`sglang_graph_index.json` is authoritative: + +```json +{ + "schema_version": 1, + "graphs": [ + { + "capture_index": 0, + "filename": "graph_000000_sglang.json", + "output_kind": "next_token_logits", + "phase": "decode", + "shape_key": { + "size": 512, + "stream_idx": null, + "variant_label": null + } + } + ] +} ``` -`state.capture_index` increments per `save_graph` call so files sort in SAVE-time order. `_GRAPH_FILENAME_RE` in `graph_ops.py` parses them on LOAD. +LOAD validates exact fields and uses `capture_index`, not lexical filename +order. `graph_manifest.json` separately records native topology groups, +templates, and on-demand members. -## Expected logs +## Expected checkpoints -SAVE (success): +SAVE logs include: -``` -[Foundry] SGLang hooks installing: mode=save workspace=foundry_archive_qwen_1.7b -[Foundry] SGLang hooks installed -Foundry SGLang integration activated from .../save_qwen_1.7b.toml -[Foundry] SGLang rank=0 workspace_dir=foundry_archive_qwen_1.7b/rank_0 -[Foundry] SGLang graph extension setup completed in 0.x s -[Foundry] SGLang skipped allocator to scratch boundary 1073741824 -[Foundry] SGLang reused saved memory pool config ← only on LOAD -[Foundry] SGLang kernel_warmup skipped in save mode -… -[Foundry] Saved SGLang CUDA graph graph_0_FULL_t512_r512_UX_pcN.json key=512 -… -[Foundry] Saved SGLang CUDA graph graph_51_FULL_t1_r1_UX_pcN.json key=1 -[foundry] Saved graph_manifest.json: 9 topology groups (9 templates, 43 on-demand, 43 stripped) -[Foundry] Saved SGLang warmup state to foundry_archive_qwen_1.7b/warmup_state.json -[Foundry] SGLang final_alloc_offset=22785556480 -… -INFO: Application startup complete. +```text +[Foundry] SGLang hooks installing: mode=save ... +[Foundry] SGLang skipped allocator to scratch boundary ... +[Foundry] SGLang runner warmup skipped in save mode +[Foundry] Saved SGLang CUDA graph graph_000000_sglang.json shape_key=ShapeKey(...) +[Foundry] Saved SGLang warmup state to .../warmup_state.json +[Foundry] SGLang final_alloc_offset=... +INFO: Application startup complete. ``` -LOAD (success): +LOAD logs include: -``` -[Foundry] SGLang hooks installing: mode=load workspace=foundry_archive_qwen_1.7b -[Foundry] SGLang hooks installed -[Foundry] SGLang graph extension setup completed in 0.x s +```text +[Foundry] SGLang hooks installing: mode=load ... [Foundry] SGLang reused saved memory pool config -[Foundry] SGLang kernel_warmup skipped in load mode -[Foundry] SGLang alloc_offset[before_preallocate]=… (… MB) -[Foundry] SGLang alloc_offset[after_preallocate]=… (… MB) -[Foundry] SGLang alloc_offset[after_pre_init]=… (… MB) -[CGE] Using graph_manifest.json (9 topology groups) -[CGE] Phase 1: 52 graphs parsed in 0.x ms, 9 topologies, 4 threads, 52 binary + 0 json -[CGE BUILD] Template 0 (...): N nodes, done in X.X ms -… -[CGE] Phase 2: 9 templates + 43 on-demand = 52 graphs built in xx.x ms -[CGE] finish_graph_loads: 52 graphs, xx.x ms -[Foundry] Loaded 52 SGLang graphs in 0.0x s -[Foundry] SGLang alloc_offset[after_load_all_graphs]=22785556480 (… MB) -… -INFO: Application startup complete. +[Foundry] SGLang runner warmup skipped in load mode +[CGE] Using graph_manifest.json ... +[CGE] finish_graph_loads: ... +[Foundry] Loaded ... SGLang graphs in ...s +INFO: Application startup complete. ``` -The `after_load_all_graphs` value **must** equal SAVE's `final_alloc_offset`. If it doesn't, see [`memory-consistency.md`](memory-consistency.md). +Compatibility validation occurs before the saved pool config is applied. +Missing index files, legacy schemas, version mismatches, backend mismatches, or +unsupported output/shape records stop LOAD and request a fresh SAVE. + +## Migration rule + +There is no archive migration path. When changing SGLang, Foundry, graph +flags, model, topology, or relevant kernel packages: + +```bash +rm -rf foundry_archive +bash serve_qwen3-mini.sh --save +``` diff --git a/recipe/sglang/README.md b/recipe/sglang/README.md index af4ce528..12f3c3ec 100644 --- a/recipe/sglang/README.md +++ b/recipe/sglang/README.md @@ -1,154 +1,192 @@ -# Foundry recipe — SGLang +# Foundry recipe — SGLang v0.5.15.post1 -End-to-end serve scripts for SAVE / LOAD of CUDA graphs through the foundry SGLang -integration. All scripts in this directory share the same pair of foundry TOML files -(`foundry_save.toml` / `foundry_load.toml`) — pick a script for your model + parallelism, -run `--save`, then `--load`, then query. The integration code is in -[`../../python/foundry/integration/sglang/`](../../python/foundry/integration/sglang/); -design notes are under [`../../docs/sglang/`](../../docs/sglang/). +These recipes persist and restore full decode CUDA graphs for exactly SGLang +`0.5.15.post1` (upstream commit `0b3bb0c`). Prefill graphs are always disabled. -## Files in this directory +| Script | Validated configuration | +|---|---| +| `serve_qwen3-mini.sh` | Single GPU, Qwen3-1.7B, FlashInfer | +| `serve_qwen3-1.7b_dp.sh` | Regular data parallelism, one full replica per GPU, FlashInfer | +| `serve_qwen3-30ba3bfp8_ep.sh` | DP-attention, DeepEP low-latency, FA3, Qwen3-30B-A3B-FP8 | -``` -recipe/sglang/ -├── README.md # this file -├── foundry_save.toml # shared SAVE config (workspace_root = "foundry_archive") -├── foundry_load.toml # shared LOAD config (same workspace_root) -├── serve_qwen3-mini.sh # Qwen3-1.7B single GPU -├── serve_qwen3-1.7b_dp.sh # Qwen3-1.7B data parallel -└── serve_qwen3-30ba3bfp8_ep.sh # Qwen3-30B-A3B FP8 expert parallel (DeepEP) -``` +Tensor parallelism without DP-attention, pipeline parallelism, speculative +decoding, LoRA/PDMux, dLLM, hidden-state capture, and memory-saver graphs are +not supported. + +## Requirements + +Use a dedicated environment with: + +- Python 3.12; +- CUDA 13; +- PyTorch 2.11.0 with cu130; +- `flashinfer_python[cu13]==0.6.12` and `flashinfer_cubin==0.6.12`; +- `sglang-kernel==0.4.4`; +- `sgl-deep-gemm==0.1.4`; +- the SGLang FA3 package, `kernels-community/sgl-flash-attn3` revision 1, + declared by SGLang through the `kernels` package. + +The expert-parallel recipe additionally requires DeepEP commit `9af0e0d` and +the NVSHMEM library supplied with the cu13 PyTorch environment. + +## Check out and install SGLang -Every script accepts the same trailing `--save` / `--load` flag. Scripts that scale -across GPUs take the parallel-size as the first positional argument: +Keep Foundry and the Foundry SGLang fork as sibling checkouts. Start the +external fork from the exact upstream release: ```bash -bash serve_qwen3-mini.sh [--save|--load] -bash serve_qwen3-1.7b_dp.sh [--save|--load] -bash serve_qwen3-30ba3bfp8_ep.sh [--save|--load] +git clone https://github.com/foundry-org/sglang.git +cd sglang +git checkout -b foundry-v0.5.15.post1 \ + 0b3bb0cbe31873994c9f989fddfe2f87ca839fdd ``` -A single SAVE pass is enough — SGLang has no startup profile-forward, so there is no -non-determinism that requires a two-pass save (unlike the vLLM recipe). +Apply the four-file shim in +[`../../docs/sglang/direct-edits.md`](../../docs/sglang/direct-edits.md), then +install both projects as packages: -Because the two TOMLs are shared (single `workspace_root = "foundry_archive"`), one -archive is written per host; run a fresh `rm -rf foundry_archive` whenever you change -model or topology before SAVE. - -| Mode | Script | Model | Notes | -|---|---|---|---| -| Single GPU | `serve_qwen3-mini.sh` | Qwen3-1.7B | FlashInfer backend | -| Data parallel | `serve_qwen3-1.7b_dp.sh` | Qwen3-1.7B | one full replica/rank; `NCCL_CUMEM_ENABLE=0`/`NCCL_NVLS_ENABLE=0` | -| Expert parallel | `serve_qwen3-30ba3bfp8_ep.sh` | Qwen3-30B-A3B-FP8 | DP-attention + DeepEP; fa3 backend | +```bash +conda create -p venv python=3.12 +conda activate venv +conda install -c conda-forge boost-cpp boost + +pip install torch==2.11.0 torchvision==0.26.0 torchaudio==2.11.0 \ + --index-url https://download.pytorch.org/whl/cu130 +pip install -e sglang/python \ + --extra-index-url https://download.pytorch.org/whl/cu130 +pushd foundry +CC=gcc CXX=g++ pip install -e . --no-build-isolation +popd + +python -c 'from importlib.metadata import version; assert version("sglang") == "0.5.15.post1"' +``` -## Installation +The editable SGLang install is supported because it creates distribution +metadata. A source checkout exposed only through `PYTHONPATH` is not supported: +SAVE and LOAD both require installed package metadata to report exactly +`0.5.15.post1`. -The recipes assume `foundry` and the SGLang fork are **pip-installed** (editable is -fine) so both import without any `PYTHONPATH`, and foundry's spawn-site patch -auto-detects `libcuda_hook.so` from its install — the scripts set no `LD_PRELOAD` -themselves. The standard workspace layout has `foundry/` (this repo) and `sglang/` -(the foundry-org SGLang fork) as siblings: +The tagged release pins the kernel stack above in `python/pyproject.toml`. +Confirm the resolved environment before capture: +```bash +python -m pip show sglang torch flashinfer-python flashinfer-cubin \ + sglang-kernel sgl-deep-gemm kernels ``` -/ -├── foundry/ # this repo -│ ├── python/foundry/ # `pip install -e .` builds libcuda_hook.so here -│ ├── recipe/sglang/ # <-- you are here -│ └── ... -└── sglang/ # foundry-org/sglang fork (with direct edits applied) + +For DeepEP on H100: + +```bash +git clone https://github.com/deepseek-ai/DeepEP.git +cd DeepEP +git checkout 9af0e0d0e74f3577af1979c9b9e1ac2cad0104ee +TORCH_CUDA_ARCH_LIST="9.0" python setup.py install +cd .. ``` -Use a dedicated env, kept separate from the vLLM env so kernel pins don't clash: +Foundry auto-detects `libcuda_hook.so` and +`site-packages/nvidia/nvshmem/lib/libnvshmem_host.so.3` before worker spawn. +Use `hook_library_path` or `nvshmem_host_path` in both TOMLs only when +auto-detection does not fit the environment. + +## Archive compatibility + +Archives created by earlier SGLang integrations are incompatible. Delete the +entire workspace and perform a new SAVE after upgrading: ```bash -conda create -p venv/ python=3.12 -conda activate venv/ -conda install -c conda-forge boost-cpp boost # foundry C++ deps - -# in-tree sglang fork, editable -pip install -e sglang/python --extra-index-url https://download.pytorch.org/whl/cu130 -# foundry -pushd foundry && pip install -e . --no-build-isolation && popd +rm -rf foundry_archive ``` -`libcuda_hook.so` finds boost via a baked rpath; if it can't, add the conda lib dir to -`LD_LIBRARY_PATH` (`export LD_LIBRARY_PATH=$CONDA_PREFIX/lib:$LD_LIBRARY_PATH`). +LOAD validates the archive schema, archived SGLang version, installed SGLang +version, and resolved graph configuration before applying the saved memory-pool +configuration. Both phases must resolve to decode `full` and prefill +`disabled`. -## Run (single GPU / DP) +## Baseline, SAVE, LOAD, query + +Run from this directory. Each script accepts an optional trailing `--save` or +`--load`; omitting it starts a baseline server. ```bash -# single GPU +cd foundry/recipe/sglang + +# 1. Baseline. Wait for startup, optionally query it, then stop the process. +bash serve_qwen3-mini.sh + +# 2. Fresh SAVE. Remove every older archive first. rm -rf foundry_archive -bash serve_qwen3-mini.sh --save # wait for "Application startup complete", then SIGTERM -bash serve_qwen3-mini.sh --load # leave running +bash serve_qwen3-mini.sh --save +# Wait for "Application startup complete", then stop the process. + +# 3. Fresh-process LOAD. Do not reuse the SAVE process. +bash serve_qwen3-mini.sh --load -# query (separate shell) -curl -s http://0.0.0.0:12000/v1/completions -H 'Content-Type: application/json' \ +# 4. Query LOAD from another shell. +curl -s http://127.0.0.1:12000/v1/completions \ + -H 'Content-Type: application/json' \ -d '{"model":"Qwen/Qwen3-1.7B","prompt":"The capital of France is","max_tokens":12,"temperature":0}' +``` + +Regular DP: -# data parallel (pick GPUs with CUDA_VISIBLE_DEVICES) +```bash rm -rf foundry_archive CUDA_VISIBLE_DEVICES=0,1 bash serve_qwen3-1.7b_dp.sh 2 --save +# Stop SAVE after startup. CUDA_VISIBLE_DEVICES=0,1 bash serve_qwen3-1.7b_dp.sh 2 --load ``` -## Run (expert parallel / DeepEP) - -EP needs three kernel packages — all SGLang-native, no vLLM involved: - -- **NVSHMEM** — already in the env. cu13 `torch` pulls the `nvidia-nvshmem-cuXX` - wheel as a dependency (`libnvshmem_host.so.3` under `site-packages/nvidia/nvshmem/lib/`). - Foundry auto-detects it from the wheel (just like `libcuda_hook.so`) and the - spawn-site patches preload it into each worker — no manual path, no TOML field. -- **DeepEP** @ `9af0e0d` — SGLang's pin. Build via SGLang's own installer - `sglang/scripts/ci/cuda/ci_install_deepep.sh` (it `git checkout`s exactly `9af0e0d` - and builds against the NVSHMEM wheel above). For a single node you can skip the - script's gdrcopy/RDMA apt steps and just build the wheel: - - ```bash - git clone https://github.com/deepseek-ai/DeepEP.git && cd DeepEP - git checkout 9af0e0d0e74f3577af1979c9b9e1ac2cad0104ee - TORCH_CUDA_ARCH_LIST="9.0" python setup.py install # Hopper; "9.0;10.0" for Blackwell - cd .. - ``` - -- **`sgl-deep-gemm >= 0.1.2`** (0.1.0 lacks `m_grouped_bf16_gemm_nt_masked`) and - **`flash-attn-3`** (the fa3 attention backend — flashinfer's ragged-prefill path - has an off-by-one under this config). +DP-attention with DeepEP: ```bash rm -rf foundry_archive CUDA_VISIBLE_DEVICES=0,1 bash serve_qwen3-30ba3bfp8_ep.sh 2 --save +# Stop SAVE after startup. CUDA_VISIBLE_DEVICES=0,1 bash serve_qwen3-30ba3bfp8_ep.sh 2 --load -curl -s http://0.0.0.0:12000/v1/completions -H 'Content-Type: application/json' \ +curl -s http://127.0.0.1:12000/v1/completions \ + -H 'Content-Type: application/json' \ -d '{"model":"Qwen/Qwen3-30B-A3B-FP8","prompt":"The capital of France is","max_tokens":12,"temperature":0}' ``` -The EP script sets `--enable-dp-attention --moe-a2a-backend deepep --deepep-mode -low_latency --moe-runner-backend deep_gemm --attention-backend fa3 ---disable-custom-all-reduce` and `SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK=256`. -DeepEP low-latency caps dispatch at that per-rank token count (and asserts -`(n+1)*2 <= NVSHMEM_QP_DEPTH`); keep it and `--chunked-prefill-size` identical between -SAVE and LOAD so the captured graphs match. +The multi-GPU scripts set `NCCL_CUMEM_ENABLE=0` and +`NCCL_NVLS_ENABLE=0` for Foundry runs. The DeepEP script also keeps +`SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK=256` and +`--chunked-prefill-size 256` identical on SAVE and LOAD. ## Archive layout -``` +```text foundry_archive/ -├── warmup_state.json # KV-block sizing + MemoryPoolConfig (rank 0) +├── warmup_state.json └── rank_/ - ├── graph_*.json + .cugraph # one pair per captured graph - ├── graph_manifest.json # topology groups + template assignments - ├── fatbin_image_packed.img # packed CUDA modules - └── final_alloc_offset.json # per-rank VMM watermark + ├── sglang_graph_index.json + ├── graph_000000_sglang.json + ├── graph_000000_sglang.cugraph + ├── graph_manifest.json + ├── fatbin_image_packed.img + ├── fatbin_entrypoint_packed.txt + └── final_alloc_offset.json ``` -For DP / EP each rank gets its own `rank_/`. +`warmup_state.json` schema 1 records `sglang_version`, +`decode_cuda_graph_backend`, `prefill_cuda_graph_backend`, +`memory_pool_config`, and `final_alloc_offset`, plus GPU/CUDA diagnostics. + +`sglang_graph_index.json` schema 1 is the authoritative graph list. Each entry +records capture order, phase, filename, output kind, and the complete +`ShapeKey(size, stream_idx, variant_label)`. Filenames are opaque; LOAD does +not derive graph identity from them. `graph_manifest.json` remains the Foundry +topology/template manifest used by the batched native loader. + +For DP and EP, each rank owns a separate `rank_/` directory. ## Troubleshooting -| Symptom | Likely cause | +| Symptom | Action | |---|---| -| `Reserved address … != requested base 0x600000000000` | VMM base collided with another allocation. Re-run; non-deterministic, the next run usually succeeds. | -| EP replay `illegal memory access` / `nvshmemx_cumodule_init not found` | `libnvshmem_host.so.3` not preloaded — foundry couldn't auto-detect the `nvidia-nvshmem` wheel. Confirm it's installed (`pip show nvidia-nvshmem-cu13`), or set `nvshmem_host_path` in both TOMLs. | -| `nvshmem_qp_depth >= (num_max_dispatch_tokens_per_rank + 1) * 2` | `SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK` too high for `NVSHMEM_QP_DEPTH`; lower it or raise the QP depth. | +| Package version is `unknown` or not `0.5.15.post1` | Install `sglang/python` with pip; do not rely on a source-path-only checkout. | +| Archive schema, version, or backend mismatch | Delete `foundry_archive` and perform a fresh SAVE with the exact environment and flags. | +| Reserved VMM address differs from `0x600000000000` | Another allocation owns the requested range; restart the process. | +| DeepEP reports missing NVSHMEM symbols | Verify the cu13 NVSHMEM wheel or set `nvshmem_host_path` in both TOMLs. | +| DeepEP token/QP-depth assertion fails | Lower `SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK` or raise `NVSHMEM_QP_DEPTH`, identically for SAVE and LOAD. | diff --git a/recipe/sglang/serve_qwen3-1.7b_dp.sh b/recipe/sglang/serve_qwen3-1.7b_dp.sh index c562a5ac..94d32149 100755 --- a/recipe/sglang/serve_qwen3-1.7b_dp.sh +++ b/recipe/sglang/serve_qwen3-1.7b_dp.sh @@ -37,7 +37,7 @@ fi # LD_PRELOAD of libcuda_hook.so is set by foundry's setup_ld_preload_env at # worker spawn time (path auto-detected; propagated to the DP controller and # every rank's scheduler child). Assumes foundry + the sglang fork are -# pip-installed (see README) so both import without PYTHONPATH. +# pip-installed (see README) so workers resolve both installed distributions. sglang serve \ --model-path "$MODEL_NAME" \ @@ -48,5 +48,6 @@ sglang serve \ --mem-fraction-static "$MEM_FRACTION_STATIC" \ --disable-radix-cache \ --attention-backend flashinfer \ - --cuda-graph-max-bs 512 \ + --cuda-graph-max-bs-decode 512 \ + --disable-prefill-cuda-graph \ "${FOUNDRY_ARGS[@]}" diff --git a/recipe/sglang/serve_qwen3-30ba3bfp8_ep.sh b/recipe/sglang/serve_qwen3-30ba3bfp8_ep.sh index a185fda3..6be24a67 100755 --- a/recipe/sglang/serve_qwen3-30ba3bfp8_ep.sh +++ b/recipe/sglang/serve_qwen3-30ba3bfp8_ep.sh @@ -2,8 +2,8 @@ # Qwen3-30B-A3B-FP8 (MoE), expert parallel via DeepEP low-latency + DP-attention. # Usage: CUDA_VISIBLE_DEVICES=0,1 bash serve_qwen3-30ba3bfp8_ep.sh [--save|--load] # -# Requires (see README §EP): deep_ep @ 9af0e0d, sgl-deep-gemm >= 0.1.2, flash-attn-3, -# and the nvshmem_host_path uncommented in foundry_{save,load}.toml. +# Requires (see README §EP): deep_ep @ 9af0e0d, sgl-deep-gemm 0.1.4, the +# SGLang FA3 package, and NVSHMEM (auto-detected or configured in the TOMLs). SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" @@ -66,5 +66,6 @@ sglang serve \ --disable-custom-all-reduce \ --chunked-prefill-size 256 \ --attention-backend fa3 \ - --cuda-graph-max-bs 128 \ + --cuda-graph-max-bs-decode 128 \ + --disable-prefill-cuda-graph \ "${FOUNDRY_ARGS[@]}" diff --git a/recipe/sglang/serve_qwen3-mini.sh b/recipe/sglang/serve_qwen3-mini.sh index 77449482..9e31aafa 100755 --- a/recipe/sglang/serve_qwen3-mini.sh +++ b/recipe/sglang/serve_qwen3-mini.sh @@ -30,7 +30,7 @@ fi # LD_PRELOAD of libcuda_hook.so is set by foundry's setup_ld_preload_env at # worker spawn time (path auto-detected from the foundry install). Baseline runs # don't need it preloaded by the shell. Assumes foundry + the sglang fork are -# pip-installed (see README) so both import without PYTHONPATH. +# pip-installed (see README) so workers resolve both installed distributions. sglang serve \ --model-path "$MODEL_NAME" \ @@ -40,5 +40,6 @@ sglang serve \ --mem-fraction-static "$MEM_FRACTION_STATIC" \ --disable-radix-cache \ --attention-backend flashinfer \ - --cuda-graph-max-bs 512 \ + --cuda-graph-max-bs-decode 512 \ + --disable-prefill-cuda-graph \ "${FOUNDRY_ARGS[@]}" diff --git a/tests/test_sglang_recipes.py b/tests/test_sglang_recipes.py new file mode 100644 index 00000000..b8856207 --- /dev/null +++ b/tests/test_sglang_recipes.py @@ -0,0 +1,24 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the Foundry project +"""Contracts for the supported SGLang v0.5.15.post1 serve recipes.""" + +import re +from pathlib import Path + +import pytest + +RECIPE_DIR = Path(__file__).parents[1] / "recipe" / "sglang" +SERVE_SCRIPTS = ( + "serve_qwen3-mini.sh", + "serve_qwen3-1.7b_dp.sh", + "serve_qwen3-30ba3bfp8_ep.sh", +) + + +@pytest.mark.parametrize("script_name", SERVE_SCRIPTS) +def test_recipe_uses_phase_specific_decode_graph_flags(script_name): + script = (RECIPE_DIR / script_name).read_text() + + assert "--cuda-graph-max-bs-decode" in script + assert "--disable-prefill-cuda-graph" in script + assert re.search(r"--cuda-graph-max-bs(?:\s|$)", script) is None From 820091a6b423df504c41eb9e84bb9b69ba36718d Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Thu, 23 Jul 2026 07:47:02 +0000 Subject: [PATCH 07/21] docs(sglang): fix editable install guidance Co-authored-by: Rahul Chalamala --- README.md | 20 ++++++--- docs/overview.md | 4 +- docs/sglang/direct-edits.md | 25 +++++++---- docs/sglang/hooks.md | 2 +- docs/sglang/memory-consistency.md | 3 +- docs/sglang/memory-lifecycle.md | 1 + docs/sglang/overview.md | 26 ++++++----- docs/sglang/save-load-workflow.md | 25 ++++++++--- recipe/sglang/README.md | 54 +++++++++++++++-------- recipe/sglang/serve_qwen3-1.7b_dp.sh | 4 +- recipe/sglang/serve_qwen3-30ba3bfp8_ep.sh | 2 +- recipe/sglang/serve_qwen3-mini.sh | 4 +- 12 files changed, 109 insertions(+), 61 deletions(-) diff --git a/README.md b/README.md index b50c9d75..d52b85b1 100644 --- a/README.md +++ b/README.md @@ -82,18 +82,24 @@ Foundry ships engine integrations under `foundry/python/foundry/integration/`. P | Engine | Single GPU | DP | TP | EP | |---|:---:|:---:|:---:|:---:| | vLLM | ✅ | ✅ | 🚧 | ✅ | -| SGLang `0.5.15.post1` | ✅ | ✅ | 🚧 | ✅ | +| SGLang `0.5.15.post1` | Targeted | Targeted | Unsupported | Targeted | | TensorRT-LLM | 🚧 | 🚧 | 🚧 | 🚧 | -✅ validated end-to-end (SAVE → LOAD → query)  ·  🚧 not yet +✅ validated end-to-end (SAVE → LOAD → query)  ·  Targeted = +supported pending H100 validation  ·  🚧 not yet The SGLang integration is latest-only for upstream `v0.5.15.post1` -(`0b3bb0c`). It supports full decode graphs with prefill graphs disabled for +(`0b3bb0c`). It targets full decode graphs with prefill graphs disabled for single-GPU FlashInfer, regular DP FlashInfer, and DP-attention with DeepEP -low-latency + FA3. Existing SGLang archives are incompatible and require a -fresh SAVE. See [`recipe/sglang/README.md`](recipe/sglang/README.md). - -The adapted vLLM / SGLang / TensorRT-LLM forks will be released alongside this repo at `foundry-org/vllm`, `foundry-org/sglang`, `foundry-org/TensorRT-LLM`. +low-latency + FA3; these modes are supported pending H100 validation. Existing +SGLang archives are incompatible and require a fresh SAVE. See +[`recipe/sglang/README.md`](recipe/sglang/README.md). + +For the current Foundry-only SGLang workflow, check out upstream +`sgl-project/sglang` at the tagged release and apply the documented four-file +patch. `foundry-org/sglang` should carry that same patch when its release branch +is published; the current instructions do not assume that branch already +exists. Adapted engine forks will be released alongside this repository. ### Performance diff --git a/docs/overview.md b/docs/overview.md index 9e65c7bf..a09fdf49 100644 --- a/docs/overview.md +++ b/docs/overview.md @@ -29,7 +29,7 @@ hooks and archive metadata are intentionally separate. | Aspect | vLLM | SGLang `0.5.15.post1` | |---|---|---| -| Supported graph scope | Configured vLLM piecewise/full integration modes | Full decode only; prefill disabled | +| Graph scope | Configured vLLM piecewise/full integration modes | Targeted: full decode only; prefill disabled | | Capture seam | `CUDAGraphWrapper.__call__` | `DecodeCudaGraphRunner.capture` orchestration and `FullCudaGraphBackend.capture_one` storage | | Run-once warmup seam | vLLM compile/kernel/sampler lifecycle hooks | `BaseRunner.warmup` | | Graph identity | vLLM piecewise graph metadata | `ShapeKey(size, stream_idx, variant_label)` | @@ -37,7 +37,7 @@ hooks and archive metadata are intentionally separate. | Attention preparation | vLLM-specific compile/capture path | Separate out-of-graph planning and in-graph GPU metadata work | | Profile behavior | Full profile forward; two-pass SAVE recipe | Pool-memory measurement only; one SAVE process | | Process model | Uniprocess, multiprocessing, or Ray executors | Multiprocessing `spawn`; explicit hook install in scheduler and DP-controller children | -| Validated parallel modes | Single GPU, DP, and documented EP configuration | Single GPU FlashInfer, regular DP FlashInfer, and DP-attention + DeepEP low-latency + FA3 | +| Parallel scope | Single GPU, DP, and documented EP configuration | Targeted: single GPU FlashInfer, regular DP FlashInfer, and DP-attention + DeepEP low-latency + FA3; H100 validation pending | For SGLang, TP without DP-attention, PP, speculative decoding, LoRA/PDMux, dLLM, hidden-state capture, memory-saver graphs, and prefill graphs are not diff --git a/docs/sglang/direct-edits.md b/docs/sglang/direct-edits.md index 86e8d1cf..35e19891 100644 --- a/docs/sglang/direct-edits.md +++ b/docs/sglang/direct-edits.md @@ -1,9 +1,13 @@ # SGLang v0.5.15.post1 direct edits -Foundry's runtime integration lives in this repository. The external -`foundry-org/sglang` checkout needs only the following four-file shim, applied -to upstream commit `0b3bb0cbe31873994c9f989fddfe2f87ca839fdd` -(`v0.5.15.post1`). +Foundry's runtime integration lives in this repository. For the current +Foundry-only workflow, clone upstream `sgl-project/sglang` at commit +`0b3bb0cbe31873994c9f989fddfe2f87ca839fdd` (`v0.5.15.post1`) and apply the +following four-file shim. + +`foundry-org/sglang` should carry the same patch once its release branch is +published. This guide does not imply that the current fork already contains +the patch. ## 1. Add `python/sglang/srt/foundry_shim.py` @@ -24,7 +28,7 @@ def apply_server_args(server_args) -> None: ``` The shim receives already resolved `cuda_graph_config` objects. It forces the -only supported graph shape: full decode and disabled prefill. It also disables +only accepted graph shape: full decode and disabled prefill. It also disables graph profiling and FlashInfer autotuning before calling Foundry's public `foundry.integration.sglang.install_hooks` entry point. @@ -142,10 +146,15 @@ After applying the four edits, install the checkout rather than exposing it as source only: ```bash -pip install -e sglang/python \ - --extra-index-url https://download.pytorch.org/whl/cu130 +SETUPTOOLS_SCM_PRETEND_VERSION=0.5.15.post1 \ + pip install -e sglang/python \ + --extra-index-url https://download.pytorch.org/whl/cu130 python -c 'from importlib.metadata import version; assert version("sglang") == "0.5.15.post1"' sglang serve --help | rg -- '--foundry-graph-extension-config-path' ``` -SAVE and LOAD reject missing or non-exact SGLang distribution metadata. +The tracked shim makes the tagged checkout dirty. Because SGLang uses +`setuptools_scm`, an unqualified editable install would derive a dirty/local +version. The pretend-version variable deterministically writes exact +`0.5.15.post1` distribution metadata at install time. SAVE and LOAD continue +to reject missing or non-exact metadata. diff --git a/docs/sglang/hooks.md b/docs/sglang/hooks.md index a4843105..e3a5f2ce 100644 --- a/docs/sglang/hooks.md +++ b/docs/sglang/hooks.md @@ -2,7 +2,7 @@ `foundry.integration.sglang.install_hooks(server_args)` is the public, idempotent entry point for SGLang `0.5.15.post1`. It loads the Foundry TOML, -validates the supported graph configuration, and installs the hooks below. +validates the targeted graph configuration, and installs the hooks below. ## Installation order diff --git a/docs/sglang/memory-consistency.md b/docs/sglang/memory-consistency.md index 017783c3..80961cef 100644 --- a/docs/sglang/memory-consistency.md +++ b/docs/sglang/memory-consistency.md @@ -100,7 +100,8 @@ DeepEP constructs its singleton NVSHMEM buffer lazily, and DeepGEMM may compile shape-specific kernels lazily. Creating either inside stream capture is invalid. -For the supported DP-attention + DeepEP low-latency + FA3 configuration: +For the targeted DP-attention + DeepEP low-latency + FA3 configuration, +supported pending H100 validation: - SAVE runs a dedicated eager orchestration pass to initialize lazy state; - SAVE and LOAD bootstrap the DeepEP buffer before capture/load; diff --git a/docs/sglang/memory-lifecycle.md b/docs/sglang/memory-lifecycle.md index 494b6a38..1a70d83b 100644 --- a/docs/sglang/memory-lifecycle.md +++ b/docs/sglang/memory-lifecycle.md @@ -142,6 +142,7 @@ schema_version sglang_version decode_cuda_graph_backend prefill_cuda_graph_backend +timestamp cuda_version gpu_name gpu_total_memory diff --git a/docs/sglang/overview.md b/docs/sglang/overview.md index 829c735f..54aebf32 100644 --- a/docs/sglang/overview.md +++ b/docs/sglang/overview.md @@ -1,16 +1,17 @@ # SGLang integration overview -Foundry supports exactly SGLang `0.5.15.post1`, upstream commit `0b3bb0c`. -It persists the release's full decode CUDA graphs and restores them in a fresh -process. Prefill CUDA graphs are disabled. +Foundry targets exactly SGLang `0.5.15.post1`, upstream commit `0b3bb0c`. +The integration is designed to persist the release's full decode CUDA graphs +and restore them in a fresh process. Prefill CUDA graphs are disabled. The +configurations below are supported pending H100 validation. -## Supported configurations +## Targeted configurations | Configuration | Attention/MoE backend | Status | |---|---|:---:| -| Single GPU | FlashInfer | Supported | -| Regular data parallelism | One full replica per rank, FlashInfer | Supported | -| DP-attention expert parallelism | DeepEP low-latency and FA3 | Supported | +| Single GPU | FlashInfer | Pending H100 validation | +| Regular data parallelism | One full replica per rank, FlashInfer | Pending H100 validation | +| DP-attention expert parallelism | DeepEP low-latency and FA3 | Pending H100 validation | Tensor parallelism without DP-attention, pipeline parallelism, speculative decoding, LoRA/PDMux, dLLM, hidden-state capture, and memory-saver graphs are @@ -18,7 +19,7 @@ rejected at startup. Foundry does not claim prefill-graph support. ## Upstream execution model -The supported SGLang release separates orchestration from graph storage: +The targeted SGLang release separates orchestration from graph storage: - `BaseRunner.warmup` owns run-once kernel warmup and autotuning. - `DecodeCudaGraphRunner.capture` prepares decode shapes and orchestrates @@ -95,13 +96,16 @@ rm -rf foundry_archive ``` SGLang must be installed as a distribution whose metadata reports exactly -`0.5.15.post1`. `pip install -e sglang/python` is supported; a checkout made -importable only through a source path is not. +`0.5.15.post1`. Apply the four-file tracked shim, then install with +`SETUPTOOLS_SCM_PRETEND_VERSION=0.5.15.post1 pip install -e sglang/python`. +The override is required because `setuptools_scm` otherwise derives a +dirty/local version from the patched tag. A checkout made importable only +through a source path is not supported. ## Run it The recipe includes baseline, SAVE, fresh-process LOAD, and query commands for -all supported configurations: +all targeted configurations: - [`../../recipe/sglang/README.md`](../../recipe/sglang/README.md) - [`direct-edits.md`](direct-edits.md) for the external four-file shim diff --git a/docs/sglang/save-load-workflow.md b/docs/sglang/save-load-workflow.md index c09439f0..30477413 100644 --- a/docs/sglang/save-load-workflow.md +++ b/docs/sglang/save-load-workflow.md @@ -25,8 +25,8 @@ sglang serve \ `--cuda-graph-max-bs-decode` is the release's phase-specific CLI field. `--disable-prefill-cuda-graph` ensures that only decode graphs are created. -The recipes use both flags in baseline, SAVE, and LOAD because they are part of -the supported configuration. +The recipes use both flags in baseline, SAVE, and LOAD because they define the +targeted configuration, which is supported pending H100 validation. The Foundry option is present only for SAVE/LOAD. It is generated from the external shim's annotated `ServerArgs` field. @@ -48,12 +48,18 @@ scratch_space_size = "1024MB" ## Verify package metadata ```bash +SETUPTOOLS_SCM_PRETEND_VERSION=0.5.15.post1 \ + pip install -e sglang/python \ + --extra-index-url https://download.pytorch.org/whl/cu130 python -c 'from importlib.metadata import version; assert version("sglang") == "0.5.15.post1"' ``` -An editable `pip install -e sglang/python` satisfies this requirement. A -source-only checkout does not. SAVE refuses to create an archive when package -metadata is missing or differs from `0.5.15.post1`. +The required tracked-file shim makes the tagged checkout dirty. SGLang uses +`setuptools_scm`, so the pretend-version variable is required to make an +editable install's distribution metadata deterministically equal +`0.5.15.post1`. A source-only checkout does not satisfy the gate. SAVE refuses +to create an archive when package metadata is missing or differs from that +exact version. ## Baseline to fresh-process LOAD @@ -82,7 +88,7 @@ curl -s http://127.0.0.1:12000/v1/completions \ Use the DP and DeepEP commands in [`../../recipe/sglang/README.md`](../../recipe/sglang/README.md) for the other -supported configurations. +targeted configurations. ## Archive files @@ -110,12 +116,17 @@ Each additional DP/EP rank has its own `rank_/`. "sglang_version": "0.5.15.post1", "decode_cuda_graph_backend": "full", "prefill_cuda_graph_backend": "disabled", + "timestamp": "2026-07-23T00:00:00", + "cuda_version": "13.0", + "gpu_name": "NVIDIA H100", + "gpu_total_memory": 0, "memory_pool_config": {}, "final_alloc_offset": 0 } ``` -The actual pool dictionary and final offset are populated by SAVE. +The timestamp, GPU values, pool dictionary, and final offset shown above are +illustrative; SAVE records the actual values. `sglang_graph_index.json` is authoritative: diff --git a/recipe/sglang/README.md b/recipe/sglang/README.md index 12f3c3ec..9e0b1875 100644 --- a/recipe/sglang/README.md +++ b/recipe/sglang/README.md @@ -1,9 +1,11 @@ # Foundry recipe — SGLang v0.5.15.post1 -These recipes persist and restore full decode CUDA graphs for exactly SGLang -`0.5.15.post1` (upstream commit `0b3bb0c`). Prefill graphs are always disabled. +These recipes target persistence and restoration of full decode CUDA graphs +for exactly SGLang `0.5.15.post1` (upstream commit `0b3bb0c`). Prefill graphs +are always disabled. The configurations below are supported pending H100 +validation. -| Script | Validated configuration | +| Script | Targeted configuration | |---|---| | `serve_qwen3-mini.sh` | Single GPU, Qwen3-1.7B, FlashInfer | | `serve_qwen3-1.7b_dp.sh` | Regular data parallelism, one full replica per GPU, FlashInfer | @@ -31,19 +33,20 @@ the NVSHMEM library supplied with the cu13 PyTorch environment. ## Check out and install SGLang -Keep Foundry and the Foundry SGLang fork as sibling checkouts. Start the -external fork from the exact upstream release: +Keep Foundry and SGLang as sibling checkouts. For this Foundry-only workflow, +clone the authoritative upstream tag: ```bash -git clone https://github.com/foundry-org/sglang.git -cd sglang -git checkout -b foundry-v0.5.15.post1 \ - 0b3bb0cbe31873994c9f989fddfe2f87ca839fdd +git clone --branch v0.5.15.post1 --depth 1 \ + https://github.com/sgl-project/sglang.git sglang +test "$(git -C sglang rev-parse HEAD)" = \ + "0b3bb0cbe31873994c9f989fddfe2f87ca839fdd" ``` Apply the four-file shim in [`../../docs/sglang/direct-edits.md`](../../docs/sglang/direct-edits.md), then -install both projects as packages: +install both projects as packages from the directory containing the sibling +checkouts: ```bash conda create -p venv python=3.12 @@ -52,8 +55,9 @@ conda install -c conda-forge boost-cpp boost pip install torch==2.11.0 torchvision==0.26.0 torchaudio==2.11.0 \ --index-url https://download.pytorch.org/whl/cu130 -pip install -e sglang/python \ - --extra-index-url https://download.pytorch.org/whl/cu130 +SETUPTOOLS_SCM_PRETEND_VERSION=0.5.15.post1 \ + pip install -e sglang/python \ + --extra-index-url https://download.pytorch.org/whl/cu130 pushd foundry CC=gcc CXX=g++ pip install -e . --no-build-isolation popd @@ -61,10 +65,21 @@ popd python -c 'from importlib.metadata import version; assert version("sglang") == "0.5.15.post1"' ``` -The editable SGLang install is supported because it creates distribution -metadata. A source checkout exposed only through `PYTHONPATH` is not supported: -SAVE and LOAD both require installed package metadata to report exactly -`0.5.15.post1`. +SGLang uses `setuptools_scm`. The required tracked-file shim makes the tagged +checkout dirty, so an unqualified editable install derives a dirty/local +version instead of exactly `0.5.15.post1`. +`SETUPTOOLS_SCM_PRETEND_VERSION=0.5.15.post1` deterministically embeds the +required distribution version during installation. The exact runtime metadata +gate remains in force. + +An editable SGLang install with that override is supported because it creates +exact distribution metadata. A source checkout exposed only through +`PYTHONPATH` is not supported: SAVE and LOAD both require installed package +metadata to report exactly `0.5.15.post1`. + +`foundry-org/sglang` should carry the same four-file patch once its +`v0.5.15.post1` release branch is published. These instructions do not assume +that branch exists or already contains the shim. The tagged release pins the kernel stack above in `python/pyproject.toml`. Confirm the resolved environment before capture: @@ -169,9 +184,10 @@ foundry_archive/ └── final_alloc_offset.json ``` -`warmup_state.json` schema 1 records `sglang_version`, -`decode_cuda_graph_backend`, `prefill_cuda_graph_backend`, -`memory_pool_config`, and `final_alloc_offset`, plus GPU/CUDA diagnostics. +`warmup_state.json` schema 1 records `schema_version`, `sglang_version`, +`decode_cuda_graph_backend`, `prefill_cuda_graph_backend`, `timestamp`, +`cuda_version`, `gpu_name`, `gpu_total_memory`, `memory_pool_config`, and +`final_alloc_offset`. `sglang_graph_index.json` schema 1 is the authoritative graph list. Each entry records capture order, phase, filename, output kind, and the complete diff --git a/recipe/sglang/serve_qwen3-1.7b_dp.sh b/recipe/sglang/serve_qwen3-1.7b_dp.sh index 94d32149..85e938fe 100755 --- a/recipe/sglang/serve_qwen3-1.7b_dp.sh +++ b/recipe/sglang/serve_qwen3-1.7b_dp.sh @@ -36,8 +36,8 @@ fi # LD_PRELOAD of libcuda_hook.so is set by foundry's setup_ld_preload_env at # worker spawn time (path auto-detected; propagated to the DP controller and -# every rank's scheduler child). Assumes foundry + the sglang fork are -# pip-installed (see README) so workers resolve both installed distributions. +# every rank's scheduler child). Assumes foundry + the patched upstream SGLang +# checkout are pip-installed (see README). sglang serve \ --model-path "$MODEL_NAME" \ diff --git a/recipe/sglang/serve_qwen3-30ba3bfp8_ep.sh b/recipe/sglang/serve_qwen3-30ba3bfp8_ep.sh index 6be24a67..c0ccca21 100755 --- a/recipe/sglang/serve_qwen3-30ba3bfp8_ep.sh +++ b/recipe/sglang/serve_qwen3-30ba3bfp8_ep.sh @@ -42,7 +42,7 @@ fi # (the hook from the foundry install, NVSHMEM from the nvidia-nvshmem wheel via # config._detect_nvshmem_host_path), so nothing is preloaded by the shell. Set # nvshmem_host_path in the TOML only to override the NVSHMEM auto-detection. -# Assumes foundry + the sglang fork are pip-installed (see README). +# Assumes foundry + the patched upstream SGLang checkout are pip-installed. # DeepEP low-latency caps tokens/rank at SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK # (default 128; (n+1)*2 must be <= NVSHMEM_QP_DEPTH). Raise to 256 ((256+1)*2=514 diff --git a/recipe/sglang/serve_qwen3-mini.sh b/recipe/sglang/serve_qwen3-mini.sh index 9e31aafa..5262bf83 100755 --- a/recipe/sglang/serve_qwen3-mini.sh +++ b/recipe/sglang/serve_qwen3-mini.sh @@ -29,8 +29,8 @@ fi # LD_PRELOAD of libcuda_hook.so is set by foundry's setup_ld_preload_env at # worker spawn time (path auto-detected from the foundry install). Baseline runs -# don't need it preloaded by the shell. Assumes foundry + the sglang fork are -# pip-installed (see README) so workers resolve both installed distributions. +# don't need it preloaded by the shell. Assumes foundry + the patched upstream +# SGLang checkout are pip-installed (see README). sglang serve \ --model-path "$MODEL_NAME" \ From 75972fc0a62351b1cff36057f97626c66c6b8d62 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Thu, 23 Jul 2026 11:00:41 +0000 Subject: [PATCH 08/21] fix(hook): preserve cursor for default VMM alignment Co-authored-by: Rahul Chalamala --- csrc/hook.cpp | 1 + tests/test_vmm_alloc.py | 40 ++++++++++++++++++++++++++++++++++++++++ 2 files changed, 41 insertions(+) diff --git a/csrc/hook.cpp b/csrc/hook.cpp index d552d08d..f9e7a209 100644 --- a/csrc/hook.cpp +++ b/csrc/hook.cpp @@ -267,6 +267,7 @@ static std::mutex hook_events_mutex; static CUdeviceptr hook_recording_start_base_addr{0}; static inline size_t align_to(size_t addr, size_t alignment) { + if (alignment == 0) return addr; return (addr + alignment - 1) & ~(alignment - 1); } diff --git a/tests/test_vmm_alloc.py b/tests/test_vmm_alloc.py index 5df909f0..4a2a46f6 100644 --- a/tests/test_vmm_alloc.py +++ b/tests/test_vmm_alloc.py @@ -1,5 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the Foundry project +import ctypes import os import subprocess import sys @@ -146,6 +147,36 @@ def _run_core_size_parsing(): print("[TEST] test_size_parsing PASSED") +def _run_core_zero_alignment_reserve(): + import foundry as fdry + + torch.cuda.init() + + base_addr = 0x7F0000000000 + region_size = 1024 * 1024 * 1024 + reserved_size = 64 * 1024 + reserve = ctypes.CDLL(None).cuMemAddressReserve + reserve.argtypes = [ + ctypes.POINTER(ctypes.c_uint64), + ctypes.c_size_t, + ctypes.c_size_t, + ctypes.c_uint64, + ctypes.c_ulonglong, + ] + reserve.restype = ctypes.c_int + + with fdry.allocation_region(base_addr, region_size): + ptr = ctypes.c_uint64() + offset_before = fdry.get_current_alloc_offset() + + result = reserve(ctypes.byref(ptr), reserved_size, 0, 0, 0) + + offset_after = fdry.get_current_alloc_offset() + assert result == 0 + assert base_addr <= ptr.value < base_addr + region_size + assert offset_before <= offset_after <= region_size + + def _spawn_with_preload(test_mode): so_path = _get_hook_so_path() env = os.environ.copy() @@ -176,6 +207,12 @@ def test_size_parsing(): _spawn_with_preload("size-parsing") +@pytest.mark.filterwarnings("ignore:TORCH_CUDA_ARCH_LIST is not set") +def test_zero_alignment_reserve_stays_in_allocation_region(): + """CUDA's default alignment must not reset the allocation cursor.""" + _spawn_with_preload("zero-alignment-reserve") + + if __name__ == "__main__": if "--core" in sys.argv: _run_core() @@ -183,7 +220,10 @@ def test_size_parsing(): _run_core_without_region() elif "--size-parsing" in sys.argv: _run_core_size_parsing() + elif "--zero-alignment-reserve" in sys.argv: + _run_core_zero_alignment_reserve() else: test_vmm_allocation() test_vmm_allocation_without_region() test_size_parsing() + test_zero_alignment_reserve_stays_in_allocation_region() From 2585c85de2c810d70fd3fe9c72a57b542ca56d55 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Thu, 23 Jul 2026 13:25:47 +0000 Subject: [PATCH 09/21] style(hook): format zero-alignment guard Co-authored-by: Rahul Chalamala --- csrc/hook.cpp | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/csrc/hook.cpp b/csrc/hook.cpp index f9e7a209..c3ede37e 100644 --- a/csrc/hook.cpp +++ b/csrc/hook.cpp @@ -267,7 +267,8 @@ static std::mutex hook_events_mutex; static CUdeviceptr hook_recording_start_base_addr{0}; static inline size_t align_to(size_t addr, size_t alignment) { - if (alignment == 0) return addr; + if (alignment == 0) + return addr; return (addr + alignment - 1) & ~(alignment - 1); } From 3a24fc3aa8fa8240dc726aa85d2fc9845a6d0a70 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Thu, 23 Jul 2026 13:25:50 +0000 Subject: [PATCH 10/21] fix(sglang): restore DeepEP allocation symmetry Co-authored-by: Rahul Chalamala --- .../foundry/integration/sglang/cuda_graph.py | 31 ++++- tests/test_sglang_cuda_graph.py | 118 ++++++++++++++++++ 2 files changed, 146 insertions(+), 3 deletions(-) diff --git a/python/foundry/integration/sglang/cuda_graph.py b/python/foundry/integration/sglang/cuda_graph.py index bb92659b..444fa5bd 100644 --- a/python/foundry/integration/sglang/cuda_graph.py +++ b/python/foundry/integration/sglang/cuda_graph.py @@ -292,11 +292,28 @@ def bootstrap_deepep_buffer(cuda_graph_runner: DecodeCudaGraphRunner) -> bool: def _deepep_eager_warmup() -> Iterator[None]: global _EAGER_DEEPEP_WARMUP previous = _EAGER_DEEPEP_WARMUP + compile_utils = _module("sglang.srt.layers.deep_gemm_wrapper.compile_utils") + original_compile = compile_utils._compile_deep_gemm_one_type_all + + @functools.wraps(original_compile) + def compile_outside_allocation_region(*args, **kwargs): + # DeepGEMM compiles all shapes only on the first local rank. Its warmup + # executor is explicitly deleted and empty_cache()'d by upstream, so + # only that executor is transient JIT state. The surrounding eager + # forward still owns persistent lazy allocations used by saved graphs. + graph_ops.cge.stop_allocation_region() + try: + return original_compile(*args, **kwargs) + finally: + graph_ops.cge.resume_allocation_region() + + compile_utils._compile_deep_gemm_one_type_all = compile_outside_allocation_region _EAGER_DEEPEP_WARMUP = True try: yield finally: _EAGER_DEEPEP_WARMUP = previous + compile_utils._compile_deep_gemm_one_type_all = original_compile def _is_flashinfer_backend(attn_backend: Any) -> bool: @@ -375,13 +392,16 @@ def reuse( def _save_capture(runner: DecodeCudaGraphRunner, original: Callable[..., Any]): if _uses_deepep(runner): start = time.perf_counter() + # The DeepEP/NVSHMEM buffer is a live graph dependency. Allocate it + # symmetrically inside Foundry's region before excluding transient JIT + # warmup allocations from the recorded cursor. + bootstrap_deepep_buffer(runner) with _deepep_eager_warmup(): original(runner) logger.info( "[Foundry] SGLang DeepEP eager lazy-init warmup completed in %.3fs", time.perf_counter() - start, ) - bootstrap_deepep_buffer(runner) if _is_flashinfer_backend(runner.attn_backend): prepared = _prepare_attention_metadata(runner) @@ -396,7 +416,10 @@ def _save_capture(runner: DecodeCudaGraphRunner, original: Callable[..., Any]): return result -def _load_capture(runner: DecodeCudaGraphRunner) -> None: +def _load_capture( + runner: DecodeCudaGraphRunner, + original: Callable[..., Any], +) -> None: if getattr(runner, "_foundry_graphs_loaded", False): raise RuntimeError( "Foundry LOAD cannot satisfy an SGLang CUDA-graph recapture request; " @@ -412,6 +435,8 @@ def _load_capture(runner: DecodeCudaGraphRunner) -> None: if _uses_deepep(runner): bootstrap_deepep_buffer(runner) + with _deepep_eager_warmup(): + original(runner) graph_ops.preallocate_for_load_mode() _prepare_attention_metadata(runner) @@ -446,7 +471,7 @@ def patched(self): if mode == CUDAGraphExtensionMode.SAVE: return _save_capture(self, original) if mode == CUDAGraphExtensionMode.LOAD: - return _load_capture(self) + return _load_capture(self, original) raise RuntimeError(f"Unsupported Foundry SGLang mode: {mode!r}") return patched diff --git a/tests/test_sglang_cuda_graph.py b/tests/test_sglang_cuda_graph.py index 5233771a..3d81e812 100644 --- a/tests/test_sglang_cuda_graph.py +++ b/tests/test_sglang_cuda_graph.py @@ -336,3 +336,121 @@ def test_pool_registration_uses_latest_split_runner_api(monkeypatch): assert registered.global_pool is pool assert registered.pynccl_pool is pool assert "sglang.srt.model_executor.cuda_graph_runner" not in sys.modules + + +def test_deepep_jit_warmup_does_not_advance_recorded_allocator(monkeypatch): + cuda_graph, _mode, _registered = _load_cuda_graph(monkeypatch) + events = [] + allocator = SimpleNamespace(enabled=True) + compile_utils = _module("sglang.srt.layers.deep_gemm_wrapper.compile_utils") + + def stop_allocation_region(): + allocator.enabled = False + events.append("stop") + + def resume_allocation_region(): + allocator.enabled = True + events.append("resume") + + def compile_all(): + events.append(("compile", allocator.enabled)) + + compile_utils._compile_deep_gemm_one_type_all = compile_all + monkeypatch.setitem( + sys.modules, + "sglang.srt.layers.deep_gemm_wrapper.compile_utils", + compile_utils, + ) + monkeypatch.setattr( + cuda_graph.graph_ops, + "cge", + SimpleNamespace( + stop_allocation_region=stop_allocation_region, + resume_allocation_region=resume_allocation_region, + ), + raising=False, + ) + monkeypatch.setattr(cuda_graph, "_uses_deepep", lambda _runner: True) + monkeypatch.setattr( + cuda_graph, + "bootstrap_deepep_buffer", + lambda _runner: events.append("bootstrap"), + ) + for name in ("save_graph_manifest", "pack_fatbins", "capture_final_alloc_offset"): + monkeypatch.setattr(cuda_graph.graph_ops, name, lambda: None, raising=False) + + runner = SimpleNamespace(attn_backend=object()) + + def capture(_runner): + eager = cuda_graph._EAGER_DEEPEP_WARMUP + events.append(("capture_start", eager, allocator.enabled)) + if eager: + compile_utils._compile_deep_gemm_one_type_all() + events.append(("capture_end", eager, allocator.enabled)) + return "captured" + + assert cuda_graph._save_capture(runner, capture) == "captured" + assert events == [ + "bootstrap", + ("capture_start", True, True), + "stop", + ("compile", False), + "resume", + ("capture_end", True, True), + ("capture_start", False, True), + ("capture_end", False, True), + ] + + +def test_deepep_load_replays_lazy_initialization_before_preallocation(monkeypatch): + cuda_graph, _mode, _registered = _load_cuda_graph(monkeypatch) + events = [] + compile_utils = _module("sglang.srt.layers.deep_gemm_wrapper.compile_utils") + compile_utils._compile_deep_gemm_one_type_all = lambda: None + monkeypatch.setitem( + sys.modules, + "sglang.srt.layers.deep_gemm_wrapper.compile_utils", + compile_utils, + ) + backend = FullCudaGraphBackend() + backend._pool = object() + runner = SimpleNamespace( + _foundry_graphs_loaded=False, + attn_backend=object(), + backend=backend, + capture_bs=[], + deepep_adapter=SimpleNamespace(capture=lambda **_kwargs: events.append("deepep_mode")), + device_module=SimpleNamespace(graph_pool_handle=lambda: object()), + ) + + monkeypatch.setattr(cuda_graph, "_uses_deepep", lambda _runner: True) + monkeypatch.setattr( + cuda_graph, + "bootstrap_deepep_buffer", + lambda _runner: events.append("bootstrap"), + ) + monkeypatch.setattr( + cuda_graph.graph_ops, + "preallocate_for_load_mode", + lambda: events.append("preallocate"), + raising=False, + ) + monkeypatch.setattr( + cuda_graph.graph_ops, + "load_all_graphs", + lambda _runner: [], + raising=False, + ) + + def eager_capture(_runner): + events.append(("eager_capture", cuda_graph._EAGER_DEEPEP_WARMUP)) + + cuda_graph._load_capture(runner, eager_capture) + + assert runner._foundry_graphs_loaded is True + assert events == [ + "bootstrap", + ("eager_capture", True), + "preallocate", + "deepep_mode", + ] From 094819fa3956ec3c7ac5f3db176130d6d308ea98 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Thu, 23 Jul 2026 13:26:00 +0000 Subject: [PATCH 11/21] fix(sglang): propagate LOAD preallocation failures Co-authored-by: Rahul Chalamala --- python/foundry/integration/sglang/runtime.py | 15 ++++++++-- tests/test_sglang_runtime.py | 30 ++++++++++++++++++++ 2 files changed, 43 insertions(+), 2 deletions(-) diff --git a/python/foundry/integration/sglang/runtime.py b/python/foundry/integration/sglang/runtime.py index b12a1205..bcbb4886 100644 --- a/python/foundry/integration/sglang/runtime.py +++ b/python/foundry/integration/sglang/runtime.py @@ -298,16 +298,27 @@ def preallocate_for_load_mode() -> None: if cfg is None or cfg.mode != CUDAGraphExtensionMode.LOAD: return final = 0 + source = "warmup_state" + path = None 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) + source = "per_rank" if final <= 0: final = load_warmup_state().final_alloc_offset - remaining = final - cge.get_current_alloc_offset() + current = cge.get_current_alloc_offset() + remaining = final - current if remaining > 0: - cge.preallocate_region(remaining) + success = cge.preallocate_region(remaining) + if not success: + rank = None if _state is None else _state.rank + raise RuntimeError( + "Foundry LOAD failed to preallocate the recorded CUDA allocation region: " + f"rank={rank}, requested={remaining}, current={current}, final={final}, " + f"source={source!r}, path={path!r}." + ) def log_alloc_offset(label: str) -> None: diff --git a/tests/test_sglang_runtime.py b/tests/test_sglang_runtime.py index 01f334c8..4c7ffef8 100644 --- a/tests/test_sglang_runtime.py +++ b/tests/test_sglang_runtime.py @@ -178,6 +178,36 @@ def test_validate_warmup_state_accepts_matching_latest_archive(monkeypatch): runtime.validate_warmup_state(_matching_state(runtime), _server_args()) +def test_preallocate_for_load_mode_raises_when_region_cannot_be_reserved( + tmp_path, + monkeypatch, +): + runtime = _load_runtime(monkeypatch) + (tmp_path / "final_alloc_offset.json").write_text(json.dumps({"final_alloc_offset": 4096})) + monkeypatch.setattr( + runtime, + "get_config", + lambda: SimpleNamespace( + mode=runtime.CUDAGraphExtensionMode.LOAD, + workspace_dir=str(tmp_path), + ), + ) + monkeypatch.setattr( + runtime, + "cge", + SimpleNamespace( + get_current_alloc_offset=lambda: 1024, + preallocate_region=lambda _size: False, + ), + ) + + with pytest.raises( + RuntimeError, + match=r"(?i)preallocat.*3072|3072.*preallocat", + ): + runtime.preallocate_for_load_mode() + + def test_save_warmup_state_replaces_existing_file_atomically(tmp_path, monkeypatch): runtime = _load_runtime(monkeypatch, tmp_path) state_path = tmp_path / "warmup_state.json" From 3f057206765236a097a4e9a7d8595b248dc24307 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Thu, 23 Jul 2026 13:41:11 +0000 Subject: [PATCH 12/21] docs(sglang): record H100 validation Co-authored-by: Rahul Chalamala --- README.md | 16 +++++++++------- docs/overview.md | 9 +++++++-- docs/sglang/hooks.md | 2 +- docs/sglang/memory-consistency.md | 7 +++++-- docs/sglang/overview.md | 30 ++++++++++++++++++++---------- docs/sglang/save-load-workflow.md | 11 +++++++++-- recipe/sglang/README.md | 14 +++++++++----- 7 files changed, 60 insertions(+), 29 deletions(-) diff --git a/README.md b/README.md index d52b85b1..f1b54bff 100644 --- a/README.md +++ b/README.md @@ -82,17 +82,19 @@ Foundry ships engine integrations under `foundry/python/foundry/integration/`. P | Engine | Single GPU | DP | TP | EP | |---|:---:|:---:|:---:|:---:| | vLLM | ✅ | ✅ | 🚧 | ✅ | -| SGLang `0.5.15.post1` | Targeted | Targeted | Unsupported | Targeted | +| SGLang `0.5.15.post1` | ✅ | ✅ | Unsupported | ✅ | | TensorRT-LLM | 🚧 | 🚧 | 🚧 | 🚧 | -✅ validated end-to-end (SAVE → LOAD → query)  ·  Targeted = -supported pending H100 validation  ·  🚧 not yet +✅ validated end-to-end (SAVE → LOAD → query)  ·  🚧 not yet The SGLang integration is latest-only for upstream `v0.5.15.post1` -(`0b3bb0c`). It targets full decode graphs with prefill graphs disabled for -single-GPU FlashInfer, regular DP FlashInfer, and DP-attention with DeepEP -low-latency + FA3; these modes are supported pending H100 validation. Existing -SGLang archives are incompatible and require a fresh SAVE. See +(`0b3bb0c`), CUDA 13.0, and PyTorch 2.11/cu130 on H100. Full decode with +prefill graphs disabled is validated for single-GPU FlashInfer and regular +DP=2 FlashInfer on Qwen3-1.7B (52 graphs/rank), plus EP=2 DP-attention with +DeepEP low-latency, DeepGEMM, and FA3 on Qwen3-30B-A3B-FP8 (20 graphs/rank). +All three passed baseline/SAVE/fresh LOAD/query with text/token parity and no +recapture. TP without DP-attention remains unsupported. Existing SGLang +archives are incompatible and require a fresh SAVE. See [`recipe/sglang/README.md`](recipe/sglang/README.md). For the current Foundry-only SGLang workflow, check out upstream diff --git a/docs/overview.md b/docs/overview.md index a09fdf49..39c572d3 100644 --- a/docs/overview.md +++ b/docs/overview.md @@ -29,7 +29,7 @@ hooks and archive metadata are intentionally separate. | Aspect | vLLM | SGLang `0.5.15.post1` | |---|---|---| -| Graph scope | Configured vLLM piecewise/full integration modes | Targeted: full decode only; prefill disabled | +| Graph scope | Configured vLLM piecewise/full integration modes | Validated: full decode only; prefill disabled | | Capture seam | `CUDAGraphWrapper.__call__` | `DecodeCudaGraphRunner.capture` orchestration and `FullCudaGraphBackend.capture_one` storage | | Run-once warmup seam | vLLM compile/kernel/sampler lifecycle hooks | `BaseRunner.warmup` | | Graph identity | vLLM piecewise graph metadata | `ShapeKey(size, stream_idx, variant_label)` | @@ -37,7 +37,12 @@ hooks and archive metadata are intentionally separate. | Attention preparation | vLLM-specific compile/capture path | Separate out-of-graph planning and in-graph GPU metadata work | | Profile behavior | Full profile forward; two-pass SAVE recipe | Pool-memory measurement only; one SAVE process | | Process model | Uniprocess, multiprocessing, or Ray executors | Multiprocessing `spawn`; explicit hook install in scheduler and DP-controller children | -| Parallel scope | Single GPU, DP, and documented EP configuration | Targeted: single GPU FlashInfer, regular DP FlashInfer, and DP-attention + DeepEP low-latency + FA3; H100 validation pending | +| Parallel scope | Single GPU, DP, and documented EP configuration | Validated on H100: single GPU FlashInfer, regular DP=2 FlashInfer, and EP=2 DP-attention + DeepEP low-latency + FA3 | + +SGLang validation used upstream commit `0b3bb0c`, CUDA 13.0, and PyTorch +2.11/cu130. Qwen3-1.7B restored 52 graphs for single GPU and 52 per rank for +DP=2; Qwen3-30B-A3B-FP8 restored 20 graphs per rank for EP=2. Every mode passed +baseline/SAVE/fresh LOAD/query with text/token parity and no recapture. For SGLang, TP without DP-attention, PP, speculative decoding, LoRA/PDMux, dLLM, hidden-state capture, memory-saver graphs, and prefill graphs are not diff --git a/docs/sglang/hooks.md b/docs/sglang/hooks.md index e3a5f2ce..bed7ab78 100644 --- a/docs/sglang/hooks.md +++ b/docs/sglang/hooks.md @@ -2,7 +2,7 @@ `foundry.integration.sglang.install_hooks(server_args)` is the public, idempotent entry point for SGLang `0.5.15.post1`. It loads the Foundry TOML, -validates the targeted graph configuration, and installs the hooks below. +validates the required graph configuration, and installs the hooks below. ## Installation order diff --git a/docs/sglang/memory-consistency.md b/docs/sglang/memory-consistency.md index 80961cef..59379d1a 100644 --- a/docs/sglang/memory-consistency.md +++ b/docs/sglang/memory-consistency.md @@ -100,14 +100,17 @@ DeepEP constructs its singleton NVSHMEM buffer lazily, and DeepGEMM may compile shape-specific kernels lazily. Creating either inside stream capture is invalid. -For the targeted DP-attention + DeepEP low-latency + FA3 configuration, -supported pending H100 validation: +For the H100-validated EP=2 DP-attention + DeepEP low-latency + DeepGEMM + +FA3 configuration: - SAVE runs a dedicated eager orchestration pass to initialize lazy state; - SAVE and LOAD bootstrap the DeepEP buffer before capture/load; - LOAD initializes queued NVSHMEM modules before native graph builds; - LOAD initializes the DeepEP adapter after backend hydration. +Qwen3-30B-A3B-FP8 validation restored 20 graphs per rank with text/token +parity, no recapture, and DeepEP/NVSHMEM/adapter/event restoration evidence. + The DeepEP token limit, chunked-prefill size, model, topology, and environment must remain identical between SAVE and LOAD. diff --git a/docs/sglang/overview.md b/docs/sglang/overview.md index 54aebf32..904c146d 100644 --- a/docs/sglang/overview.md +++ b/docs/sglang/overview.md @@ -1,25 +1,35 @@ # SGLang integration overview -Foundry targets exactly SGLang `0.5.15.post1`, upstream commit `0b3bb0c`. -The integration is designed to persist the release's full decode CUDA graphs -and restore them in a fresh process. Prefill CUDA graphs are disabled. The -configurations below are supported pending H100 validation. +Foundry supports exactly SGLang `0.5.15.post1`, upstream commit `0b3bb0c`. +Full decode CUDA graph SAVE and fresh-process LOAD are validated on H100 with +CUDA 13.0 and PyTorch 2.11/cu130. Prefill CUDA graphs are disabled. -## Targeted configurations +## Validated configurations | Configuration | Attention/MoE backend | Status | |---|---|:---:| -| Single GPU | FlashInfer | Pending H100 validation | -| Regular data parallelism | One full replica per rank, FlashInfer | Pending H100 validation | -| DP-attention expert parallelism | DeepEP low-latency and FA3 | Pending H100 validation | +| Single GPU | FlashInfer | Validated on H100 | +| Regular data parallelism | One full replica per rank, FlashInfer | Validated at DP=2 on H100 | +| DP-attention expert parallelism | DeepEP low-latency, DeepGEMM, and FA3 | Validated at EP=2 on H100 | Tensor parallelism without DP-attention, pipeline parallelism, speculative decoding, LoRA/PDMux, dLLM, hidden-state capture, and memory-saver graphs are rejected at startup. Foundry does not claim prefill-graph support. +Validation results: + +- Single-GPU Qwen3-1.7B passed baseline, SAVE, fresh LOAD, and query with + text/token parity, no recapture, and 52 restored graphs. +- Regular DP=2 Qwen3-1.7B passed on both FlashInfer ranks with text/token + parity, no recapture, and 52 restored graphs per rank. +- EP=2 Qwen3-30B-A3B-FP8 passed with DP-attention, DeepEP low-latency, + DeepGEMM, and FA3, with text/token parity, no recapture, and 20 restored + graphs per rank. Logs confirmed DeepEP/NVSHMEM initialization, adapter + restoration, and graph event restoration. + ## Upstream execution model -The targeted SGLang release separates orchestration from graph storage: +This SGLang release separates orchestration from graph storage: - `BaseRunner.warmup` owns run-once kernel warmup and autotuning. - `DecodeCudaGraphRunner.capture` prepares decode shapes and orchestrates @@ -105,7 +115,7 @@ through a source path is not supported. ## Run it The recipe includes baseline, SAVE, fresh-process LOAD, and query commands for -all targeted configurations: +all validated configurations: - [`../../recipe/sglang/README.md`](../../recipe/sglang/README.md) - [`direct-edits.md`](direct-edits.md) for the external four-file shim diff --git a/docs/sglang/save-load-workflow.md b/docs/sglang/save-load-workflow.md index 30477413..3e4e34b7 100644 --- a/docs/sglang/save-load-workflow.md +++ b/docs/sglang/save-load-workflow.md @@ -26,7 +26,7 @@ sglang serve \ `--cuda-graph-max-bs-decode` is the release's phase-specific CLI field. `--disable-prefill-cuda-graph` ensures that only decode graphs are created. The recipes use both flags in baseline, SAVE, and LOAD because they define the -targeted configuration, which is supported pending H100 validation. +selected full-decode configuration validated on H100. The Foundry option is present only for SAVE/LOAD. It is generated from the external shim's annotated `ServerArgs` field. @@ -88,7 +88,14 @@ curl -s http://127.0.0.1:12000/v1/completions \ Use the DP and DeepEP commands in [`../../recipe/sglang/README.md`](../../recipe/sglang/README.md) for the other -targeted configurations. +validated configurations. + +Validation used upstream SGLang `v0.5.15.post1` (`0b3bb0c`), CUDA 13.0, +PyTorch 2.11/cu130, and H100. Single-GPU and DP=2 Qwen3-1.7B FlashInfer +restored 52 graphs per applicable rank; EP=2 Qwen3-30B-A3B-FP8 with +DP-attention, DeepEP low-latency, DeepGEMM, and FA3 restored 20 graphs per +rank. Baseline/SAVE/fresh LOAD/query passed with text/token parity and no +recapture in every configuration. ## Archive files diff --git a/recipe/sglang/README.md b/recipe/sglang/README.md index 9e0b1875..17f42d85 100644 --- a/recipe/sglang/README.md +++ b/recipe/sglang/README.md @@ -1,16 +1,20 @@ # Foundry recipe — SGLang v0.5.15.post1 -These recipes target persistence and restoration of full decode CUDA graphs -for exactly SGLang `0.5.15.post1` (upstream commit `0b3bb0c`). Prefill graphs -are always disabled. The configurations below are supported pending H100 -validation. +These recipes persist and restore full decode CUDA graphs for exactly SGLang +`0.5.15.post1` (upstream commit `0b3bb0c`). Validation used CUDA 13.0, PyTorch +2.11/cu130, and H100. Prefill graphs are always disabled. -| Script | Targeted configuration | +| Script | Validated configuration | |---|---| | `serve_qwen3-mini.sh` | Single GPU, Qwen3-1.7B, FlashInfer | | `serve_qwen3-1.7b_dp.sh` | Regular data parallelism, one full replica per GPU, FlashInfer | | `serve_qwen3-30ba3bfp8_ep.sh` | DP-attention, DeepEP low-latency, FA3, Qwen3-30B-A3B-FP8 | +Single GPU and DP=2 passed baseline/SAVE/fresh LOAD/query with text/token +parity, no recapture, and 52 graphs per applicable rank. EP=2 passed the same +checks with 20 graphs per rank plus DeepEP/NVSHMEM/adapter/event restoration +evidence. The EP environment used DeepGEMM and FA3. + Tensor parallelism without DP-attention, pipeline parallelism, speculative decoding, LoRA/PDMux, dLLM, hidden-state capture, and memory-saver graphs are not supported. From e11c93d11b355a50da31252c4d89c7ca6458293b Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Thu, 23 Jul 2026 14:11:48 +0000 Subject: [PATCH 13/21] fix(sglang): isolate optional CUDA graph imports Co-authored-by: Rahul Chalamala --- python/foundry/integration/sglang/hooks.py | 13 ------------- tests/test_imports.py | 18 ++++++++++++++++-- tests/test_sglang_api_contract.py | 7 +++++++ 3 files changed, 23 insertions(+), 15 deletions(-) diff --git a/python/foundry/integration/sglang/hooks.py b/python/foundry/integration/sglang/hooks.py index bb3daed8..183fa2e0 100644 --- a/python/foundry/integration/sglang/hooks.py +++ b/python/foundry/integration/sglang/hooks.py @@ -83,7 +83,6 @@ def install_hooks(server_args) -> None: _patch_init_torch_distributed() _patch_init_memory_pool() - _patch_load_model() cuda_graph.install_cuda_graph_hooks() _patch_spawn_sites() @@ -175,18 +174,6 @@ def patched(self, pre_model_load_memory): cls.init_memory_pool = patched -def _patch_load_model() -> None: - mr = importlib.import_module("sglang.srt.model_executor.model_runner") - cls = mr.ModelRunner - orig = cls.load_model - - @functools.wraps(orig) - def patched(self, *args, **kwargs): - return orig(self, *args, **kwargs) - - cls.load_model = patched - - def _patch_spawn_sites() -> None: try: engine_mod = importlib.import_module("sglang.srt.entrypoints.engine") diff --git a/tests/test_imports.py b/tests/test_imports.py index c916069b..c83c3c7c 100644 --- a/tests/test_imports.py +++ b/tests/test_imports.py @@ -13,6 +13,10 @@ """ import importlib +import importlib.util +from importlib.metadata import PackageNotFoundError, version + +import pytest PUBLIC_MODULES = [ # Top-level @@ -32,7 +36,6 @@ "foundry.integration.sglang.hooks", "foundry.integration.sglang.runtime", "foundry.integration.sglang.graph_ops", - "foundry.integration.sglang.cuda_graph", ] # Names that have historically been ruff --fix victims because they live in @@ -77,10 +80,21 @@ def test_vllm_integration_public_api(): def test_sglang_integration_submodules_importable(): - for sub in ("config", "hooks", "runtime", "graph_ops", "cuda_graph"): + for sub in ("config", "hooks", "runtime", "graph_ops"): importlib.import_module(f"foundry.integration.sglang.{sub}") +def test_sglang_cuda_graph_importable_when_sglang_is_installed(): + if importlib.util.find_spec("sglang") is None: + pytest.skip("SGLang is an optional dependency") + try: + version("sglang") + except PackageNotFoundError: + pytest.skip("SGLang distribution metadata is unavailable") + + importlib.import_module("foundry.integration.sglang.cuda_graph") + + def test_sglang_integration_public_api(): mod = importlib.import_module("foundry.integration.sglang") for name in ("install_hooks", "CUDAGraphExtensionMode", "get_graph_extension_mode"): diff --git a/tests/test_sglang_api_contract.py b/tests/test_sglang_api_contract.py index 7095c4a3..b63849c4 100644 --- a/tests/test_sglang_api_contract.py +++ b/tests/test_sglang_api_contract.py @@ -84,6 +84,13 @@ def test_attention_backend_uses_split_graph_metadata_contract(): ] +def test_deep_gemm_compile_hook_is_available(): + _require_sglang() + compile_utils = importlib.import_module("sglang.srt.layers.deep_gemm_wrapper.compile_utils") + + assert callable(compile_utils._compile_deep_gemm_one_type_all) + + def test_removed_monolithic_cuda_graph_runner_is_not_part_of_contract(): _require_sglang() From c488d5adaad12ae66d91c6a5f668136a646bd410 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Thu, 23 Jul 2026 14:11:53 +0000 Subject: [PATCH 14/21] fix(sglang): require per-rank load offset Co-authored-by: Rahul Chalamala --- python/foundry/integration/sglang/runtime.py | 38 +++++++-------- tests/test_sglang_runtime.py | 51 ++++++++++++++++++++ 2 files changed, 70 insertions(+), 19 deletions(-) diff --git a/python/foundry/integration/sglang/runtime.py b/python/foundry/integration/sglang/runtime.py index bcbb4886..b4e54a3c 100644 --- a/python/foundry/integration/sglang/runtime.py +++ b/python/foundry/integration/sglang/runtime.py @@ -48,7 +48,6 @@ class WarmupState: gpu_name: str = "" gpu_total_memory: int = 0 memory_pool_config: dict = field(default_factory=dict) - final_alloc_offset: int = 0 @dataclass @@ -283,12 +282,6 @@ def capture_final_alloc_offset() -> int: path = os.path.join(cfg.workspace_dir, "final_alloc_offset.json") with open(path, "w") as f: json.dump({"final_alloc_offset": _final_alloc_offset}, f) - if cfg.workspace_root is not None: - warmup_state_path = os.path.join(cfg.workspace_root, "warmup_state.json") - if os.path.exists(warmup_state_path): - state = load_warmup_state() - state.final_alloc_offset = _final_alloc_offset - save_warmup_state(state) logger.info("[Foundry] SGLang final_alloc_offset=%d", _final_alloc_offset) return _final_alloc_offset @@ -297,17 +290,24 @@ def preallocate_for_load_mode() -> None: cfg = get_config() if cfg is None or cfg.mode != CUDAGraphExtensionMode.LOAD: return - final = 0 - source = "warmup_state" - path = None - 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) - source = "per_rank" - if final <= 0: - final = load_warmup_state().final_alloc_offset + if cfg.workspace_dir is None: + raise _compatibility_error( + "Foundry LOAD has no rank workspace for its per-rank final allocation offset." + ) + path = os.path.join(cfg.workspace_dir, "final_alloc_offset.json") + try: + with open(path) as f: + data = json.load(f) + except (OSError, json.JSONDecodeError) as exc: + raise _compatibility_error( + f"Cannot read the per-rank final allocation offset at {path}: {exc}" + ) from exc + final = data.get("final_alloc_offset") if isinstance(data, dict) else None + if isinstance(final, bool) or not isinstance(final, int) or final <= 0: + raise _compatibility_error( + f"Invalid per-rank final allocation offset at {path}: {final!r}; " + "expected a positive integer." + ) current = cge.get_current_alloc_offset() remaining = final - current if remaining > 0: @@ -317,7 +317,7 @@ def preallocate_for_load_mode() -> None: raise RuntimeError( "Foundry LOAD failed to preallocate the recorded CUDA allocation region: " f"rank={rank}, requested={remaining}, current={current}, final={final}, " - f"source={source!r}, path={path!r}." + f"source='per_rank', path={path!r}." ) diff --git a/tests/test_sglang_runtime.py b/tests/test_sglang_runtime.py index 4c7ffef8..614598ab 100644 --- a/tests/test_sglang_runtime.py +++ b/tests/test_sglang_runtime.py @@ -208,6 +208,57 @@ def test_preallocate_for_load_mode_raises_when_region_cannot_be_reserved( runtime.preallocate_for_load_mode() +@pytest.mark.parametrize( + ("contents", "case"), + [ + (None, "missing"), + ("{", "malformed"), + (json.dumps({"final_alloc_offset": "4096"}), "wrong_type"), + (json.dumps({"final_alloc_offset": 0}), "zero"), + (json.dumps({"final_alloc_offset": -1}), "negative"), + ], +) +def test_preallocate_for_load_mode_requires_positive_per_rank_final_offset( + tmp_path, + monkeypatch, + contents, + case, +): + runtime = _load_runtime(monkeypatch) + if contents is not None: + (tmp_path / "final_alloc_offset.json").write_text(contents) + monkeypatch.setattr( + runtime, + "get_config", + lambda: SimpleNamespace( + mode=runtime.CUDAGraphExtensionMode.LOAD, + workspace_dir=str(tmp_path), + ), + ) + preallocate_calls = [] + monkeypatch.setattr( + runtime, + "cge", + SimpleNamespace( + get_current_alloc_offset=lambda: 1024, + preallocate_region=lambda size: preallocate_calls.append(size), + ), + ) + monkeypatch.setattr( + runtime, + "load_warmup_state", + lambda: SimpleNamespace(final_alloc_offset=8192), + ) + + with pytest.raises( + RuntimeError, + match=r"(?is)per-rank.*final.*offset.*fresh.*SAVE|fresh.*SAVE.*per-rank.*final.*offset", + ): + runtime.preallocate_for_load_mode() + + assert not preallocate_calls, case + + def test_save_warmup_state_replaces_existing_file_atomically(tmp_path, monkeypatch): runtime = _load_runtime(monkeypatch, tmp_path) state_path = tmp_path / "warmup_state.json" From 3864a4c524debe28d42a06ff9e21ca86f7d24089 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Thu, 23 Jul 2026 14:53:33 +0000 Subject: [PATCH 15/21] docs(sglang): clarify per-rank allocation watermark Co-authored-by: Rahul Chalamala --- docs/sglang/memory-lifecycle.md | 3 ++- docs/sglang/save-load-workflow.md | 8 ++++---- recipe/sglang/README.md | 4 ++-- 3 files changed, 8 insertions(+), 7 deletions(-) diff --git a/docs/sglang/memory-lifecycle.md b/docs/sglang/memory-lifecycle.md index 1a70d83b..0652794c 100644 --- a/docs/sglang/memory-lifecycle.md +++ b/docs/sglang/memory-lifecycle.md @@ -147,9 +147,10 @@ cuda_version gpu_name gpu_total_memory memory_pool_config -final_alloc_offset ``` LOAD requires archive schema 1, SGLang `0.5.15.post1`, decode `full`, prefill `disabled`, matching installed package metadata, and matching current resolved backends. A mismatch requires a fresh SAVE; no archive migration is attempted. +The VMM watermark is rank-specific and is persisted separately in +`rank_/final_alloc_offset.json`. diff --git a/docs/sglang/save-load-workflow.md b/docs/sglang/save-load-workflow.md index 3e4e34b7..b6917be5 100644 --- a/docs/sglang/save-load-workflow.md +++ b/docs/sglang/save-load-workflow.md @@ -127,13 +127,13 @@ Each additional DP/EP rank has its own `rank_/`. "cuda_version": "13.0", "gpu_name": "NVIDIA H100", "gpu_total_memory": 0, - "memory_pool_config": {}, - "final_alloc_offset": 0 + "memory_pool_config": {} } ``` -The timestamp, GPU values, pool dictionary, and final offset shown above are -illustrative; SAVE records the actual values. +The timestamp, GPU values, and pool dictionary shown above are illustrative; +SAVE records the actual values. Each rank's VMM watermark lives in its +`rank_/final_alloc_offset.json`. `sglang_graph_index.json` is authoritative: diff --git a/recipe/sglang/README.md b/recipe/sglang/README.md index 17f42d85..fca9fdc6 100644 --- a/recipe/sglang/README.md +++ b/recipe/sglang/README.md @@ -190,8 +190,8 @@ foundry_archive/ `warmup_state.json` schema 1 records `schema_version`, `sglang_version`, `decode_cuda_graph_backend`, `prefill_cuda_graph_backend`, `timestamp`, -`cuda_version`, `gpu_name`, `gpu_total_memory`, `memory_pool_config`, and -`final_alloc_offset`. +`cuda_version`, `gpu_name`, `gpu_total_memory`, and `memory_pool_config`. +Each rank's `final_alloc_offset.json` stores its own VMM watermark. `sglang_graph_index.json` schema 1 is the authoritative graph list. Each entry records capture order, phase, filename, output kind, and the complete From fdb3eccb4cff866aaeb9ce1b74c6e1a77eabdd13 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Thu, 23 Jul 2026 22:49:47 +0000 Subject: [PATCH 16/21] feat(sglang): validate TP2 DP2 topology Co-authored-by: Rahul Chalamala --- .../foundry/integration/sglang/cuda_graph.py | 38 +++++-- tests/test_sglang_config.py | 100 ++++++++++++++++++ tests/test_sglang_cuda_graph.py | 96 +++++++++++++++++ 3 files changed, 228 insertions(+), 6 deletions(-) create mode 100644 tests/test_sglang_config.py diff --git a/python/foundry/integration/sglang/cuda_graph.py b/python/foundry/integration/sglang/cuda_graph.py index 444fa5bd..991c23fa 100644 --- a/python/foundry/integration/sglang/cuda_graph.py +++ b/python/foundry/integration/sglang/cuda_graph.py @@ -14,6 +14,7 @@ import functools import importlib import logging +import os import sys import time from collections.abc import Callable, Iterable, Iterator @@ -82,13 +83,38 @@ def validate_configuration(server_args: Any) -> None: ) tp_size = getattr(server_args, "tp_size", 1) + dp_size = getattr(server_args, "dp_size", 1) + pp_size = getattr(server_args, "pp_size", 1) + ep_size = getattr(server_args, "ep_size", 1) dp_attention = bool(getattr(server_args, "enable_dp_attention", False)) - if tp_size > 1 and not dp_attention: - raise ValueError( - "Foundry does not support tensor parallel decode without DP attention " - f"(tp_size={tp_size})." - ) - if getattr(server_args, "pp_size", 1) > 1: + regular_tp = tp_size > 1 and not dp_attention + if regular_tp: + topology = (tp_size, dp_size, pp_size, ep_size) + if topology != (2, 2, 1, 1): + raise ValueError( + "Foundry tensor-parallel decode supports only " + "tp_size=2, dp_size=2, pp_size=1, ep_size=1; " + f"got tp_size={tp_size}, dp_size={dp_size}, " + f"pp_size={pp_size}, ep_size={ep_size}." + ) + if not bool(getattr(server_args, "disable_custom_all_reduce", False)): + raise ValueError("Foundry TP=2 x DP=2 requires --disable-custom-all-reduce.") + if bool(getattr(server_args, "enable_nccl_nvls", False)): + raise ValueError( + "Foundry TP=2 x DP=2 does not support NCCL NVLS; leave --enable-nccl-nvls disabled." + ) + if bool(getattr(server_args, "enable_symm_mem", False)) or bool( + getattr(server_args, "enable_torch_symm_mem", False) + ): + raise ValueError( + "Foundry TP=2 x DP=2 does not support symmetric memory collectives; " + "leave --enable-symm-mem and --enable-torch-symm-mem disabled." + ) + for variable in ("NCCL_CUMEM_ENABLE", "NCCL_NVLS_ENABLE"): + value = os.environ.get(variable) + if value != "0": + raise ValueError(f"Foundry TP=2 x DP=2 requires {variable}='0'; got {value!r}.") + if pp_size > 1: raise ValueError("Foundry does not support pipeline parallel CUDA graphs (pp_size > 1).") if getattr(server_args, "speculative_algorithm", None) is not None: raise ValueError("Foundry does not support speculative decoding.") diff --git a/tests/test_sglang_config.py b/tests/test_sglang_config.py new file mode 100644 index 00000000..0f99afae --- /dev/null +++ b/tests/test_sglang_config.py @@ -0,0 +1,100 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the Foundry project +"""Rank-mapping contracts for hybrid SGLang TP and DP workers.""" + +import importlib.util +import sys +from pathlib import Path +from types import ModuleType, SimpleNamespace + +import pytest + +SGLANG_INTEGRATION_DIR = Path(__file__).parents[1] / "python" / "foundry" / "integration" / "sglang" + + +def _load_module(name: str, path: Path): + spec = importlib.util.spec_from_file_location(name, path) + assert spec is not None + assert spec.loader is not None + module = importlib.util.module_from_spec(spec) + sys.modules[name] = module + spec.loader.exec_module(module) + return module + + +def _module(name: str, **attributes): + module = ModuleType(name) + for attribute, value in attributes.items(): + setattr(module, attribute, value) + return module + + +def _load_hooks(monkeypatch: pytest.MonkeyPatch): + foundry = _module("foundry") + foundry.__path__ = [] + integration = _module("foundry.integration") + integration.__path__ = [] + sglang = _module("foundry.integration.sglang") + sglang.__path__ = [] + runtime = _module("foundry.integration.sglang.runtime") + config = _module( + "foundry.integration.sglang.config", + CUDAGraphExtensionMode=SimpleNamespace(NONE="none"), + get_graph_extension_mode=lambda: "none", + get_workspace_root=lambda: None, + load_graph_extension_config=lambda _path: None, + ) + for name, module in { + "torch": _module("torch"), + "foundry": foundry, + "foundry.integration": integration, + "foundry.integration.sglang": sglang, + "foundry.integration.sglang.runtime": runtime, + "foundry.integration.sglang.config": config, + }.items(): + monkeypatch.setitem(sys.modules, name, module) + return _load_module( + "foundry_sglang_hooks_rank_test", + SGLANG_INTEGRATION_DIR / "hooks.py", + ) + + +@pytest.mark.parametrize( + ("dp_rank", "tp_rank", "expected"), + [ + (0, 0, 0), + (0, 1, 1), + (1, 0, 2), + (1, 1, 3), + ], +) +def test_tp2_dp2_workers_have_distinct_workspace_ranks(monkeypatch, dp_rank, tp_rank, expected): + config = _load_module( + "foundry_sglang_config_rank_test", + SGLANG_INTEGRATION_DIR / "config.py", + ) + hooks = _load_hooks(monkeypatch) + server_args = SimpleNamespace( + tp_size=2, + pp_size=1, + dp_size=2, + enable_dp_attention=False, + ) + model_runner = SimpleNamespace( + server_args=server_args, + tp_rank=tp_rank, + dp_rank=dp_rank, + ) + + resolved_dp_rank = hooks._resolve_dp_rank(model_runner) + + assert resolved_dp_rank == dp_rank + assert ( + config.compute_workspace_rank( + server_args, + tp_rank=tp_rank, + pp_rank=0, + dp_rank=resolved_dp_rank, + ) + == expected + ) diff --git a/tests/test_sglang_cuda_graph.py b/tests/test_sglang_cuda_graph.py index 3d81e812..798bf280 100644 --- a/tests/test_sglang_cuda_graph.py +++ b/tests/test_sglang_cuda_graph.py @@ -185,9 +185,14 @@ def _server_args(**overrides): "tp_size": 1, "pp_size": 1, "dp_size": 1, + "ep_size": 1, "enable_dp_attention": False, "attention_backend": "flashinfer", "moe_a2a_backend": "none", + "disable_custom_all_reduce": False, + "enable_nccl_nvls": False, + "enable_symm_mem": False, + "enable_torch_symm_mem": False, "speculative_algorithm": None, "enable_lora": False, "enable_pdmux": False, @@ -206,6 +211,8 @@ def _server_args(**overrides): _server_args(dp_size=2), _server_args( tp_size=2, + dp_size=2, + ep_size=2, enable_dp_attention=True, attention_backend="fa3", moe_a2a_backend="deepep", @@ -218,6 +225,95 @@ def test_configuration_accepts_supported_full_decode_modes(monkeypatch, server_a cuda_graph.validate_configuration(server_args) +def test_configuration_accepts_exact_tp2_dp2_full_decode(monkeypatch): + monkeypatch.setenv("NCCL_CUMEM_ENABLE", "0") + monkeypatch.setenv("NCCL_NVLS_ENABLE", "0") + cuda_graph, _mode, _registered = _load_cuda_graph(monkeypatch) + + cuda_graph.validate_configuration( + _server_args( + tp_size=2, + dp_size=2, + ep_size=1, + disable_custom_all_reduce=True, + ) + ) + + +@pytest.mark.parametrize( + "overrides", + [ + {"tp_size": 2, "dp_size": 1}, + {"tp_size": 2, "dp_size": 3}, + {"tp_size": 2, "dp_size": 2, "ep_size": 2}, + {"tp_size": 4, "dp_size": 2}, + ], +) +def test_configuration_rejects_nearby_tp_topologies(monkeypatch, overrides): + monkeypatch.setenv("NCCL_CUMEM_ENABLE", "0") + monkeypatch.setenv("NCCL_NVLS_ENABLE", "0") + cuda_graph, _mode, _registered = _load_cuda_graph(monkeypatch) + + with pytest.raises( + ValueError, + match=r"tp_size=2, dp_size=2, pp_size=1, ep_size=1", + ): + cuda_graph.validate_configuration(_server_args(**overrides)) + + +@pytest.mark.parametrize( + ("overrides", "diagnostic"), + [ + ({"disable_custom_all_reduce": False}, "disable.custom.all.reduce"), + ({"enable_nccl_nvls": True}, "NVLS"), + ({"enable_symm_mem": True}, "symmetric memory"), + ({"enable_torch_symm_mem": True}, "symmetric memory"), + ], +) +def test_configuration_rejects_unsafe_tp2_dp2_collectives(monkeypatch, overrides, diagnostic): + monkeypatch.setenv("NCCL_CUMEM_ENABLE", "0") + monkeypatch.setenv("NCCL_NVLS_ENABLE", "0") + cuda_graph, _mode, _registered = _load_cuda_graph(monkeypatch) + values = { + "tp_size": 2, + "dp_size": 2, + "disable_custom_all_reduce": True, + } + values.update(overrides) + server_args = _server_args(**values) + + with pytest.raises(ValueError, match=diagnostic): + cuda_graph.validate_configuration(server_args) + + +@pytest.mark.parametrize( + ("variable", "value"), + [ + ("NCCL_CUMEM_ENABLE", None), + ("NCCL_CUMEM_ENABLE", "1"), + ("NCCL_NVLS_ENABLE", None), + ("NCCL_NVLS_ENABLE", "1"), + ], +) +def test_configuration_requires_disabled_nccl_vmm_fast_paths(monkeypatch, variable, value): + monkeypatch.setenv("NCCL_CUMEM_ENABLE", "0") + monkeypatch.setenv("NCCL_NVLS_ENABLE", "0") + if value is None: + monkeypatch.delenv(variable) + else: + monkeypatch.setenv(variable, value) + cuda_graph, _mode, _registered = _load_cuda_graph(monkeypatch) + + with pytest.raises(ValueError, match=variable): + cuda_graph.validate_configuration( + _server_args( + tp_size=2, + dp_size=2, + disable_custom_all_reduce=True, + ) + ) + + @pytest.mark.parametrize( ("overrides", "diagnostic"), [ From d33d9961103521f4dcd4f9b272c85fc5764eeeb1 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Thu, 23 Jul 2026 22:50:48 +0000 Subject: [PATCH 17/21] feat(sglang): add TP2 DP2 serve recipe Co-authored-by: Rahul Chalamala --- recipe/sglang/foundry_tp2_dp2_load.toml | 5 +++ recipe/sglang/foundry_tp2_dp2_save.toml | 5 +++ recipe/sglang/serve_qwen3-8b_tp2_dp2.sh | 45 +++++++++++++++++++++++++ tests/test_sglang_recipes.py | 43 +++++++++++++++++++++++ 4 files changed, 98 insertions(+) create mode 100644 recipe/sglang/foundry_tp2_dp2_load.toml create mode 100644 recipe/sglang/foundry_tp2_dp2_save.toml create mode 100755 recipe/sglang/serve_qwen3-8b_tp2_dp2.sh diff --git a/recipe/sglang/foundry_tp2_dp2_load.toml b/recipe/sglang/foundry_tp2_dp2_load.toml new file mode 100644 index 00000000..fa256abf --- /dev/null +++ b/recipe/sglang/foundry_tp2_dp2_load.toml @@ -0,0 +1,5 @@ +mode = "load" +base_addr = 0x600000000000 +region_size = "256GB" +workspace_root = "foundry_archive_sglang_qwen3_8b_tp2_dp2" +scratch_space_size = "1024MB" diff --git a/recipe/sglang/foundry_tp2_dp2_save.toml b/recipe/sglang/foundry_tp2_dp2_save.toml new file mode 100644 index 00000000..7410cef3 --- /dev/null +++ b/recipe/sglang/foundry_tp2_dp2_save.toml @@ -0,0 +1,5 @@ +mode = "save" +base_addr = 0x600000000000 +region_size = "256GB" +workspace_root = "foundry_archive_sglang_qwen3_8b_tp2_dp2" +scratch_space_size = "1024MB" diff --git a/recipe/sglang/serve_qwen3-8b_tp2_dp2.sh b/recipe/sglang/serve_qwen3-8b_tp2_dp2.sh new file mode 100755 index 00000000..806fab71 --- /dev/null +++ b/recipe/sglang/serve_qwen3-8b_tp2_dp2.sh @@ -0,0 +1,45 @@ +#!/bin/bash +# Qwen3-8B, two TP ranks in each of two regular-DP replicas. +# Usage: CUDA_VISIBLE_DEVICES=0,1,2,3 bash serve_qwen3-8b_tp2_dp2.sh [--save|--load] + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" + +# Keep Foundry's validated NCCL transport constraints identical in baseline, +# SAVE, and LOAD so phase comparisons do not change collective behavior. +export NCCL_CUMEM_ENABLE=0 +export NCCL_NVLS_ENABLE=0 + +FOUNDRY_ARGS=() +if [[ "$1" == "--save" ]]; then + FOUNDRY_TOML="${SCRIPT_DIR}/foundry_tp2_dp2_save.toml" + FOUNDRY_ARGS+=( --foundry-graph-extension-config-path "${FOUNDRY_TOML}" ) + echo "Using foundry SAVE: ${FOUNDRY_TOML}" +elif [[ "$1" == "--load" ]]; then + FOUNDRY_TOML="${SCRIPT_DIR}/foundry_tp2_dp2_load.toml" + FOUNDRY_ARGS+=( --foundry-graph-extension-config-path "${FOUNDRY_TOML}" ) + echo "Using foundry LOAD: ${FOUNDRY_TOML}" +elif [[ -n "$1" ]]; then + echo "Usage: $0 [--save|--load]" + exit 1 +else + echo "Running without foundry (baseline SGLang)" +fi + +sglang serve \ + --model-path Qwen/Qwen3-8B \ + --trust-remote-code \ + --host 0.0.0.0 \ + --port 12000 \ + --tp-size 2 \ + --dp-size 2 \ + --pp-size 1 \ + --ep-size 1 \ + --load-balance-method round_robin \ + --mem-fraction-static 0.8 \ + --disable-radix-cache \ + --disable-custom-all-reduce \ + --attention-backend flashinfer \ + --cuda-graph-backend-decode full \ + --cuda-graph-max-bs-decode 256 \ + --disable-prefill-cuda-graph \ + "${FOUNDRY_ARGS[@]}" diff --git a/tests/test_sglang_recipes.py b/tests/test_sglang_recipes.py index b8856207..f6868598 100644 --- a/tests/test_sglang_recipes.py +++ b/tests/test_sglang_recipes.py @@ -6,6 +6,7 @@ from pathlib import Path import pytest +import tomllib RECIPE_DIR = Path(__file__).parents[1] / "recipe" / "sglang" SERVE_SCRIPTS = ( @@ -13,6 +14,9 @@ "serve_qwen3-1.7b_dp.sh", "serve_qwen3-30ba3bfp8_ep.sh", ) +TP2_DP2_SCRIPT = "serve_qwen3-8b_tp2_dp2.sh" +TP2_DP2_SAVE_CONFIG = "foundry_tp2_dp2_save.toml" +TP2_DP2_LOAD_CONFIG = "foundry_tp2_dp2_load.toml" @pytest.mark.parametrize("script_name", SERVE_SCRIPTS) @@ -22,3 +26,42 @@ def test_recipe_uses_phase_specific_decode_graph_flags(script_name): assert "--cuda-graph-max-bs-decode" in script assert "--disable-prefill-cuda-graph" in script assert re.search(r"--cuda-graph-max-bs(?:\s|$)", script) is None + + +def test_tp2_dp2_recipe_has_exact_validated_topology_and_collective_guards(): + script = (RECIPE_DIR / TP2_DP2_SCRIPT).read_text() + + for argument in ( + "--model-path Qwen/Qwen3-8B", + "--tp-size 2", + "--dp-size 2", + "--pp-size 1", + "--ep-size 1", + "--load-balance-method round_robin", + "--mem-fraction-static 0.8", + "--disable-radix-cache", + "--disable-custom-all-reduce", + "--attention-backend flashinfer", + "--cuda-graph-backend-decode full", + "--cuda-graph-max-bs-decode 256", + "--disable-prefill-cuda-graph", + ): + assert argument in script + assert "export NCCL_CUMEM_ENABLE=0" in script + assert "export NCCL_NVLS_ENABLE=0" in script + assert script.index("export NCCL_CUMEM_ENABLE=0") < script.index('if [[ "$1" == "--save" ]]') + assert script.index("export NCCL_NVLS_ENABLE=0") < script.index('if [[ "$1" == "--save" ]]') + assert TP2_DP2_SAVE_CONFIG in script + assert TP2_DP2_LOAD_CONFIG in script + + +def test_tp2_dp2_save_and_load_configs_differ_only_by_mode(): + with open(RECIPE_DIR / TP2_DP2_SAVE_CONFIG, "rb") as file: + save_config = tomllib.load(file) + with open(RECIPE_DIR / TP2_DP2_LOAD_CONFIG, "rb") as file: + load_config = tomllib.load(file) + + assert save_config.pop("mode") == "save" + assert load_config.pop("mode") == "load" + assert save_config == load_config + assert save_config["workspace_root"] != "foundry_archive" From de92ab1b1f5f3caa9eedb49f8cf334d9cf8500ae Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Thu, 23 Jul 2026 22:52:18 +0000 Subject: [PATCH 18/21] fix: harden VMM IPC across processes Co-authored-by: Rahul Chalamala --- csrc/hook.cpp | 637 +++++++++++++++++++++++++++++++++++++++++++++++--- 1 file changed, 604 insertions(+), 33 deletions(-) diff --git a/csrc/hook.cpp b/csrc/hook.cpp index c3ede37e..a78a35ef 100644 --- a/csrc/hook.cpp +++ b/csrc/hook.cpp @@ -14,12 +14,20 @@ #include #include #include +#include #include #include #include #include #include #include +#include +#include +#include +#include +#include +#include +#include #include #include #include @@ -253,6 +261,27 @@ static std::once_flag default_allocation_region_flag; static boost::unordered::concurrent_flat_map global_alloc_metadata; static boost::unordered::concurrent_flat_map global_carved_reserve_metadata; +// Defined in the VMM-IPC section below (inside the extern "C" block, hence +// the matching linkage here); closes and forgets the exported fd for a VA +// when its allocation is freed (an exported fd keeps the physical pages +// alive and would otherwise serve stale memory to importers). +extern "C" { +static void vmm_ipc_invalidate_export(CUdeviceptr dptr); +} + +// Process-global mirrors of live LOAD-mode preallocated chunks. The +// authoritative state is thread-local, but cuIpcGetMemHandle may run on a +// different thread. Keep one coherent record per range so concurrent devices +// or regions cannot overwrite each other's handle/base/size tuple. +struct PreallocatedChunk { + size_t size; + CUmemGenericAllocationHandle handle; + uint64_t generation; +}; +static std::mutex preallocated_chunks_mutex; +static std::map preallocated_chunks; +static std::atomic next_preallocated_chunk_generation{1}; + struct HookAllocationEvent { enum class Type { Alloc, Free, Reserve }; Type type; @@ -2659,6 +2688,7 @@ CUresult cuMemFree_v2(CUdeviceptr dptr) { mem_addr_free_func(metadata.ptr, metadata.size); } + vmm_ipc_invalidate_export(dptr); global_alloc_metadata.erase_if( [dptr](const std::pair& kv) { return kv.first == dptr; }); @@ -2862,8 +2892,254 @@ void* dlsym(void* handle, const char* symbol) { // VMM allocations by translating to cuMemExportToShareableHandle API // ============================================================================= -// Magic marker to identify VMM IPC handles in CUipcMemHandle -static constexpr uint32_t VMM_IPC_MAGIC = 0x564D4D49; // "VMMI" +// Magic marker to identify VMM IPC handles in CUipcMemHandle. +// v2 ("VMM2"): the blob carries {magic, exporter pid, original ptr, aligned size}. +// The POSIX fd itself is transferred via SCM_RIGHTS over a per-process abstract +// unix socket (a raw fd integer is meaningless in another process - the v1 bug +// behind "cuMemImportFromShareableHandle failed with error 999"). +static constexpr uint32_t VMM_IPC_MAGIC = 0x564D4D32; // "VMM2" + +// --------------------------------------------------------------------------- +// VMM-IPC fd transport: each exporting process runs one server thread on an +// abstract unix socket "\0foundry-vmm-ipc.". Importers connect, send the +// 8-byte exporter VA they want, and receive the exported fd via SCM_RIGHTS. +// The server does no CUDA calls: fds are exported on the cuIpcGetMemHandle +// caller's thread and parked in vmm_ipc_exported_fds. +// --------------------------------------------------------------------------- + +// exporter VA -> (allocation handle it was exported from, fd). The handle is +// kept so VA reuse after free/realloc invalidates the cached fd. +static boost::unordered::concurrent_flat_map> + vmm_ipc_exported_fds; +static std::mutex vmm_ipc_server_mutex; +static int vmm_ipc_listen_fd = -1; +// pid that owns the running server thread. fork() copies this .so's state but +// not threads, so a forked child must rebind its own socket (vLLM's default +// worker start method is fork). +static pid_t vmm_ipc_server_pid = -1; +// Per-process random token carried in the handle blob and verified by the +// server: defends against pid reuse (a recycled pid + the deterministic +// shared-base VAs would otherwise let an importer silently fetch a different +// process's allocation). +static uint64_t vmm_ipc_token = 0; + +// Wire format of a fetch request. +struct VmmIpcRequest { + uint64_t ptr; + uint64_t token; +}; + +static void vmm_ipc_set_timeouts(int fd) { + struct timeval tv = {30, 0}; + setsockopt(fd, SOL_SOCKET, SO_RCVTIMEO, &tv, sizeof(tv)); + setsockopt(fd, SOL_SOCKET, SO_SNDTIMEO, &tv, sizeof(tv)); +} + +// Imported whole-chunk mappings (LOAD-mode exports: a buffer carved from the +// peer's preallocated chunk is shared by exporting the chunk handle plus an +// offset; cuMemMap cannot map a sub-range of an imported handle, so we map +// the entire peer chunk once and hand out interior pointers). The per-process +// token is part of the key so PID reuse cannot return an old exporter's mapping. +// Refcounted by open/close calls. +struct VmmIpcChunkMapping { + CUdeviceptr local_base; + size_t size; + CUmemGenericAllocationHandle handle; + int refcount; +}; +using VmmIpcChunkKey = std::tuple; +struct VmmIpcChunkSubptr { + VmmIpcChunkKey key; + int open_count; +}; +static std::mutex vmm_ipc_chunk_mutex; +static std::map vmm_ipc_chunk_mappings; +// Interior pointer handed to a caller -> owning chunk key and number of opens. +static std::map vmm_ipc_chunk_subptrs; + +// Dedicated VA zone for relocated peer imports, far below the foundry region +// (default 0x600000000000) and the large-reserve zone right above region end. +// Without this, the driver tends to place a relocated import immediately +// after the region reservation - exactly where the hooked large-reserve hint +// sends the NVSHMEM symmetric heap next, which silently breaks the heap's +// SAVE/LOAD address determinism (observed: heap displaced by a peer-chunk +// import landing at region_end). +static std::atomic vmm_ipc_import_hint{0x300000000000ULL}; + +static socklen_t vmm_ipc_socket_addr(pid_t pid, struct sockaddr_un* addr) { + memset(addr, 0, sizeof(*addr)); + addr->sun_family = AF_UNIX; + // Abstract namespace (sun_path[0] == '\0'): no filesystem entry, vanishes + // with the process. + int n = snprintf(addr->sun_path + 1, sizeof(addr->sun_path) - 2, "foundry-vmm-ipc.%d", (int)pid); + return (socklen_t)(offsetof(struct sockaddr_un, sun_path) + 1 + n); +} + +static void vmm_ipc_server_loop(int listen_fd) { + while (true) { + int conn = accept(listen_fd, nullptr, nullptr); + if (conn < 0) { + if (errno == EBADF || errno == EINVAL) + break; // socket gone + continue; // EINTR/ECONNABORTED/EMFILE/... are transient + } + vmm_ipc_set_timeouts(conn); + // Abstract sockets have no filesystem permissions: only serve same-uid peers. + struct ucred cred; + socklen_t cred_len = sizeof(cred); + if (getsockopt(conn, SOL_SOCKET, SO_PEERCRED, &cred, &cred_len) != 0 || cred.uid != getuid()) { + close(conn); + continue; + } + VmmIpcRequest req = {}; + int fd = -1; + if (recv(conn, &req, sizeof(req), MSG_WAITALL) == (ssize_t)sizeof(req) && + req.token == vmm_ipc_token) { + // Dup under the map lock: the export path may concurrently close and + // replace the parked fd (stale-entry refresh); the dup we send is ours. + vmm_ipc_exported_fds.visit( + (CUdeviceptr)req.ptr, + [&](const std::pair>& + kv) { fd = fcntl(kv.second.second, F_DUPFD_CLOEXEC, 0); }); + } + char status = (fd >= 0) ? 0 : 1; + struct iovec iov = {&status, 1}; + char cbuf[CMSG_SPACE(sizeof(int))] = {}; + struct msghdr msg = {}; + msg.msg_iov = &iov; + msg.msg_iovlen = 1; + if (fd >= 0) { + msg.msg_control = cbuf; + msg.msg_controllen = CMSG_SPACE(sizeof(int)); + struct cmsghdr* c = CMSG_FIRSTHDR(&msg); + c->cmsg_level = SOL_SOCKET; + c->cmsg_type = SCM_RIGHTS; + c->cmsg_len = CMSG_LEN(sizeof(int)); + memcpy(CMSG_DATA(c), &fd, sizeof(int)); + } + // MSG_NOSIGNAL: a peer dying mid-handshake must not SIGPIPE-kill us. + if (sendmsg(conn, &msg, MSG_NOSIGNAL) != 1) { + fprintf(stderr, "[HOOK] WARNING: VMM-IPC fd server sendmsg failed (errno=%d)\n", errno); + } + if (fd >= 0) + close(fd); + close(conn); + } +} + +static bool vmm_ipc_ensure_server() { + std::lock_guard lock(vmm_ipc_server_mutex); + pid_t pid = getpid(); + if (vmm_ipc_server_pid == pid) { + return vmm_ipc_listen_fd >= 0; + } + if (vmm_ipc_server_pid != -1) { + // fork() inherits parked descriptors even with FD_CLOEXEC. The child has + // no server thread and must not pin or serve the parent's allocations. + vmm_ipc_exported_fds.erase_if( + [](const std::pair>& kv) { + close(kv.second.second); + return true; + }); + } + // First call in this process, or first after fork (the parent's accept + // thread did not survive). Close any inherited listen fd so the parent's + // abstract name is not pinned alive by us after the parent exits. + if (vmm_ipc_listen_fd >= 0) { + close(vmm_ipc_listen_fd); + vmm_ipc_listen_fd = -1; + } + int s = socket(AF_UNIX, SOCK_STREAM | SOCK_CLOEXEC, 0); + struct sockaddr_un addr; + socklen_t len = vmm_ipc_socket_addr(pid, &addr); + if (s < 0 || bind(s, (struct sockaddr*)&addr, len) != 0 || listen(s, 64) != 0) { + fprintf(stderr, "[HOOK] ERROR: VMM-IPC fd server setup failed (errno=%d)\n", errno); + if (s >= 0) + close(s); + vmm_ipc_listen_fd = -1; + vmm_ipc_server_pid = pid; // don't retry-spam; exports in this process fail loudly + return false; + } + // Per-process token (re-derived after fork). /dev/urandom, with a clock^pid + // fallback - this is anti-accident (pid reuse), not a security boundary; + // same-uid access control is SO_PEERCRED above. + uint64_t tok = 0; + int ur = open("/dev/urandom", O_RDONLY | O_CLOEXEC); + if (ur >= 0) { + if (read(ur, &tok, sizeof(tok)) != (ssize_t)sizeof(tok)) + tok = 0; + close(ur); + } + if (tok == 0) { + struct timeval tv; + gettimeofday(&tv, nullptr); + tok = ((uint64_t)tv.tv_sec << 32) ^ (uint64_t)tv.tv_usec ^ ((uint64_t)pid << 16) ^ + 0x9e3779b97f4a7c15ULL; + } + vmm_ipc_token = tok; + vmm_ipc_listen_fd = s; + vmm_ipc_server_pid = pid; + std::thread(vmm_ipc_server_loop, s).detach(); + return true; +} + +// Importer side: fetch the fd for `original_ptr` from `exporter_pid`'s server. +// Returns a live fd in THIS process, or -1. +static int vmm_ipc_fetch_fd(pid_t exporter_pid, uint64_t token, CUdeviceptr original_ptr) { + int s = socket(AF_UNIX, SOCK_STREAM | SOCK_CLOEXEC, 0); + if (s < 0) + return -1; + vmm_ipc_set_timeouts(s); + struct sockaddr_un addr; + socklen_t len = vmm_ipc_socket_addr(exporter_pid, &addr); + if (connect(s, (struct sockaddr*)&addr, len) != 0) { + fprintf(stderr, + "[HOOK] ERROR: VMM-IPC connect to exporter pid %d failed (errno=%d) - exporter dead " + "or hook version mismatch\n", + (int)exporter_pid, errno); + close(s); + return -1; + } + VmmIpcRequest req = {(uint64_t)original_ptr, token}; + if (send(s, &req, sizeof(req), MSG_NOSIGNAL) != (ssize_t)sizeof(req)) { + close(s); + return -1; + } + char status = 1; + struct iovec iov = {&status, 1}; + char cbuf[CMSG_SPACE(sizeof(int))] = {}; + struct msghdr msg = {}; + msg.msg_iov = &iov; + msg.msg_iovlen = 1; + msg.msg_control = cbuf; + msg.msg_controllen = sizeof(cbuf); + ssize_t r = recvmsg(s, &msg, MSG_CMSG_CLOEXEC); + close(s); + if (r != 1 || status != 0) + return -1; + for (struct cmsghdr* c = CMSG_FIRSTHDR(&msg); c != nullptr; c = CMSG_NXTHDR(&msg, c)) { + if (c->cmsg_level == SOL_SOCKET && c->cmsg_type == SCM_RIGHTS) { + int fd = -1; + memcpy(&fd, CMSG_DATA(c), sizeof(int)); + return fd; + } + } + return -1; +} + +// Close and forget the parked exported fd for a VA (called from the free +// path). Without this, the fd keeps the freed allocation's physical pages +// alive and the server would hand importers a stale mapping. +static void vmm_ipc_invalidate_export(CUdeviceptr dptr) { + vmm_ipc_exported_fds.erase_if( + [dptr](const std::pair>& kv) { + if (kv.first != dptr) + return false; + close(kv.second.second); + return true; + }); +} // Hook for cuIpcGetMemHandle - intercept Driver API to support VMM allocations CUresult cuIpcGetMemHandle(CUipcMemHandle* pHandle, CUdeviceptr dptr) { @@ -2880,7 +3156,46 @@ CUresult cuIpcGetMemHandle(CUipcMemHandle* pHandle, CUdeviceptr dptr) { metadata = kv.second; }); - if (found && metadata.handle != 0) { + // Decide what to export: the allocation's own handle (SAVE-mode slow-path + // allocs), or the whole preallocated chunk plus an offset (LOAD-mode fast + // path carves, which have no individual handle). cuMemMap cannot map a + // sub-range of an imported handle, so chunk carves are shared by exporting + // the chunk handle; the importer maps the entire chunk and returns an + // interior pointer. + CUmemGenericAllocationHandle export_handle = 0; + CUdeviceptr export_key = dptr; // fd-registry key == blob lookup key + uint64_t chunk_base = 0; + uint64_t chunk_size = 0; + uint64_t chunk_generation = 0; + std::unique_lock preallocated_chunk_lock; + if (found) { + export_handle = metadata.handle; + if (export_handle == 0 && metadata.from_preallocation) { + // Keep the chunk alive and its record stable until the exported fd is + // parked below. free_preallocated_region takes the same mutex. + preallocated_chunk_lock = std::unique_lock(preallocated_chunks_mutex); + auto chunk_it = preallocated_chunks.upper_bound(dptr); + if (chunk_it != preallocated_chunks.begin()) { + --chunk_it; + CUdeviceptr candidate_base = chunk_it->first; + if (dptr >= candidate_base && (uint64_t)(dptr - candidate_base) < chunk_it->second.size) { + chunk_base = (uint64_t)candidate_base; + chunk_size = chunk_it->second.size; + chunk_generation = chunk_it->second.generation; + export_handle = chunk_it->second.handle; + export_key = candidate_base; + } + } + if (export_handle == 0) { + fprintf(stderr, + "[HOOK] ERROR: cuIpcGetMemHandle: carved ptr=0x%llx is outside all live " + "preallocated chunks\n", + (unsigned long long)dptr); + } + } + } + + if (export_handle != 0) { // This is a VMM allocation - export via shareable handle typedef CUresult (*cuMemExportToShareableHandle_t)( void*, CUmemGenericAllocationHandle, CUmemAllocationHandleType, unsigned long long); @@ -2892,32 +3207,81 @@ CUresult cuIpcGetMemHandle(CUipcMemHandle* pHandle, CUdeviceptr dptr) { return CUDA_ERROR_NOT_SUPPORTED; } - int fd = -1; - CUresult result = - export_func(&fd, metadata.handle, CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR, 0); + // The server also owns the per-process token packed into the blob, so it + // must exist before we hand out any handle. + if (!vmm_ipc_ensure_server()) { + return CUDA_ERROR_NOT_SUPPORTED; + } - if (result != CUDA_SUCCESS) { - fprintf(stderr, - "[HOOK] ERROR: cuMemExportToShareableHandle failed with error %d for ptr=0x%llx\n", - result, (unsigned long long)dptr); - return result; + // Export once per allocation and park the fd for the server thread. + // The cached entry is keyed by VA and validated against the allocation + // handle, so VA reuse after free/realloc re-exports instead of serving a + // stale fd (the free path also eagerly invalidates via + // vmm_ipc_invalidate_export). + int fd = -1; + vmm_ipc_exported_fds.visit( + export_key, + [&](const std::pair>& kv) { + if (kv.second.first == export_handle) + fd = kv.second.second; + }); + if (fd < 0) { + int new_fd = -1; + CUresult result = + export_func(&new_fd, export_handle, CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR, 0); + if (result != CUDA_SUCCESS) { + fprintf(stderr, + "[HOOK] ERROR: cuMemExportToShareableHandle failed with error %d for ptr=0x%llx\n", + result, (unsigned long long)dptr); + return result; + } + // Keep parked fds out of exec'd children. Plain fork still inherits + // them; vmm_ipc_ensure_server closes the inherited registry before the + // child starts its own server. SCM_RIGHTS transfer is unaffected. + fcntl(new_fd, F_SETFD, FD_CLOEXEC); + vmm_ipc_exported_fds.insert_or_visit( + std::make_pair(export_key, std::make_pair(export_handle, new_fd)), + [&](std::pair>& kv) { + // Entry exists: either a racing thread won (same handle - drop + // ours) or it is stale from a freed allocation (replace). + if (kv.second.first == export_handle) { + close(new_fd); + new_fd = kv.second.second; + } else { + close(kv.second.second); + kv.second = std::make_pair(export_handle, new_fd); + } + }); + fd = new_fd; } - // Pack our custom data into the handle structure - // CUipcMemHandle has 64 reserved bytes - we use them to store our info: - // - Magic marker (4 bytes) to identify VMM handles - // - File descriptor (4 bytes) - // - Original pointer address (8 bytes) - critical for CUDA graph replay! - // - Size (8 bytes) + // Pack our custom data into the handle structure. + // CUipcMemHandle has 64 reserved bytes: + // - Magic marker (4 bytes, "VMM2") + // - Exporter pid (4 bytes) - importer fetches the fd from this process's + // VMM-IPC socket via SCM_RIGHTS (a raw fd integer is not portable) + // - Original pointer address (8 bytes) - fd lookup key + placement hint + // - Aligned size (8 bytes) + // - Per-process token (8 bytes) - server rejects mismatches (pid reuse) + // - Chunk base + chunk size (8+8 bytes) - nonzero iff the pointer is a + // carve from the preallocated chunk; the importer then maps the whole + // chunk and returns base + (ptr - chunk_base) + // - Chunk generation (8 bytes) - distinguishes free/recreate at one base memset(pHandle, 0, sizeof(CUipcMemHandle)); + uint32_t pid_u32 = (uint32_t)getpid(); memcpy(pHandle->reserved, &VMM_IPC_MAGIC, sizeof(uint32_t)); - memcpy(pHandle->reserved + 4, &fd, sizeof(int)); + memcpy(pHandle->reserved + 4, &pid_u32, sizeof(uint32_t)); memcpy(pHandle->reserved + 8, &dptr, sizeof(CUdeviceptr)); memcpy(pHandle->reserved + 16, &metadata.size, sizeof(size_t)); + memcpy(pHandle->reserved + 24, &vmm_ipc_token, sizeof(uint64_t)); + memcpy(pHandle->reserved + 32, &chunk_base, sizeof(uint64_t)); + memcpy(pHandle->reserved + 40, &chunk_size, sizeof(uint64_t)); + memcpy(pHandle->reserved + 48, &chunk_generation, sizeof(uint64_t)); #ifdef HOOK_DEBUG - fprintf(stderr, "[HOOK] cuIpcGetMemHandle: VMM ptr=0x%llx, fd=%d, size=%zu\n", + fprintf(stderr, + "[HOOK] cuIpcGetMemHandle: VMM ptr=0x%llx, fd=%d (served via socket), size=%zu\n", (unsigned long long)dptr, fd, metadata.size); #endif return CUDA_SUCCESS; @@ -2942,19 +3306,75 @@ CUresult cuIpcOpenMemHandle(CUdeviceptr* pdptr, CUipcMemHandle handle, unsigned if (magic == VMM_IPC_MAGIC) { // This is a VMM IPC handle - extract the packed data - int fd; + uint32_t exporter_pid_u32; CUdeviceptr original_ptr; size_t size; + uint64_t token; + uint64_t chunk_base; + uint64_t chunk_size; + uint64_t chunk_generation; - memcpy(&fd, handle.reserved + 4, sizeof(int)); + memcpy(&exporter_pid_u32, handle.reserved + 4, sizeof(uint32_t)); memcpy(&original_ptr, handle.reserved + 8, sizeof(CUdeviceptr)); memcpy(&size, handle.reserved + 16, sizeof(size_t)); + memcpy(&token, handle.reserved + 24, sizeof(uint64_t)); + memcpy(&chunk_base, handle.reserved + 32, sizeof(uint64_t)); + memcpy(&chunk_size, handle.reserved + 40, sizeof(uint64_t)); + memcpy(&chunk_generation, handle.reserved + 48, sizeof(uint64_t)); + pid_t exporter_pid = (pid_t)exporter_pid_u32; + // For chunk carves the fd registry on the exporter is keyed by the chunk + // base, and what we import/map is the whole chunk. + CUdeviceptr fetch_key = (chunk_base != 0) ? (CUdeviceptr)chunk_base : original_ptr; + size_t map_size = (chunk_base != 0) ? (size_t)chunk_size : size; + + if (chunk_base != 0) { + // Fast path: peer chunk already mapped -> hand out an interior pointer. + std::lock_guard lock(vmm_ipc_chunk_mutex); + VmmIpcChunkKey key = std::make_tuple(exporter_pid, token, chunk_base, chunk_generation); + auto it = vmm_ipc_chunk_mappings.find(key); + if (it != vmm_ipc_chunk_mappings.end()) { + it->second.refcount++; + CUdeviceptr mapped = it->second.local_base + (original_ptr - (CUdeviceptr)chunk_base); + auto sub_it = vmm_ipc_chunk_subptrs.find(mapped); + if (sub_it == vmm_ipc_chunk_subptrs.end()) { + vmm_ipc_chunk_subptrs[mapped] = VmmIpcChunkSubptr{key, 1}; + } else if (sub_it->second.key == key) { + sub_it->second.open_count++; + } else { + it->second.refcount--; + fprintf(stderr, + "[HOOK] ERROR: VMM-IPC interior pointer collision at 0x%llx across chunks\n", + (unsigned long long)mapped); + return CUDA_ERROR_INVALID_VALUE; + } + *pdptr = mapped; + return CUDA_SUCCESS; + } + } #ifdef HOOK_DEBUG - fprintf(stderr, "[HOOK] cuIpcOpenMemHandle: VMM fd=%d, original_ptr=0x%llx, size=%zu\n", fd, - (unsigned long long)original_ptr, size); + fprintf(stderr, + "[HOOK] cuIpcOpenMemHandle: VMM exporter_pid=%d, original_ptr=0x%llx, size=%zu\n", + (int)exporter_pid, (unsigned long long)original_ptr, size); #endif + // Obtain a live fd in THIS process: dup from our own registry for + // same-process opens, SCM_RIGHTS fetch from the exporter otherwise. + int fd = -1; + if (exporter_pid == getpid() && token == vmm_ipc_token) { + vmm_ipc_exported_fds.visit( + fetch_key, + [&](const std::pair>& + kv) { fd = fcntl(kv.second.second, F_DUPFD_CLOEXEC, 0); }); + } else { + fd = vmm_ipc_fetch_fd(exporter_pid, token, fetch_key); + } + if (fd < 0) { + fprintf(stderr, "[HOOK] ERROR: VMM-IPC could not obtain fd for ptr=0x%llx from pid %d\n", + (unsigned long long)fetch_key, (int)exporter_pid); + return CUDA_ERROR_INVALID_VALUE; + } + // Import the allocation handle from file descriptor typedef CUresult (*cuMemImportFromShareableHandle_t)(CUmemGenericAllocationHandle*, void*, CUmemAllocationHandleType); @@ -2963,12 +3383,14 @@ CUresult cuIpcOpenMemHandle(CUdeviceptr* pdptr, CUipcMemHandle handle, unsigned if (import_func == nullptr) { fprintf(stderr, "[HOOK] ERROR: cuMemImportFromShareableHandle not found\n"); + close(fd); return CUDA_ERROR_NOT_SUPPORTED; } CUmemGenericAllocationHandle imported_handle; CUresult result = import_func(&imported_handle, (void*)(intptr_t)fd, CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR); + close(fd); // the driver does not take ownership of our fd copy if (result != CUDA_SUCCESS) { fprintf(stderr, "[HOOK] ERROR: cuMemImportFromShareableHandle failed with error %d\n", @@ -2976,16 +3398,58 @@ CUresult cuIpcOpenMemHandle(CUdeviceptr* pdptr, CUipcMemHandle handle, unsigned return result; } - // Map to the SAME address as original (critical for CUDA graph replay!) + typedef CUresult (*cuMemAddressReserve_t)(CUdeviceptr*, size_t, size_t, CUdeviceptr, + unsigned long long); + auto real_reserve_func = (cuMemAddressReserve_t)CUDA_DRIVER_CALL( + cuda_driver_entry_table, CUDA_ENTRY_cuMemAddressReserve); + typedef CUresult (*cuMemRelease_t)(CUmemGenericAllocationHandle); + auto mem_release_func = + (cuMemRelease_t)CUDA_DRIVER_CALL(cuda_driver_entry_table, CUDA_ENTRY_cuMemRelease); + typedef CUresult (*cuMemAddressFree_t)(CUdeviceptr, size_t); + auto real_address_free_func = + (cuMemAddressFree_t)CUDA_DRIVER_CALL(cuda_driver_entry_table, CUDA_ENTRY_cuMemAddressFree); + + // Reserve our own VA range for the peer mapping (real driver call, NOT the + // hooked wrapper - peer imports must never advance the deterministic + // cursor). The exporter's VA is passed as a hint: with per-rank disjoint + // regions it is honored and the mapping is address-stable; with today's + // shared-base symmetric layouts that range is occupied by our own region + // reservation, the driver picks elsewhere, and the caller gets a valid but + // relocated peer pointer. That is correct for table-indirect consumers + // (DeepEP buffer_ptrs_gpu, custom-allreduce RankData contents): peer VAs + // are never baked into captured graph nodes on those paths. + // First preference: the exporter's own VA (address-stable whenever it is + // locally free, e.g. future per-rank disjoint-base configs). + CUdeviceptr mapped_ptr = 0; + result = real_reserve_func(&mapped_ptr, map_size, 0, fetch_key, 0); + if (result == CUDA_SUCCESS && mapped_ptr != fetch_key) { + // Hint not honored. Do NOT keep the driver's fallback choice - it tends + // to sit immediately after existing reservations, i.e. on the NVSHMEM + // heap hint. Re-reserve inside the dedicated import zone. + real_address_free_func(mapped_ptr, map_size); + uint64_t zone_hint = + vmm_ipc_import_hint.fetch_add((map_size + ((1ULL << 30) - 1)) & ~((1ULL << 30) - 1)); + mapped_ptr = 0; + result = real_reserve_func(&mapped_ptr, map_size, 0, (CUdeviceptr)zone_hint, 0); + } + if (result != CUDA_SUCCESS) { + fprintf(stderr, "[HOOK] ERROR: cuMemAddressReserve for IPC import failed with error %d\n", + result); + mem_release_func(imported_handle); + return result; + } + typedef CUresult (*cuMemMap_t)(CUdeviceptr, size_t, size_t, CUmemGenericAllocationHandle, unsigned long long); auto map_func = (cuMemMap_t)CUDA_DRIVER_CALL(cuda_driver_entry_table, CUDA_ENTRY_cuMemMap); - result = map_func(original_ptr, size, 0, imported_handle, 0); + result = map_func(mapped_ptr, map_size, 0, imported_handle, 0); if (result != CUDA_SUCCESS) { fprintf(stderr, "[HOOK] ERROR: cuMemMap for IPC failed with error %d at addr=0x%llx size=%zu\n", - result, (unsigned long long)original_ptr, size); + result, (unsigned long long)mapped_ptr, map_size); + real_address_free_func(mapped_ptr, map_size); + mem_release_func(imported_handle); return result; } @@ -3005,27 +3469,87 @@ CUresult cuIpcOpenMemHandle(CUdeviceptr* pdptr, CUipcMemHandle handle, unsigned auto set_access_func = (cuMemSetAccess_t)CUDA_DRIVER_CALL(cuda_driver_entry_table, CUDA_ENTRY_cuMemSetAccess); - result = set_access_func(original_ptr, size, &accessDesc, 1); + result = set_access_func(mapped_ptr, map_size, &accessDesc, 1); if (result != CUDA_SUCCESS) { fprintf(stderr, "[HOOK] ERROR: cuMemSetAccess for IPC failed with error %d\n", result); + typedef CUresult (*cuMemUnmap_t)(CUdeviceptr, size_t); + auto mem_unmap_func = + (cuMemUnmap_t)CUDA_DRIVER_CALL(cuda_driver_entry_table, CUDA_ENTRY_cuMemUnmap); + mem_unmap_func(mapped_ptr, map_size); + real_address_free_func(mapped_ptr, map_size); + mem_release_func(imported_handle); return result; } - *pdptr = original_ptr; + if (mapped_ptr != fetch_key) { + fprintf(stderr, + "[HOOK] INFO: VMM-IPC import relocated: exporter VA 0x%llx -> local VA 0x%llx " + "(expected with shared region bases; peer tables must be refreshed by the caller)\n", + (unsigned long long)fetch_key, (unsigned long long)mapped_ptr); + } + + if (chunk_base != 0) { + // Register the whole-chunk mapping and hand out the interior pointer. + CUdeviceptr interior = mapped_ptr + (original_ptr - (CUdeviceptr)chunk_base); + std::lock_guard lock(vmm_ipc_chunk_mutex); + VmmIpcChunkKey key = std::make_tuple(exporter_pid, token, chunk_base, chunk_generation); + auto it = vmm_ipc_chunk_mappings.find(key); + if (it != vmm_ipc_chunk_mappings.end()) { + // Lost a race to a concurrent open of the same peer chunk: keep the + // winner's mapping, drop ours. + typedef CUresult (*cuMemUnmap_t)(CUdeviceptr, size_t); + auto mem_unmap_func = + (cuMemUnmap_t)CUDA_DRIVER_CALL(cuda_driver_entry_table, CUDA_ENTRY_cuMemUnmap); + mem_unmap_func(mapped_ptr, map_size); + real_address_free_func(mapped_ptr, map_size); + mem_release_func(imported_handle); + it->second.refcount++; + interior = it->second.local_base + (original_ptr - (CUdeviceptr)chunk_base); + } else { + vmm_ipc_chunk_mappings[key] = VmmIpcChunkMapping{mapped_ptr, map_size, imported_handle, 1}; + } + auto sub_it = vmm_ipc_chunk_subptrs.find(interior); + if (sub_it == vmm_ipc_chunk_subptrs.end()) { + vmm_ipc_chunk_subptrs[interior] = VmmIpcChunkSubptr{key, 1}; + } else if (sub_it->second.key == key) { + sub_it->second.open_count++; + } else { + auto mapping_it = vmm_ipc_chunk_mappings.find(key); + if (mapping_it != vmm_ipc_chunk_mappings.end() && --mapping_it->second.refcount == 0) { + typedef CUresult (*cuMemUnmap_t)(CUdeviceptr, size_t); + auto mem_unmap_func = + (cuMemUnmap_t)CUDA_DRIVER_CALL(cuda_driver_entry_table, CUDA_ENTRY_cuMemUnmap); + mem_unmap_func(mapping_it->second.local_base, mapping_it->second.size); + real_address_free_func(mapping_it->second.local_base, mapping_it->second.size); + mem_release_func(mapping_it->second.handle); + vmm_ipc_chunk_mappings.erase(mapping_it); + } + fprintf(stderr, + "[HOOK] ERROR: VMM-IPC interior pointer collision at 0x%llx across chunks\n", + (unsigned long long)interior); + return CUDA_ERROR_INVALID_VALUE; + } + *pdptr = interior; + return CUDA_SUCCESS; + } + + *pdptr = mapped_ptr; - // Track this imported allocation + // Track this imported allocation (close path unmaps, releases, and frees + // the reservation we created above) AllocMetadata alloc_metadata; - alloc_metadata.ptr = original_ptr; + alloc_metadata.ptr = mapped_ptr; alloc_metadata.size = size; alloc_metadata.handle = imported_handle; alloc_metadata.region_base = 0; // region_base == 0 indicates IPC-imported alloc_metadata.from_preallocation = false; - global_alloc_metadata.emplace(original_ptr, alloc_metadata); + global_alloc_metadata.emplace(mapped_ptr, alloc_metadata); #ifdef HOOK_DEBUG fprintf(stderr, - "[HOOK] cuIpcOpenMemHandle: Successfully mapped VMM fd=%d -> ptr=0x%llx, size=%zu\n", - fd, (unsigned long long)*pdptr, size); + "[HOOK] cuIpcOpenMemHandle: Successfully mapped exporter ptr=0x%llx -> ptr=0x%llx, " + "size=%zu\n", + (unsigned long long)original_ptr, (unsigned long long)*pdptr, size); #endif return CUDA_SUCCESS; } @@ -3043,6 +3567,36 @@ CUresult cuIpcCloseMemHandle(CUdeviceptr dptr) { auto real_func = (cuIpcCloseMemHandle_t)CUDA_DRIVER_CALL(cuda_driver_entry_table, CUDA_ENTRY_cuIpcCloseMemHandle); + // Interior pointer into an imported peer chunk? Refcounted: the chunk + // mapping is torn down when its last interior pointer closes. + { + std::lock_guard lock(vmm_ipc_chunk_mutex); + auto sub_it = vmm_ipc_chunk_subptrs.find(dptr); + if (sub_it != vmm_ipc_chunk_subptrs.end()) { + VmmIpcChunkKey key = sub_it->second.key; + if (--sub_it->second.open_count == 0) { + vmm_ipc_chunk_subptrs.erase(sub_it); + } + auto it = vmm_ipc_chunk_mappings.find(key); + if (it != vmm_ipc_chunk_mappings.end() && --it->second.refcount == 0) { + typedef CUresult (*cuMemUnmap_t)(CUdeviceptr, size_t); + auto mem_unmap_func = + (cuMemUnmap_t)CUDA_DRIVER_CALL(cuda_driver_entry_table, CUDA_ENTRY_cuMemUnmap); + typedef CUresult (*cuMemRelease_t)(CUmemGenericAllocationHandle); + auto mem_release_func = + (cuMemRelease_t)CUDA_DRIVER_CALL(cuda_driver_entry_table, CUDA_ENTRY_cuMemRelease); + typedef CUresult (*cuMemAddressFree_t)(CUdeviceptr, size_t); + auto addr_free_func = (cuMemAddressFree_t)CUDA_DRIVER_CALL(cuda_driver_entry_table, + CUDA_ENTRY_cuMemAddressFree); + mem_unmap_func(it->second.local_base, it->second.size); + mem_release_func(it->second.handle); + addr_free_func(it->second.local_base, it->second.size); + vmm_ipc_chunk_mappings.erase(it); + } + return CUDA_SUCCESS; + } + } + AllocMetadata metadata; bool found = false; @@ -3061,8 +3615,14 @@ CUresult cuIpcCloseMemHandle(CUdeviceptr dptr) { auto mem_release_func = (cuMemRelease_t)CUDA_DRIVER_CALL(cuda_driver_entry_table, CUDA_ENTRY_cuMemRelease); + typedef CUresult (*cuMemAddressFree_t)(CUdeviceptr, size_t); + auto real_address_free_func = + (cuMemAddressFree_t)CUDA_DRIVER_CALL(cuda_driver_entry_table, CUDA_ENTRY_cuMemAddressFree); + mem_unmap_func(dptr, metadata.size); mem_release_func(metadata.handle); + // The import path owns a dedicated VA reservation for this mapping. + real_address_free_func(dptr, metadata.size); global_alloc_metadata.erase_if( [dptr](const std::pair& kv) { return kv.first == dptr; }); @@ -3270,6 +3830,11 @@ bool preallocate_region(size_t size) { tls_storage.preallocated_start_addr = target_addr; tls_storage.preallocated_end_addr = target_addr + aligned_size; tls_storage.has_preallocation = true; + { + std::lock_guard lock(preallocated_chunks_mutex); + uint64_t generation = next_preallocated_chunk_generation.fetch_add(1); + preallocated_chunks[target_addr] = PreallocatedChunk{aligned_size, allocHandle, generation}; + } // Note: we do NOT advance current_alloc_base_addr here. // The alloc calls will advance it as they consume the preallocated memory. @@ -3292,6 +3857,12 @@ void free_preallocated_region() { size_t preallocated_size = tls_storage.preallocated_end_addr - tls_storage.preallocated_start_addr; + { + std::lock_guard lock(preallocated_chunks_mutex); + preallocated_chunks.erase(tls_storage.preallocated_start_addr); + } + vmm_ipc_invalidate_export((CUdeviceptr)tls_storage.preallocated_start_addr); + mem_unmap_func(tls_storage.preallocated_start_addr, preallocated_size); mem_release_func(tls_storage.preallocated_handle); From 6394fc91dc3d43e12ecf3ae3d0720e0d35e0c58e Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Thu, 23 Jul 2026 23:01:22 +0000 Subject: [PATCH 19/21] test(sglang): cover validated EP4 topology Co-authored-by: Rahul Chalamala --- tests/test_sglang_cuda_graph.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/tests/test_sglang_cuda_graph.py b/tests/test_sglang_cuda_graph.py index 798bf280..633fdbb7 100644 --- a/tests/test_sglang_cuda_graph.py +++ b/tests/test_sglang_cuda_graph.py @@ -217,6 +217,14 @@ def _server_args(**overrides): attention_backend="fa3", moe_a2a_backend="deepep", ), + _server_args( + tp_size=4, + dp_size=4, + ep_size=4, + enable_dp_attention=True, + attention_backend="fa3", + moe_a2a_backend="deepep", + ), ], ) def test_configuration_accepts_supported_full_decode_modes(monkeypatch, server_args): From 973e9a0f237f0693d45e437fa7760523e13adcdd Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Thu, 23 Jul 2026 23:51:34 +0000 Subject: [PATCH 20/21] docs(sglang): record four-GPU validation Co-authored-by: Rahul Chalamala --- README.md | 14 ++++-- ROADMAP.md | 4 +- docs/overview.md | 17 +++++-- docs/sglang/hooks.md | 28 ++++++++++- docs/sglang/memory-consistency.md | 26 ++++++++-- docs/sglang/overview.md | 28 ++++++++--- docs/sglang/save-load-workflow.md | 18 ++++++- recipe/sglang/README.md | 81 ++++++++++++++++++++++++++----- 8 files changed, 177 insertions(+), 39 deletions(-) diff --git a/README.md b/README.md index f1b54bff..32f7abd7 100644 --- a/README.md +++ b/README.md @@ -82,18 +82,24 @@ Foundry ships engine integrations under `foundry/python/foundry/integration/`. P | Engine | Single GPU | DP | TP | EP | |---|:---:|:---:|:---:|:---:| | vLLM | ✅ | ✅ | 🚧 | ✅ | -| SGLang `0.5.15.post1` | ✅ | ✅ | Unsupported | ✅ | +| SGLang `0.5.15.post1` | ✅ | ✅ | ✅[^sglang-tp] | ✅ | | TensorRT-LLM | 🚧 | 🚧 | 🚧 | 🚧 | ✅ validated end-to-end (SAVE → LOAD → query)  ·  🚧 not yet +[^sglang-tp]: SGLang conventional tensor parallelism is validated only for + dense TP=2 × DP=2 with PP=1 and EP=1. All other conventional TP topologies + remain unsupported. + The SGLang integration is latest-only for upstream `v0.5.15.post1` (`0b3bb0c`), CUDA 13.0, and PyTorch 2.11/cu130 on H100. Full decode with prefill graphs disabled is validated for single-GPU FlashInfer and regular -DP=2 FlashInfer on Qwen3-1.7B (52 graphs/rank), plus EP=2 DP-attention with +DP=2 FlashInfer on Qwen3-1.7B (52 graphs/rank), exact dense TP=2 × DP=2 +FlashInfer on Qwen3-8B (36 graphs/rank), and DP-attention EP=2 and EP=4 with DeepEP low-latency, DeepGEMM, and FA3 on Qwen3-30B-A3B-FP8 (20 graphs/rank). -All three passed baseline/SAVE/fresh LOAD/query with text/token parity and no -recapture. TP without DP-attention remains unsupported. Existing SGLang +All passed baseline/SAVE/fresh LOAD/query with text/token parity and no +recapture. EP=4 uses DP-attention/EP rather than conventional TP attention, +even though SGLang's process topology sets `tp_size=4`. Existing SGLang archives are incompatible and require a fresh SAVE. See [`recipe/sglang/README.md`](recipe/sglang/README.md). diff --git a/ROADMAP.md b/ROADMAP.md index aa139d65..a68941f7 100644 --- a/ROADMAP.md +++ b/ROADMAP.md @@ -30,8 +30,8 @@ - [x] Quantized MoE with DeepGemm - [x] Drop-in integration layer for CUDA graph persistence in vLLM - [x] Sync with latest SGLang release - - [ ] EP on SGLang - - [ ] TP on SGLang + - [x] EP on SGLang (DP-attention + DeepEP, validated at EP=2 and EP=4) + - [x] Exact dense TP=2 × DP=2 on SGLang (all other TP topologies unsupported) ## Stage 5: Disaggregated and Large-Scale Serving diff --git a/docs/overview.md b/docs/overview.md index 39c572d3..86253dff 100644 --- a/docs/overview.md +++ b/docs/overview.md @@ -37,16 +37,23 @@ hooks and archive metadata are intentionally separate. | Attention preparation | vLLM-specific compile/capture path | Separate out-of-graph planning and in-graph GPU metadata work | | Profile behavior | Full profile forward; two-pass SAVE recipe | Pool-memory measurement only; one SAVE process | | Process model | Uniprocess, multiprocessing, or Ray executors | Multiprocessing `spawn`; explicit hook install in scheduler and DP-controller children | -| Parallel scope | Single GPU, DP, and documented EP configuration | Validated on H100: single GPU FlashInfer, regular DP=2 FlashInfer, and EP=2 DP-attention + DeepEP low-latency + FA3 | +| Parallel scope | Single GPU, DP, and documented EP configuration | Validated on H100: single GPU FlashInfer, regular DP=2 FlashInfer, exact dense TP=2 × DP=2 FlashInfer, and DP-attention EP=2/EP=4 + DeepEP low-latency + FA3 | SGLang validation used upstream commit `0b3bb0c`, CUDA 13.0, and PyTorch 2.11/cu130. Qwen3-1.7B restored 52 graphs for single GPU and 52 per rank for -DP=2; Qwen3-30B-A3B-FP8 restored 20 graphs per rank for EP=2. Every mode passed +DP=2. Dense Qwen3-8B TP=2 × DP=2 restored 36 graphs per rank across rank +workspaces 0–3 at the identical `69822578688` final offset, with parity on +explicitly routed DP0 and DP1 queries. Qwen3-30B-A3B-FP8 restored 20 graphs per +rank for both EP=2 and EP=4; EP=4 restored 416 event nodes and showed +four-rank NVSHMEM/bootstrap/adapter restoration. Every mode passed baseline/SAVE/fresh LOAD/query with text/token parity and no recapture. -For SGLang, TP without DP-attention, PP, speculative decoding, LoRA/PDMux, -dLLM, hidden-state capture, memory-saver graphs, and prefill graphs are not -supported. +For SGLang, every conventional TP topology except exact dense TP=2 × DP=2 +with PP=1 and EP=1 is unsupported. Custom all-reduce, NCCL CUMEM/NVLS, +pipeline parallelism, speculative decoding, LoRA/PDMux, dLLM, hidden-state +capture, memory-saver graphs, and prefill graphs are not supported. EP=4 sets +`tp_size=4` as part of SGLang's DP-attention process topology; it is expert +parallelism, not conventional TP attention. ## Archive authority diff --git a/docs/sglang/hooks.md b/docs/sglang/hooks.md index bed7ab78..f04178f7 100644 --- a/docs/sglang/hooks.md +++ b/docs/sglang/hooks.md @@ -26,10 +26,20 @@ Before patching, Foundry requires: - resolved decode backend `full`; - resolved prefill backend `disabled`; - FlashInfer for single-GPU and regular-DP decode; +- for conventional tensor parallelism, exactly + `tp_size=2, dp_size=2, pp_size=1, ep_size=1` with FlashInfer, + `--disable-custom-all-reduce`, `NCCL_CUMEM_ENABLE=0`, and + `NCCL_NVLS_ENABLE=0`; - FA3 plus DeepEP low-latency for DP-attention. -It rejects unsupported TP, PP, speculative, LoRA/PDMux, dLLM, -hidden-state-output, memory-saver, and debug-graph configurations. +It rejects every other conventional TP topology, PP, speculative, LoRA/PDMux, +dLLM, hidden-state-output, memory-saver, and debug-graph configurations. It +also rejects custom/symmetric-memory collectives and NCCL NVLS for the exact +TP=2 × DP=2 mode. + +The validated DeepEP scope is DP-attention at EP=2 and EP=4. EP=4 carries +`tp_size=4` in SGLang's process topology, but this branch follows the +DP-attention path and is not conventional TP attention. ## Distributed initialization @@ -48,6 +58,20 @@ After upstream initialization and communicator warmup, it advances the VMM cursor to the configured scratch boundary. Nondeterministic startup allocations remain below that boundary. +### Dense TP VMM-IPC bridge + +The exact TP=2 × DP=2 topology shares graph-referenced VMM allocations across +scheduler processes through CUDA IPC. Foundry's driver hook exports the whole +physical VMM chunk, records the requested tensor's interior offset, and sends +the live POSIX file descriptor over a same-UID abstract Unix socket with +`SCM_RIGHTS`. Exporter PID/token and chunk-generation checks reject stale +handles; an importer can reserve a relocated VA when the exporter's address +collides locally. + +Four-H100 validation mapped DP0/TP0, DP0/TP1, DP1/TP0, and DP1/TP1 to +workspaces `rank_0` through `rank_3`. All ranks restored 36 graphs at final +offset `69822578688` without VMM-IPC, CUDA, or preallocation errors. + ## Memory-pool initialization On SAVE, upstream `ModelRunnerKVCacheMixin.init_memory_pool` runs unchanged. diff --git a/docs/sglang/memory-consistency.md b/docs/sglang/memory-consistency.md index 59379d1a..d51d161c 100644 --- a/docs/sglang/memory-consistency.md +++ b/docs/sglang/memory-consistency.md @@ -94,14 +94,29 @@ self._apply_memory_pool_config(self.memory_pool_config) Without the drain, subsequent attention initialization can hit cached segments on LOAD but request new VMM memory on SAVE, or the reverse. -## 6. Bootstrap DeepEP outside graph capture +## 6. Transfer dense-TP VMM allocations by live file descriptor + +The exact dense TP=2 × DP=2 path uses CUDA IPC between scheduler processes. +A raw exporter file-descriptor integer is meaningless in another process. +Foundry therefore exports the whole preallocated VMM chunk and transfers its +live file descriptor through a same-UID abstract Unix socket with +`SCM_RIGHTS`. The IPC handle carries exporter PID/token, chunk generation, and +the allocation's interior offset; imports use a relocated reservation if the +original VA collides. + +Qwen3-8B validation mapped DP0/TP0, DP0/TP1, DP1/TP0, and DP1/TP1 to rank +workspaces 0–3. All four ranks saved and restored 36 graphs at final offset +`69822578688`; explicitly routed DP0/DP1 queries had text/token parity and LOAD +had no IPC, CUDA, or preallocation errors. + +## 7. Bootstrap DeepEP outside graph capture DeepEP constructs its singleton NVSHMEM buffer lazily, and DeepGEMM may compile shape-specific kernels lazily. Creating either inside stream capture is invalid. -For the H100-validated EP=2 DP-attention + DeepEP low-latency + DeepGEMM + -FA3 configuration: +For the H100-validated EP=2 and EP=4 DP-attention + DeepEP low-latency + +DeepGEMM + FA3 configurations: - SAVE runs a dedicated eager orchestration pass to initialize lazy state; - SAVE and LOAD bootstrap the DeepEP buffer before capture/load; @@ -110,6 +125,11 @@ FA3 configuration: Qwen3-30B-A3B-FP8 validation restored 20 graphs per rank with text/token parity, no recapture, and DeepEP/NVSHMEM/adapter/event restoration evidence. +EP=4 restored 416 event nodes and confirmed bootstrap/adapter behavior on all +four ranks. + +EP=4 uses DP-attention and expert parallelism. Its SGLang process topology has +`tp_size=4`, but it is not conventional TP attention. The DeepEP token limit, chunked-prefill size, model, topology, and environment must remain identical between SAVE and LOAD. diff --git a/docs/sglang/overview.md b/docs/sglang/overview.md index 904c146d..316bd000 100644 --- a/docs/sglang/overview.md +++ b/docs/sglang/overview.md @@ -10,11 +10,15 @@ CUDA 13.0 and PyTorch 2.11/cu130. Prefill CUDA graphs are disabled. |---|---|:---:| | Single GPU | FlashInfer | Validated on H100 | | Regular data parallelism | One full replica per rank, FlashInfer | Validated at DP=2 on H100 | -| DP-attention expert parallelism | DeepEP low-latency, DeepGEMM, and FA3 | Validated at EP=2 on H100 | +| Dense tensor × data parallelism | FlashInfer | Validated only at TP=2 × DP=2, PP=1, EP=1 on H100 | +| DP-attention expert parallelism | DeepEP low-latency, DeepGEMM, and FA3 | Validated at EP=2 and EP=4 on H100 | -Tensor parallelism without DP-attention, pipeline parallelism, speculative -decoding, LoRA/PDMux, dLLM, hidden-state capture, and memory-saver graphs are -rejected at startup. Foundry does not claim prefill-graph support. +Every conventional TP topology except exact dense TP=2 × DP=2 with PP=1 and +EP=1 is rejected at startup. The validated dense topology also requires custom +all-reduce off and `NCCL_CUMEM_ENABLE=0`/`NCCL_NVLS_ENABLE=0`. Pipeline +parallelism, speculative decoding, LoRA/PDMux, dLLM, hidden-state capture, and +memory-saver graphs remain unsupported. Foundry does not claim prefill-graph +support. Validation results: @@ -22,10 +26,18 @@ Validation results: text/token parity, no recapture, and 52 restored graphs. - Regular DP=2 Qwen3-1.7B passed on both FlashInfer ranks with text/token parity, no recapture, and 52 restored graphs per rank. -- EP=2 Qwen3-30B-A3B-FP8 passed with DP-attention, DeepEP low-latency, - DeepGEMM, and FA3, with text/token parity, no recapture, and 20 restored - graphs per rank. Logs confirmed DeepEP/NVSHMEM initialization, adapter - restoration, and graph event restoration. +- Dense TP=2 × DP=2 Qwen3-8B passed on four H100s with explicitly routed DP0 + and DP1 parity, no recapture, rank workspaces 0–3, 36 restored graphs per + rank, and identical final offsets of `69822578688`. No VMM-IPC, CUDA, or + preallocation errors occurred. +- EP=2 and EP=4 Qwen3-30B-A3B-FP8 passed with DP-attention, DeepEP + low-latency, DeepGEMM, and FA3, with text/token parity, no recapture, and 20 + restored graphs per rank. EP=4 restored 416 event nodes and showed + four-rank DeepEP/NVSHMEM bootstrap and adapter evidence. + +EP=4 is not conventional TP attention. SGLang sets `tp_size=dp_size=ep_size=4` +for this DP-attention/EP process topology, while each rank owns its attention +work and DeepEP handles expert all-to-all. ## Upstream execution model diff --git a/docs/sglang/save-load-workflow.md b/docs/sglang/save-load-workflow.md index b6917be5..24a71220 100644 --- a/docs/sglang/save-load-workflow.md +++ b/docs/sglang/save-load-workflow.md @@ -45,6 +45,12 @@ scratch_space_size = "1024MB" `foundry_load.toml` uses the same values with `mode = "load"`. +The exact dense TP=2 × DP=2 recipe uses dedicated +`foundry_tp2_dp2_save.toml`/`foundry_tp2_dp2_load.toml` files and workspace +`foundry_archive_sglang_qwen3_8b_tp2_dp2`. Its launcher fixes PP=1, EP=1, +FlashInfer, full decode, disabled prefill, disabled custom all-reduce, and +`NCCL_CUMEM_ENABLE=0`/`NCCL_NVLS_ENABLE=0` for baseline, SAVE, and LOAD. + ## Verify package metadata ```bash @@ -92,11 +98,19 @@ validated configurations. Validation used upstream SGLang `v0.5.15.post1` (`0b3bb0c`), CUDA 13.0, PyTorch 2.11/cu130, and H100. Single-GPU and DP=2 Qwen3-1.7B FlashInfer -restored 52 graphs per applicable rank; EP=2 Qwen3-30B-A3B-FP8 with +restored 52 graphs per applicable rank. Dense TP=2 × DP=2 Qwen3-8B restored +36 graphs per rank across rank workspaces 0–3 at identical final offsets of +`69822578688`; explicitly routed DP0 and DP1 queries matched and no IPC, CUDA, +or preallocation errors occurred. EP=2 and EP=4 Qwen3-30B-A3B-FP8 with DP-attention, DeepEP low-latency, DeepGEMM, and FA3 restored 20 graphs per -rank. Baseline/SAVE/fresh LOAD/query passed with text/token parity and no +rank; EP=4 restored 416 event nodes with four-rank NVSHMEM/bootstrap/adapter +evidence. Baseline/SAVE/fresh LOAD/query passed with text/token parity and no recapture in every configuration. +Although SGLang sets `tp_size=4` for EP=4, that mode uses DP-attention and +expert parallelism rather than conventional TP attention. Conventional TP is +validated only for exact dense TP=2 × DP=2 with PP=1 and EP=1. + ## Archive files ```text diff --git a/recipe/sglang/README.md b/recipe/sglang/README.md index fca9fdc6..d638d9af 100644 --- a/recipe/sglang/README.md +++ b/recipe/sglang/README.md @@ -8,16 +8,23 @@ These recipes persist and restore full decode CUDA graphs for exactly SGLang |---|---| | `serve_qwen3-mini.sh` | Single GPU, Qwen3-1.7B, FlashInfer | | `serve_qwen3-1.7b_dp.sh` | Regular data parallelism, one full replica per GPU, FlashInfer | -| `serve_qwen3-30ba3bfp8_ep.sh` | DP-attention, DeepEP low-latency, FA3, Qwen3-30B-A3B-FP8 | +| `serve_qwen3-8b_tp2_dp2.sh` | Exact dense TP=2 × DP=2, PP=1, EP=1, Qwen3-8B, FlashInfer | +| `serve_qwen3-30ba3bfp8_ep.sh` | DP-attention EP=2/EP=4, DeepEP low-latency, DeepGEMM, FA3, Qwen3-30B-A3B-FP8 | Single GPU and DP=2 passed baseline/SAVE/fresh LOAD/query with text/token -parity, no recapture, and 52 graphs per applicable rank. EP=2 passed the same -checks with 20 graphs per rank plus DeepEP/NVSHMEM/adapter/event restoration -evidence. The EP environment used DeepGEMM and FA3. - -Tensor parallelism without DP-attention, pipeline parallelism, speculative +parity, no recapture, and 52 graphs per applicable rank. Dense TP=2 × DP=2 +passed on four H100s with explicitly routed DP0/DP1 parity, rank workspaces +0–3, 36 graphs per rank, identical final offsets of `69822578688`, and no +IPC/CUDA/preallocation errors. EP=2 and EP=4 passed with 20 graphs per rank +plus DeepEP/NVSHMEM/bootstrap/adapter/event restoration evidence; EP=4 +restored 416 event nodes. + +The dense TP recipe has dedicated `foundry_tp2_dp2_save.toml` and +`foundry_tp2_dp2_load.toml` files. Every other conventional TP topology, +custom all-reduce, NCCL CUMEM/NVLS, pipeline parallelism, speculative decoding, LoRA/PDMux, dLLM, hidden-state capture, and memory-saver graphs are -not supported. +not supported. EP=4 is DP-attention/expert parallelism, not conventional TP +attention, even though SGLang sets `tp_size=4` in that process topology. ## Requirements @@ -156,20 +163,65 @@ CUDA_VISIBLE_DEVICES=0,1 bash serve_qwen3-1.7b_dp.sh 2 --save CUDA_VISIBLE_DEVICES=0,1 bash serve_qwen3-1.7b_dp.sh 2 --load ``` -DP-attention with DeepEP: +Exact dense TP=2 × DP=2: + +```bash +# 1. Baseline. After startup, run the routed queries below, then stop it. +CUDA_VISIBLE_DEVICES=0,1,2,3 bash serve_qwen3-8b_tp2_dp2.sh + +# 2. Fresh SAVE using the dedicated workspace/config. +rm -rf foundry_archive_sglang_qwen3_8b_tp2_dp2 +CUDA_VISIBLE_DEVICES=0,1,2,3 bash serve_qwen3-8b_tp2_dp2.sh --save +# After startup, run the routed queries below, then stop SAVE. + +# 3. Fresh-process LOAD. +CUDA_VISIBLE_DEVICES=0,1,2,3 bash serve_qwen3-8b_tp2_dp2.sh --load + +# 4. Run after each ready phase; compare DP0 and DP1 across all three. +for dp_rank in 0 1; do + curl -s http://127.0.0.1:12000/generate \ + -H 'Content-Type: application/json' \ + -d "{\"text\":\"The capital of France is\",\"sampling_params\":{\"temperature\":0,\"max_new_tokens\":12},\"routed_dp_rank\":${dp_rank}}" +done +``` + +The TP launcher fixes `--tp-size 2 --dp-size 2 --pp-size 1 --ep-size 1`, +`--attention-backend flashinfer`, full decode, disabled prefill, +`--disable-custom-all-reduce`, `NCCL_CUMEM_ENABLE=0`, and +`NCCL_NVLS_ENABLE=0` in baseline, SAVE, and LOAD. Do not remove or generalize +these safety constraints. + +SGLang uses CUDA IPC to share tensors between the TP scheduler processes. +Foundry's VMM-IPC bridge exports the whole physical VMM chunk, preserves each +allocation's interior offset, and transfers the live file descriptor over a +same-UID abstract Unix socket with `SCM_RIGHTS`; it never treats a raw +file-descriptor integer as portable across processes. The validated mapping is +DP0/TP0 → `rank_0`, DP0/TP1 → `rank_1`, DP1/TP0 → `rank_2`, and DP1/TP1 → +`rank_3`. Each rank saved and restored 36 graphs at final offset +`69822578688`. + +DP-attention with DeepEP (EP=2 or EP=4): ```bash +EP_SIZE=4 +CUDA_VISIBLE_DEVICES=0,1,2,3 bash serve_qwen3-30ba3bfp8_ep.sh "$EP_SIZE" +# Stop the baseline after query. rm -rf foundry_archive -CUDA_VISIBLE_DEVICES=0,1 bash serve_qwen3-30ba3bfp8_ep.sh 2 --save +CUDA_VISIBLE_DEVICES=0,1,2,3 bash serve_qwen3-30ba3bfp8_ep.sh "$EP_SIZE" --save # Stop SAVE after startup. -CUDA_VISIBLE_DEVICES=0,1 bash serve_qwen3-30ba3bfp8_ep.sh 2 --load +CUDA_VISIBLE_DEVICES=0,1,2,3 bash serve_qwen3-30ba3bfp8_ep.sh "$EP_SIZE" --load curl -s http://127.0.0.1:12000/v1/completions \ -H 'Content-Type: application/json' \ -d '{"model":"Qwen/Qwen3-30B-A3B-FP8","prompt":"The capital of France is","max_tokens":12,"temperature":0}' ``` -The multi-GPU scripts set `NCCL_CUMEM_ENABLE=0` and -`NCCL_NVLS_ENABLE=0` for Foundry runs. The DeepEP script also keeps +Set `EP_SIZE=2` and select two GPUs for the validated EP=2 variant. EP=4's +`tp_size=4` is part of SGLang's DP-attention/EP process topology and does not +mean conventional TP attention. + +The regular-DP and EP scripts set `NCCL_CUMEM_ENABLE=0` and +`NCCL_NVLS_ENABLE=0` for Foundry runs; the dense TP launcher sets both in all +three phases. The DeepEP script also keeps `SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK=256` and `--chunked-prefill-size 256` identical on SAVE and LOAD. @@ -199,7 +251,9 @@ records capture order, phase, filename, output kind, and the complete not derive graph identity from them. `graph_manifest.json` remains the Foundry topology/template manifest used by the batched native loader. -For DP and EP, each rank owns a separate `rank_/` directory. +For regular DP, exact TP=2 × DP=2, and EP, each rank owns a separate +`rank_/` directory. The dense TP recipe writes the same layout under its +dedicated `foundry_archive_sglang_qwen3_8b_tp2_dp2` root. ## Troubleshooting @@ -207,6 +261,7 @@ For DP and EP, each rank owns a separate `rank_/` directory. |---|---| | Package version is `unknown` or not `0.5.15.post1` | Install `sglang/python` with pip; do not rely on a source-path-only checkout. | | Archive schema, version, or backend mismatch | Delete `foundry_archive` and perform a fresh SAVE with the exact environment and flags. | +| Dense TP topology/collective validation fails | Use exactly TP=2, DP=2, PP=1, EP=1 and retain disabled custom all-reduce plus both NCCL environment guards. Other TP topologies are unsupported. | | Reserved VMM address differs from `0x600000000000` | Another allocation owns the requested range; restart the process. | | DeepEP reports missing NVSHMEM symbols | Verify the cu13 NVSHMEM wheel or set `nvshmem_host_path` in both TOMLs. | | DeepEP token/QP-depth assertion fails | Lower `SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK` or raise `NVSHMEM_QP_DEPTH`, identically for SAVE and LOAD. | From 7d70c6fe5b6943efe70d12b68e255a757be06f2a Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Thu, 23 Jul 2026 23:52:19 +0000 Subject: [PATCH 21/21] fix(sglang): clarify TP attention validation error Co-authored-by: Rahul Chalamala --- python/foundry/integration/sglang/cuda_graph.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/foundry/integration/sglang/cuda_graph.py b/python/foundry/integration/sglang/cuda_graph.py index 991c23fa..c8abc003 100644 --- a/python/foundry/integration/sglang/cuda_graph.py +++ b/python/foundry/integration/sglang/cuda_graph.py @@ -148,7 +148,7 @@ def validate_configuration(server_args: Any) -> None: ) elif attention_backend != "flashinfer": raise ValueError( - "Foundry single-GPU and regular-DP decode require the FlashInfer " + "Foundry non-DP-attention decode requires the FlashInfer " f"attention backend; got {attention_backend!r}." ) elif moe_backend != "none":