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

[fix] baseline4 docker build move ex_engine into qwen3_6_scripts, remove COPY ex_engine from Dockerfile

Browse Source
This commit is contained in:
root
2026-08-17 11:51:01 +00:00
parent c655c1d29e
commit 1af45de371
255 changed files with 52015 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

231
qwen3_6_scripts/ex_engine/xllm_models/llm/qwen3_5.h Normal file
View File

@@ -0,0 +1,231 @@
/* Copyright 2025-2026 The xLLM Authors.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
https://github.com/jd-opensource/xllm/blob/main/LICENSE
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#pragma once
#include <cstdint>
#include <memory>
#include <string>
#include <unordered_set>
#include <utility>
#include <vector>
#include "models/model_registry.h"
#if defined(USE_NPU) || defined(USE_MLU) || defined(USE_MUSA) || \
defined(USE_DCU)
#include "core/layers/qwen3_5_decoder_layer.h"
#include "qwen3_next.h"
#endif
namespace xllm {
#if defined(USE_NPU) || defined(USE_MLU) || defined(USE_MUSA) || \
defined(USE_DCU)
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::Tensor get_input_embeddings(torch::Tensor input_ids) {
return get_word_embedding()(input_ids);
}
void load_model(std::unique_ptr<ModelLoader> loader) {
Qwen3NextForCausalLMImpl::load_model(
std::move(loader), "model.language_model.", "lm_head.");
}
void load_model(std::unique_ptr<ModelLoader> loader,
const std::string& model_prefix) {
Qwen3NextForCausalLMImpl::load_model(
std::move(loader), model_prefix, "lm_head.");
}
};
TORCH_MODULE(Qwen3_5ForCausalLM);
#endif
#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(), 248046})); \
LOAD_ARG_TEXT_OR_ROOT(mamba_ssm_dtype, "mamba_ssm_dtype", "float32")
#define LOAD_QWEN3_5_TEXT_TYPE_AND_DTYPE(default_model_type) \
SET_ARG(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_MODEL_BACKEND(qwen3_5_text, "llm");
#if defined(USE_NPU) || defined(USE_MLU) || defined(USE_MUSA) || \
defined(USE_DCU)
REGISTER_CAUSAL_MODEL(qwen3_5_text, Qwen3_5ForCausalLM);
#endif
REGISTER_MODEL_ARGS(qwen3_5_text, [&] {
LOAD_QWEN3_5_TEXT_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_MODEL_BACKEND(qwen3_5_moe_text, "llm");
#if defined(USE_NPU) || defined(USE_MLU) || defined(USE_MUSA) || \
defined(USE_DCU)
REGISTER_CAUSAL_MODEL(qwen3_5_moe_text, Qwen3_5ForCausalLM);
#endif
REGISTER_MODEL_ARGS(qwen3_5_moe_text, [&] {
LOAD_QWEN3_5_TEXT_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_TEXT_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

59
qwen3_6_scripts/ex_engine/xllm_models/llm/qwen3_5_mtp.h Normal file
View File

@@ -0,0 +1,59 @@
/* Copyright 2025-2026 The xLLM Authors.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
https://github.com/jd-opensource/xllm/blob/main/LICENSE
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#pragma once
#include <memory>
#include "models/llm/qwen3_5.h"
#include "models/llm/qwen3_5_mtp_base.h"
#include "models/model_registry.h"
namespace xllm {
class Qwen3_5MtpModelImpl final : public Qwen3_5MtpModelImplBase {
public:
explicit Qwen3_5MtpModelImpl(const ModelContext& context)
: Qwen3_5MtpModelImplBase(context) {}
};
class Qwen3_5MtpForCausalLMImpl final : public Qwen3_5MtpForCausalLMImplBase {
public:
explicit Qwen3_5MtpForCausalLMImpl(const ModelContext& context)
: Qwen3_5MtpForCausalLMImplBase(
context,
std::make_shared<Qwen3_5MtpModelImpl>(context)) {}
};
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 qwen3_5_mtp::load_model_args(
json, args, "qwen3_5_text", "qwen3_5_mtp");
});
REGISTER_MODEL_ARGS_LOADER(qwen3_5_moe_mtp,
[](const JsonReader& json, ModelArgs* args) {
return qwen3_5_mtp::load_model_args(
json,
args,
"qwen3_5_moe_text",
"qwen3_5_moe_mtp");
});
} // namespace xllm

299
qwen3_6_scripts/ex_engine/xllm_models/llm/qwen3_5_mtp_base.h Normal file
View File

@@ -0,0 +1,299 @@
/* 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 <utility>
#include <vector>
#include "core/layers/common/linear.h"
#include "core/layers/qwen3_5_decoder_layer.h"
#include "models/llm/qwen3_next_hybrid_base.h"
#include "models/model_registry.h"
namespace xllm {
namespace qwen3_5_mtp {
inline StateDict get_lm_head_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 std::string& prefix : kLmHeadPrefixes) {
StateDict 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({}, "");
}
inline bool load_model_args(const JsonReader& json,
ModelArgs* args,
const std::string& base_type,
const std::string& mtp_type) {
ModelArgsLoader base_loader = ModelRegistry::get_model_args_loader(base_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_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 qwen3_5_mtp
class Qwen3_5MtpModelImplBase : public Qwen3HybridModelImplBase {
public:
explicit Qwen3_5MtpModelImplBase(const ModelContext& context)
: Qwen3HybridModelImplBase(context) {
const torch::TensorOptions& 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_);
}
layer::AttentionMetadata attn_metadata =
layer::AttentionMetadataBuilder::build(
input_params,
model_args_.enable_mla(),
build_attention_mask(input_params),
/*device=*/device_);
prepare_mrope(positions, attn_metadata);
torch::Tensor embedding = embed_tokens_(tokens);
torch::Tensor hidden = input_params.embedding.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());
torch::Tensor mrope_cos_sin;
for (const layer::Qwen3HybridDecoderLayerModulePtr& 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) {
if (!input_params.synchronize_layer(static_cast<uint32_t>(i))) {
return ModelOutput();
}
mtp_hidden = layers_[i]->forward(mtp_hidden,
residual,
positions,
attn_metadata,
kv_caches[i],
input_params,
mrope_cos_sin);
#if defined(USE_NPU)
if (input_params.parallel.layer_synchronizer != nullptr &&
!input_params.parallel.layer_synchronizer->record_event(
static_cast<int64_t>(i), device_.index())) {
return ModelOutput();
}
#endif
}
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) {
StateDict 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) +
".");
}
}
protected:
virtual void prepare_mrope(const torch::Tensor& positions,
layer::AttentionMetadata& attn_metadata) const {
UNUSED_PARAMETER(positions);
UNUSED_PARAMETER(attn_metadata);
}
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_5MtpForCausalLMImplBase : public Qwen3HybridForCausalLMImplBase {
public:
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 std::unique_ptr<StateDict>& state_dict :
loader->get_state_dicts()) {
StateDict shared_embedding_state_dict =
state_dict->get_dict_with_prefix(kEmbeddingPrefixes);
StateDict 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 {
StateDict lm_head_state_dict =
qwen3_5_mtp::get_lm_head_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.");
}
protected:
Qwen3_5MtpForCausalLMImplBase(
const ModelContext& context,
std::shared_ptr<Qwen3_5MtpModelImplBase> mtp_model)
: Qwen3HybridForCausalLMImplBase(context),
mtp_model_(std::move(mtp_model)) {
set_model_module(mtp_model_);
}
private:
std::shared_ptr<Qwen3_5MtpModelImplBase> mtp_model_;
};
} // namespace xllm

