[SYCL] Flash Attention with XMX engine via oneDNN (#25222)

* [SYCL] F16 (default) Flash Attention with XMX engine via oneDNN graph API; Qwen3.6-27b-Q8_0 prefill speed up x1.21 at p=512 and x4.26 at p=80k

* [SYCL] Address review on FA oneDNN path. Result: llama-bench---pp512; 32% increase with fa1; llama-perplexity---0.11% difference; tested model: mradermacher/Meta-Llama-3.1-8B-Instruct-Q8_0.gguf

* PR-25222 revision v2: addressed audits

* [SYCL] flash-attn oneDNN SDPA KV F16 rev 3.0: add BMG gate + multi-device sync. Narrow the scrope of this PR to Battlemage only (bmg; Xe2). Other archs (e.g., alchemist) fall back to existing FA kernel. When device_count >1, apply stream -> wait_and_throw(), validated working path for multi-gpu sync fix by @maxious.

Co-authored-by: maxious <81432+maxious@users.noreply.github.com>

* updated comment on bmg gate, noted the issue

---------

Co-authored-by: scientist3 <scientist.3@users.noreply.github.com>
Co-authored-by: hmscider <hmscider@users.noreply.github.com>
Co-authored-by: maxious <81432+maxious@users.noreply.github.com>

common/arg.cpp

@@ -697,7 +697,7 @@ static bool common_params_parse_ex(int argc, char ** argv, common_params_context
         }
     };
 
-    // parse the first time to get -hf option (used for remote preset)
+    // parse all CLI args now, so that -hf is available below for remote preset resolution
     parse_cli_args();
 
     postprocess_cpu_params(params.cpuparams,       nullptr);
@@ -748,6 +748,11 @@ static bool common_params_parse_ex(int argc, char ** argv, common_params_context
         params.kv_overrides.back().key[0] = 0;
     }
 
+    if (!params.server_tools.empty() && !params.cors_origins_explicit) {
+        LOG_WRN("server tools are enabled, using localhost as default CORS origin (change via --cors-origins)\n");
+        params.cors_origins = "localhost";
+    }
+
     // pad tensor_buft_overrides for llama_params_fit:
     const size_t ntbo = llama_max_tensor_buft_overrides();
     while (params.tensor_buft_overrides.size() < ntbo) {
@@ -3047,6 +3052,42 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
             params.public_path = value;
         }
     ).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_STATIC_PATH"));
