update skill

* Initial changes for Recurrent state rollback for nemotron for cpu and cuda

* Removing CPU RS rollback. Will enable it in subsequent PRs

* addition of test case

* Removing assert and calling runtime API to check if op is supported

* removing extra API and updating the call sites for K

* replace static cuda detection to runtime fused_op api

* address review comments and fallback when SSM rollback not supprted

* Adding changes for supporting RS-rollback in CPU. Also added test-backend-ops for cpu and cuda

* removing memory manipulation as rs rollback is now supported in CPU

* removing the static probe which is not needed now

* correcting the format

* address review comments

* enabling test for all the backends, unsupported backends will fallback to CPU

* Apply suggestions from code review

Co-authored-by: Georgi Gerganov <ggerganov@gmail.com>

* choose different graph based on the result of fused_ssm_op is supported or not and also handled memory->n_rs_seq >1 case incase of op is not supported

* Support K > 1 in ssm_scan for all backends

* Fix CI Issues

---------

Co-authored-by: Georgi Gerganov <ggerganov@gmail.com>
Co-authored-by: Gaurav Garg <gaugarg@nvidia.com>

common/common.cpp

@@ -1275,6 +1275,8 @@ struct common_init_result::impl {
 
     // note: the order in which model, context, etc. are declared matters because their destructors will be called bottom-to-top
 
+    common_threadpools threadpools;
+
     llama_model_ptr   model;
     llama_context_ptr context;
 
@@ -1376,6 +1378,10 @@ common_init_result::common_init_result(common_params & params, bool model_only)
     }
 
     pimpl->context.reset(lctx);
