ref(EX): import upstream ILU kernels + xllm MoE CUDA sources into ex_engine
Copied from upstream_ref (NOT rewritten — exact upstream code):
ixformer C++ API (the authoritative header):
include/ixformer.h — ixformer::infer namespace: topk_softmax,
moe_compute_token_index_api, moe_w16a16_group_gemm, moe_expand_input,
moe_output_reduce_sum, silu_and_mul, rms_norm, xllm_paged_attention, etc.
include/ilu_ops_api.h — xllm::kernel::ilu namespace: moe_active_topk,
moe_gen_idx, moe_expand_input, group_gemm, moe_combine_result,
batch_prefill, batch_decode, rms_norm, matmul, act_and_mul, etc.
ILU kernel wrappers (call ixformer::infer directly):
csrc/ilu_kernel_fused_moe.cpp — topk routing + gen_idx + expand + combine
csrc/ilu_kernel_group_gemm.cpp — batched expert GEMM
csrc/ilu_kernel_{activation,norm,rope,matmul,attention}.cpp
ILU layer implementations (full pipeline):
csrc/ilu_layer_fused_moe.{cpp,h} — 797 lines, the complete MoE pipeline
that competitor 168 ran as corex_moe.py
csrc/ilu_layer_attention.{cpp,h} — prefill/decode attention dispatch
CUDA MoE kernels (from xllm + ds_vllm):
csrc/moe/moe_topk_softmax_kernels.cuh — CUB BlockReduce + warp topk
csrc/moe/moe_topk_sigmoid_kernels.cuh — sigmoid scoring variant
csrc/moe/moe_topk.cuh + moe_fused_topk.cu — entry points
csrc/moe/moeTopKFuncs.cuh — TRT-LLM derived vllm-compatible topk
csrc/moe/moe_ops.h + moe_align_sum_kernels.cu — alignment kernels
Common layer headers:
csrc/common_fused_moe{,_base}.h + common_moe_fused_topk.{cpp,h}
This commit is contained in:
project6-dev
54
ex_engine/csrc/common_fused_moe.h
Normal file
54
ex_engine/csrc/common_fused_moe.h
Normal file
@@ -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 <torch/torch.h>
|
||||
|
||||
#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
|
||||
27
ex_engine/csrc/common_fused_moe_base.h
Normal file
27
ex_engine/csrc/common_fused_moe_base.h
Normal file
@@ -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
|
||||
71
ex_engine/csrc/common_moe_fused_topk.cpp
Normal file
71
ex_engine/csrc/common_moe_fused_topk.cpp
Normal file
@@ -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<torch::Tensor, torch::Tensor> MoEFusedTopkImpl::forward(
|
||||
torch::Tensor& router_logits) {
|
||||
std::optional<torch::Tensor> 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
|
||||
53
ex_engine/csrc/common_moe_fused_topk.h
Normal file
53
ex_engine/csrc/common_moe_fused_topk.h
Normal file
@@ -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 <torch/torch.h>
|
||||
|
||||
#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<torch::Tensor, torch::Tensor> 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
|
||||
32
ex_engine/csrc/ilu_kernel_activation.cpp
Normal file
32
ex_engine/csrc/ilu_kernel_activation.cpp
Normal file
@@ -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
|
||||
163
ex_engine/csrc/ilu_kernel_attention.cpp
Normal file
163
ex_engine/csrc/ilu_kernel_attention.cpp
Normal file
@@ -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<torch::Tensor>& value,
|
||||
torch::Tensor& key_cache,
|
||||
std::optional<torch::Tensor>& 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<torch::Tensor>& value,
|
||||
torch::Tensor& output,
|
||||
std::optional<torch::Tensor>& output_lse,
|
||||
const std::optional<torch::Tensor>& q_cu_seq_lens,
|
||||
const std::optional<torch::Tensor>& kv_cu_seq_lens,
|
||||
const std::optional<torch::Tensor>& alibi_slope,
|
||||
const std::optional<torch::Tensor>& attn_bias,
|
||||
const std::optional<torch::Tensor>& q_quant_scale,
|
||||
const std::optional<torch::Tensor>& k_quant_scale,
|
||||
const std::optional<torch::Tensor>& 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<double>(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<torch::Tensor>& v_cache,
|
||||
std::optional<torch::Tensor>& output_lse,
|
||||
const std::optional<torch::Tensor>& q_quant_scale,
|
||||
const std::optional<torch::Tensor>& k_cache_quant_scale,
|
||||
const std::optional<torch::Tensor>& v_cache_quant_scale,
|
||||
const std::optional<torch::Tensor>& out_quant_scale,
|
||||
const std::optional<torch::Tensor>& alibi_slope,
|
||||
const std::optional<torch::Tensor>& 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
|
||||
99
ex_engine/csrc/ilu_kernel_fused_moe.cpp
Normal file
99
ex_engine/csrc/ilu_kernel_fused_moe.cpp
Normal file
@@ -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 <glog/logging.h>
|
||||
|
||||
#include "ilu_ops_api.h"
|
||||
|
||||
namespace xllm::kernel::ilu {
|
||||
|
||||
std::tuple<torch::Tensor, torch::Tensor> moe_active_topk(
|
||||
const torch::Tensor& input,
|
||||
int64_t topk,
|
||||
int64_t num_expert_group,
|
||||
int64_t topk_group,
|
||||
bool normalize,
|
||||
const std::optional<torch::Tensor>& mask,
|
||||
const std::string& normed_by,
|
||||
const std::string& scoring_func,
|
||||
double route_scale,
|
||||
const std::optional<torch::Tensor>& 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<torch::Tensor> 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
|
||||
39
ex_engine/csrc/ilu_kernel_group_gemm.cpp
Normal file
39
ex_engine/csrc/ilu_kernel_group_gemm.cpp
Normal file
@@ -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<torch::Tensor>& 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<int64_t>());
|
||||
|
||||
return output;
|
||||
}
|
||||
|
||||
} // namespace xllm::kernel::ilu
|
||||
73
ex_engine/csrc/ilu_kernel_matmul.cpp
Normal file
73
ex_engine/csrc/ilu_kernel_matmul.cpp
Normal file
@@ -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<torch::Tensor> bias) {
|
||||
int64_t act_type = -1;
|
||||
bool persistent = false;
|
||||
std::vector<int64_t> 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
|
||||
51
ex_engine/csrc/ilu_kernel_norm.cpp
Normal file
51
ex_engine/csrc/ilu_kernel_norm.cpp
Normal file
@@ -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<torch::Tensor>& residual,
|
||||
torch::Tensor& weight,
|
||||
std::optional<torch::Tensor>& bias,
|
||||
std::optional<torch::Tensor>& 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<torch::Tensor> fused_bias = std::nullopt;
|
||||
infer::rms_norm(input, weight, output, fused_bias, eps);
|
||||
}
|
||||
|
||||
} // namespace xllm::kernel::ilu
|
||||
31
ex_engine/csrc/ilu_kernel_rope.cpp
Normal file
31
ex_engine/csrc/ilu_kernel_rope.cpp
Normal file
@@ -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
|
||||
189
ex_engine/csrc/ilu_layer_attention.cpp
Normal file
189
ex_engine/csrc/ilu_layer_attention.cpp
Normal file
@@ -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<torch::Tensor, std::optional<torch::Tensor>> AttentionImpl::forward(
|
||||
const AttentionMetadata& attn_metadata,
|
||||
torch::Tensor& query,
|
||||
torch::Tensor& key,
|
||||
torch::Tensor& value,
|
||||
KVCache& kv_cache) {
|
||||
std::optional<torch::Tensor> 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<torch::Tensor> v_cache;
|
||||
std::optional<torch::Tensor> 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<torch::Tensor>& v_cache,
|
||||
const AttentionMetadata& attn_metadata) {
|
||||
int64_t head_size_v = enable_mla_ ? v_head_dim_ : head_size_;
|
||||
std::optional<torch::Tensor> 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<torch::Tensor>& 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<torch::Tensor> 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
|
||||
82
ex_engine/csrc/ilu_layer_attention.h
Normal file
82
ex_engine/csrc/ilu_layer_attention.h
Normal file
@@ -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 <torch/torch.h>
|
||||
|
||||
#include <tuple>
|
||||
|
||||
#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<torch::Tensor, std::optional<torch::Tensor>> 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<torch::Tensor>& v_cache,
|
||||
const AttentionMetadata& attn_metadata);
|
||||
|
||||
void decoder_forward(torch::Tensor& query,
|
||||
torch::Tensor& output,
|
||||
const torch::Tensor& k_cache,
|
||||
const std::optional<torch::Tensor>& 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
|
||||
797
ex_engine/csrc/ilu_layer_fused_moe.cpp
Normal file
797
ex_engine/csrc/ilu_layer_fused_moe.cpp
Normal file
@@ -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 <glog/logging.h>
|
||||
|
||||
#include <iomanip>
|
||||
|
||||
#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<int32_t>(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<int64_t>(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<int64_t>(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<int64_t> 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<torch::Tensor> 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<torch::Tensor> 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<torch::Tensor> 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<torch::Tensor> 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<torch::Tensor> 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<std::string> 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
|
||||
131
ex_engine/csrc/ilu_layer_fused_moe.h
Normal file
131
ex_engine/csrc/ilu_layer_fused_moe.h
Normal file
@@ -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 <torch/torch.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 "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<torch::Tensor> cusum_token_count;
|
||||
std::optional<torch::Tensor> 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<Stream> shared_stream_;
|
||||
std::unique_ptr<Stream> 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
|
||||
123
ex_engine/csrc/moe/fused_moe_cuda.cpp
Normal file
123
ex_engine/csrc/moe/fused_moe_cuda.cpp
Normal file
@@ -0,0 +1,123 @@
|
||||
/* Copyright 2026 The xLLM Authors. All Rights Reserved.
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
https://github.com/jd-opensource/xllm/blob/main/LICENSE
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
==============================================================================*/
|
||||
|
||||
#include "kernels/cuda/cuda_ops_api.h"
|
||||
#include "kernels/cuda/utils.h"
|
||||
#include "platform/device.h"
|
||||
|
||||
namespace xllm::kernel::cuda {
|
||||
|
||||
torch::Tensor cutlass_fused_moe(
|
||||
const torch::Tensor& input, // [num_tokens, hidden]
|
||||
const torch::Tensor& token_selected_experts, // [num_tokens, top_k]
|
||||
const torch::Tensor& token_final_scales, // [num_tokens, top_k]
|
||||
const torch::Tensor&
|
||||
fc1_expert_weights, // [num_experts, inter_dim, hidden]
|
||||
const torch::Tensor&
|
||||
fc2_expert_weights, // [num_experts, hidden, inter_dim]
|
||||
torch::ScalarType output_dtype,
|
||||
const std::vector<torch::Tensor>& quant_scales,
|
||||
int32_t tp_size,
|
||||
int32_t tp_rank,
|
||||
int32_t ep_size,
|
||||
int32_t ep_rank,
|
||||
int32_t cluster_size,
|
||||
int32_t cluster_rank,
|
||||
const std::optional<torch::Tensor>& fc1_expert_biases,
|
||||
const std::optional<torch::Tensor>& fc2_expert_biases,
|
||||
const std::optional<torch::Tensor>& input_sf,
|
||||
const std::optional<torch::Tensor>& swiglu_alpha,
|
||||
const std::optional<torch::Tensor>& swiglu_beta,
|
||||
const std::optional<torch::Tensor>& swiglu_limit,
|
||||
const std::optional<torch::Tensor>& output,
|
||||
bool enable_alltoall,
|
||||
bool use_deepseek_fp8_block_scale,
|
||||
bool use_w4_group_scaling,
|
||||
bool use_mxfp8_act_scaling,
|
||||
bool min_latency_mode,
|
||||
bool use_packed_weights,
|
||||
int32_t tune_max_num_tokens,
|
||||
ActivationType activation_type) {
|
||||
int64_t num_rows = input.size(0);
|
||||
int64_t hidden_size = fc2_expert_weights.size(1);
|
||||
|
||||
if (min_latency_mode) {
|
||||
num_rows *= fc2_expert_weights.size(0);
|
||||
}
|
||||
|
||||
std::vector<int64_t> output_shape = {num_rows, hidden_size};
|
||||
torch::Tensor result_output;
|
||||
if (output.has_value() && output.value().defined()) {
|
||||
result_output = output.value();
|
||||
} else {
|
||||
torch::TensorOptions options = input.options().dtype(output_dtype);
|
||||
result_output = torch::empty(output_shape, options);
|
||||
}
|
||||
|
||||
std::string fused_moe_uri = "fused_moe";
|
||||
if (Device::is_support_sm90a()) {
|
||||
fused_moe_uri += "_90";
|
||||
} else if (Device::is_support_sm100a() || Device::is_support_sm100f()) {
|
||||
fused_moe_uri += "_100";
|
||||
} else if (Device::is_support_sm120a()) {
|
||||
fused_moe_uri += "_120";
|
||||
} else {
|
||||
LOG(FATAL) << "FusedMoE is only supported on sm90, sm100, sm120.";
|
||||
}
|
||||
|
||||
bind_tvmffi_stream_to_current_torch_stream(input.device());
|
||||
|
||||
ffi::Module fused_moe_runner =
|
||||
get_function(fused_moe_uri, "init")(
|
||||
to_dl_data_type(input.scalar_type()),
|
||||
to_dl_data_type(fc1_expert_weights.scalar_type()),
|
||||
to_dl_data_type(output_dtype),
|
||||
use_deepseek_fp8_block_scale,
|
||||
use_w4_group_scaling,
|
||||
use_mxfp8_act_scaling,
|
||||
use_packed_weights)
|
||||
.cast<ffi::Module>();
|
||||
|
||||
fused_moe_runner->GetFunction("run_moe").value()(
|
||||
to_ffi_tensor(result_output),
|
||||
to_ffi_tensor(input),
|
||||
to_ffi_tensor(token_selected_experts),
|
||||
to_ffi_optional_tensor(token_final_scales),
|
||||
to_ffi_tensor(fc1_expert_weights),
|
||||
to_ffi_optional_tensor(fc1_expert_biases),
|
||||
to_ffi_tensor(fc2_expert_weights),
|
||||
to_ffi_optional_tensor(fc2_expert_biases),
|
||||
to_ffi_optional_array_tensors(quant_scales),
|
||||
to_ffi_optional_tensor(input_sf),
|
||||
to_ffi_optional_tensor(swiglu_alpha),
|
||||
to_ffi_optional_tensor(swiglu_beta),
|
||||
to_ffi_optional_tensor(swiglu_limit),
|
||||
tp_size,
|
||||
tp_rank,
|
||||
ep_size,
|
||||
ep_rank,
|
||||
cluster_size,
|
||||
cluster_rank,
|
||||
enable_alltoall,
|
||||
min_latency_mode,
|
||||
/*profile_ids=*/ffi::Optional<ffi::Array<int64_t>>(), // TODO: support
|
||||
// auto tuning
|
||||
// profile ids
|
||||
support_pdl(),
|
||||
activation_type);
|
||||
|
||||
return result_output;
|
||||
}
|
||||
} // namespace xllm::kernel::cuda
|
||||
257
ex_engine/csrc/moe/moeTopKFuncs.cuh
Normal file
257
ex_engine/csrc/moe/moeTopKFuncs.cuh
Normal file
@@ -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 <cooperative_groups.h>
|
||||
#include <cooperative_groups/reduce.h>
|
||||
#include <cub/cub.cuh>
|
||||
|
||||
namespace vllm {
|
||||
namespace moe {
|
||||
namespace reduce_topk {
|
||||
namespace cg = cooperative_groups;
|
||||
static constexpr int kWARP_SIZE = 32;
|
||||
|
||||
template <typename T_>
|
||||
struct TopKRedType {
|
||||
using T = T_;
|
||||
static_assert(
|
||||
std::is_same_v<T, float> || std::is_same_v<T, half> ||
|
||||
std::is_same_v<T, __nv_bfloat16> || std::is_same_v<T, int>,
|
||||
"Top K reduction only implemented for int, float, float16 and bfloat16");
|
||||
|
||||
using TypeCmp = std::conditional_t<sizeof(T) == 4, uint64_t, uint32_t>;
|
||||
using IdxT = std::conditional_t<sizeof(T) == 4, int32_t, int16_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<T>::TwiddleIn(
|
||||
reinterpret_cast<typename cub::Traits<T>::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<int32_t>((cmp & 0xFFFF));
|
||||
|
||||
auto compactTmp = cmp >> kMoveBits;
|
||||
auto valueBits = cub::Traits<T>::TwiddleOut(
|
||||
reinterpret_cast<typename cub::Traits<T>::UnsignedBits&>(compactTmp));
|
||||
value = reinterpret_cast<T&>(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<kWARP_SIZE> const& warp) {
|
||||
return cg::reduce(warp, compValIdx, cg::greater<TypeCmp>{});
|
||||
}
|
||||
};
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
template <int K_, bool Enable_>
|
||||
struct TopKIdx {
|
||||
// by default, empty
|
||||
};
|
||||
|
||||
template <int K_>
|
||||
struct TopKIdx<K_, true> {
|
||||
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 <int N, typename RedType>
|
||||
struct Sort;
|
||||
|
||||
template <typename RedType>
|
||||
struct Sort<1, RedType> {
|
||||
static __device__ void run(RedType* topK) {}
|
||||
};
|
||||
|
||||
template <typename RedType>
|
||||
struct Sort<2, RedType> {
|
||||
static __device__ void run(RedType* topK) { TOPK_SWAP(0, 1); }
|
||||
};
|
||||
|
||||
template <typename RedType>
|
||||
struct Sort<3, RedType> {
|
||||
static __device__ void run(RedType* topK) {
|
||||
TOPK_SWAP(0, 1);
|
||||
TOPK_SWAP(1, 2);
|
||||
TOPK_SWAP(0, 1);
|
||||
}
|
||||
};
|
||||
|
||||
template <typename RedType>
|
||||
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 <int K, typename Type>
|
||||
__forceinline__ __device__ void reduceTopK(
|
||||
cg::thread_block_tile<kWARP_SIZE> 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<Type>;
|
||||
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 <int K, typename Type, int N, bool IsSorted = false>
|
||||
__device__ void reduceTopKFunc(cg::thread_block_tile<kWARP_SIZE> 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<Type>;
|
||||
RedType topK[N];
|
||||
#pragma unroll
|
||||
for (int nn = 0; nn < N; ++nn) {
|
||||
topK[nn] = RedType{value[nn], idx[nn]};
|
||||
}
|
||||
|
||||
if constexpr (!IsSorted) {
|
||||
Sort<N, RedType>::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 <int K, typename Type, int N>
|
||||
__forceinline__ __device__ void reduceTopK(
|
||||
cg::thread_block_tile<kWARP_SIZE> 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<Type>;
|
||||
|
||||
if constexpr (N <= 4) {
|
||||
reduceTopKFunc<K, Type, N>(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<K, Type, 4>(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<K, Type, numResults>(warp, out, outIdx, topKBufferValue,
|
||||
topKBufferIdx, minValue, actualK);
|
||||
}
|
||||
};
|
||||
|
||||
#undef TOPK_SWAP
|
||||
|
||||
} // namespace reduce_topk
|
||||
} // namespace moe
|
||||
} // namespace vllm
|
||||
833
ex_engine/csrc/moe/moe_align_sum_kernels.cu
Normal file
833
ex_engine/csrc/moe/moe_align_sum_kernels.cu
Normal file
@@ -0,0 +1,833 @@
|
||||
#include <array>
|
||||
#include <cub/cub.cuh>
|
||||
|
||||
#include <cuda_runtime.h>
|
||||
#include <torch/csrc/stable/macros.h>
|
||||
#include <torch/csrc/stable/accelerator.h>
|
||||
#include <torch/csrc/stable/ops.h>
|
||||
#include <torch/csrc/stable/tensor.h>
|
||||
#include <torch/headeronly/core/ScalarType.h>
|
||||
|
||||
#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<int32_t, 1024>;
|
||||
__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 <typename scalar_t>
|
||||
__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<int32_t, 1024>;
|
||||
__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 <typename scalar_t, int32_t fill_threads>
|
||||
__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<int32_t>(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 <typename scalar_t>
|
||||
__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 <typename scalar_t>
|
||||
__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 <typename scalar_t>
|
||||
__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 <typename scalar_t, int TOPK>
|
||||
__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 <typename scalar_t, int32_t fill_threads>
|
||||
__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<scalar_t, fill_threads>(
|
||||
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 <typename scalar_t>
|
||||
__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 <typename scalar_t>
|
||||
__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 <typename scalar_t, int32_t fill_threads>
|
||||
__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<scalar_t, fill_threads>(
|
||||
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<torch::stable::Tensor> 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<const scalar_t*>(topk_ids.const_data_ptr()),
|
||||
reinterpret_cast<int32_t*>(sorted_token_ids.mutable_data_ptr()),
|
||||
reinterpret_cast<int32_t*>(experts_ids.mutable_data_ptr()),
|
||||
reinterpret_cast<int32_t*>(
|
||||
num_tokens_post_pad.mutable_data_ptr()),
|
||||
reinterpret_cast<int32_t*>(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<scalar_t>;
|
||||
|
||||
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<const scalar_t*>(topk_ids.const_data_ptr()),
|
||||
reinterpret_cast<int32_t*>(sorted_token_ids.mutable_data_ptr()),
|
||||
reinterpret_cast<int32_t*>(experts_ids.mutable_data_ptr()),
|
||||
reinterpret_cast<int32_t*>(
|
||||
num_tokens_post_pad.mutable_data_ptr()),
|
||||
reinterpret_cast<int32_t*>(expert_map.mutable_data_ptr()),
|
||||
num_experts, padded_num_experts, experts_per_warp, block_size,
|
||||
topk_ids.numel(),
|
||||
reinterpret_cast<int32_t*>(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<scalar_t>;
|
||||
sort_kernel<<<gridDims, block_threads, 0, stream>>>(
|
||||
reinterpret_cast<const scalar_t*>(topk_ids.const_data_ptr()),
|
||||
reinterpret_cast<int32_t*>(sorted_token_ids.mutable_data_ptr()),
|
||||
reinterpret_cast<int32_t*>(cumsum_buffer.mutable_data_ptr()),
|
||||
reinterpret_cast<int32_t*>(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<const int32_t*>(batch_num_tokens.const_data_ptr()),
|
||||
reinterpret_cast<int32_t*>(sorted_ids.mutable_data_ptr()),
|
||||
reinterpret_cast<int32_t*>(batch_ids.mutable_data_ptr()),
|
||||
reinterpret_cast<int32_t*>(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<scalar_t, 2><<<grid, block, 0, stream>>>(
|
||||
reinterpret_cast<scalar_t*>(output.mutable_data_ptr()),
|
||||
reinterpret_cast<const scalar_t*>(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<scalar_t, 3><<<grid, block, 0, stream>>>(
|
||||
reinterpret_cast<scalar_t*>(output.mutable_data_ptr()),
|
||||
reinterpret_cast<const scalar_t*>(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<scalar_t, 4><<<grid, block, 0, stream>>>(
|
||||
reinterpret_cast<scalar_t*>(output.mutable_data_ptr()),
|
||||
reinterpret_cast<const scalar_t*>(input.const_data_ptr()),
|
||||
hidden_size);
|
||||
});
|
||||
break;
|
||||
|
||||
default:
|
||||
torch::stable::sum_out(output, input, std::array<int64_t, 1>{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<torch::stable::Tensor> 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<<<max_loras + 1, blockDim, shared_mem, stream>>>(
|
||||
reinterpret_cast<scalar_t*>(topk_ids.mutable_data_ptr()),
|
||||
reinterpret_cast<int32_t*>(token_lora_mapping.mutable_data_ptr()),
|
||||
block_size,
|
||||
reinterpret_cast<int32_t*>(expert_map.mutable_data_ptr()),
|
||||
num_experts, max_loras, topk_ids.numel(), max_num_tokens_padded,
|
||||
max_num_m_blocks,
|
||||
reinterpret_cast<int32_t*>(sorted_token_ids.mutable_data_ptr()),
|
||||
reinterpret_cast<int32_t*>(expert_ids.mutable_data_ptr()),
|
||||
topk_num,
|
||||
reinterpret_cast<int32_t*>(
|
||||
num_tokens_post_pad.mutable_data_ptr()),
|
||||
reinterpret_cast<int32_t*>(adapter_enabled.mutable_data_ptr()),
|
||||
reinterpret_cast<int32_t*>(lora_ids.mutable_data_ptr()),
|
||||
reinterpret_cast<int32_t*>(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<scalar_t>;
|
||||
|
||||
// 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<scalar_t*>(topk_ids.mutable_data_ptr()),
|
||||
reinterpret_cast<int32_t*>(token_lora_mapping.mutable_data_ptr()),
|
||||
block_size,
|
||||
reinterpret_cast<int32_t*>(expert_map.mutable_data_ptr()),
|
||||
num_experts, max_loras, topk_ids.numel(), max_num_tokens_padded,
|
||||
max_num_m_blocks,
|
||||
reinterpret_cast<int32_t*>(sorted_token_ids.mutable_data_ptr()),
|
||||
reinterpret_cast<int32_t*>(expert_ids.mutable_data_ptr()),
|
||||
topk_num,
|
||||
reinterpret_cast<int32_t*>(
|
||||
num_tokens_post_pad.mutable_data_ptr()),
|
||||
reinterpret_cast<int32_t*>(adapter_enabled.mutable_data_ptr()),
|
||||
reinterpret_cast<int32_t*>(cumsum.mutable_data_ptr()), WARP_SIZE,
|
||||
padded_num_experts,
|
||||
reinterpret_cast<int32_t*>(lora_ids.mutable_data_ptr()),
|
||||
reinterpret_cast<int32_t*>(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<scalar_t>;
|
||||
|
||||
sort_kernel<<<gridDims, block_threads, 0, stream>>>(
|
||||
reinterpret_cast<const scalar_t*>(topk_ids.const_data_ptr()),
|
||||
reinterpret_cast<int32_t*>(sorted_token_ids.mutable_data_ptr()),
|
||||
reinterpret_cast<int32_t*>(cumsum.mutable_data_ptr()),
|
||||
reinterpret_cast<int32_t*>(expert_map.mutable_data_ptr()),
|
||||
topk_ids.numel(), num_experts, max_num_tokens_padded, topk_num,
|
||||
reinterpret_cast<int32_t*>(token_mask.mutable_data_ptr()),
|
||||
max_loras,
|
||||
reinterpret_cast<int32_t*>(lora_ids.mutable_data_ptr()),
|
||||
reinterpret_cast<int32_t*>(adapter_enabled.mutable_data_ptr()),
|
||||
has_expert_map);
|
||||
}
|
||||
});
|
||||
}
|
||||
56
ex_engine/csrc/moe/moe_fused_topk.cu
Normal file
56
ex_engine/csrc/moe/moe_fused_topk.cu
Normal file
@@ -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<torch::Tensor, torch::Tensor> moe_fused_topk(
|
||||
torch::Tensor& gating_output,
|
||||
int64_t topk,
|
||||
bool renormalize,
|
||||
const std::optional<torch::Tensor>& 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<torch::Tensor> 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
|
||||
87
ex_engine/csrc/moe/moe_ops.h
Normal file
87
ex_engine/csrc/moe/moe_ops.h
Normal file
@@ -0,0 +1,87 @@
|
||||
#pragma once
|
||||
|
||||
#include <torch/csrc/stable/tensor.h>
|
||||
|
||||
#include <optional>
|
||||
#include <tuple>
|
||||
|
||||
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<torch::stable::Tensor> 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<torch::stable::Tensor> 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<torch::stable::Tensor>& correction_bias,
|
||||
const std::optional<torch::stable::Tensor>& input_ids,
|
||||
const std::optional<torch::stable::Tensor>& 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<torch::stable::Tensor> 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<torch::stable::Tensor> 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<torch::stable::Tensor> b_qzeros,
|
||||
std::optional<torch::stable::Tensor> 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<torch::stable::Tensor, torch::stable::Tensor> 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
|
||||
285
ex_engine/csrc/moe/moe_topk.cuh
Normal file
285
ex_engine/csrc/moe/moe_topk.cuh
Normal file
@@ -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 <cooperative_groups.h>
|
||||
#include <cooperative_groups/reduce.h>
|
||||
|
||||
#include <cub/cub.cuh>
|
||||
|
||||
#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 <typename T_>
|
||||
struct TopKRedType {
|
||||
using T = T_;
|
||||
static_assert(
|
||||
std::is_same_v<T, float> || std::is_same_v<T, half> ||
|
||||
std::is_same_v<T, __nv_bfloat16> || std::is_same_v<T, int>,
|
||||
"Top K reduction only implemented for int, float, float16 and bfloat16");
|
||||
|
||||
using TypeCmp = std::conditional_t<sizeof(T) == 4, uint64_t, uint32_t>;
|
||||
using IdxT = std::conditional_t<sizeof(T) == 4, int32_t, int16_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<T>::TwiddleIn(
|
||||
reinterpret_cast<typename cub::Traits<T>::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<int32_t>((cmp & 0xFFFF));
|
||||
|
||||
auto compactTmp = cmp >> kMoveBits;
|
||||
auto valueBits = cub::Traits<T>::TwiddleOut(
|
||||
reinterpret_cast<typename cub::Traits<T>::UnsignedBits&>(compactTmp));
|
||||
value = reinterpret_cast<T&>(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<kWARP_SIZE> const& warp) {
|
||||
if constexpr (!kTLLM_GEN_HAS_FAST_REDUX || sizeof(TypeCmp) == 8) {
|
||||
return cg::reduce(warp, compValIdx, cg::greater<TypeCmp>{});
|
||||
} else {
|
||||
TypeCmp result;
|
||||
asm("redux.sync.max.u32 %0, %1, 0xffffffff;\n"
|
||||
: "=r"(result)
|
||||
: "r"(compValIdx));
|
||||
return result;
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
template <int K_, bool Enable_>
|
||||
struct TopKIdx {
|
||||
// by default, empty
|
||||
};
|
||||
|
||||
template <int K_>
|
||||
struct TopKIdx<K_, true> {
|
||||
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 <int N, typename RedType>
|
||||
struct Sort;
|
||||
|
||||
template <typename RedType>
|
||||
struct Sort<1, RedType> {
|
||||
static __device__ void run(RedType* topK) {}
|
||||
};
|
||||
|
||||
template <typename RedType>
|
||||
struct Sort<2, RedType> {
|
||||
static __device__ void run(RedType* topK) { TOPK_SWAP(0, 1); }
|
||||
};
|
||||
|
||||
template <typename RedType>
|
||||
struct Sort<3, RedType> {
|
||||
static __device__ void run(RedType* topK) {
|
||||
TOPK_SWAP(0, 1);
|
||||
TOPK_SWAP(1, 2);
|
||||
TOPK_SWAP(0, 1);
|
||||
}
|
||||
};
|
||||
|
||||
template <typename RedType>
|
||||
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 <int K, typename Type>
|
||||
__forceinline__ __device__ void reduceTopK(
|
||||
cg::thread_block_tile<kWARP_SIZE> 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<Type>;
|
||||
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 <int K, typename Type, int N, bool IsSorted = false>
|
||||
__device__ void reduceTopKFunc(cg::thread_block_tile<kWARP_SIZE> 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<Type>;
|
||||
RedType topK[N];
|
||||
#pragma unroll
|
||||
for (int nn = 0; nn < N; ++nn) {
|
||||
topK[nn] = RedType{value[nn], idx[nn]};
|
||||
}
|
||||
|
||||
if constexpr (!IsSorted) {
|
||||
Sort<N, RedType>::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 <int K, typename Type, int N>
|
||||
__forceinline__ __device__ void reduceTopK(
|
||||
cg::thread_block_tile<kWARP_SIZE> 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<Type>;
|
||||
|
||||
if constexpr (N <= 4) {
|
||||
reduceTopKFunc<K, Type, N>(
|
||||
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<K, Type, 4>(
|
||||
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<K, Type, numResults>(
|
||||
warp, out, outIdx, topKBufferValue, topKBufferIdx, minValue, actualK);
|
||||
}
|
||||
};
|
||||
|
||||
#undef TOPK_SWAP
|
||||
|
||||
} // namespace reduce_topk
|
||||
} // namespace xllm::kernel::cuda
|
||||
602
ex_engine/csrc/moe/moe_topk_sigmoid_kernels.cuh
Normal file
602
ex_engine/csrc/moe/moe_topk_sigmoid_kernels.cuh
Normal file
@@ -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 <ATen/cuda/CUDAContext.h>
|
||||
#include <c10/cuda/CUDAGuard.h>
|
||||
#include <torch/all.h>
|
||||
|
||||
#include <cub/util_type.cuh>
|
||||
#include <cuda/functional>
|
||||
|
||||
#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 <typename T, int TPB>
|
||||
__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<T>(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 <int TPB>
|
||||
__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<int, float>;
|
||||
using BlockReduce = cub::BlockReduce<cub_kvp, TPB>;
|
||||
__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 <typename T,
|
||||
int VPT,
|
||||
int NUM_EXPERTS,
|
||||
int WARPS_PER_CTA,
|
||||
int BYTES_PER_LDG>
|
||||
__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<T, ELTS_PER_LDG>;
|
||||
|
||||
// Finally, we pull in the data from global mem
|
||||
T row_chunk_temp[VPT];
|
||||
AccessType* row_chunk_vec_ptr =
|
||||
reinterpret_cast<AccessType*>(&row_chunk_temp);
|
||||
const AccessType* vec_thread_read_ptr =
|
||||
reinterpret_cast<const AccessType*>(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<T>(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 <typename T, int EXPERTS, int WARPS_PER_TB>
|
||||
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<T, EXPERTS, BYTES_PER_LDG>;
|
||||
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<T, VPT, EXPERTS, WARPS_PER_TB, BYTES_PER_LDG>
|
||||
<<<num_blocks, block_dim, 0, stream>>>(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<TYPE, NUM_EXPERTS, WARPS_PER_TB>( \
|
||||
gating_output, \
|
||||
nullptr, \
|
||||
topk_weights, \
|
||||
topk_indices, \
|
||||
num_tokens, \
|
||||
topk, \
|
||||
0, \
|
||||
num_experts, \
|
||||
renormalize, \
|
||||
correction_bias, \
|
||||
stream);
|
||||
|
||||
template <typename T>
|
||||
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<T, TPB><<<num_tokens, TPB, 0, stream>>>(gating_output,
|
||||
nullptr,
|
||||
sigmoid_workspace,
|
||||
num_experts,
|
||||
correction_bias);
|
||||
moe_topK<TPB><<<num_tokens, TPB, 0, stream>>>(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<torch::Tensor>& 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<int>(gating_output.size(-1));
|
||||
const int num_tokens = static_cast<int>(gating_output.size(0));
|
||||
const int topk = static_cast<int>(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<float>();
|
||||
}
|
||||
|
||||
if (dtype == at::ScalarType::Float) {
|
||||
topk_gating_sigmoid_kernel_launcher<float>(
|
||||
gating_output.data_ptr<float>(),
|
||||
topk_weights.data_ptr<float>(),
|
||||
topk_indices.data_ptr<int>(),
|
||||
sigmoid_workspace.data_ptr<float>(),
|
||||
num_tokens,
|
||||
num_experts,
|
||||
topk,
|
||||
renormalize,
|
||||
bias_ptr,
|
||||
stream);
|
||||
} else if (dtype == at::ScalarType::Half) {
|
||||
topk_gating_sigmoid_kernel_launcher<__half>(
|
||||
reinterpret_cast<const __half*>(gating_output.data_ptr<at::Half>()),
|
||||
topk_weights.data_ptr<float>(),
|
||||
topk_indices.data_ptr<int>(),
|
||||
sigmoid_workspace.data_ptr<float>(),
|
||||
num_tokens,
|
||||
num_experts,
|
||||
topk,
|
||||
renormalize,
|
||||
bias_ptr,
|
||||
stream);
|
||||
} else if (dtype == at::ScalarType::BFloat16) {
|
||||
topk_gating_sigmoid_kernel_launcher<__nv_bfloat16>(
|
||||
reinterpret_cast<const __nv_bfloat16*>(
|
||||
gating_output.data_ptr<at::BFloat16>()),
|
||||
topk_weights.data_ptr<float>(),
|
||||
topk_indices.data_ptr<int>(),
|
||||
sigmoid_workspace.data_ptr<float>(),
|
||||
num_tokens,
|
||||
num_experts,
|
||||
topk,
|
||||
renormalize,
|
||||
bias_ptr,
|
||||
stream);
|
||||
} else {
|
||||
LOG(FATAL) << "Unsupported gating_output dtype: " << dtype;
|
||||
}
|
||||
}
|
||||
} // namespace xllm::kernel::cuda
|
||||
@@ -22,11 +22,9 @@ limitations under the License.
|
||||
#include <torch/all.h>
|
||||
|
||||
#include <cub/util_type.cuh>
|
||||
#if CUDA_VERSION >= 12090
|
||||
#include <cuda/functional>
|
||||
#endif
|
||||
|
||||
#include "device_utils.cuh"
|
||||
#include "kernels/cuda/device_utils.cuh"
|
||||
|
||||
using cub_kvp = cub::KeyValuePair<int, float>;
|
||||
|
||||
@@ -707,7 +705,8 @@ void topk_gating_softmax_kernel_launcher(const T* gating_output,
|
||||
LAUNCH_SOFTMAX(T, 256, WARPS_PER_TB);
|
||||
break;
|
||||
default: {
|
||||
TORCH_CHECK(softmax_workspace != nullptr, "softmax_workspace must be provided for num_experts that are ");
|
||||
CHECK(softmax_workspace != nullptr)
|
||||
<< "softmax_workspace must be provided for num_experts that are "
|
||||
"not a power of 2.";
|
||||
static constexpr int TPB = 256;
|
||||
moe_softmax<T, TPB><<<num_tokens, TPB, 0, stream>>>(gating_output,
|
||||
@@ -753,23 +752,29 @@ void topk_softmax(torch::Tensor& topk_weights, // [num_tokens, topk]
|
||||
const double moe_softcapping,
|
||||
const std::optional<torch::Tensor>& correction_bias) {
|
||||
// Check data type
|
||||
TORCH_CHECK(gating_output.scalar_type() == at::ScalarType::Float ||
|
||||
CHECK(gating_output.scalar_type() == at::ScalarType::Float ||
|
||||
gating_output.scalar_type() == at::ScalarType::Half ||
|
||||
gating_output.scalar_type() == at::ScalarType::BFloat16,
|
||||
"gating_output must be float32, float16, or bfloat16");
|
||||
gating_output.scalar_type() == at::ScalarType::BFloat16)
|
||||
<< "gating_output must be float32, float16, or bfloat16";
|
||||
|
||||
// Check dimensions
|
||||
TORCH_CHECK(gating_output.dim() == 2, "gating_output must be 2D tensor [num_tokens, num_experts]");
|
||||
TORCH_CHECK(topk_weights.dim() == 2, "topk_weights must be 2D tensor [num_tokens, topk]");
|
||||
TORCH_CHECK(topk_indices.dim() == 2, "topk_indices must be 2D tensor [num_tokens, topk]");
|
||||
CHECK(gating_output.dim() == 2)
|
||||
<< "gating_output must be 2D tensor [num_tokens, num_experts]";
|
||||
CHECK(topk_weights.dim() == 2)
|
||||
<< "topk_weights must be 2D tensor [num_tokens, topk]";
|
||||
CHECK(topk_indices.dim() == 2)
|
||||
<< "topk_indices must be 2D tensor [num_tokens, topk]";
|
||||
|
||||
// Check shapes
|
||||
TORCH_CHECK(gating_output.size(0) == topk_weights.size(0), "First dimension of topk_weights must match num_tokens in ");
|
||||
CHECK(gating_output.size(0) == topk_weights.size(0))
|
||||
<< "First dimension of topk_weights must match num_tokens in "
|
||||
"gating_output"
|
||||
<< "First dimension of topk_indices must match num_tokens in "
|
||||
"gating_output";
|
||||
|
||||
TORCH_CHECK(topk_weights.size(-1) == topk_indices.size(-1), "Second dimension of topk_indices must match topk in topk_weights topk must be less than or equal to num_experts");
|
||||
CHECK(topk_weights.size(-1) == topk_indices.size(-1))
|
||||
<< "Second dimension of topk_indices must match topk in topk_weights"
|
||||
<< "topk must be less than or equal to num_experts";
|
||||
|
||||
const int num_experts = static_cast<int>(gating_output.size(-1));
|
||||
const int num_tokens = static_cast<int>(gating_output.size(0));
|
||||
@@ -791,9 +796,12 @@ void topk_softmax(torch::Tensor& topk_weights, // [num_tokens, topk]
|
||||
const float* bias_ptr = nullptr;
|
||||
if (correction_bias.has_value()) {
|
||||
const torch::Tensor& bias_tensor = correction_bias.value();
|
||||
TORCH_CHECK(bias_tensor.dim() == 1, "correction_bias must be 1D tensor [num_experts]");
|
||||
TORCH_CHECK(bias_tensor.size(0) == num_experts, "correction_bias size must match num_experts");
|
||||
TORCH_CHECK(bias_tensor.scalar_type() == at::ScalarType::Float, "correction_bias must be float32, got ");
|
||||
CHECK(bias_tensor.dim() == 1)
|
||||
<< "correction_bias must be 1D tensor [num_experts]";
|
||||
CHECK(bias_tensor.size(0) == num_experts)
|
||||
<< "correction_bias size must match num_experts";
|
||||
CHECK(bias_tensor.scalar_type() == at::ScalarType::Float)
|
||||
<< "correction_bias must be float32, got " << bias_tensor.scalar_type();
|
||||
bias_ptr = bias_tensor.data_ptr<float>();
|
||||
}
|
||||
|
||||
@@ -841,7 +849,7 @@ void topk_softmax(torch::Tensor& topk_weights, // [num_tokens, topk]
|
||||
bias_ptr,
|
||||
stream);
|
||||
} else {
|
||||
TORCH_CHECK(false, "Unsupported gating_output dtype");
|
||||
LOG(FATAL) << "Unsupported gating_output dtype: " << dtype;
|
||||
}
|
||||
}
|
||||
} // namespace xllm::kernel::cuda
|
||||
|
||||
153
ex_engine/include/ilu_ops_api.h
Normal file
153
ex_engine/include/ilu_ops_api.h
Normal file
@@ -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 <ATen/DynamicLibrary.h>
|
||||
#include <ATen/core/dispatch/Dispatcher.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <glog/logging.h>
|
||||
#include <torch/all.h>
|
||||
|
||||
#include <optional>
|
||||
|
||||
#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<torch::Tensor>& value, // (num_tokens, num_heads, head_size)
|
||||
torch::Tensor& key_cache, // (num_blocks, num_heads, block_size, head_size)
|
||||
std::optional<torch::Tensor>&
|
||||
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<torch::Tensor>& value,
|
||||
torch::Tensor& output,
|
||||
std::optional<torch::Tensor>& output_lse,
|
||||
const std::optional<torch::Tensor>& q_cu_seq_lens,
|
||||
const std::optional<torch::Tensor>& kv_cu_seq_lens,
|
||||
const std::optional<torch::Tensor>& alibi_slope,
|
||||
const std::optional<torch::Tensor>& attn_bias,
|
||||
const std::optional<torch::Tensor>& q_quant_scale,
|
||||
const std::optional<torch::Tensor>& k_quant_scale,
|
||||
const std::optional<torch::Tensor>& 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<torch::Tensor>& v_cache,
|
||||
std::optional<torch::Tensor>& output_lse,
|
||||
const std::optional<torch::Tensor>& q_quant_scale,
|
||||
const std::optional<torch::Tensor>& k_cache_quant_scale,
|
||||
const std::optional<torch::Tensor>& v_cache_quant_scale,
|
||||
const std::optional<torch::Tensor>& out_quant_scale,
|
||||
const std::optional<torch::Tensor>& alibi_slope,
|
||||
const std::optional<torch::Tensor>& 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<torch::Tensor>& residual,
|
||||
torch::Tensor& weight,
|
||||
std::optional<torch::Tensor>& bias,
|
||||
std::optional<torch::Tensor>& 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<torch::Tensor> bias);
|
||||
|
||||
std::tuple<torch::Tensor, torch::Tensor> moe_active_topk(
|
||||
const torch::Tensor& input,
|
||||
int64_t topk,
|
||||
int64_t num_expert_group,
|
||||
int64_t topk_group,
|
||||
bool normalize,
|
||||
const std::optional<torch::Tensor>& mask,
|
||||
const std::string& normed_by,
|
||||
const std::string& scoring_func,
|
||||
double route_scale,
|
||||
const std::optional<torch::Tensor>& e_score_correction_bias);
|
||||
|
||||
std::vector<torch::Tensor> 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<torch::Tensor>& dst_to_src,
|
||||
torch::Tensor& output);
|
||||
|
||||
torch::Tensor moe_combine_result(torch::Tensor& input, torch::Tensor& weight);
|
||||
} // namespace xllm::kernel::ilu
|
||||
63
ex_engine/include/ilu_utils.h
Normal file
63
ex_engine/include/ilu_utils.h
Normal file
@@ -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
|
||||
147
ex_engine/include/ixformer.h
Normal file
147
ex_engine/include/ixformer.h
Normal file
@@ -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 <torch/all.h>
|
||||
|
||||
#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<torch::Tensor>& alibi_slopes,
|
||||
const std::optional<torch::Tensor>& sinks,
|
||||
std::optional<torch::Tensor>& 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<torch::Tensor>& 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<torch::Tensor>& sinks);
|
||||
|
||||
torch::Tensor ixformer_linear(torch::Tensor& input,
|
||||
torch::Tensor& weight,
|
||||
int64_t act_type,
|
||||
const std::optional<torch::Tensor>& bias,
|
||||
const std::optional<torch::Tensor>& out,
|
||||
const std::optional<bool> persistent);
|
||||
|
||||
torch::Tensor ixformer_linear_ex(torch::Tensor& input,
|
||||
torch::Tensor& weight,
|
||||
const c10::optional<torch::Tensor>& bias,
|
||||
const c10::optional<torch::Tensor>& 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<torch::Tensor>& fused_bias,
|
||||
double alpha,
|
||||
double eps,
|
||||
bool is_post);
|
||||
|
||||
void rms_norm(torch::Tensor& input,
|
||||
torch::Tensor& weight,
|
||||
torch::Tensor& output,
|
||||
const std::optional<torch::Tensor>& 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<torch::Tensor>& expert_mask,
|
||||
const c10::optional<torch::Tensor>& expert_sizes_cpu,
|
||||
const c10::optional<torch::Tensor>& expand_tokens_gpu,
|
||||
int64_t start_expert_id,
|
||||
int64_t end_expert_id,
|
||||
int64_t num_experts);
|
||||
|
||||
void moe_expand_input(torch::Tensor outputs,
|
||||
torch::Tensor inputs,
|
||||
torch::Tensor dst_to_src,
|
||||
const c10::optional<torch::Tensor>& src_to_dst,
|
||||
int64_t dst_tokens,
|
||||
int64_t expand_factor);
|
||||
|
||||
void moe_w16a16_group_gemm(torch::Tensor output,
|
||||
torch::Tensor inputs,
|
||||
torch::Tensor weights,
|
||||
torch::Tensor tokens_per_experts,
|
||||
const c10::optional<torch::Tensor>& dst_to_src,
|
||||
const c10::optional<torch::Tensor>& bias,
|
||||
std::string format,
|
||||
int64_t persistent,
|
||||
int64_t output_n);
|
||||
|
||||
void moe_output_reduce_sum(torch::Tensor outputs,
|
||||
torch::Tensor inputs,
|
||||
const c10::optional<torch::Tensor>& mul_weight,
|
||||
const c10::optional<torch::Tensor>& mask,
|
||||
const c10::optional<torch::Tensor>& extra_residual,
|
||||
double scaling_factor);
|
||||
} // namespace ixformer::infer
|
||||
Reference in New Issue
Block a user