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

[init] baseline7 from project_6

Browse Source
This commit is contained in:
root
2026-08-25 07:05:49 +00:00
commit 0bb337bc00
1392 changed files with 321187 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

236
ex_engine/xllm_layers/mlu/qwen3_5_attention.cpp Normal file
View File

@@ -0,0 +1,236 @@
/* 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 "qwen3_5_attention.h"
#include <glog/logging.h>
#include <tuple>
#include "kernels/ops_api.h"
namespace xllm {
namespace layer {
Qwen3_5AttentionImpl::Qwen3_5AttentionImpl(const ModelArgs& args,
const QuantArgs& quant_args,
const ParallelArgs& parallel_args,
const torch::TensorOptions& options,
int32_t layer_id) {
const int64_t tp_size = parallel_args.tp_group_->world_size();
const int64_t total_num_heads = args.n_heads();
const int64_t total_num_kv_heads = args.n_kv_heads().value_or(args.n_heads());
layer_id_ = layer_id;
rank_ = parallel_args.tp_group_->rank();
CHECK(total_num_heads % tp_size == 0);
num_heads_ = total_num_heads / tp_size;
if (total_num_kv_heads >= tp_size) {
CHECK(total_num_kv_heads % tp_size == 0);
num_kv_heads_ = total_num_kv_heads / tp_size;
num_kv_head_replicas_ = 1;
} else {
CHECK(tp_size % total_num_kv_heads == 0);
num_kv_heads_ = 1;
num_kv_head_replicas_ = tp_size / total_num_kv_heads;
}
head_dim_ = args.head_dim();
q_size_ = num_heads_ * head_dim_;
kv_size_ = num_kv_heads_ * head_dim_;
scaling_ = 1.0f / std::sqrt(static_cast<float>(head_dim_));
attn_output_gate_ = args.attn_output_gate();
mrope_cu_seq_lens_ = torch::zeros(2, torch::kInt32).to(options.device());
// 1. QKV linear
qkv_proj_ = register_module(
"qkv_proj",
QKVParallelLinear(args.hidden_size(),
attn_output_gate_ ? num_heads_ * 2 : num_heads_,
num_kv_heads_,
args.head_dim(),
num_kv_head_replicas_,
/*bias=*/args.attention_bias(),
/*gather_output=*/false,
parallel_args,
options));
// 2. O proj
o_proj_ = register_module("o_proj",
RowParallelLinear(total_num_heads * head_dim_,
args.hidden_size(),
/*bias=*/false,
/*input_is_parallelized=*/true,
/*if_reduce_results=*/true,
quant_args,
parallel_args.tp_group_,
options));
// 3. Q norm
q_norm_ = register_module(
"q_norm", Qwen3NextRMSNorm(head_dim_, args.rms_norm_eps(), options));
// 4. K norm
k_norm_ = register_module(
"k_norm", Qwen3NextRMSNorm(head_dim_, args.rms_norm_eps(), options));
// 5. Attention
attn_ = register_module("attn",
Attention(num_heads_,
head_dim_,
scaling_,
num_kv_heads_,
args.sliding_window()));
// 6. Rotary embedding
const int32_t rotary_dim =
static_cast<int32_t>(head_dim_ * args.partial_rotary_factor());
rotary_emb_ =
register_module("rope",
MRotaryEmbedding(rotary_dim,
args.max_position_embeddings(),
args.rope_theta(),
/*interleaved=*/false,
args.rope_scaling_mrope_section(),
options));
}
void Qwen3_5AttentionImpl::rotary_emb_forward(
torch::Tensor& q,
torch::Tensor& k,
const torch::Tensor& positions,
const AttentionMetadata& attn_metadata) {
auto q_shape = q.sizes();
auto k_shape = k.sizes();
auto num_tokens = positions.size(-1);
mrope_cu_seq_lens_[1] = num_tokens;
xllm::kernel::RotaryParams rotary_params;
bool only_prefill =
(attn_metadata.is_prefill || attn_metadata.is_chunked_prefill);
if (only_prefill) {
rotary_params.sin = attn_metadata.mrope_sin;
rotary_params.cos = attn_metadata.mrope_cos;
rotary_params.position_ids = std::nullopt;
rotary_params.cu_query_lens = mrope_cu_seq_lens_;
rotary_params.interleaved = false;
rotary_params.discrete = false;
rotary_params.max_query_len = num_tokens;
rotary_params.q = q.view({num_tokens, -1, head_dim_});
xllm::kernel::apply_rotary(rotary_params);
q = rotary_params.q.reshape(q_shape);
rotary_params.q = k.view({num_tokens, -1, head_dim_});
xllm::kernel::apply_rotary(rotary_params);
k = rotary_params.q.reshape(k_shape);
} else {
if (positions.dim() == 2) {
rotary_params.position_ids = positions[0];
} else {
rotary_params.position_ids = positions;
}
rotary_params.sin = rotary_emb_->get_sin_cache();
rotary_params.cos = rotary_emb_->get_cos_cache();
rotary_params.interleaved = false;
rotary_params.discrete = true;
rotary_params.max_query_len = num_tokens;
rotary_params.q = q.view({1, num_tokens, -1, head_dim_});
xllm::kernel::apply_rotary(rotary_params);
q = rotary_params.q.reshape(q_shape);
rotary_params.q = k.view({1, num_tokens, -1, head_dim_});
xllm::kernel::apply_rotary(rotary_params);
k = rotary_params.q.reshape(k_shape);
}
}
torch::Tensor Qwen3_5AttentionImpl::forward(
const torch::Tensor& positions,
const torch::Tensor& hidden_states,
const AttentionMetadata& attn_metadata,
KVCache& kv_cache) {
// 1. qkv projection
auto qkv = qkv_proj_->forward(hidden_states);
torch::Tensor q, k, v;
torch::Tensor gate;
if (attn_output_gate_) {
// Split qkv for attn_output_gate case: [q_size*2, kv_size, kv_size]
auto q_gate = qkv.slice(/*dim=*/-1, 0, q_size_ * 2);
k = qkv.slice(/*dim=*/-1, q_size_ * 2, q_size_ * 2 + kv_size_);
v = qkv.slice(
/*dim=*/-1, q_size_ * 2 + kv_size_, q_size_ * 2 + kv_size_ * 2);
v = v.contiguous();
std::vector<int64_t> orig_shape;
for (int64_t i = 0; i < q_gate.dim() - 1; i++) {
orig_shape.push_back(q_gate.size(i));
}
std::vector<int64_t> new_shape = orig_shape;
new_shape.push_back(num_heads_);
new_shape.push_back(-1);
torch::Tensor q_gate_reshaped = q_gate.reshape(new_shape);
auto chunks = torch::chunk(q_gate_reshaped, 2, /*dim=*/-1);
q = chunks[0];
gate = chunks[1];
std::vector<int64_t> q_new_shape = orig_shape;
q_new_shape.push_back(-1);
q = q.reshape(q_new_shape);
std::vector<int64_t> gate_new_shape = orig_shape;
gate_new_shape.push_back(-1);
gate = gate.reshape(gate_new_shape);
} else {
// Normal case: [q_size, kv_size, kv_size]
q = qkv.slice(/*dim=*/-1, 0, q_size_);
k = qkv.slice(/*dim=*/-1, q_size_, q_size_ + kv_size_);
v = qkv.slice(/*dim=*/-1, q_size_ + kv_size_, q_size_ + 2 * kv_size_);
}
const int64_t T = q.size(0);
auto q_reshaped = q.reshape({T, num_heads_, head_dim_});
auto q_normed = std::get<0>(q_norm_->forward(q_reshaped));
auto k_reshaped = k.reshape({T, num_kv_heads_, head_dim_});
auto k_normed = std::get<0>(k_norm_->forward(k_reshaped));
q = q_normed.view({T, q_size_});
k = k_normed.view({T, kv_size_});
rotary_emb_forward(q, k, positions, attn_metadata);
auto out = std::get<0>(attn_->forward(attn_metadata, q, k, v, kv_cache));
if (attn_output_gate_) {
gate = torch::sigmoid(gate);
out = out * gate;
}
out = o_proj_->forward(out);
return out;
}
void Qwen3_5AttentionImpl::load_state_dict(const StateDict& state_dict) {
qkv_proj_->load_state_dict(state_dict, {"q_proj.", "k_proj.", "v_proj."});
o_proj_->load_state_dict(state_dict.get_dict_with_prefix("o_proj."));
if (auto w = state_dict.get_tensor("q_norm.weight"); w.defined()) {
q_norm_->load_state_dict(StateDict({{"weight", w}}));
}
if (auto w = state_dict.get_tensor("k_norm.weight"); w.defined()) {
k_norm_->load_state_dict(StateDict({{"weight", w}}));
}
}
} // namespace layer
} // namespace xllm

