opencl: disable FA and MoE weights repack to work around compiler issues for Adreno 850 GPU (#25745)

* opencl: workaround for A850 compiler compat

* opencl: fix DX compiler version parsing and cleanup

---------

Co-authored-by: Li He <lih@qti.qualcomm.com>

.github/workflows/build-self-hosted.yml

@@ -6,7 +6,7 @@ on:
     branches:
       - master
     paths: [
-      '.github/workflows/build.yml',
+      '.github/workflows/build-self-hosted.yml',
       '**/CMakeLists.txt',
       '**/.cmake',
       '**/*.h',
@@ -48,6 +48,8 @@ concurrency:
   cancel-in-progress: true
 
 env:
+  # note: this is dud token to avoid rate limiting (https://github.com/ggml-org/llama.cpp/pull/25706#issuecomment-4979941302)
+  HF_TOKEN: ${{ secrets.HF_TOKEN_CI }}
   GGML_NLOOP: 3
   GGML_N_THREADS: 1
   LLAMA_ARG_LOG_COLORS: 1

.github/workflows/server-self-hosted.yml

@@ -29,6 +29,8 @@ on:
     ]
 
 env:
+  # note: this is dud token to avoid rate limiting (https://github.com/ggml-org/llama.cpp/pull/25706#issuecomment-4979941302)
+  HF_TOKEN: ${{ secrets.HF_TOKEN_CI }}
   LLAMA_ARG_LOG_COLORS: 1
   LLAMA_ARG_LOG_PREFIX: 1
   LLAMA_ARG_LOG_TIMESTAMPS: 1

common/arg.cpp

@@ -2932,7 +2932,7 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
     ).set_examples({LLAMA_EXAMPLE_TOKENIZE}));
     add_opt(common_arg(
         {"--stdin"},
-        string_format("read the prompt from stdin (mutually exclusive with -f/--file and -p/--prompt) (default: %s)", params.tokenize_stdin ? "true" : "false"),
+        string_format("read the prompt from stdin (takes precedence over -f/--file and -p/--prompt) (default: %s)", params.tokenize_stdin ? "true" : "false"),
         [](common_params & params) {
             params.tokenize_stdin = true;
         }

common/speculative.cpp

@@ -260,7 +260,10 @@ struct common_speculative_impl_draft_simple : public common_speculative_impl {
     bool process(const llama_batch & batch) override {
         auto * ctx_dft = params.ctx_dft;
 
-        const int ret = llama_decode(ctx_dft, batch);
+        llama_batch batch_dft = batch;
+        batch_dft.logits = nullptr;
+
+        const int ret = llama_decode(ctx_dft, batch_dft);
 
         if (ret != 0) {
             SPC_ERR("failed to decode draft batch, ret = %d\n", ret);

docs/ops.md

@@ -120,4 +120,4 @@ Legend:
 |                              TRI | ❌ | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                            TRUNC | ❌ | ❌ | ✅ | 🟡 | 🟡 | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                          UPSCALE | ❌ | 🟡 | ✅ | ✅ | ❌ | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                            XIELU | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ |
+|                            XIELU | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |

docs/ops/SYCL.csv

@@ -11600,10 +11600,10 @@ zjy 2
 "SYCL0","CUMSUM","type=f32,ne=[242004,1,1,1]","support","1","yes","SYCL"
 "SYCL0","CUMSUM","type=f32,ne=[375960,1,1,1]","support","1","yes","SYCL"
 "SYCL0","CUMSUM","type=f32,ne=[20481,4,1,1]","support","1","yes","SYCL"
-"SYCL0","XIELU","type=f32,ne=[10,5,4,3]","support","0","no","SYCL"
-"SYCL0","XIELU","type=f16,ne=[10,5,4,3]","support","0","no","SYCL"
-"SYCL0","XIELU","type=f32,ne=[512,16,1,1]","support","0","no","SYCL"
-"SYCL0","XIELU","type=f16,ne=[512,16,1,1]","support","0","no","SYCL"
+"SYCL0","XIELU","type=f32,ne=[10,5,4,3]","support","1","yes","SYCL"
+"SYCL0","XIELU","type=f16,ne=[10,5,4,3]","support","1","yes","SYCL"
+"SYCL0","XIELU","type=f32,ne=[512,16,1,1]","support","1","yes","SYCL"
+"SYCL0","XIELU","type=f16,ne=[512,16,1,1]","support","1","yes","SYCL"
 "SYCL0","TRI","type=f32,ne=[10,10,4,3],tri_type=3","support","1","yes","SYCL"
 "SYCL0","TRI","type=f32,ne=[10,10,4,3],tri_type=2","support","1","yes","SYCL"
 "SYCL0","TRI","type=f32,ne=[10,10,4,3],tri_type=1","support","1","yes","SYCL"

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

@@ -2859,7 +2859,8 @@ struct ggml_cplan ggml_graph_plan(
                     } break;
                 case GGML_OP_OUT_PROD:
                     {
-                        if (ggml_is_quantized(node->src[0]->type)) {
+                        if (ggml_is_quantized(node->src[0]->type) ||
+                            node->src[0]->type == GGML_TYPE_F16) {
                             cur = ggml_type_size(GGML_TYPE_F32) * node->src[0]->ne[0] * n_tasks;
                         }
                     } break;

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

@@ -462,11 +462,12 @@ static bool ggml_backend_cpu_device_supports_op(ggml_backend_dev_t dev, const st
             return max_bias == 0.0f;
         }
         case GGML_OP_IM2COL_BACK:
-            return src0->type == GGML_TYPE_F32 && src1->type == GGML_TYPE_F32;
+            return src0->type == GGML_TYPE_F32 && (src1->type == GGML_TYPE_F32 || src1->type == GGML_TYPE_F16);
         case GGML_OP_GET_ROWS_BACK:
             return src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_F16;
         case GGML_OP_OUT_PROD:
-            return (src0->type == GGML_TYPE_F32 || (ggml_is_quantized(src0->type) && src0->ne[2] == src1->ne[2] && src0->ne[3] == src1->ne[3])) &&
+            return (src0->type == GGML_TYPE_F32 ||
+                    ((src0->type == GGML_TYPE_F16 || ggml_is_quantized(src0->type)) && src0->ne[2] == src1->ne[2] && src0->ne[3] == src1->ne[3])) &&
                 src1->type == GGML_TYPE_F32 && op->type == GGML_TYPE_F32;
         default:
             return true;

ggml/src/ggml-cpu/ops.cpp

@@ -2081,8 +2081,8 @@ void ggml_compute_forward_concat(
     const ggml_tensor * src1 = dst->src[1];
 
     if (ggml_is_quantized(src0->type)) {
-        GGML_ASSERT(ggml_is_contiguous(src0));
-        GGML_ASSERT(ggml_is_contiguous(src1));
+        GGML_ASSERT(ggml_is_contiguous_rows(src0));
+        GGML_ASSERT(ggml_is_contiguous_rows(src1));
         GGML_ASSERT(src0->ne[0] % ggml_blck_size(src0->type) == 0);
         GGML_ASSERT(src1->ne[0] % ggml_blck_size(src1->type) == 0);
     }
@@ -4449,6 +4449,70 @@ static void ggml_compute_forward_out_prod_q_f32(
     }
 }
 
+static void ggml_compute_forward_out_prod_f16_f32(
+        const ggml_compute_params * params,
+              ggml_tensor * dst) {
+
+    const ggml_tensor * src0 = dst->src[0];
+    const ggml_tensor * src1 = dst->src[1];
+
+    GGML_TENSOR_BINARY_OP_LOCALS;
+
+    const int ith = params->ith;
+    const int nth = params->nth;
+
+    GGML_ASSERT(src0->type == GGML_TYPE_F16);
+    GGML_ASSERT(src1->type == GGML_TYPE_F32);
+    GGML_ASSERT(dst->type == GGML_TYPE_F32);
+
+    GGML_ASSERT(ne02 == ne12);
+    GGML_ASSERT(ne03 == ne13);
+    GGML_ASSERT(ne2  == ne12);
+    GGML_ASSERT(ne3  == ne13);
+
+    GGML_ASSERT(nb00 == sizeof(ggml_fp16_t));
+    GGML_ASSERT(nb0  == sizeof(float));
+
+    GGML_ASSERT(ne0 == ne00);
+    GGML_ASSERT(ne1 == ne10);
+    GGML_ASSERT(ne2 == ne02);
+    GGML_ASSERT(ne3 == ne03);
+
+    if (ith == 0) {
+        ggml_vec_set_f32(ne0*ne1*ne2*ne3, (float *)dst->data, 0);
+    }
+    ggml_barrier(params->threadpool);
+
+    const int64_t nr = ne1*ne2*ne3;
+    const int64_t dr = (nr + nth - 1)/nth;
+    const int64_t ir0 = dr*ith;
+    const int64_t ir1 = MIN(ir0 + dr, nr);
+
+    float * wdata = (float *) params->wdata + (ne0 + CACHE_LINE_SIZE_F32) * ith;
+
+    for (int64_t ir = ir0; ir < ir1; ++ir) {
+        const int64_t i3 = ir/(ne2*ne1);
+        const int64_t i2 = (ir - i3*ne2*ne1)/ne1;
+        const int64_t i1 = (ir - i3*ne2*ne1 - i2*ne1);
+
+        const int64_t i02 = i2;
+        const int64_t i03 = i3;
+
+        const int64_t i12 = i2;
+        const int64_t i13 = i3;
+
+        float * d = (float *) ((char *) dst->data + (i1*nb1 + i2*nb2 + i3*nb3));
+
+        for (int64_t i01 = 0; i01 < ne01; ++i01) {
+            const int64_t i11 = i01;
+            ggml_fp16_t * s0 = (ggml_fp16_t *) ((char *) src0->data + (i01*nb01 + i02*nb02 + i03*nb03));
+            float * s1 = (float *) ((char *) src1->data + (i1*nb10 + i11*nb11 + i12*nb12 + i13*nb13));
+            ggml_fp16_to_fp32_row(s0, wdata, ne0);
+            ggml_vec_mad_f32(ne0, d, wdata, *s1);
+        }
+    }
+}
+
 void ggml_compute_forward_out_prod(
         const ggml_compute_params * params,
         ggml_tensor * dst) {
@@ -4486,9 +4550,8 @@ void ggml_compute_forward_out_prod(
             } break;
         case GGML_TYPE_F16:
             {
-                GGML_ABORT("fatal error"); // todo
-                // ggml_compute_forward_out_prod_f16_f32(params, dst);
-            }
+                ggml_compute_forward_out_prod_f16_f32(params, dst);
+            } break;
         case GGML_TYPE_F32:
             {
                 ggml_compute_forward_out_prod_f32(params, dst);
@@ -6469,7 +6532,7 @@ void ggml_compute_forward_im2col_back_f32(
     const ggml_tensor * src1 = dst->src[1]; // convolution kernel
 
     GGML_ASSERT(src0->type == GGML_TYPE_F32);
-    GGML_ASSERT(src1->type == GGML_TYPE_F32);
+    GGML_ASSERT(src1->type == GGML_TYPE_F32 || src1->type == GGML_TYPE_F16);
     GGML_ASSERT( dst->type == GGML_TYPE_F32);
 
     GGML_TENSOR_BINARY_OP_LOCALS;

ggml/src/ggml-cuda/concat.cu

@@ -141,27 +141,25 @@ static __global__ void __launch_bounds__(CUDA_CONCAT_BLOCK_SIZE)
 
 template <typename T>
 static void concat_cuda(const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, int dim, cudaStream_t stream) {
-    if (ggml_is_contiguous(src0) && ggml_is_contiguous(src1)) {
+    if (dim != 3 && ggml_is_contiguous_to_3(src0) && ggml_is_contiguous_to_3(src1)) {
         const T * src0_d = (const T *) src0->data;
         const T * src1_d = (const T *) src1->data;
         T *       dst_d  = (T *) dst->data;
 
-        if (dim != 3) {
-            for (int64_t i3 = 0; i3 < dst->ne[3]; i3++) {
-                concat_cont_cuda(
-                        src0_d + i3*(src0->nb[3] / sizeof(T)),
-                        src1_d + i3*(src1->nb[3] / sizeof(T)),
-                        dst_d  + i3*( dst->nb[3] / sizeof(T)),
-                        ggml_row_size(src0->type, src0->ne[0])/sizeof(T), src0->ne[1], src0->ne[2],
-                        ggml_row_size(dst->type, dst->ne[0])/sizeof(T),  dst->ne[1],  dst->ne[2], dim, stream);
-            }
-        } else {
-            const size_t size0 = ggml_nbytes(src0);
-            const size_t size1 = ggml_nbytes(src1);
-
-            CUDA_CHECK(cudaMemcpyAsync((char *) dst->data,         src0->data, size0, cudaMemcpyDeviceToDevice, stream));
-            CUDA_CHECK(cudaMemcpyAsync((char *) dst->data + size0, src1->data, size1, cudaMemcpyDeviceToDevice, stream));
+        for (int64_t i3 = 0; i3 < dst->ne[3]; i3++) {
+            concat_cont_cuda(
+                    src0_d + i3*(src0->nb[3] / sizeof(T)),
+                    src1_d + i3*(src1->nb[3] / sizeof(T)),
+                    dst_d  + i3*( dst->nb[3] / sizeof(T)),
+                    ggml_row_size(src0->type, src0->ne[0])/sizeof(T), src0->ne[1], src0->ne[2],
+                    ggml_row_size(dst->type, dst->ne[0])/sizeof(T),  dst->ne[1],  dst->ne[2], dim, stream);
         }
+    } else if (dim == 3 && ggml_is_contiguous(src0) && ggml_is_contiguous(src1)) {
+        const size_t size0 = ggml_nbytes(src0);
+        const size_t size1 = ggml_nbytes(src1);
+
+        CUDA_CHECK(cudaMemcpyAsync((char *) dst->data,         src0->data, size0, cudaMemcpyDeviceToDevice, stream));
+        CUDA_CHECK(cudaMemcpyAsync((char *) dst->data + size0, src1->data, size1, cudaMemcpyDeviceToDevice, stream));
     } else {
         GGML_ASSERT(!ggml_is_quantized(src0->type));
 
@@ -208,12 +206,17 @@ void ggml_cuda_op_concat(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
     GGML_ASSERT(dst->type  == src0->type);
 
     if (ggml_is_quantized(src0->type)) {
-        GGML_ASSERT(ggml_is_contiguous(src0));
-        GGML_ASSERT(ggml_is_contiguous(src1));
+        if (dim == 3) {
+            GGML_ASSERT(ggml_is_contiguous(src0));
+            GGML_ASSERT(ggml_is_contiguous(src1));
+        } else {
+            GGML_ASSERT(ggml_is_contiguous_to_3(src0));
+            GGML_ASSERT(ggml_is_contiguous_to_3(src1));
+        }
         GGML_ASSERT(src0->ne[0] % ggml_blck_size(src0->type) == 0);
         GGML_ASSERT(src1->ne[0] % ggml_blck_size(src1->type) == 0);
 
-        // if tensors are contiguous and ne[0] is multiple of the block size we can concat both tensors as byte tensors
+        // if first 3 dimensions are contiguous and ne[0] is multiple of the block size we can concat both tensors as byte tensors
         concat_cuda<uint8_t>(src0, src1, dst, dim, stream);
     } else {
         GGML_ASSERT(ggml_blck_size(src0->type) == 1);

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

@@ -65,6 +65,7 @@
 #include "ggml-cuda/tri.cuh"
 #include "ggml-cuda/cumsum.cuh"
 #include "ggml-cuda/fill.cuh"
+#include "ggml-cuda/lightning-indexer.cuh"
 #include "ggml.h"
 
 #include <algorithm>
@@ -2257,6 +2258,9 @@ static bool ggml_cuda_compute_forward(ggml_backend_cuda_context & ctx, struct gg
         case GGML_OP_FILL:
             ggml_cuda_op_fill(ctx, dst);
             break;
+        case GGML_OP_LIGHTNING_INDEXER:
+            ggml_cuda_lightning_indexer(ctx, dst);
+            break;
         default:
             return false;
     }
@@ -4816,13 +4820,23 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g
             {
                 ggml_type src0_type = op->src[0]->type;
                 ggml_type src1_type = op->src[1]->type;
+                const int32_t dim = op->op_params[0];
                 return src0_type == src1_type &&
                        src0_type == op->type &&
                        (
                            (
                                ggml_is_quantized(src0_type) &&
-                               ggml_is_contiguous(op->src[0]) &&
-                               ggml_is_contiguous(op->src[1]) &&
+                               (
+                                   (
+                                       dim == 3 &&
+                                       ggml_is_contiguous(op->src[0]) &&
+                                       ggml_is_contiguous(op->src[1])
+                                   ) || (
+                                       dim != 3 &&
+                                       ggml_is_contiguous_to_3(op->src[0]) &&
+                                       ggml_is_contiguous_to_3(op->src[1])
+                                   )
+                               ) &&
                                op->src[0]->ne[0] % ggml_blck_size(src0_type) == 0 &&
                                op->src[1]->ne[0] % ggml_blck_size(src0_type) == 0
                            ) || (
@@ -4977,6 +4991,8 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g
         case GGML_OP_DIAG:
         case GGML_OP_SOLVE_TRI:
             return true;
+        case GGML_OP_LIGHTNING_INDEXER:
+            return ggml_cuda_lightning_indexer_supported(dev_ctx->device, op);
 
         default:
             return false;

ggml/src/ggml-cuda/lightning-indexer.cu

@@ -0,0 +1,588 @@
+#include "common.cuh"
+#include "lightning-indexer.cuh"
+#include "fattn-common.cuh"
+#include "convert.cuh"
+
+#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
+#if defined(TURING_MMA_AVAILABLE)
+
+typedef union {
+    int2 i2;
+    half2 h2[2];
+} half4;
+
+// TODO add support for AMD cards via rocWMMA
+#include <mma.h>
+namespace wmma = nvcuda::wmma;
+
+template <int WARPS_PER_BLOCK, int K_VECS_PER_BLOCK, int64_t N_EMBD, int64_t N_HEAD, ggml_type TYPE_K>
+static __global__ void lightning_indexer_kernel_wmma(
+        const float * Q, const char * K, const float * W, const half * M, float * dst,
+        int64_t n_stream, int64_t n_batch, int64_t n_kv,
+        size_t nb1, size_t nb2, size_t nb3,
+        size_t nbq1, size_t nbq2, size_t nbq3,
+        size_t nbk1, size_t nbk2, size_t nbk3,
+        size_t nbw1, size_t nbw2, size_t nbw3,
+        size_t nbm1, size_t nbm2, size_t nbm3,
+        int64_t nem3
+    ) {
+
+    constexpr int THREADS_PER_BLOCK = WARPS_PER_BLOCK * WARP_SIZE;
+    constexpr int HEADS_PER_INNER_LOOP = 8;
+    constexpr int K_EMBD_PER_INNER_LOOP = 16;
+    constexpr int N_EMBD_PADDED = N_EMBD + 8;
+
+    const int i_batch  = blockIdx.y;
+    const int i_stream = blockIdx.z;
+    const int i_warp   = threadIdx.y;
+    const int i_lane   = threadIdx.x;
+    const int tid      = i_warp * WARP_SIZE + i_lane;
+
+    // each block processes K_VECS_PER_BLOCK K vectors
+    const int start_kv = blockIdx.x * K_VECS_PER_BLOCK;
+
+    const char  * q_base = (const char  *)                 Q + i_batch*nbq2 + i_stream*nbq3;
+    const float * w_base = (const float *) ((const char *) W + i_batch*nbw1 + i_stream*nbw3);
+
+    // phase 1 - load weights and first Q tile to shared memory
+
+    __shared__ float w_shared[N_HEAD];
+    __shared__ int2  q_shared_h[HEADS_PER_INNER_LOOP][N_EMBD_PADDED / 4];
+
+    if (tid < N_HEAD) {
+        w_shared[tid] = w_base[tid];
+    }
+
+    // total number of half4 elements in HEADS_PER_INNER_LOOP x N_EMBD Q tile
+    constexpr int N_Q_TILE = HEADS_PER_INNER_LOOP * (N_EMBD / 4);
+    // number of registers needed in each thread to store Q tile in thread block
+    constexpr int N_Q_NEXT = (N_Q_TILE + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK;
+
+#pragma unroll
+    for (int i_q = tid; i_q < N_Q_TILE; i_q += THREADS_PER_BLOCK) {
+        const int i_head = i_q / (N_EMBD / 4);
+        const int i_embd = i_q % (N_EMBD / 4);
+        const float4 q = *(const float4 *) (q_base + i_head*nbq1 + i_embd*sizeof(float4));
+        half4 q_packed;
+        q_packed.h2[0] = __float22half2_rn(make_float2(q.x, q.y));
+        q_packed.h2[1] = __float22half2_rn(make_float2(q.z, q.w));
+        q_shared_h[i_head][i_embd] = q_packed.i2;
+    }
+
+    // phase 2 - load (and dequantize if needed) K to shared mem
+
+    __shared__ half2 k_shared_h[K_VECS_PER_BLOCK][N_EMBD_PADDED / 4][2];
+
+    constexpr int n_k = K_VECS_PER_BLOCK * (N_EMBD / 4);
+
+    if constexpr (TYPE_K == GGML_TYPE_F16) {
+#pragma unroll
+        for (int i_k = tid; i_k < n_k; i_k += THREADS_PER_BLOCK) {
+            const int i_k_vec = i_k / (N_EMBD / 4);
+            const int i_embd = i_k % (N_EMBD / 4);
+            const int i_kv = start_kv + i_k_vec;
+            if (i_kv < n_kv) {
+                const int2 * k_base = (const int2 *) ((const char *) K + i_kv*nbk2 + i_stream*nbk3);
+                *(int2*) &k_shared_h[i_k_vec][i_embd] = k_base[i_embd];
+            } else {
+                *(int2*) &k_shared_h[i_k_vec][i_embd] = make_int2(0, 0);
+            }
+        }
+    } else {
+        constexpr dequantize_V_t dequantize_k = get_dequantize_V<TYPE_K, half, 4>();
+#pragma unroll
+        for (int i_k = tid; i_k < n_k; i_k += THREADS_PER_BLOCK) {
+            const int i_k_vec = i_k / (N_EMBD / 4);
+            const int i_embd = i_k % (N_EMBD / 4);
+            const int i_kv = start_kv + i_k_vec;
+            if (i_kv < n_kv) {
+                const void * k_base = (const void *) ((const char *) K + i_kv*nbk2 + i_stream*nbk3);
+                dequantize_k(k_base, &k_shared_h[i_k_vec][i_embd][0], i_embd * 4);
+            } else {
+                *(int2*) &k_shared_h[i_k_vec][i_embd] = make_int2(0, 0);
+            }
+        }
+    }
+
+    __syncthreads();
+
+    // phase 3 - calculate lightning indexer scores
+
+    __shared__ float qk_shared[WARPS_PER_BLOCK][HEADS_PER_INNER_LOOP][K_VECS_PER_BLOCK];
+
+    // load K fragment
+    wmma::fragment<wmma::matrix_b, HEADS_PER_INNER_LOOP, K_VECS_PER_BLOCK, K_EMBD_PER_INNER_LOOP, half, wmma::col_major> frag_k;
+    wmma::load_matrix_sync(frag_k, (half*) &k_shared_h[0][i_warp * K_EMBD_PER_INNER_LOOP / 4], N_EMBD_PADDED);
+
+    float score_k = 0.0f;
+
+    for (int i_head_0 = 0; i_head_0 < N_HEAD; i_head_0 += HEADS_PER_INNER_LOOP) {
+        const int i_head_next = i_head_0 + HEADS_PER_INNER_LOOP;
+
+        // we don't use accumulator for anything, fill it with zeros
+        wmma::fragment<wmma::accumulator, HEADS_PER_INNER_LOOP, K_VECS_PER_BLOCK, K_EMBD_PER_INNER_LOOP, float> frag_acc;
+        wmma::fill_fragment(frag_acc, 0.0f);
+
+        // load Q fragment
+        wmma::fragment<wmma::matrix_a, HEADS_PER_INNER_LOOP, K_VECS_PER_BLOCK, K_EMBD_PER_INNER_LOOP, half, wmma::row_major> frag_q;
+        wmma::load_matrix_sync(frag_q, (half*) &q_shared_h[0][i_warp * K_EMBD_PER_INNER_LOOP / 4], N_EMBD_PADDED);
+
+        // preload next Q tile to registers during matrix multiplication
+        float4 q_next[N_Q_NEXT];
+
+        if (i_head_next < N_HEAD) {
+#pragma unroll
+            for (int i_q = tid, i_q_next = 0; i_q < N_Q_TILE; i_q += THREADS_PER_BLOCK) {
+                const int i_head = i_head_next + i_q / (N_EMBD / 4);
+                const int i_embd =               i_q % (N_EMBD / 4);
+                q_next[i_q_next++] = *(const float4 *) (q_base + i_head*nbq1 + i_embd*sizeof(float4));
+            }
+        }
+
+        // perform matrix multiplication
+        wmma::mma_sync(frag_acc, frag_q, frag_k, frag_acc);
+        wmma::store_matrix_sync((float*) &qk_shared[i_warp][0][0], frag_acc, K_VECS_PER_BLOCK, wmma::mem_row_major);
+
+        // make sure all threads finished using q_shared_h so we can store next tile
+        __syncthreads();
+
+        // write preloaded Q tile to shared memory
+        if (i_head_next < N_HEAD) {
+#pragma unroll
+            for (int i_q = tid, i_q_next = 0; i_q < N_Q_TILE; i_q += THREADS_PER_BLOCK) {
+                const int i_head = i_q / (N_EMBD / 4);
+                const int i_embd = i_q % (N_EMBD / 4);
+                half4 q_packed;
+                q_packed.h2[0] = __float22half2_rn(make_float2(q_next[i_q_next].x, q_next[i_q_next].y));
+                q_packed.h2[1] = __float22half2_rn(make_float2(q_next[i_q_next].z, q_next[i_q_next].w));
+                q_shared_h[i_head][i_embd] = q_packed.i2;
+                ++i_q_next;
+            }
+        }
+
+        // accumulate QK multiplication results from all block warps
+        // (there are 256 threads in block and 256 matmul outputs)
+        // TODO it will break if WARP_SIZE is not 32
+        const int h = tid / K_VECS_PER_BLOCK;
+        const int k = tid % K_VECS_PER_BLOCK;
+        const float w_val = w_shared[i_head_0 + h];
+
+        float sum = 0.0f;
+#pragma unroll
+        for (int w = 0; w < WARPS_PER_BLOCK; ++w) {
+            sum += qk_shared[w][h][k];
+        }
+
+        // ReLU, weight
+        sum = sum > 0.0f ? sum : 0.0f;
+        sum *= w_val;
+
+        // wait until qk_shared[0] is no longer used
+        __syncthreads();
+
+        // reuse qk_shared[0] for storing partial results
+        qk_shared[0][h][k] = sum;
+
+        // wait until all threads write their results
+        __syncthreads();
+
+        // accumulate result over heads
+        if (tid < K_VECS_PER_BLOCK) {
+#pragma unroll
+            for (int i_head = 0; i_head < HEADS_PER_INNER_LOOP; ++i_head) {
+                score_k += qk_shared[0][i_head][tid];
+            }
+        }
+
+        // make sure all threads finished using qk_shared
+        __syncthreads();
+    }
+
+    // phase 4 - store output to VRAM
+
+    if (tid < K_VECS_PER_BLOCK) {
+        const int i_kv = start_kv + tid;
+        if (i_kv < n_kv) {
+            const half * m_base = (const half *) ((const char *) M + i_batch*nbm1 + (i_stream%nem3)*nbm3);
+            float * dst_base = (float *) ((char *) dst + i_batch*nb1 + i_stream*nb3);
+            dst_base[i_kv] = score_k + __half2float(m_base[i_kv]);
+        }
+    }
+}
+
+#else // defined(TURING_MMA_AVAILABLE)
+
+template <int WARPS_PER_BLOCK, int K_VECS_PER_BLOCK, int64_t N_EMBD, int64_t N_HEAD, ggml_type TYPE_K>
+static __global__ void lightning_indexer_kernel_wmma(
+        const float * Q, const char * K, const float * W, const half * M, float * dst,
+        int64_t n_stream, int64_t n_batch, int64_t n_kv,
+        size_t nb1, size_t nb2, size_t nb3,
+        size_t nbq1, size_t nbq2, size_t nbq3,
+        size_t nbk1, size_t nbk2, size_t nbk3,
+        size_t nbw1, size_t nbw2, size_t nbw3,
+        size_t nbm1, size_t nbm2, size_t nbm3,
+        int64_t nem3
+    ) {
+    GGML_UNUSED_VARS(Q, K, W, M, dst,
+        n_stream, n_batch, n_kv,
+        nb1, nb2, nb3,
+        nbq1, nbq2, nbq3,
+        nbk1, nbk2, nbk3,
+        nbw1, nbw2, nbw3,
+        nem3);
+    NO_DEVICE_CODE;
+}
+
+#endif // defined(TURING_MMA_AVAILABLE)
+#endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
+
+// TODO there is one ugly assumption used in this kernel - that WARP_SIZE is equal to 32
+// thanks to that one warp operating on float4 processes whole indexer K/Q vectors
+// 32 * 4 = 128 (N_EMBD)
+
+template <int WARPS_PER_BLOCK, int K_VECS_PER_BLOCK, int64_t N_EMBD, int64_t N_HEAD, ggml_type TYPE_K>
+static __global__ void lightning_indexer_kernel_vec(
+        const float * Q, const char * K, const float * W, const half * M, float * dst,
+        int64_t n_stream, int64_t n_batch, int64_t n_kv,
+        size_t nb1, size_t nb2, size_t nb3,
+        size_t nbq1, size_t nbq2, size_t nbq3,
+        size_t nbk1, size_t nbk2, size_t nbk3,
+        size_t nbw1, size_t nbw2, size_t nbw3,
+        size_t nbm1, size_t nbm2, size_t nbm3,
+        int64_t nem3
+    ) {
+
+    constexpr int K_VECS_PER_WARP = K_VECS_PER_BLOCK / WARPS_PER_BLOCK;
+    constexpr int THREADS_PER_BLOCK = WARPS_PER_BLOCK * WARP_SIZE;
+
+    const int i_batch  = blockIdx.y;
+    const int i_stream = blockIdx.z;
+    const int i_warp   = threadIdx.y;
+    const int i_lane   = threadIdx.x;
+    const int tid      = i_warp * WARP_SIZE + i_lane;
+
+    // each warp processes K_VECS_PER_WARP K vectors
+    const int start_kv_block = blockIdx.x * K_VECS_PER_BLOCK;
+    const int start_kv = start_kv_block + i_warp * K_VECS_PER_WARP;
+
+    const char  * q_base = (const char  *)                 Q + i_batch*nbq2 + i_stream*nbq3;
+    const float * w_base = (const float *) ((const char *) W + i_batch*nbw1 + i_stream*nbw3);
+
+    // phase 1 - load (and dequantize if needed) K to registers
+
+    float4 k_reg_f[K_VECS_PER_WARP];
+
+    if constexpr (TYPE_K == GGML_TYPE_F32) {
+        // direct copy of float4
+#pragma unroll
+        for (int k = 0; k < K_VECS_PER_WARP; ++k) {
+            int i_kv = start_kv + k;
+            if (i_kv < n_kv) {
+                const float4 * k_base = (const float4 *) ((const char *) K + i_kv*nbk2 + i_stream*nbk3);
+                k_reg_f[k] = k_base[i_lane];
+            } else {
+                k_reg_f[k] = make_float4(0, 0, 0, 0);
+            }
+        }
+    } else {
+        // dequantize remaining types to float
+        constexpr dequantize_V_t dequantize_k = get_dequantize_V<TYPE_K, float, 4>();
+#pragma unroll
+        for (int k = 0; k < K_VECS_PER_WARP; ++k) {
+            int i_kv = start_kv + k;
+            if (i_kv < n_kv) {
+                const void * k_base = (const void *) ((const char *) K + i_kv*nbk2 + i_stream*nbk3);
+                dequantize_k(k_base, &k_reg_f[k], i_lane * 4);
+            } else {
+                k_reg_f[k] = make_float4(0, 0, 0, 0);
+            }
+        }
+    }
+
+    float score_k[K_VECS_PER_WARP] = { 0.0f };
+
+    // load weights and Q only for N_HEAD_INNER heads at once to reduce shared memory usage
+    constexpr int N_HEAD_INNER = N_HEAD / 4;
+
+    for (int i_head_0 = 0; i_head_0 < N_HEAD; i_head_0 += N_HEAD_INNER) {
+        // phase 2 - load weights and Q to shared memory
+
+        __shared__ float  w_shared[N_HEAD_INNER];
+        __shared__ float4 q_shared_f[N_HEAD_INNER][N_EMBD / 4];
+
+        if (tid < N_HEAD_INNER) {
+            w_shared[tid] = w_base[i_head_0 + tid];
+        }
+
+        constexpr int n_q = N_HEAD_INNER * (N_EMBD / 4);
+#pragma unroll
+        for (int i_q = tid; i_q < n_q; i_q += THREADS_PER_BLOCK) {
+            const int i_head_inner = i_q / (N_EMBD / 4);
+            const int i_head = i_head_0 + i_head_inner;
+            const int i_embd = i_q % (N_EMBD / 4);
+            q_shared_f[i_head_inner][i_embd] = *(const float4 *) (q_base + i_head*nbq1 + i_embd*sizeof(float4));
+        }
+
+        __syncthreads();
+
+        // phase 3 - calculate lightning indexer scores
+
+        for (int i_head_inner = 0; i_head_inner < N_HEAD_INNER; ++i_head_inner) {
+            const float w_val = w_shared[i_head_inner];
+            float qk[K_VECS_PER_WARP] = { 0.0f };
+
+            // dot product of floats
+            const float4 q_vec = q_shared_f[i_head_inner][i_lane];
+
+#pragma unroll
+            for (int k = 0; k < K_VECS_PER_WARP; ++k) {
+                ggml_cuda_mad(qk[k], q_vec.x, k_reg_f[k].x);
+                ggml_cuda_mad(qk[k], q_vec.y, k_reg_f[k].y);
+                ggml_cuda_mad(qk[k], q_vec.z, k_reg_f[k].z);
+                ggml_cuda_mad(qk[k], q_vec.w, k_reg_f[k].w);
+            }
+
+#pragma unroll
+            for (int k = 0; k < K_VECS_PER_WARP; ++k) {
+                float sum = warp_reduce_sum(qk[k]);
+
+                // ReLU, weight
+                if (i_lane == 0) {
+                    sum = (sum > 0.0f) ? sum : 0.0f;
+                    score_k[k] += sum * w_val;
+                }
+            }
+        }
+
+        __syncthreads();
+    }
+
+    // phase 4 - store outputs to shared memory
+
+    __shared__ float dst_shared[K_VECS_PER_BLOCK];
+
+    if (i_lane == 0) {
+#pragma unroll
+        for (int k = 0; k < K_VECS_PER_WARP; ++k) {
+            dst_shared[i_warp * K_VECS_PER_WARP + k] = score_k[k];
+        }
+    }
+
+    __syncthreads();
+
+    // phase 5 - write from shared memory to VRAM in coalesced manner
+
+    if (tid < K_VECS_PER_BLOCK) {
+        int i_kv = start_kv_block + tid;
+        if (i_kv < n_kv) {
+            const half * m_base = (const half *) ((const char *) M + i_batch*nbm1 + (i_stream%nem3)*nbm3);
+            float * dst_base = (float *) ((char *) dst + i_batch*nb1 + i_stream*nb3);
+            dst_base[i_kv] = dst_shared[tid] + __half2float(m_base[i_kv]);
+        }
+    }
+}
+
+#define LIGHTNING_INDEXER_CASE(lightning_indexer_kernel, n_embd, n_head, K, type_K)         \
+    if (K->type == (type_K)) {                                                              \
+        lightning_indexer_kernel<WARPS_PER_BLOCK, K_VECS_PER_BLOCK, n_embd, n_head, type_K> \
+            <<<grid, block, 0, ctx.stream()>>>(                                             \
+            q_d, k_d, w_d, m_d, dst_d,                                                      \
+            n_stream, n_batch, n_kv,                                                        \
+            nb1, nb2, nb3,                                                                  \
+            nbq1, nbq2, nbq3,                                                               \
+            nbk1, nbk2, nbk3,                                                               \
+            nbw1, nbw2, nbw3,                                                               \
+            nbm1, nbm2, nbm3,                                                               \
+            nem3                                                                            \
+        );                                                                                  \
+    } else
+
+void ggml_cuda_lightning_indexer(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
+    const ggml_tensor * q = dst->src[0];
+    const ggml_tensor * k = dst->src[1];
+    const ggml_tensor * w = dst->src[2]; // weights
+    const ggml_tensor * m = dst->src[3]; // mask
+
+    GGML_ASSERT(dst->type == GGML_TYPE_F32);
+    GGML_ASSERT(  q->type == GGML_TYPE_F32);
+    GGML_ASSERT(  w->type == GGML_TYPE_F32);
+    GGML_ASSERT(  m->type == GGML_TYPE_F16);
+
+    GGML_TENSOR_LOCALS(int64_t, neq,  q, ne)
+    GGML_TENSOR_LOCALS(size_t,  nbq,  q, nb)
+    GGML_TENSOR_LOCALS(int64_t, nek,  k, ne)
+    GGML_TENSOR_LOCALS(size_t,  nbk,  k, nb)
+    GGML_TENSOR_LOCALS(int64_t, new,  w, ne)
+    GGML_TENSOR_LOCALS(size_t,  nbw,  w, nb)
+    GGML_TENSOR_LOCALS(int64_t, nem,  m, ne)
+    GGML_TENSOR_LOCALS(size_t,  nbm,  m, nb)
+    GGML_TENSOR_LOCALS(int64_t, ne, dst, ne)
+    GGML_TENSOR_LOCALS(size_t,  nb, dst, nb)
+
+    // input tensor rows must be contiguous
+    GGML_ASSERT(nbq0 == ggml_type_size(q->type));
+    GGML_ASSERT(nbk0 == ggml_type_size(k->type));
+    GGML_ASSERT(nbw0 == ggml_type_size(w->type));
+    GGML_ASSERT(nbm0 == ggml_type_size(m->type));
+
+    // dst cannot be transposed or permuted
+    GGML_ASSERT(nb0 == sizeof(float));
+    GGML_ASSERT(nb0 <= nb1);
+    GGML_ASSERT(nb1 <= nb2);
+    GGML_ASSERT(nb2 <= nb3);
+
+    const int n_embd   = q->ne[0];
+    const int n_head   = q->ne[1];
+    const int n_batch  = q->ne[2];
+    const int n_stream = q->ne[3];
+    const int n_kv     = k->ne[2];
+
+    const float *   q_d = (const float *)   q->data;
+    const char  *   k_d = (const char  *)   k->data;
+    const float *   w_d = (const float *)   w->data;
+    const half  *   m_d = (const half  *)   m->data;
+    float       * dst_d = (      float *) dst->data;
+
+    const int device = ggml_cuda_get_device();
+    const int cc     = ggml_cuda_info().devices[device].cc;
+
+    if (n_embd == 128 && n_head == 64) {
+#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
+        if (GGML_CUDA_CC_IS_NVIDIA(cc) && turing_mma_available(cc) && k->type != GGML_TYPE_F32 && k->type != GGML_TYPE_BF16) {
+            // use wmma kernel
+            constexpr int K_VECS_PER_BLOCK = 32;
+            constexpr int WARPS_PER_BLOCK = 8;
+
+            dim3 block(32, WARPS_PER_BLOCK);
+            int num_kv_blocks = (n_kv + (K_VECS_PER_BLOCK) - 1) / (K_VECS_PER_BLOCK);
+            dim3 grid(num_kv_blocks, n_batch, n_stream);
+
+            LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_wmma, 128, 64, k, GGML_TYPE_F16)
+            LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_wmma, 128, 64, k, GGML_TYPE_Q4_0)
+            LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_wmma, 128, 64, k, GGML_TYPE_Q4_1)
+            LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_wmma, 128, 64, k, GGML_TYPE_Q5_0)
+            LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_wmma, 128, 64, k, GGML_TYPE_Q5_1)
+            LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_wmma, 128, 64, k, GGML_TYPE_Q8_0)
+            GGML_ABORT("fatal error");
+        } else {
+#else // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
+        {
+#endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
+            // use vector kernel
+            constexpr int K_VECS_PER_WARP = 8;
+            constexpr int WARPS_PER_BLOCK = 8;
+            constexpr int K_VECS_PER_BLOCK = K_VECS_PER_WARP * WARPS_PER_BLOCK;
+
+            dim3 block(32, WARPS_PER_BLOCK);
+            int num_kv_blocks = (n_kv + (K_VECS_PER_BLOCK) - 1) / (K_VECS_PER_BLOCK);
+            dim3 grid(num_kv_blocks, n_batch, n_stream);
+
+            LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_vec, 128, 64, k, GGML_TYPE_F16)
+            LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_vec, 128, 64, k, GGML_TYPE_Q4_0)
+            LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_vec, 128, 64, k, GGML_TYPE_Q4_1)
+            LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_vec, 128, 64, k, GGML_TYPE_Q5_0)
+            LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_vec, 128, 64, k, GGML_TYPE_Q5_1)
+            LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_vec, 128, 64, k, GGML_TYPE_Q8_0)
+            LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_vec, 128, 64, k, GGML_TYPE_BF16)
+            LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_vec, 128, 64, k, GGML_TYPE_F32)
+            GGML_ABORT("fatal error");
+        }
+    } else if (n_embd == 128 && n_head == 32) {
+#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
+        if (GGML_CUDA_CC_IS_NVIDIA(cc) && turing_mma_available(cc) && k->type != GGML_TYPE_F32 && k->type != GGML_TYPE_BF16) {
+            // use wmma kernel
+            constexpr int K_VECS_PER_BLOCK = 32;
+            constexpr int WARPS_PER_BLOCK = 8;
+
+            dim3 block(32, WARPS_PER_BLOCK);
+            int num_kv_blocks = (n_kv + (K_VECS_PER_BLOCK) - 1) / (K_VECS_PER_BLOCK);
+            dim3 grid(num_kv_blocks, n_batch, n_stream);
+
+            LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_wmma, 128, 32, k, GGML_TYPE_F16)
+            LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_wmma, 128, 32, k, GGML_TYPE_Q4_0)
+            LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_wmma, 128, 32, k, GGML_TYPE_Q4_1)
+            LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_wmma, 128, 32, k, GGML_TYPE_Q5_0)
+            LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_wmma, 128, 32, k, GGML_TYPE_Q5_1)
+            LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_wmma, 128, 32, k, GGML_TYPE_Q8_0)
+            GGML_ABORT("fatal error");
+        } else {
+#else // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
+        {
+#endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
+            // use vector kernel
+            constexpr int K_VECS_PER_WARP = 8;
+            constexpr int WARPS_PER_BLOCK = 8;
+            constexpr int K_VECS_PER_BLOCK = K_VECS_PER_WARP * WARPS_PER_BLOCK;
+
+            dim3 block(32, WARPS_PER_BLOCK);
+            int num_kv_blocks = (n_kv + (K_VECS_PER_BLOCK) - 1) / (K_VECS_PER_BLOCK);
+            dim3 grid(num_kv_blocks, n_batch, n_stream);
+
+            LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_vec, 128, 32, k, GGML_TYPE_F16)
+            LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_vec, 128, 32, k, GGML_TYPE_Q4_0)
+            LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_vec, 128, 32, k, GGML_TYPE_Q4_1)
+            LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_vec, 128, 32, k, GGML_TYPE_Q5_0)
+            LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_vec, 128, 32, k, GGML_TYPE_Q5_1)
+            LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_vec, 128, 32, k, GGML_TYPE_Q8_0)
+            LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_vec, 128, 32, k, GGML_TYPE_BF16)
+            LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_vec, 128, 32, k, GGML_TYPE_F32)
+            GGML_ABORT("fatal error");
+        }
+    } else {
+        GGML_ABORT("fatal error");
+    }
+}
+
+bool ggml_cuda_lightning_indexer_supported(int device, const ggml_tensor * dst) {
+    GGML_UNUSED(device);
+
+    const ggml_tensor * q = dst->src[0];
+    const ggml_tensor * k = dst->src[1];
+    const ggml_tensor * w = dst->src[2]; // weights
+    const ggml_tensor * m = dst->src[3]; // mask
+
+    GGML_TENSOR_LOCALS(int64_t, neq,  q, ne)
+    GGML_TENSOR_LOCALS(size_t,  nbq,  q, nb)
+    GGML_TENSOR_LOCALS(int64_t, nek,  k, ne)
+    GGML_TENSOR_LOCALS(size_t,  nbk,  k, nb)
+    GGML_TENSOR_LOCALS(int64_t, new,  w, ne)
+    GGML_TENSOR_LOCALS(size_t,  nbw,  w, nb)
+    GGML_TENSOR_LOCALS(int64_t, nem,  m, ne)
+    GGML_TENSOR_LOCALS(size_t,  nbm,  m, nb)
+    GGML_TENSOR_LOCALS(int64_t, ne, dst, ne)
+    GGML_TENSOR_LOCALS(size_t,  nb, dst, nb)
+
+    if (neq0 != 128) {
+        return false;
+    }
+
+    if (neq1 != 64 && neq1 != 32) {
+        return false;
+    }
+
+    // alignment checks
+    for (const ggml_tensor * t : {q, k}) {
+        if (ggml_is_quantized(t->type)) {
+            continue;
+        }
+        for (size_t i = 1; i < GGML_MAX_DIMS; ++i) {
+            if (t->nb[i] % 16 != 0) {
+                return false;
+            }
+        }
+    }
+
+    switch(k->type) {
+        case GGML_TYPE_F32:
+        case GGML_TYPE_BF16:
+        case GGML_TYPE_F16:
+        case GGML_TYPE_Q8_0:
+        case GGML_TYPE_Q5_1:
+        case GGML_TYPE_Q5_0:
+        case GGML_TYPE_Q4_1:
+        case GGML_TYPE_Q4_0:
+            return true;
+        default:
+            return false;
+    }
+}

