[init] baseline7 from project_6
This commit is contained in:
root
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
|
||||
185
ex_engine/xllm_layers/npu_torch/qwen3_5_gated_delta_net.cpp
Normal file
185
ex_engine/xllm_layers/npu_torch/qwen3_5_gated_delta_net.cpp
Normal file
@@ -0,0 +1,185 @@
|
||||
/* 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_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_padded_inputs(
|
||||
const torch::Tensor& hidden_states,
|
||||
const AttentionMetadata& attn_metadata) {
|
||||
auto qkv = reshape_qkvz_with_pad(attn_metadata,
|
||||
in_proj_qkv_->forward(hidden_states));
|
||||
auto z_proj =
|
||||
reshape_qkvz_with_pad(attn_metadata, in_proj_z_->forward(hidden_states));
|
||||
auto b_proj =
|
||||
reshape_qkvz_with_pad(attn_metadata, in_proj_b_->forward(hidden_states));
|
||||
auto a_proj =
|
||||
reshape_qkvz_with_pad(attn_metadata, in_proj_a_->forward(hidden_states));
|
||||
return {merge_qkvz_from_split_activations(qkv, z_proj),
|
||||
merge_ba_from_split_activations(b_proj, a_proj)};
|
||||
}
|
||||
|
||||
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
|
||||
58
ex_engine/xllm_layers/npu_torch/qwen3_5_gated_delta_net.h
Normal file
58
ex_engine/xllm_layers/npu_torch/qwen3_5_gated_delta_net.h
Normal file
@@ -0,0 +1,58 @@
|
||||
/* 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 <string>
|
||||
#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_padded_inputs(
|
||||
const torch::Tensor& hidden_states,
|
||||
const AttentionMetadata& attn_metadata) override;
|
||||
|
||||
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
|
||||
576
ex_engine/xllm_layers/npu_torch/qwen3_gated_delta_net_base.cpp
Normal file
576
ex_engine/xllm_layers/npu_torch/qwen3_gated_delta_net_base.cpp
Normal file
@@ -0,0 +1,576 @@
|
||||
/* 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_gated_delta_net_base.h"
|
||||
|
||||
#include <glog/logging.h>
|
||||
#include <torch/torch.h>
|
||||
|
||||
#include <tuple>
|
||||
|
||||
#include "xllm/core/kernels/ops_api.h"
|
||||
|
||||
namespace xllm {
|
||||
namespace layer {
|
||||
|
||||
namespace {
|
||||
torch::Tensor l2norm(const torch::Tensor& x, int64_t dim, double eps = 1e-6) {
|
||||
auto norm = torch::sqrt(torch::sum(torch::square(x), dim, true) + eps);
|
||||
return x / norm;
|
||||
}
|
||||
|
||||
std::tuple<torch::Tensor, torch::Tensor> torch_recurrent_gated_delta_rule(
|
||||
torch::Tensor query,
|
||||
torch::Tensor key,
|
||||
torch::Tensor value,
|
||||
torch::Tensor g,
|
||||
torch::Tensor beta,
|
||||
std::optional<torch::Tensor> initial_state,
|
||||
bool output_final_state = true,
|
||||
bool use_qk_l2norm_in_kernel = true) {
|
||||
auto initial_dtype = query.dtype();
|
||||
|
||||
if (use_qk_l2norm_in_kernel) {
|
||||
query = l2norm(query, -1, 1e-6);
|
||||
key = l2norm(key, -1, 1e-6);
|
||||
}
|
||||
|
||||
auto to_float32_and_transpose = [](torch::Tensor x) {
|
||||
return x.transpose(1, 2).contiguous().to(torch::kFloat32);
|
||||
};
|
||||
query = to_float32_and_transpose(query);
|
||||
key = to_float32_and_transpose(key);
|
||||
value = to_float32_and_transpose(value);
|
||||
beta = to_float32_and_transpose(beta);
|
||||
g = to_float32_and_transpose(g);
|
||||
|
||||
int64_t batch_size = key.size(0);
|
||||
int64_t num_heads = key.size(1);
|
||||
int64_t sequence_length = key.size(2);
|
||||
int64_t k_head_dim = key.size(3);
|
||||
int64_t v_head_dim = value.size(3);
|
||||
|
||||
float scale_val = 1.0 / std::sqrt(static_cast<float>(query.size(-1)));
|
||||
torch::Tensor scale = torch::tensor(scale_val, query.options());
|
||||
query = query * scale;
|
||||
torch::Tensor core_attn_out = torch::zeros(
|
||||
{batch_size, num_heads, sequence_length, v_head_dim},
|
||||
torch::TensorOptions().dtype(torch::kFloat32).device(value.device()));
|
||||
torch::Tensor last_recurrent_state;
|
||||
if (!initial_state.has_value()) {
|
||||
last_recurrent_state = torch::zeros(
|
||||
{batch_size, num_heads, k_head_dim, v_head_dim},
|
||||
torch::TensorOptions().dtype(torch::kFloat32).device(value.device()));
|
||||
} else {
|
||||
last_recurrent_state =
|
||||
initial_state.value().to(value.device(), torch::kFloat32);
|
||||
}
|
||||
|
||||
for (int64_t i = 0; i < sequence_length; ++i) {
|
||||
torch::Tensor q_t = query.select(2, i);
|
||||
torch::Tensor k_t = key.select(2, i);
|
||||
torch::Tensor v_t = value.select(2, i);
|
||||
torch::Tensor g_t = g.select(2, i).exp().unsqueeze(-1).unsqueeze(-1);
|
||||
torch::Tensor beta_t = beta.select(2, i).unsqueeze(-1);
|
||||
last_recurrent_state = last_recurrent_state * g_t;
|
||||
torch::Tensor kv_mem =
|
||||
torch::sum(last_recurrent_state * k_t.unsqueeze(-1), -2);
|
||||
torch::Tensor delta = (v_t - kv_mem) * beta_t;
|
||||
last_recurrent_state =
|
||||
last_recurrent_state + k_t.unsqueeze(-1) * delta.unsqueeze(-2);
|
||||
core_attn_out.select(2, i) =
|
||||
torch::sum(last_recurrent_state * q_t.unsqueeze(-1), -2);
|
||||
}
|
||||
|
||||
core_attn_out = core_attn_out.transpose(1, 2).contiguous().to(initial_dtype);
|
||||
return std::make_tuple(core_attn_out, last_recurrent_state);
|
||||
}
|
||||
|
||||
std::tuple<torch::Tensor, torch::Tensor> torch_chunk_gated_delta_rule(
|
||||
torch::Tensor query,
|
||||
torch::Tensor key,
|
||||
torch::Tensor value,
|
||||
torch::Tensor g,
|
||||
torch::Tensor beta,
|
||||
int64_t chunk_size = 64,
|
||||
c10::optional<torch::Tensor> initial_state = c10::nullopt,
|
||||
bool output_final_state = true,
|
||||
bool use_qk_l2norm_in_kernel = true) {
|
||||
auto initial_dtype = query.dtype();
|
||||
if (use_qk_l2norm_in_kernel) {
|
||||
query = l2norm(query, -1, 1e-6);
|
||||
key = l2norm(key, -1, 1e-6);
|
||||
}
|
||||
auto to_float32 = [](torch::Tensor x) {
|
||||
return x.transpose(1, 2).contiguous().to(torch::kFloat32);
|
||||
};
|
||||
|
||||
query = to_float32(query);
|
||||
key = to_float32(key);
|
||||
value = to_float32(value);
|
||||
beta = to_float32(beta);
|
||||
g = to_float32(g);
|
||||
|
||||
auto batch_size = query.size(0);
|
||||
auto num_heads = query.size(1);
|
||||
auto sequence_length = query.size(2);
|
||||
auto k_head_dim = key.size(-1);
|
||||
auto v_head_dim = value.size(-1);
|
||||
|
||||
int64_t pad_size = (chunk_size - sequence_length % chunk_size) % chunk_size;
|
||||
query = torch::nn::functional::pad(
|
||||
query, torch::nn::functional::PadFuncOptions({0, 0, 0, pad_size}));
|
||||
key = torch::nn::functional::pad(
|
||||
key, torch::nn::functional::PadFuncOptions({0, 0, 0, pad_size}));
|
||||
value = torch::nn::functional::pad(
|
||||
value, torch::nn::functional::PadFuncOptions({0, 0, 0, pad_size}));
|
||||
beta = torch::nn::functional::pad(
|
||||
beta, torch::nn::functional::PadFuncOptions({0, pad_size}));
|
||||
g = torch::nn::functional::pad(
|
||||
g, torch::nn::functional::PadFuncOptions({0, pad_size}));
|
||||
|
||||
int64_t total_sequence_length = sequence_length + pad_size;
|
||||
float scale = 1.0 / std::sqrt(static_cast<float>(query.size(-1)));
|
||||
query = query * scale;
|
||||
auto v_beta = value * beta.unsqueeze(-1);
|
||||
auto k_beta = key * beta.unsqueeze(-1);
|
||||
auto reshape_to_chunks = [chunk_size](torch::Tensor x) {
|
||||
auto shape = x.sizes();
|
||||
std::vector<int64_t> new_shape = {
|
||||
shape[0], shape[1], shape[2] / chunk_size, chunk_size, shape[3]};
|
||||
return x.reshape(new_shape);
|
||||
};
|
||||
|
||||
query = reshape_to_chunks(query);
|
||||
key = reshape_to_chunks(key);
|
||||
value = reshape_to_chunks(value);
|
||||
k_beta = reshape_to_chunks(k_beta);
|
||||
v_beta = reshape_to_chunks(v_beta);
|
||||
|
||||
auto g_shape = g.sizes();
|
||||
std::vector<int64_t> g_new_shape = {
|
||||
g_shape[0], g_shape[1], g_shape[2] / chunk_size, chunk_size};
|
||||
g = g.reshape(g_new_shape);
|
||||
auto mask = torch::triu(
|
||||
torch::ones(
|
||||
{chunk_size, chunk_size},
|
||||
torch::TensorOptions().dtype(torch::kBool).device(query.device())),
|
||||
0);
|
||||
|
||||
g = g.cumsum(-1);
|
||||
auto g_diff = g.unsqueeze(-1) - g.unsqueeze(-2);
|
||||
auto decay_mask = g_diff.tril().exp().to(torch::kFloat32);
|
||||
decay_mask = decay_mask.tril();
|
||||
auto attn = -(torch::matmul(k_beta, key.transpose(-1, -2)) * decay_mask)
|
||||
.masked_fill(mask, 0.0);
|
||||
for (int64_t i = 1; i < chunk_size; ++i) {
|
||||
if (!attn.is_contiguous()) {
|
||||
attn = attn.contiguous();
|
||||
}
|
||||
auto row = attn.slice(-2, i, i + 1)
|
||||
.slice(-1, 0, i)
|
||||
.squeeze(-2)
|
||||
.clone()
|
||||
.contiguous();
|
||||
auto sub = attn.slice(-2, 0, i).slice(-1, 0, i).clone().contiguous();
|
||||
auto row_unsq = row.unsqueeze(-1).contiguous();
|
||||
auto row_sub_mul = (row_unsq * sub).contiguous();
|
||||
auto row_sub_sum = row_sub_mul.sum(-2).contiguous();
|
||||
auto row_final = (row + row_sub_sum).contiguous();
|
||||
attn.index_put_({torch::indexing::Ellipsis,
|
||||
torch::indexing::Slice(i, i + 1),
|
||||
torch::indexing::Slice(0, i)},
|
||||
row_final.unsqueeze(-2));
|
||||
}
|
||||
|
||||
attn = attn +
|
||||
torch::eye(
|
||||
chunk_size,
|
||||
torch::TensorOptions().dtype(attn.dtype()).device(attn.device()));
|
||||
value = torch::matmul(attn, v_beta);
|
||||
auto k_cumdecay = torch::matmul(attn, (k_beta * g.exp().unsqueeze(-1)));
|
||||
torch::Tensor last_recurrent_state;
|
||||
if (!initial_state.has_value()) {
|
||||
last_recurrent_state = torch::zeros(
|
||||
{batch_size, num_heads, k_head_dim, v_head_dim},
|
||||
torch::TensorOptions().dtype(value.dtype()).device(value.device()));
|
||||
} else {
|
||||
last_recurrent_state = initial_state.value().to(value);
|
||||
}
|
||||
auto core_attn_out = torch::zeros_like(value);
|
||||
mask = torch::triu(
|
||||
torch::ones(
|
||||
{chunk_size, chunk_size},
|
||||
torch::TensorOptions().dtype(torch::kBool).device(query.device())),
|
||||
1);
|
||||
int64_t num_chunks = total_sequence_length / chunk_size;
|
||||
for (int64_t i = 0; i < num_chunks; ++i) {
|
||||
auto q_i = query.select(2, i);
|
||||
auto k_i = key.select(2, i);
|
||||
auto v_i = value.select(2, i);
|
||||
auto attn_i =
|
||||
(torch::matmul(q_i, k_i.transpose(-1, -2)) * decay_mask.select(2, i))
|
||||
.masked_fill_(mask, 0.0);
|
||||
auto v_prime = torch::matmul(k_cumdecay.select(2, i), last_recurrent_state);
|
||||
auto v_new = v_i - v_prime;
|
||||
auto attn_inter = torch::matmul(q_i * g.select(2, i).unsqueeze(-1).exp(),
|
||||
last_recurrent_state);
|
||||
core_attn_out.select(2, i) = attn_inter + torch::matmul(attn_i, v_new);
|
||||
auto g_i_last = g.select(2, i).select(-1, -1).unsqueeze(-1);
|
||||
auto g_exp_term = (g_i_last - g.select(2, i)).exp().unsqueeze(-1);
|
||||
auto k_g_exp = (k_i * g_exp_term).transpose(-1, -2).contiguous();
|
||||
last_recurrent_state = last_recurrent_state * g_i_last.unsqueeze(-1).exp() +
|
||||
torch::matmul(k_g_exp, v_new);
|
||||
}
|
||||
auto core_attn_out_shape = core_attn_out.sizes();
|
||||
std::vector<int64_t> reshape_shape = {
|
||||
core_attn_out_shape[0],
|
||||
core_attn_out_shape[1],
|
||||
core_attn_out_shape[2] * core_attn_out_shape[3],
|
||||
core_attn_out_shape[4]};
|
||||
core_attn_out = core_attn_out.reshape(reshape_shape);
|
||||
core_attn_out = core_attn_out.slice(2, 0, sequence_length);
|
||||
core_attn_out = core_attn_out.transpose(1, 2).contiguous().to(initial_dtype);
|
||||
return std::make_tuple(core_attn_out, last_recurrent_state);
|
||||
}
|
||||
} // namespace
|
||||
|
||||
Qwen3GatedDeltaNetBaseImpl::Qwen3GatedDeltaNetBaseImpl(
|
||||
const ModelArgs& args,
|
||||
const QuantArgs& quant_args,
|
||||
const ParallelArgs& parallel_args,
|
||||
const torch::TensorOptions& options) {
|
||||
tp_size_ = parallel_args.tp_group_->world_size();
|
||||
rank_ = parallel_args.tp_group_->rank();
|
||||
num_k_heads_ = args.linear_num_key_heads();
|
||||
num_v_heads_ = args.linear_num_value_heads();
|
||||
head_k_dim_ = args.linear_key_head_dim();
|
||||
head_v_dim_ = args.linear_value_head_dim();
|
||||
k_size_ = num_k_heads_ * head_k_dim_;
|
||||
v_size_ = num_v_heads_ * head_v_dim_;
|
||||
conv_kernel_size_ = args.linear_conv_kernel_dim();
|
||||
|
||||
// Shared causal conv projection over mixed QKV states.
|
||||
conv1d_ = register_module("conv1d",
|
||||
ColumnParallelLinear(args.linear_conv_kernel_dim(),
|
||||
k_size_ * 2 + v_size_,
|
||||
/*bias=*/false,
|
||||
/*gather_output=*/false,
|
||||
quant_args,
|
||||
parallel_args.tp_group_,
|
||||
options));
|
||||
|
||||
auto opts = options.dtype(torch::kFloat32);
|
||||
dt_bias_ = register_parameter("dt_bias",
|
||||
torch::ones({num_v_heads_ / tp_size_}, opts),
|
||||
/*requires_grad=*/false);
|
||||
|
||||
A_log_ = register_parameter("A_log",
|
||||
torch::empty({num_v_heads_ / tp_size_}, opts),
|
||||
/*requires_grad=*/false);
|
||||
|
||||
// Output projection and gated RMSNorm shared by hybrid variants.
|
||||
o_proj_ = register_module("out_proj",
|
||||
RowParallelLinear(v_size_,
|
||||
args.hidden_size(),
|
||||
/*bias=*/false,
|
||||
/*input_is_parallelized=*/true,
|
||||
/*if_reduce_results=*/true,
|
||||
quant_args,
|
||||
parallel_args.tp_group_,
|
||||
options));
|
||||
|
||||
norm_ = register_module(
|
||||
"norm", RmsNormGated(head_v_dim_, args.rms_norm_eps(), options));
|
||||
}
|
||||
|
||||
void Qwen3GatedDeltaNetBaseImpl::load_common_state_dict(
|
||||
const StateDict& state_dict) {
|
||||
const int64_t rank = rank_;
|
||||
const int64_t world_size = tp_size_;
|
||||
const int32_t shard_tensor_count = 3;
|
||||
const std::vector<int64_t> shard_sizes = {
|
||||
k_size_ / tp_size_, k_size_ / tp_size_, v_size_ / tp_size_};
|
||||
|
||||
if (auto w = state_dict.get_tensor("conv1d.weight"); w.defined()) {
|
||||
conv1d_->load_state_dict(
|
||||
StateDict({{"weight", w.squeeze(1)}}), shard_tensor_count, shard_sizes);
|
||||
}
|
||||
o_proj_->load_state_dict(state_dict.get_dict_with_prefix("out_proj."));
|
||||
if (auto w = state_dict.get_tensor("norm.weight"); w.defined()) {
|
||||
norm_->load_state_dict(StateDict({{"weight", w}}));
|
||||
}
|
||||
LOAD_SHARDED_WEIGHT(dt_bias, 0);
|
||||
LOAD_SHARDED_WEIGHT(A_log, 0);
|
||||
}
|
||||
|
||||
void Qwen3GatedDeltaNetBaseImpl::verify_common_loaded_weights(
|
||||
const std::string& prefix) const {
|
||||
CHECK(dt_bias_is_loaded_)
|
||||
<< "Missing required weight after all shards loaded: " << prefix
|
||||
<< "dt_bias";
|
||||
CHECK(A_log_is_loaded_) << "Missing required weight after all shards loaded: "
|
||||
<< prefix << "A_log";
|
||||
}
|
||||
|
||||
torch::Tensor Qwen3GatedDeltaNetBaseImpl::forward(
|
||||
const torch::Tensor& hidden_states,
|
||||
const AttentionMetadata& attn_metadata,
|
||||
KVCache& kv_cache,
|
||||
const ModelInputParams& input_params) {
|
||||
auto [qkvz_padded, ba_padded] =
|
||||
project_padded_inputs(hidden_states, attn_metadata);
|
||||
int64_t batch_size = qkvz_padded.size(0);
|
||||
int64_t seq_len = qkvz_padded.size(1);
|
||||
|
||||
torch::Tensor qkvz_flat =
|
||||
qkvz_padded.view({batch_size * seq_len, qkvz_padded.size(-1)});
|
||||
torch::Tensor ba_flat =
|
||||
ba_padded.view({batch_size * seq_len, ba_padded.size(-1)});
|
||||
xllm::kernel::FusedQkvzbaSplitReshapeParams fused_params;
|
||||
fused_params.mixed_qkvz = qkvz_flat;
|
||||
fused_params.mixed_ba = ba_flat;
|
||||
fused_params.num_heads_qk = static_cast<int32_t>(num_k_heads_ / tp_size_);
|
||||
fused_params.num_heads_v = static_cast<int32_t>(num_v_heads_ / tp_size_);
|
||||
fused_params.head_qk = static_cast<int32_t>(head_k_dim_);
|
||||
fused_params.head_v = static_cast<int32_t>(head_v_dim_);
|
||||
|
||||
torch::Tensor mixed_qkv, z, b, a;
|
||||
std::tie(mixed_qkv, z, b, a) =
|
||||
xllm::kernel::fused_qkvzba_split_reshape_cat(fused_params);
|
||||
|
||||
mixed_qkv = mixed_qkv.view({batch_size, seq_len, mixed_qkv.size(-1)});
|
||||
z = z.view({batch_size, seq_len, num_v_heads_ / tp_size_, head_v_dim_});
|
||||
b = b.view({batch_size, seq_len, num_v_heads_ / tp_size_});
|
||||
a = a.view({batch_size, seq_len, num_v_heads_ / tp_size_});
|
||||
|
||||
torch::Tensor conv_cache = kv_cache.get_conv_cache();
|
||||
torch::Tensor ssm_cache = kv_cache.get_ssm_cache();
|
||||
torch::Tensor g, beta, core_attn_out, last_recurrent_state;
|
||||
auto device = mixed_qkv.device();
|
||||
auto conv_weight = conv1d_->weight();
|
||||
auto linear_state_indices = get_linear_state_indices(input_params, device);
|
||||
|
||||
if (attn_metadata.is_prefill) {
|
||||
mixed_qkv = mixed_qkv.transpose(1, 2);
|
||||
torch::Tensor conv_state =
|
||||
(seq_len < conv_kernel_size_ - 1)
|
||||
? torch::pad(mixed_qkv, {0, conv_kernel_size_ - 1 - seq_len})
|
||||
: (seq_len > conv_kernel_size_ - 1)
|
||||
? mixed_qkv.narrow(
|
||||
-1, seq_len - conv_kernel_size_ + 1, conv_kernel_size_ - 1)
|
||||
: mixed_qkv;
|
||||
conv_state = conv_state.transpose(1, 2).contiguous();
|
||||
conv_cache.index_put_({linear_state_indices},
|
||||
conv_state.to(conv_cache.dtype()));
|
||||
torch::Tensor bias;
|
||||
auto conv_output =
|
||||
torch::conv1d(mixed_qkv,
|
||||
conv_weight.unsqueeze(1).to(device),
|
||||
bias,
|
||||
/*stride=*/std::vector<int64_t>{1},
|
||||
/*padding=*/std::vector<int64_t>{3},
|
||||
/*dilation=*/std::vector<int64_t>{1},
|
||||
/*groups=*/static_cast<int64_t>(mixed_qkv.size(1)));
|
||||
mixed_qkv = torch::silu(conv_output.slice(2, 0, seq_len));
|
||||
|
||||
} else {
|
||||
xllm::kernel::CausalConv1dUpdateParams conv1d_params;
|
||||
conv1d_params.x = mixed_qkv.reshape({-1, mixed_qkv.size(-1)});
|
||||
conv1d_params.conv_state = conv_cache;
|
||||
conv1d_params.weight = conv_weight;
|
||||
conv1d_params.conv_state_indices = linear_state_indices;
|
||||
conv1d_params.block_idx_last_scheduled_token =
|
||||
std::optional<torch::Tensor>();
|
||||
conv1d_params.initial_state_idx = std::optional<torch::Tensor>();
|
||||
conv1d_params.query_start_loc = attn_metadata.q_cu_seq_lens;
|
||||
conv1d_params.max_query_len = attn_metadata.max_query_len;
|
||||
mixed_qkv = xllm::kernel::causal_conv1d_update(conv1d_params);
|
||||
// Reshape back to 3D [batch_size, dim, seq_len]
|
||||
mixed_qkv =
|
||||
mixed_qkv.view({batch_size, -1, mixed_qkv.size(-1)}).contiguous();
|
||||
mixed_qkv = mixed_qkv.transpose(1, 2);
|
||||
}
|
||||
|
||||
// Compute gated delta net decay and beta terms.
|
||||
if (attn_metadata.is_prefill) {
|
||||
xllm::kernel::FusedGdnGatingParams gdn_params;
|
||||
gdn_params.A_log = A_log_;
|
||||
gdn_params.a = a.contiguous().view({-1, a.size(-1)});
|
||||
gdn_params.b = b.contiguous().view({-1, b.size(-1)});
|
||||
gdn_params.dt_bias = dt_bias_;
|
||||
gdn_params.beta = 1.0f;
|
||||
gdn_params.threshold = 20.0f;
|
||||
std::tie(g, beta) = xllm::kernel::fused_gdn_gating(gdn_params);
|
||||
g = g.squeeze(0).contiguous().view({batch_size, seq_len, a.size(-1)});
|
||||
beta = beta.squeeze(0).contiguous().view({batch_size, seq_len, b.size(-1)});
|
||||
} else {
|
||||
xllm::kernel::FusedGdnGatingParams gdn_params;
|
||||
gdn_params.A_log = A_log_;
|
||||
gdn_params.a = a.view({-1, a.size(-1)});
|
||||
gdn_params.b = b.view({-1, b.size(-1)});
|
||||
gdn_params.dt_bias = dt_bias_;
|
||||
gdn_params.beta = 1.0f;
|
||||
gdn_params.threshold = 20.0f;
|
||||
std::tie(g, beta) = xllm::kernel::fused_gdn_gating(gdn_params);
|
||||
}
|
||||
auto [processed_q, processed_k, processed_v] = process_mixed_qkv(mixed_qkv);
|
||||
// Apply chunked or recurrent gated-delta attention and update caches.
|
||||
if (attn_metadata.is_prefill) {
|
||||
xllm::kernel::ChunkGatedDeltaRuleParams chunk_gated_delta_params;
|
||||
chunk_gated_delta_params.q = processed_q;
|
||||
chunk_gated_delta_params.k = processed_k;
|
||||
chunk_gated_delta_params.v = processed_v;
|
||||
chunk_gated_delta_params.g = g;
|
||||
chunk_gated_delta_params.beta = beta;
|
||||
// Get initial state from ssm_cache for sequences with previous state
|
||||
// Shape: [batch_size, num_heads, head_k_dim, head_v_dim]
|
||||
torch::Tensor initial_state_tensor =
|
||||
torch::index_select(ssm_cache, 0, linear_state_indices);
|
||||
// Todo: chunked-prefill/prefix-cache use initial_state
|
||||
initial_state_tensor.fill_(0.0);
|
||||
chunk_gated_delta_params.initial_state = initial_state_tensor;
|
||||
chunk_gated_delta_params.output_final_state = true;
|
||||
chunk_gated_delta_params.cu_seqlens = attn_metadata.q_cu_seq_lens;
|
||||
chunk_gated_delta_params.head_first = false;
|
||||
chunk_gated_delta_params.use_qk_l2norm_in_kernel = true;
|
||||
std::tie(core_attn_out, last_recurrent_state) =
|
||||
xllm::kernel::chunk_gated_delta_rule(chunk_gated_delta_params);
|
||||
ssm_cache.index_put_(
|
||||
{linear_state_indices},
|
||||
last_recurrent_state.transpose(-1, -2).to(ssm_cache.dtype()));
|
||||
} else {
|
||||
processed_q = xllm::kernel::l2_norm(processed_q, 1e-6);
|
||||
processed_k = xllm::kernel::l2_norm(processed_k, 1e-6);
|
||||
auto zero = torch::zeros({1}, attn_metadata.q_seq_lens.options());
|
||||
torch::Tensor actual_seq_lengths =
|
||||
torch::cat({zero, attn_metadata.q_seq_lens}, 0);
|
||||
double scale = 1.0 / std::sqrt(static_cast<float>(processed_q.size(-1)));
|
||||
core_attn_out = xllm::kernel::recurrent_gated_delta_rule(
|
||||
processed_q.reshape(
|
||||
{-1, processed_q.size(-2), processed_q.size(-1)}),
|
||||
processed_k.reshape(
|
||||
{-1, processed_k.size(-2), processed_k.size(-1)}),
|
||||
processed_v.reshape(
|
||||
{-1, processed_v.size(-2), processed_v.size(-1)}),
|
||||
ssm_cache,
|
||||
beta.squeeze(0).contiguous(),
|
||||
scale,
|
||||
actual_seq_lengths,
|
||||
linear_state_indices,
|
||||
c10::nullopt,
|
||||
g.squeeze(0).contiguous(),
|
||||
c10::nullopt)
|
||||
.unsqueeze(0)
|
||||
.contiguous();
|
||||
}
|
||||
|
||||
auto z_reshaped = z.view({-1, z.size(-1)});
|
||||
auto core_attn_out_reshaped =
|
||||
core_attn_out.view({-1, core_attn_out.size(-1)});
|
||||
auto norm_out = norm_->forward(core_attn_out_reshaped, z_reshaped);
|
||||
auto z_shape_og = z.sizes().vec();
|
||||
norm_out = norm_out.view(z_shape_og);
|
||||
norm_out = norm_out.view({-1, norm_out.size(2), norm_out.size(3)});
|
||||
|
||||
// Project the normalized attention output back to hidden size.
|
||||
auto rearranged_norm =
|
||||
norm_out.reshape({norm_out.size(0), norm_out.size(1) * norm_out.size(2)});
|
||||
rearranged_norm = reshape_qkvz_unpad(attn_metadata, rearranged_norm);
|
||||
auto attn_output = o_proj_->forward(rearranged_norm);
|
||||
return attn_output;
|
||||
}
|
||||
|
||||
torch::Tensor Qwen3GatedDeltaNetBaseImpl::reshape_qkvz_unpad(
|
||||
const AttentionMetadata& attn_metadata,
|
||||
const torch::Tensor& padded_qkvz) const {
|
||||
if (!attn_metadata.is_prefill) {
|
||||
return padded_qkvz;
|
||||
}
|
||||
std::vector<torch::Tensor> valid_batches;
|
||||
int64_t bs = attn_metadata.q_seq_lens.size(0);
|
||||
int64_t max_len = attn_metadata.max_query_len;
|
||||
const auto& ori_seq_lens = attn_metadata.q_seq_lens;
|
||||
auto reshaped_qkvz = padded_qkvz.view({bs, max_len, -1});
|
||||
for (int64_t b = 0; b < bs; ++b) {
|
||||
int64_t ori_len = ori_seq_lens[b].template item<int64_t>();
|
||||
torch::Tensor valid_batch = reshaped_qkvz[b].slice(0, 0, ori_len);
|
||||
valid_batches.push_back(valid_batch);
|
||||
}
|
||||
return torch::cat(valid_batches, 0).contiguous();
|
||||
}
|
||||
|
||||
torch::Tensor Qwen3GatedDeltaNetBaseImpl::get_linear_state_indices(
|
||||
const ModelInputParams& input_params,
|
||||
const torch::Device& device) const {
|
||||
CHECK(!input_params.linear_state_ids.empty())
|
||||
<< "linear_state_ids must be populated for gated delta net";
|
||||
if (input_params.linear_state_indices.defined()) {
|
||||
return input_params.linear_state_indices;
|
||||
}
|
||||
return torch::tensor(
|
||||
input_params.linear_state_ids,
|
||||
torch::TensorOptions().dtype(torch::kInt).device(device));
|
||||
}
|
||||
|
||||
torch::Tensor Qwen3GatedDeltaNetBaseImpl::reshape_qkvz_with_pad(
|
||||
const AttentionMetadata& attn_metadata,
|
||||
const torch::Tensor& qkvz) const {
|
||||
int64_t bs = attn_metadata.q_seq_lens.size(0);
|
||||
int64_t max_len = attn_metadata.max_query_len;
|
||||
const auto& start_loc = attn_metadata.q_seq_lens;
|
||||
if (!attn_metadata.is_prefill) {
|
||||
return qkvz.view({qkvz.size(0), -1, qkvz.size(-1)});
|
||||
}
|
||||
std::vector<torch::Tensor> batches;
|
||||
int64_t idx = 0;
|
||||
for (int64_t b = 0; b < bs; ++b) {
|
||||
int64_t cur_len = start_loc[b].template item<int64_t>();
|
||||
torch::Tensor batch = qkvz.slice(0, idx, idx + cur_len).contiguous();
|
||||
idx = idx + cur_len;
|
||||
if (batch.size(0) != max_len) {
|
||||
batch = batch.size(0) > max_len
|
||||
? batch.slice(0, 0, max_len).contiguous()
|
||||
: torch::nn::functional::pad(
|
||||
batch,
|
||||
torch::nn::functional::PadFuncOptions(
|
||||
{0, 0, 0, max_len - batch.size(0)}))
|
||||
.contiguous();
|
||||
}
|
||||
batches.push_back(batch);
|
||||
}
|
||||
auto ret = torch::stack(batches, 0).contiguous();
|
||||
return ret;
|
||||
}
|
||||
|
||||
std::tuple<torch::Tensor, torch::Tensor, torch::Tensor>
|
||||
Qwen3GatedDeltaNetBaseImpl::process_mixed_qkv(torch::Tensor& mixed_qkv) const {
|
||||
mixed_qkv = mixed_qkv.transpose(1, 2);
|
||||
int64_t batch_size = mixed_qkv.size(0);
|
||||
int64_t seq_len = mixed_qkv.size(1);
|
||||
std::vector<int64_t> split_sizes = {
|
||||
k_size_ / tp_size_, k_size_ / tp_size_, v_size_ / tp_size_};
|
||||
auto processed_qkv = torch::split(mixed_qkv, split_sizes, 2);
|
||||
auto processed_q = processed_qkv[0];
|
||||
auto processed_k = processed_qkv[1];
|
||||
auto processed_v = processed_qkv[2];
|
||||
processed_q = processed_q.view(
|
||||
{batch_size, seq_len, num_k_heads_ / tp_size_, head_k_dim_});
|
||||
processed_k = processed_k.view(
|
||||
{batch_size, seq_len, num_k_heads_ / tp_size_, head_k_dim_});
|
||||
processed_v = processed_v.view(
|
||||
{batch_size, seq_len, num_v_heads_ / tp_size_, head_v_dim_});
|
||||
return std::make_tuple(processed_q, processed_k, processed_v);
|
||||
}
|
||||
|
||||
} // namespace layer
|
||||
} // namespace xllm
|
||||
90
ex_engine/xllm_layers/npu_torch/qwen3_gated_delta_net_base.h
Normal file
90
ex_engine/xllm_layers/npu_torch/qwen3_gated_delta_net_base.h
Normal file
@@ -0,0 +1,90 @@
|
||||
/* 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 <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_padded_inputs(
|
||||
const torch::Tensor& hidden_states,
|
||||
const AttentionMetadata& attn_metadata) = 0;
|
||||
|
||||
void load_common_state_dict(const StateDict& state_dict);
|
||||
void verify_common_loaded_weights(const std::string& prefix) const;
|
||||
|
||||
torch::Tensor reshape_qkvz_with_pad(const AttentionMetadata& attn_metadata,
|
||||
const torch::Tensor& qkvz) const;
|
||||
torch::Tensor reshape_qkvz_unpad(const AttentionMetadata& attn_metadata,
|
||||
const torch::Tensor& padded_qkvz) const;
|
||||
torch::Tensor get_linear_state_indices(const ModelInputParams& input_params,
|
||||
const torch::Device& device) 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