diff --git a/CMakeLists.txt b/CMakeLists.txt index 547d22df4d..9056b1ee90 100755 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -9,6 +9,7 @@ option(USE_MUSA "Enable MUSA support" OFF) option(USE_DCU "Enable DCU support" OFF) option(USE_MACA "Enable MACA support" OFF) option(ENABLE_HA "Enable Mooncake etcd-based high availability support" OFF) +option(USE_XLITE "Enable xlite backend (NPU only)" OFF) add_compile_definitions(YLT_ENABLE_IBV) add_definitions(-DYLT_ENABLE_IBV) set(YLT_ENABLE_IBV ON) @@ -477,6 +478,12 @@ if(USE_NPU) set(CMAKE_VERBOSE_MAKEFILE ON) add_definitions(-DTORCH_HIGHER_THAN_PTA6) + if (USE_XLITE) + add_definitions(-DUSE_XLITE) + include(${CMAKE_CURRENT_SOURCE_DIR}/cmake/xlite.cmake) + message(STATUS "USE_XLITE is ON") + endif() + # Use vcpkg header files as the first priority search directory, #-> because the scope of third-party software managed by vcpkg is used throughout the entire xllm. message(STATUS "VCPKG_INCLUDE_DIR = ${CMAKE_BINARY_DIR}/vcpkg_installed/${VCPKG_TARGET_TRIPLET}/include") diff --git a/cmake/xlite.cmake b/cmake/xlite.cmake new file mode 100644 index 0000000000..1b39ec4b64 --- /dev/null +++ b/cmake/xlite.cmake @@ -0,0 +1,23 @@ +# xlite link helper: find_package(xlite) when USE_NPU AND USE_XLITE are ON. + +function(xllm_link_xlite target) + if(NOT USE_NPU OR NOT USE_XLITE) + return() + endif() + + if(NOT TARGET xlite::xlite) + execute_process( + COMMAND ${Python_EXECUTABLE} -c "import xlite; print(xlite.cmake_prefix_path)" + OUTPUT_VARIABLE _XLITE_CMAKE_PREFIX + OUTPUT_STRIP_TRAILING_WHITESPACE + RESULT_VARIABLE _XLITE_IMPORT_RESULT) + if(NOT _XLITE_IMPORT_RESULT EQUAL 0 OR _XLITE_CMAKE_PREFIX STREQUAL "") + message(FATAL_ERROR "USE_XLITE is ON but xlite not found. Install xlite or pass -DUSE_XLITE=OFF.") + endif() + + find_package(xlite REQUIRED CONFIG PATHS "${_XLITE_CMAKE_PREFIX}" NO_DEFAULT_PATH) + message(STATUS "xlite::xlite found via find_package (${_XLITE_CMAKE_PREFIX})") + endif() + + target_link_libraries(${target} PRIVATE xlite::xlite) +endfunction() \ No newline at end of file diff --git a/xllm/CMakeLists.txt b/xllm/CMakeLists.txt index 433565b7b9..bf70ebab23 100644 --- a/xllm/CMakeLists.txt +++ b/xllm/CMakeLists.txt @@ -127,6 +127,9 @@ if(USE_MUSA) ) endif() +if (USE_NPU AND USE_XLITE) + xllm_link_xlite(xllm) +endif() # install xllm install(TARGETS xllm RUNTIME DESTINATION ${CMAKE_INSTALL_BINDIR}) diff --git a/xllm/core/framework/model/causal_lm.h b/xllm/core/framework/model/causal_lm.h index 974621d60c..bfb1975f88 100644 --- a/xllm/core/framework/model/causal_lm.h +++ b/xllm/core/framework/model/causal_lm.h @@ -49,6 +49,10 @@ namespace layer { struct AttentionMetadata; } +namespace xlite { +class XliteModelHolder; +} + struct ModelGraphMetadataState { virtual ~ModelGraphMetadataState() = default; }; @@ -151,6 +155,9 @@ class CausalLM : public torch::nn::Module { NOT_IMPLEMENTED(); return false; } + + // xlite runtime access, nullptr for non-xlite models. + virtual xlite::XliteModelHolder* get_xlite_holder() { return nullptr; } #endif virtual layer::LmHead get_lm_head() { @@ -444,6 +451,14 @@ class CausalLMImpl : public CausalLM { requested_rolling_slots, model_id); } + + // Forward to inner Model. + xlite::XliteModelHolder* get_xlite_holder() override { + if constexpr (detail::has_get_xlite_holder::value) { + return model_->get_xlite_holder(); + } + return CausalLM::get_xlite_holder(); + } #endif layer::LmHead get_lm_head() override { diff --git a/xllm/core/framework/model/model_traits.h b/xllm/core/framework/model/model_traits.h index 2329d89e3f..4d27ee2197 100644 --- a/xllm/core/framework/model/model_traits.h +++ b/xllm/core/framework/model/model_traits.h @@ -246,6 +246,15 @@ struct has_init_or_refresh_rolling_runtime< std::declval(), std::declval()))>> : std::true_type {}; +// SFINAE: xlite-backend models expose get_xlite_holder(). +template +struct has_get_xlite_holder : std::false_type {}; + +template +struct has_get_xlite_holder< + T, + std::void_t()->get_xlite_holder())>> + : std::true_type {}; #endif template diff --git a/xllm/core/layers/xlite/xlite_attn_meta_builder.h b/xllm/core/layers/xlite/xlite_attn_meta_builder.h new file mode 100644 index 0000000000..b76ebd47a4 --- /dev/null +++ b/xllm/core/layers/xlite/xlite_attn_meta_builder.h @@ -0,0 +1,84 @@ +/* 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. +==============================================================================*/ + +// Build xlite XModelAttnMeta from xllm ModelInputParams. + +#pragma once + +#include + +#include +#include +#include + +#include "core/framework/model/model_input_params.h" +#include "core/layers/xlite/xlite_init_utils.h" + +namespace xllm::xlite { + +class XliteAttnMetaBuilder { + public: + static void Build(const ModelInputParams& params, + const torch::Tensor& positions, + uint32_t block_size, + XModelAttnMeta& m, + int64_t pad_count = 0) { + // version=0: xlite recomputes position from cachedLens (framework's + // positions tensor is off-by-one on decode). + m.version = 0; + m.lens.clear(); + m.cachedLens.clear(); + m.blockTables.clear(); + + const auto& host = params.attention.host; + int n = params.meta.num_sequences; + uint32_t bs = block_size; + + // block_tables may be undefined in edge cases (DP empty shard). + const bool has_real_seqs = + n > 0 && host.block_tables.defined() && host.block_tables.dim() >= 2; + if (has_real_seqs) { + auto block_acc = host.block_tables.accessor(); + for (int s = 0; s < n; ++s) { + int32_t q_len = host.q_seq_lens[s]; + int32_t kv_len = host.kv_seq_lens[s]; + m.lens.push_back(static_cast(q_len)); + m.cachedLens.push_back( + static_cast(std::max(0, kv_len - q_len))); // clamp >= 0 + int32_t nblocks = + (kv_len + static_cast(bs) - 1) / static_cast(bs); + std::vector row(nblocks); + for (int32_t b = 0; b < nblocks; ++b) { + row[b] = static_cast(block_acc[s][b]); + } + m.blockTables.push_back(std::move(row)); + } + } + + // DP padding: append dummy seq so sum(lens) aligns across DP groups. + if (pad_count > 0) { + m.lens.push_back(static_cast(pad_count)); + m.cachedLens.push_back(0); + int32_t nblocks = + (static_cast(pad_count) + static_cast(bs) - 1) / + static_cast(bs); + std::vector row(nblocks, 0); + m.blockTables.push_back(std::move(row)); + } + InitXTensor(m.vllmPosition, positions); + } +}; + +} // namespace xllm::xlite \ No newline at end of file diff --git a/xllm/core/layers/xlite/xlite_causal_lm_base.h b/xllm/core/layers/xlite/xlite_causal_lm_base.h new file mode 100644 index 0000000000..8ea107a2d6 --- /dev/null +++ b/xllm/core/layers/xlite/xlite_causal_lm_base.h @@ -0,0 +1,227 @@ +/* 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 + +#include +#include +#include + +#if defined(USE_NPU) +#include +#endif + +#include + +#include "core/framework/kv_cache/kv_cache.h" +#include "core/framework/model/causal_lm.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/framework/state_dict/state_dict.h" +#include "core/layers/xlite/xlite_freqs_cis.h" +#include "core/layers/xlite/xlite_init_utils.h" +#include "core/layers/xlite/xlite_model_adapter.h" +#include "runtime/options.h" + +namespace xllm::xlite { + +// Aggregate view over multiple StateDicts (MoE checkpoint spans many shards). +class MergedStateDict : public StateDict { + public: + explicit MergedStateDict(std::vector dicts) + : StateDict({}, ""), dicts_(std::move(dicts)) {} + + torch::Tensor get_tensor(const std::string& tensor_name) const override { + for (const StateDict* d : dicts_) { + if (!d) { + continue; + } + torch::Tensor t = d->get_tensor(tensor_name); + if (t.defined()) { + return t; + } + } + return torch::Tensor{nullptr}; + } + + private: + std::vector dicts_; +}; + +class XliteModelHolder { + public: + virtual ~XliteModelHolder() = default; + virtual XRuntime& xlite_rt() = 0; + virtual XModel& xlite_model() = 0; + virtual const XModelConfig& xlite_config() const = 0; + virtual const torch::Tensor& xlite_freqs_cis() const = 0; + virtual torch::Tensor xlite_output_buf(int64_t num_tokens) = 0; + virtual bool xlite_ready() const = 0; +}; + +class XliteCausalLMBase : public CausalLM, public XliteModelHolder { + public: + XliteCausalLMBase(const ModelContext& context, + std::unique_ptr adapter) + : options_(context.get_tensor_options()), + device_(options_.device()), + args_(context.get_model_args()), + parallel_(context.get_parallel_args()), + adapter_(std::move(adapter)), + ready_(false) { + cfg_ = adapter_->BuildConfig(context); + // ParallelArgs::tp_size() is not populated; derive from world_size/dp_size. + const uint32_t tp = XliteTpSize(parallel_); + const uint32_t tp_rank = XliteTpRank(parallel_); + cfg_.defTpSize = tp; + // EP=1 -> experts sharded by attention TP; EP>1 -> world_size/ep_size. + const uint32_t ep = cfg_.moeEpSize; + cfg_.moeTPSize = (ep > 1) ? (parallel_.world_size() / ep) : tp; + uint32_t rank = static_cast(parallel_.rank()); + rt_ = + std::make_unique(static_cast(device_.index()), + /*sizeMB=*/0, + rank, + tp, + static_cast(parallel_.dp_size()), + cfg_.moeTPSize, + static_cast(parallel_.ep_size())); + model_ = std::make_unique(cfg_, rank); + // MLA freqs_cis needs YaRN params from ModelArgs (XModelConfig has no such + // fields). + YarnConfig yarn; + yarn.rope_factor = args_.rope_scaling_factor(); + yarn.beta_fast = static_cast(args_.rope_scaling_beta_fast()); + yarn.beta_slow = static_cast(args_.rope_scaling_beta_slow()); + yarn.original_seq_len = + args_.rope_scaling_original_max_position_embeddings(); + freqs_cis_ = XliteFreqsCis::Precompute(cfg_, options_, yarn); + output_buf_ = torch::empty( + {static_cast(cfg_.maxBatchedTokens), args_.hidden_size()}, + options_); + LOG(INFO) << "[xlite] rank=" << rank << " tp=" << tp + << " tp_rank=" << tp_rank << " moeTPSize=" << cfg_.moeTPSize + << " ep=" << parallel_.ep_size(); + } + + void load_model(std::unique_ptr loader) override { + auto& state_dicts = loader->get_state_dicts(); + // Load once over merged shards, then Init once (not idempotent). + std::vector dicts; + dicts.reserve(state_dicts.size()); + for (auto& sd_ptr : state_dicts) { + if (sd_ptr && sd_ptr->size() > 0) { + dicts.push_back(sd_ptr.get()); + } + } + CHECK(!dicts.empty()) << "xlite load_model: no non-empty state dicts"; + MergedStateDict merged(std::move(dicts)); + adapter_->Load( + merged, *model_, cfg_, args_, parallel_, device_, weight_storages_); + model_->Init(); + // GetTensorPoolSize requires maxBatchedTokens>0 (set in BuildConfig). + CHECK_GT(cfg_.maxBatchedTokens, 0) << "cfg_.maxBatchedTokens is 0"; + size_t pool = model_->GetTensorPoolSize(0); + CHECK_EQ(rt_->InitTensorPool(pool), 0) << "xlite InitTensorPool failed"; + ready_ = true; + } + + ModelOutput forward(const torch::Tensor&, + const torch::Tensor&, + std::vector&, + const ModelInputParams&) override { + LOG(FATAL) << "xlite forward() must not be called; use XliteExecutorImpl"; + return ModelOutput(); + } + + torch::Tensor logits(const torch::Tensor& hidden_states, + const torch::Tensor& seleted_idxes) override { + int64_t n = + seleted_idxes.defined() ? seleted_idxes.size(0) : hidden_states.size(0); + // head is vocab-sharded by tp; all_gather -> [tp, n, vocab/tp]. + uint32_t tp = XliteTpSize(parallel_); + torch::Tensor out = + torch::empty({static_cast(tp), + n, + args_.vocab_size() / static_cast(tp)}, + options_); + ::XTensor x_in, x_idx, x_out; + InitXTensor(x_in, hidden_states); + InitXTensor(x_idx, + seleted_idxes.defined() + ? seleted_idxes + : torch::arange(n, options_.dtype(torch::kInt32))); + InitXTensor(x_out, out); +#if defined(USE_NPU) + aclrtStream ext = c10_npu::getCurrentNPUStream(device_.index()).stream(); + rt_->EventWaitCurrStream(ext); + model_->ForwardGetLogits(*rt_, x_in, x_idx, x_out); + rt_->EventRecordCurrStream(ext); +#else + LOG(FATAL) << "xlite backend requires USE_NPU"; +#endif + // head is vocab-sharded by tp; all_gather -> [tp, n, vocab/tp]. + return out.permute({1, 0, 2}).reshape({n, args_.vocab_size()}); + } + torch::Tensor logits(const torch::Tensor& hidden_states, + const torch::Tensor& seleted_idxes, + torch::Tensor& out_hidden) override { + out_hidden = seleted_idxes.defined() + ? hidden_states.index_select(0, seleted_idxes) + : hidden_states; + return logits(hidden_states, seleted_idxes); + } + + torch::Device device() const override { return device_; } + const torch::TensorOptions& options() const override { return options_; } + void prepare_expert_weight(int32_t, const std::vector&) override {} + void update_expert_weight(int32_t) override {} + + xlite::XliteModelHolder* get_xlite_holder() override { return this; } + + XRuntime& xlite_rt() override { return *rt_; } + XModel& xlite_model() override { return *model_; } + const XModelConfig& xlite_config() const override { return cfg_; } + const torch::Tensor& xlite_freqs_cis() const override { return freqs_cis_; } + torch::Tensor xlite_output_buf(int64_t n) override { + CHECK(output_buf_.defined()) << "xlite output_buf not allocated"; + CHECK_LE(n, static_cast(cfg_.maxBatchedTokens)) + << "xlite output_buf overflow: n=" << n + << " > maxBatchedTokens=" << cfg_.maxBatchedTokens; + return output_buf_.slice(0, 0, n); + } + bool xlite_ready() const override { return ready_; } + + protected: + torch::TensorOptions options_; + torch::Device device_; + ModelArgs args_; + ParallelArgs parallel_; + XModelConfig cfg_{}; + std::unique_ptr rt_; + std::unique_ptr model_; + std::unique_ptr adapter_; + std::vector weight_storages_; + torch::Tensor freqs_cis_; + torch::Tensor output_buf_; + bool ready_; +}; + +} // namespace xllm::xlite \ No newline at end of file diff --git a/xllm/core/layers/xlite/xlite_config_builder.h b/xllm/core/layers/xlite/xlite_config_builder.h new file mode 100644 index 0000000000..83ef78dd09 --- /dev/null +++ b/xllm/core/layers/xlite/xlite_config_builder.h @@ -0,0 +1,262 @@ +/* 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 + +#include +#include + +#include "core/framework/config/kv_cache_config.h" +#include "core/framework/config/scheduler_config.h" +#include "core/framework/model/model_args.h" +#include "core/framework/model_context.h" +#include "core/framework/parallel_state/parallel_args.h" +#include "core/framework/quant_args.h" // QuantArgs (IsW8A8) +#include "core/layers/common/dsa_topk_share_plan.h" // DsaTopkSharePlan +#include "core/layers/xlite/xlite_init_utils.h" // XliteTpSize + +namespace xllm::xlite { + +class XliteConfigBuilder { + private: + // W8A8 quantized model (quant_method non-empty, from + // quant_model_description.json). Hardcoding true flips shared expert layout + // for BF16 models. + static bool IsW8A8(const ModelContext& context) { + const auto& q = context.get_quant_args(); + return !q.quant_method().empty(); + } + + public: + // Qwen3 dense (MHA); also the base for Qwen3-MoE (nDenseLayers overridden). + static XModelConfig FromQwen3(const ModelContext& context) { + const ModelArgs& a = context.get_model_args(); + const ParallelArgs& p = context.get_parallel_args(); + XModelConfig c{}; + + c.vocabSize = static_cast(a.vocab_size()); + c.hiddenSize = static_cast(a.hidden_size()); + c.nLayers = static_cast(a.n_layers()); + + c.nHeads = static_cast(a.n_heads()); + c.nKvHeads = a.n_kv_heads().has_value() + ? static_cast(a.n_kv_heads().value()) + : c.nHeads; // GQA -> MHA fallback + + // head_dim: explicit if set, else hidden_size / n_heads. + uint32_t head_dim = a.head_dim() > 0 ? static_cast(a.head_dim()) + : c.hiddenSize / c.nHeads; + c.headDim = head_dim; + c.ropeHeadDim = head_dim; + c.nopeHeadDim = 0; // MLA-only + c.vHeadDim = 0; // MLA-only + c.qLoraRank = 0; + c.kvLoraRank = 0; + + c.attnType = XMODEL_ATTN_MHA; + c.ropeType = XMODEL_ROPE_NEOX; + c.addBias = a.qkv_bias(); + c.qkNorm = a.use_qk_norm(); + + c.normEps = a.rms_norm_eps(); + c.ropeTheta = a.rope_theta(); + c.softmaxScale = 1.0f / std::sqrt(static_cast(head_dim)); + + // dense: all layers dense (no MoE). + c.nDenseLayers = c.nLayers; + c.nRoutedExperts = 0; + c.nSharedExperts = 0; + c.nActExperts = 0; + c.intermediateSize = static_cast(a.intermediate_size()); + c.moeIntermediateSize = 0; + c.scoringFunc = XMODEL_SCORING_FUNC_SOFTMAX; + c.normTopKProb = false; + + // defTpSize overridden by XliteCausalLMBase ctor; best-effort default here. + c.defTpSize = XliteTpSize(p); + c.defDpSize = static_cast(p.dp_size()); + c.moeEpSize = static_cast(p.ep_size()); + c.moeTPSize = 1; + + c.maxSeqLen = static_cast(a.max_position_embeddings()); + // Runtime fields from singletons (ModelContext has no runtime::Options). + c.maxBatchedTokens = static_cast( + SchedulerConfig::get_instance().max_tokens_per_batch()); + c.maxBatch = static_cast( + SchedulerConfig::get_instance().max_seqs_per_batch()); + c.blockSize = + static_cast(KVCacheConfig::get_instance().block_size()); + c.deepstackNumLevel = 0; + c.weightNZ = false; + return c; + } + + // Qwen3-MoE: overrides MoE fields on top of FromQwen3. + static XModelConfig FromQwen3Moe(const ModelContext& context) { + XModelConfig c = FromQwen3(context); + const ModelArgs& a = context.get_model_args(); + const ParallelArgs& p = context.get_parallel_args(); + + c.nDenseLayers = static_cast(a.first_k_dense_replace()); + c.nRoutedExperts = static_cast(a.num_experts()); + c.nSharedExperts = 0; + c.nActExperts = static_cast(a.num_experts_per_tok()); + c.moeIntermediateSize = static_cast(a.moe_intermediate_size()); + c.scoringFunc = XMODEL_SCORING_FUNC_SOFTMAX; + c.normTopKProb = a.norm_topk_prob(); + // HF weights are [out, in]; LoadMoEExperts transposes to [in, out]. + c.expertsWeightTrans = true; + + c.moeEpSize = static_cast(p.ep_size()); + c.moeTPSize = p.ep_size() > 1 + ? static_cast(p.world_size() / p.ep_size()) + : XliteTpSize(p); + return c; + } + + // DeepSeek-V3/R1 (MLA + sigmoid MoE + shared expert, no DSA). + static XModelConfig FromDeepseekV3(const ModelContext& context) { + XModelConfig c = FromQwen3(context); + const ModelArgs& a = context.get_model_args(); + const ParallelArgs& p = context.get_parallel_args(); + + // MLA dims (V3: nope=128/rope=64/v=128/q_lora=1536/kv_lora=512). + c.attnType = XMODEL_ATTN_MLA; + c.qLoraRank = static_cast(a.q_lora_rank()); + c.kvLoraRank = static_cast(a.kv_lora_rank()); + c.nopeHeadDim = static_cast(a.qk_nope_head_dim()); + c.ropeHeadDim = static_cast(a.qk_rope_head_dim()); + c.vHeadDim = static_cast(a.v_head_dim()); + // MLA MQA: nKvHeads=1 (kv_lora_rank shared across heads); matches xlite ref + // + framework KVCacheShape. + c.nKvHeads = 1; + // softmaxScale = 1/sqrt(nope+rope) + YaRN mscale (applied when maxSeqLen > + // origSeqLen). + float qkHeadDim = static_cast(c.nopeHeadDim + c.ropeHeadDim); + c.softmaxScale = 1.0f / std::sqrt(qkHeadDim); + int64_t origSeqLen = a.rope_scaling_original_max_position_embeddings(); + if (origSeqLen > 0 && static_cast(c.maxSeqLen) > origSeqLen) { + float mscale = + 0.1f * a.rope_scaling_mscale() * std::log(a.rope_scaling_factor()) + + 1.0f; + c.softmaxScale *= mscale * mscale; + } + + // MoE: sigmoid + shared + group-limited (V3: 3 dense / 256 expert / 8 act / + // 1 shared). + c.nDenseLayers = static_cast(a.first_k_dense_replace()); + c.nRoutedExperts = static_cast(a.num_experts()); + c.nSharedExperts = static_cast(a.n_shared_experts()); + c.nActExperts = static_cast(a.num_experts_per_tok()); + c.moeIntermediateSize = static_cast(a.moe_intermediate_size()); + c.scoringFunc = XMODEL_SCORING_FUNC_SIGMOID; + c.normTopKProb = a.norm_topk_prob(); + c.routeScale = a.routed_scaling_factor(); + c.nExpertGroups = static_cast(a.n_group()); + c.nLimitedGroups = static_cast(a.topk_group()); + c.expertsWeightTrans = true; // HF [out,in] -> [in,out] + + c.moeEpSize = static_cast(p.ep_size()); + c.moeTPSize = p.ep_size() > 1 + ? static_cast(p.world_size() / p.ep_size()) + : XliteTpSize(p); + return c; + } + + // GLM-5/5.1 (MLA + DSA indexer + sigmoid MoE + shared expert, default rope). + // GLM5 ~= DeepSeek V3.2 + DSA + different MLA dims; no rope_scaling -> + // YaRN/mscale skipped. + static XModelConfig FromGlm5(const ModelContext& context) { + XModelConfig c = FromDeepseekV3(context); + const ModelArgs& a = context.get_model_args(); + + // DSA: attnType=DSA triggers ForwardAttnIndexer -> topkIndices for sparse + // MLA attention. + c.attnType = XMODEL_ATTN_DSA; + c.indexHeadDim = static_cast(a.index_head_dim()); + c.indexNHeads = static_cast(a.index_n_heads()); + c.indexTopK = static_cast(a.index_topk()); + c.indexRopeInterleaved = a.indexer_rope_interleave(); + // csrc computes 1/sqrt(n*head_dim) internally; set ref value for + // completeness. + c.indexSoftmaxScale = 1.0f / std::sqrt(static_cast(c.indexHeadDim)); + + // DSA top-k sharing (GLM-5.2): shared layers skip indexer, reuse prev full + // layer's topkIndices. Resolve per-layer skip via DsaTopkSharePlan; csrc + // only reads the bool vector. Empty = all layers run indexer (GLM-5.1 + // behavior). + xllm::layer::DsaTopkSharePlan sharePlan(a); + c.indexerSkipLayers.resize(c.nLayers); + for (uint32_t i = 0; i < c.nLayers; ++i) { + c.indexerSkipLayers[i] = sharePlan.decision_for(i).reuse_topk; + } + + // W8A8 flags conditional on QuantArgs (IsW8A8; BF16 mis-set flips shared + // expert layout). + if (IsW8A8(context)) { + c.quantAttnWeightTrans = true; + c.quantAttnWeightNz = true; + c.expertsWeightNZ = true; + } + return c; + } + + // GLM-4 MoE (MHA + partial rotary + sigmoid MoE + shared expert, W8A8). + // Builds on FromQwen3Moe (MHA + MoE baseline), overrides GLM4-specific + // fields: + // ropeHeadDim = headDim * partial_rotary_factor (GLM4 partial rotary, not + // full rope) nSharedExperts = n_shared_experts (GLM4 has shared, Qwen3-MoE + // does not) scoringFunc = SIGMOID + routeScale + gateCaptured=false + // (sigmoid routing, gate at runtime) + // W8A8: quantAttnWeightTrans=true; NO_QUANT path (o_proj/down_proj BF16) + // unaffected. + static XModelConfig FromGlm4Moe(const ModelContext& context) { + XModelConfig c = FromQwen3Moe(context); + const ModelArgs& a = context.get_model_args(); + + // partial rotary: ropeHeadDim = headDim * partial_rotary_factor (front dims + // get RoPE). + if (a.partial_rotary_factor() > 0.0f) { + c.ropeHeadDim = static_cast(static_cast(c.headDim) * + a.partial_rotary_factor()); + } + + // GLM4 has shared expert (FromQwen3Moe sets nSharedExperts=0). + c.nSharedExperts = static_cast(a.n_shared_experts()); + + // sigmoid routing (Qwen3-MoE is softmax). gateCaptured=false: gate computed + // at runtime. + c.scoringFunc = XMODEL_SCORING_FUNC_SIGMOID; + c.routeScale = a.routed_scaling_factor(); + c.gateCaptured = false; + c.nExpertGroups = static_cast(a.n_group()); + c.nLimitedGroups = static_cast(a.topk_group()); + + // W8A8 flags conditional on QuantArgs (IsW8A8; hardcoded true leaks to BF16 + // GLM-4.7). + if (IsW8A8(context)) { + c.quantAttnWeightTrans = true; + c.quantAttnWeightNz = true; + c.expertsWeightNZ = true; + } + // expertsWeightTrans already true from FromQwen3Moe (routed expert + // gate_up/down transpose). + return c; + } +}; + +} // namespace xllm::xlite \ No newline at end of file diff --git a/xllm/core/layers/xlite/xlite_freqs_cis.h b/xllm/core/layers/xlite/xlite_freqs_cis.h new file mode 100644 index 0000000000..77e481aba9 --- /dev/null +++ b/xllm/core/layers/xlite/xlite_freqs_cis.h @@ -0,0 +1,101 @@ +/* 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 WARRANTIES OR CONDITIONS OF ANY KIND, either express or +implied. See the License for the specific language governing permissions or +limitations under the License. +==============================================================================*/ + +// Precompute rotary freqs_cis for xlite. MHA: real [cos||sin]; MLA/DSA: +// complex64. + +#pragma once + +#include +#include + +#include + +namespace xllm::xlite { + +// DeepSeek YaRN params (from ModelArgs rope_scaling_*, MLA only). +struct YarnConfig { + float rope_factor = 0.0f; + float beta_fast = 0.0f; + float beta_slow = 0.0f; + int64_t original_seq_len = 0; +}; + +class XliteFreqsCis { + public: + static torch::Tensor Precompute(const XModelConfig& cfg, + const torch::TensorOptions& opts, + const YarnConfig& yarn = {}) { + uint32_t dim = cfg.ropeHeadDim; // MHA: headDim; MLA: qk_rope_head_dim + uint64_t end = cfg.maxSeqLen; + float theta = cfg.ropeTheta; + auto inv_freq = + 1.0 / torch::pow(theta, + torch::arange( + 0, static_cast(dim), 2, torch::kFloat32) / + static_cast(dim)); + + if (cfg.attnType == XMODEL_ATTN_MLA || cfg.attnType == XMODEL_ATTN_DSA) { + // MLA/DSA: complex64 interleaved [c0,s0,c1,s1,...] via torch.polar + + // YaRN. + if (yarn.rope_factor > 0.0f && yarn.original_seq_len > 0 && + static_cast(end) > yarn.original_seq_len) { + inv_freq = apply_yarn(inv_freq, dim, theta, yarn); + } + auto t = torch::arange(static_cast(end), torch::kFloat32); + auto freqs = torch::outer(t, inv_freq); + return torch::polar(torch::ones_like(freqs), freqs).to(opts.device()); + } + + // MHA: real [cos||sin]. + auto t = torch::arange(static_cast(end), torch::kFloat32); + auto table = torch::outer(t, inv_freq); + return torch::cat({table.cos(), table.sin()}, /*dim=*/-1) + .to(opts.dtype()) + .to(opts.device()); + } + + private: + // YaRN: freqs/factor*(1-smooth) + freqs*smooth; smooth from beta_fast/slow + // correction. + static torch::Tensor apply_yarn(torch::Tensor inv_freq, + uint32_t dim, + float theta, + const YarnConfig& yarn) { + int64_t half = dim / 2; + auto corr_dim = [&](float num_rotations) { + return static_cast(dim) * + std::log(static_cast(yarn.original_seq_len) / + (num_rotations * 2.0 * M_PI)) / + (2.0 * std::log(static_cast(theta))); + }; + int64_t low = static_cast(std::floor(corr_dim(yarn.beta_fast))); + int64_t high = static_cast(std::ceil(corr_dim(yarn.beta_slow))); + low = std::max(low, static_cast(0)); + high = std::min(high, half - 1); + if (low == high) { + high += 1; + } + auto ramp = torch::clamp( + (torch::arange(half, torch::kFloat32) - static_cast(low)) / + static_cast(high - low), + 0.0, + 1.0); + auto smooth = 1.0 - ramp; + return inv_freq / yarn.rope_factor * (1.0 - smooth) + inv_freq * smooth; + } +}; + +} // namespace xllm::xlite \ No newline at end of file diff --git a/xllm/core/layers/xlite/xlite_init_utils.h b/xllm/core/layers/xlite/xlite_init_utils.h new file mode 100644 index 0000000000..d535c38140 --- /dev/null +++ b/xllm/core/layers/xlite/xlite_init_utils.h @@ -0,0 +1,49 @@ +/* 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. +==============================================================================*/ + +// torch::Tensor -> xlite XTensor bridge (XDtypeOf/TensorPtr/InitXTensor in +// xlite/xlite.h). + +#pragma once + +#include +#include + +#include "core/framework/parallel_state/parallel_args.h" + +namespace xllm::xlite { + +// xlite symbols are global (::); bring into xllm::xlite. +using ::InitXTensor; +using ::TensorPtr; +using ::XDtypeOf; +using ::XTensor; + +// Framework ParallelArgs::tp_size() is not populated; derive from +// world_size/dp_size. +inline uint32_t XliteTpSize(const ParallelArgs& pa) { + int32_t dp = pa.dp_size(); + if (dp <= 0) { + return 1u; + } + return static_cast(pa.world_size() / dp); +} + +inline uint32_t XliteTpRank(const ParallelArgs& pa) { + uint32_t tp = XliteTpSize(pa); + return tp > 0 ? static_cast(pa.rank()) % tp : 0u; +} + +} // namespace xllm::xlite \ No newline at end of file diff --git a/xllm/core/layers/xlite/xlite_model_adapter.h b/xllm/core/layers/xlite/xlite_model_adapter.h new file mode 100644 index 0000000000..e2386d8e22 --- /dev/null +++ b/xllm/core/layers/xlite/xlite_model_adapter.h @@ -0,0 +1,56 @@ +/* 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. +==============================================================================*/ + +// Adapter interface: each model implements BuildConfig + Load. + +#pragma once + +#include +#include + +#include +#include +#include +#include + +#include "core/framework/model_context.h" +#include "core/framework/parallel_state/parallel_args.h" +#include "core/framework/state_dict/state_dict.h" + +namespace xllm::xlite { + +class XliteModelAdapter { + public: + virtual ~XliteModelAdapter() = default; + + // Must populate runtime fields (maxBatchedTokens etc.) from singletons. + virtual XModelConfig BuildConfig(const ModelContext& context) = 0; + + // Called before XModel::Init. + virtual void Load(const StateDict& sd, + XModel& m, + const XModelConfig& cfg, + const ModelArgs& args, + const ParallelArgs& pa, + const torch::Device& device, + std::vector& storages) = 0; + + // EP=1 -> tp_size; EP>1 -> 1 (experts sharded by EP). Dense returns 1. + virtual uint32_t MoeTpSize(const ParallelArgs& pa) const { return 1; } + + virtual std::string Name() const { return ""; } +}; + +} // namespace xllm::xlite \ No newline at end of file diff --git a/xllm/core/layers/xlite/xlite_register_macros.h b/xllm/core/layers/xlite/xlite_register_macros.h new file mode 100644 index 0000000000..f421fb1c32 --- /dev/null +++ b/xllm/core/layers/xlite/xlite_register_macros.h @@ -0,0 +1,40 @@ +/* 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. +==============================================================================*/ + +// Registers "_xlite" model class + args loader. + +#pragma once + +#include + +#include "core/framework/model_context.h" +#include "core/layers/xlite/xlite_causal_lm_base.h" +#include "models/model_registry.h" + +namespace xllm::xlite { + +#define XLITE_REGISTER_MODEL(ModelType, AdapterClass, ArgsLambda) \ + class ModelType##XliteForCausalLMImpl : public XliteCausalLMBase { \ + public: \ + ModelType##XliteForCausalLMImpl(const ModelContext& ctx) \ + : XliteCausalLMBase(ctx, std::make_unique()) {} \ + }; \ + TORCH_MODULE(ModelType##XliteForCausalLM); \ + REGISTER_CAUSAL_MODEL_WITH_VARNAME( \ + ModelType##_xlite, ModelType##_xlite, ModelType##XliteForCausalLM); \ + REGISTER_MODEL_ARGS_WITH_VARNAME( \ + ModelType##_xlite, ModelType##_xlite, ArgsLambda) + +} // namespace xllm::xlite \ No newline at end of file diff --git a/xllm/core/layers/xlite/xlite_weight_utils.h b/xllm/core/layers/xlite/xlite_weight_utils.h new file mode 100644 index 0000000000..60d07d753d --- /dev/null +++ b/xllm/core/layers/xlite/xlite_weight_utils.h @@ -0,0 +1,1097 @@ +/* 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. +==============================================================================*/ + +// StateDict -> xlite XModel weight binding. Fuses qkv/gate_up via torch::cat. + +#pragma once + +#include +#include // at_npu::native::npu_format_cast (NZ) +#include + +#include +#include +#include + +#include "core/framework/model/model_args.h" +#include "core/framework/parallel_state/parallel_args.h" +#include "core/framework/state_dict/state_dict.h" +#include "core/layers/xlite/xlite_init_utils.h" + +namespace xllm::xlite { + +class XliteWeightUtils { + public: + // .to(device) and retain in storages (InitXTensor only stores data_ptr). + static const torch::Tensor& ToDevice(const torch::Tensor& t, + const std::string& name, + const torch::Device& device, + std::vector& storages) { + CHECK(t.defined()) << "[xlite] missing weight: " << name; + storages.push_back(t.to(device, /*non_blocking=*/false)); + return storages.back(); + } + + // W8A8 INT8 weight NZ pre-convert (npu_format_cast). INT8 matmul only; call + // after ToDevice. In-place replaces storages.back() (avoids ND+NZ double + // copy). BF16/deqScale excluded. + static const torch::Tensor& CastNz(const torch::Tensor& t, + const std::string& name, + std::vector& storages) { + CHECK(t.defined()) << "[xlite] CastNz undefined: " << name; + CHECK(!storages.empty()) + << "[xlite] CastNz: storages empty, expected ND tensor from ToDevice: " + << name; + storages.back() = + at_npu::native::npu_format_cast(t.contiguous(), ACL_FORMAT_FRACTAL_NZ); + return storages.back(); + } + + // TP-shard a weight along dim (returns rank-th shard). GQA: pass adjusted + // kv_tp_rank/kv_tp_size when n_kv_heads < tp_size. + static torch::Tensor Shard(const StateDict& sd, + const std::string& name, + int64_t dim, + int32_t tp_rank, + int32_t tp_size) { + torch::Tensor t = sd.get_sharded_tensor(name, dim, tp_rank, tp_size); + CHECK(t.defined()) << "[xlite] missing weight: " << name; + return t; + } + + // W8A8 deqScale transform: FP32[N] -> FP32[2N] (fixpipe uint64 layout, scale + // at even indices). Applies to weight_scale (BF16 [out,1]) and deq_scale (F32 + // [out]); csrc expects this layout. + static torch::Tensor TransformDeqScale(const torch::Tensor& scale) { + torch::Tensor fp32 = scale.to(torch::kFloat32).view({-1}).contiguous(); + int64_t n = fp32.size(0); + torch::Tensor out = torch::zeros({n * 2}, fp32.options()); + out.slice(0, 0, 2 * n, 2) = fp32; // stride 2: FP32 at even, 0 at odd + return out; + } + + // Whether weight is quantized (INT8). BF16 models skip quant fields; W8A8 + // loads them. + static bool IsQuant(const torch::Tensor& w) { + return w.defined() && w.scalar_type() == torch::kInt8; + } + + // embed/norm/head + per-layer attn (qkv fuse, o_proj, qk_norm) + mlpNorm. + // TP: embed/head vocab dim0, q/k/v out dim0, o_proj in dim1. + static void LoadAttnAndEmbed(const StateDict& sd, + XModel& m, + const XModelConfig& cfg, + const ModelArgs& args, + const ParallelArgs& pa, + const torch::Device& device, + std::vector& storages) { + auto dev = [&](const torch::Tensor& t, + const std::string& name) -> const torch::Tensor& { + return ToDevice(t, name, device, storages); + }; + // W8A8 norm weight is FP32 (ATB format); csrc norm kernel reads BF16 -> + // cast to avoid garbage. BF16 models: no-op. Norm bias handled separately. + auto devBf16 = [&](const torch::Tensor& t, + const std::string& name) -> const torch::Tensor& { + if (t.scalar_type() == torch::kBFloat16) { + return dev(t, name); + } + return dev(t.to(torch::kBFloat16).contiguous(), name); + }; + + const int32_t tp = static_cast(XliteTpSize(pa)); + const int32_t rank = static_cast(XliteTpRank(pa)); + // GQA: n_kv_heads < tp_size -> kv weights replicate across tp/n_kv_heads + // ranks. + const uint32_t nKvHeads = cfg.nKvHeads; + int32_t kv_tp = tp; + int32_t kv_rank = rank; + if (tp > 1 && nKvHeads > 0 && static_cast(tp) > nKvHeads) { + const int32_t repeat = tp / static_cast(nKvHeads); + kv_tp = static_cast(nKvHeads); + kv_rank = rank / repeat; + } + + LOG(INFO) << "[xlite] LoadAttnAndEmbed: tp=" << tp << " rank=" << rank + << " kv_tp=" << kv_tp << " kv_rank=" << kv_rank; + + // embed: vocab dim0. + // dev() once and rebind embed so the tie branch can share it (zero-copy). + torch::Tensor embed = + (tp > 1) ? Shard(sd, "model.embed_tokens.weight", 0, rank, tp) + : sd.get_tensor("model.embed_tokens.weight"); + embed = dev(embed, "model.embed_tokens.weight"); + InitXTensor(m.embed, embed); + // norm: not sharded. + InitXTensor( + m.norm, + devBf16(sd.get_tensor("model.norm.weight"), "model.norm.weight")); + // W8A8 norm bias is FP32 (ATB format); csrc norm kernel reads BF16 -> must + // cast to avoid byte reinterpretation garbage. BF16 models (RMSNorm) have + // no bias -> skip. + torch::Tensor normBias = sd.get_tensor("model.norm.bias"); + if (normBias.defined()) { + torch::Tensor b = normBias.to(torch::kBFloat16).contiguous(); + InitXTensor(m.normBias, dev(b, "model.norm.bias")); + } + // head: vocab dim0. + torch::Tensor head; + if (args.tie_word_embeddings()) { + // tie: reuse embed's device storage (InitXTensor is zero-copy). Avoids a + // second .to(device) of embed (~300MB for Qwen3-0.6B). + head = embed; + InitXTensor(m.head, embed); + } else { + head = (tp > 1) ? Shard(sd, "lm_head.weight", 0, rank, tp) + : sd.get_tensor("lm_head.weight"); + // GLM-5.2-W8A8: lm_head.weight is FP32; xlite XliteOpMatmul only supports + // BF16xFP32->FP32. ForwardGetLogits localOutput is BF16 (hiddenState + // dtype) -> cast to BF16. GLM-5.1 BF16 model: no-op. + if (head.scalar_type() != torch::kBFloat16) { + head = head.to(torch::kBFloat16).contiguous(); + } + InitXTensor(m.head, dev(head, "lm_head.weight")); + } + if (tp > 1) { + LOG(INFO) << "[xlite] shard rank=" << rank << ": embed=" << embed.sizes() + << " head=" << head.sizes(); + } + + for (uint32_t i = 0; i < cfg.nLayers; ++i) { + const std::string L = "model.layers." + std::to_string(i) + "."; + // attnNorm: per-token, not sharded. + InitXTensor(m.attnNorm[i], + devBf16(sd.get_tensor(L + "input_layernorm.weight"), + L + "input_layernorm.weight")); + // W8A8 q/k_norm bias; BF16 models have none. Cast to BF16 (same as + // normBias). + torch::Tensor attnNormBias = sd.get_tensor(L + "input_layernorm.bias"); + if (attnNormBias.defined()) { + torch::Tensor b = attnNormBias.to(torch::kBFloat16).contiguous(); + InitXTensor(m.attnNormBias[i], dev(b, L + "input_layernorm.bias")); + } + + // q/k/v: out dim0. q by nHeads/TP, kv by nKvHeads/TP (GQA-adjusted). + torch::Tensor q, k, v; + if (tp > 1) { + q = Shard(sd, L + "self_attn.q_proj.weight", 0, rank, tp); + k = Shard(sd, L + "self_attn.k_proj.weight", 0, kv_rank, kv_tp); + v = Shard(sd, L + "self_attn.v_proj.weight", 0, kv_rank, kv_tp); + } else { + q = sd.get_tensor(L + "self_attn.q_proj.weight"); + k = sd.get_tensor(L + "self_attn.k_proj.weight"); + v = sd.get_tensor(L + "self_attn.v_proj.weight"); + } + // host cat then one H2D (vs dev each shard + device cat). + // BF16 (NO_QUANT): [out, in] no transpose. W8A8 (STATIC): + // isTransposed=true, csrc expects [in, out] -> .t(). deqScale/quantBias + // per-out-channel unaffected. + torch::Tensor qkv = torch::cat({q, k, v}, /*dim=*/0) + .contiguous(); // host [q+k+v, hidden] + if (IsQuant(q)) { + qkv = qkv.t().contiguous(); // W8A8: [hidden, q+k+v] = [in, out] + } + qkv = dev(qkv, L + "self_attn.qkv (mhaQKV)"); + // W8A8 INT8 NZ: quantAttnWeightNz. + if (IsQuant(q) && cfg.quantAttnWeightNz) { + qkv = CastNz(qkv, L + "self_attn.qkv (mhaQKV) NZ", storages); + } + InitXTensor(m.mhaQKV[i].weight, qkv); + storages.push_back(qkv); + + // W8A8 STATIC quant: I8 weight + + // inputScale/inputOffset/deqScale/quantBias. BF16 skips. + // inputScale/inputOffset: per-tensor [1], same for q/k/v (take q's), not + // sharded. + if (IsQuant(q)) { + torch::Tensor iscale = + sd.get_tensor(L + "self_attn.q_proj.input_scale"); + torch::Tensor ioffset = + sd.get_tensor(L + "self_attn.q_proj.input_offset"); + if (iscale.defined()) { + // ATB input_scale is scale; csrc expects 1/scale (reciprocal) as BF16 + // [hidden]. Cast to BF16 + repeat to hiddenSize (per-tensor; csrc + // reads [hidden] per-channel). + torch::Tensor recip = (1.0f / iscale.to(torch::kFloat32)) + .to(torch::kBFloat16) + .contiguous(); + recip = recip.repeat({static_cast(cfg.hiddenSize)}); + InitXTensor(m.mhaQKV[i].inputScale, + dev(recip, L + "q_proj.input_scale_reciprocal")); + } + if (ioffset.defined()) { + // input_offset is zero-point (not reciprocal); same [hidden] repeat + // as inputScale. + torch::Tensor off = + ioffset.to(torch::kBFloat16) + .contiguous() + .repeat({static_cast(cfg.hiddenSize)}); + InitXTensor(m.mhaQKV[i].inputOffset, + dev(off, L + "q_proj.input_offset")); + } + // quantBias: cat(q,k,v) + TP shard dim0 (same as weight). + torch::Tensor qb_q, qb_k, qb_v; + if (tp > 1) { + qb_q = Shard(sd, L + "self_attn.q_proj.quant_bias", 0, rank, tp); + qb_k = + Shard(sd, L + "self_attn.k_proj.quant_bias", 0, kv_rank, kv_tp); + qb_v = + Shard(sd, L + "self_attn.v_proj.quant_bias", 0, kv_rank, kv_tp); + } else { + qb_q = sd.get_tensor(L + "self_attn.q_proj.quant_bias"); + qb_k = sd.get_tensor(L + "self_attn.k_proj.quant_bias"); + qb_v = sd.get_tensor(L + "self_attn.v_proj.quant_bias"); + } + if (qb_q.defined()) { + torch::Tensor qb = + torch::cat({qb_q, qb_k, qb_v}, /*dim=*/0).contiguous(); + InitXTensor(m.mhaQKV[i].quantBias, dev(qb, L + "qkv.quant_bias")); + } + // deqScale: cat(q,k,v) -> TransformDeqScale (FP32[2N] fixpipe) -> H2D. + torch::Tensor ds_q, ds_k, ds_v; + if (tp > 1) { + ds_q = Shard(sd, L + "self_attn.q_proj.deq_scale", 0, rank, tp); + ds_k = Shard(sd, L + "self_attn.k_proj.deq_scale", 0, kv_rank, kv_tp); + ds_v = Shard(sd, L + "self_attn.v_proj.deq_scale", 0, kv_rank, kv_tp); + } else { + ds_q = sd.get_tensor(L + "self_attn.q_proj.deq_scale"); + ds_k = sd.get_tensor(L + "self_attn.k_proj.deq_scale"); + ds_v = sd.get_tensor(L + "self_attn.v_proj.deq_scale"); + } + if (ds_q.defined()) { + torch::Tensor ds = TransformDeqScale( + torch::cat({ds_q, ds_k, ds_v}, /*dim=*/0).contiguous()); + InitXTensor(m.mhaQKV[i].deqScale, dev(ds, L + "qkv.deq_scale")); + } + } + + // o_proj: in dim1. + torch::Tensor o = + (tp > 1) ? Shard(sd, L + "self_attn.o_proj.weight", 1, rank, tp) + : sd.get_tensor(L + "self_attn.o_proj.weight"); + InitXTensor(m.attnOut[i].weight, dev(o, L + "self_attn.o_proj.weight")); + + if (i == 0 && tp > 1) { + LOG(INFO) << "[xlite] shard layer0 rank=" << rank + << ": qkv=" << qkv.sizes() << " o_proj=" << o.sizes(); + } + + if (cfg.addBias) { + torch::Tensor qb, kb, vb; + if (tp > 1) { + qb = Shard(sd, L + "self_attn.q_proj.bias", 0, rank, tp); + kb = Shard(sd, L + "self_attn.k_proj.bias", 0, kv_rank, kv_tp); + vb = Shard(sd, L + "self_attn.v_proj.bias", 0, kv_rank, kv_tp); + } else { + qb = sd.get_tensor(L + "self_attn.q_proj.bias"); + kb = sd.get_tensor(L + "self_attn.k_proj.bias"); + vb = sd.get_tensor(L + "self_attn.v_proj.bias"); + } + // bias same as qkv: host cat then one H2D. + torch::Tensor qkvb = + torch::cat({qb, kb, vb}, /*dim=*/0).contiguous(); // host + qkvb = dev(qkvb, L + "self_attn.qkv_bias (mhaQKVBias)"); + InitXTensor(m.mhaQKVBias[i], qkvb); + storages.push_back(qkvb); + } + if (cfg.qkNorm) { + // q_norm/k_norm: [headDim] shared across heads — no sharding. + InitXTensor(m.mhaQNorm[i], + devBf16(sd.get_tensor(L + "self_attn.q_norm.weight"), + L + "self_attn.q_norm.weight")); + InitXTensor(m.mhaKNorm[i], + devBf16(sd.get_tensor(L + "self_attn.k_norm.weight"), + L + "self_attn.k_norm.weight")); + // W8A8 q/k_norm bias; BF16 models have none. Cast to BF16 (same as + // normBias). + torch::Tensor qNormBias = sd.get_tensor(L + "self_attn.q_norm.bias"); + if (qNormBias.defined()) { + torch::Tensor b = qNormBias.to(torch::kBFloat16).contiguous(); + InitXTensor(m.mhaQNormBias[i], dev(b, L + "self_attn.q_norm.bias")); + } + torch::Tensor kNormBias = sd.get_tensor(L + "self_attn.k_norm.bias"); + if (kNormBias.defined()) { + torch::Tensor b = kNormBias.to(torch::kBFloat16).contiguous(); + InitXTensor(m.mhaKNormBias[i], dev(b, L + "self_attn.k_norm.bias")); + } + } + + // mlpNorm: not sharded. + InitXTensor(m.mlpNorm[i], + devBf16(sd.get_tensor(L + "post_attention_layernorm.weight"), + L + "post_attention_layernorm.weight")); + // W8A8 mlpNorm bias; BF16 models have none. Cast to BF16 (same as + // normBias). + torch::Tensor mlpNormBias = + sd.get_tensor(L + "post_attention_layernorm.bias"); + if (mlpNormBias.defined()) { + torch::Tensor b = mlpNormBias.to(torch::kBFloat16).contiguous(); + InitXTensor(m.mlpNormBias[i], + dev(b, L + "post_attention_layernorm.bias")); + } + } + LOG(INFO) << "[xlite] LoadAttnAndEmbed: " << cfg.nLayers << " layers done"; + } + + // Dense FFN for layer i (gate_up fuse + down). gate_up dim0, down dim1. + static void LoadDenseMlp(const StateDict& sd, + XModel& m, + const XModelConfig& cfg, + const ParallelArgs& pa, + const torch::Device& device, + std::vector& storages, + uint32_t i) { + auto dev = [&](const torch::Tensor& t, + const std::string& name) -> const torch::Tensor& { + return ToDevice(t, name, device, storages); + }; + const std::string L = "model.layers." + std::to_string(i) + "."; + const int32_t tp = static_cast(XliteTpSize(pa)); + const int32_t rank = static_cast(XliteTpRank(pa)); + + torch::Tensor gp, up, down; + if (tp > 1) { + // gate/up: column-parallel dim0. down: row-parallel dim1. + gp = Shard(sd, L + "mlp.gate_proj.weight", 0, rank, tp); + up = Shard(sd, L + "mlp.up_proj.weight", 0, rank, tp); + down = Shard(sd, L + "mlp.down_proj.weight", 1, rank, tp); + } else { + gp = sd.get_tensor(L + "mlp.gate_proj.weight"); + up = sd.get_tensor(L + "mlp.up_proj.weight"); + down = sd.get_tensor(L + "mlp.down_proj.weight"); + } + // Dense MLP: BF16 no transpose (expects [out,in]); W8A8 transpose to + // [in,out]. mlpDown is BF16 (NO_QUANT) -> no transpose. + torch::Tensor gate_up = + torch::cat({gp, up}, /*dim=*/0).contiguous(); // host [2*inter, hidden] + if (IsQuant(gp)) { + gate_up = + gate_up.t().contiguous(); // W8A8: [hidden, 2*inter] = [in, out] + } + gate_up = dev(gate_up, L + "mlp.gate_proj+up_proj (mlpUpGate)"); + // W8A8 INT8 NZ: dense MLP gate_up (quantAttnWeightNz). mlpDown BF16 no + // transpose. + if (IsQuant(gp) && cfg.quantAttnWeightNz) { + gate_up = + CastNz(gate_up, L + "mlp.gate_proj+up_proj (mlpUpGate) NZ", storages); + } + InitXTensor(m.mlpUpGate[i].weight, gate_up); + storages.push_back(gate_up); + // W8A8 DYNAMIC: gate/up has weight_scale (BF16 [out,1]), no + // input/quant_bias/deq_scale. deqScale = TransformDeqScale(cat(gate_scale, + // up_scale) shard dim0). down_proj BF16 -> NO_QUANT. weight_offset unused. + if (IsQuant(gp)) { + torch::Tensor gs, us; + if (tp > 1) { + gs = Shard(sd, L + "mlp.gate_proj.weight_scale", 0, rank, tp); + us = Shard(sd, L + "mlp.up_proj.weight_scale", 0, rank, tp); + } else { + gs = sd.get_tensor(L + "mlp.gate_proj.weight_scale"); + us = sd.get_tensor(L + "mlp.up_proj.weight_scale"); + } + if (gs.defined()) { + torch::Tensor ds = + TransformDeqScale(torch::cat({gs, us}, /*dim=*/0).contiguous()); + InitXTensor(m.mlpUpGate[i].deqScale, + dev(ds, L + "mlp.gate_proj+up_proj.deq_scale")); + } + } + // down_proj: row-parallel (TP shard dim1, ForwardMLP AllReduce). + // GLM-5.2 down=W8A8_DYNAMIC: transpose + CastNz + deqScale. GLM-4.7 + // down=BF16 (no transpose). row-parallel deqScale not sharded (same as + // shared expert down). + torch::Tensor down_w = IsQuant(down) ? down.t().contiguous() : down; + down_w = dev(down_w, L + "mlp.down_proj.weight"); + if (IsQuant(down) && cfg.quantAttnWeightNz) { + down_w = CastNz(down_w, L + "mlp.down_proj (mlpDown) NZ", storages); + } + InitXTensor(m.mlpDown[i].weight, down_w); + storages.push_back(down_w); + if (IsQuant(down)) { + torch::Tensor ds_down = sd.get_tensor(L + "mlp.down_proj.weight_scale"); + if (ds_down.defined()) { + torch::Tensor ds = TransformDeqScale(ds_down); + InitXTensor(m.mlpDown[i].deqScale, + dev(ds, L + "mlp.down_proj.deq_scale")); + } + } + } + + // MoE FFN for layer i. HF weight [out, in] -> transpose to [in, out] + // (expertsWeightTrans=true). EP>1: local expert range; EP==1,TP>1: weight + // split. + static void LoadMoEExperts(const StateDict& sd, + XModel& m, + const XModelConfig& cfg, + const ParallelArgs& pa, + const torch::Device& device, + std::vector& storages, + uint32_t i) { + auto dev = [&](const torch::Tensor& t, + const std::string& name) -> const torch::Tensor& { + return ToDevice(t, name, device, storages); + }; + const std::string L = "model.layers." + std::to_string(i) + "."; + + // moeGate: not sharded (gate logits are all_reduce'd). + InitXTensor( + m.moeGate[i], + dev(sd.get_tensor(L + "mlp.gate.weight"), L + "mlp.gate.weight")); + // moeGateBias: sigmoid-only (e_score_correction_bias [n_routed_experts]); + // absent in softmax (Qwen3-MoE). + if (cfg.scoringFunc == XMODEL_SCORING_FUNC_SIGMOID) { + torch::Tensor gate_bias = + sd.get_tensor(L + "mlp.gate.e_score_correction_bias"); + if (gate_bias.defined()) { + InitXTensor(m.moeGateBias[i], + dev(gate_bias, L + "mlp.gate.e_score_correction_bias")); + } + } + + // EP: local expert range [start, end). + const uint32_t ep = cfg.moeEpSize; + const uint32_t moeTp = cfg.moeTPSize; + const uint32_t nLocal = cfg.nRoutedExperts / ep; + const uint32_t rank = static_cast(pa.rank()); + const uint32_t start = (ep == 1) ? 0 : (rank / moeTp) * nLocal; + const uint32_t end = start + nLocal; + + for (uint32_t e = start; e < end; ++e) { + const std::string E = L + "mlp.experts." + std::to_string(e) + "."; + // tp_rank: EP-group-internal TP rank (rank % moeTp). + const int32_t tp_rank = static_cast(rank % moeTp); + torch::Tensor gp, up, down; + if (moeTp > 1) { + // TP: gate/up dim0, down dim1. + gp = Shard(sd, + E + "gate_proj.weight", + 0, + tp_rank, + static_cast(moeTp)); + up = Shard( + sd, E + "up_proj.weight", 0, tp_rank, static_cast(moeTp)); + down = Shard(sd, + E + "down_proj.weight", + 1, + tp_rank, + static_cast(moeTp)); + } else { + gp = sd.get_tensor(E + "gate_proj.weight"); + up = sd.get_tensor(E + "up_proj.weight"); + down = sd.get_tensor(E + "down_proj.weight"); + } + // host fused cat+t+contiguous then one H2D. + torch::Tensor gate_up = torch::cat({gp, up}, /*dim=*/0).t().contiguous(); + gate_up = dev(gate_up, E + "gate_proj+up_proj (moeREUpGate)"); + // W8A8 INT8 NZ: routed expert gate_up (group_matmul). + if (IsQuant(gp) && cfg.expertsWeightNZ) { + gate_up = + CastNz(gate_up, E + "gate_proj+up_proj (moeREUpGate) NZ", storages); + } + // Store at original expert index e. + InitXTensor(m.moeREUpGate[i][e], gate_up); + storages.push_back(gate_up); + // down: transpose to [in, out] on host then one H2D. + torch::Tensor down_t = down.t().contiguous(); + down_t = dev(down_t, E + "down_proj (moeREDown)"); + // W8A8 INT8 NZ: routed expert down (group_matmul). + if (IsQuant(down) && cfg.expertsWeightNZ) { + down_t = CastNz(down_t, E + "down_proj (moeREDown) NZ", storages); + } + CHECK(down_t.defined()) + << "[xlite] missing weight: " << E + "down_proj.weight"; + InitXTensor(m.moeREDown[i][e], down_t); + storages.push_back(down_t); + // W8A8 DYNAMIC: expert gate/up/down weight_scale -> separate deqScale + // members (group_matmul). gate/up scale shard dim0; down scale + // row-parallel (not sharded). + if (IsQuant(gp)) { + torch::Tensor gs, us; + if (moeTp > 1) { + gs = Shard(sd, + E + "gate_proj.weight_scale", + 0, + tp_rank, + static_cast(moeTp)); + us = Shard(sd, + E + "up_proj.weight_scale", + 0, + tp_rank, + static_cast(moeTp)); + } else { + gs = sd.get_tensor(E + "gate_proj.weight_scale"); + us = sd.get_tensor(E + "up_proj.weight_scale"); + } + if (gs.defined()) { + torch::Tensor ds = + TransformDeqScale(torch::cat({gs, us}, /*dim=*/0).contiguous()); + InitXTensor(m.moeREUpGateDeqScale[i][e], + dev(ds, E + "gate+up.deq_scale")); + } + torch::Tensor ds_down = sd.get_tensor(E + "down_proj.weight_scale"); + if (ds_down.defined()) { + torch::Tensor ds = TransformDeqScale(ds_down); + InitXTensor(m.moeREDownDeqScale[i][e], dev(ds, E + "down.deq_scale")); + } + } + if (i == cfg.nDenseLayers && e == start) { + LOG(INFO) << "[xlite] shard layer" << i << " expert" << e + << " rank=" << rank << ": gate_up=" << gate_up.sizes() + << " down=" << down_t.sizes(); + } + } + // Shared expert (DeepSeek/GLM5, nSharedExperts>0); TP shard by moeTp (same + // as routed). + if (cfg.nSharedExperts > 0) { + const std::string S = L + "mlp.shared_experts."; + // tp_rank: same as routed expert (EP-group-internal). + const int32_t tp_rank = static_cast(rank % moeTp); + torch::Tensor sgp, sup, sdown; + if (moeTp > 1) { + sgp = Shard(sd, + S + "gate_proj.weight", + 0, + tp_rank, + static_cast(moeTp)); + sup = Shard( + sd, S + "up_proj.weight", 0, tp_rank, static_cast(moeTp)); + sdown = Shard(sd, + S + "down_proj.weight", + 1, + tp_rank, + static_cast(moeTp)); + } else { + sgp = sd.get_tensor(S + "gate_proj.weight"); + sup = sd.get_tensor(S + "up_proj.weight"); + sdown = sd.get_tensor(S + "down_proj.weight"); + } + // Shared expert: BF16 no transpose (ForwardLinear expects [out,in]); + // W8A8 transpose to [in,out] (same as dense mlpUpGate). host cat then one + // H2D. + + torch::Tensor sgate_up = torch::cat({sgp, sup}, /*dim=*/0).contiguous(); + bool seQuant = IsQuant(sgp); + if (seQuant) { + sgate_up = sgate_up.t().contiguous(); + } + sgate_up = dev(sgate_up, S + "gate_proj+up_proj (moeSEUpGate)"); + // W8A8 INT8 NZ: shared expert (single expert, not group_matmul). + if (seQuant && cfg.quantAttnWeightNz) { + sgate_up = CastNz( + sgate_up, S + "gate_proj+up_proj (moeSEUpGate) NZ", storages); + } + InitXTensor(m.moeSEUpGate[i].weight, sgate_up); + storages.push_back(sgate_up); + // shared down: W8A8 transpose to [in,out]; BF16 no transpose. + torch::Tensor sdown_t = seQuant ? sdown.t().contiguous() : sdown; + sdown_t = dev(sdown_t, S + "down_proj (moeSEDown)"); + // W8A8 INT8 NZ: shared expert down. + if (seQuant && cfg.quantAttnWeightNz) { + sdown_t = CastNz(sdown_t, S + "down_proj (moeSEDown) NZ", storages); + } + CHECK(sdown_t.defined()) + << "[xlite] missing weight: " << S + "down_proj.weight"; + InitXTensor(m.moeSEDown[i].weight, sdown_t); + storages.push_back(sdown_t); + // W8A8 DYNAMIC: shared expert gate/up/down weight_scale -> ForwardLinear + // (MatmulWeight.deqScale). gate/up scale shard dim0; down not sharded. + if (IsQuant(sgp)) { + torch::Tensor sgs, sus; + if (moeTp > 1) { + sgs = Shard(sd, + S + "gate_proj.weight_scale", + 0, + tp_rank, + static_cast(moeTp)); + sus = Shard(sd, + S + "up_proj.weight_scale", + 0, + tp_rank, + static_cast(moeTp)); + } else { + sgs = sd.get_tensor(S + "gate_proj.weight_scale"); + sus = sd.get_tensor(S + "up_proj.weight_scale"); + } + if (sgs.defined()) { + torch::Tensor ds = + TransformDeqScale(torch::cat({sgs, sus}, /*dim=*/0).contiguous()); + InitXTensor(m.moeSEUpGate[i].deqScale, + dev(ds, S + "gate+up.deq_scale")); + } + torch::Tensor ds_down = sd.get_tensor(S + "down_proj.weight_scale"); + if (ds_down.defined()) { + torch::Tensor ds = TransformDeqScale(ds_down); + InitXTensor(m.moeSEDown[i].deqScale, dev(ds, S + "down.deq_scale")); + } + } + } + if (i == cfg.nDenseLayers) { + LOG(INFO) << "[xlite] LoadMoEExperts: layer " << i << " loaded " + << (end - start) << "/" << cfg.nRoutedExperts + << " experts (ep=" << ep << " moeTp=" << moeTp << " range=[" + << start << "," << end << "))"; + } + } + + // Dense: attn + dense MLP for all layers. + static void LoadMHA(const StateDict& sd, + XModel& m, + const XModelConfig& cfg, + const ModelArgs& args, + const ParallelArgs& pa, + const torch::Device& device, + std::vector& storages) { + LoadAttnAndEmbed(sd, m, cfg, args, pa, device, storages); + for (uint32_t i = 0; i < cfg.nDenseLayers; ++i) { + LoadDenseMlp(sd, m, cfg, pa, device, storages, i); + } + } + + // FFN binding: dense MLP for i=nDenseLayers. + static void LoadMoEFFN(const StateDict& sd, + XModel& m, + const XModelConfig& cfg, + const ParallelArgs& pa, + const torch::Device& device, + std::vector& storages) { + for (uint32_t i = 0; i < cfg.nLayers; ++i) { + if (i < cfg.nDenseLayers) { + LoadDenseMlp(sd, m, cfg, pa, device, storages, i); + } else { + LoadMoEExperts(sd, m, cfg, pa, device, storages, i); + } + } + } + + // MoE: attn + dense MLP for first nDenseLayers, MoE FFN for the rest. + static void LoadMoE(const StateDict& sd, + XModel& m, + const XModelConfig& cfg, + const ModelArgs& args, + const ParallelArgs& pa, + const torch::Device& device, + std::vector& storages) { + LOG(INFO) << "[xlite] LoadMoE: nLayers=" << cfg.nLayers + << " nDense=" << cfg.nDenseLayers + << " experts=" << cfg.nRoutedExperts << " act=" << cfg.nActExperts + << " moeInter=" << cfg.moeIntermediateSize + << " trans=" << cfg.expertsWeightTrans << " ep=" << cfg.moeEpSize + << " moeTp=" << cfg.moeTPSize << " rank=" << pa.rank(); + LoadAttnAndEmbed(sd, m, cfg, args, pa, device, storages); + LoadMoEFFN(sd, m, cfg, pa, device, storages); + LOG(INFO) << "[xlite] LoadMoE: all layers done, storages=" + << storages.size(); + } + + // MLA attention (DeepSeek-V3/R1, GLM5): bind MLA weights, FFN via LoadMoEFFN. + static void LoadMLAAttn(const StateDict& sd, + XModel& m, + const XModelConfig& cfg, + const ModelArgs& args, + const ParallelArgs& pa, + const torch::Device& device, + std::vector& storages) { + auto dev = [&](const torch::Tensor& t, + const std::string& name) -> const torch::Tensor& { + return ToDevice(t, name, device, storages); + }; + // W8A8 norm weight FP32->BF16 (same as LoadAttnAndEmbed devBf16). + auto devBf16 = [&](const torch::Tensor& t, + const std::string& name) -> const torch::Tensor& { + if (t.scalar_type() == torch::kBFloat16) { + return dev(t, name); + } + return dev(t.to(torch::kBFloat16).contiguous(), name); + }; + const int32_t tp = XliteTpSize(pa); + const int32_t rank = XliteTpRank(pa); + + // embed/head: vocab dim0 (same as MHA); norm not sharded. + // dev() once and rebind embed so the tie branch can share it (zero-copy). + torch::Tensor embed = + (tp > 1) ? Shard(sd, "model.embed_tokens.weight", 0, rank, tp) + : sd.get_tensor("model.embed_tokens.weight"); + embed = dev(embed, "model.embed_tokens.weight"); + InitXTensor(m.embed, embed); + InitXTensor( + m.norm, + devBf16(sd.get_tensor("model.norm.weight"), "model.norm.weight")); + torch::Tensor head; + if (args.tie_word_embeddings()) { + // tie: reuse embed's device storage (zero-copy). Avoids a second + // .to(device). + head = embed; + InitXTensor(m.head, embed); + } else { + head = (tp > 1) ? Shard(sd, "lm_head.weight", 0, rank, tp) + : sd.get_tensor("lm_head.weight"); + // GLM-5.2-W8A8: lm_head.weight is FP32; xlite XliteOpMatmul only supports + // BF16xFP32->FP32. ForwardGetLogits localOutput is BF16 (hiddenState + // dtype) -> cast to BF16. GLM-5.1 BF16 model: no-op. + if (head.scalar_type() != torch::kBFloat16) { + head = head.to(torch::kBFloat16).contiguous(); + } + InitXTensor(m.head, dev(head, "lm_head.weight")); + } + + for (uint32_t i = 0; i < cfg.nLayers; ++i) { + const std::string L = "model.layers." + std::to_string(i) + "."; + InitXTensor(m.attnNorm[i], + devBf16(sd.get_tensor(L + "input_layernorm.weight"), + L + "input_layernorm.weight")); + InitXTensor(m.mlpNorm[i], + devBf16(sd.get_tensor(L + "post_attention_layernorm.weight"), + L + "post_attention_layernorm.weight")); + + // mlaQKVA = q_a + kv_a_with_mqa concat (q first; xlite csrc expects + // q-first). Lora compressed, not sharded. + torch::Tensor q_a = sd.get_tensor(L + "self_attn.q_a_proj.weight"); + torch::Tensor kv_a = + sd.get_tensor(L + "self_attn.kv_a_proj_with_mqa.weight"); + CHECK(q_a.defined()) << "[xlite] missing: " << L + << "self_attn.q_a_proj.weight"; + CHECK(kv_a.defined()) + << "[xlite] missing: " << L << "self_attn.kv_a_proj_with_mqa.weight"; + // host cat then one H2D. dev(qkva) already pushes the device copy into + // storages (ToDevice), so no separate push (the old + // `storages.push_back(qkva)` was a redundant host copy, ~4MB/layer). + torch::Tensor qkva = torch::cat({q_a, kv_a}, /*dim=*/0) + .contiguous(); // [q_lora+kv_lora+rope, hidden] + // W8A8 (GLM-5.2 STATIC): isTransposed=true, csrc expects [in, out] -> + // .t(). BF16 no transpose. + bool qaQuant = IsQuant(q_a); + if (qaQuant) { + qkva = qkva.t().contiguous(); // W8A8: [hidden, q_lora+kv_lora+rope] = + // [in, out] + } + qkva = dev(qkva, L + "self_attn.mlaQKVA"); + // W8A8 INT8 NZ: quantAttnWeightNz. + if (qaQuant && cfg.quantAttnWeightNz) { + qkva = CastNz(qkva, L + "self_attn.mlaQKVA NZ", storages); + } + InitXTensor(m.mlaQKVA[i].weight, qkva); + storages.push_back(qkva); + // W8A8 STATIC (GLM-5.2): input_scale/input_offset/quant_bias/deq_scale. + // inputScale: reciprocal + BF16, [1]->[hidden] repeat (per-channel read). + // Uses q_a's input_scale (q_a/kv_a verified identical). + if (qaQuant) { + torch::Tensor iscale = + sd.get_tensor(L + "self_attn.q_a_proj.input_scale"); + torch::Tensor ioffset = + sd.get_tensor(L + "self_attn.q_a_proj.input_offset"); + if (iscale.defined()) { + torch::Tensor recip = (1.0f / iscale.to(torch::kFloat32)) + .to(torch::kBFloat16) + .contiguous(); + recip = recip.repeat({static_cast(cfg.hiddenSize)}); + InitXTensor(m.mlaQKVA[i].inputScale, + dev(recip, L + "q_a_proj.input_scale_reciprocal")); + } + if (ioffset.defined()) { + torch::Tensor off = + ioffset.to(torch::kBFloat16) + .contiguous() + .repeat({static_cast(cfg.hiddenSize)}); + InitXTensor(m.mlaQKVA[i].inputOffset, + dev(off, L + "q_a_proj.input_offset")); + } + // quantBias/deqScale: cat{q_a, kv_a} not sharded (mlaQKVA + // column-parallel). + torch::Tensor qb_qa = + sd.get_tensor(L + "self_attn.q_a_proj.quant_bias"); + torch::Tensor qb_kva = + sd.get_tensor(L + "self_attn.kv_a_proj_with_mqa.quant_bias"); + if (qb_qa.defined()) { + torch::Tensor qb = + torch::cat({qb_qa, qb_kva}, /*dim=*/0).contiguous(); + InitXTensor(m.mlaQKVA[i].quantBias, + dev(qb, L + "mlaQKVA.quant_bias")); + } + torch::Tensor ds_qa = sd.get_tensor(L + "self_attn.q_a_proj.deq_scale"); + torch::Tensor ds_kva = + sd.get_tensor(L + "self_attn.kv_a_proj_with_mqa.deq_scale"); + if (ds_qa.defined()) { + torch::Tensor ds = TransformDeqScale( + torch::cat({ds_qa, ds_kva}, /*dim=*/0).contiguous()); + InitXTensor(m.mlaQKVA[i].deqScale, dev(ds, L + "mlaQKVA.deq_scale")); + } + } + + // q_a/kv_a_layernorm (RMSNorm, not sharded; DeepSeek has no bias). + InitXTensor(m.mlaQNorm[i], + devBf16(sd.get_tensor(L + "self_attn.q_a_layernorm.weight"), + L + "self_attn.q_a_layernorm.weight")); + InitXTensor(m.mlaKVNorm[i], + devBf16(sd.get_tensor(L + "self_attn.kv_a_layernorm.weight"), + L + "self_attn.kv_a_layernorm.weight")); + + // mlaQB = q_b_proj; TP shard dim0 (by n_heads/tp). + torch::Tensor q_b = + (tp > 1) ? Shard(sd, L + "self_attn.q_b_proj.weight", 0, rank, tp) + : sd.get_tensor(L + "self_attn.q_b_proj.weight"); + // W8A8 (GLM-5.2 STATIC): .t() + CastNz (same as mlaQKVA). BF16 no + // transpose. + bool qbQuant = IsQuant(q_b); + if (qbQuant) { + q_b = + q_b.t().contiguous(); // [q_lora, n_heads*(nope+rope)] = [in, out] + } + q_b = dev(q_b, L + "self_attn.q_b_proj.weight"); + if (qbQuant && cfg.quantAttnWeightNz) { + q_b = CastNz(q_b, L + "self_attn.q_b_proj NZ", storages); + } + InitXTensor(m.mlaQB[i].weight, q_b); + storages.push_back(q_b); + // W8A8 STATIC: mlaQB input_dim=qLoraRank. TP shard dim0; + // quantBias/deqScale shard dim0. + if (qbQuant) { + torch::Tensor iscale = + sd.get_tensor(L + "self_attn.q_b_proj.input_scale"); + torch::Tensor ioffset = + sd.get_tensor(L + "self_attn.q_b_proj.input_offset"); + if (iscale.defined()) { + torch::Tensor recip = (1.0f / iscale.to(torch::kFloat32)) + .to(torch::kBFloat16) + .contiguous(); + recip = recip.repeat({static_cast(cfg.qLoraRank)}); + InitXTensor(m.mlaQB[i].inputScale, + dev(recip, L + "q_b_proj.input_scale_reciprocal")); + } + if (ioffset.defined()) { + torch::Tensor off = + ioffset.to(torch::kBFloat16) + .contiguous() + .repeat({static_cast(cfg.qLoraRank)}); + InitXTensor(m.mlaQB[i].inputOffset, + dev(off, L + "q_b_proj.input_offset")); + } + torch::Tensor qb_w = + (tp > 1) + ? Shard(sd, L + "self_attn.q_b_proj.quant_bias", 0, rank, tp) + : sd.get_tensor(L + "self_attn.q_b_proj.quant_bias"); + if (qb_w.defined()) { + InitXTensor(m.mlaQB[i].quantBias, + dev(qb_w, L + "q_b_proj.quant_bias")); + } + torch::Tensor ds_w = + (tp > 1) + ? Shard(sd, L + "self_attn.q_b_proj.deq_scale", 0, rank, tp) + : sd.get_tensor(L + "self_attn.q_b_proj.deq_scale"); + if (ds_w.defined()) { + torch::Tensor ds = TransformDeqScale(ds_w); + InitXTensor(m.mlaQB[i].deqScale, dev(ds, L + "q_b_proj.deq_scale")); + } + } + + // kv_b split into WUKT (q absorb) + WUV (output proj); 0.2.0rc0 MLA + // refactor (old mlaKVB removed). TP shard dim0; host reshape+split then + // 2x H2D. + torch::Tensor kv_b = + (tp > 1) ? Shard(sd, L + "self_attn.kv_b_proj.weight", 0, rank, tp) + : sd.get_tensor(L + "self_attn.kv_b_proj.weight"); + CHECK(kv_b.defined()) + << "[xlite] missing: " << L << "self_attn.kv_b_proj.weight"; + // kv_b: [n_local_heads*(nope+v), kv_lora] -> [h, nope+v, kv_lora] + uint32_t nLocalHeadsMla = cfg.nHeads / tp; + torch::Tensor wkv_b = + kv_b.view({static_cast(nLocalHeadsMla), + static_cast(cfg.nopeHeadDim + cfg.vHeadDim), + static_cast(cfg.kvLoraRank)}) + .contiguous(); + // WUKT: [h, nope, kv_lora], no transpose (matches csrc htd layout). + torch::Tensor wuk_t = wkv_b.slice(1, 0, cfg.nopeHeadDim).contiguous(); + // WUV: [h, v, kv_lora] -> permute(0,2,1) for csrc EinsumMhtHtdMhd ([h, + // kv_lora, v]). + torch::Tensor wuv = + wkv_b.slice(1, cfg.nopeHeadDim).transpose(1, 2).contiguous(); + InitXTensor(m.mlaWUKT[i], dev(wuk_t, L + "self_attn.kv_b_proj.wuk_t")); + InitXTensor(m.mlaWUV[i], dev(wuv, L + "self_attn.kv_b_proj.wuv")); + + // attnOut = o_proj; TP shard dim1 (row parallel, AllReduce). + torch::Tensor o = + (tp > 1) ? Shard(sd, L + "self_attn.o_proj.weight", 1, rank, tp) + : sd.get_tensor(L + "self_attn.o_proj.weight"); + // W8A8 (GLM-5.2 STATIC): .t() + CastNz. row-parallel (shard dim1, + // AllReduce). + bool oQuant = IsQuant(o); + if (oQuant) { + o = o.t().contiguous(); // [n_heads*v, hidden] = [in, out] + } + o = dev(o, L + "self_attn.o_proj.weight"); + if (oQuant && cfg.quantAttnWeightNz) { + o = CastNz(o, L + "self_attn.o_proj NZ", storages); + } + InitXTensor(m.attnOut[i].weight, o); + storages.push_back(o); + // W8A8 STATIC: o_proj shard dim1 (AllReduce). quantBias rank0 only (avoid + // double sum); deqScale not sharded (distributive). inputScale/Offset + // repeat to nHeads/tp*vHeadDim. + if (oQuant) { + uint32_t nLocalHeads = cfg.nHeads / tp; + int64_t oInputDim = static_cast(nLocalHeads) * cfg.vHeadDim; + torch::Tensor iscale = + sd.get_tensor(L + "self_attn.o_proj.input_scale"); + torch::Tensor ioffset = + sd.get_tensor(L + "self_attn.o_proj.input_offset"); + if (iscale.defined()) { + torch::Tensor recip = (1.0f / iscale.to(torch::kFloat32)) + .to(torch::kBFloat16) + .contiguous(); + recip = recip.repeat({oInputDim}); + InitXTensor(m.attnOut[i].inputScale, + dev(recip, L + "o_proj.input_scale_reciprocal")); + } + if (ioffset.defined()) { + torch::Tensor off = + ioffset.to(torch::kBFloat16).contiguous().repeat({oInputDim}); + InitXTensor(m.attnOut[i].inputOffset, + dev(off, L + "o_proj.input_offset")); + } + if (rank == 0) { + torch::Tensor qb_o = sd.get_tensor(L + "self_attn.o_proj.quant_bias"); + if (qb_o.defined()) { + InitXTensor(m.attnOut[i].quantBias, + dev(qb_o, L + "o_proj.quant_bias")); + } + } + torch::Tensor ds_o = sd.get_tensor(L + "self_attn.o_proj.deq_scale"); + if (ds_o.defined()) { + torch::Tensor ds = TransformDeqScale(ds_o); + InitXTensor(m.attnOut[i].deqScale, dev(ds, L + "o_proj.deq_scale")); + } + } + + // DSA indexer (GLM-5/5.1, attnType==DSA). Weights not TP-sharded (indexer + // replicated; topkIndices must match across ranks). + // indexKWeightsProj = cat{wk, weights_proj} (wk first, no transpose). + // indexKNorm/Bias = k_norm (LayerNorm w/ bias). indexQB = wq_b (no + // transpose). + // GLM-5.2 shared layers (cfg.indexerSkipLayers[i]): skip indexer weight + // binding. Empty list (e.g. consumers without GLM-5.2 support) = no + // sharing, bind every layer. + if (cfg.attnType == XMODEL_ATTN_DSA && + (cfg.indexerSkipLayers.empty() || + (i < cfg.indexerSkipLayers.size() && !cfg.indexerSkipLayers[i]))) { + torch::Tensor wk = sd.get_tensor(L + "self_attn.indexer.wk.weight"); + torch::Tensor wproj = + sd.get_tensor(L + "self_attn.indexer.weights_proj.weight"); + CHECK(wk.defined()) + << "[xlite] missing: " << L << "self_attn.indexer.wk.weight"; + CHECK(wproj.defined()) << "[xlite] missing: " << L + << "self_attn.indexer.weights_proj.weight"; + wk = dev(wk, L + "self_attn.indexer.wk.weight"); + wproj = dev(wproj, L + "self_attn.indexer.weights_proj.weight"); + torch::Tensor kw_proj = torch::cat({wk, wproj}, /*dim=*/0).contiguous(); + InitXTensor(m.indexKWeightsProj[i], kw_proj); + storages.push_back(kw_proj); + + InitXTensor( + m.indexKNorm[i], + devBf16(sd.get_tensor(L + "self_attn.indexer.k_norm.weight"), + L + "self_attn.indexer.k_norm.weight")); + // indexKNormBias: cast to BF16 (same as normBias). + torch::Tensor indexKNormBias = + sd.get_tensor(L + "self_attn.indexer.k_norm.bias"); + if (indexKNormBias.defined()) { + torch::Tensor b = indexKNormBias.to(torch::kBFloat16).contiguous(); + InitXTensor(m.indexKNormBias[i], + dev(b, L + "self_attn.indexer.k_norm.bias")); + } + + torch::Tensor wq_b = sd.get_tensor(L + "self_attn.indexer.wq_b.weight"); + CHECK(wq_b.defined()) + << "[xlite] missing: " << L << "self_attn.indexer.wq_b.weight"; + // W8A8 (GLM-5.2 STATIC): .t() + CastNz. indexer not TP-sharded + // (replicated). + bool iqbQuant = IsQuant(wq_b); + if (iqbQuant) { + wq_b = wq_b.t().contiguous(); // [in, out] + } + wq_b = dev(wq_b, L + "self_attn.indexer.wq_b.weight"); + if (iqbQuant && cfg.quantAttnWeightNz) { + wq_b = CastNz(wq_b, L + "self_attn.indexer.wq_b NZ", storages); + } + InitXTensor(m.indexQB[i].weight, wq_b); + storages.push_back(wq_b); + // W8A8 STATIC: indexQB input_dim=indexNHeads*indexHeadDim. indexer + // replicated (no TP shard). + if (iqbQuant) { + int64_t iqInputDim = + static_cast(cfg.indexNHeads) * cfg.indexHeadDim; + torch::Tensor iscale = + sd.get_tensor(L + "self_attn.indexer.wq_b.input_scale"); + torch::Tensor ioffset = + sd.get_tensor(L + "self_attn.indexer.wq_b.input_offset"); + if (iscale.defined()) { + torch::Tensor recip = (1.0f / iscale.to(torch::kFloat32)) + .to(torch::kBFloat16) + .contiguous(); + recip = recip.repeat({iqInputDim}); + InitXTensor(m.indexQB[i].inputScale, + dev(recip, L + "indexer.wq_b.input_scale_reciprocal")); + } + if (ioffset.defined()) { + torch::Tensor off = + ioffset.to(torch::kBFloat16).contiguous().repeat({iqInputDim}); + InitXTensor(m.indexQB[i].inputOffset, + dev(off, L + "indexer.wq_b.input_offset")); + } + torch::Tensor qb_iqb = + sd.get_tensor(L + "self_attn.indexer.wq_b.quant_bias"); + if (qb_iqb.defined()) { + InitXTensor(m.indexQB[i].quantBias, + dev(qb_iqb, L + "indexer.wq_b.quant_bias")); + } + torch::Tensor ds_iqb = + sd.get_tensor(L + "self_attn.indexer.wq_b.deq_scale"); + if (ds_iqb.defined()) { + torch::Tensor ds = TransformDeqScale(ds_iqb); + InitXTensor(m.indexQB[i].deqScale, + dev(ds, L + "indexer.wq_b.deq_scale")); + } + } + } + + if (i == 0 && tp > 1) { + LOG(INFO) << "[xlite] MLA shard layer0 rank=" << rank + << ": qkva=" << qkva.sizes() << " q_b=" << q_b.sizes() + << " o_proj=" << o.sizes(); + } + } + LOG(INFO) << "[xlite] LoadMLAAttn: " << cfg.nLayers << " layers done" + << " (qLora=" << cfg.qLoraRank << " kvLora=" << cfg.kvLoraRank + << " nope=" << cfg.nopeHeadDim << " rope=" << cfg.ropeHeadDim + << " v=" << cfg.vHeadDim + << (cfg.attnType == XMODEL_ATTN_DSA + ? " DSA indexNHeads=" + std::to_string(cfg.indexNHeads) + + " indexHeadDim=" + std::to_string(cfg.indexHeadDim) + : "") + << ")"; + } + + // DeepSeek-V3/R1 (MLA + MoE): MLA attn + dense FFN + MoE experts. + static void LoadMLA(const StateDict& sd, + XModel& m, + const XModelConfig& cfg, + const ModelArgs& args, + const ParallelArgs& pa, + const torch::Device& device, + std::vector& storages) { + LoadMLAAttn(sd, m, cfg, args, pa, device, storages); + LoadMoEFFN(sd, m, cfg, pa, device, storages); + } +}; + +} // namespace xllm::xlite \ No newline at end of file diff --git a/xllm/core/runtime/CMakeLists.txt b/xllm/core/runtime/CMakeLists.txt index ee70ba1905..446a200ef3 100644 --- a/xllm/core/runtime/CMakeLists.txt +++ b/xllm/core/runtime/CMakeLists.txt @@ -19,6 +19,7 @@ cc_library( $<$:acl_graph_bucket_policy.h> $<$:acl_graph_executor_impl.h> $<$:acl_graph_persistent_param.h> + $<$,$>:xlite_executor_impl.h> $<$:mlu_graph_executor_impl.h> $<$:cuda_graph_executor_impl.h> $<$:musa_graph_executor_impl.h> @@ -53,6 +54,7 @@ cc_library( dit_executor.cpp $<$:acl_graph_executor_impl.cpp> $<$:acl_graph_persistent_param.cpp> + $<$,$>:xlite_executor_impl.cpp> $<$:mlu_graph_executor_impl.cpp> $<$:cuda_graph_executor_impl.cpp> $<$:musa_graph_executor_impl.cpp> @@ -118,3 +120,7 @@ cc_library( $<$:torch_musa> $<$:hip::host> ) + +if(USE_NPU AND USE_XLITE) + xllm_link_xlite(runtime) +endif() \ No newline at end of file diff --git a/xllm/core/runtime/executor.cpp b/xllm/core/runtime/executor.cpp index 66195935a4..930234338a 100644 --- a/xllm/core/runtime/executor.cpp +++ b/xllm/core/runtime/executor.cpp @@ -31,6 +31,8 @@ Executor::Executor(CausalLM* model, std::string backend; if (ModelConfig::is_python_model_impl(model_config.model_impl())) { backend = "python"; + } else if (options.npu_kernel_backend() == "XLITE") { + backend = "xlite"; } else if (options.backend() != "vlm" && options.enable_graph()) { backend = Platform::type_str(); } else { diff --git a/xllm/core/runtime/executor_impl_factory.cpp b/xllm/core/runtime/executor_impl_factory.cpp index edca53d945..d4a7dff249 100644 --- a/xllm/core/runtime/executor_impl_factory.cpp +++ b/xllm/core/runtime/executor_impl_factory.cpp @@ -21,6 +21,9 @@ limitations under the License. #include "runtime/vlm_executor_impl.h" #if defined(USE_NPU) #include "runtime/acl_graph_executor_impl.h" +#if defined(USE_XLITE) +#include "runtime/xlite_executor_impl.h" +#endif #elif defined(USE_MLU) #include "runtime/mlu_graph_executor_impl.h" #elif defined(USE_CUDA) diff --git a/xllm/core/runtime/xlite_executor_impl.cpp b/xllm/core/runtime/xlite_executor_impl.cpp new file mode 100644 index 0000000000..1586fdf37a --- /dev/null +++ b/xllm/core/runtime/xlite_executor_impl.cpp @@ -0,0 +1,153 @@ +/* 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. +==============================================================================*/ + +#include "core/runtime/xlite_executor_impl.h" + +#include + +#if defined(USE_NPU) +#include +#endif + +#include "core/framework/model/model_output.h" +#include "core/layers/xlite/xlite_causal_lm_base.h" + +namespace xllm { + +ModelOutput XliteExecutorImpl::run(const torch::Tensor& tokens, + const torch::Tensor& positions, + std::vector& kv_caches, + const ModelInputParams& params) { + auto* holder = model_->get_xlite_holder(); + CHECK(holder != nullptr) << "xlite backend requires xlite-backed model"; + CHECK(holder->xlite_ready()) << "xlite model not initialized"; + + XRuntime& rt = holder->xlite_rt(); + XModel& xlite_model = holder->xlite_model(); + + // DP padding: xlite ForwardMoEDispatch/Combine use fixed-shape + // AllGather/ReduceScatter requiring equal batchedTokens per DP rank. Empty DP + // shards get padded here. + const auto& dp_nums = params.parallel.dp_global_token_nums; + int64_t real_tokens = tokens.size(0); + int64_t max_tokens = real_tokens; + int64_t pad_count = 0; + bool empty_shard = false; + if (dp_nums.size() > 1) { + max_tokens = real_tokens; + for (int32_t v : dp_nums) { + if (static_cast(v) > max_tokens) { + max_tokens = static_cast(v); + } + } + // dp_rank = global_rank / attn_tp, attn_tp = world / dp_size. + const int32_t dp_size = static_cast(dp_nums.size()); + const int32_t world = + options_.world_size() > 0 ? options_.world_size() : dp_size; + const int32_t attn_tp = dp_size > 0 ? world / dp_size : 1; + const int32_t dp_rank = + attn_tp > 0 ? static_cast(options_.server_idx()) / attn_tp : 0; + if (dp_rank >= 0 && dp_rank < dp_size) { + real_tokens = static_cast(dp_nums[dp_rank]); + } + if (real_tokens == 0) { + empty_shard = true; + } + if (max_tokens > real_tokens) { + pad_count = max_tokens - real_tokens; + } + } + + // Pad tokens/positions to max with dummy tokens (token_id=0, position=0). + torch::Tensor run_tokens = tokens; + torch::Tensor run_positions = positions; + if (pad_count > 0) { + auto opts_i = tokens.options(); + auto opts_p = positions.defined() + ? positions.options() + : torch::TensorOptions().dtype(torch::kInt32); + torch::Tensor real_tok = real_tokens > 0 ? tokens.slice(0, 0, real_tokens) + : torch::empty({0}, opts_i); + torch::Tensor real_pos = (real_tokens > 0 && positions.defined()) + ? positions.slice(0, 0, real_tokens) + : torch::empty({0}, opts_p); + torch::Tensor pad_tok = torch::zeros({pad_count}, opts_i); + torch::Tensor pad_pos = torch::zeros({pad_count}, opts_p); + run_tokens = real_tokens > 0 ? torch::cat({real_tok, pad_tok}) : pad_tok; + run_positions = real_tokens > 0 ? torch::cat({real_pos, pad_pos}) : pad_pos; + run_tokens = run_tokens.contiguous(); + run_positions = run_positions.contiguous(); + } + + ::XModelAttnMeta attn_meta; + xlite::XliteAttnMetaBuilder::Build(params, + run_positions, + holder->xlite_config().blockSize, + attn_meta, + pad_count); + + ::XTensor x_in, x_out, x_freqs; + xlite::InitXTensor(x_in, run_tokens); + torch::Tensor out_slice = holder->xlite_output_buf(run_tokens.size(0)); + xlite::InitXTensor(x_out, out_slice); + xlite::InitXTensor(x_freqs, holder->xlite_freqs_cis()); + + if (kv_buf_.size() != kv_caches.size()) { + kv_buf_.resize(kv_caches.size()); + for (auto& layer_kv : kv_buf_) { + layer_kv.resize(2); + } + } + // DSA (GLM-5/5.1): kvCache[layer][2] is indexKCache; framework allocates it + // when index_n_heads>0. MLA/normal models keep 2 entries (k/v). + const bool is_dsa = holder->xlite_config().attnType == XMODEL_ATTN_DSA; + const size_t kv_per_layer = is_dsa ? 3 : 2; + for (auto& layer_kv : kv_buf_) { + layer_kv.resize(kv_per_layer); + } + for (size_t i = 0; i < kv_caches.size(); ++i) { + xlite::InitXTensor(kv_buf_[i][0], kv_caches[i].get_k_cache()); + xlite::InitXTensor(kv_buf_[i][1], kv_caches[i].get_v_cache()); + if (is_dsa) { + xlite::InitXTensor(kv_buf_[i][2], kv_caches[i].get_index_cache()); + } + } + + std::vector<::XTensor> no_deepstack; + // xlite Forward takes freqsCis as a vector (CxA/MHC models use multiple sets; + // GLM-5.x DSA uses only freqsCis[0]). Wrap the single set here. + std::vector<::XTensor> freqs_cis = {x_freqs}; +#if defined(USE_NPU) + aclrtStream ext = c10_npu::getCurrentNPUStream(device_.index()).stream(); + rt.EventWaitCurrStream(ext); + xlite_model.Forward( + rt, x_in, attn_meta, kv_buf_, no_deepstack, freqs_cis, x_out); + rt.EventRecordCurrStream(ext); +#else + LOG(FATAL) << "xlite backend requires USE_NPU"; +#endif + + // Unpad: empty shard returns undefined; real shard returns first real_tokens + // rows. + if (empty_shard) { + return ModelOutput(); + } + if (pad_count > 0 && real_tokens > 0) { + return ModelOutput(out_slice.slice(0, 0, real_tokens)); + } + return ModelOutput(out_slice); +} + +} // namespace xllm \ No newline at end of file diff --git a/xllm/core/runtime/xlite_executor_impl.h b/xllm/core/runtime/xlite_executor_impl.h new file mode 100644 index 0000000000..fb6560be39 --- /dev/null +++ b/xllm/core/runtime/xlite_executor_impl.h @@ -0,0 +1,63 @@ +/* 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. +==============================================================================*/ + +// Drives XModel::Forward via model->get_xlite_holder() (bypasses +// model_->forward). + +#pragma once + +#include + +#include + +#include "core/framework/kv_cache/kv_cache.h" +#include "core/framework/model/causal_lm.h" +#include "core/framework/model/model_input_params.h" +#include "core/framework/model/model_output.h" +#include "core/layers/xlite/xlite_attn_meta_builder.h" +#include "core/layers/xlite/xlite_causal_lm_base.h" +#include "core/layers/xlite/xlite_init_utils.h" +#include "runtime/base_executor_impl.h" +#include "runtime/options.h" + +namespace xllm { + +class XliteExecutorImpl : public BaseExecutorImpl { + public: + XliteExecutorImpl(CausalLM* model, + const ModelArgs& args, + const torch::Device& device, + const runtime::Options& options) + : BaseExecutorImpl(model, args, device, options), + model_(model), + device_(device), + options_(options) {} + + ModelOutput run(const torch::Tensor& tokens, + const torch::Tensor& positions, + std::vector& kv_caches, + const ModelInputParams& params) override; + + private: + CausalLM* model_; + torch::Device device_; + runtime::Options options_; + std::vector> kv_buf_; +}; + +// Header-local so the TU is linked. +REGISTER_EXECUTOR("xlite", XliteExecutorImpl); + +} // namespace xllm \ No newline at end of file diff --git a/xllm/models/CMakeLists.txt b/xllm/models/CMakeLists.txt index b7f90b05e5..69c3eef460 100644 --- a/xllm/models/CMakeLists.txt +++ b/xllm/models/CMakeLists.txt @@ -38,3 +38,7 @@ cc_library( torch_python Python::Python ) + +if(USE_NPU AND USE_XLITE) + xllm_link_xlite(models) +endif() diff --git a/xllm/models/llm/xlite/adapters/deepseek_v3_adapter.h b/xllm/models/llm/xlite/adapters/deepseek_v3_adapter.h new file mode 100644 index 0000000000..df46396240 --- /dev/null +++ b/xllm/models/llm/xlite/adapters/deepseek_v3_adapter.h @@ -0,0 +1,52 @@ +/* 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. +==============================================================================*/ + +// DeepseekV3Adapter: DeepSeek-V3/R1 (MLA + sigmoid MoE + shared expert, no +// DSA). + +#pragma once + +#include "core/layers/xlite/xlite_config_builder.h" +#include "core/layers/xlite/xlite_model_adapter.h" +#include "core/layers/xlite/xlite_weight_utils.h" + +namespace xllm::xlite { + +class DeepseekV3Adapter : public XliteModelAdapter { + public: + XModelConfig BuildConfig(const ModelContext& context) override { + return XliteConfigBuilder::FromDeepseekV3(context); + } + + void Load(const StateDict& sd, + XModel& m, + const XModelConfig& cfg, + const ModelArgs& args, + const ParallelArgs& pa, + const torch::Device& device, + std::vector& storages) override { + XliteWeightUtils::LoadMLA(sd, m, cfg, args, pa, device, storages); + } + + uint32_t MoeTpSize(const ParallelArgs& pa) const override { + return pa.ep_size() > 1 + ? static_cast(pa.world_size() / pa.ep_size()) + : XliteTpSize(pa); + } + + std::string Name() const override { return "DeepseekV3Adapter"; } +}; + +} // namespace xllm::xlite \ No newline at end of file diff --git a/xllm/models/llm/xlite/adapters/glm4_moe_adapter.h b/xllm/models/llm/xlite/adapters/glm4_moe_adapter.h new file mode 100644 index 0000000000..e51648dcc3 --- /dev/null +++ b/xllm/models/llm/xlite/adapters/glm4_moe_adapter.h @@ -0,0 +1,51 @@ +/* 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 + +==============================================================================*/ + +// Glm4MoeAdapter: GLM-4 MoE (MHA + partial rotary + sigmoid MoE + shared +// expert, W8A8). BuildConfig via FromGlm4Moe (W8A8 flags conditional on +// QuantArgs, see IsW8A8); Load -> XliteWeightUtils::LoadMoE (MHA attn + dense +// MLP + MoE experts + shared expert). + +#pragma once + +#include "core/layers/xlite/xlite_config_builder.h" +#include "core/layers/xlite/xlite_model_adapter.h" +#include "core/layers/xlite/xlite_weight_utils.h" + +namespace xllm::xlite { + +class Glm4MoeAdapter : public XliteModelAdapter { + public: + XModelConfig BuildConfig(const ModelContext& context) override { + return XliteConfigBuilder::FromGlm4Moe(context); + } + + void Load(const StateDict& sd, + XModel& m, + const XModelConfig& cfg, + const ModelArgs& args, + const ParallelArgs& pa, + const torch::Device& device, + std::vector& storages) override { + XliteWeightUtils::LoadMoE(sd, m, cfg, args, pa, device, storages); + } + + // Same as Qwen3Moe: EP>1 -> world_size/ep_size; EP==1 -> attention TP. + // XliteCausalLMBase ctor overrides cfg_.moeTPSize with the same formula. + uint32_t MoeTpSize(const ParallelArgs& pa) const override { + return pa.ep_size() > 1 + ? static_cast(pa.world_size() / pa.ep_size()) + : XliteTpSize(pa); + } + + std::string Name() const override { return "Glm4MoeAdapter"; } +}; + +} // namespace xllm::xlite \ No newline at end of file diff --git a/xllm/models/llm/xlite/adapters/glm5_adapter.h b/xllm/models/llm/xlite/adapters/glm5_adapter.h new file mode 100644 index 0000000000..11abc83eef --- /dev/null +++ b/xllm/models/llm/xlite/adapters/glm5_adapter.h @@ -0,0 +1,48 @@ +/* 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 + +==============================================================================*/ + +// Glm5Adapter: GLM-5/5.1 (MLA + DSA indexer + sigmoid MoE + shared expert, +// default rope). GLM5 ~= DeepSeek V3.2 + DSA + different MLA dims; weight +// mapping same as DeepseekV3Adapter. + +#pragma once + +#include "core/layers/xlite/xlite_config_builder.h" +#include "core/layers/xlite/xlite_model_adapter.h" +#include "core/layers/xlite/xlite_weight_utils.h" + +namespace xllm::xlite { + +class Glm5Adapter : public XliteModelAdapter { + public: + XModelConfig BuildConfig(const ModelContext& context) override { + return XliteConfigBuilder::FromGlm5(context); + } + + void Load(const StateDict& sd, + XModel& m, + const XModelConfig& cfg, + const ModelArgs& args, + const ParallelArgs& pa, + const torch::Device& device, + std::vector& storages) override { + XliteWeightUtils::LoadMLA(sd, m, cfg, args, pa, device, storages); + } + + uint32_t MoeTpSize(const ParallelArgs& pa) const override { + return pa.ep_size() > 1 + ? static_cast(pa.world_size() / pa.ep_size()) + : XliteTpSize(pa); + } + + std::string Name() const override { return "Glm5Adapter"; } +}; + +} // namespace xllm::xlite \ No newline at end of file diff --git a/xllm/models/llm/xlite/adapters/qwen3_adapter.h b/xllm/models/llm/xlite/adapters/qwen3_adapter.h new file mode 100644 index 0000000000..2ba11fba23 --- /dev/null +++ b/xllm/models/llm/xlite/adapters/qwen3_adapter.h @@ -0,0 +1,43 @@ +/* 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/layers/xlite/xlite_config_builder.h" +#include "core/layers/xlite/xlite_model_adapter.h" +#include "core/layers/xlite/xlite_weight_utils.h" + +namespace xllm::xlite { + +class Qwen3Adapter : public XliteModelAdapter { + public: + XModelConfig BuildConfig(const ModelContext& context) override { + return XliteConfigBuilder::FromQwen3(context); + } + + void Load(const StateDict& sd, + XModel& m, + const XModelConfig& cfg, + const ModelArgs& args, + const ParallelArgs& pa, + const torch::Device& device, + std::vector& storages) override { + XliteWeightUtils::LoadMHA(sd, m, cfg, args, pa, device, storages); + } + + std::string Name() const override { return "Qwen3Adapter"; } +}; + +} // namespace xllm::xlite \ No newline at end of file diff --git a/xllm/models/llm/xlite/adapters/qwen3_moe_adapter.h b/xllm/models/llm/xlite/adapters/qwen3_moe_adapter.h new file mode 100644 index 0000000000..acf63174e6 --- /dev/null +++ b/xllm/models/llm/xlite/adapters/qwen3_moe_adapter.h @@ -0,0 +1,51 @@ +/* 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/layers/xlite/xlite_config_builder.h" +#include "core/layers/xlite/xlite_model_adapter.h" +#include "core/layers/xlite/xlite_weight_utils.h" + +namespace xllm::xlite { + +class Qwen3MoeAdapter : public XliteModelAdapter { + public: + XModelConfig BuildConfig(const ModelContext& context) override { + return XliteConfigBuilder::FromQwen3Moe(context); + } + + void Load(const StateDict& sd, + XModel& m, + const XModelConfig& cfg, + const ModelArgs& args, + const ParallelArgs& pa, + const torch::Device& device, + std::vector& storages) override { + XliteWeightUtils::LoadMoE(sd, m, cfg, args, pa, device, storages); + } + + // EP>1 -> world_size/ep_size; EP=1 -> tp_size. XliteCausalLMBase ctor also + // sets this. + uint32_t MoeTpSize(const ParallelArgs& pa) const override { + return pa.ep_size() > 1 + ? static_cast(pa.world_size() / pa.ep_size()) + : XliteTpSize(pa); + } + + std::string Name() const override { return "Qwen3MoeAdapter"; } +}; + +} // namespace xllm::xlite \ No newline at end of file diff --git a/xllm/models/llm/xlite/register_all.h b/xllm/models/llm/xlite/register_all.h new file mode 100644 index 0000000000..1709d77e98 --- /dev/null +++ b/xllm/models/llm/xlite/register_all.h @@ -0,0 +1,321 @@ +/* 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. +==============================================================================*/ + +// Register all xlite-backed models as "_xlite". + +#pragma once + +#include "core/layers/xlite/xlite_register_macros.h" +#include "models/llm/xlite/adapters/deepseek_v3_adapter.h" +#include "models/llm/xlite/adapters/glm4_moe_adapter.h" +#include "models/llm/xlite/adapters/glm5_adapter.h" +#include "models/llm/xlite/adapters/qwen3_adapter.h" +#include "models/llm/xlite/adapters/qwen3_moe_adapter.h" + +namespace xllm::xlite { + +// Qwen3 dense (MHA). xlite-specific fields appended for XliteConfigBuilder. +XLITE_REGISTER_MODEL(qwen3, Qwen3Adapter, [&] { + LOAD_ARG_OR(model_type, "model_type", "qwen3"); + LOAD_ARG_OR(dtype, "torch_dtype", ""); + LOAD_ARG_OR(vocab_size, "vocab_size", 152064); + LOAD_ARG_OR(hidden_size, "hidden_size", 3584); + LOAD_ARG_OR(hidden_act, "hidden_act", "silu"); + LOAD_ARG_OR(n_layers, "num_hidden_layers", 28); + LOAD_ARG_OR(n_heads, "num_attention_heads", 28); + LOAD_ARG(n_kv_heads, "num_key_value_heads"); + // head_dim required by KV cache estimation; derive if absent. + LOAD_ARG_OR_FUNC(head_dim, "head_dim", [&] { + return args->hidden_size() / args->n_heads(); + }); + LOAD_ARG_OR(intermediate_size, "intermediate_size", 18944); + LOAD_ARG_OR(max_position_embeddings, "max_position_embeddings", 32768); + LOAD_ARG_OR(rms_norm_eps, "rms_norm_eps", 1e-6); + LOAD_ARG_OR(eos_token_id, "eos_token_id", 151643); + LOAD_ARG_OR(rope_theta, "rope_theta", 1000000.0f); + + // qwen3 < 7B: tie_word_embeddings = true + LOAD_ARG_OR(tie_word_embeddings, "tie_word_embeddings", false); + + LOAD_ARG_OR(use_sliding_window, "use_sliding_window", false); + LOAD_ARG_OR(max_window_layers, "max_window_layers", 28); + + // Eagle3: layer ids to capture, e.g. "layers_to_capture": [2, 14, 25]. + LOAD_ARG_OR(layers_to_capture, "layers_to_capture", std::vector{}); + + // xlite-specific fields for XliteConfigBuilder. + LOAD_ARG_OR(use_qk_norm, "use_qk_norm", true); + LOAD_ARG_OR(qkv_bias, "qkv_bias", false); + SET_ARG(enable_mla, false); + SET_ARG(use_moe, false); + SET_ARG(first_k_dense_replace, 0); + + SET_ARG(stop_token_ids, std::unordered_set({args->eos_token_id()})); +}); + +// DeepSeek-V3/R1 (MLA + sigmoid MoE + shared expert). enable_mla/use_moe via +// SET_ARG. +XLITE_REGISTER_MODEL(deepseek_v3, DeepseekV3Adapter, [&] { + LOAD_ARG_OR(model_type, "model_type", "deepseek_v3"); + LOAD_ARG_OR(dtype, "torch_dtype", ""); + LOAD_ARG_OR(vocab_size, "vocab_size", 129280); + LOAD_ARG_OR(hidden_size, "hidden_size", 7168); + LOAD_ARG_OR(hidden_act, "hidden_act", "silu"); + LOAD_ARG_OR(n_layers, "num_hidden_layers", 61); + LOAD_ARG_OR(n_heads, "num_attention_heads", 128); + LOAD_ARG_OR(n_kv_heads, "num_key_value_heads", 128); // MLA: = n_heads + LOAD_ARG_OR_FUNC(head_dim, "head_dim", [&] { + return args->qk_nope_head_dim() + + args->qk_rope_head_dim(); // MLA: nope+rope + }); + LOAD_ARG_OR(intermediate_size, "intermediate_size", 18432); + LOAD_ARG_OR(max_position_embeddings, "max_position_embeddings", 163840); + LOAD_ARG_OR(rms_norm_eps, "rms_norm_eps", 1e-6); + LOAD_ARG_OR(eos_token_id, "eos_token_id", 1); + LOAD_ARG_OR(bos_token_id, "bos_token_id", 0); + LOAD_ARG_OR(rope_theta, "rope_theta", 10000.0f); + LOAD_ARG_OR(use_sliding_window, "use_sliding_window", false); + LOAD_ARG_OR(sliding_window, "sliding_window", 4096); + LOAD_ARG_OR(max_window_layers, "max_window_layers", 61); + LOAD_ARG_OR(tie_word_embeddings, "tie_word_embeddings", false); + + // MLA fields + LOAD_ARG_OR(qk_nope_head_dim, "qk_nope_head_dim", 128); + LOAD_ARG_OR(qk_rope_head_dim, "qk_rope_head_dim", 64); + LOAD_ARG_OR(v_head_dim, "v_head_dim", 128); + LOAD_ARG_OR(q_lora_rank, "q_lora_rank", 1536); + LOAD_ARG_OR(kv_lora_rank, "kv_lora_rank", 512); + + // MoE fields + LOAD_ARG_OR(first_k_dense_replace, "first_k_dense_replace", 3); + LOAD_ARG_OR(moe_intermediate_size, "moe_intermediate_size", 2048); + LOAD_ARG_OR(num_experts, "n_routed_experts", 256); + LOAD_ARG_OR(n_shared_experts, "n_shared_experts", 1); + LOAD_ARG_OR(num_experts_per_tok, "num_experts_per_tok", 8); + LOAD_ARG_OR(norm_topk_prob, "norm_topk_prob", true); + LOAD_ARG_OR(routed_scaling_factor, "routed_scaling_factor", 2.5f); + LOAD_ARG_OR(n_group, "n_group", 8); + LOAD_ARG_OR(topk_group, "topk_group", 4); + LOAD_ARG_OR(scoring_func, "scoring_func", "sigmoid"); + + // xlite-specific fields (SET_ARG) + SET_ARG(enable_mla, true); + SET_ARG(use_moe, true); + LOAD_ARG_OR( + use_qk_norm, "use_qk_norm", false); // MLA uses mlaQNorm, not MHA qkNorm + SET_ARG(stop_token_ids, std::unordered_set({args->eos_token_id()})); + + // rope_scaling (deepseek_yarn): YaRN freq correction + mscale in + // XliteFreqsCis. + SET_ARG(rope_scaling_rope_type, "deepseek_yarn"); + LOAD_ARG_OR(rope_scaling_factor, "rope_scaling.factor", 40.0f); + LOAD_ARG_OR(rope_scaling_beta_fast, "rope_scaling.beta_fast", 32); + LOAD_ARG_OR(rope_scaling_beta_slow, "rope_scaling.beta_slow", 1); + LOAD_ARG_OR(rope_scaling_original_max_position_embeddings, + "rope_scaling.original_max_position_embeddings", + 4096); + LOAD_ARG_OR(rope_scaling_mscale, "rope_scaling.mscale", 1.0f); + LOAD_ARG_OR(rope_scaling_mscale_all_dim, "rope_scaling.mscale_all_dim", 1.0f); +}); + +// GLM-5/5.1 (MLA + DSA indexer + sigmoid MoE + shared, default rope). +// default rope: no rope_scaling_* -> FromGlm5 skips YaRN/mscale. +XLITE_REGISTER_MODEL(glm_moe_dsa, Glm5Adapter, [&] { + LOAD_ARG_OR(model_type, "model_type", "glm_moe_dsa"); + LOAD_ARG_OR(dtype, "torch_dtype", ""); + LOAD_ARG_OR(vocab_size, "vocab_size", 154880); + LOAD_ARG_OR(hidden_size, "hidden_size", 6144); + LOAD_ARG_OR(hidden_act, "hidden_act", "silu"); + LOAD_ARG_OR(n_layers, "num_hidden_layers", 78); + LOAD_ARG_OR(n_heads, "num_attention_heads", 64); + LOAD_ARG_OR(n_kv_heads, + "num_key_value_heads", + 64); // MLA: c.nKvHeads=1 (MQA); framework uses this + LOAD_ARG_OR_FUNC(head_dim, "head_dim", [&] { + return args->qk_nope_head_dim() + + args->qk_rope_head_dim(); // MLA: nope+rope + }); + LOAD_ARG_OR(intermediate_size, "intermediate_size", 12288); + LOAD_ARG_OR(max_position_embeddings, "max_position_embeddings", 202752); + LOAD_ARG_OR(rms_norm_eps, "rms_norm_eps", 1e-5); + // GLM5 eos is a list [154820,154827,154829]; use SET_ARG (macro can't parse + // braced list). + SET_ARG(eos_token_id_vec, std::vector({154820, 154827, 154829})); + LOAD_ARG_OR(bos_token_id, "bos_token_id", 0); + // rope_theta under rope_parameters (GLM5 has no top-level rope_theta). + LOAD_ARG_OR(rope_theta, "rope_parameters.rope_theta", 1000000.0f); + + LOAD_ARG_OR(use_sliding_window, "use_sliding_window", false); + LOAD_ARG_OR(sliding_window, "sliding_window", 4096); + LOAD_ARG_OR(max_window_layers, "max_window_layers", 61); + LOAD_ARG_OR(tie_word_embeddings, "tie_word_embeddings", false); + + // MLA fields (GLM5: nope=192/rope=64/v=256/q_lora=2048/kv_lora=512) + LOAD_ARG_OR(qk_nope_head_dim, "qk_nope_head_dim", 192); + LOAD_ARG_OR(qk_rope_head_dim, "qk_rope_head_dim", 64); + LOAD_ARG_OR(v_head_dim, "v_head_dim", 256); + LOAD_ARG_OR(q_lora_rank, "q_lora_rank", 2048); + LOAD_ARG_OR(kv_lora_rank, "kv_lora_rank", 512); + + // DSA indexer fields + LOAD_ARG_OR(index_head_dim, "index_head_dim", 128); + LOAD_ARG_OR(index_n_heads, "index_n_heads", 32); + LOAD_ARG_OR(index_topk, "index_topk", 2048); + LOAD_ARG_OR(indexer_rope_interleave, "indexer_rope_interleave", true); + // DSA top-k sharing (GLM-5.2): shared layers skip indexer, reuse prev full + // layer's topkIndices. freq>1 (periodic) or pattern (F/S per-layer). Empty + + // freq<=1 => all layers run indexer (GLM-5.1). + LOAD_ARG_OR(index_topk_freq, "index_topk_freq", 1); + LOAD_ARG_OR(index_topk_pattern, "index_topk_pattern", ""); + LOAD_ARG_OR(index_skip_topk_offset, "index_skip_topk_offset", 0); + + // MoE fields (GLM5: 3 dense / 256 expert / 8 act / 1 shared / group n=1 + // topk=1) + LOAD_ARG_OR(first_k_dense_replace, "first_k_dense_replace", 3); + LOAD_ARG_OR(moe_intermediate_size, "moe_intermediate_size", 2048); + LOAD_ARG_OR(num_experts, "n_routed_experts", 256); + LOAD_ARG_OR(n_shared_experts, "n_shared_experts", 1); + LOAD_ARG_OR(num_experts_per_tok, "num_experts_per_tok", 8); + LOAD_ARG_OR(norm_topk_prob, "norm_topk_prob", true); + LOAD_ARG_OR(routed_scaling_factor, "routed_scaling_factor", 2.5f); + LOAD_ARG_OR(n_group, "n_group", 1); + LOAD_ARG_OR(topk_group, "topk_group", 1); + LOAD_ARG_OR(scoring_func, "scoring_func", "sigmoid"); + + // xlite-specific fields (SET_ARG) + SET_ARG(enable_mla, true); // FromGlm5 sets attnType=DSA (MLA+indexer) + SET_ARG(use_moe, true); + LOAD_ARG_OR( + use_qk_norm, "use_qk_norm", false); // MLA uses mlaQNorm, not MHA qkNorm + SET_ARG(stop_token_ids, + std::unordered_set(args->eos_token_id_vec().begin(), + args->eos_token_id_vec().end())); +}); + +// Qwen3-MoE. xlite-specific fields appended for +// XliteConfigBuilder::FromQwen3Moe. +XLITE_REGISTER_MODEL(qwen3_moe, Qwen3MoeAdapter, [&] { + LOAD_ARG_OR(model_type, "model_type", "qwen3_moe"); + LOAD_ARG_OR(dtype, "torch_dtype", ""); + LOAD_ARG_OR(vocab_size, "vocab_size", 151936); + LOAD_ARG_OR(hidden_size, "hidden_size", 2048); + LOAD_ARG_OR(hidden_act, "hidden_act", "silu"); + LOAD_ARG_OR(n_layers, "num_hidden_layers", 48); + LOAD_ARG_OR(n_heads, "num_attention_heads", 32); + LOAD_ARG(n_kv_heads, "num_key_value_heads"); + // head_dim required by KV cache estimation; derive if absent. + LOAD_ARG_OR_FUNC(head_dim, "head_dim", [&] { + return args->hidden_size() / args->n_heads(); + }); + LOAD_ARG_OR(intermediate_size, "intermediate_size", 6144); + // MoE fields (consumed by FromQwen3Moe). + LOAD_ARG_OR(moe_intermediate_size, "moe_intermediate_size", 768); + LOAD_ARG_OR(num_experts, "num_experts", 128); + LOAD_ARG_OR(num_experts_per_tok, "num_experts_per_tok", 8); + LOAD_ARG_OR(norm_topk_prob, "norm_topk_prob", true); + LOAD_ARG_OR(decoder_sparse_step, "decoder_sparse_step", 1); + // first_k_dense_replace: drives XModelConfig.nDenseLayers. + LOAD_ARG_OR(first_k_dense_replace, "first_k_dense_replace", 0); + + LOAD_ARG_OR(max_position_embeddings, "max_position_embeddings", 40960); + LOAD_ARG_OR(rms_norm_eps, "rms_norm_eps", 1e-6); + LOAD_ARG_OR(eos_token_id, "eos_token_id", 151645); + LOAD_ARG_OR(rope_theta, "rope_theta", 1000000.0f); + + LOAD_ARG_OR(tie_word_embeddings, "tie_word_embeddings", false); + LOAD_ARG_OR(use_sliding_window, "use_sliding_window", false); + LOAD_ARG_OR(max_window_layers, "max_window_layers", 48); + + // Framework-parity fields (harmless if xlite ignores). + LOAD_ARG_OR(output_router_logits, "output_router_logits", false); + LOAD_ARG_OR(router_aux_loss_coef, "router_aux_loss_coef", 0.001f); + LOAD_ARG_OR(mlp_only_layers, "mlp_only_layers", std::vector()); + + // Eagle3: layer ids to capture, e.g. "layers_to_capture": [2, 14, 25]. + LOAD_ARG_OR(layers_to_capture, "layers_to_capture", std::vector{}); + + // xlite-specific fields for XliteConfigBuilder. + LOAD_ARG_OR(use_qk_norm, "use_qk_norm", true); + LOAD_ARG_OR(qkv_bias, "qkv_bias", false); + SET_ARG(enable_mla, false); + SET_ARG(use_moe, true); + + SET_ARG(stop_token_ids, std::unordered_set({args->eos_token_id()})); +}); + +// GLM-4 MoE (MHA + partial rotary + sigmoid MoE + shared expert, W8A8). +// Resolved name: "glm4_moe_xlite". Near Qwen3-MoE but has +// shared/partial_rotary/sigmoid. GLM-4.7 = glm4_moe (GQA MHA + qk_norm + +// qkv_bias + partial rotary + sigmoid MoE + shared). attention_bias=true -> +// qkv_bias. partial_rotary_factor=0.5 -> ropeHeadDim=64 (front 64 dims). +// eos_token_id is a list [151329,151336,151338] -> eos_token_id_vec (same as +// glm5; multi-element vector default brace list breaks macro, use SET_ARG +// parentheses). +XLITE_REGISTER_MODEL(glm4_moe, Glm4MoeAdapter, [&] { + LOAD_ARG_OR(model_type, "model_type", "glm4_moe"); + LOAD_ARG_OR(dtype, "torch_dtype", ""); + LOAD_ARG_OR(vocab_size, "vocab_size", 151552); + LOAD_ARG_OR(hidden_size, "hidden_size", 5120); + LOAD_ARG_OR(hidden_act, "hidden_act", "silu"); + LOAD_ARG_OR(n_layers, "num_hidden_layers", 92); + LOAD_ARG_OR(n_heads, "num_attention_heads", 96); + LOAD_ARG(n_kv_heads, "num_key_value_heads"); + LOAD_ARG_OR_FUNC(head_dim, "head_dim", [&] { + return args->hidden_size() / args->n_heads(); + }); + LOAD_ARG_OR(intermediate_size, + "intermediate_size", + 12288); // dense FFN (first 3 layers) + // MoE fields (consumed by FromGlm4Moe). + LOAD_ARG_OR(moe_intermediate_size, "moe_intermediate_size", 1536); + LOAD_ARG_OR(first_k_dense_replace, "first_k_dense_replace", 3); + LOAD_ARG_OR( + num_experts, "n_routed_experts", 160); // GLM config: n_routed_experts + LOAD_ARG_OR(n_shared_experts, "n_shared_experts", 1); + LOAD_ARG_OR(num_experts_per_tok, "num_experts_per_tok", 8); + LOAD_ARG_OR(norm_topk_prob, "norm_topk_prob", true); + LOAD_ARG_OR(routed_scaling_factor, "routed_scaling_factor", 2.5f); + LOAD_ARG_OR(n_group, "n_group", 1); + LOAD_ARG_OR(topk_group, "topk_group", 1); + LOAD_ARG_OR(partial_rotary_factor, "partial_rotary_factor", 0.5f); + + LOAD_ARG_OR(max_position_embeddings, "max_position_embeddings", 202752); + LOAD_ARG_OR(rms_norm_eps, "rms_norm_eps", 1e-5); + LOAD_ARG_OR(rope_theta, "rope_theta", 1000000.0f); + LOAD_ARG_OR(tie_word_embeddings, "tie_word_embeddings", false); + + LOAD_ARG_OR(use_sliding_window, "use_sliding_window", false); + LOAD_ARG_OR(max_window_layers, "max_window_layers", 92); + + // GLM-4.7: qk_norm=true, attention_bias=true (qkv_bias). config field name is + // attention_bias. + LOAD_ARG_OR(use_qk_norm, "use_qk_norm", true); + LOAD_ARG_OR(qkv_bias, "attention_bias", false); + + // eos_token_id is a list [151329,151336,151338] (generation_config). + // Multi-element vector uses SET_ARG parentheses to avoid macro comma (same as + // glm5). + SET_ARG(eos_token_id_vec, std::vector({151329, 151336, 151338})); + SET_ARG(stop_token_ids, + std::unordered_set(args->eos_token_id_vec().begin(), + args->eos_token_id_vec().end())); + + // xlite-specific fields (ATB does not load; SET_ARG). enable_mla=false (GQA + // MHA, not MLA). + SET_ARG(enable_mla, false); + SET_ARG(use_moe, true); +}); + +} // namespace xllm::xlite \ No newline at end of file diff --git a/xllm/models/model_registry.cpp b/xllm/models/model_registry.cpp index 0f9021687f..494c7c805c 100644 --- a/xllm/models/model_registry.cpp +++ b/xllm/models/model_registry.cpp @@ -68,6 +68,7 @@ namespace { constexpr char kAutoBackend[] = "AUTO"; constexpr char kAtbBackend[] = "ATB"; constexpr char kTorchBackend[] = "TORCH"; +constexpr char kXliteBackend[] = "XLITE"; bool is_torch_only_model_type(const std::string& model_type) { static const std::unordered_set kTorchOnlyModelTypes = { @@ -106,10 +107,10 @@ bool resolve_model_registration(const std::string& model_type, ? kAutoBackend : requested_npu_kernel_backend; if (backend != kAutoBackend && backend != kAtbBackend && - backend != kTorchBackend) { + backend != kTorchBackend && backend != kXliteBackend) { if (error_message != nullptr) { *error_message = "Unsupported --npu_kernel_backend=" + backend + - ". Supported values: AUTO, ATB, TORCH."; + ". Supported values: AUTO, ATB, TORCH, XLITE."; } return false; } @@ -118,6 +119,15 @@ bool resolve_model_registration(const std::string& model_type, if (backend == kAutoBackend) { effective_backend = is_torch_only_model_type(model_type) ? kTorchBackend : kAtbBackend; + } else if (backend == kXliteBackend) { + // Reject torch-only model types + if (is_torch_only_model_type(model_type)) { + if (error_message != nullptr) { + *error_message = + "Model type " + model_type + " only supports TORCH, not XLITE."; + } + return false; + } } else if (model_type == "qwen3" || model_type == "qwen3_moe" || model_type == "deepseek_v32" || model_type == "glm_moe_dsa" || model_type == "qwen3_vl" || model_type == "deepseek_v32_mtp") { @@ -143,7 +153,9 @@ bool resolve_model_registration(const std::string& model_type, if (effective_npu_kernel_backend != nullptr) { *effective_npu_kernel_backend = effective_backend; } - if (model_type == "qwen3" && effective_backend == kAtbBackend) { + if (backend == kXliteBackend) { + *resolved_name = model_type + "_xlite"; + } else if (model_type == "qwen3" && effective_backend == kAtbBackend) { *resolved_name = "qwen3_atb"; } else if (model_type == "qwen3_moe" && effective_backend == kAtbBackend) { *resolved_name = "qwen3_moe_atb"; diff --git a/xllm/models/models.h b/xllm/models/models.h index a07ba8c0ad..ea792ae9ad 100644 --- a/xllm/models/models.h +++ b/xllm/models/models.h @@ -69,6 +69,11 @@ limitations under the License. #include "vlm/qwen3_5.h" // IWYU pragma: keep #include "vlm/qwen3_vl.h" // IWYU pragma: keep +// xlite backend models. +#if defined(USE_XLITE) +#include "llm/xlite/register_all.h" // IWYU pragma: keep +#endif + #elif defined(USE_MLU) #include "dit/pipelines/pipeline_flux.h" // IWYU pragma: keep #include "dit/pipelines/pipeline_flux_control.h" // IWYU pragma: keep