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.

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; }
+# 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
+[[ -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)"

qwen3_6_scripts/cccl_moe_sort_scatter.cu

@@ -1,201 +1,103 @@
-// cccl_moe_sort_scatter.cu — CUB DeviceRadixSort-based MoE token dispatch
+// cccl_moe_sort_scatter.cu — CCCL CUB device-level MoE token dispatch
 //
-// Replaces the 3-kernel (histogram + prefix_sum + place) approach with:
+// Split compilation: this file uses CCCL headers only (no torch).
-//   1. DeviceRadixSort::SortPairs — sort (expert_id, token_idx) pairs by expert_id
+// Pybind wrapper in cccl_moe_sort_scatter_pybind.cpp links against this.
-//   2. DeviceHistogram::HistogramEven — count tokens per expert
-//   3. DeviceScan::ExclusiveSum — prefix sum for expert offsets
 //
-// Uses CCCL upstream headers (in cccl_preload/include/) instead of corex CUB
+// Build pattern (same as cccl_allocator_preload.cu):
-// to avoid BI-V100 corex CUB bugs.
+//   clang++ -I cccl_preload/include -DCCCL_IGNORE_DEPRECATED_CUDA_BELOW_12
-//
+//           -DCUB_WRAPPED_NAMESPACE=cccl_moe ...
-// Build: torch.utils.cpp_extension.load(
-//   name="cccl_moe_sort_scatter",
-//   sources=["cccl_moe_sort_scatter.cu"],
-//   extra_include_paths=["cccl_preload/include"],
-//   extra_cuda_cflags=["-O3", "-DCUB_WRAPPED_NAMESPACE=cccl_moe"],
-// )
 
-#include <torch/extension.h>
+// Suppress CUDA <12 check — corex 10.2 works for block-level CUB
-#include <c10/cuda/CUDAGuard.h>
+#define CCCL_IGNORE_DEPRECATED_CUDA_BELOW_12
-#include <c10/cuda/CUDAStream.h>
 
-// Use CCCL CUB, not corex CUB
+// Isolate from corex CUB
 #define CUB_WRAPPED_NAMESPACE cccl_moe
-#include <cub/device/device_radix_sort.cuh>
+
-#include <cub/device/device_scan.cuh>
 #include <cub/block/block_scan.cuh>
+#include <cuda_runtime.h>
+#include <cstdint>
 
 // ========================================================================
-// moe_sort_scatter: sort tokens by expert_id using CUB DeviceRadixSort
+// Kernels
-//
-// Input:
-//   expert_id: [N] int32, each in [0, num_experts)
-//   num_experts: int
-//
-// Output:
-//   sorted_indices: [N] int32 — original token indices sorted by expert
-//   expert_offsets: [num_experts+1] int32 — exclusive prefix sum
-//   expert_sizes:   [num_experts] int32 — count per expert
 // ========================================================================
 
-// Small kernel to build expert_sizes from sorted keys via boundary detection
+static constexpr int32_t kBlock = 256;
-__global__ void compute_expert_boundaries(
+
-    const int32_t* __restrict__ sorted_keys,
+__global__ void moe_histogram_kernel(
-    int32_t* __restrict__ expert_offsets,  // [num_experts + 1]
+    const int32_t* __restrict__ expert_id,
-    int64_t N,
+    int32_t* __restrict__ expert_sizes,
+    int64_t num_elements,
     int32_t num_experts) {
-  // Initialize all to 0
+  int64_t tid = int64_t(blockIdx.x) * kBlock + threadIdx.x;
-  int tid = blockIdx.x * blockDim.x + threadIdx.x;
+  if (tid < num_elements) {
-
+    int32_t eid = expert_id[tid];
-  // First pass: detect boundaries
+    if (eid >= 0 && eid < num_experts) {
-  if (tid < N) {
+      atomicAdd(&expert_sizes[eid], 1);
-    int32_t cur = sorted_keys[tid];
-    if (tid == 0) {
-      // First element starts expert cur
-      expert_offsets[cur] = 0;
-    } else {
-      int32_t prev = sorted_keys[tid - 1];
-      if (cur != prev) {
-        expert_offsets[cur] = tid;
-      }
-    }
-    // Last element
-    if (tid == N - 1) {
-      expert_offsets[num_experts] = N;
     }
   }
 }
 
-// Fill gaps in expert_offsets (experts with 0 tokens)
+__global__ void moe_prefix_sum_kernel(
-__global__ void fill_offset_gaps(
+    const int32_t* __restrict__ expert_sizes,
     int32_t* __restrict__ expert_offsets,
     int32_t num_experts) {
-  // Backward fill: if expert_offsets[i] == -1, copy from next non-(-1)
+  using BlockScan = cccl_moe::cub::BlockScan<int32_t, 256>;
-  // Single thread is fine for num_experts <= 256
+  __shared__ typename BlockScan::TempStorage s_scan;
-  if (threadIdx.x != 0) return;
 
-  // Fill from the end
+  int32_t val = (threadIdx.x < num_experts) ? expert_sizes[threadIdx.x] : 0;
-  int32_t next_offset = expert_offsets[num_experts];  // = N
+  int32_t offset;
-  for (int i = num_experts - 1; i >= 0; --i) {
+  BlockScan(s_scan).ExclusiveSum(val, offset);
-    if (expert_offsets[i] == -1) {
+  __syncthreads();
-      expert_offsets[i] = next_offset;
+
-    } else {
+  if (threadIdx.x < num_experts) {
-      next_offset = expert_offsets[i];
+    expert_offsets[threadIdx.x] = offset;
-    }
   }
 }
 
-// Compute expert_sizes from expert_offsets
+__global__ void moe_place_kernel(
-__global__ void compute_expert_sizes(
+    const int32_t* __restrict__ expert_id,
-    const int32_t* __restrict__ expert_offsets,
+    int32_t* __restrict__ expert_offsets,
-    int32_t* __restrict__ expert_sizes,
+    int32_t* __restrict__ dst_src,
+    int32_t* __restrict__ src_dst,
+    int64_t num_elements,
     int32_t num_experts) {
-  int tid = blockIdx.x * blockDim.x + threadIdx.x;
+  int64_t flat_idx = int64_t(blockIdx.x) * kBlock + threadIdx.x;
-  if (tid < num_experts) {
+  if (flat_idx >= num_elements) return;
-    expert_sizes[tid] = expert_offsets[tid + 1] - expert_offsets[tid];
+
-  }
+  int32_t eid = expert_id[flat_idx];
+  if (eid < 0 || eid >= num_experts) return;
+
+  int32_t pos = atomicAdd(&expert_offsets[eid], 1);
+  dst_src[pos] = static_cast<int32_t>(flat_idx);
+  src_dst[flat_idx] = pos;
 }
 
-
-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 device = expert_id.device();
-  auto stream = at::cuda::getCurrentCUDAStream();
-  int64_t N = expert_id.numel();
-  int32_t E = static_cast<int32_t>(num_experts);
-
-  // Ensure int32
-  auto keys_in = expert_id.to(torch::kInt32).contiguous();
-  auto opt_i32 = keys_in.options();
-
-  // Create value array: [0, 1, 2, ..., N-1]
-  auto vals_in = torch::arange(N, opt_i32);
-
-  // Allocate output
-  auto sorted_keys = torch::empty({N}, opt_i32);
-  auto sorted_vals = torch::empty({N}, opt_i32);
-
-  // CUB DeviceRadixSort needs temp storage
-  // First query size
-  size_t temp_bytes = 0;
-  cccl_moe::cub::DeviceRadixSort::SortPairs(
-      nullptr, temp_bytes,
-      keys_in.data_ptr<int32_t>(),
-      sorted_keys.data_ptr<int32_t>(),
-      vals_in.data_ptr<int32_t>(),
-      sorted_vals.data_ptr<int32_t>(),
-      static_cast<int>(N),
-      0,                      // begin_bit
-      sizeof(int32_t) * 8,    // end_bit (all bits, but only need log2(E) bits)
-      stream);
-
-  // Allocate temp storage
-  auto temp_storage = torch::empty({static_cast<int64_t>(temp_bytes)},
-                                    torch::dtype(torch::kUInt8).device(device));
-
-  // Sort
-  cccl_moe::cub::DeviceRadixSort::SortPairs(
-      temp_storage.data_ptr(), temp_bytes,
-      keys_in.data_ptr<int32_t>(),
-      sorted_keys.data_ptr<int32_t>(),
-      vals_in.data_ptr<int32_t>(),
-      sorted_vals.data_ptr<int32_t>(),
-      static_cast<int>(N),
-      0,
-      sizeof(int32_t) * 8,
-      stream);
-
-  // Compute expert offsets via boundary detection
-  // Initialize to -1
-  auto expert_offsets = torch::full({num_experts + 1}, -1, opt_i32);
-
-  int block = 256;
-  int grid = (N + block - 1) / block;
-  compute_expert_boundaries<<<grid, block, 0, stream>>>(
-      sorted_keys.data_ptr<int32_t>(),
-      expert_offsets.data_ptr<int32_t>(),
-      N, E);
-
-  // Handle empty experts
-  fill_offset_gaps<<<1, 1, 0, stream>>>(
-      expert_offsets.data_ptr<int32_t>(), E);
-
-  // Compute sizes from offsets
-  auto expert_sizes = torch::empty({num_experts}, opt_i32);
-  int grid2 = (E + block - 1) / block;
-  compute_expert_sizes<<<grid2, block, 0, stream>>>(
-      expert_offsets.data_ptr<int32_t>(),
-      expert_sizes.data_ptr<int32_t>(),
-      E);
-
-  // sorted_vals = the original token indices, sorted by expert
-  // expert_sizes = tokens per expert
-  // sorted_keys not needed by caller, but sorted_vals is "dst_src"
-  //
-  // Build src_dst: inverse mapping
-  // src_dst[sorted_vals[i]] = i
-  auto src_dst = torch::empty({N}, opt_i32);
-
-  // Simple inverse scatter kernel
-  // For now use a tiny lambda — could be another kernel
-  // But actually we can do it with scatter:
-  // src_dst.scatter_(0, sorted_vals.long(), torch.arange(N))
-  // This is a single CUDA kernel internally
-
-  auto arange_n = torch::arange(N, opt_i32);
-  src_dst.scatter_(0, sorted_vals.to(torch::kInt64), arange_n);
-
-  return std::make_tuple(src_dst, sorted_vals, expert_sizes);
-}
-
-
 // ========================================================================
-// pybind
+// C API — called from pybind wrapper
 // ========================================================================
 
-PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+extern "C" {
-  m.def("moe_sort_scatter", &moe_sort_scatter,
+
-        "CUB DeviceRadixSort-based MoE token dispatch "
+void cccl_moe_launch_histogram(
-        "(sort expert_ids, compute offsets+sizes)");
+    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);
 }
+
+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);
+}
+
+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, expert_offsets, dst_src, src_dst, N, E);
+}
+
+} // extern "C"

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)");
+}

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>

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.
 //

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>