# glm53-compile-02-trace-safe-probes.patch — libr4d capability probes must not do I/O
# inside a torch.compile region.
#
# STATUS: WIP — second of three patches in the "compile the GLM trunk" changelist.
# Inert (but harmless) without 01. Not part of the deployed shape.
#
# THE BREAK (measured: with 01 mounted and R4D_GLM_COMPILE=1, the worker dies during
# profile_run, i.e. the first — cold — forward):
#
#   Explanation: Dynamo does not know how to trace the builtin `posix.stat`
#   ... r4d_mhc_supported -> r4d_lib.r4d_mhc_enabled -> available -> os.path.exists
#
# THE NON-OBVIOUS PART: functools.lru_cache does NOT protect these probes. Dynamo
# re-traces the wrapped body to build its guard even when the cache is warm — it warns
# about exactly this ("Dynamo detected a call to a `functools.lru_cache`-wrapped
# function ... Dynamo ignores the cache wrapper and directly traces the wrapped
# function"). Verified in-image, torch 2.11.0+rocm10.0, torch.compile(backend="eager"):
#
#   cold lru_cache call       -> RAISE Unsupported
#   WARM lru_cache call       -> RAISE Unsupported     <-- warming is not enough
#   plain module-global read  -> compiles
#   uncached os.path.exists   -> RAISE Unsupported
#
# So each probe must resolve ONCE into a plain value and be read as a global. The set
# below is every no-argument probe / symbol lookup reachable from a forward pass,
# grep-verified call sites:
#   model_executor/layers/fused_moe/experts/r4d_w4a16_moe.py:79  available()
#   model_executor/layers/fused_moe/experts/r4d_w4a16_moe.py:81  has_w4a16_moe()
#   model_executor/kernels/mhc/r4d_dsv4.py:16                    r4d_mhc_enabled()
#   v1/attention/ops/rocm_aiter_mla_sparse.py:694                r4d_indexer_enabled(), mqa_max_heads()
#   v1/attention/ops/rocm_aiter_mla_sparse.py:1124               r4d_indexer_enabled()
#   model_executor/layers/fused_moe/experts/r4d_mxfp4_moe.py:683 r4d_act_quant_enabled()
#   model_executor/layers/fused_moe/experts/r4d_w4a16_moe.py _entry()  moe_w4a16()/moe_w4a16_mt1() per MoE call
# (lib() is the transitive dependency of all of them.)
#
# TARGET: vllm/model_executor/kernels/r4d_lib.py   (inside the pinned image)
#   base    sha256 b8dfeefb4acc72b67b0579a4ab0653ef5d48c40cc8649187d757c80ce6b9617a
#   patched sha256 0751852e1ca813ad2bdb1da7bf10c605130e63e04ef9ec57e70b5f675a29c08e
#
# APPLY: append the block below to the module. R4D_FREEZE_PROBES=0 restores stock
# behaviour. Cost when nothing is compiled: ten one-time probes at import (measured:
# logged once per process, e.g. "libr4d capability probes frozen: 10 warm [...]").
#
# OBSERVED EFFECT: with 01+02 the worker gets past this break and reaches the next one
# (see 03) — dynamo reports "Dynamo bytecode transform time: 1.08 s" for the trunk.
--- a/vllm/model_executor/kernels/r4d_lib.py
+++ b/vllm/model_executor/kernels/r4d_lib.py
@@ -361,3 +361,71 @@
 @functools.lru_cache
 def moe_w4a16_mt1():
     return getattr(lib(), "r4d_gemm_moe_w4a16_nt_b16_m64_mt1", None)
+
+
+# FORK-LOCAL (patch: trace-safe capability probes).  torch.compile traces the
+# BODY of a functools.lru_cache call even when the cache is warm, so a capability
+# probe called from a forward pass is a dynamo graph break.  Measured 2026-09-30
+# on the GLM-5.3-Flash trunk under @support_torch_compile:
+#   Dynamo does not know how to trace the builtin `posix.stat`
+#   r4d_mhc_supported -> r4d_mhc_enabled -> available -> os.path.exists
+# Each probe below is wrapped to resolve ONCE into a memo and is warmed at import
+# (outside any tracing region), so a traced call is a dict read.  The probe set is
+# every no-argument capability probe / symbol lookup reachable from a forward
+# pass: the W4A16 + mxfp4 MoE dispatch, the mHC op, the MLA sparse indexer gate
+# and the draft/lm-head guards.
+_PROBE_NAMES = (
+    "available",
+    "lib",
+    "has_w4a16_moe",
+    "moe_w4a16",
+    "moe_w4a16_mt1",
+    "has_mhc_dsv4",
+    "r4d_mhc_enabled",
+    "r4d_indexer_enabled",
+    "mqa_max_heads",
+    "r4d_act_quant_enabled",
+)
+
+
+def _freeze_probe(fn):
+    """Return fn bound to its first result.  Result must be a plain value."""
+
+    memo = []
+
+    def probe(*args, **kwargs):
+        if not memo:
+            memo.append(fn(*args, **kwargs))
+        return memo[0]
+
+    return probe
+
+
+def _freeze_probes() -> list[str]:
+    frozen = []
+    for name in _PROBE_NAMES:
+        fn = globals().get(name)
+        if fn is None:
+            continue
+        # Only the fast branch of a bound probe is traced; an empty memo means the
+        # first call still does I/O, so populate it here as well.
+        globals()[name] = _freeze_probe(fn)
+        frozen.append(name)
+    return frozen
+
+
+from vllm.logger import init_logger as _init_logger
+
+_probe_log = _init_logger(__name__)
+
+if os.environ.get("R4D_FREEZE_PROBES", "1") == "1":
+    _warmed, _failed = [], []
+    for _n in _freeze_probes():
+        try:
+            globals()[_n]()
+            _warmed.append(_n)
+        except Exception:  # noqa: BLE001 -- probe unavailable here; recompute lazily
+            _failed.append(_n)
+    _probe_log.info("libr4d capability probes frozen: %d warm %s%s",
+                    len(_warmed), _warmed,
+                    f" (deferred: {_failed})" if _failed else "")