增加dcu_alloc_decode_kernel实现

2025-11-04 20:27:27 +08:00
parent 46da95569f
commit c9bcffd2a5
5 changed files with 117 additions and 8 deletions
--- a/sgl-kernel/include/sgl_kernel_ops.h
+++ b/sgl-kernel/include/sgl_kernel_ops.h
@@ -538,6 +538,15 @@ void segment_packbits(
 /*
 * From csrc/kvcacheio
 */
+void dcu_alloc_decode_kernel(
+    const at::Tensor seq_lens_ptr,   
+    const at::Tensor last_loc_ptr,    
+    const at::Tensor free_page_ptr,   
+    at::Tensor out_indices, 
+    int64_t bs,          
+    int64_t bs_upper,              
+    int64_t page_size);
+
 void transfer_kv_per_layer(
    const at::Tensor src_k,
    at::Tensor dst_k,