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

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.
2026-08-13 11:35:43 +00:00
parent 05706f0d60
commit f6cf9d662e
6 changed files with 168 additions and 142 deletions
--- 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)"
--- 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 <torch/extension.h>
-#include <c10/cuda/CUDAGuard.h>
-#include <c10/cuda/CUDAStream.h>
 #include <cub/block/block_scan.cuh>
+#include <cuda_runtime.h>
+#include <cstdint>

 // ========================================================================
-// 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, 256>;
+    int32_t num_experts) {
+  using BlockScan = cccl_moe::cub::BlockScan<int32_t, 256>;
  __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<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);
+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<<<grid, kBlock, 0, stream>>>(expert_id, expert_sizes, N, E);
+}

-  // Step 1: histogram
-  moe_histogram_kernel<<<grid, kBlock, 0, stream>>>(
-      expert_id_i32.data_ptr<int32_t>(),
-      expert_sizes.data_ptr<int32_t>(),
-      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<int32_t>(),
-      expert_offsets.data_ptr<int32_t>(),
-      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<<<grid, kBlock, 0, stream>>>(
-      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);
-
-  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"
--- a/qwen3_6_scripts/cccl_moe_sort_scatter_pybind.cpp
+++ 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 <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)");
+}
--- 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 <class _T>
+  _CCCL_API inline __any(_T);
 };

 template <class _DecayedFp>
--- 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 <class... _Pred>
 _CCCL_HOST_DEVICE __expand_to_true<enable_if_t<_Pred::value>...> __and_helper(int);

 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.
 //
--- 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 <class _Tp>
  _CCCL_HOST_DEVICE static _Tp& __test(int);
  template <class _Tp>
-  _CCCL_HOST_DEVICE static false_type __test(...);
+  _CCCL_HOST_DEVICE static false_type __test(long);
 };

 template <class _Tp>