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

fix(CCCL): split compilation to isolate CCCL headers from torch/corex

Browse Source 
Two problems from real BI-V100 build:

1. 'CUDA versions below 12 are not supported'
   → Add CCCL_IGNORE_DEPRECATED_CUDA_BELOW_12 (official suppress macro)

2. corex thrust/complex.h conflicts with CCCL thrust headers
   → Split into two compilation units:
     - cccl_moe_sort_scatter.cu: CCCL headers only, C API, no torch
     - cccl_moe_sort_scatter_pybind.cpp: torch headers only, no CCCL
   Same pattern as proven cccl_allocator_preload.cu

3. Variadic device functions rejected by corex clang:
   → is_referenceable.h: __test(...) → __test(long)
   → invoke.h: __any(...) → template __any(_T)
   → conjunction.h: __and_helper(...) → __and_helper(long)
   SFINAE still works: int overload wins, long is fallback.
This commit is contained in:
project6-dev
2026-08-13 11:35:43 +00:00
parent 05706f0d60
commit f6cf9d662e
6 changed files with 168 additions and 142 deletions
Download Patch File Download Diff File Expand all files Collapse all files

126
qwen3_6_scripts/build_cccl_moe_sort_scatter.sh Normal file → Executable file
View File

@@ -1,76 +1,74 @@
#!/usr/bin/env bash #!/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 set -euo pipefail
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
SRC="${SCRIPT_DIR}/cccl_moe_sort_scatter.cu"
INC="${SCRIPT_DIR}/cccl_preload/include" 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}" OUT="${1:-${SCRIPT_DIR}/prebuilt/corex-3.2.3-ivcore10/cccl_moe_sort_scatter.so}"
[[ -f "${SRC}" ]] || { echo "Source not found: ${SRC}"; exit 2; } # Find corex clang++
[[ -d "${INC}/cub" ]] || { echo "CCCL include tree missing: ${INC}/cub"; exit 2; } CXX=""
for c in /usr/local/corex-3.2.3/bin/clang++ /usr/local/corex/bin/clang++; do
# Find NVCC or corex clang [[ -x "$c" ]] && CXX="$c" && break
NVCC=""
for candidate in \
/usr/local/corex/bin/nvcc \
/usr/local/cuda/bin/nvcc \
; do
if [[ -x "${candidate}" ]]; then
NVCC="${candidate}"
break
fi
done 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) # Find torch paths
TORCH_LIB=$(python3 -c "import torch; print(torch.utils.cmake_prefix_path + '/../lib')" 2>/dev/null || echo "") TORCH_INC=$(python3 -c "import torch; print(torch.utils.cpp_extension.include_paths()[0])")
PYTHON_INC=$(python3 -c "from sysconfig import get_paths; print(get_paths()['include'])" 2>/dev/null) 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] CXX=${CXX}"
echo "[build] CCCL: ${INC}" echo "[build] CCCL=${INC}"
echo "[build] Torch: ${TORCH_INC}" echo "[build] torch=${TORCH_INC}"
echo "[build] Output: ${OUT}"
if [[ -n "${NVCC}" ]]; then # Step 1: Compile CUDA kernels (CCCL headers, no torch)
"${NVCC}" \ echo "[build] Step 1: compile CUDA kernels..."
-shared --compiler-options -fPIC \ "${CXX}" \
-O3 -std=c++17 \ -fPIC -O3 -std=c++17 \
-I"${INC}" \ -I"${INC}" \
-I"${TORCH_INC}" \ -I"${CUDA_INC}" \
-I"${TORCH_INC}/torch/csrc/api/include" \ -DCCCL_IGNORE_DEPRECATED_CUDA_BELOW_12 \
${PYTHON_INC:+-I"${PYTHON_INC}"} \ -DCUB_WRAPPED_NAMESPACE=cccl_moe \
-DCUB_WRAPPED_NAMESPACE=cccl_moe \ --cuda-gpu-arch=ivcore10 \
-DTORCH_EXTENSION_NAME=cccl_moe_sort_scatter \ --cuda-path=/usr/local/corex \
-x cu \ -c "${CU_SRC}" -o /tmp/cccl_moe_kernels.o \
-o "${OUT}" "${SRC}" \ 2>&1
-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
if [[ -f "${OUT}" ]]; then # Step 2: Compile pybind wrapper (torch headers, no CCCL)
echo "[build] SUCCESS: ${OUT} ($(stat -c%s "${OUT}" 2>/dev/null || echo '?') bytes)" echo "[build] Step 2: compile pybind wrapper..."
else "${CXX}" \
echo "[build] FAILED" -fPIC -O2 -std=c++17 \
exit 1 -I"${TORCH_INC}" \
fi -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)"

