[Piecewise] Support PCG weak_ref_tensor cuda kernel on AMD (#17291)

2026-01-20 14:05:32 -08:00
parent 20ed3822bb
commit 38c233fd04
3 changed files with 9 additions and 2 deletions
--- a/python/sglang/srt/compilation/weak_ref_tensor.py
+++ b/python/sglang/srt/compilation/weak_ref_tensor.py
@@ -2,9 +2,9 @@ from typing import Any, Union

 import torch

-from sglang.srt.utils.common import is_cuda, is_npu
+from sglang.srt.utils.common import is_cuda, is_hip, is_npu

-if is_cuda():
+if is_cuda() or is_hip():
    from sgl_kernel import weak_ref_tensor
 elif is_npu():
    from torch_npu._C import _weak_ref_tensor as weak_ref_tensor