opencl: fix two issues on flash attention for Adreno a7x (#25697)

* opencl: route `sub_group_shuffle_xor` to qcom ext when KHR ext is unavailable

KHR `sub_group_shuffle_xor` is not defined by compiler when
`cl_qcom_subgroup_shuffle` is present, causing certain FA
kernels fail to build. Define the KHR shuffle_xor using
the qcom extension.

* opencl: skip FA kernels with mixed and quant types for A7x to avoid compiler crash

.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

@@ -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) {
@@ -1077,6 +1082,7 @@ bool common_params_parse(int argc, char ** argv, common_params & params, llama_e
             if (ctx_arg.print_usage) {
                 ctx_arg.print_usage(argc, argv);
             }
+            common_log_flush(common_log_main());
             exit(0);
         }
         if (ctx_arg.params.completion) {
@@ -1178,6 +1184,8 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
         params.sampling.temp = 0.2; // lower temp by default for better quality
     } else if (ex == LLAMA_EXAMPLE_SERVER) {
         params.n_parallel = -1;     // auto by default
+    } else if (ex == LLAMA_EXAMPLE_TOKENIZE) {
+        params.parse_special = true; // parse special tokens by default, like the old tokenize tool
     }
 
     params.use_color = tty_can_use_colors();
@@ -2745,14 +2753,14 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
         [](common_params & params, const std::string & value) {
             params.model.path = value;
         }
-    ).set_examples({LLAMA_EXAMPLE_COMMON, LLAMA_EXAMPLE_EXPORT_LORA, LLAMA_EXAMPLE_DOWNLOAD}).set_env("LLAMA_ARG_MODEL"));
+    ).set_examples({LLAMA_EXAMPLE_COMMON, LLAMA_EXAMPLE_EXPORT_LORA, LLAMA_EXAMPLE_DOWNLOAD, LLAMA_EXAMPLE_TOKENIZE}).set_env("LLAMA_ARG_MODEL"));
     add_opt(common_arg(
         {"-mu", "--model-url"}, "MODEL_URL",
         "model download url (default: unused)",
         [](common_params & params, const std::string & value) {
             params.model.url = value;
         }
-    ).set_examples({LLAMA_EXAMPLE_COMMON, LLAMA_EXAMPLE_DOWNLOAD}).set_env("LLAMA_ARG_MODEL_URL"));
+    ).set_examples({LLAMA_EXAMPLE_COMMON, LLAMA_EXAMPLE_DOWNLOAD, LLAMA_EXAMPLE_TOKENIZE}).set_env("LLAMA_ARG_MODEL_URL"));
     add_opt(common_arg(
         { "-dr", "--docker-repo" }, "[<repo>/]<model>[:quant]",
         "Docker Hub model repository. repo is optional, default to ai/. quant is optional, default to :latest.\n"
@@ -2761,7 +2769,7 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
         [](common_params & params, const std::string & value) {
             params.model.docker_repo = value;
         }
-    ).set_examples({LLAMA_EXAMPLE_COMMON, LLAMA_EXAMPLE_DOWNLOAD}).set_env("LLAMA_ARG_DOCKER_REPO"));
+    ).set_examples({LLAMA_EXAMPLE_COMMON, LLAMA_EXAMPLE_DOWNLOAD, LLAMA_EXAMPLE_TOKENIZE}).set_env("LLAMA_ARG_DOCKER_REPO"));
     add_opt(common_arg(
         {"-hf", "-hfr", "--hf-repo"}, "<user>/<model>[:quant]",
         "Hugging Face model repository; quant is optional, case-insensitive, default to Q4_K_M, or falls back to the first file in the repo if Q4_K_M doesn't exist.\n"
@@ -2771,14 +2779,14 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
         [](common_params & params, const std::string & value) {
             params.model.hf_repo = value;
         }
-    ).set_examples({LLAMA_EXAMPLE_COMMON, LLAMA_EXAMPLE_DOWNLOAD}).set_env("LLAMA_ARG_HF_REPO"));
+    ).set_examples({LLAMA_EXAMPLE_COMMON, LLAMA_EXAMPLE_DOWNLOAD, LLAMA_EXAMPLE_TOKENIZE}).set_env("LLAMA_ARG_HF_REPO"));
     add_opt(common_arg(
         {"-hff", "--hf-file"}, "FILE",
         "Hugging Face model file. If specified, it will override the quant in --hf-repo (default: unused)",
         [](common_params & params, const std::string & value) {
             params.model.hf_file = value;
         }
-    ).set_examples({LLAMA_EXAMPLE_COMMON, LLAMA_EXAMPLE_DOWNLOAD}).set_env("LLAMA_ARG_HF_FILE"));
+    ).set_examples({LLAMA_EXAMPLE_COMMON, LLAMA_EXAMPLE_DOWNLOAD, LLAMA_EXAMPLE_TOKENIZE}).set_env("LLAMA_ARG_HF_FILE"));
     add_opt(common_arg(
         {"-hfv", "-hfrv", "--hf-repo-v"}, "<user>/<model>[:quant]",
         "Hugging Face model repository for the vocoder model (default: unused)",
@@ -2799,7 +2807,7 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
         [](common_params & params, const std::string & value) {
             params.hf_token = value;
         }
-    ).set_examples({LLAMA_EXAMPLE_COMMON, LLAMA_EXAMPLE_DOWNLOAD}).set_env("HF_TOKEN"));
+    ).set_examples({LLAMA_EXAMPLE_COMMON, LLAMA_EXAMPLE_DOWNLOAD, LLAMA_EXAMPLE_TOKENIZE}).set_env("HF_TOKEN"));
     add_opt(common_arg(
         {"--mtp"},
         "also download the multi-token prediction (MTP) head, if available (default: unused)",
@@ -2915,6 +2923,41 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
             params.parse_special = true;
         }
     ).set_examples({LLAMA_EXAMPLE_IMATRIX}));
+    add_opt(common_arg(
+        {"--ids"},
+        string_format("only print the token IDs, in a Python-parseable list form like [1, 2, 3] (default: %s)", params.tokenize_ids ? "true" : "false"),
+        [](common_params & params) {
+            params.tokenize_ids = true;
+        }
+    ).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"),
+        [](common_params & params) {
+            params.tokenize_stdin = true;
+        }
+    ).set_examples({LLAMA_EXAMPLE_TOKENIZE}));
+    add_opt(common_arg(
+        {"--no-bos"},
+        string_format("do not add a BOS token to the prompt, even if the model normally uses one (default: %s)", params.tokenize_no_bos ? "true" : "false"),
+        [](common_params & params) {
+            params.tokenize_no_bos = true;
+        }
+    ).set_examples({LLAMA_EXAMPLE_TOKENIZE}));
+    add_opt(common_arg(
+        {"--no-parse-special"},
+        string_format("do not parse special tokens (chat, tool, etc) (default: %s)", !params.parse_special ? "true" : "false"),
+        [](common_params & params) {
+            params.parse_special = false;
+        }
+    ).set_examples({LLAMA_EXAMPLE_TOKENIZE}));
+    add_opt(common_arg(
+        {"--show-count"},
+        string_format("print the total number of tokens (default: %s)", params.tokenize_show_count ? "true" : "false"),
+        [](common_params & params) {
+            params.tokenize_show_count = true;
+        }
+    ).set_examples({LLAMA_EXAMPLE_TOKENIZE}));
     add_opt(common_arg(
         {"-pps"},
         string_format("is the prompt shared across parallel sequences (default: %s)", params.is_pp_shared ? "true" : "false"),
@@ -3009,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()),
@@ -3042,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);
         }
