wehub-resource-sync
78 lines
4.0 KiB
C++
78 lines
4.0 KiB
C++
// SPDX-License-Identifier: Apache-2.0
|
|
|
|
//
|
|
// Python bindings for the SYCL/XPU memory and CacheGen kernels.
|
|
// Exposed as `lmcache.xpu_ops`.
|
|
//
|
|
#include <pybind11/pybind11.h>
|
|
#include <torch/torch.h>
|
|
#include "mem_kernels_sycl.h"
|
|
#include "cachegen_kernels_sycl.h"
|
|
|
|
namespace py = pybind11;
|
|
|
|
PYBIND11_MODULE(xpu_ops, m) {
|
|
py::enum_<TransferDirection>(m, "TransferDirection")
|
|
.value("H2D", TransferDirection::H2D)
|
|
.value("D2H", TransferDirection::D2H)
|
|
.export_values();
|
|
py::enum_<EngineKVFormat>(m, "EngineKVFormat")
|
|
.value("NB_NL_TWO_BS_NH_HS", EngineKVFormat::NB_NL_TWO_BS_NH_HS)
|
|
.value("NL_X_TWO_NB_BS_NH_HS", EngineKVFormat::NL_X_TWO_NB_BS_NH_HS)
|
|
.value("NL_X_NB_TWO_BS_NH_HS", EngineKVFormat::NL_X_NB_TWO_BS_NH_HS)
|
|
.value("NL_X_NB_BS_HS", EngineKVFormat::NL_X_NB_BS_HS)
|
|
.value("TWO_X_NL_X_NBBS_NH_HS", EngineKVFormat::TWO_X_NL_X_NBBS_NH_HS)
|
|
.value("NL_X_NBBS_ONE_HS", EngineKVFormat::NL_X_NBBS_ONE_HS)
|
|
.value("NL_X_TWO_NB_NH_BS_HS", EngineKVFormat::NL_X_TWO_NB_NH_BS_HS)
|
|
.value("NL_X_NB_TWO_NH_BS_HS", EngineKVFormat::NL_X_NB_TWO_NH_BS_HS)
|
|
.value("NB_NL_TWO_NH_BS_HS", EngineKVFormat::NB_NL_TWO_NH_BS_HS)
|
|
.value("TWO_X_NL_X_NB_BS_NH_HS", EngineKVFormat::TWO_X_NL_X_NB_BS_NH_HS)
|
|
.value("NL_X_NB_NH_BS_TWO_HS", EngineKVFormat::NL_X_NB_NH_BS_TWO_HS)
|
|
.export_values();
|
|
m.def("multi_layer_kv_transfer", &multi_layer_kv_transfer,
|
|
py::arg("key_value"), py::arg("key_value_ptrs"),
|
|
py::arg("slot_mapping"), py::arg("paged_memory_device"),
|
|
py::arg("page_buffer_size"), py::arg("direction"),
|
|
py::arg("engine_kv_format"), py::arg("block_size") = 0,
|
|
py::arg("head_size") = 0, py::arg("skip_prefix_n_tokens") = 0,
|
|
py::call_guard<py::gil_scoped_release>());
|
|
m.def("single_layer_kv_transfer", &single_layer_kv_transfer,
|
|
py::arg("lmc_key_value_cache"), py::arg("vllm_key_value_cache"),
|
|
py::arg("slot_mapping"), py::arg("direction"),
|
|
py::arg("engine_kv_format"), py::arg("token_major") = false,
|
|
py::call_guard<py::gil_scoped_release>());
|
|
m.def("single_layer_kv_transfer_sgl", &single_layer_kv_transfer_sgl,
|
|
py::arg("lmc_key_value_cache"), py::arg("sgl_key_cache"),
|
|
py::arg("sgl_value_cache"), py::arg("slot_mapping"),
|
|
py::arg("direction"), py::arg("token_major") = false,
|
|
py::call_guard<py::gil_scoped_release>());
|
|
m.def("multi_layer_kv_transfer_unilateral",
|
|
&multi_layer_kv_transfer_unilateral, py::arg("key_value"),
|
|
py::arg("key_value_ptrs"), py::arg("slot_mapping"),
|
|
py::arg("paged_memory_device"), py::arg("page_buffer_size"),
|
|
py::arg("direction"), py::arg("engine_kv_format"),
|
|
py::call_guard<py::gil_scoped_release>());
|
|
m.def("load_and_reshape_flash", &load_and_reshape_flash);
|
|
m.def("reshape_and_cache_back_flash", &reshape_and_cache_back_flash);
|
|
m.def("lmcache_memcpy_async", &lmcache_memcpy_async,
|
|
py::call_guard<py::gil_scoped_release>());
|
|
|
|
// CacheGen / RoPE kernels (Intel XPU). Names match the
|
|
// lmcache.python_ops_fallback module so the backend selection in
|
|
// lmcache.v1.platform can transparently override.
|
|
m.def("calculate_cdf", &calculate_cdf_xpu, py::arg("input"),
|
|
py::arg("max_bins"));
|
|
m.def("rotary_embedding_k_fused", &rotary_embedding_k_fused_xpu,
|
|
py::arg("old_positions"), py::arg("new_positions"), py::arg("key"),
|
|
py::arg("head_size"), py::arg("cos_sin_cache"), py::arg("is_neox"));
|
|
m.def("encode_fast_new", &encode_fast_new_xpu, py::arg("cdf"),
|
|
py::arg("input_sym"), py::arg("output_buffer"),
|
|
py::arg("output_lengths"));
|
|
m.def("decode_fast_new", &decode_fast_new_xpu, py::arg("cdf"),
|
|
py::arg("bytestreams"), py::arg("lengths"), py::arg("output"));
|
|
m.def("decode_fast_prefsum", &decode_fast_prefsum_xpu, py::arg("cdf"),
|
|
py::arg("bytestreams"), py::arg("lengths_prefsum"), py::arg("output"));
|
|
// Backward-compat alias: GPUKVFormat -> EngineKVFormat
|
|
m.attr("GPUKVFormat") = m.attr("EngineKVFormat");
|
|
}
|