diff --git a/sgl-kernel/python/sgl_kernel/__init__.py b/sgl-kernel/python/sgl_kernel/__init__.py
index 1c36681e1..44fe6d45c 100644
--- a/sgl-kernel/python/sgl_kernel/__init__.py
+++ b/sgl-kernel/python/sgl_kernel/__init__.py
@@ -80,7 +80,7 @@ from sgl_kernel.marlin import (
     awq_marlin_repack,
     gptq_marlin_repack,
 )
-from sgl_kernel.memory import set_kv_buffer_kernel
+from sgl_kernel.memory import set_kv_buffer_kernel, weak_ref_tensor
 from sgl_kernel.moe import (
     apply_shuffle_mul_sum,
     cutlass_fp4_group_mm,
diff --git a/sgl-kernel/python/sgl_kernel/memory.py b/sgl-kernel/python/sgl_kernel/memory.py
index eb997db0c..9ff4f957d 100644
--- a/sgl-kernel/python/sgl_kernel/memory.py
+++ b/sgl-kernel/python/sgl_kernel/memory.py
@@ -16,3 +16,11 @@ def set_kv_buffer_kernel(
     except RuntimeError:  # ok, fallback to torch implementation
         k_cache[loc] = k
         v_cache[loc] = v
+
+
+def weak_ref_tensor(tensor):
+    return (
+        torch.ops.sgl_kernel.weak_ref_tensor(tensor)
+        if isinstance(tensor, torch.Tensor)
+        else tensor
+    )