126
qwen3_6_scripts/ex_engine/xllm_models/llm/qwen3_next.h Normal file
View File

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

364
qwen3_6_scripts/ex_engine/xllm_models/llm/qwen3_next_hybrid_base.h Normal file
View File

@@ -0,0 +1,364 @@
/* Copyright 2025-2026 The xLLM Authors.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
https://github.com/jd-opensource/xllm/blob/main/LICENSE
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#pragma once
#include <torch/torch.h>
#include <algorithm>
#include <memory>
#include <string>
#include <vector>
#include "core/common/flash_comm1_context.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/framework/parallel_state/parallel_args.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"
#if defined(USE_NPU)
#include "core/layers/npu_torch/qwen3_next_hybrid_decoder_layer_base.h"
#elif defined(USE_MLU)
#include "core/layers/mlu/qwen3_5/qwen3_5_hybrid_decoder_layer_base.h"
#endif
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()),
parallel_args_(context.get_parallel_args()),
flash_comm1_options_(context.get_flash_comm1_options()) {
if (model_args_.n_routed_experts() > 0) {
flash_comm1_options_.enable_flashcomm1 = false;
flash_comm1_options_.enable_mmrs_fusion = false;
}
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));
attn_mask_ = layer::AttentionMask(options.device(),
options.dtype().toScalarType(),
/*mask_value=*/-9984);
dense_attn_mask_ = layer::AttentionMask(options.device(),
options.dtype().toScalarType(),
/*mask_value=*/1);
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),
/*device=*/device_);
const int32_t num_tokens = static_cast<int32_t>(tokens.size(0));
const auto& batch_forward_type = input_params.meta.batch_forward_type;
const bool is_prefill_side = batch_forward_type.no_decode();
FlashComm1Context fc1_ctx = build_flash_comm1_context(
num_tokens, is_prefill_side, parallel_args_, flash_comm1_options_);
FlashComm1ContextScope fc1_scope(&fc1_ctx);
torch::Tensor h;
if (input_params.embedding.input_embedding.defined()) {
h = input_params.embedding.input_embedding;
} else {
h = embed_tokens_(tokens);
}
if (is_sequence_sharded(fc1_ctx)) {
h = shard_sequence(h, fc1_ctx);
}
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);
#if defined(USE_NPU)
if (input_params.parallel.layer_synchronizer != nullptr &&
!input_params.parallel.layer_synchronizer->record_event(
static_cast<int64_t>(i), device_.index())) {
return ModelOutput();
}
#endif
}
auto [hidden_states, residual_out] = norm_->forward(h, residual);
h = hidden_states;
if (is_sequence_sharded(fc1_ctx)) {
h = gather_sequence(h, fc1_ctx);
}
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) {
#if defined(USE_NPU)
// On NPU the hybrid path never consumes attn_metadata.attn_mask: full
// attention runs through the fused-infer / paged-attention kernels (which
// carry their own fixed fia_attn_mask or need no mask at all) and linear
// attention is mask-free by construction. Materializing a dense
// [seq_len, seq_len] mask here is pure waste and, for long sequences,
// triggers an NPU OOM. Hand the kernels an empty mask unless a graph buffer
// already supplies one.
if (input_params.graph.attn_mask.defined()) {
return input_params.graph.attn_mask;
}
return torch::Tensor();
#else
if (input_params.graph.attn_mask.defined()) {
return input_params.graph.attn_mask;
}
max_seq_len_ = std::max(input_params.meta.kv_max_seq_len, max_seq_len_);
const bool use_append_mask =
input_params.is_spec_verify ||
input_params.meta.batch_forward_type.is_mixed() ||
input_params.meta.batch_forward_type.is_chunked_prefill();
if (!use_append_mask) {
return dense_attn_mask_.get_attn_mask(max_seq_len_, dtype_, device_);
}
const int32_t num_sequences = input_params.meta.num_sequences;
if (num_sequences <= 0) {
return dense_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.attention.host.q_seq_lens[j],
input_params.attention.host.kv_seq_lens[j],
max_seq_len_,
dtype_,
device_));
}
return torch::cat(req_mask_vec, 0);
#endif
}
ModelArgs model_args_;
torch::nn::ModuleList blocks_{nullptr};
std::vector<layer::Qwen3HybridDecoderLayerModulePtr> layers_;
int32_t max_seq_len_ = 0;
int32_t dp_size_ = 1;
ParallelArgs parallel_args_;
FlashComm1Options flash_comm1_options_;
torch::Device device_;
torch::ScalarType dtype_ = torch::kFloat;
layer::Qwen3NextRMSNorm norm_{nullptr};
layer::AttentionMask attn_mask_;
layer::AttentionMask dense_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) {
load_model(std::move(loader), "model.", "lm_head.");
}
void load_model(std::unique_ptr<ModelLoader> loader,
const std::string& model_prefix) {
load_model(std::move(loader), model_prefix, "lm_head.");
}
void load_model(std::unique_ptr<ModelLoader> loader,
const std::string& model_prefix,
const std::string& lm_head_prefix) {
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_prefix);
model_->load_state_dict(model_state_dict);
auto lm_head_state_dict =
state_dict->get_dict_with_prefix(lm_head_prefix);
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_prefix);
}
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; }
bool is_hybrid_linear_attention() { return true; }
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

