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

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

sgl-kernel/csrc/common_extension_rocm.cc

@@ -231,6 +231,12 @@ TORCH_LIBRARY_EXPAND(sgl_kernel, m) {
       "float scale,"
       "int max_period) -> Tensor");
   m.impl("timestep_embedding", torch::kCUDA, &timestep_embedding);
+
+  /*
+   * From csrc/memory
+   */
+  m.def("weak_ref_tensor(Tensor tensor) -> Tensor");
+  m.impl("weak_ref_tensor", torch::kCUDA, &weak_ref_tensor);
 }
 
 REGISTER_EXTENSION(common_ops)

sgl-kernel/setup_rocm.py

@@ -53,6 +53,7 @@ sources = [
     "csrc/moe/moe_topk_sigmoid_kernels.cu",
     "csrc/speculative/eagle_utils.cu",
     "csrc/kvcacheio/transfer.cu",
+    "csrc/memory/weak_ref_tensor.cpp",
     "csrc/elementwise/pos_enc.cu",
     "csrc/sgl_diffusion/elementwise/timestep_embedding.cu",
 ]