[fix] baseline4 docker build move ex_engine into qwen3_6_scripts, remove COPY ex_engine from Dockerfile
This commit is contained in:
root
236
qwen3_6_scripts/ex_engine/xllm_layers/mlu/qwen3_5_attention.cpp
Normal file
236
qwen3_6_scripts/ex_engine/xllm_layers/mlu/qwen3_5_attention.cpp
Normal 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
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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
qwen3_6_scripts/ex_engine/xllm_layers/mlu/qwen3_5_fused_moe.cpp
Normal file
209
qwen3_6_scripts/ex_engine/xllm_layers/mlu/qwen3_5_fused_moe.cpp
Normal 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
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user