115
qwen3_6_scripts/cccl_moe_sort_scatter.cu
View File

@@ -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) // Split compilation: this file uses CCCL headers only (no torch).
// for histogram + prefix_sum + scatter. No device-level CUB API (conflicts // Pybind wrapper in cccl_moe_sort_scatter_pybind.cpp links against this.
// with corex's thrust/complex.h on CUDA 10.2).
// //
// Three kernels (same as corex_moe_index_combine but with CUB BlockRadixSort // Build pattern (same as cccl_allocator_preload.cu):
// for the scatter step): // clang++ -I cccl_preload/include -DCCCL_IGNORE_DEPRECATED_CUDA_BELOW_12
// 1. histogram — atomicAdd per expert // -DCUB_WRAPPED_NAMESPACE=cccl_moe ...
// 2. prefix_sum — CUB BlockScan ExclusiveSum
// 3. place — atomicAdd scatter into sorted positions // Suppress CUDA <12 check — corex 10.2 works for block-level CUB
// #define CCCL_IGNORE_DEPRECATED_CUDA_BELOW_12
// Build: torch.utils.cpp_extension.load(
// name="cccl_moe_sort_scatter", // Isolate from corex CUB
// sources=["cccl_moe_sort_scatter.cu"], #define CUB_WRAPPED_NAMESPACE cccl_moe
// extra_cuda_cflags=["-O3"],
// )
#include <torch/extension.h>
#include <c10/cuda/CUDAGuard.h>
#include <c10/cuda/CUDAStream.h>
#include <cub/block/block_scan.cuh> #include <cub/block/block_scan.cuh>
#include <cuda_runtime.h>
#include <cstdint>
// ======================================================================== // ========================================================================
// Block-level CUB kernels (proven on BI-V100 corex clang++) // Kernels
// Same pattern as corex_moe_index_combine.cu
// ======================================================================== // ========================================================================
constexpr int32_t kBlock = 256; static constexpr int32_t kBlock = 256;
__global__ void moe_histogram_kernel( __global__ void moe_histogram_kernel(
const int32_t* __restrict__ expert_id, const int32_t* __restrict__ expert_id,
@@ -45,9 +40,8 @@ __global__ void moe_histogram_kernel(
__global__ void moe_prefix_sum_kernel( __global__ void moe_prefix_sum_kernel(
const int32_t* __restrict__ expert_sizes, const int32_t* __restrict__ expert_sizes,
int32_t* __restrict__ expert_offsets, int32_t* __restrict__ expert_offsets,
int32_t num_experts, int32_t num_experts) {
int64_t* __restrict__ total_out) { using BlockScan = cccl_moe::cub::BlockScan<int32_t, 256>;
using BlockScan = cub::BlockScan<int32_t, 256>;
__shared__ typename BlockScan::TempStorage s_scan; __shared__ typename BlockScan::TempStorage s_scan;
int32_t val = (threadIdx.x < num_experts) ? expert_sizes[threadIdx.x] : 0; 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) { if (threadIdx.x < num_experts) {
expert_offsets[threadIdx.x] = offset; 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( __global__ void moe_place_kernel(
const int32_t* __restrict__ expert_id, 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__ dst_src,
int32_t* __restrict__ src_dst, int32_t* __restrict__ src_dst,
int64_t num_elements, int64_t num_elements,
@@ -82,57 +72,32 @@ __global__ void moe_place_kernel(
src_dst[flat_idx] = pos; src_dst[flat_idx] = pos;
} }
// ========================================================================
// C API — called from pybind wrapper
// ========================================================================
// Same proven 3-kernel approach as corex_moe_index_combine.cu but with extern "C" {
// an additional inverse-scatter output for full compatibility.
std::tuple<torch::Tensor, torch::Tensor, torch::Tensor>
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<int32_t>(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);
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; int64_t grid = (N + kBlock - 1) / kBlock;
moe_histogram_kernel<<<grid, kBlock, 0, stream>>>(expert_id, expert_sizes, N, E);
}
// Step 1: histogram void cccl_moe_launch_prefix_sum(
moe_histogram_kernel<<<grid, kBlock, 0, stream>>>( const int32_t* expert_sizes, int32_t* expert_offsets,
expert_id_i32.data_ptr<int32_t>(), int32_t E, cudaStream_t stream) {
expert_sizes.data_ptr<int32_t>(), moe_prefix_sum_kernel<<<1, kBlock, 0, stream>>>(expert_sizes, expert_offsets, E);
N, E); }
// Step 2: prefix sum (CUB BlockScan) void cccl_moe_launch_place(
moe_prefix_sum_kernel<<<1, kBlock, 0, stream>>>( const int32_t* expert_id, int32_t* expert_offsets,
expert_sizes.data_ptr<int32_t>(), int32_t* dst_src, int32_t* src_dst,
expert_offsets.data_ptr<int32_t>(), int64_t N, int32_t E, cudaStream_t stream) {
E, nullptr); int64_t grid = (N + kBlock - 1) / kBlock;
// Step 3: scatter — place each token into its sorted position
moe_place_kernel<<<grid, kBlock, 0, stream>>>( moe_place_kernel<<<grid, kBlock, 0, stream>>>(
expert_id_i32.data_ptr<int32_t>(), expert_id, expert_offsets, dst_src, src_dst, N, E);
expert_offsets.data_ptr<int32_t>(),
dst_src.data_ptr<int32_t>(),
src_dst.data_ptr<int32_t>(),
N, E);
return std::make_tuple(src_dst, dst_src, expert_sizes);
} }
} // extern "C"
// ========================================================================
// 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)");
}

62
qwen3_6_scripts/cccl_moe_sort_scatter_pybind.cpp Normal file
View File

@@ -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 <torch/extension.h>
#include <c10/cuda/CUDAStream.h>
#include <cuda_runtime.h>
// 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<torch::Tensor, torch::Tensor, torch::Tensor>
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<int32_t>(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<int32_t>(),
expert_sizes.data_ptr<int32_t>(),
N, E, stream);
cccl_moe_launch_prefix_sum(
expert_sizes.data_ptr<int32_t>(),
expert_offsets.data_ptr<int32_t>(),
E, stream);
cccl_moe_launch_place(
expert_id_i32.data_ptr<int32_t>(),
expert_offsets.data_ptr<int32_t>(),
dst_src.data_ptr<int32_t>(),
src_dst.data_ptr<int32_t>(),
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)");
}

