[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/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,