llama-batch: add n_keep_tail in split_equal for recurrent models (#25278)

* 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<typeof mcpStore.getToolDefinitionsForLLM>, claiming the
  shape was owned by mcpStore. Switch to ReturnType<typeof
  toolsStore.getEnabledToolsForLLM>, 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 <admin@serveurperso.com>

common/arg.cpp

@@ -27,6 +27,7 @@
 #include <cinttypes>
 #include <climits>
 #include <cstdarg>
+#include <filesystem>
 #include <fstream>
 #include <list>
 #include <regex>
@@ -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(

common/common.h

@@ -14,6 +14,7 @@
 #include <vector>
 #include <map>
 #include <algorithm>
+#include <fstream>
 
 #if defined(_WIN32) && !defined(_WIN32_WINNT)
 #define _WIN32_WINNT 0x0A00

common/speculative.cpp

@@ -2221,6 +2221,112 @@ 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_speculative_init_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_speculative_init_result::common_speculative_init_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_speculative_init_result::~common_speculative_init_result() = default;
+
+llama_model * common_speculative_init_result::model() {
+    return pimpl->model.get();
+}
+
+llama_context * common_speculative_init_result::context() {
+    return pimpl->context.get();
+}
+
+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<common_speculative_init_result>(params, model_tgt, ctx_tgt);
+}
+
 // initialization of the speculative decoding system
 //
 common_speculative * common_speculative_init(common_params_speculative & params, uint32_t n_seq) {

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, common_speculative_deleter> common_speculative_ptr;
+
+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();
+
+private:
+    struct impl;
+    std::unique_ptr<impl> pimpl;
+};
+
+using common_speculative_init_result_ptr = std::unique_ptr<common_speculative_init_result>;
+
+common_speculative_init_result_ptr common_speculative_init_from_params(common_params & params, llama_model * model_tgt, llama_context * ctx_tgt);

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")

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:

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

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

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;

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,

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:
@@ -5019,8 +5025,8 @@ void ggml_compute_forward_get_rows(
     //}
 }
 
-template<typename idx_t>
-static void ggml_compute_forward_set_rows_f32(
+template<typename src_t, typename idx_t>
+static void ggml_compute_forward_set_rows_impl(
         const ggml_compute_params * params,
               ggml_tensor * dst) {
 
@@ -5035,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);
 
@@ -5049,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) {
@@ -5062,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<src_t, float>) {
+                    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<src_t, ggml_fp16_t>) {
+                    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));
+                }
             }
         }
     }
@@ -5081,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<int64_t>(params, dst);
+                    ggml_compute_forward_set_rows_impl<float, int64_t>(params, dst);
                 } else if (src1->type == GGML_TYPE_I32) {
-                    ggml_compute_forward_set_rows_f32<int32_t>(params, dst);
+                    ggml_compute_forward_set_rows_impl<float, int32_t>(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<ggml_fp16_t, int64_t>(params, dst);
+                    } else if (src1->type == GGML_TYPE_I32) {
+                        ggml_compute_forward_set_rows_impl<ggml_fp16_t, int32_t>(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));
@@ -5680,6 +5711,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:

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;

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);

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;
             }

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;
 };
 

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;
                 }

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;
     }
 

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);

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);

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;
                 }
 

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);
 

ggml/src/ggml-metal/ggml-metal.metal

@@ -42,6 +42,8 @@ typedef matrix<bfloat, 4, 4> bfloat4x4;
 typedef matrix<bfloat, 2, 4> 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<typename TI, typename block_q, void (*quantize_func)(device const float *, device block_q &)>
+typedef decltype(kernel_get_rows_f<float, float>) get_rows_f_t;
+
+template [[host_name("kernel_get_rows_f32")]]  kernel get_rows_f_t kernel_get_rows_f<float, float>;
+template [[host_name("kernel_get_rows_f16")]]  kernel get_rows_f_t kernel_get_rows_f<half,  float>;
+template [[host_name("kernel_get_rows_i32")]]  kernel get_rows_f_t kernel_get_rows_f<int32_t, int32_t>;
+#if defined(GGML_METAL_HAS_BF16)
+template [[host_name("kernel_get_rows_bf16")]] kernel get_rows_f_t kernel_get_rows_f<bfloat, float>;
+#endif
+
+typedef decltype(kernel_get_rows_q<block_q4_0, 2, dequantize_q4_0>) get_rows_q_t;
+
+template [[host_name("kernel_get_rows_q1_0")]]    kernel get_rows_q_t kernel_get_rows_q<block_q1_0,    8, dequantize_q1_0>;
+template [[host_name("kernel_get_rows_q4_0")]]    kernel get_rows_q_t kernel_get_rows_q<block_q4_0,    2, dequantize_q4_0>;
+template [[host_name("kernel_get_rows_q4_1")]]    kernel get_rows_q_t kernel_get_rows_q<block_q4_1,    2, dequantize_q4_1>;
+template [[host_name("kernel_get_rows_q5_0")]]    kernel get_rows_q_t kernel_get_rows_q<block_q5_0,    2, dequantize_q5_0>;
+template [[host_name("kernel_get_rows_q5_1")]]    kernel get_rows_q_t kernel_get_rows_q<block_q5_1,    2, dequantize_q5_1>;
+template [[host_name("kernel_get_rows_q8_0")]]    kernel get_rows_q_t kernel_get_rows_q<block_q8_0,    2, dequantize_q8_0>;
+template [[host_name("kernel_get_rows_mxfp4")]]   kernel get_rows_q_t kernel_get_rows_q<block_mxfp4,   2, dequantize_mxfp4>;
+template [[host_name("kernel_get_rows_q2_K")]]    kernel get_rows_q_t kernel_get_rows_q<block_q2_K,    QK_NL, dequantize_q2_K>;
+template [[host_name("kernel_get_rows_q3_K")]]    kernel get_rows_q_t kernel_get_rows_q<block_q3_K,    QK_NL, dequantize_q3_K>;
+template [[host_name("kernel_get_rows_q4_K")]]    kernel get_rows_q_t kernel_get_rows_q<block_q4_K,    QK_NL, dequantize_q4_K>;
+template [[host_name("kernel_get_rows_q5_K")]]    kernel get_rows_q_t kernel_get_rows_q<block_q5_K,    QK_NL, dequantize_q5_K>;
+template [[host_name("kernel_get_rows_q6_K")]]    kernel get_rows_q_t kernel_get_rows_q<block_q6_K,    QK_NL, dequantize_q6_K>;
+template [[host_name("kernel_get_rows_iq2_xxs")]] kernel get_rows_q_t kernel_get_rows_q<block_iq2_xxs, QK_NL, dequantize_iq2_xxs>;
+template [[host_name("kernel_get_rows_iq2_xs")]]  kernel get_rows_q_t kernel_get_rows_q<block_iq2_xs,  QK_NL, dequantize_iq2_xs>;
+template [[host_name("kernel_get_rows_iq3_xxs")]] kernel get_rows_q_t kernel_get_rows_q<block_iq3_xxs, QK_NL, dequantize_iq3_xxs>;
+template [[host_name("kernel_get_rows_iq3_s")]]   kernel get_rows_q_t kernel_get_rows_q<block_iq3_s,   QK_NL, dequantize_iq3_s>;
+template [[host_name("kernel_get_rows_iq2_s")]]   kernel get_rows_q_t kernel_get_rows_q<block_iq2_s,   QK_NL, dequantize_iq2_s>;
+template [[host_name("kernel_get_rows_iq1_s")]]   kernel get_rows_q_t kernel_get_rows_q<block_iq1_s,   QK_NL, dequantize_iq1_s>;
+template [[host_name("kernel_get_rows_iq1_m")]]   kernel get_rows_q_t kernel_get_rows_q<block_iq1_m,   QK_NL, dequantize_iq1_m>;
+template [[host_name("kernel_get_rows_iq4_nl")]]  kernel get_rows_q_t kernel_get_rows_q<block_iq4_nl,  2,     dequantize_iq4_nl>;
+template [[host_name("kernel_get_rows_iq4_xs")]]  kernel get_rows_q_t kernel_get_rows_q<block_iq4_xs,  QK_NL, dequantize_iq4_xs>;
+
+template<typename TS, typename TI, typename block_q, void (*quantize_func)(device const float *, device block_q &)>
 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<typename T, typename TI>
