diff --git a/server/deps/llama.cpp b/server/deps/llama.cpp index 0aaf1dee4..fa6ea8268 160000 --- a/server/deps/llama.cpp +++ b/server/deps/llama.cpp @@ -1 +1 @@ -Subproject commit 0aaf1dee404c93d69abfee1ce1f03981ea5bda1e +Subproject commit fa6ea8268b9c2ccbc9baa65f85c6d35e307dbd0c diff --git a/server/src/laguna/laguna_backend.cpp b/server/src/laguna/laguna_backend.cpp index fcf18bda5..920df85f9 100644 --- a/server/src/laguna/laguna_backend.cpp +++ b/server/src/laguna/laguna_backend.cpp @@ -449,6 +449,17 @@ bool LagunaBackend::do_spec_decode(int committed, int n_gen, if (dflash_target_) { dflash_target_->set_keep_verify_logits(sampled_verify); } + static const bool g_device_draft_proj_disabled = + (std::getenv("DFLASH_LAGUNA_DEVICE_DRAFT_PROJ_DISABLE") != nullptr); + const bool can_project_draft_on_device = + !g_device_draft_proj_disabled && dflash_target_ && + draft_backend_ == backend_; + static const bool g_chain_device_draft_proj = + (std::getenv("DFLASH_LAGUNA_CHAIN_DEVICE_DRAFT_PROJ") != nullptr); + static const bool g_fused_draft_lm_head_disabled = + (std::getenv("DFLASH_LAGUNA_FUSED_DRAFT_LM_HEAD_DISABLE") != nullptr); + static const bool g_chain_fused_draft_lm_head = + (std::getenv("DFLASH_LAGUNA_CHAIN_FUSED_DRAFT_LM_HEAD") != nullptr); StepGraph draft_sg; @@ -610,7 +621,25 @@ 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_, + 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 + // the pager is identity and the step fits the resident pool; otherwise + // the slot-mapped chain verify below handles it. + const bool kvflash_tree_ok = + !kvflash_active() || + (kvflash_pager_.is_identity() && + committed + args_.ddtree_budget + 1 <= kvflash_tokens_); + const bool try_ddtree = + args_.ddtree_mode && target->supports_tree_verify() && kvflash_tree_ok && + q_len > 1 && tree_special_inactive && !sampled_verify; + const bool use_fused_draft_lm_head = + can_project_draft_on_device && !g_fused_draft_lm_head_disabled && + (try_ddtree || g_chain_fused_draft_lm_head); + ggml_tensor * draft_lm_head = use_fused_draft_lm_head + ? dflash_target_->output_weight() : nullptr; + + if (!build_draft_step(draft_sg, dw_, draft_lm_head, 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)))) { @@ -642,12 +671,23 @@ bool LagunaBackend::do_spec_decode(int committed, int n_gen, return false; } - 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()); + bool local_hidden_valid = false; + auto ensure_local_hidden = [&]() -> bool { + const size_t want = (size_t)hidden * (size_t)q_len; + if (local_hidden_valid && local_hidden.size() == want) return true; + local_hidden.resize(want); + ggml_backend_tensor_get(draft_sg.hidden_states, local_hidden.data(), 0, + sizeof(float) * local_hidden.size()); + local_hidden_valid = true; + return true; + }; bool used_domino = false; - if (dw_.domino.enabled && q_len > 1 && !sampled_verify && !args_.ddtree_mode) { + if (!try_ddtree && dw_.domino.enabled && q_len > 1 && !sampled_verify && !args_.ddtree_mode) { + if (!ensure_local_hidden()) { + step_graph_destroy(draft_sg); + return false; + } static std::atomic s_domino_logged{false}; if (!s_domino_logged.exchange(true)) { std::fprintf(stderr, @@ -668,8 +708,29 @@ bool LagunaBackend::do_spec_decode(int committed, int n_gen, } } } - if (!used_domino) { - if (!target->project_hidden_to_tokens(local_hidden.data(), q_len, draft_tok)) { + + if (!try_ddtree && !used_domino) { + bool projected = false; + if (g_chain_fused_draft_lm_head && draft_sg.argmax_tokens) { + draft_tok.resize((size_t)q_len); + ggml_backend_tensor_get(draft_sg.argmax_tokens, draft_tok.data(), 0, + sizeof(int32_t) * (size_t)q_len); + projected = true; + } + if (g_chain_device_draft_proj && can_project_draft_on_device && + dflash_target_->project_device_hidden_to_tokens( + draft_sg.hidden_states, q_len, draft_tok)) { + projected = true; + } + if (!projected) { + if (!ensure_local_hidden() || + !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); + return false; + } + } + if ((int)draft_tok.size() < q_len) { std::fprintf(stderr, "[laguna-spec] projection failed\n"); step_graph_destroy(draft_sg); return false; @@ -677,23 +738,26 @@ bool LagunaBackend::do_spec_decode(int committed, int n_gen, draft_tok[0] = last_tok; } - 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 - // the pager is identity and the step fits the resident pool; otherwise - // the slot-mapped chain verify below handles it. - const bool kvflash_tree_ok = - !kvflash_active() || - (kvflash_pager_.is_identity() && - committed + args_.ddtree_budget + 1 <= kvflash_tokens_); - if (args_.ddtree_mode && target->supports_tree_verify() && kvflash_tree_ok && - q_len > 1 && tree_special_inactive && !sampled_verify) { + if (try_ddtree) { const int L = q_len - 1; 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, - args_.ddtree_temp, top_lp, top_ids)) { + bool topk_ok = false; + if (draft_sg.logits) { + topk_ok = dflash_target_->project_device_logits_to_topk( + draft_sg.logits, q_len, K, args_.ddtree_temp, + top_lp, top_ids); + } + if (!topk_ok && can_project_draft_on_device) { + topk_ok = dflash_target_->project_device_hidden_to_topk( + draft_sg.hidden_states, q_len, K, args_.ddtree_temp, + top_lp, top_ids); + } + if (!topk_ok && + (!ensure_local_hidden() || + !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); return false; @@ -703,6 +767,83 @@ bool LagunaBackend::do_spec_decode(int committed, int n_gen, top_ids.data() + (size_t)K, L, K, args_.ddtree_budget, /*chain_seed=*/true); + + static const bool g_lazy_ddtree = + (std::getenv("DFLASH_LAGUNA_DDTREE_LAZY") != nullptr); + if (g_lazy_ddtree) { + std::vector accepted; + accepted.reserve((size_t)q_len + 1); + accepted.push_back(0); + + int current_index = 0; + int next_token = -1; + int cursor_pos = committed; + int32_t verify_tok = last_tok; + while (true) { + std::vector one_tok = { verify_tok }; + int pred = -1; + if (!target->verify_batch(one_tok, cursor_pos, pred, nullptr)) { + std::fprintf(stderr, "[laguna-spec] lazy verify_tree step failed\n"); + step_graph_destroy(draft_sg); + return false; + } + next_token = pred; + cursor_pos++; + + if ((int)accepted.size() >= need_commit_budget) break; + const auto & children = tree.child_maps[(size_t)current_index]; + auto it = children.find(next_token); + if (it == children.end()) break; + current_index = it->second; + accepted.push_back(current_index); + verify_tok = next_token; + } + if (next_token < 0) { + std::fprintf(stderr, "[laguna-spec] lazy verify_tree produced no posterior\n"); + step_graph_destroy(draft_sg); + return false; + } + + int commit_n = (int)accepted.size(); + if (commit_n > need_commit_budget) commit_n = need_commit_budget; + if (commit_n <= 0) { + step_graph_destroy(draft_sg); + break; + } + + bool hit_eos = false; + int emitted = 0; + for (int i = 0; i < commit_n; ++i) { + const int dfs = accepted[(size_t)i]; + const int32_t tok = (dfs == 0) ? last_tok : tree.token_ids[(size_t)dfs - 1]; + if (!ignore_eos && target->is_eos(tok)) { hit_eos = true; break; } + out_tokens.push_back(tok); + sample_history.push_back(tok); + io.emit(tok); + emitted++; + if (io.cancelled) break; + } + + n_accept_sum += std::max(0, emitted - 1); + n_draft_steps++; + + if (feature_mirror_.target_feat && cache_.target_feat && emitted > 0) { + draft_feature_mirror_sync_range(cache_.target_feat, cache_.target_feat_cap, + feature_mirror_, committed, emitted); + } + + committed += emitted; + cache_.cur_pos = committed; + n_generated += emitted; + last_tok = next_token; + cache_.last_tok = last_tok; + if (io.cancelled || hit_eos || emitted <= 0 || + (!ignore_eos && target->is_eos(next_token))) { + break; + } + continue; + } + const int N = args_.ddtree_budget + 1; std::vector flat_tokens((size_t)N, 0); flat_tokens[0] = last_tok; diff --git a/server/src/laguna/laguna_dflash_target.cpp b/server/src/laguna/laguna_dflash_target.cpp index e076622f7..7b6677a11 100644 --- a/server/src/laguna/laguna_dflash_target.cpp +++ b/server/src/laguna/laguna_dflash_target.cpp @@ -1,6 +1,7 @@ // LagunaDFlashTarget - DFlashTarget adapter for Poolside Laguna-XS.2. #include "laguna_dflash_target.h" +#include "../common/geometric_draft_topk_cuda.h" #include "../common/kvflash_pager.h" #include "../common/ddtree.h" @@ -30,6 +31,58 @@ int laguna_argmax_row(const float * row, int vocab) { return best; } +bool laguna_extract_logits_topk(ggml_tensor * logits, + int n_tokens, + int K, + float temperature, + std::vector & top_log_probs, + std::vector & top_token_ids) { + if (!logits || n_tokens <= 0 || K <= 0 || logits->ne[1] < n_tokens) { + return false; + } + + const int vocab = (int)logits->ne[0]; + 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); + +#ifdef DFLASH27B_HAVE_DRAFT_TOPK_CUDA + static const bool kGpuDraftTopk = []() { + const char * v = std::getenv("DFLASH_GPU_DRAFT_TOPK"); + return v == nullptr || v[0] != '0'; + }(); + if (kGpuDraftTopk && + geometric_extract_draft_topk_cuda(logits->data, n_tokens, vocab, K, + top_log_probs.data(), + top_token_ids.data(), + temperature)) { + return true; + } +#endif + + std::vector logits_host((size_t)vocab * (size_t)n_tokens); + ggml_backend_tensor_get(logits, logits_host.data(), 0, + sizeof(float) * logits_host.size()); + extract_draft_topk(logits_host.data(), n_tokens, vocab, K, + top_log_probs.data(), top_token_ids.data(), temperature); + return true; +} + +bool laguna_dflash_tensor_set_checked(const char * label, + ggml_tensor * t, + const void * data, + size_t offset, + size_t size, + bool required = true) { + if (!t || !t->buffer) { + if (required) { + std::fprintf(stderr, "[laguna-dflash] input tensor not allocated: %s\n", label); + } + return !required; + } + ggml_backend_tensor_set(t, data, offset, size); + return true; +} + } // namespace LagunaDFlashTarget::LagunaDFlashTarget( @@ -207,6 +260,247 @@ bool LagunaDFlashTarget::verify_tree( return false; } + static const bool g_tree_cache_disabled = + (std::getenv("DFLASH_LAGUNA_VERIFY_CACHE_DISABLE") != nullptr); + const bool can_reuse_tree_graph = + !g_tree_cache_disabled && logits_out == nullptr && N <= 65; + if (can_reuse_tree_graph) { + struct CachedTreeGraph { + const LagunaDFlashTarget * owner = nullptr; + const LagunaTargetWeights * w_ptr = nullptr; + LagunaTargetCache * cache_ptr = nullptr; + ggml_backend_t backend = nullptr; + ggml_tensor * target_feat = nullptr; + int n_alloc = 0; + int mk_w = 0; + int kv_pad = 0; + bool capture_features = false; + std::vector arena; + ggml_context * ctx = nullptr; + ggml_cgraph * gf = nullptr; + ggml_gallocr_t alloc = nullptr; + ggml_tensor * inp_embed = nullptr; + ggml_tensor * positions = nullptr; + ggml_tensor * kv_idx = nullptr; + ggml_tensor * mask_full = nullptr; + ggml_tensor * mask_swa = nullptr; + ggml_tensor * feat_rows = nullptr; + ggml_tensor * parent_ids = nullptr; + ggml_tensor * argmax = nullptr; + std::vector pos; + std::vector rows; + std::vector feat_idx; + std::vector parents; + std::vector full_mask; + std::vector swa_mask; + + void clear() { + if (alloc) { ggml_gallocr_free(alloc); alloc = nullptr; } + if (ctx) { ggml_free(ctx); ctx = nullptr; } + gf = nullptr; + inp_embed = positions = kv_idx = mask_full = mask_swa = feat_rows = parent_ids = argmax = nullptr; + owner = nullptr; + w_ptr = nullptr; + cache_ptr = nullptr; + backend = nullptr; + target_feat = nullptr; + n_alloc = 0; + mk_w = 0; + kv_pad = 0; + capture_features = false; + } + }; + static thread_local CachedTreeGraph cached_by_n[66]; + + CachedTreeGraph & cached = cached_by_n[N]; + const bool capture_features = cache_.target_feat && cache_.target_feat_cap > 0; + const bool rebuild = + cached.ctx == nullptr || cached.owner != this || cached.w_ptr != &w_ || + cached.cache_ptr != &cache_ || cached.backend != backend_ || + cached.target_feat != cache_.target_feat || cached.n_alloc != N || + cached.mk_w != mk_w || cached.kv_pad != kv_pad || + cached.capture_features != capture_features; + if (rebuild) { + cached.clear(); + const size_t arena_size = + ggml_tensor_overhead() * 32768 + ggml_graph_overhead() + 32 * 1024 * 1024; + cached.arena.resize(arena_size); + + ggml_init_params ip{}; + ip.mem_size = arena_size; + ip.mem_buffer = cached.arena.data(); + ip.no_alloc = true; + cached.ctx = ggml_init(ip); + if (!cached.ctx) return false; + cached.gf = ggml_new_graph_custom(cached.ctx, 32768, false); + + cached.inp_embed = ggml_new_tensor_3d(cached.ctx, GGML_TYPE_F32, w_.n_embd, N, 1); + ggml_set_input(cached.inp_embed); + cached.positions = ggml_new_tensor_1d(cached.ctx, GGML_TYPE_I32, N); + ggml_set_input(cached.positions); + cached.kv_idx = ggml_new_tensor_1d(cached.ctx, GGML_TYPE_I32, N); + ggml_set_input(cached.kv_idx); + cached.parent_ids = ggml_new_tensor_1d(cached.ctx, GGML_TYPE_I32, N); + ggml_set_input(cached.parent_ids); + + cached.mask_full = ggml_new_tensor_4d(cached.ctx, GGML_TYPE_F32, mk_w, N, 1, 1); + ggml_set_input(cached.mask_full); + ggml_tensor * mask_full_cnv = ggml_cast(cached.ctx, cached.mask_full, GGML_TYPE_F16); + cached.mask_swa = ggml_new_tensor_4d(cached.ctx, GGML_TYPE_F32, mk_w, N, 1, 1); + ggml_set_input(cached.mask_swa); + ggml_tensor * mask_swa_cnv = ggml_cast(cached.ctx, cached.mask_swa, GGML_TYPE_F16); + + if (capture_features) { + cached.feat_rows = ggml_new_tensor_1d(cached.ctx, GGML_TYPE_I32, N); + ggml_set_input(cached.feat_rows); + } + + LagunaGraphInputs gi{}; + gi.inp_embed = cached.inp_embed; + gi.positions = cached.positions; + gi.attn_mask = mask_full_cnv; + gi.attn_mask_swa = mask_swa_cnv; + gi.n_tokens = N; + gi.kv_start = 0; + gi.kv_pad = kv_pad; + gi.kv_idx = cached.kv_idx; + gi.tree_parent_ids = cached.parent_ids; + gi.output_last_only = false; + gi.output_logits = true; + gi.logits_are_output = false; + gi.capture_features = capture_features; + gi.target_feat_rows = cached.feat_rows; + gi.hybrid = nullptr; + + LagunaGraphOutputs go = build_laguna_graph(cached.ctx, cached.gf, w_, cache_, gi); + cached.argmax = ggml_argmax(cached.ctx, go.logits); + ggml_set_output(cached.argmax); + ggml_build_forward_expand(cached.gf, cached.argmax); + + cached.alloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend_)); + if (!cached.alloc || !ggml_gallocr_alloc_graph(cached.alloc, cached.gf)) { + std::fprintf(stderr, "laguna_verify_tree: cached gallocr_alloc_graph failed\n"); + cached.clear(); + return false; + } + + cached.owner = this; + cached.w_ptr = &w_; + cached.cache_ptr = &cache_; + cached.backend = backend_; + cached.target_feat = cache_.target_feat; + cached.n_alloc = N; + cached.mk_w = mk_w; + cached.kv_pad = kv_pad; + cached.capture_features = capture_features; + cached.pos.resize((size_t)N); + cached.rows.resize((size_t)N); + cached.feat_idx.resize((size_t)N); + cached.parents.resize((size_t)N); + cached.full_mask.resize((size_t)mk_w * (size_t)N); + cached.swa_mask.resize((size_t)mk_w * (size_t)N); + } + + std::vector embed((size_t)w_.n_embd * (size_t)N, 0.0f); + if (!w_.embedder.embed(flat_tokens.data(), N_actual, embed.data())) { + std::fprintf(stderr, "laguna_verify_tree: embed failed\n"); + return false; + } + if (!laguna_dflash_tensor_set_checked("verify_tree_cached.ie", cached.inp_embed, + embed.data(), 0, + ggml_nbytes(cached.inp_embed))) { + return false; + } + + std::fill(cached.pos.begin(), cached.pos.end(), 0); + cached.pos[0] = committed; + for (int i = 1; i < N_actual; ++i) { + cached.pos[(size_t)i] = committed + tree.depths[(size_t)i - 1]; + } + if (!laguna_dflash_tensor_set_checked("verify_tree_cached.positions", cached.positions, + cached.pos.data(), 0, + ggml_nbytes(cached.positions))) { + return false; + } + + for (int i = 0; i < N; ++i) cached.rows[(size_t)i] = committed + i; + if (!laguna_dflash_tensor_set_checked("verify_tree_cached.kv_idx", cached.kv_idx, + cached.rows.data(), 0, + ggml_nbytes(cached.kv_idx))) { + return false; + } + + std::fill(cached.parents.begin(), cached.parents.end(), -1); + cached.parents[0] = -1; + for (int i = 1; i < N_actual; ++i) { + cached.parents[(size_t)i] = tree.parents[(size_t)i]; + } + if (!laguna_dflash_tensor_set_checked("verify_tree_cached.parent_ids", cached.parent_ids, + cached.parents.data(), 0, + ggml_nbytes(cached.parent_ids), false)) { + return false; + } + + if (cached.feat_rows) { + for (int i = 0; i < N; ++i) { + cached.feat_idx[(size_t)i] = (committed + i) % cache_.target_feat_cap; + } + if (!laguna_dflash_tensor_set_checked("verify_tree_cached.feat_rows", + cached.feat_rows, + cached.feat_idx.data(), 0, + ggml_nbytes(cached.feat_rows), false)) { + return false; + } + } + + std::fill(cached.full_mask.begin(), cached.full_mask.end(), -INFINITY); + std::fill(cached.swa_mask.begin(), cached.swa_mask.end(), -INFINITY); + const int W = w_.sliding_window; + for (int q = 0; q < N_actual; ++q) { + const int depth_q = (q == 0) ? 0 : tree.depths[(size_t)q - 1]; + const int abs_q = committed + depth_q; + + for (int k = 0; k < committed && k < mk_w; ++k) { + cached.full_mask[(size_t)q * (size_t)mk_w + (size_t)k] = 0.0f; + } + + const int win_lo = std::max(0, abs_q - W + 1); + for (int k = win_lo; k < committed && k < mk_w; ++k) { + cached.swa_mask[(size_t)q * (size_t)mk_w + (size_t)k] = 0.0f; + } + + for (int j = 0; j < N_actual; ++j) { + if (!tree.visibility[(size_t)q * (size_t)N_actual + (size_t)j]) { + continue; + } + const int slot = committed + j; + if (slot < mk_w) { + cached.full_mask[(size_t)q * (size_t)mk_w + (size_t)slot] = 0.0f; + cached.swa_mask[(size_t)q * (size_t)mk_w + (size_t)slot] = 0.0f; + } + } + } + if (!laguna_dflash_tensor_set_checked("verify_tree_cached.mask_full", cached.mask_full, + cached.full_mask.data(), 0, + ggml_nbytes(cached.mask_full)) || + !laguna_dflash_tensor_set_checked("verify_tree_cached.mask_swa", cached.mask_swa, + cached.swa_mask.data(), 0, + ggml_nbytes(cached.mask_swa))) { + return false; + } + + if (ggml_backend_graph_compute(backend_, cached.gf) != GGML_STATUS_SUCCESS) { + std::fprintf(stderr, "laguna_verify_tree: cached graph_compute failed\n"); + return false; + } + + posterior_out.resize((size_t)N_actual); + ggml_backend_tensor_get(cached.argmax, posterior_out.data(), 0, + sizeof(int32_t) * (size_t)N_actual); + cache_.cur_pos = committed + N_actual; + return true; + } + const size_t arena_size = ggml_tensor_overhead() * 32768 + ggml_graph_overhead() + 32 * 1024 * 1024; static thread_local std::vector g_tree_arena; @@ -225,6 +519,8 @@ bool LagunaDFlashTarget::verify_tree( ggml_set_input(pp); ggml_tensor * kvi = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, N); ggml_set_input(kvi); + ggml_tensor * parent_ids = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, N); + ggml_set_input(parent_ids); ggml_tensor * mk_full = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, mk_w, N, 1, 1); ggml_set_input(mk_full); @@ -248,6 +544,7 @@ bool LagunaDFlashTarget::verify_tree( gi.kv_start = committed; gi.kv_pad = kv_pad; gi.kv_idx = kvi; + gi.tree_parent_ids = parent_ids; gi.output_last_only = false; gi.output_logits = true; gi.logits_are_output = logits_out != nullptr; @@ -285,25 +582,50 @@ bool LagunaDFlashTarget::verify_tree( ggml_free(ctx); return false; } - ggml_backend_tensor_set(ie, embed.data(), 0, ggml_nbytes(ie)); + if (!laguna_dflash_tensor_set_checked("verify_tree.ie", ie, embed.data(), 0, ggml_nbytes(ie))) { + ggml_free(ctx); + return false; + } std::vector pos((size_t)N, 0); pos[0] = committed; for (int i = 1; i < N_actual; ++i) { pos[(size_t)i] = committed + tree.depths[(size_t)i - 1]; } - ggml_backend_tensor_set(pp, pos.data(), 0, ggml_nbytes(pp)); + if (!laguna_dflash_tensor_set_checked("verify_tree.positions", pp, pos.data(), 0, ggml_nbytes(pp))) { + ggml_free(ctx); + return false; + } std::vector rows((size_t)N); for (int i = 0; i < N; ++i) rows[(size_t)i] = committed + i; - ggml_backend_tensor_set(kvi, rows.data(), 0, ggml_nbytes(kvi)); + if (!laguna_dflash_tensor_set_checked("verify_tree.kv_idx", kvi, rows.data(), 0, ggml_nbytes(kvi))) { + ggml_free(ctx); + return false; + } + + std::vector parents((size_t)N, -1); + parents[0] = -1; + for (int i = 1; i < N_actual; ++i) { + parents[(size_t)i] = tree.parents[(size_t)i]; + } + if (!laguna_dflash_tensor_set_checked("verify_tree.parent_ids", parent_ids, parents.data(), 0, + ggml_nbytes(parent_ids), false)) { + ggml_free(ctx); + return false; + } if (feat_rows) { std::vector feat_idx((size_t)N); for (int i = 0; i < N; ++i) { feat_idx[(size_t)i] = (committed + i) % cache_.target_feat_cap; } - ggml_backend_tensor_set(feat_rows, feat_idx.data(), 0, ggml_nbytes(feat_rows)); + if (!laguna_dflash_tensor_set_checked("verify_tree.feat_rows", feat_rows, + feat_idx.data(), 0, + ggml_nbytes(feat_rows), false)) { + ggml_free(ctx); + return false; + } } std::vector mfull((size_t)mk_w * (size_t)N, -INFINITY); @@ -333,8 +655,13 @@ bool LagunaDFlashTarget::verify_tree( } } } - 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)); + if (!laguna_dflash_tensor_set_checked("verify_tree.mask_full", mk_full, mfull.data(), 0, + ggml_nbytes(mk_full)) || + !laguna_dflash_tensor_set_checked("verify_tree.mask_swa", mk_swa, mswa.data(), 0, + ggml_nbytes(mk_swa))) { + ggml_free(ctx); + return false; + } if (ggml_backend_graph_compute(backend_, gf) != GGML_STATUS_SUCCESS) { std::fprintf(stderr, "laguna_verify_tree: graph_compute failed\n"); @@ -467,8 +794,13 @@ bool LagunaDFlashTarget::rollback_to_tree( committed + d; } } - ggml_backend_tensor_set(src_rows_kv, src_kv.data(), 0, ggml_nbytes(src_rows_kv)); - ggml_backend_tensor_set(dst_rows_kv, dst_kv.data(), 0, ggml_nbytes(dst_rows_kv)); + if (!laguna_dflash_tensor_set_checked("rollback_tree.src_rows_kv", src_rows_kv, + src_kv.data(), 0, ggml_nbytes(src_rows_kv)) || + !laguna_dflash_tensor_set_checked("rollback_tree.dst_rows_kv", dst_rows_kv, + dst_kv.data(), 0, ggml_nbytes(dst_rows_kv))) { + ggml_free(ctx); + return false; + } if (src_rows_feat && dst_rows_feat) { const int cap = cache_.target_feat_cap; @@ -478,10 +810,15 @@ bool LagunaDFlashTarget::rollback_to_tree( src_feat[(size_t)d] = (committed + accepted_dfs[(size_t)d]) % cap; dst_feat[(size_t)d] = (committed + d) % cap; } - ggml_backend_tensor_set(src_rows_feat, src_feat.data(), 0, - ggml_nbytes(src_rows_feat)); - ggml_backend_tensor_set(dst_rows_feat, dst_feat.data(), 0, - ggml_nbytes(dst_rows_feat)); + if (!laguna_dflash_tensor_set_checked("rollback_tree.src_rows_feat", src_rows_feat, + src_feat.data(), 0, + ggml_nbytes(src_rows_feat), false) || + !laguna_dflash_tensor_set_checked("rollback_tree.dst_rows_feat", dst_rows_feat, + dst_feat.data(), 0, + ggml_nbytes(dst_rows_feat), false)) { + ggml_free(ctx); + return false; + } } if (ggml_backend_graph_compute(backend_, gf) != GGML_STATUS_SUCCESS) { @@ -524,6 +861,53 @@ bool LagunaDFlashTarget::project_hidden_to_tokens( return laguna_project_hidden(backend_, w_, hidden, n_tokens, tokens_out); } +bool LagunaDFlashTarget::project_device_hidden_to_tokens( + ggml_tensor * hidden, + int n_tokens, + std::vector & tokens_out) { + if (!hidden || n_tokens <= 0 || + hidden->ne[0] != w_.n_embd || hidden->ne[1] < n_tokens) { + return false; + } + + ggml_init_params ip{}; + ip.mem_size = ggml_tensor_overhead() * 64 + ggml_graph_overhead() + 1024 * 1024; + ip.no_alloc = true; + ggml_context * ctx = ggml_init(ip); + if (!ctx) return false; + ggml_cgraph * gf = ggml_new_graph(ctx); + + ggml_tensor * inp = ggml_view_2d(ctx, hidden, w_.n_embd, n_tokens, + hidden->nb[1], 0); + ggml_tensor * logits = ggml_mul_mat(ctx, w_.output, inp); + ggml_tensor * argmax = ggml_argmax(ctx, logits); + ggml_set_output(argmax); + ggml_build_forward_expand(gf, argmax); + + static ggml_gallocr_t galloc_device_proj = nullptr; + if (!galloc_device_proj) { + galloc_device_proj = + ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend_)); + } + if (!ggml_gallocr_alloc_graph(galloc_device_proj, gf)) { + std::fprintf(stderr, "laguna_project_device_hidden_to_tokens: gallocr_alloc_graph failed\n"); + ggml_free(ctx); + return false; + } + + if (ggml_backend_graph_compute(backend_, gf) != GGML_STATUS_SUCCESS) { + std::fprintf(stderr, "laguna_project_device_hidden_to_tokens: graph_compute failed\n"); + ggml_free(ctx); + return false; + } + + tokens_out.resize((size_t)n_tokens); + ggml_backend_tensor_get(argmax, tokens_out.data(), 0, + sizeof(int32_t) * (size_t)n_tokens); + ggml_free(ctx); + return true; +} + bool LagunaDFlashTarget::project_hidden_to_logits( const float * hidden, int n_tokens, @@ -553,8 +937,12 @@ bool LagunaDFlashTarget::project_hidden_to_logits( return false; } - ggml_backend_tensor_set(inp, hidden, 0, - sizeof(float) * (size_t)n_tokens * (size_t)w_.n_embd); + if (!laguna_dflash_tensor_set_checked("project_hidden_to_logits.inp", inp, + hidden, 0, + sizeof(float) * (size_t)n_tokens * (size_t)w_.n_embd)) { + ggml_free(ctx); + return false; + } if (ggml_backend_graph_compute(backend_, gf) != GGML_STATUS_SUCCESS) { std::fprintf(stderr, "laguna_project_hidden_to_logits: graph_compute failed\n"); ggml_free(ctx); @@ -569,6 +957,65 @@ bool LagunaDFlashTarget::project_hidden_to_logits( return true; } +bool LagunaDFlashTarget::project_device_hidden_to_topk( + ggml_tensor * hidden, + int n_tokens, + int K, + float temperature, + std::vector & top_log_probs, + std::vector & top_token_ids) { + if (!hidden || n_tokens <= 0 || K <= 0 || + hidden->ne[0] != w_.n_embd || hidden->ne[1] < n_tokens) { + return false; + } + + ggml_init_params ip{}; + ip.mem_size = ggml_tensor_overhead() * 64 + ggml_graph_overhead() + 1024 * 1024; + ip.no_alloc = true; + ggml_context * ctx = ggml_init(ip); + if (!ctx) return false; + ggml_cgraph * gf = ggml_new_graph(ctx); + + ggml_tensor * inp = ggml_view_2d(ctx, hidden, w_.n_embd, n_tokens, + hidden->nb[1], 0); + ggml_tensor * logits = ggml_mul_mat(ctx, w_.output, inp); + ggml_set_output(logits); + ggml_build_forward_expand(gf, logits); + + static ggml_gallocr_t galloc_device_topk = nullptr; + if (!galloc_device_topk) { + galloc_device_topk = + ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend_)); + } + if (!ggml_gallocr_alloc_graph(galloc_device_topk, gf)) { + std::fprintf(stderr, "laguna_project_device_hidden_to_topk: gallocr_alloc_graph failed\n"); + ggml_free(ctx); + return false; + } + + if (ggml_backend_graph_compute(backend_, gf) != GGML_STATUS_SUCCESS) { + std::fprintf(stderr, "laguna_project_device_hidden_to_topk: graph_compute failed\n"); + ggml_free(ctx); + return false; + } + + const bool ok = laguna_extract_logits_topk(logits, n_tokens, K, temperature, + top_log_probs, top_token_ids); + ggml_free(ctx); + return ok; +} + +bool LagunaDFlashTarget::project_device_logits_to_topk( + ggml_tensor * logits, + int n_tokens, + int K, + float temperature, + std::vector & top_log_probs, + std::vector & top_token_ids) { + return laguna_extract_logits_topk(logits, n_tokens, K, temperature, + top_log_probs, top_token_ids); +} + bool LagunaDFlashTarget::project_hidden_to_topk( const float * hidden, int n_tokens, @@ -601,30 +1048,22 @@ bool LagunaDFlashTarget::project_hidden_to_topk( return false; } - ggml_backend_tensor_set(inp, hidden, 0, - sizeof(float) * (size_t)n_tokens * (size_t)w_.n_embd); + if (!laguna_dflash_tensor_set_checked("project_hidden_to_topk.inp", inp, + hidden, 0, + sizeof(float) * (size_t)n_tokens * (size_t)w_.n_embd)) { + ggml_free(ctx); + return false; + } if (ggml_backend_graph_compute(backend_, gf) != GGML_STATUS_SUCCESS) { std::fprintf(stderr, "laguna_project_hidden_to_topk: graph_compute failed\n"); ggml_free(ctx); return false; } - // CPU top-K via extract_draft_topk (shared with qwen35). The GPU ggml_top_k - // path bitonic-argsorts the full vocab, whose shared memory exceeds HIP - // shared-mem-per-block on gfx1151 (argsort.cu GGML_ASSERT). extract_draft_topk - // applies the temperature + log-softmax on host. - const int vocab = (int)logits->ne[0]; - std::vector logits_host((size_t)vocab * (size_t)n_tokens); - ggml_backend_tensor_get(logits, logits_host.data(), 0, - sizeof(float) * logits_host.size()); - - 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); - + const bool ok = laguna_extract_logits_topk(logits, n_tokens, K, temperature, + top_log_probs, top_token_ids); ggml_free(ctx); - return true; + return ok; } const std::vector & LagunaDFlashTarget::capture_layer_ids() const { diff --git a/server/src/laguna/laguna_dflash_target.h b/server/src/laguna/laguna_dflash_target.h index c3eee49d1..2bef4afa8 100644 --- a/server/src/laguna/laguna_dflash_target.h +++ b/server/src/laguna/laguna_dflash_target.h @@ -68,6 +68,26 @@ class LagunaDFlashTarget : public DFlashTarget { std::vector & top_log_probs, std::vector & top_token_ids) override; + bool project_device_hidden_to_tokens(ggml_tensor * hidden, + int n_tokens, + std::vector & tokens_out); + + bool project_device_hidden_to_topk(ggml_tensor * hidden, + int n_tokens, + int K, + float temperature, + std::vector & top_log_probs, + std::vector & top_token_ids); + + bool project_device_logits_to_topk(ggml_tensor * logits, + int n_tokens, + int K, + float temperature, + std::vector & top_log_probs, + std::vector & top_token_ids); + + ggml_tensor * output_weight() const { return w_.output; } + int hidden_size() const override { return w_.n_embd; } int mask_token_id() const override { return 12; } const std::vector & capture_layer_ids() const override; diff --git a/server/src/laguna/laguna_internal.h b/server/src/laguna/laguna_internal.h index 1fbc568f0..9f81893fc 100644 --- a/server/src/laguna/laguna_internal.h +++ b/server/src/laguna/laguna_internal.h @@ -277,6 +277,7 @@ struct LagunaGraphInputs { // CUDA-graph cache replays instead of re-launching every kernel. int kv_pad = 0; // 0 = legacy exact-length views + cpy append ggml_tensor * kv_idx = nullptr; // [n_tokens] I32 cache row indices (graph input) + ggml_tensor * tree_parent_ids = nullptr; // optional [n_tokens] I32, enables tree-attn carrier bool output_logits = true; bool logits_are_output = true; bool output_hidden_states = false; diff --git a/server/src/laguna/laguna_target_graph.cpp b/server/src/laguna/laguna_target_graph.cpp index bfeead3de..e97dce0a1 100644 --- a/server/src/laguna/laguna_target_graph.cpp +++ b/server/src/laguna/laguna_target_graph.cpp @@ -41,6 +41,24 @@ namespace dflash::common { static constexpr float LAGUNA_EPS = 1e-6f; +static bool laguna_tensor_set_checked(const char * label, + ggml_tensor * t, + const void * data, + size_t offset, + size_t size, + bool required = true) { + if (!t || !t->buffer) { + if (required) { + std::fprintf(stderr, "[laguna-graph] input tensor not allocated: %s\n", label); + set_last_error("laguna graph input tensor not allocated"); + return false; + } + return true; + } + ggml_backend_tensor_set(t, data, offset, size); + return true; +} + // ---- Cache lifecycle ---------------------------------------------------- bool create_laguna_target_cache(const LagunaTargetWeights & w, @@ -648,7 +666,8 @@ static ggml_tensor * build_laguna_attn_block( int n_tokens, bool is_full, int kv_pad = 0, - ggml_tensor * kv_idx = nullptr) + ggml_tensor * kv_idx = nullptr, + ggml_tensor * tree_parent_ids = nullptr) { const int head_dim = w.head_dim; const int n_head = w.n_head_arr[il]; @@ -802,8 +821,13 @@ static ggml_tensor * build_laguna_attn_block( const float kq_scale = 1.0f / std::sqrt((float)head_dim); // FULL -> attn_mask (causal). SWA -> attn_mask_swa (causal + sliding-window). ggml_tensor * use_mask = is_full ? attn_mask : attn_mask_swa; - ggml_tensor * attn = ggml_flash_attn_ext(ctx, Qfa, Kfa, Vfa, use_mask, - kq_scale, 0.0f, 0.0f); + static const bool g_tree_attn_op = + (std::getenv("DFLASH_LAGUNA_TREE_ATTN_OP") != nullptr); + ggml_tensor * attn = (g_tree_attn_op && tree_parent_ids && !is_full) + ? ggml_flash_attn_tree(ctx, Qfa, Kfa, Vfa, use_mask, + tree_parent_ids, positions, kq_scale) + : ggml_flash_attn_ext(ctx, Qfa, Kfa, Vfa, use_mask, + kq_scale, 0.0f, 0.0f); (void)win_start; (void)win_len; // attn: [head_dim, n_head, n_tokens] @@ -837,7 +861,8 @@ static ggml_tensor * build_laguna_layer( ggml_tensor * attn_mask_swa, const LagunaHybridMoe * hyb = nullptr, int kv_pad = 0, - ggml_tensor * kv_idx = nullptr) + ggml_tensor * kv_idx = nullptr, + ggml_tensor * tree_parent_ids = nullptr) { const LagunaTargetLayer & L = w.layers[il]; ggml_tensor * inp_f32 = graph_tensor_f32(ctx, inp); @@ -850,7 +875,7 @@ static ggml_tensor * build_laguna_layer( cur = build_laguna_attn_block(ctx, gf, w, L, il, cur, positions, cache.attn_k[il], cache.attn_v[il], attn_mask, attn_mask_swa, kv_start, n_tokens, is_full, - kv_pad, kv_idx); + kv_pad, kv_idx, tree_parent_ids); // Residual ggml_tensor * ffn_inp = ggml_add(ctx, cur, inp_f32); @@ -1044,7 +1069,8 @@ LagunaGraphOutputs build_laguna_graph( for (int il = 0; il < w.n_layer; ++il) { cur = build_laguna_layer(ctx, gf, w, cache, il, cur, in.positions, in.attn_mask, in.kv_start, in.n_tokens, - in.attn_mask_swa, in.hybrid, in.kv_pad, in.kv_idx); + in.attn_mask_swa, in.hybrid, in.kv_pad, in.kv_idx, + in.tree_parent_ids); // Feature capture for DFlash spec-decode: write residual-stream layer // outputs into the BF16 target feature ring. @@ -1274,16 +1300,29 @@ 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)); + if (!laguna_tensor_set_checked("laguna_step_cached.ie", cached.inp_embed, + embed, 0, ggml_nbytes(cached.inp_embed))) { + return false; + } 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)); + if (!laguna_tensor_set_checked("laguna_step_cached.positions", cached.positions, + &pos_val, 0, sizeof(pos_val))) { + return false; + } + if (!laguna_tensor_set_checked("laguna_step_cached.kv_idx", cached.kv_idx, + &pos_val, 0, sizeof(pos_val))) { + return false; + } std::fill(cached.full_mask.begin(), cached.full_mask.end(), -INFINITY); for (int k = 0; k <= kv_start && k < kv_len && k < mk_w; ++k) { cached.full_mask[(size_t)k] = 0.0f; } - ggml_backend_tensor_set(cached.mask_full, cached.full_mask.data(), 0, ggml_nbytes(cached.mask_full)); + if (!laguna_tensor_set_checked("laguna_step_cached.mask_full", cached.mask_full, + cached.full_mask.data(), 0, + ggml_nbytes(cached.mask_full))) { + return false; + } std::fill(cached.swa_mask.begin(), cached.swa_mask.end(), -INFINITY); const int W = w.sliding_window; @@ -1291,7 +1330,11 @@ bool laguna_step( for (int k = win_lo; k <= kv_start && k < kv_len && k < mk_w; ++k) { cached.swa_mask[(size_t)k] = 0.0f; } - ggml_backend_tensor_set(cached.mask_swa, cached.swa_mask.data(), 0, ggml_nbytes(cached.mask_swa)); + if (!laguna_tensor_set_checked("laguna_step_cached.mask_swa", cached.mask_swa, + cached.swa_mask.data(), 0, + ggml_nbytes(cached.mask_swa))) { + return false; + } if (ggml_backend_graph_compute(backend, cached.gf) != GGML_STATUS_SUCCESS) { std::fprintf(stderr, "laguna_step: cached graph_compute failed\n"); @@ -1376,10 +1419,16 @@ bool laguna_step( return false; } - ggml_backend_tensor_set(ie, embed, 0, ggml_nbytes(ie)); + if (!laguna_tensor_set_checked("laguna_step.ie", ie, embed, 0, ggml_nbytes(ie))) { + ggml_free(ctx); + return false; + } std::vector pos((size_t)n_tok); for (int i = 0; i < n_tok; ++i) pos[i] = kv_start + i; - ggml_backend_tensor_set(pp, pos.data(), 0, ggml_nbytes(pp)); + if (!laguna_tensor_set_checked("laguna_step.positions", pp, pos.data(), 0, ggml_nbytes(pp))) { + ggml_free(ctx); + return false; + } if (kvflash) { if (!kvi) { @@ -1395,12 +1444,18 @@ bool laguna_step( ggml_free(ctx); 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)); + if (!laguna_tensor_set_checked("laguna_step.kv_rows", kvi, rows.data(), 0, ggml_nbytes(kvi)) || + !laguna_tensor_set_checked("laguna_step.mask_full", mk_full, mfull.data(), 0, ggml_nbytes(mk_full)) || + !laguna_tensor_set_checked("laguna_step.mask_swa", mk_swa, mswa.data(), 0, ggml_nbytes(mk_swa))) { + ggml_free(ctx); + return false; + } } else { if (kvi) { - ggml_backend_tensor_set(kvi, pos.data(), 0, ggml_nbytes(kvi)); + if (!laguna_tensor_set_checked("laguna_step.kv_idx", kvi, pos.data(), 0, ggml_nbytes(kvi))) { + ggml_free(ctx); + return false; + } } if (!no_mask) { @@ -1413,7 +1468,11 @@ bool laguna_step( mfull[(size_t)q * mk_w + k] = 0.0f; } } - ggml_backend_tensor_set(mk_full, mfull.data(), 0, ggml_nbytes(mk_full)); + if (!laguna_tensor_set_checked("laguna_step.mask_full", mk_full, mfull.data(), 0, + ggml_nbytes(mk_full))) { + ggml_free(ctx); + return false; + } std::vector mswa((size_t)mk_w * n_tok, -INFINITY); const int W = w.sliding_window; @@ -1424,7 +1483,11 @@ bool laguna_step( mswa[(size_t)q * mk_w + k] = 0.0f; } } - ggml_backend_tensor_set(mk_swa, mswa.data(), 0, ggml_nbytes(mk_swa)); + if (!laguna_tensor_set_checked("laguna_step.mask_swa", mk_swa, mswa.data(), 0, + ggml_nbytes(mk_swa))) { + ggml_free(ctx); + return false; + } } } @@ -1466,6 +1529,216 @@ bool laguna_verify_batch( (void)token_ids; if (n_tokens <= 0) return false; + 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); + int kv_cap = 0; + for (int il = 0; il < w.n_layer; ++il) { + if (cache.attn_k[(size_t)il]) { kv_cap = (int)cache.attn_k[(size_t)il]->ne[1]; break; } + } + const int kv_pad = (!g_no_kvpad && kv_cap > 0) + ? std::min((kv_len + 255) & ~255, kv_cap) : 0; + const int mk_w = kv_pad > 0 ? kv_pad : kv_len; + + static const bool g_verify_cache_disabled = + (std::getenv("DFLASH_LAGUNA_VERIFY_CACHE_DISABLE") != nullptr); + const bool can_reuse_verify_graph = + !g_verify_cache_disabled && n_tokens <= 16 && !kvflash && out_logits == nullptr && + kv_pad > 0 && !g_pad_cpy; + if (can_reuse_verify_graph) { + struct CachedVerifyGraph { + const LagunaTargetWeights * w_ptr = nullptr; + LagunaTargetCache * cache_ptr = nullptr; + ggml_backend_t backend = nullptr; + ggml_tensor * target_feat = nullptr; + int n_tokens = 0; + int mk_w = 0; + int kv_pad = 0; + bool capture_features = false; + std::vector arena; + ggml_context * ctx = nullptr; + ggml_cgraph * gf = nullptr; + ggml_gallocr_t alloc = nullptr; + ggml_tensor * inp_embed = nullptr; + ggml_tensor * positions = nullptr; + ggml_tensor * kv_idx = nullptr; + ggml_tensor * mask_full = nullptr; + ggml_tensor * mask_swa = nullptr; + ggml_tensor * feat_rows = nullptr; + ggml_tensor * argmax = nullptr; + std::vector pos; + std::vector feat_idx; + std::vector full_mask; + std::vector swa_mask; + + void clear() { + if (alloc) { ggml_gallocr_free(alloc); alloc = nullptr; } + if (ctx) { ggml_free(ctx); ctx = nullptr; } + gf = nullptr; + inp_embed = positions = kv_idx = mask_full = mask_swa = feat_rows = argmax = nullptr; + w_ptr = nullptr; + cache_ptr = nullptr; + backend = nullptr; + target_feat = nullptr; + n_tokens = 0; + mk_w = 0; + kv_pad = 0; + capture_features = false; + } + }; + static thread_local CachedVerifyGraph cached_by_width[17]; + + CachedVerifyGraph & cached = cached_by_width[n_tokens]; + const bool capture_features = cache.target_feat && cache.target_feat_cap > 0; + const bool rebuild = + cached.ctx == nullptr || cached.w_ptr != &w || cached.cache_ptr != &cache || + cached.backend != backend || cached.target_feat != cache.target_feat || + cached.n_tokens != n_tokens || cached.mk_w != mk_w || + cached.kv_pad != kv_pad || cached.capture_features != capture_features; + if (rebuild) { + cached.clear(); + const size_t arena_size = + ggml_tensor_overhead() * 16384 + ggml_graph_overhead() + 16 * 1024 * 1024; + cached.arena.resize(arena_size); + + ggml_init_params ip{}; + ip.mem_size = arena_size; + ip.mem_buffer = cached.arena.data(); + ip.no_alloc = true; + cached.ctx = ggml_init(ip); + if (!cached.ctx) return false; + cached.gf = ggml_new_graph_custom(cached.ctx, 16384, false); + + cached.inp_embed = ggml_new_tensor_3d(cached.ctx, GGML_TYPE_F32, + w.n_embd, n_tokens, 1); + ggml_set_input(cached.inp_embed); + cached.positions = ggml_new_tensor_1d(cached.ctx, GGML_TYPE_I32, n_tokens); + ggml_set_input(cached.positions); + cached.kv_idx = ggml_new_tensor_1d(cached.ctx, GGML_TYPE_I32, n_tokens); + ggml_set_input(cached.kv_idx); + cached.mask_full = ggml_new_tensor_4d(cached.ctx, GGML_TYPE_F32, + mk_w, n_tokens, 1, 1); + ggml_set_input(cached.mask_full); + ggml_tensor * mask_full_cnv = ggml_cast(cached.ctx, cached.mask_full, GGML_TYPE_F16); + cached.mask_swa = ggml_new_tensor_4d(cached.ctx, GGML_TYPE_F32, + mk_w, n_tokens, 1, 1); + ggml_set_input(cached.mask_swa); + ggml_tensor * mask_swa_cnv = ggml_cast(cached.ctx, cached.mask_swa, GGML_TYPE_F16); + + if (capture_features) { + cached.feat_rows = ggml_new_tensor_1d(cached.ctx, GGML_TYPE_I32, n_tokens); + ggml_set_input(cached.feat_rows); + } + + LagunaGraphInputs gi{}; + gi.inp_embed = cached.inp_embed; + gi.positions = cached.positions; + gi.attn_mask = mask_full_cnv; + gi.attn_mask_swa = mask_swa_cnv; + gi.n_tokens = n_tokens; + gi.kv_start = 0; + gi.kv_pad = kv_pad; + gi.kv_idx = cached.kv_idx; + gi.output_last_only = false; + gi.output_logits = true; + gi.logits_are_output = false; + gi.capture_features = capture_features; + gi.target_feat_rows = cached.feat_rows; + gi.hybrid = nullptr; + + LagunaGraphOutputs go = build_laguna_graph(cached.ctx, cached.gf, w, cache, gi); + cached.argmax = ggml_argmax(cached.ctx, go.logits); + ggml_set_output(cached.argmax); + ggml_build_forward_expand(cached.gf, cached.argmax); + + cached.alloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); + if (!cached.alloc || !ggml_gallocr_alloc_graph(cached.alloc, cached.gf)) { + std::fprintf(stderr, "laguna_verify_batch: cached gallocr_alloc_graph failed\n"); + cached.clear(); + return false; + } + + cached.w_ptr = &w; + cached.cache_ptr = &cache; + cached.backend = backend; + cached.target_feat = cache.target_feat; + cached.n_tokens = n_tokens; + cached.mk_w = mk_w; + cached.kv_pad = kv_pad; + cached.capture_features = capture_features; + cached.pos.resize((size_t)n_tokens); + cached.feat_idx.resize((size_t)n_tokens); + cached.full_mask.resize((size_t)mk_w * (size_t)n_tokens); + cached.swa_mask.resize((size_t)mk_w * (size_t)n_tokens); + } + + if (!laguna_tensor_set_checked("laguna_verify_cached.ie", cached.inp_embed, + embed, 0, ggml_nbytes(cached.inp_embed))) { + return false; + } + + for (int i = 0; i < n_tokens; ++i) { + cached.pos[(size_t)i] = kv_start + i; + } + if (!laguna_tensor_set_checked("laguna_verify_cached.positions", cached.positions, + cached.pos.data(), 0, ggml_nbytes(cached.positions)) || + !laguna_tensor_set_checked("laguna_verify_cached.kv_idx", cached.kv_idx, + cached.pos.data(), 0, ggml_nbytes(cached.kv_idx))) { + return false; + } + + if (cached.feat_rows) { + for (int i = 0; i < n_tokens; ++i) { + cached.feat_idx[(size_t)i] = (kv_start + i) % cache.target_feat_cap; + } + if (!laguna_tensor_set_checked("laguna_verify_cached.feat_rows", cached.feat_rows, + cached.feat_idx.data(), 0, + ggml_nbytes(cached.feat_rows), false)) { + return false; + } + } + + std::fill(cached.full_mask.begin(), cached.full_mask.end(), -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 < mk_w; ++k) { + cached.full_mask[(size_t)q * (size_t)mk_w + (size_t)k] = 0.0f; + } + } + if (!laguna_tensor_set_checked("laguna_verify_cached.mask_full", cached.mask_full, + cached.full_mask.data(), 0, + ggml_nbytes(cached.mask_full))) { + return false; + } + + std::fill(cached.swa_mask.begin(), cached.swa_mask.end(), -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 < mk_w; ++k) { + cached.swa_mask[(size_t)q * (size_t)mk_w + (size_t)k] = 0.0f; + } + } + if (!laguna_tensor_set_checked("laguna_verify_cached.mask_swa", cached.mask_swa, + cached.swa_mask.data(), 0, + ggml_nbytes(cached.mask_swa))) { + return false; + } + + if (ggml_backend_graph_compute(backend, cached.gf) != GGML_STATUS_SUCCESS) { + std::fprintf(stderr, "laguna_verify_batch: cached graph_compute failed\n"); + return false; + } + + out_argmax.resize((size_t)n_tokens); + ggml_backend_tensor_get(cached.argmax, out_argmax.data(), 0, + sizeof(int32_t) * (size_t)n_tokens); + cache.cur_pos = kv_len; + cache.last_tok = out_argmax.empty() ? -1 : out_argmax.back(); + return true; + } + 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; @@ -1483,17 +1756,6 @@ bool laguna_verify_batch( 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); - int kv_cap = 0; - for (int il = 0; il < w.n_layer; ++il) { - if (cache.attn_k[(size_t)il]) { kv_cap = (int)cache.attn_k[(size_t)il]->ne[1]; break; } - } - const int kv_pad = (!g_no_kvpad && kv_cap > 0) - ? 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); @@ -1544,16 +1806,26 @@ bool laguna_verify_batch( return false; } - ggml_backend_tensor_set(ie, embed, 0, ggml_nbytes(ie)); + if (!laguna_tensor_set_checked("laguna_verify.ie", ie, embed, 0, ggml_nbytes(ie))) { + ggml_free(ctx); + return false; + } 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)); + if (!laguna_tensor_set_checked("laguna_verify.positions", pp, pos.data(), 0, ggml_nbytes(pp))) { + ggml_free(ctx); + return false; + } 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; } - ggml_backend_tensor_set(feat_rows, feat_idx.data(), 0, ggml_nbytes(feat_rows)); + if (!laguna_tensor_set_checked("laguna_verify.feat_rows", feat_rows, + feat_idx.data(), 0, ggml_nbytes(feat_rows), false)) { + ggml_free(ctx); + return false; + } } if (kvflash) { @@ -1570,12 +1842,18 @@ bool laguna_verify_batch( ggml_free(ctx); 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)); + if (!laguna_tensor_set_checked("laguna_verify.kv_rows", kvi, rows.data(), 0, ggml_nbytes(kvi)) || + !laguna_tensor_set_checked("laguna_verify.mask_full", mk_full, mfull.data(), 0, ggml_nbytes(mk_full)) || + !laguna_tensor_set_checked("laguna_verify.mask_swa", mk_swa, mswa.data(), 0, ggml_nbytes(mk_swa))) { + ggml_free(ctx); + return false; + } } else { if (kvi) { - ggml_backend_tensor_set(kvi, pos.data(), 0, ggml_nbytes(kvi)); + if (!laguna_tensor_set_checked("laguna_verify.kv_idx", kvi, pos.data(), 0, ggml_nbytes(kvi))) { + ggml_free(ctx); + return false; + } } std::vector mfull((size_t)mk_w * n_tokens, -INFINITY); @@ -1585,7 +1863,11 @@ bool laguna_verify_batch( mfull[(size_t)q * mk_w + k] = 0.0f; } } - ggml_backend_tensor_set(mk_full, mfull.data(), 0, ggml_nbytes(mk_full)); + if (!laguna_tensor_set_checked("laguna_verify.mask_full", mk_full, mfull.data(), 0, + ggml_nbytes(mk_full))) { + ggml_free(ctx); + return false; + } std::vector mswa((size_t)mk_w * n_tokens, -INFINITY); const int W = w.sliding_window; @@ -1596,7 +1878,11 @@ bool laguna_verify_batch( mswa[(size_t)q * mk_w + k] = 0.0f; } } - ggml_backend_tensor_set(mk_swa, mswa.data(), 0, ggml_nbytes(mk_swa)); + if (!laguna_tensor_set_checked("laguna_verify.mask_swa", mk_swa, mswa.data(), 0, + ggml_nbytes(mk_swa))) { + ggml_free(ctx); + return false; + } } if (ggml_backend_graph_compute(backend, gf) != GGML_STATUS_SUCCESS) { @@ -1780,10 +2066,16 @@ bool laguna_step_hybrid( return false; } - ggml_backend_tensor_set(ie, embed, 0, ggml_nbytes(ie)); + if (!laguna_tensor_set_checked("laguna_step_hybrid.ie", ie, embed, 0, ggml_nbytes(ie))) { + ggml_free(ctx); + return false; + } std::vector pos((size_t)n_tok); for (int i = 0; i < n_tok; ++i) pos[i] = kv_start + i; - ggml_backend_tensor_set(pp, pos.data(), 0, ggml_nbytes(pp)); + if (!laguna_tensor_set_checked("laguna_step_hybrid.positions", pp, pos.data(), 0, ggml_nbytes(pp))) { + ggml_free(ctx); + return false; + } if (kvflash) { if (!kvi) { @@ -1799,13 +2091,19 @@ bool laguna_step_hybrid( ggml_free(ctx); 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)); + if (!laguna_tensor_set_checked("laguna_step_hybrid.kv_rows", kvi, rows.data(), 0, ggml_nbytes(kvi)) || + !laguna_tensor_set_checked("laguna_step_hybrid.mask_full", mk_full, mfull.data(), 0, ggml_nbytes(mk_full)) || + !laguna_tensor_set_checked("laguna_step_hybrid.mask_swa", mk_swa, mswa.data(), 0, ggml_nbytes(mk_swa))) { + ggml_free(ctx); + return false; + } } else { if (kvi) { // set_rows row indices = absolute cache positions of this step's tokens - ggml_backend_tensor_set(kvi, pos.data(), 0, ggml_nbytes(kvi)); + if (!laguna_tensor_set_checked("laguna_step_hybrid.kv_idx", kvi, pos.data(), 0, ggml_nbytes(kvi))) { + ggml_free(ctx); + return false; + } } if (!no_mask) { @@ -1816,7 +2114,11 @@ bool laguna_step_hybrid( 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; } - ggml_backend_tensor_set(mk_full, mfull.data(), 0, ggml_nbytes(mk_full)); + if (!laguna_tensor_set_checked("laguna_step_hybrid.mask_full", mk_full, mfull.data(), 0, + ggml_nbytes(mk_full))) { + ggml_free(ctx); + return false; + } std::vector mswa((size_t)mk_w * n_tok, -INFINITY); const int Wsw = w.sliding_window; for (int q = 0; q < n_tok; ++q) { @@ -1824,7 +2126,11 @@ bool laguna_step_hybrid( const int lo = std::max(0, abs_q - Wsw + 1); for (int k = lo; k <= abs_q && k < kv_len; ++k) mswa[(size_t)q * mk_w + k] = 0.0f; } - ggml_backend_tensor_set(mk_swa, mswa.data(), 0, ggml_nbytes(mk_swa)); + if (!laguna_tensor_set_checked("laguna_step_hybrid.mask_swa", mk_swa, mswa.data(), 0, + ggml_nbytes(mk_swa))) { + ggml_free(ctx); + return false; + } } } @@ -1840,8 +2146,13 @@ bool laguna_step_hybrid( vldbuf[(size_t)mi * n_expert + g] = loc >= 0 ? 1.0f : 0.0f; } } - ggml_backend_tensor_set(hybm.lut_all, lutbuf.data(), 0, sizeof(int32_t) * lutbuf.size()); - ggml_backend_tensor_set(hybm.vld_all, vldbuf.data(), 0, sizeof(float) * vldbuf.size()); + if (!laguna_tensor_set_checked("laguna_step_hybrid.lut_all", hybm.lut_all, + lutbuf.data(), 0, sizeof(int32_t) * lutbuf.size()) || + !laguna_tensor_set_checked("laguna_step_hybrid.vld_all", hybm.vld_all, + vldbuf.data(), 0, sizeof(float) * vldbuf.size())) { + ggml_free(ctx); + return false; + } if (ggml_backend_graph_compute(backend, gf) != GGML_STATUS_SUCCESS) { std::fprintf(stderr, "laguna_step_hybrid: graph_compute failed\n"); diff --git a/server/test/bench_laguna_spark.cpp b/server/test/bench_laguna_spark.cpp index 4dcb93584..287f6cb54 100644 --- a/server/test/bench_laguna_spark.cpp +++ b/server/test/bench_laguna_spark.cpp @@ -7,12 +7,15 @@ // DFLASH_LAGUNA_NO_SINGLE_GRAPH=1 per-layer fallback (for trace capture) // DFLASH_LAGUNA_PREGATE_TRACE= pregate trace capture (fallback path) // -// Usage: bench_laguna_spark [prompt_N=128] [n_gen=256] +// Usage: +// bench_laguna_spark [prompt_N=128] [n_gen=256] +// bench_laguna_spark --draft --ddtree --ddtree-budget 22 #include "laguna_backend.h" #include "dflash27b.h" #include +#include #include #include #include @@ -20,17 +23,124 @@ using namespace dflash::common; +static bool env_true(const char * name) { + const char * v = std::getenv(name); + return v && v[0] != '\0' && std::strcmp(v, "0") != 0; +} + +static int env_int(const char * name, int fallback) { + const char * v = std::getenv(name); + return v ? std::atoi(v) : fallback; +} + +static float env_float(const char * name, float fallback) { + const char * v = std::getenv(name); + return v ? (float)std::atof(v) : fallback; +} + +static bool parse_kv_type(const char * s, ggml_type & out) { + if (!s) return false; + if (std::strcmp(s, "q4_0") == 0) { out = GGML_TYPE_Q4_0; return true; } + if (std::strcmp(s, "q5_0") == 0) { out = GGML_TYPE_Q5_0; return true; } + if (std::strcmp(s, "q8_0") == 0) { out = GGML_TYPE_Q8_0; return true; } + if (std::strcmp(s, "f16") == 0) { out = GGML_TYPE_F16; return true; } + return false; +} + int main(int argc, char ** argv) { if (argc < 2) { - std::fprintf(stderr, "usage: %s [prompt_N=128] [n_gen=256]\n", argv[0]); + std::fprintf(stderr, + "usage: %s [prompt_N=128] [n_gen=256] " + "[--draft draft.gguf] [--ddtree] [--ddtree-budget N]\n", + argv[0]); return 2; } - const int prompt_N = std::max(1, (argc >= 3) ? std::atoi(argv[2]) : 128); - const int n_gen = std::max(1, (argc >= 4) ? std::atoi(argv[3]) : 256); + + std::vector positional; + std::string draft_path = std::getenv("DFLASH_DRAFT") ? std::getenv("DFLASH_DRAFT") : ""; + bool ddtree_mode = env_true("DFLASH_DDTREE") || env_true("DFLASH_LAGUNA_DDTREE"); + int ddtree_budget = env_int("DFLASH_DDTREE_BUDGET", env_int("DFLASH_LAGUNA_DDTREE_BUDGET", 22)); + float ddtree_temp = env_float("DFLASH_DDTREE_TEMP", env_float("DFLASH_LAGUNA_DDTREE_TEMP", 1.0f)); + int verify_width = env_int("DFLASH_VERIFY_WIDTH", 0); + int max_ctx_arg = env_int("DFLASH_MAX_CTX", 0); + int draft_gpu = env_int("DFLASH_DRAFT_GPU", -1); + int draft_ctx_max = env_int("DFLASH_DRAFT_CTX_MAX", 4096); + ggml_type kv_type = GGML_TYPE_Q8_0; + parse_kv_type(std::getenv("DFLASH_KV_TYPE"), kv_type); + + for (int i = 1; i < argc; ++i) { + const char * a = argv[i]; + if (std::strcmp(a, "--draft") == 0 && i + 1 < argc) { + draft_path = argv[++i]; + } else if (std::strcmp(a, "--ddtree") == 0) { + ddtree_mode = true; + } else if (std::strcmp(a, "--ddtree-budget") == 0 && i + 1 < argc) { + ddtree_budget = std::max(1, std::atoi(argv[++i])); + } else if (std::strncmp(a, "--ddtree-budget=", 16) == 0) { + ddtree_budget = std::max(1, std::atoi(a + 16)); + } else if (std::strcmp(a, "--ddtree-temp") == 0 && i + 1 < argc) { + ddtree_temp = (float)std::atof(argv[++i]); + } else if (std::strncmp(a, "--ddtree-temp=", 14) == 0) { + ddtree_temp = (float)std::atof(a + 14); + } else if (std::strcmp(a, "--verify-width") == 0 && i + 1 < argc) { + verify_width = std::max(0, std::atoi(argv[++i])); + } else if (std::strncmp(a, "--verify-width=", 15) == 0) { + verify_width = std::max(0, std::atoi(a + 15)); + } else if (std::strcmp(a, "--max-ctx") == 0 && i + 1 < argc) { + max_ctx_arg = std::max(1, std::atoi(argv[++i])); + } else if (std::strncmp(a, "--max-ctx=", 10) == 0) { + max_ctx_arg = std::max(1, std::atoi(a + 10)); + } else if (std::strcmp(a, "--draft-gpu") == 0 && i + 1 < argc) { + draft_gpu = std::atoi(argv[++i]); + } else if (std::strncmp(a, "--draft-gpu=", 12) == 0) { + draft_gpu = std::atoi(a + 12); + } else if (std::strcmp(a, "--draft-ctx-max") == 0 && i + 1 < argc) { + draft_ctx_max = std::max(1, std::atoi(argv[++i])); + } else if (std::strncmp(a, "--draft-ctx-max=", 16) == 0) { + draft_ctx_max = std::max(1, std::atoi(a + 16)); + } else if (std::strcmp(a, "--kv") == 0 && i + 1 < argc) { + if (!parse_kv_type(argv[++i], kv_type)) { + std::fprintf(stderr, "unknown --kv type: %s\n", argv[i]); + return 2; + } + } else if (std::strncmp(a, "--kv=", 5) == 0) { + if (!parse_kv_type(a + 5, kv_type)) { + std::fprintf(stderr, "unknown --kv type: %s\n", a + 5); + return 2; + } + } else if (a[0] == '-') { + std::fprintf(stderr, "unknown option: %s\n", a); + return 2; + } else { + positional.push_back(a); + } + } + + if (positional.empty()) { + std::fprintf(stderr, "missing laguna.gguf\n"); + return 2; + } + const int prompt_N = std::max(1, positional.size() >= 2 ? std::atoi(positional[1]) : 128); + const int n_gen = std::max(1, positional.size() >= 3 ? std::atoi(positional[2]) : 256); LagunaBackendArgs args; - args.target_path = argv[1]; - args.max_ctx = prompt_N + n_gen + 64; + args.target_path = positional[0]; + args.draft_path = draft_path; + args.draft_gpu = draft_gpu; + args.draft_ctx_max = draft_ctx_max; + args.ddtree_mode = ddtree_mode; + args.ddtree_budget = ddtree_budget; + args.ddtree_temp = ddtree_temp > 0.0f ? ddtree_temp : 1.0f; + args.verify_width = verify_width; + args.max_ctx = max_ctx_arg > 0 ? max_ctx_arg : prompt_N + n_gen + 64; + args.kv_type = kv_type; + + std::printf("[spark-bench] target=%s draft=%s ddtree=%d budget=%d temp=%.2f " + "verify_width=%d max_ctx=%d kv=%s\n", + args.target_path.c_str(), + args.draft_path.empty() ? "(none)" : args.draft_path.c_str(), + (int)args.ddtree_mode, args.ddtree_budget, args.ddtree_temp, + args.verify_width, args.max_ctx, ggml_type_name(args.kv_type)); LagunaBackend be(args); if (!be.init()) {