Add a SimplePhaseChecker for execution-phase assertions (#26803)

python/sglang/srt/environ.py

@@ -238,6 +238,7 @@ class Envs:
     SGLANG_FORCE_SHUTDOWN = EnvBool(False)
     SGLANG_DEBUG_MEMORY_POOL = EnvBool(False)
     SGLANG_DEBUG_REVERT_PR = EnvInt(0)
+    SGLANG_PHASE_CHECKER_DEBUG = EnvBool(False)
     SGLANG_TEST_REQUEST_TIME_STATS = EnvBool(False)
     SGLANG_DISABLE_TP_MEMORY_INBALANCE_CHECK = EnvBool(False)
     SGLANG_SIMULATE_ACC_LEN = EnvFloat(-1)

python/sglang/srt/utils/phase_checker.py

@@ -0,0 +1,109 @@
+from __future__ import annotations
+
+from enum import IntEnum
+
+import torch
+import triton
+import triton.language as tl
+
+from sglang.srt.environ import envs
+
+
+def _phase_repr(phase: int | IntEnum) -> str:
+    if isinstance(phase, IntEnum):
+        return f"{phase.name}({int(phase)})"
+    return str(int(phase))
+
+
+def _host_debug(msg: str) -> None:
+    if envs.SGLANG_PHASE_CHECKER_DEBUG.get():
+        print(msg, flush=True)
+
+
+# debug=True so tl.device_assert below actually raises. Without it the assert
+# is stripped at compile time and only tl.device_print fires (the assert is
+# gated on the TRITON_DEBUG env var by default — see tl.device_assert docstring).
+@triton.jit(debug=True)
+def _phase_check_kernel(
+    phase_ptr,
+    enable_assert_ptr,
+    EXPECT_PHASE: tl.constexpr,
+    NEXT_PHASE: tl.constexpr,
+    CALLER_TAG: tl.constexpr,
+):
+    cur = tl.load(phase_ptr)
+    enable_assert = tl.load(enable_assert_ptr)
+    if enable_assert != 0:
+        if cur != EXPECT_PHASE:
+            # constexpr values get baked into the prefix string at compile time;
+            # only `cur` is runtime.
+            tl.device_print(
+                f"[SimplePhaseChecker FAIL] caller_tag={CALLER_TAG} "
+                f"expect={EXPECT_PHASE} next={NEXT_PHASE} actual=",
+                cur,
+            )
+        tl.device_assert(cur == EXPECT_PHASE, "SimplePhaseChecker: phase mismatch")
+    tl.store(phase_ptr, NEXT_PHASE)
+
+
+class SimplePhaseChecker:
+    """GPU-side state machine for any int-keyed phase sequence."""
+
+    def __init__(self, *, initial_phase: int | IntEnum, device: torch.device) -> None:
+        self._initial_phase = int(initial_phase)
+        self._phase = torch.tensor(
+            self._initial_phase, dtype=torch.int32, device=device
+        )
+        self._enable_assert_device = torch.zeros(1, dtype=torch.int32, device=device)
+        self._caller_tag_registry: dict[str, int] = {}
+        _host_debug(
+            f"[SimplePhaseChecker.__init__] device={device} "
+            f"initial_phase={_phase_repr(initial_phase)} "
+            f"enable_assert=OFF (call enable_assert() after init is done)"
+        )
+
+    def enable_assert(self) -> None:
+        """Reset phase to initial_phase, then enable the device-side assert."""
+        self._reset_to_idle()
+        self._enable_assert_device.fill_(1)
+        _host_debug(f"[SimplePhaseChecker.enable_assert] assert ENABLED")
+
+    def update(
+        self,
+        *,
+        expect_phase: int | IntEnum,
+        next_phase: int | IntEnum,
+        caller_name: str = "",
+    ) -> None:
+        caller_tag = self._resolve_caller_tag(caller_name)
+        _host_debug(
+            f"[SimplePhaseChecker.update] caller={caller_name!r} "
+            f"caller_tag={caller_tag} "
+            f"expect={_phase_repr(expect_phase)} "
+            f"next={_phase_repr(next_phase)} "
+            f"capturing={torch.cuda.is_current_stream_capturing()}"
+        )
+        _phase_check_kernel[(1,)](
+            self._phase,
+            self._enable_assert_device,
+            EXPECT_PHASE=int(expect_phase),
+            NEXT_PHASE=int(next_phase),
+            CALLER_TAG=caller_tag,
+        )
+
+    def _reset_to_idle(self) -> None:
+        self._phase.fill_(self._initial_phase)
+        _host_debug(
+            f"[SimplePhaseChecker._reset_to_idle] phase reset to "
+            f"{self._initial_phase}"
+        )
+
+    def _resolve_caller_tag(self, caller_name: str) -> int:
+        registry = self._caller_tag_registry
+        if caller_name not in registry:
+            registry[caller_name] = len(registry) + 1
+            _host_debug(
+                f"[SimplePhaseChecker] registered caller_tag "
+                f"{registry[caller_name]} <- {caller_name!r}"
+            )
+        return registry[caller_name]