+template<typename TS, typename TI, typename TD>
 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<float, int64_t, float>) set_rows_f_t;
+
+template [[host_name("kernel_set_rows_f32_i64_f32")]]   kernel set_rows_f_t kernel_set_rows_f<float, int64_t, float>;
+template [[host_name("kernel_set_rows_f32_i32_f32")]]   kernel set_rows_f_t kernel_set_rows_f<float, int32_t, float>;
+template [[host_name("kernel_set_rows_f32_i64_f16")]]   kernel set_rows_f_t kernel_set_rows_f<float, int64_t, half>;
+template [[host_name("kernel_set_rows_f32_i32_f16")]]   kernel set_rows_f_t kernel_set_rows_f<float, int32_t, half>;
+#if defined(GGML_METAL_HAS_BF16)
+template [[host_name("kernel_set_rows_f32_i64_bf16")]]  kernel set_rows_f_t kernel_set_rows_f<float, int64_t, bfloat>;
+template [[host_name("kernel_set_rows_f32_i32_bf16")]]  kernel set_rows_f_t kernel_set_rows_f<float, int32_t, bfloat>;
+#endif
+
+template [[host_name("kernel_set_rows_f16_i64_f16")]]   kernel set_rows_f_t kernel_set_rows_f<half, int64_t, half>;
+template [[host_name("kernel_set_rows_f16_i32_f16")]]   kernel set_rows_f_t kernel_set_rows_f<half, int32_t, half>;
+#if defined(GGML_METAL_HAS_BF16)
+template [[host_name("kernel_set_rows_bf16_i64_bf16")]] kernel set_rows_f_t kernel_set_rows_f<bfloat, int64_t, bfloat>;
+template [[host_name("kernel_set_rows_bf16_i32_bf16")]] kernel set_rows_f_t kernel_set_rows_f<bfloat, int32_t, bfloat>;
+#endif
+
+typedef decltype(kernel_set_rows_q32<float, int64_t, block_q8_0, quantize_q8_0>) set_rows_q32_t;
+
+template [[host_name("kernel_set_rows_f32_i64_q8_0")]]   kernel set_rows_q32_t kernel_set_rows_q32<float, int64_t, block_q8_0,   quantize_q8_0>;
+template [[host_name("kernel_set_rows_f32_i32_q8_0")]]   kernel set_rows_q32_t kernel_set_rows_q32<float, int32_t, block_q8_0,   quantize_q8_0>;
+template [[host_name("kernel_set_rows_f32_i64_q4_0")]]   kernel set_rows_q32_t kernel_set_rows_q32<float, int64_t, block_q4_0,   quantize_q4_0>;
+template [[host_name("kernel_set_rows_f32_i32_q4_0")]]   kernel set_rows_q32_t kernel_set_rows_q32<float, int32_t, block_q4_0,   quantize_q4_0>;
+template [[host_name("kernel_set_rows_f32_i64_q4_1")]]   kernel set_rows_q32_t kernel_set_rows_q32<float, int64_t, block_q4_1,   quantize_q4_1>;
+template [[host_name("kernel_set_rows_f32_i32_q4_1")]]   kernel set_rows_q32_t kernel_set_rows_q32<float, int32_t, block_q4_1,   quantize_q4_1>;
+template [[host_name("kernel_set_rows_f32_i64_q5_0")]]   kernel set_rows_q32_t kernel_set_rows_q32<float, int64_t, block_q5_0,   quantize_q5_0>;
+template [[host_name("kernel_set_rows_f32_i32_q5_0")]]   kernel set_rows_q32_t kernel_set_rows_q32<float, int32_t, block_q5_0,   quantize_q5_0>;
+template [[host_name("kernel_set_rows_f32_i64_q5_1")]]   kernel set_rows_q32_t kernel_set_rows_q32<float, int64_t, block_q5_1,   quantize_q5_1>;
+template [[host_name("kernel_set_rows_f32_i32_q5_1")]]   kernel set_rows_q32_t kernel_set_rows_q32<float, int32_t, block_q5_1,   quantize_q5_1>;
+template [[host_name("kernel_set_rows_f32_i64_iq4_nl")]] kernel set_rows_q32_t kernel_set_rows_q32<float, int64_t, block_iq4_nl, quantize_iq4_nl>;
+template [[host_name("kernel_set_rows_f32_i32_iq4_nl")]] kernel set_rows_q32_t kernel_set_rows_q32<float, int32_t, block_iq4_nl, quantize_iq4_nl>;
+
 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<float, float>) get_rows_f_t;
