fix(critical): fold max_completion_tokens + max_num_seqs=2 + max_model_len=80000 + xllm_latest layer import
Sub 655 root causes (confirmed from log analysis): 1. protocol.py: max_completion_tokens never folded into max_tokens → 162/881 replay requests rejected 400 (extra_forbidden) 2. max_num_seqs=1 → t2_n_2 test fails (needs n=2) 3. max_model_len=131072 → OOM crash at 62% replay, opencompass all 0 Fixes: - protocol.py: model_validator fold_max_completion_tokens - yaml: max_num_seqs=2, max_model_len=80000, PYTORCH_CUDA_ALLOC_CONF - topk_softmax stays =0 (corex CUB BlockReduce incompatible on BI-V100) xllm_latest import to ex_engine/: - npu_torch layers: GDN(1164L), Qwen3.5 GDN, attention, fused_moe - cuda/moe kernels: topk_softmax_kernels.cuh, moe_combine, moe_compute_index - npu kernels: causal_conv1d, recurrent_gated_delta_rule - model headers: qwen3_5.h, qwen3_next.h
This commit is contained in:
project6-dev
236
ex_engine/xllm_layers/mlu/qwen3_5_attention.cpp
Normal file
236
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
|
||||
79
ex_engine/xllm_layers/mlu/qwen3_5_attention.h
Normal file
79
ex_engine/xllm_layers/mlu/qwen3_5_attention.h
Normal 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
193
ex_engine/xllm_layers/mlu/qwen3_5_decoder_layer.cpp
Normal 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
73
ex_engine/xllm_layers/mlu/qwen3_5_decoder_layer.h
Normal 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
209
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
|
||||
47
ex_engine/xllm_layers/mlu/qwen3_5_fused_moe.h
Normal file
47
ex_engine/xllm_layers/mlu/qwen3_5_fused_moe.h
Normal 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
|
||||
28
ex_engine/xllm_layers/npu_torch/CMakeLists.txt
Executable file
28
ex_engine/xllm_layers/npu_torch/CMakeLists.txt
Executable file
@@ -0,0 +1,28 @@
|
||||
include(cc_library)
|
||||
|
||||
cc_library(
|
||||
NAME
|
||||
npu_torch_layers
|
||||
HDRS
|
||||
fused_moe.h
|
||||
attention.h
|
||||
qwen3_gated_delta_net_base.h
|
||||
qwen3_next_attention.h
|
||||
qwen3_next_gated_delta_net.h
|
||||
qwen3_5_gated_delta_net.h
|
||||
qwen3_next_hybrid_decoder_layer_base.h
|
||||
qwen3_next_decoder_layer_impl.h
|
||||
qwen3_5_decoder_layer_impl.h
|
||||
SRCS
|
||||
fused_moe.cpp
|
||||
attention.cpp
|
||||
qwen3_gated_delta_net_base.cpp
|
||||
qwen3_next_attention.cpp
|
||||
qwen3_next_gated_delta_net.cpp
|
||||
qwen3_next_hybrid_decoder_layer_base.cpp
|
||||
qwen3_5_gated_delta_net.cpp
|
||||
qwen3_next_decoder_layer_impl.cpp
|
||||
qwen3_5_decoder_layer_impl.cpp
|
||||
DEPS
|
||||
:common_layers
|
||||
)
|
||||
152
ex_engine/xllm_layers/npu_torch/attention.cpp
Normal file
152
ex_engine/xllm_layers/npu_torch/attention.cpp
Normal file
@@ -0,0 +1,152 @@
|
||||
/* 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/npu/npu_ops_api.h"
|
||||
#include "kernels/ops_api.h"
|
||||
|
||||
DECLARE_bool(enable_chunked_prefill);
|
||||
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),
|
||||
num_kv_heads_(num_kv_heads),
|
||||
sliding_window_(sliding_window),
|
||||
scale_(scale) {
|
||||
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 = 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;
|
||||
|
||||
torch::Tensor k_cache = kv_cache.get_k_cache();
|
||||
torch::Tensor v = value.view({-1, num_kv_heads_, head_size_});
|
||||
std::optional<torch::Tensor> v_cache = kv_cache.get_v_cache();
|
||||
|
||||
// Reshape and cache key/value
|
||||
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 (only_prefill) {
|
||||
prefill_forward(query, key, value, output, k_cache, v_cache, attn_metadata);
|
||||
} else {
|
||||
decoder_forward(query, output, k_cache, v_cache, attn_metadata);
|
||||
}
|
||||
|
||||
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) {
|
||||
query = query.view({-1, num_heads_, head_size_});
|
||||
output = output.view({-1, num_heads_, head_size_});
|
||||
|
||||
if (attn_metadata.is_prefill) {
|
||||
key = key.view({-1, num_kv_heads_, head_size_});
|
||||
value = value.view({-1, num_kv_heads_, head_size_});
|
||||
|
||||
xllm::kernel::npu::batch_prefill(query,
|
||||
key,
|
||||
value,
|
||||
attn_metadata.attn_mask,
|
||||
attn_metadata.kv_seq_lens_host,
|
||||
scale_,
|
||||
output);
|
||||
} else if (attn_metadata.is_chunked_prefill) {
|
||||
xllm::kernel::npu::batch_prefill(query,
|
||||
k_cache,
|
||||
v_cache.value(),
|
||||
attn_metadata.attn_mask,
|
||||
attn_metadata.kv_seq_lens_host,
|
||||
scale_,
|
||||
output);
|
||||
}
|
||||
}
|
||||
|
||||
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) {
|
||||
query = query.view({-1, 1, num_heads_, head_size_});
|
||||
output = output.view({-1, 1, num_heads_, head_size_});
|
||||
|
||||
torch::Tensor kv_seq_lens;
|
||||
if (attn_metadata.kv_seq_lens_host.defined()) {
|
||||
kv_seq_lens = attn_metadata.kv_seq_lens_host;
|
||||
} else {
|
||||
// Fallback if host tensor isn't prepared.
|
||||
kv_seq_lens = attn_metadata.kv_seq_lens;
|
||||
}
|
||||
|
||||
if (attn_metadata.paged_attention_tiling_data.defined()) {
|
||||
// Use CustomPagedAttention for ACL graph mode to avoid .to(kCPU) operations
|
||||
|
||||
xllm::kernel::npu::batch_decode_acl_graph(
|
||||
query,
|
||||
k_cache,
|
||||
v_cache.value_or(torch::Tensor()),
|
||||
scale_,
|
||||
attn_metadata.block_table,
|
||||
kv_seq_lens,
|
||||
attn_metadata.paged_attention_tiling_data,
|
||||
output);
|
||||
} else {
|
||||
// Standard PagedAttention path
|
||||
xllm::kernel::npu::batch_decode(query,
|
||||
k_cache,
|
||||
v_cache.value_or(torch::Tensor()),
|
||||
scale_,
|
||||
attn_metadata.block_table,
|
||||
kv_seq_lens,
|
||||
output);
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace layer
|
||||
} // namespace xllm
|
||||
70
ex_engine/xllm_layers/npu_torch/attention.h
Normal file
70
ex_engine/xllm_layers/npu_torch/attention.h
Normal file
@@ -0,0 +1,70 @@
|
||||
/* 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);
|
||||
|
||||
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 sliding_window_;
|
||||
};
|
||||
TORCH_MODULE(Attention);
|
||||
|
||||
} // namespace layer
|
||||
} // namespace xllm
|
||||
513
ex_engine/xllm_layers/npu_torch/fused_moe.cpp
Normal file
513
ex_engine/xllm_layers/npu_torch/fused_moe.cpp
Normal file
@@ -0,0 +1,513 @@
|
||||
/* 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 "fused_moe.h"
|
||||
|
||||
#include <glog/logging.h>
|
||||
|
||||
#include <numeric>
|
||||
#include <vector>
|
||||
|
||||
#include "framework/parallel_state/parallel_state.h"
|
||||
#include "kernels/ops_api.h"
|
||||
|
||||
namespace xllm {
|
||||
namespace layer {
|
||||
|
||||
namespace {
|
||||
// Generic local tensor helpers.
|
||||
torch::Tensor create_group_gemm_output(
|
||||
const torch::Tensor& a,
|
||||
const torch::Tensor& b,
|
||||
const torch::Tensor& group_list,
|
||||
torch::ScalarType dtype = torch::ScalarType::BFloat16) {
|
||||
torch::TensorOptions target_options = a.options().dtype(dtype);
|
||||
if (b.dim() != 2) {
|
||||
return torch::empty({a.size(0), b.size(1)}, target_options);
|
||||
}
|
||||
return torch::empty({group_list.size(0), a.size(0), b.size(0)},
|
||||
target_options);
|
||||
}
|
||||
|
||||
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();
|
||||
}
|
||||
|
||||
// Qwen3.5-MoE fused checkpoint fallback helpers.
|
||||
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
|
||||
|
||||
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_(model_args.n_routed_experts()),
|
||||
topk_(model_args.num_experts_per_tok()),
|
||||
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()),
|
||||
is_smoothquant_(false),
|
||||
quant_args_(quant_args),
|
||||
parallel_args_(parallel_args),
|
||||
options_(options),
|
||||
tp_pg_(parallel_args.tp_group_) {
|
||||
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;
|
||||
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;
|
||||
}
|
||||
|
||||
// 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) {
|
||||
/*
|
||||
The shared_experts are usually implemented using the RowParallelLinear
|
||||
layer. Typically, this output serves as the enable_result_reduction results
|
||||
for the module. If only tensor parallelism is applied, immediate
|
||||
reduction of the shared_experts output isn't necessary; instead, we perform
|
||||
the reduction once at the end of the MoE operation.
|
||||
*/
|
||||
shared_experts_ =
|
||||
register_module("shared_experts",
|
||||
DenseMLP(hidden_size_,
|
||||
intermediate_size * n_shared_experts_,
|
||||
is_gated_,
|
||||
false,
|
||||
hidden_act_,
|
||||
/*enable_result_reduction=*/false,
|
||||
quant_args,
|
||||
tp_pg_,
|
||||
options));
|
||||
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));
|
||||
}
|
||||
|
||||
// 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);
|
||||
input_smooth_ = register_parameter(
|
||||
"input_smooth",
|
||||
torch::empty({num_experts_per_rank_, 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::select_experts(
|
||||
const torch::Tensor& hidden_states_2d,
|
||||
const torch::Tensor& router_logits_2d,
|
||||
SelectedExpertInfo& selected_expert_info) {
|
||||
// prepare the parameters for select_experts
|
||||
xllm::kernel::MoeFusedTopkParams moe_active_topk_params;
|
||||
moe_active_topk_params.input = router_logits_2d;
|
||||
moe_active_topk_params.finished = torch::Tensor();
|
||||
moe_active_topk_params.topk = topk_;
|
||||
moe_active_topk_params.scoring_func = "softmax";
|
||||
auto [topk_weights, topk_ids] =
|
||||
xllm::kernel::moe_active_topk(moe_active_topk_params);
|
||||
topk_ids = topk_ids.to(torch::kInt32);
|
||||
if (renormalize_) {
|
||||
topk_weights = topk_weights / (topk_weights.sum(-1, true) + 1e-6);
|
||||
}
|
||||
|
||||
xllm::kernel::MoeInitRoutingV2Params moe_init_routing_params;
|
||||
moe_init_routing_params.x = hidden_states_2d;
|
||||
moe_init_routing_params.expert_idx = topk_ids;
|
||||
moe_init_routing_params.scale = std::nullopt;
|
||||
moe_init_routing_params.offset = std::nullopt;
|
||||
moe_init_routing_params.active_num = hidden_states_2d.size(0) * topk_;
|
||||
moe_init_routing_params.expert_capacity = 0;
|
||||
moe_init_routing_params.expert_num = num_experts_per_rank_;
|
||||
moe_init_routing_params.drop_pad_mode = 0;
|
||||
moe_init_routing_params.expert_tokens_num_type = 1;
|
||||
moe_init_routing_params.expert_tokens_num_flag = true;
|
||||
moe_init_routing_params.row_idx_type = 0;
|
||||
std::vector<int64_t> expert_range = {
|
||||
start_expert_id_, start_expert_id_ + num_experts_per_rank_};
|
||||
moe_init_routing_params.active_expert_range = expert_range;
|
||||
moe_init_routing_params.quant_mode = -1;
|
||||
// TODO: NPU moe_init_routing_v2 is equivalent to moe_gen_idx +
|
||||
// moe_expand_input (and the token_count/cusum outputs) on other backends.
|
||||
auto [expand_hidden_states, expand_row_ids, group_list, dynamic_scale] =
|
||||
xllm::kernel::moe_init_routing_v2(moe_init_routing_params);
|
||||
(void)dynamic_scale;
|
||||
|
||||
// collect the selected tensor
|
||||
selected_expert_info.reduce_weight = topk_weights;
|
||||
selected_expert_info.combine_idx = expand_row_ids;
|
||||
selected_expert_info.token_count_slice = group_list;
|
||||
selected_expert_info.cusum_token_count = group_list;
|
||||
return expand_hidden_states;
|
||||
}
|
||||
|
||||
torch::Tensor FusedMoEImpl::forward_expert(
|
||||
const torch::Tensor& hidden_states,
|
||||
const torch::Tensor& router_logits,
|
||||
const std::optional<torch::Tensor>& shared_output) {
|
||||
// prepare the parameters for MoE computation
|
||||
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)});
|
||||
|
||||
// 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);
|
||||
|
||||
// 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);
|
||||
|
||||
{
|
||||
xllm::kernel::GroupGemmParams group_gemm_params;
|
||||
group_gemm_params.a = expand_hidden_states;
|
||||
if (w13_.size(1) != expand_hidden_states.size(1)) {
|
||||
w13_ = w13_.transpose(1, 2);
|
||||
}
|
||||
group_gemm_params.b = w13_;
|
||||
group_gemm_params.group_list = selected_expert_info.token_count_slice;
|
||||
group_gemm_params.split_item = 2;
|
||||
group_gemm_params.group_type = 0;
|
||||
group_gemm_params.group_list_type = 1;
|
||||
gemm1_out = xllm::kernel::group_gemm(group_gemm_params);
|
||||
}
|
||||
|
||||
// Step 5: activation
|
||||
torch::Tensor act_out;
|
||||
|
||||
xllm::kernel::ActivationParams activation_params;
|
||||
activation_params.input = gemm1_out;
|
||||
activation_params.output = act_out;
|
||||
activation_params.act_mode = hidden_act_;
|
||||
activation_params.is_gated = is_gated_;
|
||||
xllm::kernel::active(activation_params);
|
||||
act_out = activation_params.output;
|
||||
// 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);
|
||||
|
||||
{
|
||||
xllm::kernel::GroupGemmParams group_gemm_params;
|
||||
group_gemm_params.a = act_out;
|
||||
if (w2_.size(1) != act_out.size(1)) {
|
||||
w2_ = w2_.transpose(1, 2);
|
||||
}
|
||||
group_gemm_params.b = w2_;
|
||||
group_gemm_params.group_list = selected_expert_info.token_count_slice;
|
||||
group_gemm_params.split_item = 2;
|
||||
group_gemm_params.group_type = 0;
|
||||
group_gemm_params.group_list_type = 1;
|
||||
gemm2_out = xllm::kernel::group_gemm(group_gemm_params);
|
||||
}
|
||||
|
||||
// Step 7: combine the intermediate results and get the final hidden states
|
||||
torch::Tensor final_hidden_states;
|
||||
xllm::kernel::MoeCombineResultParams moe_combine_params;
|
||||
moe_combine_params.input = gemm2_out;
|
||||
moe_combine_params.reduce_weight = selected_expert_info.reduce_weight;
|
||||
moe_combine_params.gather_ids = selected_expert_info.combine_idx;
|
||||
final_hidden_states = xllm::kernel::moe_combine_result(moe_combine_params);
|
||||
if (shared_output.has_value()) {
|
||||
final_hidden_states = final_hidden_states + shared_output.value();
|
||||
}
|
||||
// reshape the final hidden states to the original shape
|
||||
final_hidden_states = final_hidden_states.reshape(hidden_states_shape);
|
||||
|
||||
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_);
|
||||
}
|
||||
return final_hidden_states;
|
||||
}
|
||||
|
||||
torch::Tensor FusedMoEImpl::forward(const torch::Tensor& hidden_states,
|
||||
const ModelInputParams& input_params) {
|
||||
auto input = hidden_states;
|
||||
bool need_slice = false;
|
||||
if (parallel_args_.dp_size() > 1 && parallel_args_.ep_size() > 1) {
|
||||
input = parallel_state::gather(input,
|
||||
parallel_args_.dp_local_process_group_,
|
||||
input_params.dp_global_token_nums);
|
||||
need_slice = true;
|
||||
}
|
||||
|
||||
std::optional<torch::Tensor> shared_output = std::nullopt;
|
||||
if (n_shared_experts_ > 0) {
|
||||
shared_output = shared_experts_(input);
|
||||
if (shared_expert_gate_) {
|
||||
auto gate = torch::sigmoid(shared_expert_gate_->forward(input));
|
||||
if (shared_output.has_value()) {
|
||||
torch::Tensor res = gate * shared_output.value();
|
||||
shared_output = res;
|
||||
}
|
||||
}
|
||||
}
|
||||
auto router_logits = gate_(input);
|
||||
auto output = forward_expert(input, router_logits, shared_output);
|
||||
|
||||
if (need_slice) {
|
||||
const auto& dp_tokens = input_params.dp_global_token_nums;
|
||||
const int64_t dp_rank = parallel_args_.dp_local_process_group_->rank();
|
||||
auto start =
|
||||
std::accumulate(dp_tokens.begin(), dp_tokens.begin() + dp_rank, 0);
|
||||
auto end = start + dp_tokens[dp_rank];
|
||||
output = output.slice(0, start, end);
|
||||
}
|
||||
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_;
|
||||
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);
|
||||
LOAD_MOE_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);
|
||||
|
||||
// Some Qwen3.5-MoE checkpoints store expert weights in fused tensors
|
||||
// (gate_up_proj / down_proj). Fall back to this format when split
|
||||
// gate_proj/up_proj tensors are absent.
|
||||
if (!w13_is_loaded_) {
|
||||
w13_is_loaded_ = load_fused_gate_up_fallback(state_dict,
|
||||
rank,
|
||||
world_size,
|
||||
start_expert_id,
|
||||
num_experts_per_rank,
|
||||
w13_);
|
||||
}
|
||||
|
||||
if (!w2_is_loaded_) {
|
||||
w2_is_loaded_ = load_fused_down_fallback(state_dict,
|
||||
rank,
|
||||
world_size,
|
||||
start_expert_id,
|
||||
num_experts_per_rank,
|
||||
w2_);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
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_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_e_score_correction_bias(state_dict.get_dict_with_prefix("gate."));
|
||||
load_experts(state_dict.get_dict_with_prefix("experts."));
|
||||
}
|
||||
|
||||
} // namespace layer
|
||||
} // namespace xllm
|
||||
113
ex_engine/xllm_layers/npu_torch/fused_moe.h
Normal file
113
ex_engine/xllm_layers/npu_torch/fused_moe.h
Normal file
@@ -0,0 +1,113 @@
|
||||
/* 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 <optional>
|
||||
|
||||
#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/dense_mlp.h"
|
||||
#include "layers/common/fused_moe_base.h"
|
||||
#include "layers/common/linear.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_expert(
|
||||
const torch::Tensor& hidden_states,
|
||||
const torch::Tensor& router_logits,
|
||||
const std::optional<torch::Tensor>& shared_output);
|
||||
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;
|
||||
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);
|
||||
|
||||
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_;
|
||||
bool has_score_bias_;
|
||||
bool has_bias_;
|
||||
bool skip_bias_add_;
|
||||
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_;
|
||||
|
||||
ReplicatedLinear gate_{nullptr};
|
||||
DenseMLP shared_experts_{nullptr};
|
||||
torch::nn::Linear shared_expert_gate_{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);
|
||||
};
|
||||
TORCH_MODULE(FusedMoE);
|
||||
|
||||
} // namespace layer
|
||||
} // namespace xllm
|
||||
@@ -0,0 +1,32 @@
|
||||
/* Copyright 2025-2026 The xLLM Authors.
|
||||
|
||||
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_impl.h"
|
||||
|
||||
namespace xllm {
|
||||
namespace layer {
|
||||
|
||||
Qwen3_5DecoderLayerImpl::Qwen3_5DecoderLayerImpl(const ModelContext& context,
|
||||
int32_t layer_id)
|
||||
: Qwen3NextDecoderLayerImpl(context,
|
||||
layer_id,
|
||||
std::make_shared<Qwen3_5GatedDeltaNetImpl>(
|
||||
context.get_model_args(),
|
||||
context.get_quant_args(),
|
||||
context.get_parallel_args(),
|
||||
context.get_tensor_options())) {}
|
||||
|
||||
} // namespace layer
|
||||
} // namespace xllm
|
||||
32
ex_engine/xllm_layers/npu_torch/qwen3_5_decoder_layer_impl.h
Normal file
32
ex_engine/xllm_layers/npu_torch/qwen3_5_decoder_layer_impl.h
Normal file
@@ -0,0 +1,32 @@
|
||||
/* Copyright 2025-2026 The xLLM Authors.
|
||||
|
||||
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/npu_torch/qwen3_5_gated_delta_net.h"
|
||||
#include "layers/npu_torch/qwen3_next_decoder_layer_impl.h"
|
||||
|
||||
namespace xllm {
|
||||
namespace layer {
|
||||
|
||||
class Qwen3_5DecoderLayerImpl : public Qwen3NextDecoderLayerImpl {
|
||||
public:
|
||||
explicit Qwen3_5DecoderLayerImpl(const ModelContext& context,
|
||||
int32_t layer_id);
|
||||
};
|
||||
TORCH_MODULE(Qwen3_5DecoderLayer);
|
||||
|
||||
} // namespace layer
|
||||
} // namespace xllm
|
||||
219
ex_engine/xllm_layers/npu_torch/qwen3_5_gated_delta_net.cpp
Normal file
219
ex_engine/xllm_layers/npu_torch/qwen3_5_gated_delta_net.cpp
Normal file
@@ -0,0 +1,219 @@
|
||||
/* Copyright 2025-2026 The xLLM Authors.
|
||||
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_gated_delta_net.h"
|
||||
|
||||
#include <glog/logging.h>
|
||||
|
||||
namespace xllm {
|
||||
namespace layer {
|
||||
|
||||
Qwen3_5GatedDeltaNetImpl::Qwen3_5GatedDeltaNetImpl(
|
||||
const ModelArgs& args,
|
||||
const QuantArgs& quant_args,
|
||||
const ParallelArgs& parallel_args,
|
||||
const torch::TensorOptions& options)
|
||||
: Qwen3NextGatedDeltaNetImpl(args,
|
||||
quant_args,
|
||||
parallel_args,
|
||||
options,
|
||||
/*init_projections=*/false) {
|
||||
in_proj_qkv_ = register_module("in_proj_qkv",
|
||||
ColumnParallelLinear(args.hidden_size(),
|
||||
k_size_ * 2 + v_size_,
|
||||
/*bias=*/false,
|
||||
/*gather_output=*/false,
|
||||
quant_args,
|
||||
parallel_args.tp_group_,
|
||||
options));
|
||||
in_proj_z_ = register_module("in_proj_z",
|
||||
ColumnParallelLinear(args.hidden_size(),
|
||||
v_size_,
|
||||
/*bias=*/false,
|
||||
/*gather_output=*/false,
|
||||
quant_args,
|
||||
parallel_args.tp_group_,
|
||||
options));
|
||||
in_proj_b_ = register_module("in_proj_b",
|
||||
ColumnParallelLinear(args.hidden_size(),
|
||||
num_v_heads_,
|
||||
/*bias=*/false,
|
||||
/*gather_output=*/false,
|
||||
quant_args,
|
||||
parallel_args.tp_group_,
|
||||
options));
|
||||
in_proj_a_ = register_module("in_proj_a",
|
||||
ColumnParallelLinear(args.hidden_size(),
|
||||
num_v_heads_,
|
||||
/*bias=*/false,
|
||||
/*gather_output=*/false,
|
||||
quant_args,
|
||||
parallel_args.tp_group_,
|
||||
options));
|
||||
}
|
||||
|
||||
torch::Tensor Qwen3_5GatedDeltaNetImpl::merge_qkvz_from_split_activations(
|
||||
const torch::Tensor& qkv,
|
||||
const torch::Tensor& z) const {
|
||||
CHECK_EQ(qkv.dim(), 3) << "Expected qkv activation to be 3D, got "
|
||||
<< qkv.sizes();
|
||||
CHECK_EQ(z.dim(), 3) << "Expected z activation to be 3D, got " << z.sizes();
|
||||
CHECK_EQ(qkv.size(0), z.size(0)) << "qkv/z batch size mismatch.";
|
||||
CHECK_EQ(qkv.size(1), z.size(1)) << "qkv/z sequence size mismatch.";
|
||||
CHECK_EQ(qkv.size(2), (2 * k_size_ + v_size_) / tp_size_)
|
||||
<< "Unexpected qkv hidden size for Qwen3.5.";
|
||||
CHECK_EQ(z.size(2), v_size_ / tp_size_)
|
||||
<< "Unexpected z hidden size for Qwen3.5.";
|
||||
CHECK_GT(num_k_heads_, 0) << "linear_num_key_heads must be positive.";
|
||||
CHECK_EQ(num_v_heads_ % num_k_heads_, 0)
|
||||
<< "linear_num_value_heads must be divisible by linear_num_key_heads.";
|
||||
|
||||
const int64_t bs = qkv.size(0);
|
||||
const int64_t seqlen = qkv.size(1);
|
||||
const int64_t local_k_heads = num_k_heads_ / tp_size_;
|
||||
const int64_t local_v_heads = num_v_heads_ / tp_size_;
|
||||
const int64_t num_v_heads_per_k = num_v_heads_ / num_k_heads_;
|
||||
|
||||
auto qkv_split = torch::split(
|
||||
qkv, {k_size_ / tp_size_, k_size_ / tp_size_, v_size_ / tp_size_}, 2);
|
||||
auto q = qkv_split[0].view({bs, seqlen, local_k_heads, head_k_dim_});
|
||||
auto k = qkv_split[1].view({bs, seqlen, local_k_heads, head_k_dim_});
|
||||
auto v = qkv_split[2].view({bs, seqlen, local_v_heads, head_v_dim_});
|
||||
auto z_view = z.view({bs, seqlen, local_v_heads, head_v_dim_});
|
||||
|
||||
v = v.view({bs, seqlen, local_k_heads, num_v_heads_per_k * head_v_dim_});
|
||||
z_view =
|
||||
z_view.view({bs, seqlen, local_k_heads, num_v_heads_per_k * head_v_dim_});
|
||||
|
||||
return torch::cat({q, k, v, z_view}, -1).view({bs, seqlen, -1}).contiguous();
|
||||
}
|
||||
|
||||
torch::Tensor Qwen3_5GatedDeltaNetImpl::merge_ba_from_split_activations(
|
||||
const torch::Tensor& b,
|
||||
const torch::Tensor& a) const {
|
||||
CHECK_EQ(b.dim(), 3) << "Expected b activation to be 3D, got " << b.sizes();
|
||||
CHECK_EQ(a.dim(), 3) << "Expected a activation to be 3D, got " << a.sizes();
|
||||
CHECK_EQ(b.size(0), a.size(0)) << "b/a batch size mismatch.";
|
||||
CHECK_EQ(b.size(1), a.size(1)) << "b/a sequence size mismatch.";
|
||||
CHECK_EQ(b.size(2), num_v_heads_ / tp_size_)
|
||||
<< "Unexpected b hidden size for Qwen3.5.";
|
||||
CHECK_EQ(a.size(2), num_v_heads_ / tp_size_)
|
||||
<< "Unexpected a hidden size for Qwen3.5.";
|
||||
CHECK_GT(num_k_heads_, 0) << "linear_num_key_heads must be positive.";
|
||||
CHECK_EQ(num_v_heads_ % num_k_heads_, 0)
|
||||
<< "linear_num_value_heads must be divisible by linear_num_key_heads.";
|
||||
|
||||
const int64_t bs = b.size(0);
|
||||
const int64_t seqlen = b.size(1);
|
||||
const int64_t local_k_heads = num_k_heads_ / tp_size_;
|
||||
const int64_t num_v_heads_per_k = num_v_heads_ / num_k_heads_;
|
||||
|
||||
auto b_view = b.view({bs, seqlen, local_k_heads, num_v_heads_per_k});
|
||||
auto a_view = a.view({bs, seqlen, local_k_heads, num_v_heads_per_k});
|
||||
return torch::cat({b_view, a_view}, -1).view({bs, seqlen, -1}).contiguous();
|
||||
}
|
||||
|
||||
std::pair<torch::Tensor, torch::Tensor>
|
||||
Qwen3_5GatedDeltaNetImpl::project_decode_inputs(
|
||||
const torch::Tensor& hidden_states) {
|
||||
const auto reshape_projection = [](const torch::Tensor& projection) {
|
||||
return projection.view({projection.size(0), -1, projection.size(-1)});
|
||||
};
|
||||
auto qkv = reshape_projection(in_proj_qkv_->forward(hidden_states));
|
||||
auto z_proj = reshape_projection(in_proj_z_->forward(hidden_states));
|
||||
auto b_proj = reshape_projection(in_proj_b_->forward(hidden_states));
|
||||
auto a_proj = reshape_projection(in_proj_a_->forward(hidden_states));
|
||||
return {merge_qkvz_from_split_activations(qkv, z_proj),
|
||||
merge_ba_from_split_activations(b_proj, a_proj)};
|
||||
}
|
||||
|
||||
std::pair<torch::Tensor, torch::Tensor>
|
||||
Qwen3_5GatedDeltaNetImpl::project_flat_inputs(
|
||||
const torch::Tensor& hidden_states) {
|
||||
auto qkv = in_proj_qkv_->forward(hidden_states).unsqueeze(0);
|
||||
auto z_proj = in_proj_z_->forward(hidden_states).unsqueeze(0);
|
||||
auto b_proj = in_proj_b_->forward(hidden_states).unsqueeze(0);
|
||||
auto a_proj = in_proj_a_->forward(hidden_states).unsqueeze(0);
|
||||
auto qkvz = merge_qkvz_from_split_activations(qkv, z_proj);
|
||||
auto ba = merge_ba_from_split_activations(b_proj, a_proj);
|
||||
return {qkvz.view({hidden_states.size(0), qkvz.size(-1)}).contiguous(),
|
||||
ba.view({hidden_states.size(0), ba.size(-1)}).contiguous()};
|
||||
}
|
||||
|
||||
std::optional<
|
||||
std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>>
|
||||
Qwen3_5GatedDeltaNetImpl::project_split_inputs(
|
||||
const torch::Tensor& hidden_states,
|
||||
const AttentionMetadata& attn_metadata) {
|
||||
auto qkv = reshape_projected_tokens_with_pad(
|
||||
attn_metadata, in_proj_qkv_->forward(hidden_states));
|
||||
auto z_proj = reshape_projected_tokens_with_pad(
|
||||
attn_metadata, in_proj_z_->forward(hidden_states));
|
||||
auto b_proj = reshape_projected_tokens_with_pad(
|
||||
attn_metadata, in_proj_b_->forward(hidden_states));
|
||||
auto a_proj = reshape_projected_tokens_with_pad(
|
||||
attn_metadata, in_proj_a_->forward(hidden_states));
|
||||
|
||||
const int64_t batch_size = qkv.size(0);
|
||||
const int64_t seq_len = qkv.size(1);
|
||||
auto z =
|
||||
z_proj.view({batch_size, seq_len, num_v_heads_ / tp_size_, head_v_dim_});
|
||||
auto b = b_proj.view({batch_size, seq_len, num_v_heads_ / tp_size_});
|
||||
auto a = a_proj.view({batch_size, seq_len, num_v_heads_ / tp_size_});
|
||||
return std::make_tuple(qkv, z, b, a);
|
||||
}
|
||||
|
||||
void Qwen3_5GatedDeltaNetImpl::load_projection_state_dict(
|
||||
const StateDict& state_dict) {
|
||||
auto in_proj_qkv_state_dict = state_dict.get_dict_with_prefix("in_proj_qkv.");
|
||||
if (in_proj_qkv_state_dict.size() > 0 && !in_proj_qkv_->is_weight_loaded()) {
|
||||
in_proj_qkv_->load_state_dict(
|
||||
in_proj_qkv_state_dict,
|
||||
/*shard_tensor_count=*/3,
|
||||
/*shard_sizes=*/
|
||||
{k_size_ / tp_size_, k_size_ / tp_size_, v_size_ / tp_size_});
|
||||
}
|
||||
|
||||
auto in_proj_z_state_dict = state_dict.get_dict_with_prefix("in_proj_z.");
|
||||
if (in_proj_z_state_dict.size() > 0 && !in_proj_z_->is_weight_loaded()) {
|
||||
in_proj_z_->load_state_dict(in_proj_z_state_dict);
|
||||
}
|
||||
|
||||
auto in_proj_b_state_dict = state_dict.get_dict_with_prefix("in_proj_b.");
|
||||
if (in_proj_b_state_dict.size() > 0 && !in_proj_b_->is_weight_loaded()) {
|
||||
in_proj_b_->load_state_dict(in_proj_b_state_dict);
|
||||
}
|
||||
|
||||
auto in_proj_a_state_dict = state_dict.get_dict_with_prefix("in_proj_a.");
|
||||
if (in_proj_a_state_dict.size() > 0 && !in_proj_a_->is_weight_loaded()) {
|
||||
in_proj_a_->load_state_dict(in_proj_a_state_dict);
|
||||
}
|
||||
}
|
||||
|
||||
void Qwen3_5GatedDeltaNetImpl::verify_projection_weights(
|
||||
const std::string& prefix) const {
|
||||
CHECK(in_proj_qkv_ && in_proj_qkv_->is_weight_loaded())
|
||||
<< "Missing required weight after all shards loaded: " << prefix
|
||||
<< "in_proj_qkv.weight";
|
||||
CHECK(in_proj_z_ && in_proj_z_->is_weight_loaded())
|
||||
<< "Missing required weight after all shards loaded: " << prefix
|
||||
<< "in_proj_z.weight";
|
||||
CHECK(in_proj_b_ && in_proj_b_->is_weight_loaded())
|
||||
<< "Missing required weight after all shards loaded: " << prefix
|
||||
<< "in_proj_b.weight";
|
||||
CHECK(in_proj_a_ && in_proj_a_->is_weight_loaded())
|
||||
<< "Missing required weight after all shards loaded: " << prefix
|
||||
<< "in_proj_a.weight";
|
||||
}
|
||||
|
||||
} // namespace layer
|
||||
} // namespace xllm
|
||||
66
ex_engine/xllm_layers/npu_torch/qwen3_5_gated_delta_net.h
Normal file
66
ex_engine/xllm_layers/npu_torch/qwen3_5_gated_delta_net.h
Normal file
@@ -0,0 +1,66 @@
|
||||
/* Copyright 2025-2026 The xLLM Authors.
|
||||
|
||||
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 <tuple>
|
||||
#include <utility>
|
||||
|
||||
#include "qwen3_next_gated_delta_net.h"
|
||||
|
||||
namespace xllm {
|
||||
namespace layer {
|
||||
|
||||
class Qwen3_5GatedDeltaNetImpl : public Qwen3NextGatedDeltaNetImpl {
|
||||
public:
|
||||
Qwen3_5GatedDeltaNetImpl() = default;
|
||||
Qwen3_5GatedDeltaNetImpl(const ModelArgs& args,
|
||||
const QuantArgs& quant_args,
|
||||
const ParallelArgs& parallel_args,
|
||||
const torch::TensorOptions& options);
|
||||
|
||||
protected:
|
||||
std::pair<torch::Tensor, torch::Tensor> project_decode_inputs(
|
||||
const torch::Tensor& hidden_states) override;
|
||||
std::pair<torch::Tensor, torch::Tensor> project_flat_inputs(
|
||||
const torch::Tensor& hidden_states) override;
|
||||
std::optional<
|
||||
std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>>
|
||||
project_split_inputs(const torch::Tensor& hidden_states,
|
||||
const AttentionMetadata& attn_metadata) override;
|
||||
bool use_fla_ssm_state_layout() const override { return true; }
|
||||
|
||||
void load_projection_state_dict(const StateDict& state_dict) override;
|
||||
void verify_projection_weights(const std::string& prefix) const override;
|
||||
|
||||
private:
|
||||
torch::Tensor merge_qkvz_from_split_activations(const torch::Tensor& qkv,
|
||||
const torch::Tensor& z) const;
|
||||
torch::Tensor merge_ba_from_split_activations(const torch::Tensor& b,
|
||||
const torch::Tensor& a) const;
|
||||
|
||||
ColumnParallelLinear in_proj_qkv_{nullptr};
|
||||
ColumnParallelLinear in_proj_z_{nullptr};
|
||||
ColumnParallelLinear in_proj_b_{nullptr};
|
||||
ColumnParallelLinear in_proj_a_{nullptr};
|
||||
};
|
||||
TORCH_MODULE(Qwen3_5GatedDeltaNet);
|
||||
|
||||
} // namespace layer
|
||||
} // namespace xllm
|
||||
1164
ex_engine/xllm_layers/npu_torch/qwen3_gated_delta_net_base.cpp
Normal file
1164
ex_engine/xllm_layers/npu_torch/qwen3_gated_delta_net_base.cpp
Normal file
File diff suppressed because it is too large
Load Diff
112
ex_engine/xllm_layers/npu_torch/qwen3_gated_delta_net_base.h
Normal file
112
ex_engine/xllm_layers/npu_torch/qwen3_gated_delta_net_base.h
Normal file
@@ -0,0 +1,112 @@
|
||||
/* Copyright 2025-2026 The xLLM Authors.
|
||||
|
||||
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 <tuple>
|
||||
#include <utility>
|
||||
|
||||
#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 "framework/state_dict/utils.h"
|
||||
#include "layers/common/linear.h"
|
||||
#include "layers/common/rms_norm_gated.h"
|
||||
|
||||
namespace xllm {
|
||||
namespace layer {
|
||||
|
||||
class Qwen3GatedDeltaNetBaseImpl : public torch::nn::Module {
|
||||
public:
|
||||
Qwen3GatedDeltaNetBaseImpl() = default;
|
||||
Qwen3GatedDeltaNetBaseImpl(const ModelArgs& args,
|
||||
const QuantArgs& quant_args,
|
||||
const ParallelArgs& parallel_args,
|
||||
const torch::TensorOptions& options);
|
||||
|
||||
virtual void load_state_dict(const StateDict& state_dict) = 0;
|
||||
virtual void verify_loaded_weights(const std::string& prefix) const = 0;
|
||||
|
||||
torch::Tensor forward(const torch::Tensor& hidden_states,
|
||||
const AttentionMetadata& attn_metadata,
|
||||
KVCache& kv_cache,
|
||||
const ModelInputParams& input_params);
|
||||
|
||||
protected:
|
||||
virtual std::pair<torch::Tensor, torch::Tensor> project_decode_inputs(
|
||||
const torch::Tensor& hidden_states) = 0;
|
||||
virtual std::pair<torch::Tensor, torch::Tensor> project_flat_inputs(
|
||||
const torch::Tensor& hidden_states) = 0;
|
||||
// Qwen3.5 overrides this to project and reshape its separate qkv/z/b/a
|
||||
// weights in every forward mode. Qwen3Next keeps qkvz/ba packed and returns
|
||||
// nullopt to select the fused-split fallback.
|
||||
virtual std::optional<
|
||||
std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>>
|
||||
project_split_inputs(const torch::Tensor& hidden_states,
|
||||
const AttentionMetadata& attn_metadata) {
|
||||
return std::nullopt;
|
||||
}
|
||||
virtual bool use_fla_ssm_state_layout() const { return false; }
|
||||
|
||||
void load_common_state_dict(const StateDict& state_dict);
|
||||
void verify_common_loaded_weights(const std::string& prefix) const;
|
||||
|
||||
torch::Tensor get_linear_state_indices(const ModelInputParams& input_params,
|
||||
const torch::Device& device) const;
|
||||
|
||||
std::pair<torch::Tensor, torch::Tensor> project_padded_inputs(
|
||||
const torch::Tensor& hidden_states,
|
||||
const AttentionMetadata& attn_metadata);
|
||||
|
||||
torch::Tensor reshape_qkvz_unpad(const AttentionMetadata& attn_metadata,
|
||||
const torch::Tensor& padded_qkvz) const;
|
||||
|
||||
// Projection outputs are packed as [total_tokens, dim], while GDN kernels
|
||||
// consume dense [batch, max_query_len, dim] tensors. Split the packed tokens
|
||||
// by query length and pad each sequence before entering the kernels.
|
||||
torch::Tensor reshape_projected_tokens_with_pad(
|
||||
const AttentionMetadata& attn_metadata,
|
||||
const torch::Tensor& projected_tokens) const;
|
||||
|
||||
std::tuple<torch::Tensor, torch::Tensor, torch::Tensor> process_mixed_qkv(
|
||||
torch::Tensor& mixed_qkv) const;
|
||||
|
||||
int64_t num_k_heads_ = 0;
|
||||
int64_t num_v_heads_ = 0;
|
||||
int64_t head_k_dim_ = 0;
|
||||
int64_t head_v_dim_ = 0;
|
||||
int64_t k_size_ = 0;
|
||||
int64_t v_size_ = 0;
|
||||
int64_t tp_size_ = 1;
|
||||
int64_t rank_ = 0;
|
||||
int32_t conv_kernel_size_ = 0;
|
||||
|
||||
ColumnParallelLinear conv1d_{nullptr};
|
||||
RowParallelLinear o_proj_{nullptr};
|
||||
RmsNormGated norm_{nullptr};
|
||||
|
||||
DEFINE_WEIGHT(dt_bias);
|
||||
DEFINE_WEIGHT(A_log);
|
||||
};
|
||||
|
||||
} // namespace layer
|
||||
} // namespace xllm
|
||||
291
ex_engine/xllm_layers/npu_torch/qwen3_next_attention.cpp
Normal file
291
ex_engine/xllm_layers/npu_torch/qwen3_next_attention.cpp
Normal file
@@ -0,0 +1,291 @@
|
||||
/* Copyright 2025-2026 The xLLM Authors.
|
||||
|
||||
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_next_attention.h"
|
||||
|
||||
#include <glog/logging.h>
|
||||
|
||||
#include <tuple>
|
||||
#include <vector>
|
||||
|
||||
#include "common/flash_comm1_context.h"
|
||||
|
||||
namespace xllm {
|
||||
namespace layer {
|
||||
|
||||
Qwen3NextAttentionImpl::Qwen3NextAttentionImpl(
|
||||
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();
|
||||
// 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,
|
||||
quant_args));
|
||||
|
||||
// 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. Rotary embedding
|
||||
const int rotary_dim =
|
||||
static_cast<int>(head_dim_ * args.partial_rotary_factor());
|
||||
rotary_emb_ =
|
||||
register_module("rotary_emb",
|
||||
PartialRotaryEmbedding(rotary_dim,
|
||||
args.max_position_embeddings(),
|
||||
args.rope_theta(),
|
||||
head_dim_,
|
||||
true,
|
||||
false,
|
||||
options));
|
||||
|
||||
// 6. Attention
|
||||
attn_ = register_module("attn",
|
||||
Attention(num_heads_,
|
||||
head_dim_,
|
||||
scaling_,
|
||||
num_kv_heads_,
|
||||
args.sliding_window()));
|
||||
|
||||
// 7. Fused split_qkv_rmsnorm_mrope kernel setup
|
||||
rotary_dim_ = static_cast<int64_t>(head_dim_ * args.partial_rotary_factor());
|
||||
rms_norm_eps_ = args.rms_norm_eps();
|
||||
mrope_section_ = args.rope_scaling_mrope_section();
|
||||
is_interleaved_ = args.rope_scaling_mrope_interleaved();
|
||||
use_fused_qkv_ = false;
|
||||
if (attn_output_gate_ && !mrope_section_.empty() &&
|
||||
mrope_section_.size() == 3 && rotary_dim_ > 0 &&
|
||||
xllm::kernel::has_split_qkv_rmsnorm_mrope_specialization(
|
||||
num_heads_, num_kv_heads_, head_dim_)) {
|
||||
mrope_gather_pattern_ =
|
||||
xllm::kernel::build_split_qkv_rmsnorm_mrope_gather_pattern(
|
||||
rotary_dim_, mrope_section_, is_interleaved_, options.device());
|
||||
use_fused_qkv_ = true;
|
||||
LOG(INFO) << "Qwen3NextAttention layer " << layer_id_
|
||||
<< ": using fused split_qkv_rmsnorm_mrope kernel";
|
||||
}
|
||||
}
|
||||
|
||||
torch::Tensor Qwen3NextAttentionImpl::build_mrope_cos_sin(
|
||||
const torch::Tensor& positions) const {
|
||||
auto cos_sin_cache = rotary_emb_->get_cos_sin_cache();
|
||||
if (positions.dim() == 1) {
|
||||
return cos_sin_cache.index_select(0, positions).repeat({1, 3});
|
||||
}
|
||||
// positions is [3, T] for mRoPE (graph mode or VL)
|
||||
// transpose from [3, T] to [T, 3]
|
||||
auto positions_t = positions.permute({1, 0}).contiguous();
|
||||
auto gathered = cos_sin_cache.index_select(0, positions_t.view({-1}));
|
||||
// [T, 3, rope_dim]
|
||||
return gathered.view({positions.size(1), -1});
|
||||
}
|
||||
|
||||
torch::Tensor Qwen3NextAttentionImpl::forward(
|
||||
const torch::Tensor& positions,
|
||||
const torch::Tensor& hidden_states,
|
||||
const AttentionMetadata& attn_metadata,
|
||||
KVCache& kv_cache,
|
||||
const torch::Tensor& mrope_cos_sin) {
|
||||
const FlashComm1Context* fc1_ctx = get_current_flash_comm1_context();
|
||||
torch::Tensor h = hidden_states;
|
||||
|
||||
if (fc1_ctx && is_sequence_sharded(*fc1_ctx)) {
|
||||
h = gather_sequence(hidden_states, *fc1_ctx);
|
||||
}
|
||||
|
||||
auto qkv = qkv_proj_->forward(h);
|
||||
|
||||
if (use_fused_qkv_) {
|
||||
const int64_t T = qkv.size(0);
|
||||
xllm::kernel::SplitQkvRmsnormMropeParams params;
|
||||
params.qkvg = qkv;
|
||||
params.q_weight = q_norm_->weight();
|
||||
params.k_weight = k_norm_->weight();
|
||||
params.cos_sin = mrope_cos_sin;
|
||||
params.gather_pattern = mrope_gather_pattern_;
|
||||
params.eps = rms_norm_eps_;
|
||||
params.num_q_heads = num_heads_;
|
||||
params.num_kv_heads = num_kv_heads_;
|
||||
params.head_size = head_dim_;
|
||||
|
||||
auto [q, k, v, gate] = xllm::kernel::split_qkv_rmsnorm_mrope(params);
|
||||
|
||||
auto q_flat = q.view({T, q_size_});
|
||||
auto k_flat = k.view({T, kv_size_});
|
||||
auto v_flat = v.view({T, kv_size_});
|
||||
|
||||
auto out = std::get<0>(
|
||||
attn_->forward(attn_metadata, q_flat, k_flat, v_flat, kv_cache));
|
||||
out = out * torch::sigmoid(gate.view({T, q_size_}));
|
||||
|
||||
if (fc1_ctx && is_sequence_sharded(*fc1_ctx)) {
|
||||
return o_proj_->forward(out, row_parallel_reduce_mode_for_fc1(*fc1_ctx));
|
||||
}
|
||||
return o_proj_->forward(out);
|
||||
}
|
||||
|
||||
// Fallback path: weight-reordered layout [Q | G | K | V]
|
||||
torch::Tensor q, k, v;
|
||||
torch::Tensor gate;
|
||||
|
||||
if (attn_output_gate_) {
|
||||
q = qkv.slice(-1, 0, q_size_);
|
||||
gate = qkv.slice(-1, q_size_, q_size_ * 2);
|
||||
k = qkv.slice(-1, q_size_ * 2, q_size_ * 2 + kv_size_);
|
||||
v = qkv.slice(-1, q_size_ * 2 + kv_size_, q_size_ * 2 + kv_size_ * 2);
|
||||
} else {
|
||||
q = qkv.slice(-1, 0, q_size_);
|
||||
k = qkv.slice(-1, q_size_, q_size_ + kv_size_);
|
||||
v = qkv.slice(-1, q_size_ + kv_size_, q_size_ + 2 * kv_size_);
|
||||
}
|
||||
|
||||
const int64_t T = q.size(0);
|
||||
auto q_3d = q.view({T, num_heads_, head_dim_});
|
||||
q = std::get<0>(q_norm_->forward(q_3d)).view({T, q_size_});
|
||||
auto k_3d = k.view({T, num_kv_heads_, head_dim_});
|
||||
k = std::get<0>(k_norm_->forward(k_3d)).view({T, kv_size_});
|
||||
|
||||
rotary_emb_->forward(positions, q, k);
|
||||
auto out = std::get<0>(attn_->forward(attn_metadata, q, k, v, kv_cache));
|
||||
|
||||
if (attn_output_gate_) {
|
||||
out = out * torch::sigmoid(gate);
|
||||
}
|
||||
|
||||
if (fc1_ctx && is_sequence_sharded(*fc1_ctx)) {
|
||||
return o_proj_->forward(out, row_parallel_reduce_mode_for_fc1(*fc1_ctx));
|
||||
}
|
||||
return o_proj_->forward(out);
|
||||
}
|
||||
|
||||
void Qwen3NextAttentionImpl::load_state_dict(const StateDict& state_dict) {
|
||||
qkv_proj_->load_state_dict(state_dict, {"q_proj.", "k_proj.", "v_proj."});
|
||||
|
||||
if (attn_output_gate_ && qkv_proj_->is_weight_loaded() &&
|
||||
!qkv_weight_reordered_) {
|
||||
// Rearrange q_proj rows from per-head interleaved [q0,g0,q1,g1,...]
|
||||
// to grouped [q0,q1,...,g0,g1,...] so forward output is [Q|G|K|V].
|
||||
auto w = qkv_proj_->weight();
|
||||
auto qg_rows = w.slice(0, 0, q_size_ * 2);
|
||||
const int64_t hidden = w.size(1);
|
||||
auto qg_3d = qg_rows.view({num_heads_, 2 * head_dim_, hidden});
|
||||
auto q_part = qg_3d.slice(1, 0, head_dim_);
|
||||
auto g_part = qg_3d.slice(1, head_dim_, 2 * head_dim_);
|
||||
auto reordered = torch::cat(
|
||||
{q_part.reshape({q_size_, hidden}), g_part.reshape({q_size_, hidden})},
|
||||
0);
|
||||
qg_rows.copy_(reordered);
|
||||
|
||||
// Reorder weight_scale and weight_offset for W8A8 dynamic quantization.
|
||||
// These are per-channel (per output row) tensors that must match the
|
||||
// reordered weight layout for correct dequantization.
|
||||
const int64_t qg_size = q_size_ * 2;
|
||||
auto reorder_per_channel = [this, qg_size](torch::Tensor tensor) {
|
||||
if (!tensor.defined() || tensor.numel() == 0) {
|
||||
return;
|
||||
}
|
||||
auto qg_part = tensor.slice(0, 0, qg_size);
|
||||
auto qg_2d = qg_part.view({num_heads_, 2 * head_dim_});
|
||||
auto q_scale = qg_2d.slice(1, 0, head_dim_);
|
||||
auto g_scale = qg_2d.slice(1, head_dim_, 2 * head_dim_);
|
||||
auto reordered_scale = torch::cat(
|
||||
{q_scale.reshape({q_size_}), g_scale.reshape({q_size_})}, 0);
|
||||
qg_part.copy_(reordered_scale);
|
||||
};
|
||||
|
||||
if (qkv_proj_->is_weight_scale_loaded()) {
|
||||
reorder_per_channel(qkv_proj_->weight_scale());
|
||||
}
|
||||
if (qkv_proj_->is_weight_offset_loaded()) {
|
||||
reorder_per_channel(qkv_proj_->weight_offset());
|
||||
}
|
||||
|
||||
qkv_weight_reordered_ = true;
|
||||
}
|
||||
|
||||
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}}));
|
||||
}
|
||||
|
||||
// Gemma RMSNorm uses (1 + w) as the scale factor, but the fused kernel
|
||||
// uses standard RMSNorm (w only). Pre-add 1 so the fused kernel produces
|
||||
// the same result as Qwen3NextRMSNorm (gemma_rms_norm).
|
||||
if (use_fused_qkv_) {
|
||||
if (q_norm_->is_weight_loaded() && !q_norm_weight_adjusted_) {
|
||||
q_norm_->weight().add_(1.0);
|
||||
q_norm_weight_adjusted_ = true;
|
||||
}
|
||||
if (k_norm_->is_weight_loaded() && !k_norm_weight_adjusted_) {
|
||||
k_norm_->weight().add_(1.0);
|
||||
k_norm_weight_adjusted_ = true;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace layer
|
||||
} // namespace xllm
|
||||
88
ex_engine/xllm_layers/npu_torch/qwen3_next_attention.h
Normal file
88
ex_engine/xllm_layers/npu_torch/qwen3_next_attention.h
Normal file
@@ -0,0 +1,88 @@
|
||||
/* Copyright 2025-2026 The xLLM Authors.
|
||||
|
||||
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 <vector>
|
||||
|
||||
#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 "kernels/ops_api.h"
|
||||
#include "layers/common/linear.h"
|
||||
#include "layers/common/partial_rotary_embedding.h"
|
||||
#include "layers/common/qwen3_next_rms_norm.h"
|
||||
|
||||
namespace xllm {
|
||||
namespace layer {
|
||||
|
||||
class Qwen3NextAttentionImpl : public torch::nn::Module {
|
||||
public:
|
||||
Qwen3NextAttentionImpl() = default;
|
||||
Qwen3NextAttentionImpl(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,
|
||||
const torch::Tensor& mrope_cos_sin);
|
||||
|
||||
torch::Tensor build_mrope_cos_sin(const torch::Tensor& positions) const;
|
||||
|
||||
void load_state_dict(const StateDict& state_dict);
|
||||
|
||||
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_;
|
||||
int64_t rotary_dim_;
|
||||
float rms_norm_eps_;
|
||||
bool use_fused_qkv_;
|
||||
bool is_interleaved_;
|
||||
bool qkv_weight_reordered_ = false;
|
||||
bool q_norm_weight_adjusted_ = false;
|
||||
bool k_norm_weight_adjusted_ = false;
|
||||
std::vector<int64_t> mrope_section_;
|
||||
torch::Tensor mrope_gather_pattern_;
|
||||
|
||||
QKVParallelLinear qkv_proj_{nullptr};
|
||||
RowParallelLinear o_proj_{nullptr};
|
||||
|
||||
Qwen3NextRMSNorm q_norm_{nullptr};
|
||||
Qwen3NextRMSNorm k_norm_{nullptr};
|
||||
|
||||
Attention attn_{nullptr};
|
||||
PartialRotaryEmbedding rotary_emb_{nullptr};
|
||||
};
|
||||
TORCH_MODULE(Qwen3NextAttention);
|
||||
|
||||
} // namespace layer
|
||||
} // namespace xllm
|
||||
@@ -0,0 +1,41 @@
|
||||
/* Copyright 2025-2026 The xLLM Authors.
|
||||
|
||||
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_next_decoder_layer_impl.h"
|
||||
|
||||
namespace xllm {
|
||||
namespace layer {
|
||||
|
||||
Qwen3NextDecoderLayerImpl::Qwen3NextDecoderLayerImpl(
|
||||
const ModelContext& context,
|
||||
int32_t layer_id)
|
||||
: Qwen3NextDecoderLayerImpl(context,
|
||||
layer_id,
|
||||
std::make_shared<Qwen3NextGatedDeltaNetImpl>(
|
||||
context.get_model_args(),
|
||||
context.get_quant_args(),
|
||||
context.get_parallel_args(),
|
||||
context.get_tensor_options())) {}
|
||||
|
||||
Qwen3NextDecoderLayerImpl::Qwen3NextDecoderLayerImpl(
|
||||
const ModelContext& context,
|
||||
int32_t layer_id,
|
||||
std::shared_ptr<Qwen3GatedDeltaNetBaseImpl> linear_attention_module)
|
||||
: Qwen3HybridDecoderLayerImplBase(context,
|
||||
layer_id,
|
||||
std::move(linear_attention_module)) {}
|
||||
|
||||
} // namespace layer
|
||||
} // namespace xllm
|
||||
@@ -0,0 +1,38 @@
|
||||
/* Copyright 2025-2026 The xLLM Authors.
|
||||
|
||||
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/npu_torch/qwen3_next_gated_delta_net.h"
|
||||
#include "layers/npu_torch/qwen3_next_hybrid_decoder_layer_base.h"
|
||||
|
||||
namespace xllm {
|
||||
namespace layer {
|
||||
|
||||
class Qwen3NextDecoderLayerImpl : public Qwen3HybridDecoderLayerImplBase {
|
||||
public:
|
||||
explicit Qwen3NextDecoderLayerImpl(const ModelContext& context,
|
||||
int32_t layer_id);
|
||||
|
||||
protected:
|
||||
Qwen3NextDecoderLayerImpl(
|
||||
const ModelContext& context,
|
||||
int32_t layer_id,
|
||||
std::shared_ptr<Qwen3GatedDeltaNetBaseImpl> linear_attention_module);
|
||||
};
|
||||
TORCH_MODULE(Qwen3NextDecoderLayer);
|
||||
|
||||
} // namespace layer
|
||||
} // namespace xllm
|
||||
118
ex_engine/xllm_layers/npu_torch/qwen3_next_gated_delta_net.cpp
Normal file
118
ex_engine/xllm_layers/npu_torch/qwen3_next_gated_delta_net.cpp
Normal file
@@ -0,0 +1,118 @@
|
||||
/* Copyright 2025-2026 The xLLM Authors.
|
||||
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_next_gated_delta_net.h"
|
||||
|
||||
#include <glog/logging.h>
|
||||
|
||||
namespace xllm {
|
||||
namespace layer {
|
||||
|
||||
Qwen3NextGatedDeltaNetImpl::Qwen3NextGatedDeltaNetImpl(
|
||||
const ModelArgs& args,
|
||||
const QuantArgs& quant_args,
|
||||
const ParallelArgs& parallel_args,
|
||||
const torch::TensorOptions& options)
|
||||
: Qwen3NextGatedDeltaNetImpl(args,
|
||||
quant_args,
|
||||
parallel_args,
|
||||
options,
|
||||
/*init_projections=*/true) {}
|
||||
|
||||
Qwen3NextGatedDeltaNetImpl::Qwen3NextGatedDeltaNetImpl(
|
||||
const ModelArgs& args,
|
||||
const QuantArgs& quant_args,
|
||||
const ParallelArgs& parallel_args,
|
||||
const torch::TensorOptions& options,
|
||||
bool init_projections)
|
||||
: Qwen3GatedDeltaNetBaseImpl(args, quant_args, parallel_args, options) {
|
||||
if (init_projections) {
|
||||
init_next_projections(args, quant_args, parallel_args, options);
|
||||
}
|
||||
}
|
||||
|
||||
void Qwen3NextGatedDeltaNetImpl::init_next_projections(
|
||||
const ModelArgs& args,
|
||||
const QuantArgs& quant_args,
|
||||
const ParallelArgs& parallel_args,
|
||||
const torch::TensorOptions& options) {
|
||||
// QKVZ projection used by Qwen3-Next linear attention.
|
||||
qkvz_proj_ = register_module("in_proj_qkvz",
|
||||
ColumnParallelLinear(args.hidden_size(),
|
||||
k_size_ * 2 + v_size_ * 2,
|
||||
/*bias=*/false,
|
||||
/*gather_output=*/false,
|
||||
quant_args,
|
||||
parallel_args.tp_group_,
|
||||
options));
|
||||
// BA projection used to derive gating and beta terms.
|
||||
ba_proj_ = register_module("in_proj_ba",
|
||||
ColumnParallelLinear(args.hidden_size(),
|
||||
num_v_heads_ * 2,
|
||||
/*bias=*/false,
|
||||
/*gather_output=*/false,
|
||||
quant_args,
|
||||
parallel_args.tp_group_,
|
||||
options));
|
||||
}
|
||||
|
||||
std::pair<torch::Tensor, torch::Tensor>
|
||||
Qwen3NextGatedDeltaNetImpl::project_decode_inputs(
|
||||
const torch::Tensor& hidden_states) {
|
||||
auto qkvz = qkvz_proj_->forward(hidden_states);
|
||||
auto ba = ba_proj_->forward(hidden_states);
|
||||
return {qkvz.view({qkvz.size(0), -1, qkvz.size(-1)}),
|
||||
ba.view({ba.size(0), -1, ba.size(-1)})};
|
||||
}
|
||||
|
||||
std::pair<torch::Tensor, torch::Tensor>
|
||||
Qwen3NextGatedDeltaNetImpl::project_flat_inputs(
|
||||
const torch::Tensor& hidden_states) {
|
||||
return {qkvz_proj_->forward(hidden_states), ba_proj_->forward(hidden_states)};
|
||||
}
|
||||
|
||||
void Qwen3NextGatedDeltaNetImpl::load_state_dict(const StateDict& state_dict) {
|
||||
load_projection_state_dict(state_dict);
|
||||
load_common_state_dict(state_dict);
|
||||
}
|
||||
|
||||
void Qwen3NextGatedDeltaNetImpl::load_projection_state_dict(
|
||||
const StateDict& state_dict) {
|
||||
auto qkvz_state_dict = state_dict.get_dict_with_prefix("in_proj_qkvz.");
|
||||
if (qkvz_state_dict.size() > 0 && !qkvz_proj_->is_weight_loaded()) {
|
||||
qkvz_proj_->load_state_dict(qkvz_state_dict);
|
||||
}
|
||||
|
||||
auto ba_state_dict = state_dict.get_dict_with_prefix("in_proj_ba.");
|
||||
if (ba_state_dict.size() > 0 && !ba_proj_->is_weight_loaded()) {
|
||||
ba_proj_->load_state_dict(ba_state_dict);
|
||||
}
|
||||
}
|
||||
|
||||
void Qwen3NextGatedDeltaNetImpl::verify_loaded_weights(
|
||||
const std::string& prefix) const {
|
||||
verify_projection_weights(prefix);
|
||||
verify_common_loaded_weights(prefix);
|
||||
}
|
||||
|
||||
void Qwen3NextGatedDeltaNetImpl::verify_projection_weights(
|
||||
const std::string& prefix) const {
|
||||
CHECK(qkvz_proj_ && qkvz_proj_->is_weight_loaded())
|
||||
<< "Missing required weight after all shards loaded: " << prefix
|
||||
<< "in_proj_qkvz.weight";
|
||||
CHECK(ba_proj_ && ba_proj_->is_weight_loaded())
|
||||
<< "Missing required weight after all shards loaded: " << prefix
|
||||
<< "in_proj_ba.weight";
|
||||
}
|
||||
|
||||
} // namespace layer
|
||||
} // namespace xllm
|
||||
66
ex_engine/xllm_layers/npu_torch/qwen3_next_gated_delta_net.h
Normal file
66
ex_engine/xllm_layers/npu_torch/qwen3_next_gated_delta_net.h
Normal file
@@ -0,0 +1,66 @@
|
||||
/* Copyright 2025-2026 The xLLM Authors.
|
||||
|
||||
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 <string>
|
||||
#include <utility>
|
||||
|
||||
#include "qwen3_gated_delta_net_base.h"
|
||||
|
||||
namespace xllm {
|
||||
namespace layer {
|
||||
|
||||
class Qwen3NextGatedDeltaNetImpl : public Qwen3GatedDeltaNetBaseImpl {
|
||||
public:
|
||||
Qwen3NextGatedDeltaNetImpl() = default;
|
||||
Qwen3NextGatedDeltaNetImpl(const ModelArgs& args,
|
||||
const QuantArgs& quant_args,
|
||||
const ParallelArgs& parallel_args,
|
||||
const torch::TensorOptions& options);
|
||||
|
||||
void load_state_dict(const StateDict& state_dict) override;
|
||||
void verify_loaded_weights(const std::string& prefix) const override;
|
||||
|
||||
protected:
|
||||
Qwen3NextGatedDeltaNetImpl(const ModelArgs& args,
|
||||
const QuantArgs& quant_args,
|
||||
const ParallelArgs& parallel_args,
|
||||
const torch::TensorOptions& options,
|
||||
bool init_projections);
|
||||
|
||||
std::pair<torch::Tensor, torch::Tensor> project_decode_inputs(
|
||||
const torch::Tensor& hidden_states) override;
|
||||
std::pair<torch::Tensor, torch::Tensor> project_flat_inputs(
|
||||
const torch::Tensor& hidden_states) override;
|
||||
|
||||
virtual void load_projection_state_dict(const StateDict& state_dict);
|
||||
virtual void verify_projection_weights(const std::string& prefix) const;
|
||||
|
||||
void init_next_projections(const ModelArgs& args,
|
||||
const QuantArgs& quant_args,
|
||||
const ParallelArgs& parallel_args,
|
||||
const torch::TensorOptions& options);
|
||||
|
||||
private:
|
||||
ColumnParallelLinear qkvz_proj_{nullptr};
|
||||
ColumnParallelLinear ba_proj_{nullptr};
|
||||
};
|
||||
TORCH_MODULE(Qwen3NextGatedDeltaNet);
|
||||
|
||||
} // namespace layer
|
||||
} // namespace xllm
|
||||
@@ -0,0 +1,176 @@
|
||||
/* Copyright 2025-2026 The xLLM Authors.
|
||||
|
||||
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_next_hybrid_decoder_layer_base.h"
|
||||
|
||||
#include <algorithm>
|
||||
#include <optional>
|
||||
#include <tuple>
|
||||
|
||||
#include "common/flash_comm1_context.h"
|
||||
|
||||
namespace xllm {
|
||||
namespace layer {
|
||||
|
||||
Qwen3HybridDecoderLayerImplBase::Qwen3HybridDecoderLayerImplBase(
|
||||
const ModelContext& context,
|
||||
int32_t layer_id,
|
||||
std::shared_ptr<Qwen3GatedDeltaNetBaseImpl> linear_attention_module) {
|
||||
const auto& model_args = context.get_model_args();
|
||||
const auto& quant_args = context.get_quant_args();
|
||||
const auto& parallel_args = context.get_parallel_args();
|
||||
const auto& options = context.get_tensor_options();
|
||||
const bool use_full_attention = is_full_attention_layer(model_args, layer_id);
|
||||
|
||||
// Initialize attention layers
|
||||
if (use_full_attention) {
|
||||
attention_ = register_module(
|
||||
"self_attn",
|
||||
Qwen3NextAttention(
|
||||
model_args, quant_args, parallel_args, options, layer_id));
|
||||
} else {
|
||||
linear_attention_ =
|
||||
register_module("linear_attn", std::move(linear_attention_module));
|
||||
}
|
||||
|
||||
// Initialize norm layers
|
||||
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));
|
||||
|
||||
// Initialize mlp
|
||||
auto mlp_only_layers = model_args.mlp_only_layers();
|
||||
if ((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) {
|
||||
moe_mlp_ = register_module("mlp",
|
||||
FusedMoE(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 Qwen3HybridDecoderLayerImplBase::load_state_dict(
|
||||
const StateDict& state_dict) {
|
||||
if (attention_) {
|
||||
attention_->load_state_dict(state_dict.get_dict_with_prefix("self_attn."));
|
||||
} else {
|
||||
linear_attention_->load_state_dict(
|
||||
state_dict.get_dict_with_prefix("linear_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."));
|
||||
}
|
||||
}
|
||||
|
||||
void Qwen3HybridDecoderLayerImplBase::verify_loaded_weights(
|
||||
const std::string& prefix) const {
|
||||
if (linear_attention_) {
|
||||
linear_attention_->verify_loaded_weights(prefix + "linear_attn.");
|
||||
}
|
||||
}
|
||||
|
||||
torch::Tensor Qwen3HybridDecoderLayerImplBase::forward(
|
||||
torch::Tensor& x,
|
||||
std::optional<torch::Tensor>& residual,
|
||||
torch::Tensor& positions,
|
||||
const AttentionMetadata& attn_metadata,
|
||||
KVCache& kv_cache,
|
||||
const ModelInputParams& input_params,
|
||||
const torch::Tensor& mrope_cos_sin) {
|
||||
const FlashComm1Context* fc1_ctx = get_current_flash_comm1_context();
|
||||
// Pre-attention norm
|
||||
if (!residual.has_value()) {
|
||||
residual = x;
|
||||
x = std::get<0>(input_norm_->forward(x));
|
||||
} else {
|
||||
if (fc1_ctx && is_sequence_sharded(*fc1_ctx) &&
|
||||
residual.value().size(0) != x.size(0)) {
|
||||
residual = maybe_shard_residual(residual.value(), *fc1_ctx);
|
||||
}
|
||||
if (fc1_ctx && is_sequence_sharded(*fc1_ctx)) {
|
||||
CHECK_EQ(residual.value().size(0), x.size(0))
|
||||
<< "FC1 input residual and hidden states must share the same "
|
||||
<< "padded local sequence layout.";
|
||||
}
|
||||
std::tie(x, residual) = input_norm_->forward(x, residual);
|
||||
}
|
||||
|
||||
// Attention
|
||||
if (attention_) {
|
||||
x = attention_->forward(
|
||||
positions, x, attn_metadata, kv_cache, mrope_cos_sin);
|
||||
} else {
|
||||
x = linear_attention_->forward(x, attn_metadata, kv_cache, input_params);
|
||||
}
|
||||
|
||||
// Post-attention norm
|
||||
// Ensure the residual layout matches the attention output before post_norm.
|
||||
if (fc1_ctx && is_sequence_sharded(*fc1_ctx) && residual.has_value() &&
|
||||
residual.value().size(0) != x.size(0)) {
|
||||
residual = maybe_shard_residual(residual.value(), *fc1_ctx);
|
||||
CHECK_EQ(residual.value().size(0), x.size(0))
|
||||
<< "FC1 post-attention residual and hidden states must share the same "
|
||||
<< "padded local sequence layout.";
|
||||
}
|
||||
|
||||
std::tie(x, residual) = post_norm_->forward(x, residual);
|
||||
|
||||
// MLP forward
|
||||
if (moe_mlp_) {
|
||||
x = moe_mlp_(x, input_params);
|
||||
} else {
|
||||
x = mlp_(x);
|
||||
}
|
||||
|
||||
return x;
|
||||
}
|
||||
|
||||
torch::Tensor Qwen3HybridDecoderLayerImplBase::build_mrope_cos_sin(
|
||||
const torch::Tensor& positions) const {
|
||||
if (attention_) {
|
||||
return attention_->build_mrope_cos_sin(positions);
|
||||
}
|
||||
return {};
|
||||
}
|
||||
|
||||
} // namespace layer
|
||||
} // namespace xllm
|
||||
@@ -0,0 +1,90 @@
|
||||
/* Copyright 2025-2026 The xLLM Authors.
|
||||
|
||||
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 <memory>
|
||||
#include <string>
|
||||
|
||||
#include "framework/kv_cache/kv_cache.h"
|
||||
#include "framework/model/model_input_params.h"
|
||||
#include "framework/model_context.h"
|
||||
#include "framework/state_dict/state_dict.h"
|
||||
#include "layers/common/dense_mlp.h"
|
||||
#include "layers/common/qwen3_next_rms_norm.h"
|
||||
#include "layers/npu_torch/fused_moe.h"
|
||||
#include "layers/npu_torch/qwen3_gated_delta_net_base.h"
|
||||
#include "layers/npu_torch/qwen3_next_attention.h"
|
||||
|
||||
namespace xllm {
|
||||
namespace layer {
|
||||
|
||||
class Qwen3HybridDecoderLayerModule : public torch::nn::Module {
|
||||
public:
|
||||
virtual void load_state_dict(const StateDict& state_dict) = 0;
|
||||
virtual void verify_loaded_weights(const std::string& prefix) const = 0;
|
||||
virtual 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,
|
||||
const torch::Tensor& mrope_cos_sin = {}) = 0;
|
||||
virtual torch::Tensor build_mrope_cos_sin(
|
||||
const torch::Tensor& positions) const {
|
||||
return {};
|
||||
}
|
||||
};
|
||||
|
||||
using Qwen3HybridDecoderLayerModulePtr =
|
||||
std::shared_ptr<Qwen3HybridDecoderLayerModule>;
|
||||
|
||||
class Qwen3HybridDecoderLayerImplBase : public Qwen3HybridDecoderLayerModule {
|
||||
public:
|
||||
explicit Qwen3HybridDecoderLayerImplBase(
|
||||
const ModelContext& context,
|
||||
int32_t layer_id,
|
||||
std::shared_ptr<Qwen3GatedDeltaNetBaseImpl> linear_attention_module);
|
||||
|
||||
void load_state_dict(const StateDict& state_dict) override;
|
||||
|
||||
void verify_loaded_weights(const std::string& prefix) const override;
|
||||
|
||||
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,
|
||||
const torch::Tensor& mrope_cos_sin = {}) override;
|
||||
|
||||
torch::Tensor build_mrope_cos_sin(
|
||||
const torch::Tensor& positions) const override;
|
||||
|
||||
protected:
|
||||
Qwen3NextAttention attention_{nullptr};
|
||||
std::shared_ptr<Qwen3GatedDeltaNetBaseImpl> linear_attention_;
|
||||
|
||||
DenseMLP mlp_{nullptr};
|
||||
FusedMoE moe_mlp_{nullptr};
|
||||
|
||||
Qwen3NextRMSNorm input_norm_{nullptr};
|
||||
Qwen3NextRMSNorm post_norm_{nullptr};
|
||||
};
|
||||
|
||||
} // namespace layer
|
||||
} // namespace xllm
|
||||
Reference in New Issue
Block a user