[AMD][Bugfix] Fix vattn_asm HIP error 709 under CUDA graph capture on ROCm 10 (#39513)

Co-authored-by: jacky.cheng <yichiche@amd.com>

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

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