-
-template [[host_name("kernel_get_rows_f32")]]  kernel get_rows_f_t kernel_get_rows_f<float, float>;
-template [[host_name("kernel_get_rows_f16")]]  kernel get_rows_f_t kernel_get_rows_f<half,  float>;
-template [[host_name("kernel_get_rows_i32")]]  kernel get_rows_f_t kernel_get_rows_f<int32_t, int32_t>;
-#if defined(GGML_METAL_HAS_BF16)
-template [[host_name("kernel_get_rows_bf16")]] kernel get_rows_f_t kernel_get_rows_f<bfloat, float>;
-#endif
-
-typedef decltype(kernel_get_rows_q<block_q4_0, 2, dequantize_q4_0>) get_rows_q_t;
-
-template [[host_name("kernel_get_rows_q1_0")]]    kernel get_rows_q_t kernel_get_rows_q<block_q1_0,    8, dequantize_q1_0>;
-template [[host_name("kernel_get_rows_q4_0")]]    kernel get_rows_q_t kernel_get_rows_q<block_q4_0,    2, dequantize_q4_0>;
-template [[host_name("kernel_get_rows_q4_1")]]    kernel get_rows_q_t kernel_get_rows_q<block_q4_1,    2, dequantize_q4_1>;
-template [[host_name("kernel_get_rows_q5_0")]]    kernel get_rows_q_t kernel_get_rows_q<block_q5_0,    2, dequantize_q5_0>;
-template [[host_name("kernel_get_rows_q5_1")]]    kernel get_rows_q_t kernel_get_rows_q<block_q5_1,    2, dequantize_q5_1>;
-template [[host_name("kernel_get_rows_q8_0")]]    kernel get_rows_q_t kernel_get_rows_q<block_q8_0,    2, dequantize_q8_0>;
-template [[host_name("kernel_get_rows_mxfp4")]]   kernel get_rows_q_t kernel_get_rows_q<block_mxfp4,   2, dequantize_mxfp4>;
-template [[host_name("kernel_get_rows_q2_K")]]    kernel get_rows_q_t kernel_get_rows_q<block_q2_K,    QK_NL, dequantize_q2_K>;
-template [[host_name("kernel_get_rows_q3_K")]]    kernel get_rows_q_t kernel_get_rows_q<block_q3_K,    QK_NL, dequantize_q3_K>;
-template [[host_name("kernel_get_rows_q4_K")]]    kernel get_rows_q_t kernel_get_rows_q<block_q4_K,    QK_NL, dequantize_q4_K>;
-template [[host_name("kernel_get_rows_q5_K")]]    kernel get_rows_q_t kernel_get_rows_q<block_q5_K,    QK_NL, dequantize_q5_K>;
-template [[host_name("kernel_get_rows_q6_K")]]    kernel get_rows_q_t kernel_get_rows_q<block_q6_K,    QK_NL, dequantize_q6_K>;
-template [[host_name("kernel_get_rows_iq2_xxs")]] kernel get_rows_q_t kernel_get_rows_q<block_iq2_xxs, QK_NL, dequantize_iq2_xxs>;
-template [[host_name("kernel_get_rows_iq2_xs")]]  kernel get_rows_q_t kernel_get_rows_q<block_iq2_xs,  QK_NL, dequantize_iq2_xs>;
-template [[host_name("kernel_get_rows_iq3_xxs")]] kernel get_rows_q_t kernel_get_rows_q<block_iq3_xxs, QK_NL, dequantize_iq3_xxs>;
-template [[host_name("kernel_get_rows_iq3_s")]]   kernel get_rows_q_t kernel_get_rows_q<block_iq3_s,   QK_NL, dequantize_iq3_s>;
-template [[host_name("kernel_get_rows_iq2_s")]]   kernel get_rows_q_t kernel_get_rows_q<block_iq2_s,   QK_NL, dequantize_iq2_s>;
-template [[host_name("kernel_get_rows_iq1_s")]]   kernel get_rows_q_t kernel_get_rows_q<block_iq1_s,   QK_NL, dequantize_iq1_s>;
-template [[host_name("kernel_get_rows_iq1_m")]]   kernel get_rows_q_t kernel_get_rows_q<block_iq1_m,   QK_NL, dequantize_iq1_m>;
-template [[host_name("kernel_get_rows_iq4_nl")]]  kernel get_rows_q_t kernel_get_rows_q<block_iq4_nl,  2,     dequantize_iq4_nl>;
-template [[host_name("kernel_get_rows_iq4_xs")]]  kernel get_rows_q_t kernel_get_rows_q<block_iq4_xs,  QK_NL, dequantize_iq4_xs>;
-
-//
-// set rows
-//
-
-typedef decltype(kernel_set_rows_f<float, int64_t>) set_rows_f_t;
-
-template [[host_name("kernel_set_rows_f32_i64")]]  kernel set_rows_f_t kernel_set_rows_f<float, int64_t>;
-template [[host_name("kernel_set_rows_f32_i32")]]  kernel set_rows_f_t kernel_set_rows_f<float, int32_t>;
-template [[host_name("kernel_set_rows_f16_i64")]]  kernel set_rows_f_t kernel_set_rows_f<half, int64_t>;
-template [[host_name("kernel_set_rows_f16_i32")]]  kernel set_rows_f_t kernel_set_rows_f<half, int32_t>;
-#if defined(GGML_METAL_HAS_BF16)
-template [[host_name("kernel_set_rows_bf16_i64")]] kernel set_rows_f_t kernel_set_rows_f<bfloat, int64_t>;
-template [[host_name("kernel_set_rows_bf16_i32")]] kernel set_rows_f_t kernel_set_rows_f<bfloat, int32_t>;
-#endif
-
-typedef decltype(kernel_set_rows_q32<int64_t, block_q8_0, quantize_q8_0>) set_rows_q32_t;
-
-template [[host_name("kernel_set_rows_q8_0_i64")]]   kernel set_rows_q32_t kernel_set_rows_q32<int64_t, block_q8_0,   quantize_q8_0>;
-template [[host_name("kernel_set_rows_q8_0_i32")]]   kernel set_rows_q32_t kernel_set_rows_q32<int32_t, block_q8_0,   quantize_q8_0>;
-template [[host_name("kernel_set_rows_q4_0_i64")]]   kernel set_rows_q32_t kernel_set_rows_q32<int64_t, block_q4_0,   quantize_q4_0>;
-template [[host_name("kernel_set_rows_q4_0_i32")]]   kernel set_rows_q32_t kernel_set_rows_q32<int32_t, block_q4_0,   quantize_q4_0>;
-template [[host_name("kernel_set_rows_q4_1_i64")]]   kernel set_rows_q32_t kernel_set_rows_q32<int64_t, block_q4_1,   quantize_q4_1>;
-template [[host_name("kernel_set_rows_q4_1_i32")]]   kernel set_rows_q32_t kernel_set_rows_q32<int32_t, block_q4_1,   quantize_q4_1>;
-template [[host_name("kernel_set_rows_q5_0_i64")]]   kernel set_rows_q32_t kernel_set_rows_q32<int64_t, block_q5_0,   quantize_q5_0>;
-template [[host_name("kernel_set_rows_q5_0_i32")]]   kernel set_rows_q32_t kernel_set_rows_q32<int32_t, block_q5_0,   quantize_q5_0>;
-template [[host_name("kernel_set_rows_q5_1_i64")]]   kernel set_rows_q32_t kernel_set_rows_q32<int64_t, block_q5_1,   quantize_q5_1>;
-template [[host_name("kernel_set_rows_q5_1_i32")]]   kernel set_rows_q32_t kernel_set_rows_q32<int32_t, block_q5_1,   quantize_q5_1>;
-template [[host_name("kernel_set_rows_iq4_nl_i64")]] kernel set_rows_q32_t kernel_set_rows_q32<int64_t, block_iq4_nl, quantize_iq4_nl>;
-template [[host_name("kernel_set_rows_iq4_nl_i32")]] kernel set_rows_q32_t kernel_set_rows_q32<int32_t, block_iq4_nl, quantize_iq4_nl>;
-
 //
 // matrix-matrix multiplication
 //

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;
 }
 

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);

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);

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;
@@ -3917,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));
@@ -7419,6 +7428,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 +7477,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;
@@ -7739,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;

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),
 }
 
 

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
     };

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<idx_vec_t> 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<int32_t> 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) {

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);

src/llama-kv-cache-dsa.cpp

@@ -113,7 +113,7 @@ llama_memory_context_ptr llama_kv_cache_dsa::init_batch(
 
         std::vector<llama_ubatch> 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;

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) {

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<llama_ubatch> 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;

src/llama-kv-cache.cpp

@@ -706,7 +706,7 @@ llama_memory_context_ptr llama_kv_cache::init_batch(
 
         std::vector<llama_ubatch> 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;

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) {

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) {

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) {

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));

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;
 

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) {

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;
 

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<ggml_tensor *> 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);
@@ -2341,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<int64_t, 4> ne;
     const std::array<int, 2> nr23; // broadcast only dims 2 and 3
@@ -2349,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<int64_t, 4> ne,
             std::array<int, 2> 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]);
@@ -2396,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;
@@ -2419,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);
@@ -5848,19 +5902,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<int64_t, 2> 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<int64_t, 2> 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<int64_t, 2> 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 +5925,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 +5941,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 +5971,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<int64_t, 4> 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<int64_t, 4> 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<int64_t, 4> 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<int64_t, 4> 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 +6020,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;
 
