dylanyunlong/project_6
1
0
Fork 0
Code Issues Pull Requests Actions Projects Releases Wiki Activity

feat: pybind11 bindings for xllm CUDA kernels

Browse Source 
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:
claude
2026-08-14 10:50:13 +00:00
parent 415fff85f1
commit 093bfb380f
5 changed files with 98 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files

33
qwen3_6_scripts/build_xllm_kernels.sh
View File

@@ -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"