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

123
ex_engine/csrc/moe/fused_moe_cuda.cpp Normal file
View File

@@ -0,0 +1,123 @@
/* Copyright 2026 The xLLM Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
https://github.com/jd-opensource/xllm/blob/main/LICENSE
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "kernels/cuda/cuda_ops_api.h"
#include "kernels/cuda/utils.h"
#include "platform/device.h"
namespace xllm::kernel::cuda {
torch::Tensor cutlass_fused_moe(
const torch::Tensor& input, // [num_tokens, hidden]
const torch::Tensor& token_selected_experts, // [num_tokens, top_k]
const torch::Tensor& token_final_scales, // [num_tokens, top_k]
const torch::Tensor&
fc1_expert_weights, // [num_experts, inter_dim, hidden]
const torch::Tensor&
fc2_expert_weights, // [num_experts, hidden, inter_dim]
torch::ScalarType output_dtype,
const std::vector<torch::Tensor>& quant_scales,
int32_t tp_size,
int32_t tp_rank,
int32_t ep_size,
int32_t ep_rank,
int32_t cluster_size,
int32_t cluster_rank,
const std::optional<torch::Tensor>& fc1_expert_biases,
const std::optional<torch::Tensor>& fc2_expert_biases,
const std::optional<torch::Tensor>& input_sf,
const std::optional<torch::Tensor>& swiglu_alpha,
const std::optional<torch::Tensor>& swiglu_beta,
const std::optional<torch::Tensor>& swiglu_limit,
const std::optional<torch::Tensor>& output,
bool enable_alltoall,
bool use_deepseek_fp8_block_scale,
bool use_w4_group_scaling,
bool use_mxfp8_act_scaling,
bool min_latency_mode,
bool use_packed_weights,
int32_t tune_max_num_tokens,
ActivationType activation_type) {
int64_t num_rows = input.size(0);
int64_t hidden_size = fc2_expert_weights.size(1);
if (min_latency_mode) {
num_rows *= fc2_expert_weights.size(0);
}
std::vector<int64_t> output_shape = {num_rows, hidden_size};
torch::Tensor result_output;
if (output.has_value() && output.value().defined()) {
result_output = output.value();
} else {
torch::TensorOptions options = input.options().dtype(output_dtype);
result_output = torch::empty(output_shape, options);
}
std::string fused_moe_uri = "fused_moe";
if (Device::is_support_sm90a()) {
fused_moe_uri += "_90";
} else if (Device::is_support_sm100a() || Device::is_support_sm100f()) {
fused_moe_uri += "_100";
} else if (Device::is_support_sm120a()) {
fused_moe_uri += "_120";
} else {
LOG(FATAL) << "FusedMoE is only supported on sm90, sm100, sm120.";
}
bind_tvmffi_stream_to_current_torch_stream(input.device());
ffi::Module fused_moe_runner =
get_function(fused_moe_uri, "init")(
to_dl_data_type(input.scalar_type()),
to_dl_data_type(fc1_expert_weights.scalar_type()),
to_dl_data_type(output_dtype),
use_deepseek_fp8_block_scale,
use_w4_group_scaling,
use_mxfp8_act_scaling,
use_packed_weights)
.cast<ffi::Module>();
fused_moe_runner->GetFunction("run_moe").value()(
to_ffi_tensor(result_output),
to_ffi_tensor(input),
to_ffi_tensor(token_selected_experts),
to_ffi_optional_tensor(token_final_scales),
to_ffi_tensor(fc1_expert_weights),
to_ffi_optional_tensor(fc1_expert_biases),
to_ffi_tensor(fc2_expert_weights),
to_ffi_optional_tensor(fc2_expert_biases),
to_ffi_optional_array_tensors(quant_scales),
to_ffi_optional_tensor(input_sf),
to_ffi_optional_tensor(swiglu_alpha),
to_ffi_optional_tensor(swiglu_beta),
to_ffi_optional_tensor(swiglu_limit),
tp_size,
tp_rank,
ep_size,
ep_rank,
cluster_size,
cluster_rank,
enable_alltoall,
min_latency_mode,
/*profile_ids=*/ffi::Optional<ffi::Array<int64_t>>(), // TODO: support
// auto tuning
// profile ids
support_pdl(),
activation_type);
return result_output;
}
} // namespace xllm::kernel::cuda