From 0f7d98d1970c6001c5d78fbf04f40c5fe1bced9d Mon Sep 17 00:00:00 2001
From: rockyeast <rockyeast@example.local>
Date: Wed, 22 Jul 2026 00:37:50 -0400
Subject: [PATCH 01/10] feat: integrate Vortex through SGLang plugins

---
 pyproject.toml                                |  32 +-
 vortex_torch/__init__.py                      |  12 -
 vortex_torch/engine/sgl/config.py             | 146 +++++---
 vortex_torch/engine/sgl/plugin.py             | 318 ++++++++++++++++++
 .../engine/sgl/transformers_compat.py         |  57 ++++
 vortex_torch/version.py                       |   2 +-
 6 files changed, 508 insertions(+), 59 deletions(-)
 create mode 100644 vortex_torch/engine/sgl/plugin.py
 create mode 100644 vortex_torch/engine/sgl/transformers_compat.py

diff --git a/pyproject.toml b/pyproject.toml
index 1678f455..bd97e9f8 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -4,14 +4,23 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "vortex_torch"
-version = "0.5.0"
+dynamic = ["version"]
 requires-python = ">=3.12"
 dependencies = [
     "torch>=2.7",
+    "numpy",
+]
+
+[project.optional-dependencies]
+sglang = [
+    "sglang==0.5.12.post1",
+    "kernels>=0.12,<0.13",
+]
+research = [
     "lighteval[math]==0.12.2",
     "huggingface_hub==0.36.2",
     "inspect-ai==0.3.207",
-    "transformers==4.57.1",
+    "transformers>=4.57.1,<5.7",
     "s3fs==2025.9.0",
     "wonderwords",
     "langdetect",
@@ -22,6 +31,25 @@ dependencies = [
     "accelerate",
     "unitxt",
 ]
+all = [
+    "vortex_torch[sglang,research]",
+]
+
+[project.entry-points."sglang.srt.plugins"]
+vortex = "vortex_torch.engine.sgl.plugin:register"
+
+[tool.setuptools.dynamic]
+version = {attr = "vortex_torch.version.__version__"}
 
 [tool.setuptools.packages.find]
 include = ["vortex_torch*"]
+
+[tool.setuptools.package-data]
+vortex_torch = [
+    "custom_ops/**/*.json",
+    "custom_ops/**/*.cu",
+    "engine/sgl/attention_backend/csrc/*.cu",
+    "engine/sgl/attention_backend/csrc/*.cuh",
+    "kernels/topk/**/*.json",
+    "kernels/topk/**/*.cu",
+]
diff --git a/vortex_torch/__init__.py b/vortex_torch/__init__.py
index 3d56efd2..f29dea91 100644
--- a/vortex_torch/__init__.py
+++ b/vortex_torch/__init__.py
@@ -30,17 +30,6 @@
         "__version__",
 ]
 
-
-# Eagerly install the ServerArgs flat-kwargs adapter so sgl.Engine(vortex_*=...)
-# keeps working (it must be active in the *parent* before ServerArgs is built).
-# Light: imports only sglang.srt.server_args, not the engine. Best-effort.
-try:
-    from .engine.sgl.config import install_serverargs_adapter as _vx_install_adapter
-    _vx_install_adapter()
-except Exception:
-    pass
-
-
 def __getattr__(name):
     # Lazily expose ``vortex_torch.integration`` (the single sglang-integration
     # module) without making a plain ``import vortex_torch`` pull in sglang/the
@@ -53,4 +42,3 @@ def __getattr__(name):
         return integration
     raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
 
-
diff --git a/vortex_torch/engine/sgl/config.py b/vortex_torch/engine/sgl/config.py
index 06c6c299..d2d8b2f3 100644
--- a/vortex_torch/engine/sgl/config.py
+++ b/vortex_torch/engine/sgl/config.py
@@ -10,16 +10,98 @@
 
 Two entry points populate it:
   * Python: ``sgl.Engine(vortex_topk_val=..., enable_vortex_sparsity=True, ...)``