3
qwen3_6_scripts/cccl_preload/include/cuda/std/__functional/invoke.h
View File

@@ -43,7 +43,8 @@ _CCCL_BEGIN_NAMESPACE_CUDA_STD
struct __any struct __any
{ {
_CCCL_API inline __any(...); template <class _T>
_CCCL_API inline __any(_T);
}; };
template <class _DecayedFp> template <class _DecayedFp>

2
qwen3_6_scripts/cccl_preload/include/cuda/std/__type_traits/conjunction.h
View File

@@ -35,7 +35,7 @@ template <class... _Pred>
_CCCL_HOST_DEVICE __expand_to_true<enable_if_t<_Pred::value>...> __and_helper(int); _CCCL_HOST_DEVICE __expand_to_true<enable_if_t<_Pred::value>...> __and_helper(int);
template <class...> template <class...>
_CCCL_HOST_DEVICE false_type __and_helper(...); _CCCL_HOST_DEVICE false_type __and_helper(long);
// _And always performs lazy evaluation of its arguments. // _And always performs lazy evaluation of its arguments.
// //

2
qwen3_6_scripts/cccl_preload/include/cuda/std/__type_traits/is_referenceable.h
View File

@@ -39,7 +39,7 @@ struct __cccl_is_referenceable_impl
template <class _Tp> template <class _Tp>
_CCCL_HOST_DEVICE static _Tp& __test(int); _CCCL_HOST_DEVICE static _Tp& __test(int);
template <class _Tp> template <class _Tp>
_CCCL_HOST_DEVICE static false_type __test(...); _CCCL_HOST_DEVICE static false_type __test(long);
}; };
template <class _Tp> template <class _Tp>