cli : move to HTTP-based implementation (#24948 )

* cli: move to HTTP-based implementation * wip * working * remote server ok * cli support router mode Co-authored-by: Piotr Wilkin <ilintar@gmail.com> * case: router with only one model * Apply suggestions from code review Co-authored-by: Piotr Wilkin (ilintar) <piotr.wilkin@syndatis.com> * remove outdated comment * use destructor instead * add ftype * cli-view --> cli-ui * pimpl * no more json in header * nits fixes * also show model aliases --------- Co-authored-by: Piotr Wilkin <ilintar@gmail.com> Co-authored-by: Piotr Wilkin (ilintar) <piotr.wilkin@syndatis.com>
Make hip quality check run on all changes (#25403 )
2026-07-08 15:08:01 +02:00 · 2026-07-08 14:52:43 +02:00 · 2026-07-08 14:38:51 +02:00 · 2026-07-08 19:24:20 +08:00 · 2026-07-08 18:18:09 +08:00 · 2026-07-08 12:02:50 +03:00
123 changed files with 4354 additions and 2090 deletions
@@ -9,6 +9,8 @@ on:
      '.github/workflows/hip-quality-check.yml',
      '**/*.cu',
      '**/*.cuh',
+      'ggml/src/ggml-hip/CMakeLists.txt',
+      'ggml/src/ggml-cuda/vendors/hip.h',
      'scripts/hip/gcn-cdna-vgpr-check.py'
    ]

@@ -18,6 +20,8 @@ on:
      '.github/workflows/hip-quality-check.yml',
      '**/*.cu',
      '**/*.cuh',
+      'ggml/src/ggml-hip/CMakeLists.txt',
+      'ggml/src/ggml-cuda/vendors/hip.h',
      'scripts/hip/gcn-cdna-vgpr-check.py'
    ]

@@ -27,6 +27,7 @@
 #include <cinttypes>
 #include <climits>
 #include <cstdarg>
+#include <filesystem>
 #include <fstream>
 #include <list>
 #include <regex>
@@ -718,9 +719,8 @@ static bool common_params_parse_ex(int argc, char ** argv, common_params_context

        // model is required (except for server)
        // TODO @ngxson : maybe show a list of available models in CLI in this case
-        if (params.model.path.empty()
-                && !params.usage
-                && !params.completion) {
+        bool can_skip_model = params.usage || params.completion || !params.server_base.empty();
+        if (!can_skip_model && params.model.path.empty()) {
            throw std::invalid_argument("error: --model is required\n");
        }
    }
@@ -1240,6 +1240,13 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
            params.completion = true;
        }
    ));
+    add_opt(common_arg(
+        {"--server-base"}, "URL",
+        string_format("connect to this server instead of starting a new one, example: 'http://localhost:8080' (default: none)"),
+        [](common_params & params, const std::string & value) {
+            params.server_base = value;
+        }
+    ).set_examples({LLAMA_EXAMPLE_CLI}));
    add_opt(common_arg(
        {"--verbose-prompt"},
        string_format("print a verbose prompt before generation (default: %s)", params.verbose_prompt ? "true" : "false"),
@@ -3451,9 +3458,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(
@@ -14,6 +14,7 @@
 #include <vector>
 #include <map>
 #include <algorithm>
+#include <fstream>

 #if defined(_WIN32) && !defined(_WIN32_WINNT)
 #define _WIN32_WINNT 0x0A00
@@ -643,6 +644,9 @@ struct common_params {

    std::map<std::string, std::string> default_template_kwargs;

+    // CLI params
+    std::string server_base; // if set, connect to this server instead of starting a new one
+
    // UI configs
    bool ui = true;
    bool ui_mcp_proxy = false;
@@ -2,6 +2,16 @@

 #include <cpp-httplib/httplib.h>

+#ifdef _WIN32
+#include <winsock2.h>
+#include <windows.h>
+#else
+#include <sys/socket.h>
+#include <netinet/in.h>
+#include <arpa/inet.h>
+#include <unistd.h>
+#endif
+
 struct common_http_url {
    std::string scheme;
    std::string user;
@@ -119,3 +129,63 @@ static std::pair<httplib::Client, common_http_url> common_http_client(const std:
 static std::string common_http_show_masked_url(const common_http_url & parts) {
    return parts.scheme + "://" + (parts.user.empty() ? "" : "****:****@") + common_http_format_host(parts.host) + parts.path;
 }
+
+static int common_http_get_free_port() {
+#ifdef _WIN32
+    WSADATA wsaData;
+    if (WSAStartup(MAKEWORD(2, 2), &wsaData) != 0) {
+        return -1;
+    }
+    typedef SOCKET native_socket_t;
+#define INVALID_SOCKET_VAL INVALID_SOCKET
+#define CLOSE_SOCKET(s) closesocket(s)
+#else
+    typedef int native_socket_t;
+#define INVALID_SOCKET_VAL -1
+#define CLOSE_SOCKET(s) close(s)
+#endif
+
+    native_socket_t sock = socket(AF_INET, SOCK_STREAM, 0);
+    if (sock == INVALID_SOCKET_VAL) {
+#ifdef _WIN32
+        WSACleanup();
+#endif
+        return -1;
+    }
+
+    struct sockaddr_in serv_addr;
+    std::memset(&serv_addr, 0, sizeof(serv_addr));
+    serv_addr.sin_family = AF_INET;
+    serv_addr.sin_addr.s_addr = htonl(INADDR_ANY);
+    serv_addr.sin_port = htons(0);
+
+    if (bind(sock, (struct sockaddr*)&serv_addr, sizeof(serv_addr)) != 0) {
+        CLOSE_SOCKET(sock);
+#ifdef _WIN32
+        WSACleanup();
+#endif
+        return -1;
+    }
+
+#ifdef _WIN32
+    int namelen = sizeof(serv_addr);
+#else
+    socklen_t namelen = sizeof(serv_addr);
+#endif
+    if (getsockname(sock, (struct sockaddr*)&serv_addr, &namelen) != 0) {
+        CLOSE_SOCKET(sock);
+#ifdef _WIN32
+        WSACleanup();
+#endif
+        return -1;
+    }
+
+    int port = ntohs(serv_addr.sin_port);
+
+    CLOSE_SOCKET(sock);
+#ifdef _WIN32
+    WSACleanup();
+#endif
+
+    return port;
+}
@@ -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) {
@@ -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);
@@ -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")
@@ -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:
@@ -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
@@ -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
@@ -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;
@@ -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,
@@ -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:
@@ -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;
@@ -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);
@@ -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;
            }
@@ -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;
 };

@@ -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;
                }
@@ -4389,10 +4709,16 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g
            } break;
        case GGML_OP_SET_ROWS:
            {
-                return (op->type == GGML_TYPE_F32 || op->type == GGML_TYPE_F16 || op->type == GGML_TYPE_BF16 ||
-                       op->type == GGML_TYPE_Q4_0 || op->type == GGML_TYPE_Q4_1 || op->type == GGML_TYPE_Q5_0 ||
-                       op->type == GGML_TYPE_Q5_1 || op->type == GGML_TYPE_Q8_0 || op->type == GGML_TYPE_IQ4_NL) &&
-                       op->src[0]->type == GGML_TYPE_F32 &&
+                return (
+                           (
+                               (op->type == GGML_TYPE_F32 || op->type == GGML_TYPE_F16 || op->type == GGML_TYPE_BF16 ||
+                               op->type == GGML_TYPE_Q4_0 || op->type == GGML_TYPE_Q4_1 || op->type == GGML_TYPE_Q5_0 ||
+                               op->type == GGML_TYPE_Q5_1 || op->type == GGML_TYPE_Q8_0 || op->type == GGML_TYPE_IQ4_NL) &&
+                               op->src[0]->type == GGML_TYPE_F32
+                           ) || (
+                               op->type == GGML_TYPE_F16 && op->src[0]->type == GGML_TYPE_F16
+                           )
+                       ) &&
                       (op->src[1]->type == GGML_TYPE_I64 || op->src[1]->type == GGML_TYPE_I32);
            } break;
        case GGML_OP_SET:
@@ -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;
    }

@@ -322,17 +322,77 @@ static void set_rows_cuda(ggml_backend_cuda_context & ctx, const ggml_tensor * s
    }
 }

+template<>
+void set_rows_cuda<half, int32_t>(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) {
+    const half    * src0_d = (const half *)src0->data;
+    const int32_t * src1_d = (const int32_t *)src1->data;
+
+    GGML_TENSOR_BINARY_OP_LOCALS
+
+    cudaStream_t stream = ctx.stream();
+
+
+    if (dst->type == GGML_TYPE_F16) {
+        set_rows_cuda(
+            src0_d, src1_d, (half*)dst->data,
+            ne00, ne01, ne02, ne03,
+            ne10, ne11, ne12, ne13,
+            nb01, nb02, nb03,
+            nb10, nb11, nb12,
+            nb1, nb2, nb3,
+            stream
+        );
+    } else {
+        GGML_ABORT("unsupported type %s", ggml_type_name(dst->type));
+    }
+}
+
+template<>
+void set_rows_cuda<half, int64_t>(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) {
+    const half    * src0_d = (const half *)src0->data;
+    const int64_t * src1_d = (const int64_t *)src1->data;
+
+    GGML_TENSOR_BINARY_OP_LOCALS
+
+    cudaStream_t stream = ctx.stream();
+
+
+    if (dst->type == GGML_TYPE_F16) {
+        set_rows_cuda(
+            src0_d, src1_d, (half*)dst->data,
+            ne00, ne01, ne02, ne03,
+            ne10, ne11, ne12, ne13,
+            nb01, nb02, nb03,
+            nb10, nb11, nb12,
+            nb1, nb2, nb3,
+            stream
+        );
+    } else {
+        GGML_ABORT("unsupported type %s", ggml_type_name(dst->type));
+    }
+}
+

 void ggml_cuda_op_set_rows(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
    const ggml_tensor * src0 = dst->src[0];
    const ggml_tensor * src1 = dst->src[1];

-    GGML_ASSERT(src0->type == GGML_TYPE_F32);
+    GGML_ASSERT(src0->type == GGML_TYPE_F32 || (src0->type == GGML_TYPE_F16 && dst->type == GGML_TYPE_F16));
    GGML_ASSERT(src1->type == GGML_TYPE_I64 || src1->type == GGML_TYPE_I32);

-    if (src1->type == GGML_TYPE_I64) {
-        set_rows_cuda<float, int64_t>(ctx, src0, src1, dst);
+    if (src0->type == GGML_TYPE_F32) {
+        if (src1->type == GGML_TYPE_I64) {
+            set_rows_cuda<float, int64_t>(ctx, src0, src1, dst);
+        } else {
+            set_rows_cuda<float, int32_t>(ctx, src0, src1, dst);
+        }
+    } else if (src0->type == GGML_TYPE_F16) {
+        if (src1->type == GGML_TYPE_I64) {
+            set_rows_cuda<half, int64_t>(ctx, src0, src1, dst);
+        } else {
+            set_rows_cuda<half, int32_t>(ctx, src0, src1, dst);
+        }
    } else {
-        set_rows_cuda<float, int32_t>(ctx, src0, src1, dst);
+        GGML_ABORT("unsupported type %s", ggml_type_name(src0->type));
    }
 }
@@ -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);
@@ -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);
@@ -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;
                }

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

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

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


@@ -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
    };
@@ -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) {
@@ -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);
@@ -17,6 +17,7 @@
 #include <cstring>
 #include <limits>
 #include <stdexcept>
+#include <string>

 //
 // llama_context
@@ -30,6 +31,30 @@ static llm_graph_type ctx_type_to_graph_type(llama_context_type ctx_type) {
    throw std::runtime_error("Unsupported ctx type");
 }

+struct llm_fused_op_probe {
+    llm_fused_op op;
+    const char * name;
+    uint32_t n_tokens_per_seq;
+};
+
+static const llm_fused_op_probe llm_fused_op_flash_attn_probe = {
+    /*.op               =*/ LLM_FUSED_OP_FLASH_ATTN,
+    /*.name             =*/ "Flash Attention",
+    /*.n_tokens_per_seq =*/ 1,
+};
+
+static const llm_fused_op_probe llm_fused_op_gdn_ar_probe = {
+    /*.op               =*/ LLM_FUSED_OP_GDN_AR,
+    /*.name             =*/ "fused Gated Delta Net (autoregressive)",
+    /*.n_tokens_per_seq =*/ 1,
+};
+
+static const llm_fused_op_probe llm_fused_op_gdn_ch_probe = {
+    /*.op               =*/ LLM_FUSED_OP_GDN_CH,
+    /*.name             =*/ "fused Gated Delta Net (chunked)",
+    /*.n_tokens_per_seq =*/ 16,
+};
+
 llama_context::llama_context(
        const llama_model & model,
              llama_context_params params) :
@@ -436,6 +461,69 @@ llama_context::~llama_context() {
    ggml_opt_free(opt_ctx);
 }