+    add_opt(common_arg(
+        {"--cors-origins"}, "ORIGINS",
+        string_format(
+            "comma-separated list of allowed origins for CORS (default: %s)\n"
+            "if set to special value 'localhost', reflect the Origin header only if it is localhost",
+        params.cors_origins.c_str()),
+        [](common_params & params, const std::string & value) {
+            params.cors_origins = value;
+            params.cors_origins_explicit = true;
+        }
+    ).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_CORS_ORIGINS"));
+    add_opt(common_arg(
+        {"--cors-methods"}, "METHODS",
+        string_format("comma-separated list of allowed methods for CORS (default: %s)", params.cors_methods.c_str()),
+        [](common_params & params, const std::string & value) {
+            params.cors_methods = value;
+        }
+    ).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_CORS_METHODS"));
+    add_opt(common_arg(
+        {"--cors-headers"}, "HEADERS",
+        string_format("comma-separated list of allowed headers for CORS (default: %s)", params.cors_headers.c_str()),
+        [](common_params & params, const std::string & value) {
+            params.cors_headers = value;
+        }
+    ).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_CORS_HEADERS"));
+    add_opt(common_arg(
+        {"--cors-credentials"},
+        {"--no-cors-credentials"},
+        string_format(
+            "whether to allow credentials for CORS (default: %s)\n"
+            "note: if this is enabled and --cors-origins is set to * (default), the Origin header will be echoed back, and credentials will always be allowed",
+        params.cors_credentials ? "enabled" : "disabled"),
+        [](common_params & params, bool value) {
+            params.cors_credentials = value;
+        }
+    ).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_CORS_CREDENTIALS"));
     add_opt(common_arg(
         {"--api-prefix"}, "PREFIX",
         string_format("prefix path the server serves from, without the trailing slash (default: %s)", params.api_prefix.c_str()),
@@ -3080,7 +3121,8 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
         {"--tools"}, "TOOL1,TOOL2,...",
         "experimental: whether to enable built-in tools for AI agents - do not enable in untrusted environments (default: no tools)\n"
         "specify \"all\" to enable all tools\n"
-        "available tools: read_file, file_glob_search, grep_search, exec_shell_command, write_file, edit_file, get_datetime",
+        "available tools: read_file, file_glob_search, grep_search, exec_shell_command, write_file, edit_file, get_datetime\n"
+        "note: for security reasons, this will limit --cors-origins to localhost by default",
         [](common_params & params, const std::string & value) {
             params.server_tools = parse_csv_row(value);
         }
@@ -3088,7 +3130,8 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
     add_opt(common_arg(
         {"-ag", "--agent"},
         {"-no-ag", "--no-agent"},
-        "whether to enable CORS proxy and all built-in tools - do not enable in untrusted environments (default: disabled)",
+        "whether to enable CORS proxy and all built-in tools - do not enable in untrusted environments (default: disabled)\n"
+        "note: for security reasons, this will limit --cors-origins to localhost by default",
         [](common_params & params, bool value) {
             if (value) {
                 params.server_tools = {"all"};
@@ -3097,6 +3140,7 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
                 params.server_tools.clear();
                 params.ui_mcp_proxy = false;
             }
+            // note: do not modify cors_origins here, as the options are not evaluated in order (user may explicitly set --cors-origins before --agent)
         }
     ).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_AGENT"));
     add_opt(common_arg(

common/common.h

@@ -631,6 +631,14 @@ struct common_params {
     std::string api_prefix    = "";                                                                         // NOLINT
     std::string chat_template = "";                                                                         // NOLINT
     bool use_jinja = true;                                                                                  // NOLINT
+
+    // server CORS params
+    std::string cors_origins = "*";
+    std::string cors_methods = "GET, POST, DELETE, OPTIONS";
+    std::string cors_headers = "*";
+    bool cors_credentials = true;
+    bool cors_origins_explicit = false; // for --agent option
+
     bool enable_chat_template = true;
     bool force_pure_content_parser = false;
     common_reasoning_format reasoning_format = COMMON_REASONING_FORMAT_DEEPSEEK;

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-opencl/ggml-opencl.cpp

@@ -10525,10 +10525,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/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-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/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/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);

tools/server/server-context.cpp

@@ -220,8 +220,6 @@ struct server_slot {
             return false;
         }
 
-        GGML_ASSERT(prompt.data.size() == 0);
-
         const size_t cur_size_tgt =           llama_state_seq_get_size_ext(ctx_tgt, id, LLAMA_STATE_SEQ_FLAGS_NONE);
         const size_t cur_size_dft = ctx_dft ? llama_state_seq_get_size_ext(ctx_dft, id, LLAMA_STATE_SEQ_FLAGS_NONE) : 0;
 
@@ -252,11 +250,7 @@ struct server_slot {
         return res;
     }
 
-    void prompt_clear(bool allow_processing) {
-        if (!allow_processing) {
-            GGML_ASSERT(!is_processing());
-        }
-
+    void prompt_clear() {
         SLT_TRC(*this, "clearing prompt with %zu tokens\n", prompt.tokens.size());
 
         common_context_seq_rm(ctx_tgt, id, -1, -1);
@@ -264,7 +258,7 @@ struct server_slot {
             common_context_seq_rm(ctx_dft, id, -1, -1);
         }
 
-        prompt.tokens.clear();
+        prompt.clear();
     }
 
     std::vector<common_adapter_lora_info> lora;
@@ -493,7 +487,7 @@ struct server_slot {
 
             // do not keep context of the child slots - the parent's context is enough
             if (task->is_child()) {
-                prompt_clear(false);
+                prompt_clear();
             }
 
             reset();
@@ -1626,7 +1620,7 @@ private:
                 ret->prompt_save(*prompt_cache);
 
                 if (!ret->prompt_load(*prompt_cache, task.tokens)) {
-                    ret->prompt_clear(false);
+                    ret->prompt_clear();
                 }
 
                 prompt_cache->update();
@@ -1658,7 +1652,7 @@ private:
             if (slot.prompt.n_tokens() > 0) {
                 SRV_WRN("purging slot %d with %zu tokens\n", slot.id, slot.prompt.tokens.size());
 
-                slot.prompt_clear(false);
+                slot.prompt_clear();
 
                 res = true;
 
@@ -1691,7 +1685,7 @@ private:
                 // if lora has changed, check to see if the cache should be cleared
                 if (lora_should_clear_cache(slot.lora, task_loras)) {
                     SLT_TRC(slot, "clearing cache for lora change. %zu loras -> %zu loras\n", slot.lora.size(), task.params.lora.size());
-                    slot.prompt.tokens.clear();
+                    slot.prompt.clear();
                 } else {
                     SLT_TRC(slot, "keeping cache for alora. %zu target loras\n", task_loras.size());
                 }
@@ -2405,7 +2399,7 @@ private:
 
                                 if (params_base.kv_unified) {
                                     // [TAG_IDLE_SLOT_CLEAR]
-                                    slot.prompt_clear(false);
+                                    slot.prompt_clear();
                                 }
                             }
                         }
@@ -2573,12 +2567,12 @@ private:
                     size_t token_count = 0;
                     size_t nread = llama_state_seq_load_file(ctx_tgt, filepath.c_str(), slot->id, tokens.data(), tokens.size(), &token_count);
                     if (nread == 0) {
-                        slot->prompt.tokens.clear(); // KV may already been invalidated?
+                        slot->prompt.clear(); // KV may already been invalidated?
                         send_error(task, "Unable to restore slot, no available space in KV cache or invalid slot save file", ERROR_TYPE_INVALID_REQUEST);
                         break;
                     }
                     tokens.resize(token_count);
-                    slot->prompt.tokens.clear();
+                    slot->prompt.clear();
                     slot->prompt.tokens.insert(tokens);
 
                     const int64_t t_end = ggml_time_us();
@@ -2615,7 +2609,7 @@ private:
                     // Erase token cache
                     const size_t n_erased = slot->prompt.tokens.size();
 
-                    slot->prompt_clear(false);
+                    slot->prompt_clear();
 
                     auto res = std::make_unique<server_task_result_slot_erase>();
                     res->id       = task.id;
@@ -2775,6 +2769,27 @@ private:
             abort_all_slots("pre_decode() failed: " + std::string(e.what()));
         }
 
+        GGML_ASSERT(batch.slot_batched || batch.size() == 0);
+
+        if (batch.slot_batched) {
+            auto & slot_batched      = batch.slot_batched;
+            auto & alora_scale       = batch.alora_scale;
+            auto & alora_disabled_id = batch.alora_disabled_id;
+
+            // TODO @ngxson : alora handling is too messy, need to refactor it to be more clear and maintainable
+            // apply lora, only need to do it once per batch
+            common_set_adapter_lora(ctx_tgt, slot_batched->lora);
+
+            // if the lora is temporarily disabled for an alora, re-enable it
+            // for next time
+            if (alora_scale > 0.0f) {
+                SRV_DBG("re-enabling alora with scale %f\n", alora_scale);
+                slot_batched->lora[alora_disabled_id].scale = alora_scale;
+            }
+
+            llama_set_embeddings(ctx_tgt, slot_batched->need_embd());
+        }
+
         llama_batch batch_view;
         int32_t off_next = 0;
         int32_t n_batch = llama_n_batch(ctx_tgt);
@@ -2814,7 +2829,6 @@ private:
                 abort_all_slots("post_decode() failed: " + std::string(e.what()));
                 break; // stop any further processing
             }
-
         }
     }
 
@@ -2880,7 +2894,7 @@ private:
 
                     new_tokens.resize(slot.prompt.tokens.size() - n_discard);
 
-                    slot.prompt.tokens.clear();
+                    slot.prompt.clear();
                     slot.prompt.tokens.insert(new_tokens);
                 }
 
@@ -3556,25 +3570,6 @@ private:
     bool decode(int32_t & n_batch, int32_t off, llama_batch & batch_view) {
         SRV_DBG("n_batch (effective) = %d, off = %d\n", n_batch, off);
 
-        auto & slot_batched      = batch.slot_batched;
-        auto & alora_scale       = batch.alora_scale;
-        auto & alora_disabled_id = batch.alora_disabled_id;
-
-        // TODO @ngxson : alora handling is too messy, need to refactor it to be more clear and maintainable
-        if (slot_batched) {
-            // apply lora, only need to do it once per batch
-            common_set_adapter_lora(ctx_tgt, slot_batched->lora);
-
-            // if the lora is temporarily disabled for an alora, re-enable it
-            // for next time
-            if (alora_scale > 0.0f) {
-                SRV_DBG("re-enabling alora with scale %f\n", alora_scale);
-                slot_batched->lora[alora_disabled_id].scale = alora_scale;
-            }
-
-            llama_set_embeddings(ctx_tgt, slot_batched->need_embd());
-        }
-
         if (batch.size() == 0) {
             SRV_WRN("%s", "no tokens to decode\n");
 
@@ -3622,7 +3617,7 @@ private:
 
                             // note: it's complicated to keep track of how much of the current batch has been
                             //       processed before the error occurred, so we simply clear the entire context
-                            slot.prompt_clear(false);
+                            slot.prompt_clear();
                         }
                     }
 

tools/server/server-http.cpp

@@ -47,6 +47,16 @@ static void log_server_request(const httplib::Request & req, const httplib::Resp
     SRV_DBG("response: %s\n", res.body.c_str());
 }
 
+// returns true if the Origin header value's host is localhost / 127.0.0.1 / ::1 (any port)
+static bool origin_is_localhost(const std::string & origin) {
+    try {
+        const std::string host = common_http_parse_url(origin).host;
+        return host == "localhost" || host == "127.0.0.1" || host == "::1";
+    } catch (const std::exception &) {
+        return false;
+    }
+}
+
 // For Google Cloud Platform deployment compatibility
 struct gcp_params {
     bool enabled;
@@ -266,13 +276,26 @@ bool server_http_context::init(const common_params & params) {
     };
 
     // register server middlewares
-    srv->set_pre_routing_handler([middleware_validate_api_key, middleware_server_state](const httplib::Request & req, httplib::Response & res) {
-        res.set_header("Access-Control-Allow-Origin", req.get_header_value("Origin"));
+    srv->set_pre_routing_handler([&params, middleware_validate_api_key, middleware_server_state](const httplib::Request & req, httplib::Response & res) {
+        if (params.cors_credentials && params.cors_origins == "*") {
+            // special case: echo back the Origin header to allow any origin to access the server with credentials
+            res.set_header("Access-Control-Allow-Origin", req.get_header_value("Origin"));
+        } else if (params.cors_origins == "localhost") {
+            // special case: only reflect the Origin header if it is a localhost origin
+            std::string origin = req.get_header_value("Origin");
+            if (origin_is_localhost(origin)) {
+                res.set_header("Access-Control-Allow-Origin", origin);
+            } else {
+                SRV_WRN("(CORS) skip non-localhost origin: %s\n", origin.c_str());
+            }
+        } else {
+            res.set_header("Access-Control-Allow-Origin", params.cors_origins);
+        }
         // If this is OPTIONS request, skip validation because browsers don't include Authorization header
         if (req.method == "OPTIONS") {
-            res.set_header("Access-Control-Allow-Credentials", "true");
-            res.set_header("Access-Control-Allow-Methods",     "GET, POST");
-            res.set_header("Access-Control-Allow-Headers",     "*");
+            res.set_header("Access-Control-Allow-Credentials", params.cors_credentials ? "true" : "false");
+            res.set_header("Access-Control-Allow-Methods",     params.cors_methods);
+            res.set_header("Access-Control-Allow-Headers",     params.cors_headers);
             res.set_content("", "text/html"); // blank response, no data
             return httplib::Server::HandlerResponse::Handled; // skip further processing
         }

tools/server/server-task.cpp

@@ -1646,16 +1646,16 @@ size_t server_prompt_cache::n_tokens() const {
     size_t res = 0;
 
     for (const auto & state : states) {
-        res += state.n_tokens();
+        res += state.prompt.n_tokens();
     }
 
     return res;
 }
 
-server_prompt * server_prompt_cache::alloc(const server_prompt & prompt, size_t state_size_tgt, size_t state_size_dft) {
+server_prompt_cache_state * server_prompt_cache::alloc(const server_prompt & prompt, size_t state_size_tgt, size_t state_size_dft) {
     // first check if the current state is contained fully in the cache
     for (auto it = states.begin(); it != states.end(); ++it) {
-        const int cur_lcp_len = it->tokens.get_common_prefix(prompt.tokens);
+        const int cur_lcp_len = it->prompt.tokens.get_common_prefix(prompt.tokens);
 
         if (cur_lcp_len == (int) prompt.tokens.size()) {
             SRV_TRC("%s", " - prompt is already in the cache, skipping\n");
@@ -1680,9 +1680,9 @@ server_prompt * server_prompt_cache::alloc(const server_prompt & prompt, size_t
 
     // remove any cached prompts that are fully contained in the current prompt
     for (auto it = states.begin(); it != states.end();) {
-        const int len = it->tokens.get_common_prefix(prompt.tokens);
+        const int len = it->prompt.tokens.get_common_prefix(prompt.tokens);
 
-        if (len == (int) it->tokens.size()) {
+        if (len == (int) it->prompt.tokens.size()) {
             SRV_TRC(" - removing obsolete cached prompt with length %d\n", len);
 
             it = states.erase(it);
@@ -1721,12 +1721,14 @@ server_prompt * server_prompt_cache::alloc(const server_prompt & prompt, size_t
     }
 
     states.push_back({
-        /*.tokens      =*/ prompt.tokens.clone(),
-        /*.data        =*/ {
+        /*.prompt =*/ {
+            /*.tokens      =*/ prompt.tokens.clone(),
+            /*.checkpoints =*/ prompt.checkpoints,
+        },
+        /*.data   =*/ {
             /*.main =*/ std::move(state_data_tgt),
             /*.drft =*/ std::move(state_data_dft),
         },
-        /*.checkpoints =*/ prompt.checkpoints,
     });
 
     return &states.back();
@@ -1744,9 +1746,9 @@ bool server_prompt_cache::load(server_prompt & prompt, const server_tokens & tok
 
     // find the most similar cached prompt, that would also preserve the most context
     for (auto it = states.begin(); it != states.end(); ++it) {
-        const int lcp_cur = it->tokens.get_common_prefix(tokens_new);
+        const int lcp_cur = it->prompt.tokens.get_common_prefix(tokens_new);
 
-        const float f_keep_cur = float(lcp_cur) / it->tokens.size();
+        const float f_keep_cur = float(lcp_cur) / it->prompt.tokens.size();
         const float sim_cur    = float(lcp_cur) / tokens_new.size();
 
         // don't trash large prompts
@@ -1799,7 +1801,7 @@ bool server_prompt_cache::load(server_prompt & prompt, const server_tokens & tok
             }
         }
 
-        prompt = std::move(*it_best);
+        prompt = std::move(it_best->prompt);
 
         states.erase(it_best);
     }
@@ -1836,6 +1838,6 @@ void server_prompt_cache::update() {
 
     for (const auto & state : states) {
         SRV_TRC("   - prompt %p: %7d tokens, checkpoints: %2zu, %9.3f MiB\n",
-                (const void *)&state, state.n_tokens(), state.checkpoints.size(), state.size() / (1024.0 * 1024.0));
+                (const void *)&state, state.prompt.n_tokens(), state.prompt.checkpoints.size(), state.size() / (1024.0 * 1024.0));
     }
 }

tools/server/server-task.h

@@ -584,32 +584,14 @@ struct server_task_result_apply_lora : server_task_result {
     virtual json to_json() override;
 };
 
-struct server_prompt_data {
-    std::vector<uint8_t> main;
-    std::vector<uint8_t> drft;
-
-    size_t size() const {
-        return main.size() + drft.size();
-    }
-};
-
 struct server_prompt {
     server_tokens tokens;
 
-    server_prompt_data data;
-
     std::list<common_prompt_checkpoint> checkpoints;
 
-    size_t size() const {
-        size_t res = 0;
-
-        res += data.size();
-
-        for (const auto & ckpt : checkpoints) {
-            res += ckpt.size();
-        }
-
-        return res;
+    void clear() {
+        tokens.clear();
+        checkpoints.clear();
     }
 
     int n_tokens() const {
@@ -619,19 +601,42 @@ struct server_prompt {
     server_prompt clone() const {
         return server_prompt {
             tokens.clone(),
-            data,
             checkpoints,
         };
     }
 };
 
+struct server_prompt_data {
+    std::vector<uint8_t> main;
+    std::vector<uint8_t> drft;
+
+    size_t size() const {
+        return main.size() + drft.size();
+    }
+};
+
+struct server_prompt_cache_state {
+    server_prompt prompt;
+    server_prompt_data data;
+
+    size_t size() const {
+        size_t res = data.size();
+
+        for (const auto & ckpt : prompt.checkpoints) {
+            res += ckpt.size();
+        }
+
+        return res;
+    }
+};
+
 struct server_prompt_cache {
     server_prompt_cache(int32_t limit_size_mib, size_t limit_tokens) {
         this->limit_size   = 1024ull*1024ull*(limit_size_mib < 0 ? 0 : limit_size_mib);
         this->limit_tokens = limit_tokens;
     }
 
-    std::list<server_prompt> states;
+    std::list<server_prompt_cache_state> states;
 
     // in bytes, 0 = no limit
     size_t limit_size = 0;
@@ -643,7 +648,7 @@ struct server_prompt_cache {
 
     size_t n_tokens() const;
 
-    server_prompt * alloc(const server_prompt & prompt, size_t state_size_main, size_t state_size_drft);
+    server_prompt_cache_state * alloc(const server_prompt & prompt, size_t state_size_main, size_t state_size_drft);
 
     bool load(server_prompt & prompt, const server_tokens & tokens_new, llama_context * ctx_main, llama_context * ctx_drft, int32_t id_slot);
 

tools/server/server.cpp

@@ -303,14 +303,24 @@ int llama_server(common_params & params, int argc, char ** argv) {
         return res;
     };
 
+    if (params.cors_origins == "*" && params.api_keys.empty()) {
+        SRV_WRN("%s", "-----------------\n");
+        SRV_WRN("%s", "CORS is set to allow all origins ('*') and no API key is set\n");
+        SRV_WRN("%s", "this can be a security risk (cross-origin attacks)\n");
+        SRV_WRN("%s", "more info: https://github.com/ggml-org/llama.cpp/pull/25655\n");
+        SRV_WRN("%s", "-----------------\n");
+    }
+
     // CORS proxy (EXPERIMENTAL, only used by the Web UI for MCP)
+    std::vector<std::string> warn_names;
+    if (is_router_server) {
+        warn_names.push_back("router mode");
+    }
+
     if (params.ui_mcp_proxy) {
-        SRV_WRN("%s", "-----------------\n");
-        SRV_WRN("%s", "CORS proxy is enabled, do not expose server to untrusted environments\n");
-        SRV_WRN("%s", "This feature is EXPERIMENTAL and may be removed or changed in future versions\n");
-        SRV_WRN("%s", "-----------------\n");
         ctx_http.get ("/cors-proxy",      ex_wrapper(proxy_handler_get));
         ctx_http.post("/cors-proxy",      ex_wrapper(proxy_handler_post));
+        warn_names.push_back("MCP proxy (experimental)");
     } else {
         ctx_http.get ("/cors-proxy",      ex_wrapper(res_403));
         ctx_http.post("/cors-proxy",      ex_wrapper(res_403));
@@ -324,17 +334,24 @@ int llama_server(common_params & params, int argc, char ** argv) {
             SRV_ERR("tools setup failed: %s\n", e.what());
             return 1;
         }
-        SRV_WRN("%s", "-----------------\n");
-        SRV_WRN("%s", "Built-in tools are enabled, do not expose server to untrusted environments\n");
-        SRV_WRN("%s", "This feature is EXPERIMENTAL and may be changed in the future\n");
-        SRV_WRN("%s", "-----------------\n");
         ctx_http.get ("/tools",           ex_wrapper(tools.handle_get));
         ctx_http.post("/tools",           ex_wrapper(tools.handle_post));
+        warn_names.push_back("built-in tools (experimental)");
     } else {
         ctx_http.get ("/tools",           ex_wrapper(res_403));
         ctx_http.post("/tools",           ex_wrapper(res_403));
     }
 
+    if (warn_names.size() > 0) {
+        SRV_WRN("%s", "-----------------\n");
+        SRV_WRN("%s", "the following feature(s) are enabled:\n");
+        for (const auto & name : warn_names) {
+            SRV_WRN("    %s\n", name.c_str());
+        }
+        SRV_WRN("%s", "do not expose the server to untrusted environments\n");
+        SRV_WRN("%s", "-----------------\n");
+    }
+
     //
     // Handle downloading model
     //
@@ -452,9 +469,6 @@ int llama_server(common_params & params, int argc, char ** argv) {
     SRV_INF("listening on %s\n", ctx_http.listening_address.c_str());
 
     if (is_router_server) {
-        SRV_WRN("%s", "NOTE: router mode is experimental\n");
-        SRV_WRN("%s", "      it is not recommended to use this mode in untrusted environments\n");
-
         if (!params.models_preset_hf.empty()) {
             SRV_WRN(      "NOTE: using preset.ini from HF repo '%s'\n", params.models_preset_hf.c_str());
             SRV_WRN("%s", "      please only use presets that you can trust! Unknown presets may be unsafe\n");

tools/server/tests/unit/test_security.py

@@ -91,7 +91,7 @@ def test_openai_library_correct_api_key():
     ("localhost", "Access-Control-Allow-Origin", "localhost"),
     ("web.mydomain.fr", "Access-Control-Allow-Origin", "web.mydomain.fr"),
     ("origin", "Access-Control-Allow-Credentials", "true"),
-    ("web.mydomain.fr", "Access-Control-Allow-Methods", "GET, POST"),
+    ("web.mydomain.fr", "Access-Control-Allow-Methods", "GET, POST, DELETE, OPTIONS"),
     ("web.mydomain.fr", "Access-Control-Allow-Headers", "*"),
 ])
 def test_cors_options(origin: str, cors_header: str, cors_header_value: str):
@@ -107,6 +107,70 @@ def test_cors_options(origin: str, cors_header: str, cors_header_value: str):
     assert res.headers[cors_header] == cors_header_value
 
 
+@pytest.mark.parametrize("origin", [
+    "http://localhost",
+    "http://localhost:8080",
+    "http://127.0.0.1",
+    "http://127.0.0.1:3000",
+    "http://[::1]",
+    "http://[::1]:3000",
+])
+def test_cors_origins_localhost_reflects(origin: str):
+    global server
+    server = ServerPreset.router()
+    server.cors_origins = "localhost"
+    server.start()
+    res = server.make_request("OPTIONS", "/completions", headers={
+        "Origin": origin,
+        "Access-Control-Request-Method": "POST",
+        "Access-Control-Request-Headers": "Authorization",
+    })
+    assert res.status_code == 200
+    assert res.headers["Access-Control-Allow-Origin"] == origin
+
+
+@pytest.mark.parametrize("origin", [
+    "http://web.mydomain.fr",
+    "http://evil.com",
+    "http://notlocalhost",
+    "http://localhost.evil.com",
+])
+def test_cors_origins_localhost_rejects(origin: str):
+    global server
+    server = ServerPreset.router()
+    server.cors_origins = "localhost"
+    server.start()
+    res = server.make_request("OPTIONS", "/completions", headers={
+        "Origin": origin,
+        "Access-Control-Request-Method": "POST",
+        "Access-Control-Request-Headers": "Authorization",
+    })
+    assert res.status_code == 200
+    assert "Access-Control-Allow-Origin" not in res.headers
+
+
+def test_cors_origins_defaults_to_localhost_with_tools_enabled():
+    global server
+    server = ServerPreset.router()
+    server.server_tools = "all"
+    server.start()
+    res = server.make_request("OPTIONS", "/completions", headers={
+        "Origin": "http://localhost:8080",
+        "Access-Control-Request-Method": "POST",
+        "Access-Control-Request-Headers": "Authorization",
+    })
+    assert res.status_code == 200
+    assert res.headers["Access-Control-Allow-Origin"] == "http://localhost:8080"
+
+    res = server.make_request("OPTIONS", "/completions", headers={
+        "Origin": "http://evil.com",
+        "Access-Control-Request-Method": "POST",
+        "Access-Control-Request-Headers": "Authorization",
+    })
+    assert res.status_code == 200
+    assert "Access-Control-Allow-Origin" not in res.headers
+
+
 def test_cors_proxy_only_forwards_explicit_proxy_headers():
     class CaptureHeadersHandler(BaseHTTPRequestHandler):
         def do_GET(self):

tools/server/tests/utils.py

@@ -114,6 +114,7 @@ class ServerProcess:
     backend_sampling: bool = False
     gcp_compat: bool = False
     server_tools: str | None = None
+    cors_origins: str | None = None
 
     # session variables
     process: subprocess.Popen | None = None
@@ -170,6 +171,8 @@ class ServerProcess:
             server_args.extend(["--models-max", self.models_max])
         if self.models_preset:
             server_args.extend(["--models-preset", self.models_preset])
+        if self.cors_origins:
+            server_args.extend(["--cors-origins", self.cors_origins])
         if self.n_batch:
             server_args.extend(["--batch-size", self.n_batch])
         if self.n_ubatch:
@@ -359,7 +362,7 @@ class ServerProcess:
         if parse_body:
             try:
                 result.body = response.json()
-            except JSONDecodeError:
+            except (JSONDecodeError, requests.exceptions.JSONDecodeError):
                 result.body = response.text
         else:
             result.body = None

tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessageActions/ChatMessageActionCard/ChatMessageActionCardPermissionRequest.svelte

@@ -21,7 +21,7 @@
 <ChatMessageActionCard icon={ShieldQuestion}>
 	{#snippet message()}
 		Allow use of <span class="font-semibold">{toolName}</span>{#if serverLabel}
-			from <span class="font-semibold">{serverLabel}</span>{/if}?
+			&nbsp;from <span class="font-semibold">{serverLabel}</span>{/if}?
 	{/snippet}
 
 	{#snippet actions()}

tools/ui/src/lib/constants/model-id.ts

@@ -24,8 +24,10 @@ export const MODEL_CUSTOM_QUANTIZATION_PREFIX_RE = /^UD$/i;
 
 /**
  * Matches a parameter-count segment, e.g. `7B`, `1.5b`, `120M`.
+ * The optional leading `E` covers effective-parameter sizes, e.g. Gemma's
+ * `E2B`/`E4B` (MatFormer models sized by resident params).
  */
-export const MODEL_PARAMS_RE = /^\d+(\.\d+)?[BbMmKkTt]$/;
+export const MODEL_PARAMS_RE = /^[Ee]?\d+(\.\d+)?[BbMmKkTt]$/;
 
 /**
  * Matches an activated-parameter-count segment, e.g. `A10B`, `a2.4b`.

tools/ui/tests/unit/model-id-parser.test.ts

@@ -36,6 +36,11 @@ describe('parseModelId', () => {
 		expect(parseModelId('model-100b:q4_k_m')).toMatchObject({ params: '100B' });
 	});
 
+	it('extracts effective parameters correctly', () => {
+		expect(parseModelId('model-E4B-BF16')).toMatchObject({ params: 'E4B' });
+		expect(parseModelId('model-e2b:q4_k_m')).toMatchObject({ params: 'E2B' });
+	});
+
 	it('extracts activated parameters correctly', () => {
 		expect(parseModelId('model-100B-A10B-BF16')).toMatchObject({ activatedParams: 'A10B' });
 		expect(parseModelId('model-100B-A10B:Q4_K_M')).toMatchObject({ activatedParams: 'A10B' });