ggml/src/ggml-cuda/lightning-indexer.cuh

@@ -0,0 +1,4 @@
+#include "common.cuh"
+
+void ggml_cuda_lightning_indexer(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
+bool ggml_cuda_lightning_indexer_supported(int device, const ggml_tensor * dst);

ggml/src/ggml-cuda/mmq-load-tiles.cuh

@@ -39,29 +39,37 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
         }
 
         const block_q1_0 * bxi = (const block_q1_0 *) x + kbx0 + i*stride + kbx;
-        const int qs_offset = 4*kqsx;
-        const int qs0 = bxi->qs[qs_offset + 0] | (bxi->qs[qs_offset + 1] << 8) |
-                        (bxi->qs[qs_offset + 2] << 16) | (bxi->qs[qs_offset + 3] << 24);
-
-        int unpacked_bytes[8];
-#pragma unroll
-        for (int j = 0; j < 8; ++j) {
-            const int shift = j * 4;
-            const int bits4 = (qs0 >> shift) & 0x0F;
-            const int b0 = (bits4 & 0x01) ? 1 : -1;
-            const int b1 = (bits4 & 0x02) ? 1 : -1;
-            const int b2 = (bits4 & 0x04) ? 1 : -1;
-            const int b3 = (bits4 & 0x08) ? 1 : -1;
-            unpacked_bytes[j] = (b0 & 0xFF) | ((b1 & 0xFF) << 8) | ((b2 & 0xFF) << 16) | ((b3 & 0xFF) << 24);
-        }
+        const int16_t    * qxi = (const int16_t *) bxi->qs + kqsx * 2;
 
         const int dst_offset = kbx*(scale_entries_per_block*QI8_0) + kqsx*QI8_0;
 #pragma unroll
-        for (int j = 0; j < 8; ++j) {
+        for (int j = 0; j < 2; ++j) {
+            const int q  = qxi[j];
+
+            // unpack crumbs into nibble indices
+            const int n0 = __byte_perm(0x11100100, 0x11100100, q >> 0); // [0, 1, 4, 5] [ 8,  9, 12, 13]
+            const int n1 = __byte_perm(0x11100100, 0x11100100, q >> 2); // [2, 3, 6, 7] [10, 11, 14, 15]
+            // unpack nibbles into byte values
+            const int s0 = __byte_perm(0x01FF, 0x01FF, n0 >>  0);
+            const int s1 = __byte_perm(0x01FF, 0x01FF, n1 >>  0);
+            const int s2 = __byte_perm(0x01FF, 0x01FF, n0 >> 16);
+            const int s3 = __byte_perm(0x01FF, 0x01FF, n1 >> 16);
+            // unshuffle values
+            const int v0 = __byte_perm(s0, s1, 0x5410);
+            const int v1 = __byte_perm(s0, s1, 0x7632);
+            const int v2 = __byte_perm(s2, s3, 0x5410);
+            const int v3 = __byte_perm(s2, s3, 0x7632);
+
 #if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
-            x_qs[i*sram_stride           + dst_offset + j] = unpacked_bytes[j];
+            x_qs[i*sram_stride           + dst_offset + j*4+0] = v0;
+            x_qs[i*sram_stride           + dst_offset + j*4+1] = v1;
+            x_qs[i*sram_stride           + dst_offset + j*4+2] = v2;
+            x_qs[i*sram_stride           + dst_offset + j*4+3] = v3;
 #else
-            x_qs[i*(2*MMQ_TILE_NE_K + 1) + dst_offset + j] = unpacked_bytes[j];
+            x_qs[i*(2*MMQ_TILE_NE_K + 1) + dst_offset + j*4+0] = v0;
+            x_qs[i*(2*MMQ_TILE_NE_K + 1) + dst_offset + j*4+1] = v1;
+            x_qs[i*(2*MMQ_TILE_NE_K + 1) + dst_offset + j*4+2] = v2;
+            x_qs[i*(2*MMQ_TILE_NE_K + 1) + dst_offset + j*4+3] = v3;
 #endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
         }
     }

ggml/src/ggml-cuda/mmq.cu

