f28223c9da
Replaces all Python PyTorch fallback attention with native ixformer kernels: Decode path: - ≤32K: paged_attention_v1 (5D KV layout, x=8) — verified on real BI-V100 - >32K: paged_attention_v2 (5D→4D permute) — verified 65K+ on real BI-V100 - Removes _forward_decode_pytorch Python fallback entirely Prefill path (profiling): - _run_sdpa_fallback now uses ixformer.flash_attn_varlen_func - head_dim=256 verified correct (diff<0.004) and 1.7x faster than PyTorch - Falls back to Q-tiling pure-math if ixformer unavailable Also includes: MoE kernel integration, GDN C++ kernels, diagnostic scripts, xllm upstream layer/kernel references, .dockerignore cleanup. All changes verified on real BI-V100 hardware (single card).
232 lines
13 KiB
C++
232 lines
13 KiB
C++
/* 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
|