From bbc7de475178dd0535c16ad85f204a2529806c9d Mon Sep 17 00:00:00 2001 From: Kawrakow Date: Fri, 3 Jul 2026 18:40:20 +0200 Subject: [PATCH 1/5] Fix Windows build after CUB addition (#2076) --- ggml/src/CMakeLists.txt | 6 +++++- ggml/src/ggml-cuda/argsort.cu | 8 ++++++++ 2 files changed, 13 insertions(+), 1 deletion(-) diff --git a/ggml/src/CMakeLists.txt b/ggml/src/CMakeLists.txt index 1c13ef122..1d680c3c3 100644 --- a/ggml/src/CMakeLists.txt +++ b/ggml/src/CMakeLists.txt @@ -1092,7 +1092,7 @@ if (GGML_CUDA) list(APPEND CUDA_FLAGS -compress-mode=${GGML_CUDA_COMPRESSION_MODE}) endif() - + if (GGML_FATAL_WARNINGS) list(APPEND CUDA_FLAGS -Werror all-warnings) endif() @@ -1133,6 +1133,10 @@ if (GGML_CUDA) if (NOT MSVC) list(APPEND CUDA_CXX_FLAGS -Wno-pedantic) + else() + # CCCL 3.2 onwards will require a cpp-standard-compliant preprocessor for MSVC + # https://github.com/NVIDIA/cccl/pull/6827 + list(APPEND CUDA_CXX_FLAGS /Zc:preprocessor) endif() endif() diff --git a/ggml/src/ggml-cuda/argsort.cu b/ggml/src/ggml-cuda/argsort.cu index 28531cce0..0aaf5861d 100644 --- a/ggml/src/ggml-cuda/argsort.cu +++ b/ggml/src/ggml-cuda/argsort.cu @@ -431,6 +431,14 @@ static void argsort_openai_f32_f32_i32_cuda(const float * x, float * weights, in #endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) && CUDART_VERSION >= 11070 #ifdef GGML_CUDA_USE_CUB +# ifdef _WIN32 +# ifndef WIN32_LEAN_AND_MEAN +# define WIN32_LEAN_AND_MEAN +# endif +# ifndef NOMINMAX +# define NOMINMAX +# endif +# endif # include # if (CCCL_MAJOR_VERSION >= 3 && CCCL_MINOR_VERSION >= 1) # define STRIDED_ITERATOR_AVAILABLE From 92e60231da485e46871707c4db8375d011dace55 Mon Sep 17 00:00:00 2001 From: rankaiyx Date: Mon, 6 Jul 2026 15:26:45 +0800 Subject: [PATCH 2/5] fix(rpc): update ggml_backend_cuda_init to 3-arg signature (#2084) Commit 75a5f6d0 (Per model CUDA contexts) added a third parameter "model" to ggml_backend_cuda_init, but rpc-server was missed. The RPC server creates backends at startup before any model is loaded, so passing nullptr for the model parameter is correct: the RPC server acts as a pure compute proxy and does not host models locally. The nullptr serves as the context grouping key in all_ctx, which is sufficient since all RPC backends in the same process share the same group for multi-GPU reduce operations. Fixes: cmake build failure with GGML_CUDA=ON + GGML_RPC=ON --- examples/rpc/rpc-server.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/rpc/rpc-server.cpp b/examples/rpc/rpc-server.cpp index e67ca7f3a..866aa4c41 100644 --- a/examples/rpc/rpc-server.cpp +++ b/examples/rpc/rpc-server.cpp @@ -263,7 +263,7 @@ static ggml_backend_t create_gpu_backend(const rpc_server_params& params, uint32 ggml_backend_t backend = NULL; #ifdef GGML_USE_CUDA fprintf(stderr, "%s: using CUDA backend: CUDA%d\n", __func__, device); - backend = ggml_backend_cuda_init(device, nullptr); // init device + backend = ggml_backend_cuda_init(device, nullptr, nullptr); // init device if (!backend) { fprintf(stderr, "%s: ggml_backend_cuda_init() failed\n", __func__); } From 0a415bde4a3933daf1a13a04aa1499cbe4948f4c Mon Sep 17 00:00:00 2001 From: Samuel Oliveira Alves <107287165+SamuelOliveirads@users.noreply.github.com> Date: Mon, 6 Jul 2026 04:46:18 -0300 Subject: [PATCH 3/5] Feat: Llama-cli Spec (#2081) * feat: implement speculative decoding support in llama-cli * fix spec mismatchs and metrics --- examples/main/main.cpp | 502 +++++++++++++++++++++++++++++++++++------ 1 file changed, 438 insertions(+), 64 deletions(-) diff --git a/examples/main/main.cpp b/examples/main/main.cpp index 2bd2b8d48..2a45c3a67 100644 --- a/examples/main/main.cpp +++ b/examples/main/main.cpp @@ -1,4 +1,5 @@ #include "common.h" +#include "speculative.h" #include "chat.h" #include "console.h" #include "llama.h" @@ -10,6 +11,7 @@ #include #include #include +#include #include #include #include @@ -135,6 +137,8 @@ int main(int argc, char ** argv) { return 1; } + common_speculative_prepare_startup(params); + common_params_sampling & sparams = params.sparams; #ifndef LOG_DISABLE_LOGS @@ -199,6 +203,8 @@ int main(int argc, char ** argv) { llama_model * model; llama_context * ctx; llama_context * ctx_guidance = NULL; + common_speculative * spec = nullptr; + common_sampler * ctx_sampling = nullptr; std::vector chat_msgs; g_model = &model; g_ctx = &ctx; @@ -218,8 +224,100 @@ int main(int argc, char ** argv) { LOG_TEE("%s: error: unable to load model\n", __func__); return 1; } + + const bool requested_spec_user = params.speculative.has_stage_chain(); + + if (!common_speculative_finalize_startup(params, model)) { + if (ctx_guidance) { + llama_free(ctx_guidance); + } + params.speculative.clear_dft(); + llama_free(ctx); + llama_free_model(model); + llama_backend_free(); + return 1; + } + + const bool requested_spec = params.speculative.has_stage_chain(); + if (requested_spec_user && !requested_spec) { + LOG_TEE("%s: error: requested speculative decoding is not runnable with the finalized model/runtime configuration\n", __func__); + if (ctx_guidance) { + llama_free(ctx_guidance); + } + params.speculative.clear_dft(); + llama_free(ctx); + llama_free_model(model); + llama_backend_free(); + return 1; + } + + auto fail_spec = [&](const char * reason) { + LOG_TEE("%s: error: speculative decoding is not supported with %s in llama-cli yet\n", __func__, reason); + if (ctx_guidance) { + llama_free(ctx_guidance); + } + if (spec) { + common_speculative_free(spec); + } + params.speculative.clear_dft(); + llama_free(ctx); + llama_free_model(model); + llama_backend_free(); + return 1; + }; + + if (requested_spec) { + if (params.interactive || params.interactive_first || params.conversation) { + return fail_spec("interactive or conversation mode"); + } + if (sparams.cfg_scale > 1.f) { + return fail_spec("CFG guidance"); + } + if (params.grp_attn_n != 1) { + return fail_spec("self-extend"); + } + if (!params.path_prompt_cache.empty()) { + return fail_spec("prompt-cache save/load"); + } + if (llama_model_has_encoder(model)) { + return fail_spec("encoder-decoder models"); + } + } + auto chat_templates = common_chat_templates_init(model, params.chat_template); + if (params.has_mtp) { + llama_set_embeddings(ctx, true); + } + + if (requested_spec) { + if (!common_speculative_is_compat(ctx)) { + LOG_TEE("%s: error: speculative decoding is not supported by this context\n", __func__); + if (ctx_guidance) { + llama_free(ctx_guidance); + } + params.speculative.clear_dft(); + llama_free(ctx); + llama_free_model(model); + llama_backend_free(); + return 1; + } + + switch (common_speculative_try_init(params.speculative, ctx, &spec)) { + case COMMON_SPECULATIVE_INIT_READY: + LOG_TEE("%s: speculative decoding context initialized\n", __func__); + break; + case COMMON_SPECULATIVE_INIT_ERR_RECURRENT: + return fail_spec("recurrent speculative context initialization failure"); + case COMMON_SPECULATIVE_INIT_ERR_MTP: + return fail_spec("MTP speculative context initialization failure"); + case COMMON_SPECULATIVE_INIT_ERR_GENERIC: + return fail_spec("speculative context initialization failure"); + case COMMON_SPECULATIVE_INIT_SKIPPED: + break; + } + } + const int n_ctx_train = llama_n_ctx_train(model); const int n_ctx = llama_n_ctx(ctx); LOG("n_ctx: %d\n", n_ctx); @@ -306,6 +404,7 @@ int main(int argc, char ** argv) { LOG("prompt: \"%s\"\n", log_tostr(prompt)); LOG("tokens: %s\n", LOG_TOKENS_TOSTR_PRETTY(ctx, embd_inp).c_str()); + } // Should not run without any tokens @@ -533,6 +632,13 @@ int main(int argc, char ** argv) { std::ostringstream output_ss; g_output_ss = &output_ss; std::ostringstream assistant_ss; // for storing current assistant message, used in conversation mode + const int64_t t_start_process_prompt_us = ggml_time_us(); + int64_t t_start_generation_us = 0; + double t_prompt_processing_ms = 0.0; + double t_token_generation_ms = 0.0; + int n_prompt_tokens_processed = 0; + int n_decoded = 0; + // the first thing we will do is to output the prompt, so set color accordingly console::set_display(console::prompt); display = params.display_prompt; @@ -548,7 +654,7 @@ int main(int argc, char ** argv) { antiprompt_ids.emplace_back(::common_tokenize(ctx, antiprompt, false, true)); } - common_sampler * ctx_sampling = common_sampler_init(model, sparams); + ctx_sampling = common_sampler_init(model, sparams); if (!ctx_sampling) { fprintf(stderr, "%s: failed to initialize sampling subsystem\n", __func__); exit(1); @@ -572,6 +678,17 @@ int main(int argc, char ** argv) { embd_inp.push_back(decoder_start_token_id); } + bool embd_is_prompt = false; + bool emitted_generated = false; + std::vector emitted; + llama_tokens speculative_tokens = embd_inp; + bool emitted_hit_eog = false; + bool speculative_started = false; + int32_t final_prompt_output_index = -1; + llama_pos final_prompt_hidden_pos = -1; + bool have_speculative_sampled = false; + llama_token speculative_sampled = LLAMA_TOKEN_NULL; + while ((n_remain != 0 && !is_antiprompt) || params.interactive) { // predict if (!embd.empty()) { @@ -609,6 +726,14 @@ int main(int argc, char ** argv) { llama_kv_cache_seq_rm (ctx, 0, params.n_keep , params.n_keep + n_discard); llama_kv_cache_seq_add(ctx, 0, params.n_keep + n_discard, n_past, -n_discard); + if (spec != nullptr) { + common_speculative_context_shift(spec, 0, params.n_keep, n_discard, n_past); + if ((int) speculative_tokens.size() > params.n_keep) { + const size_t erase_begin = (size_t) params.n_keep; + const size_t erase_end = std::min(speculative_tokens.size(), erase_begin + (size_t) n_discard); + speculative_tokens.erase(speculative_tokens.begin() + erase_begin, speculative_tokens.begin() + erase_end); + } + } n_past -= n_discard; @@ -717,11 +842,52 @@ int main(int argc, char ** argv) { LOG("eval: %s\n", LOG_TOKENS_TOSTR_PRETTY(ctx, embd).c_str()); - if (llama_decode(ctx, llama_batch_get_one(&embd[i], n_eval, n_past, 0))) { + const bool need_prompt_target_features = + embd_is_prompt && + spec != nullptr && + (params.speculative.has_stage_type(COMMON_SPECULATIVE_TYPE_MTP) || + params.speculative.has_stage_type(COMMON_SPECULATIVE_TYPE_DFLASH)); + + llama_batch batch = {}; + if (need_prompt_target_features) { + batch = llama_batch_init(n_eval, 0, 1); + for (int j = 0; j < n_eval; ++j) { + common_batch_add(batch, embd[i + j], n_past + j, { 0 }, true); + } + } else { + batch = llama_batch_get_one(&embd[i], n_eval, n_past, 0); + } + + if (llama_decode(ctx, batch)) { + if (need_prompt_target_features) { + llama_batch_free(batch); + } LOG_TEE("%s : failed to eval\n", __func__); return 1; } + if (spec != nullptr) { + if (need_prompt_target_features) { + if (common_speculative_on_target_seq_batch(spec, ctx, batch, 0, true) != 0) { + llama_batch_free(batch); + LOG_TEE("%s : failed to warm speculative target state\n", __func__); + return 1; + } + final_prompt_output_index = n_eval - 1; + final_prompt_hidden_pos = n_past + n_eval - 1; + } else if (!embd_is_prompt) { + speculative_tokens.insert(speculative_tokens.end(), embd.begin() + i, embd.begin() + i + n_eval); + } + } + + if (need_prompt_target_features) { + llama_batch_free(batch); + } + + if (embd_is_prompt) { + n_prompt_tokens_processed += n_eval; + } + n_past += n_eval; LOG("n_past = %d\n", n_past); @@ -737,10 +903,34 @@ int main(int argc, char ** argv) { } } + emitted.clear(); embd.clear(); embd_guidance.clear(); if ((int) embd_inp.size() <= n_consumed && !is_interacting) { + if (!speculative_started) { + if (spec != nullptr) { + static const llama_tokens empty_speculative_prompt; + const llama_tokens & speculative_prompt = + params.speculative.has_stage_type(COMMON_SPECULATIVE_TYPE_MTP) && + !params.speculative.has_composite_stage_chain() + ? empty_speculative_prompt + : speculative_tokens; + common_speculative_begin(spec, speculative_prompt); + if (params.speculative.has_stage_type(COMMON_SPECULATIVE_TYPE_MTP) && + final_prompt_output_index >= 0 && + final_prompt_hidden_pos >= 0 && + !common_speculative_capture_output_hidden(spec, ctx, final_prompt_output_index, 0, final_prompt_hidden_pos)) { + LOG_TEE("%s: failed to capture final prompt hidden state for speculative init (output_index=%d, pos=%d)\n", + __func__, final_prompt_output_index, final_prompt_hidden_pos); + } + } + if (params.has_mtp) { + llama_set_embeddings(ctx, false); + } + speculative_started = true; + } + // optionally save the session on first sample (for faster prompt loading next time) if (!path_session.empty() && need_to_save_session && !params.prompt_cache_ro) { need_to_save_session = false; @@ -749,21 +939,153 @@ int main(int argc, char ** argv) { LOG("saved session to %s\n", path_session.c_str()); } - const llama_token id = common_sampler_sample_legacy(ctx_sampling, ctx, ctx_guidance); + const int n_predict_budget = n_remain < 0 ? std::numeric_limits::max() : n_remain; + bool used_speculative = false; - common_sampler_accept(ctx_sampling, ctx, id, /* apply_grammar= */ true); + if (spec != nullptr && n_predict_budget != 1) { + const bool sampled_before_from_carry = have_speculative_sampled; + llama_token sampled_before = LLAMA_TOKEN_NULL; + if (sampled_before_from_carry) { + sampled_before = speculative_sampled; + have_speculative_sampled = false; + speculative_sampled = LLAMA_TOKEN_NULL; + } else { + sampled_before = common_sampler_sample_legacy(ctx_sampling, ctx, ctx_guidance); + common_sampler_accept(ctx_sampling, ctx, sampled_before, /* apply_grammar= */ true); + } + static const llama_tokens empty_speculative_tokens; + const llama_tokens & draft_history = + params.speculative.has_stage_type(COMMON_SPECULATIVE_TYPE_MTP) && + !params.speculative.has_composite_stage_chain() + ? empty_speculative_tokens + : speculative_tokens; + auto draft_result = common_speculative_draft_ex( + spec, + ctx, + params.speculative, + draft_history, + sampled_before, + n_past, + 0); - LOG("last: %s\n", LOG_TOKENS_TOSTR_PRETTY(ctx, ctx_sampling->prev).c_str()); + auto & draft = draft_result.tokens; + int max_usable_draft = (int) draft.size(); + if (n_predict_budget >= 0 && n_predict_budget != std::numeric_limits::max()) { + max_usable_draft = std::min(max_usable_draft, std::max(0, n_predict_budget - 2)); + } + max_usable_draft = std::min(max_usable_draft, std::max(0, n_ctx - n_past - 2)); + max_usable_draft = std::min(max_usable_draft, std::max(0, (int) llama_n_batch(ctx) - 1)); + if ((int) draft.size() > max_usable_draft) { + draft.resize(max_usable_draft); + } - embd.push_back(id); + const int min_usable_draft = params.speculative.get_min_usable_stage_n_min(); + if ((int) draft.size() >= min_usable_draft && (!draft.empty() || n_predict_budget > 1)) { + if (llama_model_has_recurrent(model)) { + if (!common_speculative_before_draft( + spec, + model, + ctx, + ctx_sampling, + sparams, + 0, + n_past, + sampled_before, + (int) draft.size() + 1, + params.speculative.recurrent_ckpt_mode)) { + LOG_TEE("%s: speculative checkpoint setup failed, falling back to one-token decode\n", __func__); + draft.clear(); + } + } - // echo this to console - input_echo = true; + if (!draft.empty()) { + llama_batch verify_batch = llama_batch_init((int) draft.size() + 1, 0, 1); + std::vector verify_indices; + verify_indices.reserve(draft.size() + 1); - // decrement remaining sampling budget - --n_remain; + common_batch_add(verify_batch, sampled_before, n_past, { 0 }, true); + verify_indices.push_back(0); + for (size_t i = 0; i < draft.size(); ++i) { + common_batch_add(verify_batch, draft[i], n_past + 1 + (llama_pos) i, { 0 }, true); + verify_indices.push_back((int) i + 1); + } - LOG("n_remain: %d\n", n_remain); + if (llama_decode(ctx, verify_batch)) { + llama_batch_free(verify_batch); + LOG_TEE("%s : failed to eval speculative batch\n", __func__); + return 1; + } + + std::vector ids; + try { + ids = common_sampler_sample_and_accept_n(ctx_sampling, ctx, verify_indices, draft); + } catch (const std::exception & e) { + llama_batch_free(verify_batch); + LOG_TEE("%s: speculative sampling failed: %s\n", __func__, e.what()); + return 1; + } + + std::vector accepted_output_indices; + if (!ids.empty()) { + accepted_output_indices.assign(verify_indices.begin(), verify_indices.begin() + ids.size()); + } + + common_speculative_commit( + spec, + ctx, + ctx_sampling, + 0, + sampled_before, + ids, + (int) draft.size(), + n_past + 1, + accepted_output_indices); + + llama_batch_free(verify_batch); + + if (!ids.empty()) { + have_speculative_sampled = true; + speculative_sampled = ids.back(); + if (!sampled_before_from_carry) { + emitted.push_back(sampled_before); + } + emitted.insert(emitted.end(), ids.begin(), ids.end()); + embd_is_prompt = false; + emitted_generated = true; + used_speculative = true; + speculative_tokens.push_back(sampled_before); + if (ids.size() > 1) { + speculative_tokens.insert(speculative_tokens.end(), ids.begin(), ids.end() - 1); + } + n_past += (int) ids.size(); + n_remain -= (int) emitted.size(); + LOG("n_remain: %d\n", n_remain); + } + } + } + } + + if (!used_speculative) { + const llama_token id = common_sampler_sample_legacy(ctx_sampling, ctx, ctx_guidance); + common_sampler_accept(ctx_sampling, ctx, id, /* apply_grammar= */ true); + + LOG("last: %s\n", LOG_TOKENS_TOSTR_PRETTY(ctx, ctx_sampling->prev).c_str()); + + embd.push_back(id); + emitted = embd; + embd_is_prompt = false; + emitted_generated = true; + + // echo this to console + input_echo = true; + + // decrement remaining sampling budget + --n_remain; + + LOG("n_remain: %d\n", n_remain); + } else { + input_echo = true; + } } else { // some user input remains from prompt or interaction, forward it to processing LOG("embd_inp.size(): %d, n_consumed: %d\n", (int) embd_inp.size(), n_consumed); @@ -779,25 +1101,92 @@ int main(int argc, char ** argv) { break; } } + + emitted = embd; + embd_is_prompt = true; + emitted_generated = false; + } + + emitted_hit_eog = false; + + if (emitted_generated && !emitted.empty()) { + std::string generated_text = output_ss.str(); + std::vector emitted_visible; + emitted_visible.reserve(emitted.size()); + is_antiprompt = false; + + for (llama_token id : emitted) { + emitted_visible.push_back(id); + generated_text += common_token_to_piece(ctx, id, params.special); + + if (llama_token_is_eog(model, id)) { + emitted_hit_eog = true; + break; + } + + if (!params.antiprompt.empty()) { + for (std::string & antiprompt : params.antiprompt) { + size_t extra_padding = params.interactive ? 0 : 2; + size_t search_start_pos = generated_text.length() > static_cast(antiprompt.length() + extra_padding) + ? generated_text.length() - static_cast(antiprompt.length() + extra_padding) + : 0; + + if (generated_text.find(antiprompt, search_start_pos) != std::string::npos) { + if (params.interactive) { + is_interacting = true; + } + is_antiprompt = true; + break; + } + } + + if (!is_antiprompt) { + for (const std::vector & ids : antiprompt_ids) { + if (ids.size() == 1 && id == ids[0]) { + if (params.interactive) { + is_interacting = true; + } + is_antiprompt = true; + break; + } + } + } + } + + if (is_antiprompt) { + LOG("found antiprompt: %s\n", generated_text.c_str()); + break; + } + } + + emitted.swap(emitted_visible); + + if (emitted_generated && !emitted.empty()) { + n_decoded += (int) emitted.size(); + const int64_t t_current_us = ggml_time_us(); + if (n_decoded == (int) emitted.size()) { + t_start_generation_us = t_current_us; + t_prompt_processing_ms = (t_start_generation_us - t_start_process_prompt_us) / 1e3; + } + t_token_generation_ms = std::max(1, t_current_us - t_start_generation_us) / 1e3; + } + } // display text if (input_echo && display) { - for (auto id : embd) { + for (auto id : emitted) { const std::string token_str = common_token_to_piece(ctx, id, params.special); // Console/Stream Output fprintf(stdout, "%s", token_str.c_str()); // Record Displayed Tokens To Log - // Note: Generated tokens are created one by one hence this check - if (embd.size() > 1) { - // Incoming Requested Tokens - input_tokens.push_back(id); - } else { - // Outgoing Generated Tokens + if (emitted_generated) { output_tokens.push_back(id); output_ss << token_str; + } else { + input_tokens.push_back(id); } fflush(stdout); @@ -812,49 +1201,8 @@ int main(int argc, char ** argv) { // if not currently processing queued inputs; if ((int) embd_inp.size() <= n_consumed) { - // check for reverse prompt in the last n_prev tokens - if (!params.antiprompt.empty()) { - const int n_prev = 32; - const std::string last_output = llama_sampling_prev_str(ctx_sampling, ctx, n_prev); - - is_antiprompt = false; - // Check if each of the reverse prompts appears at the end of the output. - // If we're not running interactively, the reverse prompt might be tokenized with some following characters - // so we'll compensate for that by widening the search window a bit. - for (std::string & antiprompt : params.antiprompt) { - size_t extra_padding = params.interactive ? 0 : 2; - size_t search_start_pos = last_output.length() > static_cast(antiprompt.length() + extra_padding) - ? last_output.length() - static_cast(antiprompt.length() + extra_padding) - : 0; - - if (last_output.find(antiprompt, search_start_pos) != std::string::npos) { - if (params.interactive) { - is_interacting = true; - } - is_antiprompt = true; - break; - } - } - - // check for reverse prompt using special tokens - llama_token last_token = llama_sampling_last(ctx_sampling); - for (std::vector ids : antiprompt_ids) { - if (ids.size() == 1 && last_token == ids[0]) { - if (params.interactive) { - is_interacting = true; - } - is_antiprompt = true; - break; - } - } - - if (is_antiprompt) { - LOG("found antiprompt: %s\n", last_output.c_str()); - } - } - // deal with end of generation tokens in interactive mode - if (!waiting_for_first_input && llama_token_is_eog(model, llama_sampling_last(ctx_sampling))) { + if (!waiting_for_first_input && emitted_hit_eog) { LOG("found an EOG token\n"); if (params.interactive) { @@ -875,8 +1223,11 @@ int main(int argc, char ** argv) { // if current token is not EOG, we add it to current assistant message if (params.conversation && !waiting_for_first_input) { - auto id = llama_sampling_last(ctx_sampling); - assistant_ss << common_token_to_piece(ctx, id, false); + for (llama_token id : emitted) { + if (!llama_token_is_eog(model, id)) { + assistant_ss << common_token_to_piece(ctx, id, false); + } + } } if ((n_past > 0 || waiting_for_first_input) && is_interacting) { @@ -980,7 +1331,7 @@ int main(int argc, char ** argv) { } // end of generation - if (!embd.empty() && llama_token_is_eog(model, embd.back()) && !(params.interactive)) { + if (emitted_hit_eog && !(params.interactive)) { LOG_TEE(" [end of text]\n"); break; } @@ -998,10 +1349,33 @@ int main(int argc, char ** argv) { llama_state_save_file(ctx, path_session.c_str(), session_tokens.data(), session_tokens.size()); } - llama_print_timings(ctx); + if (n_decoded > 0) { + const double t_prompt = n_prompt_tokens_processed > 0 ? t_prompt_processing_ms / n_prompt_tokens_processed : 0.0; + const double n_prompt_second = (t_prompt_processing_ms > 0 && n_prompt_tokens_processed > 0) + ? 1e3 / t_prompt_processing_ms * n_prompt_tokens_processed + : 0.0; + + const double t_gen = t_token_generation_ms / n_decoded; + const double n_gen_second = 1e3 / t_token_generation_ms * n_decoded; + + LOG_TEE("\n"); + LOG_TEE("main: prompt eval time = %10.2f ms / %5d tokens (%8.2f ms per token, %8.2f tokens per second)\n", + t_prompt_processing_ms, n_prompt_tokens_processed, t_prompt, n_prompt_second); + LOG_TEE("main: eval time = %10.2f ms / %5d tokens (%8.2f ms per token, %8.2f tokens per second)\n", + t_token_generation_ms, n_decoded, t_gen, n_gen_second); + LOG_TEE("main: total time = %10.2f ms / %5d tokens\n", + t_prompt_processing_ms + t_token_generation_ms, n_prompt_tokens_processed + n_decoded); + + common_speculative_print_stats(spec, n_gen_second, n_decoded, n_past, ¶ms.speculative); + } else { + llama_print_timings(ctx); + common_speculative_print_stats(spec); + } write_logfile(ctx, params, model, input_tokens, output_ss.str(), output_tokens); if (ctx_guidance) { llama_free(ctx_guidance); } + if (spec) { common_speculative_free(spec); } + params.speculative.clear_dft(); llama_free(ctx); llama_free_model(model); From 4130ecc0289b74b08e2817b3b1851a01247d5b90 Mon Sep 17 00:00:00 2001 From: Kawrakow Date: Mon, 6 Jul 2026 11:15:54 +0200 Subject: [PATCH 4/5] GLM-DSA: improve TG performance even more (#2068) * GLM-DSA: improve TG performance even more * Fix crash when using MTP * Fix the fix --- ggml/src/ggml-backend.cpp | 5 +- src/graphs/build_deepseek2.cpp | 111 +++++++++++++++++++++++++-------- src/llama-build-context.h | 1 + 3 files changed, 91 insertions(+), 26 deletions(-) diff --git a/ggml/src/ggml-backend.cpp b/ggml/src/ggml-backend.cpp index 922c0f370..1135a2de2 100644 --- a/ggml/src/ggml-backend.cpp +++ b/ggml/src/ggml-backend.cpp @@ -2039,7 +2039,10 @@ static void ggml_backend_sched_copy_inputs(ggml_backend_sched_t sched, ggml_back ggml_backend_tensor_get_async(ids_backend, ids_tensor, ids.data(), 0, ggml_nbytes(ids_tensor)); ggml_backend_synchronize(ids_backend); - needs_sync[tensor_backend_id(ids_tensor)] = k_set_sync; + if (auto id = tensor_backend_id(ids_tensor); id >= 0 && id < GGML_SCHED_MAX_BACKENDS) { + needs_sync[id] = k_set_sync; + } + //needs_sync[tensor_backend_id(ids_tensor)] = k_set_sync; unique_ids.resize((n_expert + 31)/32); std::memset(unique_ids.data(), 0, unique_ids.size()*sizeof(uint32_t)); diff --git a/src/graphs/build_deepseek2.cpp b/src/graphs/build_deepseek2.cpp index 63ed75755..102304c8d 100644 --- a/src/graphs/build_deepseek2.cpp +++ b/src/graphs/build_deepseek2.cpp @@ -323,6 +323,29 @@ ggml_tensor * llm_build_context::build_deepseek2_tp_attention( return combined; } +static inline int dsa_n_top_k(const llama_context & lctx) { + int n_top_k = lctx.model.hparams.indexer_top_k; + if (lctx.cparams.dsa_top_k >= 0) n_top_k = lctx.cparams.dsa_top_k; + return n_top_k; +} +static inline bool dsa_should_use(const llama_context & lctx) { + return lctx.cparams.mtp_op_type == MTP_OP_NONE && lctx.cparams.dsa && lctx.model.arch == LLM_ARCH_GLM_DSA && !lctx.kv_self.kr_l.empty(); +} +static inline bool dsa_should_use(const llama_context & lctx, int il) { + return dsa_should_use(lctx) && lctx.model.layers[il].indexer_attn_q_b && lctx.kv_self.kr_l.size() > (size_t) il && lctx.kv_self.kr_l[il]; +} +static inline bool dsa_can_use_fast_path(const llama_context & lctx, int n_tokens, int n_kv) { + auto n_top_k = dsa_n_top_k(lctx); + if (n_tokens == 1 && n_top_k < n_kv && n_top_k%256 == 0 && lctx.cparams.flash_attn && (lctx.cparams.mla_attn == 1 || lctx.cparams.mla_attn == 3)) { + // We need this check because we need to pretend that the KV cache is F32 so we can use ggml_get_rows to extract + // the rows selected by the indexer. In order this to work, the KV cache row size must be a multiple of sizeof(float). + auto k = lctx.kv_self.k_l[0]; + auto row_size = ggml_row_size(k->type, k->ne[0]); + return row_size % sizeof(float) == 0; + } + return false; +} + // DSA lightning indexer (GLM-5.2 / DeepSeek-V3.2). CACHE-BACKED: the batch's indexer keys are // (Hadamard-rotated and) written to a persistent per-layer indexer-key cache (kv_self.kr_l[il]) at // kv_head, then the FULL [head_size, n_kv] cached key set is read back and scored against the current @@ -722,6 +745,9 @@ ggml_tensor * llm_build_context::build_deepseek2_layer_attention( ggml_tensor * sparse_mask = KQ_mask; ggml_tensor * sparse_mask_fa = KQ_mask; + bool use_dsa = dsa_should_use(lctx, il); + bool dsa_fast_path = use_dsa && dsa_can_use_fast_path(lctx, n_tokens, n_kv); + // self_attention { ggml_tensor * q = nullptr; @@ -775,27 +801,30 @@ ggml_tensor * llm_build_context::build_deepseek2_layer_attention( // for prefill AND decode (single sequence). Gate: --dsa opt-in (off by default) + // GLM_DSA arch + indexer tensors + cache. When off, the model runs the dense MLA path, // byte-identical to a build without this feature. - if (lctx.cparams.dsa && model.arch == LLM_ARCH_GLM_DSA && model.layers[il].indexer_attn_q_b - && kv_self.kr_l.size() > (size_t) il && kv_self.kr_l[il]) { + if (use_dsa) { // GLM-5.2 IndexShare: "full" layers compute their own lightning-indexer top-k; // "shared" layers reuse the previous full layer's top-k (transformers reference: // shared layer indexer=None, topk_indices=prev_topk_indices). The full/shared map is // hparams.indexer_is_full (GGUF metadata or derived config rule). At a given step all // layers share the same n_kv/n_tokens, so a full layer's argsort is valid to reuse. if (hparams.indexer_is_full[il]) { - auto sorted = build_deepseek2_dsa_indexer(gf, il, q, cur, KQ_mask, inp_pos); - if (lctx.cparams.flash_attn) { - last_sparse_mask_fa = ::build_deepseek2_dsa_fa_mask(lctx, ctx0, KQ_mask, sorted); - } else { - last_sparse_mask = build_deepseek2_dsa_sparse_mask(sorted, KQ_mask); + dsa_last_full_sorted = build_deepseek2_dsa_indexer(gf, il, q, cur, KQ_mask, inp_pos); + if (!dsa_fast_path) { + if (lctx.cparams.flash_attn) { + last_sparse_mask_fa = ::build_deepseek2_dsa_fa_mask(lctx, ctx0, KQ_mask, dsa_last_full_sorted); + } else { + last_sparse_mask = build_deepseek2_dsa_sparse_mask(dsa_last_full_sorted, KQ_mask); + } } } - if (lctx.cparams.flash_attn) { - GGML_ASSERT(last_sparse_mask_fa); - sparse_mask_fa = last_sparse_mask_fa; - } else { - GGML_ASSERT(last_sparse_mask); - sparse_mask = last_sparse_mask; + if (!dsa_fast_path) { + if (lctx.cparams.flash_attn) { + GGML_ASSERT(last_sparse_mask_fa); + sparse_mask_fa = last_sparse_mask_fa; + } else { + GGML_ASSERT(last_sparse_mask); + sparse_mask = last_sparse_mask; + } } } @@ -1009,13 +1038,27 @@ ggml_tensor * llm_build_context::build_deepseek2_layer_attention( cb(q, "q", il); if (lctx.cparams.flash_attn && (lctx.cparams.mla_attn == 1 || lctx.cparams.mla_attn == 3)) { - ggml_tensor * kv_cache_lora = ggml_view_2d(ctx0, kv_self.k_l[il], - kv_lora_rank, n_kv, - ggml_row_size(kv_self.k_l[il]->type, kv_lora_rank + n_embd_head_qk_rope), - ggml_row_size(kv_self.k_l[il]->type, n_embd_head_qk_rope)); - cb(kv_cache_lora, "kv_cache_lora", il); - kqv_compressed = ggml_flash_attn_ext(ctx0, q, kv_cache, kv_cache_lora, sparse_mask_fa, kq_scale, hparams.f_max_alibi_bias, 0.f); + if (dsa_fast_path) { + auto row_size = ggml_row_size(kv_self.k_l[il]->type, kv_self.k_l[il]->ne[0]); + GGML_ASSERT(row_size % sizeof(float) == 0); + GGML_ASSERT(dsa_last_full_sorted); + GGML_ASSERT(ggml_nrows(dsa_last_full_sorted) == 1); + auto k32 = ggml_reshape_4d_ext(ctx0, kv_self.k_l[il], GGML_TYPE_F32, row_size/sizeof(float), + kv_self.k_l[il]->ne[1], kv_self.k_l[il]->ne[2], kv_self.k_l[il]->ne[3]); + k32 = ggml_get_rows(ctx0, k32, dsa_last_full_sorted); + auto k = ggml_reshape_4d_ext(ctx0, k32, kv_self.k_l[il]->type, kv_self.k_l[il]->ne[0], k32->ne[1], k32->ne[2], k32->ne[3]); + auto v = ggml_view_2d(ctx0, k, kv_lora_rank, k->ne[1], k->nb[1], ggml_row_size(k->type, n_embd_head_qk_rope)); + kqv_compressed = ggml_flash_attn_ext(ctx0, q, k, v, dsa_tg_fast_mask, kq_scale, hparams.f_max_alibi_bias, 0.f); + } else { + ggml_tensor * kv_cache_lora = ggml_view_2d(ctx0, kv_self.k_l[il], + kv_lora_rank, n_kv, + ggml_row_size(kv_self.k_l[il]->type, kv_lora_rank + n_embd_head_qk_rope), + ggml_row_size(kv_self.k_l[il]->type, n_embd_head_qk_rope)); + cb(kv_cache_lora, "kv_cache_lora", il); + + kqv_compressed = ggml_flash_attn_ext(ctx0, q, kv_cache, kv_cache_lora, sparse_mask_fa, kq_scale, hparams.f_max_alibi_bias, 0.f); + } cb(kqv_compressed, "kqv_compressed", il); if (use_f32_attn_precision) { @@ -1202,7 +1245,7 @@ ggml_cgraph * llm_build_context::build_deepseek2() { // KQ_mask (mask for 1 head, it will be broadcasted to all heads) struct ggml_tensor * KQ_mask = build_inp_KQ_mask(); - if (lctx.cparams.dsa && model.arch == LLM_ARCH_GLM_DSA) { + if (dsa_should_use(lctx)) { static const int n_sink = []{ const char * e = getenv("DSA_SINK"); return e ? atoi(e) : 1; }(); if (n_sink > 0 && n_sink < (int) n_kv) { lctx.inp_dsa_sink = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_kv, n_tokens); @@ -1210,13 +1253,31 @@ ggml_cgraph * llm_build_context::build_deepseek2() { ggml_set_input(lctx.inp_dsa_sink); ggml_build_forward_expand(gf, lctx.inp_dsa_sink); } - auto minus_inf = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, KQ_mask->ne[0], n_tokens); - minus_inf = ggml_fill_inplace(ctx0, minus_inf, -INFINITY); - ggml_build_forward_expand(gf, minus_inf); - lctx.inp_mask_inf = minus_inf; + if (dsa_can_use_fast_path(lctx, n_tokens, n_kv)) { + // For DSA token generation (i.e., n_tokens = 1), when we have more than n_top_k entries in the KV cache, + // we can simply pick the selected n_top_k entries from the cache via ggml_get_rows, so the mask we + // require is all ON (ideally we shouldn't require a mask, but I think we have places in the FA + // implementation which assume that there is a mask). + int n_top_k = dsa_n_top_k(lctx); + int n_padded = GGML_PAD(n_top_k, 256); + auto minus_inf = ggml_new_tensor_1d(ctx0, GGML_TYPE_F32, n_padded*KQ_mask->ne[1] - n_top_k); + minus_inf = ggml_fill_inplace(ctx0, minus_inf, -INFINITY); + auto zero = ggml_new_tensor_1d(ctx0, GGML_TYPE_F32, n_top_k); + zero = ggml_fill_inplace(ctx0, zero, 0.0f); + auto mask = ggml_concat(ctx0, zero, minus_inf, 0); + mask = ggml_cast(ctx0, mask, GGML_TYPE_F16); + mask = ggml_reshape_2d(ctx0, mask, n_padded, KQ_mask->ne[1]); + ggml_build_forward_expand(gf, mask); + dsa_tg_fast_mask = mask; + } else { + auto minus_inf = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, KQ_mask->ne[0], n_tokens); + minus_inf = ggml_fill_inplace(ctx0, minus_inf, -INFINITY); + ggml_build_forward_expand(gf, minus_inf); + lctx.inp_mask_inf = minus_inf; + } } - last_sparse_mask = last_sparse_mask_fa = nullptr; + last_sparse_mask = last_sparse_mask_fa = dsa_last_full_sorted = nullptr; // whether to use n_tokens as the matrix dimension during multiplication or n_head // n_tokens is higher during prompt processing, this allows to optimize for this case diff --git a/src/llama-build-context.h b/src/llama-build-context.h index b168cc970..32a667392 100644 --- a/src/llama-build-context.h +++ b/src/llama-build-context.h @@ -102,6 +102,7 @@ struct llm_build_context { struct ggml_tensor * dsa_last_full_sorted = nullptr; struct ggml_tensor * last_sparse_mask_fa = nullptr; struct ggml_tensor * last_sparse_mask = nullptr; + struct ggml_tensor * dsa_tg_fast_mask = nullptr; // TODO: consider making the entire interface noexcept llm_build_context( From eacd450d2b57894c8c8c3e0285e755cb93e8a036 Mon Sep 17 00:00:00 2001 From: Kawrakow Date: Mon, 6 Jul 2026 11:57:40 +0200 Subject: [PATCH 5/5] GLM-DSA: much better PP performance (CPU-only) (#2074) * GLM-DSA: much better PP performance (CPU-only) * Cleanup --- ggml/src/ggml.c | 3 +- ggml/src/iqk/iqk_flash_attn.cpp | 56 +++++++++++++++++++++++++++++++-- ggml/src/iqk/iqk_mul_mat.h | 2 +- src/graphs/build_deepseek2.cpp | 6 ++++ 4 files changed, 63 insertions(+), 4 deletions(-) diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c index 53aef9869..4e78b2a6b 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -21919,7 +21919,8 @@ static void ggml_compute_forward_flash_attn_ext_f16( Dk, Dv, neq1, nek1, q->nb[1], k->nb[1], v->nb[1], mask ? mask->nb[1] : 0, q->data, k->data, v->data, mask ? mask->data : NULL, sinks ? sinks->data : NULL, scale, softcap, (float *)dst->data, - params->wdata, (barrier_t)ggml_barrier, (void *)params->shared, ith, nth, dst->op_params[4])) return; + params->wdata, (barrier_t)ggml_barrier, (void *)params->shared, ith, nth, dst->op_params[4], + dst->src[5])) return; // if (max_bias <= 0.0f && q->type == GGML_TYPE_F32 && mask && mask->type == GGML_TYPE_F16) { // //if (ith == 0) printf("k: %ld x %ld x %ld, q: %ld x %ld x %ld, v: %ld x %ld x %ld mask: %ld x %ld x %ld\n", diff --git a/ggml/src/iqk/iqk_flash_attn.cpp b/ggml/src/iqk/iqk_flash_attn.cpp index c562b0c06..e1ea5177d 100644 --- a/ggml/src/iqk/iqk_flash_attn.cpp +++ b/ggml/src/iqk/iqk_flash_attn.cpp @@ -51,6 +51,14 @@ size_t iqk_fa_work_buffer_size(const struct ggml_tensor * dst, int nth) { auto Q = dst->src[0]; auto K = dst->src[1]; auto V = dst->src[2]; + auto indexer = dst->src[5]; + if (indexer && indexer->type == GGML_TYPE_I32 && indexer->ne[0] < K->ne[1] && Q->ne[1] >= nth && + Q->ne[3] == 1 && K->ne[3] == 1 && V->ne[3] == 1 && K->ne[2] == 1) { + auto row_size_k = ggml_row_size(K->type, K->ne[0]); + auto row_size_v = ggml_row_size(V->type, V->ne[0]); + auto work_size = (row_size_k + row_size_v + 64) * indexer->ne[0]; + return work_size * nth; + } int rk2 = Q->ne[2]/K->ne[2]; size_t size = 0; if (Q->ne[1] >= 8 && K->type == GGML_TYPE_Q8_0) { @@ -163,10 +171,54 @@ extern "C" IQK_API bool iqk_flash_attn_noalibi(int type_q, int type_mask, float float softcap, // if > 0, a "soft-cap" operation is applied before softmax float * qkv, // v*softmax(scale*(k*q)) [[maybe_unused]] void * work_buffer_in, [[maybe_unused]] barrier_t barrier, [[maybe_unused]] void * barrier_data, - int ith, int nth, int n_swa) { + int ith, int nth, int n_swa, [[maybe_unused]] ggml_tensor * indexer) { if (type_q != 0 || type_mask != 1 || max_bias > 0) return false; + if (indexer && indexer->type == GGML_TYPE_I32) { + if (indexer->ne[0] < nek1 && neq1 >= nth && neq3 == 1 && nek3 == 1 && nev3 == 1 && nek2 == 1) { + int npt = (neq1 + nth - 1)/nth; + int ith_mid = nth; + int neq1_this_thread = npt; + int first = ith*npt; + if (npt*nth > neq1) { + ith_mid = neq1 - nth*(npt - 1); + if (ith >= ith_mid) { + --neq1_this_thread; + first = ith_mid*npt + (ith - ith_mid)*neq1_this_thread; + } + } + // Workbuffer: we need + // * indexer->ne[0] * sizeof(ggml_half) to extract the mask for a row + // * indexer->ne[0] * ggml_row_size(int_type_k_in, Dk) to extract the selected K cache entries + // * indexer->ne[0] * ggml_row_size(int_type_v, Dv) to extract the selected V cache entries + auto row_size_k = ggml_row_size(ggml_type(int_type_k_in), Dk); + auto row_size_v = ggml_row_size(ggml_type(int_type_v ), Dv); + auto work_size = (row_size_k + row_size_v + 64) * indexer->ne[0]; + auto work_k = (char *)work_buffer_in + ith*work_size; + auto work_v = work_k + row_size_k*indexer->ne[0]; + auto work_m = (uint16_t *)(work_v + row_size_v*indexer->ne[0]); + int nkv = indexer->ne[0]; + for (int iq = first; iq < first + neq1_this_thread; ++iq) { + auto idx = (const int *)((const char *)indexer->data + iq*indexer->nb[1]); + auto M = (const uint16_t *)((const char *)mask + iq*stride_m); + for (int j = 0; j < nkv; ++j) { + std::memcpy(work_k + row_size_k*j, ((const char *)k + idx[j]*stride_k), row_size_k); + std::memcpy(work_v + row_size_v*j, ((const char *)v + idx[j]*stride_v), row_size_v); + work_m[j] = M[idx[j]]; + } + auto this_q = (const char *)q + iq*stride_q; + auto this_qkv = qkv + iq*ne1*nb1/sizeof(float); + if (!iqk_flash_attn_impl(int_type_k_in, int_type_v, + Dk, Dv, neq2, nkv, nbq2, row_size_k, row_size_v, 0, Dv, + (const float *)this_q, work_k, work_v, work_m, nullptr, 0, + scale, softcap, + this_qkv, nullptr, nullptr)) return false; + } + return true; + } + } + if (auto type_k = ggml_type(int_type_k_in), type_v = ggml_type(int_type_v); !are_kv_types_supported(type_k, type_v)) { if (ith == 0) { fprintf(stderr, "\n==================== K cache %s coupled with V cache %s is not a supported combination on the CPU backend.\n", @@ -520,7 +572,7 @@ bool iqk_flash_attn_noalibi([[maybe_unused]] int type_q, [[maybe_unused]] int ty [[maybe_unused]] float softcap, // if > 0, a "soft-cap" operation is applied before softmax [[maybe_unused]] float * qkv, // v*softmax(scale*(k*q)) [[maybe_unused]] void * work_buffer, [[maybe_unused]] barrier_t barrier, [[maybe_unused]] void * barrier_data, - [[maybe_unused]] int ith, [[maybe_unused]] int nth, [[maybe_unused]] int n_swa) { + [[maybe_unused]] int ith, [[maybe_unused]] int nth, [[maybe_unused]] int n_swa, [[maybe_unused]] ggml_tensor * indexer) { return false; } diff --git a/ggml/src/iqk/iqk_mul_mat.h b/ggml/src/iqk/iqk_mul_mat.h index 6e372f938..a51597090 100644 --- a/ggml/src/iqk/iqk_mul_mat.h +++ b/ggml/src/iqk/iqk_mul_mat.h @@ -68,7 +68,7 @@ IQK_API bool iqk_flash_attn_noalibi(int type_q, int type_mask, float max_bias, float softcap, // if > 0, a "soft-cap" operation is applied before softmax float * qkv, // v*softmax(scale*(k*q)) void * work_buffer, barrier_t barrier, void * barrier_data, - int ith, int nth, int n_swa); + int ith, int nth, int n_swa, struct ggml_tensor * indexer); IQK_API void iqk_topk_moe(int n_experts, int n_experts_used, int nrows, const float * logits, float * weights, int32_t * ids, int ith, int nth); diff --git a/src/graphs/build_deepseek2.cpp b/src/graphs/build_deepseek2.cpp index 102304c8d..5aad17060 100644 --- a/src/graphs/build_deepseek2.cpp +++ b/src/graphs/build_deepseek2.cpp @@ -1007,6 +1007,9 @@ ggml_tensor * llm_build_context::build_deepseek2_layer_attention( if (use_f32_attn_precision || q->ne[1] <= 8) { ggml_flash_attn_ext_set_prec(kqv, GGML_PREC_F32); } + if (use_dsa && dsa_last_full_sorted) { + kqv->src[5] = dsa_last_full_sorted; + } cb(kqv, "kqv", il); if (iter == 0) { @@ -1064,6 +1067,9 @@ ggml_tensor * llm_build_context::build_deepseek2_layer_attention( if (use_f32_attn_precision) { ggml_flash_attn_ext_set_prec(kqv_compressed, GGML_PREC_F32); } + if (use_dsa && dsa_last_full_sorted) { + kqv_compressed->src[5] = dsa_last_full_sorted; + } kqv_compressed = ggml_permute(ctx0, kqv_compressed, 0, 2, 1, 3); cb(kqv_compressed, "kqv_compressed_perm", il);