From 95191ebdca6684c3bd28ea566bb268a3bf63469c Mon Sep 17 00:00:00 2001
From: Xiaoyu Zhang <35585791+BBuf@users.noreply.github.com>
Date: Sun, 2 Nov 2025 10:55:39 +0800
Subject: [PATCH] Migrate weak_ref_tensor to sgl-kernel (#12505)

---
 sgl-kernel/CMakeLists.txt                  |  1 +
 sgl-kernel/csrc/common_extension.cc        |  3 ++
 sgl-kernel/csrc/memory/weak_ref_tensor.cpp | 35 ++++++++++++++++++++++
 sgl-kernel/include/sgl_kernel_ops.h        |  1 +
 4 files changed, 40 insertions(+)
 create mode 100644 sgl-kernel/csrc/memory/weak_ref_tensor.cpp

diff --git a/sgl-kernel/CMakeLists.txt b/sgl-kernel/CMakeLists.txt
index d08b4a653..0c8bb02a4 100644
--- a/sgl-kernel/CMakeLists.txt
+++ b/sgl-kernel/CMakeLists.txt
@@ -320,6 +320,7 @@ set(SOURCES
     "csrc/kvcacheio/transfer.cu"
     "csrc/mamba/causal_conv1d.cu"
     "csrc/memory/store.cu"
+    "csrc/memory/weak_ref_tensor.cpp"
 
     "csrc/moe/cutlass_moe/w4a8/scaled_mm_entry.cu"
     "csrc/moe/cutlass_moe/w4a8/w4a8_moe_data.cu"
diff --git a/sgl-kernel/csrc/common_extension.cc b/sgl-kernel/csrc/common_extension.cc
index 980db0e38..a36e9f495 100644
--- a/sgl-kernel/csrc/common_extension.cc
+++ b/sgl-kernel/csrc/common_extension.cc
@@ -396,6 +396,9 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) {
   m.def("store_kv_cache(Tensor k_cache, Tensor v_cache, Tensor out_loc, Tensor k, Tensor v) -> ()");
   m.impl("store_kv_cache", &store_kv_cache);
 
+  m.def("weak_ref_tensor(Tensor tensor) -> Tensor");
+  m.impl("weak_ref_tensor", torch::kCUDA, &weak_ref_tensor);
+
   /*
    * From FlashInfer
    */
diff --git a/sgl-kernel/csrc/memory/weak_ref_tensor.cpp b/sgl-kernel/csrc/memory/weak_ref_tensor.cpp
new file mode 100644
index 000000000..6d7f8c4e5
--- /dev/null
+++ b/sgl-kernel/csrc/memory/weak_ref_tensor.cpp
@@ -0,0 +1,35 @@
+/* Copyright 2025 SGLang Team. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+// Adapted from: https://github.com/vllm-project/vllm/blob/main/csrc/ops.h
+
+#include <ATen/ATen.h>
+#include <ATen/Tensor.h>
+
+#include <vector>
+
+at::Tensor weak_ref_tensor(const at::Tensor& tensor) {
+  TORCH_CHECK(tensor.is_cuda(), "weak_ref_tensor expects a CUDA tensor");
+
+  void* data_ptr = tensor.data_ptr();
+  std::vector<int64_t> sizes = tensor.sizes().vec();
+  std::vector<int64_t> strides = tensor.strides().vec();
+
+  auto options = tensor.options();
+
+  auto new_tensor = at::from_blob(data_ptr, sizes, strides, options);
+
+  return new_tensor;
+}
diff --git a/sgl-kernel/include/sgl_kernel_ops.h b/sgl-kernel/include/sgl_kernel_ops.h
index 497d432a1..a022b80dc 100644
--- a/sgl-kernel/include/sgl_kernel_ops.h
+++ b/sgl-kernel/include/sgl_kernel_ops.h
@@ -665,6 +665,7 @@ void transfer_kv_all_layer_direct_lf_pf(
 /*
  * From csrc/memory
  */
+at::Tensor weak_ref_tensor(const at::Tensor& tensor);
 void store_kv_cache(at::Tensor k_cache, at::Tensor v_cache, at::Tensor out_loc, at::Tensor k, at::Tensor v);
 
 /*