From 093bfb380ff6de4a59700538aeeb562665a89f67 Mon Sep 17 00:00:00 2001 From: claude Date: Fri, 14 Aug 2026 10:50:13 +0000 Subject: [PATCH] feat: pybind11 bindings for xllm CUDA kernels MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit norm.cu compiled successfully on BI-V100 (only warning: fp8 __host__ attr). Failed at import because no PYBIND11_MODULE — now fixed. New bindings/ directory with 4 binding files: xllm_norm_bind.cpp → rms_norm, fused_add_rms_norm xllm_activation_bind.cpp → silu_and_mul, gelu_and_mul, act_and_mul xllm_rope_bind.cpp → rotary_embedding xllm_cache_bind.cpp → reshape_paged_cache, block_copy Build script updated: each .so = kernel .cu + binding .cpp --- .../cuda/bindings/xllm_activation_bind.cpp | 18 ++++++++++ .../cuda/bindings/xllm_cache_bind.cpp | 19 +++++++++++ .../cuda/bindings/xllm_norm_bind.cpp | 24 ++++++++++++++ .../cuda/bindings/xllm_rope_bind.cpp | 17 ++++++++++ qwen3_6_scripts/build_xllm_kernels.sh | 33 +++++++++++-------- 5 files changed, 98 insertions(+), 13 deletions(-) create mode 100644 ex_engine/xllm_kernels/cuda/bindings/xllm_activation_bind.cpp create mode 100644 ex_engine/xllm_kernels/cuda/bindings/xllm_cache_bind.cpp create mode 100644 ex_engine/xllm_kernels/cuda/bindings/xllm_norm_bind.cpp create mode 100644 ex_engine/xllm_kernels/cuda/bindings/xllm_rope_bind.cpp diff --git a/ex_engine/xllm_kernels/cuda/bindings/xllm_activation_bind.cpp b/ex_engine/xllm_kernels/cuda/bindings/xllm_activation_bind.cpp new file mode 100644 index 00000000..fbdaf693 --- /dev/null +++ b/ex_engine/xllm_kernels/cuda/bindings/xllm_activation_bind.cpp @@ -0,0 +1,18 @@ +// xllm_activation_bind.cpp +#include + +namespace xllm::kernel::cuda { +void act_and_mul(torch::Tensor out, torch::Tensor input, + const std::string& act_mode); +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("silu_and_mul", [](torch::Tensor out, torch::Tensor input) { + xllm::kernel::cuda::act_and_mul(out, input, "silu"); + }, "SiLU and Mul", py::arg("out"), py::arg("input")); + m.def("gelu_and_mul", [](torch::Tensor out, torch::Tensor input) { + xllm::kernel::cuda::act_and_mul(out, input, "gelu"); + }, "GELU and Mul", py::arg("out"), py::arg("input")); + m.def("act_and_mul", &xllm::kernel::cuda::act_and_mul, + "Activation and Mul", py::arg("out"), py::arg("input"), py::arg("act_mode")); +} diff --git a/ex_engine/xllm_kernels/cuda/bindings/xllm_cache_bind.cpp b/ex_engine/xllm_kernels/cuda/bindings/xllm_cache_bind.cpp new file mode 100644 index 00000000..e0539922 --- /dev/null +++ b/ex_engine/xllm_kernels/cuda/bindings/xllm_cache_bind.cpp @@ -0,0 +1,19 @@ +// xllm_cache_bind.cpp +#include + +namespace xllm::kernel::cuda { +void reshape_paged_cache(torch::Tensor slot_ids, torch::Tensor keys, + torch::Tensor values, torch::Tensor key_cache, + torch::Tensor value_cache); +void block_copy(torch::Tensor key_cache_ptrs, torch::Tensor value_cache_ptrs, + torch::Tensor src_block_indices, torch::Tensor dst_block_indices, + torch::Tensor cum_sum, int64_t numel_per_block, + torch::ScalarType cache_dtype); +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("reshape_paged_cache", &xllm::kernel::cuda::reshape_paged_cache, + "Reshape Paged KV Cache"); + m.def("block_copy", &xllm::kernel::cuda::block_copy, + "Block Copy for KV Cache"); +} diff --git a/ex_engine/xllm_kernels/cuda/bindings/xllm_norm_bind.cpp b/ex_engine/xllm_kernels/cuda/bindings/xllm_norm_bind.cpp new file mode 100644 index 00000000..dee3e003 --- /dev/null +++ b/ex_engine/xllm_kernels/cuda/bindings/xllm_norm_bind.cpp @@ -0,0 +1,24 @@ +// xllm_norm_bind.cpp — pybind11 entry point for xllm norm kernels +// Compiled together with norm.cu to produce xllm_norm.so +// +// Exports: rms_norm, fused_add_rms_norm + +#include + +namespace xllm::kernel::cuda { +void rms_norm(torch::Tensor output, torch::Tensor input, + torch::Tensor weight, double eps); +void fused_add_rms_norm(torch::Tensor& input, torch::Tensor& residual, + torch::Tensor& weight, double epsilon); +} // namespace xllm::kernel::cuda + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("rms_norm", &xllm::kernel::cuda::rms_norm, + "RMS Norm (xllm CUDA kernel)", + py::arg("output"), py::arg("input"), + py::arg("weight"), py::arg("eps") = 1e-6); + m.def("fused_add_rms_norm", &xllm::kernel::cuda::fused_add_rms_norm, + "Fused Add + RMS Norm (xllm CUDA kernel)", + py::arg("input"), py::arg("residual"), + py::arg("weight"), py::arg("epsilon") = 1e-6); +} diff --git a/ex_engine/xllm_kernels/cuda/bindings/xllm_rope_bind.cpp b/ex_engine/xllm_kernels/cuda/bindings/xllm_rope_bind.cpp new file mode 100644 index 00000000..644b4e84 --- /dev/null +++ b/ex_engine/xllm_kernels/cuda/bindings/xllm_rope_bind.cpp @@ -0,0 +1,17 @@ +// xllm_rope_bind.cpp +#include +#include + +namespace xllm::kernel::cuda { +void rotary_embedding(torch::Tensor& positions, torch::Tensor& query, + std::optional key, + torch::Tensor& cos_sin_cache, bool is_neox); +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("rotary_embedding", &xllm::kernel::cuda::rotary_embedding, + "Rotary Position Embedding (xllm CUDA kernel)", + py::arg("positions"), py::arg("query"), + py::arg("key"), py::arg("cos_sin_cache"), + py::arg("is_neox") = true); +} diff --git a/qwen3_6_scripts/build_xllm_kernels.sh b/qwen3_6_scripts/build_xllm_kernels.sh index 6b5f8108..4537b3c8 100644 --- a/qwen3_6_scripts/build_xllm_kernels.sh +++ b/qwen3_6_scripts/build_xllm_kernels.sh @@ -1,9 +1,8 @@ #!/bin/bash # build_xllm_kernels.sh — Compile xllm CUDA kernels into .so on BI-V100 # -# Uses corex's own CUB (/usr/local/corex/include/cub/) NOT cccl_upstream -# Source: ex_engine/xllm_kernels/cuda/ -# Output: qwen3_6_scripts/prebuilt/corex-3.2.3-ivcore10/ +# Uses corex CUB (/usr/local/corex/include/cub/) NOT cccl_upstream +# Each .so = kernel .cu + pybind11 binding .cpp # # Run: bash qwen3_6_scripts/build_xllm_kernels.sh @@ -12,20 +11,23 @@ SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" PROJECT_DIR="$(dirname "$SCRIPT_DIR")" CUDA_DIR="${PROJECT_DIR}/ex_engine/xllm_kernels/cuda" HEADER_DIR="${CUDA_DIR}/headers" +BIND_DIR="${CUDA_DIR}/bindings" PREBUILT_DIR="${SCRIPT_DIR}/prebuilt/corex-3.2.3-ivcore10" mkdir -p "$PREBUILT_DIR" build_kernel() { local name="$1" - local cu_file="$2" + shift + local sources="$@" echo "=== Building ${name}.so ===" python3 -c " import os, glob, shutil from torch.utils.cpp_extension import load +sources = '${sources}'.split() mod = load( name='${name}', - sources=['${cu_file}'], + sources=sources, extra_cflags=['-std=c++17'], extra_include_paths=['${HEADER_DIR}', '/usr/local/corex/include'], verbose=True, @@ -43,19 +45,24 @@ for f in glob.glob(os.path.join(build_dir, '*.so')): fns = [x for x in dir(mod) if not x.startswith('_')] print(f'Functions: {fns}') " + echo "" } echo "Building xllm CUDA kernels for BI-V100 (ivcore10)" echo "Using corex CUB: /usr/local/corex/include/cub/" echo "" -# Build each kernel -build_kernel "xllm_norm" "${CUDA_DIR}/norm.cu" -build_kernel "xllm_activation" "${CUDA_DIR}/activation.cu" -build_kernel "xllm_rope" "${CUDA_DIR}/rope.cu" -build_kernel "xllm_block_copy" "${CUDA_DIR}/block_copy.cu" -build_kernel "xllm_cache" "${CUDA_DIR}/reshape_paged_cache.cu" +build_kernel "xllm_norm" \ + "${CUDA_DIR}/norm.cu" "${BIND_DIR}/xllm_norm_bind.cpp" + +build_kernel "xllm_activation" \ + "${CUDA_DIR}/activation.cu" "${BIND_DIR}/xllm_activation_bind.cpp" + +build_kernel "xllm_rope" \ + "${CUDA_DIR}/rope.cu" "${BIND_DIR}/xllm_rope_bind.cpp" + +build_kernel "xllm_cache" \ + "${CUDA_DIR}/reshape_paged_cache.cu" "${CUDA_DIR}/block_copy.cu" "${BIND_DIR}/xllm_cache_bind.cpp" -echo "" echo "=== All kernels built ===" -ls -la "${PREBUILT_DIR}"/xllm_*.so 2>/dev/null +ls -lh "${PREBUILT_DIR}"/xllm_*.so 2>/dev/null || echo "No .so files found"