From f4e2264a833aea01c0c558bbd657170d34aae6ac Mon Sep 17 00:00:00 2001 From: project6-dev Date: Mon, 10 Aug 2026 03:59:37 +0000 Subject: [PATCH] ref(EX): import upstream ILU kernels + xllm MoE CUDA sources into ex_engine MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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} --- ex_engine/csrc/common_fused_moe.h | 54 ++ ex_engine/csrc/common_fused_moe_base.h | 27 + ex_engine/csrc/common_moe_fused_topk.cpp | 71 ++ ex_engine/csrc/common_moe_fused_topk.h | 53 ++ ex_engine/csrc/ilu_kernel_activation.cpp | 32 + ex_engine/csrc/ilu_kernel_attention.cpp | 163 ++++ ex_engine/csrc/ilu_kernel_fused_moe.cpp | 99 +++ ex_engine/csrc/ilu_kernel_group_gemm.cpp | 39 + ex_engine/csrc/ilu_kernel_matmul.cpp | 73 ++ ex_engine/csrc/ilu_kernel_norm.cpp | 51 ++ ex_engine/csrc/ilu_kernel_rope.cpp | 31 + ex_engine/csrc/ilu_layer_attention.cpp | 189 ++++ ex_engine/csrc/ilu_layer_attention.h | 82 ++ ex_engine/csrc/ilu_layer_fused_moe.cpp | 797 +++++++++++++++++ ex_engine/csrc/ilu_layer_fused_moe.h | 131 +++ ex_engine/csrc/moe/fused_moe_cuda.cpp | 123 +++ ex_engine/csrc/moe/moeTopKFuncs.cuh | 257 ++++++ ex_engine/csrc/moe/moe_align_sum_kernels.cu | 833 ++++++++++++++++++ ex_engine/csrc/moe/moe_fused_topk.cu | 56 ++ ex_engine/csrc/moe/moe_ops.h | 87 ++ ex_engine/csrc/moe/moe_topk.cuh | 285 ++++++ .../csrc/moe/moe_topk_sigmoid_kernels.cuh | 602 +++++++++++++ .../csrc/moe/moe_topk_softmax_kernels.cuh | 40 +- ex_engine/include/ilu_ops_api.h | 153 ++++ ex_engine/include/ilu_utils.h | 63 ++ ex_engine/include/ixformer.h | 147 ++++ 26 files changed, 4522 insertions(+), 16 deletions(-) create mode 100644 ex_engine/csrc/common_fused_moe.h create mode 100644 ex_engine/csrc/common_fused_moe_base.h create mode 100644 ex_engine/csrc/common_moe_fused_topk.cpp create mode 100644 ex_engine/csrc/common_moe_fused_topk.h create mode 100644 ex_engine/csrc/ilu_kernel_activation.cpp create mode 100644 ex_engine/csrc/ilu_kernel_attention.cpp create mode 100644 ex_engine/csrc/ilu_kernel_fused_moe.cpp create mode 100644 ex_engine/csrc/ilu_kernel_group_gemm.cpp create mode 100644 ex_engine/csrc/ilu_kernel_matmul.cpp create mode 100644 ex_engine/csrc/ilu_kernel_norm.cpp create mode 100644 ex_engine/csrc/ilu_kernel_rope.cpp create mode 100644 ex_engine/csrc/ilu_layer_attention.cpp create mode 100644 ex_engine/csrc/ilu_layer_attention.h create mode 100644 ex_engine/csrc/ilu_layer_fused_moe.cpp create mode 100644 ex_engine/csrc/ilu_layer_fused_moe.h create mode 100644 ex_engine/csrc/moe/fused_moe_cuda.cpp create mode 100644 ex_engine/csrc/moe/moeTopKFuncs.cuh create mode 100644 ex_engine/csrc/moe/moe_align_sum_kernels.cu create mode 100644 ex_engine/csrc/moe/moe_fused_topk.cu create mode 100644 ex_engine/csrc/moe/moe_ops.h create mode 100644 ex_engine/csrc/moe/moe_topk.cuh create mode 100644 ex_engine/csrc/moe/moe_topk_sigmoid_kernels.cuh create mode 100644 ex_engine/include/ilu_ops_api.h create mode 100644 ex_engine/include/ilu_utils.h create mode 100644 ex_engine/include/ixformer.h diff --git a/ex_engine/csrc/common_fused_moe.h b/ex_engine/csrc/common_fused_moe.h new file mode 100644 index 00000000..6e148c15 --- /dev/null +++ b/ex_engine/csrc/common_fused_moe.h @@ -0,0 +1,54 @@ +/* Copyright 2025 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. +==============================================================================*/ + +#pragma once + +#include + +#include "dense_mlp.h" +#include "framework/model/model_args.h" +#include "framework/model/model_input_params.h" +#include "framework/parallel_state/parallel_args.h" +#include "framework/quant_args.h" +#include "framework/state_dict/state_dict.h" +#include "framework/state_dict/utils.h" +#include "fused_moe_base.h" +#include "linear.h" + +namespace xllm { +namespace layer { + +// FusedMoE common implementation - placeholder for unsupported backends +// Actual implementations are in backend-specific fused_moe.h files. +class FusedMoEImpl : public torch::nn::Module { + public: + FusedMoEImpl() = default; + FusedMoEImpl(const ModelArgs& model_args, + const FusedMoEArgs& moe_args, + const QuantArgs& quant_args, + const ParallelArgs& parallel_args, + const torch::TensorOptions& options); + + torch::Tensor forward_experts(const torch::Tensor& hidden_states, + const torch::Tensor& router_logits, + bool enable_all2all_communication); + torch::Tensor forward(const torch::Tensor& hidden_states, + const ModelInputParams& input_params); + void load_state_dict(const StateDict& state_dict); +}; +TORCH_MODULE(FusedMoE); + +} // namespace layer +} // namespace xllm diff --git a/ex_engine/csrc/common_fused_moe_base.h b/ex_engine/csrc/common_fused_moe_base.h new file mode 100644 index 00000000..72e2f1cd --- /dev/null +++ b/ex_engine/csrc/common_fused_moe_base.h @@ -0,0 +1,27 @@ +/* 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. +==============================================================================*/ + +#pragma once + +namespace xllm { +namespace layer { + +struct FusedMoEArgs { + bool is_gated = true; + bool enable_result_reduction = true; +}; + +} // namespace layer +} // namespace xllm diff --git a/ex_engine/csrc/common_moe_fused_topk.cpp b/ex_engine/csrc/common_moe_fused_topk.cpp new file mode 100644 index 00000000..0c548a46 --- /dev/null +++ b/ex_engine/csrc/common_moe_fused_topk.cpp @@ -0,0 +1,71 @@ +/* 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 "moe_fused_topk.h" + +#include "kernels/ops_api.h" + +namespace xllm { +namespace layer { + +MoEFusedTopkImpl::MoEFusedTopkImpl(const ModelArgs& model_args, + const QuantArgs& quant_args, + const torch::TensorOptions& options) + : topk_(model_args.num_experts_per_tok()), + num_expert_group_(model_args.n_group()), + topk_group_(model_args.topk_group()), + route_scale_(model_args.routed_scaling_factor()), + hidden_size_(model_args.hidden_size()), + renormalize_(model_args.norm_topk_prob()), + scoring_func_(model_args.scoring_func()) { + const std::string& topk_method = model_args.topk_method(); + if (topk_method == "noaux_tc") { + e_score_correction_bias_ = register_parameter( + "e_score_correction_bias", + torch::empty({model_args.n_routed_experts()}, options), + false); + } +} + +// select the experts and return the reduce_weight and expert_id +std::tuple MoEFusedTopkImpl::forward( + torch::Tensor& router_logits) { + std::optional e_score_correction_bias = std::nullopt; + if (e_score_correction_bias_.defined()) { + e_score_correction_bias = e_score_correction_bias_; + } + + xllm::kernel::MoeFusedTopkParams moe_active_topk_params; + moe_active_topk_params.input = router_logits; + moe_active_topk_params.topk = topk_; + moe_active_topk_params.num_expert_group = num_expert_group_; + moe_active_topk_params.topk_group = topk_group_; + moe_active_topk_params.normalize = renormalize_; + moe_active_topk_params.normed_by = "topk_logit"; + moe_active_topk_params.scoring_func = scoring_func_; + moe_active_topk_params.route_scale = route_scale_; + moe_active_topk_params.e_score_correction_bias = e_score_correction_bias; + + return xllm::kernel::moe_active_topk(moe_active_topk_params); +} + +void MoEFusedTopkImpl::load_state_dict(const StateDict& state_dict) { + if (e_score_correction_bias_.defined() && + !e_score_correction_bias_is_loaded_) { + LOAD_WEIGHT(e_score_correction_bias); + } +} +} // namespace layer +} // namespace xllm diff --git a/ex_engine/csrc/common_moe_fused_topk.h b/ex_engine/csrc/common_moe_fused_topk.h new file mode 100644 index 00000000..05560a70 --- /dev/null +++ b/ex_engine/csrc/common_moe_fused_topk.h @@ -0,0 +1,53 @@ +/* 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. +==============================================================================*/ + +#pragma once + +#include + +#include "framework/model/model_args.h" +#include "framework/quant_args.h" +#include "framework/state_dict/state_dict.h" +#include "framework/state_dict/utils.h" + +namespace xllm { +namespace layer { + +class MoEFusedTopkImpl : public torch::nn::Module { + public: + MoEFusedTopkImpl(const ModelArgs& model_args, + const QuantArgs& quant_args, + const torch::TensorOptions& options); + + std::tuple forward( + torch::Tensor& router_logits); + + void load_state_dict(const StateDict& state_dict); + + private: + int64_t topk_; + int64_t num_expert_group_; + int64_t topk_group_; + double route_scale_; + int64_t hidden_size_; + bool renormalize_; + std::string scoring_func_; + + DEFINE_WEIGHT(e_score_correction_bias); +}; + +TORCH_MODULE(MoEFusedTopk); +} // namespace layer +} // namespace xllm diff --git a/ex_engine/csrc/ilu_kernel_activation.cpp b/ex_engine/csrc/ilu_kernel_activation.cpp new file mode 100644 index 00000000..ae2a16ba --- /dev/null +++ b/ex_engine/csrc/ilu_kernel_activation.cpp @@ -0,0 +1,32 @@ +/* Copyright 2025 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 "ilu_ops_api.h" + +using namespace ixformer; + +namespace xllm::kernel::ilu { + +void act_and_mul(torch::Tensor out, + torch::Tensor input, + const std::string& act_mode) { + if (act_mode == "silu") { + infer::silu_and_mul(input, out); + } else { + LOG(FATAL) << "Unsupported act mode: " << act_mode + << ", only support silu, gelu, gelu_tanh"; + } +} +} // namespace xllm::kernel::ilu diff --git a/ex_engine/csrc/ilu_kernel_attention.cpp b/ex_engine/csrc/ilu_kernel_attention.cpp new file mode 100644 index 00000000..aa257bf1 --- /dev/null +++ b/ex_engine/csrc/ilu_kernel_attention.cpp @@ -0,0 +1,163 @@ + +/* Copyright 2025 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 "ilu_ops_api.h" +#include "ixinfer.h" +#include "utils.h" + +using namespace ixformer; + +namespace xllm::kernel::ilu { + +void reshape_paged_cache(torch::Tensor& key, + std::optional& value, + torch::Tensor& key_cache, + std::optional& value_cache, + torch::Tensor& slot_mapping) { + auto value_ = value.value_or(torch::Tensor()); + auto value_cache_ = value_cache.value_or(torch::Tensor()); + + int64_t key_token_stride = key.stride(0); + int64_t value_token_stride = 0; + if (value_.defined()) { + value_token_stride = value_.stride(0); + } + slot_mapping = slot_mapping.to(at::kLong); + infer::xllm_reshape_and_cache(key, + value_, + key_cache, + value_cache_, + slot_mapping, + key_token_stride, + value_token_stride); +} + +void batch_prefill(torch::Tensor& query, + const torch::Tensor& key, + const std::optional& value, + torch::Tensor& output, + std::optional& output_lse, + const std::optional& q_cu_seq_lens, + const std::optional& kv_cu_seq_lens, + const std::optional& alibi_slope, + const std::optional& attn_bias, + const std::optional& q_quant_scale, + const std::optional& k_quant_scale, + const std::optional& v_quant_scale, + const torch::Tensor& block_tables, + int64_t max_query_len, + int64_t max_seq_len, + float scale, + bool is_causal, + int64_t window_size_left, + int64_t window_size_right, + const std::string& compute_dtype, + bool return_lse) { + double softcap = 0.0; + bool sqrt_alibi = false; + auto q_cu_seq_lens_ = q_cu_seq_lens.value_or(torch::Tensor()); + auto kv_cu_seq_lens_ = kv_cu_seq_lens.value_or(torch::Tensor()); + auto q_quant_scale_ = q_quant_scale.value_or(torch::Tensor()); + auto k_quant_scale_ = k_quant_scale.value_or(torch::Tensor()); + auto v_quant_scale_ = v_quant_scale.value_or(torch::Tensor()); + auto block_tables_ = block_tables; + auto key_ = key; + auto value_ = value.value(); + infer::ixinfer_flash_attn_unpad_with_block_tables(query, + key_, + value_, + output, + block_tables_, + q_cu_seq_lens_, + kv_cu_seq_lens_, + max_query_len, + max_seq_len, + is_causal, + window_size_left, + window_size_right, + static_cast(scale), + softcap, + sqrt_alibi, + alibi_slope, + c10::nullopt, + output_lse); +} + +void batch_decode(torch::Tensor& query, + const torch::Tensor& k_cache, + torch::Tensor& output, + const torch::Tensor& block_table, + const torch::Tensor& seq_lens, + const std::optional& v_cache, + std::optional& output_lse, + const std::optional& q_quant_scale, + const std::optional& k_cache_quant_scale, + const std::optional& v_cache_quant_scale, + const std::optional& out_quant_scale, + const std::optional& alibi_slope, + const std::optional& mask, + const std::string& compute_dtype, + int64_t max_seq_len, + int64_t window_size_left, + int64_t window_size_right, + float scale, + bool return_lse, + bool is_causal, + int64_t kv_cache_quant_bit_size) { + if (query.dim() == 4) { + query = + query + .view({query.size(0) * query.size(1), query.size(2), query.size(3)}) + .contiguous(); + } + if (output.dim() == 4) { + output = output + .view({output.size(0) * output.size(1), + output.size(2), + output.size(3)}) + .contiguous(); + ; + } + auto v_cache_ = v_cache.value_or(torch::Tensor()); + int64_t num_kv_heads = k_cache.size(1); + int64_t page_block_size = k_cache.size(2); + double softcap = 0.0; + bool enable_cuda_graph = false; + bool use_sqrt_alibi = false; + auto block_table_ = block_table; + auto k_cache_ = k_cache; + auto seq_lens_ = seq_lens; + infer::xllm_paged_attention(output, + query, + k_cache_, + v_cache_, + num_kv_heads, + scale, + block_table_, + seq_lens_, + page_block_size, + max_seq_len, + alibi_slope, + is_causal, + (int32_t)window_size_left, + (int32_t)window_size_right, + softcap, + enable_cuda_graph, + use_sqrt_alibi, + c10::nullopt); +} + +} // namespace xllm::kernel::ilu \ No newline at end of file diff --git a/ex_engine/csrc/ilu_kernel_fused_moe.cpp b/ex_engine/csrc/ilu_kernel_fused_moe.cpp new file mode 100644 index 00000000..794f9bd9 --- /dev/null +++ b/ex_engine/csrc/ilu_kernel_fused_moe.cpp @@ -0,0 +1,99 @@ +/* Copyright 2025 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 + +#include "ilu_ops_api.h" + +namespace xllm::kernel::ilu { + +std::tuple moe_active_topk( + const torch::Tensor& input, + int64_t topk, + int64_t num_expert_group, + int64_t topk_group, + bool normalize, + const std::optional& mask, + const std::string& normed_by, + const std::string& scoring_func, + double route_scale, + const std::optional& e_score_correction_bias) { + torch::Tensor input_ = input.to(torch::kFloat32); + auto reduce_weight = + torch::empty({input.size(0), topk}, + torch::dtype(torch::kFloat).device(input.device())); + auto topk_indices = + torch::empty({input.size(0), topk}, + torch::dtype(torch::kInt32).device(input.device())); + auto token_expert_indices = + torch::empty({input.size(0), topk}, + torch::dtype(torch::kInt32).device(input.device())); + + infer::topk_softmax( + reduce_weight, topk_indices, token_expert_indices, input_, false); + + auto tt = reduce_weight.sum(-1); + if (normalize) { + reduce_weight = reduce_weight / reduce_weight.sum(-1).unsqueeze(-1); + } + return std::make_tuple(reduce_weight, topk_indices); +} + +std::vector moe_gen_idx(torch::Tensor& expert_id, + int64_t expert_num) { + auto src_dst = expert_id.new_empty({expert_id.numel()}); + auto dst_src = torch::empty_like(src_dst); + auto expert_sizes_gpu = expert_id.new_empty({expert_num}); + auto expert_sizes_gpu_cumsum = expert_id.new_zeros({expert_id.numel() + 1}); + infer::moe_compute_token_index_api(expert_id, + src_dst, + dst_src, + expert_sizes_gpu, + /*expert_mask=*/std::nullopt, + /*expert_sizes_cpu*/ std::nullopt, + /*expert_sizes_gpu*/ std::nullopt, + 0, + expert_num, + expert_num); + + expert_sizes_gpu_cumsum = expert_sizes_gpu.cumsum(-1); + return {src_dst, dst_src, expert_sizes_gpu, expert_sizes_gpu_cumsum}; +} + +torch::Tensor moe_expand_input(const torch::Tensor& input, + const torch::Tensor& gather_index, + const torch::Tensor& combine_idx, + int64_t topk) { + int64_t dst_tokens = input.size(0) * topk; + auto output = input.new_empty({dst_tokens, input.size(1)}); + infer::moe_expand_input( + output, input, combine_idx, gather_index, dst_tokens, topk); + + return output; +} + +torch::Tensor moe_combine_result(torch::Tensor& input, torch::Tensor& weight) { + input = input.view({-1, weight.size(1), input.size(1)}); + auto output = input.new_empty({input.size(0), input.size(2)}); + infer::moe_output_reduce_sum(output, + input, + weight, + /*mask=*/std::nullopt, + /*extra_residual*/ std::nullopt, + /*scaling_factor=*/1.0); + return output; +} + +} // namespace xllm::kernel::ilu diff --git a/ex_engine/csrc/ilu_kernel_group_gemm.cpp b/ex_engine/csrc/ilu_kernel_group_gemm.cpp new file mode 100644 index 00000000..38743e66 --- /dev/null +++ b/ex_engine/csrc/ilu_kernel_group_gemm.cpp @@ -0,0 +1,39 @@ +/* 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 "ilu_ops_api.h" + +namespace xllm::kernel::ilu { + +torch::Tensor group_gemm(torch::Tensor& input, + torch::Tensor& weight, + torch::Tensor& tokens_per_experts, + const std::optional& dst_to_src, + torch::Tensor& output) { + infer::moe_w16a16_group_gemm( + output, + input, + weight, + tokens_per_experts, + dst_to_src, + /*bias=*/std::nullopt, + /*format=*/"TN", + /*persistent=*/0, + /*output_n=*/tokens_per_experts.sum().item()); + + return output; +} + +} // namespace xllm::kernel::ilu diff --git a/ex_engine/csrc/ilu_kernel_matmul.cpp b/ex_engine/csrc/ilu_kernel_matmul.cpp new file mode 100644 index 00000000..91b6868f --- /dev/null +++ b/ex_engine/csrc/ilu_kernel_matmul.cpp @@ -0,0 +1,73 @@ +/* Copyright 2025 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 "ilu_ops_api.h" +#include "util/env_var.h" + +namespace xllm::kernel::ilu { + +bool gemv_conditions(const torch::Tensor& input, + const torch::Tensor& weight, + const torch::Tensor& bias, + int64_t gemv_max_batch) { + // gemv input:[m,k] weight:[n,k] + // 1. m <= gemv_max_batch + // 2. k % 32 == 0 && n % 2 == 0 + // 3. bias is None + + torch::Tensor input_view = input.view({-1, input.size(-1)}); + torch::Tensor weight_view = weight.view({-1, weight.size(-1)}); + + int64_t m = input_view.size(0); + int64_t k = input_view.size(1); + int64_t n = weight_view.size(0); + + if (bias.defined() == false && m <= gemv_max_batch && k % 32 == 0 && + n % 2 == 0) { + return true; + } + return false; +} + +torch::Tensor matmul(torch::Tensor a, + torch::Tensor b, + std::optional bias) { + int64_t act_type = -1; + bool persistent = false; + std::vector output_shape = a.sizes().vec(); + if (!output_shape.empty()) { + output_shape[output_shape.size() - 1] = b.size(0); + } + torch::Tensor output = a.new_empty(output_shape); + + bool use_gemv = true; + const int64_t gemv_max_batch = 1; + const bool disable_infer_gemm_ex = + xllm::util::get_bool_env("DISABLE_INFER_GEMM_EX", false); + + use_gemv = + use_gemv && + gemv_conditions(a, b, bias.value_or(at::Tensor()), gemv_max_batch) && + !disable_infer_gemm_ex && (act_type == -1); + + if (use_gemv) { + output = infer::ixformer_linear_ex(a, b, bias, output); + } else { + output = infer::ixformer_linear(a, b, act_type, bias, output, persistent); + } + return output; +} + +} // namespace xllm::kernel::ilu diff --git a/ex_engine/csrc/ilu_kernel_norm.cpp b/ex_engine/csrc/ilu_kernel_norm.cpp new file mode 100644 index 00000000..c5a98595 --- /dev/null +++ b/ex_engine/csrc/ilu_kernel_norm.cpp @@ -0,0 +1,51 @@ +/* Copyright 2025 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 "ilu_ops_api.h" +#include "utils.h" + +using namespace ixformer; + +namespace xllm::kernel::ilu { + +void residual_layer_norm(torch::Tensor& input, + torch::Tensor& output, + std::optional& residual, + torch::Tensor& weight, + std::optional& bias, + std::optional& residual_out, + double eps) { + auto residual_ = residual.value_or(torch::zeros_like(input)); + torch::Tensor residual_out_ = residual_out.value_or(torch::zeros_like(input)); + infer::residual_rms_norm(input, + residual_, + weight, + output, + residual_out_, + bias, + /*alpha=*/1.0, + eps, + false); +} + +void rms_norm(torch::Tensor& output, + torch::Tensor& input, + torch::Tensor& weight, + double eps) { + std::optional fused_bias = std::nullopt; + infer::rms_norm(input, weight, output, fused_bias, eps); +} + +} // namespace xllm::kernel::ilu \ No newline at end of file diff --git a/ex_engine/csrc/ilu_kernel_rope.cpp b/ex_engine/csrc/ilu_kernel_rope.cpp new file mode 100644 index 00000000..89370b79 --- /dev/null +++ b/ex_engine/csrc/ilu_kernel_rope.cpp @@ -0,0 +1,31 @@ +/* Copyright 2025 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 "ilu_ops_api.h" +#include "utils.h" + +namespace xllm::kernel::ilu { + +void apply_rope_pos_ids_cos_sin_cache(torch::Tensor& query, + torch::Tensor& key, + torch::Tensor& cos_sin_cache, + torch::Tensor& positions, + bool interleave) { + const int64_t head_size = cos_sin_cache.size(-1); + infer::xllm_rotary_embedding( + positions, query, key, head_size, cos_sin_cache, !interleave); +} + +} // namespace xllm::kernel::ilu diff --git a/ex_engine/csrc/ilu_layer_attention.cpp b/ex_engine/csrc/ilu_layer_attention.cpp new file mode 100644 index 00000000..b66f28a4 --- /dev/null +++ b/ex_engine/csrc/ilu_layer_attention.cpp @@ -0,0 +1,189 @@ +/* Copyright 2025 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 "attention.h" + +#include "kernels/ilu/ilu_ops_api.h" +#include "kernels/ops_api.h" + +namespace xllm { +namespace layer { +AttentionImpl::AttentionImpl(int64_t num_heads, + int64_t head_size, + float scale, + int64_t num_kv_heads, + int64_t sliding_window) + : num_heads_(num_heads), + head_size_(head_size), + scale_(scale), + num_kv_heads_(num_kv_heads), + v_head_dim_(head_size), + use_fused_mla_qkv_(false), + enable_lighting_indexer_(false), + enable_mla_(false), + sliding_window_(sliding_window) { + if (sliding_window_ > -1) { + sliding_window_ = sliding_window_ - 1; + } +} + +AttentionImpl::AttentionImpl(int64_t num_heads, + int64_t head_size, + int64_t num_kv_heads, + int64_t v_head_dim, + int64_t sliding_window, + float scale, + bool use_fused_mla_qkv, + bool enable_lighting_indexer, + bool enable_mla) + : num_heads_(num_heads), + head_size_(head_size), + scale_(scale), + num_kv_heads_(num_kv_heads), + v_head_dim_(v_head_dim), + use_fused_mla_qkv_(use_fused_mla_qkv), + enable_lighting_indexer_(enable_lighting_indexer), + enable_mla_(enable_mla), + sliding_window_(sliding_window) { + if (sliding_window_ > -1) { + sliding_window_ = sliding_window_ - 1; + } +} + +std::tuple> AttentionImpl::forward( + const AttentionMetadata& attn_metadata, + torch::Tensor& query, + torch::Tensor& key, + torch::Tensor& value, + KVCache& kv_cache) { + std::optional output_lse = std::nullopt; + torch::Tensor output; + if (enable_mla_) { + output = torch::empty({query.size(0), num_heads_ * v_head_dim_}, + query.options()); + } else { + output = torch::empty_like(query); + } + if (attn_metadata.is_dummy) { + return std::make_tuple(output, output_lse); + } + + bool only_prefill = + attn_metadata.is_prefill || attn_metadata.is_chunked_prefill; + int64_t num_kv_heads = (enable_mla_ && !only_prefill) ? 1 : num_kv_heads_; + torch::Tensor k_cache = kv_cache.get_k_cache(); + std::optional v_cache; + std::optional v; + if (!enable_mla_) { + v = value.view({-1, num_kv_heads, head_size_}); + v_cache = kv_cache.get_v_cache(); + } + + bool skip_process_cache = enable_mla_ && (only_prefill || use_fused_mla_qkv_); + if (!skip_process_cache) { + xllm::kernel::ReshapePagedCacheParams reshape_paged_cache_params; + reshape_paged_cache_params.key = key.view({-1, num_kv_heads, head_size_}); + reshape_paged_cache_params.value = v; + reshape_paged_cache_params.k_cache = k_cache; + reshape_paged_cache_params.v_cache = v_cache; + reshape_paged_cache_params.slot_mapping = attn_metadata.slot_mapping; + xllm::kernel::reshape_paged_cache(reshape_paged_cache_params); + } + + if (enable_lighting_indexer_ || !only_prefill) { + decoder_forward(query, output, k_cache, v_cache, attn_metadata); + } else { + prefill_forward(query, key, value, output, k_cache, v_cache, attn_metadata); + } + + int64_t head_size = enable_mla_ ? v_head_dim_ : head_size_; + output = output.view({-1, num_heads_ * head_size}); + return {output, output_lse}; +} + +void AttentionImpl::prefill_forward(torch::Tensor& query, + torch::Tensor& key, + torch::Tensor& value, + torch::Tensor& output, + const torch::Tensor& k_cache, + const std::optional& v_cache, + const AttentionMetadata& attn_metadata) { + int64_t head_size_v = enable_mla_ ? v_head_dim_ : head_size_; + std::optional output_lse = std::nullopt; + query = query.view({-1, num_heads_, head_size_}); + output = output.view({-1, num_heads_, head_size_v}); + // torch::Tensor k_cache_ = k_cache; + // torch::Tensor v_cache_ = v_cache.value(); + xllm::kernel::ilu::batch_prefill(query, + k_cache, + v_cache, + output, + output_lse, + attn_metadata.q_cu_seq_lens, + attn_metadata.kv_cu_seq_lens, + /*alibi_slope=*/std::nullopt, + /*attn_bias=*/std::nullopt, + /*q_quant_scale=*/std::nullopt, + /*k_quant_scale=*/std::nullopt, + /*v_quant_scale=*/std::nullopt, + attn_metadata.block_table, + attn_metadata.max_query_len, + attn_metadata.max_seq_len, + scale_, + attn_metadata.is_causal, + sliding_window_, + /*window_size_right=*/-1, + attn_metadata.compute_dtype, + /*return_lse=*/false); +} + +void AttentionImpl::decoder_forward(torch::Tensor& query, + torch::Tensor& output, + const torch::Tensor& k_cache, + const std::optional& v_cache, + const AttentionMetadata& attn_metadata) { + int64_t head_size_v = enable_mla_ ? v_head_dim_ : head_size_; + query = query.view({-1, 1, num_heads_, head_size_}); + output = output.view({-1, 1, num_heads_, head_size_v}); + std::optional output_lse = std::nullopt; + + int64_t block_aligned_max_seq_len = + attn_metadata.block_table.size(-1) * k_cache.size(2); + + xllm::kernel::ilu::batch_decode(query, + k_cache, + output, + attn_metadata.block_table, + attn_metadata.kv_seq_lens, + v_cache, + output_lse, + /*q_quant_scale=*/std::nullopt, + /*k_quant_scale=*/std::nullopt, + /*v_quant_scale=*/std::nullopt, + /*out_quant_scale=*/std::nullopt, + /*alibi_slope=*/std::nullopt, + attn_metadata.attn_mask, + attn_metadata.compute_dtype, + block_aligned_max_seq_len, + sliding_window_, + /*window_size_right=*/-1, + scale_, + /*return_lse=*/false, + attn_metadata.is_causal, + /*kv_cache_quant_bit_size=*/-1); +} + +} // namespace layer +} // namespace xllm diff --git a/ex_engine/csrc/ilu_layer_attention.h b/ex_engine/csrc/ilu_layer_attention.h new file mode 100644 index 00000000..a971835f --- /dev/null +++ b/ex_engine/csrc/ilu_layer_attention.h @@ -0,0 +1,82 @@ +/* Copyright 2025 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. +==============================================================================*/ + +#pragma once + +#include + +#include + +#include "framework/kv_cache/kv_cache.h" +#include "framework/model/model_input_params.h" +#include "layers/common/attention_metadata.h" + +namespace xllm { +namespace layer { +class AttentionImpl : public torch::nn::Module { + public: + AttentionImpl() = default; + + AttentionImpl(int64_t num_heads, + int64_t head_size, + float scale, + int64_t num_kv_heads, + int64_t sliding_window); + AttentionImpl(int64_t num_heads, + int64_t head_size, + int64_t num_kv_heads, + int64_t v_head_dim, + int64_t sliding_window, + float scale, + bool use_fused_mla_qkv, + bool enable_lighting_indexer, + bool enable_mla); + + std::tuple> forward( + const AttentionMetadata& attn_metadata, + torch::Tensor& query, + torch::Tensor& key, + torch::Tensor& value, + KVCache& kv_cache); + + void prefill_forward(torch::Tensor& query, + torch::Tensor& key, + torch::Tensor& value, + torch::Tensor& output, + const torch::Tensor& k_cache, + const std::optional& v_cache, + const AttentionMetadata& attn_metadata); + + void decoder_forward(torch::Tensor& query, + torch::Tensor& output, + const torch::Tensor& k_cache, + const std::optional& v_cache, + const AttentionMetadata& attn_metadata); + + private: + int64_t num_heads_; + int64_t head_size_; + float scale_; + int64_t num_kv_heads_; + int64_t v_head_dim_; + bool use_fused_mla_qkv_; + bool enable_lighting_indexer_; + bool enable_mla_; + int64_t sliding_window_; +}; +TORCH_MODULE(Attention); + +} // namespace layer +} // namespace xllm diff --git a/ex_engine/csrc/ilu_layer_fused_moe.cpp b/ex_engine/csrc/ilu_layer_fused_moe.cpp new file mode 100644 index 00000000..4238012e --- /dev/null +++ b/ex_engine/csrc/ilu_layer_fused_moe.cpp @@ -0,0 +1,797 @@ +/* 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 "fused_moe.h" + +#include + +#include + +#include "common/global_flags.h" +#include "framework/parallel_state/parallel_state.h" +#include "kernels/ops_api.h" +#include "layers/common/dp_utils.h" +#include "util/utils.h" + +namespace { + +int32_t get_dtype_size(torch::ScalarType dtype) { + return static_cast(torch::elementSize(dtype)); +} + +} // namespace + +namespace xllm { +namespace layer { + +FusedMoEImpl::FusedMoEImpl(const ModelArgs& model_args, + const FusedMoEArgs& moe_args, + const QuantArgs& quant_args, + const ParallelArgs& parallel_args, + const torch::TensorOptions& options) + : num_total_experts_(static_cast(model_args.n_routed_experts())), + topk_(model_args.num_experts_per_tok()), + num_expert_group_(model_args.n_group()), + topk_group_(model_args.topk_group()), + route_scale_(model_args.routed_scaling_factor()), + hidden_size_(model_args.hidden_size()), + n_shared_experts_(model_args.n_shared_experts()), + is_gated_(moe_args.is_gated), + renormalize_(model_args.norm_topk_prob() ? 1 : 0), + hidden_act_(model_args.hidden_act()), + scoring_func_(model_args.scoring_func()), + quant_args_(quant_args), + parallel_args_(parallel_args), + options_(options), + device_(options.device()) { + const int64_t num_experts = num_total_experts_; + const int64_t intermediate_size = + static_cast(model_args.moe_intermediate_size()); + const std::string& topk_method = model_args.topk_method(); + int64_t ep_size = parallel_args.ep_size(); + int64_t ep_rank = 0; + tp_pg_ = parallel_args.tp_group_; + if (ep_size > 1) { + ep_rank = parallel_args.moe_ep_group_->rank(); + tp_pg_ = parallel_args.moe_tp_group_; + } + + // smoothquant check: If quant_method is not empty, only w8a8 smoothquant is + // supported + if (!quant_args.quant_method().empty()) { + if (quant_args.quant_method() != "smoothquant" || quant_args.bits() != 8 || + !quant_args.activation_dynamic()) { + LOG(FATAL) << "FusedMoE only supports w8a8 smoothquant quantization when " + "quant_method is set. " + << "Got quant_method=" << quant_args.quant_method() + << ", bits=" << quant_args.bits() + << ", activation_dynamic=" << quant_args.activation_dynamic(); + } + // If confirmed as smoothquant w8a8, set is_smoothquant_ to true + is_smoothquant_ = true; + } else { + is_smoothquant_ = false; + } + + // Deep EP initialization check + enable_deep_ep_ = FLAGS_expert_parallel_degree == 2 && ep_size > 1; + if (enable_deep_ep_) { + // for now, we only implement the deep ep for decode stage. + // so we will assume the max_token_num is limited to max_batch_size * (1+K) + // K is the number of speculative tokens. + int64_t dispatch_token_size; + if (quant_args.quant_method() == "smoothquant") { + // float32 is for the scale of the quantized input + dispatch_token_size = hidden_size_ * get_dtype_size(torch::kInt8) + + get_dtype_size(torch::kFloat32); + } else { + dispatch_token_size = + hidden_size_ * get_dtype_size(options_.dtype().toScalarType()); + } + torch::ScalarType combine_dtype = options_.dtype().toScalarType(); + int64_t combine_token_size = hidden_size_ * get_dtype_size(combine_dtype); + // Ensure calculation base is at least ep_size + int64_t effective_seqs = + std::max((int64_t)FLAGS_max_seqs_per_batch, (int64_t)ep_size); + // NOTE: FLAGS_max_seqs_per_batch represents the maximum total batch size, + // regardless of the dp size. To ensure robust scheduling and account + // for the worst-case scenario, we must guarantee that each rank is capable + // of handling the maximum possible number of tokens. Therefore, we define + // max_num_tokens_per_rank as the full maximum value, without dividing by + // either the rank count or the dp size. + int64_t max_num_tokens_per_rank = + (1 + FLAGS_num_speculative_tokens) * effective_seqs * topk_; + + // make sure that all layers share the same deep ep instance + // so that the memory footprint is minimized + deep_ep_ = DeepEPManager::get_instance(dispatch_token_size, + combine_token_size, + max_num_tokens_per_rank, + num_experts, + parallel_args, + options_); + + // obtain the buffer and parameters of deep ep + deep_ep_buffer_ = deep_ep_->get_buffer(); + deep_ep_params_ = deep_ep_->get_params(); + + // intermediate buffer that can be initialized once + // we place these tensor here in order to speed up forward pass + int64_t n_tokens_recv = deep_ep_params_.max_num_tokens_recv; + int64_t token_bytes = is_smoothquant_ + ? get_dtype_size(torch::kInt8) + : get_dtype_size(options_.dtype().toScalarType()); + token_bytes = token_bytes * hidden_size_; + int64_t head_size = n_tokens_recv * token_bytes; + dispatch_recv_token_tensor_head_ = + deep_ep_buffer_.combine_send_token_tensor.narrow(0, 0, head_size) + .view({n_tokens_recv, token_bytes}); + // input scale in smoothquant + if (is_smoothquant_) { + int64_t tail_size = n_tokens_recv * get_dtype_size(torch::kFloat32); + dispatch_recv_token_tensor_tail_ = + deep_ep_buffer_.combine_send_token_tensor + .narrow(0, head_size, tail_size) + .view({n_tokens_recv, -1}); + } + } + + // calculate the number of experts per rank + num_experts_per_rank_ = num_experts / ep_size; + start_expert_id_ = ep_rank * num_experts_per_rank_; + + if (topk_method == "noaux_tc") { + e_score_correction_bias_ = register_parameter( + "e_score_correction_bias", torch::empty({num_experts}, options), false); + } + + gate_ = register_module( + "gate_proj", + ReplicatedLinear(hidden_size_, num_experts, false, quant_args, options)); + if (n_shared_experts_ > 0) { + ProcessGroup* shared_expert_pg; + if (parallel_args_.ep_size() > 1) { + // we use tp=1 for shared experts computation in deep ep mode + CHECK(parallel_args_.ep_size() == parallel_args_.world_size()) + << "Models with shared experts only support ep_size equal to " + "world size for now."; + shared_expert_pg = parallel_args.moe_tp_group_; + } else { + shared_expert_pg = parallel_args.process_group_; + } + // The shared experts computation can proceed in parallel with the + // final communication step during the MoE computation, as long as it + // remains independent of any communication operations. For optimal + // performance, ensure that the shared experts layer on each rank always + // maintains its own unique weights. + shared_experts_ = + register_module("shared_experts", + DenseMLP(hidden_size_, + intermediate_size * n_shared_experts_, + is_gated_, + false, + hidden_act_, + /*enable_result_reduction=*/true, + quant_args, + shared_expert_pg, + options)); + } + + // create weight buffer + const int64_t world_size = tp_pg_->world_size(); + int64_t local_intermediate_size = intermediate_size / world_size; + if (is_smoothquant_) { + auto quant_option = options_.dtype(torch::kInt8); + auto fp_option = options_.dtype(torch::kFloat32); + w13_ = register_parameter( + "w13", + torch::empty( + {num_experts_per_rank_, local_intermediate_size * 2, hidden_size_}, + quant_option), + false); + w13_scale_ = register_parameter( + "w13_scale", + torch::empty({num_experts_per_rank_, local_intermediate_size * 2}, + fp_option), + false); + // Note: We do not check enable_deep_ep_ here, since smooth quantization + // information may be needed even when deep EP mode is disabled. This allows + // retrieving quantization parameters for any subset of experts as required. + input_smooth_ = register_parameter( + "input_smooth", + torch::empty({num_total_experts_, hidden_size_}, fp_option), + false); + w2_ = register_parameter( + "w2", + torch::empty( + {num_experts_per_rank_, hidden_size_, local_intermediate_size}, + quant_option), + false); + w2_scale_ = register_parameter( + "w2_scale", + torch::empty({num_experts_per_rank_, hidden_size_}, fp_option), + false); + act_smooth_ = register_parameter( + "act_smooth", + torch::empty({num_experts_per_rank_, local_intermediate_size}, + fp_option), + false); + + } else { + w13_ = register_parameter( + "w13", + torch::empty( + {num_experts_per_rank_, local_intermediate_size * 2, hidden_size_}, + options_), + false); + w2_ = register_parameter( + "w2", + torch::empty( + {num_experts_per_rank_, hidden_size_, local_intermediate_size}, + options_), + false); + } +} + +torch::Tensor FusedMoEImpl::create_group_gemm_output( + const torch::Tensor& a, + const torch::Tensor& b, + const torch::Tensor& group_list, + torch::ScalarType dtype, + torch::Tensor& workspace) { + // unify shape logic: define the target shape once. + bool is_3d_weight = (b.dim() != 2); + int64_t num_tokens = a.size(0); + int64_t out_dim = is_3d_weight ? b.size(1) : b.size(0); + + std::vector output_shape; + int64_t required_elements = num_tokens * out_dim; + + if (is_3d_weight) { + output_shape = {num_tokens, out_dim}; + } else { + output_shape = {group_list.size(0), num_tokens, out_dim}; + required_elements *= group_list.size(0); + } + + auto options = a.options().dtype(dtype); + + // non-smoothquant: direct allocation + if (!is_smoothquant_) { + return torch::empty(output_shape, options); + } + + // smoothquant: managed workspace logic + if (!workspace.defined()) { + // Lazy initialization: allocate max buffer for the lifecycle + // Note: accessing class members w13_ and w2_ directly for context + int64_t max_width = std::max(w13_.size(1), w2_.size(1)); + workspace = torch::empty({num_tokens * max_width}, options); + } + + // view construction + CHECK(workspace.numel() >= required_elements) + << "FusedMoE Workspace too small! Alloc: " << workspace.numel() + << ", Req: " << required_elements; + + // utilize the pre-calculated output_shape + return workspace.slice(0, 0, required_elements).view(output_shape); +} + +torch::Tensor FusedMoEImpl::select_experts( + const torch::Tensor& hidden_states_2d, + const torch::Tensor& router_logits_2d, + SelectedExpertInfo& selected_expert_info, + bool enable_all2all_communication) { + // prepare the parameters for select_experts + std::optional e_score_correction_bias = std::nullopt; + if (e_score_correction_bias_.defined()) { + e_score_correction_bias = e_score_correction_bias_; + } + int64_t expert_size = w13_.size(0); + + // Step 1: apply softmax topk or sigmoid topk / routing logic + torch::Tensor reduce_weight; + torch::Tensor expert_id; + { + xllm::kernel::MoeFusedTopkParams moe_active_topk_params; + moe_active_topk_params.input = router_logits_2d; + moe_active_topk_params.topk = topk_; + moe_active_topk_params.num_expert_group = num_expert_group_; + moe_active_topk_params.topk_group = topk_group_; + moe_active_topk_params.normalize = renormalize_; + moe_active_topk_params.normed_by = "topk_logit"; + moe_active_topk_params.scoring_func = scoring_func_; + moe_active_topk_params.route_scale = route_scale_; + moe_active_topk_params.e_score_correction_bias = e_score_correction_bias; + std::tie(reduce_weight, expert_id) = + xllm::kernel::moe_active_topk(moe_active_topk_params); + } + + // Step 2: generate expert ids + torch::Tensor gather_idx; + torch::Tensor combine_idx; + torch::Tensor token_count; + std::optional cusum_token_count; + { + xllm::kernel::MoeGenIdxParams moe_gen_idx_params; + moe_gen_idx_params.expert_id = expert_id; + moe_gen_idx_params.expert_num = num_total_experts_; + std::vector output_vec = + xllm::kernel::moe_gen_idx(moe_gen_idx_params); + gather_idx = output_vec[0]; + combine_idx = output_vec[1]; + token_count = output_vec[2]; + // during all2all communication, we do not need cusum_token_count in the + // following computation + if (enable_all2all_communication) { + cusum_token_count = std::nullopt; + } else { + cusum_token_count = output_vec[3]; + } + } + + // Step 3: expand and quantize input if needed + torch::Tensor expand_hidden_states; + torch::Tensor hidden_states_scale; + torch::Tensor token_count_slice; + // all2all related variables + torch::Tensor dispatch_send_token_tensor; + // in all2all, the input is scattered, so there is no need to slice the token + // count, and we can use the dispatch buffer directly + if (enable_all2all_communication) { + token_count_slice = token_count; + int64_t num_token_expand = hidden_states_2d.size(0) * topk_; + int64_t dispatch_bytes = + num_token_expand * deep_ep_params_.dispatch_token_size; + dispatch_send_token_tensor = + deep_ep_buffer_.dispatch_send_token_tensor.slice(0, 0, dispatch_bytes) + .view({num_token_expand, deep_ep_params_.dispatch_token_size}); + } else { + token_count_slice = + token_count.slice(0, start_expert_id_, start_expert_id_ + expert_size); + } + + if (is_smoothquant_) { + xllm::kernel::ScaledQuantizeParams scaled_quantize_params; + scaled_quantize_params.x = hidden_states_2d; + // use dispatch_send_token_tensor buffer for input + // to reduce memory footprint + if (enable_all2all_communication) { + scaled_quantize_params.smooth = input_smooth_; + scaled_quantize_params.output = + dispatch_send_token_tensor.slice(1, 0, hidden_size_); + } else { + scaled_quantize_params.smooth = input_smooth_.slice( + 0, start_expert_id_, start_expert_id_ + expert_size); + scaled_quantize_params.gather_index_start_position = + cusum_token_count.value().index({start_expert_id_}).unsqueeze(0); + } + scaled_quantize_params.token_count = token_count_slice; + scaled_quantize_params.gather_index = gather_idx; + scaled_quantize_params.act_mode = "none"; + scaled_quantize_params.active_coef = 1.0; + scaled_quantize_params.is_gated = false; + scaled_quantize_params.quant_type = torch::kChar; + std::tie(expand_hidden_states, hidden_states_scale) = + xllm::kernel::scaled_quantize(scaled_quantize_params); + if (enable_all2all_communication) { + // since view_as_dtype has not supported stride yet, + // we need to copy the scale output to the dispatch buffer + torch::Tensor dispatch_scale_slice = + dispatch_send_token_tensor.slice(1, hidden_size_); + torch::Tensor hidden_states_scale_bytes = + view_as_dtype(hidden_states_scale, torch::kInt8) + .view_as(dispatch_scale_slice); + dispatch_scale_slice.copy_(hidden_states_scale_bytes); + } + } else { + xllm::kernel::MoeExpandInputParams moe_expand_input_params; + moe_expand_input_params.input = hidden_states_2d; + moe_expand_input_params.gather_index = gather_idx; + moe_expand_input_params.combine_idx = combine_idx; + moe_expand_input_params.topk = topk_; + expand_hidden_states = + xllm::kernel::moe_expand_input(moe_expand_input_params); + if (enable_all2all_communication) { + // use copy to place the output inside the dispatch buffer + torch::Tensor dispatch_tensor = + view_as_dtype(expand_hidden_states, torch::kChar); + dispatch_send_token_tensor.copy_(dispatch_tensor); + } + } + + // collect the selected tensor + selected_expert_info.reduce_weight = reduce_weight; + selected_expert_info.combine_idx = combine_idx; + selected_expert_info.token_count_slice = token_count_slice; + selected_expert_info.cusum_token_count = cusum_token_count; + if (is_smoothquant_) { + selected_expert_info.input_scale = hidden_states_scale; + } + + return expand_hidden_states; +} + +torch::Tensor FusedMoEImpl::forward_experts(const torch::Tensor& hidden_states, + const torch::Tensor& router_logits, + bool enable_all2all_communication) { + if (!stream_initialized_) { + // update device record + device_ = xllm::Device(hidden_states.device()); + + // acquire streams from the pool again + routed_stream_ = device_.get_stream_from_pool(); + shared_stream_ = device_.get_stream_from_pool(); + stream_initialized_ = true; + } + + std::optional e_score_correction_bias = std::nullopt; + if (e_score_correction_bias_.defined()) { + e_score_correction_bias = e_score_correction_bias_; + } + + // prepare the parameters for MoE computation + torch::Tensor shared_expert_output; + torch::IntArrayRef hidden_states_shape = hidden_states.sizes(); + torch::ScalarType hidden_states_dtype = hidden_states.dtype().toScalarType(); + torch::Tensor hidden_states_2d = + hidden_states.reshape({-1, hidden_states.size(-1)}); + torch::Tensor router_logits_2d = + router_logits.reshape({-1, router_logits.size(-1)}); + int64_t group_gemm_max_dim = enable_all2all_communication + ? deep_ep_params_.max_num_tokens_recv / topk_ + : hidden_states_2d.size(0); + int64_t expert_size = w13_.size(0); + + // Step 1-3: select experts + SelectedExpertInfo selected_expert_info; + torch::Tensor expand_hidden_states = + select_experts(hidden_states_2d, + router_logits_2d, + selected_expert_info, + enable_all2all_communication); + + // Communciation Step 1: Dipatch + // intermediate outputs that are used both in dispatch and combine + torch::Tensor gather_by_rank_index; + torch::Tensor token_sum; + if (enable_all2all_communication) { + int64_t dispatch_token_num = hidden_states_2d.size(0) * topk_; + + // 1. Dispatch Step: Generate layout and send data + deep_ep_->dispatch_step(dispatch_token_num, + selected_expert_info.token_count_slice); + + // 2. Process Result: Generate indices and unpack to computation buffer + // use the buffer during initialization for the output + expand_hidden_states = dispatch_recv_token_tensor_head_; + std::optional output_tail = std::nullopt; + if (is_smoothquant_) { + output_tail = dispatch_recv_token_tensor_tail_; + // update selected_expert_info with the tail (input scale) + selected_expert_info.input_scale = output_tail; + } + + DeepEPMetaResult deep_ep_meta = deep_ep_->process_dispatch_result( + num_experts_per_rank_, expand_hidden_states, output_tail); + + // Extract metadata for subsequent steps + gather_by_rank_index = deep_ep_meta.gather_rank_index; + selected_expert_info.token_count_slice = deep_ep_meta.token_count_slice; + token_sum = deep_ep_meta.token_sum; + } + + // common gemm workspace for reduce memory footprint + torch::Tensor gemm_workspace; + + // Step 4: group gemm 1 + torch::Tensor gemm1_out = + create_group_gemm_output(expand_hidden_states, + w13_, + selected_expert_info.token_count_slice, + hidden_states_dtype, + gemm_workspace); + // ensure the lifespan of these parameters via brace + { + xllm::kernel::GroupGemmParams group_gemm_params; + torch::ScalarType a_dtype = + is_smoothquant_ ? torch::kInt8 : hidden_states_dtype; + group_gemm_params.a = + view_as_dtype(expand_hidden_states, a_dtype).view({-1, hidden_size_}); + group_gemm_params.b = w13_; + group_gemm_params.token_count = + selected_expert_info.token_count_slice.to("cpu"); + if (is_smoothquant_) { + torch::Tensor a_scale = + selected_expert_info.input_scale.value().flatten(); + selected_expert_info.input_scale = + view_as_dtype(a_scale, torch::kFloat32); + group_gemm_params.a_scale = selected_expert_info.input_scale; + group_gemm_params.b_scale = w13_scale_; + } + group_gemm_params.max_dim = group_gemm_max_dim; + group_gemm_params.trans_a = false; + group_gemm_params.trans_b = true; + group_gemm_params.a_quant_bit = is_smoothquant_ ? 8 : -1; + group_gemm_params.output = gemm1_out; + group_gemm_params.combine_idx = std::nullopt; + gemm1_out = xllm::kernel::group_gemm(group_gemm_params); + } + + // Step 5: activation or scaled quantization(fused with activation) + torch::Tensor act_out; + torch::Tensor act_out_scale; + if (is_smoothquant_) { + int64_t slice_dim = gemm1_out.size(1); + if (is_gated_) slice_dim /= 2; + // slice operation is a view, does not take up extra memory, but points to + // the same memory + act_out = expand_hidden_states.slice(1, 0, slice_dim); + act_out_scale = + selected_expert_info.input_scale.value().slice(0, 0, gemm1_out.size(0)); + // call scaled quantization kernel (also fused with activation) + xllm::kernel::ScaledQuantizeParams scaled_quantize_params; + scaled_quantize_params.x = gemm1_out; + scaled_quantize_params.smooth = act_smooth_; + scaled_quantize_params.token_count = selected_expert_info.token_count_slice; + scaled_quantize_params.output = act_out; + scaled_quantize_params.output_scale = act_out_scale; + scaled_quantize_params.act_mode = hidden_act_; + scaled_quantize_params.active_coef = 1.0; + scaled_quantize_params.is_gated = is_gated_; + scaled_quantize_params.quant_type = torch::kChar; + std::tie(act_out, act_out_scale) = + xllm::kernel::scaled_quantize(scaled_quantize_params); + } else { + act_out = is_gated_ + ? gemm1_out.slice(1, 0, gemm1_out.size(1) / 2).contiguous() + : gemm1_out; + // call activation kernel + xllm::kernel::ActivationParams activation_params; + activation_params.input = gemm1_out; + activation_params.output = act_out; + activation_params.cusum_token_count = + selected_expert_info.cusum_token_count; + activation_params.act_mode = hidden_act_; + activation_params.is_gated = is_gated_; + activation_params.start_expert_id = start_expert_id_; + activation_params.expert_size = expert_size; + xllm::kernel::active(activation_params); + } + + // Step 6: group gemm 2 + torch::Tensor gemm2_out = + create_group_gemm_output(act_out, + w2_, + selected_expert_info.token_count_slice, + hidden_states_dtype, + gemm_workspace); + // ensure the lifespan of these parameters via brace + { + xllm::kernel::GroupGemmParams group_gemm_params; + group_gemm_params.a = act_out; + group_gemm_params.b = w2_; + group_gemm_params.token_count = + selected_expert_info.token_count_slice.to("cpu"); + if (is_smoothquant_) { + group_gemm_params.a_scale = act_out_scale; + group_gemm_params.b_scale = w2_scale_; + } + group_gemm_params.max_dim = group_gemm_max_dim; + group_gemm_params.trans_a = false; + group_gemm_params.trans_b = true; + group_gemm_params.a_quant_bit = is_smoothquant_ ? 8 : -1; + group_gemm_params.output = gemm2_out; + group_gemm_params.combine_idx = selected_expert_info.combine_idx; + gemm2_out = xllm::kernel::group_gemm(group_gemm_params); + } + + // Communciation Step 2: Combine + if (enable_all2all_communication) { + int64_t num_token_expand = hidden_states_2d.size(0) * topk_; + // Delegate pack, layout generation and combine to DeepEP + torch::Tensor combine_send_layout = + deep_ep_->combine_step_pack(gemm2_out, + gather_by_rank_index, + token_sum, + hidden_size_, + hidden_states_dtype); + + // create a wait event for the current stream to finish computation + auto current_stream = device_.current_stream(); + routed_stream_->wait_stream(*current_stream); + // pure communciation kernel: dispatch + { + torch::StreamGuard stream_guard = routed_stream_->set_stream_guard(); + gemm2_out = deep_ep_->combine_step_comm(combine_send_layout, + num_token_expand, + hidden_size_, + hidden_states_dtype); + } + + // pure computation kernel: shared experts + if (n_shared_experts_ > 0) { + shared_stream_->wait_stream(*current_stream); + torch::StreamGuard stream_guard = shared_stream_->set_stream_guard(); + shared_expert_output = shared_experts_(hidden_states); + } + + // join for parallelization + current_stream->wait_stream(*routed_stream_); + if (n_shared_experts_ > 0) { + current_stream->wait_stream(*shared_stream_); + } + } + + // After group gemm is finished, some tensors are no + // longer needed. We must explicitly release the memory. + expand_hidden_states = torch::Tensor(); + selected_expert_info.input_scale = std::nullopt; + act_out = torch::Tensor(); + + // Step 7: combine the intermediate results and get the final hidden states + torch::Tensor final_hidden_states; + // ensure the lifespan of these parameters via brace + { + xllm::kernel::MoeCombineResultParams moe_combine_result_params; + moe_combine_result_params.input = gemm2_out; + moe_combine_result_params.reduce_weight = + selected_expert_info.reduce_weight; + moe_combine_result_params.gather_ids = selected_expert_info.combine_idx; + moe_combine_result_params.cusum_token_count = + selected_expert_info.cusum_token_count; + moe_combine_result_params.start_expert_id = start_expert_id_; + moe_combine_result_params.expert_size = expert_size; + moe_combine_result_params.bias = std::nullopt; + // if all2all communication is enabled and shared output is provided, + // we will fused the add up to combine result + if (enable_all2all_communication && n_shared_experts_ > 0) { + moe_combine_result_params.residual = + shared_expert_output.reshape({-1, shared_expert_output.size(-1)}); + } + + final_hidden_states = + xllm::kernel::moe_combine_result(moe_combine_result_params); + } + + // reshape the final hidden states to the original shape + final_hidden_states = final_hidden_states.reshape(hidden_states_shape); + + if (enable_all2all_communication) { + return final_hidden_states; + } + + // Communciation Step 3: AllReduce for non-all2all communication + // shared experts can be parallelized with the final communication step + // during moe computation. + auto current_stream = device_.current_stream(); + routed_stream_->wait_stream(*current_stream); + { + torch::StreamGuard stream_guard = routed_stream_->set_stream_guard(); + if (tp_pg_->world_size() > 1) { + final_hidden_states = parallel_state::reduce(final_hidden_states, tp_pg_); + } + if (parallel_args_.ep_size() > 1) { + final_hidden_states = parallel_state::reduce( + final_hidden_states, parallel_args_.moe_ep_group_); + } + } + + if (n_shared_experts_ > 0) { + shared_stream_->wait_stream(*current_stream); + torch::StreamGuard stream_guard = shared_stream_->set_stream_guard(); + // for non all2all, we compute the shared experts parallelized with the + // final communication step + shared_expert_output = shared_experts_(hidden_states); + shared_expert_output = + shared_expert_output.reshape({-1, shared_expert_output.size(-1)}); + } + + // join for parallelization + current_stream->wait_stream(*routed_stream_); + if (n_shared_experts_ > 0) { + current_stream->wait_stream(*shared_stream_); + final_hidden_states += shared_expert_output; + } + + return final_hidden_states; +} + +torch::Tensor FusedMoEImpl::forward(const torch::Tensor& hidden_states, + const ModelInputParams& input_params) { + // we only support all2all communication for decode stage for now + bool enable_all2all_communication = + enable_deep_ep_ && std::all_of(input_params.dp_is_decode.begin(), + input_params.dp_is_decode.end(), + [](int32_t val) { return val == 1; }); + + bool is_dp_ep_parallel = + parallel_args_.dp_size() > 1 && parallel_args_.ep_size() > 1; + // during all2all communication, the output has been + // gathered and sliced by dispatch and combine steps, + // so we do not need to gather input and slice output again + bool need_gather_and_slice = + is_dp_ep_parallel && !enable_all2all_communication; + + auto input = hidden_states; + if (need_gather_and_slice) { + input = parallel_state::gather(input, + parallel_args_.dp_local_process_group_, + input_params.dp_global_token_nums); + } + // MoE Gate + auto router_logits = gate_(input); + + // MoE Experts + auto output = + forward_experts(input, router_logits, enable_all2all_communication); + + if (need_gather_and_slice) { + output = get_dp_local_slice(output, input_params, parallel_args_); + } + + return output; +} + +void FusedMoEImpl::load_e_score_correction_bias(const StateDict& state_dict) { + if (e_score_correction_bias_.defined() && + !e_score_correction_bias_is_loaded_) { + LOAD_WEIGHT(e_score_correction_bias); + } +} + +void FusedMoEImpl::load_experts(const StateDict& state_dict) { + const int64_t rank = tp_pg_->rank(); + const int64_t world_size = tp_pg_->world_size(); + const int64_t start_expert_id = start_expert_id_; + const int64_t num_experts_per_rank = num_experts_per_rank_; + const int64_t num_total_experts = num_total_experts_; + std::vector prefixes = {"gate_proj.", "up_proj."}; + if (is_smoothquant_) { + LOAD_MOE_FUSED_WEIGHT("qweight", w1, w3, w13); + LOAD_MOE_FUSED_WEIGHT("per_channel_scale", w1_scale, w3_scale, w13_scale); + // When supporting DeepEP All2All mode, + // we need to load the complete set of expert weights corresponding to + // "up_proj.smooth". Note that even if deep EP mode is not enabled, it + // remains possible to retrieve the smooth quantization information for a + // subset of experts. Therefore, we intentionally do not check whether + // deep_ep_ is enabled in this case. + LOAD_MOE_ALL_EXPERT_WEIGHT("up_proj.", "smooth", input_smooth, -1); + LOAD_MOE_WEIGHT("down_proj.", "qweight", w2, 1); + LOAD_MOE_WEIGHT("down_proj.", "per_channel_scale", w2_scale, -1); + LOAD_MOE_WEIGHT("down_proj.", "smooth", act_smooth, 0); + } else { + LOAD_MOE_FUSED_WEIGHT("weight", w1, w3, w13); + LOAD_MOE_WEIGHT("down_proj.", "weight", w2, 1); + } +} + +void FusedMoEImpl::load_state_dict(const StateDict& state_dict) { + if (state_dict.size() == 0) { + return; + } + + if (n_shared_experts_ > 0) { + shared_experts_->load_state_dict( + state_dict.get_dict_with_prefix("shared_experts.")); + } + gate_->load_state_dict(state_dict.get_dict_with_prefix("gate.")); + load_e_score_correction_bias(state_dict.get_dict_with_prefix("gate.")); + load_experts(state_dict.get_dict_with_prefix("experts.")); +} + +} // namespace layer +} // namespace xllm diff --git a/ex_engine/csrc/ilu_layer_fused_moe.h b/ex_engine/csrc/ilu_layer_fused_moe.h new file mode 100644 index 00000000..3e477064 --- /dev/null +++ b/ex_engine/csrc/ilu_layer_fused_moe.h @@ -0,0 +1,131 @@ +/* Copyright 2025 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. +==============================================================================*/ + +#pragma once + +#include + +#include "framework/model/model_args.h" +#include "framework/model/model_input_params.h" +#include "framework/parallel_state/parallel_args.h" +#include "framework/quant_args.h" +#include "framework/state_dict/state_dict.h" +#include "framework/state_dict/utils.h" +#include "layers/common/deep_ep.h" +#include "layers/common/dense_mlp.h" +#include "layers/common/fused_moe_base.h" +#include "layers/common/linear.h" +#include "platform/device.h" +#include "util/tensor_helper.h" + +namespace xllm { +namespace layer { + +class FusedMoEImpl : public torch::nn::Module { + public: + FusedMoEImpl() = default; + FusedMoEImpl(const ModelArgs& model_args, + const FusedMoEArgs& moe_args, + const QuantArgs& quant_args, + const ParallelArgs& parallel_args, + const torch::TensorOptions& options); + + torch::Tensor forward_experts(const torch::Tensor& hidden_states, + const torch::Tensor& router_logits, + bool enable_all2all_communication); + torch::Tensor forward(const torch::Tensor& hidden_states, + const ModelInputParams& input_params); + void load_state_dict(const StateDict& state_dict); + + private: + // struct to store the selected expert info + struct SelectedExpertInfo { + torch::Tensor reduce_weight; + torch::Tensor combine_idx; + torch::Tensor token_count_slice; + std::optional cusum_token_count; + std::optional input_scale; + }; + + // initial steps for MoE computation, select the experts for each token + torch::Tensor select_experts(const torch::Tensor& hidden_states_2d, + const torch::Tensor& router_logits_2d, + SelectedExpertInfo& selected_expert_info, + bool enable_all2all_communication); + + private: + int64_t num_total_experts_; + int64_t topk_; + int64_t num_expert_group_; + int64_t topk_group_; + double route_scale_; + int64_t hidden_size_; + int64_t n_shared_experts_; + bool is_gated_; + int64_t renormalize_; + std::string hidden_act_; + std::string scoring_func_; + bool is_smoothquant_; + + int64_t num_experts_per_rank_; + int64_t start_expert_id_; + + // Deep EP related parameters + bool enable_deep_ep_; + DeepEPBuffer deep_ep_buffer_; + DeepEPParams deep_ep_params_; + torch::Tensor dispatch_recv_token_tensor_head_; + torch::Tensor dispatch_recv_token_tensor_tail_; + + // steams for parallel shared experts + std::unique_ptr shared_stream_; + std::unique_ptr routed_stream_; + xllm::Device device_; + bool stream_initialized_ = false; + + ReplicatedLinear gate_{nullptr}; + DenseMLP shared_experts_{nullptr}; + DeepEP deep_ep_{nullptr}; + + QuantArgs quant_args_; + ParallelArgs parallel_args_; + torch::TensorOptions options_; + ProcessGroup* tp_pg_; + + DEFINE_WEIGHT(w13); + DEFINE_FUSED_WEIGHT(w1); + DEFINE_FUSED_WEIGHT(w3); + DEFINE_FUSED_WEIGHT(w2); + DEFINE_WEIGHT(e_score_correction_bias); + DEFINE_WEIGHT(w13_scale); + DEFINE_FUSED_WEIGHT(w1_scale); + DEFINE_FUSED_WEIGHT(w3_scale); + DEFINE_FUSED_WEIGHT(w2_scale); + DEFINE_FUSED_WEIGHT(input_smooth); + DEFINE_FUSED_WEIGHT(act_smooth); + + void load_e_score_correction_bias(const StateDict& state_dict); + void load_experts(const StateDict& state_dict); + // create the group gemm output tensor with the workspace + torch::Tensor create_group_gemm_output(const torch::Tensor& a, + const torch::Tensor& b, + const torch::Tensor& group_list, + torch::ScalarType dtype, + torch::Tensor& workspace); +}; +TORCH_MODULE(FusedMoE); + +} // namespace layer +} // namespace xllm diff --git a/ex_engine/csrc/moe/fused_moe_cuda.cpp b/ex_engine/csrc/moe/fused_moe_cuda.cpp new file mode 100644 index 00000000..735e27eb --- /dev/null +++ b/ex_engine/csrc/moe/fused_moe_cuda.cpp @@ -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& 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& fc1_expert_biases, + const std::optional& fc2_expert_biases, + const std::optional& input_sf, + const std::optional& swiglu_alpha, + const std::optional& swiglu_beta, + const std::optional& swiglu_limit, + const std::optional& 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 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(); + + 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>(), // TODO: support + // auto tuning + // profile ids + support_pdl(), + activation_type); + + return result_output; +} +} // namespace xllm::kernel::cuda diff --git a/ex_engine/csrc/moe/moeTopKFuncs.cuh b/ex_engine/csrc/moe/moeTopKFuncs.cuh new file mode 100644 index 00000000..70e21cf8 --- /dev/null +++ b/ex_engine/csrc/moe/moeTopKFuncs.cuh @@ -0,0 +1,257 @@ +/* + * Adapted from + * https://github.com/NVIDIA/TensorRT-LLM/blob/v1.3.0rc2/cpp/tensorrt_llm/kernels/moeTopKFuncs.cuh + * Copyright (c) 2026, The vLLM team. + * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION. All rights + * reserved. SPDX-License-Identifier: Apache-2.0 + * + * 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 + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * 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. + */ +#pragma once + +#include +#include +#include + +namespace vllm { +namespace moe { +namespace reduce_topk { +namespace cg = cooperative_groups; +static constexpr int kWARP_SIZE = 32; + +template +struct TopKRedType { + using T = T_; + static_assert( + std::is_same_v || std::is_same_v || + std::is_same_v || std::is_same_v, + "Top K reduction only implemented for int, float, float16 and bfloat16"); + + using TypeCmp = std::conditional_t; + using IdxT = std::conditional_t; + + static constexpr int kMoveBits = (sizeof(T) == 4) ? 32 : 16; + static constexpr int kMaxIdx = 65535; + TypeCmp compValIdx; + + static __host__ __device__ inline TypeCmp makeCmpVal(T val, int32_t idx = 0) { + auto valueBits = cub::Traits::TwiddleIn( + reinterpret_cast::UnsignedBits&>(val)); + TypeCmp compactTmp = valueBits; + compactTmp = (compactTmp << kMoveBits) | (0xFFFF & (kMaxIdx - idx)); + // Use 65535 minus idx to give higher priority to elements with smaller + // indices. + return compactTmp; + } + + static __host__ __device__ void unpack(T& value, int32_t& index, + TypeCmp cmp) { + // Since “65535-idx” is always smaller than 65536 and positive, we can + // directly use it as the lower 16 bits + index = kMaxIdx - static_cast((cmp & 0xFFFF)); + + auto compactTmp = cmp >> kMoveBits; + auto valueBits = cub::Traits::TwiddleOut( + reinterpret_cast::UnsignedBits&>(compactTmp)); + value = reinterpret_cast(valueBits); + } + + __host__ __device__ TopKRedType() = default; + + __host__ __device__ TopKRedType(T val, int32_t idx) + : compValIdx(makeCmpVal(val, idx)) {} + + __host__ __device__ operator TypeCmp() const noexcept { return compValIdx; } + + __device__ inline TypeCmp reduce( + cg::thread_block_tile const& warp) { + return cg::reduce(warp, compValIdx, cg::greater{}); + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct TopKIdx { + // by default, empty +}; + +template +struct TopKIdx { + static constexpr int K = K_; + int32_t val[K]; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +#define TOPK_SWAP(I, J) \ + { \ + auto pairMin = min(topK[I].compValIdx, topK[J].compValIdx); \ + auto pairMax = max(topK[I].compValIdx, topK[J].compValIdx); \ + topK[I].compValIdx = pairMax; \ + topK[J].compValIdx = pairMin; \ + } + +template +struct Sort; + +template +struct Sort<1, RedType> { + static __device__ void run(RedType* topK) {} +}; + +template +struct Sort<2, RedType> { + static __device__ void run(RedType* topK) { TOPK_SWAP(0, 1); } +}; + +template +struct Sort<3, RedType> { + static __device__ void run(RedType* topK) { + TOPK_SWAP(0, 1); + TOPK_SWAP(1, 2); + TOPK_SWAP(0, 1); + } +}; + +template +struct Sort<4, RedType> { + static __device__ void run(RedType* topK) { + TOPK_SWAP(0, 2); + TOPK_SWAP(1, 3); + TOPK_SWAP(0, 1); + TOPK_SWAP(2, 3); + TOPK_SWAP(1, 2); + } +}; + +template +__forceinline__ __device__ void reduceTopK( + cg::thread_block_tile const& warp, Type (&out)[K], + int32_t (&outIdx)[K], Type value, int32_t idx, Type const minValue, + int actualK = K) { + static_assert(K > 0, "Top K must have K > 0"); + static_assert(K < kWARP_SIZE, "Top K must have K < kWARP_SIZE"); + using RedType = TopKRedType; + RedType topK{value, idx}; + typename RedType::TypeCmp packedMax{}; +#pragma unroll + for (int kk = 0; kk < actualK; ++kk) { + topK = + kk > 0 && packedMax == topK.compValIdx ? RedType{minValue, idx} : topK; + // get the next largest value + packedMax = topK.reduce(warp); + RedType::unpack(out[kk], outIdx[kk], packedMax); + } +}; + +template +__device__ void reduceTopKFunc(cg::thread_block_tile const& warp, + Type (&out)[K], int32_t (&outIdx)[K], + Type (&value)[N], int32_t (&idx)[N], + Type minValue, int actualK = K) { + static_assert(K > 0, "Top K must have K > 0"); + static_assert(K < kWARP_SIZE, "Top K must have K < kWARP_SIZE"); + static_assert(N > 0, "Top K must have N > 0"); + static_assert(N < 5, + "Only support candidates number less than or equal to 128"); + using RedType = TopKRedType; + RedType topK[N]; +#pragma unroll + for (int nn = 0; nn < N; ++nn) { + topK[nn] = RedType{value[nn], idx[nn]}; + } + + if constexpr (!IsSorted) { + Sort::run(topK); + } + typename RedType::TypeCmp packedMax{}; +#pragma unroll + for (int kk = 0; kk < actualK; ++kk) { + bool update = kk > 0 && packedMax == topK[0].compValIdx; +#pragma unroll + for (int nn = 0; nn < N; ++nn) { + topK[nn] = update && nn == N - 1 ? RedType{minValue, idx[nn]} + : update ? topK[nn + 1] + : topK[nn]; + } + // get the next largest value + packedMax = topK[0].reduce(warp); + RedType::unpack(out[kk], outIdx[kk], packedMax); + } +}; + +template +__forceinline__ __device__ void reduceTopK( + cg::thread_block_tile const& warp, Type (&out)[K], + int32_t (&outIdx)[K], Type (&value)[N], int32_t (&idx)[N], + Type const minValue, int actualK = K) { + static_assert(K > 0, "Top K must have K > 0"); + static_assert(K < kWARP_SIZE, "Top K must have K < kWARP_SIZE"); + static_assert(N > 0, "Top K must have N > 0"); + static_assert( + N <= 16, + "Only support candidates number less than or equal to 16*32=512"); + static_assert(N <= 4 || N % 4 == 0, + "Only support candidates number is a multiple of 4*32=128 or " + "less than or equal to 4"); + using RedType = TopKRedType; + + if constexpr (N <= 4) { + reduceTopKFunc(warp, out, outIdx, value, idx, minValue, + actualK); + } else { + constexpr int numLoops = N / 4; + constexpr int numResults = (numLoops * K - 1) / kWARP_SIZE + 1; + + Type topKBufferValue[numResults]; + int32_t topKBufferIdx[numResults]; + int32_t laneIdx = threadIdx.x % kWARP_SIZE; + + for (int ii = 0; ii < numResults; ++ii) { + topKBufferValue[ii] = minValue; + topKBufferIdx[ii] = ii * kWARP_SIZE - 1; + } + for (int loop = 0; loop < numLoops; ++loop) { + int start = loop * 4; + Type topKValue[K]; + int32_t topKIdx[K]; + Type inValue[4]; + int32_t inIdx[4]; + for (int i = 0; i < 4; ++i) { + inValue[i] = value[start + i]; + inIdx[i] = idx[start + i]; + } + reduceTopKFunc(warp, topKValue, topKIdx, inValue, inIdx, + minValue, actualK); + int inOffset = laneIdx % K; + if (laneIdx >= loop * K && laneIdx < (loop + 1) * K) { + topKBufferValue[0] = topKValue[inOffset]; + topKBufferIdx[0] = topKIdx[inOffset]; + } + if (loop == numLoops - 1 && (laneIdx < (numLoops * K - kWARP_SIZE))) { + topKBufferValue[1] = topKValue[inOffset]; + topKBufferIdx[1] = topKIdx[inOffset]; + } + } + + reduceTopKFunc(warp, out, outIdx, topKBufferValue, + topKBufferIdx, minValue, actualK); + } +}; + +#undef TOPK_SWAP + +} // namespace reduce_topk +} // namespace moe +} // namespace vllm diff --git a/ex_engine/csrc/moe/moe_align_sum_kernels.cu b/ex_engine/csrc/moe/moe_align_sum_kernels.cu new file mode 100644 index 00000000..d7c68ff2 --- /dev/null +++ b/ex_engine/csrc/moe/moe_align_sum_kernels.cu @@ -0,0 +1,833 @@ +#include +#include + +#include +#include +#include +#include +#include +#include + +#include "../../cuda_compat.h" +#include "core/math.hpp" +#include "libtorch_stable/dispatch_utils.h" +#include "libtorch_stable/torch_utils.h" + +#define CEILDIV(x, y) (((x) + (y) - 1) / (y)) + +namespace vllm { +namespace moe { +namespace batched_moe_align_block_size { + +// Note num_threads needs to be 1024 for BlockScan Reduction in the kernel. +static constexpr int32_t num_threads = 1024; +static constexpr int32_t num_blocks = 1; +__global__ void batched_moe_align_block_size_kernel( + int32_t const num_batches, int32_t const max_tokens_per_batch, + int32_t const block_size, int32_t const* __restrict__ batch_num_tokens, + int32_t* __restrict__ sorted_ids, int32_t* __restrict__ block_ids, + int32_t* __restrict__ num_tokens_post_pad) { + // TODO(varun): This is a naive implementation. Could be optimized. + + size_t const batch_id = threadIdx.x; + size_t const stride = blockDim.x * gridDim.x; + int32_t const num_blocks_per_batch = + CEILDIV(max_tokens_per_batch, block_size); + int32_t const sorted_ids_size = + num_blocks_per_batch * num_batches * block_size; + int32_t const block_ids_size = sorted_ids_size / block_size; + int32_t const SENTINEL = + num_batches * max_tokens_per_batch; // To denote invalid entries. + // Initialize sorted_ids + for (size_t i = threadIdx.x; i < sorted_ids_size; i += stride) { + sorted_ids[i] = SENTINEL; + } + // Initialize expert_ids with -1 + for (size_t i = threadIdx.x; i < block_ids_size; i += stride) { + block_ids[i] = -1; + } + + int32_t b_num_tokens = 0; + if (batch_id < num_batches) { + b_num_tokens = batch_num_tokens[batch_id]; + } + int32_t const ceil_b_num_tokens = + CEILDIV(b_num_tokens, block_size) * block_size; + + // Compute prefix sum over token counts per expert + using BlockScan = cub::BlockScan; + __shared__ typename BlockScan::TempStorage temp_storage; + int cumsum_val; + BlockScan(temp_storage).ExclusiveSum(ceil_b_num_tokens, cumsum_val); + __syncthreads(); + + bool const is_last_batch = batch_id == (num_batches - 1); + if (is_last_batch) { + *num_tokens_post_pad = cumsum_val + ceil_b_num_tokens; + } + + if (batch_id < num_batches) { + int32_t const batch_offset = batch_id * max_tokens_per_batch; + for (size_t i = 0; i < b_num_tokens; ++i) { + sorted_ids[cumsum_val + i] = batch_offset + i; + } + + int32_t const block_start = cumsum_val / block_size; + int32_t const num_blocks = ceil_b_num_tokens / block_size; + for (size_t i = 0; i < num_blocks; ++i) { + block_ids[block_start + i] = batch_id; + } + } +} +} // namespace batched_moe_align_block_size + +template +__device__ void _moe_align_block_size( + const scalar_t* __restrict__ topk_ids, + int32_t* __restrict__ sorted_token_ids, int32_t* __restrict__ expert_ids, + int32_t* __restrict__ total_tokens_post_pad, + int32_t* __restrict__ expert_map, int32_t num_experts, + int32_t padded_num_experts, int32_t experts_per_warp, int32_t block_size, + size_t numel, int32_t* __restrict__ cumsum, int32_t max_num_tokens_padded, + int32_t max_num_m_blocks, int32_t model_offset, int32_t inactive_expert_id, + int32_t topk_num, int32_t* token_mask, bool has_expert_map) { + extern __shared__ int32_t shared_counts[]; + + // Compute input buffer offsets. Typically these will all be 0, except when + // using Multi LoRA. + int sorted_token_ids_offset = max_num_tokens_padded * model_offset; + int expert_ids_offset = max_num_m_blocks * model_offset; + int cumsum_offset = (num_experts + 1) * model_offset; + + // Use separate threadblocks to fill sorted_token_ids. + // This is safe since the current kernel does not use sorted_token_ids. + if (blockIdx.x % 2) { + // Initialize sorted_token_ids with numel + for (size_t it = threadIdx.x; it < max_num_tokens_padded; + it += blockDim.x) { + sorted_token_ids[sorted_token_ids_offset + it] = numel; + } + return; + } + + const int warp_id = threadIdx.x / WARP_SIZE; + const int my_expert_start = warp_id * experts_per_warp; + + for (int i = 0; i < experts_per_warp; ++i) { + if (my_expert_start + i < padded_num_experts) { + shared_counts[warp_id * experts_per_warp + i] = 0; + } + } + + __syncthreads(); + + const size_t tid = threadIdx.x; + const size_t stride = blockDim.x; + + for (size_t i = tid; i < numel; i += stride) { + int expert_id = topk_ids[i]; + if (expert_id >= num_experts) { + continue; + } + if (has_expert_map) { + expert_id = expert_map[expert_id]; + // filter invalid experts + if (expert_id == -1) continue; + } + int warp_idx = expert_id / experts_per_warp; + int expert_offset = expert_id % experts_per_warp; + int mask = token_mask == nullptr ? 1 : token_mask[i / topk_num]; + atomicAdd(&shared_counts[warp_idx * experts_per_warp + expert_offset], + mask); + } + + __syncthreads(); + + // Compute prefix sum over token counts per expert + using BlockScan = cub::BlockScan; + __shared__ typename BlockScan::TempStorage temp_storage; + + int expert_count = 0; + int expert_id = threadIdx.x; + if (expert_id < num_experts) { + int warp_idx = expert_id / experts_per_warp; + int expert_offset = expert_id % experts_per_warp; + expert_count = shared_counts[warp_idx * experts_per_warp + expert_offset]; + expert_count = CEILDIV(expert_count, block_size) * block_size; + } + + int cumsum_val; + BlockScan(temp_storage).ExclusiveSum(expert_count, cumsum_val); + if (expert_id <= num_experts) { + cumsum[cumsum_offset + expert_id] = cumsum_val; + } + + if (expert_id == num_experts) { + total_tokens_post_pad[model_offset] = cumsum_val; + } + + __syncthreads(); + + if (threadIdx.x < num_experts) { + for (int i = cumsum[cumsum_offset + threadIdx.x]; + i < cumsum[cumsum_offset + threadIdx.x + 1]; i += block_size) { + expert_ids[expert_ids_offset + i / block_size] = threadIdx.x; + } + } + + // Fill remaining expert_ids with -1 + const size_t fill_start_idx = + cumsum[cumsum_offset + num_experts] / block_size + threadIdx.x; + for (size_t i = fill_start_idx; i < max_num_m_blocks; i += blockDim.x) { + expert_ids[expert_ids_offset + i] = inactive_expert_id; + } +} + +template +__device__ void _moe_align_block_size_small_batch_expert( + const scalar_t* __restrict__ topk_ids, + int32_t* __restrict__ sorted_token_ids, int32_t* __restrict__ expert_ids, + int32_t* __restrict__ total_tokens_post_pad, + int32_t* __restrict__ expert_map, int32_t num_experts, int32_t block_size, + size_t numel, int32_t max_num_tokens_padded, int32_t max_num_m_blocks, + int32_t inactive_expert_id, int32_t model_offset, int32_t topk_num, + int32_t* token_mask, bool has_expert_map) { + // Compute input buffer offsets. Typically these will all be 0, except when + // using Multi LoRA. + int sorted_token_ids_offset = max_num_tokens_padded * model_offset; + int expert_ids_offset = max_num_m_blocks * model_offset; + + // Use an additional group of threads to fill sorted_token_ids. + // Since the current kernel will use sorted_token_ids afterward, + // we fill sorted_token_ids within the same threadblock to make + // synchronization easier. + if (threadIdx.x < fill_threads) { + // Initialize sorted_token_ids with numel + for (size_t it = threadIdx.x; it < max_num_tokens_padded; + it += fill_threads) { + sorted_token_ids[sorted_token_ids_offset + it] = numel; + } + // Three __syncthreads() corresponding to the other threads + __syncthreads(); + __syncthreads(); + __syncthreads(); + return; + } + + const size_t tid = threadIdx.x - fill_threads; + const size_t stride = blockDim.x - fill_threads; + + extern __shared__ int32_t shared_mem[]; + int32_t* cumsum = shared_mem; + int32_t* tokens_cnts = (int32_t*)(shared_mem + num_experts + 1); + + for (int i = 0; i < num_experts; ++i) { + tokens_cnts[(tid + 1) * num_experts + i] = 0; + } + + for (size_t i = tid; i < numel; i += stride) { + int32_t expert_id = topk_ids[i]; + if (has_expert_map) { + expert_id = expert_map[expert_id]; + // filter invalid expert + if (expert_id == -1) continue; + } + int mask = token_mask == nullptr ? 1 : token_mask[i / topk_num]; + tokens_cnts[(tid + 1) * num_experts + expert_id] += mask; + } + + __syncthreads(); + + if (tid < num_experts) { + tokens_cnts[tid] = 0; + for (int i = 1; i <= stride; ++i) { + tokens_cnts[i * num_experts + tid] += + tokens_cnts[(i - 1) * num_experts + tid]; + } + } + + __syncthreads(); + + if (tid == 0) { + cumsum[0] = 0; + for (int i = 1; i <= num_experts; ++i) { + cumsum[i] = + cumsum[i - 1] + + CEILDIV(tokens_cnts[stride * num_experts + i - 1], block_size) * + block_size; + } + total_tokens_post_pad[model_offset] = + static_cast(cumsum[num_experts]); + } + + __syncthreads(); + + if (tid < num_experts) { + for (int i = cumsum[tid]; i < cumsum[tid + 1]; i += block_size) { + expert_ids[expert_ids_offset + i / block_size] = tid; + } + } + + // Fill remaining expert_ids with -1 + const size_t fill_start_idx = cumsum[num_experts] / block_size + tid; + for (size_t i = fill_start_idx; i < max_num_m_blocks; i += stride) { + expert_ids[expert_ids_offset + i] = inactive_expert_id; + } + + for (size_t i = tid; i < numel; i += stride) { + int32_t expert_id = topk_ids[i]; + if (has_expert_map) { + expert_id = expert_map[expert_id]; + // filter invalid expert + if (expert_id == -1) continue; + } + int32_t rank_post_pad = + tokens_cnts[tid * num_experts + expert_id] + cumsum[expert_id]; + + if (token_mask == nullptr || token_mask[i / topk_num]) { + sorted_token_ids[sorted_token_ids_offset + rank_post_pad] = i; + ++tokens_cnts[tid * num_experts + expert_id]; + } + } +} + +template +__device__ void _count_and_sort_expert_tokens( + const scalar_t* __restrict__ topk_ids, + int32_t* __restrict__ sorted_token_ids, int32_t* __restrict__ cumsum_buffer, + int32_t* __restrict__ expert_map, size_t numel, int32_t num_experts, + int32_t max_num_tokens_padded, int32_t* __restrict__ token_mask, + int32_t model_offset, int32_t topk_num, bool has_expert_map) { + const size_t tid = blockIdx.y * blockDim.x + threadIdx.x; + const size_t stride = blockDim.x * gridDim.y; + + for (size_t i = tid; i < numel; i += stride) { + int32_t expert_id = topk_ids[i]; + if (expert_id >= num_experts) { + continue; + } + + if (has_expert_map) { + expert_id = expert_map[expert_id]; + // filter invalid experts + if (expert_id == -1) continue; + } + + if (token_mask == nullptr || token_mask[i / topk_num]) { + int32_t rank_post_pad = atomicAdd( + &cumsum_buffer[(model_offset * (num_experts + 1)) + expert_id], 1); + sorted_token_ids[max_num_tokens_padded * model_offset + rank_post_pad] = + i; + } + } +} + +template +__global__ void moe_align_block_size_kernel( + const scalar_t* __restrict__ topk_ids, + int32_t* __restrict__ sorted_token_ids, int32_t* __restrict__ expert_ids, + int32_t* __restrict__ total_tokens_post_pad, + int32_t* __restrict__ expert_map, int32_t num_experts, + int32_t padded_num_experts, int32_t experts_per_warp, int32_t block_size, + size_t numel, int32_t* __restrict__ cumsum, int32_t max_num_tokens_padded, + int32_t topk_num, bool has_expert_map) { + _moe_align_block_size( + topk_ids, sorted_token_ids, expert_ids, total_tokens_post_pad, expert_map, + num_experts, padded_num_experts, experts_per_warp, block_size, numel, + cumsum, max_num_tokens_padded, CEILDIV(max_num_tokens_padded, block_size), + 0, -1, topk_num, nullptr, has_expert_map); +} + +template +__global__ void count_and_sort_expert_tokens_kernel( + const scalar_t* __restrict__ topk_ids, + int32_t* __restrict__ sorted_token_ids, int32_t* __restrict__ cumsum_buffer, + int32_t* __restrict__ expert_map, size_t numel, int32_t num_experts, + int32_t max_num_tokens_padded, int32_t topk_num, bool has_expert_map) { + _count_and_sort_expert_tokens( + topk_ids, sorted_token_ids, cumsum_buffer, expert_map, numel, num_experts, + max_num_tokens_padded, nullptr, 0, topk_num, has_expert_map); +} + +template +__global__ void moe_sum_kernel( + scalar_t* __restrict__ out, // [..., d] + const scalar_t* __restrict__ input, // [..., topk, d] + const int d) { + const int64_t token_idx = blockIdx.x; + for (int64_t idx = threadIdx.x; idx < d; idx += blockDim.x) { + scalar_t x = 0.0; +#pragma unroll + for (int k = 0; k < TOPK; ++k) { + x += VLLM_LDG(&input[token_idx * TOPK * d + k * d + idx]); + } + out[token_idx * d + idx] = x; + } +} + +template +__global__ void moe_align_block_size_small_batch_expert_kernel( + const scalar_t* __restrict__ topk_ids, + int32_t* __restrict__ sorted_token_ids, int32_t* __restrict__ expert_ids, + int32_t* __restrict__ total_tokens_post_pad, + int32_t* __restrict__ expert_map, int32_t num_experts, int32_t block_size, + size_t numel, int32_t max_num_tokens_padded, int32_t topk_num, + bool has_expert_map) { + _moe_align_block_size_small_batch_expert( + topk_ids, sorted_token_ids, expert_ids, total_tokens_post_pad, expert_map, + num_experts, block_size, numel, max_num_tokens_padded, + CEILDIV(max_num_tokens_padded, block_size), -1, 0, topk_num, nullptr, + has_expert_map); +} + +template +__global__ void moe_lora_align_block_size_kernel( + scalar_t* __restrict__ topk_ids, int32_t* __restrict__ token_lora_mapping, + int64_t block_size, int32_t* __restrict__ expert_map, int num_experts, + int max_loras, size_t numel, int max_num_tokens_padded, + int max_num_m_blocks, int32_t* __restrict__ sorted_token_ids, + int32_t* __restrict__ expert_ids, int32_t topk_num, + int32_t* total_tokens_post_pad, int32_t* adapter_enabled, + int32_t* __restrict__ cumsum, int32_t experts_per_warp, + int32_t padded_num_experts, int32_t* lora_ids, + int32_t* __restrict__ token_mask, bool has_expert_map) { + int lora_idx = blockIdx.x / 2; + int lora_id = lora_ids[lora_idx]; + // Output buffers are indexed by lora_id (in [0, max_loras)). The grid + // iterates one extra slot to accommodate the "-1" entry that + // active_lora_ids may hold in position 0 for mixed base + LoRA batches; + // guard against any other unexpected lora_id >= max_loras to avoid + // out-of-bounds writes. This mirrors the `lora_id >= max_loras` guard in + // the Triton _fused_moe_lora_kernel. + if (lora_id == -1 || lora_id >= max_loras || adapter_enabled[lora_id] == 0) { + return; + } + + // Populate the token_mask based on the token-LoRA mapping + int num_tokens = numel / topk_num; + if (threadIdx.x == 0) { + total_tokens_post_pad[lora_id] = 0; + + for (int i = 0; i < num_tokens; i++) { + token_mask[(lora_id * num_tokens) + i] = + (int)token_lora_mapping[i] == lora_id; + } + } + + __syncthreads(); + + _moe_align_block_size( + topk_ids, sorted_token_ids, expert_ids, total_tokens_post_pad, expert_map, + num_experts, padded_num_experts, experts_per_warp, block_size, numel, + cumsum, max_num_tokens_padded, max_num_m_blocks, lora_id, -1, topk_num, + &token_mask[(lora_id * num_tokens)], has_expert_map); +} + +template +__global__ void lora_count_and_sort_expert_tokens_kernel( + const scalar_t* __restrict__ topk_ids, + int32_t* __restrict__ sorted_token_ids, int32_t* __restrict__ cumsum_buffer, + int32_t* __restrict__ expert_map, size_t numel, int32_t num_experts, + int32_t max_num_tokens_padded, int32_t topk_num, int32_t* token_mask, + int32_t max_loras, int32_t* lora_ids, int32_t* adapter_enabled, + bool has_expert_map) { + int lora_idx = blockIdx.x; + int lora_id = lora_ids[lora_idx]; + // Same guard rationale as moe_lora_align_block_size_kernel. Additionally + // skip disabled adapter slots: moe_lora_align_block_size_kernel early-returns + // for them and leaves token_mask[lora_id, :] uninitialized (token_mask is + // allocated with torch::empty), so running the sort loop here would traverse + // garbage mask bits and pollute this slot's rows of sorted_token_ids and + // cumsum_buffer. Downstream consumers already skip disabled slots, so the + // pollution is dormant today, but the check keeps behavior symmetric with + // the other two align kernels and avoids O(numel) wasted work per disabled + // slot. Short-circuit evaluation ensures adapter_enabled is only indexed + // after lora_id is confirmed to be in [0, max_loras). + if (lora_id == -1 || lora_id >= max_loras || adapter_enabled[lora_id] == 0) { + return; + } + + int num_tokens = numel / topk_num; + + _count_and_sort_expert_tokens( + topk_ids, sorted_token_ids, cumsum_buffer, expert_map, numel, num_experts, + max_num_tokens_padded, &token_mask[(lora_id * num_tokens)], lora_id, + topk_num, has_expert_map); +} + +template +__global__ void moe_lora_align_block_size_small_batch_expert_kernel( + scalar_t* __restrict__ topk_ids, int32_t* token_lora_mapping, + int64_t block_size, int32_t* __restrict__ expert_map, int num_experts, + int max_loras, size_t numel, int max_num_tokens_padded, + int max_num_m_blocks, int32_t* __restrict__ sorted_token_ids, + int32_t* __restrict__ expert_ids, int topk_num, + int32_t* total_tokens_post_pad, int32_t* adapter_enabled, int32_t* lora_ids, + int32_t* token_mask, bool has_expert_map) { + int lora_idx = blockIdx.x; + int lora_id = lora_ids[lora_idx]; + // Same guard rationale as moe_lora_align_block_size_kernel. + if (lora_id == -1 || lora_id >= max_loras || adapter_enabled[lora_id] == 0) { + return; + } + + int num_tokens = numel / topk_num; + if (threadIdx.x == 0) { + total_tokens_post_pad[lora_id] = 0; + + for (int i = 0; i < num_tokens; i++) { + token_mask[(lora_id * num_tokens) + i] = + (int)token_lora_mapping[i] == lora_id; + } + } + + __syncthreads(); + + _moe_align_block_size_small_batch_expert( + topk_ids, sorted_token_ids, expert_ids, total_tokens_post_pad, expert_map, + num_experts, block_size, numel, max_num_tokens_padded, max_num_m_blocks, + -1, lora_id, topk_num, &token_mask[(lora_id * num_tokens)], + has_expert_map); +} + +} // namespace moe +} // namespace vllm + +// taken from +// https://github.com/sgl-project/sglang/blob/8b5f83ed3b7d2a49ad5c5cd5aa61c5d502f47dbc +void moe_align_block_size( + torch::stable::Tensor topk_ids, int64_t num_experts, int64_t block_size, + torch::stable::Tensor sorted_token_ids, torch::stable::Tensor experts_ids, + torch::stable::Tensor num_tokens_post_pad, + std::optional maybe_expert_map) { + const cudaStream_t stream = + get_current_cuda_stream(topk_ids.get_device_index()); + + int64_t padded_num_experts = + ((num_experts + WARP_SIZE - 1) / WARP_SIZE) * WARP_SIZE; + int experts_per_warp = WARP_SIZE; + int threads = 1024; + threads = ((threads + WARP_SIZE - 1) / WARP_SIZE) * WARP_SIZE; + + // BlockScan uses 1024 threads and assigns one thread per expert. + STD_TORCH_CHECK(padded_num_experts < 1024, + "padded_num_experts must be less than 1024"); + bool has_expert_map = maybe_expert_map.has_value(); + torch::stable::Tensor expert_map; + if (has_expert_map) { + expert_map = maybe_expert_map.value(); + } else { + expert_map = torch::stable::new_empty(topk_ids, {0}, + torch::headeronly::ScalarType::Int); + } + + VLLM_STABLE_DISPATCH_INTEGRAL_AND_UNSIGNED_TYPES( + topk_ids.scalar_type(), "moe_align_block_size_kernel", [&] { + // calc needed amount of shared mem for `cumsum` tensors + bool small_batch_expert_mode = + (topk_ids.numel() < 1024) && (num_experts <= 64); + + if (small_batch_expert_mode) { + const int32_t threads = max((int32_t)num_experts, WARP_SIZE); + const int32_t shared_mem_size = + ((threads + 1) * num_experts + (num_experts + 1)) * + sizeof(int32_t); + + // threadIdx.x >= fill_threads: counting experts and aligning + // threadIdx.x < fill_threads: filling sorted_token_ids + constexpr int32_t fill_threads = 256; + auto small_batch_expert_kernel = + vllm::moe::moe_align_block_size_small_batch_expert_kernel< + scalar_t, fill_threads>; + small_batch_expert_kernel<<<1, fill_threads + threads, + shared_mem_size, stream>>>( + reinterpret_cast(topk_ids.const_data_ptr()), + reinterpret_cast(sorted_token_ids.mutable_data_ptr()), + reinterpret_cast(experts_ids.mutable_data_ptr()), + reinterpret_cast( + num_tokens_post_pad.mutable_data_ptr()), + reinterpret_cast(expert_map.mutable_data_ptr()), + num_experts, block_size, topk_ids.numel(), + sorted_token_ids.size(0), topk_ids.size(1), has_expert_map); + } else { + torch::stable::Tensor cumsum_buffer = torch::stable::new_empty( + topk_ids, {num_experts + 1}, torch::headeronly::ScalarType::Int); + auto align_kernel = vllm::moe::moe_align_block_size_kernel; + + size_t num_warps = CEILDIV(padded_num_experts, experts_per_warp); + size_t shared_mem_size = + num_warps * experts_per_warp * sizeof(int32_t); + + // launch two threadblocks + // blockIdx.x == 0: counting experts and aligning + // blockIdx.x == 1: filling sorted_token_ids + align_kernel<<<2, threads, shared_mem_size, stream>>>( + reinterpret_cast(topk_ids.const_data_ptr()), + reinterpret_cast(sorted_token_ids.mutable_data_ptr()), + reinterpret_cast(experts_ids.mutable_data_ptr()), + reinterpret_cast( + num_tokens_post_pad.mutable_data_ptr()), + reinterpret_cast(expert_map.mutable_data_ptr()), + num_experts, padded_num_experts, experts_per_warp, block_size, + topk_ids.numel(), + reinterpret_cast(cumsum_buffer.mutable_data_ptr()), + sorted_token_ids.size(0), topk_ids.size(1), has_expert_map); + + const int block_threads = std::min(256, (int)threads); + const int num_blocks = + (topk_ids.numel() + block_threads - 1) / block_threads; + const int max_blocks = 65535; + const int actual_blocks = std::min(num_blocks, max_blocks); + dim3 gridDims(1, actual_blocks); + + auto sort_kernel = + vllm::moe::count_and_sort_expert_tokens_kernel; + sort_kernel<<>>( + reinterpret_cast(topk_ids.const_data_ptr()), + reinterpret_cast(sorted_token_ids.mutable_data_ptr()), + reinterpret_cast(cumsum_buffer.mutable_data_ptr()), + reinterpret_cast(expert_map.mutable_data_ptr()), + topk_ids.numel(), num_experts, sorted_token_ids.size(0), + topk_ids.size(1), has_expert_map); + } + }); +} + +void batched_moe_align_block_size(int64_t max_tokens_per_batch, + int64_t block_size, + const torch::stable::Tensor& batch_num_tokens, + torch::stable::Tensor sorted_ids, + torch::stable::Tensor batch_ids, + torch::stable::Tensor num_tokens_post_pad) { + namespace batched_kernel = vllm::moe::batched_moe_align_block_size; + + const cudaStream_t stream = + get_current_cuda_stream(batch_num_tokens.get_device_index()); + int32_t const B = batch_num_tokens.size(0); + int32_t const num_blocks_per_batch = + round_to_next_multiple_of(max_tokens_per_batch, block_size) / block_size; + int32_t const num_blocks = num_blocks_per_batch * B; + int64_t const sorted_ids_size = num_blocks * block_size; + + STD_TORCH_CHECK(sorted_ids.size(0) == sorted_ids_size); + STD_TORCH_CHECK(batch_ids.size(0) == sorted_ids_size / block_size); + STD_TORCH_CHECK(num_tokens_post_pad.size(0) == 1); + STD_TORCH_CHECK(B <= batched_kernel::num_threads); + + batched_kernel::batched_moe_align_block_size_kernel<<< + batched_kernel::num_blocks, batched_kernel::num_threads, 0, stream>>>( + B, max_tokens_per_batch, block_size, + reinterpret_cast(batch_num_tokens.const_data_ptr()), + reinterpret_cast(sorted_ids.mutable_data_ptr()), + reinterpret_cast(batch_ids.mutable_data_ptr()), + reinterpret_cast(num_tokens_post_pad.mutable_data_ptr())); +} + +void moe_sum(torch::stable::Tensor& input, // [num_tokens, topk, hidden_size] + torch::stable::Tensor& output) // [num_tokens, hidden_size] +{ + const int hidden_size = input.size(-1); + const auto num_tokens = output.numel() / hidden_size; + const int topk = input.size(1); + + dim3 grid(num_tokens); + dim3 block(std::min(hidden_size, 1024)); + const torch::stable::accelerator::DeviceGuard device_guard( + output.get_device_index()); + const cudaStream_t stream = + get_current_cuda_stream(output.get_device_index()); + + switch (topk) { + case 2: + VLLM_STABLE_DISPATCH_FLOATING_TYPES( + input.scalar_type(), "moe_sum_kernel", [&] { + vllm::moe::moe_sum_kernel<<>>( + reinterpret_cast(output.mutable_data_ptr()), + reinterpret_cast(input.const_data_ptr()), + hidden_size); + }); + break; + + case 3: + VLLM_STABLE_DISPATCH_FLOATING_TYPES( + input.scalar_type(), "moe_sum_kernel", [&] { + vllm::moe::moe_sum_kernel<<>>( + reinterpret_cast(output.mutable_data_ptr()), + reinterpret_cast(input.const_data_ptr()), + hidden_size); + }); + break; + + case 4: + VLLM_STABLE_DISPATCH_FLOATING_TYPES( + input.scalar_type(), "moe_sum_kernel", [&] { + vllm::moe::moe_sum_kernel<<>>( + reinterpret_cast(output.mutable_data_ptr()), + reinterpret_cast(input.const_data_ptr()), + hidden_size); + }); + break; + + default: + torch::stable::sum_out(output, input, std::array{1}); + break; + } +} + +void moe_lora_align_block_size( + torch::stable::Tensor topk_ids, torch::stable::Tensor token_lora_mapping, + int64_t num_experts, int64_t block_size, int64_t max_loras, + int64_t max_num_tokens_padded, int64_t max_num_m_blocks, + torch::stable::Tensor sorted_token_ids, torch::stable::Tensor expert_ids, + torch::stable::Tensor num_tokens_post_pad, + torch::stable::Tensor adapter_enabled, torch::stable::Tensor lora_ids, + std::optional maybe_expert_map) { + const int topk_num = topk_ids.size(1); + + STD_TORCH_CHECK(block_size > 0, "block_size should be greater than 0. "); + + int device_max_shared_mem; + int dev = topk_ids.get_device_index(); + cudaDeviceGetAttribute(&device_max_shared_mem, + cudaDevAttrMaxSharedMemoryPerBlockOptin, dev); + const cudaStream_t stream = get_current_cuda_stream(dev); + + int64_t padded_num_experts = + ((num_experts + WARP_SIZE - 1) / WARP_SIZE) * WARP_SIZE; + + // BlockScan uses 1024 threads and assigns one thread per expert. + STD_TORCH_CHECK(padded_num_experts < 1024, + "padded_num_experts must be less than 1024"); + + torch::stable::Tensor token_mask = + torch::stable::new_empty(topk_ids, {max_loras * topk_ids.size(0)}, + torch::headeronly::ScalarType::Int); + bool has_expert_map = maybe_expert_map.has_value(); + torch::stable::Tensor expert_map; + if (has_expert_map) { + expert_map = maybe_expert_map.value(); + } else { + expert_map = torch::stable::new_empty(topk_ids, {0}, + torch::headeronly::ScalarType::Int); + } + + VLLM_STABLE_DISPATCH_INTEGRAL_TYPES( + topk_ids.scalar_type(), "moe_lora_align_sum_kernel", [&] { + bool small_batch_expert_mode = + (topk_ids.numel() < 1024) && (num_experts <= 64); + + if (small_batch_expert_mode) { + const int32_t num_thread = max((int32_t)num_experts, 128); + const int32_t shared_mem = + (num_thread + 1) * num_experts * sizeof(int32_t) + + (num_experts + 1) * sizeof(int32_t); + if (shared_mem > device_max_shared_mem) { + STD_TORCH_CHECK(false, "Shared memory usage exceeds device limit."); + } + + // threadIdx.x >= fill_threads: counting experts and aligning + // threadIdx.x < fill_threads: filling sorted_token_ids + constexpr int32_t fill_threads = 256; + + dim3 blockDim(num_thread + fill_threads); + auto kernel = + vllm::moe::moe_lora_align_block_size_small_batch_expert_kernel< + scalar_t, fill_threads>; + STD_CUDA_CHECK(VLLM_DevFuncAttribute_SET_MaxDynamicSharedMemorySize( + (void*)kernel, shared_mem)); + // Grid size is (max_loras + 1) because active_lora_ids has length + // max_loras + 1: sorted-unique values of token_lora_mapping, which + // can include -1 (base-model tokens) in addition to up to max_loras + // real LoRA slots. Using max_loras would drop the real LoRA slot + // when -1 is present at position 0 and leave output buffers + // uninitialized, causing illegal memory accesses in downstream + // MoE-LoRA kernels. This mirrors the fix made for the Triton + // _fused_moe_lora_kernel grid in vllm-project/vllm#32277. + kernel<<>>( + reinterpret_cast(topk_ids.mutable_data_ptr()), + reinterpret_cast(token_lora_mapping.mutable_data_ptr()), + block_size, + reinterpret_cast(expert_map.mutable_data_ptr()), + num_experts, max_loras, topk_ids.numel(), max_num_tokens_padded, + max_num_m_blocks, + reinterpret_cast(sorted_token_ids.mutable_data_ptr()), + reinterpret_cast(expert_ids.mutable_data_ptr()), + topk_num, + reinterpret_cast( + num_tokens_post_pad.mutable_data_ptr()), + reinterpret_cast(adapter_enabled.mutable_data_ptr()), + reinterpret_cast(lora_ids.mutable_data_ptr()), + reinterpret_cast(token_mask.mutable_data_ptr()), + has_expert_map); + } else { + int num_thread = 1024; + dim3 blockDim(num_thread); + size_t num_warps = CEILDIV(padded_num_experts, WARP_SIZE); + + size_t shared_mem_size = num_warps * WARP_SIZE * sizeof(int32_t); + + // cumsum buffer + torch::stable::Tensor cumsum = torch::stable::new_zeros( + topk_ids, {max_loras * (num_experts + 1)}, + torch::headeronly::ScalarType::Int); + + auto align_kernel = + vllm::moe::moe_lora_align_block_size_kernel; + + // Launch two threadblocks per LoRA slot, across max_loras + 1 slots + // to cover the extra "-1" (base-model tokens) entry that + // active_lora_ids may contain in addition to up to max_loras real + // LoRA slots. Using max_loras would drop the real LoRA slot when -1 + // occupies position 0 and leave the output buffers uninitialized, + // causing illegal memory accesses downstream. Mirrors the grid fix + // applied to _fused_moe_lora_kernel in vllm-project/vllm#32277. + // blockIdx.x % 2 == 0: counting experts and aligning + // blockIdx.x % 2 == 1: filling sorted_token_ids + align_kernel<<<(max_loras + 1) * 2, blockDim, shared_mem_size, + stream>>>( + reinterpret_cast(topk_ids.mutable_data_ptr()), + reinterpret_cast(token_lora_mapping.mutable_data_ptr()), + block_size, + reinterpret_cast(expert_map.mutable_data_ptr()), + num_experts, max_loras, topk_ids.numel(), max_num_tokens_padded, + max_num_m_blocks, + reinterpret_cast(sorted_token_ids.mutable_data_ptr()), + reinterpret_cast(expert_ids.mutable_data_ptr()), + topk_num, + reinterpret_cast( + num_tokens_post_pad.mutable_data_ptr()), + reinterpret_cast(adapter_enabled.mutable_data_ptr()), + reinterpret_cast(cumsum.mutable_data_ptr()), WARP_SIZE, + padded_num_experts, + reinterpret_cast(lora_ids.mutable_data_ptr()), + reinterpret_cast(token_mask.mutable_data_ptr()), + has_expert_map); + + const int block_threads = std::min(256, (int)num_thread); + const int num_blocks = + (topk_ids.numel() + block_threads - 1) / block_threads; + + const int max_blocks = 65535; + const int actual_blocks = std::min(num_blocks, max_blocks); + + // Same rationale as align_kernel above: iterate over max_loras + 1 + // slots so the sort kernel processes the real LoRA slot even when + // active_lora_ids has -1 at position 0. + dim3 gridDims(max_loras + 1, actual_blocks); + auto sort_kernel = + vllm::moe::lora_count_and_sort_expert_tokens_kernel; + + sort_kernel<<>>( + reinterpret_cast(topk_ids.const_data_ptr()), + reinterpret_cast(sorted_token_ids.mutable_data_ptr()), + reinterpret_cast(cumsum.mutable_data_ptr()), + reinterpret_cast(expert_map.mutable_data_ptr()), + topk_ids.numel(), num_experts, max_num_tokens_padded, topk_num, + reinterpret_cast(token_mask.mutable_data_ptr()), + max_loras, + reinterpret_cast(lora_ids.mutable_data_ptr()), + reinterpret_cast(adapter_enabled.mutable_data_ptr()), + has_expert_map); + } + }); +} \ No newline at end of file diff --git a/ex_engine/csrc/moe/moe_fused_topk.cu b/ex_engine/csrc/moe/moe_fused_topk.cu new file mode 100644 index 00000000..26f2a475 --- /dev/null +++ b/ex_engine/csrc/moe/moe_fused_topk.cu @@ -0,0 +1,56 @@ +/* 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 "moe_topk_sigmoid_kernels.cuh" +#include "moe_topk_softmax_kernels.cuh" + +namespace xllm::kernel::cuda { + +std::tuple moe_fused_topk( + torch::Tensor& gating_output, + int64_t topk, + bool renormalize, + const std::optional& correction_bias, + const std::string& scoring_func) { + int64_t num_tokens = gating_output.size(0); + + torch::Tensor topk_weights = torch::empty( + {num_tokens, topk}, + torch::dtype(torch::kFloat32).device(gating_output.device())); + torch::Tensor topk_ids = + torch::empty({num_tokens, topk}, + torch::dtype(torch::kInt32).device(gating_output.device())); + + if (scoring_func == "softmax") { + std::optional none_correction_bias = std::nullopt; + topk_softmax(topk_weights, + topk_ids, + gating_output, + renormalize, + /*moe_softcapping=*/0.0, + none_correction_bias); + } else if (scoring_func == "sigmoid") { + topk_sigmoid( + topk_weights, topk_ids, gating_output, renormalize, correction_bias); + } else { + LOG(FATAL) << "Unsupported scoring function for moe topk: " << scoring_func + << "only softmax and sigmoid are supported"; + } + + return std::make_tuple(topk_weights, topk_ids); +} + +} // namespace xllm::kernel::cuda diff --git a/ex_engine/csrc/moe/moe_ops.h b/ex_engine/csrc/moe/moe_ops.h new file mode 100644 index 00000000..43cbb7f8 --- /dev/null +++ b/ex_engine/csrc/moe/moe_ops.h @@ -0,0 +1,87 @@ +#pragma once + +#include + +#include +#include + +void topk_softmax(torch::stable::Tensor& topk_weights, + torch::stable::Tensor& topk_indices, + torch::stable::Tensor& token_expert_indices, + torch::stable::Tensor& gating_output, bool renormalize, + std::optional bias); + +void topk_sigmoid(torch::stable::Tensor& topk_weights, + torch::stable::Tensor& topk_indices, + torch::stable::Tensor& token_expert_indices, + torch::stable::Tensor& gating_output, bool renormalize, + std::optional bias); + +void topk_softplus_sqrt( + torch::stable::Tensor& topk_weights, torch::stable::Tensor& topk_indices, + torch::stable::Tensor& token_expert_indices, + torch::stable::Tensor& gating_output, bool renormalize, + double routed_scaling_factor, + const std::optional& correction_bias, + const std::optional& input_ids, + const std::optional& tid2eid); + +void moe_sum(torch::stable::Tensor& input, torch::stable::Tensor& output); + +void moe_align_block_size( + torch::stable::Tensor topk_ids, int64_t num_experts, int64_t block_size, + torch::stable::Tensor sorted_token_ids, torch::stable::Tensor experts_ids, + torch::stable::Tensor num_tokens_post_pad, + std::optional maybe_expert_map); + +void batched_moe_align_block_size( + int64_t max_tokens_per_batch, int64_t block_size, + const torch::stable::Tensor& expert_num_tokens, + torch::stable::Tensor sorted_ids, torch::stable::Tensor expert_ids, + torch::stable::Tensor num_tokens_post_pad); + +void moe_lora_align_block_size( + torch::stable::Tensor topk_ids, torch::stable::Tensor token_lora_mapping, + int64_t num_experts, int64_t block_size, int64_t max_loras, + int64_t max_num_tokens_padded, int64_t max_num_m_blocks, + torch::stable::Tensor sorted_token_ids, torch::stable::Tensor expert_ids, + torch::stable::Tensor num_tokens_post_pad, + torch::stable::Tensor adapter_enabled, torch::stable::Tensor lora_ids, + std::optional maybe_expert_map); +#ifndef USE_ROCM +torch::stable::Tensor moe_wna16_gemm( + torch::stable::Tensor input, torch::stable::Tensor output, + torch::stable::Tensor b_qweight, torch::stable::Tensor b_scales, + std::optional b_qzeros, + std::optional topk_weights, + torch::stable::Tensor sorted_token_ids, torch::stable::Tensor expert_ids, + torch::stable::Tensor num_tokens_post_pad, int64_t top_k, + int64_t BLOCK_SIZE_M, int64_t BLOCK_SIZE_N, int64_t BLOCK_SIZE_K, + int64_t bit); + +std::tuple grouped_topk( + const torch::stable::Tensor& scores, int64_t n_group, int64_t topk_group, + int64_t topk, bool renormalize, double routed_scaling_factor, + const torch::stable::Tensor& bias, int64_t scoring_func); +#endif + +bool moe_permute_unpermute_supported(); + +int64_t moe_permute_sort_workspace_size(int64_t num_expanded_rows, + int64_t num_expert); + +void shuffle_rows(const torch::stable::Tensor& input_tensor, + const torch::stable::Tensor& dst2src_map, + torch::stable::Tensor& output_tensor); + +#ifndef USE_ROCM +// DeepSeek V3 optimized router GEMM kernel for SM90+ +// Computes output = mat_a @ mat_b.T where: +// mat_a: [num_tokens, hidden_dim] in bf16 +// mat_b: [num_experts, hidden_dim] in bf16 +// output: [num_tokens, num_experts] in bf16 or fp32 +// Supports num_tokens in [1, 16], num_experts in {256, 384}, hidden_dim = 7168 +void dsv3_router_gemm(torch::stable::Tensor& output, + const torch::stable::Tensor& mat_a, + const torch::stable::Tensor& mat_b); +#endif diff --git a/ex_engine/csrc/moe/moe_topk.cuh b/ex_engine/csrc/moe/moe_topk.cuh new file mode 100644 index 00000000..8d66bb21 --- /dev/null +++ b/ex_engine/csrc/moe/moe_topk.cuh @@ -0,0 +1,285 @@ + +/* + * Copyright (c) 2025, NVIDIA CORPORATION. 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 + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * 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. + */ + +// refers to +// https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/tensorrt_llm/kernels/moeTopKFuncs.cuh + +#pragma once + +#include +#include + +#include + +#include "core/kernels/cuda/arch_condition.h" + +namespace xllm::kernel::cuda { +namespace reduce_topk { +namespace cg = cooperative_groups; +static constexpr int kWARP_SIZE = 32; +static constexpr bool kTLLM_GEN_HAS_FAST_REDUX = arch::is_major_v<10>; + +template +struct TopKRedType { + using T = T_; + static_assert( + std::is_same_v || std::is_same_v || + std::is_same_v || std::is_same_v, + "Top K reduction only implemented for int, float, float16 and bfloat16"); + + using TypeCmp = std::conditional_t; + using IdxT = std::conditional_t; + + static constexpr int kMoveBits = (sizeof(T) == 4) ? 32 : 16; + static constexpr int kMaxIdx = 65535; + TypeCmp compValIdx; + + static __host__ __device__ inline TypeCmp makeCmpVal(T val, int32_t idx = 0) { + auto valueBits = cub::Traits::TwiddleIn( + reinterpret_cast::UnsignedBits&>(val)); + TypeCmp compactTmp = valueBits; + compactTmp = (compactTmp << kMoveBits) | (0xFFFF & (kMaxIdx - idx)); + // Use 65535 minus idx to give higher priority to elements with smaller + // indices. + return compactTmp; + } + + static __host__ __device__ void unpack(T& value, + int32_t& index, + TypeCmp cmp) { + // Since “65535-idx” is always smaller than 65536 and positive, we can + // directly use it as the lower 16 bits + index = kMaxIdx - static_cast((cmp & 0xFFFF)); + + auto compactTmp = cmp >> kMoveBits; + auto valueBits = cub::Traits::TwiddleOut( + reinterpret_cast::UnsignedBits&>(compactTmp)); + value = reinterpret_cast(valueBits); + } + + __host__ __device__ TopKRedType() = default; + + __host__ __device__ TopKRedType(T val, int32_t idx) + : compValIdx(makeCmpVal(val, idx)) {} + + __host__ __device__ operator TypeCmp() const noexcept { return compValIdx; } + + __device__ inline TypeCmp reduce( + cg::thread_block_tile const& warp) { + if constexpr (!kTLLM_GEN_HAS_FAST_REDUX || sizeof(TypeCmp) == 8) { + return cg::reduce(warp, compValIdx, cg::greater{}); + } else { + TypeCmp result; + asm("redux.sync.max.u32 %0, %1, 0xffffffff;\n" + : "=r"(result) + : "r"(compValIdx)); + return result; + } + } +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct TopKIdx { + // by default, empty +}; + +template +struct TopKIdx { + static constexpr int K = K_; + int32_t val[K]; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +#define TOPK_SWAP(I, J) \ + { \ + auto pairMin = min(topK[I].compValIdx, topK[J].compValIdx); \ + auto pairMax = max(topK[I].compValIdx, topK[J].compValIdx); \ + topK[I].compValIdx = pairMax; \ + topK[J].compValIdx = pairMin; \ + } + +template +struct Sort; + +template +struct Sort<1, RedType> { + static __device__ void run(RedType* topK) {} +}; + +template +struct Sort<2, RedType> { + static __device__ void run(RedType* topK) { TOPK_SWAP(0, 1); } +}; + +template +struct Sort<3, RedType> { + static __device__ void run(RedType* topK) { + TOPK_SWAP(0, 1); + TOPK_SWAP(1, 2); + TOPK_SWAP(0, 1); + } +}; + +template +struct Sort<4, RedType> { + static __device__ void run(RedType* topK) { + TOPK_SWAP(0, 2); + TOPK_SWAP(1, 3); + TOPK_SWAP(0, 1); + TOPK_SWAP(2, 3); + TOPK_SWAP(1, 2); + } +}; + +template +__forceinline__ __device__ void reduceTopK( + cg::thread_block_tile const& warp, + Type (&out)[K], + int32_t (&outIdx)[K], + Type value, + int32_t idx, + Type const minValue, + int actualK = K) { + static_assert(K > 0, "Top K must have K > 0"); + static_assert(K < kWARP_SIZE, "Top K must have K < kWARP_SIZE"); + using RedType = TopKRedType; + RedType topK{value, idx}; + typename RedType::TypeCmp packedMax{}; +#pragma unroll + for (int kk = 0; kk < actualK; ++kk) //@todo: check if actualK is correct + { + topK = + kk > 0 && packedMax == topK.compValIdx ? RedType{minValue, idx} : topK; + // get the next largest value + packedMax = topK.reduce(warp); + RedType::unpack(out[kk], outIdx[kk], packedMax); + } +}; + +template +__device__ void reduceTopKFunc(cg::thread_block_tile const& warp, + Type (&out)[K], + int32_t (&outIdx)[K], + Type (&value)[N], + int32_t (&idx)[N], + Type minValue, + int actualK = K) { + static_assert(K > 0, "Top K must have K > 0"); + static_assert(K < kWARP_SIZE, "Top K must have K < kWARP_SIZE"); + static_assert(N > 0, "Top K must have N > 0"); + static_assert(N < 5, + "Only support candidates number less than or equal to 128"); + using RedType = TopKRedType; + RedType topK[N]; +#pragma unroll + for (int nn = 0; nn < N; ++nn) { + topK[nn] = RedType{value[nn], idx[nn]}; + } + + if constexpr (!IsSorted) { + Sort::run(topK); + } + typename RedType::TypeCmp packedMax{}; +#pragma unroll + for (int kk = 0; kk < actualK; ++kk) { + bool update = kk > 0 && packedMax == topK[0].compValIdx; +#pragma unroll + for (int nn = 0; nn < N; ++nn) { + topK[nn] = update && nn == N - 1 ? RedType{minValue, idx[nn]} + : update ? topK[nn + 1] + : topK[nn]; + } + // get the next largest value + packedMax = topK[0].reduce(warp); + RedType::unpack(out[kk], outIdx[kk], packedMax); + } +}; + +template +__forceinline__ __device__ void reduceTopK( + cg::thread_block_tile const& warp, + Type (&out)[K], + int32_t (&outIdx)[K], + Type (&value)[N], + int32_t (&idx)[N], + Type const minValue, + int actualK = K) { + static_assert(K > 0, "Top K must have K > 0"); + static_assert(K < kWARP_SIZE, "Top K must have K < kWARP_SIZE"); + static_assert(N > 0, "Top K must have N > 0"); + static_assert( + N <= 16, + "Only support candidates number less than or equal to 16*32=512"); + static_assert(N <= 4 || N % 4 == 0, + "Only support candidates number is a multiple of 4*32=128 or " + "less than or equal to 4"); + using RedType = TopKRedType; + + if constexpr (N <= 4) { + reduceTopKFunc( + warp, out, outIdx, value, idx, minValue, actualK); + } else { + constexpr int numLoops = N / 4; + constexpr int numResults = (numLoops * K - 1) / kWARP_SIZE + 1; + + Type topKBufferValue[numResults]; + int32_t topKBufferIdx[numResults]; + int32_t laneIdx = threadIdx.x % kWARP_SIZE; + + // Sentinel index must be in [0, kMaxIdx] to survive makeCmpVal pack/unpack + // (kMaxIdx - idx is stored in 16 bits; -1 would become 0 and unpack to + // 65535). Use kMaxIdx so sentinel slots have smallest compValIdx for + // minValue and lose to any real candidate. + for (int ii = 0; ii < numResults; ++ii) { + topKBufferValue[ii] = minValue; + topKBufferIdx[ii] = RedType::kMaxIdx; + } + for (int loop = 0; loop < numLoops; ++loop) { + int start = loop * 4; + Type topKValue[K]; + int32_t topKIdx[K]; + Type inValue[4]; + int32_t inIdx[4]; + for (int i = 0; i < 4; ++i) { + inValue[i] = value[start + i]; + inIdx[i] = idx[start + i]; + } + reduceTopKFunc( + warp, topKValue, topKIdx, inValue, inIdx, minValue, actualK); + int inOffset = laneIdx % K; + if (laneIdx >= loop * K && laneIdx < (loop + 1) * K) { + topKBufferValue[0] = topKValue[inOffset]; + topKBufferIdx[0] = topKIdx[inOffset]; + } + if (loop == numLoops - 1 && (laneIdx < (numLoops * K - kWARP_SIZE))) { + topKBufferValue[1] = topKValue[inOffset]; + topKBufferIdx[1] = topKIdx[inOffset]; + } + } + + reduceTopKFunc( + warp, out, outIdx, topKBufferValue, topKBufferIdx, minValue, actualK); + } +}; + +#undef TOPK_SWAP + +} // namespace reduce_topk +} // namespace xllm::kernel::cuda diff --git a/ex_engine/csrc/moe/moe_topk_sigmoid_kernels.cuh b/ex_engine/csrc/moe/moe_topk_sigmoid_kernels.cuh new file mode 100644 index 00000000..a8de51c2 --- /dev/null +++ b/ex_engine/csrc/moe/moe_topk_sigmoid_kernels.cuh @@ -0,0 +1,602 @@ +// Adapt from +// https://github.com/vllm-project/vllm/blob/v0.7.3/csrc/moe/topk_softmax_kernels.cu +// which is originally adapted from +// https://github.com/NVIDIA/TensorRT-LLM/blob/v0.7.1/cpp/tensorrt_llm/kernels/mixtureOfExperts/moe_kernels.cu +/* Copyright 2025 SGLang Team. 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 + + http://www.apache.org/licenses/LICENSE-2.0 + +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 +#include +#include + +#include +#include + +#include "kernels/cuda/device_utils.cuh" + +namespace { + +using namespace xllm::kernel::cuda; + +// ====================== Sigmoid things =============================== +// We have our own implementation of sigmoid here so we can support transposing +// the output in the sigmoid kernel when we extend this module to support +// expert-choice routing. +template +__launch_bounds__(TPB) __global__ + void moe_sigmoid(const T* input, + const bool* finished, + float* output, + const int num_cols, + const float* correction_bias) { + const int thread_row_offset = blockIdx.x * num_cols; + + // Don't touch finished rows. + if ((finished != nullptr) && finished[blockIdx.x]) { + return; + } + + // First pass: Apply transformation, find max, and write transformed values to + // output + for (int ii = threadIdx.x; ii < num_cols; ii += TPB) { + const int idx = thread_row_offset + ii; + float val = convert_to_float(input[idx]); + + val = 1.0f / (1.0f + expf(-val)); + + // Apply correction bias if provided + if (correction_bias != nullptr) { + val = val + correction_bias[ii]; + } + + output[idx] = val; // Store transformed value + } +} + +template +__launch_bounds__(TPB) __global__ + void moe_topK(const float* inputs_after_sigmoid, + const bool* finished, + float* output, + int* indices, + const int num_experts, + const int k, + const int start_expert, + const int end_expert, + const bool renormalize, + const float* correction_bias) { + using cub_kvp = cub::KeyValuePair; + using BlockReduce = cub::BlockReduce; + __shared__ typename BlockReduce::TempStorage tmpStorage; + + cub_kvp thread_kvp; + cub::ArgMax arg_max; + + const int block_row = blockIdx.x; + + const bool row_is_active = finished ? !finished[block_row] : true; + const int thread_read_offset = blockIdx.x * num_experts; + float row_sum_for_renormalize = 0; + for (int k_idx = 0; k_idx < k; ++k_idx) { + thread_kvp.key = 0; + thread_kvp.value = -1.f; // This is OK because inputs are probabilities + + cub_kvp inp_kvp; + for (int expert = threadIdx.x; expert < num_experts; expert += TPB) { + const int idx = thread_read_offset + expert; + inp_kvp.key = expert; + inp_kvp.value = inputs_after_sigmoid[idx]; + + for (int prior_k = 0; prior_k < k_idx; ++prior_k) { + const int prior_winning_expert = indices[k * block_row + prior_k]; + + if (prior_winning_expert == expert) { + inp_kvp = thread_kvp; + } + } + + thread_kvp = arg_max(inp_kvp, thread_kvp); + } + + const cub_kvp result_kvp = + BlockReduce(tmpStorage).Reduce(thread_kvp, arg_max); + if (threadIdx.x == 0) { + // Ignore experts the node isn't responsible for with expert parallelism + const int expert = result_kvp.key; + const bool node_uses_expert = + expert >= start_expert && expert < end_expert; + const bool should_process_row = row_is_active && node_uses_expert; + + const int idx = k * block_row + k_idx; + float val = result_kvp.value; + if (correction_bias != nullptr) { + val -= correction_bias[expert]; + } + output[idx] = val; + indices[idx] = should_process_row ? (expert - start_expert) : num_experts; + assert(indices[idx] >= 0); + row_sum_for_renormalize += val; + } + __syncthreads(); + } + + if (renormalize && threadIdx.x == 0) { + float row_sum_for_renormalize_inv = 1.f / row_sum_for_renormalize; + for (int k_idx = 0; k_idx < k; ++k_idx) { + const int idx = k * block_row + k_idx; + output[idx] = output[idx] * row_sum_for_renormalize_inv; + } + } +} + +// ====================== TopK sigmoid things =============================== + +/* + A Top-K gating sigmoid written to exploit when the number of experts in the + MoE layers are a small power of 2. This allows us to cleanly share the rows + among the threads in a single warp and eliminate communication between warps + (so no need to use shared mem). + + It fuses the sigmoid, max and argmax into a single kernel. + + Limitations: + 1) This implementation is intended for when the number of experts is a small + power of 2. 2) This implementation assumes k is small, but will work for any + k. +*/ + +template +__launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__ + void topk_gating_sigmoid(const T* input, + const bool* finished, + float* output, + const int num_rows, + int* indices, + const int k, + const int start_expert, + const int end_expert, + const bool renormalize, + const float* correction_bias) { + // We begin by enforcing compile time assertions and setting up compile time + // constants. + static_assert(VPT == (VPT & -VPT), "VPT must be power of 2"); + static_assert(NUM_EXPERTS == (NUM_EXPERTS & -NUM_EXPERTS), + "NUM_EXPERTS must be power of 2"); + static_assert(BYTES_PER_LDG == (BYTES_PER_LDG & -BYTES_PER_LDG), + "BYTES_PER_LDG must be power of 2"); + static_assert(BYTES_PER_LDG <= 16, "BYTES_PER_LDG must be leq 16"); + + // Number of bytes each thread pulls in per load + static constexpr int ELTS_PER_LDG = BYTES_PER_LDG / sizeof(T); + static constexpr int ELTS_PER_ROW = NUM_EXPERTS; + static constexpr int THREADS_PER_ROW = ELTS_PER_ROW / VPT; + static constexpr int LDG_PER_THREAD = VPT / ELTS_PER_LDG; + + // Restrictions based on previous section. + static_assert( + VPT % ELTS_PER_LDG == 0, + "The elements per thread must be a multiple of the elements per ldg"); + static_assert(WARP_SIZE % THREADS_PER_ROW == 0, + "The threads per row must cleanly divide the threads per warp"); + static_assert(THREADS_PER_ROW == (THREADS_PER_ROW & -THREADS_PER_ROW), + "THREADS_PER_ROW must be power of 2"); + static_assert(THREADS_PER_ROW <= WARP_SIZE, + "THREADS_PER_ROW can be at most warp size"); + + // We have NUM_EXPERTS elements per row. We specialize for small #experts + static constexpr int ELTS_PER_WARP = WARP_SIZE * VPT; + static constexpr int ROWS_PER_WARP = ELTS_PER_WARP / ELTS_PER_ROW; + static constexpr int ROWS_PER_CTA = WARPS_PER_CTA * ROWS_PER_WARP; + + // Restrictions for previous section. + static_assert(ELTS_PER_WARP % ELTS_PER_ROW == 0, + "The elts per row must cleanly divide the total elt per warp"); + + // ===================== From this point, we finally start computing run-time + // variables. ======================== + + // Compute CTA and warp rows. We pack multiple rows into a single warp, and a + // block contains WARPS_PER_CTA warps. This, each block processes a chunk of + // rows. We start by computing the start row for each block. + const int cta_base_row = blockIdx.x * ROWS_PER_CTA; + + // Now, using the base row per thread block, we compute the base row per warp. + const int warp_base_row = cta_base_row + threadIdx.y * ROWS_PER_WARP; + + // The threads in a warp are split into sub-groups that will work on a row. + // We compute row offset for each thread sub-group + const int thread_row_in_warp = threadIdx.x / THREADS_PER_ROW; + const int thread_row = warp_base_row + thread_row_in_warp; + + // Threads with indices out of bounds should early exit here. + if (thread_row >= num_rows) { + return; + } + const bool row_is_active = finished ? !finished[thread_row] : true; + + // We finally start setting up the read pointers for each thread. First, each + // thread jumps to the start of the row it will read. + const T* thread_row_ptr = input + thread_row * ELTS_PER_ROW; + + // Now, we compute the group each thread belong to in order to determine the + // first column to start loads. + const int thread_group_idx = threadIdx.x % THREADS_PER_ROW; + const int first_elt_read_by_thread = thread_group_idx * ELTS_PER_LDG; + const T* thread_read_ptr = thread_row_ptr + first_elt_read_by_thread; + + // Determine the pointer type to use to read in the data depending on the + // BYTES_PER_LDG template param. In theory, this can support all powers of 2 + // up to 16. NOTE(woosuk): The original implementation uses CUTLASS aligned + // array here. We defined our own aligned array and use it here to avoid the + // dependency on CUTLASS. + using AccessType = AlignedArray; + + // Finally, we pull in the data from global mem + T row_chunk_temp[VPT]; + AccessType* row_chunk_vec_ptr = + reinterpret_cast(&row_chunk_temp); + const AccessType* vec_thread_read_ptr = + reinterpret_cast(thread_read_ptr); +#pragma unroll + // Note(Byron): interleaved loads to achieve better memory coalescing + // | thread[0] | thread[1] | thread[2] | thread[3] | thread[0] | thread[1] | + // thread[2] | thread[3] | ... + for (int ii = 0; ii < LDG_PER_THREAD; ++ii) { + row_chunk_vec_ptr[ii] = vec_thread_read_ptr[ii * THREADS_PER_ROW]; + } + + float row_chunk[VPT]; +#pragma unroll + // Note(Byron): upcast logits to float32 + for (int ii = 0; ii < VPT; ++ii) { + float val = convert_to_float(row_chunk_temp[ii]); + val = 1.0f / (1.0f + expf(-val)); + // Apply correction bias if provided + if (correction_bias != nullptr) { + /* + LDG is interleaved + |thread0 LDG| |thread1 LDG| |thread0 LDG| |thread1 LDG| + |--------- group0 --------| |----------group1 --------| + ^ local2 + */ + const int group_id = ii / ELTS_PER_LDG; + const int local_id = ii % ELTS_PER_LDG; + const int expert_idx = first_elt_read_by_thread + + group_id * THREADS_PER_ROW * ELTS_PER_LDG + + local_id; + val = val + correction_bias[expert_idx]; + } + + row_chunk[ii] = val; + } + + // Now, row_chunk contains the sigmoid of the row chunk. Now, I want to find + // the topk elements in each row, along with the max index. + int start_col = first_elt_read_by_thread; + static constexpr int COLS_PER_GROUP_LDG = ELTS_PER_LDG * THREADS_PER_ROW; + + float row_sum_for_renormalize = 0; + + for (int k_idx = 0; k_idx < k; ++k_idx) { + // First, each thread does the local argmax + float max_val = row_chunk[0]; + int expert = start_col; +#pragma unroll + for (int ldg = 0, col = start_col; ldg < LDG_PER_THREAD; + ++ldg, col += COLS_PER_GROUP_LDG) { +#pragma unroll + for (int ii = 0; ii < ELTS_PER_LDG; ++ii) { + float val = row_chunk[ldg * ELTS_PER_LDG + ii]; + + // No check on the experts here since columns with the smallest index + // are processed first and only updated if > (not >=) + if (val > max_val) { + max_val = val; + expert = col + ii; + } + } + } + +// Now, we perform the argmax reduce. We use the butterfly pattern so threads +// reach consensus about the max. This will be useful for K > 1 so that the +// threads can agree on "who" had the max value. That thread can then blank out +// their max with -inf and the warp can run more iterations... +#pragma unroll + for (int mask = THREADS_PER_ROW / 2; mask > 0; mask /= 2) { + float other_max = + XLLM_SHFL_XOR_SYNC_WIDTH(0xffffffff, max_val, mask, THREADS_PER_ROW); + int other_expert = + XLLM_SHFL_XOR_SYNC_WIDTH(0xffffffff, expert, mask, THREADS_PER_ROW); + + // We want lower indices to "win" in every thread so we break ties this + // way + if (other_max > max_val || + (other_max == max_val && other_expert < expert)) { + max_val = other_max; + expert = other_expert; + } + } + + // Write the max for this k iteration to global memory. + if (thread_group_idx == 0) { + // Add a guard to ignore experts not included by this node + const bool node_uses_expert = + expert >= start_expert && expert < end_expert; + const bool should_process_row = row_is_active && node_uses_expert; + + // The lead thread from each sub-group will write out the final results to + // global memory. (This will be a single) thread per row of the + // input/output matrices. + const int idx = k * thread_row + k_idx; + if (correction_bias != nullptr) { + max_val -= correction_bias[expert]; + } + output[idx] = max_val; + indices[idx] = should_process_row ? (expert - start_expert) : NUM_EXPERTS; + row_sum_for_renormalize += max_val; + } + + // Finally, we clear the value in the thread with the current max if there + // is another iteration to run. + if (k_idx + 1 < k) { + const int ldg_group_for_expert = expert / COLS_PER_GROUP_LDG; + const int thread_to_clear_in_group = + (expert / ELTS_PER_LDG) % THREADS_PER_ROW; + + // Only the thread in the group which produced the max will reset the + // "winning" value to -inf. + if (thread_group_idx == thread_to_clear_in_group) { + const int offset_for_expert = expert % ELTS_PER_LDG; + // Safe to set to any negative value since row_chunk values must be + // between 0 and 1. + row_chunk[ldg_group_for_expert * ELTS_PER_LDG + offset_for_expert] = + -10000.f; + } + } + } + + // Fuse renormalization of topk_weights into this kernel + if (renormalize && thread_group_idx == 0) { + float row_sum_for_renormalize_inv = 1.f / row_sum_for_renormalize; +#pragma unroll + for (int k_idx = 0; k_idx < k; ++k_idx) { + const int idx = k * thread_row + k_idx; + output[idx] = output[idx] * row_sum_for_renormalize_inv; + } + } +} + +template +void topk_gating_sigmoid_launcher_helper(const T* input, + const bool* finished, + float* output, + int* indices, + const int num_rows, + const int k, + const int start_expert, + const int end_expert, + const bool renormalize, + const float* correction_bias, + cudaStream_t stream) { + static constexpr std::size_t MAX_BYTES_PER_LDG = 16; + + static constexpr int BYTES_PER_LDG = + MIN(MAX_BYTES_PER_LDG, sizeof(T) * EXPERTS); + using Constants = TopkConstants; + static constexpr int VPT = Constants::VPT; + static constexpr int ROWS_PER_WARP = Constants::ROWS_PER_WARP; + const int num_warps = (num_rows + ROWS_PER_WARP - 1) / ROWS_PER_WARP; + const int num_blocks = (num_warps + WARPS_PER_TB - 1) / WARPS_PER_TB; + + dim3 block_dim(WARP_SIZE, WARPS_PER_TB); + topk_gating_sigmoid + <<>>(input, + finished, + output, + num_rows, + indices, + k, + start_expert, + end_expert, + renormalize, + correction_bias); +} + +#define LAUNCH_SIGMOID(TYPE, NUM_EXPERTS, WARPS_PER_TB) \ + topk_gating_sigmoid_launcher_helper( \ + gating_output, \ + nullptr, \ + topk_weights, \ + topk_indices, \ + num_tokens, \ + topk, \ + 0, \ + num_experts, \ + renormalize, \ + correction_bias, \ + stream); + +template +void topk_gating_sigmoid_kernel_launcher(const T* gating_output, + float* topk_weights, + int* topk_indices, + float* sigmoid_workspace, + const int num_tokens, + const int num_experts, + const int topk, + const bool renormalize, + const float* correction_bias, + cudaStream_t stream) { + static constexpr int WARPS_PER_TB = 4; + switch (num_experts) { + case 1: + LAUNCH_SIGMOID(T, 1, WARPS_PER_TB); + break; + case 2: + LAUNCH_SIGMOID(T, 2, WARPS_PER_TB); + break; + case 4: + LAUNCH_SIGMOID(T, 4, WARPS_PER_TB); + break; + case 8: + LAUNCH_SIGMOID(T, 8, WARPS_PER_TB); + break; + case 16: + LAUNCH_SIGMOID(T, 16, WARPS_PER_TB); + break; + case 32: + LAUNCH_SIGMOID(T, 32, WARPS_PER_TB); + break; + case 64: + LAUNCH_SIGMOID(T, 64, WARPS_PER_TB); + break; + case 128: + LAUNCH_SIGMOID(T, 128, WARPS_PER_TB); + break; + case 256: + LAUNCH_SIGMOID(T, 256, WARPS_PER_TB); + break; + default: { + TORCH_CHECK(sigmoid_workspace != nullptr, + "sigmoid_workspace must be provided for num_experts that are " + "not a power of 2."); + static constexpr int TPB = 256; + moe_sigmoid<<>>(gating_output, + nullptr, + sigmoid_workspace, + num_experts, + correction_bias); + moe_topK<<>>(sigmoid_workspace, + nullptr, + topk_weights, + topk_indices, + num_experts, + topk, + 0, + num_experts, + renormalize, + correction_bias); + } + } +} +} // namespace + +namespace xllm::kernel::cuda { +void topk_sigmoid(torch::Tensor& topk_weights, // [num_tokens, topk] + torch::Tensor& topk_indices, // [num_tokens, topk] + torch::Tensor& gating_output, // [num_tokens, num_experts] + const bool renormalize, + const std::optional& correction_bias) { + // Check data type + 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"; + + // Check dimensions + 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 + CHECK(gating_output.size(0) == topk_weights.size(0)) + << "First dimension of topk_weights must match num_tokens in " + "gating_output"; + CHECK(gating_output.size(0) == topk_indices.size(0)) + << "First dimension of topk_indices must match num_tokens in " + "gating_output"; + CHECK(topk_weights.size(-1) == topk_indices.size(-1)) + << "Second dimension of topk_indices must match topk in topk_weights"; + CHECK(topk_weights.size(-1) <= gating_output.size(-1)) + << "topk must be less than or equal to num_experts"; + + const int num_experts = static_cast(gating_output.size(-1)); + const int num_tokens = static_cast(gating_output.size(0)); + const int topk = static_cast(topk_weights.size(-1)); + + const bool is_pow_2 = + (num_experts != 0) && ((num_experts & (num_experts - 1)) == 0); + const bool needs_workspace = !is_pow_2 || num_experts > 256; + const int64_t workspace_size = needs_workspace ? num_tokens * num_experts : 0; + + const at::cuda::OptionalCUDAGuard device_guard(device_of(gating_output)); + const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + torch::Tensor sigmoid_workspace = torch::empty( + {workspace_size}, gating_output.options().dtype(at::ScalarType::Float)); + + const at::ScalarType dtype = gating_output.scalar_type(); + + // Validate correction_bias if provided - must always be float32 + const float* bias_ptr = nullptr; + if (correction_bias.has_value()) { + const torch::Tensor& bias_tensor = correction_bias.value(); + 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(); + } + + if (dtype == at::ScalarType::Float) { + topk_gating_sigmoid_kernel_launcher( + gating_output.data_ptr(), + topk_weights.data_ptr(), + topk_indices.data_ptr(), + sigmoid_workspace.data_ptr(), + num_tokens, + num_experts, + topk, + renormalize, + bias_ptr, + stream); + } else if (dtype == at::ScalarType::Half) { + topk_gating_sigmoid_kernel_launcher<__half>( + reinterpret_cast(gating_output.data_ptr()), + topk_weights.data_ptr(), + topk_indices.data_ptr(), + sigmoid_workspace.data_ptr(), + num_tokens, + num_experts, + topk, + renormalize, + bias_ptr, + stream); + } else if (dtype == at::ScalarType::BFloat16) { + topk_gating_sigmoid_kernel_launcher<__nv_bfloat16>( + reinterpret_cast( + gating_output.data_ptr()), + topk_weights.data_ptr(), + topk_indices.data_ptr(), + sigmoid_workspace.data_ptr(), + num_tokens, + num_experts, + topk, + renormalize, + bias_ptr, + stream); + } else { + LOG(FATAL) << "Unsupported gating_output dtype: " << dtype; + } +} +} // namespace xllm::kernel::cuda diff --git a/ex_engine/csrc/moe/moe_topk_softmax_kernels.cuh b/ex_engine/csrc/moe/moe_topk_softmax_kernels.cuh index 7f1d8a24..ea74ad18 100644 --- a/ex_engine/csrc/moe/moe_topk_softmax_kernels.cuh +++ b/ex_engine/csrc/moe/moe_topk_softmax_kernels.cuh @@ -22,11 +22,9 @@ limitations under the License. #include #include -#if CUDA_VERSION >= 12090 #include -#endif -#include "device_utils.cuh" +#include "kernels/cuda/device_utils.cuh" using cub_kvp = cub::KeyValuePair; @@ -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<<>>(gating_output, @@ -753,23 +752,29 @@ void topk_softmax(torch::Tensor& topk_weights, // [num_tokens, topk] const double moe_softcapping, const std::optional& 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(gating_output.size(-1)); const int num_tokens = static_cast(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(); } @@ -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 diff --git a/ex_engine/include/ilu_ops_api.h b/ex_engine/include/ilu_ops_api.h new file mode 100644 index 00000000..e4fd7853 --- /dev/null +++ b/ex_engine/include/ilu_ops_api.h @@ -0,0 +1,153 @@ +/* Copyright 2025 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. +==============================================================================*/ + +#pragma once + +#include +#include +#include +#include +#include + +#include + +#include "ATen/Tensor.h" +#include "ATen/cuda/CUDAEvent.h" +#include "c10/core/Device.h" +#include "c10/core/DeviceGuard.h" +#include "c10/core/GradMode.h" +#include "c10/core/InferenceMode.h" +#include "c10/core/MemoryFormat.h" +#include "c10/core/ScalarType.h" +#include "c10/core/TensorOptions.h" +#include "c10/cuda/CUDAFunctions.h" +#include "c10/cuda/CUDAGuard.h" +#include "c10/cuda/CUDAStream.h" +#include "ixformer.h" +#include "kernels/kernels.h" + +// #include "utils.h" +using namespace ixformer; + +namespace xllm::kernel::ilu { + +void apply_rope_pos_ids_cos_sin_cache(torch::Tensor& query, + torch::Tensor& key, + torch::Tensor& cos_sin_cache, + torch::Tensor& positions, + bool interleave); + +// act_mode only support silu, gelu, gelu_tanh +void act_and_mul(torch::Tensor out, + torch::Tensor input, + const std::string& act_mode); + +void reshape_paged_cache( + torch::Tensor& key, // (num_tokens, num_heads, head_size) + std::optional& value, // (num_tokens, num_heads, head_size) + torch::Tensor& key_cache, // (num_blocks, num_heads, block_size, head_size) + std::optional& + value_cache, // (num_blocks, num_heads, block_size, head_size) + torch::Tensor& slot_mapping); //(num_tokens) + +void batch_prefill(torch::Tensor& query, + const torch::Tensor& key, + const std::optional& value, + torch::Tensor& output, + std::optional& output_lse, + const std::optional& q_cu_seq_lens, + const std::optional& kv_cu_seq_lens, + const std::optional& alibi_slope, + const std::optional& attn_bias, + const std::optional& q_quant_scale, + const std::optional& k_quant_scale, + const std::optional& v_quant_scale, + const torch::Tensor& block_tables, + int64_t max_query_len, + int64_t max_seq_len, + float scale, + bool is_causal, + int64_t window_size_left, + int64_t window_size_right, + const std::string& compute_dtype, + bool return_lse); + +void batch_decode(torch::Tensor& query, + const torch::Tensor& k_cache, + torch::Tensor& output, + const torch::Tensor& block_table, + const torch::Tensor& seq_lens, + const std::optional& v_cache, + std::optional& output_lse, + const std::optional& q_quant_scale, + const std::optional& k_cache_quant_scale, + const std::optional& v_cache_quant_scale, + const std::optional& out_quant_scale, + const std::optional& alibi_slope, + const std::optional& mask, + const std::string& compute_dtype, + int64_t max_seq_len, + int64_t window_size_left, + int64_t window_size_right, + float scale, + bool return_lse, + bool is_causal, + int64_t kv_cache_quant_bit_size); + +void residual_layer_norm(torch::Tensor& input, + torch::Tensor& output, + std::optional& residual, + torch::Tensor& weight, + std::optional& bias, + std::optional& residual_out, + double eps); + +void rms_norm(torch::Tensor& output, + torch::Tensor& input, + torch::Tensor& weight, + double eps); + +torch::Tensor matmul(torch::Tensor a, + torch::Tensor b, + std::optional bias); + +std::tuple moe_active_topk( + const torch::Tensor& input, + int64_t topk, + int64_t num_expert_group, + int64_t topk_group, + bool normalize, + const std::optional& mask, + const std::string& normed_by, + const std::string& scoring_func, + double route_scale, + const std::optional& e_score_correction_bias); + +std::vector moe_gen_idx(torch::Tensor& expert_id, + int64_t expert_num); + +torch::Tensor moe_expand_input(const torch::Tensor& input, + const torch::Tensor& gather_index, + const torch::Tensor& combine_idx, + int64_t topk); + +torch::Tensor group_gemm(torch::Tensor& input, + torch::Tensor& weight, + torch::Tensor& tokens_per_experts, + const std::optional& dst_to_src, + torch::Tensor& output); + +torch::Tensor moe_combine_result(torch::Tensor& input, torch::Tensor& weight); +} // namespace xllm::kernel::ilu diff --git a/ex_engine/include/ilu_utils.h b/ex_engine/include/ilu_utils.h new file mode 100644 index 00000000..e8af0c3c --- /dev/null +++ b/ex_engine/include/ilu_utils.h @@ -0,0 +1,63 @@ +/* Copyright 2025 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. +==============================================================================*/ +#pragma once +namespace xllm::kernel::ilu { +#undef check_tensor_contiguous +#define check_tensor_contiguous(x, type) \ + TORCH_CHECK(x.scalar_type() == type); \ + TORCH_CHECK(x.is_cuda()); \ + TORCH_CHECK(x.is_contiguous()); + +#undef check_tensor_half_bf_float +#define check_tensor_half_bf_float(x) \ + TORCH_CHECK(x.scalar_type() == at::ScalarType::Half || \ + x.scalar_type() == at::ScalarType::Float || \ + x.scalar_type() == at::ScalarType::BFloat16); \ + TORCH_CHECK(x.is_cuda()); + +// from torchCheckMsgImpl +inline const char* ixformer_check_msg_impl(const char* msg) { return msg; } +// // If there is just 1 user-provided C-string argument, use it. + +#define IXFORMER_CHECK_MSG(cond, type, ...) \ + (ixformer_check_msg_impl( \ + "Expected " #cond \ + " to be true, but got false. " \ + "(Could this error message be improved? If so, " \ + "please report an enhancement request to ixformer.)", \ + ##__VA_ARGS__)) + +#define IXFORMER_CHECK(cond, ...) \ + { \ + if (!(cond)) { \ + std::cerr << __FILE__ << " (" << __LINE__ << ")" \ + << "-" << __FUNCTION__ << " : " \ + << IXFORMER_CHECK_MSG(cond, "", ##__VA_ARGS__) << std::endl; \ + throw std::runtime_error("IXFORMER_CHECK ERROR"); \ + } \ + } + +#undef CUINFER_CHECK +#define CUINFER_CHECK(func) \ + do { \ + cuinferStatus_t status = (func); \ + if (status != CUINFER_STATUS_SUCCESS) { \ + std::cerr << "Error in file " << __FILE__ << " on line " << __LINE__ \ + << ": " << cuinferGetErrorString(status) << std::endl; \ + throw std::runtime_error("CUINFER_CHECK ERROR"); \ + } \ + } while (0) + +} // namespace xllm::kernel::ilu \ No newline at end of file diff --git a/ex_engine/include/ixformer.h b/ex_engine/include/ixformer.h new file mode 100644 index 00000000..57ce66dc --- /dev/null +++ b/ex_engine/include/ixformer.h @@ -0,0 +1,147 @@ +/* Copyright 2025 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 + +#include "ATen/Tensor.h" +#include "utils.h" + +namespace ixformer::infer { +torch::Tensor ixinfer_flash_attn_unpad_with_block_tables( + torch::Tensor& query, + torch::Tensor& key_cache, + torch::Tensor& value_cache, + torch::Tensor& out, + torch::Tensor& block_tables, + torch::Tensor& cu_seq_q, + torch::Tensor& cu_seq_k, + int64_t max_seq_q, + int64_t max_seq_k, + bool is_causal, + int64_t window_left, + int64_t window_right, + double scale, + double softcap, + bool sqrt_alibi, + const std::optional& alibi_slopes, + const std::optional& sinks, + std::optional& lse); + +void silu_and_mul(torch::Tensor& input, torch::Tensor& output); + +torch::Tensor xllm_paged_attention( + torch::Tensor& out, + torch::Tensor& query, + torch::Tensor& key_cache, + torch::Tensor& value_cache, + int64_t num_kv_heads, + double scale, + torch::Tensor& block_tables, + torch::Tensor& context_lens, + int64_t block_size, + int64_t max_context_len, + const std::optional& alibi_slopes, + bool causal, + int32_t window_left, + int32_t window_right, + double softcap, + bool enable_cuda_graph, + bool use_sqrt_alibi, + const std::optional& sinks); + +torch::Tensor ixformer_linear(torch::Tensor& input, + torch::Tensor& weight, + int64_t act_type, + const std::optional& bias, + const std::optional& out, + const std::optional persistent); + +torch::Tensor ixformer_linear_ex(torch::Tensor& input, + torch::Tensor& weight, + const c10::optional& bias, + const c10::optional& out); + +void xllm_reshape_and_cache(torch::Tensor& key, + torch::Tensor& value, + torch::Tensor& key_cache, + torch::Tensor& value_cache, + torch::Tensor& slot_mapping, + int64_t key_token_stride, + int64_t value_token_stride); + +void xllm_rotary_embedding(torch::Tensor& positions, + torch::Tensor& query, + torch::Tensor& key, + int64_t head_size, + torch::Tensor& cos_sin_cache, + bool is_neox); + +void residual_rms_norm(torch::Tensor& input, + torch::Tensor& residual, + torch::Tensor& weight, + torch::Tensor& output, + torch::Tensor& residual_output, + const std::optional& fused_bias, + double alpha, + double eps, + bool is_post); + +void rms_norm(torch::Tensor& input, + torch::Tensor& weight, + torch::Tensor& output, + const std::optional& fused_bias, + double eps); + +void topk_softmax(torch::Tensor& topk_weights, + torch::Tensor& topk_indices, + torch::Tensor& token_expert_indices, + torch::Tensor& gating_output, + bool 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& expert_mask, + const c10::optional& expert_sizes_cpu, + const c10::optional& expand_tokens_gpu, + int64_t start_expert_id, + int64_t end_expert_id, + int64_t num_experts); + +void moe_expand_input(torch::Tensor outputs, + torch::Tensor inputs, + torch::Tensor dst_to_src, + const c10::optional& src_to_dst, + int64_t dst_tokens, + int64_t expand_factor); + +void moe_w16a16_group_gemm(torch::Tensor output, + torch::Tensor inputs, + torch::Tensor weights, + torch::Tensor tokens_per_experts, + const c10::optional& dst_to_src, + const c10::optional& bias, + std::string format, + int64_t persistent, + int64_t output_n); + +void moe_output_reduce_sum(torch::Tensor outputs, + torch::Tensor inputs, + const c10::optional& mul_weight, + const c10::optional& mask, + const c10::optional& extra_residual, + double scaling_factor); +} // namespace ixformer::infer