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

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",
 ]