fix(build): use block-level CUB only — device-level API conflicts with corex CUDA 10.2
CCCL latest requires CUDA 12+, corex is 10.2. Device-level CUB headers (DeviceRadixSort etc) pull in thrust/detail/type_traits.h which conflicts with corex's thrust/complex.h namespace. Rewrite to use block-level CUB BlockScan only (same pattern as the proven corex_moe_index_combine.cu): histogram + prefix_sum + scatter. No extra_include_paths needed — uses corex's built-in cub/block/block_scan.cuh.
This commit is contained in:
project6-dev
@@ -1,192 +1,129 @@
|
||||
// cccl_moe_sort_scatter.cu — CUB DeviceRadixSort-based MoE token dispatch
|
||||
// cccl_moe_sort_scatter.cu — Block-level CUB MoE token dispatch
|
||||
//
|
||||
// Replaces the 3-kernel (histogram + prefix_sum + place) approach with:
|
||||
// 1. DeviceRadixSort::SortPairs — sort (expert_id, token_idx) pairs by expert_id
|
||||
// 2. DeviceHistogram::HistogramEven — count tokens per expert
|
||||
// 3. DeviceScan::ExclusiveSum — prefix sum for expert offsets
|
||||
// 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).
|
||||
//
|
||||
// Uses CCCL upstream headers (in cccl_preload/include/) instead of corex CUB
|
||||
// to avoid BI-V100 corex CUB bugs.
|
||||
// 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_include_paths=["cccl_preload/include"],
|
||||
// extra_cuda_cflags=["-O3", "-DCUB_WRAPPED_NAMESPACE=cccl_moe"],
|
||||
// extra_cuda_cflags=["-O3"],
|
||||
// )
|
||||
|
||||
#include <torch/extension.h>
|
||||
#include <c10/cuda/CUDAGuard.h>
|
||||
#include <c10/cuda/CUDAStream.h>
|
||||
|
||||
// Use CCCL CUB, not 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>
|
||||
|
||||
// ========================================================================
|
||||
// moe_sort_scatter: sort tokens by expert_id using CUB DeviceRadixSort
|
||||
//
|
||||
// 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
|
||||
// Block-level CUB kernels (proven on BI-V100 corex clang++)
|
||||
// Same pattern as corex_moe_index_combine.cu
|
||||
// ========================================================================
|
||||
|
||||
// Small kernel to build expert_sizes from sorted keys via boundary detection
|
||||
__global__ void compute_expert_boundaries(
|
||||
const int32_t* __restrict__ sorted_keys,
|
||||
int32_t* __restrict__ expert_offsets, // [num_experts + 1]
|
||||
int64_t N,
|
||||
int32_t num_experts) {
|
||||
// Initialize all to 0
|
||||
int tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
constexpr int32_t kBlock = 256;
|
||||
|
||||
// First pass: detect boundaries
|
||||
if (tid < N) {
|
||||
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 fill_offset_gaps(
|
||||
int32_t* __restrict__ expert_offsets,
|
||||
int32_t num_experts) {
|
||||
// Backward fill: if expert_offsets[i] == -1, copy from next non-(-1)
|
||||
// Single thread is fine for num_experts <= 256
|
||||
if (threadIdx.x != 0) return;
|
||||
|
||||
// Fill from the end
|
||||
int32_t next_offset = expert_offsets[num_experts]; // = N
|
||||
for (int i = num_experts - 1; i >= 0; --i) {
|
||||
if (expert_offsets[i] == -1) {
|
||||
expert_offsets[i] = next_offset;
|
||||
} else {
|
||||
next_offset = expert_offsets[i];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Compute expert_sizes from expert_offsets
|
||||
__global__ void compute_expert_sizes(
|
||||
const int32_t* __restrict__ expert_offsets,
|
||||
__global__ void moe_histogram_kernel(
|
||||
const int32_t* __restrict__ expert_id,
|
||||
int32_t* __restrict__ expert_sizes,
|
||||
int64_t num_elements,
|
||||
int32_t num_experts) {
|
||||
int tid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (tid < num_experts) {
|
||||
expert_sizes[tid] = expert_offsets[tid + 1] - expert_offsets[tid];
|
||||
int64_t tid = int64_t(blockIdx.x) * kBlock + threadIdx.x;
|
||||
if (tid < num_elements) {
|
||||
int32_t eid = expert_id[tid];
|
||||
if (eid >= 0 && eid < num_experts) {
|
||||
atomicAdd(&expert_sizes[eid], 1);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
__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>;
|
||||
__shared__ typename BlockScan::TempStorage s_scan;
|
||||
|
||||
int32_t val = (threadIdx.x < num_experts) ? expert_sizes[threadIdx.x] : 0;
|
||||
int32_t offset;
|
||||
BlockScan(s_scan).ExclusiveSum(val, offset);
|
||||
__syncthreads();
|
||||
|
||||
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__ dst_src,
|
||||
int32_t* __restrict__ src_dst,
|
||||
int64_t num_elements,
|
||||
int32_t num_experts) {
|
||||
int64_t flat_idx = int64_t(blockIdx.x) * kBlock + threadIdx.x;
|
||||
if (flat_idx >= num_elements) return;
|
||||
|
||||
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;
|
||||
}
|
||||
|
||||
|
||||
// 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 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();
|
||||
auto expert_id_i32 = expert_id.to(torch::kInt32).contiguous();
|
||||
auto opt_i32 = expert_id_i32.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 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);
|
||||
|
||||
// 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
|
||||
int64_t grid = (N + kBlock - 1) / kBlock;
|
||||
|
||||
auto arange_n = torch::arange(N, opt_i32);
|
||||
src_dst.scatter_(0, sorted_vals.to(torch::kInt64), arange_n);
|
||||
// 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);
|
||||
|
||||
return std::make_tuple(src_dst, sorted_vals, expert_sizes);
|
||||
// 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
|
||||
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);
|
||||
}
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user