@@ -7769,24 +7875,28 @@ static std::vector<std::unique_ptr<test_case>> 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}) {
@@ -9202,10 +9312,15 @@ static std::vector<std::unique_ptr<test_case>> 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 +9938,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<test_case> & 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)) {

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

tools/quantize/quantize.cpp

@@ -33,6 +33,7 @@ struct quant_option {
 
 static const std::vector<quant_option> 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",  },

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 <user>/<model>[:quant]` | Same as --hf-repo, but for the draft model (default: unused)<br/>(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) |

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_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;
@@ -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,36 @@ 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);
 
-            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;
+                // progress callback
+                params_dft.load_progress_callback           = load_progress_callback;
+                params_dft.load_progress_callback_user_data = &load_progress_spec;
 
-            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;
+                spec_init = common_speculative_init_from_params(params_dft, model_tgt, ctx_tgt);
+                model_dft = spec_init->model();
+                ctx_dft   = spec_init->context();
+
+                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 +1286,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 +1302,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;
 
@@ -2362,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.get(), 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);
 
@@ -2899,8 +2844,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 +2917,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 +2954,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 +2966,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 +2983,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 +3164,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++) {
@@ -3320,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.get(), 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);
 
@@ -3395,7 +3340,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

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<json> 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;
 }
 
@@ -1631,7 +1663,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 +1691,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<uint8_t> state_data_tgt;
     std::vector<uint8_t> state_data_dft;
 
@@ -1752,12 +1809,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 +1823,7 @@ void server_prompt_cache::update() {
     const size_t limit_tokens_cur = limit_size > 0 ? std::max<size_t>(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));
 

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;

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

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": {

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",

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 @@
 			</Tooltip.Content>
 		</Tooltip.Root>
 
-		<DropdownMenu.Content align="start" class="w-48">
+		<DropdownMenu.Content align="start" class="w-52">
+			<ChatFormActionAddReasoningSubmenu />
+
+			<DropdownMenu.Separator />
+
 			<DropdownMenu.Sub>
 				<DropdownMenu.SubTrigger class="flex cursor-pointer items-center gap-2">
 					<File class="h-4 w-4" />

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 @@
 </script>
 
 {#if modelSupportsThinking}
-	<DropdownMenu.Root bind:open={subOpen}>
-		<Tooltip.Root>
-			<Tooltip.Trigger>
-				<DropdownMenu.Trigger
-					class={[
-						'flex h-6 w-6 cursor-pointer items-center justify-center rounded-full p-0 transition-colors focus:outline-none focus-visible:ring-2 focus-visible:ring-ring focus-visible:ring-offset-2',
-						thinkingEnabled ? 'bg-amber-400/10 hover:bg-amber-400/20' : 'bg-muted'
-					]}
-					aria-label={`${tooltipText}. Click to configure.`}
-				>
-					{#if thinkingEnabled}
-						<Lightbulb class="h-3 w-3 text-amber-400" />
-					{:else}
-						<LightbulbOff class="h-3 w-3 text-muted-foreground" />
-					{/if}
-				</DropdownMenu.Trigger>
-			</Tooltip.Trigger>
+	<DropdownMenu.Sub bind:open={subOpen}>
+		<DropdownMenu.SubTrigger class="flex cursor-pointer items-center gap-2">
+			{#if thinkingEnabled}
+				<Lightbulb class="h-4 w-4 shrink-0 text-amber-400" />
+			{:else}
+				<LightbulbOff class="h-4 w-4 shrink-0 text-muted-foreground" />
+			{/if}
 
-			<Tooltip.Content>
-				<p class="capitalize">{tooltipText}</p>
-			</Tooltip.Content>
-		</Tooltip.Root>
+			<span class="text-sm inline-flex gap-2 {!thinkingEnabled ? 'text-muted-foreground' : ''}">
+				Reasoning
 
-		<DropdownMenu.Content
-			align="start"
-			class="w-60 rounded-xl bg-popover p-3 text-popover-foreground shadow-md outline-none"
+				<span class="capitalize text-muted-foreground">
+					{thinkingEnabled ? currentEffort : 'off'}
+				</span>
+			</span>
+		</DropdownMenu.SubTrigger>
+
+		<DropdownMenu.SubContent
+			class="w-60 bg-popover p-1.5 text-popover-foreground shadow-md outline-none"
 		>
-			<div class="mb-2 px-2.5 text-sm font-medium">Reasoning effort</div>
-
 			{#each REASONING_EFFORT_LEVELS as level (level.value)}
 				<button
 					type="button"
-					class="flex w-full cursor-pointer items-center gap-2 rounded-lg px-2.5 py-2 text-left text-sm transition-colors hover:bg-accent"
+					class="flex w-full cursor-pointer items-center gap-3 rounded-md px-2 py-1.75 text-left text-sm transition-colors hover:bg-accent"
 					class:bg-accent={isSelected(level)}
 					onclick={() => handleSelection(level)}
 				>
@@ -140,6 +123,6 @@
 					{/if}
 				</button>
 			{/each}
-		</DropdownMenu.Content>
-	</DropdownMenu.Root>
+		</DropdownMenu.SubContent>
+	</DropdownMenu.Sub>
 {/if}

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;
+	});
 </script>
 
 <div
@@ -100,7 +136,7 @@
 	style="container-type: inline-size"
 >
 	{#if showAddButton}
-		<div class="mr-auto flex items-center gap-3">
+		<div class="mr-auto flex items-center gap-2">
 			<ChatFormActionsAdd
 				{disabled}
 				{hasAudioModality}
@@ -117,8 +153,10 @@
 		</div>
 	{/if}
 
-	<div class="flex items-center gap-2">
-		<ChatFormReasoningToggle />
+	<div class="flex items-center gap-1.5">
+		{#if hasProcessedTokens}
+			<ChatFormContextGauge />
+		{/if}
 
 		{#if showModelSelector}
 			<ChatFormActionModels

tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormReasoningEffortSubmenu.svelte

@@ -6,6 +6,7 @@
 	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';
+	import { DIALOG_SUBMENU_CONTENT } from '$lib/constants/css-classes';
 	import {
 		modelsStore,
 		checkModelSupportsThinking,
@@ -71,9 +72,7 @@
 
 {#if modelSupportsThinking}
 	<DropdownMenu.Sub bind:open={subOpen}>
-		<DropdownMenu.SubTrigger
-			class="flex cursor-pointer items-center gap-2 rounded-md px-2.5 py-1.5 text-sm transition-colors outline-none hover:bg-accent focus:bg-accent"
-		>
+		<DropdownMenu.SubTrigger class="flex cursor-pointer items-center gap-2">
 			{#if thinkingEnabled}
 				<Lightbulb class="h-4 w-4 shrink-0 text-amber-400" />
 			{:else}
@@ -89,23 +88,15 @@
 			{/if}
 		</DropdownMenu.SubTrigger>
 
-		<DropdownMenu.SubContent
-			class="w-60 rounded-xl bg-popover p-3 text-popover-foreground shadow-md outline-none data-[side=bottom]:slide-in-from-top-2 data-[side=left]:slide-in-from-right-2 data-[side=right]:slide-in-from-left-2 data-[side=top]:slide-in-from-bottom-2 data-[state=closed]:animate-out data-[state=closed]:fade-out-0 data-[state=closed]:zoom-out-95 data-[state=open]:animate-in data-[state=open]:fade-in-0 data-[state=open]:zoom-in-95"
-		>
+		<DropdownMenu.SubContent class={DIALOG_SUBMENU_CONTENT}>
 			{#each REASONING_EFFORT_LEVELS as level (level.value)}
 				<button
 					type="button"
-					class="flex w-full cursor-pointer items-center gap-2 rounded-lg px-2.5 py-2 text-left text-sm transition-colors hover:bg-accent"
+					class="flex w-full cursor-pointer items-center gap-2"
 					class:bg-accent={isSelected(level)}
 					onclick={() => handleSelection(level)}
 				>
-					{#if isSelected(level)}
-						<Check class="h-4 w-4 shrink-0 text-foreground" />
-					{:else}
-						<div class="h-4 w-4 shrink-0"></div>
-					{/if}
-
-					<span class="flex-1">{level.label}</span>
+					<span class="flex-1 text-left">{level.label}</span>
 
 					{#if !level.isOff}
 						<span class="text-[11px] text-muted-foreground opacity-60">
@@ -125,6 +116,10 @@
 							</Tooltip.Content>
 						</Tooltip.Root>
 					{/if}
+
+					{#if isSelected(level)}
+						<Check class="h-4 w-4 shrink-0 text-foreground" />
+					{/if}
 				</button>
 			{/each}
 		</DropdownMenu.SubContent>

tools/ui/src/lib/components/app/chat/ChatForm/ChatFormContextGauge/ChatFormContextGauge.svelte

@@ -0,0 +1,108 @@
+<script lang="ts">
+	import { untrack } from 'svelte';
+	import * as HoverCard from '$lib/components/ui/hover-card';
+	import { activeConversation, activeMessages } from '$lib/stores/conversations.svelte';
+	import { chatStore, isChatStreaming, isLoading } from '$lib/stores/chat.svelte';
+	import { formatParameters } from '$lib/utils/formatters';
+	import { useContextGauge } from '$lib/hooks/use-context-gauge.svelte';
+	import ContextGaugeDial from './ContextGaugeDial.svelte';
+	import ContextGaugeDetails from './ContextGaugeDetails.svelte';
+	import ContextGaugeLoadModel from './ContextGaugeLoadModel.svelte';
+	import { colorLevelBgClass, colorLevelTextClass } from './context-gauge';
+
+	const gauge = useContextGauge();
+
+	$effect(() => {
+		const conv = activeConversation();
+		untrack(() => chatStore.setActiveProcessingConversation(conv?.id ?? null));
+	});
+
+	$effect(() => {
+		const conv = activeConversation();
+		const messages = activeMessages() as DatabaseMessage[];
+		if (!conv) return;
+		if (isLoading() || isChatStreaming()) return;
+
+		if (messages.length === 0) {
+			untrack(() => chatStore.clearProcessingState(conv.id));
+			return;
+		}
+
+		untrack(() => chatStore.restoreProcessingStateFromMessages(messages, conv.id));
+	});
+
+	$effect(() => {
+		gauge.startMonitoring();
+	});
+
+	const showProgressBar = $derived(
+		gauge.contextTotal !== null &&
+			gauge.contextTotal > 0 &&
+			(gauge.activeModelId !== null || gauge.isActiveModelLoaded)
+	);
+</script>
+
+<HoverCard.Root>
+	<HoverCard.Trigger class="flex h-5 w-5 cursor-default items-center justify-center">
+		<ContextGaugeDial percent={gauge.contextPercent} level={gauge.colorLevel} />
+	</HoverCard.Trigger>
+
+	<HoverCard.Content
+		side="bottom"
+		class="z-50 w-64 rounded-lg border border-border/50 bg-popover p-3 text-popover-foreground shadow-lg"
+	>
+		<div class="flex flex-col gap-2">
+			<div class="flex items-center gap-2">
+				<span class="font-medium">Context</span>
+				<span class="text-muted-foreground">·</span>
+				<span class="font-mono text-muted-foreground">
+					{formatParameters(gauge.contextUsed)}
+					/ {gauge.contextTotal !== null ? formatParameters(gauge.contextTotal) : '-'}
+				</span>
+			</div>
+
+			{#if gauge.activeModelId !== null && !gauge.isActiveModelLoaded}
+				<ContextGaugeLoadModel
+					modelId={gauge.activeModelId}
+					isLoading={gauge.isActiveModelLoading}
+					onLoad={gauge.loadModel}
+				/>
+			{:else if showProgressBar}
+				<div class="h-1.5 w-full overflow-hidden rounded-full bg-muted">
+					<div
+						class="h-full rounded-full transition-all duration-300 {colorLevelBgClass(
+							gauge.colorLevel
+						)}"
+						style="width: {gauge.contextPercent}%"
+					></div>
+				</div>
+
+				<div class="flex justify-between text-xs text-muted-foreground">
+					<span>
+						<span class={colorLevelTextClass(gauge.colorLevel)}>{gauge.contextPercent}%</span> used
+					</span>
+					<span>
+						{formatParameters((gauge.contextTotal ?? 0) - gauge.contextUsed)} remaining
+					</span>
+				</div>
+			{:else}
+				<div class="text-xs text-muted-foreground">No context info available</div>
+			{/if}
+
+			{#if gauge.hasAnyUsage}
+				<ContextGaugeDetails
+					currentRead={gauge.currentRead}
+					currentFresh={gauge.currentFresh}
+					currentCache={gauge.currentCache}
+					currentOutput={gauge.currentOutput}
+					kvTotal={gauge.kvTotal}
+					cumulativeRead={gauge.cumulativeRead}
+					cumulativeOutput={gauge.cumulativeOutput}
+					cumulativeCacheTotal={gauge.cumulativeCacheTotal}
+					averageTokensPerSecond={gauge.averageTokensPerSecond}
+					transientDetails={gauge.transientDetails}
+				/>
+			{/if}
+		</div>
+	</HoverCard.Content>
+</HoverCard.Root>

tools/ui/src/lib/components/app/chat/ChatForm/ChatFormContextGauge/ContextGaugeDetailRow.svelte

@@ -0,0 +1,20 @@
+<script lang="ts">
+	interface Props {
+		label: string;
+		value: string;
+		subtitle?: string;
+	}
+
+	let { label, value, subtitle }: Props = $props();
+</script>
+
+<div class="grid gap-1.5">
+	<div class="flex items-baseline justify-between">
+		<span class="text-muted-foreground">{label}</span>
+		<span class="font-mono text-muted-foreground">{value}</span>
+	</div>
+
+	{#if subtitle}
+		<div class="text-[10px] leading-tight text-muted-foreground/70">{subtitle}</div>
+	{/if}
+</div>

tools/ui/src/lib/components/app/chat/ChatForm/ChatFormContextGauge/ContextGaugeDetails.svelte

@@ -0,0 +1,122 @@
+<script lang="ts">
+	import { ChevronDown } from '@lucide/svelte';
+	import * as Collapsible from '$lib/components/ui/collapsible';
+	import { STATS_UNITS } from '$lib/constants';
+	import ContextGaugeDetailRow from './ContextGaugeDetailRow.svelte';
+
+	interface Props {
+		currentRead: number;
+		currentFresh: number;
+		currentCache: number;
+		currentOutput: number;
+		kvTotal: number;
+		cumulativeRead: number;
+		cumulativeOutput: number;
+		cumulativeCacheTotal: number;
+		averageTokensPerSecond: number | null;
+		transientDetails: string[];
+	}
+
+	let {
+		currentRead,
+		currentFresh,
+		currentCache,
+		currentOutput,
+		kvTotal,
+		cumulativeRead,
+		cumulativeOutput,
+		cumulativeCacheTotal,
+		averageTokensPerSecond,
+		transientDetails
+	}: Props = $props();
+
+	let open = $state(false);
+
+	const hasCumulative = $derived(cumulativeRead > 0 || cumulativeOutput > 0);
+	const hasCurrent = $derived(currentRead > 0 || currentOutput > 0);
+</script>
+
+<Collapsible.Root bind:open class="mt-3 border-t border-border/50 pt-4">
+	<Collapsible.Trigger
+		class="flex w-full cursor-pointer items-center gap-1 text-xs text-muted-foreground hover:text-foreground"
+	>
+		<span>Token usage details</span>
+
+		<ChevronDown class={'ml-auto h-3 w-3 transition-transform' + (open ? ' rotate-180' : '')} />
+	</Collapsible.Trigger>
+
+	<Collapsible.Content class="flex flex-col gap-4 text-xs pt-4">
+		{#if hasCumulative}
+			<div>
+				<h3 class="text-[11px] font-medium uppercase tracking-wide text-muted-foreground/70 mb-2">
+					Across all turns
+				</h3>
+
+				<div class="flex flex-col gap-2">
+					{#if cumulativeRead > 0}
+						<ContextGaugeDetailRow
+							label="Prompt tokens evaluated"
+							value={`${cumulativeRead.toLocaleString()} tok`}
+							subtitle={cumulativeCacheTotal > 0
+								? `${cumulativeCacheTotal.toLocaleString()} reused from KV cache`
+								: undefined}
+						/>
+					{/if}
+					{#if cumulativeOutput > 0}
+						<ContextGaugeDetailRow
+							label="Tokens generated"
+							value={`${cumulativeOutput.toLocaleString()} tok`}
+						/>
+					{/if}
+				</div>
+			</div>
+		{/if}
+
+		{#if hasCurrent}
+			<div>
+				<h3 class="text-[11px] font-medium uppercase tracking-wide text-muted-foreground/70 mb-2">
+					This turn · KV cache
+				</h3>
+
+				<div class="flex flex-col gap-2">
+					{#if currentRead > 0}
+						<ContextGaugeDetailRow
+							label="Prompt"
+							value={`${currentRead.toLocaleString()} tok`}
+							subtitle={currentCache > 0
+								? `${currentFresh.toLocaleString()} fresh + ${currentCache.toLocaleString()} cached`
+								: undefined}
+						/>
+					{/if}
+
+					{#if currentOutput > 0}
+						<ContextGaugeDetailRow
+							label="Generated"
+							value={`${currentOutput.toLocaleString()} tok`}
+						/>
+					{/if}
+
+					<div class="pt-1 mt-0.5 border-t border-border/30">
+						<div class="flex justify-between">
+							<span class="text-muted-foreground">KV cache total</span>
+							<span class="font-mono font-medium">{kvTotal.toLocaleString()} tok</span>
+						</div>
+					</div>
+				</div>
+			</div>
+		{/if}
+
+		{#if averageTokensPerSecond !== null}
+			<div class="pt-1.5 mt-1 border-t border-border/30">
+				<ContextGaugeDetailRow
+					label="Avg speed"
+					value={`${averageTokensPerSecond.toFixed(1)}${STATS_UNITS.TOKENS_PER_SECOND}`}
+				/>
+			</div>
+		{/if}
+
+		{#each transientDetails as detail (detail)}
+			<div class="font-mono text-muted-foreground">{detail}</div>
+		{/each}
+	</Collapsible.Content>
+</Collapsible.Root>

tools/ui/src/lib/components/app/chat/ChatForm/ChatFormContextGauge/ContextGaugeDial.svelte

@@ -0,0 +1,43 @@
+<script lang="ts">
+	import type { ColorLevel } from './context-gauge';
+	import { colorLevelTextClass } from './context-gauge';
+
+	interface Props {
+		percent: number | null;
+		level: ColorLevel;
+		size?: 'sm' | 'md';
+	}
+
+	let { percent, level, size = 'sm' }: Props = $props();
+
+	const RADIUS = 11;
+	const CIRCUMFERENCE = 2 * Math.PI * RADIUS;
+
+	const strokeLevelClass = $derived(colorLevelTextClass(level));
+	const dimensions = $derived(size === 'md' ? 'h-6 w-6' : 'h-5 w-5');
+	const strokeWidth = $derived(size === 'md' ? 4 : 3);
+</script>
+
+<svg viewBox="0 0 32 32" fill="none" class={dimensions}>
+	<circle
+		cx="16"
+		cy="16"
+		r={RADIUS}
+		stroke="currentColor"
+		stroke-opacity="0.1"
+		stroke-width={strokeWidth}
+	/>
+
+	<circle
+		cx="16"
+		cy="16"
+		r={RADIUS}
+		class="transition-colors duration-300 {strokeLevelClass}"
+		stroke="currentColor"
+		stroke-width={strokeWidth}
+		stroke-linecap="round"
+		stroke-dasharray={CIRCUMFERENCE}
+		stroke-dashoffset={percent !== null ? CIRCUMFERENCE * (1 - percent / 100) : CIRCUMFERENCE}
+		transform="rotate(-90 16 16)"
+	/>
+</svg>

tools/ui/src/lib/components/app/chat/ChatForm/ChatFormContextGauge/ContextGaugeLoadModel.svelte

@@ -0,0 +1,24 @@
+<script lang="ts">
+	import { Loader2 } from '@lucide/svelte';
+	import { Button } from '$lib/components/ui/button';
+
+	interface Props {
+		modelId: string | null;
+		isLoading: boolean;
+		onLoad: () => void;
+	}
+
+	let { modelId, isLoading, onLoad }: Props = $props();
+</script>
+
+{#if modelId !== null && !isLoading}
+	<div class="flex flex-col gap-2 border-t border-border/50 pt-2 text-xs text-muted-foreground">
+		<span>Available context size is only visible once the model is loaded.</span>
+		<Button size="sm" variant="secondary" class="self-start" onclick={onLoad}>Load model</Button>
+	</div>
+{:else if isLoading}
+	<div class="flex items-center gap-2 border-t border-border/50 pt-2 text-xs text-muted-foreground">
+		<Loader2 class="h-3.5 w-3.5 animate-spin" />
+		<span>Loading model...</span>
+	</div>
+{/if}

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';
+	}
+}

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 @@
 		<ChatMessageUser
 			class={className}
 			{deletionInfo}
+			{isLastUserMessage}
 			{message}
+			{nextAssistantMessage}
 			onConfirmDelete={handleConfirmDelete}
 			onCopy={handleCopy}
 			onDelete={handleDelete}

tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageAssistant/ChatMessageAssistant.svelte

@@ -11,11 +11,10 @@
 	import { useProcessingState } from '$lib/hooks/use-processing-state.svelte';
 	import { isLoading, isChatStreaming } from '$lib/stores/chat.svelte';
 	import { copyToClipboard, deriveAgenticSections, modelLoadProgressText } from '$lib/utils';
-	import { AgenticSectionType } from '$lib/enums';
+	import { AgenticSectionType, ChatMessageStatisticsMode } from '$lib/enums';
 	import { REASONING_TAGS } from '$lib/constants/agentic';
-	import { tick } from 'svelte';
 	import { fade } from 'svelte/transition';
-	import { MessageRole, ChatMessageStatsView } from '$lib/enums';
+	import { MessageRole } from '$lib/enums';
 	import { config } from '$lib/stores/settings.svelte';
 	import { isRouterMode } from '$lib/stores/server.svelte';
 	import { modelsStore } from '$lib/stores/models.svelte';
@@ -122,62 +121,6 @@
 		return parts.join('\n\n\n');
 	});
 
-	let activeStatsView = $state<ChatMessageStatsView>(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 @@
 
 	<div class="info my-6 grid gap-4 tabular-nums">
 		{#if displayedModel}
-			<div
-				bind:this={statsContainerEl}
-				class="inline-flex flex-wrap items-start gap-2 text-xs text-muted-foreground"
-			>
+			<div class="inline-flex flex-wrap items-start gap-2 text-xs text-muted-foreground">
 				{#if isRouter}
 					<ModelsSelectorDropdown
 						currentModel={pendingModel ?? displayedModel}
@@ -347,28 +286,25 @@
 				{#if currentConfig.showMessageStats && message.timings && message.timings.predicted_n && message.timings.predicted_ms}
 					{@const agentic = message.timings.agentic}
 					<ChatMessageStatistics
+						mode={ChatMessageStatisticsMode.GENERATION}
 						promptTokens={agentic ? agentic.llm.prompt_n : message.timings.prompt_n}
 						promptMs={agentic ? agentic.llm.prompt_ms : message.timings.prompt_ms}
 						predictedTokens={agentic ? agentic.llm.predicted_n : message.timings.predicted_n}
 						predictedMs={agentic ? agentic.llm.predicted_ms : message.timings.predicted_ms}
 						agenticTimings={agentic}
-						onActiveViewChange={handleStatsViewChange}
 					/>
 				{: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}
 						<ChatMessageStatistics
+							mode={ChatMessageStatisticsMode.GENERATION}
 							isLive
-							isProcessingPrompt={!!isStillProcessingPrompt}
 							promptTokens={liveStats?.tokensProcessed}
 							promptMs={liveStats?.timeMs}
-							predictedTokens={genStats?.tokensGenerated}
-							predictedMs={genStats?.timeMs}
+							predictedTokens={genStats.tokensGenerated}
+							predictedMs={genStats.timeMs}
 						/>
 					{/if}
 				{/if}

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();
+		}
+	});
 </script>
 
 <div
@@ -60,6 +99,37 @@
 			renderMarkdown={true}
 		/>
 
+		{#if showStoredReadingStats}
+			<!-- Reading stats sourced from the assistant message that followed this turn -->
+			<div class="info my-2 grid w-full justify-items-end gap-4 tabular-nums">
+				<div
+					class="inline-flex flex-wrap items-start justify-end gap-2 text-xs text-muted-foreground"
+				>
+					<ChatMessageStatistics
+						mode={ChatMessageStatisticsMode.READING}
+						promptTokens={storedReadingStats!.promptTokens}
+						promptMs={storedReadingStats!.promptMs}
+					/>
+				</div>
+			</div>
+		{:else if showLiveReadingStats}
+			{@const liveStats = processingState.getLiveProcessingStats()}
+			{#if liveStats}
+				<div class="info my-2 grid w-full justify-items-end gap-4 tabular-nums">
+					<div
+						class="inline-flex flex-wrap items-start justify-end gap-2 text-xs text-muted-foreground"
+					>
+						<ChatMessageStatistics
+							mode={ChatMessageStatisticsMode.READING}
+							isLive
+							promptTokens={liveStats.tokensProcessed}
+							promptMs={liveStats.timeMs}
+						/>
+					</div>
+				</div>
+			{/if}
+		{/if}
+
 		{#if message.timestamp}
 			<div class="max-w-[80%]">
 				<ChatMessageActionIcons

tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageUser/ChatMessageUserPending.svelte

@@ -2,7 +2,6 @@
 	import { ActionIcon, ChatMessageEditForm, ChatMessageUserBubble } from '$lib/components/app';
 	import { fadeInView } from '$lib/actions/fade-in-view.svelte';
 	import { ArrowUp, Edit, Trash2 } from '@lucide/svelte';
-	import { getProcessingInfoContext } from '$lib/contexts';
 	import { useMessageEditContext } from '$lib/hooks/use-message-edit-context.svelte';
 
 	interface Props {
@@ -23,9 +22,6 @@
 		onDelete
 	}: Props = $props();
 
-	const processingInfoCtx = getProcessingInfoContext();
-	let showProcessingInfo = $derived(processingInfoCtx.showProcessingInfo);
-
 	const editCtx = useMessageEditContext({
 		getContent: () => content,
 		getExtras: () => extras,
@@ -36,9 +32,7 @@
 <div
 	use:fadeInView
 	aria-label="Pending user message"
-	class="group flex flex-col items-end gap-3 transition-opacity hover:opacity-80 md:gap-2 {className} sticky {showProcessingInfo
-		? 'bottom-44'
-		: 'bottom-32'}"
+	class="group flex flex-col items-end gap-3 transition-opacity hover:opacity-80 md:gap-2 {className} sticky bottom-32"
 	role="group"
 >
 	{#if editCtx.isEditing}

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<number, boolean> = $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}
 
 <div class="agentic-content">
-	{#if highlightTurns && turnGroups.length > 1}
+	{#if turnGroups.length > 1}
 		{#each turnGroups as turn, turnIndex (turnIndex)}
 			{@const turnStats = message?.timings?.agentic?.perTurn?.[turnIndex]}
-			<div class="agentic-turn my-2 hover:bg-muted/80 dark:hover:bg-muted/30">
-				<span class="agentic-turn-label">Turn {turnIndex + 1}</span>
+
+			<div class="agentic-turn group/turn grid gap-3 mb-4">
 				{#each turn.sections as section, sIdx (turn.flatIndices[sIdx])}
 					{@render renderSection(section, turn.flatIndices[sIdx])}
 				{/each}
-				{#if turnStats}
-					<div class="turn-stats">
+
+				{#if turnStats && showMessageStats}
+					<div class="turn-stats transition-opacity duration-150">
 						<ChatMessageStatistics
 							promptTokens={turnStats.llm.prompt_n}
 							promptMs={turnStats.llm.prompt_ms}
@@ -402,39 +402,21 @@
 	.agentic-content {
 		display: flex;
 		flex-direction: column;
-		gap: 0.5rem;
 		width: 100%;
 		max-width: 48rem;
+		gap: 1rem;
+	}
+
+	.agentic-content > :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);
 	}
 </style>

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}
 
 <div class="inline-flex items-center text-xs text-muted-foreground">
-	<div class="inline-flex items-center rounded-sm bg-muted-foreground/15 p-0.5">
-		{#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}
+		<div class="inline-flex items-center rounded-sm bg-muted-foreground/15 p-0.5">
+			{#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}
-	</div>
+
+			{@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}
+		</div>
+	{/if}
 
 	<div class="flex items-center gap-1 px-2">
 		{#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)}
 			<ChatMessageStatisticsBadge
 				class="bg-transparent"
 				icon={WholeWord}

tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessages.svelte

@@ -186,6 +186,8 @@
 			message: DatabaseMessage;
 			toolMessages: DatabaseMessage[];
 			isLastAssistantMessage: boolean;
+			isLastUserMessage: boolean;
+			nextAssistantMessage: DatabaseMessage | null;
 			siblingInfo: ChatMessageSiblingInfo;
 		}> = [];
 
@@ -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;
 	});
 </script>
@@ -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)}
 		<ChatMessage
 			class="mx-auto mt-12 w-full max-w-3xl"
 			{message}
 			{toolMessages}
 			{isLastAssistantMessage}
+			{isLastUserMessage}
+			{nextAssistantMessage}
 			{siblingInfo}
 		/>
 	{/each}

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}
-					<ChatScreenProcessingInfo />
-				{/if}
 			</div>
 
 			<ChatScreenForm

tools/ui/src/lib/components/app/chat/ChatScreen/ChatScreenProcessingInfo.svelte

@@ -1,127 +0,0 @@
-<script lang="ts">
-	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));
-			}
-		}
-	});
-</script>
-
-<div
-	class={[
-		'chat-processing-info-container pointer-events-none relative w-full hidden md:block',
-		processingVisible && 'visible'
-	]}
->
-	<div class="chat-processing-info-content absolute bottom-4 left-1/2 -translate-x-1/2">
-		{#each processingDetails as detail (detail)}
-			<span class="chat-processing-info-detail pointer-events-auto backdrop-blur-sm">{detail}</span>
-		{/each}
-	</div>
-</div>
-
-<style>
-	.chat-processing-info-container {
-		position: sticky;
-		top: 0;
-		z-index: 10;
-		padding: 0 1rem 0.75rem;
-		opacity: 0;
-		transform: translateY(50%);
-		transition:
-			opacity 300ms ease-out,
-			transform 300ms ease-out;
-	}
-
-	.chat-processing-info-container.visible {
-		opacity: 1;
-		transform: translateY(0);
-	}
-
-	.chat-processing-info-content {
-		display: flex;
-		flex-wrap: wrap;
-		align-items: center;
-		gap: 1rem;
-		justify-content: center;
-		max-width: 48rem;
-		margin: 0 auto;
-	}
-
-	.chat-processing-info-detail {
-		color: var(--muted-foreground);
-		font-size: 0.75rem;
-		padding: 0.25rem 0.75rem;
-		border-radius: 0.375rem;
-		font-family:
-			ui-monospace, SFMono-Regular, 'SF Mono', Consolas, 'Liberation Mono', Menlo, monospace;
-		white-space: nowrap;
-	}
-
-	@media (max-width: 768px) {
-		.chat-processing-info-content {
-			gap: 0.5rem;
-		}
-
-		.chat-processing-info-detail {
-			font-size: 0.7rem;
-			padding: 0.2rem 0.5rem;
-		}
-	}
-</style>

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.

tools/ui/src/lib/components/app/content/CollapsibleContentBlock.svelte

@@ -76,7 +76,7 @@
 		open = value;
 		onToggle?.();
 	}}
-	class={className}
+	class="{className} my-0!"
 >
 	<Card class="gap-0 border-muted bg-muted/30 py-0">
 		<Collapsible.Trigger class="flex w-full cursor-pointer items-start justify-between gap-2 p-3">

tools/ui/src/lib/components/app/content/SyntaxHighlightedCode.svelte

@@ -72,8 +72,8 @@
 </script>
 
 <div
-	class="code-preview-wrapper rounded-lg border border-border bg-muted {className}"
-	style="max-height: {maxHeight}; max-width: {maxWidth};"
+	class="code-preview-wrapper min-w-0 max-w-full overflow-x-auto rounded-lg border border-border bg-muted {className}"
+	style="max-height: {maxHeight}; {maxWidth ? `max-width: ${maxWidth};` : ''}"
 >
 	<!-- Needs to be formatted as single line for proper rendering -->
 	<pre class="m-0"><code class="hljs text-sm leading-relaxed">{@html highlightedHtml}</code></pre>

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<MCPServerSettingsEntry[]>([]);
+	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 @@
 		<Dialog.Footer>
 			<Button variant="secondary" size="sm" onclick={() => handleOpenChange(false)}>Not now</Button>
 
-			<Button variant="default" size="sm" onclick={enableSelected}>Add selected</Button>
+			<Button
+				variant="default"
+				size="sm"
+				onclick={enableSelected}
+				disabled={footerLabel === 'Continue'}>{footerLabel}</Button
+			>
 		</Dialog.Footer>
 	</Dialog.Content>
 </Dialog.Root>

tools/ui/src/lib/components/app/settings/SettingsChat/SettingsChatToolsTab.svelte

@@ -39,13 +39,17 @@
 					{@const faviconUrl = group.serverId ? mcpStore.getServerFavicon(group.serverId) : null}
 
 					<span class="inline-flex min-w-0 items-center gap-1.5 font-medium">
-						<McpServerIdentity
-							iconClass="h-4 w-4"
-							iconRounded="rounded-sm"
-							showVersion={false}
-							displayName={group.label}
-							{faviconUrl}
-						/>
+						{#if group.source === 'mcp'}
+							<McpServerIdentity
+								iconClass="h-4 w-4"
+								iconRounded="rounded-sm"
+								showVersion={false}
+								displayName={group.label}
+								{faviconUrl}
+							/>
+						{:else}
+							<TruncatedText text={group.label} class="font-medium" />
+						{/if}
 					</span>
 
 					<span class="ml-auto shrink-0 text-xs text-muted-foreground">

tools/ui/src/lib/components/ui/hover-card/hover-card-content.svelte

@@ -0,0 +1,31 @@
+<script lang="ts">
+	import { LinkPreview as HoverCardPrimitive } from 'bits-ui';
+	import { cn, type WithoutChildrenOrChild } from '$lib/components/ui/utils.js';
+	import HoverCardPortal from './hover-card-portal.svelte';
+	import type { ComponentProps } from 'svelte';
+
+	let {
+		ref = $bindable(null),
+		class: className,
+		align = 'center',
+		sideOffset = 4,
+		portalProps,
+		...restProps
+	}: HoverCardPrimitive.ContentProps & {
+		portalProps?: WithoutChildrenOrChild<ComponentProps<typeof HoverCardPortal>>;
+	} = $props();
+</script>
+
+<HoverCardPortal {...portalProps}>
+	<HoverCardPrimitive.Content
+		bind:ref
+		data-slot="hover-card-content"
+		{align}
+		{sideOffset}
+		class={cn(
+			'data-open:animate-in data-closed:animate-out data-closed:fade-out-0 data-open:fade-in-0 data-closed:zoom-out-95 data-open:zoom-in-95 data-[side=bottom]:slide-in-from-top-2 data-[side=left]:slide-in-from-right-2 data-[side=right]:slide-in-from-left-2 data-[side=top]:slide-in-from-bottom-2 ring-foreground/10 bg-popover text-popover-foreground w-64 rounded-lg p-2.5 text-sm shadow-md ring-1 duration-100 z-50 origin-(--transform-origin) outline-hidden',
+			className
+		)}
+		{...restProps}
+	/>
+</HoverCardPortal>

tools/ui/src/lib/components/ui/hover-card/hover-card-portal.svelte

@@ -0,0 +1,7 @@
+<script lang="ts">
+	import { LinkPreview as HoverCardPrimitive } from 'bits-ui';
+
+	let { ...restProps }: HoverCardPrimitive.PortalProps = $props();
+</script>
+
+<HoverCardPrimitive.Portal {...restProps} />

tools/ui/src/lib/components/ui/hover-card/hover-card-trigger.svelte

@@ -0,0 +1,7 @@
+<script lang="ts">
+	import { LinkPreview as HoverCardPrimitive } from 'bits-ui';
+
+	let { ref = $bindable(null), ...restProps }: HoverCardPrimitive.TriggerProps = $props();
+</script>
+
+<HoverCardPrimitive.Trigger bind:ref data-slot="hover-card-trigger" {...restProps} />

tools/ui/src/lib/components/ui/hover-card/hover-card.svelte

@@ -0,0 +1,7 @@
+<script lang="ts">
+	import { LinkPreview as HoverCardPrimitive } from 'bits-ui';
+
+	let { open = $bindable(false), ...restProps }: HoverCardPrimitive.RootProps = $props();
+</script>
+
+<HoverCardPrimitive.Root bind:open {...restProps} />

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
+};

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';

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';

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<string, string> = {
+	[ReasoningEffort.OFF]: 'Off',
 	[ReasoningEffort.LOW]: 'Low',
 	[ReasoningEffort.MEDIUM]: 'Medium',
 	[ReasoningEffort.HIGH]: 'High',
@@ -13,7 +14,7 @@ export const REASONING_EFFORT_LABELS: Record<string, string> = {
 };
 
 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' },

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

tools/ui/src/lib/constants/settings-registry.ts

@@ -258,18 +258,6 @@ const SETTINGS_REGISTRY: Record<string, SettingsSectionEntry> = {
 					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<string, SettingsSectionEntry> = {
 					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',

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`;

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';

tools/ui/src/lib/contexts/processing-info.context.ts

@@ -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);
-}

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.
  */