+void llama_context::resolve_fused_ops(const llama_memory_context_i * mctx, uint32_t n_seqs) {
+    const char * func = __func__;
+    auto resolve = [&](const llm_fused_op_probe & probe, bool & enabled) {
+        if (!enabled) {
+            return;
+        }
+
+        const uint32_t n_tokens_probe = probe.n_tokens_per_seq*n_seqs;
+
+        auto * gf = graph_reserve(n_tokens_probe, n_seqs, n_tokens_probe, mctx, true);
+        if (!gf) {
+            throw std::runtime_error(std::string("failed to reserve graph for ") + probe.name + " check");
+        }
+
+        bool device_mismatch = false;
+        for (const auto & node : get_gf_res_reserve()->get_fused_nodes()) {
+            if (node.op != probe.op) {
+                continue;
+            }
+
+            GGML_ASSERT(node.il >= 0);
+
+            ggml_backend_t backend_fused = ggml_backend_sched_get_tensor_backend(sched.get(), node.tensor);
+            ggml_backend_dev_t device_fused = backend_fused ? ggml_backend_get_device(backend_fused) : nullptr;
+
+            // TODO: make this descriptor-specific; model.dev_layer() preserves the current behavior,
+            // but is still wrong for cases like --no-kv-offload.
+            ggml_backend_dev_t device_layer = model.dev_layer(node.il);
+
+            if (device_fused != device_layer) {
+                LLAMA_LOG_WARN("%s: layer %d is assigned to device %s but %s "
+                        "is assigned to device %s (usually due to missing support)\n",
+                        func, node.il,
+                        device_layer ? ggml_backend_dev_name(device_layer) : "none",
+                        probe.name,
+                        device_fused ? ggml_backend_dev_name(device_fused) : "none");
+                device_mismatch = true;
+                break;
+            }
+        }
+
+        if (device_mismatch) {
+            enabled = false;
+            LLAMA_LOG_WARN("%s: %s not supported, set to disabled\n", func, probe.name);
+        } else {
+            enabled = true;
+            LLAMA_LOG_INFO("%s: %s enabled\n", func, probe.name);
+        }
+    };
+
+    if (cparams.auto_fa) {
+        resolve(llm_fused_op_flash_attn_probe, cparams.flash_attn);
+        cparams.auto_fa = false;
+    }
+
+    if (cparams.auto_fgdn) {
+        LLAMA_LOG_INFO("%s: resolving fused Gated Delta Net support:\n", func);
+        resolve(llm_fused_op_gdn_ar_probe, cparams.fused_gdn_ar);
+        resolve(llm_fused_op_gdn_ch_probe, cparams.fused_gdn_ch);
+        cparams.auto_fgdn = false;
+    }
+}
+
 void llama_context::sched_reserve() {
    if (!sched_need_reserve) {
        return;
@@ -475,128 +563,7 @@ void llama_context::sched_reserve() {

    LLAMA_LOG_DEBUG("%s: worst-case: n_tokens = %d, n_seqs = %d, n_outputs = %d\n", __func__, n_tokens, n_seqs, n_outputs);

-    // resolve automatic Flash Attention use
-    if (cparams.auto_fa) {
-        auto * gf = graph_reserve(1, n_seqs, n_outputs, mctx.get(), true);
-        if (!gf) {
-            throw std::runtime_error("failed to reserve graph for Flash Attention check");
-        }
-
-        const size_t prefix_len = strlen(LLAMA_TENSOR_NAME_FATTN) + 1;
-        bool fa_device_mismatch = false;
-        for (int i = 0; i < ggml_graph_n_nodes(gf); i++) {
-            ggml_tensor * n = ggml_graph_node(gf, i);
-            if (n->op != GGML_OP_FLASH_ATTN_EXT) {
-                continue;
-            }
-            ggml_backend_dev_t device_fa = ggml_backend_get_device(ggml_backend_sched_get_tensor_backend(sched.get(), n));
-
-            // TODO: instead of the tensor names, use a map to keep track of which (FA) tensors belong to which layer
-            GGML_ASSERT(strncmp(n->name, LLAMA_TENSOR_NAME_FATTN "-", prefix_len) == 0);
-            const int il = std::stoi(n->name + prefix_len);
-            ggml_backend_dev_t device_kv = model.dev_layer(il);
-            if (device_fa != device_kv) {
-                LLAMA_LOG_WARN("%s: layer %d is assigned to device %s but the Flash Attention tensor "
-                        "is assigned to device %s (usually due to missing support)\n",
-                        __func__, il, ggml_backend_dev_name(device_kv), ggml_backend_dev_name(device_fa));
-                // FIXME: fa_device_mismatch logic is wrong for --no-kv-offload, but this is broken anyways
-                fa_device_mismatch = true;
-                break;
-            }
-        }
-
-        if (fa_device_mismatch) {
-            cparams.flash_attn = false;
-            LLAMA_LOG_WARN("%s: Flash Attention was auto, set to disabled\n", __func__);
-        } else {
-            cparams.flash_attn = true;
-            LLAMA_LOG_INFO("%s: Flash Attention was auto, set to enabled\n", __func__);
-        }
-
-        cparams.auto_fa = false;
-    }
-
-    if (cparams.auto_fgdn) {
-        LLAMA_LOG_INFO("%s: resolving fused Gated Delta Net support:\n", __func__);
-
-        if (cparams.fused_gdn_ar) {
-            auto * gf = graph_reserve(1, n_seqs, n_outputs, mctx.get(), true);
-            if (!gf) {
-                throw std::runtime_error("failed to reserve graph for fused Gated Delta Net check (autoregressive)");
-            }
-
-            const size_t prefix_len = strlen(LLAMA_TENSOR_NAME_FGDN_AR) + 1;
-            bool gdn_device_mismatch = false;
-            for (int i = 0; i < ggml_graph_n_nodes(gf); i++) {
-                ggml_tensor * n = ggml_graph_node(gf, i);
-                if (n->op != GGML_OP_GATED_DELTA_NET) {
-                    continue;
-                }
-                ggml_backend_dev_t device_gdn = ggml_backend_get_device(ggml_backend_sched_get_tensor_backend(sched.get(), n));
-
-                GGML_ASSERT(strncmp(n->name, LLAMA_TENSOR_NAME_FGDN_AR "-", prefix_len) == 0);
-                const int il = std::stoi(n->name + prefix_len);
-                ggml_backend_dev_t device_kv = model.dev_layer(il);
-                if (device_gdn != device_kv) {
-                    LLAMA_LOG_WARN("%s: layer %d is assigned to device %s but the fused Gated Delta Net tensor "
-                            "is assigned to device %s (usually due to missing support)\n",
-                            __func__, il, ggml_backend_dev_name(device_kv), ggml_backend_dev_name(device_gdn));
-                    gdn_device_mismatch = true;
-                    break;
-                }
-            }
-
-            if (gdn_device_mismatch) {
-                cparams.fused_gdn_ar = false;
-                LLAMA_LOG_WARN("%s: fused Gated Delta Net (autoregressive) not supported, set to disabled\n", __func__);
-            } else {
-                LLAMA_LOG_INFO("%s: fused Gated Delta Net (autoregressive) enabled\n", __func__);
-            }
-        }
-
-        if (cparams.fused_gdn_ch) {
-            // more than one token in the batch per sequence in order to take the chunked path
-            // note: n_outputs must match n_tokens for embedding models with mean/rank pooling,
-            // because build_pooling creates inp_mean with shape [n_tokens, n_seqs] and multiplies
-            // it with t_embd which is reduced to [n_outputs, ...] via out_ids. if n_outputs != n_tokens,
-            // the ggml_mul_mat assertion fails.
-            const uint32_t n_tokens_ch = 16*n_seqs;
-            auto * gf = graph_reserve(n_tokens_ch, n_seqs, n_tokens_ch, mctx.get(), true);
-            if (!gf) {
-                throw std::runtime_error("failed to reserve graph for fused Gated Delta Net check (chunked)");
-            }
-
-            const size_t prefix_len = strlen(LLAMA_TENSOR_NAME_FGDN_CH) + 1;
-            bool gdn_device_mismatch = false;
-            for (int i = 0; i < ggml_graph_n_nodes(gf); i++) {
-                ggml_tensor * n = ggml_graph_node(gf, i);
-                if (n->op != GGML_OP_GATED_DELTA_NET) {
-                    continue;
-                }
-                ggml_backend_dev_t device_gdn = ggml_backend_get_device(ggml_backend_sched_get_tensor_backend(sched.get(), n));
-
-                GGML_ASSERT(strncmp(n->name, LLAMA_TENSOR_NAME_FGDN_CH "-", prefix_len) == 0);
-                const int il = std::stoi(n->name + prefix_len);
-                ggml_backend_dev_t device_kv = model.dev_layer(il);
-                if (device_gdn != device_kv) {
-                    LLAMA_LOG_WARN("%s: layer %d is assigned to device %s but the fused Gated Delta Net tensor "
-                            "is assigned to device %s (usually due to missing support)\n",
-                            __func__, il, ggml_backend_dev_name(device_kv), ggml_backend_dev_name(device_gdn));
-                    gdn_device_mismatch = true;
-                    break;
-                }
-            }
-
-            if (gdn_device_mismatch) {
-                cparams.fused_gdn_ch = false;
-                LLAMA_LOG_WARN("%s: fused Gated Delta Net (chunked) not supported, set to disabled\n", __func__);
-            } else {
-                LLAMA_LOG_INFO("%s: fused Gated Delta Net (chunked) enabled\n", __func__);
-            }
-        }
-
-        cparams.auto_fgdn = false;
-    }
+    resolve_fused_ops(mctx.get(), n_seqs);

    // reserve worst-case graph
    int n_splits_pp = -1;
@@ -262,6 +262,10 @@ private:

    llm_graph_cb graph_get_cb() const;

+    // disable auto fused ops (Flash Attention, Gated Delta Net) whose op lands on a device
+    // that differs from the layer it belongs to (usually due to missing backend support)
+    void resolve_fused_ops(const llama_memory_context_i * mctx, uint32_t n_seqs);
+
    // TODO: read/write lora adapters and cvec
    size_t state_write_data(llama_io_write_i & io);
    size_t state_read_data (llama_io_read_i  & io);
@@ -1192,6 +1192,7 @@ void llm_graph_result::reset() {
    params = {};

    inputs.clear();
+    fused_nodes.clear();

    buf_compute_meta.resize(ggml_tensor_overhead()*max_nodes + ggml_graph_overhead_custom(max_nodes, false));

@@ -1293,6 +1294,10 @@ llm_graph_input_i * llm_graph_result::add_input(llm_graph_input_ptr input) {
    return inputs.back().get();
 }

+void llm_graph_result::add_fused_node(llm_graph_fused_node result) {
+    fused_nodes.push_back(result);
+}
+
 void llm_graph_result::set_params(const llm_graph_params & params) {
    this->params = params;
 }
@@ -1352,6 +1357,8 @@ void llm_graph_context::cb(ggml_tensor * cur, const char * name, int il) const {
    }
 }

+
+
 ggml_tensor * llm_graph_context::build_cvec(
         ggml_tensor * cur,
                 int   il) const {
@@ -2402,7 +2409,7 @@ ggml_tensor * llm_graph_context::build_attn_mha(

        cur = ggml_flash_attn_ext(ctx0, q, k, v, kq_mask, kq_scale, hparams.f_max_alibi_bias,
                                  hparams.attn_soft_cap ? hparams.f_attn_logit_softcapping : 0.0f);
-        cb(cur, LLAMA_TENSOR_NAME_FATTN, il);
+        res->add_fused_node({LLM_FUSED_OP_FLASH_ATTN, cur, il});

        ggml_flash_attn_ext_add_sinks(cur, sinks);
        ggml_flash_attn_ext_set_prec (cur, GGML_PREC_F32);
@@ -38,6 +38,12 @@ enum llm_graph_type {
    LLM_GRAPH_TYPE_DECODER_MTP,
 };

+enum llm_fused_op {
+    LLM_FUSED_OP_FLASH_ATTN,
+    LLM_FUSED_OP_GDN_AR,
+    LLM_FUSED_OP_GDN_CH,
+};
+
 enum llm_ffn_op_type : int {
    LLM_FFN_NONE = 0,           // sentinel: unset; archs must assign before use
    LLM_FFN_SILU,
@@ -775,6 +781,12 @@ struct llm_graph_params {
    }
 };

+struct llm_graph_fused_node {
+    llm_fused_op op;
+    ggml_tensor * tensor;
+    int il;
+};
+
 class llm_graph_result {
 public:
    llm_graph_result(int64_t max_nodes);
@@ -808,6 +820,10 @@ public:

    llm_graph_input_i * add_input(llm_graph_input_ptr input);

+    void add_fused_node(llm_graph_fused_node result);
+
+    const std::vector<llm_graph_fused_node> & get_fused_nodes() const { return fused_nodes; }
+
    void set_params(const llm_graph_params & params);

    // important graph nodes
@@ -826,6 +842,7 @@ public:
    std::map<llama_seq_id, ggml_tensor *> t_sampled_probs;

    std::vector<llm_graph_input_ptr> inputs;
+    std::vector<llm_graph_fused_node> fused_nodes;

    ggml_context_ptr ctx_compute;

@@ -103,7 +103,3 @@ std::string llama_format_tensor_shape(const std::vector<int64_t> & ne);
 std::string llama_format_tensor_shape(const struct ggml_tensor * t);

 std::string gguf_kv_to_str(const struct gguf_context * ctx_gguf, int i);
-
-#define LLAMA_TENSOR_NAME_FATTN   "__fattn__"
-#define LLAMA_TENSOR_NAME_FGDN_AR "__fgdn_ar__"
-#define LLAMA_TENSOR_NAME_FGDN_CH "__fgdn_ch__"
@@ -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;
@@ -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) {
@@ -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;
@@ -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;
@@ -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) {
@@ -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) {
@@ -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) {
@@ -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));
@@ -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;

@@ -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) {
@@ -401,9 +401,9 @@ std::pair<ggml_tensor *, ggml_tensor *> llm_build_delta_net_base::build_delta_ne
    // K=1: output carries the final state only. state s is 4D [S_v, S_v, H_v, n_seqs].
    ggml_tensor * result = ggml_gated_delta_net(ctx0, q, k, v, g, b, s, /*K=*/1);
    if (n_tokens == 1) {
-        cb(result, LLAMA_TENSOR_NAME_FGDN_AR, il);
+        res->add_fused_node({LLM_FUSED_OP_GDN_AR, result, il});
    } else {
-        cb(result, LLAMA_TENSOR_NAME_FGDN_CH, il);
+        res->add_fused_node({LLM_FUSED_OP_GDN_CH, result, il});
    }

    ggml_tensor * output = ggml_view_4d(ctx0, result,
@@ -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;

@@ -566,9 +566,9 @@ ggml_tensor * llm_build_delta_net_base::build_recurrent_attn(
    // state s is 4D [S_v, S_v, H_v, n_seqs]; K snapshot slots are written into the output.
    ggml_tensor * gdn_out = ggml_gated_delta_net(ctx0, q, k, v, g, b, s, K);
    if (n_seq_tokens > 1) {
-        cb(gdn_out, LLAMA_TENSOR_NAME_FGDN_CH, il);
+        res->add_fused_node({LLM_FUSED_OP_GDN_CH, gdn_out, il});
    } else {
-        cb(gdn_out, LLAMA_TENSOR_NAME_FGDN_AR, il);
+        res->add_fused_node({LLM_FUSED_OP_GDN_AR, gdn_out, il});
    }

    const int64_t attn_score_elems    = S_v * H_v * n_seq_tokens * n_seqs;
@@ -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)) {
@@ -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
@@ -2,11 +2,13 @@

 set(TARGET llama-cli-impl)

-add_library(${TARGET} cli.cpp)
+add_library(${TARGET} cli.cpp
+                      cli-client.cpp
+                      cli-context.cpp)
 set_target_properties(${TARGET} PROPERTIES WINDOWS_EXPORT_ALL_SYMBOLS ON)

 target_include_directories(${TARGET} PUBLIC ${CMAKE_CURRENT_SOURCE_DIR} ../server)
-target_link_libraries(${TARGET} PUBLIC server-context llama-common ${CMAKE_THREAD_LIBS_INIT})
+target_link_libraries(${TARGET} PUBLIC llama-server-impl llama-common ${CMAKE_THREAD_LIBS_INIT})

 if(LLAMA_TOOLS_INSTALL)
    install(TARGETS ${TARGET} LIBRARY)
@@ -0,0 +1,130 @@
+#include "cli-client.h"
+
+#include "http.h"
+
+#include <algorithm>
+#include <chrono>
+#include <thread>
+
+// generation can stall for a long time during prompt processing, so the
+// read timeout must be generous
+static constexpr time_t CLI_HTTP_READ_TIMEOUT_SEC = 3600;
+
+// upper bound for the accumulated response body kept for error reporting
+static constexpr size_t CLI_HTTP_MAX_ERROR_BODY = 1024 * 1024;
+
+// returns the path with the base url's path prefix prepended (if any)
+static std::string join_path(const common_http_url & parts, const std::string & path) {
+    if (parts.path.empty() || parts.path == "/") {
+        return path;
+    }
+    std::string prefix = parts.path;
+    if (prefix.back() == '/') {
+        prefix.pop_back();
+    }
+    return prefix + path;
+}
+
+std::string cli_client::get(const std::string & path) {
+    auto [cli, parts] = common_http_client(server_base);
+    cli.set_read_timeout(CLI_HTTP_READ_TIMEOUT_SEC, 0);
+    auto path_with_model = path + (model.empty() ? "" : ("?model=" + model));
+    auto res = cli.Get(join_path(parts, path_with_model));
+    if (!res) {
+        throw std::runtime_error("failed to connect to " + server_base + ": " + httplib::to_string(res.error()));
+    }
+    if (res->status < 200 || res->status >= 300) {
+        throw std::runtime_error("GET " + path + " failed with status " + std::to_string(res->status) + ": " + res->body);
+    }
+    return res->body;
+}
+
+std::string cli_client::post(const std::string & path, const std::string & body) {
+    auto [cli, parts] = common_http_client(server_base);
+    cli.set_read_timeout(CLI_HTTP_READ_TIMEOUT_SEC, 0);
+    auto res = cli.Post(join_path(parts, path), body, "application/json");
+    if (!res) {
+        throw std::runtime_error("failed to connect to " + server_base + ": " + httplib::to_string(res.error()));
+    }
+    if (res->status < 200 || res->status >= 300) {
+        throw std::runtime_error("POST " + path + " failed with status " + std::to_string(res->status) + ": " + res->body);
+    }
+    return res->body;
+}
+
+std::string cli_client::post_sse(const std::string & path,
+                                  const std::string & body,
+                                  const std::function<bool()> & should_stop,
+                                  const std::function<void(const std::string &)> & on_data) {
+    auto [cli, parts] = common_http_client(server_base);
+    cli.set_read_timeout(CLI_HTTP_READ_TIMEOUT_SEC, 0);
+
+    std::string pending;  // buffer for incomplete SSE lines
+    std::string raw_body; // accumulated body, used only for error reporting
+
+    auto receiver = [&](const char * data, size_t len) -> bool {
+        if (should_stop()) {
+            return false; // aborts the request
+        }
+        if (raw_body.size() < CLI_HTTP_MAX_ERROR_BODY) {
+            raw_body.append(data, std::min(len, CLI_HTTP_MAX_ERROR_BODY - raw_body.size()));
+        }
+        pending.append(data, len);
+        size_t pos;
+        while ((pos = pending.find('\n')) != std::string::npos) {
+            std::string line = pending.substr(0, pos);
+            pending.erase(0, pos + 1);
+            if (!line.empty() && line.back() == '\r') {
+                line.pop_back();
+            }
+            if (line.rfind("data: ", 0) != 0) {
+                continue;
+            }
+            std::string payload = line.substr(6);
+            if (payload == "[DONE]") {
+                continue;
+            }
+            on_data(payload);
+        }
+        return true;
+    };
+
+    httplib::Headers headers = {{"Accept", "text/event-stream"}};
+    auto res = cli.Post(join_path(parts, path), headers, body, "application/json", receiver);
+
+    if (!res) {
+        if (res.error() == httplib::Error::Canceled && should_stop()) {
+            return ""; // cancelled by the user
+        }
+        return "failed to connect to " + server_base + ": " + httplib::to_string(res.error());
+    }
+    if (res->status < 200 || res->status >= 300) {
+        if (!raw_body.empty()) {
+            return raw_body;
+        }
+        return "request failed with status " + std::to_string(res->status);
+    }
+    return "";
+}
+
+bool cli_client::wait_health(const std::function<bool()> & is_aborted) {
+    int connect_attempts = 0;
+    while (!is_aborted()) {
+        auto [cli, parts] = common_http_client(server_base);
+        cli.set_connection_timeout(1, 0);
+        auto res = cli.Get(join_path(parts, "/health"));
+        if (res) {
+            if (res->status == 200) {
+                return true;
+            }
+            // any other status means the server is up but not ready yet
+            // (e.g. 503 while the model is still loading)
+        } else if (++connect_attempts >= 10) {
+            last_error = "failed to connect to " + server_base + ": " + httplib::to_string(res.error());
+            return false;
+        }
+        std::this_thread::sleep_for(std::chrono::milliseconds(300));
+    }
+    last_error = "aborted while waiting for the server to become ready";
+    return false;
+}
@@ -0,0 +1,33 @@
+#pragma once
+
+#include <functional>
+#include <string>
+
+// openai-like client for CLI
+struct cli_client {
+    std::string server_base; // base url, for example "http://127.0.0.1:8080"
+    std::string last_error;  // set when wait_health() fails
+
+    std::string model; // optional, set when the server has multiple models (router mode)
+
+    // simple GET request, returns the raw response body
+    // throws std::runtime_error on transport error or non-2xx status
+    std::string get(const std::string & path);
+
+    // simple POST request, returns the raw response body
+    // throws std::runtime_error on transport error or non-2xx status
+    std::string post(const std::string & path, const std::string & body);
+
+    // POST request with an SSE streaming response
+    // on_data is invoked per "data:" event with the raw event payload
+    // returns after the stream is finished (empty string on graceful exit)
+    // otherwise, the raw error response body
+    std::string post_sse(const std::string & path,
+                          const std::string & body,
+                          const std::function<bool()> & should_stop,
+                          const std::function<void(const std::string &)> & on_data);
+
+    // poll /health until the server is ready to accept requests
+    // returns false if is_aborted returned true or the server is unreachable
+    bool wait_health(const std::function<bool()> & is_aborted);
+};
@@ -0,0 +1,622 @@
+#include "cli-context.h"
+#include "cli-ui.h"
+
+#include "arg.h"
+#include "base64.hpp"
+#include "log.h"
+#include "console.h"
+
+#define JSON_ASSERT GGML_ASSERT
+#include <nlohmann/json.hpp>
+
+#include <algorithm>
+#include <cctype>
+#include <filesystem>
+#include <fstream>
+#include <map>
+#include <set>
+
+using json = nlohmann::ordered_json;
+
+struct cli_context_impl {
+    json messages      = json::array();
+    json pending_media = json::array(); // staged multimodal content parts
+};
+
+cli_context::cli_context(const common_params & params) : params(params), impl(new cli_context_impl()) {}
+
+cli_context::~cli_context() {
+    shutdown();
+}
+
+std::atomic<bool> & cli_context::interrupted() {
+    static std::atomic<bool> flag = false;
+    return flag;
+}
+
+static bool should_stop() {
+    return cli_context::interrupted().load();
+}
+
+static constexpr size_t FILE_GLOB_MAX_RESULTS = 100;
+
+const char * LLAMA_ASCII_LOGO = R"(
+▄▄ ▄▄
+██ ██
+██ ██  ▀▀█▄ ███▄███▄  ▀▀█▄    ▄████ ████▄ ████▄
+██ ██ ▄█▀██ ██ ██ ██ ▄█▀██    ██    ██ ██ ██ ██
+██ ██ ▀█▄██ ██ ██ ██ ▀█▄██ ██ ▀████ ████▀ ████▀
+                                    ██    ██
+                                    ▀▀    ▀▀
+)";
+
+// number of values an arg consumes on the command line
+static int arg_num_values(const common_arg & opt) {
+    if (opt.value_hint_2 != nullptr) {
+        return 2;
+    }
+    if (opt.value_hint != nullptr) {
+        return 1;
+    }
+    return 0;
+}
+
+static std::string format_error_message(const json & err) {
+    if (err.contains("error") && err.at("error").is_object()) {
+        const auto & e = err.at("error");
+        if (e.contains("message") && e.at("message").is_string()) {
+            return e.at("message").get<std::string>();
+        }
+    }
+    return err.dump();
+}
+
+// err is the raw response body of a failed request; it may or may not be JSON
+static std::string format_error_message(const std::string & err) {
+    json parsed = json::parse(err, nullptr, false);
+    if (!parsed.is_discarded()) {
+        return format_error_message(parsed);
+    }
+    return err;
+}
+
+static std::string media_type_from_ext(const std::string & fname) {
+    std::string ext = std::filesystem::path(fname).extension().string();
+    std::transform(ext.begin(), ext.end(), ext.begin(), [](unsigned char c) { return std::tolower(c); });
+    if (ext == ".wav" || ext == ".mp3") {
+        return "audio";
+    }
+    if (ext == ".mp4" || ext == ".avi" || ext == ".mkv" || ext == ".mov" || ext == ".webm") {
+        return "video";
+    }
+    return "image";
+}
+
+bool cli_context::init() {
+    ui::init(params);
+
+    std::optional<ui::spinner> spinner;
+
+    bool use_external_server = !params.server_base.empty();
+    if (use_external_server) {
+        std::string base = params.server_base;
+        while (!base.empty() && base.back() == '/') {
+            base.pop_back();
+        }
+        client.server_base = base;
+
+        spinner.emplace("Connecting to server at " + base);
+    } else {
+        if (params.model.path.empty() && params.model.url.empty() &&
+                params.model.hf_repo.empty() && params.model.docker_repo.empty()) {
+            ui::show_error(
+                "no model specified",
+                "use -m <file.gguf> or -hf <user/repo> to run a local model,\n"
+                "or --server-base <url> to connect to a running llama-server"
+            );
+            return false;
+        }
+
+        spinner.emplace("\n\nLoading model...");
+
+        server.emplace();
+        if (!server->start(params)) {
+            ui::show_error("server start failed");
+            return false;
+        }
+        if (!server->wait_ready(should_stop)) {
+            if (!should_stop()) {
+                ui::show_error("the server exited before becoming ready");
+            }
+            return false;
+        }
+        client.server_base = server->address();
+    }
+
+    // for --server-base this is the main availability check; for a spawned
+    // server it is a cheap sanity check on top of the ready signal
+    auto is_aborted = [this]() {
+        return should_stop() || (server && !server->alive());
+    };
+    bool healthy = false;
+    try {
+        healthy = client.wait_health(is_aborted);
+    } catch (const std::exception & e) {
+        client.last_error = e.what();
+    }
+    if (!healthy) {
+        if (!should_stop()) {
+            ui::show_error(client.last_error);
+        }
+        return false;
+    }
+
+    if (use_external_server) {
+        spinner.reset();
+        if (!list_and_ask_models()) {
+            return false;
+        }
+        // restore the spinner for the next step
+        spinner.emplace("Waiting for server...");
+    }
+
+    fetch_server_props();
+
+    return true;
+}
+
+void cli_context::fetch_server_props() {
+    try {
+        json props = json::parse(client.get("/props"));
+        model_name = props.value("model_alias", "");
+        if (model_name.empty()) {
+            const std::string path = props.value("model_path", "");
+            if (!path.empty()) {
+                model_name = std::filesystem::path(path).filename().string();
+            }
+        }
+        model_ftype = props.value("model_ftype", "");
+        build_info = props.value("build_info", "");
+        if (props.contains("modalities") && props.at("modalities").is_object()) {
+            const auto & modalities = props.at("modalities");
+            has_vision = modalities.value("vision", false);
+            has_audio  = modalities.value("audio", false);
+            has_video  = modalities.value("video", false);
+        }
+    } catch (const std::exception & e) {
+        // /props can be disabled on remote servers; not fatal
+        LOG_DBG("failed to fetch /props: %s\n", e.what());
+    }
+}
+
+bool cli_context::list_and_ask_models() {
+    json resp = json::parse(client.get("/v1/models"));
+    if (!resp.contains("data") || !resp.at("data").is_array()) {
+        throw std::runtime_error("invalid response from /v1/models");
+    }
+    std::vector<std::string> models;
+    std::vector<std::string> models_display;
+    for (const auto & m : resp.at("data")) {
+        if (!m.contains("id") || !m.at("id").is_string()) {
+            continue;
+        }
+        std::string name = m.at("id").get<std::string>();
+        std::string display = name;
+        if (m.contains("aliases") && m.at("aliases").is_array()) {
+            std::vector<std::string> aliases;
+            for (const auto & a : m.at("aliases")) {
+                if (a.is_string()) {
+                    aliases.push_back(a.get<std::string>());
+                }
+            }
+            if (!aliases.empty()) {
+                display += " (" + string_join(aliases, ", ") + ")";
+            }
+        }
+        models.push_back(name);
+        models_display.push_back(display);
+    }
+
+    // only one model: use it without asking
+    if (models.size() == 1) {
+        model_name = models[0];
+        client.model = model_name;
+        return true;
+    }
+
+    std::string message = "\nAvailable models:";
+    for (size_t i = 0; i < models_display.size(); ++i) {
+        message += "\n  " + std::to_string(i + 1) + ". " + models_display[i];
+    }
+    message += "\n";
+    ui::show_message(message);
+    std::string selection;
+    while (selection.empty()) {
+        if (should_stop()) {
+            return false;
+        }
+        ui::user_turn user_turn;
+        selection = user_turn.read_input(false, "Select model by number: ");
+        if (selection.empty()) {
+            continue;
+        }
+        try {
+            size_t idx = std::stoul(selection);
+            if (idx > 0 && idx <= models.size()) {
+                model_name = models[idx - 1];
+                client.model = model_name;
+                ui::show_message("Selected model: " + model_name);
+                break;
+            }
+        } catch (...) {
+            // ignore
+        }
+        ui::show_error("Invalid selection. Please enter a valid number.");
+        selection.clear();
+        continue;
+    }
+    return true;
+}
+
+void cli_context::add_system_prompt() {
+    if (!params.system_prompt.empty()) {
+        impl->messages.push_back({
+            {"role",    "system"},
+            {"content", params.system_prompt}
+        });
+    }
+}
+
+void cli_context::push_user_message(const std::string & text) {
+    json content;
+    if (impl->pending_media.empty()) {
+        content = text;
+    } else {
+        // multimodal message: media parts first, then the text
+        content = impl->pending_media;
+        content.push_back({
+            {"type", "text"},
+            {"text", text}
+        });
+        impl->pending_media = json::array();
+    }
+    impl->messages.push_back({
+        {"role",    "user"},
+        {"content", content}
+    });
+}
+
+bool cli_context::stage_media_file(const std::string & fname, const std::string & type) {
+    std::ifstream file(fname, std::ios::binary);
+    if (!file) {
+        return false;
+    }
+    std::string data((std::istreambuf_iterator<char>(file)), std::istreambuf_iterator<char>());
+    std::string encoded = base64::encode(data);
+
+    if (type == "audio") {
+        std::string ext = std::filesystem::path(fname).extension().string();
+        std::transform(ext.begin(), ext.end(), ext.begin(), [](unsigned char c) { return std::tolower(c); });
+        impl->pending_media.push_back({
+            {"type", "input_audio"},
+            {"input_audio", {
+                {"data",   encoded},
+                {"format", ext == ".mp3" ? "mp3" : "wav"}
+            }}
+        });
+    } else if (type == "video") {
+        impl->pending_media.push_back({
+            {"type", "input_video"},
+            {"input_video", {
+                {"data", encoded}
+            }}
+        });
+    } else {
+        // the server detects the actual image type from the data
+        impl->pending_media.push_back({
+            {"type", "image_url"},
+            {"image_url", {
+                {"url", "data:image/unknown;base64," + encoded}
+            }}
+        });
+    }
+    return true;
+}
+
+bool cli_context::generate_completion(std::string & assistant_content, cli_timings & timings) {
+    json body = {
+        {"messages",          impl->messages},
+        {"stream",            true},
+        // in order to get timings even when we cancel mid-way
+        {"timings_per_token", true},
+    };
+    if (!client.model.empty()) {
+        body["model"] = client.model;
+    }
+
+    bool stream_error = false;
+
+    ui::assistant_turn a;
+
+    std::string err = client.post_sse("/v1/chat/completions", body.dump(), should_stop, [&](const std::string & payload) {
+        json chunk = json::parse(payload, nullptr, false);
+        if (chunk.is_discarded()) {
+            return;
+        }
+        if (chunk.contains("error")) {
+            stream_error = true;
+            ui::show_error(format_error_message(chunk));
+            return;
+        }
+        if (chunk.contains("timings")) {
+            const auto & t = chunk.at("timings");
+            timings.prompt_per_second    = t.value("prompt_per_second",    0.0);
+            timings.predicted_per_second = t.value("predicted_per_second", 0.0);
+        }
+        if (!chunk.contains("choices") || !chunk.at("choices").is_array() || chunk.at("choices").empty()) {
+            return;
+        }
+        const auto & choice = chunk.at("choices").at(0);
+        if (!choice.contains("delta")) {
+            return;
+        }
+        const auto & delta = choice.at("delta");
+        if (delta.contains("reasoning_content") && delta.at("reasoning_content").is_string()) {
+            const std::string text = delta.at("reasoning_content").get<std::string>();
+            if (!text.empty()) {
+                a.push(ui::ASSISTANT_DISPLAY_MODE_REASONING, text);
+            }
+        }
+        if (delta.contains("content") && delta.at("content").is_string()) {
+            const std::string text = delta.at("content").get<std::string>();
+            if (!text.empty()) {
+                assistant_content += text;
+                a.push(ui::ASSISTANT_DISPLAY_MODE_CONTENT, text);
+            }
+        }
+    });
+
+    cli_context::interrupted().store(false);
+
+    if (!err.empty()) {
+        ui::show_error(format_error_message(err));
+        return false;
+    }
+    return !stream_error;
+}
+
+int cli_context::run() {
+    add_system_prompt();
+
+    std::string modalities = "text";
+    if (has_vision) {
+        modalities += ", vision";
+    }
+    if (has_audio) {
+        modalities += ", audio";
+    }
+    if (has_video) {
+        modalities += ", video";
+    }
+
+    std::string banner;
+    banner += "\n";
+    banner += LLAMA_ASCII_LOGO;
+    banner += "\n";
+    banner += "build      : " + build_info + "\n";
+    banner += "model      : " + model_name + "\n";
+    if (!model_ftype.empty()) {
+        banner += "ftype      : " + model_ftype + "\n";
+    }
+    banner += "modalities : " + modalities + "\n";
+    if (!params.system_prompt.empty()) {
+        banner += "using custom system prompt\n";
+    }
+    banner += "\n";
+    banner += "available commands:\n";
+    banner += "  /exit or Ctrl+C     stop or exit\n";
+    banner += "  /regen              regenerate the last response\n";
+    banner += "  /clear              clear the chat history\n";
+    banner += "  /read <file>        add a text file\n";
+    banner += "  /glob <pattern>     add text files using globbing pattern\n";
+    if (has_vision) {
+        banner += "  /image <file>       add an image file\n";
+    }
+    if (has_audio) {
+        banner += "  /audio <file>       add an audio file\n";
+    }
+    if (has_video) {
+        banner += "  /video <file>       add a video file\n";
+    }
+    banner += "\n";
+
+    ui::show_message(banner);
+
+    // interactive loop
+    std::string cur_msg;
+
+    auto add_text_file = [&](const std::string & fname) -> bool {
+        std::ifstream file(fname, std::ios::binary);
+        if (!file) {
+            ui::show_error(string_format("file does not exist or cannot be opened: '%s'", fname.c_str()));
+            return false;
+        }
+        std::string content((std::istreambuf_iterator<char>(file)), std::istreambuf_iterator<char>());
+        cur_msg += "--- File: ";
+        cur_msg += fname;
+        cur_msg += " ---\n";
+        cur_msg += content;
+        ui::show_message(string_format("Loaded text from '%s'", fname.c_str()));
+        return true;
+    };
+
+    while (true) {
+        std::string buffer;
+        {
+            ui::user_turn user_turn;
+
+            if (params.prompt.empty()) {
+                buffer = user_turn.read_input(params.multiline_input);
+            } else {
+                // process input prompt from args
+                for (auto & fname : params.image) {
+                    if (!stage_media_file(fname, media_type_from_ext(fname))) {
+                        ui::show_error(string_format("file does not exist or cannot be opened: '%s'", fname.c_str()));
+                        break;
+                    }
+                    ui::show_message(string_format("Loaded media from '%s'", fname.c_str()));
+                }
+                buffer = params.prompt;
+                user_turn.echo(buffer);
+                params.prompt.clear(); // only use it once
+            }
+        }
+
+        if (should_stop()) {
+            cli_context::interrupted().store(false);
+            break;
+        }
+
+        // remove trailing newline
+        if (!buffer.empty() && buffer.back() == '\n') {
+            buffer.pop_back();
+        }
+
+        // skip empty messages
+        if (buffer.empty()) {
+            continue;
+        }
+
+        bool add_user_msg = true;
+
+        // process commands
+        if (string_starts_with(buffer, "/exit")) {
+            break;
+        } else if (string_starts_with(buffer, "/regen")) {
+            if (impl->messages.size() >= 2) {
+                size_t last_idx = impl->messages.size() - 1;
+                impl->messages.erase(last_idx);
+                add_user_msg = false;
+            } else {
+                ui::show_error("No message to regenerate.");
+                continue;
+            }
+        } else if (string_starts_with(buffer, "/clear")) {
+            impl->messages.clear();
+            add_system_prompt();
+
+            impl->pending_media = json::array();
+            ui::show_message("Chat history cleared.");
+            continue;
+        } else if (
+                (string_starts_with(buffer, "/image ") && has_vision) ||
+                (string_starts_with(buffer, "/audio ") && has_audio) ||
+                (string_starts_with(buffer, "/video ") && has_video)) {
+            std::string type = buffer.substr(1, 5);
+            // just in case (bad copy-paste for example), we strip all trailing/leading spaces
+            std::string fname = string_strip(buffer.substr(7));
+            if (!stage_media_file(fname, type)) {
+                ui::show_error(string_format("file does not exist or cannot be opened: '%s'", fname.c_str()));
+                continue;
+            }
+            ui::show_message(string_format("Loaded media from '%s'", fname.c_str()));
+            continue;
+        } else if (string_starts_with(buffer, "/read ")) {
+            std::string fname = string_strip(buffer.substr(6));
+            add_text_file(fname);
+            continue;
+        } else if (string_starts_with(buffer, "/glob ")) {
+            std::error_code ec;
+            size_t count = 0;
+            auto curdir = std::filesystem::current_path();
+            std::string pattern = string_strip(buffer.substr(6));
+            std::filesystem::path rel_path;
+
+            auto startglob = pattern.find_first_of("![*?");
+            if (startglob != std::string::npos && startglob != 0) {
+                auto endpath = pattern.substr(0, startglob).find_last_of('/');
+                if (endpath != std::string::npos) {
+                    std::string rel_pattern = pattern.substr(0, endpath);
+#if !defined(_WIN32)
+                    if (string_starts_with(rel_pattern, '~')) {
+                        const char * home = std::getenv("HOME");
+                        if (home && home[0]) {
+                            rel_pattern = home + rel_pattern.substr(1);
+                        }
+                    }
+#endif
+                    rel_path = rel_pattern;
+                    pattern.erase(0, endpath + 1);
+                    curdir /= rel_path;
+                }
+            }
+
+            for (const auto & entry : std::filesystem::recursive_directory_iterator(curdir,
+                    std::filesystem::directory_options::skip_permission_denied, ec)) {
+                if (!entry.is_regular_file()) {
+                    continue;
+                }
+
+                std::string rel = std::filesystem::relative(entry.path(), curdir, ec).string();
+                if (ec) {
+                    ec.clear();
+                    continue;
+                }
+                std::replace(rel.begin(), rel.end(), '\\', '/');
+
+                if (!glob_match(pattern, rel)) {
+                    continue;
+                }
+
+                if (!add_text_file((rel_path / rel).string())) {
+                    continue;
+                }
+
+                if (++count >= FILE_GLOB_MAX_RESULTS) {
+                    ui::show_error(string_format("Maximum number of globbed files allowed (%zu) reached.", FILE_GLOB_MAX_RESULTS));
+                    break;
+                }
+            }
+            continue;
+        } else {
+            // not a command
+            cur_msg += buffer;
+        }
+
+        // generate response
+        if (add_user_msg) {
+            push_user_message(cur_msg);
+            cur_msg.clear();
+        }
+        cli_timings timings;
+        std::string assistant_content;
+        generate_completion(assistant_content, timings);
+        impl->messages.push_back({
+            {"role",    "assistant"},
+            {"content", assistant_content}
+        });
+
+        if (params.show_timings) {
+            ui::show_info(string_format(
+                "\n[ Prompt: %.1f t/s | Generation: %.1f t/s ]",
+                timings.prompt_per_second,
+                timings.predicted_per_second
+            ));
+        }
+
+        if (params.single_turn) {
+            break;
+        }
+    }
+
+    ui::show_message("\n\nExiting...");
+
+    return 0;
+}
+
+void cli_context::shutdown() {
+    if (server) {
+        server->stop();
+        server.reset();
+    }
+}
@@ -0,0 +1,66 @@
+#pragma once
+
+#include "common.h"
+
+#include "cli-client.h"
+#include "cli-server.h"
+
+#include <atomic>
+#include <memory>
+#include <optional>
+#include <string>
+
+struct cli_timings {
+    double prompt_per_second    = 0.0;
+    double predicted_per_second = 0.0;
+};
+
+struct cli_context_impl;
+
+struct cli_context {
+    common_params params;
+
+    cli_client client;                // always initialized
+    std::optional<cli_server> server; // only set when no --server-base is given
+
+    // properties of the connected server
+    // will be populated by fetch_server_props()
+    std::string model_name;
+    std::string model_ftype;
+    std::string build_info;
+    bool has_vision = false;
+    bool has_audio  = false;
+    bool has_video  = false;
+
+    cli_context(const common_params & params);
+    ~cli_context();
+
+    // connect to --server-base or spawn a local llama-server child;
+    // argc/argv are needed to forward the server-relevant args to the child
+    bool init();
+
+    // run the interactive chat loop, returns the process exit code
+    int run();
+
+    // stop the local server child (if any)
+    void shutdown();
+
+    // set by the SIGINT handler; cleared once the interrupt has been handled
+    static std::atomic<bool> & interrupted();
+
+private:
+    bool generate_completion(std::string & assistant_content, cli_timings & timings);
+    void fetch_server_props();
+    void add_system_prompt();
+    void push_user_message(const std::string & text);
+
+    // check if server have multiple models (router mode)
+    // if yes, list them then ask; do nothing otherwise
+    bool list_and_ask_models();
+
+    // read a file and stage it as a multimodal content part; type is one of
+    // "image", "audio", "video"; returns false if the file cannot be read
+    bool stage_media_file(const std::string & fname, const std::string & type);
+
+    std::unique_ptr<cli_context_impl> impl;
+};
@@ -0,0 +1,89 @@
+#pragma once
+
+#include <thread>
+
+#include "http.h"
+
+// llama_server will be available as a dynamic library symbol
+int llama_server(common_params & params, int argc, char ** argv);
+void llama_server_terminate();
+
+struct cli_server {
+    std::thread th;
+    int port = -1;
+    std::atomic<bool> is_alive = false;
+    std::atomic<bool> is_stopping = false;
+
+    ~cli_server() {
+        stop();
+    }
+
+    void stop() {
+        if (is_stopping.exchange(true)) {
+            return;
+        }
+        if (alive()) {
+            llama_server_terminate();
+        }
+        if (th.joinable()) {
+            th.join();
+        }
+    }
+
+    // spawn llama-server in a thread and interact with it via a random port
+    bool start(common_params & params) {
+        port = common_http_get_free_port();
+        if (port <= 0) {
+            fprintf(stderr, "failed to get a free port\n");
+            exit(1);
+        }
+
+        is_alive.store(true, std::memory_order_release);
+
+        common_params server_params = params; // copy
+        server_params.port = port;
+
+        th = std::thread([this, server_params]() mutable {
+            // argc / argv are only used in router mode, we can skip them for now
+            int res = llama_server(server_params, 0, nullptr);
+            if (res != 0) {
+                fprintf(stderr, "llama_server exited with code %d\n", res);
+            }
+            is_alive.store(false, std::memory_order_release);
+        });
+
+        return true;
+    }
+
+    std::string address() const {
+        return "http://127.0.0.1:" + std::to_string(port);
+    }
+
+    bool wait_ready(std::function<bool()> should_stop) {
+        if (!alive()) {
+            return false;
+        }
+        while (!should_stop()) {
+            auto [cli, parts] = common_http_client(address());
+            cli.set_connection_timeout(1, 0);
+            auto res = cli.Get("/health");
+            if (res) {
+                if (res->status == 200) {
+                    return true;
+                }
+                // any other status means the server is up but not ready yet
+                // (e.g. 503 while the model is still loading)
+            }
+            if (!alive()) {
+                // in case server die permanently
+                return false;
+            }
+            std::this_thread::sleep_for(std::chrono::milliseconds(200));
+        }
+        return true;
+    }
+
+    bool alive() const {
+        return is_alive.load(std::memory_order_acquire);
+    }
+};
@@ -0,0 +1,251 @@
+#pragma once
+
+#include "common.h"
+#include "console.h"
+
+#include <array>
+#include <algorithm>
+#include <cctype>
+#include <filesystem>
+#include <string_view>
+
+// TODO?: Make this reusable, enums, docs
+static const std::array<std::string_view, 8> cmds = {
+    "/audio ",
+    "/clear",
+    "/exit",
+    "/glob ",
+    "/image ",
+    "/read ",
+    "/regen",
+    "/video ",
+};
+
+static std::vector<std::pair<std::string, size_t>> auto_completion_callback(std::string_view line, size_t cursor_byte_pos) {
+    std::vector<std::pair<std::string, size_t>> matches;
+    std::string cmd;
+
+    if (line.length() > 1 && line.front() == '/' && !std::any_of(cmds.begin(), cmds.end(), [line](std::string_view prefix) {
+        return string_starts_with(line, prefix);
+    })) {
+        auto it = cmds.begin();
+
+        while ((it = std::find_if(it, cmds.end(), [line](std::string_view cmd_line) {
+            return string_starts_with(cmd_line, line);
+        })) != cmds.end()) {
+            matches.emplace_back(*it, it->length());
+            ++it;
+        }
+    } else {
+        auto it = std::find_if(cmds.begin(), cmds.end(), [line](std::string_view prefix) {
+            return prefix.back() == ' ' && string_starts_with(line, prefix);
+        });
+
+        if (it != cmds.end()) {
+            cmd = *it;
+        }
+    }
+
+    if (!cmd.empty() && cmd != "/glob " && line.length() >= cmd.length() && cursor_byte_pos >= cmd.length()) {
+        const std::string path_prefix  = std::string(line.substr(cmd.length(), cursor_byte_pos - cmd.length()));
+        const std::string path_postfix = std::string(line.substr(cursor_byte_pos));
+        auto cur_dir = std::filesystem::current_path();
+        std::string cur_dir_str = cur_dir.string();
+        std::string expanded_prefix = path_prefix;
+
+#if !defined(_WIN32)
+        if (string_starts_with(path_prefix, '~')) {
+            const char * home = std::getenv("HOME");
+            if (home && home[0]) {
+                expanded_prefix = home + path_prefix.substr(1);
+            }
+        }
+        if (string_starts_with(expanded_prefix, '/')) {
+#else
+        if (std::isalpha(static_cast<unsigned char>(expanded_prefix[0])) && expanded_prefix.find(':') == 1) {
+#endif
+            cur_dir = std::filesystem::path(expanded_prefix).parent_path();
+            cur_dir_str.clear();
+        } else if (!path_prefix.empty()) {
+            cur_dir /= std::filesystem::path(path_prefix).parent_path();
+        }
+
+        std::error_code ec;
+        for (const auto & entry : std::filesystem::directory_iterator(cur_dir, ec)) {
+            if (ec) {
+                break;
+            }
+            if (!entry.exists(ec)) {
+                ec.clear();
+                continue;
+            }
+
+            const std::string path_full = entry.path().string();
+            std::string path_entry = !cur_dir_str.empty() && string_starts_with(path_full, cur_dir_str) ? path_full.substr(cur_dir_str.length() + 1) : path_full;
+
+            if (entry.is_directory(ec)) {
+                path_entry.push_back(std::filesystem::path::preferred_separator);
+            }
+
+            if (expanded_prefix.empty() || string_starts_with(path_entry, expanded_prefix)) {
+                const std::string updated_line = cmd + path_entry;
+                matches.emplace_back(updated_line + path_postfix, updated_line.length());
+            }
+
+            if (ec) {
+                ec.clear();
+            }
+        }
+
+        if (matches.empty()) {
+            const std::string updated_line = cmd + path_prefix;
+            matches.emplace_back(updated_line + path_postfix, updated_line.length());
+        }
+
+        // Add the longest common prefix
+        if (!expanded_prefix.empty() && matches.size() > 1) {
+            const std::string_view match0(matches[0].first);
+            const std::string_view match1(matches[1].first);
+            auto it = std::mismatch(match0.begin(), match0.end(), match1.begin(), match1.end());
+            size_t len = it.first - match0.begin();
+
+            for (size_t i = 2; i < matches.size(); ++i) {
+                const std::string_view matchi(matches[i].first);
+                auto cmp = std::mismatch(match0.begin(), match0.end(), matchi.begin(), matchi.end());
+                len = std::min(len, static_cast<size_t>(cmp.first - match0.begin()));
+            }
+
+            const std::string updated_line = std::string(match0.substr(0, len));
+            matches.emplace_back(updated_line + path_postfix, updated_line.length());
+        }
+
+        std::sort(matches.begin(), matches.end(), [](const auto & a, const auto & b) {
+            return a.first.compare(0, a.second, b.first, 0, b.second) < 0;
+        });
+    }
+
+    return matches;
+}
+
+// note: make this view implementation generic, so that we can move to TUI in the future if we want to
+namespace ui {
+    static void init(const common_params & params) {
+        // TODO: avoid using atexit() here by making `console` a singleton
+        console::init(params.simple_io, params.use_color);
+        atexit([]() { console::cleanup(); });
+
+        console::set_completion_callback(auto_completion_callback);
+    }
+
+    struct spinner {
+        spinner(const std::string & message) {
+            if (!message.empty()) {
+                console::log("%s ", message.c_str());
+            }
+            console::spinner::start();
+        }
+        ~spinner() {
+            console::spinner::stop();
+        }
+    };
+
+    struct user_turn {
+        user_turn() {
+            console::set_display(DISPLAY_TYPE_USER_INPUT);
+        }
+        ~user_turn() {
+            console::set_display(DISPLAY_TYPE_RESET);
+        }
+        void echo(const std::string & buffer) {
+            if (buffer.size() > 500) {
+                console::log("\n> %s ... (truncated)\n", buffer.substr(0, 500).c_str());
+            } else {
+                console::log("\n> %s\n", buffer.c_str());
+            }
+        }
+        std::string read_input(bool multiline_input, const char * prompt = nullptr) {
+            if (prompt) {
+                console::log("%s", prompt);
+            } else {
+                console::log("\n> ");
+            }
+            std::string buffer;
+            std::string line;
+            bool another_line = true;
+            do {
+                another_line = console::readline(line, multiline_input);
+                buffer += line;
+            } while (another_line);
+            return buffer;
+        }
+    };
+
+    enum assistant_display_mode {
+        ASSISTANT_DISPLAY_MODE_REASONING,
+        ASSISTANT_DISPLAY_MODE_CONTENT,
+    };
+    struct assistant_turn {
+        assistant_display_mode mode = ASSISTANT_DISPLAY_MODE_CONTENT;
+        bool trailing_newline = true;
+        bool is_inside_reasoning = false;
+        assistant_turn() {
+            console::set_display(DISPLAY_TYPE_RESET);
+        }
+        ~assistant_turn() {
+            console::set_display(DISPLAY_TYPE_RESET);
+            add_newline_if_needed();
+        }
+        void push(assistant_display_mode m, const std::string & buffer) {
+            if (m != mode) {
+                add_newline_if_needed();
+                switch (m) {
+                    case ASSISTANT_DISPLAY_MODE_CONTENT:
+                        {
+                            if (is_inside_reasoning) {
+                                console::log("[End thinking]\n\n");
+                                is_inside_reasoning = false;
+                            }
+                            console::set_display(DISPLAY_TYPE_RESET);
+                        } break;
+                    case ASSISTANT_DISPLAY_MODE_REASONING:
+                        {
+                            console::set_display(DISPLAY_TYPE_REASONING);
+                            is_inside_reasoning = true;
+                            console::log("\n[Start thinking]\n\n");
+                        } break;
+                }
+            }
+            mode = m;
+            if (buffer.empty()) {
+                return;
+            }
+            trailing_newline = buffer.back() == '\n';
+            console::log("%s", buffer.c_str());
+            console::flush();
+        }
+        void add_newline_if_needed() {
+            if (!trailing_newline) {
+                console::log("\n");
+                console::flush();
+            }
+        }
+    };
+
+    static void show_error(const std::string & title, const std::string & message = "") {
+        console::spinner::stop();
+        console::error("Error: %s\n", title.c_str());
+        if (!message.empty()) {
+            console::log("%s\n", message.c_str());
+        }
+    }
+
+    static void show_message(const std::string & message) {
+        console::log("%s\n", message.c_str());
+    }
+
+    static void show_info(const std::string & message) {
+        console::set_display(DISPLAY_TYPE_INFO);
+        console::log("%s\n", message.c_str());
+        console::set_display(DISPLAY_TYPE_RESET);
+    }
+}
@@ -1,20 +1,9 @@
-#include "chat.h"
-#include "common.h"
 #include "arg.h"
-#include "console.h"
-#include "fit.h"
-// #include "log.h"
+#include "common.h"
+#include "log.h"

-#include "server-common.h"
-#include "server-context.h"
-#include "server-task.h"
+#include "cli-context.h"

-#include <array>
-#include <atomic>
-#include <algorithm>
-#include <filesystem>
-#include <fstream>
-#include <thread>
 #include <signal.h>

 #if defined(_WIN32)
@@ -25,342 +14,19 @@
 #include <windows.h>
 #endif

-const char * LLAMA_ASCII_LOGO = R"(
-▄▄ ▄▄
-██ ██
-██ ██  ▀▀█▄ ███▄███▄  ▀▀█▄    ▄████ ████▄ ████▄
-██ ██ ▄█▀██ ██ ██ ██ ▄█▀██    ██    ██ ██ ██ ██
-██ ██ ▀█▄██ ██ ██ ██ ▀█▄██ ██ ▀████ ████▀ ████▀
-                                    ██    ██
-                                    ▀▀    ▀▀
-)";
-
-static std::atomic<bool> g_is_interrupted = false;
-static bool should_stop() {
-    return g_is_interrupted.load();
-}
-
 #if defined (__unix__) || (defined (__APPLE__) && defined (__MACH__)) || defined (_WIN32)
 static void signal_handler(int) {
-    if (g_is_interrupted.load()) {
+    if (cli_context::interrupted().load()) {
        // second Ctrl+C - exit immediately
        // make sure to clear colors before exiting (not using LOG or console.cpp here to avoid deadlock)
        fprintf(stdout, "\033[0m\n");
        fflush(stdout);
        std::exit(130);
    }
-    g_is_interrupted.store(true);
+    cli_context::interrupted().store(true);
 }
 #endif

-struct cli_context {
-    server_context ctx_server;
-    json messages = json::array();
-    std::vector<raw_buffer> input_files;
-    task_params defaults;
-    bool verbose_prompt;
-
-    // thread for showing "loading" animation
-    std::atomic<bool> loading_show;
-
-    cli_context(const common_params & params) {
-        defaults.sampling    = params.sampling;
-        defaults.speculative = params.speculative;
-        defaults.n_keep      = params.n_keep;
-        defaults.n_predict   = params.n_predict;
-        defaults.antiprompt  = params.antiprompt;
-
-        defaults.stream = true; // make sure we always use streaming mode
-        defaults.timings_per_token = true; // in order to get timings even when we cancel mid-way
-        // defaults.return_progress = true; // TODO: show progress
-
-        verbose_prompt = params.verbose_prompt;
-    }
-
-    std::string generate_completion(result_timings & out_timings) {
-        server_response_reader rd = ctx_server.get_response_reader();
-        auto chat_params = format_chat();
-        {
-            // TODO: reduce some copies here in the future
-            server_task task = server_task(SERVER_TASK_TYPE_COMPLETION);
-            task.id         = rd.get_new_id();
-            task.index      = 0;
-            task.params     = defaults;           // copy
-            task.cli_prompt = chat_params.prompt; // copy
-            task.cli_files  = input_files;        // copy
-            task.cli        = true;
-
-            // chat template settings
-            task.params.chat_parser_params = common_chat_parser_params(chat_params);
-            task.params.chat_parser_params.reasoning_format = COMMON_REASONING_FORMAT_DEEPSEEK;
-            if (!chat_params.parser.empty()) {
-                task.params.chat_parser_params.parser.load(chat_params.parser);
-            }
-
-            // Copy the preserved tokens into the sampling params
-            const llama_vocab * vocab = llama_model_get_vocab(
-                llama_get_model(ctx_server.get_llama_context()));
-            for (const auto & token : chat_params.preserved_tokens) {
-                auto ids = common_tokenize(vocab, token, false, true);
-                if (ids.size() == 1) {
-                    task.params.sampling.preserved_tokens.insert(ids[0]);
-                }
-            }
-
-            // reasoning budget sampler
-            if (!chat_params.thinking_end_tag.empty()) {
-                task.params.sampling.reasoning_budget_tokens = defaults.sampling.reasoning_budget_tokens;
-                task.params.sampling.generation_prompt = chat_params.generation_prompt;
-
-                if (!chat_params.thinking_start_tag.empty()) {
-                    task.params.sampling.reasoning_budget_start =
-                        common_tokenize(vocab, chat_params.thinking_start_tag, false, true);
-                }
-                task.params.sampling.reasoning_budget_end =
-                    common_tokenize(vocab, chat_params.thinking_end_tag, false, true);
-                task.params.sampling.reasoning_budget_forced =
-                    common_tokenize(vocab, defaults.sampling.reasoning_budget_message + chat_params.thinking_end_tag, false, true);
-            }
-
-            rd.post_task({std::move(task)});
-        }
-
-        if (verbose_prompt) {
-            console::set_display(DISPLAY_TYPE_PROMPT);
-            console::log("%s\n\n", chat_params.prompt.c_str());
-            console::set_display(DISPLAY_TYPE_RESET);
-        }
-
-        // wait for first result
-        console::spinner::start();
-        server_task_result_ptr result = rd.next(should_stop);
-
-        while (true) {
-            auto res_partial = dynamic_cast<server_task_result_cmpl_partial *>(result.get());
-            if (res_partial && res_partial->is_begin) {
-                // this is the "send 200 status to client" signal in streaming mode
-                // skip, do not stop the spinner
-                result = rd.next(should_stop);
-            } else {
-                console::spinner::stop();
-                break;
-            }
-        }
-
-        std::string curr_content;
-        bool is_thinking = false;
-
-        while (result) {
-            if (should_stop()) {
-                break;
-            }
-            if (result->is_error()) {
-                json err_data = result->to_json();
-                if (err_data.contains("message")) {
-                    console::error("Error: %s\n", err_data["message"].get<std::string>().c_str());
-                } else {
-                    console::error("Error: %s\n", err_data.dump().c_str());
-                }
-                return curr_content;
-            }
-            auto res_partial = dynamic_cast<server_task_result_cmpl_partial *>(result.get());
-            if (res_partial) {
-                out_timings = std::move(res_partial->timings);
-                for (const auto & diff : res_partial->oaicompat_msg_diffs) {
-                    if (!diff.content_delta.empty()) {
-                        if (is_thinking) {
-                            console::log("\n[End thinking]\n\n");
-                            console::set_display(DISPLAY_TYPE_RESET);
-                            is_thinking = false;
-                        }
-                        curr_content += diff.content_delta;
-                        console::log("%s", diff.content_delta.c_str());
-                        console::flush();
-                    }
-                    if (!diff.reasoning_content_delta.empty()) {
-                        console::set_display(DISPLAY_TYPE_REASONING);
-                        if (!is_thinking) {
-                            console::log("[Start thinking]\n");
-                        }
-                        is_thinking = true;
-                        console::log("%s", diff.reasoning_content_delta.c_str());
-                        console::flush();
-                    }
-                }
-            }
-            auto res_final = dynamic_cast<server_task_result_cmpl_final *>(result.get());
-            if (res_final) {
-                out_timings = std::move(res_final->timings);
-                break;
-            }
-            result = rd.next(should_stop);
-        }
-        g_is_interrupted.store(false);
-        // server_response_reader automatically cancels pending tasks upon destruction
-        return curr_content;
-    }
-
-    // TODO: support remote files in the future (http, https, etc)
-    std::string load_input_file(const std::string & fname, bool is_media) {
-        std::ifstream file = fs_open_ifstream(fname, std::ios::binary);
-        if (!file) {
-            return "";
-        }
-        if (is_media) {
-            raw_buffer buf;
-            buf.assign((std::istreambuf_iterator<char>(file)), std::istreambuf_iterator<char>());
-            input_files.push_back(std::move(buf));
-            return get_media_marker();
-        } else {
-            std::string content((std::istreambuf_iterator<char>(file)), std::istreambuf_iterator<char>());
-            return content;
-        }
-    }
-
-    common_chat_params format_chat() {
-        auto meta = ctx_server.get_meta();
-        auto & chat_params = meta.chat_params;
-
-        auto caps = common_chat_templates_get_caps(chat_params.tmpls.get());
-
-        common_chat_templates_inputs inputs;
-        inputs.messages              = common_chat_msgs_parse_oaicompat(messages);
-        inputs.tools                 = {}; // TODO
-        inputs.tool_choice           = COMMON_CHAT_TOOL_CHOICE_NONE;
-        inputs.json_schema           = ""; // TODO
-        inputs.grammar               = ""; // TODO
-        inputs.use_jinja             = chat_params.use_jinja;
-        inputs.parallel_tool_calls   = caps["supports_parallel_tool_calls"];
-        inputs.add_generation_prompt = true;
-        inputs.reasoning_format      = COMMON_REASONING_FORMAT_DEEPSEEK;
-        inputs.force_pure_content    = chat_params.force_pure_content;
-        inputs.enable_thinking       = chat_params.enable_thinking ? common_chat_templates_support_enable_thinking(chat_params.tmpls.get()) : false;
-
-        // Apply chat template to the list of messages
-        return common_chat_templates_apply(chat_params.tmpls.get(), inputs);
-    }
-};
-
-// TODO?: Make this reusable, enums, docs
-static const std::array<std::string_view, 8> cmds = {
-    "/audio ",
-    "/clear",
-    "/exit",
-    "/glob ",
-    "/image ",
-    "/read ",
-    "/regen",
-    "/video ",
-};
-
-static std::vector<std::pair<std::string, size_t>> auto_completion_callback(std::string_view line, size_t cursor_byte_pos) {
-    std::vector<std::pair<std::string, size_t>> matches;
-    std::string cmd;
-
-    if (line.length() > 1 && line.front() == '/' && !std::any_of(cmds.begin(), cmds.end(), [line](std::string_view prefix) {
-        return string_starts_with(line, prefix);
-    })) {
-        auto it = cmds.begin();
-
-        while ((it = std::find_if(it, cmds.end(), [line](std::string_view cmd_line) {
-            return string_starts_with(cmd_line, line);
-        })) != cmds.end()) {
-            matches.emplace_back(*it, it->length());
-            ++it;
-        }
-    } else {
-        auto it = std::find_if(cmds.begin(), cmds.end(), [line](std::string_view prefix) {
-            return prefix.back() == ' ' && string_starts_with(line, prefix);
-        });
-
-        if (it != cmds.end()) {
-            cmd = *it;
-        }
-    }
-
-    if (!cmd.empty() && cmd != "/glob " && line.length() >= cmd.length() && cursor_byte_pos >= cmd.length()) {
-        const std::string path_prefix  = std::string(line.substr(cmd.length(), cursor_byte_pos - cmd.length()));
-        const std::string path_postfix = std::string(line.substr(cursor_byte_pos));
-        auto cur_dir = std::filesystem::current_path();
-        std::string cur_dir_str = cur_dir.string();
-        std::string expanded_prefix = path_prefix;
-
-#if !defined(_WIN32)
-        if (string_starts_with(path_prefix, '~')) {
-            const char * home = std::getenv("HOME");
-            if (home && home[0]) {
-                expanded_prefix = home + path_prefix.substr(1);
-            }
-        }
-        if (string_starts_with(expanded_prefix, '/')) {
-#else
-        if (std::isalpha(expanded_prefix[0]) && expanded_prefix.find(':') == 1) {
-#endif
-            cur_dir = std::filesystem::path(expanded_prefix).parent_path();
-            cur_dir_str.clear();
-        } else if (!path_prefix.empty()) {
-            cur_dir /= std::filesystem::path(path_prefix).parent_path();
-        }
-
-        std::error_code ec;
-        for (const auto & entry : std::filesystem::directory_iterator(cur_dir, ec)) {
-            if (ec) {
-                break;
-            }
-            if (!entry.exists(ec)) {
-                ec.clear();
-                continue;
-            }
-
-            const std::string path_full = entry.path().string();
-            std::string path_entry = !cur_dir_str.empty() && string_starts_with(path_full, cur_dir_str) ? path_full.substr(cur_dir_str.length() + 1) : path_full;
-
-            if (entry.is_directory(ec)) {
-                path_entry.push_back(std::filesystem::path::preferred_separator);
-            }
-
-            if (expanded_prefix.empty() || string_starts_with(path_entry, expanded_prefix)) {
-                const std::string updated_line = cmd + path_entry;
-                matches.emplace_back(updated_line + path_postfix, updated_line.length());
-            }
-
-            if (ec) {
-                ec.clear();
-            }
-        }
-
-        if (matches.empty()) {
-            const std::string updated_line = cmd + path_prefix;
-            matches.emplace_back(updated_line + path_postfix, updated_line.length());
-        }
-
-        // Add the longest common prefix
-        if (!expanded_prefix.empty() && matches.size() > 1) {
-            const std::string_view match0(matches[0].first);
-            const std::string_view match1(matches[1].first);
-            auto it = std::mismatch(match0.begin(), match0.end(), match1.begin(), match1.end());
-            size_t len = it.first - match0.begin();
-
-            for (size_t i = 2; i < matches.size(); ++i) {
-                const std::string_view matchi(matches[i].first);
-                auto cmp = std::mismatch(match0.begin(), match0.end(), matchi.begin(), matchi.end());
-                len = std::min(len, static_cast<size_t>(cmp.first - match0.begin()));
-            }
-
-            const std::string updated_line = std::string(match0.substr(0, len));
-            matches.emplace_back(updated_line + path_postfix, updated_line.length());
-        }
-
-        std::sort(matches.begin(), matches.end(), [](const auto & a, const auto & b) {
-            return a.first.compare(0, a.second, b.first, 0, b.second) < 0;
-        });
-    }
-
-    return matches;
-}
-
-static constexpr size_t FILE_GLOB_MAX_RESULTS = 100;
-
 // satisfies -Wmissing-declarations
 int llama_cli(int argc, char ** argv);

@@ -375,25 +41,6 @@ int llama_cli(int argc, char ** argv) {
        return 1;
    }

-    // TODO: maybe support it later?
-    if (params.conversation_mode == COMMON_CONVERSATION_MODE_DISABLED) {
-        console::error("--no-conversation is not supported by llama-cli\n");
-        console::error("please use llama-completion instead\n");
-    }
-
-    // struct that contains llama context and inference
-    cli_context ctx_cli(params);
-
-    llama_backend_init();
-    llama_numa_init(params.numa);
-
-    // TODO: avoid using atexit() here by making `console` a singleton
-    console::init(params.simple_io, params.use_color);
-    atexit([]() { console::cleanup(); });
-
-    console::set_display(DISPLAY_TYPE_RESET);
-    console::set_completion_callback(auto_completion_callback);
-
 #if defined (__unix__) || (defined (__APPLE__) && defined (__MACH__))
    struct sigaction sigint_action;
    sigint_action.sa_handler = signal_handler;
@@ -408,276 +55,11 @@ int llama_cli(int argc, char ** argv) {
    SetConsoleCtrlHandler(reinterpret_cast<PHANDLER_ROUTINE>(console_ctrl_handler), true);
 #endif

-    console::log("\nLoading model... "); // followed by loading animation
-    console::spinner::start();
-    if (!ctx_cli.ctx_server.load_model(params)) {
-        console::spinner::stop();
-        console::error("\nFailed to load the model\n");
+    cli_context ctx_cli(params);
+
+    if (!ctx_cli.init()) {
        return 1;
    }

-    ctx_cli.defaults.sampling = params.sampling;
-
-    console::spinner::stop();
-    console::log("\n");
-
-    std::thread inference_thread([&ctx_cli]() {
-        ctx_cli.ctx_server.start_loop();
-    });
-
-    auto inf = ctx_cli.ctx_server.get_meta();
-    std::string modalities = "text";
-    if (inf.has_inp_image) {
-        modalities += ", vision";
-    }
-    if (inf.has_inp_audio) {
-        modalities += ", audio";
-    }
-
-    auto add_system_prompt = [&]() {
-        if (!params.system_prompt.empty()) {
-            ctx_cli.messages.push_back({
-                {"role",    "system"},
-                {"content", params.system_prompt}
-            });
-        }
-    };
-    add_system_prompt();
-
-    console::log("\n");
-    console::log("%s\n", LLAMA_ASCII_LOGO);
-    console::log("build      : %s\n", inf.build_info.c_str());
-    console::log("model      : %s\n", inf.model_name.c_str());
-    if (!inf.model_ftype.empty()) {
-        console::log("ftype      : %s\n", inf.model_ftype.c_str());
-    }
-    console::log("modalities : %s\n", modalities.c_str());
-    if (!params.system_prompt.empty()) {
-        console::log("using custom system prompt\n");
-    }
-    console::log("\n");
-    console::log("available commands:\n");
-    console::log("  /exit or Ctrl+C     stop or exit\n");
-    console::log("  /regen              regenerate the last response\n");
-    console::log("  /clear              clear the chat history\n");
-    console::log("  /read <file>        add a text file\n");
-    console::log("  /glob <pattern>     add text files using globbing pattern\n");
-    if (inf.has_inp_image) {
-        console::log("  /image <file>       add an image file\n");
-    }
-    if (inf.has_inp_audio) {
-        console::log("  /audio <file>       add an audio file\n");
-    }
-    if (inf.has_inp_video) {
-        console::log("  /video <file>       add a video file\n");
-    }
-    console::log("\n");
-
-    // interactive loop
-    std::string cur_msg;
-
-    auto add_text_file = [&](const std::string & fname) -> bool {
-        std::string marker = ctx_cli.load_input_file(fname, false);
-        if (marker.empty()) {
-            console::error("file does not exist or cannot be opened: '%s'\n", fname.c_str());
-            return false;
-        }
-        if (inf.fim_sep_token != LLAMA_TOKEN_NULL) {
-            cur_msg += common_token_to_piece(ctx_cli.ctx_server.get_llama_context(), inf.fim_sep_token, true);
-            cur_msg += fname;
-            cur_msg.push_back('\n');
-        } else {
-            cur_msg += "--- File: ";
-            cur_msg += fname;
-            cur_msg += " ---\n";
-        }
-        cur_msg += marker;
-        console::log("Loaded text from '%s'\n", fname.c_str());
-        return true;
-    };
-
-    while (true) {
-        std::string buffer;
-        console::set_display(DISPLAY_TYPE_USER_INPUT);
-        if (params.prompt.empty()) {
-            console::log("\n> ");
-            std::string line;
-            bool another_line = true;
-            do {
-                another_line = console::readline(line, params.multiline_input);
-                buffer += line;
-            } while (another_line);
-        } else {
-            // process input prompt from args
-            for (auto & fname : params.image) {
-                std::string marker = ctx_cli.load_input_file(fname, true);
-                if (marker.empty()) {
-                    console::error("file does not exist or cannot be opened: '%s'\n", fname.c_str());
-                    break;
-                }
-                console::log("Loaded media from '%s'\n", fname.c_str());
-                cur_msg += marker;
-            }
-            buffer = params.prompt;
-            if (buffer.size() > 500) {
-                console::log("\n> %s ... (truncated)\n", buffer.substr(0, 500).c_str());
-            } else {
-                console::log("\n> %s\n", buffer.c_str());
-            }
-            params.prompt.clear(); // only use it once
-        }
-        console::set_display(DISPLAY_TYPE_RESET);
-        console::log("\n");
-
-        if (should_stop()) {
-            g_is_interrupted.store(false);
-            break;
-        }
-
-        // remove trailing newline
-        if (!buffer.empty() &&buffer.back() == '\n') {
-            buffer.pop_back();
-        }
-
-        // skip empty messages
-        if (buffer.empty()) {
-            continue;
-        }
-
-        bool add_user_msg = true;
-
-        // process commands
-        if (string_starts_with(buffer, "/exit")) {
-            break;
-        } else if (string_starts_with(buffer, "/regen")) {
-            if (ctx_cli.messages.size() >= 2) {
-                size_t last_idx = ctx_cli.messages.size() - 1;
-                ctx_cli.messages.erase(last_idx);
-                add_user_msg = false;
-            } else {
-                console::error("No message to regenerate.\n");
-                continue;
-            }
-        } else if (string_starts_with(buffer, "/clear")) {
-            ctx_cli.messages.clear();
-            add_system_prompt();
-
-            ctx_cli.input_files.clear();
-            console::log("Chat history cleared.\n");
-            continue;
-        } else if (
-                (string_starts_with(buffer, "/image ") && inf.has_inp_image) ||
-                (string_starts_with(buffer, "/audio ") && inf.has_inp_audio) ||
-                (string_starts_with(buffer, "/video ") && inf.has_inp_video)) {
-            // just in case (bad copy-paste for example), we strip all trailing/leading spaces
-            std::string fname = string_strip(buffer.substr(7));
-            std::string marker = ctx_cli.load_input_file(fname, true);
-            if (marker.empty()) {
-                console::error("file does not exist or cannot be opened: '%s'\n", fname.c_str());
-                continue;
-            }
-            cur_msg += marker;
-            console::log("Loaded media from '%s'\n", fname.c_str());
-            continue;
-        } else if (string_starts_with(buffer, "/read ")) {
-            std::string fname = string_strip(buffer.substr(6));
-            add_text_file(fname);
-            continue;
-        } else if (string_starts_with(buffer, "/glob ")) {
-            std::error_code ec;
-            size_t count = 0;
-            auto curdir = std::filesystem::current_path();
-            std::string pattern = string_strip(buffer.substr(6));
-            std::filesystem::path rel_path;
-
-            auto startglob = pattern.find_first_of("![*?");
-            if (startglob != std::string::npos && startglob != 0) {
-                auto endpath = pattern.substr(0, startglob).find_last_of('/');
-                if (endpath != std::string::npos) {
-                    std::string rel_pattern = pattern.substr(0, endpath);
-#if !defined(_WIN32)
-                    if (string_starts_with(rel_pattern, '~')) {
-                        const char * home = std::getenv("HOME");
-                        if (home && home[0]) {
-                            rel_pattern = home + rel_pattern.substr(1);
-                        }
-                    }
-#endif
-                    rel_path = rel_pattern;
-                    pattern.erase(0, endpath + 1);
-                    curdir /= rel_path;
-                }
-            }
-
-            for (const auto & entry : std::filesystem::recursive_directory_iterator(curdir,
-                    std::filesystem::directory_options::skip_permission_denied, ec)) {
-                if (!entry.is_regular_file()) {
-                    continue;
-                }
-
-                std::string rel = std::filesystem::relative(entry.path(), curdir, ec).string();
-                if (ec) {
-                    ec.clear();
-                    continue;
-                }
-                std::replace(rel.begin(), rel.end(), '\\', '/');
-
-                if (!glob_match(pattern, rel)) {
-                    continue;
-                }
-
-                if (!add_text_file((rel_path / rel).string())) {
-                    continue;
-                }
-
-                if (++count >= FILE_GLOB_MAX_RESULTS) {
-                    console::error("Maximum number of globbed files allowed (%zu) reached.\n", FILE_GLOB_MAX_RESULTS);
-                    break;
-                }
-            }
-            continue;
-        } else {
-            // not a command
-            cur_msg += buffer;
-        }
-
-        // generate response
-        if (add_user_msg) {
-            ctx_cli.messages.push_back({
-                {"role",    "user"},
-                {"content", cur_msg}
-            });
-            cur_msg.clear();
-        }
-        result_timings timings;
-        std::string assistant_content = ctx_cli.generate_completion(timings);
-        ctx_cli.messages.push_back({
-            {"role",    "assistant"},
-            {"content", assistant_content}
-        });
-        console::log("\n");
-
-        if (params.show_timings) {
-            console::set_display(DISPLAY_TYPE_INFO);
-            console::log("\n");
-            console::log("[ Prompt: %.1f t/s | Generation: %.1f t/s ]\n", timings.prompt_per_second, timings.predicted_per_second);
-            console::set_display(DISPLAY_TYPE_RESET);
-        }
-
-        if (params.single_turn) {
-            break;
-        }
-    }
-
-    console::set_display(DISPLAY_TYPE_RESET);
-
-    console::log("\nExiting...\n");
-    ctx_cli.ctx_server.terminate();
-    inference_thread.join();
-
-    // bump the log level to display timings
-    common_log_set_verbosity_thold(LOG_LEVEL_INFO);
-    common_memory_breakdown_print(ctx_cli.ctx_server.get_llama_context());
-
-    return 0;
+    return ctx_cli.run();
 }
@@ -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",  },
@@ -57,7 +57,7 @@ The core architecture consists of the following components:
 - `server_tokens`: Unified representation of token sequences (supports both text and multimodal tokens); used by `server_task` and `server_slot`.
 - `server_prompt_checkpoint`: For recurrent (e.g., RWKV) and SWA models, stores snapshots of KV cache state. Enables reuse when subsequent requests share the same prompt prefix, saving redundant computation.
 - `server_models`: Standalone component for managing multiple backend instances (used in router mode). It is completely independent of `server_context`.
- `stream_session_manager`: Process wide owner of resumable SSE stream sessions (`g_stream_sessions`), keyed by conversation id. Backs the replay buffer that lets a client reattach to a generation after an HTTP disconnect. See the "Resumable streaming" section below.
+- `stream_session_manager`: process wide owner of resumable SSE stream sessions, keyed by conversation id. A file-static singleton inside `server-stream.cpp`, driven through `server_stream_session_manager_start/stop`. Backs the replay buffer that lets a client reattach to a generation after an HTTP disconnect. See the "Resumable streaming" section below.

 ```mermaid
 graph TD
@@ -127,10 +127,12 @@ It is opt in via the `X-Conversation-Id` header on `POST /v1/chat/completions`.
 The feature lives entirely in `server-stream.{h,cpp}` and rests on three types:

 - `stream_session`: a bounded ring buffer (4 MiB cap, oldest bytes drop first) plus a condvar. `append` pushes raw SSE bytes, `read_from` drains from any offset and blocks for live bytes or finalize, `finalize` wakes readers, `cancel` stops the producer. One conv maps to at most one live session.
- `stream_session_manager` (`g_stream_sessions`): owns all sessions keyed by conv id, enforces the one conv one session invariant via `create_or_replace`, and runs a GC thread that drops completed sessions past their TTL.
+- `stream_session_manager`: a file-static singleton (`g_stream_sessions`) inside `server-stream.cpp`, owns all sessions keyed by conv id, enforces the one conv one session invariant via `create_or_replace`, and runs a GC thread that drops completed sessions past their TTL. Exposed to main only through `server_stream_session_manager_start/stop`.
 - `stream_pipe_producer` / `stream_pipe_consumer`: the write and read ends. The producer owns the session lifetime and finalizes it on destruction; the consumer is read only and never finalizes, so a reader detaching cannot kill a running generation.

-Producer side: `server_res_generator` attaches a producer pipe when the header is present. The HTTP content provider mirrors every chunk into the ring before writing it to the socket. While a pipe is attached, `stream_aware_should_stop` ignores peer disconnect, so a dropped socket does not stop generation: only an explicit `DELETE` does. When the peer leaves early, `on_complete` calls `close()`, which drains the rest of the generation into the ring on the http worker.
+The implementation is hidden in `server-stream.cpp` (pimpl). The header exposes only the route handler factories, `server_stream_session_attach_pipe`, `server_stream_aware_should_stop`, `server_stream_conv_id_from_headers` and the GC lifecycle; the session, manager and consumer types stay in the `.cpp`.
+
+Producer side: `server_res_generator` attaches a producer pipe when the header is present. The HTTP content provider mirrors every chunk into the ring before writing it to the socket. While a pipe is attached, `server_stream_aware_should_stop` ignores peer disconnect, so a dropped socket does not stop generation: only an explicit `DELETE` does. When the peer leaves early, `on_complete` calls `close()`, which drains the rest of the generation into the ring on the http worker.

 Lifetime safety: the producer pipe holds a shared `alive` flag also captured by the session cancel hook. `~server_res_generator` calls `cleanup()` to clear that hook while the reader is still alive, so a `cancel` arriving during teardown can never call `stop()` on a freed response. This ordering is the most fragile part of the feature: finalizing or destroying the producer before `cleanup()` runs reintroduces a use after free.

@@ -144,7 +146,7 @@ Routes:

 Router mode binds the same paths to proxy handlers. A `conv_id -> child` map (`conv_models`), populated when a POST is routed, resolves the owning child in one lookup with no polling. The lookup groups ids per child; GET and DELETE proxy straight to the owner. This loopback REST hop is expected to move to a websocket IPC later, swapping only the transport.

-Lifecycle: `g_stream_sessions.start_gc()` runs in main after common init, `stop_gc()` runs first in `clean_up()` and finalizes every live session so no reader hangs. Reader blocking and the post drop drain both run on httplib worker threads, which block on a condvar rather than spin.
+Lifecycle: `server_stream_session_manager_start()` runs in main after common init, `server_stream_session_manager_stop()` runs first in `clean_up()` and finalizes every live session so no reader hangs. Reader blocking and the post drop drain both run on httplib worker threads, which block on a condvar rather than spin.

 | Constant | Value | Role |
 | --- | --- | --- |
@@ -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) |
@@ -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
@@ -4243,7 +4188,7 @@ std::unique_ptr<server_res_generator> server_routes::handle_completions_impl(
                }
            };

-            auto effective_should_stop = stream_aware_should_stop(res_this, req.should_stop);
+            auto effective_should_stop = server_stream_aware_should_stop(res_this, req.should_stop);

            try {
                if (effective_should_stop()) {
@@ -4339,7 +4284,7 @@ std::unique_ptr<server_res_generator> server_routes::handle_completions_impl(

    // attach a producer pipe to the response when X-Conversation-Id is present.
    // the pipe mirrors SSE chunks into the ring buffer and wires up the cancel hook.
-    stream_session_attach_pipe(*res, req.headers);
+    server_stream_session_attach_pipe(*res, req.headers);

    return res;
 }
@@ -4576,6 +4521,7 @@ void server_routes::init_routes() {
            { "default_generation_settings", default_generation_settings_for_props },
            { "total_slots",                 params.n_parallel },
            { "model_alias",                 meta->model_name },
+            { "model_ftype",                 meta->model_ftype },
            { "model_path",                  meta->model_path },
            { "modalities",                  json {
                {"vision", meta->has_inp_image},
@@ -7,6 +7,7 @@
 #include "build-info.h"
 #include "preset.h"
 #include "download.h"
+#include "http.h"

 #include <cpp-httplib/httplib.h> // TODO: remove this once we use HTTP client from download.h
 #include <optional>
@@ -28,14 +29,7 @@
 #include <sstream>
 #include <cstring>

-#ifdef _WIN32
-#include <winsock2.h>
-#include <windows.h>
-#else
-#include <sys/socket.h>
-#include <netinet/in.h>
-#include <arpa/inet.h>
-#include <unistd.h>
+#ifndef _WIN32
 extern char **environ;
 #endif

@@ -716,66 +710,6 @@ std::optional<server_model_meta> server_models::get_meta(const std::string & nam
    return std::nullopt;
 }

-static int get_free_port() {
-#ifdef _WIN32
-    WSADATA wsaData;
-    if (WSAStartup(MAKEWORD(2, 2), &wsaData) != 0) {
-        return -1;
-    }
-    typedef SOCKET native_socket_t;
-#define INVALID_SOCKET_VAL INVALID_SOCKET
-#define CLOSE_SOCKET(s) closesocket(s)
-#else
-    typedef int native_socket_t;
-#define INVALID_SOCKET_VAL -1
-#define CLOSE_SOCKET(s) close(s)
-#endif
-
-    native_socket_t sock = socket(AF_INET, SOCK_STREAM, 0);
-    if (sock == INVALID_SOCKET_VAL) {
-#ifdef _WIN32
-        WSACleanup();
-#endif
-        return -1;
-    }
-
-    struct sockaddr_in serv_addr;
-    std::memset(&serv_addr, 0, sizeof(serv_addr));
-    serv_addr.sin_family = AF_INET;
-    serv_addr.sin_addr.s_addr = htonl(INADDR_ANY);
-    serv_addr.sin_port = htons(0);
-
-    if (bind(sock, (struct sockaddr*)&serv_addr, sizeof(serv_addr)) != 0) {
-        CLOSE_SOCKET(sock);
-#ifdef _WIN32
-        WSACleanup();
-#endif
-        return -1;
-    }
-
-#ifdef _WIN32
-    int namelen = sizeof(serv_addr);
-#else
-    socklen_t namelen = sizeof(serv_addr);
-#endif
-    if (getsockname(sock, (struct sockaddr*)&serv_addr, &namelen) != 0) {
-        CLOSE_SOCKET(sock);
-#ifdef _WIN32
-        WSACleanup();
-#endif
-        return -1;
-    }
-
-    int port = ntohs(serv_addr.sin_port);
-
-    CLOSE_SOCKET(sock);
-#ifdef _WIN32
-    WSACleanup();
-#endif
-
-    return port;
-}
-
 // helper to convert vector<string> to char **
 // pointers are only valid as long as the original vector is valid
 static std::vector<char *> to_char_ptr_array(const std::vector<std::string> & vec) {
@@ -879,7 +813,7 @@ void server_models::load(const std::string & name, const load_options & opts) {
    // prepare new instance info
    instance_t inst;
    inst.meta             = meta;
-    inst.meta.port        = get_free_port();
+    inst.meta.port        = common_http_get_free_port();
    inst.meta.status      = SERVER_MODEL_STATUS_LOADING;
    inst.meta.loaded_info = json{};
    inst.meta.last_used   = ggml_time_ms();
@@ -1681,7 +1615,7 @@ void server_models_routes::init_routes() {
        }
        // remember which child serves this conversation so the stream routes can route straight
        // to it without polling, keyed on the exact conv id from the header
-        std::string conv_id = stream_conv_id_from_headers(req.headers);
+        std::string conv_id = server_stream_conv_id_from_headers(req.headers);
        if (!conv_id.empty()) {
            models.conv_models.remember(conv_id, name);
        }
@@ -1896,7 +1830,7 @@ void server_models_routes::init_routes() {
        if (!from.empty()) {
            child_path += "?from=" + from;
        }
-        SRV_INF("proxying stream resume to model %s on port %d, path=%s\n",
+        SRV_TRC("proxying stream resume to model %s on port %d, path=%s\n",
                owner->name.c_str(), owner->port, child_path.c_str());
        auto proxy = std::make_unique<server_http_proxy>(
                "GET",
@@ -6,6 +6,12 @@
 #include <chrono>
 #include <memory>
 #include <utility>
+#include <shared_mutex>
+
+enum class stream_read_status {
+    OK,
+    OFFSET_LOST,
+};

 namespace {
 constexpr int64_t STREAM_SESSION_TTL_SECONDS         = 300;
@@ -13,7 +19,6 @@ constexpr size_t  STREAM_SESSION_MAX_BYTES           = 4 * 1024 * 1024;
 constexpr int64_t STREAM_SESSION_GC_INTERVAL_SECONDS = 60;
 constexpr int64_t STREAM_READ_WAKE_INTERVAL_MS       = 200;

-// returns unix time in seconds
 int64_t now_seconds() {
    return std::chrono::duration_cast<std::chrono::seconds>(
        std::chrono::system_clock::now().time_since_epoch()
@@ -21,6 +26,91 @@ int64_t now_seconds() {
 }
 }

+// owns all live sessions keyed by conversation_id, one conv = at most one live session.
+// a periodic GC evicts expired ones
+class stream_session_manager {
+public:
+    stream_session_manager();
+    ~stream_session_manager();
+
+    stream_session_manager(const stream_session_manager &)             = delete;
+    stream_session_manager & operator=(const stream_session_manager &) = delete;
+
+    // install a new session, evicting and cancelling any previous one. conversation_id must be non empty
+    stream_session_ptr create_or_replace(const std::string & conversation_id);
+
+    stream_session_ptr get(const std::string & conversation_id);
+
+    std::vector<stream_session_ptr> list_all() const;
+
+    void evict(const std::string & conversation_id);
+
+    void evict_and_cancel(const std::string & conversation_id);
+
+    void start_gc();
+    void stop_gc();
+
+private:
+    void gc_loop();
+
+    mutable std::shared_mutex                           map_mu;
+    std::unordered_map<std::string, stream_session_ptr> sessions; // key: conversation_id
+    std::thread                                         gc_thread;
+    bool                                                running;
+    std::mutex                                          gc_wake_mu;
+    std::condition_variable                             gc_wake_cv;
+};
+
+// process wide manager, lifecycle controlled by llama-server main() via start_gc/stop_gc
+static stream_session_manager g_stream_sessions;
+
+void server_stream_session_manager_start() {
+    g_stream_sessions.start_gc();
+}
+
+void server_stream_session_manager_stop() {
+    g_stream_sessions.stop_gc();
+}
+
+struct stream_session {
+    std::string conversation_id;
+    int64_t     started_ts; // unix seconds at construction
+
+    stream_session(std::string conversation_id_, size_t max_bytes_);
+    stream_session(const stream_session &)             = delete;
+    stream_session & operator=(const stream_session &) = delete;
+
+    bool append(const char * data, size_t len);
+
+    void finalize();
+
+    // drain from offset into sink, blocking for more bytes or finalize. OFFSET_LOST if offset
+    // fell below the dropped prefix
+    stream_read_status read_from(size_t offset,
+        const std::function<bool(const char *, size_t)> & sink,
+        const std::function<bool()> & should_stop);
+
+    bool    is_done() const;
+    bool    is_cancelled() const;
+    size_t  total_size() const;     // bytes that ever entered the session
+    size_t  dropped_prefix() const; // bytes evicted from the front due to cap
+    int64_t completed_at() const;   // 0 while alive, unix seconds after finalize
+
+    void set_stop_producer(std::function<void()> fn);
+
+    void cancel();
+
+private:
+    mutable std::mutex      mu;
+    std::condition_variable cv;
+    std::vector<char>       buffer;
+    size_t                  prefix_dropped;
+    size_t                  cap_bytes;
+    bool                    done;
+    std::atomic<bool>       cancelled; // polled lock-free by the should_stop closure, no mu
+    int64_t                 completed_ts;
+    std::function<void()>   stop_producer;
+};
 stream_session::stream_session(std::string conversation_id_, size_t max_bytes_)
    : conversation_id(std::move(conversation_id_))
    , started_ts(now_seconds())
@@ -38,7 +128,7 @@ bool stream_session::append(const char * data, size_t len) {
    }
    {
        std::lock_guard<std::mutex> lock(mu);
-        if (done.load(std::memory_order_relaxed)) {
+        if (done) {
            return false;
        }
        if (len >= cap_bytes) {
@@ -62,11 +152,14 @@ bool stream_session::append(const char * data, size_t len) {
 }

 void stream_session::finalize() {
-    bool was_done = done.exchange(true, std::memory_order_acq_rel);
-    if (was_done) {
-        return;
+    {
+        std::lock_guard<std::mutex> lock(mu);
+        if (done) {
+            return;
+        }
+        done         = true;
+        completed_ts = now_seconds();
    }
-    completed_ts.store(now_seconds(), std::memory_order_release);
    cv.notify_all();
 }

@@ -96,7 +189,7 @@ stream_read_status stream_session::read_from(size_t offset,
            lock.lock();
            continue;
        }
-        if (done.load(std::memory_order_acquire)) {
+        if (done) {
            return stream_read_status::OK;
        }
        // wait for new bytes, finalize, or a periodic wake to re check should_stop
@@ -105,7 +198,8 @@ stream_read_status stream_session::read_from(size_t offset,
 }

 bool stream_session::is_done() const {
-    return done.load(std::memory_order_acquire);
+    std::lock_guard<std::mutex> lock(mu);
+    return done;
 }

 size_t stream_session::total_size() const {
@@ -119,7 +213,8 @@ size_t stream_session::dropped_prefix() const {
 }

 int64_t stream_session::completed_at() const {
-    return completed_ts.load(std::memory_order_acquire);
+    std::lock_guard<std::mutex> lock(mu);
+    return completed_ts;
 }

 void stream_session::set_stop_producer(std::function<void()> fn) {
@@ -128,7 +223,7 @@ void stream_session::set_stop_producer(std::function<void()> fn) {
 }

 void stream_session::cancel() {
-    // flip cancelled first so the producer-side stream_aware_should_stop can break out of the
+    // flip cancelled first so the producer-side server_stream_aware_should_stop can break out of the
    // recv() wait even if remove_waiting_task_ids does not notify the condvar (the cancel task
    // posted by rd.stop() will eventually notify, but we do not want to depend on that timing)
    cancelled.store(true, std::memory_order_release);
@@ -237,18 +332,24 @@ void stream_session_manager::evict_and_cancel(const std::string & conversation_i
 }

 void stream_session_manager::start_gc() {
-    if (running.exchange(true)) {
-        return;
+    {
+        std::lock_guard<std::mutex> lock(gc_wake_mu);
+        if (running) {
+            return;
+        }
+        running = true;
    }
    gc_thread = std::thread([this] { gc_loop(); });
 }

 void stream_session_manager::stop_gc() {
-    bool was_running = running.exchange(false);
+    bool was_running;
+    {
+        std::lock_guard<std::mutex> lock(gc_wake_mu);
+        was_running = running;
+        running = false;
+    }
    if (was_running) {
-        {
-            std::lock_guard<std::mutex> lock(gc_wake_mu);
-        }
        gc_wake_cv.notify_all();
        if (gc_thread.joinable()) {
            gc_thread.join();
@@ -270,15 +371,15 @@ void stream_session_manager::stop_gc() {
 }

 void stream_session_manager::gc_loop() {
-    while (running.load(std::memory_order_acquire)) {
+    while (true) {
        {
            std::unique_lock<std::mutex> lock(gc_wake_mu);
            gc_wake_cv.wait_for(lock,
                std::chrono::seconds(STREAM_SESSION_GC_INTERVAL_SECONDS),
-                [this] { return !running.load(std::memory_order_acquire); });
-        }
-        if (!running.load(std::memory_order_acquire)) {
-            return;
+                [this] { return !running; });
+            if (!running) {
+                return;
+            }
        }
        int64_t cutoff = now_seconds() - STREAM_SESSION_TTL_SECONDS;
        std::vector<stream_session_ptr> to_drop;
@@ -301,10 +402,19 @@ void stream_session_manager::gc_loop() {
    }
 }

-// process wide manager, lifecycle controlled by llama-server main() via start_gc/stop_gc
-stream_session_manager g_stream_sessions;
+// stream_pipe

-// stream_pipe ---------------------------------------------------------------------------------
+// consumer end: read-only replay of the ring buffer, the destructor does not finalize the session
+struct stream_pipe_consumer : stream_pipe {
+    stream_read_status read(size_t & offset,
+        const std::function<bool(const char *, size_t)> & sink,
+        const std::function<bool()> & should_stop);
+
+    static std::shared_ptr<stream_pipe_consumer> create(stream_session_ptr session);
+
+private:
+    explicit stream_pipe_consumer(stream_session_ptr session);
+};

 stream_pipe::stream_pipe(stream_session_ptr session)
    : session_(std::move(session)) {
@@ -408,12 +518,10 @@ static server_http_res_ptr make_error_response(int status, const std::string & m
    return res;
 }

-server_http_context::handler_t make_stream_get_handler() {
+server_http_context::handler_t server_stream_make_get_handler() {
    return [](const server_http_req & req) -> server_http_res_ptr {
-        // GET /v1/stream/<conv_id>?from=N replays the SSE bytes already buffered for the
-        // session, blocks for more bytes when the session is still running, returns when
-        // the session is finalized. the body is streamed back as text/event-stream so the
-        // browser EventSource can attach to it like a fresh request
+        // GET /v1/stream/<conv_id>?from=N replays buffered SSE bytes then blocks for live
+        // bytes until the session finalizes, streamed as text/event-stream for EventSource
        std::string conv_id = req.get_param("conv_id");
        if (conv_id.empty()) {
            return make_error_response(400, "Missing conversation id in path", ERROR_TYPE_INVALID_REQUEST);
@@ -459,11 +567,10 @@ server_http_context::handler_t make_stream_get_handler() {
    };
 }

-server_http_context::handler_t make_streams_lookup_handler() {
+server_http_context::handler_t server_stream_make_lookup_handler() {
    return [](const server_http_req & req) -> server_http_res_ptr {
-        // POST /v1/streams/lookup with body {"conversation_ids": ["X", "Y", ...]} returns the
-        // matching sessions, only for ids the caller already knows. each id matches the exact key
-        // and any "<id>::<model>" variant, so one lookup covers every per model session for a conv
+        // POST /v1/streams/lookup returns the matching sessions, only for ids the caller already
+        // knows. each id matches the exact key and any "<id>::<model>" per model variant
        std::vector<std::string> requested;
        try {
            json body = json::parse(req.body);
@@ -518,11 +625,10 @@ server_http_context::handler_t make_streams_lookup_handler() {
    };
 }

-server_http_context::handler_t make_stream_delete_handler() {
+server_http_context::handler_t server_stream_make_delete_handler() {
    return [](const server_http_req & req) -> server_http_res_ptr {
-        // DELETE /v1/stream/<conv_id> is the explicit user Stop, cancels the producer hook
-        // wired by handle_completions_impl and evicts the buffer. idempotent, a session that
-        // already finalized or was never created returns 204 either way
+        // DELETE /v1/stream/<conv_id> is the explicit user Stop, cancels the producer and evicts
+        // the buffer. idempotent, returns 204 even if the session was already gone
        std::string conv_id = req.get_param("conv_id");
        if (conv_id.empty()) {
            return make_error_response(400, "Missing conversation id in path", ERROR_TYPE_INVALID_REQUEST);
@@ -536,7 +642,7 @@ server_http_context::handler_t make_stream_delete_handler() {
    };
 }

-std::string stream_conv_id_from_headers(const std::map<std::string, std::string> & headers) {
+std::string server_stream_conv_id_from_headers(const std::map<std::string, std::string> & headers) {
    // case-insensitive scan for x-conversation-id
    static constexpr char   target[]   = "x-conversation-id";
    static constexpr size_t target_len = sizeof(target) - 1;
@@ -555,8 +661,8 @@ std::string stream_conv_id_from_headers(const std::map<std::string, std::string>
    return std::string();
 }

-void stream_session_attach_pipe(server_http_res & res, const std::map<std::string, std::string> & headers) {
-    std::string conversation_id = stream_conv_id_from_headers(headers);
+void server_stream_session_attach_pipe(server_http_res & res, const std::map<std::string, std::string> & headers) {
+    std::string conversation_id = server_stream_conv_id_from_headers(headers);
    SRV_TRC("conv_id=%s (empty=%d)\n", conversation_id.c_str(), conversation_id.empty() ? 1 : 0);
    if (conversation_id.empty()) {
        return;
@@ -565,7 +671,7 @@ void stream_session_attach_pipe(server_http_res & res, const std::map<std::strin
    res.spipe = stream_pipe_producer::create(session, res);
 }

-std::function<bool()> stream_aware_should_stop(server_http_res * res, std::function<bool()> fallback) {
+std::function<bool()> server_stream_aware_should_stop(server_http_res * res, std::function<bool()> fallback) {
    return [res, fallback = std::move(fallback)]() -> bool {
        if (res->spipe) {
            return res->spipe->is_cancelled();
@@ -3,81 +3,23 @@
 #include "server-http.h"

 #include <atomic>
-#include <condition_variable>
 #include <cstddef>
-#include <cstdint>
 #include <functional>
 #include <memory>
-#include <mutex>
-#include <shared_mutex>
 #include <string>
-#include <thread>
-#include <unordered_map>
-#include <vector>

-enum class stream_read_status {
-    OK,
-    OFFSET_LOST,
-};
+// streaming buffer for one generation, survives HTTP disconnect. the producer appends SSE bytes,
+// readers drain from any offset via read_from. keyed by conversation_id, one conv = one live session

-// streaming buffer for one generation, survives HTTP disconnect. the producer appends raw SSE
-// bytes, readers drain from any offset via read_from and block until more bytes or finalize.
-// keyed by conversation_id: one conv = at most one live session
-struct stream_session {
-    std::string conversation_id;
-    int64_t     started_ts; // unix seconds at construction, used by /v1/streams listing
-
-    stream_session(std::string conversation_id_, size_t max_bytes_);
-    stream_session(const stream_session &)             = delete;
-    stream_session & operator=(const stream_session &) = delete;
-
-    // append raw bytes, drops from the front if the cap is reached.
-    // returns false if the session is already finalized
-    bool append(const char * data, size_t len);
-
-    // mark the session as complete, wakes all pending readers
-    void finalize();
-
-    // drain bytes from offset, calling sink for each chunk. blocks until more
-    // bytes arrive or finalize is called. returns OK on clean exit, OFFSET_LOST
-    // if offset falls below the dropped prefix
-    stream_read_status read_from(size_t offset,
-        const std::function<bool(const char *, size_t)> & sink,
-        const std::function<bool()> & should_stop);
-
-    bool    is_done() const;
-    bool    is_cancelled() const;
-    size_t  total_size() const;     // bytes that ever entered the session
-    size_t  dropped_prefix() const; // bytes evicted from the front due to cap
-    int64_t completed_at() const;   // 0 while alive, unix seconds after finalize
-
-    // attach the producer stop hook used to cancel its reader, pass an empty function to detach
-    void set_stop_producer(std::function<void()> fn);
-
-    // signal the producer to abort its inference asap via the stop hook, idempotent
-    void cancel();
-
-private:
-    mutable std::mutex      mu;
-    std::condition_variable cv;
-    std::vector<char>       buffer;
-    size_t                  prefix_dropped;
-    size_t                  cap_bytes;
-    std::atomic<bool>       done;
-    std::atomic<bool>       cancelled;
-    std::atomic<int64_t>    completed_ts;
-    std::function<void()>   stop_producer; // protected by mu
-};
+struct stream_session;

 using stream_session_ptr = std::shared_ptr<stream_session>;

-// one end of a stream_session pipe. the base holds the session and the shared query, the
-// producer and consumer ends derive from it. virtual dtor so each end runs its own teardown:
+// base of the producer/consumer pipe ends. virtual dtor so each runs its own teardown:
 // the producer finalizes the session, the consumer leaves it untouched
 struct stream_pipe {
    virtual ~stream_pipe() = default;

-    // true if the session was cancelled (e.g. via DELETE /v1/stream/<conv_id>)
    bool is_cancelled() const;

 protected:
@@ -95,7 +37,6 @@ protected:
 struct stream_pipe_producer : stream_pipe {
    ~stream_pipe_producer() override;

-    // append raw bytes to the session's ring buffer, returns false if already finalized
    bool write(const char * data, size_t len);

    // mark the natural end on the wire so a later close() is a no-op
@@ -121,83 +62,21 @@ private:
    server_http_res *                   res_ = nullptr;
 };

-// consumer end: read-only replay of the ring buffer, the destructor does not finalize the session
-struct stream_pipe_consumer : stream_pipe {
-    // drain bytes from offset, calling sink for each available chunk. blocks until more data
-    // arrives or the session finalizes. should_stop is polled, returns OFFSET_LOST if offset
-    // fell below the dropped prefix
-    stream_read_status read(size_t & offset,
-        const std::function<bool(const char *, size_t)> & sink,
-        const std::function<bool()> & should_stop);
+void server_stream_session_manager_start();
+void server_stream_session_manager_stop();

-    static std::shared_ptr<stream_pipe_consumer> create(stream_session_ptr session);
+// route handler factories wired under /v1/stream/* by server.cpp
+server_http_context::handler_t server_stream_make_get_handler();
+server_http_context::handler_t server_stream_make_lookup_handler();
+server_http_context::handler_t server_stream_make_delete_handler();

-private:
-    explicit stream_pipe_consumer(stream_session_ptr session);
-};
+// extract the X-Conversation-Id header value (case-insensitive), empty when absent
+std::string server_stream_conv_id_from_headers(const std::map<std::string, std::string> & headers);

-// owns all live sessions, runs a periodic GC to evict expired ones.
-// the map is keyed by conversation_id, so the invariant "one conv = at most one
-// live session" is enforced at the type level
-class stream_session_manager {
-public:
-    stream_session_manager();
-    ~stream_session_manager();
-
-    stream_session_manager(const stream_session_manager &)             = delete;
-    stream_session_manager & operator=(const stream_session_manager &) = delete;
-
-    // install a new session for this conversation, evicting and cancelling any previous one.
-    // the conversation_id must be non empty, the caller is responsible for that check.
-    // returns the new session
-    stream_session_ptr create_or_replace(const std::string & conversation_id);
-
-    // lookup, returns null if unknown or already evicted
-    stream_session_ptr get(const std::string & conversation_id);
-
-    // list every live or recently completed session, used by GET /v1/streams without filter
-    std::vector<stream_session_ptr> list_all() const;
-
-    // remove from the map and finalize, wakes any pending readers
-    void evict(const std::string & conversation_id);
-
-    // signal the producer to cancel asap then evict, used by the explicit user Stop path
-    void evict_and_cancel(const std::string & conversation_id);
-
-    void start_gc();
-    void stop_gc();
-
-private:
-    void gc_loop();
-
-    mutable std::shared_mutex                           map_mu;
-    std::unordered_map<std::string, stream_session_ptr> sessions; // key: conversation_id
-    std::thread                                         gc_thread;
-    std::atomic<bool>                                   running;
-    std::mutex                                          gc_wake_mu;
-    std::condition_variable                             gc_wake_cv;
-};
-
-// process wide manager, linked by both llama-server and llama-cli. llama-server main() drives
-// start_gc/stop_gc, llama-cli leaves it idle. the dtor calls stop_gc() unconditionally so exit
-// is safe whether or not the GC thread ran
-extern stream_session_manager g_stream_sessions;
-
-// route handler factories operating on g_stream_sessions, wired under /v1/stream/* by server.cpp.
-// keeps the resumable stream surface confined to server-stream
-server_http_context::handler_t make_stream_get_handler();
-server_http_context::handler_t make_streams_lookup_handler();
-server_http_context::handler_t make_stream_delete_handler();
-
-// extract the X-Conversation-Id header value (case-insensitive), empty when absent. exposed so
-// the router can track which child serves a forwarded POST
-std::string stream_conv_id_from_headers(const std::map<std::string, std::string> & headers);
-
-// on an X-Conversation-Id header, create or replace the session and attach a producer pipe to
-// res. no-op when absent, called from the server_res_generator constructor
-void stream_session_attach_pipe(server_http_res & res, const std::map<std::string, std::string> & headers);
+// on an X-Conversation-Id header, create or replace the session and attach a producer pipe to res
+void server_stream_session_attach_pipe(server_http_res & res, const std::map<std::string, std::string> & headers);

 // should_stop closure that ignores peer disconnect when a pipe is attached, so only an explicit
 // DELETE stops the producer and generation keeps flowing into the ring buffer. without a pipe it
 // delegates to fallback, the legacy non-resumable flow
-std::function<bool()> stream_aware_should_stop(server_http_res * res, std::function<bool()> fallback);
+std::function<bool()> server_stream_aware_should_stop(server_http_res * res, std::function<bool()> fallback);
@@ -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));

@@ -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;
@@ -36,6 +36,19 @@ static inline void signal_handler(int signal) {
    shutdown_handler(signal);
 }

+// satisfies -Wmissing-declarations (used by llama command)
+int llama_server(int argc, char ** argv);
+
+// to be used via CLI (argc / argv are used by router mode only)
+int llama_server(common_params & params, int argc, char ** argv);
+void llama_server_terminate();
+void llama_server_terminate() {
+    if (shutdown_handler) {
+        shutdown_handler(0);
+    }
+}
+
+
 // wrapper function that handles exceptions and logs errors
 // this is to make sure handler_t never throws exceptions; instead, it returns an error response
 static server_http_context::handler_t ex_wrapper(server_http_context::handler_t func) {
@@ -72,9 +85,6 @@ static server_http_context::handler_t ex_wrapper(server_http_context::handler_t
    };
 }

-// satisfies -Wmissing-declarations
-int llama_server(int argc, char ** argv);
-
 int llama_server(int argc, char ** argv) {
    std::setlocale(LC_NUMERIC, "C");

@@ -85,7 +95,7 @@ int llama_server(int argc, char ** argv) {

    // start the stream session manager GC right after common init, before any HTTP route can
    // touch it. lifecycle is symmetric, stop_gc() runs in clean_up() before backend free
-    g_stream_sessions.start_gc();
+    server_stream_session_manager_start();

    if (!common_params_parse(argc, argv, params, LLAMA_EXAMPLE_SERVER)) {
        return 1;
@@ -94,16 +104,26 @@ int llama_server(int argc, char ** argv) {
    llama_backend_init();
    llama_numa_init(params.numa);

+    return llama_server(params, argc, argv);
+}
+
+int llama_server(common_params & params, int argc, char ** argv) {
+    bool is_run_by_cli = (argv == nullptr);
+
    common_models_handler models_handler;
-    try {
-        models_handler = common_models_handler_init(params, LLAMA_EXAMPLE_SERVER);
-        if (common_models_handler_is_preset_repo(models_handler)) {
-            // apply the preset and start the server in router mode
-            common_models_handler_apply(models_handler, params);
+
+    // note: router mode also accepts -hf remote-preset, so we need to check that first
+    if (!is_run_by_cli && !params.model.hf_repo.empty()) {
+        try {
+            models_handler = common_models_handler_init(params, LLAMA_EXAMPLE_SERVER);
+            if (common_models_handler_is_preset_repo(models_handler)) {
+                // apply the preset and start the server in router mode
+                common_models_handler_apply(models_handler, params);
+            }
+        } catch (const std::exception & e) {
+            SRV_ERR("failed to fetch model metadata: %s\n", e.what());
+            return 1;
        }
-    } catch (const std::exception & e) {
-        SRV_ERR("failed to fetch model metadata: %s\n", e.what());
-        return 1;
    }

    // router server never loads a model and must not touch the GPU
@@ -245,8 +265,8 @@ int llama_server(int argc, char ** argv) {
    ctx_http.post("/slots/:id_slot",           ex_wrapper(routes.post_slots));

    // resumable streaming, the conversation_id is the session identity end to end. router and
-    // child wire different handlers under the same paths: a child binds the local g_stream_sessions
-    // backed factories, the router binds proxies that resolve the owning child through the
+    // child wire different handlers under the same paths: a child binds the local session
+    // factories, the router binds proxies that resolve the owning child through the
    // conv_id -> model map
    server_http_context::handler_t stream_get_h;
    server_http_context::handler_t streams_lookup_h;
@@ -256,9 +276,9 @@ int llama_server(int argc, char ** argv) {
        streams_lookup_h = models_routes->router_streams_lookup;
        stream_delete_h  = models_routes->router_stream_delete;
    } else {
-        stream_get_h     = make_stream_get_handler();
-        streams_lookup_h = make_streams_lookup_handler();
-        stream_delete_h  = make_stream_delete_handler();
+        stream_get_h     = server_stream_make_get_handler();
+        streams_lookup_h = server_stream_make_lookup_handler();
+        stream_delete_h  = server_stream_make_delete_handler();
    }
    ctx_http.get ("/v1/stream/:conv_id",       ex_wrapper(stream_get_h));
    // POST /v1/streams/lookup with body {"conversation_ids": [...]}. you can only ask for ids
@@ -321,8 +341,9 @@ int llama_server(int argc, char ** argv) {

    if (child.is_child() && child.get_mode() == SERVER_CHILD_MODE_DOWNLOAD) {
        return child.run_download(params);
-    } else if (!is_router_server) {
+    } else if (!is_router_server && !is_run_by_cli) {
        // single-model mode (NOT spawned by router)
+        // if this is invoked by CLI, model downloading should be already handled
        try {
            common_models_handler_apply(models_handler, params);
        } catch (const std::exception & e) {
@@ -343,7 +364,7 @@ int llama_server(int argc, char ** argv) {
        clean_up = [&models_routes]() {
            SRV_INF("%s: cleaning up before exit...\n", __func__);
            // stop the session GC first, it finalizes live sessions and wakes pending readers
-            g_stream_sessions.stop_gc();
+            server_stream_session_manager_stop();
            if (models_routes.has_value()) {
                models_routes->stopping.store(true); // maybe redundant, but just to be safe
                models_routes->models.unload_all();
@@ -371,7 +392,7 @@ int llama_server(int argc, char ** argv) {
        clean_up = [&ctx_http, &ctx_server]() {
            SRV_INF("%s: cleaning up before exit...\n", __func__);
            // stop the session GC first, it finalizes live sessions and wakes pending readers
-            g_stream_sessions.stop_gc();
+            server_stream_session_manager_stop();
            ctx_http.stop();
            ctx_server.terminate();
            llama_backend_free();
@@ -411,20 +432,22 @@ int llama_server(int argc, char ** argv) {
        };
    }

-    // TODO: refactor in common/console
+    // register signal handler if not running by CLI
+    if (!is_run_by_cli) {
 #if defined (__unix__) || (defined (__APPLE__) && defined (__MACH__))
-    struct sigaction sigint_action;
-    sigint_action.sa_handler = signal_handler;
-    sigemptyset (&sigint_action.sa_mask);
-    sigint_action.sa_flags = 0;
-    sigaction(SIGINT, &sigint_action, NULL);
-    sigaction(SIGTERM, &sigint_action, NULL);
+        struct sigaction sigint_action;
+        sigint_action.sa_handler = signal_handler;
+        sigemptyset (&sigint_action.sa_mask);
+        sigint_action.sa_flags = 0;
+        sigaction(SIGINT, &sigint_action, NULL);
+        sigaction(SIGTERM, &sigint_action, NULL);
 #elif defined (_WIN32)
-    auto console_ctrl_handler = +[](DWORD ctrl_type) -> BOOL {
-        return (ctrl_type == CTRL_C_EVENT) ? (signal_handler(SIGINT), true) : false;
-    };
-    SetConsoleCtrlHandler(reinterpret_cast<PHANDLER_ROUTINE>(console_ctrl_handler), true);
+        auto console_ctrl_handler = +[](DWORD ctrl_type) -> BOOL {
+            return (ctrl_type == CTRL_C_EVENT) ? (signal_handler(SIGINT), true) : false;
+        };
+        SetConsoleCtrlHandler(reinterpret_cast<PHANDLER_ROUTINE>(console_ctrl_handler), true);
 #endif
+    }

    SRV_INF("listening on %s\n", ctx_http.listening_address.c_str());

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