test(cp_shared_kv): load real sgl_kernel first to stop CPU-CI stub leaking
test_cp_shared_kv_runtime installs CPU-CI sgl_kernel stubs at import time via sys.modules.setdefault + torch.library FRAGMENT defs. On a GPU box, when this module is collected before the real sgl_kernel loads, the empty stub permanently shadows it process-wide, so a later real-kernel test (e.g. fast_topk_transform_ragged_fused in test_nsa_topk_transform) calls a None-returning lambda and crashes. Importing the real sgl_kernel first makes setdefault keep the real module (setattr only fills missing names) and the FRAGMENT defs hit the already-registered path; on CPU-CI the import fails and stubbing proceeds as before. Verified on GPU: test_cp_shared_kv_runtime alone 120 passed; combined with test_nsa_topk_transform 125 passed (previously 1 failed / segfault on cleanup). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -6,6 +6,20 @@ from unittest.mock import patch
|
||||
|
||||
import torch
|
||||
|
||||
# Prefer the real compiled sgl_kernel when it is available (e.g. on a GPU box).
|
||||
# The CPU-CI stubs below register torch.library ops and a sys.modules module stub
|
||||
# at import time. If this module is imported before the real sgl_kernel loads, the
|
||||
# empty stub permanently shadows it process-wide, so a later test that needs the
|
||||
# genuine kernel (e.g. fast_topk_transform_ragged_fused) calls a None-returning
|
||||
# lambda. Loading the real module first makes the FRAGMENT defines hit the
|
||||
# already-registered path and keeps sys.modules["sgl_kernel"] pointing at the real
|
||||
# module (setattr only fills genuinely-missing names). On CPU-CI the import fails
|
||||
# and the stubs are built exactly as before.
|
||||
try:
|
||||
import sgl_kernel # noqa: F401
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
_sgl_kernel_lib = torch.library.Library("sgl_kernel", "FRAGMENT")
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user