Revert "[4/4] build: align training with the pinned SGLang revision"

.github/workflows/test.yaml

@@ -39,7 +39,7 @@ jobs:
       cancel-in-progress: false
       queue: max
     container:
-      image: lmsysorg/sglang:v0.5.18-cu130 # CUDA build-toolchain base; the venvs install the pinned SGLang revision below
+      image: lmsysorg/sglang:v0.5.18-cu130 # CUDA build-toolchain base; the capture venv installs the pinned SGLang revision below
       options: --gpus all --shm-size=2g --rm -v /dev/shm --privileged --pid=host
     # The self-hosted runner exports HTTP(S)_PROXY into the container but no
     # longer a NO_PROXY, so local requests (sglang's own warmup to 127.0.0.1)
@@ -51,8 +51,6 @@ jobs:
       # v2 deliberately ignores the legacy venv cache with stale dist-info files.
       SPECFORGE_CI_CACHE: /github/home/specforge-ci-v2
       FLASH_ATTN_VERSION: 2.8.3.post2
-      # Both CI environments need the plugin's forward-observer/completion APIs.
-      SGLANG_COMMIT: c757c4bb7c55d1eb69fdb2bdc6e6fb8507acb7a0
     steps:
       - name: Checkout code
         uses: actions/checkout@v4
@@ -96,11 +94,12 @@ jobs:
             uv venv sf -p 3.11
           fi
           source sf/bin/activate
-          uv pip install wheel setuptools setuptools-rust setuptools-scm psutil packaging ninja
-          uv pip install torch==2.14.1 --index-url https://download.pytorch.org/whl/cu130
-          uv pip install -v . --no-build-isolation --prerelease=allow \
-            "sglang @ git+https://github.com/sgl-project/sglang.git@${SGLANG_COMMIT}#subdirectory=python"
-          uv pip install --no-deps plugins/sglang-spec-capture
+          uv pip install wheel setuptools psutil packaging ninja
+          uv pip install torch==2.13.0 --index-url https://download.pytorch.org/whl/cu130
+          uv pip install -v . --prerelease=allow
+          # SGLang 0.5.18 installs cubin separately and requires it to match
+          # flashinfer-python; replace the 0.6.12 copy in cached venvs.
+          uv pip install flashinfer-cubin==0.6.17 --index-url https://flashinfer.ai/whl
           # The live server-capture gate needs the Mooncake client AND the
           # mooncake_master console script; the isolated venv cannot see the
           # container's site-packages. cu130 image -> the cuda13 wheel.
@@ -113,11 +112,12 @@ jobs:
           if [ ! -d sf-capture ]; then
             uv venv sf-capture -p 3.12
           fi
-          # Use the same CI-owned revision as sf, with ordinary cache reuse.
+          # Keep native plugin APIs isolated from the training environment.
+          sglang_requirement="sglang @ git+https://github.com/sgl-project/sglang.git@c757c4bb7c55d1eb69fdb2bdc6e6fb8507acb7a0#subdirectory=python"
           uv pip install --python sf-capture/bin/python wheel setuptools setuptools-rust setuptools-scm
           uv pip install --python sf-capture/bin/python torch==2.14.1 --index-url https://download.pytorch.org/whl/cu130
           uv pip install --python sf-capture/bin/python --no-build-isolation --prerelease=allow \
-            "sglang @ git+https://github.com/sgl-project/sglang.git@${SGLANG_COMMIT}#subdirectory=python" mooncake-transfer-engine-cuda13
+            "$sglang_requirement" mooncake-transfer-engine-cuda13
           uv pip install --python sf-capture/bin/python --no-deps plugins/sglang-spec-capture
           sf-capture/bin/python -c 'from sglang.srt.model_executor.forward_observer import ForwardObserver; from sglang.srt.managers.deferred_output import ResponseCompletion'
 

docs/sections/get_started/installation.md

@@ -15,26 +15,10 @@ cd SpecForge
 uv venv -p 3.11 --seed
 source .venv/bin/activate
 
