Removing patching for SGLang versions

src/hyperloom/agents/kernel/scripts/install.sh

@@ -156,7 +156,7 @@ INFERENCEX_PATH="${INFERENCEX_PATH:-}"
 # The internal extension is used ONLY when $TRACELENS_INTERNAL_ROOT is set
 # (env / .env); leave it unset for the base-only report. No separate toggle.
 TRACELENS_REPO="https://github.com/AMD-AGI/TraceLens.git"
-TRACELENS_REF="c74d4d2ca48d6fcd7e7e829b409446000fe4300f"
+TRACELENS_REF="9fc0dc6487bde554c6ed314a15b61022e5ec62ea"
 # Operator override iff TRACELENS_ROOT points OUTSIDE the pod-local default.
 # The persistent kernel-agent env re-exports the resolved default path, so a
 # presence-only check (${VAR:+1}) would misclassify it as an override and skip

src/hyperloom/agents/kernel/tools/tracelens_analysis.py

@@ -5685,7 +5685,7 @@ def run_command(
 # Defaults kept in sync with src/hyperloom/agents/kernel/scripts/install.sh (TRACELENS_REPO /
 # TRACELENS_REF). Overridable via env so a run can pin its own SHA.
 _TRACELENS_REPO_DEFAULT = "https://github.com/AMD-AGI/TraceLens.git"
-_TRACELENS_REF_DEFAULT = "c74d4d2ca48d6fcd7e7e829b409446000fe4300f"
+_TRACELENS_REF_DEFAULT = "9fc0dc6487bde554c6ed314a15b61022e5ec62ea"
 
 
 def _default_tracelens_root() -> Path:

src/hyperloom/inference_optimizer/multi_node/commands/infera.py

@@ -90,6 +90,17 @@ def _collect_forward_env() -> dict[str, str]:
     trace_dir = os.environ.get("HYPERLOOM_MN_PROFILE_TRACE_DIR", "").strip()
     if trace_dir and "SGLANG_TORCH_PROFILER_DIR" not in fwd:
         fwd["SGLANG_TORCH_PROFILER_DIR"] = trace_dir
+    # Forward no-patch shape-discovery config; the pod-side launcher sets
+    # PYTHONPATH itself (it is blocked from SSH forwarding).
+    for _shape_key in (
+        "TRACELENS_ROOT",
+        "TRACELENS_SHAPE_DISCOVERY",
+        "HYPERLOOM_SGLANG_SHAPE_MODE",
+        "HYPERLOOM_SGLANG_VERSION_PIN",
+    ):
+        _shape_val = os.environ.get(_shape_key, "").strip()
+        if _shape_val and _shape_key not in fwd:
+            fwd[_shape_key] = _shape_val
     unset_fwd = os.environ.get("HYPERLOOM_MN_UNSET_FWD_ENV", "").strip()
     if unset_fwd:
         try:

src/hyperloom/inference_optimizer/multi_node/scripts/launch_infera_node.py

@@ -6,6 +6,7 @@
 import argparse
 import json
 import os
+import re
 import shlex
 import subprocess
 import sys
@@ -411,6 +412,46 @@ def _reap_stale_engine_ports() -> None:
                     break
 
 
+# SGLang >= 0.5.18 no-patch shape tool; gate mirrored inline (no hyperloom import).
+_KERNEL_SHAPE_TOOL_REL = ("TraceLens", "TraceUtils", "kernel_shape_tool")
+_SGLANG_SITECUSTOMIZE_MIN_VERSION = (0, 5, 18)
+
+
+def _sglang_shape_mode() -> str:
+    """Pod-side mirror of hyperloom's SGLang shape-mode gate (no hyperloom import)."""
+    override = os.environ.get("HYPERLOOM_SGLANG_SHAPE_MODE", "auto").strip().lower()
+    if override in {"patch", "patched"}:
+        return "patched"
+    if override == "sitecustomize":
+        return "sitecustomize"
+    version = ""
+    try:
+        import sglang  # type: ignore
+
+        version = (getattr(sglang, "__version__", "") or "").strip()
+    except Exception:  # noqa: BLE001
+        version = os.environ.get("HYPERLOOM_SGLANG_VERSION_PIN", "").strip()
+    m = re.match(r"^\s*v?(\d+(?:\.\d+)*)", version)
+    if not m:
+        return "patched"
+    vt = tuple(int(p) for p in m.group(1).split("."))
+    return "sitecustomize" if vt >= _SGLANG_SITECUSTOMIZE_MIN_VERSION else "patched"
+
+
+def _maybe_activate_kernel_shape_tool(env: dict[str, str]) -> None:
+    """SGLang >= 0.5.18: put the no-patch kernel_shape_tool on PYTHONPATH."""
+    root = (env.get("TRACELENS_ROOT") or os.environ.get("TRACELENS_ROOT") or "").strip()
+    if not root or _sglang_shape_mode() != "sitecustomize":
+        return
+    tool = Path(root).joinpath(*_KERNEL_SHAPE_TOOL_REL)
+    if not tool.is_dir():
+        _log(f"WARN kernel_shape_tool not found at {tool}; SGLang shape discovery disabled")
+        return
+    existing = (env.get("PYTHONPATH") or "").strip()
+    env["PYTHONPATH"] = f"{tool}{os.pathsep}{existing}" if existing else str(tool)
+    env.setdefault("TRACELENS_SHAPE_DISCOVERY", "1")
+
+
 def _build_sglang_cmd(
     a: argparse.Namespace,
     node_rank: int,
@@ -724,6 +765,7 @@ def main() -> int:
     if _shared_log_dir.startswith("/") and "$" not in _shared_log_dir:
         log_file = Path(_shared_log_dir) / f"mn_infera_server_{advertise_host}_r{node_rank}.log"
     if args.framework == "sglang":
+        _maybe_activate_kernel_shape_tool(env)
         cmd = _build_sglang_cmd(args, node_rank, leader, advertise_host=advertise_host)
         pid = _detach_launch(cmd, log_file, pid_file, env)
     else:

src/hyperloom/inference_optimizer/multi_node/scripts/launch_multinode.py

@@ -9,6 +9,7 @@
 import argparse
 import json
 import os
+import re
 import shlex
 import subprocess
 import pathlib
@@ -310,6 +311,46 @@ def _probe_mec_firmware_lt_177() -> bool:
     return False
 
 
+# SGLang >= 0.5.18 no-patch shape tool; gate mirrored inline (no hyperloom import).
+_KERNEL_SHAPE_TOOL_REL = ("TraceLens", "TraceUtils", "kernel_shape_tool")
+_SGLANG_SITECUSTOMIZE_MIN_VERSION = (0, 5, 18)
+
+
+def _sglang_shape_mode() -> str:
+    """Pod-side mirror of hyperloom's SGLang shape-mode gate (no hyperloom import)."""
+    override = os.environ.get("HYPERLOOM_SGLANG_SHAPE_MODE", "auto").strip().lower()
+    if override in {"patch", "patched"}:
+        return "patched"
+    if override == "sitecustomize":
+        return "sitecustomize"
+    version = ""
+    try:
+        import sglang  # type: ignore
+
+        version = (getattr(sglang, "__version__", "") or "").strip()
+    except Exception:  # noqa: BLE001
+        version = os.environ.get("HYPERLOOM_SGLANG_VERSION_PIN", "").strip()
+    m = re.match(r"^\s*v?(\d+(?:\.\d+)*)", version)
+    if not m:
+        return "patched"
+    vt = tuple(int(p) for p in m.group(1).split("."))
+    return "sitecustomize" if vt >= _SGLANG_SITECUSTOMIZE_MIN_VERSION else "patched"
+
+
+def _maybe_activate_kernel_shape_tool(sub_env: dict[str, str]) -> None:
+    """SGLang >= 0.5.18: put the no-patch kernel_shape_tool on PYTHONPATH."""
+    root = os.environ.get("TRACELENS_ROOT", "").strip()
+    if not root or _sglang_shape_mode() != "sitecustomize":
+        return
+    tool = Path(root).joinpath(*_KERNEL_SHAPE_TOOL_REL)
+    if not tool.is_dir():
+        sys.stderr.write(f"WARN kernel_shape_tool not found at {tool}; SGLang shape discovery disabled\n")
+        return
+    existing = sub_env.get("PYTHONPATH", "").strip()
+    sub_env["PYTHONPATH"] = f"{tool}{os.pathsep}{existing}" if existing else str(tool)
+    sub_env.setdefault("TRACELENS_SHAPE_DISCOVERY", "1")
+
+
 def _subprocess_env() -> dict[str, str]:
     """Build the framework launcher subprocess env."""
     env = dict(os.environ)
@@ -443,6 +484,7 @@ def _spawn_remote(
     dist_init_addr = f"{head_ip}:{dist_init_port}"
     fw = framework.lower()
     if fw == "sglang":
+        _maybe_activate_kernel_shape_tool(sub_env)
         cmd = _build_sglang_cmd(
             model=model,
             tp=tp,

src/hyperloom/orchestrator/actions/executors/_multi_node_server_lifecycle.py

@@ -372,12 +372,20 @@ async def restart_server_for_round(
 
         # Multi-node TraceLens SGLang patch fan-out (fail-soft).
         try:
-            from ._server_patcher import _tracelens_patch_enabled
+            from ._server_patcher import _tracelens_patch_enabled, resolve_sglang_shape_mode
         except Exception:  # noqa: BLE001
             _tracelens_patch_enabled_fn = lambda: True  # noqa: E731 - safe default
+            _sglang_shape_mode_val = "patched"
         else:
             _tracelens_patch_enabled_fn = _tracelens_patch_enabled
-        if _tracelens_patch_enabled_fn() and (os.environ.get("TRACELENS_ROOT", "").strip()):
+            _sglang_shape_mode_val = resolve_sglang_shape_mode()
+        if _sglang_shape_mode_val == "sitecustomize":
+            # sitecustomize mode: shapes come from the no-patch tool; skip the patch fan-out.
+            log.info(
+                "restart_server_for_round: SGLang shape mode=sitecustomize; "
+                "skipping TraceLens patch fan-out (shapes via kernel_shape_tool)."
+            )
+        elif _tracelens_patch_enabled_fn() and (os.environ.get("TRACELENS_ROOT", "").strip()):
             try:
                 from hyperloom.inference_optimizer.multi_node.cli import cmd_apply_tracelens_patch
 

src/hyperloom/orchestrator/actions/executors/_server_patcher.py

@@ -253,6 +253,64 @@ def ensure_sglang_patched_for_tracelens(
     return _ensure_patched(plan)
 
 
+# SGLang shape-discovery gate: >= 0.5.18 uses the no-patch TraceLens tool
+# (PYTHONPATH + sitecustomize + TRACELENS_SHAPE_DISCOVERY); older uses git-apply.
+_SGLANG_SITECUSTOMIZE_MIN_VERSION: tuple[int, ...] = (0, 5, 18)
+_SGLANG_SHAPE_MODE_ENV = "HYPERLOOM_SGLANG_SHAPE_MODE"
+# No-patch tool location, relative to TRACELENS_ROOT.
+_KERNEL_SHAPE_TOOL_REL: tuple[str, ...] = ("TraceLens", "TraceUtils", "kernel_shape_tool")
+
+
+def sglang_shape_mode(version: str) -> str:
+    """Return the shape-discovery mechanism for an SGLang version.
+
+    ``"sitecustomize"`` (>= 0.5.18) uses the no-patch TraceLens tool;
+    ``"patched"`` (< 0.5.18) uses the legacy ``git apply`` flow.
+    ``HYPERLOOM_SGLANG_SHAPE_MODE=patch|sitecustomize`` overrides the gate
+    (``auto`` / unset = version-based).
+    """
+    override = os.environ.get(_SGLANG_SHAPE_MODE_ENV, "auto").strip().lower()
+    if override in {"patch", "patched"}:
+        return "patched"
+    if override == "sitecustomize":
+        return "sitecustomize"
+    vt = _version_tuple(version)
+    if vt is None:
+        # Unparseable version: keep the safe legacy mechanism.
+        return "patched"
+    return "sitecustomize" if vt >= _SGLANG_SITECUSTOMIZE_MIN_VERSION else "patched"
+
+
+def kernel_shape_tool_dir(tracelens_root: Path | str | None = None) -> Path | None:
+    """Resolve the no-patch ``kernel_shape_tool`` dir under TRACELENS_ROOT, or None."""
+    root = _resolve_tracelens_root(tracelens_root)
+    if root is None:
+        return None
+    tool = root.joinpath(*_KERNEL_SHAPE_TOOL_REL)
+    return tool if tool.is_dir() else None
+
+
+def _detect_installed_sglang_version() -> str | None:
+    """Return the locally-installed SGLang version, or ``None`` if unimportable."""
+    try:
+        import sglang  # type: ignore  # noqa: I001 - runtime probe
+    except Exception:  # noqa: BLE001
+        return None
+    return (getattr(sglang, "__version__", "") or "").strip() or None
+
+
+def resolve_sglang_shape_mode() -> str:
+    """Resolve the SGLang shape mode from override -> local install -> MN version pin.
+
+    ``HYPERLOOM_SGLANG_SHAPE_MODE`` wins; otherwise the version comes from the
+    locally-installed SGLang (single-node / sandbox) or, when SGLang is not
+    importable in the controller (multi-node), ``HYPERLOOM_SGLANG_VERSION_PIN``.
+    Falls back to ``"patched"`` (legacy) when the version cannot be determined.
+    """
+    version = _detect_installed_sglang_version() or os.environ.get("HYPERLOOM_SGLANG_VERSION_PIN", "").strip()
+    return sglang_shape_mode(version)
+
+
 def ensure_sglang_patched_for_ck_blockscale(
     kernelforge_root: Path | str | None = None,
 ) -> bool:

src/hyperloom/orchestrator/actions/executors/_workload_envs.py

@@ -1513,7 +1513,13 @@ def materialize_config_with_envs(
         # try to patch, fall back to the safe set on failure. Default-on
         # (HYPERLOOM_ENABLE_PATCH=0 disables); skip for atom.
         tracelens_patch_ok = False
-        patch_attempted = _tracelens_patch_enabled() and not is_atom
+        # Function-local import to stay out of the module-level import cycle
+        # (matches _multi_node_server_lifecycle).
+        from ._server_patcher import kernel_shape_tool_dir, resolve_sglang_shape_mode
+
+        is_sglang = "sglang" in fw
+        sglang_sitecustomize = is_sglang and resolve_sglang_shape_mode() == "sitecustomize"
+        patch_attempted = _tracelens_patch_enabled() and not is_atom and not sglang_sitecustomize
         # Written in every branch, not only the failing one. "No status" used to mean both "patched fine" and
         # "never tried because the image already carries it", and those two call for different reactions when a
         # trace later turns up without annotations.
@@ -1612,39 +1618,49 @@ def materialize_config_with_envs(
                             "imprecise.",
                             _model,
                         )
-            # Both capture options are annotation-only and need TraceLens
-            # server-side support to land: without it the trace carries no
-            # ``kernel_shape_profiler`` events (trace-health check 5), so asking
-            # for them pays the capture cost for data nothing downstream reads.
-            # Keyed on the degraded *reason* rather than ``tracelens_patch_ok``:
-            # a patch that was never attempted (HYPERLOOM_ENABLE_PATCH=0) can
-            # still be baked into the image, and must keep the annotations.
-            _patch_degraded = envs.get("HYPERLOOM_PROFILE_DEGRADED_REASON") == _TRACELENS_PATCH_UNAVAILABLE
-            if _patch_degraded:
-                _shape_disc = False
-            extra_body["shape_discovery"] = _shape_disc
-            if _patch_degraded:
-                extra_body["detailed_annotations"] = False
-            else:
+            if sglang_sitecustomize:
+                # No-patch path: shapes come from the tool via PYTHONPATH, not a
+                # request-body flag or CUDA-graph arg (unpatched SGLang rejects both).
+                extra_body.pop("shape_discovery", None)
                 extra_body.setdefault("detailed_annotations", True)
-            # NOTE: this write happens before the per-task ``extra_envs`` merge, so
-            # an ``extra_envs`` entry for PROFILE_EXTRA_BODY can still drop
-            # start_step/num_steps the way ``args_mode="replace"`` used to drop
-            # vLLM's --profiler-config bounds. The vLLM side is re-asserted at the
-            # end of this function; SGLang is NOT, because deciding whether a
-            # non-positive num_steps means "unbounded" or "no capture" needs a
-            # SGLang-side answer this layer does not have. Every OOM observed so
-            # far was vLLM.
-            envs["PROFILE_EXTRA_BODY"] = _json.dumps(extra_body)
-            if tracelens_patch_ok and _shape_disc:
-                # TraceLens-patched SGLang exposes
-                # --enable-shape-discovery-for-cuda-graph-profile; unpatched
-                # SGLang errors on it.
-                existing_sglang = str(envs.get("EXTRA_SGLANG_ARGS", ""))
-                if "shape-discovery-for-cuda-graph-profile" not in existing_sglang:
-                    envs["EXTRA_SGLANG_ARGS"] = (
-                        f"{existing_sglang} --enable-shape-discovery-for-cuda-graph-profile"
-                    ).strip()
+                _tool_dir = kernel_shape_tool_dir()
+                if _shape_disc and _tool_dir is not None:
+                    _existing_pp = str(envs.get("PYTHONPATH", "")).strip()
+                    envs["PYTHONPATH"] = f"{_tool_dir}{os.pathsep}{_existing_pp}" if _existing_pp else str(_tool_dir)
+                    envs["TRACELENS_SHAPE_DISCOVERY"] = "1"
+                else:
+                    envs["TRACELENS_SHAPE_DISCOVERY"] = "0"
+                    if _shape_disc and _tool_dir is None:
+                        log.warning(
+                            "SGLang shape mode=sitecustomize but kernel_shape_tool "
+                            "not found under TRACELENS_ROOT; shapes will be absent "
+                            "(set TRACELENS_ROOT to an NFS path visible to the server).",
+                        )
+                envs["PROFILE_EXTRA_BODY"] = _json.dumps(extra_body)
+            else:
+                # Legacy patched path: capture options need the git-apply patch to
+                # land. Keyed on the degraded reason, not tracelens_patch_ok, since a
+                # patch may be baked into the image without being attempted here.
+                _patch_degraded = envs.get("HYPERLOOM_PROFILE_DEGRADED_REASON") == _TRACELENS_PATCH_UNAVAILABLE
+                if _patch_degraded:
+                    _shape_disc = False
+                extra_body["shape_discovery"] = _shape_disc
+                if _patch_degraded:
+                    extra_body["detailed_annotations"] = False
+                else:
+                    extra_body.setdefault("detailed_annotations", True)
+                # Written before the per-task extra_envs merge, so an extra_envs
+                # PROFILE_EXTRA_BODY can still drop start_step/num_steps. Not
+                # re-asserted for SGLang (unlike vLLM): "unbounded" vs "no capture"
+                # for non-positive num_steps needs a SGLang-side answer.
+                envs["PROFILE_EXTRA_BODY"] = _json.dumps(extra_body)
+                if tracelens_patch_ok and _shape_disc:
+                    # Patched SGLang exposes this arg; unpatched errors on it.
+                    existing_sglang = str(envs.get("EXTRA_SGLANG_ARGS", ""))
+                    if "shape-discovery-for-cuda-graph-profile" not in existing_sglang:
+                        envs["EXTRA_SGLANG_ARGS"] = (
+                            f"{existing_sglang} --enable-shape-discovery-for-cuda-graph-profile"
+                        ).strip()
 
     if not _is_scriptable_profile:
         # NUM_PROMPTS / NUM_WARMUPS are serving-request concepts; xDiT drives its

src/hyperloom/orchestrator/actions/executors/profile.py

@@ -660,20 +660,25 @@ def _note_check(
             skip_reason="main trace could not be sampled",
         )
     else:
+        # Shape markers left by either mechanism: the sglang_profiler:: op
+        # namespace, or the kernel_shape_profiler frame (when with_stack is on).
+        _shape_markers = ("sglang_profiler::", "kernel_shape_profiler")
+        _shape_present = any(m in main_text for m in _shape_markers)
         _note_check(
             CHECK_SGLANG_SHAPE_PROFILER,
-            status="passed" if "kernel_shape_profiler" in main_text else "failed",
+            status="passed" if _shape_present else "failed",
             sampled_file=main_traces[0].name,
             sampled_bytes=_TRACE_INSPECT_BYTES,
         )
-        if "kernel_shape_profiler" not in main_text:
+        if not _shape_present:
             issues.append(
                 f"[5] sglang main trace ({main_traces[0].name}, sampled "
                 f"first {_TRACE_INSPECT_BYTES // 1_000_000} MB) lacks "
-                "kernel_shape_profiler events — shape-discovery "
-                "patch didn't reach the live SGLang. Verify "
-                "_server_patcher (PR #207) succeeded for the "
-                "deployed SGLang version (check log warnings)."
+                "kernel-shape events — shape discovery didn't reach the live "
+                "SGLang. For SGLang < 0.5.18 verify the _server_patcher "
+                "git-apply succeeded; for >= 0.5.18 verify the kernel_shape_tool "
+                "is on the server PYTHONPATH and TRACELENS_SHAPE_DISCOVERY=1 "
+                "(check log warnings)."
             )
 
     if issues: