[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/include/sgl_kernel_ops.h
+++ b/sgl-kernel/include/sgl_kernel_ops.h
@@ -699,3 +699,8 @@ void qserve_w4a8_per_group_gemm(
 * From csrc/spatial
 */
 std::vector<int64_t> create_greenctx_stream_by_value(int64_t smA, int64_t smB, int64_t device);
+
+/*
+ * From csrc/memory
+ */
+void store_kv_cache(at::Tensor k_cache, at::Tensor v_cache, at::Tensor out_loc, at::Tensor k, at::Tensor v);