From 42e9621500518f220aba427a4aa0925769f473ff Mon Sep 17 00:00:00 2001 From: cchh05 Date: Thu, 23 Jul 2026 17:06:14 +0800 Subject: [PATCH 1/3] feat(lora): multi-tenant LoRA framework + attention/MLP wire-up (PR 1/2) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Introduces the plumbing for multi-tenant LoRA serving on xllm: composition-based Linear wrappers that ride the existing attention / MLP forward paths and pick up per-batch adapter routing through a thread-local LoRA context. This PR is the first half of a two-PR split. It lands the framework + attention/MLP wire-up so the wrapper code path exists in the tree. PR 2/2 will wire the request path (RequestParams -> Sequence -> BatchInputBuilder -> ModelInputParams.adapter_ids) and the HTTP endpoints (/v1/load_lora_adapter, /v1/unload_lora_adapter, /v1/lora_adapters, /v1/lora_stats). New subsystems -------------- xllm/core/framework/lora/ * LoRARuntime — process singleton: adapter loader, per-proj device pool, hot-swap executor thread pinned to the model device (needed for CANN 8.5's forward-thread-only CPU->NPU copy restriction), per-request per-projection delta lookup. * LoRARegistry — name<->int_id map, pin/unpin lifecycle so /v1/unload_lora_adapter can drain in-flight requests gracefully. * LoRAAdapterLoader — PEFT (adapter_config.json + safetensors) reader with target_modules whitelist and a canonical key parser (base_model.model.*.layers...lora_A|B). * LoRAContext — thread-local frame holding pointers into ModelInputParams.adapter_ids / adapter_ids_per_token so LoRA- wrapped Linear layers can find the per-seq / per-token routing without changing Linear::forward's signature. * LoRAMetrics — bvar Prometheus counters per adapter (TTFT / e2e / tokens generated / errors / QPS). * lora_config — gflags: enable_lora, max_loras, max_lora_rank, lora_target_modules whitelist, lora_modules static preload, allow_runtime_lora_updating, enable_lora_row_parallel_all_reduce (defaults on; correctness), enable_lora_row_parallel_fused_ar (defaults off; a per-layer optimisation that fuses the LoRA delta into the base's own row- parallel all-reduce, S-LoRA / vLLM RowParallelLinearWithShardedLoRA pattern). xllm/core/layers/common/lora/ * LoRAQKVParallelLinear — composition wrapper around QKVParallelLinear. Fast path shares one shrink+expand across the whole single-adapter batch; slow path (mixed base + adapter batch) walks per-seq. Correctly handles GQA replica sharding of the k/v LoRA-B slices — B.size(0) is num_kv_heads * head_dim (not tp_size times that) so the shard-index derivation must divide by the shard count in B, not tp_world_size. Also accepts a q_has_gate parameter for Qwen3-Next-style attn_output_gate where q + gate are fused into the q lane. * LoRAColumnParallelLinear — wrapper for gate_up_proj (fused gate/up) and any future column-parallel projection. * LoRARowParallelLinear — wrapper for o_proj / down_proj. Ships two TP-correct paths: (a) explicit rank-dim all-reduce of the shrink output before the expand, (b) fused mode where the base row- parallel's own all-reduce covers both the base output and the LoRA delta partial-sum. Fused mode reclaims ~3.6 pp single-adapter and ~8.3 pp mixed-batch throughput on Ascend NPU relative to the naive rank-dim AR path, at the cost of holding LoRA-B replicated across TP ranks (cheap because r is small). Wire-up ------- xllm/core/framework/model/model_input_params.h * Adds ModelInputParams::adapter_ids (one entry per sequence in the batch, index-aligned with attention.host.q_seq_lens) and adapter_ids_per_token (device tensor built lazily in to(device) via repeat_interleave from adapter_ids and q_seq_lens). Empty adapter_ ids is a no-op — a pure-base batch does not touch either field. xllm/models/llm/llm_model_base.h * Pushes a LoRAContextFrame at the top of forward, calls set_lora_context_layer on each layer of the decoder loop. Cost when LoRA is off is one atomic pointer copy — the wrappers early- return whenever the frame is null. xllm/core/layers/common/qwen2_attention.{h,cpp}, xllm/core/layers/common/dense_mlp.{h,cpp} * Swaps QKVParallelLinear / RowParallelLinear / ColumnParallelLinear member types for their LoRA-wrapped drop-in equivalents. Base checkpoint keys (qkv_proj.weight, o_proj.weight, gate_up_proj. weight, down_proj.weight) are unchanged because the base linear is held as a plain member inside the wrapper (not register_module'd on the wrapper), so no checkpoint compat break. Not in this PR -------------- * Qwen3-Next hybrid attention (Qwen3.5-122B) wire-up and 122B model install — separate CL, more model-specific work. * Fused MoE expert LoRA delta (Phase 1+2 grouped-gemm injection into fused_moe.cpp) — separate CL, non-trivial MoE-side refactor. * Request path: sequence.adapter_id propagation from RequestParams through BatchInputBuilder to ModelInputParams.adapter_ids — PR 2/2. * HTTP endpoints /v1/load_lora_adapter, /v1/unload_lora_adapter, /v1/lora_adapters, /v1/lora_stats, and the two-field routing in ChatServiceImpl — PR 2/2. Validation (from prior fork builds) ----------------------------------- * End-to-end verified on Qwen3-30B-A3B-Instruct-2507 TP=8 with a real Megatron-trained PEFT LoRA (r=16, alpha=32, target={q,k,v,o}_proj): 200-sample business eval, +9pp compliance-rate and +0.04 gt Jaccard vs base, base output pixel-perfect unchanged (fix is non-invasive). * 30-min sustained load stress: 5363 requests, err=0.00%, HBM +3 MB drift. * Divergence sweep N=50 with a public Qwen3.5-122B LoRA: 70% diverge rate under temperature=0 (systematic delta, not noise). --- third_party/Mooncake | 2 +- third_party/minja | 2 +- third_party/torch_npu_ops | 2 +- third_party/xllm_atb_layers | 2 +- third_party/xllm_ops | 2 +- xllm/core/framework/CMakeLists.txt | 1 + xllm/core/framework/lora/CMakeLists.txt | 38 + xllm/core/framework/lora/adapter_loader.cpp | 248 ++++ xllm/core/framework/lora/adapter_loader.h | 119 ++ xllm/core/framework/lora/lora_config.cpp | 178 +++ xllm/core/framework/lora/lora_config.h | 102 ++ xllm/core/framework/lora/lora_context.cpp | 48 + xllm/core/framework/lora/lora_context.h | 104 ++ xllm/core/framework/lora/lora_metrics.cpp | 295 +++++ xllm/core/framework/lora/lora_metrics.h | 171 +++ xllm/core/framework/lora/lora_registry.cpp | 171 +++ xllm/core/framework/lora/lora_registry.h | 146 +++ xllm/core/framework/lora/lora_request.h | 50 + xllm/core/framework/lora/lora_runtime.cpp | 1058 +++++++++++++++++ xllm/core/framework/lora/lora_runtime.h | 402 +++++++ .../core/framework/model/model_input_params.h | 35 + xllm/core/layers/common/CMakeLists.txt | 7 + xllm/core/layers/common/dense_mlp.cpp | 39 +- xllm/core/layers/common/dense_mlp.h | 6 +- .../lora/lora_column_parallel_linear.cpp | 222 ++++ .../common/lora/lora_column_parallel_linear.h | 86 ++ .../common/lora/lora_qkv_parallel_linear.cpp | 299 +++++ .../common/lora/lora_qkv_parallel_linear.h | 141 +++ .../common/lora/lora_row_parallel_linear.cpp | 248 ++++ .../common/lora/lora_row_parallel_linear.h | 100 ++ xllm/core/layers/common/qwen2_attention.cpp | 40 +- xllm/core/layers/common/qwen2_attention.h | 6 +- xllm/models/llm/llm_model_base.h | 13 + 33 files changed, 4337 insertions(+), 46 deletions(-) create mode 100644 xllm/core/framework/lora/CMakeLists.txt create mode 100644 xllm/core/framework/lora/adapter_loader.cpp create mode 100644 xllm/core/framework/lora/adapter_loader.h create mode 100644 xllm/core/framework/lora/lora_config.cpp create mode 100644 xllm/core/framework/lora/lora_config.h create mode 100644 xllm/core/framework/lora/lora_context.cpp create mode 100644 xllm/core/framework/lora/lora_context.h create mode 100644 xllm/core/framework/lora/lora_metrics.cpp create mode 100644 xllm/core/framework/lora/lora_metrics.h create mode 100644 xllm/core/framework/lora/lora_registry.cpp create mode 100644 xllm/core/framework/lora/lora_registry.h create mode 100644 xllm/core/framework/lora/lora_request.h create mode 100644 xllm/core/framework/lora/lora_runtime.cpp create mode 100644 xllm/core/framework/lora/lora_runtime.h create mode 100644 xllm/core/layers/common/lora/lora_column_parallel_linear.cpp create mode 100644 xllm/core/layers/common/lora/lora_column_parallel_linear.h create mode 100644 xllm/core/layers/common/lora/lora_qkv_parallel_linear.cpp create mode 100644 xllm/core/layers/common/lora/lora_qkv_parallel_linear.h create mode 100644 xllm/core/layers/common/lora/lora_row_parallel_linear.cpp create mode 100644 xllm/core/layers/common/lora/lora_row_parallel_linear.h diff --git a/third_party/Mooncake b/third_party/Mooncake index 1adcb2fb39..8c31e74842 160000 --- a/third_party/Mooncake +++ b/third_party/Mooncake @@ -1 +1 @@ -Subproject commit 1adcb2fb399661e69df0bc4c342c3f4a7028b8e8 +Subproject commit 8c31e74842584c08ce1312b3a0b3347ba597263e diff --git a/third_party/minja b/third_party/minja index 1e65bd7254..1d9d12a48e 160000 --- a/third_party/minja +++ b/third_party/minja @@ -1 +1 @@ -Subproject commit 1e65bd7254aa974d4bc01f9819acec0a0f7bb2d3 +Subproject commit 1d9d12a48e9f7a26b39aca82f4758ccb4893cd02 diff --git a/third_party/torch_npu_ops b/third_party/torch_npu_ops index adc29dc9d2..bf90ef22cc 160000 --- a/third_party/torch_npu_ops +++ b/third_party/torch_npu_ops @@ -1 +1 @@ -Subproject commit adc29dc9d2e69875524b8a54bfc251463275d66e +Subproject commit bf90ef22cc789be1a89541da11d2813ef2c8dd4c diff --git a/third_party/xllm_atb_layers b/third_party/xllm_atb_layers index 015b561ca8..1147537e6d 160000 --- a/third_party/xllm_atb_layers +++ b/third_party/xllm_atb_layers @@ -1 +1 @@ -Subproject commit 015b561ca8df1861eda7f6ccf2ff60a4db8c86b1 +Subproject commit 1147537e6d3137542d6a1fe4c0858f5558e5fd0f diff --git a/third_party/xllm_ops b/third_party/xllm_ops index 818f04b2fc..cc2845f243 160000 --- a/third_party/xllm_ops +++ b/third_party/xllm_ops @@ -1 +1 @@ -Subproject commit 818f04b2fc5aed617fcd2e750936ef7d19956ad5 +Subproject commit cc2845f2432731bc3af7d16edbc5be171ca30c8a diff --git a/xllm/core/framework/CMakeLists.txt b/xllm/core/framework/CMakeLists.txt index 770b5d5345..03bd1d76e0 100644 --- a/xllm/core/framework/CMakeLists.txt +++ b/xllm/core/framework/CMakeLists.txt @@ -24,6 +24,7 @@ add_subdirectory(state_dict) add_subdirectory(tokenizer) add_subdirectory(eplb) add_subdirectory(dit_cache) +add_subdirectory(lora) cc_library( diff --git a/xllm/core/framework/lora/CMakeLists.txt b/xllm/core/framework/lora/CMakeLists.txt new file mode 100644 index 0000000000..c1798c5287 --- /dev/null +++ b/xllm/core/framework/lora/CMakeLists.txt @@ -0,0 +1,38 @@ +include(cc_library) +include(cc_test) + +# Multi-tenant LoRA support. +# +# M1 lora_config - CLI flags + LoRAConfig struct +# M3 lora_request - LoRARequest immutable descriptor +# M3 lora_registry - name<->int_id map + refcount drain semantics +# M2 adapter_loader - PEFT (adapter_config.json + safetensors) reader +# MVP lora_runtime - process singleton wiring loader + registry + +# "currently active" whole-block A/B tensors so +# forward paths can pick them up in one lookup. +# Replaced by proper LoraCacheManager (M4) once +# per-adapter and per-layer routing lands. +cc_library( + NAME + lora + HDRS + lora_config.h + lora_request.h + lora_registry.h + adapter_loader.h + lora_runtime.h + lora_context.h + lora_metrics.h + SRCS + lora_config.cpp + lora_registry.cpp + adapter_loader.cpp + lora_runtime.cpp + lora_context.cpp + lora_metrics.cpp + DEPS + :state_dict + glog::glog + gflags::gflags + torch +) diff --git a/xllm/core/framework/lora/adapter_loader.cpp b/xllm/core/framework/lora/adapter_loader.cpp new file mode 100644 index 0000000000..8306751b5b --- /dev/null +++ b/xllm/core/framework/lora/adapter_loader.cpp @@ -0,0 +1,248 @@ +/* Copyright 2025 The xLLM Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "adapter_loader.h" + +#include + +#include +#include +#include +#include +#include + +#include "framework/state_dict/state_dict.h" + +namespace xllm { + +namespace { +// Fixed prefix HuggingFace PEFT prepends to every weight in +// adapter_model.safetensors. We strip it and then further rewrite the +// remainder into an xllm-friendly key. +constexpr const char* kHfPrefix = "base_model.model."; +constexpr const char* kLoraATag = ".lora_A.weight"; +constexpr const char* kLoraBTag = ".lora_B.weight"; + +bool contains(const std::vector& v, const std::string& item) { + return std::find(v.begin(), v.end(), item) != v.end(); +} +} // namespace + +std::optional LoRAAdapterLoader::load(const LoRARequest& req) { + namespace fs = std::filesystem; + const fs::path root = req.lora_path; + const fs::path config_path = root / "adapter_config.json"; + const fs::path weights_path = root / "adapter_model.safetensors"; + + if (!fs::exists(config_path)) { + LOG(ERROR) << "[LoRAAdapterLoader] '" << req.lora_name + << "': missing adapter_config.json at " << config_path; + return std::nullopt; + } + if (!fs::exists(weights_path)) { + LOG(ERROR) << "[LoRAAdapterLoader] '" << req.lora_name + << "': missing adapter_model.safetensors at " << weights_path; + return std::nullopt; + } + + LoRAAdapter adapter; + adapter.request = req; + + if (!parse_config(config_path.string(), req.base_model_name, &adapter)) { + return std::nullopt; + } + if (!load_weights(weights_path.string(), &adapter)) { + return std::nullopt; + } + + LOG(INFO) << "[LoRAAdapterLoader] loaded '" << req.lora_name + << "' r=" << adapter.r << " scaling=" << adapter.scaling + << " target_modules=" << adapter.target_modules.size() + << " tensors=" << adapter.tensors.size(); + return adapter; +} + +bool LoRAAdapterLoader::parse_config( + const std::string& path, + const std::string& expected_base_model_name, + LoRAAdapter* out) const { + std::ifstream f(path); + if (!f) { + LOG(ERROR) << "[LoRAAdapterLoader] cannot open " << path; + return false; + } + nlohmann::json j; + try { + f >> j; + } catch (const std::exception& e) { + LOG(ERROR) << "[LoRAAdapterLoader] " << path + << ": parse error: " << e.what(); + return false; + } + + // Reject non-LoRA PEFT variants explicitly. DoRA and IA3 need extra + // machinery we don't have yet. + if (j.value("peft_type", std::string("LORA")) != "LORA") { + LOG(ERROR) << "[LoRAAdapterLoader] " << path << ": peft_type != LORA (got '" + << j.value("peft_type", "?") << "'), refusing"; + return false; + } + if (j.value("bias", std::string("none")) != "none") { + LOG(ERROR) << "[LoRAAdapterLoader] " << path << ": bias != none, refusing"; + return false; + } + if (j.value("use_dora", false)) { + LOG(ERROR) << "[LoRAAdapterLoader] " << path << ": use_dora=true, refusing"; + return false; + } + + const int32_t r = j.value("r", 0); + const int32_t lora_alpha = j.value("lora_alpha", 0); + if (r <= 0) { + LOG(ERROR) << "[LoRAAdapterLoader] " << path << ": r=" << r + << " must be > 0"; + return false; + } + if (r > config_.max_lora_rank) { + LOG(ERROR) << "[LoRAAdapterLoader] " << path << ": r=" << r + << " exceeds max_lora_rank=" << config_.max_lora_rank; + return false; + } + out->r = r; + out->scaling = lora_alpha > 0 ? static_cast(lora_alpha) / r : 1.0f; + + // target_modules whitelist check. + if (j.contains("target_modules")) { + for (const auto& m : j["target_modules"]) { + const std::string mod = m.get(); + if (!contains(config_.lora_target_modules, mod)) { + LOG(ERROR) << "[LoRAAdapterLoader] " << path << ": target_modules " + << "contains '" << mod + << "' which is not in --lora-target-modules whitelist"; + return false; + } + out->target_modules.push_back(mod); + } + } + + // Optional base-model check: only enforced if the caller passed one. + if (!expected_base_model_name.empty() && + j.contains("base_model_name_or_path")) { + const std::string got = j["base_model_name_or_path"].get(); + if (got != expected_base_model_name) { + LOG(ERROR) << "[LoRAAdapterLoader] " << path + << ": base_model_name_or_path='" << got + << "' does not match expected='" << expected_base_model_name + << "'"; + return false; + } + } + + return true; +} + +std::optional LoRAAdapterLoader::canonicalize_weight_name( + const std::string& raw) const { + // Strip HuggingFace fixed prefix. + std::string s; + if (raw.rfind(kHfPrefix, 0) == 0) { + s = raw.substr(std::string(kHfPrefix).size()); + } else { + // Some adapters skip the prefix. Tolerate it silently. + s = raw; + } + + // Look for the .lora_A.weight / .lora_B.weight suffix. + const auto a_pos = s.rfind(kLoraATag); + const auto b_pos = s.rfind(kLoraBTag); + const bool is_a = a_pos != std::string::npos && + a_pos + std::string(kLoraATag).size() == s.size(); + const bool is_b = b_pos != std::string::npos && + b_pos + std::string(kLoraBTag).size() == s.size(); + if (!is_a && !is_b) { + return std::nullopt; // Not a LoRA weight (could be embedding LoRA which + // we skip in v1). + } + + // Canonical form: "#A" or "#B" where subkey is the + // dot-path of the module (e.g. "layers.5.self_attn.q_proj"). Using '#' + // as the separator keeps subkey addressable via ordinary regex splits. + const auto pos = is_a ? a_pos : b_pos; + const std::string subkey = s.substr(0, pos); + return subkey + (is_a ? "#A" : "#B"); +} + +bool LoRAAdapterLoader::load_weights(const std::string& safetensors_path, + LoRAAdapter* out) const { + auto sd = StateDictFromSafeTensor::load(safetensors_path); + if (!sd) { + LOG(ERROR) << "[LoRAAdapterLoader] safetensors load failed: " + << safetensors_path; + return false; + } + + int32_t skipped = 0; + int32_t taken = 0; + for (const auto& [raw_name, tensor] : *sd) { + auto canon = canonicalize_weight_name(raw_name); + if (!canon) { + ++skipped; + LOG_EVERY_N(WARNING, 32) + << "[LoRAAdapterLoader] skip unrecognised tensor '" << raw_name + << "' in " << safetensors_path; + continue; + } + // Sanity check: A is [r, ?], B is [?, r]. + const bool is_a = + canon->size() >= 2 && canon->substr(canon->size() - 2) == "#A"; + const auto sizes = tensor.sizes(); + if (sizes.size() != 2) { + LOG(ERROR) << "[LoRAAdapterLoader] '" << raw_name + << "' expected 2D tensor, got sizes.size()=" << sizes.size(); + return false; + } + if (is_a && sizes[0] != out->r) { + LOG(ERROR) << "[LoRAAdapterLoader] '" << raw_name + << "' first dim=" << sizes[0] << " != r=" << out->r; + return false; + } + if (!is_a && sizes[1] != out->r) { + LOG(ERROR) << "[LoRAAdapterLoader] '" << raw_name + << "' second dim=" << sizes[1] << " != r=" << out->r; + return false; + } + // MMAP_CLONE_APPLIED: `tensor` is at::from_blob(mmap_ptr,...) — the mmap + // is owned by `sd` which is a local unique_ptr. Once this function + // returns and sd destructs, the mmap is unmapped and `tensor` points + // into freed memory. Materialize a heap-allocated copy right now. + auto owned = torch::empty( + tensor.sizes(), + torch::TensorOptions().dtype(tensor.dtype()).device(torch::kCPU)); + owned.copy_(tensor); + out->tensors.emplace(*canon, std::move(owned)); + ++taken; + } + + if (taken == 0) { + LOG(ERROR) << "[LoRAAdapterLoader] no LoRA weights found in " + << safetensors_path; + return false; + } + LOG(INFO) << "[LoRAAdapterLoader] " << safetensors_path << " taken=" << taken + << " skipped=" << skipped; + return true; +} + +} // namespace xllm diff --git a/xllm/core/framework/lora/adapter_loader.h b/xllm/core/framework/lora/adapter_loader.h new file mode 100644 index 0000000000..6951a94936 --- /dev/null +++ b/xllm/core/framework/lora/adapter_loader.h @@ -0,0 +1,119 @@ +/* Copyright 2025 The xLLM Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#pragma once + +#include + +#include +#include +#include +#include +#include + +#include "lora_config.h" +#include "lora_request.h" + +namespace xllm { + +// LoRAAdapter is the concrete in-memory result of loading a HuggingFace PEFT +// adapter. It owns: +// - the parsed adapter_config.json (r, alpha, target_modules, ...) +// - one A/B tensor pair per (target_module, layer_idx) +// +// Weights are held in **CPU tensors** at load time. Movement to device is +// LoraCacheManager's job (M4). Keeping the raw CPU state here means we can +// re-materialise on any device without re-parsing safetensors. +struct LoRAAdapter { + // Descriptor mirrored from the load call. + LoRARequest request; + + // Rank r as reported by adapter_config.json. + int32_t r = 0; + + // scale factor pushed into the delta path. Per PEFT convention this is + // `lora_alpha / r`. Stored pre-computed so the forward path just does + // one multiply. + float scaling = 0.0f; + + // Modules the adapter targets, e.g. {"q_proj", "k_proj", "v_proj", + // "o_proj"}. Validated against LoRAConfig::lora_target_modules at load. + std::vector target_modules; + + // Per-(module, layer) A / B weights. + // + // Keyed by "layers.{L}.self_attn.q_proj" (or the closest xllm-native form + // after the module-name mapping step below). Values are the raw torch + // tensors as read from safetensors; A is [r, in_features], B is + // [out_features, r] following HuggingFace convention. + // + // The map is deliberately unstructured (name -> tensor) so future work + // like packed layer merging (q/k/v -> qkv) can rewrite the keys without + // fighting a rigid data structure. + std::unordered_map tensors; +}; + +// LoRAAdapterLoader turns a filesystem path into a LoRAAdapter. It never +// touches device memory and never registers anything. The registry / cache +// call this and route the result themselves. +class LoRAAdapterLoader { + public: + explicit LoRAAdapterLoader(const LoRAConfig& config) : config_(config) {} + + // Load and validate a PEFT adapter from `req.lora_path`. On success the + // returned struct owns all the CPU tensors; on failure the returned + // optional is empty and the reason is logged. + // + // Validation performed: + // 1. adapter_config.json exists and parses. + // 2. r > 0 and r <= config.max_lora_rank. + // 3. base_model_name_or_path (if present) matches req.base_model_name + // when the caller provided one. + // 4. target_modules is a subset of config.lora_target_modules. + // 5. peft_type == "LORA" (DoRA / IA3 refused for now). + // 6. bias == "none" (bias-augmented LoRA refused for now). + // 7. adapter_model.safetensors present and readable. + // + // Any tensor whose name does not fit the expected pattern + // "base_model.model..lora_A.weight" is skipped with a warning + // rather than failing outright, so a slightly non-standard adapter still + // loads whatever it can. + std::optional load(const LoRARequest& req); + + private: + // adapter_config.json parser. Returns false on any structural problem. + bool parse_config(const std::string& path, + const std::string& expected_base_model_name, + LoRAAdapter* out) const; + + // Read `adapter_model.safetensors`, apply the module-name mapping, and + // populate out->tensors. + bool load_weights(const std::string& safetensors_path, + LoRAAdapter* out) const; + + // HF LoRA weight names look like + // base_model.model.model.layers.5.self_attn.q_proj.lora_A.weight + // We strip the fixed "base_model.model." prefix and remap ".lora_A." / + // ".lora_B." tails to a stable canonical form we can match on later. + // + // Returns std::nullopt if the name is unrecognised (skipped rather than + // fatal, per point 7 above). + std::optional canonicalize_weight_name( + const std::string& raw) const; + + const LoRAConfig& config_; +}; + +} // namespace xllm diff --git a/xllm/core/framework/lora/lora_config.cpp b/xllm/core/framework/lora/lora_config.cpp new file mode 100644 index 0000000000..2b2794c311 --- /dev/null +++ b/xllm/core/framework/lora/lora_config.cpp @@ -0,0 +1,178 @@ +/* Copyright 2025 The xLLM Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "lora_config.h" + +#include + +#include +#include +#include +#include + +// Master switch. When off, no other LoRA flag has any effect and the LoRA +// HTTP endpoints return 400. +DEFINE_bool(enable_lora, + false, + "Enable multi-tenant LoRA serving. When false, the engine runs as " + "a plain single-model server."); + +// Number of adapters that can live on GPU simultaneously. Directly caps the +// batch-level active adapter count. +DEFINE_int32(max_loras, + 16, + "Maximum number of LoRA adapters kept on device at the same " + "time. Also caps how many distinct adapters can appear in one " + "batch."); + +// Size of the pinned CPU pool. Should be a proper superset of max_loras so +// that reload from CPU is much cheaper than reload from disk. +DEFINE_int32(max_cpu_loras, + 64, + "Maximum number of LoRA adapters kept in the pinned CPU pool."); + +// Ceiling on r. Loads with r > this are refused. Chosen so 8 / 16 / 32 all +// fit; the largest common practical rank in production is 32. +DEFINE_int32(max_lora_rank, 32, "Upper bound on any single adapter's rank r."); + +// Whitelist of PEFT target_modules the engine will accept. HuggingFace +// adapters that name modules outside this set are refused so we never +// silently ignore a Q/K/V-only vs full attention mismatch. +DEFINE_string( + lora_target_modules, + "qkv_proj,o_proj,gate_up_proj,down_proj,q_proj,k_proj,v_proj,gate_proj,up_" + "proj", + "Comma-separated whitelist of PEFT target_modules that adapters may " + "target. Anything else is refused at load time."); + +// name=path,name=path. Static adapters registered before the server begins +// serving so /v1/models is populated on first request. +DEFINE_string( + lora_modules, + "", + "Comma-separated list of static adapters: name1=path1,name2=path2. " + "Each is loaded before the HTTP server starts."); + +// If false, only the static set is available and POST /v1/load_lora_adapter +// returns 403. Useful for tenants that want a frozen deployment. +DEFINE_bool( + allow_runtime_lora_updating, + true, + "Whether HTTP endpoints /v1/load_lora_adapter and /v1/unload_lora_adapter " + "are honoured at runtime."); + +// Row-parallel LoRA (o_proj / down_proj) requires an extra rank-dim +// all-reduce per layer per forward to keep the delta TP-correct. Under NPU +// HCCL this launch cost is non-trivial and can subtract ~15-20% throughput +// even though the bytes moved are tiny. Set to false for pure-tps deployments +// where the adapters are known to target only column-parallel projections +// (q_proj / k_proj / v_proj / gate_proj / up_proj); the wrapper then skips +// the collective and the row-parallel delta silently no-ops (same behaviour +// as pre-fix). Default true because most PEFT adapters do include o_proj / +// down_proj and precision improvement of ~9pp compliance typically outweighs +// the throughput cost. +DEFINE_bool( + enable_lora_row_parallel_all_reduce, + true, + "Whether LoRARowParallelLinear applies its delta under TP>1 via an extra " + "rank-dim all-reduce. Setting false restores pre-fix behaviour: row-" + "parallel LoRA deltas (o_proj, down_proj) are silently skipped and only " + "column-parallel deltas (q/k/v_proj, gate/up_proj) apply. Trades ~+9pp " + "compliance-rate precision for ~+17% throughput headroom."); + +// Fused variant: sum the LoRA delta into the base's partial [T, out] BEFORE +// the base row-parallel's all-reduce, so both base output and delta ride the +// same collective. Cuts per-layer AR count from 2 (rank-dim + base) to 1 +// (base only), same as vLLM RowParallelLinearWithShardedLoRA / LoRAX / +// HF-TGI. Requires B to be replicated on every rank and A to be sliced on +// in-dim (matches base's shard layout). +// +// When true, this SUPERSEDES enable_lora_row_parallel_all_reduce — the base +// linear is constructed with enable_result_reduction=false and the wrapper +// owns the collective for the fused output. +// +// Default false while we validate. Enabling it should reclaim most of the +// ~17% throughput cost of enable_lora_row_parallel_all_reduce=true without +// giving up the +9pp precision benefit. +DEFINE_bool( + enable_lora_row_parallel_fused_ar, + false, + "Fuse LoRA row-parallel delta into the base's all-reduce (S-LoRA / vLLM " + "RowParallelLinearWithShardedLoRA pattern). Cuts per-layer AR count " + "from 2 to 1. Supersedes enable_lora_row_parallel_all_reduce when true."); + +namespace xllm { + +namespace { +// Split "a,b,,c" into {"a", "b", "c"}. Empty tokens are dropped so a +// dangling trailing comma is tolerated. +std::vector split_comma(const std::string& raw) { + std::vector out; + std::string tok; + std::stringstream ss(raw); + while (std::getline(ss, tok, ',')) { + if (!tok.empty()) { + out.push_back(tok); + } + } + return out; +} +} // namespace + +std::vector> LoRAConfig::parse_modules( + const std::string& raw) { + std::vector> out; + for (const auto& tok : split_comma(raw)) { + // Accept both name=path (vLLM style) and name:path (business style). + // Whichever separator appears FIRST wins so paths containing the other + // char still parse. + size_t sep = std::string::npos; + for (size_t i = 0; i < tok.size(); ++i) { + if (tok[i] == '=' || tok[i] == ':') { + sep = i; + break; + } + } + if (sep == std::string::npos || sep == 0 || sep == tok.size() - 1) { + LOG(ERROR) << "[LoRAConfig] malformed --lora-modules entry '" << tok + << "', expected name=path or name:path"; + continue; + } + out.emplace_back(tok.substr(0, sep), tok.substr(sep + 1)); + } + return out; +} + +void LoRAConfig::load_from_flags() { + enable_lora = FLAGS_enable_lora; + max_loras = FLAGS_max_loras; + max_cpu_loras = FLAGS_max_cpu_loras; + max_lora_rank = FLAGS_max_lora_rank; + lora_target_modules = split_comma(FLAGS_lora_target_modules); + lora_modules = parse_modules(FLAGS_lora_modules); + allow_runtime_lora_updating = FLAGS_allow_runtime_lora_updating; + + if (enable_lora) { + LOG(INFO) << "[LoRAConfig] enabled: max_loras=" << max_loras + << " max_cpu_loras=" << max_cpu_loras + << " max_lora_rank=" << max_lora_rank + << " target_modules=" << lora_target_modules.size() + << " static_modules=" << lora_modules.size() + << " runtime_updates=" + << (allow_runtime_lora_updating ? "on" : "off"); + } +} + +} // namespace xllm diff --git a/xllm/core/framework/lora/lora_config.h b/xllm/core/framework/lora/lora_config.h new file mode 100644 index 0000000000..18dcdcca94 --- /dev/null +++ b/xllm/core/framework/lora/lora_config.h @@ -0,0 +1,102 @@ +/* Copyright 2025 The xLLM Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#pragma once + +#include + +#include +#include +#include +#include + +namespace xllm { + +// --------------------------------------------------------------------------- +// LoRA runtime configuration. +// +// Mirrors the vLLM CLI surface (see reference_multilora_engines memory) so +// the same launch flags work on business gateways that already integrate +// with vLLM. +// +// The struct is populated once at startup from gflags and then held on the +// LLMEngine as an immutable read-only view. Adapter registration / eviction +// happens inside LoRARegistry / LoraCacheManager, not here. +// --------------------------------------------------------------------------- +struct LoRAConfig { + // Master switch. When false, all LoRA APIs report 400 and the model runs + // exactly as baseline (no delta path). + bool enable_lora = false; + + // Maximum number of adapters that may be materialised on a single device + // simultaneously. Directly caps the GPU-side slot pool. vLLM default 4; + // we set 16 since business is 5-20 tenants. + int32_t max_loras = 16; + + // Maximum number of adapters the CPU-pinned pool can hold. Should be a + // superset of max_loras so cache hits are cheap when a request lands on + // an adapter that was recently evicted from GPU. + int32_t max_cpu_loras = 64; + + // Upper bound on r for any adapter. Weights above this are refused at + // load time. Keeps per-slot memory predictable. + int32_t max_lora_rank = 32; + + // Comma-separated module names that adapters may target. The loader + // rejects an adapter whose adapter_config.json target_modules is not a + // subset of this. Kept as a std::vector for cheap membership checks. + std::vector lora_target_modules = { + "qkv_proj", + "o_proj", + "gate_up_proj", + "down_proj", + "q_proj", + "k_proj", + "v_proj", + "gate_proj", + "up_proj", + }; + + // Static adapters registered at startup, one per name=path pair. Each + // pair is loaded before the HTTP server starts serving so /v1/models + // returns them immediately. + std::vector> lora_modules; + + // When true, POST /v1/load_lora_adapter is accepted at runtime. When + // false, only the static --lora-modules set is served. + bool allow_runtime_lora_updating = true; + + // Parse the raw --lora-modules CLI value ("name1=path1,name2=path2") into + // the vector form. Bad tokens are logged and skipped. + static std::vector> parse_modules( + const std::string& raw); + + // Fill this instance from the current gflags values. Idempotent. + void load_from_flags(); +}; + +} // namespace xllm + +// gflags declarations. The DEFINE_* live in lora_config.cpp so the ODR +// stays consistent with the rest of xllm. +DECLARE_bool(enable_lora); +DECLARE_int32(max_loras); +DECLARE_int32(max_cpu_loras); +DECLARE_int32(max_lora_rank); +DECLARE_string(lora_target_modules); +DECLARE_string(lora_modules); +DECLARE_bool(allow_runtime_lora_updating); +DECLARE_bool(enable_lora_row_parallel_all_reduce); +DECLARE_bool(enable_lora_row_parallel_fused_ar); diff --git a/xllm/core/framework/lora/lora_context.cpp b/xllm/core/framework/lora/lora_context.cpp new file mode 100644 index 0000000000..bb1594bab1 --- /dev/null +++ b/xllm/core/framework/lora/lora_context.cpp @@ -0,0 +1,48 @@ +/* Copyright 2026 The xLLM Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "framework/lora/lora_context.h" + +#include + +namespace xllm { +namespace { + +thread_local std::vector tls_stack; + +} // namespace + +void push_lora_context(LoRAContextFrame frame) { + tls_stack.push_back(std::move(frame)); +} + +void pop_lora_context() { + if (!tls_stack.empty()) { + tls_stack.pop_back(); + } +} + +const LoRAContextFrame* current_lora_context() { + if (tls_stack.empty()) return nullptr; + return &tls_stack.back(); +} + +void set_lora_context_layer(int32_t layer_index) { + if (!tls_stack.empty()) { + tls_stack.back().layer_index = layer_index; + } +} + +} // namespace xllm diff --git a/xllm/core/framework/lora/lora_context.h b/xllm/core/framework/lora/lora_context.h new file mode 100644 index 0000000000..995bb0ec9a --- /dev/null +++ b/xllm/core/framework/lora/lora_context.h @@ -0,0 +1,104 @@ +/* Copyright 2026 The xLLM Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +// Thread-local context propagating per-forward LoRA routing info from the +// model forward loop down into individual Linear wrappers. The Linear +// wrappers themselves have a fixed forward(input) signature that can't be +// extended without breaking base-class ABI, so we thread the info through +// TLS instead. +// +// Contract: +// * LlmModelImplBase::forward pushes a Frame at the top of each forward, +// with adapter_ids (per-seq) and q_seq_lens_vec (per-seq). +// * Inside the layer for-loop, it updates Frame::layer_index each +// iteration. +// * LoRA*ParallelLinearImpl::forward reads the current Frame and, if +// any adapter_id is non-zero, pulls the corresponding per-layer +// per-proj (A,B,scaling) from LoRARuntime and applies the delta. +// +// Frames stack (a batch can call into a submodule that itself launches +// more forwards, though in practice we're always one deep). LoRAScope +// RAII pushes/pops. + +#pragma once + +#include + +#include +#include +#include + +namespace xllm { + +struct LoRAContextFrame { + // Per-sequence adapter routing. Empty when the batch is pure-base or + // when LoRA is disabled globally. + const std::vector* adapter_ids = nullptr; + + // Per-sequence lengths, index-aligned with adapter_ids. Used to slice + // the [num_tokens, hidden] hidden state into per-seq chunks. + const std::vector* q_seq_lens_vec = nullptr; + + // Phase A W2 v2: per-token adapter id, device-side int64 [total_tokens]. + // Prepared by ModelInputParams::to(device) at batch-build time by + // expanding adapter_ids according to q_seq_lens_vec. Pointer is valid + // for the duration of the forward pass. nullptr or undefined() when the + // batch is pure-base (all adapter_ids == 0) or LoRA is disabled - + // consumers treat that as the fast path and skip delta entirely. + // + // Consumers (fused_moe.cpp, LoRA wrappers) may safely combine this with + // per_layer combine_idx / cu_seq_lens to derive per-expanded-row adapter + // labels via index_select - all device-side, no CPU->NPU copy in the + // forward loop (CANN 8.5 + torch_npu 2.7.1 forbid that; see 07-06 notes). + const torch::Tensor* adapter_ids_per_token = nullptr; + + // Which decoder layer the forward is currently in (0-indexed). Updated + // by the model's layer loop each iteration. LoRA wrapper uses this to + // pull the right per-layer weight from LoRARuntime. + int32_t layer_index = -1; + + // Which architectural family this model is. For Qwen3 (and Qwen2/ + // oxygen) the PEFT canonical form is "layers.{L}.self_attn.q_proj", + // "layers.{L}.mlp.gate_proj", etc. Other archs may vary; the wrapper + // uses this + a proj-name hint to compose the LoRARuntime lookup key. + std::string arch = "qwen3"; +}; + +// Push a new frame. Returns a token used to pop it. +void push_lora_context(LoRAContextFrame frame); + +// Pop the top frame. Must match the last push (LIFO). +void pop_lora_context(); + +// Read the top frame. Returns nullptr if the stack is empty. +const LoRAContextFrame* current_lora_context(); + +// RAII helper for push/pop. +class LoRAScope { + public: + explicit LoRAScope(LoRAContextFrame frame) { + push_lora_context(std::move(frame)); + } + ~LoRAScope() { pop_lora_context(); } + LoRAScope(const LoRAScope&) = delete; + LoRAScope& operator=(const LoRAScope&) = delete; +}; + +// Update the layer_index on the top frame in-place. Cheaper than a full +// push/pop and matches how the layer loop iterates. No-op if the stack +// is empty. +void set_lora_context_layer(int32_t layer_index); + +} // namespace xllm diff --git a/xllm/core/framework/lora/lora_metrics.cpp b/xllm/core/framework/lora/lora_metrics.cpp new file mode 100644 index 0000000000..101279828d --- /dev/null +++ b/xllm/core/framework/lora/lora_metrics.cpp @@ -0,0 +1,295 @@ +/* Copyright 2026 The xLLM Authors. All Rights Reserved. +Licensed under the Apache License, Version 2.0. +==============================================================================*/ + +#include "lora_metrics.h" + +#include + +namespace xllm { + +namespace { +// Sanitise arbitrary user-provided lora_name into a valid bvar identifier. +// bvar names allow [A-Za-z0-9_], so replace anything else with '_'. +std::string sanitise_bvar_name(const std::string& in) { + std::string out; + out.reserve(in.size()); + for (char c : in) { + if ((c >= 'a' && c <= 'z') || (c >= 'A' && c <= 'Z') || + (c >= '0' && c <= '9') || c == '_') { + out.push_back(c); + } else { + out.push_back('_'); + } + } + return out; +} +} // namespace + +LoRAMetrics& LoRAMetrics::instance() { + static LoRAMetrics g; + return g; +} + +LoRAMetrics::LoRAMetrics() { + ttft_ms_ = std::make_unique>( + "xllm_lora_ttft_milliseconds", std::list{"lora_name"}); + inter_token_ms_ = + std::make_unique>( + "xllm_lora_inter_token_latency_milliseconds", + std::list{"lora_name"}); + e2e_ms_ = std::make_unique>( + "xllm_lora_end_2_end_latency_milliseconds", + std::list{"lora_name"}); + // vLLM #45030 -compatible info gauge with 2 labels. + requests_info_ = std::make_unique>>( + "xllm_lora_requests_info", std::list{"lora_name", "state"}); + active_adapters_count_ = std::make_unique>( + "xllm_lora_active_adapters_count", 0); + // W4-A4 batch-level metrics (global). + batch_fast_path_total_ = + std::make_unique>("xllm_lora_batch_fast_path_total"); + batch_slow_path_total_ = + std::make_unique>("xllm_lora_batch_slow_path_total"); + batch_distinct_adapters_ = std::make_unique( + "xllm_lora_batch_distinct_adapters"); + LOG(INFO) << "[LoRAMetrics] initialised"; +} + +LoRAMetrics::~LoRAMetrics() = default; + +void LoRAMetrics::add_adapter(uint64_t int_id, const std::string& lora_name) { + const std::string safe = sanitise_bvar_name(lora_name); + std::unique_lock lock(mu_); + id_to_name_[int_id] = lora_name; + if (per_adapter_.count(lora_name)) { + // Idempotent: same-name re-register (e.g. after unload+load) reuses + // the existing bvars so historical counters/histograms carry over. + return; + } + LoRAAdapterBvars b; + b.requests_total = std::make_unique>("xllm_lora_" + safe + + "_requests_total"); + b.tokens_prompt_total = std::make_unique>( + "xllm_lora_" + safe + "_tokens_prompt_total"); + b.tokens_generated_total = std::make_unique>( + "xllm_lora_" + safe + "_tokens_generated_total"); + b.errors_total = std::make_unique>("xllm_lora_" + safe + + "_errors_total"); + b.active_state = std::make_unique>( + "xllm_lora_" + safe + "_active_state", 0.0); + b.device_slots = std::make_unique>( + "xllm_lora_" + safe + "_device_slots", 0.0); + per_adapter_.emplace(lora_name, std::move(b)); + LOG(INFO) << "[LoRAMetrics] registered bvars for '" << lora_name + << "' int_id=" << int_id; +} + +void LoRAMetrics::remove_adapter(uint64_t int_id) { + std::unique_lock lock(mu_); + auto id_it = id_to_name_.find(int_id); + if (id_it == id_to_name_.end()) return; + const std::string lora_name = id_it->second; + id_to_name_.erase(id_it); + auto it = per_adapter_.find(lora_name); + if (it == per_adapter_.end()) return; + // Prune the labelled sub-recorders in the shared histograms too so + // /vars stops emitting stale zero-observation rows for this adapter. + if (ttft_ms_) ttft_ms_->delete_stats({lora_name}); + if (inter_token_ms_) inter_token_ms_->delete_stats({lora_name}); + if (e2e_ms_) e2e_ms_->delete_stats({lora_name}); + if (requests_info_) { + requests_info_->delete_stats({lora_name, "running"}); + requests_info_->delete_stats({lora_name, "waiting"}); + } + per_adapter_.erase(it); + LOG(INFO) << "[LoRAMetrics] removed bvars for '" << lora_name + << "' int_id=" << int_id; +} + +// ---- Hot path helpers ----------------------------------------------------- + +namespace { +// Look up name for an adapter_id under shared_lock. Returns empty string +// on miss (caller should treat as "not tracked"). +inline std::string name_of( + const std::unordered_map& id_to_name, + uint64_t adapter_id) { + auto it = id_to_name.find(adapter_id); + if (it == id_to_name.end()) return std::string(); + return it->second; +} +} // namespace + +void LoRAMetrics::observe_ttft(uint64_t adapter_id, int64_t ms) { + if (adapter_id == 0) return; + std::shared_lock lock(mu_); + std::string name = name_of(id_to_name_, adapter_id); + if (name.empty() || !ttft_ms_) return; + auto* rec = ttft_ms_->get_stats({name}); + if (rec) *rec << ms; +} + +void LoRAMetrics::observe_inter_token(uint64_t adapter_id, int64_t ms) { + if (adapter_id == 0) return; + std::shared_lock lock(mu_); + std::string name = name_of(id_to_name_, adapter_id); + if (name.empty() || !inter_token_ms_) return; + auto* rec = inter_token_ms_->get_stats({name}); + if (rec) *rec << ms; +} + +void LoRAMetrics::observe_e2e(uint64_t adapter_id, int64_t ms) { + if (adapter_id == 0) return; + std::shared_lock lock(mu_); + std::string name = name_of(id_to_name_, adapter_id); + if (name.empty() || !e2e_ms_) return; + auto* rec = e2e_ms_->get_stats({name}); + if (rec) *rec << ms; +} + +void LoRAMetrics::inc_requests(uint64_t adapter_id) { + if (adapter_id == 0) return; + std::shared_lock lock(mu_); + std::string name = name_of(id_to_name_, adapter_id); + if (name.empty()) return; + auto it = per_adapter_.find(name); + if (it == per_adapter_.end() || !it->second.requests_total) return; + *it->second.requests_total << 1; +} + +void LoRAMetrics::add_tokens_prompt(uint64_t adapter_id, int64_t n) { + if (adapter_id == 0 || n <= 0) return; + std::shared_lock lock(mu_); + std::string name = name_of(id_to_name_, adapter_id); + if (name.empty()) return; + auto it = per_adapter_.find(name); + if (it == per_adapter_.end() || !it->second.tokens_prompt_total) return; + *it->second.tokens_prompt_total << static_cast(n); +} + +void LoRAMetrics::add_tokens_generated(uint64_t adapter_id, int64_t n) { + if (adapter_id == 0 || n <= 0) return; + std::shared_lock lock(mu_); + std::string name = name_of(id_to_name_, adapter_id); + if (name.empty()) return; + auto it = per_adapter_.find(name); + if (it == per_adapter_.end() || !it->second.tokens_generated_total) return; + *it->second.tokens_generated_total << static_cast(n); +} + +void LoRAMetrics::inc_errors(uint64_t adapter_id) { + if (adapter_id == 0) return; + std::shared_lock lock(mu_); + std::string name = name_of(id_to_name_, adapter_id); + if (name.empty()) return; + auto it = per_adapter_.find(name); + if (it == per_adapter_.end() || !it->second.errors_total) return; + *it->second.errors_total << 1; +} + +void LoRAMetrics::set_state(uint64_t adapter_id, int state) { + if (adapter_id == 0) return; + std::shared_lock lock(mu_); + std::string name = name_of(id_to_name_, adapter_id); + if (name.empty()) return; + auto it = per_adapter_.find(name); + if (it == per_adapter_.end() || !it->second.active_state) return; + it->second.active_state->set_value(static_cast(state)); + // Recompute active_adapters_count_ (state == 1) via a full pass. This + // is O(N adapters) but state changes are rare (load / drain). + int64_t active = 0; + for (const auto& [_, bv] : per_adapter_) { + if (bv.active_state && bv.active_state->get_value() == 1.0) ++active; + } + if (active_adapters_count_) active_adapters_count_->set_value(active); +} + +void LoRAMetrics::set_device_slots(uint64_t adapter_id, int64_t slots) { + if (adapter_id == 0) return; + std::shared_lock lock(mu_); + std::string name = name_of(id_to_name_, adapter_id); + if (name.empty()) return; + auto it = per_adapter_.find(name); + if (it == per_adapter_.end() || !it->second.device_slots) return; + it->second.device_slots->set_value(static_cast(slots)); +} + +void LoRAMetrics::inc_requests_info(uint64_t adapter_id, + const std::string& state) { + if (adapter_id == 0) return; + std::shared_lock lock(mu_); + std::string name = name_of(id_to_name_, adapter_id); + if (name.empty() || !requests_info_) return; + auto* adder = requests_info_->get_stats({name, state}); + if (adder) *adder << 1; +} + +void LoRAMetrics::dec_requests_info(uint64_t adapter_id, + const std::string& state) { + if (adapter_id == 0) return; + std::shared_lock lock(mu_); + std::string name = name_of(id_to_name_, adapter_id); + if (name.empty() || !requests_info_) return; + auto* adder = requests_info_->get_stats({name, state}); + if (adder) *adder << -1; +} + +std::vector LoRAMetrics::snapshot_all() const { + std::vector out; + std::shared_lock lock(mu_); + out.reserve(per_adapter_.size()); + for (const auto& [name, bv] : per_adapter_) { + Snapshot s; + s.lora_name = name; + if (bv.requests_total) s.requests_total = bv.requests_total->get_value(); + if (bv.tokens_prompt_total) + s.tokens_prompt_total = bv.tokens_prompt_total->get_value(); + if (bv.tokens_generated_total) + s.tokens_generated_total = bv.tokens_generated_total->get_value(); + if (bv.errors_total) s.errors_total = bv.errors_total->get_value(); + if (bv.active_state) s.active_state = bv.active_state->get_value(); + if (bv.device_slots) s.device_slots = bv.device_slots->get_value(); + // Latency percentiles: read the MultiDimension sub-recorder if it + // materialised (i.e. at least one observation happened). Skip on miss. + if (ttft_ms_) { + auto* r = ttft_ms_->get_stats({name}); + if (r) { + s.ttft_p50_ms = r->latency_percentile(0.5); + s.ttft_p99_ms = r->latency_percentile(0.99); + } + } + if (e2e_ms_) { + auto* r = e2e_ms_->get_stats({name}); + if (r) { + s.e2e_p50_ms = r->latency_percentile(0.5); + s.e2e_p99_ms = r->latency_percentile(0.99); + s.qps = r->qps(); + } + } + out.push_back(std::move(s)); + } + return out; +} + +size_t LoRAMetrics::adapter_count() const { + std::shared_lock lock(mu_); + return per_adapter_.size(); +} + +// W4-A4: observe a scheduler batch's LoRA distribution. Fast path is +// (0 or 1) distinct adapters, slow path is 2+. total_seqs unused today +// but included for future hit-rate weighting. +void LoRAMetrics::observe_batch_lora(size_t distinct_count, size_t total_seqs) { + (void)total_seqs; + if (batch_distinct_adapters_) { + (*batch_distinct_adapters_) << static_cast(distinct_count); + } + if (distinct_count <= 1) { + if (batch_fast_path_total_) (*batch_fast_path_total_) << 1; + } else { + if (batch_slow_path_total_) (*batch_slow_path_total_) << 1; + } +} + +} // namespace xllm diff --git a/xllm/core/framework/lora/lora_metrics.h b/xllm/core/framework/lora/lora_metrics.h new file mode 100644 index 0000000000..cac460220d --- /dev/null +++ b/xllm/core/framework/lora/lora_metrics.h @@ -0,0 +1,171 @@ +/* Copyright 2026 The xLLM Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +==============================================================================*/ + +// Per-adapter observability metrics for multi-LoRA (P1-D). +// +// Design (2026-07-09 explore report): +// - Latency (TTFT / inter-token / E2E) uses +// bvar::MultiDimension +// with a single "lora_name" label. get_stats({name}) lazily creates a +// sub-recorder on first observation; delete_stats({name}) removes it when +// the adapter is drained. Mirrors the existing dp_rank labelled histograms +// in metrics.cpp so /vars scrape format stays consistent. +// +// - Counters (requests_total, tokens_total, errors_total) and gauges +// (active_state, queue_depth, cache_hit_rate) use heap-allocated +// bvar::Adder / bvar::Status stored in shared_mutex-guarded +// std::unordered_map keyed by lora_name. On adapter load a full set of +// named bvars is created (auto-exposed to brpc built-in /vars endpoint); +// on final removal from LoRARegistry (drain complete) the map slot is +// erased and the bvars are destroyed (auto-hidden from /vars). +// +// - Info line "xllm:lora_requests_info{lora_name,state}" mirrors vLLM #45030 +// for llm-d / router integrations. Emitted via a single global +// MultiDimension keyed by (lora_name, state) so scrapers can see +// both running and waiting counts per adapter without a JSON parse. +// +// All access is thread-safe. Reads take a shared_lock; add_adapter and +// remove_adapter take a unique_lock. Hot-path metric observations use +// the sub-bvars directly under a shared_lock -- <100ns overhead when the +// adapter's entry already exists. + +#pragma once + +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace xllm { + +// Per-adapter mutable counter/gauge bvars. One instance per loaded adapter, +// destroyed on unload_final_removal. +struct LoRAAdapterBvars { + // Counters -- monotonically increasing. + std::unique_ptr> requests_total; + std::unique_ptr> tokens_prompt_total; + std::unique_ptr> tokens_generated_total; + std::unique_ptr> errors_total; + // Gauges -- point-in-time values. + // active_state: 0 = unloaded, 1 = active (device pool populated), 2 = + // draining. + std::unique_ptr> active_state; + // device_slots: how many per_proj slots the adapter currently owns + // (36 layers x 7 projs = 252 for a full Qwen adapter). + std::unique_ptr> device_slots; +}; + +class LoRAMetrics { + public: + static LoRAMetrics& instance(); + + // ---- Lifecycle: called by LoRARegistry on register / final_removal ---- + + // Create the per-adapter bvars and expose them under snake_case names of + // the form "xllm_lora__". Also binds int_id -> name so + // hot-path emit points can index by adapter_id without another map lookup + // through LoRARegistry. + void add_adapter(uint64_t int_id, const std::string& lora_name); + + // Destroy the per-adapter bvars AND the labelled sub-recorders in the + // shared MultiDimension histograms. Idempotent. + void remove_adapter(uint64_t int_id); + + // ---- Hot-path emit points (called from scheduler / response processor) ---- + + // TTFT: observed once per sequence when its first output token materialises. + void observe_ttft(uint64_t adapter_id, int64_t ms); + // Inter-token latency (TPOT): observed for each token after the first. + void observe_inter_token(uint64_t adapter_id, int64_t ms); + // End-to-end latency: observed on request completion. + void observe_e2e(uint64_t adapter_id, int64_t ms); + + // Bump the request/token/error counters. Safe if adapter_id == 0 + // (non-LoRA request) -- becomes a no-op. + void inc_requests(uint64_t adapter_id); + void add_tokens_prompt(uint64_t adapter_id, int64_t n); + void add_tokens_generated(uint64_t adapter_id, int64_t n); + void inc_errors(uint64_t adapter_id); + + // ---- State transitions (called from LoRARuntime) ---- + + // 0 = pending / 1 = active / 2 = draining. Mirrors LoRARegistry entry + // lifetime; used by ops dashboards. + void set_state(uint64_t adapter_id, int state); + void set_device_slots(uint64_t adapter_id, int64_t slots); + + // ---- vLLM #45030 -style info counters ---- + // Emit "xllm:lora_requests_info{lora_name=name,state=running|waiting}". + // Called by scheduler when a sequence transitions running <-> waiting. + void inc_requests_info(uint64_t adapter_id, const std::string& state); + void dec_requests_info(uint64_t adapter_id, const std::string& state); + + // ---- Dump for /v1/lora_stats HTTP endpoint (JSON-friendly snapshot) ---- + + struct Snapshot { + std::string lora_name; + double requests_total = 0; + double tokens_prompt_total = 0; + double tokens_generated_total = 0; + double errors_total = 0; + double active_state = 0; + double device_slots = 0; + int64_t ttft_p50_ms = 0; + int64_t ttft_p99_ms = 0; + int64_t e2e_p50_ms = 0; + int64_t e2e_p99_ms = 0; + int64_t qps = 0; + }; + std::vector snapshot_all() const; + + // W4-A4: batch-level LoRA metrics (global, not per-adapter). Called from + // scheduler prepare_batch once per assembled batch. + // distinct_count = number of distinct non-zero adapter_ids in the batch + // total_seqs = total sequences in the batch (for hit-rate ratio) + // fast_path = (distinct_count <= 1), slow_path = (distinct_count >= 2) + void observe_batch_lora(size_t distinct_count, size_t total_seqs); + + // For test / debug: how many adapters are currently tracked. + size_t adapter_count() const; + + private: + LoRAMetrics(); + ~LoRAMetrics(); + LoRAMetrics(const LoRAMetrics&) = delete; + LoRAMetrics& operator=(const LoRAMetrics&) = delete; + + // Adapter-id -> name lookup so hot-path emit can avoid the LoRARegistry + // mutex. Populated by add_adapter, cleared by remove_adapter. + mutable std::shared_mutex mu_; + std::unordered_map id_to_name_; + std::unordered_map per_adapter_; + + // Shared labelled histograms. Sub-recorders lazily materialised by + // MultiDimension::get_stats({name}) on first observation, and pruned in + // remove_adapter. + std::unique_ptr> ttft_ms_; + std::unique_ptr> inter_token_ms_; + std::unique_ptr> e2e_ms_; + // Requests-info gauge with 2 labels (lora_name, state). + std::unique_ptr>> requests_info_; + + // Global gauge: how many adapters are currently in `active` state. + std::unique_ptr> active_adapters_count_; + + // W4-A4 batch-level counters (global, not per-adapter). + std::unique_ptr> batch_fast_path_total_; + std::unique_ptr> batch_slow_path_total_; + std::unique_ptr batch_distinct_adapters_; +}; + +} // namespace xllm diff --git a/xllm/core/framework/lora/lora_registry.cpp b/xllm/core/framework/lora/lora_registry.cpp new file mode 100644 index 0000000000..e0d9b32e31 --- /dev/null +++ b/xllm/core/framework/lora/lora_registry.cpp @@ -0,0 +1,171 @@ +/* Copyright 2025 The xLLM Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "lora_registry.h" + +#include + +#include +#include + +namespace xllm { + +std::optional LoRARegistry::register_adapter(const LoRARequest& req) { + if (req.lora_name.empty()) { + LOG(ERROR) << "[LoRARegistry] refuse to register empty name"; + return std::nullopt; + } + std::unique_lock lock(mu_); + + auto it = name_to_id_.find(req.lora_name); + if (it != name_to_id_.end()) { + // Name already known. Idempotent path. + const uint64_t existing_id = it->second; + auto& entry = id_to_entry_.at(existing_id); + if (entry.unloading) { + // Drain in progress. Re-registering during drain is legal: revive + // the existing entry (skip renumbering) so ongoing requests keep + // working. This matches vLLM behaviour when the same name is + // reloaded before drain completes. + entry.unloading = false; + LOG(INFO) << "[LoRARegistry] revived '" << req.lora_name + << "' (was draining) id=" << existing_id; + return existing_id; + } + if (entry.req.lora_path != req.lora_path) { + LOG(ERROR) << "[LoRARegistry] name '" << req.lora_name + << "' already bound to '" << entry.req.lora_path + << "', refuse to rebind to '" << req.lora_path + << "' (use load_inplace to swap, P1)"; + return std::nullopt; + } + // Same name, same path -> idempotent, no state change. + return existing_id; + } + + const uint64_t id = next_id_.fetch_add(1); + name_to_id_[req.lora_name] = id; + LoRARequest stored = req; + stored.lora_int_id = id; + id_to_entry_[id] = Entry{stored, /*pin_count=*/0, /*unloading=*/false}; + LOG(INFO) << "[LoRARegistry] registered '" << req.lora_name << "' id=" << id + << " path=" << req.lora_path; + // P1-D: notify observability layer of the new adapter binding. + if (on_register_) on_register_(id, req.lora_name); + return id; +} + +std::optional LoRARegistry::lookup_and_pin( + const std::string& lora_name) { + std::unique_lock lock(mu_); + auto it = name_to_id_.find(lora_name); + if (it == name_to_id_.end()) return std::nullopt; + const uint64_t id = it->second; + auto& entry = id_to_entry_.at(id); + if (entry.unloading) return std::nullopt; + ++entry.pin_count; + return PinnedAdapter{id, entry.req}; +} + +void LoRARegistry::unpin(uint64_t int_id) { + std::unique_lock lock(mu_); + auto it = id_to_entry_.find(int_id); + if (it == id_to_entry_.end()) { + // Should never happen: caller pinned so the entry existed. + LOG(ERROR) << "[LoRARegistry] unpin of unknown id=" << int_id; + return; + } + auto& entry = it->second; + if (entry.pin_count <= 0) { + LOG(ERROR) << "[LoRARegistry] pin_count already <= 0 for id=" << int_id; + return; + } + --entry.pin_count; + if (entry.pin_count == 0 && entry.unloading) { + // Deferred removal completes now. + LOG(INFO) << "[LoRARegistry] drained '" << entry.req.lora_name + << "' id=" << int_id << ", removing"; + // P1-A.3: fire on_final_removal so downstream pools free adapter state. + auto cb = on_final_removal_; + name_to_id_.erase(entry.req.lora_name); + id_to_entry_.erase(it); + if (cb) cb(int_id); + } +} + +std::optional LoRARegistry::lookup( + const std::string& lora_name) const { + std::shared_lock lock(mu_); + auto it = name_to_id_.find(lora_name); + if (it == name_to_id_.end()) return std::nullopt; + const auto& entry = id_to_entry_.at(it->second); + if (entry.unloading) return std::nullopt; + return entry.req; +} + +std::optional LoRARegistry::lookup(uint64_t int_id) const { + std::shared_lock lock(mu_); + auto it = id_to_entry_.find(int_id); + if (it == id_to_entry_.end()) return std::nullopt; + if (it->second.unloading) return std::nullopt; + return it->second.req; +} + +bool LoRARegistry::unregister(const std::string& lora_name) { + std::unique_lock lock(mu_); + auto it = name_to_id_.find(lora_name); + if (it == name_to_id_.end()) return false; + const uint64_t id = it->second; + auto& entry = id_to_entry_.at(id); + if (entry.pin_count == 0) { + // No in-flight work; free immediately. + LOG(INFO) << "[LoRARegistry] unregistered '" << lora_name << "' id=" << id + << " (no pins)"; + // P1-A.3: fire on_final_removal so downstream pools free adapter state. + auto cb = on_final_removal_; + name_to_id_.erase(it); + id_to_entry_.erase(id); + if (cb) cb(id); + return true; + } + // Requests are still using it; caller (unload_lora_adapter handler) is + // expected to poll until the entry disappears. + entry.unloading = true; + LOG(INFO) << "[LoRARegistry] draining '" << lora_name << "' id=" << id + << " pins=" << entry.pin_count; + return true; +} + +bool LoRARegistry::contains(uint64_t int_id) const { + std::shared_lock lock(mu_); + return id_to_entry_.find(int_id) != id_to_entry_.end(); +} + +std::vector LoRARegistry::list() const { + std::shared_lock lock(mu_); + std::vector out; + out.reserve(id_to_entry_.size()); + for (const auto& [_, entry] : id_to_entry_) { + if (!entry.unloading) out.push_back(entry.req); + } + return out; +} + +size_t LoRARegistry::size() const { + std::shared_lock lock(mu_); + return id_to_entry_.size(); +} + +} // namespace xllm diff --git a/xllm/core/framework/lora/lora_registry.h b/xllm/core/framework/lora/lora_registry.h new file mode 100644 index 0000000000..0121904cc2 --- /dev/null +++ b/xllm/core/framework/lora/lora_registry.h @@ -0,0 +1,146 @@ +/* Copyright 2025 The xLLM Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "lora_request.h" + +namespace xllm { + +// LoRARegistry owns the name<->int_id map of loaded adapters and refcounts +// them so unload can drain in-flight requests before releasing weights. +// +// Threading model: +// - Reads (lookup / list / has) hold a shared_lock. +// - Writes (register / unregister) hold a unique_lock. +// This lets a hot request path resolve model="biz-A" -> id concurrently +// while the API handler serialises adapter mutations. +// +// Reference-count semantics (drain-on-unload, mirrors vLLM behaviour): +// - lookup_and_pin() atomically resolves the name and increments refcount. +// Callers MUST call unpin() when the request finishes. +// - unregister() marks the adapter as "unloading"; new lookup_and_pin() +// for that name returns nullopt. The actual weight removal is postponed +// to the moment refcount drops to zero. +class LoRARegistry { + public: + LoRARegistry() = default; + + // Register a name -> id binding. Returns the assigned int_id. + // + // Idempotency (matches vLLM /v1/load_lora_adapter): if `lora_name` is + // already registered and points to the same path, we return the existing + // int_id and no state changes. If the same name is registered with a + // different path, this is a caller error and we log + reject unless the + // caller opted into "load_inplace" (P1). + // + // int_id numbering is monotonically increasing and never reuses freed ids + // for the lifetime of the process, so an adapter that was unloaded and + // later reloaded gets a fresh id. This keeps stale references safe. + std::optional register_adapter(const LoRARequest& req); + + // Return the int_id + full LoRARequest for a given name if the adapter is + // currently loaded and NOT marked for unload. Also increments the pin + // count so the adapter cannot be freed underneath the caller. + // + // Returns std::nullopt if the name is unknown, or if the adapter is being + // drained. + struct PinnedAdapter { + uint64_t int_id; + LoRARequest req; + }; + std::optional lookup_and_pin(const std::string& lora_name); + + // Decrement the pin count of an adapter previously returned by + // lookup_and_pin. When the count hits zero AND the adapter is marked for + // unload, the entry is finally removed. + void unpin(uint64_t int_id); + + // Read-only lookup; does NOT modify refcount. Used by the /v1/models and + // /v1/lora_adapters endpoints to enumerate. + std::optional lookup(const std::string& lora_name) const; + std::optional lookup(uint64_t int_id) const; + + // Mark adapter as unloading. If refcount is zero, remove immediately; + // otherwise defer removal until unpin() drops the last reference. + // + // Returns true if the adapter existed (whether removed synchronously or + // scheduled for drain). + bool unregister(const std::string& lora_name); + + // P1-A.3: register a callback fired at the moment an adapter entry is + // finally erased from the map (either synchronously from unregister when + // refcount is 0, or deferred from unpin when the last ref drops after + // an unloading mark). The callback receives the int_id so downstream + // pools (per_proj_device_pool_, host cache, metric bvars) can free + // their per-adapter state at the same moment the registry forgets it. + // + // Called with the registry mutex held; keep the callback body short. + using FinalRemovalCallback = std::function; + void set_on_final_removal(FinalRemovalCallback cb) { + std::unique_lock lock(mu_); + on_final_removal_ = std::move(cb); + } + + // P1-D: symmetric hook fired when register_adapter first creates an + // entry (idempotent no-ops on re-register skip the callback). Used by + // LoRAMetrics to allocate its per-adapter bvars. Called with the + // registry mutex held; keep the callback body short. + using RegisterCallback = + std::function; + void set_on_register(RegisterCallback cb) { + std::unique_lock lock(mu_); + on_register_ = std::move(cb); + } + + // Enumerate currently-loaded adapter names. Used by /v1/models. + std::vector list() const; + + size_t size() const; + + // P1-A.4: probe whether an int_id still has an entry (may be unloading). + // Used by /v1/unload_lora_adapter drain polling to know when the last + // in-flight request finished and the entry was erased. + bool contains(uint64_t int_id) const; + + private: + struct Entry { + LoRARequest req; + // Live in-flight requests. When >0, the entry cannot be freed. + int32_t pin_count = 0; + // A caller has requested unregister but requests are still using it. + bool unloading = false; + }; + + mutable std::shared_mutex mu_; + std::unordered_map name_to_id_; + std::unordered_map id_to_entry_; + std::atomic next_id_{1}; + FinalRemovalCallback on_final_removal_; + RegisterCallback on_register_; +}; + +} // namespace xllm diff --git a/xllm/core/framework/lora/lora_request.h b/xllm/core/framework/lora/lora_request.h new file mode 100644 index 0000000000..69a9053db5 --- /dev/null +++ b/xllm/core/framework/lora/lora_request.h @@ -0,0 +1,50 @@ +/* Copyright 2025 The xLLM Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#pragma once + +#include +#include + +namespace xllm { + +// LoRARequest is the immutable descriptor of one adapter as it enters the +// engine. It is generated by the HTTP handler (or by the static +// --lora-modules CLI) and lives inside LoRARegistry. +// +// Field naming mirrors vLLM's `vllm/lora/request.py:LoRARequest` so business +// gateways using vLLM contracts port over 1:1. +struct LoRARequest { + // Human-readable adapter name. This is what the user puts in the OpenAI + // `model` field (e.g. "biz-A-v3"). Used for hashing / equality. + std::string lora_name; + + // Engine-assigned monotonically increasing integer id. > 0 always; 0 is + // reserved to mean "no adapter" so batches with mixed base+adapter can + // encode base rows as int_id == 0. + uint64_t lora_int_id = 0; + + // Filesystem or object-storage path where the adapter's PEFT files live. + // The engine expects `adapter_config.json` and `adapter_model.safetensors` + // under this path; validation happens in LoRAAdapterLoader, not here. + std::string lora_path; + + // The name of the base model this adapter was trained against. Compared + // against the running engine's base model at load time so a business + // team cannot accidentally serve a Qwen3 adapter on a Llama base. + std::string base_model_name; +}; + +} // namespace xllm diff --git a/xllm/core/framework/lora/lora_runtime.cpp b/xllm/core/framework/lora/lora_runtime.cpp new file mode 100644 index 0000000000..abd0ce6413 --- /dev/null +++ b/xllm/core/framework/lora/lora_runtime.cpp @@ -0,0 +1,1058 @@ +/* Copyright 2025 The xLLM Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "lora_runtime.h" + +#include + +#include + +#include "lora_metrics.h" + +#if defined(USE_NPU) +#include +#endif + +namespace xllm { + +LoRARuntime& LoRARuntime::instance() { + static LoRARuntime g; + return g; +} + +void LoRARuntime::init(const LoRAConfig& config) { + { + std::lock_guard g(materialise_mu_); + config_ = config; + loader_ = std::make_unique(config_); + } + // P1-A.3: bind the final-removal callback so per_proj_device_pool_ frees + // its per-adapter tensors at the exact moment the registry drops the last + // reference (either synchronous unregister with pin=0, or the deferred + // unpin after a drain). + registry_.set_on_register([](uint64_t int_id, const std::string& lora_name) { + LoRAMetrics::instance().add_adapter(int_id, lora_name); + // Newly registered but not yet installed -> pending (0). Actual + // active_state flip to 1 happens when install_*_on_device succeeds. + LoRAMetrics::instance().set_state(int_id, 0); + }); + registry_.set_on_final_removal([this](uint64_t int_id) { + { + std::lock_guard g(materialise_mu_); + auto it = per_proj_device_pool_.find(int_id); + if (it != per_proj_device_pool_.end()) { + const size_t slot_count = it->second.size(); + per_proj_device_pool_.erase(it); + LOG(INFO) << "[LoRARuntime] freed per_proj device pool for id=" + << int_id << " slots=" << slot_count; + } + // Also drop legacy P0-A dummy pool entry if the same int_id sits there. + auto dp = device_pool_.find(int_id); + if (dp != device_pool_.end()) { + device_pool_.erase(dp); + } + // P1 hot-swap HBM leak fix (C4 stress test 07-21): also erase the + // MoE expert pool. Without this, each hot-swap of a MoE-experts + // adapter (Qwen3-30B-A3B) accumulates ~1.6 GB per rank; ~2 churns + // exhaust NPU 61 GB and crash the process with PTA memory error. + auto mp = moe_expert_lora_pool_.find(int_id); + if (mp != moe_expert_lora_pool_.end()) { + const size_t layers = mp->second.size(); + moe_expert_lora_pool_.erase(mp); + LOG(INFO) << "[LoRARuntime] freed MoE expert pool for id=" << int_id + << " layers=" << layers; + } + } + // Drop metrics bvars AFTER releasing materialise_mu_ (LoRAMetrics + // takes its own shared_mutex; keep locks disjoint). + LoRAMetrics::instance().remove_adapter(int_id); + }); + LOG(INFO) << "[LoRARuntime] initialised, enable=" << config_.enable_lora; +} + +bool LoRARuntime::enabled() const { return config_.enable_lora; } + +void LoRARuntime::set_model_device_dtype(torch::Device device, + torch::ScalarType dtype) { + { + std::lock_guard g(materialise_mu_); + model_device_ = device; + model_dtype_ = dtype; + } + LOG(INFO) << "[LoRARuntime] model registered device=" << device + << " dtype=" << dtype; + + // Path C prod v3 hot-swap: spawn the pinned executor thread here so it + // inherits our aclrtSetDevice state (ModelContext ctor called + // aclrtSetDevice(device_id) earlier on this same thread). + if (!executor_started_) { + executor_started_ = true; + const int32_t device_index = device.index(); + executor_thread_ = + std::thread(&LoRARuntime::executor_loop, this, device_index, dtype); + LOG(INFO) << "[LoRARuntime] spawned hot-swap executor thread for " + << device; + } +} + +// W4-A3: set the hot-swap install config so the pinned executor thread +// knows whether to run per-proj-only or per-proj + moe-experts (and with +// what TP + MoE args). Called from model ctor. +void LoRARuntime::set_hotswap_config(const HotswapConfig& cfg) { + std::lock_guard lk(materialise_mu_); + hotswap_config_ = cfg; + hotswap_config_set_ = true; + LOG(INFO) << "[LoRARuntime] hotswap config set: tp={" << cfg.tp.tp_size << "," + << cfg.tp.tp_rank << "} install_per_proj=" << cfg.install_per_proj + << " install_moe=" << cfg.install_moe_experts + << " num_experts_total=" << cfg.num_experts_total + << " per_rank=" << cfg.num_experts_per_rank + << " start_expert_id=" << cfg.start_expert_id + << " moe_inter=" << cfg.moe_intermediate_size; +} + +void LoRARuntime::executor_loop(int32_t device_index, torch::ScalarType dtype) { +#if defined(USE_NPU) + aclError ret = aclrtSetDevice(device_index); + if (ret != 0) { + LOG(ERROR) << "[LoRARuntime] executor_loop aclrtSetDevice(" << device_index + << ") failed, err=" << ret; + return; + } +#endif + LOG(INFO) << "[LoRARuntime] executor_loop entered on device=" << device_index; + + while (!executor_stop_.load()) { + LoadTask task; + { + std::unique_lock lk(task_mu_); + task_cv_.wait( + lk, [this] { return executor_stop_.load() || !task_queue_.empty(); }); + if (executor_stop_.load() && task_queue_.empty()) break; + task = std::move(task_queue_.front()); + task_queue_.pop(); + } + + LOG(INFO) << "[LoRARuntime] executor picked up task name='" << task.name + << "'"; + torch::Device dev(torch::kPrivateUse1, device_index); + + // W4-A3: read hotswap config, prefer per-proj (+ optional moe experts) + // install. Fall back to whole-block install if config not set (legacy + // Path C prod v3 path). + HotswapConfig cfg; + bool have_cfg = false; + { + std::lock_guard lk(materialise_mu_); + cfg = hotswap_config_; + have_cfg = hotswap_config_set_; + } + + std::optional id_opt; + if (have_cfg && cfg.install_per_proj) { + id_opt = install_static_adapter_on_device_per_proj( + task.name, task.path, task.base_model_name, dev, dtype, cfg.tp); + if (id_opt.has_value()) { + LOG(INFO) << "[LoRARuntime] hot-swap per-proj installed '" << task.name + << "' id=" << *id_opt; + } else { + LOG(WARNING) << "[LoRARuntime] hot-swap per-proj install returned " + << "nullopt for '" << task.name + << "'; MoE-only adapters can still succeed via next step"; + } + // MoE experts install: shares int_id with per-proj via registry + // idempotency. Real MoE fine-tuned adapters have both attn (per-proj) + // and experts tensors; attn-only adapters are handled by per-proj only. + if (cfg.install_moe_experts) { + auto moe_id = + install_static_adapter_on_moe_experts(task.name, + task.path, + task.base_model_name, + dev, + dtype, + cfg.tp, + cfg.num_experts_total, + cfg.num_experts_per_rank, + cfg.start_expert_id, + cfg.moe_intermediate_size); + if (moe_id.has_value()) { + LOG(INFO) << "[LoRARuntime] hot-swap MoE-experts installed '" + << task.name << "' id=" << *moe_id; + // If per-proj skipped (e.g. attn-only refused because adapter + // contains only experts), surface moe id as the result. + if (!id_opt.has_value()) id_opt = moe_id; + } + } + } else { + id_opt = install_static_adapter_on_device( + task.name, task.path, task.base_model_name, dev, dtype); + } + task.result.set_value(id_opt); + } + + LOG(INFO) << "[LoRARuntime] executor_loop exiting"; +} + +std::optional LoRARuntime::load_and_activate_hotswap( + const std::string& lora_name, + const std::string& lora_path, + const std::string& base_model_name) { + if (!enabled()) { + LOG(ERROR) << "[LoRARuntime] not enabled; refuse hot-swap for '" + << lora_name << "'"; + return std::nullopt; + } + if (!executor_started_) { + LOG(ERROR) << "[LoRARuntime] executor not started; the model has not " + "called set_model_device_dtype yet"; + return std::nullopt; + } + + LoadTask task; + task.name = lora_name; + task.path = lora_path; + task.base_model_name = base_model_name; + auto fut = task.result.get_future(); + + { + std::lock_guard lk(task_mu_); + task_queue_.push(std::move(task)); + } + task_cv_.notify_one(); + + LOG(INFO) << "[LoRARuntime] hot-swap enqueued '" << lora_name << "', waiting"; + auto id_opt = fut.get(); + LOG(INFO) << "[LoRARuntime] hot-swap completed '" << lora_name << "' id=" + << (id_opt.has_value() ? std::to_string(*id_opt) : "nullopt"); + return id_opt; +} + +bool LoRARuntime::pick_whole_block_ab(const LoRAAdapter& adapter, + torch::ScalarType dtype, + torch::Tensor* A_out, + torch::Tensor* B_out) const { + // Path C is a *whole-decoder-block* delta -- one A/B for the whole + // model. A real PEFT adapter has A/B per (layer, module). For MVP we + // pick a first (module, layer=0) pair whose shapes match [rank, hidden] + // / [hidden, rank] and treat that as the model-level delta. Real + // per-proj / per-layer routing is P0-C. + // + // Preference order: q_proj > gate_proj > any first pair. + static const std::vector kPreferSubstr = { + "layers.0.self_attn.q_proj", + "self_attn.q_proj", + "layers.0", + }; + auto find_pair = [&](const std::string& subkey_hint, + torch::Tensor* A, + torch::Tensor* B) -> bool { + for (const auto& [key, tensor] : adapter.tensors) { + if (key.size() < 2) continue; + const std::string subkey = key.substr(0, key.size() - 2); + const std::string tail = key.substr(key.size() - 2); + if (!subkey_hint.empty() && subkey.find(subkey_hint) == std::string::npos) + continue; + if (tail == "#A") { + auto it = adapter.tensors.find(subkey + "#B"); + if (it == adapter.tensors.end()) continue; + *A = tensor; + *B = it->second; + return true; + } + } + return false; + }; + for (const auto& hint : kPreferSubstr) { + if (find_pair(hint, A_out, B_out)) break; + } + if (!A_out->defined() || !B_out->defined()) { + if (!find_pair("", A_out, B_out)) { + LOG(ERROR) << "[LoRARuntime] adapter '" << adapter.request.lora_name + << "' has no usable A/B pair"; + return false; + } + } + + // PICK_NO_CAST_APPLIED: skipping in-place .to(dtype) here — safetensors + // mmap tensor + torch::to on CPU crashes torch_npu's opapi allocator on + // ctor thread. Callers (install_static / load_and_activate) do their + // own clone + cast now. + return true; +} + +std::optional LoRARuntime::load_and_activate( + const std::string& lora_name, + const std::string& lora_path, + const std::string& base_model_name) { + if (!enabled()) { + LOG(ERROR) << "[LoRARuntime] not enabled; refuse to load '" << lora_name + << "'"; + return std::nullopt; + } + torch::ScalarType dtype = torch::kFloat32; + { + std::lock_guard g(materialise_mu_); + if (!model_dtype_.has_value()) { + LOG(ERROR) << "[LoRARuntime] model has not registered dtype yet; " + "reject load '" + << lora_name << "'"; + return std::nullopt; + } + dtype = *model_dtype_; + } + + LoRARequest req{lora_name, /*int_id=*/0, lora_path, base_model_name}; + if (!loader_) { + LOG(ERROR) << "[LoRARuntime] loader not initialised"; + return std::nullopt; + } + auto adapter_opt = loader_->load(req); + if (!adapter_opt) return std::nullopt; + + torch::Tensor A_cpu, B_cpu; + if (!pick_whole_block_ab(*adapter_opt, dtype, &A_cpu, &B_cpu)) { + return std::nullopt; + } + + const auto id_opt = registry_.register_adapter(req); + if (!id_opt) return std::nullopt; + + { + std::lock_guard g(materialise_mu_); + pending_ = + PendingDelta{A_cpu, B_cpu, adapter_opt->scaling, lora_name, *id_opt}; + // Clear the previous active adapter -- next forward will materialise + // the new one on device. + active_.reset(); + } + LOG(INFO) << "[LoRARuntime] queued '" << lora_name << "' id=" << *id_opt + << " A_cpu.shape=" << A_cpu.sizes() + << " B_cpu.shape=" << B_cpu.sizes() + << " scaling=" << adapter_opt->scaling + << " (device migration deferred to first forward)"; + return id_opt; +} + +std::optional LoRARuntime::install_static_adapter_on_device( + const std::string& lora_name, + const std::string& lora_path, + const std::string& base_model_name, + torch::Device device, + torch::ScalarType dtype) { + if (!enabled()) { + LOG(ERROR) << "[LoRARuntime] not enabled; refuse to install '" << lora_name + << "'"; + return std::nullopt; + } + if (!loader_) { + LOG(ERROR) << "[LoRARuntime] loader not initialised"; + return std::nullopt; + } + + LoRARequest req{lora_name, /*int_id=*/0, lora_path, base_model_name}; + auto adapter_opt = loader_->load(req); + if (!adapter_opt) return std::nullopt; + + // PICK_BRACKET_APPLIED: log before/after to isolate the crash location. + LOG(INFO) << "[LoRARuntime] BEFORE pick_whole_block_ab for '" << lora_name + << "'"; + torch::Tensor A_cpu, B_cpu; + if (!pick_whole_block_ab(*adapter_opt, dtype, &A_cpu, &B_cpu)) { + LOG(ERROR) << "[LoRARuntime] pick_whole_block_ab returned false"; + return std::nullopt; + } + LOG(INFO) << "[LoRARuntime] AFTER pick_whole_block_ab: A=" << A_cpu.sizes() + << " B=" << B_cpu.sizes() << " A.dtype=" << A_cpu.dtype() + << " A.device=" << A_cpu.device(); + + // KEY DIFFERENCE vs load_and_activate: perform the CPU->NPU migration + // on this thread, which is the model ctor thread and thus has a valid + // NPU context (V60 experiment 2026-07-06 confirmed). + // STATIC_TO_CLONE_APPLIED: safetensors-loaded tensors have mmap backing + // that crashes torch_npu's opapi copy from ctor thread. Clone to a fresh + // cpu allocator buffer, then .to(device). + LOG(INFO) << "[LoRARuntime] static preload '" << lora_name + << "' A_cpu.dtype=" << A_cpu.dtype() << " sizes=" << A_cpu.sizes() + << " storage=" << A_cpu.storage().nbytes() << "B"; + // DIRECT_TO_APPLIED: even .clone() crashes on safetensors mmap tensors on + // ctor thread. Try single-step .to(device, dtype). + // + // Strategy: build a fresh CPU tensor via torch::empty (same allocator as + // V60 randn) and copy_ from mmap tensor - this materializes into normal + // heap. Then .to(device) proceeds. + torch::Tensor A_dev, B_dev; + try { + LOG(INFO) << "[LoRARuntime] materializing A_cpu via torch::empty+copy_"; + auto cpu_opts = + torch::TensorOptions().dtype(A_cpu.dtype()).device(torch::kCPU); + torch::Tensor A_cpu_owned = torch::empty(A_cpu.sizes(), cpu_opts); + A_cpu_owned.copy_(A_cpu); + torch::Tensor B_cpu_owned = torch::empty(B_cpu.sizes(), cpu_opts); + B_cpu_owned.copy_(B_cpu); + LOG(INFO) << "[LoRARuntime] materialized; casting + moving to " << device; + A_dev = A_cpu_owned.to(device, dtype); + LOG(INFO) << "[LoRARuntime] A_dev on " << A_dev.device() + << " dtype=" << A_dev.dtype() << " sizes=" << A_dev.sizes(); + B_dev = B_cpu_owned.to(device, dtype); + LOG(INFO) << "[LoRARuntime] B_dev on " << B_dev.device() + << " dtype=" << B_dev.dtype() << " sizes=" << B_dev.sizes(); + } catch (const std::exception& e) { + LOG(ERROR) << "[LoRARuntime] .to(device) failed for static adapter '" + << lora_name << "': " << e.what(); + return std::nullopt; + } + + const auto id_opt = registry_.register_adapter(req); + if (!id_opt) return std::nullopt; + + { + std::lock_guard g(materialise_mu_); + ActiveDelta ad; + ad.A = A_dev; + ad.B = B_dev; + ad.scaling = adapter_opt->scaling; + ad.name = lora_name; + ad.int_id = *id_opt; + device_pool_[*id_opt] = ad; + active_ = std::move(ad); + // Static preload wins over any pending -- unload() and future + // load_and_activate() calls behave normally. + pending_.reset(); + } + LOG(INFO) << "[LoRARuntime] installed static adapter '" << lora_name + << "' id=" << *id_opt << " A_device=" << A_dev.device() + << " A.shape=" << A_dev.sizes() << " B.shape=" << B_dev.sizes() + << " scaling=" << adapter_opt->scaling; + return id_opt; +} + +// TP-aware overload. For the current whole-block delta path we don't +// actually shard the A/B tensors -- every rank keeps the full copy and +// runs the delta locally on its own slice of h (which atb has already +// all_reduce'd back to full hidden by the time we get here). The +// TPInfo is stored on the resulting ActiveDelta so downstream code +// (and future M10 per-proj work) can see it, but for now the code +// path is identical to the non-TP overload. +std::optional LoRARuntime::install_static_adapter_on_device( + const std::string& lora_name, + const std::string& lora_path, + const std::string& base_model_name, + torch::Device device, + torch::ScalarType dtype, + TPInfo tp) { + if (tp.tp_size > 1) { + LOG(WARNING) << "[LoRARuntime] TP shard skeleton engaged (tp_size=" + << tp.tp_size << " tp_rank=" << tp.tp_rank + << ") but whole-block delta ignores it -- delta stays " + "kReplicated. Per-proj sharding is P1 (M10)."; + } + auto id = install_static_adapter_on_device( + lora_name, lora_path, base_model_name, device, dtype); + if (id.has_value()) { + std::lock_guard g(materialise_mu_); + auto it = device_pool_.find(*id); + if (it != device_pool_.end()) { + it->second.tp = tp; + it->second.shard = LoRAShardStrategy::kReplicated; + } + } + return id; +} + +bool LoRARuntime::unload(const std::string& lora_name) { + bool ok = registry_.unregister(lora_name); + { + std::lock_guard g(materialise_mu_); + if (active_ && active_->name == lora_name) { + active_.reset(); + LOG(INFO) << "[LoRARuntime] deactivated active '" << lora_name << "'"; + } + if (pending_ && pending_->name == lora_name) { + pending_.reset(); + LOG(INFO) << "[LoRARuntime] deactivated pending '" << lora_name << "'"; + } + } + return ok; +} + +std::optional LoRARuntime::active_delta() { + std::lock_guard g(materialise_mu_); + + // Fast path: nothing to do. + if (!pending_ && !active_) return std::nullopt; + + // Promote pending -> active. We DELIBERATELY leave the tensors on CPU + // here even though the model wants them on device: the actual .to(npu) + // has to happen on the worker's forward thread AND under the atb path + // which has aclrtSetDevice set. Doing it here (from the LoRARuntime + // mutex) even on the forward thread crashes with aclrtMemcpy 107017 + // because torch_npu's opapi copy stream is not attached. + // + // The caller (qwen3.h forward loop) does the .to(h.device()) inline: + // that copy runs in the atb-managed NPU stream and works. + if (pending_) { + ActiveDelta ad; + ad.A = pending_->A_cpu; // still on CPU + ad.B = pending_->B_cpu; // still on CPU + ad.scaling = pending_->scaling; + ad.name = pending_->name; + ad.int_id = pending_->int_id; + active_ = std::move(ad); + pending_.reset(); + LOG(INFO) << "[LoRARuntime] promoted '" << active_->name + << "' id=" << active_->int_id + << " (CPU-side; caller performs .to(device))"; + } + return active_; +} + +std::optional LoRARuntime::get_delta_by_int_id( + uint64_t int_id) { + if (int_id == 0) return std::nullopt; + std::lock_guard g(materialise_mu_); + auto it = device_pool_.find(int_id); + if (it == device_pool_.end()) return std::nullopt; + return it->second; +} + +// M10 per-proj installer. Loads full adapter, then walks every tensor in +// adapter.tensors (canonicalized keys like "layers.5.self_attn.q_proj#A"), +// parses (layer_index, proj_name), .to(device), and populates the pool. +std::optional LoRARuntime::install_static_adapter_on_device_per_proj( + const std::string& lora_name, + const std::string& lora_path, + const std::string& base_model_name, + torch::Device device, + torch::ScalarType dtype, + TPInfo tp) { + if (!enabled()) { + LOG(ERROR) << "[LoRARuntime] not enabled; refuse per-proj install '" + << lora_name << "'"; + return std::nullopt; + } + if (!loader_) { + LOG(ERROR) << "[LoRARuntime] loader not initialised"; + return std::nullopt; + } + + LoRARequest req{lora_name, /*int_id=*/0, lora_path, base_model_name}; + auto adapter_opt = loader_->load(req); + if (!adapter_opt) return std::nullopt; + + const auto id_opt = registry_.register_adapter(req); + if (!id_opt) return std::nullopt; + const uint64_t int_id = *id_opt; + + // Phase A W1 A1.3: idempotent guard (symmetric to MoE variant). If this + // int_id already has per-proj attention slots, skip the expensive re-load. + { + std::lock_guard g(materialise_mu_); + if (per_proj_device_pool_.count(int_id) > 0) { + LOG(INFO) << "[LoRARuntime] per-proj install: idempotent skip '" + << lora_name << "' id=" << int_id << " (already installed)"; + return int_id; + } + } + + // Parse canonical keys "layers.{L}.{module_path}.{proj}#A|#B" into + // (layer_index, proj_name) and pair the A/B tensors up. + // + // For Qwen family, module_path values we care about are: + // self_attn.q_proj / self_attn.k_proj / self_attn.v_proj / self_attn.o_proj + // mlp.gate_proj / mlp.up_proj / mlp.down_proj + // We keep only the final proj token in the ProjKey; the parent + // (self_attn vs mlp) is implied by the proj name. + std::unordered_map, + ProjKeyHash> + pairs; + + auto parse_key = [](const std::string& canon) + -> std::optional> { + // Trailing "#A" / "#B" + if (canon.size() < 2) return std::nullopt; + const std::string tail = canon.substr(canon.size() - 2); + if (tail != "#A" && tail != "#B") return std::nullopt; + const bool is_a = (tail == "#A"); + const std::string prefix = canon.substr(0, canon.size() - 2); + // Expect "layers.{L}.." or "model.layers.{L}.." + // (canonicalize_weight_name strips "base_model.model." but the + // remaining "model." from PEFT paths like + // "base_model.model.model.layers.5.self_attn.q_proj.lora_A.weight" + // is left as-is). VL / hybrid models add "model.language_model." prefix, + // e.g. Qwen3.5-122B PEFT adapters use + // "base_model.model.model.language_model.layers.3.self_attn.q_proj...". + const std::string kLayersPrefix = "layers."; + const std::string kModelLayersPrefix = "model.layers."; + const std::string kModelLangLayersPrefix = "model.language_model.layers."; + const std::string kLangLayersPrefix = "language_model.layers."; + std::string after_layers; + if (prefix.rfind(kModelLangLayersPrefix, 0) == 0) { + after_layers = prefix.substr(kModelLangLayersPrefix.size()); + } else if (prefix.rfind(kLangLayersPrefix, 0) == 0) { + after_layers = prefix.substr(kLangLayersPrefix.size()); + } else if (prefix.rfind(kModelLayersPrefix, 0) == 0) { + after_layers = prefix.substr(kModelLayersPrefix.size()); + } else if (prefix.rfind(kLayersPrefix, 0) == 0) { + after_layers = prefix.substr(kLayersPrefix.size()); + } else { + return std::nullopt; + } + const auto dot = after_layers.find('.'); + if (dot == std::string::npos) return std::nullopt; + int32_t layer_index; + try { + layer_index = std::stoi(after_layers.substr(0, dot)); + } catch (...) { + return std::nullopt; + } + // last-dot split: proj is the tail after the final '.' + const auto last_dot = prefix.rfind('.'); + if (last_dot == std::string::npos) return std::nullopt; + const std::string proj = prefix.substr(last_dot + 1); + return std::make_tuple(layer_index, proj, is_a); + }; + + // fix #5: experts LoRA keys must NOT go through the per-proj installer. + // xllm's per_proj_device_pool_ is keyed by (layer, proj_name) and stores + // dense-MLP wrapper deltas. If we let an "layers.N.mlp.experts.M.gate_proj" + // key parse to proj_name="gate_proj", 128 experts would all collide on the + // same key and get silently overwritten, and the FusedMoE forward path + // has no wrapper to consume them anyway --- classic silent no-op. + // Refuse loud so callers know to either drop experts from target_modules + // or route the adapter through install_static_adapter_on_moe_experts. + { + int32_t experts_keys = 0; + std::string first_bad_key; + for (const auto& [key, unused_tensor] : adapter_opt->tensors) { + if (key.find(".experts.") != std::string::npos) { + ++experts_keys; + if (first_bad_key.empty()) first_bad_key = key; + } + } + if (experts_keys > 0) { + LOG(ERROR) << "[LoRARuntime] per-proj installer refuses adapter '" + << lora_name << "': contains " << experts_keys + << " MoE-expert key(s), e.g. '" << first_bad_key << "'. " + << "Per-proj installer only handles dense attention/MLP LoRA. " + << "Either drop experts from adapter target_modules " + << "(attention-only: q/k/v/o_proj) or use " + << "install_static_adapter_on_moe_experts (Phase D)."; + return std::nullopt; + } + } + + for (const auto& [key, tensor] : adapter_opt->tensors) { + auto parsed = parse_key(key); + if (!parsed) continue; + auto [layer_index, proj_name, is_a] = *parsed; + ProjKey pk{layer_index, proj_name}; + auto& slot = pairs[pk]; + if (is_a) + slot.first = tensor; + else + slot.second = tensor; + } + + // Materialize each (A, B) pair on device. + int32_t installed = 0, skipped = 0; + std::unordered_map per_proj; + for (auto& [pk, ab] : pairs) { + if (!ab.first.defined() || !ab.second.defined()) { + ++skipped; + continue; + } + torch::Tensor A_cpu = ab.first; + torch::Tensor B_cpu = ab.second; + + // TP-shard skeleton: for now we keep A/B replicated (whole hidden). + // Per-proj Column/Row sharding is a P1 optimization; the current + // ATB backend doesn't run per-proj anyway. + (void)tp; + + // .to(device, dtype) + torch::Tensor A_dev, B_dev; + try { + A_dev = A_cpu.to(device, dtype).contiguous(); + B_dev = B_cpu.to(device, dtype).contiguous(); + } catch (const std::exception& e) { + LOG(ERROR) << "[LoRARuntime] per-proj .to(device) failed for '" + << lora_name << "' key=layers." << pk.layer_index << "." + << pk.proj_name << ": " << e.what(); + ++skipped; + continue; + } + + ProjDelta pd; + pd.A = A_dev; + pd.B = B_dev; + pd.scaling = adapter_opt->scaling; + pd.r = adapter_opt->r; + per_proj.emplace(pk, std::move(pd)); + ++installed; + } + + { + std::lock_guard g(materialise_mu_); + per_proj_device_pool_[int_id] = std::move(per_proj); + } + LOG(INFO) << "[LoRARuntime] installed per-proj adapter '" << lora_name + << "' id=" << int_id << " slots=" << installed + << " skipped=" << skipped << " r=" << adapter_opt->r + << " scaling=" << adapter_opt->scaling; + // P1-D: adapter is now active on device. + LoRAMetrics::instance().set_state(int_id, 1); + LoRAMetrics::instance().set_device_slots(int_id, installed); + return int_id; +} + +const LoRARuntime::ProjDelta* LoRARuntime::get_per_proj_delta( + uint64_t int_id, + int32_t layer_index, + const std::string& proj_name) { + if (int_id == 0) return nullptr; + std::lock_guard g(materialise_mu_); + auto it = per_proj_device_pool_.find(int_id); + if (it == per_proj_device_pool_.end()) return nullptr; + auto pit = it->second.find(ProjKey{layer_index, proj_name}); + if (pit == it->second.end()) return nullptr; + return &pit->second; +} + +// --------------------------------------------------------------------- +// Day 1 Phase 1: MoE expert LoRA install path. +// --------------------------------------------------------------------- +// +// Parses adapter tensors keyed like +// layers.{L}.mlp.experts.{E}.{proj}#A or #B +// (proj in {gate_proj, up_proj, down_proj}), stacks them across E into +// 3D tensors laid out to match fused_moe.cpp's group_gemm inputs after +// its lazy [E, out, in] -> [E, in, out] transpose, applies TP shard on +// out_dim for gate/up B and in_dim for down A, moves to device, and +// stores in moe_expert_lora_pool_. +// +// Only this rank's experts (start_expert_id .. start_expert_id + +// num_experts_per_rank) are kept. Other experts' A/B tensors are +// dropped after materialize. + +namespace { + +// Parse "layers.{L}.mlp.experts.{E}.{proj}#{A|B}" into (L, E, proj, is_a). +// Returns nullopt for non-expert keys. +struct MoeKey { + int32_t layer_index; + int32_t expert_index; + std::string proj; + bool is_a; +}; + +std::optional parse_moe_expert_key(const std::string& canon) { + if (canon.size() < 2) return std::nullopt; + const std::string tail = canon.substr(canon.size() - 2); + if (tail != "#A" && tail != "#B") return std::nullopt; + const bool is_a = (tail == "#A"); + const std::string prefix = canon.substr(0, canon.size() - 2); + + // Strip optional "model." / "model.language_model." / "language_model." + // leader (VL / hybrid models add these). + std::string p = prefix; + if (p.rfind("model.language_model.layers.", 0) == 0) { + p = p.substr(std::string("model.language_model.").size()); + } else if (p.rfind("language_model.layers.", 0) == 0) { + p = p.substr(std::string("language_model.").size()); + } else if (p.rfind("model.layers.", 0) == 0) { + p = p.substr(std::string("model.").size()); + } + if (p.rfind("layers.", 0) != 0) return std::nullopt; + p = p.substr(std::string("layers.").size()); // "{L}.mlp.experts.{E}.{proj}" + + auto dot1 = p.find('.'); + if (dot1 == std::string::npos) return std::nullopt; + int32_t layer; + try { + layer = std::stoi(p.substr(0, dot1)); + } catch (...) { + return std::nullopt; + } + const std::string rest = p.substr(dot1 + 1); // "mlp.experts.{E}.{proj}" + + // We require the exact "mlp.experts." prefix. Some adapters may use + // "mlp.experts." vs "mlp.experts_" depending on training toolchain; + // the mainline PEFT format uses the former. + const std::string kExpertsPrefix = "mlp.experts."; + if (rest.rfind(kExpertsPrefix, 0) != 0) return std::nullopt; + const std::string rest2 = rest.substr(kExpertsPrefix.size()); // "{E}.{proj}" + + auto dot2 = rest2.find('.'); + if (dot2 == std::string::npos) return std::nullopt; + int32_t expert; + try { + expert = std::stoi(rest2.substr(0, dot2)); + } catch (...) { + return std::nullopt; + } + const std::string proj = rest2.substr(dot2 + 1); + if (proj != "gate_proj" && proj != "up_proj" && proj != "down_proj") { + return std::nullopt; + } + + return MoeKey{layer, expert, proj, is_a}; +} + +} // namespace + +std::optional LoRARuntime::install_static_adapter_on_moe_experts( + const std::string& lora_name, + const std::string& lora_path, + const std::string& base_model_name, + torch::Device device, + torch::ScalarType dtype, + TPInfo tp, + int32_t num_experts_total, + int32_t num_experts_per_rank, + int32_t start_expert_id, + int32_t intermediate_size) { + if (!enabled()) { + LOG(ERROR) << "[LoRARuntime] not enabled; refuse MoE install '" << lora_name + << "'"; + return std::nullopt; + } + if (!loader_) { + LOG(ERROR) << "[LoRARuntime] loader not initialised"; + return std::nullopt; + } + + LoRARequest req{lora_name, /*int_id=*/0, lora_path, base_model_name}; + auto adapter_opt = loader_->load(req); + if (!adapter_opt) return std::nullopt; + + // Register (or reuse) the adapter name so per_proj + moe paths share + // the same int_id. + // register_adapter is idempotent when name+path match. + const auto id_opt = registry_.register_adapter(req); + if (!id_opt) { + LOG(ERROR) << "[LoRARuntime] MoE install: register failed for '" + << lora_name << "'"; + return std::nullopt; + } + const uint64_t int_id = *id_opt; + + // Phase A W1 A1.3: idempotent guard. If this int_id already has MoE experts + // installed, skip the expensive re-load. Same name+path caller (via + // register_adapter idempotency) sees the pool intact; genuine re-load + // requires unload first. + { + std::lock_guard g(materialise_mu_); + if (moe_expert_lora_pool_.count(int_id) > 0) { + LOG(INFO) << "[LoRARuntime] MoE install: idempotent skip '" << lora_name + << "' id=" << int_id << " (already installed)"; + return int_id; + } + } + + // Bucket adapter tensors by (layer, proj) -> vec. + struct Slot { + int32_t expert_index = -1; + bool is_a = false; + torch::Tensor tensor; + }; + std::unordered_map< + int32_t /*layer*/, + std::unordered_map>> + by_layer_proj; + + int32_t total_expert_tensors = 0; + for (const auto& [key, tensor] : adapter_opt->tensors) { + auto parsed = parse_moe_expert_key(key); + if (!parsed) continue; + ++total_expert_tensors; + // Keep only this rank's experts. + if (parsed->expert_index < start_expert_id || + parsed->expert_index >= start_expert_id + num_experts_per_rank) { + continue; + } + Slot slot{parsed->expert_index, parsed->is_a, tensor}; + by_layer_proj[parsed->layer_index][parsed->proj].push_back(std::move(slot)); + } + + if (by_layer_proj.empty()) { + LOG(INFO) << "[LoRARuntime] MoE install '" << lora_name + << "': no experts.* tensors matched this rank (start_expert_id=" + << start_expert_id << " num_per_rank=" << num_experts_per_rank + << " total_expert_tensors_seen=" << total_expert_tensors << ")"; + return int_id; + } + + const int32_t r = adapter_opt->r; + const float scaling = adapter_opt->scaling; + const int32_t hidden = -1; // inferred from A tensor row count below. + const int32_t tp_world = std::max(1, tp.tp_size); + const int32_t tp_rank = std::max(0, tp.tp_rank); + const int32_t local_intermediate = intermediate_size / tp_world; + + auto stack_experts = [&](const std::vector& slots, + bool is_a) -> torch::Tensor { + // Build a map expert_index -> tensor, ordered by expert_index within + // this rank (we already filtered to this rank's expert range). + std::vector filtered; + for (const auto& s : slots) { + if (s.is_a != is_a) continue; + filtered.push_back(s); + } + if (filtered.empty()) return torch::Tensor(); + std::sort( + filtered.begin(), filtered.end(), [](const Slot& a, const Slot& b) { + return a.expert_index < b.expert_index; + }); + if (static_cast(filtered.size()) != num_experts_per_rank) { + LOG(WARNING) << "[LoRARuntime] MoE install '" << lora_name << "': got " + << filtered.size() + << " expert tensors for a slot but expected " + << num_experts_per_rank << " (is_a=" << is_a + << "); skipping this slot"; + return torch::Tensor(); + } + std::vector tensors; + tensors.reserve(filtered.size()); + for (auto& s : filtered) tensors.push_back(s.tensor); + // Stack across expert dim -> [E_per_rank, ...]. + return torch::stack(tensors, /*dim=*/0); + }; + + int32_t layers_installed = 0; + int32_t layers_skipped = 0; + std::unordered_map pool_layers; + + for (auto& [layer, projs] : by_layer_proj) { + MoeExpertDelta med; + med.r = r; + med.scaling = scaling; + + // gate_proj. + if (auto it = projs.find("gate_proj"); it != projs.end()) { + auto A_stacked = stack_experts(it->second, /*is_a=*/true); + auto B_stacked = stack_experts(it->second, /*is_a=*/false); + if (A_stacked.defined() && B_stacked.defined()) { + // PEFT stored A as [r, hidden]; base group_gemm wants + // A_gate: [E_per_rank, hidden, r]. Transpose(1,2). + // PEFT stored B as [out_full, r]; we need [E_per_rank, r, out_local]. + // Shard on dim=1 (out) then transpose to put r first. + A_stacked = A_stacked.transpose(1, 2).contiguous(); + // B_stacked: [E, out_full, r] -> slice on out -> [E, out_local, r] + // -> transpose to [E, r, out_local] + auto B_shard = B_stacked + .slice(1, + tp_rank * local_intermediate, + (tp_rank + 1) * local_intermediate) + .transpose(1, 2) + .contiguous(); + try { + med.A_gate = A_stacked.to(device, dtype).contiguous(); + med.B_gate = B_shard.to(device, dtype).contiguous(); + } catch (const std::exception& e) { + LOG(ERROR) << "[LoRARuntime] MoE .to(device) gate_proj L=" << layer + << ": " << e.what(); + med.A_gate = torch::Tensor(); + med.B_gate = torch::Tensor(); + } + } + } + + // up_proj. + if (auto it = projs.find("up_proj"); it != projs.end()) { + auto A_stacked = stack_experts(it->second, /*is_a=*/true); + auto B_stacked = stack_experts(it->second, /*is_a=*/false); + if (A_stacked.defined() && B_stacked.defined()) { + A_stacked = A_stacked.transpose(1, 2).contiguous(); + auto B_shard = B_stacked + .slice(1, + tp_rank * local_intermediate, + (tp_rank + 1) * local_intermediate) + .transpose(1, 2) + .contiguous(); + try { + med.A_up = A_stacked.to(device, dtype).contiguous(); + med.B_up = B_shard.to(device, dtype).contiguous(); + } catch (const std::exception& e) { + LOG(ERROR) << "[LoRARuntime] MoE .to(device) up_proj L=" << layer + << ": " << e.what(); + med.A_up = torch::Tensor(); + med.B_up = torch::Tensor(); + } + } + } + + // down_proj. + if (auto it = projs.find("down_proj"); it != projs.end()) { + auto A_stacked = stack_experts(it->second, /*is_a=*/true); + auto B_stacked = stack_experts(it->second, /*is_a=*/false); + if (A_stacked.defined() && B_stacked.defined()) { + // A: [r, in_full] -> shard on in -> [r, in_local] -> [E, in_local, r] + auto A_shard = A_stacked + .slice(2, + tp_rank * local_intermediate, + (tp_rank + 1) * local_intermediate) + .transpose(1, 2) + .contiguous(); + // B: [hidden, r] -> [E, r, hidden] + B_stacked = B_stacked.transpose(1, 2).contiguous(); + try { + med.A_down = A_shard.to(device, dtype).contiguous(); + med.B_down = B_stacked.to(device, dtype).contiguous(); + } catch (const std::exception& e) { + LOG(ERROR) << "[LoRARuntime] MoE .to(device) down_proj L=" << layer + << ": " << e.what(); + med.A_down = torch::Tensor(); + med.B_down = torch::Tensor(); + } + } + } + + // Only keep the layer if at least one proj is defined. + if (med.A_gate.defined() || med.A_up.defined() || med.A_down.defined()) { + pool_layers.emplace(layer, std::move(med)); + ++layers_installed; + } else { + ++layers_skipped; + } + } + + { + std::lock_guard g(materialise_mu_); + moe_expert_lora_pool_[int_id] = std::move(pool_layers); + } + LOG(INFO) << "[LoRARuntime] installed MoE expert adapter '" << lora_name + << "' id=" << int_id << " layers=" << layers_installed + << " skipped=" << layers_skipped << " r=" << r + << " scaling=" << scaling << " E_per_rank=" << num_experts_per_rank + << " start_expert=" << start_expert_id + << " intermediate_local=" << local_intermediate + << " tp=" << tp_world; + return int_id; +} + +const LoRARuntime::MoeExpertDelta* LoRARuntime::get_moe_expert_delta( + uint64_t int_id, + int32_t layer_index) { + if (int_id == 0) return nullptr; + std::lock_guard g(materialise_mu_); + auto it = moe_expert_lora_pool_.find(int_id); + if (it == moe_expert_lora_pool_.end()) return nullptr; + auto lit = it->second.find(layer_index); + if (lit == it->second.end()) return nullptr; + return &lit->second; +} + +} // namespace xllm diff --git a/xllm/core/framework/lora/lora_runtime.h b/xllm/core/framework/lora/lora_runtime.h new file mode 100644 index 0000000000..d9038a06f8 --- /dev/null +++ b/xllm/core/framework/lora/lora_runtime.h @@ -0,0 +1,402 @@ +/* Copyright 2025 The xLLM Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#pragma once + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "adapter_loader.h" +#include "lora_config.h" +#include "lora_registry.h" + +namespace xllm { + +// LoRARuntime is the process-global singleton that ties the LoRA modules +// together for the model forward path. It is deliberately narrow-featured +// so P0-A can end-to-end demonstrate load-via-API without waiting on M4 +// (LoraCache) or M8 (Scheduler). +// +// Concretely: +// - A single model-level dummy A/B pair (Path C semantics) can be +// dynamically installed by loading an adapter. +// - The A/B tensors are moved to the model's inference device once at +// load time; there is no eviction / hot-swap yet. +// - forward() paths call `active_delta()` and, if it returns a valid +// value, use those tensors as the whole-block LoRA delta. +// +// Multi-tenant, per-request routing is P0-B / P0-C work. This singleton +// only holds "the currently applied adapter" -- the last one loaded wins. +// That's enough to prove the full end-to-end HTTP -> loader -> registry -> +// forward path from a single curl call. +// Path C prod v3 M10 TP shard skeleton. When TP > 1, per-proj LoRA deltas +// need to shard A/B differently depending on which xllm Linear the delta +// hooks into: +// - ColumnParallelLinear (Q/K/V/gate/up): +// A [rank, hidden] replicated on every rank +// B [hidden/tp, rank] column-sharded (each rank owns a slice) +// - RowParallelLinear (o_proj, down_proj): +// A [rank, hidden/tp] row-sharded +// B [hidden, rank] replicated, delta gets all_reduce'd +// +// Whole-block delta on the current NPU path (atb output is already all- +// reduced back to full hidden) uses REPLICATED for both A and B; each +// rank keeps the full tensor and computes the same delta locally. No +// cross-rank comm needed. +// +// This enum is the skeleton; the actual per-proj sharding is P1 work. +enum class LoRAShardStrategy { + kReplicated = 0, // whole tensor on every rank (current whole-block) + kColumnSharded = 1, // dim-shard on last dim (for Column-parallel proj) + kRowSharded = 2, // dim-shard on second-to-last (for Row-parallel proj) +}; + +struct TPInfo { + int32_t tp_size = 1; + int32_t tp_rank = 0; +}; + +class LoRARuntime { + public: + static LoRARuntime& instance(); + + // One-time init; must be called before any adapter API is used. Copies + // config by value so later registry / cache modules can inspect it + // without another lookup. + void init(const LoRAConfig& config); + + bool enabled() const; + + // Called once by the model (QWen3ModelImpl ctor) so subsequent HTTP + // calls to /v1/load_lora_adapter know what device / dtype to + // materialise weights on. Idempotent; last caller wins (single model + // per engine today). + void set_model_device_dtype(torch::Device device, torch::ScalarType dtype); + + // W4-A3: hot-swap install config. Passed from model ctor so that + // load_and_activate_hotswap knows whether to run per-proj + moe-experts + // installers (and with what TP + MoE args) or only per-proj (dense). + // Idempotent; last caller wins (single model per engine today). + struct HotswapConfig { + TPInfo tp; // TP {world_size, rank} + bool install_per_proj = true; // always install per-proj attn/MLP + bool install_moe_experts = false; // set true for qwen3_moe + // MoE args (ignored when install_moe_experts=false) + int32_t num_experts_total = 0; + int32_t num_experts_per_rank = 0; + int32_t start_expert_id = 0; + int32_t moe_intermediate_size = 0; + }; + void set_hotswap_config(const HotswapConfig& cfg); + + // Load an adapter from a filesystem path, parse its PEFT files, pick a + // whole-block A/B pair, cast to the model's dtype, and register it. + // + // The tensors stay on CPU here -- the actual device migration happens + // lazily on the first active_delta() call from the model forward + // thread, which owns the correct NPU context. + // + // The most recently loaded adapter becomes the pending one and will + // be promoted to active on the next forward. + std::optional load_and_activate(const std::string& lora_name, + const std::string& lora_path, + const std::string& base_model_name); + + // Path C prod v3 static preload path. + // + // Load a PEFT adapter from disk, pick a whole-block A/B pair, cast to + // the given dtype, .to(device) IN THIS THREAD, register it with the + // registry, and install it as the currently active adapter with + // *device-side* tensors. Callable ONLY from a thread with a valid NPU + // context (in practice: QWen3ModelImpl ctor). If called after + // model init has completed the .to(device) will crash with + // aclrtMemcpy 107017 -- see docs/lora_investigation. + // + // Distinct from load_and_activate() which keeps tensors on CPU and + // defers migration; that path is broken on CANN 8.5 + torch_npu 2.7.1 + // and only survives via ctor-time dummy fill. + std::optional install_static_adapter_on_device( + const std::string& lora_name, + const std::string& lora_path, + const std::string& base_model_name, + torch::Device device, + torch::ScalarType dtype); + + // TP-aware overload. Same as above but the caller passes tp_size/tp_rank + // so we can (in future M10 work) shard A/B before .to(device). For the + // current whole-block delta path we ignore the info and store the full + // tensor on every rank (kReplicated). This overload exists so callers + // can start passing the info now and the internal sharding lights up + // when M10 lands, without another API break. + std::optional install_static_adapter_on_device( + const std::string& lora_name, + const std::string& lora_path, + const std::string& base_model_name, + torch::Device device, + torch::ScalarType dtype, + TPInfo tp); + + // Path C prod v3 multi-adapter: look up the device-resident A/B for a + // specific adapter by its int_id. Returns std::nullopt if int_id unknown + // or the adapter was installed without a device pool entry. + // + // Populated by install_static_adapter_on_device when it succeeds. + bool unload(const std::string& lora_name); + + // Path C prod v3 hot-swap: enqueue a load task to the pinned executor + // thread which owns the NPU device context. Blocks until the executor + // completes the install. Returns the assigned int_id, or nullopt on + // failure. Callable from any thread (HTTP handler etc.). + std::optional load_and_activate_hotswap( + const std::string& lora_name, + const std::string& lora_path, + const std::string& base_model_name); + + LoRARegistry& registry() { return registry_; } + const LoRARegistry& registry() const { return registry_; } + const LoRAConfig& config() const { return config_; } + + // The active whole-block delta tensors. std::nullopt = no adapter, the + // forward path should just skip its delta step. + // + // Note: tensors are populated on CPU by the load path (which typically + // runs on the API thread, without an NPU context) and lazily migrated + // to device on the first active_delta() call from the model forward + // thread (which owns the correct NPU context). This avoids the + // aclrtMemcpy-invalid-handle failure you hit when a background thread + // tries to allocate device memory it does not own. + struct ActiveDelta { + torch::Tensor A; // [rank, hidden] (or shard thereof) on model_device + torch::Tensor B; // [hidden, rank] (or shard thereof) on model_device + float scaling; + std::string name; + uint64_t int_id; + // TP shard tag. kReplicated is the current whole-block default; the + // sharded variants light up when M10 per-proj delta ships. + LoRAShardStrategy shard = LoRAShardStrategy::kReplicated; + TPInfo tp{}; + }; + std::optional active_delta(); + + // Path C prod v3 multi-adapter: look up the device-resident A/B for a + // specific adapter by its int_id. Populated by + // install_static_adapter_on_device on ctor thread. Returns nullopt if the + // int_id is not known (base-model request or adapter installed without + // device weights). + std::optional get_delta_by_int_id(uint64_t int_id); + + // ---- M10 per-proj/per-layer delta support ---- + // + // ProjKey identifies a specific weight in a decoder layer, e.g. + // (layer=5, proj="q_proj"). PEFT convention: proj is one of + // {q_proj,k_proj,v_proj,o_proj,gate_proj,up_proj,down_proj}. + // + // ProjDelta holds the device-resident A/B for that specific slot. + // scaling is per-adapter (alpha/rank, or alpha/sqrt(rank) for rslora). + struct ProjKey { + int32_t layer_index; + std::string proj_name; + bool operator==(const ProjKey& other) const { + return layer_index == other.layer_index && proj_name == other.proj_name; + } + }; + struct ProjKeyHash { + size_t operator()(const ProjKey& k) const noexcept { + return std::hash()(k.layer_index) ^ + (std::hash()(k.proj_name) << 1); + } + }; + struct ProjDelta { + torch::Tensor A; // [rank, in_features] or shard thereof + torch::Tensor B; // [out_features, rank] or shard thereof + float scaling = 0.0f; + int32_t r = 0; + }; + + // Materialize every (layer, proj) tensor of the adapter onto device and + // register in the per-proj pool. Call from the model ctor thread (V60 + // rules). + std::optional install_static_adapter_on_device_per_proj( + const std::string& lora_name, + const std::string& lora_path, + const std::string& base_model_name, + torch::Device device, + torch::ScalarType dtype, + TPInfo tp); + + // Look up a single per-proj delta. Called from the Linear wrapper + // forward path once per (batch, layer, proj). Returns nullptr if the + // adapter doesn't have anything at that slot (e.g. adapter only + // targets q_proj/k_proj/v_proj but not gate_proj). + const ProjDelta* get_per_proj_delta(uint64_t int_id, + int32_t layer_index, + const std::string& proj_name); + + // ---- MoE expert LoRA (Day 1 Phase 1) ---- + // + // For MoE models like Qwen3-30B-A3B, adapter safetensors carry per-expert + // LoRA weights, e.g. + // base_model.model.model.layers.{L}.mlp.experts.{E}.gate_proj.lora_A.weight + // These cannot use ProjDelta (single 2D pair per slot) -- we need to stack + // over the E dim so FusedMoEImpl::forward_expert can consume them via one + // grouped-gemm call per proj. + // + // MoeExpertDelta holds device-resident 3D tensors: + // A_gate/A_up: [E_per_rank, hidden, r] + // B_gate/B_up: [E_per_rank, r, out_local] (out_local = intermediate/tp) + // A_down: [E_per_rank, in_local, r] (in_local = intermediate/tp) + // B_down: [E_per_rank, r, hidden] + // Empty (undefined) tensor means "adapter did not target that proj". + struct MoeExpertDelta { + torch::Tensor A_gate; + torch::Tensor B_gate; + torch::Tensor A_up; + torch::Tensor B_up; + torch::Tensor A_down; + torch::Tensor B_down; + float scaling = 0.0f; + int32_t r = 0; + }; + + // Materialize an adapter's experts.{E}.{proj} tensors onto device and + // register in the MoE expert pool. Must be called from the model ctor + // thread (V60 rules) -- xllm can only do CPU->NPU .to() there. + // + // Preconditions: (1) the same adapter has already been installed via + // install_static_adapter_on_device_per_proj (which parses adapter_config, + // registers name->int_id, and canonicalises tensor keys); or a fresh + // adapter with no attention LoRA tensors is being registered by this call. + // + // Only the experts.{E}.{proj} keys are consumed here. Attention keys + // (self_attn.q_proj etc.) are ignored -- the sister call handles them. + // + // Fails silently (returns nullopt) if the adapter has zero experts.* + // tensors. This lets callers install attention-only adapters unchanged. + // + // num_experts_total: the base model's total expert count (128 for + // Qwen3-30B-A3B). num_experts_per_rank: how many experts this rank owns + // (32 when tp=4, ep=1). start_expert_id: this rank's first expert + // (rank * num_experts_per_rank when ep=1). intermediate_size: base + // model's per-expert intermediate dim (768 for Qwen3-30B-A3B, before + // TP shard). tp: as usual. + std::optional install_static_adapter_on_moe_experts( + const std::string& lora_name, + const std::string& lora_path, + const std::string& base_model_name, + torch::Device device, + torch::ScalarType dtype, + TPInfo tp, + int32_t num_experts_total, + int32_t num_experts_per_rank, + int32_t start_expert_id, + int32_t intermediate_size); + + // Look up the MoE expert delta for a decoder layer. Called from + // FusedMoEImpl::forward_expert once per forward per layer. Returns + // nullptr if the adapter has no MoE tensors for this layer (or the + // adapter id is unknown). + const MoeExpertDelta* get_moe_expert_delta(uint64_t int_id, + int32_t layer_index); + + private: + LoRARuntime() = default; + + // Called with materialise_mu_ held. Picks a plausible A/B pair from the + // adapter's canonicalised tensor set. Result is CPU-side, dtype-cast to + // the model dtype but kept off-device. + bool pick_whole_block_ab(const LoRAAdapter& adapter, + torch::ScalarType dtype, + torch::Tensor* A_out, + torch::Tensor* B_out) const; + + // Pending / not-yet-migrated CPU tensors, seeded by load_and_activate. + // active_delta() moves these to device on the forward thread. + struct PendingDelta { + torch::Tensor A_cpu; + torch::Tensor B_cpu; + float scaling; + std::string name; + uint64_t int_id; + }; + std::optional pending_; + + mutable std::mutex materialise_mu_; + LoRAConfig config_; + LoRARegistry registry_; + std::unique_ptr loader_; + + // Recorded by the model at forward-init time so the HTTP handler knows + // where to place freshly-loaded LoRA tensors. std::nullopt means no + // model has registered yet. + std::optional model_device_; + std::optional model_dtype_; + + // Currently-active adapter's device tensors. Guarded by materialise_mu_. + std::optional active_; + + // Path C prod v3 hot-swap: pinned executor thread + task queue. + struct LoadTask { + std::string name; + std::string path; + std::string base_model_name; + std::promise> result; + }; + std::mutex task_mu_; + std::condition_variable task_cv_; + std::queue task_queue_; + std::atomic executor_stop_{false}; + std::thread executor_thread_; + bool executor_started_ = false; + + // W4-A3: hot-swap install config. Guarded by materialise_mu_ (writes + // are rare from model ctor; reads from executor_loop task pickup). + HotswapConfig hotswap_config_{}; + bool hotswap_config_set_ = false; + + void executor_loop(int32_t device_index, torch::ScalarType dtype); + + // Path C prod v3 multi-adapter: int_id -> device-resident A/B/scaling. + // Written by install_static_adapter_on_device from ctor thread. + // Read by get_delta_by_int_id from forward thread. Guarded by + // materialise_mu_. + std::unordered_map device_pool_; + + // M10 per-proj device pool: int_id -> ((layer, proj) -> ProjDelta). + // Populated by install_static_adapter_on_device_per_proj at ctor time + // (V60 window). Read by get_per_proj_delta from Linear wrappers. + // Guarded by materialise_mu_. + std::unordered_map> + per_proj_device_pool_; + + // Day 1 Phase 1: MoE expert LoRA pool. int_id -> (layer_index -> + // MoeExpertDelta). Populated by install_static_adapter_on_moe_experts + // at ctor time. Read by get_moe_expert_delta from FusedMoEImpl::forward. + // Guarded by materialise_mu_. + std::unordered_map> + moe_expert_lora_pool_; +}; + +} // namespace xllm diff --git a/xllm/core/framework/model/model_input_params.h b/xllm/core/framework/model/model_input_params.h index b52eedb206..f962a2e872 100644 --- a/xllm/core/framework/model/model_input_params.h +++ b/xllm/core/framework/model/model_input_params.h @@ -972,6 +972,28 @@ struct ModelInputParams { params.rec_params = llmrec->to(device); } + // Multi-tenant LoRA routing. Propagate adapter_ids and, if not empty, + // materialise adapter_ids_per_token on device via repeat_interleave so + // MoE expert grouped-gemm delta / slow-path per-token masking can index + // straight into it. Pure-base batch (empty adapter_ids) is a no-op. + params.adapter_ids = adapter_ids; + if (!adapter_ids.empty()) { + std::vector host; + const auto& q_lens = attention.host.q_seq_lens; + host.reserve(1024); + for (size_t si = 0; si < adapter_ids.size(); ++si) { + const int64_t seq_len = + (si < q_lens.size()) ? static_cast(q_lens[si]) : 0; + for (int64_t t = 0; t < seq_len; ++t) { + host.push_back(static_cast(adapter_ids[si])); + } + } + if (!host.empty()) { + auto host_t = torch::tensor(host, torch::kInt64); + params.adapter_ids_per_token = safe_to(host_t, device, true); + } + } + return params; } @@ -1145,6 +1167,19 @@ struct ModelInputParams { // Flag for graph capture/replay mode. bool enable_graph = false; + + // Multi-tenant LoRA routing: per-sequence adapter int_id (0 = base model) + // populated from Sequence::adapter_id() and index-aligned with + // attention.host.q_seq_lens (one entry per sequence in the batch). + // Consumed by LoRA wrappers via current_lora_context() to pick the right + // adapter A/B tensors for each seq. + std::vector adapter_ids; + + // Per-token adapter id tensor, built lazily on the device in to(device) + // from (adapter_ids, q_seq_lens_vec) via repeat_interleave. Used by MoE + // expert LoRA (fused grouped-gemm delta) and the slow-path per-token + // masking. Undefined when adapter_ids is empty (pure-base batch). + torch::Tensor adapter_ids_per_token; }; } // namespace xllm diff --git a/xllm/core/layers/common/CMakeLists.txt b/xllm/core/layers/common/CMakeLists.txt index 5314192fb4..3103164a55 100755 --- a/xllm/core/layers/common/CMakeLists.txt +++ b/xllm/core/layers/common/CMakeLists.txt @@ -29,6 +29,9 @@ cc_library( dp_utils.h add_matmul.h moe_fused_topk.h + lora/lora_qkv_parallel_linear.h + lora/lora_row_parallel_linear.h + lora/lora_column_parallel_linear.h SRCS oxygen_vision_attention.cpp qwen2_attention.cpp @@ -52,10 +55,14 @@ cc_library( dp_utils.cpp add_matmul.cpp moe_fused_topk.cpp + lora/lora_qkv_parallel_linear.cpp + lora/lora_row_parallel_linear.cpp + lora/lora_column_parallel_linear.cpp DEPS "-Wl,--whole-archive" "-Wl,--no-whole-archive" :kv_cache + :lora :block :prefix_cache :parallel_state diff --git a/xllm/core/layers/common/dense_mlp.cpp b/xllm/core/layers/common/dense_mlp.cpp index 78ed058521..d06f055421 100644 --- a/xllm/core/layers/common/dense_mlp.cpp +++ b/xllm/core/layers/common/dense_mlp.cpp @@ -65,14 +65,15 @@ DenseMLPImpl::DenseMLPImpl(int64_t hidden_size, int64_t out_feature = is_gated_ ? intermediate_size_ * 2 : intermediate_size_; gate_up_proj_ = register_module("gate_up_proj", - ColumnParallelLinear(hidden_size, - out_feature, - /*bias=*/has_bias, - /*gather_output=*/false, - quant_args, - process_group_, - options, - gate_up_proj_extra_args)); + LoRAColumnParallelLinear(hidden_size, + out_feature, + /*bias=*/has_bias, + /*gather_output=*/false, + quant_args, + process_group_, + options, + /*proj_name=*/"gate_up_proj", + gate_up_proj_extra_args)); act_ = register_module("act", Activation(hidden_act_, is_gated_, swiglu_limit_)); @@ -82,16 +83,18 @@ DenseMLPImpl::DenseMLPImpl(int64_t hidden_size, module_prefix.empty() ? quant_args : quant_args.for_module(module_prefix + ".down_proj"); - down_proj_ = register_module("down_proj", - RowParallelLinear(intermediate_size_, - hidden_size, - /*bias=*/has_bias, - /*input_is_parallelized=*/true, - enable_result_reduction, - down_proj_quant_args, - process_group_, - options, - down_proj_extra_args)); + down_proj_ = + register_module("down_proj", + LoRARowParallelLinear(intermediate_size_, + hidden_size, + /*bias=*/has_bias, + /*input_is_parallelized=*/true, + enable_result_reduction, + down_proj_quant_args, + process_group_, + options, + /*proj_name=*/"down_proj", + down_proj_extra_args)); } torch::Tensor DenseMLPImpl::forward(const torch::Tensor& hidden_states) { diff --git a/xllm/core/layers/common/dense_mlp.h b/xllm/core/layers/common/dense_mlp.h index a6f6222ae4..b31eab96b6 100644 --- a/xllm/core/layers/common/dense_mlp.h +++ b/xllm/core/layers/common/dense_mlp.h @@ -25,6 +25,8 @@ limitations under the License. #include "framework/quant_args.h" #include "framework/state_dict/state_dict.h" #include "linear.h" +#include "lora/lora_column_parallel_linear.h" +#include "lora/lora_row_parallel_linear.h" namespace xllm { namespace layer { @@ -58,8 +60,8 @@ class DenseMLPImpl : public torch::nn::Module { bool is_gated_; int64_t intermediate_size_; ProcessGroup* process_group_; - ColumnParallelLinear gate_up_proj_{nullptr}; - RowParallelLinear down_proj_{nullptr}; + LoRAColumnParallelLinear gate_up_proj_{nullptr}; // LoRA-wrapped drop-in + LoRARowParallelLinear down_proj_{nullptr}; // LoRA-wrapped drop-in Activation act_{nullptr}; bool is_smoothquant_; std::string hidden_act_; diff --git a/xllm/core/layers/common/lora/lora_column_parallel_linear.cpp b/xllm/core/layers/common/lora/lora_column_parallel_linear.cpp new file mode 100644 index 0000000000..f10b52aeab --- /dev/null +++ b/xllm/core/layers/common/lora/lora_column_parallel_linear.cpp @@ -0,0 +1,222 @@ +/* Copyright 2026 The xLLM Authors. All Rights Reserved. +Licensed under the Apache License, Version 2.0. +==============================================================================*/ + +#include "layers/common/lora/lora_column_parallel_linear.h" + +#include +#include + +#include "framework/lora/lora_context.h" +#include "framework/lora/lora_runtime.h" + +namespace xllm { +namespace layer { + +LoRAColumnParallelLinearImpl::LoRAColumnParallelLinearImpl( + int64_t in_features, + int64_t out_features, + bool bias, + bool gather_output, + const QuantArgs& quant_args, + ProcessGroup* process_group, + const torch::TensorOptions& options, + const std::string& proj_name, + const LinearExtraArgs& linear_extra_args) + : proj_name_(proj_name), + in_features_(in_features), + out_features_(out_features), + is_fused_gate_up_(proj_name == "gate_up_proj") { + base_ = ColumnParallelLinear(in_features, + out_features, + bias, + gather_output, + quant_args, + process_group, + options, + linear_extra_args); + + // out_features here is the total; per-rank it is out_features / world_size. + tp_rank_ = process_group ? process_group->rank() : 0; + tp_world_size_ = process_group ? process_group->world_size() : 1; + const int64_t world_size = tp_world_size_; + out_size_local_ = out_features / std::max(1, world_size); + inter_size_local_ = is_fused_gate_up_ ? out_size_local_ / 2 : out_size_local_; +} + +torch::Tensor LoRAColumnParallelLinearImpl::forward(torch::Tensor input) { + auto y = base_->forward(input); + + const auto* ctx = current_lora_context(); + if (ctx == nullptr || ctx->adapter_ids == nullptr || + ctx->q_seq_lens_vec == nullptr || ctx->layer_index < 0) { + return y; + } + const auto& adapter_ids = *ctx->adapter_ids; + const auto& q_seq_lens = *ctx->q_seq_lens_vec; + if (adapter_ids.empty() || adapter_ids.size() != q_seq_lens.size()) { + return y; + } + bool any_nonzero = false; + for (auto id : adapter_ids) { + if (id != 0) { + any_nonzero = true; + break; + } + } + if (!any_nonzero) return y; + + // P1a Phase 0 fast path: single-adapter batch (99% prod). Skip per-seq + // slice/matmul loop; do 2 matmuls per proj on full [T, hidden]. + { + uint64_t sole_aid = 0; + bool single_adapter = true; + for (auto id : adapter_ids) { + if (id == 0) { + single_adapter = false; // base-only seq disqualifies fast path + break; + } + if (sole_aid == 0) { + sole_aid = id; + } else if (id != sole_aid) { + single_adapter = false; + break; + } + } + if (single_adapter && sole_aid != 0) { + auto& runtime = LoRARuntime::instance(); + if (is_fused_gate_up_) { + const auto* gate_pd = + runtime.get_per_proj_delta(sole_aid, ctx->layer_index, "gate_proj"); + const auto* up_pd = + runtime.get_per_proj_delta(sole_aid, ctx->layer_index, "up_proj"); + if (gate_pd == nullptr && up_pd == nullptr) return y; + auto shrink_expand = + [&](const LoRARuntime::ProjDelta* pd) -> torch::Tensor { + if (pd == nullptr) { + return torch::zeros({input.size(0), inter_size_local_}, + input.options()); + } + auto tmp = torch::matmul(input, pd->A.transpose(0, 1)); + torch::Tensor B_local = pd->B; + if (tp_world_size_ > 1 && pd->B.size(0) > inter_size_local_) { + const int64_t start = tp_rank_ * inter_size_local_; + B_local = pd->B.slice(0, start, start + inter_size_local_); + } + auto d = torch::matmul(tmp, B_local.transpose(0, 1)); + return (d * pd->scaling).to(input.dtype()); + }; + auto gate_delta = shrink_expand(gate_pd); + auto up_delta = shrink_expand(up_pd); + auto fused_delta = torch::cat({gate_delta, up_delta}, /*dim=*/-1); + y.add_(fused_delta); + return y; + } else { + // Single-proj branch (reserved). Not exercised in Qwen family + // currently, but mirror the fast path for symmetry. + const auto* pd = + runtime.get_per_proj_delta(sole_aid, ctx->layer_index, proj_name_); + if (pd == nullptr) return y; + auto tmp = torch::matmul(input, pd->A.transpose(0, 1)); + torch::Tensor B_local = pd->B; + if (tp_world_size_ > 1 && pd->B.size(0) > out_size_local_) { + const int64_t start = tp_rank_ * out_size_local_; + B_local = pd->B.slice(0, start, start + out_size_local_); + } + auto delta = torch::matmul(tmp, B_local.transpose(0, 1)); + delta = (delta * pd->scaling).to(y.dtype()); + y.add_(delta); + return y; + } + } + } + + auto& runtime = LoRARuntime::instance(); + int64_t tok_off = 0; + + for (size_t seq_idx = 0; seq_idx < adapter_ids.size(); ++seq_idx) { + const int32_t seq_len = q_seq_lens[seq_idx]; + if (seq_len <= 0) continue; + const uint64_t aid = adapter_ids[seq_idx]; + if (aid == 0) { + tok_off += seq_len; + continue; + } + + auto x_seq = input.slice(0, tok_off, tok_off + seq_len); + + if (is_fused_gate_up_) { + // Fused gate_up: lookup gate_proj + up_proj deltas, concat. + const auto* gate_pd = + runtime.get_per_proj_delta(aid, ctx->layer_index, "gate_proj"); + const auto* up_pd = + runtime.get_per_proj_delta(aid, ctx->layer_index, "up_proj"); + if (gate_pd == nullptr && up_pd == nullptr) { + tok_off += seq_len; + continue; + } + auto make_delta = [&](const LoRARuntime::ProjDelta* pd) -> torch::Tensor { + if (pd == nullptr) { + return torch::zeros({x_seq.size(0), inter_size_local_}, + x_seq.options()); + } + auto tmp = torch::matmul(x_seq, pd->A.transpose(0, 1)); + // TP shard: adapter's B is [inter_full, rank]. This rank owns + // rows [tp_rank * inter_size_local_, +inter_size_local_). + torch::Tensor B_local = pd->B; + if (tp_world_size_ > 1 && pd->B.size(0) > inter_size_local_) { + const int64_t start = tp_rank_ * inter_size_local_; + B_local = pd->B.slice(0, start, start + inter_size_local_); + } + auto d = torch::matmul(tmp, B_local.transpose(0, 1)); + return (d * pd->scaling).to(x_seq.dtype()); + }; + auto gate_delta = make_delta(gate_pd); + auto up_delta = make_delta(up_pd); + auto fused_delta = torch::cat({gate_delta, up_delta}, /*dim=*/-1); + y.slice(0, tok_off, tok_off + seq_len).add_(fused_delta); + } else { + // Single proj (future use — reserved). Not applicable to Qwen for now. + const auto* pd = + runtime.get_per_proj_delta(aid, ctx->layer_index, proj_name_); + if (pd == nullptr) { + tok_off += seq_len; + continue; + } + auto tmp = torch::matmul(x_seq, pd->A.transpose(0, 1)); + // TP shard: adapter's B is [out_full, rank]. This rank owns + // rows [tp_rank * out_size_local_, +out_size_local_). + torch::Tensor B_local = pd->B; + if (tp_world_size_ > 1 && pd->B.size(0) > out_size_local_) { + const int64_t start = tp_rank_ * out_size_local_; + B_local = pd->B.slice(0, start, start + out_size_local_); + } + auto delta = torch::matmul(tmp, B_local.transpose(0, 1)); + delta = (delta * pd->scaling).to(y.dtype()); + y.slice(0, tok_off, tok_off + seq_len).add_(delta); + } + tok_off += seq_len; + } + return y; +} + +void LoRAColumnParallelLinearImpl::load_state_dict( + const StateDict& state_dict) { + base_->load_state_dict(state_dict); +} + +void LoRAColumnParallelLinearImpl::load_state_dict( + const StateDict& state_dict, + const std::vector& prefixes) { + base_->load_state_dict(state_dict, prefixes); +} + +void LoRAColumnParallelLinearImpl::load_state_dict( + const StateDict& state_dict, + int32_t shard_tensor_count, + const std::vector& shard_sizes) { + base_->load_state_dict(state_dict, shard_tensor_count, shard_sizes); +} + +} // namespace layer +} // namespace xllm diff --git a/xllm/core/layers/common/lora/lora_column_parallel_linear.h b/xllm/core/layers/common/lora/lora_column_parallel_linear.h new file mode 100644 index 0000000000..55ed594d60 --- /dev/null +++ b/xllm/core/layers/common/lora/lora_column_parallel_linear.h @@ -0,0 +1,86 @@ +/* Copyright 2026 The xLLM Authors. All Rights Reserved. +Licensed under the Apache License, Version 2.0. +==============================================================================*/ + +// Composition-based LoRA wrapper around ColumnParallelLinear. +// +// Used for gate_up_proj (fused gate+up in Qwen family). LoRA A/B come from +// PEFT under two distinct target modules: "gate_proj" and "up_proj". This +// wrapper reads both per-proj deltas from LoRARuntime and concatenates them +// along the output dim to match the base's fused output. +// +// LoRA math (fused gate_up): +// gate_delta = B_gate @ A_gate @ x A: [r, hidden] B: [inter, r] +// up_delta = B_up @ A_up @ x +// y = base(x) # shape [T, 2*inter_local] +// y += cat([gate_delta, up_delta], -1) * scaling + +#pragma once + +#include + +#include + +#include "framework/parallel_state/parallel_args.h" +#include "framework/quant_args.h" +#include "framework/state_dict/state_dict.h" +#include "layers/common/linear.h" + +namespace xllm { +namespace layer { + +class LoRAColumnParallelLinearImpl : public torch::nn::Module { + public: + LoRAColumnParallelLinearImpl() = default; + + // Signature mirrors ColumnParallelLinearImpl (with LinearExtraArgs). + // Extra parameter: + // fused: "gate_up_proj" -> lookup A/B for "gate_proj" + "up_proj", + // concat delta along dim=-1. + // single proj name (e.g. "o_proj" not applicable here since that's Row) + // - reserved for future use. + LoRAColumnParallelLinearImpl( + int64_t in_features, + int64_t out_features, + bool bias, + bool gather_output, + const QuantArgs& quant_args, + ProcessGroup* process_group, + const torch::TensorOptions& options, + const std::string& proj_name, + const LinearExtraArgs& linear_extra_args = LinearExtraArgs()); + + torch::Tensor forward(torch::Tensor input); + + void load_state_dict(const StateDict& state_dict); + void load_state_dict(const StateDict& state_dict, + const std::vector& prefixes); + void load_state_dict(const StateDict& state_dict, + int32_t shard_tensor_count, + const std::vector& shard_sizes); + + torch::Tensor weight() const { return base_->weight(); } + std::optional get_input_scale() const { + return base_->get_input_scale(); + } + + void pretty_print(std::ostream& stream) const { + stream << name() << " (LoRA-wrapped/" << proj_name_ + << ") base_weight=" << base_->weight().sizes(); + } + + private: + ColumnParallelLinear base_{nullptr}; + std::string proj_name_; // "gate_up_proj" for fused, or single + int64_t in_features_ = 0; + int64_t out_features_ = 0; // per-rank local, total (fused = 2*inter) + int64_t out_size_local_ = 0; // per-rank, matches base + int64_t inter_size_local_ = 0; // per-rank, half of out_size_local_ if fused + bool is_fused_gate_up_ = false; + int64_t tp_rank_ = 0; + int64_t tp_world_size_ = 1; +}; +TORCH_MODULE(LoRAColumnParallelLinear); + +} // namespace layer +} // namespace xllm diff --git a/xllm/core/layers/common/lora/lora_qkv_parallel_linear.cpp b/xllm/core/layers/common/lora/lora_qkv_parallel_linear.cpp new file mode 100644 index 0000000000..aebb9a39a7 --- /dev/null +++ b/xllm/core/layers/common/lora/lora_qkv_parallel_linear.cpp @@ -0,0 +1,299 @@ +/* Copyright 2026 The xLLM Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "layers/common/lora/lora_qkv_parallel_linear.h" + +#include +#include + +#include "framework/lora/lora_context.h" +#include "framework/lora/lora_runtime.h" + +namespace xllm { +namespace layer { + +LoRAQKVParallelLinearImpl::LoRAQKVParallelLinearImpl( + int64_t hidden_size, + int64_t num_heads, + int64_t num_kv_heads, + int64_t head_size, + int64_t num_kv_head_replicas, + bool bias, + bool gather_output, + const ParallelArgs& parallel_args, + const torch::TensorOptions& options, + const QuantArgs& quant_args, + bool q_has_gate) + : hidden_size_(hidden_size) { + // Cache TP topology for LoRA weight sharding logic. + tp_rank_ = parallel_args.tp_group_->rank(); + tp_world_size_ = parallel_args.tp_group_->world_size(); + + // Match base's per-partition sizing exactly. When q_has_gate=true + // (Qwen3-Next attn_output_gate), num_heads passed in already accounts + // for the fused q+gate lanes (num_heads_attn * 2), so q_size_local is + // naturally the fused lane width. The PEFT adapter for such a model + // must also carry a B_q sized to the fused lane; if it does not, the + // shape check in set_lora_weights / per-proj slicing will trip. + // q_size_local = num_heads * head_size (already tp-partitioned) + // kv_size_local = num_kv_heads * head_size (with kv_head_replicas) + q_size_local_ = num_heads * head_size; + kv_size_local_ = num_kv_heads * head_size; + out_size_local_ = q_size_local_ + 2 * kv_size_local_; + q_has_gate_ = q_has_gate; + + // Base linear owns the vanilla forward + weight loading path. NOT + // register_module'd here on purpose — see the note in the header. Base + // is held as a plain member; attention's load_state_dict("qkv_proj.") + // routes through this wrapper's load_state_dict, which forwards to + // base_->load_state_dict, so the on-disk state_dict key layout stays + // "qkv_proj.weight" (identical to vanilla xllm). + base_ = QKVParallelLinear(hidden_size, + num_heads, + num_kv_heads, + head_size, + num_kv_head_replicas, + bias, + gather_output, + parallel_args, + options, + quant_args); +} + +torch::Tensor LoRAQKVParallelLinearImpl::forward(torch::Tensor input) { + auto y = base_->forward(input); + + // Legacy hardcoded path (Spike Day 5b). Kept behind lora_active_ so + // unit tests still work without a LoRARuntime. Real production path + // uses the per-request routing below. + if (lora_active_ && lora_rank_ > 0) { + auto lora_intermediate = torch::matmul(input, lora_a_.transpose(0, 1)); + auto delta = torch::matmul(lora_intermediate, lora_b_.transpose(0, 1)); + y = y + (delta * lora_scaling_).to(y.dtype()); + } + + // M10 per-request per-proj real LoRA. QKV wrapper concatenates Q/K/V + // deltas along the last dim to match the base's fused output. + const auto* ctx = current_lora_context(); + // DEBUG P1G: dump wrapper's view of context (rate limited) + LOG_EVERY_N(ERROR, 200) << "[P1G_QKV_FWD] ctx=" << (ctx ? "ok" : "null") + << " aids_ptr=" + << (ctx && ctx->adapter_ids ? "ok" : "null") + << " layer_idx=" << (ctx ? ctx->layer_index : -99) + << " aids_size=" + << (ctx && ctx->adapter_ids ? ctx->adapter_ids->size() + : 0); + if (ctx == nullptr || ctx->adapter_ids == nullptr || + ctx->q_seq_lens_vec == nullptr || ctx->layer_index < 0) { + return y; + } + const auto& adapter_ids = *ctx->adapter_ids; + const auto& q_seq_lens = *ctx->q_seq_lens_vec; + if (adapter_ids.empty() || adapter_ids.size() != q_seq_lens.size()) { + return y; + } + + // Fast path: batch is pure-base. + bool any_nonzero = false; + for (auto id : adapter_ids) { + if (id != 0) { + any_nonzero = true; + break; + } + } + if (!any_nonzero) return y; + + // P1a Phase 0 fast path: batch uses one non-zero adapter (99% prod). + // Skip per-seq slice/matmul loop; do 2 matmuls on full [T, hidden]: + // tmp = x @ A^T [T, r] + // delta_q = tmp @ B_q_local^T [T, q_local] + // delta_k = tmp @ B_k_local^T [T, kv_local] + // delta_v = tmp @ B_v_local^T [T, kv_local] + // y[:, q_local .. q_local+2*kv_local] += cat([q,k,v], -1) + // For base-only sequences interleaved with adapter sequences we still + // fall back to per-seq (see below); this fast path requires *all* + // non-zero ids match one adapter and no base-only seq is present. + { + uint64_t sole_aid = 0; + bool single_adapter = true; + for (auto id : adapter_ids) { + if (id == 0) { + single_adapter = false; // any base-only seq disqualifies fast path + break; + } + if (sole_aid == 0) { + sole_aid = id; + } else if (id != sole_aid) { + single_adapter = false; + break; + } + } + if (single_adapter && sole_aid != 0) { + auto& runtime = LoRARuntime::instance(); + const auto* q_pd = + runtime.get_per_proj_delta(sole_aid, ctx->layer_index, "q_proj"); + const auto* k_pd = + runtime.get_per_proj_delta(sole_aid, ctx->layer_index, "k_proj"); + const auto* v_pd = + runtime.get_per_proj_delta(sole_aid, ctx->layer_index, "v_proj"); + if (q_pd == nullptr && k_pd == nullptr && v_pd == nullptr) { + return y; + } + // A is the same for Q/K/V? No — three independent A matrices. But + // each proj has its own A. We do 3 shrinks + 3 expands (each on + // the full [T, hidden] input, no slicing). + auto shrink_expand = [&](const LoRARuntime::ProjDelta* pd, + int64_t out_local) -> torch::Tensor { + if (pd == nullptr) { + return torch::zeros({input.size(0), out_local}, input.options()); + } + // tmp: [T, hidden] @ [hidden, r] -> [T, r] + auto tmp = torch::matmul(input, pd->A.transpose(0, 1)); + // B [out_full, r] -> slice rows to local: [out_local, r]. + // For q_proj under TP: out_full = q_hidden_total, sliced evenly + // per tp_rank. For k/v_proj when TP > num_kv_heads (GQA replicas), + // out_full = kv_hidden_total which is smaller than out_local * + // tp_world_size; several ranks share the same kv head. Compute + // the actual shard count from B.size(0) / out_local so each rank + // picks the correct slice. + torch::Tensor B_local = pd->B; + if (tp_world_size_ > 1 && pd->B.size(0) > out_local) { + const int64_t num_shards = pd->B.size(0) / out_local; + const int64_t shard_idx = + (num_shards > 0) ? (tp_rank_ * num_shards / tp_world_size_) : 0; + const int64_t start = shard_idx * out_local; + B_local = pd->B.slice(0, start, start + out_local); + } + // d: [T, r] @ [r, out_local] -> [T, out_local] + auto d = torch::matmul(tmp, B_local.transpose(0, 1)); + return (d * pd->scaling).to(input.dtype()); + }; + auto q_delta = shrink_expand(q_pd, q_size_local_); + auto k_delta = shrink_expand(k_pd, kv_size_local_); + auto v_delta = shrink_expand(v_pd, kv_size_local_); + auto qkv_delta = torch::cat({q_delta, k_delta, v_delta}, /*dim=*/-1); + y.add_(qkv_delta); + return y; + } + } + + // Fallback per-seq apply: slice y along dim=0 by q_seq_lens, look up + // (int_id, layer, {q_proj,k_proj,v_proj}), stack Q/K/V deltas into + // one [seq_tokens, q_size + 2*kv_size] slab, add to that slice. + auto& runtime = LoRARuntime::instance(); + int64_t tok_off = 0; + for (size_t seq_idx = 0; seq_idx < adapter_ids.size(); ++seq_idx) { + const int32_t seq_len = q_seq_lens[seq_idx]; + if (seq_len <= 0) continue; + const uint64_t aid = adapter_ids[seq_idx]; + if (aid == 0) { + tok_off += seq_len; + continue; + } + + const auto* q_pd = + runtime.get_per_proj_delta(aid, ctx->layer_index, "q_proj"); + const auto* k_pd = + runtime.get_per_proj_delta(aid, ctx->layer_index, "k_proj"); + const auto* v_pd = + runtime.get_per_proj_delta(aid, ctx->layer_index, "v_proj"); + if (q_pd == nullptr && k_pd == nullptr && v_pd == nullptr) { + tok_off += seq_len; + continue; + } + + auto x_seq = input.slice(0, tok_off, tok_off + seq_len); + // Compute per-proj delta on the fused output. Use zeros as filler + // for missing proj slots so the concat shape stays correct. + // TP shard: dummy adapter's B is stored at full out size, but this + // rank only handles [tp_rank * out_size, (tp_rank+1) * out_size). + // Slice B to the local shard before matmul so the delta shape + // matches base linear's local output. + auto make_delta = [&](const LoRARuntime::ProjDelta* pd, int64_t out_size) { + if (pd == nullptr) { + return torch::zeros({x_seq.size(0), out_size}, x_seq.options()); + } + auto tmp = torch::matmul(x_seq, pd->A.transpose(0, 1)); + // B is [out_full, rank]. Slice per shard so this rank picks its + // portion of the output. For k/v_proj with GQA replicas the shard + // count is smaller than tp_world_size, so derive it from B.size(0). + torch::Tensor B_local = pd->B; + if (tp_world_size_ > 1 && pd->B.size(0) > out_size) { + const int64_t num_shards = pd->B.size(0) / out_size; + const int64_t shard_idx = + (num_shards > 0) ? (tp_rank_ * num_shards / tp_world_size_) : 0; + const int64_t start = shard_idx * out_size; + B_local = pd->B.slice(0, start, start + out_size); + } + auto d = torch::matmul(tmp, B_local.transpose(0, 1)); + return (d * pd->scaling).to(x_seq.dtype()); + }; + auto q_delta = make_delta(q_pd, q_size_local_); + auto k_delta = make_delta(k_pd, kv_size_local_); + auto v_delta = make_delta(v_pd, kv_size_local_); + auto qkv_delta = torch::cat({q_delta, k_delta, v_delta}, /*dim=*/-1); + + y.slice(0, tok_off, tok_off + seq_len).add_(qkv_delta); + tok_off += seq_len; + } + return y; +} + +void LoRAQKVParallelLinearImpl::load_state_dict( + const StateDict& state_dict, + const std::vector& prefixes) { + // Passthrough: base handles the fused q/k/v load. We do not touch the + // LoRA A/B tensors here — those are set separately via + // set_lora_weights (Spike) or load_lora_state_dict (adapter manager). + base_->load_state_dict(state_dict, prefixes); +} + +void LoRAQKVParallelLinearImpl::load_state_dict(const StateDict& state_dict) { + base_->load_state_dict(state_dict); +} + +void LoRAQKVParallelLinearImpl::set_lora_weights(const torch::Tensor& lora_a, + const torch::Tensor& lora_b, + double scaling) { + // Sanity checks — surface shape mismatches early instead of getting + // an opaque matmul dispatch error at forward time. + CHECK_EQ(lora_a.dim(), 2) + << "lora_a must be 2-D [rank, hidden], got dim=" << lora_a.dim(); + CHECK_EQ(lora_b.dim(), 2) + << "lora_b must be 2-D [out_local, rank], got dim=" << lora_b.dim(); + CHECK_EQ(lora_a.size(1), hidden_size_) + << "lora_a hidden dim mismatch: got " << lora_a.size(1) << ", expected " + << hidden_size_; + CHECK_EQ(lora_b.size(0), out_size_local_) + << "lora_b out_local dim mismatch: got " << lora_b.size(0) + << ", expected " << out_size_local_ << " (q_local=" << q_size_local_ + << " + 2*kv_local=" << kv_size_local_ << ")"; + CHECK_EQ(lora_a.size(0), lora_b.size(1)) + << "lora_a rank " << lora_a.size(0) << " != lora_b rank " + << lora_b.size(1); + + lora_a_ = lora_a.contiguous(); + lora_b_ = lora_b.contiguous(); + lora_rank_ = lora_a.size(0); + lora_scaling_ = scaling; + lora_active_ = true; + + LOG(INFO) << "LoRAQKVParallelLinear: activated adapter with rank=" + << lora_rank_ << ", scaling=" << lora_scaling_ + << ", tp_rank=" << tp_rank_ << "/" << tp_world_size_; +} + +} // namespace layer +} // namespace xllm diff --git a/xllm/core/layers/common/lora/lora_qkv_parallel_linear.h b/xllm/core/layers/common/lora/lora_qkv_parallel_linear.h new file mode 100644 index 0000000000..3f34a13689 --- /dev/null +++ b/xllm/core/layers/common/lora/lora_qkv_parallel_linear.h @@ -0,0 +1,141 @@ +/* Copyright 2026 The xLLM Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +// Composition-based LoRA wrapper around QKVParallelLinear. +// +// Design decisions (from Spike Day 2 code reading): +// Base linear forward is non-virtual + private weight members, so we wrap +// by composition rather than inheritance. +// Wrapper owns a base QKVParallelLinear registered as "base" plus the +// LoRA A/B tensors. Base weight loading path is unchanged; adapter A/B +// are set via set_lora_weights (Spike phase) or load_lora_state_dict +// (real adapter manager, later phases). +// Delta path: y = base(x); delta = B(A(x)) * scaling; return y + delta. +// Fused Q/K/V is handled by treating lora_A as a single [rank, hidden] +// projection and lora_B as [q_size + 2*kv_size_local, rank] (col-shard +// under TP, matching base's fused output layout). + +#pragma once + +#include + +#include "framework/parallel_state/parallel_args.h" +#include "framework/quant_args.h" +#include "framework/state_dict/state_dict.h" +#include "layers/common/linear.h" + +namespace xllm { +namespace layer { + +class LoRAQKVParallelLinearImpl : public torch::nn::Module { + public: + LoRAQKVParallelLinearImpl() = default; + + // Signature mirrors QKVParallelLinearImpl exactly so we can drop-in + // replace the member type in Qwen2Attention with zero call-site changes. + // + // q_has_gate: when true, the base linear was constructed with + // num_heads_effective = num_heads_attn * 2 (Qwen3-Next attn_output_gate + // fuses q + gate into the q lane). The wrapper then sizes q_size_local + // as num_heads_effective * head_size to match base's fused output, and + // the PEFT adapter is expected to have been trained on the same fused + // lane (B_q shape [q_size_fused, rank]). + LoRAQKVParallelLinearImpl(int64_t hidden_size, + int64_t num_heads, + int64_t num_kv_heads, + int64_t head_size, + int64_t num_kv_head_replicas, + bool bias, + bool gather_output, + const ParallelArgs& parallel_args, + const torch::TensorOptions& options, + const QuantArgs& quant_args = QuantArgs{}, + bool q_has_gate = false); + + // Same signature as base: y = base(x); delta = B(A(x)) * scaling. + torch::Tensor forward(torch::Tensor input); + + // Base weight loading -- direct passthrough. + void load_state_dict(const StateDict& state_dict, + const std::vector& prefixes); + void load_state_dict(const StateDict& state_dict); + + // Passthrough accessors -- keeps wrapper drop-in compatible with + // QKVParallelLinearImpl for callers that peek at the base state. + void pretty_print(std::ostream& stream) const { + stream << name() + << " (LoRA-wrapped) base_weight=" << base_->weight().sizes() + << " lora_active=" << (lora_active_ ? "true" : "false"); + } + torch::Tensor weight() const { return base_->weight(); } + std::optional get_input_scale() const { + return base_->get_input_scale(); + } + + // ---- LoRA-specific interface (Spike + future adapter manager) ---- + + // Set LoRA weights directly (Spike hardcoded flow). + // lora_a: [rank, hidden_size] replicated under TP + // lora_b: [q_size + 2 * kv_size_local, rank] col-sharded under TP + // scaling: alpha / rank (typically pre-multiplied into B for real loads) + void set_lora_weights(const torch::Tensor& lora_a, + const torch::Tensor& lora_b, + double scaling); + + // Disable LoRA delta (return base output only). + void clear_lora() { lora_active_ = false; } + + bool has_active_lora() const { return lora_active_; } + int64_t lora_rank() const { return lora_rank_; } + + private: + // Base linear held via ModuleHolder (shared_ptr semantics) but NOT + // register_module'd on this wrapper. Rationale: keeping the base out of + // this wrapper's module tree preserves the original state_dict key + // layout — a checkpoint written for vanilla QKVParallelLinear (path + // "qkv_proj.weight") continues to load unchanged. If we registered the + // base as a submodule, keys would become "qkv_proj.base.weight" and + // every existing checkpoint would break. + // + // Base construction happens inside the wrapper's ctor; forward() and + // load_state_dict() delegate directly. + QKVParallelLinear base_{nullptr}; + + // LoRA delta parameters. Held as raw tensors (not registered params) so + // they can be swapped/cleared at runtime without perturbing the module + // tree. Real adapter manager will register a slot pool separately. + torch::Tensor lora_a_; // [rank, hidden] replicated + torch::Tensor lora_b_; // [q_size + 2 * kv_size_local, rank] col-shard + double lora_scaling_ = 1.0; + int64_t lora_rank_ = 0; + bool lora_active_ = false; + + // Cached shape info for TP-aware LoRA slicing (populated in ctor). + int64_t hidden_size_ = 0; + int64_t q_size_local_ = 0; // q_size / tp_world_size + int64_t kv_size_local_ = 0; // kv_size / tp_world_size (with replication) + int64_t out_size_local_ = 0; // q_size_local + 2 * kv_size_local + int64_t tp_rank_ = 0; + int64_t tp_world_size_ = 1; + + // True for Qwen3-Next-style attn_output_gate: q_size_local includes the + // fused gate lane, and the PEFT adapter's B_q is expected at fused width. + // Currently informational — sizing already uses fused num_heads. + bool q_has_gate_ = false; +}; +TORCH_MODULE(LoRAQKVParallelLinear); + +} // namespace layer +} // namespace xllm diff --git a/xllm/core/layers/common/lora/lora_row_parallel_linear.cpp b/xllm/core/layers/common/lora/lora_row_parallel_linear.cpp new file mode 100644 index 0000000000..1a5132eccc --- /dev/null +++ b/xllm/core/layers/common/lora/lora_row_parallel_linear.cpp @@ -0,0 +1,248 @@ +/* Copyright 2026 The xLLM Authors. All Rights Reserved. +Licensed under the Apache License, Version 2.0. +==============================================================================*/ + +#include "layers/common/lora/lora_row_parallel_linear.h" + +#include +#include + +#include "framework/lora/lora_config.h" +#include "framework/lora/lora_context.h" +#include "framework/lora/lora_runtime.h" +#include "framework/parallel_state/parallel_state.h" + +namespace xllm { +namespace layer { + +LoRARowParallelLinearImpl::LoRARowParallelLinearImpl( + int64_t in_features, + int64_t out_features, + bool bias, + bool input_is_parallelized, + bool enable_result_reduction, + const QuantArgs& quant_args, + ProcessGroup* process_group, + const torch::TensorOptions& options, + const std::string& proj_name, + const LinearExtraArgs& linear_extra_args) + : proj_name_(proj_name), + in_features_(in_features), + out_features_(out_features) { + // Base's own AR is disabled when the fused-AR path is on so the wrapper + // can fold the LoRA delta into the base's partial output and reduce both + // together in a single collective. + fused_ar_ = FLAGS_enable_lora_row_parallel_fused_ar; + const bool base_reduce = fused_ar_ ? false : enable_result_reduction; + // NOT register_module'd on this wrapper: keeps checkpoint keys unchanged. + // (Same rationale as LoRAQKVParallelLinearImpl.) + base_ = RowParallelLinear(in_features, + out_features, + bias, + input_is_parallelized, + base_reduce, + quant_args, + process_group, + options, + linear_extra_args); + auto* pg = base_->process_group(); + tp_world_size_ = (pg != nullptr) ? pg->world_size() : 1; + tp_rank_ = (pg != nullptr) ? pg->rank() : 0; + in_features_local_ = + (tp_world_size_ > 1) ? (in_features_ / tp_world_size_) : in_features_; + // Cache the intent to all-reduce the base output ourselves in fused mode. + // Under TP=1 there is no reduction to do at all. + wrapper_owns_reduction_ = fused_ar_ && tp_world_size_ > 1; +} + +torch::Tensor LoRARowParallelLinearImpl::forward(torch::Tensor input) { + auto y = base_->forward(input); + + auto* pg = base_->process_group(); + + // Fast-out: TP>1 with row-parallel AR disabled and NOT in fused mode. + // The wrapper skips the delta entirely so o_proj / down_proj LoRA silently + // no-ops, matching pre-a9d6ad74 behaviour. + if (tp_world_size_ > 1 && !fused_ar_ && + !FLAGS_enable_lora_row_parallel_all_reduce) { + return y; + } + + // In fused-AR mode the base was constructed with enable_result_reduction= + // false, so `y` here is a per-rank partial-sum on the out-dim. The wrapper + // will add the LoRA delta (also a per-rank partial on out-dim) and + // all-reduce the sum in one collective. Under TP=1 there is no partial + // and nothing to reduce; the code below still works because + // wrapper_owns_reduction_ = false in that case. + + // Row-parallel LoRA delta. + // + // Fused-AR path (default when enable_lora_row_parallel_fused_ar=true): + // * A is sliced on in-dim per rank -> A_local [r, in_local] + // * B is replicated at full [out, r] + // * tmp_local = x_local @ A_local^T [T, r] partial-A + // * local_delta = tmp_local @ B^T * scaling [T, out] partial on + // out-dim + // (mathematically B @ A_local @ x_local; the sum-over-ranks is deferred) + // * combined = y (base partial) + local_delta + // * output = all_reduce(combined) -> full base + full delta in ONE + // collective + // + // Legacy path (fallback, enable_lora_row_parallel_all_reduce=true, + // fused_ar_=false): + // * A slice as above, tmp_local same + // * all-reduce on rank-dim tmp [T, r=16] + // * delta = tmp_full @ B^T * scaling (replicated) + // * y is already reduced by base, y += delta + // + // Cost win of fused vs legacy: 1 AR per proj vs 2. On NPU HCCL where + // launch latency dominates, this halves the collective count and reclaims + // most of the ~17% overhead the legacy path incurs. + const auto* ctx = current_lora_context(); + if (ctx == nullptr || ctx->adapter_ids == nullptr || + ctx->q_seq_lens_vec == nullptr || ctx->layer_index < 0) { + // No LoRA context: nothing to add. Still need to reduce y if fused_ar. + if (wrapper_owns_reduction_ && pg != nullptr) { + y = xllm::parallel_state::reduce(y, pg); + } + return y; + } + const auto& adapter_ids = *ctx->adapter_ids; + const auto& q_seq_lens = *ctx->q_seq_lens_vec; + if (adapter_ids.empty() || adapter_ids.size() != q_seq_lens.size()) { + if (wrapper_owns_reduction_ && pg != nullptr) { + y = xllm::parallel_state::reduce(y, pg); + } + return y; + } + + bool any_nonzero = false; + for (auto id : adapter_ids) { + if (id != 0) { + any_nonzero = true; + break; + } + } + if (!any_nonzero) { + // Pure-base batch: still need to reduce y in fused mode. + if (wrapper_owns_reduction_ && pg != nullptr) { + y = xllm::parallel_state::reduce(y, pg); + } + return y; + } + + // Fast path: batch uses a single adapter and no base-only seq. Shrink+ + // expand on the full batch, one all-reduce for the whole batch instead + // of one per seq. + { + uint64_t sole_aid = 0; + bool single_adapter = true; + for (auto id : adapter_ids) { + if (id == 0) { + single_adapter = false; + break; + } + if (sole_aid == 0) { + sole_aid = id; + } else if (id != sole_aid) { + single_adapter = false; + break; + } + } + if (single_adapter && sole_aid != 0) { + auto& runtime = LoRARuntime::instance(); + const auto* pd = + runtime.get_per_proj_delta(sole_aid, ctx->layer_index, proj_name_); + if (pd == nullptr) { + if (wrapper_owns_reduction_ && pg != nullptr) { + y = xllm::parallel_state::reduce(y, pg); + } + return y; + } + // A: [r, in_full]; slice to local in-shard for TP>1. + torch::Tensor A_local = pd->A; + if (tp_world_size_ > 1 && pd->A.size(1) > in_features_local_) { + const int64_t start = tp_rank_ * in_features_local_; + A_local = pd->A.slice(1, start, start + in_features_local_); + } + // tmp_local: [T, in_local] @ [in_local, r] -> [T, r] (partial) + auto tmp = torch::matmul(input, A_local.transpose(0, 1)); + torch::Tensor delta; + if (wrapper_owns_reduction_) { + // Fused mode: keep tmp partial. Compute per-rank partial delta on + // out-dim and add to y (also partial on out-dim), then reduce once. + delta = torch::matmul(tmp, pd->B.transpose(0, 1)); + } else { + // Legacy: reduce rank-dim tmp first, then expand to out-dim. + if (tp_world_size_ > 1 && pg != nullptr) { + tmp = xllm::parallel_state::reduce(tmp, pg); + } + delta = torch::matmul(tmp, pd->B.transpose(0, 1)); + } + delta = (delta * pd->scaling).to(y.dtype()); + y.add_(delta); + if (wrapper_owns_reduction_ && pg != nullptr) { + y = xllm::parallel_state::reduce(y, pg); + } + return y; + } + } + + // Slow path: per-seq, one all-reduce per adapter-bearing seq (legacy) or + // deferred to a single reduction at the end (fused). Fused mode wins big + // here — legacy issues N collectives per proj per layer for an N-seq + // batch; fused issues 1. Interleaved base + adapter seqs are still the + // main cost driver; adapter-affinity batching at the gateway helps. + auto& runtime = LoRARuntime::instance(); + int64_t tok_off = 0; + for (size_t seq_idx = 0; seq_idx < adapter_ids.size(); ++seq_idx) { + const int32_t seq_len = q_seq_lens[seq_idx]; + if (seq_len <= 0) continue; + const uint64_t aid = adapter_ids[seq_idx]; + if (aid == 0) { + tok_off += seq_len; + continue; + } + + const auto* pd = + runtime.get_per_proj_delta(aid, ctx->layer_index, proj_name_); + if (pd == nullptr) { + tok_off += seq_len; + continue; + } + + auto x_seq = input.slice(0, tok_off, tok_off + seq_len); + torch::Tensor A_local = pd->A; + if (tp_world_size_ > 1 && pd->A.size(1) > in_features_local_) { + const int64_t start = tp_rank_ * in_features_local_; + A_local = pd->A.slice(1, start, start + in_features_local_); + } + auto tmp = torch::matmul(x_seq, A_local.transpose(0, 1)); + torch::Tensor delta; + if (wrapper_owns_reduction_) { + // Fused: keep tmp partial; delta stays partial on out-dim; combined + // reduce happens after the loop. + delta = torch::matmul(tmp, pd->B.transpose(0, 1)); + } else { + if (tp_world_size_ > 1 && pg != nullptr) { + tmp = xllm::parallel_state::reduce(tmp, pg); + } + delta = torch::matmul(tmp, pd->B.transpose(0, 1)); + } + delta = (delta * pd->scaling).to(y.dtype()); + + y.slice(0, tok_off, tok_off + seq_len).add_(delta); + tok_off += seq_len; + } + if (wrapper_owns_reduction_ && pg != nullptr) { + y = xllm::parallel_state::reduce(y, pg); + } + return y; +} + +void LoRARowParallelLinearImpl::load_state_dict(const StateDict& state_dict) { + base_->load_state_dict(state_dict); +} + +} // namespace layer +} // namespace xllm diff --git a/xllm/core/layers/common/lora/lora_row_parallel_linear.h b/xllm/core/layers/common/lora/lora_row_parallel_linear.h new file mode 100644 index 0000000000..9da9460b91 --- /dev/null +++ b/xllm/core/layers/common/lora/lora_row_parallel_linear.h @@ -0,0 +1,100 @@ +/* Copyright 2026 The xLLM Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://github.com/jd-opensource/xllm/blob/main/LICENSE +==============================================================================*/ + +// Composition-based LoRA wrapper around RowParallelLinear. +// +// Used for o_proj (attention output projection) and down_proj (MLP output). +// Both are Row-parallel (input already TP-sharded, output reduced). +// +// LoRA math (single proj): +// y = base(x) +// delta = B @ (A @ x) * scaling A: [r, in_full] B: [out, r] +// return y + delta +// +// TP handling (mirrors SGLang RowParallelLinearWithLoRA): +// * A is slice-sharded on in-dim at forward time: A_local = A[:, rank_slice] +// Input arrives already sharded on in-dim, so this alignment is natural. +// * B is kept full-width (replicated) — cheap because r is tiny. +// * forward: tmp_local = x_local @ A_local^T -> [T, r] (partial per rank) +// tmp_full = all_reduce(tmp_local) (rank-dim, cheap) +// delta = tmp_full @ B^T * scaling (replicated on ranks) +// y = base_output (already reduced by base) + delta +// +// Cost win over naive "all_reduce(delta)": we reduce a rank-dim tensor +// [T, r=16] instead of hidden-dim [T, out]. For Qwen3-30B r=16 vs out=2048 +// that is 128x smaller. + +#pragma once + +#include + +#include +#include + +#include "framework/parallel_state/parallel_args.h" +#include "framework/quant_args.h" +#include "framework/state_dict/state_dict.h" +#include "layers/common/linear.h" + +namespace xllm { +namespace layer { + +class LoRARowParallelLinearImpl : public torch::nn::Module { + public: + LoRARowParallelLinearImpl() = default; + + // Signature mirrors RowParallelLinearImpl exactly (drop-in replace). + // Extra parameters: + // proj_name: which PEFT target module this wrapper serves. Passed to + // LoRARuntime::get_per_proj_delta at forward time. Must be one of + // "o_proj" or "down_proj" for Qwen family. + LoRARowParallelLinearImpl( + int64_t in_features, + int64_t out_features, + bool bias, + bool input_is_parallelized, + bool enable_result_reduction, + const QuantArgs& quant_args, + ProcessGroup* process_group, + const torch::TensorOptions& options, + const std::string& proj_name, + const LinearExtraArgs& linear_extra_args = LinearExtraArgs()); + + torch::Tensor forward(torch::Tensor input); + + void load_state_dict(const StateDict& state_dict); + + void pretty_print(std::ostream& stream) const { + stream << name() << " (LoRA-wrapped/" << proj_name_ + << ") base_weight=" << base_->weight().sizes(); + } + + private: + RowParallelLinear base_{nullptr}; + std::string proj_name_; + int64_t in_features_ = 0; + int64_t out_features_ = 0; + + // Cached TP topology so forward's hot path avoids a virtual dispatch + // into ProcessGroup for the tp_rank / world_size on every call. + int32_t tp_rank_ = 0; + int32_t tp_world_size_ = 1; + int64_t in_features_local_ = 0; + + // Fused-AR mode: base is constructed with enable_result_reduction=false + // and the wrapper owns the collective. Set once in the ctor from the + // enable_lora_row_parallel_fused_ar flag. + bool fused_ar_ = false; + // True iff we actually issue an AR in forward (fused_ar_ && tp>1). + bool wrapper_owns_reduction_ = false; +}; +TORCH_MODULE(LoRARowParallelLinear); + +} // namespace layer +} // namespace xllm diff --git a/xllm/core/layers/common/qwen2_attention.cpp b/xllm/core/layers/common/qwen2_attention.cpp index 1def5d5d36..b652cbee63 100644 --- a/xllm/core/layers/common/qwen2_attention.cpp +++ b/xllm/core/layers/common/qwen2_attention.cpp @@ -79,27 +79,29 @@ Qwen2AttentionImpl::Qwen2AttentionImpl(const ModelContext& context) { // 1. QKV parallel linear qkv_proj_ = register_module("qkv_proj", - QKVParallelLinear(args.hidden_size(), - num_heads_, - num_kv_heads_, - args.head_dim(), - num_kv_head_replicas_, - qkv_bias, - /*gather_output=*/false, - parallel_args, - options, - quant_args)); + LoRAQKVParallelLinear(args.hidden_size(), + num_heads_, + num_kv_heads_, + args.head_dim(), + num_kv_head_replicas_, + qkv_bias, + /*gather_output=*/false, + parallel_args, + options, + quant_args)); // 2. Output projection - o_proj_ = register_module("o_proj", - RowParallelLinear(total_num_heads * args.head_dim(), - args.hidden_size(), - /*bias=*/false, - /*input_is_parallelized=*/true, - /*enable_result_reduction=*/true, - quant_args, - parallel_args.tp_group_, - options)); + o_proj_ = + register_module("o_proj", + LoRARowParallelLinear(total_num_heads * args.head_dim(), + args.hidden_size(), + /*bias=*/false, + /*input_is_parallelized=*/true, + /*enable_result_reduction=*/true, + quant_args, + parallel_args.tp_group_, + options, + /*proj_name=*/"o_proj")); // 3. RMSNorm if (is_qwen3_style_) { diff --git a/xllm/core/layers/common/qwen2_attention.h b/xllm/core/layers/common/qwen2_attention.h index e2b8ea545f..dc8147f7b1 100644 --- a/xllm/core/layers/common/qwen2_attention.h +++ b/xllm/core/layers/common/qwen2_attention.h @@ -25,6 +25,8 @@ limitations under the License. #include "framework/state_dict/state_dict.h" #include "layers/common/rms_norm.h" #include "linear.h" +#include "lora/lora_qkv_parallel_linear.h" +#include "lora/lora_row_parallel_linear.h" #include "rotary_embedding.h" namespace xllm { @@ -56,8 +58,8 @@ class Qwen2AttentionImpl : public torch::nn::Module { bool is_qwen3_style_; bool can_use_fused_qk_norm_rope_; - QKVParallelLinear qkv_proj_{nullptr}; - RowParallelLinear o_proj_{nullptr}; + LoRAQKVParallelLinear qkv_proj_{nullptr}; // LoRA-wrapped drop-in + LoRARowParallelLinear o_proj_{nullptr}; // LoRA-wrapped drop-in RMSNorm q_norm_{nullptr}; RMSNorm k_norm_{nullptr}; Attention attn_{nullptr}; diff --git a/xllm/models/llm/llm_model_base.h b/xllm/models/llm/llm_model_base.h index ce20d9f92e..8d42fa756b 100644 --- a/xllm/models/llm/llm_model_base.h +++ b/xllm/models/llm/llm_model_base.h @@ -23,6 +23,7 @@ limitations under the License. #include "core/common/interruption_bus.h" #include "core/framework/kv_cache/kv_cache.h" +#include "core/framework/lora/lora_context.h" #include "core/framework/model/model_input_params.h" #include "core/framework/model/model_output.h" #include "core/framework/model_context.h" @@ -103,8 +104,20 @@ class LlmModelImplBase : public torch::nn::Module { << "Rec multi-round mode requires unshared_v_caches per layer."; } + // Multi-tenant LoRA: push a LoRAContext frame so LoRA-wrapped + // Linear layers (QKV / o_proj / gate_up / down) can read the + // per-batch adapter routing without extending Linear::forward. + // set_lora_context_layer updates layer_index on the top frame + // each iteration; LoRAScope pops on scope exit. + LoRAContextFrame lora_frame; + lora_frame.adapter_ids = &modified_input_params.adapter_ids; + lora_frame.adapter_ids_per_token = + &modified_input_params.adapter_ids_per_token; + LoRAScope _lora_scope(lora_frame); + std::optional residual; for (size_t i = 0; i < layers_.size(); i++) { + set_lora_context_layer(static_cast(i)); if (llmrec_params != nullptr) { attn_metadata.full_k_cache = llmrec_params->full_k_caches[i]; attn_metadata.full_v_cache = llmrec_params->full_v_caches[i]; From 0c4850b40fb8a222de1e6ba64990295735eb40d9 Mon Sep 17 00:00:00 2001 From: cchh05 Date: Thu, 23 Jul 2026 17:17:14 +0800 Subject: [PATCH 2/3] feat(lora): request path adapter_id propagation + chat routing (PR 2/2) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Wire the LoRA adapter_id from ChatServiceImpl all the way to ModelInputParams.adapter_ids so the LoRA-wrapped Linear layers landed in PR 1/2 can pick up the correct per-sequence routing at forward time. Depends on PR #2014 (framework + attention/MLP wire-up) for: - `ModelInputParams::adapter_ids` / `adapter_ids_per_token` fields - `LoRARuntime` singleton + `LoRARegistry` - `LoRAContextFrame` push in `LlmModelImplBase::forward` Chain (rank0 path) ------------------ ``` ChatServiceImpl::process_async_impl | | lookup_and_pin(model) // model = adapter name | -> lora_pinned->int_id | v RequestParams { adapter_id, adapter_name } | v LLMMaster::generate_request -> RequestState { adapter_id } | v Request::init -> SequenceParams { adapter_id } | v Sequence { adapter_id_ } | v (per forward pass) BatchInputBuilder::build_forward_input -> BuilderState { adapter_ids } -> input_params.adapter_ids = state_.adapter_ids | v ModelInputParams::to(device) -> adapter_ids_per_token = repeat_interleave(adapter_ids, q_seq_lens) | v LlmModelImplBase::forward -> LoRAContextFrame captures &input_params.adapter_ids | v LoRA-wrapped Linear (in PR #2014): reads current_lora_context() ``` Files touched ------------- Request / batch path: - `framework/request/request_params.h` — adds `std::optional adapter_id` + `std::string adapter_name` - `framework/request/request_state.h` — adds `uint64_t adapter_id = 0` - `framework/request/sequence.h` — adds `SequenceParams::adapter_id`, `Sequence::adapter_id()`, member - `framework/request/sequence.cpp` — ctor init from `seq_params.adapter_id` - `framework/request/request.cpp` — `sequence_params.adapter_id = state_.adapter_id` - `framework/batch/batch_input_builder.h` — adds `BuilderState::adapter_ids` - `framework/batch/batch_input_builder.cpp` — push_back per seq + per-thread state merge + `input_params.adapter_ids = std::move(...)` - `runtime/forward_params.h` — adds `ForwardInput::adapter_ids` Distributed runtime: - `distributed_runtime/llm_master.cpp` — `if (sp.adapter_id.has_value()) req_state.adapter_id = sp.adapter_id.value();` Chat routing: - `api_service/chat_service_impl.cpp` — look up `model` against `LoRARegistry` via `lookup_and_pin`; if the name resolves to a registered adapter, populate `request_params.adapter_id` / `adapter_name`. Pure-base requests keep both unset and the whole chain no-ops. Not in this PR (follow-up CLs) ------------------------------ - HTTP endpoints `/v1/load_lora_adapter`, `/v1/unload_lora_adapter`, `/v1/lora_adapters`, `/v1/lora_stats`. Requires `master->load_lora_broadcast()` and `unload_lora_broadcast()` engine APIs that don't exist upstream yet — separate PR. - Non-rank0 broadcast of `adapter_ids` via shm channel. Rank0 sees the correct `adapter_ids` today; non-rank0 workers will need proto and `params_utils.cpp` / `forward_shared_memory_manager.cpp` changes to receive them. In practice for TP=1 this PR is already end-to-end usable; TP>1 will silent-fallback to base on non-rank0 workers until the broadcast lands. - Full drain lifecycle: unpin adapter on chat request finish (both success and error). Correct behaviour today is best-effort because unload endpoint is a follow-up. Rollout ------- Depends on PR #2014 landing first. Static adapter preload via `--lora_modules==,=` (existing gflag in PR #2014) works end-to-end with this CL for TP=1; dynamic hot-load / TP>1 broadcast follows in later CLs. --- xllm/api_service/chat_service_impl.cpp | 16 ++++++++++++++++ xllm/core/distributed_runtime/llm_master.cpp | 1 + .../core/framework/batch/batch_input_builder.cpp | 5 +++++ xllm/core/framework/batch/batch_input_builder.h | 4 ++++ xllm/core/framework/request/request.cpp | 1 + xllm/core/framework/request/request_params.h | 10 ++++++++++ xllm/core/framework/request/request_state.h | 3 +++ xllm/core/framework/request/sequence.cpp | 3 ++- xllm/core/framework/request/sequence.h | 9 +++++++++ xllm/core/runtime/forward_params.h | 6 ++++++ 10 files changed, 57 insertions(+), 1 deletion(-) diff --git a/xllm/api_service/chat_service_impl.cpp b/xllm/api_service/chat_service_impl.cpp index 876ecb99de..d71ca10d15 100644 --- a/xllm/api_service/chat_service_impl.cpp +++ b/xllm/api_service/chat_service_impl.cpp @@ -36,6 +36,7 @@ limitations under the License. #include "core/distributed_runtime/llm_master.h" #include "core/distributed_runtime/rec_master.h" #include "core/distributed_runtime/vlm_master.h" +#include "core/framework/lora/lora_runtime.h" #include "core/framework/request/rec_type.h" #include "core/framework/request/request_params.h" #include "core/util/utils.h" @@ -524,8 +525,23 @@ void ChatServiceImpl::process_rec_chat_request(std::shared_ptr call) { return; } + // LoRA lookup: `model` field may name a registered LoRA adapter + // instead of a base model. If so, pin it (drain-safe) and set the + // request's adapter_id / adapter_name. Base-model requests keep + // adapter_id unset. + std::optional lora_pinned; + if (LoRARuntime::instance().enabled()) { + lora_pinned = LoRARuntime::instance().registry().lookup_and_pin(model); + } + const uint64_t pinned_adapter_id_local = + lora_pinned.has_value() ? lora_pinned->int_id : 0; + RequestParams request_params( rpc_request, call->get_x_request_id(), call->get_x_request_time()); + if (lora_pinned.has_value()) { + request_params.adapter_id = lora_pinned->int_id; + request_params.adapter_name = lora_pinned->req.lora_name; + } // Build messages from rpc request (inline code, not extracted to helper) std::vector messages; diff --git a/xllm/core/distributed_runtime/llm_master.cpp b/xllm/core/distributed_runtime/llm_master.cpp index 0e905fe044..0e3c49920f 100644 --- a/xllm/core/distributed_runtime/llm_master.cpp +++ b/xllm/core/distributed_runtime/llm_master.cpp @@ -506,6 +506,7 @@ std::shared_ptr LLMMaster::generate_request( sp.decode_address, call); req_state.sample_slots = sp.sample_slots; + if (sp.adapter_id.has_value()) req_state.adapter_id = sp.adapter_id.value(); auto request = std::make_shared(sp.request_id, sp.x_request_id, diff --git a/xllm/core/framework/batch/batch_input_builder.cpp b/xllm/core/framework/batch/batch_input_builder.cpp index 85ed4ac4d5..6d14482d52 100644 --- a/xllm/core/framework/batch/batch_input_builder.cpp +++ b/xllm/core/framework/batch/batch_input_builder.cpp @@ -550,6 +550,9 @@ void BatchInputBuilder::process_sequences_multithreaded() { state_.q_seq_lens.insert(state_.q_seq_lens.end(), state.q_seq_lens.begin(), state.q_seq_lens.end()); + state_.adapter_ids.insert(state_.adapter_ids.end(), + state.adapter_ids.begin(), + state.adapter_ids.end()); state_.kv_cache_tokens_nums.insert(state_.kv_cache_tokens_nums.end(), state.kv_cache_tokens_nums.begin(), state.kv_cache_tokens_nums.end()); @@ -679,6 +682,7 @@ void BatchInputBuilder::process_single_sequence( #if defined(USE_NPU) state.seq_lens.push_back(seq_len); state.q_seq_lens.push_back(padded_q_seq_len); + state.adapter_ids.push_back(sequence->adapter_id()); #elif defined(USE_MLU) || defined(USE_CUDA) || defined(USE_ILU) || \ defined(USE_DCU) state.seq_lens.push_back(state.seq_lens.back() + seq_len); @@ -1209,6 +1213,7 @@ ForwardInput BatchInputBuilder::state_to_forward_input() { input_params.attention.host.kv_seq_lens = std::move(state_.seq_lens); input_params.attention.host.q_cu_seq_lens = std::move(q_cu_seq_lens); input_params.attention.host.q_seq_lens = std::move(state_.q_seq_lens); + input_params.adapter_ids = std::move(state_.adapter_ids); input_params.attention.device.new_cache_slots = torch::tensor(state_.new_token_slot_ids, torch::kInt); diff --git a/xllm/core/framework/batch/batch_input_builder.h b/xllm/core/framework/batch/batch_input_builder.h index cf95d4f8c1..2a1d9fcb1f 100644 --- a/xllm/core/framework/batch/batch_input_builder.h +++ b/xllm/core/framework/batch/batch_input_builder.h @@ -112,6 +112,10 @@ class BatchInputBuilder { #if defined(USE_NPU) || defined(USE_MUSA) std::vector seq_lens; std::vector q_seq_lens; + + // LoRA adapter int_id per sequence, index-aligned with q_seq_lens. + // Consumed by ModelInputParams::to(device) to build adapter_ids_per_token. + std::vector adapter_ids; #elif defined(USE_MLU) || defined(USE_CUDA) || defined(USE_ILU) || \ defined(USE_DCU) std::vector seq_lens = {0}; // cu_seq_lens diff --git a/xllm/core/framework/request/request.cpp b/xllm/core/framework/request/request.cpp index 798d61c5d4..303ee76bd9 100644 --- a/xllm/core/framework/request/request.cpp +++ b/xllm/core/framework/request/request.cpp @@ -60,6 +60,7 @@ void Request::create_sequences_group() { sequence_params.rec_type = state_.rec_type; sequence_params.bos_token_id = state_.bos_token_id; sequence_params.request_id = request_id_; + sequence_params.adapter_id = state_.adapter_id; sequence_params.sample_slots = &(state_.sample_slots); sequence_params.sampling_param = &(state_.sampling_param); sequence_params.stopping_checker = &(state_.stopping_checker); diff --git a/xllm/core/framework/request/request_params.h b/xllm/core/framework/request/request_params.h index 5ddfdf8bbf..8f70b5ca6a 100644 --- a/xllm/core/framework/request/request_params.h +++ b/xllm/core/framework/request/request_params.h @@ -170,6 +170,16 @@ struct RequestParams { bool is_sample_request = false; std::vector sample_slots; + + // ---- Multi-tenant LoRA (M7) ---- + // Engine-assigned int_id of the adapter this request should route to. + // Kept as std::optional so the scheduler can partition base vs adapter + // requests. Set by ChatServiceImpl after resolving the model / lora_name + // fields against LoRARegistry. + std::optional adapter_id; + + // Human-readable adapter name kept for logging / metrics only. + std::string adapter_name; }; } // namespace xllm diff --git a/xllm/core/framework/request/request_state.h b/xllm/core/framework/request/request_state.h index d7b09dd4cf..3e63842eb5 100644 --- a/xllm/core/framework/request/request_state.h +++ b/xllm/core/framework/request/request_state.h @@ -175,6 +175,9 @@ struct RequestState final { std::optional call_; std::vector sample_slots; + // LoRA adapter int_id for this request. 0 = no adapter (base model). + // Set by LLMMaster::generate_request from RequestParams::adapter_id. + uint64_t adapter_id = 0; }; } // namespace xllm diff --git a/xllm/core/framework/request/sequence.cpp b/xllm/core/framework/request/sequence.cpp index 786429b71f..12d1baad97 100644 --- a/xllm/core/framework/request/sequence.cpp +++ b/xllm/core/framework/request/sequence.cpp @@ -236,7 +236,8 @@ Sequence::Sequence(size_t index, sequence_params_(seq_params), decoder_(std::move(decoder)), termination_flag_(std::make_shared>(INT32_MAX)), - request_id_(seq_params.request_id) { + request_id_(seq_params.request_id), + adapter_id_(seq_params.adapter_id) { if (is_onerec_model()) { init_onerec_sequence(prompt_token_ids, std::move(input_embedding)); return; diff --git a/xllm/core/framework/request/sequence.h b/xllm/core/framework/request/sequence.h index d1248cce15..3a91927806 100644 --- a/xllm/core/framework/request/sequence.h +++ b/xllm/core/framework/request/sequence.h @@ -90,6 +90,10 @@ struct SequenceParams { int32_t bos_token_id = 0; + // LoRA adapter int_id for this sequence (0 = base model, no adapter). + // Propagated from RequestState::adapter_id via Request::init. + uint64_t adapter_id = 0; + // request id for suffix-decoding request identity std::string request_id; @@ -235,6 +239,10 @@ class Sequence final { } Block copy_block(BlockType type) const { return kv_state_.copy_block(type); } const std::string& request_id() const { return request_id_; } + + // LoRA adapter int_id assigned to this sequence. + uint64_t adapter_id() const { return adapter_id_; } + void set_adapter_id(uint64_t v) { adapter_id_ = v; } // get input embedding torch::Tensor get_input_embedding() const { return input_embedding_; } @@ -631,6 +639,7 @@ class Sequence final { std::atomic last_token_handled_{false}; std::string request_id_; + uint64_t adapter_id_ = 0; // Multi-round beam search result caching int32_t beam_width_cached_ = 0; diff --git a/xllm/core/runtime/forward_params.h b/xllm/core/runtime/forward_params.h index fd6b1037d3..80516f3213 100644 --- a/xllm/core/runtime/forward_params.h +++ b/xllm/core/runtime/forward_params.h @@ -89,6 +89,12 @@ struct ForwardInputBufferEntry { torch::Tensor* target = nullptr; uint64_t offset = 0; uint64_t aligned_bytes = 0; + + // LoRA adapter int_id per sequence, index-aligned with sequence order. + // Populated by BatchInputBuilder from Sequence::adapter_id() and carried + // through the shm channel to non-rank0 workers. Consumed by + // ModelInputParams::to(device) to build adapter_ids_per_token. + std::vector adapter_ids; }; struct ForwardInputBufferPlan { From 33b72ab863f246e43f9d9ba6804d366498d3c289 Mon Sep 17 00:00:00 2001 From: cchh05 Date: Thu, 23 Jul 2026 19:56:16 +0800 Subject: [PATCH 3/3] fix(lora): move BuilderState::adapter_ids outside NPU/MUSA ifdef so ILU/MLU/CUDA/DCU builds see it --- xllm/core/framework/batch/batch_input_builder.h | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/xllm/core/framework/batch/batch_input_builder.h b/xllm/core/framework/batch/batch_input_builder.h index 2a1d9fcb1f..2616c89500 100644 --- a/xllm/core/framework/batch/batch_input_builder.h +++ b/xllm/core/framework/batch/batch_input_builder.h @@ -112,16 +112,18 @@ class BatchInputBuilder { #if defined(USE_NPU) || defined(USE_MUSA) std::vector seq_lens; std::vector q_seq_lens; - - // LoRA adapter int_id per sequence, index-aligned with q_seq_lens. - // Consumed by ModelInputParams::to(device) to build adapter_ids_per_token. - std::vector adapter_ids; #elif defined(USE_MLU) || defined(USE_CUDA) || defined(USE_ILU) || \ defined(USE_DCU) std::vector seq_lens = {0}; // cu_seq_lens std::vector q_seq_lens = {0}; // q_cu_seq_len #endif + // LoRA adapter int_id per sequence, index-aligned with q_seq_lens. + // Consumed by ModelInputParams::to(device) to build adapter_ids_per_token. + // Platform-independent — kept outside the seq_lens ifdef branches so + // ILU/MLU/CUDA/DCU builds see it too. + std::vector adapter_ids; + // Cache and block data std::vector new_token_slot_ids; std::vector> block_tables_vec;