@@ -122,11 +122,12 @@ void ggml_cuda_mul_mat_q(
 
     const bool fallback = ne01 % 128 != 0;
 
-    // TODO: tighter pool buffer size vs q8 path
     const bool use_native_fp4 = blackwell_mma_available(cc) && (src0->type == GGML_TYPE_MXFP4 || src0->type == GGML_TYPE_NVFP4);
+    const size_t y_block_size       = use_native_fp4 ? sizeof(block_fp4_mmq) : sizeof(block_q8_1_mmq);
+    const size_t y_values_per_block = use_native_fp4 ? QK_FP4_MMQ            : QK8_1_MMQ;
 
     if (!ids) {
-        const size_t nbytes_src1_q8_1 = ne13*ne12 * ne11*ne10_padded * sizeof(block_q8_1)/QK8_1 +
+        const size_t nbytes_src1_q8_1 = ne13*ne12 * ne11*ne10_padded * y_block_size/y_values_per_block +
             ggml_cuda_mmq_get_J_max(src0->type, fallback, cc, ne11) * sizeof(block_q8_1_mmq);
         ggml_cuda_pool_alloc<char> src1_q8_1(ctx.pool(), nbytes_src1_q8_1);
 
@@ -148,7 +149,7 @@ void ggml_cuda_mul_mat_q(
 
         // Stride depends on quantization format
         const int64_t s12 = use_native_fp4 ?
-                                ne11 * ne10_padded * sizeof(block_fp4_mmq) / (QK_K * sizeof(int)) :  // block_fp4_mmq holds 256 values
+                                ne11 * ne10_padded * sizeof(block_fp4_mmq) / (QK_FP4_MMQ * sizeof(int)) :
                                 ne11 * ne10_padded * sizeof(block_q8_1) / (QK8_1 * sizeof(int));
         const int64_t s13 = ne12*s12;
 
@@ -184,7 +185,7 @@ void ggml_cuda_mul_mat_q(
         CUDA_CHECK(cudaGetLastError());
     }
 
-    const size_t nbytes_src1_q8_1 = ne12*n_expert_used*ne10_padded * sizeof(block_q8_1)/QK8_1 +
+    const size_t nbytes_src1_q8_1 = ne12*n_expert_used*ne10_padded * y_block_size/y_values_per_block +
         ggml_cuda_mmq_get_J_max(src0->type, fallback, cc, ne11) * sizeof(block_q8_1_mmq);
     ggml_cuda_pool_alloc<char> src1_q8_1(ctx.pool(), nbytes_src1_q8_1);
 
@@ -207,8 +208,8 @@ void ggml_cuda_mul_mat_q(
         CUDA_CHECK(cudaGetLastError());
     }
 
-    static_assert(QK_K == 8 * QK_MXFP4, "QK_K needs to be 8 * QK_MXFP4");
-    const int64_t s12 = use_native_fp4 ? ne11 * ne10_padded * sizeof(block_fp4_mmq) / (QK_K * sizeof(int)) :
+    static_assert(QK_FP4_MMQ == 8 * QK_MXFP4, "QK_FP4_MMQ needs to be 8 * QK_MXFP4");
+    const int64_t s12 = use_native_fp4 ? ne11 * ne10_padded * sizeof(block_fp4_mmq) / (QK_FP4_MMQ * sizeof(int)) :
                                          ne11 * ne10_padded * sizeof(block_q8_1) / (QK8_1 * sizeof(int));
     const int64_t s13 = ne12*s12;
 

ggml/src/ggml-cuda/mmq.cuh

@@ -21,6 +21,9 @@ enum mmq_q8_1_ds_layout {
     MMQ_Q8_1_DS_LAYOUT_D2S6,
 };
 
+static constexpr int QK8_1_MMQ  = 4*QK8_1;
+static constexpr int QK_FP4_MMQ = 2*QK8_1_MMQ;
+
 struct block_q8_1_mmq {
     // The y float data is converted to a data layout that can simply be copied to shared memory as a contiguous block.
     // The y float data is first grouped as blocks of 128 values.
@@ -39,7 +42,7 @@ struct block_q8_1_mmq {
         half  d2s6[8];  // 1 16 bit scale per 64 values + 1 16 bit partial sum per 16 values for the first 96 values,
                         //     stored as d0,d1,s1,s2,s3,s4,s5
     };
-    int8_t qs[4*QK8_1]; // 128 values quantized to 8 bit each
+    int8_t qs[QK8_1_MMQ];
 };
 
 // this struct is used for fp4 data types (currently only used for Blackwell)
@@ -47,10 +50,10 @@ struct block_q8_1_mmq {
 // nvfp4 has block size 16, each int32 of d4 contains 4 ue4m3 scales
 struct block_fp4_mmq {
     uint32_t d4[4];
-    int8_t   qs[4 * 32];  // 256 FP4 values packed as 4-bit pairs (2 per byte)
+    int8_t   qs[QK_FP4_MMQ / 2];
 };
 
-static_assert(sizeof(block_q8_1_mmq) == 4*QK8_1 + 4*sizeof(half2), "Unexpected block_q8_1_mmq size");
+static_assert(sizeof(block_q8_1_mmq) == QK8_1_MMQ + 4*sizeof(half2), "Unexpected block_q8_1_mmq size");
 static_assert(sizeof(block_q8_1_mmq) == 4*sizeof(block_q8_1),      "Unexpected block_q8_1_mmq size");
 static_assert(sizeof(block_fp4_mmq)  == sizeof(block_q8_1_mmq),    "Unexpected block_fp4_mmq size");
 
@@ -833,9 +836,9 @@ static __device__ __forceinline__ void mul_mat_q_process_tile(
 
 #if defined(BLACKWELL_MMA_AVAILABLE)
     // FP4 tile stores 8 blocks
-    constexpr int ne_block = (type == GGML_TYPE_MXFP4 || type == GGML_TYPE_NVFP4) ? QK_K : 4 * QK8_1;
+    constexpr int ne_block = (type == GGML_TYPE_MXFP4 || type == GGML_TYPE_NVFP4) ? QK_FP4_MMQ : QK8_1_MMQ;
 #else
-    constexpr int ne_block = 4 * QK8_1;
+    constexpr int ne_block = QK8_1_MMQ;
 #endif  // defined(BLACKWELL_MMA_AVAILABLE)
 
     constexpr int ITER_K          = ggml_cuda_mmq_get_K_vram(type, J, fallback);

ggml/src/ggml-cuda/quantize.cu

@@ -90,8 +90,8 @@ static __global__ void quantize_mmq_nvfp4(
     const int64_t i2 = blockIdx.z % ne2;
     const int64_t i3 = blockIdx.z / ne2;
     const int64_t i01 = ids ? ids[i1] : i1;
-    const int64_t k_block = i0_base / QK_K;
-    const int64_t blocks_per_col = (ne0 + QK_K - 1) / QK_K;
+    const int64_t k_block = i0_base / QK_FP4_MMQ;
+    const int64_t blocks_per_col = (ne0 + QK_FP4_MMQ - 1) / QK_FP4_MMQ;
     if (k_block >= blocks_per_col) {
         return;
     }
@@ -100,7 +100,7 @@ static __global__ void quantize_mmq_nvfp4(
     block_fp4_mmq * y = (block_fp4_mmq *) vy;
     block_fp4_mmq * yb = y + ib;
 
-    const int sub = (i0_base % QK_K) / QK_NVFP4_SUB;
+    const int sub = (i0_base % QK_FP4_MMQ) / QK_NVFP4_SUB;
 
     float vals_raw[QK_NVFP4_SUB];
     float amax_raw = 0.0f;
@@ -207,7 +207,7 @@ static __global__ void quantize_mmq_mxfp4(const float * __restrict__ x,
 
     block_fp4_mmq * y = (block_fp4_mmq *) vy;
 
-    const int64_t block_fp4_mmq_size = 8 * QK_MXFP4;  // 256 values
+    const int64_t block_fp4_mmq_size = QK_FP4_MMQ;
     const int64_t ib0                = blockIdx.z * ((int64_t) ne1 * (ne0 / block_fp4_mmq_size));
     const int64_t ib = ib0 + (warp_start_offset / block_fp4_mmq_size) * ne1 + blockIdx.x;
     const int64_t quad_idx_in_block  = (warp_start_offset % block_fp4_mmq_size) / vals_per_warp;
@@ -303,8 +303,8 @@ static __global__ void quantize_mmq_q8_1(
     block_q8_1_mmq * y = (block_q8_1_mmq *) vy;
 
     const int64_t ib0 = blockIdx.z*((int64_t)gridDim.x*gridDim.y*blockDim.x/QK8_1); // first block of channel
-    const int64_t ib  = ib0 + (i0 / (4*QK8_1))*ne1 + blockIdx.x;                    // block index in channel
-    const int64_t iqs = i0 % (4*QK8_1);                                             // quant index in block
+    const int64_t ib  = ib0 + (i0 / QK8_1_MMQ)*ne1 + blockIdx.x;                    // block index in channel
+    const int64_t iqs = i0 % QK8_1_MMQ;                                             // quant index in block
 
     // Load 4 floats per thread and calculate max. abs. value between them:
     const float4 xi = i0 < ne00 ? x4[(i03*s03 + i02*s02 + i01*s01 + i00)/4] : make_float4(0.0f, 0.0f, 0.0f, 0.0f);
@@ -394,7 +394,7 @@ void quantize_mmq_q8_1_cuda(
         const int64_t ne00, const int64_t s01, const int64_t s02, const int64_t s03,
         const int64_t ne0, const int64_t ne1, const int64_t ne2, const int64_t ne3, cudaStream_t stream) {
     GGML_ASSERT(ne00 % 4 == 0);
-    GGML_ASSERT(ne0 % (4*QK8_1) == 0);
+    GGML_ASSERT(ne0 % QK8_1_MMQ == 0);
 
     // ne1 tends to assume the highest values, therefore use it as the "x" dimension of the CUDA grid:
     const int64_t block_num_y = (ne0 + 4*CUDA_QUANTIZE_BLOCK_SIZE_MMQ - 1) / (4*CUDA_QUANTIZE_BLOCK_SIZE_MMQ);

ggml/src/ggml-cuda/vecdotq.cuh

@@ -681,35 +681,40 @@ static __device__ __forceinline__ float vec_dot_q1_0_q8_1(
     // Q8_1: 32 elements per block with individual scales
     // iqs selects which of the 4 chunks of 32 elements to process (0-3)
 
-    const float d1 = bq1_0->d;
+    const float     d1 = bq1_0->d;
+    const int16_t * qs = (const int16_t *) bq1_0->qs + iqs * 2;
 
     // Process only the chunk specified by iqs
     const block_q8_1 * bq8_1_chunk = bq8_1 + iqs;
 
-    // Load 32 bits (4 bytes) for this chunk from Q1_0
-    const int offset = iqs * 4;
-    const int v = bq1_0->qs[offset + 0] | (bq1_0->qs[offset + 1] << 8) |
-                  (bq1_0->qs[offset + 2] << 16) | (bq1_0->qs[offset + 3] << 24);
-
-    // Unpack 32 bits into 32 signed values (-1 or +1)
-    int vi_bytes[8];
-#pragma unroll
-    for (int j = 0; j < 8; ++j) {
-        const int shift = j * 4;
-        const int bits4 = (v >> shift) & 0x0F;
-        const int b0 = (bits4 & 0x01) ? 1 : -1;
-        const int b1 = (bits4 & 0x02) ? 1 : -1;
-        const int b2 = (bits4 & 0x04) ? 1 : -1;
-        const int b3 = (bits4 & 0x08) ? 1 : -1;
-        vi_bytes[j] = (b0 & 0xFF) | ((b1 & 0xFF) << 8) | ((b2 & 0xFF) << 16) | ((b3 & 0xFF) << 24);
-    }
-
-    // Compute dot product for this 32-element chunk
     int sumi = 0;
 #pragma unroll
-    for (int j = 0; j < 8; ++j) {
-        const int u = get_int_b4(bq8_1_chunk->qs, j);
-        sumi = ggml_cuda_dp4a(vi_bytes[j], u, sumi);
+    for (int j = 0; j < 2; ++j) {
+        const int q  = qs[j];
+
+        const int u0 = get_int_b4(bq8_1_chunk->qs, j*4+0);
+        const int u1 = get_int_b4(bq8_1_chunk->qs, j*4+1);
+        const int u2 = get_int_b4(bq8_1_chunk->qs, j*4+2);
+        const int u3 = get_int_b4(bq8_1_chunk->qs, j*4+3);
+
+        // unpack crumbs into nibble indices
+        const int n0 = __byte_perm(0x11100100, 0x11100100, q >> 0); // [0, 1, 4, 5] [ 8,  9, 12, 13]
+        const int n1 = __byte_perm(0x11100100, 0x11100100, q >> 2); // [2, 3, 6, 7] [10, 11, 14, 15]
+        // unpack nibbles into byte values
+        const int s0 = __byte_perm(0x01FF, 0x01FF, n0 >>  0);
+        const int s1 = __byte_perm(0x01FF, 0x01FF, n1 >>  0);
+        const int s2 = __byte_perm(0x01FF, 0x01FF, n0 >> 16);
+        const int s3 = __byte_perm(0x01FF, 0x01FF, n1 >> 16);
+        // unshuffle values
+        const int v0 = __byte_perm(s0, s1, 0x5410);
+        const int v1 = __byte_perm(s0, s1, 0x7632);
+        const int v2 = __byte_perm(s2, s3, 0x5410);
+        const int v3 = __byte_perm(s2, s3, 0x7632);
+
+        sumi = ggml_cuda_dp4a(v0, u0, sumi);
+        sumi = ggml_cuda_dp4a(v1, u1, sumi);
+        sumi = ggml_cuda_dp4a(v2, u2, sumi);
+        sumi = ggml_cuda_dp4a(v3, u3, sumi);
     }
 
     // Apply Q1_0's single scale and this chunk's Q8_1 scale

ggml/src/ggml-hexagon/htp/hmx-queue.c

@@ -38,7 +38,7 @@ static inline void hmx_queue_process(struct hmx_queue *q, bool* killed) {
         if (!d->done) {
             FARF(HIGH, "hmx-queue-process: ir %u func %p data %p", ir, d->func, d->data);
 
-            enum hmx_queue_signal sig = (enum hmx_queue_signal) (unsigned int) d->func;
+            uintptr_t sig = (uintptr_t) d->func;
             switch (sig) {
                 case HMX_QUEUE_NOOP:    /* noop */;     break;
                 case HMX_QUEUE_KILL:    *killed = true; break;

ggml/src/ggml-metal/ggml-metal-device.cpp

@@ -1834,6 +1834,23 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_col2im_1d(ggml_m
     return res;
 }
 
+ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_snake(ggml_metal_library_t lib, enum ggml_type type) {
+    GGML_ASSERT(type == GGML_TYPE_F32 || type == GGML_TYPE_F16 || type == GGML_TYPE_BF16);
+
+    char base[256];
+    char name[256];
+
+    snprintf(base, 256, "kernel_snake_%s", ggml_type_name(type));
+    snprintf(name, 256, "%s", base);
+
+    ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
+    if (!res.pipeline) {
+        res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr);
+    }
+
+    return res;
+}
+
 ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_conv_transpose_2d(ggml_metal_library_t lib, const ggml_tensor * op) {
     assert(op->op == GGML_OP_CONV_TRANSPOSE_2D);
 

ggml/src/ggml-metal/ggml-metal-device.h

@@ -151,6 +151,7 @@ struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_im2col
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_conv_transpose_1d (ggml_metal_library_t lib, const struct ggml_tensor * op);
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_conv_transpose_2d (ggml_metal_library_t lib, const struct ggml_tensor * op);
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_col2im_1d         (ggml_metal_library_t lib, const struct ggml_tensor * op);
+struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_snake             (ggml_metal_library_t lib, enum ggml_type type);
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_conv_2d           (ggml_metal_library_t lib, const struct ggml_tensor * op);
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_conv_2d_dw        (ggml_metal_library_t lib, const struct ggml_tensor * op, bool tiled);
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_conv_3d           (ggml_metal_library_t lib, const struct ggml_tensor * op);

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

@@ -616,6 +616,11 @@ typedef struct {
     int32_t  p0;
 } ggml_metal_kargs_col2im_1d;
 
+typedef struct {
+    int32_t T;
+    int32_t C;
+} ggml_metal_kargs_snake;
+
 typedef struct {
     int32_t  IC;
     int32_t  IH;

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

@@ -3077,7 +3077,58 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) {
     return 1;
 }
 
+// Snake activation autofuse: mul -> sin -> sqr -> mul -> add
+static bool ggml_metal_op_can_fuse_snake(ggml_metal_op_t ctx, int idx) {
+    static constexpr ggml_op snake_ops[5] = { GGML_OP_MUL, GGML_OP_SIN, GGML_OP_SQR, GGML_OP_MUL, GGML_OP_ADD };
+
+    if (ctx->node(idx)->op != GGML_OP_MUL || !ctx->can_fuse(idx, snake_ops, 5)) {
+        return false;
+    }
+
+    const ggml_tensor * mul0     = ctx->node(idx + 0);
+    const ggml_tensor * sin_node = ctx->node(idx + 1);
+    const ggml_tensor * sqr      = ctx->node(idx + 2);
+    const ggml_tensor * mul1     = ctx->node(idx + 3);
+    const ggml_tensor * add      = ctx->node(idx + 4);
+
+    // x carries the full activation shape, a is the broadcast operand
+    const ggml_tensor * x = ggml_are_same_shape(mul0, mul0->src[0]) ? mul0->src[0] : mul0->src[1];
+    const ggml_tensor * a = (x == mul0->src[0]) ? mul0->src[1] : mul0->src[0];
+
+    // mul1 reads sqr and inv_b in either operand order
+    const ggml_tensor * inv_b    = (mul1->src[0] == sqr) ? mul1->src[1] : mul1->src[0];
+
+    // closure check: the trailing add reads the same x as the leading mul
+    const ggml_tensor * x_in_add = (add->src[0] == mul1) ? add->src[1] : add->src[0];
+
+    // x is in the supported whitelist and every chain intermediate shares x's type.
+    // a and inv_b bind as device const float * in the kernel, so they stay F32.
+    const bool types_ok =
+        (x->type == GGML_TYPE_F32 || x->type == GGML_TYPE_F16 || x->type == GGML_TYPE_BF16) &&
+        (a->type    == GGML_TYPE_F32) && (inv_b->type    == GGML_TYPE_F32) &&
+        (mul0->type == x->type)       && (sin_node->type == x->type) &&
+        (sqr->type  == x->type)       && (mul1->type     == x->type) &&
+        (add->type  == x->type);
+    // a / inv_b collapse to [1, C, 1, 1], x and add stay 2D
+    const bool shape_ok = ggml_are_same_shape(a, inv_b) && a->ne[0] == 1 && a->ne[1] == x->ne[1];
+    const bool dim_ok =
+        (x->ne[2]     == 1) && (x->ne[3]     == 1) &&
+        (add->ne[2]   == 1) && (add->ne[3]   == 1) &&
+        (a->ne[2]     == 1) && (a->ne[3]     == 1) &&
+        (inv_b->ne[2] == 1) && (inv_b->ne[3] == 1);
+    // kernel reads x[idx] and a[c] / inv_b[c] linearly, so every operand is contiguous
+    const bool contig_ok =
+        ggml_is_contiguous(x) && ggml_is_contiguous(add) &&
+        ggml_is_contiguous(a) && ggml_is_contiguous(inv_b);
+
+    return types_ok && shape_ok && dim_ok && contig_ok && x_in_add == x;
+}
+
 int ggml_metal_op_bin(ggml_metal_op_t ctx, int idx) {
+    if (ctx->use_fusion && ggml_metal_op_can_fuse_snake(ctx, idx)) {
+        return ggml_metal_op_snake_fused(ctx, idx);
+    }
+
     ggml_tensor * op = ctx->node(idx);
 
     ggml_metal_library_t lib = ctx->lib;
@@ -3984,6 +4035,55 @@ int ggml_metal_op_col2im_1d(ggml_metal_op_t ctx, int idx) {
     return 1;
 }
 
+// Dispatch the fused snake kernel from the matched mul -> sin -> sqr -> mul -> add chain.
+// idx points at the leading mul. The caller has validated the chain.
+int ggml_metal_op_snake_fused(ggml_metal_op_t ctx, int idx) {
+    ggml_metal_library_t lib = ctx->lib;
+    ggml_metal_encoder_t enc = ctx->enc;
+
+    const ggml_tensor * mul0 = ctx->node(idx + 0);
+    const ggml_tensor * sqr  = ctx->node(idx + 2);
+    const ggml_tensor * mul1 = ctx->node(idx + 3);
+    ggml_tensor *       add  = ctx->node(idx + 4);
+
+    const ggml_tensor * x = ggml_are_same_shape(mul0, mul0->src[0]) ? mul0->src[0] : mul0->src[1];
+    const ggml_tensor * a = (x == mul0->src[0]) ? mul0->src[1] : mul0->src[0];
+    const ggml_tensor * inv_b = (mul1->src[0] == sqr) ? mul1->src[1] : mul1->src[0];
+
+    const int T     = (int) x->ne[0];
+    const int C     = (int) x->ne[1];
+    const int total = T * C;
+
+    // the encode loop pre-checked the leading mul only, check the rest of the chain
+    for (int i = 1; i < 5; ++i) {
+        if (!ggml_metal_op_concurrency_check(ctx, ctx->node(idx + i))) {
+            ggml_metal_op_concurrency_reset(ctx);
+
+            break;
+        }
+    }
+
+    auto pipeline = ggml_metal_library_get_pipeline_snake(lib, x->type);
+
+    ggml_metal_kargs_snake args = {
+        /*.T =*/ T,
+        /*.C =*/ C,
+    };
+
+    ggml_metal_encoder_set_pipeline(enc, pipeline);
+    ggml_metal_encoder_set_bytes   (enc, &args, sizeof(args), 0);
+    ggml_metal_encoder_set_buffer  (enc, ggml_metal_get_buffer_id(x),     1);
+    ggml_metal_encoder_set_buffer  (enc, ggml_metal_get_buffer_id(a),     2);
+    ggml_metal_encoder_set_buffer  (enc, ggml_metal_get_buffer_id(inv_b), 3);
+    ggml_metal_encoder_set_buffer  (enc, ggml_metal_get_buffer_id(add),   4);
+
+    const int nth = 256;
+    const int ntg = (total + nth - 1) / nth;
+    ggml_metal_encoder_dispatch_threadgroups(enc, ntg, 1, 1, nth, 1, 1);
+
+    return 5;
+}
+
 int ggml_metal_op_conv_transpose_2d(ggml_metal_op_t ctx, int idx) {
     ggml_tensor * op = ctx->node(idx);
 

ggml/src/ggml-metal/ggml-metal-ops.h

@@ -80,6 +80,7 @@ int ggml_metal_op_conv_3d           (ggml_metal_op_t ctx, int idx);
 int ggml_metal_op_conv_transpose_1d (ggml_metal_op_t ctx, int idx);
 int ggml_metal_op_conv_transpose_2d (ggml_metal_op_t ctx, int idx);
 int ggml_metal_op_col2im_1d         (ggml_metal_op_t ctx, int idx);
+int ggml_metal_op_snake_fused       (ggml_metal_op_t ctx, int idx);
 int ggml_metal_op_upscale           (ggml_metal_op_t ctx, int idx);
 int ggml_metal_op_pad               (ggml_metal_op_t ctx, int idx);
 int ggml_metal_op_pad_reflect_1d    (ggml_metal_op_t ctx, int idx);

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

@@ -5406,6 +5406,35 @@ template [[host_name("kernel_col2im_1d_bf16")]] kernel void kernel_col2im_1d<bfl
 #endif
 
 
+template <typename T>
+kernel void kernel_snake(
+        constant ggml_metal_kargs_snake & args,
+        device const T     * x,
+        device const float * a,
+        device const float * inv_b,
+        device       T     * dst,
+        uint         tgpig [[threadgroup_position_in_grid]],
+        uint         tpitg [[thread_position_in_threadgroup]],
+        uint         ntg   [[threads_per_threadgroup]]) {
+
+    const int idx = tgpig * ntg + tpitg;
+    if (idx >= args.T * args.C) {
+        return;
+    }
+
+    const int   c  = idx / args.T;  // x is [T, C], a / inv_b collapse to [1, C]
+    const float xi = float(x[idx]);
+    const float si = sin(a[c] * xi);
+    dst[idx] = T(xi + si * si * inv_b[c]);
+}
+
+template [[host_name("kernel_snake_f32")]]  kernel void kernel_snake<float>(constant ggml_metal_kargs_snake &, device const float *, device const float *, device const float *, device float *, uint, uint, uint);
+template [[host_name("kernel_snake_f16")]]  kernel void kernel_snake<half>(constant ggml_metal_kargs_snake &, device const half *, device const float *, device const float *, device half *, uint, uint, uint);
+#if defined(GGML_METAL_HAS_BF16)
+template [[host_name("kernel_snake_bf16")]] kernel void kernel_snake<bfloat>(constant ggml_metal_kargs_snake &, device const bfloat *, device const float *, device const float *, device bfloat *, uint, uint, uint);
+#endif
+
+
 typedef void (conv_transpose_2d_t)(
         constant ggml_metal_kargs_conv_transpose_2d & args,
         device const float * src0,

ggml/src/ggml-opencl/ggml-opencl.cpp

@@ -114,6 +114,7 @@ enum GPU_FAMILY {
 
 enum ADRENO_GPU_GEN {
     ADRENO_UNKNOWN,
+    A6X,
     A7X,
     A8X,
     X1E,
@@ -122,6 +123,7 @@ enum ADRENO_GPU_GEN {
 
 enum ADRENO_CL_COMPILER_TYPE {
     E031,
+    E17,
     DX,
 };
 
@@ -243,6 +245,19 @@ static ggml_cl_version get_opencl_c_version(ggml_cl_version platform_version, cl
 }
 
 static ADRENO_GPU_GEN get_adreno_gpu_gen(const char *device_name) {
+    if (strstr(device_name, "610") || strstr(device_name, "612") ||
+        strstr(device_name, "613") || strstr(device_name, "615") ||
+        strstr(device_name, "616") || strstr(device_name, "618") ||
+        strstr(device_name, "619") || strstr(device_name, "620") ||
+        strstr(device_name, "630") || strstr(device_name, "640") ||
+        strstr(device_name, "642") || strstr(device_name, "643") ||
+        strstr(device_name, "644") || strstr(device_name, "650") ||
+        strstr(device_name, "660") || strstr(device_name, "663") ||
+        strstr(device_name, "680") || strstr(device_name, "685") ||
+        strstr(device_name, "690")) {
+        return ADRENO_GPU_GEN::A6X;
+    }
+
     if (strstr(device_name, "730") ||
         strstr(device_name, "740") ||
         strstr(device_name, "750")) {
@@ -250,7 +265,8 @@ static ADRENO_GPU_GEN get_adreno_gpu_gen(const char *device_name) {
     }
 
     if (strstr(device_name, "830") ||
-        strstr(device_name, "840")) {
+        strstr(device_name, "840") ||
+        strstr(device_name, "850")) {
         return ADRENO_GPU_GEN::A8X;
     }
 
@@ -274,6 +290,17 @@ static ggml_cl_compiler_version get_adreno_cl_compiler_version(const char *drive
     size_t compiler_minor_offset = 8;
     size_t compiler_patch_offset = 11;
 
+    if (compiler_ver_pos == std::string::npos) {
+        compiler_ver_pos = driver_ver_str.find("E17");
+        if (compiler_ver_pos != std::string::npos) {
+            type = ADRENO_CL_COMPILER_TYPE::E17;
+            compiler_ver_len = 12;
+            compiler_major_offset = 4;
+            compiler_minor_offset = 7;
+            compiler_patch_offset = 10;
+        }
+    }
+
     if (compiler_ver_pos == std::string::npos) {
         compiler_ver_pos = driver_ver_str.find("DX");
         if (compiler_ver_pos == std::string::npos) {
@@ -282,6 +309,8 @@ static ggml_cl_compiler_version get_adreno_cl_compiler_version(const char *drive
         type = ADRENO_CL_COMPILER_TYPE::DX;
         compiler_ver_len = 11;
         compiler_major_offset = 3;
+        compiler_minor_offset = 6;
+        compiler_patch_offset = 9;
     }
 
     std::string compiler_ver_str = driver_ver_str.substr(compiler_ver_pos, compiler_ver_len);
@@ -1641,6 +1670,7 @@ static void load_cl_kernels(ggml_backend_opencl_context *backend_ctx) {
     // those compiler versions since it is anyway not used for Adreno.
     if (backend_ctx->gpu_family != ADRENO ||
         backend_ctx->adreno_cl_compiler_version.newer_than_or_same(E031, 38, 11, 0) ||
+        backend_ctx->adreno_cl_compiler_version.type == E17 ||
         backend_ctx->adreno_cl_compiler_version.type == DX) {
 #ifdef GGML_OPENCL_EMBED_KERNELS
         const std::string kernel_src {
@@ -6932,8 +6962,31 @@ inline bool use_adreno_kernels(const ggml_backend_opencl_context *backend_ctx, c
     return threashold_ok;
 }
 
+static bool adreno_e17_compiler_quirks(const ggml_backend_opencl_context *backend_ctx) {
+    if (!backend_ctx || backend_ctx->gpu_family != GPU_FAMILY::ADRENO ||
+        backend_ctx->adreno_cl_compiler_version.type != ADRENO_CL_COMPILER_TYPE::E17) {
+        return false;
+    }
+    const char * env = getenv("GGML_OPENCL_ADRENO_E17_QUIRKS");
+    return !(env && env[0] == '0');
+}
+
 inline bool use_adreno_moe_kernels(const ggml_backend_opencl_context *backend_ctx, const ggml_tensor *tensor) {
-    GGML_UNUSED(backend_ctx);
+    // The moe weight repack kernels *_trans4_ns alias a private ushort8 through a uchar*.
+    // Certain compilers (found with some A7x and A6x) miscompiles this, corrupting the weights.
+    // So, exclude A6x and A7x from using Adreno MoE kernels for now.
+    // The quants that have a general mul_mat_id kernel fallback to the general version; the
+    // rest fallback to CPU.
+    if (backend_ctx && (backend_ctx->adreno_gen == ADRENO_GPU_GEN::A6X ||
+                        backend_ctx->adreno_gen == ADRENO_GPU_GEN::A7X ||
+                        backend_ctx->adreno_gen == ADRENO_GPU_GEN::ADRENO_UNKNOWN)) {
+        return false;
+    }
+
+    if (adreno_e17_compiler_quirks(backend_ctx)) {
+        return false;
+    }
+
     int ne01 = tensor->ne[1];
     return (((strstr(tensor->name, "ffn") != NULL) && (strstr(tensor->name, "exps") != NULL)) || (strstr(tensor->name, "as") != NULL)) && (ne01 % 32 == 0);
 }
@@ -7266,6 +7319,10 @@ static bool ggml_opencl_supports_op(ggml_backend_dev_t dev, const struct ggml_te
         case GGML_OP_MEAN:
             return op->src[0]->type == GGML_TYPE_F32;
         case GGML_OP_FLASH_ATTN_EXT: {
+            // The E17 compilers segfault while building FA kernels, skip E17 for now
+            if (adreno_e17_compiler_quirks(backend_ctx)) {
+                return false;
+            }
             const ggml_tensor * q = op->src[0];
             const ggml_tensor * k = op->src[1];
             const ggml_tensor * v = op->src[2];
@@ -7319,6 +7376,14 @@ static bool ggml_opencl_supports_op(ggml_backend_dev_t dev, const struct ggml_te
                 return false;
             }
 
+            // Some compilers for A7x (Adreno 740, compiler E031.41) crashes when
+            // building FA kernels with mixed or quant types (f32_f16, f32_q8_0, f32_q4_0)
+            // Here we skip all A7x for these kernels to avoid crash
+            if (backend_ctx->adreno_gen == ADRENO_GPU_GEN::A7X &&
+                (is_f32_f16 || is_f32_q8_0 || is_f32_q4_0)) {
+                return false;
+            }
+
             if (dk == 512) {
                 if (backend_ctx->gpu_family == INTEL) {
                     return false;
@@ -10525,10 +10590,16 @@ static ggml_backend_buffer_t ggml_backend_opencl_buffer_type_alloc_buffer(ggml_b
 
     cl_int err;
     cl_mem mem = clCreateBuffer(backend_ctx->context, CL_MEM_READ_WRITE, size, NULL, &err);
+#if GGML_OPENCL_TARGET_VERSION >= 300
+    // clCreateBufferWithProperties and cl_mem_properties are OpenCL 3.0. Drivers older than
+    // that do not export the symbol, so a build targeting them fails to link. The large
+    // buffer extension is only ever enabled on drivers that are well past 3.0, so this path
+    // is dead there anyway.
     if (err != CL_SUCCESS && backend_ctx->adreno_use_large_buffer) {
         cl_mem_properties props[] = { 0x41A6 /* CL_LARGE_BUFFER_QCOM */, 1, 0 };
         mem = clCreateBufferWithProperties(backend_ctx->context, props, CL_MEM_READ_WRITE, size, NULL, &err);
     }
+#endif
 
     if (err != CL_SUCCESS) {
         GGML_LOG_INFO("%s: failed to allocate %.2f MiB\n", __func__, size / 1024.0 / 1024.0);

ggml/src/ggml-opencl/kernels/flash_attn_f32_f16.cl

@@ -30,6 +30,10 @@
 #elif defined(cl_qcom_subgroup_shuffle)
 #pragma OPENCL EXTENSION cl_qcom_subgroup_shuffle : enable
 #define HAS_SUBGROUP_SHUFFLE 1
+// Adreno compilers that expose only cl_qcom_subgroup_shuffle do not declare the KHR
+// name, so calling it is an implicit declaration and the program fails to build.
+// Route it to the qcom builtin.
+#define sub_group_shuffle_xor(val, mask) qcom_sub_group_shuffle_xor((val), (mask), CLK_SUB_GROUP_SHUFFLE_WIDTH_WAVE_SIZE_QCOM, 0.0f)
 #endif
 
 #define ACC_TYPE float

ggml/src/ggml-opencl/kernels/flash_attn_f32_q4_0.cl

@@ -10,6 +10,10 @@
 #elif defined(cl_qcom_subgroup_shuffle)
 #pragma OPENCL EXTENSION cl_qcom_subgroup_shuffle : enable
 #define HAS_SUBGROUP_SHUFFLE 1
+// Adreno compilers that expose only cl_qcom_subgroup_shuffle do not declare the KHR
+// name, so calling it is an implicit declaration and the program fails to build.
+// Route it to the qcom builtin.
+#define sub_group_shuffle_xor(val, mask) qcom_sub_group_shuffle_xor((val), (mask), CLK_SUB_GROUP_SHUFFLE_WIDTH_WAVE_SIZE_QCOM, 0.0f)
 #endif
 
 // Flash attention: Q=f32, K=q4_0, V=q4_0.

ggml/src/ggml-opencl/kernels/flash_attn_f32_q8_0.cl

@@ -10,6 +10,10 @@
 #elif defined(cl_qcom_subgroup_shuffle)
 #pragma OPENCL EXTENSION cl_qcom_subgroup_shuffle : enable
 #define HAS_SUBGROUP_SHUFFLE 1
+// Adreno compilers that expose only cl_qcom_subgroup_shuffle do not declare the KHR
+// name, so calling it is an implicit declaration and the program fails to build.
+// Route it to the qcom builtin.
+#define sub_group_shuffle_xor(val, mask) qcom_sub_group_shuffle_xor((val), (mask), CLK_SUB_GROUP_SHUFFLE_WIDTH_WAVE_SIZE_QCOM, 0.0f)
 #endif
 
 // Flash attention: Q=f32, K=q8_0, V=q8_0.

ggml/src/ggml-opencl/kernels/gemm_moe_mxfp4_f32_ns.cl

@@ -274,8 +274,9 @@ kernel void kernel_gemm_moe_mxfp4_f32_ns(
         shared_b[b_local_offset.y] = bx8_f16.hi;
 
         // Dequantization
-        reg_a.lo = mxfp4_to_fp16_packed8(as_ushort2(mxfp4x16.lo)) * s;
-        reg_a.hi = mxfp4_to_fp16_packed8(as_ushort2(mxfp4x16.hi)) * s;
+        // Cast the e8m0 scale to half to satisfy E17 compilers
+        reg_a.lo = mxfp4_to_fp16_packed8(as_ushort2(mxfp4x16.lo)) * (half)s;
+        reg_a.hi = mxfp4_to_fp16_packed8(as_ushort2(mxfp4x16.hi)) * (half)s;
 
         sub_group_barrier(CLK_LOCAL_MEM_FENCE);
 
@@ -304,8 +305,9 @@ kernel void kernel_gemm_moe_mxfp4_f32_ns(
         shared_b[b_local_offset.y] = bx8_f16.hi;
 
         // Dequantization
-        reg_a.lo = mxfp4_to_fp16_packed8(as_ushort2(mxfp4x16.lo)) * s;
-        reg_a.hi = mxfp4_to_fp16_packed8(as_ushort2(mxfp4x16.hi)) * s;
+        // Cast the e8m0 scale to half to satisfy E17 compilers
+        reg_a.lo = mxfp4_to_fp16_packed8(as_ushort2(mxfp4x16.lo)) * (half)s;
+        reg_a.hi = mxfp4_to_fp16_packed8(as_ushort2(mxfp4x16.hi)) * (half)s;
 
         sub_group_barrier(CLK_LOCAL_MEM_FENCE);
 

ggml/src/ggml-opencl/kernels/gemv_noshuffle_iq4_nl_f32.cl

@@ -296,7 +296,12 @@ kernel void kernel_gemv_noshuffle_iq4_nl_f32(
     // 2 outputs per fiber in wave 0
     if (groupId == 0) {
         dst = (global float*)((global char*)dst + offsetd);
-        vstore2(totalSum, 0, &(dst[gid * 2]));
+        // Guard the two output rows. The x-grid is padded to CEIL_DIV(ne01/2,64)*64,
+        // so when ne01 is not a multiple of 128 the tail row-pairs run past row ne01
+        // and would overrun dst into the adjacent tensor. No-op / byte-identical when
+        // ne01 % 128 == 0 (M/2 already a multiple of 64 -> no padding).
+        if (gid * 2 + 0 < M) dst[gid * 2 + 0] = totalSum.s0;
+        if (gid * 2 + 1 < M) dst[gid * 2 + 1] = totalSum.s1;
     }
 
 }

ggml/src/ggml-opencl/kernels/gemv_noshuffle_q1_0_f32.cl

@@ -116,6 +116,10 @@ __kernel void kernel_gemv_noshuffle_q1_0_f32(
 
     if (groupId == 0) {
         dst = (global float*)((global char*)dst + offsetd);
-        dst[gid] = totalSum;
+        // Guard the output row. The x-grid is padded to CEIL_DIV(M,wavesize)*wavesize,
+        // so when ne01 is not a multiple of the wave size the tail work-items run past
+        // row ne01 and would overrun dst into the adjacent tensor. No-op / byte-identical
+        // when ne01 is wave-aligned (no padding).
+        if (gid < M) dst[gid] = totalSum;
     }
 }

ggml/src/ggml-opencl/kernels/gemv_noshuffle_q4_0_f32.cl

@@ -268,7 +268,12 @@ __kernel void kernel_gemv_noshuffle_q4_0_f32(
     // 2 outputs per fiber in wave 0
     if (groupId == 0) {
         dst = (global float*)((global char*)dst + offsetd);
-        vstore2(totalSum, 0, &(dst[gid * 2]));
+        // Guard the two output rows. The x-grid is padded to CEIL_DIV(ne01/2,64)*64,
+        // so when ne01 is not a multiple of 128 the tail row-pairs run past row ne01
+        // and would overrun dst into the adjacent tensor. No-op / byte-identical when
+        // ne01 % 128 == 0 (M/2 already a multiple of 64 -> no padding).
+        if (gid * 2 + 0 < M) dst[gid * 2 + 0] = totalSum.s0;
+        if (gid * 2 + 1 < M) dst[gid * 2 + 1] = totalSum.s1;
     }
 
 }

ggml/src/ggml-opencl/kernels/gemv_noshuffle_q4_0_f32_spec.cl

@@ -262,7 +262,11 @@ __kernel void kernel_gemv_noshuffle_q4_0_f32(
     // 2 outputs per fiber in wave 0
     if (groupId == 0) {
         dst = (global float*)((global char*)dst + offsetd);
-        vstore2(totalSum, 0, &(dst[gid * 2]));
+        // Guard the two output rows against the padded x-grid tail overrunning dst.
+        // The current shape specializations are all ne01 % 128 == 0 (no padding), so
+        // this is a no-op / byte-identical today; keep it in lockstep with the base kernel.
+        if (gid * 2 + 0 < ne01) dst[gid * 2 + 0] = totalSum.s0;
+        if (gid * 2 + 1 < ne01) dst[gid * 2 + 1] = totalSum.s1;
     }
 
 }

ggml/src/ggml-opencl/kernels/gemv_noshuffle_q4_1_f32.cl

@@ -277,7 +277,12 @@ kernel void kernel_gemv_noshuffle_q4_1_f32(
     // 2 outputs per fiber in wave 0
     if (groupId == 0) {
         dst = (global float*)((global char*)dst + offsetd);
-        vstore2(totalSum, 0, &(dst[gid * 2]));
+        // Guard the two output rows. The x-grid is padded to CEIL_DIV(ne01/2,64)*64,
+        // so when ne01 is not a multiple of 128 the tail row-pairs run past row ne01
+        // and would overrun dst into the adjacent tensor. No-op / byte-identical when
+        // ne01 % 128 == 0 (M/2 already a multiple of 64 -> no padding).
+        if (gid * 2 + 0 < M) dst[gid * 2 + 0] = totalSum.s0;
+        if (gid * 2 + 1 < M) dst[gid * 2 + 1] = totalSum.s1;
     }
 
 }

ggml/src/ggml-opencl/kernels/gemv_noshuffle_q4_k_f32.cl

@@ -312,7 +312,12 @@ kernel void kernel_gemv_noshuffle_q4_k_f32(
     // 2 outputs per fiber in wave 0
     if (groupId == 0) {
         dst = (global float*)((global char*)dst + offsetd);
-        vstore2(totalSum, 0, &(dst[gid * 2]));
+        // Guard the two output rows. The x-grid is padded to CEIL_DIV(ne01/2,64)*64,
+        // so when ne01 is not a multiple of 128 the tail row-pairs run past row ne01
+        // and would overrun dst into the adjacent tensor. No-op / byte-identical when
+        // ne01 % 128 == 0 (M/2 already a multiple of 64 -> no padding).
+        if (gid * 2 + 0 < M) dst[gid * 2 + 0] = totalSum.s0;
+        if (gid * 2 + 1 < M) dst[gid * 2 + 1] = totalSum.s1;
     }
 
 }

ggml/src/ggml-opencl/kernels/gemv_noshuffle_q5_0_f32.cl

@@ -285,7 +285,12 @@ __kernel void kernel_gemv_noshuffle_q5_0_f32(
     // 2 outputs per fiber in wave 0
     if (groupId == 0) {
         dst = (global float*)((global char*)dst + offsetd);
-        vstore2(totalSum, 0, &(dst[gid * 2]));
+        // Guard the two output rows. The x-grid is padded to CEIL_DIV(ne01/2,64)*64,
+        // so when ne01 is not a multiple of 128 the tail row-pairs run past row ne01
+        // and would overrun dst into the adjacent tensor. No-op / byte-identical when
+        // ne01 % 128 == 0 (M/2 already a multiple of 64 -> no padding).
+        if (gid * 2 + 0 < M) dst[gid * 2 + 0] = totalSum.s0;
+        if (gid * 2 + 1 < M) dst[gid * 2 + 1] = totalSum.s1;
     }
 
 }

ggml/src/ggml-opencl/kernels/gemv_noshuffle_q5_1_f32.cl

@@ -288,7 +288,12 @@ __kernel void kernel_gemv_noshuffle_q5_1_f32(
     // 2 outputs per fiber in wave 0
     if (groupId == 0) {
         dst = (global float*)((global char*)dst + offsetd);
-        vstore2(totalSum, 0, &(dst[gid * 2]));
+        // Guard the two output rows. The x-grid is padded to CEIL_DIV(ne01/2,64)*64,
+        // so when ne01 is not a multiple of 128 the tail row-pairs run past row ne01
+        // and would overrun dst into the adjacent tensor. No-op / byte-identical when
+        // ne01 % 128 == 0 (M/2 already a multiple of 64 -> no padding).
+        if (gid * 2 + 0 < M) dst[gid * 2 + 0] = totalSum.s0;
+        if (gid * 2 + 1 < M) dst[gid * 2 + 1] = totalSum.s1;
     }
 
 }

ggml/src/ggml-opencl/kernels/gemv_noshuffle_q5_k_f32.cl

@@ -321,6 +321,11 @@ kernel void kernel_gemv_noshuffle_q5_k_f32(
     // 2 outputs per fiber in wave 0
     if (groupId == 0) {
         dst = (global float*)((global char*)dst + offsetd);
-        vstore2(totalSum, 0, &(dst[gid * 2]));
+        // Guard the two output rows. The x-grid is padded to CEIL_DIV(ne01/2,64)*64,
+        // so when ne01 is not a multiple of 128 the tail row-pairs run past row ne01
+        // and would overrun dst into the adjacent tensor. No-op / byte-identical when
+        // ne01 % 128 == 0 (M/2 already a multiple of 64 -> no padding).
+        if (gid * 2 + 0 < M) dst[gid * 2 + 0] = totalSum.s0;
+        if (gid * 2 + 1 < M) dst[gid * 2 + 1] = totalSum.s1;
     }
 }

ggml/src/ggml-opencl/kernels/gemv_noshuffle_q6_k_f32.cl

@@ -288,6 +288,11 @@ kernel void kernel_gemv_noshuffle_q6_K_f32(
 
     if (grp == 0) {
         dst = (global float*)((global char*)dst + offsetd);
-        vstore2(total_sum, 0, &(dst[gid * 2]));
+        // Guard the two output rows. The x-grid is padded to CEIL_DIV(ne01/2,64)*64,
+        // so when ne01 is not a multiple of 128 the tail row-pairs run past row ne01
+        // and would overrun dst into the adjacent tensor (garbage downstream).
+        // No-op / byte-identical when ne01 % 128 == 0 (no padding).
+        if (gid * 2 + 0 < ne01) dst[gid * 2 + 0] = total_sum.s0;
+        if (gid * 2 + 1 < ne01) dst[gid * 2 + 1] = total_sum.s1;
     }
 }

ggml/src/ggml-opencl/kernels/gemv_noshuffle_q8_0_f32.cl

@@ -190,6 +190,10 @@ __kernel void kernel_gemv_noshuffle_q8_0_f32(
     // 1 outputs per fiber in wave 0
     if (groupId == 0) {
         dst = (global float*)((global char*)dst + offsetd);
-        dst[gid] = totalSum;
+        // Guard the output row. The x-grid is padded to CEIL_DIV(M,wavesize)*wavesize,
+        // so when ne01 is not a multiple of the wave size the tail work-items run past
+        // row ne01 and would overrun dst into the adjacent tensor. No-op / byte-identical
+        // when ne01 is wave-aligned (no padding).
+        if (gid < M) dst[gid] = totalSum;
     }
 }

ggml/src/ggml-opencl/kernels/mul_mv_f16_f16.cl

@@ -64,7 +64,14 @@ kernel void kernel_mul_mat_f16_f16(
 
     global half * x = (global half *) (src0 + offset_src0);
 
-    if (ne00 < 128) {
+    // The vector path below casts the row pointers to half4, which must be 8-byte aligned.
+    // A row address is r0*nb01 + ..., and a permuted or strided src leaves nb01/nb11
+    // unconstrained -- an odd ne00, say, gives a row that is only 2-byte aligned. Every
+    // src1 row this work-item walks is src1_base + r1*nb11, so require both.
+    const ulong src1_base  = (ulong) (src1 + (i12)*nb12 + (i13)*nb13);
+    const bool  row_aligned = (((ulong) x) & 7) == 0 && (src1_base & 7) == 0 && (nb11 & 7) == 0;
+
+    if (ne00 < 128 || !row_aligned) {
         for (int row = 0; row < N_F16_F16; ++row) {
             int r1 = rb + row;
             if (r1 >= ne11) {

ggml/src/ggml-opencl/kernels/mul_mv_f16_f32.cl

@@ -64,7 +64,14 @@ kernel void kernel_mul_mat_f16_f32(
 
     global half * x = (global half *) (src0 + offset_src0);
 
-    if (ne00 < 128) {
+    // The vector path below casts the row pointers to half4/float4, which must be 8- and
+    // 16-byte aligned. A row address is r0*nb01 + ..., and a permuted or strided src leaves
+    // nb01/nb11 unconstrained -- an odd ne00, say, gives a row that is only 2-byte aligned.
+    // Every src1 row this work-item walks is src1_base + r1*nb11, so require both.
+    const ulong src1_base  = (ulong) (src1 + (i12)*nb12 + (i13)*nb13);
+    const bool  row_aligned = (((ulong) x) & 7) == 0 && (src1_base & 15) == 0 && (nb11 & 15) == 0;
+
+    if (ne00 < 128 || !row_aligned) {
         for (int row = 0; row < N_F16_F32; ++row) {
             int r1 = rb + row;
             if (r1 >= ne11) {

ggml/src/ggml-opencl/kernels/mul_mv_f16_f32_1row.cl

@@ -64,8 +64,15 @@ kernel void kernel_mul_mat_f16_f32_1row(
     global half  * x = (global half  *) (src0 + offset_src0);
     global float * y = (global float *) (src1 + offset_src1);
 
+    // The vector path below casts the row pointers to half4/float4, which must be 8- and
+    // 16-byte aligned. A row address is r0*nb01 + ..., and a permuted or strided src leaves
+    // nb01/nb11 unconstrained -- an odd ne00, say, gives a row that is only 2-byte aligned.
+    // Take the vector path only when the rows this work-item touches are actually aligned;
+    // the scalar loop has no such requirement.
+    const bool row_aligned = (((ulong) x) & 7) == 0 && (((ulong) y) & 15) == 0;
+
     float sumf = 0;
-    if (ne00 < 128) {
+    if (ne00 < 128 || !row_aligned) {
         for (int i = get_sub_group_local_id(); i < ne00; i += get_max_sub_group_size()) {
             sumf += (float) x[i] * (float) y[i];
         }

ggml/src/ggml-opencl/kernels/mul_mv_f16_f32_l4.cl

@@ -24,6 +24,10 @@
 #elif defined(cl_qcom_subgroup_shuffle)
 #pragma OPENCL EXTENSION cl_qcom_subgroup_shuffle : enable
 #define HAS_SUBGROUP_SHUFFLE 1
+// Adreno compilers that expose only cl_qcom_subgroup_shuffle do not declare the KHR
+// name, so calling it is an implicit declaration and the program fails to build.
+// Route it to the qcom builtin.
+#define sub_group_shuffle_xor(val, mask) qcom_sub_group_shuffle_xor((val), (mask), CLK_SUB_GROUP_SHUFFLE_WIDTH_WAVE_SIZE_QCOM, 0.0f)
 #endif
 
 // Assumes row size (ne00) is a multiple of 4

ggml/src/ggml-opencl/kernels/mul_mv_q4_k_f32.cl

@@ -1,3 +1,5 @@
+#pragma OPENCL EXTENSION cl_khr_fp16 : enable
+
 #ifdef cl_intel_required_subgroup_size
 #pragma OPENCL EXTENSION cl_intel_required_subgroup_size : enable
 #define INTEL_GPU 1

ggml/src/ggml-sycl/common.hpp

@@ -64,6 +64,7 @@ extern int g_ggml_sycl_enable_fusion;
 extern int g_ggml_sycl_prioritize_dmmv;
 extern int g_ggml_sycl_enable_flash_attention;
 extern int g_ggml_sycl_dev2dev_memcpy;
+extern int g_ggml_sycl_fa_onednn;
 
 
 #if defined(__clang__) && __has_builtin(__builtin_expect)

ggml/src/ggml-sycl/conv2d-dw.cpp

@@ -71,8 +71,8 @@ struct dw_cwhn_layout {
     }
 };
 
-template <typename Layout>
-static void conv2d_dw_kernel(const float * input, const float * kernel, float * output,
+template <typename KernelT, typename Layout>
+static void conv2d_dw_kernel(const float * input, const KernelT * kernel, float * output,
                              const conv2d_dw_params p, const sycl::nd_item<3> & item_ct1) {
     const int global_idx     = item_ct1.get_local_id(2) +
                                item_ct1.get_group(2) * item_ct1.get_local_range(2);
@@ -93,15 +93,15 @@ static void conv2d_dw_kernel(const float * input, const float * kernel, float *
         for (int kx = bounds.x_min; kx < bounds.x_max; ++kx) {
             const int in_x = dw_calculate_input_coord(out_x, kx, p.stride_x, p.dilation_x, p.padding_x);
             acc += input[Layout::input_index(n, c, in_y, in_x, p)] *
-                   kernel[Layout::kernel_index(c, ky, kx, p)];
+                   static_cast<float>(kernel[Layout::kernel_index(c, ky, kx, p)]);
         }
     }
 
     output[Layout::output_index(n, c, out_y, out_x, p)] = acc;
 }
 
-template <typename Layout>
-static void conv2d_dw_sycl(const float * x_d, const float * w_d, float * y_d,
+template <typename KernelT, typename Layout>
+static void conv2d_dw_sycl(const float * x_d, const KernelT * w_d, float * y_d,
                             const conv2d_dw_params p, const queue_ptr & stream) {
     const int total      = p.batches * p.channels * p.out_h * p.out_w;
     const int num_blocks = (total + SYCL_CONV2D_DW_BLOCK_SIZE - 1) / SYCL_CONV2D_DW_BLOCK_SIZE;
@@ -109,7 +109,7 @@ static void conv2d_dw_sycl(const float * x_d, const float * w_d, float * y_d,
     const sycl::range<3> block_nums(1, 1, num_blocks);
     stream->parallel_for(sycl::nd_range<3>(block_nums * block_dims, block_dims),
         [=](sycl::nd_item<3> item_ct1) {
-            conv2d_dw_kernel<Layout>(x_d, w_d, y_d, p, item_ct1);
+            conv2d_dw_kernel<KernelT, Layout>(x_d, w_d, y_d, p, item_ct1);
         });
 }
 
@@ -119,9 +119,9 @@ void ggml_sycl_op_conv2d_dw(ggml_backend_sycl_context & ctx, ggml_tensor * dst)
     const ggml_tensor * kernel = dst->src[0];
     const ggml_tensor * input  = dst->src[1];
 
-    GGML_ASSERT(kernel->type == GGML_TYPE_F32 && input->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32);
+    GGML_ASSERT((kernel->type == GGML_TYPE_F32 || kernel->type == GGML_TYPE_F16) &&
+                input->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32);
 
-    const float * w_d = (const float *) kernel->data;
     const float * x_d = (const float *) input->data;
     float *       y_d = (float *) dst->data;
 
@@ -148,11 +148,23 @@ void ggml_sycl_op_conv2d_dw(ggml_backend_sycl_context & ctx, ggml_tensor * dst)
 
     const queue_ptr stream = ctx.stream();
 
-    if (ggml_is_contiguous(input)) {
-        conv2d_dw_sycl<dw_whcn_layout>(x_d, w_d, y_d, params, stream);
-    } else if (ggml_is_contiguous_channels(input)) {
-        conv2d_dw_sycl<dw_cwhn_layout>(x_d, w_d, y_d, params, stream);
+    if (kernel->type == GGML_TYPE_F16) {
+        const sycl::half * w_d = (const sycl::half *) kernel->data;
+        if (ggml_is_contiguous(input)) {
+            conv2d_dw_sycl<sycl::half, dw_whcn_layout>(x_d, w_d, y_d, params, stream);
+        } else if (ggml_is_contiguous_channels(input)) {
+            conv2d_dw_sycl<sycl::half, dw_cwhn_layout>(x_d, w_d, y_d, params, stream);
+        } else {
+            GGML_ABORT("Unsupported memory layout for conv2d_dw");
+        }
     } else {
-        GGML_ABORT("Unsupported memory layout for conv2d_dw");
+        const float * w_d = (const float *) kernel->data;
+        if (ggml_is_contiguous(input)) {
+            conv2d_dw_sycl<float, dw_whcn_layout>(x_d, w_d, y_d, params, stream);
+        } else if (ggml_is_contiguous_channels(input)) {
+            conv2d_dw_sycl<float, dw_cwhn_layout>(x_d, w_d, y_d, params, stream);
+        } else {
+            GGML_ABORT("Unsupported memory layout for conv2d_dw");
+        }
     }
 }

ggml/src/ggml-sycl/dequantize.hpp

@@ -19,6 +19,7 @@
 typedef void (*dequantize_kernel_t)(const void * vx, const int64_t ib, const int iqs, dfloat2 & v);
 typedef void (*dequantize_kernel_t_reorder)(const void *d, const int64_t ib, const void *qs,
                                             const int iqs, dfloat2 &v);
+typedef void (*dequantize_kernel_f32_t)(const void * vx, const int64_t ib, const int iqs, float & v0, float & v1);
 
 #if QK_K == 256
 static inline void get_scale_min_k4(int j, const uint8_t * q, uint8_t & d, uint8_t & m);
@@ -85,6 +86,21 @@ static __dpct_inline__ void dequantize_q1_0_reorder(const void *d_ptr, const int
     v.y() = (2 * bit_1 - 1) * d;
 }
 
+static __dpct_inline__ void dequantize_q1_0(const void *vx, const int64_t ib,
+                                            const int iqs, dfloat2 &v) {
+    const block_q1_0 * x = (const block_q1_0 *) vx;
+    const dfloat d = x[ib].d;
+
+    const int bit_index_0 = iqs + 0;
+    const int bit_index_1 = iqs + 1;
+
+    const int bit_0 = (x[ib].qs[bit_index_0 / 8] >> (bit_index_0 % 8)) & 1;
+    const int bit_1 = (x[ib].qs[bit_index_1 / 8] >> (bit_index_1 % 8)) & 1;
+
+    v.x() = (2 * bit_0 - 1) * d;
+    v.y() = (2 * bit_1 - 1) * d;
+}
+
 static __dpct_inline__ void dequantize_q4_1(const void *vx, const int64_t ib,
                                             const int iqs, dfloat2 &v) {
     const block_q4_1 * x = (const block_q4_1 *) vx;
@@ -140,6 +156,39 @@ static __dpct_inline__ void dequantize_q4_K(const void *vx, const int64_t ib,
 #endif
 }
 
+static __dpct_inline__ void dequantize_q4_K_f32(const void *vx, const int64_t ib,
+                                                const int iqs, float &v0, float &v1) {
+#if QK_K == 256
+    const block_q4_K * x = (const block_q4_K *) vx;
+    const sycl::half2 dm = x[ib].dm;
+    const float dall = dm[0];
+    const float dmin = dm[1];
+
+    auto dequantize_one = [&](const int idx) -> float {
+        const int il = idx / 64;
+        const int in = idx % 64;
+        const int is = 2 * il + (in >= 32 ? 1 : 0);
+        const int qsi = 32 * il + (in & 31);
+
+        uint8_t sc;
+        uint8_t m;
+        get_scale_min_k4(is, x[ib].scales, sc, m);
+
+        const float d = dall * sc;
+        const float mn = dmin * m;
+        const uint8_t q = x[ib].qs[qsi];
+        const uint8_t qv = (in >= 32) ? (q >> 4) : (q & 0xF);
+
+        return d * qv - mn;
+    };
+
+    v0 = dequantize_one(iqs + 0);
+    v1 = dequantize_one(iqs + 1);
+#else
+    GGML_ABORT("Q4_K dequantize not supported for QK_K != 256");
+#endif
+}
+
 static __dpct_inline__ void dequantize_q2_K(const void *vx, const int64_t ib,
                                             const int iqs, dfloat2 &v) {
 #if QK_K == 256
@@ -159,7 +208,7 @@ static __dpct_inline__ void dequantize_q2_K(const void *vx, const int64_t ib,
         const float d = dall * (sc & 0xF);
         const float m = dmin * (sc >> 4);
 
-        return sycl::fma((dfloat) ((q >> (2 * g)) & 3), (dfloat) d, (dfloat) (-m));
+        return (dfloat) d * (dfloat) ((q >> (2 * g)) & 3) - (dfloat) m;
     };
 
     v.x() = dequantize_one(iqs + 0);
@@ -169,6 +218,35 @@ static __dpct_inline__ void dequantize_q2_K(const void *vx, const int64_t ib,
 #endif
 }
 
+static __dpct_inline__ void dequantize_q2_K_f32(const void *vx, const int64_t ib,
+                                                const int iqs, float &v0, float &v1) {
+#if QK_K == 256
+    const block_q2_K * x = (const block_q2_K *) vx;
+    const float dall = x[ib].dm[0];
+    const float dmin = x[ib].dm[1];
+
+    auto dequantize_one = [&](const int idx) -> float {
+        const int n = idx / 128;
+        const int r = idx % 128;
+        const int g = r / 32;
+        const int l = r % 32;
+        const int is = 8 * n + l / 16;
+
+        const uint8_t q = x[ib].qs[32 * n + l];
+        const uint8_t sc = x[ib].scales[is + 2 * g];
+        const float d = dall * (sc & 0xF);
+        const float m = dmin * (sc >> 4);
+
+        return d * ((q >> (2 * g)) & 3) - m;
+    };
+
+    v0 = dequantize_one(iqs + 0);
+    v1 = dequantize_one(iqs + 1);
+#else
+    GGML_ABORT("Q2_K dequantize not supported for QK_K != 256");
+#endif
+}
+
 static __dpct_inline__ void dequantize_q3_K(const void *vx, const int64_t ib,
                                             const int iqs, dfloat2 &v) {
 #if QK_K == 256
@@ -242,6 +320,42 @@ static __dpct_inline__ void dequantize_q5_K(const void *vx, const int64_t ib,
 #endif
 }
 
+static __dpct_inline__ void dequantize_q5_K_f32(const void *vx, const int64_t ib,
+                                                const int iqs, float &v0, float &v1) {
+#if QK_K == 256
+    const block_q5_K * x = (const block_q5_K *) vx;
+    const float dall = x[ib].dm[0];
+    const float dmin = x[ib].dm[1];
+
+    auto dequantize_one = [&](const int idx) -> float {
+        const int il = idx / 64;
+        const int in = idx % 64;
+        const int is = 2 * il + (in >= 32 ? 1 : 0);
+        const int ir = (in & 31) / 2;
+        const int iq = in & 1;
+
+        const uint8_t q = x[ib].qs[32 * il + 2 * ir + iq];
+        const uint8_t h = x[ib].qh[2 * ir + iq];
+        const uint8_t qv = (in >= 32) ? (q >> 4) : (q & 0xF);
+
+        uint8_t sc;
+        uint8_t m;
+        get_scale_min_k4(is, x[ib].scales, sc, m);
+
+        const float d = dall * sc;
+        const float mn = dmin * m;
+        const uint8_t hm = 1 << (2 * il + (in >= 32 ? 1 : 0));
+
+        return (qv + ((h & hm) ? 16 : 0)) * d - mn;
+    };
+
+    v0 = dequantize_one(iqs + 0);
+    v1 = dequantize_one(iqs + 1);
+#else
+    GGML_ABORT("Q5_K dequantize not supported for QK_K != 256");
+#endif
+}
+
 static __dpct_inline__ void dequantize_q6_K(const void *vx, const int64_t ib,
                                             const int iqs, dfloat2 &v) {
 #if QK_K == 256
@@ -296,21 +410,6 @@ static __dpct_inline__ void dequantize_mxfp4(const void *vx, const int64_t ib,
     v.y() = d * kvalues_mxfp4[q >> 4] * 0.5f;
 }
 
-static __dpct_inline__ void dequantize_q1_0(const void *vx, const int64_t ib,
-                                            const int iqs, dfloat2 &v) {
-    const block_q1_0 * x = (const block_q1_0 *) vx;
-    const dfloat d = x[ib].d;
-
-    const int bit_index_0 = iqs + 0;
-    const int bit_index_1 = iqs + 1;
-
-    const int bit_0 = (x[ib].qs[bit_index_0 / 8] >> (bit_index_0 % 8)) & 1;
-    const int bit_1 = (x[ib].qs[bit_index_1 / 8] >> (bit_index_1 % 8)) & 1;
-
-    v.x() = (2 * bit_0 - 1) * d;
-    v.y() = (2 * bit_1 - 1) * d;
-}
-
 static __dpct_inline__ void dequantize_nvfp4(const void *vx, const int64_t ib,
                                              const int iqs, dfloat2 &v) {
     const block_nvfp4 & xb = ((const block_nvfp4 *) vx)[ib];

ggml/src/ggml-sycl/element_wise.cpp

@@ -247,6 +247,17 @@ static __dpct_inline__ T op_leaky_relu(T x, float negative_slope) {
     }
 }
 
+template<typename T>
+static __dpct_inline__ T op_xielu(T x, float alpha_n, float alpha_p, float beta, float eps) {
+    const float xi        = static_cast<float>(x);
+    const float gate_pos  = (xi > 0.0f);
+    const float y_pos     = alpha_p * xi * xi + beta * xi;
+    const float min_v_eps = sycl::fmin(xi, eps);
+    const float y_neg     = (sycl::expm1(min_v_eps) - xi) * alpha_n + beta * xi;
+    const float out       = gate_pos * y_pos + (1.0f - gate_pos) * y_neg;
+    return static_cast<T>(out);
+}
+
 template<typename T>
 static __dpct_inline__ T op_sqr(T x) {
     return x * x;
@@ -359,6 +370,13 @@ static void unary_op_leaky_relu_kernel(const T * x, T * dst, const int k, float
     }
 }
 
+template<typename T>
+static void unary_op_xielu_kernel(const T * x, T * dst, const int k, float alpha_n, float alpha_p, float beta, float eps, const sycl::nd_item<1> &item_ct1) {
+    SYCL_GLOBAL_ID_LOOP(k, item_ct1) {
+        dst[i] = op_xielu(x[i], alpha_n, alpha_p, beta, eps);
+    }
+}
+
 template<typename T>
 static void unary_op_sqr_kernel(const T * x, T * dst, const int k, const sycl::nd_item<1> &item_ct1) {
     SYCL_GLOBAL_ID_LOOP(k, item_ct1) {
@@ -836,6 +854,23 @@ static inline void ggml_sycl_op_clamp(ggml_backend_sycl_context & ctx, ggml_tens
         }, min_val, max_val);
 }
 
+static inline void ggml_sycl_op_xielu(ggml_backend_sycl_context & ctx, ggml_tensor * dst) {
+    const float alpha_n = ggml_get_op_params_f32(dst, 1);
+    const float alpha_p = ggml_get_op_params_f32(dst, 2);
+    const float beta    = ggml_get_op_params_f32(dst, 3);
+    const float eps     = ggml_get_op_params_f32(dst, 4);
+    ggml_sycl_detail::dispatch_ggml_sycl_op_unary(ctx, dst,
+        [](const auto* src, auto* dst_ptr, int k_elements, queue_ptr stream, float alpha_n_arg, float alpha_p_arg, float beta_arg, float eps_arg) {
+            const int num_blocks = ceil_div(k_elements, SYCL_RELU_BLOCK_SIZE);
+            stream->parallel_for(
+                sycl::nd_range<1>(sycl::range<1>(num_blocks) * sycl::range<1>(SYCL_RELU_BLOCK_SIZE),
+                                  sycl::range<1>(SYCL_RELU_BLOCK_SIZE)),
+                [=](sycl::nd_item<1> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]] {
+                    unary_op_xielu_kernel(src, dst_ptr, k_elements, alpha_n_arg, alpha_p_arg, beta_arg, eps_arg, item_ct1);
+                });
+        }, alpha_n, alpha_p, beta, eps);
+}
+
 static inline void ggml_sycl_op_floor(ggml_backend_sycl_context & ctx, ggml_tensor * dst) {
     ggml_sycl_detail::ggml_sycl_op_unary(ctx, dst, [](auto x) {
         return op_floor(x);
@@ -1153,6 +1188,11 @@ void ggml_sycl_clamp(ggml_backend_sycl_context & ctx, ggml_tensor * dst) {
     ggml_sycl_op_clamp(ctx, dst);
 }
 
+void ggml_sycl_xielu(ggml_backend_sycl_context & ctx, ggml_tensor * dst) {
+    scope_op_debug_print scope_dbg_print(__func__, dst, /*num_src=*/1);
+    ggml_sycl_op_xielu(ctx, dst);
+}
+
 void ggml_sycl_sgn(ggml_backend_sycl_context & ctx, ggml_tensor * dst) {
     scope_op_debug_print scope_dbg_print(__func__, dst, /*num_src=*/1);
     ggml_sycl_op_sgn(ctx, dst);

ggml/src/ggml-sycl/element_wise.hpp

@@ -75,6 +75,8 @@ void ggml_sycl_sqr(ggml_backend_sycl_context & ctx, ggml_tensor * dst);
 
 void ggml_sycl_clamp(ggml_backend_sycl_context & ctx, ggml_tensor * dst);
 
+void ggml_sycl_xielu(ggml_backend_sycl_context & ctx, ggml_tensor * dst);
+
 void ggml_sycl_sgn(ggml_backend_sycl_context & ctx, ggml_tensor * dst);
 
 void ggml_sycl_abs(ggml_backend_sycl_context & ctx, ggml_tensor * dst);

ggml/src/ggml-sycl/fattn-onednn.cpp

@@ -0,0 +1,265 @@
+#include <cstdint>
+#include <cstdio>
+#include <cstring>
+#include <string>
+#include <unordered_map>
+#include <vector>
+
+#include "fattn-onednn.hpp"
+#include "fattn-tile.hpp"
+
+// set minimum query length to treat as prefill (32)
+#define GGML_SYCL_FA_ONEDNN_MIN_Q 32
+
+bool ggml_sycl_flash_attn_ext_onednn_supported(const ggml_tensor * dst) {
+#if !GGML_SYCL_DNNL
+    GGML_UNUSED(dst);
+    return false;
+#else
+    if (!g_ggml_sycl_fa_onednn) {
+        return false;
+    }
+    // Battlemage (Xe2) only, for now. On other Intel archs oneDNN's fused SDPA returns wrong results
+    // for some shapes (e.g. head_dim=64 on Arc / xe_hpg) -- an oneDNN bug tracked upstream at
+    // https://github.com/uxlfoundation/oneDNN/issues/5510. Remove this hardware limitation once that
+    // is fixed; until then non-BMG archs fall back to the existing FA kernel.
+    const gpu_arch arch = ggml_sycl_info().devices[ggml_sycl_get_device()].hw_info.arch;
+    if (arch != gpu_arch::intel_gpu_bmg_g21 && arch != gpu_arch::intel_gpu_bmg_g31) {
+        return false;
+    }
+    const ggml_tensor * Q     = dst->src[0];
+    const ggml_tensor * K     = dst->src[1];
+    const ggml_tensor * V     = dst->src[2];
+    const ggml_tensor * mask  = dst->src[3];
+    const ggml_tensor * sinks = dst->src[4];
+
+    // gate for f16 KV only for now
+    // need to implement quantized KV
+    if (K->type != GGML_TYPE_F16 || V->type != GGML_TYPE_F16) {
+        return false;
+    }
+    // gate for the following cases
+    // 1. if the oneDNN graph Add node has no input --> skip
+    // 2. types other than f16 need different logical_tensor declaration
+    // 3. the mask must be shape [1, 1, q, seq]
+    // 4. sinks: excludes attention sink (Xiao et al., 2024) that can't be modeled by oneDNN graph
+    if (!mask || mask->type != GGML_TYPE_F16 || mask->ne[2] != 1 || mask->ne[3] != 1 || sinks) {
+        return false;
+    }
+    float max_bias = 0.0f, logit_softcap = 0.0f;
+    memcpy(&max_bias,      (const float *) dst->op_params + 1, sizeof(float));
+    memcpy(&logit_softcap, (const float *) dst->op_params + 2, sizeof(float));
+    if (max_bias != 0.0f || logit_softcap != 0.0f) {
+        return false;
+    }
+    // K and V must share head_dim: the SDPA graph uses a single `d` for both.
+    const int64_t d = K->ne[0];
+    if (V->ne[0] != d || Q->ne[3] != 1) {
+        return false;
+    }
+    // GQA must divide evenly.
+    if (K->ne[2] == 0 || Q->ne[2] % K->ne[2] != 0) {
+        return false;
+    }
+    // Prefill only.
+    if (Q->ne[1] < GGML_SYCL_FA_ONEDNN_MIN_Q) {
+        return false;
+    }
+    return true;
+#endif
+}
+
+#if GGML_SYCL_DNNL
+
+#include "dnnl.hpp"
+#include "dnnl_sycl.hpp"
+#include "oneapi/dnnl/dnnl_graph.hpp"   // graph API lives only under oneapi/dnnl/, not at the include root
+
+using namespace dnnl;
+using namespace dnnl::graph;
+
+// strided src (f16 or f32) -> contiguous f16 [ne0,ne1,ne2,ne3] (ne0 innermost). nb* are BYTE strides.
+template <typename src_t>
+static void cont_to_f16_sycl(const char * src, sycl::half * dst,
+        int64_t ne0, int64_t ne1, int64_t ne2, int64_t ne3,
+        size_t nb1, size_t nb2, size_t nb3, dpct::queue_ptr stream) {
+    const int64_t n = ne0 * ne1 * ne2 * ne3;
+    stream->parallel_for(sycl::range<1>(n), [=](sycl::id<1> ix) {
+        const int64_t gid = ix[0];
+        int64_t       i   = gid;
+        const int64_t i0 = i % ne0; i /= ne0;
+        const int64_t i1 = i % ne1; i /= ne1;
+        const int64_t i2 = i % ne2; const int64_t i3 = i / ne2;
+        const src_t * p = (const src_t *) (src + i1 * nb1 + i2 * nb2 + i3 * nb3) + i0;
+        dst[gid] = (sycl::half) (*p);
+    });
+}
+
+// oneDNN SDPA out (f16 contiguous [mb,H,q,d]) -> ggml dst (f32 [head_dim,H,n_tok,mb], contiguous).
+static void permute_sdpa_out_sycl(const sycl::half * out, float * dst,
+        int64_t mb, int64_t H, int64_t q, int64_t d, dpct::queue_ptr stream) {
+    const int64_t n = mb * H * q * d;
+    stream->parallel_for(sycl::range<1>(n), [=](sycl::id<1> ix) {
+        const int64_t gid = ix[0];
+        int64_t       i   = gid;
+        const int64_t e = i % d; i /= d;
+        const int64_t t = i % q; i /= q;
+        const int64_t h = i % H; const int64_t b = i / H;
+        dst[e + h * d + t * d * H + b * d * H * q] = (float) out[gid];
+    });
+}
+
+struct sdpa_partition {
+    compiled_partition          cp;
+    std::vector<logical_tensor> ins;
+    logical_tensor              out;
+    size_t id_q = 0, id_k = 0, id_v = 0, id_scale = 0, id_mask = 0;
+    bool   ok = false;
+};
+
+// Build + compile the contiguous-input GQA SDPA graph (MatMul->Divide->Add->SoftMax->MatMul), f32 out.
+// Mirrors the hardware-verified scratch/onednn_sdpa_probe.cpp build_gqa (partitions=1, sdp_primitive_kernel_t).
+static sdpa_partition build_sdpa(const engine & eng, int H, int Hkv, int q, int seq, int d) {
+    using ltype = logical_tensor::layout_type;
+    using dt    = logical_tensor::data_type;
+    using ldims = logical_tensor::dims;
+    const dt    fi = dt::f32, t = dt::f16;
+    const int   rep = H / Hkv;
+    const ldims q_sz = {1, Hkv, rep, q, d}, kv_sz = {1, Hkv, 1, seq, d}, s_sz = {1, Hkv, rep, q, seq},
+                sc = {1, 1, 1, 1, 1}, msk = {1, 1, 1, q, seq}, o_sz = {1, Hkv, rep, q, d};
+    int64_t        id = 0;
+    sdpa_partition E;
+
+    auto query  = logical_tensor(id++, t,  q_sz, ltype::strided);
+    auto key    = logical_tensor(id++, t,  kv_sz, ltype::strided);
+    auto score  = logical_tensor(id++, fi, s_sz, ltype::strided);
+    auto bmm1   = op(id++, op::kind::MatMul, "bmm1");
+    bmm1.set_attr<bool>(op::attr::transpose_b, true);          // key is [.., seq, d]
+    bmm1.add_inputs({query, key}); bmm1.add_outputs({score});
+
+    auto scale  = logical_tensor(id++, t,  sc,   ltype::strided);
+    auto scaled = logical_tensor(id++, fi, s_sz, ltype::strided);
+    auto sdiv   = op(id++, op::kind::Divide, "scale_div");     // score / (1/kq_scale) == score * kq_scale
+    sdiv.add_inputs({score, scale}); sdiv.add_outputs({scaled});
+
+    auto mask   = logical_tensor(id++, t,  msk,  ltype::strided);
+    auto masked = logical_tensor(id++, fi, s_sz, ltype::strided);
+    auto madd   = op(id++, op::kind::Add, "mask_add");
+    madd.add_inputs({scaled, mask}); madd.add_outputs({masked});
+
+    auto probs  = logical_tensor(id++, t,  s_sz, ltype::strided);
+    auto smax   = op(id++, op::kind::SoftMax, "softmax");
+    smax.set_attr<int64_t>(op::attr::axis, -1);
+    smax.set_attr<std::string>(op::attr::mode, "inf_as_zero");
+    smax.add_inputs({masked}); smax.add_outputs({probs});
+
+    auto value  = logical_tensor(id++, t,  kv_sz, ltype::strided);
+    // f16 output is REQUIRED to hit sdp_primitive_kernel_t (the systolic micro-kernel); an f32 output
+    // falls to larger_partition_kernel_t which materializes N^2 (confirmed: scratch/onednn_sdpa_kernel_probe.cpp).
+    // converted to the f32 ggml dst in the permute below.
+    auto output = logical_tensor(id++, t,  o_sz, ltype::strided);   // f16 contiguous [mb,Hkv,rep,q,d]
+    auto bmm2   = op(id++, op::kind::MatMul, "bmm2");
+    bmm2.add_inputs({probs, value}); bmm2.add_outputs({output});
+
+    dnnl::graph::graph g(eng.get_kind());
+    g.add_op(bmm1); g.add_op(sdiv); g.add_op(madd); g.add_op(smax); g.add_op(bmm2);
+    g.finalize();
+
+    auto parts = g.get_partitions();
+    if (parts.size() != 1 || !parts[0].is_supported()) {
+        return E;   // ok stays false -> caller falls back to TILE
+    }
+    E.ins      = parts[0].get_input_ports();
+    E.out      = parts[0].get_output_ports()[0];
+    E.cp       = parts[0].compile(E.ins, {E.out}, eng);
+    E.out      = E.cp.query_logical_tensor(E.out.get_id());
+    E.id_q     = query.get_id(); E.id_k = key.get_id(); E.id_v = value.get_id();
+    E.id_scale = scale.get_id(); E.id_mask = mask.get_id();
+    E.ok       = true;
+    return E;
+}
+
+void ggml_sycl_flash_attn_ext_onednn(ggml_backend_sycl_context & ctx, ggml_tensor * dst) try {
+    const ggml_tensor * Q    = dst->src[0];
+    const ggml_tensor * K    = dst->src[1];
+    const ggml_tensor * V    = dst->src[2];
+    const ggml_tensor * mask = dst->src[3];
+
+    const int64_t d   = K->ne[0];   // head_dim
+    const int64_t seq = K->ne[1];   // n_kv
+    const int64_t Hkv = K->ne[2];   // n_head_kv
+    const int64_t H   = Q->ne[2];   // n_head
+    const int64_t q   = Q->ne[1];   // n_tok
+    const int64_t mb  = Q->ne[3];   // batch (== 1, gated)
+
+    float kq_scale = 1.0f;
+    memcpy(&kq_scale, (const float *) dst->op_params + 0, sizeof(float));
+
+    dpct::queue_ptr stream = ctx.stream();
+    dnnl::engine    eng    = ctx.engine_dnnl(stream);
+    dnnl::stream    strm   = ctx.stream_dnnl(stream);
+
+    // cont/cast inputs to contiguous f16 (head-major) -- the layout the fast systolic path wants.
+    ggml_sycl_pool_alloc<sycl::half> Qf(ctx.pool(), (size_t) H   * q   * d);
+    ggml_sycl_pool_alloc<sycl::half> Kf(ctx.pool(), (size_t) Hkv * seq * d);
+    ggml_sycl_pool_alloc<sycl::half> Vf(ctx.pool(), (size_t) Hkv * seq * d);
+    cont_to_f16_sycl<float>     ((const char *) Q->data, Qf.get(), d, q,   H,   mb, Q->nb[1], Q->nb[2], Q->nb[3], stream);
+    cont_to_f16_sycl<sycl::half>((const char *) K->data, Kf.get(), d, seq, Hkv, mb, K->nb[1], K->nb[2], K->nb[3], stream);
+    cont_to_f16_sycl<sycl::half>((const char *) V->data, Vf.get(), d, seq, Hkv, mb, V->nb[1], V->nb[2], V->nb[3], stream);
+
+    // divide-by-(1/scale) reproduces ggml's score *= kq_scale on the proven probe graph.
+    const sycl::half scale_h = (sycl::half) (1.0f / kq_scale);
+    ggml_sycl_pool_alloc<sycl::half> scbuf(ctx.pool(), 1);
+    stream->memcpy(scbuf.get(), &scale_h, sizeof(sycl::half));
+
+    ggml_sycl_pool_alloc<sycl::half> outf(ctx.pool(), (size_t) H * q * d);   // f16 contiguous SDPA out [mb,H,q,d]
+
+    // compile once per (device, shape), reuse across layers/calls.
+    static std::unordered_map<std::string, sdpa_partition> cache;
+    char keyb[96];
+    snprintf(keyb, sizeof(keyb), "%d:%lld:%lld:%lld:%lld:%lld", ggml_sycl_get_device(),
+             (long long) H, (long long) Hkv, (long long) q, (long long) seq, (long long) d);
+    auto it = cache.find(keyb);
+    if (it == cache.end()) {
+        it = cache.emplace(keyb, build_sdpa(eng, (int) H, (int) Hkv, (int) q, (int) seq, (int) d)).first;
+    }
+    sdpa_partition & E = it->second;
+    // _supported() is authoritative: if it accepted this op the partition must build.
+    // A failure here is a gap in _supported() -- surface it, don't mask it with a fallback.
+    GGML_ASSERT(E.ok && "oneDNN SDPA partition failed to build for a _supported() shape");
+
+    auto id2ptr = [&](size_t r) -> void * {
+        if (r == E.id_q)     return Qf.get();
+        if (r == E.id_k)     return Kf.get();
+        if (r == E.id_v)     return Vf.get();
+        if (r == E.id_scale) return scbuf.get();
+        if (r == E.id_mask)  return (void *) mask->data;
+        return nullptr;
+    };
+    std::vector<tensor> ti;
+    ti.reserve(E.ins.size());
+    for (auto & lt : E.ins) {
+        ti.emplace_back(lt, eng, id2ptr(lt.get_id()));
+    }
+    tensor to(E.out, eng, outf.get());
+    E.cp.execute(strm, ti, {to});
+
+    permute_sdpa_out_sycl(outf.get(), (float *) dst->data, mb, H, q, d, stream);
+    // Single device: no sync is required, and actually PP perf is ~6% > wait_and_throw() (tested on llama-3.1-8b & qwen3.6-27b, both Q8_0, with Arc B70).
+    // Any future multi-GPU refactor MUST re-measure this single-device path and keep the best
+    // single-device PP speed. Otherwise (multiple devices/streams can race the reuse):
+    if (ggml_sycl_info().device_count > 1) {
+        // cont_to_f16 -> oneDNN execute -> permute is async on this stream, but the
+        // pool_alloc*s above free their device buffers at host return. Without this wait the next
+        // scheduler op re-acquires those bytes while the GPU is still computing the SDPA, turning
+        // it into garbage and collapsing multi-turn output to a single repeated token ("GGGGG...").
+        stream->wait_and_throw();
+    }
+}
+catch (const std::exception & e) {
+    // any oneDNN/SYCL failure is non-fatal: fall back to the existing kernel (strictly additive).
+    GGML_LOG_WARN("%s: oneDNN SDPA failed (%s); falling back to TILE kernel\n", __func__, e.what());
+    ggml_sycl_flash_attn_ext_tile(ctx, dst);
+}
+
+#endif // GGML_SYCL_DNNL

ggml/src/ggml-sycl/fattn-onednn.hpp

@@ -0,0 +1,14 @@
+#ifndef GGML_SYCL_FATTN_ONEDNN_HPP
+#define GGML_SYCL_FATTN_ONEDNN_HPP
+
+#include "common.hpp"
+
+// Static-only check: fused-XMX oneDNN Graph SDPA path==flash-attn op
+// (f16 KV, no softcap/ALiBi, single stream, tuned head_dim, prefill-sized q.)
+bool ggml_sycl_flash_attn_ext_onednn_supported(const ggml_tensor * dst);
+
+// Run flash attention through oneDNN's fused xmx SDPA
+// execute the cached SDPA partition, write the f32 dst. Falls back to the TILE kernel on any failure.
+void ggml_sycl_flash_attn_ext_onednn(ggml_backend_sycl_context & ctx, ggml_tensor * dst);
+
+#endif // GGML_SYCL_FATTN_ONEDNN_HPP

ggml/src/ggml-sycl/fattn.cpp

@@ -18,6 +18,7 @@
 #include "fattn-tile.hpp"
 #include "fattn-vec.hpp"
 #include "fattn.hpp"
+#include "fattn-onednn.hpp"
 
 
 #define FATTN_VEC_CASE(D, type_K, type_V)                                                                        \
@@ -96,6 +97,7 @@ static void ggml_sycl_flash_attn_ext_vec(ggml_backend_sycl_context & ctx, ggml_t
 enum best_fattn_kernel {
     BEST_FATTN_KERNEL_NONE     =   0,
     BEST_FATTN_KERNEL_VEC      = 100,
+    BEST_FATTN_KERNEL_ONEDNN   = 150, // added enum for onednn==150
     BEST_FATTN_KERNEL_TILE     = 200,
 };
 
@@ -189,7 +191,11 @@ static best_fattn_kernel ggml_sycl_get_best_fattn_kernel(const int device, const
     // For small batch sizes the vector kernel may be preferable over the kernels optimized for large batch sizes:
     const bool can_use_vector_kernel = Q->ne[0] <= 512 && Q->ne[0] % 64 == 0 && K->ne[1] % FATTN_KQ_STRIDE == 0;
 
-    // Todo: Use the XMX kernel if possible:
+    // Fused-XMX path: oneDNN Graph SDPA (flash attention). Strictly
+    // additive -- taken only when statically supported, otherwise falls through to VEC/TILE below.
+    if (ggml_sycl_flash_attn_ext_onednn_supported(dst)) {
+        return BEST_FATTN_KERNEL_ONEDNN;
+    }
 
     // If there are no tensor cores available, use the generic tile kernel:
     if (can_use_vector_kernel) {
@@ -213,6 +219,13 @@ void ggml_sycl_flash_attn_ext(ggml_backend_sycl_context & ctx, ggml_tensor * dst
     switch (ggml_sycl_get_best_fattn_kernel(ggml_sycl_get_device(), dst)) {
         case BEST_FATTN_KERNEL_NONE:
             GGML_ABORT("Not support Flash-Attention");
+        case BEST_FATTN_KERNEL_ONEDNN:
+            // guarded: ggml_sycl_flash_attn_ext_onednn() is only defined under GGML_SYCL_DNNL;
+            // the reference must be compiled out here or the GGML_SYCL_DNNL=0 build fails to link.
+#if GGML_SYCL_DNNL
+            ggml_sycl_flash_attn_ext_onednn(ctx, dst);
+#endif
+            break;
         case BEST_FATTN_KERNEL_TILE:
             ggml_sycl_flash_attn_ext_tile(ctx, dst);
             break;

ggml/src/ggml-sycl/getrows.cpp

@@ -60,6 +60,50 @@ static void k_get_rows(
     dst_row[iybs + iqs + y_offset] = v.y();
 }
 
+template<int qk, int qr, dequantize_kernel_f32_t dequantize_kernel, typename dst_t>
+static void k_get_rows_f32(
+            const void * src0, const int32_t * src1, dst_t * dst,
+            int64_t ne00,
+            int64_t ne12,
+            size_t s1, size_t s2, size_t s3,
+            size_t nb01, size_t nb02, size_t nb03,
+            size_t s10, size_t s11, size_t s12,
+            const sycl::nd_item<3> &item_ct1) {
+
+    const int i00 = (item_ct1.get_group(2) * item_ct1.get_local_range(2) +
+                     item_ct1.get_local_id(2)) *
+                    2;
+    const int i10 = item_ct1.get_local_range(1) * item_ct1.get_group(1) +
+                    item_ct1.get_local_id(1);
+    const int i11 = (item_ct1.get_group(0) * item_ct1.get_local_range(0) +
+                     item_ct1.get_local_id(0)) /
+                    ne12;
+    const int i12 = (item_ct1.get_group(0) * item_ct1.get_local_range(0) +
+                     item_ct1.get_local_id(0)) %
+                    ne12;
+
+    if (i00 >= ne00) {
+        return;
+    }
+
+    const int i01 = src1[i10*s10 + i11*s11 + i12*s12];
+
+    dst_t * dst_row = dst + i10*s1 + i11*s2 + i12*s3;
+    const void * src0_row = (const char *)src0 + i01*nb01 + i11*nb02 + i12*nb03;
+
+    const int ib = i00/qk;
+    const int iqs = (i00%qk)/qr;
+    const int iybs = i00 - i00%qk;
+    const int y_offset = qr == 1 ? 1 : qk/2;
+
+    float v0;
+    float v1;
+    dequantize_kernel(src0_row, ib, iqs, v0, v1);
+
+    dst_row[iybs + iqs + 0] = (dst_t) v0;
+    dst_row[iybs + iqs + y_offset] = (dst_t) v1;
+}
+
 template<typename src0_t, typename dst_t>
 static void k_get_rows_float(
             const src0_t * src0, const int32_t * src1, dst_t * dst,
@@ -129,6 +173,39 @@ static void get_rows_sycl(ggml_backend_sycl_context & ctx, const ggml_tensor *sr
     GGML_UNUSED(ctx);
 }
 
+template <int qk, int qr, dequantize_kernel_f32_t dq>
+static void get_rows_sycl_f32(ggml_backend_sycl_context & ctx, const ggml_tensor *src0, const ggml_tensor *src1,
+                              ggml_tensor *dst, const void *src0_dd,
+                              const int32_t *src1_dd, float *dst_dd,
+                              queue_ptr stream) {
+
+    GGML_TENSOR_BINARY_OP_LOCALS
+
+    const sycl::range<3> block_dims(1, 1, SYCL_GET_ROWS_BLOCK_SIZE);
+    const int block_num_x = (ne00 + 2*SYCL_GET_ROWS_BLOCK_SIZE - 1) / (2*SYCL_GET_ROWS_BLOCK_SIZE);
+    const sycl::range<3> block_nums(ne11 * ne12, ne10, block_num_x);
+
+    const size_t s1 = nb1 / ggml_element_size(dst);
+    const size_t s2 = nb2 / ggml_element_size(dst);
+    const size_t s3 = nb3 / ggml_element_size(dst);
+
+    const size_t s10 = nb10 / ggml_element_size(src1);
+    const size_t s11 = nb11 / ggml_element_size(src1);
+    const size_t s12 = nb12 / ggml_element_size(src1);
+
+    GGML_ASSERT(ne00 % 2 == 0);
+
+    stream->parallel_for(sycl::nd_range<3>(block_nums * block_dims, block_dims),
+                         [=](sycl::nd_item<3> item_ct1) {
+                             k_get_rows_f32<qk, qr, dq>(
+                                 src0_dd, src1_dd, dst_dd, ne00, ne12, s1, s2,
+                                 s3, nb01, nb02, nb03, s10, s11, s12, item_ct1);
+                         });
+
+    GGML_UNUSED(dst);
+    GGML_UNUSED(ctx);
+}
+
 template <typename src0_t, typename dst_t>
 static void get_rows_sycl_float(ggml_backend_sycl_context & ctx, const ggml_tensor *src0,
                                 const ggml_tensor *src1, ggml_tensor *dst,
@@ -244,7 +321,7 @@ void ggml_sycl_op_get_rows(ggml_backend_sycl_context & ctx, ggml_tensor * dst) {
             src1_i32, (float *)dst->data, ctx.stream());
             break;
         case GGML_TYPE_Q2_K:
-            get_rows_sycl<QK_K, 1, dequantize_q2_K>(ctx, dst->src[0], dst->src[1], dst, (const float *)dst->src[0]->data,
+            get_rows_sycl_f32<QK_K, 1, dequantize_q2_K_f32>(ctx, dst->src[0], dst->src[1], dst, (const float *)dst->src[0]->data,
             src1_i32, (float *)dst->data, ctx.stream());
             break;
         case GGML_TYPE_Q3_K:
@@ -260,7 +337,7 @@ void ggml_sycl_op_get_rows(ggml_backend_sycl_context & ctx, ggml_tensor * dst) {
             src1_i32, (float *)dst->data, ctx.stream());
             break;
         case GGML_TYPE_Q4_K:
-            get_rows_sycl<QK_K, 1, dequantize_q4_K>(ctx, dst->src[0], dst->src[1], dst, (const float *)dst->src[0]->data,
+            get_rows_sycl_f32<QK_K, 1, dequantize_q4_K_f32>(ctx, dst->src[0], dst->src[1], dst, (const float *)dst->src[0]->data,
             src1_i32, (float *)dst->data, ctx.stream());
             break;
         case GGML_TYPE_Q5_0:
@@ -272,7 +349,7 @@ void ggml_sycl_op_get_rows(ggml_backend_sycl_context & ctx, ggml_tensor * dst) {
             src1_i32, (float *)dst->data, ctx.stream());
             break;
         case GGML_TYPE_Q5_K:
-            get_rows_sycl<QK_K, 1, dequantize_q5_K>(ctx, dst->src[0], dst->src[1], dst, (const float *)dst->src[0]->data,
+            get_rows_sycl_f32<QK_K, 1, dequantize_q5_K_f32>(ctx, dst->src[0], dst->src[1], dst, (const float *)dst->src[0]->data,
             src1_i32, (float *)dst->data, ctx.stream());
             break;
         case GGML_TYPE_Q6_K:

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

@@ -84,6 +84,7 @@ int g_ggml_sycl_debug = 0;
 int g_ggml_sycl_enable_optimize = 1;
 int g_ggml_sycl_enable_graph = 0;
 int g_ggml_sycl_enable_dnn = 1;
+int g_ggml_sycl_fa_onednn = 1;
 int g_ggml_sycl_enable_vmm = 1;
 int g_ggml_sycl_enable_fusion = 1;
 int g_ggml_sycl_prioritize_dmmv = 0;
@@ -285,6 +286,7 @@ static void ggml_check_sycl() try {
         g_ggml_sycl_enable_optimize = ggml_sycl_get_env("GGML_SYCL_ENABLE_OPT", 1);
         g_ggml_sycl_enable_graph = ggml_sycl_get_env("GGML_SYCL_ENABLE_GRAPH", 0);
         g_ggml_sycl_enable_dnn = ggml_sycl_get_env("GGML_SYCL_ENABLE_DNN", 1);
+        g_ggml_sycl_fa_onednn = ggml_sycl_get_env("GGML_SYCL_FA_ONEDNN", 1);
         g_ggml_sycl_enable_vmm = ggml_sycl_get_env("GGML_SYCL_ENABLE_VMM", 1);
         g_ggml_sycl_enable_fusion = ggml_sycl_get_env("GGML_SYCL_ENABLE_FUSION", 1);
         g_ggml_sycl_prioritize_dmmv = ggml_sycl_get_env("GGML_SYCL_PRIORITIZE_DMMV", 0);
@@ -352,8 +354,10 @@ static void ggml_check_sycl() try {
 
 #if defined(GGML_SYCL_DNNL)
         GGML_LOG_INFO("  GGML_SYCL_ENABLE_DNN: %d\n", g_ggml_sycl_enable_dnn);
+        GGML_LOG_INFO("  GGML_SYCL_FA_ONEDNN: %d\n", g_ggml_sycl_fa_onednn);
 #else
         GGML_LOG_INFO("  GGML_SYCL_ENABLE_DNN: DNN disabled by compile flag\n");
+        GGML_LOG_INFO("  GGML_SYCL_FA_ONEDNN: %d\n", g_ggml_sycl_fa_onednn);
 #endif
 #ifdef SYCL_FLASH_ATTN
         GGML_LOG_INFO("  GGML_SYCL_ENABLE_FLASH_ATTN: %d\n", g_ggml_sycl_enable_flash_attention);
@@ -839,7 +843,7 @@ static const char * ggml_backend_sycl_buffer_type_get_name(ggml_backend_buffer_t
 }
 
 static bool check_usm_system(int device, size_t size) {
-    bool use_usm_system = g_ggml_sycl_usm_system && size >= MEM_SIZE_1G;
+    bool use_usm_system = g_ggml_sycl_usm_system && size >= ((size_t)4 * MEM_SIZE_1G);
 
     if (use_usm_system && !ggml_sycl_info().devices[device].usm_system_support) {
         GGML_LOG_INFO("Device does not support USM system allocations\n");
@@ -878,6 +882,7 @@ ggml_backend_sycl_buffer_type_alloc_buffer(ggml_backend_buffer_type_t buft,
 
     void * dev_ptr;
     if (use_usm_system) {
+        GGML_SYCL_DEBUG("[SYCL] allocating %lu Bytes with USM system\n", size);
         dev_ptr = (void *)aligned_malloc_host(alignment, aligned_size);
         if (!dev_ptr) {
             GGML_LOG_ERROR("%s: can't allocate %lu Bytes of memory on host\n", __func__, size);
@@ -5006,6 +5011,9 @@ static bool ggml_sycl_compute_forward(ggml_backend_sycl_context & ctx, struct gg
                 case GGML_UNARY_OP_ELU:
                     ggml_sycl_elu(ctx, dst);
                     break;
+                case GGML_UNARY_OP_XIELU:
+                    ggml_sycl_xielu(ctx, dst);
+                    break;
                 case GGML_UNARY_OP_FLOOR:
                     ggml_sycl_floor(ctx, dst);
                     break;
@@ -5668,6 +5676,7 @@ static bool do_ggml_backend_sycl_device_supports_op(ggml_backend_dev_t dev, cons
                 case GGML_UNARY_OP_EXPM1:
                 case GGML_UNARY_OP_SOFTPLUS:
                 case GGML_UNARY_OP_ELU:
+                case GGML_UNARY_OP_XIELU:
                 case GGML_UNARY_OP_CEIL:
                     return true;
                 case GGML_UNARY_OP_FLOOR:

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

@@ -961,6 +961,7 @@ struct vk_device_struct {
     vk_pipeline pipeline_col2im_1d_f32;
     vk_pipeline pipeline_col2im_1d_f16;
     vk_pipeline pipeline_col2im_1d_bf16;
+    vk_pipeline pipeline_out_prod_f32;
     vk_pipeline pipeline_snake_f32;
     vk_pipeline pipeline_snake_f16;
     vk_pipeline pipeline_snake_bf16;
@@ -5479,6 +5480,8 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
     ggml_vk_create_pipeline(device, device->pipeline_col2im_1d_f16,  "col2im_1d_f16",  col2im_1d_f16_len,  col2im_1d_f16_data,  "main", 2, sizeof(vk_op_col2im_1d_push_constants), {256, 1, 1}, {}, 1, true);
     ggml_vk_create_pipeline(device, device->pipeline_col2im_1d_bf16, "col2im_1d_bf16", col2im_1d_bf16_len, col2im_1d_bf16_data, "main", 2, sizeof(vk_op_col2im_1d_push_constants), {256, 1, 1}, {}, 1, true);
 
+    ggml_vk_create_pipeline(device, device->pipeline_out_prod_f32, "out_prod_f32", out_prod_f32_len, out_prod_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {256, 1, 1}, {}, 1);
+
     ggml_vk_create_pipeline(device, device->pipeline_snake_f32,  "snake_f32",  snake_f32_len,  snake_f32_data,  "main", 4, sizeof(vk_op_snake_push_constants), {256, 1, 1}, {}, 1);
     ggml_vk_create_pipeline(device, device->pipeline_snake_f16,  "snake_f16",  snake_f16_len,  snake_f16_data,  "main", 4, sizeof(vk_op_snake_push_constants), {256, 1, 1}, {}, 1);
     ggml_vk_create_pipeline(device, device->pipeline_snake_bf16, "snake_bf16", snake_bf16_len, snake_bf16_data, "main", 4, sizeof(vk_op_snake_push_constants), {256, 1, 1}, {}, 1);
@@ -10745,6 +10748,11 @@ static vk_pipeline ggml_vk_op_get_pipeline(ggml_backend_vk_context * ctx, const
             return ctx->device->pipeline_add_id_f32;
         }
         return nullptr;
+    case GGML_OP_OUT_PROD:
+        if (src0->type == GGML_TYPE_F32 && src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) {
+            return ctx->device->pipeline_out_prod_f32;
+        }
+        return nullptr;
     case GGML_OP_CONCAT: {
         if (src0->type != src1->type || src0->type != dst->type) {
             return nullptr;
@@ -11701,6 +11709,7 @@ static void ggml_vk_op_f32(ggml_backend_vk_context * ctx, vk_context& subctx, co
     case GGML_OP_DIV:
     case GGML_OP_MUL:
     case GGML_OP_ADD1:
+    case GGML_OP_OUT_PROD:
     case GGML_OP_ARANGE:
     case GGML_OP_FILL:
     case GGML_OP_SCALE:
@@ -12014,6 +12023,24 @@ static void ggml_vk_add(ggml_backend_vk_context * ctx, vk_context& subctx, const
     });
 }
 
+static void ggml_vk_out_prod(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) {
+    const uint32_t src0_type_size = ggml_type_size(src0->type);
+    const uint32_t src1_type_size = ggml_type_size(src1->type);
+    const uint32_t dst_type_size = ggml_type_size(dst->type);
+
+    ggml_vk_op_f32<vk_op_binary_push_constants>(ctx, subctx, src0, src1, nullptr, nullptr, dst, GGML_OP_OUT_PROD, {
+        (uint32_t)ggml_nelements(dst),
+        (uint32_t)src0->ne[0], (uint32_t)src0->ne[1], (uint32_t)src0->ne[2],(uint32_t)src0->ne[3],
+        (uint32_t)src0->nb[0] / src0_type_size, (uint32_t)src0->nb[1] / src0_type_size, (uint32_t)src0->nb[2] / src0_type_size, (uint32_t)src0->nb[3] / src0_type_size,
+        (uint32_t)src1->ne[0], (uint32_t)src1->ne[1], (uint32_t)src1->ne[2],(uint32_t)src1->ne[3],
+        (uint32_t)src1->nb[0] / src1_type_size, (uint32_t)src1->nb[1] / src1_type_size, (uint32_t)src1->nb[2] / src1_type_size, (uint32_t)src1->nb[3] / src1_type_size,
+        (uint32_t) dst->ne[0], (uint32_t) dst->ne[1], (uint32_t) dst->ne[2],(uint32_t) dst->ne[3],
+        (uint32_t) dst->nb[0] /  dst_type_size, (uint32_t) dst->nb[1] /  dst_type_size, (uint32_t) dst->nb[2] /  dst_type_size, (uint32_t) dst->nb[3] /  dst_type_size,
+        0,
+        0.0f, 0.0f, 0,
+    });
+}
+
 static void ggml_vk_sub(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) {
     const uint32_t src0_type_size = ggml_type_size(src0->type);
     const uint32_t src1_type_size = ggml_type_size(src1->type);
@@ -14801,6 +14828,9 @@ static bool ggml_vk_build_graph(ggml_backend_vk_context * ctx, ggml_cgraph * cgr
             ggml_vk_add(ctx, compute_ctx, src0, src1, node);
         }
         break;
+    case GGML_OP_OUT_PROD:
+        ggml_vk_out_prod(ctx, compute_ctx, src0, src1, node);
+        break;
     case GGML_OP_SUB:
         ggml_vk_sub(ctx, compute_ctx, src0, src1, node);
 
@@ -17655,6 +17685,10 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm
         case GGML_OP_OPT_STEP_ADAMW:
         case GGML_OP_OPT_STEP_SGD:
             return ggml_is_contiguous(op->src[0]) && op->src[0]->type == GGML_TYPE_F32;
+        case GGML_OP_OUT_PROD:
+            return ggml_is_contiguous(op->src[0]) && op->src[0]->type == GGML_TYPE_F32
+                && ggml_is_contiguous(op->src[1]) && op->src[1]->type == GGML_TYPE_F32
+                && op->type == GGML_TYPE_F32;
         case GGML_OP_LOG:
         case GGML_OP_TRI:
         case GGML_OP_DIAG:

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

@@ -0,0 +1,59 @@
+#version 450
+
+#extension GL_EXT_shader_16bit_storage : require
+
+layout (push_constant) uniform parameter
+{
+    uint ne;
+    uint ne00; uint ne01; uint ne02; uint ne03; uint nb00; uint nb01; uint nb02; uint nb03;
+    uint ne10; uint ne11; uint ne12; uint ne13; uint nb10; uint nb11; uint nb12; uint nb13;
+    uint ne20; uint ne21; uint ne22; uint ne23; uint nb20; uint nb21; uint nb22; uint nb23;
+    uint misalign_offsets;
+    float param1; float param2; int param3;
+} p;
+
+layout (binding = 0) readonly buffer A {float data_a[];};
+layout (binding = 1) readonly buffer B {float data_b[];};
+layout (binding = 2) writeonly buffer D {float data_d[];};
+
+uint get_idx() {
+    return gl_GlobalInvocationID.z * 262144 + gl_GlobalInvocationID.y * 512 + gl_GlobalInvocationID.x;
+}
+
+uint get_aoffset() { return p.misalign_offsets >> 16; }
+uint get_boffset() { return (p.misalign_offsets >> 8) & 0xFF; }
+uint get_doffset() { return p.misalign_offsets & 0xFF; }
+
+layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in;
+
+void main() {
+    uint idx = get_idx();
+    if (idx >= p.ne) {
+        return;
+    }
+
+    uint tmp = idx;
+    uint i0 = tmp % p.ne20; tmp /= p.ne20;
+    uint i1 = tmp % p.ne21; tmp /= p.ne21;
+    uint i2 = tmp % p.ne22; tmp /= p.ne22;
+    uint i3 = tmp;
+
+    uint a_i0 = i0 % p.ne00;
+    uint a_i2 = i2 / (p.ne22 / p.ne02);
+    uint a_i3 = i3 / (p.ne23 / p.ne03);
+
+    uint b_i0 = i1 % p.ne10;
+    uint b_i2 = i2;
+    uint b_i3 = i3;
+
+    float sum = 0.0f;
+    uint K = p.ne01;
+    for (uint k = 0; k < K; k++) {
+        uint aoff = get_aoffset() + a_i3*p.nb03 + a_i2*p.nb02 + k*p.nb01 + a_i0*p.nb00;
+        uint boff = get_boffset() + b_i3*p.nb13 + b_i2*p.nb12 + k*p.nb11 + b_i0*p.nb10;
+        sum += data_a[aoff] * data_b[boff];
+    }
+
+    uint doff = get_doffset() + i3*p.nb23 + i2*p.nb22 + i1*p.nb21 + i0*p.nb20;
+    data_d[doff] = sum;
+}

ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp

@@ -1036,6 +1036,8 @@ void process_shaders() {
         }
     }
 
+    string_to_spv("out_prod_f32", "out_prod.comp", {});
+
     string_to_spv("timestep_embedding_f32", "timestep_embedding.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "float"}}));
 
     string_to_spv("conv_transpose_1d_f32", "conv_transpose_1d.comp", {{"A_TYPE", "float"},  {"B_TYPE", "float"}, {"D_TYPE", "float"}});

src/models/deepseek4.cpp

@@ -435,27 +435,29 @@ ggml_tensor * llama_model_deepseek4::graph::build_overlap_compressed_kv_from_sta
     kv_state    = dsv4_append_zero_row(ctx0, kv_state,    false);
     score_state = dsv4_append_zero_row(ctx0, score_state, true);
 
-    ggml_tensor * prev_idxs = dsv4_view_1d(ctx0, state_read_idxs, ratio*n_blocks, 0);
-    ggml_tensor * cur_idxs  = dsv4_view_1d(ctx0, state_read_idxs, ratio*n_blocks, ratio*n_blocks);
+    const int64_t n_read = ratio*n_blocks;
 
-    ggml_tensor * kv_prev = ggml_get_rows(ctx0, kv_state, prev_idxs);
-    kv_prev = ggml_cont(ctx0, ggml_view_2d(ctx0, kv_prev, n_embd_head, ratio*n_blocks, kv_prev->nb[1], 0));
+    ggml_tensor * kv_rows = ggml_get_rows(ctx0, kv_state, state_read_idxs);
+    ggml_tensor * score_rows = ggml_get_rows(ctx0, score_state, state_read_idxs);
+
+    ggml_tensor * kv_prev = ggml_cont(ctx0,
+            ggml_view_2d(ctx0, kv_rows, n_embd_head, n_read, kv_rows->nb[1], 0));
     kv_prev = ggml_reshape_3d(ctx0, kv_prev, n_embd_head, ratio, n_blocks);
     cb(kv_prev, name, il);
 
-    ggml_tensor * score_prev = ggml_get_rows(ctx0, score_state, prev_idxs);
-    score_prev = ggml_cont(ctx0, ggml_view_2d(ctx0, score_prev, n_embd_head, ratio*n_blocks, score_prev->nb[1], 0));
+    ggml_tensor * score_prev = ggml_cont(ctx0,
+            ggml_view_2d(ctx0, score_rows, n_embd_head, n_read, score_rows->nb[1], 0));
     score_prev = ggml_reshape_3d(ctx0, score_prev, n_embd_head, ratio, n_blocks);
     cb(score_prev, name, il);
 
-    ggml_tensor * kv_cur = ggml_get_rows(ctx0, kv_state, cur_idxs);
-    kv_cur = ggml_cont(ctx0, ggml_view_2d(ctx0, kv_cur, n_embd_head, ratio*n_blocks, kv_cur->nb[1],
-            ggml_row_size(kv_cur->type, n_embd_head)));
+    ggml_tensor * kv_cur = ggml_cont(ctx0,
+            ggml_view_2d(ctx0, kv_rows, n_embd_head, n_read, kv_rows->nb[1],
+                n_read*kv_rows->nb[1] + ggml_row_size(kv_rows->type, n_embd_head)));
     kv_cur = ggml_reshape_3d(ctx0, kv_cur, n_embd_head, ratio, n_blocks);
 
-    ggml_tensor * score_cur = ggml_get_rows(ctx0, score_state, cur_idxs);
-    score_cur = ggml_cont(ctx0, ggml_view_2d(ctx0, score_cur, n_embd_head, ratio*n_blocks, score_cur->nb[1],
-            ggml_row_size(score_cur->type, n_embd_head)));
+    ggml_tensor * score_cur = ggml_cont(ctx0,
+            ggml_view_2d(ctx0, score_rows, n_embd_head, n_read, score_rows->nb[1],
+                n_read*score_rows->nb[1] + ggml_row_size(score_rows->type, n_embd_head)));
     score_cur = ggml_reshape_3d(ctx0, score_cur, n_embd_head, ratio, n_blocks);
 
     ggml_tensor * values = ggml_concat(ctx0, kv_prev, kv_cur, 1);

tests/test-backend-ops.cpp

@@ -5555,7 +5555,7 @@ struct test_concat : public test_case {
     const std::array<int64_t, 4> ne_a;
     const int64_t ne_b_d;
     const int dim;
-    const int v; // view (1 << 0: non-cont a, 1 << 1: non-cont b)
+    const int v; // view (1 << 0: non-cont a (first 3 dim), 1 << 1: non-cont b (first 3 dim), 1 << 2: non-cont a (last 2 dim), 1 << 3: non-cont b (last 2 dim))
 
     std::string vars() override {
         return VARS_TO_STR5(type, ne_a, ne_b_d, dim, v);
@@ -5576,6 +5576,13 @@ struct test_concat : public test_case {
             a = ggml_new_tensor(ctx, type, 4, ne.data());
             ggml_set_name(a, "a");
 
+            a = ggml_view_4d(ctx, a, ne_a[0], ne_a[1], ne_a[2], ne_a[3], a->nb[1], a->nb[2], a->nb[3], 0);
+            ggml_set_name(a, "view_of_a");
+        } else if (v & 4) {
+            auto ne = ne_a; ne[2] *= 2; ne[3] *= 4;
+            a = ggml_new_tensor(ctx, type, 4, ne.data());
+            ggml_set_name(a, "a");
+
             a = ggml_view_4d(ctx, a, ne_a[0], ne_a[1], ne_a[2], ne_a[3], a->nb[1], a->nb[2], a->nb[3], 0);
             ggml_set_name(a, "view_of_a");
         } else {
@@ -5588,6 +5595,13 @@ struct test_concat : public test_case {
             b = ggml_new_tensor(ctx, type, 4, ne.data());
             ggml_set_name(b, "b");
 
+            b = ggml_view_4d(ctx, b, ne_b[0], ne_b[1], ne_b[2], ne_b[3], b->nb[1], b->nb[2], b->nb[3], 0);
+            ggml_set_name(b, "view_of_b");
+        } else if (v & 8) {
+            auto ne = ne_b; ne[2] *= 3; ne[3] *= 2;
+            b = ggml_new_tensor(ctx, type, 4, ne.data());
+            ggml_set_name(b, "b");
+
             b = ggml_view_4d(ctx, b, ne_b[0], ne_b[1], ne_b[2], ne_b[3], b->nb[1], b->nb[2], b->nb[3], 0);
             ggml_set_name(b, "view_of_b");
         } else {
@@ -9089,8 +9103,10 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
     }
 
     for (ggml_type type_a : { GGML_TYPE_Q4_0, GGML_TYPE_Q4_1, GGML_TYPE_Q5_0, GGML_TYPE_Q5_1, GGML_TYPE_Q8_0 }) {
-        for (int dim : { 0, 1, 2, 3, }) {
-            test_cases.emplace_back(new test_concat(type_a, {128, 12, 13, 14}, dim == 0 ? 256 : 7, dim, 0));
+        for (int v : { 0, 4, 8, 12 }) {
+            for (int dim : { 0, 1, 2, 3, }) {
+                test_cases.emplace_back(new test_concat(type_a, {128, 12, 13, 14}, dim == 0 ? 256 : 7, dim, v));
+            }
         }
     }
 

tools/server/server-tools.cpp

@@ -289,7 +289,7 @@ struct server_tool_read_file : server_tool {
             {"function", {
                 {"name", name},
                 {"description", "Read the contents of a file. Optionally specify a 1-based line range. "
-                                "If append_loc is true, each line is prefixed with its line number (e.g. \"1\u2192 ...\")."},
+                                "If append_loc is true, each line is prefixed with its line number (e.g. \"1\u2192...\")."},
                 {"parameters", {
                     {"type", "object"},
                     {"properties", {
@@ -339,7 +339,7 @@ struct server_tool_read_file : server_tool {
 
             std::string out_line;
             if (append_loc) {
-                out_line = std::to_string(lineno) + "\u2192 " + line + "\n";
+                out_line = std::to_string(lineno) + "\u2192" + line + "\n";
             } else {
                 out_line = line + "\n";
             }

tools/server/tests/unit/test_speculative.py

@@ -12,8 +12,9 @@ def create_server():
     server = ServerPreset.stories15m_moe()
     # set default values
     server.model_draft = download_file(MODEL_DRAFT_FILE_URL)
-    server.draft_min = 4
-    server.draft_max = 8
+    server.spec_type = "draft-simple"
+    server.spec_draft_n_min = 4
+    server.spec_draft_n_max = 8
     server.fa = "off"
 
 
@@ -25,6 +26,7 @@ def fixture_create_server():
 def test_with_and_without_draft():
     global server
     server.model_draft = None  # disable draft model
+    server.spec_type = None
     server.start()
     res = server.make_request("POST", "/completion", data={
         "prompt": "I believe the meaning of life is",
@@ -46,6 +48,7 @@ def test_with_and_without_draft():
         "n_predict": 16,
     })
     assert res.status_code == 200
+    assert res.body["timings"]["draft_n"] > 0
     content_draft = res.body["content"]
 
     assert content_no_draft == content_draft
@@ -63,8 +66,8 @@ def test_different_draft_min_draft_max():
     last_content = None
     for draft_min, draft_max in test_values:
         server.stop()
-        server.draft_min = draft_min
-        server.draft_max = draft_max
+        server.spec_draft_n_min = draft_min
+        server.spec_draft_n_max = draft_max
         server.start()
         res = server.make_request("POST", "/completion", data={
             "prompt": "I believe the meaning of life is",

tools/server/tests/utils.py

@@ -95,6 +95,7 @@ class ServerProcess:
     no_models_autoload: bool | None = None
     lora_files: List[str] | None = None
     enable_ctx_shift: int | None = False
+    spec_type: str | None = None
     spec_draft_n_min: int | None = None
     spec_draft_n_max: int | None = None
     no_ui: bool | None = None
@@ -226,6 +227,8 @@ class ServerProcess:
                 server_args.extend(["--lora", lora_file])
         if self.enable_ctx_shift:
             server_args.append("--context-shift")
+        if self.spec_type:
+            server_args.extend(["--spec-type", self.spec_type])
         if self.api_key:
             server_args.extend(["--api-key", self.api_key])
         if self.spec_draft_n_max:

tools/tokenize/tokenize.cpp

@@ -103,19 +103,11 @@ int main(int argc, char ** argv) {
         return 1;
     }
 
-    // which prompt source was requested?
-    // -p/--prompt and -f/--file both end up in params.prompt (common's -f also
-    // strips a single trailing newline), but -f additionally records the path
-    // in params.prompt_file, so we use that to tell them apart.
+    // -f and -p both land in params.prompt; -f also sets prompt_file. -f and -p
+    // resolve like the other tools (no mutual exclusion), --stdin takes precedence.
     const bool use_stdin = params.tokenize_stdin;
     const bool use_file  = !params.prompt_file.empty();
 
-    // sanity check: --stdin is mutually exclusive with -f/--file and -p/--prompt
-    if (use_stdin && (use_file || !params.prompt.empty())) {
-        LOG_ERR("error: --stdin is mutually exclusive with --file and --prompt\n");
-        return 1;
-    }
-
     // must have some prompt
     if (!use_stdin && !use_file && params.prompt.empty()) {
         LOG_ERR("error: must specify one of: --stdin, --file or --prompt\n");

tools/ui/eslint.config.js

@@ -29,6 +29,13 @@ export default ts.config(
 			// This app uses hash-based routing (#/) where resolve() from $app/paths does not apply
 			'svelte/no-navigation-without-resolve': 'off',
 
+			// Snippet bodies often ignore one or more of the parent's params
+			// (e.g. `{#snippet children(_meta, ctx)}` when only ctx is read).
+			'@typescript-eslint/no-unused-vars': [
+				'error',
+				{ argsIgnorePattern: '^_', varsIgnorePattern: '^_' }
+			],
+
 			// Enforce empty line at end of file
 			'eol-last': 'error'
 		}

tools/ui/src/app.css

@@ -193,6 +193,33 @@
 		-ms-overflow-style: none;
 		scrollbar-width: none;
 	}
+
+	.shimmer-text {
+		background: linear-gradient(
+			90deg,
+			var(--muted-foreground),
+			var(--foreground),
+			var(--muted-foreground)
+		);
+		background-size: 200% 100%;
+		background-clip: text;
+		-webkit-background-clip: text;
+		-webkit-text-fill-color: transparent;
+		font-weight: 500;
+		animation: shimmer 1s linear infinite;
+	}
+
+	@keyframes shimmer {
+		to {
+			background-position: -200% 0;
+		}
+	}
+
+	@media (prefers-reduced-motion: reduce) {
+		.shimmer-text {
+			animation: none;
+		}
+	}
 }
 
 .mermaidTooltip {

tools/ui/src/lib/components/app/actions/ActionIconCopyToClipboard.svelte

@@ -1,4 +1,5 @@
 <script lang="ts">
+	import { ICON_CLASS_DEFAULT } from '$lib/constants/css-classes';
 	import { Copy } from '@lucide/svelte';
 	import { copyToClipboard } from '$lib/utils';
 	import ActionIcon from './ActionIcon.svelte';
@@ -11,7 +12,7 @@
 <ActionIcon
 	icon={Copy}
 	tooltip={ariaLabel}
-	iconSize="h-4 w-4"
+	iconSize={ICON_CLASS_DEFAULT}
 	disabled={!canCopy}
 	onclick={() => canCopy && copyToClipboard(text)}
 />

tools/ui/src/lib/components/app/chat/ChatAttachments/ChatAttachmentsList/ChatAttachmentsListItem/ChatAttachmentsListItemThumbnailFile.svelte

@@ -1,4 +1,5 @@
 <script lang="ts">
+	import { ICON_CLASS_DEFAULT } from '$lib/constants/css-classes';
 	import { X, Music, Video } from '@lucide/svelte';
 	import {
 		formatFileSize,
@@ -109,9 +110,9 @@
 		class="flex h-8 w-8 items-center justify-center rounded bg-primary/10 text-xs font-medium text-primary"
 	>
 		{#if isAudio}
-			<Music class="h-4 w-4 text-white/70" />
+			<Music class="{ICON_CLASS_DEFAULT} text-white/70" />
 		{:else if isVideo}
-			<Video class="h-4 w-4 text-white/70" />
+			<Video class="{ICON_CLASS_DEFAULT} text-white/70" />
 		{:else}
 			{fileTypeLabel}
 		{/if}

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

@@ -1,4 +1,5 @@
 <script lang="ts">
+	import { ICON_CLASS_DEFAULT } from '$lib/constants/css-classes';
 	import type { ChatAttachmentDisplayItem } from '$lib/types';
 	import { FileText, Eye, Info } from '@lucide/svelte';
 	import { Button } from '$lib/components/ui/button';
@@ -88,7 +89,7 @@
 		onclick={() => (pdfViewMode = PdfViewMode.TEXT)}
 		disabled={pdfImagesLoading}
 	>
-		<FileText class="mr-1 h-4 w-4" />
+		<FileText class="mr-1 {ICON_CLASS_DEFAULT}" />
 		Text
 	</Button>
 
@@ -100,10 +101,10 @@
 	>
 		{#if pdfImagesLoading}
 			<div
-				class="mr-1 h-4 w-4 animate-spin rounded-full border-2 border-current border-t-transparent"
+				class="mr-1 {ICON_CLASS_DEFAULT} animate-spin rounded-full border-2 border-current border-t-transparent"
 			></div>
 		{:else}
-			<Eye class="mr-1 h-4 w-4" />
+			<Eye class="mr-1 {ICON_CLASS_DEFAULT}" />
 		{/if}
 		Pages
 	</Button>
@@ -111,7 +112,7 @@
 
 {#if !hasVisionModality && activeModelId && currentItem}
 	<Alert.Root class="mb-4 max-w-4xl">
-		<Info class="h-4 w-4" />
+		<Info class={ICON_CLASS_DEFAULT} />
 		<Alert.Title>Preview only</Alert.Title>
 		<Alert.Description>
 			<span class="inline-flex">

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

@@ -1,4 +1,5 @@
 <script lang="ts">
+	import { ICON_CLASS_DEFAULT } from '$lib/constants/css-classes';
 	import { Music, Video, FileText } from '@lucide/svelte';
 	import { HorizontalScrollCarousel } from '$lib/components/app/misc';
 
@@ -49,11 +50,11 @@
 							class="bg-foreground-muted/50 flex h-12 w-12 flex-col items-center justify-center gap-0.5 py-1"
 						>
 							{#if item.isAudio}
-								<Music class="h-4 w-4 text-white/70" />
+								<Music class="{ICON_CLASS_DEFAULT} text-white/70" />
 							{:else if item.isVideo}
-								<Video class="h-4 w-4 text-white/70" />
+								<Video class="{ICON_CLASS_DEFAULT} text-white/70" />
 							{:else}
-								<FileText class="h-4 w-4 text-white/70" />
+								<FileText class="{ICON_CLASS_DEFAULT} text-white/70" />
 							{/if}
 
 							<span class="font-mono text-[9px] text-white/60">{getFileExtension(item.name)}</span>

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

@@ -1,4 +1,5 @@
 <script lang="ts">
+	import { ICON_CLASS_DEFAULT } from '$lib/constants/css-classes';
 	import { Plus } from '@lucide/svelte';
 	import { Button } from '$lib/components/ui/button';
 	import * as Tooltip from '$lib/components/ui/tooltip';
@@ -23,7 +24,7 @@
 		>
 			<span class="sr-only">{ATTACHMENT_TOOLTIP_TEXT}</span>
 
-			<Plus class="h-4 w-4" />
+			<Plus class={ICON_CLASS_DEFAULT} />
 		</Button>
 	</Tooltip.Trigger>
 

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

@@ -1,4 +1,5 @@
 <script lang="ts">
+	import { ICON_CLASS_DEFAULT } from '$lib/constants/css-classes';
 	import { Plus, File, MessageSquare, Zap, FolderOpen } from '@lucide/svelte';
 	import * as DropdownMenu from '$lib/components/ui/dropdown-menu';
 	import * as Tooltip from '$lib/components/ui/tooltip';
@@ -83,7 +84,7 @@
 					>
 						<span class="sr-only">{ATTACHMENT_TOOLTIP_TEXT}</span>
 
-						<Plus class="h-4 w-4" />
+						<Plus class={ICON_CLASS_DEFAULT} />
 					</DropdownMenu.Trigger>
 				{/snippet}
 			</Tooltip.Trigger>
@@ -100,7 +101,7 @@
 
 			<DropdownMenu.Sub>
 				<DropdownMenu.SubTrigger class="flex cursor-pointer items-center gap-2">
-					<File class="h-4 w-4" />
+					<File class={ICON_CLASS_DEFAULT} />
 
 					<span>Add files</span>
 				</DropdownMenu.SubTrigger>
@@ -113,7 +114,7 @@
 								class="{item.class ?? ''} flex cursor-pointer items-center gap-2"
 								onclick={() => attachmentMenu.callbacks[item.action]()}
 							>
-								<item.icon class="h-4 w-4" />
+								<item.icon class={ICON_CLASS_DEFAULT} />
 
 								<span>{item.label}</span>
 							</DropdownMenu.Item>
@@ -126,7 +127,7 @@
 												class="{item.class ?? ''} flex items-center gap-2"
 												disabled
 											>
-												<item.icon class="h-4 w-4" />
+												<item.icon class={ICON_CLASS_DEFAULT} />
 
 												<span>{item.label}</span>
 											</DropdownMenu.Item>
@@ -147,7 +148,7 @@
 				class="flex cursor-pointer items-center gap-2"
 				onclick={onSystemPromptClick}
 			>
-				<MessageSquare class="h-4 w-4" />
+				<MessageSquare class={ICON_CLASS_DEFAULT} />
 
 				<span>System Message</span>
 			</DropdownMenu.Item>
@@ -163,7 +164,7 @@
 					class="flex cursor-pointer items-center gap-2"
 					onclick={onMcpPromptClick}
 				>
-					<Zap class="h-4 w-4" />
+					<Zap class={ICON_CLASS_DEFAULT} />
 
 					<span>MCP Prompt</span>
 				</DropdownMenu.Item>
@@ -174,7 +175,7 @@
 					class="flex cursor-pointer items-center gap-2"
 					onclick={onMcpResourcesClick}
 				>
-					<FolderOpen class="h-4 w-4" />
+					<FolderOpen class={ICON_CLASS_DEFAULT} />
 
 					<span>MCP Resources</span>
 				</DropdownMenu.Item>

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

@@ -1,4 +1,5 @@
 <script lang="ts">
+	import { ICON_CLASS_DEFAULT } from '$lib/constants/css-classes';
 	import { Settings, Plus } from '@lucide/svelte';
 	import { Switch } from '$lib/components/ui/switch';
 	import * as DropdownMenu from '$lib/components/ui/dropdown-menu';
@@ -60,7 +61,7 @@
 <DropdownMenu.Root>
 	<DropdownMenu.Sub onOpenChange={handleMcpSubMenuOpen}>
 		<DropdownMenu.SubTrigger class="flex cursor-pointer items-center gap-2">
-			<McpLogo class="h-4 w-4" />
+			<McpLogo class={ICON_CLASS_DEFAULT} />
 
 			<span>MCP Servers</span>
 		</DropdownMenu.SubTrigger>
@@ -92,7 +93,7 @@
 										<McpServerIdentity
 											{displayName}
 											{faviconUrl}
-											iconClass="h-4 w-4"
+											iconClass={ICON_CLASS_DEFAULT}
 											iconRounded="rounded-sm"
 											showVersion={false}
 											nameClass="text-sm"
@@ -123,7 +124,7 @@
 							class="flex cursor-pointer items-center gap-2"
 							onclick={handleMcpSettingsClick}
 						>
-							<Settings class="h-4 w-4" />
+							<Settings class={ICON_CLASS_DEFAULT} />
 
 							<span>Manage MCP Servers</span>
 						</DropdownMenu.Item>
@@ -140,7 +141,7 @@
 					class="flex cursor-pointer items-center gap-2"
 					onclick={handleMcpSettingsClick}
 				>
-					<Plus class="h-4 w-4" />
+					<Plus class={ICON_CLASS_DEFAULT} />
 
 					<span>Add MCP Servers</span>
 				</DropdownMenu.Item>

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

@@ -1,4 +1,5 @@
 <script lang="ts">
+	import { ICON_CLASS_DEFAULT } from '$lib/constants/css-classes';
 	import { Lightbulb, LightbulbOff, Check, Info } from '@lucide/svelte';
 	import * as DropdownMenu from '$lib/components/ui/dropdown-menu';
 	import * as Tooltip from '$lib/components/ui/tooltip';
@@ -13,9 +14,9 @@
 	<DropdownMenu.Sub bind:open={subOpen}>
 		<DropdownMenu.SubTrigger class="flex cursor-pointer items-center gap-2">
 			{#if reasoning.thinkingEnabled}
-				<Lightbulb class="h-4 w-4 shrink-0 text-amber-400" />
+				<Lightbulb class="{ICON_CLASS_DEFAULT} shrink-0 text-amber-400" />
 			{:else}
-				<LightbulbOff class="h-4 w-4 shrink-0 text-muted-foreground" />
+				<LightbulbOff class="{ICON_CLASS_DEFAULT} shrink-0 text-muted-foreground" />
 			{/if}
 
 			<span
@@ -46,9 +47,9 @@
 					}}
 				>
 					{#if reasoning.isSelected(level)}
-						<Check class="h-4 w-4 shrink-0 text-foreground" />
+						<Check class="{ICON_CLASS_DEFAULT} shrink-0 text-foreground" />
 					{:else}
-						<div class="h-4 w-4 shrink-0"></div>
+						<div class="{ICON_CLASS_DEFAULT} shrink-0"></div>
 					{/if}
 
 					<span class="flex-1">{level.label}</span>

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

@@ -1,4 +1,5 @@
 <script lang="ts">
+	import { ICON_CLASS_DEFAULT } from '$lib/constants/css-classes';
 	import type { Snippet } from 'svelte';
 	import * as Tooltip from '$lib/components/ui/tooltip';
 	import * as Sheet from '$lib/components/ui/sheet';
@@ -108,15 +109,15 @@
 					>
 						<Collapsible.Trigger class={sheetItemClass}>
 							{#if reasoningExpanded}
-								<ChevronDown class="h-4 w-4 shrink-0" />
+								<ChevronDown class="{ICON_CLASS_DEFAULT} shrink-0" />
 							{:else}
-								<ChevronRight class="h-4 w-4 shrink-0" />
+								<ChevronRight class="{ICON_CLASS_DEFAULT} shrink-0" />
 							{/if}
 
 							{#if reasoning.thinkingEnabled}
-								<Lightbulb class="h-4 w-4 shrink-0 text-amber-400" />
+								<Lightbulb class="{ICON_CLASS_DEFAULT} shrink-0 text-amber-400" />
 							{:else}
-								<LightbulbOff class="h-4 w-4 shrink-0 text-muted-foreground" />
+								<LightbulbOff class="{ICON_CLASS_DEFAULT} shrink-0 text-muted-foreground" />
 							{/if}
 
 							<span class="flex-1">Reasoning</span>
@@ -138,9 +139,9 @@
 									>
 										<div class="flex min-w-0 items-center gap-3">
 											{#if reasoning.isSelected(level)}
-												<Check class="h-4 w-4 shrink-0 text-foreground" />
+												<Check class="{ICON_CLASS_DEFAULT} shrink-0 text-foreground" />
 											{:else}
-												<div class="h-4 w-4 shrink-0"></div>
+												<div class="{ICON_CLASS_DEFAULT} shrink-0"></div>
 											{/if}
 
 											<span class="text-sm">{level.label}</span>
@@ -161,12 +162,12 @@
 				<Collapsible.Root open={filesExpanded} onOpenChange={(open) => (filesExpanded = open)}>
 					<Collapsible.Trigger class={sheetItemClass}>
 						{#if filesExpanded}
-							<ChevronDown class="h-4 w-4 shrink-0" />
+							<ChevronDown class="{ICON_CLASS_DEFAULT} shrink-0" />
 						{:else}
-							<ChevronRight class="h-4 w-4 shrink-0" />
+							<ChevronRight class="{ICON_CLASS_DEFAULT} shrink-0" />
 						{/if}
 
-						<File class="h-4 w-4 shrink-0" />
+						<File class="{ICON_CLASS_DEFAULT} shrink-0" />
 
 						<span class="flex-1">Add files</span>
 					</Collapsible.Trigger>
@@ -181,7 +182,7 @@
 										class={sheetItemClass}
 										onclick={() => attachmentMenu.callbacks[item.action]()}
 									>
-										<item.icon class="h-4 w-4 shrink-0" />
+										<item.icon class="{ICON_CLASS_DEFAULT} shrink-0" />
 
 										<span>{item.label}</span>
 									</button>
@@ -189,7 +190,7 @@
 									<Tooltip.Root delayDuration={TOOLTIP_DELAY_DURATION}>
 										<Tooltip.Trigger>
 											<button type="button" class={sheetItemClass} disabled>
-												<item.icon class="h-4 w-4 shrink-0" />
+												<item.icon class="{ICON_CLASS_DEFAULT} shrink-0" />
 
 												<span>{item.label}</span>
 											</button>
@@ -208,12 +209,12 @@
 				<Collapsible.Root open={mcpExpanded} onOpenChange={(open) => (mcpExpanded = open)}>
 					<Collapsible.Trigger class={sheetItemClass}>
 						{#if mcpExpanded}
-							<ChevronDown class="h-4 w-4 shrink-0" />
+							<ChevronDown class="{ICON_CLASS_DEFAULT} shrink-0" />
 						{:else}
-							<ChevronRight class="h-4 w-4 shrink-0" />
+							<ChevronRight class="{ICON_CLASS_DEFAULT} shrink-0" />
 						{/if}
 
-						<McpLogo class="inline h-4 w-4 shrink-0" />
+						<McpLogo class="inline {ICON_CLASS_DEFAULT} shrink-0" />
 
 						<span class="flex-1">MCP Servers</span>
 
@@ -242,7 +243,7 @@
 											<img
 												src={faviconUrl}
 												alt=""
-												class="h-4 w-4 shrink-0 rounded-sm"
+												class="{ICON_CLASS_DEFAULT} shrink-0 rounded-sm"
 												onerror={(e) => {
 													(e.currentTarget as HTMLImageElement).style.display = 'none';
 												}}
@@ -280,12 +281,12 @@
 					<Collapsible.Root open={toolsExpanded} onOpenChange={(open) => (toolsExpanded = open)}>
 						<Collapsible.Trigger class={sheetItemClass}>
 							{#if toolsExpanded}
-								<ChevronDown class="h-4 w-4 shrink-0" />
+								<ChevronDown class="{ICON_CLASS_DEFAULT} shrink-0" />
 							{:else}
-								<ChevronRight class="h-4 w-4 shrink-0" />
+								<ChevronRight class="{ICON_CLASS_DEFAULT} shrink-0" />
 							{/if}
 
-							<PencilRuler class="inline h-4 w-4 shrink-0" />
+							<PencilRuler class="inline {ICON_CLASS_DEFAULT} shrink-0" />
 
 							<span class="flex-1">Tools</span>
 
@@ -310,7 +311,7 @@
 											<img
 												src={favicon}
 												alt=""
-												class="h-4 w-4 shrink-0 rounded-sm"
+												class="{ICON_CLASS_DEFAULT} shrink-0 rounded-sm"
 												onerror={(e) => {
 													(e.currentTarget as HTMLImageElement).style.display = 'none';
 												}}
@@ -325,7 +326,7 @@
 
 										<Checkbox
 											{checked}
-											class="h-4 w-4 shrink-0"
+											class="{ICON_CLASS_DEFAULT} shrink-0"
 											onclick={(e) => e.stopPropagation()}
 											onCheckedChange={() => toolsPanel.toggleGroupByLabel(group.label)}
 										/>
@@ -341,7 +342,7 @@
 					class={sheetItemClass}
 					onclick={() => attachmentMenu.callbacks[AttachmentAction.SYSTEM_PROMPT_CLICK]()}
 				>
-					<MessageSquare class="h-4 w-4 shrink-0" />
+					<MessageSquare class="{ICON_CLASS_DEFAULT} shrink-0" />
 
 					<span>System Message</span>
 				</button>
@@ -352,7 +353,7 @@
 						class={sheetItemClass}
 						onclick={() => attachmentMenu.callbacks[AttachmentAction.MCP_PROMPT_CLICK]()}
 					>
-						<Zap class="h-4 w-4 shrink-0" />
+						<Zap class="{ICON_CLASS_DEFAULT} shrink-0" />
 
 						<span>MCP Prompt</span>
 					</button>
@@ -364,7 +365,7 @@
 						class={sheetItemClass}
 						onclick={() => attachmentMenu.callbacks[AttachmentAction.MCP_RESOURCES_CLICK]()}
 					>
-						<FolderOpen class="h-4 w-4 shrink-0" />
+						<FolderOpen class="{ICON_CLASS_DEFAULT} shrink-0" />
 
 						<span>MCP Resources</span>
 					</button>

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

@@ -1,4 +1,5 @@
 <script lang="ts">
+	import { ICON_CLASS_DEFAULT } from '$lib/constants/css-classes';
 	import { PencilRuler, ChevronDown, ChevronRight, Loader2, Info, Check } from '@lucide/svelte';
 	import { Checkbox } from '$lib/components/ui/checkbox';
 	import * as Collapsible from '$lib/components/ui/collapsible';
@@ -15,7 +16,7 @@
 
 <DropdownMenu.Sub onOpenChange={(open) => open && toolsPanel.handleOpen()}>
 	<DropdownMenu.SubTrigger class="flex cursor-pointer items-center gap-2">
-		<PencilRuler class="h-4 w-4" />
+		<PencilRuler class={ICON_CLASS_DEFAULT} />
 
 		<span>Tools</span>
 	</DropdownMenu.SubTrigger>
@@ -24,14 +25,14 @@
 		{#if toolsPanel.totalToolCount === 0}
 			{#if toolsStore.loading}
 				<div class="px-3 py-4 text-center text-sm text-muted-foreground">
-					<Loader2 class="mx-auto mb-1 h-4 w-4 animate-spin" />
+					<Loader2 class="mx-auto mb-1 {ICON_CLASS_DEFAULT} animate-spin" />
 
 					Loading tools...
 				</div>
 			{:else if toolsStore.isToolsEndpointUnreachable}
 				<div class="grid gap-2.5 px-3 py-4 text-sm text-muted-foreground">
 					<span class="flex gap-2">
-						<Info class="mt-0.5 h-4 w-4 shrink-0" />
+						<Info class="mt-0.5 {ICON_CLASS_DEFAULT} shrink-0" />
 
 						<span>
 							Run llama-server with <code>{CLI_FLAGS.TOOLS}</code> flag to enable
@@ -41,7 +42,7 @@
 					</span>
 
 					<span class="flex gap-2">
-						<Info class="mt-0.5 h-4 w-4 shrink-0" />
+						<Info class="mt-0.5 {ICON_CLASS_DEFAULT} shrink-0" />
 
 						<span>
 							{hasMcpServersAvailable ? 'Enable' : 'Add'} MCP Server(s) to access
@@ -54,7 +55,7 @@
 				<div class="px-3 py-4 text-center text-sm text-muted-foreground">Failed to load tools</div>
 			{:else if toolsPanel.noToolsInfoMessage}
 				<div class="flex gap-2 px-3 py-4 text-sm text-muted-foreground">
-					<Info class="mt-0.5 h-4 w-4 shrink-0" />
+					<Info class="mt-0.5 {ICON_CLASS_DEFAULT} shrink-0" />
 
 					<span>{toolsPanel.noToolsInfoMessage}</span>
 				</div>
@@ -87,7 +88,7 @@
 										<img
 											src={favicon}
 											alt=""
-											class="h-4 w-4 shrink-0 rounded-sm"
+											class="{ICON_CLASS_DEFAULT} shrink-0 rounded-sm"
 											onerror={(e) => {
 												(e.currentTarget as HTMLImageElement).style.display = 'none';
 											}}
@@ -109,7 +110,7 @@
 											{...props}
 											{checked}
 											onCheckedChange={() => toolsPanel.toggleGroupByLabel(group.label)}
-											class="mr-2 h-4 w-4 shrink-0"
+											class="mr-2 {ICON_CLASS_DEFAULT} shrink-0"
 										/>
 									{/snippet}
 								</Tooltip.Trigger>

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

@@ -1,4 +1,5 @@
 <script lang="ts">
+	import { ICON_CLASS_DEFAULT } from '$lib/constants/css-classes';
 	import { Mic, Square } from '@lucide/svelte';
 	import { Button } from '$lib/components/ui/button';
 	import * as Tooltip from '$lib/components/ui/tooltip';
@@ -36,9 +37,9 @@
 				<span class="sr-only">{isRecording ? 'Stop recording' : 'Start recording'}</span>
 
 				{#if isRecording}
-					<Square class="h-4 w-4 animate-pulse fill-white" />
+					<Square class="{ICON_CLASS_DEFAULT} animate-pulse fill-white" />
 				{:else}
-					<Mic class="h-4 w-4" />
+					<Mic class={ICON_CLASS_DEFAULT} />
 				{/if}
 			</Button>
 		</Tooltip.Trigger>

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

@@ -1,4 +1,5 @@
 <script lang="ts">
+	import { ICON_CLASS_DEFAULT } from '$lib/constants/css-classes';
 	import { Square, SkipForward } from '@lucide/svelte';
 	import { Button } from '$lib/components/ui/button';
 	import { ChatService } from '$lib/services';
@@ -185,7 +186,9 @@
 		>
 			<span class="sr-only">Skip reasoning</span>
 
-			<SkipForward class="h-4 w-4 stroke-muted-foreground group-hover:stroke-foreground" />
+			<SkipForward
+				class="{ICON_CLASS_DEFAULT} stroke-muted-foreground group-hover:stroke-foreground"
+			/>
 		</Button>
 	{/if}
 

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

@@ -6,7 +6,7 @@
 	import { REASONING_EFFORT_TOKENS } from '$lib/constants/reasoning-effort-tokens';
 	import { REASONING_EFFORT_LEVELS } from '$lib/constants/reasoning-effort';
 	import type { ReasoningEffortLevel } from '$lib/types';
-	import { DIALOG_SUBMENU_CONTENT } from '$lib/constants/css-classes';
+	import { DIALOG_SUBMENU_CONTENT, ICON_CLASS_DEFAULT } from '$lib/constants/css-classes';
 	import {
 		modelsStore,
 		checkModelSupportsThinking,
@@ -74,9 +74,9 @@
 	<DropdownMenu.Sub bind:open={subOpen}>
 		<DropdownMenu.SubTrigger class="flex cursor-pointer items-center gap-2">
 			{#if thinkingEnabled}
-				<Lightbulb class="h-4 w-4 shrink-0 text-amber-400" />
+				<Lightbulb class="{ICON_CLASS_DEFAULT} shrink-0 text-amber-400" />
 			{:else}
-				<LightbulbOff class="h-4 w-4 shrink-0 text-muted-foreground" />
+				<LightbulbOff class="{ICON_CLASS_DEFAULT} shrink-0 text-muted-foreground" />
 			{/if}
 
 			<span class="flex-1">Thinking</span>
@@ -118,7 +118,7 @@
 					{/if}
 
 					{#if isSelected(level)}
-						<Check class="h-4 w-4 shrink-0 text-foreground" />
+						<Check class="{ICON_CLASS_DEFAULT} shrink-0 text-foreground" />
 					{/if}
 				</button>
 			{/each}

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

@@ -384,7 +384,6 @@
 			{isLastAssistantMessage}
 			{message}
 			{toolMessages}
-			messageContent={message.content}
 			onConfirmDelete={handleConfirmDelete}
 			onContinue={handleContinue}
 			onCopy={handleCopy}

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

@@ -2,23 +2,20 @@
 	import {
 		ChatMessageAgenticContent,
 		ChatMessageActionIcons,
-		ChatMessageEditForm,
-		ChatMessageStatistics,
-		ModelBadge,
-		ModelsSelectorDropdown
+		ChatMessageAssistantModel,
+		ChatMessageAssistantProcessingInfo,
+		ChatMessageAssistantRawOutput,
+		ChatMessageAssistantStatistics,
+		ChatMessageEditForm
 	} from '$lib/components/app';
 	import { getMessageEditContext } from '$lib/contexts';
 	import { useProcessingState } from '$lib/hooks/use-processing-state.svelte';
 	import { isLoading, isChatStreaming } from '$lib/stores/chat.svelte';
-	import { copyToClipboard, deriveAgenticSections, modelLoadProgressText } from '$lib/utils';
-	import { AgenticSectionType, ChatMessageStatisticsMode } from '$lib/enums';
-	import { REASONING_TAGS } from '$lib/constants/agentic';
-	import { fade } from 'svelte/transition';
+	import { modelLoadProgressText } from '$lib/utils';
 	import { MessageRole } from '$lib/enums';
 	import { config } from '$lib/stores/settings.svelte';
 	import { isRouterMode } from '$lib/stores/server.svelte';
 	import { modelsStore } from '$lib/stores/models.svelte';
-	import { ServerModelStatus } from '$lib/enums';
 
 	import { hasAgenticContent } from '$lib/utils';
 
@@ -33,7 +30,6 @@
 		isLastAssistantMessage?: boolean;
 		message: DatabaseMessage;
 		toolMessages?: DatabaseMessage[];
-		messageContent: string | undefined;
 		onCopy: () => void;
 		onConfirmDelete: () => void;
 		onContinue?: () => void;
@@ -54,7 +50,6 @@
 		isLastAssistantMessage = false,
 		message,
 		toolMessages = [],
-		messageContent,
 		onConfirmDelete,
 		onContinue,
 		onCopy,
@@ -77,55 +72,11 @@
 
 	let currentConfig = $derived(config());
 	let isRouter = $derived(isRouterMode());
+
 	let showRawOutput = $state(false);
 
-	let rawOutputContent = $derived.by(() => {
-		const sections = deriveAgenticSections(message, toolMessages, [], false);
-		const parts: string[] = [];
-
-		for (const section of sections) {
-			switch (section.type) {
-				case AgenticSectionType.REASONING:
-				case AgenticSectionType.REASONING_PENDING:
-					parts.push(`${REASONING_TAGS.START}\n${section.content}\n${REASONING_TAGS.END}`);
-					break;
-
-				case AgenticSectionType.TEXT:
-					parts.push(section.content);
-					break;
-
-				case AgenticSectionType.TOOL_CALL:
-				case AgenticSectionType.TOOL_CALL_PENDING:
-				case AgenticSectionType.TOOL_CALL_STREAMING: {
-					const callObj: Record<string, unknown> = { name: section.toolName };
-
-					if (section.toolArgs) {
-						try {
-							callObj.arguments = JSON.parse(section.toolArgs);
-						} catch {
-							callObj.arguments = section.toolArgs;
-						}
-					}
-
-					parts.push(JSON.stringify(callObj, null, 2));
-
-					if (section.toolResult) {
-						parts.push(`[Tool Result]\n${section.toolResult}`);
-					}
-
-					break;
-				}
-			}
-		}
-
-		return parts.join('\n\n\n');
-	});
-
 	let displayedModel = $derived(message.model ?? null);
 
-	// model being switched to while it loads, so the selector bar tracks it
-	let pendingModel = $state<string | null>(null);
-
 	let isCurrentlyLoading = $derived(isLoading());
 	let isStreaming = $derived(isChatStreaming());
 	let hasNoContent = $derived(!message?.content?.trim());
@@ -189,10 +140,6 @@
 		};
 	});
 
-	function handleCopyModel() {
-		void copyToClipboard(displayedModel ?? '');
-	}
-
 	$effect(() => {
 		if (showProcessingInfoTop || showProcessingInfoBottom) {
 			processingState.startMonitoring();
@@ -211,23 +158,14 @@
 	aria-label="Assistant message with actions"
 >
 	{#if showProcessingInfoTop}
-		<div class="mt-6 w-full max-w-3xl" in:fade>
-			<div class="processing-container">
-				<span class="processing-text">
-					{modelLoadingText ??
-						processingState.getPromptProgressText() ??
-						processingState.getProcessingMessage() ??
-						'Processing...'}
-				</span>
-			</div>
-		</div>
+		<ChatMessageAssistantProcessingInfo {modelLoadingText} {processingState} position="top" />
 	{/if}
 
 	{#if editCtx.isEditing}
 		<ChatMessageEditForm />
-	{:else if message.role === MessageRole.ASSISTANT}
+	{:else}
 		{#if showRawOutput}
-			<pre class="raw-output">{rawOutputContent || ''}</pre>
+			<ChatMessageAssistantRawOutput {message} {toolMessages} />
 		{:else}
 			<ChatMessageAgenticContent
 				{message}
@@ -236,78 +174,28 @@
 				{isLastAssistantMessage}
 			/>
 		{/if}
-	{:else}
-		<div class="text-sm whitespace-pre-wrap">
-			{messageContent}
-		</div>
 	{/if}
 
 	{#if showProcessingInfoBottom}
-		<div class="mt-4 w-full max-w-3xl" in:fade>
-			<div class="processing-container">
-				<span class="processing-text">
-					{modelLoadingText ??
-						processingState.getPromptProgressText() ??
-						processingState.getProcessingMessage() ??
-						'Processing...'}
-				</span>
-			</div>
-		</div>
+		<ChatMessageAssistantProcessingInfo {modelLoadingText} {processingState} position="bottom" />
 	{/if}
 
 	<div class="info my-6 grid gap-4 tabular-nums">
 		{#if displayedModel}
 			<div class="inline-flex flex-wrap items-start gap-2 text-xs text-muted-foreground">
-				{#if isRouter}
-					<ModelsSelectorDropdown
-						currentModel={pendingModel ?? displayedModel}
-						disabled={isLoading()}
-						onModelChange={async (modelId: string, modelName: string) => {
-							const status = modelsStore.getModelStatus(modelId);
+				<ChatMessageAssistantModel
+					{displayedModel}
+					isLoading={isLoading()}
+					{isRouter}
+					{onRegenerate}
+				/>
 
-							if (status !== ServerModelStatus.LOADED) {
-								pendingModel = modelId;
-
-								try {
-									await modelsStore.loadModel(modelId);
-								} finally {
-									pendingModel = null;
-								}
-							}
-
-							onRegenerate(modelName);
-							return true;
-						}}
-					/>
-				{:else}
-					<ModelBadge model={displayedModel || undefined} onclick={handleCopyModel} />
-				{/if}
-
-				{#if currentConfig.showMessageStats && message.timings && message.timings.predicted_n && message.timings.predicted_ms}
-					{@const agentic = message.timings.agentic}
-					<ChatMessageStatistics
-						mode={ChatMessageStatisticsMode.GENERATION}
-						promptTokens={agentic ? agentic.llm.prompt_n : message.timings.prompt_n}
-						promptMs={agentic ? agentic.llm.prompt_ms : message.timings.prompt_ms}
-						predictedTokens={agentic ? agentic.llm.predicted_n : message.timings.predicted_n}
-						predictedMs={agentic ? agentic.llm.predicted_ms : message.timings.predicted_ms}
-						agenticTimings={agentic}
-					/>
-				{:else if isLoading() && currentConfig.showMessageStats}
-					{@const liveStats = processingState.getLiveProcessingStats()}
-					{@const genStats = processingState.getLiveGenerationStats()}
-
-					{#if genStats}
-						<ChatMessageStatistics
-							mode={ChatMessageStatisticsMode.GENERATION}
-							isLive
-							promptTokens={liveStats?.tokensProcessed}
-							promptMs={liveStats?.timeMs}
-							predictedTokens={genStats.tokensGenerated}
-							predictedMs={genStats.timeMs}
-						/>
-					{/if}
-				{/if}
+				<ChatMessageAssistantStatistics
+					{message}
+					isLoading={isLoading()}
+					{processingState}
+					showMessageStats={currentConfig.showMessageStats}
+				/>
 			</div>
 		{/if}
 	</div>
@@ -353,47 +241,4 @@
 			);
 		}
 	}
-
-	.processing-container {
-		display: flex;
-		flex-direction: column;
-		align-items: flex-start;
-		gap: 0.5rem;
-	}
-
-	.processing-text {
-		background: linear-gradient(
-			90deg,
-			var(--muted-foreground),
-			var(--foreground),
-			var(--muted-foreground)
-		);
-		background-size: 200% 100%;
-		background-clip: text;
-		-webkit-background-clip: text;
-		-webkit-text-fill-color: transparent;
-		animation: shine 1s linear infinite;
-		font-weight: 500;
-		font-size: 0.875rem;
-	}
-
-	@keyframes shine {
-		to {
-			background-position: -200% 0;
-		}
-	}
-
-	.raw-output {
-		width: 100%;
-		max-width: 48rem;
-		margin-top: 1.5rem;
-		padding: 1rem 1.25rem;
-		border-radius: 1rem;
-		background: hsl(var(--muted) / 0.3);
-		color: var(--foreground);
-		font-size: 0.875rem;
-		line-height: 1.6;
-		white-space: pre-wrap;
-		word-break: break-word;
-	}
 </style>

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

@@ -0,0 +1,46 @@
+<script lang="ts">
+	import { ModelBadge, ModelsSelectorDropdown } from '$lib/components/app';
+	import { copyToClipboard } from '$lib/utils';
+	import { modelsStore } from '$lib/stores/models.svelte';
+	import { ServerModelStatus } from '$lib/enums';
+
+	interface Props {
+		displayedModel: string | null;
+		isRouter: boolean;
+		isLoading: boolean;
+		onRegenerate: (modelOverride?: string) => void;
+	}
+
+	let { displayedModel, isRouter, isLoading, onRegenerate }: Props = $props();
+
+	let pendingModel = $state<string | null>(null);
+
+	function handleCopyModel() {
+		void copyToClipboard(displayedModel ?? '');
+	}
+</script>
+
+{#if isRouter}
+	<ModelsSelectorDropdown
+		currentModel={pendingModel ?? displayedModel}
+		disabled={isLoading}
+		onModelChange={async (modelId: string, modelName: string) => {
+			const status = modelsStore.getModelStatus(modelId);
+
+			if (status !== ServerModelStatus.LOADED) {
+				pendingModel = modelId;
+
+				try {
+					await modelsStore.loadModel(modelId);
+				} finally {
+					pendingModel = null;
+				}
+			}
+
+			onRegenerate(modelName);
+			return true;
+		}}
+	/>
+{:else}
+	<ModelBadge model={displayedModel || undefined} onclick={handleCopyModel} />
+{/if}

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

@@ -0,0 +1,25 @@
+<script lang="ts">
+	import { fade } from 'svelte/transition';
+	import type { UseProcessingStateReturn } from '$lib/hooks/use-processing-state.svelte';
+
+	interface Props {
+		modelLoadingText: string | null;
+		processingState: UseProcessingStateReturn;
+		position: 'top' | 'bottom';
+	}
+
+	let { modelLoadingText, processingState, position }: Props = $props();
+
+	const marginClass = position === 'top' ? 'mt-6' : 'mt-4';
+</script>
+
+<div class="{marginClass} w-full max-w-3xl" in:fade>
+	<div class="flex flex-col items-start gap-2">
+		<span class="shimmer-text text-sm">
+			{modelLoadingText ??
+				processingState.getPromptProgressText() ??
+				processingState.getProcessingMessage() ??
+				'Processing...'}
+		</span>
+	</div>
+</div>

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

@@ -0,0 +1,33 @@
+<script lang="ts">
+	import { deriveAgenticSections, buildAssistantRawOutput } from '$lib/utils';
+
+	interface Props {
+		message: DatabaseMessage;
+		toolMessages?: DatabaseMessage[];
+	}
+
+	let { message, toolMessages = [] }: Props = $props();
+
+	let rawOutputContent = $derived.by(() => {
+		const sections = deriveAgenticSections(message, toolMessages, [], false);
+		return buildAssistantRawOutput(sections);
+	});
+</script>
+
+<pre class="raw-output">{rawOutputContent || ''}</pre>
+
+<style>
+	.raw-output {
+		width: 100%;
+		max-width: 48rem;
+		margin-top: 1.5rem;
+		padding: 1rem 1.25rem;
+		border-radius: 1rem;
+		background: hsl(var(--muted) / 0.3);
+		color: var(--foreground);
+		font-size: 0.875rem;
+		line-height: 1.6;
+		white-space: pre-wrap;
+		word-break: break-word;
+	}
+</style>

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

@@ -0,0 +1,40 @@
+<script lang="ts">
+	import { ChatMessageStatistics } from '$lib/components/app';
+	import { ChatMessageStatisticsMode } from '$lib/enums';
+	import type { UseProcessingStateReturn } from '$lib/hooks/use-processing-state.svelte';
+
+	interface Props {
+		message: DatabaseMessage;
+		isLoading: boolean;
+		processingState: UseProcessingStateReturn;
+		showMessageStats: boolean;
+	}
+
+	let { message, isLoading, processingState, showMessageStats }: Props = $props();
+</script>
+
+{#if showMessageStats && message.timings && message.timings.predicted_n && message.timings.predicted_ms}
+	{@const agentic = message.timings.agentic}
+	<ChatMessageStatistics
+		mode={ChatMessageStatisticsMode.GENERATION}
+		promptTokens={agentic ? agentic.llm.prompt_n : message.timings.prompt_n}
+		promptMs={agentic ? agentic.llm.prompt_ms : message.timings.prompt_ms}
+		predictedTokens={agentic ? agentic.llm.predicted_n : message.timings.predicted_n}
+		predictedMs={agentic ? agentic.llm.predicted_ms : message.timings.predicted_ms}
+		agenticTimings={agentic}
+	/>
+{:else if isLoading && showMessageStats}
+	{@const liveStats = processingState.getLiveProcessingStats()}
+	{@const genStats = processingState.getLiveGenerationStats()}
+
+	{#if genStats}
+		<ChatMessageStatistics
+			mode={ChatMessageStatisticsMode.GENERATION}
+			isLive
+			promptTokens={liveStats?.tokensProcessed}
+			promptMs={liveStats?.timeMs}
+			predictedTokens={genStats.tokensGenerated}
+			predictedMs={genStats.timeMs}
+		/>
+	{/if}
+{/if}

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

@@ -0,0 +1,66 @@
+<script lang="ts">
+	import { BuiltInTool } from '$lib/enums';
+	import {
+		extractSearchQuery,
+		extractSearchResults,
+		isWebSearchToolName,
+		type AgenticSection
+	} from '$lib/utils';
+	import type { DatabaseMessageExtra } from '$lib/types';
+	import ChatMessageToolCallBlockDefault from './ChatMessageToolCallBlockDefault.svelte';
+	import ChatMessageToolCallBlockEditFile from './ChatMessageToolCallBlockEditFile.svelte';
+	import ChatMessageToolCallBlockExecShellCommand from './ChatMessageToolCallBlockExecShellCommand.svelte';
+	import ChatMessageToolCallBlockFileGlobSearch from './ChatMessageToolCallBlockFileGlobSearch.svelte';
+	import ChatMessageToolCallBlockGetDatetime from './ChatMessageToolCallBlockGetDatetime.svelte';
+	import ChatMessageToolCallBlockGrepSearch from './ChatMessageToolCallBlockGrepSearch.svelte';
+	import ChatMessageToolCallBlockReadFile from './ChatMessageToolCallBlockReadFile.svelte';
+	import ChatMessageToolCallBlockRunJavascript from './ChatMessageToolCallBlockRunJavascript.svelte';
+	import ChatMessageToolCallBlockSearchResults from './ChatMessageToolCallBlockSearchResults.svelte';
+	import ChatMessageToolCallBlockWriteFile from './ChatMessageToolCallBlockWriteFile.svelte';
+
+	interface Props {
+		section: AgenticSection;
+		attachments?: DatabaseMessageExtra[];
+		open: boolean;
+		isStreaming: boolean;
+		isExecuting?: boolean;
+		onToggle?: () => void;
+	}
+
+	let { section, attachments, open, isStreaming, isExecuting, onToggle }: Props = $props();
+
+	const searchResults = $derived(extractSearchResults(section.toolResult));
+	const searchQuery = $derived(extractSearchQuery(section.toolArgs));
+	const isSearchCall = $derived(
+		searchResults.length > 0 || (searchQuery.length > 0 && isWebSearchToolName(section.toolName))
+	);
+</script>
+
+{#if isSearchCall}
+	<ChatMessageToolCallBlockSearchResults {section} {open} {isStreaming} {onToggle} />
+{:else if section.toolName === BuiltInTool.GET_DATETIME}
+	<ChatMessageToolCallBlockGetDatetime {section} {isStreaming} />
+{:else if section.toolName === BuiltInTool.READ_FILE}
+	<ChatMessageToolCallBlockReadFile {section} {open} {isStreaming} {onToggle} />
+{:else if section.toolName === BuiltInTool.EDIT_FILE}
+	<ChatMessageToolCallBlockEditFile {section} {open} {isStreaming} {onToggle} />
+{:else if section.toolName === BuiltInTool.WRITE_FILE}
+	<ChatMessageToolCallBlockWriteFile {section} {open} {isStreaming} {onToggle} />
+{:else if section.toolName === BuiltInTool.EXEC_SHELL_COMMAND}
+	<ChatMessageToolCallBlockExecShellCommand
+		{section}
+		{open}
+		{isStreaming}
+		{isExecuting}
+		{attachments}
+		{onToggle}
+	/>
+{:else if section.toolName === BuiltInTool.FILE_GLOB_SEARCH}
+	<ChatMessageToolCallBlockFileGlobSearch {section} {open} {isStreaming} {onToggle} />
+{:else if section.toolName === BuiltInTool.GREP_SEARCH}
+	<ChatMessageToolCallBlockGrepSearch {section} {open} {isStreaming} {onToggle} />
+{:else if section.toolName === BuiltInTool.RUN_JAVASCRIPT}
+	<ChatMessageToolCallBlockRunJavascript {section} {open} {isStreaming} {onToggle} />
+{:else}
+	<ChatMessageToolCallBlockDefault {section} {open} {isStreaming} {attachments} {onToggle} />
+{/if}

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

@@ -0,0 +1,124 @@
+<script lang="ts">
+	// Fall-through renderer for tool calls without a dedicated block.
+	// Renders section.toolArgs / section.toolResult directly using the
+	// shared chrome shell.
+
+	import { Loader2 } from '@lucide/svelte';
+	import { MarkdownContent, SyntaxHighlightedCode } from '$lib/components/app';
+	import { FileTypeText, ToolResultKind } from '$lib/enums';
+	import { MAX_HEIGHT_CODE_BLOCK } from '$lib/constants';
+	import {
+		classifyToolResult,
+		formatJsonPretty,
+		parseToolResultWithImages,
+		type AgenticSection,
+		type ToolResultLine
+	} from '$lib/utils';
+	import { getBuiltinToolUi } from '$lib/constants/built-in-tools';
+	import type { DatabaseMessageExtra } from '$lib/types';
+	import ToolCallBlock from './ToolCallBlock.svelte';
+
+	interface Props {
+		section: AgenticSection;
+		open: boolean;
+		isStreaming: boolean;
+		attachments?: DatabaseMessageExtra[];
+		onToggle?: () => void;
+	}
+
+	let { section, open, isStreaming, attachments, onToggle }: Props = $props();
+
+	const title = $derived(getBuiltinToolUi(section.toolName)?.label ?? section.toolName ?? '');
+
+	const parsedLines: ToolResultLine[] = $derived(
+		section.toolResult ? parseToolResultWithImages(section.toolResult, attachments) : []
+	);
+	const outputKind = $derived(classifyToolResult(section.toolResult));
+</script>
+
+<ToolCallBlock {section} {open} {isStreaming} meta={null} {title} {onToggle}>
+	{#snippet children(_meta, ctx)}
+		{#if ctx.isStreamingCall}
+			<div class="mb-2 flex items-center gap-2 text-xs text-muted-foreground/70">
+				<span>Input</span>
+				{#if ctx.isStreaming}
+					<Loader2 class="h-3 w-3 animate-spin" />
+				{/if}
+			</div>
+			{#if section.toolArgs}
+				<SyntaxHighlightedCode
+					code={formatJsonPretty(section.toolArgs)}
+					language={FileTypeText.JSON}
+					maxHeight={MAX_HEIGHT_CODE_BLOCK}
+					streaming={ctx.isCodeStreaming}
+				/>
+			{:else if ctx.isStreaming}
+				<div class="rounded bg-muted/20 p-2 text-xs text-muted-foreground/70 italic">
+					Receiving arguments...
+				</div>
+			{:else}
+				<div
+					class="rounded bg-yellow-500/10 p-2 text-xs text-yellow-600 italic dark:text-yellow-400"
+				>
+					Response was truncated
+				</div>
+			{/if}
+		{:else}
+			{@const showInput = Boolean(section.toolArgs)}
+			{#if showInput}
+				<div class="mb-1.5 flex items-center gap-2 text-xs text-muted-foreground/70">
+					<span>Input</span>
+				</div>
+				<SyntaxHighlightedCode
+					code={formatJsonPretty(section.toolArgs ?? '')}
+					language={FileTypeText.JSON}
+					maxHeight={MAX_HEIGHT_CODE_BLOCK}
+					streaming={ctx.isCodeStreaming}
+				/>
+			{/if}
+			<div
+				class={showInput
+					? 'mt-4 mb-1.5 flex items-center gap-2 text-xs text-muted-foreground/70'
+					: 'mb-1.5 flex items-center gap-2 text-xs text-muted-foreground/70'}
+			>
+				<span>Output</span>
+				{#if ctx.isPending}
+					<Loader2 class="h-3 w-3 animate-spin" />
+				{/if}
+			</div>
+			{#if ctx.isPending}
+				<div class="rounded bg-muted/20 p-2 text-xs text-muted-foreground/70 italic">
+					Waiting for result...
+				</div>
+			{:else if section.toolResult}
+				{#if outputKind === ToolResultKind.JSON}
+					<SyntaxHighlightedCode
+						code={formatJsonPretty(section.toolResult)}
+						language={FileTypeText.JSON}
+						maxHeight={MAX_HEIGHT_CODE_BLOCK}
+					/>
+				{:else if outputKind === ToolResultKind.MARKDOWN}
+					<MarkdownContent content={section.toolResult} {attachments} />
+				{:else}
+					<div class="overflow-auto">
+						{#each parsedLines as line, i (i)}
+							<div class="font-mono text-[11px] leading-relaxed whitespace-pre-wrap">
+								{line.text}
+							</div>
+							{#if line.image}
+								<img
+									src={line.image.base64Url}
+									alt={line.image.name}
+									class="mt-2 mb-2 h-auto max-w-full rounded-lg"
+									loading="lazy"
+								/>
+							{/if}
+						{/each}
+					</div>
+				{/if}
+			{:else}
+				<div class="rounded bg-muted/20 p-2 text-xs text-muted-foreground/70 italic">No output</div>
+			{/if}
+		{/if}
+	{/snippet}
+</ToolCallBlock>

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

@@ -0,0 +1,166 @@
+<script lang="ts">
+	import { XCircle } from '@lucide/svelte';
+	import { MAX_HEIGHT_CODE_BLOCK, RESULT_STAT_SEPARATOR } from '$lib/constants';
+	import { computeLineDiff, prefixFor, type AgenticSection } from '$lib/utils';
+	import { parseEditFileMeta } from './parsers/edit-file';
+	import ToolCallBlock from './ToolCallBlock.svelte';
+
+	interface Props {
+		section: AgenticSection;
+		open: boolean;
+		isStreaming: boolean;
+		onToggle?: () => void;
+	}
+
+	let { section, open, isStreaming, onToggle }: Props = $props();
+
+	const editFileMeta = $derived(parseEditFileMeta(section));
+
+	const editDiffs = $derived(
+		(editFileMeta?.edits ?? []).map((edit) => computeLineDiff(edit.oldText, edit.newText))
+	);
+</script>
+
+<ToolCallBlock {section} {open} {isStreaming} meta={editFileMeta} {onToggle}>
+	{#snippet titleSnippet()}
+		<span class="text-muted-foreground">Edit file </span>
+		<span class="font-mono">{editFileMeta?.filePath}</span>
+		{#if editFileMeta?.errorMessage}
+			<span class="ml-1 text-xs italic text-muted-foreground/70">(failed)</span>
+		{/if}
+	{/snippet}
+
+	{#snippet children(meta, _ctx)}
+		{#if meta?.errorMessage}
+			<div
+				class="flex items-start gap-2 rounded bg-red-500/10 p-2 text-xs text-red-600 italic dark:text-red-400"
+			>
+				<XCircle class="mt-0.5 h-3 w-3 shrink-0" />
+				<span>{meta.errorMessage}</span>
+			</div>
+		{:else if meta && meta.edits.length > 0}
+			{#each editDiffs as diffLines, ei (ei)}
+				<div class={ei === 0 ? '' : 'mt-3'}>
+					<div class="mb-1.5 text-xs text-muted-foreground/70 italic">
+						Edit {ei + 1}&nbsp;of&nbsp;{meta.edits.length}
+					</div>
+					<div class="diff-block" style:max-height={MAX_HEIGHT_CODE_BLOCK}>
+						<div class="diff-pre">
+							{#each diffLines as line, li (li)}
+								<div class="diff-line diff-{line.kind}">
+									<span class="diff-old-num">{line.oldLine ?? ''}</span>
+									<span class="diff-marker">{prefixFor(line.kind)}</span>
+									<span class="diff-new-num">{line.newLine ?? ''}</span>
+									<span class="diff-text">{line.text || ' '}</span>
+								</div>
+							{/each}
+						</div>
+					</div>
+				</div>
+			{/each}
+			<div class="mt-1.5 text-xs text-muted-foreground/70 italic">
+				{#if meta.resultMessage}
+					{meta.resultMessage}{meta.editsApplied != null ? RESULT_STAT_SEPARATOR : ''}{/if}
+				{#if meta.editsApplied != null}
+					<span class="font-mono">{meta.editsApplied}</span>
+					{meta.editsApplied === 1 ? 'edit' : 'edits'}&nbsp;applied
+				{/if}
+			</div>
+		{:else}
+			<div class="rounded bg-muted/20 p-2 text-xs text-muted-foreground/70 italic">No edits</div>
+		{/if}
+	{/snippet}
+</ToolCallBlock>
+
+<style>
+	.diff-block {
+		overflow: auto;
+		border-radius: 0.75rem;
+		border-width: 1px;
+		border-color: color-mix(in oklch, var(--border) 30%, transparent);
+		background: var(--code-background);
+		box-shadow: 0 1px 2px 0 rgb(0 0 0 / 0.05);
+	}
+
+	:global(.dark) .diff-block {
+		border-color: color-mix(in oklch, var(--border) 20%, transparent);
+	}
+
+	/* Each row is a 4-column grid: old-line#, marker, new-line#, text.
+	 * The gutters stay fixed-width so the text column lines up unversally. */
+	.diff-line {
+		display: grid;
+		grid-template-columns: 3.25rem 1.5rem 3.25rem 1fr;
+		font-family: var(--font-mono);
+		font-size: 11px;
+		line-height: 1.65;
+		align-items: stretch;
+	}
+
+	.diff-old-num,
+	.diff-new-num {
+		text-align: right;
+		padding-right: 0.5rem;
+		user-select: none;
+		color: color-mix(in oklch, var(--muted-foreground) 70%, transparent);
+		font-variant-numeric: tabular-nums;
+	}
+
+	.diff-marker {
+		text-align: center;
+		color: color-mix(in oklch, var(--muted-foreground) 70%, transparent);
+		user-select: none;
+	}
+
+	.diff-line.diff-add {
+		background-color: #f0fff4;
+		color: #22863a;
+	}
+	.diff-line.diff-add .diff-new-num,
+	.diff-line.diff-add .diff-marker {
+		color: #22863a;
+	}
+
+	.diff-line.diff-remove {
+		background-color: #ffeef0;
+		color: #b31d28;
+	}
+	.diff-line.diff-remove .diff-old-num,
+	.diff-line.diff-remove .diff-marker {
+		color: #b31d28;
+	}
+
+	.diff-line.diff-add .diff-old-num,
+	.diff-line.diff-remove .diff-new-num {
+		/* Empty gutter columns for add/remove rows mirror git unification
+		 * (added lines don't have an old number, removed lines don't have a
+		 * new number). Keep them visible so columns stay aligned across
+		 * mixed rows. */
+		opacity: 0;
+	}
+
+	.diff-text {
+		padding-left: 0.4rem;
+		padding-right: 0.5rem;
+		white-space: pre;
+		overflow-x: auto;
+		min-width: 0;
+	}
+
+	:global(.dark) .diff-line.diff-add {
+		background-color: #033a16;
+		color: #aff5b4;
+	}
+	:global(.dark) .diff-line.diff-add .diff-new-num,
+	:global(.dark) .diff-line.diff-add .diff-marker {
+		color: #aff5b4;
+	}
+	:global(.dark) .diff-line.diff-remove {
+		background-color: #67060c;
+		color: #ffdcd7;
+	}
+	:global(.dark) .diff-line.diff-remove .diff-old-num,
+	:global(.dark) .diff-line.diff-remove .diff-marker {
+		color: #ffdcd7;
+	}
+</style>

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

@@ -0,0 +1,293 @@
+<script lang="ts">
+	// Block for `exec_shell_command`. Unlike the other tools, this
+	// renderer uses CollapsibleTerminalBlock (terminal-style frame)
+	// and treats "live" output chunks as active even after the call
+	// resolved, so the spinner stays on while stdout is still flowing.
+	// The scroll-to-bottom auto-scroll logic mirrors what was here
+	// before extraction.
+
+	import { Check, Loader2, XCircle, AlertTriangle } from '@lucide/svelte';
+	import { CollapsibleTerminalBlock } from '$lib/components/app';
+	import { SETTINGS_KEYS } from '$lib/constants';
+	import { config } from '$lib/stores/settings.svelte';
+	import { TOOL_RUNTIME_SCROLL_AT_BOTTOM_THRESHOLD_PX } from '$lib/constants/auto-scroll';
+	import {
+		highlightCode,
+		isExitCodeSummaryLine,
+		parseExecShellCommandError,
+		parseExecShellCommandExitStatus,
+		parseToolResultWithImages,
+		type AgenticSection,
+		type ExecShellExitStatus,
+		type ToolResultLine
+	} from '$lib/utils';
+	import { parseExecShellCommandMeta } from './parsers/exec-shell-command';
+	import type { DatabaseMessageExtra } from '$lib/types';
+	import ToolCallBlock from './ToolCallBlock.svelte';
+
+	interface Props {
+		section: AgenticSection;
+		open: boolean;
+		isStreaming: boolean;
+		/** True while the agentic loop is streaming output chunks for THIS
+		 *  tool call. Drives max-height + auto-scroll while true; releases
+		 *  them when the loop reports this call as done. */
+		isExecuting?: boolean;
+		attachments?: DatabaseMessageExtra[];
+		onToggle?: () => void;
+	}
+
+	let { section, open, isStreaming, isExecuting = false, attachments, onToggle }: Props = $props();
+
+	// `isLive` covers all in-flight phases: pre-chunk spinner and
+	// streaming itself. Frozen output (tool done while agent continues)
+	// is not live.
+	const isLive = $derived(isExecuting);
+
+	const execShellMeta = $derived(parseExecShellCommandMeta(section));
+	const execShellError = $derived(parseExecShellCommandError(section.toolResult));
+	const execShellExitStatus: ExecShellExitStatus | undefined = $derived(
+		parseExecShellCommandExitStatus(section.toolResult)
+	);
+
+	const parsedLines: ToolResultLine[] = $derived(
+		section.toolResult ? parseToolResultWithImages(section.toolResult, attachments) : []
+	);
+
+	// Drop the trailing "[exit code: N]" line - rendered as a colored
+	// badge below. During streaming we keep it so a partial stream still
+	// shows the status once the final chunk lands.
+	const outputLines: ToolResultLine[] = $derived(
+		execShellExitStatus && parsedLines.length > 0
+			? parsedLines.slice(0, parsedLines.length - 1)
+			: parsedLines
+	);
+
+	const isExitCodeFinalLine = $derived(
+		execShellExitStatus !== undefined &&
+			parsedLines.length > 0 &&
+			isExitCodeSummaryLine(parsedLines[parsedLines.length - 1].text, execShellExitStatus)
+	);
+
+	// Highlight just the command for the title; the (typically large)
+	// output blob uses bare monospace to skip hljs per-line highlighting.
+	const highlightedCommandHtml = $derived(
+		execShellMeta ? highlightCode(execShellMeta.command, 'bash') : ''
+	);
+
+	const exitBadgeClass = $derived(
+		execShellExitStatus?.timedOut
+			? 'exit-badge warning'
+			: execShellExitStatus?.code === 0
+				? 'exit-badge success'
+				: 'exit-badge failure'
+	);
+
+	const useFullHeightCodeBlocks = $derived(
+		Boolean(config()[SETTINGS_KEYS.FULL_HEIGHT_CODE_BLOCKS])
+	);
+
+	const autoScroll = $derived(isLive && !useFullHeightCodeBlocks);
+
+	const SCROLL_BOTTOM_THRESHOLD_PX = TOOL_RUNTIME_SCROLL_AT_BOTTOM_THRESHOLD_PX;
+
+	let scrollEl: HTMLDivElement | undefined = $state();
+	let userScrolledUp = $state(false);
+	let lastScrollTop = 0;
+	let pendingFrame: number | null = null;
+
+	function isAtBottom(): boolean {
+		if (!scrollEl) return false;
+		return (
+			scrollEl.scrollHeight - scrollEl.clientHeight - scrollEl.scrollTop <=
+			SCROLL_BOTTOM_THRESHOLD_PX
+		);
+	}
+
+	function scrollToBottomOnFrame() {
+		if (pendingFrame !== null || !scrollEl || userScrolledUp) return;
+		pendingFrame = requestAnimationFrame(() => {
+			pendingFrame = null;
+
+			// Re-check on rAF - user may scroll between scheduling and paint.
+			if (scrollEl && !userScrolledUp) {
+				scrollEl.scrollTop = scrollEl.scrollHeight;
+			}
+		});
+	}
+
+	function handleScrollEvent() {
+		if (!scrollEl) return;
+		const isScrollingUp = scrollEl.scrollTop < lastScrollTop;
+		if (isScrollingUp && !isAtBottom()) {
+			userScrolledUp = true;
+		} else if (isAtBottom()) {
+			userScrolledUp = false;
+		}
+		lastScrollTop = scrollEl.scrollTop;
+	}
+
+	$effect(() => {
+		void section.toolResult;
+		if (!scrollEl || !autoScroll) return;
+		scrollToBottomOnFrame();
+	});
+
+	$effect(() => {
+		// Catch layout changes that don't touch toolResult (line-wrap
+		// reflow, image attaches, hljs settle).
+		if (!scrollEl || !autoScroll) return;
+
+		const observer = new MutationObserver(() => scrollToBottomOnFrame());
+		observer.observe(scrollEl, {
+			childList: true,
+			subtree: true,
+			characterData: true
+		});
+
+		return () => observer.disconnect();
+	});
+
+	$effect(() => {
+		// Reset on stream end so the next render (full-height) starts
+		// pinned.
+		if (!isLive) {
+			userScrolledUp = false;
+			lastScrollTop = 0;
+		}
+	});
+</script>
+
+{#snippet execShellTitle()}
+	{#if highlightedCommandHtml}
+		<span class="font-mono">{@html highlightedCommandHtml}</span>
+	{:else}
+		<span class="font-mono">{execShellMeta?.command}</span>
+	{/if}
+{/snippet}
+
+<ToolCallBlock
+	{section}
+	{open}
+	{isStreaming}
+	meta={execShellMeta ? { errorMessage: execShellError } : null}
+	wrapper={CollapsibleTerminalBlock}
+	extraLiveStreaming={isLive}
+	spinIconWhenActive={true}
+	{onToggle}
+>
+	{#snippet titleSnippet()}
+		{@render execShellTitle()}
+	{/snippet}
+
+	{#snippet children(_meta, ctx)}
+		{#if ctx.isPending}
+			<div class="flex items-start gap-2 text-xs text-muted-foreground/70">
+				<Loader2 class="h-3 w-3 animate-spin" />
+				Running...
+			</div>
+		{:else if execShellError}
+			<div class="flex items-start gap-2 text-xs text-red-600 italic dark:text-red-400">
+				<XCircle class="mt-0.5 h-3 w-3 shrink-0" />
+				<span>{execShellError}</span>
+			</div>
+		{:else if section.toolResult}
+			<div
+				bind:this={scrollEl}
+				class="terminal-output"
+				class:is-clamped={!useFullHeightCodeBlocks}
+				onscroll={handleScrollEvent}
+			>
+				{#each outputLines as line, i (i)}
+					<div class="font-mono text-[11px] leading-relaxed whitespace-pre-wrap">{line.text}</div>
+					{#if line.image}
+						<img
+							src={line.image.base64Url}
+							alt={line.image.name}
+							class="mt-2 mb-2 h-auto max-w-full rounded-lg"
+							loading="lazy"
+						/>
+					{/if}
+				{/each}
+
+				{#if isExitCodeFinalLine && execShellExitStatus}
+					<div class={exitBadgeClass}>
+						{#if execShellExitStatus.timedOut}
+							<AlertTriangle class="h-3 w-3" />
+							<span>timed out</span>
+							<span class="exit-sep">&middot;</span>
+							<span>exit {execShellExitStatus.code}</span>
+						{:else if execShellExitStatus.code === 0}
+							<Check class="h-3 w-3" />
+							<span>exit 0</span>
+						{:else}
+							<XCircle class="h-3 w-3" />
+							<span>exit {execShellExitStatus.code}</span>
+						{/if}
+					</div>
+				{/if}
+			</div>
+		{/if}
+	{/snippet}
+</ToolCallBlock>
+
+<style>
+	.terminal-output {
+		overscroll-behavior: contain;
+	}
+
+	.terminal-output.is-clamped {
+		max-height: 28rem;
+		overflow-y: auto;
+		scrollbar-gutter: stable;
+		padding-right: 0.25rem;
+	}
+
+	.exit-badge {
+		display: inline-flex;
+		align-items: center;
+		gap: 0.35rem;
+		margin-top: 0.5rem;
+		padding: 0.2rem 0.55rem;
+		border-radius: 0.375rem;
+		font-family: var(--font-mono);
+		font-size: 11px;
+		font-weight: 500;
+		letter-spacing: 0.01em;
+		line-height: 1;
+	}
+
+	.exit-badge.success {
+		background: color-mix(in oklch, var(--color-green-500, #22c55e) 14%, transparent);
+		color: var(--color-green-700, #15803d);
+	}
+
+	:global(.dark) .exit-badge.success {
+		background: color-mix(in oklch, var(--color-green-400, #4ade80) 18%, transparent);
+		color: var(--color-green-300, #86efac);
+	}
+
+	.exit-badge.failure {
+		background: color-mix(in oklch, var(--color-red-500, #ef4444) 14%, transparent);
+		color: var(--color-red-700, #b91c1c);
+	}
+
+	:global(.dark) .exit-badge.failure {
+		background: color-mix(in oklch, var(--color-red-400, #f87171) 18%, transparent);
+		color: var(--color-red-300, #fca5a5);
+	}
+
+	.exit-badge.warning {
+		background: color-mix(in oklch, var(--color-amber-500, #f59e0b) 14%, transparent);
+		color: var(--color-amber-700, #b45309);
+	}
+
+	:global(.dark) .exit-badge.warning {
+		background: color-mix(in oklch, var(--color-amber-400, #fbbf24) 18%, transparent);
+		color: var(--color-amber-300, #fcd34d);
+	}
+
+	.exit-sep {
+		opacity: 0.45;
+	}
+</style>

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

@@ -0,0 +1,61 @@
+<script lang="ts">
+	import { XCircle } from '@lucide/svelte';
+	import { type AgenticSection } from '$lib/utils';
+	import { parseFileGlobSearchMeta } from './parsers/file-glob-search';
+	import ToolCallBlock from './ToolCallBlock.svelte';
+
+	interface Props {
+		section: AgenticSection;
+		open: boolean;
+		isStreaming: boolean;
+		onToggle?: () => void;
+	}
+
+	let { section, open, isStreaming, onToggle }: Props = $props();
+
+	const fileGlobMeta = $derived(parseFileGlobSearchMeta(section));
+</script>
+
+<ToolCallBlock {section} {open} {isStreaming} meta={fileGlobMeta} {onToggle}>
+	{#snippet titleSnippet()}
+		{#if fileGlobMeta}
+			<span class="text-muted-foreground"
+				>{fileGlobMeta.include === '**' ? 'List files' : 'Search files'}&nbsp;</span
+			>
+			{#if fileGlobMeta.include !== '**'}
+				<span class="font-mono">{fileGlobMeta.include}</span>
+			{/if}
+			<span class="text-muted-foreground">&nbsp;in&nbsp;</span>
+			<span class="font-mono">{fileGlobMeta.path}</span>
+		{/if}
+	{/snippet}
+
+	{#snippet children(meta, ctx)}
+		{#if ctx.isPending}
+			<div class="rounded bg-muted/20 p-2 text-xs text-muted-foreground/70 italic">
+				Searching...
+			</div>
+		{:else if meta?.errorMessage}
+			<div
+				class="flex items-start gap-2 rounded bg-red-500/10 p-2 text-xs text-red-600 italic dark:text-red-400"
+			>
+				<XCircle class="mt-0.5 h-3 w-3 shrink-0" />
+				<span>{meta.errorMessage}</span>
+			</div>
+		{:else if meta && meta.matches.length > 0}
+			<div class="max-h-96 overflow-auto">
+				{#each meta.matches as match, i (i)}
+					<div class="font-mono text-[11px] leading-relaxed whitespace-pre-wrap">{match}</div>
+				{/each}
+			</div>
+			<div class="mt-1.5 text-xs text-muted-foreground/70 italic">
+				Total matches: <span class="font-mono">{meta.totalMatches ?? meta.matches.length}</span>
+			</div>
+		{:else}
+			<div class="text-xs text-muted-foreground/70 italic">No matches</div>
+			<div class="mt-1.5 text-xs text-muted-foreground/70 italic">
+				Total matches: <span class="font-mono">{meta?.totalMatches ?? 0}</span>
+			</div>
+		{/if}
+	{/snippet}
+</ToolCallBlock>

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

@@ -0,0 +1,57 @@
+<script lang="ts">
+	import { Clock, Loader2 } from '@lucide/svelte';
+	import { AgenticSectionType } from '$lib/enums';
+	import type { AgenticSection } from '$lib/utils';
+
+	interface Props {
+		section: AgenticSection;
+		isStreaming?: boolean;
+	}
+
+	let { section, isStreaming = false }: Props = $props();
+
+	const isPending = $derived(section.type === AgenticSectionType.TOOL_CALL_PENDING);
+	const isStreamingCall = $derived(section.type === AgenticSectionType.TOOL_CALL_STREAMING);
+	const showSpinner = $derived(isPending || (isStreamingCall && isStreaming));
+
+	type GetDatetimeMeta = {
+		dateString?: string;
+		errorMessage?: string;
+	};
+
+	function parseGetDatetimeMeta(toolResultString: string | undefined): GetDatetimeMeta {
+		if (!toolResultString) return {};
+
+		try {
+			const parsed: unknown = JSON.parse(toolResultString);
+			if (parsed && typeof parsed === 'object' && !Array.isArray(parsed)) {
+				const obj = parsed as Record<string, unknown>;
+				if (typeof obj.error === 'string') return { errorMessage: obj.error };
+				if (typeof obj.result === 'string') return { dateString: obj.result.trim() };
+			}
+		} catch {
+			return { dateString: toolResultString.trim() };
+		}
+
+		return {};
+	}
+
+	const dateMeta = $derived(parseGetDatetimeMeta(section.toolResult));
+</script>
+
+<div class="text-muted-foreground flex items-center gap-2 py-1.5">
+	<Clock class="text-muted-foreground/60 h-3.5 w-3.5 shrink-0" />
+	{#if showSpinner}
+		<span class="text-foreground/80 text-sm font-medium">Current time</span>
+		<Loader2 class="text-muted-foreground/70 h-3 w-3 animate-spin" />
+	{:else if dateMeta.errorMessage}
+		<span class="text-foreground/80 text-sm font-medium">Current time&nbsp;</span>
+		<span class="text-red-600 text-xs italic dark:text-red-400">-&nbsp;{dateMeta.errorMessage}</span
+		>
+	{:else if dateMeta.dateString}
+		<span class="text-foreground/80 text-sm font-medium">Current time is&nbsp;</span>
+		<span class="font-mono text-foreground/90 text-sm">{dateMeta.dateString}</span>
+	{:else}
+		<span class="text-foreground/80 text-sm font-medium">Current time</span>
+	{/if}
+</div>

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

@@ -0,0 +1,67 @@
+<script lang="ts">
+	import { XCircle } from '@lucide/svelte';
+	import { type AgenticSection } from '$lib/utils';
+	import { parseGrepSearchMeta } from './parsers/grep-search';
+	import ToolCallBlock from './ToolCallBlock.svelte';
+
+	interface Props {
+		section: AgenticSection;
+		open: boolean;
+		isStreaming: boolean;
+		onToggle?: () => void;
+	}
+
+	let { section, open, isStreaming, onToggle }: Props = $props();
+
+	const grepMeta = $derived(parseGrepSearchMeta(section));
+</script>
+
+<ToolCallBlock {section} {open} {isStreaming} meta={grepMeta} {onToggle}>
+	{#snippet titleSnippet()}
+		{#if grepMeta}
+			<span class="text-muted-foreground">Search for&nbsp;</span>
+			<span class="font-mono">{grepMeta.pattern}</span>
+			<span class="text-muted-foreground">&nbsp;in&nbsp;</span>
+			<span class="font-mono">{grepMeta.path}</span>
+		{/if}
+	{/snippet}
+
+	{#snippet children(meta, ctx)}
+		{#if ctx.isPending}
+			<div class="rounded bg-muted/20 p-2 text-xs text-muted-foreground/70 italic">
+				Searching...
+			</div>
+		{:else if meta?.errorMessage}
+			<div
+				class="flex items-start gap-2 rounded bg-red-500/10 p-2 text-xs text-red-600 italic dark:text-red-400"
+			>
+				<XCircle class="mt-0.5 h-3 w-3 shrink-0" />
+				<span>{meta.errorMessage}</span>
+			</div>
+		{:else if meta && meta.matches.length > 0}
+			<div class="max-h-96 overflow-auto">
+				{#each meta.matches as match, mi (mi)}
+					<div class="font-mono text-[11px] leading-relaxed">
+						<span class="text-muted-foreground/70">{match.file}</span>
+						{#if meta.showLineNumbers && match.line != null}
+							<span class="text-muted-foreground/70">:{match.line}</span>
+						{/if}
+						<span class="text-muted-foreground/70">:</span>
+						<span>{match.content}</span>
+					</div>
+				{/each}
+			</div>
+			<div class="mt-1.5 text-xs text-muted-foreground/70 italic">
+				Total matches: <span class="font-mono">{meta.totalMatches ?? meta.matches.length}</span>
+				{#if meta.showLineNumbers}
+					&nbsp;<span class="italic">(with line numbers)</span>
+				{/if}
+			</div>
+		{:else}
+			<div class="text-xs text-muted-foreground/70 italic">No matches</div>
+			<div class="mt-1.5 text-xs text-muted-foreground/70 italic">
+				Total matches: <span class="font-mono">{meta?.totalMatches ?? 0}</span>
+			</div>
+		{/if}
+	{/snippet}
+</ToolCallBlock>

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

@@ -0,0 +1,44 @@
+<script lang="ts">
+	import { SyntaxHighlightedCode } from '$lib/components/app';
+	import { DEFAULT_LANGUAGE, MAX_HEIGHT_CODE_BLOCK } from '$lib/constants';
+	import { type AgenticSection } from '$lib/utils';
+	import { parseReadFileMeta } from './parsers/read-file';
+	import ToolCallBlock from './ToolCallBlock.svelte';
+
+	interface Props {
+		section: AgenticSection;
+		open: boolean;
+		isStreaming: boolean;
+		onToggle?: () => void;
+	}
+
+	let { section, open, isStreaming, onToggle }: Props = $props();
+
+	const readFileMeta = $derived(parseReadFileMeta(section));
+</script>
+
+<ToolCallBlock {section} {open} {isStreaming} meta={readFileMeta} {onToggle}>
+	{#snippet titleSnippet()}
+		<span class="text-muted-foreground">Read file </span>
+		<span class="font-mono">{readFileMeta?.fileName}</span>
+		{#if readFileMeta?.lineRange}
+			<span class="text-muted-foreground"
+				>&nbsp;(lines {readFileMeta.lineRange.start}-{readFileMeta.lineRange.end})</span
+			>
+		{/if}
+	{/snippet}
+
+	{#snippet children(_meta, _ctx)}
+		{#if section.toolResult}
+			<SyntaxHighlightedCode
+				code={section.toolResult}
+				language={readFileMeta?.language ?? DEFAULT_LANGUAGE}
+				maxHeight={MAX_HEIGHT_CODE_BLOCK}
+			/>
+		{:else}
+			<div class="rounded bg-muted/20 p-2 text-xs text-muted-foreground/70 italic">
+				Waiting for file content...
+			</div>
+		{/if}
+	{/snippet}
+</ToolCallBlock>

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

@@ -0,0 +1,69 @@
+<script lang="ts">
+	import { XCircle, Terminal } from '@lucide/svelte';
+	import { SyntaxHighlightedCode } from '$lib/components/app';
+	import { FileTypeText } from '$lib/enums';
+	import { MAX_HEIGHT_CODE_BLOCK } from '$lib/constants';
+	import { getBuiltinToolUi, type AgenticSection } from '$lib/utils';
+	import { parseRunJavascriptMeta } from './parsers/run-javascript';
+	import ToolCallBlock from './ToolCallBlock.svelte';
+
+	interface Props {
+		section: AgenticSection;
+		open: boolean;
+		isStreaming: boolean;
+		onToggle?: () => void;
+	}
+
+	let { section, open, isStreaming, onToggle }: Props = $props();
+
+	const runJsMeta = $derived(parseRunJavascriptMeta(section));
+	const title = $derived(getBuiltinToolUi(section.toolName)?.label ?? section.toolName ?? '');
+</script>
+
+<ToolCallBlock {section} {open} {isStreaming} meta={runJsMeta} {title} {onToggle}>
+	{#snippet children(meta, ctx)}
+		{#if ctx.isPending}
+			<div class="rounded bg-muted/20 p-2 text-xs text-muted-foreground/70 italic">Running...</div>
+		{:else if meta?.errorMessage}
+			<div
+				class="flex items-start gap-2 rounded bg-red-500/10 p-2 text-xs text-red-600 italic dark:text-red-400"
+			>
+				<XCircle class="mt-0.5 h-3 w-3 shrink-0" />
+				<span>{meta.errorMessage}</span>
+			</div>
+			<div class="mt-3">
+				<SyntaxHighlightedCode
+					code={meta.code}
+					language={FileTypeText.JAVASCRIPT}
+					maxHeight={MAX_HEIGHT_CODE_BLOCK}
+					streaming={ctx.isCodeStreaming}
+				/>
+			</div>
+		{:else if meta}
+			<SyntaxHighlightedCode
+				code={meta.code}
+				language={FileTypeText.JAVASCRIPT}
+				maxHeight={MAX_HEIGHT_CODE_BLOCK}
+				streaming={ctx.isCodeStreaming}
+			/>
+			<div class="mb-2 mt-3 flex items-center gap-2 text-xs text-muted-foreground/70">
+				<Terminal class="h-3 w-3" />
+				<span>Console</span>
+				{#if meta.timeoutMs != null}
+					<span class="font-mono">&middot;&nbsp;timeout&nbsp;{meta.timeoutMs}&nbsp;ms</span>
+				{/if}
+			</div>
+			{#if section.toolResult}
+				<div class="mt-1">
+					<SyntaxHighlightedCode
+						code={section.toolResult}
+						language={FileTypeText.JAVASCRIPT}
+						maxHeight={MAX_HEIGHT_CODE_BLOCK}
+					/>
+				</div>
+			{:else}
+				<div class="rounded bg-muted/20 p-2 text-xs text-muted-foreground/70 italic">No output</div>
+			{/if}
+		{/if}
+	{/snippet}
+</ToolCallBlock>

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

@@ -0,0 +1,167 @@
+<script lang="ts">
+	import { ICON_CLASS_DEFAULT, ICON_CLASS_SPIN } from '$lib/constants/css-classes';
+	import { Globe, Loader2 } from '@lucide/svelte';
+	import { CollapsibleContentBlock } from '$lib/components/app';
+	import * as HoverCard from '$lib/components/ui/hover-card';
+	import { AgenticSectionType } from '$lib/enums';
+	import { mcpStore } from '$lib/stores/mcp.svelte';
+	import {
+		extractSearchResults,
+		extractSearchQuery,
+		faviconForUrl,
+		sanitizeExternalUrl,
+		type SearchResult,
+		type AgenticSection
+	} from '$lib/utils';
+
+	interface Props {
+		section: AgenticSection;
+		open?: boolean;
+		isStreaming?: boolean;
+		onToggle?: () => void;
+	}
+
+	let { section, open = $bindable(false), isStreaming = false, onToggle }: Props = $props();
+
+	const isPending = $derived(section.type === AgenticSectionType.TOOL_CALL_PENDING);
+	const isStreamingCall = $derived(section.type === AgenticSectionType.TOOL_CALL_STREAMING);
+	const showSpinner = $derived(isPending || (isStreamingCall && isStreaming));
+
+	const results: SearchResult[] = $derived(extractSearchResults(section.toolResult));
+	const query = $derived(extractSearchQuery(section.toolArgs));
+
+	// Same icon-resolution chain as ChatMessageToolCallBlockDefault so
+	// MCP-server branding is consistent across both views. Spinner wins
+	// while the call is in flight so the user sees execution status.
+	const iconUrl = $derived(showSpinner ? null : mcpStore.getServerFaviconForTool(section.toolName));
+	const icon = $derived(showSpinner ? Loader2 : undefined);
+	const iconClass = $derived(showSpinner ? ICON_CLASS_SPIN : ICON_CLASS_DEFAULT);
+
+	// Verb reflects state: "Searching" while the call is in flight, "Searched"
+	// once results (or a definitive empty response) have arrived. Lets the
+	// heading read as a live progress indicator rather than a completed
+	// retrospective.
+	const title = $derived.by(() => {
+		const verb = showSpinner ? 'Searching' : 'Searched';
+		return query ? `${verb} web for "${query}"` : `${verb} web`;
+	});
+
+	function hideBrokenIcon(event: Event) {
+		(event.currentTarget as HTMLImageElement).style.display = 'none';
+	}
+
+	function formatPublishDate(iso: string | undefined): string | null {
+		if (!iso) return null;
+		try {
+			const date = new Date(iso);
+			if (Number.isNaN(date.getTime())) return iso;
+			return date.toLocaleDateString(undefined, {
+				year: 'numeric',
+				month: 'short',
+				day: 'numeric'
+			});
+		} catch {
+			return iso;
+		}
+	}
+
+	function hostFor(url: string): string | null {
+		try {
+			return new URL(url).host;
+		} catch {
+			return null;
+		}
+	}
+
+	function hasDetails(result: SearchResult): boolean {
+		return Boolean(result.highlights || result.published || result.author);
+	}
+</script>
+
+{#snippet pill(result: SearchResult)}
+	{@const faviconUrl = faviconForUrl(result.url)}
+	{@const safeUrl = sanitizeExternalUrl(result.url)}
+	{@const showHoverCard = safeUrl !== null && hasDetails(result)}
+	{#if safeUrl}
+		<HoverCard.Root openDelay={150} closeDelay={100}>
+			<HoverCard.Trigger
+				href={safeUrl}
+				target="_blank"
+				rel="noopener noreferrer"
+				class="hover:bg-muted/80 focus-visible:ring-ring inline-flex max-w-full items-center gap-1.5 rounded-full border bg-muted px-2.5 py-1 text-xs transition-colors outline-none focus-visible:ring-2"
+			>
+				{#if faviconUrl}
+					<img
+						src={faviconUrl}
+						alt=""
+						class="h-3 w-3 shrink-0 rounded-sm"
+						onerror={hideBrokenIcon}
+					/>
+				{:else}
+					<Globe class="text-muted-foreground/70 h-3 w-3 shrink-0" />
+				{/if}
+				<span class="truncate font-medium text-foreground/80">{result.title}</span>
+			</HoverCard.Trigger>
+			{#if showHoverCard}
+				{@const publishDate = formatPublishDate(result.published)}
+				{@const host = hostFor(safeUrl)}
+				<HoverCard.Content
+					side="top"
+					align="start"
+					sideOffset={6}
+					class="bg-popover text-popover-foreground z-50 w-80 max-w-[90vw] rounded-lg border p-0 shadow-lg"
+				>
+					<div class="flex flex-col gap-2 p-3">
+						<a
+							href={safeUrl}
+							target="_blank"
+							rel="noopener noreferrer"
+							class="line-clamp-3 text-sm font-medium leading-snug hover:underline"
+							>{result.title}</a
+						>
+						{#if publishDate || result.author}
+							<div class="text-muted-foreground flex items-center gap-1.5 text-[11px]">
+								{#if publishDate}
+									<span>{publishDate}</span>
+								{/if}
+								{#if publishDate && result.author}
+									<span class="opacity-50">&middot;</span>
+								{/if}
+								{#if result.author}
+									<span class="truncate">{result.author}</span>
+								{/if}
+							</div>
+						{/if}
+						{#if result.highlights}
+							<p
+								class="text-popover-foreground/85 line-clamp-5 text-xs leading-relaxed whitespace-pre-line"
+							>
+								{result.highlights}
+							</p>
+						{/if}
+						{#if host}
+							<div class="text-muted-foreground/80 truncate text-[11px]">{host}</div>
+						{/if}
+					</div>
+				</HoverCard.Content>
+			{/if}
+		</HoverCard.Root>
+	{/if}
+{/snippet}
+
+<CollapsibleContentBlock {open} class="my-2" {icon} {iconClass} {iconUrl} {title} {onToggle}>
+	{#if results.length > 0}
+		<div class="flex flex-wrap items-center gap-2 pb-1">
+			{#each results as result (result.url)}
+				{@render pill(result)}
+			{/each}
+		</div>
+	{:else if showSpinner}
+		<div class="text-muted-foreground/70 flex items-center gap-2 py-1 text-xs italic">
+			<Loader2 class="h-3 w-3 animate-spin" />
+			<span>Searching...</span>
+		</div>
+	{:else}
+		<div class="text-muted-foreground/70 py-1 text-xs italic">No results</div>
+	{/if}
+</CollapsibleContentBlock>

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

@@ -0,0 +1,55 @@
+<script lang="ts">
+	import { XCircle } from '@lucide/svelte';
+	import { SyntaxHighlightedCode } from '$lib/components/app';
+	import { MAX_HEIGHT_CODE_BLOCK, RESULT_STAT_SEPARATOR } from '$lib/constants';
+	import { type AgenticSection } from '$lib/utils';
+	import { parseWriteFileMeta } from './parsers/write-file';
+	import ToolCallBlock from './ToolCallBlock.svelte';
+
+	interface Props {
+		section: AgenticSection;
+		open: boolean;
+		isStreaming: boolean;
+		onToggle?: () => void;
+	}
+
+	let { section, open, isStreaming, onToggle }: Props = $props();
+
+	const writeFileMeta = $derived(parseWriteFileMeta(section));
+</script>
+
+<ToolCallBlock {section} {open} {isStreaming} meta={writeFileMeta} {onToggle}>
+	{#snippet titleSnippet()}
+		<span class="text-muted-foreground">Write file </span>
+		<span class="font-mono">{writeFileMeta?.filePath}</span>
+		{#if writeFileMeta?.errorMessage}
+			<span class="ml-1 text-xs italic text-muted-foreground/70">(failed)</span>
+		{/if}
+	{/snippet}
+
+	{#snippet children(meta, ctx)}
+		{#if meta?.errorMessage}
+			<div
+				class="flex items-start gap-2 rounded bg-red-500/10 p-2 text-xs text-red-600 italic dark:text-red-400"
+			>
+				<XCircle class="mt-0.5 h-3 w-3 shrink-0" />
+				<span>{meta.errorMessage}</span>
+			</div>
+		{:else if meta}
+			<SyntaxHighlightedCode
+				code={meta.content}
+				language={meta.language}
+				maxHeight={MAX_HEIGHT_CODE_BLOCK}
+				streaming={ctx.isCodeStreaming}
+			/>
+			<div class="mt-1.5 text-xs text-muted-foreground/70 italic">
+				{#if meta.resultMessage}
+					{meta.resultMessage}{meta.bytesWritten != null ? RESULT_STAT_SEPARATOR : ''}{/if}
+				{#if meta.bytesWritten != null}
+					<span class="font-mono">{meta.bytesWritten}</span>
+					bytes
+				{/if}
+			</div>
+		{/if}
+	{/snippet}
+</ToolCallBlock>

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

@@ -0,0 +1,129 @@
+<script lang="ts" generics="TMeta">
+	// Generic chrome shell shared by every per-tool block under
+	// `ChatMessageToolCall/`. Owns:
+	//   - the collapsible wrapper (defaults to CollapsibleContentBlock;
+	//     `exec_shell_command` swaps in CollapsibleTerminalBlock via the
+	//     `wrapper` prop);
+	//   - the icon, spinner state, and MCP favicon fallback chain;
+	//   - the status subtitle pill.
+	// Components supply only their `meta`, a title snippet, and a body
+	// snippet - everything around them is this single source of truth.
+
+	import { Loader2, Wrench } from '@lucide/svelte';
+	import { CollapsibleContentBlock } from '$lib/components/app';
+	import { ICON_CLASS_DEFAULT, ICON_CLASS_SPIN } from '$lib/constants/css-classes';
+	import { AgenticSectionType } from '$lib/enums';
+	import { getBuiltinToolUi } from '$lib/constants/built-in-tools';
+	import { mcpStore } from '$lib/stores/mcp.svelte';
+	import type { Component, Snippet } from 'svelte';
+	import type { AgenticSection, BuiltinToolUiEntry } from '$lib/utils';
+
+	type ToolCallBlockMetaWithError = TMeta & { errorMessage?: string };
+
+	interface ToolCallCtx {
+		isStreaming: boolean;
+		isPending: boolean;
+		isStreamingCall: boolean;
+		isCodeStreaming: boolean;
+	}
+
+	interface Props {
+		section: AgenticSection;
+		open: boolean;
+		isStreaming: boolean;
+		/**
+		 * The per-tool meta, including any `errorMessage` field that the
+		 * shared chrome uses to compute the status pill subtitle.
+		 */
+		meta: ToolCallBlockMetaWithError | null | undefined;
+		/**
+		 * True while the tool's process is actively producing output
+		 * chunks after its args finished streaming (used by
+		 * `exec_shell_command`'s stdout feed).
+		 */
+		extraLiveStreaming?: boolean;
+		/**
+		 * Swap the title-row icon for a spinning `Loader2` while the
+		 * spinner is showing. Only meaningful for tools where "live"
+		 * is interesting (e.g. exec_shell_command showing the in-flight
+		 * process). Other tools leave it off and render the spinner
+		 * inline within the body.
+		 */
+		spinIconWhenActive?: boolean;
+		/**
+		 * Wrapper component that renders the title row and the body
+		 * children. Defaults to CollapsibleContentBlock;
+		 * `exec_shell_command` uses CollapsibleTerminalBlock for its
+		 * terminal-style frame.
+		 */
+		wrapper?: typeof CollapsibleContentBlock;
+		title?: string;
+		titleSnippet?: Snippet;
+		onToggle?: () => void;
+		children: Snippet<[TMeta | null | undefined, ToolCallCtx]>;
+	}
+
+	let {
+		section,
+		open,
+		isStreaming,
+		meta,
+		extraLiveStreaming = false,
+		spinIconWhenActive = false,
+		wrapper: Wrapper = CollapsibleContentBlock,
+		title,
+		titleSnippet,
+		onToggle,
+		children
+	}: Props = $props();
+
+	const isPending = $derived(section.type === AgenticSectionType.TOOL_CALL_PENDING);
+	const isStreamingCall = $derived(section.type === AgenticSectionType.TOOL_CALL_STREAMING);
+	const showSpinner = $derived(isPending || (isStreamingCall && isStreaming) || extraLiveStreaming);
+	const isCodeStreaming = $derived(isStreaming && (isPending || isStreamingCall));
+
+	const toolUi: BuiltinToolUiEntry | null = $derived(getBuiltinToolUi(section.toolName));
+	const toolIcon: Component = $derived(
+		spinIconWhenActive && showSpinner ? Loader2 : (toolUi?.icon ?? Wrench)
+	);
+	const toolIconClass = $derived(
+		spinIconWhenActive && showSpinner ? ICON_CLASS_SPIN : ICON_CLASS_DEFAULT
+	);
+	// Drop the MCP favicon while the spinner is on so the title row
+	// signals "in flight" without being overwritten by server branding.
+	const mcpServerFavicon = $derived(
+		showSpinner ? null : mcpStore.getServerFaviconForTool(section.toolName)
+	);
+	const iconUrl = $derived(
+		showSpinner || (toolUi?.icon ?? null) || !mcpServerFavicon ? null : mcpServerFavicon
+	);
+
+	function subtitleFor(errorMessage?: string): string | undefined {
+		if (extraLiveStreaming) return 'streaming...';
+		if (showSpinner) return 'executing...';
+		if (errorMessage) return 'failed';
+		if (isStreamingCall && !isStreaming) return 'incomplete';
+		return undefined;
+	}
+
+	const subtitle = $derived(subtitleFor(meta?.errorMessage));
+</script>
+
+<Wrapper
+	{open}
+	class="my-2"
+	icon={toolIcon}
+	iconClass={toolIconClass}
+	{iconUrl}
+	{title}
+	{titleSnippet}
+	{subtitle}
+	{onToggle}
+>
+	{@render children(meta, {
+		isStreaming,
+		isPending,
+		isStreamingCall,
+		isCodeStreaming
+	})}
+</Wrapper>

tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/parsers/_shared.ts

@@ -0,0 +1,49 @@
+// Helpers shared by the per-tool meta parsers under
+// `src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/parsers/`.
+// Each tool needs the same first three steps (tool-name check,
+// args-present check, JSON parse) - keeping them here lets each parser
+// stay focused on its own format quirks.
+
+import { BuiltInTool } from '$lib/enums';
+import { parsePartialJsonArgs } from '$lib/utils/parse-partial-json-args';
+import type { AgenticSection } from '$lib/utils/agentic';
+
+/**
+ * Strict (final-state) JSON parser for a tool-args blob. Mirrors the
+ * behaviour the per-tool components used before extraction: an
+ * invalid JSON blob, a JSON array, or a JSON primitive all map to
+ * `null` so callers don't have to guard against surprise shapes.
+ */
+function parseFinalToolArgs(blob: string): Record<string, unknown> | null {
+	try {
+		const parsed: unknown = JSON.parse(blob);
+		if (parsed && typeof parsed === 'object' && !Array.isArray(parsed)) {
+			return parsed as Record<string, unknown>;
+		}
+		return null;
+	} catch {
+		return null;
+	}
+}
+
+/**
+ * Parse a section's toolArgs against an expected tool name. Returns
+ * `null` when:
+ *   - the section's toolName doesn't match (component isn't for this
+ *     tool);
+ *   - the section has no args yet (call hasn't started streaming);
+ *   - or the args blob can't be parsed.
+ *
+ * Pass `{ partial: true }` for tools that need to render incrementally
+ * as each token lands (read_file, edit_file, write_file).
+ */
+export function parseToolArgs(
+	expected: BuiltInTool,
+	section: AgenticSection,
+	options: { partial?: boolean } = {}
+): Record<string, unknown> | null {
+	if (section.toolName !== expected || !section.toolArgs) return null;
+	return options.partial
+		? parsePartialJsonArgs(section.toolArgs)
+		: parseFinalToolArgs(section.toolArgs);
+}

tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/parsers/edit-file.ts

@@ -0,0 +1,71 @@
+// Meta parser for `edit_file` tool calls. Reads the file path and the
+// array of edits from the streamed args (partial JSON for incremental
+// rendering), plus the result blob for `result` / `edits_applied` /
+// `error` fields.
+
+import { BuiltInTool } from '$lib/enums';
+import { FILE_PATH_SEPARATOR_REGEX } from '$lib/constants';
+import { tryParseToolResultObject, type AgenticSection } from '$lib/utils';
+import { parseToolArgs } from './_shared';
+
+export type EditFileEdit = {
+	oldText: string;
+	newText: string;
+};
+
+export type EditFileMeta = {
+	fileName: string;
+	filePath: string;
+	edits: EditFileEdit[];
+	resultMessage?: string;
+	editsApplied?: number;
+	errorMessage?: string;
+};
+
+export function parseEditFileMeta(section: AgenticSection): EditFileMeta | null {
+	const args = parseToolArgs(BuiltInTool.EDIT_FILE, section, { partial: true });
+	if (!args) return null;
+
+	const rawPath = args.path ?? args.file_path ?? args.filePath;
+	if (typeof rawPath !== 'string' || !rawPath) return null;
+
+	const fileName = rawPath.split(FILE_PATH_SEPARATOR_REGEX).pop() || rawPath;
+
+	// Filter the streamed edits array strictly: each entry must be an
+	// object with a non-empty `old_text`. Edits without an old_text
+	// would diff against empty and render as a full re-write.
+	const rawEdits = Array.isArray(args.edits) ? args.edits : [];
+	const edits: EditFileEdit[] = [];
+	for (const e of rawEdits) {
+		if (!e || typeof e !== 'object' || Array.isArray(e)) continue;
+		const obj = e as Record<string, unknown>;
+		const oldText = typeof obj.old_text === 'string' ? obj.old_text : '';
+		if (!oldText) continue;
+		const newText = typeof obj.new_text === 'string' ? obj.new_text : '';
+		edits.push({ oldText, newText });
+	}
+
+	const resultObj = tryParseToolResultObject(section.toolResult);
+	let resultMessage: string | undefined;
+	let editsApplied: number | undefined;
+	let errorMessage: string | undefined;
+	if (typeof resultObj?.error === 'string') {
+		errorMessage = resultObj.error;
+	} else if (resultObj) {
+		if (typeof resultObj.result === 'string') {
+			resultMessage = resultObj.result;
+		}
+		if (Number.isFinite(Number(resultObj.edits_applied))) {
+			editsApplied = Number(resultObj.edits_applied);
+		}
+	}
+
+	return {
+		fileName,
+		filePath: rawPath,
+		edits,
+		resultMessage,
+		editsApplied,
+		errorMessage
+	};
+}