EngineX-Hygon/sglang
5
0
Fork 0
Code Issues Pull Requests Actions 7 Projects Releases Wiki Activity

Optional extension for green context (#9231)

Browse Source
This commit is contained in:
Liangsheng Yin
2025-08-15 21:33:52 +08:00
committed by GitHub
parent c186feed7f
commit 0c8594e67d
6 changed files with 73 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files

20
sgl-kernel/python/sgl_kernel/spatial.py
View File

@@ -1,6 +1,17 @@
import torch
from torch.cuda.streams import ExternalStream
try:
from . import spatial_ops # triggers TORCH extension registration
except Exception as _e:
_spatial_import_error = _e
else:
_spatial_import_error = None
_IMPORT_ERROR = ImportError(
"Failed to load sgl_kernel.spatial_ops extension. Ensure CUDA Driver >= 12.4"
)
def create_greenctx_stream_by_value(
SM_a: int, SM_b: int, device_id: int = None
@@ -14,11 +25,8 @@ def create_greenctx_stream_by_value(
Returns:
tuple[ExternalStream, ExternalStream]: The two streams.
"""
if torch.version.cuda < "12.4":
raise RuntimeError(
"Green Contexts feature requires CUDA Toolkit 12.4 or newer."
)
if _spatial_import_error is not None:
raise _IMPORT_ERROR from _spatial_import_error
if device_id is None:
device_id = torch.cuda.current_device()
@@ -42,6 +50,8 @@ def get_sm_available(device_id: int = None) -> int:
Returns:
int: The SMs available.
"""
if _spatial_import_error is not None:
raise _IMPORT_ERROR from _spatial_import_error
if device_id is None:
device_id = torch.cuda.current_device()