diff --git a/qwen3_6_scripts/build_cccl_moe_sort_scatter.sh b/qwen3_6_scripts/build_cccl_moe_sort_scatter.sh old mode 100644 new mode 100755 index 0511b51c..532fa0c7 --- a/qwen3_6_scripts/build_cccl_moe_sort_scatter.sh +++ b/qwen3_6_scripts/build_cccl_moe_sort_scatter.sh @@ -1,76 +1,74 @@ #!/usr/bin/env bash -# Build cccl_moe_sort_scatter.so using CCCL upstream headers +# Build cccl_moe_sort_scatter — split compilation +# +# Step 1: Compile .cu with CCCL headers (no torch) → .o +# Step 2: Compile _pybind.cpp with torch headers (no CCCL) → .o +# Step 3: Link both → .so set -euo pipefail SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" -SRC="${SCRIPT_DIR}/cccl_moe_sort_scatter.cu" INC="${SCRIPT_DIR}/cccl_preload/include" +CU_SRC="${SCRIPT_DIR}/cccl_moe_sort_scatter.cu" +PY_SRC="${SCRIPT_DIR}/cccl_moe_sort_scatter_pybind.cpp" OUT="${1:-${SCRIPT_DIR}/prebuilt/corex-3.2.3-ivcore10/cccl_moe_sort_scatter.so}" -[[ -f "${SRC}" ]] || { echo "Source not found: ${SRC}"; exit 2; } -[[ -d "${INC}/cub" ]] || { echo "CCCL include tree missing: ${INC}/cub"; exit 2; } - -# Find NVCC or corex clang -NVCC="" -for candidate in \ - /usr/local/corex/bin/nvcc \ - /usr/local/cuda/bin/nvcc \ - ; do - if [[ -x "${candidate}" ]]; then - NVCC="${candidate}" - break - fi +# Find corex clang++ +CXX="" +for c in /usr/local/corex-3.2.3/bin/clang++ /usr/local/corex/bin/clang++; do + [[ -x "$c" ]] && CXX="$c" && break done +[[ -n "${CXX}" ]] || { echo "no corex clang++"; exit 2; } -TORCH_INC=$(python3 -c "import torch; print(torch.utils.cpp_extension.include_paths()[0])" 2>/dev/null) -TORCH_LIB=$(python3 -c "import torch; print(torch.utils.cmake_prefix_path + '/../lib')" 2>/dev/null || echo "") -PYTHON_INC=$(python3 -c "from sysconfig import get_paths; print(get_paths()['include'])" 2>/dev/null) +# Find torch paths +TORCH_INC=$(python3 -c "import torch; print(torch.utils.cpp_extension.include_paths()[0])") +TORCH_LIB=$(python3 -c "import torch.utils.cpp_extension as e; import os; print(os.path.join(os.path.dirname(e.__file__), '..', '..', 'lib'))" | xargs realpath) +PYTHON_INC=$(python3 -c "from sysconfig import get_paths; print(get_paths()['include'])") +CUDA_INC="/usr/local/corex/include" -echo "[build] NVCC: ${NVCC:-not found}" -echo "[build] CCCL: ${INC}" -echo "[build] Torch: ${TORCH_INC}" -echo "[build] Output: ${OUT}" +echo "[build] CXX=${CXX}" +echo "[build] CCCL=${INC}" +echo "[build] torch=${TORCH_INC}" -if [[ -n "${NVCC}" ]]; then - "${NVCC}" \ - -shared --compiler-options -fPIC \ - -O3 -std=c++17 \ - -I"${INC}" \ - -I"${TORCH_INC}" \ - -I"${TORCH_INC}/torch/csrc/api/include" \ - ${PYTHON_INC:+-I"${PYTHON_INC}"} \ - -DCUB_WRAPPED_NAMESPACE=cccl_moe \ - -DTORCH_EXTENSION_NAME=cccl_moe_sort_scatter \ - -x cu \ - -o "${OUT}" "${SRC}" \ - -ltorch -lc10 -ltorch_cuda -ltorch_cpu \ - ${TORCH_LIB:+-L"${TORCH_LIB}"} \ - 2>&1 -else - echo "[build] No nvcc found, trying torch JIT at runtime" - python3 -c " -from torch.utils.cpp_extension import load -mod = load( - name='cccl_moe_sort_scatter', - sources=['${SRC}'], - extra_include_paths=['${INC}'], - extra_cuda_cflags=['-O3', '-DCUB_WRAPPED_NAMESPACE=cccl_moe'], - verbose=True, -) -print('[build] JIT compiled successfully') -import shutil, os -# Copy to output -src_so = os.path.join(os.path.dirname(mod.__file__), 'cccl_moe_sort_scatter.so') -if os.path.exists(src_so): - os.makedirs(os.path.dirname('${OUT}'), exist_ok=True) - shutil.copy2(src_so, '${OUT}') - print(f'[build] Copied to ${OUT}') -" 2>&1 -fi +# Step 1: Compile CUDA kernels (CCCL headers, no torch) +echo "[build] Step 1: compile CUDA kernels..." +"${CXX}" \ + -fPIC -O3 -std=c++17 \ + -I"${INC}" \ + -I"${CUDA_INC}" \ + -DCCCL_IGNORE_DEPRECATED_CUDA_BELOW_12 \ + -DCUB_WRAPPED_NAMESPACE=cccl_moe \ + --cuda-gpu-arch=ivcore10 \ + --cuda-path=/usr/local/corex \ + -c "${CU_SRC}" -o /tmp/cccl_moe_kernels.o \ + 2>&1 -if [[ -f "${OUT}" ]]; then - echo "[build] SUCCESS: ${OUT} ($(stat -c%s "${OUT}" 2>/dev/null || echo '?') bytes)" -else - echo "[build] FAILED" - exit 1 -fi +# Step 2: Compile pybind wrapper (torch headers, no CCCL) +echo "[build] Step 2: compile pybind wrapper..." +"${CXX}" \ + -fPIC -O2 -std=c++17 \ + -I"${TORCH_INC}" \ + -I"${TORCH_INC}/torch/csrc/api/include" \ + -I"${PYTHON_INC}" \ + -I"${CUDA_INC}" \ + -D_GLIBCXX_USE_CXX11_ABI=0 \ + -DTORCH_EXTENSION_NAME=cccl_moe_sort_scatter \ + -x c++ \ + -c "${PY_SRC}" -o /tmp/cccl_moe_pybind.o \ + 2>&1 + +# Step 3: Link +echo "[build] Step 3: link..." +mkdir -p "$(dirname "${OUT}")" +"${CXX}" \ + -shared -fPIC \ + /tmp/cccl_moe_kernels.o \ + /tmp/cccl_moe_pybind.o \ + -L"${TORCH_LIB}" \ + -ltorch -lc10 -ltorch_cpu -ltorch_cuda \ + -L/usr/local/corex/lib64 -lcudart \ + -Wl,-rpath,"${TORCH_LIB}" \ + -o "${OUT}" \ + 2>&1 + +SIZE=$(stat -c%s "${OUT}" 2>/dev/null || echo "?") +echo "[build] SUCCESS: ${OUT} (${SIZE} bytes)" diff --git a/qwen3_6_scripts/cccl_moe_sort_scatter.cu b/qwen3_6_scripts/cccl_moe_sort_scatter.cu index 87c50790..83ca2fa0 100644 --- a/qwen3_6_scripts/cccl_moe_sort_scatter.cu +++ b/qwen3_6_scripts/cccl_moe_sort_scatter.cu @@ -1,32 +1,27 @@ -// cccl_moe_sort_scatter.cu — Block-level CUB MoE token dispatch +// cccl_moe_sort_scatter.cu — CCCL CUB device-level MoE token dispatch // -// Uses CUB BlockScan (already proven on BI-V100 in corex_moe_index_combine.cu) -// for histogram + prefix_sum + scatter. No device-level CUB API (conflicts -// with corex's thrust/complex.h on CUDA 10.2). +// Split compilation: this file uses CCCL headers only (no torch). +// Pybind wrapper in cccl_moe_sort_scatter_pybind.cpp links against this. // -// Three kernels (same as corex_moe_index_combine but with CUB BlockRadixSort -// for the scatter step): -// 1. histogram — atomicAdd per expert -// 2. prefix_sum — CUB BlockScan ExclusiveSum -// 3. place — atomicAdd scatter into sorted positions -// -// Build: torch.utils.cpp_extension.load( -// name="cccl_moe_sort_scatter", -// sources=["cccl_moe_sort_scatter.cu"], -// extra_cuda_cflags=["-O3"], -// ) +// Build pattern (same as cccl_allocator_preload.cu): +// clang++ -I cccl_preload/include -DCCCL_IGNORE_DEPRECATED_CUDA_BELOW_12 +// -DCUB_WRAPPED_NAMESPACE=cccl_moe ... + +// Suppress CUDA <12 check — corex 10.2 works for block-level CUB +#define CCCL_IGNORE_DEPRECATED_CUDA_BELOW_12 + +// Isolate from corex CUB +#define CUB_WRAPPED_NAMESPACE cccl_moe -#include -#include -#include #include +#include +#include // ======================================================================== -// Block-level CUB kernels (proven on BI-V100 corex clang++) -// Same pattern as corex_moe_index_combine.cu +// Kernels // ======================================================================== -constexpr int32_t kBlock = 256; +static constexpr int32_t kBlock = 256; __global__ void moe_histogram_kernel( const int32_t* __restrict__ expert_id, @@ -45,9 +40,8 @@ __global__ void moe_histogram_kernel( __global__ void moe_prefix_sum_kernel( const int32_t* __restrict__ expert_sizes, int32_t* __restrict__ expert_offsets, - int32_t num_experts, - int64_t* __restrict__ total_out) { - using BlockScan = cub::BlockScan; + int32_t num_experts) { + using BlockScan = cccl_moe::cub::BlockScan; __shared__ typename BlockScan::TempStorage s_scan; int32_t val = (threadIdx.x < num_experts) ? expert_sizes[threadIdx.x] : 0; @@ -58,15 +52,11 @@ __global__ void moe_prefix_sum_kernel( if (threadIdx.x < num_experts) { expert_offsets[threadIdx.x] = offset; } - if (threadIdx.x == 0 && total_out != nullptr) { - // Last thread's offset + val = total - *total_out = offset + val; - } } __global__ void moe_place_kernel( const int32_t* __restrict__ expert_id, - int32_t* __restrict__ expert_offsets, // modified in-place by atomicAdd + int32_t* __restrict__ expert_offsets, int32_t* __restrict__ dst_src, int32_t* __restrict__ src_dst, int64_t num_elements, @@ -82,57 +72,32 @@ __global__ void moe_place_kernel( src_dst[flat_idx] = pos; } +// ======================================================================== +// C API — called from pybind wrapper +// ======================================================================== -// Same proven 3-kernel approach as corex_moe_index_combine.cu but with -// an additional inverse-scatter output for full compatibility. - -std::tuple -moe_sort_scatter(const torch::Tensor& expert_id, int64_t num_experts) { - TORCH_CHECK(expert_id.is_cuda(), "expert_id must be on CUDA"); - auto stream = at::cuda::getCurrentCUDAStream(); - int64_t N = expert_id.numel(); - int32_t E = static_cast(num_experts); - - auto expert_id_i32 = expert_id.to(torch::kInt32).contiguous(); - auto opt_i32 = expert_id_i32.options(); - - auto expert_sizes = torch::zeros({num_experts}, opt_i32); - auto expert_offsets = torch::empty({num_experts}, opt_i32); - auto dst_src = torch::empty({N}, opt_i32); - auto src_dst = torch::empty({N}, opt_i32); +extern "C" { +void cccl_moe_launch_histogram( + const int32_t* expert_id, int32_t* expert_sizes, + int64_t N, int32_t E, cudaStream_t stream) { int64_t grid = (N + kBlock - 1) / kBlock; + moe_histogram_kernel<<>>(expert_id, expert_sizes, N, E); +} - // Step 1: histogram - moe_histogram_kernel<<>>( - expert_id_i32.data_ptr(), - expert_sizes.data_ptr(), - N, E); +void cccl_moe_launch_prefix_sum( + const int32_t* expert_sizes, int32_t* expert_offsets, + int32_t E, cudaStream_t stream) { + moe_prefix_sum_kernel<<<1, kBlock, 0, stream>>>(expert_sizes, expert_offsets, E); +} - // Step 2: prefix sum (CUB BlockScan) - moe_prefix_sum_kernel<<<1, kBlock, 0, stream>>>( - expert_sizes.data_ptr(), - expert_offsets.data_ptr(), - E, nullptr); - - // Step 3: scatter — place each token into its sorted position +void cccl_moe_launch_place( + const int32_t* expert_id, int32_t* expert_offsets, + int32_t* dst_src, int32_t* src_dst, + int64_t N, int32_t E, cudaStream_t stream) { + int64_t grid = (N + kBlock - 1) / kBlock; moe_place_kernel<<>>( - expert_id_i32.data_ptr(), - expert_offsets.data_ptr(), - dst_src.data_ptr(), - src_dst.data_ptr(), - N, E); - - return std::make_tuple(src_dst, dst_src, expert_sizes); + expert_id, expert_offsets, dst_src, src_dst, N, E); } - -// ======================================================================== -// pybind -// ======================================================================== - -PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { - m.def("moe_sort_scatter", &moe_sort_scatter, - "CUB DeviceRadixSort-based MoE token dispatch " - "(sort expert_ids, compute offsets+sizes)"); -} +} // extern "C" diff --git a/qwen3_6_scripts/cccl_moe_sort_scatter_pybind.cpp b/qwen3_6_scripts/cccl_moe_sort_scatter_pybind.cpp new file mode 100644 index 00000000..fc2b0e1c --- /dev/null +++ b/qwen3_6_scripts/cccl_moe_sort_scatter_pybind.cpp @@ -0,0 +1,62 @@ +// cccl_moe_sort_scatter_pybind.cpp — Torch pybind wrapper +// +// Links against cccl_moe_sort_scatter.so (built separately with CCCL headers). +// This file only includes torch headers — no CCCL, no namespace conflict. + +#include +#include +#include + +// C API from cccl_moe_sort_scatter.so +extern "C" { +void cccl_moe_launch_histogram( + const int32_t* expert_id, int32_t* expert_sizes, + int64_t N, int32_t E, cudaStream_t stream); +void cccl_moe_launch_prefix_sum( + const int32_t* expert_sizes, int32_t* expert_offsets, + int32_t E, cudaStream_t stream); +void cccl_moe_launch_place( + const int32_t* expert_id, int32_t* expert_offsets, + int32_t* dst_src, int32_t* src_dst, + int64_t N, int32_t E, cudaStream_t stream); +} + +std::tuple +moe_sort_scatter(const torch::Tensor& expert_id, int64_t num_experts) { + TORCH_CHECK(expert_id.is_cuda(), "expert_id must be on CUDA"); + auto stream = at::cuda::getCurrentCUDAStream(); + int64_t N = expert_id.numel(); + int32_t E = static_cast(num_experts); + + auto expert_id_i32 = expert_id.to(torch::kInt32).contiguous(); + auto opt_i32 = expert_id_i32.options(); + + auto expert_sizes = torch::zeros({num_experts}, opt_i32); + auto expert_offsets = torch::empty({num_experts}, opt_i32); + auto dst_src = torch::empty({N}, opt_i32); + auto src_dst = torch::empty({N}, opt_i32); + + cccl_moe_launch_histogram( + expert_id_i32.data_ptr(), + expert_sizes.data_ptr(), + N, E, stream); + + cccl_moe_launch_prefix_sum( + expert_sizes.data_ptr(), + expert_offsets.data_ptr(), + E, stream); + + cccl_moe_launch_place( + expert_id_i32.data_ptr(), + expert_offsets.data_ptr(), + dst_src.data_ptr(), + src_dst.data_ptr(), + N, E, stream); + + return std::make_tuple(src_dst, dst_src, expert_sizes); +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("moe_sort_scatter", &moe_sort_scatter, + "CCCL CUB-based MoE token dispatch (histogram+prefix_sum+scatter)"); +} diff --git a/qwen3_6_scripts/cccl_preload/include/cuda/std/__functional/invoke.h b/qwen3_6_scripts/cccl_preload/include/cuda/std/__functional/invoke.h index 0fc242a2..0cc92b1c 100644 --- a/qwen3_6_scripts/cccl_preload/include/cuda/std/__functional/invoke.h +++ b/qwen3_6_scripts/cccl_preload/include/cuda/std/__functional/invoke.h @@ -43,7 +43,8 @@ _CCCL_BEGIN_NAMESPACE_CUDA_STD struct __any { - _CCCL_API inline __any(...); + template + _CCCL_API inline __any(_T); }; template diff --git a/qwen3_6_scripts/cccl_preload/include/cuda/std/__type_traits/conjunction.h b/qwen3_6_scripts/cccl_preload/include/cuda/std/__type_traits/conjunction.h index 996cb6b6..e3b5cc0d 100644 --- a/qwen3_6_scripts/cccl_preload/include/cuda/std/__type_traits/conjunction.h +++ b/qwen3_6_scripts/cccl_preload/include/cuda/std/__type_traits/conjunction.h @@ -35,7 +35,7 @@ template _CCCL_HOST_DEVICE __expand_to_true...> __and_helper(int); template -_CCCL_HOST_DEVICE false_type __and_helper(...); +_CCCL_HOST_DEVICE false_type __and_helper(long); // _And always performs lazy evaluation of its arguments. // diff --git a/qwen3_6_scripts/cccl_preload/include/cuda/std/__type_traits/is_referenceable.h b/qwen3_6_scripts/cccl_preload/include/cuda/std/__type_traits/is_referenceable.h index d8949b8f..3524096a 100644 --- a/qwen3_6_scripts/cccl_preload/include/cuda/std/__type_traits/is_referenceable.h +++ b/qwen3_6_scripts/cccl_preload/include/cuda/std/__type_traits/is_referenceable.h @@ -39,7 +39,7 @@ struct __cccl_is_referenceable_impl template _CCCL_HOST_DEVICE static _Tp& __test(int); template - _CCCL_HOST_DEVICE static false_type __test(...); + _CCCL_HOST_DEVICE static false_type __test(long); }; template