Integrate Vortex with official SGLang 0.5.12 plugins

.gitignore

@@ -64,7 +64,9 @@ MANIFEST
 !vortex_torch/custom_ops/**/meta.json
 !vortex_torch/custom_ops/**/config.json
 
-tests/
+tests/*
+!tests/test_codegen_io.py
+!tests/test_sglang_runtime_contract.py
 *_backup.py
 # PyInstaller
 #  Usually these files are written by a python script from a template

README.md

@@ -55,17 +55,15 @@ This makes Vortex a platform for **autonomous algorithm discovery**: AI agents g
 
 ```bash
 git clone --recursive https://github.com/Infini-AI-Lab/vortex_torch.git
-
-# Install SGLang dependency
-cd third_party/sglang/v0.5.9/sglang
-pip install -e "python"
-cd ../../../../
-
-# Install Vortex
 cd vortex_torch
-pip install -e .
+pip install -e ".[sglang]"
 ```
 
+This installs the official `sglang==0.5.12.post1` package. SGLang discovers
+Vortex through its standard `sglang.srt.plugins` entry point; no patched
+SGLang checkout is required. Add the research dependencies used by the
+benchmark scripts with `pip install -e ".[sglang,research]"`.
+
 ---
 
 ## 🤖 Innovate Sparse Attention with OpenHands
@@ -223,15 +221,13 @@ class CustomSparseAttention(vFlow):
 
 ## 🏃 Launch it with SGLang
 
-The launch script is a **separate file** from the flow. It imports
-sglang and vortex_torch, then starts the engine. Importing `vortex_torch`
-is what wires vortex into sglang's decode loop (it installs the
-`ServerArgs` ↔ `VortexConfig` adapter), so the import is required even
-though you don't call it directly.
+The launch script is a **separate file** from the flow. The installed SGLang
+plugin wires Vortex into the decode loop automatically. Import `VortexConfig`
+to construct the Vortex configuration; no import-for-side-effect ordering is
+required.
 
 ```python
 import sglang as sgl
-import vortex_torch  # noqa: F401 — import for side effect: installs the VortexConfig adapter
 from vortex_torch.engine.sgl.config import VortexConfig
 
 llm = sgl.Engine(
@@ -471,12 +467,11 @@ examples/misc/server_launch.sh Qwen/Qwen3-4B 1
 
 Two details make server mode work:
 
-1. **`import vortex_torch` must run first.** The script doesn't call
-   `python -m sglang.launch_server` directly — that builds `ServerArgs` in
-   the parent before vortex is imported, so the adapter that folds the
-   config wouldn't be installed yet. Instead it imports `vortex_torch`,
-   then calls sglang's `run_server`, so the `ServerArgs` ↔ `VortexConfig`
-   adapter is in place before the args are pickled to the scheduler worker.
+1. **Load plugins before parsing server arguments.** The custom launcher calls
+   SGLang's `load_plugins()` before `prepare_server_args()`. This discovers the
+   installed `sglang.srt.plugins` entry point and runs Vortex's `register()`
+   function before `--vortex-config` is parsed. The standard SGLang CLI already
+   performs the same step.
 2. **Knobs are passed as JSON via `--vortex-config`.** The per-knob
    `--vortex-*` CLI flags no longer exist; the script writes the
    `VortexConfig` fields (prefix stripped) to a temp JSON file and feeds it
@@ -509,5 +504,3 @@ If you find Vortex useful in your research, please cite:
       url={https://arxiv.org/abs/2606.06453}, 
 }
 ```
-
-

docs/examples.md

@@ -197,11 +197,9 @@ examples/misc/server_launch.sh Qwen/Qwen3-4B 1
 
 Two details make server mode work:
 
-1. **`import vortex_torch` must run first.** The script doesn't call
-   `python -m sglang.launch_server` directly — that builds `ServerArgs` before
-   Vortex is imported, so the adapter wouldn't be installed yet. It imports
-   `vortex_torch`, then calls SGLang's `run_server`, so the `ServerArgs` ↔
-   `VortexConfig` adapter is in place before the args are pickled to the worker.
+1. **Vortex must be installed with the `sglang` extra.** SGLang discovers the
+   `vortex_torch` plugin before parsing `ServerArgs`, and the same config is
+   preserved when the scheduler worker is spawned.
 2. **Knobs are passed as JSON via `--vortex-config`.** The per-knob `--vortex-*`
    flags no longer exist; the script writes the `VortexConfig` fields (prefix
    stripped) to a temp JSON file and feeds it through `--vortex-config '<json>'`.

docs/installation.md

@@ -1,54 +1,46 @@
 # Installation
 
-Vortex plugs into a **vendored SGLang** (under `third_party/`). Install SGLang
-in editable mode first, then Vortex. Installation is CPU-only — all kernels are
-prebuilt wheels or JIT-compiled at runtime — so it works even while the GPUs are
-busy.
+Vortex is an out-of-tree plugin for the official `sglang==0.5.12.post1`
+package. Installation does not require a GPU; Vortex's custom kernels are
+JIT-compiled when a GPU server starts.
 
 ## From source
 
 ```bash
-git clone --recursive https://github.com/Infini-AI-Lab/vortex_torch.git
+git clone https://github.com/Infini-AI-Lab/vortex_torch.git
 cd vortex_torch
-
-# 1. SGLang dependency (vendored, editable)
-cd third_party/sglang/v0.5.9/sglang
-pip install -e "python"
-cd ../../../../
-
-# 2. Vortex (editable)
-pip install -e .
+pip install -e ".[sglang]"
 ```
 
-If you cloned without `--recursive`, pull the submodules first:
+SGLang discovers the installed Vortex hook through the
+`sglang.srt.plugins` entry-point group. The old `third_party/sglang/v0.5.9`
+tree remains only as historical reference material. It is not installed,
+tested, or supported by the current runtime.
+
+Install the benchmark and dataset dependencies as a separate extra:
 
 ```bash
-git submodule update --init --recursive
+pip install -e ".[sglang,research]"
 ```
 
 ## Reproducible conda environment (recommended)
 
-The repo ships a one-shot script that builds the exact tested environment —
-Python 3.12, torch 2.9.1+cu128, flashinfer 0.6.3, transformers 4.57.1, plus
-editable SGLang and Vortex:
+The repo ships a one-shot script that creates a Python 3.12 environment and
+installs the same official SGLang/Vortex plugin stack:
 
 ```bash
 bash install_vortex.sh          # creates the `vortex_v1` conda env
 conda activate vortex_v1
 ```
 
-```{note}
-For **MLA models** such as GLM-4.7-Flash (HF type `glm4_moe_lite`, which
-requires `transformers >= 5.0`), use `install_vortex_glm.sh` instead — it builds
-a separate `vortex_glm` env that overrides transformers with a GLM-supporting
-build.
-```
+`install_vortex_glm.sh` now delegates to the same stack under the historical
+`vortex_glm` environment name; SGLang 0.5.12 already uses Transformers 5.x.
 
 ## Verify
 
 ```bash
-python -c "import torch, sglang, vortex_torch, flashinfer; print('vortex ok')"
+python -c "from importlib.metadata import entry_points; assert any(e.name == 'vortex' for e in entry_points(group='sglang.srt.plugins')); print('vortex plugin ok')"
 ```
 
-You should see `vortex ok` with no import errors. You're ready for the
+You should see `vortex plugin ok`. You're ready for the
 [Quick Start](quickstart.md).

docs/landing/index.html

@@ -388,13 +388,8 @@ <h3>Install</h3>
         <pre><code>git clone --recursive https://github.com/Infini-AI-Lab/vortex_torch.git
 cd vortex_torch
 
-<span class="cmt"># SGLang dependency (vendored)</span>
-cd third_party/sglang/v0.5.9/sglang
-pip install -e "python"
-cd ../../../../
-
-<span class="cmt"># Vortex</span>
-pip install -e .</code></pre>
+<span class="cmt"># Official SGLang + Vortex plugin</span>
+pip install -e ".[sglang]"</code></pre>
       </div>
       <div class="panel cite reveal">
         <div class="panel-head">

docs/quickstart.md

@@ -53,14 +53,13 @@ class CustomSparseAttention(vFlow):
 
 ## 2. Launch SGLang with your flow
 
-The launch script is a **separate file**. Importing `vortex_torch` is what wires
-Vortex into SGLang (it installs the `ServerArgs` ↔ `VortexConfig` adapter), so the
-import is required even though you don't call it directly. Every Vortex knob lives
-in a single [`VortexConfig`](examples.md); passing it turns sparsity **on**.
+The launch script is a **separate file**. The installed SGLang plugin wires
+Vortex in automatically. Import `vortex_torch` here to construct the
+[`VortexConfig`](examples.md); passing that config turns sparsity **on**.
 
 ```python
 import sglang as sgl
-import vortex_torch  # noqa: F401 — installs the VortexConfig adapter
+import vortex_torch
 from vortex_torch.engine.sgl.config import VortexConfig
 
 llm = sgl.Engine(

examples/misc/server_launch.sh

@@ -7,7 +7,7 @@
 # The vortex_* hyper-parameters are no longer individual CLI flags: sglang's
 # ServerArgs now carries a single aggregated `vortex` field, exposed on the CLI
 # as `--vortex-config '<json>'` (see vortex_torch/engine/sgl/config.py and
-# third_party/.../server_args.py). Passing the old per-knob `--vortex-*` flags
+# vortex_torch/engine/sgl/plugin.py). Passing the old per-knob `--vortex-*` flags
 # fails argparse. We therefore write the knobs to a JSON file and feed it
 # through `--vortex-config`. Keys are the VortexConfig field names (the
 # `vortex_` prefix is stripped). Providing a non-null vortex config implicitly
@@ -38,21 +38,18 @@ cat > "$VORTEX_CONFIG_FILE" <<'JSON'
 }
 JSON
 
-# NOTE: we cannot use `python -m sglang.launch_server` directly. That entrypoint
-# builds `ServerArgs` in the parent process before anything imports vortex_torch,
-# so the `--vortex-config` JSON string is never folded into a VortexConfig (the
-# `ServerArgs.__init__` adapter that does this is installed by `import
-# vortex_torch`). The raw string then gets pickled to the spawned scheduler
-# worker, where `server_args.vortex_block_size` -> `getattr(str, "block_size")`
-# raises AttributeError. Importing vortex_torch FIRST, in this parent process,
-# installs the adapter so the conversion happens before ServerArgs is pickled.
+# This custom Python launcher calls `prepare_server_args()` directly, so it must
+# mirror SGLang's standard CLI and call `load_plugins()` first. That discovers
+# the installed Vortex entry point, registers the ServerArgs hooks, and makes
+# `--vortex-config` available before argument parsing.
 python -c '
 import os, sys
-import vortex_torch  # installs the ServerArgs adapter + backend integration
 from sglang.launch_server import run_server
+from sglang.srt.plugins import load_plugins
 from sglang.srt.server_args import prepare_server_args
 from sglang.srt.utils import kill_process_tree
 
+load_plugins()
 server_args = prepare_server_args(sys.argv[1:])
 try:
     run_server(server_args)

install_vortex.sh

@@ -1,86 +1,43 @@
 #!/usr/bin/env bash
-#
-# Reproducible build of the `vortex_v1` conda environment — the default env for
-# this project (all the slash commands, RULER, and AIME24 runners expect it).
-#
-# Unlike `vortex_glm` (see install_vortex_glm.sh), this env KEEPS the pinned
-# transformers==4.57.1 that both sglang and vortex_torch require — there is NO
-# transformers override step. It therefore does NOT load GLM-4.7-Flash
-# (`glm4_moe_lite` needs transformers >= 5.0); use install_vortex_glm.sh for that.
-#
-# Captured from the working env: python 3.12, torch 2.9.1+cu128, torchvision
-# 0.24.1, torchaudio 2.9.1, flashinfer 0.6.3, transformers 4.57.1, sglang
-# (editable, vendored v0.5.9), vortex_torch (editable).
-#
-# Usage:
-#   bash install_vortex.sh            # create the env
-#   FORCE=1 bash install_vortex.sh    # remove an existing env first
-#   ENV_NAME=vortex2 bash install_vortex.sh   # build under a different name
-#
-# Install only needs CPU (all kernels are prebuilt wheels or JIT-compiled at
-# runtime), so it works even while the GPUs are busy.
+# Build a reproducible Vortex + official SGLang 0.5.12 environment.
 
 set -euo pipefail
 
 ENV_NAME="${ENV_NAME:-vortex_v1}"
 PY_VER="${PY_VER:-3.12}"
-
 REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
-SGLANG_DIR="$REPO_ROOT/third_party/sglang/v0.5.9/sglang/python"
-[ -d "$SGLANG_DIR" ] || { echo "ERROR: vendored sglang not found at $SGLANG_DIR" >&2; exit 1; }
 
-# ---- conda bootstrap -------------------------------------------------------
 source "$(conda info --base)/etc/profile.d/conda.sh"
 
 if conda env list | awk '{print $1}' | grep -qx "$ENV_NAME"; then
     if [ "${FORCE:-0}" = "1" ]; then
-        echo ">>> removing existing env '$ENV_NAME' (FORCE=1)"
         conda env remove -y -n "$ENV_NAME"
     else
-        echo "ERROR: conda env '$ENV_NAME' already exists. Re-run with FORCE=1 to recreate." >&2
+        echo "ERROR: conda env '$ENV_NAME' already exists. Re-run with FORCE=1." >&2
         exit 1
     fi
 fi
 
-echo ">>> [1/4] creating conda env '$ENV_NAME' (python $PY_VER)"
+echo ">>> [1/2] creating conda env '$ENV_NAME' (python $PY_VER)"
 conda create -y -n "$ENV_NAME" python="$PY_VER"
 conda activate "$ENV_NAME"
-python -m pip install --upgrade pip
+python -m pip install --upgrade pip setuptools wheel
 
-# ---- 2. torch (CUDA 12.8 build) pinned first ------------------------------
-# Pinned before sglang sees `torch==2.9.1` so the exact CUDA build is locked in
-# and torchvision/torchaudio match. (Default PyPI torch 2.9.1 is the cu128 wheel.)
-echo ">>> [2/4] installing torch 2.9.1 + torchvision 0.24.1 + torchaudio 2.9.1 (cu128)"
-pip install torch==2.9.1 torchvision==0.24.1 torchaudio==2.9.1
+echo ">>> [2/2] installing Vortex, official SGLang, and research tools"
+python -m pip install -e "$REPO_ROOT[sglang,research]"
 
-# ---- 3. sglang (editable, vendored) ---------------------------------------
-# Pulls the runtime tree: flashinfer_python/cubin 0.6.3, sgl-kernel, xgrammar,
-# outlines, cuda-python, etc. (and the pinned transformers 4.57.1 — KEPT here).
-echo ">>> [3/4] installing vendored sglang (editable) from $SGLANG_DIR"
-pip install -e "$SGLANG_DIR"
+python - <<'PY'
+from importlib.metadata import entry_points, version
 
-# ---- 4. vortex_torch (editable) -------------------------------------------
-echo ">>> [4/4] installing vortex_torch (editable) from $REPO_ROOT"
-pip install -e "$REPO_ROOT"
+import vortex_torch
 
-# ---- verify ----------------------------------------------------------------
-echo ">>> verifying the environment"
-python - <<'PY'
-import torch, transformers, sglang, vortex_torch
-import flashinfer
-print(f"  python        : {__import__('sys').version.split()[0]}")
-print(f"  torch         : {torch.__version__}  (cuda {torch.version.cuda})")
-print(f"  transformers  : {transformers.__version__}")
-print(f"  flashinfer    : {flashinfer.__version__}")
-print(f"  sglang        : {sglang.__version__}")
-print(f"  vortex_torch  : {getattr(vortex_torch, '__version__', '?')}")
-ok = transformers.__version__.startswith("4.57")
-print(f"  transformers 4.57.x (sglang/vortex_torch pin): {ok}")
-assert ok, f"expected transformers 4.57.x, got {transformers.__version__}"
+plugins = {ep.name: ep.value for ep in entry_points(group="sglang.srt.plugins")}
+assert plugins.get("vortex") == "vortex_torch.engine.sgl.plugin:register", plugins
+assert version("sglang") == "0.5.12.post1"
+assert version("vortex-torch") == vortex_torch.__version__
+print(f"sglang       : {version('sglang')}")
+print(f"vortex_torch : {vortex_torch.__version__}")
+print(f"plugin       : {plugins['vortex']}")
 PY
 
-echo ""
-echo ">>> done. Activate with:  conda activate $ENV_NAME"
-echo ">>> sanity run (needs a free GPU):"
-echo "    conda activate $ENV_NAME"
-echo "    CUDA_VISIBLE_DEVICES=<gpu> python algorithm_scientist/run_ruler.py --config <submission>.json"
+echo ">>> done. Activate with: conda activate $ENV_NAME"