Migrate weak_ref_tensor to sgl-kernel (#12505)

2025-11-02 10:55:39 +08:00
parent 9a512cf95b
commit 95191ebdca
4 changed files with 40 additions and 0 deletions
--- 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);

 /*