feat: implement 5 missing MoE ops — topk_softmax + token_index + expand + group_gemm + combine
Symbol dump from real device confirms: libixformer.so has 0 MoE symbols. topk_softmax, moe_compute_token_index_api, moe_expand_input, moe_w16a16_group_gemm, moe_output_reduce_sum — all missing. Non-MoE symbols (silu_and_mul, rms_norm, flash_attn, reshape_and_cache, rotary_embedding) are present and working. Implementation strategy — use available primitives: - topk_softmax: pure CUDA kernel (64-expert, shared-mem argmax) - moe_compute_token_index: histogram + prefix_sum + scatter (3 kernels) - moe_expand_input: gather kernel - moe_w16a16_group_gemm: per-expert loop calling cuinferCustomGemm (confirmed in libcuinfer.so symbol dump: cuinferCustomGemm exists) - moe_output_reduce_sum: weighted combine kernel All in ixformer::infer namespace so ix_full_bridge_v2.cpp links directly. Compile: nvcc moe_ops_impl.cu + ix_full_bridge_v2.cpp → single .so
This commit is contained in:
project_6
489
ex_engine/csrc/moe_ops_impl.cu
Normal file
489
ex_engine/csrc/moe_ops_impl.cu
Normal file
@@ -0,0 +1,489 @@
|
||||
// moe_ops_impl.cu — Implement the 5 missing MoE functions
|
||||
//
|
||||
// These functions are declared in ixformer.h (from xllm upstream)
|
||||
// but NOT present in the base image's libixformer.so.
|
||||
//
|
||||
// We implement them using available primitives:
|
||||
// - cuinferCustomGemm (from libcuinfer.so) for group_gemm
|
||||
// - Pure CUDA kernels for topk_softmax, moe_compute_index, expand, combine
|
||||
// - ixformer::functions::cuinfer_gemm (from libixformer.so) as fallback
|
||||
//
|
||||
// Reference AST chain:
|
||||
// xllm/core/kernels/ilu/fused_moe.cpp → calls these 5 functions
|
||||
// xllm/core/kernels/ilu/group_gemm.cpp → calls moe_w16a16_group_gemm
|
||||
// xllm/core/kernels/ilu/ixformer.h → declares them in ixformer::infer
|
||||
//
|
||||
// We provide them in the SAME namespace so ix_full_bridge_v2.cpp links cleanly.
|
||||
|
||||
#include <torch/extension.h>
|
||||
#include <c10/cuda/CUDAStream.h>
|
||||
#include <c10/cuda/CUDAGuard.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <cuda_fp16.h>
|
||||
#include <optional>
|
||||
#include <vector>
|
||||
#include <numeric>
|
||||
|
||||
// ============================================================================
|
||||
// Forward-declare cuinfer C API (from libcuinfer.so, confirmed in symbol dump)
|
||||
// ============================================================================
|
||||
extern "C" {
|
||||
|
||||
typedef struct cuinferContext* cuinferHandle_t;
|
||||
typedef enum { CUINFER_STATUS_SUCCESS = 0 } cuinferStatus_t;
|
||||
typedef enum {
|
||||
CUINFER_OP_TENSOR_OP_N = 0,
|
||||
CUINFER_OP_TENSOR_OP_T = 1,
|
||||
} cuinferOperation_t;
|
||||
typedef enum {
|
||||
CUINFER_GEMM_DEFAULT = 0,
|
||||
} cuinferGEMMCustomOption_t;
|
||||
typedef enum {
|
||||
CUINFER_POINTER_MODE_HOST = 0,
|
||||
} cuinferPointerMode_t;
|
||||
|
||||
cuinferStatus_t cuinferCreate(cuinferHandle_t* handle);
|
||||
cuinferStatus_t cuinferDestroy(cuinferHandle_t handle);
|
||||
cuinferStatus_t cuinferSetStream(cuinferHandle_t handle, cudaStream_t stream);
|
||||
|
||||
cuinferStatus_t cuinferCustomGemm(
|
||||
cuinferHandle_t handle, cudaStream_t stream,
|
||||
cuinferPointerMode_t ptrMode,
|
||||
cuinferOperation_t transa, cuinferOperation_t transb,
|
||||
int m, int n, int k,
|
||||
const void* alpha,
|
||||
const void* A, cudaDataType_t Atype, int lda, long long int strideA,
|
||||
const void* B, cudaDataType_t Btype, int ldb, long long int strideB,
|
||||
const void* beta,
|
||||
void* C, cudaDataType_t Ctype, int ldc, long long int strideC,
|
||||
int batchCount,
|
||||
cudaDataType_t computeType, cudaDataType_t scaleType,
|
||||
const void* customHostPtr, const void* customDevicePtr,
|
||||
cuinferGEMMCustomOption_t customOption);
|
||||
|
||||
} // extern "C"
|
||||
|
||||
|
||||
// ============================================================================
|
||||
// Kernel 1: topk_softmax
|
||||
// Adapted from moe_topk_softmax_v3.cu (already working, 64-expert specialized)
|
||||
// ============================================================================
|
||||
|
||||
static constexpr int MOE_EXPERTS = 64;
|
||||
static constexpr int MOE_BLOCK = 64;
|
||||
|
||||
__device__ float smem_reduce_max(float val, float* smem) {
|
||||
int tid = threadIdx.x;
|
||||
smem[tid] = val;
|
||||
__syncthreads();
|
||||
for (int s = MOE_BLOCK / 2; s > 0; s >>= 1) {
|
||||
if (tid < s) smem[tid] = fmaxf(smem[tid], smem[tid + s]);
|
||||
__syncthreads();
|
||||
}
|
||||
return smem[0];
|
||||
}
|
||||
|
||||
__device__ float smem_reduce_sum(float val, float* smem) {
|
||||
int tid = threadIdx.x;
|
||||
smem[tid] = val;
|
||||
__syncthreads();
|
||||
for (int s = MOE_BLOCK / 2; s > 0; s >>= 1) {
|
||||
if (tid < s) smem[tid] += smem[tid + s];
|
||||
__syncthreads();
|
||||
}
|
||||
return smem[0];
|
||||
}
|
||||
|
||||
__device__ void smem_argmax(float val, int idx, float* s_val, int* s_idx) {
|
||||
int tid = threadIdx.x;
|
||||
s_val[tid] = val;
|
||||
s_idx[tid] = idx;
|
||||
__syncthreads();
|
||||
for (int s = MOE_BLOCK / 2; s > 0; s >>= 1) {
|
||||
if (tid < s && s_val[tid + s] > s_val[tid]) {
|
||||
s_val[tid] = s_val[tid + s];
|
||||
s_idx[tid] = s_idx[tid + s];
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
}
|
||||
|
||||
__global__ void topk_softmax_kernel(
|
||||
const float* __restrict__ input,
|
||||
float* __restrict__ topk_weights,
|
||||
int32_t* __restrict__ topk_indices,
|
||||
int32_t* __restrict__ token_expert_indices,
|
||||
int num_tokens, int topk, bool renormalize
|
||||
) {
|
||||
int row = blockIdx.x;
|
||||
if (row >= num_tokens) return;
|
||||
int tid = threadIdx.x;
|
||||
|
||||
__shared__ float smem[MOE_BLOCK];
|
||||
__shared__ int smem_idx[MOE_BLOCK];
|
||||
|
||||
float val = (tid < MOE_EXPERTS) ? input[row * MOE_EXPERTS + tid] : -1e30f;
|
||||
|
||||
// Softmax
|
||||
float row_max = smem_reduce_max(val, smem);
|
||||
val = (tid < MOE_EXPERTS) ? expf(val - row_max) : 0.0f;
|
||||
float row_sum = smem_reduce_sum(val, smem);
|
||||
val *= (1.0f / row_sum);
|
||||
|
||||
float* out_w = topk_weights + row * topk;
|
||||
int32_t* out_idx = topk_indices + row * topk;
|
||||
int32_t* out_src = token_expert_indices + row * topk;
|
||||
|
||||
float my_val = val;
|
||||
float topk_sum = 0.0f;
|
||||
|
||||
for (int ki = 0; ki < topk; ki++) {
|
||||
smem_argmax(my_val, tid, smem, smem_idx);
|
||||
float winner_val = smem[0];
|
||||
int winner_idx = smem_idx[0];
|
||||
__syncthreads();
|
||||
|
||||
if (tid == 0) {
|
||||
out_w[ki] = winner_val;
|
||||
out_idx[ki] = winner_idx;
|
||||
out_src[ki] = row;
|
||||
}
|
||||
topk_sum += winner_val;
|
||||
if (tid == winner_idx) my_val = -1.0f;
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
if (renormalize && tid == 0) {
|
||||
float inv = 1.0f / (topk_sum + 1e-8f);
|
||||
for (int ki = 0; ki < topk; ki++)
|
||||
out_w[ki] *= inv;
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
// ============================================================================
|
||||
// Kernel 2: moe_compute_token_index
|
||||
// Histogram + prefix sum + scatter — from xllm_kernels/cuda/moe_compute_index.cu
|
||||
// ============================================================================
|
||||
|
||||
__global__ void histogram_kernel(
|
||||
const int32_t* __restrict__ expert_ids,
|
||||
int32_t* __restrict__ expert_sizes,
|
||||
int num_elements, int num_experts
|
||||
) {
|
||||
int idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (idx < num_elements) {
|
||||
int eid = expert_ids[idx];
|
||||
if (eid >= 0 && eid < num_experts) {
|
||||
atomicAdd(&expert_sizes[eid], 1);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
__global__ void place_indices_kernel(
|
||||
const int32_t* __restrict__ expert_ids,
|
||||
int32_t* __restrict__ expert_offsets, // will be atomicAdd'd
|
||||
int32_t* __restrict__ src_dst,
|
||||
int32_t* __restrict__ dst_src,
|
||||
int num_elements
|
||||
) {
|
||||
int idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (idx < num_elements) {
|
||||
int eid = expert_ids[idx];
|
||||
int pos = atomicAdd(&expert_offsets[eid], 1);
|
||||
src_dst[idx] = pos; // where token idx goes in sorted order
|
||||
dst_src[pos] = idx; // reverse mapping
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
// ============================================================================
|
||||
// Kernel 3: moe_expand_input
|
||||
// Gather-based expand: output[i] = input[gather_index[i]]
|
||||
// ============================================================================
|
||||
|
||||
template <typename scalar_t>
|
||||
__global__ void expand_input_kernel(
|
||||
scalar_t* __restrict__ output,
|
||||
const scalar_t* __restrict__ input,
|
||||
const int32_t* __restrict__ dst_to_src,
|
||||
int num_output_tokens, int hidden_size
|
||||
) {
|
||||
int token = blockIdx.x;
|
||||
if (token >= num_output_tokens) return;
|
||||
|
||||
int src_token = dst_to_src[token];
|
||||
const scalar_t* src = input + (int64_t)src_token * hidden_size;
|
||||
scalar_t* dst = output + (int64_t)token * hidden_size;
|
||||
|
||||
for (int h = threadIdx.x; h < hidden_size; h += blockDim.x) {
|
||||
dst[h] = src[h];
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
// ============================================================================
|
||||
// Kernel 4: moe_combine_result (weighted sum of expert outputs)
|
||||
// output[t] = sum_k( weight[t][k] * gemm2_output[flat_index(t,k)] )
|
||||
// ============================================================================
|
||||
|
||||
template <typename scalar_t>
|
||||
__global__ void combine_result_kernel(
|
||||
scalar_t* __restrict__ output, // [N, H]
|
||||
const scalar_t* __restrict__ input, // [N*topk, H]
|
||||
const float* __restrict__ weights, // [N, topk]
|
||||
int num_tokens, int topk, int hidden_size
|
||||
) {
|
||||
int token = blockIdx.x;
|
||||
if (token >= num_tokens) return;
|
||||
|
||||
for (int h = threadIdx.x; h < hidden_size; h += blockDim.x) {
|
||||
float acc = 0.0f;
|
||||
for (int k = 0; k < topk; k++) {
|
||||
int flat = token * topk + k;
|
||||
float w = weights[token * topk + k];
|
||||
acc += w * __half2float(input[flat * hidden_size + h]);
|
||||
}
|
||||
output[token * hidden_size + h] = __float2half(acc);
|
||||
}
|
||||
}
|
||||
|
||||
// Float specialization
|
||||
template <>
|
||||
__global__ void combine_result_kernel<float>(
|
||||
float* __restrict__ output,
|
||||
const float* __restrict__ input,
|
||||
const float* __restrict__ weights,
|
||||
int num_tokens, int topk, int hidden_size
|
||||
) {
|
||||
int token = blockIdx.x;
|
||||
if (token >= num_tokens) return;
|
||||
|
||||
for (int h = threadIdx.x; h < hidden_size; h += blockDim.x) {
|
||||
float acc = 0.0f;
|
||||
for (int k = 0; k < topk; k++) {
|
||||
int flat = token * topk + k;
|
||||
float w = weights[token * topk + k];
|
||||
acc += w * input[flat * hidden_size + h];
|
||||
}
|
||||
output[token * hidden_size + h] = acc;
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
// ============================================================================
|
||||
// C++ wrapper functions — ixformer::infer namespace
|
||||
// These provide the MISSING symbols that ix_full_bridge_v2.cpp needs.
|
||||
// ============================================================================
|
||||
|
||||
namespace ixformer { namespace infer {
|
||||
|
||||
void topk_softmax(
|
||||
torch::Tensor& topk_weights,
|
||||
torch::Tensor& topk_indices,
|
||||
torch::Tensor& token_expert_indices,
|
||||
torch::Tensor& gating_output,
|
||||
bool renormalize
|
||||
) {
|
||||
int num_tokens = gating_output.size(0);
|
||||
int topk = topk_weights.size(1);
|
||||
auto stream = c10::cuda::getCurrentCUDAStream();
|
||||
|
||||
auto input_f32 = gating_output.to(torch::kFloat32).contiguous();
|
||||
|
||||
topk_softmax_kernel<<<num_tokens, MOE_BLOCK, 0, stream>>>(
|
||||
input_f32.data_ptr<float>(),
|
||||
topk_weights.data_ptr<float>(),
|
||||
topk_indices.data_ptr<int32_t>(),
|
||||
token_expert_indices.data_ptr<int32_t>(),
|
||||
num_tokens, topk, renormalize);
|
||||
}
|
||||
|
||||
void moe_compute_token_index_api(
|
||||
torch::Tensor& topk_ids,
|
||||
torch::Tensor& src_dst,
|
||||
torch::Tensor& dst_src,
|
||||
torch::Tensor& expert_sizes_gpu,
|
||||
const c10::optional<torch::Tensor>& expert_mask,
|
||||
const c10::optional<torch::Tensor>& expert_sizes_cpu,
|
||||
const c10::optional<torch::Tensor>& expand_tokens_gpu,
|
||||
int64_t start_expert_id,
|
||||
int64_t end_expert_id,
|
||||
int64_t num_experts
|
||||
) {
|
||||
auto stream = c10::cuda::getCurrentCUDAStream();
|
||||
int num_elements = topk_ids.numel();
|
||||
|
||||
// Zero expert_sizes
|
||||
cudaMemsetAsync(expert_sizes_gpu.data_ptr<int32_t>(), 0,
|
||||
num_experts * sizeof(int32_t), stream);
|
||||
|
||||
// Phase 1: histogram
|
||||
int blocks1 = (num_elements + 255) / 256;
|
||||
histogram_kernel<<<blocks1, 256, 0, stream>>>(
|
||||
topk_ids.data_ptr<int32_t>(),
|
||||
expert_sizes_gpu.data_ptr<int32_t>(),
|
||||
num_elements, num_experts);
|
||||
|
||||
// Phase 2: prefix sum for offsets (exclusive scan on GPU)
|
||||
// Use a separate buffer for offsets, then reset for place_indices
|
||||
auto expert_offsets = torch::zeros({num_experts}, topk_ids.options().dtype(torch::kInt32));
|
||||
// Copy sizes → do exclusive scan on CPU (small: 64 experts)
|
||||
auto sizes_cpu = expert_sizes_gpu.to(torch::kCPU);
|
||||
auto offsets_cpu = torch::zeros({num_experts}, torch::dtype(torch::kInt32));
|
||||
int32_t* s = sizes_cpu.data_ptr<int32_t>();
|
||||
int32_t* o = offsets_cpu.data_ptr<int32_t>();
|
||||
int32_t running = 0;
|
||||
for (int i = 0; i < num_experts; i++) {
|
||||
o[i] = running;
|
||||
running += s[i];
|
||||
}
|
||||
expert_offsets = offsets_cpu.to(topk_ids.device());
|
||||
|
||||
// Phase 3: place indices
|
||||
int blocks3 = (num_elements + 255) / 256;
|
||||
place_indices_kernel<<<blocks3, 256, 0, stream>>>(
|
||||
topk_ids.data_ptr<int32_t>(),
|
||||
expert_offsets.data_ptr<int32_t>(),
|
||||
src_dst.data_ptr<int32_t>(),
|
||||
dst_src.data_ptr<int32_t>(),
|
||||
num_elements);
|
||||
}
|
||||
|
||||
void moe_expand_input(
|
||||
torch::Tensor outputs,
|
||||
torch::Tensor inputs,
|
||||
torch::Tensor dst_to_src,
|
||||
const c10::optional<torch::Tensor>& src_to_dst,
|
||||
int64_t dst_tokens,
|
||||
int64_t expand_factor
|
||||
) {
|
||||
auto stream = c10::cuda::getCurrentCUDAStream();
|
||||
int hidden_size = inputs.size(1);
|
||||
int block = std::min(hidden_size, 256);
|
||||
|
||||
AT_DISPATCH_FLOATING_TYPES_AND_HALF(inputs.scalar_type(), "expand_input", [&] {
|
||||
expand_input_kernel<scalar_t><<<dst_tokens, block, 0, stream>>>(
|
||||
outputs.data_ptr<scalar_t>(),
|
||||
inputs.data_ptr<scalar_t>(),
|
||||
dst_to_src.data_ptr<int32_t>(),
|
||||
dst_tokens, hidden_size);
|
||||
});
|
||||
}
|
||||
|
||||
void moe_w16a16_group_gemm(
|
||||
torch::Tensor output,
|
||||
torch::Tensor inputs,
|
||||
torch::Tensor weights,
|
||||
torch::Tensor tokens_per_experts,
|
||||
const c10::optional<torch::Tensor>& dst_to_src,
|
||||
const c10::optional<torch::Tensor>& bias,
|
||||
std::string format,
|
||||
int64_t persistent,
|
||||
int64_t output_n
|
||||
) {
|
||||
// Implementation: loop over experts, call cuinferCustomGemm for each
|
||||
// weights: [num_experts, N, K] with format "TN" means transB
|
||||
// For each expert e with count tokens:
|
||||
// A = inputs[offset:offset+count, :] (count × K, row-major)
|
||||
// B = weights[e, :, :] (N × K, needs transB)
|
||||
// C = output[offset:offset+count, :] (count × N, row-major)
|
||||
// GEMM: C = A × B^T → (count, K) × (K, N) = (count, N)
|
||||
|
||||
auto stream = c10::cuda::getCurrentCUDAStream();
|
||||
int num_experts = weights.size(0);
|
||||
int N = weights.size(1); // output dim
|
||||
int K = weights.size(2); // input dim
|
||||
|
||||
// Get token counts on CPU
|
||||
auto counts_cpu = tokens_per_experts.to(torch::kCPU).to(torch::kInt32);
|
||||
int32_t* counts = counts_cpu.data_ptr<int32_t>();
|
||||
|
||||
// Create cuinfer handle
|
||||
cuinferHandle_t handle;
|
||||
cuinferCreate(&handle);
|
||||
cuinferSetStream(handle, stream);
|
||||
|
||||
float alpha = 1.0f, beta = 0.0f;
|
||||
|
||||
int offset = 0;
|
||||
for (int e = 0; e < num_experts; e++) {
|
||||
int M = counts[e];
|
||||
if (M <= 0) continue;
|
||||
|
||||
// A: inputs[offset : offset+M, :] → M × K
|
||||
// B: weights[e, :, :] → N × K (transposed: compute A × B^T)
|
||||
// C: output[offset : offset+M, :] → M × N
|
||||
const void* A_ptr = (const char*)inputs.data_ptr() +
|
||||
(int64_t)offset * K * inputs.element_size();
|
||||
const void* B_ptr = (const char*)weights.data_ptr() +
|
||||
(int64_t)e * N * K * weights.element_size();
|
||||
void* C_ptr = (char*)output.data_ptr() +
|
||||
(int64_t)offset * N * output.element_size();
|
||||
|
||||
cudaDataType_t dtype = (inputs.scalar_type() == torch::kFloat16)
|
||||
? CUDA_R_16F : CUDA_R_32F;
|
||||
|
||||
// cuinferCustomGemm: row-major convention
|
||||
// We want C = A × B^T
|
||||
// In cuinfer (column-major internally): transa=N, transb=T
|
||||
// M_gemm = M (rows of C), N_gemm = N (cols of C), K_gemm = K
|
||||
cuinferCustomGemm(
|
||||
handle, stream,
|
||||
CUINFER_POINTER_MODE_HOST,
|
||||
CUINFER_OP_TENSOR_OP_N, // transa = no transpose
|
||||
CUINFER_OP_TENSOR_OP_T, // transb = transpose (TN format)
|
||||
M, N, K,
|
||||
&alpha,
|
||||
A_ptr, dtype, K, 0, // lda=K for row-major A
|
||||
B_ptr, dtype, K, 0, // ldb=K for row-major B (will be transposed)
|
||||
&beta,
|
||||
C_ptr, dtype, N, 0, // ldc=N for row-major C
|
||||
1, // batchCount=1
|
||||
CUDA_R_32F, // computeType
|
||||
CUDA_R_32F, // scaleType
|
||||
nullptr, nullptr, // custom pointers
|
||||
CUINFER_GEMM_DEFAULT);
|
||||
|
||||
offset += M;
|
||||
}
|
||||
|
||||
cuinferDestroy(handle);
|
||||
}
|
||||
|
||||
void moe_output_reduce_sum(
|
||||
torch::Tensor outputs,
|
||||
torch::Tensor inputs,
|
||||
const c10::optional<torch::Tensor>& mul_weight,
|
||||
const c10::optional<torch::Tensor>& mask,
|
||||
const c10::optional<torch::Tensor>& extra_residual,
|
||||
double scaling_factor
|
||||
) {
|
||||
// inputs: [N, topk, H] — expert outputs per token
|
||||
// mul_weight: [N, topk] — router weights
|
||||
// outputs: [N, H] — weighted sum
|
||||
auto stream = c10::cuda::getCurrentCUDAStream();
|
||||
int num_tokens = inputs.size(0);
|
||||
int topk = inputs.size(1);
|
||||
int hidden_size = inputs.size(2);
|
||||
int block = std::min(hidden_size, 256);
|
||||
|
||||
// Reshape inputs to [N*topk, H] for the kernel
|
||||
auto input_flat = inputs.reshape({num_tokens * topk, hidden_size});
|
||||
|
||||
if (inputs.scalar_type() == torch::kFloat16) {
|
||||
combine_result_kernel<__half><<<num_tokens, block, 0, stream>>>(
|
||||
reinterpret_cast<__half*>(outputs.data_ptr()),
|
||||
reinterpret_cast<const __half*>(input_flat.data_ptr()),
|
||||
mul_weight.value().data_ptr<float>(),
|
||||
num_tokens, topk, hidden_size);
|
||||
} else {
|
||||
combine_result_kernel<float><<<num_tokens, block, 0, stream>>>(
|
||||
outputs.data_ptr<float>(),
|
||||
input_flat.data_ptr<float>(),
|
||||
mul_weight.value().data_ptr<float>(),
|
||||
num_tokens, topk, hidden_size);
|
||||
}
|
||||
}
|
||||
|
||||
}} // namespace ixformer::infer
|
||||
Reference in New Issue
Block a user