From 547de8c7747811de1e1a6cc460116b28deced225 Mon Sep 17 00:00:00 2001
From: Xiaoyu Zhang <35585791+BBuf@users.noreply.github.com>
Date: Mon, 10 Nov 2025 22:12:59 +0800
Subject: [PATCH] [1 / 2] register weak_ref_tensor in sgl-kernel (#12999)

---
 sgl-kernel/python/sgl_kernel/__init__.py | 2 +-
 sgl-kernel/python/sgl_kernel/memory.py   | 8 ++++++++
 2 files changed, 9 insertions(+), 1 deletion(-)

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
+    )