diff --git a/python/sglang/_mps_stub.py b/python/sglang/_mps_stub.py index aaaebe27e..8a090b398 100644 --- a/python/sglang/_mps_stub.py +++ b/python/sglang/_mps_stub.py @@ -69,6 +69,19 @@ class Event: return 0.0 +class StreamContext: + """Minimal stand-in for ``torch.cuda.StreamContext``.""" + + def __init__(self, stream: Any = None) -> None: + pass + + def __enter__(self) -> "StreamContext": + return self + + def __exit__(self, *args: Any) -> None: + pass + + _default_stream = Stream() @@ -235,6 +248,7 @@ def install() -> None: # Only patch attributes that are actually missing for name, obj in [ ("Stream", Stream), + ("StreamContext", StreamContext), ("Event", Event), ("current_stream", current_stream), ("stream", stream),