From c1a411fb1b5550309ab59785b17672b7cb69a885 Mon Sep 17 00:00:00 2001 From: zhangrunda Date: Tue, 7 Jul 2026 21:23:53 +0800 Subject: [PATCH 01/17] common : add missing include in common.h (#25220) Signed-off-by: zhangrunda --- common/common.h | 1 + 1 file changed, 1 insertion(+) diff --git a/common/common.h b/common/common.h index 2adb310b8..b7d85b5dd 100644 --- a/common/common.h +++ b/common/common.h @@ -14,6 +14,7 @@ #include #include #include +#include #if defined(_WIN32) && !defined(_WIN32_WINNT) #define _WIN32_WINNT 0x0A00 From 6c487e2f79dea747d70325250121e750ed364b2b Mon Sep 17 00:00:00 2001 From: Thiago Padilha Date: Tue, 7 Jul 2026 10:24:35 -0300 Subject: [PATCH 02/17] server: enforce prompt cache RAM limit (#25070) Before this commit, --cache-ram was not a hard limit: - The cache always kept at least one entry, even if that entry exceeded the RAM/token limits. - Old entries were only evicted for the RAM/token limits after saving the new one, which could cause the cache to temporarily exceed the RAM/token limits even if individual entries were below the limit. Now, ensure that the RAM limit is strict with these changes: - Skip saving state to cache if by itself it exceeds the RAM limit. - Evict old entries as necessary to make the new entry fit. Additionally, token-limit cleanup may now evict the last remaining cache entry instead of always preserving one. --- tools/server/server-task.cpp | 40 +++++++++++++++++++++++++----------- 1 file changed, 28 insertions(+), 12 deletions(-) diff --git a/tools/server/server-task.cpp b/tools/server/server-task.cpp index 775f50baf..54fa949df 100644 --- a/tools/server/server-task.cpp +++ b/tools/server/server-task.cpp @@ -1631,7 +1631,22 @@ server_prompt * server_prompt_cache::alloc(const server_prompt & prompt, size_t } } - // next, remove any cached prompts that are fully contained in the current prompt + // calculate checkpoints size to see if it will fit with the prompt + size_t checkpoints_size = 0; + for (const auto & ckpt : prompt.checkpoints) { + checkpoints_size += ckpt.size(); + } + + const size_t state_size_new = state_size_tgt + state_size_dft + checkpoints_size; + + // skip over-limit entries to avoid disturbing the cache + if (limit_size > 0 && state_size_new > limit_size) { + SRV_WRN(" - prompt state size %.3f MiB exceeds cache size limit %.3f MiB, skipping\n", + state_size_new / (1024.0 * 1024.0), limit_size / (1024.0 * 1024.0)); + return nullptr; + } + + // remove any cached prompts that are fully contained in the current prompt for (auto it = states.begin(); it != states.end();) { const int len = it->tokens.get_common_prefix(prompt.tokens); @@ -1644,6 +1659,16 @@ server_prompt * server_prompt_cache::alloc(const server_prompt & prompt, size_t } } + if (limit_size > 0) { + // make room before allocating the new vectors to avoid breaching the limit + while (!states.empty() && size() + state_size_new > limit_size) { + SRV_WRN(" - making room for prompt cache entry, removing oldest entry (size = %.3f MiB)\n", + states.front().size() / (1024.0 * 1024.0)); + + states.pop_front(); + } + } + std::vector state_data_tgt; std::vector state_data_dft; @@ -1752,12 +1777,7 @@ bool server_prompt_cache::load(server_prompt & prompt, const server_tokens & tok void server_prompt_cache::update() { if (limit_size > 0) { - // always keep at least one state, regardless of the limits - while (states.size() > 1 && size() > limit_size) { - if (states.empty()) { - break; - } - + while (!states.empty() && size() > limit_size) { SRV_WRN(" - cache size limit reached, removing oldest entry (size = %.3f MiB)\n", states.front().size() / (1024.0 * 1024.0)); states.pop_front(); @@ -1771,11 +1791,7 @@ void server_prompt_cache::update() { const size_t limit_tokens_cur = limit_size > 0 ? std::max(limit_tokens, limit_size/size_per_token) : limit_tokens; if (limit_tokens > 0) { - while (states.size() > 1 && n_tokens() > limit_tokens_cur) { - if (states.empty()) { - break; - } - + while (!states.empty() && n_tokens() > limit_tokens_cur) { SRV_WRN(" - cache token limit (%zu, est: %zu) reached, removing oldest entry (size = %.3f MiB)\n", limit_tokens, limit_tokens_cur, states.front().size() / (1024.0 * 1024.0)); From 5eca4e3cabad9c92ccbf290d8406fbe7c110a680 Mon Sep 17 00:00:00 2001 From: Thomas LECONTE <161708450+thom-dev-fr@users.noreply.github.com> Date: Tue, 7 Jul 2026 16:13:03 +0200 Subject: [PATCH 03/17] server : add timings and progress to /responses API stream (#25348) --- tools/server/server-task.cpp | 34 ++++++++++++++- tools/server/server-task.h | 2 + .../tests/unit/test_compat_oai_responses.py | 41 +++++++++++++++++++ 3 files changed, 76 insertions(+), 1 deletion(-) diff --git a/tools/server/server-task.cpp b/tools/server/server-task.cpp index 54fa949df..8d611e520 100644 --- a/tools/server/server-task.cpp +++ b/tools/server/server-task.cpp @@ -730,6 +730,10 @@ json server_task_result_cmpl_final::to_json_oaicompat_resp_stream() { }} }); + if (timings.prompt_n >= 0) { + server_sent_events.back().at("data").push_back({"timings", timings.to_json()}); + } + return server_sent_events; } @@ -1016,6 +1020,7 @@ void server_task_result_cmpl_partial::update(task_result_state & state) { thinking_block_started = state.thinking_block_started; text_block_started = state.text_block_started; + oai_resp_created = state.oai_resp_created; oai_resp_id = state.oai_resp_id; oai_resp_reasoning_id = state.oai_resp_reasoning_id; oai_resp_message_id = state.oai_resp_message_id; @@ -1024,6 +1029,10 @@ void server_task_result_cmpl_partial::update(task_result_state & state) { // track if the accumulated message has any reasoning content anthropic_has_reasoning = !state.chat_msg.reasoning_content.empty(); + if (res_type == TASK_RESPONSE_TYPE_OAI_RESP && !state.oai_resp_created && (is_progress || n_decoded == 1)) { + state.oai_resp_created = true; + } + // Pre-compute state updates based on diffs (for next chunk) for (const common_chat_msg_diff & diff : oaicompat_msg_diffs) { if (!diff.reasoning_content_delta.empty() && !state.thinking_block_started) { @@ -1181,7 +1190,7 @@ json server_task_result_cmpl_partial::to_json_oaicompat_chat() { json server_task_result_cmpl_partial::to_json_oaicompat_resp() { std::vector events; - if (n_decoded == 1) { + if (!oai_resp_created) { events.push_back(json { {"event", "response.created"}, {"data", json { @@ -1204,6 +1213,18 @@ json server_task_result_cmpl_partial::to_json_oaicompat_resp() { }}, }}, }); + } else if (is_progress) { + events.push_back(json { + {"event", "response.in_progress"}, + {"data", json { + {"type", "response.in_progress"}, + {"response", json { + {"id", oai_resp_id}, + {"object", "response"}, + {"status", "in_progress"}, + }}, + }}, + }); } for (const common_chat_msg_diff & diff : oaicompat_msg_diffs) { @@ -1302,6 +1323,17 @@ json server_task_result_cmpl_partial::to_json_oaicompat_resp() { }); } } + + if (!events.empty()) { + json & data = events.back().at("data"); + if (timings.prompt_n >= 0) { + data.push_back({"timings", timings.to_json()}); + } + if (is_progress) { + data.push_back({"prompt_progress", progress.to_json()}); + } + } + return events; } diff --git a/tools/server/server-task.h b/tools/server/server-task.h index 49f62d386..dc6b2dac1 100644 --- a/tools/server/server-task.h +++ b/tools/server/server-task.h @@ -117,6 +117,7 @@ struct task_result_state { bool text_block_started = false; // for OpenAI Responses streaming API + bool oai_resp_created = false; const std::string oai_resp_id; const std::string oai_resp_reasoning_id; const std::string oai_resp_message_id; @@ -440,6 +441,7 @@ struct server_task_result_cmpl_partial : server_task_result { bool text_block_started = false; // for OpenAI Responses API + bool oai_resp_created = false; std::string oai_resp_id; std::string oai_resp_reasoning_id; std::string oai_resp_message_id; diff --git a/tools/server/tests/unit/test_compat_oai_responses.py b/tools/server/tests/unit/test_compat_oai_responses.py index 7aab4a8ba..14528b487 100644 --- a/tools/server/tests/unit/test_compat_oai_responses.py +++ b/tools/server/tests/unit/test_compat_oai_responses.py @@ -71,3 +71,44 @@ def test_responses_stream_with_openai_library(): assert r.response.output[0].id.startswith("msg_") assert gathered_text == r.response.output_text assert match_regex("(Suddenly)+", r.response.output_text) + + +def test_responses_stream_with_llama_telemetry(): + global server + server.n_ctx = 256 + server.n_batch = 32 + server.n_slots = 1 + server.start() + + saw_progress = False + saw_delta_timings = False + completed = None + + res = server.make_stream_request("POST", "/responses", data={ + "input": "This is a test" * 10, + "max_output_tokens": 8, + "temperature": 0.8, + "stream": True, + "timings_per_token": True, + "return_progress": True, + }) + + for data in res: + if "prompt_progress" in data: + assert data["type"] == "response.in_progress" + assert data["prompt_progress"]["total"] > 0 + assert data["prompt_progress"]["processed"] >= data["prompt_progress"]["cache"] + saw_progress = True + if "timings" in data: + assert "prompt_per_second" in data["timings"] + assert "predicted_per_second" in data["timings"] + if data["type"] == "response.output_text.delta": + saw_delta_timings = True + if data["type"] == "response.completed": + completed = data + + assert saw_progress + assert saw_delta_timings + assert completed is not None + assert "usage" in completed["response"] + assert "timings" in completed From f5525f7e7a7e7cbecd386144299493ea40499bd3 Mon Sep 17 00:00:00 2001 From: Alex <18387287+wadealexc@users.noreply.github.com> Date: Tue, 7 Jul 2026 10:20:42 -0400 Subject: [PATCH 04/17] server : fix draft model fit vs load inconsistency (#25056) * fix: draft model fit vs load inconsistency * refactor(server): unify draft/mtp parameter initialization, model, and context load - moves speculative init to speculative.cpp - changes server_context_impl model_dft and ctx_dft to use raw pointers - fix: don't throttle progress callback when loading draft model - refactor: rename draft model/ctx load method * fix: valign --- common/speculative.cpp | 107 +++++++++++++++++++++ common/speculative.h | 18 ++++ tools/server/server-context.cpp | 158 +++++++++++--------------------- 3 files changed, 176 insertions(+), 107 deletions(-) diff --git a/common/speculative.cpp b/common/speculative.cpp index 5b26597f6..1d2977e42 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -2221,6 +2221,113 @@ int32_t common_speculative_n_max(const common_params_speculative * spec) { return n_max; } +common_params common_base_params_to_speculative(const common_params & params) { + const bool has_draft = params.speculative.has_dft(); + + const auto & params_spec = params.speculative.draft; + common_params result = params; + + if (has_draft) { + result.devices = params_spec.devices; + result.model = params_spec.mparams; + result.n_gpu_layers = params_spec.n_gpu_layers; + result.tensor_buft_overrides = params_spec.tensor_buft_overrides; + + if (params_spec.cpuparams.n_threads > 0) { + result.cpuparams.n_threads = params_spec.cpuparams.n_threads; + result.cpuparams_batch.n_threads = params_spec.cpuparams_batch.n_threads; + } + } + + result.cache_type_k = params_spec.cache_type_k; + result.cache_type_v = params_spec.cache_type_v; + result.n_outputs_max = params.n_parallel; + + return result; +} + +struct common_init_speculative_result::impl { + impl() = default; + ~impl() = default; + + // note: the order in which model, context, etc. are declared matters because their destructors will be called bottom-to-top + llama_model_ptr model; + llama_context_ptr context; +}; + +common_init_speculative_result::common_init_speculative_result( + common_params & params, + llama_model * model_tgt, + llama_context * ctx_tgt) : + pimpl(new impl{}) { + const bool has_draft = params.speculative.has_dft(); + const bool spec_mtp = std::find(params.speculative.types.begin(), + params.speculative.types.end(), + COMMON_SPECULATIVE_TYPE_DRAFT_MTP) != params.speculative.types.end(); + GGML_ASSERT(has_draft || spec_mtp); + + auto mparams = common_model_params_to_llama(params); + auto cparams = common_context_params_to_llama(params); + + if (spec_mtp) { + cparams.ctx_type = LLAMA_CONTEXT_TYPE_MTP; + } + + // note: for small models maybe we can set this to the maximum possible draft from all speculative types + // the extra memory for small models is likely negligible? + cparams.n_rs_seq = 0; + cparams.ctx_other = ctx_tgt; + + std::string model_path; + if (has_draft) { + model_path = params.speculative.draft.mparams.path; + LOG_TRC("%s: loading draft model '%s'\n", __func__, model_path.c_str()); + + llama_model * model_dft = llama_model_load_from_file(params.model.path.c_str(), mparams); + if (model_dft == NULL) { + LOG_ERR("%s: failed to load draft model, '%s'\n", __func__, model_path.c_str()); + return; + } + + pimpl->model.reset(model_dft); + + llama_context * ctx_dft = llama_init_from_model(model_dft, cparams); + if (ctx_dft == nullptr) { + LOG_ERR("%s: failed to create MTP context\n", __func__); + return; + } + + pimpl->context.reset(ctx_dft); + } else if (spec_mtp) { + model_path = params.model.path; + + LOG_TRC("%s: creating MTP draft context against the target model '%s'\n", + __func__, model_path.c_str()); + + llama_context * ctx_dft = llama_init_from_model(model_tgt, cparams); + if (ctx_dft == nullptr) { + LOG_ERR("%s: failed to create MTP context\n", __func__); + return; + } + + pimpl->context.reset(ctx_dft); + } +} + +common_init_speculative_result::~common_init_speculative_result() = default; + +llama_model * common_init_speculative_result::model() { + return pimpl->model.get(); +} + +llama_context * common_init_speculative_result::context() { + return pimpl->context.get(); +} + +common_init_speculative_result_ptr common_init_speculative_from_params(common_params & params, llama_model * model_tgt, llama_context * ctx_tgt) { + return std::make_unique(params, model_tgt, ctx_tgt); +} + // initialization of the speculative decoding system // common_speculative * common_speculative_init(common_params_speculative & params, uint32_t n_seq) { diff --git a/common/speculative.h b/common/speculative.h index c58fac3cc..2f6d225b6 100644 --- a/common/speculative.h +++ b/common/speculative.h @@ -23,6 +23,8 @@ std::string common_speculative_type_to_str(enum common_speculative_type type); // return the max number of draft tokens based on the speculative parameters int32_t common_speculative_n_max(const common_params_speculative * spec); +common_params common_base_params_to_speculative(const common_params & params); + common_speculative * common_speculative_init(common_params_speculative & params, uint32_t n_seq); void common_speculative_free(common_speculative * spec); @@ -80,3 +82,19 @@ struct common_speculative_deleter { }; typedef std::unique_ptr common_speculative_ptr; + +struct common_init_speculative_result { + common_init_speculative_result(common_params & params, llama_model * model_tgt, llama_context * ctx_tgt); + ~common_init_speculative_result(); + + llama_model * model(); + llama_context * context(); + +private: + struct impl; + std::unique_ptr pimpl; +}; + +using common_init_speculative_result_ptr = std::unique_ptr; + +common_init_speculative_result_ptr common_init_speculative_from_params(common_params & params, llama_model * model_tgt, llama_context * ctx_tgt); diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index bb3b91ab5..3890fe5f0 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -897,8 +897,10 @@ private: server_batch batch; - llama_model_ptr model_dft; - llama_context_ptr ctx_dft; + llama_model * model_dft = nullptr; + llama_context * ctx_dft = nullptr; + + common_init_speculative_result_ptr spec_init; common_context_seq_rm_type ctx_tgt_seq_rm_type = COMMON_CONTEXT_SEQ_RM_TYPE_NO; common_context_seq_rm_type ctx_dft_seq_rm_type = COMMON_CONTEXT_SEQ_RM_TYPE_NO; @@ -939,8 +941,10 @@ private: void destroy() { spec.reset(); - ctx_dft.reset(); - model_dft.reset(); + spec_init.reset(); + + ctx_dft = nullptr; + model_dft = nullptr; llama_init.reset(); @@ -1084,30 +1088,15 @@ private: // optionally reserve VRAM for the draft / MTP context before fitting the target model if (params_base.fit_params) { if (has_spec) { - common_params params_dft = params_base; - bool measure_model_bytes = true; + // MTP draft context lives on the target model, only context+compute are new + bool measure_model_bytes = has_draft; - if (has_draft) { - const auto & params_spec = params_base.speculative.draft; - params_dft.devices = params_spec.devices; - params_dft.model = params_spec.mparams; - params_dft.n_gpu_layers = params_spec.n_gpu_layers; - params_dft.cache_type_k = params_spec.cache_type_k; - params_dft.cache_type_v = params_spec.cache_type_v; - params_dft.tensor_buft_overrides = params_spec.tensor_buft_overrides; - } else { - // MTP draft context lives on the target model, only context+compute are new - measure_model_bytes = false; - } - - params_dft.n_outputs_max = params_base.n_parallel; + common_params params_dft = common_base_params_to_speculative(params_base); auto mparams_dft = common_model_params_to_llama(params_dft); auto cparams_dft = common_context_params_to_llama(params_dft); if (spec_mtp) { cparams_dft.ctx_type = LLAMA_CONTEXT_TYPE_MTP; - cparams_dft.type_k = params_base.speculative.draft.cache_type_k; - cparams_dft.type_v = params_base.speculative.draft.cache_type_v; } cparams_dft.n_rs_seq = 0; @@ -1175,82 +1164,35 @@ private: add_bos_token = llama_vocab_get_add_bos(vocab); - if (has_draft) { - // TODO speculative: move to common/speculative.cpp? - const auto & params_spec = params_base.speculative.draft; - - SRV_TRC("loading draft model '%s'\n", params_spec.mparams.path.c_str()); - - auto params_dft = params_base; - - params_dft.devices = params_spec.devices; - params_dft.model = params_spec.mparams; - params_dft.n_gpu_layers = params_spec.n_gpu_layers; - params_dft.cache_type_k = params_spec.cache_type_k; - params_dft.cache_type_v = params_spec.cache_type_v; - - if (params_spec.cpuparams.n_threads > 0) { - params_dft.cpuparams.n_threads = params_spec.cpuparams.n_threads; - params_dft.cpuparams_batch.n_threads = params_spec.cpuparams_batch.n_threads; - } - - params_dft.tensor_buft_overrides = params_spec.tensor_buft_overrides; - - auto mparams_dft = common_model_params_to_llama(params_dft); - - // progress callback - mparams_dft.progress_callback = load_progress_callback; - mparams_dft.progress_callback_user_data = &load_progress_spec; - - model_dft.reset(llama_model_load_from_file(params_dft.model.path.c_str(), mparams_dft)); - if (model_dft == nullptr) { - SRV_ERR("failed to load draft model, '%s'\n", params_dft.model.path.c_str()); - return false; - } - - auto cparams = common_context_params_to_llama(params_dft); - - if (spec_mtp) { - cparams.ctx_type = LLAMA_CONTEXT_TYPE_MTP; - } - - // note: for small models maybe we can set this to the maximum possible draft from all speculative types - // the extra memory for small models is likely negligible? - cparams.n_rs_seq = 0; - cparams.ctx_other = ctx_tgt; - - ctx_dft.reset(llama_init_from_model(model_dft.get(), cparams)); - if (ctx_dft == nullptr) { - SRV_ERR("%s", "failed to create draft context\n"); - return false; - } - - params_base.speculative.draft.ctx_tgt = ctx_tgt; - params_base.speculative.draft.ctx_dft = ctx_dft.get(); - } else if (spec_mtp) { - // no new model load, so we simply report 0.0 and 1.0 progress + if (has_spec) { + // spec_mtp doesn't use load a model internally, so we report 0.0 and 1.0 manually load_progress_callback(0.0f, &load_progress_spec); + load_progress_spec.t_last_load_progress_ms = 0; // reset so internal cbs aren't delayed - SRV_TRC("creating MTP draft context against the target model '%s'\n", - params_base.model.path.c_str()); + { + common_params params_dft = common_base_params_to_speculative(params_base); + // progress callback + params_dft.load_progress_callback = load_progress_callback; + params_dft.load_progress_callback_user_data = &load_progress_spec; - auto cparams_mtp = common_context_params_to_llama(params_base); - cparams_mtp.ctx_type = LLAMA_CONTEXT_TYPE_MTP; - cparams_mtp.type_k = params_base.speculative.draft.cache_type_k; - cparams_mtp.type_v = params_base.speculative.draft.cache_type_v; - cparams_mtp.n_rs_seq = 0; - cparams_mtp.n_outputs_max = params_base.n_parallel; - cparams_mtp.ctx_other = ctx_tgt; + spec_init = common_init_speculative_from_params(params_dft, model_tgt, ctx_tgt); + model_dft = spec_init->model(); + ctx_dft = spec_init->context(); - ctx_dft.reset(llama_init_from_model(model_tgt, cparams_mtp)); - if (ctx_dft == nullptr) { - SRV_ERR("%s", "failed to create MTP context\n"); - return false; + if (has_draft && model_dft == nullptr) { + SRV_ERR("failed to load draft model, '%s'\n", params_dft.model.path.c_str()); + return false; + } + + if (ctx_dft == nullptr) { + SRV_ERR("%s", "failed to create MTP context\n"); + return false; + } + + params_base.speculative.draft.ctx_tgt = ctx_tgt; + params_base.speculative.draft.ctx_dft = ctx_dft; } - params_base.speculative.draft.ctx_tgt = ctx_tgt; - params_base.speculative.draft.ctx_dft = ctx_dft.get(); - load_progress_callback(1.0f, &load_progress_spec); } @@ -1343,13 +1285,15 @@ private: } if (ctx_dft) { - ctx_dft_seq_rm_type = common_context_can_seq_rm(ctx_dft.get()); + ctx_dft_seq_rm_type = common_context_can_seq_rm(ctx_dft); } if (spec) { SRV_TRC("%s", "speculative decoding context initialized\n"); } else { - ctx_dft.reset(); + spec_init.reset(); + ctx_dft = nullptr; + model_dft = nullptr; } for (int i = 0; i < params_base.n_parallel; i++) { @@ -1357,7 +1301,7 @@ private: slot.id = i; slot.ctx_tgt = ctx_tgt; - slot.ctx_dft = ctx_dft.get(); + slot.ctx_dft = ctx_dft; slot.spec = spec.get(); slot.n_ctx = n_ctx_slot; @@ -2363,7 +2307,7 @@ private: cur.update_pos(slot.prompt.n_tokens() - n_tokens_cur, pos_min, pos_max); cur.update_tgt(ctx_tgt, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); - cur.update_dft(ctx_dft.get(), slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); + cur.update_dft(ctx_dft, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); // stash the draft's speculative state with the checkpoint common_speculative_get_state(spec.get(), slot.id, cur.data_spec); @@ -2899,8 +2843,8 @@ private: common_context_seq_add(ctx_tgt, slot.id, n_keep + n_discard, slot.prompt.n_tokens(), -n_discard); if (ctx_dft) { - common_context_seq_rm (ctx_dft.get(), slot.id, n_keep , n_keep + n_discard); - common_context_seq_add(ctx_dft.get(), slot.id, n_keep + n_discard, slot.prompt.tokens.pos_next(), -n_discard); + common_context_seq_rm (ctx_dft, slot.id, n_keep , n_keep + n_discard); + common_context_seq_add(ctx_dft, slot.id, n_keep + n_discard, slot.prompt.tokens.pos_next(), -n_discard); } // add generated tokens to cache @@ -2972,7 +2916,7 @@ private: llama_memory_seq_pos_max(llama_get_memory(ctx_tgt), slot.id)); if (use_ckpt_dft) { - slot.spec_ckpt.update_dft(ctx_dft.get(), slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); + slot.spec_ckpt.update_dft(ctx_dft, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); } slot.spec_prompt = slot.prompt.tokens.get_text_tokens(); @@ -3009,10 +2953,10 @@ private: if (ctx_dft) { if (use_ckpt_dft) { - ckpt.load_dft(ctx_dft.get(), slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); + ckpt.load_dft(ctx_dft, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); } - common_context_seq_rm(ctx_dft.get(), slot.id, ckpt.pos_max + 1, -1); + common_context_seq_rm(ctx_dft, slot.id, ckpt.pos_max + 1, -1); } if (!draft.empty()) { @@ -3021,7 +2965,7 @@ private: (ctx_tgt_seq_rm_type == COMMON_CONTEXT_SEQ_RM_TYPE_RS && draft.size() > llama_n_rs_seq(ctx_tgt)); const bool use_ckpt_dft = - (ctx_dft_seq_rm_type == COMMON_CONTEXT_SEQ_RM_TYPE_RS && draft.size() > llama_n_rs_seq(ctx_dft.get())); + (ctx_dft_seq_rm_type == COMMON_CONTEXT_SEQ_RM_TYPE_RS && draft.size() > llama_n_rs_seq(ctx_dft)); if (use_ckpt_tgt) { //const int64_t t_start = ggml_time_us(); @@ -3038,7 +2982,7 @@ private: } if (use_ckpt_dft) { - ckpt.update_dft(ctx_dft.get(), slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); + ckpt.update_dft(ctx_dft, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); } } }); @@ -3219,8 +3163,8 @@ private: common_context_seq_add(ctx_tgt, slot.id, head_c, head_c + n_match, kv_shift); if (ctx_dft) { - common_context_seq_rm (ctx_dft.get(), slot.id, head_p, head_c); - common_context_seq_add(ctx_dft.get(), slot.id, head_c, head_c + n_match, kv_shift); + common_context_seq_rm (ctx_dft, slot.id, head_p, head_c); + common_context_seq_add(ctx_dft, slot.id, head_c, head_c + n_match, kv_shift); } for (size_t i = 0; i < n_match; i++) { @@ -3321,7 +3265,7 @@ private: if (!do_reset) { // restore the context checkpoint it->load_tgt(ctx_tgt, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); - it->load_dft(ctx_dft.get(), slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); + it->load_dft(ctx_dft, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); // restore the draft's speculative state common_speculative_set_state(spec.get(), slot.id, it->data_spec); @@ -3395,7 +3339,7 @@ private: common_context_seq_rm(ctx_tgt, slot.id, p0, -1); if (ctx_dft) { - common_context_seq_rm(ctx_dft.get(), slot.id, p0, -1); + common_context_seq_rm(ctx_dft, slot.id, p0, -1); } // If using an alora, there may be uncached tokens that come From 3899b39ce2acc2e019f149b7107f24b6ca297390 Mon Sep 17 00:00:00 2001 From: Oliver Simons Date: Tue, 7 Jul 2026 17:12:19 +0200 Subject: [PATCH 05/17] CUDA: Fuse MMVQ post-scale for NVFP4 (#24481) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * CUDA: Fuse MMVQ for NVFP4 and BS 1 TODO: 1. Add tests to test-backend-ops (did verify correctness manually for one model) 2. Reorder bias/scale once PRs for NVFP4 are merged/landed * Add dense MMVQ fusion as well Perf numbers on B4500. Note qwen35 is FP8->Q8 + ./scripts/compare-llama-bench.py -b master -c osimons/nvfp4_fuse_mmvq --tool llama-bench -i llama-bench.sqlite | Model | Test | t/s master | t/s osimons/nvfp4_fuse_mmvq | Speedup | |:-------------------------|:-------------|-------------:|------------------------------:|----------:| | qwen35moe 35B.A3B NVFP4 | tg128@d32768 | 150.15 | 156.29 | 1.04 | | qwen35moe 35B.A3B Q4_K_M | tg128@d32768 | 157.91 | 157.64 | 1.00 | Perf numbers on DGX Spark + ./scripts/compare-llama-bench.py -b master -c osimons/nvfp4_fuse_mmvq --tool llama-bench -i llama-bench.sqlite | Model | Test | t/s master | t/s osimons/nvfp4_fuse_mmvq | Speedup | |:-------------------------|:-------------|-------------:|------------------------------:|----------:| | qwen35moe 35B.A3B NVFP4 | tg128@d32768 | 58.31 | 59.69 | 1.02 | | qwen35moe 35B.A3B Q4_K_M | tg128@d32768 | 54.94 | 54.79 | 1.00 | * Add tests for the added fusion ops * Cleanup test-backend-ops * Cleanup ggml-cuda/mmvq 1. Unrestrict post-scale fusion 2. Rename names accordingly 3. Remove env variable to disable fusion * Merge old mul_mat patterns into the lane-based approach * Enable fusion for MoE in shared MMVQ * Restrict scale_view_nodes, enroll MM + ADD into lane-matcher * Refactor mmvq loads, still does not help non-nvfp4 kernels * Restrict scale-fusion to NVFP4 This is necessary, as the prolog is quite heavy in GEMV for some quants/model configs, leading to net perf regression. We should really be looking to refactor this such that ratio of prologue/hot-loop/epilogue is better on the hot-loop front: + ./scripts/compare-llama-bench.py -b master -c c1b9381d327e063cc846b46b59708444b66dc4d8 --tool llama-bench -i llama-bench.sqlite | CPU | Model | Test | t/s master | t/s c1b9381d3 | Speedup | |:----------------------------|:-------------------------|:-------------|-------------:|----------------:|----------:| | INTEL(R) XEON(R) GOLD 6542Y | gemma4 26B.A4B NVFP4 | tg128@d32768 | 151.70 | 154.32 | 1.02 | | INTEL(R) XEON(R) GOLD 6542Y | gemma4 26B.A4B Q4_K_M | tg128@d32768 | 187.95 | 185.73 | 0.99 | | INTEL(R) XEON(R) GOLD 6542Y | gpt-oss 20B MXFP4 MoE | tg128@d32768 | 304.62 | 300.69 | 0.99 | | INTEL(R) XEON(R) GOLD 6542Y | qwen35moe 35B.A3B NVFP4 | tg128@d32768 | 193.72 | 211.99 | 1.09 | | INTEL(R) XEON(R) GOLD 6542Y | qwen35moe 35B.A3B Q4_K_M | tg128@d32768 | 217.76 | 218.15 | 1.00 * Reorder scale & bias-add to adhere to #24331 * Restrict lane scale to NVFP4 Don't need to test unfused combinations * Cleanup * Merge single-lane mm-fusion helpers * Refactor and clean-up host-side fusion logic * Move gate_bias and scale into the same active-thread guard Latest perf numbers: B6000 build: 5b7d9f272 (9578) + ./scripts/compare-llama-bench.py -b master -c osimons/nvfp4_fuse_mmvq --tool llama-bench -i llama-bench.sqlite | CPU | Model | Test | t/s master | t/s osimons/nvfp4_fuse_mmvq | Speedup | |:----------------------------|:-------------------------|:-------------|-------------:|------------------------------:|----------:| | INTEL(R) XEON(R) GOLD 6542Y | gemma4 26B.A4B NVFP4 | tg128@d32768 | 151.79 | 154.10 | 1.02 | | INTEL(R) XEON(R) GOLD 6542Y | gemma4 26B.A4B Q4_K_M | tg128@d32768 | 187.90 | 187.27 | 1.00 | | INTEL(R) XEON(R) GOLD 6542Y | gpt-oss 20B MXFP4 MoE | tg128@d32768 | 303.77 | 306.56 | 1.01 | | INTEL(R) XEON(R) GOLD 6542Y | qwen35moe 35B.A3B NVFP4 | tg128@d32768 | 193.41 | 207.99 | 1.08 | | INTEL(R) XEON(R) GOLD 6542Y | qwen35moe 35B.A3B Q4_K_M | tg128@d32768 | 217.60 | 218.58 | 1.00 | DGX Spark build: 5b7d9f272 (9578) + ./scripts/compare-llama-bench.py -b master -c osimons/nvfp4_fuse_mmvq --tool llama-bench -i llama-bench.sqlite | CPU | Model | Test | t/s master | t/s osimons/nvfp4_fuse_mmvq | Speedup | |:------|:-------------------------|:-------------|-------------:|------------------------------:|----------:| | CPU | gemma4 26B.A4B NVFP4 | tg128@d32768 | 34.61 | 34.84 | 1.01 | | CPU | gemma4 26B.A4B Q4_K_M | tg128@d32768 | 46.95 | 46.90 | 1.00 | | CPU | gpt-oss 20B MXFP4 MoE | tg128@d32768 | 64.84 | 64.62 | 1.00 | | CPU | qwen35moe 35B.A3B NVFP4 | tg128@d32768 | 59.63 | 60.72 | 1.02 | | CPU | qwen35moe 35B.A3B Q4_K_M | tg128@d32768 | 56.53 | 56.55 | 1.00 | PPL values for 5 chunks: this PR model mode ppl uncertainty log /mnt/share/gguf/unsloth/Qwen3.6-35B-A3B-GGUF/Qwen3.6-35B-A3B-UD-Q4_K_M.gguf fusion_enabled 5.2892 0.35389 ppl-value-checks/Qwen3.6-35B-A3B-UD-Q4_K_M.fusion_enabled.log /mnt/share/gguf/unsloth/Qwen3.6-35B-A3B-GGUF/Qwen3.6-35B-A3B-UD-Q4_K_M.gguf fusion_disabled 5.2742 0.35215 ppl-value-checks/Qwen3.6-35B-A3B-UD-Q4_K_M.fusion_disabled.log /mnt/share/gguf/nvidia/Qwen3.6-35B-A3B-2.06GB-per-token-CT/Qwen3.6-35B-A3B-2.06GB-per-token-CT_fp8_q8.gguf fusion_enabled 5.4487 0.36866 ppl-value-checks/Qwen3.6-35B-A3B-2.06GB-per-token-CT_fp8_q8.fusion_enabled.log /mnt/share/gguf/nvidia/Qwen3.6-35B-A3B-2.06GB-per-token-CT/Qwen3.6-35B-A3B-2.06GB-per-token-CT_fp8_q8.gguf fusion_disabled 5.4403 0.36782 ppl-value-checks/Qwen3.6-35B-A3B-2.06GB-per-token-CT_fp8_q8.fusion_disabled.log /mnt/share/gguf/nvidia/Gemma-4-26B-A4B-NVFP4/Gemma-4-26B-A4B-NVFP4_fp8_q8.gguf fusion_enabled 17342.4348 3703.13932 ppl-value-checks/Gemma-4-26B-A4B-NVFP4_fp8_q8.fusion_enabled.log /mnt/share/gguf/nvidia/Gemma-4-26B-A4B-NVFP4/Gemma-4-26B-A4B-NVFP4_fp8_q8.gguf fusion_disabled 18627.0624 3998.42475 ppl-value-checks/Gemma-4-26B-A4B-NVFP4_fp8_q8.fusion_disabled.log /mnt/share/gguf/ggml-org/gpt-oss-20b-GGUF/gpt-oss-20b-mxfp4.gguf fusion_enabled 363.8913 33.14007 ppl-value-checks/gpt-oss-20b-mxfp4.fusion_enabled.log /mnt/share/gguf/ggml-org/gpt-oss-20b-GGUF/gpt-oss-20b-mxfp4.gguf fusion_disabled 363.8913 33.14007 ppl-value-checks/gpt-oss-20b-mxfp4.fusion_disabled.log /mnt/share/gguf/unsloth/gemma-4-26B-A4B-it-GGUF/gemma-4-26B-A4B-it-UD-Q4_K_XL.gguf fusion_enabled 17330.3926 3716.70472 ppl-value-checks/gemma-4-26B-A4B-it-UD-Q4_K_XL.fusion_enabled.log /mnt/share/gguf/unsloth/gemma-4-26B-A4B-it-GGUF/gemma-4-26B-A4B-it-UD-Q4_K_XL.gguf fusion_disabled 17933.9524 3883.17066 ppl-value-checks/gemma-4-26B-A4B-it-UD-Q4_K_XL.fusion_disabled.log master: summary: ppl-value-checks/summary.tsv model mode ppl uncertainty log /mnt/share/gguf/unsloth/Qwen3.6-35B-A3B-GGUF/Qwen3.6-35B-A3B-UD-Q4_K_M.gguf fusion_enabled 5.2892 0.35389 ppl-value-checks/Qwen3.6-35B-A3B-UD-Q4_K_M.fusion_enabled.log /mnt/share/gguf/unsloth/Qwen3.6-35B-A3B-GGUF/Qwen3.6-35B-A3B-UD-Q4_K_M.gguf fusion_disabled 5.2742 0.35215 ppl-value-checks/Qwen3.6-35B-A3B-UD-Q4_K_M.fusion_disabled.log /mnt/share/gguf/nvidia/Qwen3.6-35B-A3B-2.06GB-per-token-CT/Qwen3.6-35B-A3B-2.06GB-per-token-CT_fp8_q8.gguf fusion_enabled 5.4487 0.36866 ppl-value-checks/Qwen3.6-35B-A3B-2.06GB-per-token-CT_fp8_q8.fusion_enabled.log /mnt/share/gguf/nvidia/Qwen3.6-35B-A3B-2.06GB-per-token-CT/Qwen3.6-35B-A3B-2.06GB-per-token-CT_fp8_q8.gguf fusion_disabled 5.4403 0.36782 ppl-value-checks/Qwen3.6-35B-A3B-2.06GB-per-token-CT_fp8_q8.fusion_disabled.log /mnt/share/gguf/nvidia/Gemma-4-26B-A4B-NVFP4/Gemma-4-26B-A4B-NVFP4_fp8_q8.gguf fusion_enabled 17342.4348 3703.13932 ppl-value-checks/Gemma-4-26B-A4B-NVFP4_fp8_q8.fusion_enabled.log /mnt/share/gguf/nvidia/Gemma-4-26B-A4B-NVFP4/Gemma-4-26B-A4B-NVFP4_fp8_q8.gguf fusion_disabled 18627.0624 3998.42475 ppl-value-checks/Gemma-4-26B-A4B-NVFP4_fp8_q8.fusion_disabled.log /mnt/share/gguf/ggml-org/gpt-oss-20b-GGUF/gpt-oss-20b-mxfp4.gguf fusion_enabled 363.8913 33.14007 ppl-value-checks/gpt-oss-20b-mxfp4.fusion_enabled.log /mnt/share/gguf/ggml-org/gpt-oss-20b-GGUF/gpt-oss-20b-mxfp4.gguf fusion_disabled 363.8913 33.14007 ppl-value-checks/gpt-oss-20b-mxfp4.fusion_disabled.log /mnt/share/gguf/unsloth/gemma-4-26B-A4B-it-GGUF/gemma-4-26B-A4B-it-UD-Q4_K_XL.gguf fusion_enabled 17330.3926 3716.70472 ppl-value-checks/gemma-4-26B-A4B-it-UD-Q4_K_XL.fusion_enabled.log /mnt/share/gguf/unsloth/gemma-4-26B-A4B-it-GGUF/gemma-4-26B-A4B-it-UD-Q4_K_XL.gguf fusion_disabled 17933.9524 3883.17066 ppl-value-checks/gemma-4-26B-A4B-it-UD-Q4_K_XL.fusion_disabled.log * Allow views to weights in ggml_can_fuse_subgraph * Remove gate_first from test_mul_mat_vec_fusion * Ditch lane-parsing approach in favor of hard-coded patterns * Apply suggestions from code review Co-authored-by: Georgi Gerganov * Rename ggml_is_constant_view_src to ggml_is_constant * Finish renaming of 0905129e9d12e2bc6f16d6d3cc4e6b40606fc893 * Readd descriptive prints for fusion debugging * Add weight-buffer pre-allocation to `test_case` This is required so we correctly test fusion of NVFP4. * Update ggml/src/ggml.c Co-authored-by: Johannes Gäßler * Add 2nd context for weights as suggested by @JohannesGaessler This reflects more natural use of ggml compared to artifically pre-allocating weights into the same context * Exclude fused tests from gradient mode I'm unsure of the current state, but naively every fusion pattern should require its own backpropagation implementation. I don't see these implemented for the CUDA backend, so we can disable tests to avoid triggering GGML_ASSERT for ggml_tensor * build_graph(ggml_context * ctx) override { GGML_ASSERT(!use_weight_context()); return build_graph(ctx, nullptr); } * Apply suggestions from code review Co-authored-by: Johannes Gäßler --------- Co-authored-by: Georgi Gerganov Co-authored-by: Johannes Gäßler --- ggml/src/ggml-cuda/common.cuh | 4 + ggml/src/ggml-cuda/ggml-cuda.cu | 396 +++++++++++++++++++++++++++++--- ggml/src/ggml-cuda/mmvq.cu | 75 ++++-- ggml/src/ggml.c | 9 +- tests/test-backend-ops.cpp | 180 ++++++++++++--- 5 files changed, 576 insertions(+), 88 deletions(-) diff --git a/ggml/src/ggml-cuda/common.cuh b/ggml/src/ggml-cuda/common.cuh index e6e50e041..290dc4aff 100644 --- a/ggml/src/ggml-cuda/common.cuh +++ b/ggml/src/ggml-cuda/common.cuh @@ -1505,12 +1505,16 @@ struct ggml_cuda_mm_fusion_args_host { const ggml_tensor * x_bias = nullptr; const ggml_tensor * gate = nullptr; const ggml_tensor * gate_bias = nullptr; + const ggml_tensor * x_scale = nullptr; + const ggml_tensor * gate_scale = nullptr; ggml_glu_op glu_op; }; struct ggml_cuda_mm_fusion_args_device { const void * x_bias = nullptr; const void * gate = nullptr; const void * gate_bias = nullptr; + const void * x_scale = nullptr; + const void * gate_scale = nullptr; ggml_glu_op glu_op; }; diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 71cc2a619..674381537 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -1582,12 +1582,18 @@ static bool ggml_cuda_should_fuse_mul_mat(const ggml_tensor * ffn_up, const ggml_tensor * ffn_gate, const ggml_tensor * glu, const ggml_tensor * ffn_up_bias = nullptr, - const ggml_tensor * ffn_gate_bias = nullptr) { + const ggml_tensor * ffn_gate_bias = nullptr, + const ggml_tensor * ffn_up_scale = nullptr, + const ggml_tensor * ffn_gate_scale = nullptr) { const bool has_bias = ffn_up_bias != nullptr || ffn_gate_bias != nullptr; + const bool has_scale = ffn_up_scale != nullptr || ffn_gate_scale != nullptr; if (has_bias && (!ffn_up_bias || !ffn_gate_bias)) { return false; } + if (has_scale && (!ffn_up_scale || !ffn_gate_scale)) { + return false; + } const bool is_mul_mat = ffn_up->op == GGML_OP_MUL_MAT && ffn_gate->op == GGML_OP_MUL_MAT && glu->op == GGML_OP_GLU; const bool is_mul_mat_id = ffn_up->op == GGML_OP_MUL_MAT_ID && ffn_gate->op == GGML_OP_MUL_MAT_ID && glu->op == GGML_OP_GLU; @@ -1599,34 +1605,45 @@ static bool ggml_cuda_should_fuse_mul_mat(const ggml_tensor * ffn_up, } const ggml_op expected_bias_op = is_mul_mat ? GGML_OP_ADD : GGML_OP_ADD_ID; + const ggml_tensor * ffn_up_bias_src = has_scale ? ffn_up_scale : ffn_up; + const ggml_tensor * ffn_gate_bias_src = has_scale ? ffn_gate_scale : ffn_gate; + const ggml_tensor * ffn_up_out = has_bias ? ffn_up_bias : ffn_up_bias_src; + const ggml_tensor * ffn_gate_out = has_bias ? ffn_gate_bias : ffn_gate_bias_src; + + if (glu->src[0] != ffn_gate_out || glu->src[1] != ffn_up_out) { + return false; + } + + if (has_scale) { + if (ffn_up_scale->op != GGML_OP_MUL || ffn_gate_scale->op != GGML_OP_MUL) { + return false; + } + const bool up_has_mm = ffn_up_scale->src[0] == ffn_up || ffn_up_scale->src[1] == ffn_up; + const bool gate_has_mm = ffn_gate_scale->src[0] == ffn_gate || ffn_gate_scale->src[1] == ffn_gate; + if (!up_has_mm || !gate_has_mm) { + return false; + } + } if (has_bias) { if (ffn_up_bias->op != expected_bias_op || ffn_gate_bias->op != expected_bias_op) { return false; } - if (glu->src[0] != ffn_gate_bias || glu->src[1] != ffn_up_bias) { - return false; - } - if (expected_bias_op == GGML_OP_ADD) { - const bool up_has_mul = ffn_up_bias->src[0] == ffn_up || ffn_up_bias->src[1] == ffn_up; - const bool gate_has_mul = ffn_gate_bias->src[0] == ffn_gate || ffn_gate_bias->src[1] == ffn_gate; + const bool up_has_mul = ffn_up_bias->src[0] == ffn_up_bias_src || ffn_up_bias->src[1] == ffn_up_bias_src; + const bool gate_has_mul = ffn_gate_bias->src[0] == ffn_gate_bias_src || ffn_gate_bias->src[1] == ffn_gate_bias_src; if (!up_has_mul || !gate_has_mul) { return false; } } else { // GGML_OP_ADD_ID - if (ffn_up_bias->src[0] != ffn_up || ffn_gate_bias->src[0] != ffn_gate) { + if (ffn_up_bias->src[0] != ffn_up_bias_src || ffn_gate_bias->src[0] != ffn_gate_bias_src) { return false; } if (ffn_up_bias->src[2] != ffn_up->src[2] || ffn_gate_bias->src[2] != ffn_gate->src[2]) { return false; } } - } else { - if (glu->src[0] != ffn_gate && glu->src[1] != ffn_up) { - return false; - } } if (ffn_up->src[0]->type != ffn_gate->src[0]->type || !ggml_are_same_shape(ffn_up->src[0], ffn_gate->src[0]) || @@ -1638,7 +1655,7 @@ static bool ggml_cuda_should_fuse_mul_mat(const ggml_tensor * ffn_up, return false; } - if (ffn_up->src[2] && (ffn_up->src[2] != ffn_gate->src[2])) { + if (is_mul_mat_id && ffn_up->src[2] != ffn_gate->src[2]) { return false; } @@ -3204,10 +3221,240 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph bool fused_mul_mat_vec = false; int fused_node_count = 0; - // gate + glu + up + auto get_mul_mat_scale = [](const ggml_tensor * scale_node, const ggml_tensor * mm_node) -> const ggml_tensor * { + const bool scale_lhs_mm = scale_node->src[0] == mm_node; + const bool scale_rhs_mm = scale_node->src[1] == mm_node; + if (!scale_lhs_mm && !scale_rhs_mm) { + return nullptr; + } + + const ggml_tensor * scale = scale_lhs_mm ? scale_node->src[1] : scale_node->src[0]; + if (mm_node->src[0]->type != GGML_TYPE_NVFP4 || scale_node->type != GGML_TYPE_F32 || + scale->type != GGML_TYPE_F32 || !ggml_is_contiguous(scale) || ggml_nelements(scale) != 1 || + !ggml_are_same_shape(scale_node, mm_node)) { + return nullptr; + } + + return scale; + }; + + auto get_mul_mat_id_scale = [](const ggml_tensor * reshape, const ggml_tensor * repeat, const ggml_tensor * getrows, + const ggml_tensor * scale_node, const ggml_tensor * mm_node) -> const ggml_tensor * { + if (repeat->src[0] != reshape || getrows->src[0] != repeat || getrows->src[1] != mm_node->src[2]) { + return nullptr; + } + if (!((scale_node->src[0] == mm_node && scale_node->src[1] == getrows) || + (scale_node->src[0] == getrows && scale_node->src[1] == mm_node))) { + return nullptr; + } + + const ggml_tensor * scale = reshape->src[0]; + if (mm_node->src[0]->type != GGML_TYPE_NVFP4 || scale_node->type != GGML_TYPE_F32 || + scale->type != GGML_TYPE_F32 || !ggml_is_contiguous(scale) || ggml_nelements(scale) != mm_node->src[0]->ne[2] || + !ggml_are_same_shape(scale_node, mm_node)) { + return nullptr; + } + + return scale; + }; + + auto get_bias_tensor = [](const ggml_tensor * bias_node, const ggml_tensor * mul_node, ggml_op op_bias) -> const ggml_tensor * { + if (op_bias == GGML_OP_ADD) { + if (bias_node->src[0] == mul_node) { + return bias_node->src[1]; + } + if (bias_node->src[1] == mul_node) { + return bias_node->src[0]; + } + return nullptr; + } + GGML_ASSERT(op_bias == GGML_OP_ADD_ID); + GGML_ASSERT(bias_node->src[0] == mul_node); + return bias_node->src[1]; + }; + + // gate + glu + up, with optional scale/bias on both lanes. for (ggml_op op : { GGML_OP_MUL_MAT, GGML_OP_MUL_MAT_ID }) { const ggml_op bias_op = op == GGML_OP_MUL_MAT ? GGML_OP_ADD : GGML_OP_ADD_ID; + if (op == GGML_OP_MUL_MAT) { + for (const bool with_bias : { false, true }) { + const int gate_idx = i; + const int gate_scale_idx = i + 1; + const int gate_bias_idx = with_bias ? i + 2 : -1; + const int up_idx = with_bias ? i + 3 : i + 2; + const int up_scale_idx = up_idx + 1; + const int up_bias_idx = with_bias ? up_idx + 2 : -1; + const int glu_idx = with_bias ? up_idx + 3 : up_idx + 2; + + const int out_nodes[] = { glu_idx }; + ggml_op ops[7]; + if (with_bias) { + ops[0] = op; + ops[1] = GGML_OP_MUL; + ops[2] = bias_op; + ops[3] = op; + ops[4] = GGML_OP_MUL; + ops[5] = bias_op; + ops[6] = GGML_OP_GLU; + } else { + ops[0] = op; + ops[1] = GGML_OP_MUL; + ops[2] = op; + ops[3] = GGML_OP_MUL; + ops[4] = GGML_OP_GLU; + } + const int n_ops = with_bias ? 7 : 5; + + if (!ggml_can_fuse_subgraph(cgraph, i, n_ops, ops, out_nodes, 1) || + !ggml_cuda_check_fusion_memory_ranges(cgraph, i, n_ops, out_nodes, 1)) { + continue; + } + + ggml_tensor * gate_n = cgraph->nodes[gate_idx]; + ggml_tensor * gate_scale_n = cgraph->nodes[gate_scale_idx]; + ggml_tensor * gate_out_n = with_bias ? cgraph->nodes[gate_bias_idx] : gate_scale_n; + ggml_tensor * up_n = cgraph->nodes[up_idx]; + ggml_tensor * up_scale_n = cgraph->nodes[up_scale_idx]; + ggml_tensor * up_out_n = with_bias ? cgraph->nodes[up_bias_idx] : up_scale_n; + const ggml_tensor * glu = cgraph->nodes[glu_idx]; + + if (!ggml_cuda_should_fuse_mul_mat(up_n, gate_n, glu, + with_bias ? up_out_n : nullptr, with_bias ? gate_out_n : nullptr, up_scale_n, gate_scale_n)) { + continue; + } + + const ggml_tensor * gate_scale = get_mul_mat_scale(gate_scale_n, gate_n); + const ggml_tensor * up_scale = get_mul_mat_scale(up_scale_n, up_n); + if (!gate_scale || !up_scale) { + continue; + } + + const ggml_tensor * up_bias = with_bias ? get_bias_tensor(up_out_n, up_scale_n, bias_op) : nullptr; + const ggml_tensor * gate_bias = with_bias ? get_bias_tensor(gate_out_n, gate_scale_n, bias_op) : nullptr; + if (with_bias && (!ggml_are_same_shape(gate_out_n->src[0], gate_out_n->src[1]) || + !ggml_are_same_shape(up_out_n->src[0], up_out_n->src[1]))) { + continue; + } + + const ggml_tensor * src0 = up_n->src[0]; + const ggml_tensor * src1 = up_n->src[1]; + const ggml_tensor * ids = up_n->src[2]; + + ggml_cuda_mm_fusion_args_host fusion_data{}; + fusion_data.gate = gate_n->src[0]; + fusion_data.x_bias = up_bias; + fusion_data.gate_bias = gate_bias; + fusion_data.x_scale = up_scale; + fusion_data.gate_scale = gate_scale; + fusion_data.glu_op = ggml_get_glu_op(glu); + + if (ggml_cuda_should_fuse_mul_mat_vec_q(up_n)) { + ggml_cuda_mul_mat_vec_q(*cuda_ctx, src0, src1, ids, cgraph->nodes[glu_idx], &fusion_data); + fused_mul_mat_vec = true; + fused_node_count = n_ops; + break; + } + } + + if (fused_mul_mat_vec) { + break; + } + } else { + for (const bool with_bias : { false, true }) { + const int gate_idx = i; + const int gate_scale_idx = i + 4; + const int gate_bias_idx = with_bias ? i + 5 : -1; + const int up_idx = with_bias ? i + 6 : i + 5; + const int up_scale_idx = up_idx + 4; + const int up_bias_idx = with_bias ? up_idx + 5 : -1; + const int glu_idx = with_bias ? up_idx + 6 : up_idx + 5; + + const int out_nodes[] = { glu_idx }; + ggml_op ops[13]; + if (with_bias) { + ops[0] = op; + ops[1] = GGML_OP_RESHAPE; + ops[2] = GGML_OP_REPEAT; + ops[3] = GGML_OP_GET_ROWS; + ops[4] = GGML_OP_MUL; + ops[5] = bias_op; + ops[6] = op; + ops[7] = GGML_OP_RESHAPE; + ops[8] = GGML_OP_REPEAT; + ops[9] = GGML_OP_GET_ROWS; + ops[10] = GGML_OP_MUL; + ops[11] = bias_op; + ops[12] = GGML_OP_GLU; + } else { + ops[0] = op; + ops[1] = GGML_OP_RESHAPE; + ops[2] = GGML_OP_REPEAT; + ops[3] = GGML_OP_GET_ROWS; + ops[4] = GGML_OP_MUL; + ops[5] = op; + ops[6] = GGML_OP_RESHAPE; + ops[7] = GGML_OP_REPEAT; + ops[8] = GGML_OP_GET_ROWS; + ops[9] = GGML_OP_MUL; + ops[10] = GGML_OP_GLU; + } + const int n_ops = with_bias ? 13 : 11; + + if (!ggml_can_fuse_subgraph(cgraph, i, n_ops, ops, out_nodes, 1) || + !ggml_cuda_check_fusion_memory_ranges(cgraph, i, n_ops, out_nodes, 1)) { + continue; + } + + ggml_tensor * gate_n = cgraph->nodes[gate_idx]; + ggml_tensor * gate_scale_n = cgraph->nodes[gate_scale_idx]; + ggml_tensor * gate_out_n = with_bias ? cgraph->nodes[gate_bias_idx] : gate_scale_n; + ggml_tensor * up_n = cgraph->nodes[up_idx]; + ggml_tensor * up_scale_n = cgraph->nodes[up_scale_idx]; + ggml_tensor * up_out_n = with_bias ? cgraph->nodes[up_bias_idx] : up_scale_n; + const ggml_tensor * glu = cgraph->nodes[glu_idx]; + + if (!ggml_cuda_should_fuse_mul_mat(up_n, gate_n, glu, + with_bias ? up_out_n : nullptr, with_bias ? gate_out_n : nullptr, up_scale_n, gate_scale_n)) { + continue; + } + + const ggml_tensor * gate_scale = get_mul_mat_id_scale(cgraph->nodes[gate_idx + 1], cgraph->nodes[gate_idx + 2], + cgraph->nodes[gate_idx + 3], gate_scale_n, gate_n); + const ggml_tensor * up_scale = get_mul_mat_id_scale(cgraph->nodes[up_idx + 1], cgraph->nodes[up_idx + 2], + cgraph->nodes[up_idx + 3], up_scale_n, up_n); + if (!gate_scale || !up_scale) { + continue; + } + + const ggml_tensor * up_bias = with_bias ? get_bias_tensor(up_out_n, up_scale_n, bias_op) : nullptr; + const ggml_tensor * gate_bias = with_bias ? get_bias_tensor(gate_out_n, gate_scale_n, bias_op) : nullptr; + + const ggml_tensor * src0 = up_n->src[0]; + const ggml_tensor * src1 = up_n->src[1]; + const ggml_tensor * ids = up_n->src[2]; + + ggml_cuda_mm_fusion_args_host fusion_data{}; + fusion_data.gate = gate_n->src[0]; + fusion_data.x_bias = up_bias; + fusion_data.gate_bias = gate_bias; + fusion_data.x_scale = up_scale; + fusion_data.gate_scale = gate_scale; + fusion_data.glu_op = ggml_get_glu_op(glu); + + if (ggml_cuda_should_fuse_mul_mat_vec_q(up_n)) { + ggml_cuda_mul_mat_vec_q(*cuda_ctx, src0, src1, ids, cgraph->nodes[glu_idx], &fusion_data); + fused_mul_mat_vec = true; + fused_node_count = n_ops; + break; + } + } + + if (fused_mul_mat_vec) { + break; + } + } + if (ggml_cuda_can_fuse(cgraph, i, { op, bias_op, op, bias_op, GGML_OP_GLU }, {})) { ggml_tensor * glu = cgraph->nodes[i + 4]; ggml_tensor * gate_bias_n = glu->src[0]; @@ -3227,23 +3474,8 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph continue; } - auto get_bias_tensor = [](const ggml_tensor * bias_node, const ggml_tensor * mul_node, ggml_op op_bias) { - if (op_bias == GGML_OP_ADD) { - if (bias_node->src[0] == mul_node) { - return bias_node->src[1]; - } - if (bias_node->src[1] == mul_node) { - return bias_node->src[0]; - } - return (ggml_tensor *) nullptr; - } - GGML_ASSERT(op_bias == GGML_OP_ADD_ID); - GGML_ASSERT(bias_node->src[0] == mul_node); - return bias_node->src[1]; - }; - - ggml_tensor * up_bias_tensor = get_bias_tensor(up_bias_n, up_n, bias_op); - ggml_tensor * gate_bias_tensor = get_bias_tensor(gate_bias_n, gate_n, bias_op); + const ggml_tensor * up_bias_tensor = get_bias_tensor(up_bias_n, up_n, bias_op); + const ggml_tensor * gate_bias_tensor = get_bias_tensor(gate_bias_n, gate_n, bias_op); if (!up_bias_tensor || !gate_bias_tensor) { continue; @@ -3331,7 +3563,95 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph fused_mul_mat_vec = false; fused_node_count = 0; - // gate + add + glu + up + add + // mul_mat + scale + optional bias + for (ggml_op op : { GGML_OP_MUL_MAT, GGML_OP_MUL_MAT_ID }) { + const ggml_op bias_op = op == GGML_OP_MUL_MAT ? GGML_OP_ADD : GGML_OP_ADD_ID; + + for (const bool with_bias : { false, true }) { + const int n_ops = op == GGML_OP_MUL_MAT ? (with_bias ? 3 : 2) : (with_bias ? 6 : 5); + const int out_nodes[] = { i + n_ops - 1 }; + ggml_op ops[6]; + if (op == GGML_OP_MUL_MAT) { + if (with_bias) { + ops[0] = op; + ops[1] = GGML_OP_MUL; + ops[2] = bias_op; + } else { + ops[0] = op; + ops[1] = GGML_OP_MUL; + } + } else { + if (with_bias) { + ops[0] = op; + ops[1] = GGML_OP_RESHAPE; + ops[2] = GGML_OP_REPEAT; + ops[3] = GGML_OP_GET_ROWS; + ops[4] = GGML_OP_MUL; + ops[5] = bias_op; + } else { + ops[0] = op; + ops[1] = GGML_OP_RESHAPE; + ops[2] = GGML_OP_REPEAT; + ops[3] = GGML_OP_GET_ROWS; + ops[4] = GGML_OP_MUL; + } + } + + if (!ggml_can_fuse_subgraph(cgraph, i, n_ops, ops, out_nodes, 1) || + !ggml_cuda_check_fusion_memory_ranges(cgraph, i, n_ops, out_nodes, 1)) { + continue; + } + + ggml_tensor * mm_node = cgraph->nodes[i]; + ggml_tensor * scale_node = op == GGML_OP_MUL_MAT ? cgraph->nodes[i + 1] : cgraph->nodes[i + 4]; + ggml_tensor * out_node = with_bias ? cgraph->nodes[i + n_ops - 1] : scale_node; + + const ggml_tensor * scale = nullptr; + if (op == GGML_OP_MUL_MAT) { + scale = get_mul_mat_scale(scale_node, mm_node); + } else { + scale = get_mul_mat_id_scale(cgraph->nodes[i + 1], cgraph->nodes[i + 2], cgraph->nodes[i + 3], scale_node, mm_node); + } + if (!scale) { + continue; + } + + const ggml_tensor * bias = with_bias ? get_bias_tensor(out_node, scale_node, bias_op) : nullptr; + if (with_bias && !bias) { + continue; + } + if (with_bias && bias_op == GGML_OP_ADD && !ggml_are_same_shape(out_node->src[0], out_node->src[1])) { + continue; + } + if (with_bias && bias_op == GGML_OP_ADD_ID && out_node->src[2] != mm_node->src[2]) { + continue; + } + + const ggml_tensor * src0 = mm_node->src[0]; + const ggml_tensor * src1 = mm_node->src[1]; + const ggml_tensor * ids = mm_node->src[2]; + + ggml_cuda_mm_fusion_args_host fusion_data{}; + fusion_data.x_bias = bias; + fusion_data.x_scale = scale; + + if (ggml_cuda_should_fuse_mul_mat_vec_q(mm_node)) { + ggml_cuda_mul_mat_vec_q(*cuda_ctx, src0, src1, ids, out_node, &fusion_data); + fused_mul_mat_vec = true; + fused_node_count = n_ops; + break; + } + } + if (fused_mul_mat_vec) { + break; + } + } + + if (fused_mul_mat_vec) { + return fused_node_count - 1; + } + + // mul_mat + add for (ggml_op op : { GGML_OP_MUL_MAT, GGML_OP_MUL_MAT_ID }) { const ggml_op bias_op = op == GGML_OP_MUL_MAT ? GGML_OP_ADD : GGML_OP_ADD_ID; @@ -3562,12 +3882,6 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud } } -#ifdef GGML_CUDA_DEBUG - const int nodes_fused = i - prev_i - 1; - if (nodes_fused > 0) { - GGML_LOG_INFO("nodes_fused: %d\n", nodes_fused); - } -#endif prev_i = i; if (ggml_cuda_is_view_or_noop(node)) { @@ -3581,6 +3895,12 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud int nodes_to_skip = ggml_cuda_try_fuse(cuda_ctx, cgraph, i); if (nodes_to_skip != 0) { +#ifdef GGML_CUDA_DEBUG + const int last_fused = i + nodes_to_skip; + GGML_LOG_INFO("nodes_fused: %d, first: %s (%s), last: %s (%s)\n", + nodes_to_skip + 1, ggml_op_name(node->op), node->name, + ggml_op_name(cgraph->nodes[last_fused]->op), cgraph->nodes[last_fused]->name); +#endif i += nodes_to_skip; continue; } diff --git a/ggml/src/ggml-cuda/mmvq.cu b/ggml/src/ggml-cuda/mmvq.cu index 2278c7d9b..a48cc48b2 100644 --- a/ggml/src/ggml-cuda/mmvq.cu +++ b/ggml/src/ggml-cuda/mmvq.cu @@ -521,9 +521,13 @@ static __global__ void mul_mat_vec_q( bool use_gate = false; bool use_bias = false; bool use_gate_bias = false; + bool use_scale = false; + bool use_gate_scale = false; [[maybe_unused]] const void * vgate = nullptr; const float * x_bias = nullptr; const float * gate_bias = nullptr; + const float * x_scale = nullptr; + const float * gate_scale = nullptr; ggml_glu_op active_glu; if constexpr (has_fusion) { @@ -534,34 +538,47 @@ static __global__ void mul_mat_vec_q( x_bias = (const float *) fusion.x_bias; gate_bias = (const float *) fusion.gate_bias; active_glu = fusion.glu_op; + if constexpr (type == GGML_TYPE_NVFP4) { + use_scale = fusion.x_scale != nullptr; + use_gate_scale = fusion.gate_scale != nullptr && use_gate; + x_scale = (const float *) fusion.x_scale; + gate_scale = (const float *) fusion.gate_scale; + } } [[maybe_unused]] float x_biases[ncols_dst] = { 0.0f }; [[maybe_unused]] float gate_biases[ncols_dst] = { 0.0f }; + [[maybe_unused]] float x_scales; + [[maybe_unused]] float gate_scales; if constexpr (has_fusion) { + // 1. Hide latency by prefetching bias, gates and scales here + // 2. load only on threads that won't die after partial sum calculation const uint32_t channel_bias = ids ? channel_x : channel_dst; - if (use_bias) { - x_bias = x_bias + sample_dst*stride_sample_dst + channel_bias*stride_channel_dst + row0; - // 1. Hide latency by prefetching bias and gate here - // 2. load only on threads that won't die after partial sum calculation - if (threadIdx.x < rows_per_cuda_block && threadIdx.y == 0 && - (rows_per_cuda_block == 1 || uint32_t(row0 + threadIdx.x) < stride_col_dst)) { + if (threadIdx.x < rows_per_cuda_block && threadIdx.y == 0 && + (rows_per_cuda_block == 1 || uint32_t(row0 + threadIdx.x) < stride_col_dst)) { + if (use_bias) { + x_bias = x_bias + sample_dst * stride_sample_dst + channel_bias * stride_channel_dst + row0; #pragma unroll for (int j = 0; j < ncols_dst; ++j) { x_biases[j] = x_bias[j * stride_col_dst + threadIdx.x]; } } - } - if (use_gate_bias) { - gate_bias = gate_bias + sample_dst*stride_sample_dst + channel_bias*stride_channel_dst + row0; - if (threadIdx.x < rows_per_cuda_block && threadIdx.y == 0 && - (rows_per_cuda_block == 1 || uint32_t(row0 + threadIdx.x) < stride_col_dst)) { + if (use_gate_bias) { + gate_bias = gate_bias + sample_dst * stride_sample_dst + channel_bias * stride_channel_dst + row0; #pragma unroll for (int j = 0; j < ncols_dst; ++j) { gate_biases[j] = gate_bias[j * stride_col_dst + threadIdx.x]; } } + if constexpr (type == GGML_TYPE_NVFP4) { + if (use_scale) { + x_scales = x_scale[ids ? channel_x : 0]; + } + if (use_gate_scale) { + gate_scales = gate_scale[ids ? channel_x : 0]; + } + } } } @@ -643,11 +660,21 @@ static __global__ void mul_mat_vec_q( if (threadIdx.x < rows_per_cuda_block && (rows_per_cuda_block == 1 || uint32_t(row0 + threadIdx.x) < stride_col_dst)) { float result = tmp[j][threadIdx.x]; if constexpr (has_fusion) { + if constexpr (type == GGML_TYPE_NVFP4) { + if (use_scale) { + result *= x_scales; + } + } if (use_bias) { result += x_biases[j]; } if (use_gate) { float gate_value = tmp_gate[j][threadIdx.x]; + if constexpr (type == GGML_TYPE_NVFP4) { + if (use_gate_scale) { + gate_value *= gate_scales; + } + } if (use_gate_bias) { gate_value += gate_biases[j]; } @@ -673,7 +700,10 @@ static __global__ void mul_mat_vec_q( } if constexpr (!has_fusion) { - GGML_UNUSED_VARS(use_gate, use_bias, use_gate_bias, active_glu, gate_bias, x_bias, tmp_gate); + GGML_UNUSED_VARS(use_gate, use_bias, use_gate_bias, use_scale, use_gate_scale, active_glu, gate_bias, x_bias, x_scale, gate_scale, tmp_gate); + } + if constexpr (type != GGML_TYPE_NVFP4) { + GGML_UNUSED_VARS(use_scale, use_gate_scale, x_scale, gate_scale, x_scales, gate_scales); } } @@ -769,7 +799,8 @@ static void mul_mat_vec_q_switch_fusion( const dim3 & block_nums, const dim3 & block_dims, const int nbytes_shared, const uint32_t ids_stride, cudaStream_t stream) { - const bool has_fusion = fusion.gate != nullptr || fusion.x_bias != nullptr || fusion.gate_bias != nullptr; + const bool has_fusion = fusion.gate != nullptr || fusion.x_bias != nullptr || fusion.gate_bias != nullptr || + fusion.x_scale != nullptr || fusion.gate_scale != nullptr; if constexpr (c_ncols_dst == 1) { if (has_fusion) { const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(block_nums, block_dims, nbytes_shared, stream); @@ -834,7 +865,6 @@ static void mul_mat_vec_q_switch_ncols_dst( const int warp_size = ggml_cuda_info().devices[device].warp_size; const mmvq_parameter_table_id table_id = get_device_table_id(cc); - const bool has_fusion = fusion.gate != nullptr || fusion.x_bias != nullptr || fusion.gate_bias != nullptr; const bool has_ids = ids != nullptr; const auto should_use_small_k = [&](int c_ncols_dst) { @@ -973,8 +1003,6 @@ static void mul_mat_vec_q_switch_ncols_dst( GGML_ABORT("fatal error"); break; } - - GGML_UNUSED(has_fusion); } static void mul_mat_vec_q_switch_type( const void * vx, const ggml_type type_x, const void * vy, const int32_t * ids, const ggml_cuda_mm_fusion_args_device fusion, float * dst, @@ -1154,6 +1182,9 @@ void ggml_cuda_mul_mat_vec_q( if (fusion) { GGML_ASSERT( !ids || dst->ne[2] == 1); GGML_ASSERT( ids || dst->ne[1] == 1); + // Scale fusion is only allowed for NVFP4 currently as the cost of checking this at run-time in the prologue is + // non-negligible for some models such as gpt-oss-20b + GGML_ASSERT((fusion->x_scale == nullptr && fusion->gate_scale == nullptr) || src0->type == GGML_TYPE_NVFP4); if (fusion->x_bias) { GGML_ASSERT(fusion->x_bias->type == GGML_TYPE_F32); @@ -1171,6 +1202,18 @@ void ggml_cuda_mul_mat_vec_q( GGML_ASSERT(!ids || fusion->gate_bias->ne[1] == src0->ne[2]); fusion_local.gate_bias = fusion->gate_bias->data; } + if (fusion->x_scale) { + GGML_ASSERT(fusion->x_scale->type == GGML_TYPE_F32); + GGML_ASSERT(ggml_is_contiguous(fusion->x_scale)); + GGML_ASSERT(ggml_nelements(fusion->x_scale) == (ids ? src0->ne[2] : 1)); + fusion_local.x_scale = fusion->x_scale->data; + } + if (fusion->gate_scale) { + GGML_ASSERT(fusion->gate_scale->type == GGML_TYPE_F32); + GGML_ASSERT(ggml_is_contiguous(fusion->gate_scale)); + GGML_ASSERT(ggml_nelements(fusion->gate_scale) == (ids ? src0->ne[2] : 1)); + fusion_local.gate_scale = fusion->gate_scale->data; + } fusion_local.glu_op = fusion->glu_op; } diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c index 440095f80..fc46ab0aa 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -7419,6 +7419,10 @@ static int ggml_node_list_find_tensor(const struct ggml_cgraph * cgraph, return -1; } +static bool ggml_is_constant(const struct ggml_tensor * tensor) { + return tensor->buffer != NULL && ggml_backend_buffer_get_usage(tensor->buffer) == GGML_BACKEND_BUFFER_USAGE_WEIGHTS && (tensor->flags & GGML_TENSOR_FLAG_PARAM) == 0; +} + bool ggml_can_fuse_subgraph_ext(const struct ggml_cgraph * cgraph, const int * node_idxs, int count, @@ -7464,10 +7468,11 @@ bool ggml_can_fuse_subgraph_ext(const struct ggml_cgraph * cgraph, return false; } - // if node is a view, check if the view_src and all it's parent view_srcs are within the subgraph + // if node is a view, check if the view_src and all its parent view_srcs are within the subgraph. + // external view sources are allowed only for weight tensors, which are constant for this graph execution. struct ggml_tensor * view_src = node->view_src; while (view_src) { - if (ggml_node_list_find_tensor(cgraph, node_idxs, count, view_src) == -1) { + if (ggml_node_list_find_tensor(cgraph, node_idxs, count, view_src) == -1 && !ggml_is_constant(view_src)) { return false; } view_src = view_src->view_src; diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index c21675ef5..82a9d9546 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -1137,6 +1137,10 @@ struct test_case { } virtual ggml_tensor * build_graph(ggml_context * ctx) = 0; + virtual ggml_tensor * build_graph(ggml_context * ctx, ggml_context * ctx_weights) { + GGML_UNUSED(ctx_weights); + return build_graph(ctx); + } virtual double max_nmse_err() { return 1e-7; @@ -1213,6 +1217,7 @@ struct test_case { virtual bool run_whole_graph() { return false; } virtual std::vector fusion_test_nodes() { return {}; } + virtual bool use_weight_context() { return false; } ggml_cgraph * gf = nullptr; ggml_cgraph * gb = nullptr; @@ -1319,20 +1324,28 @@ struct test_case { /* .mem_base = */ NULL, /* .no_alloc = */ true, }; + const bool use_weights = use_weight_context(); + ggml_context * ctx = ggml_init(params); GGML_ASSERT(ctx); + ggml_context * ctx_weights = use_weights ? ggml_init(params) : nullptr; + GGML_ASSERT(!use_weights || ctx_weights); gf = ggml_new_graph(ctx); // pre-graph sentinel add_sentinel(ctx); + if (ctx_weights) { + add_sentinel(ctx_weights); + } - ggml_tensor * out = build_graph(ctx); + ggml_tensor * out = build_graph(ctx, ctx_weights); current_op_name = op_desc(out); check_for_f16_tensor(ctx); if (!matches_filter(out, op_names_filter)) { //printf(" %s: skipping\n", op_desc(out).c_str()); + ggml_free(ctx_weights); ggml_free(ctx); return test_status_t::SKIPPED; } @@ -1355,18 +1368,36 @@ struct test_case { print_test_result_locked(output_printer, result); + ggml_free(ctx_weights); ggml_free(ctx); return test_status_t::NOT_SUPPORTED; } // post-graph sentinel add_sentinel(ctx); + if (ctx_weights) { + add_sentinel(ctx_weights); + } + + ggml_backend_buffer_t buf_weights = nullptr; + if (ctx_weights) { + buf_weights = ggml_backend_alloc_ctx_tensors(ctx_weights, backend1); + if (buf_weights == NULL) { + printf("failed to allocate weight tensors [%s] ", ggml_backend_name(backend1)); + ggml_free(ctx_weights); + ggml_free(ctx); + return test_status_t::FAIL; + } + ggml_backend_buffer_set_usage(buf_weights, GGML_BACKEND_BUFFER_USAGE_WEIGHTS); + } // allocate ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors(ctx, backend1); if (buf == NULL) { printf("failed to allocate tensors [%s] ", ggml_backend_name(backend1)); + ggml_backend_buffer_free(buf_weights); + ggml_free(ctx_weights); ggml_free(ctx); return test_status_t::FAIL; } @@ -1381,6 +1412,9 @@ struct test_case { // randomize tensors initialize_tensors(ctx); + if (ctx_weights) { + initialize_tensors(ctx_weights); + } // compare struct callback_userdata { @@ -1466,7 +1500,8 @@ struct test_case { fused_nodes_to_verify.size()); ggml_backend_buffer_free(buf); - + ggml_backend_buffer_free(buf_weights); + ggml_free(ctx_weights); ggml_free(ctx); // Create test result @@ -1490,10 +1525,14 @@ struct test_case { /* .mem_base = */ NULL, /* .no_alloc = */ true, }; + const bool use_weights = use_weight_context(); + ggml_context_ptr ctx(ggml_init(params)); // smart ptr GGML_ASSERT(ctx); + ggml_context_ptr ctx_weights(use_weights ? ggml_init(params) : nullptr); + GGML_ASSERT(!use_weights || ctx_weights); - ggml_tensor * out = build_graph(ctx.get()); + ggml_tensor * out = build_graph(ctx.get(), ctx_weights.get()); current_op_name = op_desc(out); if (!matches_filter(out, op_names_filter)) { //printf(" %s: skipping\n", op_desc(out).c_str()); @@ -1510,6 +1549,16 @@ struct test_case { return true; } + ggml_backend_buffer_ptr buf_weights(nullptr); + if (ctx_weights) { + buf_weights.reset(ggml_backend_alloc_ctx_tensors(ctx_weights.get(), backend)); + if (buf_weights == NULL) { + printf("failed to allocate weight tensors\n"); + return false; + } + ggml_backend_buffer_set_usage(buf_weights.get(), GGML_BACKEND_BUFFER_USAGE_WEIGHTS); + } + // allocate ggml_backend_buffer_ptr buf(ggml_backend_alloc_ctx_tensors(ctx.get(), backend)); // smart ptr @@ -1520,6 +1569,9 @@ struct test_case { // randomize tensors initialize_tensors(ctx.get()); + if (ctx_weights) { + initialize_tensors(ctx_weights.get()); + } // build graph ggml_cgraph * gf = ggml_new_graph_custom(ctx.get(), graph_nodes, false); @@ -5848,19 +5900,21 @@ struct test_mul_mat_vec_fusion : public test_case { const bool b; // broadcast b matrix (only for use_id) const bool with_bias; const bool with_gate; + const bool with_lane_scale; std::array batch_dims; test_mul_mat_vec_fusion(ggml_type type, ggml_glu_op op, int64_t m, int64_t n, int64_t k, bool use_id = false, int n_mats = 1, int n_used = 1, bool b = false, bool with_bias = false, bool with_gate = true, - std::array batch_dims = {4, 2}) - : type(type), glu_op(op), m(m), n(n), k(k), use_id(use_id), n_mats(n_mats), n_used(n_used), b(b), with_bias(with_bias), with_gate(with_gate), batch_dims(batch_dims) { + bool with_lane_scale = false, std::array batch_dims = {4, 2}) + : type(type), glu_op(op), m(m), n(n), k(k), use_id(use_id), n_mats(n_mats), n_used(n_used), b(b), with_bias(with_bias), + with_gate(with_gate), with_lane_scale(with_lane_scale), batch_dims(batch_dims) { if (use_id) { GGML_ASSERT(n_used <= n_mats); } } std::string vars() override { - return VARS_TO_STR12(type, glu_op, m, n, k, use_id, n_mats, n_used, b, with_bias, with_gate, batch_dims); + return VARS_TO_STR13(type, glu_op, m, n, k, use_id, n_mats, n_used, b, with_bias, with_gate, with_lane_scale, batch_dims); } std::string op_desc(ggml_tensor * t) override { @@ -5869,6 +5923,7 @@ struct test_mul_mat_vec_fusion : public test_case { } bool run_whole_graph() override { return true; } + bool use_weight_context() override { return use_id && with_lane_scale; } ggml_tensor * build_gate(ggml_context * ctx, ggml_tensor * ffn_gate, ggml_tensor * ffn_up) { ggml_tensor * out = nullptr; @@ -5884,7 +5939,26 @@ struct test_mul_mat_vec_fusion : public test_case { return out; } + ggml_tensor * build_lane_scale_dense(ggml_context * ctx, ggml_tensor * out) { + ggml_tensor * scale = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 1); + return ggml_mul(ctx, out, scale); + } + + ggml_tensor * build_lane_scale_id(ggml_context * ctx, ggml_context * ctx_weights, ggml_tensor * out, ggml_tensor * ids) { + GGML_ASSERT(ctx_weights); + ggml_tensor * scale = ggml_new_tensor_1d(ctx_weights, GGML_TYPE_F32, n_mats); + ggml_tensor * s = ggml_reshape_3d(ctx, scale, 1, n_mats, 1); + s = ggml_repeat_4d(ctx, s, 1, n_mats, m, 1); + s = ggml_get_rows(ctx, s, ids); + return ggml_mul(ctx, out, s); + } + ggml_tensor * build_graph(ggml_context * ctx) override { + GGML_ASSERT(!use_weight_context()); + return build_graph(ctx, nullptr); + } + + ggml_tensor * build_graph(ggml_context * ctx, ggml_context * ctx_weights) override { if (!use_id) { const int channels = batch_dims[0]; const int samples = batch_dims[1]; @@ -5895,19 +5969,34 @@ struct test_mul_mat_vec_fusion : public test_case { ggml_tensor * gate = with_gate ? ggml_new_tensor(ctx, type, 4, ne0.data()) : nullptr; ggml_tensor * up = ggml_new_tensor(ctx, type, 4, ne0.data()); - ggml_tensor * ffn_up = ggml_mul_mat(ctx, up, cur); - if (with_bias) { - std::array bias_ne = { ffn_up->ne[0], 1, channels, samples }; - ggml_tensor * up_bias = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, bias_ne.data()); - ffn_up = ggml_add(ctx, ffn_up, up_bias); - } + auto build_lane_up = [&]() { + ggml_tensor * ffn_up = ggml_mul_mat(ctx, up, cur); + if (with_lane_scale) { + ffn_up = build_lane_scale_dense(ctx, ffn_up); + } + if (with_bias) { + std::array bias_ne = { ffn_up->ne[0], 1, channels, samples }; + ggml_tensor * up_bias = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, bias_ne.data()); + ffn_up = ggml_add(ctx, ffn_up, up_bias); + } + return ffn_up; + }; - ggml_tensor * ffn_gate = with_gate ? ggml_mul_mat(ctx, gate, cur) : nullptr; - if (with_bias && with_gate) { - std::array bias_ne = { ffn_gate->ne[0], 1, channels, samples }; - ggml_tensor * gate_bias = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, bias_ne.data()); - ffn_gate = ggml_add(ctx, ffn_gate, gate_bias); - } + auto build_lane_gate = [&]() { + ggml_tensor * ffn_gate = ggml_mul_mat(ctx, gate, cur); + if (with_lane_scale) { + ffn_gate = build_lane_scale_dense(ctx, ffn_gate); + } + if (with_bias) { + std::array bias_ne = { ffn_gate->ne[0], 1, channels, samples }; + ggml_tensor * gate_bias = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, bias_ne.data()); + ffn_gate = ggml_add(ctx, ffn_gate, gate_bias); + } + return ffn_gate; + }; + + ggml_tensor * ffn_up = build_lane_up(); + ggml_tensor * ffn_gate = with_gate ? build_lane_gate() : nullptr; ggml_tensor * out = with_gate ? build_gate(ctx, ffn_gate, ffn_up) : ffn_up; @@ -5929,17 +6018,32 @@ struct test_mul_mat_vec_fusion : public test_case { ggml_tensor * cur = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, k, this->b ? 1 : n_used, m); ggml_set_name(cur, "cur"); - ggml_tensor * ffn_up = ggml_mul_mat_id(ctx, ups, cur, ids); - if (with_bias) { - ggml_tensor * up_bias_param = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, ffn_up->ne[0], n_mats); - ffn_up = ggml_add_id(ctx, ffn_up, up_bias_param, ids); - } + auto build_lane_up = [&]() { + ggml_tensor * ffn_up = ggml_mul_mat_id(ctx, ups, cur, ids); + if (with_lane_scale) { + ffn_up = build_lane_scale_id(ctx, ctx_weights, ffn_up, ids); + } + if (with_bias) { + ggml_tensor * up_bias_param = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, ffn_up->ne[0], n_mats); + ffn_up = ggml_add_id(ctx, ffn_up, up_bias_param, ids); + } + return ffn_up; + }; - ggml_tensor * ffn_gate = with_gate? ggml_mul_mat_id(ctx, gates, cur, ids) : nullptr; - if (with_bias && with_gate) { - ggml_tensor * gate_bias_param = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, ffn_gate->ne[0], n_mats); - ffn_gate = ggml_add_id(ctx, ffn_gate, gate_bias_param, ids); - } + auto build_lane_gate = [&]() { + ggml_tensor * ffn_gate = ggml_mul_mat_id(ctx, gates, cur, ids); + if (with_lane_scale) { + ffn_gate = build_lane_scale_id(ctx, ctx_weights, ffn_gate, ids); + } + if (with_bias) { + ggml_tensor * gate_bias_param = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, ffn_gate->ne[0], n_mats); + ffn_gate = ggml_add_id(ctx, ffn_gate, gate_bias_param, ids); + } + return ffn_gate; + }; + + ggml_tensor * ffn_up = build_lane_up(); + ggml_tensor * ffn_gate = with_gate ? build_lane_gate() : nullptr; ggml_tensor * out = with_gate ? build_gate(ctx, ffn_gate, ffn_up) : ffn_up; @@ -9202,10 +9306,15 @@ static std::vector> make_test_cases_eval() { if (!with_gate && glu_op != GGML_GLU_OP_SWIGLU) { continue; } - test_cases.emplace_back(new test_mul_mat_vec_fusion(type, glu_op, 1, 32, 256, - use_id, 16, 8, b, with_bias, with_gate)); - test_cases.emplace_back(new test_mul_mat_vec_fusion(type, glu_op, 1, 32, 256, - use_id, 16, 8, b, with_bias, with_gate, {1, 1})); + for (bool with_lane_scale : {false, true}) { + if (with_lane_scale && type != GGML_TYPE_NVFP4) { + continue; + } + test_cases.emplace_back(new test_mul_mat_vec_fusion(type, glu_op, 1, 32, 256, + use_id, 16, 8, b, with_bias, with_gate, with_lane_scale)); + test_cases.emplace_back(new test_mul_mat_vec_fusion(type, glu_op, 1, 32, 256, + use_id, 16, 8, b, with_bias, with_gate, with_lane_scale, {1, 1})); + } } } } @@ -9823,6 +9932,13 @@ static bool test_backend(ggml_backend_t backend, ggml_backend_dev_t dev, test_mo } if (mode == MODE_GRAD) { + test_cases.erase( + std::remove_if(test_cases.begin(), test_cases.end(), [](const std::unique_ptr & tc) { + return tc->run_whole_graph(); + }), + test_cases.end() + ); + size_t n_ok = 0; for (auto & test : test_cases) { if (test->eval_grad(backend, op_names_filter, output_printer)) { From c198af4dc24f8e0ab8a569a60f931e03a192fd79 Mon Sep 17 00:00:00 2001 From: Georgi Gerganov Date: Tue, 7 Jul 2026 18:52:30 +0300 Subject: [PATCH 06/17] spec : fix naming, spacing (#25410) --- common/speculative.cpp | 17 ++++++++--------- common/speculative.h | 10 +++++----- tools/server/server-context.cpp | 13 +++++++------ 3 files changed, 20 insertions(+), 20 deletions(-) diff --git a/common/speculative.cpp b/common/speculative.cpp index 1d2977e42..580728a20 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -2246,7 +2246,7 @@ common_params common_base_params_to_speculative(const common_params & params) { return result; } -struct common_init_speculative_result::impl { +struct common_speculative_init_result::impl { impl() = default; ~impl() = default; @@ -2255,7 +2255,7 @@ struct common_init_speculative_result::impl { llama_context_ptr context; }; -common_init_speculative_result::common_init_speculative_result( +common_speculative_init_result::common_speculative_init_result( common_params & params, llama_model * model_tgt, llama_context * ctx_tgt) : @@ -2301,8 +2301,7 @@ common_init_speculative_result::common_init_speculative_result( } else if (spec_mtp) { model_path = params.model.path; - LOG_TRC("%s: creating MTP draft context against the target model '%s'\n", - __func__, model_path.c_str()); + LOG_TRC("%s: creating MTP draft context against the target model '%s'\n", __func__, model_path.c_str()); llama_context * ctx_dft = llama_init_from_model(model_tgt, cparams); if (ctx_dft == nullptr) { @@ -2314,18 +2313,18 @@ common_init_speculative_result::common_init_speculative_result( } } -common_init_speculative_result::~common_init_speculative_result() = default; +common_speculative_init_result::~common_speculative_init_result() = default; -llama_model * common_init_speculative_result::model() { +llama_model * common_speculative_init_result::model() { return pimpl->model.get(); } -llama_context * common_init_speculative_result::context() { +llama_context * common_speculative_init_result::context() { return pimpl->context.get(); } -common_init_speculative_result_ptr common_init_speculative_from_params(common_params & params, llama_model * model_tgt, llama_context * ctx_tgt) { - return std::make_unique(params, model_tgt, ctx_tgt); +common_speculative_init_result_ptr common_speculative_init_from_params(common_params & params, llama_model * model_tgt, llama_context * ctx_tgt) { + return std::make_unique(params, model_tgt, ctx_tgt); } // initialization of the speculative decoding system diff --git a/common/speculative.h b/common/speculative.h index 2f6d225b6..062bf2093 100644 --- a/common/speculative.h +++ b/common/speculative.h @@ -83,9 +83,9 @@ struct common_speculative_deleter { typedef std::unique_ptr common_speculative_ptr; -struct common_init_speculative_result { - common_init_speculative_result(common_params & params, llama_model * model_tgt, llama_context * ctx_tgt); - ~common_init_speculative_result(); +struct common_speculative_init_result { + common_speculative_init_result(common_params & params, llama_model * model_tgt, llama_context * ctx_tgt); + ~common_speculative_init_result(); llama_model * model(); llama_context * context(); @@ -95,6 +95,6 @@ private: std::unique_ptr pimpl; }; -using common_init_speculative_result_ptr = std::unique_ptr; +using common_speculative_init_result_ptr = std::unique_ptr; -common_init_speculative_result_ptr common_init_speculative_from_params(common_params & params, llama_model * model_tgt, llama_context * ctx_tgt); +common_speculative_init_result_ptr common_speculative_init_from_params(common_params & params, llama_model * model_tgt, llama_context * ctx_tgt); diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 3890fe5f0..df5666581 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -900,7 +900,7 @@ private: llama_model * model_dft = nullptr; llama_context * ctx_dft = nullptr; - common_init_speculative_result_ptr spec_init; + common_speculative_init_result_ptr spec_init; common_context_seq_rm_type ctx_tgt_seq_rm_type = COMMON_CONTEXT_SEQ_RM_TYPE_NO; common_context_seq_rm_type ctx_dft_seq_rm_type = COMMON_CONTEXT_SEQ_RM_TYPE_NO; @@ -1171,11 +1171,12 @@ private: { common_params params_dft = common_base_params_to_speculative(params_base); + // progress callback params_dft.load_progress_callback = load_progress_callback; params_dft.load_progress_callback_user_data = &load_progress_spec; - spec_init = common_init_speculative_from_params(params_dft, model_tgt, ctx_tgt); + spec_init = common_speculative_init_from_params(params_dft, model_tgt, ctx_tgt); model_dft = spec_init->model(); ctx_dft = spec_init->context(); @@ -2306,8 +2307,8 @@ private: // this is not true for SWA models: https://github.com/ggml-org/llama.cpp/pull/24411#issuecomment-4677983225 cur.update_pos(slot.prompt.n_tokens() - n_tokens_cur, pos_min, pos_max); - cur.update_tgt(ctx_tgt, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); - cur.update_dft(ctx_dft, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); + cur.update_tgt(ctx_tgt, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); + cur.update_dft(ctx_dft, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); // stash the draft's speculative state with the checkpoint common_speculative_get_state(spec.get(), slot.id, cur.data_spec); @@ -3264,8 +3265,8 @@ private: if (!do_reset) { // restore the context checkpoint - it->load_tgt(ctx_tgt, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); - it->load_dft(ctx_dft, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); + it->load_tgt(ctx_tgt, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); + it->load_dft(ctx_dft, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); // restore the draft's speculative state common_speculative_set_state(spec.get(), slot.id, it->data_spec); From bec4772f6a2527d371557b5d2032641e5ff7619c Mon Sep 17 00:00:00 2001 From: Pasha Khosravi Date: Tue, 7 Jul 2026 12:05:47 -0700 Subject: [PATCH 07/17] Add Q2_0 quantization: type definition and CPU backend (#24448) --- ggml/include/ggml.h | 4 +- ggml/src/ggml-common.h | 10 ++++ ggml/src/ggml-cpu/arch-fallback.h | 7 +++ ggml/src/ggml-cpu/arch/arm/quants.c | 74 ++++++++++++++++++++++++++++ ggml/src/ggml-cpu/ggml-cpu.c | 6 +++ ggml/src/ggml-cpu/ops.cpp | 7 +++ ggml/src/ggml-cpu/quants.c | 51 +++++++++++++++++++ ggml/src/ggml-cpu/quants.h | 3 ++ ggml/src/ggml-quants.c | 76 +++++++++++++++++++++++++++++ ggml/src/ggml-quants.h | 3 ++ ggml/src/ggml.c | 10 ++++ gguf-py/gguf/constants.py | 3 ++ include/llama.h | 1 + src/llama-model-loader.cpp | 2 + src/llama-quant.cpp | 4 +- tests/test-quantize-fns.cpp | 3 +- tools/quantize/quantize.cpp | 1 + 17 files changed, 262 insertions(+), 3 deletions(-) diff --git a/ggml/include/ggml.h b/ggml/include/ggml.h index d6807b6dd..ac133665d 100644 --- a/ggml/include/ggml.h +++ b/ggml/include/ggml.h @@ -429,7 +429,8 @@ extern "C" { GGML_TYPE_MXFP4 = 39, // MXFP4 (1 block) GGML_TYPE_NVFP4 = 40, // NVFP4 (4 blocks, E4M3 scale) GGML_TYPE_Q1_0 = 41, - GGML_TYPE_COUNT = 42, + GGML_TYPE_Q2_0 = 42, + GGML_TYPE_COUNT = 43, }; // precision @@ -473,6 +474,7 @@ extern "C" { GGML_FTYPE_MOSTLY_MXFP4 = 25, // except 1d tensors GGML_FTYPE_MOSTLY_NVFP4 = 26, // except 1d tensors GGML_FTYPE_MOSTLY_Q1_0 = 27, // except 1d tensors + GGML_FTYPE_MOSTLY_Q2_0 = 28, // except 1d tensors }; // available tensor operations: diff --git a/ggml/src/ggml-common.h b/ggml/src/ggml-common.h index a51b37dcd..83f9118da 100644 --- a/ggml/src/ggml-common.h +++ b/ggml/src/ggml-common.h @@ -96,6 +96,9 @@ typedef sycl::half2 ggml_half2; #define QI1_0 (QK1_0 / 32) #define QR1_0 1 +#define QI2_0 (QK2_0 / 32) +#define QR2_0 1 + #define QI4_0 (QK4_0 / (4 * QR4_0)) #define QR4_0 2 @@ -181,6 +184,13 @@ typedef struct { } block_q1_0; static_assert(sizeof(block_q1_0) == sizeof(ggml_half) + QK1_0 / 8, "wrong q1_0 block size/padding"); +#define QK2_0 64 +typedef struct { + ggml_half d; // delta (scale) + uint8_t qs[QK2_0 / 4]; // 2 bits per element +} block_q2_0; +static_assert(sizeof(block_q2_0) == sizeof(ggml_half) + QK2_0 / 4, "wrong q2_0 block size/padding"); + #define QK4_0 32 typedef struct { ggml_half d; // delta diff --git a/ggml/src/ggml-cpu/arch-fallback.h b/ggml/src/ggml-cpu/arch-fallback.h index da2ac98a4..152e0bac9 100644 --- a/ggml/src/ggml-cpu/arch-fallback.h +++ b/ggml/src/ggml-cpu/arch-fallback.h @@ -17,6 +17,7 @@ #define ggml_vec_dot_mxfp4_q8_0_generic ggml_vec_dot_mxfp4_q8_0 #define ggml_vec_dot_nvfp4_q8_0_generic ggml_vec_dot_nvfp4_q8_0 #define ggml_vec_dot_q1_0_q8_0_generic ggml_vec_dot_q1_0_q8_0 +#define ggml_vec_dot_q2_0_q8_0_generic ggml_vec_dot_q2_0_q8_0 #define ggml_vec_dot_tq1_0_q8_K_generic ggml_vec_dot_tq1_0_q8_K #define ggml_vec_dot_tq2_0_q8_K_generic ggml_vec_dot_tq2_0_q8_K #define ggml_vec_dot_q2_K_q8_K_generic ggml_vec_dot_q2_K_q8_K @@ -82,6 +83,7 @@ #define ggml_gemm_q2_K_8x8_q8_K_generic ggml_gemm_q2_K_8x8_q8_K #elif defined(__x86_64__) || defined(__i386__) || defined(_M_IX86) || defined(_M_X64) // quants.c +#define ggml_vec_dot_q2_0_q8_0_generic ggml_vec_dot_q2_0_q8_0 // repack.cpp #define ggml_quantize_mat_q8_0_4x4_generic ggml_quantize_mat_q8_0_4x4 #define ggml_quantize_mat_q8_K_4x4_generic ggml_quantize_mat_q8_K_4x4 @@ -113,6 +115,7 @@ #define quantize_row_q8_K_generic quantize_row_q8_K #define ggml_vec_dot_nvfp4_q8_0_generic ggml_vec_dot_nvfp4_q8_0 #define ggml_vec_dot_q1_0_q8_0_generic ggml_vec_dot_q1_0_q8_0 +#define ggml_vec_dot_q2_0_q8_0_generic ggml_vec_dot_q2_0_q8_0 #define ggml_vec_dot_tq1_0_q8_K_generic ggml_vec_dot_tq1_0_q8_K #define ggml_vec_dot_tq2_0_q8_K_generic ggml_vec_dot_tq2_0_q8_K #define ggml_vec_dot_iq1_m_q8_K_generic ggml_vec_dot_iq1_m_q8_K @@ -162,6 +165,7 @@ #define ggml_vec_dot_mxfp4_q8_0_generic ggml_vec_dot_mxfp4_q8_0 #define ggml_vec_dot_nvfp4_q8_0_generic ggml_vec_dot_nvfp4_q8_0 #define ggml_vec_dot_q1_0_q8_0_generic ggml_vec_dot_q1_0_q8_0 +#define ggml_vec_dot_q2_0_q8_0_generic ggml_vec_dot_q2_0_q8_0 // repack.cpp #define ggml_quantize_mat_q8_0_4x4_generic ggml_quantize_mat_q8_0_4x4 #define ggml_quantize_mat_q8_0_4x8_generic ggml_quantize_mat_q8_0_4x8 @@ -202,6 +206,7 @@ #elif defined(__riscv) // quants.c #define ggml_vec_dot_nvfp4_q8_0_generic ggml_vec_dot_nvfp4_q8_0 +#define ggml_vec_dot_q2_0_q8_0_generic ggml_vec_dot_q2_0_q8_0 // repack.cpp #define ggml_quantize_mat_q8_0_4x1_generic ggml_quantize_mat_q8_0_4x1 #define ggml_quantize_mat_q8_0_4x4_generic ggml_quantize_mat_q8_0_4x4 @@ -243,6 +248,7 @@ #define quantize_row_q8_K_generic quantize_row_q8_K #define ggml_vec_dot_nvfp4_q8_0_generic ggml_vec_dot_nvfp4_q8_0 #define ggml_vec_dot_q1_0_q8_0_generic ggml_vec_dot_q1_0_q8_0 +#define ggml_vec_dot_q2_0_q8_0_generic ggml_vec_dot_q2_0_q8_0 #define ggml_vec_dot_tq1_0_q8_K_generic ggml_vec_dot_tq1_0_q8_K #define ggml_vec_dot_tq2_0_q8_K_generic ggml_vec_dot_tq2_0_q8_K #define ggml_vec_dot_q2_K_q8_K_generic ggml_vec_dot_q2_K_q8_K @@ -306,6 +312,7 @@ #define ggml_vec_dot_mxfp4_q8_0_generic ggml_vec_dot_mxfp4_q8_0 #define ggml_vec_dot_nvfp4_q8_0_generic ggml_vec_dot_nvfp4_q8_0 #define ggml_vec_dot_q1_0_q8_0_generic ggml_vec_dot_q1_0_q8_0 +#define ggml_vec_dot_q2_0_q8_0_generic ggml_vec_dot_q2_0_q8_0 // repack.cpp #define ggml_quantize_mat_q8_0_4x4_generic ggml_quantize_mat_q8_0_4x4 #define ggml_quantize_mat_q8_0_4x8_generic ggml_quantize_mat_q8_0_4x8 diff --git a/ggml/src/ggml-cpu/arch/arm/quants.c b/ggml/src/ggml-cpu/arch/arm/quants.c index 445366797..636d7be12 100644 --- a/ggml/src/ggml-cpu/arch/arm/quants.c +++ b/ggml/src/ggml-cpu/arch/arm/quants.c @@ -219,6 +219,80 @@ void ggml_vec_dot_q1_0_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const voi #endif } +void ggml_vec_dot_q2_0_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) { + const int qk = QK2_0; + const int nb = n / qk; + + assert(n % qk == 0); + assert(nrc == 1); + UNUSED(nrc); + UNUSED(bx); + UNUSED(by); + UNUSED(bs); + + const block_q2_0 * GGML_RESTRICT x = vx; + const block_q8_0 * GGML_RESTRICT y = vy; + + float sumf = 0.0f; + +#if defined(__ARM_NEON) + // Replicate pattern: each byte repeated 4 times + static const uint8_t tbl_idx_lo[16] = {0,0,0,0, 1,1,1,1, 2,2,2,2, 3,3,3,3}; + static const uint8_t tbl_idx_hi[16] = {4,4,4,4, 5,5,5,5, 6,6,6,6, 7,7,7,7}; + // Right-shift amounts: 0,2,4,6 repeated for each group of 4 + static const int8_t shift_vals[16] = {0,-2,-4,-6, 0,-2,-4,-6, 0,-2,-4,-6, 0,-2,-4,-6}; + + const uint8x16_t idx_lo = vld1q_u8(tbl_idx_lo); + const uint8x16_t idx_hi = vld1q_u8(tbl_idx_hi); + const int8x16_t shifts = vld1q_s8(shift_vals); + const uint8x16_t mask2 = vdupq_n_u8(0x03); + const int8x16_t one = vdupq_n_s8(1); + + float32x4_t sumv = vdupq_n_f32(0.0f); + + for (int i = 0; i < nb; i++) { + const float d0 = GGML_CPU_FP16_TO_FP32(x[i].d); + + // group 64: one Q2_0 block (64 weights) maps to two Q8_0 blocks (2 * 32 = 64) + for (int k = 0; k < 2; k++) { + const block_q8_0 * GGML_RESTRICT yb = &y[i * 2 + k]; + const float d1 = GGML_CPU_FP16_TO_FP32(yb->d); + + // Load 8 bytes of packed 2-bit values + const uint8x8_t raw = vld1_u8(&x[i].qs[k * 8]); + const uint8x16_t raw16 = vcombine_u8(raw, raw); + + // First 16 elements: replicate bytes 0-3, shift, mask, subtract 1 + uint8x16_t bytes0 = vqtbl1q_u8(raw16, idx_lo); + int8x16_t qv0 = vsubq_s8( + vreinterpretq_s8_u8(vandq_u8(vshlq_u8(bytes0, shifts), mask2)), + one); + + // Second 16 elements: replicate bytes 4-7, shift, mask, subtract 1 + uint8x16_t bytes1 = vqtbl1q_u8(raw16, idx_hi); + int8x16_t qv1 = vsubq_s8( + vreinterpretq_s8_u8(vandq_u8(vshlq_u8(bytes1, shifts), mask2)), + one); + + // Load Q8_0 values and dot product + const int8x16_t y0 = vld1q_s8(yb->qs); + const int8x16_t y1 = vld1q_s8(yb->qs + 16); + + int32x4_t p0 = ggml_vdotq_s32(vdupq_n_s32(0), qv0, y0); + int32x4_t p1 = ggml_vdotq_s32(p0, qv1, y1); + + sumv = vmlaq_n_f32(sumv, vcvtq_f32_s32(p1), d0 * d1); + } + } + + sumf = vaddvq_f32(sumv); +#else + ggml_vec_dot_q2_0_q8_0_generic(n, s, bs, vx, bx, vy, by, nrc); + return; +#endif + + *s = sumf; +} void ggml_vec_dot_q4_0_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) { const int qk = QK8_0; diff --git a/ggml/src/ggml-cpu/ggml-cpu.c b/ggml/src/ggml-cpu/ggml-cpu.c index 8a0e5c255..a82842fcf 100644 --- a/ggml/src/ggml-cpu/ggml-cpu.c +++ b/ggml/src/ggml-cpu/ggml-cpu.c @@ -230,6 +230,12 @@ static const struct ggml_type_traits_cpu type_traits_cpu[GGML_TYPE_COUNT] = { .vec_dot_type = GGML_TYPE_Q8_0, .nrows = 1, }, + [GGML_TYPE_Q2_0] = { + .from_float = quantize_row_q2_0, + .vec_dot = ggml_vec_dot_q2_0_q8_0, + .vec_dot_type = GGML_TYPE_Q8_0, + .nrows = 1, + }, [GGML_TYPE_Q4_0] = { .from_float = quantize_row_q4_0, .vec_dot = ggml_vec_dot_q4_0_q8_0, diff --git a/ggml/src/ggml-cpu/ops.cpp b/ggml/src/ggml-cpu/ops.cpp index c555831ce..a735efff4 100644 --- a/ggml/src/ggml-cpu/ops.cpp +++ b/ggml/src/ggml-cpu/ops.cpp @@ -665,6 +665,7 @@ void ggml_compute_forward_add( ggml_compute_forward_add_non_quantized(params, dst); } break; case GGML_TYPE_Q1_0: + case GGML_TYPE_Q2_0: case GGML_TYPE_Q4_0: case GGML_TYPE_Q4_1: case GGML_TYPE_Q5_0: @@ -1115,6 +1116,7 @@ void ggml_compute_forward_add1( } } break; case GGML_TYPE_Q1_0: + case GGML_TYPE_Q2_0: case GGML_TYPE_Q4_0: case GGML_TYPE_Q4_1: case GGML_TYPE_Q5_0: @@ -1245,6 +1247,7 @@ void ggml_compute_forward_acc( case GGML_TYPE_F16: case GGML_TYPE_BF16: case GGML_TYPE_Q1_0: + case GGML_TYPE_Q2_0: case GGML_TYPE_Q4_0: case GGML_TYPE_Q4_1: case GGML_TYPE_Q5_0: @@ -4454,6 +4457,7 @@ void ggml_compute_forward_out_prod( switch (src0->type) { case GGML_TYPE_Q1_0: + case GGML_TYPE_Q2_0: case GGML_TYPE_Q4_0: case GGML_TYPE_Q4_1: case GGML_TYPE_Q5_0: @@ -4730,6 +4734,7 @@ void ggml_compute_forward_set( case GGML_TYPE_F16: case GGML_TYPE_BF16: case GGML_TYPE_Q1_0: + case GGML_TYPE_Q2_0: case GGML_TYPE_Q4_0: case GGML_TYPE_Q4_1: case GGML_TYPE_Q5_0: @@ -4954,6 +4959,7 @@ void ggml_compute_forward_get_rows( switch (src0->type) { case GGML_TYPE_Q1_0: + case GGML_TYPE_Q2_0: case GGML_TYPE_Q4_0: case GGML_TYPE_Q4_1: case GGML_TYPE_Q5_0: @@ -5680,6 +5686,7 @@ void ggml_compute_forward_clamp( } break; case GGML_TYPE_BF16: case GGML_TYPE_Q1_0: + case GGML_TYPE_Q2_0: case GGML_TYPE_Q4_0: case GGML_TYPE_Q4_1: case GGML_TYPE_Q5_0: diff --git a/ggml/src/ggml-cpu/quants.c b/ggml/src/ggml-cpu/quants.c index e5f9a4083..5e36459f8 100644 --- a/ggml/src/ggml-cpu/quants.c +++ b/ggml/src/ggml-cpu/quants.c @@ -26,6 +26,10 @@ void quantize_row_q1_0(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, in quantize_row_q1_0_ref(x, y, k); } +void quantize_row_q2_0(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k) { + quantize_row_q2_0_ref(x, y, k); +} + void quantize_row_q4_0(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k) { quantize_row_q4_0_ref(x, y, k); } @@ -170,6 +174,53 @@ void ggml_vec_dot_q1_0_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, c *s = sumf; } +void ggml_vec_dot_q2_0_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) { + const int qk = QK2_0; + const int nb = n / qk; + + assert(n % qk == 0); + assert(nrc == 1); + UNUSED(nrc); + UNUSED(bx); + UNUSED(by); + UNUSED(bs); + + const block_q2_0 * GGML_RESTRICT x = vx; + const block_q8_0 * GGML_RESTRICT y = vy; + + float sumf = 0.0f; + + for (int i = 0; i < nb; i++) { + const float d0 = GGML_CPU_FP16_TO_FP32(x[i].d); + + float sumi = 0.0f; + + // group 64: one Q2_0 block (64 weights) maps to two Q8_0 blocks (2 * 32 = 64) + for (int k = 0; k < 2; k++) { + const block_q8_0 * GGML_RESTRICT yb = &y[i * 2 + k]; + const float d1 = GGML_CPU_FP16_TO_FP32(yb->d); + int sumi_block = 0; + + const uint8_t * GGML_RESTRICT qs = &x[i].qs[k * 8]; + const int8_t * GGML_RESTRICT qy = yb->qs; + + for (int b = 0; b < 8; ++b) { + const uint8_t byte = qs[b]; + // Extract 4 two-bit values, map {0,1,2,3} -> {-1,0,1,2} + sumi_block += ((int)((byte >> 0) & 3) - 1) * qy[b*4 + 0]; + sumi_block += ((int)((byte >> 2) & 3) - 1) * qy[b*4 + 1]; + sumi_block += ((int)((byte >> 4) & 3) - 1) * qy[b*4 + 2]; + sumi_block += ((int)((byte >> 6) & 3) - 1) * qy[b*4 + 3]; + } + + sumi += d1 * sumi_block; + } + + sumf += d0 * sumi; + } + + *s = sumf; +} void ggml_vec_dot_q4_0_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) { const int qk = QK8_0; diff --git a/ggml/src/ggml-cpu/quants.h b/ggml/src/ggml-cpu/quants.h index d4bc87a1c..93ea7eeff 100644 --- a/ggml/src/ggml-cpu/quants.h +++ b/ggml/src/ggml-cpu/quants.h @@ -13,6 +13,7 @@ extern "C" { // Quantization void quantize_row_q1_0(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +void quantize_row_q2_0(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); void quantize_row_q4_0(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); void quantize_row_q4_1(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); void quantize_row_q5_0(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); @@ -38,6 +39,7 @@ void quantize_row_iq4_xs (const float * GGML_RESTRICT x, void * GGML_RESTRICT y, // Dot product void ggml_vec_dot_q1_0_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); +void ggml_vec_dot_q2_0_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); void ggml_vec_dot_q4_0_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); void ggml_vec_dot_q4_1_q8_1(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); void ggml_vec_dot_q5_0_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); @@ -71,6 +73,7 @@ void quantize_row_q8_0_generic(const float * GGML_RESTRICT x, void * GGML_RESTRI void quantize_row_q8_1_generic(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k); void quantize_row_q8_K_generic(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); void ggml_vec_dot_q1_0_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); +void ggml_vec_dot_q2_0_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); void ggml_vec_dot_q4_0_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); void ggml_vec_dot_q4_1_q8_1_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); void ggml_vec_dot_q5_0_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); diff --git a/ggml/src/ggml-quants.c b/ggml/src/ggml-quants.c index 15d231f70..1ebc50a76 100644 --- a/ggml/src/ggml-quants.c +++ b/ggml/src/ggml-quants.c @@ -71,6 +71,44 @@ void quantize_row_q1_0_ref(const float * GGML_RESTRICT x, block_q1_0 * GGML_REST } } +void quantize_row_q2_0_ref(const float * GGML_RESTRICT x, block_q2_0 * GGML_RESTRICT y, int64_t k) { + static const int qk = QK2_0; + + assert(k % qk == 0); + + const int nb = k / qk; + + for (int i = 0; i < nb; i++) { + // Compute scale as max absolute value in the block + float amax = 0.0f; + for (int j = 0; j < qk; j++) { + const float a = fabsf(x[i*qk + j]); + if (a > amax) amax = a; + } + const float d = amax; + const float id = d > 0.0f ? 1.0f / d : 0.0f; + + y[i].d = GGML_FP32_TO_FP16(d); + + // Clear quant bytes + for (int j = 0; j < qk / 4; ++j) { + y[i].qs[j] = 0; + } + + // Encode 2-bit values: round(w/d) clamped to [-1, 2], then add 1 + // 00 (-1) = -scale, 01 (0) = 0, 10 (+1) = +scale, 11 (+2) = 2*scale + for (int j = 0; j < qk; ++j) { + const float w = x[i*qk + j]; + int q = (int)roundf(w * id) + 1; + if (q < 0) q = 0; + if (q > 3) q = 3; + const int byte_index = j / 4; + const int bit_offset = (j % 4) * 2; + y[i].qs[byte_index] |= ((uint8_t)q << bit_offset); + } + } +} + // reference implementation for deterministic creation of model files void quantize_row_q4_0_ref(const float * GGML_RESTRICT x, block_q4_0 * GGML_RESTRICT y, int64_t k) { static const int qk = QK4_0; @@ -398,6 +436,26 @@ void dequantize_row_q1_0(const block_q1_0 * GGML_RESTRICT x, float * GGML_RESTRI } } +void dequantize_row_q2_0(const block_q2_0 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) { + static const int qk = QK2_0; + + assert(k % qk == 0); + + const int nb = k / qk; + + for (int i = 0; i < nb; i++) { + const float d = GGML_FP16_TO_FP32(x[i].d); + + for (int j = 0; j < qk; ++j) { + const int byte_index = j / 4; + const int bit_offset = (j % 4) * 2; + const uint8_t q = (x[i].qs[byte_index] >> bit_offset) & 0x03; + // 00=-1, 01=0, 10=+1, 11=+2 + y[i*qk + j] = ((int)q - 1) * d; + } + } +} + void dequantize_row_q4_0(const block_q4_0 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) { static const int qk = QK4_0; @@ -2052,6 +2110,20 @@ size_t quantize_q1_0(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, return nrow * row_size; } +size_t quantize_q2_0(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) { + if (!quant_weights) { + quantize_row_q2_0_ref(src, dst, (int64_t)nrow*n_per_row); + return nrow * ggml_row_size(GGML_TYPE_Q2_0, n_per_row); + } + size_t row_size = ggml_row_size(GGML_TYPE_Q2_0, n_per_row); + char * qrow = (char *)dst; + for (int64_t row = 0; row < nrow; ++row) { + quantize_row_q2_0_ref(src, (block_q2_0*)qrow, n_per_row); + src += n_per_row; + qrow += row_size; + } + return nrow * row_size; +} size_t quantize_q4_0(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrow, int64_t n_per_row, const float * quant_weights) { if (!quant_weights) { @@ -5461,6 +5533,10 @@ bool ggml_validate_row_data(enum ggml_type type, const void * data, size_t nbyte { VALIDATE_ROW_DATA_D_F16_IMPL(block_q1_0, data, nb); } break; + case GGML_TYPE_Q2_0: + { + VALIDATE_ROW_DATA_D_F16_IMPL(block_q2_0, data, nb); + } break; case GGML_TYPE_Q4_0: { VALIDATE_ROW_DATA_D_F16_IMPL(block_q4_0, data, nb); diff --git a/ggml/src/ggml-quants.h b/ggml/src/ggml-quants.h index d56c86da8..75188f1af 100644 --- a/ggml/src/ggml-quants.h +++ b/ggml/src/ggml-quants.h @@ -15,6 +15,7 @@ extern "C" { // Quantization GGML_API void quantize_row_q1_0_ref(const float * GGML_RESTRICT x, block_q1_0 * GGML_RESTRICT y, int64_t k); +GGML_API void quantize_row_q2_0_ref(const float * GGML_RESTRICT x, block_q2_0 * GGML_RESTRICT y, int64_t k); GGML_API void quantize_row_q4_0_ref(const float * GGML_RESTRICT x, block_q4_0 * GGML_RESTRICT y, int64_t k); GGML_API void quantize_row_q4_1_ref(const float * GGML_RESTRICT x, block_q4_1 * GGML_RESTRICT y, int64_t k); GGML_API void quantize_row_q5_0_ref(const float * GGML_RESTRICT x, block_q5_0 * GGML_RESTRICT y, int64_t k); @@ -43,6 +44,7 @@ GGML_API void quantize_row_iq2_s_ref (const float * GGML_RESTRICT x, block_iq2_ // Dequantization GGML_API void dequantize_row_q1_0(const block_q1_0 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); +GGML_API void dequantize_row_q2_0(const block_q2_0 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); GGML_API void dequantize_row_q4_0(const block_q4_0 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); GGML_API void dequantize_row_q4_1(const block_q4_1 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); GGML_API void dequantize_row_q5_0(const block_q5_0 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); @@ -93,6 +95,7 @@ GGML_API size_t quantize_q4_K(const float * GGML_RESTRICT src, void * GGML_RESTR GGML_API size_t quantize_q5_K(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix); GGML_API size_t quantize_q6_K(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix); GGML_API size_t quantize_q1_0(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix); +GGML_API size_t quantize_q2_0(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix); GGML_API size_t quantize_q4_0(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix); GGML_API size_t quantize_q4_1(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix); GGML_API size_t quantize_q5_0(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix); diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c index fc46ab0aa..3e78d3dc0 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -681,6 +681,14 @@ static const struct ggml_type_traits type_traits[GGML_TYPE_COUNT] = { .to_float = (ggml_to_float_t) dequantize_row_q1_0, .from_float_ref = (ggml_from_float_t) quantize_row_q1_0_ref, }, + [GGML_TYPE_Q2_0] = { + .type_name = "q2_0", + .blck_size = QK2_0, + .type_size = sizeof(block_q2_0), + .is_quantized = true, + .to_float = (ggml_to_float_t) dequantize_row_q2_0, + .from_float_ref = (ggml_from_float_t) quantize_row_q2_0_ref, + }, [GGML_TYPE_Q4_0] = { .type_name = "q4_0", .blck_size = QK4_0, @@ -1417,6 +1425,7 @@ enum ggml_type ggml_ftype_to_ggml_type(enum ggml_ftype ftype) { case GGML_FTYPE_MOSTLY_Q4_0: wtype = GGML_TYPE_Q4_0; break; case GGML_FTYPE_MOSTLY_Q4_1: wtype = GGML_TYPE_Q4_1; break; case GGML_FTYPE_MOSTLY_Q1_0: wtype = GGML_TYPE_Q1_0; break; + case GGML_FTYPE_MOSTLY_Q2_0: wtype = GGML_TYPE_Q2_0; break; case GGML_FTYPE_MOSTLY_Q5_0: wtype = GGML_TYPE_Q5_0; break; case GGML_FTYPE_MOSTLY_Q5_1: wtype = GGML_TYPE_Q5_1; break; case GGML_FTYPE_MOSTLY_Q8_0: wtype = GGML_TYPE_Q8_0; break; @@ -7744,6 +7753,7 @@ size_t ggml_quantize_chunk( switch (type) { case GGML_TYPE_Q1_0: result = quantize_q1_0 (src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; + case GGML_TYPE_Q2_0: result = quantize_q2_0 (src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; case GGML_TYPE_Q4_0: result = quantize_q4_0 (src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; case GGML_TYPE_Q4_1: result = quantize_q4_1 (src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; case GGML_TYPE_Q5_0: result = quantize_q5_0 (src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index cd4cdef89..869e436ac 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -4533,6 +4533,7 @@ class GGMLQuantizationType(IntEnum): MXFP4 = 39 NVFP4 = 40 Q1_0 = 41 + Q2_0 = 42 class ExpertGatingFuncType(IntEnum): @@ -4588,6 +4589,7 @@ class LlamaFileType(IntEnum): MOSTLY_MXFP4_MOE = 38 # except 1d tensors MOSTLY_NVFP4 = 39 # except 1d tensors MOSTLY_Q1_0 = 40 # except 1d tensors + MOSTLY_Q2_0 = 41 # except 1d tensors GUESSED = 1024 # not specified in the model file @@ -4713,6 +4715,7 @@ GGML_QUANT_SIZES: dict[GGMLQuantizationType, tuple[int, int]] = { GGMLQuantizationType.MXFP4: (32, 1 + 16), GGMLQuantizationType.NVFP4: (64, 4 + 32), GGMLQuantizationType.Q1_0: (128, 2 + 16), + GGMLQuantizationType.Q2_0: (64, 2 + 16), } diff --git a/include/llama.h b/include/llama.h index b89ab758b..a311ac202 100644 --- a/include/llama.h +++ b/include/llama.h @@ -155,6 +155,7 @@ extern "C" { LLAMA_FTYPE_MOSTLY_MXFP4_MOE = 38, // except 1d tensors LLAMA_FTYPE_MOSTLY_NVFP4 = 39, // except 1d tensors LLAMA_FTYPE_MOSTLY_Q1_0 = 40, // except 1d tensors + LLAMA_FTYPE_MOSTLY_Q2_0 = 41, // except 1d tensors LLAMA_FTYPE_GUESSED = 1024, // not specified in the model file }; diff --git a/src/llama-model-loader.cpp b/src/llama-model-loader.cpp index 55554735d..28f8bb793 100644 --- a/src/llama-model-loader.cpp +++ b/src/llama-model-loader.cpp @@ -37,6 +37,7 @@ const char * llama_ftype_name(llama_ftype ftype) { case LLAMA_FTYPE_MOSTLY_F16: name = LLAMA_FTYPE_PREFIX "F16"; break; case LLAMA_FTYPE_MOSTLY_BF16: name = LLAMA_FTYPE_PREFIX "BF16"; break; case LLAMA_FTYPE_MOSTLY_Q1_0: name = LLAMA_FTYPE_PREFIX "Q1_0"; break; + case LLAMA_FTYPE_MOSTLY_Q2_0: name = LLAMA_FTYPE_PREFIX "Q2_0"; break; case LLAMA_FTYPE_MOSTLY_Q4_0: name = LLAMA_FTYPE_PREFIX "Q4_0"; break; case LLAMA_FTYPE_MOSTLY_Q4_1: name = LLAMA_FTYPE_PREFIX "Q4_1"; break; case LLAMA_FTYPE_MOSTLY_Q5_0: name = LLAMA_FTYPE_PREFIX "Q5_0"; break; @@ -767,6 +768,7 @@ llama_model_loader::llama_model_loader( case GGML_TYPE_IQ3_S: ftype = LLAMA_FTYPE_MOSTLY_IQ3_S; break; case GGML_TYPE_NVFP4: ftype = LLAMA_FTYPE_MOSTLY_NVFP4; break; case GGML_TYPE_Q1_0: ftype = LLAMA_FTYPE_MOSTLY_Q1_0; break; + case GGML_TYPE_Q2_0: ftype = LLAMA_FTYPE_MOSTLY_Q2_0; break; default: { LLAMA_LOG_WARN("%s: unknown type %s\n", __func__, ggml_type_name(type_max)); diff --git a/src/llama-quant.cpp b/src/llama-quant.cpp index 847e79f46..aebbc1ffb 100644 --- a/src/llama-quant.cpp +++ b/src/llama-quant.cpp @@ -380,6 +380,7 @@ static ggml_type tensor_type_fallback(quantize_state_impl & qs, const ggml_tenso case GGML_TYPE_IQ3_XXS: case GGML_TYPE_IQ3_S: // types on the right: block size 32 case GGML_TYPE_IQ4_XS: return_type = GGML_TYPE_IQ4_NL; break; + case GGML_TYPE_Q2_0: case GGML_TYPE_Q2_K: case GGML_TYPE_Q3_K: case GGML_TYPE_TQ1_0: @@ -480,7 +481,7 @@ static ggml_type llama_tensor_get_type_impl(quantize_state_impl & qs, ggml_type else if (ftype == LLAMA_FTYPE_MOSTLY_IQ3_XXS) { new_type = GGML_TYPE_IQ3_S; } - else if (ftype == LLAMA_FTYPE_MOSTLY_TQ1_0 || ftype == LLAMA_FTYPE_MOSTLY_TQ2_0) { + else if (ftype == LLAMA_FTYPE_MOSTLY_TQ1_0 || ftype == LLAMA_FTYPE_MOSTLY_TQ2_0 || ftype == LLAMA_FTYPE_MOSTLY_Q2_0) { new_type = GGML_TYPE_Q4_K; } } @@ -800,6 +801,7 @@ ggml_type llama_ftype_get_default_type(llama_ftype ftype) { case LLAMA_FTYPE_MOSTLY_BF16: return GGML_TYPE_BF16; case LLAMA_FTYPE_ALL_F32: return GGML_TYPE_F32; case LLAMA_FTYPE_MOSTLY_Q1_0: return GGML_TYPE_Q1_0; + case LLAMA_FTYPE_MOSTLY_Q2_0: return GGML_TYPE_Q2_0; case LLAMA_FTYPE_MOSTLY_MXFP4_MOE: return GGML_TYPE_MXFP4; diff --git a/tests/test-quantize-fns.cpp b/tests/test-quantize-fns.cpp index 74978e446..9510ac14c 100644 --- a/tests/test-quantize-fns.cpp +++ b/tests/test-quantize-fns.cpp @@ -158,6 +158,7 @@ static int test_vec_dot_q(bool verbose) { type == GGML_TYPE_Q1_0 ? MAX_QUANTIZATION_TOTAL_ERROR_BINARY : type == GGML_TYPE_TQ1_0 ? MAX_QUANTIZATION_TOTAL_ERROR_TERNARY : type == GGML_TYPE_TQ2_0 ? MAX_QUANTIZATION_TOTAL_ERROR_TERNARY : + type == GGML_TYPE_Q2_0 ? MAX_QUANTIZATION_TOTAL_ERROR_TERNARY : type == GGML_TYPE_Q2_K ? MAX_QUANTIZATION_TOTAL_ERROR_2BITS : type == GGML_TYPE_IQ2_S ? MAX_QUANTIZATION_TOTAL_ERROR_2BITS : type == GGML_TYPE_Q3_K ? MAX_QUANTIZATION_TOTAL_ERROR_3BITS : @@ -183,7 +184,7 @@ static int test_vec_dot_q(bool verbose) { ? MAX_DOT_PRODUCT_ERROR_LOWBIT : type == GGML_TYPE_Q1_0 ? MAX_DOT_PRODUCT_ERROR_BINARY - : type == GGML_TYPE_TQ1_0 || type == GGML_TYPE_TQ2_0 + : type == GGML_TYPE_TQ1_0 || type == GGML_TYPE_TQ2_0 || type == GGML_TYPE_Q2_0 ? MAX_DOT_PRODUCT_ERROR_TERNARY : type == GGML_TYPE_NVFP4 ? MAX_DOT_PRODUCT_ERROR_FP4 diff --git a/tools/quantize/quantize.cpp b/tools/quantize/quantize.cpp index 840eefc2f..15ef64c4b 100644 --- a/tools/quantize/quantize.cpp +++ b/tools/quantize/quantize.cpp @@ -33,6 +33,7 @@ struct quant_option { static const std::vector QUANT_OPTIONS = { { "Q1_0", LLAMA_FTYPE_MOSTLY_Q1_0, " 1.125 bpw quantization", }, + { "Q2_0", LLAMA_FTYPE_MOSTLY_Q2_0, " 2.25 bpw quantization (group 64)", }, { "Q4_0", LLAMA_FTYPE_MOSTLY_Q4_0, " 4.34G, +0.4685 ppl @ Llama-3-8B", }, { "Q4_1", LLAMA_FTYPE_MOSTLY_Q4_1, " 4.78G, +0.4511 ppl @ Llama-3-8B", }, { "MXFP4_MOE",LLAMA_FTYPE_MOSTLY_MXFP4_MOE," MXFP4 MoE", }, From 931ca30bef7ef930b1fecb573b2f44dc60c1b01a Mon Sep 17 00:00:00 2001 From: lhez Date: Tue, 7 Jul 2026 20:34:29 -0700 Subject: [PATCH 08/17] opencl: fix potential crash in aos reconstruct (#25383) --- ggml/src/ggml-opencl/ggml-opencl.cpp | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/ggml/src/ggml-opencl/ggml-opencl.cpp b/ggml/src/ggml-opencl/ggml-opencl.cpp index 568eef313..9d1d10678 100644 --- a/ggml/src/ggml-opencl/ggml-opencl.cpp +++ b/ggml/src/ggml-opencl/ggml-opencl.cpp @@ -16653,6 +16653,7 @@ static cl_mem ggml_cl_mul_mat_dequant_quant_to_f16( ? ggml_cl_is_q4_0_soa(tensor) : ggml_cl_is_q8_0_soa(tensor); + cl_mem aos = nullptr; if (is_soa) { // Reconstruct full parent AoS; view's own nb[] then index it correctly. const ggml_tensor * parent = tensor->view_src ? tensor->view_src : tensor; @@ -16664,7 +16665,7 @@ static cl_mem ggml_cl_mul_mat_dequant_quant_to_f16( const size_t parent_nbytes = (size_t) ggml_nelements(parent) / blck_size * block_bytes; cl_int err; - cl_mem aos = clCreateBuffer(backend_ctx->context, CL_MEM_READ_WRITE, parent_nbytes, NULL, &err); + aos = clCreateBuffer(backend_ctx->context, CL_MEM_READ_WRITE, parent_nbytes, NULL, &err); CL_CHECK(err); // large q4_0/q8_0 WEIGHTS are stored transposed and small weights @@ -16751,9 +16752,6 @@ static cl_mem ggml_cl_mul_mat_dequant_quant_to_f16( if (extra_reconstruct) { *extra_reconstruct = aos; - } else { - // OpenCL retains the memobj while queued kernels reference it. - CL_CHECK(clReleaseMemObject(aos)); } } else { auto * extra = (ggml_tensor_extra_cl *) tensor->extra; @@ -16817,6 +16815,13 @@ static cl_mem ggml_cl_mul_mat_dequant_quant_to_f16( size_t lws[3] = { 1, 1, 1 }; CL_CHECK(clEnqueueNDRangeKernel(backend_ctx->queue, dq_kernel, 3, NULL, gws, lws, 0, NULL, NULL)); + // release the reconstructed aos if + // 1. it was actually reconstructed + // 2. the caller didn't request it to be returned + // src_buf may refer to aos, so we should release after this enqueue + if (aos && !extra_reconstruct) { + CL_CHECK(clReleaseMemObject(aos)); + } return out; } From 68a521b591edd2f36a456809230d63aa81003dfc Mon Sep 17 00:00:00 2001 From: fairydreaming <166155368+fairydreaming@users.noreply.github.com> Date: Wed, 8 Jul 2026 05:46:28 +0200 Subject: [PATCH 09/17] ggml : add support for CPU f16->f16 GGML_OP_SET_ROWS (#25344) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * ggml : add support for CPU f16->f16 GGML_OP_SET_ROWS * ggml : add missing type checks in f16 GGML_OP_SET_ROWS * ggml : merge ggml_compute_forward_set_rows_f32() and ggml_compute_forward_set_rows_f16() into ggml_compute_forward_set_rows_impl() * chore : replace assert() with GGML_ASSERT() --------- Co-authored-by: Stanisław Szymczyk --- ggml/src/ggml-cpu/ops.cpp | 41 ++++++++++++++++++++++++++++------- ggml/src/ggml.c | 2 +- tests/test-backend-ops.cpp | 44 ++++++++++++++++++++++---------------- 3 files changed, 59 insertions(+), 28 deletions(-) diff --git a/ggml/src/ggml-cpu/ops.cpp b/ggml/src/ggml-cpu/ops.cpp index a735efff4..fde939b4a 100644 --- a/ggml/src/ggml-cpu/ops.cpp +++ b/ggml/src/ggml-cpu/ops.cpp @@ -5025,8 +5025,8 @@ void ggml_compute_forward_get_rows( //} } -template -static void ggml_compute_forward_set_rows_f32( +template +static void ggml_compute_forward_set_rows_impl( const ggml_compute_params * params, ggml_tensor * dst) { @@ -5041,7 +5041,7 @@ static void ggml_compute_forward_set_rows_f32( assert(ne0 == nc); assert(ne2 == ne02); assert(ne3 == ne03); - assert(src0->type == GGML_TYPE_F32); + GGML_ASSERT(src0->type == GGML_TYPE_F32 || (src0->type == GGML_TYPE_F16 && dst->type == GGML_TYPE_F16)); assert(ne02 % ne11 == 0); assert(ne03 % ne12 == 0); @@ -5055,6 +5055,8 @@ static void ggml_compute_forward_set_rows_f32( const int64_t ir0 = dr*ith; const int64_t ir1 = std::min(ir0 + dr, nr); + const size_t rs = ggml_row_size(src0->type, nc); + ggml_from_float_t const from_float = ggml_get_type_traits_cpu(dst->type)->from_float; for (int64_t i03 = 0; i03 < ne03; ++i03) { @@ -5068,9 +5070,18 @@ static void ggml_compute_forward_set_rows_f32( GGML_ASSERT(i1 >= 0 && i1 < ne1); - from_float( - (const float *) ((char *) src0->data + i*nb01 + i02*nb02 + i03*nb03), - ((char *) dst->data + i1*nb1 + i02*nb2 + i03*nb3), nc); + if constexpr (std::is_same_v) { + from_float( + (const float *) ((char *) src0->data + i*nb01 + i02*nb02 + i03*nb03), + ((char *) dst->data + i1*nb1 + i02*nb2 + i03*nb3), nc); + } else if constexpr (std::is_same_v) { + memcpy( + ((char *) dst->data + i1*nb1 + i02*nb2 + i03*nb3), + ((char *) src0->data + i*nb01 + i02*nb02 + i03*nb03), + rs); + } else { + GGML_ABORT("src0->type = %d (%s) not supported", src0->type, ggml_type_name(src0->type)); + } } } } @@ -5087,13 +5098,27 @@ void ggml_compute_forward_set_rows( case GGML_TYPE_F32: { if (src1->type == GGML_TYPE_I64) { - ggml_compute_forward_set_rows_f32(params, dst); + ggml_compute_forward_set_rows_impl(params, dst); } else if (src1->type == GGML_TYPE_I32) { - ggml_compute_forward_set_rows_f32(params, dst); + ggml_compute_forward_set_rows_impl(params, dst); } else { GGML_ABORT("src1->type = %d (%s) not supported", src1->type, ggml_type_name(src1->type)); } } break; + case GGML_TYPE_F16: + { + if (dst->type == GGML_TYPE_F16) { + if (src1->type == GGML_TYPE_I64) { + ggml_compute_forward_set_rows_impl(params, dst); + } else if (src1->type == GGML_TYPE_I32) { + ggml_compute_forward_set_rows_impl(params, dst); + } else { + GGML_ABORT("src1->type = %d (%s) not supported", src1->type, ggml_type_name(src1->type)); + } + } else { + GGML_ABORT("dst->type = %d (%s) not supported with src0->type = %d (%s)", dst->type, ggml_type_name(dst->type), src0->type, ggml_type_name(src0->type)); + } + } break; default: { GGML_ABORT("src0->type = %d (%s) not supported", src0->type, ggml_type_name(src0->type)); diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c index 3e78d3dc0..de0321d9f 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -3926,7 +3926,7 @@ struct ggml_tensor * ggml_set_rows( GGML_ASSERT(b->ne[2] % c->ne[1] == 0); GGML_ASSERT(b->ne[3] % c->ne[2] == 0); GGML_ASSERT(c->ne[3] == 1); - GGML_ASSERT(b->type == GGML_TYPE_F32); + GGML_ASSERT(b->type == GGML_TYPE_F32 || b->type == GGML_TYPE_F16); GGML_ASSERT(c->type == GGML_TYPE_I64 || c->type == GGML_TYPE_I32); GGML_ASSERT(ggml_is_contiguous_rows(a)); diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 82a9d9546..0c0aac2bb 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -2393,7 +2393,8 @@ static void init_set_rows_row_ids(ggml_tensor * t, int num_rows) { // GGML_OP_SET_ROWS struct test_set_rows : public test_case { - const ggml_type type; + const ggml_type type_src; + const ggml_type type_dst; const ggml_type type_idx; const std::array ne; const std::array nr23; // broadcast only dims 2 and 3 @@ -2401,21 +2402,22 @@ struct test_set_rows : public test_case { const bool v; // view (non-contiguous src1) std::string vars() override { - return VARS_TO_STR6(type, type_idx, ne, nr23, r, v); + return VARS_TO_STR7(type_src, type_dst, type_idx, ne, nr23, r, v); } - test_set_rows(ggml_type type, + test_set_rows(ggml_type type_src, + ggml_type type_dst, ggml_type type_idx, std::array ne, std::array nr23, int r, bool v = false) - : type(type), type_idx(type_idx), ne(ne), nr23(nr23), r(r), v(v) {} + : type_src(type_src), type_dst(type_dst), type_idx(type_idx), ne(ne), nr23(nr23), r(r), v(v) {} ggml_tensor * build_graph(ggml_context * ctx) override { - ggml_tensor * dst = ggml_new_tensor_4d(ctx, type, ne[0], ne[1], ne[2]*nr23[0], ne[3]*nr23[1]); + ggml_tensor * dst = ggml_new_tensor_4d(ctx, type_dst, ne[0], ne[1], ne[2]*nr23[0], ne[3]*nr23[1]); ggml_set_name(dst, "dst"); - ggml_tensor * src = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, ne[0], r, ne[2]*nr23[0], ne[3]*nr23[1]); + ggml_tensor * src = ggml_new_tensor_4d(ctx, type_src, ne[0], r, ne[2]*nr23[0], ne[3]*nr23[1]); ggml_set_name(src, "src"); ggml_tensor * row_idxs = ggml_new_tensor_3d(ctx, type_idx, r, ne[2], ne[3]); @@ -2448,17 +2450,17 @@ struct test_set_rows : public test_case { } double max_nmse_err() override { - if (type == GGML_TYPE_Q4_0 || type == GGML_TYPE_Q4_1 || type == GGML_TYPE_IQ4_NL || - type == GGML_TYPE_Q5_0 || type == GGML_TYPE_Q5_1 || type == GGML_TYPE_Q8_0) { + if (type_dst == GGML_TYPE_Q4_0 || type_dst == GGML_TYPE_Q4_1 || type_dst == GGML_TYPE_IQ4_NL || + type_dst == GGML_TYPE_Q5_0 || type_dst == GGML_TYPE_Q5_1 || type_dst == GGML_TYPE_Q8_0) { // estimate what the max nmse error would be if one quantized value is // off by one. The test values are distributed in [-1,1], so it'll be // roughly (2.0 / 2^bits)^2, divided by the mean square value of the reference, // which is roughly 0.25 times the number of elements. double err_estimate = 1.0f/8.0f; - if (type == GGML_TYPE_Q5_0 || type == GGML_TYPE_Q5_1) { + if (type_dst == GGML_TYPE_Q5_0 || type_dst == GGML_TYPE_Q5_1) { err_estimate /= 2.0f; } - if (type == GGML_TYPE_Q8_0) { + if (type_dst == GGML_TYPE_Q8_0) { err_estimate /= 8.0f; } err_estimate *= err_estimate; @@ -2471,7 +2473,7 @@ struct test_set_rows : public test_case { // See dicussion here: https://github.com/ggml-org/llama.cpp/pull/23760#issuecomment-4566312209 double max_nmse_err(ggml_backend_t backend) override { ggml_backend_reg_t reg = ggml_backend_dev_backend_reg(ggml_backend_get_device(backend)); - if (type == GGML_TYPE_Q8_0 && strcmp(ggml_backend_reg_name(reg), "WebGPU") == 0) { + if (type_dst == GGML_TYPE_Q8_0 && strcmp(ggml_backend_reg_name(reg), "WebGPU") == 0) { return std::max(test_case::max_nmse_err(backend), 2e-7); } return test_case::max_nmse_err(backend); @@ -7873,24 +7875,28 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_get_rows_back(GGML_TYPE_I32, 256, 5, 4, 1, v)); } - test_cases.emplace_back(new test_set_rows(GGML_TYPE_F32, GGML_TYPE_I64, { 1, 8, 1, 3 }, { 1, 1 }, 2, false)); - test_cases.emplace_back(new test_set_rows(GGML_TYPE_F32, GGML_TYPE_I32, { 1, 8, 1, 3 }, { 1, 1 }, 2, false)); - test_cases.emplace_back(new test_set_rows(GGML_TYPE_Q8_0, GGML_TYPE_I32, { 256, 5, 1, 3 }, { 1, 1, }, 1, false)); + test_cases.emplace_back(new test_set_rows(GGML_TYPE_F32, GGML_TYPE_F32, GGML_TYPE_I64, { 1, 8, 1, 3 }, { 1, 1 }, 2, false)); + test_cases.emplace_back(new test_set_rows(GGML_TYPE_F32, GGML_TYPE_F32, GGML_TYPE_I32, { 1, 8, 1, 3 }, { 1, 1 }, 2, false)); + test_cases.emplace_back(new test_set_rows(GGML_TYPE_F32, GGML_TYPE_Q8_0, GGML_TYPE_I32, { 256, 5, 1, 3 }, { 1, 1, }, 1, false)); for (ggml_type type : all_types) { for (int b : {1, 7}) { for (bool v : {false, true}) { - test_cases.emplace_back(new test_set_rows(type, GGML_TYPE_I64, { 256, 5, b, 3 }, { 1, 1, }, 1, v)); - test_cases.emplace_back(new test_set_rows(type, GGML_TYPE_I64, { 256, 11, 1, b }, { 2, 3, }, 7, v)); + test_cases.emplace_back(new test_set_rows(GGML_TYPE_F32, type, GGML_TYPE_I64, { 256, 5, b, 3 }, { 1, 1, }, 1, v)); + test_cases.emplace_back(new test_set_rows(GGML_TYPE_F32, type, GGML_TYPE_I64, { 256, 11, 1, b }, { 2, 3, }, 7, v)); - test_cases.emplace_back(new test_set_rows(type, GGML_TYPE_I64, { 3*ggml_blck_size(type), 3, b, 1 }, { 2, 3, }, 2, v)); + test_cases.emplace_back(new test_set_rows(GGML_TYPE_F32, type, GGML_TYPE_I64, { 3*ggml_blck_size(type), 3, b, 1 }, { 2, 3, }, 2, v)); if (ggml_blck_size(type) == 1) { - test_cases.emplace_back(new test_set_rows(type, GGML_TYPE_I64, { 31, 3, b, 1 }, { 2, 3, }, 2, v)); - test_cases.emplace_back(new test_set_rows(type, GGML_TYPE_I64, { 33, 5, 1, b }, { 2, 3, }, 1, v)); + test_cases.emplace_back(new test_set_rows(GGML_TYPE_F32, type, GGML_TYPE_I64, { 31, 3, b, 1 }, { 2, 3, }, 2, v)); + test_cases.emplace_back(new test_set_rows(GGML_TYPE_F32, type, GGML_TYPE_I64, { 33, 5, 1, b }, { 2, 3, }, 1, v)); } } } } + test_cases.emplace_back(new test_set_rows(GGML_TYPE_F16, GGML_TYPE_F16, GGML_TYPE_I64, { 1, 8, 1, 3 }, { 1, 1 }, 2, false)); + test_cases.emplace_back(new test_set_rows(GGML_TYPE_F16, GGML_TYPE_F16, GGML_TYPE_I32, { 1, 8, 1, 3 }, { 1, 1 }, 2, false)); + test_cases.emplace_back(new test_set_rows(GGML_TYPE_F16, GGML_TYPE_F16, GGML_TYPE_I64, { 1, 8, 1, 3 }, { 1, 1 }, 2, true)); + test_cases.emplace_back(new test_set_rows(GGML_TYPE_F16, GGML_TYPE_F16, GGML_TYPE_I32, { 1, 8, 1, 3 }, { 1, 1 }, 2, true)); for (int mode : { GGML_ROPE_TYPE_NORMAL, GGML_ROPE_TYPE_NEOX, GGML_ROPE_TYPE_MROPE, GGML_ROPE_TYPE_VISION }) { for (ggml_type type : {GGML_TYPE_F16, GGML_TYPE_F32}) { From 57b50e1f6b50e01eb14c43fc1253602af74c1870 Mon Sep 17 00:00:00 2001 From: tyronecai Date: Wed, 8 Jul 2026 13:00:05 +0800 Subject: [PATCH 10/17] ggml : fix A indexing in simd_gemm scalar tail-column path (#25390) `simd_gemm()` has an incorrect A-matrix index in the scalar tail-column path for full row blocks. --- ggml/src/ggml-cpu/simd-gemm.h | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ggml/src/ggml-cpu/simd-gemm.h b/ggml/src/ggml-cpu/simd-gemm.h index 4119d04f8..2ebd10051 100644 --- a/ggml/src/ggml-cpu/simd-gemm.h +++ b/ggml/src/ggml-cpu/simd-gemm.h @@ -78,7 +78,7 @@ static void simd_gemm( for (int64_t i = 0; i < GEMM_RM; i++) { float a = C[i * N + jj]; for (int64_t kk = 0; kk < K; kk++) { - a += A[i + kk] * B[kk * N + jj]; + a += A[i * K + kk] * B[kk * N + jj]; } C[i * N + jj] = a; } From 4a7ee3126d529618b330269f7a68c53d1ca012c1 Mon Sep 17 00:00:00 2001 From: hourhl <67227355+hourhl@users.noreply.github.com> Date: Wed, 8 Jul 2026 13:02:09 +0800 Subject: [PATCH 11/17] fix: OOB reads in UGM tokenizer (precompiled_charsmap handling) (#18750) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix: OOB reads in UGM tokenizer (precompiled_charsmap handling) - Validate minimum size (4 bytes) before reading xcda_blob_size - Use strnlen with bounds check instead of unsafe strlen Both issues allow heap-buffer-overflow from malicious T5/UGM GGUF files. * Replace unsafe strnlen() with a bounds-checked loop that scans for \0 within the remaining array size. * move bounds checks to load * typo merge fix --------- Co-authored-by: hourhl Co-authored-by: Sigbjørn Skjæret <1629204+CISC@users.noreply.github.com> --- src/llama-vocab.cpp | 26 +++++++++++++++++++------- 1 file changed, 19 insertions(+), 7 deletions(-) diff --git a/src/llama-vocab.cpp b/src/llama-vocab.cpp index 6e78a3f6c..fdd447147 100644 --- a/src/llama-vocab.cpp +++ b/src/llama-vocab.cpp @@ -887,9 +887,6 @@ struct llm_tokenizer_ugm : llm_tokenizer { // blob containing XOR-compressed compact double array (XCDA) entries uint32_t xcda_blob_size = *(const uint32_t *) &precompiled_charsmap[0]; charsmap_offset += sizeof(xcda_blob_size); - if (xcda_blob_size + charsmap_offset >= precompiled_charsmap.size()) { - throw std::runtime_error("Index out of array bounds in precompiled charsmap!"); - } // Next xcda_blob_size bytes contain entries of XOR-compressed compact // double array (XCDA). Each entry is bit-packed into a 32-bit integer. @@ -1205,7 +1202,15 @@ private: throw std::runtime_error("Index out of array bounds in precompiled charsmap!"); } const char * prefix_replacement = &(tokenizer.prefix_replacements)[longest_prefix_offset]; - return { prefix_replacement, strlen(prefix_replacement), longest_prefix_length }; + size_t max_len = tokenizer.prefix_replacements_size - longest_prefix_offset; + size_t repl_len = 0; + while (repl_len < max_len && prefix_replacement[repl_len] != '\0') { + repl_len++; + } + if (repl_len == max_len) { + throw std::runtime_error("Unterminated string in precompiled charsmap!"); + } + return { prefix_replacement, repl_len, longest_prefix_length }; } // check if the input prefix contains a valid sequence of UTF-8 code units @@ -2018,11 +2023,18 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) { const size_t n_precompiled_charsmap = gguf_get_arr_n(ctx, precompiled_charsmap_keyidx); const char * pc = (const char *) gguf_get_arr_data(ctx, precompiled_charsmap_keyidx); precompiled_charsmap.assign(pc, pc + n_precompiled_charsmap); + if (precompiled_charsmap.size() < sizeof(uint32_t)) { + throw std::runtime_error("precompiled_charsmap too small for xcda_blob_size header!"); + } + uint32_t * xcda_blob_size = (uint32_t *) &precompiled_charsmap[0]; +#if defined(__BYTE_ORDER__) && defined(__ORDER_BIG_ENDIAN__) && __BYTE_ORDER__ == __ORDER_BIG_ENDIAN__ + *xcda_blob_size = __builtin_bswap32(*xcda_blob_size); +#endif + if (*xcda_blob_size + sizeof(uint32_t) >= precompiled_charsmap.size()) { + throw std::runtime_error("Index out of array bounds in precompiled charsmap!"); + } #if defined(__BYTE_ORDER__) && defined(__ORDER_BIG_ENDIAN__) && __BYTE_ORDER__ == __ORDER_BIG_ENDIAN__ // correct endianness of data in precompiled_charsmap binary blob - uint32_t * xcda_blob_size = (uint32_t *) &precompiled_charsmap[0]; - *xcda_blob_size = __builtin_bswap32(*xcda_blob_size); - assert(*xcda_blob_size + sizeof(uint32_t) < n_precompiled_charsmap); size_t xcda_array_size = *xcda_blob_size / sizeof(uint32_t); uint32_t * xcda_array = (uint32_t *) &precompiled_charsmap[sizeof(uint32_t)]; for (size_t i = 0; i < xcda_array_size; ++i) { From 0512ef1e5a7f76297c72beae1383b70632048008 Mon Sep 17 00:00:00 2001 From: fairydreaming <166155368+fairydreaming@users.noreply.github.com> Date: Wed, 8 Jul 2026 08:49:07 +0200 Subject: [PATCH 12/17] metal : add set_rows with src0 f16 (#25434) Co-authored-by: Georgi Gerganov --- ggml/src/ggml-metal/ggml-metal-device.cpp | 8 +- ggml/src/ggml-metal/ggml-metal-device.h | 2 +- ggml/src/ggml-metal/ggml-metal-device.m | 2 +- ggml/src/ggml-metal/ggml-metal-ops.cpp | 2 +- ggml/src/ggml-metal/ggml-metal.metal | 149 +++++++++++----------- 5 files changed, 83 insertions(+), 80 deletions(-) diff --git a/ggml/src/ggml-metal/ggml-metal-device.cpp b/ggml/src/ggml-metal/ggml-metal-device.cpp index b2478fe97..e7ac21a2b 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.cpp +++ b/ggml/src/ggml-metal/ggml-metal-device.cpp @@ -160,11 +160,15 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_get_rows(ggml_me return res; } -ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_set_rows(ggml_metal_library_t lib, ggml_type tidx, ggml_type tdst) { +ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_set_rows(ggml_metal_library_t lib, const ggml_tensor * op) { char base[256]; char name[256]; - snprintf(base, 256, "kernel_set_rows_%s_%s", ggml_type_name(tdst), ggml_type_name(tidx)); + const auto tsrc = op->src[0]->type; + const auto tidx = op->src[1]->type; + const auto tdst = op->type; + + snprintf(base, 256, "kernel_set_rows_%s_%s_%s", ggml_type_name(tsrc), ggml_type_name(tidx), ggml_type_name(tdst)); snprintf(name, 256, "%s", base); ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name); diff --git a/ggml/src/ggml-metal/ggml-metal-device.h b/ggml/src/ggml-metal/ggml-metal-device.h index 7c8bde362..dc75a34b2 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.h +++ b/ggml/src/ggml-metal/ggml-metal-device.h @@ -112,7 +112,7 @@ struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_cpy struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_pool_1d (ggml_metal_library_t lib, const struct ggml_tensor * op, enum ggml_op_pool op_pool); struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_pool_2d (ggml_metal_library_t lib, const struct ggml_tensor * op, enum ggml_op_pool op_pool); struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_get_rows (ggml_metal_library_t lib, enum ggml_type tsrc); -struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_set_rows (ggml_metal_library_t lib, enum ggml_type tidx, enum ggml_type tdst); +struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_set_rows (ggml_metal_library_t lib, const struct ggml_tensor * op); struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_diag (ggml_metal_library_t lib, const struct ggml_tensor * op); struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_repeat (ggml_metal_library_t lib, enum ggml_type tsrc); struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_concat (ggml_metal_library_t lib, enum ggml_type tsrc); diff --git a/ggml/src/ggml-metal/ggml-metal-device.m b/ggml/src/ggml-metal/ggml-metal-device.m index c742692ca..962a001f3 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.m +++ b/ggml/src/ggml-metal/ggml-metal-device.m @@ -1334,7 +1334,7 @@ bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_te return op->src[0]->type != GGML_TYPE_NVFP4; case GGML_OP_SET_ROWS: { - if (op->src[0]->type != GGML_TYPE_F32) { + if (op->src[0]->type != GGML_TYPE_F32 && op->src[0]->type != GGML_TYPE_F16) { return false; } diff --git a/ggml/src/ggml-metal/ggml-metal-ops.cpp b/ggml/src/ggml-metal/ggml-metal-ops.cpp index d2bc1254a..d196bae4f 100644 --- a/ggml/src/ggml-metal/ggml-metal-ops.cpp +++ b/ggml/src/ggml-metal/ggml-metal-ops.cpp @@ -1202,7 +1202,7 @@ int ggml_metal_op_set_rows(ggml_metal_op_t ctx, int idx) { GGML_TENSOR_LOCALS( int32_t, ne, op, ne); GGML_TENSOR_LOCALS(uint64_t, nb, op, nb); - auto pipeline = ggml_metal_library_get_pipeline_set_rows(lib, op->src[1]->type, op->type); + auto pipeline = ggml_metal_library_get_pipeline_set_rows(lib, op); const int32_t nk0 = ne0/ggml_blck_size(op->type); diff --git a/ggml/src/ggml-metal/ggml-metal.metal b/ggml/src/ggml-metal/ggml-metal.metal index ad50372f4..dcb6803f5 100644 --- a/ggml/src/ggml-metal/ggml-metal.metal +++ b/ggml/src/ggml-metal/ggml-metal.metal @@ -42,6 +42,8 @@ typedef matrix bfloat4x4; typedef matrix bfloat2x4; #endif +#define QK_NL 16 + constexpr constant static float kvalues_iq4nl_f[16] = { -127.f, -104.f, -83.f, -65.f, -49.f, -35.f, -22.f, -10.f, 1.f, 13.f, 25.f, 38.f, 53.f, 69.f, 89.f, 113.f }; @@ -9386,7 +9388,40 @@ kernel void kernel_get_rows_f( } } -template +typedef decltype(kernel_get_rows_f) get_rows_f_t; + +template [[host_name("kernel_get_rows_f32")]] kernel get_rows_f_t kernel_get_rows_f; +template [[host_name("kernel_get_rows_f16")]] kernel get_rows_f_t kernel_get_rows_f; +template [[host_name("kernel_get_rows_i32")]] kernel get_rows_f_t kernel_get_rows_f; +#if defined(GGML_METAL_HAS_BF16) +template [[host_name("kernel_get_rows_bf16")]] kernel get_rows_f_t kernel_get_rows_f; +#endif + +typedef decltype(kernel_get_rows_q) get_rows_q_t; + +template [[host_name("kernel_get_rows_q1_0")]] kernel get_rows_q_t kernel_get_rows_q; +template [[host_name("kernel_get_rows_q4_0")]] kernel get_rows_q_t kernel_get_rows_q; +template [[host_name("kernel_get_rows_q4_1")]] kernel get_rows_q_t kernel_get_rows_q; +template [[host_name("kernel_get_rows_q5_0")]] kernel get_rows_q_t kernel_get_rows_q; +template [[host_name("kernel_get_rows_q5_1")]] kernel get_rows_q_t kernel_get_rows_q; +template [[host_name("kernel_get_rows_q8_0")]] kernel get_rows_q_t kernel_get_rows_q; +template [[host_name("kernel_get_rows_mxfp4")]] kernel get_rows_q_t kernel_get_rows_q; +template [[host_name("kernel_get_rows_q2_K")]] kernel get_rows_q_t kernel_get_rows_q; +template [[host_name("kernel_get_rows_q3_K")]] kernel get_rows_q_t kernel_get_rows_q; +template [[host_name("kernel_get_rows_q4_K")]] kernel get_rows_q_t kernel_get_rows_q; +template [[host_name("kernel_get_rows_q5_K")]] kernel get_rows_q_t kernel_get_rows_q; +template [[host_name("kernel_get_rows_q6_K")]] kernel get_rows_q_t kernel_get_rows_q; +template [[host_name("kernel_get_rows_iq2_xxs")]] kernel get_rows_q_t kernel_get_rows_q; +template [[host_name("kernel_get_rows_iq2_xs")]] kernel get_rows_q_t kernel_get_rows_q; +template [[host_name("kernel_get_rows_iq3_xxs")]] kernel get_rows_q_t kernel_get_rows_q; +template [[host_name("kernel_get_rows_iq3_s")]] kernel get_rows_q_t kernel_get_rows_q; +template [[host_name("kernel_get_rows_iq2_s")]] kernel get_rows_q_t kernel_get_rows_q; +template [[host_name("kernel_get_rows_iq1_s")]] kernel get_rows_q_t kernel_get_rows_q; +template [[host_name("kernel_get_rows_iq1_m")]] kernel get_rows_q_t kernel_get_rows_q; +template [[host_name("kernel_get_rows_iq4_nl")]] kernel get_rows_q_t kernel_get_rows_q; +template [[host_name("kernel_get_rows_iq4_xs")]] kernel get_rows_q_t kernel_get_rows_q; + +template kernel void kernel_set_rows_q32( constant ggml_metal_kargs_set_rows & args, device const void * src0, @@ -9410,14 +9445,14 @@ kernel void kernel_set_rows_q32( const TI i1 = ((const device TI *) ((const device char *) src1 + i10*args.nb10 + i11*args.nb11 + i12*args.nb12))[0]; device block_q * dst_row = ( device block_q *) (( device char *) dst + i1*args.nb1 + i02*args.nb2 + i03*args.nb3); - const device float * src_row = (const device float *) ((const device char *) src0 + i01*args.nb01 + i02*args.nb02 + i03*args.nb03); + const device TS * src_row = (const device TS *) ((const device char *) src0 + i01*args.nb01 + i02*args.nb02 + i03*args.nb03); for (int ind = tiitg%tptg.x; ind < args.nk0; ind += tptg.x) { quantize_func(src_row + 32*ind, dst_row[ind]); } } -template +template kernel void kernel_set_rows_f( constant ggml_metal_kargs_set_rows & args, device const void * src0, @@ -9440,14 +9475,47 @@ kernel void kernel_set_rows_f( const int32_t i10 = i01; const TI i1 = ((const device TI *) ((const device char *) src1 + i10*args.nb10 + i11*args.nb11 + i12*args.nb12))[0]; - device T * dst_row = ( device T *) (( device char *) dst + i1*args.nb1 + i02*args.nb2 + i03*args.nb3); - const device float * src_row = (const device float *) ((const device char *) src0 + i01*args.nb01 + i02*args.nb02 + i03*args.nb03); + device TD * dst_row = ( device TD *) (( device char *) dst + i1*args.nb1 + i02*args.nb2 + i03*args.nb3); + const device TS * src_row = (const device TS *) ((const device char *) src0 + i01*args.nb01 + i02*args.nb02 + i03*args.nb03); for (int ind = tiitg%tptg.x; ind < args.nk0; ind += tptg.x) { - dst_row[ind] = (T) src_row[ind]; + dst_row[ind] = (TD) src_row[ind]; } } +typedef decltype(kernel_set_rows_f) set_rows_f_t; + +template [[host_name("kernel_set_rows_f32_i64_f32")]] kernel set_rows_f_t kernel_set_rows_f; +template [[host_name("kernel_set_rows_f32_i32_f32")]] kernel set_rows_f_t kernel_set_rows_f; +template [[host_name("kernel_set_rows_f32_i64_f16")]] kernel set_rows_f_t kernel_set_rows_f; +template [[host_name("kernel_set_rows_f32_i32_f16")]] kernel set_rows_f_t kernel_set_rows_f; +#if defined(GGML_METAL_HAS_BF16) +template [[host_name("kernel_set_rows_f32_i64_bf16")]] kernel set_rows_f_t kernel_set_rows_f; +template [[host_name("kernel_set_rows_f32_i32_bf16")]] kernel set_rows_f_t kernel_set_rows_f; +#endif + +template [[host_name("kernel_set_rows_f16_i64_f16")]] kernel set_rows_f_t kernel_set_rows_f; +template [[host_name("kernel_set_rows_f16_i32_f16")]] kernel set_rows_f_t kernel_set_rows_f; +#if defined(GGML_METAL_HAS_BF16) +template [[host_name("kernel_set_rows_bf16_i64_bf16")]] kernel set_rows_f_t kernel_set_rows_f; +template [[host_name("kernel_set_rows_bf16_i32_bf16")]] kernel set_rows_f_t kernel_set_rows_f; +#endif + +typedef decltype(kernel_set_rows_q32) set_rows_q32_t; + +template [[host_name("kernel_set_rows_f32_i64_q8_0")]] kernel set_rows_q32_t kernel_set_rows_q32; +template [[host_name("kernel_set_rows_f32_i32_q8_0")]] kernel set_rows_q32_t kernel_set_rows_q32; +template [[host_name("kernel_set_rows_f32_i64_q4_0")]] kernel set_rows_q32_t kernel_set_rows_q32; +template [[host_name("kernel_set_rows_f32_i32_q4_0")]] kernel set_rows_q32_t kernel_set_rows_q32; +template [[host_name("kernel_set_rows_f32_i64_q4_1")]] kernel set_rows_q32_t kernel_set_rows_q32; +template [[host_name("kernel_set_rows_f32_i32_q4_1")]] kernel set_rows_q32_t kernel_set_rows_q32; +template [[host_name("kernel_set_rows_f32_i64_q5_0")]] kernel set_rows_q32_t kernel_set_rows_q32; +template [[host_name("kernel_set_rows_f32_i32_q5_0")]] kernel set_rows_q32_t kernel_set_rows_q32; +template [[host_name("kernel_set_rows_f32_i64_q5_1")]] kernel set_rows_q32_t kernel_set_rows_q32; +template [[host_name("kernel_set_rows_f32_i32_q5_1")]] kernel set_rows_q32_t kernel_set_rows_q32; +template [[host_name("kernel_set_rows_f32_i64_iq4_nl")]] kernel set_rows_q32_t kernel_set_rows_q32; +template [[host_name("kernel_set_rows_f32_i32_iq4_nl")]] kernel set_rows_q32_t kernel_set_rows_q32; + kernel void kernel_diag_f32( constant ggml_metal_kargs_diag & args, device const char * src0, @@ -10190,75 +10258,6 @@ kernel void kernel_mul_mm_id( } } -#define QK_NL 16 - -// -// get rows -// - -typedef decltype(kernel_get_rows_f) get_rows_f_t; - -template [[host_name("kernel_get_rows_f32")]] kernel get_rows_f_t kernel_get_rows_f; -template [[host_name("kernel_get_rows_f16")]] kernel get_rows_f_t kernel_get_rows_f; -template [[host_name("kernel_get_rows_i32")]] kernel get_rows_f_t kernel_get_rows_f; -#if defined(GGML_METAL_HAS_BF16) -template [[host_name("kernel_get_rows_bf16")]] kernel get_rows_f_t kernel_get_rows_f; -#endif - -typedef decltype(kernel_get_rows_q) get_rows_q_t; - -template [[host_name("kernel_get_rows_q1_0")]] kernel get_rows_q_t kernel_get_rows_q; -template [[host_name("kernel_get_rows_q4_0")]] kernel get_rows_q_t kernel_get_rows_q; -template [[host_name("kernel_get_rows_q4_1")]] kernel get_rows_q_t kernel_get_rows_q; -template [[host_name("kernel_get_rows_q5_0")]] kernel get_rows_q_t kernel_get_rows_q; -template [[host_name("kernel_get_rows_q5_1")]] kernel get_rows_q_t kernel_get_rows_q; -template [[host_name("kernel_get_rows_q8_0")]] kernel get_rows_q_t kernel_get_rows_q; -template [[host_name("kernel_get_rows_mxfp4")]] kernel get_rows_q_t kernel_get_rows_q; -template [[host_name("kernel_get_rows_q2_K")]] kernel get_rows_q_t kernel_get_rows_q; -template [[host_name("kernel_get_rows_q3_K")]] kernel get_rows_q_t kernel_get_rows_q; -template [[host_name("kernel_get_rows_q4_K")]] kernel get_rows_q_t kernel_get_rows_q; -template [[host_name("kernel_get_rows_q5_K")]] kernel get_rows_q_t kernel_get_rows_q; -template [[host_name("kernel_get_rows_q6_K")]] kernel get_rows_q_t kernel_get_rows_q; -template [[host_name("kernel_get_rows_iq2_xxs")]] kernel get_rows_q_t kernel_get_rows_q; -template [[host_name("kernel_get_rows_iq2_xs")]] kernel get_rows_q_t kernel_get_rows_q; -template [[host_name("kernel_get_rows_iq3_xxs")]] kernel get_rows_q_t kernel_get_rows_q; -template [[host_name("kernel_get_rows_iq3_s")]] kernel get_rows_q_t kernel_get_rows_q; -template [[host_name("kernel_get_rows_iq2_s")]] kernel get_rows_q_t kernel_get_rows_q; -template [[host_name("kernel_get_rows_iq1_s")]] kernel get_rows_q_t kernel_get_rows_q; -template [[host_name("kernel_get_rows_iq1_m")]] kernel get_rows_q_t kernel_get_rows_q; -template [[host_name("kernel_get_rows_iq4_nl")]] kernel get_rows_q_t kernel_get_rows_q; -template [[host_name("kernel_get_rows_iq4_xs")]] kernel get_rows_q_t kernel_get_rows_q; - -// -// set rows -// - -typedef decltype(kernel_set_rows_f) set_rows_f_t; - -template [[host_name("kernel_set_rows_f32_i64")]] kernel set_rows_f_t kernel_set_rows_f; -template [[host_name("kernel_set_rows_f32_i32")]] kernel set_rows_f_t kernel_set_rows_f; -template [[host_name("kernel_set_rows_f16_i64")]] kernel set_rows_f_t kernel_set_rows_f; -template [[host_name("kernel_set_rows_f16_i32")]] kernel set_rows_f_t kernel_set_rows_f; -#if defined(GGML_METAL_HAS_BF16) -template [[host_name("kernel_set_rows_bf16_i64")]] kernel set_rows_f_t kernel_set_rows_f; -template [[host_name("kernel_set_rows_bf16_i32")]] kernel set_rows_f_t kernel_set_rows_f; -#endif - -typedef decltype(kernel_set_rows_q32) set_rows_q32_t; - -template [[host_name("kernel_set_rows_q8_0_i64")]] kernel set_rows_q32_t kernel_set_rows_q32; -template [[host_name("kernel_set_rows_q8_0_i32")]] kernel set_rows_q32_t kernel_set_rows_q32; -template [[host_name("kernel_set_rows_q4_0_i64")]] kernel set_rows_q32_t kernel_set_rows_q32; -template [[host_name("kernel_set_rows_q4_0_i32")]] kernel set_rows_q32_t kernel_set_rows_q32; -template [[host_name("kernel_set_rows_q4_1_i64")]] kernel set_rows_q32_t kernel_set_rows_q32; -template [[host_name("kernel_set_rows_q4_1_i32")]] kernel set_rows_q32_t kernel_set_rows_q32; -template [[host_name("kernel_set_rows_q5_0_i64")]] kernel set_rows_q32_t kernel_set_rows_q32; -template [[host_name("kernel_set_rows_q5_0_i32")]] kernel set_rows_q32_t kernel_set_rows_q32; -template [[host_name("kernel_set_rows_q5_1_i64")]] kernel set_rows_q32_t kernel_set_rows_q32; -template [[host_name("kernel_set_rows_q5_1_i32")]] kernel set_rows_q32_t kernel_set_rows_q32; -template [[host_name("kernel_set_rows_iq4_nl_i64")]] kernel set_rows_q32_t kernel_set_rows_q32; -template [[host_name("kernel_set_rows_iq4_nl_i32")]] kernel set_rows_q32_t kernel_set_rows_q32; - // // matrix-matrix multiplication // From da46e59cbf899282a70eb48a3502f7811e9ede29 Mon Sep 17 00:00:00 2001 From: Georgi Gerganov Date: Wed, 8 Jul 2026 10:00:03 +0300 Subject: [PATCH 13/17] llama-eval : fix crash when answer is None in HTML dump (#25435) dict.get("key", default) returns None (not default) when the key exists but its value is explicitly None. This caused an AttributeError in _escape_html() when a task errored before grading and answer was set to None. Assisted-by: pi:llama.cpp/Qwen3.6-27B --- examples/llama-eval/llama-eval.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/examples/llama-eval/llama-eval.py b/examples/llama-eval/llama-eval.py index 4bdd239c0..61bdbddd8 100755 --- a/examples/llama-eval/llama-eval.py +++ b/examples/llama-eval/llama-eval.py @@ -362,7 +362,7 @@ class EvalState: case = cases.get(task_id, {}) status = case.get("status", "pending") expected = case.get("expected", "") - answer = case.get("answer", "") if status == "ok" else "" + answer = case.get("answer") or "" if status == "ok" else "" is_correct = case.get("correct", False) if status == "ok" else False response = case.get("response", "") or "" prompt = case.get("prompt", "") or "" @@ -647,7 +647,7 @@ class EvalState: question, prompt, expected = self.get_case(i) case = cases.get(task_id, {}) status = case.get("status", "pending") - answer = case.get("answer", "N/A") if status == "ok" else "N/A" + answer = case.get("answer") or "N/A" if status == "ok" else "N/A" tokens = case.get("tokens") tokens_str = str(tokens) if tokens is not None else "N/A" tps_gen = case.get("tps_gen") From f1161b15f2bbca15defaea3b8639b987af28ce8d Mon Sep 17 00:00:00 2001 From: Aleksander Grygier Date: Wed, 8 Jul 2026 09:22:35 +0200 Subject: [PATCH 14/17] ui: Context usage gauge and panel (#25340) * feat: WIP * feat: Retire ChatScreenProcessingInfo component, context, and keepStatsVisible settings * feat: Always-on gauge with active-model /props, conversation stats and live-reactive reading/output/avg * feat: Add /tokenize endpoint, TokenizeService, FNV-1a and JSON Schema utilities * feat: Surface enabled-tools token count in context hover card * refactor(tools): make toolsStore the sole owner of the OpenAI wire format Previously mcpStore.getToolDefinitionsForLLM() owned the MCP->OpenAI shape conversion (plus normalizeSchemaProperties). That created two sources of truth for what gets sent to the LLM, with the duplication-prone risk of the deduplicated enabled list (which feeds the token-count cache) drifting from the bytes actually shipped on chat. Now: - mcpStore: pure protocol state + routing. Drop getToolDefinitionsForLLM and the inline OpenAIToolDefinition conversion + normalizeSchemaProperties. Doc comment adjusted to declare wire-format ownership as belonging to toolsStore. Connection lifecycle, health checks, executeTool, and the connections/toolsIndex remain. - toolsStore: owns the wire shape (added earlier this series). mcpEntries() inlines the MCP tool conversion; uses normalizeJsonSchema (the JSON Schema util extracted in the prior commit) so missing 'type' fields are inferred from defaults. mcpTools getter iterates mcpEntries() so the Settings UI and the deduplicated enabled list see the same definitions. getEnabledToolsForLLM iterates mcpEntries() instead of calling mcpStore, so the JSON sent to the LLM is identical to what toolsStore.refreshEnabledToolsTokenCount tokenizes. - agentic: the chat-completion tools field's type was annotated as ReturnType, claiming the shape was owned by mcpStore. Switch to ReturnType, the actual source. Assisted-by: Claude * feat: UI WIP * feat: UI WIP * feat: UI WIP * feat: Adjust reasoning submenu layout and spacing * feat: Adjust context usage gauge thresholds and styling * feat: Split context usage gauge stats into current and cumulative breakdowns * chore: Format * refactor: Cleanup * refactor: Cleanup * feat: improve token gauge accuracy and display * refactor: remove MCP recommendation gating and simplify server visibility * feat: add token audit logging to ChatStore for debugging * refactor: Simplify context token reading to use server promptTokens directly * feat: Replace last-known token tracking with live server-derived stats for accurate streaming gauges * feat: UI Improvements * feat: Move prompt processing stats to the preceding user message * feat: Fix context token double-counting and refine gauge layout * refactor: remove always-show-agentic-turns setting and simplify agentic turn display * feat: track and display cache tokens in context gauge * feat: add diagnostic logging for chat completion requests * refactor: improve token audit console output with fresh/cached breakdown * fix: invalidate enabled tools token count cache on tool changes * test: add unit tests for tools store token count invalidation * refactor: Remove tools token counting infrastructure * refactor: Update ChatFormContextGauge to use simplified token tracking * refactor: Update ChatStore to remove tools token counting * chore: Formatting * feat: Improve UI text * feat: simplify context usage derivation and refine gauge labels * refactor: cleanup logs * cleaning * fix: UI * refactor: Enums * refactor: Extract context gauge logic into hook and split UI into sub-components * refactor: Cleanup comments --------- Co-authored-by: Pascal --- tools/ui/package-lock.json | 8 +- tools/ui/package.json | 2 +- .../ChatFormActionAddDropdown.svelte | 9 +- .../ChatFormActionAddReasoningSubmenu.svelte} | 75 ++--- .../ChatFormActions/ChatFormActions.svelte | 50 ++- .../ChatFormReasoningEffortSubmenu.svelte | 23 +- .../ChatFormContextGauge.svelte | 108 +++++++ .../ContextGaugeDetailRow.svelte | 20 ++ .../ContextGaugeDetails.svelte | 122 ++++++++ .../ContextGaugeDial.svelte | 43 +++ .../ContextGaugeLoadModel.svelte | 24 ++ .../ChatFormContextGauge/context-gauge.ts | 37 +++ .../ChatMessage/ChatMessage.svelte | 6 + .../ChatMessageAssistant.svelte | 80 +---- .../ChatMessageUser/ChatMessageUser.svelte | 72 ++++- .../ChatMessageUserPending.svelte | 8 +- .../ChatMessageAgenticContent.svelte | 46 +-- .../ChatMessageStatistics.svelte | 97 +++--- .../app/chat/ChatMessages/ChatMessages.svelte | 26 +- .../app/chat/ChatScreen/ChatScreen.svelte | 20 +- .../ChatScreenProcessingInfo.svelte | 127 -------- tools/ui/src/lib/components/app/chat/index.ts | 21 +- .../content/CollapsibleContentBlock.svelte | 2 +- .../app/content/SyntaxHighlightedCode.svelte | 4 +- .../DialogMcpServerRecommendations.svelte | 36 ++- .../SettingsChat/SettingsChatToolsTab.svelte | 18 +- .../ui/hover-card/hover-card-content.svelte | 31 ++ .../ui/hover-card/hover-card-portal.svelte | 7 + .../ui/hover-card/hover-card-trigger.svelte | 7 + .../ui/hover-card/hover-card.svelte | 7 + .../src/lib/components/ui/hover-card/index.ts | 15 + tools/ui/src/lib/constants/context-keys.ts | 1 - tools/ui/src/lib/constants/css-classes.ts | 1 + .../ui/src/lib/constants/reasoning-effort.ts | 3 +- tools/ui/src/lib/constants/settings-keys.ts | 2 - .../ui/src/lib/constants/settings-registry.ts | 24 -- tools/ui/src/lib/constants/storage.ts | 1 - tools/ui/src/lib/contexts/index.ts | 6 - .../lib/contexts/processing-info.context.ts | 16 - tools/ui/src/lib/enums/chat.enums.ts | 6 + tools/ui/src/lib/enums/index.ts | 1 + .../src/lib/enums/reasoning-effort.enums.ts | 1 + .../src/lib/hooks/use-context-gauge.svelte.ts | 295 ++++++++++++++++++ .../hooks/use-mcp-recommendations.svelte.ts | 5 - .../lib/hooks/use-processing-state.svelte.ts | 10 +- tools/ui/src/lib/services/chat.service.ts | 3 +- .../ui/src/lib/services/migration.service.ts | 6 +- tools/ui/src/lib/stores/agentic.svelte.ts | 2 +- tools/ui/src/lib/stores/chat.svelte.ts | 4 +- .../ui/src/lib/stores/conversations.svelte.ts | 78 ++--- tools/ui/src/lib/stores/mcp.svelte.ts | 106 +------ tools/ui/src/lib/stores/tools.svelte.ts | 160 +++++++--- tools/ui/src/lib/utils/branching.ts | 11 +- 53 files changed, 1234 insertions(+), 659 deletions(-) rename tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/{ChatFormReasoningToggle.svelte => ChatFormActionAdd/ChatFormActionAddReasoningSubmenu.svelte} (63%) create mode 100644 tools/ui/src/lib/components/app/chat/ChatForm/ChatFormContextGauge/ChatFormContextGauge.svelte create mode 100644 tools/ui/src/lib/components/app/chat/ChatForm/ChatFormContextGauge/ContextGaugeDetailRow.svelte create mode 100644 tools/ui/src/lib/components/app/chat/ChatForm/ChatFormContextGauge/ContextGaugeDetails.svelte create mode 100644 tools/ui/src/lib/components/app/chat/ChatForm/ChatFormContextGauge/ContextGaugeDial.svelte create mode 100644 tools/ui/src/lib/components/app/chat/ChatForm/ChatFormContextGauge/ContextGaugeLoadModel.svelte create mode 100644 tools/ui/src/lib/components/app/chat/ChatForm/ChatFormContextGauge/context-gauge.ts delete mode 100644 tools/ui/src/lib/components/app/chat/ChatScreen/ChatScreenProcessingInfo.svelte create mode 100644 tools/ui/src/lib/components/ui/hover-card/hover-card-content.svelte create mode 100644 tools/ui/src/lib/components/ui/hover-card/hover-card-portal.svelte create mode 100644 tools/ui/src/lib/components/ui/hover-card/hover-card-trigger.svelte create mode 100644 tools/ui/src/lib/components/ui/hover-card/hover-card.svelte create mode 100644 tools/ui/src/lib/components/ui/hover-card/index.ts delete mode 100644 tools/ui/src/lib/contexts/processing-info.context.ts create mode 100644 tools/ui/src/lib/hooks/use-context-gauge.svelte.ts diff --git a/tools/ui/package-lock.json b/tools/ui/package-lock.json index 9dce3a0c9..7216de682 100644 --- a/tools/ui/package-lock.json +++ b/tools/ui/package-lock.json @@ -11,7 +11,7 @@ "@chromatic-com/storybook": "5.0.0", "@eslint/compat": "1.4.1", "@eslint/js": "9.39.2", - "@internationalized/date": "3.10.1", + "@internationalized/date": "3.12.2", "@lucide/svelte": "0.515.0", "@modelcontextprotocol/sdk": "1.26.0", "@playwright/test": "1.56.1", @@ -2981,9 +2981,9 @@ } }, "node_modules/@internationalized/date": { - "version": "3.10.1", - "resolved": "https://registry.npmjs.org/@internationalized/date/-/date-3.10.1.tgz", - "integrity": "sha512-oJrXtQiAXLvT9clCf1K4kxp3eKsQhIaZqxEyowkBcsvZDdZkbWrVmnGknxs5flTD0VGsxrxKgBCZty1EzoiMzA==", + "version": "3.12.2", + "resolved": "https://registry.npmjs.org/@internationalized/date/-/date-3.12.2.tgz", + "integrity": "sha512-FY1Y+H64NDs+HAF6omlnWxm3mEpfgaCSWtL5l551ZZfImA+kGjPFgrnJrGjH6lfmLL0g8Z/mBu1R3kufeCp6Jw==", "dev": true, "license": "Apache-2.0", "dependencies": { diff --git a/tools/ui/package.json b/tools/ui/package.json index bcb4165d1..8b3516a02 100644 --- a/tools/ui/package.json +++ b/tools/ui/package.json @@ -30,7 +30,7 @@ "@chromatic-com/storybook": "5.0.0", "@eslint/compat": "1.4.1", "@eslint/js": "9.39.2", - "@internationalized/date": "3.10.1", + "@internationalized/date": "3.12.2", "@lucide/svelte": "0.515.0", "@modelcontextprotocol/sdk": "1.26.0", "@playwright/test": "1.56.1", diff --git a/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActionAdd/ChatFormActionAddDropdown.svelte b/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActionAdd/ChatFormActionAddDropdown.svelte index 479540321..433a3662d 100644 --- a/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActionAdd/ChatFormActionAddDropdown.svelte +++ b/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActionAdd/ChatFormActionAddDropdown.svelte @@ -11,7 +11,8 @@ } from '$lib/constants'; import { ChatFormActionAddToolsSubmenu, - ChatFormActionAddMcpServersSubmenu + ChatFormActionAddMcpServersSubmenu, + ChatFormActionAddReasoningSubmenu } from '$lib/components/app'; import { useAttachmentMenu } from '$lib/hooks/use-attachment-menu.svelte'; @@ -92,7 +93,11 @@ - + + + + + diff --git a/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormReasoningToggle.svelte b/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActionAdd/ChatFormActionAddReasoningSubmenu.svelte similarity index 63% rename from tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormReasoningToggle.svelte rename to tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActionAdd/ChatFormActionAddReasoningSubmenu.svelte index f6bcbcb09..070fd3ac6 100644 --- a/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormReasoningToggle.svelte +++ b/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActionAdd/ChatFormActionAddReasoningSubmenu.svelte @@ -2,7 +2,7 @@ import { Lightbulb, LightbulbOff, Check, Info } from '@lucide/svelte'; import * as DropdownMenu from '$lib/components/ui/dropdown-menu'; import * as Tooltip from '$lib/components/ui/tooltip'; - import { ReasoningEffort, MessageRole } from '$lib/enums'; + import { ReasoningEffort } from '$lib/enums'; import { REASONING_EFFORT_TOKENS } from '$lib/constants/reasoning-effort-tokens'; import { REASONING_EFFORT_LEVELS } from '$lib/constants/reasoning-effort'; import type { ReasoningEffortLevel } from '$lib/types'; @@ -18,31 +18,23 @@ import { isRouterMode } from '$lib/stores/server.svelte'; import type { DatabaseMessage } from '$lib/types/database'; - let thinkingEnabled = $derived(conversationsStore.getThinkingEnabled()); - let currentEffort = $derived(conversationsStore.getReasoningEffort()); - let isOff = $derived(!thinkingEnabled); - let tooltipText = $derived(thinkingEnabled ? `${currentEffort} Reasoning` : 'Disabled Reasoning'); let subOpen = $state(false); - // Get conversation model from message history let conversationModel = $derived( chatStore.getConversationModel(activeMessages() as DatabaseMessage[]) ); - // Fallback: if model props aren't available, check if any assistant messages - // for this model in the active conversation have reasoning content. let modelSupportsThinkingFromMessages = $derived.by(() => { const modelId = isRouterMode() ? modelsStore.selectedModelName || conversationModel : null; if (!modelId) return false; + const messages = conversationsStore.activeMessages; + return messages.some( - (m: DatabaseMessage) => - m.role === MessageRole.ASSISTANT && m.model === modelId && !!m.reasoningContent + (m) => m.role === 'assistant' && m.model === modelId && !!m.reasoningContent ); }); - // Check if model supports thinking. Primary: chat template from /props. - // Fallback: message history (reasoning content in assistant messages). let modelSupportsThinking = $derived.by(() => { loadedModelIds(); propsCacheVersion(); @@ -52,15 +44,15 @@ return checkModelSupportsThinking(modelId ?? '') || modelSupportsThinkingFromMessages; } - // In non-router mode, use the built-in supportsThinking return supportsThinking() || modelSupportsThinkingFromMessages; }); - // Check if current item is selected + let thinkingEnabled = $derived(conversationsStore.getThinkingEnabled()); + let currentEffort = $derived(conversationsStore.getReasoningEffort()); + let isOff = $derived(!thinkingEnabled); + function isSelected(item: ReasoningEffortLevel): boolean { - if (item.isOff) { - return isOff; - } + if (item.isOff) return isOff; return thinkingEnabled && currentEffort === item.value; } @@ -76,39 +68,30 @@ {#if modelSupportsThinking} - - - - - {#if thinkingEnabled} - - {:else} - - {/if} - - + + + {#if thinkingEnabled} + + {:else} + + {/if} - -

{tooltipText}

- - + + Reasoning - + {thinkingEnabled ? currentEffort : 'off'} + + + + + -
Reasoning effort
- {#each REASONING_EFFORT_LEVELS as level (level.value)} {/each} - - + + {/if} diff --git a/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActions.svelte b/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActions.svelte index a80f00bc6..7be356f2a 100644 --- a/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActions.svelte +++ b/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActions.svelte @@ -7,14 +7,20 @@ ChatFormActionModels, ChatFormActionRecord, ChatFormActionSubmit, - ChatFormReasoningToggle + ChatFormContextGauge } from '$lib/components/app'; - import { FileTypeCategory } from '$lib/enums'; + import { FileTypeCategory, MessageRole } from '$lib/enums'; import { mcpStore } from '$lib/stores/mcp.svelte'; import { config } from '$lib/stores/settings.svelte'; - import { conversationsStore } from '$lib/stores/conversations.svelte'; + import { activeMessages, conversationsStore } from '$lib/stores/conversations.svelte'; + import { + activeProcessingState, + isChatStreaming, + isLoading as chatIsLoading + } from '$lib/stores/chat.svelte'; import { getFileTypeCategory } from '$lib/utils'; import { goto } from '$app/navigation'; + import { page } from '$app/state'; import { ROUTES } from '$lib/constants/routes'; interface Props { @@ -93,6 +99,36 @@ let activeMessage = $derived( conversationsStore.activeMessages[conversationsStore.activeMessages.length - 1] ); + + let hasProcessedTokens = $derived.by(() => { + if (!page.params.id) return false; + + const messages = activeMessages() as DatabaseMessage[]; + let totalHistoricalTokens = 0; + for (const m of messages) { + if (m.role !== MessageRole.ASSISTANT) continue; + const timings = m.timings; + if (!timings) continue; + const agenticLlm = timings.agentic?.llm; + if (agenticLlm?.prompt_n != null || agenticLlm?.predicted_n != null) { + totalHistoricalTokens += (agenticLlm?.prompt_n ?? 0) + (agenticLlm?.predicted_n ?? 0); + } else { + totalHistoricalTokens += (timings.prompt_n ?? 0) + (timings.predicted_n ?? 0); + } + } + if (totalHistoricalTokens > 0) return true; + + if (!chatIsLoading() && !isChatStreaming()) return false; + + const processingState = activeProcessingState(); + if (!processingState) return false; + const livePromptTokens = Math.max( + processingState.promptTokens ?? 0, + processingState.promptProgress?.processed ?? 0 + ); + const liveOutputTokens = processingState.outputTokensUsed ?? 0; + return livePromptTokens > 0 || liveOutputTokens > 0; + });
{#if showAddButton} -
+
{/if} -
- +
+ {#if hasProcessedTokens} + + {/if} {#if showModelSelector} - + {#if thinkingEnabled} {:else} @@ -89,23 +88,15 @@ {/if} - + {#each REASONING_EFFORT_LEVELS as level (level.value)} {/each} diff --git a/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormContextGauge/ChatFormContextGauge.svelte b/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormContextGauge/ChatFormContextGauge.svelte new file mode 100644 index 000000000..855cf6ce7 --- /dev/null +++ b/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormContextGauge/ChatFormContextGauge.svelte @@ -0,0 +1,108 @@ + + + + + + + + +
+
+ Context + · + + {formatParameters(gauge.contextUsed)} + / {gauge.contextTotal !== null ? formatParameters(gauge.contextTotal) : '-'} + +
+ + {#if gauge.activeModelId !== null && !gauge.isActiveModelLoaded} + + {:else if showProgressBar} +
+
+
+ +
+ + {gauge.contextPercent}% used + + + {formatParameters((gauge.contextTotal ?? 0) - gauge.contextUsed)} remaining + +
+ {:else} +
No context info available
+ {/if} + + {#if gauge.hasAnyUsage} + + {/if} +
+ + diff --git a/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormContextGauge/ContextGaugeDetailRow.svelte b/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormContextGauge/ContextGaugeDetailRow.svelte new file mode 100644 index 000000000..71c6a33cb --- /dev/null +++ b/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormContextGauge/ContextGaugeDetailRow.svelte @@ -0,0 +1,20 @@ + + +
+
+ {label} + {value} +
+ + {#if subtitle} +
{subtitle}
+ {/if} +
diff --git a/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormContextGauge/ContextGaugeDetails.svelte b/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormContextGauge/ContextGaugeDetails.svelte new file mode 100644 index 000000000..fdec5aca5 --- /dev/null +++ b/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormContextGauge/ContextGaugeDetails.svelte @@ -0,0 +1,122 @@ + + + + + Token usage details + + + + + + {#if hasCumulative} +
+

+ Across all turns +

+ +
+ {#if cumulativeRead > 0} + 0 + ? `${cumulativeCacheTotal.toLocaleString()} reused from KV cache` + : undefined} + /> + {/if} + {#if cumulativeOutput > 0} + + {/if} +
+
+ {/if} + + {#if hasCurrent} +
+

+ This turn · KV cache +

+ +
+ {#if currentRead > 0} + 0 + ? `${currentFresh.toLocaleString()} fresh + ${currentCache.toLocaleString()} cached` + : undefined} + /> + {/if} + + {#if currentOutput > 0} + + {/if} + +
+
+ KV cache total + {kvTotal.toLocaleString()} tok +
+
+
+
+ {/if} + + {#if averageTokensPerSecond !== null} +
+ +
+ {/if} + + {#each transientDetails as detail (detail)} +
{detail}
+ {/each} + + diff --git a/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormContextGauge/ContextGaugeDial.svelte b/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormContextGauge/ContextGaugeDial.svelte new file mode 100644 index 000000000..6e2616d36 --- /dev/null +++ b/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormContextGauge/ContextGaugeDial.svelte @@ -0,0 +1,43 @@ + + + + + + + diff --git a/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormContextGauge/ContextGaugeLoadModel.svelte b/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormContextGauge/ContextGaugeLoadModel.svelte new file mode 100644 index 000000000..24a67cfdf --- /dev/null +++ b/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormContextGauge/ContextGaugeLoadModel.svelte @@ -0,0 +1,24 @@ + + +{#if modelId !== null && !isLoading} +
+ Available context size is only visible once the model is loaded. + +
+{:else if isLoading} +
+ + Loading model... +
+{/if} diff --git a/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormContextGauge/context-gauge.ts b/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormContextGauge/context-gauge.ts new file mode 100644 index 000000000..5b0010015 --- /dev/null +++ b/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormContextGauge/context-gauge.ts @@ -0,0 +1,37 @@ +export type ColorLevel = 'ok' | 'warning' | 'critical' | 'neutral'; + +const WARNING_THRESHOLD = 80; +const CRITICAL_THRESHOLD = 95; + +export function colorLevelFromPercent(percent: number | null): ColorLevel { + if (percent === null) return 'neutral'; + if (percent >= CRITICAL_THRESHOLD) return 'critical'; + if (percent >= WARNING_THRESHOLD) return 'warning'; + return 'ok'; +} + +export function colorLevelTextClass(level: ColorLevel): string { + switch (level) { + case 'critical': + return 'text-red-400'; + case 'warning': + return 'text-amber-400'; + case 'ok': + return 'text-muted-foreground'; + default: + return 'text-muted-foreground'; + } +} + +export function colorLevelBgClass(level: ColorLevel): string { + switch (level) { + case 'critical': + return 'bg-red-500'; + case 'warning': + return 'bg-amber-500'; + case 'ok': + return 'bg-green-500'; + default: + return 'bg-muted'; + } +} diff --git a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessage.svelte b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessage.svelte index dadcae0c4..2b4a3f9d3 100644 --- a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessage.svelte +++ b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessage.svelte @@ -24,6 +24,8 @@ message: DatabaseMessage; toolMessages?: DatabaseMessage[]; isLastAssistantMessage?: boolean; + isLastUserMessage?: boolean; + nextAssistantMessage?: DatabaseMessage | null; siblingInfo?: ChatMessageSiblingInfo | null; } @@ -32,6 +34,8 @@ message, toolMessages = [], isLastAssistantMessage = false, + isLastUserMessage = false, + nextAssistantMessage = null, siblingInfo = null }: Props = $props(); @@ -359,7 +363,9 @@ (ChatMessageStatsView.GENERATION); - let statsContainerEl: HTMLDivElement | undefined = $state(); - - function getScrollParent(el: HTMLElement): HTMLElement | null { - let parent = el.parentElement; - while (parent) { - const style = getComputedStyle(parent); - if (/(auto|scroll)/.test(style.overflowY)) { - return parent; - } - parent = parent.parentElement; - } - return null; - } - - async function handleStatsViewChange(view: ChatMessageStatsView) { - const el = statsContainerEl; - if (!el) { - activeStatsView = view; - - return; - } - - const scrollParent = getScrollParent(el); - if (!scrollParent) { - activeStatsView = view; - - return; - } - - const yBefore = el.getBoundingClientRect().top; - - activeStatsView = view; - - await tick(); - - const delta = el.getBoundingClientRect().top - yBefore; - if (delta !== 0) { - scrollParent.scrollTop += delta; - } - - // Correct any drift after browser paint - requestAnimationFrame(() => { - const drift = el.getBoundingClientRect().top - yBefore; - - if (Math.abs(drift) > 1) { - scrollParent.scrollTop += drift; - } - }); - } - - let highlightAgenticTurns = $derived( - isAgentic && - (currentConfig.alwaysShowAgenticTurns || activeStatsView === ChatMessageStatsView.SUMMARY) - ); - let displayedModel = $derived(message.model ?? null); // model being switched to while it loads, so the selector bar tracks it @@ -291,7 +234,6 @@ {toolMessages} isStreaming={isChatStreaming()} {isLastAssistantMessage} - highlightTurns={highlightAgenticTurns} /> {/if} {:else} @@ -315,10 +257,7 @@
{#if displayedModel} -
+
{#if isRouter} {:else if isLoading() && currentConfig.showMessageStats} {@const liveStats = processingState.getLiveProcessingStats()} {@const genStats = processingState.getLiveGenerationStats()} - {@const promptProgress = processingState.processingState?.promptProgress} - {@const isStillProcessingPrompt = - promptProgress && promptProgress.processed < promptProgress.total} - {#if liveStats || genStats} + {#if genStats} {/if} {/if} diff --git a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageUser/ChatMessageUser.svelte b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageUser/ChatMessageUser.svelte index 80a0183e6..f7590c3a3 100644 --- a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageUser/ChatMessageUser.svelte +++ b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageUser/ChatMessageUser.svelte @@ -2,10 +2,14 @@ import { ChatMessageActionIcons, ChatMessageEditForm, + ChatMessageStatistics, ChatMessageUserBubble } from '$lib/components/app/chat'; import { getMessageEditContext } from '$lib/contexts'; - import { MessageRole } from '$lib/enums'; + import { useProcessingState } from '$lib/hooks/use-processing-state.svelte'; + import { isLoading } from '$lib/stores/chat.svelte'; + import { MessageRole, ChatMessageStatisticsMode } from '$lib/enums'; + import { config } from '$lib/stores/settings.svelte'; interface Props { class?: string; @@ -17,6 +21,8 @@ assistantMessages: number; messageTypes: string[]; } | null; + isLastUserMessage?: boolean; + nextAssistantMessage?: DatabaseMessage | null; showDeleteDialog: boolean; onEdit: () => void; onDelete: () => void; @@ -32,6 +38,8 @@ message, siblingInfo = null, deletionInfo, + isLastUserMessage = false, + nextAssistantMessage = null, showDeleteDialog, onEdit, onDelete, @@ -44,6 +52,37 @@ // Get contexts const editCtx = getMessageEditContext(); + const processingState = useProcessingState(); + + const currentConfig = $derived(config()); + const isActivelyProcessing = $derived(isLastUserMessage && isLoading()); + + // For agentic turns, prefer the cumulative agentic.llm totals over per-call timings. + let storedReadingStats = $derived.by(() => { + const timings = nextAssistantMessage?.timings; + if (!timings?.prompt_n || !timings?.prompt_ms) return null; + + const agentic = timings.agentic; + + return { + promptTokens: agentic ? agentic.llm.prompt_n : timings.prompt_n, + promptMs: agentic ? agentic.llm.prompt_ms : timings.prompt_ms + }; + }); + + let showStoredReadingStats = $derived( + Boolean(currentConfig.showMessageStats) && storedReadingStats !== null + ); + + let showLiveReadingStats = $derived( + Boolean(currentConfig.showMessageStats) && isActivelyProcessing && storedReadingStats === null + ); + + $effect(() => { + if (showLiveReadingStats) { + processingState.startMonitoring(); + } + });
+ {#if showStoredReadingStats} + +
+
+ +
+
+ {:else if showLiveReadingStats} + {@const liveStats = processingState.getLiveProcessingStats()} + {#if liveStats} +
+
+ +
+
+ {/if} + {/if} + {#if message.timestamp}
content, getExtras: () => extras, @@ -36,9 +32,7 @@
{#if editCtx.isEditing} diff --git a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessageAgenticContent.svelte b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessageAgenticContent.svelte index 07b0489cc..f0d03f547 100644 --- a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessageAgenticContent.svelte +++ b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessageAgenticContent.svelte @@ -41,15 +41,13 @@ toolMessages?: DatabaseMessage[]; isStreaming?: boolean; isLastAssistantMessage?: boolean; - highlightTurns?: boolean; } let { message, toolMessages = [], isStreaming = false, - isLastAssistantMessage = false, - highlightTurns = false + isLastAssistantMessage = false }: Props = $props(); let expandedStates: Record = $state({}); @@ -57,6 +55,7 @@ const showToolCallInProgress = $derived(config().showToolCallInProgress as boolean); const showThoughtInProgress = $derived(config().showThoughtInProgress as boolean); const renderThinkingAsMarkdown = $derived(config().renderThinkingAsMarkdown as boolean); + const showMessageStats = $derived(config().showMessageStats as boolean); const hasReasoningError = $derived( isLastAssistantMessage ? !!agenticLastError(message.convId) : false @@ -354,16 +353,17 @@ {/snippet}
- {#if highlightTurns && turnGroups.length > 1} + {#if turnGroups.length > 1} {#each turnGroups as turn, turnIndex (turnIndex)} {@const turnStats = message?.timings?.agentic?.perTurn?.[turnIndex]} -
- Turn {turnIndex + 1} + +
{#each turn.sections as section, sIdx (turn.flatIndices[sIdx])} {@render renderSection(section, turn.flatIndices[sIdx])} {/each} - {#if turnStats} -
+ + {#if turnStats && showMessageStats} +
:global(*), + .agentic-turn > :global(*) { + min-width: 0; } .agentic-text { width: 100%; } - .agentic-turn { - position: relative; - border: 1.5px dashed var(--muted-foreground); - border-radius: 0.75rem; - padding: 1rem; - transition: background 0.1s; - } - - .agentic-turn-label { - position: absolute; - top: -1rem; - left: 0.75rem; - padding: 0 0.375rem; - background: var(--background); - font-size: 0.7rem; - font-weight: 500; - color: var(--muted-foreground); - text-transform: uppercase; - letter-spacing: 0.05em; - } - .turn-stats { - margin-top: 0.75rem; - padding-top: 0.5rem; border-top: 1px solid hsl(var(--muted) / 0.5); } diff --git a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessageStatistics/ChatMessageStatistics.svelte b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessageStatistics/ChatMessageStatistics.svelte index 6906adbb1..7ef73a499 100644 --- a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessageStatistics/ChatMessageStatistics.svelte +++ b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessageStatistics/ChatMessageStatistics.svelte @@ -2,7 +2,7 @@ import { Clock, Gauge, WholeWord, BookOpenText, Sparkles, Wrench, Layers } from '@lucide/svelte'; import { ChatMessageStatisticsBadge } from '$lib/components/app'; import * as Tooltip from '$lib/components/ui/tooltip'; - import { ChatMessageStatsView } from '$lib/enums'; + import { ChatMessageStatsView, ChatMessageStatisticsMode } from '$lib/enums'; import type { ChatMessageAgenticTimings } from '$lib/types/chat'; import { formatPerformanceTime } from '$lib/utils'; import { MS_PER_SECOND, DEFAULT_PERFORMANCE_TIME } from '$lib/constants'; @@ -19,6 +19,7 @@ agenticTimings?: ChatMessageAgenticTimings; onActiveViewChange?: (view: ChatMessageStatsView) => void; hideSummary?: boolean; + mode?: ChatMessageStatisticsMode; } let { @@ -31,19 +32,30 @@ initialView = ChatMessageStatsView.GENERATION, agenticTimings, onActiveViewChange, - hideSummary = false + hideSummary = false, + mode = ChatMessageStatisticsMode.SWITCHABLE }: Props = $props(); - let activeView: ChatMessageStatsView = $derived(initialView); + let isSwitchable = $derived(mode === ChatMessageStatisticsMode.SWITCHABLE); + + let activeView: ChatMessageStatsView = $derived( + mode === ChatMessageStatisticsMode.READING + ? ChatMessageStatsView.READING + : mode === ChatMessageStatisticsMode.GENERATION + ? ChatMessageStatsView.GENERATION + : initialView + ); let hasAutoSwitchedToGeneration = $state(false); $effect(() => { - onActiveViewChange?.(activeView); + if (isSwitchable) { + onActiveViewChange?.(activeView); + } }); // In live mode: auto-switch to GENERATION tab when prompt processing completes $effect(() => { - if (isLive) { + if (isLive && isSwitchable) { // Auto-switch to generation tab only when prompt processing is done (once) if ( !hasAutoSwitchedToGeneration && @@ -91,8 +103,7 @@ formattedPromptTime !== undefined ); - // In live mode, generation tab is disabled until we have generation stats - let isGenerationDisabled = $derived(isLive && !hasGenerationStats); + let isGenerationDisabled = $derived(isLive && isSwitchable && !hasGenerationStats); let hasAgenticStats = $derived(agenticTimings !== undefined && agenticTimings.toolCallsCount > 0); @@ -153,44 +164,44 @@ {/snippet}
-
- {#if hasPromptStats || isLive} - {@render viewButton({ - view: ChatMessageStatsView.READING, - icon: BookOpenText, - label: 'Reading', - tooltipText: 'Reading (prompt processing)' - })} - {/if} - - {@render viewButton({ - view: ChatMessageStatsView.GENERATION, - icon: Sparkles, - label: 'Generation', - tooltipText: isGenerationDisabled - ? 'Generation (waiting for tokens...)' - : 'Generation (token output)', - disabled: isGenerationDisabled - })} - - {#if hasAgenticStats} - {@render viewButton({ - view: ChatMessageStatsView.TOOLS, - icon: Wrench, - label: 'Tools', - tooltipText: 'Tool calls' - })} - - {#if !hideSummary} + {#if isSwitchable} +
+ {#if hasPromptStats || isLive} {@render viewButton({ - view: ChatMessageStatsView.SUMMARY, - icon: Layers, - label: 'Summary', - tooltipText: 'Agentic summary' + view: ChatMessageStatsView.READING, + icon: BookOpenText, + label: 'Reading', + tooltipText: 'Processing' })} {/if} - {/if} -
+ + {@render viewButton({ + view: ChatMessageStatsView.GENERATION, + icon: Sparkles, + label: 'Generation', + tooltipText: isGenerationDisabled ? 'Waiting for tokens...' : 'Generation', + disabled: isGenerationDisabled + })} + + {#if hasAgenticStats} + {@render viewButton({ + view: ChatMessageStatsView.TOOLS, + icon: Wrench, + label: 'Tools', + tooltipText: 'Tool calls' + })} + + {#if !hideSummary} + {@render viewButton({ + view: ChatMessageStatsView.SUMMARY, + icon: Layers, + label: 'Summary', + tooltipText: 'Agentic summary' + })} + {/if} + {/if} +
+ {/if}
{#if activeView === ChatMessageStatsView.GENERATION && hasGenerationStats} @@ -256,7 +267,7 @@ value={formattedAgenticTotalTime} tooltipLabel="Total time (LLM + tools)" /> - {:else if hasPromptStats} + {:else if hasPromptStats && (mode === ChatMessageStatisticsMode.READING || isSwitchable)} = []; @@ -236,18 +238,36 @@ message: msg, toolMessages, isLastAssistantMessage: false, + isLastUserMessage: false, + nextAssistantMessage: null, siblingInfo }); } - // Mark the last assistant message + let lastAssistantIdx = -1; for (let i = result.length - 1; i >= 0; i--) { if (result[i].message.role === MessageRole.ASSISTANT) { result[i].isLastAssistantMessage = true; + lastAssistantIdx = i; break; } } + if (lastAssistantIdx > 0 && result[lastAssistantIdx - 1].message.role === MessageRole.USER) { + result[lastAssistantIdx - 1].isLastUserMessage = true; + } + + for (let i = 0; i < result.length; i++) { + if (result[i].message.role !== MessageRole.USER) continue; + + for (let j = i + 1; j < result.length; j++) { + if (result[j].message.role === MessageRole.ASSISTANT) { + result[i].nextAssistantMessage = result[j].message; + break; + } + } + } + return result; }); @@ -257,12 +277,14 @@ {isVisible ? 'opacity-100' : 'opacity-0'} {previousRouteId === '/(chat)/chat/[id]' ? '' : 'delay-300'}" > - {#each displayMessages as { message, toolMessages, isLastAssistantMessage, siblingInfo } (message.id)} + {#each displayMessages as { message, toolMessages, isLastAssistantMessage, isLastUserMessage, nextAssistantMessage, siblingInfo } (message.id)} {/each} diff --git a/tools/ui/src/lib/components/app/chat/ChatScreen/ChatScreen.svelte b/tools/ui/src/lib/components/app/chat/ChatScreen/ChatScreen.svelte index d2fba5a93..a0d66c642 100644 --- a/tools/ui/src/lib/components/app/chat/ChatScreen/ChatScreen.svelte +++ b/tools/ui/src/lib/components/app/chat/ChatScreen/ChatScreen.svelte @@ -4,12 +4,10 @@ ChatScreenForm, ChatMessages, ChatScreenDragOverlay, - ChatScreenProcessingInfo, ChatScreenStreamResumeStatus, ServerLoadingSplash, ChatScreenServerError } from '$lib/components/app'; - import { setProcessingInfoContext } from '$lib/contexts'; import { createAutoScrollController } from '$lib/hooks/use-auto-scroll.svelte'; import { useChatScreenActiveModel } from '$lib/hooks/use-chat-screen-active-model.svelte'; import { useChatScreenDragAndDrop } from '$lib/hooks/use-chat-screen-drag-and-drop.svelte'; @@ -23,8 +21,7 @@ errorDialog, isLoading, isChatStreaming, - isEditing, - activeProcessingState + isEditing } from '$lib/stores/chat.svelte'; import { conversationsStore, @@ -42,12 +39,6 @@ let { showCenteredEmpty = false } = $props(); - setProcessingInfoContext({ - get showProcessingInfo() { - return showProcessingInfo; - } - }); - let disableAutoScroll = $derived(Boolean(config().disableAutoScroll) || isMobile.current); let isMobileUserScrolledUp = $state(false); let mobileScrollDownHint = $state(false); @@ -63,11 +54,6 @@ let isServerLoading = $derived(serverLoading()); let hasPropsError = $derived(!!serverError()); let isCurrentConversationLoading = $derived(isLoading() || isChatStreaming()); - let showProcessingInfo = $derived( - isCurrentConversationLoading || - (config().keepStatsVisible && !!page.params.id) || - activeProcessingState() !== null - ); let chatFormBottomPosition = $derived.by(() => { if (!isMobile.current) return '1rem'; if (device.isStandalone) return '1.5rem'; @@ -298,10 +284,6 @@ }} /> {/if} - - {#if showProcessingInfo} - - {/if}
- import { untrack } from 'svelte'; - import { PROCESSING_INFO_TIMEOUT } from '$lib/constants'; - import { useProcessingState } from '$lib/hooks/use-processing-state.svelte'; - import { chatStore, isLoading, isChatStreaming } from '$lib/stores/chat.svelte'; - import { activeMessages, activeConversation } from '$lib/stores/conversations.svelte'; - import { config } from '$lib/stores/settings.svelte'; - - const processingState = useProcessingState(); - - let isCurrentConversationLoading = $derived(isLoading()); - let isStreaming = $derived(isChatStreaming()); - let processingDetails = $derived(processingState.getTechnicalDetails()); - - let processingVisible = $derived(processingDetails.length > 0); - - let { onVisibilityChange }: { onVisibilityChange?: (visible: boolean) => void } = $props(); - - $effect(() => { - onVisibilityChange?.(processingVisible); - }); - - $effect(() => { - const conversation = activeConversation(); - - untrack(() => chatStore.setActiveProcessingConversation(conversation?.id ?? null)); - }); - - $effect(() => { - const keepStatsVisible = config().keepStatsVisible; - const shouldMonitor = keepStatsVisible || isCurrentConversationLoading || isStreaming; - - if (shouldMonitor) { - processingState.startMonitoring(); - } - - if (!isCurrentConversationLoading && !isStreaming && !keepStatsVisible) { - const timeout = setTimeout(() => { - if (!config().keepStatsVisible && !isChatStreaming()) { - processingState.stopMonitoring(); - } - }, PROCESSING_INFO_TIMEOUT); - - return () => clearTimeout(timeout); - } - }); - - $effect(() => { - const conversation = activeConversation(); - const messages = activeMessages() as DatabaseMessage[]; - const keepStatsVisible = config().keepStatsVisible; - - if (keepStatsVisible && conversation) { - if (messages.length === 0) { - untrack(() => chatStore.clearProcessingState(conversation.id)); - return; - } - - if (!isCurrentConversationLoading && !isStreaming) { - untrack(() => chatStore.restoreProcessingStateFromMessages(messages, conversation.id)); - } - } - }); - - - - - diff --git a/tools/ui/src/lib/components/app/chat/index.ts b/tools/ui/src/lib/components/app/chat/index.ts index d7004d3ae..4f826841e 100644 --- a/tools/ui/src/lib/components/app/chat/index.ts +++ b/tools/ui/src/lib/components/app/chat/index.ts @@ -241,13 +241,18 @@ export { default as ChatFormActionAddToolsSubmenu } from './ChatForm/ChatFormAct export { default as ChatFormActionAddMcpServersSubmenu } from './ChatForm/ChatFormActions/ChatFormActionAdd/ChatFormActionAddMcpServersSubmenu.svelte'; /** - * **ChatFormReasoningToggle** - Thinking toggle button with effort dropdown + * Dropdown submenu for selecting reasoning effort level. * - * A toggle button with lightbulb icon that indicates thinking status. - * Shows the reasoning effort dropdown when clicked. + * Shows a "Reasoning" sub-menu item with a lightbulb icon indicating + * thinking status, and a nested list of effort levels. * Only visible when the current model supports thinking. */ -export { default as ChatFormReasoningToggle } from './ChatForm/ChatFormActions/ChatFormReasoningToggle.svelte'; +export { default as ChatFormActionAddReasoningSubmenu } from './ChatForm/ChatFormActions/ChatFormActionAdd/ChatFormActionAddReasoningSubmenu.svelte'; + +/** + * Compact context-usage gauge with per-turn and cumulative breakdown in the tooltip. + */ +export { default as ChatFormContextGauge } from './ChatForm/ChatFormContextGauge/ChatFormContextGauge.svelte'; /** * Hidden file input element for programmatic file selection. @@ -669,14 +674,6 @@ export { default as ChatScreenDragOverlay } from './ChatScreen/ChatScreenDragOve */ export { default as ChatScreenForm } from './ChatScreen/ChatScreenForm.svelte'; -/** - * Processing info display during generation. Shows real-time statistics: - * tokens per second, prompt/completion token counts, and elapsed time. - * Data sourced from slotsService polling during active generation. - * Only visible when `isCurrentConversationLoading` is true. - */ -export { default as ChatScreenProcessingInfo } from './ChatScreen/ChatScreenProcessingInfo.svelte'; - /** * Server error alert displayed when the server is unreachable. * Shows the error message with a retry button. diff --git a/tools/ui/src/lib/components/app/content/CollapsibleContentBlock.svelte b/tools/ui/src/lib/components/app/content/CollapsibleContentBlock.svelte index 8bab55d19..3875b449a 100644 --- a/tools/ui/src/lib/components/app/content/CollapsibleContentBlock.svelte +++ b/tools/ui/src/lib/components/app/content/CollapsibleContentBlock.svelte @@ -76,7 +76,7 @@ open = value; onToggle?.(); }} - class={className} + class="{className} my-0!" > diff --git a/tools/ui/src/lib/components/app/content/SyntaxHighlightedCode.svelte b/tools/ui/src/lib/components/app/content/SyntaxHighlightedCode.svelte index c4d1706bf..d8dfe0ad9 100644 --- a/tools/ui/src/lib/components/app/content/SyntaxHighlightedCode.svelte +++ b/tools/ui/src/lib/components/app/content/SyntaxHighlightedCode.svelte @@ -72,8 +72,8 @@ 
{@html highlightedHtml}
diff --git a/tools/ui/src/lib/components/app/dialogs/DialogMcpServerRecommendations.svelte b/tools/ui/src/lib/components/app/dialogs/DialogMcpServerRecommendations.svelte index 9b4489b82..cdbc055ee 100644 --- a/tools/ui/src/lib/components/app/dialogs/DialogMcpServerRecommendations.svelte +++ b/tools/ui/src/lib/components/app/dialogs/DialogMcpServerRecommendations.svelte @@ -4,9 +4,10 @@ import * as Dialog from '$lib/components/ui/dialog'; import { fly } from 'svelte/transition'; import { McpServerCardCompact, McpServerForm } from '$lib/components/app/mcp'; - import { RECOMMENDED_MCP_SERVERS } from '$lib/constants'; + import { RECOMMENDED_MCP_SERVERS, SETTINGS_KEYS } from '$lib/constants'; import { conversationsStore } from '$lib/stores/conversations.svelte'; import { mcpStore } from '$lib/stores/mcp.svelte'; + import { settingsStore } from '$lib/stores/settings.svelte'; import { uuid } from '$lib/utils'; import { MCP_SERVERS_ADDED_TO_CHAT_LOCALSTORAGE_KEY, MCP_SERVER_ID_PREFIX } from '$lib/constants'; import type { MCPServerSettingsEntry } from '$lib/types'; @@ -24,6 +25,22 @@ ); let addedServers = $state([]); + let didAddAny = $state(false); + + let selectedRecommendedCount = $derived.by( + () => RECOMMENDED_MCP_SERVERS.filter((server) => selected[server.id]).length + ); + + let footerLabel = $derived.by(() => { + const recommended = selectedRecommendedCount; + const custom = addedServers.length; + const total = recommended + custom; + + if (total === 0) return 'Continue'; + if (recommended === 0) return custom === 1 ? 'Add server' : `Add ${custom} servers`; + if (custom === 0) return recommended === 1 ? 'Add server' : `Add ${recommended} servers`; + return `Add ${recommended} servers and ${custom} custom`; + }); let showAddForm = $state(false); let newServerUrl = $state(''); @@ -44,9 +61,14 @@ showAddForm = false; newServerUrl = ''; newServerHeaders = ''; - addedServers = []; + + if (!didAddAny) { + settingsStore.updateConfig(SETTINGS_KEYS.MCP_SERVERS, []); + } localStorage.setItem(MCP_SERVERS_ADDED_TO_CHAT_LOCALSTORAGE_KEY, 'true'); + addedServers = []; + didAddAny = false; } open = value; onOpenChange?.(value); @@ -59,6 +81,7 @@ } function enableSelected() { + didAddAny = true; localStorage.setItem(MCP_SERVERS_ADDED_TO_CHAT_LOCALSTORAGE_KEY, 'true'); for (const server of RECOMMENDED_MCP_SERVERS) { @@ -83,6 +106,8 @@ function saveNewServer() { if (newServerUrlError) return; + didAddAny = true; + const newServerId = uuid() ?? `${MCP_SERVER_ID_PREFIX}-${Date.now()}`; localStorage.setItem(MCP_SERVERS_ADDED_TO_CHAT_LOCALSTORAGE_KEY, 'true'); @@ -174,7 +199,12 @@ - + diff --git a/tools/ui/src/lib/components/app/settings/SettingsChat/SettingsChatToolsTab.svelte b/tools/ui/src/lib/components/app/settings/SettingsChat/SettingsChatToolsTab.svelte index b56832496..8bb49a095 100644 --- a/tools/ui/src/lib/components/app/settings/SettingsChat/SettingsChatToolsTab.svelte +++ b/tools/ui/src/lib/components/app/settings/SettingsChat/SettingsChatToolsTab.svelte @@ -39,13 +39,17 @@ {@const faviconUrl = group.serverId ? mcpStore.getServerFavicon(group.serverId) : null} - + {#if group.source === 'mcp'} + + {:else} + + {/if} diff --git a/tools/ui/src/lib/components/ui/hover-card/hover-card-content.svelte b/tools/ui/src/lib/components/ui/hover-card/hover-card-content.svelte new file mode 100644 index 000000000..db1406a7b --- /dev/null +++ b/tools/ui/src/lib/components/ui/hover-card/hover-card-content.svelte @@ -0,0 +1,31 @@ + + + + + diff --git a/tools/ui/src/lib/components/ui/hover-card/hover-card-portal.svelte b/tools/ui/src/lib/components/ui/hover-card/hover-card-portal.svelte new file mode 100644 index 000000000..9dc25827d --- /dev/null +++ b/tools/ui/src/lib/components/ui/hover-card/hover-card-portal.svelte @@ -0,0 +1,7 @@ + + + diff --git a/tools/ui/src/lib/components/ui/hover-card/hover-card-trigger.svelte b/tools/ui/src/lib/components/ui/hover-card/hover-card-trigger.svelte new file mode 100644 index 000000000..2d42f89c3 --- /dev/null +++ b/tools/ui/src/lib/components/ui/hover-card/hover-card-trigger.svelte @@ -0,0 +1,7 @@ + + + diff --git a/tools/ui/src/lib/components/ui/hover-card/hover-card.svelte b/tools/ui/src/lib/components/ui/hover-card/hover-card.svelte new file mode 100644 index 000000000..ffc075dcd --- /dev/null +++ b/tools/ui/src/lib/components/ui/hover-card/hover-card.svelte @@ -0,0 +1,7 @@ + + + diff --git a/tools/ui/src/lib/components/ui/hover-card/index.ts b/tools/ui/src/lib/components/ui/hover-card/index.ts new file mode 100644 index 000000000..098f69176 --- /dev/null +++ b/tools/ui/src/lib/components/ui/hover-card/index.ts @@ -0,0 +1,15 @@ +import Root from './hover-card.svelte'; +import Content from './hover-card-content.svelte'; +import Trigger from './hover-card-trigger.svelte'; +import Portal from './hover-card-portal.svelte'; + +export { + Root, + Content, + Trigger, + Portal, + Root as HoverCard, + Content as HoverCardContent, + Trigger as HoverCardTrigger, + Portal as HoverCardPortal +}; diff --git a/tools/ui/src/lib/constants/context-keys.ts b/tools/ui/src/lib/constants/context-keys.ts index 12de0d0bc..0bd733b37 100644 --- a/tools/ui/src/lib/constants/context-keys.ts +++ b/tools/ui/src/lib/constants/context-keys.ts @@ -1,4 +1,3 @@ export const CONTEXT_KEY_MESSAGE_EDIT = 'chat-message-edit'; export const CONTEXT_KEY_CHAT_ACTIONS = 'chat-actions'; export const CONTEXT_KEY_CHAT_SETTINGS_CONFIG = 'chat-settings-config'; -export const CONTEXT_KEY_PROCESSING_INFO = 'processing-info'; diff --git a/tools/ui/src/lib/constants/css-classes.ts b/tools/ui/src/lib/constants/css-classes.ts index ca5386fcd..009b5c52e 100644 --- a/tools/ui/src/lib/constants/css-classes.ts +++ b/tools/ui/src/lib/constants/css-classes.ts @@ -17,3 +17,4 @@ export const PANEL_CLASSES = ` `; export const CHAT_FORM_POPOVER_MAX_HEIGHT = 'max-h-80'; +export const DIALOG_SUBMENU_CONTENT = 'w-60'; diff --git a/tools/ui/src/lib/constants/reasoning-effort.ts b/tools/ui/src/lib/constants/reasoning-effort.ts index d854e912a..28a24420e 100644 --- a/tools/ui/src/lib/constants/reasoning-effort.ts +++ b/tools/ui/src/lib/constants/reasoning-effort.ts @@ -6,6 +6,7 @@ import type { ReasoningEffortLevel } from '$lib/types'; * Keys match the ReasoningEffort enum values for type-safe lookups. */ export const REASONING_EFFORT_LABELS: Record = { + [ReasoningEffort.OFF]: 'Off', [ReasoningEffort.LOW]: 'Low', [ReasoningEffort.MEDIUM]: 'Medium', [ReasoningEffort.HIGH]: 'High', @@ -13,7 +14,7 @@ export const REASONING_EFFORT_LABELS: Record = { }; export const REASONING_EFFORT_LEVELS: ReasoningEffortLevel[] = [ - { value: 'off', label: 'Off', isOff: true }, + { value: ReasoningEffort.OFF, label: 'Off', isOff: true }, { value: ReasoningEffort.LOW, label: 'Low' }, { value: ReasoningEffort.MEDIUM, label: 'Medium' }, { value: ReasoningEffort.HIGH, label: 'High' }, diff --git a/tools/ui/src/lib/constants/settings-keys.ts b/tools/ui/src/lib/constants/settings-keys.ts index d4d782705..ea8963044 100644 --- a/tools/ui/src/lib/constants/settings-keys.ts +++ b/tools/ui/src/lib/constants/settings-keys.ts @@ -22,7 +22,6 @@ export const SETTINGS_KEYS = { // Display SHOW_MESSAGE_STATS: 'showMessageStats', SHOW_THOUGHT_IN_PROGRESS: 'showThoughtInProgress', - KEEP_STATS_VISIBLE: 'keepStatsVisible', AUTO_MIC_ON_EMPTY: 'autoMicOnEmpty', RENDER_USER_CONTENT_AS_MARKDOWN: 'renderUserContentAsMarkdown', DISABLE_AUTO_SCROLL: 'disableAutoScroll', @@ -61,7 +60,6 @@ export const SETTINGS_KEYS = { MCP_REQUEST_TIMEOUT_SECONDS: 'mcpRequestTimeoutSeconds', MCP_DEFAULT_SERVER_OVERRIDES: 'mcpDefaultServerOverrides', AGENTIC_MAX_TURNS: 'agenticMaxTurns', - ALWAYS_SHOW_AGENTIC_TURNS: 'alwaysShowAgenticTurns', AGENTIC_MAX_TOOL_PREVIEW_LINES: 'agenticMaxToolPreviewLines', SHOW_TOOL_CALL_IN_PROGRESS: 'showToolCallInProgress', // Performance diff --git a/tools/ui/src/lib/constants/settings-registry.ts b/tools/ui/src/lib/constants/settings-registry.ts index b314510ab..f8657d673 100644 --- a/tools/ui/src/lib/constants/settings-registry.ts +++ b/tools/ui/src/lib/constants/settings-registry.ts @@ -258,18 +258,6 @@ const SETTINGS_REGISTRY: Record = { paramType: SyncableParameterType.BOOLEAN } }, - { - key: SETTINGS_KEYS.KEEP_STATS_VISIBLE, - label: 'Keep stats visible after generation', - help: 'Keep processing statistics visible after generation finishes.', - defaultValue: false, - type: SettingsFieldType.CHECKBOX, - section: SETTINGS_SECTION_SLUGS.DISPLAY, - sync: { - serverKey: SETTINGS_KEYS.KEEP_STATS_VISIBLE, - paramType: SyncableParameterType.BOOLEAN - } - }, { key: SETTINGS_KEYS.AUTO_MIC_ON_EMPTY, label: 'Show microphone on empty input', @@ -379,18 +367,6 @@ const SETTINGS_REGISTRY: Record = { paramType: SyncableParameterType.BOOLEAN } }, - { - key: SETTINGS_KEYS.ALWAYS_SHOW_AGENTIC_TURNS, - label: 'Always show agentic turns in conversation', - help: 'Always expand and display agentic loop turns in conversation messages.', - defaultValue: false, - type: SettingsFieldType.CHECKBOX, - section: SETTINGS_SECTION_SLUGS.DISPLAY, - sync: { - serverKey: SETTINGS_KEYS.ALWAYS_SHOW_AGENTIC_TURNS, - paramType: SyncableParameterType.BOOLEAN - } - }, { key: SETTINGS_KEYS.SHOW_BUILD_VERSION, label: 'Show build version information', diff --git a/tools/ui/src/lib/constants/storage.ts b/tools/ui/src/lib/constants/storage.ts index d0c5f2eff..eca9739ba 100644 --- a/tools/ui/src/lib/constants/storage.ts +++ b/tools/ui/src/lib/constants/storage.ts @@ -21,7 +21,6 @@ export const DISABLED_TOOLS_LOCALSTORAGE_KEY = `${STORAGE_APP_NAME}.disabledTool /** Disabled tools keyed by stable selection identity, no migration from the name based key */ export const DISABLED_TOOL_KEYS_LOCALSTORAGE_KEY = `${STORAGE_APP_NAME}.disabledToolKeys`; export const FAVORITE_MODELS_LOCALSTORAGE_KEY = `${STORAGE_APP_NAME}.favoriteModels`; -export const THINKING_ENABLED_DEFAULT_LOCALSTORAGE_KEY = `${STORAGE_APP_NAME}.thinkingEnabledDefault`; export const REASONING_EFFORT_DEFAULT_LOCALSTORAGE_KEY = `${STORAGE_APP_NAME}.reasoningEffortDefault`; /** Set when user has interacted with the MCP server recommendations dialog (checked servers, added custom server, or dismissed) */ export const MCP_SERVERS_ADDED_TO_CHAT_LOCALSTORAGE_KEY = `${STORAGE_APP_NAME}.mcpServersSetupDone`; diff --git a/tools/ui/src/lib/contexts/index.ts b/tools/ui/src/lib/contexts/index.ts index 01cd1d4b7..c6719fa9e 100644 --- a/tools/ui/src/lib/contexts/index.ts +++ b/tools/ui/src/lib/contexts/index.ts @@ -17,9 +17,3 @@ export { setChatSettingsConfigContext, type ChatSettingsConfigContext } from './chat-settings-config.context'; - -export { - getProcessingInfoContext, - setProcessingInfoContext, - type ProcessingInfoContext -} from './processing-info.context'; diff --git a/tools/ui/src/lib/contexts/processing-info.context.ts b/tools/ui/src/lib/contexts/processing-info.context.ts deleted file mode 100644 index 0cf43336f..000000000 --- a/tools/ui/src/lib/contexts/processing-info.context.ts +++ /dev/null @@ -1,16 +0,0 @@ -import { getContext, setContext } from 'svelte'; -import { CONTEXT_KEY_PROCESSING_INFO } from '$lib/constants'; - -export interface ProcessingInfoContext { - readonly showProcessingInfo: boolean; -} - -const PROCESSING_INFO_KEY = Symbol.for(CONTEXT_KEY_PROCESSING_INFO); - -export function setProcessingInfoContext(ctx: ProcessingInfoContext): ProcessingInfoContext { - return setContext(PROCESSING_INFO_KEY, ctx); -} - -export function getProcessingInfoContext(): ProcessingInfoContext { - return getContext(PROCESSING_INFO_KEY); -} diff --git a/tools/ui/src/lib/enums/chat.enums.ts b/tools/ui/src/lib/enums/chat.enums.ts index 278e4af84..f4994bb8e 100644 --- a/tools/ui/src/lib/enums/chat.enums.ts +++ b/tools/ui/src/lib/enums/chat.enums.ts @@ -5,6 +5,12 @@ export enum ChatMessageStatsView { SUMMARY = 'summary' } +export enum ChatMessageStatisticsMode { + SWITCHABLE = 'switchable', + READING = 'reading', + GENERATION = 'generation' +} + /** * Connection state of a streamed completion, drives the resume status indicator. */ diff --git a/tools/ui/src/lib/enums/index.ts b/tools/ui/src/lib/enums/index.ts index d8e47958b..847014e78 100644 --- a/tools/ui/src/lib/enums/index.ts +++ b/tools/ui/src/lib/enums/index.ts @@ -10,6 +10,7 @@ export { AgenticSectionType, ContinueIntentKind, ToolCallType } from './agentic. export { ChatMessageStatsView, + ChatMessageStatisticsMode, StreamConnectionState, ContentPartType, ConversationSelectionMode, diff --git a/tools/ui/src/lib/enums/reasoning-effort.enums.ts b/tools/ui/src/lib/enums/reasoning-effort.enums.ts index dadb0c726..172f118e7 100644 --- a/tools/ui/src/lib/enums/reasoning-effort.enums.ts +++ b/tools/ui/src/lib/enums/reasoning-effort.enums.ts @@ -3,6 +3,7 @@ * These values are sent to the server and mapped to token budgets. */ export enum ReasoningEffort { + OFF = 'off', LOW = 'low', MEDIUM = 'medium', HIGH = 'high', diff --git a/tools/ui/src/lib/hooks/use-context-gauge.svelte.ts b/tools/ui/src/lib/hooks/use-context-gauge.svelte.ts new file mode 100644 index 000000000..e11e2f7ab --- /dev/null +++ b/tools/ui/src/lib/hooks/use-context-gauge.svelte.ts @@ -0,0 +1,295 @@ +/** + * Reactive state for the context usage gauge: resolves the active model, + * fetches its cached props, parses live server stats, and exposes per-turn + * read / fresh / cache / output and cumulative token counts. + */ + +import { + modelsStore, + modelOptions, + selectedModelId, + singleModelName +} from '$lib/stores/models.svelte'; +import { chatStore } from '$lib/stores/chat.svelte'; +import { activeMessages } from '$lib/stores/conversations.svelte'; +import { isRouterMode } from '$lib/stores/server.svelte'; +import { MessageRole } from '$lib/enums'; +import { STATS_UNITS } from '$lib/constants'; +import type { ChatMessageTimings, DatabaseMessage } from '$lib/types'; +import { useProcessingState } from './use-processing-state.svelte'; +import { + colorLevelFromPercent, + type ColorLevel +} from '$lib/components/app/chat/ChatForm/ChatFormContextGauge/context-gauge'; + +interface LiveStats { + freshTokens: number; + promptTokens: number; + cacheTokens: number; + outputTokens: number; +} + +export interface UseContextGaugeReturn { + readonly activeModelId: string | null; + readonly isActiveModelLoaded: boolean; + readonly isActiveModelLoading: boolean; + readonly contextTotal: number | null; + readonly contextUsed: number; + readonly currentRead: number; + readonly currentFresh: number; + readonly currentCache: number; + readonly currentOutput: number; + readonly kvTotal: number; + readonly cumulativeRead: number; + readonly cumulativeOutput: number; + readonly cumulativeCacheTotal: number; + readonly averageTokensPerSecond: number | null; + readonly contextPercent: number | null; + readonly colorLevel: ColorLevel; + readonly transientDetails: string[]; + readonly hasAnyUsage: boolean; + loadModel(): Promise; + startMonitoring(): void; +} + +function lastAssistantTimings(messages: DatabaseMessage[]): ChatMessageTimings | undefined { + for (let i = messages.length - 1; i >= 0; i--) { + const m = messages[i]; + if (m.role === MessageRole.ASSISTANT && m.timings) return m.timings; + } + return undefined; +} + +function deriveLiveStats( + state: ReturnType['processingState'] +): LiveStats | null { + if (!state || (state.status !== 'preparing' && state.status !== 'generating')) { + return null; + } + const promptTokens = state.promptTokens ?? 0; + const cacheTokens = state.cacheTokens ?? 0; + return { + freshTokens: promptTokens, + promptTokens: promptTokens + cacheTokens, + cacheTokens, + outputTokens: state.outputTokensUsed ?? 0 + }; +} + +const TRANSIENT_DETAILS_EXCLUDED_PREFIXES = ['Context:', 'Output:']; + +function filterTransientDetails(raw: string[]): string[] { + return raw.filter((detail) => { + if (TRANSIENT_DETAILS_EXCLUDED_PREFIXES.some((prefix) => detail.startsWith(prefix))) { + return false; + } + return !detail.includes(STATS_UNITS.TOKENS_PER_SECOND); + }); +} + +export function useContextGauge(): UseContextGaugeReturn { + const processingState = useProcessingState(); + + // Resolve the model the gauge reports context for: explicit selection > + // last assistant model > single-model mode (mirrors useChatScreenActiveModel). + const activeModelId = $derived.by(() => { + if (!isRouterMode()) { + return singleModelName(); + } + + const selectedId = selectedModelId(); + if (selectedId) { + const model = modelOptions().find((m) => m.id === selectedId); + if (model) return model.model; + } + + return chatStore.getConversationModel(activeMessages() as DatabaseMessage[]); + }); + + const isActiveModelLoaded = $derived( + activeModelId !== null && modelsStore.isModelLoaded(activeModelId) + ); + + const isActiveModelLoading = $derived( + activeModelId !== null && modelsStore.isModelOperationInProgress(activeModelId) + ); + + // Pull /props on demand so n_ctx surfaces before the first chat request. + $effect(() => { + if (activeModelId && isActiveModelLoaded) { + const cached = modelsStore.getModelProps(activeModelId); + if (!cached) { + void modelsStore.fetchModelProps(activeModelId); + } + } + }); + + const contextTotal = $derived.by(() => { + void modelsStore.propsCacheVersion; + return activeModelId ? modelsStore.getModelContextSize(activeModelId) : null; + }); + + const liveStats = $derived(deriveLiveStats(processingState.processingState)); + + const currentRead = $derived.by(() => { + const timings = lastAssistantTimings(activeMessages() as DatabaseMessage[]); + let read = 0; + if (timings) { + read = (timings.prompt_n ?? 0) + (timings.cache_n ?? 0); + } + // live.promptTokens is already the combined reading (prompt + cache), + // so do not also add live.cacheTokens. + if (liveStats && liveStats.promptTokens > 0) { + read = Math.max(read, liveStats.promptTokens); + } + return read; + }); + + const currentFresh = $derived.by(() => { + const timings = lastAssistantTimings(activeMessages() as DatabaseMessage[]); + const fresh = timings?.prompt_n ?? 0; + return Math.max(fresh, liveStats?.freshTokens ?? 0); + }); + + const currentCache = $derived.by(() => { + const timings = lastAssistantTimings(activeMessages() as DatabaseMessage[]); + const cached = timings?.cache_n ?? 0; + if (liveStats && liveStats.promptTokens > 0) { + return Math.max(cached, liveStats.cacheTokens); + } + return cached; + }); + + const currentOutput = $derived.by(() => { + if (liveStats && liveStats.outputTokens > 0) return liveStats.outputTokens; + const timings = lastAssistantTimings(activeMessages() as DatabaseMessage[]); + return timings?.predicted_n ?? 0; + }); + + const kvTotal = $derived(currentRead + currentOutput); + const contextUsed = $derived(currentRead + currentOutput); + + const cumulative = $derived.by(() => { + const messages = activeMessages() as DatabaseMessage[]; + + // Agentic sessions stamp the same agentic.llm totals onto every + // assistant message; cache_n is never per-turn so cache_total stays 0. + const agenticMessages = messages.filter( + (m) => m.role === MessageRole.ASSISTANT && m.timings?.agentic?.llm?.predicted_n != null + ); + + if (agenticMessages.length > 0) { + const llm = agenticMessages[agenticMessages.length - 1].timings!.agentic!.llm; + const output = llm.predicted_n ?? 0; + const outputMs = llm.predicted_ms ?? 0; + const averageTokensPerSecond = outputMs > 0 && output > 0 ? (output / outputMs) * 1000 : null; + return { + read: llm.prompt_n ?? 0, + output, + cacheTotal: 0, + averageTokensPerSecond + }; + } + + let read = 0; + let output = 0; + let outputMs = 0; + let cacheTotal = 0; + for (const m of messages) { + if (m.role !== MessageRole.ASSISTANT || !m.timings) continue; + read += m.timings.prompt_n ?? 0; + cacheTotal += m.timings.cache_n ?? 0; + output += m.timings.predicted_n ?? 0; + outputMs += m.timings.predicted_ms ?? 0; + } + const averageTokensPerSecond = outputMs > 0 && output > 0 ? (output / outputMs) * 1000 : null; + return { read, output, cacheTotal, averageTokensPerSecond }; + }); + + const contextPercent = $derived.by(() => { + if (contextTotal === null || contextTotal <= 0) return null; + return Math.round((contextUsed / contextTotal) * 100); + }); + + const colorLevel = $derived(colorLevelFromPercent(contextPercent)); + + // Drop lines the surrounding Context / Output / speed rows already render. + const transientDetails = $derived(filterTransientDetails(processingState.getTechnicalDetails())); + + const hasAnyUsage = $derived( + cumulative.read > 0 || + cumulative.output > 0 || + currentRead > 0 || + currentOutput > 0 || + cumulative.averageTokensPerSecond !== null || + transientDetails.length > 0 + ); + + async function loadModel() { + if (!activeModelId || isActiveModelLoading) return; + try { + await modelsStore.loadModel(activeModelId); + } catch { + // toast already surfaced by modelsStore.loadModel + } + } + + return { + get activeModelId() { + return activeModelId; + }, + get isActiveModelLoaded() { + return isActiveModelLoaded; + }, + get isActiveModelLoading() { + return isActiveModelLoading; + }, + get contextTotal() { + return contextTotal; + }, + get contextUsed() { + return contextUsed; + }, + get currentRead() { + return currentRead; + }, + get currentFresh() { + return currentFresh; + }, + get currentCache() { + return currentCache; + }, + get currentOutput() { + return currentOutput; + }, + get kvTotal() { + return kvTotal; + }, + get cumulativeRead() { + return cumulative.read; + }, + get cumulativeOutput() { + return cumulative.output; + }, + get cumulativeCacheTotal() { + return cumulative.cacheTotal; + }, + get averageTokensPerSecond() { + return cumulative.averageTokensPerSecond; + }, + get contextPercent() { + return contextPercent; + }, + get colorLevel() { + return colorLevel; + }, + get transientDetails() { + return transientDetails; + }, + get hasAnyUsage() { + return hasAnyUsage; + }, + loadModel, + startMonitoring: () => processingState.startMonitoring() + }; +} diff --git a/tools/ui/src/lib/hooks/use-mcp-recommendations.svelte.ts b/tools/ui/src/lib/hooks/use-mcp-recommendations.svelte.ts index c8a85fa8f..4f4c2c782 100644 --- a/tools/ui/src/lib/hooks/use-mcp-recommendations.svelte.ts +++ b/tools/ui/src/lib/hooks/use-mcp-recommendations.svelte.ts @@ -54,11 +54,6 @@ export function useMcpRecommendations() { // effect, and we must not wipe the timeout that was just scheduled. if (checked) return; - if (mcpStore.optedInRecommendationIds.size > 0) { - checked = true; - return; - } - const hasRecommendations = mcpStore .getServers() .some((server) => RECOMMENDED_MCP_SERVER_IDS.has(server.id)); diff --git a/tools/ui/src/lib/hooks/use-processing-state.svelte.ts b/tools/ui/src/lib/hooks/use-processing-state.svelte.ts index f28031972..9fbda75d6 100644 --- a/tools/ui/src/lib/hooks/use-processing-state.svelte.ts +++ b/tools/ui/src/lib/hooks/use-processing-state.svelte.ts @@ -1,5 +1,4 @@ import { activeProcessingState } from '$lib/stores/chat.svelte'; -import { config } from '$lib/stores/settings.svelte'; import { STATS_UNITS } from '$lib/constants'; import type { ApiProcessingState, LiveProcessingStats, LiveGenerationStats } from '$lib/types'; @@ -46,7 +45,6 @@ export function useProcessingState(): UseProcessingStateReturn { return activeProcessingState(); }); - // Track last known state for keepStatsVisible functionality $effect(() => { if (processingState && isMonitoring) { lastKnownState = processingState; @@ -88,14 +86,8 @@ export function useProcessingState(): UseProcessingStateReturn { function stopMonitoring(): void { if (!isMonitoring) return; - isMonitoring = false; - // Only clear last known state if keepStatsVisible is disabled - const currentConfig = config(); - if (!currentConfig.keepStatsVisible) { - lastKnownState = null; - lastKnownProcessingStats = null; - } + isMonitoring = false; } function getProcessingMessage(): string { diff --git a/tools/ui/src/lib/services/chat.service.ts b/tools/ui/src/lib/services/chat.service.ts index dbd0a94f1..15f4cc3b5 100644 --- a/tools/ui/src/lib/services/chat.service.ts +++ b/tools/ui/src/lib/services/chat.service.ts @@ -340,6 +340,7 @@ export class ChatService { if (stream && conversationId) { headers['X-Conversation-Id'] = streamIdentity(conversationId, options.model); } + const response = await fetch(API_CHAT.COMPLETIONS, { method: 'POST', headers, @@ -1011,7 +1012,7 @@ export class ChatService { * * @param response - The fetch Response object containing the JSON data * @param onComplete - Optional callback invoked when response is successfully parsed - * @param onError - Optional callback invoked if an error occurs during parsing + * @param onError - Optional callback invoked if an error occurs while parsing * @returns {Promise} Promise that resolves to the generated content string * @throws {Error} if the response cannot be parsed or is malformed */ diff --git a/tools/ui/src/lib/services/migration.service.ts b/tools/ui/src/lib/services/migration.service.ts index d7709bc6b..981283be9 100644 --- a/tools/ui/src/lib/services/migration.service.ts +++ b/tools/ui/src/lib/services/migration.service.ts @@ -564,9 +564,9 @@ const configTypesMigration: Migration = { const config = JSON.parse(configRaw); let changed = false; - // Pre-schema configs persisted booleans as the strings "true"/"false", which the - // strict server schema now rejects. Coerce those back to real booleans. No config - // string field holds exactly "true"/"false", so the match is unambiguous. + // Pre-schema configs persisted booleans as "true"/"false" strings; the strict server + // schema rejects them. No config string field holds exactly "true"/"false", so the + // match is unambiguous. for (const key of Object.keys(config)) { if (config[key] === 'true') { config[key] = true; diff --git a/tools/ui/src/lib/stores/agentic.svelte.ts b/tools/ui/src/lib/stores/agentic.svelte.ts index 27491257a..1a677602f 100644 --- a/tools/ui/src/lib/stores/agentic.svelte.ts +++ b/tools/ui/src/lib/stores/agentic.svelte.ts @@ -477,7 +477,7 @@ class AgenticStore { conversationId: string; messages: ApiChatMessageData[]; options: AgenticFlowOptions; - tools: ReturnType; + tools: ReturnType; agenticConfig: AgenticConfig; callbacks: AgenticFlowCallbacks; signal?: AbortSignal; diff --git a/tools/ui/src/lib/stores/chat.svelte.ts b/tools/ui/src/lib/stores/chat.svelte.ts index 6949047b5..fcd07c4fd 100644 --- a/tools/ui/src/lib/stores/chat.svelte.ts +++ b/tools/ui/src/lib/stores/chat.svelte.ts @@ -60,6 +60,7 @@ import { ErrorDialogType, MessageRole, MessageType, + ReasoningEffort, StreamConnectionState } from '$lib/enums'; @@ -2334,7 +2335,8 @@ class ChatStore { if (currentConfig.excludeReasoningFromContext) apiOptions.excludeReasoningFromContext = true; apiOptions.enableThinking = conversationsStore.getThinkingEnabled(); - apiOptions.reasoningEffort = conversationsStore.getReasoningEffort(); + const effort = conversationsStore.getReasoningEffort(); + if (effort !== ReasoningEffort.OFF) apiOptions.reasoningEffort = effort; if (hasValue(currentConfig.temperature)) apiOptions.temperature = Number(currentConfig.temperature); diff --git a/tools/ui/src/lib/stores/conversations.svelte.ts b/tools/ui/src/lib/stores/conversations.svelte.ts index 486202207..b29a900fe 100644 --- a/tools/ui/src/lib/stores/conversations.svelte.ts +++ b/tools/ui/src/lib/stores/conversations.svelte.ts @@ -47,7 +47,6 @@ import { NON_ALPHANUMERIC_REGEX, MULTIPLE_UNDERSCORE_REGEX, SETTINGS_KEYS, - THINKING_ENABLED_DEFAULT_LOCALSTORAGE_KEY, REASONING_EFFORT_DEFAULT_LOCALSTORAGE_KEY } from '$lib/constants'; @@ -84,11 +83,16 @@ class ConversationsStore { /** Pending MCP server overrides for new conversations (before first message) */ pendingMcpServerOverrides = $state(ConversationsStore.loadMcpDefaults()); - /** Global (non-conversation-specific) thinking toggle default */ - pendingThinkingEnabled = $state(ConversationsStore.loadThinkingDefaults()); + /** Global (non-conversation-specific) thinking toggle default, derived from reasoning effort */ + pendingThinkingEnabled = $state(false); /** Global (non-conversation-specific) reasoning effort default */ - pendingReasoningEffort = $state(ConversationsStore.loadReasoningEffortDefault()); + pendingReasoningEffort = $state( + ConversationsStore.loadReasoningEffortDefault() + ); + + /** Last non-off reasoning effort, restored when re-enabling thinking globally */ + private lastNonOffEffort: ReasoningEffort | null = null; private static loadMcpDefaults(): McpServerOverride[] { const raw = config()[SETTINGS_KEYS.MCP_DEFAULT_SERVER_OVERRIDES]; @@ -112,35 +116,14 @@ class ConversationsStore { settingsStore.updateConfig(SETTINGS_KEYS.MCP_DEFAULT_SERVER_OVERRIDES, JSON.stringify(plain)); } - /** Load thinking-enabled default from localStorage */ - private static loadThinkingDefaults(): boolean { - if (typeof globalThis.localStorage === 'undefined') return true; - try { - const raw = localStorage.getItem(THINKING_ENABLED_DEFAULT_LOCALSTORAGE_KEY); - if (!raw) return true; - return raw === 'true'; - } catch { - return true; - } - } - - /** Persist thinking-enabled default to localStorage */ - private saveThinkingDefaults(): void { - if (typeof globalThis.localStorage === 'undefined') return; - localStorage.setItem( - THINKING_ENABLED_DEFAULT_LOCALSTORAGE_KEY, - this.pendingThinkingEnabled ? 'true' : 'false' - ); - } - /** Load reasoning effort default from localStorage */ - private static loadReasoningEffortDefault(): ReasoningEffort { - if (typeof globalThis.localStorage === 'undefined') return ReasoningEffort.MEDIUM; + private static loadReasoningEffortDefault(): ReasoningEffort | ReasoningEffort.OFF { + if (typeof globalThis.localStorage === 'undefined') return ReasoningEffort.OFF; try { const raw = localStorage.getItem(REASONING_EFFORT_DEFAULT_LOCALSTORAGE_KEY); - return (raw as ReasoningEffort) || ReasoningEffort.MEDIUM; + return (raw as ReasoningEffort | ReasoningEffort.OFF) || ReasoningEffort.OFF; } catch { - return ReasoningEffort.MEDIUM; + return ReasoningEffort.OFF; } } @@ -303,10 +286,17 @@ class ConversationsStore { this.pendingMcpServerOverrides = []; } - // Inherit global thinking default into the new conversation - conversation.thinkingEnabled = this.pendingThinkingEnabled; + // Inherit global thinking/reasoning defaults into the new conversation + const thinkingEnabled = this.getThinkingEnabled(); + conversation.thinkingEnabled = thinkingEnabled; + conversation.reasoningEffort = + this.pendingReasoningEffort === ReasoningEffort.OFF ? undefined : this.pendingReasoningEffort; await DatabaseService.updateConversation(conversation.id, { - thinkingEnabled: this.pendingThinkingEnabled + thinkingEnabled, + reasoningEffort: + this.pendingReasoningEffort === ReasoningEffort.OFF + ? undefined + : this.pendingReasoningEffort }); this.conversations = [conversation, ...this.conversations]; @@ -332,7 +322,6 @@ class ConversationsStore { } this.pendingMcpServerOverrides = []; - this.pendingThinkingEnabled = ConversationsStore.loadThinkingDefaults(); this.activeConversation = conversation; if (conversation.currNode) { @@ -363,7 +352,7 @@ class ConversationsStore { this.activeMessages = []; // reload defaults so new chats inherit persisted state this.pendingMcpServerOverrides = ConversationsStore.loadMcpDefaults(); - this.pendingThinkingEnabled = ConversationsStore.loadThinkingDefaults(); + this.pendingReasoningEffort = ConversationsStore.loadReasoningEffortDefault(); } /** @@ -794,9 +783,11 @@ class ConversationsStore { */ getThinkingEnabled(): boolean { if (this.activeConversation) { - return this.activeConversation.thinkingEnabled ?? this.pendingThinkingEnabled; + if (this.activeConversation.thinkingEnabled !== undefined) { + return this.activeConversation.thinkingEnabled; + } } - return this.pendingThinkingEnabled; + return this.getReasoningEffort() !== ReasoningEffort.OFF; } /** @@ -806,8 +797,17 @@ class ConversationsStore { */ async setThinkingEnabled(enabled: boolean): Promise { if (!this.activeConversation) { - this.pendingThinkingEnabled = enabled; - this.saveThinkingDefaults(); + if (enabled) { + const effort = this.lastNonOffEffort ?? ReasoningEffort.LOW; + this.pendingReasoningEffort = effort; + this.saveReasoningEffortDefaults(); + } else { + if (this.pendingReasoningEffort !== ReasoningEffort.OFF) { + this.lastNonOffEffort = this.pendingReasoningEffort; + } + this.pendingReasoningEffort = ReasoningEffort.OFF; + this.saveReasoningEffortDefaults(); + } return; } @@ -831,7 +831,7 @@ class ConversationsStore { * Gets the effective reasoning effort for the active conversation. * Returns the conversation override if set, otherwise the global default. */ - getReasoningEffort(): ReasoningEffort { + getReasoningEffort(): ReasoningEffort | ReasoningEffort.OFF { if (this.activeConversation) { return this.activeConversation.reasoningEffort ?? this.pendingReasoningEffort; } diff --git a/tools/ui/src/lib/stores/mcp.svelte.ts b/tools/ui/src/lib/stores/mcp.svelte.ts index 37eb56332..a53463ee9 100644 --- a/tools/ui/src/lib/stores/mcp.svelte.ts +++ b/tools/ui/src/lib/stores/mcp.svelte.ts @@ -12,21 +12,22 @@ * - Lifecycle management (initialize, shutdown) * - Multi-server coordination * - Tool name conflict detection and resolution - * - OpenAI-compatible tool definition generation * - Automatic tool-to-server routing * - Health checks * + * MCP connection state and raw `Tool[]` per server are owned here; the + * OpenAI-compatible wire format for those tools is built in `toolsStore` + * (see {@link toolsStore.mcpEntries} / {@link toolsStore.getEnabledToolsForLLM}). + * * @see MCPService in services/mcp.service.ts for protocol operations */ import { browser } from '$app/environment'; -import { SvelteSet } from 'svelte/reactivity'; import { SETTINGS_KEYS } from '$lib/constants'; import { MCPService } from '$lib/services/mcp.service'; import { config, settingsStore } from '$lib/stores/settings.svelte'; import { mcpResourceStore } from '$lib/stores/mcp-resources.svelte'; import { serverStore } from '$lib/stores/server.svelte'; -import { conversationsStore } from '$lib/stores/conversations.svelte'; import { mode } from 'mode-watcher'; import { parseMcpServerSettings, @@ -40,9 +41,7 @@ import { HealthCheckStatus, MCPRefType, ColorMode, - UrlProtocol, - JsonSchemaType, - ToolCallType + UrlProtocol } from '$lib/enums'; import { DEFAULT_CACHE_TTL_MS, @@ -53,12 +52,10 @@ import { MCP_RECONNECT_BACKOFF_MULTIPLIER, MCP_RECONNECT_INITIAL_DELAY, MCP_RECONNECT_MAX_DELAY, - MCP_RECONNECT_ATTEMPT_TIMEOUT_MS, - RECOMMENDED_MCP_SERVER_IDS + MCP_RECONNECT_ATTEMPT_TIMEOUT_MS } from '$lib/constants'; import type { MCPToolCall, - OpenAIToolDefinition, ServerStatus, ToolExecutionResult, MCPClientConfig, @@ -582,30 +579,10 @@ class MCPStore { } /** - * Recommended MCP server IDs the user opted in to via per-chat overrides. - * Single source of truth for "which recommendations has the user accepted", - * shared by the recommendations hook and the visible-servers getter. - */ - get optedInRecommendationIds(): ReadonlySet { - const ids = new SvelteSet(); - for (const override of conversationsStore.pendingMcpServerOverrides) { - if (RECOMMENDED_MCP_SERVER_IDS.has(override.serverId) && override.enabled) { - ids.add(override.serverId); - } - } - return ids; - } - - /** - * MCP servers selectable in chat-add UIs and the settings page: - * enabled in settings and either non-recommended or explicitly opted in. + * MCP servers selectable in chat-add UIs and the settings page. */ get visibleMcpServers(): MCPServerSettingsEntry[] { - const optedIn = this.optedInRecommendationIds; - return this.getServersSorted().filter( - (server) => - server.enabled && (!RECOMMENDED_MCP_SERVER_IDS.has(server.id) || optedIn.has(server.id)) - ); + return this.getServersSorted().filter((server) => server.enabled); } async ensureInitialized(perChatOverrides?: McpServerOverride[]): Promise { @@ -979,73 +956,6 @@ class MCPStore { } } - getToolDefinitionsForLLM(): OpenAIToolDefinition[] { - const tools: OpenAIToolDefinition[] = []; - - for (const connection of this.connections.values()) { - for (const tool of connection.tools) { - const rawSchema = (tool.inputSchema as Record) ?? { - type: JsonSchemaType.OBJECT, - properties: {}, - required: [] - }; - - tools.push({ - type: ToolCallType.FUNCTION as const, - function: { - name: tool.name, - description: tool.description, - parameters: this.normalizeSchemaProperties(rawSchema) - } - }); - } - } - - return tools; - } - - private normalizeSchemaProperties(schema: Record): Record { - if (!schema || typeof schema !== 'object') { - return schema; - } - - const normalized = { ...schema }; - if (normalized.properties && typeof normalized.properties === 'object') { - const props = normalized.properties as Record>; - const normalizedProps: Record> = {}; - for (const [key, prop] of Object.entries(props)) { - if (!prop || typeof prop !== 'object') { - normalizedProps[key] = prop; - continue; - } - const normalizedProp = { ...prop }; - if (!normalizedProp.type && normalizedProp.default !== undefined) { - const defaultVal = normalizedProp.default; - if (typeof defaultVal === 'string') normalizedProp.type = 'string'; - else if (typeof defaultVal === 'number') - normalizedProp.type = Number.isInteger(defaultVal) ? 'integer' : 'number'; - else if (typeof defaultVal === 'boolean') normalizedProp.type = 'boolean'; - else if (Array.isArray(defaultVal)) normalizedProp.type = 'array'; - else if (typeof defaultVal === 'object' && defaultVal !== null) - normalizedProp.type = 'object'; - } - if (normalizedProp.properties) - Object.assign( - normalizedProp, - this.normalizeSchemaProperties(normalizedProp as Record) - ); - if (normalizedProp.items && typeof normalizedProp.items === 'object') - normalizedProp.items = this.normalizeSchemaProperties( - normalizedProp.items as Record - ); - normalizedProps[key] = normalizedProp; - } - normalized.properties = normalizedProps; - } - - return normalized; - } - getToolNames(): string[] { return Array.from(this.toolsIndex.keys()); } diff --git a/tools/ui/src/lib/stores/tools.svelte.ts b/tools/ui/src/lib/stores/tools.svelte.ts index a63781988..dcaab5f42 100644 --- a/tools/ui/src/lib/stores/tools.svelte.ts +++ b/tools/ui/src/lib/stores/tools.svelte.ts @@ -13,33 +13,6 @@ import { import { SvelteMap, SvelteSet } from 'svelte/reactivity'; /** Stable selection identity for a tool, shared by the disabled set and the permission store */ -function toolKey(source: ToolSource, name: string, serverId?: string): string { - switch (source) { - case ToolSource.MCP: - return serverId ? `mcp-${serverId}:${name}` : `mcp:${name}`; - case ToolSource.CUSTOM: - return `custom:${name}`; - case ToolSource.FRONTEND: - return `frontend:${name}`; - default: - return `builtin:${name}`; - } -} - -function mcpDefinition( - name: string, - description: string | undefined, - schema?: Record -): OpenAIToolDefinition { - return { - type: ToolCallType.FUNCTION, - function: { - name, - description, - parameters: schema ?? { type: JsonSchemaType.OBJECT, properties: {}, required: [] } - } - }; -} class ToolsStore { private _builtinTools = $state([]); @@ -77,12 +50,96 @@ class ToolsStore { } } + private toolKey(source: ToolSource, name: string, serverId?: string): string { + switch (source) { + case ToolSource.MCP: + return serverId ? `mcp-${serverId}:${name}` : `mcp:${name}`; + case ToolSource.CUSTOM: + return `custom:${name}`; + case ToolSource.FRONTEND: + return `frontend:${name}`; + default: + return `builtin:${name}`; + } + } + + private inferTypeFromDefault(value: unknown): string | undefined { + if (typeof value === 'string') return 'string'; + if (typeof value === 'boolean') return 'boolean'; + if (typeof value === 'number') return Number.isInteger(value) ? 'integer' : 'number'; + if (Array.isArray(value)) return 'array'; + if (value !== null && typeof value === 'object') return 'object'; + return undefined; + } + + /** + * Recursively normalize a JSON Schema object: infers `type` from `default` + * for properties / items that omit it, and descends into nested `properties` + * and `items`. Returns a new object -- does not mutate the input. + */ + private normalizeJsonSchema(schema: Record): Record { + if (!schema || typeof schema !== 'object') return schema; + + const normalized: Record = { ...schema }; + + if (normalized.properties && typeof normalized.properties === 'object') { + const props = normalized.properties as Record>; + const normalizedProps: Record> = {}; + for (const [key, prop] of Object.entries(props)) { + if (!prop || typeof prop !== 'object') { + normalizedProps[key] = prop; + continue; + } + + const normalizedProp: Record = { ...prop }; + + if (!normalizedProp.type && normalizedProp.default !== undefined) { + const inferred = this.inferTypeFromDefault(normalizedProp.default); + if (inferred) normalizedProp.type = inferred; + } + + if (normalizedProp.properties) { + Object.assign( + normalizedProp, + this.normalizeJsonSchema(normalizedProp as Record) + ); + } + + if (normalizedProp.items && typeof normalizedProp.items === 'object') { + normalizedProp.items = this.normalizeJsonSchema( + normalizedProp.items as Record + ); + } + + normalizedProps[key] = normalizedProp; + } + normalized.properties = normalizedProps; + } + + return normalized; + } + + private mcpDefinition( + name: string, + description: string | undefined, + schema?: Record + ): OpenAIToolDefinition { + return { + type: ToolCallType.FUNCTION, + function: { + name, + description, + parameters: schema ?? { type: JsonSchemaType.OBJECT, properties: {}, required: [] } + } + }; + } + get builtinTools(): OpenAIToolDefinition[] { return this._builtinTools; } get mcpTools(): OpenAIToolDefinition[] { - return mcpStore.getToolDefinitionsForLLM(); + return this.mcpEntries().map((e) => e.definition); } get frontendTools(): OpenAIToolDefinition[] { @@ -124,11 +181,22 @@ class ToolsStore { for (const [serverId, connection] of connections) { const serverName = mcpStore.getServerDisplayName(serverId); for (const tool of connection.tools) { - const schema = (tool.inputSchema as Record) ?? undefined; + const rawSchema = (tool.inputSchema as Record) ?? { + type: JsonSchemaType.OBJECT, + properties: {}, + required: [] + }; out.push({ serverId, serverName, - definition: mcpDefinition(tool.name, tool.description, schema) + definition: { + type: ToolCallType.FUNCTION, + function: { + name: tool.name, + description: tool.description, + parameters: this.normalizeJsonSchema(rawSchema) + } + } }); } } @@ -138,7 +206,7 @@ class ToolsStore { out.push({ serverId, serverName, - definition: mcpDefinition(tool.name, tool.description) + definition: this.mcpDefinition(tool.name, tool.description) }); } } @@ -160,14 +228,18 @@ class ToolsStore { for (const def of this._builtinTools) { const name = def.function.name; - push({ source: ToolSource.BUILTIN, key: toolKey(ToolSource.BUILTIN, name), definition: def }); + push({ + source: ToolSource.BUILTIN, + key: this.toolKey(ToolSource.BUILTIN, name), + definition: def + }); } for (const def of this.frontendTools) { const name = def.function.name; push({ source: ToolSource.FRONTEND, - key: toolKey(ToolSource.FRONTEND, name), + key: this.toolKey(ToolSource.FRONTEND, name), definition: def }); } @@ -178,14 +250,18 @@ class ToolsStore { source: ToolSource.MCP, serverId, serverName, - key: toolKey(ToolSource.MCP, name, serverId), + key: this.toolKey(ToolSource.MCP, name, serverId), definition }); } for (const def of this.customTools) { const name = def.function.name; - push({ source: ToolSource.CUSTOM, key: toolKey(ToolSource.CUSTOM, name), definition: def }); + push({ + source: ToolSource.CUSTOM, + key: this.toolKey(ToolSource.CUSTOM, name), + definition: def + }); } return entries; @@ -233,7 +309,8 @@ class ToolsStore { /** * Enabled tool definitions for sending to the LLM. - * MCP tools keep their normalized schemas from mcpStore. + * MCP tool schemas are normalized here so the wire payload is consistent + * across all four sources (built-in, frontend/sandbox, MCP, custom JSON). * The API identifies tools by name, so a name is sent at most once. */ getEnabledToolsForLLM(): OpenAIToolDefinition[] { @@ -256,7 +333,8 @@ class ToolsStore { for (const def of this._builtinTools) take(def); for (const def of this.frontendTools) take(def); - for (const def of mcpStore.getToolDefinitionsForLLM()) take(def); + // mcpEntries() over mcpStore directly so wire shape stays normalized and aligned with the tools UI. + for (const entry of this.mcpEntries()) take(entry.definition); for (const def of this.customTools) take(def); return result; @@ -308,15 +386,17 @@ class ToolsStore { const connection = mcpStore.getConnections().get(serverId); if (!connection) return; for (const tool of connection.tools) { - this._disabledTools.delete(toolKey(ToolSource.MCP, tool.name, serverId)); + this._disabledTools.delete(this.toolKey(ToolSource.MCP, tool.name, serverId)); } this.persistDisabledTools(); } toggleGroup(group: ToolGroup): void { const allEnabled = group.tools.every((t) => this.isToolEnabled(t.key)); + const target = !allEnabled; for (const tool of group.tools) { - this.setToolEnabled(tool.key, !allEnabled); + if (target) this._disabledTools.delete(tool.key); + else this._disabledTools.add(tool.key); } this.persistDisabledTools(); } @@ -332,7 +412,7 @@ class ToolsStore { tools: { name: string; description?: string }[]; }[] { const result: ReturnType = []; - for (const server of mcpStore.getServersSorted().filter((s) => s.enabled)) { + for (const server of mcpStore.visibleMcpServers) { const health = mcpStore.getHealthCheckState(server.id); if (health.status === HealthCheckStatus.SUCCESS && health.tools.length > 0) { result.push({ diff --git a/tools/ui/src/lib/utils/branching.ts b/tools/ui/src/lib/utils/branching.ts index c40abbdd6..6ff701318 100644 --- a/tools/ui/src/lib/utils/branching.ts +++ b/tools/ui/src/lib/utils/branching.ts @@ -111,15 +111,7 @@ function findLeafNodeInMap( } /** - * Convenience wrapper around {@link findLeafNodeInMap} for callers that only have - * a flat message array. - * - * Finds the leaf node (message with no children) for a given message branch. - * Traverses down the tree following the last child until reaching a leaf. - * - * @param messages - All messages in the conversation - * @param messageId - Starting message ID to find leaf for - * @returns The leaf node ID, or the original messageId if no children + * Convenience wrapper around {@link findLeafNodeInMap} for callers that have a flat message array. */ export function findLeafNode(messages: readonly DatabaseMessage[], messageId: string): string { const nodeMap = new Map(messages.map((msg) => [msg.id, msg] as const)); @@ -225,7 +217,6 @@ export function getMessageSiblings( /** * Builds sibling information for every message in a conversation. - * A single node map is shared across all lookups for O(1) access. * * @param messages - All messages in the conversation * @returns Map of message ID to its sibling information From f296fdfbed71e900a3e0d6579673960e6a560654 Mon Sep 17 00:00:00 2001 From: rankaiyx Date: Wed, 8 Jul 2026 15:45:28 +0800 Subject: [PATCH 15/17] common: auto-create prompts-log-dir at argument parsing, so all tools using the flag benefit (#25322) --- common/arg.cpp | 8 +++++++- tools/server/README.md | 2 +- 2 files changed, 8 insertions(+), 2 deletions(-) diff --git a/common/arg.cpp b/common/arg.cpp index fdf58614b..a26093199 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -27,6 +27,7 @@ #include #include #include +#include #include #include #include @@ -3451,9 +3452,14 @@ common_params_context common_params_parser_init(common_params & params, llama_ex ).set_env("LLAMA_ARG_LOG_FILE")); add_opt(common_arg( {"--log-prompts-dir"}, "PATH", - "Log prompts to directory (only used for debugging, default: disabled)", + "Log prompts to directory (auto-created if not present; only used for debugging, default: disabled)", [](common_params & params, const std::string & value) { params.path_prompts_log_dir = value; + std::error_code ec; + std::filesystem::create_directories(value, ec); + if (ec) { + fprintf(stderr, "warning: failed to create prompts-log-dir '%s': %s\n", value.c_str(), ec.message().c_str()); + } } ).set_examples({LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI})); add_opt(common_arg( diff --git a/tools/server/README.md b/tools/server/README.md index 501b66123..365a94465 100644 --- a/tools/server/README.md +++ b/tools/server/README.md @@ -228,7 +228,7 @@ For the full list of features, please refer to [server's changelog](https://gith | `-sps, --slot-prompt-similarity SIMILARITY` | how much the prompt of a request must match the prompt of a slot in order to use that slot (default: 0.10, 0.0 = disabled) | | `--lora-init-without-apply` | load LoRA adapters without applying them (apply later via POST /lora-adapters) (default: disabled) | | `--sleep-idle-seconds SECONDS` | number of seconds of idleness after which the server will sleep (default: -1; -1 = disabled) | -| `--log-prompts-dir PATH` | Log prompts to directory (only used for debugging, default: disabled) | +| `--log-prompts-dir PATH` | Log prompts to directory (auto-created if not present; only used for debugging, default: disabled) | | `--spec-draft-hf, -hfd, -hfrd, --hf-repo-draft /[:quant]` | Same as --hf-repo, but for the draft model (default: unused)
(env: LLAMA_ARG_SPEC_DRAFT_HF_REPO) | | `--spec-draft-threads, -td, --threads-draft N` | number of threads to use during generation (default: same as --threads) | | `--spec-draft-threads-batch, -tbd, --threads-batch-draft N` | number of threads to use during batch and prompt processing (default: same as --threads-draft) | From 230ea9d214320c5e79cc8166ed708ac60514c71e Mon Sep 17 00:00:00 2001 From: Aman Gupta Date: Wed, 8 Jul 2026 15:55:19 +0800 Subject: [PATCH 16/17] llama-batch: add n_keep_tail in split_equal for recurrent models (#25278) --- src/llama-batch.cpp | 74 ++++++++++++++++++++++++++++++-- src/llama-batch.h | 3 +- src/llama-kv-cache-dsa.cpp | 2 +- src/llama-kv-cache-dsv4.cpp | 2 +- src/llama-kv-cache-iswa.cpp | 2 +- src/llama-kv-cache.cpp | 2 +- src/llama-memory-hybrid-iswa.cpp | 18 ++++---- src/llama-memory-hybrid.cpp | 18 ++++---- src/llama-memory-recurrent.cpp | 15 +++---- src/models/delta-net-base.cpp | 4 +- 10 files changed, 102 insertions(+), 38 deletions(-) diff --git a/src/llama-batch.cpp b/src/llama-batch.cpp index 6bf76939c..5436717c4 100644 --- a/src/llama-batch.cpp +++ b/src/llama-batch.cpp @@ -505,7 +505,7 @@ llama_ubatch llama_batch_allocr::split_simple(uint32_t n_ubatch) { return ubatch_add(idxs, idxs.size(), false); } -llama_ubatch llama_batch_allocr::split_equal(uint32_t n_ubatch, bool sequential) { +llama_ubatch llama_batch_allocr::split_equal(uint32_t n_ubatch, bool sequential, uint32_t n_keep_tail) { if (sequential && has_cpl) { LLAMA_LOG_ERROR("%s: sequential split is not supported when there are coupled sequences in the input batch (you may need to use the -kvu flag)\n", __func__); @@ -548,7 +548,7 @@ llama_ubatch llama_batch_allocr::split_equal(uint32_t n_ubatch, bool sequential) } } - const uint32_t n_seqs = cur_seq_set.size(); + uint32_t n_seqs = cur_seq_set.size(); // we are done if (n_seqs == 0) { @@ -569,7 +569,7 @@ llama_ubatch llama_batch_allocr::split_equal(uint32_t n_ubatch, bool sequential) std::vector idxs_per_seq(n_seqs); while (true) { - // we can only add new n_seq_tokens tokens if all the sequence sets have at least one more unused token and + // we can only add new n_seq_tokens tokens if all the sequence sets have at least 1 more unused tokens and // if we haven't reached n_ubatch bool can_expand = true; @@ -600,6 +600,72 @@ llama_ubatch llama_batch_allocr::split_equal(uint32_t n_ubatch, bool sequential) } } + // if n_keep_tail > 0, keep only the seqs that either finish in this ubatch or have at least + // n_keep_tail tokens remaining for a future ubatch, so that the trailing n_keep_tail tokens + // of each seq are never split across ubatches + if (n_keep_tail > 0) { + GGML_ASSERT(n_ubatch > n_keep_tail); + + auto n_remaining = [&](uint32_t s) { + return (uint32_t) (seq_set_map[cur_seq_set[s]].size() - cur_idx[s]); + }; + + // keep the longest prefix of seqs that satisfy the constraint, to preserve sequential seq ids + uint32_t n_keep = 0; + while (n_keep < n_seqs) { + const uint32_t remaining = n_remaining(n_keep); + + if (remaining != 0 && remaining < n_keep_tail) { + break; + } + + n_keep++; + } + + // all seqs violate the constraint - resolve the first one directly and emit it alone + if (n_keep == 0) { + auto & idxs = idxs_per_seq[0]; + + const auto & seq_idxs = seq_set_map[cur_seq_set[0]]; + + if (idxs.size() + n_remaining(0) <= n_ubatch) { + // extend the seq to completion + while (n_remaining(0) > 0) { + const int32_t idx = seq_idxs[cur_idx[0]]; + + idxs.push_back(idx); + + used[idx] = true; + ++n_used; + + ++cur_idx[0]; + } + } else { + // truncate the seq so that at least n_keep_tail tokens remain + while (n_remaining(0) < n_keep_tail) { + used[idxs.back()] = false; + --n_used; + + idxs.pop_back(); + + --cur_idx[0]; + } + } + + n_keep = 1; + } + + // return the tokens of the deferred seqs back to the pool + for (uint32_t s = n_keep; s < n_seqs; ++s) { + for (const int32_t idx : idxs_per_seq[s]) { + used[idx] = false; + --n_used; + } + } + + n_seqs = n_keep; + } + // concat the per-sequence-set lists std::vector idxs; @@ -814,7 +880,7 @@ void llama_batch_allocr::ubatch_print(const llama_ubatch & ubatch, int debug) { LLAMA_LOG_DEBUG("%s: output = %p\n", __func__, (void *) ubatch.output); LLAMA_LOG_DEBUG("%s: n_outputs = %d\n", __func__, n_outputs); - if (debug > 1) { + if (debug > 0) { int seq_id_max = 0; for (uint32_t i = 0; i < ubatch.n_tokens; ++i) { for (int s = 0; s < ubatch.n_seq_id[i]; ++s) { diff --git a/src/llama-batch.h b/src/llama-batch.h index f77520e86..a3d1889d4 100644 --- a/src/llama-batch.h +++ b/src/llama-batch.h @@ -104,7 +104,8 @@ public: // make ubatches of equal-length sequences sets // if sequential == true, the tokens in the ubatch will have increasing sequential sequence ids - llama_ubatch split_equal(uint32_t n_ubatch, bool sequential); + // n_keep_tail = minimum trailing tokens of a seq that must land in the same ubatch + llama_ubatch split_equal(uint32_t n_ubatch, bool sequential, uint32_t n_keep_tail); // sequence-set-wise split - each ubatch contains a single sequence-set llama_ubatch split_seq(uint32_t n_ubatch); diff --git a/src/llama-kv-cache-dsa.cpp b/src/llama-kv-cache-dsa.cpp index 916ab6537..241c50365 100644 --- a/src/llama-kv-cache-dsa.cpp +++ b/src/llama-kv-cache-dsa.cpp @@ -113,7 +113,7 @@ llama_memory_context_ptr llama_kv_cache_dsa::init_batch( std::vector ubatches; while (true) { - auto ubatch = n_stream == 1 ? balloc.split_simple(n_ubatch) : balloc.split_equal(n_ubatch, true); + auto ubatch = n_stream == 1 ? balloc.split_simple(n_ubatch) : balloc.split_equal(n_ubatch, true, 0); if (ubatch.n_tokens == 0) { break; diff --git a/src/llama-kv-cache-dsv4.cpp b/src/llama-kv-cache-dsv4.cpp index 3a698d719..9fccf347e 100644 --- a/src/llama-kv-cache-dsv4.cpp +++ b/src/llama-kv-cache-dsv4.cpp @@ -1110,7 +1110,7 @@ llama_memory_context_ptr llama_kv_cache_dsv4::init_batch( if (has_coupled) { ubatch = balloc.split_seq(n_ubatch); } else { - ubatch = balloc.split_equal(n_ubatch, raw_per_seq || comp_per_seq); + ubatch = balloc.split_equal(n_ubatch, raw_per_seq || comp_per_seq, 0); } if (ubatch.n_tokens == 0) { diff --git a/src/llama-kv-cache-iswa.cpp b/src/llama-kv-cache-iswa.cpp index 2fcf238d9..e91866469 100644 --- a/src/llama-kv-cache-iswa.cpp +++ b/src/llama-kv-cache-iswa.cpp @@ -206,7 +206,7 @@ llama_memory_context_ptr llama_kv_cache_iswa::init_batch(llama_batch_allocr & ba std::vector ubatches; while (true) { - auto ubatch = balloc.split_equal(n_ubatch, !unified); + auto ubatch = balloc.split_equal(n_ubatch, !unified, 0); if (ubatch.n_tokens == 0) { break; diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index 680de5144..e70583e64 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -706,7 +706,7 @@ llama_memory_context_ptr llama_kv_cache::init_batch( std::vector ubatches; while (true) { - auto ubatch = n_stream == 1 ? balloc.split_simple(n_ubatch) : balloc.split_equal(n_ubatch, true); + auto ubatch = n_stream == 1 ? balloc.split_simple(n_ubatch) : balloc.split_equal(n_ubatch, true, 0); if (ubatch.n_tokens == 0) { break; diff --git a/src/llama-memory-hybrid-iswa.cpp b/src/llama-memory-hybrid-iswa.cpp index c7d4bcd41..06f7fd542 100644 --- a/src/llama-memory-hybrid-iswa.cpp +++ b/src/llama-memory-hybrid-iswa.cpp @@ -77,15 +77,15 @@ llama_memory_context_ptr llama_memory_hybrid_iswa::init_batch(llama_batch_allocr // if all tokens are output, split by sequence ubatch = balloc.split_seq(n_ubatch); } else { - if (mem_recr->n_rs_seq > 0) { - // [TAG_RECURRENT_ROLLBACK_SPLITS] - // TODO: recurrent state rollback does not support equal splits - ubatch = balloc.split_seq(n_ubatch); - } else { - // Use non-sequential split when KV cache is unified (needed for hellaswag/winogrande/multiple-choice) - const bool unified = (mem_attn->get_base()->get_n_stream() == 1); - ubatch = balloc.split_equal(n_ubatch, !unified); - } + // Use non-sequential split when KV cache is unified (needed for hellaswag/winogrande/multiple-choice) + const bool unified = (mem_attn->get_base()->get_n_stream() == 1); + + // [TAG_RECURRENT_ROLLBACK_SPLITS] + // the trailing (1 + n_rs_seq) tokens of each seq must stay in the same ubatch + // so that the rollback snapshots remain valid + const uint32_t n_rs_seq = mem_recr->n_rs_seq; + + ubatch = balloc.split_equal(n_ubatch, !unified, n_rs_seq > 0 ? n_rs_seq + 1 : 0); } if (ubatch.n_tokens == 0) { diff --git a/src/llama-memory-hybrid.cpp b/src/llama-memory-hybrid.cpp index f2d49cbce..42c7381a9 100644 --- a/src/llama-memory-hybrid.cpp +++ b/src/llama-memory-hybrid.cpp @@ -78,15 +78,15 @@ llama_memory_context_ptr llama_memory_hybrid::init_batch(llama_batch_allocr & ba // if all tokens are output, split by sequence ubatch = balloc.split_seq(n_ubatch); } else { - if (mem_recr->n_rs_seq > 0) { - // [TAG_RECURRENT_ROLLBACK_SPLITS] - // TODO: recurrent state rollback does not support equal splits - ubatch = balloc.split_seq(n_ubatch); - } else { - // Use non-sequential split when KV cache is unified (needed for hellaswag/winogrande/multiple-choice) - const bool unified = (mem_attn->get_n_stream() == 1); - ubatch = balloc.split_equal(n_ubatch, !unified); - } + // Use non-sequential split when KV cache is unified (needed for hellaswag/winogrande/multiple-choice) + const bool unified = (mem_attn->get_n_stream() == 1); + + // [TAG_RECURRENT_ROLLBACK_SPLITS] + // the trailing (1 + n_rs_seq) tokens of each seq must stay in the same ubatch + // so that the rollback snapshots remain valid + const uint32_t n_rs_seq = mem_recr->n_rs_seq; + + ubatch = balloc.split_equal(n_ubatch, !unified, n_rs_seq > 0 ? n_rs_seq + 1 : 0); } if (ubatch.n_tokens == 0) { diff --git a/src/llama-memory-recurrent.cpp b/src/llama-memory-recurrent.cpp index 6a4892fb4..3d6c6db87 100644 --- a/src/llama-memory-recurrent.cpp +++ b/src/llama-memory-recurrent.cpp @@ -416,15 +416,12 @@ llama_memory_context_ptr llama_memory_recurrent::init_batch(llama_batch_allocr & // if all tokens are output, split by sequence ubatch = balloc.split_seq(n_ubatch); } else { - if (n_rs_seq > 0) { - // [TAG_RECURRENT_ROLLBACK_SPLITS] - // TODO: recurrent state rollback does not support equal splits - ubatch = balloc.split_seq(n_ubatch); - } else { - // TODO: non-sequential equal split can be done if using unified KV cache - // for simplicity, we always use sequential equal split for now - ubatch = balloc.split_equal(n_ubatch, true); - } + // TODO: non-sequential equal split can be done if using unified KV cache + // for simplicity, we always use sequential equal split for now + // [TAG_RECURRENT_ROLLBACK_SPLITS] + // the trailing (1 + n_rs_seq) tokens of each seq must stay in the same ubatch + // so that the rollback snapshots remain valid + ubatch = balloc.split_equal(n_ubatch, true, n_rs_seq > 0 ? n_rs_seq + 1 : 0); } if (ubatch.n_tokens == 0) { diff --git a/src/models/delta-net-base.cpp b/src/models/delta-net-base.cpp index ad9ce7714..cf5e38095 100644 --- a/src/models/delta-net-base.cpp +++ b/src/models/delta-net-base.cpp @@ -496,8 +496,8 @@ ggml_tensor * llm_build_delta_net_base::build_conv_state( ggml_build_forward_expand(gf, ggml_cpy(ctx0, conv_state_last, conv_state_update)); } else { // [TAG_RECURRENT_ROLLBACK_SPLITS] - // TODO: this logic incorrectly assumes that the last (n_rs_seq + 1) tokens of a sequence in a batch are - // inside the same ubatch. currently with `split_equal()` this is not correct + // this logic assumes that the last (n_rs_seq + 1) tokens of a sequence in a batch are inside + // the same ubatch, which `split_equal()` guarantees via its n_keep_tail argument const int64_t K = (int64_t) cparams.n_rs_seq + 1; From bbebeec4a87355896e3faac0c2baca8130c91b6a Mon Sep 17 00:00:00 2001 From: Pascal Date: Wed, 8 Jul 2026 11:02:50 +0200 Subject: [PATCH 17/17] server-stream: follow-up on SSE Replay Buffer (#23226) (#25047) * server-stream : pimpl * server-stream: prefix free functions with server_stream_ address review from ggerganov: scope the public stream functions under the server_stream_ prefix, matching server_stream_session_manager_start/stop. * server-stream: guard session and manager state with the mutex address review from ggerganov: make done, completed_ts and the GC running flag plain members under their mutex and set the condvar predicates under the lock. keep cancelled atomic for the lock-free should_stop poll. * server-stream: trim comments to the non-obvious address review from ggerganov: drop comments that restate the code, keep the concurrency, lifetime and ordering rationale. de-stale a few comments left by the pimpl: g_stream_sessions is now internal and the /v1/streams listing is gone. * server-stream: update dev docs for the pimpl and prefix reflect server_stream_session_manager_start/stop and the server_stream_ prefix, note the manager is now a file-static singleton hidden in the .cpp * server-stream: move stream traces to debug level keep the bring-up traces for diagnostics but off the default log: skip drain, draining, drain ended, DELETE evict, attach_pipe, and the router stream resume proxy. * server-stream: align router stream resume proxy trace with upstream the child-side bring-up traces are already SRV_TRC on master, move the router stream resume proxy trace to the same level. * server-stream: move stream_read_status enum to the cpp it is only used by the hidden session and consumer types, so it belongs with them behind the pimpl boundary, not on the public header surface. --------- Co-authored-by: Georgi Gerganov --- tools/server/README-dev.md | 10 +- tools/server/server-context.cpp | 4 +- tools/server/server-models.cpp | 4 +- tools/server/server-stream.cpp | 188 +++++++++++++++++++++++++------- tools/server/server-stream.h | 151 +++---------------------- tools/server/server.cpp | 16 +-- 6 files changed, 180 insertions(+), 193 deletions(-) diff --git a/tools/server/README-dev.md b/tools/server/README-dev.md index dfc9004de..882adca09 100644 --- a/tools/server/README-dev.md +++ b/tools/server/README-dev.md @@ -57,7 +57,7 @@ The core architecture consists of the following components: - `server_tokens`: Unified representation of token sequences (supports both text and multimodal tokens); used by `server_task` and `server_slot`. - `server_prompt_checkpoint`: For recurrent (e.g., RWKV) and SWA models, stores snapshots of KV cache state. Enables reuse when subsequent requests share the same prompt prefix, saving redundant computation. - `server_models`: Standalone component for managing multiple backend instances (used in router mode). It is completely independent of `server_context`. -- `stream_session_manager`: Process wide owner of resumable SSE stream sessions (`g_stream_sessions`), keyed by conversation id. Backs the replay buffer that lets a client reattach to a generation after an HTTP disconnect. See the "Resumable streaming" section below. +- `stream_session_manager`: process wide owner of resumable SSE stream sessions, keyed by conversation id. A file-static singleton inside `server-stream.cpp`, driven through `server_stream_session_manager_start/stop`. Backs the replay buffer that lets a client reattach to a generation after an HTTP disconnect. See the "Resumable streaming" section below. ```mermaid graph TD @@ -127,10 +127,12 @@ It is opt in via the `X-Conversation-Id` header on `POST /v1/chat/completions`. The feature lives entirely in `server-stream.{h,cpp}` and rests on three types: - `stream_session`: a bounded ring buffer (4 MiB cap, oldest bytes drop first) plus a condvar. `append` pushes raw SSE bytes, `read_from` drains from any offset and blocks for live bytes or finalize, `finalize` wakes readers, `cancel` stops the producer. One conv maps to at most one live session. -- `stream_session_manager` (`g_stream_sessions`): owns all sessions keyed by conv id, enforces the one conv one session invariant via `create_or_replace`, and runs a GC thread that drops completed sessions past their TTL. +- `stream_session_manager`: a file-static singleton (`g_stream_sessions`) inside `server-stream.cpp`, owns all sessions keyed by conv id, enforces the one conv one session invariant via `create_or_replace`, and runs a GC thread that drops completed sessions past their TTL. Exposed to main only through `server_stream_session_manager_start/stop`. - `stream_pipe_producer` / `stream_pipe_consumer`: the write and read ends. The producer owns the session lifetime and finalizes it on destruction; the consumer is read only and never finalizes, so a reader detaching cannot kill a running generation. -Producer side: `server_res_generator` attaches a producer pipe when the header is present. The HTTP content provider mirrors every chunk into the ring before writing it to the socket. While a pipe is attached, `stream_aware_should_stop` ignores peer disconnect, so a dropped socket does not stop generation: only an explicit `DELETE` does. When the peer leaves early, `on_complete` calls `close()`, which drains the rest of the generation into the ring on the http worker. +The implementation is hidden in `server-stream.cpp` (pimpl). The header exposes only the route handler factories, `server_stream_session_attach_pipe`, `server_stream_aware_should_stop`, `server_stream_conv_id_from_headers` and the GC lifecycle; the session, manager and consumer types stay in the `.cpp`. + +Producer side: `server_res_generator` attaches a producer pipe when the header is present. The HTTP content provider mirrors every chunk into the ring before writing it to the socket. While a pipe is attached, `server_stream_aware_should_stop` ignores peer disconnect, so a dropped socket does not stop generation: only an explicit `DELETE` does. When the peer leaves early, `on_complete` calls `close()`, which drains the rest of the generation into the ring on the http worker. Lifetime safety: the producer pipe holds a shared `alive` flag also captured by the session cancel hook. `~server_res_generator` calls `cleanup()` to clear that hook while the reader is still alive, so a `cancel` arriving during teardown can never call `stop()` on a freed response. This ordering is the most fragile part of the feature: finalizing or destroying the producer before `cleanup()` runs reintroduces a use after free. @@ -144,7 +146,7 @@ Routes: Router mode binds the same paths to proxy handlers. A `conv_id -> child` map (`conv_models`), populated when a POST is routed, resolves the owning child in one lookup with no polling. The lookup groups ids per child; GET and DELETE proxy straight to the owner. This loopback REST hop is expected to move to a websocket IPC later, swapping only the transport. -Lifecycle: `g_stream_sessions.start_gc()` runs in main after common init, `stop_gc()` runs first in `clean_up()` and finalizes every live session so no reader hangs. Reader blocking and the post drop drain both run on httplib worker threads, which block on a condvar rather than spin. +Lifecycle: `server_stream_session_manager_start()` runs in main after common init, `server_stream_session_manager_stop()` runs first in `clean_up()` and finalizes every live session so no reader hangs. Reader blocking and the post drop drain both run on httplib worker threads, which block on a condvar rather than spin. | Constant | Value | Role | | --- | --- | --- | diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index df5666581..aa5d0a2ab 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -4188,7 +4188,7 @@ std::unique_ptr server_routes::handle_completions_impl( } }; - auto effective_should_stop = stream_aware_should_stop(res_this, req.should_stop); + auto effective_should_stop = server_stream_aware_should_stop(res_this, req.should_stop); try { if (effective_should_stop()) { @@ -4284,7 +4284,7 @@ std::unique_ptr server_routes::handle_completions_impl( // attach a producer pipe to the response when X-Conversation-Id is present. // the pipe mirrors SSE chunks into the ring buffer and wires up the cancel hook. - stream_session_attach_pipe(*res, req.headers); + server_stream_session_attach_pipe(*res, req.headers); return res; } diff --git a/tools/server/server-models.cpp b/tools/server/server-models.cpp index 6a8eb2a2b..0cbf520af 100644 --- a/tools/server/server-models.cpp +++ b/tools/server/server-models.cpp @@ -1681,7 +1681,7 @@ void server_models_routes::init_routes() { } // remember which child serves this conversation so the stream routes can route straight // to it without polling, keyed on the exact conv id from the header - std::string conv_id = stream_conv_id_from_headers(req.headers); + std::string conv_id = server_stream_conv_id_from_headers(req.headers); if (!conv_id.empty()) { models.conv_models.remember(conv_id, name); } @@ -1896,7 +1896,7 @@ void server_models_routes::init_routes() { if (!from.empty()) { child_path += "?from=" + from; } - SRV_INF("proxying stream resume to model %s on port %d, path=%s\n", + SRV_TRC("proxying stream resume to model %s on port %d, path=%s\n", owner->name.c_str(), owner->port, child_path.c_str()); auto proxy = std::make_unique( "GET", diff --git a/tools/server/server-stream.cpp b/tools/server/server-stream.cpp index c2bba8ec4..553ac26b1 100644 --- a/tools/server/server-stream.cpp +++ b/tools/server/server-stream.cpp @@ -6,6 +6,12 @@ #include #include #include +#include + +enum class stream_read_status { + OK, + OFFSET_LOST, +}; namespace { constexpr int64_t STREAM_SESSION_TTL_SECONDS = 300; @@ -13,7 +19,6 @@ constexpr size_t STREAM_SESSION_MAX_BYTES = 4 * 1024 * 1024; constexpr int64_t STREAM_SESSION_GC_INTERVAL_SECONDS = 60; constexpr int64_t STREAM_READ_WAKE_INTERVAL_MS = 200; -// returns unix time in seconds int64_t now_seconds() { return std::chrono::duration_cast( std::chrono::system_clock::now().time_since_epoch() @@ -21,6 +26,91 @@ int64_t now_seconds() { } } +// owns all live sessions keyed by conversation_id, one conv = at most one live session. +// a periodic GC evicts expired ones +class stream_session_manager { +public: + stream_session_manager(); + ~stream_session_manager(); + + stream_session_manager(const stream_session_manager &) = delete; + stream_session_manager & operator=(const stream_session_manager &) = delete; + + // install a new session, evicting and cancelling any previous one. conversation_id must be non empty + stream_session_ptr create_or_replace(const std::string & conversation_id); + + stream_session_ptr get(const std::string & conversation_id); + + std::vector list_all() const; + + void evict(const std::string & conversation_id); + + void evict_and_cancel(const std::string & conversation_id); + + void start_gc(); + void stop_gc(); + +private: + void gc_loop(); + + mutable std::shared_mutex map_mu; + std::unordered_map sessions; // key: conversation_id + std::thread gc_thread; + bool running; + std::mutex gc_wake_mu; + std::condition_variable gc_wake_cv; +}; + +// process wide manager, lifecycle controlled by llama-server main() via start_gc/stop_gc +static stream_session_manager g_stream_sessions; + +void server_stream_session_manager_start() { + g_stream_sessions.start_gc(); +} + +void server_stream_session_manager_stop() { + g_stream_sessions.stop_gc(); +} + +struct stream_session { + std::string conversation_id; + int64_t started_ts; // unix seconds at construction + + stream_session(std::string conversation_id_, size_t max_bytes_); + stream_session(const stream_session &) = delete; + stream_session & operator=(const stream_session &) = delete; + + bool append(const char * data, size_t len); + + void finalize(); + + // drain from offset into sink, blocking for more bytes or finalize. OFFSET_LOST if offset + // fell below the dropped prefix + stream_read_status read_from(size_t offset, + const std::function & sink, + const std::function & should_stop); + + bool is_done() const; + bool is_cancelled() const; + size_t total_size() const; // bytes that ever entered the session + size_t dropped_prefix() const; // bytes evicted from the front due to cap + int64_t completed_at() const; // 0 while alive, unix seconds after finalize + + void set_stop_producer(std::function fn); + + void cancel(); + +private: + mutable std::mutex mu; + std::condition_variable cv; + std::vector buffer; + size_t prefix_dropped; + size_t cap_bytes; + bool done; + std::atomic cancelled; // polled lock-free by the should_stop closure, no mu + int64_t completed_ts; + std::function stop_producer; +}; stream_session::stream_session(std::string conversation_id_, size_t max_bytes_) : conversation_id(std::move(conversation_id_)) , started_ts(now_seconds()) @@ -38,7 +128,7 @@ bool stream_session::append(const char * data, size_t len) { } { std::lock_guard lock(mu); - if (done.load(std::memory_order_relaxed)) { + if (done) { return false; } if (len >= cap_bytes) { @@ -62,11 +152,14 @@ bool stream_session::append(const char * data, size_t len) { } void stream_session::finalize() { - bool was_done = done.exchange(true, std::memory_order_acq_rel); - if (was_done) { - return; + { + std::lock_guard lock(mu); + if (done) { + return; + } + done = true; + completed_ts = now_seconds(); } - completed_ts.store(now_seconds(), std::memory_order_release); cv.notify_all(); } @@ -96,7 +189,7 @@ stream_read_status stream_session::read_from(size_t offset, lock.lock(); continue; } - if (done.load(std::memory_order_acquire)) { + if (done) { return stream_read_status::OK; } // wait for new bytes, finalize, or a periodic wake to re check should_stop @@ -105,7 +198,8 @@ stream_read_status stream_session::read_from(size_t offset, } bool stream_session::is_done() const { - return done.load(std::memory_order_acquire); + std::lock_guard lock(mu); + return done; } size_t stream_session::total_size() const { @@ -119,7 +213,8 @@ size_t stream_session::dropped_prefix() const { } int64_t stream_session::completed_at() const { - return completed_ts.load(std::memory_order_acquire); + std::lock_guard lock(mu); + return completed_ts; } void stream_session::set_stop_producer(std::function fn) { @@ -128,7 +223,7 @@ void stream_session::set_stop_producer(std::function fn) { } void stream_session::cancel() { - // flip cancelled first so the producer-side stream_aware_should_stop can break out of the + // flip cancelled first so the producer-side server_stream_aware_should_stop can break out of the // recv() wait even if remove_waiting_task_ids does not notify the condvar (the cancel task // posted by rd.stop() will eventually notify, but we do not want to depend on that timing) cancelled.store(true, std::memory_order_release); @@ -237,18 +332,24 @@ void stream_session_manager::evict_and_cancel(const std::string & conversation_i } void stream_session_manager::start_gc() { - if (running.exchange(true)) { - return; + { + std::lock_guard lock(gc_wake_mu); + if (running) { + return; + } + running = true; } gc_thread = std::thread([this] { gc_loop(); }); } void stream_session_manager::stop_gc() { - bool was_running = running.exchange(false); + bool was_running; + { + std::lock_guard lock(gc_wake_mu); + was_running = running; + running = false; + } if (was_running) { - { - std::lock_guard lock(gc_wake_mu); - } gc_wake_cv.notify_all(); if (gc_thread.joinable()) { gc_thread.join(); @@ -270,15 +371,15 @@ void stream_session_manager::stop_gc() { } void stream_session_manager::gc_loop() { - while (running.load(std::memory_order_acquire)) { + while (true) { { std::unique_lock lock(gc_wake_mu); gc_wake_cv.wait_for(lock, std::chrono::seconds(STREAM_SESSION_GC_INTERVAL_SECONDS), - [this] { return !running.load(std::memory_order_acquire); }); - } - if (!running.load(std::memory_order_acquire)) { - return; + [this] { return !running; }); + if (!running) { + return; + } } int64_t cutoff = now_seconds() - STREAM_SESSION_TTL_SECONDS; std::vector to_drop; @@ -301,10 +402,19 @@ void stream_session_manager::gc_loop() { } } -// process wide manager, lifecycle controlled by llama-server main() via start_gc/stop_gc -stream_session_manager g_stream_sessions; +// stream_pipe -// stream_pipe --------------------------------------------------------------------------------- +// consumer end: read-only replay of the ring buffer, the destructor does not finalize the session +struct stream_pipe_consumer : stream_pipe { + stream_read_status read(size_t & offset, + const std::function & sink, + const std::function & should_stop); + + static std::shared_ptr create(stream_session_ptr session); + +private: + explicit stream_pipe_consumer(stream_session_ptr session); +}; stream_pipe::stream_pipe(stream_session_ptr session) : session_(std::move(session)) { @@ -408,12 +518,10 @@ static server_http_res_ptr make_error_response(int status, const std::string & m return res; } -server_http_context::handler_t make_stream_get_handler() { +server_http_context::handler_t server_stream_make_get_handler() { return [](const server_http_req & req) -> server_http_res_ptr { - // GET /v1/stream/?from=N replays the SSE bytes already buffered for the - // session, blocks for more bytes when the session is still running, returns when - // the session is finalized. the body is streamed back as text/event-stream so the - // browser EventSource can attach to it like a fresh request + // GET /v1/stream/?from=N replays buffered SSE bytes then blocks for live + // bytes until the session finalizes, streamed as text/event-stream for EventSource std::string conv_id = req.get_param("conv_id"); if (conv_id.empty()) { return make_error_response(400, "Missing conversation id in path", ERROR_TYPE_INVALID_REQUEST); @@ -459,11 +567,10 @@ server_http_context::handler_t make_stream_get_handler() { }; } -server_http_context::handler_t make_streams_lookup_handler() { +server_http_context::handler_t server_stream_make_lookup_handler() { return [](const server_http_req & req) -> server_http_res_ptr { - // POST /v1/streams/lookup with body {"conversation_ids": ["X", "Y", ...]} returns the - // matching sessions, only for ids the caller already knows. each id matches the exact key - // and any "::" variant, so one lookup covers every per model session for a conv + // POST /v1/streams/lookup returns the matching sessions, only for ids the caller already + // knows. each id matches the exact key and any "::" per model variant std::vector requested; try { json body = json::parse(req.body); @@ -518,11 +625,10 @@ server_http_context::handler_t make_streams_lookup_handler() { }; } -server_http_context::handler_t make_stream_delete_handler() { +server_http_context::handler_t server_stream_make_delete_handler() { return [](const server_http_req & req) -> server_http_res_ptr { - // DELETE /v1/stream/ is the explicit user Stop, cancels the producer hook - // wired by handle_completions_impl and evicts the buffer. idempotent, a session that - // already finalized or was never created returns 204 either way + // DELETE /v1/stream/ is the explicit user Stop, cancels the producer and evicts + // the buffer. idempotent, returns 204 even if the session was already gone std::string conv_id = req.get_param("conv_id"); if (conv_id.empty()) { return make_error_response(400, "Missing conversation id in path", ERROR_TYPE_INVALID_REQUEST); @@ -536,7 +642,7 @@ server_http_context::handler_t make_stream_delete_handler() { }; } -std::string stream_conv_id_from_headers(const std::map & headers) { +std::string server_stream_conv_id_from_headers(const std::map & headers) { // case-insensitive scan for x-conversation-id static constexpr char target[] = "x-conversation-id"; static constexpr size_t target_len = sizeof(target) - 1; @@ -555,8 +661,8 @@ std::string stream_conv_id_from_headers(const std::map return std::string(); } -void stream_session_attach_pipe(server_http_res & res, const std::map & headers) { - std::string conversation_id = stream_conv_id_from_headers(headers); +void server_stream_session_attach_pipe(server_http_res & res, const std::map & headers) { + std::string conversation_id = server_stream_conv_id_from_headers(headers); SRV_TRC("conv_id=%s (empty=%d)\n", conversation_id.c_str(), conversation_id.empty() ? 1 : 0); if (conversation_id.empty()) { return; @@ -565,7 +671,7 @@ void stream_session_attach_pipe(server_http_res & res, const std::map stream_aware_should_stop(server_http_res * res, std::function fallback) { +std::function server_stream_aware_should_stop(server_http_res * res, std::function fallback) { return [res, fallback = std::move(fallback)]() -> bool { if (res->spipe) { return res->spipe->is_cancelled(); diff --git a/tools/server/server-stream.h b/tools/server/server-stream.h index ff363bb4c..c0c3e924f 100644 --- a/tools/server/server-stream.h +++ b/tools/server/server-stream.h @@ -3,81 +3,23 @@ #include "server-http.h" #include -#include #include -#include #include #include -#include -#include #include -#include -#include -#include -enum class stream_read_status { - OK, - OFFSET_LOST, -}; +// streaming buffer for one generation, survives HTTP disconnect. the producer appends SSE bytes, +// readers drain from any offset via read_from. keyed by conversation_id, one conv = one live session -// streaming buffer for one generation, survives HTTP disconnect. the producer appends raw SSE -// bytes, readers drain from any offset via read_from and block until more bytes or finalize. -// keyed by conversation_id: one conv = at most one live session -struct stream_session { - std::string conversation_id; - int64_t started_ts; // unix seconds at construction, used by /v1/streams listing - - stream_session(std::string conversation_id_, size_t max_bytes_); - stream_session(const stream_session &) = delete; - stream_session & operator=(const stream_session &) = delete; - - // append raw bytes, drops from the front if the cap is reached. - // returns false if the session is already finalized - bool append(const char * data, size_t len); - - // mark the session as complete, wakes all pending readers - void finalize(); - - // drain bytes from offset, calling sink for each chunk. blocks until more - // bytes arrive or finalize is called. returns OK on clean exit, OFFSET_LOST - // if offset falls below the dropped prefix - stream_read_status read_from(size_t offset, - const std::function & sink, - const std::function & should_stop); - - bool is_done() const; - bool is_cancelled() const; - size_t total_size() const; // bytes that ever entered the session - size_t dropped_prefix() const; // bytes evicted from the front due to cap - int64_t completed_at() const; // 0 while alive, unix seconds after finalize - - // attach the producer stop hook used to cancel its reader, pass an empty function to detach - void set_stop_producer(std::function fn); - - // signal the producer to abort its inference asap via the stop hook, idempotent - void cancel(); - -private: - mutable std::mutex mu; - std::condition_variable cv; - std::vector buffer; - size_t prefix_dropped; - size_t cap_bytes; - std::atomic done; - std::atomic cancelled; - std::atomic completed_ts; - std::function stop_producer; // protected by mu -}; +struct stream_session; using stream_session_ptr = std::shared_ptr; -// one end of a stream_session pipe. the base holds the session and the shared query, the -// producer and consumer ends derive from it. virtual dtor so each end runs its own teardown: +// base of the producer/consumer pipe ends. virtual dtor so each runs its own teardown: // the producer finalizes the session, the consumer leaves it untouched struct stream_pipe { virtual ~stream_pipe() = default; - // true if the session was cancelled (e.g. via DELETE /v1/stream/) bool is_cancelled() const; protected: @@ -95,7 +37,6 @@ protected: struct stream_pipe_producer : stream_pipe { ~stream_pipe_producer() override; - // append raw bytes to the session's ring buffer, returns false if already finalized bool write(const char * data, size_t len); // mark the natural end on the wire so a later close() is a no-op @@ -121,83 +62,21 @@ private: server_http_res * res_ = nullptr; }; -// consumer end: read-only replay of the ring buffer, the destructor does not finalize the session -struct stream_pipe_consumer : stream_pipe { - // drain bytes from offset, calling sink for each available chunk. blocks until more data - // arrives or the session finalizes. should_stop is polled, returns OFFSET_LOST if offset - // fell below the dropped prefix - stream_read_status read(size_t & offset, - const std::function & sink, - const std::function & should_stop); +void server_stream_session_manager_start(); +void server_stream_session_manager_stop(); - static std::shared_ptr create(stream_session_ptr session); +// route handler factories wired under /v1/stream/* by server.cpp +server_http_context::handler_t server_stream_make_get_handler(); +server_http_context::handler_t server_stream_make_lookup_handler(); +server_http_context::handler_t server_stream_make_delete_handler(); -private: - explicit stream_pipe_consumer(stream_session_ptr session); -}; +// extract the X-Conversation-Id header value (case-insensitive), empty when absent +std::string server_stream_conv_id_from_headers(const std::map & headers); -// owns all live sessions, runs a periodic GC to evict expired ones. -// the map is keyed by conversation_id, so the invariant "one conv = at most one -// live session" is enforced at the type level -class stream_session_manager { -public: - stream_session_manager(); - ~stream_session_manager(); - - stream_session_manager(const stream_session_manager &) = delete; - stream_session_manager & operator=(const stream_session_manager &) = delete; - - // install a new session for this conversation, evicting and cancelling any previous one. - // the conversation_id must be non empty, the caller is responsible for that check. - // returns the new session - stream_session_ptr create_or_replace(const std::string & conversation_id); - - // lookup, returns null if unknown or already evicted - stream_session_ptr get(const std::string & conversation_id); - - // list every live or recently completed session, used by GET /v1/streams without filter - std::vector list_all() const; - - // remove from the map and finalize, wakes any pending readers - void evict(const std::string & conversation_id); - - // signal the producer to cancel asap then evict, used by the explicit user Stop path - void evict_and_cancel(const std::string & conversation_id); - - void start_gc(); - void stop_gc(); - -private: - void gc_loop(); - - mutable std::shared_mutex map_mu; - std::unordered_map sessions; // key: conversation_id - std::thread gc_thread; - std::atomic running; - std::mutex gc_wake_mu; - std::condition_variable gc_wake_cv; -}; - -// process wide manager, linked by both llama-server and llama-cli. llama-server main() drives -// start_gc/stop_gc, llama-cli leaves it idle. the dtor calls stop_gc() unconditionally so exit -// is safe whether or not the GC thread ran -extern stream_session_manager g_stream_sessions; - -// route handler factories operating on g_stream_sessions, wired under /v1/stream/* by server.cpp. -// keeps the resumable stream surface confined to server-stream -server_http_context::handler_t make_stream_get_handler(); -server_http_context::handler_t make_streams_lookup_handler(); -server_http_context::handler_t make_stream_delete_handler(); - -// extract the X-Conversation-Id header value (case-insensitive), empty when absent. exposed so -// the router can track which child serves a forwarded POST -std::string stream_conv_id_from_headers(const std::map & headers); - -// on an X-Conversation-Id header, create or replace the session and attach a producer pipe to -// res. no-op when absent, called from the server_res_generator constructor -void stream_session_attach_pipe(server_http_res & res, const std::map & headers); +// on an X-Conversation-Id header, create or replace the session and attach a producer pipe to res +void server_stream_session_attach_pipe(server_http_res & res, const std::map & headers); // should_stop closure that ignores peer disconnect when a pipe is attached, so only an explicit // DELETE stops the producer and generation keeps flowing into the ring buffer. without a pipe it // delegates to fallback, the legacy non-resumable flow -std::function stream_aware_should_stop(server_http_res * res, std::function fallback); +std::function server_stream_aware_should_stop(server_http_res * res, std::function fallback); diff --git a/tools/server/server.cpp b/tools/server/server.cpp index eafef86ba..9e8603be6 100644 --- a/tools/server/server.cpp +++ b/tools/server/server.cpp @@ -85,7 +85,7 @@ int llama_server(int argc, char ** argv) { // start the stream session manager GC right after common init, before any HTTP route can // touch it. lifecycle is symmetric, stop_gc() runs in clean_up() before backend free - g_stream_sessions.start_gc(); + server_stream_session_manager_start(); if (!common_params_parse(argc, argv, params, LLAMA_EXAMPLE_SERVER)) { return 1; @@ -245,8 +245,8 @@ int llama_server(int argc, char ** argv) { ctx_http.post("/slots/:id_slot", ex_wrapper(routes.post_slots)); // resumable streaming, the conversation_id is the session identity end to end. router and - // child wire different handlers under the same paths: a child binds the local g_stream_sessions - // backed factories, the router binds proxies that resolve the owning child through the + // child wire different handlers under the same paths: a child binds the local session + // factories, the router binds proxies that resolve the owning child through the // conv_id -> model map server_http_context::handler_t stream_get_h; server_http_context::handler_t streams_lookup_h; @@ -256,9 +256,9 @@ int llama_server(int argc, char ** argv) { streams_lookup_h = models_routes->router_streams_lookup; stream_delete_h = models_routes->router_stream_delete; } else { - stream_get_h = make_stream_get_handler(); - streams_lookup_h = make_streams_lookup_handler(); - stream_delete_h = make_stream_delete_handler(); + stream_get_h = server_stream_make_get_handler(); + streams_lookup_h = server_stream_make_lookup_handler(); + stream_delete_h = server_stream_make_delete_handler(); } ctx_http.get ("/v1/stream/:conv_id", ex_wrapper(stream_get_h)); // POST /v1/streams/lookup with body {"conversation_ids": [...]}. you can only ask for ids @@ -343,7 +343,7 @@ int llama_server(int argc, char ** argv) { clean_up = [&models_routes]() { SRV_INF("%s: cleaning up before exit...\n", __func__); // stop the session GC first, it finalizes live sessions and wakes pending readers - g_stream_sessions.stop_gc(); + server_stream_session_manager_stop(); if (models_routes.has_value()) { models_routes->stopping.store(true); // maybe redundant, but just to be safe models_routes->models.unload_all(); @@ -371,7 +371,7 @@ int llama_server(int argc, char ** argv) { clean_up = [&ctx_http, &ctx_server]() { SRV_INF("%s: cleaning up before exit...\n", __func__); // stop the session GC first, it finalizes live sessions and wakes pending readers - g_stream_sessions.stop_gc(); + server_stream_session_manager_stop(); ctx_http.stop(); ctx_server.terminate(); llama_backend_free();