-# Install CUDA 13 training/server dependencies, then SpecForge and its plugin.
-# Building the pinned SGLang revision requires Rust 1.92, C++ and CUDA development tools.
-uv pip install torch==2.14.1 --index-url https://download.pytorch.org/whl/cu130
-uv pip install wheel setuptools setuptools-rust setuptools-scm
-uv pip install -v -e . --no-build-isolation --prerelease=allow
-# Install SGLang explicitly; SpecForge's package dependencies do not install it.
-uv pip install --no-build-isolation --prerelease=allow \
-  "sglang @ git+https://github.com/sgl-project/sglang.git@c757c4bb7c55d1eb69fdb2bdc6e6fb8507acb7a0#subdirectory=python"
-uv pip install --no-deps plugins/sglang-spec-capture
-uv pip install mooncake-transfer-engine-cuda13
+# install specforge
+uv pip install -v -e . --prerelease=allow
 ```
 
-Install SGLang commit `c757c4bb7c55d1eb69fdb2bdc6e6fb8507acb7a0` manually
-using the command above for capture workloads. SpecForge pins Torch `2.14.1`
-and Transformers `5.19.0`, but does not install SGLang automatically.
-Normal installer caches are reused. Online
-training and offline feature preparation use the capture plugin; see the
-[plugin guide](../../../plugins/sglang-spec-capture/README.md) for server setup.
-Offline training from existing features does not require a capture server.
-
 - **Install from PyPI**
 
 ::: code-group
@@ -49,10 +33,6 @@ pip install specforge
 
 :::
 
-PyPI installation also requires the explicit SGLang installation above when
-using a capture server. Install SGLang and the capture plugin in the server's
-environment; an existing compatible server can be used separately.
-
 ## Accelerator-specific environments
 
 ### NVIDIA CUDA

examples/offline/README.md

@@ -13,7 +13,7 @@ config/tokenizer and the same draft config used for training.
 
 Use an SGLang environment with the native APIs described in the
 [plugin installation guide](../../plugins/sglang-spec-capture/README.md), and
-install the current plugin there. Use the SGLang commit pinned in `pyproject.toml`; released 0.5.23
+install the current plugin there. Use SGLang commit `c757c4bb7c55d1eb69fdb2bdc6e6fb8507acb7a0`; released 0.5.23
 still lacks the capture APIs. The server and SpecForge client may use the same
 environment with a compatible Mooncake wheel, or separate environments.
 

plugins/sglang-spec-capture/README.md

@@ -1,8 +1,8 @@
 # specforge-sglang-capture
 
 An SGLang plugin that writes a target model's prefill hidden states straight
-into Mooncake for SpecForge's disaggregated training and offline preparation. It requires an SGLang
-build with these generic extension points:
+into Mooncake for SpecForge's disaggregated training and offline preparation.
+It requires an SGLang build with these generic extension points:
 
 - `--enable-return-aux-hidden-states` / `--return-aux-hidden-state-layer-ids`:
   aux hidden-state capture on a target without a draft. `--speculative-algorithm`
@@ -24,24 +24,20 @@ build with these generic extension points:
 This version requires SGLang's [aux-capture changes in PR #43449](https://github.com/sgl-project/sglang/pull/43449)
 and the completion-based `defer_response` revision of
 [PR #43450](https://github.com/sgl-project/sglang/pull/43450), or a branch containing
-both. Install SGLang manually at commit
-`c757c4bb7c55d1eb69fdb2bdc6e6fb8507acb7a0`; the released 0.5.23 wheel still lacks the forward
-observer, completion and auxiliary-capture APIs required here.
+both. CI uses SGLang commit `c757c4bb7c55d1eb69fdb2bdc6e6fb8507acb7a0` in a separate
+cached `sf-capture` environment. These extension points are not available in stock SGLang 0.5.18.
 
 The aux-capture API returns `last_hidden_states` after the final norm and
 before LM-head-only scaling, including Inkling's muP division. The plugin
 preserves those values; SpecForge folds `logits_mup_width_multiplier` into
 its frozen target head. Older builds that return scaled states are unsupported;
 upgrade the capture environment to a build with the required API.
 
-The supported capture setup uses this SGLang revision with
-Torch 2.14.1 and Transformers 5.19.0. SpecForge does not install SGLang as a
-package dependency. They can run in the same environment;
-the server remains a separate process. Separate environments are also supported
-through `capture_servers[].python_executable`. The plugin itself does not import
-the main SpecForge package. CI uses a separate `sf-capture` environment for
-the live server. Both CI environments install the revision specified by
-`SGLANG_COMMIT` in `.github/workflows/test.yaml`.
+Use separate environments for the SGLang server and SpecForge training.
+SpecForge currently pins SGLang 0.5.18 and different Torch/Transformers versions;
+installing the main SpecForge package into the server environment would replace
+the server's dependencies. The server only needs this plugin, its SGLang build,
+and Mooncake; the plugin does not import the main SpecForge package.
 
 Install into the SGLang server's environment. The plugin declares no
 dependencies, so the server's pins stay intact:
@@ -59,24 +55,20 @@ Source installation requires the Rust toolchain pinned by SGLang, a C/C++
 compiler, and CUDA development tools. An existing compatible SGLang server
 environment can also be used.
 
-For example, from the SpecForge checkout on a CUDA 13 host with the Rust,
-C/C++ and CUDA build tools required by the pinned SGLang revision:
+For example, from the SpecForge checkout on a CUDA 13 host, with an SGLang
+checkout containing both changes:
 
 ```bash
