From d2e687bb83b3b21e108dd3f6e904923ad2248cbe Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Mon, 29 Jun 2026 02:23:43 +0200 Subject: [PATCH] Add draft LoRA variants for Laguna --- server/src/common/backend_factory.cpp | 1 + server/src/common/backend_factory.h | 3 + server/src/common/model_backend.h | 3 + server/src/draft/draft_gguf_loader.cpp | 333 ++++++++++++++++++++++++- server/src/internal.h | 15 +- server/src/laguna/laguna_backend.cpp | 192 ++++++++++---- server/src/laguna/laguna_backend.h | 13 +- server/src/server/http_server.cpp | 2 + server/src/server/http_server.h | 16 ++ server/src/server/server_main.cpp | 53 +++- 10 files changed, 577 insertions(+), 54 deletions(-) diff --git a/server/src/common/backend_factory.cpp b/server/src/common/backend_factory.cpp index e93c15df8..a44817198 100644 --- a/server/src/common/backend_factory.cpp +++ b/server/src/common/backend_factory.cpp @@ -146,6 +146,7 @@ std::unique_ptr create_backend(const BackendArgs & args) { LagunaBackendArgs lcfg; lcfg.target_path = args.model_path; lcfg.draft_path = args.draft_path ? args.draft_path : ""; + lcfg.draft_loras = args.draft_loras; lcfg.draft_gpu = args.draft_device.gpu; lcfg.draft_ctx_max = args.draft_ctx_max; lcfg.ddtree_mode = args.ddtree_mode; diff --git a/server/src/common/backend_factory.h b/server/src/common/backend_factory.h index c5dc454f3..8d7e05fe4 100644 --- a/server/src/common/backend_factory.h +++ b/server/src/common/backend_factory.h @@ -11,12 +11,14 @@ #pragma once #include "model_backend.h" +#include "internal.h" #include "placement/placement_config.h" #include "placement/remote_draft_config.h" #include "placement/remote_target_shard_config.h" #include #include +#include namespace dflash::common { @@ -29,6 +31,7 @@ struct BackendArgs { // Optional: speculative decode draft model (qwen35 only) const char * draft_path = nullptr; + std::vector draft_loras; // Device placement DevicePlacement device; diff --git a/server/src/common/model_backend.h b/server/src/common/model_backend.h index b88f6f21f..03c18b46b 100644 --- a/server/src/common/model_backend.h +++ b/server/src/common/model_backend.h @@ -124,6 +124,9 @@ struct GenerateRequest { // path returns success but emits no tokens, so each backend can route the // retry through its existing AR path without copying retry policy. bool force_ar_decode = false; + // Optional request-selected DFlash draft LoRA variant. Empty uses the + // backend default; "base" selects the unadapted draft when available. + std::string draft_lora; }; struct GenerateResult { diff --git a/server/src/draft/draft_gguf_loader.cpp b/server/src/draft/draft_gguf_loader.cpp index 29efc543d..aa8bc6c3b 100644 --- a/server/src/draft/draft_gguf_loader.cpp +++ b/server/src/draft/draft_gguf_loader.cpp @@ -26,11 +26,15 @@ #include "internal.h" #include "common/derived_scalars.h" +#include #include #include #include #include +#include #include +#include +#include #if !defined(_WIN32) #include @@ -115,6 +119,276 @@ int count_swa_layers(const DraftWeights & w) { return n_swa; } +float get_f32_or(const gguf_context * g, const char * key, float fallback) { + int64_t id = gguf_find_key(g, key); + if (id < 0 || gguf_get_kv_type(g, id) != GGUF_TYPE_FLOAT32) return fallback; + return gguf_get_val_f32(g, id); +} + +std::string get_str_or_empty(const gguf_context * g, const char * key) { + int64_t id = gguf_find_key(g, key); + if (id < 0 || gguf_get_kv_type(g, id) != GGUF_TYPE_STRING) return {}; + const char * value = gguf_get_val_str(g, id); + return value ? std::string(value) : std::string(); +} + +bool ends_with(const std::string & value, const char * suffix) { + const size_t n = std::strlen(suffix); + return value.size() >= n && value.compare(value.size() - n, n, suffix) == 0; +} + +struct DraftLoraTensorRef { + ggml_tensor * tensor = nullptr; + size_t offset = 0; + size_t size = 0; +}; + +struct DraftLoraPair { + DraftLoraTensorRef a; + DraftLoraTensorRef b; + bool applied = false; +}; + +struct DraftLoraAdapter { + std::string path; + float scale = 1.0f; + float alpha = 1.0f; + size_t data_start = 0; + Mmap mmap; + gguf_context * gctx = nullptr; + ggml_context * meta_ctx = nullptr; + std::unordered_map pairs; + + ~DraftLoraAdapter() { + if (gctx) gguf_free(gctx); + if (meta_ctx) ggml_free(meta_ctx); + } +}; + +bool read_lora_tensor_f32(const DraftLoraAdapter & adapter, + const DraftLoraTensorRef & ref, + std::vector & out, + std::string & err) { + if (!ref.tensor) { + err = "missing LoRA tensor"; + return false; + } + if (adapter.data_start + ref.offset + ref.size > adapter.mmap.len) { + err = "LoRA tensor overflows file: " + std::string(ref.tensor->name); + return false; + } + const int64_t n = ggml_nelements(ref.tensor); + if (n <= 0) { + err = "LoRA tensor has invalid element count: " + std::string(ref.tensor->name); + return false; + } + out.resize((size_t)n); + const uint8_t * src = (const uint8_t *)adapter.mmap.addr + adapter.data_start + ref.offset; + if (ref.tensor->type == GGML_TYPE_F32) { + if (ref.size != (size_t)n * sizeof(float)) { + err = "LoRA F32 tensor byte size mismatch: " + std::string(ref.tensor->name); + return false; + } + const float * p = reinterpret_cast(src); + std::copy(p, p + n, out.begin()); + return true; + } + if (ref.tensor->type == GGML_TYPE_F16) { + if (ref.size != (size_t)n * sizeof(ggml_fp16_t)) { + err = "LoRA F16 tensor byte size mismatch: " + std::string(ref.tensor->name); + return false; + } + const ggml_fp16_t * p = reinterpret_cast(src); + for (int64_t i = 0; i < n; ++i) { + out[(size_t)i] = ggml_fp16_to_fp32(p[i]); + } + return true; + } + err = "unsupported LoRA tensor type for " + std::string(ref.tensor->name) + + ": " + ggml_type_name(ref.tensor->type); + return false; +} + +bool load_draft_lora_adapter(const DraftLoraSpec & spec, + std::unique_ptr & out, + std::string & err) { + auto adapter = std::make_unique(); + adapter->path = spec.path; + adapter->scale = spec.scale; + + gguf_init_params gip{}; + gip.no_alloc = true; + gip.ctx = &adapter->meta_ctx; + adapter->gctx = gguf_init_from_file(spec.path.c_str(), gip); + if (!adapter->gctx) { + err = "gguf_init_from_file failed for draft LoRA: " + spec.path; + return false; + } + + const std::string general_type = get_str_or_empty(adapter->gctx, "general.type"); + const std::string adapter_type = get_str_or_empty(adapter->gctx, "adapter.type"); + if (!general_type.empty() && general_type != "adapter") { + err = "draft LoRA general.type must be 'adapter', got: " + general_type; + return false; + } + if (!adapter_type.empty() && adapter_type != "lora") { + err = "draft LoRA adapter.type must be 'lora', got: " + adapter_type; + return false; + } + adapter->alpha = get_f32_or(adapter->gctx, "adapter.lora.alpha", 1.0f); + + if (!adapter->mmap.open_ro(spec.path, err)) { + return false; + } + adapter->data_start = gguf_get_data_offset(adapter->gctx); + + const int64_t n_tensors = gguf_get_n_tensors(adapter->gctx); + for (int64_t tid = 0; tid < n_tensors; ++tid) { + const char * raw_name = gguf_get_tensor_name(adapter->gctx, tid); + if (!raw_name) continue; + std::string name(raw_name); + const bool is_a = ends_with(name, ".lora_a"); + const bool is_b = ends_with(name, ".lora_b"); + if (!is_a && !is_b) { + if (ends_with(name, "_norm.weight")) continue; + err = "unexpected draft LoRA tensor suffix: " + name; + return false; + } + + const char * suffix = is_a ? ".lora_a" : ".lora_b"; + std::string base = name.substr(0, name.size() - std::strlen(suffix)); + ggml_tensor * t = ggml_get_tensor(adapter->meta_ctx, raw_name); + if (!t) { + err = "draft LoRA tensor descriptor missing: " + name; + return false; + } + + DraftLoraTensorRef ref; + ref.tensor = t; + ref.offset = gguf_get_tensor_offset(adapter->gctx, tid); + ref.size = gguf_get_tensor_size(adapter->gctx, tid); + DraftLoraPair & pair = adapter->pairs[base]; + if (is_a) pair.a = ref; + else pair.b = ref; + } + + for (const auto & it : adapter->pairs) { + if (!it.second.a.tensor || !it.second.b.tensor) { + err = "draft LoRA tensor pair is incomplete for base tensor: " + it.first; + return false; + } + } + if (adapter->pairs.empty()) { + err = "draft LoRA contained no lora_a/lora_b tensor pairs: " + spec.path; + return false; + } + + out = std::move(adapter); + return true; +} + +bool merge_lora_into_tensor_bytes( + const char * tname, + const ggml_tensor * base, + const uint8_t * base_bytes, + size_t base_size, + std::vector> & adapters, + std::vector & merged_bytes, + std::string & err) { + std::vector matching; + for (auto & adapter : adapters) { + if (adapter->pairs.find(tname) != adapter->pairs.end()) { + matching.push_back(adapter.get()); + } + } + if (matching.empty()) return false; + + if (base->type != GGML_TYPE_F16 && base->type != GGML_TYPE_F32) { + err = "draft LoRA merge only supports F16/F32 base tensors; tensor " + + std::string(tname) + " has type " + ggml_type_name(base->type); + return false; + } + if (base->ne[0] <= 0 || base->ne[1] <= 0 || base->ne[2] != 1 || base->ne[3] != 1) { + err = "draft LoRA merge expects a 2D base tensor: " + std::string(tname); + return false; + } + + const int64_t in_dim = base->ne[0]; + const int64_t out_dim = base->ne[1]; + const int64_t n = in_dim * out_dim; + std::vector merged((size_t)n); + if (base->type == GGML_TYPE_F32) { + if (base_size != (size_t)n * sizeof(float)) { + err = "draft LoRA base F32 byte size mismatch: " + std::string(tname); + return false; + } + const float * p = reinterpret_cast(base_bytes); + std::copy(p, p + n, merged.begin()); + } else { + if (base_size != (size_t)n * sizeof(ggml_fp16_t)) { + err = "draft LoRA base F16 byte size mismatch: " + std::string(tname); + return false; + } + const ggml_fp16_t * p = reinterpret_cast(base_bytes); + for (int64_t i = 0; i < n; ++i) { + merged[(size_t)i] = ggml_fp16_to_fp32(p[i]); + } + } + + std::vector a; + std::vector b; + for (DraftLoraAdapter * adapter : matching) { + DraftLoraPair & pair = adapter->pairs[tname]; + if (pair.a.tensor->ne[0] != in_dim || + pair.b.tensor->ne[1] != out_dim || + pair.a.tensor->ne[1] != pair.b.tensor->ne[0]) { + char buf[384]; + std::snprintf(buf, sizeof(buf), + "draft LoRA shape mismatch for %s: base=[%lld,%lld] " + "lora_a=[%lld,%lld] lora_b=[%lld,%lld]", + tname, + (long long)in_dim, (long long)out_dim, + (long long)pair.a.tensor->ne[0], (long long)pair.a.tensor->ne[1], + (long long)pair.b.tensor->ne[0], (long long)pair.b.tensor->ne[1]); + err = buf; + return false; + } + if (!read_lora_tensor_f32(*adapter, pair.a, a, err) || + !read_lora_tensor_f32(*adapter, pair.b, b, err)) { + return false; + } + + const int64_t rank = pair.a.tensor->ne[1]; + if (rank <= 0) { + err = "draft LoRA rank must be positive for tensor: " + std::string(tname); + return false; + } + const float factor = adapter->scale * adapter->alpha / (float)rank; + for (int64_t o = 0; o < out_dim; ++o) { + float * dst_col = merged.data() + o * in_dim; + for (int64_t r = 0; r < rank; ++r) { + const float br = b[(size_t)(r + o * rank)] * factor; + const float * a_col = a.data() + r * in_dim; + for (int64_t i = 0; i < in_dim; ++i) { + dst_col[i] += a_col[i] * br; + } + } + } + pair.applied = true; + } + + merged_bytes.resize(base_size); + if (base->type == GGML_TYPE_F32) { + std::memcpy(merged_bytes.data(), merged.data(), base_size); + } else { + ggml_fp16_t * p = reinterpret_cast(merged_bytes.data()); + for (int64_t i = 0; i < n; ++i) { + p[i] = ggml_fp32_to_fp16(merged[(size_t)i]); + } + } + return true; +} + bool check_shape_1d(const ggml_tensor * t, int64_t ne0, const char * name, char * buf, size_t buf_sz) { if (!t || t->ne[0] != ne0) { std::snprintf(buf, buf_sz, "draft GGUF: Domino tensor %s shape mismatch: got [%lld], expected [%lld]", @@ -143,7 +417,8 @@ bool check_shape_2d(const ggml_tensor * t, int64_t ne0, int64_t ne1, bool load_draft_gguf(const std::string & path, ggml_backend_t backend, DraftWeights & out, - const TargetWeights * target) { + const TargetWeights * target, + const DraftLoadOptions * options) { // ── 1. Parse metadata + create ggml_context with tensor descriptors ── ggml_context * meta_ctx = nullptr; @@ -410,6 +685,25 @@ bool load_draft_gguf(const std::string & path, n_swa, out.n_layer, out.swa_window); } + std::vector> lora_adapters; + if (options && !options->loras.empty()) { + lora_adapters.reserve(options->loras.size()); + for (const DraftLoraSpec & spec : options->loras) { + if (spec.path.empty()) continue; + std::unique_ptr adapter; + std::string lora_err; + if (!load_draft_lora_adapter(spec, adapter, lora_err)) { + set_last_error(lora_err); + gguf_free(gctx); + return false; + } + std::fprintf(stderr, + "[draft GGUF] loaded LoRA adapter: %s (pairs=%zu scale=%.3f alpha=%.3f)\n", + spec.path.c_str(), adapter->pairs.size(), adapter->scale, adapter->alpha); + lora_adapters.emplace_back(std::move(adapter)); + } + } + // ── 3. Allocate CUDA buffer for all tensors ────────────────────────── out.buf = ggml_backend_alloc_ctx_tensors(meta_ctx, backend); if (!out.buf) { @@ -426,6 +720,8 @@ bool load_draft_gguf(const std::string & path, const int64_t n_tensors = gguf_get_n_tensors(gctx); size_t total = 0; + size_t merged_lora_tensors = 0; + std::vector merged_bytes; for (int64_t tid = 0; tid < n_tensors; tid++) { const char * tname = gguf_get_tensor_name(gctx, tid); ggml_tensor * t = ggml_get_tensor(meta_ctx, tname); @@ -437,10 +733,39 @@ bool load_draft_gguf(const std::string & path, gguf_free(gctx); return false; } - ggml_backend_tensor_set(t, (const uint8_t *)mm.addr + off, 0, sz); + const uint8_t * tensor_bytes = (const uint8_t *)mm.addr + off; + if (!lora_adapters.empty()) { + std::string merge_err; + const bool merged = merge_lora_into_tensor_bytes( + tname, t, tensor_bytes, sz, lora_adapters, merged_bytes, merge_err); + if (!merge_err.empty()) { + set_last_error(merge_err); + gguf_free(gctx); + return false; + } + if (merged) { + ggml_backend_tensor_set(t, merged_bytes.data(), 0, merged_bytes.size()); + merged_lora_tensors++; + } else { + ggml_backend_tensor_set(t, tensor_bytes, 0, sz); + } + } else { + ggml_backend_tensor_set(t, tensor_bytes, 0, sz); + } total += sz; } + for (const auto & adapter : lora_adapters) { + for (const auto & it : adapter->pairs) { + if (!it.second.applied) { + set_last_error("draft LoRA tensor did not match any draft base tensor: " + + it.first + " from " + adapter->path); + gguf_free(gctx); + return false; + } + } + } + gguf_free(gctx); // Structural defense: derive head_dim / n_head / n_head_kv from weight @@ -495,8 +820,8 @@ bool load_draft_gguf(const std::string & path, char summary[192]; std::snprintf(summary, sizeof(summary), - "draft GGUF loaded: %" PRId64 " tensors, %.2f GiB on GPU", - n_tensors, total / (1024.0 * 1024.0 * 1024.0)); + "draft GGUF loaded: %" PRId64 " tensors, %.2f GiB on GPU, LoRA merged tensors=%zu", + n_tensors, total / (1024.0 * 1024.0 * 1024.0), merged_lora_tensors); set_last_error(summary); return true; diff --git a/server/src/internal.h b/server/src/internal.h index 8e93532fb..f8e39d8de 100644 --- a/server/src/internal.h +++ b/server/src/internal.h @@ -303,6 +303,18 @@ struct DraftWeights { DraftDominoWeights domino; }; +struct DraftLoraSpec { + // Empty name means this spec belongs to the default LoRA stack. Named + // specs create switchable request-selectable variants. + std::string name; + std::string path; + float scale = 1.0f; +}; + +struct DraftLoadOptions { + std::vector loras; +}; + bool load_draft_safetensors(const std::string & path, ggml_backend_t backend, DraftWeights & out, @@ -315,7 +327,8 @@ bool load_draft_safetensors(const std::string & path, bool load_draft_gguf(const std::string & path, ggml_backend_t backend, DraftWeights & out, - const TargetWeights * target = nullptr); + const TargetWeights * target = nullptr, + const DraftLoadOptions * options = nullptr); void free_draft_weights(DraftWeights & w); diff --git a/server/src/laguna/laguna_backend.cpp b/server/src/laguna/laguna_backend.cpp index f4de8fd8e..227f8ee3a 100644 --- a/server/src/laguna/laguna_backend.cpp +++ b/server/src/laguna/laguna_backend.cpp @@ -360,12 +360,15 @@ bool LagunaBackend::do_spec_decode(int committed, int n_gen, bool * forced_close_out, float * accept_rate_out, const std::vector * sample_history_prefix) { + DraftWeights * active_dw = active_dw_; + if (!active_dw) return false; + DraftWeights & dw = *active_dw; const int hidden = w_.n_embd; int32_t last_tok = cache_.last_tok; if (last_tok < 0) return false; DFlashTarget * target = dflash_target_; - const int block_size = dw_.block_size; + const int block_size = dw.block_size; // [TAG_LAGUNA_VERIFY_WIDTH] Speculative verify width (chain). On this MoE // target the batched verify forward's cost grows with the verify width: it // reads the union of experts the batch routes to (bandwidth-bound), and above @@ -398,7 +401,7 @@ bool LagunaBackend::do_spec_decode(int committed, int n_gen, StepGraph draft_sg; // The draft graph is always block_size-wide (build_draft_step uses - // dw_.block_size); chain reads/verifies only its first q_len outputs. + // dw.block_size); chain reads/verifies only its first q_len outputs. std::vector noise_embed((size_t)hidden * (size_t)block_size); std::vector noise_ids((size_t)block_size); std::vector draft_tok((size_t)q_len); @@ -555,7 +558,7 @@ bool LagunaBackend::do_spec_decode(int committed, int n_gen, const bool use_mirror_view = draft_feature_mirror_can_view(feature_mirror_, committed, draft_ctx, mirror_slot0); - if (!build_draft_step(draft_sg, dw_, /*lm_head=*/nullptr, draft_backend_, + if (!build_draft_step(draft_sg, dw, /*lm_head=*/nullptr, draft_backend_, draft_ctx, use_mirror_view ? &feature_mirror_ : nullptr, committed, std::min(ring_cap, std::max(DRAFT_CTX_MAX_DEFAULT, args_.draft_ctx_max)))) { @@ -592,15 +595,15 @@ bool LagunaBackend::do_spec_decode(int committed, int n_gen, sizeof(float) * local_hidden.size()); bool used_domino = false; - if (dw_.domino.enabled && q_len > 1 && !sampled_verify && !args_.ddtree_mode) { + if (dw.domino.enabled && q_len > 1 && !sampled_verify && !args_.ddtree_mode) { static std::atomic s_domino_logged{false}; if (!s_domino_logged.exchange(true)) { std::fprintf(stderr, "[laguna-spec] Domino GRU head active for greedy chain decode " "(H=%d E=%d)\n", - dw_.domino.gru_hidden_dim, dw_.domino.emb_dim); + dw.domino.gru_hidden_dim, dw.domino.emb_dim); } - if (domino_correct_greedy_chain(dw_, draft_backend_, *target, + if (domino_correct_greedy_chain(dw, draft_backend_, *target, local_hidden.data(), q_len, last_tok, draft_tok)) { used_domino = true; @@ -987,11 +990,18 @@ GenerateResult LagunaBackend::generate_impl(const GenerateRequest & req, cache_.last_tok = argmax(last_logits); result.tokens.reserve(req.n_gen); const bool sampled_verify = laguna_sampled_verify_enabled(sampler_, req.do_sample); + if (!req.draft_lora.empty() && + (args_.draft_path.empty() || !dflash_target_ || + !select_decode_draft(req.draft_lora))) { + result.error = "draft_lora"; + return result; + } const bool can_spec = req.n_gen > 0 && !req.force_ar_decode && !args_.draft_path.empty() && dflash_target_ + && select_decode_draft(req.draft_lora) && !draft_parked_ && feature_mirror_.target_feat && cache_.target_feat @@ -1206,11 +1216,18 @@ GenerateResult LagunaBackend::restore_and_generate_impl(int slot, cache_.last_tok = argmax(last_logits); result.tokens.reserve(req.n_gen); const bool sampled_verify = laguna_sampled_verify_enabled(sampler_, req.do_sample); + if (!req.draft_lora.empty() && + (args_.draft_path.empty() || !dflash_target_ || + !select_decode_draft(req.draft_lora))) { + result.error = "draft_lora"; + return result; + } const bool can_spec = req.n_gen > 0 && !req.force_ar_decode && !args_.draft_path.empty() && dflash_target_ + && select_decode_draft(req.draft_lora) && !draft_parked_ && feature_mirror_.target_feat && cache_.target_feat @@ -2731,7 +2748,7 @@ void LagunaBackend::maybe_post_request_swap() { bool LagunaBackend::load_decode_draft() { if (args_.draft_path.empty()) return false; - if (draft_backend_ && feature_mirror_.target_feat) { + if (draft_backend_ && feature_mirror_.target_feat && !draft_variants_.empty()) { draft_parked_ = false; return true; } @@ -2747,57 +2764,101 @@ bool LagunaBackend::load_decode_draft() { draft_backend_ = backend_; } - if (!load_draft_gguf(args_.draft_path, draft_backend_, dw_, nullptr)) { - std::fprintf(stderr, "[laguna] draft load failed: %s\n", dflash27b_last_error()); - if (draft_backend_ && draft_backend_ != backend_) { - ggml_backend_free(draft_backend_); + draft_variants_.clear(); + draft_variants_.push_back(LagunaDraftVariant{}); + draft_variants_.back().name = "base"; + + for (const DraftLoraSpec & spec : args_.draft_loras) { + const std::string name = spec.name.empty() ? "default" : spec.name; + auto it = std::find_if(draft_variants_.begin(), draft_variants_.end(), + [&](const LagunaDraftVariant & v) { return v.name == name; }); + if (it == draft_variants_.end()) { + draft_variants_.push_back(LagunaDraftVariant{}); + draft_variants_.back().name = name; + it = draft_variants_.end() - 1; } - draft_backend_ = nullptr; - return false; + it->loras.push_back(spec); } - dw_.mask_token_id = 12; - const int draft_hidden = (int)dw_.fc->ne[1]; - const int fc_in = (int)dw_.fc->ne[0]; - const int n_capture = fc_in / w_.n_embd; + int base_fc_in = 0; + int base_draft_hidden = 0; + int base_n_capture = 0; - if (draft_hidden != dw_.n_embd) { - std::printf("[laguna] draft: overriding n_embd %d -> %d (from fc weight)\n", - dw_.n_embd, draft_hidden); - dw_.n_embd = draft_hidden; - } - if (dw_.n_layer > 0 && dw_.layers[0].wq) { - const int q_dim = (int)dw_.layers[0].wq->ne[1]; - const int inferred_n_head = q_dim / dw_.head_dim; - if (inferred_n_head != dw_.n_head) { - std::printf("[laguna] draft: overriding n_head %d -> %d\n", - dw_.n_head, inferred_n_head); - dw_.n_head = inferred_n_head; + for (LagunaDraftVariant & variant : draft_variants_) { + DraftLoadOptions load_options; + load_options.loras = variant.loras; + const DraftLoadOptions * options = + load_options.loras.empty() ? nullptr : &load_options; + if (!load_draft_gguf(args_.draft_path, draft_backend_, variant.weights, + nullptr, options)) { + std::fprintf(stderr, "[laguna] draft load failed for variant '%s': %s\n", + variant.name.c_str(), dflash27b_last_error()); + free_decode_draft(); + return false; } - } - if (dw_.n_layer > 0 && dw_.layers[0].w_gate) { - const int inferred_ff = (int)dw_.layers[0].w_gate->ne[1]; - if (inferred_ff != dw_.n_ff) { - std::printf("[laguna] draft: overriding n_ff %d -> %d\n", - dw_.n_ff, inferred_ff); - dw_.n_ff = inferred_ff; + + DraftWeights & dw = variant.weights; + dw.mask_token_id = 12; + const int draft_hidden = (int)dw.fc->ne[1]; + const int fc_in = (int)dw.fc->ne[0]; + const int n_capture = fc_in / w_.n_embd; + + if (draft_hidden != dw.n_embd) { + std::printf("[laguna] draft[%s]: overriding n_embd %d -> %d (from fc weight)\n", + variant.name.c_str(), dw.n_embd, draft_hidden); + dw.n_embd = draft_hidden; + } + if (dw.n_layer > 0 && dw.layers[0].wq) { + const int q_dim = (int)dw.layers[0].wq->ne[1]; + const int inferred_n_head = q_dim / dw.head_dim; + if (inferred_n_head != dw.n_head) { + std::printf("[laguna] draft[%s]: overriding n_head %d -> %d\n", + variant.name.c_str(), dw.n_head, inferred_n_head); + dw.n_head = inferred_n_head; + } } - } - dw_.n_target_layers = n_capture; - dw_.swa_window = 2048; - for (int i = 0; i < dw_.n_layer - 1 && i < (int)dw_.layers.size(); i++) { - dw_.layers[(size_t)i].is_swa = true; + if (dw.n_layer > 0 && dw.layers[0].w_gate) { + const int inferred_ff = (int)dw.layers[0].w_gate->ne[1]; + if (inferred_ff != dw.n_ff) { + std::printf("[laguna] draft[%s]: overriding n_ff %d -> %d\n", + variant.name.c_str(), dw.n_ff, inferred_ff); + dw.n_ff = inferred_ff; + } + } + dw.n_target_layers = n_capture; + dw.swa_window = 2048; + for (int i = 0; i < dw.n_layer - 1 && i < (int)dw.layers.size(); i++) { + dw.layers[(size_t)i].is_swa = true; + } + + if (base_fc_in == 0) { + base_fc_in = fc_in; + base_draft_hidden = draft_hidden; + base_n_capture = n_capture; + } else if (fc_in != base_fc_in || draft_hidden != base_draft_hidden || + n_capture != base_n_capture) { + std::fprintf(stderr, + "[laguna] draft LoRA variant '%s' changed draft dimensions " + "(fc_in=%d hidden=%d capture=%d, base fc_in=%d hidden=%d capture=%d)\n", + variant.name.c_str(), fc_in, draft_hidden, n_capture, + base_fc_in, base_draft_hidden, base_n_capture); + free_decode_draft(); + return false; + } + + std::printf("[laguna] draft variant loaded: name=%s loras=%zu fc_in=%d " + "target_hidden=%d draft_hidden=%d n_capture_layers=%d swa=%d\n", + variant.name.c_str(), variant.loras.size(), fc_in, w_.n_embd, + draft_hidden, n_capture, dw.swa_window); } - std::printf("[laguna] draft loaded: fc_in=%d target_hidden=%d " - "draft_hidden=%d n_capture_layers=%d swa=%d\n", - fc_in, w_.n_embd, draft_hidden, n_capture, dw_.swa_window); + const int n_capture = base_n_capture; constexpr int TARGET_FEAT_CAP = 4096; const int feat_cap = std::min(args_.max_ctx, TARGET_FEAT_CAP); if (!cache_.target_feat && !create_laguna_target_feat(backend_, cache_, n_capture, w_.n_embd, feat_cap, - dw_.capture_layer_ids)) { + draft_variants_[0].weights.capture_layer_ids)) { std::fprintf(stderr, "[laguna] target_feat alloc failed\n"); free_decode_draft(); return false; @@ -2812,6 +2873,16 @@ bool LagunaBackend::load_decode_draft() { return false; } + const bool only_default_loras = + !args_.draft_loras.empty() && + std::all_of(args_.draft_loras.begin(), args_.draft_loras.end(), + [](const DraftLoraSpec & s) { return s.name.empty(); }); + default_draft_lora_ = only_default_loras ? "default" : "base"; + if (!select_decode_draft(default_draft_lora_)) { + free_decode_draft(); + return false; + } + delete dflash_target_; dflash_target_ = new LagunaDFlashTarget(w_, cache_, backend_); if (kvflash_active()) dflash_target_->set_kvflash_pager(&kvflash_pager_); @@ -2828,14 +2899,41 @@ bool LagunaBackend::load_decode_draft() { return true; } +bool LagunaBackend::select_decode_draft(const std::string & name) { + std::string wanted = name; + if (wanted.empty()) { + wanted = default_draft_lora_; + } + for (LagunaDraftVariant & variant : draft_variants_) { + if (variant.name == wanted) { + if (active_dw_ != &variant.weights) { + std::fprintf(stderr, "[laguna] selected draft LoRA variant: %s\n", + variant.name.c_str()); + } + active_dw_ = &variant.weights; + active_draft_lora_ = variant.name; + return true; + } + } + std::fprintf(stderr, "[laguna] unknown draft LoRA variant '%s'\n", + wanted.c_str()); + return false; +} + void LagunaBackend::free_decode_draft() { delete dflash_target_; dflash_target_ = nullptr; draft_feature_mirror_free(feature_mirror_); free_laguna_target_feat(cache_); - if (dw_.ctx) { - free_draft_weights(dw_); + for (LagunaDraftVariant & variant : draft_variants_) { + if (variant.weights.ctx) { + free_draft_weights(variant.weights); + } } + draft_variants_.clear(); + active_dw_ = nullptr; + active_draft_lora_.clear(); + default_draft_lora_ = "base"; if (draft_backend_ && draft_backend_ != backend_) { ggml_backend_free(draft_backend_); } diff --git a/server/src/laguna/laguna_backend.h b/server/src/laguna/laguna_backend.h index 9a651bf8a..c5c824dc0 100644 --- a/server/src/laguna/laguna_backend.h +++ b/server/src/laguna/laguna_backend.h @@ -36,6 +36,7 @@ namespace dflash::common { struct LagunaBackendArgs { std::string target_path; std::string draft_path; + std::vector draft_loras; int draft_gpu = -1; int draft_ctx_max = 2048; bool ddtree_mode = false; @@ -48,6 +49,12 @@ struct LagunaBackendArgs { ggml_type kv_type = GGML_TYPE_Q8_0; }; +struct LagunaDraftVariant { + std::string name; + std::vector loras; + DraftWeights weights; +}; + class LagunaBackend : public ModelBackend { public: explicit LagunaBackend(const LagunaBackendArgs & args); @@ -100,7 +107,10 @@ class LagunaBackend : public ModelBackend { // DFlash speculative decode ggml_backend_t draft_backend_ = nullptr; - DraftWeights dw_{}; + std::vector draft_variants_; + DraftWeights * active_dw_ = nullptr; + std::string active_draft_lora_; + std::string default_draft_lora_ = "base"; DraftFeatureMirror feature_mirror_{}; LagunaDFlashTarget * dflash_target_ = nullptr; bool draft_parked_ = false; @@ -178,6 +188,7 @@ class LagunaBackend : public ModelBackend { void maybe_post_request_swap(); bool load_decode_draft(); + bool select_decode_draft(const std::string & name); void free_decode_draft(); bool do_spec_decode(int committed, int n_gen, std::vector & out_tokens, diff --git a/server/src/server/http_server.cpp b/server/src/server/http_server.cpp index 8d4bdab24..bf82388a2 100644 --- a/server/src/server/http_server.cpp +++ b/server/src/server/http_server.cpp @@ -1688,6 +1688,7 @@ bool HttpServer::route_request(int fd, const HttpRequest & hr) { // Bandit: parse session_id from extra_body (opt-in adaptive keep_ratio) req.session_id = parse_session_id_from_body(body); + req.draft_lora = parse_draft_lora_from_body(body); // Serialize tools JSON for template injection. std::string tools_json; @@ -2392,6 +2393,7 @@ void HttpServer::worker_loop() { gen_req.sampler = req.sampler; gen_req.do_sample = req.sampler.needs_logit_processing(); gen_req.stream = false; // we handle streaming via on_token callback + gen_req.draft_lora = req.draft_lora; // Level 2 force-close: when thinking is opted in, the server is // configured with a hard-limit reply budget, and we resolved the diff --git a/server/src/server/http_server.h b/server/src/server/http_server.h index 957741b4d..732047b4d 100644 --- a/server/src/server/http_server.h +++ b/server/src/server/http_server.h @@ -225,6 +225,8 @@ struct ParsedRequest { std::vector stop_sequences; // Bandit: per-session adaptive keep_ratio opt-in std::string session_id; + // Optional DFlash drafter LoRA variant selected by extra_body.draft_lora. + std::string draft_lora; DiskPrefixCachePolicy disk_cache_policy; }; @@ -399,4 +401,18 @@ inline std::string parse_session_id_from_body(const json & body) { return {}; } +inline std::string parse_draft_lora_from_body(const json & body) { + if (body.contains("extra_body")) { + const auto & eb = body["extra_body"]; + if (eb.is_object() && eb.contains("draft_lora") && + eb["draft_lora"].is_string()) { + return eb["draft_lora"].get(); + } + } + if (body.contains("draft_lora") && body["draft_lora"].is_string()) { + return body["draft_lora"].get(); + } + return {}; +} + } // namespace dflash::common diff --git a/server/src/server/server_main.cpp b/server/src/server/server_main.cpp index 99abfb99d..f575999d7 100644 --- a/server/src/server/server_main.cpp +++ b/server/src/server/server_main.cpp @@ -33,6 +33,7 @@ #include #include #include +#include #include #ifdef _WIN32 @@ -195,7 +196,10 @@ static void print_usage(const char * prog) { "Usage: %s [options]\n" "\n" "Options:\n" - " --draft Draft model for speculative decode (qwen35 only)\n" + " --draft Draft model for speculative decode\n" + " --draft-lora \n" + " LoRA GGUF adapter merged into a switchable draft variant\n" + " Request selects a named variant with extra_body.draft_lora\n" " --port Listen port (default: 8080)\n" " --host Bind address (default: 0.0.0.0)\n" " --max-ctx Max context length (default: 131072)\n" @@ -309,6 +313,20 @@ static void print_usage(const char * prog) { "\n", prog); } +static DraftLoraSpec parse_draft_lora_spec(const char * arg) { + DraftLoraSpec spec; + std::string s = arg ? arg : ""; + const size_t eq = s.find('='); + if (eq != std::string::npos) { + spec.name = s.substr(0, eq); + spec.path = s.substr(eq + 1); + } else { + spec.path = std::move(s); + } + spec.scale = 1.0f; + return spec; +} + int main(int argc, char ** argv) { if (argc < 2 || argv[1][0] == '-') { print_usage(argv[0]); @@ -354,6 +372,8 @@ int main(int argc, char ** argv) { for (int i = 2; i < argc; i++) { if (std::strcmp(argv[i], "--draft") == 0 && i + 1 < argc) { bargs.draft_path = argv[++i]; + } else if (std::strcmp(argv[i], "--draft-lora") == 0 && i + 1 < argc) { + bargs.draft_loras.push_back(parse_draft_lora_spec(argv[++i])); } else if (std::strcmp(argv[i], "--port") == 0 && i + 1 < argc) { sconfig.port = std::atoi(argv[++i]); } else if (std::strcmp(argv[i], "--host") == 0 && i + 1 < argc) { @@ -628,6 +648,26 @@ int main(int argc, char ** argv) { if (!validate_server_placement(bargs, sconfig)) return 2; + if (!bargs.draft_loras.empty() && !bargs.draft_path) { + std::fprintf(stderr, "[server] --draft-lora requires --draft\n"); + return 2; + } + if (!bargs.draft_loras.empty() && bargs.remote_draft.enabled()) { + std::fprintf(stderr, + "[server] --draft-lora is not supported with --draft-ipc-bin yet\n"); + return 2; + } + for (const DraftLoraSpec & spec : bargs.draft_loras) { + if (spec.path.empty()) { + std::fprintf(stderr, "[server] --draft-lora path must not be empty\n"); + return 2; + } + if (spec.name == "base") { + std::fprintf(stderr, "[server] draft LoRA name 'base' is reserved\n"); + return 2; + } + } + if (bargs.remote_draft.enabled() && bargs.draft_path) { const std::string arch = detect_arch(bargs.model_path); if (arch.empty()) { @@ -995,6 +1035,17 @@ int main(int argc, char ** argv) { std::fprintf(stderr, "[server] │ port = %d\n", sconfig.port); std::fprintf(stderr, "[server] │ model = %s\n", bargs.model_path); std::fprintf(stderr, "[server] │ draft = %s\n", bargs.draft_path ? bargs.draft_path : "(none)"); + if (!bargs.draft_loras.empty()) { + for (size_t i = 0; i < bargs.draft_loras.size(); ++i) { + std::fprintf(stderr, "[server] │ draft_lora[%zu] = %s%s%s\n", + i, + bargs.draft_loras[i].name.empty() + ? "" + : bargs.draft_loras[i].name.c_str(), + bargs.draft_loras[i].name.empty() ? "" : "=", + bargs.draft_loras[i].path.c_str()); + } + } std::fprintf(stderr, "[server] │ model_name = %s\n", sconfig.model_name.c_str()); std::fprintf(stderr, "[server] │ max_ctx = %d\n", sconfig.max_ctx); // max_tokens default for requests that omit the field. The request