[1 / 2] register weak_ref_tensor in sgl-kernel (#12999)

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,

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