79
ex_engine/xllm_layers/mlu/qwen3_5_attention.h Normal file
View File

@@ -0,0 +1,79 @@
/* 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 "attention.h"
#include "framework/kv_cache/kv_cache.h"
#include "framework/model/model_args.h"
#include "framework/parallel_state/parallel_args.h"
#include "framework/quant_args.h"
#include "framework/state_dict/state_dict.h"
#include "layers/common/linear.h"
#include "layers/common/partial_rotary_embedding.h"
#include "layers/common/qwen3_next_rms_norm.h"
#include "layers/common/rotary_embedding.h"
namespace xllm {
namespace layer {
class Qwen3_5AttentionImpl : public torch::nn::Module {
public:
Qwen3_5AttentionImpl() = default;
Qwen3_5AttentionImpl(const ModelArgs& args,
const QuantArgs& quant_args,
const ParallelArgs& parallel_args,
const torch::TensorOptions& options,
int32_t layer_id);
torch::Tensor forward(const torch::Tensor& positions,
const torch::Tensor& hidden_states,
const AttentionMetadata& attn_metadata,
KVCache& kv_cache);
void load_state_dict(const StateDict& state_dict);
void rotary_emb_forward(torch::Tensor& q,
torch::Tensor& k,
const torch::Tensor& positions,
const AttentionMetadata& attn_metadata);
private:
int64_t num_heads_;
int64_t num_kv_heads_;
int64_t num_kv_head_replicas_;
int64_t head_dim_;
int64_t q_size_;
int64_t kv_size_;
float scaling_;
bool attn_output_gate_;
int32_t layer_id_;
int32_t rank_;
QKVParallelLinear qkv_proj_{nullptr};
RowParallelLinear o_proj_{nullptr};
Qwen3NextRMSNorm q_norm_{nullptr};
Qwen3NextRMSNorm k_norm_{nullptr};
Attention attn_{nullptr};
MRotaryEmbedding rotary_emb_{nullptr};
torch::Tensor mrope_cu_seq_lens_;
};
TORCH_MODULE(Qwen3_5Attention);
} // namespace layer
} // namespace xllm

193
ex_engine/xllm_layers/mlu/qwen3_5_decoder_layer.cpp Normal file
View File

@@ -0,0 +1,193 @@
/* 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 "qwen3_5_decoder_layer.h"
#include <glog/logging.h>
#include "common/global_flags.h"
#include "layers/common/dp_utils.h"
namespace xllm {
namespace layer {
namespace {
bool use_moe_all2all(bool enable_deep_ep,
const ModelInputParams& input_params) {
return enable_deep_ep && all_dp_ranks_are_decode(input_params);
}
bool is_moe_layer(const ModelArgs& model_args, int32_t layer_id) {
const auto& mlp_only_layers = model_args.mlp_only_layers();
return std::count(mlp_only_layers.begin(), mlp_only_layers.end(), layer_id) ==
0 &&
model_args.n_routed_experts() > 0 &&
(layer_id + 1) % model_args.decoder_sparse_step() == 0;
}
} // namespace
Qwen3_5DecoderLayerImpl::Qwen3_5DecoderLayerImpl(const ModelContext& context,
int32_t layer_id)
: parallel_args_(context.get_parallel_args()) {
const auto& model_args = context.get_model_args();
const auto& quant_args = context.get_quant_args();
const auto& options = context.get_tensor_options();
const bool use_moe = is_moe_layer(model_args, layer_id);
enable_deep_ep_ = use_moe && FLAGS_expert_parallel_degree == 2;
if (enable_deep_ep_) {
CHECK_EQ(parallel_args_.dp_size(), parallel_args_.world_size())
<< "Qwen3.5 MoE only support deep ep all2all when dp_size == "
"world_size";
CHECK_EQ(parallel_args_.dp_size(), parallel_args_.ep_size())
<< "Qwen3.5 MoE only support deep ep all2all when dp_size == ep_size";
}
auto layer_types = model_args.layer_types();
if (layer_types.empty()) {
int32_t interval = model_args.full_attention_interval();
for (int32_t i = 0; i < model_args.n_layers(); i++) {
layer_types.push_back((i + 1) % interval == 0 ? "full_attention"
: "linear_attention");
}
}
if (layer_id >= 0 && layer_id < static_cast<int32_t>(layer_types.size())) {
layer_type_ = layer_types[layer_id];
} else {
layer_type_ = "full_attention";
}
if (layer_type_ == "linear_attention") {
// TODO: support linear attention
} else {
full_attention_ = register_module(
"self_attn",
Qwen3_5Attention(
model_args, quant_args, parallel_args_, options, layer_id));
}
input_norm_ = register_module(
"input_layernorm",
Qwen3NextRMSNorm(
model_args.hidden_size(), model_args.rms_norm_eps(), options));
post_norm_ = register_module(
"post_attention_layernorm",
Qwen3NextRMSNorm(
model_args.hidden_size(), model_args.rms_norm_eps(), options));
if (use_moe) {
moe_mlp_ = register_module("mlp",
Qwen3_5FusedMoE(model_args,
FusedMoEArgs{.is_gated = true},
quant_args,
parallel_args_,
options));
} else {
mlp_ = register_module("mlp",
DenseMLP(model_args.hidden_size(),
model_args.intermediate_size(),
true,
false,
model_args.hidden_act(),
/*enable_result_reduction=*/true,
quant_args,
parallel_args_.tp_group_,
options));
}
}
void Qwen3_5DecoderLayerImpl::load_state_dict(const StateDict& state_dict) {
if (layer_type_ == "linear_attention") {
// TODO: support linear attention
} else {
full_attention_->load_state_dict(
state_dict.get_dict_with_prefix("self_attn."));
}
input_norm_->load_state_dict(
state_dict.get_dict_with_prefix("input_layernorm."));
post_norm_->load_state_dict(
state_dict.get_dict_with_prefix("post_attention_layernorm."));
if (moe_mlp_) {
moe_mlp_->load_state_dict(state_dict.get_dict_with_prefix("mlp."));
} else {
mlp_->load_state_dict(state_dict.get_dict_with_prefix("mlp."));
}
}
torch::Tensor Qwen3_5DecoderLayerImpl::run_moe(
torch::Tensor x,
const ModelInputParams& input_params) {
const bool enable_moe_all2all =
use_moe_all2all(enable_deep_ep_, input_params);
if (need_dp_moe_gather(parallel_args_, enable_moe_all2all)) {
x = gather_dp_tokens(x, input_params, parallel_args_);
x = moe_mlp_->forward_experts(x, enable_moe_all2all);
return get_dp_local_slice(x, input_params, parallel_args_);
}
return moe_mlp_->forward_experts(x, enable_moe_all2all);
}
std::tuple<torch::Tensor, std::optional<torch::Tensor>>
Qwen3_5DecoderLayerImpl::apply_norm(Qwen3NextRMSNorm& norm,
torch::Tensor& input,
std::optional<torch::Tensor>& residual) {
if (!residual.has_value()) {
auto new_residual = input;
auto output = std::get<0>(norm->forward(input));
return {output, new_residual};
}
auto orig_dtype = input.dtype();
input = input + residual.value();
auto new_residual = input;
input = input.to(orig_dtype);
auto output = std::get<0>(norm->forward(input));
return {output, new_residual};
}
torch::Tensor Qwen3_5DecoderLayerImpl::forward(
torch::Tensor& x,
std::optional<torch::Tensor>& residual,
torch::Tensor& positions,
const AttentionMetadata& attn_metadata,
KVCache& kv_cache,
const ModelInputParams& input_params) {
// Pre-attention norm
std::tie(x, residual) = apply_norm(input_norm_, x, residual);
// Attention
if (full_attention_) {
x = full_attention_->forward(positions, x, attn_metadata, kv_cache);
} else {
// TODO: support linear attention
}
auto orig_dtype = x.dtype();
// Post-attention norm
std::tie(x, residual) = apply_norm(post_norm_, x, residual);
// MLP/MoE
if (moe_mlp_) {
x = run_moe(x, input_params);
} else {
x = mlp_->forward(x);
}
x = x.to(orig_dtype);
return x;
}
} // namespace layer
} // namespace xllm