-uv venv .venv -p 3.12
-source .venv/bin/activate
-uv pip install torch==2.14.1 --index-url https://download.pytorch.org/whl/cu130
-uv pip install wheel setuptools setuptools-rust setuptools-scm
-uv pip install --no-build-isolation --prerelease=allow -e .
-uv pip install --no-build-isolation --prerelease=allow \
-  "sglang @ git+https://github.com/sgl-project/sglang.git@c757c4bb7c55d1eb69fdb2bdc6e6fb8507acb7a0#subdirectory=python"
-uv pip install --no-deps plugins/sglang-spec-capture
-uv pip install mooncake-transfer-engine-cuda13
+python3.12 -m venv .venv-sglang .venv-training
+.venv-sglang/bin/python -m pip install torch==2.14.1 --index-url https://download.pytorch.org/whl/cu130
+.venv-sglang/bin/python -m pip install -e /path/to/sglang/python mooncake-transfer-engine-cuda13
+.venv-sglang/bin/python -m pip install --no-deps plugins/sglang-spec-capture
+.venv-training/bin/python -m pip install torch==2.13.0 --index-url https://download.pytorch.org/whl/cu130
+.venv-training/bin/python -m pip install -e . mooncake-transfer-engine-cuda13
+.venv-sglang/bin/python -m pip check
+.venv-training/bin/python -m pip check
 ```
 
-The SGLang requirement uses the exact commit above; no cache refresh is forced. Rebuild
-optional compiled training kernels against the installed Torch when upgrading.
-
 SGLang discovers the plugin through the `sglang.srt.plugins` entry point. If
 `SGLANG_PLUGINS` is set, it must include `specforge_spec_capture`.
 
@@ -86,7 +78,7 @@ SGLang discovers the plugin through the `sglang.srt.plugins` entry point. If
 SPECFORGE_SPEC_CAPTURE=1 \
 MOONCAKE_MASTER_SERVER_ADDR=127.0.0.1:50051 \
 MOONCAKE_METADATA_SERVER=http://127.0.0.1:8080/metadata \
-python -m sglang.launch_server --model-path <target> --skip-tokenizer-init \
+.venv-sglang/bin/python -m sglang.launch_server --model-path <target> --skip-tokenizer-init \
   --chunked-prefill-size -1 \
   --speculative-algorithm DFLASH --enable-return-aux-hidden-states \
   --return-aux-hidden-state-layer-ids 1 16 31 46 61 \
@@ -108,9 +100,7 @@ The default, `auto`, finds it without configuration:
 - External servers: before the first request, the producer reads each
   server's `/server_info`. The enabled plugin adds `spec_capture_backend:
   plugin` to every scheduler's `internal_states`, together with
-  `spec_capture_method` and `spec_capture_aux_layer_ids`. Offline preparation
-  validates these fields before submitting prompts; reinstall this plugin on
-  existing servers when updating the preparation client.
+  `spec_capture_method` and the ordered `spec_capture_aux_layer_ids`.
 
 `server_capture: plugin`, or `DISAGG_SERVER_CAPTURE=plugin` when the config
 leaves it at `auto`, skips detection.

pyproject.toml

@@ -11,8 +11,8 @@ description = "SpecForge: Speculative Decoding Training Framework"
 authors = [{name = "SGLang Team"}]
 urls = {Homepage = "https://github.com/sgl-project/SpecForge"}
 dependencies = [
-    "torch==2.14.1",
-    "transformers==5.19.0",
+    "torch==2.13.0",
+    "transformers==5.12.1",
     "datasets",
     "tqdm",
     "accelerate",
@@ -28,6 +28,7 @@ dependencies = [
     "psutil",
     "pyyaml",
     "safetensors",
+    "sglang==0.5.18",
     "requests",
     "tensorboard",
     "typing-extensions",

pyproject_xpu.toml

@@ -6,13 +6,13 @@ build-backend = "setuptools.build_meta"
 name = "specforge"
 dynamic = ["version"]
 readme = "README.md"
-requires-python = ">=3.11"
+requires-python = ">=3.10"
 description = "SpecForge: Speculative Decoding Training Framework"
 authors = [{name = "SGLang Team"}]
 urls = {Homepage = "https://github.com/sgl-project/SpecForge"}
 dependencies = [
-    "torch==2.14.1+xpu",
-    "transformers==5.19.0",
+    "torch==2.13.0+xpu",
+    "transformers==5.12.1",
     "datasets",
     "tqdm",
     "accelerate",
@@ -24,7 +24,7 @@ dependencies = [
     "psutil",
     "pyyaml",
     "safetensors",
-    # Install an XPU-compatible SGLang main build and the capture plugin separately.
+    #"sglang==0.5.18", # commented due to CUDA dependencies
     "requests",
     "tensorboard",
     "typing-extensions",

requirements-rocm.txt

@@ -1,9 +1,9 @@
-# PyTorch 2.14.1 publishes Linux ROCm 7.2 wheels on this index.
+# PyTorch 2.13 publishes Linux ROCm 7.2 wheels on this index.
 --extra-index-url https://download.pytorch.org/whl/rocm7.2
 
 pre-commit
-torch==2.14.1+rocm7.2
-transformers==5.19.0
+torch==2.13.0+rocm7.2
+transformers==5.12.1
 datasets
 setuptools
 tqdm
@@ -16,8 +16,7 @@ pydantic
 pyyaml
 requests
 safetensors
-# Install a ROCm-compatible SGLang main build and the capture plugin separately;
-# the default SGLang source package resolves CUDA dependencies.
+sglang[all]==0.5.18
 openai-harmony
 tensorboard
 typing-extensions

specforge/config/schema.py

@@ -64,7 +64,7 @@
     ),
     "--config": "SGLang merges that file unchecked; list its flags instead",
 }
-#: Pinned SGLang CLI spellings that write the same option, checked against its
+#: SGLang v0.5.18 spellings that write the same option, checked against its
 #: parser in tests. Positive/negative boolean flags belong to one group too:
 #: they conflict under the no-overrides rule, but their argv is never rewritten.
 _SGLANG_FLAG_ALIASES: Mapping[str, tuple] = {
@@ -85,26 +85,33 @@
     "--experts-shared-outer-loras": ("--no-experts-shared-outer-loras",),
     "--lora-strict-loading": ("--no-lora-strict-loading",),
     "--dllm-fdfo": ("--no-dllm-fdfo",),
-    "--enable-tf32-matmul": ("--no-enable-tf32-matmul",),
-    "--attn-dp-size": ("--attention-data-parallel-size",),
-    "--cuda-graph-prefill-max-context": ("--context-bucket",),
-    "--enable-fused-qk-norm-rope": ("--no-enable-fused-qk-norm-rope",),
-    "--speculative-enable-w4a4-mxfp4-megamoe": (
-        "--no-speculative-enable-w4a4-mxfp4-megamoe",
+    "--mamba-radix-cache-strategy": ("--mamba-scheduler-strategy",),
+    "--stream-output": ("--incremental-streaming-output",),
+    "--cuda-graph-bs-decode": ("--cuda-graph-bs",),
+    "--cuda-graph-max-bs-decode": ("--cuda-graph-max-bs",),
+    "--cuda-graph-bs-prefill": ("--piecewise-cuda-graph-tokens",),
+    "--cuda-graph-max-bs-prefill": ("--piecewise-cuda-graph-max-tokens",),
+    "--cuda-graph-backend-prefill": (
+        "--enable-breakable-cuda-graph",
+        "--disable-piecewise-cuda-graph",
+        "--enforce-piecewise-cuda-graph",
     ),
-    "--enable-hierarchical-cache": ("--enable-hicache",),
-    "--mm-enable-dp-encoder": ("--no-mm-enable-dp-encoder",),
-    "--ple-offload-embedding": ("--no-ple-offload-embedding",),
-    "--enable-decoder-swa-bounded-replay": ("--no-enable-decoder-swa-bounded-replay",),
-    "--enable-megamoe-for-layernorm-sp": ("--no-enable-megamoe-for-layernorm-sp",),
-    "--enable-dsa-fused-indexer": ("--no-enable-dsa-fused-indexer",),
+    "--cuda-graph-tc-compiler": ("--piecewise-cuda-graph-compiler",),
+    "--dsa-prefill-cp-mode": ("--nsa-prefill-cp-mode",),
+    "--dsa-prefill-backend": ("--nsa-prefill-backend",),
+    "--dsa-decode-backend": ("--nsa-decode-backend",),
+    "--enable-dsa-prefill-context-parallel": ("--enable-nsa-prefill-context-parallel",),
+    "--enable-linear-replayssm-spec": ("--enable-gdn-replayssm-spec",),
+    "--speculative-draft-window-size": ("--speculative-dflash-draft-window-size",),
 }
-#: Complete pinned SGLang CLI options that are also prefixes of other options.
+#: Complete SGLang v0.5.18 options that are also prefixes of other options.
 #: argparse matches these exactly instead of expanding them. Keep this small
-#: table in sync with the installed parser (covered by the launch-plan tests),
+#: table in sync with the pinned parser (covered by the launch-plan tests),
 #: without importing SGLang/PyTorch during configuration validation.
 _SGLANG_EXACT_PREFIX_FLAGS = frozenset(
     {
+        "--cuda-graph-bs",
+        "--cuda-graph-max-bs",
         "--disable-cuda-graph",
         "--dllm-algorithm",
         "--enable-dp-attention",
@@ -119,9 +126,7 @@
         "--log-level",
         "--log-requests",
         "--model",
-        "--prefill-decode-interval",
         "--quantization",
-        "--radix-eviction-policy",
         "--sidecar",
         "--speculative-adaptive",
         "--speculative-draft-model",

specforge/inference/sglang_integration.md

@@ -5,7 +5,7 @@
 Online capture uses [specforge-sglang-capture](../../plugins/sglang-spec-capture/README.md).
 Install it in an SGLang environment with native auxiliary hidden-state capture,
 `ModelRunner.forward_observer` and completion-based `Scheduler.defer_response`.
-Use the SGLang commit pinned in `pyproject.toml`; released 0.5.23 does not expose those APIs. SpecForge does not ship source
+Use SGLang commit `c757c4bb7c55d1eb69fdb2bdc6e6fb8507acb7a0`; released 0.5.23 does not expose those APIs. SpecForge does not ship source
 patches or fall back to patched servers.
 
 The launcher accepts `server_capture: auto` or `plugin`. Managed-local `auto`

tests/test_scripts/test_ci_dependencies.py

@@ -228,36 +228,6 @@ def test_workflow_orders_install_check_and_snapshot_before_tests(self):
         self.assertNotIn("flash-attn", self.steps["Install dependencies"]["run"])
         self.assertNotIn("/github/home/sf", self.steps["Restore cache"]["run"])
 
-    def test_both_environments_install_ci_pin_without_package_manifest(self):
-        self.env["SGLANG_COMMIT"] = self.job["env"]["SGLANG_COMMIT"]
-        uv = self.bin / "uv"
-        uv.write_text(
-            f"#!{sys.executable}\n"
-            "import json, os, sys\n"
-            "from pathlib import Path\n"
-            "with (Path(os.environ['TEST_ROOT']) / 'calls').open('a') as log:\n"
-            "    log.write(json.dumps(sys.argv[1:]) + '\\n')\n"
-        )
-        uv.chmod(0o755)
-        for environment in ("sf", "sf-capture"):
-            bin_dir = self.root / environment / "bin"
-            bin_dir.mkdir(parents=True)
-            (bin_dir / "activate").touch()
-            # Stub only the capture environment's post-install import check.
-            python = bin_dir / "python"
-            python.write_text("#!/bin/sh\nexit 0\n")
-            python.chmod(0o755)
-        self.run_step("Install dependencies")
-        self.run_step("Install pinned plugin capture environment")
-        requirement = (
-            "sglang @ git+https://github.com/sgl-project/sglang.git@"
-            f"{self.env['SGLANG_COMMIT']}#subdirectory=python"
-        )
-        installs = [args for args in self.calls() if requirement in args]
-        self.assertEqual(len(installs), 2)
-        self.assertIn(".", installs[0])
-        self.assertIn("sf-capture/bin/python", installs[1])
-
 
 if __name__ == "__main__":
     unittest.main()