440
qwen3_6_scripts/ex_engine/xllm_models/vlm/qwen3_5.h Normal file
View File

@@ -0,0 +1,440 @@
/* Copyright 2025-2026 The xLLM Authors.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
https://github.com/jd-opensource/xllm/blob/main/LICENSE
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#pragma once
#include "core/framework/model/model_output.h"
#include "core/layers/common/lm_head.h"
#include "core/layers/common/rotary_embedding_util.h"
#include "models/model_registry.h"
#include "models/vlm/mposition/mposition.h"
#include "models/vlm/qwen3_vl_base.h"
#include "processors/multimodal_processor.h"
#include "processors/qwen2_vl_image_processor.h"
#include "processors/qwen3_vl_prompt_processor.h"
#include "processors/qwen3_vl_video_processor.h"
#if defined(USE_NPU)
#include "models/llm/qwen3_5.h"
#include "models/vlm/npu/qwen3_vl.h"
#elif defined(USE_MLU) || defined(USE_DCU)
#include "core/layers/common/qwen3_next_rms_norm.h"
#include "core/layers/common/rms_norm.h"
#include "core/layers/qwen3_5_decoder_layer.h"
#include "core/layers/qwen3_vision_layer.h"
#include "models/llm/llm_model_base.h"
#include "qwen3_vl.h"
#endif
namespace xllm {
#if !defined(USE_NPU)
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 {
return layer::rotary::apply_mrope(cos_sin_, positions, mrope_section_);
}
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.parallel.dp_global_token_nums;
std::replace(dp_token_nums.begin(), dp_token_nums.end(), 0, 1);
}
auto inputs_embeds = input_params.embedding.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);
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,
/*enable_mla=*/false,
/*attn_mask=*/{},
h.device());
// Init batch and token_block_offset for GDN attention
if (attn_metadata.is_prefill || attn_metadata.is_chunked_prefill) {
constexpr int32_t kBlockM = 64;
constexpr int64_t pad_slot_id = -1;
constexpr int64_t default_max_num_programs = 1024;
constexpr int64_t chunk_size = 64;
auto seqlens = attn_metadata.q_cu_seq_lens.diff();
auto nums = (seqlens + kBlockM - 1) / kBlockM;
nums = nums.to(torch::kLong);
int32_t tot = nums.sum().item<int32_t>();
torch::Tensor range_batch = torch::arange(nums.size(0), nums.options());
torch::Tensor mlist_tensor = torch::repeat_interleave(range_batch, nums);
int64_t mlist_len = mlist_tensor.size(0);
int64_t max_num_programs =
std::max(default_max_num_programs, mlist_len) * 2;
torch::Tensor batch_ptr =
torch::full({max_num_programs},
pad_slot_id,
torch::dtype(torch::kInt32).device(seqlens.device()));
torch::Tensor token_block_offset_ptr =
torch::full({max_num_programs},
pad_slot_id,
torch::dtype(torch::kInt32).device(seqlens.device()));
std::vector<torch::Tensor> vec;
vec.reserve(nums.size(0));
for (int64_t i = 0; i < nums.size(0); ++i) {
vec.emplace_back(
torch::arange(nums[i].item<int64_t>(), nums.options()));
}
torch::Tensor offsetlist_tensor = torch::cat(vec, -1).to(torch::kInt32);
batch_ptr.narrow(0, 0, mlist_len).copy_(mlist_tensor);
token_block_offset_ptr.narrow(0, 0, mlist_len).copy_(offsetlist_tensor);
// Compute chunk indices for the chunked GDN kernel
{
torch::Tensor lengths = seqlens;
torch::Tensor num_chunks = (lengths + chunk_size - 1) / chunk_size;
num_chunks = num_chunks.to(torch::kLong);
torch::Tensor cumsum = torch::cumsum(num_chunks, 0);
int64_t total_chunks = cumsum[-1].item<int64_t>();
torch::Tensor arange_total =
torch::arange(total_chunks, attn_metadata.q_cu_seq_lens.options());
torch::Tensor zeros = torch::zeros({1}, cumsum.options());
torch::Tensor prefix = torch::cat(
{zeros, cumsum.slice(/*dim=*/0, /*start=*/0, /*end=*/-1)});
torch::Tensor repeats_prefix =
torch::repeat_interleave(prefix, num_chunks);
torch::Tensor indices = arange_total - repeats_prefix;
torch::Tensor mask = indices == 0;
torch::Tensor col0 = mask.cumsum(0) - 1;
attn_metadata.chunk_indices = torch::stack({col0, indices}, /*dim=*/1)
.to(attn_metadata.q_cu_seq_lens)
.to(torch::kInt32);
}
attn_metadata.tot = tot;
attn_metadata.batch = batch_ptr;
attn_metadata.token_block_offset = token_block_offset_ptr;
}
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);
#endif // !defined(USE_NPU)
#if defined(USE_NPU)
using Qwen3_5_VisionTransformer = npu::model::Qwen3_VisionTransformer;
#else
using Qwen3_5_VisionTransformer = Qwen3_VisionTransformer;
#endif
using Qwen3_5ForConditionalGenerationImpl =
Qwen3VLForConditionalGenerationBase<Qwen3_5_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(bos_token_id, "text_config.bos_token_id", 151643); \
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", \
static_cast<int32_t>(args->n_heads() * 2)); \
LOAD_ARG_OR( \
full_attention_interval, "text_config.full_attention_interval", 4); \
LOAD_ARG_OR(attn_output_gate, "text_config.attn_output_gate", true); \
LOAD_ARG_OR( \
num_nextn_predict_layers, "text_config.mtp_num_hidden_layers", 0); \
LOAD_ARG_OR(num_nextn_predict_layers, \
"text_config.num_nextn_predict_layers", \
args->num_nextn_predict_layers()); \
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_OR(rope_scaling_mrope_section, \
"text_config.rope_parameters.mrope_section", \
std::vector<int64_t>({11, 11, 10})); \
LOAD_ARG_OR(rope_scaling_mrope_interleaved, \
"text_config.rope_parameters.mrope_interleaved", \
true); \
LOAD_ARG_OR(rope_scaling_rope_type, \
"text_config.rope_parameters.rope_type", \
"default"); \
if (args->rope_scaling_rope_type() == "default") { \
args->rope_scaling_rope_type() = "mrope"; \
} \
LOAD_ARG_OR(partial_rotary_factor, \
"text_config.rope_parameters.partial_rotary_factor", \
0.25f); \
LOAD_ARG_OR(mamba_ssm_dtype, "text_config.mamba_ssm_dtype", "float32")
#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_OR(mm_deepstack_visual_indexes, \
"vision_config.deepstack_visual_indexes", \
std::vector<int64_t>()); \
if (!args->mm_deepstack_visual_indexes().empty()) { \
LOG(FATAL) << "qwen3_5 VLM does not support 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", \
args->hidden_size()); \
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(); \
})
// qwen3_5/qwen3_5_moe are multimodal entry points. On NPU, text-only serving
// uses qwen3_5_text/qwen3_5_moe_text from llm/qwen3_5.h because the VLM
// request protocol currently requires array-form chat content.
REGISTER_CAUSAL_VLM_MODEL(qwen3_5, Qwen3_5ForConditionalGeneration);
REGISTER_MPOSITION_GENERATOR(qwen3_5, Qwen3VLMPositionGenerator);
using Qwen35MultimodalProcessor = MultimodalProcessor<Qwen3VLPromptProcessor,
Qwen2VLImageProcessor,
Qwen3VLVideoProcessor>;
REGISTER_MULTIMODAL_PROCESSOR(qwen3_5, Qwen35MultimodalProcessor);
REGISTER_MODEL_ARGS(qwen3_5, [&] {
LOAD_QWEN3_5_COMMON_ARGS();
LOAD_QWEN3_5_VISION_ARGS();
SET_ARG(num_experts, 0);
SET_ARG(n_routed_experts, 0);
SET_ARG(n_shared_experts, 0);
SET_ARG(stop_token_ids,
std::unordered_set<int32_t>({args->eos_token_id(), 248046}));
});
REGISTER_CAUSAL_VLM_MODEL(qwen3_5_moe, Qwen3_5ForConditionalGeneration);
REGISTER_MPOSITION_GENERATOR(qwen3_5_moe, Qwen3VLMPositionGenerator);
REGISTER_MULTIMODAL_PROCESSOR(qwen3_5_moe, Qwen35MultimodalProcessor);
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(), 248046}));
});
// Text-only model registrations. On NPU these are handled by llm/qwen3_5.h.
#if !defined(USE_NPU)
// qwen3_5 without vision config (text-only serving).
// Model args are already registered by the VLM registration above.
REGISTER_CAUSAL_MODEL_WITH_VARNAME(qwen3_5_lm, qwen3_5, Qwen3_5ForCausalLM);
REGISTER_CAUSAL_MODEL_WITH_VARNAME(qwen3_5_moe_lm,
qwen3_5_moe,
Qwen3_5ForCausalLM);
REGISTER_CAUSAL_MODEL(qwen3_5_text, Qwen3_5ForCausalLM);
REGISTER_MODEL_ARGS(qwen3_5_text, [&] {
LOAD_QWEN3_5_COMMON_ARGS();
SET_ARG(num_experts, 0);
SET_ARG(n_routed_experts, 0);
SET_ARG(n_shared_experts, 0);
SET_ARG(decoder_sparse_step, 1);
SET_ARG(stop_token_ids,
std::unordered_set<int32_t>({args->eos_token_id(), 248046}));
});
REGISTER_CAUSAL_MODEL(qwen3_5_moe_text, Qwen3_5ForCausalLM);
REGISTER_MODEL_ARGS(qwen3_5_moe_text, [&] {
LOAD_QWEN3_5_COMMON_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(), 248046}));
});
#endif // !defined(USE_NPU)
#undef LOAD_QWEN3_5_VISION_ARGS
#undef LOAD_QWEN3_5_COMMON_ARGS
} // namespace xllm