feat: pybind11 bindings for xllm CUDA kernels
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
This commit is contained in:
@@ -0,0 +1,18 @@
|
||||
// xllm_activation_bind.cpp
|
||||
#include <torch/extension.h>
|
||||
|
||||
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"));
|
||||
}
|
||||
19
ex_engine/xllm_kernels/cuda/bindings/xllm_cache_bind.cpp
Normal file
19
ex_engine/xllm_kernels/cuda/bindings/xllm_cache_bind.cpp
Normal file
@@ -0,0 +1,19 @@
|
||||
// xllm_cache_bind.cpp
|
||||
#include <torch/extension.h>
|
||||
|
||||
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");
|
||||
}
|
||||
24
ex_engine/xllm_kernels/cuda/bindings/xllm_norm_bind.cpp
Normal file
24
ex_engine/xllm_kernels/cuda/bindings/xllm_norm_bind.cpp
Normal file
@@ -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 <torch/extension.h>
|
||||
|
||||
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);
|
||||
}
|
||||
17
ex_engine/xllm_kernels/cuda/bindings/xllm_rope_bind.cpp
Normal file
17
ex_engine/xllm_kernels/cuda/bindings/xllm_rope_bind.cpp
Normal file
@@ -0,0 +1,17 @@
|
||||
// xllm_rope_bind.cpp
|
||||
#include <torch/extension.h>
|
||||
#include <optional>
|
||||
|
||||
namespace xllm::kernel::cuda {
|
||||
void rotary_embedding(torch::Tensor& positions, torch::Tensor& query,
|
||||
std::optional<torch::Tensor> 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);
|
||||
}
|
||||
@@ -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"
|
||||
|
||||
Reference in New Issue
Block a user