Files
project_6/upstream_ref/xllm_latest/models/vlm/qwen3_5.h
project6-dev 87a19d2d00 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
2026-08-11 02:32:06 +00:00

313 lines
14 KiB
C++

/* 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