feat(CRITICAL): 从 GitHub 扫描搬运 ixformer SDK + xllm 完整 GDN/MoE 代码
来源:
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
28
upstream_ref/xllm_latest/core/layers/npu_torch/CMakeLists.txt
Executable file
28
upstream_ref/xllm_latest/core/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
upstream_ref/xllm_latest/core/layers/npu_torch/attention.cpp
Normal file
152
upstream_ref/xllm_latest/core/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
upstream_ref/xllm_latest/core/layers/npu_torch/attention.h
Normal file
70
upstream_ref/xllm_latest/core/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
upstream_ref/xllm_latest/core/layers/npu_torch/fused_moe.cpp
Normal file
513
upstream_ref/xllm_latest/core/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
upstream_ref/xllm_latest/core/layers/npu_torch/fused_moe.h
Normal file
113
upstream_ref/xllm_latest/core/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 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
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
218
upstream_ref/xllm_latest/models/llm/qwen3_5.h
Normal file
218
upstream_ref/xllm_latest/models/llm/qwen3_5.h
Normal file
@@ -0,0 +1,218 @@
|
||||
/* 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 <cstdint>
|
||||
#include <string>
|
||||
#include <unordered_set>
|
||||
#include <vector>
|
||||
|
||||
#include "core/layers/npu_torch/qwen3_5_decoder_layer_impl.h"
|
||||
#include "models/model_registry.h"
|
||||
#include "qwen3_next.h"
|
||||
|
||||
namespace xllm {
|
||||
|
||||
class Qwen3_5ModelImpl : public Qwen3NextModelImpl {
|
||||
public:
|
||||
explicit Qwen3_5ModelImpl(const ModelContext& context)
|
||||
: Qwen3NextModelImpl(context, /*init_decoder_layers=*/false) {
|
||||
const int32_t n_layers = context.get_model_args().n_layers();
|
||||
for (int32_t layer_id = 0; layer_id < n_layers; ++layer_id) {
|
||||
add_decoder_layer(
|
||||
std::make_shared<layer::Qwen3_5DecoderLayerImpl>(context, layer_id));
|
||||
}
|
||||
}
|
||||
};
|
||||
TORCH_MODULE(Qwen3_5Model);
|
||||
|
||||
class Qwen3_5ForCausalLMImpl : public Qwen3NextForCausalLMImpl {
|
||||
public:
|
||||
explicit Qwen3_5ForCausalLMImpl(const ModelContext& context)
|
||||
: Qwen3NextForCausalLMImpl(context, /*init_model=*/false) {
|
||||
set_model_module(std::make_shared<Qwen3_5ModelImpl>(context));
|
||||
}
|
||||
};
|
||||
TORCH_MODULE(Qwen3_5ForCausalLM);
|
||||
|
||||
#define LOAD_ARG_TEXT_OR_ROOT(arg_name, json_key, default_value) \
|
||||
LOAD_ARG_OR(arg_name, "text_config." json_key, default_value); \
|
||||
LOAD_ARG_OR(arg_name, json_key, args->arg_name())
|
||||
|
||||
#define LOAD_ARG_TEXT_OR_ROOT_CHAIN(arg_name, json_key, default_value) \
|
||||
LOAD_ARG_TEXT_OR_ROOT(arg_name, json_key, default_value)
|
||||
|
||||
#define LOAD_QWEN3_5_ROPE_ARG(arg_name, default_value) \
|
||||
LOAD_ARG_OR(arg_name, "text_config." #arg_name, default_value); \
|
||||
LOAD_ARG_OR(arg_name, #arg_name, args->arg_name()); \
|
||||
LOAD_ARG_OR( \
|
||||
arg_name, "text_config.rope_scaling." #arg_name, args->arg_name()); \
|
||||
LOAD_ARG_OR(arg_name, "rope_scaling." #arg_name, args->arg_name()); \
|
||||
LOAD_ARG_OR( \
|
||||
arg_name, "text_config.rope_parameters." #arg_name, args->arg_name()); \
|
||||
LOAD_ARG_OR(arg_name, "rope_parameters." #arg_name, args->arg_name())
|
||||
|
||||
#define LOAD_QWEN3_5_NEXT_COMPAT_ARGS(default_moe_intermediate_size, \
|
||||
default_num_experts, \
|
||||
default_num_experts_per_tok, \
|
||||
default_shared_expert_intermediate_size) \
|
||||
LOAD_ARG_TEXT_OR_ROOT(attention_bias, "attention_bias", false); \
|
||||
LOAD_ARG_TEXT_OR_ROOT(attention_dropout, "attention_dropout", 0.0f); \
|
||||
LOAD_ARG_TEXT_OR_ROOT(bos_token_id, "bos_token_id", 151643); \
|
||||
LOAD_ARG_TEXT_OR_ROOT(decoder_sparse_step, "decoder_sparse_step", 1); \
|
||||
LOAD_ARG_TEXT_OR_ROOT(eos_token_id, "eos_token_id", 151645); \
|
||||
LOAD_ARG_TEXT_OR_ROOT(head_dim, "head_dim", 256); \
|
||||
LOAD_ARG_TEXT_OR_ROOT(hidden_act, "hidden_act", "silu"); \
|
||||
LOAD_ARG_TEXT_OR_ROOT(hidden_size, "hidden_size", 2048); \
|
||||
LOAD_ARG_TEXT_OR_ROOT(initializer_range, "initializer_range", 0.02f); \
|
||||
LOAD_ARG_TEXT_OR_ROOT(intermediate_size, "intermediate_size", 5120); \
|
||||
LOAD_ARG_TEXT_OR_ROOT( \
|
||||
max_position_embeddings, "max_position_embeddings", 262144); \
|
||||
LOAD_ARG_TEXT_OR_ROOT(max_window_layers, "max_window_layers", 28); \
|
||||
LOAD_ARG_TEXT_OR_ROOT(moe_intermediate_size, \
|
||||
"moe_intermediate_size", \
|
||||
default_moe_intermediate_size); \
|
||||
LOAD_ARG_TEXT_OR_ROOT(norm_topk_prob, "norm_topk_prob", true); \
|
||||
LOAD_ARG_TEXT_OR_ROOT(n_heads, "num_attention_heads", 16); \
|
||||
LOAD_ARG_TEXT_OR_ROOT(num_experts, "num_experts", default_num_experts); \
|
||||
LOAD_ARG_TEXT_OR_ROOT(num_experts_per_tok, \
|
||||
"num_experts_per_tok", \
|
||||
default_num_experts_per_tok); \
|
||||
LOAD_ARG_TEXT_OR_ROOT(n_layers, "num_hidden_layers", 48); \
|
||||
LOAD_ARG_OR(n_kv_heads, "text_config.num_key_value_heads", 2); \
|
||||
LOAD_ARG_OR( \
|
||||
n_kv_heads, "num_key_value_heads", args->n_kv_heads().value_or(2)); \
|
||||
LOAD_ARG_TEXT_OR_ROOT(output_router_logits, "output_router_logits", false); \
|
||||
LOAD_ARG_TEXT_OR_ROOT(rms_norm_eps, "rms_norm_eps", 1e-6); \
|
||||
LOAD_QWEN3_5_ROPE_ARG(rope_theta, 10000000.0f); \
|
||||
LOAD_ARG_TEXT_OR_ROOT(router_aux_loss_coef, "router_aux_loss_coef", 0.001f); \
|
||||
LOAD_ARG_TEXT_OR_ROOT(use_sliding_window, "use_sliding_window", false); \
|
||||
LOAD_ARG_TEXT_OR_ROOT(sliding_window, "sliding_window", 4096); \
|
||||
LOAD_ARG_TEXT_OR_ROOT(tie_word_embeddings, "tie_word_embeddings", false); \
|
||||
LOAD_ARG_TEXT_OR_ROOT(vocab_size, "vocab_size", 151936); \
|
||||
LOAD_ARG_TEXT_OR_ROOT( \
|
||||
mlp_only_layers, "mlp_only_layers", std::vector<int32_t>()); \
|
||||
LOAD_ARG_TEXT_OR_ROOT(attn_output_gate, "attn_output_gate", true); \
|
||||
LOAD_ARG_TEXT_OR_ROOT( \
|
||||
full_attention_interval, "full_attention_interval", 4); \
|
||||
LOAD_ARG_TEXT_OR_ROOT(linear_conv_kernel_dim, "linear_conv_kernel_dim", 4); \
|
||||
LOAD_ARG_TEXT_OR_ROOT(linear_key_head_dim, "linear_key_head_dim", 128); \
|
||||
LOAD_ARG_TEXT_OR_ROOT(linear_num_key_heads, "linear_num_key_heads", 16); \
|
||||
LOAD_ARG_TEXT_OR_ROOT(linear_num_value_heads, "linear_num_value_heads", 32); \
|
||||
LOAD_ARG_TEXT_OR_ROOT(linear_value_head_dim, "linear_value_head_dim", 128); \
|
||||
LOAD_QWEN3_5_ROPE_ARG(partial_rotary_factor, 0.25f); \
|
||||
LOAD_ARG_OR(rope_scaling_mrope_section, \
|
||||
"text_config.rope_scaling.mrope_section", \
|
||||
std::vector<int64_t>()); \
|
||||
LOAD_ARG_OR(rope_scaling_mrope_section, \
|
||||
"text_config.rope_parameters.mrope_section", \
|
||||
args->rope_scaling_mrope_section()); \
|
||||
LOAD_ARG_OR(rope_scaling_mrope_section, \
|
||||
"rope_parameters.mrope_section", \
|
||||
args->rope_scaling_mrope_section()); \
|
||||
LOAD_ARG_OR(rope_scaling_mrope_interleaved, \
|
||||
"text_config.rope_scaling.mrope_interleaved", \
|
||||
false); \
|
||||
LOAD_ARG_OR(rope_scaling_mrope_interleaved, \
|
||||
"text_config.rope_parameters.mrope_interleaved", \
|
||||
args->rope_scaling_mrope_interleaved()); \
|
||||
LOAD_ARG_OR(rope_scaling_mrope_interleaved, \
|
||||
"rope_parameters.mrope_interleaved", \
|
||||
args->rope_scaling_mrope_interleaved()); \
|
||||
LOAD_ARG_TEXT_OR_ROOT(shared_expert_intermediate_size, \
|
||||
"shared_expert_intermediate_size", \
|
||||
default_shared_expert_intermediate_size); \
|
||||
LOAD_ARG_OR( \
|
||||
num_nextn_predict_layers, "text_config.mtp_num_hidden_layers", 0); \
|
||||
LOAD_ARG_OR(num_nextn_predict_layers, \
|
||||
"mtp_num_hidden_layers", \
|
||||
args->num_nextn_predict_layers()); \
|
||||
LOAD_ARG_OR(num_nextn_predict_layers, \
|
||||
"text_config.num_nextn_predict_layers", \
|
||||
args->num_nextn_predict_layers()); \
|
||||
LOAD_ARG_OR(num_nextn_predict_layers, \
|
||||
"num_nextn_predict_layers", \
|
||||
args->num_nextn_predict_layers()); \
|
||||
LOAD_ARG_OR( \
|
||||
layer_types, "text_config.layer_types", std::vector<std::string>()); \
|
||||
LOAD_ARG_OR(layer_types, "layer_types", args->layer_types()); \
|
||||
LOAD_ARG_OR( \
|
||||
layer_types, "text_config.layers_block_type", args->layer_types()); \
|
||||
LOAD_ARG_OR(layer_types, "layers_block_type", args->layer_types()); \
|
||||
LOAD_ARG_OR( \
|
||||
n_routed_experts, "text_config.n_routed_experts", args->num_experts()); \
|
||||
LOAD_ARG_OR(n_routed_experts, "n_routed_experts", args->num_experts()); \
|
||||
SET_ARG(n_shared_experts, \
|
||||
args->shared_expert_intermediate_size() > 0 ? 1 : 0); \
|
||||
SET_ARG(scoring_func, "softmax"); \
|
||||
SET_ARG(topk_method, ""); \
|
||||
SET_ARG(n_group, -1); \
|
||||
SET_ARG(topk_group, 0); \
|
||||
SET_ARG(routed_scaling_factor, 1.0f); \
|
||||
SET_ARG(stop_token_ids, \
|
||||
std::unordered_set<int32_t>({args->eos_token_id()})); \
|
||||
LOAD_ARG_TEXT_OR_ROOT(mamba_ssm_dtype, "mamba_ssm_dtype", "float32")
|
||||
|
||||
#define LOAD_QWEN3_5_TYPE_AND_DTYPE(default_model_type) \
|
||||
LOAD_ARG_OR(model_type, "model_type", default_model_type); \
|
||||
LOAD_ARG_OR(dtype, "text_config.dtype", "bfloat16"); \
|
||||
LOAD_ARG_OR(dtype, "dtype", args->dtype()); \
|
||||
LOAD_ARG_OR(dtype, "text_config.torch_dtype", args->dtype()); \
|
||||
LOAD_ARG_OR(dtype, "torch_dtype", args->dtype())
|
||||
|
||||
REGISTER_CAUSAL_MODEL(qwen3_5, Qwen3_5ForCausalLM);
|
||||
REGISTER_MODEL_ARGS(qwen3_5, [&] {
|
||||
LOAD_QWEN3_5_TYPE_AND_DTYPE("qwen3_5");
|
||||
LOAD_QWEN3_5_NEXT_COMPAT_ARGS(/*moe_intermediate_size=*/0,
|
||||
/*num_experts=*/0,
|
||||
/*num_experts_per_tok=*/0,
|
||||
/*shared_expert_intermediate_size=*/0);
|
||||
});
|
||||
|
||||
REGISTER_CAUSAL_MODEL(qwen3_5_text, Qwen3_5ForCausalLM);
|
||||
REGISTER_MODEL_ARGS(qwen3_5_text, [&] {
|
||||
LOAD_QWEN3_5_TYPE_AND_DTYPE("qwen3_5_text");
|
||||
LOAD_QWEN3_5_NEXT_COMPAT_ARGS(/*moe_intermediate_size=*/0,
|
||||
/*num_experts=*/0,
|
||||
/*num_experts_per_tok=*/0,
|
||||
/*shared_expert_intermediate_size=*/0);
|
||||
});
|
||||
|
||||
REGISTER_CAUSAL_MODEL(qwen3_5_moe, Qwen3_5ForCausalLM);
|
||||
REGISTER_MODEL_ARGS(qwen3_5_moe, [&] {
|
||||
LOAD_QWEN3_5_TYPE_AND_DTYPE("qwen3_5_moe");
|
||||
LOAD_QWEN3_5_NEXT_COMPAT_ARGS(/*moe_intermediate_size=*/512,
|
||||
/*num_experts=*/512,
|
||||
/*num_experts_per_tok=*/10,
|
||||
/*shared_expert_intermediate_size=*/512);
|
||||
});
|
||||
|
||||
REGISTER_CAUSAL_MODEL(qwen3_5_moe_text, Qwen3_5ForCausalLM);
|
||||
REGISTER_MODEL_ARGS(qwen3_5_moe_text, [&] {
|
||||
LOAD_QWEN3_5_TYPE_AND_DTYPE("qwen3_5_moe_text");
|
||||
LOAD_QWEN3_5_NEXT_COMPAT_ARGS(/*moe_intermediate_size=*/512,
|
||||
/*num_experts=*/512,
|
||||
/*num_experts_per_tok=*/10,
|
||||
/*shared_expert_intermediate_size=*/512);
|
||||
});
|
||||
|
||||
#undef LOAD_QWEN3_5_TYPE_AND_DTYPE
|
||||
#undef LOAD_QWEN3_5_NEXT_COMPAT_ARGS
|
||||
#undef LOAD_QWEN3_5_ROPE_ARG
|
||||
#undef LOAD_ARG_TEXT_OR_ROOT_CHAIN
|
||||
#undef LOAD_ARG_TEXT_OR_ROOT
|
||||
|
||||
} // namespace xllm
|
||||
280
upstream_ref/xllm_latest/models/llm/qwen3_5_mtp.h
Normal file
280
upstream_ref/xllm_latest/models/llm/qwen3_5_mtp.h
Normal file
@@ -0,0 +1,280 @@
|
||||
/* 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 <glog/logging.h>
|
||||
|
||||
#include <algorithm>
|
||||
#include <memory>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
#include "core/framework/model/model_input_params.h"
|
||||
#include "core/layers/common/linear.h"
|
||||
#include "models/model_registry.h"
|
||||
#include "qwen3_5.h"
|
||||
|
||||
namespace xllm {
|
||||
|
||||
namespace {
|
||||
|
||||
StateDict find_lm_head_state_dict(const StateDict& state_dict) {
|
||||
static const std::vector<std::string> kLmHeadPrefixes = {
|
||||
"lm_head.",
|
||||
"model.lm_head.",
|
||||
"language_model.lm_head.",
|
||||
"model.language_model.lm_head."};
|
||||
for (const auto& prefix : kLmHeadPrefixes) {
|
||||
auto sub_dict = state_dict.get_dict_with_prefix(prefix);
|
||||
if (sub_dict.get_tensor("weight").defined() ||
|
||||
sub_dict.get_tensor("qweight").defined()) {
|
||||
return sub_dict;
|
||||
}
|
||||
}
|
||||
return StateDict({}, "");
|
||||
}
|
||||
|
||||
bool load_qwen3_5_mtp_model_args(const JsonReader& json,
|
||||
ModelArgs* args,
|
||||
const std::string& base_model_type,
|
||||
const std::string& mtp_model_type) {
|
||||
auto base_loader = ModelRegistry::get_model_args_loader(base_model_type);
|
||||
if (base_loader == nullptr || base_loader(json, args) == false) {
|
||||
return false;
|
||||
}
|
||||
|
||||
int32_t mtp_num_layers = args->num_nextn_predict_layers();
|
||||
if (mtp_num_layers <= 0) {
|
||||
mtp_num_layers = 1;
|
||||
}
|
||||
args->model_type(mtp_model_type);
|
||||
args->num_nextn_predict_layers(mtp_num_layers);
|
||||
args->n_layers(mtp_num_layers);
|
||||
args->layer_types(std::vector<std::string>(
|
||||
static_cast<size_t>(mtp_num_layers), "full_attention"));
|
||||
return true;
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
class Qwen3_5MtpModelImpl : public Qwen3HybridModelImplBase {
|
||||
public:
|
||||
explicit Qwen3_5MtpModelImpl(const ModelContext& context)
|
||||
: Qwen3HybridModelImplBase(context) {
|
||||
const auto& options = context.get_tensor_options();
|
||||
const int32_t n_layers =
|
||||
std::max<int32_t>(static_cast<int32_t>(model_args_.n_layers()), 1);
|
||||
|
||||
pre_fc_norm_embedding_ = register_module(
|
||||
"pre_fc_norm_embedding",
|
||||
layer::Qwen3NextRMSNorm(
|
||||
model_args_.hidden_size(), model_args_.rms_norm_eps(), options));
|
||||
pre_fc_norm_hidden_ = register_module(
|
||||
"pre_fc_norm_hidden",
|
||||
layer::Qwen3NextRMSNorm(
|
||||
model_args_.hidden_size(), model_args_.rms_norm_eps(), options));
|
||||
fc_ = register_module("fc",
|
||||
layer::ReplicatedLinear(model_args_.hidden_size() * 2,
|
||||
model_args_.hidden_size(),
|
||||
/*bias=*/false,
|
||||
QuantArgs(),
|
||||
options));
|
||||
|
||||
layers_.reserve(n_layers);
|
||||
for (int32_t layer_id = 0; layer_id < n_layers; ++layer_id) {
|
||||
add_decoder_layer(
|
||||
std::make_shared<layer::Qwen3_5DecoderLayerImpl>(context, layer_id));
|
||||
}
|
||||
}
|
||||
|
||||
ModelOutput forward(torch::Tensor tokens,
|
||||
torch::Tensor positions,
|
||||
std::vector<KVCache>& kv_caches,
|
||||
const ModelInputParams& input_params) override {
|
||||
torch::NoGradGuard no_grad;
|
||||
|
||||
if (dp_size_ > 1 && tokens.sizes() == 0) {
|
||||
tokens = torch::tensor({1}).to(torch::kInt32).to(device_);
|
||||
positions = torch::tensor({0}).to(torch::kInt32).to(device_);
|
||||
}
|
||||
|
||||
auto attn_metadata = layer::AttentionMetadataBuilder::build(
|
||||
input_params,
|
||||
model_args_.enable_mla(),
|
||||
build_attention_mask(input_params));
|
||||
|
||||
torch::Tensor embedding = embed_tokens_(tokens);
|
||||
torch::Tensor hidden = input_params.input_embedding;
|
||||
if (hidden.defined() == false) {
|
||||
hidden = embedding;
|
||||
}
|
||||
|
||||
embedding = std::get<0>(pre_fc_norm_embedding_->forward(embedding));
|
||||
hidden = std::get<0>(pre_fc_norm_hidden_->forward(hidden));
|
||||
torch::Tensor mtp_hidden = fc_(torch::cat({embedding, hidden}, -1));
|
||||
|
||||
CHECK_EQ(kv_caches.size(), layers_.size());
|
||||
std::optional<torch::Tensor> residual = std::nullopt;
|
||||
for (size_t i = 0; i < layers_.size(); ++i) {
|
||||
mtp_hidden = layers_[i]->forward(mtp_hidden,
|
||||
residual,
|
||||
positions,
|
||||
attn_metadata,
|
||||
kv_caches[i],
|
||||
input_params);
|
||||
}
|
||||
auto [new_mtp_hidden, new_res] = norm_->forward(mtp_hidden, residual);
|
||||
mtp_hidden = new_mtp_hidden;
|
||||
return ModelOutput(mtp_hidden);
|
||||
}
|
||||
|
||||
void load_state_dict(const StateDict& state_dict) override {
|
||||
load_shared_embeddings(state_dict);
|
||||
load_mtp_state_dict(state_dict);
|
||||
}
|
||||
|
||||
void load_shared_embeddings(const StateDict& state_dict) {
|
||||
auto embedding_state_dict =
|
||||
state_dict.get_dict_with_prefix("embed_tokens.");
|
||||
if (embedding_state_dict.get_tensor("weight").defined()) {
|
||||
shared_embedding_loaded_ = true;
|
||||
}
|
||||
embed_tokens_->load_state_dict(embedding_state_dict);
|
||||
}
|
||||
|
||||
void load_mtp_state_dict(const StateDict& state_dict) {
|
||||
if (state_dict.get_tensor("pre_fc_norm_embedding.weight").defined()) {
|
||||
pre_fc_norm_embedding_loaded_ = true;
|
||||
}
|
||||
if (state_dict.get_tensor("pre_fc_norm_hidden.weight").defined()) {
|
||||
pre_fc_norm_hidden_loaded_ = true;
|
||||
}
|
||||
if (state_dict.get_tensor("fc.weight").defined() ||
|
||||
state_dict.get_tensor("fc.qweight").defined()) {
|
||||
fc_loaded_ = true;
|
||||
}
|
||||
if (state_dict.get_tensor("norm.weight").defined()) {
|
||||
norm_loaded_ = true;
|
||||
}
|
||||
|
||||
pre_fc_norm_embedding_->load_state_dict(
|
||||
state_dict.get_dict_with_prefix("pre_fc_norm_embedding."));
|
||||
pre_fc_norm_hidden_->load_state_dict(
|
||||
state_dict.get_dict_with_prefix("pre_fc_norm_hidden."));
|
||||
fc_->load_state_dict(state_dict.get_dict_with_prefix("fc."));
|
||||
for (size_t i = 0; i < layers_.size(); ++i) {
|
||||
layers_[i]->load_state_dict(
|
||||
state_dict.get_dict_with_prefix("layers." + std::to_string(i) + "."));
|
||||
}
|
||||
norm_->load_state_dict(state_dict.get_dict_with_prefix("norm."));
|
||||
}
|
||||
|
||||
void verify_loaded_weights(const std::string& prefix) const override {
|
||||
CHECK(shared_embedding_loaded_)
|
||||
<< "Failed to find shared embedding weights for qwen3.5 mtp draft "
|
||||
"model";
|
||||
CHECK(pre_fc_norm_embedding_loaded_)
|
||||
<< "Failed to find mtp pre_fc_norm_embedding weights for qwen3.5 mtp "
|
||||
"draft model";
|
||||
CHECK(pre_fc_norm_hidden_loaded_)
|
||||
<< "Failed to find mtp pre_fc_norm_hidden weights for qwen3.5 mtp "
|
||||
"draft model";
|
||||
CHECK(fc_loaded_) << "Failed to find mtp fc weights for qwen3.5 mtp draft "
|
||||
"model";
|
||||
CHECK(norm_loaded_)
|
||||
<< "Failed to find mtp norm weights for qwen3.5 mtp draft model";
|
||||
for (size_t i = 0; i < layers_.size(); ++i) {
|
||||
layers_[i]->verify_loaded_weights(prefix + "layers." + std::to_string(i) +
|
||||
".");
|
||||
}
|
||||
}
|
||||
|
||||
private:
|
||||
layer::Qwen3NextRMSNorm pre_fc_norm_embedding_{nullptr};
|
||||
layer::Qwen3NextRMSNorm pre_fc_norm_hidden_{nullptr};
|
||||
layer::ReplicatedLinear fc_{nullptr};
|
||||
bool shared_embedding_loaded_ = false;
|
||||
bool pre_fc_norm_embedding_loaded_ = false;
|
||||
bool pre_fc_norm_hidden_loaded_ = false;
|
||||
bool fc_loaded_ = false;
|
||||
bool norm_loaded_ = false;
|
||||
};
|
||||
|
||||
class Qwen3_5MtpForCausalLMImpl : public Qwen3HybridForCausalLMImplBase {
|
||||
public:
|
||||
explicit Qwen3_5MtpForCausalLMImpl(const ModelContext& context)
|
||||
: Qwen3HybridForCausalLMImplBase(context) {
|
||||
mtp_model_ = std::make_shared<Qwen3_5MtpModelImpl>(context);
|
||||
set_model_module(mtp_model_);
|
||||
}
|
||||
|
||||
void load_model(std::unique_ptr<ModelLoader> loader) {
|
||||
static const std::vector<std::string> kEmbeddingPrefixes = {
|
||||
"model.language_model.", "language_model.model.", "model.", ""};
|
||||
static const std::vector<std::string> kMtpPrefixes = {"mtp.", "model.mtp."};
|
||||
bool lm_head_loaded = false;
|
||||
|
||||
for (const auto& state_dict : loader->get_state_dicts()) {
|
||||
auto shared_embedding_state_dict =
|
||||
state_dict->get_dict_with_prefix(kEmbeddingPrefixes);
|
||||
auto mtp_state_dict = state_dict->get_dict_with_prefix(kMtpPrefixes);
|
||||
|
||||
mtp_model_->load_shared_embeddings(shared_embedding_state_dict);
|
||||
mtp_model_->load_mtp_state_dict(mtp_state_dict);
|
||||
|
||||
if (tie_word_embeddings_) {
|
||||
lm_head_->load_state_dict(
|
||||
shared_embedding_state_dict.get_dict_with_prefix("embed_tokens."));
|
||||
if (shared_embedding_state_dict.get_tensor("embed_tokens.weight")
|
||||
.defined()) {
|
||||
lm_head_loaded = true;
|
||||
}
|
||||
} else {
|
||||
auto lm_head_state_dict = find_lm_head_state_dict(*state_dict);
|
||||
lm_head_->load_state_dict(lm_head_state_dict);
|
||||
if (lm_head_state_dict.get_tensor("weight").defined() ||
|
||||
lm_head_state_dict.get_tensor("qweight").defined()) {
|
||||
lm_head_loaded = true;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
CHECK(lm_head_loaded)
|
||||
<< "Failed to find lm_head weights for qwen3.5 mtp draft model";
|
||||
mtp_model_->verify_loaded_weights("mtp.");
|
||||
}
|
||||
|
||||
private:
|
||||
std::shared_ptr<Qwen3_5MtpModelImpl> mtp_model_;
|
||||
};
|
||||
TORCH_MODULE(Qwen3_5MtpForCausalLM);
|
||||
|
||||
REGISTER_CAUSAL_MODEL(qwen3_5_mtp, Qwen3_5MtpForCausalLM);
|
||||
REGISTER_CAUSAL_MODEL(qwen3_5_moe_mtp, Qwen3_5MtpForCausalLM);
|
||||
|
||||
REGISTER_MODEL_ARGS_LOADER(qwen3_5_mtp,
|
||||
[](const JsonReader& json, ModelArgs* args) {
|
||||
return load_qwen3_5_mtp_model_args(
|
||||
json, args, "qwen3_5", "qwen3_5_mtp");
|
||||
});
|
||||
|
||||
REGISTER_MODEL_ARGS_LOADER(qwen3_5_moe_mtp,
|
||||
[](const JsonReader& json, ModelArgs* args) {
|
||||
return load_qwen3_5_mtp_model_args(
|
||||
json, args, "qwen3_5_moe", "qwen3_5_moe_mtp");
|
||||
});
|
||||
|
||||
} // namespace xllm
|
||||
126
upstream_ref/xllm_latest/models/llm/qwen3_next.h
Normal file
126
upstream_ref/xllm_latest/models/llm/qwen3_next.h
Normal file
@@ -0,0 +1,126 @@
|
||||
/* 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 <string>
|
||||
#include <unordered_set>
|
||||
#include <vector>
|
||||
|
||||
#include "core/layers/npu_torch/qwen3_next_decoder_layer_impl.h"
|
||||
#include "models/model_registry.h"
|
||||
#include "qwen3_next_hybrid_base.h"
|
||||
|
||||
namespace xllm {
|
||||
|
||||
class Qwen3NextModelImpl : public Qwen3HybridModelImplBase {
|
||||
public:
|
||||
explicit Qwen3NextModelImpl(const ModelContext& context)
|
||||
: Qwen3NextModelImpl(context, /*init_decoder_layers=*/true) {}
|
||||
|
||||
protected:
|
||||
explicit Qwen3NextModelImpl(const ModelContext& context,
|
||||
bool init_decoder_layers)
|
||||
: Qwen3HybridModelImplBase(context) {
|
||||
if (init_decoder_layers) {
|
||||
const int32_t n_layers = context.get_model_args().n_layers();
|
||||
for (int32_t layer_id = 0; layer_id < n_layers; ++layer_id) {
|
||||
add_decoder_layer(std::make_shared<layer::Qwen3NextDecoderLayerImpl>(
|
||||
context, layer_id));
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
TORCH_MODULE(Qwen3NextModel);
|
||||
|
||||
class Qwen3NextForCausalLMImpl : public Qwen3HybridForCausalLMImplBase {
|
||||
public:
|
||||
explicit Qwen3NextForCausalLMImpl(const ModelContext& context)
|
||||
: Qwen3NextForCausalLMImpl(context, /*init_model=*/true) {}
|
||||
|
||||
protected:
|
||||
explicit Qwen3NextForCausalLMImpl(const ModelContext& context,
|
||||
bool init_model)
|
||||
: Qwen3HybridForCausalLMImplBase(context) {
|
||||
if (init_model) {
|
||||
set_model_module(std::make_shared<Qwen3NextModelImpl>(context));
|
||||
}
|
||||
}
|
||||
};
|
||||
TORCH_MODULE(Qwen3NextForCausalLM);
|
||||
|
||||
// register the causal model
|
||||
REGISTER_CAUSAL_MODEL(qwen3_next, Qwen3NextForCausalLM);
|
||||
|
||||
// register the model args
|
||||
REGISTER_MODEL_ARGS(qwen3_next, [&] {
|
||||
LOAD_ARG_OR(model_type, "model_type", "qwen3_next");
|
||||
LOAD_ARG_OR(dtype, "torch_dtype", "");
|
||||
LOAD_ARG_OR(attention_bias, "attention_bias", false);
|
||||
LOAD_ARG_OR(attention_dropout, "attention_dropout", 0.0f);
|
||||
LOAD_ARG_OR(bos_token_id, "bos_token_id", 151643);
|
||||
LOAD_ARG_OR(decoder_sparse_step, "decoder_sparse_step", 1);
|
||||
LOAD_ARG_OR(eos_token_id, "eos_token_id", 151645);
|
||||
LOAD_ARG_OR(head_dim, "head_dim", 256);
|
||||
LOAD_ARG_OR(hidden_act, "hidden_act", "silu");
|
||||
LOAD_ARG_OR(hidden_size, "hidden_size", 2048);
|
||||
LOAD_ARG_OR(initializer_range, "initializer_range", 0.02f);
|
||||
LOAD_ARG_OR(intermediate_size, "intermediate_size", 5120);
|
||||
LOAD_ARG_OR(max_position_embeddings, "max_position_embeddings", 262144);
|
||||
LOAD_ARG_OR(max_window_layers, "max_window_layers", 28);
|
||||
LOAD_ARG_OR(moe_intermediate_size, "moe_intermediate_size", 512);
|
||||
LOAD_ARG_OR(norm_topk_prob, "norm_topk_prob", true);
|
||||
LOAD_ARG_OR(n_heads, "num_attention_heads", 16);
|
||||
LOAD_ARG_OR(num_experts, "num_experts", 512);
|
||||
LOAD_ARG_OR(num_experts_per_tok, "num_experts_per_tok", 10);
|
||||
LOAD_ARG_OR(n_layers, "num_hidden_layers", 48);
|
||||
LOAD_ARG_OR(n_kv_heads, "num_key_value_heads", 2);
|
||||
LOAD_ARG_OR(output_router_logits, "output_router_logits", false);
|
||||
LOAD_ARG_OR(rms_norm_eps, "rms_norm_eps", 1e-6);
|
||||
LOAD_ARG_OR(rope_theta, "rope_theta", 10000000.0f);
|
||||
LOAD_ARG_OR(router_aux_loss_coef, "router_aux_loss_coef", 0.001f);
|
||||
LOAD_ARG_OR(use_sliding_window, "use_sliding_window", false);
|
||||
LOAD_ARG_OR(sliding_window, "sliding_window", 4096);
|
||||
LOAD_ARG_OR(tie_word_embeddings, "tie_word_embeddings", false);
|
||||
LOAD_ARG_OR(vocab_size, "vocab_size", 151936);
|
||||
LOAD_ARG_OR(mlp_only_layers, "mlp_only_layers", std::vector<int>());
|
||||
|
||||
// Additional parameters for Qwen3-Next architecture
|
||||
LOAD_ARG_OR(attn_output_gate, "attn_output_gate", true);
|
||||
LOAD_ARG_OR(full_attention_interval, "full_attention_interval", 4);
|
||||
LOAD_ARG_OR(linear_conv_kernel_dim, "linear_conv_kernel_dim", 4);
|
||||
LOAD_ARG_OR(linear_key_head_dim, "linear_key_head_dim", 128);
|
||||
LOAD_ARG_OR(linear_num_key_heads, "linear_num_key_heads", 16);
|
||||
LOAD_ARG_OR(linear_num_value_heads, "linear_num_value_heads", 32);
|
||||
LOAD_ARG_OR(linear_value_head_dim, "linear_value_head_dim", 128);
|
||||
LOAD_ARG_OR(partial_rotary_factor, "partial_rotary_factor", 0.25f);
|
||||
LOAD_ARG_OR(
|
||||
shared_expert_intermediate_size, "shared_expert_intermediate_size", 512);
|
||||
LOAD_ARG_OR(layer_types, "layer_types", std::vector<std::string>());
|
||||
|
||||
// MoE compatibility with fused_moe implementation.
|
||||
LOAD_ARG_OR(n_routed_experts, "n_routed_experts", args->num_experts());
|
||||
SET_ARG(n_shared_experts,
|
||||
args->shared_expert_intermediate_size() > 0 ? 1 : 0);
|
||||
SET_ARG(scoring_func, "softmax");
|
||||
SET_ARG(topk_method, "");
|
||||
SET_ARG(n_group, -1);
|
||||
SET_ARG(topk_group, 0);
|
||||
SET_ARG(routed_scaling_factor, 1.0);
|
||||
|
||||
SET_ARG(stop_token_ids, std::unordered_set<int32_t>({args->eos_token_id()}));
|
||||
});
|
||||
|
||||
} // namespace xllm
|
||||
323
upstream_ref/xllm_latest/models/llm/qwen3_next_hybrid_base.h
Normal file
323
upstream_ref/xllm_latest/models/llm/qwen3_next_hybrid_base.h
Normal file
@@ -0,0 +1,323 @@
|
||||
/* 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 <algorithm>
|
||||
#include <memory>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
#include "core/common/global_flags.h"
|
||||
#include "core/framework/kv_cache/kv_cache.h"
|
||||
#include "core/framework/model/model_input_params.h"
|
||||
#include "core/framework/model/model_output.h"
|
||||
#include "core/framework/model_context.h"
|
||||
#include "core/framework/model_loader.h"
|
||||
#include "core/layers/common/attention_mask.h"
|
||||
#include "core/layers/common/attention_metadata_builder.h"
|
||||
#include "core/layers/common/lm_head.h"
|
||||
#include "core/layers/common/qwen3_next_rms_norm.h"
|
||||
#include "core/layers/common/word_embedding.h"
|
||||
#include "core/layers/npu_torch/qwen3_next_hybrid_decoder_layer_base.h"
|
||||
|
||||
namespace xllm {
|
||||
|
||||
class Qwen3HybridModelModule : public torch::nn::Module {
|
||||
public:
|
||||
virtual ModelOutput forward(torch::Tensor tokens,
|
||||
torch::Tensor positions,
|
||||
std::vector<KVCache>& kv_caches,
|
||||
const ModelInputParams& input_params) = 0;
|
||||
virtual void load_state_dict(const StateDict& state_dict) = 0;
|
||||
virtual void verify_loaded_weights(const std::string& prefix) const = 0;
|
||||
virtual layer::WordEmbedding get_word_embedding() = 0;
|
||||
virtual void set_word_embedding(layer::WordEmbedding& word_embedding) = 0;
|
||||
};
|
||||
|
||||
using Qwen3HybridModelModulePtr = std::shared_ptr<Qwen3HybridModelModule>;
|
||||
|
||||
class Qwen3HybridModelImplBase : public Qwen3HybridModelModule {
|
||||
public:
|
||||
explicit Qwen3HybridModelImplBase(const ModelContext& context)
|
||||
: device_(context.get_tensor_options().device()),
|
||||
model_args_(context.get_model_args()) {
|
||||
auto options = context.get_tensor_options();
|
||||
auto parallel_args = context.get_parallel_args();
|
||||
|
||||
blocks_ = register_module("layers", torch::nn::ModuleList());
|
||||
layers_.reserve(model_args_.n_layers());
|
||||
device_ = options.device();
|
||||
dtype_ = options.dtype().toScalarType();
|
||||
norm_ = register_module(
|
||||
"norm",
|
||||
xllm::layer::Qwen3NextRMSNorm(
|
||||
model_args_.hidden_size(), model_args_.rms_norm_eps(), options));
|
||||
embed_tokens_ =
|
||||
register_module("embed_tokens", layer::WordEmbedding(context));
|
||||
int32_t mask_value = FLAGS_enable_chunked_prefill ? -9984 : 1;
|
||||
attn_mask_ = layer::AttentionMask(options.device(),
|
||||
options.dtype().toScalarType(),
|
||||
/*mask_value=*/mask_value);
|
||||
dp_size_ = parallel_args.dp_size();
|
||||
}
|
||||
|
||||
// tokens: [num_tokens]
|
||||
// positions: [num_tokens] token pos in the sequence
|
||||
ModelOutput forward(torch::Tensor tokens,
|
||||
torch::Tensor positions,
|
||||
std::vector<KVCache>& kv_caches,
|
||||
const ModelInputParams& input_params) override {
|
||||
// Disable gradient computation to reduce memory usage during inference
|
||||
torch::NoGradGuard no_grad;
|
||||
if (dp_size_ > 1) {
|
||||
if (tokens.sizes() == 0) {
|
||||
tokens = torch::tensor({1}).to(torch::kInt32).to(device_);
|
||||
positions = torch::tensor({0}).to(torch::kInt32).to(device_);
|
||||
}
|
||||
}
|
||||
|
||||
layer::AttentionMetadata attn_metadata =
|
||||
layer::AttentionMetadataBuilder::build(
|
||||
input_params,
|
||||
model_args_.enable_mla(),
|
||||
build_attention_mask(input_params));
|
||||
torch::Tensor h = embed_tokens_(tokens);
|
||||
|
||||
torch::Tensor mrope_cos_sin;
|
||||
for (const auto& layer : layers_) {
|
||||
mrope_cos_sin = layer->build_mrope_cos_sin(positions);
|
||||
if (mrope_cos_sin.defined()) break;
|
||||
}
|
||||
|
||||
std::optional<torch::Tensor> residual = std::nullopt;
|
||||
for (size_t i = 0; i < layers_.size(); i++) {
|
||||
auto& layer = layers_[i];
|
||||
h = layer->forward(h,
|
||||
residual,
|
||||
positions,
|
||||
attn_metadata,
|
||||
kv_caches[i],
|
||||
input_params,
|
||||
mrope_cos_sin);
|
||||
}
|
||||
auto [hidden_states, residual_out] = norm_->forward(h, residual);
|
||||
h = hidden_states;
|
||||
return ModelOutput(h);
|
||||
}
|
||||
|
||||
// load the weight from the checkpoint
|
||||
void load_state_dict(const StateDict& state_dict) override {
|
||||
embed_tokens_->load_state_dict(
|
||||
state_dict.get_dict_with_prefix("embed_tokens."));
|
||||
for (int i = 0; i < static_cast<int>(layers_.size()); i++) {
|
||||
layers_[i]->load_state_dict(
|
||||
state_dict.get_dict_with_prefix("layers." + std::to_string(i) + "."));
|
||||
}
|
||||
norm_->load_state_dict(state_dict.get_dict_with_prefix("norm."));
|
||||
}
|
||||
|
||||
void verify_loaded_weights(const std::string& prefix) const override {
|
||||
for (size_t i = 0; i < layers_.size(); ++i) {
|
||||
layers_[i]->verify_loaded_weights(prefix + "layers." + std::to_string(i) +
|
||||
".");
|
||||
}
|
||||
}
|
||||
|
||||
layer::WordEmbedding get_word_embedding() override { return embed_tokens_; }
|
||||
|
||||
void set_word_embedding(layer::WordEmbedding& word_embedding) override {
|
||||
embed_tokens_ = word_embedding;
|
||||
}
|
||||
|
||||
void add_decoder_layer(layer::Qwen3HybridDecoderLayerModulePtr layer) {
|
||||
layers_.push_back(layer);
|
||||
blocks_->push_back(layer);
|
||||
}
|
||||
|
||||
int32_t num_hidden_layers() const {
|
||||
return static_cast<int32_t>(layers_.size());
|
||||
}
|
||||
|
||||
protected:
|
||||
torch::Tensor build_attention_mask(const ModelInputParams& input_params) {
|
||||
max_seq_len_ = std::max(input_params.kv_max_seq_len, max_seq_len_);
|
||||
if (!FLAGS_enable_chunked_prefill) {
|
||||
return attn_mask_.get_attn_mask(max_seq_len_, dtype_, device_);
|
||||
}
|
||||
|
||||
const int32_t num_sequences = input_params.num_sequences;
|
||||
if (num_sequences <= 0) {
|
||||
return attn_mask_.get_attn_mask(max_seq_len_, dtype_, device_);
|
||||
}
|
||||
|
||||
std::vector<torch::Tensor> req_mask_vec;
|
||||
req_mask_vec.reserve(num_sequences);
|
||||
for (int32_t j = 0; j < num_sequences; ++j) {
|
||||
req_mask_vec.emplace_back(
|
||||
attn_mask_.gen_append_mask(input_params.q_seq_lens_vec[j],
|
||||
input_params.kv_seq_lens_vec[j],
|
||||
max_seq_len_,
|
||||
dtype_,
|
||||
device_));
|
||||
}
|
||||
return torch::cat(req_mask_vec, 0);
|
||||
}
|
||||
|
||||
ModelArgs model_args_;
|
||||
torch::nn::ModuleList blocks_{nullptr};
|
||||
std::vector<layer::Qwen3HybridDecoderLayerModulePtr> layers_;
|
||||
int32_t max_seq_len_ = 0;
|
||||
int32_t dp_size_ = 1;
|
||||
torch::Device device_;
|
||||
torch::ScalarType dtype_ = torch::kFloat;
|
||||
layer::Qwen3NextRMSNorm norm_{nullptr};
|
||||
layer::AttentionMask attn_mask_;
|
||||
layer::WordEmbedding embed_tokens_{nullptr};
|
||||
};
|
||||
|
||||
class Qwen3HybridForCausalLMImplBase : public torch::nn::Module {
|
||||
public:
|
||||
explicit Qwen3HybridForCausalLMImplBase(const ModelContext& context) {
|
||||
tie_word_embeddings_ = context.get_model_args().tie_word_embeddings();
|
||||
lm_head_ = register_module("lm_head", layer::LmHead(context));
|
||||
}
|
||||
|
||||
// tokens: [num_tokens]
|
||||
// positions: [num_tokens] token pos in the sequence
|
||||
// returns: [num_tokens, hidden_size]
|
||||
ModelOutput forward(const torch::Tensor& tokens,
|
||||
const torch::Tensor& positions,
|
||||
std::vector<KVCache>& kv_caches,
|
||||
const ModelInputParams& input_params) {
|
||||
return model_->forward(tokens, positions, kv_caches, input_params);
|
||||
}
|
||||
|
||||
// hidden_states: [num_tokens, hidden_size]
|
||||
// seleted_idxes: [num_tokens]
|
||||
// returns: [num_tokens, vocab_size]
|
||||
torch::Tensor logits(const torch::Tensor& hidden_states,
|
||||
const torch::Tensor& seleted_idxes) {
|
||||
auto h = hidden_states;
|
||||
if (seleted_idxes.defined()) {
|
||||
h = h.index_select(/*dim=*/0, seleted_idxes);
|
||||
}
|
||||
return lm_head_(h);
|
||||
}
|
||||
|
||||
// hidden_states: [num_tokens, hidden_size]
|
||||
// seleted_idxes: [num_tokens]
|
||||
torch::Tensor pooler(const torch::Tensor& hidden_states,
|
||||
const torch::Tensor& seleted_idxes) {
|
||||
auto h = hidden_states;
|
||||
if (seleted_idxes.defined()) {
|
||||
h = h.index_select(/*dim=*/0, seleted_idxes);
|
||||
}
|
||||
namespace F = torch::nn::functional;
|
||||
return F::normalize(h, F::NormalizeFuncOptions().p(2).dim(1));
|
||||
}
|
||||
|
||||
void load_model(std::unique_ptr<ModelLoader> loader) {
|
||||
auto has_model_weights = [](const StateDict& dict) {
|
||||
return dict.get_tensor("embed_tokens.weight").defined() ||
|
||||
dict.get_dict_with_prefix("layers.").size() > 0 ||
|
||||
dict.get_tensor("norm.weight").defined();
|
||||
};
|
||||
auto has_lm_head_weights = [](const StateDict& dict) {
|
||||
return dict.get_tensor("weight").defined() ||
|
||||
dict.get_tensor("qweight").defined();
|
||||
};
|
||||
|
||||
for (const auto& state_dict : loader->get_state_dicts()) {
|
||||
auto model_state_dict = state_dict->get_dict_with_prefix("model.");
|
||||
if (!has_model_weights(model_state_dict)) {
|
||||
auto language_model_state_dict =
|
||||
state_dict->get_dict_with_prefix("language_model.model.");
|
||||
if (has_model_weights(language_model_state_dict)) {
|
||||
model_state_dict = language_model_state_dict;
|
||||
} else {
|
||||
auto wrapped_language_model_state_dict =
|
||||
state_dict->get_dict_with_prefix("model.language_model.");
|
||||
if (has_model_weights(wrapped_language_model_state_dict)) {
|
||||
model_state_dict = wrapped_language_model_state_dict;
|
||||
}
|
||||
}
|
||||
}
|
||||
model_->load_state_dict(model_state_dict);
|
||||
|
||||
auto lm_head_state_dict = state_dict->get_dict_with_prefix("lm_head.");
|
||||
if (!has_lm_head_weights(lm_head_state_dict)) {
|
||||
auto language_model_lm_head_state_dict =
|
||||
state_dict->get_dict_with_prefix("language_model.lm_head.");
|
||||
if (has_lm_head_weights(language_model_lm_head_state_dict)) {
|
||||
lm_head_state_dict = language_model_lm_head_state_dict;
|
||||
} else {
|
||||
auto wrapped_language_model_lm_head_state_dict =
|
||||
state_dict->get_dict_with_prefix("model.language_model.lm_head.");
|
||||
if (has_lm_head_weights(wrapped_language_model_lm_head_state_dict)) {
|
||||
lm_head_state_dict = wrapped_language_model_lm_head_state_dict;
|
||||
} else {
|
||||
auto wrapped_lm_head_state_dict =
|
||||
state_dict->get_dict_with_prefix("model.lm_head.");
|
||||
if (has_lm_head_weights(wrapped_lm_head_state_dict)) {
|
||||
lm_head_state_dict = wrapped_lm_head_state_dict;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
if (!has_lm_head_weights(lm_head_state_dict) && tie_word_embeddings_) {
|
||||
auto tied_lm_head_state_dict =
|
||||
model_state_dict.get_dict_with_prefix("embed_tokens.");
|
||||
if (has_lm_head_weights(tied_lm_head_state_dict)) {
|
||||
lm_head_state_dict = tied_lm_head_state_dict;
|
||||
}
|
||||
}
|
||||
lm_head_->load_state_dict(lm_head_state_dict);
|
||||
}
|
||||
|
||||
model_->verify_loaded_weights("model.");
|
||||
}
|
||||
|
||||
virtual void prepare_expert_weight(int32_t layer_id,
|
||||
const std::vector<int32_t>& expert_ids) {
|
||||
return;
|
||||
}
|
||||
virtual void update_expert_weight(int32_t layer_id) { return; }
|
||||
|
||||
layer::LmHead get_lm_head() { return lm_head_; }
|
||||
|
||||
void set_lm_head(layer::LmHead& head) { lm_head_ = head; }
|
||||
|
||||
layer::WordEmbedding get_word_embedding() {
|
||||
return model_->get_word_embedding();
|
||||
}
|
||||
|
||||
void set_word_embedding(layer::WordEmbedding& word_embedding) {
|
||||
model_->set_word_embedding(word_embedding);
|
||||
}
|
||||
|
||||
void set_model_module(Qwen3HybridModelModulePtr model) {
|
||||
model_ = register_module("model", std::move(model));
|
||||
}
|
||||
|
||||
protected:
|
||||
bool tie_word_embeddings_{false};
|
||||
layer::LmHead lm_head_{nullptr};
|
||||
Qwen3HybridModelModulePtr model_;
|
||||
};
|
||||
|
||||
} // namespace xllm
|
||||
312
upstream_ref/xllm_latest/models/vlm/qwen3_5.h
Normal file
312
upstream_ref/xllm_latest/models/vlm/qwen3_5.h
Normal file
@@ -0,0 +1,312 @@
|
||||
/* 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 "core/framework/model/model_output.h"
|
||||
#include "core/layers/common/lm_head.h"
|
||||
#include "core/layers/common/qwen3_next_rms_norm.h"
|
||||
#include "core/layers/common/rms_norm.h"
|
||||
#include "core/layers/mlu/qwen3_5_decoder_layer.h"
|
||||
#include "core/layers/qwen3_vision_layer.h"
|
||||
#include "models/llm/llm_model_base.h"
|
||||
#include "models/model_registry.h"
|
||||
#include "models/vlm/qwen3_vl_base.h"
|
||||
#include "processors/input_processor.h"
|
||||
#include "processors/qwen2_vl_image_processor.h"
|
||||
#include "qwen3_vl.h"
|
||||
|
||||
namespace xllm {
|
||||
class Qwen3_5ModelImpl final
|
||||
: public LlmModelImplBase<layer::Qwen3_5DecoderLayer> {
|
||||
public:
|
||||
Qwen3_5ModelImpl(const ModelContext& context)
|
||||
: LlmModelImplBase<layer::Qwen3_5DecoderLayer>("qwen3_5",
|
||||
context.get_model_args()) {
|
||||
auto model_args = context.get_model_args();
|
||||
auto options = context.get_tensor_options();
|
||||
auto parallel_args = context.get_parallel_args();
|
||||
dp_size_ = parallel_args.dp_size();
|
||||
|
||||
if (!mrope_section_.empty()) {
|
||||
int64_t rotary_dim = static_cast<int64_t>(
|
||||
model_args.head_dim() * model_args.partial_rotary_factor());
|
||||
cos_sin_ = layer::rotary::get_concat_rotary_embedding(
|
||||
rotary_dim,
|
||||
model_args.max_position_embeddings(),
|
||||
model_args.rope_theta(),
|
||||
options);
|
||||
}
|
||||
|
||||
layers_.reserve(model_args.n_layers());
|
||||
rms_norm_ = register_module(
|
||||
"norm",
|
||||
layer::Qwen3NextRMSNorm(
|
||||
model_args.hidden_size(), model_args.rms_norm_eps(), options));
|
||||
embed_tokens_ =
|
||||
register_module("embed_tokens", layer::WordEmbedding(context));
|
||||
|
||||
for (int32_t i = 0; i < model_args.n_layers(); i++) {
|
||||
auto layer = layer::Qwen3_5DecoderLayer(context, i);
|
||||
layers_.push_back(layer);
|
||||
}
|
||||
}
|
||||
|
||||
void load_state_dict(const StateDict& state_dict) override {
|
||||
embed_tokens_->load_state_dict(
|
||||
state_dict.get_dict_with_prefix("embed_tokens."));
|
||||
|
||||
// call each layer's load_state_dict function
|
||||
for (size_t i = 0; i < layers_.size(); i++) {
|
||||
layers_[i]->load_state_dict(
|
||||
state_dict.get_dict_with_prefix("layers." + std::to_string(i) + "."));
|
||||
}
|
||||
rms_norm_->load_state_dict(state_dict.get_dict_with_prefix("norm."));
|
||||
}
|
||||
|
||||
std::pair<torch::Tensor, torch::Tensor> apply_mrope(
|
||||
const torch::Tensor positions) override {
|
||||
auto target_cos_sin = cos_sin_.index({positions});
|
||||
auto target_cos_sin_chunks = target_cos_sin.chunk(/*chunks=*/2, /*dim=*/-1);
|
||||
auto cos_pos = target_cos_sin_chunks[0].contiguous();
|
||||
auto sin_pos = target_cos_sin_chunks[1].contiguous();
|
||||
auto apply = [this](torch::Tensor x) {
|
||||
auto freqs_t = x[0].clone();
|
||||
int64_t mrop_length = static_cast<int64_t>(freqs_t.size(-1) / 2);
|
||||
|
||||
for (int32_t dim_idx = 1; dim_idx <= 2; ++dim_idx) {
|
||||
int64_t offset = dim_idx;
|
||||
int64_t section_len = mrope_section_[dim_idx];
|
||||
int64_t length = section_len * 3;
|
||||
|
||||
auto idx_first_half = torch::arange(offset, length, 3, torch::kLong);
|
||||
auto idx_second_half = torch::arange(
|
||||
offset + mrop_length, length + mrop_length, 3, torch::kLong);
|
||||
|
||||
auto idx_tensor =
|
||||
torch::cat({idx_first_half, idx_second_half}, 0).to(x.device());
|
||||
auto src = x[dim_idx].index_select(-1, idx_tensor);
|
||||
freqs_t.index_copy_(-1, idx_tensor, src);
|
||||
}
|
||||
return freqs_t;
|
||||
};
|
||||
cos_pos = apply(cos_pos.reshape({positions.size(0), -1, cos_pos.size(-1)}));
|
||||
sin_pos = apply(sin_pos.reshape({positions.size(0), -1, sin_pos.size(-1)}));
|
||||
return std::make_pair(cos_pos, sin_pos);
|
||||
}
|
||||
|
||||
virtual ModelOutput forward(torch::Tensor tokens,
|
||||
torch::Tensor positions,
|
||||
std::vector<KVCache>& kv_caches,
|
||||
const ModelInputParams& input_params) {
|
||||
ModelInputParams& input_params_new =
|
||||
const_cast<ModelInputParams&>(input_params);
|
||||
std::vector<torch::Tensor> deep_stacks;
|
||||
|
||||
if (dp_size_ > 1) {
|
||||
if (tokens.numel() == 0) {
|
||||
tokens = torch::tensor({1}).to(torch::kInt32).to(tokens.device());
|
||||
positions = torch::tensor({1}).to(torch::kInt32).to(positions.device());
|
||||
}
|
||||
auto& dp_token_nums = input_params_new.dp_global_token_nums;
|
||||
std::replace(dp_token_nums.begin(), dp_token_nums.end(), 0, 1);
|
||||
}
|
||||
|
||||
auto inputs_embeds = input_params.input_embedding;
|
||||
torch::Tensor h;
|
||||
if (inputs_embeds.defined()) {
|
||||
h = inputs_embeds;
|
||||
} else {
|
||||
h = embed_tokens_(tokens);
|
||||
}
|
||||
|
||||
if (!input_params_new.attn_metadata) {
|
||||
input_params_new.attn_metadata =
|
||||
std::make_shared<layer::AttentionMetadata>(
|
||||
get_attention_metadata(input_params_new, h));
|
||||
}
|
||||
|
||||
auto& attn_metadata = *(input_params_new.attn_metadata);
|
||||
bool only_prefill =
|
||||
(attn_metadata.is_prefill || attn_metadata.is_chunked_prefill);
|
||||
if (positions.dim() == 2 && only_prefill && !mrope_section_.empty()) {
|
||||
std::tie(attn_metadata.mrope_cos, attn_metadata.mrope_sin) =
|
||||
apply_mrope(positions);
|
||||
}
|
||||
|
||||
std::optional<torch::Tensor> residual;
|
||||
for (size_t i = 0; i < layers_.size(); i++) {
|
||||
auto& layer = layers_[i];
|
||||
h = layer(h,
|
||||
residual,
|
||||
positions,
|
||||
attn_metadata,
|
||||
kv_caches[i],
|
||||
input_params_new);
|
||||
}
|
||||
if (residual.has_value()) {
|
||||
h = h + residual.value();
|
||||
}
|
||||
auto hidden_states = std::get<0>(rms_norm_(h));
|
||||
return ModelOutput(hidden_states);
|
||||
}
|
||||
|
||||
private:
|
||||
int32_t dp_size_ = 1;
|
||||
layer::Qwen3NextRMSNorm rms_norm_{nullptr};
|
||||
layer::AttentionMetadata get_attention_metadata(
|
||||
const ModelInputParams& params,
|
||||
const torch::Tensor& h) {
|
||||
auto attn_metadata = layer::AttentionMetadataBuilder::build(params, false);
|
||||
// TODO: support linear attention
|
||||
return attn_metadata;
|
||||
}
|
||||
};
|
||||
TORCH_MODULE(Qwen3_5Model);
|
||||
|
||||
class Qwen3_5ForCausalLMImpl : public LlmForCausalLMImplBase<Qwen3_5Model> {
|
||||
public:
|
||||
Qwen3_5ForCausalLMImpl(const ModelContext& context)
|
||||
: LlmForCausalLMImplBase<Qwen3_5Model>(context) {}
|
||||
|
||||
torch::Tensor pooler(const torch::Tensor& hidden_states,
|
||||
const torch::Tensor& seleted_idxes) {
|
||||
auto h = hidden_states;
|
||||
if (seleted_idxes.defined()) {
|
||||
h = h.index_select(/*dim=*/0, seleted_idxes);
|
||||
}
|
||||
namespace F = torch::nn::functional;
|
||||
return F::normalize(h, F::NormalizeFuncOptions().p(2).dim(1));
|
||||
}
|
||||
};
|
||||
TORCH_MODULE(Qwen3_5ForCausalLM);
|
||||
|
||||
using Qwen3_5ForConditionalGenerationImpl =
|
||||
Qwen3VLForConditionalGenerationBase<Qwen3_VisionTransformer,
|
||||
Qwen3_5ForCausalLM>;
|
||||
TORCH_MODULE(Qwen3_5ForConditionalGeneration);
|
||||
|
||||
#define LOAD_QWEN3_5_COMMON_ARGS() \
|
||||
LOAD_ARG_OR(model_type, "model_type", "qwen3_5"); \
|
||||
LOAD_ARG_OR(dtype, "text_config.dtype", "bfloat16"); \
|
||||
LOAD_ARG_OR(vocab_size, "text_config.vocab_size", 248320); \
|
||||
LOAD_ARG_OR(hidden_size, "text_config.hidden_size", 5120); \
|
||||
LOAD_ARG_OR(hidden_act, "text_config.hidden_act", "silu"); \
|
||||
LOAD_ARG_OR(intermediate_size, "text_config.intermediate_size", 17408); \
|
||||
LOAD_ARG_OR(n_layers, "text_config.num_hidden_layers", 64); \
|
||||
LOAD_ARG_OR(n_heads, "text_config.num_attention_heads", 24); \
|
||||
LOAD_ARG(n_kv_heads, "text_config.num_key_value_heads"); \
|
||||
LOAD_ARG_OR( \
|
||||
max_position_embeddings, "text_config.max_position_embeddings", 262144); \
|
||||
LOAD_ARG_OR(rms_norm_eps, "text_config.rms_norm_eps", 1e-6); \
|
||||
LOAD_ARG_OR(eos_token_id, "text_config.eos_token_id", 248044); \
|
||||
LOAD_ARG_OR( \
|
||||
rope_theta, "text_config.rope_parameters.rope_theta", 10000000.0f); \
|
||||
LOAD_ARG_OR(head_dim, "text_config.head_dim", 256); \
|
||||
LOAD_ARG_OR(tie_word_embeddings, "tie_word_embeddings", false); \
|
||||
LOAD_ARG(layer_types, "text_config.layer_types"); \
|
||||
LOAD_ARG_OR( \
|
||||
linear_conv_kernel_dim, "text_config.linear_conv_kernel_dim", 4); \
|
||||
LOAD_ARG_OR(linear_key_head_dim, "text_config.linear_key_head_dim", 128); \
|
||||
LOAD_ARG_OR( \
|
||||
linear_value_head_dim, "text_config.linear_value_head_dim", 128); \
|
||||
LOAD_ARG_OR(linear_num_key_heads, "text_config.linear_num_key_heads", 16); \
|
||||
LOAD_ARG_OR( \
|
||||
linear_num_value_heads, "text_config.linear_num_value_heads", 48); \
|
||||
LOAD_ARG_OR( \
|
||||
full_attention_interval, "text_config.full_attention_interval", 4); \
|
||||
LOAD_ARG_OR(attn_output_gate, "text_config.attn_output_gate", false); \
|
||||
LOAD_ARG_OR( \
|
||||
num_nextn_predict_layers, "text_config.mtp_num_hidden_layers", 0); \
|
||||
LOAD_ARG_OR(attention_bias, "text_config.attention_bias", false); \
|
||||
LOAD_ARG_OR(attention_dropout, "text_config.attention_dropout", 0.0f); \
|
||||
LOAD_ARG_OR(initializer_range, "text_config.initializer_range", 0.02f); \
|
||||
LOAD_ARG_OR( \
|
||||
mlp_only_layers, "text_config.mlp_only_layers", std::vector<int32_t>()); \
|
||||
LOAD_ARG(rope_scaling_mrope_section, \
|
||||
"text_config.rope_parameters.mrope_section"); \
|
||||
LOAD_ARG_OR(rope_scaling_rope_type, \
|
||||
"text_config.rope_parameters.rope_type", \
|
||||
"default"); \
|
||||
LOAD_ARG_OR(partial_rotary_factor, \
|
||||
"text_config.rope_parameters.partial_rotary_factor", \
|
||||
0.25f)
|
||||
|
||||
#define LOAD_QWEN3_5_VISION_ARGS() \
|
||||
LOAD_ARG_OR(image_token_id, "image_token_id", 248056); \
|
||||
LOAD_ARG_OR(video_token_id, "video_token_id", 248057); \
|
||||
LOAD_ARG_OR(vision_start_token_id, "vision_start_token_id", 248053); \
|
||||
LOAD_ARG_OR(vision_end_token_id, "vision_end_token_id", 248054); \
|
||||
LOAD_ARG(mm_deepstack_visual_indexes, \
|
||||
"vision_config.deepstack_visual_indexes"); \
|
||||
LOAD_ARG_OR(mm_num_hidden_layers, "vision_config.depth", 27); \
|
||||
LOAD_ARG_OR(mm_hidden_act, "vision_config.hidden_act", "gelu_pytorch_tanh"); \
|
||||
LOAD_ARG_OR(mm_hidden_size, "vision_config.hidden_size", 1152); \
|
||||
LOAD_ARG_OR(mm_num_channels, "vision_config.in_channels", 3); \
|
||||
LOAD_ARG_OR(mm_initializer_range, "vision_config.initializer_range", 0.02f); \
|
||||
LOAD_ARG_OR(mm_intermediate_size, "vision_config.intermediate_size", 4304); \
|
||||
LOAD_ARG_OR(mm_num_attention_heads, "vision_config.num_heads", 16); \
|
||||
LOAD_ARG_OR(mm_num_position_embeddings, \
|
||||
"vision_config.num_position_embeddings", \
|
||||
2304); \
|
||||
LOAD_ARG_OR(mm_projection_dim, "vision_config.out_hidden_size", 5120); \
|
||||
LOAD_ARG_OR(mm_patch_size, "vision_config.patch_size", 16); \
|
||||
LOAD_ARG_OR(mm_spatial_merge_size, "vision_config.spatial_merge_size", 2); \
|
||||
LOAD_ARG_OR(mm_temporal_patch_size, "vision_config.temporal_patch_size", 2); \
|
||||
LOAD_ARG_OR_FUNC(mm_head_dim, "head_dim", [&] { \
|
||||
return args->mm_hidden_size() / args->mm_num_attention_heads(); \
|
||||
}); \
|
||||
LOAD_ARG_OR( \
|
||||
rope_scaling_rope_type, "vision_config.rope_scaling.type", "mrope")
|
||||
|
||||
REGISTER_INPUT_PROCESSOR(qwen3_5, Qwen2_5_VLInputProcessor);
|
||||
REGISTER_CAUSAL_VLM_MODEL(qwen3_5, Qwen3_5ForConditionalGeneration);
|
||||
REGISTER_IMAGE_PROCESSOR(qwen3_5, Qwen2VLImageProcessor);
|
||||
REGISTER_MODEL_ARGS(qwen3_5, [&] {
|
||||
LOAD_QWEN3_5_COMMON_ARGS();
|
||||
LOAD_QWEN3_5_VISION_ARGS();
|
||||
SET_ARG(stop_token_ids, std::unordered_set<int32_t>({args->eos_token_id()}));
|
||||
});
|
||||
|
||||
REGISTER_INPUT_PROCESSOR(qwen3_5_moe, Qwen2_5_VLInputProcessor);
|
||||
REGISTER_CAUSAL_VLM_MODEL(qwen3_5_moe, Qwen3_5ForConditionalGeneration);
|
||||
REGISTER_IMAGE_PROCESSOR(qwen3_5_moe, Qwen2VLImageProcessor);
|
||||
REGISTER_MODEL_ARGS(qwen3_5_moe, [&] {
|
||||
LOAD_QWEN3_5_COMMON_ARGS();
|
||||
LOAD_QWEN3_5_VISION_ARGS();
|
||||
|
||||
LOAD_ARG_OR(decoder_sparse_step, "text_config.decoder_sparse_step", 1);
|
||||
LOAD_ARG_OR(moe_intermediate_size, "text_config.moe_intermediate_size", 512);
|
||||
LOAD_ARG_OR(num_experts, "text_config.num_experts", 512);
|
||||
LOAD_ARG_OR(num_experts_per_tok, "text_config.num_experts_per_tok", 10);
|
||||
LOAD_ARG_OR(shared_expert_intermediate_size,
|
||||
"text_config.shared_expert_intermediate_size",
|
||||
512);
|
||||
LOAD_ARG_OR(norm_topk_prob, "text_config.norm_topk_prob", true);
|
||||
|
||||
LOAD_ARG_OR(
|
||||
n_routed_experts, "text_config.n_routed_experts", args->num_experts());
|
||||
SET_ARG(n_shared_experts,
|
||||
args->shared_expert_intermediate_size() > 0 ? 1 : 0);
|
||||
SET_ARG(scoring_func, "softmax");
|
||||
SET_ARG(topk_method, "");
|
||||
SET_ARG(n_group, -1);
|
||||
SET_ARG(topk_group, 0);
|
||||
SET_ARG(routed_scaling_factor, 1.0f);
|
||||
|
||||
SET_ARG(stop_token_ids, std::unordered_set<int32_t>({args->eos_token_id()}));
|
||||
});
|
||||
|
||||
} // namespace xllm
|
||||
Reference in New Issue
Block a user