dylanyunlong/project_6
1
0
Fork 0
Code Issues Pull Requests Actions Projects Releases Wiki Activity

ref(EX): import upstream ILU kernels + xllm MoE CUDA sources into ex_engine

Browse Source 
Copied from upstream_ref (NOT rewritten — exact upstream code):

ixformer C++ API (the authoritative header):
  include/ixformer.h — ixformer::infer namespace: topk_softmax,
    moe_compute_token_index_api, moe_w16a16_group_gemm, moe_expand_input,
    moe_output_reduce_sum, silu_and_mul, rms_norm, xllm_paged_attention, etc.
  include/ilu_ops_api.h — xllm::kernel::ilu namespace: moe_active_topk,
    moe_gen_idx, moe_expand_input, group_gemm, moe_combine_result,
    batch_prefill, batch_decode, rms_norm, matmul, act_and_mul, etc.

ILU kernel wrappers (call ixformer::infer directly):
  csrc/ilu_kernel_fused_moe.cpp — topk routing + gen_idx + expand + combine
  csrc/ilu_kernel_group_gemm.cpp — batched expert GEMM
  csrc/ilu_kernel_{activation,norm,rope,matmul,attention}.cpp

ILU layer implementations (full pipeline):
  csrc/ilu_layer_fused_moe.{cpp,h} — 797 lines, the complete MoE pipeline
    that competitor 168 ran as corex_moe.py
  csrc/ilu_layer_attention.{cpp,h} — prefill/decode attention dispatch

CUDA MoE kernels (from xllm + ds_vllm):
  csrc/moe/moe_topk_softmax_kernels.cuh — CUB BlockReduce + warp topk
  csrc/moe/moe_topk_sigmoid_kernels.cuh — sigmoid scoring variant
  csrc/moe/moe_topk.cuh + moe_fused_topk.cu — entry points
  csrc/moe/moeTopKFuncs.cuh — TRT-LLM derived vllm-compatible topk
  csrc/moe/moe_ops.h + moe_align_sum_kernels.cu — alignment kernels

Common layer headers:
  csrc/common_fused_moe{,_base}.h + common_moe_fused_topk.{cpp,h}
This commit is contained in:
project6-dev
2026-08-10 03:59:37 +00:00
parent dba027fded
commit f4e2264a83
26 changed files with 4522 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files

40
ex_engine/csrc/moe/moe_topk_softmax_kernels.cuh
View File

@@ -22,11 +22,9 @@ limitations under the License.
#include <torch/all.h>
#include <cub/util_type.cuh>
#if CUDA_VERSION >= 12090
#include <cuda/functional>
#endif
#include "device_utils.cuh"
#include "kernels/cuda/device_utils.cuh"
using cub_kvp = cub::KeyValuePair<int, float>;
@@ -707,7 +705,8 @@ void topk_gating_softmax_kernel_launcher(const T* gating_output,
LAUNCH_SOFTMAX(T, 256, WARPS_PER_TB);
break;
default: {
TORCH_CHECK(softmax_workspace != nullptr, "softmax_workspace must be provided for num_experts that are ");
CHECK(softmax_workspace != nullptr)
<< "softmax_workspace must be provided for num_experts that are "
"not a power of 2.";
static constexpr int TPB = 256;
moe_softmax<T, TPB><<<num_tokens, TPB, 0, stream>>>(gating_output,
@@ -753,23 +752,29 @@ void topk_softmax(torch::Tensor& topk_weights, // [num_tokens, topk]
const double moe_softcapping,
const std::optional<torch::Tensor>& correction_bias) {
// Check data type
TORCH_CHECK(gating_output.scalar_type() == at::ScalarType::Float ||
CHECK(gating_output.scalar_type() == at::ScalarType::Float ||
gating_output.scalar_type() == at::ScalarType::Half ||
gating_output.scalar_type() == at::ScalarType::BFloat16,
"gating_output must be float32, float16, or bfloat16");
gating_output.scalar_type() == at::ScalarType::BFloat16)
<< "gating_output must be float32, float16, or bfloat16";
// Check dimensions
TORCH_CHECK(gating_output.dim() == 2, "gating_output must be 2D tensor [num_tokens, num_experts]");
TORCH_CHECK(topk_weights.dim() == 2, "topk_weights must be 2D tensor [num_tokens, topk]");
TORCH_CHECK(topk_indices.dim() == 2, "topk_indices must be 2D tensor [num_tokens, topk]");
CHECK(gating_output.dim() == 2)
<< "gating_output must be 2D tensor [num_tokens, num_experts]";
CHECK(topk_weights.dim() == 2)
<< "topk_weights must be 2D tensor [num_tokens, topk]";
CHECK(topk_indices.dim() == 2)
<< "topk_indices must be 2D tensor [num_tokens, topk]";
// Check shapes
TORCH_CHECK(gating_output.size(0) == topk_weights.size(0), "First dimension of topk_weights must match num_tokens in ");
CHECK(gating_output.size(0) == topk_weights.size(0))
<< "First dimension of topk_weights must match num_tokens in "
"gating_output"
<< "First dimension of topk_indices must match num_tokens in "
"gating_output";
TORCH_CHECK(topk_weights.size(-1) == topk_indices.size(-1), "Second dimension of topk_indices must match topk in topk_weights topk must be less than or equal to num_experts");
CHECK(topk_weights.size(-1) == topk_indices.size(-1))
<< "Second dimension of topk_indices must match topk in topk_weights"
<< "topk must be less than or equal to num_experts";
const int num_experts = static_cast<int>(gating_output.size(-1));
const int num_tokens = static_cast<int>(gating_output.size(0));
@@ -791,9 +796,12 @@ void topk_softmax(torch::Tensor& topk_weights, // [num_tokens, topk]
const float* bias_ptr = nullptr;
if (correction_bias.has_value()) {
const torch::Tensor& bias_tensor = correction_bias.value();
TORCH_CHECK(bias_tensor.dim() == 1, "correction_bias must be 1D tensor [num_experts]");
TORCH_CHECK(bias_tensor.size(0) == num_experts, "correction_bias size must match num_experts");
TORCH_CHECK(bias_tensor.scalar_type() == at::ScalarType::Float, "correction_bias must be float32, got ");
CHECK(bias_tensor.dim() == 1)
<< "correction_bias must be 1D tensor [num_experts]";
CHECK(bias_tensor.size(0) == num_experts)
<< "correction_bias size must match num_experts";
CHECK(bias_tensor.scalar_type() == at::ScalarType::Float)
<< "correction_bias must be float32, got " << bias_tensor.scalar_type();
bias_ptr = bias_tensor.data_ptr<float>();
}
@@ -841,7 +849,7 @@ void topk_softmax(torch::Tensor& topk_weights, // [num_tokens, topk]
bias_ptr,
stream);
} else {
TORCH_CHECK(false, "Unsupported gating_output dtype");
LOG(FATAL) << "Unsupported gating_output dtype: " << dtype;
}
}
} // namespace xllm::kernel::cuda