-    still works — :func:`install_serverargs_adapter` folds those flat kwargs into
-    a ``VortexConfig`` at the ``ServerArgs`` boundary.
+    still works — the official SGLang plugin folds those flat kwargs into a
+    ``VortexConfig`` at the ``ServerArgs`` boundary.
   * Explicit: ``sgl.Engine(vortex=VortexConfig(topk_val=..., ...))``.
 """
 from __future__ import annotations
 
+import json
 from dataclasses import dataclass, fields
 from typing import Any, Dict, List, Optional, Tuple
 
 
+VORTEX_SGLANG_ABI = 1
+_REQUIRED_SGLANG_HOOKS = frozenset(
+    {
+        "server_args.vortex_config",
+        "model_runner.build_sparse_flow",
+        "kv_cache.kv_cell_size",
+        "kv_cache.make_kv_pool",
+        "model_utils.fused_kv_opt_out",
+        "disaggregation.rebuild_aux",
+    }
+)
+VORTEX_SGLANG_PLUGIN_TARGETS = frozenset(
+    {
+        "sglang.srt.server_args.ServerArgs.__init__",
+        "sglang.srt.server_args.ServerArgs.__post_init__",
+        "sglang.srt.server_args.ServerArgs.add_cli_args",
+        "sglang.srt.model_executor.model_runner.ModelRunner.configure_kv_cache_dtype",
+        "sglang.srt.model_executor.pool_configurator.DefaultPoolConfigurator._compute_cell_size",
+        "sglang.srt.model_executor.model_runner_kv_cache_mixin.ModelRunnerKVCacheMixin._init_pools",
+        "sglang.srt.models.utils.enable_fused_set_kv_buffer",
+        "sglang.srt.disaggregation.decode.DecodeTransferQueue.pop_transferred",
+    }
+)
+
+
+def validate_sglang_runtime_contract() -> None:
+    """Fail fast unless the matching Vortex plugin and hooks are active."""
+
+    try:
+        from sglang.srt.plugins import load_plugins
+
+        load_plugins()
+        from sglang.srt.server_args import ServerArgs
+    except (ImportError, ModuleNotFoundError) as exc:
+        raise RuntimeError(
+            "Vortex's SGLang plugin is not active. Install vortex_torch "
+            "with its 'sglang' extra (sglang==0.5.12.post1)."
+        ) from exc
+
+    actual_abi = getattr(ServerArgs, "_vortex_sglang_abi", None)
+    actual_hooks = frozenset(getattr(ServerArgs, "_vortex_sglang_hooks", ()))
+
+    missing_hooks = sorted(_REQUIRED_SGLANG_HOOKS - actual_hooks)
+    server_arg_fields = {field.name for field in fields(ServerArgs)}
+
+    errors = []
+    if actual_abi != VORTEX_SGLANG_ABI:
+        errors.append(
+            f"ABI {actual_abi!r} is installed; ABI {VORTEX_SGLANG_ABI} is required"
+        )
+    if missing_hooks:
+        errors.append(f"missing hooks: {', '.join(missing_hooks)}")
+    if "vortex" not in server_arg_fields:
+        errors.append("ServerArgs has no 'vortex' field")
+    from sglang.srt.plugins.hook_registry import HookRegistry
+
+    declared_targets = frozenset(
+        getattr(ServerArgs, "_vortex_sglang_targets", ())
+    )
+    missing_declarations = sorted(
+        VORTEX_SGLANG_PLUGIN_TARGETS - declared_targets
+    )
+    missing_applied = sorted(declared_targets - HookRegistry._patched)
+    if missing_declarations:
+        errors.append(
+            "plugin did not declare targets: "
+            + ", ".join(missing_declarations)
+        )
+    if missing_applied:
+        errors.append(
+            "plugin hooks were not applied: " + ", ".join(missing_applied)
+        )
+
+    if errors:
+        raise RuntimeError(
+            "Incompatible official SGLang plugin runtime for Vortex ("
+            + "; ".join(errors)
+            + "). Reinstall matching vortex_torch and SGLang packages."
+        )
+
+
 @dataclass
 class VortexConfig:
     """All vortex sparse-attention hyper-parameters (defaults mirror the former
@@ -55,10 +137,23 @@ def from_flat(cls, flat: Dict[str, Any]) -> "VortexConfig":
         return cls(**kw)
 
 
-# The legacy ``vortex_<name>`` -> default map, used by the ServerArgs shim when
-# vortex is disabled (so a stray read returns the historical default). Kept in
-# sync with the dataclass defaults above; duplicated into server_args.py as a
-# plain literal to avoid sglang importing vortex_torch.
+def normalize_vortex_config(value: Any) -> Optional[VortexConfig]:
+    """Normalize JSON, dict, and explicit object inputs to ``VortexConfig``."""
+    if value is None or isinstance(value, VortexConfig):
+        return value
+    if isinstance(value, str):
+        try:
+            value = json.loads(value)
+        except json.JSONDecodeError as exc:
+            raise ValueError(f"--vortex-config must be valid JSON: {exc}") from exc
+    if not isinstance(value, dict):
+        raise TypeError(
+            "vortex must be a VortexConfig, JSON object, dict, or None; "
+            f"got {type(value).__name__}"
+        )
+    return VortexConfig.from_flat(value)
+
+
 def legacy_defaults() -> Dict[str, Any]:
     return {f.name: f.default for f in fields(VortexConfig)}
 
@@ -76,44 +171,7 @@ def split_flat_kwargs(kwargs: Dict[str, Any]) -> Tuple[Optional[VortexConfig], D
     return cfg, kwargs
 
 
-def install_serverargs_adapter() -> bool:
-    """Wrap ``ServerArgs.__init__`` so flat ``vortex_*`` kwargs fold into the
-    single ``vortex`` field. Idempotent; parent-process only (the spawned worker
-    unpickles ``ServerArgs`` and never re-runs ``__init__``). Returns False if
-    sglang is unavailable.
-    """
-    try:
-        from sglang.srt.server_args import ServerArgs
-    except Exception:
-        return False
-    if getattr(ServerArgs, "_vortex_adapter_installed", False):
-        return True
-
-    _orig_init = ServerArgs.__init__
-
-    def __init__(self, *args, **kwargs):
-        v = kwargs.get("vortex")
-        if isinstance(v, VortexConfig):
-            # Explicit object wins; drop any stray flat vortex_* / enable flag.
-            for k in [k for k in kwargs if k.startswith("vortex_")]:
-                kwargs.pop(k)
-            kwargs.pop("enable_vortex_sparsity", None)
-        elif isinstance(v, str):
-            # CLI path: --vortex-config '<json>' arrives as a JSON string.
-            import json
-            kwargs["vortex"] = VortexConfig.from_flat(json.loads(v))
-        else:
-            # Python path: fold flat vortex_* kwargs (gated by enable flag).
-            cfg, kwargs = split_flat_kwargs(kwargs)
-            kwargs["vortex"] = cfg
-        _orig_init(self, *args, **kwargs)
-
-    ServerArgs.__init__ = __init__
-    ServerArgs._vortex_adapter_installed = True
-    return True
-
-
 def cfg(model_runner_or_server_args) -> Optional[VortexConfig]:
-    """Accessor: return the VortexConfig from a ModelRunner or ServerArgs."""
+    """Return the Vortex config from a ModelRunner or ServerArgs object."""
     sa = getattr(model_runner_or_server_args, "server_args", model_runner_or_server_args)
     return getattr(sa, "vortex", None)
diff --git a/vortex_torch/engine/sgl/plugin.py b/vortex_torch/engine/sgl/plugin.py
new file mode 100644
index 00000000..4a7ef8df
--- /dev/null
+++ b/vortex_torch/engine/sgl/plugin.py
@@ -0,0 +1,318 @@
+"""Out-of-tree Vortex integration for SGLang 0.5.12.
+
+SGLang discovers :func:`register` through the ``sglang.srt.plugins`` entry
+point. Every hook is dormant unless ``ServerArgs.vortex`` is populated, so a
+normal SGLang launch is unchanged.
+"""
+from __future__ import annotations
+
+import logging
+from dataclasses import fields, make_dataclass
+from typing import Any, Optional
+
+import torch
+from sglang.srt.server_args import ServerArgs as _BaseServerArgs
+
+from .config import (
+    _REQUIRED_SGLANG_HOOKS,
+    VORTEX_SGLANG_ABI,
+    VORTEX_SGLANG_PLUGIN_TARGETS,
+    VortexConfig,
+    legacy_defaults,
+    normalize_vortex_config,
+    split_flat_kwargs,
+    validate_sglang_runtime_contract,
+)
+
+logger = logging.getLogger(__name__)
+
+_VORTEX_FIELD = fields(
+    make_dataclass(
+        "_VortexFieldHolder", [("vortex", Optional[VortexConfig], None)]
+    )
+)[0]
+_VORTEX_LEGACY_DEFAULTS = legacy_defaults()
+
+
+def _server_args_getattr(self, name: str):
+    if name == "enable_vortex_sparsity":
+        return self.__dict__.get("vortex") is not None
+    if name.startswith("vortex_"):
+        key = name[len("vortex_") :]
+        vortex = self.__dict__.get("vortex")
+        if vortex is not None and hasattr(vortex, key):
+            return getattr(vortex, key)
+        if key in _VORTEX_LEGACY_DEFAULTS:
+            return _VORTEX_LEGACY_DEFAULTS[key]
+    raise AttributeError(
+        f"{type(self).__name__!r} object has no attribute {name!r}"
+    )
+
+
+def _install_server_args_surface() -> None:
+    """Add one dynamic dataclass field without replacing SGLang's class.
+
+    The generated upstream ``__init__`` remains intact and is wrapped below.
+    Adding the Field to ``__dataclass_fields__`` makes ``fields()``, ``asdict()``,
+    ``replace()``, CLI reconstruction, and multiprocessing all preserve the
+    Vortex config.
+    """
+    if "vortex" not in _BaseServerArgs.__dataclass_fields__:
+        _BaseServerArgs.__dataclass_fields__["vortex"] = _VORTEX_FIELD
+    _BaseServerArgs.vortex = None
+    _BaseServerArgs.__getattr__ = _server_args_getattr
+    _BaseServerArgs._vortex_sglang_abi = VORTEX_SGLANG_ABI
+    _BaseServerArgs._vortex_sglang_hooks = _REQUIRED_SGLANG_HOOKS
+    _BaseServerArgs._vortex_sglang_targets = VORTEX_SGLANG_PLUGIN_TARGETS
+
+
+def _around_server_args_init(original, self, *args, **kwargs):
+    vortex = kwargs.pop("vortex", None)
+    if vortex is not None:
+        for key in list(kwargs):
+            if key.startswith("vortex_"):
+                kwargs.pop(key)
+        kwargs.pop("enable_vortex_sparsity", None)
+        vortex = normalize_vortex_config(vortex)
+    else:
+        vortex, kwargs = split_flat_kwargs(kwargs)
+    self.vortex = vortex
+    return original(self, *args, **kwargs)
+
+
+def _around_server_args_post_init(original, self):
+    self.vortex = normalize_vortex_config(self.__dict__.get("vortex"))
+    if self.vortex is not None:
+        if self.vortex.block_size <= 0:
+            raise ValueError("vortex.block_size must be positive")
+        if self.vortex.topk_val <= 0:
+            raise ValueError("vortex.topk_val must be positive")
+        if self.vortex.layers_skip is None:
+            self.vortex.layers_skip = []
+        if self.page_size is None:
+            self.page_size = self.vortex.block_size
+        if self.disaggregation_mode in ("prefill", "decode"):
+            self.disable_overlap_schedule = True
+        if self.vortex.attention_backend in (
+            "cuda_mla",
+            "cuda_mla_profile",
+        ):
+            self.attention_backend = self.vortex.attention_backend
+
+    result = original(self)
+    if self.vortex is not None and self.page_size % self.vortex.block_size:
+        raise ValueError(
+            "SGLang page_size must be a multiple of vortex.block_size; "
+            f"got page_size={self.page_size}, "
+            f"block_size={self.vortex.block_size}"
+        )
+    return result
+
+
+def _after_add_cli_args(result, parser):
+    if any(
+        "--vortex-config" in action.option_strings
+        for action in parser._actions
+    ):
+        return None
+    parser.add_argument(
+        "--vortex-config",
+        dest="vortex",
+        type=str,
+        default=None,
+        help="JSON object configuring the Vortex sparse-attention plugin.",
+    )
+    return None
+
+
+def _validate_runner(runner) -> None:
+    server_args = runner.server_args
+    if not server_args.enable_vortex_sparsity:
+        return
+
+    from sglang.srt.configs.model_config import is_deepseek_nsa, is_deepseek_v4
+    from sglang.srt.platforms import current_platform
+
+    unsupported = []
+    if current_platform.is_out_of_tree():
+        unsupported.append("an out-of-tree hardware platform")
+    if getattr(runner, "is_hybrid_swa", False):
+        unsupported.append("hybrid/sliding-window KV pools")
+    if getattr(runner, "mambaish_config", None) is not None:
+        unsupported.append("Mamba/hybrid-linear KV pools")
+    if is_deepseek_nsa(runner.model_config.hf_config):
+        unsupported.append("native NSA KV pools")
+    if is_deepseek_v4(runner.model_config.hf_config):
+        unsupported.append("DeepSeek-V4 compressed KV pools")
+    if getattr(server_args, "prefill_only_disable_kv_cache", False):
+        unsupported.append("prefill-only cache elision")
+    if getattr(server_args, "cpu_offload_gb", 0):
+        unsupported.append("KV CPU offload")
+    if not runner.spec_algorithm.is_none():
+        unsupported.append("speculative decoding")
+
+    if unsupported:
+        raise RuntimeError(
+            "Vortex does not yet support " + ", ".join(unsupported) + "."
+        )
+
+    config = server_args.vortex
+    if not config.module_name:
+        raise ValueError("vortex.module_name is required")
+    if runner.page_size % config.block_size:
+        raise ValueError(
+            "runner.page_size must be a multiple of vortex.block_size; "
+            f"got {runner.page_size} and {config.block_size}"
+        )
+    if runner.use_mla_backend:
+        if runner.kv_cache_dtype != torch.bfloat16:
+            raise ValueError("Vortex MLA currently requires a bf16 KV cache")
+    elif runner.model_config.v_head_dim != runner.model_config.head_dim:
+        raise ValueError(
+            "Vortex MHA/GQA currently requires equal K and V head dimensions"
+        )
+
+
+def _after_configure_kv_cache_dtype(result, runner, *args, **kwargs):
+    if not runner.server_args.enable_vortex_sparsity:
+        return None
+    validate_sglang_runtime_contract()
+    _validate_runner(runner)
+    runner.block_size = runner.server_args.vortex.block_size
+    from .integration import build_sparse_flow
+
+    runner.sparse_attention = build_sparse_flow(runner)
+    return None
+
+
+def _around_compute_cell_size(original, configurator, runner, num_layers):
+    if not runner.server_args.enable_vortex_sparsity:
+        return original(configurator, runner, num_layers)
+    if getattr(runner, "sparse_attention", None) is None:
+        raise RuntimeError("Vortex sparse flow was not initialized before KV sizing")
+    from .integration import kv_cell_size
+
+    element_size = torch._utils._element_size(runner.kv_cache_dtype)
+    return kv_cell_size(runner, num_layers, element_size)
+
+
+def _around_init_pools(original, runner, *args, **kwargs):
+    if not runner.server_args.enable_vortex_sparsity:
+        return original(runner, *args, **kwargs)
+
+    # Upstream selects these classes inside one large pool-initialization
+    # method. Replace only its local class bindings while that method runs, so
+    # the Vortex pool is allocated directly and no duplicate dense pool exists.
+    globals_dict = original.__globals__
+    original_mha = globals_dict["MHATokenToKVPool"]
+    original_mla = globals_dict["MLATokenToKVPool"]
+
+    def make_pool(*_args, **_kwargs):
+        from .integration import make_kv_pool
+
+        return make_kv_pool(runner)
+
+    globals_dict["MHATokenToKVPool"] = make_pool
+    globals_dict["MLATokenToKVPool"] = make_pool
+    try:
+        result = original(runner, *args, **kwargs)
+    finally:
+        globals_dict["MHATokenToKVPool"] = original_mha
+        globals_dict["MLATokenToKVPool"] = original_mla
+
+    if not getattr(runner.token_to_kv_pool, "is_vortex_pool", False):
+        raise RuntimeError(
+            "SGLang selected a KV pool family that the Vortex plugin did not replace"
+        )
+    return result
+
+
+def _around_enable_fused_set_kv_buffer(original, forward_batch):
+    pool = getattr(forward_batch, "token_to_kv_pool", None)
+    if getattr(pool, "supports_fused_set_kv_buffer", True) is False:
+        return False
+    return original(forward_batch)
+
+
+def _after_pop_transferred(result, queue, *args, **kwargs):
+    if not result:
+        return None
+    pool = queue.scheduler.token_to_kv_pool_allocator.get_kvcache()
+    if not hasattr(pool, "rebuild_aux"):
+        return None
+    table = queue.scheduler.req_to_token_pool.req_to_token
+    for req in result:
+        loc = table[req.req_pool_idx, : req.kv_committed_len]
+        pool.rebuild_aux(loc)
+    return None
+
+
+def register() -> None:
+    """Register Vortex's conditional hooks with SGLang."""
+    from sglang.srt.plugins.hook_registry import HookRegistry, HookType
+    from sglang.srt.server_args import (
+        ATTENTION_BACKEND_CHOICES,
+        add_attention_backend_choices,
+    )
+
+    extra_backends = [
+        name
+        for name in ("cuda_mla", "cuda_mla_profile")
+        if name not in ATTENTION_BACKEND_CHOICES
+    ]
+    if extra_backends:
+        add_attention_backend_choices(extra_backends)
+
+    from .transformers_compat import patch_transformers_560_flash_attention
+
+    patch_transformers_560_flash_attention()
+    _install_server_args_surface()
+
+    hooks = (
+        (
+            "sglang.srt.server_args.ServerArgs.__init__",
+            _around_server_args_init,
+            HookType.AROUND,
+        ),
+        (
+            "sglang.srt.server_args.ServerArgs.__post_init__",
+            _around_server_args_post_init,
+            HookType.AROUND,
+        ),
+        (
+            "sglang.srt.server_args.ServerArgs.add_cli_args",
+            _after_add_cli_args,
+            HookType.AFTER,
+        ),
+        (
+            "sglang.srt.model_executor.model_runner.ModelRunner.configure_kv_cache_dtype",
+            _after_configure_kv_cache_dtype,
+            HookType.AFTER,
+        ),
+        (
+            "sglang.srt.model_executor.pool_configurator.DefaultPoolConfigurator._compute_cell_size",
+            _around_compute_cell_size,
+            HookType.AROUND,
+        ),
+        (
+            "sglang.srt.model_executor.model_runner_kv_cache_mixin.ModelRunnerKVCacheMixin._init_pools",
+            _around_init_pools,
+            HookType.AROUND,
+        ),
+        (
+            "sglang.srt.models.utils.enable_fused_set_kv_buffer",
+            _around_enable_fused_set_kv_buffer,
+            HookType.AROUND,
+        ),
+        (
+            "sglang.srt.disaggregation.decode.DecodeTransferQueue.pop_transferred",
+            _after_pop_transferred,
+            HookType.AFTER,
+        ),
+    )
+    if {target for target, _, _ in hooks} != VORTEX_SGLANG_PLUGIN_TARGETS:
+        raise RuntimeError("Vortex SGLang hook target list is out of sync")
+    for target, hook, hook_type in hooks:
+        HookRegistry.register(target, hook, hook_type)
+
+    logger.info("Registered Vortex hooks for SGLang 0.5.12")
diff --git a/vortex_torch/engine/sgl/transformers_compat.py b/vortex_torch/engine/sgl/transformers_compat.py
new file mode 100644
index 00000000..b581659a
--- /dev/null
+++ b/vortex_torch/engine/sgl/transformers_compat.py
@@ -0,0 +1,57 @@
+"""Compatibility fixes required by SGLang's pinned Transformers release."""
+from __future__ import annotations
+
+from functools import wraps
+
+
+class _NullAttentionSink:
+    def to(self, *_args, **_kwargs):
+        return None
+
+
+_NULL_ATTENTION_SINK = _NullAttentionSink()
+
+
+def patch_transformers_560_flash_attention() -> bool:
+    """Backport the sole 5.6.1 code fix onto SGLang's pinned 5.6.0.
+
+    Transformers 5.6.0 unconditionally calls ``s_aux.to(...)`` even for
+    models without attention sinks. 5.6.1 only adds a None guard. Supplying a
+    sentinel whose ``to`` returns None is behaviorally identical and avoids
+    copying the upstream attention implementation.
+    """
+    try:
+        import transformers
+        from transformers.integrations import flash_attention
+    except ImportError:
+        return False
+
+    if transformers.__version__ != "5.6.0":
+        return False
+    original = flash_attention.flash_attention_forward
+    if getattr(original, "_vortex_transformers_560_fix", False):
+        return True
+
+    @wraps(original)
+    def fixed(*args, s_aux=None, **kwargs):
+        return original(
+            *args,
+            s_aux=_NULL_ATTENTION_SINK if s_aux is None else s_aux,
+            **kwargs,
+        )
+
+    fixed._vortex_transformers_560_fix = True
+    flash_attention.flash_attention_forward = fixed
+
+    # modeling_utils copies the function into a class-wide dispatch table.
+    # Updating that table covers modules imported both before and after this
+    # compatibility hook.
+    from transformers.modeling_utils import AttentionInterface
+
+    for name in (
+        "flash_attention_2",
+        "flash_attention_3",
+        "flash_attention_4",
+    ):
+        AttentionInterface.register(name, fixed)
+    return True
diff --git a/vortex_torch/version.py b/vortex_torch/version.py
index 5becc17c..3d187266 100644
--- a/vortex_torch/version.py
+++ b/vortex_torch/version.py
@@ -1 +1 @@
-__version__ = "1.0.0"
+__version__ = "0.5.0"

From cacc9081f326ded6f3ed74356683816acc757886 Mon Sep 17 00:00:00 2001
From: rockyeast <rockyeast@example.local>
Date: Wed, 22 Jul 2026 00:38:01 -0400
Subject: [PATCH 02/10] test: cover SGLang plugin runtime contract

---
 .gitignore                            |  3 +-
 tests/test_sglang_runtime_contract.py | 61 +++++++++++++++++++++++++++
 2 files changed, 63 insertions(+), 1 deletion(-)
 create mode 100644 tests/test_sglang_runtime_contract.py

diff --git a/.gitignore b/.gitignore
index a214fe8b..373010bd 100644
--- a/.gitignore
+++ b/.gitignore
@@ -64,7 +64,8 @@ MANIFEST
 !vortex_torch/custom_ops/**/meta.json
 !vortex_torch/custom_ops/**/config.json
 
-tests/
+tests/*
+!tests/test_sglang_runtime_contract.py
 *_backup.py
 # PyInstaller
 #  Usually these files are written by a python script from a template
diff --git a/tests/test_sglang_runtime_contract.py b/tests/test_sglang_runtime_contract.py
new file mode 100644
index 00000000..d055a7aa
--- /dev/null
+++ b/tests/test_sglang_runtime_contract.py
@@ -0,0 +1,61 @@
+import sys
+from dataclasses import make_dataclass
+from types import ModuleType
+import unittest
+from unittest.mock import patch
+
+from vortex_torch.engine.sgl.config import (
+    _REQUIRED_SGLANG_HOOKS,
+    VORTEX_SGLANG_ABI,
+    VORTEX_SGLANG_PLUGIN_TARGETS,
+    validate_sglang_runtime_contract,
+)
+
+
+def _fake_plugin_sglang(*, abi=VORTEX_SGLANG_ABI, hooks=(), fields=()):
+    sglang = ModuleType("sglang")
+    srt = ModuleType("sglang.srt")
+    plugins = ModuleType("sglang.srt.plugins")
+    hook_registry = ModuleType("sglang.srt.plugins.hook_registry")
+    server_args = ModuleType("sglang.srt.server_args")
+
+    server_args.ServerArgs = make_dataclass("ServerArgs", fields)
+    server_args.ServerArgs._vortex_sglang_abi = abi
+    server_args.ServerArgs._vortex_sglang_hooks = frozenset(hooks)
+    server_args.ServerArgs._vortex_sglang_targets = VORTEX_SGLANG_PLUGIN_TARGETS
+    hook_registry.HookRegistry = type(
+        "HookRegistry", (), {"_patched": set(VORTEX_SGLANG_PLUGIN_TARGETS)}
+    )
+    plugins.load_plugins = lambda: None
+    sglang.srt = srt
+
+    return patch.dict(
+        sys.modules,
+        {
+            "sglang": sglang,
+            "sglang.srt": srt,
+            "sglang.srt.plugins": plugins,
+            "sglang.srt.plugins.hook_registry": hook_registry,
+            "sglang.srt.server_args": server_args,
+        },
+    )
+
+
+class SGLangRuntimeContractTest(unittest.TestCase):
+    def test_contract_accepts_complete_plugin_runtime(self):
+        with _fake_plugin_sglang(
+            hooks=_REQUIRED_SGLANG_HOOKS,
+            fields=[("vortex", object)],
+        ):
+            validate_sglang_runtime_contract()
+
+    def test_contract_rejects_incomplete_runtime(self):
+        with _fake_plugin_sglang(abi=None):
+            with self.assertRaisesRegex(
+                RuntimeError, "Incompatible official SGLang plugin runtime"
+            ):
+                validate_sglang_runtime_contract()
+
+
+if __name__ == "__main__":
+    unittest.main()

From e47bf183c8341f1f06fd88c691734f64d35d3012 Mon Sep 17 00:00:00 2001
From: rockyeast <rockyeast@example.local>
Date: Wed, 22 Jul 2026 00:38:11 -0400
Subject: [PATCH 03/10] fix: align Vortex runtime with SGLang 0.5.12

---
 .../sgl/attention_backend/flashinfer.py       | 15 ++++++++++-----
 .../sgl/attention_backend/triton_mla.py       |  1 +
 .../engine/sgl/attention_backend/trtllm.py    | 11 ++++++++---
 .../sgl/attention_backend/trtllm_mla.py       |  1 +
 vortex_torch/engine/sgl/memory_pool.py        | 19 ++++++++++++-------
 vortex_torch/engine/sgl/memory_pool_mla.py    |  1 +
 vortex_torch/indexer/utils_sglang.py          | 11 +++++++++++
 7 files changed, 44 insertions(+), 15 deletions(-)

diff --git a/vortex_torch/engine/sgl/attention_backend/flashinfer.py b/vortex_torch/engine/sgl/attention_backend/flashinfer.py
index 0c48478b..0287c89c 100644
--- a/vortex_torch/engine/sgl/attention_backend/flashinfer.py
+++ b/vortex_torch/engine/sgl/attention_backend/flashinfer.py
@@ -22,8 +22,9 @@
     get_chunkwise_nh2hn_transpose,
     get_decode_planner,
     get_prefill_planner,
+    normalize_prefill_seq_lens,
 )
-if os.environ["SGLANG_ENABLE_TORCH_COMPILE"] == "1":
+if os.environ.get("SGLANG_ENABLE_TORCH_COMPILE") == "1":
     import logging
 
     torch._logging.set_logs(dynamo=logging.ERROR)
@@ -327,7 +328,10 @@ def init_forward_metadata(self, forward_batch: ForwardBatch):
 
         elif forward_batch.forward_mode.is_extend():
             
-            prefix_lens = forward_batch.extend_prefix_lens
+            prefix_lens, input_seq_lens = normalize_prefill_seq_lens(
+                forward_batch.seq_lens,
+                forward_batch.extend_prefix_lens,
+            )
             extend_no_prefix = not any(forward_batch.extend_prefix_lens_cpu)
             bs = len(forward_batch.req_pool_indices)
             
@@ -335,7 +339,7 @@ def init_forward_metadata(self, forward_batch: ForwardBatch):
                 cached_seq_lens=prefix_lens,
                 dense_kv_indptr=self.kv_indptr_prefill[:bs*self.num_kv_heads+1],
                 dense_kv_indices=self.kv_indices_prefill,
-                input_seq_lens=(forward_batch.seq_lens.to(torch.int32) - prefix_lens),
+                input_seq_lens=input_seq_lens,
                 qo_indptr_ragged=self.qo_indptr[0][:bs+1],
                 qo_indptr_paged=self.qo_indptr[1][:bs*self.num_kv_heads+1],
                 kv_last_page_len=self.kv_last_page_len_prefill[:bs*self.num_kv_heads],
@@ -644,7 +648,8 @@ def forward_decode(
                 q=q,
                 o=self.forward_metadata.decode_wrappers[1]._paged_kv_indices_buf,
                 cache=cache,
-                ctx=self.ctx
+                ctx=self.ctx,
+                cur_layer=layer.layer_id,
             )
 
             # Sparse attention compute
@@ -669,4 +674,4 @@ def forward_decode(
             )
 
         # Restore to merged head dimension
-        return o.view(-1, layer.tp_q_head_num * layer.head_dim)
\ No newline at end of file
+        return o.view(-1, layer.tp_q_head_num * layer.head_dim)
diff --git a/vortex_torch/engine/sgl/attention_backend/triton_mla.py b/vortex_torch/engine/sgl/attention_backend/triton_mla.py
index 80304bd7..f0d87ba8 100644
--- a/vortex_torch/engine/sgl/attention_backend/triton_mla.py
+++ b/vortex_torch/engine/sgl/attention_backend/triton_mla.py
@@ -225,6 +225,7 @@ def forward_decode(
         cache = forward_batch.token_to_kv_pool.get_cache(layer.layer_id)
         self.compiled_indexer.forward(
             q=query, o=md.sparse_block_tables, cache=cache, ctx=self.ctx,
+            cur_layer=layer.layer_id,
         )
 
         # 3) block-sparse MLA decode in Triton over the fused latent.
diff --git a/vortex_torch/engine/sgl/attention_backend/trtllm.py b/vortex_torch/engine/sgl/attention_backend/trtllm.py
index 9090b841..352530ab 100644
--- a/vortex_torch/engine/sgl/attention_backend/trtllm.py
+++ b/vortex_torch/engine/sgl/attention_backend/trtllm.py
@@ -20,6 +20,7 @@
     get_chunkwise_nh2hn_transpose,
     get_decode_planner_trtllm,
     get_prefill_planner,
+    normalize_prefill_seq_lens,
 )
 if os.environ["SGLANG_ENABLE_TORCH_COMPILE"] == "1":
     import logging
@@ -319,7 +320,10 @@ def init_forward_metadata(self, forward_batch: ForwardBatch):
 
         elif forward_batch.forward_mode.is_extend():
             
-            prefix_lens = forward_batch.extend_prefix_lens
+            prefix_lens, input_seq_lens = normalize_prefill_seq_lens(
+                forward_batch.seq_lens,
+                forward_batch.extend_prefix_lens,
+            )
             extend_no_prefix = not any(forward_batch.extend_prefix_lens_cpu)
             bs = len(forward_batch.req_pool_indices)
             
@@ -327,7 +331,7 @@ def init_forward_metadata(self, forward_batch: ForwardBatch):
                 cached_seq_lens=prefix_lens,
                 dense_kv_indptr=self.kv_indptr_prefill[:bs*self.num_kv_heads+1],
                 dense_kv_indices=self.kv_indices_prefill,
-                input_seq_lens=(forward_batch.seq_lens.to(torch.int32) - prefix_lens),
+                input_seq_lens=input_seq_lens,
                 qo_indptr_ragged=self.qo_indptr[0][:bs+1],
                 qo_indptr_paged=self.qo_indptr[1][:bs*self.num_kv_heads+1],
                 kv_last_page_len=self.kv_last_page_len_prefill[:bs*self.num_kv_heads],
@@ -584,6 +588,7 @@ def forward_decode(
                 o=self.ctx.metadata.sparse_block_tables,
                 cache=cache,
                 ctx=self.ctx,
+                cur_layer=layer.layer_id,
             )
             o = trtllm_batch_decode_with_kv_cache(
                 query=q,
@@ -611,4 +616,4 @@ def forward_decode(
             )
 
         # Restore to merged head dimension
-        return o.view(-1, layer.tp_q_head_num * layer.head_dim)
\ No newline at end of file
+        return o.view(-1, layer.tp_q_head_num * layer.head_dim)
diff --git a/vortex_torch/engine/sgl/attention_backend/trtllm_mla.py b/vortex_torch/engine/sgl/attention_backend/trtllm_mla.py
index 82b9560d..2a1799f4 100644
--- a/vortex_torch/engine/sgl/attention_backend/trtllm_mla.py
+++ b/vortex_torch/engine/sgl/attention_backend/trtllm_mla.py
@@ -268,6 +268,7 @@ def forward_decode(
         cache = forward_batch.token_to_kv_pool.get_cache(layer.layer_id)
         self.compiled_indexer.forward(
             q=query, o=md.sparse_block_tables, cache=cache, ctx=self.ctx,
+            cur_layer=layer.layer_id,
         )
 
         # 3) MLA decode over the selected pages, on the fused latent.
diff --git a/vortex_torch/engine/sgl/memory_pool.py b/vortex_torch/engine/sgl/memory_pool.py
index 35c26c87..cd56f318 100644
--- a/vortex_torch/engine/sgl/memory_pool.py
+++ b/vortex_torch/engine/sgl/memory_pool.py
@@ -57,6 +57,8 @@
 
 class VortexCachePool(KVCache):
 
+    is_vortex_pool = True
+
     # Vortex stores K/V in a block-interleaved layout (see
     # vortex_torch/cache/triton_kernels/set_kv.py — position is mapped to
     # ``(token//page) * (page * num_kv_head) + head * page + token%page``).
@@ -205,8 +207,7 @@ def get_cache_size_bytes(self) -> int:
         return total_bytes
     
     def get_kv_size_bytes(self):
-        
-        raise NotImplementedError
+        return self.get_cache_size_bytes()
     
     # for disagg (PD disaggregation, Option B)
     def get_contiguous_buf_infos(self):
@@ -267,17 +268,18 @@ def rebuild_aux(self, loc: torch.Tensor):
             if layer_id in self.layers_skip:
                 continue
             self.compiled_cache.forward(
-                self.cache[layer_id - self.start_layer], loc, ctx=self.ctx
+                self.cache[layer_id - self.start_layer], loc, ctx=self.ctx,
+                cur_layer=layer_id,
             )
 
     def maybe_get_custom_mem_pool(self):
         return self.custom_mem_pool
 
-    def get_cpu_copy(self, indices):
+    def get_cpu_copy(self, indices, mamba_indices=None):
         
         raise NotImplementedError
 
-    def load_cpu_copy(self, kv_cache_cpu, indices):
+    def load_cpu_copy(self, kv_cache_cpu, indices, mamba_indices=None):
         
         raise NotImplementedError
 
@@ -359,8 +361,11 @@ def set_kv_buffer(
         )
         if layer_id in self.layers_skip:
             return
-        self.compiled_cache.forward(self.cache[layer_id - self.start_layer], loc, ctx=self.ctx)
+        self.compiled_cache.forward(
+            self.cache[layer_id - self.start_layer], loc, ctx=self.ctx,
+            cur_layer=layer_id,
+        )
         
     def move_kv_cache(self, tgt_loc: torch.Tensor, src_loc: torch.Tensor):
         
-        raise NotImplementedError
\ No newline at end of file
+        raise NotImplementedError
diff --git a/vortex_torch/engine/sgl/memory_pool_mla.py b/vortex_torch/engine/sgl/memory_pool_mla.py
index fe02b6f0..0ed55aff 100644
--- a/vortex_torch/engine/sgl/memory_pool_mla.py
+++ b/vortex_torch/engine/sgl/memory_pool_mla.py
@@ -33,6 +33,7 @@
 
 class VortexMLACachePool(MLATokenToKVPool):
 
+    is_vortex_pool = True
     supports_fused_set_kv_buffer = False
 
     def __init__(
diff --git a/vortex_torch/indexer/utils_sglang.py b/vortex_torch/indexer/utils_sglang.py
index 3f5eb11e..e36fb655 100644
--- a/vortex_torch/indexer/utils_sglang.py
+++ b/vortex_torch/indexer/utils_sglang.py
@@ -5,6 +5,17 @@
 from .prefill_sglang import get_sglang_prefill_module
 
 
+def normalize_prefill_seq_lens(
+    seq_lens: torch.Tensor,
+    prefix_lens: torch.Tensor,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Return the int32 lengths required by the prefill planner ABI."""
+
+    prefix_lens_i32 = prefix_lens.to(torch.int32)
+    input_seq_lens_i32 = seq_lens.to(torch.int32) - prefix_lens_i32
+    return prefix_lens_i32, input_seq_lens_i32
+
+
 def get_decode_planner(policy: str = None):
 
     module = get_sglang_plan_decode_v2_module(

From 4501d0426044ba44870e202ec651068073b3ab53 Mon Sep 17 00:00:00 2001
From: rockyeast <rockyeast@example.local>
Date: Wed, 22 Jul 2026 00:38:17 -0400
Subject: [PATCH 04/10] fix: register transpose outputs in the compiler graph

---
 vortex_torch/indexer/transpose.py | 14 +++++++++-----
 1 file changed, 9 insertions(+), 5 deletions(-)

diff --git a/vortex_torch/indexer/transpose.py b/vortex_torch/indexer/transpose.py
index c1019973..5d3a0db2 100644
--- a/vortex_torch/indexer/transpose.py
+++ b/vortex_torch/indexer/transpose.py
@@ -45,18 +45,16 @@ def profile(self, x: vTensor, ctx: Context) -> vTensor:
             FORMAT.BATCHED if x._format == FORMAT.BATCHED else FORMAT.RAGGED
         )
 
-        # Allocate output buffer: [S, D1, D0]
-        # S is derived from runtime context (number of pages/tokens in the pipeline)
-        S = ctx.max_num_pages
+        # The leading axis is a dynamic batch/page placeholder at trace time.
         D0, D1 = x.shape[1], x.shape[2]
         # Pure-metadata vTensor — no real allocation. The compiled code
         # supplies storage; we only need shape/dtype/device for codegen.
         self.output_buffer = vTensor(
-            shape=(S, D1, D0),
+            shape=(0, D1, D0),
             dtype=ctx.vortex_dtype,
             device=x.device,
             _format=self.output_format,
-            tensor_id=-1,  # set by caller / graph registration if needed
+            tensor_id=len(ctx.tensor_list),
         )
 
         for t in [x]:
@@ -65,4 +63,10 @@ def profile(self, x: vTensor, ctx: Context) -> vTensor:
                     t.shape[1] * t.shape[2]
                 )
 
+        ctx.tensor_list.append(self.output_buffer)
+        ctx.output_tensor_to_op_list.append(len(ctx.op_list))
+        ctx.op_list.append(self)
+        ctx.op_to_input_tensor_list.append([x.tensor_id])
+        ctx.op_to_output_tensor_list.append([self.output_buffer.tensor_id])
+
         return self.output_buffer

From 61f19cc41c840ed4137b19b37e3d37bf34dc1ae5 Mon Sep 17 00:00:00 2001
From: rockyeast <rockyeast@example.local>
Date: Wed, 22 Jul 2026 00:38:32 -0400
Subject: [PATCH 05/10] fix: make generated modules collision-safe

---
 .gitignore                                 |  1 +
 tests/test_codegen_io.py                   | 55 ++++++++++++++++
 vortex_torch/_codegen_io.py                | 74 ++++++++++++++++++++++
 vortex_torch/cache/compiler/interface.py   | 24 +++----
 vortex_torch/indexer/compiler/interface.py | 29 ++++-----
 5 files changed, 149 insertions(+), 34 deletions(-)
 create mode 100644 tests/test_codegen_io.py
 create mode 100644 vortex_torch/_codegen_io.py

diff --git a/.gitignore b/.gitignore
index 373010bd..698d83f5 100644
--- a/.gitignore
+++ b/.gitignore
@@ -65,6 +65,7 @@ MANIFEST
 !vortex_torch/custom_ops/**/config.json
 
 tests/*
+!tests/test_codegen_io.py
 !tests/test_sglang_runtime_contract.py
 *_backup.py
 # PyInstaller
diff --git a/tests/test_codegen_io.py b/tests/test_codegen_io.py
new file mode 100644
index 00000000..6bde2ffd
--- /dev/null
+++ b/tests/test_codegen_io.py
@@ -0,0 +1,55 @@
+from concurrent.futures import ThreadPoolExecutor
+from pathlib import Path
+import tempfile
+import unittest
+
+from vortex_torch._codegen_io import write_generated_module
+
+
+class CodegenIOTest(unittest.TestCase):
+    def test_generated_modules_are_namespaced_and_content_addressed(self):
+        with tempfile.TemporaryDirectory() as directory:
+            common = {
+                "cache_dir": directory,
+                "flow_name": "same/flow",
+            }
+
+            indexer = write_generated_module(
+                **common, namespace="indexer", source="VALUE = 1\n"
+            )
+            cache = write_generated_module(
+                **common, namespace="cache", source="VALUE = 1\n"
+            )
+            variant = write_generated_module(
+                **common, namespace="indexer", source="VALUE = 2\n"
+            )
+
+            self.assertEqual(len({indexer, cache, variant}), 3)
+            self.assertEqual(Path(indexer).read_text(), "VALUE = 1\n")
+            self.assertEqual(Path(cache).read_text(), "VALUE = 1\n")
+            self.assertEqual(Path(variant).read_text(), "VALUE = 2\n")
+
+    def test_concurrent_writers_only_publish_complete_files(self):
+        source = "VALUE = 'complete'\n" * 4096
+
+        with tempfile.TemporaryDirectory() as directory:
+            root = Path(directory)
+
+            def write_once(_):
+                return write_generated_module(
+                    cache_dir=directory,
+                    flow_name="shared",
+                    namespace="indexer",
+                    source=source,
+                )
+
+            with ThreadPoolExecutor(max_workers=16) as executor:
+                paths = list(executor.map(write_once, range(64)))
+
+            self.assertEqual(len(set(paths)), 1)
+            self.assertEqual(Path(paths[0]).read_text(), source)
+            self.assertFalse(list(root.glob("*.tmp")))
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/vortex_torch/_codegen_io.py b/vortex_torch/_codegen_io.py
new file mode 100644
index 00000000..cabb87c9
--- /dev/null
+++ b/vortex_torch/_codegen_io.py
@@ -0,0 +1,74 @@
+"""Safe filesystem helpers for generated Python modules."""
+
+from __future__ import annotations
+
+import hashlib
+import os
+from pathlib import Path
+import re
+import tempfile
+from typing import Optional
+
+
+def _safe_stem(value: str) -> str:
+    stem = re.sub(r"[^A-Za-z0-9_.-]+", "_", value).strip("._")
+    return stem or "vortex"
+
+
+def _cache_root(configured: Optional[str]) -> Path:
+    if configured:
+        root = Path(configured).expanduser().resolve()
+        root.mkdir(parents=True, exist_ok=True)
+        return root
+
+    env_root = os.environ.get("VORTEX_CODEGEN_CACHE_DIR")
+    root = (
+        Path(env_root).expanduser()
+        if env_root
+        else Path.home() / ".cache" / "vortex_torch" / "generated"
+    )
+    try:
+        root.mkdir(parents=True, exist_ok=True)
+    except OSError:
+        root = Path(tempfile.gettempdir()) / "vortex_torch" / "generated"
+        root.mkdir(parents=True, exist_ok=True)
+    return root.resolve()
+
+
+def write_generated_module(
+    *,
+    cache_dir: Optional[str],
+    flow_name: str,
+    namespace: str,
+    source: str,
+) -> str:
+    """Atomically publish a content-addressed generated Python module.
+
+    Separate namespaces prevent indexer and cache modules from overwriting each
+    other. The source digest keeps different model/config variants apart, while
+    ``os.replace`` ensures concurrent compiler processes only observe complete
+    files.
+    """
+
+    root = _cache_root(cache_dir)
+    digest = hashlib.sha256(source.encode("utf-8")).hexdigest()[:16]
+    filename = (
+        f"{_safe_stem(flow_name)}_{_safe_stem(namespace)}_"
+        f"{digest}_compiled_func.py"
+    )
+    destination = root / filename
+
+    fd, temporary = tempfile.mkstemp(
+        prefix=f".{filename}.", suffix=".tmp", dir=str(root)
+    )
+    try:
+        with os.fdopen(fd, "w", encoding="utf-8", newline="\n") as handle:
+            handle.write(source)
+            handle.flush()
+            os.fsync(handle.fileno())
+        os.replace(temporary, destination)
+    finally:
+        if os.path.exists(temporary):
+            os.unlink(temporary)
+
+    return str(destination)
diff --git a/vortex_torch/cache/compiler/interface.py b/vortex_torch/cache/compiler/interface.py
index 485ce6fd..4dfb0afd 100644
--- a/vortex_torch/cache/compiler/interface.py
+++ b/vortex_torch/cache/compiler/interface.py
@@ -16,25 +16,13 @@
 from ..context import Context
 from ...utils import INDENT, indent_block
 from ...abs import FORMAT
+from ..._codegen_io import write_generated_module
 from .impl import AVAILABLE_IMPL_BACKENDS
-import os
 
 
 def generate_interface(full_graph: Graph, sub_graphs: List[Graph], ctx: Context) -> Tuple[str, str]:
     """Emit the compiled module to disk and return ``(file_path, class_name)``."""
 
-    cache_dir = ctx.compilation_cache_dir or os.path.dirname(__file__)
-    cache_dir = os.path.expanduser(cache_dir)
-    cache_dir = os.path.abspath(cache_dir)
-
-    if not os.path.exists(cache_dir):
-        os.makedirs(cache_dir, exist_ok=True)
-    dst = os.path.join(
-        cache_dir,
-        f"{ctx.sparse_attention_name}_compiled_func.py",
-    )
-    print(f"Generating compiled cache function interface at {dst}")
-
     body_parts: List[str] = []
 
     for sub_graph_id, sub_graph in enumerate(sub_graphs):
@@ -48,9 +36,13 @@ def generate_interface(full_graph: Graph, sub_graphs: List[Graph], ctx: Context)
     body_str = "\n".join(body_parts)
 
     final_str = header_str + "\n\n" + auxilary_func_def_str + "\n\n" + body_str
-
-    with open(dst, "w") as f:
-        f.write(final_str)
+    dst = write_generated_module(
+        cache_dir=ctx.compilation_cache_dir,
+        flow_name=ctx.sparse_attention_name,
+        namespace="cache",
+        source=final_str,
+    )
+    print(f"Generating compiled cache function interface at {dst}")
 
     return dst, f"{ctx.sparse_attention_name}_CompiledFunc"
 
diff --git a/vortex_torch/indexer/compiler/interface.py b/vortex_torch/indexer/compiler/interface.py
index 220fd172..887b2883 100644
--- a/vortex_torch/indexer/compiler/interface.py
+++ b/vortex_torch/indexer/compiler/interface.py
@@ -3,22 +3,11 @@
 from ..context import Context
 from ...utils import Schedule, INDENT, indent_block
 from ...abs import FORMAT
+from ..._codegen_io import write_generated_module
 from .impl import AVAILABLE_IMPL_BACKENDS
-import os
-def generate_interface(full_graph: Graph, sub_graphs: list[Graph], ctx: Context) -> str:
 
-    cache_dir = ctx.compilation_cache_dir or os.path.dirname(__file__)
-    cache_dir = os.path.expanduser(cache_dir)
-    cache_dir = os.path.abspath(cache_dir)
-    
-    if not os.path.exists(cache_dir):
-        os.makedirs(cache_dir, exist_ok=True)
-    dst = os.path.join(
-        cache_dir,
-        f"{ctx.sparse_attention_name}_compiled_func.py"
-    )
-    print(f"Generating compiled function interface at {dst}")
-    
+
+def generate_interface(full_graph: Graph, sub_graphs: list[Graph], ctx: Context) -> Tuple[str, str]:
     body_parts: list[str] = []
 
     for sub_graph_id, sub_graph in enumerate(sub_graphs):
@@ -34,9 +23,13 @@ def generate_interface(full_graph: Graph, sub_graphs: list[Graph], ctx: Context)
     body_str = "\n".join(body_parts)
 
     final_str = header_str + "\n\n" + auxilary_func_def_str + "\n\n" + body_str
-
-    with open(dst, "w") as f:
-        f.write(final_str)
+    dst = write_generated_module(
+        cache_dir=ctx.compilation_cache_dir,
+        flow_name=ctx.sparse_attention_name,
+        namespace="indexer",
+        source=final_str,
+    )
+    print(f"Generating compiled function interface at {dst}")
 
     return dst, f"{ctx.sparse_attention_name}_CompiledFunc"
 
@@ -182,4 +175,4 @@ class {ctx.sparse_attention_name}_CompiledFunc:
 {INDENT}def forward(self, {entry_point_arg_str}):
 {entry_point_impl_str}
 """ 
-    return entry_cls_str
\ No newline at end of file
+    return entry_cls_str

From 7ae579811d6a48af34ef4f435f605e8f8f1c3430 Mon Sep 17 00:00:00 2001
From: rockyeast <rockyeast@example.local>
Date: Wed, 22 Jul 2026 00:38:40 -0400
Subject: [PATCH 06/10] docs: install Vortex with official SGLang

---
 README.md               | 36 ++++++----------
 docs/examples.md        |  8 ++--
 docs/installation.md    | 44 ++++++++-----------
 docs/landing/index.html |  9 +---
 docs/quickstart.md      |  9 ++--
 install_vortex.sh       | 77 ++++++++-------------------------
 install_vortex_glm.sh   | 95 ++---------------------------------------
 7 files changed, 62 insertions(+), 216 deletions(-)

diff --git a/README.md b/README.md
index 283dbb60..e5cbaf12 100644
--- a/README.md
+++ b/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
+`vortex_torch` here because the launch code constructs `VortexConfig`.
 
 ```python
 import sglang as sgl
-import vortex_torch  # noqa: F401 — import for side effect: installs the VortexConfig adapter
+import vortex_torch
 from vortex_torch.engine.sgl.config import VortexConfig
 
 llm = sgl.Engine(
@@ -471,12 +467,10 @@ 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. **Vortex is an SGLang plugin.** Installed SGLang discovers it from the
+   `sglang.srt.plugins` entry point. The script also imports `vortex_torch`
+   explicitly so it can run the compatibility preflight before starting
+   scheduler workers.
 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 +503,3 @@ If you find Vortex useful in your research, please cite:
       url={https://arxiv.org/abs/2606.06453}, 
 }
 ```
-
-
diff --git a/docs/examples.md b/docs/examples.md
index 2410180e..12e8fcd8 100644
--- a/docs/examples.md
+++ b/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>'`.
diff --git a/docs/installation.md b/docs/installation.md
index ebd57840..2d27d3f0 100644
--- a/docs/installation.md
+++ b/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).
diff --git a/docs/landing/index.html b/docs/landing/index.html
index 86f81e64..25c83b74 100644
--- a/docs/landing/index.html
+++ b/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">
diff --git a/docs/quickstart.md b/docs/quickstart.md
index d7553fe7..f0671f98 100644
--- a/docs/quickstart.md
+++ b/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(
diff --git a/install_vortex.sh b/install_vortex.sh
index cc416c8e..aeb1df14 100755
--- a/install_vortex.sh
+++ b/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"
diff --git a/install_vortex_glm.sh b/install_vortex_glm.sh
index 6b0e868b..1301ca57 100755
--- a/install_vortex_glm.sh
+++ b/install_vortex_glm.sh
@@ -1,95 +1,8 @@
 #!/usr/bin/env bash
-#
-# Reproducible build of the `vortex_glm` conda environment.
-#
-# `vortex_glm` is the env for GLM-4.7-Flash (HF model type `glm4_moe_lite`), which
-# REQUIRES transformers >= 5.0 — `vortex_v1` (transformers 4.57.1) cannot load it.
-# Both sglang and vortex_torch pin `transformers==4.57.1`, so the build installs
-# them first and then OVERRIDES transformers with the pinned GLM-supporting git
-# commit as the final step. flashinfer 0.6.3 (used by the MLA prefill) comes in as
-# an sglang base dep; flash-attention is never installed (not a dependency).
-#
-# Captured from the working env: python 3.12, torch 2.9.1+cu128, flashinfer 0.6.3,
-# sglang (editable, vendored v0.5.9), transformers @ 76732b4 (5.0.0.dev0).
-#
-# Usage:
-#   bash install_vortex_glm.sh            # create the env
-#   FORCE=1 bash install_vortex_glm.sh    # remove an existing env first
-#   ENV_NAME=glm2 bash install_vortex_glm.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.
+# Compatibility wrapper: SGLang 0.5.12 already carries Transformers 5.x.
 
 set -euo pipefail
 
-ENV_NAME="${ENV_NAME:-vortex_glm}"
-PY_VER="${PY_VER:-3.12}"
-# transformers commit that ships glm4_moe_lite support (== 5.0.0.dev0 in the env).
-TRANSFORMERS_COMMIT="${TRANSFORMERS_COMMIT:-76732b4e7120808ff989edbd16401f61fa6a0afa}"
-
-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
-        exit 1
-    fi
-fi
-
-echo ">>> [1/5] 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
-
-# ---- 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/5] 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
-
-# ---- 3. sglang (editable, vendored) ---------------------------------------
-# Pulls the runtime tree: flashinfer_python/cubin 0.6.3, sgl-kernel, xgrammar,
-# outlines, cuda-python, etc. (and transformers 4.57.1, overridden in step 5).
-echo ">>> [3/5] installing vendored sglang (editable) from $SGLANG_DIR"
-pip install -e "$SGLANG_DIR"
-
-# ---- 4. vortex_torch (editable) -------------------------------------------
-echo ">>> [4/5] installing vortex_torch (editable) from $REPO_ROOT"
-pip install -e "$REPO_ROOT"
-
-# ---- 5. transformers override (GLM-4.7 / glm4_moe_lite) -------------------
-# LAST so it wins over the transformers==4.57.1 pins above. pip will print a
-# dependency-conflict warning about that pin — expected and harmless here.
-echo ">>> [5/5] overriding transformers with glm4_moe_lite commit ${TRANSFORMERS_COMMIT:0:12}"
-pip install --upgrade "git+https://github.com/huggingface/transformers.git@${TRANSFORMERS_COMMIT}"
-
-# ---- verify ----------------------------------------------------------------
-echo ">>> verifying the environment"
-python - <<'PY'
-import torch, transformers, sglang, vortex_torch
-import flashinfer
-from transformers.models.auto.configuration_auto import CONFIG_MAPPING_NAMES
-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 = "glm4_moe_lite" in CONFIG_MAPPING_NAMES
-print(f"  glm4_moe_lite recognized: {ok}")
-assert ok, "transformers does not recognize glm4_moe_lite — GLM-4.7-Flash will NOT load"
-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 examples/ruler/run_ruler_mla.py --n 20 --dump"
+export ENV_NAME="${ENV_NAME:-vortex_glm}"
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+exec bash "$SCRIPT_DIR/install_vortex.sh"

From b3bfc7ab285466bedfba215106ea58a224df9886 Mon Sep 17 00:00:00 2001
From: rockyeast <rockyeast@example.local>
Date: Wed, 22 Jul 2026 02:03:44 -0400
Subject: [PATCH 07/10] fix: load SGLang plugins before parsing server args

---
 README.md                      | 15 ++++++++-------
 examples/misc/server_launch.sh | 17 +++++++----------
 2 files changed, 15 insertions(+), 17 deletions(-)

diff --git a/README.md b/README.md
index e5cbaf12..14c22551 100644
--- a/README.md
+++ b/README.md
@@ -222,12 +222,12 @@ class CustomSparseAttention(vFlow):
 ## 🏃 Launch it with SGLang
 
 The launch script is a **separate file** from the flow. The installed SGLang
-plugin wires Vortex into the decode loop automatically. Import
-`vortex_torch` here because the launch code constructs `VortexConfig`.
+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
 from vortex_torch.engine.sgl.config import VortexConfig
 
 llm = sgl.Engine(
@@ -467,10 +467,11 @@ examples/misc/server_launch.sh Qwen/Qwen3-4B 1
 
 Two details make server mode work:
 
-1. **Vortex is an SGLang plugin.** Installed SGLang discovers it from the
-   `sglang.srt.plugins` entry point. The script also imports `vortex_torch`
-   explicitly so it can run the compatibility preflight before starting
-   scheduler workers.
+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
diff --git a/examples/misc/server_launch.sh b/examples/misc/server_launch.sh
index 342d584d..f915ce76 100644
--- a/examples/misc/server_launch.sh
+++ b/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)

From 039af77f8214ab9ce2cd7fa6d0ca8d2b3ee36361 Mon Sep 17 00:00:00 2001
From: rockyeast <rockyeast@example.local>
Date: Wed, 22 Jul 2026 02:03:58 -0400
Subject: [PATCH 08/10] fix: validate GLM installer support

---
 install_vortex_glm.sh | 22 ++++++++++++++++++++--
 1 file changed, 20 insertions(+), 2 deletions(-)

diff --git a/install_vortex_glm.sh b/install_vortex_glm.sh
index 1301ca57..831fdceb 100755
--- a/install_vortex_glm.sh
+++ b/install_vortex_glm.sh
@@ -1,8 +1,26 @@
 #!/usr/bin/env bash
-# Compatibility wrapper: SGLang 0.5.12 already carries Transformers 5.x.
+# Compatibility wrapper around the shared installer, with a GLM capability
+# check so dependency resolution cannot silently select an incompatible
+# Transformers build.
 
 set -euo pipefail
 
 export ENV_NAME="${ENV_NAME:-vortex_glm}"
 SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
-exec bash "$SCRIPT_DIR/install_vortex.sh"
+bash "$SCRIPT_DIR/install_vortex.sh"
+
+source "$(conda info --base)/etc/profile.d/conda.sh"
+conda activate "$ENV_NAME"
+
+python - <<'PY'
+import transformers
+from transformers.models.auto.configuration_auto import CONFIG_MAPPING_NAMES
+
+supported = "glm4_moe_lite" in CONFIG_MAPPING_NAMES
+print(f"transformers          : {transformers.__version__}")
+print(f"glm4_moe_lite support : {supported}")
+assert supported, (
+    "Transformers does not recognize glm4_moe_lite; "
+    "the GLM-4.7-Flash environment is incomplete."
+)
+PY

From f55b13485391b27bf2bf5c1e64e2ae4594cbc320 Mon Sep 17 00:00:00 2001
From: rockyeast <rockyeast@example.local>
Date: Wed, 22 Jul 2026 02:22:05 -0400
Subject: [PATCH 09/10] fix: preserve Vortex planner batch limit

---
 tests/test_sglang_runtime_contract.py | 50 +++++++++++++++++++++++++++
 vortex_torch/engine/sgl/config.py     |  2 ++
 vortex_torch/engine/sgl/plugin.py     | 14 ++++++++
 3 files changed, 66 insertions(+)

diff --git a/tests/test_sglang_runtime_contract.py b/tests/test_sglang_runtime_contract.py
index d055a7aa..d9ad62d3 100644
--- a/tests/test_sglang_runtime_contract.py
+++ b/tests/test_sglang_runtime_contract.py
@@ -1,3 +1,4 @@
+import importlib
 import sys
 from dataclasses import make_dataclass
 from types import ModuleType
@@ -56,6 +57,55 @@ def test_contract_rejects_incomplete_runtime(self):
             ):
                 validate_sglang_runtime_contract()
 
+    def test_plugin_caps_vortex_request_pool_at_planner_limit(self):
+        server_args = ModuleType("sglang.srt.server_args")
+        server_args.ServerArgs = make_dataclass("ServerArgs", [])
+        sglang = ModuleType("sglang")
+        srt = ModuleType("sglang.srt")
+
+        module_name = "vortex_torch.engine.sgl.plugin"
+        previous = sys.modules.pop(module_name, None)
+        try:
+            with patch.dict(
+                sys.modules,
+                {
+                    "sglang": sglang,
+                    "sglang.srt": srt,
+                    "sglang.srt.server_args": server_args,
+                },
+            ):
+                plugin = importlib.import_module(module_name)
+
+                class Runner:
+                    pass
+
+                runner = Runner()
+                runner.server_args = type(
+                    "Args", (), {"enable_vortex_sparsity": True}
+                )()
+                original = lambda _runner, _capacity: 2048
+
+                self.assertEqual(
+                    plugin._around_resolve_max_num_reqs(original, runner, 8192),
+                    1024,
+                )
+                self.assertEqual(
+                    plugin._around_resolve_max_num_reqs(
+                        lambda _runner, _capacity: 512, runner, 8192
+                    ),
+                    512,
+                )
+
+                runner.server_args.enable_vortex_sparsity = False
+                self.assertEqual(
+                    plugin._around_resolve_max_num_reqs(original, runner, 8192),
+                    2048,
+                )
+        finally:
+            sys.modules.pop(module_name, None)
+            if previous is not None:
+                sys.modules[module_name] = previous
+
 
 if __name__ == "__main__":
     unittest.main()
diff --git a/vortex_torch/engine/sgl/config.py b/vortex_torch/engine/sgl/config.py
index d2d8b2f3..db83da6f 100644
--- a/vortex_torch/engine/sgl/config.py
+++ b/vortex_torch/engine/sgl/config.py
@@ -27,6 +27,7 @@
         "server_args.vortex_config",
         "model_runner.build_sparse_flow",
         "kv_cache.kv_cell_size",
+        "kv_cache.max_num_reqs",
         "kv_cache.make_kv_pool",
         "model_utils.fused_kv_opt_out",
         "disaggregation.rebuild_aux",
@@ -39,6 +40,7 @@
         "sglang.srt.server_args.ServerArgs.add_cli_args",
         "sglang.srt.model_executor.model_runner.ModelRunner.configure_kv_cache_dtype",
         "sglang.srt.model_executor.pool_configurator.DefaultPoolConfigurator._compute_cell_size",
+        "sglang.srt.model_executor.model_runner_kv_cache_mixin.ModelRunnerKVCacheMixin._resolve_max_num_reqs",
         "sglang.srt.model_executor.model_runner_kv_cache_mixin.ModelRunnerKVCacheMixin._init_pools",
         "sglang.srt.models.utils.enable_fused_set_kv_buffer",
         "sglang.srt.disaggregation.decode.DecodeTransferQueue.pop_transferred",
diff --git a/vortex_torch/engine/sgl/plugin.py b/vortex_torch/engine/sgl/plugin.py
index 4a7ef8df..cbf3a350 100644
--- a/vortex_torch/engine/sgl/plugin.py
+++ b/vortex_torch/engine/sgl/plugin.py
@@ -32,6 +32,7 @@
     )
 )[0]
 _VORTEX_LEGACY_DEFAULTS = legacy_defaults()
+_VORTEX_MAX_BATCH_SIZE = 1024
 
 
 def _server_args_getattr(self, name: str):
@@ -196,6 +197,14 @@ def _around_compute_cell_size(original, configurator, runner, num_layers):
     return kv_cell_size(runner, num_layers, element_size)
 
 
+def _around_resolve_max_num_reqs(original, runner, token_capacity):
+    max_num_reqs = original(runner, token_capacity)
+    if not runner.server_args.enable_vortex_sparsity:
+        return max_num_reqs
+    # Both SGL planner variants launch one 1024-thread block over the batch.
+    return min(max_num_reqs, _VORTEX_MAX_BATCH_SIZE)
+
+
 def _around_init_pools(original, runner, *args, **kwargs):
     if not runner.server_args.enable_vortex_sparsity:
         return original(runner, *args, **kwargs)
@@ -294,6 +303,11 @@ def register() -> None:
             _around_compute_cell_size,
             HookType.AROUND,
         ),
+        (
+            "sglang.srt.model_executor.model_runner_kv_cache_mixin.ModelRunnerKVCacheMixin._resolve_max_num_reqs",
+            _around_resolve_max_num_reqs,
+            HookType.AROUND,
+        ),
         (
             "sglang.srt.model_executor.model_runner_kv_cache_mixin.ModelRunnerKVCacheMixin._init_pools",
             _around_init_pools,

From af1954295fda4470f89b4745b2222561d22a0ee7 Mon Sep 17 00:00:00 2001
From: rockyeast <rockyeast@example.local>
Date: Thu, 23 Jul 2026 04:10:32 -0400
Subject: [PATCH 10/10] fix: avoid partial package imports in SGLang workers

---
 vortex_torch/engine/sgl/attention_backend/flashinfer.py | 2 +-
 vortex_torch/engine/sgl/attention_backend/trtllm.py     | 2 +-
 vortex_torch/engine/sgl/integration.py                  | 8 +++-----
 3 files changed, 5 insertions(+), 7 deletions(-)

diff --git a/vortex_torch/engine/sgl/attention_backend/flashinfer.py b/vortex_torch/engine/sgl/attention_backend/flashinfer.py
index 0287c89c..2ddb0678 100644
--- a/vortex_torch/engine/sgl/attention_backend/flashinfer.py
+++ b/vortex_torch/engine/sgl/attention_backend/flashinfer.py
@@ -13,7 +13,7 @@
 from typing import TYPE_CHECKING, Callable, List, Optional, Union, Dict, Tuple
 from functools import partial
 import torch
-from vortex_torch import is_hopper
+from vortex_torch.utils import is_hopper
 from vortex_torch.abs import as_vtensor, FORMAT
 from vortex_torch.indexer import Context, MetaData
 from vortex_torch.indexer.compiler.compile import compile as compile_indexer
diff --git a/vortex_torch/engine/sgl/attention_backend/trtllm.py b/vortex_torch/engine/sgl/attention_backend/trtllm.py
index 352530ab..c1d53ba9 100644
--- a/vortex_torch/engine/sgl/attention_backend/trtllm.py
+++ b/vortex_torch/engine/sgl/attention_backend/trtllm.py
@@ -11,7 +11,7 @@
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, List, Optional, Union, Dict
 import torch
-from vortex_torch import is_hopper
+from vortex_torch.utils import is_hopper
 from vortex_torch.abs import as_vtensor, FORMAT
 from vortex_torch.indexer import Context, MetaData
 from vortex_torch.indexer.compiler.compile import compile as compile_indexer
diff --git a/vortex_torch/engine/sgl/integration.py b/vortex_torch/engine/sgl/integration.py
index 2e18a813..d25ce400 100644
--- a/vortex_torch/engine/sgl/integration.py
+++ b/vortex_torch/engine/sgl/integration.py
@@ -152,12 +152,10 @@ def build_sparse_flow(runner) -> Optional[Any]:
     if not sa.enable_vortex_sparsity:
         return None
 
-    import vortex_torch
+    from vortex_torch.flow import build_vflow, vFlowMLA
 
-    flow = vortex_torch.flow.build_vflow(
-        sa.vortex_module_name, user_file=sa.vortex_module_path
-    )
-    if isinstance(flow, vortex_torch.flow.vFlowMLA):
+    flow = build_vflow(sa.vortex_module_name, user_file=sa.vortex_module_path)
+    if isinstance(flow, vFlowMLA):
         # MLA flow: latent geometry instead of a single head_dim.
         flow.initialize(
             block_size=runner.block_size,