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

feat(CRITICAL): 从 GitHub 扫描搬运 ixformer SDK + xllm 完整 GDN/MoE 代码

Browse Source 
来源:
  1. Chranos/ixformer (GitHub) → ixformer_sdk/ (230 files, 70K lines)
     - inference/functions/vllm.py: vllm_moe_topk_softmax 完整实现 (2033 lines)
     - inference/functions/moe.py: MoE ops 完整实现 (1380 lines)
     - contrib/vllm_flash_attn/: FA2 Python 接口 (1018 lines)
     - contrib/tgi/fused_moe.py: TGI fused MoE (429 lines)
     - csrc/include/ixformer/: C++ kernel headers + cmake

  2. Deep-Spark/xllm (GitHub) → upstream_ref/xllm_latest/ (+15 files)
     - npu_torch/qwen3_5_decoder_layer_impl.cpp/.h
     - npu_torch/qwen3_5_gated_delta_net.cpp/.h
     - npu_torch/qwen3_next_*.cpp/.h (6 files)
     - npu_torch/attention.cpp/.h + fused_moe.cpp/.h + CMakeLists.txt
     - models/llm/qwen3_5.h + qwen3_5_mtp.h + qwen3_next.h
     - models/vlm/qwen3_5.h

调用链完整性:
  ixformer_sdk/inference/functions/vllm.py
    → ops.infer.moe_topk_softmax() (C++ 层)
    → 这就是 base 镜像 libixformer.so 里的实现

  upstream_ref/xllm_latest/core/layers/ilu/fused_moe.cpp
    → ixformer::infer::topk_softmax() (直接 C++ 调用)
    → ixformer::infer::group_gemm() → 完整 7-step MoE pipeline
This commit is contained in:
project6-dev
2026-08-11 02:31:56 +00:00
parent a8b16da5da
commit 87a19d2d00
250 changed files with 76690 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

28
upstream_ref/xllm_latest/core/layers/npu_torch/CMakeLists.txt Executable file
View 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
upstream_ref/xllm_latest/core/layers/npu_torch/attention.cpp Normal file
View 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
upstream_ref/xllm_latest/core/layers/npu_torch/attention.h Normal file
View 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
upstream_ref/xllm_latest/core/layers/npu_torch/fused_moe.cpp Normal file
View 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
upstream_ref/xllm_latest/core/layers/npu_torch/fused_moe.h Normal file
View 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

32
upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_5_decoder_layer_impl.cpp Normal file
View File

@@ -0,0 +1,32 @@
/* Copyright 2026 The xLLM Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
https://github.com/jd-opensource/xllm/blob/main/LICENSE
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "qwen3_5_decoder_layer_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
upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_5_decoder_layer_impl.h Normal file
View File

@@ -0,0 +1,32 @@
/* Copyright 2026 The xLLM Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
https://github.com/jd-opensource/xllm/blob/main/LICENSE
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#pragma once
#include "layers/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

240
upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_next_attention.cpp Normal file
View File

@@ -0,0 +1,240 @@
/* 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_next_attention.h"
#include <glog/logging.h>
#include <tuple>
#include <vector>
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));
// 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) {
auto qkv = qkv_proj_->forward(hidden_states);
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_}));
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);
}
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_) {
// 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);
}
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_) {
q_norm_->weight().add_(1.0);
k_norm_->weight().add_(1.0);
}
}
} // namespace layer
} // namespace xllm

85
upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_next_attention.h Normal file
View File

@@ -0,0 +1,85 @@
/* 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 <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_;
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

41
upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_next_decoder_layer_impl.cpp Normal file
View File

@@ -0,0 +1,41 @@
/* 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_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

38
upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_next_decoder_layer_impl.h Normal file
View File

@@ -0,0 +1,38 @@
/* Copyright 2026 The xLLM Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
https://github.com/jd-opensource/xllm/blob/main/LICENSE
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#pragma once
#include "layers/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

113
upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_next_gated_delta_net.cpp Normal file
View File

@@ -0,0 +1,113 @@
/* 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_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_padded_inputs(
const torch::Tensor& hidden_states,
const AttentionMetadata& attn_metadata) {
auto qkvz = qkvz_proj_->forward(hidden_states);
auto ba = ba_proj_->forward(hidden_states);
return {reshape_qkvz_with_pad(attn_metadata, qkvz),
reshape_qkvz_with_pad(attn_metadata, ba)};
}
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

65
upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_next_gated_delta_net.h Normal file
View File

@@ -0,0 +1,65 @@
/* 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_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_padded_inputs(
const torch::Tensor& hidden_states,
const AttentionMetadata& attn_metadata) 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

155
upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_next_hybrid_decoder_layer_base.cpp Normal file
View File

@@ -0,0 +1,155 @@
/* 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_next_hybrid_decoder_layer_base.h"
#include <algorithm>
#include <optional>
#include <tuple>
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) {
// Pre-attention norm
if (!residual.has_value()) {
residual = x;
x = std::get<0>(input_norm_->forward(x));
} else {
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
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

90
upstream_ref/xllm_latest/core/layers/npu_torch/qwen3_next_hybrid_decoder_layer_base.h Normal file
View 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 <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