Add a SimplePhaseChecker for execution-phase assertions (#26803)
This commit is contained in:
fzyzcjy
@@ -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)
|
||||
|
||||
@@ -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]
|
||||
@@ -0,0 +1,483 @@
|
||||
"""Unit tests for srt/utils/phase_checker.py — SimplePhaseChecker."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import subprocess
|
||||
import sys
|
||||
import textwrap
|
||||
import unittest
|
||||
from enum import IntEnum
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.utils.phase_checker import SimplePhaseChecker
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
register_cuda_ci(est_time=120, stage="base-b", runner_config="1-gpu-small")
|
||||
|
||||
|
||||
class _Phase(IntEnum):
|
||||
IDLE = 0
|
||||
A = 1
|
||||
B = 2
|
||||
C = 3
|
||||
|
||||
|
||||
def _phase_value(checker: SimplePhaseChecker) -> int:
|
||||
"""Read the device-side phase tensor as a host int."""
|
||||
return int(checker._phase.item())
|
||||
|
||||
|
||||
def _assert_flag(checker: SimplePhaseChecker) -> int:
|
||||
"""Read the device-side assert-enable flag as a host int."""
|
||||
return int(checker._enable_assert_device.item())
|
||||
|
||||
|
||||
class TestConstruction(CustomTestCase):
|
||||
"""Constructor sets phase and leaves assert OFF so init-time work is tolerated."""
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls) -> None:
|
||||
cls.device = torch.device("cuda:0")
|
||||
|
||||
def test_init_stores_initial_phase_int(self) -> None:
|
||||
checker = SimplePhaseChecker(initial_phase=7, device=self.device)
|
||||
self.assertEqual(_phase_value(checker), 7)
|
||||
|
||||
def test_init_stores_initial_phase_intenum(self) -> None:
|
||||
checker = SimplePhaseChecker(initial_phase=_Phase.B, device=self.device)
|
||||
self.assertEqual(_phase_value(checker), int(_Phase.B))
|
||||
|
||||
def test_init_assert_flag_is_off(self) -> None:
|
||||
"""assert flag must default to OFF so warmup / cuda-graph-capture don't trip."""
|
||||
checker = SimplePhaseChecker(initial_phase=_Phase.IDLE, device=self.device)
|
||||
self.assertEqual(_assert_flag(checker), 0)
|
||||
|
||||
def test_init_caller_registry_is_empty(self) -> None:
|
||||
checker = SimplePhaseChecker(initial_phase=_Phase.IDLE, device=self.device)
|
||||
self.assertEqual(checker._caller_tag_registry, {})
|
||||
|
||||
def test_init_phase_tensor_is_on_requested_device(self) -> None:
|
||||
checker = SimplePhaseChecker(initial_phase=0, device=self.device)
|
||||
self.assertEqual(checker._phase.device, self.device)
|
||||
self.assertEqual(checker._enable_assert_device.device, self.device)
|
||||
self.assertEqual(checker._phase.dtype, torch.int32)
|
||||
self.assertEqual(checker._enable_assert_device.dtype, torch.int32)
|
||||
|
||||
|
||||
class TestUpdateAssertDisabled(CustomTestCase):
|
||||
"""With the assert flag OFF, update always advances the phase and never raises."""
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls) -> None:
|
||||
cls.device = torch.device("cuda:0")
|
||||
|
||||
def test_update_advances_phase_on_match(self) -> None:
|
||||
checker = SimplePhaseChecker(initial_phase=_Phase.IDLE, device=self.device)
|
||||
checker.update(expect_phase=_Phase.IDLE, next_phase=_Phase.A, caller_name="t")
|
||||
torch.cuda.synchronize()
|
||||
self.assertEqual(_phase_value(checker), int(_Phase.A))
|
||||
|
||||
def test_update_advances_phase_on_mismatch(self) -> None:
|
||||
"""assert OFF tolerates mismatches — store still happens unconditionally."""
|
||||
checker = SimplePhaseChecker(initial_phase=_Phase.IDLE, device=self.device)
|
||||
checker.update(expect_phase=_Phase.C, next_phase=_Phase.B, caller_name="t")
|
||||
torch.cuda.synchronize()
|
||||
self.assertEqual(_phase_value(checker), int(_Phase.B))
|
||||
|
||||
def test_init_time_lifecycle_violations_tolerated(self) -> None:
|
||||
"""Documented use case: pre-enable_assert work may freely violate the lifecycle."""
|
||||
checker = SimplePhaseChecker(initial_phase=_Phase.IDLE, device=self.device)
|
||||
# Bogus sequence — would all mismatch if assert were on.
|
||||
checker.update(expect_phase=_Phase.B, next_phase=_Phase.C, caller_name="warmup")
|
||||
checker.update(
|
||||
expect_phase=_Phase.IDLE, next_phase=_Phase.A, caller_name="warmup"
|
||||
)
|
||||
checker.update(
|
||||
expect_phase=_Phase.B, next_phase=_Phase.IDLE, caller_name="warmup"
|
||||
)
|
||||
torch.cuda.synchronize() # no raise
|
||||
self.assertEqual(_phase_value(checker), int(_Phase.IDLE))
|
||||
|
||||
|
||||
class TestUpdateAssertEnabled(CustomTestCase):
|
||||
"""With assert ON, matching updates pass and walk a multi-state lifecycle."""
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls) -> None:
|
||||
cls.device = torch.device("cuda:0")
|
||||
|
||||
def test_update_advances_phase_on_match(self) -> None:
|
||||
checker = SimplePhaseChecker(initial_phase=_Phase.IDLE, device=self.device)
|
||||
checker.enable_assert()
|
||||
checker.update(expect_phase=_Phase.IDLE, next_phase=_Phase.A, caller_name="t")
|
||||
torch.cuda.synchronize()
|
||||
self.assertEqual(_phase_value(checker), int(_Phase.A))
|
||||
|
||||
def test_full_4_state_lifecycle_round_trip(self) -> None:
|
||||
"""IDLE -> A -> B -> C -> IDLE under assert ON, no synchronization mid-cycle."""
|
||||
checker = SimplePhaseChecker(initial_phase=_Phase.IDLE, device=self.device)
|
||||
checker.enable_assert()
|
||||
for _ in range(3): # multiple cycles to exercise stable behavior
|
||||
checker.update(
|
||||
expect_phase=_Phase.IDLE, next_phase=_Phase.A, caller_name="p1"
|
||||
)
|
||||
checker.update(expect_phase=_Phase.A, next_phase=_Phase.B, caller_name="p2")
|
||||
checker.update(expect_phase=_Phase.B, next_phase=_Phase.C, caller_name="p3")
|
||||
checker.update(
|
||||
expect_phase=_Phase.C, next_phase=_Phase.IDLE, caller_name="p4"
|
||||
)
|
||||
torch.cuda.synchronize()
|
||||
self.assertEqual(_phase_value(checker), int(_Phase.IDLE))
|
||||
|
||||
def test_update_mismatch_after_enable_raises_in_subprocess(self) -> None:
|
||||
"""A mismatched update with assert ON must fire device_assert at the next sync.
|
||||
|
||||
Run in a subprocess because device-side asserts poison the CUDA context.
|
||||
"""
|
||||
script = textwrap.dedent("""
|
||||
import sys
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.utils.phase_checker import SimplePhaseChecker
|
||||
|
||||
device = torch.device("cuda:0")
|
||||
checker = SimplePhaseChecker(initial_phase=0, device=device)
|
||||
checker.enable_assert()
|
||||
# phase=0 but we claim expect=99 — kernel must fire device_assert.
|
||||
checker.update(expect_phase=99, next_phase=1, caller_name="bad")
|
||||
try:
|
||||
torch.cuda.synchronize()
|
||||
except RuntimeError as e:
|
||||
msg = str(e).lower()
|
||||
if "device-side assert" in msg or "phase mismatch" in msg:
|
||||
sys.exit(0)
|
||||
print(f"Unexpected RuntimeError: {e}", file=sys.stderr)
|
||||
sys.exit(2)
|
||||
print("expected RuntimeError but none was raised", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
""")
|
||||
result = subprocess.run(
|
||||
[sys.executable, "-c", script],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=180,
|
||||
)
|
||||
# Accept both exit-0 (RuntimeError caught) and SIGABRT (-6, the kernel's
|
||||
# device_assert killed the process directly before sync could raise).
|
||||
# Either way, the kernel-side check fired — which is what we're verifying.
|
||||
# The presence of the SimplePhaseChecker FAIL line in stdout confirms it.
|
||||
self.assertIn(
|
||||
"SimplePhaseChecker FAIL",
|
||||
result.stdout,
|
||||
f"stdout: {result.stdout}\nstderr: {result.stderr}",
|
||||
)
|
||||
self.assertIn(
|
||||
result.returncode,
|
||||
(0, -6),
|
||||
f"unexpected returncode {result.returncode}; "
|
||||
f"stdout: {result.stdout}\nstderr: {result.stderr}",
|
||||
)
|
||||
|
||||
|
||||
class TestEnableAssert(CustomTestCase):
|
||||
"""enable_assert flips the device flag and resets phase to the initial value."""
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls) -> None:
|
||||
cls.device = torch.device("cuda:0")
|
||||
|
||||
def test_enable_assert_sets_flag_to_one(self) -> None:
|
||||
checker = SimplePhaseChecker(initial_phase=_Phase.IDLE, device=self.device)
|
||||
self.assertEqual(_assert_flag(checker), 0)
|
||||
checker.enable_assert()
|
||||
self.assertEqual(_assert_flag(checker), 1)
|
||||
|
||||
def test_enable_assert_resets_phase_to_initial(self) -> None:
|
||||
"""Phase advanced during init must be wiped by enable_assert.
|
||||
|
||||
Otherwise post-init lifecycle would start from whatever captured kernels
|
||||
left in the phase tensor — the explicit purpose of enable_assert.
|
||||
"""
|
||||
checker = SimplePhaseChecker(initial_phase=_Phase.IDLE, device=self.device)
|
||||
# Mess with the phase as init-time work might.
|
||||
checker.update(
|
||||
expect_phase=_Phase.IDLE, next_phase=_Phase.C, caller_name="warmup"
|
||||
)
|
||||
torch.cuda.synchronize()
|
||||
self.assertEqual(_phase_value(checker), int(_Phase.C))
|
||||
|
||||
checker.enable_assert()
|
||||
self.assertEqual(_phase_value(checker), int(_Phase.IDLE))
|
||||
|
||||
def test_enable_assert_with_nonzero_initial_phase(self) -> None:
|
||||
"""Reset target tracks the original initial_phase, not 0."""
|
||||
checker = SimplePhaseChecker(initial_phase=42, device=self.device)
|
||||
checker.update(expect_phase=42, next_phase=7, caller_name="t")
|
||||
torch.cuda.synchronize()
|
||||
self.assertEqual(_phase_value(checker), 7)
|
||||
|
||||
checker.enable_assert()
|
||||
self.assertEqual(_phase_value(checker), 42)
|
||||
|
||||
def test_enable_assert_is_idempotent(self) -> None:
|
||||
checker = SimplePhaseChecker(initial_phase=_Phase.IDLE, device=self.device)
|
||||
checker.enable_assert()
|
||||
checker.enable_assert()
|
||||
self.assertEqual(_assert_flag(checker), 1)
|
||||
self.assertEqual(_phase_value(checker), int(_Phase.IDLE))
|
||||
|
||||
|
||||
class TestResetToIdle(CustomTestCase):
|
||||
"""_reset_to_idle only touches the phase tensor; the assert flag is untouched."""
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls) -> None:
|
||||
cls.device = torch.device("cuda:0")
|
||||
|
||||
def test_reset_after_update_restores_initial_phase(self) -> None:
|
||||
checker = SimplePhaseChecker(initial_phase=_Phase.IDLE, device=self.device)
|
||||
checker.update(expect_phase=_Phase.IDLE, next_phase=_Phase.B, caller_name="t")
|
||||
torch.cuda.synchronize()
|
||||
self.assertEqual(_phase_value(checker), int(_Phase.B))
|
||||
|
||||
checker._reset_to_idle()
|
||||
self.assertEqual(_phase_value(checker), int(_Phase.IDLE))
|
||||
|
||||
def test_reset_does_not_toggle_assert_flag_when_off(self) -> None:
|
||||
checker = SimplePhaseChecker(initial_phase=_Phase.IDLE, device=self.device)
|
||||
self.assertEqual(_assert_flag(checker), 0)
|
||||
checker._reset_to_idle()
|
||||
self.assertEqual(_assert_flag(checker), 0)
|
||||
|
||||
def test_reset_does_not_toggle_assert_flag_when_on(self) -> None:
|
||||
checker = SimplePhaseChecker(initial_phase=_Phase.IDLE, device=self.device)
|
||||
checker.enable_assert()
|
||||
self.assertEqual(_assert_flag(checker), 1)
|
||||
checker._reset_to_idle()
|
||||
self.assertEqual(_assert_flag(checker), 1)
|
||||
|
||||
def test_reset_with_nonzero_initial_phase(self) -> None:
|
||||
checker = SimplePhaseChecker(initial_phase=5, device=self.device)
|
||||
checker.update(expect_phase=5, next_phase=9, caller_name="t")
|
||||
torch.cuda.synchronize()
|
||||
checker._reset_to_idle()
|
||||
self.assertEqual(_phase_value(checker), 5)
|
||||
|
||||
|
||||
class TestCallerTagRegistry(CustomTestCase):
|
||||
"""Caller tags are interned 1-indexed integers, assigned in first-seen order."""
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls) -> None:
|
||||
cls.device = torch.device("cuda:0")
|
||||
|
||||
def test_first_caller_gets_tag_one(self) -> None:
|
||||
checker = SimplePhaseChecker(initial_phase=_Phase.IDLE, device=self.device)
|
||||
self.assertEqual(checker._resolve_caller_tag("first"), 1)
|
||||
|
||||
def test_distinct_callers_get_distinct_increasing_tags(self) -> None:
|
||||
checker = SimplePhaseChecker(initial_phase=_Phase.IDLE, device=self.device)
|
||||
self.assertEqual(checker._resolve_caller_tag("a"), 1)
|
||||
self.assertEqual(checker._resolve_caller_tag("b"), 2)
|
||||
self.assertEqual(checker._resolve_caller_tag("c"), 3)
|
||||
|
||||
def test_repeated_caller_returns_same_tag(self) -> None:
|
||||
checker = SimplePhaseChecker(initial_phase=_Phase.IDLE, device=self.device)
|
||||
first = checker._resolve_caller_tag("repeat")
|
||||
# Mix other callers in between.
|
||||
checker._resolve_caller_tag("other")
|
||||
second = checker._resolve_caller_tag("repeat")
|
||||
self.assertEqual(first, second)
|
||||
|
||||
def test_empty_caller_name_is_interned_like_any_other(self) -> None:
|
||||
"""default caller_name="" must be a legal key, not crash on use."""
|
||||
checker = SimplePhaseChecker(initial_phase=_Phase.IDLE, device=self.device)
|
||||
checker.enable_assert()
|
||||
checker.update(expect_phase=_Phase.IDLE, next_phase=_Phase.A) # caller_name=""
|
||||
torch.cuda.synchronize()
|
||||
self.assertIn("", checker._caller_tag_registry)
|
||||
self.assertEqual(_phase_value(checker), int(_Phase.A))
|
||||
|
||||
def test_update_populates_registry(self) -> None:
|
||||
checker = SimplePhaseChecker(initial_phase=_Phase.IDLE, device=self.device)
|
||||
checker.update(
|
||||
expect_phase=_Phase.IDLE, next_phase=_Phase.A, caller_name="alpha"
|
||||
)
|
||||
checker.update(
|
||||
expect_phase=_Phase.A, next_phase=_Phase.IDLE, caller_name="beta"
|
||||
)
|
||||
torch.cuda.synchronize()
|
||||
self.assertEqual(checker._caller_tag_registry, {"alpha": 1, "beta": 2})
|
||||
|
||||
|
||||
class TestMultipleInstances(CustomTestCase):
|
||||
"""Independent checkers must not share phase, assert flag, or caller registry."""
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls) -> None:
|
||||
cls.device = torch.device("cuda:0")
|
||||
|
||||
def test_phase_tensors_are_independent(self) -> None:
|
||||
a = SimplePhaseChecker(initial_phase=_Phase.IDLE, device=self.device)
|
||||
b = SimplePhaseChecker(initial_phase=_Phase.IDLE, device=self.device)
|
||||
a.update(expect_phase=_Phase.IDLE, next_phase=_Phase.B, caller_name="a")
|
||||
torch.cuda.synchronize()
|
||||
self.assertEqual(_phase_value(a), int(_Phase.B))
|
||||
self.assertEqual(_phase_value(b), int(_Phase.IDLE))
|
||||
|
||||
def test_assert_flags_are_independent(self) -> None:
|
||||
a = SimplePhaseChecker(initial_phase=_Phase.IDLE, device=self.device)
|
||||
b = SimplePhaseChecker(initial_phase=_Phase.IDLE, device=self.device)
|
||||
a.enable_assert()
|
||||
self.assertEqual(_assert_flag(a), 1)
|
||||
self.assertEqual(_assert_flag(b), 0)
|
||||
|
||||
def test_caller_registries_are_independent(self) -> None:
|
||||
a = SimplePhaseChecker(initial_phase=_Phase.IDLE, device=self.device)
|
||||
b = SimplePhaseChecker(initial_phase=_Phase.IDLE, device=self.device)
|
||||
a._resolve_caller_tag("shared_name")
|
||||
# Same string, different checker — should still get tag=1, not 2.
|
||||
self.assertEqual(b._resolve_caller_tag("shared_name"), 1)
|
||||
|
||||
|
||||
class TestCudaGraphCapture(CustomTestCase):
|
||||
"""The kernel is launched unconditionally so it is capture-safe; the device flag
|
||||
decides at replay time whether the assert fires.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls) -> None:
|
||||
cls.device = torch.device("cuda:0")
|
||||
|
||||
def _capture_one_update(
|
||||
self,
|
||||
checker: SimplePhaseChecker,
|
||||
*,
|
||||
expect_phase: int,
|
||||
next_phase: int,
|
||||
caller_name: str,
|
||||
) -> torch.cuda.CUDAGraph:
|
||||
"""Warm up + capture a single update() call into a fresh CUDAGraph.
|
||||
|
||||
Both warmup and capture launches actually execute on the GPU and advance
|
||||
the phase tensor. Caller must keep assert OFF here — otherwise the second
|
||||
launch (capture) would see phase=next_phase and mismatch the expected.
|
||||
"""
|
||||
assert _assert_flag(checker) == 0, "capture helper requires assert OFF"
|
||||
stream = torch.cuda.Stream(self.device)
|
||||
# Warmup launch on a side stream — required before torch.cuda.graph.
|
||||
stream.wait_stream(torch.cuda.current_stream(self.device))
|
||||
with torch.cuda.stream(stream):
|
||||
checker.update(
|
||||
expect_phase=expect_phase,
|
||||
next_phase=next_phase,
|
||||
caller_name=caller_name,
|
||||
)
|
||||
torch.cuda.current_stream(self.device).wait_stream(stream)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
graph = torch.cuda.CUDAGraph()
|
||||
with torch.cuda.graph(graph, stream=stream):
|
||||
checker.update(
|
||||
expect_phase=expect_phase,
|
||||
next_phase=next_phase,
|
||||
caller_name=caller_name,
|
||||
)
|
||||
return graph
|
||||
|
||||
def test_captured_update_advances_phase_on_replay(self) -> None:
|
||||
"""Replay reissues the captured kernel; phase tensor is updated each replay."""
|
||||
checker = SimplePhaseChecker(initial_phase=_Phase.IDLE, device=self.device)
|
||||
# Capture under assert OFF (the documented init-time pattern).
|
||||
graph = self._capture_one_update(
|
||||
checker,
|
||||
expect_phase=int(_Phase.IDLE),
|
||||
next_phase=int(_Phase.B),
|
||||
caller_name="captured",
|
||||
)
|
||||
|
||||
# Enable assert (resets phase -> IDLE) and replay — captured expect=IDLE matches.
|
||||
checker.enable_assert()
|
||||
graph.replay()
|
||||
torch.cuda.synchronize()
|
||||
self.assertEqual(_phase_value(checker), int(_Phase.B))
|
||||
|
||||
# Reset + replay again — same result, no raise.
|
||||
checker._reset_to_idle()
|
||||
graph.replay()
|
||||
torch.cuda.synchronize()
|
||||
self.assertEqual(_phase_value(checker), int(_Phase.B))
|
||||
|
||||
def test_assert_flag_toggle_visible_to_replayed_graph(self) -> None:
|
||||
"""Key contract: capture with assert OFF, enable assert later, replay sees the new flag.
|
||||
|
||||
We don't try to make the replay raise (that would poison the context); we
|
||||
verify the *positive* side — capture under assert OFF (with deliberate
|
||||
phase mismatches tolerated) and confirm the same captured kernel still
|
||||
works correctly when assert is later turned on.
|
||||
"""
|
||||
checker = SimplePhaseChecker(initial_phase=_Phase.IDLE, device=self.device)
|
||||
graph = self._capture_one_update(
|
||||
checker,
|
||||
expect_phase=int(_Phase.IDLE),
|
||||
next_phase=int(_Phase.A),
|
||||
caller_name="captured",
|
||||
)
|
||||
self.assertEqual(_assert_flag(checker), 0)
|
||||
|
||||
# Replay with assert OFF tolerates a deliberately diverged phase.
|
||||
checker._phase.fill_(999)
|
||||
graph.replay()
|
||||
torch.cuda.synchronize() # no raise — flag is OFF
|
||||
self.assertEqual(_phase_value(checker), int(_Phase.A))
|
||||
|
||||
# Now turn on asserts (also resets phase -> IDLE) and replay.
|
||||
checker.enable_assert()
|
||||
self.assertEqual(_assert_flag(checker), 1)
|
||||
self.assertEqual(_phase_value(checker), int(_Phase.IDLE))
|
||||
|
||||
graph.replay()
|
||||
torch.cuda.synchronize() # no raise — phase matched expect
|
||||
self.assertEqual(_phase_value(checker), int(_Phase.A))
|
||||
|
||||
|
||||
class TestPhaseReprNoCrash(CustomTestCase):
|
||||
"""Debug logging path (host_debug) tolerates both int and IntEnum phase args."""
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls) -> None:
|
||||
cls.device = torch.device("cuda:0")
|
||||
|
||||
def test_update_with_int_phases_does_not_crash(self) -> None:
|
||||
checker = SimplePhaseChecker(initial_phase=0, device=self.device)
|
||||
checker.enable_assert()
|
||||
checker.update(expect_phase=0, next_phase=1, caller_name="ints")
|
||||
torch.cuda.synchronize()
|
||||
self.assertEqual(_phase_value(checker), 1)
|
||||
|
||||
def test_update_with_intenum_phases_does_not_crash(self) -> None:
|
||||
checker = SimplePhaseChecker(initial_phase=_Phase.IDLE, device=self.device)
|
||||
checker.enable_assert()
|
||||
checker.update(
|
||||
expect_phase=_Phase.IDLE, next_phase=_Phase.A, caller_name="enums"
|
||||
)
|
||||
torch.cuda.synchronize()
|
||||
self.assertEqual(_phase_value(checker), int(_Phase.A))
|
||||
|
||||
def test_update_mixing_int_and_intenum_phases(self) -> None:
|
||||
"""expect_phase is an IntEnum, next_phase is a plain int (and vice versa)."""
|
||||
checker = SimplePhaseChecker(initial_phase=_Phase.IDLE, device=self.device)
|
||||
checker.enable_assert()
|
||||
checker.update(expect_phase=_Phase.IDLE, next_phase=5, caller_name="mix1")
|
||||
torch.cuda.synchronize()
|
||||
self.assertEqual(_phase_value(checker), 5)
|
||||
checker.update(expect_phase=5, next_phase=_Phase.IDLE, caller_name="mix2")
|
||||
torch.cuda.synchronize()
|
||||
self.assertEqual(_phase_value(checker), int(_Phase.IDLE))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user