@@ -3050,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"};
@@ -3059,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(
@@ -3505,7 +3587,7 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
         [](common_params & params) {
             params.offline = true;
         }
-    ).set_examples({LLAMA_EXAMPLE_COMMON, LLAMA_EXAMPLE_DOWNLOAD}).set_env("LLAMA_ARG_OFFLINE"));
+    ).set_examples({LLAMA_EXAMPLE_COMMON, LLAMA_EXAMPLE_DOWNLOAD, LLAMA_EXAMPLE_TOKENIZE}).set_env("LLAMA_ARG_OFFLINE"));
     add_opt(common_arg(
         {"-lv", "--verbosity", "--log-verbosity"}, "N",
         string_format("Set the verbosity threshold. Messages with a higher verbosity will be ignored. Values:\n"

common/common.h

@@ -105,6 +105,7 @@ enum llama_example {
     LLAMA_EXAMPLE_RESULTS,
     LLAMA_EXAMPLE_EXPORT_GRAPH_OPS,
     LLAMA_EXAMPLE_DOWNLOAD,
+    LLAMA_EXAMPLE_TOKENIZE,
 
     LLAMA_EXAMPLE_COUNT,
 };
@@ -630,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;
@@ -716,6 +725,12 @@ struct common_params {
     // batched-bench params
     bool batched_bench_output_jsonl = false;
 
+    // tokenize params
+    bool tokenize_ids        = false; // if true, only print the token IDs
+    bool tokenize_stdin      = false; // if true, read the prompt from stdin
+    bool tokenize_no_bos     = false; // if true, do not add the BOS token
+    bool tokenize_show_count = false; // if true, print the total token count
+
     // common params
     std::string out_file; // output filename for all example programs
     // optional callback for model loading progress and cancellation:

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

conversion/base.py

@@ -109,7 +109,9 @@ class ModelBase:
     sentence_transformers_dense_modules: bool = False
 
     # MTP (multi-token prediction) export modes; set by main() before instantiation.
-    # Architectures opt in by overriding the handling (see _Qwen35MtpMixin).
+    # Architectures that implement the filtering/export behavior opt in by
+    # setting supports_mtp_export = True on their model class or a mixin.
+    supports_mtp_export: bool = False
     mtp_only: bool = False
     no_mtp: bool = False
 

conversion/hunyuan.py

@@ -361,6 +361,7 @@ class HunyuanVLTextModel(HunYuanModel):
 @ModelBase.register("HYV3ForCausalLM")
 class HYV3Model(TextModel):
     model_arch = gguf.MODEL_ARCH.HY_V3
+    supports_mtp_export = True
 
     # Trunk layer count, stashed before indexing so the classmethod
     # filter_tensors can identify the appended MTP block(s) (mirrors

conversion/qwen.py

@@ -541,6 +541,7 @@ class _Qwen35MtpMixin:
     `mtp.*` to the standard layer-indexed nextn naming so the existing
     tensor_map handles them."""
 
+    supports_mtp_export = True
     hparams: dict[str, Any]
     model_arch: gguf.MODEL_ARCH
     gguf_writer: gguf.GGUFWriter

conversion/step3.py

@@ -98,6 +98,7 @@ class Step3VLTextModel(Qwen3Model):
 @ModelBase.register("Step3p5ForCausalLM", "Step3p7ForConditionalGeneration")
 class Step35Model(TextModel):
     model_arch = gguf.MODEL_ARCH.STEP35
+    supports_mtp_export = True
 
     # The --mtp / --no-mtp toggles are ModelBase.mtp_only / no_mtp (set in
     # convert_hf_to_gguf.py main()). Unlike Qwen3.5, which stores MTP under a

convert_hf_to_gguf.py

@@ -259,10 +259,8 @@ def main() -> None:
             sys.exit(1)
 
         if args.mtp or args.no_mtp:
-            from conversion.qwen import _Qwen35MtpMixin
-            from conversion.step3 import Step35Model
-            if not (issubclass(model_class, _Qwen35MtpMixin) or issubclass(model_class, Step35Model)):
-                logger.error("--mtp / --no-mtp are only supported for Qwen3.5/3.6 and Step3.5 text variants today")
+            if not model_class.supports_mtp_export:
+                logger.error("--mtp / --no-mtp are not supported for %s", model_architecture)
                 sys.exit(1)
             if args.no_mtp:
                 model_class.no_mtp = True

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/include/ggml.h

@@ -780,6 +780,10 @@ extern "C" {
     GGML_API bool ggml_is_contiguous_1(const struct ggml_tensor * tensor); // contiguous for dims >= 1
     GGML_API bool ggml_is_contiguous_2(const struct ggml_tensor * tensor); // contiguous for dims >= 2
 
+    GGML_API bool ggml_is_contiguous_to_1(const struct ggml_tensor * tensor); // contiguous for dims < 1
+    GGML_API bool ggml_is_contiguous_to_2(const struct ggml_tensor * tensor); // contiguous for dims < 2
+    GGML_API bool ggml_is_contiguous_to_3(const struct ggml_tensor * tensor); // contiguous for dims < 3
+
     // returns whether the tensor elements are allocated as one contiguous block of memory (no gaps, but permutation ok)
     GGML_API bool ggml_is_contiguously_allocated(const struct ggml_tensor * tensor);
 

ggml/src/ggml-cpu/CMakeLists.txt

@@ -638,6 +638,7 @@ function(ggml_add_cpu_backend_variant_impl tag_name)
             ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_qai8dxp_qsi8cxp/
             ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_fp32_bf16p_bf16p/
             ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_f16p_qsi4c32p/
+            ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_f32p_f32p/
             ${KLEIDIAI_SRC}/kai/ukernels/matmul/pack/)
 
         set(ARCH_FLAGS_TEMP "${ARCH_FLAGS}")
@@ -687,9 +688,15 @@ function(ggml_add_cpu_backend_variant_impl tag_name)
                 ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_fp32_bf16p_bf16p/kai_matmul_clamp_f32_bf16p2vlx2_bf16p2vlx2_2vlx2vl_sme2_mopa_asm.S
                 ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_f16p_qsi4c32p/kai_matmul_clamp_f32_f16p1vlx2_qsi4c32p4vlx2_1vlx4vl_sme2_mopa.c
                 ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_f16p_qsi4c32p/kai_matmul_clamp_f32_f16p1vlx2_qsi4c32p4vlx2_1vlx4vl_sme2_mopa_asm.S
+                ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_f32p_f32p/kai_matmul_clamp_f32_f32p2vlx1_f32p2vlx1biasf32_sme2_mopa.c
+                ${KLEIDIAI_SRC}/kai/ukernels/matmul/matmul_clamp_f32_f32p_f32p/kai_matmul_clamp_f32_f32p2vlx1_f32p2vlx1biasf32_sme2_mopa_asm.S
                 ${KLEIDIAI_SRC}/kai/ukernels/matmul/pack/kai_lhs_pack_bf16p2vlx2_f32_sme.c
                 ${KLEIDIAI_SRC}/kai/ukernels/matmul/pack/kai_rhs_pack_kxn_bf16p2vlx2b_f32_x32_sme.c
                 ${KLEIDIAI_SRC}/kai/ukernels/matmul/pack/kai_lhs_pack_f16pmrx2_f32_neon.c
+                ${KLEIDIAI_SRC}/kai/ukernels/matmul/pack/kai_lhs_pack_f32p2vlx1_f32_sme.c
+                ${KLEIDIAI_SRC}/kai/ukernels/matmul/pack/kai_lhs_pack_f32p2vlx1_f32_sme_asm.S
+                ${KLEIDIAI_SRC}/kai/ukernels/matmul/pack/kai_rhs_pack_nxk_f32p2vlx1biasf32_f32_f32_sme.c
+                ${KLEIDIAI_SRC}/kai/ukernels/matmul/pack/kai_rhs_pack_nxk_f32p2vlx1biasf32_f32_f32_sme_asm.S
                 ${KLEIDIAI_SRC}/kai/kai_common_sme_asm.S)
             set(PRIVATE_ARCH_FLAGS "-fno-tree-vectorize;${PRIVATE_ARCH_FLAGS}+sve+sve2+sme2+fp16")
         endif()

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

@@ -2859,7 +2859,14 @@ 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;
+                case GGML_OP_SET_ROWS:
+                    {
+                        if (node->src[0]->type == GGML_TYPE_F16 && node->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/kleidiai/kernels.cpp

@@ -20,14 +20,17 @@
 #include "kai_matmul_clamp_f32_qsi8d32p4x8_qsi4c32p8x8_16x8_sve_i8mm.h"
 #include "kai_matmul_clamp_f32_qsi8d32p1x8_qsi4c32p8x8_1x8_sve_dotprod.h"
 #include "kai_matmul_clamp_f32_f16p1vlx2_qsi4c32p4vlx2_1vlx4vl_sme2_mopa.h"
+#include "kai_matmul_clamp_f32_f32p2vlx1_f32p2vlx1biasf32_sme2_mopa.h"
 
 #include "kai_lhs_pack_bf16p2vlx2_f32_sme.h"
+#include "kai_lhs_pack_f32p2vlx1_f32_sme.h"
 #include "kai_lhs_quant_pack_qsi8d32p_f32.h"
 #include "kai_lhs_quant_pack_qsi8d32p4x8sb_f32_neon.h"
 #include "kai_lhs_quant_pack_qsi8d32p_f32_neon.h"
 #include "kai_lhs_quant_pack_qai8dxp_f32.h"
 
 #include "kai_rhs_pack_kxn_bf16p2vlx2b_f32_x32_sme.h"
+#include "kai_rhs_pack_nxk_f32p2vlx1biasf32_f32_f32_sme.h"
 #include "kai_rhs_pack_nxk_qsi4c32pscalef16_qsu4c32s16s0.h"
 #include "kai_rhs_pack_nxk_qsi4c32ps1s0scalef16_qsu4c32s16s0_neon.h"
 #include "kai_rhs_pack_nxk_qsi8cxp_qsi8cx_neon.h"
@@ -865,6 +868,65 @@ static ggml_kleidiai_kernels gemm_gemv_kernels_q8[] = {
     { /* Sentinel */ }
 };
 
+static ggml_kleidiai_kernels ggml_kleidiai_kernels_f32[] = {
+#if defined(__ARM_FEATURE_SME)
+    {
+        /* SME GEMM */
+        {
+            /* .get_m_step            = */ kai_get_m_step_matmul_clamp_f32_f32p2vlx1_f32p2vlx1biasf32_sme2_mopa,
+            /* .get_n_step            = */ kai_get_n_step_matmul_clamp_f32_f32p2vlx1_f32p2vlx1biasf32_sme2_mopa,
+            /* .get_mr                = */ kai_get_mr_matmul_clamp_f32_f32p2vlx1_f32p2vlx1biasf32_sme2_mopa,
+            /* .get_nr                = */ kai_get_nr_matmul_clamp_f32_f32p2vlx1_f32p2vlx1biasf32_sme2_mopa,
+            /* .get_kr                = */ kai_get_kr_matmul_clamp_f32_f32p2vlx1_f32p2vlx1biasf32_sme2_mopa,
+            /* .get_sr                = */ kai_get_sr_matmul_clamp_f32_f32p2vlx1_f32p2vlx1biasf32_sme2_mopa,
+            /* .get_dst_offset        = */ kai_get_dst_offset_matmul_clamp_f32_f32p2vlx1_f32p2vlx1biasf32_sme2_mopa,
+            /* .get_dst_size          = */ kai_get_dst_size_matmul_clamp_f32_f32p2vlx1_f32p2vlx1biasf32_sme2_mopa,
+            /* .get_lhs_offset_ex     = */ &kernel_offs_fn2<kai_get_lhs_packed_offset_matmul_clamp_f32_f32p2vlx1_f32p2vlx1biasf32_sme2_mopa>,
+            /* .get_rhs_packed_offset_ex = */ &kernel_offs_fn2<kai_get_rhs_packed_offset_matmul_clamp_f32_f32p2vlx1_f32p2vlx1biasf32_sme2_mopa>,
+            /* .run_kernel_ex         = */ &kernel_run_fn10<kai_run_matmul_clamp_f32_f32p2vlx1_f32p2vlx1biasf32_sme2_mopa>,
+        },
+        /* .gemm_lhs_info = */ {
+            /* .get_offset            = */ kai_get_lhs_offset_lhs_pack_f32p2vlx1_f32_sme,
+            /* .get_packed_offset_ex  = */ &lhs_offs_fn5<kai_get_lhs_packed_offset_lhs_pack_f32p2vlx1_f32_sme>,
+            /* .packed_size_ex        = */ &lhs_ps_fn5<kai_get_lhs_packed_size_lhs_pack_f32p2vlx1_f32_sme>,
+            /* .pack_func_ex          = */ &lhs_pack_void_fn9<kai_run_lhs_pack_f32p2vlx1_f32_sme>,
+        },
+        /* SME GEMV */
+        {
+            /* .get_m_step            = */ kai_get_m_step_matmul_clamp_f32_f32p2vlx1_f32p2vlx1biasf32_sme2_mopa,
+            /* .get_n_step            = */ kai_get_n_step_matmul_clamp_f32_f32p2vlx1_f32p2vlx1biasf32_sme2_mopa,
+            /* .get_mr                = */ kai_get_mr_matmul_clamp_f32_f32p2vlx1_f32p2vlx1biasf32_sme2_mopa,
+            /* .get_nr                = */ kai_get_nr_matmul_clamp_f32_f32p2vlx1_f32p2vlx1biasf32_sme2_mopa,
+            /* .get_kr                = */ kai_get_kr_matmul_clamp_f32_f32p2vlx1_f32p2vlx1biasf32_sme2_mopa,
+            /* .get_sr                = */ kai_get_sr_matmul_clamp_f32_f32p2vlx1_f32p2vlx1biasf32_sme2_mopa,
+            /* .get_dst_offset        = */ kai_get_dst_offset_matmul_clamp_f32_f32p2vlx1_f32p2vlx1biasf32_sme2_mopa,
+            /* .get_dst_size          = */ kai_get_dst_size_matmul_clamp_f32_f32p2vlx1_f32p2vlx1biasf32_sme2_mopa,
+            /* .get_lhs_offset_ex     = */ nullptr,
+            /* .get_rhs_packed_offset_ex = */ nullptr,
+            /* .run_kernel_ex         = */ nullptr,
+        },
+        /* .gemv_lhs_info = */ {
+            /* .get_offset            = */ kai_get_lhs_offset_lhs_pack_f32p2vlx1_f32_sme,
+            /* .get_packed_offset_ex  = */ &lhs_offs_fn5<kai_get_lhs_packed_offset_lhs_pack_f32p2vlx1_f32_sme>,
+            /* .packed_size_ex        = */ &lhs_ps_fn5<kai_get_lhs_packed_size_lhs_pack_f32p2vlx1_f32_sme>,
+            /* .pack_func_ex          = */ &lhs_pack_void_fn9<kai_run_lhs_pack_f32p2vlx1_f32_sme>,
+        },
+        /* .rhs_info = */ {
+            /* .packed_stride         = */ nullptr,
+            /* .to_float              = */ nullptr,
+            /* .packed_size_ex        = */ &rhs_ps_fn2<kai_get_rhs_packed_size_rhs_pack_nxk_f32p2vlx1biasf32_f32_f32_sme>,
+            /* .packed_stride_ex      = */ &rhs_stride_fn1<kai_get_rhs_packed_stride_rhs_pack_nxk_f32p2vlx1biasf32_f32_f32_sme>,
+            /* .pack_func_ex          = */ &rhs_pack_fn13<kai_run_rhs_pack_nxk_f32p2vlx1biasf32_f32_f32_sme>,
+        },
+        /* .required_cpu       = */ CPU_FEATURE_SME,
+        /* .lhs_type           = */ GGML_TYPE_F32,
+        /* .rhs_type           = */ GGML_TYPE_F32,
+        /* .op_type            = */ GGML_TYPE_F32,
+    },
+#endif
+    { /* Sentinel */ }
+};
+
 ggml_kleidiai_kernels * ggml_kleidiai_select_kernels(cpu_feature cpu_features, const ggml_tensor * tensor) {
     ggml_kleidiai_kernels * kernel = nullptr;
 
@@ -888,12 +950,15 @@ ggml_kleidiai_kernels * ggml_kleidiai_select_kernels(cpu_feature cpu_features, c
 
         if (tensor->src[0]->type == GGML_TYPE_Q8_0) {
             try_table(gemm_gemv_kernels_q8);
+        } else if (tensor->src[0]->type == GGML_TYPE_F32) {
+            try_table(ggml_kleidiai_kernels_f32);
         } else {
             try_table(gemm_gemv_kernels);
         }
 #else
     GGML_UNUSED(gemm_gemv_kernels);
     GGML_UNUSED(gemm_gemv_kernels_q8);
+    GGML_UNUSED(ggml_kleidiai_kernels_f32);
     GGML_UNUSED(cpu_features);
 #endif
     }
@@ -937,3 +1002,20 @@ ggml_kleidiai_kernels * ggml_kleidiai_select_kernels_q8_0(cpu_feature features)
 
     return kernels;
 }
+
+ggml_kleidiai_kernels * ggml_kleidiai_select_kernels_f32(cpu_feature features) {
+    ggml_kleidiai_kernels * kernels = nullptr;
+
+#if defined(__ARM_FEATURE_SME)
+    for (size_t i = 0; i < NELEMS(ggml_kleidiai_kernels_f32) - 1; ++i) {
+        if ((features & ggml_kleidiai_kernels_f32[i].required_cpu) == ggml_kleidiai_kernels_f32[i].required_cpu) {
+            kernels = &ggml_kleidiai_kernels_f32[i];
+            break;
+        }
+    }
+#else
+    GGML_UNUSED(features);
+#endif
+
+    return kernels;
+}

ggml/src/ggml-cpu/kleidiai/kernels.h

@@ -55,6 +55,12 @@ struct lhs_packing_info {
         size_t m_idx_start, const void * lhs, size_t lhs_stride, void * lhs_packed);
 };
 
+enum rhs_repack_mode {
+    RHS_REPACK_PER_KERNEL,
+    RHS_REPACK_SHARED,
+    RHS_REPACK_SINGLE_ONLY,
+};
+
 struct rhs_packing_info {
     size_t (*packed_stride)(size_t k, size_t nr, size_t kr, size_t bl);
 
@@ -68,6 +74,8 @@ struct rhs_packing_info {
 
     void (*pack_func_ex)(size_t num_groups, size_t n, size_t k, size_t nr, size_t kr, size_t sr, size_t bl,
         size_t rhs_stride, const void * rhs, const void * bias, const void * scale, void * rhs_packed, size_t extra_bytes, const void * params);
+
+    rhs_repack_mode repack_mode = RHS_REPACK_PER_KERNEL;
 };
 
 struct ggml_kleidiai_kernels {
@@ -88,3 +96,4 @@ struct ggml_kleidiai_kernels {
 ggml_kleidiai_kernels * ggml_kleidiai_select_kernels(cpu_feature cpu_features, const ggml_tensor * tensor);
 ggml_kleidiai_kernels * ggml_kleidiai_select_kernels_q4_0(cpu_feature features);
 ggml_kleidiai_kernels * ggml_kleidiai_select_kernels_q8_0(cpu_feature features);
+ggml_kleidiai_kernels * ggml_kleidiai_select_kernels_f32(cpu_feature features);

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

@@ -60,10 +60,11 @@ struct ggml_kleidiai_context {
     cpu_feature features;
     ggml_kleidiai_kernels * kernels_q4;
     ggml_kleidiai_kernels * kernels_q8;
+    ggml_kleidiai_kernels * kernels_f32;
     int sme_thread_cap; // <= 0 means “SME disabled/unknown”;
     int thread_hint;    // <= 0 means “no hint”
     int chunk_multiplier;
-} static ctx = { CPU_FEATURE_NONE, nullptr, nullptr, 0, -1, 4 };
+} static ctx = { CPU_FEATURE_NONE, nullptr, nullptr, nullptr, 0, -1, 4 };
 
 static const char* cpu_feature_to_string(cpu_feature f) {
     if (f == CPU_FEATURE_NONE) {
@@ -156,10 +157,10 @@ static size_t detect_num_smcus() {
             }
         }
     }
-    return 1;
+    return 0;
 
 #else
-    return 1;
+    return 0;
 #endif
 }
 
@@ -192,7 +193,6 @@ static void init_kleidiai_context(void) {
         const char *env_threads     = getenv("GGML_TOTAL_THREADS");
         const char *env_chunk_mult  = getenv("GGML_KLEIDIAI_CHUNK_MULTIPLIER");
 
-        const bool cpu_has_sme = ggml_cpu_has_sme();
         size_t detected_smcus = 0;
 
         ctx.features  = (ggml_cpu_has_dotprod()     ? CPU_FEATURE_DOTPROD : CPU_FEATURE_NONE) |
@@ -216,56 +216,47 @@ static void init_kleidiai_context(void) {
         }
 
         // SME policy:
-        // - If CPU doesn't support SME: SME always off.
-        // - Else:
-        //   - env unset => auto-detect cores; enable if detected > 0.
-        //   - env=0     => force off.
-        //   - env>0     => force N cores (skip detection).
+        // - env unset => auto-detect SMCUs; enable SME only if detected > 0.
+        // - env=0     => force off.
+        // - env>0     => force N cores, if the binary was built with SME.
         int sme_cores = 0;
         bool sme_env_ok = false;
         bool sme_env_set = (env_sme != nullptr);
 
-        if (!cpu_has_sme) {
-            if (sme_env_set) {
-                bool ok = false;
-                int req = parse_uint_env(env_sme, "GGML_KLEIDIAI_SME", &ok);
-                if (ok && req > 0) {
-                    GGML_LOG_WARN("kleidiai: GGML_KLEIDIAI_SME=%d but SME is not supported on this CPU; disabling SME\n", req);
-                }
-            }
-            sme_cores = 0;
-        } else {
-            if (sme_env_set) {
-                bool ok = false;
-                int v = parse_uint_env(env_sme, "GGML_KLEIDIAI_SME", &ok);
-                sme_env_ok = ok;
+        if (sme_env_set) {
+            bool ok = false;
+            int v = parse_uint_env(env_sme, "GGML_KLEIDIAI_SME", &ok);
+            sme_env_ok = ok;
 
-                if (!ok) {
-                    GGML_LOG_WARN("kleidiai: GGML_KLEIDIAI_SME set but parsing failed; falling back to runtime SME-core detection\n");
-                    detected_smcus = detect_num_smcus();
-                    sme_cores = detected_smcus > 0 ? (int)detected_smcus : 0;
-                } else if (v == 0) {
-                    sme_cores = 0;
-                } else {
-                    sme_cores = v;
-                }
-            } else {
+            if (!ok) {
+                GGML_LOG_WARN("kleidiai: GGML_KLEIDIAI_SME set but parsing failed; falling back to runtime SME-core detection\n");
                 detected_smcus = detect_num_smcus();
                 sme_cores = detected_smcus > 0 ? (int)detected_smcus : 0;
+            } else if (v == 0) {
+                sme_cores = 0;
+            } else if (!ggml_cpu_has_sme()) {
+                GGML_LOG_WARN("kleidiai: GGML_KLEIDIAI_SME=%d but the binary was not built with SME; disabling SME\n", v);
+                sme_cores = 0;
+            } else {
+                sme_cores = v;
             }
+        } else {
+            detected_smcus = detect_num_smcus();
+            sme_cores = detected_smcus > 0 ? (int)detected_smcus : 0;
+        }
 
-            if (!sme_env_set && sme_cores == 0) {
-                GGML_LOG_WARN("kleidiai: SME supported but runtime SME-core detection returned 0; falling back to NEON\n");
-            }
+        if (!sme_env_set && ggml_cpu_has_sme() && sme_cores == 0) {
+            GGML_LOG_WARN("kleidiai: runtime SME-core detection returned 0; falling back to NEON\n");
+        }
 
-            if (sme_cores > 0) {
-                ctx.features |= CPU_FEATURE_SME;
-            }
+        if (sme_cores > 0) {
+            ctx.features |= CPU_FEATURE_SME;
         }
 
         // Kernel selection
-        ctx.kernels_q4 = ggml_kleidiai_select_kernels_q4_0(ctx.features);
-        ctx.kernels_q8 = ggml_kleidiai_select_kernels_q8_0(ctx.features);
+        ctx.kernels_q4  = ggml_kleidiai_select_kernels_q4_0(ctx.features);
+        ctx.kernels_q8  = ggml_kleidiai_select_kernels_q8_0(ctx.features);
+        ctx.kernels_f32 = ggml_kleidiai_select_kernels_f32(ctx.features);
 
         if (!ctx.kernels_q4) {
             GGML_LOG_INFO("kleidiai: no compatible q4 kernels found for CPU features mask %d\n", (int)ctx.features);
@@ -279,6 +270,12 @@ static void init_kleidiai_context(void) {
             GGML_LOG_INFO("kleidiai: primary q8 kernel feature %s\n", cpu_feature_to_string(ctx.kernels_q8->required_cpu));
         }
 
+        if (!ctx.kernels_f32) {
+            GGML_LOG_INFO("kleidiai: no compatible f32 kernels found for CPU features mask %d\n", (int)ctx.features);
+        } else {
+            GGML_LOG_INFO("kleidiai: primary f32 kernel feature %s\n", cpu_feature_to_string(ctx.kernels_f32->required_cpu));
+        }
+
         ctx.sme_thread_cap = (ctx.features & CPU_FEATURE_SME) ? sme_cores : 0;
 
         if (ctx.features & CPU_FEATURE_SME) {
@@ -334,6 +331,13 @@ static inline size_t ceil_div_size(size_t a, size_t b) {
     return b == 0 ? 0 : (a + b - 1) / b;
 }
 
+static inline size_t kleidiai_chunk_cols(size_t n, int nth_total, bool disable_chunking, size_t n_step) {
+    const size_t multiplier = (nth_total == 1 || disable_chunking) ? 1 : std::max<size_t>(1, (size_t) ctx.chunk_multiplier);
+    const size_t divisor = std::max<size_t>(1, (size_t) nth_total * multiplier);
+    const size_t chunk_cols = align_up(std::max<size_t>(1, ceil_div_size(n, divisor)), n_step);
+    return chunk_cols ? chunk_cols : n_step;
+}
+
 struct kleidiai_block_args {
     size_t lhs_bl;
     size_t rhs_bl;
@@ -418,6 +422,10 @@ static inline ggml_kleidiai_kernels * kleidiai_primary_kernel_q8() {
     return ctx.kernels_q8;
 }
 
+static inline ggml_kleidiai_kernels * kleidiai_primary_kernel_f32() {
+    return ctx.kernels_f32;
+}
+
 template <typename SelectFallback>
 static int kleidiai_collect_kernel_chain_common(
         ggml_kleidiai_kernels * primary,
@@ -430,11 +438,16 @@ static int kleidiai_collect_kernel_chain_common(
     }
     out[count++] = primary;
 
+    if (primary->rhs_info.repack_mode == RHS_REPACK_SINGLE_ONLY) {
+        return count;
+    }
+
     if ((primary->required_cpu & CPU_FEATURE_SME) == CPU_FEATURE_SME) {
         const cpu_feature fallback_mask = static_cast<cpu_feature>(features & ~CPU_FEATURE_SME);
         if (fallback_mask != CPU_FEATURE_NONE) {
             ggml_kleidiai_kernels * fallback = select_fallback(fallback_mask);
             if (fallback && fallback != primary &&
+                fallback->rhs_info.repack_mode != RHS_REPACK_SINGLE_ONLY &&
                 fallback->lhs_type == primary->lhs_type &&
                 fallback->rhs_type == primary->rhs_type &&
                 fallback->op_type  == primary->op_type) {
@@ -465,6 +478,12 @@ static int kleidiai_collect_q8_chain(std::array<ggml_kleidiai_kernels *, GGML_KL
         [&](cpu_feature mask) { return ggml_kleidiai_select_kernels_q8_0(mask); });
 }
 
+static int kleidiai_collect_f32_chain(std::array<ggml_kleidiai_kernels *, GGML_KLEIDIAI_MAX_KERNEL_SLOTS> & out) {
+    ggml_kleidiai_kernels * primary = kleidiai_primary_kernel_f32();
+    return kleidiai_collect_kernel_chain_common(primary, ctx.features, out,
+        [&](cpu_feature mask) { return ggml_kleidiai_select_kernels_f32(mask); });
+}
+
 static inline int64_t ggml_ne(const ggml_tensor * tensor, int dim) {
     GGML_ASSERT(dim >= 0 && dim < GGML_MAX_DIMS);
     return tensor->ne[dim];
@@ -539,6 +558,36 @@ class tensor_traits : public ggml::cpu::tensor_traits {
             return true;
         }
 
+        if (op->src[0]->type == GGML_TYPE_F32) {
+            size_t cursor = 0;
+            bool any_slot = false;
+
+            for (int slot = 0; slot < slot_count; ++slot) {
+                ggml_kleidiai_kernels * kernels = kernel_chain[slot];
+                lhs_packing_info * lhs_info = &kernels->gemm_lhs_info;
+                kernel_info * kernel        = &kernels->gemm;
+
+                if (!lhs_info || !lhs_info->packed_size_ex || !kernel) {
+                    return false;
+                }
+
+                const size_t mr = kernel->get_mr();
+                const size_t kr = kernel->get_kr();
+                const size_t sr = kernel->get_sr();
+
+                cursor  = align_up(cursor, GGML_KLEIDIAI_PACK_ALIGN);
+                cursor += lhs_info->packed_size_ex(m, k, 0, mr, kr, sr);
+                any_slot = true;
+            }
+
+            if (!any_slot) {
+                return false;
+            }
+
+            size = cursor;
+            return true;
+        }
+
         if (op->src[0]->type == GGML_TYPE_F16) {
             const int64_t lhs_batch_size0 = op->src[1]->ne[2];
             const int64_t rhs_batch_size0 = op->src[0]->ne[2];
@@ -595,6 +644,8 @@ class tensor_traits : public ggml::cpu::tensor_traits {
         if (dst->op == GGML_OP_MUL_MAT) {
             if (dst->src[0]->type == GGML_TYPE_Q4_0 || dst->src[0]->type == GGML_TYPE_Q8_0) {
                 return compute_forward_qx(params, dst);
+            } else if (dst->src[0]->type == GGML_TYPE_F32) {
+                return compute_forward_f32(params, dst);
             } else if (dst->src[0]->type == GGML_TYPE_F16) {
                 return compute_forward_fp16(params, dst);
             }
@@ -606,6 +657,144 @@ class tensor_traits : public ggml::cpu::tensor_traits {
         return false;
     }
 
+    bool compute_forward_f32(ggml_compute_params * params, struct ggml_tensor * dst) {
+        GGML_ASSERT(dst->src[0]->type == GGML_TYPE_F32);
+
+        const ggml_tensor * src0 = dst->src[0];
+        const ggml_tensor * src1 = dst->src[1];
+
+        GGML_TENSOR_BINARY_OP_LOCALS
+
+        if (src1->type != GGML_TYPE_F32 || dst->type != GGML_TYPE_F32) {
+            return false;
+        }
+
+        ggml_kleidiai_kernels * kernels = kleidiai_primary_kernel_f32();
+        if (!kernels) {
+            return false;
+        }
+
+        kernel_info * kernel        = &kernels->gemm;
+        lhs_packing_info * lhs_info = &kernels->gemm_lhs_info;
+
+        if (!kernel || !lhs_info || !lhs_info->get_offset || !lhs_info->get_packed_offset_ex ||
+            !lhs_info->packed_size_ex || !lhs_info->pack_func_ex ||
+            !kernel->get_rhs_packed_offset_ex || !kernel->run_kernel_ex || !kernel->get_dst_offset) {
+            return false;
+        }
+
+        const kleidiai_weight_header * header = kleidiai_weight_header_from_ptr(src0->data);
+        const bool has_header = kleidiai_is_weight_header_valid(header);
+
+        const uint8_t * rhs_base = has_header ? kleidiai_weight_slot_ptr(header, 0)
+                                              : static_cast<const uint8_t *>(src0->data);
+        if (!rhs_base) {
+            return false;
+        }
+
+        const int nth = params->nth > 0 ? params->nth : 1;
+        const int ith = params->ith;
+
+        const size_t k = ne00;
+        const size_t m = ne11;
+        const size_t n = ne01;
+
+        const size_t mr = kernel->get_mr();
+        const size_t kr = kernel->get_kr();
+        const size_t sr = kernel->get_sr();
+
+        const size_t lhs_packed_size = lhs_info->packed_size_ex(m, k, 0, mr, kr, sr);
+        GGML_ASSERT(lhs_packed_size <= params->wsize);
+
+        uint8_t * lhs_packed   = static_cast<uint8_t *>(params->wdata);
+        const size_t dst_stride = dst->nb[1];
+        const size_t n_step = kernel->get_n_step() ? kernel->get_n_step() : 1;
+        const bool disable_chunking = ggml_is_numa();
+        GGML_ASSERT(n <= (size_t) INT_MAX);
+
+        for (int64_t batch_idx = 0; batch_idx < ne12; ++batch_idx) {
+            const uint8_t * lhs_batch_base = static_cast<const uint8_t *>(src1->data) + batch_idx * src1->nb[2];
+            uint8_t * dst_batch_base = static_cast<uint8_t *>(dst->data) + batch_idx * dst->nb[2];
+
+            {
+                const int64_t m_roundup_mr = kai_roundup((int64_t)m, (int64_t)mr);
+                int64_t max_threads = mr ? (m_roundup_mr / (int64_t)mr) : nth;
+                max_threads = std::max<int64_t>(1, max_threads);
+                const int64_t use_threads = std::min<int64_t>(nth, max_threads);
+
+                if (ith < use_threads) {
+                    const int64_t num_m_per_thread0   = round_down((size_t)(m_roundup_mr / use_threads), mr);
+                    const int64_t num_m_per_threadN_1 = (int64_t)m - (use_threads - 1) * num_m_per_thread0;
+
+                    const int64_t m_start = (int64_t)ith * num_m_per_thread0;
+                    const int64_t m_count = (ith == use_threads - 1) ? num_m_per_threadN_1 : num_m_per_thread0;
+
+                    const size_t base_packed_off  = lhs_info->get_packed_offset_ex(m_start, k, 0, mr, kr, sr);
+                    const size_t next_block_off   = lhs_info->get_packed_offset_ex(m_start + mr, k, 0, mr, kr, sr);
+                    const size_t row_stride_bytes = mr ? (next_block_off - base_packed_off) / mr : 0;
+
+                    int64_t remaining = m_count;
+                    int64_t cur       = m_start;
+
+                    while (remaining > 0) {
+                        const int64_t take = std::min<int64_t>((int64_t)m - cur, remaining);
+                        const size_t src_off = lhs_info->get_offset(cur, src1->nb[1]);
+                        const void * src_ptr = lhs_batch_base + src_off;
+                        const size_t dst_off = base_packed_off + (size_t)(cur - m_start) * row_stride_bytes;
+                        void * dst_ptr       = lhs_packed + dst_off;
+
+                        lhs_info->pack_func_ex(take, k, 0, mr, kr, sr, 0, src_ptr, src1->nb[1], dst_ptr);
+
+                        cur       += take;
+                        remaining -= take;
+                    }
+                }
+            }
+
+            if (ith == 0) {
+                ggml_threadpool_chunk_set(params->threadpool, 0);
+            }
+
+            ggml_barrier(params->threadpool);
+
+            const size_t chunk_cols = kleidiai_chunk_cols(n, nth, disable_chunking, n_step);
+            GGML_ASSERT(chunk_cols <= (size_t) INT_MAX);
+
+            int current_col = ggml_threadpool_chunk_add(params->threadpool, (int) chunk_cols);
+            while ((size_t) current_col < n) {
+                const size_t n_start = (size_t) current_col;
+                const size_t n_to_process = std::min(chunk_cols, n - n_start);
+
+                if (n_to_process > 0) {
+                    const size_t lhs_packed_offset = lhs_info->get_packed_offset_ex(0, k, 0, mr, kr, sr);
+                    const size_t rhs_packed_offset = kernel->get_rhs_packed_offset_ex(n_start, k, 0);
+                    const size_t dst_offset        = kernel->get_dst_offset(0, n_start, dst_stride);
+
+                    const void * lhs_ptr = lhs_packed + lhs_packed_offset;
+                    const void * rhs_ptr = rhs_base + rhs_packed_offset;
+                    float * dst_ptr      = reinterpret_cast<float *>(dst_batch_base + dst_offset);
+
+                    kernel->run_kernel_ex(m, n_to_process, k, 0,
+                                          lhs_ptr,
+                                          rhs_ptr,
+                                          dst_ptr,
+                                          dst_stride,
+                                          sizeof(float),
+                                          -FLT_MAX,
+                                          FLT_MAX);
+                }
+
+                current_col = ggml_threadpool_chunk_add(params->threadpool, (int) chunk_cols);
+            }
+
+            if (batch_idx != ne12 - 1) {
+                ggml_barrier(params->threadpool);
+            }
+        }
+
+        return true;
+    }
+
     bool compute_forward_fp16(ggml_compute_params * params, struct ggml_tensor * dst) {
         const ggml_tensor * src0 = dst->src[0];
         const ggml_tensor * src1 = dst->src[1];
@@ -1214,7 +1403,7 @@ class tensor_traits : public ggml::cpu::tensor_traits {
 
 public:
     int repack(struct ggml_tensor * tensor, const void * data, size_t data_size) {
-        GGML_ASSERT(tensor->type == GGML_TYPE_Q4_0 || tensor->type == GGML_TYPE_Q8_0);
+        GGML_ASSERT(tensor->type == GGML_TYPE_Q4_0 || tensor->type == GGML_TYPE_Q8_0 || tensor->type == GGML_TYPE_F32);
         const size_t n = tensor->ne[1];
         const size_t k = tensor->ne[0];
 
@@ -1233,12 +1422,15 @@ public:
 
         std::array<ggml_kleidiai_kernels *, GGML_KLEIDIAI_MAX_KERNEL_SLOTS> kernel_chain;
         const bool want_q8 = tensor->type == GGML_TYPE_Q8_0;
-        const int slot_total = want_q8 ? kleidiai_collect_q8_chain(kernel_chain)
-                                       : kleidiai_collect_q4_chain(kernel_chain);
+        const bool want_f32 = tensor->type == GGML_TYPE_F32;
+        const int slot_total = want_f32 ? kleidiai_collect_f32_chain(kernel_chain)
+                                        : want_q8 ? kleidiai_collect_q8_chain(kernel_chain)
+                                                  : kleidiai_collect_q4_chain(kernel_chain);
         const bool allow_fallback = kleidiai_pack_fallback_allowed();
 
         std::vector<int8_t> qdata;
         std::vector<float>  scales;
+        std::vector<float>  bias;
 
         if (want_q8 && slot_total > 0) {
             qdata.resize(n * k, 0);
@@ -1286,6 +1478,10 @@ public:
             }
         }
 
+        if (want_f32 && slot_total > 0) {
+            bias.resize(n, 0.0f);
+        }
+
         for (int slot = 0; slot < slot_total && slot < GGML_KLEIDIAI_MAX_KERNEL_SLOTS; ++slot) {
             if (!allow_fallback && slot > 0) {
                 break;
@@ -1302,8 +1498,9 @@ public:
             const size_t sr = kernel->get_sr();
             const ggml_type rhs_type = kernels->rhs_type;
             const size_t block_len = rhs_type == GGML_TYPE_Q8_0 ? QK8_0 :
-                                     rhs_type == GGML_TYPE_Q4_0 ? QK4_0 : 0;
-            if (block_len == 0) {
+                                     rhs_type == GGML_TYPE_Q4_0 ? QK4_0 :
+                                     rhs_type == GGML_TYPE_F32 ? 0 : SIZE_MAX;
+            if (block_len == SIZE_MAX) {
                 continue;
             }
 
@@ -1326,6 +1523,10 @@ public:
                 rhs_info->pack_func_ex(1, n, k, nr, kr, sr, 0, 0,
                                        qdata.data(), nullptr, scales.data(),
                                        dst_ptr, 0, &params);
+            } else if (rhs_type == GGML_TYPE_F32) {
+                rhs_info->pack_func_ex(1, n, k, nr, kr, sr, 0, tensor->nb[1],
+                                       data, bias.data(), nullptr,
+                                       dst_ptr, 0, nullptr);
             } else {
                 continue;
             }
@@ -1400,7 +1601,7 @@ static size_t ggml_backend_cpu_kleidiai_buffer_type_get_alignment(ggml_backend_b
 static size_t ggml_backend_cpu_kleidiai_buffer_type_get_alloc_size(ggml_backend_buffer_type_t buft, const struct ggml_tensor * tensor) {
     GGML_UNUSED(buft);
 
-    if (tensor->type != GGML_TYPE_Q4_0 && tensor->type != GGML_TYPE_Q8_0) {
+    if (tensor->type != GGML_TYPE_Q4_0 && tensor->type != GGML_TYPE_Q8_0 && tensor->type != GGML_TYPE_F32) {
         return ggml_nbytes(tensor);
     }
 
@@ -1412,8 +1613,10 @@ static size_t ggml_backend_cpu_kleidiai_buffer_type_get_alloc_size(ggml_backend_
 
     std::array<ggml_kleidiai_kernels *, GGML_KLEIDIAI_MAX_KERNEL_SLOTS> kernel_chain;
     const bool want_q8 = tensor->type == GGML_TYPE_Q8_0;
-    const int slot_total = want_q8 ? kleidiai_collect_q8_chain(kernel_chain)
-                                   : kleidiai_collect_q4_chain(kernel_chain);
+    const bool want_f32 = tensor->type == GGML_TYPE_F32;
+    const int slot_total = want_f32 ? kleidiai_collect_f32_chain(kernel_chain)
+                                    : want_q8 ? kleidiai_collect_q8_chain(kernel_chain)
+                                              : kleidiai_collect_q4_chain(kernel_chain);
     const bool allow_fallback = kleidiai_pack_fallback_allowed();
 
     size_t slot_count = 0;
@@ -1433,8 +1636,9 @@ static size_t ggml_backend_cpu_kleidiai_buffer_type_get_alloc_size(ggml_backend_
 
         const ggml_type rhs_type = kernels->rhs_type;
         const size_t block_len = rhs_type == GGML_TYPE_Q4_0 ? QK4_0 :
-                                 rhs_type == GGML_TYPE_Q8_0 ? QK8_0 : 0;
-        if (block_len == 0) {
+                                 rhs_type == GGML_TYPE_Q8_0 ? QK8_0 :
+                                 rhs_type == GGML_TYPE_F32 ? 0 : SIZE_MAX;
+        if (block_len == SIZE_MAX) {
             continue;
         }
 
@@ -1455,25 +1659,41 @@ class extra_buffer_type : ggml::cpu::extra_buffer_type {
     bool supports_op(ggml_backend_dev_t, const struct ggml_tensor * op) override {
         std::array<ggml_kleidiai_kernels *, GGML_KLEIDIAI_MAX_KERNEL_SLOTS> kernel_chain;
         const int slot_total = kleidiai_collect_kernel_chain(op, kernel_chain);
-        if ((op->op == GGML_OP_MUL_MAT || op->op == GGML_OP_GET_ROWS) &&
-            (op->src[0]->type == GGML_TYPE_Q4_0 || op->src[0]->type == GGML_TYPE_Q8_0) &&
+        const bool src0_is_kleidiai =
             op->src[0]->buffer &&
             (ggml_n_dims(op->src[0]) == 2) &&
             op->src[0]->buffer->buft == ggml_backend_cpu_kleidiai_buffer_type() &&
-            slot_total > 0) {
+            slot_total > 0;
+
+        if ((op->op == GGML_OP_MUL_MAT || op->op == GGML_OP_GET_ROWS) &&
+            (op->src[0]->type == GGML_TYPE_Q4_0 || op->src[0]->type == GGML_TYPE_Q8_0 || op->src[0]->type == GGML_TYPE_F32) &&
+            src0_is_kleidiai) {
             if (op->src[0]->type == GGML_TYPE_Q4_0 && ctx.kernels_q4 == nullptr) {
                 return false;
             }
             if (op->src[0]->type == GGML_TYPE_Q8_0 && ctx.kernels_q8 == nullptr) {
                 return false;
             }
+            if (op->src[0]->type == GGML_TYPE_F32 && ctx.kernels_f32 == nullptr) {
+                return false;
+            }
             if (op->src[1]->buffer && !ggml_backend_buft_is_host(op->src[1]->buffer->buft)) {
                 return false;
             }
-            if ((op->src[1]->type == GGML_TYPE_F32 || op->src[1]->type == GGML_TYPE_I32) &&
-                ggml_ne(op->src[1], 3) == 1) {
-                return true;
+
+            if (op->src[0]->type == GGML_TYPE_Q4_0 || op->src[0]->type == GGML_TYPE_Q8_0) {
+                if ((op->src[1]->type == GGML_TYPE_F32 || op->src[1]->type == GGML_TYPE_I32) &&
+                    ggml_ne(op->src[1], 3) == 1) {
+                    return true;
+                }
+                return false;
             }
+
+            if (op->op != GGML_OP_MUL_MAT || op->src[1]->type != GGML_TYPE_F32 || op->type != GGML_TYPE_F32) {
+                return false;
+            }
+
+            return true;
         }
         return false;
     }

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);
@@ -5041,7 +5104,7 @@ static void ggml_compute_forward_set_rows_impl(
     assert(ne0  == nc);
     assert(ne2  == ne02);
     assert(ne3  == ne03);
-    GGML_ASSERT(src0->type == GGML_TYPE_F32 || (src0->type == GGML_TYPE_F16 && dst->type == GGML_TYPE_F16));
+    GGML_ASSERT(src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_F16);
     assert(ne02 % ne11 == 0);
     assert(ne03 % ne12 == 0);
 
@@ -5075,10 +5138,19 @@ static void ggml_compute_forward_set_rows_impl(
                             (const float *) ((char *) src0->data +  i*nb01 + i02*nb02 + i03*nb03),
                                             ((char *)  dst->data + i1*nb1  + i02*nb2  + i03*nb3), nc);
                 } else if constexpr (std::is_same_v<src_t, ggml_fp16_t>) {
-                    memcpy(
+                    if (dst->type == GGML_TYPE_F16) {
+                        memcpy(
                                             ((char *)  dst->data + i1*nb1  + i02*nb2  + i03*nb3),
                                             ((char *) src0->data +  i*nb01 + i02*nb02 + i03*nb03),
                                             rs);
+                    } else {
+                        float * wdata = (float *) params->wdata + (nc + CACHE_LINE_SIZE_F32) * ith;
+                        ggml_fp16_to_fp32_row(
+                                (const ggml_fp16_t *) ((char *) src0->data + i*nb01 + i02*nb02 + i03*nb03),
+                                wdata, nc);
+                        from_float(wdata,
+                                ((char *) dst->data + i1*nb1 + i02*nb2 + i03*nb3), nc);
+                    }
                 } else {
                     GGML_ABORT("src0->type = %d (%s) not supported", src0->type, ggml_type_name(src0->type));
                 }
@@ -5107,16 +5179,12 @@ void ggml_compute_forward_set_rows(
             } break;
         case GGML_TYPE_F16:
             {
-                if (dst->type == GGML_TYPE_F16) {
-                    if (src1->type == GGML_TYPE_I64) {
-                        ggml_compute_forward_set_rows_impl<ggml_fp16_t, int64_t>(params, dst);
-                    } else if (src1->type == GGML_TYPE_I32) {
-                        ggml_compute_forward_set_rows_impl<ggml_fp16_t, int32_t>(params, dst);
-                    } else {
-                        GGML_ABORT("src1->type = %d (%s) not supported", src1->type, ggml_type_name(src1->type));
-                    }
+                if (src1->type == GGML_TYPE_I64) {
+                    ggml_compute_forward_set_rows_impl<ggml_fp16_t, int64_t>(params, dst);
+                } else if (src1->type == GGML_TYPE_I32) {
+                    ggml_compute_forward_set_rows_impl<ggml_fp16_t, int32_t>(params, dst);
                 } else {
-                    GGML_ABORT("dst->type = %d (%s) not supported with src0->type = %d (%s)", dst->type, ggml_type_name(dst->type), src0->type, ggml_type_name(src0->type));
+                    GGML_ABORT("src1->type = %d (%s) not supported", src1->type, ggml_type_name(src1->type));
                 }
             } break;
         default:
@@ -6362,7 +6430,6 @@ static void ggml_compute_forward_im2col_f16(
     const ggml_tensor * src0 = dst->src[0];
     const ggml_tensor * src1 = dst->src[1];
 
-    GGML_ASSERT(src0->type == GGML_TYPE_F16);
     GGML_ASSERT(src1->type == GGML_TYPE_F16 || src1->type == GGML_TYPE_F32);
     GGML_ASSERT( dst->type == GGML_TYPE_F16);
 
@@ -6393,7 +6460,6 @@ static void ggml_compute_forward_im2col_f16(
     int ofs0 = is_2D ? nb13 : nb12;
     int ofs1 = is_2D ? nb12 : nb11;
 
-    GGML_ASSERT(nb00 == sizeof(ggml_fp16_t));
     GGML_ASSERT(nb10 == ggml_type_size(src1->type));
 
     // im2col: [N, IC, IH, IW] => [N, OH, OW, IC*KH*KW]
@@ -6466,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;
@@ -6563,7 +6629,6 @@ static void ggml_compute_forward_im2col_3d_f16(
     const ggml_tensor * src0 = dst->src[0];
     const ggml_tensor * src1 = dst->src[1];
 
-    GGML_ASSERT(src0->type == GGML_TYPE_F16);
     GGML_ASSERT(src1->type == GGML_TYPE_F32);
     GGML_ASSERT( dst->type == GGML_TYPE_F16);
 

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

@@ -4816,13 +4816,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
                            ) || (

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

@@ -805,6 +805,11 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv(ggml_meta
                 nsg = N_SG_Q1_0;
                 nr0 = N_R0_Q1_0;
             } break;
+        case GGML_TYPE_Q2_0:
+            {
+                nsg = N_SG_Q2_0;
+                nr0 = N_R0_Q2_0;
+            } break;
         case GGML_TYPE_Q4_0:
             {
                 nsg = N_SG_Q4_0;
@@ -1029,6 +1034,11 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id(ggml_m
                 nsg = N_SG_Q1_0;
                 nr0 = N_R0_Q1_0;
             } break;
+        case GGML_TYPE_Q2_0:
+            {
+                nsg = N_SG_Q2_0;
+                nr0 = N_R0_Q2_0;
+            } break;
         case GGML_TYPE_Q4_0:
             {
                 nsg = N_SG_Q4_0;
@@ -1824,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-device.m

@@ -1289,6 +1289,7 @@ bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_te
                            case GGML_TYPE_BF16:
                            case GGML_TYPE_Q8_0:
                            case GGML_TYPE_Q1_0:
+                           case GGML_TYPE_Q2_0:
                            case GGML_TYPE_Q4_0:
                            case GGML_TYPE_Q4_1:
                            case GGML_TYPE_Q5_0:
@@ -1316,6 +1317,7 @@ bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_te
                                 return false;
                         }
                     case GGML_TYPE_Q1_0:
+                    case GGML_TYPE_Q2_0:
                     case GGML_TYPE_Q4_0:
                     case GGML_TYPE_Q4_1:
                     case GGML_TYPE_Q5_0:
@@ -1338,7 +1340,11 @@ bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_te
             return op->src[0]->type != GGML_TYPE_NVFP4;
         case GGML_OP_SET_ROWS:
             {
-                if (op->src[0]->type != GGML_TYPE_F32 && op->src[0]->type != GGML_TYPE_F16) {
+                if (op->src[0]->type == GGML_TYPE_F16) {
+                    return op->type == GGML_TYPE_F16;
+                }
+
+                if (op->src[0]->type != GGML_TYPE_F32) {
                     return false;
                 }
 

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

@@ -24,6 +24,9 @@
 #define N_R0_Q1_0 8
 #define N_SG_Q1_0 2
 
+#define N_R0_Q2_0 8
+#define N_SG_Q2_0 2
+
 #define N_R0_Q4_0 4
 #define N_SG_Q4_0 2
 
@@ -613,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

@@ -2077,6 +2077,7 @@ int ggml_metal_op_mul_mat(ggml_metal_op_t ctx, int idx) {
            op->src[0]->type == GGML_TYPE_F16  ||
            op->src[0]->type == GGML_TYPE_BF16 ||
            op->src[0]->type == GGML_TYPE_Q1_0 ||
+           op->src[0]->type == GGML_TYPE_Q2_0 ||
            op->src[0]->type == GGML_TYPE_Q4_0 ||
            op->src[0]->type == GGML_TYPE_Q4_1 ||
            op->src[0]->type == GGML_TYPE_Q5_0 ||
@@ -3076,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;
@@ -3983,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

@@ -170,6 +170,39 @@ void dequantize_q1_0_t4(device const block_q1_0 * xb, short il, thread type4 & r
     reg = (type4) reg_f;
 }
 
+template <typename type4x4>
+void dequantize_q2_0(device const block_q2_0 * xb, short il, thread type4x4 & reg) {
+    device const uint8_t * qs = xb->qs;
+    const float d = xb->d;
+
+    const int byte_offset = il * 4;  // il*16 elements = il*4 bytes (4 elements per byte)
+    float4x4 reg_f;
+
+    for (int i = 0; i < 4; i++) {
+        const uint8_t b = qs[byte_offset + i];
+        reg_f[i][0] = ((float)((b >> 0) & 3) - 1.0f) * d;
+        reg_f[i][1] = ((float)((b >> 2) & 3) - 1.0f) * d;
+        reg_f[i][2] = ((float)((b >> 4) & 3) - 1.0f) * d;
+        reg_f[i][3] = ((float)((b >> 6) & 3) - 1.0f) * d;
+    }
+
+    reg = (type4x4) reg_f;
+}
+
+template <typename type4>
+void dequantize_q2_0_t4(device const block_q2_0 * xb, short il, thread type4 & reg) {
+    const float d = xb->d;
+    const uint8_t b = xb->qs[il];
+
+    float4 reg_f;
+    reg_f[0] = ((float)((b >> 0) & 3) - 1.0f) * d;
+    reg_f[1] = ((float)((b >> 2) & 3) - 1.0f) * d;
+    reg_f[2] = ((float)((b >> 4) & 3) - 1.0f) * d;
+    reg_f[3] = ((float)((b >> 6) & 3) - 1.0f) * d;
+
+    reg = (type4) reg_f;
+}
+
 template <typename type4x4>
 void dequantize_q4_0(device const block_q4_0 * xb, short il, thread type4x4 & reg) {
     device const uint16_t * qs = ((device const uint16_t *)xb + 1);
@@ -221,6 +254,27 @@ void quantize_q1_0(device const float * src, device block_q1_0 & dst) {
     }
 }
 
+void quantize_q2_0(device const float * src, device block_q2_0 & dst) {
+    float amax = 0.0f;
+    for (int j = 0; j < QK2_0; j++) {
+        float a = fabs(src[j]);
+        if (a > amax) amax = a;
+    }
+    const float d = amax;
+    dst.d = d;
+
+    const float id = d > 0.0f ? 1.0f / d : 0.0f;
+
+    for (int j = 0; j < QK2_0 / 4; j++) {
+        dst.qs[j] = 0;
+    }
+    for (int j = 0; j < QK2_0; j++) {
+        int q = (int)round(src[j] * id) + 1;
+        q = max(0, min(3, q));
+        dst.qs[j / 4] |= (q << (2 * (j % 4)));
+    }
+}
+
 void quantize_q4_0(device const float * src, device block_q4_0 & dst) {
 #pragma METAL fp math_mode(safe)
     float amax = 0.0f; // absolute max
@@ -3289,6 +3343,55 @@ inline float block_q_n_dot_y(device const block_q1_0 * qb_curr, float sumy, thre
     return qb_curr->d * (2.0f * acc - sumy);
 }
 
+// Q2_0 dot: d * (sum_lo(y) + 2*sum_hi(y) - sumy) via per-bit conditional adds
+inline float block_q_n_dot_y(device const block_q2_0 * qb_curr, float sumy, thread float * yl, int il) {
+    device const uint8_t * qs = qb_curr->qs + (il / 4);
+    const uint8_t b0 = qs[0];
+    const uint8_t b1 = qs[1];
+    const uint8_t b2 = qs[2];
+    const uint8_t b3 = qs[3];
+
+    // Accumulate where low bit is set (bits 0,2,4,6 of each byte)
+    float acc_lo = 0.0f;
+    acc_lo += select(0.0f, yl[ 0], bool(b0 & 0x01));
+    acc_lo += select(0.0f, yl[ 1], bool(b0 & 0x04));
+    acc_lo += select(0.0f, yl[ 2], bool(b0 & 0x10));
+    acc_lo += select(0.0f, yl[ 3], bool(b0 & 0x40));
+    acc_lo += select(0.0f, yl[ 4], bool(b1 & 0x01));
+    acc_lo += select(0.0f, yl[ 5], bool(b1 & 0x04));
+    acc_lo += select(0.0f, yl[ 6], bool(b1 & 0x10));
+    acc_lo += select(0.0f, yl[ 7], bool(b1 & 0x40));
+    acc_lo += select(0.0f, yl[ 8], bool(b2 & 0x01));
+    acc_lo += select(0.0f, yl[ 9], bool(b2 & 0x04));
+    acc_lo += select(0.0f, yl[10], bool(b2 & 0x10));
+    acc_lo += select(0.0f, yl[11], bool(b2 & 0x40));
+    acc_lo += select(0.0f, yl[12], bool(b3 & 0x01));
+    acc_lo += select(0.0f, yl[13], bool(b3 & 0x04));
+    acc_lo += select(0.0f, yl[14], bool(b3 & 0x10));
+    acc_lo += select(0.0f, yl[15], bool(b3 & 0x40));
+
+    // Accumulate where high bit is set (bits 1,3,5,7 of each byte)
+    float acc_hi = 0.0f;
+    acc_hi += select(0.0f, yl[ 0], bool(b0 & 0x02));
+    acc_hi += select(0.0f, yl[ 1], bool(b0 & 0x08));
+    acc_hi += select(0.0f, yl[ 2], bool(b0 & 0x20));
+    acc_hi += select(0.0f, yl[ 3], bool(b0 & 0x80));
+    acc_hi += select(0.0f, yl[ 4], bool(b1 & 0x02));
+    acc_hi += select(0.0f, yl[ 5], bool(b1 & 0x08));
+    acc_hi += select(0.0f, yl[ 6], bool(b1 & 0x20));
+    acc_hi += select(0.0f, yl[ 7], bool(b1 & 0x80));
+    acc_hi += select(0.0f, yl[ 8], bool(b2 & 0x02));
+    acc_hi += select(0.0f, yl[ 9], bool(b2 & 0x08));
+    acc_hi += select(0.0f, yl[10], bool(b2 & 0x20));
+    acc_hi += select(0.0f, yl[11], bool(b2 & 0x80));
+    acc_hi += select(0.0f, yl[12], bool(b3 & 0x02));
+    acc_hi += select(0.0f, yl[13], bool(b3 & 0x08));
+    acc_hi += select(0.0f, yl[14], bool(b3 & 0x20));
+    acc_hi += select(0.0f, yl[15], bool(b3 & 0x80));
+
+    return qb_curr->d * (acc_lo + 2.0f * acc_hi - sumy);
+}
+
 // function for calculate inner product between half a q4_0 block and 16 floats (yl), sumy is SUM(yl[i])
 // il indicates where the q4 quants begin (0 or QK4_0/4)
 // we assume that the yl's have been multiplied with the appropriate scale factor
@@ -3592,6 +3695,86 @@ kernel void kernel_mul_mv_q1_0_f32(
     kernel_mul_mv_q1_0_f32_impl<N_R0_Q1_0, constant ggml_metal_kargs_mul_mv &>(args, src0, src1, dst, nullptr, tgpig, tiisg, sgitg);
 }
 
+template<int nr0, typename args_t>
+void kernel_mul_mv_q2_0_f32_impl(
+        args_t args,
+        device const char * src0,
+        device const char * src1,
+        device       char * dst,
+        threadgroup  char * shmem,
+        uint3  tgpig,
+        ushort tiisg,
+        ushort sgitg) {
+    const short NSG = FC_mul_mv_nsg;
+
+    const int nb = args.ne00/QK2_0;
+
+    const int r0 = tgpig.x;
+    const int r1 = tgpig.y;
+    const int im = tgpig.z;
+
+    const int first_row = (r0 * NSG + sgitg) * nr0;
+
+    const uint i12 = im%FC_mul_mv_ne12;
+    const uint i13 = im/FC_mul_mv_ne12;
+
+    const uint64_t offset1 = r1*args.nb11 + (i12)*args.nb12 + (i13)*args.nb13;
+
+    device const float * y = (device const float *) (src1 + offset1);
+
+    device const block_q2_0 * ax[nr0];
+    for (int row = 0; row < nr0; ++row) {
+        const uint64_t offset0 = (first_row + row)*args.nb01 + (i12/FC_mul_mv_r2)*args.nb02 + (i13/FC_mul_mv_r3)*args.nb03;
+        ax[row] = (device const block_q2_0 *) ((device char *) src0 + offset0);
+    }
+
+    float yl[16];
+    float sumf[nr0] = {0.f};
+
+    // group 64: 4 sub-blocks of 16 weights per Q2_0 block
+    const short ix = (tiisg/4);
+    const short il = (tiisg%4)*16;
+
+    device const float * yb = y + ix*QK2_0 + il;
+
+    for (int ib = ix; ib < nb; ib += N_SIMDWIDTH/4) {
+        float sumy = 0.f;
+
+        FOR_UNROLL (short i = 0; i < 16; i++) {
+            yl[i] = yb[i];
+            sumy += yb[i];
+        }
+
+        FOR_UNROLL (short row = 0; row < nr0; row++) {
+            sumf[row] += block_q_n_dot_y(ax[row] + ib, sumy, yl, il);
+        }
+
+        yb += QK2_0 * (N_SIMDWIDTH/4);
+    }
+
+    device float * dst_f32 = (device float *) dst + (uint64_t)im*args.ne0*args.ne1 + (uint64_t)r1*args.ne0;
+
+    for (int row = 0; row < nr0; ++row) {
+        const float tot = simd_sum(sumf[row]);
+
+        if (tiisg == 0 && first_row + row < args.ne01) {
+            dst_f32[first_row + row] = tot;
+        }
+    }
+}
+
+[[host_name("kernel_mul_mv_q2_0_f32")]]
+kernel void kernel_mul_mv_q2_0_f32(
+        constant ggml_metal_kargs_mul_mv & args,
+        device const char * src0,
+        device const char * src1,
+        device       char * dst,
+        uint3  tgpig[[threadgroup_position_in_grid]],
+        ushort tiisg[[thread_index_in_simdgroup]],
+        ushort sgitg[[simdgroup_index_in_threadgroup]]) {
+    kernel_mul_mv_q2_0_f32_impl<N_R0_Q2_0, constant ggml_metal_kargs_mul_mv &>(args, src0, src1, dst, nullptr, tgpig, tiisg, sgitg);
+}
+
 kernel void kernel_mul_mv_q4_0_f32(
         constant ggml_metal_kargs_mul_mv & args,
         device const char * src0,
@@ -3989,6 +4172,11 @@ template [[host_name("kernel_mul_mv_ext_q1_0_f32_r1_3")]]   kernel mul_mv_ext_q4
 template [[host_name("kernel_mul_mv_ext_q1_0_f32_r1_4")]]   kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<4, block_q1_0,   128, dequantize_q1_0_t4>;
 template [[host_name("kernel_mul_mv_ext_q1_0_f32_r1_5")]]   kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<5, block_q1_0,   128, dequantize_q1_0_t4>;
 
+template [[host_name("kernel_mul_mv_ext_q2_0_f32_r1_2")]]   kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<2, block_q2_0,    64, dequantize_q2_0_t4>;
+template [[host_name("kernel_mul_mv_ext_q2_0_f32_r1_3")]]   kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<3, block_q2_0,    64, dequantize_q2_0_t4>;
+template [[host_name("kernel_mul_mv_ext_q2_0_f32_r1_4")]]   kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<4, block_q2_0,    64, dequantize_q2_0_t4>;
+template [[host_name("kernel_mul_mv_ext_q2_0_f32_r1_5")]]   kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<5, block_q2_0,    64, dequantize_q2_0_t4>;
+
 template [[host_name("kernel_mul_mv_ext_q4_0_f32_r1_2")]]   kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<2, block_q4_0,   32, dequantize_q4_0_t4>;
 template [[host_name("kernel_mul_mv_ext_q4_0_f32_r1_3")]]   kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<3, block_q4_0,   32, dequantize_q4_0_t4>;
 template [[host_name("kernel_mul_mv_ext_q4_0_f32_r1_4")]]   kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<4, block_q4_0,   32, dequantize_q4_0_t4>;
@@ -5218,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,
@@ -7700,6 +7917,7 @@ typedef decltype(kernel_cpy_f32_q<QK8_0,  block_q8_0,  quantize_q8_0>)  cpy_f_q_
 
 template [[host_name("kernel_cpy_f32_q8_0")]]   kernel cpy_f_q_t kernel_cpy_f32_q<QK8_0,  block_q8_0,   quantize_q8_0>;
 template [[host_name("kernel_cpy_f32_q1_0")]]   kernel cpy_f_q_t kernel_cpy_f32_q<QK1_0,  block_q1_0,   quantize_q1_0>;
+template [[host_name("kernel_cpy_f32_q2_0")]]   kernel cpy_f_q_t kernel_cpy_f32_q<QK2_0,  block_q2_0,   quantize_q2_0>;
 template [[host_name("kernel_cpy_f32_q4_0")]]   kernel cpy_f_q_t kernel_cpy_f32_q<QK4_0,  block_q4_0,   quantize_q4_0>;
 template [[host_name("kernel_cpy_f32_q4_1")]]   kernel cpy_f_q_t kernel_cpy_f32_q<QK4_1,  block_q4_1,   quantize_q4_1>;
 template [[host_name("kernel_cpy_f32_q5_0")]]   kernel cpy_f_q_t kernel_cpy_f32_q<QK5_0,  block_q5_0,   quantize_q5_0>;
@@ -7745,6 +7963,7 @@ kernel void kernel_cpy_q_f32(
 typedef decltype(kernel_cpy_q_f32<float4x4, block_q4_0, 2, dequantize_q4_0>) cpy_q_f_t;
 
 template [[host_name("kernel_cpy_q1_0_f32")]] kernel cpy_q_f_t kernel_cpy_q_f32<float4x4, block_q1_0, 8, dequantize_q1_0>;
+template [[host_name("kernel_cpy_q2_0_f32")]] kernel cpy_q_f_t kernel_cpy_q_f32<float4x4, block_q2_0, 4, dequantize_q2_0>;
 template [[host_name("kernel_cpy_q4_0_f32")]] kernel cpy_q_f_t kernel_cpy_q_f32<float4x4, block_q4_0, 2, dequantize_q4_0>;
 template [[host_name("kernel_cpy_q4_1_f32")]] kernel cpy_q_f_t kernel_cpy_q_f32<float4x4, block_q4_1, 2, dequantize_q4_1>;
 template [[host_name("kernel_cpy_q5_0_f32")]] kernel cpy_q_f_t kernel_cpy_q_f32<float4x4, block_q5_0, 2, dequantize_q5_0>;
@@ -7752,6 +7971,7 @@ template [[host_name("kernel_cpy_q5_1_f32")]] kernel cpy_q_f_t kernel_cpy_q_f32<
 template [[host_name("kernel_cpy_q8_0_f32")]] kernel cpy_q_f_t kernel_cpy_q_f32<float4x4, block_q8_0, 2, dequantize_q8_0>;
 
 template [[host_name("kernel_cpy_q1_0_f16")]] kernel cpy_q_f_t kernel_cpy_q_f32<half4x4, block_q1_0, 8, dequantize_q1_0>;
+template [[host_name("kernel_cpy_q2_0_f16")]] kernel cpy_q_f_t kernel_cpy_q_f32<half4x4, block_q2_0, 4, dequantize_q2_0>;
 template [[host_name("kernel_cpy_q4_0_f16")]] kernel cpy_q_f_t kernel_cpy_q_f32<half4x4, block_q4_0, 2, dequantize_q4_0>;
 template [[host_name("kernel_cpy_q4_1_f16")]] kernel cpy_q_f_t kernel_cpy_q_f32<half4x4, block_q4_1, 2, dequantize_q4_1>;
 template [[host_name("kernel_cpy_q5_0_f16")]] kernel cpy_q_f_t kernel_cpy_q_f32<half4x4, block_q5_0, 2, dequantize_q5_0>;
@@ -9596,6 +9816,7 @@ template [[host_name("kernel_get_rows_bf16")]] kernel get_rows_f_t kernel_get_ro
 typedef decltype(kernel_get_rows_q<block_q4_0, 2, dequantize_q4_0>) get_rows_q_t;
 
 template [[host_name("kernel_get_rows_q1_0")]]    kernel get_rows_q_t kernel_get_rows_q<block_q1_0,    8, dequantize_q1_0>;
+template [[host_name("kernel_get_rows_q2_0")]]    kernel get_rows_q_t kernel_get_rows_q<block_q2_0,    4, dequantize_q2_0>;
 template [[host_name("kernel_get_rows_q4_0")]]    kernel get_rows_q_t kernel_get_rows_q<block_q4_0,    2, dequantize_q4_0>;
 template [[host_name("kernel_get_rows_q4_1")]]    kernel get_rows_q_t kernel_get_rows_q<block_q4_1,    2, dequantize_q4_1>;
 template [[host_name("kernel_get_rows_q5_0")]]    kernel get_rows_q_t kernel_get_rows_q<block_q5_0,    2, dequantize_q5_0>;
@@ -10466,6 +10687,7 @@ template [[host_name("kernel_mul_mm_f16_f32")]]     kernel mul_mm_t kernel_mul_m
 template [[host_name("kernel_mul_mm_bf16_f32")]]    kernel mul_mm_t kernel_mul_mm<bfloat, bfloat4x4, simdgroup_bfloat8x8, bfloat, bfloat2x4, simdgroup_bfloat8x8, bfloat4x4,     1,     dequantize_bf16,    bfloat, bfloat4x4, float, float2x4>;
 #endif
 template [[host_name("kernel_mul_mm_q1_0_f32")]]    kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q1_0,    8,     dequantize_q1_0,    float,  float4x4,  float, float2x4>;
+template [[host_name("kernel_mul_mm_q2_0_f32")]]    kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q2_0,    4,     dequantize_q2_0,    float,  float4x4,  float, float2x4>;
 template [[host_name("kernel_mul_mm_q4_0_f32")]]    kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q4_0,    2,     dequantize_q4_0,    float,  float4x4,  float, float2x4>;
 template [[host_name("kernel_mul_mm_q4_1_f32")]]    kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q4_1,    2,     dequantize_q4_1,    float,  float4x4,  float, float2x4>;
 template [[host_name("kernel_mul_mm_q5_0_f32")]]    kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q5_0,    2,     dequantize_q5_0,    float,  float4x4,  float, float2x4>;
@@ -10490,6 +10712,7 @@ template [[host_name("kernel_mul_mm_iq4_xs_f32")]]  kernel mul_mm_t kernel_mul_m
 template [[host_name("kernel_mul_mm_f32_f16")]]     kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   float4x4,      1,     dequantize_f32,     float,  float4x4,  half, half2x4>;
 template [[host_name("kernel_mul_mm_f16_f16")]]     kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   half4x4,       1,     dequantize_f16,     half,   half4x4,   half, half2x4>;
 template [[host_name("kernel_mul_mm_q1_0_f16")]]    kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q1_0,    8,     dequantize_q1_0,    float,  float4x4,  half, half2x4>;
+template [[host_name("kernel_mul_mm_q2_0_f16")]]    kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q2_0,    4,     dequantize_q2_0,    float,  float4x4,  half, half2x4>;
 template [[host_name("kernel_mul_mm_q4_0_f16")]]    kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q4_0,    2,     dequantize_q4_0,    float,  float4x4,  half, half2x4>;
 template [[host_name("kernel_mul_mm_q4_1_f16")]]    kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q4_1,    2,     dequantize_q4_1,    float,  float4x4,  half, half2x4>;
 template [[host_name("kernel_mul_mm_q5_0_f16")]]    kernel mul_mm_t kernel_mul_mm<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q5_0,    2,     dequantize_q5_0,    float,  float4x4,  half, half2x4>;
@@ -10523,6 +10746,7 @@ template [[host_name("kernel_mul_mm_id_f16_f32")]]     kernel mul_mm_id kernel_m
 template [[host_name("kernel_mul_mm_id_bf16_f32")]]    kernel mul_mm_id kernel_mul_mm_id<bfloat, bfloat4x4, simdgroup_bfloat8x8, bfloat, bfloat2x4, simdgroup_bfloat8x8, bfloat4x4,     1,     dequantize_bf16,    bfloat, bfloat4x4, float, float2x4>;
 #endif
 template [[host_name("kernel_mul_mm_id_q1_0_f32")]]    kernel mul_mm_id kernel_mul_mm_id<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q1_0,    8,     dequantize_q1_0,    float,  float4x4,  float, float2x4>;
+template [[host_name("kernel_mul_mm_id_q2_0_f32")]]    kernel mul_mm_id kernel_mul_mm_id<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q2_0,    4,     dequantize_q2_0,    float,  float4x4,  float, float2x4>;
 template [[host_name("kernel_mul_mm_id_q4_0_f32")]]    kernel mul_mm_id kernel_mul_mm_id<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q4_0,    2,     dequantize_q4_0,    float,  float4x4,  float, float2x4>;
 template [[host_name("kernel_mul_mm_id_q4_1_f32")]]    kernel mul_mm_id kernel_mul_mm_id<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q4_1,    2,     dequantize_q4_1,    float,  float4x4,  float, float2x4>;
 template [[host_name("kernel_mul_mm_id_q5_0_f32")]]    kernel mul_mm_id kernel_mul_mm_id<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q5_0,    2,     dequantize_q5_0,    float,  float4x4,  float, float2x4>;
@@ -10547,6 +10771,7 @@ template [[host_name("kernel_mul_mm_id_iq4_xs_f32")]]  kernel mul_mm_id kernel_m
 template [[host_name("kernel_mul_mm_id_f32_f16")]]     kernel mul_mm_id kernel_mul_mm_id<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   float4x4,      1,     dequantize_f32,     float,  float4x4,  half, half2x4>;
 template [[host_name("kernel_mul_mm_id_f16_f16")]]     kernel mul_mm_id kernel_mul_mm_id<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   half4x4,       1,     dequantize_f16,     half,   half4x4,   half, half2x4>;
 template [[host_name("kernel_mul_mm_id_q1_0_f16")]]    kernel mul_mm_id kernel_mul_mm_id<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q1_0,    8,     dequantize_q1_0,    float,  float4x4,  half, half2x4>;
+template [[host_name("kernel_mul_mm_id_q2_0_f16")]]    kernel mul_mm_id kernel_mul_mm_id<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q2_0,    4,     dequantize_q2_0,    float,  float4x4,  half, half2x4>;
 template [[host_name("kernel_mul_mm_id_q4_0_f16")]]    kernel mul_mm_id kernel_mul_mm_id<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q4_0,    2,     dequantize_q4_0,    float,  float4x4,  half, half2x4>;
 template [[host_name("kernel_mul_mm_id_q4_1_f16")]]    kernel mul_mm_id kernel_mul_mm_id<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q4_1,    2,     dequantize_q4_1,    float,  float4x4,  half, half2x4>;
 template [[host_name("kernel_mul_mm_id_q5_0_f16")]]    kernel mul_mm_id kernel_mul_mm_id<half,   half4x4,   simdgroup_half8x8,   half,   half2x4,   simdgroup_half8x8,   block_q5_0,    2,     dequantize_q5_0,    float,  float4x4,  half, half2x4>;
@@ -10702,6 +10927,7 @@ template [[host_name("kernel_mul_mv_id_bf16_f32_4")]]  kernel kernel_mul_mv_id_4
 template [[host_name("kernel_mul_mv_id_q8_0_f32")]]    kernel kernel_mul_mv_id_t kernel_mul_mv_id<mmv_fn<kernel_mul_mv_q8_0_f32_impl<N_R0_Q8_0>>>;
 
 template [[host_name("kernel_mul_mv_id_q1_0_f32")]]    kernel kernel_mul_mv_id_t kernel_mul_mv_id<mmv_fn<kernel_mul_mv_q1_0_f32_impl<N_R0_Q1_0>>>;
+template [[host_name("kernel_mul_mv_id_q2_0_f32")]]    kernel kernel_mul_mv_id_t kernel_mul_mv_id<mmv_fn<kernel_mul_mv_q2_0_f32_impl<N_R0_Q2_0>>>;
 template [[host_name("kernel_mul_mv_id_q4_0_f32")]]    kernel kernel_mul_mv_id_t kernel_mul_mv_id<mmv_fn<mul_vec_q_n_f32_impl<block_q4_0, N_R0_Q4_0>>>;
 template [[host_name("kernel_mul_mv_id_q4_1_f32")]]    kernel kernel_mul_mv_id_t kernel_mul_mv_id<mmv_fn<mul_vec_q_n_f32_impl<block_q4_1, N_R0_Q4_1>>>;
 template [[host_name("kernel_mul_mv_id_q5_0_f32")]]    kernel kernel_mul_mv_id_t kernel_mul_mv_id<mmv_fn<mul_vec_q_n_f32_impl<block_q5_0, N_R0_Q5_0>>>;

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

@@ -532,6 +532,7 @@ struct ggml_backend_opencl_context {
     bool fp16_support;
     bool has_vector_subgroup_broadcast;
     bool has_subgroup_shuffle = false;       // cl_khr_subgroup_shuffle or cl_qcom_subgroup_shuffle
+    bool has_integer_dot      = false;       // cl_khr_integer_dot_product or cl_qcom_dot_product8
     bool has_qcom_subgroup_shuffle = false;  // specifically cl_qcom_subgroup_shuffle
     bool disable_fusion;
 
@@ -834,7 +835,7 @@ struct ggml_backend_opencl_context {
     cl_kernel kernel_gemv_moe_q5_1_f32_ns, kernel_gemm_moe_q5_1_f32_ns;
     cl_kernel kernel_gemv_moe_q4_k_f32_ns, kernel_gemm_moe_q4_k_f32_ns, kernel_gemm_moe_q4_k_f32_ns_bin;
     cl_kernel kernel_gemv_moe_q4_k_f32_ns_wimg = nullptr;  // weight-as-texture MoE decode GEMV (opt-in)
-    cl_kernel kernel_gemm_moe_q4_k_q8_1_dp4a;    // dp4a (int8) prefill GEMM variant
+    cl_kernel kernel_gemm_moe_q4_k_q8_1_dp4a = nullptr;    // dp4a (int8) prefill GEMM variant
     cl_kernel kernel_moe_reorder_quant_a_q8_1;   // fused reorder + q8_1 quant for the dp4a GEMM
     cl_kernel kernel_gemm_moe_q8_1_dp4a_q80 = nullptr;   // generic dp4a MoE GEMM (MOE_QT=80), opt-in
     cl_kernel kernel_moe_expand_scale_q8_0 = nullptr;    // q8_0 per-block d -> uniform scale[16]
@@ -844,12 +845,12 @@ struct ggml_backend_opencl_context {
     cl_kernel kernel_moe_expand_scale_q5_K = nullptr;    // q5_K 6-bit s[] -> uniform scale[16]/min[8]
     cl_kernel kernel_gemv_moe_q5_k_f32_ns, kernel_gemm_moe_q5_k_f32_ns;
     cl_kernel kernel_gemv_moe_q6_k_f32_ns, kernel_gemm_moe_q6_k_f32_ns;
-    cl_kernel kernel_gemm_moe_q6_k_q8_1_dp4a;    // dp4a (int8) q6_K MoE prefill GEMM
+    cl_kernel kernel_gemm_moe_q6_k_q8_1_dp4a = nullptr;    // dp4a (int8) q6_K MoE prefill GEMM
     cl_kernel kernel_gemv_moe_mxfp4_f32, kernel_gemm_moe_mxfp4_f32;
     cl_kernel kernel_gemv_moe_mxfp4_f32_ns, kernel_gemm_moe_mxfp4_f32_ns, kernel_gemm_moe_mxfp4_f32_ns_bin;
     cl_kernel kernel_gemv_moe_mxfp4_f32_ns_wimg = nullptr;      // weight-as-texture MoE decode GEMV
-    cl_kernel kernel_gemm_moe_mxfp4_q8_1_dp4a;   // dp4a (int8) mxfp4 MoE prefill GEMM
-    cl_kernel kernel_gemm_moe_q4_0_q8_1_dp4a;    // dp4a (int8) q4_0 MoE prefill GEMM
+    cl_kernel kernel_gemm_moe_mxfp4_q8_1_dp4a = nullptr;   // dp4a (int8) mxfp4 MoE prefill GEMM
+    cl_kernel kernel_gemm_moe_q4_0_q8_1_dp4a = nullptr;    // dp4a (int8) q4_0 MoE prefill GEMM
     cl_kernel kernel_moe_reorder_b;
     cl_kernel kernel_moe_histogram, kernel_moe_scan, kernel_moe_fill, kernel_moe_scatter;
     cl_kernel kernel_moe_combine_f32 = nullptr;   // fused router-weight mul + cross-expert sum
@@ -1037,10 +1038,10 @@ struct ggml_backend_opencl_context {
     cl_kernel kernel_gemv_noshuffle_q1_0_f32;
     cl_kernel kernel_gemv_noshuffle_q4_k_f32;
     cl_kernel kernel_gemm_noshuffle_q4_k_f32;
-    cl_kernel kernel_gemm_noshuffle_q4_k_q8_1_dp4a;  // dp4a (int8) dense prefill GEMM
-    cl_kernel kernel_gemm_noshuffle_q4_k_q8_1_dp4a_wimg;  // dp4a dense prefill GEMM, weights via texture (X1 opt-in)
-    cl_kernel kernel_gemm_noshuffle_q5_k_q8_1_dp4a;  // dp4a (int8) dense q5_K prefill GEMM
-    cl_kernel kernel_gemm_noshuffle_q6_k_q8_1_dp4a;  // dp4a (int8) dense q6_K prefill GEMM
+    cl_kernel kernel_gemm_noshuffle_q4_k_q8_1_dp4a = nullptr;  // dp4a (int8) dense prefill GEMM
+    cl_kernel kernel_gemm_noshuffle_q4_k_q8_1_dp4a_wimg = nullptr;  // dp4a dense prefill GEMM, weights via texture (X1 opt-in)
+    cl_kernel kernel_gemm_noshuffle_q5_k_q8_1_dp4a = nullptr;  // dp4a (int8) dense q5_K prefill GEMM
+    cl_kernel kernel_gemm_noshuffle_q6_k_q8_1_dp4a = nullptr;  // dp4a (int8) dense q6_K prefill GEMM
     cl_kernel kernel_quant_a_q8_1;                    // plain activation q8_1 pre-pass
     cl_kernel kernel_gemv_noshuffle_q6_K_f32;
     cl_kernel kernel_gemm_noshuffle_q6_K_f32;
@@ -3490,7 +3491,7 @@ static void load_cl_kernels(ggml_backend_opencl_context *backend_ctx) {
     }
 
     // gemm_noshuffle_q5_0_q8_1_dp4a (dp4a dense q5_0 prefill GEMM)
-    {
+    if (backend_ctx->has_integer_dot) {
 #ifdef GGML_OPENCL_EMBED_KERNELS
         const std::string kernel_src {
             #include "gemm_noshuffle_q5_0_q8_1_dp4a.cl.h"
@@ -3580,7 +3581,7 @@ static void load_cl_kernels(ggml_backend_opencl_context *backend_ctx) {
     }
 
     // gemm_noshuffle_iq4_nl_q8_1_dp4a (dp4a dense IQ4_NL prefill GEMM)
-    {
+    if (backend_ctx->has_integer_dot) {
 #ifdef GGML_OPENCL_EMBED_KERNELS
         const std::string kernel_src {
             #include "gemm_noshuffle_iq4_nl_q8_1_dp4a.cl.h"
@@ -3595,7 +3596,7 @@ static void load_cl_kernels(ggml_backend_opencl_context *backend_ctx) {
     }
 
     // gemm_noshuffle_q4_0_q8_1_dp4a (dp4a dense q4_0 prefill GEMM)
-    {
+    if (backend_ctx->has_integer_dot) {
 #ifdef GGML_OPENCL_EMBED_KERNELS
         const std::string kernel_src {
             #include "gemm_noshuffle_q4_0_q8_1_dp4a.cl.h"
@@ -3708,7 +3709,7 @@ static void load_cl_kernels(ggml_backend_opencl_context *backend_ctx) {
     }
 
     // gemm_noshuffle_q4_k_q8_1_dp4a (dp4a dense prefill GEMM)
-    {
+    if (backend_ctx->has_integer_dot) {
 #ifdef GGML_OPENCL_EMBED_KERNELS
         const std::string kernel_src {
             #include "gemm_noshuffle_q4_k_q8_1_dp4a.cl.h"
@@ -3730,7 +3731,7 @@ static void load_cl_kernels(ggml_backend_opencl_context *backend_ctx) {
     }
 
     // gemm_noshuffle_q8_0_q8_1_dp4a (dp4a dense q8_0 prefill GEMM)
-    {
+    if (backend_ctx->has_integer_dot) {
 #ifdef GGML_OPENCL_EMBED_KERNELS
         const std::string kernel_src {
             #include "gemm_noshuffle_q8_0_q8_1_dp4a.cl.h"
@@ -3746,7 +3747,7 @@ static void load_cl_kernels(ggml_backend_opencl_context *backend_ctx) {
     }
 
     // gemm_noshuffle_q5_k_q8_1_dp4a (dp4a dense prefill GEMM for q5_K)
-    {
+    if (backend_ctx->has_integer_dot) {
 #ifdef GGML_OPENCL_EMBED_KERNELS
         const std::string kernel_src {
             #include "gemm_noshuffle_q5_k_q8_1_dp4a.cl.h"
@@ -3761,7 +3762,7 @@ static void load_cl_kernels(ggml_backend_opencl_context *backend_ctx) {
     }
 
     // gemm_noshuffle_q6_k_q8_1_dp4a (dp4a dense prefill GEMM for q6_K ffn_down/output)
-    {
+    if (backend_ctx->has_integer_dot) {
 #ifdef GGML_OPENCL_EMBED_KERNELS
         const std::string kernel_src {
             #include "gemm_noshuffle_q6_k_q8_1_dp4a.cl.h"
@@ -4091,7 +4092,7 @@ static void load_cl_kernels(ggml_backend_opencl_context *backend_ctx) {
     }
 
     // gemm_moe_q4_k_q8_1_dp4a (dp4a prefill GEMM)
-    {
+    if (backend_ctx->has_integer_dot) {
 #ifdef GGML_OPENCL_EMBED_KERNELS
         const std::string kernel_src {
             #include "gemm_moe_q4_k_q8_1_dp4a.cl.h"
@@ -4108,7 +4109,7 @@ static void load_cl_kernels(ggml_backend_opencl_context *backend_ctx) {
     }
 
     // gemm_moe_mxfp4_q8_1_dp4a (dp4a prefill GEMM)
-    {
+    if (backend_ctx->has_integer_dot) {
 #ifdef GGML_OPENCL_EMBED_KERNELS
         const std::string kernel_src {
             #include "gemm_moe_mxfp4_q8_1_dp4a.cl.h"
@@ -4125,7 +4126,7 @@ static void load_cl_kernels(ggml_backend_opencl_context *backend_ctx) {
     }
 
     // gemm_moe_q4_0_q8_1_dp4a (dp4a prefill GEMM)
-    {
+    if (backend_ctx->has_integer_dot) {
 #ifdef GGML_OPENCL_EMBED_KERNELS
         const std::string kernel_src {
             #include "gemm_moe_q4_0_q8_1_dp4a.cl.h"
@@ -4142,7 +4143,7 @@ static void load_cl_kernels(ggml_backend_opencl_context *backend_ctx) {
     }
 
     // gemm_moe_q8_1_dp4a (generic dp4a MoE GEMM; MOE_QT=80 -> q8_0 expert variant)
-    {
+    if (backend_ctx->has_integer_dot) {
 #ifdef GGML_OPENCL_EMBED_KERNELS
         const std::string kernel_src {
             #include "gemm_moe_q8_1_dp4a.cl.h"
@@ -4256,7 +4257,7 @@ static void load_cl_kernels(ggml_backend_opencl_context *backend_ctx) {
     }
 
     // gemm_moe_q6_k_q8_1_dp4a (dp4a q6_K MoE prefill GEMM)
-    {
+    if (backend_ctx->has_integer_dot) {
 #ifdef GGML_OPENCL_EMBED_KERNELS
         const std::string kernel_src {
             #include "gemm_moe_q6_k_q8_1_dp4a.cl.h"
@@ -5602,6 +5603,8 @@ static void ggml_opencl_print_backend_info(ggml_backend_opencl_device_context *
         backend_ctx->has_subgroup_shuffle ? "true" : "false");
     GGML_LOG_INFO("ggml_opencl: device FP16 support: %s\n",
         backend_ctx->fp16_support ? "true" : "false");
+    GGML_LOG_INFO("ggml_opencl: khr dot product support: %s\n",
+        backend_ctx->has_integer_dot ? "true" : "false");
     GGML_LOG_INFO("ggml_opencl: mem base addr align: %u\n",
         backend_ctx->alignment);
     GGML_LOG_INFO("ggml_opencl: global mem size: %zu MB\n",
@@ -5810,6 +5813,12 @@ static ggml_backend_opencl_context * ggml_cl_init(ggml_backend_dev_t dev) {
         strstr(ext_buffer, "cl_khr_subgroup_shuffle") != NULL ||
         backend_ctx->has_qcom_subgroup_shuffle;
 
+    // check for cl_khr_integer_dot_product
+    // cl_qcom_dot_product8 uses signed * unsigned
+    // while cl_khr_integer_dot_product uses signed * signed -- we stick with khr for now
+    backend_ctx->has_integer_dot =
+        strstr(ext_buffer, "cl_khr_integer_dot_product") != NULL;
+
     cl_uint base_align_in_bits;
     CL_CHECK(clGetDeviceInfo(device, CL_DEVICE_MEM_BASE_ADDR_ALIGN, sizeof(cl_uint), &base_align_in_bits, NULL));
     GGML_ASSERT(base_align_in_bits % 8u == 0);
@@ -7310,6 +7319,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;
@@ -10516,10 +10533,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);
@@ -15819,18 +15842,14 @@ static void ggml_cl_mul_mat_q4_0_f32_adreno(ggml_backend_t backend, const ggml_t
         CL_CHECK(clReleaseMemObject(b_sub_buf));
         CL_CHECK(clReleaseMemObject(b_img));
     } else {
-        // dp4a (int8) dense prefill GEMM: quant activations to q8_1, then the int8
-        // dp4a inner-loop GEMM, in place of the transpose + f16 half-dot kernel.
-        // q4_0 = d*(q-8); mirrors the IQ4_NL/q8_0 dense dp4a paths (+ the sum term).
-        // OPT-IN / DEFAULT OFF: correct, but neutral on X2E. q4_0's dequant
-        // ((q-8)*scale) is already trivial so the f16 GEMM is weight-BW-bound and the
-        // int8 ALU win has nothing to beat -- same as q5_0 dense (unlike IQ4_NL, whose
-        // codebook dequant is expensive enough for dp4a to help). Kept for A/B; force
-        // on with GGML_OPENCL_Q4_0_DENSE_DP4A=1. Needs N>8, K%32==0, M%64==0.
+        // dp4a (int8) dense prefill GEMM, default off
         static const char * q4_0_dense_dp4a_env = getenv("GGML_OPENCL_Q4_0_DENSE_DP4A");
-        const bool q4_0_dense_dp4a_on = q4_0_dense_dp4a_env
+        bool q4_0_dense_dp4a_on = q4_0_dense_dp4a_env
             ? (atoi(q4_0_dense_dp4a_env) != 0)
             : false;
+        // dot prod has to be available
+        q4_0_dense_dp4a_on = backend_ctx->has_integer_dot && q4_0_dense_dp4a_on;
+
         if (q4_0_dense_dp4a_on && backend_ctx->kernel_gemm_noshuffle_q4_0_q8_1_dp4a
                 && N > 8 && (K % 32 == 0) && (M % 64 == 0)) {
             cl_mem a_sub = nullptr;
@@ -16253,27 +16272,16 @@ static void ggml_cl_mul_mat_q5_0_f32_adreno(ggml_backend_t backend, const ggml_t
         CL_CHECK(clReleaseMemObject(b_sub_buf));
         CL_CHECK(clReleaseMemObject(b_img));
     } else {
-        // dp4a (int8) dense q5_0 prefill GEMM. Quantizes the [N,K] activations to
-        // q8_1 and runs the int8 dot instead of the f16 half-dot. Large-batch
-        // (ne1>8) only. q5_0 weight = (x-16)*d (x = nibble | hi<<4); x packed as a
-        // 0..31 byte (dp4a), the -16 centering folded into a single min term
-        // (d*16) via the q8_1 block sum. Reads the qs/qh/d buffers byte-identically
-        // to the f16 kernel (greedy byte-identical, MUL_MAT NMSE-OK).
-        //
-        // OPT-IN / DEFAULT OFF. Unlike q8_0/q4_K dense, dp4a is not a win for q5_0 on
-        // X2E: the q5_0 model is bottlenecked elsewhere, so the dense-GEMM int8 win
-        // has nothing to surface and the q8_1 prepass slightly hurts. Kept correct +
-        // opt-in for the X1 A/B (different texture-cache dynamic) and the
-        // weight-texture variant. Env: GGML_OPENCL_Q5_DENSE_DP4A=1.
-        // Weight-as-texture variant (X1 lever): routes the dominant qs nibble plane
-        // through an image1d_buffer (qh stays a buffer). Opt-in
-        // GGML_OPENCL_Q5_DENSE_DP4A_WIMG; when set it also forces the dp4a path on.
+        // dp4a (int8) dense q5_0 prefill GEMM, default off
         static const char * q5_dense_dp4a_env = getenv("GGML_OPENCL_Q5_DENSE_DP4A");
         static const char * q5_dense_wimg_env = getenv("GGML_OPENCL_Q5_DENSE_DP4A_WIMG");
         const bool q5_dense_wimg_on = q5_dense_wimg_env && (atoi(q5_dense_wimg_env) != 0);
-        const bool q5_dense_dp4a_on = q5_dense_wimg_on
+              bool q5_dense_dp4a_on = q5_dense_wimg_on
             ? true
             : (q5_dense_dp4a_env && (atoi(q5_dense_dp4a_env) != 0));
+        // dot prod has to be available
+        q5_dense_dp4a_on = backend_ctx->has_integer_dot && q5_dense_dp4a_on;
+
         if (q5_dense_dp4a_on && backend_ctx->kernel_gemm_noshuffle_q5_0_q8_1_dp4a
                 && N > 8 && (K % 32 == 0) && (M % 64 == 0)) {
             cl_mem a_sub = nullptr;
@@ -16708,15 +16716,14 @@ static void ggml_cl_mul_mat_iq4_nl_f32_adreno(ggml_backend_t backend, const ggml
     } else {
         // dp4a (int8) dense IQ4_NL prefill GEMM. Quantizes the [N,K] activations to
         // q8_1 and runs the int8 dot instead of the f16 half-dot. Large-batch
-        // (ne1>8) only. IQ4_NL weight = kvalues[nibble]*d; the codebook value IS the
-        // int8 (no min term), so this is the q8_0 dense case plus a nibble->int8 LUT
-        // unpack. Reads the q/d buffers byte-identically to the f16 kernel. No bin
-        // kernel for IQ4_NL -> baseline is f16, default ON for X2E (like q4_K/q6_K
-        // dense dp4a). X1 stays on f16. Env: GGML_OPENCL_IQ4NL_DENSE_DP4A.
+        // (ne1>8) only
         static const char * iq4nl_dense_dp4a_env = getenv("GGML_OPENCL_IQ4NL_DENSE_DP4A");
-        const bool iq4nl_dense_dp4a_on = iq4nl_dense_dp4a_env
+        bool iq4nl_dense_dp4a_on = iq4nl_dense_dp4a_env
             ? (atoi(iq4nl_dense_dp4a_env) != 0)
             : (backend_ctx->adreno_gen == ADRENO_GPU_GEN::X2E);
+        // dot prod has to be available
+        iq4nl_dense_dp4a_on = backend_ctx->has_integer_dot && iq4nl_dense_dp4a_on;
+
         if (iq4nl_dense_dp4a_on && backend_ctx->kernel_gemm_noshuffle_iq4_nl_q8_1_dp4a
                 && N > 8 && (K % 32 == 0) && (M % 64 == 0)) {
             cl_mem a_sub = nullptr;
@@ -16966,13 +16973,15 @@ static void ggml_cl_mul_mat_q8_0_f32_adreno(ggml_backend_t backend, const ggml_t
         static const char * q8_dense_wimg_env = getenv("GGML_OPENCL_Q8_DENSE_DP4A_WIMG");
         const bool q8_dense_wimg_on = q8_dense_wimg_env && (atoi(q8_dense_wimg_env) != 0);
 
-        const bool q8_bin_loaded   = (backend_ctx->kernel_gemm_noshuffle_q8_0_f32_bin != nullptr);
+        const bool q8_bin_loaded = (backend_ctx->kernel_gemm_noshuffle_q8_0_f32_bin != nullptr);
         // bin kernel takes precedence
-        const bool q8_dense_dp4a_on = q8_dense_wimg_on
+        bool q8_dense_dp4a_on = q8_dense_wimg_on
             ? true
             : q8_dense_dp4a_env
             ? (atoi(q8_dense_dp4a_env) != 0)
             : (backend_ctx->adreno_gen == ADRENO_GPU_GEN::X2E && !q8_bin_loaded);
+        // dot prod has to be available
+        q8_dense_dp4a_on = backend_ctx->has_integer_dot && q8_dense_dp4a_on;
 
         if (q8_dense_dp4a_on && backend_ctx->kernel_gemm_noshuffle_q8_0_q8_1_dp4a
                 && N > 8 && (K % 32 == 0) && (M % 64 == 0)) {
@@ -17379,13 +17388,16 @@ static void ggml_cl_mul_mat_q4_k_f32_adreno(ggml_backend_t backend, const ggml_t
         static const char * q4k_dense_dp4a_env = getenv("GGML_OPENCL_Q4K_DENSE_DP4A");
         static const char * q4k_dense_wimg_env = getenv("GGML_OPENCL_Q4K_DENSE_DP4A_WIMG");
 
-        const bool          q4k_dense_wimg_on  = q4k_dense_wimg_env && (atoi(q4k_dense_wimg_env) != 0);
-        const bool          q4k_dense_dp4a_on  = q4k_dense_wimg_on
+        const bool q4k_dense_wimg_on = q4k_dense_wimg_env && (atoi(q4k_dense_wimg_env) != 0);
+              bool q4k_dense_dp4a_on = q4k_dense_wimg_on
             ? true
             : q4k_dense_dp4a_env
             ? (atoi(q4k_dense_dp4a_env) != 0)
             : (backend_ctx->adreno_gen == ADRENO_GPU_GEN::X2E);
 
+        // dp4 has to be available
+        q4k_dense_dp4a_on = backend_ctx->has_integer_dot && q4k_dense_dp4a_on;
+
         // Min N for the dp4a prefill GEMM, default 9, i.e., ne1 > 8
         static const char * q4k_dp4a_minn_env = getenv("GGML_OPENCL_Q4K_DP4A_MINN");
         const int           q4k_dp4a_minn     = q4k_dp4a_minn_env ? atoi(q4k_dp4a_minn_env) : 9;
@@ -17608,9 +17620,11 @@ static void ggml_cl_mul_mat_q6_K_f32_adreno(ggml_backend_t backend, const ggml_t
 
         // dp4a (int8) dense q6_K prefill GEMM
         static const char * q6k_dense_dp4a_env = getenv("GGML_OPENCL_Q6K_DENSE_DP4A");
-        static const bool   q6k_dense_dp4a_on  = (q6k_dense_dp4a_env != nullptr)
+                     bool   q6k_dense_dp4a_on  = (q6k_dense_dp4a_env != nullptr)
                                                    ? (atoi(q6k_dense_dp4a_env) != 0)
                                                    : (backend_ctx->adreno_gen != ADRENO_GPU_GEN::X1E);
+        // dot prod has to be available
+        q6k_dense_dp4a_on = backend_ctx->has_integer_dot && q6k_dense_dp4a_on;
 
         const bool is_output_w_dp4a = strncmp(src0->name, "output", 6) == 0 ||
                                       strncmp(src0->name, "token_embd", 10) == 0;
@@ -17901,9 +17915,11 @@ static void ggml_cl_mul_mat_q5_K_f32_adreno(ggml_backend_t backend, const ggml_t
 
         // dp4a (int8) dense q5_K prefill GEMM
         static const char * q5k_dense_dp4a_env = getenv("GGML_OPENCL_Q5K_DENSE_DP4A");
-        const bool          q5k_dense_dp4a_on  = q5k_dense_dp4a_env
+                     bool   q5k_dense_dp4a_on  = q5k_dense_dp4a_env
             ? (atoi(q5k_dense_dp4a_env) != 0)
             : (backend_ctx->adreno_gen == ADRENO_GPU_GEN::X2E);
+        // dot prod has to be available
+        q5k_dense_dp4a_on = backend_ctx->has_integer_dot && q5k_dense_dp4a_on;
 
         if (q5k_dense_dp4a_on && ne1 > 8 && (ne00 % 32 == 0) && (ne01 % 64 == 0)) {
             const int Mm = ne01, Nn = ne1, Kk = ne00;
@@ -20640,6 +20656,8 @@ static void ggml_cl_mul_mat_id(ggml_backend_t backend, const ggml_tensor * src0,
                     bool use_moe_dp4a = q4_0_moe_dp4a_env
                         ? (atoi(q4_0_moe_dp4a_env) != 0)
                         : (backend_ctx->adreno_gen == ADRENO_GPU_GEN::X2E);
+                    // dot prod has to be available
+                    use_moe_dp4a = backend_ctx->has_integer_dot && use_moe_dp4a;
                     // bin kernel takes precedence
                     use_moe_dp4a = use_moe_dp4a && backend_ctx->kernel_gemm_moe_q4_0_f32_ns_bin == nullptr;
 
@@ -21815,6 +21833,8 @@ static void ggml_cl_mul_mat_id(ggml_backend_t backend, const ggml_tensor * src0,
                     bool  use_moe_dp4a = (q4k_moe_dp4a_env != nullptr)
                                          ? (atoi(q4k_moe_dp4a_env) != 0)
                                          : (backend_ctx->adreno_gen == ADRENO_GPU_GEN::X2E || backend_ctx->adreno_gen == ADRENO_GPU_GEN::X1E);
+                    // dot prod has to be available
+                    use_moe_dp4a = backend_ctx->has_integer_dot && use_moe_dp4a;
                     // bin kernel takes precedence
                     use_moe_dp4a = use_moe_dp4a && backend_ctx->kernel_gemm_moe_q4_k_f32_ns_bin == nullptr;
 
@@ -22316,10 +22336,12 @@ static void ggml_cl_mul_mat_id(ggml_backend_t backend, const ggml_tensor * src0,
 
                     // dp4a (int8) q6_K MoE prefill GEMM
                     static const char * q6k_moe_dp4a_env = getenv("GGML_OPENCL_Q6K_MOE_DP4A");
-                    static const bool   use_moe_dp4a = (q6k_moe_dp4a_env != nullptr)
+                                 bool   use_moe_dp4a = (q6k_moe_dp4a_env != nullptr)
                                                          ? (atoi(q6k_moe_dp4a_env) != 0)
                                                          : (backend_ctx->adreno_gen == ADRENO_GPU_GEN::X2E
                                                             || backend_ctx->adreno_gen == ADRENO_GPU_GEN::X1E);
+                    // dot prod has to be available
+                    use_moe_dp4a = backend_ctx->has_integer_dot && use_moe_dp4a;
 
                     cl_buffer_region region;
                     region.origin = 0;
@@ -22569,6 +22591,8 @@ static void ggml_cl_mul_mat_id(ggml_backend_t backend, const ggml_tensor * src0,
                     bool use_moe_dp4a = mxfp4_moe_dp4a_env
                         ? (atoi(mxfp4_moe_dp4a_env) != 0)
                         : (backend_ctx->adreno_gen == ADRENO_GPU_GEN::X2E);
+                    // dot prod has to be available
+                    use_moe_dp4a = backend_ctx->has_integer_dot && use_moe_dp4a;
                     // bin kernel takes precedence
                     use_moe_dp4a = use_moe_dp4a && backend_ctx->kernel_gemm_moe_mxfp4_f32_ns_bin == nullptr;
 

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/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-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-vec.hpp

@@ -15,13 +15,11 @@
 
 namespace syclex = sycl::ext::oneapi::experimental;
 
-static int ggml_sycl_fattn_vec_get_nthreads_host(const int cc) {
-    return 128;
-    GGML_UNUSED(cc);
-}
-
-static constexpr int ggml_sycl_fattn_vec_get_nthreads_device() {
-    return 128;
+static int ggml_sycl_fattn_vec_get_nthreads_device(gpu_arch arch) {
+    // Xe2 (Battlemage, Lunar Lake) runs the flash-attention vec kernel best with a 256-thread work group.
+    return (arch == gpu_arch::intel_gpu_bmg_g21 ||
+            arch == gpu_arch::intel_gpu_bmg_g31 ||
+            arch == gpu_arch::intel_gpu_lnl_m) ? 256 : 128;
 }
 
 // Currenlty llvm with the amdgcn target dose not support unrolling loops
@@ -36,7 +34,8 @@ template <int D,
           int type_K,
           int type_V,
           bool use_logit_softcap,
-          int warp_size>  // D == head size
+          int warp_size,
+          int nthreads>  // D == head size
 static void flash_attn_ext_vec(const char* __restrict__ Q,
                         const char* __restrict__ K,
                         const char* __restrict__ V,
@@ -99,7 +98,6 @@ static void flash_attn_ext_vec(const char* __restrict__ Q,
     constexpr int nthreads_KQ_q = (D/4 < warp_size ? D/4 : warp_size);
     constexpr int nthreads_V_q  = (D/4 < warp_size ? D/4 : warp_size);
 
-    constexpr int nthreads    = ggml_sycl_fattn_vec_get_nthreads_device();
     constexpr int nthreads_KQ = type_K == GGML_TYPE_F16 ? 128 / cpy_nb : nthreads_KQ_q;
     constexpr int nthreads_V  = type_V == GGML_TYPE_F16 ? 128 / cpy_nb : nthreads_V_q;
 
@@ -581,24 +579,34 @@ static void flash_attn_ext_vec(const char* __restrict__ Q,
 #endif // __clang__
 
 
+
 template <int D, int cols_per_block, int type_K, int type_V, bool use_logit_softcap>
 void ggml_sycl_flash_attn_ext_vec_case_impl(ggml_backend_sycl_context & ctx, ggml_tensor * dst) {
 
-    const int warp_size = WARP_16_SIZE; //better performance than WARP_32_SIZE
-
-    const int cc = ggml_sycl_info().devices[ggml_sycl_get_device()].cc;
-
-    const int nthreads = ggml_sycl_fattn_vec_get_nthreads_host(cc);
-    const int nwarps   = nthreads / warp_size;
+    constexpr int warp_size = WARP_16_SIZE; //better performance than WARP_32_SIZE
 
     const bool need_f16_K = type_K == GGML_TYPE_F16;
     const bool need_f16_V = type_V == GGML_TYPE_F16;
     constexpr size_t nbytes_shared = 0;
 
-    launch_fattn<D, cols_per_block, 1,
-                 flash_attn_ext_vec<D, cols_per_block, type_K, type_V,
-                                    use_logit_softcap, warp_size>, warp_size>(
-        ctx, dst, nwarps, nbytes_shared, D, need_f16_K, need_f16_V, false);
+    const auto arch = ggml_sycl_info().devices[ctx.device].hw_info.arch;
+    const int nthreads = ggml_sycl_fattn_vec_get_nthreads_device(arch);
+    // 256 threads would overflow the 64 KB work-group local memory at D == 512, so keep 128 there.
+    if (D <= 256 && nthreads == 256) {
+        constexpr int nthreads_hw = 256;
+        constexpr int nwarps = nthreads_hw / warp_size;
+        launch_fattn<D, cols_per_block, 1,
+                     flash_attn_ext_vec<D, cols_per_block, type_K, type_V,
+                                        use_logit_softcap, warp_size, nthreads_hw>, warp_size>(
+            ctx, dst, nwarps, nbytes_shared, D, need_f16_K, need_f16_V, false);
+    } else {
+        constexpr int nthreads_hw = 128;
+        constexpr int nwarps = nthreads_hw / warp_size;
+        launch_fattn<D, cols_per_block, 1,
+                     flash_attn_ext_vec<D, cols_per_block, type_K, type_V,
+                                        use_logit_softcap, warp_size, nthreads_hw>, warp_size>(
+            ctx, dst, nwarps, nbytes_shared, D, need_f16_K, need_f16_V, false);
+    }
 }
 
 template <int D, int type_K, int type_V>

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

@@ -723,6 +723,7 @@ struct vk_device_struct {
     bool uma;
     bool prefer_host_memory;
     bool float_controls_rte_fp16;
+    bool float_controls_denorm_preserve_fp16;
     bool subgroup_basic;
     bool subgroup_arithmetic;
     bool subgroup_shuffle;
@@ -868,8 +869,9 @@ struct vk_device_struct {
     vk_pipeline pipeline_cpy_f32_quant[GGML_TYPE_COUNT];
     vk_pipeline pipeline_cpy_quant_f32[GGML_TYPE_COUNT];
     vk_pipeline pipeline_cpy_transpose_16, pipeline_cpy_transpose_32;
-    vk_pipeline pipeline_set_rows_i32[GGML_TYPE_COUNT];
-    vk_pipeline pipeline_set_rows_i64[GGML_TYPE_COUNT];
+    // [src0 0=fp32,1=fp16][dst]
+    vk_pipeline pipeline_set_rows_i32[2][GGML_TYPE_COUNT];
+    vk_pipeline pipeline_set_rows_i64[2][GGML_TYPE_COUNT];
     vk_pipeline pipeline_norm_f32;
     vk_pipeline pipeline_group_norm_f32;
     vk_pipeline pipeline_rms_norm_f32;
@@ -959,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;
@@ -2595,10 +2598,10 @@ static void ggml_vk_create_pipeline_func(vk_device& device, vk_pipeline& pipelin
 
     vk::ShaderModuleCreateInfo shader_module_create_info({}, spv_size, reinterpret_cast<const uint32_t *>(spv_data));
 
-    // Patch SPIR-V to enable RTE rounding for FP16, avoiding the need for
-    // separate shader variants compiled with -DRTE16.
+    // Patch SPIR-V to enable supported FP16 float controls, avoiding the need
+    // for separate shader variants.
     std::vector<uint32_t> spirv;
-    if (device->float_controls_rte_fp16) {
+    if (device->float_controls_rte_fp16 || device->float_controls_denorm_preserve_fp16) {
         const uint32_t* spv_words = reinterpret_cast<const uint32_t *>(spv_data);
         size_t word_count = spv_size / sizeof(uint32_t);
         spirv.assign(spv_words, spv_words + word_count);
@@ -2635,9 +2638,17 @@ static void ggml_vk_create_pipeline_func(vk_device& device, vk_pipeline& pipelin
 
         // Insert from latest position first so earlier indices stay valid.
 
-        // OpExecutionMode %entrypoint RoundingModeRTE 16
-        uint32_t exec_mode[] = { (4u << spv::WordCountShift) | spv::OpExecutionMode, entry_point_id, spv::ExecutionModeRoundingModeRTE, 16 };
-        spirv.insert(spirv.begin() + exec_insert_pos, std::begin(exec_mode), std::end(exec_mode));
+        if (device->float_controls_rte_fp16) {
+            // OpExecutionMode %entrypoint RoundingModeRTE 16
+            uint32_t exec_mode[] = { (4u << spv::WordCountShift) | spv::OpExecutionMode, entry_point_id, spv::ExecutionModeRoundingModeRTE, 16 };
+            spirv.insert(spirv.begin() + exec_insert_pos, std::begin(exec_mode), std::end(exec_mode));
+        }
+
+        if (device->float_controls_denorm_preserve_fp16) {
+            // OpExecutionMode %entrypoint DenormPreserve 16
+            uint32_t exec_mode[] = { (4u << spv::WordCountShift) | spv::OpExecutionMode, entry_point_id, spv::ExecutionModeDenormPreserve, 16 };
+            spirv.insert(spirv.begin() + exec_insert_pos, std::begin(exec_mode), std::end(exec_mode));
+        }
 
         // OpExtension "SPV_KHR_float_controls"
         const char ext_str[] = "SPV_KHR_float_controls";
@@ -2647,9 +2658,17 @@ static void ggml_vk_create_pipeline_func(vk_device& device, vk_pipeline& pipelin
         memcpy(&extension[1], ext_str, sizeof(ext_str));
         spirv.insert(spirv.begin() + ext_insert_pos, extension.begin(), extension.end());
 
-        // OpCapability RoundingModeRTE
-        uint32_t capability[] = { (2u << spv::WordCountShift) | spv::OpCapability, spv::CapabilityRoundingModeRTE };
-        spirv.insert(spirv.begin() + cap_insert_pos, std::begin(capability), std::end(capability));
+        if (device->float_controls_rte_fp16) {
+            // OpCapability RoundingModeRTE
+            uint32_t capability[] = { (2u << spv::WordCountShift) | spv::OpCapability, spv::CapabilityRoundingModeRTE };
+            spirv.insert(spirv.begin() + cap_insert_pos, std::begin(capability), std::end(capability));
+        }
+
+        if (device->float_controls_denorm_preserve_fp16) {
+            // OpCapability DenormPreserve
+            uint32_t capability[] = { (2u << spv::WordCountShift) | spv::OpCapability, spv::CapabilityDenormPreserve };
+            spirv.insert(spirv.begin() + cap_insert_pos, std::begin(capability), std::end(capability));
+        }
 
         shader_module_create_info = vk::ShaderModuleCreateInfo({}, spirv.size() * sizeof(uint32_t), spirv.data());
     }
@@ -5187,20 +5206,22 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
     ggml_vk_create_pipeline(device, device->pipeline_cpy_f32_quant[GGML_TYPE_Q8_0], "cpy_f32_q8_0", cpy_f32_q8_0_len, cpy_f32_q8_0_data, "main", 2, sizeof(vk_op_unary_push_constants), {32, 1, 1}, {}, 1);
     ggml_vk_create_pipeline(device, device->pipeline_cpy_f32_quant[GGML_TYPE_IQ4_NL], "cpy_f32_iq4_nl", cpy_f32_iq4_nl_len, cpy_f32_iq4_nl_data, "main", 2, sizeof(vk_op_unary_push_constants), {32, 1, 1}, {}, 1);
 
-#define SET_ROWS(itype) \
-        ggml_vk_create_pipeline(device, device->pipeline_set_rows ## itype [GGML_TYPE_F32],  "set_rows_f32" #itype,  set_rows_f32 ## itype ## _len,  set_rows_f32 ## itype ## _data,  "main", 3, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {1}, 1, true); \
-        ggml_vk_create_pipeline(device, device->pipeline_set_rows ## itype [GGML_TYPE_F16],  "set_rows_f16" #itype,  set_rows_f16 ## itype ## _len,  set_rows_f16 ## itype ## _data,  "main", 3, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {1}, 1, true); \
-        ggml_vk_create_pipeline(device, device->pipeline_set_rows ## itype [GGML_TYPE_BF16], "set_rows_bf16" #itype, set_rows_bf16 ## itype ## _len, set_rows_bf16 ## itype ## _data, "main", 3, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {1}, 1, true); \
-        ggml_vk_create_pipeline(device, device->pipeline_set_rows ## itype [GGML_TYPE_Q1_0], "set_rows_q1_0" #itype, set_rows_q1_0 ## itype ## _len, set_rows_q1_0 ## itype ## _data, "main", 3, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {1}, 1, true); \
-        ggml_vk_create_pipeline(device, device->pipeline_set_rows ## itype [GGML_TYPE_Q4_0], "set_rows_q4_0" #itype, set_rows_q4_0 ## itype ## _len, set_rows_q4_0 ## itype ## _data, "main", 3, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {1}, 1, true); \
-        ggml_vk_create_pipeline(device, device->pipeline_set_rows ## itype [GGML_TYPE_Q4_1], "set_rows_q4_1" #itype, set_rows_q4_1 ## itype ## _len, set_rows_q4_1 ## itype ## _data, "main", 3, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {1}, 1, true); \
-        ggml_vk_create_pipeline(device, device->pipeline_set_rows ## itype [GGML_TYPE_Q5_0], "set_rows_q5_0" #itype, set_rows_q5_0 ## itype ## _len, set_rows_q5_0 ## itype ## _data, "main", 3, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {1}, 1, true); \
-        ggml_vk_create_pipeline(device, device->pipeline_set_rows ## itype [GGML_TYPE_Q5_1], "set_rows_q5_1" #itype, set_rows_q5_1 ## itype ## _len, set_rows_q5_1 ## itype ## _data, "main", 3, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {1}, 1, true); \
-        ggml_vk_create_pipeline(device, device->pipeline_set_rows ## itype [GGML_TYPE_Q8_0], "set_rows_q8_0" #itype, set_rows_q8_0 ## itype ## _len, set_rows_q8_0 ## itype ## _data, "main", 3, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {1}, 1, true); \
-        ggml_vk_create_pipeline(device, device->pipeline_set_rows ## itype [GGML_TYPE_IQ4_NL], "set_rows_iq4_nl" #itype, set_rows_iq4_nl ## itype ## _len, set_rows_iq4_nl ## itype ## _data, "main", 3, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {1}, 1, true);
+#define SET_ROWS(src_idx, src, itype) \
+        ggml_vk_create_pipeline(device, device->pipeline_set_rows ## itype [src_idx][GGML_TYPE_F32],  "set_rows_" #src "_f32" #itype,  set_rows_ ## src ## _f32 ## itype ## _len,  set_rows_ ## src ## _f32 ## itype ## _data,  "main", 3, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {1}, 1, true); \
+        ggml_vk_create_pipeline(device, device->pipeline_set_rows ## itype [src_idx][GGML_TYPE_F16],  "set_rows_" #src "_f16" #itype,  set_rows_ ## src ## _f16 ## itype ## _len,  set_rows_ ## src ## _f16 ## itype ## _data,  "main", 3, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {1}, 1, true); \
+        ggml_vk_create_pipeline(device, device->pipeline_set_rows ## itype [src_idx][GGML_TYPE_BF16], "set_rows_" #src "_bf16" #itype, set_rows_ ## src ## _bf16 ## itype ## _len, set_rows_ ## src ## _bf16 ## itype ## _data, "main", 3, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {1}, 1, true); \
+        ggml_vk_create_pipeline(device, device->pipeline_set_rows ## itype [src_idx][GGML_TYPE_Q1_0], "set_rows_" #src "_q1_0" #itype, set_rows_ ## src ## _q1_0 ## itype ## _len, set_rows_ ## src ## _q1_0 ## itype ## _data, "main", 3, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {1}, 1, true); \
+        ggml_vk_create_pipeline(device, device->pipeline_set_rows ## itype [src_idx][GGML_TYPE_Q4_0], "set_rows_" #src "_q4_0" #itype, set_rows_ ## src ## _q4_0 ## itype ## _len, set_rows_ ## src ## _q4_0 ## itype ## _data, "main", 3, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {1}, 1, true); \
+        ggml_vk_create_pipeline(device, device->pipeline_set_rows ## itype [src_idx][GGML_TYPE_Q4_1], "set_rows_" #src "_q4_1" #itype, set_rows_ ## src ## _q4_1 ## itype ## _len, set_rows_ ## src ## _q4_1 ## itype ## _data, "main", 3, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {1}, 1, true); \
+        ggml_vk_create_pipeline(device, device->pipeline_set_rows ## itype [src_idx][GGML_TYPE_Q5_0], "set_rows_" #src "_q5_0" #itype, set_rows_ ## src ## _q5_0 ## itype ## _len, set_rows_ ## src ## _q5_0 ## itype ## _data, "main", 3, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {1}, 1, true); \
+        ggml_vk_create_pipeline(device, device->pipeline_set_rows ## itype [src_idx][GGML_TYPE_Q5_1], "set_rows_" #src "_q5_1" #itype, set_rows_ ## src ## _q5_1 ## itype ## _len, set_rows_ ## src ## _q5_1 ## itype ## _data, "main", 3, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {1}, 1, true); \
+        ggml_vk_create_pipeline(device, device->pipeline_set_rows ## itype [src_idx][GGML_TYPE_Q8_0], "set_rows_" #src "_q8_0" #itype, set_rows_ ## src ## _q8_0 ## itype ## _len, set_rows_ ## src ## _q8_0 ## itype ## _data, "main", 3, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {1}, 1, true); \
+        ggml_vk_create_pipeline(device, device->pipeline_set_rows ## itype [src_idx][GGML_TYPE_IQ4_NL], "set_rows_" #src "_iq4_nl" #itype, set_rows_ ## src ## _iq4_nl ## itype ## _len, set_rows_ ## src ## _iq4_nl ## itype ## _data, "main", 3, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {1}, 1, true);
 
-    SET_ROWS(_i32)
-    SET_ROWS(_i64)
+    SET_ROWS(0, f32, _i32)
+    SET_ROWS(0, f32, _i64)
+    SET_ROWS(1, f16, _i32)
+    SET_ROWS(1, f16, _i64)
 #undef SET_ROWS
 
 
@@ -5459,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);
@@ -6031,6 +6054,7 @@ static vk_device ggml_vk_get_device(size_t idx) {
             device->shader_core_count = 0;
         }
         device->float_controls_rte_fp16 = vk12_props.shaderRoundingModeRTEFloat16;
+        device->float_controls_denorm_preserve_fp16 = vk12_props.shaderDenormPreserveFloat16;
 
         device->subgroup_basic = (vk11_props.subgroupSupportedStages & vk::ShaderStageFlagBits::eCompute) &&
                                  (vk11_props.subgroupSupportedOperations & vk::SubgroupFeatureFlagBits::eBasic);
@@ -10724,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;
@@ -10843,10 +10872,17 @@ static vk_pipeline ggml_vk_op_get_pipeline(ggml_backend_vk_context * ctx, const
     case GGML_OP_DUP:
         return ggml_vk_get_cpy_pipeline(ctx, src0, dst, dst->type);
     case GGML_OP_SET_ROWS:
-        if (src1->type == GGML_TYPE_I64) {
-            return ctx->device->pipeline_set_rows_i64[dst->type];
-        } else {
-            return ctx->device->pipeline_set_rows_i32[dst->type];
+        {
+            if (src0->type != GGML_TYPE_F32 && src0->type != GGML_TYPE_F16) {
+                return nullptr;
+            }
+            const int src_idx = src0->type == GGML_TYPE_F16;
+            if (src1->type == GGML_TYPE_I64) {
+                return ctx->device->pipeline_set_rows_i64[src_idx][dst->type];
+            } else if (src1->type == GGML_TYPE_I32) {
+                return ctx->device->pipeline_set_rows_i32[src_idx][dst->type];
+            }
+            return nullptr;
         }
     case GGML_OP_SILU_BACK:
         if (src0->type == GGML_TYPE_F32 && src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) {
@@ -11673,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:
@@ -11986,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);
@@ -14773,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);
 
@@ -17500,24 +17558,25 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm
             return op->type == GGML_TYPE_F32 && op->src[0]->type == GGML_TYPE_F32;
         case GGML_OP_SET_ROWS:
             {
-                if (op->src[0]->type == GGML_TYPE_F32) {
-                    switch (op->type) {
-                        case GGML_TYPE_F32:
-                        case GGML_TYPE_F16:
-                        case GGML_TYPE_BF16:
-                        case GGML_TYPE_Q1_0:
-                        case GGML_TYPE_Q4_0:
-                        case GGML_TYPE_Q4_1:
-                        case GGML_TYPE_Q5_0:
-                        case GGML_TYPE_Q5_1:
-                        case GGML_TYPE_Q8_0:
-                        case GGML_TYPE_IQ4_NL:
-                            return true;
-                        default:
-                            return false;
-                    }
+                if ((op->src[0]->type != GGML_TYPE_F32 && op->src[0]->type != GGML_TYPE_F16) ||
+                    (op->src[1]->type != GGML_TYPE_I32 && op->src[1]->type != GGML_TYPE_I64)) {
+                    return false;
+                }
+                switch (op->type) {
+                    case GGML_TYPE_F32:
+                    case GGML_TYPE_F16:
+                    case GGML_TYPE_BF16:
+                    case GGML_TYPE_Q1_0:
+                    case GGML_TYPE_Q4_0:
+                    case GGML_TYPE_Q4_1:
+                    case GGML_TYPE_Q5_0:
+                    case GGML_TYPE_Q5_1:
+                    case GGML_TYPE_Q8_0:
+                    case GGML_TYPE_IQ4_NL:
+                        return true;
+                    default:
+                        return false;
                 }
-                return false;
             }
         case GGML_OP_CONT:
         case GGML_OP_CPY:
@@ -17626,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/copy_to_quant.comp

@@ -10,7 +10,7 @@ layout(local_size_x = 32, local_size_y = 1, local_size_z = 1) in;
 const uint BLOCK_SIZE = 32;
 #endif
 
-layout (binding = 0) readonly buffer S {float data_s[];};
+layout (binding = 0) readonly buffer S {S_TYPE data_s[];};
 
 #if defined(SET_ROWS)
 #include "generic_binary_head.glsl"
@@ -35,7 +35,7 @@ void quantize(uint dst_idx, uint src_idx)
     float vmax = 0.0;
 
     [[unroll]] for (int j = 0; j < QUANT_K_Q4_0; ++j) {
-        const float v = data_s[src_idx + j];
+        const float v = float(data_s[src_idx + j]);
         if (amax < abs(v)) {
             amax = abs(v);
             vmax = v;
@@ -48,8 +48,8 @@ void quantize(uint dst_idx, uint src_idx)
     data_q[dst_idx].d = float16_t(d);
 
     [[unroll]] for (int j = 0; j < QUANT_K_Q4_0/2; ++j) {
-        const float x0 = data_s[src_idx + 0              + j]*id;
-        const float x1 = data_s[src_idx + QUANT_K_Q4_0/2 + j]*id;
+        const float x0 = float(data_s[src_idx + 0              + j])*id;
+        const float x1 = float(data_s[src_idx + QUANT_K_Q4_0/2 + j])*id;
 
         const uint xi0 = min(15, int(x0 + 8.5));
         const uint xi1 = min(15, int(x1 + 8.5));
@@ -66,7 +66,7 @@ void quantize(uint dst_idx, uint src_idx)
     float vmax = -vmin;
 
     [[unroll]] for (int j = 0; j < QUANT_K_Q4_1; ++j) {
-        const float v = data_s[src_idx + j];
+        const float v = float(data_s[src_idx + j]);
 
         if (v < vmin) vmin = v;
         if (v > vmax) vmax = v;
@@ -79,8 +79,8 @@ void quantize(uint dst_idx, uint src_idx)
     data_q[dst_idx].m = float16_t(vmin);
 
     [[unroll]] for (int j = 0; j < QUANT_K_Q4_1/2; ++j) {
-        const float x0 = (data_s[src_idx + 0              + j] - vmin)*id;
-        const float x1 = (data_s[src_idx + QUANT_K_Q4_1/2 + j] - vmin)*id;
+        const float x0 = (float(data_s[src_idx + 0              + j]) - vmin)*id;
+        const float x1 = (float(data_s[src_idx + QUANT_K_Q4_1/2 + j]) - vmin)*id;
 
         const uint xi0 = min(15, int(x0 + 0.5));
         const uint xi1 = min(15, int(x1 + 0.5));
@@ -97,7 +97,7 @@ void quantize(uint dst_idx, uint src_idx)
     float vmax = 0.0;
 
     [[unroll]] for (int j = 0; j < QUANT_K_Q5_0; ++j) {
-        const float v = data_s[src_idx + j];
+        const float v = float(data_s[src_idx + j]);
         if (amax < abs(v)) {
             amax = abs(v);
             vmax = v;
@@ -111,8 +111,8 @@ void quantize(uint dst_idx, uint src_idx)
 
     uint32_t qh = 0;
     [[unroll]] for (int j = 0; j < QUANT_K_Q5_0/2; ++j) {
-        const float x0 = data_s[src_idx + 0              + j]*id;
-        const float x1 = data_s[src_idx + QUANT_K_Q5_0/2 + j]*id;
+        const float x0 = float(data_s[src_idx + 0              + j])*id;
+        const float x1 = float(data_s[src_idx + QUANT_K_Q5_0/2 + j])*id;
 
         const uint xi0 = min(31, int(x0 + 16.5));
         const uint xi1 = min(31, int(x1 + 16.5));
@@ -129,11 +129,11 @@ void quantize(uint dst_idx, uint src_idx)
 #if defined(DATA_A_Q5_1)
 void quantize(uint dst_idx, uint src_idx)
 {
-    float min = data_s[src_idx + 0];
+    float min = float(data_s[src_idx + 0]);
     float max = min;
 
     [[unroll]] for (int j = 1; j < QUANT_K_Q5_1; ++j) {
-        const float v = data_s[src_idx + j];
+        const float v = float(data_s[src_idx + j]);
         min = v < min ? v : min;
         max = v > max ? v : max;
     }
@@ -146,8 +146,8 @@ void quantize(uint dst_idx, uint src_idx)
 
     uint32_t qh = 0;
     [[unroll]] for (int j = 0; j < QUANT_K_Q5_1/2; ++j) {
-        const float x0 = (data_s[src_idx + 0              + j] - min)*id;
-        const float x1 = (data_s[src_idx + QUANT_K_Q5_1/2 + j] - min)*id;
+        const float x0 = (float(data_s[src_idx + 0              + j]) - min)*id;
+        const float x1 = (float(data_s[src_idx + QUANT_K_Q5_1/2 + j]) - min)*id;
 
         const uint xi0 = uint(x0 + 0.5);
         const uint xi1 = uint(x1 + 0.5);
@@ -166,7 +166,7 @@ void quantize(uint dst_idx, uint src_idx)
     float amax = 0.0; // absolute max
 
     [[unroll]] for (int j = 0; j < QUANT_K_Q8_0; j++) {
-        const float v = data_s[src_idx + j];
+        const float v = float(data_s[src_idx + j]);
         amax = max(amax, abs(v));
     }
 
@@ -176,7 +176,7 @@ void quantize(uint dst_idx, uint src_idx)
     data_q[dst_idx].d = float16_t(d);
 
     [[unroll]] for (int j = 0; j < QUANT_K_Q8_0; ++j) {
-        const float x0 = data_s[src_idx + j]*id;
+        const float x0 = float(data_s[src_idx + j])*id;
 
         data_q[dst_idx].qs[j] = int8_t(round(x0));
     }
@@ -189,7 +189,7 @@ void quantize(uint dst_idx, uint src_idx)
     float sum_abs = 0.0;
 
     [[unroll]] for (int j = 0; j < QUANT_K_Q1_0; j++) {
-        sum_abs += abs(data_s[src_idx + j]);
+        sum_abs += abs(float(data_s[src_idx + j]));
     }
 
     const float d = sum_abs / QUANT_K_Q1_0;
@@ -201,7 +201,7 @@ void quantize(uint dst_idx, uint src_idx)
     }
 
     [[unroll]] for (int j = 0; j < QUANT_K_Q1_0; ++j) {
-        if (data_s[src_idx + j] >= 0.0) {
+        if (float(data_s[src_idx + j]) >= 0.0) {
             data_q[dst_idx].qs[j / 8] |= uint8_t(1 << (j % 8));
         }
     }
@@ -226,7 +226,7 @@ void quantize(uint dst_idx, uint src_idx)
     float vmax = 0.0;
 
     [[unroll]] for (int j = 0; j < QUANT_K_IQ4_NL; ++j) {
-        const float v = data_s[src_idx + j];
+        const float v = float(data_s[src_idx + j]);
         if (amax < abs(v)) {
             amax = abs(v);
             vmax = v;
@@ -238,16 +238,16 @@ void quantize(uint dst_idx, uint src_idx)
 
     float sumqx = 0, sumq2 = 0;
     [[unroll]] for (int j = 0; j < QUANT_K_IQ4_NL/2; ++j) {
-        const float x0 = data_s[src_idx + 0                + j]*id;
-        const float x1 = data_s[src_idx + QUANT_K_IQ4_NL/2 + j]*id;
+        const float x0 = float(data_s[src_idx + 0                + j])*id;
+        const float x1 = float(data_s[src_idx + QUANT_K_IQ4_NL/2 + j])*id;
         const uint xi0 = best_index(x0);
         const uint xi1 = best_index(x1);
         data_q[dst_idx].qs[j] = uint8_t(xi0 | (xi1 << 4));
         const float v0 = kvalues_iq4nl[xi0];
         const float v1 = kvalues_iq4nl[xi1];
-        const float w0 = data_s[src_idx + 0                + j]*data_s[src_idx + 0                + j];
-        const float w1 = data_s[src_idx + QUANT_K_IQ4_NL/2 + j]*data_s[src_idx + QUANT_K_IQ4_NL/2 + j];
-        sumqx += w0*v0*data_s[src_idx + j] + w1*v1*data_s[src_idx + QUANT_K_IQ4_NL/2 + j];
+        const float w0 = float(data_s[src_idx + 0                + j])*float(data_s[src_idx + 0                + j]);
+        const float w1 = float(data_s[src_idx + QUANT_K_IQ4_NL/2 + j])*float(data_s[src_idx + QUANT_K_IQ4_NL/2 + j]);
+        sumqx += w0*v0*float(data_s[src_idx + j]) + w1*v1*float(data_s[src_idx + QUANT_K_IQ4_NL/2 + j]);
         sumq2 += w0*v0*v0 + w1*v1*v1;
     }
 
@@ -259,14 +259,14 @@ void quantize(uint dst_idx, uint src_idx)
 #if defined(DATA_A_F32) || defined(DATA_A_F16)
 void quantize(uint dst_idx, uint src_idx)
 {
-    data_q[dst_idx] = A_TYPE(data_s[src_idx]);
+    data_q[dst_idx] = A_TYPE(float(data_s[src_idx]));
 }
 #endif
 
 #if defined(DATA_A_BF16)
 void quantize(uint dst_idx, uint src_idx)
 {
-    data_q[dst_idx] = A_TYPE(fp32_to_bf16(data_s[src_idx]));
+    data_q[dst_idx] = A_TYPE(fp32_to_bf16(float(data_s[src_idx])));
 }
 #endif
 

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

@@ -824,13 +824,15 @@ void process_shaders() {
     string_to_spv("cpy_transpose_32", "copy_transpose.comp", {{"A_TYPE", "uint"}, {"D_TYPE", "uint"}});
 
     for (std::string t : {"q1_0", "q4_0", "q4_1", "q5_0", "q5_1", "q8_0", "iq4_nl"}) {
-        string_to_spv("cpy_f32_" + t, "copy_to_quant.comp", {{"DATA_A_" + to_uppercase(t), "1"}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}});
+        string_to_spv("cpy_f32_" + t, "copy_to_quant.comp", {{"DATA_A_" + to_uppercase(t), "1"}, {"S_TYPE", "float"}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}});
         string_to_spv("cpy_" + t + "_f32", "copy_from_quant.comp", {{"DATA_A_" + to_uppercase(t), "1"}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}});
     }
 
-    for (std::string t : {"f32", "f16", "bf16", "q1_0", "q4_0", "q4_1", "q5_0", "q5_1", "q8_0", "iq4_nl"}) {
-        string_to_spv("set_rows_" + t + "_i32", "copy_to_quant.comp", {{"SET_ROWS", "1"}, {"DATA_A_" + to_uppercase(t), "1"}, {"B_TYPE", "uint"}, {"B_SIZE", "32"}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}});
-        string_to_spv("set_rows_" + t + "_i64", "copy_to_quant.comp", {{"SET_ROWS", "1"}, {"DATA_A_" + to_uppercase(t), "1"}, {"B_TYPE", "uvec2"}, {"B_SIZE", "64"}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}});
+    for (auto src : {std::pair{"f32", "float"}, std::pair{"f16", "float16_t"}}) {
+        for (std::string dst : {"f32", "f16", "bf16", "q1_0", "q4_0", "q4_1", "q5_0", "q5_1", "q8_0", "iq4_nl"}) {
+            string_to_spv("set_rows_" + std::string(src.first) + "_" + dst + "_i32", "copy_to_quant.comp", {{"SET_ROWS", "1"}, {"DATA_A_" + to_uppercase(dst), "1"}, {"B_TYPE", "uint"}, {"B_SIZE", "32"}, {"S_TYPE", src.second}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}});
+            string_to_spv("set_rows_" + std::string(src.first) + "_" + dst + "_i64", "copy_to_quant.comp", {{"SET_ROWS", "1"}, {"DATA_A_" + to_uppercase(dst), "1"}, {"B_TYPE", "uvec2"}, {"B_SIZE", "64"}, {"S_TYPE", src.second}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}});
+        }
     }
 
     auto get_type_str = [](bool f16) {
@@ -1034,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"}});

ggml/src/ggml.c

@@ -1464,14 +1464,14 @@ bool ggml_is_transposed(const struct ggml_tensor * tensor) {
     return tensor->nb[0] > tensor->nb[1];
 }
 
-static bool ggml_is_contiguous_n(const struct ggml_tensor * tensor, int n) {
+static bool ggml_is_contiguous_m_n(const struct ggml_tensor * tensor, int m, int n) {
     size_t next_nb = ggml_type_size(tensor->type);
     if (tensor->ne[0] != ggml_blck_size(tensor->type) && tensor->nb[0] != next_nb) {
         return false;
     }
     next_nb *= tensor->ne[0]/ggml_blck_size(tensor->type);
-    for (int i = 1; i < GGML_MAX_DIMS; i++) {
-        if (i > n) {
+    for (int i = 1; i < n; i++) {
+        if (i > m) {
             if (tensor->ne[i] != 1 && tensor->nb[i] != next_nb) {
                 return false;
             }
@@ -1489,15 +1489,27 @@ bool ggml_is_contiguous(const struct ggml_tensor * tensor) {
 }
 
 bool ggml_is_contiguous_0(const struct ggml_tensor * tensor) {
-    return ggml_is_contiguous_n(tensor, 0);
+    return ggml_is_contiguous_m_n(tensor, 0, GGML_MAX_DIMS);
 }
 
 bool ggml_is_contiguous_1(const struct ggml_tensor * tensor) {
-    return ggml_is_contiguous_n(tensor, 1);
+    return ggml_is_contiguous_m_n(tensor, 1, GGML_MAX_DIMS);
 }
 
 bool ggml_is_contiguous_2(const struct ggml_tensor * tensor) {
-    return ggml_is_contiguous_n(tensor, 2);
+    return ggml_is_contiguous_m_n(tensor, 2, GGML_MAX_DIMS);
+}
+
+bool ggml_is_contiguous_to_1(const struct ggml_tensor * tensor) {
+    return ggml_is_contiguous_m_n(tensor, 0, 1);
+}
+
+bool ggml_is_contiguous_to_2(const struct ggml_tensor * tensor) {
+    return ggml_is_contiguous_m_n(tensor, 0, 2);
+}
+
+bool ggml_is_contiguous_to_3(const struct ggml_tensor * tensor) {
+    return ggml_is_contiguous_m_n(tensor, 0, 3);
 }
 
 bool ggml_is_contiguously_allocated(const struct ggml_tensor * tensor) {
@@ -4507,7 +4519,7 @@ struct ggml_tensor * ggml_conv_1d(
         int                   s0,
         int                   p0,
         int                   d0) {
-    struct ggml_tensor * im2col = ggml_im2col(ctx, a, b, s0, 0, p0, 0, d0, 0, false, GGML_TYPE_F16); // [N, OL, IC * K]
+    struct ggml_tensor * im2col = ggml_im2col(ctx, a, b, s0, 0, p0, 0, d0, 0, false, a->type == GGML_TYPE_BF16 ? GGML_TYPE_F32 : GGML_TYPE_F16); // [N, OL, IC * K]
 
     struct ggml_tensor * result =
         ggml_mul_mat(ctx,
@@ -4541,7 +4553,7 @@ struct ggml_tensor * ggml_conv_1d_dw(
         int                   d0) {
     struct ggml_tensor * new_b = ggml_reshape_4d(ctx, b, b->ne[0], 1, b->ne[1], b->ne[2]);
 
-    struct ggml_tensor * im2col = ggml_im2col(ctx, a, new_b, s0, 0, p0, 0, d0, 0, false, GGML_TYPE_F16);
+    struct ggml_tensor * im2col = ggml_im2col(ctx, a, new_b, s0, 0, p0, 0, d0, 0, false, a->type == GGML_TYPE_BF16 ? GGML_TYPE_F32 : GGML_TYPE_F16);
 
     struct ggml_tensor * result = ggml_mul_mat(ctx, im2col, a);
 
@@ -4647,7 +4659,7 @@ struct ggml_tensor * ggml_conv_2d(
         int                   p1,
         int                   d0,
         int                   d1) {
-    struct ggml_tensor * im2col = ggml_im2col(ctx, a, b, s0, s1, p0, p1, d0, d1, true, a->type); // [N, OH, OW, IC * KH * KW]
+    struct ggml_tensor * im2col = ggml_im2col(ctx, a, b, s0, s1, p0, p1, d0, d1, true, a->type == GGML_TYPE_BF16 ? GGML_TYPE_F32 : GGML_TYPE_F16); // [N, OH, OW, IC * KH * KW]
 
     struct ggml_tensor * result =
         ggml_mul_mat(ctx,
@@ -4729,7 +4741,7 @@ struct ggml_tensor * ggml_conv_3d(
         int                   d1, // dilation height
         int                   d2  // dilation depth
         ) {
-    struct ggml_tensor * im2col = ggml_im2col_3d(ctx, a, b, IC, s0, s1, s2, p0, p1, p2, d0, d1, d2, a->type); // [N*OD, OH, OW, IC * KD * KH * KW]
+    struct ggml_tensor * im2col = ggml_im2col_3d(ctx, a, b, IC, s0, s1, s2, p0, p1, p2, d0, d1, d2, a->type == GGML_TYPE_BF16 ? GGML_TYPE_F32 : GGML_TYPE_F16); // [N*OD, OH, OW, IC * KD * KH * KW]
 
     int64_t OC = a->ne[3] / IC;
     int64_t N = b->ne[3] / IC;
@@ -4779,7 +4791,7 @@ struct ggml_tensor * ggml_conv_2d_dw(
     struct ggml_tensor * new_a = ggml_reshape_4d(ctx, a, a->ne[0], a->ne[1], 1, a->ne[2] * a->ne[3]);
     struct ggml_tensor * im2col = ggml_im2col(ctx, new_a,
                                         ggml_reshape_4d(ctx, b, b->ne[0], b->ne[1], 1, b->ne[2] * b->ne[3]),
-                                        s0, s1, p0, p1, d0, d1, true, GGML_TYPE_F16); // [N * IC, OH, OW, KH * KW]
+                                        s0, s1, p0, p1, d0, d1, true, a->type == GGML_TYPE_BF16 ? GGML_TYPE_F32 : GGML_TYPE_F16); // [N * IC, OH, OW, KH * KW]
     struct ggml_tensor * new_b = ggml_reshape_4d(ctx, im2col, im2col->ne[0], im2col->ne[2] * im2col->ne[1], b->ne[2], b->ne[3]); // [N * IC, OH, OW, KH * KW] => [N, IC, OH * OW, KH * KW]
 
     new_a = ggml_reshape_4d(ctx, new_a, (new_a->ne[0] * new_a->ne[1]), new_a->ne[2],  new_a->ne[3], 1);                       // [OC，1, KH, KW] => [1, OC, 1, KH * KW]

src/llama-kv-cache-dsv4.cpp

@@ -720,7 +720,7 @@ llama_dsv4_comp_state::llama_dsv4_comp_state(
         auto it = ctx_map.find(buft);
         if (it == ctx_map.end()) {
             ggml_init_params params = {
-                /*.mem_size   =*/ size_t(2u*hparams.n_layer()*ggml_tensor_overhead()),
+                /*.mem_size   =*/ size_t(2u*(1 + n_stream)*hparams.n_layer()*ggml_tensor_overhead()),
                 /*.mem_buffer =*/ NULL,
                 /*.no_alloc   =*/ true,
             };
@@ -767,9 +767,17 @@ llama_dsv4_comp_state::llama_dsv4_comp_state(
         ggml_format_name(kv,    "dsv4_%s_state_kv_l%d",    name, il);
         ggml_format_name(score, "dsv4_%s_state_score_l%d", name, il);
 
+        std::vector<ggml_tensor *> kv_stream;
+        std::vector<ggml_tensor *> score_stream;
+
+        for (uint32_t s = 0; s < n_stream; ++s) {
+            kv_stream.push_back(ggml_view_2d(ctx, kv, n_embd_state, state_size, kv->nb[1], s*kv->nb[2]));
+            score_stream.push_back(ggml_view_2d(ctx, score, n_embd_state, state_size, score->nb[1], s*score->nb[2]));
+        }
+
         map_layer_ids[il] = layers.size();
 
-        layers.push_back({ il, kv, score });
+        layers.push_back({ il, kv, score, std::move(kv_stream), std::move(score_stream) });
     }
 
     for (auto & [buft, ctx] : ctx_map) {
@@ -809,6 +817,30 @@ void llama_dsv4_comp_state::clear(llama_seq_id seq_id, bool data) {
     }
 }
 
+void llama_dsv4_comp_state::seq_cp(llama_seq_id seq_id_src, llama_seq_id seq_id_dst) {
+    GGML_ASSERT(seq_id_src >= 0 && (uint32_t) seq_id_src < n_stream);
+    GGML_ASSERT(seq_id_dst >= 0 && (uint32_t) seq_id_dst < n_stream);
+
+    if (seq_id_src == seq_id_dst) {
+        return;
+    }
+
+    sc_info.ssrc.push_back((uint32_t) seq_id_src);
+    sc_info.sdst.push_back((uint32_t) seq_id_dst);
+}
+
+void llama_dsv4_comp_state::apply_copies(const stream_copy_info & sc_info) const {
+    for (size_t i = 0; i < sc_info.ssrc.size(); ++i) {
+        const uint32_t ssrc = sc_info.ssrc[i];
+        const uint32_t sdst = sc_info.sdst[i];
+
+        for (const auto & layer : layers) {
+            ggml_backend_tensor_copy(layer.kv_stream[ssrc], layer.kv_stream[sdst]);
+            ggml_backend_tensor_copy(layer.score_stream[ssrc], layer.score_stream[sdst]);
+        }
+    }
+}
+
 uint32_t llama_dsv4_comp_state::get_ratio() const {
     return ratio;
 }
@@ -1154,7 +1186,13 @@ llama_memory_context_ptr llama_kv_cache_dsv4::init_full() {
 }
 
 llama_memory_context_ptr llama_kv_cache_dsv4::init_update(llama_context * lctx, bool optimize) {
-    return std::make_unique<llama_kv_cache_dsv4_context>(this, lctx, optimize);
+    return std::make_unique<llama_kv_cache_dsv4_context>(
+            this,
+            lctx,
+            optimize,
+            std::move(csa_state->sc_info),
+            std::move(hca_state->sc_info),
+            std::move(lid_state->sc_info));
 }
 
 bool llama_kv_cache_dsv4::get_can_shift() const {
@@ -1174,14 +1212,19 @@ bool llama_kv_cache_dsv4::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1
     }
 
     if (p0 > 0) {
-        // DSV4 compressed cache rows are derived from running compressor state,
-        // so arbitrary rollback is not reconstructible from the raw cache alone.
-        // Allow the common prompt-cache cleanup no-op: remove [end, infinity).
-        if (seq_id >= 0 && p0 > kv_raw->seq_pos_max(seq_id)) {
-            return true;
+        if (seq_id < 0 || (uint32_t) seq_id >= n_seq_max ||
+                p0 <= kv_raw->seq_pos_max(seq_id)) {
+            return false;
         }
 
-        return false;
+        bool res = true;
+
+        res = res & kv_raw->seq_rm(seq_id, p0, -1);
+        res = res & kv_csa->seq_rm(seq_id, p0/DSV4_CSA_RATIO, -1);
+        res = res & kv_hca->seq_rm(seq_id, p0/DSV4_HCA_RATIO, -1);
+        res = res & kv_lid->seq_rm(seq_id, p0/DSV4_CSA_RATIO, -1);
+
+        return res;
     }
 
     const bool res = kv_raw->seq_rm(seq_id, p0, p1);
@@ -1194,7 +1237,16 @@ bool llama_kv_cache_dsv4::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1
 }
 
 void llama_kv_cache_dsv4::seq_cp(llama_seq_id seq_id_src, llama_seq_id seq_id_dst, llama_pos p0, llama_pos p1) {
+    GGML_ASSERT(p0 <= 0 && p1 < 0 && "DSV4 only supports full sequence copies");
+
     kv_raw->seq_cp(seq_id_src, seq_id_dst, p0, p1);
+    kv_csa->seq_cp(seq_id_src, seq_id_dst, -1, -1);
+    kv_hca->seq_cp(seq_id_src, seq_id_dst, -1, -1);
+    kv_lid->seq_cp(seq_id_src, seq_id_dst, -1, -1);
+
+    csa_state->seq_cp(seq_id_src, seq_id_dst);
+    hca_state->seq_cp(seq_id_src, seq_id_dst);
+    lid_state->seq_cp(seq_id_src, seq_id_dst);
 }
 
 void llama_kv_cache_dsv4::seq_keep(llama_seq_id seq_id) {
@@ -1639,20 +1691,26 @@ llama_kv_cache_dsv4_context::llama_kv_cache_dsv4_context(
 llama_kv_cache_dsv4_context::llama_kv_cache_dsv4_context(
         llama_kv_cache_dsv4 * kv,
         llama_context * lctx,
-        bool optimize) :
+        bool optimize,
+        stream_copy_info sc_info_csa,
+        stream_copy_info sc_info_hca,
+        stream_copy_info sc_info_lid) :
     ctx_raw(std::make_unique<llama_kv_cache_dsv4_raw_context>(kv->get_raw(), lctx, optimize)),
     ctx_csa_mem(kv->get_csa()->init_update(lctx, optimize)),
     ctx_hca_mem(kv->get_hca()->init_update(lctx, optimize)),
     ctx_lid_mem(kv->get_lid()->init_update(lctx, optimize)),
-    ctx_csa(std::make_unique<llama_kv_cache_dsv4_comp_context>(kv->get_csa())),
-    ctx_hca(std::make_unique<llama_kv_cache_dsv4_comp_context>(kv->get_hca())),
-    ctx_lid(std::make_unique<llama_kv_cache_dsv4_comp_context>(kv->get_lid())),
     csa_state(kv->get_csa_state()),
     hca_state(kv->get_hca_state()),
     lid_state(kv->get_lid_state()),
+    sc_info_csa(std::move(sc_info_csa)),
+    sc_info_hca(std::move(sc_info_hca)),
+    sc_info_lid(std::move(sc_info_lid)),
     status(llama_memory_status_combine(
-                llama_memory_status_combine(ctx_raw->get_status(), ctx_csa_mem->get_status()),
-                llama_memory_status_combine(ctx_hca_mem->get_status(), ctx_lid_mem->get_status()))) {
+                llama_memory_status_combine(
+                    llama_memory_status_combine(ctx_raw->get_status(), ctx_csa_mem->get_status()),
+                    llama_memory_status_combine(ctx_hca_mem->get_status(), ctx_lid_mem->get_status())),
+                this->sc_info_csa.empty() && this->sc_info_hca.empty() && this->sc_info_lid.empty() ?
+                    LLAMA_MEMORY_STATUS_NO_UPDATE : LLAMA_MEMORY_STATUS_SUCCESS)) {
 }
 
 llama_kv_cache_dsv4_context::llama_kv_cache_dsv4_context(
@@ -1720,6 +1778,18 @@ bool llama_kv_cache_dsv4_context::apply() {
 
     res = res & ctx_raw->apply();
 
+    if (ctx_csa_mem) {
+        res = res & ctx_csa_mem->apply();
+        res = res & ctx_hca_mem->apply();
+        res = res & ctx_lid_mem->apply();
+    }
+
+    if (ubatches.empty()) {
+        csa_state->apply_copies(sc_info_csa);
+        hca_state->apply_copies(sc_info_hca);
+        lid_state->apply_copies(sc_info_lid);
+    }
+
     return res;
 }
 

src/llama-kv-cache-dsv4.h

@@ -10,6 +10,10 @@
 
 class llama_dsv4_comp_state {
 public:
+    using stream_copy_info = llama_kv_cache::stream_copy_info;
+
+    stream_copy_info sc_info;
+
     llama_dsv4_comp_state(
             const llama_model & model,
             bool            offload,
@@ -22,6 +26,8 @@ public:
         const llama_memory_i::layer_filter_cb & filter);
 
     void clear(llama_seq_id seq_id, bool data);
+    void seq_cp(llama_seq_id seq_id_src, llama_seq_id seq_id_dst);
+    void apply_copies(const stream_copy_info & sc_info) const;
 
     uint32_t get_ratio()    const;
     uint32_t get_state_size() const;
@@ -44,6 +50,9 @@ private:
 
         ggml_tensor * kv;
         ggml_tensor * score;
+
+        std::vector<ggml_tensor *> kv_stream;
+        std::vector<ggml_tensor *> score_stream;
     };
 
     const uint32_t ratio;
@@ -245,6 +254,7 @@ private:
 class llama_kv_cache_dsv4_context : public llama_memory_context_i {
 public:
     using slot_info_vec_t = llama_kv_cache::slot_info_vec_t;
+    using stream_copy_info = llama_kv_cache::stream_copy_info;
 
     struct comp_plan {
         // Per-ubatch recipe for updating compressor state, committing completed
@@ -291,7 +301,10 @@ public:
     llama_kv_cache_dsv4_context(
             llama_kv_cache_dsv4 * kv,
             llama_context * lctx,
-            bool optimize);
+            bool optimize,
+            stream_copy_info sc_info_csa,
+            stream_copy_info sc_info_hca,
+            stream_copy_info sc_info_lid);
 
     llama_kv_cache_dsv4_context(
             llama_kv_cache_dsv4 * kv,
@@ -351,9 +364,13 @@ private:
     const std::unique_ptr<llama_kv_cache_dsv4_comp_context> ctx_hca;
     const std::unique_ptr<llama_kv_cache_dsv4_comp_context> ctx_lid;
 
-    const llama_dsv4_comp_state * csa_state = nullptr;
-    const llama_dsv4_comp_state * hca_state = nullptr;
-    const llama_dsv4_comp_state * lid_state = nullptr;
+    llama_dsv4_comp_state * csa_state = nullptr;
+    llama_dsv4_comp_state * hca_state = nullptr;
+    llama_dsv4_comp_state * lid_state = nullptr;
+
+    stream_copy_info sc_info_csa;
+    stream_copy_info sc_info_hca;
+    stream_copy_info sc_info_lid;
 
     bool reserve_plans = false;
     mutable comp_plan reserve_plan_csa;

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

@@ -2423,11 +2423,10 @@ struct test_set_rows : public test_case {
 
     void initialize_tensors(ggml_context * ctx) override {
         for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) {
+            if (ggml_is_view_op(t->op)) {
+                continue;
+            }
             if (t->type == GGML_TYPE_I64 || t->type == GGML_TYPE_I32) {
-                if (ggml_is_view_op(t->op)) {
-                    continue;
-                }
-
                 init_set_rows_row_ids(t, ne[1]);
             } else {
                 init_tensor_uniform(t);
@@ -2450,6 +2449,9 @@ struct test_set_rows : public test_case {
                 err_estimate /= 8.0f;
             }
             err_estimate *= err_estimate;
+            if (type_src == GGML_TYPE_F16) {
+                err_estimate *= 16.0f;
+            }
             err_estimate /= 0.25f*float(ne[0] * r * ne[2]*nr23[0] * ne[3]*nr23[1]);
             return err_estimate;
         }
@@ -5553,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);
@@ -5574,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 {
@@ -5586,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 {
@@ -7928,17 +7944,19 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
     test_cases.emplace_back(new test_set_rows(GGML_TYPE_F32, GGML_TYPE_F32, GGML_TYPE_I64, { 1, 8, 1, 3 }, { 1, 1 }, 2, false));
     test_cases.emplace_back(new test_set_rows(GGML_TYPE_F32, GGML_TYPE_F32, GGML_TYPE_I32, { 1, 8, 1, 3 }, { 1, 1 }, 2, false));
     test_cases.emplace_back(new test_set_rows(GGML_TYPE_F32, GGML_TYPE_Q8_0, GGML_TYPE_I32, { 256, 5, 1, 3 }, { 1, 1, }, 1, false));
-    for (ggml_type type : all_types) {
-        for (int b : {1, 7}) {
-            for (bool v : {false, true}) {
-                test_cases.emplace_back(new test_set_rows(GGML_TYPE_F32, type, GGML_TYPE_I64, { 256, 5,  b, 3 }, { 1, 1, }, 1, v));
-                test_cases.emplace_back(new test_set_rows(GGML_TYPE_F32, type, GGML_TYPE_I64, { 256, 11, 1, b }, { 2, 3, }, 7, v));
+    for (ggml_type src_type : {GGML_TYPE_F16, GGML_TYPE_F32}) {
+        for (ggml_type type : all_types) {
+            for (int b : {1, 7}) {
+                for (bool v : {false, true}) {
+                    test_cases.emplace_back(new test_set_rows(src_type, type, GGML_TYPE_I64, { 256, 5,  b, 3 }, { 1, 1, }, 1, v));
+                    test_cases.emplace_back(new test_set_rows(src_type, type, GGML_TYPE_I64, { 256, 11, 1, b }, { 2, 3, }, 7, v));
 
-                test_cases.emplace_back(new test_set_rows(GGML_TYPE_F32, type, GGML_TYPE_I64, { 3*ggml_blck_size(type), 3, b, 1 }, { 2, 3, }, 2, v));
+                    test_cases.emplace_back(new test_set_rows(src_type, type, GGML_TYPE_I64, { 3*ggml_blck_size(type), 3, b, 1 }, { 2, 3, }, 2, v));
 
-                if (ggml_blck_size(type) == 1) {
-                    test_cases.emplace_back(new test_set_rows(GGML_TYPE_F32, type, GGML_TYPE_I64, { 31, 3, b, 1 }, { 2, 3, }, 2, v));
-                    test_cases.emplace_back(new test_set_rows(GGML_TYPE_F32, type, GGML_TYPE_I64, { 33, 5, 1, b }, { 2, 3, }, 1, v));
+                    if (ggml_blck_size(type) == 1) {
+                        test_cases.emplace_back(new test_set_rows(src_type, type, GGML_TYPE_I64, { 31, 3, b, 1 }, { 2, 3, }, 2, v));
+                        test_cases.emplace_back(new test_set_rows(src_type, type, GGML_TYPE_I64, { 33, 5, 1, b }, { 2, 3, }, 1, v));
+                    }
                 }
             }
         }
@@ -8013,6 +8031,7 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
 
     // im2col 2D
     test_cases.emplace_back(new test_im2col(GGML_TYPE_F32, GGML_TYPE_F32, GGML_TYPE_F32));
+    test_cases.emplace_back(new test_im2col(GGML_TYPE_F32, GGML_TYPE_F32, GGML_TYPE_F16));
     test_cases.emplace_back(new test_im2col(GGML_TYPE_F32, GGML_TYPE_F16, GGML_TYPE_F32));
     test_cases.emplace_back(new test_im2col(GGML_TYPE_F32, GGML_TYPE_F16, GGML_TYPE_F16));
     for (int s0 : {1, 3}) {
@@ -9084,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));
+            }
         }
     }
 

tests/test-export-graph-ops.cpp

@@ -152,6 +152,10 @@ int main(int argc, char ** argv) {
         init_result = common_init_from_params(params);
 
         ctx = init_result->context();
+        if (!ctx) {
+            LOG_ERR("failed to initialize params\n");
+            return 1;
+        }
     } else {
 #ifdef LLAMA_HF_FETCH
         auto [hf_repo, hf_quant] = common_download_split_repo_tag(params.model.hf_repo);

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-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/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/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
@@ -114,6 +115,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 +172,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:
@@ -223,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:
@@ -359,7 +365,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/tokenize/tokenize.cpp

@@ -1,5 +1,6 @@
+#include "arg.h"
 #include "common.h"
-//#include "log.h" // TODO: start using log.h
+#include "log.h"
 #include "llama.h"
 
 #include <clocale>
@@ -8,115 +9,22 @@
 #include <fstream>
 #include <string>
 #include <vector>
-#include <iostream> // TODO: remove me
+#include <iostream>
+#include <sstream>
 
 #if defined(_WIN32)
 #define WIN32_LEAN_AND_MEAN
 #include <windows.h>
-#include <shellapi.h>   // For CommandLineToArgvW
 #endif
 
-static void print_usage_information(const char * argv0) {
-    printf("usage: %s [options]\n\n", argv0);
-    printf("The tokenize program tokenizes a prompt using a given model,\n");
-    printf("and prints the resulting tokens to standard output.\n\n");
-    printf("It needs a model file, a prompt, and optionally other flags\n");
-    printf("to control the behavior of the tokenizer.\n\n");
-    printf("    The possible options are:\n");
-    printf("\n");
-    printf("    -h, --help                           print this help and exit\n");
-    printf("    -m MODEL_PATH, --model MODEL_PATH    path to model.\n");
-    printf("    --ids                                if given, only print numerical token IDs, and not token strings.\n");
-    printf("                                         The output format looks like [1, 2, 3], i.e. parseable by Python.\n");
-    printf("    -f PROMPT_FNAME, --file PROMPT_FNAME read prompt from a file.\n");
-    printf("    -p PROMPT, --prompt PROMPT           read prompt from the argument.\n");
-    printf("    --stdin                              read prompt from standard input.\n");
-    printf("    --no-bos                             do not ever add a BOS token to the prompt, even if normally the model uses a BOS token.\n");
-    printf("    --no-escape                          do not escape input (such as \\n, \\t, etc.).\n");
-    printf("    --no-parse-special                   do not parse control tokens.\n");
-    printf("    --log-disable                        disable logs. Makes stderr quiet when loading the model.\n");
-    printf("    --show-count                         print the total number of tokens.\n");
-}
+static void print_usage(int argc, char ** argv) {
+    (void) argc;
 
-static void llama_log_callback_null(ggml_log_level level, const char * text, void * user_data) {
-    (void) level;
-    (void) text;
-    (void) user_data;
-}
-
-static std::string read_prompt_from_file(const char * filepath, bool & success) {
-    success = false;
-
-    std::ifstream in(filepath, std::ios::binary);
-    if (!in) {
-        fprintf(stderr, "%s: could not open file '%s' for reading: %s\n", __func__, filepath, strerror(errno));
-        return std::string();
-    }
-    // do not assume the file is seekable (e.g. /dev/stdin)
-    std::stringstream buffer;
-    buffer << in.rdbuf();
-    if (in.fail()) {
-        fprintf(stderr, "%s: could not read the entire file '%s': %s\n", __func__, filepath, strerror(errno));
-        return std::string();
-    }
-
-    success = true;
-    return buffer.str();
-}
-
-//
-// Function: ingest_args(...) -> vector<string>
-//
-//  Takes argc and argv arguments, and converts them to a vector of UTF-8 encoded
-//  strings, as an STL vector<string>.
-//
-//  In particular, it handles character encoding shenanigans on Windows.
-//
-// Note: raw_argc and raw_argv are not actually read at all on Windows.
-//       On Windows we call GetCommandLineW to get the arguments in wchar_t
-//       format, ignoring the regular argc/argv arguments to main().
-//
-// TODO: potential opportunity to roll common stuff into common/console.cpp
-//       in relation to Windows wchar_t shenanigans.
-static std::vector<std::string> ingest_args(int raw_argc, char ** raw_argv) {
-    std::vector<std::string> argv;
-
-    // Handle Windows, if given non-ASCII arguments.
-    // We convert wchar_t arguments into UTF-8 char* on this platform.
-    // Lets you invoke 'tokenize' on Windows cmd.exe with non-ASCII characters
-    // without throwing tantrums.
-#if defined(_WIN32)
-    int argc;
-    const LPWSTR cmdline_wargv = GetCommandLineW();
-    LPWSTR * wargv = CommandLineToArgvW(cmdline_wargv, &argc);
-
-    // silence unused arg warnings
-    (void) raw_argc;
-    (void) raw_argv;
-
-    for (int i = 0; i < argc; ++i) {
-        int length_needed = WideCharToMultiByte(CP_UTF8, 0, wargv[i], wcslen(wargv[i]), 0, 0, NULL, NULL);
-        char * output_buf = (char *) calloc(length_needed+1, sizeof(char));
-        GGML_ASSERT(output_buf);
-
-        WideCharToMultiByte(CP_UTF8, 0, wargv[i], wcslen(wargv[i]), output_buf, length_needed, NULL, NULL);
-        output_buf[length_needed] = '\0';
-
-        argv.push_back(output_buf);
-        free(output_buf);
-    }
-
-    LocalFree((HLOCAL) wargv);
-#else
-    int argc = raw_argc;
-    for (int i = 0; i < argc; ++i) {
-        argv.push_back(raw_argv[i]);
-    }
-#endif
-
-    GGML_ASSERT((unsigned int) argc == argv.size());
-
-    return argv;
+    LOG("\nexample usage:\n");
+    LOG("\n  %s -m your_model.gguf -p \"Hello world\"\n", argv[0]);
+    LOG("\n  %s -m your_model.gguf -f prompt.txt --ids\n", argv[0]);
+    LOG("\n  cat prompt.txt | %s -m your_model.gguf --stdin --show-count\n", argv[0]);
+    LOG("\n");
 }
 
 //
@@ -184,166 +92,69 @@ static void write_utf8_cstr_to_stdout(const char * str, bool & invalid_utf8) {
 #endif
 }
 
-int main(int raw_argc, char ** raw_argv) {
+int main(int argc, char ** argv) {
     std::setlocale(LC_NUMERIC, "C");
 
-    const std::vector<std::string> argv = ingest_args(raw_argc, raw_argv);
-    const int argc = argv.size();
+    common_params params;
 
-    if (argc <= 1) {
-        print_usage_information(argv[0].c_str());
+    common_init();
+
+    if (!common_params_parse(argc, argv, params, LLAMA_EXAMPLE_TOKENIZE, print_usage)) {
         return 1;
     }
 
-    //////
-    // Read out all the command line arguments.
-    //////
+    // 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.
+    const bool use_stdin = params.tokenize_stdin;
+    const bool use_file  = !params.prompt_file.empty();
 
-    // variables where to put any arguments we see.
-    bool printing_ids = false;
-    bool no_bos = false;
-    bool no_escape = false;
-    bool no_parse_special = false;
-    bool disable_logging = false;
-    bool show_token_count = false;
-    const char * model_path = NULL;
-    const char * prompt_path = NULL;
-    const char * prompt_arg = NULL;
-
-    // track which arguments were explicitly given
-    // used for sanity checking down the line
-    bool model_path_set = false;
-    bool prompt_path_set = false;
-    bool prompt_set = false;
-    bool stdin_set = false;
-
-    int iarg = 1;
-    for (; iarg < argc; ++iarg) {
-        std::string arg{argv[iarg]};
-        if (arg == "-h" || arg == "--help") {
-            print_usage_information(argv[0].c_str());
-            return 0;
-        }
-        else if (arg == "--ids") {
-            printing_ids = true;
-        }
-        else if (arg == "-m" || arg == "--model") {
-            if (model_path_set) {
-                fprintf(stderr, "Error: -m or --model specified multiple times.\n");
-                return 1;
-            }
-            model_path = argv[++iarg].c_str();
-            model_path_set = true;
-        }
-        else if (arg == "--no-bos") {
-            no_bos = true;
-        }
-        else if (arg == "--no-escape") {
-            no_escape = true;
-        }
-        else if (arg == "--no-parse-special") {
-            no_parse_special = true;
-        }
-        else if (arg == "-p" || arg == "--prompt") {
-            if (prompt_set) {
-                fprintf(stderr, "Error: -p or --prompt specified multiple times.\n");
-                return 1;
-            }
-            prompt_arg = argv[++iarg].c_str();
-            prompt_set = true;
-        }
-        else if (arg == "-f" || arg == "--file") {
-            if (prompt_path_set) {
-                fprintf(stderr, "Error: -f or --file specified multiple times.\n");
-                return 1;
-            }
-            prompt_path = argv[++iarg].c_str();
-            prompt_path_set = true;
-        }
-        else if (arg == "--stdin") {
-            stdin_set = true;
-        }
-        else if (arg == "--log-disable") {
-            disable_logging = true;
-        }
-        else if (arg == "--show-count") {
-            show_token_count = true;
-        }
-        else {
-            fprintf(stderr, "Error: unknown option '%s'\n", argv[iarg].c_str());
-            return 1;
-        }
-    }
-
-    //////
-    // Sanity check the command line arguments.
-    //////
-
-    // Check that we have the required stuff set.
-    if (model_path_set && model_path == NULL) {
-        fprintf(stderr, "Error: --model requires an argument.\n");
-        return 1;
-    }
-    if (!model_path_set) {
-        fprintf(stderr, "Error: must specify --model.\n");
-        return 1;
-    }
-    if (prompt_path_set && prompt_path == NULL) {
-        fprintf(stderr, "Error: --file requires an argument.\n");
-        return 1;
-    }
-    if (prompt_set && prompt_arg == NULL) {
-        fprintf(stderr, "Error: --prompt requires an argument.\n");
-        return 1;
-    }
-    const int prompts_set = !!(prompt_path_set) + !!(prompt_set) + !!(stdin_set);
-    if (prompts_set > 1) {
-        fprintf(stderr, "Error: --stdin, --file and --prompt are mutually exclusive.\n");
-        return 1;
-    }
-    // Must have some prompt.
-    if (prompts_set == 0) {
-        fprintf(stderr, "Error: must specify one of: --stdin, --file or --prompt.\n");
+    // 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;
     }
 
-    GGML_ASSERT(model_path);
-    GGML_ASSERT(prompt_path || prompt_arg || stdin_set);
-
-    //////
-    // Figure out where will the prompt come from.
-    //////
+    // must have some prompt
+    if (!use_stdin && !use_file && params.prompt.empty()) {
+        LOG_ERR("error: must specify one of: --stdin, --file or --prompt\n");
+        return 1;
+    }
 
     std::string prompt;
-    if (prompt_path_set) {
-        bool success = false;
-        prompt = read_prompt_from_file(prompt_path, success);
-        if (!success) {
+    if (use_file) {
+        // read the file verbatim: common's -f handler strips a single trailing
+        // newline, but for a tokenizer the input bytes must be preserved exactly
+        // (a trailing newline is itself a token). escapes are applied locally
+        // to match the behavior of -p/--prompt and --stdin.
+        std::ifstream in(params.prompt_file, std::ios::binary);
+        if (!in) {
+            LOG_ERR("error: could not open file '%s' for reading\n", params.prompt_file.c_str());
             return 1;
         }
-    } else if (prompt_set) {
-        prompt = prompt_arg;
-    } else {
-        GGML_ASSERT(stdin_set);
-        // we read stdin *after* loading model (early exit if model cannot
-        // be loaded, which can be a nicer user experience)
-    }
-
-    //////
-    // Start actually doing the tokenizing stuff.
-    //////
-
-    if (disable_logging) {
-        llama_log_set(llama_log_callback_null, NULL);
+        std::stringstream ss;
+        ss << in.rdbuf();
+        prompt = ss.str();
+        if (params.escape) {
+            string_process_escapes(prompt);
+        }
+    } else if (!use_stdin) {
+        // -p/--prompt is already escape-processed by common_params_parse()
+        // (controlled by --escape/--no-escape), so use it verbatim here.
+        prompt = params.prompt;
     }
+    // else: we read stdin *after* loading the model (early exit if the
+    // model cannot be loaded, which is a nicer user experience)
 
     llama_backend_init();
 
+    // load only the vocabulary (no weights), since tokenizing does not need them
     llama_model_params model_params = llama_model_default_params();
     model_params.vocab_only = true;
-    llama_model * model = llama_model_load_from_file(model_path, model_params);
+    llama_model * model = llama_model_load_from_file(params.model.path.c_str(), model_params);
     if (!model) {
-        fprintf(stderr, "Error: could not load model from file '%s'.\n", model_path);
+        LOG_ERR("error: could not load model from file '%s'.\n", params.model.path.c_str());
         return 1;
     }
 
@@ -352,42 +163,41 @@ int main(int raw_argc, char ** raw_argv) {
     llama_context_params ctx_params = llama_context_default_params();
     llama_context * ctx = llama_init_from_model(model, ctx_params);
     if (!ctx) {
-        fprintf(stderr, "Error: could not create context.\n");
+        LOG_ERR("error: could not create context.\n");
         return 1;
     }
 
     // read entire prompt from stdin?
-    if (stdin_set) {
-        GGML_ASSERT(!prompt_path_set && !prompt_set);
-
+    if (params.tokenize_stdin) {
         std::stringstream stdin_buffer;
         stdin_buffer << std::cin.rdbuf();
         if (std::cin.fail()) {
-            fprintf(stderr, "Error: could not read the entire standard input.\n");
+            LOG_ERR("error: could not read the entire standard input.\n");
             return 1;
         }
 
         prompt = stdin_buffer.str();
+
+        // stdin is not seen by common_params_parse(), so apply escape handling
+        // here to match the behavior of -p/--prompt and -f/--file.
+        if (params.escape) {
+            string_process_escapes(prompt);
+        }
     }
 
     const bool model_wants_add_bos = llama_vocab_get_add_bos(vocab);
-    const bool add_bos = model_wants_add_bos && !no_bos;
-    const bool parse_special = !no_parse_special;
-    const bool escape = !no_escape;
-
-    if (escape) {
-        string_process_escapes(prompt);
-    }
+    const bool add_bos      = model_wants_add_bos && !params.tokenize_no_bos;
+    const bool parse_special = params.parse_special;
 
     std::vector<llama_token> tokens;
     tokens = common_tokenize(vocab, prompt, add_bos, parse_special);
 
-    if (printing_ids) {
+    if (params.tokenize_ids) {
         printf("[");
     }
 
     for (int i = 0; i < (int) tokens.size(); i++) {
-        if (printing_ids) {
+        if (params.tokenize_ids) {
             if (i > 0) {
                 printf(", ");
             }
@@ -404,13 +214,14 @@ int main(int raw_argc, char ** raw_argv) {
         }
     }
 
-    if (printing_ids) {
+    if (params.tokenize_ids) {
         printf("]\n");
     }
 
-    if (show_token_count) {
+    if (params.tokenize_show_count) {
         printf("Total number of tokens: %zu\n", tokens.size());
     }
+
     // silence valgrind
     llama_free(ctx);
     llama_model_free(model);

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

@@ -17,10 +17,10 @@
 	let { onMcpSettingsClick }: Props = $props();
 
 	let mcpSearchQuery = $state('');
-	let allMcpServers = $derived(mcpStore.getServers());
-	let mcpServers = $derived(mcpStore.visibleMcpServers);
+	// Every configured server is listed; `enabled` is an on/off state,
+	// not a visibility filter, so a disabled server stays toggleable.
+	let mcpServers = $derived(mcpStore.getServers());
 	let hasMcpServers = $derived(mcpServers.length > 0);
-	// let hasAnyMcpServers = $derived(allMcpServers.length > 0);
 	let filteredMcpServers = $derived.by(() => {
 		const query = mcpSearchQuery.toLowerCase().trim();
 		if (!query) return mcpServers;
@@ -46,7 +46,7 @@
 	function handleMcpSubMenuOpen(open: boolean) {
 		if (open) {
 			mcpSearchQuery = '';
-			mcpStore.runHealthChecksForServers(allMcpServers);
+			mcpStore.runHealthChecksForServers(mcpServers);
 		}
 	}
 

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

@@ -2,85 +2,31 @@
 	import { Lightbulb, LightbulbOff, Check, Info } from '@lucide/svelte';
 	import * as DropdownMenu from '$lib/components/ui/dropdown-menu';
 	import * as Tooltip from '$lib/components/ui/tooltip';
-	import { ReasoningEffort } from '$lib/enums';
-	import { REASONING_EFFORT_TOKENS } from '$lib/constants/reasoning-effort-tokens';
-	import { REASONING_EFFORT_LEVELS } from '$lib/constants/reasoning-effort';
-	import type { ReasoningEffortLevel } from '$lib/types';
-	import {
-		modelsStore,
-		checkModelSupportsThinking,
-		supportsThinking,
-		propsCacheVersion,
-		loadedModelIds
-	} from '$lib/stores/models.svelte';
-	import { chatStore } from '$lib/stores/chat.svelte';
-	import { conversationsStore, activeMessages } from '$lib/stores/conversations.svelte';
-	import { isRouterMode } from '$lib/stores/server.svelte';
-	import type { DatabaseMessage } from '$lib/types/database';
+	import { useReasoningMenu } from '$lib/hooks/use-reasoning-menu.svelte';
 
 	let subOpen = $state(false);
 
-	let conversationModel = $derived(
-		chatStore.getConversationModel(activeMessages() as DatabaseMessage[])
-	);
-
-	let modelSupportsThinkingFromMessages = $derived.by(() => {
-		const modelId = isRouterMode() ? modelsStore.selectedModelName || conversationModel : null;
-		if (!modelId) return false;
-
-		const messages = conversationsStore.activeMessages;
-
-		return messages.some(
-			(m) => m.role === 'assistant' && m.model === modelId && !!m.reasoningContent
-		);
-	});
-
-	let modelSupportsThinking = $derived.by(() => {
-		loadedModelIds();
-		propsCacheVersion();
-
-		if (isRouterMode()) {
-			const modelId = modelsStore.selectedModelName || conversationModel;
-			return checkModelSupportsThinking(modelId ?? '') || modelSupportsThinkingFromMessages;
-		}
-
-		return supportsThinking() || modelSupportsThinkingFromMessages;
-	});
-
-	let thinkingEnabled = $derived(conversationsStore.getThinkingEnabled());
-	let currentEffort = $derived(conversationsStore.getReasoningEffort());
-	let isOff = $derived(!thinkingEnabled);
-
-	function isSelected(item: ReasoningEffortLevel): boolean {
-		if (item.isOff) return isOff;
-		return thinkingEnabled && currentEffort === item.value;
-	}
-
-	function handleSelection(item: ReasoningEffortLevel) {
-		if (item.isOff) {
-			conversationsStore.setThinkingEnabled(false);
-		} else {
-			conversationsStore.setThinkingEnabled(true);
-			conversationsStore.setReasoningEffort(item.value as ReasoningEffort);
-		}
-		subOpen = false;
-	}
+	const reasoning = useReasoningMenu();
 </script>
 
-{#if modelSupportsThinking}
+{#if reasoning.modelSupportsThinking}
 	<DropdownMenu.Sub bind:open={subOpen}>
 		<DropdownMenu.SubTrigger class="flex cursor-pointer items-center gap-2">
-			{#if thinkingEnabled}
+			{#if reasoning.thinkingEnabled}
 				<Lightbulb class="h-4 w-4 shrink-0 text-amber-400" />
 			{:else}
 				<LightbulbOff class="h-4 w-4 shrink-0 text-muted-foreground" />
 			{/if}
 
-			<span class="text-sm inline-flex gap-2 {!thinkingEnabled ? 'text-muted-foreground' : ''}">
+			<span
+				class="text-sm inline-flex gap-2 {!reasoning.thinkingEnabled
+					? 'text-muted-foreground'
+					: ''}"
+			>
 				Reasoning
 
 				<span class="capitalize text-muted-foreground">
-					{thinkingEnabled ? currentEffort : 'off'}
+					{reasoning.thinkingEnabled ? reasoning.currentEffort : 'off'}
 				</span>
 			</span>
 		</DropdownMenu.SubTrigger>
@@ -88,14 +34,18 @@
 		<DropdownMenu.SubContent
 			class="w-60 bg-popover p-1.5 text-popover-foreground shadow-md outline-none"
 		>
-			{#each REASONING_EFFORT_LEVELS as level (level.value)}
+			{#each reasoning.levels as level (level.value)}
+				{@const tokenLabel = reasoning.tokenLabel(level)}
 				<button
 					type="button"
 					class="flex w-full cursor-pointer items-center gap-3 rounded-md px-2 py-1.75 text-left text-sm transition-colors hover:bg-accent"
-					class:bg-accent={isSelected(level)}
-					onclick={() => handleSelection(level)}
+					class:bg-accent={reasoning.isSelected(level)}
+					onclick={() => {
+						reasoning.select(level);
+						subOpen = false;
+					}}
 				>
-					{#if isSelected(level)}
+					{#if reasoning.isSelected(level)}
 						<Check class="h-4 w-4 shrink-0 text-foreground" />
 					{:else}
 						<div class="h-4 w-4 shrink-0"></div>
@@ -103,11 +53,9 @@
 
 					<span class="flex-1">{level.label}</span>
 
-					{#if !level.isOff}
+					{#if tokenLabel}
 						<span class="text-[11px] text-muted-foreground opacity-60">
-							{REASONING_EFFORT_TOKENS[level.value] === -1
-								? 'Unlimited'
-								: `Max ${REASONING_EFFORT_TOKENS[level.value].toLocaleString()} tokens`}
+							{tokenLabel}
 						</span>
 					{/if}
 

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

@@ -10,10 +10,18 @@
 	import { ATTACHMENT_FILE_ITEMS } from '$lib/constants/attachment-menu';
 	import { useAttachmentMenu } from '$lib/hooks/use-attachment-menu.svelte';
 	import { useToolsPanel } from '$lib/hooks/use-tools-panel.svelte';
+	import { useReasoningMenu } from '$lib/hooks/use-reasoning-menu.svelte';
 	import { conversationsStore } from '$lib/stores/conversations.svelte';
 	import { mcpStore } from '$lib/stores/mcp.svelte';
 	import { McpLogo } from '$lib/components/app';
-	import { PencilRuler, ChevronDown, ChevronRight } from '@lucide/svelte';
+	import {
+		PencilRuler,
+		ChevronDown,
+		ChevronRight,
+		Lightbulb,
+		LightbulbOff,
+		Check
+	} from '@lucide/svelte';
 	import { HealthCheckStatus } from '$lib/enums';
 	import { AttachmentAction } from '$lib/enums/attachment.enums';
 
@@ -48,6 +56,7 @@
 	}: Props = $props();
 
 	let sheetOpen = $state(false);
+	let reasoningExpanded = $state(false);
 	let filesExpanded = $state(true);
 	let toolsExpanded = $state(false);
 	let mcpExpanded = $state(false);
@@ -67,6 +76,7 @@
 	);
 
 	const toolsPanel = useToolsPanel();
+	const reasoning = useReasoningMenu();
 
 	const sheetItemClass =
 		'flex w-full items-center gap-3 rounded-md px-3 py-2.5 text-left text-sm transition-colors hover:bg-accent active:bg-accent disabled:cursor-not-allowed disabled:opacity-50';
@@ -74,7 +84,7 @@
 	const sheetItemRowClass =
 		'flex w-full items-center justify-between gap-2 rounded-md px-3 py-2 text-left text-sm transition-colors hover:bg-accent';
 
-	let visibleMcpServers = $derived(mcpStore.visibleMcpServers);
+	let mcpServers = $derived(mcpStore.getServers());
 </script>
 
 <div class="flex items-center gap-1 {className}">
@@ -91,6 +101,63 @@
 			</Sheet.Header>
 
 			<div class="flex flex-col gap-1 px-1.5 pb-2">
+				{#if reasoning.modelSupportsThinking}
+					<Collapsible.Root
+						open={reasoningExpanded}
+						onOpenChange={(open) => (reasoningExpanded = open)}
+					>
+						<Collapsible.Trigger class={sheetItemClass}>
+							{#if reasoningExpanded}
+								<ChevronDown class="h-4 w-4 shrink-0" />
+							{:else}
+								<ChevronRight class="h-4 w-4 shrink-0" />
+							{/if}
+
+							{#if reasoning.thinkingEnabled}
+								<Lightbulb class="h-4 w-4 shrink-0 text-amber-400" />
+							{:else}
+								<LightbulbOff class="h-4 w-4 shrink-0 text-muted-foreground" />
+							{/if}
+
+							<span class="flex-1">Reasoning</span>
+
+							<span class="text-xs capitalize text-muted-foreground">
+								{reasoning.thinkingEnabled ? reasoning.currentEffort : 'off'}
+							</span>
+						</Collapsible.Trigger>
+
+						<Collapsible.Content>
+							<div class="flex flex-col gap-0.5 pl-4">
+								{#each reasoning.levels as level (level.value)}
+									{@const tokenLabel = reasoning.tokenLabel(level)}
+									<button
+										type="button"
+										class={sheetItemRowClass}
+										class:bg-accent={reasoning.isSelected(level)}
+										onclick={() => reasoning.select(level)}
+									>
+										<div class="flex min-w-0 items-center gap-3">
+											{#if reasoning.isSelected(level)}
+												<Check class="h-4 w-4 shrink-0 text-foreground" />
+											{:else}
+												<div class="h-4 w-4 shrink-0"></div>
+											{/if}
+
+											<span class="text-sm">{level.label}</span>
+										</div>
+
+										{#if tokenLabel}
+											<span class="shrink-0 text-[11px] text-muted-foreground opacity-60">
+												{tokenLabel}
+											</span>
+										{/if}
+									</button>
+								{/each}
+							</div>
+						</Collapsible.Content>
+					</Collapsible.Root>
+				{/if}
+
 				<Collapsible.Root open={filesExpanded} onOpenChange={(open) => (filesExpanded = open)}>
 					<Collapsible.Trigger class={sheetItemClass}>
 						{#if filesExpanded}
@@ -151,13 +218,13 @@
 						<span class="flex-1">MCP Servers</span>
 
 						<span class="text-xs text-muted-foreground">
-							{visibleMcpServers.length} server{visibleMcpServers.length !== 1 ? 's' : ''}
+							{mcpServers.length} server{mcpServers.length !== 1 ? 's' : ''}
 						</span>
 					</Collapsible.Trigger>
 
 					<Collapsible.Content>
 						<div class="flex flex-col gap-0.5 pl-4">
-							{#each visibleMcpServers as server (server.id)}
+							{#each mcpServers as server (server.id)}
 								{@const healthState = mcpStore.getHealthCheckState(server.id)}
 								{@const hasError = healthState.status === HealthCheckStatus.ERROR}
 								{@const displayName = mcpStore.getServerLabel(server)}
@@ -200,7 +267,7 @@
 								</button>
 							{/each}
 
-							{#if visibleMcpServers.length === 0}
+							{#if mcpServers.length === 0}
 								<div class="px-3 py-2 text-center text-sm text-muted-foreground">
 									No MCP servers configured
 								</div>

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/components/app/dialogs/DialogMcpServerAddNew.svelte

@@ -16,6 +16,7 @@
 
 	let newServerUrl = $state('');
 	let newServerHeaders = $state('');
+	let newServerUseProxy = $state(false);
 	let newServerUrlError = $derived.by(() => {
 		if (!newServerUrl.trim()) return 'URL is required';
 		try {
@@ -35,6 +36,7 @@
 		if (!value) {
 			newServerUrl = '';
 			newServerHeaders = '';
+			newServerUseProxy = false;
 		}
 		open = value;
 		onOpenChange?.(value);
@@ -49,7 +51,8 @@
 			id: newServerId,
 			enabled: true,
 			url: newServerUrl.trim(),
-			headers: newServerHeaders.trim() || undefined
+			headers: newServerHeaders.trim() || undefined,
+			useProxy: newServerUseProxy
 		});
 
 		conversationsStore.setMcpServerOverride(newServerId, true);
@@ -74,8 +77,10 @@
 				<McpServerForm
 					url={newServerUrl}
 					headers={newServerHeaders}
+					useProxy={newServerUseProxy}
 					onUrlChange={(v) => (newServerUrl = v)}
 					onHeadersChange={(v) => (newServerHeaders = v)}
+					onUseProxyChange={(v) => (newServerUseProxy = v)}
 					urlError={newServerUrl ? newServerUrlError : null}
 					id="new-server"
 				/>

tools/ui/src/lib/components/app/mcp/McpServerCard/McpServerCard.svelte

@@ -32,7 +32,9 @@
 	let isHealthChecking = $derived(healthState.status === HealthCheckStatus.CONNECTING);
 	let isConnected = $derived(healthState.status === HealthCheckStatus.SUCCESS);
 	let isError = $derived(healthState.status === HealthCheckStatus.ERROR);
-	let showSkeleton = $derived(isIdle || isHealthChecking);
+	// Disabled servers stay IDLE (no startup health check), so the body
+	// skeleton only applies while a check is running or expected to run.
+	let showSkeleton = $derived(isHealthChecking || (isIdle && server.enabled));
 	let errorMessage = $derived(
 		healthState.status === HealthCheckStatus.ERROR ? healthState.message : undefined
 	);

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

@@ -22,7 +22,9 @@
 
 	let { class: className }: Props = $props();
 
-	let servers = $derived(mcpStore.visibleMcpServers);
+	// Every configured server is listed; `enabled` is an on/off state,
+	// not a visibility filter, so a disabled server stays toggleable.
+	let servers = $derived(mcpStore.getServers());
 
 	let isAddingServer = $state(false);
 
@@ -58,9 +60,14 @@
 	// Each card decides for itself whether to render based on its own
 	// health-check state, so adding a server only flashes the new card
 	// (not every other already-loaded card) until its health check resolves.
-	function isServerPending(serverId: string): boolean {
+	// Disabled servers never receive a startup health check, so IDLE only
+	// counts as pending when the server is enabled; otherwise the real card
+	// renders and keeps the enable toggle reachable.
+	function isServerPending(serverId: string, enabled: boolean): boolean {
 		const status = mcpStore.getHealthCheckState(serverId).status;
-		return status === HealthCheckStatus.IDLE || status === HealthCheckStatus.CONNECTING;
+		return (
+			status === HealthCheckStatus.CONNECTING || (status === HealthCheckStatus.IDLE && enabled)
+		);
 	}
 </script>
 
@@ -109,7 +116,7 @@
 			style="grid-template-columns: repeat(auto-fill, minmax(min(32rem, calc(100dvw - 2rem)), 1fr));"
 		>
 			{#each servers as server (server.id)}
-				{#if isServerPending(server.id)}
+				{#if isServerPending(server.id, server.enabled)}
 					<McpServerCardSkeleton />
 				{:else}
 					<McpServerCard

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

@@ -6,8 +6,7 @@ export const NEWLINE_SEPARATOR = '\n';
 
 export const DEFAULT_AGENTIC_CONFIG: AgenticConfig = {
 	enabled: true,
-	maxTurns: 100,
-	maxToolPreviewLines: 25
+	maxTurns: 100
 } as const;
 
 export const REASONING_TAGS = {

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/src/lib/constants/routes.ts

@@ -8,7 +8,6 @@ export const SETTINGS_SECTION_SLUGS = {
 	PENALTIES: 'penalties',
 	AGENTIC: 'agentic',
 	DEVELOPER: 'developer',
-	MCP: 'mcp',
 	TOOLS: 'tools',
 	IMPORT_EXPORT: 'import-export'
 } as const;

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

@@ -59,7 +59,6 @@ export const SETTINGS_KEYS = {
 	MCP_SERVERS: 'mcpServers',
 	MCP_REQUEST_TIMEOUT_SECONDS: 'mcpRequestTimeoutSeconds',
 	AGENTIC_MAX_TURNS: 'agenticMaxTurns',
-	AGENTIC_MAX_TOOL_PREVIEW_LINES: 'agenticMaxToolPreviewLines',
 	SHOW_TOOL_CALL_IN_PROGRESS: 'showToolCallInProgress',
 	// Performance
 	PRE_ENCODE_CONVERSATION: 'preEncodeConversation',

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

@@ -24,7 +24,6 @@ import type {
 	SettingsSection
 } from '$lib/types';
 import { CLI_FLAGS, DEFAULT_MCP_CONFIG } from '$lib/constants';
-import McpLogo from '$lib/components/app/mcp/McpLogo.svelte';
 import { SETTINGS_KEYS } from './settings-keys';
 import { ROUTES, SETTINGS_SECTION_SLUGS } from './routes';
 import { TITLE_GENERATION } from './title-generation';
@@ -36,7 +35,6 @@ export const SETTINGS_SECTION_TITLES = {
 	PENALTIES: 'Penalties',
 	AGENTIC: 'Agentic',
 	TOOLS: 'Tools',
-	MCP: 'MCP',
 	IMPORT_EXPORT: 'Import/Export',
 	DEVELOPER: 'Developer'
 } as const;
@@ -635,15 +633,15 @@ const SETTINGS_REGISTRY: Record<string, SettingsSectionEntry> = {
 				}
 			},
 			{
-				key: SETTINGS_KEYS.AGENTIC_MAX_TOOL_PREVIEW_LINES,
-				label: 'Max lines per tool preview',
-				help: 'Number of lines shown in tool output previews (last N lines). Only these previews and the final LLM response persist after the agentic loop completes.',
-				defaultValue: 25,
+				key: SETTINGS_KEYS.MCP_REQUEST_TIMEOUT_SECONDS,
+				label: 'MCP request timeout (seconds)',
+				help: 'Timeout for individual MCP tool calls.',
+				defaultValue: DEFAULT_MCP_CONFIG.requestTimeoutSeconds,
 				type: SettingsFieldType.INPUT,
 				section: SETTINGS_SECTION_SLUGS.AGENTIC,
 				isPositiveInteger: true,
 				sync: {
-					serverKey: SETTINGS_KEYS.AGENTIC_MAX_TOOL_PREVIEW_LINES,
+					serverKey: SETTINGS_KEYS.MCP_REQUEST_TIMEOUT_SECONDS,
 					paramType: SyncableParameterType.NUMBER
 				}
 			}
@@ -735,26 +733,6 @@ const SETTINGS_REGISTRY: Record<string, SettingsSectionEntry> = {
 				}
 			}
 		]
-	},
-	[SETTINGS_SECTION_SLUGS.MCP]: {
-		title: SETTINGS_SECTION_TITLES.MCP,
-		slug: SETTINGS_SECTION_SLUGS.MCP,
-		icon: McpLogo,
-		settings: [
-			{
-				key: SETTINGS_KEYS.MCP_REQUEST_TIMEOUT_SECONDS,
-				label: 'Request timeout (seconds)',
-				help: 'Default timeout for individual MCP tool calls. Can be overridden per server.',
-				defaultValue: DEFAULT_MCP_CONFIG.requestTimeoutSeconds,
-				type: SettingsFieldType.INPUT,
-				section: SETTINGS_SECTION_SLUGS.MCP,
-				isPositiveInteger: true,
-				sync: {
-					serverKey: SETTINGS_KEYS.MCP_REQUEST_TIMEOUT_SECONDS,
-					paramType: SyncableParameterType.NUMBER
-				}
-			}
-		]
 	}
 } as const;
 

tools/ui/src/lib/hooks/use-reasoning-menu.svelte.ts

@@ -0,0 +1,96 @@
+import { ReasoningEffort } from '$lib/enums';
+import { REASONING_EFFORT_LEVELS } from '$lib/constants/reasoning-effort';
+import { REASONING_EFFORT_TOKENS } from '$lib/constants/reasoning-effort-tokens';
+import type { ReasoningEffortLevel } from '$lib/types';
+import type { DatabaseMessage } from '$lib/types/database';
+import {
+	modelsStore,
+	checkModelSupportsThinking,
+	supportsThinking,
+	propsCacheVersion,
+	loadedModelIds
+} from '$lib/stores/models.svelte';
+import { chatStore } from '$lib/stores/chat.svelte';
+import { conversationsStore, activeMessages } from '$lib/stores/conversations.svelte';
+import { isRouterMode } from '$lib/stores/server.svelte';
+
+export interface UseReasoningMenuReturn {
+	readonly modelSupportsThinking: boolean;
+	readonly thinkingEnabled: boolean;
+	readonly currentEffort: ReasoningEffort;
+	readonly levels: ReasoningEffortLevel[];
+	isSelected(level: ReasoningEffortLevel): boolean;
+	tokenLabel(level: ReasoningEffortLevel): string | null;
+	select(level: ReasoningEffortLevel): void;
+}
+
+/**
+ * Shared reactive state and helpers for the reasoning effort menu.
+ *
+ * Used by both the desktop dropdown (`ChatFormActionAddReasoningSubmenu`)
+ * and the mobile sheet (`ChatFormActionAddSheet`) to avoid duplicating the
+ * thinking-support derivation and the effort selection logic.
+ */
+export function useReasoningMenu(): UseReasoningMenuReturn {
+	const conversationModel = $derived(
+		chatStore.getConversationModel(activeMessages() as DatabaseMessage[])
+	);
+
+	// a router chat can carry reasoning from an earlier turn before the props
+	// cache is primed, so a model that already produced thinking still qualifies
+	const modelSupportsThinkingFromMessages = $derived.by(() => {
+		const modelId = isRouterMode() ? modelsStore.selectedModelName || conversationModel : null;
+		if (!modelId) return false;
+
+		return conversationsStore.activeMessages.some(
+			(m) => m.role === 'assistant' && m.model === modelId && !!m.reasoningContent
+		);
+	});
+
+	const modelSupportsThinking = $derived.by(() => {
+		loadedModelIds();
+		propsCacheVersion();
+
+		if (isRouterMode()) {
+			const modelId = modelsStore.selectedModelName || conversationModel;
+			return checkModelSupportsThinking(modelId ?? '') || modelSupportsThinkingFromMessages;
+		}
+
+		return supportsThinking() || modelSupportsThinkingFromMessages;
+	});
+
+	const thinkingEnabled = $derived(conversationsStore.getThinkingEnabled());
+	const currentEffort = $derived(conversationsStore.getReasoningEffort());
+
+	return {
+		get modelSupportsThinking() {
+			return modelSupportsThinking;
+		},
+		get thinkingEnabled() {
+			return thinkingEnabled;
+		},
+		get currentEffort() {
+			return currentEffort;
+		},
+		get levels() {
+			return REASONING_EFFORT_LEVELS;
+		},
+		isSelected(level: ReasoningEffortLevel): boolean {
+			if (level.isOff) return !thinkingEnabled;
+			return thinkingEnabled && currentEffort === level.value;
+		},
+		tokenLabel(level: ReasoningEffortLevel): string | null {
+			if (level.isOff) return null;
+			const tokens = REASONING_EFFORT_TOKENS[level.value];
+			return tokens === -1 ? 'Unlimited' : `Max ${tokens.toLocaleString()} tokens`;
+		},
+		select(level: ReasoningEffortLevel): void {
+			if (level.isOff) {
+				conversationsStore.setThinkingEnabled(false);
+				return;
+			}
+			conversationsStore.setThinkingEnabled(true);
+			conversationsStore.setReasoningEffort(level.value as ReasoningEffort);
+		}
+	};
+}

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

@@ -692,8 +692,31 @@ export class MCPService {
 			this.createLog(MCPConnectionPhase.INITIALIZING, 'Sending initialize request...')
 		);
 
+		// The SDK timeout only covers the initialize request, not transport.start(),
+		// which can hang forever on an unreachable host (SSE endpoint wait, WebSocket
+		// handshake, proxied fetch). This race bounds the whole handshake and closes
+		// the transport on expiry so the underlying fetch or socket is aborted.
+		const handshakeTimeoutMs =
+			serverConfig.handshakeTimeoutMs ?? DEFAULT_MCP_CONFIG.connectionTimeoutMs;
+
 		try {
-			await client.connect(transport);
+			let handshakeTimer: ReturnType<typeof setTimeout> | undefined;
+			const handshakeDeadline = new Promise<never>((_, reject) => {
+				handshakeTimer = setTimeout(() => {
+					void transport.close().catch(() => {});
+					reject(new Error(`Connection timed out after ${Math.round(handshakeTimeoutMs / 1000)}s`));
+				}, handshakeTimeoutMs);
+			});
+
+			try {
+				await Promise.race([
+					client.connect(transport, { timeout: handshakeTimeoutMs }),
+					handshakeDeadline
+				]);
+			} finally {
+				clearTimeout(handshakeTimer);
+			}
+
 			// Transport diagnostics are only for the initial handshake, not long-lived traffic.
 			stopPhaseLogging();
 			client.onerror = runtimeErrorHandler;

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

@@ -280,16 +280,13 @@ class AgenticStore {
 
 	getConfig(settings: SettingsConfigType, perChatOverrides?: McpServerOverride[]): AgenticConfig {
 		const maxTurns = Number(settings.agenticMaxTurns) || DEFAULT_AGENTIC_CONFIG.maxTurns;
-		const maxToolPreviewLines =
-			Number(settings.agenticMaxToolPreviewLines) || DEFAULT_AGENTIC_CONFIG.maxToolPreviewLines;
 		const hasTools =
 			mcpStore.hasEnabledServers(perChatOverrides) ||
 			toolsStore.builtinTools.length > 0 ||
 			toolsStore.customTools.length > 0;
 		return {
 			enabled: hasTools && DEFAULT_AGENTIC_CONFIG.enabled,
-			maxTurns,
-			maxToolPreviewLines
+			maxTurns
 		};
 	}
 

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

@@ -1334,7 +1334,7 @@ class ChatStore {
 			}
 		};
 
-		const perChatOverrides = conversationsStore.activeConversation?.mcpServerOverrides;
+		const perChatOverrides = conversationsStore.getAllMcpServerOverrides();
 
 		{
 			const agenticResult = await agenticStore.runAgenticFlow({

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

@@ -243,9 +243,9 @@ class ConversationsStore {
 		const conversationName = name || `Chat ${new Date().toLocaleString()}`;
 		const conversation = await DatabaseService.createConversation(conversationName);
 
-		// New conversations inherit per-server enabled defaults directly from
-		// `mcpServers[i].enabled` (see #checkServerEnabled). No per-conversation
-		// override list needs to be seeded.
+		// No MCP override list is seeded: getAllMcpServerOverrides resolves
+		// servers without a per-conversation override to `mcpServers[i].enabled`,
+		// and only explicit toggles are stored on the conversation.
 
 		// Inherit global thinking/reasoning defaults into the new conversation
 		const thinkingEnabled = this.getThinkingEnabled();
@@ -601,48 +601,41 @@ class ConversationsStore {
 	 */
 
 	/**
-	/**
-	 * Resolve the per-server enabled value when no active conversation exists.
-	 * The default for new chats is the server's own `enabled` flag in `mcpServers`.
+	 * Resolve the default enabled value for a server: its own `enabled`
+	 * flag in `mcpServers`, so the global on/off state lives in one place.
 	 */
-	#getDefaultOverrideForNoConversation(serverId: string): McpServerOverride | undefined {
+	#getDefaultOverride(serverId: string): McpServerOverride | undefined {
 		const server = mcpStore.getServers().find((s) => s.id === serverId);
 		if (!server) return undefined;
 		return { serverId, enabled: server.enabled };
 	}
 
 	/**
-	 * Default overrides for new chats are derived from `mcpServers[i].enabled`,
-	 * so the global on/off state lives in one place.
-	 */
-	#getAllDefaultOverridesForNoConversation(): McpServerOverride[] {
-		return mcpStore.getServers().map((s) => ({ serverId: s.id, enabled: s.enabled }));
-	}
-
-	/**
-	 * Gets MCP server override for a specific server in the active conversation.
-	 * Falls back to `mcpServers[i].enabled` if no active conversation exists.
+	 * Gets the effective MCP server override for a specific server.
+	 * A per-conversation override wins when present; a server without one
+	 * resolves to its `mcpServers[i].enabled` default.
 	 * @param serverId - The server ID to check
-	 * @returns The override if set, undefined if no matching server
+	 * @returns The effective override, undefined if no matching server
 	 */
 	getMcpServerOverride(serverId: string): McpServerOverride | undefined {
-		if (this.activeConversation) {
-			return this.activeConversation.mcpServerOverrides?.find(
-				(o: McpServerOverride) => o.serverId === serverId
-			);
-		}
-		return this.#getDefaultOverrideForNoConversation(serverId);
+		const override = this.activeConversation?.mcpServerOverrides?.find(
+			(o: McpServerOverride) => o.serverId === serverId
+		);
+		if (override) return override;
+		return this.#getDefaultOverride(serverId);
 	}
 
 	/**
-	 * Get all MCP server overrides for the current conversation.
-	 * When no active conversation, derives from `mcpServers[i].enabled`.
+	 * Gets the effective override list for the current conversation:
+	 * one entry per configured server, resolved per server. The stored
+	 * per-conversation list is sparse and only holds explicit toggles.
 	 */
 	getAllMcpServerOverrides(): McpServerOverride[] {
-		if (this.activeConversation?.mcpServerOverrides) {
-			return this.activeConversation.mcpServerOverrides;
-		}
-		return this.#getAllDefaultOverridesForNoConversation();
+		const overrides = this.activeConversation?.mcpServerOverrides;
+		return mcpStore.getServers().map((s) => {
+			const override = overrides?.find((o: McpServerOverride) => o.serverId === s.id);
+			return { serverId: s.id, enabled: override?.enabled ?? s.enabled };
+		});
 	}
 
 	/**