From ebae8ee21ef4c4fe452140a68b231b3eb693d49c Mon Sep 17 00:00:00 2001
From: Alex Nails <alex.nails@radixark.ai>
Date: Thu, 3 Sep 2026 22:56:14 -0700
Subject: [PATCH] [Fix] Register triton.runtime.cache.triton_key in the MPS
 stub so torch.compile keeps working (#37937)

Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
---
 python/sglang/_platform_stubs.py              | 10 +++++++++-
 .../unit/platforms/test_mps_triton_stub.py    | 19 +++++++++++++++++++
 2 files changed, 28 insertions(+), 1 deletion(-)

diff --git a/python/sglang/_platform_stubs.py b/python/sglang/_platform_stubs.py
index 8e20db7a1..b62ba7395 100644
--- a/python/sglang/_platform_stubs.py
+++ b/python/sglang/_platform_stubs.py
@@ -105,6 +105,10 @@ def _next_power_of_2(n: int) -> int:
     return 1 << (n - 1).bit_length() if n > 0 else 1
 
 
+def _triton_key() -> str:
+    return "triton-stub"
+
+
 class _Config:
     """Minimal stand-in for ``triton.Config`` used in ``@triton.autotune``."""
 
@@ -438,6 +442,10 @@ def install_platform_stubs() -> None:
         jit_mod.KernelInterface = _KernelInterface
         runtime.jit = jit_mod
 
+        cache_mod = _make_mock("triton.runtime.cache")
+        cache_mod.triton_key = _triton_key
+        runtime.cache = cache_mod
+
         # Torch 2.13 imports these as classes while initializing Inductor, even on
         # MPS where no Triton kernel is compiled. Define them explicitly so the
         # catch-all meta-path finder does not materialize class names as modules.
@@ -465,7 +473,7 @@ def install_platform_stubs() -> None:
             pass
 
         compiler_impl.ASTSource = _ASTSource
-        compiler_impl.triton_key = lambda: "triton-stub"
+        compiler_impl.triton_key = _triton_key
         compiler_root.compiler = compiler_impl
         triton.compiler = compiler_root
 
diff --git a/test/registered/unit/platforms/test_mps_triton_stub.py b/test/registered/unit/platforms/test_mps_triton_stub.py
index 1fc8733bd..808c10cc1 100644
--- a/test/registered/unit/platforms/test_mps_triton_stub.py
+++ b/test/registered/unit/platforms/test_mps_triton_stub.py
@@ -40,6 +40,25 @@ assert _KernelType is not None
             msg=f"stdout={completed.stdout}\nstderr={completed.stderr}",
         )
 
+    def test_torch_compile_works_after_sglang_installs_stub(self):
+        script = """
+import sglang
+import torch
+torch.compile(lambda x: x * 2 + 1)(torch.randn(8))
+"""
+        completed = subprocess.run(
+            [sys.executable, "-c", script],
+            capture_output=True,
+            text=True,
+            timeout=60,
+            check=False,
+        )
+        self.assertEqual(
+            completed.returncode,
+            0,
+            msg=f"stdout={completed.stdout}\nstderr={completed.stderr}",
+        )
+
 
 if __name__ == "__main__":
     unittest.main()