diff --git a/python/sglang/kernels/ops/attention/vattn_asm_gfx950/__init__.py b/python/sglang/kernels/ops/attention/vattn_asm_gfx950/__init__.py index b1b38f202..56576724d 100644 --- a/python/sglang/kernels/ops/attention/vattn_asm_gfx950/__init__.py +++ b/python/sglang/kernels/ops/attention/vattn_asm_gfx950/__init__.py @@ -22,6 +22,8 @@ import tempfile import torch +from sglang.srt.distributed.device_communicators.cuda_wrapper import find_loaded_library + _DIR = os.path.dirname(os.path.abspath(__file__)) _LOG2E = 1.4426950408889634 @@ -97,7 +99,15 @@ def _declared_kernarg_size(source_file): def _hip_lib(): global _hip if _hip is None: - _hip = ctypes.CDLL("libamdhip64.so") + # ROCm 10 ships the HIP runtime (_rocm_sdk_core) and the toolchain + # (_rocm_sdk_devel) as separate wheels. Torch maps core's + # libamdhip64.so.7, which an unversioned CDLL("libamdhip64.so") does not + # match, so the loader takes devel's copy off LD_LIBRARY_PATH as a + # second HIP runtime and every launch on a torch stream then fails with + # hipErrorContextIsDestroyed (709). Initialize CUDA first so torch's + # copy is mapped, then bind to that one. + torch.cuda.current_device() + _hip = ctypes.CDLL(find_loaded_library("libamdhip64") or "libamdhip64.so") _hip.hipModuleLoad.restype = ctypes.c_int _hip.hipModuleLoad.argtypes = [ctypes.c_void_p, ctypes.c_char_p] _hip.hipModuleGetFunction.restype = ctypes.c_int diff --git a/test/registered/amd/test_vattn_segplan_mi35x.py b/test/registered/amd/test_vattn_segplan_mi35x.py index 8f99a16be..e776ac857 100644 --- a/test/registered/amd/test_vattn_segplan_mi35x.py +++ b/test/registered/amd/test_vattn_segplan_mi35x.py @@ -165,6 +165,39 @@ class TestVattnSegPlan(CustomTestCase): self.assertLessEqual(e_plan, max(2 * e_leg, 0.02)) torch.cuda.empty_cache() + def test_launches_on_a_side_stream_and_under_graph_capture(self): + # ROCm 10 images carry two libamdhip64 of the same SONAME (torch loads + # _rocm_sdk_core, LD_LIBRARY_PATH points at _rocm_sdk_devel). Binding + # the wrong one still loads the module and still launches on the default + # stream, so only a torch stream catches it: every launch there fails + # with hipErrorContextIsDestroyed (709). Decode captures graphs, so this + # is the path serving actually takes. + V = self.V + lens, qlens, hq, hkv = [4096, 3777], [4, 2], 16, 1 + k, v, bt, q, cu_q, seq_lens, kd, vd = make(lens, qlens, hq, hkv) + scale = 1.0 / math.sqrt(HD) + r = ref(k, v, bt, q, cu_q, seq_lens, kd, vd, hq, hkv) + + stream = torch.cuda.Stream() + with torch.cuda.stream(stream): + eager = V.mtp_verify_attn_fwd_asm( + q, k, v, bt, seq_lens, cu_q, kd, vd, scale + ) + stream.synchronize() + self.assertLess((eager.float() - r).abs().max().item(), 0.05) + + out = torch.empty_like(eager) + g = torch.cuda.CUDAGraph() + V.reset_seg_plan_cache() + with torch.cuda.graph(g, stream=stream): + V.mtp_verify_attn_fwd_asm( + q, k, v, bt, seq_lens, cu_q, kd, vd, scale, out=out + ) + out.zero_() + g.replay() + torch.cuda.synchronize() + self.assertLess((out.float() - r).abs().max().item(), 0.05) + def test_plan_cache_per_forward(self): V = self.V lens, qlens, hq, hkv = [248000, 60000, 9000, 500, 17, 0], [4] * 6, 16, 1