[Feature] Support custom set kv buffer kernel (#8884)

2025-08-12 16:56:51 -07:00
parent 0edda32001
commit 86a0be65d8
6 changed files with 178 additions and 0 deletions
--- a/sgl-kernel/python/sgl_kernel/init.py
+++ b/sgl-kernel/python/sgl_kernel/init.py
@@ -67,6 +67,7 @@ from sgl_kernel.marlin import (
    awq_marlin_repack,
    gptq_marlin_repack,
 )
+from sgl_kernel.memory import set_kv_buffer_kernel
 from sgl_kernel.moe import (
    apply_shuffle_mul_sum,
    cutlass_fp4_group_mm,