diff --git a/python/sglang/srt/compilation/weak_ref_tensor.py b/python/sglang/srt/compilation/weak_ref_tensor.py
index 83cb7e64b..30f59d46f 100644
--- 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
diff --git a/sgl-kernel/csrc/common_extension_rocm.cc b/sgl-kernel/csrc/common_extension_rocm.cc
index d0147e5f3..add827787 100644
--- a/sgl-kernel/csrc/common_extension_rocm.cc
+++ b/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)
diff --git a/sgl-kernel/setup_rocm.py b/sgl-kernel/setup_rocm.py
index a7336a103..16fdbbd2f 100644
--- a/sgl-kernel/setup_rocm.py
+++ b/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",
 ]