diff --git a/qwen3_6_scripts/cccl_moe_sort_scatter.cu b/qwen3_6_scripts/cccl_moe_sort_scatter.cu index dd57d92a..87c50790 100644 --- a/qwen3_6_scripts/cccl_moe_sort_scatter.cu +++ b/qwen3_6_scripts/cccl_moe_sort_scatter.cu @@ -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 #include #include - -// Use CCCL CUB, not corex CUB -#define CUB_WRAPPED_NAMESPACE cccl_moe -#include -#include #include // ======================================================================== -// 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; + __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(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 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(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(), - sorted_keys.data_ptr(), - vals_in.data_ptr(), - sorted_vals.data_ptr(), - static_cast(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(temp_bytes)}, - torch::dtype(torch::kUInt8).device(device)); - - // Sort - cccl_moe::cub::DeviceRadixSort::SortPairs( - temp_storage.data_ptr(), temp_bytes, - keys_in.data_ptr(), - sorted_keys.data_ptr(), - vals_in.data_ptr(), - sorted_vals.data_ptr(), - static_cast(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<<>>( - sorted_keys.data_ptr(), - expert_offsets.data_ptr(), - N, E); - - // Handle empty experts - fill_offset_gaps<<<1, 1, 0, stream>>>( - expert_offsets.data_ptr(), E); - - // Compute sizes from offsets - auto expert_sizes = torch::empty({num_experts}, opt_i32); - int grid2 = (E + block - 1) / block; - compute_expert_sizes<<>>( - expert_offsets.data_ptr(), - expert_sizes.data_ptr(), - 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<<>>( + expert_id_i32.data_ptr(), + expert_sizes.data_ptr(), + 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(), + expert_offsets.data_ptr(), + E, nullptr); + + // Step 3: scatter — place each token into its sorted position + 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); }