"docs/vscode:/vscode.git/clone" did not exist on "2ba586d5b30948dd7c1a45c969ad9b9ec12b7788"
Unverified Commit fe68c148 authored by Hubert Lu's avatar Hubert Lu Committed by GitHub
Browse files

Fix errors of hicache kernels in sgl-kernel for ROCm (#10339)

parent 70c0c1f9
......@@ -163,6 +163,14 @@ TORCH_LIBRARY_EXPAND(sgl_kernel, m) {
"transfer_kv_direct(Tensor[] src_layers, Tensor[] dst_layers, Tensor src_indices, Tensor dst_indices, int "
"page_size) -> ()");
m.impl("transfer_kv_direct", torch::kCUDA, &transfer_kv_direct);
m.def(
"transfer_kv_per_layer_direct_pf_lf(Tensor[] src_ptrs, Tensor[] dst_ptrs, Tensor src_indices, "
"Tensor dst_indices, int layer_id, int page_size)->() ");
m.impl("transfer_kv_per_layer_direct_pf_lf", torch::kCUDA, &transfer_kv_per_layer_direct_pf_lf);
m.def(
"transfer_kv_all_layer_direct_lf_pf(Tensor[] src_ptrs, Tensor[] dst_ptrs, Tensor src_indices, "
"Tensor dst_indices, int page_size) ->() ");
m.impl("transfer_kv_all_layer_direct_lf_pf", torch::kCUDA, &transfer_kv_all_layer_direct_lf_pf);
}
REGISTER_EXTENSION(common_ops)
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment