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
@@ -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
   */