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

upstream(xllm): sync to jd-opensource/xllm latest + revert serving_chat.py

Browse Source 
搬运 jd-opensource/xllm 最新代码到 upstream_ref/xllm_latest/:
- core/kernels/ilu/ 10 files (ixformer.h API 不变)
- core/layers/ilu/ 4 files (fused_moe.cpp config 访问从 FLAGS→singleton)
- core/layers/npu_torch/ 14 files (qwen3_gated_delta_net_base.cpp 576→1164行,
  新增 repeat_tensor_heads, checkpoint_stride, spec_verify 等 GDN 功能)
- models/llm/ 5 files (qwen3_5.h 模型注册重构, 新增 qwen3_5_mtp_base.h)
- models/vlm/ 1 file (qwen3_5.h 218→440行)

serving_chat.py: 还原到 8030a11b 原版,删掉 6dcf3590 的语法错误 min(8192,
(缺右括号导致 py_compile 失败)
This commit is contained in:
project6-dev
2026-08-12 04:22:34 +00:00
parent 6dcf3590d5
commit d025b08a95
35 changed files with 1595 additions and 593 deletions
Download Patch File Download Diff File Expand all files Collapse all files

65
upstream_ref/xllm_latest/models/llm/qwen3_5.h
View File

@@ -1,4 +1,4 @@
/* Copyright 2025 The xLLM Authors. All Rights Reserved.
/* 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.
@@ -16,16 +16,23 @@ limitations under the License.
#pragma once
#include <cstdint>
#include <memory>
#include <string>
#include <unordered_set>
#include <utility>
#include <vector>
#include "core/layers/npu_torch/qwen3_5_decoder_layer_impl.h"
#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)
@@ -45,8 +52,24 @@ class Qwen3_5ForCausalLMImpl : public Qwen3NextForCausalLMImpl {
: 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); \
@@ -163,53 +186,43 @@ TORCH_MODULE(Qwen3_5ForCausalLM);
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()})); \
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_TYPE_AND_DTYPE(default_model_type) \
LOAD_ARG_OR(model_type, "model_type", default_model_type); \
#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_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_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_TYPE_AND_DTYPE("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_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_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_TYPE_AND_DTYPE("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_TYPE_AND_DTYPE
#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

253
upstream_ref/xllm_latest/models/llm/qwen3_5_mtp.h
View File

@@ -1,4 +1,4 @@
/* Copyright 2026 The xLLM Authors. All Rights Reserved.
/* 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.
@@ -15,250 +15,26 @@ 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/llm/qwen3_5.h"
#include "models/llm/qwen3_5_mtp_base.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 {
class Qwen3_5MtpModelImpl final : public Qwen3_5MtpModelImplBase {
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;
: Qwen3_5MtpModelImplBase(context) {}
};
class Qwen3_5MtpForCausalLMImpl : public Qwen3HybridForCausalLMImplBase {
class Qwen3_5MtpForCausalLMImpl final : public Qwen3_5MtpForCausalLMImplBase {
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_;
: Qwen3_5MtpForCausalLMImplBase(
context,
std::make_shared<Qwen3_5MtpModelImpl>(context)) {}
};
TORCH_MODULE(Qwen3_5MtpForCausalLM);
@@ -267,14 +43,17 @@ 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");
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 load_qwen3_5_mtp_model_args(
json, args, "qwen3_5_moe", "qwen3_5_moe_mtp");
return qwen3_5_mtp::load_model_args(
json,
args,
"qwen3_5_moe_text",
"qwen3_5_moe_mtp");
});
} // namespace xllm

299
upstream_ref/xllm_latest/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

2
upstream_ref/xllm_latest/models/llm/qwen3_next.h
View File

@@ -1,4 +1,4 @@
/* Copyright 2026 The xLLM Authors. All Rights Reserved.
/* 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.

151
upstream_ref/xllm_latest/models/llm/qwen3_next_hybrid_base.h
View File

@@ -1,4 +1,4 @@
/* Copyright 2026 The xLLM Authors. All Rights Reserved.
/* 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.
@@ -22,18 +22,23 @@ limitations under the License.
#include <string>
#include <vector>
#include "core/common/global_flags.h"
#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 {
@@ -55,7 +60,14 @@ class Qwen3HybridModelImplBase : public Qwen3HybridModelModule {
public:
explicit Qwen3HybridModelImplBase(const ModelContext& context)
: device_(context.get_tensor_options().device()),
model_args_(context.get_model_args()) {
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();
@@ -69,10 +81,12 @@ class Qwen3HybridModelImplBase : public Qwen3HybridModelModule {
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);
/*mask_value=*/-9984);
dense_attn_mask_ = layer::AttentionMask(options.device(),
options.dtype().toScalarType(),
/*mask_value=*/1);
dp_size_ = parallel_args.dp_size();
}
@@ -95,8 +109,25 @@ class Qwen3HybridModelImplBase : public Qwen3HybridModelModule {
layer::AttentionMetadataBuilder::build(
input_params,
model_args_.enable_mla(),
build_attention_mask(input_params));
torch::Tensor h = embed_tokens_(tokens);
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_) {
@@ -114,9 +145,19 @@ class Qwen3HybridModelImplBase : public Qwen3HybridModelModule {
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);
}
@@ -155,27 +196,48 @@ class Qwen3HybridModelImplBase : public Qwen3HybridModelModule {
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_);
#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.num_sequences;
const int32_t num_sequences = input_params.meta.num_sequences;
if (num_sequences <= 0) {
return attn_mask_.get_attn_mask(max_seq_len_, dtype_, device_);
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.q_seq_lens_vec[j],
input_params.kv_seq_lens_vec[j],
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_;
@@ -183,10 +245,13 @@ class Qwen3HybridModelImplBase : public Qwen3HybridModelModule {
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};
};
@@ -232,53 +297,28 @@ class Qwen3HybridForCausalLMImplBase : public torch::nn::Module {
}
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();
};
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.");
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;
}
}
}
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.");
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;
}
}
}
}
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.");
@@ -288,8 +328,7 @@ class Qwen3HybridForCausalLMImplBase : public torch::nn::Module {
}
lm_head_->load_state_dict(lm_head_state_dict);
}
model_->verify_loaded_weights("model.");
model_->verify_loaded_weights(model_prefix);
}
virtual void prepare_expert_weight(int32_t layer_id,
@@ -298,6 +337,8 @@ class Qwen3HybridForCausalLMImplBase : public torch::nn::Module {
}
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; }