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

2025-11-10 22:12:59 +08:00
parent 37c40a87a8
commit 547de8c774
2 changed files with 9 additions and 1 deletions
--- 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
+    )