ui: move Models, MCP and Settings to a sidebar footer (#32022 )

* ui : move Models, MCP and Settings to a sidebar footer The sidebar's actions split in two: New chat, Search and Projects stay at its top, and Models, MCP and Settings move to a footer at its bottom, under a thin rule, in the expanded sidebar and its collapsed strip alike. On a phone the footer keeps clear of the home indicator, and Models stays desktop-only as before. The search dialog still lists every action. Assisted-by: Claude (Opus 5.5 Medium) * ui : group the sidebar footer in an inset surface The design system gains inset-group-surface, for rows grouped inside a glass panel: white on the light glass and a faint tint of the text colour on the dark one, with the sidebar's own ambient edge, no inner highlight, a faint wide shade under it and the card radius. It is matte-glass pointed at its own component tokens, like the sidebar and chat form surfaces, and reads as part of the panel: no tight contact shadow or solid border like a button, no inset shade like a field. The sidebar footer uses it in place of the rule above Models, MCP and Settings. DESIGN.md and the Surfaces page document it. Assisted-by: Claude (Opus 5.5 Medium) * ui : give the sidebar's logo button the add button's key look Whenever the sidebar shows, open or peeking, its logo button takes the same key as the chat form's add button (matte-key, white in light mode, a light tint in dark), in place of its own muted fill. While the sidebar is hidden the logo has no fill and no shadow, and hover and focus never change the fill. It keeps its round shape and its own size. Assisted-by: Claude (Opus 5.5 Medium) * ui : put search in a round key at the top of the phone sidebar On a phone the sidebar's Search row moves to a round key at the right end of its top row, across from the logo, with the same key look as the logo and the chat form's add button. It opens search as the row did. Desktop keeps the Search row. Assisted-by: Claude (Opus 5.5 Medium) * ui : size the phone sidebar for a thumb and list its footer with the actions On a phone the sidebar's actions and rows are 48px tall at 17px, flush one under the other with 16px corners, its section headings 15px and grey, and its groups 28px apart. MCP and Settings join the top actions there: the footer is a desktop pattern. Assisted-by: Claude (Opus 5.5 Medium) * ui : mark the selected sidebar row with a tinted inset surface Assisted-by: Claude (Opus 5.5 Medium) * ui: keep the phone sidebar's round keys white on the grey page Assisted-by: Claude (Opus 5.5 Medium) * ui : open Settings and MCP as pages on a phone, dialogs elsewhere On a phone the sidebar opens them at #/settings and #/mcp-servers, with the menu key and a title like any phone page. Pinned under the status bar, Settings' section tabs fill the screen above them in the page's colour. Assisted-by: Claude (Opus 5.5 Medium) * ui : open a phone sidebar's pages behind it, settled before they slide back A tap in the phone sidebar now opens its page while the current one is still aside and slides back once it has rendered, instead of swapping pages mid-slide. The slid page's fixed bars are pinned again as the page under them changes, the lock ends with the slide itself, and a phone's new chat shows its form and greeting at once, without the move and entry they have on a desktop. Assisted-by: Claude (Opus 5.5 Medium) * ui: style the sidebar footer group and the selected row with master's classes until the design system lands Assisted-by: Claude (Opus 5.5 Medium) * ui: add GlassKey with master's styles for the phone sidebar's round keys Assisted-by: Claude (Opus 5.5 Medium) * ui: center the top edge of a new chat's form on the screen The form rested 42px below the middle up to 2xl and 6px above it past that; one drop, from the resting card's height, now puts its top edge on the middle at every width from md up. Assisted-by: Claude (Opus 5.5 Medium) * ui: round the sidebar's action rows concentric with the footer group The sidebar stacks four surfaces: the page, the sidebar, the footer group inset in it and a row in the group. app.css now names the group's radius and padding, and derives the row radius from them (16px less 4px), so the footer's rows run parallel to its corners; the actions at the top of the sidebar share that radius. The footer moves to its own component, the README documents the pattern, and a client test holds the rule on the rendered footer. A selected action row was bg-accent, the same grey as a hovered row in light and darker than the footer's fill in dark, so it read as hovered or not at all; it is now raised and edged, while a hover only tints. Assisted-by: Claude (Opus 5.5 Medium) * ui: import demo chats through the Settings page on a phone A phone opens Settings as a page now, not a dialog, so the demo harness waited for a dialog that never came and every phone demo with traces failed; it now imports through whichever of the two opens. Assisted-by: Claude (Opus 5.5 Medium) * ui: count the sidebar's rows from their items in the surface test The test expected three footer rows and two top actions, which the stacked PRs change as they add Projects, Library and Servers; it now expects as many rows as the sidebar has items. Assisted-by: Claude (Opus 5.5 Medium) * ui: add a line on Meta+Enter off macOS too The chat form takes Meta+Enter as Ctrl+Enter and leaves the newline to the editor, but the editor bound only Mod-Enter, which is Ctrl off macOS, so on Linux Meta+Enter neither sent nor added a line. This failed the client test on CI's Linux runners, on master too. Assisted-by: Claude (Opus 5.5 Medium)
cuda : remove unused ids and type_src0 parameters (#32018 )
2026-11-10 12:57:40 -06:00 · 2026-11-10 18:58:57 +01:00 · 2026-11-10 18:26:53 +01:00 · 2026-11-10 18:08:32 +02:00 · 2026-11-10 17:12:27 +02:00 · 2026-11-10 17:10:11 +02:00
240 changed files with 9070 additions and 5239 deletions
@@ -1,4 +1,4 @@
-ARG ONEAPI_VERSION=2026.1.1-devel-ubuntu24.04
+ARG ONEAPI_VERSION=2026.2.0-devel-ubuntu24.04
 ARG BUILD_DATE=N/A
 ARG APP_VERSION=N/A
 ARG APP_REVISION=N/A
@@ -22,20 +22,20 @@ RUN LLAMA_BUILD_NUMBER="$APP_VERSION" npm run build
 FROM docker.io/intel/deep-learning-essentials:$ONEAPI_VERSION AS build

 RUN echo "Download oneDNN installation package" && \
-    wget https://registrationcenter-download.intel.com/akdlm/IRC_NAS/94c3dbac-0852-45be-a57d-21c204cada3e/intel-onednn-2026.0.2.46_offline.sh -O intel-onednn_offline.sh && \
+    wget https://registrationcenter-download.intel.com/akdlm/IRC_NAS/46452a4a-7b7f-49f2-be4c-602a900180bb/intel-onednn-2026.1.0.625_offline.sh -O intel-onednn_offline.sh && \
    echo "Install oneDNN components" && \
    bash intel-onednn_offline.sh -s -a --silent --eula accept

 ARG GGML_SYCL_F16=ON
-ARG LEVEL_ZERO_VERSION=1.28.2
+ARG LEVEL_ZERO_VERSION=1.34.0
 ARG LEVEL_ZERO_UBUNTU_VERSION=u24.04
 RUN apt-get update && \
    apt-get install -y git libssl-dev wget ca-certificates && \
    cd /tmp && \
-    wget -q "https://github.com/oneapi-src/level-zero/releases/download/v${LEVEL_ZERO_VERSION}/level-zero_${LEVEL_ZERO_VERSION}%2B${LEVEL_ZERO_UBUNTU_VERSION}_amd64.deb" -O level-zero.deb && \
-    wget -q "https://github.com/oneapi-src/level-zero/releases/download/v${LEVEL_ZERO_VERSION}/level-zero-devel_${LEVEL_ZERO_VERSION}%2B${LEVEL_ZERO_UBUNTU_VERSION}_amd64.deb" -O level-zero-devel.deb && \
-    apt-get -o Dpkg::Options::="--force-overwrite" install -y ./level-zero.deb ./level-zero-devel.deb && \
-    rm -f /tmp/level-zero.deb /tmp/level-zero-devel.deb
+    wget -q "https://github.com/oneapi-src/level-zero/releases/download/v${LEVEL_ZERO_VERSION}/libze1_${LEVEL_ZERO_VERSION}%2B${LEVEL_ZERO_UBUNTU_VERSION}_amd64.deb" -O libze1.deb && \
+    wget -q "https://github.com/oneapi-src/level-zero/releases/download/v${LEVEL_ZERO_VERSION}/libze-dev_${LEVEL_ZERO_VERSION}%2B${LEVEL_ZERO_UBUNTU_VERSION}_amd64.deb" -O libze-dev.deb && \
+    apt-get -o Dpkg::Options::="--force-overwrite" install -y ./libze1.deb ./libze-dev.deb && \
+    rm -f /tmp/libze1.deb /tmp/libze-dev.deb

 WORKDIR /app

@@ -91,11 +91,11 @@ COPY --from=build /opt/intel/oneapi/dnnl /opt/intel/oneapi/dnnl
 #ARG IGDGMM_VERSION=22.8.2


-ARG IGC_VERSION=v2.34.4
-ARG IGC_VERSION_FULL=2_2.34.4+21428
-ARG COMPUTE_RUNTIME_VERSION=26.18.38308.1
-ARG COMPUTE_RUNTIME_VERSION_FULL=26.18.38308.1-0
-ARG IGDGMM_VERSION=22.10.0
+ARG IGC_VERSION=v2.42.6
+ARG IGC_VERSION_FULL=2_2.42.6+23070
+ARG COMPUTE_RUNTIME_VERSION=26.40.40196.8
+ARG COMPUTE_RUNTIME_VERSION_FULL=26.40.40196.8-0
+ARG IGDGMM_VERSION=22.10.2
 RUN mkdir /tmp/neo/ && cd /tmp/neo/ \
  && wget https://github.com/intel/intel-graphics-compiler/releases/download/$IGC_VERSION/intel-igc-core-${IGC_VERSION_FULL}_amd64.deb \
  && wget https://github.com/intel/intel-graphics-compiler/releases/download/$IGC_VERSION/intel-igc-opencl-${IGC_VERSION_FULL}_amd64.deb \
@@ -37,7 +37,7 @@ env:
  LLAMA_ARG_LOG_PREFIX: 1
  LLAMA_ARG_LOG_TIMESTAMPS: 1
  # TODO: fix failing tests on OpenVINO backend
-  CTEST_EXCLUDE: "test-llama-archs|test-llama-context"
+  CTEST_EXCLUDE: "test-llama-archs|test-llama-context|test-fp8-kv"

 jobs:
  ubuntu-24-openvino:
@@ -52,8 +52,8 @@ jobs:

    env:
      ONEAPI_ROOT: /opt/intel/oneapi/
-      ONEAPI_INSTALLER_VERSION: "2026.1"
-      LEVEL_ZERO_VERSION: "1.33.1"
+      ONEAPI_INSTALLER_VERSION: "2026.2"
+      LEVEL_ZERO_VERSION: "1.34.0"
      LEVEL_ZERO_UBUNTU_VERSION: "u24.04"

    continue-on-error: true
@@ -68,12 +68,12 @@ jobs:
        run: |
          cd /tmp
          echo "Download oneAPI Deep Learning Essentials installation package"
-          wget https://registrationcenter-download.intel.com/akdlm/IRC_NAS/b5c100a7-8052-4e88-90e2-63d9916a998b/intel-deep-learning-essentials-2026.1.4.17_offline.sh -O intel-deep-learning-essentials_offline.sh
+          wget https://registrationcenter-download.intel.com/akdlm/IRC_NAS/aef232d2-aac5-4cd6-bceb-911019bbb6a8/intel-deep-learning-essentials-2026.2.0.640_offline.sh -O intel-deep-learning-essentials_offline.sh
          echo "Install oneAPI Deep Learning Essentials components"
          sudo bash intel-deep-learning-essentials_offline.sh -s -a --silent --eula accept

          echo "Download oneDNN installation package"
-          wget https://registrationcenter-download.intel.com/akdlm/IRC_NAS/94c3dbac-0852-45be-a57d-21c204cada3e/intel-onednn-2026.0.2.46_offline.sh -O intel-onednn_offline.sh
+          wget https://registrationcenter-download.intel.com/akdlm/IRC_NAS/46452a4a-7b7f-49f2-be4c-602a900180bb/intel-onednn-2026.1.0.625_offline.sh -O intel-onednn_offline.sh
          echo "Install oneDNN components"
          sudo bash intel-onednn_offline.sh -s -a --silent --eula accept

@@ -150,13 +150,13 @@ jobs:
        shell: bash

    env:
-      WINDOWS_DLE_URL: https://registrationcenter-download.intel.com/akdlm/IRC_NAS/a51104eb-6f4a-4a0e-9dae-1b81b0c81ae1/intel-deep-learning-essentials-2026.1.4.16_offline.exe
-      WINDOWS_ONEDNN_URL: https://registrationcenter-download.intel.com/akdlm/IRC_NAS/34e28083-d04b-46a3-826c-8c1d00bf1e1c/intel-onednn-2026.0.2.47_offline.exe
+      WINDOWS_DLE_URL: https://registrationcenter-download.intel.com/akdlm/IRC_NAS/b177a006-5cd9-4cb2-807f-8199cc73d53c/intel-deep-learning-essentials-2026.2.0.638_offline.exe
+      WINDOWS_ONEDNN_URL: https://registrationcenter-download.intel.com/akdlm/IRC_NAS/8f905666-c85c-43fb-9bd0-2618c7ecd33b/intel-onednn-2026.1.0.622_offline.exe
      WINDOWS_DPCPP_MKL: intel.oneapi.win.cpp-dpcpp-common:intel.oneapi.win.mkl.devel:intel.oneapi.win.tbb.devel
      WINDOWS_DNNL: intel.oneapi.win.dnnl
-      LEVEL_ZERO_SDK_URL: https://github.com/oneapi-src/level-zero/releases/download/v1.33.1/level-zero-win-sdk-1.33.1.zip
+      LEVEL_ZERO_SDK_URL: https://github.com/oneapi-src/level-zero/releases/download/v1.34.0/level-zero-win-sdk-1.34.0.zip
      ONEAPI_ROOT: "C:/Program Files (x86)/Intel/oneAPI"
-      ONEAPI_INSTALLER_VERSION: "2026.1"
+      ONEAPI_INSTALLER_VERSION: "2026.2"
    steps:
      - name: Clone
        id: checkout
@@ -1307,14 +1307,14 @@ jobs:
        shell: bash

    env:
-      WINDOWS_DLE_URL: https://registrationcenter-download.intel.com/akdlm/IRC_NAS/a51104eb-6f4a-4a0e-9dae-1b81b0c81ae1/intel-deep-learning-essentials-2026.1.4.16_offline.exe
-      WINDOWS_ONEDNN_URL: https://registrationcenter-download.intel.com/akdlm/IRC_NAS/34e28083-d04b-46a3-826c-8c1d00bf1e1c/intel-onednn-2026.0.2.47_offline.exe
+      WINDOWS_DLE_URL: https://registrationcenter-download.intel.com/akdlm/IRC_NAS/b177a006-5cd9-4cb2-807f-8199cc73d53c/intel-deep-learning-essentials-2026.2.0.638_offline.exe
+      WINDOWS_ONEDNN_URL: https://registrationcenter-download.intel.com/akdlm/IRC_NAS/8f905666-c85c-43fb-9bd0-2618c7ecd33b/intel-onednn-2026.1.0.622_offline.exe
      WINDOWS_DPCPP_MKL: intel.oneapi.win.cpp-dpcpp-common:intel.oneapi.win.mkl.devel:intel.oneapi.win.tbb.devel
      WINDOWS_DNNL: intel.oneapi.win.dnnl

-      LEVEL_ZERO_SDK_URL: https://github.com/oneapi-src/level-zero/releases/download/v1.33.1/level-zero-win-sdk-1.33.1.zip
+      LEVEL_ZERO_SDK_URL: https://github.com/oneapi-src/level-zero/releases/download/v1.34.0/level-zero-win-sdk-1.34.0.zip
      ONEAPI_ROOT: "C:/Program Files (x86)/Intel/oneAPI"
-      ONEAPI_INSTALLER_VERSION: "2026.1"
+      ONEAPI_INSTALLER_VERSION: "2026.2"

    steps:
      - name: Clone
@@ -1366,8 +1366,7 @@ jobs:

          cp "${{ env.ONEAPI_ROOT }}/compiler/latest/bin/ur_adapter_level_zero.dll" ./build/bin
          cp "${{ env.ONEAPI_ROOT }}/compiler/latest/bin/ur_adapter_level_zero_v2.dll" ./build/bin
-          cp "${{ env.ONEAPI_ROOT }}/compiler/latest/bin/ur_adapter_opencl.dll" ./build/bin
-          cp "${{ env.ONEAPI_ROOT }}/compiler/latest/bin/ur_loader.dll" ./build/bin
+          cp "${{ env.ONEAPI_ROOT }}/compiler/latest/bin/ur_loader_shared.dll" ./build/bin
          cp "${{ env.ONEAPI_ROOT }}/compiler/latest/bin/ur_win_proxy_loader.dll" ./build/bin
          ZE_LOADER_DLL=$(find "${{ env.ONEAPI_ROOT }}" "$LEVEL_ZERO_V1_SDK_PATH" -iname ze_loader.dll -print -quit 2>/dev/null || true)
          if [ -n "$ZE_LOADER_DLL" ]; then
@@ -1421,8 +1420,8 @@ jobs:

    env:
      ONEAPI_ROOT: /opt/intel/oneapi/
-      ONEAPI_INSTALLER_VERSION: "2026.1"
-      LEVEL_ZERO_VERSION: "1.33.1"
+      ONEAPI_INSTALLER_VERSION: "2026.2"
+      LEVEL_ZERO_VERSION: "1.34.0"
      LEVEL_ZERO_UBUNTU_VERSION: "u24.04"

    steps:
@@ -1437,12 +1436,12 @@ jobs:
        run: |
          cd /tmp
          echo "Download oneAPI Deep Learning Essentials installation package"
-          wget https://registrationcenter-download.intel.com/akdlm/IRC_NAS/b5c100a7-8052-4e88-90e2-63d9916a998b/intel-deep-learning-essentials-2026.1.4.17_offline.sh -O intel-deep-learning-essentials_offline.sh
+          wget https://registrationcenter-download.intel.com/akdlm/IRC_NAS/aef232d2-aac5-4cd6-bceb-911019bbb6a8/intel-deep-learning-essentials-2026.2.0.640_offline.sh -O intel-deep-learning-essentials_offline.sh
          echo "Install oneAPI Deep Learning Essentials components"
          sudo bash intel-deep-learning-essentials_offline.sh -s -a --silent --eula accept

          echo "Download oneDNN installation package"
-          wget https://registrationcenter-download.intel.com/akdlm/IRC_NAS/94c3dbac-0852-45be-a57d-21c204cada3e/intel-onednn-2026.0.2.46_offline.sh -O intel-onednn_offline.sh
+          wget https://registrationcenter-download.intel.com/akdlm/IRC_NAS/46452a4a-7b7f-49f2-be4c-602a900180bb/intel-onednn-2026.1.0.625_offline.sh -O intel-onednn_offline.sh
          echo "Install oneDNN components"
          sudo bash intel-onednn_offline.sh -s -a --silent --eula accept

@@ -180,7 +180,7 @@ if [ ! -z ${GG_BUILD_OPENVINO} ]; then
    CMAKE_EXTRA="${CMAKE_EXTRA} -DGGML_OPENVINO=ON"

    # TODO: fix failing tests on OpenVINO backend
-    CTEST_EXTRA="-E test-llama-archs|test-llama-context"
+    CTEST_EXTRA="-E test-llama-archs|test-llama-context|test-fp8-kv"
 fi

 ## helpers
@@ -282,6 +282,7 @@ const std::vector<ggml_type> kv_cache_types = {
    GGML_TYPE_F32,
    GGML_TYPE_F16,
    GGML_TYPE_BF16,
+    GGML_TYPE_F8_E4M3,
    GGML_TYPE_Q8_0,
    GGML_TYPE_Q4_0,
    GGML_TYPE_Q4_1,
@@ -4491,7 +4492,7 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
        [](common_params & params, const std::string & value) {
            params.tts_lang = value;
        }
-    ).set_examples({LLAMA_EXAMPLE_TTS}));
+    ).set_examples({LLAMA_EXAMPLE_TTS, LLAMA_EXAMPLE_SERVER}));
    add_opt(common_arg(
        {"--tts-voice"}, "VOICE",
        "voice for audio generation (model-dependent)\n"
@@ -4499,7 +4500,7 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
        [](common_params & params, const std::string & value) {
            params.tts_voice = value;
        }
-    ).set_examples({LLAMA_EXAMPLE_TTS}));
+    ).set_examples({LLAMA_EXAMPLE_TTS, LLAMA_EXAMPLE_SERVER}));
    add_opt(common_arg(
        {"--tts-speaker-file"}, "FNAME",
        "speaker file path for audio generation",
@@ -131,6 +131,12 @@ static std::string normalize_quotes_to_json(const std::string & input) {
            continue;
        }

+        // Models can emit raw control characters inside string literals; escape them to keep the JSON valid.
+        if ((in_single_quoted || in_double_quoted) && static_cast<unsigned char>(c) < 0x20) {
+            result += escape_json_string_inner(std::string(1, c));
+            continue;
+        }
+
        // Handle quote characters
        if (c == '"') {
            if (in_single_quoted) {
@@ -819,21 +819,29 @@ static bool is_lfm3_template(const std::string & src) {
    return src.find("<additional_tools>") != std::string::npos;
 }

-// LFM3-Audio: a known language code maps to its name, no language lets the model detect it
-static std::string lfm3_asr_system_prompt(const std::string & language) {
+static std::string string_lower(std::string s) {
+    std::transform(s.begin(), s.end(), s.begin(), [](unsigned char c) { return (char) std::tolower(c); });
+    return s;
+}
+
+// LFM3-Audio: language code -> name
+static std::string lfm3_language_name(const std::string & language) {
    static const std::map<std::string, std::string> names = {
        { "ar", "Arabic" }, { "da", "Danish" }, { "de", "German" }, { "el", "Greek" }, { "en", "English" },
        { "es", "Spanish" }, { "fi", "Finnish" }, { "fr", "French" }, { "he", "Hebrew" }, { "it", "Italian" },
        { "ja", "Japanese" }, { "ko", "Korean" }, { "nl", "Dutch" }, { "pl", "Polish" }, { "pt", "Portuguese" },
        { "sv", "Swedish" }, { "tr", "Turkish" }, { "zh", "Chinese" },
    };
+    const auto it = names.find(string_lower(language));
+    return it != names.end() ? it->second : language;
+}
+
+// LFM3-Audio: no language = auto detect
+static std::string lfm3_asr_system_prompt(const std::string & language) {
    if (language.empty()) {
        return "Perform ASR.";
    }
-    std::string code = language;
-    std::transform(code.begin(), code.end(), code.begin(), [](unsigned char c) { return (char) std::tolower(c); });
-    const auto it = names.find(code);
-    return "Perform ASR in " + (it != names.end() ? it->second : language) + ".";
+    return "Perform ASR in " + lfm3_language_name(language) + ".";
 }

 common_chat_prompt_preset common_chat_get_asr_prompt(const common_chat_templates * chat_templates, const std::string & language) {
@@ -868,6 +876,27 @@ std::string common_chat_get_tts_prompt(const common_chat_templates * chat_templa
    });
 }

+std::string common_chat_get_speech_system_prompt(const common_chat_templates * chat_templates,
+                                                 const std::string & language, const std::string & voice) {
+    if (!chat_templates || !chat_templates->template_default || !is_lfm3_template(chat_templates->template_default->source())) {
+        return "";
+    }
+    // same voices as the "tts" template
+    const std::string v = string_lower(voice);
+    const bool has_voice = v == "female" || v == "male";
+    if (language.empty() && !has_voice) {
+        return "";
+    }
+    std::string out = "Speak";
+    if (!language.empty()) {
+        out += " in " + lfm3_language_name(language);
+    }
+    if (has_voice) {
+        out += " using a " + v + " voice";
+    }
+    return out + ".";
+}
+
 std::string common_chat_templates_source(const struct common_chat_templates * tmpls, const std::string & variant) {
    if (!variant.empty()) {
        if (variant == "tool_use") {
@@ -500,3 +500,7 @@ common_chat_prompt_preset common_chat_get_asr_prompt(const common_chat_templates
 // language and voice are optional, the template defines their defaults
 std::string common_chat_get_tts_prompt(const common_chat_templates * chat_templates, const std::string & text,
                                       const std::string & language = "", const std::string & voice = "");
+
+// system prompt line for the speech language and voice (e.g. LFM3-Audio), empty if not used
+std::string common_chat_get_speech_system_prompt(const common_chat_templates * chat_templates,
+                                                 const std::string & language, const std::string & voice);
@@ -860,6 +860,7 @@ public:
        std::unordered_set<std::string> required;
        std::vector<std::pair<std::string, const common_chat_schema *>> properties;
        std::map<std::string, size_t> enum_values;
+        std::vector<const common_chat_schema *> leaves;
        std::function<void(const common_chat_schema &, bool, int)> add_component = [&](const common_chat_schema & comp, bool is_required, int depth) {
            if (depth >= COMMON_RECURSIVE_LIMIT) {
                throw std::invalid_argument("JSON schema is nested more than " + std::to_string(COMMON_RECURSIVE_LIMIT) + " levels deep");
@@ -879,6 +880,10 @@ public:
                for (const auto & v : as<common_chat_schema_enum>(comp).values) {
                    enum_values[_generate_constant_rule(v)] += 1;
                }
+            } else if (is_required && comp.kind() != common_chat_schema::KIND_ANY) {
+                // Direct non-structural children are used for merging. anyOf alternatives are not merged.
+                // Empty schemas are skipped.
+                leaves.push_back(&comp);
            }
        };
        for (const auto & child : schema.children) {
@@ -901,6 +906,38 @@ public:
                return _add_rule(rule_name, "(" + string_join(enum_intersection, " | ") + ")");
            }
        }
+
+        if (properties.empty() && !leaves.empty()) {
+            if (leaves.size() == 1) {
+                return visit(*leaves[0], name);
+            }
+            bool all_numeric = true;
+            bool any_integer = false;
+            int64_t lo = std::numeric_limits<int64_t>::min();
+            int64_t hi = std::numeric_limits<int64_t>::max();
+            for (const auto * leaf : leaves) {
+                if (leaf->kind() == common_chat_schema::KIND_INTEGER) {
+                    const auto & i = as<common_chat_schema_integer>(*leaf);
+                    any_integer = true;
+                    lo = std::max(lo, i.minimum);
+                    hi = std::min(hi, i.maximum);
+                } else if (leaf->kind() != common_chat_schema::KIND_NUMBER) {
+                    all_numeric = false;
+                }
+            }
+            if (all_numeric) {
+                if (lo > hi) {
+                    throw std::invalid_argument("allOf has an empty numeric range");
+                }
+                if (any_integer) {
+                    common_chat_schema_integer merged;
+                    merged.minimum = lo;
+                    merged.maximum = hi;
+                    return visit(merged, name);
+                }
+                return visit(*leaves[0], name);
+            }
+        }
        return _add_rule(rule_name, _build_object_rule(properties, required, name, nullptr));
    }

@@ -300,7 +300,7 @@ common_chat_params common_chat_params_init_deepseek_v3_2(const common_chat_templ
    });

    if (include_grammar) {
-        data.grammar_lazy = has_tools && !require_tools;
+        data.grammar_lazy = !has_response_format && has_tools && !require_tools;
        data.grammar      = build_grammar([&](const common_grammar_builder & builder) {
            data.parser.build_grammar(builder, data.grammar_lazy);
        });
@@ -109,7 +109,7 @@ common_chat_params common_chat_params_init_ministral_3(const common_chat_templat
    });

    if (include_grammar) {
-        data.grammar_lazy = has_tools && inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_AUTO;
+        data.grammar_lazy = !has_response_format && has_tools && inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_AUTO;

        data.grammar = build_grammar([&](const common_grammar_builder & builder) {
            data.parser.build_grammar(builder, data.grammar_lazy);
@@ -191,7 +191,7 @@ common_chat_params common_chat_params_init_qwen3_coder(const common_chat_templat
    });

    if (include_grammar) {
-        data.grammar_lazy = has_tools && inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_AUTO;
+        data.grammar_lazy = !has_response_format && has_tools && inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_AUTO;

        data.grammar = build_grammar([&](const common_grammar_builder & builder) {
            data.parser.build_grammar(builder, data.grammar_lazy);
@@ -371,6 +371,7 @@ MMPROJ_MODEL_MAP: dict[str, str] = {
    "InternVisionModel": "internvl",
    "JanusForConditionalGeneration": "januspro",
    "KimiK25ForConditionalGeneration": "kimivl",
+    "KimiK31ForConditionalGeneration": "kimivl",
    "KimiVLForConditionalGeneration": "kimivl",
    "Lfm2AudioForConditionalGeneration": "lfm2",
    "D1OmniModel": "lfm2",
@@ -1,5 +1,6 @@
 from __future__ import annotations

+import re
 from typing import Callable, Iterable, TYPE_CHECKING

 import torch
@@ -198,6 +199,18 @@ class SarvamMoEModel(BailingMoeV2Model):
    #  - full rotary (no partial_rotary_factor)
    #  - expert bias is zero-mean normalized at load time

+    def set_vocab(self):
+        # BPE with byte fallback: mark <0xXX> tokens as BYTE so they decode to raw bytes
+        tokens, toktypes, tokpre = self.get_vocab_base()
+        toktypes = [gguf.TokenType.BYTE if re.fullmatch(r"<0x[0-9A-F]{2}>", tok) else typ for tok, typ in zip(tokens, toktypes)]
+        self.gguf_writer.add_tokenizer_model("gpt2")
+        self.gguf_writer.add_tokenizer_pre(tokpre)
+        self.gguf_writer.add_token_list(tokens)
+        self.gguf_writer.add_token_types(toktypes)
+
+        special_vocab = gguf.SpecialVocab(self.dir_model, load_merges=True)
+        special_vocab.add_to_gguf(self.gguf_writer)
+
    def set_gguf_parameters(self):
        super().set_gguf_parameters()
        hparams = self.hparams
@@ -7,6 +7,7 @@ import ast
 import logging
 import contextlib
 import json
+import math
 import os
 import re
 import sys
@@ -343,6 +344,36 @@ class ModelBase:
            logger.info(f"  + {scale_name} (per-expert scale, shape [{len(scales)}])")
            self.gguf_writer.add_tensor(scale_name, scale_vals)

+    def _prepare_kv_cache_scales(self):
+        for name in list(self.model_tensors.keys()):
+            if not name.endswith((".k_scale", ".v_scale")):
+                continue
+
+            mapped = self.tensor_map.get_type_and_name(key=name, try_suffixes=(".k_scale", ".v_scale"))
+            if mapped is None:
+                continue
+
+            tensor_type, new_name = mapped
+            if tensor_type not in (gguf.MODEL_TENSOR.ATTN_K, gguf.MODEL_TENSOR.ATTN_V, gguf.MODEL_TENSOR.ATTN_QKV):
+                continue
+
+            bid = int(new_name.split(".")[1])
+            scale_type = gguf.MODEL_TENSOR.ATTN_K if name.endswith(".k_scale") else gguf.MODEL_TENSOR.ATTN_V
+            suffix = ".k_scale" if name.endswith(".k_scale") else ".v_scale"
+            new_name = gguf.TENSOR_NAMES[scale_type].format(bid=bid) + suffix
+
+            scale = LazyTorchTensor.to_eager(self.model_tensors.pop(name)())
+            if scale.dtype not in (torch.bfloat16, torch.float32) or scale.numel() != 1:
+                raise ValueError(f"KV cache scale {name!r} must be a scalar BF16 or FP32 tensor")
+
+            scale = scale.float()
+            value = float(scale.item())
+            if not math.isfinite(value) or value == 0.0:
+                raise ValueError(f"KV cache scale {name!r} must be finite and nonzero")
+
+            logger.info(f"  + {new_name} (KV cache scale, shape [1])")
+            self.gguf_writer.add_tensor(new_name, scale.flatten().numpy())
+
    def dequant_model(self):
        # If all quantized tensors were already handled (e.g. pure NVFP4), skip
        if self._is_nvfp4 and not any(k.endswith((".weight_scale", ".weight_scale_inv", ".input_scale", ".activation_scale", "_activation_scale", ".k_scale", ".v_scale")) for k in self.model_tensors):
@@ -1455,6 +1486,8 @@ class ModelBase:
            unmatched = layer_algo.keys() - matched
            assert not unmatched, f"quantized_layers entries not in the model tensors: {sorted(unmatched)}"

+        self._prepare_kv_cache_scales()
+
        # NVFP4 weights are repacked and written directly to gguf_writer.
        # This must run before dequant_model so NVFP4 tensors are removed
        # from model_tensors, leaving only non-NVFP4 (e.g. FP8) for dequant.
@@ -85,6 +85,16 @@ class GptOssModel(TextModel):
    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
        name, gen = item

+        # Normalize erroneous GPT-OSS MXFP4 tensor naming.
+        # https://github.com/p-e-w/heretic/pull/211#issuecomment-5094370390
+        if "mlp.experts." in name:
+            for proj in ("down_proj", "gate_up_proj"):
+                name = name.replace(
+                    f"{proj}_blocks_blocks", f"{proj}_blocks"
+                ).replace(
+                    f"{proj}_blocks_scales", f"{proj}_scales"
+                )
+
        if "sinks" in name:
            name += ".weight"

@@ -57,6 +57,8 @@ class KimiVLModel(MmprojModel):
 class KimiK25Model(MmprojModel):
    """Kimi-K2.5 with MoonViT3d vision encoder"""

+    projector_type = gguf.VisionProjectorType.KIMIK25
+
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

@@ -78,7 +80,7 @@ class KimiK25Model(MmprojModel):
        super().set_gguf_parameters()
        assert self.hparams_vision is not None

-        self.gguf_writer.add_clip_projector_type(gguf.VisionProjectorType.KIMIK25)
+        self.gguf_writer.add_clip_projector_type(self.projector_type)

        # Position embedding parameters (for interpolation)
        self.gguf_writer.add_uint32("vision.pos_emb_height", self.hparams_vision.get("init_pos_emb_height", 64))
@@ -87,7 +89,11 @@ class KimiK25Model(MmprojModel):

        # Projector parameters
        self.gguf_writer.add_vision_use_gelu(self.hparams_vision.get("projector_hidden_act", "gelu") == "gelu")
-        self.gguf_writer.add_vision_attention_layernorm_eps(self.hparams_vision.get("projector_ln_eps", 1e-5))
+        if self.hparams_vision.get("norm_type", "layernorm") == "rmsnorm":
+            # nn.RMSNorm without eps uses finfo(float32).eps, also for bf16 inputs
+            self.gguf_writer.add_vision_attention_layernorm_eps(torch.finfo(torch.float32).eps)
+        else:
+            self.gguf_writer.add_vision_attention_layernorm_eps(self.hparams_vision.get("projector_ln_eps", 1e-5))
        self.gguf_writer.add_vision_projector_scale_factor(self.merge_kernel_size[0])

        # Image size limits
@@ -120,7 +126,7 @@ class KimiK25Model(MmprojModel):

    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
        assert self.hparams_vision is not None
-        n_head = self.hparams_vision.get("num_attention_heads", 16)
+        n_head = self.find_vparam(["num_attention_heads", "vt_num_attention_heads"], optional=True) or 16

        # Permute Q/K weights/biases from interleaved to split RoPE format
        # This allows using build_rope_2d at runtime without post-permutation.
@@ -171,3 +177,19 @@ class Glm5vModel(KimiK25Model):
            name = name.replace("mm_projector.linear_", "mm_projector.proj.linear_", 1)

        yield from super().modify_tensors(data_torch, name, bid)
+
+
+@ModelBase.register("KimiK31ForConditionalGeneration")
+@ModelBase.example("moonshotai/Kimi-K3.1-Pro")
+class KimiK31VisionModel(KimiK25Model):
+    """Kimi-K2.5 MoonViT3d with RMSNorm, a wider QKV and the patchmergerv2 projector"""
+
+    projector_type = gguf.VisionProjectorType.KIMIK31
+
+    def set_gguf_parameters(self):
+        super().set_gguf_parameters()
+        assert self.hparams_vision is not None
+        assert self.hparams_vision["norm_type"] == "rmsnorm"
+        assert self.hparams_vision["mm_projector_type"] == "patchmergerv2"
+        n_head = self.find_vparam(["vt_num_attention_heads"])
+        self.gguf_writer.add_vision_head_dim(self.hparams_vision["qkv_hidden_size"] // n_head)
@@ -52,6 +52,9 @@ The packages for FP32 and FP16 would have different accuracy and performance on

 ## News

+- 2026.11
+  - Update the CI and release build for oneAPI 2026.2 and Level Zero SDK 1.34.0. oneDNN is updated to the 2026.1 stand-alone package.
+
 - 2026.09
  - Update the CI build environment for oneAPI 2026.1 (unified oneAPI Toolkit). oneDNN is removed from the Deep Learning Essentials package in 2026.0, so the CI now uses the oneAPI Toolkit installer which still includes oneDNN.
  - oneAPI 2026.1 improves the SYCL build performance: measured with the same code on Arc B570, prompt processing 1331 vs 434 t/s (3.1x) vs the 2025.3-based release build.
@@ -285,6 +288,7 @@ Upon a successful installation, SYCL is enabled for the available Intel devices,

 |Verified release|
 |-|
+|2026.2 |
 |2026.1 |
 |2025.3.3 |
 |2025.2.1|
@@ -825,7 +829,8 @@ User can use the device management in [docs/multi-gpu.md](https://github.com/ggm
 | GGML_SYCL_ENABLE_XMX_ESIMD | 0 or 1 (default) | Use the XMX (DPAS) multi-column mat-vec kernels for reordered Q2_K to Q6_K and Q8_0 weights on Xe2 GPUs. Speeds up speculative decoding and short prompt batches. |
 | GGML_SYCL_XMX_MAX_COLS | non-negative integer, 80 (default) | Most activation columns (batch size) the `GGML_SYCL_ENABLE_XMX_ESIMD` kernels take; wider batches use the other mat-mul kernels. A larger value keeps XMX for bigger batches (as repeated launches), `0` never uses it. |
 | GGML_SYCL_XMX_GLU_MAX_COLS | 0 to 16, 16 (default) | Most columns the fused gate and up XMX kernel takes; wider batches run gate and up as separate mat-vecs. 16 is the most the kernel is built for, larger values are clamped to 16 and `0` turns the fused kernel off. |
-| GGML_SYCL_XMX_GATHER_TYPES | decimal bitmask, all bits set (default) | Weight formats that may use the XMX dequant-GEMM paths, which dequantize weights straight into the XMX tiles. This speeds up prompt processing of MoE models on GPUs with XMX units (Arc A- and B-series, Arc Pro, Data Center GPU Max), for example pp512 of Qwen3-30B-A3B UD-IQ3_XXS by about 50% on an Arc Pro B60. Bits:<br>* 1: IQ4_NL, 2: IQ3_S, 4: IQ4_XS, 8: IQ3_XXS, 16: IQ2_XXS, 32: IQ2_XS, 64: IQ2_S, 128: IQ1_S, 256: IQ1_M<br>* 512: Q8_0, 1024: Q4_K, 2048: Q5_K, 4096: Q6_K (MoE `MUL_MAT_ID` only)<br>Add values to combine them, for example `3` for IQ4_NL and IQ3_S; `0` disables the paths. A set bit does not force the path: batches of more than 64 tokens per expert or row lengths that are not a multiple of 256 (32 for IQ4_NL and Q8_0) use the library GEMM. |
+| GGML_SYCL_XMX_GATHER_TYPES | decimal bitmask, all bits set (default) | Weight formats that may use the XMX dequant-GEMM paths, which dequantize weights straight into the XMX tiles. This speeds up prompt processing of MoE models on GPUs with XMX units (Arc A- and B-series, Arc Pro, Data Center GPU Max), for example pp512 of Qwen3-30B-A3B UD-IQ3_XXS by about 50% on an Arc Pro B60. Bits:<br>* 1: IQ4_NL, 2: IQ3_S, 4: IQ4_XS, 8: IQ3_XXS, 16: IQ2_XXS, 32: IQ2_XS, 64: IQ2_S, 128: IQ1_S, 256: IQ1_M<br>* 512: Q8_0, 1024: Q4_K, 2048: Q5_K, 4096: Q6_K, 8192: Q2_K, 16384: Q3_K (MoE `MUL_MAT_ID` only)<br>Add values to combine them, for example `3` for IQ4_NL and IQ3_S; `0` disables the paths. A set bit does not force the path: batches of more than 64 tokens per expert use the library GEMM unless the device schedule width heuristic is overridden with `GGML_SYCL_MMID_SCHED=3`. Row lengths must be a multiple of 256 (32 for IQ4_NL and Q8_0). |
+| GGML_SYCL_MMID_SCHED | decimal bitmask, 1 (default) | Build the grouped MoE routing and tile schedule on the GPU, avoiding the expert-ID host readback and CPU sort.<br>* 1: enable the device schedule.<br>* 2: override the fallback for models with few routed experts by disabling the performance heuristic that rejects batches averaging more than 64 routed rows per expert. This may reduce performance; type, shape, precision and device capability checks still apply.<br>Use `3` to enable the device schedule and override this heuristic, or `0` to use the host schedule. Bit 2 alone does not enable the device schedule. Higher bits are ignored. |
 | GGML_SYCL_XMX_GATHER_SHAPES | decimal bitmask, 255 (default) | XMX `joint_matrix` combinations the paths of `GGML_SYCL_XMX_GATHER_TYPES` may use; the operand type comes from `GGML_SYCL_DYNAMIC_PRECISION` and the best supported combination is picked automatically (logged as `fg_pick_combo`). Bits:<br>* Xe2, Xe3, Xe-HPC: 1: f16 8x16x16, 2: f16 16x16x16, 4: f16 32x64x16, 8: f16 32x64x32, 32: tf32 8x16x8, 64: bf16 8x16x16<br>* Xe-HPG (Arc A770, ARL-H): 16: f16 8x8x16, 128: bf16 8x8x16<br>Clear a bit to exclude a combination, or set a single bit to force one for testing. |
 | GGML_SYCL_DYNAMIC_PRECISION | `F16` (default with `GGML_SYCL_F16=ON`), `BF16`, `TF32` or `F32` (default otherwise) | Operand type of the XMX dequant-GEMM paths (`GGML_SYCL_XMX_GATHER_TYPES`); accumulation is always f32. `F16` is the fastest, but activations above 65504 overflow. `BF16` keeps the f32 range at a 7-bit mantissa, `TF32` keeps the range and the f16 mantissa but is about 30% slower and needs Xe2, Xe3 or Xe-HPC, and `F32` turns the XMX paths off. Ops that request a higher src1 precision ([TAG_GGML_PREC]) get it regardless of this setting. |
 | GGML_SYCL_DYNAMIC_REQUIRED_PRECISION | `F32` (default), `TF32`, `BF16` or `F16` | Lowest type the XMX paths may use for an op that requests an F32 src1, such as Mistral 4 `ffn_down_exps`. The default runs such ops on the library f32 GEMM; `TF32` or `BF16` trade mantissa for speed while keeping the f32 range. `F16` ignores the request and can overflow; it is meant for testing only. |
@@ -131,15 +131,9 @@ endif()
 # POSIX conformance
 #

-# clock_gettime came in POSIX.1b (1993)
-# CLOCK_MONOTONIC came in POSIX.1-2001 / SUSv3 as optional
-# posix_memalign came in POSIX.1-2001 / SUSv3
-# M_PI is an XSI extension since POSIX.1-2001 / SUSv3, came in XPG1 (1985)
-
-# Somehow in OpenBSD whenever POSIX conformance is specified
-# some string functions rely on locale_t availability,
-# which was introduced in POSIX.1-2008, forcing us to go higher
-if (CMAKE_SYSTEM_NAME MATCHES "OpenBSD")
+# POSIX.1-2001 with XSI extensions provides clock_gettime, CLOCK_MONOTONIC, posix_memalign and M_PI.
+# OpenBSD needs POSIX.1-2008 for locale_t. FreeBSD needs it for F_DUPFD_CLOEXEC and strnlen.
+if (CMAKE_SYSTEM_NAME MATCHES "OpenBSD|FreeBSD")
    add_compile_definitions(_XOPEN_SOURCE=700)
 elseif (CMAKE_SYSTEM_NAME MATCHES "AIX")
    # Don't define _XOPEN_SOURCE.  We need _ALL_SOURCE, which is the default,
@@ -1398,34 +1398,35 @@ static void ggml_backend_cann_buffer_get_tensor(ggml_backend_buffer_t buffer,
 static bool ggml_backend_cann_buffer_cpy_tensor(ggml_backend_buffer_t buffer,
                                                const ggml_tensor *   src,
                                                ggml_tensor *         dst) {
-    if (ggml_backend_buft_is_cann(src->buffer->buft)) {
-        ggml_backend_cann_buffer_context * src_ctx = (ggml_backend_cann_buffer_context *) src->buffer->context;
-        ggml_backend_cann_buffer_context * dst_ctx = (ggml_backend_cann_buffer_context *) buffer->context;
+    if (!ggml_backend_buft_is_cann(src->buffer->buft)) {
+        return false;
+    }

-        size_t memcpy_size = ggml_nbytes(src);
-        // Same device.
-        if (src_ctx->device == dst_ctx->device) {
-            ACL_CHECK(aclrtMemcpy((char *) dst->data, memcpy_size, (const char *) src->data, memcpy_size,
-                                  ACL_MEMCPY_DEVICE_TO_DEVICE));
-            return true;
-        } else {
+    ggml_backend_cann_buffer_context * src_ctx = (ggml_backend_cann_buffer_context *) src->buffer->context;
+    ggml_backend_cann_buffer_context * dst_ctx = (ggml_backend_cann_buffer_context *) buffer->context;
+
+    size_t memcpy_size = ggml_nbytes(src);
+    // Same device.
+    if (src_ctx->device == dst_ctx->device) {
+        ACL_CHECK(aclrtMemcpy((char *) dst->data, memcpy_size, (const char *) src->data, memcpy_size,
+                                ACL_MEMCPY_DEVICE_TO_DEVICE));
+        return true;
+    } else {
 #ifdef ASCEND_310P
-            // TODO: Support 310p P2P copy
-            return false;
+        // TODO: Support 310p P2P copy
+        return false;
 #endif
-            // Different device but can access by peer.
-            int32_t canAccessPeer = 0;
-            ACL_CHECK(aclrtDeviceCanAccessPeer(&canAccessPeer, src_ctx->device, dst_ctx->device));
-            if (canAccessPeer) {
-                ggml_cann_set_device(src_ctx->device);
-                ACL_CHECK(aclrtDeviceEnablePeerAccess(dst_ctx->device, 0));
-                ACL_CHECK(aclrtMemcpy((char *) dst->data, memcpy_size, (const char *) src->data, memcpy_size,
-                                      ACL_MEMCPY_DEVICE_TO_DEVICE));
-                return true;
-            }
+        // Different device but can access by peer.
+        int32_t canAccessPeer = 0;
+        ACL_CHECK(aclrtDeviceCanAccessPeer(&canAccessPeer, src_ctx->device, dst_ctx->device));
+        if (canAccessPeer) {
+            ggml_cann_set_device(src_ctx->device);
+            ACL_CHECK(aclrtDeviceEnablePeerAccess(dst_ctx->device, 0));
+            ACL_CHECK(aclrtMemcpy((char *) dst->data, memcpy_size, (const char *) src->data, memcpy_size,
+                                    ACL_MEMCPY_DEVICE_TO_DEVICE));
+            return true;
        }
    }
-    return false;
 }

 /**
@@ -2446,6 +2447,9 @@ static bool ggml_backend_cann_supports_op(ggml_backend_dev_t dev, const ggml_ten
            break;
        case GGML_OP_MUL_MAT:
            {
+                if (ggml_prec_get_src(op, 1) == GGML_PREC_BF16 || ggml_prec_get_src(op, 1) == GGML_PREC_F32) {
+                    return false;
+                }
                switch (op->src[0]->type) {
 #ifndef ASCEND_310P
                    case GGML_TYPE_BF16:
@@ -2864,12 +2868,13 @@ static ggml_backend_t ggml_backend_cann_device_init(ggml_backend_dev_t dev, cons
 *              otherwise false.
 */
 static bool ggml_backend_cann_supports_buft(ggml_backend_dev_t dev, ggml_backend_buffer_type_t buft) {
-    if (ggml_backend_buft_is_cann(buft)) {
-        ggml_backend_cann_device_context *      dev_ctx  = (ggml_backend_cann_device_context *) dev->context;
-        ggml_backend_cann_buffer_type_context * buft_ctx = (ggml_backend_cann_buffer_type_context *) buft->context;
-        return buft_ctx->device == dev_ctx->device;
+    if (!ggml_backend_buft_is_cann(buft)) {
+        return false;
    }
-    return false;
+
+    ggml_backend_cann_device_context *      dev_ctx  = (ggml_backend_cann_device_context *) dev->context;
+    ggml_backend_cann_buffer_type_context * buft_ctx = (ggml_backend_cann_buffer_type_context *) buft->context;
+    return buft_ctx->device == dev_ctx->device;
 }

 static ggml_backend_buffer_type_t ggml_backend_cann_device_get_buffer_type(ggml_backend_dev_t dev) {
@@ -9569,8 +9569,8 @@ static void ggml_compute_forward_flash_attn_ext_f16(
 void ggml_compute_forward_flash_attn_ext(
        const ggml_compute_params * params,
        ggml_tensor * dst) {
-    switch (dst->op_params[3]) {
-        case GGML_PREC_DEFAULT:
+    switch (ggml_prec_get_acc(dst)) {
+        case GGML_PREC_UNDEFINED:
        case GGML_PREC_F32:
        case GGML_PREC_BF16:
            {
@@ -1180,7 +1180,7 @@ class tensor_traits_common : public tensor_traits_base {
        const int64_t DK = nek0;
        const int64_t DV = nev0;

-        const bool supported_prec  = (dst->op_params[3] == GGML_PREC_F32 || dst->op_params[3] == GGML_PREC_BF16 || dst->op_params[3] == GGML_PREC_DEFAULT);
+        const bool supported_prec  = (ggml_prec_get_acc(dst) == GGML_PREC_F32 || ggml_prec_get_acc(dst) == GGML_PREC_BF16 || ggml_prec_get_acc(dst) == GGML_PREC_UNDEFINED);
        const bool supported_types = (q->type == GGML_TYPE_F32 && k->type == GGML_TYPE_F16 && v->type == GGML_TYPE_F16);
        const bool supported_shape = (DK > 0 && DK <= 128 && DV > 0 && DV <= 128);
        const bool supported_vlen  = (__riscv_vlenb() == 128);
@@ -7,7 +7,6 @@
 #include <cstdint>
 #include <cstdlib>
 #include <memory>
-#include <mutex>

 #if defined(GGML_USE_HIP)
 #define GGML_COMMON_DECL_HIP
@@ -1271,186 +1270,9 @@ struct ggml_tensor_extra_gpu {
 #define USE_CUDA_GRAPH
 #endif

-struct ggml_cuda_graph {
-#ifdef USE_CUDA_GRAPH
-    ~ggml_cuda_graph() {
-        for (cudaGraphExec_t instance : instances) {
-            CUDA_CHECK(cudaGraphExecDestroy(instance));
-        }
-    }
-    // the ggml graph is captured as a series of CUDA graphs, so the GPU can run one chunk while the next is enqueued
-    std::vector<cudaGraphExec_t> instances;
-    size_t num_nodes = 0;
-    std::vector<cudaGraphNode_t> nodes;
-    bool disable_due_to_gpu_arch = false;
-    bool warmup_complete = false;
-    uint64_t uid = 0;
-    int64_t last_used_time = 0;
-    struct node_properties {
-        ggml_tensor node;
-        void *   node_src_data_ptrs[GGML_MAX_SRC];
-        int64_t  node_src_ne[GGML_MAX_SRC][GGML_MAX_DIMS];
-        size_t   node_src_nb[GGML_MAX_SRC][GGML_MAX_DIMS];
-    };
-    std::vector<node_properties> node_props;
-
-    bool is_enabled() const {
-        static const bool disable_cuda_graphs_due_to_env = (getenv("GGML_CUDA_DISABLE_GRAPHS") != nullptr);
-        return !(disable_due_to_gpu_arch || disable_cuda_graphs_due_to_env);
-    }
-#endif
-};
-
-struct ggml_cuda_concurrent_event {
-    std::vector<cudaEvent_t> join_events;
-    cudaEvent_t              fork_event = nullptr;
-
-    int                                          n_streams = 0;
-    std::unordered_map<const ggml_tensor *, int> stream_mapping;
-
-    // Original order of nodes in this concurrent region (before interleaving)
-    // Used to restore grouping for fusion within streams
-    std::vector<const ggml_tensor *> original_order;
-
-    const ggml_tensor * join_node;
-
-    ggml_cuda_concurrent_event() = default;
-
-    ggml_cuda_concurrent_event(const ggml_cuda_concurrent_event &) = delete;
-    ggml_cuda_concurrent_event & operator=(const ggml_cuda_concurrent_event &) = delete;
-
-    explicit ggml_cuda_concurrent_event(int n_streams) : n_streams(n_streams) {
-        join_events.resize(n_streams);
-
-        for (size_t i = 0; i < join_events.size(); ++i) {
-            CUDA_CHECK(cudaEventCreateWithFlags(&join_events[i], cudaEventDisableTiming));
-        }
-
-        CUDA_CHECK(cudaEventCreateWithFlags(&fork_event, cudaEventDisableTiming));
-    }
-
-    ggml_cuda_concurrent_event(ggml_cuda_concurrent_event && other) noexcept
-    : join_events(std::move(other.join_events))
-    , fork_event(other.fork_event)
-    , n_streams(other.n_streams)
-    , stream_mapping(std::move(other.stream_mapping))
-    , original_order(std::move(other.original_order))
-    , join_node(other.join_node) {
-        other.fork_event = nullptr;
-    }
-
-    // 1. check if any branches write to overlapping memory ranges (except the join node)
-    // 2. check whether all srcs are either within the branch or outside the nodes covered by ggml_cuda_concurrent_event
-    // we assume all nodes have the same buffer
-    bool is_valid() const {
-        std::vector<std::vector<std::pair<int64_t, int64_t>>> write_ranges;
-        write_ranges.resize(n_streams);
-
-        // get join_node's memory range to exclude from overlap checking.
-        // multiple nodes can use join_node's buffer; we synchronize on the join node.
-        const ggml_tensor * join_t     = join_node->view_src ? join_node->view_src : join_node;
-        const int64_t       join_start = (int64_t) join_t->data;
-        const int64_t       join_end   = join_start + ggml_nbytes(join_t);
-
-        for (const auto & [tensor, stream] : stream_mapping) {
-            const ggml_tensor * t = tensor->view_src ? tensor->view_src : tensor;
-            const int64_t       t_start = (int64_t) t->data;
-            const int64_t       t_end   = t_start + ggml_nbytes(t);
-
-            // skip tensors that overlap with join_node's buffer.
-            if ((t_start <= join_start && join_start < t_end) || (join_start <= t_start && t_start < join_end)) {
-                continue;
-            }
-
-            // concurrent streams begin from 1
-            write_ranges[stream - 1].emplace_back(t_start, t_end);
-        }
-
-        for (int i = 0; i < n_streams; ++i) {
-            // sorts first by start then by end of write range
-            std::sort(write_ranges[i].begin(), write_ranges[i].end());
-        }
-
-        bool writes_overlap = false;
-        bool dependent_srcs = false;
-        for (const auto & [tensor, stream] : stream_mapping) {
-            const ggml_tensor * t = tensor->view_src ? tensor->view_src : tensor;
-            const int64_t       t_start = (int64_t) t->data;
-            const int64_t       t_end   = t_start + ggml_nbytes(t);
-
-            // skip tensors that overlap with join_node's buffer
-            if ((t_start <= join_start && join_start < t_end) || (join_start <= t_start && t_start < join_end)) {
-                continue;
-            }
-
-            // check if this buffer's write data overlaps with another stream's
-            std::pair<int64_t, int64_t> data_range = std::make_pair(t_start, t_end);
-            for (int i = 0; i < n_streams; ++i) {
-                if (i == stream - 1) {
-                    continue;
-                }
-                auto it = std::lower_bound(write_ranges[i].begin(), write_ranges[i].end(), data_range);
-
-                if (it != write_ranges[i].end()) {
-                    const std::pair<int64_t, int64_t> & other = *it;
-
-                    // std::lower_bound returns the first element where other >= data_range (lexicographically).
-                    // This guarantees other.first >= data_range.first.
-                    // Therefore, overlap occurs iff other.first < data_range.second
-                    // (i.e., the other range starts before this range ends).
-                    if (other.first < data_range.second) {
-                        GGML_LOG_DEBUG("Writes overlap for %s", tensor->name);
-                        writes_overlap = true;
-                        break;
-                    }
-                }
-            }
-
-            //check if all srcs are either in branch or don't have a branch
-            for (int i = 0; i < GGML_MAX_SRC; ++i) {
-                if (!tensor->src[i]) {
-                    continue;
-                }
-
-                auto it = stream_mapping.find(tensor->src[i]);
-
-                if (it == stream_mapping.end()) {
-                    continue;
-                }
-
-                if (it->second != stream) {
-                    dependent_srcs = true;
-                    break;
-                }
-            }
-
-            if (dependent_srcs || writes_overlap) {
-                break;
-            }
-        }
-
-        return !writes_overlap && !dependent_srcs;
-    }
-
-    ~ggml_cuda_concurrent_event() {
-        if (fork_event != nullptr) {
-            CUDA_CHECK(cudaEventDestroy(fork_event));
-        }
-        for (cudaEvent_t e : join_events) {
-            if (e != nullptr) {
-                CUDA_CHECK(cudaEventDestroy(e));
-            }
-        }
-    }
-};
-
-struct ggml_cuda_stream_context {
-    std::unordered_map<const ggml_tensor *, ggml_cuda_concurrent_event> concurrent_events;
-
-    void reset() {
-        concurrent_events.clear();
-    }
-};
+// defined in graph.cuh
+struct ggml_cuda_graph;
+struct ggml_cuda_stream_context;

 struct ggml_backend_cuda_context {
    int device;
@@ -1472,61 +1294,20 @@ struct ggml_backend_cuda_context {
    int64_t next_graph_eviction_time = 0;
    int64_t last_graph_use_time = 0;

-    ggml_cuda_graph * cuda_graph(const void * first_node_ptr, int64_t time_now = ggml_time_us()) {
-        // Give all cached graphs 10s to be reused after an idle period, without updating each graph.
-        if (!cuda_graphs.empty() && time_now - last_graph_use_time >= 10'000'000) {
-            next_graph_eviction_time = time_now + 10'000'000;
-        }
-        last_graph_use_time = time_now;
-
-        // Sweep every 5s during active use, evicting graphs unused for at least 10s.
-        if (time_now >= next_graph_eviction_time) {
-            next_graph_eviction_time = time_now + 5'000'000;
-            for (auto it = cuda_graphs.begin(); it != cuda_graphs.end(); ) {
-                if (it->first != first_node_ptr && time_now - it->second->last_used_time >= 10'000'000) {
-                    it = cuda_graphs.erase(it);
-                } else {
-                    ++it;
-                }
-            }
-        }
-
-        auto it = cuda_graphs.find(first_node_ptr);
-        if (it == cuda_graphs.end()) {
-            it = cuda_graphs.emplace(first_node_ptr, std::make_unique<ggml_cuda_graph>()).first;
-        }
-        it->second->last_used_time = time_now;
-        return it->second.get();
-    }
+    // defined in ggml-cuda.cu, which includes graph.cuh
+    ggml_cuda_graph * cuda_graph(const void * first_node_ptr, int64_t time_now = ggml_time_us());

    // Check if any CUDA graph is enabled for this context (used by kernels that need to know
    // if graphs are in use without having access to the specific graph key)
-    bool any_cuda_graph_enabled() const {
-        for (const auto & [key, graph] : cuda_graphs) {
-            if (graph && graph->is_enabled()) {
-                return true;
-            }
-        }
-        return false;
-    }
+    bool any_cuda_graph_enabled() const;

    // Check if any CUDA graph has an instance for this context
-    bool any_cuda_graph_has_instance() const {
-        for (const auto & [key, graph] : cuda_graphs) {
-            if (graph && !graph->instances.empty()) {
-                return true;
-            }
-        }
-        return false;
-    }
+    bool any_cuda_graph_has_instance() const;
 #endif // USE_CUDA_GRAPH

-    explicit ggml_backend_cuda_context(int device) :
-        device(device),
-        name(GGML_CUDA_NAME + std::to_string(device)) {
-    }
+    explicit ggml_backend_cuda_context(int device);

-    ggml_cuda_stream_context concurrent_stream_context;
+    std::unique_ptr<ggml_cuda_stream_context> concurrent_stream_context;

    ~ggml_backend_cuda_context();

@@ -1540,7 +1321,7 @@ struct ggml_backend_cuda_context {

    cudaStream_t stream() { return stream(device, curr_stream_no); }

-    ggml_cuda_stream_context & stream_context() { return concurrent_stream_context; }
+    ggml_cuda_stream_context & stream_context() { return *concurrent_stream_context; }

    cublasHandle_t cublas_handle() {
        if (cublas_handles[device][curr_stream_no] == nullptr) {
@@ -1644,60 +1425,8 @@ struct ggml_cuda_pdl_config {

 };

-static bool ggml_cuda_kernel_can_use_pdl(const void * kernel) {
-    const int device = ggml_cuda_get_device();
-
-    struct cache_key {
-        int          device;
-        const void * kernel;
-
-        bool operator==(const cache_key & other) const { return device == other.device && kernel == other.kernel; }
-    };
-
-    struct cache_key_hash {
-        // MurmurHash3 mixing function for better hash distribution (vs. just std::hash which in some implementations simply returns the identity)
-        static size_t hash_mix(size_t x) {
-            std::uint64_t       y = x;
-            const std::uint64_t m = 0xe9846af9b1a615d;
-
-            y ^= y >> 32;
-            y *= m;
-            y ^= y >> 32;
-            y *= m;
-            y ^= y >> 28;
-
-            return static_cast<size_t>(y);
-        }
-
-        size_t operator()(const cache_key & key) const {
-            // Use a nonzero seed to avoid mapping all-zero keys to zero
-            size_t h = 42;
-            h        = hash_mix(h + key.device);
-            h        = hash_mix(h + reinterpret_cast<size_t>(key.kernel));
-            return h;
-        }
-    };
-
-    static std::mutex                                          cache_mutex;
-    static std::unordered_map<cache_key, bool, cache_key_hash> cache;
-
-    const cache_key             key = { device, kernel };
-    std::lock_guard<std::mutex> lock(cache_mutex);
-    const auto                  it = cache.find(key);
-    if (it != cache.end()) {
-        return it->second;
-    }
-
-    cudaFuncAttributes attr = {};
-    CUDA_CHECK(cudaFuncGetAttributes(&attr, kernel));
-
-    // PDL device-side primitives are emitted only for PTX versions >= 90.
-    // We have to guard on a loaded kernel's PTX version so a kernel forward-JIT'ed
-    // from pre-Hopper PTX to a Hopper-or-newer GPU does not opt into PDL.
-    const bool can_use_pdl = attr.ptxVersion >= 90;
-    cache.emplace(key, can_use_pdl);
-    return can_use_pdl;
-}
+// whether kernel can be launched with PDL on the current device (also honors GGML_CUDA_PDL=0), defined in ggml-cuda.cu
+bool ggml_cuda_kernel_can_use_pdl(const void * kernel);

 #endif //defined(GGML_CUDA_USE_PDL)

@@ -1711,13 +1440,7 @@ static bool ggml_cuda_kernel_can_use_pdl(const void * kernel) {
 template<typename Kernel, typename... Args>
 static __inline__ void ggml_cuda_kernel_launch(Kernel kernel, const ggml_cuda_kernel_launch_params & launch_params, Args&&... args) {
 #if defined(GGML_CUDA_USE_PDL)
-
-    static const bool env_pdl_enabled = []() {
-        const char * env = getenv("GGML_CUDA_PDL");
-        return env == nullptr || std::atoi(env) != 0;
-    }();
-
-    if (env_pdl_enabled && ggml_cuda_kernel_can_use_pdl(reinterpret_cast<const void *>(kernel))) {
+    if (ggml_cuda_kernel_can_use_pdl(reinterpret_cast<const void *>(kernel))) {
        auto pdl_cfg = ggml_cuda_pdl_config(launch_params);

        CUDA_CHECK(cudaLaunchKernelEx(&pdl_cfg.cfg, kernel, std::forward<Args>(args)... ));
@@ -4,6 +4,7 @@

 #include "ggml-cuda/allreduce.cuh"
 #include "ggml-cuda/common.cuh"
+#include "ggml-cuda/graph.cuh"
 #include "ggml-cuda/acc.cuh"
 #include "ggml-cuda/add-id.cuh"
 #include "ggml-cuda/arange.cuh"
@@ -583,6 +584,7 @@ struct ggml_cuda_pool_vmm : public ggml_cuda_pool {
    CUdeviceptr pool_addr = 0;
    size_t pool_used = 0;
    size_t pool_size = 0;
+    size_t max_size = CUDA_POOL_VMM_MAX_SIZE;
    size_t granularity;
 #if defined(GGML_USE_HIP)
    std::vector<std::pair<CUdeviceptr, size_t>> mappings;
@@ -604,7 +606,7 @@ struct ggml_cuda_pool_vmm : public ggml_cuda_pool {
 #else
            CU_CHECK(cuMemUnmap(pool_addr, pool_size));
 #endif
-            CU_CHECK(cuMemAddressFree(pool_addr, CUDA_POOL_VMM_MAX_SIZE));
+            CU_CHECK(cuMemAddressFree(pool_addr, max_size));
        }
    }

@@ -620,7 +622,21 @@ struct ggml_cuda_pool_vmm : public ggml_cuda_pool {
            size_t reserve_size = size - avail;
            reserve_size = granularity * ((reserve_size + granularity - 1) / granularity);

-            GGML_ASSERT(pool_size + reserve_size <= CUDA_POOL_VMM_MAX_SIZE);
+            // reserve virtual address space (if not already reserved)
+            if (pool_addr == 0) {
+                CUresult err = cuMemAddressReserve(&pool_addr, max_size, 0, 0, 0);
+                while (err == CUDA_ERROR_OUT_OF_MEMORY && (max_size >> 1) >= reserve_size && (max_size >> 1) >= granularity) {
+                    max_size >>= 1;
+                    err = cuMemAddressReserve(&pool_addr, max_size, 0, 0, 0);
+                }
+                CU_CHECK(err);
+                if (max_size < CUDA_POOL_VMM_MAX_SIZE) {
+                    GGML_LOG_WARN(GGML_CUDA_NAME " pool[%d]: reserved %.2f GiB of virtual memory instead of %.2f GiB\n",
+                                  device, max_size / (1024.0 * 1024.0 * 1024.0), CUDA_POOL_VMM_MAX_SIZE / (1024.0 * 1024.0 * 1024.0));
+                }
+            }
+
+            GGML_ASSERT(pool_size + reserve_size <= max_size);

            // allocate more physical memory
            CUmemAllocationProp prop = {};
@@ -630,11 +646,6 @@ struct ggml_cuda_pool_vmm : public ggml_cuda_pool {
            CUmemGenericAllocationHandle handle;
            CU_CHECK(cuMemCreate(&handle, reserve_size, &prop, 0));

-            // reserve virtual address space (if not already reserved)
-            if (pool_addr == 0) {
-                CU_CHECK(cuMemAddressReserve(&pool_addr, CUDA_POOL_VMM_MAX_SIZE, 0, 0, 0));
-            }
-
            // map at the end of the pool
            CUdeviceptr start_ptr = (CUdeviceptr)((char *)(pool_addr) + pool_size);
            CU_CHECK(cuMemMap(start_ptr, reserve_size, 0, handle, 0));
@@ -735,6 +746,124 @@ std::unique_ptr<ggml_cuda_pool> ggml_backend_cuda_context::new_pool_for_device(g
    return std::unique_ptr<ggml_cuda_pool>(new ggml_cuda_pool_leg(ctx, device));
 }

+ggml_backend_cuda_context::ggml_backend_cuda_context(int device) :
+    device(device),
+    name(GGML_CUDA_NAME + std::to_string(device)),
+    concurrent_stream_context(std::make_unique<ggml_cuda_stream_context>()) {
+}
+
+#ifdef USE_CUDA_GRAPH
+ggml_cuda_graph * ggml_backend_cuda_context::cuda_graph(const void * first_node_ptr, int64_t time_now) {
+    // Give all cached graphs 10s to be reused after an idle period, without updating each graph.
+    if (!cuda_graphs.empty() && time_now - last_graph_use_time >= 10'000'000) {
+        next_graph_eviction_time = time_now + 10'000'000;
+    }
+    last_graph_use_time = time_now;
+
+    // Sweep every 5s during active use, evicting graphs unused for at least 10s.
+    if (time_now >= next_graph_eviction_time) {
+        next_graph_eviction_time = time_now + 5'000'000;
+        for (auto it = cuda_graphs.begin(); it != cuda_graphs.end(); ) {
+            if (it->first != first_node_ptr && time_now - it->second->last_used_time >= 10'000'000) {
+                it = cuda_graphs.erase(it);
+            } else {
+                ++it;
+            }
+        }
+    }
+
+    auto it = cuda_graphs.find(first_node_ptr);
+    if (it == cuda_graphs.end()) {
+        it = cuda_graphs.emplace(first_node_ptr, std::make_unique<ggml_cuda_graph>()).first;
+    }
+    it->second->last_used_time = time_now;
+    return it->second.get();
+}
+
+bool ggml_backend_cuda_context::any_cuda_graph_enabled() const {
+    for (const auto & [key, graph] : cuda_graphs) {
+        if (graph && graph->is_enabled()) {
+            return true;
+        }
+    }
+    return false;
+}
+
+bool ggml_backend_cuda_context::any_cuda_graph_has_instance() const {
+    for (const auto & [key, graph] : cuda_graphs) {
+        if (graph && !graph->instances.empty()) {
+            return true;
+        }
+    }
+    return false;
+}
+#endif // USE_CUDA_GRAPH
+
+#if defined(GGML_CUDA_USE_PDL)
+bool ggml_cuda_kernel_can_use_pdl(const void * kernel) {
+    static const bool env_pdl_enabled = []() {
+        const char * env = getenv("GGML_CUDA_PDL");
+        return env == nullptr || std::atoi(env) != 0;
+    }();
+    if (!env_pdl_enabled) {
+        return false;
+    }
+
+    const int device = ggml_cuda_get_device();
+
+    struct cache_key {
+        int          device;
+        const void * kernel;
+
+        bool operator==(const cache_key & other) const { return device == other.device && kernel == other.kernel; }
+    };
+
+    struct cache_key_hash {
+        // MurmurHash3 mixing function for better hash distribution (vs. just std::hash which in some implementations simply returns the identity)
+        static size_t hash_mix(size_t x) {
+            std::uint64_t       y = x;
+            const std::uint64_t m = 0xe9846af9b1a615d;
+
+            y ^= y >> 32;
+            y *= m;
+            y ^= y >> 32;
+            y *= m;
+            y ^= y >> 28;
+
+            return static_cast<size_t>(y);
+        }
+
+        size_t operator()(const cache_key & key) const {
+            // Use a nonzero seed to avoid mapping all-zero keys to zero
+            size_t h = 42;
+            h        = hash_mix(h + key.device);
+            h        = hash_mix(h + reinterpret_cast<size_t>(key.kernel));
+            return h;
+        }
+    };
+
+    static std::mutex                                          cache_mutex;
+    static std::unordered_map<cache_key, bool, cache_key_hash> cache;
+
+    const cache_key             key = { device, kernel };
+    std::lock_guard<std::mutex> lock(cache_mutex);
+    const auto                  it = cache.find(key);
+    if (it != cache.end()) {
+        return it->second;
+    }
+
+    cudaFuncAttributes attr = {};
+    CUDA_CHECK(cudaFuncGetAttributes(&attr, kernel));
+
+    // PDL device-side primitives are emitted only for PTX versions >= 90.
+    // We have to guard on a loaded kernel's PTX version so a kernel forward-JIT'ed
+    // from pre-Hopper PTX to a Hopper-or-newer GPU does not opt into PDL.
+    const bool can_use_pdl = attr.ptxVersion >= 90;
+    cache.emplace(key, can_use_pdl);
+    return can_use_pdl;
+}
+#endif // defined(GGML_CUDA_USE_PDL)
+
 // destroying a cuBLAS handle while a graph is being captured in a different thread can result in a CUDA error
 // this lock is used to ensure that no cuBLAS handle is destroyed while a graph is being captured

@@ -1869,7 +1998,7 @@ static void ggml_cuda_mul_mat_cublas_impl(ggml_backend_cuda_context & ctx, const

 static void ggml_cuda_mul_mat_cublas(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) {
    const int cc = ggml_cuda_info().devices[ctx.device].cc;
-    const ggml_prec prec = (ggml_prec) ggml_get_op_params_i32(dst, 0);
+    const ggml_prec prec = ggml_prec_get_acc(dst);
    ggml_type compute_type = src0->type;
    if (ggml_is_quantized(compute_type)) {
        compute_type = fast_fp16_hardware_available(cc) ? GGML_TYPE_F16 : GGML_TYPE_F32;
@@ -5143,7 +5272,9 @@ static enum ggml_status ggml_backend_cuda_graph_compute(ggml_backend_t backend,

    ggml_cuda_graph * graph = cuda_ctx->cuda_graph(graph_key);
    if (graph->is_enabled()) {
-        const bool graph_compatible = ggml_cuda_graph_check_compability(cgraph);
+        // graph->uid is set only for a compatible graph (see ggml_cuda_graph_update_required), so a reused graph needs no new check
+        const bool graph_reused     = cgraph->uid != 0 && cgraph->uid == graph->uid;
+        const bool graph_compatible = graph_reused || ggml_cuda_graph_check_compability(cgraph);
        if (graph_compatible) {
            const bool properties_changed = ggml_cuda_graph_update_required(cuda_ctx, cgraph);

@@ -5978,7 +6109,13 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g
                if (b->type == GGML_TYPE_F16 && a->type != GGML_TYPE_F16 && !ggml_cuda_op_mul_mat_use_fwht(op)) {
                    return false;
                }
-                if (op->op == GGML_OP_MUL_MAT_ID && (ggml_get_op_params_i32(op, 3) == GGML_PREC_BF16 || ggml_get_op_params_i32(op, 3) == GGML_PREC_F32)) {
+                const int cc_dev = ggml_cuda_info().devices[dev_ctx->device].cc;
+                const bool f32_src1_ok = GGML_CUDA_CC_IS_NVIDIA(cc_dev) && cc_dev < GGML_CUDA_CC_AMPERE;
+                const enum ggml_prec prec_src1 = ggml_prec_get_src(op, 1);
+                if (prec_src1 == GGML_PREC_F32 && (op->op == GGML_OP_MUL_MAT_ID || !f32_src1_ok || a->type != GGML_TYPE_F32)) {
+                    return false;
+                }
+                if (prec_src1 == GGML_PREC_BF16 && (op->op == GGML_OP_MUL_MAT_ID || a->type != GGML_TYPE_F32)) {
                    return false;
                }
 #ifdef GGML_USE_MUSA
@@ -0,0 +1,191 @@
+#pragma once
+
+// CUDA graph capture and concurrent stream state, only needed by ggml-cuda.cu
+
+#include "common.cuh"
+
+#include <algorithm>
+#include <unordered_map>
+#include <utility>
+#include <vector>
+
+struct ggml_cuda_graph {
+#ifdef USE_CUDA_GRAPH
+    ~ggml_cuda_graph() {
+        for (cudaGraphExec_t instance : instances) {
+            CUDA_CHECK(cudaGraphExecDestroy(instance));
+        }
+    }
+    // the ggml graph is captured as a series of CUDA graphs, so the GPU can run one chunk while the next is enqueued
+    std::vector<cudaGraphExec_t> instances;
+    size_t num_nodes = 0;
+    std::vector<cudaGraphNode_t> nodes;
+    bool disable_due_to_gpu_arch = false;
+    bool warmup_complete = false;
+    uint64_t uid = 0;
+    int64_t last_used_time = 0;
+    struct node_properties {
+        ggml_tensor node;
+        void *   node_src_data_ptrs[GGML_MAX_SRC];
+        int64_t  node_src_ne[GGML_MAX_SRC][GGML_MAX_DIMS];
+        size_t   node_src_nb[GGML_MAX_SRC][GGML_MAX_DIMS];
+    };
+    std::vector<node_properties> node_props;
+
+    bool is_enabled() const {
+        static const bool disable_cuda_graphs_due_to_env = (getenv("GGML_CUDA_DISABLE_GRAPHS") != nullptr);
+        return !(disable_due_to_gpu_arch || disable_cuda_graphs_due_to_env);
+    }
+#endif
+};
+
+struct ggml_cuda_concurrent_event {
+    std::vector<cudaEvent_t> join_events;
+    cudaEvent_t              fork_event = nullptr;
+
+    int                                          n_streams = 0;
+    std::unordered_map<const ggml_tensor *, int> stream_mapping;
+
+    // Original order of nodes in this concurrent region (before interleaving)
+    // Used to restore grouping for fusion within streams
+    std::vector<const ggml_tensor *> original_order;
+
+    const ggml_tensor * join_node;
+
+    ggml_cuda_concurrent_event() = default;
+
+    ggml_cuda_concurrent_event(const ggml_cuda_concurrent_event &) = delete;
+    ggml_cuda_concurrent_event & operator=(const ggml_cuda_concurrent_event &) = delete;
+
+    explicit ggml_cuda_concurrent_event(int n_streams) : n_streams(n_streams) {
+        join_events.resize(n_streams);
+
+        for (size_t i = 0; i < join_events.size(); ++i) {
+            CUDA_CHECK(cudaEventCreateWithFlags(&join_events[i], cudaEventDisableTiming));
+        }
+
+        CUDA_CHECK(cudaEventCreateWithFlags(&fork_event, cudaEventDisableTiming));
+    }
+
+    ggml_cuda_concurrent_event(ggml_cuda_concurrent_event && other) noexcept
+    : join_events(std::move(other.join_events))
+    , fork_event(other.fork_event)
+    , n_streams(other.n_streams)
+    , stream_mapping(std::move(other.stream_mapping))
+    , original_order(std::move(other.original_order))
+    , join_node(other.join_node) {
+        other.fork_event = nullptr;
+    }
+
+    // 1. check if any branches write to overlapping memory ranges (except the join node)
+    // 2. check whether all srcs are either within the branch or outside the nodes covered by ggml_cuda_concurrent_event
+    // we assume all nodes have the same buffer
+    bool is_valid() const {
+        std::vector<std::vector<std::pair<int64_t, int64_t>>> write_ranges;
+        write_ranges.resize(n_streams);
+
+        // get join_node's memory range to exclude from overlap checking.
+        // multiple nodes can use join_node's buffer; we synchronize on the join node.
+        const ggml_tensor * join_t     = join_node->view_src ? join_node->view_src : join_node;
+        const int64_t       join_start = (int64_t) join_t->data;
+        const int64_t       join_end   = join_start + ggml_nbytes(join_t);
+
+        for (const auto & [tensor, stream] : stream_mapping) {
+            const ggml_tensor * t = tensor->view_src ? tensor->view_src : tensor;
+            const int64_t       t_start = (int64_t) t->data;
+            const int64_t       t_end   = t_start + ggml_nbytes(t);
+
+            // skip tensors that overlap with join_node's buffer.
+            if ((t_start <= join_start && join_start < t_end) || (join_start <= t_start && t_start < join_end)) {
+                continue;
+            }
+
+            // concurrent streams begin from 1
+            write_ranges[stream - 1].emplace_back(t_start, t_end);
+        }
+
+        for (int i = 0; i < n_streams; ++i) {
+            // sorts first by start then by end of write range
+            std::sort(write_ranges[i].begin(), write_ranges[i].end());
+        }
+
+        bool writes_overlap = false;
+        bool dependent_srcs = false;
+        for (const auto & [tensor, stream] : stream_mapping) {
+            const ggml_tensor * t = tensor->view_src ? tensor->view_src : tensor;
+            const int64_t       t_start = (int64_t) t->data;
+            const int64_t       t_end   = t_start + ggml_nbytes(t);
+
+            // skip tensors that overlap with join_node's buffer
+            if ((t_start <= join_start && join_start < t_end) || (join_start <= t_start && t_start < join_end)) {
+                continue;
+            }
+
+            // check if this buffer's write data overlaps with another stream's
+            std::pair<int64_t, int64_t> data_range = std::make_pair(t_start, t_end);
+            for (int i = 0; i < n_streams; ++i) {
+                if (i == stream - 1) {
+                    continue;
+                }
+                auto it = std::lower_bound(write_ranges[i].begin(), write_ranges[i].end(), data_range);
+
+                if (it != write_ranges[i].end()) {
+                    const std::pair<int64_t, int64_t> & other = *it;
+
+                    // std::lower_bound returns the first element where other >= data_range (lexicographically).
+                    // This guarantees other.first >= data_range.first.
+                    // Therefore, overlap occurs iff other.first < data_range.second
+                    // (i.e., the other range starts before this range ends).
+                    if (other.first < data_range.second) {
+                        GGML_LOG_DEBUG("Writes overlap for %s", tensor->name);
+                        writes_overlap = true;
+                        break;
+                    }
+                }
+            }
+
+            //check if all srcs are either in branch or don't have a branch
+            for (int i = 0; i < GGML_MAX_SRC; ++i) {
+                if (!tensor->src[i]) {
+                    continue;
+                }
+
+                auto it = stream_mapping.find(tensor->src[i]);
+
+                if (it == stream_mapping.end()) {
+                    continue;
+                }
+
+                if (it->second != stream) {
+                    dependent_srcs = true;
+                    break;
+                }
+            }
+
+            if (dependent_srcs || writes_overlap) {
+                break;
+            }
+        }
+
+        return !writes_overlap && !dependent_srcs;
+    }
+
+    ~ggml_cuda_concurrent_event() {
+        if (fork_event != nullptr) {
+            CUDA_CHECK(cudaEventDestroy(fork_event));
+        }
+        for (cudaEvent_t e : join_events) {
+            if (e != nullptr) {
+                CUDA_CHECK(cudaEventDestroy(e));
+            }
+        }
+    }
+};
+
+struct ggml_cuda_stream_context {
+    std::unordered_map<const ggml_tensor *, ggml_cuda_concurrent_event> concurrent_events;
+
+    void reset() {
+        concurrent_events.clear();
+    }
+};
@@ -616,7 +616,7 @@ static void mul_mat_vec_f_cuda(
        const int64_t ids_stride, enum ggml_prec prec, cudaStream_t stream) {

    if constexpr(std::is_same_v<T, half>) {
-        if (prec == GGML_PREC_DEFAULT) {
+        if (prec == GGML_PREC_UNDEFINED) {
            mul_mat_vec_f_cuda_switch_ncols_dst<T, half>
                (x, y, ids, fusion, dst, ncols, nrows, ncols_dst, stride_row, stride_col_y, stride_col_dst,
                nchannels_x, nchannels_y, nchannels_dst, stride_channel_x, stride_channel_y,
@@ -651,7 +651,8 @@ void ggml_cuda_mul_mat_vec_f(ggml_backend_cuda_context & ctx, const ggml_tensor
    GGML_ASSERT(        nb0        == ts_dst);

    const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc;
-    const enum ggml_prec prec = fast_fp16_available(cc) ? ggml_prec(dst->op_params[0]) : GGML_PREC_F32;
+    // the fused GLU path gets the GLU node as dst, which carries no matmul precision, so it accumulates in F32
+    const enum ggml_prec prec = fast_fp16_available(cc) && !(fusion && fusion->gate) ? ggml_prec_get_acc(dst) : GGML_PREC_F32;

    const float   * src1_d =       (const float   *) src1->data;
    const int32_t *  ids_d = ids ? (const int32_t *)  ids->data : nullptr;
@@ -743,7 +744,7 @@ void ggml_cuda_op_mul_mat_vec_f(

    const int id = ggml_cuda_get_device();
    const int cc = ggml_cuda_info().devices[id].cc;
-    const enum ggml_prec prec = fast_fp16_available(cc) ? ggml_prec(dst->op_params[0]) : GGML_PREC_F32;
+    const enum ggml_prec prec = fast_fp16_available(cc) ? ggml_prec_get_acc(dst) : GGML_PREC_F32;

    // ggml_cuda_op provides single, contiguous matrices
    const int64_t stride_row         = ne00;
@@ -1555,7 +1555,7 @@ void ggml_cuda_mul_mat_vec_q(
        const int64_t s11 = src1->nb[1] / ts_src1;
        const int64_t s12 = src1->nb[2] / ts_src1;
        const int64_t s13 = src1->nb[3] / ts_src1;
-        quantize_row_q8_1_cuda(src1_d, nullptr, src1_q8_1.get(), src0->type, ne10, s11, s12, s13, ne10_padded, ne11, ne12, ne13, stream);
+        quantize_row_q8_1_cuda(src1_d, src1_q8_1.get(), ne10, s11, s12, s13, ne10_padded, ne11, ne12, ne13, stream);
    }

    const int64_t s01 = src0->nb[1] / ts_src0;
@@ -557,10 +557,9 @@ static __global__ void quantize_mmq_q8_1(
 }

 void quantize_row_q8_1_cuda(
-        const float * x, const int32_t * ids, void * vy, const ggml_type type_src0,
+        const float * x, void * vy,
        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(!ids);
    GGML_ASSERT(ne0 % QK8_1 == 0);

    const uint3 ne2_fastdiv = init_fastdiv_values(ne2);
@@ -570,7 +569,6 @@ void quantize_row_q8_1_cuda(
    const dim3 block_size(CUDA_QUANTIZE_BLOCK_SIZE, 1, 1);
    const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(num_blocks, block_size, 0, stream);
    ggml_cuda_kernel_launch(quantize_q8_1, launch_params, x, vy, ne00, s01, s02, s03, ne0, ne1, ne2_fastdiv);
-    GGML_UNUSED(type_src0);
 }

 void quantize_mmq_q8_1_cuda(
@@ -11,14 +11,9 @@
 static_assert(MATRIX_ROW_PADDING %    CUDA_QUANTIZE_BLOCK_SIZE      == 0, "Risk of out-of-bounds access.");
 static_assert(MATRIX_ROW_PADDING % (4*CUDA_QUANTIZE_BLOCK_SIZE_MMQ) == 0, "Risk of out-of-bounds access.");

-typedef void (*quantize_cuda_t)(
-        const float * x, const int32_t * ids, void * vy,
-        ggml_type type_src0, int64_t ne00, int64_t s01, int64_t s02, int64_t s03,
-        int64_t ne0, int64_t ne1, int64_t ne2, int64_t ne3, cudaStream_t stream);
-
 void quantize_row_q8_1_cuda(
-        const float * x, const int32_t * ids, void * vy,
-        ggml_type type_src0, int64_t ne00, int64_t s01, int64_t s02, int64_t s03,
+        const float * x, void * vy,
+        int64_t ne00, int64_t s01, int64_t s02, int64_t s03,
        int64_t ne0, int64_t ne1, int64_t ne2, int64_t ne3, cudaStream_t stream);

 void quantize_mmq_q8_1_cuda(
@@ -57,7 +57,9 @@
 #define cudaDeviceProp hipDeviceProp_t
 #define cudaDeviceSynchronize hipDeviceSynchronize
 #define cudaError_t hipError_t
+#define CUresult hipError_t
 #define cudaErrorMemoryAllocation hipErrorOutOfMemory
+#define CUDA_ERROR_OUT_OF_MEMORY hipErrorOutOfMemory
 #define cudaErrorPeerAccessAlreadyEnabled hipErrorPeerAccessAlreadyEnabled
 #define cudaErrorPeerAccessNotEnabled hipErrorPeerAccessNotEnabled
 #define cudaEventCreateWithFlags hipEventCreateWithFlags
@@ -129,6 +129,7 @@

 // Additional mappings for MUSA graphs
 #define CUDA_SUCCESS MUSA_SUCCESS
+#define CUDA_ERROR_OUT_OF_MEMORY MUSA_ERROR_OUT_OF_MEMORY
 #define CUresult MUresult
 #define cuGetErrorString muGetErrorString
 #define cudaErrorGraphExecUpdateFailure musaErrorGraphExecUpdateFailure
@@ -1523,7 +1523,7 @@ bool ggml_et_op_flash_attn_ext(ggml_backend_et_device_context * dev_ctx, const g
    }

    const ggml_prec prec = ggml_flash_attn_ext_get_prec(node);
-    if (prec != GGML_PREC_F32 && prec != GGML_PREC_BF16 && prec != GGML_PREC_DEFAULT) {
+    if (prec != GGML_PREC_F32 && prec != GGML_PREC_BF16 && prec != GGML_PREC_UNDEFINED) {
        GGML_LOG_ERROR("ET: FLASH_ATTN_EXT baseline kernel only supports F32 precision\n");
        return false;
    }
@@ -941,6 +941,9 @@ static bool ggml_backend_et_device_supports_op(ggml_backend_dev_t dev, const ggm
                        op->nb[1] == op->ne[0] * sizeof(float);
            break;
        case GGML_OP_MUL_MAT:
+            if (ggml_prec_get_src(op, 1) == GGML_PREC_BF16 || ggml_prec_get_src(op, 1) == GGML_PREC_F32) {
+                break;
+            }
            // Support Q8_0 x F32 -> F32, F16 x F32 -> F32, F16 x F16 -> F32, and F32 x F32 -> F32 matrix multiplication
            // Stride requirements: first dimension must be contiguous for all tensors
            if (op->type == GGML_TYPE_F32 &&
@@ -1299,7 +1302,7 @@ static bool ggml_backend_et_device_supports_op(ggml_backend_dev_t dev, const ggm
                const bool me_eligible = op->src[1]->type == GGML_TYPE_F16 && op->src[2]->type == GGML_TYPE_F16 &&
                                         (op->src[0]->ne[0] % 32) == 0;

-                supported = me_eligible && mask_ok && (prec == GGML_PREC_F32 || prec == GGML_PREC_BF16 || prec == GGML_PREC_DEFAULT) &&
+                supported = me_eligible && mask_ok && (prec == GGML_PREC_F32 || prec == GGML_PREC_BF16 || prec == GGML_PREC_UNDEFINED) &&
                            max_bias == 0.0f && logit_softcap == 0.0f && op->src[0]->nb[0] == sizeof(float) &&
                            op->src[1]->nb[0] == k_elem && op->src[2]->nb[0] == v_elem && op->nb[0] == sizeof(float) &&
                            op->src[0]->ne[0] == op->src[1]->ne[0] &&  // dk matches
@@ -78,7 +78,7 @@ typedef int (*fastrpc_munmap_pfn_t)(int domain, int fd, void *addr, size_t lengt

 typedef int (*remote_handle64_open_pfn_t)(const char* name, remote_handle64 *ph);
 typedef int (*remote_handle64_invoke_pfn_t)(remote_handle64 h, uint32_t dwScalars, remote_arg *pra);
-typedef int (*remote_handle64_close_pfn_t)(remote_handle h);
+typedef int (*remote_handle64_close_pfn_t)(remote_handle64 h);
 typedef int (*remote_handle_control_pfn_t)(uint32_t req, void* data, uint32_t datalen);
 typedef int (*remote_handle64_control_pfn_t)(remote_handle64 h, uint32_t req, void* data, uint32_t datalen);
 typedef int (*remote_session_control_pfn_t)(uint32_t req, void *data, uint32_t datalen);
@@ -939,14 +939,15 @@ static bool htp_matmul_check_types(
        return false;
    }

+    if (ggml_prec_get_src(dst, 1) == GGML_PREC_BF16 || ggml_prec_get_src(dst, 1) == GGML_PREC_F32) {
+        return false;
+    }
+
    if (is_matmul_id) {
        const struct ggml_tensor * src2 = dst->src[2];
        if (src1->type != GGML_TYPE_F32 || !src2 || src2->type != GGML_TYPE_I32) {
            return false;
        }
-        if (ggml_get_op_params_i32(dst, 3) == GGML_PREC_BF16 || ggml_get_op_params_i32(dst, 3) == GGML_PREC_F32) {
-            return false;
-        }
        switch (src0->type) {
            case GGML_TYPE_Q1_0:
            case GGML_TYPE_Q2_0:
@@ -1510,6 +1510,18 @@ void ggml_metal_device_get_memory(ggml_metal_device_t dev, size_t * free, size_t
    }
 }

+static bool ggml_metal_supports_mul_mat_src1_prec(const struct ggml_tensor * op, bool has_tensor) {
+    const enum ggml_prec prec = ggml_prec_get_src(op, 1);
+
+    if (prec == GGML_PREC_F32) {
+        return false;
+    }
+    if (prec == GGML_PREC_BF16) {
+        return has_tensor && op->op == GGML_OP_MUL_MAT && op->src[0]->type != GGML_TYPE_F32;
+    }
+    return true;
+}
+
 static bool ggml_metal_supports_mul_mat_op(
        bool has_simdgroup_reduction,
        const struct ggml_tensor * op,
@@ -1894,6 +1906,9 @@ bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_te
        case GGML_OP_SOLVE_TRI:
            return has_simdgroup_reduction && op->src[0]->type == GGML_TYPE_F32;
        case GGML_OP_MUL_MAT:
+            if (!ggml_metal_supports_mul_mat_src1_prec(op, dev->props.has_tensor)) {
+                return false;
+            }
            // the FWHT kernels read an F16 source directly; every other F16 src1 path
            // still goes through ggml_metal_supports_mul_mat_op
            if (op->src[0]->type == GGML_TYPE_F32 && op->src[1]->type == GGML_TYPE_F16 &&
@@ -1904,6 +1919,9 @@ bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_te
                    has_simdgroup_reduction, op, true,
                    ggml_metal_op_mul_mat_use_mm(op, has_simdgroup_mm));
        case GGML_OP_MUL_MAT_ID:
+            if (!ggml_metal_supports_mul_mat_src1_prec(op, dev->props.has_tensor)) {
+                return false;
+            }
            return ggml_metal_supports_mul_mat_op(
                    has_simdgroup_reduction, op, false,
                    ggml_metal_op_mul_mat_id_use_mm(op, has_simdgroup_mm));
@@ -11937,6 +11937,9 @@ static bool ggml_opencl_supports_op(ggml_backend_dev_t dev, const struct ggml_te
        case GGML_OP_GROUP_NORM:
            return ggml_is_contiguous(op->src[0]);
        case GGML_OP_MUL_MAT:
+            if (ggml_prec_get_src(op, 1) == GGML_PREC_BF16 || ggml_prec_get_src(op, 1) == GGML_PREC_F32) {
+                return false;
+            }
            if (ggml_cl_mul_mat_can_use_fwht(backend_ctx, op)) {
                return true;
            }
@@ -12056,7 +12059,7 @@ static bool ggml_opencl_supports_op(ggml_backend_dev_t dev, const struct ggml_te
                       op->src[1]->ne[1] <= INT_MAX && op->ne[1] <= INT_MAX && op->ne[2] <= INT_MAX - 64;
            }
 #ifdef GGML_OPENCL_USE_ADRENO_KERNELS
-            if ((ggml_get_op_params_i32(op, 3) == GGML_PREC_BF16 || ggml_get_op_params_i32(op, 3) == GGML_PREC_F32) &&
+            if ((ggml_prec_get_src(op, 1) == GGML_PREC_BF16 || ggml_prec_get_src(op, 1) == GGML_PREC_F32) &&
                use_adreno_moe_kernels(backend_ctx, op->src[0])) {
                return false;
            }
@@ -1581,6 +1581,9 @@ static ggml_openvino_op_support is_op_supported_case(const ggml_tensor * op) {
        break;
    }
    case GGML_OP_MUL_MAT: {
+        if (ggml_prec_get_src(op, 1) == GGML_PREC_BF16 || ggml_prec_get_src(op, 1) == GGML_PREC_F32) {
+            return {false, "MUL_MAT with a BF16 or F32 src1 precision is not supported"};
+        }
        if (ggml_openvino_is_gpu() && op->src[0] != nullptr && op->src[1] != nullptr &&
            ggml_is_quantized(op->src[0]->type) && strcmp(op->src[0]->name, "a") == 0 &&
            strcmp(op->src[1]->name, "b") == 0 && op->src[0]->ne[1] == 1 && op->src[1]->ne[1] == 64 &&
@@ -444,7 +444,10 @@ static bool negotiate_hello(const std::shared_ptr<socket_t> & sock, rpc_dispatch
    sock->get_caps(request.conn_caps);

    bool status = send_rpc_cmd(sock, RPC_CMD_HELLO, &request, sizeof(request), &response, sizeof(response));
-    RPC_STATUS_ASSERT(status);
+    if (!status) {
+        LOG_ERROR("RPC hello exchange with the server failed\n");
+        return false;
+    }

    if (response.major != RPC_PROTO_MAJOR_VERSION || response.minor > RPC_PROTO_MINOR_VERSION) {
        LOG_ERROR("RPC server version mismatch: %d.%d.%d\n",
@@ -526,7 +529,7 @@ public:
    void busy_spin_release();
    void graph_compute(uint32_t device, const ggml_cgraph * cgraph);

-    void start(const std::string & endpoint);
+    bool start(const std::string & endpoint);
    void work();

    bool get_send_tensor_hash() const { return flags & RPC_DISPATCHER_FLAGS_SEND_TENSOR_HASH; }
@@ -641,26 +644,31 @@ void rpc_dispatcher::busy_spin_release() {
    GGML_ASSERT(previous > 0);
 }

-void rpc_dispatcher::start(const std::string & endpoint) {
+bool rpc_dispatcher::start(const std::string & endpoint) {
    std::string host;
    int port;
    if (!parse_endpoint(endpoint, host, port)) {
-        GGML_ABORT("Failed to parse endpoint: %s\n", endpoint.c_str());
+        LOG_ERROR("Failed to parse endpoint: %s\n", endpoint.c_str());
+        return false;
    }
    if (!rpc_transport_init()) {
-        GGML_ABORT("RPC transport initialization failed\n");
+        LOG_ERROR("RPC transport initialization failed\n");
+        return false;
    }

    sock = socket_t::connect(host.c_str(), port);
    if (sock == nullptr) {
-        GGML_ABORT("Failed to connect to %s\n", endpoint.c_str());
+        LOG_ERROR("Failed to connect to %s\n", endpoint.c_str());
+        return false;
    }
    if (!negotiate_hello(sock, flags)) {
-        GGML_ABORT("RPC handshake failed for %s\n", endpoint.c_str());
+        LOG_ERROR("RPC handshake failed for %s\n", endpoint.c_str());
+        return false;
    }
    LOG_DBG("[%s] connected to %s\n", __func__, endpoint.c_str());
    running = true;
    thread = std::thread(rpc_dispatcher_trampoline, this);
+    return true;
 }

 void rpc_dispatcher::work() {
@@ -724,7 +732,9 @@ static std::shared_ptr<rpc_dispatcher> get_dispatcher(const std::string & endpoi
    }

    auto dispatcher = std::make_shared<rpc_dispatcher>();
-    dispatcher->start(endpoint);
+    if (!dispatcher->start(endpoint)) {
+        return nullptr;
+    }
    dispatchers[endpoint] = dispatcher;
    return dispatcher;
 }
@@ -999,6 +1009,9 @@ static ggml_backend_buffer_t ggml_backend_rpc_buffer_type_alloc_buffer(ggml_back
    rpc_msg_alloc_buffer_rsp response;

    auto dispatcher = get_dispatcher(buft_ctx->endpoint);
+    if (dispatcher == nullptr) {
+        return nullptr;
+    }
    dispatcher->send(RPC_CMD_ALLOC_BUFFER, request, sizeof(*request), nullptr, 0, &response, sizeof(response));
    if (response.remote_ptr != 0) {
        LOG_DBG("[%s] endpoint: %s, device: %u, size: %s -> remote_ptr: 0x%" PRIx64 ", remote_size: %s\n",
@@ -1119,6 +1132,9 @@ static size_t ggml_backend_rpc_buffer_type_get_alloc_size(ggml_backend_buffer_ty
        LOG_DBG2("[%s] cache miss: %s, querying server\n", __func__, rpc_log_format_tensor(tensor).c_str());
        rpc_msg_get_alloc_size_rsp response;
        auto dispatcher = get_dispatcher(buft_ctx->endpoint);
+        if (dispatcher == nullptr) {
+            return min_size;
+        }
        dispatcher->send(RPC_CMD_GET_ALLOC_SIZE, request, sizeof(*request), nullptr, 0, &response, sizeof(response));

        {
@@ -1339,6 +1355,9 @@ ggml_backend_buffer_type_t ggml_backend_rpc_buffer_type(const char * endpoint, u
        return it->second;
    }
    auto dispatcher = get_dispatcher(endpoint);
+    if (dispatcher == nullptr) {
+        return nullptr;
+    }
    size_t alignment = get_alignment(dispatcher, device);
    size_t max_size = get_max_size(dispatcher, device);
    ggml_backend_rpc_buffer_type_context * buft_ctx = new ggml_backend_rpc_buffer_type_context {
@@ -1363,6 +1382,9 @@ ggml_backend_buffer_type_t ggml_backend_rpc_buffer_type(const char * endpoint, u
 ggml_backend_t ggml_backend_rpc_init(const char * endpoint, uint32_t device) {
    std::string dev_name = "RPC" + std::to_string(device) + "[" + std::string(endpoint) + "]";
    auto dispatcher = get_dispatcher(endpoint);
+    if (dispatcher == nullptr) {
+        return nullptr;
+    }
    ggml_backend_rpc_context * ctx = new ggml_backend_rpc_context {
        /* .dispatcher = */ dispatcher,
        /* .endpoint   = */ endpoint,
@@ -1386,6 +1408,11 @@ bool ggml_backend_is_rpc(ggml_backend_t backend) {

 void ggml_backend_rpc_get_device_memory(const char * endpoint, uint32_t device, size_t * free, size_t * total) {
    auto dispatcher = get_dispatcher(endpoint);
+    if (dispatcher == nullptr) {
+        *free  = 0;
+        *total = 0;
+        return;
+    }
    auto request = std::make_shared<rpc_msg_get_device_memory_req>();
    request->device = device;
    rpc_msg_get_device_memory_rsp response;
@@ -3103,18 +3130,27 @@ static bool ggml_backend_rpc_device_supports_buft(ggml_backend_dev_t dev, ggml_b
 static ggml_backend_event_t ggml_backend_rpc_device_event_new(ggml_backend_dev_t dev) {
    ggml_backend_rpc_device_context * ctx = (ggml_backend_rpc_device_context *)dev->context;
    auto dispatcher = get_dispatcher(ctx->endpoint);
+    if (dispatcher == nullptr) {
+        return nullptr;
+    }
    return dispatcher->event_new(dev);
 }

 static void ggml_backend_rpc_device_event_free(ggml_backend_dev_t dev, ggml_backend_event_t event) {
    ggml_backend_rpc_device_context * ctx = (ggml_backend_rpc_device_context *)dev->context;
    auto dispatcher = get_dispatcher(ctx->endpoint);
+    if (dispatcher == nullptr) {
+        return;
+    }
    dispatcher->event_free(event);
 }

 static void ggml_backend_rpc_device_event_synchronize(ggml_backend_dev_t dev, ggml_backend_event_t event) {
    ggml_backend_rpc_device_context * ctx = (ggml_backend_rpc_device_context *)dev->context;
    auto dispatcher = get_dispatcher(ctx->endpoint);
+    if (dispatcher == nullptr) {
+        return;
+    }
    dispatcher->event_synchronize(event);
 }

@@ -3334,6 +3370,9 @@ ggml_backend_reg_t ggml_backend_rpc_reg(void) {

 static uint32_t ggml_backend_rpc_get_device_count(const char * endpoint) {
    auto dispatcher = get_dispatcher(endpoint);
+    if (dispatcher == nullptr) {
+        return 0;
+    }
    rpc_msg_device_count_rsp response;
    dispatcher->send(RPC_CMD_DEVICE_COUNT, nullptr, 0, nullptr, 0, &response, sizeof(response));
    return response.device_count;
@@ -92,9 +92,17 @@ enum ggml_sycl_xmx_gather_type {
    GGML_SYCL_XMX_GATHER_Q4_K     = 1 << 10,
    GGML_SYCL_XMX_GATHER_Q5_K     = 1 << 11,
    GGML_SYCL_XMX_GATHER_Q6_K     = 1 << 12,
+    GGML_SYCL_XMX_GATHER_Q2_K     = 1 << 13,
+    GGML_SYCL_XMX_GATHER_Q3_K     = 1 << 14,
 };
 static constexpr int GGML_SYCL_XMX_GATHER_TYPES_DEFAULT = ~0;
 extern int g_ggml_sycl_xmx_gather_types;
+
+enum ggml_sycl_mmid_sched_bit {
+    GGML_SYCL_MMID_SCHED_DEVICE    = 1 << 0,
+    GGML_SYCL_MMID_SCHED_UNBOUNDED = 1 << 1,
+};
+extern int g_ggml_sycl_mmid_sched;
 // Which joint_matrix combinations the XMX dequant-GEMM paths may use, one bit each (see fused-gemm.cpp).
 // GGML_SYCL_DYNAMIC_PRECISION picks the operand type, this mask the combinations of that type.
 static constexpr int GGML_SYCL_XMX_GATHER_SHAPES_DEFAULT = 0xff;
@@ -121,7 +129,7 @@ extern int g_ggml_sycl_dynamic_required_precision;
 // [TAG_GGML_PREC] src1 precision request of the MUL_MAT/MUL_MAT_ID op dst
 static inline int32_t ggml_sycl_src1_prec(const ggml_tensor * dst) {
    return g_ggml_sycl_dynamic_required_precision == GGML_SYCL_DYNAMIC_PRECISION_F16 ? GGML_PREC_UNDEFINED :
-                                                                                       dst->op_params[3];
+                                                                                       ggml_prec_get_src(dst, 1);
 }

 // [TAG_GGML_PREC] the library GEMM and dmmv may convert src1 of the MUL_MAT/MUL_MAT_ID op dst to f16
@@ -443,7 +451,7 @@ struct mmid_row_mapping {
    int32_t i2;
 };

-struct ggml_sycl_gg_tile {
+struct ggml_sycl_grouped_gemm_tile {
    int32_t expert;
    int32_t n0;
    int32_t n1;
@@ -552,7 +560,7 @@ struct ggml_backend_sycl_context {
    std::unique_ptr<ggml_sycl_pool> host_pools[GGML_SYCL_MAX_DEVICES];

    std::vector<mmid_row_mapping> mmid_row_mapping_host;
-    std::vector<ggml_sycl_gg_tile> mmid_tile_schedule_host;
+    std::vector<ggml_sycl_grouped_gemm_tile> mmid_tile_schedule_host;

    static std::unique_ptr<ggml_sycl_pool> new_pool_for_device(queue_ptr qptr, int device);

@@ -69,17 +69,6 @@ bool ggml_sycl_flash_attn_ext_onednn_supported(const ggml_tensor * dst, bool use
            }
        }
    }
-    // This is the improved SPDA gate. Rather than gating Alchemist GPUs from all SPDA features, we instead target only the failing shapes.
-    // If the GPU being assessed isn't in the grouping below, it has full access to all SPDA shapes. Otherwise, if it's an Alchemist GPU, we block only the shapes with head sizes that fail.
-    // It is much easier to compare the device to a small list of failing cases than to define all the passing ones.
-    const gpu_arch arch = ggml_sycl_info().devices[ggml_sycl_get_device()].hw_info.arch;
-    bool support_spda = !(arch == gpu_arch::intel_gpu_dg2_g10 ||
-                         arch == gpu_arch::intel_gpu_dg2_g11 ||
-                         arch == gpu_arch::intel_gpu_dg2_g12);
-
-    if (!support_spda && K->ne[0] == 64) {
-        return false;
-    }
    // Optional KV-length ceiling (GGML_SYCL_FA_ONEDNN_MAX_KV, 0 = unlimited). Escape hatch:
    // very long sequences make the fused SDPA slow enough to risk the xe driver watchdog on
    // some stacks; past the cap we fall back to the native FA kernel instead.
@@ -411,6 +411,8 @@ template <> struct fg_block_traits<block_iq2_s>   { static constexpr int qk = QK
 template <> struct fg_block_traits<block_iq1_s>   { static constexpr int qk = QK_K; };
 template <> struct fg_block_traits<block_iq1_m>   { static constexpr int qk = QK_K; };
 template <> struct fg_block_traits<block_q8_0>    { static constexpr int qk = QK8_0; };
+template <> struct fg_block_traits<block_q2_K>    { static constexpr int qk = QK_K; };
+template <> struct fg_block_traits<block_q3_K>    { static constexpr int qk = QK_K; };
 template <> struct fg_block_traits<block_q4_K>    { static constexpr int qk = QK_K; };
 template <> struct fg_block_traits<block_q5_K>    { static constexpr int qk = QK_K; };
 template <> struct fg_block_traits<block_q6_K>    { static constexpr int qk = QK_K; };
@@ -595,6 +597,49 @@ static __dpct_inline__ void fg_scale_min_k4(const int j, const uint8_t * __restr
    }
 }

+template <typename E>
+static __dpct_inline__ void fg_decode_q2_K(const uint8_t * __restrict__ qs, const uint8_t * __restrict__ scales,
+                                           const float d, const float dmin, const int ib, typename E::pair * a) {
+    const uint8_t * q = qs + 32 * (ib / 4);
+    const int shift = 2 * (ib % 4);
+#pragma unroll
+    for (int j = 0; j < 16; ++j) {
+        const uint8_t sc = scales[2 * ib + j / 8];
+        const float dl = d * (sc & 0xF);
+        const float ml = dmin * (sc >> 4);
+        a[j] = E::make(dl * ((q[2 * j] >> shift) & 3) - ml,
+                       dl * ((q[2 * j + 1] >> shift) & 3) - ml);
+    }
+}
+
+// scales of the two 16-value halves of sub-block ib
+static __dpct_inline__ void fg_q3_K_dl(const uint8_t * __restrict__ scales, const float d, const int ib, float * dl) {
+#pragma unroll
+    for (int i = 0; i < 2; ++i) {
+        const int s  = 2 * ib + i;
+        const int sc = ((scales[s % 8] >> (4 * (s / 8))) & 0xF) | (((scales[8 + s % 4] >> (2 * (s / 4))) & 3) << 4);
+        dl[i] = d * (sc - 32);
+    }
+}
+
+template <typename E>
+static __dpct_inline__ void fg_decode_q3_K(const uint8_t * __restrict__ qs, const uint8_t * __restrict__ hmask,
+                                           const uint8_t * __restrict__ scales, const float d, const int ib,
+                                           typename E::pair * a) {
+    // block_q3_K is 110 bytes, so qs and hmask are only 2-byte aligned
+    const uint16_t * q = (const uint16_t *) (qs + 32 * (ib / 4));
+    const uint16_t * h = (const uint16_t *) hmask;
+    const int shift = 2 * (ib % 4);
+    float dl[2];
+    fg_q3_K_dl(scales, d, ib, dl);
+#pragma unroll
+    for (int j = 0; j < 16; ++j) {
+        // per byte: 2 low bits, plus 4 if the hmask bit is set, minus 4
+        const uint32_t v = ((q[j] >> shift) & 0x0303u) | (((h[j] >> ib) & 0x0101u) << 2);
+        a[j] = E::make(dl[j / 8] * ((int) (v & 0xFF) - 4), dl[j / 8] * ((int) (v >> 8) - 4));
+    }
+}
+
 // q4_K sub-block ib (0..7) uses scale/min pair ib and the low (even ib) or high (odd ib) nibbles of
 // qs[32 * (ib / 2) ...], as in dequantize_row_q4_K
 template <typename E>
@@ -658,6 +703,18 @@ static __dpct_inline__ void fg_stage_a(const block_q8_0 * __restrict__ xrow, con
    fg_decode_q8_0<E>(xrow[kb].qs, (float) xrow[kb].d, a);
 }

+template <typename E>
+static __dpct_inline__ void fg_stage_a(const block_q2_K * __restrict__ xrow, const int kb, typename E::pair * a) {
+    FG_SUPERBLOCK(block_q2_K);
+    fg_decode_q2_K<E>(blk->qs, blk->scales, (float) blk->dm[0], (float) blk->dm[1], ib, a);
+}
+
+template <typename E>
+static __dpct_inline__ void fg_stage_a(const block_q3_K * __restrict__ xrow, const int kb, typename E::pair * a) {
+    FG_SUPERBLOCK(block_q3_K);
+    fg_decode_q3_K<E>(blk->qs, blk->hmask, blk->scales, (float) blk->d, ib, a);
+}
+
 template <typename E>
 static __dpct_inline__ void fg_stage_a(const block_q4_K * __restrict__ xrow, const int kb, typename E::pair * a) {
    FG_SUPERBLOCK(block_q4_K);
@@ -693,6 +750,50 @@ template <> struct fg_soa<block_q8_0> {
    }
 };

+template <> struct fg_soa<block_q2_K> {
+    static constexpr bool supported = true;
+
+    // [qs][scales][dm]
+    template <typename E>
+    static __dpct_inline__ void stage(const uint8_t * x, const size_t nblocks, const size_t ib, const int kb,
+                                      typename E::pair * a) {
+        const uint8_t * scales = x + nblocks * (QK_K / 4);
+        const sycl::half2 dm = ((const sycl::half2 *) (scales + nblocks * (QK_K / 16)))[ib];
+        fg_decode_q2_K<E>(x + ib * (QK_K / 4), scales + ib * (QK_K / 16), (float) dm[0], (float) dm[1],
+                          kb % (QK_K / 32), a);
+    }
+};
+
+template <> struct fg_soa<block_q3_K> {
+    static constexpr bool supported = true;
+
+    // [qs][hmask][scales][d]
+    template <typename E>
+    static __dpct_inline__ void stage(const uint8_t * x, const size_t nblocks, const size_t ib, const int kb,
+                                      typename E::pair * a) {
+        const uint8_t * hmask = x + nblocks * (QK_K / 4);
+        const uint8_t * scales = hmask + nblocks * (QK_K / 8);
+        const float d = (float) ((const sycl::half *) (scales + nblocks * 12))[ib];
+        const int sub = kb % (QK_K / 32);
+        // 16-byte aligned, see opt_for_reorder_id
+        using v4 = sycl::vec<uint32_t, 4>;
+        const v4 * q = (const v4 *) (x + ib * (QK_K / 4) + 32 * (sub / 4));
+        const v4 * h = (const v4 *) (hmask + ib * (QK_K / 8));
+        const v4 qv[2] = { q[0], q[1] };
+        const v4 hv[2] = { h[0], h[1] };
+        const int shift = 2 * (sub % 4);
+        float dl[2];
+        fg_q3_K_dl(scales + ib * 12, d, sub, dl);
+#pragma unroll
+        for (int w = 0; w < 8; ++w) {
+            const uint32_t v = ((qv[w / 4][w % 4] >> shift) & 0x03030303u) | (((hv[w / 4][w % 4] >> sub) & 0x01010101u) << 2);
+            const float    f = dl[w / 4];
+            a[2 * w]     = E::make(f * ((int) (v & 0xFF) - 4), f * ((int) ((v >> 8) & 0xFF) - 4));
+            a[2 * w + 1] = E::make(f * ((int) ((v >> 16) & 0xFF) - 4), f * ((int) (v >> 24) - 4));
+        }
+    }
+};
+
 template <> struct fg_soa<block_q4_K> {
    static constexpr bool supported = true;
    // [qs][scales][dm]
@@ -745,7 +846,7 @@ static void fused_dequant_gemm_tile(
    const int b0, const int n0, const int n1,
    sycl::local_accessor<typename S::tsa, 1> tile_a,
    sycl::local_accessor<float, 1> tile_c,
-    const sycl::nd_item<2> & item) {
+    const sycl::nd_item<2> & item, ggml_sycl_grouped_gemm_rows dst_rows = {}) {
    if constexpr (fg_built<S::SG>()) {
        using EA = typename S::EA;
        using TA = typename S::ta;
@@ -854,7 +955,13 @@ static void fused_dequant_gemm_tile(
                for (int s = 0; s < FG_KSPLIT; ++s) {
                    sum += tile_c[s * S::SG_ROWS * S::BN + r * S::BN + c];
                }
-                dst[(size_t) n * ldd + m] = sum;
+                if (dst_rows.map) {
+                    const mmid_row_mapping row = dst_rows.map[n];
+                    float * out = (float *) (dst_rows.base + (row.i1 % dst_rows.ne1) * dst_rows.nb1 + row.i2 * dst_rows.nb2);
+                    out[m] = sum;
+                } else {
+                    dst[(size_t) n * ldd + m] = sum;
+                }
            }
        }
    }
@@ -880,20 +987,24 @@ static void fused_dequant_gemm_launch(const void * src0, const typename S::tsb *
 // grouped: work-group (t, mt) is tile t of the schedule; its B columns sit at t * BN
 template <typename S, typename block_q_t, bool reordered>
 static void grouped_dequant_gemm_launch(const char * src0_dd, const size_t expert_stride,
-                                        const ggml_sycl_gg_tile * tiles_ptr, const typename S::tsb * packed, float * dst,
-                                        const int M, const int Npad, const int K, const int64_t n_tiles,
-                                        const int64_t groups_m, dpct::queue_ptr stream) {
+                                        const ggml_sycl_grouped_gemm_tile * tiles_ptr, const typename S::tsb * packed,
+                                        float * dst, const int M, const int Npad, const int K, const int64_t n_tiles,
+                                        const int64_t groups_m, dpct::queue_ptr stream,
+                                        ggml_sycl_grouped_gemm_rows dst_rows = {}) {
    stream->submit([&](sycl::handler & cgh) {
        sycl::local_accessor<typename S::tsa, 1> tile_a(FG_KSPLIT * S::SG_ROWS * FG_BK, cgh);
        sycl::local_accessor<float, 1>          tile_c(FG_KSPLIT * S::SG_ROWS * S::BN, cgh);
        cgh.parallel_for(
            sycl::nd_range<2>(sycl::range<2>(n_tiles, groups_m * S::WG_SIZE), sycl::range<2>(1, S::WG_SIZE)),
            [=](sycl::nd_item<2> item) [[sycl::reqd_sub_group_size(S::SG)]] {
-                const int               t    = item.get_group(0);
-                const ggml_sycl_gg_tile tile = tiles_ptr[t];
+                const int                         t    = item.get_group(0);
+                const ggml_sycl_grouped_gemm_tile tile = tiles_ptr[t];
+                if (tile.n0 == tile.n1) {
+                    return;
+                }
                const block_q_t *       x    = (const block_q_t *) (src0_dd + (size_t) tile.expert * expert_stride);
                fused_dequant_gemm_tile<S, block_q_t, reordered>(x, packed, dst, M, Npad, K, M, t * S::BN, tile.n0,
-                                                                 tile.n1, tile_a, tile_c, item);
+                                                                 tile.n1, tile_a, tile_c, item, dst_rows);
            });
    });
 }
@@ -901,17 +1012,17 @@ static void grouped_dequant_gemm_launch(const char * src0_dd, const size_t exper
 // src1 f32 rows -> packed [K/V][n_tiles*BN][V], tile t holds its rows [n0, n1) at columns t*BN..,
 // zero past n1. The column runs fastest so a sub-group writes one contiguous run.
 template <typename S>
-static void grouped_gemm_pack_b(const float * y, typename S::tsb * packed, const ggml_sycl_gg_tile * tiles, int Npad,
-                                int K, dpct::queue_ptr stream) {
+static void grouped_gemm_pack_b(const float * y, typename S::tsb * packed, const ggml_sycl_grouped_gemm_tile * tiles,
+                                int Npad, int K, dpct::queue_ptr stream) {
    using E        = typename S::EB;
    constexpr int V = S::VNNI;
    const int kqs   = K / V;
    stream->parallel_for(sycl::range<1>((size_t) Npad * kqs), [=](sycl::id<1> id) {
-        const size_t idx = id[0];
-        const int    kq  = idx / Npad;
-        const int    n   = idx - (size_t) kq * Npad;
-        const ggml_sycl_gg_tile tile = tiles[n / S::BN];
-        const int    row = tile.n0 + n % S::BN;
+        const size_t                      idx  = id[0];
+        const int                         kq   = idx / Npad;
+        const int                         n    = idx - (size_t) kq * Npad;
+        const ggml_sycl_grouped_gemm_tile tile = tiles[n / S::BN];
+        const int                         row  = tile.n0 + n % S::BN;
        // one guarded load run per work-item, as a per-element select costs ~1.5% prefill
        typename S::tsb vals[V] = {};
        if (row < tile.n1) {
@@ -932,7 +1043,8 @@ static void grouped_gemm_pack_b(const float * y, typename S::tsb * packed, const
 // q8_0 and the k-quants take only the grouped path. The plain kernel decodes A again for every BN columns
 // of a dense batch, and for these formats that costs more than the one dequantization of the library GEMM.
 template <typename T> static constexpr bool fg_plain_ok() {
-    return !std::is_same_v<T, block_q8_0> && !std::is_same_v<T, block_q4_K> && !std::is_same_v<T, block_q5_K> &&
+    return !std::is_same_v<T, block_q8_0> && !std::is_same_v<T, block_q2_K> && !std::is_same_v<T, block_q3_K> &&
+           !std::is_same_v<T, block_q4_K> && !std::is_same_v<T, block_q5_K> &&
           !std::is_same_v<T, block_q6_K>;
 }

@@ -966,6 +1078,8 @@ template <typename F> static bool fg_visit_type(ggml_type type, bool reordered,
        case GGML_TYPE_IQ1_S:   return fg_visit_layout<block_iq1_s>(reordered, f);
        case GGML_TYPE_IQ1_M:   return fg_visit_layout<block_iq1_m>(reordered, f);
        case GGML_TYPE_Q8_0:    return fg_visit_layout<block_q8_0>(reordered, f);
+        case GGML_TYPE_Q2_K:    return fg_visit_layout<block_q2_K>(reordered, f);
+        case GGML_TYPE_Q3_K:    return fg_visit_layout<block_q3_K>(reordered, f);
        case GGML_TYPE_Q4_K:    return fg_visit_layout<block_q4_K>(reordered, f);
        case GGML_TYPE_Q5_K:    return fg_visit_layout<block_q5_K>(reordered, f);
        case GGML_TYPE_Q6_K:    return fg_visit_layout<block_q6_K>(reordered, f);
@@ -1037,7 +1151,7 @@ bool ggml_sycl_fused_dequant_gemm(ggml_type src0_type, bool src0_reordered, cons
 template <typename S>
 static bool fg_grouped_run(ggml_type src0_type, bool reordered, const void * src0_base, size_t expert_stride,
                           const float * src1, float * dst, const int64_t * expert_row_offsets, int64_t n_as, int64_t M,
-                           int64_t K, std::vector<ggml_sycl_gg_tile> & tiles, ggml_sycl_pool & pool,
+                           int64_t K, std::vector<ggml_sycl_grouped_gemm_tile> & tiles, ggml_sycl_pool & pool,
                           dpct::queue_ptr stream) {
    // the host knows every slice, so it lays out the work-groups: no search on the device
    tiles.clear();
@@ -1053,15 +1167,16 @@ static bool fg_grouped_run(ggml_type src0_type, bool reordered, const void * src

    const int64_t max_tiles = grouped_gemm_max_tiles(expert_row_offsets[n_as], n_as, S::BN);
    GGML_ASSERT(n_tiles <= max_tiles);
-    ggml_sycl_pool_alloc<ggml_sycl_gg_tile> tiles_dev(pool, max_tiles);
-    SYCL_CHECK(CHECK_TRY_ERROR(stream->memcpy(tiles_dev.get(), tiles.data(), n_tiles * sizeof(ggml_sycl_gg_tile))));
+    ggml_sycl_pool_alloc<ggml_sycl_grouped_gemm_tile> tiles_dev(pool, max_tiles);
+    SYCL_CHECK(CHECK_TRY_ERROR(
+        stream->memcpy(tiles_dev.get(), tiles.data(), n_tiles * sizeof(ggml_sycl_grouped_gemm_tile))));

    ggml_sycl_pool_alloc<typename S::tsb> packed_b(pool, (size_t) K * max_tiles * S::BN);
    grouped_gemm_pack_b<S>(src1, packed_b.get(), tiles_dev.get(), Npad, (int) K, stream);

-    const typename S::tsb *   packed    = packed_b.get();
-    const ggml_sycl_gg_tile * tiles_ptr = tiles_dev.get();
-    const char *              src0_dd   = (const char *) src0_base;
+    const typename S::tsb *             packed    = packed_b.get();
+    const ggml_sycl_grouped_gemm_tile * tiles_ptr = tiles_dev.get();
+    const char *                        src0_dd   = (const char *) src0_base;
    return fg_visit_type(src0_type, reordered, [&](auto tag) {
        using T = decltype(tag);
        grouped_dequant_gemm_launch<S, typename T::type, T::reordered>(src0_dd, expert_stride, tiles_ptr, packed, dst,
@@ -1073,7 +1188,7 @@ static bool fg_grouped_run(ggml_type src0_type, bool reordered, const void * src
 bool ggml_sycl_grouped_dequant_gemm(ggml_type src0_type, bool reordered, const void * src0_base, size_t expert_stride,
                                    const float * src1, int32_t src1_prec, float * dst,
                                    const int64_t * expert_row_offsets, int64_t n_as, int64_t M, int64_t K,
-                                    int64_t total_rows, std::vector<ggml_sycl_gg_tile> & tiles,
+                                    int64_t total_rows, std::vector<ggml_sycl_grouped_gemm_tile> & tiles,
                                    ggml_sycl_pool & pool, dpct::queue_ptr stream) {
    int64_t n_active = 0;
    for (int64_t e = 0; e < n_as; ++e) {
@@ -1097,3 +1212,227 @@ bool ggml_sycl_grouped_dequant_gemm(ggml_type src0_type, bool reordered, const v
    });
    return launched;
 }
+
+// One work-group sorts routes by expert and builds the tile table.
+static constexpr int MMID_SCHED_WG = 256;
+
+using mmid_sched_atomic = sycl::atomic_ref<uint32_t, sycl::memory_order::relaxed,
+                                           sycl::memory_scope::work_group,
+                                           sycl::access::address_space::local_space>;
+
+bool ggml_sycl_build_mmid_schedule(const int32_t * ids_dev, size_t ids_token_stride, int64_t n_as, int64_t n_ids,
+                                   int64_t n_tokens, int64_t n_tiles_max, mmid_row_mapping * row_mapping,
+                                   ggml_sycl_grouped_gemm_tile * tiles, uint32_t * expert_offsets,
+                                   dpct::queue_ptr stream) {
+    if (n_as <= 0 || n_as > GGML_SYCL_MMID_SCHED_MAX_EXPERTS || n_ids <= 0 || n_tokens <= 0 || n_ids > INT32_MAX / n_tokens) {
+        return false;
+    }
+    const int64_t n_routes = n_tokens * n_ids;
+    if (n_tiles_max < grouped_gemm_max_tiles(n_routes, n_as, GGML_SYCL_FG_BN)) {
+        return false;
+    }
+
+    stream->submit([&](sycl::handler & cgh) {
+        // l_cnt is reused: route counts, then the write cursors, then the tile count per expert
+        sycl::local_accessor<uint32_t, 1> l_cnt(n_as + 1, cgh);
+        sycl::local_accessor<uint32_t, 1> l_off(n_as + 1, cgh);
+        sycl::local_accessor<uint32_t, 1> l_toff(n_as + 1, cgh);
+        cgh.parallel_for(
+            sycl::nd_range<1>(MMID_SCHED_WG, MMID_SCHED_WG),
+            [=](sycl::nd_item<1> item) {
+                const auto     grp = item.get_group();
+                const int64_t  tid = item.get_local_id(0);
+                uint32_t * cnt  = l_cnt.get_multi_ptr<sycl::access::decorated::no>().get();
+                uint32_t * off  = l_off.get_multi_ptr<sycl::access::decorated::no>().get();
+                uint32_t * toff = l_toff.get_multi_ptr<sycl::access::decorated::no>().get();
+
+                for (int64_t e = tid; e < n_as; e += MMID_SCHED_WG) {
+                    cnt[e] = 0;
+                }
+                // a dropped route must leave a row that names (0, 0) rather than stale memory
+                for (int64_t r = tid; r < n_routes; r += MMID_SCHED_WG) {
+                    row_mapping[r] = { 0, 0 };
+                }
+                sycl::group_barrier(grp);
+
+                for (int64_t r = tid; r < n_routes; r += MMID_SCHED_WG) {
+                    const int64_t   token  = r / n_ids;
+                    const int64_t   slot   = r - token * n_ids;
+                    const int32_t * id_row = (const int32_t *) ((const char *) ids_dev + token * ids_token_stride);
+                    const int32_t   e      = id_row[slot];
+                    if (e >= 0 && e < n_as) {
+                        mmid_sched_atomic(cnt[e]).fetch_add(1u);
+                    }
+                }
+                sycl::group_barrier(grp);
+
+                sycl::joint_exclusive_scan(grp, cnt, cnt + n_as, off, 0u, sycl::plus<uint32_t>());
+                sycl::group_barrier(grp);
+                if (tid == 0) {
+                    off[n_as] = off[n_as - 1] + cnt[n_as - 1];
+                }
+                sycl::group_barrier(grp);
+
+                for (int64_t e = tid; e < n_as; e += MMID_SCHED_WG) {
+                    cnt[e] = off[e];
+                }
+                sycl::group_barrier(grp);
+
+                // Atomic arrival order can differ from the host sort; each expert gets the same rows.
+                for (int64_t r = tid; r < n_routes; r += MMID_SCHED_WG) {
+                    const int64_t   token  = r / n_ids;
+                    const int64_t   slot   = r - token * n_ids;
+                    const int32_t * id_row = (const int32_t *) ((const char *) ids_dev + token * ids_token_stride);
+                    const int32_t   e      = id_row[slot];
+                    if (e >= 0 && e < n_as) {
+                        const uint32_t pos = mmid_sched_atomic(cnt[e]).fetch_add(1u);
+                        row_mapping[pos]   = { (int32_t) slot, (int32_t) token };
+                    }
+                }
+                sycl::group_barrier(grp);
+
+                for (int64_t e = tid; e < n_as; e += MMID_SCHED_WG) {
+                    cnt[e] = (off[e + 1] - off[e] + GGML_SYCL_FG_BN - 1) / GGML_SYCL_FG_BN;
+                }
+                sycl::group_barrier(grp);
+
+                sycl::joint_exclusive_scan(grp, cnt, cnt + n_as, toff, 0u, sycl::plus<uint32_t>());
+                sycl::group_barrier(grp);
+                if (tid == 0) {
+                    toff[n_as] = toff[n_as - 1] + cnt[n_as - 1];
+                }
+                sycl::group_barrier(grp);
+
+                const uint32_t n_tiles = toff[n_as];
+                for (int64_t e = tid; e < n_as; e += MMID_SCHED_WG) {
+                    const uint32_t base = off[e];
+                    const uint32_t end  = off[e + 1];
+                    for (uint32_t j = 0; j < cnt[e]; ++j) {
+                        const uint32_t n0 = base + j * GGML_SYCL_FG_BN;
+                        const uint32_t n1 = n0 + GGML_SYCL_FG_BN < end ? n0 + GGML_SYCL_FG_BN : end;
+                        tiles[toff[e] + j] = { (int32_t) e, (int32_t) n0, (int32_t) n1 };
+                    }
+                }
+                // the launch is bounded, so the unused tail must be tiles the GEMM skips
+                for (int64_t t = n_tiles + tid; t < n_tiles_max; t += MMID_SCHED_WG) {
+                    tiles[t] = { 0, 0, 0 };
+                }
+                for (int64_t e = tid; e <= n_as; e += MMID_SCHED_WG) {
+                    expert_offsets[e] = off[e];
+                }
+            });
+    });
+    return true;
+}
+
+template <typename S>
+static bool fg_grouped_device_run(ggml_type src0_type, bool reordered, const void * src0_base, size_t expert_stride,
+                                   ggml_sycl_grouped_gemm_rows src1, ggml_sycl_grouped_gemm_rows dst,
+                                   const int32_t * ids_dev, size_t ids_token_stride, int64_t n_as, int64_t n_ids,
+                                   int64_t n_tokens, int64_t M, int64_t K, ggml_sycl_pool & pool,
+                                   dpct::queue_ptr stream) {
+    if constexpr (S::BN != GGML_SYCL_FG_BN) {
+        return false;
+    } else {
+        const int64_t total_rows = n_tokens * n_ids;
+        const int64_t max_tiles = grouped_gemm_max_tiles(total_rows, n_as, S::BN);
+        if (max_tiles > INT32_MAX / S::BN) {
+            return false;
+        }
+        const int Npad = (int) (max_tiles * S::BN);
+        ggml_sycl_pool_alloc<mmid_row_mapping> rows(pool, total_rows);
+        ggml_sycl_pool_alloc<ggml_sycl_grouped_gemm_tile> tiles(pool, max_tiles);
+        ggml_sycl_pool_alloc<uint32_t> offsets(pool, n_as + 1);
+        if (!ggml_sycl_build_mmid_schedule(ids_dev, ids_token_stride, n_as, n_ids, n_tokens, max_tiles,
+                                           rows.get(), tiles.get(), offsets.get(), stream)) {
+            return false;
+        }
+        src1.map = rows.get();
+        dst.map = rows.get();
+        ggml_sycl_pool_alloc<typename S::tsb> packed_b(pool, (size_t) K * Npad);
+        const ggml_sycl_grouped_gemm_tile * tile_ptr = tiles.get();
+        typename S::tsb * packed = packed_b.get();
+        using E = typename S::EB;
+        constexpr int V = S::VNNI;
+        stream->parallel_for(sycl::range<1>((size_t) Npad * (K / V)), [=](sycl::id<1> id) {
+            const size_t idx = id[0];
+            const int kq = idx / Npad;
+            const int n = idx - (size_t) kq * Npad;
+            const ggml_sycl_grouped_gemm_tile tile = tile_ptr[n / S::BN];
+            const int row = tile.n0 + n % S::BN;
+            typename S::tsb vals[V] = {};
+            if (row < tile.n1) {
+                const mmid_row_mapping mapping = src1.map[row];
+                const float * in = (const float *) (src1.base + (mapping.i1 % src1.ne1) * src1.nb1 + mapping.i2 * src1.nb2) + V * kq;
+#pragma unroll
+                for (int v = 0; v < V; ++v) {
+                    vals[v] = E::cvt(in[v]);
+                }
+            }
+            typename S::tsb * out = packed + ((size_t) kq * Npad + n) * V;
+#pragma unroll
+            for (int v = 0; v < V; ++v) {
+                out[v] = vals[v];
+            }
+        });
+        const int64_t groups_m = (M + S::SG_ROWS - 1) / S::SG_ROWS;
+        return fg_visit_type(src0_type, reordered, [&](auto tag) {
+            using T = decltype(tag);
+            grouped_dequant_gemm_launch<S, typename T::type, T::reordered>((const char *) src0_base, expert_stride,
+                tile_ptr, packed, (float *) dst.base, (int) M, Npad, (int) K, max_tiles, groups_m, stream, dst);
+        });
+    }
+}
+
+bool ggml_sycl_grouped_dequant_gemm_device_supported(ggml_type src0_type, bool reordered, int32_t src1_prec,
+                                                      int64_t n_as, int64_t n_ids, int64_t n_tokens,
+                                                      int64_t M, int64_t K, dpct::queue_ptr stream) {
+    if (n_as <= 0 || n_as > GGML_SYCL_MMID_SCHED_MAX_EXPERTS || n_ids <= 0 || n_tokens <= 0 ||
+        n_ids > INT32_MAX / n_tokens) {
+        return false;
+    }
+    const int64_t total_rows = n_ids * n_tokens;
+    // The active expert count is unknown until the GPU builds the schedule.
+    if (g_ggml_sycl_dynamic_precision == GGML_SYCL_DYNAMIC_PRECISION_F32 ||
+        !ggml_sycl_xmx_gather_type_enabled(src0_type) || !fg_visit_type(src0_type, reordered, [](auto) {}) ||
+        !ggml_sycl_grouped_dequant_gemm_shape_ok(src0_type, M, K, total_rows, std::min(n_as, total_rows),
+                                                (g_ggml_sycl_mmid_sched & GGML_SYCL_MMID_SCHED_UNBOUNDED) != 0)) {
+        return false;
+    }
+    const sycl::device device = stream->get_device();
+    if (device.get_info<sycl::info::device::max_work_group_size>() < MMID_SCHED_WG ||
+        device.get_info<sycl::info::device::local_mem_size>() < 3 * (n_as + 1) * sizeof(uint32_t)) {
+        return false;
+    }
+    if (grouped_gemm_max_tiles(total_rows, n_as, GGML_SYCL_FG_BN) > INT32_MAX / GGML_SYCL_FG_BN) {
+        return false;
+    }
+    const int combo = fg_pick_combo(stream, GGML_TYPE_F32, src1_prec);
+    if (combo < 0) {
+        return false;
+    }
+    // the device schedule is built for GGML_SYCL_FG_BN-wide tiles, see fg_grouped_device_run()
+    bool bn_ok = false;
+    fg_visit_combo(combo, [&](auto s) { bn_ok = decltype(s)::BN == GGML_SYCL_FG_BN; });
+    return bn_ok;
+}
+
+bool ggml_sycl_grouped_dequant_gemm_device(ggml_type src0_type, bool reordered, const void * src0_base,
+                                           size_t expert_stride, const ggml_sycl_grouped_gemm_rows & src1,
+                                           int32_t src1_prec, const ggml_sycl_grouped_gemm_rows & dst,
+                                           const int32_t * ids_dev, size_t ids_token_stride, int64_t n_as,
+                                           int64_t n_ids, int64_t n_tokens, int64_t M, int64_t K, ggml_sycl_pool & pool,
+                                           dpct::queue_ptr stream) {
+    if (src1.ne1 <= 0 || dst.ne1 <= 0 ||
+        !ggml_sycl_grouped_dequant_gemm_device_supported(src0_type, reordered, src1_prec, n_as, n_ids,
+                                                        n_tokens, M, K, stream)) {
+        return false;
+    }
+    const int combo = fg_pick_combo(stream, GGML_TYPE_F32, src1_prec);
+    bool launched = false;
+    fg_visit_combo(combo, [&](auto s) {
+        launched = fg_grouped_device_run<decltype(s)>(src0_type, reordered, src0_base, expert_stride, src1, dst,
+            ids_dev, ids_token_stride, n_as, n_ids, n_tokens, M, K, pool, stream);
+    });
+    return launched;
+}
@@ -4,6 +4,18 @@
 #include "common.hpp"


+// Routed row n is at base + (map[n].i1 % ne1)*nb1 + map[n].i2*nb2.
+struct ggml_sycl_grouped_gemm_rows {
+    char *                   base;
+    const mmid_row_mapping * map;
+    int64_t                  ne1;
+    size_t                   nb1;
+    size_t                   nb2;
+};
+
+static constexpr int GGML_SYCL_FG_BN = 32;
+static constexpr int64_t GGML_SYCL_MMID_SCHED_MAX_EXPERTS = 4096;
+
 // Shape and type gates for the kernels below. Device capability is separate: it needs a queue to ask.
 static constexpr int GGML_SYCL_FG_MAX_N = 64; // widest N taken; each shape covers it in BN-wide tiles

@@ -15,6 +27,8 @@ constexpr bool ggml_sycl_fused_dequant_gemm_type_ok(ggml_type src0_type, int64_t
        return K % 32 == 0;
    }
    const bool superblock =
+           src0_type == GGML_TYPE_Q2_K ||
+           src0_type == GGML_TYPE_Q3_K ||
           src0_type == GGML_TYPE_Q4_K ||
           src0_type == GGML_TYPE_Q5_K ||
           src0_type == GGML_TYPE_Q6_K ||
@@ -39,9 +53,9 @@ constexpr bool ggml_sycl_fused_dequant_gemm_shape_ok(ggml_type src0_type, int64_
 // grouped variant: the per-expert fused kernel is only worth it while each expert is narrow,
 // so wider average slices are left to the per-expert library GEMM loop
 constexpr bool ggml_sycl_grouped_dequant_gemm_shape_ok(ggml_type src0_type, int64_t M, int64_t K,
-                                                       int64_t total_rows, int64_t n_active) {
+                                                       int64_t total_rows, int64_t n_active, bool unbounded = false) {
    return ggml_sycl_fused_dequant_gemm_shape_ok(src0_type, M, 1, K, M) && total_rows > 0 &&
-           total_rows <= INT32_MAX && total_rows <= n_active * GGML_SYCL_FG_MAX_N;
+           total_rows <= INT32_MAX && (unbounded || total_rows <= n_active * GGML_SYCL_FG_MAX_N);
 }

 // Runtime type gate, kept out of the constexpr predicates above so those stay pure.
@@ -57,6 +71,8 @@ inline bool ggml_sycl_xmx_gather_type_enabled(ggml_type src0_type) {
        case GGML_TYPE_IQ1_S:   return (g_ggml_sycl_xmx_gather_types & GGML_SYCL_XMX_GATHER_IQ1_S   ) != 0;
        case GGML_TYPE_IQ1_M:   return (g_ggml_sycl_xmx_gather_types & GGML_SYCL_XMX_GATHER_IQ1_M   ) != 0;
        case GGML_TYPE_Q8_0:    return (g_ggml_sycl_xmx_gather_types & GGML_SYCL_XMX_GATHER_Q8_0    ) != 0;
+        case GGML_TYPE_Q2_K:    return (g_ggml_sycl_xmx_gather_types & GGML_SYCL_XMX_GATHER_Q2_K    ) != 0;
+        case GGML_TYPE_Q3_K:    return (g_ggml_sycl_xmx_gather_types & GGML_SYCL_XMX_GATHER_Q3_K    ) != 0;
        case GGML_TYPE_Q4_K:    return (g_ggml_sycl_xmx_gather_types & GGML_SYCL_XMX_GATHER_Q4_K    ) != 0;
        case GGML_TYPE_Q5_K:    return (g_ggml_sycl_xmx_gather_types & GGML_SYCL_XMX_GATHER_Q5_K    ) != 0;
        case GGML_TYPE_Q6_K:    return (g_ggml_sycl_xmx_gather_types & GGML_SYCL_XMX_GATHER_Q6_K    ) != 0;
@@ -82,7 +98,25 @@ bool ggml_sycl_fused_dequant_gemm(ggml_type src0_type, bool src0_reordered, cons
 bool ggml_sycl_grouped_dequant_gemm(ggml_type src0_type, bool reordered, const void * src0_base, size_t expert_stride,
                                    const float * src1, int32_t src1_prec, float * dst,
                                    const int64_t * expert_row_offsets, int64_t n_as, int64_t M, int64_t K,
-                                    int64_t total_rows, std::vector<ggml_sycl_gg_tile> & tiles,
+                                    int64_t total_rows, std::vector<ggml_sycl_grouped_gemm_tile> & tiles,
                                    ggml_sycl_pool & pool, dpct::queue_ptr stream);

+// Build routing and tile tables on the device; unused tile slots are empty.
+bool ggml_sycl_build_mmid_schedule(const int32_t * ids_dev, size_t ids_token_stride, int64_t n_as, int64_t n_ids,
+                                   int64_t n_tokens, int64_t n_tiles_max, mmid_row_mapping * row_mapping,
+                                   ggml_sycl_grouped_gemm_tile * tiles, uint32_t * expert_offsets,
+                                   dpct::queue_ptr stream);
+
+bool ggml_sycl_grouped_dequant_gemm_device_supported(ggml_type src0_type, bool reordered, int32_t src1_prec,
+                                                      int64_t n_as, int64_t n_ids, int64_t n_tokens,
+                                                      int64_t M, int64_t K, dpct::queue_ptr stream);
+
+// Returns false before queuing work when the type, shape, layout or device is unsupported.
+bool ggml_sycl_grouped_dequant_gemm_device(ggml_type src0_type, bool reordered, const void * src0_base,
+                                           size_t expert_stride, const ggml_sycl_grouped_gemm_rows & src1,
+                                           int32_t src1_prec, const ggml_sycl_grouped_gemm_rows & dst,
+                                           const int32_t * ids_dev, size_t ids_token_stride, int64_t n_as,
+                                           int64_t n_ids, int64_t n_tokens, int64_t M, int64_t K, ggml_sycl_pool & pool,
+                                           dpct::queue_ptr stream);
+
 #endif // GGML_SYCL_FUSED_GEMM_HPP
@@ -118,6 +118,7 @@ int g_ggml_sycl_xmx_glu_max_cols = GGML_SYCL_XMX_GLU_MAX_COLS;
 int g_ggml_sycl_prioritize_dmmv = 0;
 int g_ggml_sycl_xmx_gather_types = GGML_SYCL_XMX_GATHER_TYPES_DEFAULT;
 int g_ggml_sycl_xmx_gather_shapes = GGML_SYCL_XMX_GATHER_SHAPES_DEFAULT;
+int g_ggml_sycl_mmid_sched = GGML_SYCL_MMID_SCHED_DEVICE;
 int g_ggml_sycl_dynamic_precision = GGML_SYCL_DYNAMIC_PRECISION_DEFAULT;
 int g_ggml_sycl_dynamic_required_precision = GGML_SYCL_DYNAMIC_PRECISION_F32;
 static const char * ggml_sycl_dynamic_precision_names[] = { "F16", "BF16", "TF32", "F32" };
@@ -451,6 +452,7 @@ static void ggml_check_sycl() try {
        g_ggml_sycl_prioritize_dmmv = ggml_sycl_get_env("GGML_SYCL_PRIORITIZE_DMMV", 0);
        g_ggml_sycl_graph_eviction_timeout = ggml_sycl_get_env("GGML_SYCL_GRAPH_EVICTION_TIMEOUT", 10);
        g_ggml_sycl_xmx_gather_types = ggml_sycl_get_env("GGML_SYCL_XMX_GATHER_TYPES", GGML_SYCL_XMX_GATHER_TYPES_DEFAULT);
+        g_ggml_sycl_mmid_sched = ggml_sycl_get_env("GGML_SYCL_MMID_SCHED", GGML_SYCL_MMID_SCHED_DEVICE) & 3;
        g_ggml_sycl_xmx_gather_shapes = ggml_sycl_get_env("GGML_SYCL_XMX_GATHER_SHAPES", GGML_SYCL_XMX_GATHER_SHAPES_DEFAULT);
        g_ggml_sycl_dynamic_precision =
            ggml_sycl_get_env_precision("GGML_SYCL_DYNAMIC_PRECISION", GGML_SYCL_DYNAMIC_PRECISION_DEFAULT);
@@ -575,6 +577,7 @@ static void ggml_check_sycl() try {

        GGML_LOG_INFO("  GGML_SYCL_ENABLE_OPT: %d\n", g_ggml_sycl_enable_optimize);
        GGML_LOG_INFO("  GGML_SYCL_XMX_GATHER_TYPES: %d\n", g_ggml_sycl_xmx_gather_types);
+        GGML_LOG_INFO("  GGML_SYCL_MMID_SCHED: %d\n", g_ggml_sycl_mmid_sched);
        GGML_LOG_INFO("  GGML_SYCL_XMX_GATHER_SHAPES: %d\n", g_ggml_sycl_xmx_gather_shapes);
        GGML_LOG_INFO("  GGML_SYCL_DYNAMIC_PRECISION: %s\n",
                      ggml_sycl_dynamic_precision_names[g_ggml_sycl_dynamic_precision]);
@@ -812,7 +815,9 @@ ggml_backend_sycl_buffer_init_tensor(ggml_backend_buffer_t buffer,
        return GGML_STATUS_SUCCESS;
    }

-    if (g_ggml_sycl_enable_optimize) {
+    // a compute buffer copy of a host weight (e.g. -ncmoe experts offloaded for prompt processing) is
+    // refilled in the normal layout before every run, so reordering it is wasted work
+    if (g_ggml_sycl_enable_optimize && ggml_backend_buffer_get_usage(buffer) != GGML_BACKEND_BUFFER_USAGE_COMPUTE) {
        // set reorder extra buffer based on supported type
        switch (tensor->type) {
            case GGML_TYPE_Q4_0:
@@ -3310,7 +3315,7 @@ inline void ggml_sycl_op_mul_mat_sycl(
    // the f16 route converts src1 to f16 [TAG_GGML_PREC]
    use_fp16 = use_fp16 && ggml_sycl_src1_f16_ok(dst);
    if ((src0->type == GGML_TYPE_F16 || ggml_is_quantized(src0->type)) && use_fp16 && ggml_is_contiguous(src0) &&
-        row_diff == src0->ne[1] && dst->op_params[0] == GGML_PREC_DEFAULT) {
+        row_diff == src0->ne[1] && ggml_prec_get_acc(dst) == GGML_PREC_UNDEFINED) {
        ggml_sycl_pool_alloc<sycl::half> src1_as_f16(ctx.pool());
        if (src1->type != GGML_TYPE_F16) {
            scope_op_debug_print scope_dbg_print(__func__, "/to_fp16_sycl", dst, /*num_src=*/2,
@@ -4474,36 +4479,6 @@ static bool ggml_sycl_supports_dmmv(enum ggml_type type) {
    }
 }

-static bool ggml_sycl_supports_mmvq_id(enum ggml_type type) {
-    switch (type) {
-        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_Q2_0:
-        case GGML_TYPE_Q2_K:
-        case GGML_TYPE_Q3_K:
-        case GGML_TYPE_Q4_K:
-        case GGML_TYPE_Q5_K:
-        case GGML_TYPE_Q6_K:
-        case GGML_TYPE_MXFP4:
-        case GGML_TYPE_NVFP4:
-        case GGML_TYPE_IQ2_XXS:
-        case GGML_TYPE_IQ2_XS:
-        case GGML_TYPE_IQ2_S:
-        case GGML_TYPE_IQ3_XXS:
-        case GGML_TYPE_IQ3_S:
-        case GGML_TYPE_IQ1_S:
-        case GGML_TYPE_IQ1_M:
-        case GGML_TYPE_IQ4_NL:
-        case GGML_TYPE_IQ4_XS:
-            return true;
-        default:
-            return false;
-    }
-}
-
 // Helper functions to unify device memory allocation for both async and sync paths
 static inline void * sycl_ext_malloc_device(dpct::queue_ptr stream, size_t size) {
    bool use_async = g_ggml_sycl_use_async_mem_op;
@@ -4979,6 +4954,16 @@ static bool reorder_qw_q3_k(uint8_t * data_device, size_t size, size_t offset, d
        }, stream);
 }

+// Reorder each Q3_K expert slice on its own.
+static bool reorder_qw_q3_k_moe(uint8_t * data_device, size_t expert_bytes, int64_t n_expert, dpct::queue_ptr stream) {
+    for (int64_t e = 0; e < n_expert; ++e) {
+        if (!reorder_qw_q3_k(data_device + e * expert_bytes, expert_bytes, 0, stream)) {
+            return false;
+        }
+    }
+    return true;
+}
+
 static bool reorder_qw_q5_k(uint8_t * data_device, size_t size, size_t offset, dpct::queue_ptr stream) {
    GGML_ASSERT(offset == 0);
    return reorder_qw_planar<block_q5_K>(__func__, data_device, size, offsetof(block_q5_K, qs), QK_K / 2,
@@ -5050,6 +5035,8 @@ static bool reorder_qw(const ggml_tensor * src0, dpct::queue_ptr stream) {
        switch (src0->type) {
            case GGML_TYPE_Q8_0:
                return reorder_qw_q8_0_moe(data_device, src0->nb[2], src0->ne[2], stream);
+            case GGML_TYPE_Q3_K:
+                return reorder_qw_q3_k_moe(data_device, src0->nb[2], src0->ne[2], stream);
            case GGML_TYPE_Q4_K:
                return reorder_qw_q4_k_moe(data_device, src0->nb[2], src0->ne[2], stream);
            case GGML_TYPE_Q5_K:
@@ -5135,7 +5122,7 @@ static void opt_for_reorder_id(ggml_backend_sycl_context * ctx, const ggml_tenso
    if (!g_ggml_sycl_enable_optimize || !ctx->opt_feature.reorder) {
        return;
    }
-    if (src0->type != GGML_TYPE_Q8_0 && src0->type != GGML_TYPE_Q4_K &&
+    if (src0->type != GGML_TYPE_Q8_0 && src0->type != GGML_TYPE_Q3_K && src0->type != GGML_TYPE_Q4_K &&
        src0->type != GGML_TYPE_Q5_K && src0->type != GGML_TYPE_Q6_K && src0->type != GGML_TYPE_MXFP4) {
        return;
    }
@@ -5143,6 +5130,10 @@ static void opt_for_reorder_id(ggml_backend_sycl_context * ctx, const ggml_tenso
    if (src0->type == GGML_TYPE_MXFP4 && (src0->nb[2] % 16 != 0 || (uintptr_t) src0->data % 16 != 0)) {
        return;
    }
+    // The grouped GEMM loads reordered Q3_K qs and hmask as 16-byte vectors.
+    if (src0->type == GGML_TYPE_Q3_K && ((src0->nb[2] / sizeof(block_q3_K)) % 16 != 0 || (uintptr_t) src0->data % 16 != 0)) {
+        return;
+    }
    ggml_tensor_extra_gpu * extra = static_cast<ggml_tensor_extra_gpu *>(src0->extra);
    if (!extra || extra->optimized_feature.reorder) {
        return;
@@ -5641,26 +5632,31 @@ __dpct_inline__ static void k_zero_dst_rows(
    }
 }

-// returns true when ggml_sycl_mul_mat_id takes the fallback path that requires stream synchronization
-static bool ggml_sycl_mul_mat_id_needs_sync(const ggml_tensor * dst) {
-    const ggml_tensor * src0 = dst->src[0];
-    const ggml_tensor * src1 = dst->src[1];
-    const ggml_tensor * ids  = dst->src[2];
-
+// layout checks for the fused MMVQ path; the type check is ggml_sycl_mul_mat_vec_q_id_supports_type
+static bool ggml_sycl_mul_mat_id_mmvq_eligible(
+    const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * ids, const ggml_tensor * dst) {
    const int64_t ne10 = src1->ne[0];
    const int64_t ne11 = src1->ne[1];
    const int64_t ne12 = src1->ne[2];

-    if (ne12 != 1) return true;
-    if (src1->type != GGML_TYPE_F32 || dst->type != GGML_TYPE_F32) return true;
-    if (ne10 != src0->ne[0] || ne10 % QK8_1 != 0) return true;
-    if (!ggml_is_contiguous(src1)) return true;
+    if (ne12 != 1) return false;
+    if (src1->type != GGML_TYPE_F32 || dst->type != GGML_TYPE_F32) return false;
+    if (ne10 != src0->ne[0] || ne10 % QK8_1 != 0) return false;
+    if (!ggml_is_contiguous(src1)) return false;

    const int64_t n_ids_per_group = ids->ne[0];
-    if (ids->ne[1] != 1) return true;
-    if (ne11 != 1 && ne11 != n_ids_per_group) return true;
+    if (ids->type != GGML_TYPE_I32 || ids->nb[0] != sizeof(int32_t)) return false;
+    if (ids->ne[1] != 1) return false;
+    if (ne11 != 1 && ne11 != n_ids_per_group) return false;

-    return !ggml_sycl_supports_mmvq_id(src0->type);
+    return true;
+}
+
+// returns true when ggml_sycl_mul_mat_id takes the fused MMVQ path
+static bool ggml_sycl_mul_mat_id_should_use_mmvq(const ggml_tensor * dst) {
+    const ggml_tensor * src0 = dst->src[0];
+    return ggml_sycl_mul_mat_id_mmvq_eligible(src0, dst->src[1], dst->src[2], dst) &&
+           ggml_sycl_mul_mat_vec_q_id_supports_type(src0->type);
 }

 // Fused MoE TG fast path. Returns false to fall back to the per-expert loop below.
@@ -5668,13 +5664,11 @@ static bool ggml_sycl_mul_mat_id_mmvq_fused(
    ggml_backend_sycl_context & ctx, const ggml_tensor * src0,
    const ggml_tensor * src1, const ggml_tensor * ids, ggml_tensor * dst)
 {
-    if (ggml_sycl_mul_mat_id_needs_sync(dst)) return false;
+    if (!ggml_sycl_mul_mat_id_should_use_mmvq(dst)) return false;

    const int64_t ne10 = src1->ne[0];
    const int64_t ne11 = src1->ne[1];
-
    const int64_t n_ids_per_group = ids->ne[0];
-
    const queue_ptr stream           = ctx.stream();
    const int       src1_padded_cols = GGML_PAD((int) ne10, MATRIX_ROW_PADDING);
    const int       n_experts_used   = (int) n_ids_per_group;
@@ -5772,6 +5766,18 @@ static void mmid_counting_sort_rows(
    routed_row_src.resize(expert_row_offsets[n_as]);
 }

+// true when the device MoE schedule may handle this MUL_MAT_ID: enabled, prefill, and an
+// f32/i32 layout it accepts; ggml_sycl_grouped_dequant_gemm_device decides type and width
+static bool ggml_sycl_mul_mat_id_device_sched_eligible(const ggml_tensor * dst) {
+    const ggml_tensor * src0 = dst->src[0];
+    const ggml_tensor * src1 = dst->src[1];
+    const ggml_tensor * ids = dst->src[2];
+    return (g_ggml_sycl_mmid_sched & GGML_SYCL_MMID_SCHED_DEVICE) && src1->ne[2] != 1 && ggml_is_contiguous(src0) &&
+           ids->type == GGML_TYPE_I32 && ids->nb[0] == sizeof(int32_t) &&
+           src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32 &&
+           src1->nb[1] == sizeof(float)*src1->ne[0] && dst->nb[1] == sizeof(float)*dst->ne[0];
+}
+
 static void ggml_sycl_mul_mat_id(ggml_backend_sycl_context & ctx,
                                 ggml_tensor *dst) try {
    scope_op_debug_print scope_dbg_print(__func__, dst, /*num_src=*/3);
@@ -5792,13 +5798,46 @@ static void ggml_sycl_mul_mat_id(ggml_backend_sycl_context & ctx,
            return;
        }
    }
-    GGML_ASSERT(ggml_sycl_mul_mat_id_needs_sync(dst));
+    GGML_ASSERT(!ggml_sycl_mul_mat_id_should_use_mmvq(dst));

    // The per-expert loop below reads the experts in whatever layout they have: reorder MXFP4 here as well, so prompt processing does not depend on a single-token decode having run first.
    if (src0->type == GGML_TYPE_MXFP4) {
        opt_for_reorder_id(&ctx, src0);
    }

+    if (ggml_sycl_mul_mat_id_device_sched_eligible(dst)) {
+        if (g_ggml_sycl_dynamic_precision != GGML_SYCL_DYNAMIC_PRECISION_F32 &&
+            ggml_sycl_xmx_gather_type_enabled(src0->type)) {
+            opt_for_reorder_id(&ctx, src0);
+        }
+        const bool reordered = src0->extra && ((const ggml_tensor_extra_gpu *) src0->extra)->optimized_feature.reorder;
+        const ggml_sycl_grouped_gemm_rows src1_rows = { (char *) src1->data, nullptr, ne11, nb11, nb12 };
+        const ggml_sycl_grouped_gemm_rows dst_rows = { (char *) dst->data, nullptr, ne1, nb1, nb2 };
+        if (ggml_sycl_grouped_dequant_gemm_device(src0->type, reordered, src0->data, nb02, src1_rows,
+                ggml_sycl_src1_prec(dst), dst_rows, (const int32_t *) ids->data, ids->nb[1], n_as, n_ids,
+                ids->ne[1], ne01, ne10, ctx.pool(), stream)) {
+            // the device schedule drops skipped slots (id == -1), so their dst rows are zeroed here
+            const char * ids_data = (const char *) ids->data;
+            const size_t ids_nb1  = ids->nb[1];
+            char *       dst_data = (char *) dst->data;
+            const int    nth      = (int) std::min<int64_t>(ne0, 256);
+            stream->parallel_for(
+                sycl::nd_range<2>(sycl::range<2>(ids->ne[1], n_ids * nth), sycl::range<2>(1, nth)),
+                [=](sycl::nd_item<2> item) {
+                    const int64_t iid1 = item.get_group(0);
+                    const int64_t id   = item.get_group(1);
+                    if (*(const int32_t *) (ids_data + iid1*ids_nb1 + id*sizeof(int32_t)) != -1) {
+                        return;
+                    }
+                    float * row = (float *) (dst_data + id*nb1 + iid1*nb2);
+                    for (int64_t j = item.get_local_id(1); j < ne0; j += nth) {
+                        row[j] = 0.0f;
+                    }
+                });
+            return;
+        }
+    }
+
    std::vector<char> ids_host(ggml_nbytes(ids));
    const char * ids_dev = (const char *) ids->data;

@@ -6890,7 +6929,7 @@ static bool mul_mat_uses_onednn(int device, ggml_tensor * dst) {
 #endif
 #ifdef GGML_SYCL_F16
    if ((src0->type == GGML_TYPE_F16 || ggml_is_quantized(src0->type)) && ggml_is_contiguous(src0) &&
-        dst->op_params[0] == GGML_PREC_DEFAULT) {
+        ggml_prec_get_acc(dst) == GGML_PREC_UNDEFINED) {
        return true;
    }
 #endif
@@ -6904,6 +6943,54 @@ static bool mul_mat_uses_onednn(int device, ggml_tensor * dst) {
 #endif
 }

+// returns true when ggml_sycl_mul_mat_id takes the fallback path that requires stream synchronization:
+// neither the fused MMVQ path nor the device-scheduled grouped GEMM (GGML_SYCL_MMID_SCHED) handles dst
+static bool ggml_sycl_mul_mat_id_needs_sync(ggml_backend_sycl_context * ctx, const ggml_tensor * dst) {
+    const ggml_tensor * src0 = dst->src[0];
+    const ggml_tensor * src1 = dst->src[1];
+    const ggml_tensor * ids = dst->src[2];
+    if (!src0 || !src1 || !ids || ggml_backend_buffer_is_sycl_split(src0->buffer)) {
+        return true;
+    }
+    const bool reordered = src0->extra && ((const ggml_tensor_extra_gpu *) src0->extra)->optimized_feature.reorder;
+    if (ggml_sycl_mul_mat_id_mmvq_eligible(src0, src1, ids, dst)) {
+        return reordered ? !ggml_sycl_mul_mat_vec_q_id_reorder_supports_type(src0->type) :
+                           !ggml_sycl_mul_mat_vec_q_id_supports_type(src0->type);
+    }
+    return !ggml_sycl_mul_mat_id_device_sched_eligible(dst) ||
+           !ggml_sycl_grouped_dequant_gemm_device_supported(src0->type, reordered, ggml_sycl_src1_prec(dst),
+               src0->ne[2], ids->ne[0], ids->ne[1], src0->ne[1], src0->ne[0], ctx->stream());
+}
+
+// mirrors the paths in ggml_sycl_mul_mat_id() that call opt_for_reorder_id()
+static bool mul_mat_id_reorders_src0(const ggml_tensor * dst) {
+    const ggml_tensor * src0 = dst->src[0];
+    if (ggml_sycl_mul_mat_id_should_use_mmvq(dst) || src0->type == GGML_TYPE_MXFP4) {
+        return true;
+    }
+    return ggml_sycl_mul_mat_id_device_sched_eligible(dst) &&
+           g_ggml_sycl_dynamic_precision != GGML_SYCL_DYNAMIC_PRECISION_F32 &&
+           ggml_sycl_xmx_gather_type_enabled(src0->type);
+}
+
+// true when the node and its SYCL sources live in this device's buffers (no split or other-device buffers)
+static bool graph_node_is_in_local_device(ggml_backend_buffer_type_t device_buft, const ggml_tensor * node) {
+    if (!node->buffer || node->buffer->buft != device_buft) {
+        return false;
+    }
+    for (int j = 0; j < GGML_MAX_SRC; ++j) {
+        const ggml_tensor * src = node->src[j];
+        if (!src || !src->buffer) {
+            continue;
+        }
+        if (ggml_backend_buffer_is_sycl_split(src->buffer) ||
+            (ggml_backend_buffer_is_sycl(src->buffer) && src->buffer->buft != device_buft)) {
+            return false;
+        }
+    }
+    return true;
+}
+
 static bool check_graph_compatibility(ggml_backend_sycl_context * ctx, ggml_cgraph * cgraph) {
    if (ggml_sycl_info().active_device_count > 1) {
        // A sycl_ex::command_graph object can only be created for a single device
@@ -6911,6 +6998,8 @@ static bool check_graph_compatibility(ggml_backend_sycl_context * ctx, ggml_cgra
        return false;
    }

+    const ggml_backend_buffer_type_t device_buft = ggml_backend_sycl_buffer_type(ctx->device);
+
    for (int i = 0; i < cgraph->n_nodes; i++) {
        ggml_tensor * node = cgraph->nodes[i];

@@ -6921,6 +7010,18 @@ static bool check_graph_compatibility(ggml_backend_sycl_context * ctx, ggml_cgra

        const ggml_op node_op = node->op;

+        if (!graph_node_is_in_local_device(device_buft, node)) {
+            GGML_LOG_DEBUG("%s: disabling SYCL graphs due to node %s not being device local\n", __func__,
+                           ggml_op_name(node_op));
+            return false;
+        }
+
+        // Reorder expert weights before recording, so the reorder state that
+        // ggml_sycl_mul_mat_id_needs_sync() sees is the one the recorded graph runs with.
+        if (node_op == GGML_OP_MUL_MAT_ID && mul_mat_id_reorders_src0(node)) {
+            opt_for_reorder_id(ctx, node->src[0]);
+        }
+
        switch (node_op) {
            default:
                break;
@@ -6928,7 +7029,7 @@ static bool check_graph_compatibility(ggml_backend_sycl_context * ctx, ggml_cgra
                // The ggml_sycl_mul_mat_id() fallback path does a blocking host wait on the sycl
                // queue after submitting a memcpy operation, but wait() can't be called on a queue
                // that is recording to a graph.
-                if (ggml_sycl_mul_mat_id_needs_sync(node)) {
+                if (ggml_sycl_mul_mat_id_needs_sync(ctx, node)) {
                    GGML_LOG_DEBUG("%s: disabling SYCL graphs due to host wait in %s\n", __func__,
                                   ggml_op_name(node_op));
                    return false;
@@ -6968,6 +7069,13 @@ static bool check_graph_compatibility(ggml_backend_sycl_context * ctx, ggml_cgra
                }
 #endif
                break;
+            case GGML_OP_OUT_PROD:
+            case GGML_OP_CONV_3D:
+            case GGML_OP_SOLVE_TRI:
+                // these call oneMKL and wait on the host
+                GGML_LOG_DEBUG("%s: disabling SYCL graphs due to unsupported node type %s\n", __func__,
+                               ggml_op_name(node_op));
+                return false;
        }
    }
    return true;
@@ -7321,6 +7429,11 @@ static bool do_ggml_backend_sycl_device_supports_op(ggml_backend_dev_t dev, cons
                    return false;
                }

+                if (op->op == GGML_OP_MUL_MAT &&
+                    (ggml_prec_get_src(op, 1) == GGML_PREC_BF16 || ggml_prec_get_src(op, 1) == GGML_PREC_F32)) {
+                    return false;
+                }
+
                return true;
            }
        case GGML_OP_OUT_PROD:
@@ -2706,6 +2706,50 @@ static void launch_mul_mat_vec_q_moe(
    });
 }

+bool ggml_sycl_mul_mat_vec_q_id_supports_type(enum ggml_type src0_type) {
+    switch (src0_type) {
+        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_Q2_0:
+        case GGML_TYPE_Q2_K:
+        case GGML_TYPE_Q3_K:
+        case GGML_TYPE_Q4_K:
+        case GGML_TYPE_Q5_K:
+        case GGML_TYPE_Q6_K:
+        case GGML_TYPE_MXFP4:
+        case GGML_TYPE_NVFP4:
+        case GGML_TYPE_IQ2_XXS:
+        case GGML_TYPE_IQ2_XS:
+        case GGML_TYPE_IQ2_S:
+        case GGML_TYPE_IQ3_XXS:
+        case GGML_TYPE_IQ3_S:
+        case GGML_TYPE_IQ1_S:
+        case GGML_TYPE_IQ1_M:
+        case GGML_TYPE_IQ4_NL:
+        case GGML_TYPE_IQ4_XS:
+            return true;
+        default:
+            return false;
+    }
+}
+
+bool ggml_sycl_mul_mat_vec_q_id_reorder_supports_type(enum ggml_type src0_type) {
+    switch (src0_type) {
+        case GGML_TYPE_Q8_0:
+        case GGML_TYPE_Q3_K:
+        case GGML_TYPE_Q4_K:
+        case GGML_TYPE_Q5_K:
+        case GGML_TYPE_Q6_K:
+        case GGML_TYPE_MXFP4:
+            return true;
+        default:
+            return false;
+    }
+}
+
 bool ggml_sycl_mul_mat_vec_q_id(
    enum ggml_type     src0_type,
    const void *       vx_base,
@@ -2719,7 +2763,9 @@ bool ggml_sycl_mul_mat_vec_q_id(
    size_t             dst_row_stride,
    size_t             src1_row_stride,
    dpct::queue_ptr    stream) {
-    // if adding a new type, ensure it is also added to ggml_sycl_supports_mmvq_id
+    if (!ggml_sycl_mul_mat_vec_q_id_supports_type(src0_type)) {
+        return false;
+    }
    switch (src0_type) {
        case GGML_TYPE_Q4_0:
            launch_mul_mat_vec_q_moe<QK4_0, QI4_0, block_q4_0, VDR_Q4_0_Q8_1_MMVQ, vec_dot_q4_0_q8_1>(
@@ -2952,12 +2998,20 @@ bool ggml_sycl_mul_mat_vec_q_id_reorder(
    size_t             dst_row_stride,
    size_t             src1_row_stride,
    dpct::queue_ptr    stream) {
+    if (!ggml_sycl_mul_mat_vec_q_id_reorder_supports_type(src0_type)) {
+        return false;
+    }
    switch (src0_type) {
        case GGML_TYPE_Q8_0:
            launch_mul_mat_vec_q_moe_reorder<reorder_vec_dot_q_sycl<GGML_TYPE_Q8_0>>(
                vx_base, vy, ids_dev, dst_base, ncols, nrows, n_experts_used,
                expert_weight_stride, dst_row_stride, src1_row_stride, stream);
            return true;
+        case GGML_TYPE_Q3_K:
+            launch_mul_mat_vec_q_moe_reorder<reorder_vec_dot_q_sycl<GGML_TYPE_Q3_K>>(
+                vx_base, vy, ids_dev, dst_base, ncols, nrows, n_experts_used,
+                expert_weight_stride, dst_row_stride, src1_row_stride, stream);
+            return true;
        case GGML_TYPE_Q4_K:
            launch_mul_mat_vec_q_moe_reorder<reorder_vec_dot_q_sycl<GGML_TYPE_Q4_K>>(
                vx_base, vy, ids_dev, dst_base, ncols, nrows, n_experts_used,
@@ -117,4 +117,7 @@ bool ggml_sycl_mul_mat_vec_q_glu_plain(
    int                stride_col_dst,       // floats between output columns in dst
    dpct::queue_ptr    stream);

+bool ggml_sycl_mul_mat_vec_q_id_supports_type(enum ggml_type src0_type);
+bool ggml_sycl_mul_mat_vec_q_id_reorder_supports_type(enum ggml_type src0_type);
+
 #endif // GGML_SYCL_MMVQ_HPP
@@ -18,7 +18,7 @@ static bool ggml_vk_comm_needs_f32(const ggml_tensor * t, int depth = 0) {
    switch (t->op) {
        case GGML_OP_MUL_MAT:
        case GGML_OP_MUL_MAT_ID:
-            return ggml_get_op_params_i32(t, 0) == GGML_PREC_BF16 || ggml_get_op_params_i32(t, 0) == GGML_PREC_F32;
+            return ggml_prec_get_acc(t) == GGML_PREC_BF16 || ggml_prec_get_acc(t) == GGML_PREC_F32;
        case GGML_OP_ADD:
            return ggml_vk_comm_needs_f32(t->src[0], depth + 1) || ggml_vk_comm_needs_f32(t->src[1], depth + 1);
        case GGML_OP_MUL:
@@ -829,7 +829,7 @@ struct vk_op_upscale_push_constants {
 struct vk_op_sum_rows_push_constants
 {
    uint32_t n_cols;
-    uint32_t ne01, ne02;
+    uint32_t ne01, ne02, ne03;
    uint32_t nb01, nb02, nb03;
    uint32_t nb11, nb12, nb13;
    float weight;
@@ -844,6 +844,7 @@ static vk_op_sum_rows_push_constants vk_op_sum_rows_push_constants_init(const gg
    p.n_cols = (uint32_t)n_cols;
    p.ne01 = (uint32_t)src->ne[1];
    p.ne02 = (uint32_t)src->ne[2];
+    p.ne03 = (uint32_t)src->ne[3];
    p.nb01 = (uint32_t)src->nb[1] / type_size;
    p.nb02 = (uint32_t)src->nb[2] / type_size;
    p.nb03 = (uint32_t)src->nb[3] / type_size;
@@ -837,6 +837,7 @@ struct vk_device_struct {
    vk_pipeline pipeline_add_id_f32;

    vk_pipeline pipeline_concat_i8, pipeline_concat_i16, pipeline_concat_i32, pipeline_concat_i64;
+    vk_pipeline pipeline_concat_transpose_i8, pipeline_concat_transpose_i16, pipeline_concat_transpose_i32;
    vk_pipeline pipeline_upscale_nearest_f32, pipeline_upscale_bilinear_f32, pipeline_upscale_bicubic_f32, pipeline_upscale_bilinear_antialias_f32;
    vk_pipeline pipeline_scale_f32;
    vk_pipeline pipeline_log[2];
@@ -1260,6 +1260,12 @@ static vk_fa_tuning_params get_fa_tuning_params_scalar(const vk_device& device,
        result.block_rows /= 2;
    }

+    // With head sizes above 128, 16 rows per workgroup spill registers on Qualcomm and run about
+    // 40x slower; at large KV a single dispatch can then exceed the Windows GPU timeout.
+    if (device->vendor_id == VK_VENDOR_ID_QUALCOMM && (hsk > 128 || hsv > 128)) {
+        result.block_rows = std::min(result.block_rows, 8u);
+    }
+
    // On AMD RDNA, for small head sizes and big batch size the shader uses few registers, so too many subgroups get scheduled
    // at once and end up thrashing the cache. Fix this by setting a large (unused) shmem buffer that reduces occupancy.
    // This targets an occupancy of 4 subgroups per SIMD.
@@ -3271,7 +3277,7 @@ void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
    }
    ggml_vk_create_pipeline(device, device->pipeline_mul_mat_vec_nc_f16_f32, "mul_mat_vec_nc_f16_f32", mul_mat_vec_nc_f16_f32_len, mul_mat_vec_nc_f16_f32_data, "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_nc_push_constants), {1, 1, 1}, {}, 1);

-    ggml_vk_create_pipeline(device, device->pipeline_norm_f32, "norm_f32", norm_f32_len, norm_f32_data, "main", 2, sizeof(vk_op_unary_push_constants), {1, 1, 1}, {}, 1);
+    ggml_vk_create_pipeline(device, device->pipeline_norm_f32, "norm_f32", norm_f32_len, norm_f32_data, "main", 2, sizeof(vk_op_unary_push_constants), {1, 1, 1}, {}, 1, true);
    ggml_vk_create_pipeline(device, device->pipeline_group_norm_f32, "group_norm_f32", group_norm_f32_len, group_norm_f32_data, "main", 2, sizeof(vk_op_push_constants), {1, 1, 1}, {}, 1);

    const bool rms_subgroups = (device->vendor_id == VK_VENDOR_ID_INTEL || device->vendor_id == VK_VENDOR_ID_NVIDIA || device->vendor_id == VK_VENDOR_ID_AMD) && device->subgroup_arithmetic && device->subgroup_require_full_support;
@@ -3304,7 +3310,7 @@ void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
 #undef RMS_SHADER

    ggml_vk_create_pipeline(device, device->pipeline_rms_norm_back_f32, "rms_norm_back_f32", rms_norm_back_f32_len, rms_norm_back_f32_data, "main", 3, sizeof(vk_op_push_constants), {1, 1, 1}, {}, 1);
-    ggml_vk_create_pipeline(device, device->pipeline_l2_norm_f32, "l2_norm_f32", l2_norm_f32_len, l2_norm_f32_data, "main", 2, sizeof(vk_op_unary_push_constants), {1, 1, 1}, {}, 1);
+    ggml_vk_create_pipeline(device, device->pipeline_l2_norm_f32, "l2_norm_f32", l2_norm_f32_len, l2_norm_f32_data, "main", 2, sizeof(vk_op_unary_push_constants), {1, 1, 1}, {}, 1, true);

    ggml_vk_create_pipeline(device, device->pipeline_cpy_32, "cpy_32", cpy_32_len, cpy_32_data, "main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1);
    ggml_vk_create_pipeline(device, device->pipeline_cpy_f32_f16, "cpy_f32_f16", cpy_f32_f16_len, cpy_f32_f16_data, "main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1);
@@ -3410,6 +3416,9 @@ void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
    ggml_vk_create_pipeline(device, device->pipeline_concat_i8, "concat_i8", concat_i8_len, concat_i8_data, "main", 3, sizeof(vk_op_binary_push_constants), {512, 1, 1}, {}, 1);
    ggml_vk_create_pipeline(device, device->pipeline_concat_i16, "concat_i16", concat_i16_len, concat_i16_data, "main", 3, sizeof(vk_op_binary_push_constants), {512, 1, 1}, {}, 1);
    ggml_vk_create_pipeline(device, device->pipeline_concat_i32, "concat_i32", concat_i32_len, concat_i32_data, "main", 3, sizeof(vk_op_binary_push_constants), {512, 1, 1}, {}, 1);
+    ggml_vk_create_pipeline(device, device->pipeline_concat_transpose_i8, "concat_transpose_i8", concat_transpose_i8_len, concat_transpose_i8_data, "main", 3, sizeof(vk_op_binary_push_constants), {32, 32, 1}, {}, 1);
+    ggml_vk_create_pipeline(device, device->pipeline_concat_transpose_i16, "concat_transpose_i16", concat_transpose_i16_len, concat_transpose_i16_data, "main", 3, sizeof(vk_op_binary_push_constants), {32, 32, 1}, {}, 1);
+    ggml_vk_create_pipeline(device, device->pipeline_concat_transpose_i32, "concat_transpose_i32", concat_transpose_i32_len, concat_transpose_i32_data, "main", 3, sizeof(vk_op_binary_push_constants), {32, 32, 1}, {}, 1);
    ggml_vk_create_pipeline(device, device->pipeline_concat_i64, "concat_i64", concat_i64_len, concat_i64_data, "main", 3, sizeof(vk_op_binary_push_constants), {512, 1, 1}, {}, 1);

    ggml_vk_create_pipeline(device, device->pipeline_upscale_nearest_f32, "upscale_f32", upscale_f32_len, upscale_f32_data, "main", 2, sizeof(vk_op_upscale_push_constants), {512, 1, 1}, {GGML_SCALE_MODE_NEAREST}, 1);
@@ -3584,7 +3593,7 @@ void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {

    ggml_vk_create_pipeline(device, device->pipeline_argmax_f32, "argmax_f32", argmax_f32_len, argmax_f32_data, "main", 2, sizeof(vk_op_push_constants), {1, 1, 1}, { device->subgroup_size }, 1);

-    ggml_vk_create_pipeline(device, device->pipeline_sum_rows_f32, "sum_rows_f32", sum_rows_f32_len, sum_rows_f32_data, "main", 2, sizeof(vk_op_sum_rows_push_constants), {1, 1, 1}, { device->subgroup_size }, 1);
+    ggml_vk_create_pipeline(device, device->pipeline_sum_rows_f32, "sum_rows_f32", sum_rows_f32_len, sum_rows_f32_data, "main", 2, sizeof(vk_op_sum_rows_push_constants), {1, 1, 1}, { device->subgroup_size }, 1, true);
    ggml_vk_create_pipeline(device, device->pipeline_cross_entropy_loss_f32, "cross_entropy_loss_f32", cross_entropy_loss_f32_len, cross_entropy_loss_f32_data, "main", 3, sizeof(vk_op_push_constants), {1, 1, 1}, { device->subgroup_size }, 1);
    ggml_vk_create_pipeline(device, device->pipeline_cross_entropy_loss_f32_wg512, "cross_entropy_loss_f32_wg512", cross_entropy_loss_f32_len, cross_entropy_loss_f32_data, "main", 3, sizeof(vk_op_push_constants), {1, 1, 1}, { 512 }, 1);
    ggml_vk_create_pipeline(device, device->pipeline_cross_entropy_loss_back_f32, "cross_entropy_loss_back_f32", cross_entropy_loss_back_f32_len, cross_entropy_loss_back_f32_data, "main", 4, sizeof(vk_op_push_constants), {1, 1, 1}, { device->subgroup_size }, 1);
@@ -8847,6 +8856,26 @@ static vk_pipeline ggml_vk_rms_norm_pipeline(const std::array<vk_pipeline, 3> &
    return pipelines[2];
 }

+static bool ggml_vk_concat_is_transposed(const vk_device & device, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * dst) {
+    if (ggml_get_op_params_i32(dst, 0) != 0) {           // dim 0 only
+        return false;
+    }
+    if (src0->ne[2] != 1 || src0->ne[3] != 1 || src1->ne[2] != 1 || src1->ne[3] != 1) {
+        return false;
+    }
+    const size_t ts = ggml_type_size(src0->type);
+    if (src0->nb[0] != ts || dst->nb[0] != ts) {          // src0 and dst rows must be contiguous
+        return false;
+    }
+    if (src1->nb[0] <= src1->nb[1]) {                     // src1 must actually be transposed
+        return false;
+    }
+    if (CEIL_DIV(src1->ne[1], 32) > device->properties.limits.maxComputeWorkGroupCount[0]) {
+        return false;
+    }
+    return src0->ne[1] == src1->ne[1] && dst->ne[1] == src1->ne[1];
+}
+
 static vk_pipeline ggml_vk_op_get_pipeline(ggml_backend_vk_context * ctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * src2, const ggml_tensor * dst, ggml_op op) {
    switch (op) {
    case GGML_OP_GET_ROWS:
@@ -8939,6 +8968,18 @@ static vk_pipeline ggml_vk_op_get_pipeline(ggml_backend_vk_context * ctx, const
        if (!ggml_vk_concat_supported(src0, src1, dst)) {
            return nullptr;
        }
+        if (!ggml_is_quantized(src0->type) && ggml_vk_concat_is_transposed(ctx->device, src0, src1, dst)) {
+            switch (ggml_vk_concat_unit_size(src0->type)) {
+            case 1:
+                return ctx->device->pipeline_concat_transpose_i8;
+            case 2:
+                return ctx->device->pipeline_concat_transpose_i16;
+            case 4:
+                return ctx->device->pipeline_concat_transpose_i32;
+            default:
+                break;
+            }
+        }
        switch (ggml_vk_concat_unit_size(src0->type)) {
        case 1:
            return ctx->device->pipeline_concat_i8;
@@ -9706,8 +9747,13 @@ static void ggml_vk_op_f32(ggml_backend_vk_context * ctx, vk_context& subctx, co
        break;
    case GGML_OP_GROUP_NORM:
        {
-            const uint32_t num_groups = dst->op_params[0];
-            elements = { num_groups * (uint32_t)src0->ne[3], 1, 1 };
+            const uint32_t total_workgroups = dst->op_params[0] * (uint32_t)src0->ne[3];
+            const uint32_t groups_x = std::min(total_workgroups, ctx->device->properties.limits.maxComputeWorkGroupCount[0]);
+            const uint64_t groups_y_needed = CEIL_DIV(uint64_t(total_workgroups), groups_x);
+            const uint32_t groups_y = std::min<uint64_t>(groups_y_needed, ctx->device->properties.limits.maxComputeWorkGroupCount[1]);
+            const uint64_t groups_xy = uint64_t(groups_x) * groups_y;
+            elements = { groups_x, groups_y, uint32_t(CEIL_DIV(uint64_t(total_workgroups), groups_xy)) };
+            GGML_ASSERT(elements[2] <= ctx->device->properties.limits.maxComputeWorkGroupCount[2]);
        } break;
    case GGML_OP_DIAG_MASK_INF:
        elements = { (uint32_t)ggml_nrows(src0), (uint32_t)ne00, 1 };
@@ -9871,6 +9917,16 @@ static void ggml_vk_op_f32(ggml_backend_vk_context * ctx, vk_context& subctx, co
    case GGML_OP_GLU:
    case GGML_OP_CONV_2D_DW:
        {
+            // dispatched per 32x32 tile, not per element
+            if (op == GGML_OP_CONCAT &&
+                (pipeline == ctx->device->pipeline_concat_transpose_i8 ||
+                 pipeline == ctx->device->pipeline_concat_transpose_i16 ||
+                 pipeline == ctx->device->pipeline_concat_transpose_i32)) {
+                const uint32_t tiles_y = CEIL_DIV(dst->ne[0], 32);
+                const uint32_t groups_y = std::min(tiles_y, ctx->device->properties.limits.maxComputeWorkGroupCount[1]);
+                elements = { (uint32_t)src1->ne[1], groups_y * 32, CEIL_DIV(tiles_y, groups_y) };
+                break;
+            }
            uint32_t ne = ggml_nelements(dst);
            if (op == GGML_OP_CPY && ggml_is_quantized(src0->type) && ggml_is_quantized(dst->type)) {
                // Convert from number of logical elements to 2- or 4-byte units.
@@ -10927,8 +10983,9 @@ void ggml_vk_group_norm(ggml_backend_vk_context * ctx, vk_context& subctx, const
    const uint32_t num_groups = int_op_params[0];
    const float eps = float_op_params[1];
    const uint32_t group_size = src0->ne[0] * src0->ne[1] * ((src0->ne[2] + num_groups - 1) / num_groups);
+    const uint32_t num_workgroups = num_groups * src0->ne[3];

-    ggml_vk_op_f32<vk_op_push_constants>(ctx, subctx, src0, nullptr, nullptr, nullptr, dst, GGML_OP_GROUP_NORM, { group_size, 0, eps, 0.0f, 0.0f, 0.0f });
+    ggml_vk_op_f32<vk_op_push_constants>(ctx, subctx, src0, nullptr, nullptr, nullptr, dst, GGML_OP_GROUP_NORM, { group_size, num_workgroups, eps, 0.0f, 0.0f, 0.0f });
 }

 static uint32_t ggml_vk_rms_num_partials(ggml_backend_vk_context * ctx, const ggml_tensor *node) {
@@ -11186,7 +11243,7 @@ void ggml_vk_rms_norm(ggml_backend_vk_context * ctx, vk_context& subctx, const s

 void ggml_vk_rms_norm_back(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) {
    float * op_params = (float *)dst->op_params;
-    ggml_vk_op_f32<vk_op_push_constants>(ctx, subctx, src0, src1, nullptr, nullptr, dst, GGML_OP_RMS_NORM_BACK, { (uint32_t)src0->ne[0], (uint32_t)src0->ne[1], op_params[0], 0.0f, 0.0f, 0.0f });
+    ggml_vk_op_f32<vk_op_push_constants>(ctx, subctx, src0, src1, nullptr, nullptr, dst, GGML_OP_RMS_NORM_BACK, { (uint32_t)src0->ne[0], (uint32_t)ggml_nrows(src0), op_params[0], 0.0f, 0.0f, 0.0f });
 }

 void ggml_vk_l2_norm(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, ggml_tensor * dst) {
@@ -11850,6 +11907,7 @@ void ggml_vk_cross_entropy_loss(ggml_backend_vk_context * ctx, vk_context& subct
    sp.n_cols = nrows;
    sp.ne01 = 1;
    sp.ne02 = 1;
+    sp.ne03 = 1;
    sp.weight = 1.0f;
    init_pushconst_fastdiv(sp);
    sp.misalign_offsets = get_misalign_bytes(ctx, dst) / ggml_type_size(dst->type);
@@ -0,0 +1,39 @@
+#version 450
+
+#include "types.glsl"
+#include "generic_binary_head.glsl"
+
+#define TILE_DIM 32
+layout(local_size_x = TILE_DIM, local_size_y = 8, local_size_z = 1) in;
+
+// 33 columns (not 32) pads away shared-memory bank conflicts.
+shared A_TYPE tile[TILE_DIM][TILE_DIM + 1];
+
+void main() {
+    const uint tid_x = gl_LocalInvocationID.x;
+    const uint tid_y = gl_LocalInvocationID.y;
+
+    // workgroup x walks rows (src1's contiguous dim1) so consecutive workgroups stream src1, y walks dst columns
+    const uint row0 = gl_WorkGroupID.x * TILE_DIM;
+    const uint col0 = (gl_WorkGroupID.z * gl_NumWorkGroups.y + gl_WorkGroupID.y) * TILE_DIM;
+
+    const uint i1 = row0 + tid_x;
+    [[unroll]] for (uint j = 0; j < TILE_DIM; j += gl_WorkGroupSize.y) {
+        const uint i0 = col0 + tid_y + j;
+        if (i0 >= p.ne00 && i0 - p.ne00 < p.ne10 && i1 < p.ne11) {
+            tile[tid_y + j][tid_x] = A_TYPE(data_b[get_boffset() + (i0 - p.ne00)*p.nb10 + i1*p.nb11]);
+        }
+    }
+
+    barrier();
+
+    const uint d0 = col0 + tid_x;
+    [[unroll]] for (uint j = 0; j < TILE_DIM; j += gl_WorkGroupSize.y) {
+        const uint d1 = row0 + tid_y + j;
+        if (d0 < p.ne20 && d1 < p.ne21) {
+            data_d[get_doffset() + d1*p.nb21 + d0*p.nb20] = d0 < p.ne00
+                ? D_TYPE(data_a[get_aoffset() + d1*p.nb01 + d0*p.nb00])
+                : D_TYPE(tile[tid_x][tid_y + j]);
+        }
+    }
+}
@@ -25,6 +25,8 @@ void main() {
    const uint row = gl_WorkGroupID.z * 262144 + gl_WorkGroupID.y * 512 + gl_WorkGroupID.x;
    const uint tid = gl_LocalInvocationID.x;

+    if (row >= p.ne01 * p.ne02 * p.ne03) { return; }
+
    const uint i03 = fastdiv(row, p.ne0_12mp, p.ne0_12L);
    const uint i03_offset = i03 * p.ne01*p.ne02;
    const uint i02 = fastdiv(row - i03_offset, p.ne0_1mp, p.ne0_1L);
@@ -18,8 +18,12 @@ void main() {
    const float eps = p.param1;

    const uint tid = gl_LocalInvocationID.x;
-    const uint start = gl_WorkGroupID.x * group_size + tid;
-    const uint end = (gl_WorkGroupID.x + 1) * group_size;
+    const uint group = gl_WorkGroupID.x + gl_WorkGroupID.y * gl_NumWorkGroups.x + gl_WorkGroupID.z * gl_NumWorkGroups.x * gl_NumWorkGroups.y;
+    if (group >= p.KY) {
+        return;
+    }
+    const uint start = group * group_size + tid;
+    const uint end = (group + 1) * group_size;

    tmp[tid] = 0.0f;

@@ -14,6 +14,8 @@ void main() {
    const uint row = gl_WorkGroupID.z * 262144 + gl_WorkGroupID.y * 512 + gl_WorkGroupID.x;
    const uint tid = gl_LocalInvocationID.x;

+    if (row >= p.ne01 * p.ne02 * p.ne03) { return; }
+
    const uint a_base = get_aoffset() + src0_idx(row * p.ne00);
    const uint d_base = get_doffset() + dst_idx(row * p.ne10);

@@ -14,6 +14,8 @@ void main() {
    const uint row = gl_WorkGroupID.z * 262144 + gl_WorkGroupID.y * 512 + gl_WorkGroupID.x;
    const uint tid = gl_LocalInvocationID.x;

+    if (row >= p.ne01 * p.ne02 * p.ne03) { return; }
+
    const uint a_base = get_aoffset() + src0_idx(row * p.ne00);
    const uint d_base = get_doffset() + dst_idx(row * p.ne10);

@@ -19,6 +19,8 @@ void main() {
    const uint row = gl_WorkGroupID.z * 262144 + gl_WorkGroupID.y * 512 + gl_WorkGroupID.x;
    const uint tid = gl_LocalInvocationID.x;

+    if (row >= p.KY) { return; }
+
    // Compute derivative of x[i]/norm(x) = g[i]/norm(x) - x[i] dot(x,g)/KX / norm(x)^1.5

    // partial sums for thread in warp
@@ -19,6 +19,8 @@ void main() {
    const uint col = gl_LocalInvocationID.x;
    const float weight = p.weight;

+    if (row >= p.ne01 * p.ne02 * p.ne03) { return; }
+
    const uint i03 = fastdiv(row, p.ne0_12mp, p.ne0_12L);
    const uint i03_offset = i03 * p.ne01*p.ne02;
    const uint i02 = fastdiv(row - i03_offset, p.ne0_1mp, p.ne0_1L);
@@ -5,7 +5,7 @@
 layout (push_constant) uniform parameter
 {
    uint n_cols;
-    uint ne01, ne02;
+    uint ne01, ne02, ne03;
    uint nb01, nb02, nb03;
    uint nb11, nb12, nb13;
    float weight;
@@ -969,6 +969,9 @@ void process_shaders() {
    string_to_spv("concat_i16", "concat.comp", {{"A_TYPE", "uint16_t"}, {"B_TYPE", "uint16_t"}, {"D_TYPE", "uint16_t"}});
    string_to_spv("concat_i32", "concat.comp", {{"A_TYPE", "uint"}, {"B_TYPE", "uint"}, {"D_TYPE", "uint"}});
    string_to_spv("concat_i64", "concat.comp", {{"A_TYPE", "uvec2"}, {"B_TYPE", "uvec2"}, {"D_TYPE", "uvec2"}});
+    string_to_spv("concat_transpose_i8", "concat_transpose.comp", {{"A_TYPE", "uint8_t"}, {"B_TYPE", "uint8_t"}, {"D_TYPE", "uint8_t"}});
+    string_to_spv("concat_transpose_i16", "concat_transpose.comp", {{"A_TYPE", "uint16_t"}, {"B_TYPE", "uint16_t"}, {"D_TYPE", "uint16_t"}});
+    string_to_spv("concat_transpose_i32", "concat_transpose.comp", {{"A_TYPE", "uint"}, {"B_TYPE", "uint"}, {"D_TYPE", "uint"}});

    string_to_spv("upscale_f32", "upscale.comp", {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}});

@@ -4910,6 +4910,9 @@ static bool ggml_backend_webgpu_device_supports_op(ggml_backend_dev_t dev, const
                    default:
                        break;
                }
+                if (ggml_prec_get_src(op, 1) == GGML_PREC_BF16 || ggml_prec_get_src(op, 1) == GGML_PREC_F32) {
+                    supports_op = false;
+                }
                break;
            }
        case GGML_OP_MUL_MAT_ID:
@@ -4953,7 +4956,7 @@ static bool ggml_backend_webgpu_device_supports_op(ggml_backend_dev_t dev, const
                default:
                    break;
            }
-            if (ggml_get_op_params_i32(op, 3) == GGML_PREC_BF16 || ggml_get_op_params_i32(op, 3) == GGML_PREC_F32) {
+            if (ggml_prec_get_src(op, 1) == GGML_PREC_BF16 || ggml_prec_get_src(op, 1) == GGML_PREC_F32) {
                supports_op = false;
            }
            break;
@@ -119,6 +119,10 @@ static bool ggml_zdnn_supports_op(const ggml_backend_zdnn_device_context * ctx_d

                const int64_t max_batch = ctx_dev->max_size;

+                if (ggml_prec_get_src(op, 1) == GGML_PREC_BF16 || ggml_prec_get_src(op, 1) == GGML_PREC_F32) {
+                    return false;
+                }
+
                if (!ggml_is_matrix(weights) || !ggml_is_matrix(inputs) ||
                    !ggml_is_contiguous(weights) || !ggml_is_contiguous(inputs) ||
                    weights->view_src != nullptr ||
@@ -720,6 +720,10 @@ static bool ggml_backend_zendnn_device_supports_op(ggml_backend_dev_t dev, const
                return false;
            }

+            if (ggml_prec_get_src(op, 1) == GGML_PREC_BF16 || ggml_prec_get_src(op, 1) == GGML_PREC_F32) {
+                return false;
+            }
+
            if (ggml_zendnn_adaptive_fallback_enabled()) {
                const int64_t K = inputs->ne[0];
                const int64_t N = (inputs->ne[1]*inputs->ne[2]*inputs->ne[3]);
@@ -7044,6 +7044,7 @@ class VisionProjectorType:
    KIMIVL = "kimivl"
    PADDLEOCR = "paddleocr"
    KIMIK25 = "kimik25"
+    KIMIK31 = "kimik31"
    LIGHTONOCR = "lightonocr"
    COGVLM = "cogvlm"
    JANUS_PRO = "janus_pro"
@@ -2282,6 +2282,7 @@ class TensorNameMap:
            "vision_tower.post_trunk_norm", # dots.ocr
            "vision_encoder.post_trunk_norm", # dots3note
            "vit.perceive.after_rms", # HunyuanVL
+            "mm_projector.post_norm", # Kimi-K3.1
        ),

        MODEL_TENSOR.V_MM_INP_PROJ: (
@@ -783,12 +783,21 @@ static int llama_graph_n_input_tensors(ggml_cgraph * gf) {
    return (int) users.size();
 }

-void llama_context::sched_new(size_t max_nodes, bool parallel) {
+void llama_context::sched_new(const sched_params & params) {
+    // reusing the scheduler avoids reallocating its buffers (incl. the compute buffers)
+    if (sched && params.can_reuse(sched_params_cur)) {
+        LLAMA_LOG_DEBUG("%s: reusing scheduler: max_nodes = %zu (current max_nodes = %zu)\n",
+                __func__, params.max_nodes, sched_params_cur.max_nodes);
+        return;
+    }
+
    sched.reset(ggml_backend_sched_new(
-            backend_ptrs.data(), backend_buft.data(), backend_ptrs.size(), max_nodes, parallel, cparams.op_offload));
+            backend_ptrs.data(), backend_buft.data(), backend_ptrs.size(), params.max_nodes, params.parallel, params.op_offload));

    ggml_backend_sched_set_copy_callback(sched.get(), sched_copy_experts, this);
    ggml_backend_sched_set_realloc_callback(sched.get(), sched_graph_realloc, this);
+
+    sched_params_cur = params;
 }

 void llama_context::sched_reserve() {
@@ -817,7 +826,11 @@ void llama_context::sched_reserve() {
    gf_res_reserve.reset(new llm_graph_result(max_nodes));
    gf_res_prev_active = nullptr;

-    sched_new(max_nodes, cparams.pipeline_parallel);
+    sched_new({
+        /*.max_nodes  =*/ max_nodes,
+        /*.parallel   =*/ cparams.pipeline_parallel,
+        /*.op_offload =*/ cparams.op_offload,
+    });

    llama_memory_context_ptr mctx;
    if (memory) {
@@ -856,7 +869,11 @@ void llama_context::sched_reserve() {
            if (cparams.pipeline_parallel) {
                LLAMA_LOG_WARN("%s: compute buffer allocation failed, retrying without pipeline parallelism\n", __func__);
                cparams.pipeline_parallel = false;
-                sched_new(max_nodes, false);
+                sched_new({
+                    /*.max_nodes  =*/ max_nodes,
+                    /*.parallel   =*/ false,
+                    /*.op_offload =*/ cparams.op_offload,
+                });
                gf = graph_reserve(n_tokens, n_seqs, n_outputs_pp, mctx.get());
            }
            if (!gf) {
@@ -4514,6 +4531,45 @@ llama_context * llama_init_from_model(
        return nullptr;
    }

+    const bool fp8_cache = params.type_k == GGML_TYPE_F8_E4M3 || params.type_v == GGML_TYPE_F8_E4M3;
+    // Direct cache consumers do not use the shared attention scaling helpers.
+    if (fp8_cache && (model->hparams.is_mla() ||
+            model->arch == LLM_ARCH_DEEPSEEK32 ||
+            model->arch == LLM_ARCH_DEEPSEEK4 ||
+            model->arch == LLM_ARCH_DEEPSEEK41 ||
+            model->arch == LLM_ARCH_GLM_DSA ||
+            model->arch == LLM_ARCH_GLM5_NEXT ||
+            model->arch == LLM_ARCH_DOTS3NOTE ||
+            model->arch == LLM_ARCH_MISTRAL_LARGE4 ||
+            model->arch == LLM_ARCH_HY_V4 ||
+            model->arch == LLM_ARCH_MINIMAX_M3 ||
+            model->arch == LLM_ARCH_QWEN4EXP ||
+            model->arch == LLM_ARCH_DFLASH)) {
+        LLAMA_LOG_ERROR("%s: FP8 cache is not supported for architecture '%s'\n", __func__, llm_arch_name(model->arch));
+        return nullptr;
+    }
+
+    if (fp8_cache) {
+        for (uint32_t il = 0; il < model->hparams.n_layer_all; ++il) {
+            if (!model->hparams.has_kv(il)) {
+                continue;
+            }
+
+            for (ggml_tensor * scale : { model->layers[il].k_cache_s, model->layers[il].v_cache_s }) {
+                if (!scale || !scale->data) {
+                    continue;
+                }
+
+                float value;
+                ggml_backend_tensor_get(scale, &value, 0, sizeof(value));
+                if (!std::isfinite(value) || value == 0.0f) {
+                    LLAMA_LOG_ERROR("%s: FP8 cache scale '%s' must be finite and nonzero\n", __func__, ggml_get_name(scale));
+                    return nullptr;
+                }
+            }
+        }
+    }
+
    if (ggml_is_quantized(params.type_v) && params.flash_attn_type != LLAMA_FLASH_ATTN_TYPE_ENABLED) {
        if (params.flash_attn_type == LLAMA_FLASH_ATTN_TYPE_AUTO) {
            LLAMA_LOG_INFO("%s: enabling flash_attn since it is required for quantized V cache\n", __func__);
@@ -297,8 +297,22 @@ private:

    llm_graph_cb graph_get_cb() const;

-    // create a new backend scheduler for a graph of the specified size and register the callbacks below
-    void sched_new(size_t max_nodes, bool parallel);
+    // parameters defining the backend scheduler instance
+    struct sched_params {
+        size_t max_nodes  = 0;
+        bool   parallel   = false;
+        bool   op_offload = false;
+
+        // a scheduler can serve any graph of up to max_nodes nodes
+        bool can_reuse(const sched_params & other) const {
+            return parallel   == other.parallel
+                && op_offload == other.op_offload
+                && max_nodes  <= other.max_nodes;
+        }
+    };
+
+    // create a new backend scheduler with the given parameters, unless the current one can be reused
+    void sched_new(const sched_params & params);

    // ggml_backend_sched copy callback, copies only the experts used by MUL_MAT_ID and updates the MoE cache
    static bool sched_copy_experts(ggml_backend_t backend, const ggml_tensor * src, ggml_tensor * dst, ggml_cgraph * graph, void * user_data);
@@ -390,6 +404,8 @@ private:

    ggml_backend_sched_ptr sched;

+    sched_params sched_params_cur;
+
    bool sched_need_reserve = true;

    // whether the reserve graph has a logits tensor, set in sched_reserve()
@@ -1428,6 +1428,7 @@ void llm_graph_result::reset() {
    t_sampled_probs.clear();
    t_sampled_logits.clear();
    t_candidates.clear();
+    t_sampling_all.clear();

    params = {};

@@ -1476,22 +1477,8 @@ void llm_graph_result::set_outputs(const llm_graph_params & params) {
            }
        }
    }
-    for (auto * tensor : t_sampled) {
-        if (tensor != nullptr) {
-            ggml_set_output(tensor);
-        }
-    }
-    for (auto * tensor : t_sampled_probs) {
-        if (tensor != nullptr) {
-            ggml_set_output(tensor);
-        }
-    }
-    for (auto * tensor : t_sampled_logits) {
-        if (tensor != nullptr) {
-            ggml_set_output(tensor);
-        }
-    }
-    for (auto * tensor : t_candidates) {
+    // the chains are read on the host after compute, so keep every one alive
+    for (auto * tensor : t_sampling_all) {
        if (tensor != nullptr) {
            ggml_set_output(tensor);
        }
@@ -3328,6 +3315,43 @@ llm_graph_input_attn_kv * llm_graph_context::build_attn_inp_kv() const {
    return (llm_graph_input_attn_kv *) res->add_input(std::move(inp));
 }

+static void build_attn_scale_kv_inputs(
+        ggml_context * ctx,
+        const llama_kv_cache_context * mctx,
+        ggml_tensor * & q,
+        ggml_tensor * & k,
+        ggml_tensor * & v,
+        int32_t il) {
+    if (ggml_tensor * scale = mctx->get_k_scale(il)) {
+        // Store K/s and move its dequant scale to Q: (Q*s) dot (K/s) = Q dot K.
+        q = ggml_mul(ctx, q, scale);
+        if (k) {
+            k = ggml_div(ctx, k, scale);
+        }
+    }
+
+    if (v) {
+        // Store V/s here and restore its dequant scale after the weighted sum.
+        ggml_tensor * scale = mctx->get_v_scale(il);
+        if (scale) {
+            v = ggml_div(ctx, v, scale);
+        }
+    }
+}
+
+static ggml_tensor * build_attn_scale_kv_output(
+        ggml_context * ctx,
+        const llama_kv_cache_context * mctx,
+        ggml_tensor * cur,
+        int32_t il) {
+    if (ggml_tensor * scale = mctx->get_v_scale(il)) {
+        // P*(V/s)*s = P*V.
+        cur = ggml_mul(ctx, cur, scale);
+    }
+
+    return cur;
+}
+
 ggml_tensor * llm_graph_context::build_attn(
        llm_graph_input_attn_kv * inp,
        ggml_tensor * wo,
@@ -3354,6 +3378,10 @@ ggml_tensor * llm_graph_context::build_attn(
        v_cur = llama_mul_mat_hadamard(ctx0, v_cur, inp->self_v_rot[il]);
    }

+    const auto * mctx_cur = inp->mctx;
+
+    build_attn_scale_kv_inputs(ctx0, mctx_cur, q_cur, k_cur, v_cur, il);
+
    // these nodes are added to the graph together so that they are not reordered
    // by doing so, the number of splits in the graph is reduced
    // expand k later to enable rope fusion which directly writes into k-v cache
@@ -3361,8 +3389,6 @@ ggml_tensor * llm_graph_context::build_attn(
    ggml_build_forward_expand(gf, v_cur);
    ggml_build_forward_expand(gf, k_cur);

-    const auto * mctx_cur = inp->mctx;
-
    ggml_tensor * q = q_cur;
    ggml_tensor * k;
    ggml_tensor * v;
@@ -3389,6 +3415,7 @@ ggml_tensor * llm_graph_context::build_attn(
    const auto & kv_rows = inp->self_attn_mask.rows;

    ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask, kv_rows, sinks, v_mla, 0, kq_scale, il);
+    cur = build_attn_scale_kv_output(ctx0, mctx_cur, cur, il);
    cb(cur, "kqv_out", il);

    if (inp->self_v_rot[il]) {
@@ -3648,6 +3675,11 @@ ggml_tensor * llm_graph_context::build_attn(
        }
    }

+    const auto * mctx_iswa = inp->mctx;
+    const auto * mctx_cur = is_swa ? mctx_iswa->get_swa() : mctx_iswa->get_base();
+
+    build_attn_scale_kv_inputs(ctx0, mctx_cur, q_cur, k_cur, v_cur, il);
+
    // these nodes are added to the graph together so that they are not reordered
    // by doing so, the number of splits in the graph is reduced
    ggml_build_forward_expand(gf, q_cur);
@@ -3660,10 +3692,6 @@ ggml_tensor * llm_graph_context::build_attn(
        ggml_build_forward_expand(gf, v_cur);
    }

-    const auto * mctx_iswa = inp->mctx;
-
-    const auto * mctx_cur = is_swa ? mctx_iswa->get_swa() : mctx_iswa->get_base();
-
    // whole seq fits into batch in training mode
    const bool use_kv_cur = cparams.training && k_cur && v_cur;
    if (use_kv_cur) {
@@ -3693,6 +3721,7 @@ ggml_tensor * llm_graph_context::build_attn(
    ggml_tensor * v = use_kv_cur ? v_cur : mctx_cur->get_v(ctx0, il);

    ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask, kv_rows, sinks, v_mla, 0, kq_scale, il);
+    cur = build_attn_scale_kv_output(ctx0, mctx_cur, cur, il);
    cb(cur, "kqv_out", il);

    if (v_rot) {
@@ -4507,6 +4536,11 @@ void llm_graph_context::build_sampling() const {
            assert(sampler->iface->backend_apply);
            sampler->iface->backend_apply(sampler, ctx0, gf, &data);

+            if (data.sampled     != nullptr) { res->t_sampling_all.push_back(data.sampled);     }
+            if (data.probs       != nullptr) { res->t_sampling_all.push_back(data.probs);       }
+            if (data.logits      != nullptr) { res->t_sampling_all.push_back(data.logits);      }
+            if (data.candidates  != nullptr) { res->t_sampling_all.push_back(data.candidates);  }
+
            if (data.sampled != nullptr) {
                if (active) {
                    res->t_sampled[row] = data.sampled;
@@ -984,6 +984,8 @@ public:
    std::vector<ggml_tensor *> t_sampled_logits;
    std::vector<ggml_tensor *> t_candidates;

+    std::vector<ggml_tensor *> t_sampling_all;
+
    std::vector<llm_graph_input_ptr> inputs;
    std::vector<llm_graph_fused_node> fused_nodes;

@@ -246,7 +246,7 @@ llama_kv_cache::llama_kv_cache(

        map_layer_ids[il] = layers.size();

-        layers.push_back({ il, k, v, k_stream, v_stream, });
+        layers.push_back({ il, k, v, model.layers[il].k_cache_s, model.layers[il].v_cache_s, k_stream, v_stream, });
    }

    if (reuse) {
@@ -1470,6 +1470,30 @@ ggml_tensor * llama_kv_cache::get_v(ggml_context * ctx, int32_t il, uint32_t n_k
            ggml_row_size(v->type, kv_size*n_embd_v_gqa)*sinfo.s0);
 }

+static bool is_type_scaled(ggml_type type) {
+    return type == GGML_TYPE_F8_E4M3;
+}
+
+ggml_tensor * llama_kv_cache::get_k_scale(int32_t il) const {
+    if (!is_type_scaled(type_k())) {
+        return nullptr;
+    }
+
+    const int32_t ikv = map_layer_ids.at(il);
+
+    return layers[ikv].k_scale;
+}
+
+ggml_tensor * llama_kv_cache::get_v_scale(int32_t il) const {
+    if (!is_type_scaled(type_v())) {
+        return nullptr;
+    }
+
+    const int32_t ikv = map_layer_ids.at(il);
+
+    return layers[ikv].v_scale;
+}
+
 ggml_tensor * llama_kv_cache::cpy_k(ggml_context * ctx, ggml_tensor * k_cur, ggml_tensor * k_idxs, int32_t il, const slot_info & sinfo) const {
    GGML_UNUSED(sinfo);

@@ -2109,7 +2133,7 @@ ggml_tensor * llama_kv_cache::build_rope_shift(
    const int64_t n_embd_head_k = hparams.n_embd_head_k(il);
    const int64_t n_cells       = ggml_nelements(k)/n_embd_k_gqa;

-    if (!ggml_is_quantized(k->type) && !rot && !rp.k_norm) {
+    if (!ggml_is_quantized(k->type) && k->type != GGML_TYPE_F8_E4M3 && !rot && !rp.k_norm) {
        // we rotate only the n_rot dimensions
        ggml_tensor * k_rope =
            ggml_view_3d(ctx, k,
@@ -3049,6 +3073,14 @@ ggml_tensor * llama_kv_cache_context::get_v(ggml_context * ctx, int32_t il) cons
    return kv->get_v(ctx, il, n_kv, sinfos[i_cur]);
 }

+ggml_tensor * llama_kv_cache_context::get_k_scale(int32_t il) const {
+    return kv->get_k_scale(il);
+}
+
+ggml_tensor * llama_kv_cache_context::get_v_scale(int32_t il) const {
+    return kv->get_v_scale(il);
+}
+
 ggml_tensor * llama_kv_cache_context::cpy_k(ggml_context * ctx, ggml_tensor * k_cur, ggml_tensor * k_idxs, int32_t il) const {
    return kv->cpy_k(ctx, k_cur, k_idxs, il, sinfos[i_cur]);
 }
@@ -207,6 +207,9 @@ public:
    ggml_tensor * get_k(ggml_context * ctx, int32_t il, uint32_t n_kv, const slot_info & sinfo) const;
    ggml_tensor * get_v(ggml_context * ctx, int32_t il, uint32_t n_kv, const slot_info & sinfo) const;

+    ggml_tensor * get_k_scale(int32_t il) const;
+    ggml_tensor * get_v_scale(int32_t il) const;
+
    // store k_cur and v_cur in the cache based on the provided head location
    ggml_tensor * cpy_k(ggml_context * ctx, ggml_tensor * k_cur, ggml_tensor * k_idxs, int32_t il, const slot_info & sinfo) const;
    ggml_tensor * cpy_v(ggml_context * ctx, ggml_tensor * v_cur, ggml_tensor * v_idxs, int32_t il, const slot_info & sinfo) const;
@@ -270,6 +273,9 @@ private:
        ggml_tensor * k;
        ggml_tensor * v;

+        ggml_tensor * k_scale;
+        ggml_tensor * v_scale;
+
        std::vector<ggml_tensor *> k_stream;
        std::vector<ggml_tensor *> v_stream;

@@ -420,6 +426,9 @@ public:
    ggml_tensor * get_k(ggml_context * ctx, int32_t il) const;
    ggml_tensor * get_v(ggml_context * ctx, int32_t il) const;

+    ggml_tensor * get_k_scale(int32_t il) const;
+    ggml_tensor * get_v_scale(int32_t il) const;
+
    // store k_cur and v_cur in the cache based on the provided head location
    // note: the heads in k_cur and v_cur should be laid out contiguously in memory
    //   - k_cur  [n_embd_head_k, n_head_k, n_tokens]
@@ -1697,8 +1697,8 @@ bool llama_model_base::load_tensors(llama_model_loader & ml) {
        }
    }

-    // resolve AUTO on systems without mmap support (e.g. iGPUs): fall back to OFF; see #28160
-    if (ml.lazy.mode == LLAMA_LAZY_MODE_AUTO) {
+    // resolve AUTO on systems without mmap support (e.g. iGPUs): fall back to OFF unless mmap is in use; see #28160, #31943
+    if (ml.lazy.mode == LLAMA_LAZY_MODE_AUTO && !ml.use_mmap) {
        for (const auto & dev : devices) {
            ggml_backend_dev_props props;
            ggml_backend_dev_get_props(dev.dev, &props);
@@ -1813,6 +1813,16 @@ bool llama_model_base::load_tensors(llama_model_loader & ml) {
        for (int i = 0; i < n_layer_all; ++i) {
            auto & layer = layers[i];

+            if (hparams.has_kv(i)) {
+                layer.k_cache_s = create_tensor(tn(LLM_TENSOR_ATTN_K, "k_scale", i), {1}, TENSOR_NOT_REQUIRED);
+                layer.v_cache_s = create_tensor(tn(LLM_TENSOR_ATTN_V, "v_scale", i), {1}, TENSOR_NOT_REQUIRED);
+
+                if ((layer.k_cache_s && layer.k_cache_s->type != GGML_TYPE_F32) ||
+                    (layer.v_cache_s && layer.v_cache_s->type != GGML_TYPE_F32)) {
+                    throw std::runtime_error(format("KV cache scales for layer %d must be F32", i));
+                }
+            }
+
            // attention weight scales (per-tensor, shape {1})
            if (!layer.wq_s && layer.wq) {
                layer.wq_s = create_tensor(tn(LLM_TENSOR_ATTN_Q,   "scale", i), {1}, TENSOR_NOT_REQUIRED);
@@ -290,6 +290,8 @@ struct llama_layer {
    struct ggml_tensor * wq        = nullptr;
    struct ggml_tensor * wk        = nullptr;
    struct ggml_tensor * wv        = nullptr;
+    struct ggml_tensor * k_cache_s = nullptr;
+    struct ggml_tensor * v_cache_s = nullptr;
    struct ggml_tensor * wo        = nullptr;
    struct ggml_tensor * wqkv      = nullptr;
    struct ggml_tensor * wg        = nullptr;
@@ -318,7 +318,6 @@ ggml_tensor * llama_model_plamo2::graph::build_plamo2_mamba_layer(llm_graph_inpu
                                                            (d_conv - 1) * (d_inner + 2 * n_group * d_state) * (n_seqs),
                                                            kv_head * (d_conv - 1) * (d_inner + 2 * n_group * d_state) *
                                                                ggml_element_size(conv_states_all))));
-        cb(conv_states_all, "mamba_conv1d_state", il);

        // 1D convolution
        x = ggml_ssm_conv(ctx0, conv_x, model.layers[il].ssm_conv1d);
@@ -390,7 +389,6 @@ ggml_tensor * llama_model_plamo2::graph::build_plamo2_mamba_layer(llm_graph_inpu
                                 n_heads * head_dim * n_seq_tokens * n_seqs * ggml_element_size(y_ssm)),
                    ggml_view_1d(ctx0, ssm_states_all, n_heads * head_dim * d_state * n_seqs,
                                 kv_head * n_seqs * n_heads * head_dim * d_state * ggml_element_size(ssm_states_all))));
-        cb(ssm_states_all, "mamba_ssm_states", il);

        ggml_tensor * y = ggml_view_4d(ctx0, y_ssm, head_dim, n_heads, n_seq_tokens, n_seqs,
                                       head_dim * ggml_element_size(x), head_dim * n_heads * ggml_element_size(x),
@@ -190,6 +190,8 @@ if (NOT WIN32 OR NOT BUILD_SHARED_LIBS)
    llama_build_and_test(test-grammar-integration.cpp  LABEL behavior)
    llama_build_and_test(test-llama-grammar.cpp  LABEL behavior  LINK llama)
    llama_build_and_test(test-batch-alloc.cpp  LABEL behavior  LINK llama)
+    llama_build_and_test(test-kv-cache.cpp LABEL behavior)
+    target_include_directories(test-kv-cache PRIVATE ${PROJECT_SOURCE_DIR}/src)
    llama_build_and_test(test-chat.cpp  LABEL behavior  WORKING_DIRECTORY ${PROJECT_SOURCE_DIR})
    target_include_directories(test-chat PRIVATE ${PROJECT_SOURCE_DIR}/tools/server)
    target_link_libraries(test-chat PRIVATE server-context)
@@ -237,7 +239,9 @@ if (GGML_CPU AND (NOT WIN32 OR NOT BUILD_SHARED_LIBS))
    set_tests_properties(test-generate-fp8-models PROPERTIES FIXTURES_SETUP fp8-models)
    llama_test(test-fp8 NAME test-fp8-requantize LABEL "main behavior backend"
        ARGS --models "${FP8_MODEL_DIR}")
-    set_tests_properties(test-fp8-requantize PROPERTIES FIXTURES_REQUIRED fp8-models)
+    llama_test(test-fp8 NAME test-fp8-kv LABEL "main behavior backend"
+        ARGS --models "${FP8_MODEL_DIR}" --kv)
+    set_tests_properties(test-fp8-requantize test-fp8-kv PROPERTIES FIXTURES_REQUIRED fp8-models)
 endif()

 set(MODEL_DIR "${CMAKE_CURRENT_BINARY_DIR}/test-models/")
@@ -296,6 +300,12 @@ llama_test_cmd(python3
    WORKING_DIRECTORY ${PROJECT_SOURCE_DIR}
    ARGS -m unittest discover -s tests -p test_conversion_fp8.py -v
 )
+llama_test_cmd(python3
+    NAME test-kv-cache-scales
+    LABEL "python behavior"
+    WORKING_DIRECTORY ${PROJECT_SOURCE_DIR}
+    ARGS -m unittest discover -s tests -p test_kv_cache_scales.py -v
+)
 llama_build_and_test(test-json-schema.cpp  LABEL behavior)
 llama_build_and_test(test-chat-auto-parser.cpp  LABEL behavior  WORKING_DIRECTORY ${PROJECT_SOURCE_DIR})
 llama_build_and_test(test-chat-template.cpp  LABEL behavior)
@@ -5677,6 +5677,91 @@ struct test_mul_mat_w4a8 : public test_mul_mat {
    }
 };

+// src1 GGML_PREC_BF16: the implementation may keep src1 in F32 or BF16, but must not convert it down to F16,
+// which cannot hold the src1 values that the BF16 range allows
+struct test_mul_mat_prec_src1_bf16 : public test_mul_mat {
+    test_mul_mat_prec_src1_bf16(ggml_type type_a = GGML_TYPE_F32, int64_t m = 64, int64_t n = 32, int64_t k = 256,
+            std::array<int64_t, 2> bs = {1, 1},
+            std::array<int64_t, 2> nr = {1, 1})
+        : test_mul_mat(type_a, GGML_TYPE_F32, m, n, k, bs, nr) {}
+    ggml_tensor * build_graph(ggml_context * ctx) override {
+        ggml_tensor * out = test_mul_mat::build_graph(ctx);
+        for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) {
+            if (t->op == GGML_OP_MUL_MAT) {
+                ggml_prec_set_src(t, GGML_PREC_BF16, 1);
+            }
+        }
+        return out;
+    }
+    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 (strcmp(t->name, "b") == 0) {
+                init_tensor_uniform(t, -1e6f, 1e6f);
+            } else {
+                init_tensor_uniform(t);
+            }
+        }
+    }
+    std::string op_desc(ggml_tensor * t) override {
+        GGML_UNUSED(t);
+        return "MUL_MAT_PREC_SRC1_BF16";
+    }
+};
+
+// src1 GGML_PREC_F32: the implementation can not convert src1 or the F32 src0 tiles down to F16 or TF32
+struct test_mul_mat_prec_src1_f32 : public test_mul_mat {
+    test_mul_mat_prec_src1_f32(ggml_type type_a = GGML_TYPE_F32, int64_t m = 64, int64_t n = 32, int64_t k = 256,
+            std::array<int64_t, 2> bs = {1, 1},
+            std::array<int64_t, 2> nr = {1, 1})
+        : test_mul_mat(type_a, GGML_TYPE_F32, m, n, k, bs, nr) {}
+    double max_nmse_err() override {
+        return 1e-10;
+    }
+    ggml_tensor * build_graph(ggml_context * ctx) override {
+        ggml_tensor * out = test_mul_mat::build_graph(ctx);
+        for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) {
+            if (t->op == GGML_OP_MUL_MAT) {
+                ggml_prec_set_src(t, GGML_PREC_F32, 1);
+            }
+        }
+        return out;
+    }
+    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 (strcmp(t->name, "b") == 0) {
+                init_tensor_uniform(t, -1e6f, 1e6f);
+            } else {
+                init_tensor_uniform(t);
+            }
+        }
+    }
+    std::string op_desc(ggml_tensor * t) override {
+        GGML_UNUSED(t);
+        return "MUL_MAT_PREC_SRC1_F32";
+    }
+};
+
+// src1 GGML_PREC_BF16 with src0 values outside the F16 range: the F32 src0 tiles must not be downcast to F16
+struct test_mul_mat_prec_src1_bf16_src0 : public test_mul_mat_prec_src1_bf16 {
+    test_mul_mat_prec_src1_bf16_src0(ggml_type type_a = GGML_TYPE_F32, int64_t m = 64, int64_t n = 32, int64_t k = 256,
+            std::array<int64_t, 2> bs = {1, 1},
+            std::array<int64_t, 2> nr = {1, 1})
+        : test_mul_mat_prec_src1_bf16(type_a, m, n, k, bs, nr) {}
+    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 (strcmp(t->name, "a") == 0) {
+                init_tensor_uniform(t, -1e6f, 1e6f);
+            } else {
+                init_tensor_uniform(t);
+            }
+        }
+    }
+    std::string op_desc(ggml_tensor * t) override {
+        GGML_UNUSED(t);
+        return "MUL_MAT_PREC_SRC1_BF16_SRC0";
+    }
+};
+
 // FP4 native W4A4 path (default precision)
 struct test_mul_mat_w4a4 : public test_mul_mat {
    test_mul_mat_w4a4(ggml_type type_a = GGML_TYPE_NVFP4, ggml_type type_b = GGML_TYPE_F32,
@@ -11634,6 +11719,36 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
    test_cases.emplace_back(new test_mul_mat_hadamard(GGML_TYPE_F32, GGML_TYPE_F16, 8192, 1, 8192));
    test_cases.emplace_back(new test_mul_mat_hadamard(GGML_TYPE_F32, GGML_TYPE_F16, 1024, 7, 1024));  // many rows

+    // src1 precision: few-row paths, aligned tiles, partial tiles (K and output), broadcast batches
+    for (ggml_type type_a : {GGML_TYPE_F32, GGML_TYPE_F16, GGML_TYPE_Q8_0}) {
+        // Q8_0 requires K to be a multiple of its block size (32)
+        const int64_t k_partial = type_a == GGML_TYPE_Q8_0 ? 96 : 100;
+        test_cases.emplace_back(new test_mul_mat_prec_src1_bf16(type_a,  64,   1, 256));
+        test_cases.emplace_back(new test_mul_mat_prec_src1_bf16(type_a,  64,   8, 256));
+        test_cases.emplace_back(new test_mul_mat_prec_src1_bf16(type_a,  64,  32, 256));
+        test_cases.emplace_back(new test_mul_mat_prec_src1_bf16(type_a, 128, 512, 512));
+        test_cases.emplace_back(new test_mul_mat_prec_src1_bf16(type_a,  67,  33, k_partial));
+        test_cases.emplace_back(new test_mul_mat_prec_src1_bf16(type_a,  64,  32, 256, {2, 3}, {2, 1}));
+    }
+
+    // src1 precision: same shapes with GGML_PREC_F32, which the backends must not round down to F16 or TF32
+    for (ggml_type type_a : {GGML_TYPE_F32, GGML_TYPE_F16}) {
+        test_cases.emplace_back(new test_mul_mat_prec_src1_f32(type_a,  64,   1, 256));
+        test_cases.emplace_back(new test_mul_mat_prec_src1_f32(type_a,  64,   8, 256));
+        test_cases.emplace_back(new test_mul_mat_prec_src1_f32(type_a,  64,  32, 256));
+        test_cases.emplace_back(new test_mul_mat_prec_src1_f32(type_a, 128, 512, 512));
+        test_cases.emplace_back(new test_mul_mat_prec_src1_f32(type_a,  67,  33, 100));
+        test_cases.emplace_back(new test_mul_mat_prec_src1_f32(type_a,  64,  32, 256, {2, 3}, {2, 1}));
+    }
+
+    // src0 range: an F32 src0 must not be downcast to F16 when src1 requests BF16
+    test_cases.emplace_back(new test_mul_mat_prec_src1_bf16_src0(GGML_TYPE_F32,  64,   1, 256));
+    test_cases.emplace_back(new test_mul_mat_prec_src1_bf16_src0(GGML_TYPE_F32,  64,   8, 256));
+    test_cases.emplace_back(new test_mul_mat_prec_src1_bf16_src0(GGML_TYPE_F32,  64,  32, 256));
+    test_cases.emplace_back(new test_mul_mat_prec_src1_bf16_src0(GGML_TYPE_F32, 128, 512, 512));
+    test_cases.emplace_back(new test_mul_mat_prec_src1_bf16_src0(GGML_TYPE_F32,  67,  33, 100));
+    test_cases.emplace_back(new test_mul_mat_prec_src1_bf16_src0(GGML_TYPE_F32,  64,  32, 256, {2, 3}, {2, 1}));
+
    // FP4 activation precision (default = native W4A4, src1 GGML_PREC_Q8 = W4A8)
    test_cases.emplace_back(new test_mul_mat_w4a8(GGML_TYPE_NVFP4, GGML_TYPE_F32, 32,  1, 256));
    test_cases.emplace_back(new test_mul_mat_w4a8(GGML_TYPE_NVFP4, GGML_TYPE_F32, 32, 32, 256));
@@ -11931,6 +12046,23 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
    // the batches of a are padded, nb[2] is not a multiple of nb[1]
    test_cases.emplace_back(new test_mul_mat(GGML_TYPE_F32, GGML_TYPE_F32, 8, 1, 64, {8, 1}, {1, 1}, {0, 1, 2, 3}, 0, 1, false, 16, 16));
    test_cases.emplace_back(new test_mul_mat(GGML_TYPE_F16, GGML_TYPE_F32, 8, 16, 64, {8, 1}, {1, 1}, {0, 1, 2, 3}, 0, 1, false, 16, 16));
+    // Ragged expert slices, strided routing IDs and broadcast activations.
+    for (ggml_type type : {GGML_TYPE_Q8_0, GGML_TYPE_IQ4_NL, GGML_TYPE_IQ3_S, GGML_TYPE_Q2_K, GGML_TYPE_Q3_K}) {
+        for (int n : {3, 17, 33}) {
+            for (bool broadcast : {false, true}) {
+                test_cases.emplace_back(new test_mul_mat_id(type, GGML_TYPE_F32, 37, 2, broadcast, 65, n, 256));
+            }
+            test_cases.emplace_back(new test_mul_mat_id(type, GGML_TYPE_F32, 37, 2, false, 65, n, 256, 0, false, true));
+        }
+    }
+
+    // Both experts receive every token
+    for (ggml_type type : {GGML_TYPE_Q2_K, GGML_TYPE_Q3_K, GGML_TYPE_Q4_K}) {
+        for (int n : {64, 65}) {
+            test_cases.emplace_back(new test_mul_mat_id(type, GGML_TYPE_F32, 2, 2, false, 33, n, 512));
+        }
+    }
+
    // as is a view whose experts are strided by more rows than it uses
    test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_F32, GGML_TYPE_F32, 4, 2, false, 8,  1, 64, 64));
    test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_F16, GGML_TYPE_F32, 4, 2, false, 8, 16, 64, 64));
@@ -12413,6 +12545,15 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
        test_cases.emplace_back(new test_concat(GGML_TYPE_Q8_0, {8192, 2, 1, 1}, dim == 0 ? 4096 : 1, dim, 0));
    }

+    // GDN conv_input
+    for (int64_t ne_b_d : { 31, 32, 33, 64 }) {
+        test_cases.emplace_back(new test_concat(GGML_TYPE_F32, {3, 10240, 1, 1}, ne_b_d, 0, 16));
+        test_cases.emplace_back(new test_concat(GGML_TYPE_F32, {1, 37, 1, 1}, ne_b_d, 0, 16));
+    }
+    test_cases.emplace_back(new test_concat(GGML_TYPE_F32, {64, 37, 1, 1}, 32, 0, 16));
+    // Empty src0
+    test_cases.emplace_back(new test_concat(GGML_TYPE_F32, {0, 64, 1, 1}, 33, 0, 16));
+
    for (ggml_sort_order order : {GGML_SORT_ORDER_ASC, GGML_SORT_ORDER_DESC}) {
        for (uint32_t i = 4; i <= 1024*1024; i *= 2) {
            test_cases.emplace_back(new test_argsort(GGML_TYPE_F32, {i-1, 1, 1, 1}));
@@ -12776,6 +12917,9 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
    test_cases.emplace_back(new test_flash_attn_ext(128, 64, 4, {1, 1}, 128, 2, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q2_0, GGML_TYPE_Q4_0));
    test_cases.emplace_back(new test_flash_attn_ext(64, 128, 4, {1, 1}, 128, 2, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q4_0, GGML_TYPE_Q2_0));
    test_cases.emplace_back(new test_flash_attn_ext(128, 64, 4, {1, 1}, 64, 2, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q2_0, GGML_TYPE_F16));
+    test_cases.emplace_back(new test_flash_attn_ext(64, 64, 4, {1, 1}, 128, 2, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F8_E4M3, GGML_TYPE_F8_E4M3));
+    test_cases.emplace_back(new test_flash_attn_ext(64, 64, 4, {1, 1}, 128, 2, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F8_E4M3, GGML_TYPE_F16));
+    test_cases.emplace_back(new test_flash_attn_ext(64, 64, 4, {1, 1}, 128, 2, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F8_E4M3));

    // q8_0 KV cases: decode and prompt batches, KV pad, permuted KV, feature flags, and long context
    test_cases.emplace_back(new test_flash_attn_ext(256, 256, 2, {16, 1},   113,   1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0));
@@ -13181,9 +13325,6 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_perf() {
        test_cases.emplace_back(new test_arange(GGML_TYPE_F32, 0.0f, n, 1.0f));
    }

-    // delta-net conv input: conv state followed by the transposed qkv of the ubatch
-    test_cases.emplace_back(new test_concat(GGML_TYPE_F32, {3, 8192, 1, 1}, 512, 0, 16));
-
    // LoRA at prefill: B is [rank, n_out], so K = rank
    for (int64_t k : {16, 32, 48, 64}) {
        test_cases.emplace_back(new test_mul_mat(GGML_TYPE_F16, GGML_TYPE_F32, 2048, 512, k, {1, 1}, {1, 1}));
@@ -13216,6 +13357,13 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_perf() {
        }
    }

+    // GDN conv_input shapes
+    for (int64_t channels : {8192, 10240}) {
+        for (int64_t n_tokens : {512, 2048, 4096}) {
+            test_cases.emplace_back(new test_concat(GGML_TYPE_F32, {3, channels, 1, 1}, n_tokens, 0, 16));
+        }
+    }
+
    // CONT of a 0<->2 permute at DeepSeek-V4 lightning-indexer shapes:
    // indexer_kq is [n_kv, n_tokens, n_head=64] and gets ggml_cont(ggml_permute(.., 2,1,0,3)).
    for (int64_t n_kv : { 1024, 1280, 2048, 2304 }) {
@@ -13452,7 +13600,7 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_perf() {

    // qwen3-30b-a3b
    for (int bs : {1, 4, 8, 32, 64, 128, 256, 512}) {
-        for (ggml_type type_a : {GGML_TYPE_F32, GGML_TYPE_F16, GGML_TYPE_Q4_0, GGML_TYPE_Q8_0, GGML_TYPE_Q4_K, GGML_TYPE_Q6_K, GGML_TYPE_IQ2_XS, GGML_TYPE_IQ4_XS}) {
+        for (ggml_type type_a : {GGML_TYPE_F32, GGML_TYPE_F16, GGML_TYPE_Q4_0, GGML_TYPE_Q8_0, GGML_TYPE_Q2_K, GGML_TYPE_Q3_K, GGML_TYPE_Q4_K, GGML_TYPE_Q6_K, GGML_TYPE_IQ2_XS, GGML_TYPE_IQ4_XS}) {
            for (ggml_type type_b : {GGML_TYPE_F32}) {
                test_cases.emplace_back(new test_mul_mat_id(type_a, type_b, 128, 8, false, 768, bs, 2048));
                test_cases.emplace_back(new test_mul_mat_id_fusion(type_a, type_b, 128, 8, false, 768, bs, 2048, 1));
@@ -2160,6 +2160,12 @@ static void test_lfm2_parser(const std::string & template_path, bool detailed_de
        })
        .run();

+    // Raw control characters in pythonic string arguments must be escaped into valid JSON.
+    tst.test("<|tool_call_start|>[python(code='a\nb')]<|tool_call_end|>")
+        .tools({ python_tool })
+        .expect_tool_calls({ { "python", R"#({"code": "a\nb"})#", "" } })
+        .run();
+
    // String escape sequences decode to their actual characters (newline + tab here),
    // so a "write a two line file" style call produces real line breaks, not literal "\n".
    tst.test("<|tool_call_start|>[python(code=\"First line\\nSecond line\\tindented\")]<|tool_call_end|>")
@@ -2397,6 +2403,18 @@ static void test_template_output_peg_parsers(bool detailed_debug) {
            .expect_content(R"({"amount": 123.45, "date": "2025-12-03"})")
            .run();

+        tst.test(
+               "I need to output the invoice details in JSON\n"
+               "</think>\n\n"
+               R"({"amount": 123.45, "date": "2025-12-03"})")
+            .reasoning_format(COMMON_REASONING_FORMAT_AUTO)
+            .enable_thinking(true)
+            .tools({ get_time_tool })
+            .json_schema(invoice_schema)
+            .expect_reasoning("I need to output the invoice details in JSON")
+            .expect_content(R"({"amount": 123.45, "date": "2025-12-03"})")
+            .run();
+
        // a tool call ends the prefilled thinking block, with or without a closing </think>
        tst.test(
               "<tool_call>\n"
@@ -2741,6 +2759,18 @@ static void test_template_output_peg_parsers(bool detailed_debug) {
            .expect_content(R"({"amount": 123.45, "date": "2025-12-03"})")
            .run();

+        tst.test(
+               "[THINK]I need to output the invoice details in JSON[/THINK]"
+               "```json\n"
+               R"({"amount": 123.45, "date": "2025-12-03"})"
+               "\n```")
+            .reasoning_format(COMMON_REASONING_FORMAT_AUTO)
+            .tools({ special_function_tool })
+            .json_schema(invoice_schema)
+            .expect_reasoning("I need to output the invoice details in JSON")
+            .expect_content(R"({"amount": 123.45, "date": "2025-12-03"})")
+            .run();
+
        // fake tool call marker in reasoning
        tst.test(
               "[THINK]Let me think about [TOOL_CALLS]special_function[ARGS]{\"arg1\":1} and more[/THINK]"
@@ -4326,6 +4356,17 @@ static void test_template_output_peg_parsers(bool detailed_debug) {
            .expect(message_assist)
            .run();

+        tst.test(
+               "```json\n"
+               R"({"amount": 123.45, "date": "2025-12-03"})"
+               "\n```")
+            .enable_thinking(false)
+            .reasoning_format(COMMON_REASONING_FORMAT_DEEPSEEK)
+            .tools({ get_time_tool })
+            .json_schema(invoice_schema)
+            .expect_content(R"({"amount": 123.45, "date": "2025-12-03"})")
+            .run();
+
        // Thinking + content
        tst.test("I'm\nthinking</think>Hello, world!\nWhat's up?")
            .enable_thinking(true)
@@ -7,8 +7,12 @@
 #include "gguf.h"
 #include "llama.h"
 #include "llama-cpp.h"
+#include "../src/llama-context.h"
+#include "../src/llama-kv-cache.h"
+#include "../src/llama-model.h"

 #include <algorithm>
+#include <cmath>
 #include <cstdio>
 #include <cstring>
 #include <filesystem>
@@ -286,6 +290,299 @@ static void check_fp8_requantize(
    }
 }

+static void check_fp8_attention(llama_model * model, bool offload_kqv = true) {
+    // TODO: enable CPU KV cache tests with tensor split after the meta backend handles the attention transfers.
+    if (model->split_mode() == LLAMA_SPLIT_MODE_TENSOR && !offload_kqv) {
+        LOG("FP8 attention: tensor split with CPU KV cache is not supported, skipping\n");
+        return;
+    }
+
+    ggml_context_ptr scale_ctx(ggml_init({ 4096, nullptr, true }));
+    ggml_tensor * k_scale = ggml_new_tensor_1d(scale_ctx.get(), GGML_TYPE_F32, 1);
+    ggml_tensor * v_scale = ggml_new_tensor_1d(scale_ctx.get(), GGML_TYPE_F32, 1);
+    ggml_set_name(k_scale, "blk.0.attn_k.k_scale");
+    ggml_set_name(v_scale, "blk.0.attn_v.v_scale");
+    ggml_backend_buffer_ptr scale_buf(ggml_backend_alloc_ctx_tensors_from_buft(scale_ctx.get(), ggml_backend_dev_buffer_type(model->dev_layer(0))));
+    ggml_tensor * old_k_scale = model->layers[0].k_cache_s;
+    ggml_tensor * old_v_scale = model->layers[0].v_cache_s;
+
+    const int64_t head_size = model->hparams.n_embd_head_k();
+    const int64_t n_head = model->hparams.n_head();
+    const int64_t n_tokens = 4;
+    GGML_ASSERT(n_head == model->hparams.n_head_kv());
+    GGML_ASSERT(head_size == model->hparams.n_embd_head_v());
+
+    for (llama_flash_attn_type flash : { LLAMA_FLASH_ATTN_TYPE_DISABLED, LLAMA_FLASH_ATTN_TYPE_ENABLED }) {
+        if (model->split_mode() == LLAMA_SPLIT_MODE_TENSOR && flash == LLAMA_FLASH_ATTN_TYPE_DISABLED) {
+            LOG("FP8 attention: tensor split requires flash attention, skipping disabled cases\n");
+            continue;
+        }
+        for (int scaled : { 1, 2, 3 }) {
+            const float sk = scaled & 1 ? 0.03125f : 1.0f;
+            const float sv = scaled & 2 ? 0.125f : 1.0f;
+            ggml_backend_tensor_set(k_scale, &sk, 0, sizeof(sk));
+            ggml_backend_tensor_set(v_scale, &sv, 0, sizeof(sv));
+            model->layers[0].k_cache_s = scaled & 1 ? k_scale : nullptr;
+            model->layers[0].v_cache_s = scaled & 2 ? v_scale : nullptr;
+
+            std::vector<float> tensors[4];
+            auto params = llama_context_default_params();
+            params.n_ctx = 256;
+            params.n_batch = n_tokens;
+            params.n_ubatch = n_tokens;
+            params.n_threads = 1;
+            params.n_threads_batch = 1;
+            params.type_k = scaled & 1 ? GGML_TYPE_F8_E4M3 : GGML_TYPE_F16;
+            params.type_v = scaled & 2 ? GGML_TYPE_F8_E4M3 : GGML_TYPE_F16;
+            params.flash_attn_type = flash;
+            params.offload_kqv = offload_kqv;
+            if (model->split_mode() == LLAMA_SPLIT_MODE_TENSOR) {
+                ggml_context_ptr ctx(ggml_init({ 8*ggml_tensor_overhead(), nullptr, true }));
+                ggml_tensor * q = ggml_new_tensor_4d(ctx.get(), GGML_TYPE_F32, head_size, n_tokens, n_head, 1);
+                ggml_tensor * k = ggml_new_tensor_4d(ctx.get(), params.type_k, head_size, params.n_ctx, n_head, 1);
+                ggml_tensor * v = ggml_new_tensor_4d(ctx.get(), params.type_v, head_size, params.n_ctx, n_head, 1);
+                ggml_tensor * mask = ggml_new_tensor_4d(ctx.get(), GGML_TYPE_F16, params.n_ctx, params.n_ctx, 1, 1);
+                ggml_tensor * attn = ggml_flash_attn_ext(ctx.get(), q, k, v, mask, 1.0f/sqrtf(head_size), 0.0f, 0.0f);
+                auto * dev = model->dev_layer(0);
+                if (!ggml_backend_dev_supports_op(dev, attn)) {
+                    LOG("FP8 attention: scaled=%d offload=%d: tensor-split flash attention not supported on %s, skipping\n",
+                            scaled, offload_kqv, ggml_backend_dev_name(dev));
+                    continue;
+                }
+            }
+            if (offload_kqv) {
+                // Cache writes cannot fall back from a preallocated device buffer.
+                llama_kv_cache probe(*model, false, model->hparams, params.type_k, params.type_v,
+                        flash == LLAMA_FLASH_ATTN_TYPE_DISABLED, false, true, params.n_ctx, 1,
+                        0, LLAMA_SWA_TYPE_NONE, nullptr, nullptr, nullptr, nullptr);
+                ggml_context_ptr ctx(ggml_init({ 64*ggml_tensor_overhead(), nullptr, true }));
+                llama_kv_cache::slot_info slots;
+                bool supported = true;
+                for (uint32_t il = 0; il < model->hparams.n_layer(); ++il) {
+                    ggml_tensor * k = ggml_new_tensor_3d(ctx.get(), GGML_TYPE_F32, head_size, n_head, n_tokens);
+                    ggml_tensor * v = ggml_new_tensor_3d(ctx.get(), GGML_TYPE_F32, head_size, n_head, n_tokens);
+                    ggml_tensor * k_idxs = ggml_new_tensor_1d(ctx.get(), GGML_TYPE_I64, n_tokens);
+                    ggml_tensor * v_idxs = ggml_new_tensor_1d(ctx.get(), GGML_TYPE_I64,
+                            flash == LLAMA_FLASH_ATTN_TYPE_DISABLED ? n_tokens*head_size*n_head : n_tokens);
+                    auto * dev = model->dev_layer(il);
+                    if (!ggml_backend_dev_supports_op(dev, probe.cpy_k(ctx.get(), k, k_idxs, il, slots)) ||
+                        !ggml_backend_dev_supports_op(dev, probe.cpy_v(ctx.get(), v, v_idxs, il, slots))) {
+                        LOG("FP8 attention: flash=%s scaled=%d offload=1: cache writes not supported on %s, skipping\n",
+                                llama_flash_attn_type_name(flash), scaled, ggml_backend_dev_name(dev));
+                        supported = false;
+                        break;
+                    }
+                }
+                if (!supported) {
+                    continue;
+                }
+            }
+            params.cb_eval_user_data = tensors;
+            params.cb_eval = [](ggml_tensor * t, bool ask, void * user_data) {
+                const char * names[] = { "Qcur-0", "Kcur-0", "Vcur-0", "kqv_out-0" };
+                for (int i = 0; i < 4; ++i) {
+                    if (strcmp(t->name, names[i]) != 0) {
+                        continue;
+                    }
+                    if (!ask) {
+                        GGML_ASSERT(t->type == GGML_TYPE_F32 && ggml_is_contiguous(t));
+                        auto & data = static_cast<std::vector<float> *>(user_data)[i];
+                        data.resize(ggml_nelements(t));
+                        ggml_backend_tensor_get(t, data.data(), 0, data.size()*sizeof(float));
+                    }
+                    return true;
+                }
+                return false;
+            };
+            llama_context_ptr lctx(llama_init_from_model(model, params));
+            GGML_ASSERT(lctx);
+            GGML_ASSERT(lctx->get_cparams().flash_attn == (flash == LLAMA_FLASH_ATTN_TYPE_ENABLED));
+            get_logits(model, lctx.get(), { 1, 2, 3, 4 });
+            llama_synchronize(lctx.get());
+            for (const auto & data : tensors) {
+                GGML_ASSERT(data.size() == (size_t) (head_size*n_head*n_tokens));
+            }
+
+            // Round the original K/V independently of the cache write graph.
+            for (int i : { 1, 2 }) {
+                const float scale = i == 1 ? sk : sv;
+                const ggml_type type = i == 1 ? params.type_k : params.type_v;
+                auto & data = tensors[i];
+                if (type == GGML_TYPE_F8_E4M3) {
+                    for (float & value : data) {
+                        value /= scale;
+                    }
+                    std::vector<uint8_t> quantized(data.size());
+                    ggml_get_type_traits(type)->from_float_ref(data.data(), quantized.data(), data.size());
+                    ggml_get_type_traits(type)->to_float(quantized.data(), data.data(), data.size());
+                    for (float & value : data) {
+                        value *= scale;
+                    }
+                } else {
+                    for (float & value : data) {
+                        value = ggml_fp16_to_fp32(ggml_fp32_to_fp16(value));
+                    }
+                }
+            }
+
+            std::vector<float> expected(tensors[3].size(), 0.0f);
+            for (int64_t token = 0; token < n_tokens; ++token) {
+                for (int64_t head = 0; head < n_head; ++head) {
+                    const int64_t q_offset = (token*n_head + head)*head_size;
+                    std::vector<double> scores(token + 1);
+                    double max_score = -INFINITY;
+                    for (int64_t past = 0; past <= token; ++past) {
+                        const int64_t kv_offset = (past*n_head + head)*head_size;
+                        double dot = 0.0;
+                        for (int64_t d = 0; d < head_size; ++d) {
+                            dot += (double) tensors[0][q_offset + d]*tensors[1][kv_offset + d];
+                        }
+                        scores[past] = dot/sqrt((double) head_size);
+                        max_score = std::max(max_score, scores[past]);
+                    }
+                    double sum = 0.0;
+                    for (double & score : scores) {
+                        score = exp(score - max_score);
+                        sum += score;
+                    }
+                    for (int64_t past = 0; past <= token; ++past) {
+                        const int64_t kv_offset = (past*n_head + head)*head_size;
+                        for (int64_t d = 0; d < head_size; ++d) {
+                            expected[q_offset + d] += scores[past]/sum*tensors[2][kv_offset + d];
+                        }
+                    }
+                }
+            }
+            const double error = nmse(expected, tensors[3]);
+            if (!std::isfinite(error) || error >= 1e-5) {
+                LOG_ERR("FP8 attention: flash=%s scaled=%d offload=%d nmse=%.2e\n", llama_flash_attn_type_name(flash), scaled, offload_kqv, error);
+            }
+            GGML_ASSERT(std::isfinite(error) && error < 1e-5);
+        }
+    }
+    model->layers[0].k_cache_s = old_k_scale;
+    model->layers[0].v_cache_s = old_v_scale;
+}
+
+static void check_fp8_context_shift(llama_model * model) {
+    ggml_context_ptr scale_ctx(ggml_init({ 4096, nullptr, true }));
+    ggml_tensor * scale = ggml_new_tensor_1d(scale_ctx.get(), GGML_TYPE_F32, 1);
+    ggml_backend_buffer_ptr scale_buf(ggml_backend_alloc_ctx_tensors_from_buft(scale_ctx.get(), ggml_backend_cpu_buffer_type()));
+    ggml_tensor * old_scale = model->layers[0].k_cache_s;
+
+    for (float scale_value : { 1.0f, 0.5f }) {
+        ggml_backend_tensor_set(scale, &scale_value, 0, sizeof(scale_value));
+        model->layers[0].k_cache_s = scale;
+
+        auto params = llama_context_default_params();
+        params.n_ctx = 256;
+        params.n_batch = 4;
+        params.n_ubatch = 4;
+        params.n_threads = 1;
+        params.n_threads_batch = 1;
+        params.type_k = GGML_TYPE_F8_E4M3;
+        params.type_v = GGML_TYPE_F8_E4M3;
+        params.flash_attn_type = LLAMA_FLASH_ATTN_TYPE_DISABLED;
+        params.offload_kqv = false;
+        llama_context_ptr lctx(llama_init_from_model(model, params));
+        GGML_ASSERT(lctx);
+        get_logits(model, lctx.get(), { 1, 2, 3, 4 });
+        llama_synchronize(lctx.get());
+
+        auto * kv = dynamic_cast<llama_kv_cache *>(llama_get_memory(lctx.get()));
+        GGML_ASSERT(kv && kv->get_can_shift());
+        ggml_tensor * k = kv->get_k_storage(0);
+        std::vector<uint8_t> before(ggml_nbytes(k));
+        ggml_backend_tensor_get(k, before.data(), 0, before.size());
+
+        ggml_context_ptr ctx(ggml_init({ 2*1024*1024, nullptr, false }));
+        const auto & hp = model->hparams;
+        const auto & cp = lctx->get_cparams();
+        ggml_tensor * keys = ggml_new_tensor_3d(ctx.get(), GGML_TYPE_F32, hp.n_embd_head_k(), hp.n_head_kv(), kv->get_size());
+        ggml_get_type_traits(GGML_TYPE_F8_E4M3)->to_float(before.data(), (float *) keys->data, ggml_nelements(keys));
+        ggml_tensor * shifts = ggml_new_tensor_1d(ctx.get(), GGML_TYPE_I32, kv->get_size());
+        ggml_set_i32(shifts, 0);
+        for (int i = 1; i < 4; ++i) {
+            ((int32_t *) shifts->data)[i] = -1;
+        }
+        ggml_tensor * expected_f32 = ggml_rope_ext(ctx.get(), keys, shifts, model->get_rope_factors(cp, 0), hp.n_rot(), hp.rope_type,
+                cp.n_ctx_orig_yarn, model->get_rope_freq_base(cp, 0), model->get_rope_freq_scale(cp, 0),
+                cp.yarn_ext_factor, cp.yarn_attn_factor, cp.yarn_beta_fast, cp.yarn_beta_slow);
+        ggml_cgraph * gf = ggml_new_graph(ctx.get());
+        ggml_build_forward_expand(gf, expected_f32);
+        GGML_ASSERT(ggml_graph_compute_with_ctx(ctx.get(), gf, 1) == GGML_STATUS_SUCCESS);
+        std::vector<uint8_t> expected(before.size());
+        ggml_quantize_chunk(GGML_TYPE_F8_E4M3, (float *) expected_f32->data, expected.data(), 0,
+                ggml_nelements(expected_f32)/keys->ne[0], keys->ne[0], nullptr);
+
+        llama_memory_seq_add(kv, 0, 1, 4, -1);
+        GGML_ASSERT(kv->get_has_shift());
+        GGML_ASSERT(lctx->memory_update(false));
+        llama_synchronize(lctx.get());
+        GGML_ASSERT(!kv->get_has_shift());
+        std::vector<uint8_t> after(before.size());
+        ggml_backend_tensor_get(k, after.data(), 0, after.size());
+        GGML_ASSERT(after != before);
+        GGML_ASSERT(after == expected);
+        lctx.reset();
+    }
+
+    model->layers[0].k_cache_s = old_scale;
+}
+
+enum class test_status {
+    PASS,
+    FAIL,
+    SKIP,
+};
+
+static test_status test_fp8_kv(const std::string & path, const char * target_backend) {
+    ggml_backend_dev_t cpu_devices[] = {nullptr};
+    auto params = llama_model_default_params();
+    params.devices = cpu_devices;
+    params.progress_callback = silent_model_load_progress;
+    {
+        llama_model_ptr model = load_model(path, params);
+        if (model->arch != LLM_ARCH_LLAMA || model->hparams.n_expert > 0) {
+            LOG("FP8 KV: only dense Llama fixtures are supported, skipping %s\n", path.c_str());
+            return test_status::SKIP;
+        }
+        LOG("FP8 KV: CPU attention and context shift\n");
+        check_fp8_attention(model.get());
+        check_fp8_context_shift(model.get());
+    }
+
+    std::vector<ggml_backend_dev_t> devices_meta;
+    for (size_t i = 0; i < ggml_backend_dev_count(); ++i) {
+        auto * dev = ggml_backend_dev_get(i);
+        if (target_backend && strcmp(target_backend, ggml_backend_dev_name(dev)) != 0) {
+            continue;
+        }
+        if (ggml_backend_dev_buffer_type(dev) != ggml_backend_cpu_buffer_type()) {
+            devices_meta.push_back(dev);
+        }
+        if (ggml_backend_dev_type(dev) == GGML_BACKEND_DEVICE_TYPE_CPU) {
+            continue;
+        }
+        ggml_backend_dev_t devices[] = {dev, nullptr};
+        params.devices = devices;
+        LOG("FP8 KV: device=%s\n", ggml_backend_dev_name(dev));
+        llama_model_ptr model = load_model(path, params);
+        check_fp8_attention(model.get());
+        check_fp8_attention(model.get(), false);
+    }
+    if (!target_backend && !devices_meta.empty()) {
+        devices_meta.push_back(nullptr);
+        params.devices = devices_meta.data();
+        params.split_mode = LLAMA_SPLIT_MODE_TENSOR;
+        LOG("FP8 KV: tensor split\n");
+        llama_model_ptr model = load_model(path, params);
+        check_fp8_attention(model.get());
+        check_fp8_attention(model.get(), false);
+    }
+    return test_status::PASS;
+}
+
 static void test_fp8(const std::string & model_path) {
    const auto dir = std::filesystem::temp_directory_path() / ("llama-fp8-" + std::to_string(std::random_device{}()));
    if (!std::filesystem::create_directory(dir)) {
@@ -324,51 +621,93 @@ static void test_fp8(const std::string & model_path) {
    check_fp8_requantize(input, reference, output, reference_output, n_embd, model_params);
 }

+static const char * test_status_str(test_status status) {
+    switch (status) {
+        case test_status::PASS: return "PASS";
+        case test_status::FAIL: return "FAIL";
+        case test_status::SKIP: return "SKIP";
+    }
+    GGML_ABORT("invalid test status");
+}
+
 int main(int argc, char ** argv) {
    if (argc == 2 && (strcmp(argv[1], "-h") == 0 || strcmp(argv[1], "--help") == 0)) {
-        printf("Usage: %s -m model.gguf | --models DIR\n", argv[0]);
+        printf("Usage: %s (-m model.gguf | --models DIR) [--kv [-b backend]]\n", argv[0]);
        return 0;
    }
-    if (argc != 3 || (strcmp(argv[1], "-m") != 0 && strcmp(argv[1], "--model") != 0 && strcmp(argv[1], "--models") != 0)) {
-        fprintf(stderr, "Usage: %s -m model.gguf | --models DIR\n", argv[0]);
+    std::string model_file;
+    std::string models_dir;
+    const char * target_backend = nullptr;
+    bool kv = false;
+    bool valid = true;
+    for (int i = 1; i < argc; ++i) {
+        if ((strcmp(argv[i], "-m") == 0 || strcmp(argv[i], "--model") == 0 || strcmp(argv[i], "--models") == 0) && i + 1 < argc) {
+            if (!model_file.empty() || !models_dir.empty()) {
+                valid = false;
+                break;
+            }
+            const bool directory = strcmp(argv[i], "--models") == 0;
+            if (directory) {
+                models_dir = argv[++i];
+            } else {
+                model_file = argv[++i];
+            }
+        } else if (strcmp(argv[i], "--kv") == 0) {
+            kv = true;
+        } else if ((strcmp(argv[i], "-b") == 0 || strcmp(argv[i], "--backend") == 0) && i + 1 < argc) {
+            target_backend = argv[++i];
+        } else {
+            valid = false;
+            break;
+        }
+    }
+    if (!valid || (model_file.empty() && models_dir.empty()) || (target_backend && !kv)) {
+        fprintf(stderr, "Usage: %s (-m model.gguf | --models DIR) [--kv [-b backend]]\n", argv[0]);
        return 1;
    }
    common_init();
    llama_backend_init();
    int result = 0;
    try {
+        if (target_backend && !ggml_backend_dev_by_name(target_backend)) {
+            throw std::runtime_error("unknown backend device: " + std::string(target_backend));
+        }
        std::vector<std::string> models;
-        if (strcmp(argv[1], "--models") == 0) {
-            if (!std::filesystem::is_directory(argv[2])) {
-                throw std::runtime_error("models directory does not exist: " + std::string(argv[2]));
+        if (!models_dir.empty()) {
+            if (!std::filesystem::is_directory(models_dir)) {
+                throw std::runtime_error("models directory does not exist: " + models_dir);
            }
-            for (const auto & entry : std::filesystem::directory_iterator(argv[2])) {
+            for (const auto & entry : std::filesystem::directory_iterator(models_dir)) {
                if (entry.is_regular_file() && entry.path().extension() == ".gguf") {
                    models.push_back(entry.path().string());
                }
            }
            std::sort(models.begin(), models.end());
            if (models.empty()) {
-                throw std::runtime_error("no .gguf models found in: " + std::string(argv[2]));
+                throw std::runtime_error("no .gguf models found in: " + models_dir);
            }
        } else {
-            models.push_back(argv[2]);
+            models.push_back(model_file);
        }

-        std::vector<bool> passed;
+        std::vector<test_status> results;
        for (const auto & model : models) {
            try {
-                test_fp8(model);
-                passed.push_back(true);
+                if (kv) {
+                    results.push_back(test_fp8_kv(model, target_backend));
+                } else {
+                    test_fp8(model);
+                    results.push_back(test_status::PASS);
+                }
            } catch (const std::exception & err) {
                fprintf(stderr, "FP8 test failed for '%s': %s\n", model.c_str(), err.what());
-                passed.push_back(false);
+                results.push_back(test_status::FAIL);
                result = 1;
            }
        }
        printf("\n| Model | Result |\n|-------|--------|\n");
        for (size_t i = 0; i < models.size(); ++i) {
-            printf("| %s | %s |\n", std::filesystem::path(models[i]).filename().string().c_str(), passed[i] ? "PASS" : "FAIL");
+            printf("| %s | %s |\n", std::filesystem::path(models[i]).filename().string().c_str(), test_status_str(results[i]));
        }
    } catch (const std::exception & err) {
        fprintf(stderr, "FP8 test failed: %s\n", err.what());
@@ -1325,6 +1325,66 @@ static void test_all(const std::string & title, std::function<void(const TestCas
        )"""
    });

+    test({
+        FAILURE,
+        "allOf with an empty integer range",
+        R"""({
+            "allOf": [
+                {"type": "integer", "minimum": 9},
+                {"type": "integer", "maximum": 5}
+            ]
+        })""",
+        ""
+    });
+
+    // NOTE: The number child has a minimum and maximum but the parser does not store bounds for number.
+    // NOTE: So the grammar is an unbounded integer, so the issue's 0-100 isn't enforced. This is an existing limitation of the parser.
+    test({
+        SUCCESS,
+        "allOf with plain integer and number schemas",
+        R"""({
+            "allOf": [
+                {"type": "integer"},
+                {"type": "number", "minimum": 0, "maximum": 100}
+            ]
+        })""",
+        R"""(
+            integral-part ::= [0] | [1-9] [0-9]{0,15}
+            root ::= ("-"? integral-part)
+            space ::= | " " | "\n"{1,2} [ \t]{0,20}
+        )"""
+    });
+
+    test({
+        SUCCESS,
+        "allOf with a single integer schema",
+        R"""({
+            "allOf": [
+                {"type": "integer"}
+            ]
+        })""",
+        R"""(
+            integral-part ::= [0] | [1-9] [0-9]{0,15}
+            root ::= ("-"? integral-part)
+            space ::= | " " | "\n"{1,2} [ \t]{0,20}
+        )"""
+    });
+
+    test({
+        SUCCESS,
+        "allOf with bounded integer and number schemas",
+        R"""({
+            "allOf": [
+                {"type": "integer", "minimum": 0, "maximum": 100},
+                {"type": "number"}
+            ]
+        })""",
+        R"""(
+            root ::= ([0-9] | ([1-8] [0-9] | [9] [0-9]) | "100")
+            space ::= | " " | "\n"{1,2} [ \t]{0,20}
+        )"""
+    });
+
    test({
        SUCCESS,
        "conflicting names",
@@ -0,0 +1,105 @@
+#include "llama-kv-cache.h"
+#include "llama-model.h"
+
+#include <numeric>
+#include <string>
+
+static void init_model(llama_model & model, uint32_t n_layer) {
+    model.hparams = {};
+    model.hparams.n_layer_all = n_layer;
+    model.hparams.n_embd_head_k_full = 4;
+    model.hparams.n_embd_head_v_full = 4;
+    model.hparams.n_head_arr.fill(1);
+    model.hparams.n_head_kv_arr.fill(1);
+    std::iota(model.hparams.kv_source_impl.begin(), model.hparams.kv_source_impl.end(), 0);
+    model.layers.resize(n_layer);
+}
+
+int main() {
+    llama_backend_init();
+
+    {
+        std::string errors;
+        llama_log_set([](ggml_log_level level, const char * text, void * user_data) {
+            if (level == GGML_LOG_LEVEL_ERROR) {
+                *static_cast<std::string *>(user_data) += text;
+            }
+        }, &errors);
+
+        for (llm_arch arch : { LLM_ARCH_DEEPSEEK2, LLM_ARCH_DEEPSEEK32, LLM_ARCH_DEEPSEEK4,
+                LLM_ARCH_DEEPSEEK41, LLM_ARCH_GLM_DSA, LLM_ARCH_GLM5_NEXT, LLM_ARCH_DOTS3NOTE,
+                LLM_ARCH_MISTRAL_LARGE4, LLM_ARCH_HY_V4, LLM_ARCH_MINIMAX_M3, LLM_ARCH_QWEN4EXP,
+                LLM_ARCH_DFLASH }) {
+            std::unique_ptr<llama_model> model(llama_model_create(arch, llama_model_default_params()));
+            init_model(*model, 1);
+            if (arch == LLM_ARCH_DEEPSEEK2) {
+                model->hparams.n_embd_head_k_mla_impl = 4;
+                model->hparams.n_embd_head_v_mla_impl = 4;
+            }
+
+            for (ggml_type type_k : { GGML_TYPE_F16, GGML_TYPE_F8_E4M3 }) {
+                for (ggml_type type_v : { GGML_TYPE_F16, GGML_TYPE_F8_E4M3 }) {
+                    if (type_k == GGML_TYPE_F16 && type_v == GGML_TYPE_F16) {
+                        continue;
+                    }
+                    auto params = llama_context_default_params();
+                    params.type_k = type_k;
+                    params.type_v = type_v;
+                    errors.clear();
+                    llama_context * ctx = llama_init_from_model(model.get(), params);
+                    GGML_ASSERT(ctx == nullptr);
+                    if (type_k == type_v || (!model->hparams.is_mla() && arch != LLM_ARCH_DEEPSEEK4 && arch != LLM_ARCH_DEEPSEEK41)) {
+                        GGML_ASSERT(errors.find("FP8 cache is not supported for architecture") != std::string::npos);
+                    } else {
+                        GGML_ASSERT(errors.find("model does not support different K") != std::string::npos);
+                    }
+                }
+            }
+        }
+
+        llama_log_set(nullptr, nullptr);
+    }
+
+    {
+        std::unique_ptr<llama_model> target(llama_model_create(LLM_ARCH_GEMMA4, llama_model_default_params()));
+        std::unique_ptr<llama_model> assistant(llama_model_create(LLM_ARCH_GEMMA4_ASSISTANT, llama_model_default_params()));
+        init_model(*target, 3);
+        init_model(*assistant, 2);
+
+        ggml_context_ptr ctx(ggml_init({ 4096, nullptr, false }));
+        target->layers[2].k_cache_s = ggml_new_tensor_1d(ctx.get(), GGML_TYPE_F32, 1);
+        target->layers[2].v_cache_s = ggml_new_tensor_1d(ctx.get(), GGML_TYPE_F32, 1);
+        assistant->layers[0].k_cache_s = ggml_new_tensor_1d(ctx.get(), GGML_TYPE_F32, 1);
+        assistant->layers[0].v_cache_s = ggml_new_tensor_1d(ctx.get(), GGML_TYPE_F32, 1);
+        assistant->layers[1].k_cache_s = assistant->layers[0].k_cache_s;
+        assistant->layers[1].v_cache_s = assistant->layers[0].v_cache_s;
+        ggml_set_f32(target->layers[2].k_cache_s, 0.5f);
+        ggml_set_f32(target->layers[2].v_cache_s, 0.25f);
+        ggml_set_f32(assistant->layers[0].k_cache_s, 2.0f);
+        ggml_set_f32(assistant->layers[0].v_cache_s, 4.0f);
+
+        for (ggml_type type_k : { GGML_TYPE_F16, GGML_TYPE_F8_E4M3 }) {
+            for (ggml_type type_v : { GGML_TYPE_F16, GGML_TYPE_F8_E4M3 }) {
+                llama_kv_cache source(*target, false, target->hparams, type_k, type_v,
+                        false, false, true, 32, 1, 0, LLAMA_SWA_TYPE_NONE, nullptr, nullptr,
+                        [](uint32_t il) { return il == 0 ? 2 : -1; }, nullptr);
+                llama_kv_cache shared(*assistant, false, assistant->hparams, type_k, type_v,
+                        false, false, true, 32, 1, 0, LLAMA_SWA_TYPE_NONE, &source, nullptr, nullptr,
+                        [](int32_t il) { return il == 0 ? 2 : 1; });
+
+                GGML_ASSERT(source.get_k_scale(2) == (type_k == GGML_TYPE_F8_E4M3 ? target->layers[2].k_cache_s : nullptr));
+                GGML_ASSERT(source.get_v_scale(2) == (type_v == GGML_TYPE_F8_E4M3 ? target->layers[2].v_cache_s : nullptr));
+                GGML_ASSERT(source.get_k_scale(0) == source.get_k_scale(2));
+                GGML_ASSERT(source.get_v_scale(0) == source.get_v_scale(2));
+                GGML_ASSERT(shared.get_k_storage(0) == source.get_k_storage(2));
+                GGML_ASSERT(shared.get_k_scale(0) == source.get_k_scale(2));
+                GGML_ASSERT(shared.get_v_scale(0) == source.get_v_scale(2));
+                GGML_ASSERT(shared.get_k_scale(1) == nullptr);
+                GGML_ASSERT(shared.get_v_scale(1) == nullptr);
+            }
+        }
+    }
+
+    llama_backend_free();
+    return 0;
+}
@@ -0,0 +1,33 @@
+#include "ggml-backend.h"
+#include "ggml-rpc.h"
+#include "ggml.h"
+
+#include <cstdio>
+
+// Regression test: an unreachable or malformed RPC endpoint must be reported as an
+// error instead of aborting the calling process.
+//
+// Before the fix, ggml_backend_rpc_add_server() called GGML_ABORT (SIGABRT) when the
+// endpoint could not be parsed, the transport failed to initialize, the socket could
+// not connect, or the HELLO handshake failed - so a typo in an endpoint or a node
+// that was still booting killed the whole application, and this test would die
+// instead of failing cleanly.
+//
+// After the fix it returns nullptr.
+//
+// ref: https://github.com/ggml-org/llama.cpp/pull/30415
+
+int main(void) {
+    ggml_backend_load_all();
+
+    // 1) nothing listens on this port: the connect must fail, not abort
+    ggml_backend_reg_t reg = ggml_backend_rpc_add_server("127.0.0.1:1");
+    GGML_ASSERT(reg == nullptr);
+
+    // 2) a malformed endpoint must take the parse-failure path and also return nullptr
+    reg = ggml_backend_rpc_add_server("this-is-not-an-endpoint");
+    GGML_ASSERT(reg == nullptr);
+
+    printf("ok\n");
+    return 0;
+}
@@ -0,0 +1,57 @@
+"""Test KV cache scales with synthetic tensors; no GPU or model downloads required.
+
+Install dependencies and run from the repository root:
+    python -m pip install -r requirements/requirements-convert_hf_to_gguf.txt
+    python -m pip install -e ./gguf-py
+    python -m unittest discover -s tests -p 'test_kv_cache_scales.py' -v
+
+Use unittest discovery so the repository root is available for imports.
+"""
+
+import tempfile
+import unittest
+from pathlib import Path
+
+import torch
+
+from conversion.base import ModelBase
+import gguf
+
+
+class TestKVCacheScales(unittest.TestCase):
+    def test_projection_names(self):
+        cases = (
+            (gguf.MODEL_ARCH.PHI3, "model.layers.2.self_attn.qkv_proj", "model.layers.2.self_attn.qkv_proj"),
+            (gguf.MODEL_ARCH.GPT2, "transformer.h.2.attn.c_attn", "transformer.h.2.attn.c_attn"),
+            (gguf.MODEL_ARCH.LLAMA, "model.layers.2.self_attn.k_proj", "model.layers.2.self_attn.v_proj"),
+        )
+        for arch, k_name, v_name in cases:
+            for dtype in (torch.float32, torch.bfloat16):
+                with self.subTest(arch=arch, dtype=dtype), tempfile.TemporaryDirectory() as tmp:
+                    path = Path(tmp) / "scales.gguf"
+                    model = object.__new__(ModelBase)
+                    model.tensor_map = gguf.get_tensor_name_map(arch, 3)
+                    model.gguf_writer = gguf.GGUFWriter(path, gguf.MODEL_ARCH_NAMES[arch])
+                    model.model_tensors = {
+                        k_name + ".k_scale": lambda: torch.tensor(0.5, dtype=dtype),
+                        v_name + ".v_scale": lambda: torch.tensor(1.0, dtype=dtype),
+                    }
+                    model._prepare_kv_cache_scales()
+                    self.assertFalse(model.model_tensors)
+                    model.gguf_writer.write_header_to_file()
+                    model.gguf_writer.write_kv_data_to_file()
+                    model.gguf_writer.write_tensors_to_file()
+                    model.gguf_writer.close()
+
+                    # Release the reader's memory map before Windows removes the temporary directory.
+                    self._check_scales(path)
+
+    def _check_scales(self, path):
+        reader = gguf.GGUFReader(path)
+        self.assertEqual([t.name for t in reader.tensors], ["blk.2.attn_k.k_scale", "blk.2.attn_v.v_scale"])
+        self.assertEqual([t.tensor_type for t in reader.tensors], [gguf.GGMLQuantizationType.F32] * 2)
+        self.assertEqual([float(t.data[0]) for t in reader.tensors], [0.5, 1.0])
+
+
+if __name__ == "__main__":
+    unittest.main()
@@ -500,6 +500,9 @@ static ggml_type ggml_type_from_name(const std::string & s) {
    if (s == "bf16") {
        return GGML_TYPE_BF16;
    }
+    if (s == "f8_e4m3") {
+        return GGML_TYPE_F8_E4M3;
+    }
    if (s == "q8_0") {
        return GGML_TYPE_Q8_0;
    }
@@ -23,18 +23,21 @@ Overview of the `libmtmd` structure:
 - `mtmd_context` (`mtmd.cpp`): the public handle
    - `mtmd_models`: one mmproj file, loaded by `mtmd_model_load()`
        - `mtmd_model` `v` / `a` / `gen_a` (`mtmd-model.h`): one per modality, holds hparams, weights (`ctx_data` + `buf`) and requirements
-    - `mtmd_subcontext` `ctx_v` / `ctx_a` / `ctx_gen_a` (`mtmd-subcontext.h`): one per model, holds a const ref to it
-        - backends + `sched`: GPU/CPU backends, compute buffers, resolved flash attn
+    - `mtmd_backends` (`mtmd-subcontext.h`): GPU/CPU backend instances, CPU threadpool and one scheduler, shared by the subcontexts
+    - `compute_mutex`: held by every public function that runs a compute pass, so only one pass at a time runs on the backends
+    - `mtmd_subcontext` `ctx_v` / `ctx_a` / `ctx_gen_a` (`mtmd-subcontext.h`): one per model, holds a const ref to it and a ref to the backends
+        - `buf_compute_meta` and the resolved flash attn of its graphs
        - `mtmd_graph` (`mtmd-graph.h`, `models/`): built for each `encode()` call, temporary
    - `mtmd_image_preprocessor` / `mtmd_audio_preprocessor` (`mtmd-image.h`, `mtmd-audio.h`): read the hparams of their model
 - `mtmd_gen_state` / `mtmd_audio_stream`: public, one per sequence / stream
    - `mtmd_subctx_gen_state` / `mtmd_subctx_enc_state`: allocated on the backend of a subcontext, see `gen_state_init()` / `enc_state_init()`

 Rules of thumb:
- `mtmd_model` is read-only once loaded. Anything that depends on the backend (scheduler, compute buffers, flash attn resolution, warmup) belongs to `mtmd_subcontext`
+- `mtmd_model` is read-only once loaded. Anything that depends on the backend (flash attn resolution, warmup) belongs to `mtmd_subcontext`
+- Like the pp/tg graphs of `llama_context`, the subcontexts share one scheduler: each warmup reserves its worst-case graph on it, so the compute buffers fit the largest one, and switching modality re-plans inside them
 - A model declares its runtime needs (e.g. `max_nodes`, `flash_attn_auto`) in `load_hparams`, the subcontext applies them
- Models must outlive their subcontexts, `mtmd_context` declares `models` first for that reason
- Backends are owned by each subcontext for now; they will move to `mtmd_context` so that subcontexts can share them
+- Models and backends must outlive the subcontexts, `mtmd_context` declares them first for that reason
+- Locking happens only at the public API boundary (`mtmd.cpp`): a new public function that runs a graph must take `compute_mutex`. The internals (subcontexts, preprocessors, `mtmd_audio_log_mel()`) assume it is held and never lock

 ## Pipeline

@@ -40,7 +40,9 @@ ggml_cgraph * mtmd_graph_kimik25::build() {
    ggml_set_name(pos_w, "pos_w");
    ggml_set_input(pos_w);

-    ggml_tensor * learned_pos_embd = resize_position_embeddings_3d(GGML_SCALE_MODE_BICUBIC);
+    const bool is_k31 = proj_type == PROJECTOR_TYPE_KIMIK31;
+
+    ggml_tensor * learned_pos_embd = resize_position_embeddings_3d(is_k31 ? GGML_SCALE_MODE_BILINEAR : GGML_SCALE_MODE_BICUBIC);

    // Kimi-K2.5 uses interleaved 2D RoPE pattern natively, but
    // Q / K are permuted during conversion to use split format.
@@ -57,14 +59,25 @@ ggml_cgraph * mtmd_graph_kimik25::build() {

    ggml_tensor * cur = build_vit(
                            inp, n_patches,
-                            NORM_TYPE_NORMAL,
+                            is_k31 ? NORM_TYPE_RMS : NORM_TYPE_NORMAL,
                            hparams.ffn_op,
                            nullptr,
                            add_pos);

    cb(cur, "vit_out", -1);

-    {
+    if (is_k31) {
+        // patchmergerv2: no pre-norm, RMS norm after the MLP
+        cur = build_patch_merge_permute(cur, hparams.n_merge);
+        cur = build_ffn(cur,
+            model.mm_1_w, nullptr,
+            nullptr, nullptr,
+            model.mm_2_w, nullptr,
+            FFN_GELU_ERF,
+            -1);
+        cur = build_norm(cur, model.mm_post_norm_w, nullptr, NORM_TYPE_RMS, 1e-5f, -1);
+        cb(cur, "proj_out", -1);
+    } else {
        // patch_merger
        const int scale_factor = model.hparams.n_merge;
        cur = build_patch_merge_permute(cur, scale_factor);
@@ -88,7 +101,7 @@ ggml_cgraph * mtmd_graph_kimik25::build() {
            model.mm_1_w, model.mm_1_b,
            nullptr, nullptr,
            model.mm_2_w, model.mm_2_b,
-            FFN_GELU,
+            FFN_GELU_ERF,
            -1);

        cb(cur, "proj_out", -1);
@@ -885,6 +885,7 @@ std::unique_ptr<mtmd_graph> mtmd_get_graph_builder(mtmd_subcontext * ctx, const
                builder = std::make_unique<mtmd_graph_paddleocr>(ctx, img);
            } break;
        case PROJECTOR_TYPE_KIMIK25:
+        case PROJECTOR_TYPE_KIMIK31:
            {
                builder = std::make_unique<mtmd_graph_kimik25>(ctx, img);
            } break;
@@ -818,6 +818,45 @@ mtmd_image_preproc_out mtmd_image_preprocessor_dyn_size::preprocess(const mtmd_i
    return output;
 }

+//
+// mtmd_image_preprocessor_kimik25
+//
+
+mtmd_image_preproc_out mtmd_image_preprocessor_kimik25::preprocess(const mtmd_image_u8 & img) const {
+    GGML_ASSERT(hparams.image_max_pixels > 0);
+    // patch_limit_on_one_side from the preprocessor config, same for K2.5 and K3.1
+    const int     side_limit     = 512;
+    const int     ps             = hparams.patch_size;
+    const int64_t in_patch_limit = hparams.image_max_pixels / (ps * ps);
+    const int     width          = img.get_size().width;
+    const int     height         = img.get_size().height;
+
+    const double s1 = std::sqrt((double) in_patch_limit
+        / (std::max(1.0, (double) (width / ps)) * std::max(1.0, (double) (height / ps))));
+    const double s2 = (double) side_limit * ps / width;
+    const double s3 = (double) side_limit * ps / height;
+    const double scale = std::min({1.0, s1, s2, s3});
+    const int new_w = std::min(std::max(1, (int) (width  * scale)), side_limit * ps);
+    const int new_h = std::min(std::max(1, (int) (height * scale)), side_limit * ps);
+
+    mtmd_image_u8 content;
+    if (new_w == width && new_h == height) {
+        content = img;
+    } else {
+        img_tool::resize(img, content, mtmd_image_size{new_w, new_h}, hparams.image_resize_algo, PAD_NONE);
+    }
+
+    const int align = ps * hparams.n_merge;
+    mtmd_image_u8 canvas;
+    canvas.set_size(mtmd_image_size{MTMD_ALIGN(new_w, align), MTMD_ALIGN(new_h, align)}, img.is_placeholder());
+    img_tool::fill(canvas, {0, 0, 0});
+    img_tool::composite(canvas, content, 0, 0);
+
+    mtmd_image_preproc_out output;
+    output.append(hparams, canvas, true);
+    return output;
+}
+
 //
 // mtmd_image_preprocessor_qwen3vl
 //
@@ -143,6 +143,12 @@ struct mtmd_image_preprocessor_dyn_size : mtmd_image_preprocessor {
    virtual mtmd_image_size target_size(const mtmd_image_size & original_size) const { return get_target_size(hparams, original_size); }
 };

+// Kimi-K2.5 / K3.1 navit_resize_image: scale down to the patch limits, then pad bottom-right with black to a multiple of patch_size*n_merge
+struct mtmd_image_preprocessor_kimik25 : mtmd_image_preprocessor {
+    mtmd_image_preprocessor_kimik25(const mtmd_model * model) : mtmd_image_preprocessor(model) {}
+    mtmd_image_preproc_out preprocess(const mtmd_image_u8 & img) const override;
+};
+
 // Qwen3-VL, same as dyn_size for images, video frames are sized as Qwen3VLVideoProcessor
 struct mtmd_image_preprocessor_qwen3vl : mtmd_image_preprocessor_dyn_size {
    using mtmd_image_preprocessor_dyn_size::mtmd_image_preprocessor_dyn_size;
@@ -548,6 +548,7 @@ enum projector_type {
    PROJECTOR_TYPE_YOUTUVL,
    PROJECTOR_TYPE_YASA2,
    PROJECTOR_TYPE_KIMIK25,
+    PROJECTOR_TYPE_KIMIK31,
    PROJECTOR_TYPE_NEMOTRON_V2_VL,
    PROJECTOR_TYPE_NEMOTRON_OMNI,
    PROJECTOR_TYPE_HUNYUANVL,
@@ -634,6 +635,7 @@ static std::map<projector_type, std::string> PROJECTOR_TYPE_NAMES = {
    { PROJECTOR_TYPE_YOUTUVL,           "youtuvl"},
    { PROJECTOR_TYPE_YASA2,             "yasa2"},
    { PROJECTOR_TYPE_KIMIK25,           "kimik25"},
+    { PROJECTOR_TYPE_KIMIK31,           "kimik31"},
    { PROJECTOR_TYPE_NEMOTRON_V2_VL,    "nemotron_v2_vl"},
    { PROJECTOR_TYPE_NEMOTRON_OMNI,     "nemotron_omni"},
    { PROJECTOR_TYPE_EXAONE4_5,         "exaone4_5"},
@@ -546,6 +546,7 @@ struct mtmd_model_loader {
                        hparams.set_warmup_n_tokens(256); // avoid OOM on warmup
                    } break;
                case PROJECTOR_TYPE_KIMIK25:
+                case PROJECTOR_TYPE_KIMIK31:
                    {
                        hparams.image_resize_algo = RESIZE_ALGO_BICUBIC;
                        hparams.rope_theta = 10000.0f;
@@ -561,6 +562,9 @@ struct mtmd_model_loader {
                        } else {
                            hparams.set_limit_image_tokens(2, 4096);
                        }
+                        if (model.proj_type == PROJECTOR_TYPE_KIMIK31) {
+                            hparams.set_warmup_n_tokens(46*46); // avoid OOM on warmup
+                        }
                    } break;
                case PROJECTOR_TYPE_DEEPSEEK4V:
                    {
@@ -2143,6 +2147,12 @@ struct mtmd_model_loader {
                    model.mm_2_w = get_tensor(string_format(TN_LLAVA_PROJ, 2, "weight"));
                    model.mm_2_b = get_tensor(string_format(TN_LLAVA_PROJ, 2, "bias"));
                } break;
+            case PROJECTOR_TYPE_KIMIK31:
+                {
+                    model.mm_1_w         = get_tensor(string_format(TN_LLAVA_PROJ, 1, "weight"));
+                    model.mm_2_w         = get_tensor(string_format(TN_LLAVA_PROJ, 2, "weight"));
+                    model.mm_post_norm_w = get_tensor(string_format(TN_MM_POST_NORM, "weight"));
+                } break;
            case PROJECTOR_TYPE_DEEPSEEK4V:
                {
                    model.mm_1_w = get_tensor(string_format(TN_LLAVA_PROJ, 1, "weight"));
@@ -3937,6 +3947,7 @@ int mtmd_model::n_output_tokens(const mtmd_image_f32 * img) const {
        case PROJECTOR_TYPE_D1OMNI_V:
        case PROJECTOR_TYPE_KIMIVL:
        case PROJECTOR_TYPE_KIMIK25:
+        case PROJECTOR_TYPE_KIMIK31:
            {
                // dynamic size
                int out_patch_size = params.patch_size * hparams.n_merge;
@@ -4278,6 +4289,7 @@ int mtmd_model::n_mmproj_embd() const {
        case PROJECTOR_TYPE_KIMIVL:
        case PROJECTOR_TYPE_PADDLEOCR:
        case PROJECTOR_TYPE_KIMIK25:
+        case PROJECTOR_TYPE_KIMIK31:
        case PROJECTOR_TYPE_YASA2:
        case PROJECTOR_TYPE_DEEPSEEK4V:
        case PROJECTOR_TYPE_DEEPSEEK41V:
@@ -32,7 +32,7 @@ static void mtmd_backend_cpu_set_n_threads(ggml_backend_t backend_cpu, int n_thr
    }
 }

-std::shared_ptr<ggml_threadpool> mtmd_threadpool_init(int n_threads) {
+static std::shared_ptr<ggml_threadpool> mtmd_threadpool_init(int n_threads) {
    ggml_backend_dev_t dev = ggml_backend_dev_by_type(GGML_BACKEND_DEVICE_TYPE_CPU);
    ggml_backend_reg_t reg = dev ? ggml_backend_dev_backend_reg(dev) : nullptr;
    if (!reg || n_threads < 2) {
@@ -63,14 +63,9 @@ ggml_backend_dev_t mtmd_select_device(bool use_gpu, ggml_backend_dev_t device) {
    return dev ? dev : ggml_backend_dev_by_type(GGML_BACKEND_DEVICE_TYPE_IGPU);
 }

-mtmd_subcontext::mtmd_subcontext(const mtmd_model & model, const mtmd_subctx_params & params) :
-        model(model),
-        threadpool(params.threadpool),
-        max_nodes(model.max_nodes),
-        flash_attn_type(params.flash_attn_type) {
-    if (flash_attn_type == MTMD_FLASH_ATTN_TYPE_AUTO && !model.flash_attn_auto) {
-        flash_attn_type = MTMD_FLASH_ATTN_TYPE_DISABLED;
-    }
+mtmd_backends::mtmd_backends(ggml_backend_dev_t device, int n_threads, int max_nodes,
+                             ggml_backend_sched_eval_callback cb_eval, void * cb_eval_user_data) {
+    threadpool = mtmd_threadpool_init(n_threads);

    backend_cpu = ggml_backend_init_by_type(GGML_BACKEND_DEVICE_TYPE_CPU, nullptr);
    if (!backend_cpu) {
@@ -83,57 +78,70 @@ mtmd_subcontext::mtmd_subcontext(const mtmd_model & model, const mtmd_subctx_par
            set_threadpool_fn(backend_cpu, threadpool.get());
        }
    }
-    if (params.device != nullptr) {
-        backend = ggml_backend_dev_init(params.device, nullptr);
+    if (device != nullptr) {
+        backend = ggml_backend_dev_init(device, nullptr);
        if (!backend) {
            ggml_backend_free(backend_cpu);
            throw std::runtime_error(string_format("%s: failed to initialize \"%s\" backend\n",
-                                                   __func__, ggml_backend_dev_name(params.device)));
+                                                   __func__, ggml_backend_dev_name(device)));
        }
    }

    if (backend) {
        LOG_INF("%s: using %s backend\n", __func__, ggml_backend_name(backend));
-        backend_ptrs.push_back(backend);
-        backend_buft.push_back(ggml_backend_get_default_buffer_type(backend));
+        ptrs.push_back(backend);
+        bufts.push_back(ggml_backend_get_default_buffer_type(backend));
    } else {
        backend = backend_cpu;
        LOG_INF("%s: using CPU backend\n", __func__);
    }

-    backend_ptrs.push_back(backend_cpu);
-    backend_buft.push_back(ggml_backend_get_default_buffer_type(backend_cpu));
+    ptrs.push_back(backend_cpu);
+    bufts.push_back(ggml_backend_get_default_buffer_type(backend_cpu));

-    sched.reset(
-        ggml_backend_sched_new(backend_ptrs.data(), backend_buft.data(), backend_ptrs.size(), max_nodes, false, true)
-    );
+    sched.reset(ggml_backend_sched_new(ptrs.data(), bufts.data(), ptrs.size(), max_nodes, false, true));
+    if (cb_eval != nullptr) {
+        ggml_backend_sched_set_eval_callback(sched.get(), cb_eval, cb_eval_user_data);
+    }
+}

-    if (params.cb_eval != nullptr) {
-        ggml_backend_sched_set_eval_callback(sched.get(), params.cb_eval, params.cb_eval_user_data);
+mtmd_backends::~mtmd_backends() {
+    // before the backends it uses
+    sched.reset();
+    ggml_backend_free(backend);
+    if (backend != backend_cpu) {
+        ggml_backend_free(backend_cpu);
+    }
+}
+
+std::map<ggml_backend_dev_t, size_t> mtmd_backends::get_mem_compute() const {
+    std::map<ggml_backend_dev_t, size_t> res;
+    for (ggml_backend_t b : ptrs) {
+        res[ggml_backend_get_device(b)] += ggml_backend_sched_get_buffer_size(sched.get(), b);
+    }
+    return res;
+}
+
+mtmd_subcontext::mtmd_subcontext(const mtmd_model & model, mtmd_backends & backends, const mtmd_subctx_params & params) :
+        model(model),
+        backends(backends),
+        max_nodes(model.max_nodes),
+        flash_attn_type(params.flash_attn_type) {
+    if (flash_attn_type == MTMD_FLASH_ATTN_TYPE_AUTO && !model.flash_attn_auto) {
+        flash_attn_type = MTMD_FLASH_ATTN_TYPE_DISABLED;
    }

    debug_output_embeddings = std::getenv("MTMD_DEBUG_EMBEDDINGS") != nullptr;

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

-    // TODO: fix warmup for audio gen
+    auto batch = get_dummy_batch();
    if (model.modality != MTMD_MODALITY_GEN_AUDIO) {
        // check batching support
-        auto batch = get_dummy_batch();
        support_batch = mtmd_get_graph_builder(this, batch)->support_batch();
-
-        if (params.warmup) {
-            warmup(batch);
-        }
    }
-}
-
-mtmd_subcontext::~mtmd_subcontext() {
-    // the scheduler references the backends
-    sched.reset();
-    ggml_backend_free(backend);
-    if (backend != backend_cpu) {
-        ggml_backend_free(backend_cpu);
+    if (params.warmup) {
+        warmup(batch);
    }
 }

@@ -146,6 +154,11 @@ mtmd_image_f32_batch mtmd_subcontext::get_dummy_batch() const {
        const int sz = hparams.warmup_image_size;
        img.set_size({sz, sz}, false, false);
        LOG_INF("%s: warmup with image size = %d x %d\n", __func__, sz, sz);
+    } else if (model.modality == MTMD_MODALITY_GEN_AUDIO) {
+        // TODO: improve this, move to per-model impl
+        const int n_embd = model.n_mmproj_embd();
+        img.set_size({n_embd, 1}, false, true);
+        LOG_INF("%s: warmup with one gen-audio step\n", __func__);
    } else {
        // GEMMA4UA uses n_mel_bins as a raw-waveform frame size (640), not a mel-bin count,
        // so the [1, 256] bound only applies to FFT-based models.
@@ -170,7 +183,7 @@ void mtmd_subcontext::warmup(const mtmd_image_f32_batch & batch) {
        if (!info.fattn && info.fattn_op) {
            auto op = info.fattn_op;
            LOG_WRN("%s: *****************************************************************\n", __func__);
-            LOG_WRN("%s: WARNING: flash attention not supported by %s, memory usage will increase\n", __func__, ggml_backend_name(backend));
+            LOG_WRN("%s: WARNING: flash attention not supported by %s, memory usage will increase\n", __func__, ggml_backend_name(backends.backend));
            LOG_WRN("%s: op params: \n", __func__);
            static auto print_shape = [](const char * fn, const char * name, ggml_tensor * t) {
                LOG_WRN("%s:   %s: type = %s, ne = [%d %d %d %d], nb = [%d %d %d %d]\n", fn,
@@ -200,7 +213,7 @@ void mtmd_subcontext::warmup(const mtmd_image_f32_batch & batch) {
        (flash_attn_type == MTMD_FLASH_ATTN_TYPE_ENABLED) ? "enabled" : "disabled");

    // print ops that are not supported by the GPU backend (if there is one)
-    if (backend && backend != backend_cpu) {
+    if (backends.backend != backends.backend_cpu) {
        std::vector<support_info_op> unsupported_ops;
        for (const auto & op : info.ops) {
            if (!op.is_accel) {
@@ -211,7 +224,7 @@ void mtmd_subcontext::warmup(const mtmd_image_f32_batch & batch) {
            LOG_WRN("%s: *****************************************************************\n", __func__);
            LOG_WRN("%s: WARNING: the mmproj graph uses unsupported operators by the backend\n", __func__);
            LOG_WRN("%s:          the performance will be suboptimal                      \n", __func__);
-            LOG_WRN("%s:          list of unsupported ops (backend=%s):\n", __func__, ggml_backend_name(backend));
+            LOG_WRN("%s:          list of unsupported ops (backend=%s):\n", __func__, ggml_backend_name(backends.backend));
            for (const auto & op : unsupported_ops) {
                LOG_WRN("%s: %16s: type = %s, ne = [%d %d %d %d]\n", __func__,
                        ggml_op_name(op.op->op),
@@ -229,22 +242,9 @@ void mtmd_subcontext::warmup(const mtmd_image_f32_batch & batch) {

 mtmd_subcontext::support_info_graph mtmd_subcontext::reserve_compute_meta(const mtmd_image_f32_batch & batch) {
    ggml_cgraph * gf = mtmd_get_graph_builder(this, batch)->build();
-    ggml_backend_sched_reserve(sched.get(), gf);
+    ggml_backend_sched_reserve(backends.sched.get(), gf);

-    mem_compute.clear();
-    for (size_t i = 0; i < backend_ptrs.size(); ++i) {
-        ggml_backend_t backend = backend_ptrs[i];
-        ggml_backend_buffer_type_t buft = backend_buft[i];
-        size_t size = ggml_backend_sched_get_buffer_size(sched.get(), backend);
-        if (size > 1) {
-            LOG_INF("%s: %10s compute buffer size = %8.2f MiB\n", __func__,
-                    ggml_backend_buft_name(buft),
-                    size / 1024.0 / 1024.0);
-        }
-        mem_compute[ggml_backend_get_device(backend)] += size;
-    }
-
-    const int n_splits = ggml_backend_sched_get_n_splits(sched.get());
+    const int n_splits = ggml_backend_sched_get_n_splits(backends.sched.get());
    const int n_nodes  = ggml_graph_n_nodes(gf);

    LOG_INF("%s: graph splits = %d, nodes = %d\n", __func__,  n_splits, n_nodes);
@@ -259,7 +259,7 @@ mtmd_subcontext::support_info_graph mtmd_subcontext::reserve_compute_meta(const
    for (int i = 0; i < ggml_graph_n_nodes(gf); i++) {
        ggml_tensor * node = ggml_graph_node(gf, i);
        res.ops.push_back({node, true});
-        if (!ggml_backend_supports_op(backend, node)) {
+        if (!ggml_backend_supports_op(backends.backend, node)) {
            res.ops.back().is_accel = false;
            if (node->op == GGML_OP_FLASH_ATTN_EXT) {
                res.fattn    = false;
@@ -271,46 +271,6 @@ mtmd_subcontext::support_info_graph mtmd_subcontext::reserve_compute_meta(const
    return res;
 }

-// TODO: remove this and use ggml_mul_mat with GGML_PREC_F32 once all backends honor it for F32 x F32
-// ggml_mul_mat for 2D F32 tensors, always in F32 on CPU: some GPU backends round the inputs to F16 or TF32, which loses the quiet bins of a spectrum
-static void mtmd_mul_mat_prec_op(ggml_tensor * dst, int ith, int nth, void * /*userdata*/) {
-    const ggml_tensor * a = dst->src[0];
-    const ggml_tensor * b = dst->src[1];
-    const int64_t n_k = a->ne[0];
-
-    for (int64_t n = ith; n < b->ne[1]; n += nth) {
-        const float * b_row = (const float *) ((const char *) b->data + n * b->nb[1]);
-        float * dst_row = (float *) ((char *) dst->data + n * dst->nb[1]);
-        for (int64_t m = 0; m < a->ne[1]; m++) {
-            const float * a_row = (const float *) ((const char *) a->data + m * a->nb[1]);
-            // independent partial sums, so that the compiler can vectorize the loop
-            float acc[16] = {};
-            int64_t k = 0;
-            for (; k + 16 <= n_k; k += 16) {
-                for (int l = 0; l < 16; l++) {
-                    acc[l] += a_row[k + l] * b_row[k + l];
-                }
-            }
-            float sum = 0.0f;
-            for (; k < n_k; k++) {
-                sum += a_row[k] * b_row[k];
-            }
-            for (int l = 0; l < 16; l++) {
-                sum += acc[l];
-            }
-            dst_row[m] = sum;
-        }
-    }
-}
-
-static ggml_tensor * mtmd_mul_mat_prec(ggml_context * ctx, ggml_tensor * a, ggml_tensor * b) {
-    GGML_ASSERT(a->type == GGML_TYPE_F32 && b->type == GGML_TYPE_F32);
-    GGML_ASSERT(ggml_n_dims(a) <= 2 && ggml_n_dims(b) <= 2 && a->ne[0] == b->ne[0]);
-    GGML_ASSERT(a->nb[0] == sizeof(float) && b->nb[0] == sizeof(float));
-    ggml_tensor * args[] = { a, b };
-    return ggml_custom_4d(ctx, GGML_TYPE_F32, a->ne[1], b->ne[1], 1, 1, args, 2, mtmd_mul_mat_prec_op, GGML_N_TASKS_MAX, nullptr);
-}
-
 bool mtmd_audio_log_mel(const mtmd_audio_mel_config & config, const float * samples, const float * window, int64_t n_frames, float * out) {
    ggml_backend_sched_t sched = config.sched;
    ggml_backend_t backend_cpu = ggml_backend_sched_get_backend(sched, ggml_backend_sched_get_n_backends(sched) - 1);
@@ -356,7 +316,8 @@ bool mtmd_audio_log_mel(const mtmd_audio_mel_config & config, const float * samp
        ggml_set_input(inp_filters);

        ggml_tensor * cur = ggml_cont(ctx0, ggml_view_2d(ctx0, inp_samples, frame_size, n_cur, config.hop_length * ggml_element_size(inp_samples), 0));
-        cur = mtmd_mul_mat_prec(ctx0, inp_dft, cur); // [2 * n_fft_bins, n_cur]
+        cur = ggml_mul_mat(ctx0, inp_dft, cur); // [2 * n_fft_bins, n_cur]
+        ggml_prec_set_src(cur, GGML_PREC_F32, 1);

        cur = ggml_sqr(ctx0, cur);
        cur = ggml_add(ctx0,
@@ -366,7 +327,8 @@ bool mtmd_audio_log_mel(const mtmd_audio_mel_config & config, const float * samp
            cur = ggml_sqrt(ctx0, cur);
        }

-        cur = mtmd_mul_mat_prec(ctx0, cur, inp_filters); // [n_cur, n_mel]
+        cur = ggml_mul_mat(ctx0, cur, inp_filters); // [n_cur, n_mel]
+        ggml_prec_set_src(cur, GGML_PREC_F32, 1);
        cur = config.mel_floor_add
            ? ggml_scale_bias(ctx0, cur, 1.0f, config.mel_floor)
            : ggml_clamp(ctx0, cur, config.mel_floor, FLT_MAX);
@@ -475,7 +437,7 @@ mtmd_subctx_gen_state * mtmd_subcontext::gen_state_init(uint32_t seed) {
        ggml_format_name(t, "gen_state_%s", slot.name.c_str());
        st->slots.push_back(t);
    }
-    st->buf.reset(ggml_backend_alloc_ctx_tensors(st->ctx_data.get(), backend));
+    st->buf.reset(ggml_backend_alloc_ctx_tensors(st->ctx_data.get(), backends.backend));
    if (!st->buf) {
        LOG_ERR("%s: failed to allocate the gen state\n", __func__);
        return nullptr;
@@ -503,7 +465,7 @@ mtmd_subctx_enc_state * mtmd_subcontext::enc_state_init(int32_t n_frames_max) {
            ggml_format_name(t, "enc_state_%s", slot.name.c_str());
            st->slots.push_back(t);
        }
-        st->buf.reset(ggml_backend_alloc_ctx_tensors(st->ctx_data.get(), backend));
+        st->buf.reset(ggml_backend_alloc_ctx_tensors(st->ctx_data.get(), backends.backend));
        if (!st->buf) {
            LOG_ERR("%s: failed to allocate the encoder state\n", __func__);
            return nullptr;
@@ -543,7 +505,7 @@ mtmd_subctx_enc_state * mtmd_subcontext::enc_state_init(int32_t n_frames_max) {
        st->k.push_back(k);
        st->v.push_back(v);
    }
-    st->buf.reset(ggml_backend_alloc_ctx_tensors(st->ctx_data.get(), backend));
+    st->buf.reset(ggml_backend_alloc_ctx_tensors(st->ctx_data.get(), backends.backend));
    if (!st->buf) {
        LOG_ERR("%s: failed to allocate the encoder state\n", __func__);
        return nullptr;
@@ -588,14 +550,14 @@ bool mtmd_subcontext::encode(mtmd_subctx_encode_params * params) {
    }

    // build the inference graph
-    ggml_backend_sched_reset(sched.get());
+    ggml_backend_sched_reset(backends.sched.get());
    ggml_cgraph * gf = mtmd_get_graph_builder(this, imgs, params)->build();

    // the last node is the embedding tensor, code2wav has no out_embd
    // read it before scheduling: backends may reorder the nodes of graphs with several outputs (Metal)
    ggml_tensor * embeddings = params->out_embd ? ggml_graph_node(gf, -1) : nullptr;

-    if (!ggml_backend_sched_alloc_graph(sched.get(), gf)) {
+    if (!ggml_backend_sched_alloc_graph(backends.sched.get(), gf)) {
        LOG_ERR("%s: failed to allocate compute graph\n", __func__);
        return false;
    }
@@ -1362,6 +1324,7 @@ bool mtmd_subcontext::encode(mtmd_subctx_encode_params * params) {
        case PROJECTOR_TYPE_PIXTRAL:
        case PROJECTOR_TYPE_KIMIVL:
        case PROJECTOR_TYPE_KIMIK25:
+        case PROJECTOR_TYPE_KIMIK31:
        case PROJECTOR_TYPE_LIGHTONOCR:
            {
                // set the 2D positions
@@ -2308,9 +2271,9 @@ bool mtmd_subcontext::encode(mtmd_subctx_encode_params * params) {
            GGML_ABORT("Unknown projector type");
    }

-    mtmd_backend_cpu_set_n_threads(backend_cpu, params->n_threads);
+    mtmd_backend_cpu_set_n_threads(backends.backend_cpu, params->n_threads);

-    auto status = ggml_backend_sched_graph_compute(sched.get(), gf);
+    auto status = ggml_backend_sched_graph_compute(backends.sched.get(), gf);
    if (status != GGML_STATUS_SUCCESS) {
        LOG_ERR("%s: ggml_backend_sched_graph_compute failed with error %d\n", __func__, status);
        return false;
@@ -2708,11 +2671,3 @@ bool mtmd_subcontext::encode(mtmd_subctx_encode_params * params) {

    return true;
 }
-
-std::map<ggml_backend_dev_t, size_t> mtmd_subcontext::get_mem_usage() const {
-    std::map<ggml_backend_dev_t, size_t> result = model.mem_weights;
-    for (auto & [dev, size] : mem_compute) {
-        result[dev] += size;
-    }
-    return result;
-}
@@ -16,8 +16,7 @@

 // !!! Internal header, to be used by mtmd only !!!

-// a subcontext runs the graphs of one mtmd_model (one modality)
-// it owns the backends and the compute buffers, the weights are owned by the model
+// runs the graphs of one mtmd_model on the shared backends and scheduler

 enum mtmd_flash_attn_type {
    MTMD_FLASH_ATTN_TYPE_AUTO     = -1,
@@ -26,22 +25,40 @@ enum mtmd_flash_attn_type {
 };

 struct mtmd_subctx_params {
-    ggml_backend_dev_t device; // nullptr means CPU, see mtmd_select_device()
-    std::shared_ptr<ggml_threadpool> threadpool; // CPU threadpool shared by the subcontexts, can be nullptr
    enum mtmd_flash_attn_type flash_attn_type;
    bool warmup;
-    ggml_backend_sched_eval_callback cb_eval;
-    void * cb_eval_user_data;
 };

 // the device the weights and the graphs are placed on, nullptr means CPU
 ggml_backend_dev_t mtmd_select_device(bool use_gpu, ggml_backend_dev_t device);

-// nullptr if n_threads < 2 or the CPU backend has no threadpool support
-std::shared_ptr<ggml_threadpool> mtmd_threadpool_init(int n_threads);
+// shared by the subcontexts, not thread-safe
+struct mtmd_backends {
+    ggml_backend_t backend     = nullptr; // backend_cpu if no accelerator
+    ggml_backend_t backend_cpu = nullptr;
+
+    // for ggml_backend_sched_new(), backend_cpu last
+    std::vector<ggml_backend_t>             ptrs;
+    std::vector<ggml_backend_buffer_type_t> bufts;
+
+    std::shared_ptr<ggml_threadpool> threadpool; // can be nullptr
+
+    // its buffers fit the largest graph reserved by the subcontexts
+    ggml_backend_sched_ptr sched;
+
+    // device nullptr means CPU
+    mtmd_backends(ggml_backend_dev_t device, int n_threads, int max_nodes,
+                  ggml_backend_sched_eval_callback cb_eval, void * cb_eval_user_data);
+    ~mtmd_backends();
+
+    std::map<ggml_backend_dev_t, size_t> get_mem_compute() const;
+
+    mtmd_backends(const mtmd_backends &) = delete;
+    mtmd_backends & operator=(const mtmd_backends &) = delete;
+};

 // compute the log-mel of n_frames frames, the waveform must already be padded and window must hold one whole frame
-// out has the layout [n_frames, n_mel]; like the encoders, must not run at the same time as another pass on config.sched
+// out has the layout [n_frames, n_mel]; must not run at the same time as another pass
 struct mtmd_audio_mel_config;
 bool mtmd_audio_log_mel(const mtmd_audio_mel_config & config, const float * samples, const float * window, int64_t n_frames, float * out);

@@ -134,38 +151,26 @@ struct mtmd_subctx_encode_params {

 struct mtmd_subcontext {
    const mtmd_model & model;
+    mtmd_backends & backends;

    std::vector<uint8_t> buf_compute_meta;

-    std::vector<ggml_backend_t> backend_ptrs;
-    std::vector<ggml_backend_buffer_type_t> backend_buft;
-
-    std::shared_ptr<ggml_threadpool> threadpool; // shared by the subcontexts of a model
-
-    ggml_backend_t backend = nullptr;
-    ggml_backend_t backend_cpu = nullptr;
-
    int max_nodes = 8192;
-    ggml_backend_sched_ptr sched;
    mtmd_flash_attn_type flash_attn_type = MTMD_FLASH_ATTN_TYPE_AUTO;
    bool is_allocated = false;

    bool debug_output_embeddings = false;

-    // for measuring memory usage
-    std::map<ggml_backend_dev_t, size_t> mem_compute;
-
    bool support_batch = false;

-    // the model must outlive the subcontext
-    mtmd_subcontext(const mtmd_model & model, const mtmd_subctx_params & params);
-    ~mtmd_subcontext();
+    // model and backends must outlive the subcontext
+    mtmd_subcontext(const mtmd_model & model, mtmd_backends & backends, const mtmd_subctx_params & params);
+
+    bool encode(mtmd_subctx_encode_params * params);

    // reserve the compute buffers for the given batch
    void warmup(const mtmd_image_f32_batch & batch);

-    bool encode(mtmd_subctx_encode_params * params);
-
    // TODO: remove encode_image() and always use batched version
    bool encode_image(int n_threads, const mtmd_image_f32 * img, std::vector<float> & out_vec);
    bool encode_batch(int n_threads, const mtmd_image_f32_batch * imgs, std::vector<float> & out_batch_embd);
@@ -176,9 +181,6 @@ struct mtmd_subcontext {
    // seed UINT32_MAX means random, free the result with delete
    mtmd_subctx_gen_state * gen_state_init(uint32_t seed);

-    // weights and compute buffers
-    std::map<ggml_backend_dev_t, size_t> get_mem_usage() const;
-
 private:
    struct support_info_op {
        ggml_tensor * op;
@@ -28,6 +28,7 @@
 #include <cstring>
 #include <climits>
 #include <map>
+#include <mutex>
 #include <type_traits>
 #include <vector>

@@ -593,6 +594,10 @@ static bool cbx_is_mtl(const mtmd_model & model) {
 struct mtmd_context {
    // must outlive the subcontexts below
    mtmd_models models;
+    std::unique_ptr<mtmd_backends> backends; // shared by the subcontexts
+
+    // held by every public function that runs a graph, the internals never lock
+    std::mutex compute_mutex;

    std::unique_ptr<mtmd_subcontext> ctx_v; // vision
    std::unique_ptr<mtmd_subcontext> ctx_a; // audio
@@ -635,7 +640,10 @@ struct mtmd_context {
    bool ov_img_first      = false;

    // MiniCPM-V 4.6/4.7 prepends an <image_id>N</image_id> tag before <image>
-    bool use_image_id = false;
+    bool has_prompt_img_id = false;
+
+    // Kimi-K3.1 writes the original image size after img_beg
+    bool has_prompt_img_size = false;

    // string template for slice image delimiters with row/col (idefics3)
    std::string sli_img_start_tmpl;
@@ -706,23 +714,37 @@ struct mtmd_context {
            throw std::runtime_error(string_format("Failed to load mmproj model from %s\n", mmproj_fname));
        }

+        int max_nodes = 0;
+        for (const auto * m : { models.v.get(), models.a.get(), models.gen_a.get() }) {
+            if (m) {
+                max_nodes = std::max(max_nodes, m->max_nodes);
+            }
+        }
+
+        // no_alloc only measures memory, nothing is computed
+        backends = std::make_unique<mtmd_backends>(device, no_alloc ? 1 : ctx_params.n_threads, max_nodes,
+                                                   ctx_params.cb_eval, ctx_params.cb_eval_user_data);
+
+        // like llama_context::sched_reserve(), each warmup reserves its graph on the shared scheduler
        mtmd_subctx_params sctx_params {
-            /* device            */ device,
-            // no_alloc only measures memory, nothing is computed
-            /* threadpool        */ mtmd_threadpool_init(no_alloc ? 1 : ctx_params.n_threads),
            /* flash_attn_type   */ mtmd_to_flash_attn_type(ctx_params.flash_attn_type),
            /* warmup            */ ctx_params.warmup,
-            /* cb_eval           */ ctx_params.cb_eval,
-            /* cb_eval_user_data */ ctx_params.cb_eval_user_data,
        };
        if (models.v) {
-            ctx_v = std::make_unique<mtmd_subcontext>(*models.v, sctx_params);
+            ctx_v = std::make_unique<mtmd_subcontext>(*models.v, *backends, sctx_params);
        }
        if (models.a) {
-            ctx_a = std::make_unique<mtmd_subcontext>(*models.a, sctx_params);
+            ctx_a = std::make_unique<mtmd_subcontext>(*models.a, *backends, sctx_params);
        }
        if (models.gen_a) {
-            ctx_gen_a = std::make_unique<mtmd_subcontext>(*models.gen_a, sctx_params);
+            ctx_gen_a = std::make_unique<mtmd_subcontext>(*models.gen_a, *backends, sctx_params);
+        }
+        for (size_t i = 0; i < backends->ptrs.size(); ++i) {
+            const size_t size = ggml_backend_sched_get_buffer_size(backends->sched.get(), backends->ptrs[i]);
+            if (size > 1) {
+                LOG_INF("%s: %10s compute buffer size = %8.2f MiB\n", __func__,
+                        ggml_backend_buft_name(backends->bufts[i]), size / 1024.0 / 1024.0);
+            }
        }

        // if both vision and audio mmproj are present, we need to validate their n_embd
@@ -830,7 +852,7 @@ struct mtmd_context {
                    tok_row_end       = {lookup_token("\n")};
                    tok_row_end_trail = false; // no trailing end-of-row token
                    ov_img_first      = true;
-                    use_image_id      = true;
+                    has_prompt_img_id = true;
                    if (ctx_v->model.proj_type == PROJECTOR_TYPE_MINICPM_OCR_1_0) {
                        image_preproc = std::make_unique<mtmd_image_preprocessor_minicpm_ocr_1_0>(models.v.get());
                    } else {
@@ -995,12 +1017,21 @@ struct mtmd_context {
                        // <|begin_of_image|> ... (image embeddings) ... <|end_of_image|>
                        img_beg = "<|begin_of_image|>";
                        img_end = "<|end_of_image|>";
+                        image_preproc = std::make_unique<mtmd_image_preprocessor_dyn_size>(models.v.get());
                    } else {
                        // <|media_begin|> ... (image embeddings) ... <|media_end|>
                        img_beg = "<|media_begin|>";
                        img_end = "<|media_end|>";
+                        image_preproc = std::make_unique<mtmd_image_preprocessor_kimik25>(models.v.get());
                    }
-                    image_preproc = std::make_unique<mtmd_image_preprocessor_dyn_size>(models.v.get());
+                } break;
+            case PROJECTOR_TYPE_KIMIK31:
+                {
+                    // <|media_begin|>image WxH<|media_content|> ... (image embeddings) ... <|media_end|>
+                    img_beg = "<|media_begin|>";
+                    img_end = "<|media_end|>";
+                    has_prompt_img_size = true;
+                    image_preproc = std::make_unique<mtmd_image_preprocessor_kimik25>(models.v.get());
                } break;
            case PROJECTOR_TYPE_LIGHTONOCR:
                {
@@ -1249,7 +1280,7 @@ struct mtmd_context {
        // initialize audio preprocessor
        if (audio_preproc) {
            audio_preproc->mel_config.n_threads = n_threads;
-            audio_preproc->mel_config.sched     = ctx_a->sched.get();
+            audio_preproc->mel_config.sched     = backends->sched.get();
            audio_preproc->initialize();
        }
    }
@@ -1622,6 +1653,9 @@ struct mtmd_tokenizer {
            if (!ctx->img_beg.empty()) {
                add_text(ctx->img_beg, true); // add image begin token
            }
+            if (ctx->has_prompt_img_size) {
+                add_text(string_format("image %ux%u<|media_content|>", bitmaps[0]->nx, bitmaps[0]->ny), true);
+            }

            // TODO @ngxson : this is quite hacky because preprocessor only support batch with one single element, that need to be fixed in the future (e.g. by changing the preprocessor interface always take single input)

@@ -1676,7 +1710,7 @@ struct mtmd_tokenizer {

            if (has_tiling_grid && ctx->proj_type_v() == PROJECTOR_TYPE_MINICPMV4_7) {
                GGML_ASSERT(bitmaps.size() == 1);
-                if (ctx->use_image_id) {
+                if (ctx->has_prompt_img_id) {
                    add_text("<image_id>" + std::to_string(n_images_added) + "</image_id>", true);
                }
                add_text(ctx->tok_ov_img_start);
@@ -1701,7 +1735,7 @@ struct mtmd_tokenizer {

                // add overview image (first)
                if (ctx->ov_img_first) {
-                    if (ctx->use_image_id) {
+                    if (ctx->has_prompt_img_id) {
                        add_text("<image_id>" + std::to_string(n_images_added) + "</image_id>", true);
                    }
                    add_text(ctx->tok_ov_img_start);
@@ -2346,6 +2380,7 @@ static int32_t mtmd_encode_chunk_impl(mtmd_context * ctx, const mtmd_input_chunk

 int32_t mtmd_encode_chunk(mtmd_context * ctx, const mtmd_input_chunk * chunk) {
    // this is the non-batching version
+    std::lock_guard<std::mutex> lock(ctx->compute_mutex);
    try {
        return mtmd_encode_chunk_impl(ctx, chunk, ctx->out_embd);
    } catch (const std::exception & e) {
@@ -2355,6 +2390,7 @@ int32_t mtmd_encode_chunk(mtmd_context * ctx, const mtmd_input_chunk * chunk) {
 }

 int32_t mtmd_encode(mtmd_context * ctx, const mtmd_image_tokens * image_tokens) {
+    std::lock_guard<std::mutex> lock(ctx->compute_mutex);
    try {
        return mtmd_encode_impl(ctx, image_tokens, ctx->out_embd);
    } catch (const std::exception & e) {
@@ -3125,6 +3161,7 @@ static int32_t mtmd_gen_audio_process_impl(mtmd_context * ctx, mtmd_gen_state *
 }

 int32_t mtmd_gen_audio_process(mtmd_context * ctx, mtmd_gen_state * state, const struct mtmd_gen_inp * inp, struct mtmd_gen_out * out) {
+    std::lock_guard<std::mutex> lock(ctx->compute_mutex);
    try {
        return mtmd_gen_audio_process_impl(ctx, state, inp, out);
    } catch (const std::exception & e) {
@@ -3233,6 +3270,7 @@ static int32_t mtmd_batch_encode_impl(mtmd_batch * batch) {
 }

 int32_t mtmd_batch_encode(mtmd_batch * batch) {
+    std::lock_guard<std::mutex> lock(batch->ctx->compute_mutex);
    try {
        return mtmd_batch_encode_impl(batch);
    } catch (const std::exception & e) {
@@ -3846,6 +3884,7 @@ static void mtmd_debug_encode_impl(mtmd_context * ctx, mtmd_subcontext * sctx, m
 }

 void mtmd_debug_encode_image(mtmd_context * ctx, const std::vector<std::vector<float>> & image) {
+    std::lock_guard<std::mutex> lock(ctx->compute_mutex);
    if (!ctx->ctx_v) {
        LOG_ERR("%s: model does not support vision input\n", __func__);
        return;
@@ -3864,6 +3903,7 @@ void mtmd_debug_encode_image(mtmd_context * ctx, const std::vector<std::vector<f
 }

 void mtmd_debug_encode_audio(mtmd_context * ctx, const std::vector<float> & input) {
+    std::lock_guard<std::mutex> lock(ctx->compute_mutex);
    if (!ctx->ctx_a) {
        LOG_ERR("%s: model does not support audio input\n", __func__);
        return;
@@ -3908,6 +3948,7 @@ void mtmd_debug_preprocess_image(mtmd_context * ctx, const std::vector<uint8_t>
 }

 void mtmd_debug_preprocess_audio(mtmd_context * ctx, const std::vector<float> & samples) {
+    std::lock_guard<std::mutex> lock(ctx->compute_mutex);
    if (!ctx->ctx_a || !ctx->audio_preproc) {
        LOG_ERR("%s: model does not support audio input\n", __func__);
        return;
@@ -3948,20 +3989,14 @@ mtmd_memory_usage mtmd_get_memory_usage(const char * mmproj_fname,
        mtmd_log_set(stub_log_callback, nullptr); // suppress logging
        ctx.reset(new mtmd_context(mmproj_fname, nullptr, ctx_params, true));
        mtmd_log_set(saved_log_callback, saved_log_user_data); // restore log callback
-        std::map<ggml_backend_dev_t, size_t> total_mem;
-        auto merge = [&](const struct mtmd_subcontext * c) {
-            for (auto & [dev, size] : c->get_mem_usage()) {
-                total_mem[dev] += size;
+        // compute buffers are shared, weights are per model
+        std::map<ggml_backend_dev_t, size_t> total_mem = ctx->backends->get_mem_compute();
+        for (const auto * m : { ctx->models.v.get(), ctx->models.a.get(), ctx->models.gen_a.get() }) {
+            if (m) {
+                for (auto & [dev, size] : m->mem_weights) {
+                    total_mem[dev] += size;
+                }
            }
-        };
-        if (ctx->ctx_v) {
-            merge(ctx->ctx_v.get());
-        }
-        if (ctx->ctx_a) {
-            merge(ctx->ctx_a.get());
-        }
-        if (ctx->ctx_gen_a) {
-            merge(ctx->ctx_gen_a.get());
        }
        mtmd_memory_usage res;
        res.backend_mem_usage = std::move(total_mem);
@@ -4185,10 +4220,12 @@ int32_t mtmd_audio_stream_flush(mtmd_audio_stream * s, bool end) {
    if (s->ended) {
        return 1;
    }
+    std::lock_guard<std::mutex> lock(s->ctx->compute_mutex);
    return s->flush(end);
 }

 int32_t mtmd_audio_stream_step(mtmd_audio_stream * s, llama_token last_token) {
+    std::lock_guard<std::mutex> lock(s->ctx->compute_mutex);
    return s->step(last_token);
 }

@@ -353,6 +353,7 @@ DEPRECATED(MTMD_API int32_t mtmd_encode(mtmd_context * ctx, const mtmd_image_tok
           "use mtmd_encode_chunk() instead");

 // text chunk will be ignored silently, only media chunk will be encoded
+// the output is stored in ctx: do not call it from several threads at once
 // returns 0 on success
 // returns 1 on generic error
 MTMD_API int32_t mtmd_encode_chunk(mtmd_context * ctx,
@@ -377,6 +378,7 @@ MTMD_API void         mtmd_batch_free(mtmd_batch * batch);
 // returns 3 if it cannot be batched with the existing chunks in the batch
 MTMD_API int32_t mtmd_batch_add_chunk(mtmd_batch * batch, const mtmd_input_chunk * chunk);

+// thread-safe, the output is stored in the batch
 // returns 0 on success
 // returns 1 on generic error
 MTMD_API int32_t mtmd_batch_encode(mtmd_batch * batch);
@@ -431,11 +433,13 @@ MTMD_API int32_t mtmd_audio_stream_push(mtmd_audio_stream * stream, const float
 // pad the stream so that all pending rows can be produced
 // end = false: the stream stays open (use it on a pause, to get a partial transcript)
 // end = true:  no more push() is accepted
+// thread-safe across streams
 // returns 0 on success, 1 if the stream has ended
 MTMD_API int32_t mtmd_audio_stream_flush(mtmd_audio_stream * stream, bool end);

 // produce the next decoder rows
 // last_token is the token sampled from the previous row, LLAMA_TOKEN_NULL before the first sample
+// thread-safe across streams
 // return value:
 //   > 0 number of rows ready, read them with mtmd_audio_stream_get_output_embd()
 //   0   more audio is needed
@@ -557,6 +561,7 @@ MTMD_API struct mtmd_gen_inp mtmd_gen_inp_default(const mtmd_context * ctx);

 // run one step and update the state in place
 // the caller drives the backbone and collects the audio frames
+// thread-safe across gen states
 MTMD_API int32_t mtmd_gen_audio_process(mtmd_context * ctx,
                                mtmd_gen_state * state,
                                const struct mtmd_gen_inp * inp,
@@ -13,6 +13,11 @@ if (LLAMA_BUILD_TESTS AND UNIX AND NOT GGML_BACKEND_DL)
            $<TARGET_FILE:ggml-rpc-server>
            $<TARGET_FILE:test-rpc-multi-server>)
    set_property(TEST test-rpc-multi-server PROPERTY LABELS main)
+    add_executable(test-rpc-unreachable ${PROJECT_SOURCE_DIR}/tests/test-rpc-unreachable.cpp)
+    target_link_libraries(test-rpc-unreachable PRIVATE ggml ggml-rpc)
+    target_include_directories(test-rpc-unreachable PRIVATE ${PROJECT_SOURCE_DIR}/ggml/src)
+    add_test(NAME test-rpc-unreachable COMMAND test-rpc-unreachable)
+    set_property(TEST test-rpc-unreachable PROPERTY LABELS main)
 endif()

 if(LLAMA_TOOLS_INSTALL)
@@ -6186,7 +6186,7 @@ void server_routes::init_routes() {
        // live transcription: Voxtral Realtime (no chat), or an "asr" model in a session opened with ?intent=transcription
        const bool live_asr = !ctx_server.asr_prompt.empty() && req.get_param("intent") == "transcription";
        if (ctx_server.mtmd_cap.audio_stream && ctx_server.mtmd_cap.chat && !live_asr) {
-            server_realtime_conversation({queue_tasks, queue_results, *meta, params.sampling, ctx_server.mctx, ctx_server.mtmd_cap, ctx_server.vocab}, conn);
+            server_realtime_conversation({queue_tasks, queue_results, *meta, params, ctx_server.mctx, ctx_server.mtmd_cap, ctx_server.vocab}, conn);
            return;
        }

@@ -7073,8 +7073,10 @@ std::unique_ptr<server_res_generator> server_routes::handle_tts_impl(const serve
    server_schema::eval_tts_schema(params, task, body);
    const bool stream = task.params.stream;
    // a model with a "tts" chat template (e.g. LFM3-Audio) speaks the prompt rendered from it
-    const std::string lang       = json_value(body, "lang", std::string());
-    const std::string voice      = json_value(body, "voice", std::string());
+    const std::string lang       = json_value(body, "lang", params.tts_lang);
+    // speaker_ref in the request overrides the default voice
+    const bool has_speaker_ref   = req.files.count("speaker_ref") > 0 || !json_value(body, "speaker_ref_b64", std::string()).empty();
+    const std::string voice      = json_value(body, "voice", has_speaker_ref ? std::string() : params.tts_voice);
    const std::string tts_prompt = common_chat_get_tts_prompt(meta->chat_params.tmpls.get(), prompt, lang, voice);
    // a model with predefined voices only speaks with one of them
    mtmd::bitmap_ptr voice_ref;
@@ -130,6 +130,8 @@ struct realtime_conversation {
    std::string                    type = "realtime";
    std::string                    instructions;
    std::string                    language; // transcription
+    std::string                    out_language;
+    std::string                    out_voice;
    mtmd_helper_gen_audio_out_mode out_mode          = MTMD_HELPER_GEN_AUDIO_OUT_AUTO;
    int32_t                        max_output_tokens = -1; // "inf"
    int                            input_rate        = 16000;
@@ -176,7 +178,8 @@ struct realtime_conversation {

    std::atomic<bool> closed{false};

-    realtime_conversation(const server_realtime_deps & deps, server_http_context::ws_conn & conn) : deps(deps), conn(conn) {}
+    realtime_conversation(const server_realtime_deps & deps, server_http_context::ws_conn & conn)
+        : deps(deps), conn(conn), out_language(deps.params.tts_lang), out_voice(deps.params.tts_voice) {}

    //
    // events
@@ -207,7 +210,11 @@ struct realtime_conversation {
        }
        json audio = {{"input", audio_in}};
        if (deps.cap.audio_gen) {
-            audio["output"] = {{"format", {{"type", "audio/pcm"}, {"rate", output_rate}}}};
+            audio["output"] = {
+                {"format",   {{"type", "audio/pcm"}, {"rate", output_rate}}},
+                {"voice",    out_voice.empty()    ? json(nullptr) : json(out_voice)},
+                {"language", out_language.empty() ? json(nullptr) : json(out_language)},
+            };
        }
        return {
            {"id",                id},
@@ -245,7 +252,7 @@ struct realtime_conversation {

    // sampling of a response: the server defaults, with the tool call grammar of chat
    common_params_sampling response_sampling(const common_chat_session * chat) const {
-        common_params_sampling sampling = deps.sampling;
+        common_params_sampling sampling = deps.params.sampling;
        if (chat) {
            chat->apply_sampling(sampling);
        }
@@ -544,6 +551,8 @@ struct realtime_conversation {
            return;
        }
        response.out_mode = mode;
+        response.lang     = out_language;
+        response.voice    = out_voice;

        // with tools, the output goes through the tool call parser of the chat template and its grammar
        std::optional<common_chat_session> chat;
@@ -661,6 +670,27 @@ struct realtime_conversation {
        if (in.contains("transcription") && in["transcription"].is_object()) {
            language = json_value(in["transcription"], "language", language);
        }
+        const json out = audio.value("output", json::object());
+        if (out.contains("voice") || out.contains("language")) {
+            std::string voice = out_voice;
+            if (out.contains("voice")) {
+                // string or {"id": name}
+                const json & v = out["voice"];
+                voice = v.is_object() ? json_value(v, "id", std::string()) : v.is_string() ? v.get<std::string>() : std::string();
+            }
+            const std::string lang = json_value(out, "language", out_language);
+            const auto & voices = deps.meta.tts_voices;
+            if (!voice.empty() && !voices.empty() && std::find(voices.begin(), voices.end(), voice) == voices.end()) {
+                send_error("session.update: unknown voice \"" + voice + "\", see \"voices\" in GET /props");
+                return;
+            }
+            if (started && (voice != out_voice || lang != out_language)) {
+                send_error("session.update: the output voice and language cannot change after the first audio or message");
+                return;
+            }
+            out_voice    = voice;
+            out_language = lang;
+        }
        if (!started && !configure()) {
            send_error("the chat template does not fit the realtime turn format", "server_error");
            return;
@@ -674,6 +704,11 @@ struct realtime_conversation {
        if (type == "transcription") {
            const auto preset = common_chat_get_asr_prompt(deps.meta.chat_params.tmpls.get(), language);
            instructions = preset.system;
+        } else if (deps.cap.audio_gen) {
+            const std::string speech = common_chat_get_speech_system_prompt(deps.meta.chat_params.tmpls.get(), out_language, out_voice);
+            if (!speech.empty()) {
+                instructions = instructions.empty() ? speech : speech + "\n" + instructions;
+            }
        }
        const bool ok = derive_turn_text();
        instructions = saved;
@@ -869,7 +904,7 @@ struct realtime_conversation {
        {
            server_task task(SERVER_TASK_TYPE_REALTIME);
            task.id                       = rd.get_new_id();
-            task.params.sampling          = deps.sampling;
+            task.params.sampling          = deps.params.sampling;
            task.params.stream            = true;
            task.params.oaicompat_cmpl_id = id;
            task.params.oaicompat_model   = deps.meta.model_name;
@@ -1052,6 +1087,15 @@ bool server_realtime_slot::response_start(const server_realtime_op & op) {
        inp.data.stream   = true;
        inp.data.out_type = MTMD_HELPER_GEN_AUDIO_OUTTYPE_PCM;
        inp.data.out_mode = op.response.out_mode;
+        inp.set_lang(op.response.lang);
+        // predefined voice -> speaker ref, LFM3 uses the system prompt instead
+        if (!op.response.voice.empty() && mtmd_gen_audio_get_info(mctx).n_voices > 0) {
+            mtmd::bitmap_ptr voice(mtmd_gen_audio_get_voice(mctx, op.response.voice.c_str()));
+            if (!voice) {
+                return fail("unknown voice \"" + op.response.voice + "\"", ERROR_TYPE_INVALID_REQUEST);
+            }
+            inp.set_speaker_ref(std::move(voice));
+        }
        if (gen.set_input(inp.get()) != 0) {
            return fail("failed to start the response");
        }
@@ -15,7 +15,7 @@
 #include <vector>

 struct server_context_meta;
-struct common_params_sampling;
+struct common_params;
 struct mtmd_context;
 struct mtmd_helper_cap;
 struct llama_vocab;
@@ -24,7 +24,7 @@ struct server_realtime_deps {
    server_queue &                 queue_tasks;
    server_response &              queue_results;
    const server_context_meta &    meta;
-    const common_params_sampling & sampling;
+    const common_params &          params;
    mtmd_context *                 mctx;
    const mtmd_helper_cap &        cap;
    const llama_vocab *            vocab;
@@ -70,6 +70,8 @@ struct server_control_response {
    mtmd_helper_gen_audio_out_mode out_mode  = MTMD_HELPER_GEN_AUDIO_OUT_AUTO;
    common_params_sampling         sampling;       // of the text tokens: tool call grammar and preserved tokens
    int32_t                        n_predict = -1; // text tokens and audio frames, -1: until the context is full
+    std::string                    lang;           // empty: model default
+    std::string                    voice;          // empty: model default
 };

 struct task_params {
@@ -1,3 +1,6 @@
 engine-strict=true
 ignore-scripts=true
 min-release-age=7
+# The single-bundle production build needs more than the default heap on CI runners;
+# npm passes this to every script it runs, on every platform
+node-options=--max-old-space-size=6144
--- a/Show More
+++ b/Show More