+
+    set_process_priority(params.cpuparams.priority);
+
+    pimpl->threadpools.init(lctx, params);
 }
 
 llama_model * common_init_result::model() {
@@ -1724,6 +1730,10 @@ struct llama_context_params common_context_params_to_llama(const common_params &
     return cparams;
 }
 
+//
+// Threadpool utils
+//
+
 struct ggml_threadpool_params ggml_threadpool_params_from_cpu_params(const common_cpu_params & params) {
     struct ggml_threadpool_params tpp;
 
@@ -1740,6 +1750,56 @@ struct ggml_threadpool_params ggml_threadpool_params_from_cpu_params(const commo
     return tpp;
 }
 
+common_threadpools::~common_threadpools() {
+    if (!free_fn) {
+        return;
+    }
+    free_fn(threadpool);
+    free_fn(threadpool_batch);
+}
+
+void common_threadpools::init(llama_context * ctx, const common_params & params) {
+    GGML_ASSERT(!threadpool);
+    GGML_ASSERT(!threadpool_batch);
+
+    COM_INF("llama threadpool init, n_threads = %d\n", (int) params.cpuparams.n_threads);
+
+    auto * cpu_dev = ggml_backend_dev_by_type(GGML_BACKEND_DEVICE_TYPE_CPU);
+    if (!cpu_dev) {
+        COM_WRN("%s", "no CPU backend found\n");
+        return;
+    }
+    auto * reg = ggml_backend_dev_backend_reg(cpu_dev);
+    auto * ggml_threadpool_new_fn = (decltype(ggml_threadpool_new) *) ggml_backend_reg_get_proc_address(reg, "ggml_threadpool_new");
+    free_fn = (decltype(ggml_threadpool_free) *) ggml_backend_reg_get_proc_address(reg, "ggml_threadpool_free");
+
+    struct ggml_threadpool_params tpp_batch =
+            ggml_threadpool_params_from_cpu_params(params.cpuparams_batch);
+    struct ggml_threadpool_params tpp =
+            ggml_threadpool_params_from_cpu_params(params.cpuparams);
+
+    if (!ggml_threadpool_params_match(&tpp, &tpp_batch)) {
+        threadpool_batch = ggml_threadpool_new_fn(&tpp_batch);
+        if (!threadpool_batch) {
+            COM_WRN("batch threadpool create failed : n_threads %d\n", tpp_batch.n_threads);
+            return;
+        }
+
+        // start the non-batch threadpool in the paused state
+        tpp.paused = true;
+    }
+
+    threadpool = ggml_threadpool_new_fn(&tpp);
+    if (!threadpool) {
+        COM_WRN("threadpool create failed : n_threads %d\n", tpp.n_threads);
+        free_fn(threadpool_batch);
+        threadpool_batch = nullptr;
+        return;
+    }
+
+    llama_attach_threadpool(ctx, threadpool, threadpool_batch);
+}
+
 //
 // Batch utils
 //

common/common.h

@@ -929,9 +929,8 @@ using common_init_result_ptr = std::unique_ptr<common_init_result>;
 
 common_init_result_ptr common_init_from_params(common_params & params, bool model_only = false);
 
-struct llama_model_params     common_model_params_to_llama  (      common_params & params);
-struct llama_context_params   common_context_params_to_llama(const common_params & params);
-struct ggml_threadpool_params ggml_threadpool_params_from_cpu_params(const common_cpu_params & params);
+struct llama_model_params   common_model_params_to_llama  (      common_params & params);
+struct llama_context_params common_context_params_to_llama(const common_params & params);
 
 // clear LoRA adapters from context, then apply new list of adapters
 void common_set_adapter_lora(struct llama_context * ctx, std::vector<common_adapter_lora_info> & lora);
@@ -942,6 +941,28 @@ std::string common_get_model_endpoint();
 // for testing purposes
 char * common_get_model_or_exit(int, char*[]);
 
+//
+// Threadpool utils
+//
+
+struct ggml_threadpool_params ggml_threadpool_params_from_cpu_params(const common_cpu_params & params);
+
+struct common_threadpools {
+    common_threadpools() = default;
+    ~common_threadpools();
+
+    common_threadpools(const common_threadpools &) = delete;
+    common_threadpools & operator=(const common_threadpools &) = delete;
+
+    void init(llama_context * ctx, const common_params & params);
+
+private:
+    ggml_threadpool * threadpool       = nullptr;
+    ggml_threadpool * threadpool_batch = nullptr;
+
+    decltype(ggml_threadpool_free) * free_fn = nullptr;
+};
+
 //
 // Context utils
 //

common/imatrix-loader.cpp

@@ -102,7 +102,8 @@ bool common_imatrix_load(const std::string & fname, common_imatrix & imatrix) {
     const int64_t chunk_count_key = gguf_find_key(ctx_gguf, LLM_KV_IMATRIX_CHUNK_COUNT);
     const int64_t chunk_size_key  = gguf_find_key(ctx_gguf, LLM_KV_IMATRIX_CHUNK_SIZE);
 
-    if (datasets_key != -1 && gguf_get_arr_type(ctx_gguf, datasets_key) == GGUF_TYPE_STRING) {
+    if (datasets_key != -1 && gguf_get_kv_type(ctx_gguf, datasets_key) == GGUF_TYPE_ARRAY &&
+        gguf_get_arr_type(ctx_gguf, datasets_key) == GGUF_TYPE_STRING) {
         const int64_t n = gguf_get_arr_n(ctx_gguf, datasets_key);
         imatrix.datasets.reserve(imatrix.datasets.size() + n);
         for (int64_t i = 0; i < n; ++i) {
@@ -143,6 +144,13 @@ bool common_imatrix_load(const std::string & fname, common_imatrix & imatrix) {
             return false;
         }
 
+        if (in_sum2->type != GGML_TYPE_F32 || counts->type != GGML_TYPE_F32) {
+            LOG_ERR("%s: sums and counts for %s must be F32\n", __func__, name.c_str());
+            gguf_free(ctx_gguf);
+            ggml_free(ctx);
+            return false;
+        }
+
         auto & e = imatrix.entries[name];
 
         const int64_t nval    = ggml_nelements(in_sum2);

docs/backend/SYCL.md

@@ -804,7 +804,7 @@ User can use the device management in [docs/multi-gpu.md](https://github.com/ggm
 | GGML_SYCL_ENABLE_MKL_FA | 1 (default) or 0 | Enable oneMKL GEMM flash attention for XMX-accelerated prompt processing with quantized KV cache. Automatically activates during prefill (prompt processing) when all conditions are met: (1) flash-attn enabled (`-fa` or `--flash-attn on`), (2) KV cache quantized (`--cache-type-k q8_0 --cache-type-v q8_0` or other `*_0/*_1` types), (3) batch size ≥ 1024 (`--batch-size 1024`), (4) prompt length ≥ 1024 tokens. Set to 0 to force the TILE kernel for A/B testing. Example minimum command: `llama-cli -m model.gguf -fa -ngl 99 --cache-type-k q8_0 --cache-type-v q8_0 --batch-size 1024 -p "your prompt"` |
 | GGML_SYCL_MKL_FA_DEBUG | 0 (default) or 1 | Enable per-call diagnostic logging for MKL flash attention: GEMM/softmax timings, interleaved-head detection, and buffer memory usage. |
 | GGML_SYCL_MKL_FA_DIAG | 0 (default) or 1 | Enable output fingerprinting for MKL flash attention. Dumps the first 64 float output values for the first 6 FA calls with n_kv ≥ 1024, labeled with kernel type (MKL/TILE/VEC) for cross-kernel comparison. |
-| GGML_SYCL_ENABLE_FUSION | 0 or 1 (default) | Enable fused-kernel dispatch in graph compute (currently top-k MoE gating). |
+| GGML_SYCL_ENABLE_FUSION | 0 or 1 (default) | Enable fused-kernel dispatch in graph compute. |
 | GGML_SYCL_ENABLE_ESIMD | 0 or 1 (default)| Enable ESIMD kernels when available. |
 | ZES_ENABLE_SYSMAN | 0 (default) or 1 | Support to get free memory of GPU by sycl::aspect::ext_intel_free_memory.<br>Recommended to use when --split-mode = layer |
 | UR_L0_ENABLE_RELAXED_ALLOCATION_LIMITS | 0 (default) or 1 | Allow SYCL/Unified Runtime Level Zero device allocations larger than 4 GiB. llama.cpp's direct Level Zero allocation path requests the relaxed maximum-size limit itself when GGML_SYCL_ENABLE_LEVEL_ZERO=1. |

examples/convert-llama2c-to-ggml/convert-llama2c-to-ggml.cpp

@@ -549,20 +549,34 @@ static void load_vocab(const char * filename, const Config * config, struct my_l
 
         const int token_idx = gguf_find_key(ctx, KV_TOKENIZER_LIST);
         GGML_ASSERT(token_idx >= 0);
-
-        const int score_idx = gguf_find_key(ctx, KV_TOKENIZER_SCORES);
-        GGML_ASSERT(score_idx >= 0);
-        const float * scores = (const float * ) gguf_get_arr_data(ctx, score_idx);
-
-        const int toktype_idx = gguf_find_key(ctx, KV_TOKENIZER_TOKEN_TYPE);
-        GGML_ASSERT(toktype_idx >= 0);
-        const int * toktypes = (const int * ) gguf_get_arr_data(ctx, toktype_idx);
+        if (gguf_get_kv_type(ctx, token_idx) != GGUF_TYPE_ARRAY ||
+            gguf_get_arr_type(ctx, token_idx) != GGUF_TYPE_STRING) {
+            die_fmt("invalid gguf type for %s", KV_TOKENIZER_LIST);
+        }
 
         const uint32_t n_vocab = gguf_get_arr_n(ctx, token_idx);
         if (n_vocab != static_cast<uint32_t>(config->vocab_size)) {
             die_fmt("vocab size mismatch: (gguf) %u != (llama2c) %d", n_vocab, config->vocab_size);
         }
 
+        const int score_idx = gguf_find_key(ctx, KV_TOKENIZER_SCORES);
+        GGML_ASSERT(score_idx >= 0);
+        if (gguf_get_kv_type(ctx, score_idx) != GGUF_TYPE_ARRAY ||
+            gguf_get_arr_type(ctx, score_idx) != GGUF_TYPE_FLOAT32 ||
+            gguf_get_arr_n(ctx, score_idx) < n_vocab) {
+            die_fmt("invalid gguf type or size for %s", KV_TOKENIZER_SCORES);
+        }
+        const float * scores = (const float * ) gguf_get_arr_data(ctx, score_idx);
+
+        const int toktype_idx = gguf_find_key(ctx, KV_TOKENIZER_TOKEN_TYPE);
+        GGML_ASSERT(toktype_idx >= 0);
+        if (gguf_get_kv_type(ctx, toktype_idx) != GGUF_TYPE_ARRAY ||
+            gguf_get_arr_type(ctx, toktype_idx) != GGUF_TYPE_INT32 ||
+            gguf_get_arr_n(ctx, toktype_idx) < n_vocab) {
+            die_fmt("invalid gguf type or size for %s", KV_TOKENIZER_TOKEN_TYPE);
+        }
+        const int * toktypes = (const int * ) gguf_get_arr_data(ctx, toktype_idx);
+
         vocab->id_to_token.resize(n_vocab);
 
         for (uint32_t i = 0; i < n_vocab; i++) {

examples/test-cmake/README.md

@@ -27,10 +27,10 @@ Build/run this project using the installation created above:
 (venv) $ ./build.sh
 -- Configuring done (0.0s)
 -- Generating done (0.0s)
--- Build files have been written to: /home/danbev/work/ai/llama.cpp/examples/test-cmake/build
+-- Build files have been written to: /path/to/llama.cpp/examples/test-cmake/build
 [100%] Built target test-cmake
 [test-cmake] Using llama.cpp version 0.1.0-dev-b10335
 [test-cmake] Initializing backend...
-load_backend: loaded CPU backend from /home/danbev/work/ai/llama.cpp/examples/test-cmake/install/lib/llama.cpp/libggml-cpu-alderlake.so
+load_backend: loaded CPU backend from /path/to/llama.cpp/examples/test-cmake/install/lib/llama.cpp/libggml-cpu-alderlake.so
 [test-cmake] Backend initialized.
 ```

ggml/include/ggml.h

@@ -2459,7 +2459,8 @@ extern "C" {
             struct ggml_tensor  * A,
             struct ggml_tensor  * B,
             struct ggml_tensor  * C,
-            struct ggml_tensor  * ids);
+            struct ggml_tensor  * ids,
+            int64_t               K);
 
     // partition into non-overlapping windows with padding if needed
     // example:

ggml/src/ggml-cpu/ggml-cpu.c

@@ -2795,6 +2795,11 @@ struct ggml_cplan ggml_graph_plan(
     n_threads = 1;
 #endif
 
+#if defined(__wasi__)
+    // WASI doesn't support parallelism yet
+    n_threads = 1;
+#endif
+
     size_t work_size = 0;
 
     struct ggml_cplan cplan;

ggml/src/ggml-cpu/ggml-cpu.cpp

@@ -472,6 +472,8 @@ static bool ggml_backend_cpu_device_supports_op(ggml_backend_dev_t dev, const st
                 src1->type == GGML_TYPE_F32 && op->type == GGML_TYPE_F32;
         case GGML_OP_CONV_2D:
             return ggml_is_contiguous(op->src[0]);
+        case GGML_OP_SSM_SCAN:
+            return ggml_get_op_params_i32(op, 0) == 1 || op->src[3]->ne[0] == 1;
         default:
             return true;
     }

ggml/src/ggml-cpu/ops.cpp

@@ -9644,11 +9644,13 @@ static void ggml_compute_forward_ssm_scan_f32(
     const int64_t ng = src4->ne[1];
     const int64_t nt = src1->ne[2]; // number of tokens per sequence
     const int64_t ns = src1->ne[3]; // number of sequences in the batch
+    const int64_t K  = ggml_get_op_params_i32(dst, 0);
 
     // can't use ggml_nbytes because src1 is not necessarily contiguous
     const int64_t s_off = ggml_nelements(src1) * ggml_element_size(src1);
 
-    GGML_ASSERT(ggml_nelements(src1) + nc*nr*nh*ns == ggml_nelements(dst));
+    GGML_ASSERT(K >= 1);
+    GGML_ASSERT(ggml_nelements(src1) + K*nc*nr*nh*ns == ggml_nelements(dst));
     GGML_ASSERT(src0->nb[0] == sizeof(float));
     GGML_ASSERT(src1->nb[0] == sizeof(float));
     GGML_ASSERT(src2->nb[0] == sizeof(float));
@@ -9657,6 +9659,7 @@ static void ggml_compute_forward_ssm_scan_f32(
     GGML_ASSERT(src5->nb[0] == sizeof(float));
     GGML_ASSERT(src6->nb[0] == sizeof(int32_t));
     GGML_ASSERT(nh % ng == 0);
+    GGML_ASSERT(src3->ne[0] == 1 || K == 1);
 
     // heads per thread
     const int dh = (nh + nth - 1)/nth;
@@ -9831,6 +9834,13 @@ static void ggml_compute_forward_ssm_scan_f32(
                     }
                 }
             }
+            const int64_t slot = nt - 1 - i2;
+            if (K > 1 && slot > 0 && slot < K) {
+                float * s_snapshot = (float *) ((char *) dst->data + s_off + (slot*ns + i3)*(src0->nb[3]));
+                for (int h = ih0; h < ih1; ++h) {
+                    memcpy((char *) s_snapshot + h*src0->nb[2], (char *) s + h*src0->nb[2], src0->nb[2]);
+                }
+            }
             // use the output as the source when it's not the first token-wise iteration
             s0 = s;
         }

ggml/src/ggml-cuda/ggml-cuda.cu

@@ -5189,11 +5189,17 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g
                    (op->src[1]->type == GGML_TYPE_F32 || op->src[1]->type == GGML_TYPE_F16) &&
                    (op->type         == GGML_TYPE_F32 || op->type         == GGML_TYPE_F16);
         case GGML_OP_SSM_SCAN: {
+            const int32_t K = ggml_get_op_params_i32(op, 0);
+
             if (op->src[3]->ne[0] == 1) {
                 // Mamba2
                 // (kernel only supports (d_state == 128 || d_state == 256) && d_head % 16 == 0)
                 return (op->src[0]->ne[0] == 128 || op->src[0]->ne[0] == 256) && op->src[0]->ne[1] % 16 == 0;
             } else {
+                if (K > 1) {
+                    return false;
+                }
+
                 // Mamba
                 // (kernel only supports d_state == 16, d_head == 1, n_head % 128 == 0, n_group == 1)
                 return op->src[0]->ne[0] == 16 && op->src[0]->ne[1] == 1 && op->src[0]->ne[2] % 128 == 0 && op->src[4]->ne[1] == 1;

ggml/src/ggml-cuda/ssm-scan.cu

@@ -149,7 +149,7 @@ __global__ void __launch_bounds__(d_state, 1)
         const int src0_nb2, const int src0_nb3, const int src1_nb2, const int src1_nb3,
         const int src2_nb1, const int src2_nb2, const int src3_nb1,
         const int src4_nb2, const int src4_nb3, const int src5_nb2, const int src5_nb3,
-        const int64_t s_off, const int64_t n_head, const int64_t d_head, const int64_t n_group, const int64_t n_tok) {
+        const int64_t s_off, const int64_t n_head, const int64_t d_head, const int64_t n_group, const int64_t n_tok, const int64_t K) {
     const float   * GGML_CUDA_RESTRICT src0 = src0_ptr;
     const float   * GGML_CUDA_RESTRICT src1 = src1_ptr;
     const float   * GGML_CUDA_RESTRICT src2 = src2_ptr;
@@ -217,6 +217,16 @@ __global__ void __launch_bounds__(d_state, 1)
         if (lane == 0) {
             y_warp[i * stride_y] = state_sum;
         }
+
+        // Slot 0 is the final state written below; slots 1..K-1 are rollback snapshots.
+        const int64_t slot = n_tok - 1 - i;
+        if (K > 1 && slot > 0 && slot < K) {
+            float * s_snapshot_warp = (float *) ((char *) dst + s_off + (slot * gridDim.y + seq_idx) * src0_nb3 + head_idx * src0_nb2 + head_off * d_state);
+#pragma unroll
+            for (int j = 0; j < c_factor; j++) {
+                s_snapshot_warp[WARP_SIZE * j + lane] = state[j];
+            }
+        }
     }
 
     // write back the state
@@ -232,7 +242,7 @@ static void ssm_scan_f32_cuda(const float * src0, const float * src1, const floa
                               const int src2_nb2, const int src3_nb1, const int src4_nb2, const int src4_nb3, const int src5_nb2,
                               const int src5_nb3, const int64_t s_off, const int64_t d_state, const int64_t head_dim,
                               const int64_t n_head, const int64_t n_group, const int64_t n_tok, const int64_t n_seq,
-                              cudaStream_t stream) {
+                              const int64_t K, cudaStream_t stream) {
     // NOTE: if you change conditions here, be sure to update the corresponding supports_op condition!
     if (src3_nb1 == sizeof(float)) {
         // Mamba-2
@@ -245,7 +255,7 @@ static void ssm_scan_f32_cuda(const float * src0, const float * src1, const floa
             ggml_cuda_kernel_launch(ssm_scan_f32_group<128/WARP_SIZE, 128>, launch_params,
                     src0, src1, src2, src3, src4, src5, src6, dst,
                     src0_nb2, src0_nb3, src1_nb2, src1_nb3, src2_nb1, src2_nb2, src3_nb1,
-                    src4_nb2, src4_nb3, src5_nb2, src5_nb3, s_off, n_head, head_dim, n_group, n_tok);
+                    src4_nb2, src4_nb3, src5_nb2, src5_nb3, s_off, n_head, head_dim, n_group, n_tok, K);
         } else if (d_state == 256) { // Falcon-H1
             constexpr int threads   = 256;
             constexpr int num_warps = threads/WARP_SIZE;
@@ -255,12 +265,13 @@ static void ssm_scan_f32_cuda(const float * src0, const float * src1, const floa
             ggml_cuda_kernel_launch(ssm_scan_f32_group<256/WARP_SIZE, 256>, launch_params,
                     src0, src1, src2, src3, src4, src5, src6, dst,
                     src0_nb2, src0_nb3, src1_nb2, src1_nb3, src2_nb1, src2_nb2, src3_nb1,
-                    src4_nb2, src4_nb3, src5_nb2, src5_nb3, s_off, n_head, head_dim, n_group, n_tok);
+                    src4_nb2, src4_nb3, src5_nb2, src5_nb3, s_off, n_head, head_dim, n_group, n_tok, K);
         } else {
             GGML_ABORT("doesn't support d_state!=(128 or 256).");
         }
     } else {
         // Mamba-1
+        GGML_ASSERT(K == 1);
         constexpr int threads = 128;
         GGML_ASSERT(n_head % threads == 0);
         GGML_ASSERT(head_dim == 1);
@@ -769,10 +780,12 @@ void ggml_cuda_op_ssm_scan(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
     const int64_t ng  = src4->ne[1];  // n_group
     const int64_t n_t = src1->ne[2];  // number of tokens per sequence
     const int64_t n_s = src1->ne[3];  // number of sequences in the batch
+    const int32_t K_param = ggml_get_op_params_i32(dst, 0);
+    const int64_t K = K_param > 0 ? K_param : 1;
 
     const int64_t s_off = ggml_nelements(src1) * sizeof(float);
 
-    GGML_ASSERT(ggml_nelements(src1) + nc*nr*nh*n_s == ggml_nelements(dst));
+    GGML_ASSERT(ggml_nelements(src1) + K*nc*nr*nh*n_s == ggml_nelements(dst));
     GGML_ASSERT(src0->nb[0] == sizeof(float));
     GGML_ASSERT(src1->nb[0] == sizeof(float));
     GGML_ASSERT(src2->nb[0] == sizeof(float));
@@ -780,6 +793,7 @@ void ggml_cuda_op_ssm_scan(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
     GGML_ASSERT(src4->nb[0] == sizeof(float));
     GGML_ASSERT(src5->nb[0] == sizeof(float));
     GGML_ASSERT(src6->nb[0] == sizeof(int32_t));
+    GGML_ASSERT(src3->ne[0] == 1 || K == 1);
 
     const float * src0_d = (const float *) src0->data;
     const float * src1_d = (const float *) src1->data;
@@ -814,6 +828,7 @@ void ggml_cuda_op_ssm_scan(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
     const bool is_mamba2 = (src3->nb[1] == sizeof(float));
     const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc;
     const bool use_ssd = is_mamba2 && n_t > SSM_SSD_MIN_TOKENS
+                      && K == 1
                       && n_t <= SSM_SSD_MAX_TOKENS
                       && GGML_CUDA_CC_IS_NVIDIA(cc)
                       && cc >= GGML_CUDA_CC_TURING
@@ -841,5 +856,5 @@ void ggml_cuda_op_ssm_scan(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
     ssm_scan_f32_cuda(src0_d, src1_d, src2_d, src3_d, src4_d, src5_d, src6_d, dst_d,
                       src0->nb[2], src0->nb[3], src1->nb[2], src1->nb[3], src2->nb[1], src2->nb[2],
                       src3->nb[1], src4->nb[2], src4->nb[3], src5->nb[2], src5->nb[3],
-                      s_off, nc, nr, nh, ng, n_t, n_s, stream);
+                      s_off, nc, nr, nh, ng, n_t, n_s, K, stream);
 }

ggml/src/ggml-et/et-kernels/src/ssm_scan_f32.c

@@ -12,7 +12,8 @@ struct ggml_et_ssm_scan_params {
     struct ggml_tensor src4;  // B:   [d_state, n_group, n_seq_tokens, n_seqs]
     struct ggml_tensor src5;  // C:   [d_state, n_group, n_seq_tokens, n_seqs]
     struct ggml_tensor src6;  // ids: [n_seqs] i32
-    struct ggml_tensor dst;   // packed [y, final_state]
+    struct ggml_tensor dst;   // packed [y, states]
+    int32_t            K;
 };
 
 static inline float softplus_f32(float x) {
@@ -72,6 +73,7 @@ int entry_point(struct ggml_et_ssm_scan_params * params, void * env) {
     const int64_t n_seq_tokens = src1->ne[2];
     const int64_t n_seqs       = src1->ne[3];
     const int64_t y_elems      = src1->ne[0] * src1->ne[1] * src1->ne[2] * src1->ne[3];
+    const int64_t K            = params->K;
 
     if (src0->nb[0] != sizeof(float) || src1->nb[0] != sizeof(float) || src2->nb[0] != sizeof(float) ||
         src3->nb[0] != sizeof(float) || src4->nb[0] != sizeof(float) || src5->nb[0] != sizeof(float) ||
@@ -79,7 +81,7 @@ int entry_point(struct ggml_et_ssm_scan_params * params, void * env) {
         return -1;
     }
 
-    if (n_group <= 0 || n_head % n_group != 0) {
+    if (K < 1 || n_group <= 0 || n_head % n_group != 0) {
         return -1;
     }
 
@@ -260,6 +262,15 @@ int entry_point(struct ggml_et_ssm_scan_params * params, void * env) {
                         sumf += st * C_row[state_idx];
                     }
 
+                    const int64_t slot = n_seq_tokens - 1 - token_idx;
+                    if (slot > 0 && slot < K) {
+                        float * state_snapshot =
+                            (float *) ((char *) state_dst + (size_t) slot * n_seqs * src0->nb[3]);
+                        for (int64_t i = 0; i < d_state; ++i) {
+                            state_snapshot[i] = state_dst[i];
+                        }
+                    }
+
                     dst_data[seq_idx * (n_seq_tokens * n_head * head_dim) + token_idx * (n_head * head_dim) +
                              head_idx * head_dim + dim_idx] = sumf;
                 }

ggml/src/ggml-et/ggml-et-ops.cpp

@@ -2064,6 +2064,7 @@ bool ggml_et_op_ssm_scan(ggml_backend_et_device_context * dev_ctx, const ggml_te
     params.src5 = *node->src[5];
     params.src6 = *node->src[6];
     params.dst  = *node;
+    params.K    = ggml_get_op_params_i32(node, 0);
 
     bool kernel_result = ggml_et_launch_kernel(dev_ctx, "ssm_scan_f32", &params, sizeof(params), 0xFFFFFFFF);
 

ggml/src/ggml-et/ggml-et-ops.h

@@ -218,7 +218,8 @@ struct ggml_et_ssm_scan_params {
     ggml_tensor src4;  // B:   [d_state, n_group, n_seq_tokens, n_seqs]
     ggml_tensor src5;  // C:   [d_state, n_group, n_seq_tokens, n_seqs]
     ggml_tensor src6;  // ids: [n_seqs] i32
-    ggml_tensor dst;   // [y, final_state] packed output from ggml_ssm_scan()
+    ggml_tensor dst;   // [y, states] packed output from ggml_ssm_scan()
+    int32_t     K;
 };
 
 struct ggml_et_rwkv_wkv6_params {

ggml/src/ggml-metal/ggml-metal-device.m

@@ -1376,9 +1376,10 @@ bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_te
                 ggml_is_contiguous_rows(op->src[1]) &&
                 ggml_is_contiguous_rows(op->src[2]) &&
                 ggml_is_contiguous_rows(op->src[3]);
-        case GGML_OP_SSM_CONV:
         case GGML_OP_SSM_SCAN:
             return has_simdgroup_reduction;
+        case GGML_OP_SSM_CONV:
+            return has_simdgroup_reduction;
         case GGML_OP_RWKV_WKV6:
         case GGML_OP_RWKV_WKV7:
             return true;

ggml/src/ggml-metal/ggml-metal-impl.h

@@ -880,6 +880,7 @@ typedef struct {
     int64_t  n_group;
     int64_t  n_seq_tokens;
     int64_t  n_seqs;
+    int64_t  K;
     uint64_t s_off;
     uint64_t nb00;
     uint64_t nb01;

ggml/src/ggml-metal/ggml-metal-ops.cpp

@@ -1710,6 +1710,10 @@ int ggml_metal_op_ssm_scan(ggml_metal_op_t ctx, int idx) {
     const int64_t n_group      = ne41;
     const int64_t n_seq_tokens = ne12;
     const int64_t n_seqs       = ne13;
+    const int64_t K            = ggml_get_op_params_i32(op, 0);
+
+    GGML_ASSERT(K >= 1);
+    GGML_ASSERT(ggml_nelements(op->src[1]) + K*d_state*d_inner*n_head*n_seqs == ggml_nelements(op));
 
     ggml_metal_kargs_ssm_scan args = {
         /*.d_state      =*/ d_state,
@@ -1718,6 +1722,7 @@ int ggml_metal_op_ssm_scan(ggml_metal_op_t ctx, int idx) {
         /*.n_group      =*/ n_group,
         /*.n_seq_tokens =*/ n_seq_tokens,
         /*.n_seqs       =*/ n_seqs,
+        /*.K            =*/ K,
         /*.s_off        =*/ ggml_nelements(op->src[1]) * sizeof(float),
         /*.nb00         =*/ nb00,
         /*.nb01         =*/ nb01,

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

@@ -2429,6 +2429,8 @@ kernel void kernel_ssm_scan_f32(
     const int32_t nh  = args.n_head;
     const int32_t ng  = args.n_group;
     const int32_t n_t = args.n_seq_tokens;
+    const int32_t n_s = args.n_seqs;
+    const int32_t K   = args.K;
 
     const int32_t s_off = args.s_off;
 
@@ -2487,6 +2489,12 @@ kernel void kernel_ssm_scan_f32(
             // recurse
             s0 = s;
 
+            const int32_t slot = n_t - 1 - (i2 + t);
+            if (slot > 0 && slot < K) {
+                device float * s_snapshot = (device float *) ((device char *) s_buff + (int64_t) slot*n_s*args.nb03);
+                s_snapshot[i] = s;
+            }
+
             B  += args.ns42;
             C  += args.ns52;
         }

ggml/src/ggml-sycl/element_wise.cpp

@@ -81,43 +81,6 @@ static __dpct_inline__ T op_elu(T x) {
     return (x > static_cast<T>(0.f)) ? x : op_expm1(x);
 }
 
-template<typename T>
-static __dpct_inline__ T op_tanh(T x) {
-    if constexpr (std::is_same_v<T, sycl::ext::oneapi::bfloat16>) {
-        constexpr int ver = __INTEL_LLVM_COMPILER;
-#if defined(__INTEL_LLVM_COMPILER) && (__INTEL_LLVM_COMPILER >= 20260000)
-            return sycl::ext::oneapi::experimental::tanh(x);
-#else
-            return static_cast<T>(sycl::tanh(static_cast<float>(x)));
-#endif
-    } else {
-        return sycl::tanh(x);
-    }
-}
-
-template<typename T>
-static __dpct_inline__ T op_gelu(T x) {
-    const T GELU_COEF_A    = static_cast<T>(0.044715f);
-    const T SQRT_2_OVER_PI = static_cast<T>(0.79788456080286535587989211986876f);
-    return static_cast<T>(0.5f) * x *
-           (static_cast<T>(1.0f) +
-            op_tanh(SQRT_2_OVER_PI * x * (static_cast<T>(1.0f) + GELU_COEF_A * x * x)));
-}
-
-template<typename T>
-static __dpct_inline__ T op_exp(T x) {
-    if constexpr (std::is_same_v<T, sycl::ext::oneapi::bfloat16>) {
-        return sycl::ext::oneapi::experimental::exp(x);
-    } else {
-        return sycl::exp(x);
-    }
-}
-
-template<typename T>
-static __dpct_inline__ T op_silu(T x) {
-    return x / (static_cast<T>(1.0f) + op_exp(-x));
-}
-
 template<typename T>
 static __dpct_inline__ T op_erf(T x) {
     if constexpr (std::is_same_v<T, sycl::ext::oneapi::bfloat16>) {

ggml/src/ggml-sycl/element_wise.hpp

@@ -28,6 +28,39 @@ typed_data<T_Dst, T_Src> cast_data(ggml_tensor * dst) {
 
 const float GELU_QUICK_COEF = -1.702f;
 
+// Single-element activations, shared with the mat-vec kernels that fuse a GLU epilogue
+// (mmvq.cpp), so both apply the same formula.
+template <typename T> static __dpct_inline__ T op_tanh(T x) {
+    if constexpr (std::is_same_v<T, sycl::ext::oneapi::bfloat16>) {
+#if defined(__INTEL_LLVM_COMPILER) && (__INTEL_LLVM_COMPILER >= 20260000)
+        return sycl::ext::oneapi::experimental::tanh(x);
+#else
+        return static_cast<T>(sycl::tanh(static_cast<float>(x)));
+#endif
+    } else {
+        return sycl::tanh(x);
+    }
+}
+
+template <typename T> static __dpct_inline__ T op_gelu(T x) {
+    const T GELU_COEF_A    = static_cast<T>(0.044715f);
+    const T SQRT_2_OVER_PI = static_cast<T>(0.79788456080286535587989211986876f);
+    return static_cast<T>(0.5f) * x *
+           (static_cast<T>(1.0f) +
+            op_tanh(SQRT_2_OVER_PI * x * (static_cast<T>(1.0f) + GELU_COEF_A * x * x)));
+}
+
+template <typename T> static __dpct_inline__ T op_exp(T x) {
+    if constexpr (std::is_same_v<T, sycl::ext::oneapi::bfloat16>) {
+        return sycl::ext::oneapi::experimental::exp(x);
+    } else {
+        return sycl::exp(x);
+    }
+}
+
+template <typename T> static __dpct_inline__ T op_silu(T x) {
+    return x / (static_cast<T>(1.0f) + op_exp(-x));
+}
 
 void ggml_sycl_sqrt(ggml_backend_sycl_context & ctx, ggml_tensor * dst);
 

ggml/src/ggml-sycl/fusion.cpp

@@ -2,6 +2,61 @@
 
 #include <algorithm>
 
+// mul_mat(gate) + mul_mat(up) + GLU: graph shape and tensor properties only. Backend state
+// (weight layout, split buffers, DMMV) is checked by ggml_sycl_mul_mat_glu_mmvq_fused().
+static bool ggml_sycl_should_fuse_mul_mat_glu(const ggml_tensor * gate, const ggml_tensor * up,
+                                              const ggml_tensor * glu) {
+    // the fused epilogue implements these two; the rest fall back to the standalone GLU kernels
+    const ggml_glu_op glu_op = ggml_get_glu_op(glu);
+    if (glu_op != GGML_GLU_OP_SWIGLU && glu_op != GGML_GLU_OP_GEGLU) {
+        return false;
+    }
+
+    // the kernel always treats src[0] as the activated operand and src[1] as the multiplier
+    if (ggml_get_op_params_i32(glu, 1) /* swapped */) {
+        return false;
+    }
+
+    const ggml_tensor * wu  = up->src[0];
+    const ggml_tensor * wg  = gate->src[0];
+    const ggml_tensor * act = up->src[1];
+
+    // one set of block offsets and one quantized activation must serve both weights
+    if (wu->type != wg->type || !ggml_are_same_shape(wu, wg) || !ggml_are_same_stride(wu, wg)) {
+        return false;
+    }
+    if (act != gate->src[1]) {
+        return false;
+    }
+
+    // only q4_K has a fused reorder GEMV so far, and it walks whole super-blocks
+    if (wu->type != GGML_TYPE_Q4_K || wu->ne[0] % QK_K != 0) {
+        return false;
+    }
+
+    // one 2D reorder-layout matrix in, a plain column stride out: no broadcast or padding
+    if (!ggml_is_contiguous(wu) || !ggml_is_contiguous(wg) || !ggml_is_contiguous(act) ||
+        !ggml_is_contiguous(glu)) {
+        return false;
+    }
+    if (act->type != GGML_TYPE_F32 || glu->type != GGML_TYPE_F32) {
+        return false;
+    }
+    if (act->ne[2] != 1 || act->ne[3] != 1 || wu->ne[2] != 1 || wu->ne[3] != 1) {
+        return false;
+    }
+    // the kernel writes rows [0, wu->ne[1]) of each glu column, strided by glu->ne[0]
+    if (glu->ne[0] != wu->ne[1] || glu->ne[1] != act->ne[1]) {
+        return false;
+    }
+    // mat-vec only: one column per decoded token, up to the batch the reorder kernels cover
+    if (act->ne[1] > MMVQ_MAX_BATCH_SIZE) {
+        return false;
+    }
+
+    return true;
+}
+
 bool ggml_sycl_can_fuse(const ggml_cgraph * cgraph, int node_idx, std::initializer_list<enum ggml_op> ops,
                         std::initializer_list<enum ggml_unary_op> unary_ops) {
 #ifndef NDEBUG
@@ -13,6 +68,28 @@ bool ggml_sycl_can_fuse(const ggml_cgraph * cgraph, int node_idx, std::initializ
         return false;
     }
 
+    // gate and up are siblings, not a chain, so ggml_can_fuse cannot express this: use the
+    // subgraph form with the GLU as the only materialised output.
+    if (ops.size() == 3 && ops.begin()[0] == GGML_OP_MUL_MAT && ops.begin()[1] == GGML_OP_MUL_MAT &&
+        ops.begin()[2] == GGML_OP_GLU) {
+        if (!ggml_can_fuse_subgraph(cgraph, node_idx, ops, { node_idx + 2 })) {
+            return false;
+        }
+
+        const ggml_tensor * glu  = cgraph->nodes[node_idx + 2];
+        const ggml_tensor * gate = glu->src[0];
+        const ggml_tensor * up   = glu->src[1];
+
+        // don't assume which of the two mat-muls is the gate; infer it from the GLU's operands
+        const bool ok = (gate == cgraph->nodes[node_idx] && up == cgraph->nodes[node_idx + 1]) ||
+                        (gate == cgraph->nodes[node_idx + 1] && up == cgraph->nodes[node_idx]);
+        if (!ok) {
+            return false;
+        }
+
+        return ggml_sycl_should_fuse_mul_mat_glu(gate, up, glu);
+    }
+
     if (!ggml_can_fuse(cgraph, node_idx, ops)) {
         return false;
     }

ggml/src/ggml-sycl/gated_delta_net.cpp

@@ -14,9 +14,9 @@ void gated_delta_net_sycl(const float *     q,
                           const float *     beta,
                           const float *     curr_state,
                           float *           dst,
+                          float *           state,
                           int64_t           H,
                           int64_t           n_tokens,
-                          int64_t           n_seqs,
                           int64_t           sq1,
                           int64_t           sq2,
                           int64_t           sq3,
@@ -29,6 +29,7 @@ void gated_delta_net_sycl(const float *     q,
                           const sycl::uint3 neqk1_magic,
                           const sycl::uint3 rq3_magic,
                           float             scale,
+                          int64_t           state_slot_stride,
                           int               K) {
     auto           item_ct1 = sycl::ext::oneapi::this_work_item::get_nd_item<3>();
     const uint32_t h_idx    = item_ct1.get_group(2);
@@ -40,15 +41,12 @@ void gated_delta_net_sycl(const float *     q,
     const uint32_t iq1 = fastmodulo(h_idx, neqk1_magic);
     const uint32_t iq3 = fastdiv(sequence, rq3_magic);
 
-    const int64_t attn_score_elems = S_v * H * n_tokens * n_seqs;
     float *       attn_data        = dst;
-    float *       state            = dst + attn_score_elems;
 
     // input state holds s0 only [S_v, S_v, H, n_seqs] — seq stride is D = H * S_v * S_v.
     // output state layout (per-slot D * n_seqs) — same per-(seq,head) offset as before.
     const int64_t state_in_offset      = sequence * H * S_v * S_v + h_idx * S_v * S_v;
     const int64_t state_out_offset     = (sequence * H + h_idx) * S_v * S_v;
-    const int64_t state_size_per_token = S_v * S_v * H * n_seqs; // per-slot stride in output
     state += state_out_offset;
     curr_state += state_in_offset + col * S_v;
     attn_data += (sequence * n_tokens * H + h_idx) * S_v;
@@ -145,7 +143,7 @@ void gated_delta_net_sycl(const float *     q,
         if constexpr (keep_rs_t) {
             const int target_slot = (int) n_tokens - 1 - t;
             if (target_slot >= 0 && target_slot < K) {
-                float * curr_state = (dst + attn_score_elems) + target_slot * state_size_per_token + state_out_offset;
+                float * curr_state = state + target_slot * state_slot_stride;
 #pragma unroll
                 for (int r = 0; r < rows_per_lane; r++) {
                     const int i = r * warp_size + lane;
@@ -172,6 +170,7 @@ static void launch_gated_delta_net(const float *   q_d,
                                    const float *   b_d,
                                    const float *   s_d,
                                    float *         dst_d,
+                                   float *         state_d,
                                    int64_t         S_v,
                                    int64_t         H,
                                    int64_t         n_tokens,
@@ -188,6 +187,7 @@ static void launch_gated_delta_net(const float *   q_d,
                                    int64_t         neqk1,
                                    int64_t         rq3,
                                    float           scale,
+                                   int64_t         state_slot_stride,
                                    int             K,
                                    dpct::queue_ptr stream) {
     //TODO: Add chunked kernel for even faster pre-fill
@@ -206,9 +206,9 @@ static void launch_gated_delta_net(const float *   q_d,
                 constexpr int sv = 16;
                 stream->parallel_for(sycl::nd_range<3>(grid_dims * block_dims, block_dims),
                                      [=](sycl::nd_item<3> /*item_ct1*/) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
-                                         gated_delta_net_sycl<sv, KDA, keep_rs_t>(q_d, k_d, v_d, g_d, b_d, s_d, dst_d, H, n_tokens,
-                                                                       n_seqs, sq1, sq2, sq3, sv1, sv2, sv3, sb1, sb2,
-                                                                       sb3, neqk1_magic, rq3_magic, scale, K);
+                                         gated_delta_net_sycl<sv, KDA, keep_rs_t>(q_d, k_d, v_d, g_d, b_d, s_d, dst_d, state_d, H, n_tokens,
+                                                                       sq1, sq2, sq3, sv1, sv2, sv3, sb1, sb2,
+                                                                       sb3, neqk1_magic, rq3_magic, scale, state_slot_stride, K);
                                      });
             }
             break;
@@ -217,9 +217,9 @@ static void launch_gated_delta_net(const float *   q_d,
                 constexpr int sv = 32;
                 stream->parallel_for(sycl::nd_range<3>(grid_dims * block_dims, block_dims),
                                      [=](sycl::nd_item<3> /*item_ct1*/) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
-                                         gated_delta_net_sycl<sv, KDA, keep_rs_t>(q_d, k_d, v_d, g_d, b_d, s_d, dst_d, H, n_tokens,
-                                                                       n_seqs, sq1, sq2, sq3, sv1, sv2, sv3, sb1, sb2,
-                                                                       sb3, neqk1_magic, rq3_magic, scale, K);
+                                         gated_delta_net_sycl<sv, KDA, keep_rs_t>(q_d, k_d, v_d, g_d, b_d, s_d, dst_d, state_d, H, n_tokens,
+                                                                       sq1, sq2, sq3, sv1, sv2, sv3, sb1, sb2,
+                                                                       sb3, neqk1_magic, rq3_magic, scale, state_slot_stride, K);
                                      });
             }
             break;
@@ -229,8 +229,8 @@ static void launch_gated_delta_net(const float *   q_d,
                 stream->parallel_for(sycl::nd_range<3>(grid_dims * block_dims, block_dims),
                                         [=](sycl::nd_item<3> /*item_ct1*/) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
                                             gated_delta_net_sycl<sv, KDA, keep_rs_t>(
-                                                q_d, k_d, v_d, g_d, b_d, s_d, dst_d, H, n_tokens, n_seqs, sq1, sq2,
-                                                sq3, sv1, sv2, sv3, sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, K);
+                                                q_d, k_d, v_d, g_d, b_d, s_d, dst_d, state_d, H, n_tokens, sq1, sq2,
+                                                sq3, sv1, sv2, sv3, sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, state_slot_stride, K);
                                         });
             }
             break;
@@ -241,8 +241,8 @@ static void launch_gated_delta_net(const float *   q_d,
                 stream->parallel_for(sycl::nd_range<3>(grid_dims * block_dims, block_dims),
                                         [=](sycl::nd_item<3> /*item_ct1*/) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
                                             gated_delta_net_sycl<sv, KDA, keep_rs_t>(
-                                                q_d, k_d, v_d, g_d, b_d, s_d, dst_d, H, n_tokens, n_seqs, sq1, sq2,
-                                                sq3, sv1, sv2, sv3, sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, K);
+                                                q_d, k_d, v_d, g_d, b_d, s_d, dst_d, state_d, H, n_tokens, sq1, sq2,
+                                                sq3, sv1, sv2, sv3, sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, state_slot_stride, K);
                                         });
             }
             break;
@@ -253,7 +253,8 @@ static void launch_gated_delta_net(const float *   q_d,
     }
 }
 
-void ggml_sycl_op_gated_delta_net(ggml_backend_sycl_context & ctx, ggml_tensor * dst) {
+static void ggml_sycl_op_gated_delta_net_impl(ggml_backend_sycl_context & ctx, ggml_tensor * dst,
+                                              const ggml_sycl_gated_delta_net_fused_cache * cache) {
     ggml_tensor * src_q     = dst->src[0];
     ggml_tensor * src_k     = dst->src[1];
     ggml_tensor * src_v     = dst->src[2];
@@ -318,30 +319,48 @@ void ggml_sycl_op_gated_delta_net(ggml_backend_sycl_context & ctx, ggml_tensor *
     const int K = ggml_get_op_params_i32(dst, 0);
     const bool keep_rs = K > 1;
 
+    // recurrent state -> dst tail (after attention scores), or the cache when fusing
+    float * state_d           = dst_d + S_v * H * n_tokens * n_seqs;
+    int64_t state_slot_stride = S_v * S_v * H * n_seqs;
+    if (cache != nullptr) {
+        state_d           = cache->data;
+        state_slot_stride = cache->slot_stride;
+    }
+
     if (kda) {
         if (keep_rs) {
-            launch_gated_delta_net<true, true>(q_d, k_d, v_d, g_d, b_d, s_d, dst_d,
+            launch_gated_delta_net<true, true>(q_d, k_d, v_d, g_d, b_d, s_d, dst_d, state_d,
                 S_v, H, n_tokens, n_seqs, sq1, sq2, sq3, sv1, sv2, sv3,
-                sb1, sb2, sb3, neqk1, rq3, scale, K, stream);
+                sb1, sb2, sb3, neqk1, rq3, scale, state_slot_stride, K, stream);
         } else {
-            launch_gated_delta_net<true, false>(q_d, k_d, v_d, g_d, b_d, s_d, dst_d,
+            launch_gated_delta_net<true, false>(q_d, k_d, v_d, g_d, b_d, s_d, dst_d, state_d,
                 S_v, H, n_tokens, n_seqs, sq1, sq2, sq3, sv1, sv2, sv3,
-                sb1, sb2, sb3, neqk1, rq3, scale, K, stream);
+                sb1, sb2, sb3, neqk1, rq3, scale, state_slot_stride, K, stream);
         }
     } else {
         if (keep_rs) {
-            launch_gated_delta_net<false, true>(q_d, k_d, v_d, g_d, b_d, s_d, dst_d,
+            launch_gated_delta_net<false, true>(q_d, k_d, v_d, g_d, b_d, s_d, dst_d, state_d,
                 S_v, H, n_tokens, n_seqs, sq1, sq2, sq3, sv1, sv2, sv3,
-                sb1, sb2, sb3, neqk1, rq3, scale, K, stream);
+                sb1, sb2, sb3, neqk1, rq3, scale, state_slot_stride, K, stream);
         } else {
-            launch_gated_delta_net<false, false>(q_d, k_d, v_d, g_d, b_d, s_d, dst_d,
+            launch_gated_delta_net<false, false>(q_d, k_d, v_d, g_d, b_d, s_d, dst_d, state_d,
                 S_v, H, n_tokens, n_seqs, sq1, sq2, sq3, sv1, sv2, sv3,
-                sb1, sb2, sb3, neqk1, rq3, scale, K, stream);
+                sb1, sb2, sb3, neqk1, rq3, scale, state_slot_stride, K, stream);
         }
     }
 }
 
+void ggml_sycl_op_gated_delta_net(ggml_backend_sycl_context & ctx, ggml_tensor * dst) {
+    ggml_sycl_op_gated_delta_net_impl(ctx, dst, nullptr);
+}
+
 void ggml_sycl_gated_delta_net(ggml_backend_sycl_context & ctx, ggml_tensor * dst) {
     scope_op_debug_print scope_dbg_print(__func__, dst, /*num_src=*/6);
     ggml_sycl_op_gated_delta_net(ctx, dst);
 }
+
+void ggml_sycl_op_gated_delta_net_fused_cache(ggml_backend_sycl_context & ctx, ggml_tensor * dst,
+                                              ggml_sycl_gated_delta_net_fused_cache cache) {
+    scope_op_debug_print scope_dbg_print(__func__, dst, /*num_src=*/6);
+    ggml_sycl_op_gated_delta_net_impl(ctx, dst, &cache);
+}

ggml/src/ggml-sycl/gated_delta_net.hpp

@@ -5,5 +5,15 @@
 #include "common.hpp"
 #include "ggml.h"
 
+// fused-kernel recurrent-state output; strides in elements (per-seq stride is always D, set in-kernel)
+struct ggml_sycl_gated_delta_net_fused_cache {
+    float * data;        // rollback slot 0
+    int64_t slot_stride; // between rollback slots (0 when K==1)
+};
+
 void ggml_sycl_op_gated_delta_net(ggml_backend_sycl_context & ctx, ggml_tensor * dst);
 void ggml_sycl_gated_delta_net(ggml_backend_sycl_context & ctx, ggml_tensor * dst);
+
+// same op, but writes the snapshot(s) into the cache instead of dst (see ggml_sycl_try_gdn_cache_fusion)
+void ggml_sycl_op_gated_delta_net_fused_cache(ggml_backend_sycl_context & ctx, ggml_tensor * dst,
+                                              ggml_sycl_gated_delta_net_fused_cache cache);

ggml/src/ggml-sycl/ggml-sycl.cpp

@@ -11,6 +11,7 @@
 //
 
 #include <algorithm>
+#include <array>
 #include <assert.h>
 #include <atomic>
 #include <cinttypes>
@@ -4560,6 +4561,66 @@ static void ggml_sycl_mul_mat(ggml_backend_sycl_context & ctx, const ggml_tensor
     }
 }
 
+// Fused dense-FFN mat-vec for the {mul_mat(gate), mul_mat(up), GLU} subgraph at node_idx.
+// Returns false if it declined, in which case the caller runs the three nodes normally.
+static bool ggml_sycl_mul_mat_glu_mmvq_fused(ggml_backend_sycl_context & ctx, ggml_cgraph * cgraph, int node_idx) {
+    if (!ggml_sycl_can_fuse(cgraph, node_idx, { GGML_OP_MUL_MAT, GGML_OP_MUL_MAT, GGML_OP_GLU }, {})) {
+        return false;
+    }
+
+    ggml_tensor *       glu  = cgraph->nodes[node_idx + 2];
+    ggml_tensor *       gate = glu->src[0];
+    ggml_tensor *       up   = glu->src[1];
+    const ggml_tensor * wu   = up->src[0];
+    const ggml_tensor * wg   = gate->src[0];
+    const ggml_tensor * act  = up->src[1];
+
+    // this writes glu->data directly rather than the per-device row slices that
+    // ggml_sycl_op_mul_mat() stitches back together, so it cannot serve split weights
+    if (ggml_backend_buffer_is_sycl_split(wu->buffer) || ggml_backend_buffer_is_sycl_split(wg->buffer)) {
+        return false;
+    }
+
+    // with DMMV prioritised the unfused path would not have gone through mmvq at all
+    if (g_ggml_sycl_prioritize_dmmv) {
+        return false;
+    }
+
+    // install the reorder (SoA) layout the fused kernel needs, as the unfused mmvq path would;
+    // a no-op once done. after the bail checks so a declined op does not pay for it.
+    opt_for_reorder(&ctx, wu, act, up, mul_mat_algo::MMVQ);
+    opt_for_reorder(&ctx, wg, act, gate, mul_mat_algo::MMVQ);
+
+    const auto * extra_u = static_cast<const ggml_tensor_extra_gpu *>(wu->extra);
+    const auto * extra_g = static_cast<const ggml_tensor_extra_gpu *>(wg->extra);
+    if (!extra_u || !extra_g || !extra_u->optimized_feature.reorder || !extra_g->optimized_feature.reorder) {
+        return false;
+    }
+
+    // log the up mat-mul: glu's own srcs are the two intermediates the fusion never materialises
+    scope_op_debug_print scope_dbg_print(__func__, up, /*num_src=*/2, " : fused with gate + GLU");
+
+    const int64_t ne00 = wu->ne[0];
+    const int64_t ne11 = act->ne[1];
+
+    const queue_ptr stream           = ctx.stream();
+    const int       src1_padded_cols = GGML_PAD((int) ne00, MATRIX_ROW_PADDING);
+
+    // one activation, quantized once and fully consumed into src1_ddq before the GEMV on this
+    // in-order queue, so glu->data aliasing the dead activation needs no memory-range check
+    ggml_sycl_pool_alloc<char> src1_q8_alloc(ctx.pool(),
+                                             (size_t) ne11 * src1_padded_cols * sizeof(block_q8_1) / QK8_1);
+    char *                     src1_ddq = src1_q8_alloc.get();
+
+    quantize_row_q8_1_sycl<quantize_and_reorder_q8_1_soa>((const float *) act->data, src1_ddq, (int) ne00, (int) ne11,
+                                                          src1_padded_cols, stream);
+
+    return ggml_sycl_mul_mat_vec_q_glu_reorder(wu->type, ggml_get_glu_op(glu), wu->data, wg->data, src1_ddq,
+                                               (float *) glu->data, (int) ne00, (int) wu->ne[1], (int) ne11,
+                                               /*stride_col_y_bytes=*/src1_padded_cols * (int) sizeof(block_q8_1) /
+                                                   QK8_1,
+                                               /*stride_col_dst=*/(int) glu->ne[0], stream);
+}
 
 __dpct_inline__ static void k_copy_src1_to_contiguous(
     const char *__restrict__ src1_original, char *__restrict__ src1_contiguous,
@@ -5464,12 +5525,90 @@ catch (sycl::exception const &exc) {
   std::exit(1);
 }
 
+static bool ggml_sycl_is_view_or_noop(const ggml_tensor * t) {
+    return ggml_is_empty(t) || t->op == GGML_OP_RESHAPE || t->op == GGML_OP_TRANSPOSE ||
+           t->op == GGML_OP_VIEW || t->op == GGML_OP_PERMUTE || t->op == GGML_OP_NONE;
+}
+
+// match gated_delta_net + the strided cpy that scatters its state snapshots into the cache
+// (slot i -> rollback group i, slot 0 newest), so the kernel can write them and skip the cpy.
+// returns the number of following nodes to skip (0 = no fusion)
+// ported from ggml_cuda_try_gdn_cache_fusion - pure graph inspection, backend-agnostic
+static int ggml_sycl_try_gdn_cache_fusion(const ggml_cgraph * cgraph, int node_idx,
+                                          ggml_sycl_gated_delta_net_fused_cache & fused_state_cpy) {
+    if (!g_ggml_sycl_enable_fusion) {
+        return 0;
+    }
+
+    const ggml_tensor * gdn = cgraph->nodes[node_idx];
+    // the kernel skips the snapshot tail, so the gdn output must not be a graph output, and the cpy
+    // found below is taken to be its only reader, as it is in every graph that builds this op
+    if (gdn->op != GGML_OP_GATED_DELTA_NET || gdn->type != GGML_TYPE_F32 ||
+        (gdn->flags & GGML_TENSOR_FLAG_OUTPUT)) {
+        return 0;
+    }
+
+    const ggml_tensor * src_v     = gdn->src[2];
+    const int64_t       S_v       = src_v->ne[0];
+    const int64_t       H         = src_v->ne[1];
+    const int64_t       n_tokens  = src_v->ne[2];
+    const int64_t       n_seqs    = src_v->ne[3];
+    const int64_t       D         = S_v * S_v * H;
+    const int64_t       K         = ggml_get_op_params_i32(gdn, 0); // snapshot slot count
+    const int64_t       n_written = std::min<int64_t>(n_tokens, K); // newest n_written slots are written
+
+    // snapshot tail starts right after the attention scores
+    const size_t tail_off = ggml_row_size(GGML_TYPE_F32, S_v * H * n_tokens * n_seqs);
+
+    // the cpy must be the first node the compute loop below runs, so nothing can read the cache first.
+    // skip exactly what that loop skips: views, no-ops, and nodes the graph does not compute.
+    const ggml_tensor * cpy  = nullptr;
+    int                 skip = 0;
+    for (int j = node_idx + 1; j < cgraph->n_nodes && cpy == nullptr; ++j) {
+        const ggml_tensor * n = cgraph->nodes[j];
+        if (ggml_sycl_is_view_or_noop(n) || (n->flags & GGML_TENSOR_FLAG_COMPUTE) == 0) {
+            continue;
+        }
+        if (n->op != GGML_OP_CPY || (n->flags & GGML_TENSOR_FLAG_OUTPUT)) {
+            return 0;
+        }
+        cpy  = n;
+        skip = j - node_idx;
+    }
+    if (cpy == nullptr) {
+        return 0;
+    }
+
+    const ggml_tensor * src = cpy->src[0]; // view of the gdn snapshot tail
+    const ggml_tensor * dst = cpy->src[1]; // cache view the kernel writes to
+
+    // src must be this gdn's snapshot tail (contiguous, at the tail offset)
+    if (src->op != GGML_OP_VIEW || src->view_src != gdn || src->view_offs != tail_off ||
+        !ggml_is_contiguous(src)) {
+        return 0;
+    }
+
+    // dst is the [D, n_seqs, n_written] cache view, with the per-seq stride D that the kernel assumes.
+    // ggml_cpy pins src to the same element count, so src needs no shape check of its own.
+    const std::array<int64_t, GGML_MAX_DIMS> expected_ne = { D, n_seqs, n_written, 1 };
+    if (dst->op != GGML_OP_VIEW || dst->type != GGML_TYPE_F32 || dst->data == nullptr ||
+        !std::equal(expected_ne.begin(), expected_ne.end(), dst->ne) ||
+        dst->nb[0] != ggml_type_size(GGML_TYPE_F32) ||
+        dst->nb[1] != (size_t) ggml_row_size(GGML_TYPE_F32, D)) {
+        return 0;
+    }
+
+    fused_state_cpy.data        = (float *) dst->data; // rollback group 0 (newest)
+    fused_state_cpy.slot_stride = K > 1 ? (int64_t) (dst->nb[2] / sizeof(float)) : 0;
+    return skip;
+}
+
 static void ggml_backend_sycl_graph_compute_impl(ggml_backend_sycl_context * sycl_ctx, ggml_cgraph * cgraph) {
     ggml_sycl_set_main_device(sycl_ctx->device);
 
     for (int i = 0; i < cgraph->n_nodes; i++) {
         ggml_tensor * node = cgraph->nodes[i];
-        if (ggml_is_empty(node) || node->op == GGML_OP_RESHAPE || node->op == GGML_OP_TRANSPOSE || node->op == GGML_OP_VIEW || node->op == GGML_OP_PERMUTE || node->op == GGML_OP_NONE) {
+        if (ggml_sycl_is_view_or_noop(node)) {
             continue;
         }
         if ((node->flags & GGML_TENSOR_FLAG_COMPUTE) == 0) {
@@ -5489,6 +5628,16 @@ static void ggml_backend_sycl_graph_compute_impl(ggml_backend_sycl_context * syc
             }
         }
 #endif
+        // gated_delta_net -> cpy: scatter recurrent-state snapshots into the cache
+        if (node->op == GGML_OP_GATED_DELTA_NET) {
+            ggml_sycl_gated_delta_net_fused_cache fused_state_cpy;
+            const int gdn_nodes_to_skip = ggml_sycl_try_gdn_cache_fusion(cgraph, i, fused_state_cpy);
+            if (gdn_nodes_to_skip > 0) {
+                ggml_sycl_op_gated_delta_net_fused_cache(*sycl_ctx, node, fused_state_cpy);
+                i += gdn_nodes_to_skip;
+                continue;
+            }
+        }
         if (node->op == GGML_OP_RMS_NORM &&
             ggml_sycl_can_fuse(cgraph, i, { GGML_OP_RMS_NORM, GGML_OP_MUL }, {})) {
             ggml_sycl_op_rms_norm_fused(*sycl_ctx, node, cgraph->nodes[i + 1]);
@@ -5502,6 +5651,11 @@ static void ggml_backend_sycl_graph_compute_impl(ggml_backend_sycl_context * syc
             continue;
         }
 
+        if (node->op == GGML_OP_MUL_MAT && ggml_sycl_mul_mat_glu_mmvq_fused(*sycl_ctx, cgraph, i)) {
+            i += 2;
+            continue;
+        }
+
         bool ok = ggml_sycl_compute_forward(*sycl_ctx, node);
         if (!ok) {
             GGML_LOG_ERROR("%s: error: op not supported %s (%s)\n", __func__, node->name, ggml_op_name(node->op));

ggml/src/ggml-sycl/mmvq.cpp

@@ -2,6 +2,7 @@
 
 #include "ggml.h"
 #include "common.hpp"
+#include "element_wise.hpp"
 #include "quants.hpp"
 #include "vecdotq.hpp"
 
@@ -56,11 +57,13 @@ static void mul_mat_vec_q_reorder(const void * __restrict__ vx, const void * __r
     }
 }
 
-template <typename reorder_vec_dot_q_sycl, int ncols_dst>
-static void mul_mat_vec_q_reorder_ncols(const void * __restrict__ vx, const void * __restrict__ vy,
-                                        float * __restrict__ dst, const int ncols, const int nrows,
-                                        const int stride_col_y_bytes, const int stride_col_dst,
-                                        const sycl::nd_item<3> & nd_item) {
+// With has_fusion, `vgate` is a second weight matrix sharing vx's shape, stride and reorder
+// layout: one pass computes both row dot products and the epilogue writes glu(gate, up).
+template <typename reorder_vec_dot_q_sycl, int ncols_dst, bool has_fusion = false>
+static void mul_mat_vec_q_reorder_ncols(const void * __restrict__ vx, const void * __restrict__ vgate,
+                                        const void * __restrict__ vy, float * __restrict__ dst, const int ncols,
+                                        const int nrows, const int stride_col_y_bytes, const int stride_col_dst,
+                                        const ggml_glu_op glu_op, const sycl::nd_item<3> & nd_item) {
     using block_type   = ggml_sycl_reordered::block_q_t<reorder_vec_dot_q_sycl::gtype>;
     using block_traits = typename block_type::traits;
 
@@ -70,6 +73,8 @@ static void mul_mat_vec_q_reorder_ncols(const void * __restrict__ vx, const void
     const int  sg_id        = sg.get_group_linear_id();
     const int  row          = workgroup_id * sg_range + sg_id;
 
+    // row is sub-group uniform, so this retires whole sub-groups and the collectives below
+    // stay convergent
     if (row >= nrows) {
         return;
     }
@@ -82,10 +87,15 @@ static void mul_mat_vec_q_reorder_ncols(const void * __restrict__ vx, const void
     static_assert(blocks_per_subgroup > 0);
     static_assert(block_elements_per_subgroup > 0);
 
-    float partial_sum[ncols_dst] = {0.0f};
+    float partial_sum[ncols_dst] = { 0.0f };
+    // sized 1 rather than 0 when unused: zero-length arrays are not standard C++, and the
+    // array is dead and eliminated in that case
+    [[maybe_unused]] float partial_gate[has_fusion ? ncols_dst : 1] = { 0.0f };
     for (int i = sg.get_local_linear_id() / block_elements_per_subgroup; i < blocks_per_row; i += blocks_per_subgroup) {
         const int ibx = row * blocks_per_row + i;
 
+        // the offsets depend only on the block index and the matrix shape, never on the base
+        // pointer, which is what lets vgate reuse them
         const auto bx_offset = block_type::get_block_offset(ibx, nblocks);
         const auto d_offset  = block_type::get_d_offset(nrows, ncols, ibx);
         const int  iby       = i * block_type::block_to_q8_1_ratio();
@@ -96,11 +106,16 @@ static void mul_mat_vec_q_reorder_ncols(const void * __restrict__ vx, const void
 
 #pragma unroll
             for (int j = 0; j < ncols_dst; ++j) {
-                const char       * vy_j           = (const char *)vy + j * stride_col_y_bytes;
-                const int8_t     * q8_1_quant_ptr = (const int8_t *)vy_j + iby * QK8_1;
-                const sycl::half2* q8_1_ds_ptr    = (const sycl::half2 *)(vy_j + ncols + iby * sizeof(sycl::half2));
+                const char        * vy_j           = (const char *) vy + j * stride_col_y_bytes;
+                const int8_t      * q8_1_quant_ptr = (const int8_t *) vy_j + iby * QK8_1;
+                const sycl::half2 * q8_1_ds_ptr    = (const sycl::half2 *) (vy_j + ncols + iby * sizeof(sycl::half2));
 
                 partial_sum[j] += reorder_vec_dot_q_sycl()(vx, bx_offset, d_offset, q8_1_quant_ptr, q8_1_ds_ptr, iqs);
+
+                if constexpr (has_fusion) {
+                    partial_gate[j] +=
+                        reorder_vec_dot_q_sycl()(vgate, bx_offset, d_offset, q8_1_quant_ptr, q8_1_ds_ptr, iqs);
+                }
             }
         }
     }
@@ -109,6 +124,13 @@ static void mul_mat_vec_q_reorder_ncols(const void * __restrict__ vx, const void
     for (int j = 0; j < ncols_dst; ++j) {
         float sum = sycl::reduce_over_group(nd_item.get_sub_group(), partial_sum[j], std::plus<>());
 
+        if constexpr (has_fusion) {
+            const float gate = sycl::reduce_over_group(nd_item.get_sub_group(), partial_gate[j], std::plus<>());
+
+            // uniform across the launch; the launcher only instantiates SWIGLU and GEGLU
+            sum *= glu_op == GGML_GLU_OP_SWIGLU ? op_silu(gate) : op_gelu(gate);
+        }
+
         if (sg.leader()) {
             dst[j * stride_col_dst + row] = sum;
         }
@@ -691,7 +713,8 @@ static void reorder_mul_mat_vec_q4_0_q8_1_sycl_ncols(
         cgh.parallel_for(sycl::nd_range<3>(block_nums * block_dims, block_dims),
                          [=](sycl::nd_item<3> nd_item) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
                              mul_mat_vec_q_reorder_ncols<reorder_vec_dot_q_sycl<GGML_TYPE_Q4_0>, ncols_dst>(
-                                 vx, vy, dst, ncols, nrows, stride_col_y_bytes, stride_col_dst, nd_item);
+                                 vx, /*vgate=*/ nullptr, vy, dst, ncols, nrows, stride_col_y_bytes, stride_col_dst,
+                                 /*glu_op=*/ GGML_GLU_OP_SWIGLU, nd_item);
                          });
     });
 }
@@ -1108,7 +1131,8 @@ static void reorder_mul_mat_vec_q8_0_q8_1_sycl_ncols(
         cgh.parallel_for(sycl::nd_range<3>(block_nums * block_dims, block_dims),
                          [=](sycl::nd_item<3> nd_item) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
                              mul_mat_vec_q_reorder_ncols<reorder_vec_dot_q_sycl<GGML_TYPE_Q8_0>, ncols_dst>(
-                                 vx, vy, dst, ncols, nrows, stride_col_y_bytes, stride_col_dst, nd_item);
+                                 vx, /*vgate=*/ nullptr, vy, dst, ncols, nrows, stride_col_y_bytes, stride_col_dst,
+                                 /*glu_op=*/ GGML_GLU_OP_SWIGLU, nd_item);
                          });
     });
 }
@@ -1436,7 +1460,8 @@ static void reorder_mul_mat_vec_q3_k_q8_1_sycl_ncols(
         cgh.parallel_for(sycl::nd_range<3>(block_nums * block_dims, block_dims),
                          [=](sycl::nd_item<3> nd_item) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
                              mul_mat_vec_q_reorder_ncols<reorder_vec_dot_q_sycl<GGML_TYPE_Q3_K>, ncols_dst>(
-                                 vx, vy, dst, ncols, nrows, stride_col_y_bytes, stride_col_dst, nd_item);
+                                 vx, /*vgate=*/ nullptr, vy, dst, ncols, nrows, stride_col_y_bytes, stride_col_dst,
+                                 /*glu_op=*/ GGML_GLU_OP_SWIGLU, nd_item);
                          });
     });
 }
@@ -1604,7 +1629,8 @@ static void reorder_mul_mat_vec_q4_k_q8_1_sycl_ncols(
         cgh.parallel_for(sycl::nd_range<3>(block_nums * block_dims, block_dims),
                          [=](sycl::nd_item<3> nd_item) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
                              mul_mat_vec_q_reorder_ncols<reorder_vec_dot_q_sycl<GGML_TYPE_Q4_K>, ncols_dst>(
-                                 vx, vy, dst, ncols, nrows, stride_col_y_bytes, stride_col_dst, nd_item);
+                                 vx, /*vgate=*/ nullptr, vy, dst, ncols, nrows, stride_col_y_bytes, stride_col_dst,
+                                 /*glu_op=*/ GGML_GLU_OP_SWIGLU, nd_item);
                          });
     });
 }
@@ -1731,7 +1757,8 @@ static void reorder_mul_mat_vec_q5_k_q8_1_sycl_ncols(
         cgh.parallel_for(sycl::nd_range<3>(block_nums * block_dims, block_dims),
                          [=](sycl::nd_item<3> nd_item) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
                              mul_mat_vec_q_reorder_ncols<reorder_vec_dot_q_sycl<GGML_TYPE_Q5_K>, ncols_dst>(
-                                 vx, vy, dst, ncols, nrows, stride_col_y_bytes, stride_col_dst, nd_item);
+                                 vx, /*vgate=*/ nullptr, vy, dst, ncols, nrows, stride_col_y_bytes, stride_col_dst,
+                                 /*glu_op=*/ GGML_GLU_OP_SWIGLU, nd_item);
                          });
     });
 }
@@ -1789,7 +1816,8 @@ static void reorder_mul_mat_vec_q6_k_q8_1_sycl_ncols(
         cgh.parallel_for(sycl::nd_range<3>(block_nums * block_dims, block_dims),
                          [=](sycl::nd_item<3> nd_item) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
                              mul_mat_vec_q_reorder_ncols<reorder_vec_dot_q_sycl<GGML_TYPE_Q6_K>, ncols_dst>(
-                                 vx, vy, dst, ncols, nrows, stride_col_y_bytes, stride_col_dst, nd_item);
+                                 vx, /*vgate=*/ nullptr, vy, dst, ncols, nrows, stride_col_y_bytes, stride_col_dst,
+                                 /*glu_op=*/ GGML_GLU_OP_SWIGLU, nd_item);
                          });
     });
 }
@@ -2736,3 +2764,77 @@ bool ggml_sycl_mul_mat_vec_q_id_reorder(
             return false;
     }
 }
+
+template <typename reorder_vec_dot_q_sycl, int ncols_dst>
+static void launch_mul_mat_vec_q_reorder_glu(const void * vx, const void * vgate, const void * vy, float * dst,
+                                             const int ncols, const int nrows, const int stride_col_y_bytes,
+                                             const int stride_col_dst, const ggml_glu_op glu_op,
+                                             dpct::queue_ptr stream) {
+    GGML_ASSERT(ncols % QK_K == 0);
+
+    constexpr size_t num_subgroups = WARP_SIZE;
+
+    const int            block_num_y = ceil_div(nrows, GGML_SYCL_MMV_Y * (int) num_subgroups);
+    const sycl::range<3> block_nums(1, 1, block_num_y);
+    const sycl::range<3> block_dims(1, GGML_SYCL_MMV_Y, num_subgroups * WARP_SIZE);
+
+    stream->submit([&](sycl::handler & cgh) {
+        cgh.parallel_for(sycl::nd_range<3>(block_nums * block_dims, block_dims),
+                         [=](sycl::nd_item<3> nd_item) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
+                             mul_mat_vec_q_reorder_ncols<reorder_vec_dot_q_sycl, ncols_dst, /*has_fusion=*/ true>(
+                                 vx, vgate, vy, dst, ncols, nrows, stride_col_y_bytes, stride_col_dst, glu_op,
+                                 nd_item);
+                         });
+    });
+}
+
+bool ggml_sycl_mul_mat_vec_q_glu_reorder(enum ggml_type src0_type, enum ggml_glu_op glu_op, const void * vx,
+                                         const void * vgate, const void * vy, float * dst, int ncols, int nrows,
+                                         int ncols_dst, int stride_col_y_bytes, int stride_col_dst,
+                                         dpct::queue_ptr stream) {
+    if (src0_type != GGML_TYPE_Q4_K) {
+        return false;
+    }
+    if (glu_op != GGML_GLU_OP_SWIGLU && glu_op != GGML_GLU_OP_GEGLU) {
+        return false;
+    }
+
+    using vec_dot = reorder_vec_dot_q_sycl<GGML_TYPE_Q4_K>;
+
+    switch (ncols_dst) {
+        case 1:
+            launch_mul_mat_vec_q_reorder_glu<vec_dot, 1>(vx, vgate, vy, dst, ncols, nrows, stride_col_y_bytes,
+                                                         stride_col_dst, glu_op, stream);
+            return true;
+        case 2:
+            launch_mul_mat_vec_q_reorder_glu<vec_dot, 2>(vx, vgate, vy, dst, ncols, nrows, stride_col_y_bytes,
+                                                         stride_col_dst, glu_op, stream);
+            return true;
+        case 3:
+            launch_mul_mat_vec_q_reorder_glu<vec_dot, 3>(vx, vgate, vy, dst, ncols, nrows, stride_col_y_bytes,
+                                                         stride_col_dst, glu_op, stream);
+            return true;
+        case 4:
+            launch_mul_mat_vec_q_reorder_glu<vec_dot, 4>(vx, vgate, vy, dst, ncols, nrows, stride_col_y_bytes,
+                                                         stride_col_dst, glu_op, stream);
+            return true;
+        case 5:
+            launch_mul_mat_vec_q_reorder_glu<vec_dot, 5>(vx, vgate, vy, dst, ncols, nrows, stride_col_y_bytes,
+                                                         stride_col_dst, glu_op, stream);
+            return true;
+        case 6:
+            launch_mul_mat_vec_q_reorder_glu<vec_dot, 6>(vx, vgate, vy, dst, ncols, nrows, stride_col_y_bytes,
+                                                         stride_col_dst, glu_op, stream);
+            return true;
+        case 7:
+            launch_mul_mat_vec_q_reorder_glu<vec_dot, 7>(vx, vgate, vy, dst, ncols, nrows, stride_col_y_bytes,
+                                                         stride_col_dst, glu_op, stream);
+            return true;
+        case 8:
+            launch_mul_mat_vec_q_reorder_glu<vec_dot, 8>(vx, vgate, vy, dst, ncols, nrows, stride_col_y_bytes,
+                                                         stride_col_dst, glu_op, stream);
+            return true;
+        default:
+            return false;
+    }
+}

ggml/src/ggml-sycl/mmvq.hpp

@@ -57,4 +57,20 @@ bool ggml_sycl_mul_mat_vec_q_id_reorder(
     size_t             src1_row_stride,
     dpct::queue_ptr    stream);
 
+// Fused dense-FFN GEMV: writes glu(gate . y, up . y) instead of the two mat-vec results.
+// vx / vgate must share shape, stride and reorder layout. Returns false if unhandled.
+bool ggml_sycl_mul_mat_vec_q_glu_reorder(
+    enum ggml_type     src0_type,
+    enum ggml_glu_op   glu_op,
+    const void *       vx,
+    const void *       vgate,
+    const void *       vy,
+    float *            dst,
+    int                ncols,                // K, shared by both weights
+    int                nrows,                // output rows, i.e. weight ne[1]
+    int                ncols_dst,            // activation columns, 1..MMVQ_MAX_BATCH_SIZE
+    int                stride_col_y_bytes,   // bytes between activation columns in vy
+    int                stride_col_dst,       // floats between output columns in dst
+    dpct::queue_ptr    stream);
+
 #endif // GGML_SYCL_MMVQ_HPP

ggml/src/ggml-sycl/ssm_scan.cpp

@@ -10,6 +10,7 @@ static void ssm_scan_f32_group(
         const int src2_nb1, const int src2_nb2, const int src3_nb1,
         const int src4_nb2, const int src4_nb3, const int src5_nb2, const int src5_nb3,
         const int64_t s_off, const int64_t n_head, const int64_t d_head, const int64_t n_group, const int64_t n_tok,
+        const int64_t K,
         const sycl::nd_item<2> & item) {
 
     const int lane     = item.get_local_id(1) % WARP_SIZE;
@@ -64,6 +65,15 @@ static void ssm_scan_f32_group(
         if (lane == 0) {
             y_warp[i * stride_y] = state_sum;
         }
+
+        const int64_t slot = n_tok - 1 - i;
+        if (K > 1 && slot > 0 && slot < K) {
+            float * s_snapshot_warp = (float *) ((char *) dst + s_off + (slot * item.get_group_range(0) + seq_idx) * src0_nb3 + head_idx * src0_nb2 + head_off * d_state);
+#pragma unroll
+            for (int j = 0; j < c_factor; j++) {
+                s_snapshot_warp[WARP_SIZE * j + lane] = state[j];
+            }
+        }
     }
 
 #pragma unroll
@@ -79,6 +89,7 @@ static void ssm_scan_f32_sycl(
         const int src2_nb2, const int src3_nb1, const int src4_nb2, const int src4_nb3, const int src5_nb2,
         const int src5_nb3, const int64_t s_off, const int64_t d_state, const int64_t head_dim,
         const int64_t n_head, const int64_t n_group, const int64_t n_tok, const int64_t n_seq,
+        const int64_t K,
         dpct::queue_ptr stream) {
 
     // NOTE: if you change conditions here, be sure to update the corresponding supports_op condition!
@@ -94,7 +105,7 @@ static void ssm_scan_f32_sycl(
                 ssm_scan_f32_group<128 / WARP_SIZE, 128>(
                     src0, src1, src2, src3, src4, src5, src6, dst,
                     src0_nb2, src0_nb3, src1_nb2, src1_nb3, src2_nb1, src2_nb2, src3_nb1,
-                    src4_nb2, src4_nb3, src5_nb2, src5_nb3, s_off, n_head, head_dim, n_group, n_tok, item);
+                    src4_nb2, src4_nb3, src5_nb2, src5_nb3, s_off, n_head, head_dim, n_group, n_tok, K, item);
             });
     } else if (d_state == 256) {
         constexpr int threads   = 256;
@@ -107,7 +118,7 @@ static void ssm_scan_f32_sycl(
                 ssm_scan_f32_group<256 / WARP_SIZE, 256>(
                     src0, src1, src2, src3, src4, src5, src6, dst,
                     src0_nb2, src0_nb3, src1_nb2, src1_nb3, src2_nb1, src2_nb2, src3_nb1,
-                    src4_nb2, src4_nb3, src5_nb2, src5_nb3, s_off, n_head, head_dim, n_group, n_tok, item);
+                    src4_nb2, src4_nb3, src5_nb2, src5_nb3, s_off, n_head, head_dim, n_group, n_tok, K, item);
             });
     } else {
         GGML_ABORT("ssm_scan: unsupported d_state (must be 128 or 256)");
@@ -133,9 +144,12 @@ inline void ggml_sycl_op_ssm_scan(ggml_backend_sycl_context & ctx, ggml_tensor *
     const int64_t ng  = src4->ne[1];
     const int64_t n_t = src1->ne[2];
     const int64_t n_s = src1->ne[3];
+    const int64_t K   = ggml_get_op_params_i32(dst, 0);
     const int64_t s_off = ggml_nelements(src1) * sizeof(float);
 
-    GGML_ASSERT(ggml_nelements(src1) + nc * nr * nh * n_s == ggml_nelements(dst));
+    GGML_ASSERT(K >= 1);
+    GGML_ASSERT(ggml_nelements(src1) + K * nc * nr * nh * n_s == ggml_nelements(dst));
+    GGML_ASSERT(src3->ne[0] == 1 || K == 1);
 
     dpct::queue_ptr stream = ctx.stream();
     SYCL_CHECK(ggml_sycl_set_device(ctx.device));
@@ -147,7 +161,7 @@ inline void ggml_sycl_op_ssm_scan(ggml_backend_sycl_context & ctx, ggml_tensor *
         static_cast<const int32_t *>(src6->data), static_cast<float *>(dst->data),
         src0->nb[2], src0->nb[3], src1->nb[2], src1->nb[3], src2->nb[1], src2->nb[2],
         src3->nb[1], src4->nb[2], src4->nb[3], src5->nb[2], src5->nb[3],
-        s_off, nc, nr, nh, ng, n_t, n_s, stream);
+        s_off, nc, nr, nh, ng, n_t, n_s, K, stream);
 }
 
 void ggml_sycl_ssm_scan(ggml_backend_sycl_context & ctx, ggml_tensor * dst) {

ggml/src/ggml-vulkan/ggml-vulkan.cpp

@@ -1861,6 +1861,7 @@ struct vk_op_ssm_scan_push_constants {
     uint32_t nb42, nb43, nb52, nb53;
     uint32_t s_off;
     uint32_t n_head, d_head, n_group, n_tok;
+    uint32_t n_seq, K;
 };
 struct vk_op_ssm_conv_push_constants {
     uint32_t nb01, nb02;
@@ -12731,7 +12732,8 @@ static void ggml_vk_ssm_scan(ggml_backend_vk_context * ctx, vk_context& subctx,
         (uint32_t)src4->nb[2], (uint32_t)src4->nb[3],
         (uint32_t)src5->nb[2], (uint32_t)src5->nb[3],
         (uint32_t)s_off,
-        n_head, head_dim, n_group, n_tok
+        n_head, head_dim, n_group, n_tok,
+        n_seq, (uint32_t) ggml_get_op_params_i32(dst, 0)
     };
 
     vk_subbuffer dst_buf = ggml_vk_tensor_subbuffer(ctx, dst);
@@ -19417,8 +19419,9 @@ static void ggml_vk_check_results_0(ggml_backend_vk_context * ctx, ggml_cgraph *
         } else if (tensor->op == GGML_OP_ADD_ID) {
             tensor_clone = ggml_add_id(ggml_ctx, src_clone[0], src_clone[1], src_clone[2]);
         } else if (tensor->op == GGML_OP_SSM_SCAN) {
+            const int32_t K = ggml_get_op_params_i32(tensor, 0);
             tensor_clone = ggml_ssm_scan(ggml_ctx, src_clone[0], src_clone[1], src_clone[2],
-                                         src_clone[3], src_clone[4], src_clone[5], src_clone[6]);
+                                         src_clone[3], src_clone[4], src_clone[5], src_clone[6], K);
         } else if (tensor->op == GGML_OP_SSM_CONV) {
             tensor_clone = ggml_ssm_conv(ggml_ctx, src_clone[0], src_clone[1]);
         } else if (tensor->op == GGML_OP_ROLL) {

ggml/src/ggml-vulkan/vulkan-shaders/ssm_scan.comp

@@ -33,6 +33,8 @@ layout(push_constant) uniform PushConstants {
     uint d_head;
     uint n_group;
     uint n_tok;
+    uint n_seq;
+    uint K;
 };
 
 float softplus(float x) {
@@ -114,6 +116,14 @@ void main() {
         if (lane == 0) {
             d[y_base_idx + i * stride_y] = state_sum;
         }
+
+        const uint slot = n_tok - 1u - i;
+        if (slot > 0u && slot < K) {
+            const uint snapshot_base_idx = s_base_idx + slot * n_seq * (nb03 / 4u);
+            [[unroll]] for (uint j = 0; j < c_factor; j++) {
+                d[snapshot_base_idx + SUBGROUP_SIZE * j + lane] = state[j];
+            }
+        }
     }
 
     // write back the state

ggml/src/ggml-webgpu/ggml-webgpu.cpp

@@ -1327,6 +1327,7 @@ static webgpu_encoded_op ggml_webgpu_ssm_scan(webgpu_context & ctx,
         (uint32_t) src4->ne[1],
         (uint32_t) src1->ne[2],
         (uint32_t) ggml_nelements(src1),
+        (uint32_t) ggml_get_op_params_i32(dst, 0),
     };
 
     std::vector<wgpu::BindGroupEntry> entries = {

ggml/src/ggml-webgpu/wgsl-shaders/ssm_scan.wgsl

@@ -41,6 +41,7 @@ struct Params {
     n_seq_tokens: u32,
 
     y_elems: u32,
+    K: u32,
 };
 
 @group(0) @binding(0) var<storage, read_write> s_in: array<f32>;
@@ -123,6 +124,7 @@ fn main(
     let head_seq = wg_linear / params.d_inner;
     let ir = head_seq % params.n_head;
     let i3 = head_seq / params.n_head;
+    let n_seqs = params.y_elems / (params.n_seq_tokens * params.n_head * params.d_inner);
 
     let state_slot = read_state_slot(i3);
     let g = ir / (params.n_head / params.n_group);
@@ -179,6 +181,15 @@ fn main(
 #endif
             s_prev = s;
 
+            let slot = params.n_seq_tokens - 1u - token;
+            if (slot > 0u && slot < params.K) {
+                let snapshot_idx =
+                    params.offset_dst + params.y_elems + tid + i1 * params.d_state +
+                    ir * (params.d_state * params.d_inner) +
+                    (slot * n_seqs + i3) * (params.d_state * params.d_inner * params.n_head);
+                dst[snapshot_idx] = s;
+            }
+
 #ifdef USE_SUBGROUP_REDUCTION
 #ifdef XBC_OVERLAP
             let subgroup_partial = subgroupAdd(s * read_merged_f32(c_idx));

ggml/src/ggml.c

@@ -5588,7 +5588,10 @@ struct ggml_tensor * ggml_ssm_scan(
         struct ggml_tensor  * A,
         struct ggml_tensor  * B,
         struct ggml_tensor  * C,
-        struct ggml_tensor  * ids) {
+        struct ggml_tensor  * ids,
+        int64_t               K) {
+    GGML_ASSERT(K >= 1);
+    GGML_ASSERT(K <= INT32_MAX);
     GGML_ASSERT(ggml_is_contiguous(s));
     GGML_ASSERT(ggml_is_contiguous(dt));
     GGML_ASSERT(ggml_is_contiguous(A));
@@ -5625,11 +5628,12 @@ struct ggml_tensor * ggml_ssm_scan(
         if (A->ne[0] != 1) {
             // Mamba-1 has more granular decay factors
             GGML_ASSERT(A->ne[0] == d_state);
+            GGML_ASSERT(K == 1);
         }
     }
 
     // concatenated y + ssm_states
-    struct ggml_tensor * result = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, ggml_nelements(x) + s->ne[0]*s->ne[1]*s->ne[2]*ids->ne[0]);
+    struct ggml_tensor * result = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, ggml_nelements(x) + K*s->ne[0]*s->ne[1]*s->ne[2]*ids->ne[0]);
 
     result->op   = GGML_OP_SSM_SCAN;
     result->src[0] = s;
@@ -5640,6 +5644,8 @@ struct ggml_tensor * ggml_ssm_scan(
     result->src[5] = C;
     result->src[6] = ids;
 
+    ggml_set_op_params_i32(result, 0, (int32_t) K);
+
     return result;
 }
 

skills/code-review/SKILL.md

@@ -46,6 +46,8 @@ Mandatory on every review; any finding here is **blocking**. Rule of thumb: GGUF
 
 - **Sizes/counts from tensor dims:** validate before allocating. Products like `ne[i]*nb[i]`/nbytes can overflow on crafted dims into an undersized alloc then heap overflow. Overflow checks must run BEFORE the arithmetic they guard - padding/alignment macros wrap to 0 near `SIZE_MAX`, so a guard after the pad passes.
 - **GGUF strings/arrays:** cap declared lengths and element counts before using them to size a loop or buffer; validate element type and length before casting an array to a pointer or reading fixed indices (`[i+1]`, `[0..2]`).
+- **Element-type confusion:** casting `gguf_get_arr_data()` or `tensor->data` to `float *`/`int32_t *` needs an element-type check first (`gguf_get_kv_type() == GGUF_TYPE_ARRAY` then `gguf_get_arr_type()`; `type == GGML_TYPE_F32` for tensors). A `UINT8` array or `I8` tensor passes every length check, then gets read 4 bytes per element - a nearby length check is not a type check.
+- **Loaders:** `GGML_ASSERT` on a file-derived value aborts the process; throw instead where the caller already catches (vocab, model loader, clip).
 - **File-supplied counts indexing fixed arrays:** bound any count (e.g. layer/block count into a `LLAMA_MAX_*` array) before indexing; watch checks that only fire when an optional key is present.
 - **Declared vs actual array length:** check the declared length of a GGUF array against the count actually read, not just against a buffer size.
 - **Bounds comparisons:** flag narrowing casts (`size_t`->`int32_t`) and signed/unsigned mixing that can bypass a length check and copy past a buffer.

src/llama-arch.cpp

@@ -1001,6 +1001,8 @@ bool llm_arch_supports_rs_rollback(const llm_arch & arch) {
         case LLM_ARCH_QWEN35:
         case LLM_ARCH_QWEN35MOE:
         case LLM_ARCH_DEEPSEEK4:
+        case LLM_ARCH_NEMOTRON_H:
+        case LLM_ARCH_NEMOTRON_H_MOE:
             return true;
         default:
             return false;

src/llama-context.cpp

@@ -103,7 +103,7 @@ llama_context::llama_context(
 
     cparams.n_rs_seq = params.n_rs_seq;
     if (cparams.n_rs_seq > 0 && !llm_arch_supports_rs_rollback(model.arch)) {
-        LLAMA_LOG_DEBUG("%s: n_rs_seq=%u requested but model arch does not support recurrent partial rollback; clamping to 0\n",
+        LLAMA_LOG_DEBUG("%s: n_rs_seq=%u requested but model does not support recurrent partial rollback; clamping to 0\n",
                         __func__, cparams.n_rs_seq);
         cparams.n_rs_seq = 0;
     }

src/llama-model-loader.cpp

@@ -316,15 +316,19 @@ namespace GGUFMeta {
         struct GGUFMeta::ArrayInfo arr_info =
             GGUFMeta::GKV<GGUFMeta::ArrayInfo>::get_kv(ctx, kid);
 
+        bool type_ok = false;
         switch (arr_info.gt) {
             case GGUF_TYPE_UINT32:
-            case GGUF_TYPE_INT32:   GGML_ASSERT((std::is_same<T,     int32_t>::value) ||
-                                                (std::is_same<T,    uint32_t>::value)); break;
-            case GGUF_TYPE_FLOAT32: GGML_ASSERT((std::is_same<T,       float>::value)); break;
-            case GGUF_TYPE_STRING:  GGML_ASSERT((std::is_same<T, std::string>::value)); break;
+            case GGUF_TYPE_INT32:   type_ok = (std::is_same<T,     int32_t>::value) ||
+                                              (std::is_same<T,    uint32_t>::value); break;
+            case GGUF_TYPE_FLOAT32: type_ok = (std::is_same<T,       float>::value); break;
+            case GGUF_TYPE_STRING:  type_ok = (std::is_same<T, std::string>::value); break;
             default:
                 throw std::runtime_error(format("%s is not a string/float32/uint32/int32 array", key.c_str()));
         }
+        if (!type_ok) {
+            throw std::runtime_error(format("%s has wrong array element type %s", key.c_str(), gguf_type_name(arr_info.gt)));
+        }
 
         if constexpr (std::is_same<T, std::string>::value) {
             const size_t n_items = gguf_get_arr_n(ctx, kid);
@@ -357,16 +361,20 @@ namespace GGUFMeta {
         struct GGUFMeta::ArrayInfo arr_info =
             GGUFMeta::GKV<GGUFMeta::ArrayInfo>::get_kv(ctx, kid);
 
+        bool type_ok = false;
         switch (arr_info.gt) {
             case GGUF_TYPE_BOOL:
             case GGUF_TYPE_UINT32:
-            case GGUF_TYPE_INT32:   GGML_ASSERT((std::is_same<T,     int32_t>::value) ||
-                                                (std::is_same<T,    uint32_t>::value)); break;
-            case GGUF_TYPE_FLOAT32: GGML_ASSERT((std::is_same<T,       float>::value)); break;
-            case GGUF_TYPE_STRING:  GGML_ASSERT((std::is_same<T, std::string>::value)); break;
+            case GGUF_TYPE_INT32:   type_ok = (std::is_same<T,     int32_t>::value) ||
+                                              (std::is_same<T,    uint32_t>::value); break;
+            case GGUF_TYPE_FLOAT32: type_ok = (std::is_same<T,       float>::value); break;
+            case GGUF_TYPE_STRING:  type_ok = (std::is_same<T, std::string>::value); break;
             default:
                 throw std::runtime_error(format("%s is not a string/float32/uint32/int32 array", key.c_str()));
         }
+        if (!type_ok) {
+            throw std::runtime_error(format("%s has wrong array element type %s", key.c_str(), gguf_type_name(arr_info.gt)));
+        }
 
         if (arr_info.length > N_MAX) {
             throw std::runtime_error(format("array length %u for key %s exceeds max %u", (uint32_t) arr_info.length, key.c_str(), (uint32_t) N_MAX));
@@ -1002,7 +1010,7 @@ static bool weight_buft_supported(const llama_hparams & hparams, ggml_tensor * w
                 ggml_tensor * B   = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, d_state, n_group, n_seq_tokens, n_seqs);
                 ggml_tensor * C   = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, d_state, n_group, n_seq_tokens, n_seqs);
                 ggml_tensor * ids = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, n_seqs);
-                op_tensor = ggml_ssm_scan(ctx, s, x, dt, w, B, C, ids);
+                op_tensor = ggml_ssm_scan(ctx, s, x, dt, w, B, C, ids, /*K=*/1);
             } break;
         case GGML_OP_RWKV_WKV6:
             {

src/llama-vocab.cpp

@@ -1989,6 +1989,10 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) {
                 // Kimi-K2 doesn't need merges, skip
                 LLAMA_LOG_INFO("%s: Kimi-K2 tokenizer detected, skipping BPE merges\n", __func__);
             } else {
+                if (gguf_get_kv_type(ctx, merges_keyidx) != GGUF_TYPE_ARRAY ||
+                    gguf_get_arr_type(ctx, merges_keyidx) != GGUF_TYPE_STRING) {
+                    throw std::runtime_error(format("invalid gguf type for %s", kv(LLM_KV_TOKENIZER_MERGES).c_str()));
+                }
                 const int n_merges = gguf_get_arr_n(ctx, merges_keyidx);
                 for (int i = 0; i < n_merges; i++) {
                     const std::string word = gguf_get_arr_str(ctx, merges_keyidx, i);
@@ -2028,8 +2032,13 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) {
 
             const int precompiled_charsmap_keyidx = gguf_find_key(ctx, kv(LLM_KV_TOKENIZER_PRECOMPILED_CHARSMAP).c_str());
             if (precompiled_charsmap_keyidx != -1) {
+                if (gguf_get_kv_type(ctx, precompiled_charsmap_keyidx) != GGUF_TYPE_ARRAY) {
+                    throw std::runtime_error(format("invalid gguf type for %s", kv(LLM_KV_TOKENIZER_PRECOMPILED_CHARSMAP).c_str()));
+                }
                 const gguf_type pc_type = gguf_get_arr_type(ctx, precompiled_charsmap_keyidx);
-                GGML_ASSERT(pc_type == GGUF_TYPE_INT8 || pc_type == GGUF_TYPE_UINT8);
+                if (pc_type != GGUF_TYPE_INT8 && pc_type != GGUF_TYPE_UINT8) {
+                    throw std::runtime_error(format("invalid gguf type for %s", kv(LLM_KV_TOKENIZER_PRECOMPILED_CHARSMAP).c_str()));
+                }
 
                 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);
@@ -2081,6 +2090,10 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) {
                 throw std::runtime_error("cannot find tokenizer merges in model file\n");
             }
             {
+                if (gguf_get_kv_type(ctx, merges_keyidx) != GGUF_TYPE_ARRAY ||
+                    gguf_get_arr_type(ctx, merges_keyidx) != GGUF_TYPE_STRING) {
+                    throw std::runtime_error(format("invalid gguf type for %s", kv(LLM_KV_TOKENIZER_MERGES).c_str()));
+                }
                 const int n_merges = gguf_get_arr_n(ctx, merges_keyidx);
                 for (int i = 0; i < n_merges; i++) {
                     const std::string word = gguf_get_arr_str(ctx, merges_keyidx, i);
@@ -2407,11 +2420,20 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) {
         throw std::runtime_error("cannot find tokenizer vocab in model file\n");
     }
 
+    if (gguf_get_kv_type(ctx, token_idx) != GGUF_TYPE_ARRAY ||
+        gguf_get_arr_type(ctx, token_idx) != GGUF_TYPE_STRING) {
+        throw std::runtime_error(format("invalid gguf type for %s", kv(LLM_KV_TOKENIZER_LIST).c_str()));
+    }
+
     const uint32_t n_tokens = gguf_get_arr_n(ctx, token_idx);
 
     const float * scores = nullptr;
     const int score_idx = gguf_find_key(ctx, kv(LLM_KV_TOKENIZER_SCORES).c_str());
     if (score_idx != -1) {
+        if (gguf_get_kv_type(ctx, score_idx) != GGUF_TYPE_ARRAY ||
+            gguf_get_arr_type(ctx, score_idx) != GGUF_TYPE_FLOAT32) {
+            throw std::runtime_error(format("invalid gguf type for %s", kv(LLM_KV_TOKENIZER_SCORES).c_str()));
+        }
         const uint32_t n_scores = gguf_get_arr_n(ctx, score_idx);
         if (n_scores < n_tokens) {
             throw std::runtime_error("Index out of array bounds for scores (" + std::to_string(n_scores) + " < " + std::to_string(n_tokens) + ")\n");
@@ -2422,6 +2444,10 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) {
     const int * toktypes = nullptr;
     const int toktype_idx = gguf_find_key(ctx, kv(LLM_KV_TOKENIZER_TOKEN_TYPE).c_str());
     if (toktype_idx != -1) {
+        if (gguf_get_kv_type(ctx, toktype_idx) != GGUF_TYPE_ARRAY ||
+            gguf_get_arr_type(ctx, toktype_idx) != GGUF_TYPE_INT32) {
+            throw std::runtime_error(format("invalid gguf type for %s", kv(LLM_KV_TOKENIZER_TOKEN_TYPE).c_str()));
+        }
         const uint32_t n_toktypes = gguf_get_arr_n(ctx, toktype_idx);
         if (n_toktypes < n_tokens) {
             throw std::runtime_error("Index out of array bounds for toktypes (" + std::to_string(n_toktypes) + " < " + std::to_string(n_tokens) + ")\n");
@@ -2584,6 +2610,10 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) {
         {
             const int suppress_idx = gguf_find_key(ctx, kv(LLM_KV_TOKENIZER_SUPPRESS_TOKENS).c_str());
             if (suppress_idx != -1) {
+                if (gguf_get_kv_type(ctx, suppress_idx) != GGUF_TYPE_ARRAY ||
+                    gguf_get_arr_type(ctx, suppress_idx) != GGUF_TYPE_INT32) {
+                    throw std::runtime_error(format("invalid gguf type for %s", kv(LLM_KV_TOKENIZER_SUPPRESS_TOKENS).c_str()));
+                }
                 const int n = gguf_get_arr_n(ctx, suppress_idx);
                 const int32_t * data = (const int32_t *) gguf_get_arr_data(ctx, suppress_idx);
                 // drop out-of-range ids

src/llama.cpp

@@ -257,7 +257,11 @@ static bool llama_prepare_model_devices(const llama_model_params & params, llama
                     }
 
                     case GGML_BACKEND_DEVICE_TYPE_IGPU:
-                        if (igpus.empty()) {
+                        // igpus.empty() - workaround for integrated devices seen by multiple backends
+                        // ref: https://github.com/ggml-org/llama.cpp/pull/23897
+                        // ggml_backend_dev_backend_reg - allow devices of the same backend regardless if integrated
+                        // ref: https://github.com/ggml-org/llama.cpp/pull/23897#issuecomment-5264222997
+                        if (igpus.empty() || ggml_backend_dev_backend_reg(dev) == ggml_backend_dev_backend_reg(igpus.back().dev)) {
                             igpus.push_back({false, dev});
                         }
                         break;

src/models/dflash.cpp

@@ -14,11 +14,14 @@ void llama_model_dflash::load_arch_hparams(llama_model_loader & ml) {
 
     hparams.n_embd_inp_enc_impl = (uint32_t) target_layer_ids.size() * hparams.n_embd;
 
-    LLAMA_LOG_INFO("%s: DFlash extract_layers = [", __func__);
-    for (size_t i = 0; i < target_layer_ids.size(); ++i) {
-        LLAMA_LOG_INFO("%d%s", target_layer_ids[i], i + 1 < target_layer_ids.size() ? ", " : "");
+    std::string layers;
+    const char * sep = "";
+    for (const auto id : target_layer_ids) {
+        layers += sep;
+        layers += std::to_string(id);
+        sep = ", ";
     }
-    LLAMA_LOG_INFO("]\n");
+    LLAMA_LOG_INFO("%s: DFlash extract_layers = [%s]\n", __func__, layers.c_str());
 
     // DeepSeek-V4 DSpark backbone: stages are full DSV4 blocks, uniform sliding window (the draft KV ring)
     ml.get_key(LLM_KV_HYPER_CONNECTION_COUNT, hparams.dsv4_hc_mult, false);

src/models/mamba-base.cpp

@@ -2,6 +2,8 @@
 
 #include "llama-memory-recurrent.h"
 
+#include <algorithm>
+
 llm_build_mamba_base::llm_build_mamba_base(const llm_graph_params & params) : llm_graph_context(params) {}
 
 ggml_tensor * llm_build_mamba_base::build_mamba_layer(llm_graph_input_rs * inp,
@@ -118,7 +120,7 @@ ggml_tensor * llm_build_mamba_base::build_mamba_layer(llm_graph_input_rs * inp,
             // Custom operator to optimize the parallel associative scan
             // as described in the Annex D of the Mamba paper.
             // => {d_inner, n_seq_tokens, n_seqs} and {d_state, d_inner, n_seqs}
-            return ggml_ssm_scan(ctx, ssm, x, dt, A, B, C, ids);
+            return ggml_ssm_scan(ctx, ssm, x, dt, A, B, C, ids, /*K=*/1);
         };
 
         ggml_tensor * y_ssm = build_rs(inp, ssm_states_all, hparams.n_embd_s(), ubatch.n_seqs, get_ssm_rows);
@@ -153,7 +155,8 @@ ggml_tensor * llm_build_mamba_base::build_mamba2_layer(llm_graph_input_rs * inp,
                                                           int                  il) const {
     const auto * mctx_cur = inp->mctx;
 
-    const auto kv_head = mctx_cur->get_head();
+    const auto kv_head  = mctx_cur->get_head();
+    const auto mem_size = mctx_cur->get_size();
 
     const int64_t d_conv   = hparams.ssm_d_conv;
     const int64_t d_inner  = hparams.ssm_d_inner;
@@ -164,6 +167,7 @@ ggml_tensor * llm_build_mamba_base::build_mamba2_layer(llm_graph_input_rs * inp,
     const int64_t n_seqs   = ubatch.n_seqs;
 
     const int64_t n_seq_tokens = ubatch.n_seq_tokens;
+    const int64_t K            = cparams.n_rs_seq > 0 ? (int64_t) cparams.n_rs_seq + 1 : 1;
 
     GGML_ASSERT(n_seqs != 0);
     GGML_ASSERT(ubatch.equal_seqs());
@@ -173,6 +177,7 @@ ggml_tensor * llm_build_mamba_base::build_mamba2_layer(llm_graph_input_rs * inp,
 
     ggml_tensor * conv_states_all = mctx_cur->get_r_l(il);
     ggml_tensor * ssm_states_all  = mctx_cur->get_s_l(il);
+    const int64_t state_slots     = ssm_states_all->ne[1];
 
     ggml_tensor * conv = build_rs(inp, conv_states_all, hparams.n_embd_r(), n_seqs);
     conv               = ggml_reshape_3d(ctx0, conv, d_conv - 1, d_inner + 2 * n_group * d_state, n_seqs);
@@ -198,15 +203,19 @@ ggml_tensor * llm_build_mamba_base::build_mamba2_layer(llm_graph_input_rs * inp,
         // => {d_conv - 1 + n_seq_tokens, d_inner + 2*n_group*d_state, n_seqs}
         ggml_tensor * conv_x = ggml_concat(ctx0, conv, ggml_transpose(ctx0, xBC), 0);
 
-        // copy last (d_conv - 1) columns back into the state cache
-        ggml_tensor * last_conv = ggml_view_3d(ctx0, conv_x, d_conv - 1, d_inner + 2 * n_group * d_state, n_seqs,
-                                               conv_x->nb[1], conv_x->nb[2], n_seq_tokens * (conv_x->nb[0]));
+        const int64_t row_count = (d_conv - 1) * (d_inner + 2 * n_group * d_state);
+        const size_t  row_size  = ggml_row_size(conv_states_all->type, row_count);
+        const int64_t n_written = std::min<int64_t>(n_seq_tokens, K);
 
-        ggml_build_forward_expand(gf, ggml_cpy(ctx0, last_conv,
-                                               ggml_view_1d(ctx0, conv_states_all,
-                                                            (d_conv - 1) * (d_inner + 2 * n_group * d_state) * (n_seqs),
-                                                            kv_head * (d_conv - 1) * (d_inner + 2 * n_group * d_state) *
-                                                                ggml_element_size(conv_states_all))));
+        for (int64_t slot = 0; slot < n_written; ++slot) {
+            ggml_tensor * last_conv = ggml_view_3d(ctx0, conv_x, d_conv - 1, d_inner + 2 * n_group * d_state, n_seqs,
+                                                   conv_x->nb[1], conv_x->nb[2], (n_seq_tokens - slot) * conv_x->nb[0]);
+
+            ggml_build_forward_expand(gf, ggml_cpy(ctx0, last_conv,
+                                                   ggml_view_2d(ctx0, conv_states_all, row_count, n_seqs,
+                                                                conv_states_all->nb[1],
+                                                                ((size_t) slot * mem_size + kv_head) * row_size)));
+        }
 
         // 1D convolution
         // The equivalent is to make a self-overlapping view of conv_x
@@ -244,20 +253,27 @@ ggml_tensor * llm_build_mamba_base::build_mamba2_layer(llm_graph_input_rs * inp,
         // (this is necessary in order to properly use the states before they are overwritten,
         //  while avoiding to make unnecessary copies of the states)
         auto get_ssm_rows = [&](ggml_context * ctx, ggml_tensor * states, ggml_tensor * ids) {
-            ggml_tensor * ssm = ggml_reshape_4d(ctx, states, d_state, head_dim, n_head, mctx_cur->get_size());
+            ggml_tensor * ssm = ggml_reshape_4d(ctx, states, d_state, head_dim, n_head, state_slots);
 
             // TODO: use semistructured matrices to implement state-space duality
             // => {d_inner, n_seq_tokens, n_seqs} and {d_state, d_inner, n_seqs}
-            return ggml_ssm_scan(ctx, ssm, x, dt, A, B, C, ids);
+            // K > 1 asks the backend to return rollback snapshots in addition to the final state.
+            return ggml_ssm_scan(ctx, ssm, x, dt, A, B, C, ids, K);
         };
 
         ggml_tensor * y_ssm = build_rs(inp, ssm_states_all, hparams.n_embd_s(), ubatch.n_seqs, get_ssm_rows);
+        const int64_t D            = d_state * d_inner;
+        const int64_t n_written    = std::min<int64_t>(n_seq_tokens, K);
+        const size_t  row_size     = ggml_row_size(ssm_states_all->type, D);
+        const size_t  y_row_size   = ggml_row_size(y_ssm->type, D);
+        const size_t  state_offset = ggml_nelements(x) * ggml_element_size(x);
 
-        // store last states
         ggml_build_forward_expand(
-            gf, ggml_cpy(ctx0, ggml_view_1d(ctx0, y_ssm, d_state * d_inner * n_seqs, ggml_nelements(x) * x->nb[0]),
-                         ggml_view_1d(ctx0, ssm_states_all, d_state * d_inner * n_seqs,
-                                      kv_head * d_state * d_inner * ggml_element_size(ssm_states_all))));
+            gf, ggml_cpy(ctx0,
+                         ggml_view_3d(ctx0, y_ssm, D, n_seqs, n_written,
+                                      y_row_size, y_row_size * n_seqs, state_offset),
+                         ggml_view_3d(ctx0, ssm_states_all, D, n_seqs, n_written,
+                                      ssm_states_all->nb[1], (size_t) mem_size * row_size, kv_head * row_size)));
 
         ggml_tensor * y = ggml_view_4d(ctx0, y_ssm, head_dim, n_head, n_seq_tokens, n_seqs, x->nb[1], n_head * x->nb[1],
                                        n_seq_tokens * n_head * x->nb[1], 0);

src/models/plamo2.cpp

@@ -382,7 +382,7 @@ ggml_tensor * llama_model_plamo2::graph::build_plamo2_mamba_layer(llm_graph_inpu
             // Custom operator to optimize the parallel associative scan
             // as described in the Annex D of the Mamba paper.
             // => {d_inner, n_seq_tokens, n_seqs} and {d_state, d_inner, n_seqs}
-            return ggml_ssm_scan(ctx, ssm, x, dt, A, B, C, ids);
+            return ggml_ssm_scan(ctx, ssm, x, dt, A, B, C, ids, /*K=*/1);
         };
 
         ggml_tensor * y_ssm = build_rs(inp, ssm_states_all, hparams.n_embd_s(), ubatch.n_seqs, get_ssm_rows);

tests/CMakeLists.txt

@@ -217,6 +217,16 @@ if (NOT WIN32 OR NOT BUILD_SHARED_LIBS)
     set_tests_properties(test-recurrent-state-rollback PROPERTIES
         FIXTURES_REQUIRED generate-models
     )
+
+    llama_test(
+        test-recurrent-state-rollback
+        NAME test-recurrent-state-rollback-nemotron-h
+        LABEL main
+        ARGS -m "${MODEL_DIR}/nemotron_h-dense.gguf"
+    )
+    set_tests_properties(test-recurrent-state-rollback-nemotron-h PROPERTIES
+        FIXTURES_REQUIRED generate-models
+    )
 endif()
 
 llama_build_and_test(test-chat-peg-parser.cpp peg-parser/simple-tokenize.cpp)

tests/test-backend-ops.cpp

@@ -4111,9 +4111,10 @@ struct test_ssm_scan : public test_case {
     const int64_t n_seq_tokens;
     const int64_t n_seqs;
     const bool    xbc_overlap;
+    const int64_t K;
 
     std::string vars() override {
-        return VARS_TO_STR8(type, d_state, head_dim, n_head, n_group, n_seq_tokens, n_seqs, xbc_overlap);
+        return VARS_TO_STR9(type, d_state, head_dim, n_head, n_group, n_seq_tokens, n_seqs, xbc_overlap, K);
     }
 
     test_ssm_scan(ggml_type type = GGML_TYPE_F32,
@@ -4123,8 +4124,9 @@ struct test_ssm_scan : public test_case {
             int64_t n_group = 1,
             int64_t n_seq_tokens = 32,
             int64_t n_seqs = 32,
-            bool xbc_overlap = false)
-        : type(type), d_state(d_state), head_dim(head_dim), n_head(n_head), n_group(n_group), n_seq_tokens(n_seq_tokens), n_seqs(n_seqs), xbc_overlap(xbc_overlap) {}
+            bool xbc_overlap = false,
+            int64_t K = 1)
+        : type(type), d_state(d_state), head_dim(head_dim), n_head(n_head), n_group(n_group), n_seq_tokens(n_seq_tokens), n_seqs(n_seqs), xbc_overlap(xbc_overlap), K(K) {}
 
     double max_nmse_err() override {
         // SSD path (head_dim > 1) uses FP16 intermediates (M matrix, X_dt); Mamba-1 is pure FP32.
@@ -4153,7 +4155,7 @@ struct test_ssm_scan : public test_case {
             C = ggml_new_tensor_4d(ctx, type, d_state,  n_group, n_seq_tokens, n_seqs);
         }
         ggml_tensor * ids = ggml_new_tensor_1d(ctx, GGML_TYPE_I32,  n_seqs);
-        ggml_tensor * out = ggml_ssm_scan(ctx, s, x, dt, A, B, C, ids);
+        ggml_tensor * out = ggml_ssm_scan(ctx, s, x, dt, A, B, C, ids, K);
         return out;
     }
 
@@ -4185,6 +4187,114 @@ struct test_ssm_scan : public test_case {
     }
 };
 
+struct test_ssm_scan_rollback : public test_case {
+    const ggml_type type;
+
+    const int64_t d_state;
+    const int64_t head_dim;
+    const int64_t n_head;
+    const int64_t n_group;
+    const int64_t n_seq_tokens;
+    const int64_t n_seqs;
+    const int64_t K;
+
+    std::string vars() override {
+        return VARS_TO_STR8(type, d_state, head_dim, n_head, n_group, n_seq_tokens, n_seqs, K);
+    }
+
+    std::string op_desc(ggml_tensor * t) override {
+        GGML_UNUSED(t);
+        return "SSM_SCAN_ROLLBACK";
+    }
+
+    bool run_whole_graph() override {
+        return true;
+    }
+
+    double max_err() override {
+        return 1e-6;
+    }
+
+    double err(const float * a, const float * b, size_t n) override {
+        double result = 0.0;
+        for (size_t i = 0; i < n; ++i) {
+            result = std::max(result, (double) fabsf(a[i]));
+            result = std::max(result, (double) fabsf(b[i]));
+        }
+        return result;
+    }
+
+    test_ssm_scan_rollback(ggml_type type = GGML_TYPE_F32,
+            int64_t d_state = 32,
+            int64_t head_dim = 64,
+            int64_t n_head  = 16,
+            int64_t n_group = 2,
+            int64_t n_seq_tokens = 8,
+            int64_t n_seqs = 2,
+            int64_t K = 3)
+        : type(type), d_state(d_state), head_dim(head_dim), n_head(n_head), n_group(n_group),
+          n_seq_tokens(n_seq_tokens), n_seqs(n_seqs), K(K) {}
+
+    ggml_tensor * build_graph(ggml_context * ctx) override {
+        ggml_tensor * s   = ggml_new_tensor_4d(ctx, type, d_state,  head_dim,     n_head,       n_seqs);
+        ggml_tensor * x   = ggml_new_tensor_4d(ctx, type, head_dim, n_head,       n_seq_tokens, n_seqs);
+        ggml_tensor * dt  = ggml_new_tensor_3d(ctx, type, n_head,   n_seq_tokens, n_seqs);
+        ggml_tensor * A   = ggml_new_tensor_2d(ctx, type, 1,        n_head);
+        ggml_tensor * B   = ggml_new_tensor_4d(ctx, type, d_state,  n_group,      n_seq_tokens, n_seqs);
+        ggml_tensor * C   = ggml_new_tensor_4d(ctx, type, d_state,  n_group,      n_seq_tokens, n_seqs);
+        ggml_tensor * ids = ggml_new_tensor_1d(ctx, GGML_TYPE_I32,  n_seqs);
+
+        ggml_tensor * full = ggml_ssm_scan(ctx, s, x, dt, A, B, C, ids, K);
+
+        const int64_t y_elems     = head_dim * n_head * n_seq_tokens * n_seqs;
+        const int64_t state_elems = d_state  * head_dim * n_head      * n_seqs;
+
+        ggml_tensor * out = nullptr;
+        for (int64_t slot = 0; slot < K; ++slot) {
+            const int64_t prefix_tokens = n_seq_tokens - slot;
+
+            ggml_tensor * x_prefix  = ggml_cont(ctx, ggml_view_4d(ctx, x,  head_dim, n_head,  prefix_tokens, n_seqs, x->nb[1],  x->nb[2],  x->nb[3],  0));
+            ggml_tensor * dt_prefix = ggml_cont(ctx, ggml_view_3d(ctx, dt, n_head,   prefix_tokens, n_seqs, dt->nb[1], dt->nb[2], 0));
+            ggml_tensor * B_prefix  = ggml_cont(ctx, ggml_view_4d(ctx, B,  d_state,  n_group, prefix_tokens, n_seqs, B->nb[1],  B->nb[2],  B->nb[3],  0));
+            ggml_tensor * C_prefix  = ggml_cont(ctx, ggml_view_4d(ctx, C,  d_state,  n_group, prefix_tokens, n_seqs, C->nb[1],  C->nb[2],  C->nb[3],  0));
+
+            ggml_tensor * prefix = ggml_ssm_scan(ctx, s, x_prefix, dt_prefix, A, B_prefix, C_prefix, ids, /*K=*/1);
+
+            ggml_tensor * full_state   = ggml_view_1d(ctx, full,   state_elems, (y_elems + slot*state_elems)*ggml_element_size(full));
+            ggml_tensor * prefix_state = ggml_view_1d(ctx, prefix, state_elems, (head_dim*n_head*prefix_tokens*n_seqs)*ggml_element_size(prefix));
+            ggml_tensor * diff         = ggml_sum(ctx, ggml_sqr(ctx, ggml_sub(ctx, full_state, prefix_state)));
+
+            out = out == nullptr ? diff : ggml_add(ctx, out, diff);
+        }
+
+        return out;
+    }
+
+    void initialize_tensors(ggml_context * ctx) override {
+        std::random_device rd;
+        std::default_random_engine rng(rd());
+        for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) {
+            if (t->type == GGML_TYPE_I32) {
+                if (ggml_is_view_op(t->op)) { continue; }
+                for (int64_t r = 0; r < ggml_nrows(t); r++) {
+                    std::vector<int32_t> data(t->ne[0]);
+                    for (int i = 0; i < t->ne[0]; i++) {
+                        data[i] = i;
+                    }
+                    std::shuffle(data.begin(), data.end(), rng);
+                    ggml_backend_tensor_set(t, data.data(), r * t->nb[1], t->ne[0] * sizeof(int32_t));
+                }
+            } else if (ggml_is_view_op(t->op)) {
+                continue;
+            } else if (t->ne[1] == n_head && t->ne[2] == 1) {
+                init_tensor_uniform(t, -1.0f, -0.5f);
+            } else {
+                init_tensor_uniform(t);
+            }
+        }
+    }
+};
+
 // GGML_OP_RWKV_WKV6
 struct test_rwkv_wkv6 : public test_case {
     const ggml_type type;
@@ -8952,6 +9062,9 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
     test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 128, 80, 128, 1, 256, 1)); // Nemotron-9B SSD path
     test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 128, 80, 128, 1, 512, 1)); // Nemotron-9B SSD multi-chunk (2 aligned chunks)
     test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 128, 64, 80, 8, 300, 2)); // Mamba-2 SSD multi-chunk (partial 2nd chunk, 2 seqs)
+    test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 128, 64, 16, 2, 4, 2, false, /*K=*/4)); // Mamba-2 rollback snapshots
+    test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 128, 64, 16, 2, 8, 2, false, /*K=*/3)); // Mamba-2 rollback overflow
+    test_cases.emplace_back(new test_ssm_scan_rollback(GGML_TYPE_F32, 128, 64, 16, 2, 8, 2, /*K=*/3)); // rollback snapshots match prefix states
 
     test_cases.emplace_back(new test_rwkv_wkv6(GGML_TYPE_F32, 32, 64, 1, 1));
     test_cases.emplace_back(new test_rwkv_wkv6(GGML_TYPE_F32, 32, 64, 32, 1));
@@ -9804,6 +9917,13 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
                                     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}));
+                                if (!use_id && with_gate && !with_bias) {
+                                    // small multi-token batches (speculative decoding / MTP verify)
+                                    for (int64_t m_batch : { 2, 4, 8 }) {
+                                        test_cases.emplace_back(new test_mul_mat_vec_fusion(type, glu_op, m_batch, 32, 256,
+                                            use_id, 16, 8, b, with_bias, with_gate, with_lane_scale, {1, 1}));
+                                    }
+                                }
                             }
                         }
                     }

tests/test-chat.cpp

@@ -4618,7 +4618,7 @@ static void test_template_output_peg_parsers(bool detailed_debug) {
 
         // Real life test - execute_command
         tst.test("<|tool_call_begin|>functions.execute_command:0<|tool_call_argument_begin|>{\"command\": \"ls -lah\""
-            ", \"cwd\": \"/home/jarvis/development/exllamav3\", \"timeout\": 10}")
+            ", \"cwd\": \"/home/user/development/exllamav3\", \"timeout\": 10}")
             .reasoning_format(COMMON_REASONING_FORMAT_AUTO)
             .parallel_tool_calls(true)
             .tools({
@@ -4648,7 +4648,7 @@ static void test_template_output_peg_parsers(bool detailed_debug) {
             expect_tool_calls({
                 {
                     "execute_command",
-                    R"({"command": "ls -lah", "cwd": "/home/jarvis/development/exllamav3", "timeout": 10})",
+                    R"({"command": "ls -lah", "cwd": "/home/user/development/exllamav3", "timeout": 10})",
                     "functions.execute_command:0"
                 }
             })

tools/completion/completion.cpp

@@ -160,47 +160,6 @@ int llama_completion(int argc, char ** argv) {
     // start measuring performance timings from here
     llama_perf_context_reset(ctx);
 
-    LOG_INF("%s: llama threadpool init, n_threads = %d\n", __func__, (int) params.cpuparams.n_threads);
-
-    auto * cpu_dev = ggml_backend_dev_by_type(GGML_BACKEND_DEVICE_TYPE_CPU);
-    if (!cpu_dev) {
-        LOG_ERR("%s: no CPU backend found\n", __func__);
-        return 1;
-    }
-    auto * reg = ggml_backend_dev_backend_reg(cpu_dev);
-    auto * ggml_threadpool_new_fn = (decltype(ggml_threadpool_new) *) ggml_backend_reg_get_proc_address(reg, "ggml_threadpool_new");
-    auto * ggml_threadpool_free_fn = (decltype(ggml_threadpool_free) *) ggml_backend_reg_get_proc_address(reg, "ggml_threadpool_free");
-
-    struct ggml_threadpool_params tpp_batch =
-            ggml_threadpool_params_from_cpu_params(params.cpuparams_batch);
-    struct ggml_threadpool_params tpp =
-            ggml_threadpool_params_from_cpu_params(params.cpuparams);
-
-    if (!set_process_priority(params.cpuparams.priority)) {
-        LOG_ERR("%s: error: failed to set process priority\n", __func__);
-        return 1;
-    }
-
-    struct ggml_threadpool * threadpool_batch = NULL;
-    if (!ggml_threadpool_params_match(&tpp, &tpp_batch)) {
-        threadpool_batch = ggml_threadpool_new_fn(&tpp_batch);
-        if (!threadpool_batch) {
-            LOG_ERR("%s: batch threadpool create failed : n_threads %d\n", __func__, tpp_batch.n_threads);
-            return 1;
-        }
-
-        // start the non-batch threadpool in the paused state
-        tpp.paused = true;
-    }
-
-    struct ggml_threadpool * threadpool = ggml_threadpool_new_fn(&tpp);
-    if (!threadpool) {
-        LOG_ERR("%s: threadpool create failed : n_threads %d\n", __func__, tpp.n_threads);
-        return 1;
-    }
-
-    llama_attach_threadpool(ctx, threadpool, threadpool_batch);
-
     const int n_ctx_train = llama_model_n_ctx_train(model);
     const int n_ctx = llama_n_ctx(ctx);
 
@@ -993,8 +952,5 @@ int llama_completion(int argc, char ** argv) {
 
     llama_backend_free();
 
-    ggml_threadpool_free_fn(threadpool);
-    ggml_threadpool_free_fn(threadpool_batch);
-
     return 0;
 }

tools/mtmd/clip.cpp

@@ -3734,6 +3734,9 @@ struct clip_model_loader {
             }
             return;
         }
+        if (gguf_get_kv_type(ctx_gguf.get(), i) != GGUF_TYPE_ARRAY) {
+            throw std::runtime_error(string_format("%s: key '%s' is not an array\n", __func__, key.c_str()));
+        }
         const auto type = gguf_get_arr_type(ctx_gguf.get(), i);
         if (type != GGUF_TYPE_FLOAT32) {
             throw std::runtime_error(string_format("%s: array '%s' has type %d, expected %d (GGUF_TYPE_FLOAT32)\n", __func__, key.c_str(), type, GGUF_TYPE_FLOAT32));
@@ -3768,6 +3771,9 @@ struct clip_model_loader {
             }
             return;
         }
+        if (gguf_get_kv_type(ctx_gguf.get(), i) != GGUF_TYPE_ARRAY) {
+            throw std::runtime_error(string_format("%s: key '%s' is not an array\n", __func__, key.c_str()));
+        }
         const auto type = gguf_get_arr_type(ctx_gguf.get(), i);
         if (type != GGUF_TYPE_INT32) {
             throw std::runtime_error(string_format("%s: array '%s' has type %d, expected %d (GGUF_TYPE_INT32)\n", __func__, key.c_str(), type, GGUF_TYPE_INT32));

tools/server/server-context.cpp

@@ -688,97 +688,99 @@ struct server_slot {
         other.prompt = prompt.clone();
         other.init_sampler();
     }
-
-    // returns 0 on success
-    // caller need to update prompt.tokens after a successful call to keep track of the processing progress
-    int process_mtmd_chunk(size_t idx, size_t & n_tokens_out) {
-        GGML_ASSERT(mctx);
-        const auto & input_tokens = task->tokens;
-        const auto & chunk = input_tokens.find_chunk(idx);
-        int32_t res = 0;
-
-        auto try_decode = [&]() -> int32_t {
-            if (mbatch) {
-                float * embd = mtmd_batch_get_output_embd(mbatch.get(), chunk.get());
-                if (embd) {
-                    void * cb_data = spec;
-                    static auto cb = [](llama_batch batch, void * user_data) {
-                        common_speculative * spec = static_cast<common_speculative *>(user_data);
-                        if (!common_speculative_process(spec, batch)) {
-                            return 1;
-                        }
-                        return 0;
-                    };
-
-                    llama_pos new_n_past; // unused for now
-                    res = mtmd_helper_decode_image_chunk(
-                        mctx,
-                        ctx_tgt,
-                        chunk.get(),
-                        embd,
-                        prompt.tokens.pos_next(),
-                        id,
-                        llama_n_batch(ctx_tgt),
-                        &new_n_past,
-                        cb,
-                        cb_data
-                    );
-                    if (res != 0) {
-                        SLT_ERR(*this, "failed to decode mtmd chunk, idx = %zu, res = %d\n", idx, res);
-                        return -1;
-                    }
-                    n_tokens_out = mtmd_input_chunk_get_n_tokens(chunk.get());
-                    return 0; // success
-                }
-            }
-            return 1; // (non-error) need to create & encode batch
-        };
-
-        // if the batch is already exist, try searching & encode
-        res = try_decode();
-        if (res == 0) {
-            return 0;
-        }
-        if (res < 0) {
-            // fatal error
-            return res;
-        }
-
-        // otherwise, the batch is either uninitialized or is used up
-        // we need to create & encode a new batch
-        mbatch.reset(mtmd_batch_init(mctx));
-        res = mtmd_batch_add_chunk(mbatch.get(), chunk.get());
-        GGML_ASSERT(res == 0); // we should never have an empty batch
-
-        // try batching as much as possible
-        int n_added = 1;
-        size_t idx_cur = idx;
-        while (res == 0) {
-            auto [next_chunk, next_idx] = input_tokens.find_next_media_chunk(idx_cur);
-            if (next_chunk == nullptr) {
-                break;
-            }
-            res = mtmd_batch_add_chunk(mbatch.get(), next_chunk->get());
-            n_added += (res == 0 ? 1 : 0);
-            idx_cur = next_idx;
-            SLT_DBG(*this, "try adding media chunk idx = %zu to batch, res = %d\n", next_idx, res);
-            // if res != 0, batch is full or chunk is not compatible -> this loop breaks
-        }
-
-        // TODO @ngxson : move this log line to debug when it become more stable
-        SLT_TRC(*this, "encoding mtmd batch from idx = %zu, n_chunks = %d\n", idx, n_added);
-
-        res = mtmd_batch_encode(mbatch.get());
-        if (res != 0) {
-            SLT_ERR(*this, "failed to encode mtmd batch for chunk idx = %zu, res = %d\n", idx, res);
-            return -1;
-        }
-
-        return try_decode();
-    }
 };
 
+// returns 0 on success
+// caller need to update prompt.tokens after a successful call to keep track of the processing progress
+// note: this is not a member of server_slot because we want to run it inside yield_to_queue
+//       slot is passed as const to avoid accidental modification of the slot state
+//       some pointers are allowed to be used, they are not used by to_json()
+static int process_mtmd_chunk(const server_slot & slot, mtmd::batch_ptr & mbatch, size_t idx, size_t & n_tokens_out) {
+    GGML_ASSERT(slot.mctx);
+    const auto & mctx = slot.mctx;
+    const auto & input_tokens = slot.task->tokens;
+    const auto & chunk = input_tokens.find_chunk(idx);
+    int32_t res = 0;
 
+    auto try_decode = [&]() -> int32_t {
+        if (mbatch) {
+            float * embd = mtmd_batch_get_output_embd(mbatch.get(), chunk.get());
+            if (embd) {
+                void * cb_data = slot.spec;
+                static auto cb = [](llama_batch batch, void * user_data) {
+                    common_speculative * spec = static_cast<common_speculative *>(user_data);
+                    if (!common_speculative_process(spec, batch)) {
+                        return 1;
+                    }
+                    return 0;
+                };
+
+                llama_pos new_n_past; // unused for now
+                res = mtmd_helper_decode_image_chunk(
+                    mctx,
+                    slot.ctx_tgt,
+                    chunk.get(),
+                    embd,
+                    slot.prompt.tokens.pos_next(),
+                    slot.id,
+                    llama_n_batch(slot.ctx_tgt),
+                    &new_n_past,
+                    cb,
+                    cb_data
+                );
+                if (res != 0) {
+                    SLT_ERR(slot, "failed to decode mtmd chunk, idx = %zu, res = %d\n", idx, res);
+                    return -1;
+                }
+                n_tokens_out = mtmd_input_chunk_get_n_tokens(chunk.get());
+                return 0; // success
+            }
+        }
+        return 1; // (non-error) need to create & encode batch
+    };
+
+    // if the batch is already exist, try searching & encode
+    res = try_decode();
+    if (res == 0) {
+        return 0;
+    }
+    if (res < 0) {
+        // fatal error
+        return res;
+    }
+
+    // otherwise, the batch is either uninitialized or is used up
+    // we need to create & encode a new batch
+    mbatch.reset(mtmd_batch_init(mctx));
+    res = mtmd_batch_add_chunk(mbatch.get(), chunk.get());
+    GGML_ASSERT(res == 0); // we should never have an empty batch
+
+    // try batching as much as possible
+    int n_added = 1;
+    size_t idx_cur = idx;
+    while (res == 0) {
+        auto [next_chunk, next_idx] = input_tokens.find_next_media_chunk(idx_cur);
+        if (next_chunk == nullptr) {
+            break;
+        }
+        res = mtmd_batch_add_chunk(mbatch.get(), next_chunk->get());
+        n_added += (res == 0 ? 1 : 0);
+        idx_cur = next_idx;
+        SLT_DBG(slot, "try adding media chunk idx = %zu to batch, res = %d\n", next_idx, res);
+        // if res != 0, batch is full or chunk is not compatible -> this loop breaks
+    }
+
+    // TODO @ngxson : move this log line to debug when it become more stable
+    SLT_TRC(slot, "encoding mtmd batch from idx = %zu, n_chunks = %d\n", idx, n_added);
+
+    res = mtmd_batch_encode(mbatch.get());
+    if (res != 0) {
+        SLT_ERR(slot, "failed to encode mtmd batch for chunk idx = %zu, res = %d\n", idx, res);
+        return -1;
+    }
+
+    return try_decode();
+}
 
 //
 // server_context_impl (private implementation)
@@ -1354,8 +1356,8 @@ private:
         GGML_ASSERT(!sleeping);
 
         // wiring up server queues
-        queue_tasks.on_new_task([this](server_task && task) {
-            process_single_task(std::move(task));
+        queue_tasks.on_new_task([this](server_task && task, bool is_yielding) {
+            return process_single_task(std::move(task), is_yielding);
         });
         queue_tasks.on_update_slots([this]() {
             update_slots();
@@ -2286,7 +2288,14 @@ private:
                 cur.pos_max, cur.n_tokens, (float) cur.size() / 1024 / 1024);
     }
 
-    void process_single_task(server_task && task) {
+    // returns false to decline the task, it is offered again after the decode is done
+    bool process_single_task(server_task && task, bool is_yielding) {
+        // while yielding, an encode / decode is running and only accessing metrics is safe
+        if (is_yielding && task.type != SERVER_TASK_TYPE_METRICS) {
+            SRV_DBG("decoding, decline task, id_task = %d\n", task.id);
+            return false;
+        }
+
         switch (task.type) {
             case SERVER_TASK_TYPE_COMPLETION:
             case SERVER_TASK_TYPE_INFILL:
@@ -2620,6 +2629,8 @@ private:
                     queue_results.send(std::move(res));
                 } break;
         }
+
+        return true;
     }
 
     void iterate(std::vector<server_slot> & slots, std::function<void(server_slot &)> callback) {
@@ -3382,8 +3393,13 @@ private:
                         //       so the timing is queued and flushed on the next sync
                         metrics_pre_decode();
 
+                        // encode on the worker thread, so we can still handle metrics tasks
                         size_t n_tokens_out = 0;
-                        int32_t res = slot.process_mtmd_chunk(cur_token_idx, n_tokens_out);
+                        int32_t res = 0;
+                        queue_tasks.yield_to_queue([&]() {
+                            res = process_mtmd_chunk(slot, slot.mbatch, cur_token_idx, n_tokens_out);
+                        });
+
                         if (res != 0) {
                             SLT_ERR(slot, "failed to process mtmd chunk, res = %d\n", res);
                             send_error(slot, "failed to process mtmd chunk", ERROR_TYPE_SERVER);
@@ -3557,7 +3573,20 @@ private:
             }
         }
 
-        const int ret = llama_decode(ctx_tgt, batch_view);
+        bool has_output = false;
+        for (int i = off; i < off + batch_view.n_tokens; ++i) {
+            has_output |= batch.tokens[i].output;
+        }
+
+        // decode on the worker thread, so we can still handle metrics tasks while waiting
+        // note: the sync is done here too, so that the wait also happens off the main thread
+        int ret = 0;
+        queue_tasks.yield_to_queue([&]() {
+            ret = llama_decode(ctx_tgt, batch_view);
+            if (ret == 0 && has_output) {
+                llama_synchronize(ctx_tgt);
+            }
+        });
 
         if (ret != 0) {
             {
@@ -3609,7 +3638,7 @@ private:
             return false; // retry with the updated n_batch
         } else {
             // success, apply batch metrics
-            metrics_post_decode(off, batch_view.n_tokens);
+            metrics_post_decode(off, batch_view.n_tokens, has_output);
         }
 
         // TODO: avoid restoring the draft context and re-evaluating the drafted tokens when not needed [TAG_SPEC_AVOID_DRAFT_REEVAL]
@@ -3922,7 +3951,8 @@ private:
         n_prompt_queued = 0;
     }
 
-    void metrics_post_decode(int32_t off, int32_t n_tokens) {
+    // has_output is computed by the caller, which also already synchronized the context if it is set
+    void metrics_post_decode(int32_t off, int32_t n_tokens, bool has_output) {
         metrics.n_decode++;
         for (const auto & slot : slots) {
             if (slot.is_processing()) {
@@ -3935,13 +3965,10 @@ private:
         // note: a slot can be released before we get here, which clears its stats
         //       the tokens were still computed, counted in the global metrics, not in slot
         uint64_t n_prompt_tokens = 0;
-        bool     has_output      = false;
 
         for (int i = off; i < off + n_tokens; ++i) {
             const auto & t = batch.tokens[i];
 
-            has_output |= t.output;
-
             if (!t.is_prompt) {
                 continue; // generated tokens are handled after sampling
             }
@@ -3957,14 +3984,12 @@ private:
         metrics_queue_prompt(n_prompt_tokens);
 
         if (has_output) {
-            // sync if we have at least one output in batch
-            // so that we can calculate the timings correctly
-            llama_synchronize(ctx_tgt);
+            // the context is already synchronized, so the timings are correct
             metrics_flush_prompt();
         }
 
         // advance the prompt timing of the slots that had tokens in this batch
-        // note: a second pass, it must run after the sync above to reflect the compute
+        // note: a second pass, it must run after the sync to reflect the compute
         const int64_t t_now = ggml_time_us();
         for (int i = off; i < off + n_tokens; ++i) {
             const auto & t = batch.tokens[i];

tools/server/server-queue.cpp

@@ -4,6 +4,7 @@
 #include "log.h"
 
 #include <chrono>
+#include <thread>
 
 #define QUE_INF(fmt, ...) LOG_INF("que  %12.*s: " fmt, 12, __func__, __VA_ARGS__)
 #define QUE_WRN(fmt, ...) LOG_WRN("que  %12.*s: " fmt, 12, __func__, __VA_ARGS__)
@@ -122,10 +123,135 @@ void server_queue::terminate() {
     condition_tasks.notify_all();
 }
 
+bool server_queue::process_new_tasks(bool is_yielding) {
+    while (true) {
+        std::unique_lock<std::mutex> lock(mutex_tasks);
+        if (!running) {
+            QUE_DBG("%s", "terminate\n");
+            return true;
+        }
+        if (queue_tasks.empty()) {
+            return false;
+        }
+        server_task task = std::move(queue_tasks.front());
+        queue_tasks.pop_front();
+        lock.unlock();
+
+        QUE_DBG("processing task, id = %d\n", task.id);
+        if (!callback_new_task(std::move(task), is_yielding)) {
+            // set it aside, do not put it back in the queue, else we offer it again in a loop
+            GGML_ASSERT(is_yielding && "a task can only be declined while yielding");
+            QUE_DBG("task declined, id = %d\n", task.id);
+            lock.lock();
+            queue_tasks_unhandled.push_back(std::move(task));
+        }
+    }
+}
+
+void server_queue::worker_loop() {
+    while (true) {
+        std::function<void()> work;
+        {
+            std::unique_lock<std::mutex> lock(mutex_tasks);
+            worker.cv.wait(lock, [&]{
+                return worker.stop || worker.work != nullptr;
+            });
+            if (worker.stop) {
+                return;
+            }
+            work = std::move(worker.work);
+            worker.work = nullptr;
+        }
+
+        // note: do not hold any lock here, work() may post new tasks
+        std::exception_ptr exception;
+        try {
+            work();
+        } catch (...) {
+            exception = std::current_exception();
+        }
+
+        // signal completion to yield_to_queue()
+        std::unique_lock<std::mutex> lock(mutex_tasks);
+        worker.exception = std::move(exception);
+        worker.busy = false;
+        condition_tasks.notify_all();
+    }
+}
+
+void server_queue::worker_stop() {
+    if (!worker.thread.joinable()) {
+        return;
+    }
+    {
+        std::unique_lock<std::mutex> lock(mutex_tasks);
+        worker.stop = true;
+    }
+    worker.cv.notify_one();
+    worker.thread.join();
+}
+
+void server_queue::yield_to_queue(std::function<void()> && work) {
+    GGML_ASSERT(worker.thread.joinable() && "yield_to_queue() requires start_loop() to be running");
+
+    QUE_DBG("%s", "yielding to queue\n");
+
+    {
+        std::unique_lock<std::mutex> lock(mutex_tasks);
+        GGML_ASSERT(!worker.busy && "yield_to_queue() cannot be nested");
+        worker.busy = true;
+        worker.work = std::move(work);
+    }
+    worker.cv.notify_one();
+
+    while (true) {
+        // note: on terminate this is a no-op, but we still wait for the work to finish
+        process_new_tasks(true);
+
+        std::unique_lock<std::mutex> lock(mutex_tasks);
+        // declined tasks are moved to queue_tasks_unhandled, so a non-empty queue always has something new
+        condition_tasks.wait(lock, [&]{
+            return !worker.busy || (running && !queue_tasks.empty());
+        });
+        if (!worker.busy) {
+            break;
+        }
+    }
+
+    std::exception_ptr exception;
+    {
+        std::unique_lock<std::mutex> lock(mutex_tasks);
+
+        // put the declined tasks back, keeping their order
+        while (!queue_tasks_unhandled.empty()) {
+            queue_tasks.push_front(std::move(queue_tasks_unhandled.back()));
+            queue_tasks_unhandled.pop_back();
+        }
+
+        // make sure to avoid idle timeout here
+        time_last_task = ggml_time_ms();
+
+        // the worker is idle now, take the exception it may have left behind
+        std::swap(exception, worker.exception);
+    }
+
+    QUE_DBG("%s", "done yielding to queue\n");
+
+    // note: rethrow only after the declined tasks are back in the queue, so they are not lost
+    if (exception) {
+        std::rethrow_exception(exception);
+    }
+}
+
 void server_queue::start_loop(int64_t idle_sleep_ms) {
     running = true;
     time_last_task = ggml_time_ms();
 
+    // spawn the worker thread used by yield_to_queue()
+    GGML_ASSERT(!worker.thread.joinable() && "start_loop() is already running");
+    worker.stop = false;
+    worker.thread = std::thread([this]() { worker_loop(); });
+
     constexpr auto max_wait_time = std::chrono::seconds(1);
     auto should_sleep = [&]() -> bool {
         // caller must hold mutex_tasks
@@ -138,24 +264,10 @@ void server_queue::start_loop(int64_t idle_sleep_ms) {
 
     while (true) {
         QUE_DBG("%s", "processing new tasks\n");
-
-        while (true) {
-            std::unique_lock<std::mutex> lock(mutex_tasks);
-            if (!running) {
-                QUE_DBG("%s", "terminate\n");
-                return;
-            }
-            if (queue_tasks.empty()) {
-                lock.unlock();
-                break;
-            }
-            server_task task = std::move(queue_tasks.front());
-            queue_tasks.pop_front();
-            lock.unlock();
-
-            QUE_DBG("processing task, id = %d\n", task.id);
-            callback_new_task(std::move(task));
+        if (process_new_tasks(false)) {
+            break; // terminate
         }
+
         // all tasks in the current loop is processed, slots data is now ready
         QUE_DBG("%s", "update slots\n");
 
@@ -206,6 +318,8 @@ void server_queue::start_loop(int64_t idle_sleep_ms) {
             }
         }
     }
+
+    worker_stop();
 }
 
 void server_queue::cleanup_pending_task(int id_target) {
@@ -214,11 +328,15 @@ void server_queue::cleanup_pending_task(int id_target) {
         return task.id == id_target;
     };
     queue_tasks.erase(
-        std::remove_if(queue_tasks.begin(),          queue_tasks.end(),          rm_func),
+        std::remove_if(queue_tasks.begin(),           queue_tasks.end(),           rm_func),
         queue_tasks.end());
     queue_tasks_deferred.erase(
-        std::remove_if(queue_tasks_deferred.begin(), queue_tasks_deferred.end(), rm_func),
+        std::remove_if(queue_tasks_deferred.begin(),  queue_tasks_deferred.end(),  rm_func),
         queue_tasks_deferred.end());
+    // a task declined while yielding is not in queue_tasks yet, but it can still be cancelled
+    queue_tasks_unhandled.erase(
+        std::remove_if(queue_tasks_unhandled.begin(), queue_tasks_unhandled.end(), rm_func),
+        queue_tasks_unhandled.end());
 }
 
 //

tools/server/server-queue.h

@@ -4,7 +4,9 @@
 
 #include <condition_variable>
 #include <deque>
+#include <exception>
 #include <mutex>
+#include <thread>
 #include <vector>
 #include <unordered_set>
 
@@ -21,16 +23,32 @@ private:
     // queues
     std::deque<server_task> queue_tasks;
     std::deque<server_task> queue_tasks_deferred;
+    // tasks declined while yielding, put back in queue_tasks once the yield is done
+    // note: kept as a member so that cleanup_pending_task() can also reach them
+    std::deque<server_task> queue_tasks_unhandled;
 
     std::mutex mutex_tasks;
     std::condition_variable condition_tasks;
 
+    // used by yield_to_queue, all fields are guarded by mutex_tasks
+    struct worker_t {
+        std::thread             thread;
+        std::condition_variable cv;   // the worker sleeps on this until there is work
+        std::function<void()>   work; // pending work, picked up by the thread
+        std::exception_ptr      exception; // exception thrown by work(), if any
+        bool stop = false;
+        bool busy = false;
+    };
+    worker_t worker;
+
     // callback functions
-    std::function<void(server_task &&)> callback_new_task;
-    std::function<void(void)>           callback_update_slots;
-    std::function<void(bool)>           callback_sleeping_state;
+    std::function<bool(server_task &&, bool)> callback_new_task;
+    std::function<void(void)>                 callback_update_slots;
+    std::function<void(bool)>                 callback_sleeping_state;
 
 public:
+    ~server_queue() { worker_stop(); }
+
     // Add a new task to the end of the queue
     int post(server_task && task, bool front = false);
 
@@ -75,6 +93,15 @@ public:
      */
     void start_loop(int64_t idle_sleep_ms = -1);
 
+    // run work() on a separate thread, while the current thread calls process_new_tasks
+    // returns once work() is done (may throw exceptions)
+    // must be called from start_loop() thread (ideally inside callback_update_slots)
+    // use case: return metrics while encode/decode is running
+    // ref: https://github.com/ggml-org/llama.cpp/pull/27041
+    //
+    // tasks declined by callback_new_task are put back in the queue once this returns
+    void yield_to_queue(std::function<void()> && work);
+
     // for metrics
     size_t queue_tasks_deferred_size() {
         std::unique_lock<std::mutex> lock(mutex_tasks);
@@ -86,7 +113,10 @@ public:
     //
 
     // Register function to process a new task
-    void on_new_task(std::function<void(server_task &&)> callback) {
+    // the second argument tells whether the queue is currently yielding (see yield_to_queue)
+    // only then may the callback return false to decline the task, and it must leave it
+    // untouched, so that it can be put back in the queue later
+    void on_new_task(std::function<bool(server_task &&, bool)> callback) {
         callback_new_task = std::move(callback);
     }
 
@@ -112,6 +142,15 @@ public:
 
 private:
     void cleanup_pending_task(int id_target);
+
+    // process all pending tasks in the queue
+    // returns true if the queue is terminated, false if there is no more task to process
+    // while yielding, declined tasks are moved to queue_tasks_unhandled
+    bool process_new_tasks(bool is_yielding);
+
+    // for worker_t
+    void worker_loop();
+    void worker_stop();
 };
 
 // struct for managing server responses

tools/ui/src/lib/components/app/chat/ChatAttachments/ChatAttachmentsPreview/ChatAttachmentsPreview.svelte

@@ -5,6 +5,7 @@
 		ChatAttachmentsPreviewNavButtons,
 		ChatAttachmentsPreviewThumbnailStrip
 	} from '$lib/components/app';
+	import { UI_DATA_ATTRS } from '$lib/constants';
 	import { modelsStore } from '$lib/stores';
 	import {
 		createBase64DataUrl,
@@ -90,7 +91,7 @@
 		const index = currentIndex;
 
 		setTimeout(() => {
-			const thumbnail = document.querySelector(`[data-thumbnail-index="${index}"]`);
+			const thumbnail = document.querySelector(`[${UI_DATA_ATTRS.THUMBNAIL_INDEX}="${index}"]`);
 
 			thumbnail?.scrollIntoView({ behavior: 'smooth', block: 'nearest', inline: 'center' });
 		}, 0);

tools/ui/src/lib/components/app/chat/ChatAttachments/ChatAttachmentsPreview/ChatAttachmentsPreviewThumbnailStrip.svelte

@@ -1,7 +1,7 @@
 <script lang="ts">
 	import { FileText, Music, Video } from '@lucide/svelte';
 	import { HorizontalScrollCarousel } from '$lib/components/app/misc';
-	import { ICON_CLASS_DEFAULT } from '$lib/constants';
+	import { ICON_CLASS_DEFAULT, UI_DATA_ATTRS } from '$lib/constants';
 
 	interface PreviewItem {
 		id: string;
@@ -36,7 +36,7 @@
 		<HorizontalScrollCarousel class="max-w-full">
 			{#each items as item, index (item.id)}
 				<button
-					data-thumbnail-index={index}
+					{...{ [UI_DATA_ATTRS.THUMBNAIL_INDEX]: index }}
 					class={[
 						'relative flex-shrink-0 cursor-pointer overflow-hidden rounded border-2 bg-black/80 backdrop-blur-sm transition-all hover:opacity-90',
 						index === currentIndex ? 'border-white' : 'border-transparent opacity-60',

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

@@ -3,12 +3,11 @@
 	import {
 		ChatAttachmentsList,
 		ChatFormActions,
-		ChatFormContentEditable,
-		ChatFormFileInputInvisible,
+		ChatFormCurrentWorkingDirectory,
+		ChatFormInput,
+		ChatFormInputFileInputInvisible,
 		ChatFormMcpResourcesList,
 		ChatFormPickers,
-		ChatFormTextarea,
-		ChatFormWorkingDirectory,
 		DialogMcpResourcesBrowser
 	} from '$lib/components/app';
 	import {
@@ -110,7 +109,7 @@
 	}: Props = $props();
 
 	// Component References
-	// Shared handle of the two input renderers (textarea + contenteditable).
+	// Shared handle of the two input renderers (plain textarea + rich chat form input).
 	type ChatInputHandle = {
 		focus(): void;
 		resetHeight(): void;
@@ -121,16 +120,16 @@
 
 	let audioRecorder: AudioRecorder | undefined;
 	let chatFormActionsRef: ChatFormActions | undefined = $state(undefined);
-	let fileInputRef: ChatFormFileInputInvisible | undefined = $state(undefined);
+	let fileInputRef: ChatFormInputFileInputInvisible | undefined = $state(undefined);
 	let pickersRef: { handleKeydown: (event: KeyboardEvent) => boolean } | undefined =
 		$state(undefined);
 	let inputRef: ChatInputHandle | undefined = $state(undefined);
 
-	// Render-mode gate: the plain textarea by default, the contenteditable
+	// Render-mode gate: the plain textarea by default, the rich chat form input
 	// while the buffer carries a `file://` mention link or a complete code
 	// span (badges and code chips need a DOM the textarea cannot provide).
 	// Demotes back once neither remains.
-	let useContenteditable = $state(false);
+	let useRichInput = $state(false);
 
 	// Audio Recording State
 	let isRecording = $state(false);
@@ -242,16 +241,15 @@
 	}
 
 	$effect(() => {
-		const wantContenteditable =
-			containsFileMentionLink(value ?? '') || containsCodeSpan(value ?? '');
+		const wantRichInput = containsFileMentionLink(value ?? '') || containsCodeSpan(value ?? '');
 
-		if (useContenteditable === wantContenteditable) return;
+		if (useRichInput === wantRichInput) return;
 
 		if (!caretOffsetPinned) {
 			pendingCaretOffset = inputRef?.getCaretOffset() ?? (value ?? '').length;
 		}
 
-		useContenteditable = wantContenteditable;
+		useRichInput = wantRichInput;
 		queueCaretRestore();
 	});
 
@@ -315,7 +313,7 @@
 
 			// Caret inside a fenced code block (closed, or still open
 			// while being typed): Enter adds a line, never submits. The
-			// contenteditable consumes this case locally; this gate
+			// rich chat form input consumes this case locally; this gate
 			// covers the plain textarea, where skipping submit lets the
 			// native newline through.
 			if (!isModifier && isOffsetInCodeBlock(value ?? '', inputRef?.getCaretOffset() ?? 0)) {
@@ -508,9 +506,9 @@
 		value = built.newValue;
 		onValueChange?.(built.newValue);
 
-		// Already in contenteditable mode: no renderer flip, so the swap
+		// Already in rich chat form input mode: no renderer flip, so the swap
 		// effect's caret restore never runs.
-		if (useContenteditable) {
+		if (useRichInput) {
 			queueCaretRestore();
 		}
 	}
@@ -544,7 +542,7 @@
 	}
 </script>
 
-<ChatFormFileInputInvisible bind:this={fileInputRef} onFileSelect={handleFileSelect} />
+<ChatFormInputFileInputInvisible bind:this={fileInputRef} onFileSelect={handleFileSelect} />
 
 <form
 	class="relative grid {className}"
@@ -603,35 +601,20 @@
 		<div
 			class="flex-column relative min-h-12 items-center rounded-4xl md:rounded-3xl py-2 pb-2.25 shadow-sm transition-all focus-within:shadow-md md:py-3!"
 		>
-			{#if useContenteditable}
-				<ChatFormContentEditable
-					class="px-5 py-1.5 md:pt-0 mb-0.5"
-					bind:this={inputRef}
-					bind:value
-					onKeydown={handleKeydown}
-					onInput={() => {
-						pickers.handleInput();
-						onValueChange?.(value);
-					}}
-					onPaste={handlePaste}
-					{disabled}
-					{placeholder}
-				/>
-			{:else}
-				<ChatFormTextarea
-					class="px-5 py-1.5 md:pt-0"
-					bind:this={inputRef}
-					bind:value
-					onKeydown={handleKeydown}
-					onInput={() => {
-						pickers.handleInput();
-						onValueChange?.(value);
-					}}
-					onPaste={handlePaste}
-					{disabled}
-					{placeholder}
-				/>
-			{/if}
+			<ChatFormInput
+				class="px-5 py-1.5 md:pt-0"
+				bind:this={inputRef}
+				bind:value
+				onKeydown={handleKeydown}
+				onInput={() => {
+					pickers.handleInput();
+					onValueChange?.(value);
+				}}
+				onPaste={handlePaste}
+				{disabled}
+				{placeholder}
+				{useRichInput}
+			/>
 
 			{#if mcpResourceStore.hasAttachments}
 				<ChatFormMcpResourcesList
@@ -667,7 +650,7 @@
 	<ContextGaugePopup />
 
 	{#if toolsStore.hasEnabledCwdTools}
-		<ChatFormWorkingDirectory
+		<ChatFormCurrentWorkingDirectory
 			directory={cwd}
 			isOpen={pickers.isWorkingDirectoryPickerOpen}
 			bind:query={pickers.workingDirectoryQuery}

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

@@ -1,10 +1,10 @@
 <script lang="ts">
-	import ChatFormWorkingDirectoryChip from './ChatFormWorkingDirectoryChip.svelte';
-	import ChatFormWorkingDirectoryResultsList from './ChatFormWorkingDirectoryResultsList.svelte';
+	import ChatFormCurrentWorkingDirectoryChip from './ChatFormCurrentWorkingDirectoryChip.svelte';
+	import ChatFormCurrentWorkingDirectoryResultsList from './ChatFormCurrentWorkingDirectoryResultsList.svelte';
 	import { FolderOpen } from '@lucide/svelte';
 	import SearchInput from '$lib/components/app/forms/SearchInput.svelte';
 	import * as Popover from '$lib/components/ui/popover';
-	import { DEFAULT_MOBILE_BREAKPOINT, HOME_TILDE, SEARCH } from '$lib/constants';
+	import { DEFAULT_MOBILE_BREAKPOINT, HOME_TILDE, SEARCH, UI_DATA_ATTRS } from '$lib/constants';
 	import { BuiltInTool, GlobSearchType, KeyboardKey } from '$lib/enums';
 	import { useDebouncedSearch } from '$lib/hooks/use-debounced-search.svelte';
 	import { usePickerNavigation } from '$lib/hooks/use-picker-navigation.svelte';
@@ -120,7 +120,7 @@
 	});
 
 	useScrollActiveRow({
-		dataIndex: 'result',
+		dataAttr: UI_DATA_ATTRS.RESULT_INDEX,
 		getContainer: () => listContainer,
 		getCount: () => queryResults.length,
 		getIndex: () => nav.hoveredIndex,
@@ -250,7 +250,7 @@
 			// user cancelled - silently ignore; other errors are logged
 			if (err instanceof DOMException && err.name === 'AbortError') return;
 
-			console.error('[ChatFormWorkingDirectory] showDirectoryPicker failed:', err);
+			console.error('[ChatFormCurrentWorkingDirectory] showDirectoryPicker failed:', err);
 		}
 	}
 
@@ -331,7 +331,7 @@
 	onclick={onOpen}
 	{disabled}
 >
-	<ChatFormWorkingDirectoryChip
+	<ChatFormCurrentWorkingDirectoryChip
 		{directory}
 		{homeBase}
 		{disabled}
@@ -372,7 +372,7 @@
 			{#if !fileSearchEnabled}
 				<div class="px-2 py-1.5 text-sm text-muted-foreground">{searchUnavailableMessage}</div>
 			{:else if query.trim() && (search.isSearching || queryResults.length > 0 || searchError)}
-				<ChatFormWorkingDirectoryResultsList
+				<ChatFormCurrentWorkingDirectoryResultsList
 					results={queryResults}
 					hoveredIndex={nav.hoveredIndex}
 					isSearching={search.isSearching}

tools/ui/src/lib/components/app/chat/ChatForm/ChatFormCurrentWorkingDirectory/ChatFormCurrentWorkingDirectoryResultsList.svelte

@@ -1,6 +1,7 @@
 <script lang="ts">
 	import { Folder } from '@lucide/svelte';
 	import { cn } from '$lib/components/ui/utils';
+	import { UI_DATA_ATTRS } from '$lib/constants';
 	import { highlightMatch } from '$lib/utils';
 	import { fly } from 'svelte/transition';
 
@@ -46,7 +47,7 @@
 		{#each results as path, index (path)}
 			<button
 				type="button"
-				data-result-index={index}
+				{...{ [UI_DATA_ATTRS.RESULT_INDEX]: index }}
 				data-highlighted={index === hoveredIndex ? '' : undefined}
 				class={cn(
 					'relative flex w-full cursor-pointer items-center gap-2 rounded-sm px-2 py-1.5 text-sm outline-hidden select-none data-highlighted:bg-accent data-highlighted:text-accent-foreground'

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

@@ -0,0 +1,78 @@
+<script lang="ts">
+	import ChatFormInputBasic from './ChatFormInputBasic.svelte';
+	import ChatFormInputRich from './ChatFormInputRich.svelte';
+
+	interface Props {
+		class?: string;
+		disabled?: boolean;
+		onInput?: () => void;
+		onKeydown?: (event: KeyboardEvent) => void;
+		onPaste?: (event: ClipboardEvent) => void;
+		placeholder?: string;
+		value?: string;
+		useRichInput?: boolean;
+	}
+
+	let {
+		class: className = '',
+		disabled = false,
+		onInput,
+		onKeydown,
+		onPaste,
+		placeholder = 'Ask anything...',
+		useRichInput = false,
+		value = $bindable('')
+	}: Props = $props();
+
+	let basicRef: ChatFormInputBasic | undefined = $state();
+	let richRef: ChatFormInputRich | undefined = $state();
+
+	// The two renderers share one imperative handle (focus/caret/height), so
+	// the parent can drive whichever variant is mounted through this one.
+	export function getElement() {
+		return useRichInput ? richRef?.getElement() : basicRef?.getElement();
+	}
+
+	export function focus() {
+		if (useRichInput) richRef?.focus();
+		else basicRef?.focus();
+	}
+
+	export function resetHeight() {
+		if (useRichInput) richRef?.resetHeight();
+		else basicRef?.resetHeight();
+	}
+
+	export function getCaretOffset(): number {
+		return useRichInput ? (richRef?.getCaretOffset() ?? 0) : (basicRef?.getCaretOffset() ?? 0);
+	}
+
+	export function setCaretOffset(offset: number) {
+		if (useRichInput) richRef?.setCaretOffset(offset);
+		else basicRef?.setCaretOffset(offset);
+	}
+</script>
+
+{#if useRichInput}
+	<ChatFormInputRich
+		bind:this={richRef}
+		class={className}
+		{disabled}
+		{onInput}
+		{onKeydown}
+		{onPaste}
+		{placeholder}
+		bind:value
+	/>
+{:else}
+	<ChatFormInputBasic
+		bind:this={basicRef}
+		class={className}
+		{disabled}
+		{onInput}
+		{onKeydown}
+		{onPaste}
+		{placeholder}
+		bind:value
+	/>
+{/if}

tools/ui/src/lib/components/app/chat/ChatForm/ChatFormInput/ChatFormInputBasic.svelte

@@ -48,7 +48,7 @@
 		}
 	}
 
-	// Plain-text caret offsets, shared with the contenteditable variant so
+	// Plain-text caret offsets, shared with the rich chat form input variant so
 	// the picker/paste flows can address either renderer through one handle.
 	export function getCaretOffset(): number {
 		if (!textareaElement) return 0;

tools/ui/src/lib/components/app/chat/ChatForm/ChatFormInput/ChatFormInputRich.svelte

@@ -1,8 +1,8 @@
 <script lang="ts">
-	import { CODE_BLOCK } from '$lib/constants';
-	import { ColorMode } from '$lib/enums';
+	import { CODE_BLOCK, CODE_TOKEN_ATTR, UI_DATA_ATTRS } from '$lib/constants';
+	import { BooleanString, ChatFormInputRichTokenKind, ColorMode } from '$lib/enums';
 	import { isMobile } from '$lib/stores';
-	import type { ContentEditableToken } from '$lib/types';
+	import type { ChatFormInputRichToken } from '$lib/types';
 	import type { SourceHistoryEntry } from '$lib/utils';
 	import {
 		badgeAwareWordJump,
@@ -53,7 +53,7 @@
 	// browser's native undo stack.
 	const history = new SourceHistory();
 
-	// Browsers disagree on what an empty contenteditable contains (`<br>`,
+	// Browsers disagree on what an empty rich chat form input contains (`<br>`,
 	// `<div><br></div>`, or nothing), so emptiness is decided by the
 	// serialized source, not the DOM shape.
 	function syncEmptyState(serialized?: string) {
@@ -61,15 +61,15 @@
 
 		const source = serialized ?? serializeContent(rootElement);
 
-		rootElement.dataset.empty = source.length === 0 ? 'true' : 'false';
+		rootElement.dataset.empty = source.length === 0 ? BooleanString.TRUE : BooleanString.FALSE;
 	}
 
-	function renderTokens(tokens: ContentEditableToken[]) {
+	function renderTokens(tokens: ChatFormInputRichToken[]) {
 		if (!rootElement) return;
 
 		const caret = rangeToTextOffset(rootElement, safeRange());
 
-		// eslint-disable-next-line svelte/no-dom-manipulating -- the token layer is owned imperatively; Svelte renders only the contenteditable host, never its children
+		// eslint-disable-next-line svelte/no-dom-manipulating -- the token layer is owned imperatively; Svelte renders only the rich chat form input host, never its children
 		rootElement.replaceChildren(buildFragment(tokens));
 
 		syncCodeBlockHatches(rootElement);
@@ -127,7 +127,9 @@
 	}
 
 	function highlightCodeBlocks(root: HTMLElement) {
-		for (const el of root.querySelectorAll<HTMLElement>('code[data-code-token="block"]')) {
+		for (const el of root.querySelectorAll<HTMLElement>(
+			`code[${CODE_TOKEN_ATTR}="${ChatFormInputRichTokenKind.CODE_BLOCK}"]`
+		)) {
 			highlightCodeBlockElement(el);
 		}
 	}
@@ -151,7 +153,10 @@
 		}
 
 		while (node && node !== rootElement) {
-			if (node instanceof HTMLElement && node.dataset.codeToken === 'block') {
+			if (
+				node instanceof HTMLElement &&
+				node.getAttribute(CODE_TOKEN_ATTR) === ChatFormInputRichTokenKind.CODE_BLOCK
+			) {
 				const caret = rangeToTextOffset(rootElement, range);
 
 				if (highlightCodeBlockElement(node)) {
@@ -189,11 +194,13 @@
 	 * (deduped via the data attribute) swapped on mode change.
 	 */
 	function loadHighlightTheme(isDark: boolean) {
-		document.querySelectorAll('style[data-highlight-theme-preview]').forEach((s) => s.remove());
+		document
+			.querySelectorAll(`style[${UI_DATA_ATTRS.HIGHLIGHT_THEME_PREVIEW}]`)
+			.forEach((s) => s.remove());
 
 		const style = document.createElement('style');
 
-		style.setAttribute('data-highlight-theme-preview', 'true');
+		style.setAttribute(UI_DATA_ATTRS.HIGHLIGHT_THEME_PREVIEW, BooleanString.TRUE);
 		style.textContent = isDark ? githubDarkCss : githubLightCss;
 
 		document.head.appendChild(style);
@@ -311,7 +318,7 @@
 					source[source.length - 2] !== '\n' &&
 					last?.nodeType === Node.TEXT_NODE
 				) {
-					// eslint-disable-next-line svelte/no-dom-manipulating -- the token layer is owned imperatively; Svelte renders only the contenteditable host, never its children
+					// eslint-disable-next-line svelte/no-dom-manipulating -- the token layer is owned imperatively; Svelte renders only the rich chat form input host, never its children
 					rootElement.appendChild(document.createTextNode('\n'));
 					restoreCaret(source.length);
 					resizeHeight();
@@ -404,7 +411,10 @@
 			let node: Node | null = container.parentNode;
 
 			while (node && node !== rootElement) {
-				if (node instanceof HTMLElement && node.dataset.codeToken === 'block') {
+				if (
+					node instanceof HTMLElement &&
+					node.getAttribute(CODE_TOKEN_ATTR) === ChatFormInputRichTokenKind.CODE_BLOCK
+				) {
 					const tail = document.createRange();
 
 					tail.setStart(container, offset);
@@ -462,7 +472,11 @@
 
 		const first = rootElement.firstChild;
 
-		if (!(first instanceof HTMLElement) || first.dataset.codeToken !== 'block') return false;
+		if (
+			!(first instanceof HTMLElement) ||
+			first.getAttribute(CODE_TOKEN_ATTR) !== ChatFormInputRichTokenKind.CODE_BLOCK
+		)
+			return false;
 
 		const range = safeRange();
 
@@ -483,7 +497,7 @@
 			if (firstLineEnd !== -1 && caret > firstLineEnd) return false;
 		}
 
-		// eslint-disable-next-line svelte/no-dom-manipulating -- the token layer is owned imperatively; Svelte renders only the contenteditable host, never its children
+		// eslint-disable-next-line svelte/no-dom-manipulating -- the token layer is owned imperatively; Svelte renders only the rich chat form input host, never its children
 		rootElement.prepend(document.createElement('br'));
 		restoreCaret(0, extend);
 
@@ -507,7 +521,11 @@
 
 		const second = first.nextSibling;
 
-		if (!(second instanceof HTMLElement) || second.dataset.codeToken !== 'block') return;
+		if (
+			!(second instanceof HTMLElement) ||
+			second.getAttribute(CODE_TOKEN_ATTR) !== ChatFormInputRichTokenKind.CODE_BLOCK
+		)
+			return;
 
 		const range = safeRange();
 		const onHatch =
@@ -787,7 +805,7 @@
 	}
 </script>
 
-<div class="flex-1 {className}">
+<div class="flex-1 {className} mb-0.5">
 	<div
 		bind:this={rootElement}
 		contenteditable={!disabled}
@@ -798,7 +816,7 @@
 		data-placeholder={placeholder}
 		tabindex={disabled ? -1 : 0}
 		class={[
-			'chat-form-contenteditable text-md min-h-12 w-full overflow-y-auto whitespace-pre-wrap wrap-break-word border-0 bg-transparent p-0 leading-6 outline-none focus-visible:ring-0 focus-visible:ring-offset-0',
+			'chat-form-input-rich text-md min-h-12 w-full overflow-y-auto whitespace-pre-wrap wrap-break-word border-0 bg-transparent p-0 leading-6 outline-none focus-visible:ring-0 focus-visible:ring-offset-0',
 			disabled && 'cursor-not-allowed'
 		]}
 		style="max-height: var(--max-message-height);"
@@ -815,18 +833,18 @@
 <style>
 	/* pre-wrap is load-bearing: without it Chromium collapses \n in
 	   text nodes and converts them to spaces while typing */
-	.chat-form-contenteditable {
+	.chat-form-input-rich {
 		white-space: pre-wrap;
 	}
 
-	.chat-form-contenteditable:global([data-empty='true'])::before {
+	.chat-form-input-rich:global([data-empty='true'])::before {
 		content: attr(data-placeholder);
 		color: var(--muted-foreground);
 		pointer-events: none;
 	}
 
 	/* Inline code - mirrors markdown-content.css */
-	.chat-form-contenteditable :global(code[data-code-token='inline']) {
+	.chat-form-input-rich :global(code[data-code-token='code_inline']) {
 		background: var(--muted);
 		color: var(--muted-foreground);
 		padding: 0.125rem 0.375rem;
@@ -835,7 +853,7 @@
 	}
 
 	/* Fenced code block - mirrors .code-block-wrapper in markdown-content.css */
-	.chat-form-contenteditable :global(code[data-code-token='block']) {
+	.chat-form-input-rich :global(code[data-code-token='code_block']) {
 		display: block;
 		margin: 0.25rem 0;
 		padding: 0.75rem 1rem;

tools/ui/src/lib/components/app/chat/ChatForm/ChatFormPickers/ChatFormPicker/ChatFormPickerList.svelte

@@ -1,7 +1,7 @@
 <script lang="ts" generics="T">
 	import { SearchInput } from '$lib/components/app';
 	import ScrollArea from '$lib/components/ui/scroll-area/scroll-area.svelte';
-	import { CHAT_FORM_POPOVER_MAX_HEIGHT } from '$lib/constants';
+	import { CHAT_FORM_POPOVER_MAX_HEIGHT, UI_DATA_ATTRS } from '$lib/constants';
 	import { useScrollActiveRow } from '$lib/hooks/use-scroll-active-row.svelte';
 	import type { Snippet } from 'svelte';
 
@@ -55,7 +55,7 @@
 	// selectedIndex/items.length are untracked so hover and result replacement
 	// never re-fire the scroll; keyboard nav is the only path that bumps the trigger.
 	useScrollActiveRow({
-		dataIndex: 'picker',
+		dataAttr: UI_DATA_ATTRS.PICKER_INDEX,
 		getContainer: () => listContainer,
 		getCount: () => items.length,
 		getIndex: () => selectedIndex,

tools/ui/src/lib/components/app/chat/ChatForm/ChatFormPickers/ChatFormPicker/ChatFormPickerListItem.svelte

@@ -1,4 +1,5 @@
 <script lang="ts">
+	import { UI_DATA_ATTRS } from '$lib/constants';
 	import type { Snippet } from 'svelte';
 
 	interface Props {
@@ -24,7 +25,7 @@
 
 <button
 	type="button"
-	data-picker-index={dataIndex}
+	{...{ [UI_DATA_ATTRS.PICKER_INDEX]: dataIndex }}
 	{disabled}
 	{onclick}
 	{onmouseenter}

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

@@ -1,7 +1,7 @@
 <script lang="ts">
-	import ChatFormCommandPicker from './ChatFormCommandPicker.svelte';
-	import ChatFormMentionPicker from './ChatFormMentionPicker.svelte';
+	import ChatFormPickerCommand from './ChatFormPickerCommand.svelte';
 	import ChatFormPickerMcpPrompts from './ChatFormPickerMcpPrompts/ChatFormPickerMcpPrompts.svelte';
+	import ChatFormPickerMention from './ChatFormPickerMention.svelte';
 	import type {
 		ChatFormCommand,
 		FileMentionEntry,
@@ -55,9 +55,9 @@
 		scopePath
 	}: Props = $props();
 
-	let commandPickerRef: ChatFormCommandPicker | undefined = $state(undefined);
+	let commandPickerRef: ChatFormPickerCommand | undefined = $state(undefined);
 	let promptPickerRef: ChatFormPickerMcpPrompts | undefined = $state(undefined);
-	let mentionPickerRef: ChatFormMentionPicker | undefined = $state(undefined);
+	let mentionPickerRef: ChatFormPickerMention | undefined = $state(undefined);
 
 	/** Delegate keyboard events to the active picker child; true if handled. */
 	export function handleKeydown(event: KeyboardEvent): boolean {
@@ -77,7 +77,7 @@
 	}
 </script>
 
-<ChatFormCommandPicker
+<ChatFormPickerCommand
 	bind:this={commandPickerRef}
 	isOpen={isCommandPickerOpen ?? false}
 	query={commandQuery ?? ''}
@@ -96,7 +96,7 @@
 	{onPromptLoadError}
 />
 
-<ChatFormMentionPicker
+<ChatFormPickerMention
 	bind:this={mentionPickerRef}
 	isOpen={isMentionPickerOpen ?? false}
 	query={mentionQuery ?? ''}

tools/ui/src/lib/components/app/chat/index.ts

@@ -91,7 +91,7 @@ export { default as ChatAttachmentsListItemThumbnailImage } from './ChatAttachme
  * preview without carousel, or a gallery/carousel view when multiple items exist.
  * Uses ChatAttachmentPreviewSingle internally for each item's content.
  */
-export { default as ChatAttachmentsPreview } from './ChatAttachments/ChatAttachmentsPreview.svelte';
+export { default as ChatAttachmentsPreview } from './ChatAttachments/ChatAttachmentsPreview/ChatAttachmentsPreview.svelte';
 export { default as ChatAttachmentsPreviewNavButtons } from './ChatAttachments/ChatAttachmentsPreview/ChatAttachmentsPreviewNavButtons.svelte';
 export { default as ChatAttachmentsPreviewFileInfo } from './ChatAttachments/ChatAttachmentsPreview/ChatAttachmentsPreviewFileInfo.svelte';
 export { default as ChatAttachmentsPreviewThumbnailStrip } from './ChatAttachments/ChatAttachmentsPreview/ChatAttachmentsPreviewThumbnailStrip.svelte';
@@ -120,8 +120,8 @@ export { default as ChatAttachmentsPreviewCurrentItem } from './ChatAttachments/
  * Used by ChatScreenForm and ChatMessageEditForm for both new conversations and message editing.
  *
  * **Architecture:**
- * - Composes ChatFormTextarea (or ChatFormContentEditable for messages with
- *   file mention links), ChatFormActions, and ChatFormPickerMcpPrompts
+ * - Composes ChatFormInput (a plain textarea, or a ChatFormInputRich for
+ *   messages with file mention links), ChatFormActions, and ChatFormPickerMcpPrompts
  * - Manages file upload state via `uploadedFiles` bindable prop
  * - Integrates with ModelsSelectorDropdown for model selection in router mode
  * - Communicates with parent via callbacks (onSubmit, onFilesAdd, onStop, etc.)
@@ -258,7 +258,7 @@ export { default as ChatFormContextGauge } from './ChatForm/ChatFormContextGauge
 /**
  * Hidden file input element for programmatic file selection.
  */
-export { default as ChatFormFileInputInvisible } from './ChatForm/ChatFormFileInputInvisible.svelte';
+export { default as ChatFormInputFileInputInvisible } from './ChatForm/ChatFormInput/ChatFormInputFileInputInvisible.svelte';
 
 /**
  * Displays MCP Resource attachments as a horizontal carousel.
@@ -267,18 +267,13 @@ export { default as ChatFormFileInputInvisible } from './ChatForm/ChatFormFileIn
 export { default as ChatFormMcpResourcesList } from './ChatForm/ChatFormMcpResourcesList.svelte';
 
 /**
- * Auto-resizing contenteditable input that renders `[name](file://...)`
- * mention links as inline chips while keeping the value as the markdown
- * source string. ChatForm swaps it in once a mention link lands in the
- * buffer. Shares the focus()/resetHeight()/caret handle with the textarea.
+ * The message editor. Renders a plain auto-resizing textarea by default,
+ * or a ChatFormInputRich that renders `[name](file://...)` mention links as
+ * inline chips (keeping the value as the markdown source string) once a
+ * mention link lands in the buffer. The variant is selected via the
+ * `useRichInput` prop; both share one imperative handle.
  */
-export { default as ChatFormContentEditable } from './ChatForm/ChatFormContentEditable.svelte';
-
-/**
- * Plain auto-resizing textarea with IME composition support. Default input
- * renderer inside ChatForm until a file mention lands.
- */
-export { default as ChatFormTextarea } from './ChatForm/ChatFormTextarea.svelte';
+export { default as ChatFormInput } from './ChatForm/ChatFormInput/ChatFormInput.svelte';
 
 /**
  * Working directory selector for agent mode. Renders a chip below the chat
@@ -288,7 +283,7 @@ export { default as ChatFormTextarea } from './ChatForm/ChatFormTextarea.svelte'
  * synthetic "Set working directory to ..." user message into chat history
  * and is enforced on tool calls via the `x-tool-cwd` request header.
  */
-export { default as ChatFormWorkingDirectory } from './ChatForm/ChatFormWorkingDirectory.svelte';
+export { default as ChatFormCurrentWorkingDirectory } from './ChatForm/ChatFormCurrentWorkingDirectory/ChatFormCurrentWorkingDirectory.svelte';
 
 /**
  * **ChatFormPickerMcpPrompts** - MCP prompt selection interface
@@ -359,14 +354,14 @@ export { default as ChatFormPickerPopover } from './ChatForm/ChatFormPickers/Cha
  * Generic scrollable list for picker popovers. Provides search input,
  * scroll-into-view for keyboard navigation, loading skeletons, empty state,
  * and optional footer. Uses Svelte 5 snippets for item/skeleton/footer rendering.
- * Shared by ChatFormPickerMcpPrompts and ChatFormMentionPicker.
+ * Shared by ChatFormPickerMcpPrompts and ChatFormPickerMention.
  */
 export { default as ChatFormPickerList } from './ChatForm/ChatFormPickers/ChatFormPicker/ChatFormPickerList.svelte';
 
 /**
  * Generic button wrapper for picker list items. Provides consistent styling,
  * hover/selected states, and data-picker-index attribute for scroll-into-view.
- * Shared by ChatFormPickerMcpPrompts and ChatFormMentionPicker.
+ * Shared by ChatFormPickerMcpPrompts and ChatFormPickerMention.
  */
 export { default as ChatFormPickerListItem } from './ChatForm/ChatFormPickers/ChatFormPicker/ChatFormPickerListItem.svelte';
 
@@ -389,14 +384,14 @@ export { default as ChatFormPickerListItemSkeleton } from './ChatForm/ChatFormPi
  * tool, scoped to the conversation cwd (or server home when unset).
  * Selection splices a `[name](file:///<abs path>)` link into the input.
  */
-export { default as ChatFormMentionPicker } from './ChatForm/ChatFormPickers/ChatFormMentionPicker.svelte';
+export { default as ChatFormPickerMention } from './ChatForm/ChatFormPickers/ChatFormPickerMention.svelte';
 
 /**
  * `/`-triggered slash-command picker. Lists the available slash commands
  * (`/prompt`, `/cwd`, `/model`) filtered by the typed query; selection
  * hands the command to the parent for dispatch.
  */
-export { default as ChatFormCommandPicker } from './ChatForm/ChatFormPickers/ChatFormCommandPicker.svelte';
+export { default as ChatFormPickerCommand } from './ChatForm/ChatFormPickers/ChatFormPickerCommand.svelte';
 
 /**
  * Hosts the chat-form pickers (slash-command, MCP prompt, file mention)

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

@@ -25,14 +25,12 @@
 		DialogMermaidPreview
 	} from '$lib/components/app';
 	import {
-		BOOL_TRUE_STRING,
 		CODE_BLOCK_CLASS,
-		DATA_ERROR_BOUND_ATTR,
-		DATA_ERROR_HANDLED_ATTR,
 		DIAGRAM_VIEW_MODE_ATTR,
 		DIAGRAM_VIEW_RENDERED,
 		DIAGRAM_VIEW_SOURCE,
 		IMAGE_NOT_ERROR_BOUND_SELECTOR,
+		MARKDOWN_DATA_ATTRS,
 		MERMAID_BLOCK_CLASS,
 		MERMAID_LANGUAGE,
 		MERMAID_RENDERED_ATTR,
@@ -42,7 +40,7 @@
 		SVG,
 		TOGGLE_SOURCE_BTN_CLASS
 	} from '$lib/constants';
-	import { ColorMode, UrlProtocol } from '$lib/enums';
+	import { BooleanString, ColorMode, UrlProtocol } from '$lib/enums';
 	import { FileTypeText } from '$lib/enums/files.enums';
 	import { createAutoScrollController } from '$lib/hooks/use-auto-scroll.svelte';
 	import { settingsStore } from '$lib/stores';
@@ -486,13 +484,19 @@
 			const copyButton = wrapper.querySelector<HTMLButtonElement>('.copy-code-btn');
 			const previewButton = wrapper.querySelector<HTMLButtonElement>('.preview-code-btn');
 
-			if (copyButton && copyButton.dataset.listenerBound !== 'true') {
-				copyButton.dataset.listenerBound = 'true';
+			if (
+				copyButton &&
+				copyButton.getAttribute(MARKDOWN_DATA_ATTRS.LISTENER_BOUND) !== BooleanString.TRUE
+			) {
+				copyButton.setAttribute(MARKDOWN_DATA_ATTRS.LISTENER_BOUND, BooleanString.TRUE);
 				copyButton.addEventListener('click', handleCopyClick);
 			}
 
-			if (previewButton && previewButton.dataset.listenerBound !== 'true') {
-				previewButton.dataset.listenerBound = 'true';
+			if (
+				previewButton &&
+				previewButton.getAttribute(MARKDOWN_DATA_ATTRS.LISTENER_BOUND) !== BooleanString.TRUE
+			) {
+				previewButton.setAttribute(MARKDOWN_DATA_ATTRS.LISTENER_BOUND, BooleanString.TRUE);
 				previewButton.addEventListener('click', handlePreviewClick);
 			}
 		}
@@ -508,7 +512,7 @@
 		const images = containerRef.querySelectorAll<HTMLImageElement>(IMAGE_NOT_ERROR_BOUND_SELECTOR);
 
 		for (const img of images) {
-			img.dataset[DATA_ERROR_BOUND_ATTR] = BOOL_TRUE_STRING;
+			img.setAttribute(MARKDOWN_DATA_ATTRS.ERROR_BOUND, BooleanString.TRUE);
 			img.addEventListener('error', handleImageError);
 		}
 	}
@@ -691,7 +695,7 @@
 
 		// Mark nodes immediately to prevent duplicate renders if called again during streaming.
 		// This avoids needing a guard that would block node discovery.
-		nodes.forEach((node) => node.setAttribute(MERMAID_RENDERED_ATTR, 'true'));
+		nodes.forEach((node) => node.setAttribute(MERMAID_RENDERED_ATTR, BooleanString.TRUE));
 
 		// Read mode before await so Svelte tracks it reactively.
 		const isDark = mode.current === ColorMode.DARK;
@@ -738,7 +742,7 @@
 		if (nodes.length === 0) return;
 
 		nodes.forEach((node) => {
-			node.setAttribute(SVG.RENDERED_ATTR, 'true');
+			node.setAttribute(SVG.RENDERED_ATTR, BooleanString.TRUE);
 
 			const source = node.getAttribute(SVG.SOURCE_ATTR) ?? node.textContent ?? '';
 			const clean = sanitizeSvg(source);
@@ -765,11 +769,11 @@
 		// Don't handle data URLs or already-handled images
 		if (
 			img.src.startsWith(UrlProtocol.DATA) ||
-			img.dataset[DATA_ERROR_HANDLED_ATTR] === BOOL_TRUE_STRING
+			img.getAttribute(MARKDOWN_DATA_ATTRS.ERROR_HANDLED) === BooleanString.TRUE
 		)
 			return;
 
-		img.dataset[DATA_ERROR_HANDLED_ATTR] = BOOL_TRUE_STRING;
+		img.setAttribute(MARKDOWN_DATA_ATTRS.ERROR_HANDLED, BooleanString.TRUE);
 
 		const src = img.src;
 		// Create fallback element
@@ -869,13 +873,16 @@
 		: ''}"
 >
 	{#each renderedBlocks as block (block.id)}
-		<div class="markdown-block" data-block-id={block.id}>
+		<div class="markdown-block" {...{ [MARKDOWN_DATA_ATTRS.BLOCK_ID]: block.id }}>
 			{@html block.html}
 		</div>
 	{/each}
 
 	{#if unstableBlockHtml}
-		<div class="markdown-block markdown-block--unstable" data-block-id="unstable">
+		<div
+			class="markdown-block markdown-block--unstable"
+			{...{ [MARKDOWN_DATA_ATTRS.BLOCK_ID]: 'unstable' }}
+		>
 			<!-- eslint-disable-next-line no-at-html-tags -->
 			{@html unstableBlockHtml}
 		</div>

tools/ui/src/lib/components/app/content/MarkdownContent/markdown-handlers.ts

@@ -3,7 +3,14 @@
  * Uses dependency injection pattern to avoid direct component state access.
  */
 
-import { MERMAID_BLOCK_CLASS, MERMAID_SYNTAX_ATTR, MERMAID_WRAPPER_CLASS } from '$lib/constants';
+import {
+	CODE_BLOCK_CLASS,
+	MARKDOWN_DATA_ATTRS,
+	MERMAID_BLOCK_CLASS,
+	MERMAID_SYNTAX_ATTR,
+	MERMAID_WRAPPER_CLASS
+} from '$lib/constants';
+import { BooleanString } from '$lib/enums';
 import { copyCodeToClipboard, copyToClipboard } from '$lib/utils';
 
 export interface PreviewState {
@@ -40,11 +47,11 @@ export function createHandleCopyClick() {
 
 		if (!target) return;
 
-		const wrapper = target.closest('.code-block-wrapper');
+		const wrapper = target.closest(`.${CODE_BLOCK_CLASS.WRAPPER}`);
 
 		if (!wrapper) return;
 
-		const codeElement = wrapper.querySelector<HTMLElement>('code[data-code-id]');
+		const codeElement = wrapper.querySelector<HTMLElement>(`code[${MARKDOWN_DATA_ATTRS.CODE_ID}]`);
 
 		if (!codeElement) return;
 
@@ -86,16 +93,16 @@ export function createHandlePreviewClick(previewState: PreviewState) {
 
 		if (!target) return;
 
-		const wrapper = target.closest('.code-block-wrapper');
+		const wrapper = target.closest(`.${CODE_BLOCK_CLASS.WRAPPER}`);
 
 		if (!wrapper) return;
 
-		const codeElement = wrapper.querySelector<HTMLElement>('code[data-code-id]');
+		const codeElement = wrapper.querySelector<HTMLElement>(`code[${MARKDOWN_DATA_ATTRS.CODE_ID}]`);
 
 		if (!codeElement) return;
 
 		const rawCode = codeElement.textContent ?? '';
-		const languageLabel = wrapper.querySelector<HTMLElement>('.code-language');
+		const languageLabel = wrapper.querySelector<HTMLElement>(`.${CODE_BLOCK_CLASS.LANGUAGE}`);
 		const language = languageLabel?.textContent?.trim() || 'text';
 
 		previewState.setPreviewCode(rawCode);
@@ -112,8 +119,8 @@ export function createHandleMermaidClick(mermaidState: MermaidPreviewState) {
 	return async function handleMermaidClick(event: MouseEvent) {
 		const target = event.target as HTMLElement;
 		// Check if clicking on copy or preview button in mermaid block
-		const copyBtn = target.closest(`.${MERMAID_WRAPPER_CLASS} .copy-code-btn`);
-		const previewBtn = target.closest(`.${MERMAID_WRAPPER_CLASS} .preview-code-btn`);
+		const copyBtn = target.closest(`.${MERMAID_WRAPPER_CLASS} .${CODE_BLOCK_CLASS.COPY_BTN}`);
+		const previewBtn = target.closest(`.${MERMAID_WRAPPER_CLASS} .${CODE_BLOCK_CLASS.PREVIEW_BTN}`);
 
 		if (copyBtn || previewBtn) {
 			const wrapper = target.closest(`.${MERMAID_WRAPPER_CLASS}`);
@@ -189,15 +196,17 @@ export function createHandleMermaidPreviewOpenChange(mermaidState: MermaidPrevie
 export function createHandleImageError(
 	renderedBlocksState: RenderedBlocksState,
 	IMAGE_NOT_ERROR_BOUND_SELECTOR: string,
-	DATA_ERROR_BOUND_ATTR: string,
-	BOOL_TRUE_STRING: string
+	errorBoundAttr: string,
+	booleanString: BooleanString
 ) {
 	return async function handleImageError(event: Event) {
 		const img = event.target as HTMLImageElement;
 
 		if (!img) return;
 
-		const blockId = img.closest('[data-block-id]')?.getAttribute('data-block-id');
+		const blockId = img
+			.closest(`[${MARKDOWN_DATA_ATTRS.BLOCK_ID}]`)
+			?.getAttribute(MARKDOWN_DATA_ATTRS.BLOCK_ID);
 
 		if (!blockId) return;
 
@@ -206,19 +215,22 @@ export function createHandleImageError(
 		if (!block) return;
 
 		// Skip if already handled
-		if (img.dataset[DATA_ERROR_BOUND_ATTR] === BOOL_TRUE_STRING) return;
+		if (img.getAttribute(errorBoundAttr) === booleanString) return;
 
-		img.dataset[DATA_ERROR_BOUND_ATTR] = BOOL_TRUE_STRING;
+		img.setAttribute(errorBoundAttr, booleanString);
 
 		// Get the fallback HTML and replace the image
-		const fallbackHtml = `<div class="image-error-placeholder" data-original-src="${img.src}">
+		const fallbackHtml = `<div class="image-error-placeholder" ${MARKDOWN_DATA_ATTRS.ORIGINAL_SRC}="${img.src}">
 			<span class="image-error-icon">⚠️</span>
 			<span class="image-error-text">Failed to load image</span>
 		</div>`;
 		// Replace the img element with fallback in the block's HTML
 		const newHtml = block.html.replace(/img[^>]*src=["']([^"']*)[^>]*>/g, (match, src) => {
 			if (src === img.src) {
-				return fallbackHtml.replace('data-original-src=""', `data-original-src="${src}"`);
+				return fallbackHtml.replace(
+					`${MARKDOWN_DATA_ATTRS.ORIGINAL_SRC}=""`,
+					`${MARKDOWN_DATA_ATTRS.ORIGINAL_SRC}="${src}"`
+				);
 			}
 
 			return match;
@@ -243,19 +255,27 @@ export function createSetupCodeBlockActions(
 	return function setupCodeBlockActions(containerRef: HTMLElement | null) {
 		if (!containerRef) return;
 
-		const wrappers = containerRef.querySelectorAll<HTMLElement>('.code-block-wrapper');
+		const wrappers = containerRef.querySelectorAll<HTMLElement>(`.${CODE_BLOCK_CLASS.WRAPPER}`);
 
 		for (const wrapper of wrappers) {
-			const copyButton = wrapper.querySelector<HTMLButtonElement>('.copy-code-btn');
-			const previewButton = wrapper.querySelector<HTMLButtonElement>('.preview-code-btn');
+			const copyButton = wrapper.querySelector<HTMLButtonElement>(`.${CODE_BLOCK_CLASS.COPY_BTN}`);
+			const previewButton = wrapper.querySelector<HTMLButtonElement>(
+				`.${CODE_BLOCK_CLASS.PREVIEW_BTN}`
+			);
 
-			if (copyButton && copyButton.dataset.listenerBound !== 'true') {
-				copyButton.dataset.listenerBound = 'true';
+			if (
+				copyButton &&
+				copyButton.getAttribute(MARKDOWN_DATA_ATTRS.LISTENER_BOUND) !== BooleanString.TRUE
+			) {
+				copyButton.setAttribute(MARKDOWN_DATA_ATTRS.LISTENER_BOUND, BooleanString.TRUE);
 				copyButton.addEventListener('click', handleCopyClick);
 			}
 
-			if (previewButton && previewButton.dataset.listenerBound !== 'true') {
-				previewButton.dataset.listenerBound = 'true';
+			if (
+				previewButton &&
+				previewButton.getAttribute(MARKDOWN_DATA_ATTRS.LISTENER_BOUND) !== BooleanString.TRUE
+			) {
+				previewButton.setAttribute(MARKDOWN_DATA_ATTRS.LISTENER_BOUND, BooleanString.TRUE);
 				previewButton.addEventListener('click', handlePreviewClick);
 			}
 		}
@@ -269,8 +289,8 @@ export function createSetupCodeBlockActions(
 export function createSetupImageErrorHandlers(
 	handleImageError: (event: Event) => void,
 	IMAGE_NOT_ERROR_BOUND_SELECTOR: string,
-	DATA_ERROR_BOUND_ATTR: string,
-	BOOL_TRUE_STRING: string
+	errorBoundAttr: string,
+	booleanString: BooleanString
 ) {
 	return function setupImageErrorHandlers(containerRef: HTMLElement | null) {
 		if (!containerRef) return;
@@ -278,7 +298,7 @@ export function createSetupImageErrorHandlers(
 		const images = containerRef.querySelectorAll<HTMLImageElement>(IMAGE_NOT_ERROR_BOUND_SELECTOR);
 
 		for (const img of images) {
-			img.dataset[DATA_ERROR_BOUND_ATTR] = BOOL_TRUE_STRING;
+			img.setAttribute(errorBoundAttr, booleanString);
 			img.addEventListener('error', handleImageError);
 		}
 	};

tools/ui/src/lib/components/app/content/MarkdownContent/markdown-utils.ts

@@ -2,6 +2,7 @@
  * Utility functions for markdown processing in MarkdownContent component.
  */
 
+import { MARKDOWN_DATA_ATTRS } from '$lib/constants';
 import type { RootContent as HastRootContent } from 'hast';
 
 /**
@@ -69,7 +70,7 @@ export function getCodeInfoFromTarget(target: HTMLElement): CodeInfo | null {
 		return null;
 	}
 
-	const codeElement = wrapper.querySelector<HTMLElement>('code[data-code-id]');
+	const codeElement = wrapper.querySelector<HTMLElement>(`code[${MARKDOWN_DATA_ATTRS.CODE_ID}]`);
 
 	if (!codeElement) {
 		console.error('No code element found in wrapper');

tools/ui/src/lib/components/app/content/MarkdownContent/plugins/rehype/enhance-code-blocks.ts

@@ -17,7 +17,7 @@ import {
 	createWrapper,
 	generateBlockId
 } from './code-block-utils';
-import { CODE_BLOCK_CLASS } from '$lib/constants';
+import { CODE_BLOCK_CLASS, MARKDOWN_DATA_ATTRS } from '$lib/constants';
 import type { Element, ElementContent, Root } from 'hast';
 import type { Plugin } from 'unified';
 import { visit } from 'unist-util-visit';
@@ -65,16 +65,18 @@ export const rehypeEnhanceCodeBlocks: Plugin<[], Root> = () => {
 
 			codeElement.properties = {
 				...codeElement.properties,
-				'data-code-id': codeId
+				[MARKDOWN_DATA_ATTRS.CODE_ID]: codeId
 			};
 
-			const actions: Element[] = [createCopyButton(codeId, 'data-code-id', 'Copy code')];
+			const actions: Element[] = [
+				createCopyButton(codeId, MARKDOWN_DATA_ATTRS.CODE_ID, 'Copy code')
+			];
 
 			if (language.toLowerCase() === 'html') {
-				actions.push(createPreviewButton(codeId, 'data-code-id', 'Preview code'));
+				actions.push(createPreviewButton(codeId, MARKDOWN_DATA_ATTRS.CODE_ID, 'Preview code'));
 			}
 
-			const header = createBlockHeader(language, codeId, 'data-code-id', actions);
+			const header = createBlockHeader(language, codeId, MARKDOWN_DATA_ATTRS.CODE_ID, actions);
 			const wrapper = createWrapper(
 				header,
 				node,

tools/ui/src/lib/components/app/content/MarkdownContent/plugins/rehype/file-badge.ts

@@ -1,6 +1,6 @@
 /**
  * Rehype plugin that rewrites `file://` markdown anchors into the inline
- * mention chip, sharing the class string with the contenteditable
+ * mention chip, sharing the class string with the ChatFormInputRich
  * tokenizer via `$lib/constants`.
  *
  * The chip is presentational: `file://` navigation is blocked from

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

@@ -1,7 +1,7 @@
 <script lang="ts">
 	import { browser } from '$app/environment';
-	import { SYNTAX_CODE_SCROLL_AT_BOTTOM_THRESHOLD_PX } from '$lib/constants';
-	import { ColorMode } from '$lib/enums';
+	import { SYNTAX_CODE_SCROLL_AT_BOTTOM_THRESHOLD_PX, UI_DATA_ATTRS } from '$lib/constants';
+	import { BooleanString, ColorMode } from '$lib/enums';
 	import { highlightCode } from '$lib/utils';
 	import githubLightCss from 'highlight.js/styles/github.css?inline';
 	import githubDarkCss from 'highlight.js/styles/github-dark.css?inline';
@@ -38,13 +38,15 @@
 	function loadHighlightTheme(isDark: boolean) {
 		if (!browser) return;
 
-		const existingThemes = document.querySelectorAll('style[data-highlight-theme-preview]');
+		const existingThemes = document.querySelectorAll(
+			`style[${UI_DATA_ATTRS.HIGHLIGHT_THEME_PREVIEW}]`
+		);
 
 		existingThemes.forEach((style) => style.remove());
 
 		const style = document.createElement('style');
 
-		style.setAttribute('data-highlight-theme-preview', 'true');
+		style.setAttribute(UI_DATA_ATTRS.HIGHLIGHT_THEME_PREVIEW, BooleanString.TRUE);
 		style.textContent = isDark ? githubDarkCss : githubLightCss;
 
 		document.head.appendChild(style);

tools/ui/src/lib/components/app/dialogs/DialogMcpServerAddNew.svelte

@@ -4,14 +4,12 @@
 	import { Button } from '$lib/components/ui/button';
 	import * as Dialog from '$lib/components/ui/dialog';
 	import {
-		BOOL_FALSE_STRING,
-		BOOL_TRUE_STRING,
 		DISMISSED_RECOMMENDED_MCP_SERVERS_LOCALSTORAGE_KEY,
 		HEADERS,
 		MCP_SERVER_ID_PREFIX,
 		RECOMMENDED_MCP_SERVERS
 	} from '$lib/constants';
-	import { HealthCheckStatus } from '$lib/enums';
+	import { BooleanString, HealthCheckStatus } from '$lib/enums';
 	import { conversationsStore, mcpStore } from '$lib/stores';
 	import { canonicalizeServerUrl, parseHeadersToArray, uuid } from '$lib/utils';
 
@@ -97,9 +95,9 @@
 
 		if (!raw) return false;
 
-		if (raw === BOOL_TRUE_STRING) return true;
+		if (raw === BooleanString.TRUE) return true;
 
-		if (raw === BOOL_FALSE_STRING) return false;
+		if (raw === BooleanString.FALSE) return false;
 
 		try {
 			const parsed = JSON.parse(raw);
@@ -116,7 +114,7 @@
 		if (browser) {
 			localStorage.setItem(
 				DISMISSED_RECOMMENDED_MCP_SERVERS_LOCALSTORAGE_KEY,
-				dismissed ? BOOL_TRUE_STRING : BOOL_FALSE_STRING
+				dismissed ? BooleanString.TRUE : BooleanString.FALSE
 			);
 		}
 	}

tools/ui/src/lib/components/app/misc/ConversationSelection.svelte

@@ -3,6 +3,7 @@
 	import { Button } from '$lib/components/ui/button';
 	import { Checkbox } from '$lib/components/ui/checkbox';
 	import { ScrollArea } from '$lib/components/ui/scroll-area';
+	import { UI_DATA_ATTRS } from '$lib/constants';
 	import { useMarqueeSelection } from '$lib/hooks/use-marquee-selection.svelte';
 	import { SvelteSet } from 'svelte/reactivity';
 
@@ -138,7 +139,7 @@
 								class="cursor-pointer border-b transition-colors hover:bg-muted/50 {checked
 									? 'bg-muted/75'
 									: ''}"
-								data-conversation-row={conv.id}
+								{...{ [UI_DATA_ATTRS.CONVERSATION_ROW]: conv.id }}
 								onmousedown={(event) => marquee.rowMouseDown(conv.id, event)}
 								onclick={(event) => marquee.rowClick(conv.id, event.shiftKey)}
 							>

tools/ui/src/lib/components/app/navigation/SidebarNavigation/SidebarNavigationConversationItem.svelte

@@ -15,7 +15,7 @@
 	import { TruncatedText } from '$lib/components/app';
 	import { Checkbox } from '$lib/components/ui/checkbox';
 	import * as Tooltip from '$lib/components/ui/tooltip';
-	import { FORK_TREE_DEPTH_PADDING, ICON_CLASS_DEFAULT } from '$lib/constants';
+	import { FORK_TREE_DEPTH_PADDING, ICON_CLASS_DEFAULT, UI_DATA_ATTRS } from '$lib/constants';
 	import { RouterService } from '$lib/services/router.service';
 	import { chatStore, conversationsStore } from '$lib/stores';
 	import { onMount } from 'svelte';
@@ -154,14 +154,13 @@
 	});
 </script>
 
-<!-- svelte-ignore a11y_mouse_events_have_key_events -->
 <button
 	class="group flex min-h-9 w-full cursor-pointer items-center justify-between space-x-3 rounded-lg py-1.5 text-left transition-colors hover:bg-foreground/10 {isActive
 		? 'bg-foreground/5 text-accent-foreground'
 		: ''} {isSelected ? 'bg-primary/10 hover:bg-primary/15' : ''} {isSelectionMode
 		? 'is-selection-mode'
 		: ''} px-2"
-	data-conversation-row={conversation.id}
+	{...{ [UI_DATA_ATTRS.CONVERSATION_ROW]: conversation.id }}
 	onclick={(e) => handleSelect(e)}
 	onmouseover={handleMouseOver}
 	onmouseleave={handleMouseLeave}

tools/ui/src/lib/components/app/settings/SettingsChatMobileHeader.svelte

@@ -1,6 +1,7 @@
 <script lang="ts">
 	import { ChevronLeft, ChevronRight, Settings } from '@lucide/svelte';
-	import { ICON_CLASS_DEFAULT } from '$lib/constants';
+	import { ICON_CLASS_DEFAULT, UI_DATA_ATTRS } from '$lib/constants';
+	import { BooleanString } from '$lib/enums';
 	import { useScrollCarousel } from '$lib/hooks/use-scroll-carousel.svelte';
 	import type { SettingsSection, SettingsSectionTitle } from '$lib/types';
 	import { onMount, tick } from 'svelte';
@@ -20,7 +21,9 @@
 		await tick();
 
 		if (carousel.scrollContainer) {
-			const activeTab = carousel.scrollContainer.querySelector('[data-active="true"]');
+			const activeTab = carousel.scrollContainer.querySelector(
+				`[${UI_DATA_ATTRS.ACTIVE}="${BooleanString.TRUE}"]`
+			);
 
 			if (activeTab instanceof HTMLElement) {
 				carousel.scrollToCenter(activeTab);
@@ -66,7 +69,7 @@
 								)
 									? 'bg-accent text-accent-foreground'
 									: 'text-muted-foreground'}"
-								data-active={isActive(section)}
+								{...{ [UI_DATA_ATTRS.ACTIVE]: isActive(section) }}
 								href={getHref(section)}
 								onclick={(e: MouseEvent) => {
 									carousel.scrollToCenter(e.currentTarget as HTMLElement);
@@ -82,7 +85,7 @@
 								)
 									? 'bg-accent text-accent-foreground'
 									: 'text-muted-foreground'}"
-								data-active={isActive(section)}
+								{...{ [UI_DATA_ATTRS.ACTIVE]: isActive(section) }}
 								onclick={(e: MouseEvent) => {
 									onSectionChange?.(section.title);
 									carousel.scrollToCenter(e.currentTarget as HTMLElement);

tools/ui/src/lib/constants/chat-form.constants.ts

@@ -1,3 +1,6 @@
+/** Data attribute that tags ChatFormInputRich code spans and blocks. */
+export const CODE_TOKEN_ATTR = 'data-code-token';
+
 export const INITIAL_FILE_SIZE = 0;
 export const PROMPT_CONTENT_SEPARATOR = '\n\n';
 export const CLIPBOARD_CONTENT_QUOTE_PREFIX = '"';

tools/ui/src/lib/constants/headers.constants.ts

@@ -1,3 +1,6 @@
+/** Number of trailing characters to keep visible when partially redacting mcp-session-id */
+const MCP_SESSION_ID_VISIBLE_CHARS = 5;
+
 /** HTTP header handling for API and MCP requests. */
 export const HEADERS = {
 	/** Canonical casing for the Authorization header (RFC 7235) */
@@ -7,7 +10,7 @@ export const HEADERS = {
 	/** Content-Type HTTP header name */
 	CONTENT_TYPE: 'Content-Type',
 	/** Partial-redaction rules for MCP headers: header name -> visible trailing chars */
-	PARTIAL_REDACT: new Map<string, number>([['mcp-session-id', 5]]),
+	PARTIAL_REDACT: new Map<string, number>([['mcp-session-id', MCP_SESSION_ID_VISIBLE_CHARS]]),
 
 	/** Header names whose values should be redacted in diagnostic logs */
 	REDACTED: new Set([

tools/ui/src/lib/constants/markdown.constants.ts

@@ -1,8 +1,14 @@
 export const IMAGE_NOT_ERROR_BOUND_SELECTOR = 'img:not([data-error-bound])';
-export const DATA_ERROR_BOUND_ATTR = 'errorBound';
-export const DATA_ERROR_HANDLED_ATTR = 'errorHandled';
-export const BOOL_TRUE_STRING = 'true';
-export const BOOL_FALSE_STRING = 'false';
+
+/** Data attributes for the markdown renderer DOM contract. */
+export const MARKDOWN_DATA_ATTRS = {
+	BLOCK_ID: 'data-block-id',
+	CODE_ID: 'data-code-id',
+	ERROR_BOUND: 'data-error-bound',
+	ERROR_HANDLED: 'data-error-handled',
+	LISTENER_BOUND: 'data-listener-bound',
+	ORIGINAL_SRC: 'data-original-src'
+} as const;
 
 /** Markdown structural markers used by `looksLikeMarkdown`. Inline / line-level. */
 export const MARKDOWN = {

tools/ui/src/lib/constants/mention-badge.constants.ts

@@ -1,6 +1,6 @@
 /**
  * Shared visual contract between the two DOM-only badge paths (the
- * contenteditable tokenizer + the rehype plugin). Svelte cannot be
+ * ChatFormInputRich tokenizer + the rehype plugin). Svelte cannot be
  * mounted at the per-keystroke tokenizer hot path nor from a hast tree,
  * so both emit the badge with the same class string literal; Tailwind's
  * scanner picks it up in both sources.
@@ -10,6 +10,13 @@ export const MENTION_BADGE_CLASSNAME =
 
 export const MENTION_BADGE_ICON_CLASSNAME = 'h-3 w-3 shrink-0';
 
+/** Full `data-*` attribute names that tag ChatFormInputRich mention badges. */
+export const MENTION_BADGE_DATA_ATTRS = {
+	BADGE: 'data-mention-badge',
+	NAME: 'data-mention-name',
+	PATH: 'data-mention-path'
+} as const;
+
 /** Regex flag that makes the mention scanner walk every link in a message instead of the first. */
 export const MENTION_LINK_SCAN_FLAGS = 'g';
 

tools/ui/src/lib/constants/ui.constants.ts

@@ -7,6 +7,16 @@ import type { DesktopIconStripItem } from '$lib/types';
 export const FORK_TREE_DEPTH_PADDING = 8;
 export const SYSTEM_MESSAGE_PLACEHOLDER = 'System message';
 
+/** Data attributes for app-level DOM contracts. */
+export const UI_DATA_ATTRS = {
+	ACTIVE: 'data-active',
+	CONVERSATION_ROW: 'data-conversation-row',
+	HIGHLIGHT_THEME_PREVIEW: 'data-highlight-theme-preview',
+	PICKER_INDEX: 'data-picker-index',
+	RESULT_INDEX: 'data-result-index',
+	THUMBNAIL_INDEX: 'data-thumbnail-index'
+} as const;
+
 export const TOOL_GROUP_LABELS = {
 	[ToolSource.BUILTIN]: 'Built-in',
 	[ToolSource.CUSTOM]: 'JSON Schema',

tools/ui/src/lib/enums/boolean-string.enums.ts

@@ -0,0 +1,5 @@
+/** String representation of a boolean used in data attributes and persisted values. */
+export enum BooleanString {
+	TRUE = 'true',
+	FALSE = 'false'
+}

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

@@ -91,11 +91,11 @@ export enum FileMentionEntryType {
 }
 
 /**
- * Kinds of tokens the chat-form contenteditable produces.
+ * Kinds of tokens the chat-form-input-rich produces.
  */
-export enum ContentEditableTokenKind {
+export enum ChatFormInputRichTokenKind {
 	TEXT = 'text',
 	BADGE = 'badge',
-	INLINE_CODE = 'inlineCode',
-	CODE_BLOCK = 'codeBlock'
+	CODE_INLINE = 'code_inline',
+	CODE_BLOCK = 'code_block'
 }

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

@@ -28,11 +28,13 @@ export {
 	ReasoningFormat,
 	ChatFormCommandAction,
 	FileMentionEntryType,
-	ContentEditableTokenKind
+	ChatFormInputRichTokenKind
 } from './chat.enums';
 
 export { SessionRecordType } from './conversation-import.enums';
 
+export { BooleanString } from './boolean-string.enums';
+
 export { ReasoningEffort } from './reasoning-effort.enums';
 
 export {

tools/ui/src/lib/hooks/use-marquee-selection.svelte.ts

@@ -9,6 +9,7 @@
  * matches what the user sees on screen.
  */
 
+import { UI_DATA_ATTRS } from '$lib/constants';
 import { SvelteSet } from 'svelte/reactivity';
 
 interface UseMarqueeSelectionOptions {
@@ -18,8 +19,8 @@ interface UseMarqueeSelectionOptions {
 	orderedIds: () => string[];
 	/** Document listeners attach only while the getter returns true. */
 	enabled: () => boolean;
-	/** DOM attribute key (after the `data-` prefix) that marks selectable rows. */
-	attributeName?: () => string;
+	/** Full `data-*` attribute that marks selectable rows. */
+	dataAttr?: () => string;
 	/** Minimum pixel distance before a press becomes a marquee drag. */
 	dragThresholdPx?: number;
 }
@@ -36,16 +37,8 @@ export function useMarqueeSelection(options: UseMarqueeSelectionOptions) {
 	let dragMode: 'add' | 'remove' | null = null;
 	let suppressNextClick = false;
 
-	function resolveAttributeName(): string {
-		return options.attributeName?.() ?? 'conversation-row';
-	}
-
-	/**
-	 * `dataset` keys are camelCased. `data-conversation-row` -> `conversationRow`.
-	 * We resolve the attribute name once per call and read via the camelCase key.
-	 */
-	function datasetKey(key: string = resolveAttributeName()): string {
-		return key.replace(/-([a-z])/g, (_, c) => c.toUpperCase());
+	function resolveDataAttr(): string {
+		return options.dataAttr?.() ?? UI_DATA_ATTRS.CONVERSATION_ROW;
 	}
 
 	function decideDragMode(startingRowId: string | null, currentlySelected: ReadonlySet<string>) {
@@ -78,9 +71,8 @@ export function useMarqueeSelection(options: UseMarqueeSelectionOptions) {
 	}
 
 	function findRowAtPoint(x: number, y: number): string | null {
-		const attr = resolveAttributeName();
-		const selector = `[data-${attr}]`;
-		const key = datasetKey(attr);
+		const attr = resolveDataAttr();
+		const selector = `[${attr}]`;
 
 		let bestMatch: HTMLElement | null = null;
 		let bestCenterDistance = Infinity;
@@ -89,7 +81,7 @@ export function useMarqueeSelection(options: UseMarqueeSelectionOptions) {
 			const rect = row.getBoundingClientRect();
 
 			if (y >= rect.top && y <= rect.bottom && x >= rect.left && x <= rect.right) {
-				return row.dataset[key] ?? null;
+				return row.getAttribute(attr);
 			}
 
 			if (x >= rect.left && x <= rect.right) {
@@ -102,13 +94,12 @@ export function useMarqueeSelection(options: UseMarqueeSelectionOptions) {
 			}
 		}
 
-		return bestMatch ? (bestMatch.dataset[key] ?? null) : null;
+		return bestMatch ? bestMatch.getAttribute(attr) : null;
 	}
 
 	function updateMarqueeRect(currentX: number, currentY: number) {
-		const attr = resolveAttributeName();
-		const selector = `[data-${attr}]`;
-		const key = datasetKey(attr);
+		const attr = resolveDataAttr();
+		const selector = `[${attr}]`;
 		const selected = options.selectedIds();
 		const left = Math.min(dragStartX, currentX);
 		const top = Math.min(dragStartY, currentY);
@@ -117,7 +108,7 @@ export function useMarqueeSelection(options: UseMarqueeSelectionOptions) {
 		const visibleIds = new SvelteSet(options.orderedIds());
 
 		for (const row of document.querySelectorAll<HTMLElement>(selector)) {
-			const id = row.dataset[key];
+			const id = row.getAttribute(attr);
 
 			if (!id || !visibleIds.has(id)) continue;
 

tools/ui/src/lib/hooks/use-scroll-active-row.svelte.ts

@@ -11,8 +11,8 @@ export interface UseScrollActiveRowOptions {
 	getContainer: () => HTMLDivElement | null;
 	getIndex: () => number;
 	getCount: () => number;
-	/** Attribute prefix, e.g. 'picker' for `[data-picker-index="0"]`. */
-	dataIndex: string;
+	/** Full data attribute marking the row, e.g. `data-picker-index`. */
+	dataAttr: string;
 }
 
 export function useScrollActiveRow(opts: UseScrollActiveRowOptions) {
@@ -41,9 +41,7 @@ export function useScrollActiveRow(opts: UseScrollActiveRowOptions) {
 
 			if (!container || index < 0 || index >= opts.getCount()) return;
 
-			const row = container.querySelector(
-				`[data-${opts.dataIndex}-index="${index}"]`
-			) as HTMLElement | null;
+			const row = container.querySelector(`[${opts.dataAttr}="${index}"]`) as HTMLElement | null;
 
 			row?.scrollIntoView({ block: 'nearest', inline: 'nearest' });
 		});

tools/ui/src/lib/services/migration.service.ts

@@ -29,7 +29,7 @@ import {
 	STORAGE_APP_NAME,
 	STORAGE_APP_NAME_DEPRECATED
 } from '$lib/constants';
-import { MessageRole } from '$lib/enums';
+import { BooleanString, MessageRole } from '$lib/enums';
 import Dexie from 'dexie';
 
 // Types
@@ -613,10 +613,10 @@ const configTypesMigration: Migration = {
 		// schema rejects them. No config string field holds exactly "true"/"false", so the
 		// match is unambiguous.
 		for (const key of Object.keys(config)) {
-			if (config[key] === 'true') {
+			if (config[key] === BooleanString.TRUE) {
 				config[key] = true;
 				changed = true;
-			} else if (config[key] === 'false') {
+			} else if (config[key] === BooleanString.FALSE) {
 				config[key] = false;
 				changed = true;
 			}

tools/ui/src/lib/types/chat-form-input-rich.d.ts

@@ -0,0 +1,11 @@
+import { ChatFormInputRichTokenKind } from '$lib/enums';
+
+/**
+ * A single token produced by the chat-form-input-rich tokenizer:
+ * plain text, a file/folder mention badge, or an inline/fenced code span.
+ */
+export type ChatFormInputRichToken =
+	| { kind: ChatFormInputRichTokenKind.TEXT; text: string }
+	| { kind: ChatFormInputRichTokenKind.BADGE; name: string; path: string }
+	| { kind: ChatFormInputRichTokenKind.CODE_INLINE; text: string }
+	| { kind: ChatFormInputRichTokenKind.CODE_BLOCK; text: string };

tools/ui/src/lib/types/contenteditable.d.ts

@@ -1,11 +0,0 @@
-import { ContentEditableTokenKind } from '$lib/enums';
-
-/**
- * A single token produced by the chat-form contenteditable tokenizer:
- * plain text, a file/folder mention badge, or an inline/fenced code span.
- */
-export type ContentEditableToken =
-	| { kind: ContentEditableTokenKind.TEXT; text: string }
-	| { kind: ContentEditableTokenKind.BADGE; name: string; path: string }
-	| { kind: ContentEditableTokenKind.INLINE_CODE; text: string }
-	| { kind: ContentEditableTokenKind.CODE_BLOCK; text: string };

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

@@ -182,8 +182,8 @@ export type {
 	GlobSearchChildResult
 } from './glob';
 
-// Contenteditable token types (chat form)
-export type { ContentEditableToken } from './contenteditable';
+// ChatFormInputRich token types (chat form)
+export type { ChatFormInputRichToken } from './chat-form-input-rich';
 
 // Agentic types
 export type {

tools/ui/src/lib/utils/chat-form-input-rich-tokenizer.ts

@@ -1,5 +1,5 @@
 /**
- * Maps between the chat-form contenteditable's markdown source and the
+ * Maps between the chat-form-input-rich's markdown source and the
  * badge/code/text token stream the DOM is built from. A badge is one
  * opaque source contribution (`[name](file://path)`); its own subtree
  * is never walked, and the caret cannot land inside it, so offsets
@@ -30,15 +30,17 @@ import {
 	getMentionBadgeLabel
 } from './mention-badge';
 import {
+	CODE_TOKEN_ATTR,
 	MENTION_BADGE_CLASSNAME,
+	MENTION_BADGE_DATA_ATTRS,
 	MENTION_BADGE_ICON_CLASSNAME,
 	MENTION_BADGE_SVG_ATTRIBUTES,
 	SETTINGS_KEYS
 } from '$lib/constants';
-import { ContentEditableTokenKind } from '$lib/enums';
+import { BooleanString, ChatFormInputRichTokenKind } from '$lib/enums';
 import { settingsStore } from '$lib/stores/settings.svelte';
 import { toolsStore } from '$lib/stores/tools.svelte';
-import type { ContentEditableToken } from '$lib/types/contenteditable';
+import type { ChatFormInputRichToken } from '$lib/types/chat-form-input-rich';
 
 // Block wrappers browsers insert for newlines; each folds back into a
 // single `\n` during serialization.
@@ -66,7 +68,7 @@ const CODE_SPAN_RE = /(```[\s\S]*?```)|(`[^`\n]+`)/g;
 /**
  * Cheap gate check for `ChatForm`: does the buffer contain a
  * complete code span (inline or fenced)? Used to promote the plain
- * textarea to the contenteditable renderer.
+ * textarea to the chat-form-input-rich renderer.
  */
 export function containsCodeSpan(value: string): boolean {
 	CODE_SPAN_RE.lastIndex = 0;
@@ -102,14 +104,14 @@ export function isOffsetInCodeBlock(source: string, offset: number): boolean {
 
 /**
  * Tokenize a markdown source value into the segments the
- * contenteditable will render. Code spans are carved out first
+ * chat-form-input-rich will render. Code spans are carved out first
  * (their content is literal - a `file://` link inside backticks
  * must NOT render as a badge), then plain text and badges
  * interleave in the remaining gaps. Any whitespace after a badge
  * stays in a plain text token so the round trip is byte-exact.
  */
-export function tokenizeContent(input: string): ContentEditableToken[] {
-	const tokens: ContentEditableToken[] = [];
+export function tokenizeContent(input: string): ChatFormInputRichToken[] {
+	const tokens: ChatFormInputRichToken[] = [];
 
 	let cursor = 0;
 
@@ -126,8 +128,8 @@ export function tokenizeContent(input: string): ContentEditableToken[] {
 
 		tokens.push(
 			match[1] !== undefined
-				? { kind: ContentEditableTokenKind.CODE_BLOCK, text: match[1] }
-				: { kind: ContentEditableTokenKind.INLINE_CODE, text: match[2] }
+				? { kind: ChatFormInputRichTokenKind.CODE_BLOCK, text: match[1] }
+				: { kind: ChatFormInputRichTokenKind.CODE_INLINE, text: match[2] }
 		);
 		cursor = start + match[0].length;
 	}
@@ -142,7 +144,7 @@ export function tokenizeContent(input: string): ContentEditableToken[] {
 /**
  * Tokenize a code-free segment into text and badge tokens.
  */
-function pushTextAndBadgeTokens(input: string, tokens: ContentEditableToken[]) {
+function pushTextAndBadgeTokens(input: string, tokens: ChatFormInputRichToken[]) {
 	let cursor = 0;
 
 	MENTION_BADGE_RE.lastIndex = 0;
@@ -154,24 +156,27 @@ function pushTextAndBadgeTokens(input: string, tokens: ContentEditableToken[]) {
 		const start = match.index;
 
 		if (start > cursor) {
-			tokens.push({ kind: ContentEditableTokenKind.TEXT, text: input.slice(cursor, start) });
+			tokens.push({ kind: ChatFormInputRichTokenKind.TEXT, text: input.slice(cursor, start) });
 		}
 
-		tokens.push({ kind: ContentEditableTokenKind.BADGE, name, path });
+		tokens.push({ kind: ChatFormInputRichTokenKind.BADGE, name, path });
 		cursor = start + whole.length;
 	}
 
 	if (cursor < input.length) {
-		tokens.push({ kind: ContentEditableTokenKind.TEXT, text: input.slice(cursor) });
+		tokens.push({ kind: ChatFormInputRichTokenKind.TEXT, text: input.slice(cursor) });
 	}
 }
 
 function isCodeBlockElement(node: Node | null): node is HTMLElement {
-	return node instanceof HTMLElement && node.dataset.codeToken === 'block';
+	return (
+		node instanceof HTMLElement &&
+		node.getAttribute(CODE_TOKEN_ATTR) === ChatFormInputRichTokenKind.CODE_BLOCK
+	);
 }
 
 /**
- * Serialize a contenteditable subtree back to source. `<br>` and block
+ * Serialize a chat-form-input-rich subtree back to source. `<br>` and block
  * wrappers the browser inserted for newlines fold back into `\n` (a
  * trailing `<br>` is the browser's caret placeholder, not a newline);
  * any other element is transparent. Code spans serialize their
@@ -208,9 +213,9 @@ export function serializeContent(root: HTMLElement): string {
 
 			const el = child as HTMLElement;
 
-			if (el.dataset.mentionBadge === 'true') {
-				const name = el.dataset.mentionName ?? '';
-				const path = el.dataset.mentionPath ?? '';
+			if (el.getAttribute(MENTION_BADGE_DATA_ATTRS.BADGE) === BooleanString.TRUE) {
+				const name = el.getAttribute(MENTION_BADGE_DATA_ATTRS.NAME) ?? '';
+				const path = el.getAttribute(MENTION_BADGE_DATA_ATTRS.PATH) ?? '';
 
 				if (name && path) {
 					if (pendingBlockBoundary) {
@@ -225,8 +230,10 @@ export function serializeContent(root: HTMLElement): string {
 				continue;
 			}
 
-			if (el.dataset.codeToken !== undefined) {
-				const isBlock = el.dataset.codeToken === 'block';
+			const codeToken = el.getAttribute(CODE_TOKEN_ATTR);
+
+			if (codeToken !== null) {
+				const isBlock = codeToken === ChatFormInputRichTokenKind.CODE_BLOCK;
 
 				if (isBlock && (pendingBlockBoundary || !first)) out += '\n';
 
@@ -285,8 +292,8 @@ export function serializeContent(root: HTMLElement): string {
  * A mismatch means token boundaries shifted (a code span was just
  * completed or broken) and the DOM needs a rebuild to restyle.
  */
-export function domMatchesTokens(root: HTMLElement, tokens: ContentEditableToken[]): boolean {
-	const expected = tokens.filter((token) => token.kind !== ContentEditableTokenKind.TEXT);
+export function domMatchesTokens(root: HTMLElement, tokens: ChatFormInputRichToken[]): boolean {
+	const expected = tokens.filter((token) => token.kind !== ChatFormInputRichTokenKind.TEXT);
 
 	let index = 0;
 
@@ -295,8 +302,8 @@ export function domMatchesTokens(root: HTMLElement, tokens: ContentEditableToken
 			if (child.nodeType !== Node.ELEMENT_NODE) continue;
 
 			const el = child as HTMLElement;
-			const isBadge = el.dataset.mentionBadge === 'true';
-			const isCode = el.dataset.codeToken !== undefined;
+			const isBadge = el.getAttribute(MENTION_BADGE_DATA_ATTRS.BADGE) === BooleanString.TRUE;
+			const isCode = el.getAttribute(CODE_TOKEN_ATTR) !== null;
 
 			if (!isBadge && !isCode) {
 				if (!walk(el)) return false;
@@ -309,25 +316,25 @@ export function domMatchesTokens(root: HTMLElement, tokens: ContentEditableToken
 			if (!token) return false;
 
 			if (isBadge) {
-				if (token.kind !== ContentEditableTokenKind.BADGE) return false;
+				if (token.kind !== ChatFormInputRichTokenKind.BADGE) return false;
 
-				if (token.name !== (el.dataset.mentionName ?? '')) return false;
+				if (token.name !== (el.getAttribute(MENTION_BADGE_DATA_ATTRS.NAME) ?? '')) return false;
 
-				if (token.path !== (el.dataset.mentionPath ?? '')) return false;
+				if (token.path !== (el.getAttribute(MENTION_BADGE_DATA_ATTRS.PATH) ?? '')) return false;
 
 				continue;
 			}
 
-			const codeKind: ContentEditableTokenKind =
-				el.dataset.codeToken === 'block'
-					? ContentEditableTokenKind.CODE_BLOCK
-					: ContentEditableTokenKind.INLINE_CODE;
+			const codeKind: ChatFormInputRichTokenKind =
+				el.getAttribute(CODE_TOKEN_ATTR) === ChatFormInputRichTokenKind.CODE_BLOCK
+					? ChatFormInputRichTokenKind.CODE_BLOCK
+					: ChatFormInputRichTokenKind.CODE_INLINE;
 
 			if (token.kind !== codeKind) return false;
 
 			if (
-				token.kind === ContentEditableTokenKind.INLINE_CODE ||
-				token.kind === ContentEditableTokenKind.CODE_BLOCK
+				token.kind === ChatFormInputRichTokenKind.CODE_INLINE ||
+				token.kind === ChatFormInputRichTokenKind.CODE_BLOCK
 			) {
 				if (token.text !== (el.textContent ?? '')) return false;
 			}
@@ -428,8 +435,11 @@ export function rangeToTextOffset(root: HTMLElement, range: Range | null): numbe
 				total += 1;
 			}
 
-			if (el.dataset.mentionBadge === 'true') {
-				const len = badgeSourceLength(el.dataset.mentionName ?? '', el.dataset.mentionPath ?? '');
+			if (el.getAttribute(MENTION_BADGE_DATA_ATTRS.BADGE) === BooleanString.TRUE) {
+				const len = badgeSourceLength(
+					el.getAttribute(MENTION_BADGE_DATA_ATTRS.NAME) ?? '',
+					el.getAttribute(MENTION_BADGE_DATA_ATTRS.PATH) ?? ''
+				);
 
 				if (len === 0) continue;
 
@@ -445,8 +455,10 @@ export function rangeToTextOffset(root: HTMLElement, range: Range | null): numbe
 				continue;
 			}
 
-			if (el.dataset.codeToken !== undefined) {
-				const isBlock = el.dataset.codeToken === 'block';
+			const codeToken = el.getAttribute(CODE_TOKEN_ATTR);
+
+			if (codeToken !== null) {
+				const isBlock = codeToken === ChatFormInputRichTokenKind.CODE_BLOCK;
 
 				if (isBlock && !first) {
 					if (!atOrBeforeCaret(el, 0)) {
@@ -521,26 +533,29 @@ export function rangeToTextOffset(root: HTMLElement, range: Range | null): numbe
  * string + inline SVG are shared with the rehype plugin via
  * `$lib/constants`.
  */
-export function buildFragment(tokens: ContentEditableToken[]): DocumentFragment {
+export function buildFragment(tokens: ChatFormInputRichToken[]): DocumentFragment {
 	const fragment = document.createDocumentFragment();
 
 	for (let index = 0; index < tokens.length; index++) {
 		const token = tokens[index];
 
-		if (token.kind === ContentEditableTokenKind.TEXT) {
+		if (token.kind === ChatFormInputRichTokenKind.TEXT) {
 			let text = token.text;
 
 			// The separator \n at a fenced-block boundary is synthesized
 			// at serialization time; keeping it in the DOM would render a
 			// phantom empty line next to the block.
 			if (
-				tokens[index - 1]?.kind === ContentEditableTokenKind.CODE_BLOCK &&
+				tokens[index - 1]?.kind === ChatFormInputRichTokenKind.CODE_BLOCK &&
 				text.startsWith('\n')
 			) {
 				text = text.slice(1);
 			}
 
-			if (tokens[index + 1]?.kind === ContentEditableTokenKind.CODE_BLOCK && text.endsWith('\n')) {
+			if (
+				tokens[index + 1]?.kind === ChatFormInputRichTokenKind.CODE_BLOCK &&
+				text.endsWith('\n')
+			) {
 				text = text.slice(0, -1);
 			}
 
@@ -552,13 +567,12 @@ export function buildFragment(tokens: ContentEditableToken[]): DocumentFragment
 		}
 
 		if (
-			token.kind === ContentEditableTokenKind.INLINE_CODE ||
-			token.kind === ContentEditableTokenKind.CODE_BLOCK
+			token.kind === ChatFormInputRichTokenKind.CODE_INLINE ||
+			token.kind === ChatFormInputRichTokenKind.CODE_BLOCK
 		) {
 			const code = document.createElement('code');
 
-			code.dataset.codeToken =
-				token.kind === ContentEditableTokenKind.CODE_BLOCK ? 'block' : 'inline';
+			code.setAttribute(CODE_TOKEN_ATTR, token.kind);
 			code.textContent = token.text;
 			fragment.appendChild(code);
 
@@ -574,9 +588,9 @@ export function buildFragment(tokens: ContentEditableToken[]): DocumentFragment
 
 		const badge = document.createElement('span');
 
-		badge.dataset.mentionBadge = 'true';
-		badge.dataset.mentionName = token.name;
-		badge.dataset.mentionPath = token.path;
+		badge.setAttribute(MENTION_BADGE_DATA_ATTRS.BADGE, BooleanString.TRUE);
+		badge.setAttribute(MENTION_BADGE_DATA_ATTRS.NAME, token.name);
+		badge.setAttribute(MENTION_BADGE_DATA_ATTRS.PATH, token.path);
 		badge.title = decodeFileLinkPath(token.path);
 		badge.className = MENTION_BADGE_CLASSNAME;
 		badge.contentEditable = 'false';
@@ -734,14 +748,14 @@ export function badgeAwareWordJump(
 
 	for (const token of tokenizeContent(source)) {
 		const len =
-			token.kind === ContentEditableTokenKind.BADGE
+			token.kind === ChatFormInputRichTokenKind.BADGE
 				? badgeSourceLength(token.name, token.path)
 				: token.text.length;
 
-		if (token.kind === ContentEditableTokenKind.BADGE)
+		if (token.kind === ChatFormInputRichTokenKind.BADGE)
 			badgeSpans.push([masked.length, masked.length + len]);
 
-		masked += token.kind === ContentEditableTokenKind.BADGE ? 'a'.repeat(len) : token.text;
+		masked += token.kind === ChatFormInputRichTokenKind.BADGE ? 'a'.repeat(len) : token.text;
 	}
 
 	if (badgeSpans.length === 0) return null;
@@ -804,7 +818,7 @@ export function badgeAwareWordJump(
 export function leadingBadgeEdgeOffset(source: string, caret: number): number | null {
 	const [first] = tokenizeContent(source);
 
-	if (!first || first.kind !== ContentEditableTokenKind.BADGE) return null;
+	if (!first || first.kind !== ChatFormInputRichTokenKind.BADGE) return null;
 
 	return caret === badgeSourceLength(first.name, first.path) ? 0 : null;
 }
@@ -872,8 +886,11 @@ export function textOffsetToRange(root: HTMLElement, offset: number): Range {
 
 			const el = child as HTMLElement;
 
-			if (el.dataset.mentionBadge === 'true') {
-				const len = badgeSourceLength(el.dataset.mentionName ?? '', el.dataset.mentionPath ?? '');
+			if (el.getAttribute(MENTION_BADGE_DATA_ATTRS.BADGE) === BooleanString.TRUE) {
+				const len = badgeSourceLength(
+					el.getAttribute(MENTION_BADGE_DATA_ATTRS.NAME) ?? '',
+					el.getAttribute(MENTION_BADGE_DATA_ATTRS.PATH) ?? ''
+				);
 
 				if (len === 0) continue;
 
@@ -911,8 +928,10 @@ export function textOffsetToRange(root: HTMLElement, offset: number): Range {
 				continue;
 			}
 
-			if (el.dataset.codeToken !== undefined) {
-				const isBlock = el.dataset.codeToken === 'block';
+			const codeToken = el.getAttribute(CODE_TOKEN_ATTR);
+
+			if (codeToken !== null) {
+				const isBlock = codeToken === ChatFormInputRichTokenKind.CODE_BLOCK;
 
 				if (isBlock && (pendingBlockBoundary || !first)) {
 					pendingBlockBoundary = false;

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

@@ -207,7 +207,7 @@ export {
 	type CommandDismissSnapshot
 } from './command-token';
 
-// Tokenization for the chat-form contenteditable (mention links + code spans <-> chip DOM)
+// Tokenization for the ChatFormInputRich (mention links + code spans <-> chip DOM)
 export {
 	tokenizeContent,
 	containsCodeSpan,
@@ -221,12 +221,12 @@ export {
 	textOffsetToRange,
 	badgeAwareWordJump,
 	leadingBadgeEdgeOffset
-} from './contenteditable-tokenizer';
+} from './chat-form-input-rich-tokenizer';
 
-// Source-space undo/redo history for the chat-form contenteditable
+// Source-space undo/redo history for the ChatFormInputRich
 export { SourceHistory, type SourceHistoryEntry } from './source-history';
 
-// Mention-badge visual contract (used by the contenteditable / rehype
+// Mention-badge visual contract (used by the ChatFormInputRich / rehype
 // DOM paths that build the same chip without a Svelte mount)
 export {
 	containsFileMentionLink,

tools/ui/src/lib/utils/source-history.ts

@@ -1,5 +1,5 @@
 /**
- * Source-space undo/redo history for the chat-form contenteditable, whose
+ * Source-space undo/redo history for the ChatFormInputRich, whose
  * imperative DOM rebuilds destroy the browser's native undo stack.
  * Entries record the state BEFORE an edit; edits within `groupWindowMs`
  * extend the open group so a typing burst undoes as a unit, while

tools/ui/tests/client/chat-form-enter-code-block.svelte.test.ts

@@ -2,8 +2,8 @@
 // fenced-code-block flow: while the caret sits inside a fenced
 // block region - closed, or still OPEN while the user is typing
 // one - plain Enter adds a line instead of submitting the message.
-// The textarea path is covered here end-to-end (the contenteditable
-// consumes the same case locally; see chat-form-contenteditable).
+// The textarea path is covered here end-to-end (the ChatFormInputRich
+// consumes the same case locally; see chat-form-input-rich).
 
 import ChatFormTestWrapper from './components/ChatFormTestWrapper.svelte';
 import { SETTINGS_KEYS } from '$lib/constants';

tools/ui/tests/client/chat-form-input-rich-blocks.svelte.test.ts

@@ -1,9 +1,9 @@
-// Guards the newline contract of the chat-form contenteditable: browsers
+// Guards the newline contract of the ChatFormInputRich: browsers
 // restructure the flat DOM on Enter (`<div>` wrappers, `<br>` shapes) and
 // serialization must fold those back into `\n` so the emitted value never
 // diverges from what is on screen.
 
-import ChatFormContentEditableHarness from './components/ChatFormContentEditableHarness.svelte';
+import ChatFormInputRichHarness from './components/ChatFormInputRichHarness.svelte';
 import { tick } from 'svelte';
 import { describe, expect, it } from 'vitest';
 import { render } from 'vitest-browser-svelte';
@@ -13,7 +13,7 @@ const SOURCE = 'see [docs](file:///a/b) here';
 function editableIn(container: HTMLElement): HTMLElement {
 	const el = container.querySelector('[role="textbox"]');
 
-	if (!(el instanceof HTMLElement)) throw new Error('contenteditable not rendered');
+	if (!(el instanceof HTMLElement)) throw new Error('ChatFormInputRich not rendered');
 
 	return el;
 }
@@ -35,9 +35,9 @@ function setCaret(node: Node, offset: number) {
 	selection.addRange(range);
 }
 
-describe('ChatFormContentEditable browser newline shapes', () => {
+describe('ChatFormInputRich browser newline shapes', () => {
 	it('serializes a Chromium Enter <div> wrapper as a newline', async () => {
-		const screen = render(ChatFormContentEditableHarness, { value: SOURCE });
+		const screen = render(ChatFormInputRichHarness, { value: SOURCE });
 
 		await tick();
 
@@ -53,7 +53,7 @@ describe('ChatFormContentEditable browser newline shapes', () => {
 	});
 
 	it('serializes a Firefox full <div> wrap as lines, badge included', async () => {
-		const screen = render(ChatFormContentEditableHarness, { value: SOURCE });
+		const screen = render(ChatFormInputRichHarness, { value: SOURCE });
 
 		await tick();
 
@@ -73,7 +73,7 @@ describe('ChatFormContentEditable browser newline shapes', () => {
 	});
 
 	it('serializes a <br> as a newline', async () => {
-		const screen = render(ChatFormContentEditableHarness, { value: 'here' });
+		const screen = render(ChatFormInputRichHarness, { value: 'here' });
 
 		await tick();
 
@@ -88,7 +88,7 @@ describe('ChatFormContentEditable browser newline shapes', () => {
 	});
 
 	it('ignores a trailing <br> (browser caret placeholder)', async () => {
-		const screen = render(ChatFormContentEditableHarness, { value: 'abc' });
+		const screen = render(ChatFormInputRichHarness, { value: 'abc' });
 
 		await tick();
 
@@ -102,7 +102,7 @@ describe('ChatFormContentEditable browser newline shapes', () => {
 	});
 
 	it('serializes one newline per empty-line <div><br></div>', async () => {
-		const screen = render(ChatFormContentEditableHarness, { value: 'abc' });
+		const screen = render(ChatFormInputRichHarness, { value: 'abc' });
 
 		await tick();
 
@@ -121,7 +121,7 @@ describe('ChatFormContentEditable browser newline shapes', () => {
 	});
 
 	it('treats a <div><br></div>-only buffer as empty for the placeholder', async () => {
-		const screen = render(ChatFormContentEditableHarness, { value: 'abc' });
+		const screen = render(ChatFormInputRichHarness, { value: 'abc' });
 
 		await tick();
 
@@ -138,7 +138,7 @@ describe('ChatFormContentEditable browser newline shapes', () => {
 	});
 
 	it('maps the caret across block boundaries in both directions', async () => {
-		const screen = render(ChatFormContentEditableHarness, { value: 'abc\ndef' });
+		const screen = render(ChatFormInputRichHarness, { value: 'abc\ndef' });
 
 		await tick();
 

tools/ui/tests/client/chat-form-input-rich-undo.svelte.test.ts

@@ -1,9 +1,9 @@
-// Guards the editing-key contract of the chat-form contenteditable:
+// Guards the editing-key contract of the ChatFormInputRich:
 // undo/redo is replayed from source snapshots (the token rebuilds destroy
 // the native undo stack), and Tab is NOT intercepted (WCAG 2.1.2 no
 // keyboard trap), matching the plain textarea.
 
-import ChatFormContentEditableHarness from './components/ChatFormContentEditableHarness.svelte';
+import ChatFormInputRichHarness from './components/ChatFormInputRichHarness.svelte';
 import { tick } from 'svelte';
 import { describe, expect, it } from 'vitest';
 import { render } from 'vitest-browser-svelte';
@@ -13,7 +13,7 @@ const SOURCE = 'see [docs](file:///a/b)';
 function editableIn(container: HTMLElement): HTMLElement {
 	const el = container.querySelector('[role="textbox"]');
 
-	if (!(el instanceof HTMLElement)) throw new Error('contenteditable not rendered');
+	if (!(el instanceof HTMLElement)) throw new Error('ChatFormInputRich not rendered');
 
 	return el;
 }
@@ -31,9 +31,9 @@ function keydown(root: HTMLElement, init: KeyboardEventInit) {
 	return event;
 }
 
-describe('ChatFormContentEditable undo/redo', () => {
+describe('ChatFormInputRich undo/redo', () => {
 	it('undoes and redoes an edit across a badge-containing buffer', async () => {
-		const screen = render(ChatFormContentEditableHarness, { value: SOURCE });
+		const screen = render(ChatFormInputRichHarness, { value: SOURCE });
 
 		await tick();
 
@@ -57,7 +57,7 @@ describe('ChatFormContentEditable undo/redo', () => {
 	});
 
 	it('redoes with Ctrl+Y as well', async () => {
-		const screen = render(ChatFormContentEditableHarness, { value: SOURCE });
+		const screen = render(ChatFormInputRichHarness, { value: SOURCE });
 
 		await tick();
 
@@ -75,7 +75,7 @@ describe('ChatFormContentEditable undo/redo', () => {
 	});
 
 	it('coalesces a typing burst into one undo step', async () => {
-		const screen = render(ChatFormContentEditableHarness, { value: 'abc' });
+		const screen = render(ChatFormInputRichHarness, { value: 'abc' });
 
 		await tick();
 
@@ -92,7 +92,7 @@ describe('ChatFormContentEditable undo/redo', () => {
 	});
 
 	it('keeps a newline as its own undo step', async () => {
-		const screen = render(ChatFormContentEditableHarness, { value: 'abc' });
+		const screen = render(ChatFormInputRichHarness, { value: 'abc' });
 
 		await tick();
 
@@ -113,7 +113,7 @@ describe('ChatFormContentEditable undo/redo', () => {
 	});
 
 	it('is a no-op when there is nothing to undo', async () => {
-		const screen = render(ChatFormContentEditableHarness, { value: 'abc' });
+		const screen = render(ChatFormInputRichHarness, { value: 'abc' });
 
 		await tick();
 
@@ -127,7 +127,7 @@ describe('ChatFormContentEditable undo/redo', () => {
 	});
 
 	it('abandons the redo branch after a fresh edit', async () => {
-		const screen = render(ChatFormContentEditableHarness, { value: 'abc' });
+		const screen = render(ChatFormInputRichHarness, { value: 'abc' });
 
 		await tick();
 
@@ -147,9 +147,9 @@ describe('ChatFormContentEditable undo/redo', () => {
 	});
 });
 
-describe('ChatFormContentEditable Tab key', () => {
+describe('ChatFormInputRich Tab key', () => {
 	it('does not trap Tab (focus can leave the editable)', async () => {
-		const screen = render(ChatFormContentEditableHarness, { value: SOURCE });
+		const screen = render(ChatFormInputRichHarness, { value: SOURCE });
 
 		await tick();
 

tools/ui/tests/client/chat-form-input-rich.svelte.test.ts

@@ -1,9 +1,9 @@
-// Guards the clipboard contract of the chat-form contenteditable:
+// Guards the clipboard contract of the ChatFormInputRich:
 // copy/cut expose the markdown SOURCE of the selection (each badge
 // contributes its full `[name](file://...)` link) and pasting such
 // markdown re-renders the badges.
 
-import ChatFormContentEditable from '$lib/components/app/chat/ChatForm/ChatFormContentEditable.svelte';
+import ChatFormInputRich from '$lib/components/app/chat/ChatForm/ChatFormInput/ChatFormInputRich.svelte';
 import { rangeToTextOffset, serializeContent, textOffsetToRange } from '$lib/utils';
 import { tick } from 'svelte';
 import { describe, expect, it, vi } from 'vitest';
@@ -16,7 +16,7 @@ const BADGE_SELECTOR = '[data-mention-badge="true"]';
 function editableIn(container: HTMLElement): HTMLElement {
 	const el = container.querySelector('[role="textbox"]');
 
-	if (!(el instanceof HTMLElement)) throw new Error('contenteditable not rendered');
+	if (!(el instanceof HTMLElement)) throw new Error('ChatFormInputRich not rendered');
 
 	return el;
 }
@@ -43,9 +43,9 @@ function clipboardEvent(type: 'copy' | 'cut' | 'paste', text = '') {
 	return { data, event };
 }
 
-describe('ChatFormContentEditable clipboard', () => {
+describe('ChatFormInputRich clipboard', () => {
 	it('copy exposes the markdown source of the selection', async () => {
-		const { container } = render(ChatFormContentEditable, { value: SOURCE });
+		const { container } = render(ChatFormInputRich, { value: SOURCE });
 
 		await tick();
 
@@ -62,7 +62,7 @@ describe('ChatFormContentEditable clipboard', () => {
 	});
 
 	it('cut exposes the markdown source and removes the slice', async () => {
-		const { container } = render(ChatFormContentEditable, { value: SOURCE });
+		const { container } = render(ChatFormInputRich, { value: SOURCE });
 
 		await tick();
 
@@ -88,7 +88,7 @@ describe('ChatFormContentEditable clipboard', () => {
 	});
 
 	it('paste of markdown mention links re-renders badges', async () => {
-		const { container } = render(ChatFormContentEditable, { value: 'hello ' });
+		const { container } = render(ChatFormInputRich, { value: 'hello ' });
 
 		await tick();
 
@@ -114,7 +114,7 @@ describe('ChatFormContentEditable clipboard', () => {
 	});
 
 	it('paste without mention links keeps the DOM untouched', async () => {
-		const { container } = render(ChatFormContentEditable, { value: 'hello ' });
+		const { container } = render(ChatFormInputRich, { value: 'hello ' });
 
 		await tick();
 
@@ -138,14 +138,14 @@ describe('ChatFormContentEditable clipboard', () => {
 	});
 });
 
-describe('ChatFormContentEditable code spans', () => {
+describe('ChatFormInputRich code spans', () => {
 	it('renders inline code from the initial value', async () => {
-		const { container } = render(ChatFormContentEditable, { value: 'run `npm test` now' });
+		const { container } = render(ChatFormInputRich, { value: 'run `npm test` now' });
 
 		await tick();
 
 		const root = editableIn(container);
-		const code = root.querySelector('code[data-code-token="inline"]');
+		const code = root.querySelector('code[data-code-token="code_inline"]');
 
 		expect(code).not.toBeNull();
 		expect(code!.textContent).toBe('`npm test`');
@@ -153,12 +153,12 @@ describe('ChatFormContentEditable code spans', () => {
 
 	it('renders a fenced code block with a language', async () => {
 		const source = 'before\n```js\nconst a = 1;\n```\nafter';
-		const { container } = render(ChatFormContentEditable, { value: source });
+		const { container } = render(ChatFormInputRich, { value: source });
 
 		await tick();
 
 		const root = editableIn(container);
-		const code = root.querySelector('code[data-code-token="block"]');
+		const code = root.querySelector('code[data-code-token="code_block"]');
 
 		expect(code).not.toBeNull();
 		expect(code!.textContent).toBe('```js\nconst a = 1;\n```');
@@ -166,7 +166,7 @@ describe('ChatFormContentEditable code spans', () => {
 
 	it('copy exposes the markdown source of a selection spanning code', async () => {
 		const source = 'run `npm test` now';
-		const { container } = render(ChatFormContentEditable, { value: source });
+		const { container } = render(ChatFormInputRich, { value: source });
 
 		await tick();
 
@@ -183,7 +183,7 @@ describe('ChatFormContentEditable code spans', () => {
 	});
 
 	it('paste of a code span renders the styled element', async () => {
-		const { container } = render(ChatFormContentEditable, { value: 'run ' });
+		const { container } = render(ChatFormInputRich, { value: 'run ' });
 
 		await tick();
 
@@ -201,7 +201,7 @@ describe('ChatFormContentEditable code spans', () => {
 		await tick();
 
 		expect(event.defaultPrevented).toBe(true);
-		const code = root.querySelector('code[data-code-token="inline"]');
+		const code = root.querySelector('code[data-code-token="code_inline"]');
 
 		expect(code).not.toBeNull();
 		expect(code!.textContent).toBe('`npm test`');
@@ -210,12 +210,12 @@ describe('ChatFormContentEditable code spans', () => {
 
 	it('highlights a fenced block content and stays byte-exact', async () => {
 		const source = '```js\nconst a = 1;\n```';
-		const { container } = render(ChatFormContentEditable, { value: source });
+		const { container } = render(ChatFormInputRich, { value: source });
 
 		await tick();
 
 		const root = editableIn(container);
-		const code = root.querySelector('code[data-code-token="block"]');
+		const code = root.querySelector('code[data-code-token="code_block"]');
 
 		expect(code).not.toBeNull();
 		expect(code!.querySelector('.hljs-keyword')).not.toBeNull();
@@ -223,7 +223,7 @@ describe('ChatFormContentEditable code spans', () => {
 	});
 
 	it('does not highlight inline code', async () => {
-		const { container } = render(ChatFormContentEditable, { value: 'run `const` now' });
+		const { container } = render(ChatFormInputRich, { value: 'run `const` now' });
 
 		await tick();
 
@@ -233,9 +233,9 @@ describe('ChatFormContentEditable code spans', () => {
 	});
 });
 
-describe('ChatFormContentEditable code block escape hatches', () => {
+describe('ChatFormInputRich code block escape hatches', () => {
 	const BLOCK_SOURCE = '```js\nconst a = 1;\n```';
-	const BLOCK_SELECTOR = 'code[data-code-token="block"]';
+	const BLOCK_SELECTOR = 'code[data-code-token="code_block"]';
 
 	function blockIn(root: HTMLElement): HTMLElement {
 		const el = root.querySelector(BLOCK_SELECTOR);
@@ -273,7 +273,7 @@ describe('ChatFormContentEditable code block escape hatches', () => {
 	}
 
 	it('pads a trailing code block with a br hatch that stays invisible to copy', async () => {
-		const { container } = render(ChatFormContentEditable, { value: BLOCK_SOURCE });
+		const { container } = render(ChatFormInputRich, { value: BLOCK_SOURCE });
 
 		await tick();
 
@@ -292,7 +292,7 @@ describe('ChatFormContentEditable code block escape hatches', () => {
 	});
 
 	it('escapes a trailing code block with ArrowDown and types after it', async () => {
-		const { container } = render(ChatFormContentEditable, { value: BLOCK_SOURCE });
+		const { container } = render(ChatFormInputRich, { value: BLOCK_SOURCE });
 
 		await tick();
 
@@ -318,7 +318,7 @@ describe('ChatFormContentEditable code block escape hatches', () => {
 	});
 
 	it('escapes a leading code block with ArrowUp and types before it', async () => {
-		const { container } = render(ChatFormContentEditable, { value: BLOCK_SOURCE });
+		const { container } = render(ChatFormInputRich, { value: BLOCK_SOURCE });
 
 		await tick();
 
@@ -343,7 +343,7 @@ describe('ChatFormContentEditable code block escape hatches', () => {
 	});
 
 	it('escapes a leading code block with ArrowLeft from its first character', async () => {
-		const { container } = render(ChatFormContentEditable, { value: BLOCK_SOURCE });
+		const { container } = render(ChatFormInputRich, { value: BLOCK_SOURCE });
 
 		await tick();
 
@@ -358,7 +358,7 @@ describe('ChatFormContentEditable code block escape hatches', () => {
 	});
 
 	it('removes the transient leading hatch when the caret moves back into the block', async () => {
-		const { container } = render(ChatFormContentEditable, { value: BLOCK_SOURCE });
+		const { container } = render(ChatFormInputRich, { value: BLOCK_SOURCE });
 
 		await tick();
 
@@ -378,7 +378,7 @@ describe('ChatFormContentEditable code block escape hatches', () => {
 	});
 
 	it('extends the selection out of the block with Shift+ArrowDown', async () => {
-		const { container } = render(ChatFormContentEditable, { value: BLOCK_SOURCE });
+		const { container } = render(ChatFormInputRich, { value: BLOCK_SOURCE });
 
 		await tick();
 
@@ -397,7 +397,7 @@ describe('ChatFormContentEditable code block escape hatches', () => {
 	});
 
 	it('line-separates text typed right after the closing fence', async () => {
-		const { container } = render(ChatFormContentEditable, { value: BLOCK_SOURCE });
+		const { container } = render(ChatFormInputRich, { value: BLOCK_SOURCE });
 
 		await tick();
 
@@ -419,7 +419,7 @@ describe('ChatFormContentEditable code block escape hatches', () => {
 	});
 
 	it('does not double the newline when Shift+Enter already added one', async () => {
-		const { container } = render(ChatFormContentEditable, { value: BLOCK_SOURCE });
+		const { container } = render(ChatFormInputRich, { value: BLOCK_SOURCE });
 
 		await tick();
 
@@ -437,7 +437,7 @@ describe('ChatFormContentEditable code block escape hatches', () => {
 	});
 
 	it('moves a caret stuck before the inserted newline onto the new line', async () => {
-		const { container } = render(ChatFormContentEditable, {
+		const { container } = render(ChatFormInputRich, {
 			value: BLOCK_SOURCE + '\ntext after the code block'
 		});
 
@@ -469,7 +469,7 @@ describe('ChatFormContentEditable code block escape hatches', () => {
 	});
 
 	it('appends the artificial trailing newline when the browser did not add one', async () => {
-		const { container } = render(ChatFormContentEditable, {
+		const { container } = render(ChatFormInputRich, {
 			value: BLOCK_SOURCE + '\ntext after the code block'
 		});
 
@@ -501,7 +501,7 @@ describe('ChatFormContentEditable code block escape hatches', () => {
 	});
 
 	it('lands the caret on the new line with a single Shift+Enter after text below a block', async () => {
-		const { container } = render(ChatFormContentEditable, {
+		const { container } = render(ChatFormInputRich, {
 			value: BLOCK_SOURCE + '\ntext after the code block'
 		});
 
@@ -534,7 +534,7 @@ describe('ChatFormContentEditable code block escape hatches', () => {
 	});
 
 	it('lets Backspace at the text start move into the block without a source fight', async () => {
-		const { container } = render(ChatFormContentEditable, { value: BLOCK_SOURCE });
+		const { container } = render(ChatFormInputRich, { value: BLOCK_SOURCE });
 
 		await tick();
 
@@ -560,7 +560,7 @@ describe('ChatFormContentEditable code block escape hatches', () => {
 	});
 
 	it('lets forward Delete eat the text after a block normally', async () => {
-		const { container } = render(ChatFormContentEditable, { value: BLOCK_SOURCE });
+		const { container } = render(ChatFormInputRich, { value: BLOCK_SOURCE });
 
 		await tick();
 
@@ -581,7 +581,7 @@ describe('ChatFormContentEditable code block escape hatches', () => {
 	});
 
 	it('renders text after a block without a phantom empty line', async () => {
-		const { container } = render(ChatFormContentEditable, {
+		const { container } = render(ChatFormInputRich, {
 			value: BLOCK_SOURCE + '\nhello'
 		});
 
@@ -594,7 +594,7 @@ describe('ChatFormContentEditable code block escape hatches', () => {
 	});
 
 	it('keeps an intentional blank line after a block out of the separator', async () => {
-		const { container } = render(ChatFormContentEditable, {
+		const { container } = render(ChatFormInputRich, {
 			value: BLOCK_SOURCE + '\n\nhello'
 		});
 
@@ -607,7 +607,7 @@ describe('ChatFormContentEditable code block escape hatches', () => {
 	});
 
 	it('re-highlights while typing inside a block and keeps the caret', async () => {
-		const { container } = render(ChatFormContentEditable, { value: BLOCK_SOURCE });
+		const { container } = render(ChatFormInputRich, { value: BLOCK_SOURCE });
 
 		await tick();
 
@@ -639,12 +639,12 @@ describe('ChatFormContentEditable code block escape hatches', () => {
 	});
 });
 
-describe('ChatFormContentEditable Enter in code blocks', () => {
+describe('ChatFormInputRich Enter in code blocks', () => {
 	const BLOCK_SOURCE = '```js\nconst a = 1;\n```';
 
 	it('adds a line instead of submitting on plain Enter inside a block', async () => {
 		const onKeydown = vi.fn();
-		const { container } = render(ChatFormContentEditable, {
+		const { container } = render(ChatFormInputRich, {
 			onKeydown,
 			value: BLOCK_SOURCE
 		});
@@ -670,7 +670,7 @@ describe('ChatFormContentEditable Enter in code blocks', () => {
 		expect(onKeydown).not.toHaveBeenCalled();
 		expect(serializeContent(root)).toBe('```js\n\nconst a = 1;\n```');
 
-		const code = root.querySelector('code[data-code-token="block"]');
+		const code = root.querySelector('code[data-code-token="code_block"]');
 		const selection = window.getSelection();
 
 		expect(code!.contains(selection!.getRangeAt(0).startContainer)).toBe(true);
@@ -679,7 +679,7 @@ describe('ChatFormContentEditable Enter in code blocks', () => {
 
 	it('adds a line after a still-open fence (no closing ``` yet)', async () => {
 		const onKeydown = vi.fn();
-		const { container } = render(ChatFormContentEditable, {
+		const { container } = render(ChatFormInputRich, {
 			onKeydown,
 			value: '```js\nconst a = 1;'
 		});
@@ -707,7 +707,7 @@ describe('ChatFormContentEditable Enter in code blocks', () => {
 
 	it('forwards plain Enter to the parent when the caret is outside a block', async () => {
 		const onKeydown = vi.fn();
-		const { container } = render(ChatFormContentEditable, {
+		const { container } = render(ChatFormInputRich, {
 			onKeydown,
 			value: BLOCK_SOURCE + '\nafter'
 		});
@@ -730,7 +730,7 @@ describe('ChatFormContentEditable Enter in code blocks', () => {
 
 	it('forwards plain Enter on the trailing hatch line after a block', async () => {
 		const onKeydown = vi.fn();
-		const { container } = render(ChatFormContentEditable, {
+		const { container } = render(ChatFormInputRich, {
 			onKeydown,
 			value: BLOCK_SOURCE
 		});
@@ -753,7 +753,7 @@ describe('ChatFormContentEditable Enter in code blocks', () => {
 
 	it('forwards Ctrl+Enter inside a block so explicit submit survives', async () => {
 		const onKeydown = vi.fn();
-		const { container } = render(ChatFormContentEditable, {
+		const { container } = render(ChatFormInputRich, {
 			onKeydown,
 			value: BLOCK_SOURCE
 		});
@@ -780,7 +780,7 @@ describe('ChatFormContentEditable Enter in code blocks', () => {
 
 	it('forwards Enter inside an inline code span', async () => {
 		const onKeydown = vi.fn();
-		const { container } = render(ChatFormContentEditable, {
+		const { container } = render(ChatFormInputRich, {
 			onKeydown,
 			value: 'run `npm test` now'
 		});
@@ -790,7 +790,7 @@ describe('ChatFormContentEditable Enter in code blocks', () => {
 		const root = editableIn(container);
 
 		root.focus();
-		const code = root.querySelector('code[data-code-token="inline"]')!;
+		const code = root.querySelector('code[data-code-token="code_inline"]')!;
 
 		setSelection(root, (range) => {
 			range.setStart(code.firstChild!, 3);

tools/ui/tests/client/chat-form-mention-picker-gate.svelte.test.ts

@@ -3,7 +3,7 @@
 // it, the picker still opens but explains why instead of firing searches
 // that would only fail with "Search failed".
 
-import ChatFormMentionPicker from '$lib/components/app/chat/ChatForm/ChatFormPickers/ChatFormMentionPicker.svelte';
+import ChatFormPickerMention from '$lib/components/app/chat/ChatForm/ChatFormPickers/ChatFormPickerMention.svelte';
 import { DISABLED_TOOL_KEYS_LOCALSTORAGE_KEY } from '$lib/constants';
 import { BuiltInTool } from '$lib/enums';
 import { toolsStore } from '$lib/stores/tools.svelte';
@@ -25,7 +25,7 @@ function setBuiltinTools(defs: OpenAIToolDefinition[]) {
 }
 
 function renderPicker() {
-	return render(ChatFormMentionPicker, {
+	return render(ChatFormPickerMention, {
 		isOpen: true,
 		onClose: () => {},
 		onSelect: () => {},
@@ -39,7 +39,7 @@ afterEach(() => {
 	localStorage.removeItem(DISABLED_TOOL_KEYS_LOCALSTORAGE_KEY);
 });
 
-describe('ChatFormMentionPicker file_glob_search gate', () => {
+describe('ChatFormPickerMention file_glob_search gate', () => {
 	it('explains that file search is unavailable when the server has no tools', async () => {
 		setBuiltinTools([]);
 		renderPicker();

tools/ui/tests/client/components/ChatFormInputRichHarness.svelte

@@ -1,5 +1,5 @@
 <script lang="ts">
-	import ChatFormContentEditable from '$lib/components/app/chat/ChatForm/ChatFormContentEditable.svelte';
+	import ChatFormInputRich from '$lib/components/app/chat/ChatForm/ChatFormInput/ChatFormInputRich.svelte';
 	import { untrack } from 'svelte';
 
 	interface Props {
@@ -9,7 +9,7 @@
 	let { value: initial = '' }: Props = $props();
 
 	let value = $state(untrack(() => initial));
-	let inputRef: ChatFormContentEditable | undefined = $state(undefined);
+	let inputRef: ChatFormInputRich | undefined = $state(undefined);
 
 	export function getValue() {
 		return value;
@@ -24,4 +24,4 @@
 	}
 </script>
 
-<ChatFormContentEditable bind:this={inputRef} bind:value />
+<ChatFormInputRich bind:this={inputRef} bind:value />

tools/ui/tests/unit/chat-form-input-rich-tokenizer.test.ts

@@ -97,7 +97,7 @@ describe('tokenizeContent', () => {
 	it('tokenizes inline code with the backticks included', () => {
 		expect(tokenizeContent('run `npm test` now')).toEqual([
 			{ kind: 'text', text: 'run ' },
-			{ kind: 'inlineCode', text: '`npm test`' },
+			{ kind: 'code_inline', text: '`npm test`' },
 			{ kind: 'text', text: ' now' }
 		]);
 	});
@@ -107,7 +107,7 @@ describe('tokenizeContent', () => {
 
 		expect(tokenizeContent(source)).toEqual([
 			{ kind: 'text', text: 'before\n' },
-			{ kind: 'codeBlock', text: '```\nconst a = 1;\n```' },
+			{ kind: 'code_block', text: '```\nconst a = 1;\n```' },
 			{ kind: 'text', text: '\nafter' }
 		]);
 	});
@@ -116,15 +116,15 @@ describe('tokenizeContent', () => {
 		const source = '```js\nconst a = 1;\n```';
 
 		expect(tokenizeContent(source)).toEqual([
-			{ kind: 'codeBlock', text: '```js\nconst a = 1;\n```' }
+			{ kind: 'code_block', text: '```js\nconst a = 1;\n```' }
 		]);
 	});
 
 	it('prefers the fenced block over inline spans at triple backticks', () => {
 		expect(tokenizeContent('```a``` ```b```')).toEqual([
-			{ kind: 'codeBlock', text: '```a```' },
+			{ kind: 'code_block', text: '```a```' },
 			{ kind: 'text', text: ' ' },
-			{ kind: 'codeBlock', text: '```b```' }
+			{ kind: 'code_block', text: '```b```' }
 		]);
 	});
 
@@ -140,7 +140,7 @@ describe('tokenizeContent', () => {
 
 	it('does not recognize badges inside code spans', () => {
 		expect(tokenizeContent('`[a](file:///p)`')).toEqual([
-			{ kind: 'inlineCode', text: '`[a](file:///p)`' }
+			{ kind: 'code_inline', text: '`[a](file:///p)`' }
 		]);
 	});
 
@@ -148,7 +148,7 @@ describe('tokenizeContent', () => {
 		expect(tokenizeContent('[a](file:///p) `x`')).toEqual([
 			{ kind: 'badge', name: 'a', path: '/p' },
 			{ kind: 'text', text: ' ' },
-			{ kind: 'inlineCode', text: '`x`' }
+			{ kind: 'code_inline', text: '`x`' }
 		]);
 	});
 });