tools/ui/src/lib/enums/index.ts

@@ -10,6 +10,7 @@ export { AgenticSectionType, ContinueIntentKind, ToolCallType } from './agentic.
 
 export {
 	ChatMessageStatsView,
+	ChatMessageStatisticsMode,
 	StreamConnectionState,
 	ContentPartType,
 	ConversationSelectionMode,

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',

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<void>;
+	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<typeof useProcessingState>['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()
+	};
+}

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));

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 {

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<string>} Promise that resolves to the generated content string
 	 * @throws {Error} if the response cannot be parsed or is malformed
 	 */

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;

tools/ui/src/lib/stores/agentic.svelte.ts

@@ -477,7 +477,7 @@ class AgenticStore {
 		conversationId: string;
 		messages: ApiChatMessageData[];
 		options: AgenticFlowOptions;
-		tools: ReturnType<typeof mcpStore.getToolDefinitionsForLLM>;
+		tools: ReturnType<typeof toolsStore.getEnabledToolsForLLM>;
 		agenticConfig: AgenticConfig;
 		callbacks: AgenticFlowCallbacks;
 		signal?: AbortSignal;

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);

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<McpServerOverride[]>(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<ReasoningEffort>(ConversationsStore.loadReasoningEffortDefault());
+	pendingReasoningEffort = $state<ReasoningEffort | ReasoningEffort.OFF>(
+		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<void> {
 		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;
 		}

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<string> {
-		const ids = new SvelteSet<string>();
-		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<boolean> {
@@ -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<string, unknown>) ?? {
-					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<string, unknown>): Record<string, unknown> {
-		if (!schema || typeof schema !== 'object') {
-			return schema;
-		}
-
-		const normalized = { ...schema };
-		if (normalized.properties && typeof normalized.properties === 'object') {
-			const props = normalized.properties as Record<string, Record<string, unknown>>;
-			const normalizedProps: Record<string, Record<string, unknown>> = {};
-			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<string, unknown>)
-					);
-				if (normalizedProp.items && typeof normalizedProp.items === 'object')
-					normalizedProp.items = this.normalizeSchemaProperties(
-						normalizedProp.items as Record<string, unknown>
-					);
-				normalizedProps[key] = normalizedProp;
-			}
-			normalized.properties = normalizedProps;
-		}
-
-		return normalized;
-	}
-
 	getToolNames(): string[] {
 		return Array.from(this.toolsIndex.keys());
 	}