73
ex_engine/xllm_layers/mlu/qwen3_5_decoder_layer.h Normal file
View File

@@ -0,0 +1,73 @@
/* 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 <optional>
#include <string>
#include "framework/kv_cache/kv_cache.h"
#include "framework/model/model_args.h"
#include "framework/model/model_input_params.h"
#include "framework/model_context.h"
#include "framework/parallel_state/parallel_args.h"
#include "framework/state_dict/state_dict.h"
#include "layers/common/dense_mlp.h"
#include "layers/common/qwen3_next_rms_norm.h"
#include "layers/mlu/qwen3_5_attention.h"
#include "layers/mlu/qwen3_5_fused_moe.h"
namespace xllm {
namespace layer {
class Qwen3_5DecoderLayerImpl final : public torch::nn::Module {
public:
Qwen3_5DecoderLayerImpl(const ModelContext& context, int32_t layer_id);
void load_state_dict(const StateDict& state_dict);
torch::Tensor forward(torch::Tensor& x,
std::optional<torch::Tensor>& residual,
torch::Tensor& positions,
const AttentionMetadata& attn_metadata,
KVCache& kv_cache,
const ModelInputParams& input_params);
private:
std::tuple<torch::Tensor, std::optional<torch::Tensor>> apply_norm(
Qwen3NextRMSNorm& norm,
torch::Tensor& input,
std::optional<torch::Tensor>& residual);
torch::Tensor run_moe(torch::Tensor x, const ModelInputParams& input_params);
std::string layer_type_;
Qwen3_5Attention full_attention_{nullptr};
// TODO: support linear attention
// Qwen3_5GatedDeltaNet linear_attention_{nullptr};
DenseMLP mlp_{nullptr};
Qwen3_5FusedMoE moe_mlp_{nullptr};
Qwen3NextRMSNorm input_norm_{nullptr};
Qwen3NextRMSNorm post_norm_{nullptr};
ParallelArgs parallel_args_;
bool enable_deep_ep_ = false;
};
TORCH_MODULE(Qwen3_5DecoderLayer);
} // namespace layer
} // namespace xllm

209
ex_engine/xllm_layers/mlu/qwen3_5_fused_moe.cpp Normal file
View File

@@ -0,0 +1,209 @@
/* 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 "qwen3_5_fused_moe.h"
#include <glog/logging.h>
#include "framework/parallel_state/parallel_state.h"
#include "framework/state_dict/utils.h"
namespace xllm {
namespace layer {
namespace {
torch::Tensor get_tensor_with_weight_suffix(const StateDict& state_dict,
const std::string& tensor_name) {
auto tensor = state_dict.get_tensor(tensor_name);
if (!tensor.defined()) {
tensor = state_dict.get_tensor(tensor_name + ".weight");
}
return tensor;
}
torch::Tensor slice_expert_weights(const torch::Tensor& weight,
int64_t start_expert_id,
int64_t num_experts_per_rank) {
return weight
.slice(0, start_expert_id, start_expert_id + num_experts_per_rank)
.contiguous();
}
bool load_fused_gate_up_fallback(const StateDict& state_dict,
int64_t rank,
int64_t world_size,
int64_t start_expert_id,
int64_t num_experts_per_rank,
torch::Tensor& w13) {
auto fused_gate_up =
get_tensor_with_weight_suffix(state_dict, "gate_up_proj");
if (!fused_gate_up.defined()) {
return false;
}
if (world_size > 1) {
CHECK_EQ(fused_gate_up.size(1) % 2, 0)
<< "gate_up_proj dim1 must be even, got " << fused_gate_up.size(1);
const int64_t full_intermediate = fused_gate_up.size(1) / 2;
CHECK_EQ(full_intermediate % world_size, 0)
<< "gate_up_proj intermediate dim is not divisible by world_size";
const int64_t inter_shard = full_intermediate / world_size;
auto gate_full = fused_gate_up.slice(1, 0, full_intermediate);
auto up_full =
fused_gate_up.slice(1, full_intermediate, full_intermediate * 2);
auto gate_shard =
gate_full.slice(1, rank * inter_shard, (rank + 1) * inter_shard);
auto up_shard =
up_full.slice(1, rank * inter_shard, (rank + 1) * inter_shard);
fused_gate_up = torch::cat({gate_shard, up_shard}, 1);
}
auto gate_up_slice = slice_expert_weights(
fused_gate_up, start_expert_id, num_experts_per_rank);
CHECK_EQ(w13.sizes(), gate_up_slice.sizes())
<< "weight size mismatch for " << state_dict.prefix()
<< "experts.gate_up_proj";
w13.copy_(gate_up_slice);
return true;
}
bool load_fused_down_fallback(const StateDict& state_dict,
int64_t rank,
int64_t world_size,
int64_t start_expert_id,
int64_t num_experts_per_rank,
torch::Tensor& w2) {
auto fused_down = get_tensor_with_weight_suffix(state_dict, "down_proj");
if (!fused_down.defined()) {
return false;
}
if (world_size > 1) {
CHECK_EQ(fused_down.size(2) % world_size, 0)
<< "down_proj dim2 is not divisible by world_size";
const int64_t down_shard = fused_down.size(2) / world_size;
fused_down =
fused_down.slice(2, rank * down_shard, (rank + 1) * down_shard);
}
auto down_slice =
slice_expert_weights(fused_down, start_expert_id, num_experts_per_rank);
CHECK_EQ(w2.sizes(), down_slice.sizes())
<< "weight size mismatch for " << state_dict.prefix()
<< "experts.down_proj";
w2.copy_(down_slice);
return true;
}
} // namespace
Qwen3_5FusedMoEImpl::Qwen3_5FusedMoEImpl(const ModelArgs& model_args,
const FusedMoEArgs& moe_args,
const QuantArgs& quant_args,
const ParallelArgs& parallel_args,
const torch::TensorOptions& options)
: FusedMoEImpl(model_args, moe_args, quant_args, parallel_args, options) {
if (n_shared_experts_ > 0) {
shared_expert_gate_ = register_module(
"shared_expert_gate",
torch::nn::Linear(
torch::nn::LinearOptions(hidden_size_, 1).bias(false)));
shared_expert_gate_->weight.set_data(
shared_expert_gate_->weight.to(options));
}
}
void Qwen3_5FusedMoEImpl::load_experts(const StateDict& state_dict) {
FusedMoEImpl::load_experts(state_dict);
if (!is_smoothquant_) {
if (!w13_is_loaded_) {
w13_is_loaded_ = load_fused_gate_up_fallback(state_dict,
tp_pg_->rank(),
tp_pg_->world_size(),
start_expert_id_,
num_experts_per_rank_,
w13_);
}
if (!w2_is_loaded_) {
w2_is_loaded_ = load_fused_down_fallback(state_dict,
tp_pg_->rank(),
tp_pg_->world_size(),
start_expert_id_,
num_experts_per_rank_,
w2_);
}
}
}
void Qwen3_5FusedMoEImpl::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_expert."));
auto weight = state_dict.get_tensor("shared_expert_gate.weight");
if (weight.defined()) {
weight = weight.reshape({weight.size(0), -1});
DCHECK_EQ(shared_expert_gate_->weight.sizes(), weight.sizes())
<< "proj weight size mismatch for " << name();
shared_expert_gate_->weight.data().copy_(weight);
}
}
gate_->load_state_dict(state_dict.get_dict_with_prefix("gate."));
load_experts(state_dict.get_dict_with_prefix("experts."));
}
void Qwen3_5FusedMoEImpl::final_comm_allreduce(
torch::Tensor& final_hidden_states,
const torch::Tensor& hidden_states,
torch::Tensor& shared_expert_output) {
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();
shared_expert_output = shared_experts_(hidden_states);
if (shared_expert_gate_) {
auto gate = torch::sigmoid(shared_expert_gate_->forward(hidden_states));
shared_expert_output = gate * shared_expert_output;
}
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;
}
}
} // namespace layer
} // namespace xllm

47
ex_engine/xllm_layers/mlu/qwen3_5_fused_moe.h Normal file
View File

@@ -0,0 +1,47 @@
/* 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 "layers/mlu/fused_moe.h"
namespace xllm {
namespace layer {
class Qwen3_5FusedMoEImpl final : public FusedMoEImpl {
public:
Qwen3_5FusedMoEImpl() = default;
Qwen3_5FusedMoEImpl(const ModelArgs& model_args,
const FusedMoEArgs& moe_args,
const QuantArgs& quant_args,
const ParallelArgs& parallel_args,
const torch::TensorOptions& options);
void load_state_dict(const StateDict& state_dict) override;
protected:
void final_comm_allreduce(torch::Tensor& final_hidden_states,
const torch::Tensor& hidden_states,
torch::Tensor& shared_expert_output) override;
private:
void load_experts(const StateDict& state_dict);
torch::nn::Linear shared_expert_gate_{nullptr};
};
TORCH_MODULE(Qwen3_5FusedMoE);
} // namespace layer
} // namespace xllm