diff --git a/server/CMakeLists.txt b/server/CMakeLists.txt index 1b894c0af..0ea784c1c 100644 --- a/server/CMakeLists.txt +++ b/server/CMakeLists.txt @@ -258,6 +258,7 @@ add_library(dflash_common STATIC src/laguna/laguna_dflash_target.cpp src/common/backend_ipc.cpp src/common/domino_head.cpp + src/common/dspark_head.cpp src/common/target_shard_ipc.cpp src/common/target_shard_ipc_daemon.cpp src/common/dflash_feature_ring.cpp diff --git a/server/deps/llama.cpp b/server/deps/llama.cpp index ac06e5431..02cc5286c 160000 --- a/server/deps/llama.cpp +++ b/server/deps/llama.cpp @@ -1 +1 @@ -Subproject commit ac06e5431da1768e6838a5616d60c6053f6af19d +Subproject commit 02cc5286c431273efb7f4177f733aee042020d00 diff --git a/server/scripts/convert_dflash_to_gguf.py b/server/scripts/convert_dflash_to_gguf.py index de26810bb..5d71a393b 100644 --- a/server/scripts/convert_dflash_to_gguf.py +++ b/server/scripts/convert_dflash_to_gguf.py @@ -244,6 +244,17 @@ def bytes_to_np(raw: bytes, dtype: str, shape: list[int]) -> np.ndarray: } +DSPARK_TENSOR_MAP = { + "dspark_markov_head.markov_w1.weight": ("dflash.dspark.markov.w1", gguf.GGMLQuantizationType.F16), + "dspark_markov_head.markov_w2.weight": ("dflash.dspark.markov.w2", gguf.GGMLQuantizationType.F16), +} + +DSPARK_CONFIDENCE_TENSOR_MAP = { + "dspark_confidence_head.weight": ("dflash.dspark.confidence.weight", gguf.GGMLQuantizationType.F16), + "dspark_confidence_head.bias": ("dflash.dspark.confidence.bias", gguf.GGMLQuantizationType.F32), +} + + def add_domino_aux_heads(writer, arch: str, aux_path: Path | None): if aux_path is None: return @@ -265,6 +276,9 @@ def add_domino_aux_heads(writer, arch: str, aux_path: Path | None): print(f"[error] Domino aux heads file is not a tensor dict: {aux_path}", file=sys.stderr) sys.exit(1) + if not any(k in state for k in DOMINO_TENSOR_MAP): + return + missing = [k for k in DOMINO_TENSOR_MAP if k not in state] if missing: print(f"[warn] incomplete Domino aux heads; missing {missing}; skipping Domino tensors") @@ -294,6 +308,79 @@ def add_domino_aux_heads(writer, arch: str, aux_path: Path | None): print(f"[tensor] {gguf_name:50s} aux ->{raw_dtype.name:4s} {tuple(arr.shape)}") +def add_dspark_aux_heads(writer, arch: str, aux_path: Path | None): + if aux_path is None: + return + if not aux_path.exists(): + return + + try: + import torch + except ImportError as exc: + print(f"[error] --aux-heads requires torch: {exc}", file=sys.stderr) + sys.exit(1) + + state = torch.load(aux_path, map_location="cpu") + if isinstance(state, dict) and "state_dict" in state and isinstance(state["state_dict"], dict): + state = state["state_dict"] + if not isinstance(state, dict): + print(f"[error] DSpark aux heads file is not a tensor dict: {aux_path}", file=sys.stderr) + sys.exit(1) + + missing = [k for k in DSPARK_TENSOR_MAP if k not in state] + if missing: + return + + print(f"[info] reading DSpark aux heads from {aux_path}") + w1 = state["dspark_markov_head.markov_w1.weight"] + w2 = state["dspark_markov_head.markov_w2.weight"] + vocab = int(w1.shape[0]) + rank = int(w1.shape[1]) + if tuple(w2.shape) != (vocab, rank): + print(f"[error] DSpark markov_w2 shape {tuple(w2.shape)} != {(vocab, rank)}", file=sys.stderr) + sys.exit(1) + + writer.add_uint32(f"{arch}.dflash.dspark.enabled", 1) + writer.add_uint32(f"{arch}.dflash.dspark.markov_rank", rank) + writer.add_uint32(f"{arch}.dflash.dspark.vocab_size", vocab) + + for st_name, (gguf_name, raw_dtype) in DSPARK_TENSOR_MAP.items(): + t = state[st_name] + if hasattr(t, "detach"): + t = t.detach().cpu() + arr = t.float().numpy().astype("{raw_dtype.name:4s} {tuple(arr.shape)}") + + conf_missing = [k for k in DSPARK_CONFIDENCE_TENSOR_MAP if k not in state] + if conf_missing: + print(f"[warn] incomplete DSpark confidence head; missing {conf_missing}; Markov head will still load") + return + + conf_w = state["dspark_confidence_head.weight"] + conf_b = state["dspark_confidence_head.bias"] + confidence_dim = int(conf_w.shape[1]) + if int(conf_w.shape[0]) != 1 or tuple(conf_b.shape) != (1,): + print( + f"[error] DSpark confidence shapes weight={tuple(conf_w.shape)} bias={tuple(conf_b.shape)}", + file=sys.stderr, + ) + sys.exit(1) + writer.add_uint32(f"{arch}.dflash.dspark.confidence_dim", confidence_dim) + writer.add_uint32(f"{arch}.dflash.dspark.confidence.enabled", 1) + for st_name, (gguf_name, raw_dtype) in DSPARK_CONFIDENCE_TENSOR_MAP.items(): + t = state[st_name] + if hasattr(t, "detach"): + t = t.detach().cpu() + arr = t.float().numpy() + if raw_dtype == gguf.GGMLQuantizationType.F16: + arr = arr.astype("{raw_dtype.name:4s} {tuple(arr.shape)}") + + # ────────────────────────────────────────────────────────────────────── # Main # ────────────────────────────────────────────────────────────────────── @@ -400,6 +487,7 @@ def sort_key(t): if not args.no_aux_heads: aux_path = args.aux_heads if args.aux_heads is not None else args.safetensors.parent / "dflash_aux_heads.pt" add_domino_aux_heads(writer, ARCH, aux_path) + add_dspark_aux_heads(writer, ARCH, aux_path) print(f"[info] writing {args.out_gguf}") writer.write_header_to_file() diff --git a/server/scripts/quantize_dflash_draft.py b/server/scripts/quantize_dflash_draft.py new file mode 100755 index 000000000..b81370750 --- /dev/null +++ b/server/scripts/quantize_dflash_draft.py @@ -0,0 +1,110 @@ +#!/usr/bin/env python3 +"""Requantize a DFlash draft GGUF (f16) for faster spec-decode drafting. + +Schemes: + q8_0 - all 2D matmul weights to q8_0 (~47% smaller, acceptance-neutral) + q4-mix - backbone (blk.*) weights to q4_0, dflash.* head weights kept at + q8_0 (~67% smaller; the head split protects the Markov/projection + bias precision that near-tie corrections depend on). + +Norms, biases and non-f16 tensors are copied as-is. Metadata is preserved, +so the output loads anywhere the input does. + +Measured on Laguna-XS.2 + v24 drafter (RTX 3090, HumanEval, verify width 6): +f16 236.6 tok/s -> q8_0 241.3 -> q4-mix 249.6, acceptance unchanged. +""" +import argparse +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "deps" / "llama.cpp" / "gguf-py")) + +import numpy as np # noqa: E402 +import gguf # noqa: E402 +from gguf import GGUFReader, GGUFWriter, GGMLQuantizationType # noqa: E402 +from gguf.quants import quantize # noqa: E402 + + +def pick_type(name: str, scheme: str) -> GGMLQuantizationType: + if scheme == "q8_0": + return GGMLQuantizationType.Q8_0 + # q4-mix: protect the dflash.* heads + if name.startswith("dflash."): + return GGMLQuantizationType.Q8_0 + return GGMLQuantizationType.Q4_0 + + +def copy_metadata(r: GGUFReader, w: GGUFWriter) -> None: + skip = {"GGUF.version", "GGUF.tensor_count", "GGUF.kv_count", "general.architecture"} + T = gguf.GGUFValueType + for f in r.fields.values(): + if f.name in skip: + continue + ftype = f.types[0] + val = f.parts[f.data[0]] + if ftype == T.STRING: + w.add_string(f.name, bytes(val).decode()) + elif ftype == T.ARRAY: + sub = f.types[1] + vals = [f.parts[i] for i in f.data] + if sub == T.STRING: + w.add_array(f.name, [bytes(v).decode() for v in vals]) + else: + w.add_array(f.name, [np.asarray(v)[0].item() for v in vals]) + elif ftype == T.BOOL: + w.add_bool(f.name, bool(val[0])) + elif ftype == T.FLOAT32: + w.add_float32(f.name, float(val[0])) + elif ftype == T.FLOAT64: + w.add_float64(f.name, float(val[0])) + else: + fn = {T.UINT32: w.add_uint32, T.INT32: w.add_int32, + T.UINT64: w.add_uint64, T.INT64: w.add_int64, + T.UINT8: w.add_uint8, T.INT8: w.add_int8, + T.UINT16: w.add_uint16, T.INT16: w.add_int16}[ftype] + fn(f.name, val[0].item()) + + +def main() -> int: + ap = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + ap.add_argument("input", help="f16 draft GGUF") + ap.add_argument("output", help="output GGUF path") + ap.add_argument("--scheme", choices=["q8_0", "q4-mix"], default="q4-mix") + args = ap.parse_args() + + r = GGUFReader(args.input) + arch = None + for f in r.fields.values(): + if f.name == "general.architecture": + arch = bytes(f.parts[f.data[0]]).decode() + if not arch: + print("error: no general.architecture in input", file=sys.stderr) + return 1 + + w = GGUFWriter(args.output, arch) + copy_metadata(r, w) + + n_q = n_keep = 0 + for t in r.tensors: + shape = [int(x) for x in t.shape] + if (t.tensor_type == GGMLQuantizationType.F16 and len(shape) == 2 + and shape[0] % 32 == 0 and "norm" not in t.name): + qt = pick_type(t.name, args.scheme) + arr = np.array(t.data, dtype=np.float16).reshape(shape[::-1]).astype(np.float32) + w.add_tensor(t.name, quantize(arr, qt), raw_dtype=qt) + n_q += 1 + else: + w.add_tensor(t.name, np.array(t.data), raw_dtype=t.tensor_type) + n_keep += 1 + + w.write_header_to_file() + w.write_kv_data_to_file() + w.write_tensors_to_file() + w.close() + print(f"{args.scheme}: quantized {n_q} tensors, kept {n_keep} as-is -> {args.output}") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/server/src/common/backend_factory.h b/server/src/common/backend_factory.h index c5dc454f3..73ed8698c 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 { diff --git a/server/src/common/dspark_head.cpp b/server/src/common/dspark_head.cpp new file mode 100644 index 000000000..4aea171b7 --- /dev/null +++ b/server/src/common/dspark_head.cpp @@ -0,0 +1,397 @@ +#include "dspark_head.h" + +#include "ggml-alloc.h" +#include "ddtree.h" + +#include +#include +#include +#include + +namespace dflash::common { + +namespace { + +bool dspark_step(const DraftWeights & dw, + ggml_backend_t backend, + int32_t prev_token, + const float * draft_hidden, + const float * base_logits, + int vocab, + int32_t & out_token, + float * confidence_out) { + const int hidden = dw.n_embd; + const int rank = dw.dspark.markov_rank; + if (hidden <= 0 || rank <= 0 || vocab <= 0) return false; + if (!dw.dspark.markov_w1 || !dw.dspark.markov_w2) return false; + + const bool want_conf = + confidence_out != nullptr && + dw.dspark.confidence_w != nullptr && + dw.dspark.confidence_b != nullptr && + dw.dspark.confidence_dim > 0; + + const size_t arena_size = + ggml_tensor_overhead() * 256 + ggml_graph_overhead() + 2 * 1024 * 1024; + static thread_local std::vector g_arena; + if (g_arena.size() < arena_size) g_arena.resize(arena_size); + + ggml_init_params ip{}; + ip.mem_size = arena_size; + ip.mem_buffer = g_arena.data(); + ip.no_alloc = true; + ggml_context * ctx = ggml_init(ip); + if (!ctx) return false; + ggml_cgraph * gf = ggml_new_graph_custom(ctx, 256, false); + + ggml_tensor * inp_prev = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 1); + ggml_tensor * inp_base = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, vocab, 1); + ggml_set_input(inp_prev); + ggml_set_input(inp_base); + + ggml_tensor * prev_emb = ggml_get_rows(ctx, dw.dspark.markov_w1, inp_prev); + ggml_tensor * bias = ggml_mul_mat(ctx, dw.dspark.markov_w2, prev_emb); + ggml_tensor * corrected = ggml_add(ctx, inp_base, bias); + ggml_tensor * tok = ggml_argmax(ctx, corrected); + ggml_set_output(tok); + ggml_build_forward_expand(gf, tok); + + ggml_tensor * conf = nullptr; + ggml_tensor * inp_hidden = nullptr; + if (want_conf) { + inp_hidden = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hidden, 1); + ggml_set_input(inp_hidden); + ggml_tensor * conf_in = inp_hidden; + if (dw.dspark.confidence_dim == hidden + rank) { + conf_in = ggml_concat(ctx, inp_hidden, prev_emb, 0); + } else if (dw.dspark.confidence_dim != hidden) { + ggml_free(ctx); + return false; + } + conf = ggml_mul_mat(ctx, dw.dspark.confidence_w, conf_in); + conf = ggml_add(ctx, conf, ggml_reshape_2d(ctx, dw.dspark.confidence_b, 1, 1)); + conf = ggml_sigmoid(ctx, conf); + ggml_set_output(conf); + ggml_build_forward_expand(gf, conf); + } + + static thread_local ggml_gallocr_t galloc = nullptr; + if (!galloc) { + galloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); + } + if (!ggml_gallocr_alloc_graph(galloc, gf)) { + std::fprintf(stderr, "dspark_step: gallocr_alloc_graph failed\n"); + ggml_free(ctx); + return false; + } + + ggml_backend_tensor_set(inp_prev, &prev_token, 0, sizeof(prev_token)); + ggml_backend_tensor_set(inp_base, base_logits, 0, sizeof(float) * (size_t)vocab); + if (want_conf) { + ggml_backend_tensor_set(inp_hidden, draft_hidden, 0, + sizeof(float) * (size_t)hidden); + } + + if (ggml_backend_graph_compute(backend, gf) != GGML_STATUS_SUCCESS) { + std::fprintf(stderr, "dspark_step: graph_compute failed\n"); + ggml_free(ctx); + return false; + } + + ggml_backend_tensor_get(tok, &out_token, 0, sizeof(out_token)); + if (want_conf) { + ggml_backend_tensor_get(conf, confidence_out, 0, sizeof(float)); + } + ggml_free(ctx); + return true; +} + +} // namespace + +bool dspark_markov_correct_greedy_chain(const DraftWeights & dw, + ggml_backend_t backend, + DFlashTarget & target, + const float * local_hidden, + int q_len, + int32_t last_tok, + float confidence_threshold, + std::vector & draft_tok) { + if (!dw.dspark.enabled || q_len <= 1 || !local_hidden) return false; + const int hidden = dw.n_embd; + const int n_candidates = q_len - 1; + if (hidden <= 0 || n_candidates <= 0) return false; + if (confidence_threshold < 0.0f) confidence_threshold = 0.0f; + if (confidence_threshold > 1.0f) confidence_threshold = 1.0f; + const bool use_confidence_gate = + confidence_threshold > 0.0f && + dw.dspark.confidence_w != nullptr && + dw.dspark.confidence_b != nullptr && + dw.dspark.confidence_dim > 0; + + std::vector candidate_hidden((size_t)n_candidates * (size_t)hidden); + for (int i = 0; i < n_candidates; ++i) { + const float * src = local_hidden + (size_t)(i + 1) * (size_t)hidden; + std::memcpy(candidate_hidden.data() + (size_t)i * (size_t)hidden, + src, sizeof(float) * (size_t)hidden); + } + + std::vector base_logits; + if (!target.project_hidden_to_logits(candidate_hidden.data(), n_candidates, base_logits)) { + return false; + } + if (base_logits.size() % (size_t)n_candidates != 0) return false; + const int vocab = (int)(base_logits.size() / (size_t)n_candidates); + if (dw.dspark.vocab_size > 0 && vocab != dw.dspark.vocab_size) { + std::fprintf(stderr, "dspark_markov_correct_greedy_chain: vocab mismatch target=%d dspark=%d\n", + vocab, dw.dspark.vocab_size); + return false; + } + + draft_tok.clear(); + draft_tok.reserve((size_t)q_len); + draft_tok.push_back(last_tok); + int32_t prefix_tok = last_tok; + for (int i = 0; i < n_candidates; ++i) { + int32_t tok = -1; + float confidence = 0.0f; + float * confidence_ptr = use_confidence_gate ? &confidence : nullptr; + if (!dspark_step(dw, backend, prefix_tok, + candidate_hidden.data() + (size_t)i * (size_t)hidden, + base_logits.data() + (size_t)i * (size_t)vocab, + vocab, + tok, + confidence_ptr)) { + return false; + } + if (use_confidence_gate && confidence < confidence_threshold) { + break; + } + draft_tok.push_back(tok); + prefix_tok = tok; + } + return true; +} + +namespace { + +struct MarkovChainGraph { + ggml_context * ctx = nullptr; + ggml_cgraph * gf = nullptr; + ggml_tensor * inp_hidden = nullptr; + ggml_tensor * inp_seed = nullptr; + ggml_tensor * base = nullptr; // [vocab, n_positions] + std::vector toks; // corrected argmax per depth + std::vector corrected; // corrected logits per depth +}; + +// Guards shared by the fused Markov paths: head present, usable inputs, and +// the target lm_head vocab matching the head's training vocab. +bool dspark_fused_usable(const DraftWeights & dw, ggml_backend_t backend, + ggml_tensor * lm_head, const float * hidden, + const char * who) { + if (!dw.dspark.enabled || !hidden || !backend || !lm_head) return false; + if (!dw.dspark.markov_w1 || !dw.dspark.markov_w2) return false; + if (dw.n_embd <= 0 || dw.dspark.markov_rank <= 0) return false; + const int vocab = (int)lm_head->ne[1]; + if (vocab <= 0) return false; + if (dw.dspark.vocab_size > 0 && vocab != dw.dspark.vocab_size) { + static bool s_vocab_warned = false; + if (!s_vocab_warned) { + s_vocab_warned = true; + std::fprintf(stderr, "%s: vocab mismatch lm_head=%d dspark=%d; falling back\n", + who, vocab, dw.dspark.vocab_size); + } + return false; + } + return true; +} + +// One graph: base logits for all n_positions hidden columns (a single lm_head +// matmul), then for rows [first_corrected, n_positions) the low-rank Markov +// correction chained along the main path - +// bias_i = markov_w2 . markov_w1[prev] +// corrected_i = base_i + bias_i +// tok_i = argmax(corrected_i) (feeds the next step's get_rows) +// The chain seed is an I32 graph input; markov_w1 doubles as the previous- +// token embedding table. Rows below first_corrected keep the uncorrected base. +bool build_markov_chain_graph(const DraftWeights & dw, + ggml_tensor * lm_head, + int n_positions, int first_corrected, + bool corrected_are_outputs, + std::vector & arena, + MarkovChainGraph & out) { + const int hdim = dw.n_embd; + const int vocab = (int)lm_head->ne[1]; + const int n_corr = n_positions - first_corrected; + if (n_positions <= 0 || n_corr <= 0) return false; + + const size_t arena_size = ggml_tensor_overhead() * (size_t)(64 + 16 * n_corr) + + ggml_graph_overhead_custom(512, false) + 2 * 1024 * 1024; + if (arena.size() < arena_size) arena.resize(arena_size); + + ggml_init_params ip{}; + ip.mem_size = arena.size(); + ip.mem_buffer = arena.data(); + ip.no_alloc = true; + out.ctx = ggml_init(ip); + if (!out.ctx) return false; + out.gf = ggml_new_graph_custom(out.ctx, 512, false); + + out.inp_hidden = ggml_new_tensor_2d(out.ctx, GGML_TYPE_F32, hdim, n_positions); + out.inp_seed = ggml_new_tensor_1d(out.ctx, GGML_TYPE_I32, 1); + ggml_set_input(out.inp_hidden); + ggml_set_input(out.inp_seed); + + out.base = ggml_mul_mat(out.ctx, lm_head, out.inp_hidden); + if (first_corrected > 0) { + // The uncorrected rows are read back by the caller. + ggml_set_output(out.base); + ggml_build_forward_expand(out.gf, out.base); + } + + ggml_tensor * prev_ids = out.inp_seed; + out.toks.assign((size_t)n_corr, nullptr); + out.corrected.assign((size_t)n_corr, nullptr); + for (int i = 0; i < n_corr; ++i) { + const int row = first_corrected + i; + ggml_tensor * prev_emb = ggml_get_rows(out.ctx, dw.dspark.markov_w1, prev_ids); + ggml_tensor * bias = ggml_mul_mat(out.ctx, dw.dspark.markov_w2, prev_emb); + ggml_tensor * base_i = ggml_view_2d(out.ctx, out.base, vocab, 1, + out.base->nb[1], (size_t)row * out.base->nb[1]); + ggml_tensor * corrected = ggml_add(out.ctx, base_i, bias); + if (corrected_are_outputs) { + ggml_set_output(corrected); + ggml_build_forward_expand(out.gf, corrected); + } + ggml_tensor * tok = ggml_argmax(out.ctx, corrected); + ggml_set_output(tok); + ggml_build_forward_expand(out.gf, tok); + out.corrected[(size_t)i] = corrected; + out.toks[(size_t)i] = tok; + prev_ids = tok; + } + return true; +} + +} // namespace + +bool dspark_markov_correct_greedy_chain_fused(const DraftWeights & dw, + ggml_backend_t backend, + ggml_tensor * lm_head, + const float * local_hidden, + int q_len, + int32_t last_tok, + std::vector & draft_tok) { + if (q_len <= 1) return false; + if (!dspark_fused_usable(dw, backend, lm_head, local_hidden, "dspark_fused")) return false; + const int hdim = dw.n_embd; + const int n_cand = q_len - 1; + + static thread_local std::vector g_arena_chain; + MarkovChainGraph g; + if (!build_markov_chain_graph(dw, lm_head, n_cand, /*first_corrected=*/0, + /*corrected_are_outputs=*/false, g_arena_chain, g)) { + return false; + } + + static thread_local ggml_gallocr_t galloc_chain = nullptr; + if (!galloc_chain) { + galloc_chain = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); + } + if (!ggml_gallocr_alloc_graph(galloc_chain, g.gf)) { + std::fprintf(stderr, "dspark_fused: gallocr_alloc_graph failed\n"); + ggml_free(g.ctx); + return false; + } + + // Candidate hidden states start at position 1 (position 0 is the seed). + ggml_backend_tensor_set(g.inp_hidden, local_hidden + (size_t)hdim, 0, + sizeof(float) * (size_t)hdim * (size_t)n_cand); + ggml_backend_tensor_set(g.inp_seed, &last_tok, 0, sizeof(int32_t)); + + if (ggml_backend_graph_compute(backend, g.gf) != GGML_STATUS_SUCCESS) { + std::fprintf(stderr, "dspark_fused: graph_compute failed\n"); + ggml_free(g.ctx); + return false; + } + + draft_tok.assign((size_t)q_len, 0); + draft_tok[0] = last_tok; + // One synchronize instead of n_cand blocking readbacks. + int32_t t_out[16]; + const int n_get = n_cand < 16 ? n_cand : 16; + for (int i = 0; i < n_get; ++i) { + ggml_backend_tensor_get_async(backend, g.toks[(size_t)i], &t_out[i], 0, sizeof(int32_t)); + } + ggml_backend_synchronize(backend); + for (int i = 0; i < n_get; ++i) { + draft_tok[(size_t)i + 1] = t_out[i]; + } + ggml_free(g.ctx); + return true; +} + +bool dspark_markov_project_topk(const DraftWeights & dw, + ggml_backend_t backend, + ggml_tensor * lm_head, + const float * hidden, + int n_tokens, int K, float temperature, + int32_t last_tok, + std::vector & top_log_probs, + std::vector & top_token_ids) { + if (n_tokens <= 1 || K <= 0) return false; + if (!dspark_fused_usable(dw, backend, lm_head, hidden, "dspark_topk")) return false; + const int hdim = dw.n_embd; + const int vocab = (int)lm_head->ne[1]; + + static thread_local std::vector g_arena_topk; + MarkovChainGraph g; + if (!build_markov_chain_graph(dw, lm_head, n_tokens, /*first_corrected=*/1, + /*corrected_are_outputs=*/true, g_arena_topk, g)) { + return false; + } + + static thread_local ggml_gallocr_t galloc_topk = nullptr; + if (!galloc_topk) { + galloc_topk = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); + } + if (!ggml_gallocr_alloc_graph(galloc_topk, g.gf)) { + std::fprintf(stderr, "dspark_topk: gallocr_alloc_graph failed\n"); + ggml_free(g.ctx); + return false; + } + + ggml_backend_tensor_set(g.inp_hidden, hidden, 0, + sizeof(float) * (size_t)hdim * (size_t)n_tokens); + ggml_backend_tensor_set(g.inp_seed, &last_tok, 0, sizeof(int32_t)); + + if (ggml_backend_graph_compute(backend, g.gf) != GGML_STATUS_SUCCESS) { + std::fprintf(stderr, "dspark_topk: graph_compute failed\n"); + ggml_free(g.ctx); + return false; + } + + // Corrected logits per row (row 0 keeps the uncorrected base), then the + // same host top-K extraction as project_hidden_to_topk for identical + // budget-allocation semantics in build_ddtree. + std::vector logits_host((size_t)vocab * (size_t)n_tokens); + ggml_backend_tensor_get_async(backend, g.base, logits_host.data(), 0, + sizeof(float) * (size_t)vocab); + const int n_corr = n_tokens - 1; + for (int i = 0; i < n_corr; ++i) { + ggml_backend_tensor_get_async(backend, g.corrected[(size_t)i], + logits_host.data() + (size_t)(i + 1) * (size_t)vocab, + 0, sizeof(float) * (size_t)vocab); + } + ggml_backend_synchronize(backend); + + top_log_probs.assign((size_t)n_tokens * (size_t)K, 0.0f); + top_token_ids.assign((size_t)n_tokens * (size_t)K, 0); + extract_draft_topk(logits_host.data(), n_tokens, vocab, K, + top_log_probs.data(), top_token_ids.data(), temperature); + + ggml_free(g.ctx); + return true; +} + +} // namespace dflash::common diff --git a/server/src/common/dspark_head.h b/server/src/common/dspark_head.h new file mode 100644 index 000000000..b76815ce5 --- /dev/null +++ b/server/src/common/dspark_head.h @@ -0,0 +1,47 @@ +#pragma once + +#include "dflash_target.h" +#include "internal.h" + +#include +#include + +namespace dflash::common { + +bool dspark_markov_correct_greedy_chain(const DraftWeights & dw, + ggml_backend_t backend, + DFlashTarget & target, + const float * local_hidden, + int q_len, + int32_t last_tok, + float confidence_threshold, + std::vector & draft_tok); + +// Fused variant: base logits (one lm_head matmul over all candidates) + +// unrolled Markov correction chain + in-graph argmax feeding the next +// step's get_rows, all in ONE graph on the draft backend. No host logits +// round-trip. Does not implement the confidence gate; callers wanting +// confidence-prefix truncation must use the unfused path. +bool dspark_markov_correct_greedy_chain_fused(const DraftWeights & dw, + ggml_backend_t backend, + ggml_tensor * lm_head, + const float * local_hidden, + int q_len, + int32_t last_tok, + std::vector & draft_tok); + +// DDTree candidate generation with the Markov correction: base logits for +// all n_tokens positions in ONE lm_head matmul; rows 1..n-1 get the low-rank +// previous-token bias chained along the main (argmax) path; top-K extracted +// on host via extract_draft_topk. Output contract matches +// DFlashTarget::project_hidden_to_topk (row 0 = seed position, uncorrected). +bool dspark_markov_project_topk(const DraftWeights & dw, + ggml_backend_t backend, + ggml_tensor * lm_head, + const float * hidden, + int n_tokens, int K, float temperature, + int32_t last_tok, + std::vector & top_log_probs, + std::vector & top_token_ids); + +} // namespace dflash::common diff --git a/server/src/draft/draft_gguf_loader.cpp b/server/src/draft/draft_gguf_loader.cpp index e8a097531..2694636ec 100644 --- a/server/src/draft/draft_gguf_loader.cpp +++ b/server/src/draft/draft_gguf_loader.cpp @@ -28,11 +28,15 @@ #include "common/gguf_mmap.h" #include "common/gguf_bounds.h" +#include #include #include #include #include +#include #include +#include +#include #if !defined(_WIN32) #include @@ -149,6 +153,10 @@ bool load_draft_gguf(const std::string & path, const uint32_t domino_meta_gru = read_u32("dflash.domino.gru_hidden_dim", 0); const uint32_t domino_meta_emb = read_u32("dflash.domino.emb_dim", 0); const uint32_t domino_meta_vocab = read_u32("dflash.domino.vocab_size", 0); + const uint32_t dspark_meta_enabled = read_u32("dflash.dspark.enabled", 0); + const uint32_t dspark_meta_rank = read_u32("dflash.dspark.markov_rank", 0); + const uint32_t dspark_meta_vocab = read_u32("dflash.dspark.vocab_size", 0); + const uint32_t dspark_meta_conf = read_u32("dflash.dspark.confidence_dim", 0); // Explicit captured target-layer ids (data-driven). Lets any DFlash drafter // load without a hardcoded per-arch set; the array length also backstops // n_target_layers when the scalar KV is absent. @@ -335,6 +343,61 @@ bool load_draft_gguf(const std::string & path, out.domino.vocab_size); } + out.dspark = DraftDSparkWeights{}; + out.dspark.markov_w1 = g("dflash.dspark.markov.w1"); + out.dspark.markov_w2 = g("dflash.dspark.markov.w2"); + out.dspark.confidence_w = g("dflash.dspark.confidence.weight"); + out.dspark.confidence_b = g("dflash.dspark.confidence.bias"); + + const bool dspark_any = + out.dspark.markov_w1 || out.dspark.markov_w2 || + out.dspark.confidence_w || out.dspark.confidence_b || + dspark_meta_enabled != 0; + if (dspark_any) { + if (!out.dspark.markov_w1 || !out.dspark.markov_w2) { + set_last_error("draft GGUF: incomplete DSpark Markov tensors"); + gguf_free(gctx); + return false; + } + out.dspark.markov_rank = + dspark_meta_rank != 0 ? (int)dspark_meta_rank : (int)out.dspark.markov_w1->ne[0]; + out.dspark.vocab_size = + dspark_meta_vocab != 0 ? (int)dspark_meta_vocab : (int)out.dspark.markov_w1->ne[1]; + + const int64_t R = out.dspark.markov_rank; + const int64_t V = out.dspark.vocab_size; + char shape_err[320]; + if (!check_shape_2d(out.dspark.markov_w1, R, V, "dspark.markov.w1", shape_err, sizeof(shape_err)) || + !check_shape_2d(out.dspark.markov_w2, R, V, "dspark.markov.w2", shape_err, sizeof(shape_err))) { + set_last_error(shape_err); + gguf_free(gctx); + return false; + } + + const bool conf_any = out.dspark.confidence_w || out.dspark.confidence_b || dspark_meta_conf != 0; + if (conf_any) { + if (!out.dspark.confidence_w || !out.dspark.confidence_b) { + set_last_error("draft GGUF: incomplete DSpark confidence tensors"); + gguf_free(gctx); + return false; + } + out.dspark.confidence_dim = + dspark_meta_conf != 0 ? (int)dspark_meta_conf : (int)out.dspark.confidence_w->ne[0]; + const int64_t C = out.dspark.confidence_dim; + if (!check_shape_2d(out.dspark.confidence_w, C, 1, "dspark.confidence.weight", shape_err, sizeof(shape_err)) || + !check_shape_1d(out.dspark.confidence_b, 1, "dspark.confidence.bias", shape_err, sizeof(shape_err))) { + set_last_error(shape_err); + gguf_free(gctx); + return false; + } + } + + out.dspark.enabled = true; + std::fprintf(stderr, "[draft GGUF] DSpark Markov head enabled: rank=%d vocab=%d confidence_dim=%d\n", + out.dspark.markov_rank, out.dspark.vocab_size, + out.dspark.confidence_dim); + } + // GGUF Qwen3.6 drafters carry SWA metadata emitted by the converter: // dflash-draft.attention.sliding_window = 2048 // dflash-draft.attention.sliding_window_pattern = [true,true,true,true,false] @@ -386,7 +449,8 @@ bool load_draft_gguf(const std::string & path, gguf_free(gctx); return false; } - ggml_backend_tensor_set(t, mm_addr + data_start + rel_off, 0, sz); + const uint8_t * tensor_bytes = mm_addr + data_start + rel_off; + ggml_backend_tensor_set(t, tensor_bytes, 0, sz); total += sz; } diff --git a/server/src/internal.h b/server/src/internal.h index 89b0f162e..478b66621 100644 --- a/server/src/internal.h +++ b/server/src/internal.h @@ -263,6 +263,18 @@ struct DraftDominoWeights { ggml_tensor * head_b2 = nullptr; // [vocab_size] f32 }; +struct DraftDSparkWeights { + bool enabled = false; + int markov_rank = 0; + int vocab_size = 0; + int confidence_dim = 0; + + ggml_tensor * markov_w1 = nullptr; // [markov_rank, vocab_size] + ggml_tensor * markov_w2 = nullptr; // [markov_rank, vocab_size] + ggml_tensor * confidence_w = nullptr; // [confidence_dim, 1] + ggml_tensor * confidence_b = nullptr; // [1] f32 +}; + struct DraftWeights { ggml_context * ctx = nullptr; ggml_backend_t backend = nullptr; @@ -301,6 +313,10 @@ struct DraftWeights { // speculative decode corrects each draft token with a lightweight GRU // conditioned on the realized prefix before target verification. DraftDominoWeights domino; + + // Optional DSpark/DeepSpec-style Markov correction head. When present, + // greedy chain decode adds a low-rank previous-token bias before argmax. + DraftDSparkWeights dspark; }; bool load_draft_safetensors(const std::string & path, diff --git a/server/src/kv_quant.cpp b/server/src/kv_quant.cpp index 676a103d2..996794a5d 100644 --- a/server/src/kv_quant.cpp +++ b/server/src/kv_quant.cpp @@ -157,6 +157,19 @@ bool is_supported_kv_pair(ggml_type k, ggml_type v) { // ─── Environment-variable resolution ──────────────────────────────────────── +void validate_kv_pair_or_abort(ggml_type k, ggml_type v, const char * who) { + if (is_supported_kv_pair(k, v)) return; + std::fprintf(stderr, + "%s KV pair (K=%s, V=%s) not supported by fattn-cuda. Supported pairs:\n", + who, kv_type_name(k), kv_type_name(v)); + for (int i = 0; i < N_SUPPORTED_PAIRS; ++i) { + std::fprintf(stderr, " K=%-6s V=%s\n", + kv_type_name(SUPPORTED_PAIRS[i].k), + kv_type_name(SUPPORTED_PAIRS[i].v)); + } + std::abort(); +} + void resolve_kv_types(ggml_type & k_out, ggml_type & v_out) { ggml_type k = GGML_TYPE_Q4_0; ggml_type v = GGML_TYPE_Q4_0; @@ -191,17 +204,7 @@ void resolve_kv_types(ggml_type & k_out, ggml_type & v_out) { } // Validate the resolved (K, V) pair - if (!is_supported_kv_pair(k, v)) { - std::fprintf(stderr, - "[dflash] KV pair (K=%s, V=%s) not supported by fattn-cuda. Supported pairs:\n", - kv_type_name(k), kv_type_name(v)); - for (int i = 0; i < N_SUPPORTED_PAIRS; ++i) { - std::fprintf(stderr, " K=%-6s V=%s\n", - kv_type_name(SUPPORTED_PAIRS[i].k), - kv_type_name(SUPPORTED_PAIRS[i].v)); - } - std::abort(); - } + validate_kv_pair_or_abort(k, v, "[dflash]"); k_out = k; v_out = v; diff --git a/server/src/kv_quant.h b/server/src/kv_quant.h index 9bd9b9f70..e18e5f0cc 100644 --- a/server/src/kv_quant.h +++ b/server/src/kv_quant.h @@ -19,6 +19,10 @@ const char * kv_type_name(ggml_type t); // GGML_CUDA_FA_ALL_QUANTS=ON, which is now forced ON in dflash/CMakeLists.txt). bool is_supported_kv_pair(ggml_type k, ggml_type v); +// Aborts with the supported-pairs listing when (k, v) is not supported by +// the compiled fattn kernels. `who` is the log prefix (e.g. "[laguna]"). +void validate_kv_pair_or_abort(ggml_type k, ggml_type v, const char * who); + // Resolves K and V types from environment variables. // Precedence (high -> low): // 1. DFLASH27B_KV_K= / DFLASH27B_KV_V= (independent override) diff --git a/server/src/laguna/laguna_backend.cpp b/server/src/laguna/laguna_backend.cpp index 4268da74c..2a1ea32fb 100644 --- a/server/src/laguna/laguna_backend.cpp +++ b/server/src/laguna/laguna_backend.cpp @@ -12,6 +12,7 @@ #include "dflash27b.h" #include "common/ddtree.h" #include "common/domino_head.h" +#include "common/dspark_head.h" #include "common/dflash_feature_ring.h" #include "common/dflash_draft_graph.h" #include "kv_quant.h" @@ -46,21 +47,35 @@ namespace dflash::common { namespace { -static bool laguna_kv_type_env_present() { - return std::getenv("DFLASH27B_KV_K") || - std::getenv("DFLASH27B_KV_V") || - std::getenv("DFLASH27B_KV_F16") || - std::getenv("DFLASH27B_KV_Q4") || - std::getenv("DFLASH27B_KV_TQ3"); -} - +// Laguna honors only the explicit per-axis --cache-type-k/v overrides. +// The DFLASH27B_KV_F16/_Q4/_TQ3 shorthands are qwen-family toggles - the +// server auto-sets _KV_TQ3 for max_ctx > 6144 - and must not displace +// laguna's Q8_0 default (a TQ3_0/Q4_0 KV cache garbles laguna output). static void resolve_laguna_kv_types(const LagunaBackendArgs & args, ggml_type & k_type, ggml_type & v_type) { k_type = args.kv_type; v_type = args.kv_type; - if (laguna_kv_type_env_present()) { - dflash::resolve_kv_types(k_type, v_type); + if (const char * s = std::getenv("DFLASH27B_KV_K")) { + const ggml_type parsed = dflash::parse_kv_type(s); + if (parsed == GGML_TYPE_COUNT) { + std::fprintf(stderr, "[laguna] Unknown KV K type: \"%s\"\n", s); + std::abort(); + } + k_type = parsed; + } + if (const char * s = std::getenv("DFLASH27B_KV_V")) { + const ggml_type parsed = dflash::parse_kv_type(s); + if (parsed == GGML_TYPE_COUNT) { + std::fprintf(stderr, "[laguna] Unknown KV V type: \"%s\"\n", s); + std::abort(); + } + v_type = parsed; + } + if (k_type != args.kv_type || v_type != args.kv_type) { + dflash::validate_kv_pair_or_abort(k_type, v_type, "[laguna]"); + std::fprintf(stderr, "[laguna] KV cache types overridden: K=%s V=%s\n", + dflash::kv_type_name(k_type), dflash::kv_type_name(v_type)); } } @@ -90,6 +105,26 @@ static bool laguna_sampled_verify_enabled(const SamplerCfg & sampler, bool do_sa return kSampledVerify && do_sample && sampler.needs_logit_processing(); } +static bool laguna_dspark_enabled() { + static const bool kEnabled = []() { + const char * e = std::getenv("DFLASH_LAGUNA_DSPARK"); + return e == nullptr || std::string(e) != "0"; + }(); + return kEnabled; +} + +static float laguna_dspark_confidence_threshold() { + static const float kThreshold = []() { + const char * e = std::getenv("DFLASH_LAGUNA_DSPARK_CONFIDENCE_THRESHOLD"); + if (!e) return 0.0f; + float threshold = std::atof(e); + if (threshold < 0.0f) threshold = 0.0f; + if (threshold > 1.0f) threshold = 1.0f; + return threshold; + }(); + return kThreshold; +} + // ── Construction / initialisation ─────────────────────────────────────── LagunaBackend::LagunaBackend(const LagunaBackendArgs & args) @@ -415,12 +450,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 @@ -454,7 +492,7 @@ bool LagunaBackend::do_spec_decode(int committed, int n_gen, if (chain_w < 2) chain_w = 2; if (chain_w > std::min(block_size, 8)) chain_w = std::min(block_size, 8); // DDTree sizes its batch via its budget; chain uses the width chosen above. - const int q_len = args_.ddtree_mode ? block_size : chain_w; + const int base_q_len = args_.ddtree_mode ? block_size : chain_w; const bool ignore_eos = (std::getenv("DFLASH_IGNORE_EOS") != nullptr); const bool sampled_verify = laguna_sampled_verify_enabled(sampler_, true); @@ -465,11 +503,11 @@ 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); - std::vector target_tok((size_t)q_len); + std::vector draft_tok((size_t)base_q_len); + std::vector target_tok((size_t)base_q_len); std::vector verify_logits; std::vector verify_history; std::vector pos_q((size_t)block_size); @@ -481,6 +519,7 @@ bool LagunaBackend::do_spec_decode(int committed, int n_gen, int n_generated = 0; int n_draft_steps = 0; int n_accept_sum = 0; + int n_draft_pos_sum = 0; auto argmax_logits = [](const std::vector & ll) { int best = 0; @@ -577,8 +616,20 @@ bool LagunaBackend::do_spec_decode(int committed, int n_gen, } auto t_dec0 = std::chrono::steady_clock::now(); + static const bool step_prof = std::getenv("DFLASH_LAGUNA_STEP_PROF") != nullptr; + double prof_draft_ms = 0.0, prof_heads_ms = 0.0, prof_verify_ms = 0.0; + auto prof_now = std::chrono::steady_clock::now(); + auto prof_lap = [&]() { + auto t = std::chrono::steady_clock::now(); + double ms = std::chrono::duration(t - prof_now).count(); + prof_now = t; + return ms; + }; while (n_generated < n_gen) { + int q_len = base_q_len; + draft_tok.resize((size_t)q_len); + target_tok.resize((size_t)q_len); const int need_commit_budget = n_gen - n_generated; if (budget_hook && !budget_hook->close_token_ids.empty()) { @@ -592,7 +643,7 @@ bool LagunaBackend::do_spec_decode(int committed, int n_gen, const bool ok = run_ar_tail(need_commit_budget); auto t_dec1 = std::chrono::steady_clock::now(); const double decode_s = std::chrono::duration(t_dec1 - t_dec0).count(); - const int total_draft_pos = std::max(1, n_draft_steps * q_len); + const int total_draft_pos = std::max(1, n_draft_pos_sum); const double accept_pct = 100.0 * (double)n_accept_sum / (double)total_draft_pos; std::fprintf(stderr, "[laguna-spec] tail-off-stats tokens=%d time=%.3f s speed=%.2f tok/s " @@ -626,7 +677,7 @@ bool LagunaBackend::do_spec_decode(int committed, int n_gen, const char * e = std::getenv("DFLASH_LAGUNA_DRAFT_PAD"); return !(e && e[0] == '0' && e[1] == '\0'); }(); - 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)), @@ -650,6 +701,7 @@ bool LagunaBackend::do_spec_decode(int committed, int n_gen, for (int i = 0; i < block_size; i++) pos_q[(size_t)i] = draft_ctx + i; for (int i = 0; i < kctx; i++) pos_k[(size_t)i] = (i < draft_ctx) ? i : 0; for (int j = 0; j < block_size; j++) pos_k[(size_t)kctx + j] = draft_ctx + j; + if (step_prof) prof_lap(); ggml_backend_tensor_set(draft_sg.positions, pos_q.data(), 0, sizeof(int32_t) * pos_q.size()); ggml_backend_tensor_set(draft_sg.positions_k, pos_k.data(), 0, @@ -664,15 +716,16 @@ bool LagunaBackend::do_spec_decode(int committed, int n_gen, local_hidden.resize((size_t)hidden * (size_t)q_len); ggml_backend_tensor_get(draft_sg.hidden_states, local_hidden.data(), 0, sizeof(float) * local_hidden.size()); + if (step_prof) prof_draft_ms += prof_lap(); 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); } static const bool fused_domino = []() { const char * e = std::getenv("DFLASH_LAGUNA_FUSED_DOMINO"); @@ -682,7 +735,7 @@ bool LagunaBackend::do_spec_decode(int committed, int n_gen, // Run on the draft backend: same stream as the draft forward, // second graph key in the multi-key CUDA-graph cache. if (domino_correct_greedy_chain_fused( - dw_, draft_backend_, target->lm_head_tensor(), + dw, draft_backend_, target->lm_head_tensor(), target->gpu_embd_table(), local_hidden.data(), q_len, last_tok, draft_tok)) { used_domino = true; @@ -690,7 +743,7 @@ bool LagunaBackend::do_spec_decode(int committed, int n_gen, } if (used_domino) { // fused path done - } else if (domino_correct_greedy_chain(dw_, draft_backend_, *target, + } else if (domino_correct_greedy_chain(dw, draft_backend_, *target, local_hidden.data(), q_len, last_tok, draft_tok)) { used_domino = true; @@ -703,7 +756,49 @@ bool LagunaBackend::do_spec_decode(int committed, int n_gen, } } } - if (!used_domino) { + bool used_dspark = false; + if (!used_domino && laguna_dspark_enabled() && dw.dspark.enabled && + q_len > 1 && !sampled_verify && !args_.ddtree_mode) { + static std::atomic s_dspark_logged{false}; + if (!s_dspark_logged.exchange(true)) { + std::fprintf(stderr, + "[laguna-spec] DSpark Markov head active for greedy chain decode " + "(rank=%d vocab=%d confidence_dim=%d)\n", + dw.dspark.markov_rank, dw.dspark.vocab_size, dw.dspark.confidence_dim); + } + static const bool fused_dspark = []() { + const char * e = std::getenv("DFLASH_LAGUNA_FUSED_DSPARK"); + return !(e && e[0] == '0' && e[1] == '\0'); + }(); + bool ds_ok = false; + if (fused_dspark && laguna_dspark_confidence_threshold() <= 0.0f) { + // One graph on the draft stream: lm_head + markov chain + + // in-graph argmax; no host logits round-trip. + ds_ok = dspark_markov_correct_greedy_chain_fused( + dw, draft_backend_, target->lm_head_tensor(), + local_hidden.data(), q_len, last_tok, draft_tok); + } + if (!ds_ok) { + ds_ok = dspark_markov_correct_greedy_chain(dw, draft_backend_, *target, + local_hidden.data(), q_len, + last_tok, + laguna_dspark_confidence_threshold(), + draft_tok); + } + if (ds_ok) { + used_dspark = true; + q_len = (int)draft_tok.size(); + target_tok.resize((size_t)q_len); + } else { + static std::atomic s_dspark_warned{false}; + if (!s_dspark_warned.exchange(true)) { + std::fprintf(stderr, + "[laguna-spec] DSpark Markov head failed; falling back to base " + "DFlash projection\n"); + } + } + } + if (!used_domino && !used_dspark) { if (!target->project_hidden_to_tokens(local_hidden.data(), q_len, draft_tok)) { std::fprintf(stderr, "[laguna-spec] projection failed\n"); step_graph_destroy(draft_sg); @@ -712,6 +807,7 @@ bool LagunaBackend::do_spec_decode(int committed, int n_gen, draft_tok[0] = last_tok; } + if (step_prof) prof_heads_ms += prof_lap(); const bool tree_special_inactive = !(budget_hook && !budget_hook->close_token_ids.empty()); // kvflash: the tree graph is position-indexed, so only take it while @@ -727,7 +823,25 @@ bool LagunaBackend::do_spec_decode(int committed, int n_gen, const int K = (args_.ddtree_budget > L) ? 8 : 1; std::vector top_lp; std::vector top_ids; - if (!target->project_hidden_to_topk(local_hidden.data(), q_len, K, + static const bool dspark_tree = []() { + const char * e = std::getenv("DFLASH_LAGUNA_DSPARK_TREE"); + return !(e && e[0] == '0' && e[1] == '\0'); + }(); + bool topk_ok = false; + if (dspark_tree && laguna_dspark_enabled() && dw.dspark.enabled) { + static std::atomic s_dstree_logged{false}; + if (!s_dstree_logged.exchange(true)) { + std::fprintf(stderr, + "[laguna-spec] DSpark Markov head active for DDTree candidates\n"); + } + topk_ok = dspark_markov_project_topk(dw, draft_backend_, + target->lm_head_tensor(), + local_hidden.data(), q_len, K, + args_.ddtree_temp, last_tok, + top_lp, top_ids); + } + if (!topk_ok && + !target->project_hidden_to_topk(local_hidden.data(), q_len, K, args_.ddtree_temp, top_lp, top_ids)) { std::fprintf(stderr, "[laguna-spec] ddtree topk projection failed\n"); step_graph_destroy(draft_sg); @@ -780,6 +894,7 @@ bool LagunaBackend::do_spec_decode(int committed, int n_gen, n_accept_sum += std::max(0, emitted - 1); n_draft_steps++; + n_draft_pos_sum += q_len; if (io.cancelled || hit_eos || emitted <= 0 || next_token < 0 || (!ignore_eos && target->is_eos(next_token))) { committed += emitted; @@ -812,12 +927,14 @@ bool LagunaBackend::do_spec_decode(int committed, int n_gen, } int verify_last_tok = -1; + if (step_prof) prof_lap(); if (!target->verify_batch(draft_tok, committed, verify_last_tok, &target_tok)) { std::fprintf(stderr, "[laguna-spec] verify failed\n"); step_graph_destroy(draft_sg); return false; } + if (step_prof) prof_verify_ms += prof_lap(); int accept_n = 1; int bonus_tok = -1; int verify_vocab = 0; @@ -943,6 +1060,7 @@ bool LagunaBackend::do_spec_decode(int committed, int n_gen, n_generated += emitted; n_accept_sum += std::min(accept_n, emitted); n_draft_steps++; + n_draft_pos_sum += q_len; if (io.cancelled) break; if (hit_eos) break; } @@ -951,8 +1069,17 @@ bool LagunaBackend::do_spec_decode(int committed, int n_gen, auto t_dec1 = std::chrono::steady_clock::now(); const double decode_s = std::chrono::duration(t_dec1 - t_dec0).count(); - const int total_draft_pos = std::max(1, n_draft_steps * q_len); + const int total_draft_pos = std::max(1, n_draft_pos_sum); const double accept_pct = 100.0 * (double)n_accept_sum / (double)total_draft_pos; + if (step_prof && n_draft_steps > 0) { + std::fprintf(stderr, + "[step-prof] per-step ms: draft=%.2f heads=%.2f verify=%.2f " + "other=%.2f total=%.2f (steps=%d)\n", + prof_draft_ms / n_draft_steps, prof_heads_ms / n_draft_steps, + prof_verify_ms / n_draft_steps, + (decode_s * 1000.0 - prof_draft_ms - prof_heads_ms - prof_verify_ms) / n_draft_steps, + decode_s * 1000.0 / n_draft_steps, n_draft_steps); + } std::fprintf(stderr, "[laguna-spec] tokens=%d time=%.3f s speed=%.2f tok/s " "steps=%d accepted=%d/%d (%.1f%%) avg_commit=%.2f\n", n_generated, decode_s, @@ -1095,11 +1222,11 @@ 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); - const bool can_spec = req.n_gen > 0 && !req.force_ar_decode && !args_.draft_path.empty() && dflash_target_ + && select_decode_draft(std::string()) && !draft_parked_ && feature_mirror_.target_feat && cache_.target_feat @@ -1319,11 +1446,11 @@ 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); - const bool can_spec = req.n_gen > 0 && !req.force_ar_decode && !args_.draft_path.empty() && dflash_target_ + && select_decode_draft(std::string()) && !draft_parked_ && feature_mirror_.target_feat && cache_.target_feat @@ -2849,7 +2976,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; } @@ -2865,57 +2992,85 @@ 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_backend_ = nullptr; - return false; - } + draft_variants_.clear(); + draft_variants_.push_back(LagunaDraftVariant{}); + draft_variants_.back().name = "base"; - 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_) { + if (!load_draft_gguf(args_.draft_path, draft_backend_, variant.weights, + nullptr)) { + 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 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 fc_in=%d " + "target_hidden=%d draft_hidden=%d n_capture_layers=%d swa=%d\n", + variant.name.c_str(), 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; @@ -2930,6 +3085,12 @@ bool LagunaBackend::load_decode_draft() { return false; } + default_draft_variant_ = "base"; + if (!select_decode_draft(default_draft_variant_)) { + 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_); @@ -2946,14 +3107,39 @@ 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_variant_; + } + for (LagunaDraftVariant & variant : draft_variants_) { + if (variant.name == wanted) { + if (active_dw_ != &variant.weights) { + std::fprintf(stderr, "[laguna] selected draft variant: %s\n", + variant.name.c_str()); + } + active_dw_ = &variant.weights; + return true; + } + } + std::fprintf(stderr, "[laguna] unknown draft 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; + default_draft_variant_ = "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..4dddfd950 100644 --- a/server/src/laguna/laguna_backend.h +++ b/server/src/laguna/laguna_backend.h @@ -48,6 +48,11 @@ struct LagunaBackendArgs { ggml_type kv_type = GGML_TYPE_Q8_0; }; +struct LagunaDraftVariant { + std::string name; + DraftWeights weights; +}; + class LagunaBackend : public ModelBackend { public: explicit LagunaBackend(const LagunaBackendArgs & args); @@ -100,7 +105,9 @@ 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 default_draft_variant_ = "base"; DraftFeatureMirror feature_mirror_{}; LagunaDFlashTarget * dflash_target_ = nullptr; bool draft_parked_ = false; @@ -178,6 +185,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/laguna/laguna_target_graph.cpp b/server/src/laguna/laguna_target_graph.cpp index b4a23761f..0c7687e4c 100644 --- a/server/src/laguna/laguna_target_graph.cpp +++ b/server/src/laguna/laguna_target_graph.cpp @@ -1203,6 +1203,31 @@ LagunaGraphOutputs build_laguna_graph( // ---- Public turnkey forward step ---------------------------------------- // +// Pinned host staging for per-step graph inputs. ggml_backend_tensor_set from +// pageable memory costs a staged DMA plus a stream synchronize PER CALL, and +// the decode/verify hot loops issue ~5 uploads per step - those syncs, not +// kernels, dominated step wall time. Staging in pinned memory and uploading +// with tensor_set_async on the backend stream leaves the single sync inside +// ggml_backend_graph_compute() as the only per-step synchronization. +static uint8_t * laguna_host_stage(ggml_backend_t backend, + ggml_backend_buffer_t & buf, + size_t need) { + if (buf && ggml_backend_buffer_get_size(buf) >= need) { + return (uint8_t *) ggml_backend_buffer_get_base(buf); + } + if (buf) { + ggml_backend_buffer_free(buf); + buf = nullptr; + } + ggml_backend_buffer_type_t buft = nullptr; + if (ggml_backend_dev_t dev = ggml_backend_get_device(backend)) { + buft = ggml_backend_dev_host_buffer_type(dev); + } + if (!buft) return nullptr; + buf = ggml_backend_buft_alloc_buffer(buft, need); + return buf ? (uint8_t *) ggml_backend_buffer_get_base(buf) : nullptr; +} + // Wires the FULL + SWA causal masks, runs the backend graph, and returns // last-token logits and/or a GPU-computed argmax on the host. Updates // cache.cur_pos. The common greedy decode path reuses a per-thread graph while @@ -1260,9 +1285,11 @@ bool laguna_step( ggml_tensor * argmax = nullptr; std::vector full_mask; std::vector swa_mask; + ggml_backend_buffer_t stage_buf = nullptr; void clear() { if (alloc) { ggml_gallocr_free(alloc); alloc = nullptr; } + if (stage_buf) { ggml_backend_buffer_free(stage_buf); stage_buf = nullptr; } if (ctx) { ggml_free(ctx); ctx = nullptr; } gf = nullptr; inp_embed = positions = kv_idx = mask_full = mask_swa = argmax = nullptr; @@ -1340,24 +1367,36 @@ bool laguna_step( cached.swa_mask.resize((size_t)mk_w); } - ggml_backend_tensor_set(cached.inp_embed, embed, 0, ggml_nbytes(cached.inp_embed)); - int32_t pos_val = kv_start; - ggml_backend_tensor_set(cached.positions, &pos_val, 0, sizeof(pos_val)); - ggml_backend_tensor_set(cached.kv_idx, &pos_val, 0, sizeof(pos_val)); - - std::fill(cached.full_mask.begin(), cached.full_mask.end(), -INFINITY); + const size_t embed_sz = ggml_nbytes(cached.inp_embed); + const size_t mask_sz = (size_t)mk_w * sizeof(float); + uint8_t * st = laguna_host_stage(backend, cached.stage_buf, + embed_sz + sizeof(int32_t) + 2 * mask_sz); + if (!st) { + std::fprintf(stderr, "laguna_step: host stage alloc failed\n"); + return false; + } + float * st_embed = (float *) st; + int32_t * st_pos = (int32_t *) (st + embed_sz); + float * st_mfull = (float *) (st + embed_sz + sizeof(int32_t)); + float * st_mswa = (float *) (st + embed_sz + sizeof(int32_t) + mask_sz); + + std::memcpy(st_embed, embed, embed_sz); + *st_pos = kv_start; + std::fill(st_mfull, st_mfull + mk_w, -INFINITY); for (int k = 0; k <= kv_start && k < kv_len && k < mk_w; ++k) { - cached.full_mask[(size_t)k] = 0.0f; + st_mfull[(size_t)k] = 0.0f; } - ggml_backend_tensor_set(cached.mask_full, cached.full_mask.data(), 0, ggml_nbytes(cached.mask_full)); - - std::fill(cached.swa_mask.begin(), cached.swa_mask.end(), -INFINITY); const int W = w.sliding_window; const int win_lo = std::max(0, kv_start - W + 1); + std::fill(st_mswa, st_mswa + mk_w, -INFINITY); for (int k = win_lo; k <= kv_start && k < kv_len && k < mk_w; ++k) { - cached.swa_mask[(size_t)k] = 0.0f; + st_mswa[(size_t)k] = 0.0f; } - ggml_backend_tensor_set(cached.mask_swa, cached.swa_mask.data(), 0, ggml_nbytes(cached.mask_swa)); + ggml_backend_tensor_set_async(backend, cached.inp_embed, st_embed, 0, embed_sz); + ggml_backend_tensor_set_async(backend, cached.positions, st_pos, 0, sizeof(int32_t)); + ggml_backend_tensor_set_async(backend, cached.kv_idx, st_pos, 0, sizeof(int32_t)); + ggml_backend_tensor_set_async(backend, cached.mask_full, st_mfull, 0, mask_sz); + ggml_backend_tensor_set_async(backend, cached.mask_swa, st_mswa, 0, mask_sz); if (ggml_backend_graph_compute(backend, cached.gf) != GGML_STATUS_SUCCESS) { std::fprintf(stderr, "laguna_step: cached graph_compute failed\n"); @@ -1532,23 +1571,6 @@ bool laguna_verify_batch( (void)token_ids; if (n_tokens <= 0) return false; - const size_t arena_size = ggml_tensor_overhead() * 16384 + ggml_graph_overhead() + 16 * 1024 * 1024; - static thread_local std::vector g_arena_block; - static thread_local std::vector g_arena_bonus; - std::vector & g_arena = (n_tokens == 1) ? g_arena_bonus : g_arena_block; - if (g_arena.size() < arena_size) g_arena.resize(arena_size); - ggml_init_params ip{}; - ip.mem_size = arena_size; - ip.mem_buffer = g_arena.data(); - ip.no_alloc = true; - ggml_context * ctx = ggml_init(ip); - ggml_cgraph * gf = ggml_new_graph_custom(ctx, 16384, false); - - ggml_tensor * ie = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, w.n_embd, n_tokens, 1); - ggml_set_input(ie); - ggml_tensor * pp = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, n_tokens); - ggml_set_input(pp); - const int kv_len = kv_start + n_tokens; static const bool g_no_kvpad = (std::getenv("DFLASH_LAGUNA_NO_KVPAD") != nullptr); static const bool g_pad_cpy = (std::getenv("DFLASH_LAGUNA_PAD_CPY") != nullptr); @@ -1560,114 +1582,207 @@ bool laguna_verify_batch( ? std::min((kv_len + 255) & ~255, kv_cap) : 0; const int mk_w = kv_pad > 0 ? kv_pad : kv_len; - ggml_tensor * kvi = nullptr; - if (kv_pad > 0 && !g_pad_cpy) { - kvi = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, n_tokens); - ggml_set_input(kvi); - } + // Persistent step graph: within a (n_tokens, mk_w) window the built graph + // is structurally identical every step - the kv position enters only via + // input DATA (positions / kv_idx rows / mask contents / feat rows), which + // is also what makes the ggml-cuda graph cache replay it. Reuse the built + // graph and skip the per-step host rebuild + allocator pass (~0.5ms/step). + // DFLASH_LAGUNA_PERSIST_VERIFY=0 restores the rebuild-every-step path. + static const bool g_persist = []() { + const char * e = std::getenv("DFLASH_LAGUNA_PERSIST_VERIFY"); + return !(e && e[0] == '0' && e[1] == '\0'); + }(); + struct VerifySlot { + std::vector arena; + ggml_context * ctx = nullptr; + ggml_cgraph * gf = nullptr; + ggml_gallocr_t galloc = nullptr; + ggml_tensor *ie = nullptr, *pp = nullptr, *kvi = nullptr; + ggml_tensor *mk_full = nullptr, *mk_swa = nullptr, *feat_rows = nullptr; + ggml_tensor *argmax = nullptr, *logits = nullptr; + int n_tokens = 0, mk_w = 0, kv_pad = 0, kv_cap = 0; + bool feat = false, want_logits = false; + // Identity the graph was built against - never reuse across a + // different backend, weight set, or cache instance. + ggml_backend_t backend_id = nullptr; + const LagunaTargetWeights * w_id = nullptr; + const LagunaTargetCache * cache_id = nullptr; + ggml_backend_buffer_t stage_buf = nullptr; + }; + static thread_local VerifySlot g_slot_block, g_slot_bonus; + VerifySlot & S = (n_tokens == 1) ? g_slot_bonus : g_slot_block; + const bool want_logits = out_logits != nullptr; + const bool want_feat = cache.target_feat && cache.target_feat_cap > 0; + // Reuse requires the kv_idx input path (kv_pad > 0, no PAD_CPY): those + // graphs take the kv position purely as input data. The NO_KVPAD and + // PAD_CPY fallbacks bake kv_start into the graph structure, so a reused + // graph would read/write KV at stale offsets. + const bool reuse = g_persist && (kv_pad > 0 && !g_pad_cpy) && S.ctx != nullptr && + S.backend_id == backend && S.w_id == &w && S.cache_id == &cache && + S.n_tokens == n_tokens && S.mk_w == mk_w && S.kv_pad == kv_pad && + S.kv_cap == kv_cap && S.feat == want_feat && S.want_logits == want_logits; + + if (!reuse) { + const size_t arena_size = ggml_tensor_overhead() * 16384 + ggml_graph_overhead() + 16 * 1024 * 1024; + if (S.arena.size() < arena_size) S.arena.resize(arena_size); + if (S.ctx) { ggml_free(S.ctx); S.ctx = nullptr; } + ggml_init_params ip{}; + ip.mem_size = arena_size; + ip.mem_buffer = S.arena.data(); + ip.no_alloc = true; + ggml_context * ctx = ggml_init(ip); + ggml_cgraph * gf = ggml_new_graph_custom(ctx, 16384, false); + + ggml_tensor * ie = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, w.n_embd, n_tokens, 1); + ggml_set_input(ie); + ggml_tensor * pp = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, n_tokens); + ggml_set_input(pp); + + ggml_tensor * kvi = nullptr; + if (kv_pad > 0 && !g_pad_cpy) { + kvi = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, n_tokens); + ggml_set_input(kvi); + } - ggml_tensor * mk_full = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, mk_w, n_tokens, 1, 1); - ggml_set_input(mk_full); - ggml_tensor * mk_full_cnv = ggml_cast(ctx, mk_full, GGML_TYPE_F16); - ggml_tensor * mk_swa = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, mk_w, n_tokens, 1, 1); - ggml_set_input(mk_swa); - ggml_tensor * mk_swa_cnv = ggml_cast(ctx, mk_swa, GGML_TYPE_F16); + ggml_tensor * mk_full = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, mk_w, n_tokens, 1, 1); + ggml_set_input(mk_full); + ggml_tensor * mk_full_cnv = ggml_cast(ctx, mk_full, GGML_TYPE_F16); + ggml_tensor * mk_swa = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, mk_w, n_tokens, 1, 1); + ggml_set_input(mk_swa); + ggml_tensor * mk_swa_cnv = ggml_cast(ctx, mk_swa, GGML_TYPE_F16); - ggml_tensor * feat_rows = nullptr; - if (cache.target_feat && cache.target_feat_cap > 0) { - feat_rows = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, n_tokens); - ggml_set_input(feat_rows); - } + ggml_tensor * feat_rows = nullptr; + if (want_feat) { + feat_rows = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, n_tokens); + ggml_set_input(feat_rows); + } - LagunaGraphInputs gi{}; - gi.inp_embed = ie; - gi.positions = pp; - gi.attn_mask = mk_full_cnv; - gi.attn_mask_swa = mk_swa_cnv; - gi.n_tokens = n_tokens; - gi.kv_start = kv_start; - gi.kv_pad = kv_pad; - gi.kv_idx = kvi; - gi.output_last_only = false; - gi.output_logits = true; - gi.logits_are_output = out_logits != nullptr; - gi.capture_features = true; - gi.target_feat_rows = feat_rows; - gi.hybrid = nullptr; + LagunaGraphInputs gi{}; + gi.inp_embed = ie; + gi.positions = pp; + gi.attn_mask = mk_full_cnv; + gi.attn_mask_swa = mk_swa_cnv; + gi.n_tokens = n_tokens; + gi.kv_start = kv_start; + gi.kv_pad = kv_pad; + gi.kv_idx = kvi; + gi.output_last_only = false; + gi.output_logits = true; + gi.logits_are_output = out_logits != nullptr; + gi.capture_features = true; + gi.target_feat_rows = feat_rows; + gi.hybrid = nullptr; + + LagunaGraphOutputs go = build_laguna_graph(ctx, gf, w, cache, gi); + ggml_tensor * argmax = ggml_argmax(ctx, go.logits); + ggml_set_output(argmax); + ggml_build_forward_expand(gf, argmax); - LagunaGraphOutputs go = build_laguna_graph(ctx, gf, w, cache, gi); - ggml_tensor * argmax = ggml_argmax(ctx, go.logits); - ggml_set_output(argmax); - ggml_build_forward_expand(gf, argmax); - - static ggml_gallocr_t galloc_verify_block = nullptr; - static ggml_gallocr_t galloc_verify_bonus = nullptr; - ggml_gallocr_t & galloc_verify = (n_tokens == 1) ? galloc_verify_bonus : galloc_verify_block; - if (!galloc_verify) galloc_verify = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); - if (!ggml_gallocr_alloc_graph(galloc_verify, gf)) { - std::fprintf(stderr, "laguna_verify_batch: gallocr_alloc_graph failed\n"); - ggml_free(ctx); + if (S.galloc && S.backend_id != backend) { + ggml_gallocr_free(S.galloc); + S.galloc = nullptr; + if (S.stage_buf) { ggml_backend_buffer_free(S.stage_buf); S.stage_buf = nullptr; } + } + if (!S.galloc) S.galloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); + if (!ggml_gallocr_alloc_graph(S.galloc, gf)) { + std::fprintf(stderr, "laguna_verify_batch: gallocr_alloc_graph failed\n"); + ggml_free(ctx); + return false; + } + S.ctx = ctx; S.gf = gf; + S.ie = ie; S.pp = pp; S.kvi = kvi; + S.mk_full = mk_full; S.mk_swa = mk_swa; S.feat_rows = feat_rows; + S.argmax = argmax; S.logits = go.logits; + S.n_tokens = n_tokens; S.mk_w = mk_w; S.kv_pad = kv_pad; S.kv_cap = kv_cap; + S.feat = want_feat; S.want_logits = want_logits; + S.backend_id = backend; S.w_id = &w; S.cache_id = &cache; + } + + ggml_cgraph * gf = S.gf; + ggml_tensor * ie = S.ie; + ggml_tensor * pp = S.pp; + ggml_tensor * kvi = S.kvi; + ggml_tensor * mk_full = S.mk_full; + ggml_tensor * mk_swa = S.mk_swa; + ggml_tensor * feat_rows = S.feat_rows; + ggml_tensor * argmax = S.argmax; + + const size_t embed_sz = ggml_nbytes(ie); + const size_t pos_sz = (size_t)n_tokens * sizeof(int32_t); + const size_t mask_sz = (size_t)mk_w * n_tokens * sizeof(float); + uint8_t * st = laguna_host_stage(backend, S.stage_buf, + embed_sz + 3 * pos_sz + 2 * mask_sz); + if (!st) { + std::fprintf(stderr, "laguna_verify_batch: host stage alloc failed\n"); return false; } - - ggml_backend_tensor_set(ie, embed, 0, ggml_nbytes(ie)); - std::vector pos((size_t)n_tokens); - for (int i = 0; i < n_tokens; ++i) pos[(size_t)i] = kv_start + i; - ggml_backend_tensor_set(pp, pos.data(), 0, ggml_nbytes(pp)); + float * st_embed = (float *) st; + int32_t * st_pos = (int32_t *) (st + embed_sz); + int32_t * st_kvi = (int32_t *) (st + embed_sz + pos_sz); + int32_t * st_feat = (int32_t *) (st + embed_sz + 2 * pos_sz); + float * st_mfull = (float *) (st + embed_sz + 3 * pos_sz); + float * st_mswa = (float *) (st + embed_sz + 3 * pos_sz + mask_sz); + + std::memcpy(st_embed, embed, embed_sz); + ggml_backend_tensor_set_async(backend, ie, st_embed, 0, embed_sz); + for (int i = 0; i < n_tokens; ++i) st_pos[(size_t)i] = kv_start + i; + ggml_backend_tensor_set_async(backend, pp, st_pos, 0, pos_sz); if (feat_rows) { - std::vector feat_idx((size_t)n_tokens); for (int i = 0; i < n_tokens; ++i) { - feat_idx[(size_t)i] = (kv_start + i) % cache.target_feat_cap; + st_feat[(size_t)i] = (kv_start + i) % cache.target_feat_cap; } - ggml_backend_tensor_set(feat_rows, feat_idx.data(), 0, ggml_nbytes(feat_rows)); + ggml_backend_tensor_set_async(backend, feat_rows, st_feat, 0, pos_sz); } if (kvflash) { if (!kvi) { std::fprintf(stderr, "laguna_verify_batch: kvflash requires the kv_pad " "set_rows path (NO_KVPAD / PAD_CPY are incompatible)\n"); - ggml_free(ctx); + ggml_free(S.ctx); S.ctx = nullptr; return false; } std::vector rows; std::vector mfull, mswa; if (!kvflash_fill_rows_and_masks(*kvflash, kv_start, n_tokens, mk_w, w.sliding_window, rows, &mfull, &mswa)) { - ggml_free(ctx); + ggml_free(S.ctx); S.ctx = nullptr; return false; } - ggml_backend_tensor_set(kvi, rows.data(), 0, ggml_nbytes(kvi)); - ggml_backend_tensor_set(mk_full, mfull.data(), 0, ggml_nbytes(mk_full)); - ggml_backend_tensor_set(mk_swa, mswa.data(), 0, ggml_nbytes(mk_swa)); + std::memcpy(st_kvi, rows.data(), pos_sz); + std::memcpy(st_mfull, mfull.data(), mask_sz); + std::memcpy(st_mswa, mswa.data(), mask_sz); + ggml_backend_tensor_set_async(backend, kvi, st_kvi, 0, pos_sz); + ggml_backend_tensor_set_async(backend, mk_full, st_mfull, 0, mask_sz); + ggml_backend_tensor_set_async(backend, mk_swa, st_mswa, 0, mask_sz); } else { if (kvi) { - ggml_backend_tensor_set(kvi, pos.data(), 0, ggml_nbytes(kvi)); + ggml_backend_tensor_set_async(backend, kvi, st_pos, 0, pos_sz); } - std::vector mfull((size_t)mk_w * n_tokens, -INFINITY); + std::fill(st_mfull, st_mfull + (size_t)mk_w * n_tokens, -INFINITY); for (int q = 0; q < n_tokens; ++q) { const int abs_q = kv_start + q; for (int k = 0; k <= abs_q && k < kv_len; ++k) { - mfull[(size_t)q * mk_w + k] = 0.0f; + st_mfull[(size_t)q * mk_w + k] = 0.0f; } } - ggml_backend_tensor_set(mk_full, mfull.data(), 0, ggml_nbytes(mk_full)); + ggml_backend_tensor_set_async(backend, mk_full, st_mfull, 0, mask_sz); - std::vector mswa((size_t)mk_w * n_tokens, -INFINITY); + std::fill(st_mswa, st_mswa + (size_t)mk_w * n_tokens, -INFINITY); const int W = w.sliding_window; for (int q = 0; q < n_tokens; ++q) { const int abs_q = kv_start + q; const int win_lo = std::max(0, abs_q - W + 1); for (int k = win_lo; k <= abs_q && k < kv_len; ++k) { - mswa[(size_t)q * mk_w + k] = 0.0f; + st_mswa[(size_t)q * mk_w + k] = 0.0f; } } - ggml_backend_tensor_set(mk_swa, mswa.data(), 0, ggml_nbytes(mk_swa)); + ggml_backend_tensor_set_async(backend, mk_swa, st_mswa, 0, mask_sz); } if (ggml_backend_graph_compute(backend, gf) != GGML_STATUS_SUCCESS) { std::fprintf(stderr, "laguna_verify_batch: graph_compute failed\n"); - ggml_free(ctx); + ggml_free(S.ctx); S.ctx = nullptr; return false; } @@ -1677,13 +1792,13 @@ bool laguna_verify_batch( if (out_logits) { const int vocab = (int)w.embedder.n_vocab; out_logits->resize((size_t)vocab * (size_t)n_tokens); - ggml_backend_tensor_get(go.logits, out_logits->data(), 0, + ggml_backend_tensor_get(S.logits, out_logits->data(), 0, sizeof(float) * out_logits->size()); } cache.cur_pos = kv_len; cache.last_tok = out_argmax.empty() ? -1 : out_argmax.back(); - ggml_free(ctx); + if (!g_persist) { ggml_free(S.ctx); S.ctx = nullptr; } return true; } diff --git a/server/src/server/server_main.cpp b/server/src/server/server_main.cpp index 99abfb99d..317fdbed7 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,7 @@ 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" " --port Listen port (default: 8080)\n" " --host Bind address (default: 0.0.0.0)\n" " --max-ctx Max context length (default: 131072)\n"